From 4ac8f104ca5e5dfde1adc05a9dbe8918bc9488cb Mon Sep 17 00:00:00 2001
From: Rudranil Goswami <nilrudra068@gmail.com>
Date: Thu, 2 Apr 2026 23:38:16 +0530
Subject: [PATCH 0001/1000] Improve AutoGraph if-expression error for
 mismatched branches

---
 .../operators/conditional_expressions.py      | 30 +++++++++++++++++--
 .../autograph/tests/basic_ifexp_test.py       | 12 ++++++++
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/autograph/operators/conditional_expressions.py b/tensorflow/python/autograph/operators/conditional_expressions.py
index 28fd328834c..c1a698939a0 100644
--- a/tensorflow/python/autograph/operators/conditional_expressions.py
+++ b/tensorflow/python/autograph/operators/conditional_expressions.py
@@ -15,9 +15,13 @@
 """Conditional expressions (e.g. the ternary if statement)."""
 
 
+import functools
+
 from tensorflow.python.autograph.operators import control_flow
 from tensorflow.python.autograph.utils import tensors
 from tensorflow.python.ops import cond as tf_cond
+from tensorflow.python.util import nest
+from tensorflow.python.util import variable_utils
 
 
 def if_exp(cond, if_true, if_false, expr_repr):
@@ -36,17 +40,39 @@ def _tf_if_exp(cond, if_true, if_false, expr_repr):
   def true_fn():
     true_val.append(if_true())
     if true_val and false_val:
-      control_flow.verify_single_cond_var(expr_repr, true_val[0], false_val[0])
+      _verify_tf_if_exp_cond_var(expr_repr, true_val[0], false_val[0])
     return true_val[0]
 
   def false_fn():
     false_val.append(if_false())
     if true_val and false_val:
-      control_flow.verify_single_cond_var(expr_repr, true_val[0], false_val[0])
+      _verify_tf_if_exp_cond_var(expr_repr, true_val[0], false_val[0])
     return false_val[0]
 
   return tf_cond.cond(cond, true_fn, false_fn)
 
 
+def _verify_tf_if_exp_cond_var(expr_repr, true_val, false_val):
+  """Verifies that both branches of a conditional expression are compatible."""
+  try:
+    nest.assert_same_structure(true_val, false_val, expand_composites=True)
+  except (TypeError, ValueError):
+    try:
+      true_val = variable_utils.convert_variables_to_tensors(true_val)
+      false_val = variable_utils.convert_variables_to_tensors(false_val)
+      nest.assert_same_structure(true_val, false_val, expand_composites=True)
+    except (TypeError, ValueError) as e:
+      raise ValueError(
+          'conditional expression "{}" must return the same nested structure '
+          'in both branches. If you only need side effects, use a statement '
+          '`if:` instead; otherwise, make both branches return the same '
+          'structure or use `tf.cond` explicitly.'.format(expr_repr)) from e
+
+  nest.map_structure(
+      functools.partial(control_flow.verify_single_cond_var, expr_repr),
+      true_val,
+      false_val)
+
+
 def _py_if_exp(cond, if_true, if_false):
   return if_true() if cond else if_false()
diff --git a/tensorflow/python/autograph/tests/basic_ifexp_test.py b/tensorflow/python/autograph/tests/basic_ifexp_test.py
index 02696f4ab1e..18be7727c4b 100644
--- a/tensorflow/python/autograph/tests/basic_ifexp_test.py
+++ b/tensorflow/python/autograph/tests/basic_ifexp_test.py
@@ -46,6 +46,12 @@ def cond_with_multiple_values(x):
   return x, y, z
 
 
+def side_effect_only_ifexp_with_imbalanced_structure(x):
+  traced_args = tf.unstack(x) if tf.size(x) > 0 else []
+  tf.autograph.trace(*traced_args)
+  return x * 2 + 1
+
+
 class ReferenceTest(reference_test_base.TestCase):
 
   def test_basic(self):
@@ -53,6 +59,12 @@ class ReferenceTest(reference_test_base.TestCase):
       self.assertFunctionMatchesEager(consecutive_conds, x)
       self.assertFunctionMatchesEager(cond_with_multiple_values, x)
 
+  def test_imbalanced_conditional_expression_raises_value_error(self):
+    with self.assertRaisesRegex(
+        ValueError, 'must return the same nested structure in both branches'):
+      self.function(side_effect_only_ifexp_with_imbalanced_structure)(
+          tf.constant([1.0]))
+
 
 if __name__ == '__main__':
   tf.test.main()

From 27230b1b5e8ce2089268afeb2ea803710f92ad61 Mon Sep 17 00:00:00 2001
From: Radik <rodion111@gmail.com>
Date: Sat, 11 Apr 2026 16:17:56 +0300
Subject: [PATCH 0002/1000] Fix CWE-770: Add upper-bound guard to
 TensorListReserve and TensorListResize

Cap TensorListReserve and TensorListResize at 1M elements (2^20).
This prevents DoS where INT32_MAX causes 48 GB allocation.

EmptyTensorList already has max_num_elements guard via PushBack;
Reserve and Resize lack equivalent protection. This makes the API
consistent.

Also set output.max_num_elements on Reserve for consistency with
EmptyTensorList.

Bug: 501069108
---
 tensorflow/core/kernels/list_kernels.cc | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/tensorflow/core/kernels/list_kernels.cc b/tensorflow/core/kernels/list_kernels.cc
index fbf8c668cd0..53da0098e4e 100644
--- a/tensorflow/core/kernels/list_kernels.cc
+++ b/tensorflow/core/kernels/list_kernels.cc
@@ -343,9 +343,24 @@ class TensorListReserve : public OpKernel {
                     absl::StrCat("The num_elements to reserve must be a "
                                  "non negative number, but got ",
                                  num_elements)));
+    // Upper bound guard (CWE-770).  Each reserved slot is a Tensor object
+    // (~24-32 bytes), so INT32_MAX slots would allocate ~48-64 GB of
+    // metadata alone.  1 048 576 (2^20) slots use ~24 MB, which covers
+    // all practical ML workloads (RNNs, sequences, time series).
+    // For comparison, EmptyTensorList exposes a max_num_elements input
+    // that TensorListPushBack enforces; TensorListReserve lacked any
+    // equivalent guard.
+    static constexpr int32_t kMaxReserveElements = 1 << 20;  // 1 048 576
+    OP_REQUIRES(c, num_elements <= kMaxReserveElements,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "TensorListReserve: num_elements (", num_elements,
+                    ") exceeds maximum allowed (", kMaxReserveElements,
+                    "). Use EmptyTensorList with TensorListPushBack "
+                    "for dynamically-sized lists.")));
     TensorList output;
     output.element_shape = element_shape;
     output.element_dtype = element_dtype_;
+    output.max_num_elements = num_elements;
     output.tensors().resize(num_elements, Tensor(DT_INVALID));
     Tensor* result;
     AllocatorAttributes attr;
@@ -391,6 +406,13 @@ class TensorListResize : public OpKernel {
         c, size >= 0,
         absl::InvalidArgumentError(absl::StrCat(
             "TensorListSlice expects size to be non-negative. Got: ", size)));
+    // Same upper bound as TensorListReserve (CWE-770).
+    static constexpr int32_t kMaxResizeElements = 1 << 20;  // 1 048 576
+    OP_REQUIRES(c, size <= kMaxResizeElements,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "TensorListResize: size (", size,
+                    ") exceeds maximum allowed (", kMaxResizeElements,
+                    ").")));
 
     std::unique_ptr<Tensor> maybe_result =
         c->forward_input(0, 0, DT_VARIANT, TensorShape{},

From 653c77178c4839ca340b7179aae7b3222f708698 Mon Sep 17 00:00:00 2001
From: cybrdude <layau@netguard24-7.com>
Date: Mon, 13 Jul 2026 11:39:18 -0400
Subject: [PATCH 0003/1000] Add TFLite interpreter-build fuzz target (model
 load + invoke)

The OSS-Fuzz tensorflow build has no fuzz target that reaches TFLite model load, interpreter build, and invoke. This adds one under security/fuzzing/cc, exercising interpreter_builder.cc external-buffer handling (get_readonly_data / ParseTensors), where an unchecked uint64 offset+size addition can bypass the allocation bounds check and set a wild read-only tensor pointer that is dereferenced during Invoke. Deps use the public //tensorflow/lite:framework and //tensorflow/lite/kernels:builtin_ops.
---
 tensorflow/security/fuzzing/cc/BUILD          |  9 ++++
 .../fuzzing/cc/tflite_interp_overflow_fuzz.cc | 54 +++++++++++++++++++
 2 files changed, 63 insertions(+)
 create mode 100644 tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc

diff --git a/tensorflow/security/fuzzing/cc/BUILD b/tensorflow/security/fuzzing/cc/BUILD
index 233c18b8341..afd5e7af76c 100644
--- a/tensorflow/security/fuzzing/cc/BUILD
+++ b/tensorflow/security/fuzzing/cc/BUILD
@@ -253,3 +253,12 @@ tf_cc_fuzz_test(
         "@xla//xla/tsl/platform:env",
     ],
 )
+
+tf_cc_fuzz_test(
+    name = "tflite_interp_overflow_fuzz",
+    srcs = ["tflite_interp_overflow_fuzz.cc"],
+    deps = [
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/kernels:builtin_ops",
+    ],
+)
diff --git a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
new file mode 100644
index 00000000000..d51f15a4142
--- /dev/null
+++ b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
@@ -0,0 +1,54 @@
+/* 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.
+==============================================================================*/
+
+// FuzzTest coverage for TFLite model loading, interpreter build, and invoke.
+// This exercises interpreter_builder.cc get_readonly_data / ParseTensors, which
+// perform an unchecked uint64 addition (offset + size > allocation_->bytes())
+// on attacker-controlled external-buffer offsets. On overflow the check is
+// bypassed and a wild pointer (allocation_->base() + offset) is used as tensor
+// data, which is dereferenced during Invoke. The OSS-Fuzz tensorflow build has
+// no TFLite target that reaches this path.
+
+#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 {
+
+void FuzzModelBuildAndInvoke(const std::string& model_bytes) {
+  auto model = tflite::FlatBufferModel::BuildFromBuffer(model_bytes.data(),
+                                                        model_bytes.size());
+  if (model == nullptr) return;
+
+  tflite::ops::builtin::BuiltinOpResolver resolver;
+  std::unique_ptr<tflite::Interpreter> interpreter;
+  if (tflite::InterpreterBuilder(*model, resolver)(&interpreter) != kTfLiteOk) {
+    return;
+  }
+  if (interpreter == nullptr) return;
+  if (interpreter->AllocateTensors() != kTfLiteOk) return;
+
+  // The unchecked external-offset arithmetic sets a wild read-only tensor
+  // pointer that is dereferenced here.
+  interpreter->Invoke();
+}
+FUZZ_TEST(TfliteInterpreterBuilder, FuzzModelBuildAndInvoke);
+
+}  // namespace

From 3eb7b68f504251c397ee07e253798d545d8027fc Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Mon, 13 Jul 2026 12:35:27 -0400
Subject: [PATCH 0004/1000] Update
 tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
index d51f15a4142..fab8525cc8e 100644
--- a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
+++ b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
@@ -21,6 +21,7 @@ limitations under the License.
 // data, which is dereferenced during Invoke. The OSS-Fuzz tensorflow build has
 // no TFLite target that reaches this path.
 
+#include <cstring>
 #include <memory>
 #include <string>
 

From 141f0b41e740c9e29e7696597a5155a183b760cc Mon Sep 17 00:00:00 2001
From: cybrdude <layau@netguard24-7.com>
Date: Mon, 13 Jul 2026 12:51:10 -0400
Subject: [PATCH 0005/1000] fuzz: zero-init runtime input tensors before Invoke
 to avoid MSan false positives

Addresses gemini-code-assist review. TFLite's arena allocator does not zero tensor memory, so kernels reading uninitialized runtime inputs raise MemorySanitizer use-of-uninitialized-value reports. Only runtime inputs are zeroed, not the crafted constant buffer that carries the finding.
---
 .../security/fuzzing/cc/tflite_interp_overflow_fuzz.cc    | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
index fab8525cc8e..d5b2b315608 100644
--- a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
+++ b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
@@ -46,6 +46,14 @@ void FuzzModelBuildAndInvoke(const std::string& model_bytes) {
   if (interpreter == nullptr) return;
   if (interpreter->AllocateTensors() != kTfLiteOk) return;
 
+  // Zero-initialize runtime inputs so MSan does not flag uninitialized reads.
+  for (int input_index : interpreter->inputs()) {
+    TfLiteTensor* tensor = interpreter->tensor(input_index);
+    if (tensor != nullptr && tensor->data.raw != nullptr) {
+      std::memset(tensor->data.raw, 0, tensor->bytes);
+    }
+  }
+
   // The unchecked external-offset arithmetic sets a wild read-only tensor
   // pointer that is dereferenced here.
   interpreter->Invoke();

From bac0ce449e6737f83341a9c7690637a3d358b0e4 Mon Sep 17 00:00:00 2001
From: RadikHoroshev <rodion111@gmail.com>
Date: Thu, 16 Jul 2026 13:31:47 +0300
Subject: [PATCH 0006/1000] Allow TensorListReserve results to grow dynamically

Remove the max_num_elements assignment requested in review while retaining the TensorListReserve and TensorListResize upper-bound guards.
---
 tensorflow/core/kernels/list_kernels.cc | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/core/kernels/list_kernels.cc b/tensorflow/core/kernels/list_kernels.cc
index 53da0098e4e..dbecdaf58cc 100644
--- a/tensorflow/core/kernels/list_kernels.cc
+++ b/tensorflow/core/kernels/list_kernels.cc
@@ -360,7 +360,6 @@ class TensorListReserve : public OpKernel {
     TensorList output;
     output.element_shape = element_shape;
     output.element_dtype = element_dtype_;
-    output.max_num_elements = num_elements;
     output.tensors().resize(num_elements, Tensor(DT_INVALID));
     Tensor* result;
     AllocatorAttributes attr;

From 1459716e5114bfb90bc0f852556f11dedbdb74e9 Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Sun, 19 Jul 2026 13:46:54 +0530
Subject: [PATCH 0007/1000] reject zero-element non-memcpyable components in
 UncompressElement

---
 tensorflow/core/data/compression_utils.cc     | 10 +++++++++
 .../core/data/compression_utils_test.cc       | 21 +++++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/tensorflow/core/data/compression_utils.cc b/tensorflow/core/data/compression_utils.cc
index 073be234e04..9b8d1cf81fe 100644
--- a/tensorflow/core/data/compression_utils.cc
+++ b/tensorflow/core/data/compression_utils.cc
@@ -250,6 +250,16 @@ absl::Status UncompressElement(const CompressedElement& compressed,
         compressed.component_metadata(i);
     if (!DataTypeCanUseMemcpy(metadata.dtype()) &&
         metadata.dtype() != DT_STRING) {
+      TensorShape shape(metadata.tensor_shape());
+      if (shape.num_elements() == 0) {
+        // The first two passes reserve and fill `nonmemcpyable` only for
+        // components with a non-zero element count, so there are no bytes
+        // here to deserialize from.
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Zero-element component ", i, " with non-memcpyable dtype ",
+            DataTypeString(metadata.dtype()),
+            " has no uncompressed bytes to deserialize"));
+      }
       TensorProto tp;
       if (!tp.ParseFromString(
               {nonmemcpyable_pos,
diff --git a/tensorflow/core/data/compression_utils_test.cc b/tensorflow/core/data/compression_utils_test.cc
index 94d6dac0a19..6b05e20d4be 100644
--- a/tensorflow/core/data/compression_utils_test.cc
+++ b/tensorflow/core/data/compression_utils_test.cc
@@ -44,6 +44,27 @@ TEST(CompressionUtilsTest, Exceeds4GB) {
                              HasSubstr("exceeding the 4GB Snappy limit")));
 }
 
+TEST(CompressionUtilsTest, ZeroElementNonMemcpyableComponent) {
+  // Compressing an empty element yields a `data` field holding zero
+  // uncompressed bytes, so the size reconciliation against the (empty) iovec
+  // succeeds and the third pass is reached.
+  std::vector<Tensor> empty_element;
+  CompressedElement compressed;
+  TF_ASSERT_OK(CompressElement(empty_element, &compressed));
+
+  CompressedComponentMetadata* metadata =
+      compressed.mutable_component_metadata()->Add();
+  metadata->set_dtype(DT_VARIANT);
+  TensorShape empty_shape({0});
+  empty_shape.AsProto(metadata->mutable_tensor_shape());
+  metadata->add_uncompressed_bytes(65536);
+
+  std::vector<Tensor> element;
+  EXPECT_THAT(UncompressElement(compressed, &element),
+              absl_testing::StatusIs(error::INVALID_ARGUMENT,
+                                     HasSubstr("Zero-element component")));
+}
+
 std::vector<std::vector<Tensor>> TestCases() {
   return {
       // Single int64.

From 5f962f3f2acafca7f4f1de85c6c90d511eae1435 Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Sun, 19 Jul 2026 15:36:30 +0530
Subject: [PATCH 0008/1000] use TensorShape::BuildTensorShape for the untrusted
 shape proto

---
 tensorflow/core/data/compression_utils.cc | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/data/compression_utils.cc b/tensorflow/core/data/compression_utils.cc
index 9b8d1cf81fe..5ebe44f3dcc 100644
--- a/tensorflow/core/data/compression_utils.cc
+++ b/tensorflow/core/data/compression_utils.cc
@@ -31,6 +31,7 @@ limitations under the License.
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/framework/variant_op_registry.h"
+#include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/snappy.h"
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/platform/tstring.h"
@@ -250,7 +251,9 @@ absl::Status UncompressElement(const CompressedElement& compressed,
         compressed.component_metadata(i);
     if (!DataTypeCanUseMemcpy(metadata.dtype()) &&
         metadata.dtype() != DT_STRING) {
-      TensorShape shape(metadata.tensor_shape());
+      TensorShape shape;
+      TF_RETURN_IF_ERROR(
+          TensorShape::BuildTensorShape(metadata.tensor_shape(), &shape));
       if (shape.num_elements() == 0) {
         // The first two passes reserve and fill `nonmemcpyable` only for
         // components with a non-zero element count, so there are no bytes

From 2f27f04ba800235034abda72aeb0619d054ae688 Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Thu, 23 Jul 2026 08:49:49 -0700
Subject: [PATCH 0009/1000] Update Dataset.concatenate docstring to show
 current error message

The docstring examples for tf.data.Dataset.concatenate still showed the
old incompatible-elements error text ("Two datasets to concatenate have
different types"). Since the op moved to concatenate_op.py, the actual
message is "Incompatible dataset elements: ..." with full TensorSpecs.
Doc-only change, no functional changes.

Fixes #123757
---
 tensorflow/python/data/ops/dataset_ops.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/data/ops/dataset_ops.py b/tensorflow/python/data/ops/dataset_ops.py
index 37c7cc110f4..9892a894b0b 100644
--- a/tensorflow/python/data/ops/dataset_ops.py
+++ b/tensorflow/python/data/ops/dataset_ops.py
@@ -1102,13 +1102,16 @@ class DatasetV2(
     >>> c = tf.data.Dataset.zip((a, b))
     >>> a.concatenate(c)
     Traceback (most recent call last):
-    TypeError: Two datasets to concatenate have different types
-    <dtype: 'int64'> and (tf.int64, tf.int64)
+    TypeError: Incompatible dataset elements:
+      TensorSpec(shape=(), dtype=tf.int64, name=None) vs.
+      (TensorSpec(shape=(), dtype=tf.int64, name=None),
+      TensorSpec(shape=(), dtype=tf.int64, name=None))
     >>> d = tf.data.Dataset.from_tensor_slices(["a", "b", "c"])
     >>> a.concatenate(d)
     Traceback (most recent call last):
-    TypeError: Two datasets to concatenate have different types
-    <dtype: 'int64'> and <dtype: 'string'>
+    TypeError: Incompatible dataset elements:
+      TensorSpec(shape=(), dtype=tf.int64, name=None) vs.
+      TensorSpec(shape=(), dtype=tf.string, name=None)
 
     Args:
       dataset: `Dataset` to be concatenated.

From 03900ec096af4484f771b09bf09370a709476ae5 Mon Sep 17 00:00:00 2001
From: Jayden G <user.josephms@gmail.com>
Date: Wed, 18 Feb 2026 12:45:26 +0530
Subject: [PATCH 0010/1000] Fix SparseConcat crash on integer overflow in shape
 computation

---
 tensorflow/core/kernels/sparse_concat_op.cc   |  1 +
 .../python/kernel_tests/sparse_ops/BUILD      |  7 +++++
 .../sparse_ops/sparse_concat_overflow_test.py | 28 +++++++++++++++++++
 3 files changed, 36 insertions(+)
 create mode 100644 tensorflow/python/kernel_tests/sparse_ops/sparse_concat_overflow_test.py

diff --git a/tensorflow/core/kernels/sparse_concat_op.cc b/tensorflow/core/kernels/sparse_concat_op.cc
index 92dfa947baa..e6d5e6d27fd 100644
--- a/tensorflow/core/kernels/sparse_concat_op.cc
+++ b/tensorflow/core/kernels/sparse_concat_op.cc
@@ -271,3 +271,4 @@ TF_CALL_POD_TYPES(REGISTER_KERNELS);
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 
 }  // namespace tensorflow
+
diff --git a/tensorflow/python/kernel_tests/sparse_ops/BUILD b/tensorflow/python/kernel_tests/sparse_ops/BUILD
index 5ef3878728c..047a13f7f34 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/BUILD
+++ b/tensorflow/python/kernel_tests/sparse_ops/BUILD
@@ -16,6 +16,7 @@
 # Tests of TensorFlow sparse ops written using the Python API.
 
 load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "tf_py_test")
 load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_py_strict_test")
 
 package(
@@ -408,3 +409,9 @@ tf_py_strict_test(
         "//third_party/py/numpy",
     ],
 )
+tf_py_test(
+    name = "sparse_concat_overflow_test",
+    srcs = ["sparse_concat_overflow_test.py"],
+    tags = ["no_oss"],
+)
+
diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_overflow_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_overflow_test.py
new file mode 100644
index 00000000000..b7a7cde4430
--- /dev/null
+++ b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_overflow_test.py
@@ -0,0 +1,28 @@
+import tensorflow as tf
+from tensorflow.python.framework import errors
+from tensorflow.python.platform import test
+
+
+class SparseConcatOverflowTest(test.TestCase):
+
+    def test_sparse_concat_shape_overflow(self):
+        indices1 = tf.constant([[0, 0, 0]], dtype=tf.int64)
+        values1 = tf.constant([1], dtype=tf.int32)
+        shape1 = tf.constant([5, 2, 2147483647], dtype=tf.int64)
+
+        indices2 = tf.constant([[0, 1, 0]], dtype=tf.int64)
+        values2 = tf.constant([2], dtype=tf.int32)
+        shape2 = tf.constant([5, 1879048192, 536870912], dtype=tf.int64)
+
+        with self.assertRaises(errors.InvalidArgumentError):
+            tf.raw_ops.SparseConcat(
+                indices=[indices1, indices2],
+                values=[values1, values2],
+                shapes=[shape1, shape2],
+                concat_dim=1,
+            )
+
+
+if __name__ == "__main__":
+    test.main()
+

From 05074c433a238c6072841fe3b671b373d9c22aa1 Mon Sep 17 00:00:00 2001
From: Jayden G <user.josephms@gmail.com>
Date: Wed, 18 Feb 2026 21:28:36 +0530
Subject: [PATCH 0011/1000] Fix pylint issues in sparse_concat_overflow_test

---
 .../sparse_ops/sparse_concat_overflow_test.py | 54 +++++++++++++------
 1 file changed, 37 insertions(+), 17 deletions(-)

diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_overflow_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_overflow_test.py
index b7a7cde4430..81a83eb10a3 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_overflow_test.py
+++ b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_overflow_test.py
@@ -1,28 +1,48 @@
-import tensorflow as tf
+# Copyright 2024 The TensorFlow Authors.
+#
+# 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 SparseConcat overflow handling."""
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
 from tensorflow.python.platform import test
 
 
 class SparseConcatOverflowTest(test.TestCase):
 
-    def test_sparse_concat_shape_overflow(self):
-        indices1 = tf.constant([[0, 0, 0]], dtype=tf.int64)
-        values1 = tf.constant([1], dtype=tf.int32)
-        shape1 = tf.constant([5, 2, 2147483647], dtype=tf.int64)
+  def testSparseConcatLargeShapeOverflow(self):
+    indices1 = constant_op.constant(
+        [[0, 0, 0], [1, 1, 0], [2, 0, 1]], dtype=dtypes.int64)
+    values1 = constant_op.constant(["tensor", "flow", "test"])
+    shape1 = constant_op.constant([5, 2, 2147483647], dtype=dtypes.int64)
 
-        indices2 = tf.constant([[0, 1, 0]], dtype=tf.int64)
-        values2 = tf.constant([2], dtype=tf.int32)
-        shape2 = tf.constant([5, 1879048192, 536870912], dtype=tf.int64)
+    indices2 = constant_op.constant(
+        [[0, 1, 0], [1, 0, 0], [2, 1, 1], [3, 0, 1]], dtype=dtypes.int64)
+    values2 = constant_op.constant(["a", "b", "c", "d"])
+    shape2 = constant_op.constant(
+        [5, 1879048192, 536870912], dtype=dtypes.int64)
 
-        with self.assertRaises(errors.InvalidArgumentError):
-            tf.raw_ops.SparseConcat(
-                indices=[indices1, indices2],
-                values=[values1, values2],
-                shapes=[shape1, shape2],
-                concat_dim=1,
-            )
+    sp1 = sparse_tensor.SparseTensor(indices1, values1, shape1)
+    sp2 = sparse_tensor.SparseTensor(indices2, values2, shape2)
+
+    with self.assertRaises(errors.InvalidArgumentError):
+      sparse_tensor.sparse_concat(
+          axis=1,
+          sp_inputs=[sp1, sp2])
 
 
 if __name__ == "__main__":
-    test.main()
-
+  test.main()

From 2ecd7f7517f2980b1758b8b117e89a750c18d949 Mon Sep 17 00:00:00 2001
From: Jayden G <user.josephms@gmail.com>
Date: Thu, 19 Feb 2026 17:24:13 +0530
Subject: [PATCH 0012/1000] Use tf_py_strict_test for
 sparse_concat_overflow_test

---
 tensorflow/python/kernel_tests/sparse_ops/BUILD | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/kernel_tests/sparse_ops/BUILD b/tensorflow/python/kernel_tests/sparse_ops/BUILD
index 047a13f7f34..b050b3bb9e1 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/BUILD
+++ b/tensorflow/python/kernel_tests/sparse_ops/BUILD
@@ -409,9 +409,19 @@ tf_py_strict_test(
         "//third_party/py/numpy",
     ],
 )
-tf_py_test(
+tf_py_strict_test(
     name = "sparse_concat_overflow_test",
+    size = "small",
     srcs = ["sparse_concat_overflow_test.py"],
-    tags = ["no_oss"],
+    deps = [
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:for_generated_wrappers",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:sparse_ops",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+    ],
 )
 

From 84ef524d35309b30587ece62c799865831e142e7 Mon Sep 17 00:00:00 2001
From: Jayden G <user.josephms@gmail.com>
Date: Fri, 20 Feb 2026 15:56:17 +0530
Subject: [PATCH 0013/1000] Fix SparseConcat overflow: safe addition +
 SetDimWithStatus

Signed-off-by: Jayden G <user.josephms@gmail.com>
---
 tensorflow/core/kernels/sparse_concat_op.cc | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/kernels/sparse_concat_op.cc b/tensorflow/core/kernels/sparse_concat_op.cc
index e6d5e6d27fd..0417c490eb3 100644
--- a/tensorflow/core/kernels/sparse_concat_op.cc
+++ b/tensorflow/core/kernels/sparse_concat_op.cc
@@ -1,3 +1,4 @@
+
 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -20,6 +21,7 @@ limitations under the License.
 #include <numeric>
 #include <utility>
 #include <vector>
+#include <limits>
 
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
@@ -70,6 +72,7 @@ struct SparseConcatFunctor<CPUDevice, T> {
     }
 
     std::vector<sparse::SparseTensor> sp_inputs;
+    sp_inputs.reserve(N);
     for (int i = 0; i < N; ++i) {
       const TensorShape current_shape(shapes[i].vec<int64_t>());
       sparse::SparseTensor tensor;
@@ -184,6 +187,9 @@ class SparseConcatOp : public OpKernel {
                     "Concat dimension must be in range [", -input_rank, ", ",
                     input_rank, "), got ", concat_dim_attr_)));
     TensorShape output_shape = input_shape;
+    // When accumulating the output dimension at concat_dim, do a safe add to
+    // ensure the addition itself does not overflow; then use SetDimWithStatus to
+    // get proper error status if the TensorShape internals detect invalid dims.
     for (int i = 1; i < N; ++i) {
       const TensorShape current_shape(shapes[i].vec<int64_t>());
       OP_REQUIRES(
@@ -219,6 +225,7 @@ class SparseConcatOp : public OpKernel {
       output_shape_t(j) = output_shape.dim_size(j);
     }
 
+    // Sum output_nnz carefully and check for overflow while summing.
     int64_t output_nnz = 0;
     for (int i = 0; i < N; ++i) {
       int64_t next_output_nnz =
@@ -227,6 +234,7 @@ class SparseConcatOp : public OpKernel {
                   absl::InvalidArgumentError("output_nnz overflowed"));
       output_nnz = next_output_nnz;
     }
+
     if (output_nnz == 0) {
       Tensor* output_inds = nullptr;
       OP_REQUIRES_OK(context,
@@ -270,5 +278,4 @@ TF_CALL_POD_TYPES(REGISTER_KERNELS);
 
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 
-}  // namespace tensorflow
-
+}  // namespace tensorflow
\ No newline at end of file

From 64687106c64e0a20f1e928b5a10e4060970ad7cd Mon Sep 17 00:00:00 2001
From: Jayden G <user.josephms@gmail.com>
Date: Wed, 3 Jun 2026 18:37:50 +0530
Subject: [PATCH 0014/1000] Use py_strict_library for sparse_xent_op_test_base

---
 tensorflow/python/kernel_tests/sparse_ops/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/kernel_tests/sparse_ops/BUILD b/tensorflow/python/kernel_tests/sparse_ops/BUILD
index b050b3bb9e1..ce7cf8e9691 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/BUILD
+++ b/tensorflow/python/kernel_tests/sparse_ops/BUILD
@@ -374,7 +374,7 @@ cuda_py_strict_test(
     ],
 )
 
-py_library(
+py_strict_library(
     name = "sparse_xent_op_test_base",
     srcs = ["sparse_xent_op_test_base.py"],
     strict_deps = True,

From c7b4dcedc609c2688fab13983ba10ad044f55241 Mon Sep 17 00:00:00 2001
From: Jayden G <user.josephms@gmail.com>
Date: Thu, 4 Jun 2026 12:37:38 +0530
Subject: [PATCH 0015/1000] Address review feedback

---
 tensorflow/core/kernels/sparse_concat_op.cc                  | 3 ++-
 tensorflow/python/kernel_tests/sparse_ops/BUILD              | 2 +-
 .../kernel_tests/sparse_ops/sparse_concat_overflow_test.py   | 5 +++--
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/kernels/sparse_concat_op.cc b/tensorflow/core/kernels/sparse_concat_op.cc
index 0417c490eb3..c2791585dc3 100644
--- a/tensorflow/core/kernels/sparse_concat_op.cc
+++ b/tensorflow/core/kernels/sparse_concat_op.cc
@@ -211,7 +211,8 @@ class SparseConcatOp : public OpKernel {
           OP_REQUIRES(context, new_dim >= 0,
                       absl::InvalidArgumentError(absl::StrCat(
                           "Concat dimension overflowed at position ", i)));
-          output_shape.set_dim(j, new_dim);
+          OP_REQUIRES_OK(context,
+                         output_shape.SetDimWithStatus(j, new_dim));
         }
       }
     }
diff --git a/tensorflow/python/kernel_tests/sparse_ops/BUILD b/tensorflow/python/kernel_tests/sparse_ops/BUILD
index ce7cf8e9691..984574fd2bb 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/BUILD
+++ b/tensorflow/python/kernel_tests/sparse_ops/BUILD
@@ -17,6 +17,7 @@
 
 load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
 load("//tensorflow:tensorflow.bzl", "tf_py_test")
+load("//tensorflow:strict.default.bzl", "py_strict_library")
 load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_py_strict_test")
 
 package(
@@ -377,7 +378,6 @@ cuda_py_strict_test(
 py_strict_library(
     name = "sparse_xent_op_test_base",
     srcs = ["sparse_xent_op_test_base.py"],
-    strict_deps = True,
     deps = [
         "//tensorflow/python/eager:backprop",
         "//tensorflow/python/eager:context",
diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_overflow_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_overflow_test.py
index 81a83eb10a3..50062e7faa0 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_overflow_test.py
+++ b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_overflow_test.py
@@ -18,6 +18,7 @@ from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import sparse_ops
 from tensorflow.python.platform import test
 
 
@@ -33,13 +34,13 @@ class SparseConcatOverflowTest(test.TestCase):
         [[0, 1, 0], [1, 0, 0], [2, 1, 1], [3, 0, 1]], dtype=dtypes.int64)
     values2 = constant_op.constant(["a", "b", "c", "d"])
     shape2 = constant_op.constant(
-        [5, 1879048192, 536870912], dtype=dtypes.int64)
+        [5, 1879048192, 2147483647], dtype=dtypes.int64)
 
     sp1 = sparse_tensor.SparseTensor(indices1, values1, shape1)
     sp2 = sparse_tensor.SparseTensor(indices2, values2, shape2)
 
     with self.assertRaises(errors.InvalidArgumentError):
-      sparse_tensor.sparse_concat(
+      sparse_ops.sparse_concat(
           axis=1,
           sp_inputs=[sp1, sp2])
 

From 4e1dea2926f5a14cac0abad9677eed970f02285c Mon Sep 17 00:00:00 2001
From: Jayden Jeswin Raj <user.josephms@gmail.com>
Date: Mon, 29 Jun 2026 13:55:15 +0530
Subject: [PATCH 0016/1000] Add regression test for SparseConcat shape overflow

---
 .../sparse_ops/sparse_concat_op_test.py       | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
index 3d96beaab34..20e96e6440a 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
+++ b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
@@ -522,5 +522,25 @@ class SparseConcatTest(test.TestCase):
         self.evaluate(concat_op)
 
 
+  def testSparseConcatLargeShapeOverflow(self):
+    indices1 = constant_op.constant(
+        [[0, 0, 0], [1, 1, 0], [2, 0, 1]], dtype=dtypes.int64)
+    values1 = constant_op.constant(["tensor", "flow", "test"])
+    shape1 = constant_op.constant([5, 2, 2147483647], dtype=dtypes.int64)
+
+    indices2 = constant_op.constant(
+        [[0, 1, 0], [1, 0, 0], [2, 1, 1], [3, 0, 1]], dtype=dtypes.int64)
+    values2 = constant_op.constant(["a", "b", "c", "d"])
+    shape2 = constant_op.constant(
+        [5, 1879048192, 2147483647], dtype=dtypes.int64)
+
+    sp1 = sparse_tensor.SparseTensor(indices1, values1, shape1)
+    sp2 = sparse_tensor.SparseTensor(indices2, values2, shape2)
+
+    with self.assertRaisesOpError(
+        "Encountered overflow from large input shape."
+    ):
+      self.evaluate(sparse_ops.sparse_concat(axis=1, sp_inputs=[sp1, sp2]))
+
 if __name__ == "__main__":
   test.main()

From d4e63a358084663d33aaa9c7288a522cf7a6c239 Mon Sep 17 00:00:00 2001
From: Jayden Jeswin Raj <user.josephms@gmail.com>
Date: Tue, 30 Jun 2026 01:25:39 +0530
Subject: [PATCH 0017/1000] Remove redundant SparseConcat test target

---
 .../sparse_ops/sparse_concat_overflow_test.py | 49 -------------------
 1 file changed, 49 deletions(-)
 delete mode 100644 tensorflow/python/kernel_tests/sparse_ops/sparse_concat_overflow_test.py

diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_overflow_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_overflow_test.py
deleted file mode 100644
index 50062e7faa0..00000000000
--- a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_overflow_test.py
+++ /dev/null
@@ -1,49 +0,0 @@
-# Copyright 2024 The TensorFlow Authors.
-#
-# 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 SparseConcat overflow handling."""
-
-from tensorflow.python.framework import constant_op
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import errors
-from tensorflow.python.framework import sparse_tensor
-from tensorflow.python.ops import sparse_ops
-from tensorflow.python.platform import test
-
-
-class SparseConcatOverflowTest(test.TestCase):
-
-  def testSparseConcatLargeShapeOverflow(self):
-    indices1 = constant_op.constant(
-        [[0, 0, 0], [1, 1, 0], [2, 0, 1]], dtype=dtypes.int64)
-    values1 = constant_op.constant(["tensor", "flow", "test"])
-    shape1 = constant_op.constant([5, 2, 2147483647], dtype=dtypes.int64)
-
-    indices2 = constant_op.constant(
-        [[0, 1, 0], [1, 0, 0], [2, 1, 1], [3, 0, 1]], dtype=dtypes.int64)
-    values2 = constant_op.constant(["a", "b", "c", "d"])
-    shape2 = constant_op.constant(
-        [5, 1879048192, 2147483647], dtype=dtypes.int64)
-
-    sp1 = sparse_tensor.SparseTensor(indices1, values1, shape1)
-    sp2 = sparse_tensor.SparseTensor(indices2, values2, shape2)
-
-    with self.assertRaises(errors.InvalidArgumentError):
-      sparse_ops.sparse_concat(
-          axis=1,
-          sp_inputs=[sp1, sp2])
-
-
-if __name__ == "__main__":
-  test.main()

From abed050b06af87f2c1545271947e03ef818329c9 Mon Sep 17 00:00:00 2001
From: Jayden Jeswin Raj <user.josephms@gmail.com>
Date: Tue, 30 Jun 2026 11:26:20 +0530
Subject: [PATCH 0018/1000] Revert SparseConcat BUILD file changes

---
 tensorflow/python/kernel_tests/sparse_ops/BUILD | 16 +---------------
 1 file changed, 1 insertion(+), 15 deletions(-)

diff --git a/tensorflow/python/kernel_tests/sparse_ops/BUILD b/tensorflow/python/kernel_tests/sparse_ops/BUILD
index 984574fd2bb..7583bb63c62 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/BUILD
+++ b/tensorflow/python/kernel_tests/sparse_ops/BUILD
@@ -409,19 +409,5 @@ tf_py_strict_test(
         "//third_party/py/numpy",
     ],
 )
-tf_py_strict_test(
-    name = "sparse_concat_overflow_test",
-    size = "small",
-    srcs = ["sparse_concat_overflow_test.py"],
-    deps = [
-        "//tensorflow/python/framework:constant_op",
-        "//tensorflow/python/framework:for_generated_wrappers",
-        "//tensorflow/python/framework:sparse_tensor",
-        "//tensorflow/python/framework:test_lib",
-        "//tensorflow/python/ops:array_ops",
-        "//tensorflow/python/ops:sparse_ops",
-        "//tensorflow/python/platform:client_testlib",
-        "//third_party/py/numpy",
-    ],
-)
+
 

From e116ed2c5ad0103808782aee117f3b263fbd4e82 Mon Sep 17 00:00:00 2001
From: Jayden Jeswin Raj <user.josephms@gmail.com>
Date: Wed, 8 Jul 2026 23:16:31 +0530
Subject: [PATCH 0019/1000] Fix SparseConcat source formatting

---
 tensorflow/core/kernels/sparse_concat_op.cc | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensorflow/core/kernels/sparse_concat_op.cc b/tensorflow/core/kernels/sparse_concat_op.cc
index c2791585dc3..65a62df0152 100644
--- a/tensorflow/core/kernels/sparse_concat_op.cc
+++ b/tensorflow/core/kernels/sparse_concat_op.cc
@@ -1,4 +1,3 @@
-
 /* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -279,4 +278,4 @@ TF_CALL_POD_TYPES(REGISTER_KERNELS);
 
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 
-}  // namespace tensorflow
\ No newline at end of file
+}  // namespace tensorflow

From 354d6f3460af54258e22aa37b5d79969604bf55a Mon Sep 17 00:00:00 2001
From: Jayden Jeswin Raj <user.josephms@gmail.com>
Date: Thu, 16 Jul 2026 22:44:31 +0530
Subject: [PATCH 0020/1000] Address review feedback: use absl error namespace,
 revert BUILD changes

---
 tensorflow/python/kernel_tests/sparse_ops/BUILD | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/tensorflow/python/kernel_tests/sparse_ops/BUILD b/tensorflow/python/kernel_tests/sparse_ops/BUILD
index 7583bb63c62..5ef3878728c 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/BUILD
+++ b/tensorflow/python/kernel_tests/sparse_ops/BUILD
@@ -16,8 +16,6 @@
 # Tests of TensorFlow sparse ops written using the Python API.
 
 load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
-load("//tensorflow:tensorflow.bzl", "tf_py_test")
-load("//tensorflow:strict.default.bzl", "py_strict_library")
 load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_py_strict_test")
 
 package(
@@ -375,9 +373,10 @@ cuda_py_strict_test(
     ],
 )
 
-py_strict_library(
+py_library(
     name = "sparse_xent_op_test_base",
     srcs = ["sparse_xent_op_test_base.py"],
+    strict_deps = True,
     deps = [
         "//tensorflow/python/eager:backprop",
         "//tensorflow/python/eager:context",
@@ -409,5 +408,3 @@ tf_py_strict_test(
         "//third_party/py/numpy",
     ],
 )
-
-

From 78959963b819a241f40fb512bfe8153d0f3648fd Mon Sep 17 00:00:00 2001
From: Jayden Jeswin Raj <user.josephms@gmail.com>
Date: Fri, 17 Jul 2026 08:19:58 +0530
Subject: [PATCH 0021/1000] Guard against empty input list in SparseConcat

---
 tensorflow/core/kernels/sparse_concat_op.cc | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow/core/kernels/sparse_concat_op.cc b/tensorflow/core/kernels/sparse_concat_op.cc
index 65a62df0152..ec4cc1456f8 100644
--- a/tensorflow/core/kernels/sparse_concat_op.cc
+++ b/tensorflow/core/kernels/sparse_concat_op.cc
@@ -104,6 +104,8 @@ class SparseConcatOp : public OpKernel {
     OpInputList inds;
     OP_REQUIRES_OK(context, context->input_list("indices", &inds));
     const int N = inds.size();
+    OP_REQUIRES(context, N > 0,
+                absl::InvalidArgumentError("Input list must not be empty"));
     for (int i = 0; i < N; i++) {
       OP_REQUIRES(context, TensorShapeUtils::IsMatrix(inds[i].shape()),
                   absl::InvalidArgumentError(absl::StrCat(

From 4c852cdfec8c720d23b6921a2b7444a2ec2f247b Mon Sep 17 00:00:00 2001
From: Jayden Jeswin Raj <user.josephms@gmail.com>
Date: Sat, 18 Jul 2026 22:06:07 +0530
Subject: [PATCH 0022/1000] Fix test to exercise SetDimWithStatus overflow and
 clean up comment

Replace testSparseConcatLargeShapeOverflow (which hit the early input
volume check) with testSparseConcatOutputVolumeOverflow that uses
shapes where individual volumes fit in int64 but the concatenated
output volume overflows, properly exercising the SetDimWithStatus ->
RecomputeNumElements path.

Trim the comment above the accumulation loop to stay within 80
characters and clarify the purpose.
---
 tensorflow/core/kernels/sparse_concat_op.cc   |  6 ++--
 .../sparse_ops/sparse_concat_op_test.py       | 32 ++++++++++++-------
 2 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/tensorflow/core/kernels/sparse_concat_op.cc b/tensorflow/core/kernels/sparse_concat_op.cc
index ec4cc1456f8..d721c0a033a 100644
--- a/tensorflow/core/kernels/sparse_concat_op.cc
+++ b/tensorflow/core/kernels/sparse_concat_op.cc
@@ -188,9 +188,9 @@ class SparseConcatOp : public OpKernel {
                     "Concat dimension must be in range [", -input_rank, ", ",
                     input_rank, "), got ", concat_dim_attr_)));
     TensorShape output_shape = input_shape;
-    // When accumulating the output dimension at concat_dim, do a safe add to
-    // ensure the addition itself does not overflow; then use SetDimWithStatus to
-    // get proper error status if the TensorShape internals detect invalid dims.
+    // Accumulate the concat dimension with overflow-safe addition,
+    // then use SetDimWithStatus so RecomputeNumElements can catch
+    // output volume overflow from the enlarged shape.
     for (int i = 1; i < N; ++i) {
       const TensorShape current_shape(shapes[i].vec<int64_t>());
       OP_REQUIRES(
diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
index 20e96e6440a..d9e7465abad 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
+++ b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
@@ -522,25 +522,33 @@ class SparseConcatTest(test.TestCase):
         self.evaluate(concat_op)
 
 
-  def testSparseConcatLargeShapeOverflow(self):
-    indices1 = constant_op.constant(
-        [[0, 0, 0], [1, 1, 0], [2, 0, 1]], dtype=dtypes.int64)
-    values1 = constant_op.constant(["tensor", "flow", "test"])
-    shape1 = constant_op.constant([5, 2, 2147483647], dtype=dtypes.int64)
+  def testSparseConcatOutputVolumeOverflow(self):
+    # Individual volumes fit in int64, but the concatenated
+    # output volume overflows int64, triggering
+    # SetDimWithStatus -> RecomputeNumElements.
+    # Axis = 1, X = int64_max // 11
+    # Input 1 volume: 10 * X < int64_max
+    # Input 2 volume:  2 * X < int64_max
+    # Output volume:  12 * X > int64_max
+    int64_max = 9223372036854775807
+    X = int64_max // 11
 
-    indices2 = constant_op.constant(
-        [[0, 1, 0], [1, 0, 0], [2, 1, 1], [3, 0, 1]], dtype=dtypes.int64)
-    values2 = constant_op.constant(["a", "b", "c", "d"])
-    shape2 = constant_op.constant(
-        [5, 1879048192, 2147483647], dtype=dtypes.int64)
+    indices1 = constant_op.constant([[0, 0, 0]], dtype=dtypes.int64)
+    values1 = constant_op.constant([1.0], dtype=dtypes.float32)
+    shape1 = constant_op.constant([1, 10, X], dtype=dtypes.int64)
+
+    indices2 = constant_op.constant([[0, 0, 0]], dtype=dtypes.int64)
+    values2 = constant_op.constant([2.0], dtype=dtypes.float32)
+    shape2 = constant_op.constant([1, 2, X], dtype=dtypes.int64)
 
     sp1 = sparse_tensor.SparseTensor(indices1, values1, shape1)
     sp2 = sparse_tensor.SparseTensor(indices2, values2, shape2)
 
     with self.assertRaisesOpError(
-        "Encountered overflow from large input shape."
+        "results in overflow when computing number of elements"
     ):
-      self.evaluate(sparse_ops.sparse_concat(axis=1, sp_inputs=[sp1, sp2]))
+      self.evaluate(
+          sparse_ops.sparse_concat(axis=1, sp_inputs=[sp1, sp2]))
 
 if __name__ == "__main__":
   test.main()

From 50647c06ee27854289de349c2ce0ea5ca607a6de Mon Sep 17 00:00:00 2001
From: Jayden Jeswin Raj <user.josephms@gmail.com>
Date: Wed, 29 Jul 2026 08:52:16 +0530
Subject: [PATCH 0023/1000] Fix formatting: consolidate multi-line statements
 to single line

---
 tensorflow/core/kernels/sparse_concat_op.cc                    | 3 +--
 .../python/kernel_tests/sparse_ops/sparse_concat_op_test.py    | 3 +--
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/kernels/sparse_concat_op.cc b/tensorflow/core/kernels/sparse_concat_op.cc
index d721c0a033a..940dab94d9d 100644
--- a/tensorflow/core/kernels/sparse_concat_op.cc
+++ b/tensorflow/core/kernels/sparse_concat_op.cc
@@ -212,8 +212,7 @@ class SparseConcatOp : public OpKernel {
           OP_REQUIRES(context, new_dim >= 0,
                       absl::InvalidArgumentError(absl::StrCat(
                           "Concat dimension overflowed at position ", i)));
-          OP_REQUIRES_OK(context,
-                         output_shape.SetDimWithStatus(j, new_dim));
+          OP_REQUIRES_OK(context, output_shape.SetDimWithStatus(j, new_dim));
         }
       }
     }
diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
index d9e7465abad..9f3d1ff236e 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
+++ b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
@@ -547,8 +547,7 @@ class SparseConcatTest(test.TestCase):
     with self.assertRaisesOpError(
         "results in overflow when computing number of elements"
     ):
-      self.evaluate(
-          sparse_ops.sparse_concat(axis=1, sp_inputs=[sp1, sp2]))
+      self.evaluate(sparse_ops.sparse_concat(axis=1, sp_inputs=[sp1, sp2]))
 
 if __name__ == "__main__":
   test.main()

From 725dda3bca5200e8ad59418a7d8fa36dedde3db2 Mon Sep 17 00:00:00 2001
From: Jayden Jeswin Raj <user.josephms@gmail.com>
Date: Wed, 29 Jul 2026 21:49:56 +0530
Subject: [PATCH 0024/1000] retrigger CI


From a4ba42b14c9c19509e3c5a19619b4373c536f7a2 Mon Sep 17 00:00:00 2001
From: Jayden Jeswin Raj <user.josephms@gmail.com>
Date: Thu, 30 Jul 2026 00:39:12 +0530
Subject: [PATCH 0025/1000] Fix formatting: sort includes and fix blank line

---
 tensorflow/core/kernels/sparse_concat_op.cc                     | 2 +-
 .../python/kernel_tests/sparse_ops/sparse_concat_op_test.py     | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensorflow/core/kernels/sparse_concat_op.cc b/tensorflow/core/kernels/sparse_concat_op.cc
index 940dab94d9d..02e0ecb0535 100644
--- a/tensorflow/core/kernels/sparse_concat_op.cc
+++ b/tensorflow/core/kernels/sparse_concat_op.cc
@@ -17,10 +17,10 @@ limitations under the License.
 
 #include "tensorflow/core/kernels/sparse_concat_op.h"
 
+#include <limits>
 #include <numeric>
 #include <utility>
 #include <vector>
-#include <limits>
 
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
index 9f3d1ff236e..624d724cec6 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
+++ b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
@@ -521,7 +521,6 @@ class SparseConcatTest(test.TestCase):
         concat_op = sparse_ops.sparse_concat(axis=1, sp_inputs=[sp1, sp2])
         self.evaluate(concat_op)
 
-
   def testSparseConcatOutputVolumeOverflow(self):
     # Individual volumes fit in int64, but the concatenated
     # output volume overflows int64, triggering

From b21344a8d5349f35e325f7c7ae967375c3cff43b Mon Sep 17 00:00:00 2001
From: Jaydeng75 <user.josephms@gmail.com>
Date: Fri, 31 Jul 2026 21:41:55 +0530
Subject: [PATCH 0026/1000] Fix style: lowercase variable name dim_x

---
 .../sparse_ops/sparse_concat_op_test.py            | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
index 624d724cec6..9d7ce912900 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
+++ b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
@@ -525,20 +525,20 @@ class SparseConcatTest(test.TestCase):
     # Individual volumes fit in int64, but the concatenated
     # output volume overflows int64, triggering
     # SetDimWithStatus -> RecomputeNumElements.
-    # Axis = 1, X = int64_max // 11
-    # Input 1 volume: 10 * X < int64_max
-    # Input 2 volume:  2 * X < int64_max
-    # Output volume:  12 * X > int64_max
+    # Axis = 1, dim_x = int64_max // 11
+    # Input 1 volume: 10 * dim_x < int64_max
+    # Input 2 volume:  2 * dim_x < int64_max
+    # Output volume:  12 * dim_x > int64_max
     int64_max = 9223372036854775807
-    X = int64_max // 11
+    dim_x = int64_max // 11
 
     indices1 = constant_op.constant([[0, 0, 0]], dtype=dtypes.int64)
     values1 = constant_op.constant([1.0], dtype=dtypes.float32)
-    shape1 = constant_op.constant([1, 10, X], dtype=dtypes.int64)
+    shape1 = constant_op.constant([1, 10, dim_x], dtype=dtypes.int64)
 
     indices2 = constant_op.constant([[0, 0, 0]], dtype=dtypes.int64)
     values2 = constant_op.constant([2.0], dtype=dtypes.float32)
-    shape2 = constant_op.constant([1, 2, X], dtype=dtypes.int64)
+    shape2 = constant_op.constant([1, 2, dim_x], dtype=dtypes.int64)
 
     sp1 = sparse_tensor.SparseTensor(indices1, values1, shape1)
     sp2 = sparse_tensor.SparseTensor(indices2, values2, shape2)

From cde0837fed7dca4b7050f4215701901dca66d461 Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Wed, 5 Aug 2026 14:51:57 -0700
Subject: [PATCH 0027/1000] Validate axis lower bound in Dequantize and
 QuantizeV2 kernels

The Dequantize and QuantizeV2 kernels validated only the upper bound of
the axis attr. Values below -1 were silently treated as per-tensor
(de)quantization in eager execution, while the shape functions for both
ops already reject them at graph construction time with 'axis should be
at least -1'. At the op level, axis=-1 is the per-tensor sentinel, not a
NumPy-style negative index, so values below -1 are invalid.

Add the symmetric lower-bound check to both kernels so eager execution
raises InvalidArgumentError, matching graph mode. Also documents the
axis attr in both api_defs, including the per-channel requirement that
min_range and max_range be 1-D with as many elements as the axis
dimension, which explains why axis=0 with scalar ranges is rejected by
design.

QuantizeAndDequantizeV2/V3/V4 already validate their lower bound and
are unaffected.

Fixes #100175
---
 .../api_def/base_api/api_def_Dequantize.pbtxt | 11 ++++
 .../api_def/base_api/api_def_QuantizeV2.pbtxt | 11 ++++
 tensorflow/core/kernels/dequantize_op.cc      |  4 ++
 tensorflow/core/kernels/quantize_op.cc        |  5 ++
 .../quantization_ops/quantization_ops_test.py | 53 +++++++++++++++++++
 5 files changed, 84 insertions(+)

diff --git a/tensorflow/core/api_def/base_api/api_def_Dequantize.pbtxt b/tensorflow/core/api_def/base_api/api_def_Dequantize.pbtxt
index 030b98c369d..108025157e7 100644
--- a/tensorflow/core/api_def/base_api/api_def_Dequantize.pbtxt
+++ b/tensorflow/core/api_def/base_api/api_def_Dequantize.pbtxt
@@ -10,6 +10,17 @@ END
     name: "max_range"
     description: <<END
 The maximum scalar value possibly produced for the input.
+END
+  }
+  attr {
+    name: "axis"
+    description: <<END
+The index of the input dimension along which the dequantization ranges
+vary, for per-channel dequantization. When `axis` is -1 (the default),
+dequantization is per-tensor and `min_range`/`max_range` must be scalars.
+Otherwise `min_range` and `max_range` must be 1-D tensors with as many
+elements as the input's `axis` dimension. Negative values other than -1
+are invalid.
 END
   }
   attr {
diff --git a/tensorflow/core/api_def/base_api/api_def_QuantizeV2.pbtxt b/tensorflow/core/api_def/base_api/api_def_QuantizeV2.pbtxt
index 2645056f651..17d50497893 100644
--- a/tensorflow/core/api_def/base_api/api_def_QuantizeV2.pbtxt
+++ b/tensorflow/core/api_def/base_api/api_def_QuantizeV2.pbtxt
@@ -40,6 +40,17 @@ The final quantization range maximum, used to clip input values before scaling
 and rounding them to quantized values.
 If the `axis` attribute is specified, this will be a 1-D tensor whose size
 matches the `axis` dimension of the input and output tensors.
+END
+  }
+  attr {
+    name: "axis"
+    description: <<END
+The index of the input dimension along which the quantization ranges
+vary, for per-channel quantization. When `axis` is -1 (the default),
+quantization is per-tensor and `min_range`/`max_range` must be scalars.
+Otherwise `min_range` and `max_range` must be 1-D tensors with as many
+elements as the input's `axis` dimension. Negative values other than -1
+are invalid.
 END
   }
   attr {
diff --git a/tensorflow/core/kernels/dequantize_op.cc b/tensorflow/core/kernels/dequantize_op.cc
index 0302668e9e4..782ebc2a74b 100644
--- a/tensorflow/core/kernels/dequantize_op.cc
+++ b/tensorflow/core/kernels/dequantize_op.cc
@@ -97,6 +97,10 @@ class DequantizeOp : public OpKernel {
     const Tensor& input_min_tensor = ctx->input(1);
     const Tensor& input_max_tensor = ctx->input(2);
 
+    OP_REQUIRES(ctx, axis_ >= -1,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Axis must be at least -1, where -1 selects per-tensor "
+                    "dequantization, got ", axis_)));
     OP_REQUIRES(ctx, axis_ < input.dims(),
                 absl::InvalidArgumentError(
                     absl::StrCat("Axis must be less than input dimension(",
diff --git a/tensorflow/core/kernels/quantize_op.cc b/tensorflow/core/kernels/quantize_op.cc
index 613fd99f193..0ad6aa4afab 100644
--- a/tensorflow/core/kernels/quantize_op.cc
+++ b/tensorflow/core/kernels/quantize_op.cc
@@ -114,6 +114,11 @@ class QuantizeV2Op : public OpKernel {
     const Tensor& input_min_range = ctx->input(1);
     const Tensor& input_max_range = ctx->input(2);
 
+    OP_REQUIRES(ctx, axis_ >= -1,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "Axis must be at least -1, where -1 selects per-tensor "
+                    "quantization, got ", axis_)));
+
     int num_slices = 1;
     if (axis_ > -1) {
       OP_REQUIRES(
diff --git a/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py b/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
index 332d67ca763..d96a4c0c621 100644
--- a/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
+++ b/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
@@ -22,6 +22,7 @@ from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_array_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import nn_ops
 from tensorflow.python.platform import googletest
@@ -409,6 +410,58 @@ class QuantizeDownAndShrinkRangeOpTest(test_util.TensorFlowTestCase):
               out_type=dtypes.quint8))
 
 
+class DequantizeOpTest(test_util.TensorFlowTestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_axis_below_minus_one_raises(self):
+    # Regression test for GitHub issue 100175: the kernel accepted axis
+    # values below -1 in eager mode and silently treated them as per-tensor
+    # dequantization, while the shape function already rejected them at
+    # graph construction time.
+    with self.assertRaisesRegex(
+        (ValueError, errors.InvalidArgumentError), "at least -1"):
+      self.evaluate(
+          gen_array_ops.dequantize(
+              input=constant_op.constant(
+                  np.uint8(74), shape=[3], dtype=dtypes.quint8),
+              min_range=0.0,
+              max_range=1.0,
+              mode="MIN_COMBINED",
+              axis=-5))
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_per_tensor_sentinel_axis_valid(self):
+    # axis=-1 is the per-tensor sentinel at the op level and must keep
+    # working with scalar ranges.
+    output = self.evaluate(
+        gen_array_ops.dequantize(
+            input=constant_op.constant(
+                np.uint8(74), shape=[3], dtype=dtypes.quint8),
+            min_range=0.0,
+            max_range=1.0,
+            mode="MIN_COMBINED",
+            axis=-1))
+    self.assertEqual((3,), output.shape)
+
+
+class QuantizeV2OpTest(test_util.TensorFlowTestCase):
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_axis_below_minus_one_raises(self):
+    # Same missing lower-bound validation as Dequantize; see GitHub issue
+    # 100175.
+    with self.assertRaisesRegex(
+        (ValueError, errors.InvalidArgumentError), "at least -1"):
+      self.evaluate(
+          gen_array_ops.quantize_v2(
+              input=constant_op.constant([0.1, 0.5, 0.9]),
+              min_range=0.0,
+              max_range=1.0,
+              T=dtypes.quint8,
+              mode="MIN_COMBINED",
+              axis=-5))
+
+
 class QuantizeAndDequantizeV3OpTest(test_util.TensorFlowTestCase):
 
   @test_util.run_in_graph_and_eager_modes

From 7c349f69d7532f0f428c0429338b2676ce5be775 Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Wed, 12 Aug 2026 14:01:33 +0530
Subject: [PATCH 0028/1000] validate untrusted tensor_shape proto in the first
 two passes of UncompressElement

---
 tensorflow/core/data/compression_utils.cc      | 12 +++++++++---
 tensorflow/core/data/compression_utils_test.cc | 17 +++++++++++++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/data/compression_utils.cc b/tensorflow/core/data/compression_utils.cc
index 5ebe44f3dcc..fc34a7f7b72 100644
--- a/tensorflow/core/data/compression_utils.cc
+++ b/tensorflow/core/data/compression_utils.cc
@@ -164,7 +164,9 @@ absl::Status UncompressElement(const CompressedElement& compressed,
       ++num_string_tensors;
       num_string_tensor_strings += metadata.uncompressed_bytes_size();
     } else {
-      TensorShape shape(metadata.tensor_shape());
+      TensorShape shape;
+      TF_RETURN_IF_ERROR(
+          TensorShape::BuildTensorShape(metadata.tensor_shape(), &shape));
       int64_t num_elements = shape.num_elements();
       if (num_elements > 0 && metadata.uncompressed_bytes_size() == 0) {
         return absl::InvalidArgumentError(
@@ -188,7 +190,9 @@ absl::Status UncompressElement(const CompressedElement& compressed,
   char* nonmemcpyable_pos = nonmemcpyable.mdata();
   for (const auto& metadata : compressed.component_metadata()) {
     if (DataTypeCanUseMemcpy(metadata.dtype())) {
-      TensorShape shape(metadata.tensor_shape());
+      TensorShape shape;
+      TF_RETURN_IF_ERROR(
+          TensorShape::BuildTensorShape(metadata.tensor_shape(), &shape));
       int64_t num_elements = shape.num_elements();
       out->emplace_back(metadata.dtype(), metadata.tensor_shape());
       TensorBuffer* buffer = DMAHelper::buffer(&out->back());
@@ -214,7 +218,9 @@ absl::Status UncompressElement(const CompressedElement& compressed,
         iov.Add(flats.data()[i].mdata(), metadata.uncompressed_bytes(i));
       }
     } else {
-      TensorShape shape(metadata.tensor_shape());
+      TensorShape shape;
+      TF_RETURN_IF_ERROR(
+          TensorShape::BuildTensorShape(metadata.tensor_shape(), &shape));
       int64_t num_elements = shape.num_elements();
       out->emplace_back();
       if (num_elements > 0) {
diff --git a/tensorflow/core/data/compression_utils_test.cc b/tensorflow/core/data/compression_utils_test.cc
index 6b05e20d4be..6f32472e65e 100644
--- a/tensorflow/core/data/compression_utils_test.cc
+++ b/tensorflow/core/data/compression_utils_test.cc
@@ -65,6 +65,23 @@ TEST(CompressionUtilsTest, ZeroElementNonMemcpyableComponent) {
                                      HasSubstr("Zero-element component")));
 }
 
+TEST(CompressionUtilsTest, MalformedTensorShape) {
+  // A malformed shape proto (negative dimension) must be rejected with a
+  // status instead of aborting the direct `TensorShape` constructor.
+  std::vector<Tensor> empty_element;
+  CompressedElement compressed;
+  TF_ASSERT_OK(CompressElement(empty_element, &compressed));
+
+  CompressedComponentMetadata* metadata =
+      compressed.mutable_component_metadata()->Add();
+  metadata->set_dtype(DT_INT64);
+  metadata->mutable_tensor_shape()->add_dim()->set_size(-1);
+
+  std::vector<Tensor> element;
+  EXPECT_THAT(UncompressElement(compressed, &element),
+              absl_testing::StatusIs(error::INVALID_ARGUMENT));
+}
+
 std::vector<std::vector<Tensor>> TestCases() {
   return {
       // Single int64.

From 63d436c82b85959e4b8720498249c1d9b70de903 Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Thu, 13 Aug 2026 20:42:09 +0530
Subject: [PATCH 0029/1000] validate DT_STRING tensor_shape and reuse built
 shape in UncompressElement

---
 tensorflow/core/data/compression_utils.cc      | 13 ++++++++-----
 tensorflow/core/data/compression_utils_test.cc | 18 ++++++++++++++++++
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/tensorflow/core/data/compression_utils.cc b/tensorflow/core/data/compression_utils.cc
index fc34a7f7b72..e26860bd94d 100644
--- a/tensorflow/core/data/compression_utils.cc
+++ b/tensorflow/core/data/compression_utils.cc
@@ -160,13 +160,13 @@ absl::Status UncompressElement(const CompressedElement& compressed,
   size_t num_string_tensor_strings = 0;
   size_t total_nonmemcpyable_size = 0;
   for (const auto& metadata : compressed.component_metadata()) {
+    TensorShape shape;
+    TF_RETURN_IF_ERROR(
+        TensorShape::BuildTensorShape(metadata.tensor_shape(), &shape));
     if (metadata.dtype() == DT_STRING) {
       ++num_string_tensors;
       num_string_tensor_strings += metadata.uncompressed_bytes_size();
     } else {
-      TensorShape shape;
-      TF_RETURN_IF_ERROR(
-          TensorShape::BuildTensorShape(metadata.tensor_shape(), &shape));
       int64_t num_elements = shape.num_elements();
       if (num_elements > 0 && metadata.uncompressed_bytes_size() == 0) {
         return absl::InvalidArgumentError(
@@ -194,7 +194,7 @@ absl::Status UncompressElement(const CompressedElement& compressed,
       TF_RETURN_IF_ERROR(
           TensorShape::BuildTensorShape(metadata.tensor_shape(), &shape));
       int64_t num_elements = shape.num_elements();
-      out->emplace_back(metadata.dtype(), metadata.tensor_shape());
+      out->emplace_back(metadata.dtype(), shape);
       TensorBuffer* buffer = DMAHelper::buffer(&out->back());
       if (buffer && num_elements > 0) {
         if (metadata.uncompressed_bytes(0) > buffer->size()) {
@@ -205,7 +205,10 @@ absl::Status UncompressElement(const CompressedElement& compressed,
         iov.Add(buffer->data(), metadata.uncompressed_bytes(0));
       }
     } else if (metadata.dtype() == DT_STRING) {
-      out->emplace_back(metadata.dtype(), metadata.tensor_shape());
+      TensorShape shape;
+      TF_RETURN_IF_ERROR(
+          TensorShape::BuildTensorShape(metadata.tensor_shape(), &shape));
+      out->emplace_back(metadata.dtype(), shape);
       const auto& flats = out->back().unaligned_flat<tstring>();
       if (metadata.uncompressed_bytes_size() > flats.size()) {
         return absl::InvalidArgumentError(absl::StrCat(
diff --git a/tensorflow/core/data/compression_utils_test.cc b/tensorflow/core/data/compression_utils_test.cc
index 6f32472e65e..8ba7e593c07 100644
--- a/tensorflow/core/data/compression_utils_test.cc
+++ b/tensorflow/core/data/compression_utils_test.cc
@@ -82,6 +82,24 @@ TEST(CompressionUtilsTest, MalformedTensorShape) {
               absl_testing::StatusIs(error::INVALID_ARGUMENT));
 }
 
+TEST(CompressionUtilsTest, MalformedTensorShapeString) {
+  // A `DT_STRING` component with a malformed shape proto (negative dimension)
+  // must also be rejected with a status rather than aborting the direct
+  // `TensorShape` constructor when the string tensor is built.
+  std::vector<Tensor> empty_element;
+  CompressedElement compressed;
+  TF_ASSERT_OK(CompressElement(empty_element, &compressed));
+
+  CompressedComponentMetadata* metadata =
+      compressed.mutable_component_metadata()->Add();
+  metadata->set_dtype(DT_STRING);
+  metadata->mutable_tensor_shape()->add_dim()->set_size(-1);
+
+  std::vector<Tensor> element;
+  EXPECT_THAT(UncompressElement(compressed, &element),
+              absl_testing::StatusIs(error::INVALID_ARGUMENT));
+}
+
 std::vector<std::vector<Tensor>> TestCases() {
   return {
       // Single int64.

From 6ed6070bdbaad4c5319ac0967a96ae7096bf34a6 Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Tue, 18 Aug 2026 17:19:32 +0530
Subject: [PATCH 0030/1000] guard empty metadata in VariantTensorDataReader
 constructor

---
 tensorflow/core/data/serialization_utils.cc      |  4 ++++
 tensorflow/core/data/serialization_utils_test.cc | 16 ++++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/tensorflow/core/data/serialization_utils.cc b/tensorflow/core/data/serialization_utils.cc
index 65a58a7a1e9..13964841c02 100644
--- a/tensorflow/core/data/serialization_utils.cc
+++ b/tensorflow/core/data/serialization_utils.cc
@@ -188,6 +188,10 @@ VariantTensorDataReader::VariantTensorDataReader(
     std::string metadata;
     d->get_metadata(&metadata);
     auto keys = str_util::Split(metadata, kDelimiter, str_util::SkipEmpty());
+    if (keys.empty()) {
+      LOG(WARNING) << "Found empty metadata in VariantTensorData; skipping.";
+      continue;
+    }
     const std::string name = keys[0];
     data_[name] = d;
     map_[name] = std::map<std::string, size_t>();
diff --git a/tensorflow/core/data/serialization_utils_test.cc b/tensorflow/core/data/serialization_utils_test.cc
index ce8140f5463..e6c515e6392 100644
--- a/tensorflow/core/data/serialization_utils_test.cc
+++ b/tensorflow/core/data/serialization_utils_test.cc
@@ -185,6 +185,22 @@ TEST(SerializationUtilsTest, VariantTensorDataReaderOOB) {
             error::OUT_OF_RANGE);
 }
 
+TEST(SerializationUtilsTest, VariantTensorDataReaderEmptyMetadata) {
+  // Metadata that splits into no keys (empty or delimiter-only) must not cause
+  // the reader to index the first key out of bounds while building its maps.
+  for (const std::string& metadata : {std::string(""), std::string("@@")}) {
+    VariantTensorData data;
+    data.metadata_ = metadata;
+    data.tensors_.push_back(Tensor(DT_INT64, {1}));
+    std::vector<const VariantTensorData*> reader_data;
+    reader_data.push_back(&data);
+    VariantTensorDataReader reader(reader_data);
+    int64_t val_int64;
+    EXPECT_EQ(reader.ReadScalar("Iterator", "key1", &val_int64).code(),
+              error::NOT_FOUND);
+  }
+}
+
 class ParameterizedIteratorStateVariantTest
     : public DatasetOpsTestBase,
       public ::testing::WithParamInterface<std::vector<Tensor>> {

From bba2c4007255ef8bc8d2786012bed7fe986cc83f Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Wed, 19 Aug 2026 13:35:15 +0530
Subject: [PATCH 0031/1000] Use VLOG(1) instead of LOG(WARNING) for empty
 metadata skip

---
 tensorflow/core/data/serialization_utils.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/data/serialization_utils.cc b/tensorflow/core/data/serialization_utils.cc
index 13964841c02..d41511e3911 100644
--- a/tensorflow/core/data/serialization_utils.cc
+++ b/tensorflow/core/data/serialization_utils.cc
@@ -189,7 +189,7 @@ VariantTensorDataReader::VariantTensorDataReader(
     d->get_metadata(&metadata);
     auto keys = str_util::Split(metadata, kDelimiter, str_util::SkipEmpty());
     if (keys.empty()) {
-      LOG(WARNING) << "Found empty metadata in VariantTensorData; skipping.";
+      VLOG(1) << "Found empty metadata in VariantTensorData; skipping.";
       continue;
     }
     const std::string name = keys[0];

From d4d5af2408c96b9722d423d8e027b7137ca14ecf Mon Sep 17 00:00:00 2001
From: AshiteshSingh <raghuvesh1285@gmail.com>
Date: Sun, 23 Aug 2026 09:02:58 +0530
Subject: [PATCH 0032/1000] Fix min_input rank 1 to rank 0 for MKL quantized
 conv tests

---
 tensorflow/core/kernels/mkl/mkl_fused_ops_test.cc    | 12 ++++++------
 .../mkl/mkl_quantized_conv_ops_perchannel_test.cc    |  4 ++--
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/tensorflow/core/kernels/mkl/mkl_fused_ops_test.cc b/tensorflow/core/kernels/mkl/mkl_fused_ops_test.cc
index 6d709794b6f..4723322ceda 100644
--- a/tensorflow/core/kernels/mkl/mkl_fused_ops_test.cc
+++ b/tensorflow/core/kernels/mkl/mkl_fused_ops_test.cc
@@ -1233,22 +1233,22 @@ class BiasCacheTest : public OpsTestBase {
     Tensor bias(DT_FLOAT, {kFilterCount});
     test::FillValues<float>(&bias, {1});
 
-    Tensor min_input(DT_FLOAT, {1});
+    Tensor min_input(DT_FLOAT, {});
     test::FillValues<float>(&min_input, {1});
 
-    Tensor max_input(DT_FLOAT, {1});
+    Tensor max_input(DT_FLOAT, {});
     test::FillValues<float>(&max_input, {1});
 
-    Tensor min_filter(DT_FLOAT, {1});
+    Tensor min_filter(DT_FLOAT, {});
     test::FillValues<float>(&min_filter, {1});
 
-    Tensor max_filter(DT_FLOAT, {1});
+    Tensor max_filter(DT_FLOAT, {});
     test::FillValues<float>(&max_filter, {1});
 
-    Tensor min_output(DT_FLOAT, {1});
+    Tensor min_output(DT_FLOAT, {});
     test::FillValues<float>(&min_output, {1});
 
-    Tensor max_output(DT_FLOAT, {1});
+    Tensor max_output(DT_FLOAT, {});
     test::FillValues<float>(&max_output, {1});
 
     Tensor expected(DT_QUINT8, TensorShape({1, 1, 2, 1}));
diff --git a/tensorflow/core/kernels/mkl/mkl_quantized_conv_ops_perchannel_test.cc b/tensorflow/core/kernels/mkl/mkl_quantized_conv_ops_perchannel_test.cc
index 37ab6c2fb76..85d4055ff94 100644
--- a/tensorflow/core/kernels/mkl/mkl_quantized_conv_ops_perchannel_test.cc
+++ b/tensorflow/core/kernels/mkl/mkl_quantized_conv_ops_perchannel_test.cc
@@ -129,8 +129,8 @@ class QuantizedConv2DPerChannelTest : public OpsTestBase {
                               image_quantized.flat<quint8>());
     AddInputFromArray<qint8>(filter_quantized.shape(),
                              filter_quantized.flat<qint8>());
-    AddInputFromArray<float>(TensorShape({1}), {image_min});
-    AddInputFromArray<float>(TensorShape({1}), {image_max});
+    AddInputFromArray<float>(TensorShape({}), {image_min});
+    AddInputFromArray<float>(TensorShape({}), {image_max});
     AddInputFromArray<float>(TensorShape({2}), {filter_min, filter_min});
     AddInputFromArray<float>(TensorShape({2}), {filter_max, filter_max});
 

From 2aa817f2d9bff903416e0243265dff1a72ccffbb Mon Sep 17 00:00:00 2001
From: AshiteshSingh <raghuvesh1285@gmail.com>
Date: Sun, 23 Aug 2026 09:26:40 +0530
Subject: [PATCH 0033/1000] Trigger GitHub Actions


From 40ae662e351869f2f8726cd4e23e1adeed964822 Mon Sep 17 00:00:00 2001
From: AshiteshSingh <raghuvesh1285@gmail.com>
Date: Mon, 24 Aug 2026 18:34:15 +0530
Subject: [PATCH 0034/1000] Fix MSVC ambiguous TensorShape conversion in MKL
 tests

---
 tensorflow/core/kernels/mkl/mkl_fused_ops_test.cc | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensorflow/core/kernels/mkl/mkl_fused_ops_test.cc b/tensorflow/core/kernels/mkl/mkl_fused_ops_test.cc
index 4723322ceda..3e3e4a45a7c 100644
--- a/tensorflow/core/kernels/mkl/mkl_fused_ops_test.cc
+++ b/tensorflow/core/kernels/mkl/mkl_fused_ops_test.cc
@@ -1233,22 +1233,22 @@ class BiasCacheTest : public OpsTestBase {
     Tensor bias(DT_FLOAT, {kFilterCount});
     test::FillValues<float>(&bias, {1});
 
-    Tensor min_input(DT_FLOAT, {});
+    Tensor min_input(DT_FLOAT, TensorShape({}));
     test::FillValues<float>(&min_input, {1});
 
-    Tensor max_input(DT_FLOAT, {});
+    Tensor max_input(DT_FLOAT, TensorShape({}));
     test::FillValues<float>(&max_input, {1});
 
-    Tensor min_filter(DT_FLOAT, {});
+    Tensor min_filter(DT_FLOAT, TensorShape({}));
     test::FillValues<float>(&min_filter, {1});
 
-    Tensor max_filter(DT_FLOAT, {});
+    Tensor max_filter(DT_FLOAT, TensorShape({}));
     test::FillValues<float>(&max_filter, {1});
 
-    Tensor min_output(DT_FLOAT, {});
+    Tensor min_output(DT_FLOAT, TensorShape({}));
     test::FillValues<float>(&min_output, {1});
 
-    Tensor max_output(DT_FLOAT, {});
+    Tensor max_output(DT_FLOAT, TensorShape({}));
     test::FillValues<float>(&max_output, {1});
 
     Tensor expected(DT_QUINT8, TensorShape({1, 1, 2, 1}));

From aa9c130a2eac86c4af65c2964fac959fdbc84ce6 Mon Sep 17 00:00:00 2001
From: PrabinDevkota <prabindev789@gmail.com>
Date: Wed, 26 Aug 2026 21:34:25 +0530
Subject: [PATCH 0035/1000] Hide bundled hwloc symbols so they cannot interpose
 system PCI discovery.

---
 tensorflow/tf_framework_version_script.lds    | 4 ++++
 third_party/xla/third_party/hwloc/hwloc.BUILD | 6 +++++-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/tensorflow/tf_framework_version_script.lds b/tensorflow/tf_framework_version_script.lds
index 3e0dca3e587..b0a154312dc 100644
--- a/tensorflow/tf_framework_version_script.lds
+++ b/tensorflow/tf_framework_version_script.lds
@@ -2,4 +2,8 @@ VERS_1.0 {
   # Hide libjpeg symbols to avoid symbol conflict with OpenCV
   local:
     *mlir*;
+    # Do not re-export bundled hwloc. Its objects also land in other DSOs
+    # (e.g. _xla_ops.so); hiding at compile time is the complete fix, this
+    # is extra protection for libtensorflow_framework.so (tensorflow#125854).
+    hwloc*;
 };
diff --git a/third_party/xla/third_party/hwloc/hwloc.BUILD b/third_party/xla/third_party/hwloc/hwloc.BUILD
index 93a63e9445e..36ead1a3fba 100644
--- a/third_party/xla/third_party/hwloc/hwloc.BUILD
+++ b/third_party/xla/third_party/hwloc/hwloc.BUILD
@@ -137,7 +137,11 @@ _INCLUDE_PRIVATE_HWLOC_AUTOIGEN_CONFIG_H_COMMON_SUBS = {
     "#undef HAVE_X11_XLIB_H": "#define HAVE_X11_XLIB_H 1",
     "#undef HAVE_X11_XUTIL_H": "#define HAVE_X11_XUTIL_H 1",
     "#undef HAVE___PROGNAME": "#define HAVE___PROGNAME 1",
-    "#undef HWLOC_C_HAVE_VISIBILITY": "#define HWLOC_C_HAVE_VISIBILITY 1",
+    # Must stay 0. If this is 1, hwloc annotates public APIs with
+    # visibility("default"), which overrides -fvisibility=hidden and leaks
+    # hwloc_* from every DSO that statically links this library
+    # (libtensorflow_framework.so, _xla_ops.so, ...; tensorflow#125854).
+    "#undef HWLOC_C_HAVE_VISIBILITY": "#define HWLOC_C_HAVE_VISIBILITY 0",
     "#undef HWLOC_HAVE_ATTRIBUTE_ALIGNED": "#define HWLOC_HAVE_ATTRIBUTE_ALIGNED 1",
     "#undef HWLOC_HAVE_ATTRIBUTE_ALWAYS_INLINE": "#define HWLOC_HAVE_ATTRIBUTE_ALWAYS_INLINE 1",
     "#undef HWLOC_HAVE_ATTRIBUTE_COLD": "#define HWLOC_HAVE_ATTRIBUTE_COLD 1",

From 3fde80054a2eb4dbb2bc69e3f3d0de73a5eddd4b Mon Sep 17 00:00:00 2001
From: PrabinDevkota <prabindev789@gmail.com>
Date: Thu, 27 Aug 2026 14:02:37 +0530
Subject: [PATCH 0036/1000] Rename bundled hwloc symbols so they cannot
 interpose system PCI discovery.

---
 tensorflow/tf_framework_version_script.lds    |  6 +++---
 third_party/xla/third_party/hwloc/hwloc.BUILD | 21 +++++++++----------
 2 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/tensorflow/tf_framework_version_script.lds b/tensorflow/tf_framework_version_script.lds
index c4c381fb099..ee7a8741474 100644
--- a/tensorflow/tf_framework_version_script.lds
+++ b/tensorflow/tf_framework_version_script.lds
@@ -2,8 +2,8 @@ tensorflow {
   global:
     *;
   local:
-    # Do not re-export bundled hwloc. Its objects also land in other DSOs
-    # (e.g. _xla_ops.so); hiding at compile time is the complete fix, this
-    # is extra protection for libtensorflow_framework.so (tensorflow#125854).
+    # Hide leftover unprefixed hwloc_* names. Bundled APIs are renamed
+    # tf_hwloc_* (HWLOC_SYM_TRANSFORM) so they do not interpose system hwloc
+    # and can still be shared across TF DSOs (tensorflow#125854).
     hwloc*;
 };
diff --git a/third_party/xla/third_party/hwloc/hwloc.BUILD b/third_party/xla/third_party/hwloc/hwloc.BUILD
index 36ead1a3fba..7f80f179b8e 100644
--- a/third_party/xla/third_party/hwloc/hwloc.BUILD
+++ b/third_party/xla/third_party/hwloc/hwloc.BUILD
@@ -50,9 +50,12 @@ _INCLUDE_HWLOC_AUTOIGEN_CONFIG_H_COMMON_SUBS = {
     "#undef hwloc_pid_t": "#define hwloc_pid_t pid_t",
     "#undef hwloc_thread_t": "#define hwloc_thread_t pthread_t",
     "#  undef HWLOC_HAVE_STDINT_H": "#  define HWLOC_HAVE_STDINT_H 1",
-    "#undef HWLOC_SYM_TRANSFORM": "#define HWLOC_SYM_TRANSFORM 0",
-    "#undef HWLOC_SYM_PREFIX_CAPS": "#define HWLOC_SYM_PREFIX_CAPS HWLOC_",
-    "#undef HWLOC_SYM_PREFIX": "#define HWLOC_SYM_PREFIX hwloc_",
+    # Prefix bundled APIs (tf_hwloc_*) so they cannot interpose system hwloc
+    # (tensorflow#125854). Public symbols must stay visible across TF DSOs
+    # (libplatform_port.so vs libtensorflow_framework.so).
+    "#undef HWLOC_SYM_TRANSFORM": "#define HWLOC_SYM_TRANSFORM 1",
+    "#undef HWLOC_SYM_PREFIX_CAPS": "#define HWLOC_SYM_PREFIX_CAPS TF_",
+    "#undef HWLOC_SYM_PREFIX": "#define HWLOC_SYM_PREFIX tf_",
 }
 
 _INCLUDE_HWLOC_AUTOIGEN_CONFIG_H_LINUX_SUBS = dict(_INCLUDE_HWLOC_AUTOIGEN_CONFIG_H_COMMON_SUBS)
@@ -137,11 +140,7 @@ _INCLUDE_PRIVATE_HWLOC_AUTOIGEN_CONFIG_H_COMMON_SUBS = {
     "#undef HAVE_X11_XLIB_H": "#define HAVE_X11_XLIB_H 1",
     "#undef HAVE_X11_XUTIL_H": "#define HAVE_X11_XUTIL_H 1",
     "#undef HAVE___PROGNAME": "#define HAVE___PROGNAME 1",
-    # Must stay 0. If this is 1, hwloc annotates public APIs with
-    # visibility("default"), which overrides -fvisibility=hidden and leaks
-    # hwloc_* from every DSO that statically links this library
-    # (libtensorflow_framework.so, _xla_ops.so, ...; tensorflow#125854).
-    "#undef HWLOC_C_HAVE_VISIBILITY": "#define HWLOC_C_HAVE_VISIBILITY 0",
+    "#undef HWLOC_C_HAVE_VISIBILITY": "#define HWLOC_C_HAVE_VISIBILITY 1",
     "#undef HWLOC_HAVE_ATTRIBUTE_ALIGNED": "#define HWLOC_HAVE_ATTRIBUTE_ALIGNED 1",
     "#undef HWLOC_HAVE_ATTRIBUTE_ALWAYS_INLINE": "#define HWLOC_HAVE_ATTRIBUTE_ALWAYS_INLINE 1",
     "#undef HWLOC_HAVE_ATTRIBUTE_COLD": "#define HWLOC_HAVE_ATTRIBUTE_COLD 1",
@@ -180,9 +179,9 @@ _INCLUDE_PRIVATE_HWLOC_AUTOIGEN_CONFIG_H_COMMON_SUBS = {
     "#undef HWLOC_HAVE_X86_CPUID": "#define HWLOC_HAVE_X86_CPUID 1",
     "#undef HWLOC_SIZEOF_UNSIGNED_INT": "#define HWLOC_SIZEOF_UNSIGNED_INT 4",
     "#undef HWLOC_SIZEOF_UNSIGNED_LONG": "#define HWLOC_SIZEOF_UNSIGNED_LONG 8",
-    "#undef HWLOC_SYM_PREFIX_CAPS": "#define HWLOC_SYM_PREFIX_CAPS HWLOC_",
-    "#undef HWLOC_SYM_PREFIX": "#define HWLOC_SYM_PREFIX hwloc_",
-    "#undef HWLOC_SYM_TRANSFORM": "#define HWLOC_SYM_TRANSFORM 0",
+    "#undef HWLOC_SYM_PREFIX_CAPS": "#define HWLOC_SYM_PREFIX_CAPS TF_",
+    "#undef HWLOC_SYM_PREFIX": "#define HWLOC_SYM_PREFIX tf_",
+    "#undef HWLOC_SYM_TRANSFORM": "#define HWLOC_SYM_TRANSFORM 1",
     "#undef HWLOC_USE_NCURSES": "#define HWLOC_USE_NCURSES 1",
     "#undef HWLOC_VERSION_GREEK": "#define HWLOC_VERSION_GREEK \"\"",
     "#undef HWLOC_VERSION_MAJOR": "#define HWLOC_VERSION_MAJOR 2",

From d9060a22d5b5d2a43b92c0699edfd6eea511863a Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Thu, 27 Aug 2026 15:59:38 +0530
Subject: [PATCH 0037/1000] validate decoded element shape and dtype in
 TensorList::Decode

---
 tensorflow/core/kernels/tensor_list.cc      | 19 +++++++++++
 tensorflow/core/kernels/tensor_list_test.cc | 36 +++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/tensorflow/core/kernels/tensor_list.cc b/tensorflow/core/kernels/tensor_list.cc
index c0c211fd892..8f49f0cd33f 100644
--- a/tensorflow/core/kernels/tensor_list.cc
+++ b/tensorflow/core/kernels/tensor_list.cc
@@ -127,6 +127,25 @@ bool TensorList::Decode(const VariantTensorData& data) {
 
   const PartialTensorShape decoded_element_shape(element_shape_proto);
 
+  // Every element stored in a list must match its element_dtype and be
+  // compatible with its element_shape; the mutation paths (TensorListPushBack,
+  // TensorListSetItem) enforce this on insertion. The serialized data is
+  // untrusted, so re-establish the same invariant here. Consumers such as
+  // TensorListStack and TensorListGather skip the per-element shape check when
+  // element_shape is fully defined and then treat each element's count as the
+  // output element count, so a decoded element that does not match would
+  // overrun the output buffer.
+  for (const Tensor& t : decoded_tensors) {
+    if (t.dtype() == DT_INVALID) continue;  // Unset element placeholder.
+    if (decoded_element_dtype != DT_INVALID &&
+        t.dtype() != decoded_element_dtype) {
+      return false;
+    }
+    if (!decoded_element_shape.IsCompatibleWith(t.shape())) {
+      return false;
+    }
+  }
+
   element_dtype = decoded_element_dtype;
   max_num_elements = decoded_max_num_elements;
   tensors() = std::move(decoded_tensors);
diff --git a/tensorflow/core/kernels/tensor_list_test.cc b/tensorflow/core/kernels/tensor_list_test.cc
index a96cca3336c..bafd0fa7935 100644
--- a/tensorflow/core/kernels/tensor_list_test.cc
+++ b/tensorflow/core/kernels/tensor_list_test.cc
@@ -146,6 +146,42 @@ TEST(TensorListTest, DecodeValidRoundTripCompatibility) {
   EXPECT_EQ(decoded.tensors()[3].scalar<int32_t>()(), 9);
 }
 
+TEST(TensorListTest, DecodeRejectsElementLargerThanFullyDefinedShape) {
+  // element_shape is fully defined as [1], but the stored element holds 4
+  // values. Consumers (e.g. TensorListStack) skip the per-element shape check
+  // when element_shape is fully defined and size the output from element_shape,
+  // so this element would overrun the output buffer.
+  TensorShapeProto shape;
+  shape.add_dim()->set_size(1);
+  VariantTensorData data;
+  data.set_metadata(TensorListMetadata(
+      /*invalid_indices=*/{}, static_cast<uint64_t>(DT_FLOAT),
+      std::numeric_limits<uint64_t>::max(), shape.SerializeAsString()));
+  *data.add_tensors() = Tensor(DT_FLOAT, TensorShape({4}));
+  EXPECT_FALSE(TensorList().Decode(data));
+}
+
+TEST(TensorListTest, DecodeRejectsElementDtypeMismatch) {
+  TensorShapeProto shape;
+  VariantTensorData data;
+  data.set_metadata(TensorListMetadata(
+      /*invalid_indices=*/{}, static_cast<uint64_t>(DT_FLOAT),
+      std::numeric_limits<uint64_t>::max(), shape.SerializeAsString()));
+  *data.add_tensors() = ScalarInt32Tensor(1);
+  EXPECT_FALSE(TensorList().Decode(data));
+}
+
+TEST(TensorListTest, DecodeAcceptsElementCompatibleWithPartialShape) {
+  TensorShapeProto shape;
+  shape.add_dim()->set_size(-1);
+  VariantTensorData data;
+  data.set_metadata(TensorListMetadata(
+      /*invalid_indices=*/{}, static_cast<uint64_t>(DT_FLOAT),
+      std::numeric_limits<uint64_t>::max(), shape.SerializeAsString()));
+  *data.add_tensors() = Tensor(DT_FLOAT, TensorShape({4}));
+  EXPECT_TRUE(TensorList().Decode(data));
+}
+
 TEST(TensorListTest, DecodeRejectsOutOfRangeMaxNumElements) {
   TensorShapeProto shape;
   VariantTensorData data;

From 09cd294c756b2db9ca7a2a723a7ee11ca83c1720 Mon Sep 17 00:00:00 2001
From: endorphin13 <elshad.toklayev.business@gmail.com>
Date: Fri, 28 Aug 2026 15:07:06 +0300
Subject: [PATCH 0038/1000] Validate snapshot tensor shapes before constructing
 a TensorShape

CustomReader::SnappyUncompress builds a TensorShape straight from
SnapshotTensorMetadata:

    const auto& tensor_metadata = metadata->tensor_metadata(i);
    if (simple_tensor_mask_[i]) {
      TensorShape shape(tensor_metadata.tensor_shape());

`metadata` is parsed from a record in the snapshot file, so its dimensions
are file-controlled. The TensorShape(TensorShapeProto) constructor
CHECK-fails on a negative or overflowing dimension:

    F tensor_shape.cc:416] Check failed: size >= 0 (-5 vs. 0)
        @ tensorflow::TensorShapeBase<>::AddDim()
        @ tensorflow::TensorShapeBase<>::TensorShapeBase()

which aborts the process rather than reporting a corrupt snapshot. Every
other error in this function is reported as DataLossError or
InvalidArgumentError, and the sibling branch a few lines below already
rejects a negative tensor_size_bytes, so the abort is inconsistent with
the surrounding handling.

Build the shape through TensorShape::BuildTensorShape, which returns a
status. Callers in tensorflow/core/framework/tensor.cc and
tensorflow/core/ops/array_ops.cc already validate the proto before
constructing a TensorShape.

The same pattern in tensorflow/core/data/compression_utils.cc is being
addressed separately in #123572; this change covers the snapshot reader,
which that PR does not touch.
---
 tensorflow/core/data/snapshot_utils.cc | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/data/snapshot_utils.cc b/tensorflow/core/data/snapshot_utils.cc
index 9bca1b5e301..748faf51db6 100644
--- a/tensorflow/core/data/snapshot_utils.cc
+++ b/tensorflow/core/data/snapshot_utils.cc
@@ -966,7 +966,13 @@ absl::Status CustomReader::SnappyUncompress(
   for (int i = 0, end = simple_tensor_mask_.size(); i < end; ++i) {
     const auto& tensor_metadata = metadata->tensor_metadata(i);
     if (simple_tensor_mask_[i]) {
-      TensorShape shape(tensor_metadata.tensor_shape());
+      // The shape comes from the snapshot file. Build it through the checked
+      // API: constructing a TensorShape directly from an unvalidated proto
+      // CHECK-fails on a negative or overflowing dimension, which aborts the
+      // process instead of reporting a corrupt snapshot.
+      TensorShape shape;
+      TF_RETURN_IF_ERROR(
+          TensorShape::BuildTensorShape(tensor_metadata.tensor_shape(), &shape));
       Tensor simple_tensor(dtypes_[i], shape);
       TensorBuffer* buffer = DMAHelper::buffer(&simple_tensor);
       iov[index].iov_base = buffer->data();

From f3c26798b91e06385919290e9312f3e61a6fd0e0 Mon Sep 17 00:00:00 2001
From: alwaysprince05 <prince.p24@medhaviskillsuniversity.edu.in>
Date: Sat, 22 Aug 2026 17:35:05 +0530
Subject: [PATCH 0039/1000] Fix device placement validation for
 jit_compile=True

When using tf.device() with a non-existent device under jit_compile=True,
the device constraint was silently ignored instead of raising an error.
This adds placer validation in XlaCompiler::CompileGraph to match the
behavior of eager execution and autoclustering.

Also adds a regression test to verify all three execution modes
(eager, jit_compile=True, autoclustering) consistently validate
device constraints.

Fixes #124880
---
 tensorflow/compiler/tf2xla/BUILD              |  1 +
 tensorflow/compiler/tf2xla/xla_compiler.cc    | 16 +++++++++
 .../polymorphic_function_xla_jit_test.py      | 33 +++++++++++++++++++
 3 files changed, 50 insertions(+)

diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD
index 8f1d7610cf9..aa318f8f0e5 100644
--- a/tensorflow/compiler/tf2xla/BUILD
+++ b/tensorflow/compiler/tf2xla/BUILD
@@ -656,6 +656,7 @@ cc_library(
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
         "//tensorflow/core:ops",
+        "//tensorflow/core/common_runtime:placer",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/tpu:tpu_defs",
         "//tensorflow/core/util:overflow",
diff --git a/tensorflow/compiler/tf2xla/xla_compiler.cc b/tensorflow/compiler/tf2xla/xla_compiler.cc
index 1646fad2ceb..2c48a1ad3b4 100644
--- a/tensorflow/compiler/tf2xla/xla_compiler.cc
+++ b/tensorflow/compiler/tf2xla/xla_compiler.cc
@@ -74,6 +74,8 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/function_def_utils.h"
 #include "tensorflow/core/common_runtime/graph_constructor.h"
 #include "tensorflow/core/common_runtime/graph_optimizer.h"
+#include "tensorflow/core/common_runtime/placer.h"
+#include "tensorflow/core/common_runtime/device_set.h"
 #include "tensorflow/core/framework/attr_value_util.h"
 #include "tensorflow/core/framework/function.h"
 #include "tensorflow/core/framework/graph_debug_info.pb.h"
@@ -1548,6 +1550,20 @@ absl::Status XlaCompiler::CompileGraph(
   // Report the error here if initialization failed.
   TF_RETURN_IF_ERROR(initialization_status_);
 
+  // Run the placer to validate device constraints before XLA compilation.
+  // This ensures that tf.device() constraints are validated even for
+  // XLA-compiled functions (jit_compile=True).
+  if (device_mgr_ != nullptr) {
+    DeviceSet device_set;
+    for (Device* device : device_mgr_->ListDevices()) {
+      device_set.AddDevice(device);
+    }
+    Placer placer(graph.get(), name, options_.flib_def, &device_set,
+                  device_.get(), /*allow_soft_placement=*/true,
+                  /*log_device_placement=*/false);
+    TF_RETURN_IF_ERROR(placer.Run());
+  }
+
   // Detect invalid nodes.
   // FunctionalizeControlFlow may remove some nodes from the graph.
   TF_RETURN_IF_ERROR(ValidateGraph(graph.get(), *options_.flib_def,
diff --git a/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py b/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
index 844b2f4d4db..53149de23bf 100644
--- a/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
+++ b/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
@@ -1256,6 +1256,39 @@ class FunctionTest(xla_test.XLATestCase):
       x, y = constant_op.constant([2, 3]), constant_op.constant([2, 3])
       self._compareTwoMethodsCompilerIROutput(f, [x, y], {})
 
+  def testDevicePlacementValidationWithJitCompile(self):
+    """Test that jit_compile=True validates tf.device() constraints.
+
+    When using tf.device() with a non-existent device (e.g., GPU on CPU-only
+    machine), eager execution and autoclustering correctly fail with
+    "Could not satisfy device specification". This test ensures jit_compile=True
+    also validates device constraints instead of silently ignoring them.
+    """
+    # Use a device specification that doesn't exist on this machine.
+    # The test runs on CPU, so GPU device should fail.
+    invalid_device = '/job:worker/replica:0/task:0/device:GPU:0'
+
+    def compute(x):
+      with ops.device(invalid_device):
+        return math_ops.add(x, x)
+
+    x = constant_op.constant([1.0, 2.0])
+
+    # Eager execution should fail
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                 'Could not satisfy device specification'):
+      compute(x).numpy()
+
+    # jit_compile=True should also fail (this is the fix for #124880)
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                 'Could not satisfy device specification'):
+      polymorphic_function.function(compute, jit_compile=True)(x).numpy()
+
+    # Autoclustering (jit_compile=False with XLA auto-jit) should also fail
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                 'Could not satisfy device specification'):
+      polymorphic_function.function(compute)(x).numpy()
+
 
 if __name__ == '__main__':
   ops.enable_eager_execution()

From 25b76c881871c084101a994ebcb56115957d1028 Mon Sep 17 00:00:00 2001
From: alwaysprince05 <prince.p24@medhaviskillsuniversity.edu.in>
Date: Fri, 28 Aug 2026 17:44:50 +0530
Subject: [PATCH 0040/1000] Fix device placement validation for
 jit_compile=True

Reverts the Placer-in-CompileGraph approach which had compilation errors
(device_mgr_ is a value type, not a pointer; device_ is already a raw
pointer) and would crash at runtime due to incomplete device names.

Instead, adds device validation in the eager execution path (execute.cc)
where real devices are available. Created ValidateFunctionDeviceConstraints
utility that checks all nodes in a function definition against available
devices before the function is handed to XLA compilation.

Fixes #124880
---
 tensorflow/compiler/tf2xla/BUILD              |   1 -
 tensorflow/compiler/tf2xla/xla_compiler.cc    |  16 ---
 tensorflow/core/common_runtime/eager/BUILD    |  14 +++
 .../core/common_runtime/eager/execute.cc      |  23 ++++
 .../eager/validate_function_devices.cc        | 107 ++++++++++++++++++
 .../eager/validate_function_devices.h         |  40 +++++++
 .../polymorphic_function_xla_jit_test.py      |   9 +-
 7 files changed, 187 insertions(+), 23 deletions(-)
 create mode 100644 tensorflow/core/common_runtime/eager/validate_function_devices.cc
 create mode 100644 tensorflow/core/common_runtime/eager/validate_function_devices.h

diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD
index aa318f8f0e5..8f1d7610cf9 100644
--- a/tensorflow/compiler/tf2xla/BUILD
+++ b/tensorflow/compiler/tf2xla/BUILD
@@ -656,7 +656,6 @@ cc_library(
         "//tensorflow/core:lib",
         "//tensorflow/core:lib_internal",
         "//tensorflow/core:ops",
-        "//tensorflow/core/common_runtime:placer",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/tpu:tpu_defs",
         "//tensorflow/core/util:overflow",
diff --git a/tensorflow/compiler/tf2xla/xla_compiler.cc b/tensorflow/compiler/tf2xla/xla_compiler.cc
index 2c48a1ad3b4..1646fad2ceb 100644
--- a/tensorflow/compiler/tf2xla/xla_compiler.cc
+++ b/tensorflow/compiler/tf2xla/xla_compiler.cc
@@ -74,8 +74,6 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/function_def_utils.h"
 #include "tensorflow/core/common_runtime/graph_constructor.h"
 #include "tensorflow/core/common_runtime/graph_optimizer.h"
-#include "tensorflow/core/common_runtime/placer.h"
-#include "tensorflow/core/common_runtime/device_set.h"
 #include "tensorflow/core/framework/attr_value_util.h"
 #include "tensorflow/core/framework/function.h"
 #include "tensorflow/core/framework/graph_debug_info.pb.h"
@@ -1550,20 +1548,6 @@ absl::Status XlaCompiler::CompileGraph(
   // Report the error here if initialization failed.
   TF_RETURN_IF_ERROR(initialization_status_);
 
-  // Run the placer to validate device constraints before XLA compilation.
-  // This ensures that tf.device() constraints are validated even for
-  // XLA-compiled functions (jit_compile=True).
-  if (device_mgr_ != nullptr) {
-    DeviceSet device_set;
-    for (Device* device : device_mgr_->ListDevices()) {
-      device_set.AddDevice(device);
-    }
-    Placer placer(graph.get(), name, options_.flib_def, &device_set,
-                  device_.get(), /*allow_soft_placement=*/true,
-                  /*log_device_placement=*/false);
-    TF_RETURN_IF_ERROR(placer.Run());
-  }
-
   // Detect invalid nodes.
   // FunctionalizeControlFlow may remove some nodes from the graph.
   TF_RETURN_IF_ERROR(ValidateGraph(graph.get(), *options_.flib_def,
diff --git a/tensorflow/core/common_runtime/eager/BUILD b/tensorflow/core/common_runtime/eager/BUILD
index 8935b2fc2a4..f1ee1f4419e 100644
--- a/tensorflow/core/common_runtime/eager/BUILD
+++ b/tensorflow/core/common_runtime/eager/BUILD
@@ -665,6 +665,7 @@ cc_library(
         ":small_constants_optimizer",
         ":summary_optimizer",
         ":tensor_handle",
+        ":validate_function_devices",
         "//tensorflow/c:tf_tensor_internal",
         "//tensorflow/compiler/jit:common",
         "//tensorflow/core/profiler/lib:scoped_memory_debug_annotation",
@@ -705,6 +706,19 @@ cc_library(
     ]),
 )
 
+cc_library(
+    name = "validate_function_devices",
+    srcs = ["validate_function_devices.cc"],
+    hdrs = ["validate_function_devices.h"],
+    deps = [
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
 tf_cc_test(
     name = "execute_test",
     srcs = ["execute_test.cc"],
diff --git a/tensorflow/core/common_runtime/eager/execute.cc b/tensorflow/core/common_runtime/eager/execute.cc
index 3e8df80f1d2..f71217dc711 100644
--- a/tensorflow/core/common_runtime/eager/execute.cc
+++ b/tensorflow/core/common_runtime/eager/execute.cc
@@ -37,6 +37,7 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/eager/eager_operation.h"
 #include "tensorflow/core/common_runtime/eager/small_constants_optimizer.h"
 #include "tensorflow/core/common_runtime/eager/summary_optimizer.h"
+#include "tensorflow/core/common_runtime/eager/validate_function_devices.h"
 #include "tensorflow/core/common_runtime/int32_fulltype.h"
 #include "tensorflow/core/framework/cancellation.h"
 #include "tensorflow/core/framework/full_type.pb.h"
@@ -1405,6 +1406,28 @@ absl::Status GetOrCreateKernelAndDevice(
       TF_RETURN_IF_ERROR(
           UpdateCompileCounter(op, ctx, compile_with_xla, has_tpu_replication));
       if (compile_with_xla && !has_tpu_replication) {
+        // Validate device constraints when jit_compile=True was explicitly
+        // set by the user. When jit_compile=True, the function is compiled
+        // by XLA as a whole, bypassing the Placer which normally validates
+        // device constraints. We validate here to ensure tf.device()
+        // constraints are enforced.
+        //
+        // Only validate when _XlaMustCompile was explicitly set (user's
+        // jit_compile=True), not when compile_with_xla is inferred from
+        // the function being on an XLA device (TPU/XLA_GPU/XLA_CPU).
+        bool has_xla_must_compile = false;
+        GetFuncAttr(op, ctx, kXlaMustCompileAttr, &has_xla_must_compile)
+            .IgnoreError();
+        if (has_xla_must_compile) {
+          const FunctionLibraryDefinition* func_lib = ctx.FuncLibDef();
+          const FunctionDef* fdef = func_lib->Find(op->Name());
+          if (fdef != nullptr) {
+            std::vector<DeviceAttributes> device_attrs;
+            ctx.ListDevices(&device_attrs);
+            TF_RETURN_IF_ERROR(ValidateFunctionDeviceConstraints(
+                *fdef, device_attrs));
+          }
+        }
         if (ctx.JitCompileRewrite()) {
           xla_compile_device_type = op->GetDeviceParsedName().type;
         } else {
diff --git a/tensorflow/core/common_runtime/eager/validate_function_devices.cc b/tensorflow/core/common_runtime/eager/validate_function_devices.cc
new file mode 100644
index 00000000000..441f63133fc
--- /dev/null
+++ b/tensorflow/core/common_runtime/eager/validate_function_devices.cc
@@ -0,0 +1,107 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/core/common_runtime/eager/validate_function_devices.h"
+
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/node_def_util.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/util/device_name_utils.h"
+
+namespace tensorflow {
+
+absl::Status ValidateFunctionDeviceConstraints(
+    const FunctionDef& fdef,
+    const std::vector<DeviceAttributes>& available_devices) {
+  // Build a set of available device types and full device names.
+  std::vector<DeviceType> supported_device_types;
+  std::vector<std::string> available_device_names;
+  for (const auto& dev : available_devices) {
+    DeviceNameUtils::ParsedName parsed;
+    if (DeviceNameUtils::ParseFullName(dev.name(), &parsed)) {
+      supported_device_types.push_back(DeviceType(parsed.type));
+    }
+    available_device_names.push_back(dev.name());
+  }
+
+  for (const NodeDef& node : fdef.node_def()) {
+    const std::string& device = node.device();
+    if (device.empty()) {
+      continue;  // No device constraint on this node.
+    }
+
+    DeviceNameUtils::ParsedName parsed_device;
+    if (!DeviceNameUtils::ParseFullName(device, &parsed_device)) {
+      // If the device name can't be parsed, skip it. Other validation passes
+      // will catch malformed device names.
+      continue;
+    }
+
+    // Check if the requested device type is available.
+    bool device_type_available = false;
+    for (const DeviceType& avail_type : supported_device_types) {
+      if (absl::StrEqualsIgnoreCase(parsed_device.type, avail_type.type_string())) {
+        device_type_available = true;
+        break;
+      }
+    }
+
+    if (!device_type_available) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Could not satisfy device specification '", device,
+          "' for operation ", node.name(), " (", node.op(),
+          "). Available devices [",
+          absl::StrJoin(available_device_names, ", "), "]."));
+    }
+
+    // For fully-specified device names, check if an exact match exists.
+    if (parsed_device.has_job && parsed_device.has_replica &&
+        parsed_device.has_task && parsed_device.has_type &&
+        parsed_device.has_id) {
+      bool exact_match = false;
+      for (const auto& dev : available_devices) {
+        DeviceNameUtils::ParsedName avail_parsed;
+        if (DeviceNameUtils::ParseFullName(dev.name(), &avail_parsed)) {
+          if (DeviceNameUtils::IsCompleteSpecification(parsed_device,
+                                                       avail_parsed)) {
+            exact_match = true;
+            break;
+          }
+        }
+      }
+
+      if (!exact_match) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Could not satisfy device specification '", device,
+            "' for operation ", node.name(), " (", node.op(),
+            "). Available devices [",
+            absl::StrJoin(available_device_names, ", "), "]."));
+      }
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace tensorflow
diff --git a/tensorflow/core/common_runtime/eager/validate_function_devices.h b/tensorflow/core/common_runtime/eager/validate_function_devices.h
new file mode 100644
index 00000000000..e75d30c7c3e
--- /dev/null
+++ b/tensorflow/core/common_runtime/eager/validate_function_devices.h
@@ -0,0 +1,40 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_COMMON_RUNTIME_EAGER_VALIDATE_FUNCTION_DEVICES_H_
+#define TENSORFLOW_CORE_COMMON_RUNTIME_EAGER_VALIDATE_FUNCTION_DEVICES_H_
+
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/device_attributes.pb.h"
+
+namespace tensorflow {
+
+// Validates that device constraints specified on nodes within a FunctionDef
+// can be satisfied by the available devices. This is used to ensure that
+// tf.device() constraints are validated even for XLA-compiled functions
+// (jit_compile=True), where the Placer is bypassed.
+//
+// Returns OK if all device constraints are satisfiable, or an error status
+// describing the first unsatisfiable device constraint.
+absl::Status ValidateFunctionDeviceConstraints(
+    const FunctionDef& fdef,
+    const std::vector<DeviceAttributes>& available_devices);
+
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_COMMON_RUNTIME_EAGER_VALIDATE_FUNCTION_DEVICES_H_
diff --git a/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py b/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
index 53149de23bf..0e121b0c07b 100644
--- a/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
+++ b/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
@@ -1260,9 +1260,11 @@ class FunctionTest(xla_test.XLATestCase):
     """Test that jit_compile=True validates tf.device() constraints.
 
     When using tf.device() with a non-existent device (e.g., GPU on CPU-only
-    machine), eager execution and autoclustering correctly fail with
+    machine), eager execution correctly fails with
     "Could not satisfy device specification". This test ensures jit_compile=True
     also validates device constraints instead of silently ignoring them.
+
+    See https://github.com/tensorflow/tensorflow/issues/124880
     """
     # Use a device specification that doesn't exist on this machine.
     # The test runs on CPU, so GPU device should fail.
@@ -1284,11 +1286,6 @@ class FunctionTest(xla_test.XLATestCase):
                                  'Could not satisfy device specification'):
       polymorphic_function.function(compute, jit_compile=True)(x).numpy()
 
-    # Autoclustering (jit_compile=False with XLA auto-jit) should also fail
-    with self.assertRaisesRegex(errors.InvalidArgumentError,
-                                 'Could not satisfy device specification'):
-      polymorphic_function.function(compute)(x).numpy()
-
 
 if __name__ == '__main__':
   ops.enable_eager_execution()

From 26e33de30f4e3294dd07c40f745e7168f6a46be3 Mon Sep 17 00:00:00 2001
From: alwaysprince05 <prince.p24@medhaviskillsuniversity.edu.in>
Date: Sat, 29 Aug 2026 01:09:31 +0530
Subject: [PATCH 0041/1000] Fix absl function name: StrEqualsIgnoreCase to
 EqualsIgnoreCase

---
 .../core/common_runtime/eager/validate_function_devices.cc      | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/common_runtime/eager/validate_function_devices.cc b/tensorflow/core/common_runtime/eager/validate_function_devices.cc
index 441f63133fc..1fcf96d4b1d 100644
--- a/tensorflow/core/common_runtime/eager/validate_function_devices.cc
+++ b/tensorflow/core/common_runtime/eager/validate_function_devices.cc
@@ -61,7 +61,7 @@ absl::Status ValidateFunctionDeviceConstraints(
     // Check if the requested device type is available.
     bool device_type_available = false;
     for (const DeviceType& avail_type : supported_device_types) {
-      if (absl::StrEqualsIgnoreCase(parsed_device.type, avail_type.type_string())) {
+      if (absl::EqualsIgnoreCase(parsed_device.type, avail_type.type_string())) {
         device_type_available = true;
         break;
       }

From 6136b12b7e3cb0cb04d1f3df59d2f581d1adcd04 Mon Sep 17 00:00:00 2001
From: alwaysprince05 <prince.p24@medhaviskillsuniversity.edu.in>
Date: Sun, 30 Aug 2026 00:28:59 +0530
Subject: [PATCH 0042/1000] Add validate_function_devices to mkl_core deps

---
 tensorflow/core/common_runtime/eager/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/core/common_runtime/eager/BUILD b/tensorflow/core/common_runtime/eager/BUILD
index f1ee1f4419e..2879718a356 100644
--- a/tensorflow/core/common_runtime/eager/BUILD
+++ b/tensorflow/core/common_runtime/eager/BUILD
@@ -847,6 +847,7 @@ cc_library(
         ":small_constants_optimizer",
         ":summary_optimizer",
         ":tensor_handle",
+        ":validate_function_devices",
         "//tensorflow/c:c_api_internal",
         "//tensorflow/c:tf_tensor_internal",
         "//tensorflow/c/eager:abstract_function",

From b348f86f63cca95b21cf5adb81c274284f65f07f Mon Sep 17 00:00:00 2001
From: alwaysprince05 <prince.p24@medhaviskillsuniversity.edu.in>
Date: Sun, 30 Aug 2026 01:55:59 +0530
Subject: [PATCH 0043/1000] trigger CI re-run


From d0064ef68ec4b8c97aec4bafe5bbfaa2935e7ace Mon Sep 17 00:00:00 2001
From: endorphin13 <elshad.toklayev.business@gmail.com>
Date: Sat, 29 Aug 2026 11:29:42 +0300
Subject: [PATCH 0044/1000] Add a regression test for an invalid shape in
 snapshot metadata Writes a valid snappy snapshot, then rewrites its metadata
 record with a negative dimension and asserts that reading it returns an
 error. The record framing is an 8-byte little-endian length followed by the
 payload, and for a snappy snapshot the metadata is the first record, so it
 can be patched without touching the compressed data that follows. int64
 tensors are used so `simple_tensor_mask_` is set and the read takes the path
 in `SnappyUncompress` that builds the shape. Without this change the test
 aborts the process on the CHECK inside the TensorShape proto constructor
 rather than failing.

---
 tensorflow/core/data/snapshot_utils_test.cc | 64 +++++++++++++++++++++
 1 file changed, 64 insertions(+)

diff --git a/tensorflow/core/data/snapshot_utils_test.cc b/tensorflow/core/data/snapshot_utils_test.cc
index 47b1719aaef..8df36d155af 100644
--- a/tensorflow/core/data/snapshot_utils_test.cc
+++ b/tensorflow/core/data/snapshot_utils_test.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "tensorflow/core/data/snapshot_utils.h"
 
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <vector>
@@ -24,6 +25,7 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor.pb.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
 #include "tensorflow/core/lib/io/compression.h"
+#include "tensorflow/core/platform/coding.h"
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/test.h"
@@ -150,6 +152,68 @@ TEST(SnapshotUtilTest, SnappyTensorSizeMismatch) {
   TF_ASSERT_OK(Env::Default()->DeleteFile(filename));
 }
 
+TEST(SnapshotUtilTest, SnappyInvalidTensorShapeIsAnError) {
+  // A corrupt snapshot whose metadata carries a negative dimension must be
+  // reported as an error. Building a TensorShape directly from the proto
+  // CHECK-fails instead, which aborts the process.
+  //
+  // Uses a memcpy-able dtype so the tensors take the simple-tensor path in
+  // CustomReader::SnappyUncompress, which is where the shape is built.
+  std::vector<Tensor> tensors;
+  tensorflow::DataTypeVector dtypes;
+  for (int i = 0; i < 2; ++i) {
+    Tensor t(DT_INT64, TensorShape({4}));
+    t.flat<int64_t>().setZero();
+    dtypes.push_back(DT_INT64);
+    tensors.push_back(t);
+  }
+
+  std::string filename;
+  EXPECT_TRUE(Env::Default()->LocalTempFilename(&filename));
+
+  std::unique_ptr<Writer> writer;
+  TF_ASSERT_OK(Writer::Create(Env::Default(), filename, io::compression::kSnappy,
+                              /*version=*/1, dtypes, &writer));
+  TF_ASSERT_OK(writer->WriteTensors(tensors));
+  TF_ASSERT_OK(writer->Close());
+
+  // Records are an 8-byte little-endian length followed by the payload; the
+  // metadata is the first record written for a snappy snapshot. Rewrite it
+  // with a negative dimension.
+  constexpr size_t kHeaderSize = sizeof(uint64_t);
+  std::string contents;
+  TF_ASSERT_OK(ReadFileToString(Env::Default(), filename, &contents));
+  ASSERT_GT(contents.size(), kHeaderSize);
+  const uint64_t metadata_size = core::DecodeFixed64(contents.data());
+  ASSERT_GE(contents.size(), kHeaderSize + metadata_size);
+
+  experimental::SnapshotTensorMetadata metadata;
+  ASSERT_TRUE(
+      metadata.ParseFromString(contents.substr(kHeaderSize, metadata_size)));
+  ASSERT_GT(metadata.tensor_metadata_size(), 0);
+  ASSERT_GT(metadata.tensor_metadata(0).tensor_shape().dim_size(), 0);
+  metadata.mutable_tensor_metadata(0)
+      ->mutable_tensor_shape()
+      ->mutable_dim(0)
+      ->set_size(-1);
+
+  const std::string patched = metadata.SerializeAsString();
+  char header[kHeaderSize];
+  core::EncodeFixed64(header, patched.size());
+  TF_ASSERT_OK(WriteStringToFile(
+      Env::Default(), filename,
+      std::string(header, kHeaderSize) + patched +
+          contents.substr(kHeaderSize + metadata_size)));
+
+  std::unique_ptr<Reader> reader;
+  TF_ASSERT_OK(Reader::Create(Env::Default(), filename, io::compression::kSnappy,
+                              /*version=*/1, dtypes, &reader));
+  std::vector<Tensor> read_tensors;
+  EXPECT_FALSE(reader->ReadTensors(&read_tensors).ok());
+
+  TF_ASSERT_OK(Env::Default()->DeleteFile(filename));
+}
+
 void SnapshotReaderBenchmarkLoop(::testing::benchmark::State& state,
                                  std::string compression_type, int version) {
   tensorflow::DataTypeVector dtypes;

From 7f55b61bb74cb21fca9f2c135c295fdb00fb46a8 Mon Sep 17 00:00:00 2001
From: PrabinDevkota <prabindev789@gmail.com>
Date: Mon, 31 Aug 2026 00:29:08 +0530
Subject: [PATCH 0045/1000] Reject ExpandDims, Pack, and StridedSlice when
 output rank exceeds TensorShape's limit.

---
 tensorflow/compiler/tf2xla/kernels/pack_op.cc |  9 +-
 .../compiler/tf2xla/kernels/shape_op.cc       |  7 ++
 tensorflow/core/kernels/BUILD                 | 22 +++++
 tensorflow/core/kernels/pack_op.cc            | 10 ++-
 tensorflow/core/kernels/pack_op_test.cc       | 86 +++++++++++++++++++
 tensorflow/core/kernels/shape_ops.h           | 11 ++-
 tensorflow/core/kernels/shape_ops_test.cc     | 57 ++++++++++++
 tensorflow/core/ops/array_ops.cc              | 15 ++++
 tensorflow/core/util/strided_slice_op.cc      | 25 ++++--
 tensorflow/core/util/strided_slice_op_test.cc | 62 +++++++++++++
 .../kernel_tests/array_ops/array_ops_test.py  | 46 ++++++++++
 .../kernel_tests/array_ops/stack_op_test.py   | 17 ++++
 12 files changed, 359 insertions(+), 8 deletions(-)
 create mode 100644 tensorflow/core/kernels/pack_op_test.cc

diff --git a/tensorflow/compiler/tf2xla/kernels/pack_op.cc b/tensorflow/compiler/tf2xla/kernels/pack_op.cc
index ba4e8bbef7b..3b1e83e5c4e 100644
--- a/tensorflow/compiler/tf2xla/kernels/pack_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/pack_op.cc
@@ -61,10 +61,17 @@ class PackOp : public XlaOpKernel {
                                         -expanded_num_dims, ", ",
                                         expanded_num_dims, ")"));
 
+    OP_REQUIRES(
+        ctx, shapes[0].dims() < TensorShape::MaxDimensions(),
+        errors::InvalidArgument(
+            "Cannot pack tensors of rank ", shapes[0].dims(),
+            ": packing would exceed the maximum supported rank of ",
+            TensorShape::MaxDimensions(), "."));
+
     std::vector<xla::XlaOp> reshaped_inputs(num);
 
     TensorShape child_shape(shapes[0]);
-    child_shape.InsertDim(axis, 1);
+    OP_REQUIRES_OK(ctx, child_shape.InsertDimWithStatus(axis, 1));
 
     for (int i = 0; i < num; ++i) {
       // Reshape the inputs to have an extra dimension of size 1.
diff --git a/tensorflow/compiler/tf2xla/kernels/shape_op.cc b/tensorflow/compiler/tf2xla/kernels/shape_op.cc
index c37bda13193..253538082db 100644
--- a/tensorflow/compiler/tf2xla/kernels/shape_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/shape_op.cc
@@ -289,6 +289,13 @@ class ExpandDimsOp : public XlaOpKernel {
                                         " for tensor with ", input_shape.dims(),
                                         " dimensions."));
 
+    OP_REQUIRES(
+        ctx, input_shape.dims() < TensorShape::MaxDimensions(),
+        errors::InvalidArgument(
+            "Cannot expand a tensor of rank ", input_shape.dims(),
+            ": expanding would exceed the maximum supported rank of ",
+            TensorShape::MaxDimensions(), "."));
+
     auto existing_dims = input_shape.dim_sizes();
     // Safe - # elements in tensor dims bounded.
     const int existing_dims_size = static_cast<int>(existing_dims.size());
diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 5bc3c390c96..5a35dff9186 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -1136,6 +1136,26 @@ tf_kernel_library(
     deps = ARRAY_DEPS,
 )
 
+tf_cc_test(
+    name = "pack_op_test",
+    size = "small",
+    srcs = ["pack_op_test.cc"],
+    deps = [
+        ":ops_testutil",
+        ":ops_util",
+        ":pack_op",
+        "//tensorflow/core:core_cpu",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+        "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest",
+    ],
+)
+
 tf_kernel_library(
     name = "pad_op",
     prefix = "pad_op",
@@ -2390,6 +2410,8 @@ tf_cc_test(
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
         "//tensorflow/core:testlib",
+        "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest",
     ],
 )
 
diff --git a/tensorflow/core/kernels/pack_op.cc b/tensorflow/core/kernels/pack_op.cc
index 95331ce42dc..b5248401d21 100644
--- a/tensorflow/core/kernels/pack_op.cc
+++ b/tensorflow/core/kernels/pack_op.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/tensor_types.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/kernels/concat_lib.h"
@@ -65,8 +66,15 @@ class PackOp : public OpKernel {
                     "axis = ", axis_, " not in [", -expanded_num_dims, ", ",
                     expanded_num_dims, ")")));
 
+    OP_REQUIRES(
+        c, first_input.dims() < TensorShape::MaxDimensions(),
+        absl::InvalidArgumentError(absl::StrCat(
+            "Cannot pack tensors of rank ", first_input.dims(),
+            ": packing would exceed the maximum supported rank of ",
+            TensorShape::MaxDimensions(), ".")));
+
     TensorShape output_shape(first_input.shape());
-    output_shape.InsertDim(axis, num);
+    OP_REQUIRES_OK(c, output_shape.InsertDimWithStatus(axis, num));
 
     // In the num = 1 case, just reshape the input
     if (num == 1) {
diff --git a/tensorflow/core/kernels/pack_op_test.cc b/tensorflow/core/kernels/pack_op_test.cc
new file mode 100644
index 00000000000..30008065875
--- /dev/null
+++ b/tensorflow/core/kernels/pack_op_test.cc
@@ -0,0 +1,86 @@
+/* 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.
+==============================================================================*/
+
+#include <cstdint>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "tensorflow/core/framework/fake_input.h"
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/kernels/ops_testutil.h"
+#include "tensorflow/core/kernels/ops_util.h"
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/status_matchers.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace {
+
+using ::tensorflow::testing::StatusIs;
+using ::testing::HasSubstr;
+
+absl::Status MakeZeroElementShape(int rank, TensorShape* shape) {
+  std::vector<int64_t> dims(rank, 1);
+  if (rank > 0) {
+    dims[0] = 0;
+  }
+  return TensorShape::BuildTensorShape(dims, shape);
+}
+
+class PackOpTest : public OpsTestBase {
+ protected:
+  void MakeOp() {
+    TF_ASSERT_OK(NodeDefBuilder("pack", "Pack")
+                     .Input(FakeInput(1, DT_FLOAT))
+                     .Attr("N", 1)
+                     .Attr("T", DT_FLOAT)
+                     .Attr("axis", 0)
+                     .Finalize(node_def()));
+    TF_ASSERT_OK(InitOp());
+  }
+};
+
+TEST_F(PackOpTest, PackingMaxRankSucceeds) {
+  MakeOp();
+  TensorShape input_shape;
+  TF_ASSERT_OK(
+      MakeZeroElementShape(TensorShape::MaxDimensions() - 1, &input_shape));
+  AddInput<float>(input_shape, [](int) { return 0.0f; });
+
+  TF_ASSERT_OK(RunOpKernel());
+  EXPECT_EQ(GetOutput(0)->dims(), TensorShape::MaxDimensions());
+  EXPECT_EQ(GetOutput(0)->dim_size(0), 1);
+}
+
+TEST_F(PackOpTest, PackingBeyondMaxRankFails) {
+  MakeOp();
+  TensorShape input_shape;
+  TF_ASSERT_OK(MakeZeroElementShape(TensorShape::MaxDimensions(), &input_shape));
+  AddInput<float>(input_shape, [](int) { return 0.0f; });
+
+  EXPECT_THAT(RunOpKernel(),
+              StatusIs(error::INVALID_ARGUMENT,
+                       HasSubstr("maximum supported rank")));
+}
+
+}  // namespace
+}  // namespace tensorflow
diff --git a/tensorflow/core/kernels/shape_ops.h b/tensorflow/core/kernels/shape_ops.h
index 7b9c940542a..f4a07570257 100644
--- a/tensorflow/core/kernels/shape_ops.h
+++ b/tensorflow/core/kernels/shape_ops.h
@@ -168,6 +168,13 @@ class ExpandDimsOp : public OpKernel {
       dim = std::min<Tdim>(dim + input_dims + 1, input_dims);
     }
 
+    OP_REQUIRES(
+        ctx, input_dims < TensorShape::MaxDimensions(),
+        absl::InvalidArgumentError(absl::StrCat(
+            "Cannot expand a tensor of rank ", input_dims,
+            ": expanding would exceed the maximum supported rank of ",
+            TensorShape::MaxDimensions(), ".")));
+
     // Compute new shape with an additional dimension.
     absl::InlinedVector<int64_t, 8> output_shape_vec(input_dims + 1);
     for (int64_t i = 0; i < dim; ++i) {
@@ -177,7 +184,9 @@ class ExpandDimsOp : public OpKernel {
     for (int64_t i = dim + 1; i < input_dims + 1; ++i) {
       output_shape_vec[i] = input_shape.dim_size(i - 1);
     }
-    TensorShape output_shape(output_shape_vec);
+    TensorShape output_shape;
+    OP_REQUIRES_OK(
+        ctx, TensorShape::BuildTensorShape(output_shape_vec, &output_shape));
 
     Tensor output_t;
     if (!output_t.CopyFrom(input_t, output_shape)) {
diff --git a/tensorflow/core/kernels/shape_ops_test.cc b/tensorflow/core/kernels/shape_ops_test.cc
index 1cee9c25366..2e639e088a6 100644
--- a/tensorflow/core/kernels/shape_ops_test.cc
+++ b/tensorflow/core/kernels/shape_ops_test.cc
@@ -13,13 +13,21 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <cstdint>
 #include <functional>
 #include <memory>
+#include <vector>
 
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
 #include "tensorflow/core/common_runtime/kernel_benchmark_testlib.h"
 #include "tensorflow/core/framework/allocator.h"
+#include "tensorflow/core/framework/fake_input.h"
+#include "tensorflow/core/framework/node_def_builder.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/graph/algorithm.h"
@@ -27,13 +35,62 @@ limitations under the License.
 #include "tensorflow/core/graph/testlib.h"
 #include "tensorflow/core/kernels/ops_testutil.h"
 #include "tensorflow/core/kernels/ops_util.h"
+#include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/status_matchers.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/platform/test_benchmark.h"
 
 namespace tensorflow {
 namespace {
 
+using ::tensorflow::testing::StatusIs;
+using ::testing::HasSubstr;
+
+absl::Status MakeZeroElementShape(int rank, TensorShape* shape) {
+  std::vector<int64_t> dims(rank, 1);
+  if (rank > 0) {
+    dims[0] = 0;
+  }
+  return TensorShape::BuildTensorShape(dims, shape);
+}
+
+class ExpandDimsOpTest : public OpsTestBase {
+ protected:
+  void MakeOp() {
+    TF_ASSERT_OK(NodeDefBuilder("expand", "ExpandDims")
+                     .Input(FakeInput(DT_FLOAT))
+                     .Input(FakeInput(DT_INT32))
+                     .Finalize(node_def()));
+    TF_ASSERT_OK(InitOp());
+  }
+};
+
+TEST_F(ExpandDimsOpTest, ExpandingMaxRankSucceeds) {
+  MakeOp();
+  TensorShape input_shape;
+  TF_ASSERT_OK(
+      MakeZeroElementShape(TensorShape::MaxDimensions() - 1, &input_shape));
+  AddInput<float>(input_shape, [](int) { return 0.0f; });
+  AddInputFromList<int32_t>(TensorShape({}), {0});
+
+  TF_ASSERT_OK(RunOpKernel());
+  EXPECT_EQ(GetOutput(0)->dims(), TensorShape::MaxDimensions());
+  EXPECT_EQ(GetOutput(0)->dim_size(0), 1);
+}
+
+TEST_F(ExpandDimsOpTest, ExpandingBeyondMaxRankFails) {
+  MakeOp();
+  TensorShape input_shape;
+  TF_ASSERT_OK(MakeZeroElementShape(TensorShape::MaxDimensions(), &input_shape));
+  AddInput<float>(input_shape, [](int) { return 0.0f; });
+  AddInputFromList<int32_t>(TensorShape({}), {0});
+
+  EXPECT_THAT(RunOpKernel(),
+              StatusIs(error::INVALID_ARGUMENT,
+                       HasSubstr("maximum supported rank")));
+}
+
 static void BM_ExpandDims(::testing::benchmark::State& state) {
   Graph* g = new Graph(OpRegistry::Global());
 
diff --git a/tensorflow/core/ops/array_ops.cc b/tensorflow/core/ops/array_ops.cc
index a97ed5d9bd0..e9a772a395a 100644
--- a/tensorflow/core/ops/array_ops.cc
+++ b/tensorflow/core/ops/array_ops.cc
@@ -24,6 +24,7 @@ limitations under the License.
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/shape_inference.h"
 #include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/lib/core/errors.h"
@@ -377,6 +378,13 @@ REGISTER_OP("Pack")
       int32_t axis;
       TF_RETURN_IF_ERROR(GetAxisForPackAndUnpack(c, rank + 1, &axis));
 
+      if (rank >= TensorShape::MaxDimensions()) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Cannot pack tensors of rank ", rank,
+            ": packing would exceed the maximum supported rank of ",
+            TensorShape::MaxDimensions(), "."));
+      }
+
       // Copy all dimensions over, inserting a dimension of value #inputs
       // at <axis>.
       std::vector<DimensionHandle> dims;
@@ -2150,6 +2158,13 @@ REGISTER_OP("ExpandDims")
         dim += rank + 1;
       }
 
+      if (rank >= TensorShape::MaxDimensions()) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Cannot expand a tensor of rank ", rank,
+            ": expanding would exceed the maximum supported rank of ",
+            TensorShape::MaxDimensions(), "."));
+      }
+
       ShapeHandle end;
       TF_RETURN_IF_ERROR(c->Subshape(input, dim, &end));
 
diff --git a/tensorflow/core/util/strided_slice_op.cc b/tensorflow/core/util/strided_slice_op.cc
index eb02f5f6e3f..d9981276bdd 100644
--- a/tensorflow/core/util/strided_slice_op.cc
+++ b/tensorflow/core/util/strided_slice_op.cc
@@ -309,7 +309,8 @@ absl::Status ValidateStridedSliceOp(
     }
     bool shrink_i = (dense_spec.shrink_axis_mask & (1 << i));
     if (dim_i == -1) {
-      processing_shape->AddDim(shrink_i ? 1 : -1);
+      TF_RETURN_IF_ERROR(
+          processing_shape->AddDimWithStatus(shrink_i ? 1 : -1));
       continue;
     }
 
@@ -397,9 +398,9 @@ absl::Status ValidateStridedSliceOp(
         size_i = interval_length / stride_i +
                  (interval_length % stride_i != 0 ? 1 : 0);
       }
-      processing_shape->AddDim(size_i);
+      TF_RETURN_IF_ERROR(processing_shape->AddDimWithStatus(size_i));
     } else {
-      processing_shape->AddDim(-1);
+      TF_RETURN_IF_ERROR(processing_shape->AddDimWithStatus(-1));
     }
   }
 
@@ -408,6 +409,19 @@ absl::Status ValidateStridedSliceOp(
   // new_axis will increase dimension by 1 (with a one-size dimension)
   // slices like foo[3,...] will reduce dimension by 1.
   // This cannot be done earlier, because it depends on Step 3.
+  int output_rank = 0;
+  for (int32_t gather_index : dense_spec.final_shape_gather_indices) {
+    if (gather_index != kShrinkAxis) {
+      ++output_rank;
+    }
+  }
+  if (output_rank > TensorShape::MaxDimensions()) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "StridedSlice would produce a tensor of rank ", output_rank,
+        ", which exceeds the maximum supported rank of ",
+        TensorShape::MaxDimensions(), "."));
+  }
+
   final_shape->Clear();
   if (shape_spec != nullptr) {
     shape_spec->output_to_sparse_mapping.clear();
@@ -427,13 +441,14 @@ absl::Status ValidateStridedSliceOp(
     int64_t sparse_index =
         dense_spec.final_shape_gather_indices_sparse[dense_dim];
     if (gather_index >= 0) {
-      final_shape->AddDim(processing_shape->dim_size(gather_index));
+      TF_RETURN_IF_ERROR(final_shape->AddDimWithStatus(
+          processing_shape->dim_size(gather_index)));
       if (shape_spec != nullptr) {
         shape_spec->output_to_sparse_mapping.push_back(sparse_index);
         shape_spec->output_to_processing_mapping.push_back(gather_index);
       }
     } else if (gather_index == kNewAxis) {
-      final_shape->AddDim(1);
+      TF_RETURN_IF_ERROR(final_shape->AddDimWithStatus(1));
       if (shape_spec != nullptr) {
         shape_spec->output_to_sparse_mapping.push_back(-1);
         shape_spec->output_to_processing_mapping.push_back(-1);
diff --git a/tensorflow/core/util/strided_slice_op_test.cc b/tensorflow/core/util/strided_slice_op_test.cc
index ee96537d009..cd666e7b275 100644
--- a/tensorflow/core/util/strided_slice_op_test.cc
+++ b/tensorflow/core/util/strided_slice_op_test.cc
@@ -19,9 +19,11 @@ limitations under the License.
 #include <ostream>
 #include <tuple>
 #include <utility>
+#include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/status/status.h"
 #include "xla/tsl/lib/core/status_test_util.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
@@ -356,6 +358,30 @@ TensorShape AsTensorShape(absl::Span<const int64_t> dim_sizes) {
   return out;
 }
 
+TensorShape ZeroElementRank(int rank) {
+  std::vector<int64_t> dims(rank, 1);
+  if (rank > 0) {
+    dims[0] = 0;
+  }
+  return AsTensorShape(dims);
+}
+
+absl::Status ValidateNewAxis(const TensorShape& input_shape,
+                             TensorShape* processing_shape,
+                             TensorShape* final_shape) {
+  Tensor begin_tensor = test::AsTensor<int32_t>({0});
+  Tensor end_tensor = test::AsTensor<int32_t>({0});
+  Tensor strides_tensor = test::AsTensor<int32_t>({1});
+  bool is_identity, is_simple_slice, slice_dim0;
+  IntVector begin, end, strides;
+  return ValidateStridedSliceOp(
+      &begin_tensor, &end_tensor, strides_tensor, input_shape,
+      /*begin_mask_spec=*/0, /*end_mask_spec=*/0, /*ellipsis_mask=*/0,
+      /*new_axis_mask=*/1, /*shrink_axis_mask=*/0, processing_shape,
+      final_shape, &is_identity, &is_simple_slice, &slice_dim0, &begin, &end,
+      &strides);
+}
+
 TEST(ValidateStridedSliceOpTest, BasicStride) {
   Tensor begin_tensor = test::AsTensor<int32_t>({1, 1});
   Tensor end_tensor = test::AsTensor<int32_t>({7, 7});
@@ -771,5 +797,41 @@ TEST(ValidateStridedSliceOpTest, PartialInputShape) {
       &is_simple_slice, &slice_dim0, &begin, &end, &strides, &shape_spec));
 }
 
+TEST(ValidateStridedSliceOpTest, NewAxisAtMaxRankSucceeds) {
+  TensorShape processing_shape, final_shape;
+  TF_EXPECT_OK(ValidateNewAxis(ZeroElementRank(TensorShape::MaxDimensions() - 1),
+                               &processing_shape, &final_shape));
+  EXPECT_EQ(final_shape.dims(), TensorShape::MaxDimensions());
+  EXPECT_EQ(final_shape.dim_size(0), 1);
+}
+
+TEST(ValidateStridedSliceOpTest, NewAxisBeyondMaxRankFails) {
+  TensorShape processing_shape, final_shape;
+  EXPECT_THAT(
+      ValidateNewAxis(ZeroElementRank(TensorShape::MaxDimensions()),
+                      &processing_shape, &final_shape),
+      absl_testing::StatusIs(tsl::error::Code::INVALID_ARGUMENT,
+                             ::testing::HasSubstr("maximum supported rank")));
+}
+
+TEST(ValidateStridedSliceOpTest, ManyNewAxesBeyondMaxRankFails) {
+  // Mirrors the public StridedSlice repro: rank 225 plus 30 new axes.
+  Tensor begin_tensor = test::AsTensor<int32_t>(std::vector<int32_t>(31, 0));
+  Tensor end_tensor = test::AsTensor<int32_t>(std::vector<int32_t>(31, 0));
+  Tensor strides_tensor = test::AsTensor<int32_t>(std::vector<int32_t>(31, 1));
+  TensorShape processing_shape, final_shape;
+  bool is_identity, is_simple_slice, slice_dim0;
+  IntVector begin, end, strides;
+  EXPECT_THAT(
+      ValidateStridedSliceOp(
+          &begin_tensor, &end_tensor, strides_tensor, ZeroElementRank(225),
+          /*begin_mask_spec=*/0, /*end_mask_spec=*/0,
+          /*ellipsis_mask=*/(1 << 30), /*new_axis_mask=*/(1 << 30) - 1,
+          /*shrink_axis_mask=*/0, &processing_shape, &final_shape,
+          &is_identity, &is_simple_slice, &slice_dim0, &begin, &end, &strides),
+      absl_testing::StatusIs(tsl::error::Code::INVALID_ARGUMENT,
+                             ::testing::HasSubstr("maximum supported rank")));
+}
+
 }  // namespace
 }  // namespace tensorflow
diff --git a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
index 566908b69f2..580749a93d7 100644
--- a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
@@ -392,6 +392,13 @@ class BooleanMaskTest(test_util.TensorFlowTestCase):
 @test_util.run_all_in_graph_and_eager_modes
 class OperatorShapeTest(test_util.TensorFlowTestCase):
 
+  def _zero_high_rank(self, rank):
+    # tf.zeros(high_rank) in graph mode calls np.empty, and NumPy rejects
+    # rank > 32. Fill takes a 1-D shape tensor instead.
+    return array_ops.fill(
+        [0] + [1] * (rank - 1),
+        constant_op.constant(0.0, dtype=dtypes.float32))
+
   def testExpandScalar(self):
     scalar = "hello"
     scalar_expanded = array_ops.expand_dims(scalar, [0])
@@ -430,6 +437,18 @@ class OperatorShapeTest(test_util.TensorFlowTestCase):
               shape=constant_op.constant([1 for i in range(254)],
                                          dtype=dtypes.int64)))
 
+  def testExpandDimsAtMaxRank(self):
+    # TensorShape::MaxDimensions() is 254.
+    value = self._zero_high_rank(253)
+    expanded = array_ops.expand_dims(value, 0)
+    self.assertEqual(254, self.evaluate(array_ops.rank(expanded)))
+
+  def testExpandDimsExceedsMaxRank(self):
+    value = self._zero_high_rank(254)
+    with self.assertRaisesRegex(
+        (errors.InvalidArgumentError, ValueError), "maximum supported rank"):
+      self.evaluate(array_ops.expand_dims(value, 0))
+
 
 @test_util.with_eager_op_as_function
 class ReverseV2Test(test_util.TensorFlowTestCase):
@@ -757,6 +776,33 @@ class StridedSliceTest(test_util.TensorFlowTestCase):
       s = array_ops.strided_slice(x, begin, end, strides)
       self.assertAllEqual([3.], self.evaluate(s))
 
+  def testNewAxisAtMaxRank(self):
+    value = array_ops.fill(
+        [0] + [1] * 252,
+        constant_op.constant(0.0, dtype=dtypes.float32))
+    sliced = gen_array_ops.strided_slice(
+        value,
+        begin=[0],
+        end=[0],
+        strides=[1],
+        new_axis_mask=1)
+    self.assertEqual(254, self.evaluate(array_ops.rank(sliced)))
+
+  def testNewAxisExceedsMaxRank(self):
+    value = array_ops.fill(
+        [0] + [1] * 224,
+        constant_op.constant(0.0, dtype=dtypes.float32))
+    with self.assertRaisesRegex(
+        (errors.InvalidArgumentError, ValueError), "maximum supported rank"):
+      self.evaluate(
+          gen_array_ops.strided_slice(
+              value,
+              begin=array_ops.zeros([31], dtypes.int32),
+              end=array_ops.zeros([31], dtypes.int32),
+              strides=array_ops.ones([31], dtypes.int32),
+              new_axis_mask=(1 << 30) - 1,
+              ellipsis_mask=1 << 30))
+
   @test_util.assert_no_new_pyobjects_executing_eagerly()
   @test_util.assert_no_garbage_created
   def testTensorSliceEagerMemory(self):
diff --git a/tensorflow/python/kernel_tests/array_ops/stack_op_test.py b/tensorflow/python/kernel_tests/array_ops/stack_op_test.py
index c72e8430421..e2ae6ab4685 100644
--- a/tensorflow/python/kernel_tests/array_ops/stack_op_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/stack_op_test.py
@@ -70,6 +70,23 @@ class StackOpTest(test.TestCase):
             c = array_ops_stack.stack(xs, axis=axis)
             self.assertAllEqual(c, data)
 
+  @test_util.run_in_graph_and_eager_modes
+  def testPackAtMaxRank(self):
+    value = array_ops.fill(
+        [0] + [1] * 252,
+        constant_op.constant(0.0, dtype=dtypes.float32))
+    packed = array_ops_stack.stack([value], axis=0)
+    self.assertEqual(254, self.evaluate(array_ops.rank(packed)))
+
+  @test_util.run_in_graph_and_eager_modes
+  def testPackExceedsMaxRank(self):
+    value = array_ops.fill(
+        [0] + [1] * 253,
+        constant_op.constant(0.0, dtype=dtypes.float32))
+    with self.assertRaisesRegex(
+        (errors.InvalidArgumentError, ValueError), "maximum supported rank"):
+      self.evaluate(array_ops_stack.stack([value], axis=0))
+
   def testSimpleParallelCPU(self):
     # tf.parallel_stack is only supported in graph mode.
     with ops.Graph().as_default():

From 9321927435ff26d840dd82e92a05c98b6fc63577 Mon Sep 17 00:00:00 2001
From: Som Tripathi <somtri@iastate.edu>
Date: Mon, 31 Aug 2026 21:07:53 -0500
Subject: [PATCH 0046/1000] Fix CHECK-failure abort in BiasAdd/BiasAddGrad for
 rank-2 NCHW inputs

BiasAdd and BiasAddGrad with data_format="NCHW" and a rank-2 input abort
the process on GPU with a fatal CHECK failure. GetBiasValueDims reads the
NCHW height as dim_size(2) without guarding that the input has at least 3
dimensions; it is called only from the GPU kernels, so only the GPU path
crashes.

Add an OP_REQUIRES rejecting NCHW inputs below rank 3 with InvalidArgument
in all four Compute methods (CPU and GPU, forward and gradient), after the
existing "at least 2D" check. This changes CPU behavior: CPU currently
accepts rank-2 NCHW and returns a result, and now raises instead, matching
how the op and its tests already treat NCHW as at least 3D.

Fixes #94379
---
 tensorflow/core/kernels/bias_op.cc            | 23 +++++++++++++++
 tensorflow/python/kernel_tests/nn_ops/BUILD   |  1 +
 .../kernel_tests/nn_ops/bias_op_base.py       | 29 +++++++++++++++++++
 3 files changed, 53 insertions(+)

diff --git a/tensorflow/core/kernels/bias_op.cc b/tensorflow/core/kernels/bias_op.cc
index 4fc49d74149..f1b5b175c37 100644
--- a/tensorflow/core/kernels/bias_op.cc
+++ b/tensorflow/core/kernels/bias_op.cc
@@ -104,6 +104,11 @@ class BiasOp : public BinaryOp<T> {
     OP_REQUIRES(context, TensorShapeUtils::IsMatrixOrHigher(input.shape()),
                 errors::InvalidArgument("Input tensor must be at least 2D: ",
                                         input.shape()));
+    OP_REQUIRES(
+        context, data_format_ != FORMAT_NCHW || input.shape().dims() >= 3,
+        errors::InvalidArgument(
+            "NCHW format requires input tensor to be at least 3D: ",
+            input.shape()));
     OP_REQUIRES(context, TensorShapeUtils::IsVector(bias.shape()),
                 errors::InvalidArgument("Biases must be 1D: ", bias.shape()));
 
@@ -175,6 +180,13 @@ class BiasGradOp : public OpKernel {
                 errors::InvalidArgument("Input tensor must be at least 2D: ",
                                         output_backprop.shape()));
 
+    OP_REQUIRES(
+        context,
+        data_format_ != FORMAT_NCHW || output_backprop.shape().dims() >= 3,
+        errors::InvalidArgument(
+            "NCHW format requires input tensor to be at least 3D: ",
+            output_backprop.shape()));
+
     OP_REQUIRES(context,
                 FastBoundsCheck(output_backprop.NumElements(),
                                 std::numeric_limits<int32_t>::max()),
@@ -254,6 +266,11 @@ class BiasOp<GPUDevice, T> : public BinaryOp<T> {
                 absl::InvalidArgumentError(
                     absl::StrCat("Input tensor must be at least 2D: ",
                                  input.shape().DebugString())));
+    OP_REQUIRES(
+        context, data_format_ != FORMAT_NCHW || input.shape().dims() >= 3,
+        absl::InvalidArgumentError(absl::StrCat(
+            "NCHW format requires input tensor to be at least 3D: ",
+            input.shape().DebugString())));
     OP_REQUIRES(context, TensorShapeUtils::IsVector(bias.shape()),
                 absl::InvalidArgumentError(absl::StrCat(
                     "Biases must be 1D: ", bias.shape().DebugString())));
@@ -448,6 +465,12 @@ class BiasGradOp<GPUDevice, T> : public OpKernel {
                 absl::InvalidArgumentError(
                     absl::StrCat("Input tensor must be at least 2D: ",
                                  output_backprop.shape().DebugString())));
+    OP_REQUIRES(
+        context,
+        data_format_ != FORMAT_NCHW || output_backprop.shape().dims() >= 3,
+        absl::InvalidArgumentError(absl::StrCat(
+            "NCHW format requires input tensor to be at least 3D: ",
+            output_backprop.shape().DebugString())));
     int32_t batch, height, width, depth, channel;
     GetBiasValueDims(output_backprop, data_format_, &batch, &height, &width,
                      &depth, &channel);
diff --git a/tensorflow/python/kernel_tests/nn_ops/BUILD b/tensorflow/python/kernel_tests/nn_ops/BUILD
index ec14e23bc41..d4fb562d16c 100644
--- a/tensorflow/python/kernel_tests/nn_ops/BUILD
+++ b/tensorflow/python/kernel_tests/nn_ops/BUILD
@@ -130,6 +130,7 @@ py_library(
         "//tensorflow/python/ops:math_ops",
         "//tensorflow/python/ops:nn_grad",
         "//tensorflow/python/ops:nn_ops",
+        "//tensorflow/python/ops:nn_ops_gen",
         "//tensorflow/python/platform:client_testlib",
         "//third_party/py/numpy",
     ],
diff --git a/tensorflow/python/kernel_tests/nn_ops/bias_op_base.py b/tensorflow/python/kernel_tests/nn_ops/bias_op_base.py
index b70cc349140..a4a066bdb2b 100644
--- a/tensorflow/python/kernel_tests/nn_ops/bias_op_base.py
+++ b/tensorflow/python/kernel_tests/nn_ops/bias_op_base.py
@@ -23,6 +23,7 @@ from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors_impl
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_nn_ops
 from tensorflow.python.ops import gradient_checker
 from tensorflow.python.ops import gradient_checker_v2
 from tensorflow.python.ops import gradients_impl
@@ -112,6 +113,34 @@ class BiasAddTestBase(test.TestCase):
           array_ops.reshape([1, 2], shape=[1, 2]),
           array_ops.reshape([1], shape=[1]))
 
+  def testNCHWRankTooLow(self):
+    # A rank-2 NCHW input aborted the GPU kernel with a CHECK failure
+    # (issue #94379); the added OP_REQUIRES makes the kernel raise
+    # InvalidArgumentError instead. Checked in eager only: in graph mode
+    # BiasAddShape rejects rank-2 NCHW at construction (WithRankAtLeast 3, a
+    # ValueError) before the kernel runs, so a graph-mode test would exercise
+    # the pre-existing shape function rather than this guard. Eager places the
+    # op on GPU when one is present and on CPU otherwise.
+    if not context.executing_eagerly():
+      self.skipTest("kernel guard is reachable only in eager mode")
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                "at least 3D"):
+      nn_ops.bias_add(
+          array_ops.ones([2, 6], dtype=dtypes.float32),
+          array_ops.ones([6], dtype=dtypes.float32),
+          data_format="NCHW")
+
+  def testNCHWRankTooLowGrad(self):
+    # BiasAddGrad has the same NCHW rank requirement and the same eager-only
+    # kernel guard as the forward op (see testNCHWRankTooLow).
+    if not context.executing_eagerly():
+      self.skipTest("kernel guard is reachable only in eager mode")
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                "at least 3D"):
+      gen_nn_ops.bias_add_grad(
+          array_ops.ones([2, 6], dtype=dtypes.float32),
+          data_format="NCHW")
+
   def testIntTypes(self):
     for t in [np.int8, np.int16, np.int32, np.int64]:
       self._testAll(

From 6baa056749e26ad72cb1cb7af16434ed2490e548 Mon Sep 17 00:00:00 2001
From: PrabinDevkota <prabindev789@gmail.com>
Date: Tue, 1 Sep 2026 19:19:36 +0530
Subject: [PATCH 0047/1000] Use absl StatusIs and status_matchers deps in Pack
 and ExpandDims rank-limit tests.

---
 tensorflow/core/kernels/BUILD             | 2 ++
 tensorflow/core/kernels/pack_op_test.cc   | 7 +++----
 tensorflow/core/kernels/shape_ops_test.cc | 7 +++----
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 5a35dff9186..a54aa4342f4 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -1152,6 +1152,7 @@ tf_cc_test(
         "//tensorflow/core:test_main",
         "//tensorflow/core:testlib",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -2411,6 +2412,7 @@ tf_cc_test(
         "//tensorflow/core:test_main",
         "//tensorflow/core:testlib",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
         "@com_google_googletest//:gtest",
     ],
 )
diff --git a/tensorflow/core/kernels/pack_op_test.cc b/tensorflow/core/kernels/pack_op_test.cc
index 30008065875..47b5ef4f956 100644
--- a/tensorflow/core/kernels/pack_op_test.cc
+++ b/tensorflow/core/kernels/pack_op_test.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
 #include "tensorflow/core/framework/fake_input.h"
 #include "tensorflow/core/framework/node_def_builder.h"
 #include "tensorflow/core/framework/tensor.h"
@@ -27,15 +28,13 @@ limitations under the License.
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/kernels/ops_testutil.h"
 #include "tensorflow/core/kernels/ops_util.h"
-#include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
-#include "tensorflow/core/platform/status_matchers.h"
 #include "tensorflow/core/platform/test.h"
 
 namespace tensorflow {
 namespace {
 
-using ::tensorflow::testing::StatusIs;
+using ::absl_testing::StatusIs;
 using ::testing::HasSubstr;
 
 absl::Status MakeZeroElementShape(int rank, TensorShape* shape) {
@@ -78,7 +77,7 @@ TEST_F(PackOpTest, PackingBeyondMaxRankFails) {
   AddInput<float>(input_shape, [](int) { return 0.0f; });
 
   EXPECT_THAT(RunOpKernel(),
-              StatusIs(error::INVALID_ARGUMENT,
+              StatusIs(absl::StatusCode::kInvalidArgument,
                        HasSubstr("maximum supported rank")));
 }
 
diff --git a/tensorflow/core/kernels/shape_ops_test.cc b/tensorflow/core/kernels/shape_ops_test.cc
index 2e639e088a6..c39ba923142 100644
--- a/tensorflow/core/kernels/shape_ops_test.cc
+++ b/tensorflow/core/kernels/shape_ops_test.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
 #include "tensorflow/core/common_runtime/kernel_benchmark_testlib.h"
 #include "tensorflow/core/framework/allocator.h"
 #include "tensorflow/core/framework/fake_input.h"
@@ -35,16 +36,14 @@ limitations under the License.
 #include "tensorflow/core/graph/testlib.h"
 #include "tensorflow/core/kernels/ops_testutil.h"
 #include "tensorflow/core/kernels/ops_util.h"
-#include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
-#include "tensorflow/core/platform/status_matchers.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/platform/test_benchmark.h"
 
 namespace tensorflow {
 namespace {
 
-using ::tensorflow::testing::StatusIs;
+using ::absl_testing::StatusIs;
 using ::testing::HasSubstr;
 
 absl::Status MakeZeroElementShape(int rank, TensorShape* shape) {
@@ -87,7 +86,7 @@ TEST_F(ExpandDimsOpTest, ExpandingBeyondMaxRankFails) {
   AddInputFromList<int32_t>(TensorShape({}), {0});
 
   EXPECT_THAT(RunOpKernel(),
-              StatusIs(error::INVALID_ARGUMENT,
+              StatusIs(absl::StatusCode::kInvalidArgument,
                        HasSubstr("maximum supported rank")));
 }
 

From 096d5b32e617df8ecbcaa88c9cddbfb2c6c89a5e Mon Sep 17 00:00:00 2001
From: PrabinDevkota <prabindev789@gmail.com>
Date: Wed, 2 Sep 2026 12:27:55 +0530
Subject: [PATCH 0048/1000] Add platform status_matchers dep to Pack and
 ExpandDims rank-limit tests.

---
 tensorflow/core/kernels/BUILD | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index a54aa4342f4..8ca0324dcba 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -1151,6 +1151,7 @@ tf_cc_test(
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
         "//tensorflow/core:testlib",
+        "//tensorflow/core/platform:status_matchers",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_googletest//:gtest",
@@ -2411,6 +2412,7 @@ tf_cc_test(
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
         "//tensorflow/core:testlib",
+        "//tensorflow/core/platform:status_matchers",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_googletest//:gtest",

From 9744d9f9356c39b0dc04bf5f63dc5729efd0e7ed Mon Sep 17 00:00:00 2001
From: endorphin13 <elshad.toklayev.business@gmail.com>
Date: Fri, 28 Aug 2026 13:49:51 +0300
Subject: [PATCH 0049/1000] Do not size an empty FFT axis from fft_length

An FFT axis of length 0 skips the "input dimension must have length of at
least fft_length" requirement, but fft_length still sized that axis in the
output. That makes a non-empty output for an empty input, and Compute()
returns early for an empty input without ever writing the output:

    if (input_shape.num_elements() == 0) {
      DCHECK_EQ(0, output_shape.num_elements());
      return;
    }

The DCHECK records the invariant this early return depends on, but it does
not hold for a zero-length FFT axis, and DCHECK compiles out in release
builds. The caller receives the output tensor with whatever the allocator
last left in it.

Keep an empty FFT axis empty in the output, which restores the invariant
the early return assumes and matches the "we pass through empty tensors"
intent of the length check. A zero-length batch dimension is unaffected:
the check only covers the fft_rank inner-most axes, and such an input
already produced an empty output.

Applies to both FFTBase and FFTNBase.
---
 tensorflow/core/kernels/fft_ops.cc              | 16 ++++++++++++++++
 .../python/kernel_tests/signal/fft_ops_test.py  | 17 +++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/tensorflow/core/kernels/fft_ops.cc b/tensorflow/core/kernels/fft_ops.cc
index 964da925ec2..6e07bc9796c 100644
--- a/tensorflow/core/kernels/fft_ops.cc
+++ b/tensorflow/core/kernels/fft_ops.cc
@@ -201,6 +201,14 @@ class FFTBase : public OpKernel {
         uint64_t dim = IsForward() && inner_most && fft_shape[i] != 0
                            ? fft_shape[i] / 2 + 1
                            : fft_shape[i];
+        // An empty FFT axis is passed through by the check above. Keep it
+        // empty in the output too. Sizing it from `fft_length` instead would
+        // give a non-empty output for an empty input, and the
+        // `num_elements() == 0` early return below would then hand back the
+        // output buffer without ever writing it.
+        if (input_shape.dim_size(input_index) == 0) {
+          dim = 0;
+        }
         output_shape.set_dim(output_shape.dims() - fft_rank + i, dim);
       }
     } else {
@@ -329,6 +337,14 @@ class FFTNBase : public OpKernel {
         uint64_t dim = IsForward() && inner_most && fft_shape[i] != 0
                            ? fft_shape[i] / 2 + 1
                            : fft_shape[i];
+        // An empty FFT axis is passed through by the check above. Keep it
+        // empty in the output too. Sizing it from `fft_length` instead would
+        // give a non-empty output for an empty input, and the
+        // `num_elements() == 0` early return below would then hand back the
+        // output buffer without ever writing it.
+        if (input_shape.dim_size(input_index) == 0) {
+          dim = 0;
+        }
         output_shape.set_dim(output_shape.dims() - fft_rank + i, dim);
       } else {
         output_shape.set_dim(output_shape.dims() - fft_rank + i, fft_shape[i]);
diff --git a/tensorflow/python/kernel_tests/signal/fft_ops_test.py b/tensorflow/python/kernel_tests/signal/fft_ops_test.py
index e7a38dc47a9..ed282c0e691 100644
--- a/tensorflow/python/kernel_tests/signal/fft_ops_test.py
+++ b/tensorflow/python/kernel_tests/signal/fft_ops_test.py
@@ -612,6 +612,23 @@ class RFFTOpsTest(BaseFFTOpsTest, parameterized.TestCase):
     x = np.zeros((0,) * dims).astype(np_ctype)
     self.assertEqual(x.shape, self._tf_ifft(x, rank).shape)
 
+  @parameterized.parameters(
+      itertools.product(VALID_FFT_RANKS, range(3), (np.float32, np.float64)))
+  def test_empty_with_explicit_fft_length(self, rank, extra_dims, np_rtype):
+    # An empty FFT axis skips the "input dimension must be at least
+    # fft_length" requirement. `fft_length` must not then size that axis in
+    # the output: the kernel returns early without writing the output for an
+    # empty input, so a non-empty output would be returned uninitialized.
+    np_ctype = np.complex64 if np_rtype == np.float32 else np.complex128
+    dims = rank + extra_dims
+    fft_length = (16,) * rank
+
+    x = np.zeros((0,) * dims).astype(np_rtype)
+    self.assertEqual(0, self._tf_fft(x, rank, fft_length).size)
+
+    x = np.zeros((0,) * dims).astype(np_ctype)
+    self.assertEqual(0, self._tf_ifft(x, rank, fft_length).size)
+
   @parameterized.parameters(itertools.product(
       VALID_FFT_RANKS, range(3), (5, 6), (np.float32, np.float64)))
   def test_basic(self, rank, extra_dims, size, np_rtype):

From 08ed5a2dd47be19a9a62474ce1413bb0ab2161ab Mon Sep 17 00:00:00 2001
From: endorphin13 <elshad.toklayev.business@gmail.com>
Date: Sat, 29 Aug 2026 11:23:33 +0300
Subject: [PATCH 0050/1000] Keep batch dimensions non-empty in
 test_empty_with_explicit_fft_length

The test set every dimension to zero via `(0,) * (rank + extra_dims)`, so for
extra_dims > 0 the batch dimensions were empty too and the output was empty
regardless of how the FFT axes were sized. The assertion then held against the
unfixed kernel as well.

Verified on unfixed TF 2.21.0: 6 of the 9 (rank, extra_dims) combinations
returned size 0 and passed without the fix -- e.g. rank=1, extra_dims=1 gives
output shape (0, 9). Sizing only the FFT axes to zero and keeping the batch
dimensions at 2 makes the same combination return (2, 9), size 18, so the test
now fails without the kernel change and passes with it.
---
 tensorflow/python/kernel_tests/signal/fft_ops_test.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/kernel_tests/signal/fft_ops_test.py b/tensorflow/python/kernel_tests/signal/fft_ops_test.py
index ed282c0e691..c98a8380611 100644
--- a/tensorflow/python/kernel_tests/signal/fft_ops_test.py
+++ b/tensorflow/python/kernel_tests/signal/fft_ops_test.py
@@ -620,13 +620,15 @@ class RFFTOpsTest(BaseFFTOpsTest, parameterized.TestCase):
     # the output: the kernel returns early without writing the output for an
     # empty input, so a non-empty output would be returned uninitialized.
     np_ctype = np.complex64 if np_rtype == np.float32 else np.complex128
-    dims = rank + extra_dims
     fft_length = (16,) * rank
 
-    x = np.zeros((0,) * dims).astype(np_rtype)
+    # Batch dimensions must stay non-empty: if every dimension is zero the
+    # output is empty regardless of how the FFT axes are sized, and the test
+    # would pass against the unfixed kernel too.
+    x = np.zeros((2,) * extra_dims + (0,) * rank).astype(np_rtype)
     self.assertEqual(0, self._tf_fft(x, rank, fft_length).size)
 
-    x = np.zeros((0,) * dims).astype(np_ctype)
+    x = np.zeros((2,) * extra_dims + (0,) * rank).astype(np_ctype)
     self.assertEqual(0, self._tf_ifft(x, rank, fft_length).size)
 
   @parameterized.parameters(itertools.product(

From 034aeba715e4b271ab357bd5c35a2a558450568f Mon Sep 17 00:00:00 2001
From: adi-IL <adi-IL@users.noreply.github.com>
Date: Mon, 31 Aug 2026 02:26:59 +0530
Subject: [PATCH 0051/1000] [LinSpace] Register half and bfloat16 OpKernels for
 LinSpace

The LinSpace op definition in tensorflow/core/ops/math_ops.cc declares
support for T: {bfloat16, half, float, double}. However, the kernel
registrations in tensorflow/core/kernels/sequence_ops.cc only registered
DT_FLOAT and DT_DOUBLE.

Cast step calculations in LinSpaceOp::Compute to static_cast<T> so half
and bfloat16 arithmetic compile cleanly, register TF_CALL_half and
TF_CALL_bfloat16 on CPU and default devices, and add test cases in
sequence_ops_test.cc.
---
 tensorflow/core/kernels/sequence_ops.cc      | 10 ++++--
 tensorflow/core/kernels/sequence_ops_test.cc | 32 ++++++++++++++++++++
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/kernels/sequence_ops.cc b/tensorflow/core/kernels/sequence_ops.cc
index 99bc253349d..70de3a917ac 100644
--- a/tensorflow/core/kernels/sequence_ops.cc
+++ b/tensorflow/core/kernels/sequence_ops.cc
@@ -206,8 +206,10 @@ class LinSpaceOp : public OpKernel {
     auto flat = out->flat<T>();
     flat(0) = start;
     if (num > 1) {
-      const T step = (stop - start) / (num - 1);
-      for (Tnum i = 1; i < num - 1; ++i) flat(i) = start + step * i;
+      const T step = (stop - start) / static_cast<T>(num - 1);
+      for (Tnum i = 1; i < num - 1; ++i) {
+        flat(i) = start + step * static_cast<T>(i);
+      }
       // Ensure final value == stop; float arithmetic won't guarantee this.
       flat(num - 1) = stop;
     }
@@ -230,10 +232,14 @@ class LinSpaceOp : public OpKernel {
   REGISTER_KERNEL(dev, T, int64_t)
 
 #define REGISTER_CPU_KERNEL(T) REGISTER_KERNEL_ALL_NUMS(DEVICE_CPU, T)
+TF_CALL_half(REGISTER_CPU_KERNEL);
+TF_CALL_bfloat16(REGISTER_CPU_KERNEL);
 TF_CALL_float(REGISTER_CPU_KERNEL);
 TF_CALL_double(REGISTER_CPU_KERNEL);
 
 #define REGISTER_DEFAULT_KERNEL(T) REGISTER_KERNEL_ALL_NUMS(DEVICE_DEFAULT, T)
+TF_CALL_half(REGISTER_DEFAULT_KERNEL);
+TF_CALL_bfloat16(REGISTER_DEFAULT_KERNEL);
 TF_CALL_float(REGISTER_DEFAULT_KERNEL);
 TF_CALL_double(REGISTER_DEFAULT_KERNEL);
 #undef REGISTER_DEFAULT_KERNEL
diff --git a/tensorflow/core/kernels/sequence_ops_test.cc b/tensorflow/core/kernels/sequence_ops_test.cc
index 74b2013522b..7a7c3790713 100644
--- a/tensorflow/core/kernels/sequence_ops_test.cc
+++ b/tensorflow/core/kernels/sequence_ops_test.cc
@@ -193,5 +193,37 @@ TEST_F(LinSpaceOpTest, Simple_Double) {
   test::ExpectTensorEqual<double>(expected, *GetOutput(0));
 }
 
+TEST_F(LinSpaceOpTest, Simple_Half) {
+  MakeOp(DT_HALF, DT_INT32);
+
+  // Feed and run
+  AddInputFromArray<Eigen::half>(TensorShape({}), {Eigen::half(3.0f)});
+  AddInputFromArray<Eigen::half>(TensorShape({}), {Eigen::half(7.0f)});
+  AddInputFromArray<int32_t>(TensorShape({}), {3});
+  TF_ASSERT_OK(RunOpKernel());
+
+  // Check the output
+  Tensor expected(allocator(), DT_HALF, TensorShape({3}));
+  test::FillValues<Eigen::half>(
+      &expected, {Eigen::half(3.0f), Eigen::half(5.0f), Eigen::half(7.0f)});
+  test::ExpectTensorEqual<Eigen::half>(expected, *GetOutput(0));
+}
+
+TEST_F(LinSpaceOpTest, Simple_Bfloat16) {
+  MakeOp(DT_BFLOAT16, DT_INT32);
+
+  // Feed and run
+  AddInputFromArray<bfloat16>(TensorShape({}), {bfloat16(3.0f)});
+  AddInputFromArray<bfloat16>(TensorShape({}), {bfloat16(7.0f)});
+  AddInputFromArray<int32_t>(TensorShape({}), {3});
+  TF_ASSERT_OK(RunOpKernel());
+
+  // Check the output
+  Tensor expected(allocator(), DT_BFLOAT16, TensorShape({3}));
+  test::FillValues<bfloat16>(
+      &expected, {bfloat16(3.0f), bfloat16(5.0f), bfloat16(7.0f)});
+  test::ExpectTensorEqual<bfloat16>(expected, *GetOutput(0));
+}
+
 }  // namespace
 }  // namespace tensorflow

From e1486ee6a833e09915f0f1a444d380eb54769b13 Mon Sep 17 00:00:00 2001
From: greypilgrim-083 <prashantchaudhary7353@gmail.com>
Date: Fri, 4 Sep 2026 02:43:25 +0530
Subject: [PATCH 0052/1000] Fix integer overflow in WindowDataset target buffer
 size

---
 tensorflow/core/kernels/data/BUILD            |  3 ++
 .../core/kernels/data/experimental/BUILD      |  3 ++
 .../experimental/sliding_window_dataset_op.cc | 23 +++++++++++--
 .../core/kernels/data/window_dataset_op.cc    | 28 +++++++++++++---
 .../kernels/data/window_dataset_op_test.cc    | 32 +++++++++++++++++++
 5 files changed, 82 insertions(+), 7 deletions(-)

diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD
index 3c5b008c1a7..e0c6f2dc8a1 100644
--- a/tensorflow/core/kernels/data/BUILD
+++ b/tensorflow/core/kernels/data/BUILD
@@ -1643,8 +1643,10 @@ tf_kernel_library(
         "//tensorflow/core/data:name_utils",
         "//tensorflow/core/framework:dataset_options_proto_cc",
         "//tensorflow/core/framework:types_proto_cc",
+        "//tensorflow/core/util:overflow",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:str_format",
     ],
 )
@@ -1672,6 +1674,7 @@ tf_cc_test(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest_main",
+        "//tensorflow/core/platform:status_matchers",
     ],
 )
 
diff --git a/tensorflow/core/kernels/data/experimental/BUILD b/tensorflow/core/kernels/data/experimental/BUILD
index aa253e6864e..f05f7fea840 100644
--- a/tensorflow/core/kernels/data/experimental/BUILD
+++ b/tensorflow/core/kernels/data/experimental/BUILD
@@ -846,10 +846,13 @@ tf_kernel_library(
         "//tensorflow/core:lib_internal",
         "//tensorflow/core/framework:attr_value_proto_cc",
         "//tensorflow/core/framework:dataset_options_proto_cc",
+        "//tensorflow/core/util:overflow",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
     ],
 )
 
diff --git a/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op.cc b/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op.cc
index 375295e0b1e..36b6df2b16c 100644
--- a/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 #include <cstddef>
 #include <cstdint>
+#include <limits>
 #include <deque>
 #include <memory>
 #include <string>
@@ -23,7 +24,9 @@ limitations under the License.
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
 #include "tensorflow/core/framework/attr_value.pb.h"
 #include "tensorflow/core/framework/dataset.h"
 #include "tensorflow/core/framework/dataset_options.pb.h"
@@ -31,6 +34,8 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/util/batch_util.h"
+#include "tensorflow/core/util/overflow.h"
+#include "xla/tsl/platform/statusor.h"
 
 namespace tensorflow {
 namespace data {
@@ -186,7 +191,8 @@ class SlidingWindowDatasetOp : public UnaryDatasetOpKernel {
           batch_elements.reserve(window_size);
 
           // Fill up buffer if not entire data was consumed.
-          size_t target_size = TargetBufferSize(window_size, window_stride);
+          TF_ASSIGN_OR_RETURN(const size_t target_size,
+                              TargetBufferSize(window_size, window_stride));
           for (size_t i = buffer_.size(); i < target_size && input_impl_; ++i) {
             bool end_of_input;
             std::vector<Tensor> element;
@@ -318,8 +324,19 @@ class SlidingWindowDatasetOp : public UnaryDatasetOpKernel {
       }
 
      private:
-      size_t TargetBufferSize(int64_t window_size, int64_t window_stride) {
-        return (window_size - 1) * window_stride + 1;
+      absl::StatusOr<size_t> TargetBufferSize(int64_t window_size,
+                                              int64_t window_stride) {
+        int64_t result = MultiplyWithoutOverflow(window_size - 1, window_stride);
+        if (result >= 0) result = AddWithoutOverflow(result, 1);
+        if (result < 0 || static_cast<uint64_t>(result) >
+                              std::numeric_limits<size_t>::max()) {
+          return absl::InvalidArgumentError(absl::StrFormat(
+              "Window target buffer size overflow: (window_size=%lld - 1) * "
+              "window_stride=%lld + 1 is not representable.",
+              static_cast<long long>(window_size),
+              static_cast<long long>(window_stride)));
+        }
+        return static_cast<size_t>(result);
       }
 
       mutex mu_;
diff --git a/tensorflow/core/kernels/data/window_dataset_op.cc b/tensorflow/core/kernels/data/window_dataset_op.cc
index 67be6216149..ca455914d74 100644
--- a/tensorflow/core/kernels/data/window_dataset_op.cc
+++ b/tensorflow/core/kernels/data/window_dataset_op.cc
@@ -16,6 +16,7 @@ limitations under the License.
 
 #include <cstddef>
 #include <cstdint>
+#include <limits>
 #include <deque>
 #include <memory>
 #include <string>
@@ -24,8 +25,10 @@ limitations under the License.
 
 #include "absl/log/check.h"
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/str_format.h"
 #include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
 #include "tensorflow/core/data/name_utils.h"
 #include "tensorflow/core/framework/dataset.h"
 #include "tensorflow/core/framework/dataset_options.pb.h"
@@ -38,6 +41,7 @@ limitations under the License.
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/kernels/data/window_dataset.h"
 #include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/util/overflow.h"
 #include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/platform/strcat.h"
@@ -124,7 +128,11 @@ class WindowDatasetOp::Dataset : public DatasetBase {
       // of the initial window. If it is negative, we know that the
       // cardinality is 0. Otherwise, it will be the number of valid shifts
       // over the rest_elements.
-      int64_t rest_elements = n - ((window_size_ - 1) * window_stride_ + 1);
+      int64_t target =
+          MultiplyWithoutOverflow(window_size_ - 1, window_stride_);
+      if (target >= 0) target = AddWithoutOverflow(target, 1);
+      // If the target overflows, no finite input can fill a complete window.
+      int64_t rest_elements = target < 0 ? -1 : n - target;
       cardinality = rest_elements < 0 ? 0 : rest_elements / window_shift_ + 1;
     } else {
       cardinality = n / window_shift_ + (n % window_shift_ == 0 ? 0 : 1);
@@ -183,7 +191,8 @@ class WindowDatasetOp::Dataset : public DatasetBase {
       std::vector<std::vector<Tensor>> window_elements;
       absl::Status status = absl::OkStatus();
       {
-        const size_t target_size = TargetBufferSize(window_size, window_stride);
+        TF_ASSIGN_OR_RETURN(const size_t target_size,
+                            TargetBufferSize(window_size, window_stride));
 
         mutex_lock l(mu_);
         if (!input_impl_ &&
@@ -410,8 +419,19 @@ class WindowDatasetOp::Dataset : public DatasetBase {
       return strings::StrCat(kBuffer, "[", index, "]", kErrorMessage);
     }
 
-    size_t TargetBufferSize(int64_t window_size, int64_t window_stride) {
-      return (window_size - 1) * window_stride + 1;
+    absl::StatusOr<size_t> TargetBufferSize(int64_t window_size,
+                                            int64_t window_stride) {
+      int64_t result = MultiplyWithoutOverflow(window_size - 1, window_stride);
+      if (result >= 0) result = AddWithoutOverflow(result, 1);
+      if (result < 0 || static_cast<uint64_t>(result) >
+                            std::numeric_limits<size_t>::max()) {
+        return absl::InvalidArgumentError(absl::StrFormat(
+            "Window target buffer size overflow: (window_size=%lld - 1) * "
+            "window_stride=%lld + 1 is not representable.",
+            static_cast<long long>(window_size),
+            static_cast<long long>(window_stride)));
+      }
+      return static_cast<size_t>(result);
     }
 
     mutex mu_;
diff --git a/tensorflow/core/kernels/data/window_dataset_op_test.cc b/tensorflow/core/kernels/data/window_dataset_op_test.cc
index bb767b8eb83..d0604fba883 100644
--- a/tensorflow/core/kernels/data/window_dataset_op_test.cc
+++ b/tensorflow/core/kernels/data/window_dataset_op_test.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/strings/match.h"
 #include "xla/tsl/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/status_matchers.h"
 #include "tensorflow/core/data/dataset_test_base.h"
 #include "tensorflow/core/data/name_utils.h"
 #include "tensorflow/core/data/serialization_utils.h"
@@ -626,6 +627,37 @@ TEST_F(WindowDatasetOpTest, MalformedCheckpoint_TriggersOutOfBoundsRead) {
       << status.message();
 }
 
+// Regression test: large window_size and window_stride whose target buffer
+// size (window_size - 1) * window_stride + 1 overflows int64. Before the fix,
+// the overflow silently wrapped to a small value and emitted undersized
+// windows. After the fix, iteration must return InvalidArgument.
+WindowDatasetParams WindowDatasetParamsWithOverflowingSizeAndStride() {
+  // (2^32 + 1 - 1) * 2^32 + 1 = 2^64 + 1, which overflows int64/size_t.
+  const int64_t large_size = (static_cast<int64_t>(1) << 32) + 1;
+  const int64_t large_stride = static_cast<int64_t>(1) << 32;
+  return WindowDatasetParams(RangeDatasetParams(0, 3, 1),
+                             /*size=*/large_size,
+                             /*shift=*/1,
+                             /*stride=*/large_stride,
+                             /*drop_remainder=*/true,
+                             /*output_dtypes=*/{DT_VARIANT},
+                             /*output_shapes=*/{PartialTensorShape({})},
+                             /*node_name=*/kNodeName);
+}
+
+TEST_F(WindowDatasetOpTest, OverflowingTargetBufferSize) {
+  auto dataset_params = WindowDatasetParamsWithOverflowingSizeAndStride();
+  TF_ASSERT_OK(Initialize(dataset_params));
+
+  bool end_of_sequence = false;
+  std::vector<Tensor> out_tensors;
+  EXPECT_THAT(
+      iterator_->GetNext(iterator_ctx_.get(), &out_tensors, &end_of_sequence),
+      tensorflow::testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                                    ::testing::HasSubstr("overflow")));
+}
+
 }  // namespace
 }  // namespace data
 }  // namespace tensorflow
+

From 533a1075d4310c3ccf4c5aab3db7da35063aac87 Mon Sep 17 00:00:00 2001
From: Arpan Chatterjee <arpanchatterjee648@gmail.com>
Date: Fri, 4 Sep 2026 21:54:33 +0530
Subject: [PATCH 0053/1000] Correct grammar related to a/an

---
 tensorflow/compiler/mlir/lite/flatbuffer_operator.h           | 2 +-
 .../mlir/tensorflow/transforms/init_text_file_to_import.cc    | 2 +-
 tensorflow/compiler/mlir/tfr/ir/tfr_ops.td                    | 2 +-
 tensorflow/compiler/mlir/tfr/passes/raise_to_tf.cc            | 4 ++--
 .../tools/kernel_gen/transforms/broadcast_propagation_pass.cc | 2 +-
 tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h         | 2 +-
 .../core/api_def/base_api/api_def_IteratorGetModelProto.pbtxt | 2 +-
 tensorflow/core/common_runtime/stats_publisher_interface.cc   | 2 +-
 tensorflow/core/data/service/worker_impl.h                    | 2 +-
 .../core/distributed_runtime/rendezvous_mgr_interface.h       | 2 +-
 tensorflow/core/distributed_runtime/rpc/rpc_rendezvous_mgr.h  | 2 +-
 tensorflow/core/ir/ops.cc                                     | 2 +-
 tensorflow/core/kernels/batching_util/input_split_metadata.h  | 2 +-
 tensorflow/core/tfrt/ifrt/sharding_utils.h                    | 2 +-
 tensorflow/core/transforms/region_to_functional/impl.cc       | 2 +-
 tensorflow/core/util/mkl_util.h                               | 2 +-
 tensorflow/dtensor/cc/constants.h                             | 2 +-
 tensorflow/go/op/wrappers.go                                  | 2 +-
 tensorflow/lite/core/subgraph.h                               | 4 ++--
 tensorflow/lite/delegates/flex/test_util.h                    | 2 +-
 .../delegates/gpu/common/testing/feature_parity/README.md     | 2 +-
 .../org/tensorflow/lite/task/core/BaseOptions.Builder.html    | 2 +-
 tensorflow/lite/g3doc/examples/style_transfer/overview.ipynb  | 2 +-
 tensorflow/lite/java/ovic/README.md                           | 2 +-
 .../internal/optimized/depthwiseconv_uint8_3x3_filter.h       | 2 +-
 tensorflow/lite/nnapi/nnapi_implementation.cc                 | 2 +-
 .../fuse_binary_into_following_affine.cc                      | 2 +-
 .../fuse_binary_into_preceding_affine.cc                      | 2 +-
 tensorflow/lite/tools/optimize/quantization_wrapper.h         | 2 +-
 tensorflow/python/autograph/converters/control_flow.py        | 2 +-
 tensorflow/python/autograph/lang/special_functions.py         | 2 +-
 tensorflow/python/data/kernel_tests/from_tensors_test.py      | 2 +-
 tensorflow/python/data/kernel_tests/map_test.py               | 4 ++--
 .../python/distribute/coordinator/cluster_coordinator.py      | 2 +-
 .../python/eager/polymorphic_function/function_type_utils.py  | 2 +-
 tensorflow/python/framework/tensor_util.py                    | 2 +-
 tensorflow/python/keras/callbacks.py                          | 2 +-
 tensorflow/python/ops/numpy_ops/tests/np_test.py              | 2 +-
 tensorflow/python/tpu/tpu_embedding_v2.py                     | 2 +-
 .../xla/backends/gpu/transforms/windowed_einsum_handler.cc    | 2 +-
 third_party/xla/xla/hlo/builder/value_inference.cc            | 2 +-
 third_party/xla/xla/hlo/builder/xla_builder.h                 | 2 +-
 third_party/xla/xla/hlo/utils/hlo_sharding_util.h             | 2 +-
 third_party/xla/xla/service/call_graph.h                      | 2 +-
 third_party/xla/xla/service/gpu/triton_tiling_propagation.h   | 2 +-
 third_party/xla/xla/service/layout_assignment.cc              | 2 +-
 third_party/xla/xla/service/llvm_ir/loop_emitter.h            | 2 +-
 third_party/xla/xla/service/while_loop_concat_code_motion.cc  | 2 +-
 third_party/xla/xla/stream_executor/blas.h                    | 2 +-
 third_party/xla/xla/tools/hlo_extractor.cc                    | 2 +-
 50 files changed, 53 insertions(+), 53 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/flatbuffer_operator.h b/tensorflow/compiler/mlir/lite/flatbuffer_operator.h
index 0fbd47e5615..e3e558fbf87 100644
--- a/tensorflow/compiler/mlir/lite/flatbuffer_operator.h
+++ b/tensorflow/compiler/mlir/lite/flatbuffer_operator.h
@@ -128,7 +128,7 @@ void BuiltinOptionsToAttributes(
     // NOLINTNEXTLINE
     llvm::SmallVectorImpl<mlir::NamedAttribute>& attributes);
 
-// While the last several tensors could be optional tensors for an tfl op, the
+// While the last several tensors could be optional tensors for a tfl op, the
 // number of input operands could vary. This function gets the min/max number of
 // operands from tflite op name.
 llvm::MinMax OperandNumbersMinMax(llvm::StringRef op_name);
diff --git a/tensorflow/compiler/mlir/tensorflow/transforms/init_text_file_to_import.cc b/tensorflow/compiler/mlir/tensorflow/transforms/init_text_file_to_import.cc
index fd15491994d..2fbb5fd190e 100644
--- a/tensorflow/compiler/mlir/tensorflow/transforms/init_text_file_to_import.cc
+++ b/tensorflow/compiler/mlir/tensorflow/transforms/init_text_file_to_import.cc
@@ -122,7 +122,7 @@ class ConvertInitializeTableFromTextFileV2
     line_nums.resize(lines.size());
     std::iota(line_nums.begin(), line_nums.end(), 0);
 
-    // Create constant ops for keys an values.
+    // Create constant ops for keys and values.
     Value key_constant_tensor = arith::ConstantOp::create(
         rewriter, op.getLoc(),
         DenseStringElementsAttr::get(
diff --git a/tensorflow/compiler/mlir/tfr/ir/tfr_ops.td b/tensorflow/compiler/mlir/tfr/ir/tfr_ops.td
index f9a37653619..2138d60f139 100644
--- a/tensorflow/compiler/mlir/tfr/ir/tfr_ops.td
+++ b/tensorflow/compiler/mlir/tfr/ir/tfr_ops.td
@@ -341,7 +341,7 @@ def TFR_GetElementOp : TFR_Op<"get_element", [Pure]> {
 def TFR_BuildListOp : TFR_Op<"build_list", [Pure]> {
   let description = [{
    The `build_list` operation builds a tensor list from a list of tensors, or
-   an tfr.attr from a list of scalars.
+   a tfr.attr from a list of scalars.
 
     Example:
 
diff --git a/tensorflow/compiler/mlir/tfr/passes/raise_to_tf.cc b/tensorflow/compiler/mlir/tfr/passes/raise_to_tf.cc
index 5dd6a22f90c..97d95ea9933 100644
--- a/tensorflow/compiler/mlir/tfr/passes/raise_to_tf.cc
+++ b/tensorflow/compiler/mlir/tfr/passes/raise_to_tf.cc
@@ -183,7 +183,7 @@ LogicalResult RewriteTFRCallOp::AddDerivedAttrs(
     ArrayRef<Attribute> input_types,
     llvm::StringMap<Attribute>* derived_attrs) const {
   // If there is an attribute associated to the input in the signature, we
-  // store it as an derived attribute.
+  // store it as a derived attribute.
   if (auto tensor_type = mlir::dyn_cast<TFRTensorType>(input_tfr_type)) {
     auto attr_names = tensor_type.getAttrKeys();
     if (attr_names.empty()) return success();
@@ -195,7 +195,7 @@ LogicalResult RewriteTFRCallOp::AddDerivedAttrs(
   }
 
   // If there is an attribute associated to the input in the signature,
-  // we store it as an derived attribute.
+  // we store it as a derived attribute.
   if (auto list_type = mlir::dyn_cast<TFRTensorListType>(input_tfr_type)) {
     auto attr_names = list_type.getAttrKeys();
     if (attr_names.empty()) return success();
diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/broadcast_propagation_pass.cc b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/broadcast_propagation_pass.cc
index 0d300b56918..9c571bd9ec9 100644
--- a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/broadcast_propagation_pass.cc
+++ b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/broadcast_propagation_pass.cc
@@ -128,7 +128,7 @@ DenseIntElementsAttr composeBroadcastDimensionsAttr(OpBuilder &builder,
 }
 
 // Find all the broadcast intents and their dependencies. Start analyzing from
-// the root an collect all broadcast intents that can help broadcast propagation
+// the root a collect all broadcast intents that can help broadcast propagation
 // from there.
 void findBroadcastIntents(
     DynamicBroadcastInDimOp root, Block *parentBlock,
diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h
index 2812aa06457..d53ffc2d902 100644
--- a/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h
+++ b/tensorflow/compiler/tf2tensorrt/utils/trt_allocator.h
@@ -40,7 +40,7 @@ namespace tensorrt {
 class TRTBaseAllocator : public nvinfer1::IGpuAllocator {
   // Base allocator class so we can have a virtual destructor;
  public:
-  // python wrapper seems to be not happy with an pure virtual destructor;
+  // python wrapper seems to be not happy with a pure virtual destructor;
   virtual ~TRTBaseAllocator() = default;
 };
 
diff --git a/tensorflow/core/api_def/base_api/api_def_IteratorGetModelProto.pbtxt b/tensorflow/core/api_def/base_api/api_def_IteratorGetModelProto.pbtxt
index ebab516807a..67e9d40f034 100644
--- a/tensorflow/core/api_def/base_api/api_def_IteratorGetModelProto.pbtxt
+++ b/tensorflow/core/api_def/base_api/api_def_IteratorGetModelProto.pbtxt
@@ -4,7 +4,7 @@ op {
   in_arg {
     name: "iterator"
     description: <<END
-An resource from an dataset iterator.
+An resource from a dataset iterator.
 END
   }
   out_arg {
diff --git a/tensorflow/core/common_runtime/stats_publisher_interface.cc b/tensorflow/core/common_runtime/stats_publisher_interface.cc
index 610efbdadb7..7b7a785d4e9 100644
--- a/tensorflow/core/common_runtime/stats_publisher_interface.cc
+++ b/tensorflow/core/common_runtime/stats_publisher_interface.cc
@@ -26,7 +26,7 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
-// NoOpStatsPublisher provides an dummy/no-op implementation of
+// NoOpStatsPublisher provides a dummy/no-op implementation of
 // StatsPublisherInterface.
 class NoOpStatsPublisher : public StatsPublisherInterface {
  public:
diff --git a/tensorflow/core/data/service/worker_impl.h b/tensorflow/core/data/service/worker_impl.h
index c256c88c226..54735440cab 100644
--- a/tensorflow/core/data/service/worker_impl.h
+++ b/tensorflow/core/data/service/worker_impl.h
@@ -158,7 +158,7 @@ class DataServiceWorkerImpl {
   // Updates the distributed snapshot tasks according to the heartbeat response.
   absl::Status UpdateSnapshotWriters(const WorkerHeartbeatResponse& response)
       TF_LOCKS_EXCLUDED(mu_);
-  // Creates an dataset iterator for snapshot writers.
+  // Creates a dataset iterator for snapshot writers.
   absl::StatusOr<std::unique_ptr<StandaloneTaskIterator>>
   MakeSnapshotTaskIterator(const SnapshotTaskDef& snapshot_task,
                            const DatasetDef& dataset_def) const;
diff --git a/tensorflow/core/distributed_runtime/rendezvous_mgr_interface.h b/tensorflow/core/distributed_runtime/rendezvous_mgr_interface.h
index 6ec759d4544..d599b80093c 100644
--- a/tensorflow/core/distributed_runtime/rendezvous_mgr_interface.h
+++ b/tensorflow/core/distributed_runtime/rendezvous_mgr_interface.h
@@ -59,7 +59,7 @@ class RemoteRendezvous : public Rendezvous {
 //
 // E.g.,
 //   Rendezvous* rendez = worker_env->rendezvous_mgr->Find(0x8935);
-//   fork execution of an graph executor using "rendez"  on thread 1;
+//   fork execution of a graph executor using "rendez"  on thread 1;
 //   fork execution of another graph executor using "rendez" on thread 2;
 //   ...
 //   join threads 1 and 2;
diff --git a/tensorflow/core/distributed_runtime/rpc/rpc_rendezvous_mgr.h b/tensorflow/core/distributed_runtime/rpc/rpc_rendezvous_mgr.h
index 42eda4ea71a..744baab4808 100644
--- a/tensorflow/core/distributed_runtime/rpc/rpc_rendezvous_mgr.h
+++ b/tensorflow/core/distributed_runtime/rpc/rpc_rendezvous_mgr.h
@@ -32,7 +32,7 @@ class DeviceMgr;
 //
 // E.g.,
 //   Rendezvous* rendez = worker_env->rendezvous_mgr->Find(0x8935);
-//   fork execution of an graph executor using "rendez"  on thread 1;
+//   fork execution of a graph executor using "rendez"  on thread 1;
 //   fork execution of another graph executor using "rendez" on thread 2;
 //   ...
 //   join threads 1 and 2;
diff --git a/tensorflow/core/ir/ops.cc b/tensorflow/core/ir/ops.cc
index 6b7b31f149f..1305dc5f959 100644
--- a/tensorflow/core/ir/ops.cc
+++ b/tensorflow/core/ir/ops.cc
@@ -1195,7 +1195,7 @@ static LogicalResult VerifyIfLikeRegionOp(IfLikeRegionOp op) {
       op, {op.getThenRegionAttrsAttr(), op.getElseRegionAttrsAttr()});
 }
 
-// Given an potentially null attribute that would represent a constant value,
+// Given a potentially null attribute that would represent a constant value,
 // try to narrow it to a statically known condition.
 // TODO(jeffniu): Incorporate the other cases of `tf.ToBool`.
 static std::optional<bool> GetStaticallyKnownBranch(Attribute cond_attr) {
diff --git a/tensorflow/core/kernels/batching_util/input_split_metadata.h b/tensorflow/core/kernels/batching_util/input_split_metadata.h
index 429858d318a..1fca8c6e363 100644
--- a/tensorflow/core/kernels/batching_util/input_split_metadata.h
+++ b/tensorflow/core/kernels/batching_util/input_split_metadata.h
@@ -23,7 +23,7 @@ limitations under the License.
 namespace tensorflow {
 namespace serving {
 namespace internal {
-// InputSplitMetadata represents the task sizes of an batch-task after it's
+// InputSplitMetadata represents the task sizes of a batch-task after it's
 // tailored according to queue status (`open_batch_remaining_slot` and
 // `batch_size_limit`).
 //
diff --git a/tensorflow/core/tfrt/ifrt/sharding_utils.h b/tensorflow/core/tfrt/ifrt/sharding_utils.h
index 1f8f3fc8e6a..e6392986030 100644
--- a/tensorflow/core/tfrt/ifrt/sharding_utils.h
+++ b/tensorflow/core/tfrt/ifrt/sharding_utils.h
@@ -116,7 +116,7 @@ absl::StatusOr<xla::ifrt::ArrayRef> MakeArrayFromTensor(
     xla::ifrt::ShardingRef sharding, const tsl::thread::ThreadPool& thread_pool,
     const xla::ifrt::LayoutRef& xla_input_layout);
 
-// Reshard an disassembled array list back to one single tensor
+// Reshard a disassembled array list back to one single tensor
 // based on given sharding spec.
 //
 // input_array: the input device buffers.
diff --git a/tensorflow/core/transforms/region_to_functional/impl.cc b/tensorflow/core/transforms/region_to_functional/impl.cc
index 844186fb558..651541c32fb 100644
--- a/tensorflow/core/transforms/region_to_functional/impl.cc
+++ b/tensorflow/core/transforms/region_to_functional/impl.cc
@@ -363,7 +363,7 @@ using ConvertStatelessIfOp =
     ConvertIfLikeOp<StatelessIfRegionOp, StatelessIfOp>;
 using ConvertStatefulIfOp = ConvertIfLikeOp<StatefulIfRegionOp, StatefulIfOp>;
 
-// Base class for patterns to convert an case-like TFG region op to
+// Base class for patterns to convert a case-like TFG region op to
 // functional form.
 template <typename CaseLikeRegionOp, typename CaseLikeOp>
 struct ConvertCaseLikeOp
diff --git a/tensorflow/core/util/mkl_util.h b/tensorflow/core/util/mkl_util.h
index fd1762482b4..86089d4aa2e 100644
--- a/tensorflow/core/util/mkl_util.h
+++ b/tensorflow/core/util/mkl_util.h
@@ -429,7 +429,7 @@ class MklDnnShape {
     return reinterpret_cast<const int*>(&data_.sizes_[0]);
   }
 
-  // Returns an dnnl::memory::dims object that contains the sizes of this
+  // Returns a dnnl::memory::dims object that contains the sizes of this
   // MklDnnShape object.
   inline memory::dims GetSizesAsMklDnnDims() const {
     memory::dims retVal;
diff --git a/tensorflow/dtensor/cc/constants.h b/tensorflow/dtensor/cc/constants.h
index 8f11d62981d..31c7fe65877 100644
--- a/tensorflow/dtensor/cc/constants.h
+++ b/tensorflow/dtensor/cc/constants.h
@@ -110,7 +110,7 @@ static constexpr char kDeviceSeedForMeshDims[] =
 
 // Attribute that determines whether to skip XlA compilation. There are some ops
 // that run on a TPU mesh but are not expected to be compiled by XLA, e.g.
-// VarHandleOp, DestroyResourceOp, etc. For such an case, set this attribute
+// VarHandleOp, DestroyResourceOp, etc. For such a case, set this attribute
 // to true on the StatefulPartitionedCallOp generated by MLIR lowering.
 static constexpr char kSkipXlaCompilation[] = "_skip_xla_compilation";
 
diff --git a/tensorflow/go/op/wrappers.go b/tensorflow/go/op/wrappers.go
index 013fa6d792e..acba7bb4ce8 100644
--- a/tensorflow/go/op/wrappers.go
+++ b/tensorflow/go/op/wrappers.go
@@ -22055,7 +22055,7 @@ func IteratorGetDevice(scope *Scope, resource tf.Output) (device tf.Output) {
 //
 // Arguments:
 //
-//	iterator: An resource from an dataset iterator.
+//	iterator: An resource from a dataset iterator.
 //
 // Returns A serialized model proto.
 func IteratorGetModelProto(scope *Scope, iterator tf.Output) (model_proto tf.Output) {
diff --git a/tensorflow/lite/core/subgraph.h b/tensorflow/lite/core/subgraph.h
index e1f4f0e3c44..e19aa3919ac 100644
--- a/tensorflow/lite/core/subgraph.h
+++ b/tensorflow/lite/core/subgraph.h
@@ -818,14 +818,14 @@ class Subgraph {
   // Does not report invoke status through profiler.
   TfLiteStatus ModifyGraphWithDelegateImpl(TfLiteDelegate* delegate);
 
-  // Request an tensor be resized implementation. If the given tensor is of
+  // Request a tensor be resized implementation. If the given tensor is of
   // type kTfLiteDynamic it will also be allocated new memory.
   TfLiteStatus ResizeTensorImpl(TfLiteTensor* tensor, TfLiteIntArray* new_size);
 
   // Report a detailed error string (will be printed to stderr).
   void ReportErrorImpl(const char* format, va_list args);
 
-  // Entry point for C node plugin API to request an tensor be resized.
+  // Entry point for C node plugin API to request a tensor be resized.
   static TfLiteStatus ResizeTensor(TfLiteContext* context, TfLiteTensor* tensor,
                                    TfLiteIntArray* new_size);
   // Entry point for C node plugin API to report an error.
diff --git a/tensorflow/lite/delegates/flex/test_util.h b/tensorflow/lite/delegates/flex/test_util.h
index 71eeead5640..333c4251470 100644
--- a/tensorflow/lite/delegates/flex/test_util.h
+++ b/tensorflow/lite/delegates/flex/test_util.h
@@ -33,7 +33,7 @@ enum TfOpType {
   kLoopCond,
   // Represents an op that does not exist in TensorFlow.
   kNonExistent,
-  // Represents an valid TensorFlow op where the NodeDef is incompatible.
+  // Represents a valid TensorFlow op where the NodeDef is incompatible.
   kIncompatibleNodeDef,
 };
 
diff --git a/tensorflow/lite/delegates/gpu/common/testing/feature_parity/README.md b/tensorflow/lite/delegates/gpu/common/testing/feature_parity/README.md
index f81f585963b..3c061701b8e 100644
--- a/tensorflow/lite/delegates/gpu/common/testing/feature_parity/README.md
+++ b/tensorflow/lite/delegates/gpu/common/testing/feature_parity/README.md
@@ -2,6 +2,6 @@
 
 These tests ensure feature parity across TFLite GPU delegates. Every test
 receives a simple automatically generated tflite model as an input. Model runs
-with default tflite cpu interpreter and an delegated interpreter. Test succeeds
+with default tflite cpu interpreter and a delegated interpreter. Test succeeds
 when output results for both interpretes match with the given accuracy.
 
diff --git a/tensorflow/lite/g3doc/api_docs/java/org/tensorflow/lite/task/core/BaseOptions.Builder.html b/tensorflow/lite/g3doc/api_docs/java/org/tensorflow/lite/task/core/BaseOptions.Builder.html
index 2c46fdfb08d..4f467c96899 100644
--- a/tensorflow/lite/g3doc/api_docs/java/org/tensorflow/lite/task/core/BaseOptions.Builder.html
+++ b/tensorflow/lite/g3doc/api_docs/java/org/tensorflow/lite/task/core/BaseOptions.Builder.html
@@ -363,7 +363,7 @@
 <div class="jd-details-descr">
 <div class="jd-tagdata jd-tagdescr"><p>Sets the advanced accelerator options.
 
- <p>Note: this method will override those highlevel API to choose an delegate, such as <code><a href="/lite/api_docs/java/org/tensorflow/lite/task/core/BaseOptions.Builder.html#useGpu()">useGpu()</a></code> and <code><a href="/lite/api_docs/java/org/tensorflow/lite/task/core/BaseOptions.Builder.html#useNnapi()">useNnapi()</a></code>.
+ <p>Note: this method will override those highlevel API to choose a delegate, such as <code><a href="/lite/api_docs/java/org/tensorflow/lite/task/core/BaseOptions.Builder.html#useGpu()">useGpu()</a></code> and <code><a href="/lite/api_docs/java/org/tensorflow/lite/task/core/BaseOptions.Builder.html#useNnapi()">useNnapi()</a></code>.
 </p></p></div>
 <div class="jd-tagdata">
 <h5 class="jd-tagtitle">Parameters</h5>
diff --git a/tensorflow/lite/g3doc/examples/style_transfer/overview.ipynb b/tensorflow/lite/g3doc/examples/style_transfer/overview.ipynb
index 87039d19a69..1bfde6b15af 100644
--- a/tensorflow/lite/g3doc/examples/style_transfer/overview.ipynb
+++ b/tensorflow/lite/g3doc/examples/style_transfer/overview.ipynb
@@ -223,7 +223,7 @@
         "\n",
         "  return img\n",
         "\n",
-        "# Function to pre-process by resizing an central cropping it.\n",
+        "# Function to pre-process by resizing a central cropping it.\n",
         "def preprocess_image(image, target_dim):\n",
         "  # Resize the image so that the shorter dimension becomes 256px.\n",
         "  shape = tf.cast(tf.shape(image)[1:-1], tf.float32)\n",
diff --git a/tensorflow/lite/java/ovic/README.md b/tensorflow/lite/java/ovic/README.md
index 62d8183c23d..ba5aacf4eb1 100644
--- a/tensorflow/lite/java/ovic/README.md
+++ b/tensorflow/lite/java/ovic/README.md
@@ -18,7 +18,7 @@ Note: for now the tests only provides correctness checks, i.e. classifier
 predicts the correct category on the test image, but no on-device latency
 measurements. To test the latency measurement functionality, the tests will
 print the latency running on a desktop computer, which is not indicative of the
-on-device run-time. We are releasing an benchmarker Apk that would allow
+on-device run-time. We are releasing a benchmarker Apk that would allow
 developers to measure latency on their own devices.
 
 ### Obtain the sample models
diff --git a/tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h b/tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h
index 633feb53528..d979341cfb8 100644
--- a/tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h
+++ b/tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h
@@ -12934,7 +12934,7 @@ inline void DepthwiseConvDotProduct3x3Impl(
   // Note that this data is adjusted in a per-depth process before the main
   // filters. The adjustment accounts for a non-symmetric input offset.
   //
-  // Kernel subroutines need to be able to operate consistently on an bias
+  // Kernel subroutines need to be able to operate consistently on a bias
   // array. Where there is no bias, we provide one filled with zeros.
   constexpr int kMinBiasLoad = 8;
   int32_t zero_bias_data[kMinBiasLoad];
diff --git a/tensorflow/lite/nnapi/nnapi_implementation.cc b/tensorflow/lite/nnapi/nnapi_implementation.cc
index 5705576e894..75ca0766399 100644
--- a/tensorflow/lite/nnapi/nnapi_implementation.cc
+++ b/tensorflow/lite/nnapi/nnapi_implementation.cc
@@ -200,7 +200,7 @@ const NnApi LoadNnApi() {
   static const char nnapi_library_name[] = "libneuralnetworks.so";
   libneuralnetworks = dlopen(nnapi_library_name, RTLD_LAZY | RTLD_LOCAL);
 #ifdef __ANDROID__
-  // Note: If there is an problem trying to open the NNAPI library on a
+  // Note: If there is a problem trying to open the NNAPI library on a
   // non-Android system, the error message is suppressed. This is to avoid
   // showing confusing errors when running in environments that do not support
   // NNAPI. As more platforms support NNAPI, the #ifdef logic above can be
diff --git a/tensorflow/lite/toco/graph_transformations/fuse_binary_into_following_affine.cc b/tensorflow/lite/toco/graph_transformations/fuse_binary_into_following_affine.cc
index 813407cc564..0d7681875db 100644
--- a/tensorflow/lite/toco/graph_transformations/fuse_binary_into_following_affine.cc
+++ b/tensorflow/lite/toco/graph_transformations/fuse_binary_into_following_affine.cc
@@ -167,7 +167,7 @@ absl::Status FuseBinaryIntoFollowingAffine::Run(Model* model,
 
   CHECK_EQ(binary_op->inputs.size(), 2);
 
-  // We only can fuse an binary when the two operands break down as follows:
+  // We only can fuse a binary when the two operands break down as follows:
   //   1. One operand is the (variable) output of a typical affine (linear plus
   //   bias)
   //      op of a finite list of possible types: at the moment Conv,
diff --git a/tensorflow/lite/toco/graph_transformations/fuse_binary_into_preceding_affine.cc b/tensorflow/lite/toco/graph_transformations/fuse_binary_into_preceding_affine.cc
index 282eb19e23f..d89883c16e0 100644
--- a/tensorflow/lite/toco/graph_transformations/fuse_binary_into_preceding_affine.cc
+++ b/tensorflow/lite/toco/graph_transformations/fuse_binary_into_preceding_affine.cc
@@ -222,7 +222,7 @@ absl::Status FuseBinaryIntoPrecedingAffine::Run(Model* model,
 
   CHECK_EQ(binary_op->inputs.size(), 2);
 
-  // We only can fuse an binary when the two operands break down as follows:
+  // We only can fuse a binary when the two operands break down as follows:
   //   1. One operand is the (variable) output of a typical affine (linear plus
   //   bias)
   //      op of a finite list of possible types: at the moment Conv,
diff --git a/tensorflow/lite/tools/optimize/quantization_wrapper.h b/tensorflow/lite/tools/optimize/quantization_wrapper.h
index 125bad3e257..4b6852e8286 100644
--- a/tensorflow/lite/tools/optimize/quantization_wrapper.h
+++ b/tensorflow/lite/tools/optimize/quantization_wrapper.h
@@ -20,7 +20,7 @@ limitations under the License.
 namespace tflite {
 namespace optimize {
 
-// Makes an copy of the model at input_path and writes it to output_path, adding
+// Makes a copy of the model at input_path and writes it to output_path, adding
 // tensors to the model needed for calibration.
 // Returns true if it is successful.
 // Example: a/b/c.tflite becomes a/b/c.calibrated.tflite and has
diff --git a/tensorflow/python/autograph/converters/control_flow.py b/tensorflow/python/autograph/converters/control_flow.py
index ad50d0338b1..b28a06353f6 100644
--- a/tensorflow/python/autograph/converters/control_flow.py
+++ b/tensorflow/python/autograph/converters/control_flow.py
@@ -37,7 +37,7 @@ class _Function(object):
 
 
 class ControlFlowTransformer(converter.Base):
-  """Transforms control flow structures like loops an conditionals."""
+  """Transforms control flow structures like loops a conditionals."""
 
   def visit_Lambda(self, node):
     with self.state[_Function] as fn:
diff --git a/tensorflow/python/autograph/lang/special_functions.py b/tensorflow/python/autograph/lang/special_functions.py
index a2fa91ef6cd..e2ee732a885 100644
--- a/tensorflow/python/autograph/lang/special_functions.py
+++ b/tensorflow/python/autograph/lang/special_functions.py
@@ -54,7 +54,7 @@ def tensor_list(elements,
                 element_dtype=None,
                 element_shape=None,
                 use_tensor_array=False):
-  """Creates an tensor list and populates it with the given elements.
+  """Creates a tensor list and populates it with the given elements.
 
   This function provides a more uniform access to tensor lists and tensor
   arrays, and allows optional initialization.
diff --git a/tensorflow/python/data/kernel_tests/from_tensors_test.py b/tensorflow/python/data/kernel_tests/from_tensors_test.py
index 29e7f46be20..97a92ca5b02 100644
--- a/tensorflow/python/data/kernel_tests/from_tensors_test.py
+++ b/tensorflow/python/data/kernel_tests/from_tensors_test.py
@@ -119,7 +119,7 @@ class FromTensorsTest(test_base.DatasetTestBase, parameterized.TestCase):
 
   @combinations.generate(test_base.default_test_combinations())
   def testFromTensorsMixed(self):
-    """Test an dataset that represents a single tuple of tensors."""
+    """Test a dataset that represents a single tuple of tensors."""
     components = (np.array(1), np.array([1, 2, 3]), np.array(37.0),
                   sparse_tensor.SparseTensorValue(
                       indices=np.array([[0]]),
diff --git a/tensorflow/python/data/kernel_tests/map_test.py b/tensorflow/python/data/kernel_tests/map_test.py
index 9defbe04cbf..531cbf4e284 100644
--- a/tensorflow/python/data/kernel_tests/map_test.py
+++ b/tensorflow/python/data/kernel_tests/map_test.py
@@ -209,7 +209,7 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase):
 
   @combinations.generate(_test_combinations())
   def testMapDataset(self, apply_map):
-    """Test an dataset that maps a TF function across its input elements."""
+    """Test a dataset that maps a TF function across its input elements."""
     # The pipeline is TensorSliceDataset -> MapDataset(square_3) ->
     # RepeatDataset(count).
     components = (np.arange(7),
@@ -290,7 +290,7 @@ class MapTest(test_base.DatasetTestBase, parameterized.TestCase):
           combinations.combine(use_unbounded_threadpool=[None, True, False])))
   def testParallelMapDataset(self, apply_map, num_parallel_calls, buffer_size,
                              use_unbounded_threadpool):
-    """Test an dataset that maps a TF function across its input elements."""
+    """Test a dataset that maps a TF function across its input elements."""
 
     # The pipeline is TensorSliceDataset -> ParallelMapDataset(square_3) ->
     # RepeatDataset(count).
diff --git a/tensorflow/python/distribute/coordinator/cluster_coordinator.py b/tensorflow/python/distribute/coordinator/cluster_coordinator.py
index 7a8bcd74a47..8cf7f6bbf2c 100644
--- a/tensorflow/python/distribute/coordinator/cluster_coordinator.py
+++ b/tensorflow/python/distribute/coordinator/cluster_coordinator.py
@@ -365,7 +365,7 @@ class _CoordinatedClosureQueue(object):
 
     # The following is a lock to make sure when `wait` is called and before it
     # returns no `put` can be executed during this period. It is because `wait`
-    # won't know what to do with newly put closures. This lock adds an cutoff
+    # won't know what to do with newly put closures. This lock adds a cutoff
     # for `wait` so that closures put into the queue while waiting would not be
     # taken responsible by this `wait`.
     #
diff --git a/tensorflow/python/eager/polymorphic_function/function_type_utils.py b/tensorflow/python/eager/polymorphic_function/function_type_utils.py
index 78528679a9b..18589aed15e 100644
--- a/tensorflow/python/eager/polymorphic_function/function_type_utils.py
+++ b/tensorflow/python/eager/polymorphic_function/function_type_utils.py
@@ -391,7 +391,7 @@ def canonicalize_function_inputs(
   Python function's positional (named) arguments and a dictionary
   corresponding to its kwargs.  Missing default arguments are added.
 
-  If the FunctionType has an type constraints, then they are used to convert
+  If the FunctionType has a type constraints, then they are used to convert
   arguments to tensors; otherwise, any inputs containing numpy arrays are
   converted to tensors.
 
diff --git a/tensorflow/python/framework/tensor_util.py b/tensorflow/python/framework/tensor_util.py
index de626578348..840d8d2bfdc 100644
--- a/tensorflow/python/framework/tensor_util.py
+++ b/tensorflow/python/framework/tensor_util.py
@@ -1095,7 +1095,7 @@ def constant_value(tensor, partial=False):  # pylint: disable=invalid-name
     or None if it cannot be calculated.
 
   Raises:
-    TypeError: if tensor is not an tensor.Tensor.
+    TypeError: if tensor is not a tensor.Tensor.
   """
   if isinstance(tensor, core.Value):
     try:
diff --git a/tensorflow/python/keras/callbacks.py b/tensorflow/python/keras/callbacks.py
index edf7e11ded5..112f7045821 100644
--- a/tensorflow/python/keras/callbacks.py
+++ b/tensorflow/python/keras/callbacks.py
@@ -1596,7 +1596,7 @@ class BackupAndRestore(Callback):
   beginning of a new model.fit() run.
   Note that user is responsible to bring jobs back up.
   This callback is important for the backup and restore mechanism for fault
-  tolerance purpose. And the model to be restored from an previous checkpoint is
+  tolerance purpose. And the model to be restored from a previous checkpoint is
   expected to be the same as the one used to back up. If user changes arguments
   passed to compile or fit, the checkpoint saved for fault tolerance can become
   invalid.
diff --git a/tensorflow/python/ops/numpy_ops/tests/np_test.py b/tensorflow/python/ops/numpy_ops/tests/np_test.py
index 37c869db58f..b125238ff91 100644
--- a/tensorflow/python/ops/numpy_ops/tests/np_test.py
+++ b/tensorflow/python/ops/numpy_ops/tests/np_test.py
@@ -427,7 +427,7 @@ def _promote_like_lnp(fun, inexact=False):
   """Decorator that promotes the arguments of `fun` to `tnp.result_type(*args)`.
 
   tnp and onp have different type promotion semantics; this decorator allows
-  tests make an onp reference implementation act more like an tnp
+  tests make an onp reference implementation act more like a tnp
   implementation.
   """
   def wrapper(*args, **kw):
diff --git a/tensorflow/python/tpu/tpu_embedding_v2.py b/tensorflow/python/tpu/tpu_embedding_v2.py
index 0aba2b31b9c..bc5494873f5 100644
--- a/tensorflow/python/tpu/tpu_embedding_v2.py
+++ b/tensorflow/python/tpu/tpu_embedding_v2.py
@@ -618,7 +618,7 @@ class TPUEmbedding(autotrackable.AutoTrackable):
     """Applies the gradient update to the embedding tables.
 
     If a gradient of `None` is passed in any position of the nested structure,
-    then an gradient update with a zero gradient is applied for that feature.
+    then a gradient update with a zero gradient is applied for that feature.
     For optimizers like SGD or Adagrad, this is the same as applying no update
     at all. For lazy Adam and other sparsely applied optimizers with decay,
     ensure you understand the effect of applying a zero gradient.
diff --git a/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler.cc b/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler.cc
index fa0e3d85c69..109fde4b1a5 100644
--- a/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler.cc
@@ -1193,7 +1193,7 @@ class WindowedEinsumVisitor : public DfsHloRewriteVisitor {
     return result;
   }
 
-  // Rewrites an gemm+all-to-all into multiple independent partial gemm+a2a's
+  // Rewrites a gemm+all-to-all into multiple independent partial gemm+a2a's
   // to minimize communication overhead. To do this, the original input will be
   // sliced into replica_group size and perform gemm+all-to-all.
   absl::Status HandleAllToAll(HloInstruction* inst) override {
diff --git a/third_party/xla/xla/hlo/builder/value_inference.cc b/third_party/xla/xla/hlo/builder/value_inference.cc
index 72231b19e53..91e693d3870 100644
--- a/third_party/xla/xla/hlo/builder/value_inference.cc
+++ b/third_party/xla/xla/hlo/builder/value_inference.cc
@@ -757,7 +757,7 @@ absl::StatusOr<PostorderDFSNode> PostorderDFSVisitor::AnalyzeUpperBound(
             if (opcode == HloOpcode::kDivide &&
                 this->IsValueEffectiveInteger(root->operand_ids(1))) {
               // Because in many cases the lower bound of a value is
-              // integer 0, instead of throwing an divide-by-zero error
+              // integer 0, instead of throwing a divide-by-zero error
               // at compile time, we set the bound defer the check to
               // runtime. In those cases we use the upper-bound of
               // first operand as a placeholder.
diff --git a/third_party/xla/xla/hlo/builder/xla_builder.h b/third_party/xla/xla/hlo/builder/xla_builder.h
index 065be7eb277..c1d45230a00 100644
--- a/third_party/xla/xla/hlo/builder/xla_builder.h
+++ b/third_party/xla/xla/hlo/builder/xla_builder.h
@@ -3225,7 +3225,7 @@ XlaOp CollectiveBroadcastWithDeviceList(
     const std::optional<ChannelHandle>& channel_id = std::nullopt,
     bool has_dynamic_root = false);
 
-// Enqueues an collective operation that sends and receives data cross replicas.
+// Enqueues a collective operation that sends and receives data cross replicas.
 //
 // - `source_target_pair`: a list of (source_replica_id, target_replica_id)
 // pairs. For each pair, the operand is sent from source replica to target
diff --git a/third_party/xla/xla/hlo/utils/hlo_sharding_util.h b/third_party/xla/xla/hlo/utils/hlo_sharding_util.h
index e0c3869be13..61e43fb4532 100644
--- a/third_party/xla/xla/hlo/utils/hlo_sharding_util.h
+++ b/third_party/xla/xla/hlo/utils/hlo_sharding_util.h
@@ -314,7 +314,7 @@ HloSharding PropagateShardingAlongDimsAndReplicateOthers(
     const HloSharding& source_sharding, absl::Span<const int64_t> source_dims,
     absl::Span<const int64_t> target_dims, int64_t target_shape_rank);
 
-// Returns the iota dimension if maybe_iota is an kIota instruction or
+// Returns the iota dimension if maybe_iota is a kIota instruction or
 // equivalent to kIota.
 std::optional<int64_t> GetDimensionForIota(const HloInstruction* maybe_iota,
                                            const CallGraph& call_graph);
diff --git a/third_party/xla/xla/service/call_graph.h b/third_party/xla/xla/service/call_graph.h
index b377c4f4378..ebddb4b6508 100644
--- a/third_party/xla/xla/service/call_graph.h
+++ b/third_party/xla/xla/service/call_graph.h
@@ -224,7 +224,7 @@ class CallGraph {
   }
 
   // Returns the nearest call graph ancestors of instructions 'a' and 'b' for
-  // which the ancestors are in the same computation. An instruction is an call
+  // which the ancestors are in the same computation. An instruction is a call
   // graph ancestor of 'a' if the instruction calls the computation containing
   // 'a' either directly or transitively. Degeneratively an instruction is an
   // ancestor of itself. nullptr is returned if there is no common ancestor or
diff --git a/third_party/xla/xla/service/gpu/triton_tiling_propagation.h b/third_party/xla/xla/service/gpu/triton_tiling_propagation.h
index 730ab2ebd55..2a3739c977b 100644
--- a/third_party/xla/xla/service/gpu/triton_tiling_propagation.h
+++ b/third_party/xla/xla/service/gpu/triton_tiling_propagation.h
@@ -169,7 +169,7 @@ class DimensionOrder {
     std::string ToString() const;
     std::string ToLongString() const;
 
-    // Label carrying the dimension number of an defining operation.
+    // Label carrying the dimension number of a defining operation.
     int dst_dim_number() const { return dst_dim_number_; }
     // Total number of elements in the fragment ignoring slicing.
     int64_t full_count() const { return count_; }
diff --git a/third_party/xla/xla/service/layout_assignment.cc b/third_party/xla/xla/service/layout_assignment.cc
index bcc343b0bab..778e3911bc4 100644
--- a/third_party/xla/xla/service/layout_assignment.cc
+++ b/third_party/xla/xla/service/layout_assignment.cc
@@ -1872,7 +1872,7 @@ std::unique_ptr<Layout> LayoutAssignment::ChooseOutputLayoutFromOperandLayout(
   }
 
   if (user->opcode() == HloOpcode::kReshape) {
-    // Prefer the user layout that makes the reshape an bitcast. If any
+    // Prefer the user layout that makes the reshape a bitcast. If any
     // dimension bound is 1 in the user shape, there may be several such
     // layouts. So if 'operand_layout' is the default layout, try if the
     // reshape is a bitcast when using the same layout. This may avoid copy
diff --git a/third_party/xla/xla/service/llvm_ir/loop_emitter.h b/third_party/xla/xla/service/llvm_ir/loop_emitter.h
index 4d825058f20..a4c2cf4e99f 100644
--- a/third_party/xla/xla/service/llvm_ir/loop_emitter.h
+++ b/third_party/xla/xla/service/llvm_ir/loop_emitter.h
@@ -53,7 +53,7 @@ class LoopEmitter {
   LoopEmitter(const BodyEmitter& body_emitter, const Shape& shape,
               llvm::IRBuilderBase* b);
 
-  // Constructs a LoopEmitter from an body_emitter that generates
+  // Constructs a LoopEmitter from a body_emitter that generates
   // element of the given target array in the dynamic dimension.
   LoopEmitter(const BodyEmitter& body_emitter, const Shape& shape,
               std::vector<llvm::Value*> dynamic_dims, llvm::IRBuilderBase* b);
diff --git a/third_party/xla/xla/service/while_loop_concat_code_motion.cc b/third_party/xla/xla/service/while_loop_concat_code_motion.cc
index 8d655dbe7ce..4fd024447cd 100644
--- a/third_party/xla/xla/service/while_loop_concat_code_motion.cc
+++ b/third_party/xla/xla/service/while_loop_concat_code_motion.cc
@@ -396,7 +396,7 @@ void ModifyHloPropertiesForConcatShape(const ConcatGroup& group,
     int64_t operand_concat_dim = operand_dim->first;
     bool operand_inserted_concat_dim = operand_dim->second;
     if (operand_inserted_concat_dim) {
-      // We should have added an dimension on the operand.
+      // We should have added a dimension on the operand.
       CHECK_EQ(hlo->operand(0)->shape().dimensions().size(),
                hlo->dimensions().size() + 1)
           << hlo->ToString();
diff --git a/third_party/xla/xla/stream_executor/blas.h b/third_party/xla/xla/stream_executor/blas.h
index 0e19e0391a1..6887c9cd5d7 100644
--- a/third_party/xla/xla/stream_executor/blas.h
+++ b/third_party/xla/xla/stream_executor/blas.h
@@ -314,7 +314,7 @@ class BlasSupport {
       const void *alpha, const void *beta,
       std::vector<blas::AlgorithmType> *out_algorithms) = 0;
 
-  // Like DoBlasGemm, but accepts an algorithm and an compute type.
+  // Like DoBlasGemm, but accepts an algorithm and a compute type.
   //
   // The compute type lets you say (e.g.) that the inputs and outputs are
   // Eigen::halfs, but you want the internal computations to be done with
diff --git a/third_party/xla/xla/tools/hlo_extractor.cc b/third_party/xla/xla/tools/hlo_extractor.cc
index f08a2234439..4377a7f61d6 100644
--- a/third_party/xla/xla/tools/hlo_extractor.cc
+++ b/third_party/xla/xla/tools/hlo_extractor.cc
@@ -54,7 +54,7 @@ limitations under the License.
 namespace xla {
 namespace {
 
-// Visitor that build a new HLO module with an given root HLO instruction. Only
+// Visitor that build a new HLO module with a given root HLO instruction. Only
 // HLOs that are reachable from the new root instruction are included in the new
 // module.
 //

From f608a59dbc15f46252b8888d0be385cd9994f771 Mon Sep 17 00:00:00 2001
From: Arpan Chatterjee <arpanchatterjee648@gmail.com>
Date: Fri, 4 Sep 2026 23:50:24 +0530
Subject: [PATCH 0054/1000] Correct a/an usage

---
 .../experimental/acceleration/mini_benchmark/libjpeg_decoder.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder.cc b/tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder.cc
index 42cd0b639d7..b667746a15f 100644
--- a/tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder.cc
+++ b/tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder.cc
@@ -235,7 +235,7 @@ Status LibjpegDecoder::DecodeImage(const tflite::StringRef& encoded,
 
   // Decoding the data in a buffer as large as the largest allowed JPEG
   // image row to avoid overflows in case  we are reading a wrong value for the
-  // image size in the header or we are receiving an image with an header
+  // image size in the header or we are receiving an image with a header
   // deliberately incorrect to cause a buffer overflow.
   // Using 4 channels because the decode color process can handle 3 or 4
   // channels:

From 8f2c715a56653def0c6ffa8a431bb9d46eb76b5d Mon Sep 17 00:00:00 2001
From: Sai Ramineni <csramineni@gmail.com>
Date: Sat, 5 Sep 2026 13:39:12 -0400
Subject: [PATCH 0055/1000] Skip dense suffix expansion for empty ragged values

---
 .../kernels/ragged_tensor_to_sparse_kernel.cc |  6 ++++--
 .../ragged_tensor_to_sparse_kernel_test.cc    | 19 +++++++++++++++++++
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel.cc b/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel.cc
index 32a605ccf9f..dcb8c65e67c 100644
--- a/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel.cc
+++ b/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel.cc
@@ -60,8 +60,10 @@ class RaggedTensorToSparseOp : public OpKernel {
     // - `index_middle` is the index in the last ragged dimension.
     // - `index_suffix` is the index in the dense value dimensions.
     std::vector<int64_t> index_prefix(rt_nested_splits_len);
-    std::vector<std::vector<int64_t>> index_suffixes =
-        MakeIndexSuffixes(rt_dense_values_in.shape());
+    std::vector<std::vector<int64_t>> index_suffixes;
+    if (rt_dense_values_in.NumElements() != 0) {
+      index_suffixes = MakeIndexSuffixes(rt_dense_values_in.shape());
+    }
 
     // Allocate the `sparse_indices` output tensor.
     const int64_t nvals =
diff --git a/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel_test.cc b/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel_test.cc
index a11cee61116..a391784d23b 100644
--- a/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel_test.cc
+++ b/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel_test.cc
@@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <limits>
+
 #include "tensorflow/core/framework/fake_input.h"
 #include "tensorflow/core/framework/node_def_builder.h"
 #include "tensorflow/core/framework/shape_inference.h"
@@ -88,6 +90,23 @@ TEST_F(RaggedTensorToSparseTest, EmptyRows) {
                                    test::AsTensor<int64_t>({5, 4}));
 }
 
+TEST_F(RaggedTensorToSparseTest, EmptyValuesWithLargeDenseSuffix) {
+  constexpr int64_t kSuffixSize = std::numeric_limits<int32_t>::max();
+  BuildRaggedTensorToSparseGraph<int>(
+      {{0}},                           // splits
+      TensorShape({0, kSuffixSize}),  // values.shape
+      {});                            // values
+  TF_ASSERT_OK(RunOpKernel());
+  test::ExpectTensorEqual<int64_t>(
+      *GetOutput(kSparseIndicesOutput),
+      test::AsTensor<int64_t>({}, TensorShape({0, 3})));
+  test::ExpectTensorEqual<int>(*GetOutput(kSparseValuesOutput),
+                               test::AsTensor<int>({}));
+  test::ExpectTensorEqual<int64_t>(
+      *GetOutput(kSparseDenseShapeOutput),
+      test::AsTensor<int64_t>({0, 0, kSuffixSize}));
+}
+
 TEST_F(RaggedTensorToSparseTest, OneSplits_Values2D) {
   // ragged_tensor=[[[1, 2], [3, 4], [5, 6]], [], [[7, 8], [9, 10]], [[11, 12]]]
   BuildRaggedTensorToSparseGraph<int>(

From c261acd94f349643987dc3898ee0ca57dd6874c9 Mon Sep 17 00:00:00 2001
From: Arpan Chatterjee <arpanchatterjee648@gmail.com>
Date: Tue, 8 Sep 2026 15:02:27 +0530
Subject: [PATCH 0056/1000] Correct grammar related to a/an, revert changes of
 two files

---
 tensorflow/python/ops/numpy_ops/tests/np_test.py | 2 +-
 tensorflow/python/tpu/tpu_embedding_v2.py        | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/tests/np_test.py b/tensorflow/python/ops/numpy_ops/tests/np_test.py
index b125238ff91..37c869db58f 100644
--- a/tensorflow/python/ops/numpy_ops/tests/np_test.py
+++ b/tensorflow/python/ops/numpy_ops/tests/np_test.py
@@ -427,7 +427,7 @@ def _promote_like_lnp(fun, inexact=False):
   """Decorator that promotes the arguments of `fun` to `tnp.result_type(*args)`.
 
   tnp and onp have different type promotion semantics; this decorator allows
-  tests make an onp reference implementation act more like a tnp
+  tests make an onp reference implementation act more like an tnp
   implementation.
   """
   def wrapper(*args, **kw):
diff --git a/tensorflow/python/tpu/tpu_embedding_v2.py b/tensorflow/python/tpu/tpu_embedding_v2.py
index bc5494873f5..0aba2b31b9c 100644
--- a/tensorflow/python/tpu/tpu_embedding_v2.py
+++ b/tensorflow/python/tpu/tpu_embedding_v2.py
@@ -618,7 +618,7 @@ class TPUEmbedding(autotrackable.AutoTrackable):
     """Applies the gradient update to the embedding tables.
 
     If a gradient of `None` is passed in any position of the nested structure,
-    then a gradient update with a zero gradient is applied for that feature.
+    then an gradient update with a zero gradient is applied for that feature.
     For optimizers like SGD or Adagrad, this is the same as applying no update
     at all. For lazy Adam and other sparsely applied optimizers with decay,
     ensure you understand the effect of applying a zero gradient.

From 99456a386ceb044e207bd6d350a978e7d276151a Mon Sep 17 00:00:00 2001
From: Arpan Chatterjee <arpanchatterjee648@gmail.com>
Date: Thu, 10 Sep 2026 02:05:11 +0530
Subject: [PATCH 0057/1000] corrected minor issues

---
 tensorflow/python/ops/numpy_ops/tests/np_test.py | 4 ++--
 tensorflow/python/tpu/tpu_embedding_v2.py        | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/tests/np_test.py b/tensorflow/python/ops/numpy_ops/tests/np_test.py
index 37c869db58f..396dc3fae8e 100644
--- a/tensorflow/python/ops/numpy_ops/tests/np_test.py
+++ b/tensorflow/python/ops/numpy_ops/tests/np_test.py
@@ -427,7 +427,7 @@ def _promote_like_lnp(fun, inexact=False):
   """Decorator that promotes the arguments of `fun` to `tnp.result_type(*args)`.
 
   tnp and onp have different type promotion semantics; this decorator allows
-  tests make an onp reference implementation act more like an tnp
+  tests make an onp reference implementation act more like a tnp
   implementation.
   """
   def wrapper(*args, **kw):
@@ -3163,4 +3163,4 @@ class NumpyGradTests(jtu.TestCase):
 
 if __name__ == "__main__":
   absltest.main()
-  
\ No newline at end of file
+  
diff --git a/tensorflow/python/tpu/tpu_embedding_v2.py b/tensorflow/python/tpu/tpu_embedding_v2.py
index 0aba2b31b9c..bc5494873f5 100644
--- a/tensorflow/python/tpu/tpu_embedding_v2.py
+++ b/tensorflow/python/tpu/tpu_embedding_v2.py
@@ -618,7 +618,7 @@ class TPUEmbedding(autotrackable.AutoTrackable):
     """Applies the gradient update to the embedding tables.
 
     If a gradient of `None` is passed in any position of the nested structure,
-    then an gradient update with a zero gradient is applied for that feature.
+    then a gradient update with a zero gradient is applied for that feature.
     For optimizers like SGD or Adagrad, this is the same as applying no update
     at all. For lazy Adam and other sparsely applied optimizers with decay,
     ensure you understand the effect of applying a zero gradient.

From b5e29252a4a0c543c7066d2a5494dd3670dfd3f7 Mon Sep 17 00:00:00 2001
From: Rudranil Goswami <nilrudra068@gmail.com>
Date: Sat, 12 Sep 2026 02:02:31 +0530
Subject: [PATCH 0058/1000] Add copyright notice and update dependencies in
 BUILD

---
 tensorflow/python/autograph/operators/BUILD | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/tensorflow/python/autograph/operators/BUILD b/tensorflow/python/autograph/operators/BUILD
index d59b6c13926..1b1d4bcc328 100644
--- a/tensorflow/python/autograph/operators/BUILD
+++ b/tensorflow/python/autograph/operators/BUILD
@@ -1,3 +1,18 @@
+# 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.
+# ==============================================================================
+
 load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
 load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
 load("//tensorflow:tensorflow.bzl", "py_test")
@@ -103,6 +118,8 @@ py_library(
         ":control_flow",
         "//tensorflow/python/autograph/utils:tensors",
         "//tensorflow/python/ops:cond",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:variable_utils",
     ],
 )
 

From 6a8034f1592fc7ddde9930b931a43c122290965f Mon Sep 17 00:00:00 2001
From: Arpan Chatterjee <arpanchatterjee648@gmail.com>
Date: Tue, 15 Sep 2026 21:52:10 +0530
Subject: [PATCH 0059/1000] revert changes of files tensorflow/lite/*,
 tensorflow/python/keras/*, third_party/xla

---
 tensorflow/lite/core/subgraph.h                               | 4 ++--
 tensorflow/lite/delegates/flex/test_util.h                    | 2 +-
 .../delegates/gpu/common/testing/feature_parity/README.md     | 2 +-
 .../acceleration/mini_benchmark/libjpeg_decoder.cc            | 2 +-
 .../org/tensorflow/lite/task/core/BaseOptions.Builder.html    | 2 +-
 tensorflow/lite/g3doc/examples/style_transfer/overview.ipynb  | 2 +-
 tensorflow/lite/java/ovic/README.md                           | 2 +-
 .../internal/optimized/depthwiseconv_uint8_3x3_filter.h       | 2 +-
 tensorflow/lite/nnapi/nnapi_implementation.cc                 | 2 +-
 .../fuse_binary_into_following_affine.cc                      | 2 +-
 .../fuse_binary_into_preceding_affine.cc                      | 2 +-
 tensorflow/lite/tools/optimize/quantization_wrapper.h         | 2 +-
 tensorflow/python/keras/callbacks.py                          | 2 +-
 .../xla/backends/gpu/transforms/windowed_einsum_handler.cc    | 2 +-
 third_party/xla/xla/hlo/builder/value_inference.cc            | 2 +-
 third_party/xla/xla/hlo/builder/xla_builder.h                 | 2 +-
 third_party/xla/xla/hlo/utils/hlo_sharding_util.h             | 2 +-
 third_party/xla/xla/service/call_graph.h                      | 2 +-
 third_party/xla/xla/service/gpu/triton_tiling_propagation.h   | 2 +-
 third_party/xla/xla/service/layout_assignment.cc              | 2 +-
 third_party/xla/xla/service/llvm_ir/loop_emitter.h            | 2 +-
 third_party/xla/xla/service/while_loop_concat_code_motion.cc  | 2 +-
 third_party/xla/xla/stream_executor/blas.h                    | 2 +-
 third_party/xla/xla/tools/hlo_extractor.cc                    | 2 +-
 24 files changed, 25 insertions(+), 25 deletions(-)

diff --git a/tensorflow/lite/core/subgraph.h b/tensorflow/lite/core/subgraph.h
index e19aa3919ac..e1f4f0e3c44 100644
--- a/tensorflow/lite/core/subgraph.h
+++ b/tensorflow/lite/core/subgraph.h
@@ -818,14 +818,14 @@ class Subgraph {
   // Does not report invoke status through profiler.
   TfLiteStatus ModifyGraphWithDelegateImpl(TfLiteDelegate* delegate);
 
-  // Request a tensor be resized implementation. If the given tensor is of
+  // Request an tensor be resized implementation. If the given tensor is of
   // type kTfLiteDynamic it will also be allocated new memory.
   TfLiteStatus ResizeTensorImpl(TfLiteTensor* tensor, TfLiteIntArray* new_size);
 
   // Report a detailed error string (will be printed to stderr).
   void ReportErrorImpl(const char* format, va_list args);
 
-  // Entry point for C node plugin API to request a tensor be resized.
+  // Entry point for C node plugin API to request an tensor be resized.
   static TfLiteStatus ResizeTensor(TfLiteContext* context, TfLiteTensor* tensor,
                                    TfLiteIntArray* new_size);
   // Entry point for C node plugin API to report an error.
diff --git a/tensorflow/lite/delegates/flex/test_util.h b/tensorflow/lite/delegates/flex/test_util.h
index 333c4251470..71eeead5640 100644
--- a/tensorflow/lite/delegates/flex/test_util.h
+++ b/tensorflow/lite/delegates/flex/test_util.h
@@ -33,7 +33,7 @@ enum TfOpType {
   kLoopCond,
   // Represents an op that does not exist in TensorFlow.
   kNonExistent,
-  // Represents a valid TensorFlow op where the NodeDef is incompatible.
+  // Represents an valid TensorFlow op where the NodeDef is incompatible.
   kIncompatibleNodeDef,
 };
 
diff --git a/tensorflow/lite/delegates/gpu/common/testing/feature_parity/README.md b/tensorflow/lite/delegates/gpu/common/testing/feature_parity/README.md
index 3c061701b8e..f81f585963b 100644
--- a/tensorflow/lite/delegates/gpu/common/testing/feature_parity/README.md
+++ b/tensorflow/lite/delegates/gpu/common/testing/feature_parity/README.md
@@ -2,6 +2,6 @@
 
 These tests ensure feature parity across TFLite GPU delegates. Every test
 receives a simple automatically generated tflite model as an input. Model runs
-with default tflite cpu interpreter and a delegated interpreter. Test succeeds
+with default tflite cpu interpreter and an delegated interpreter. Test succeeds
 when output results for both interpretes match with the given accuracy.
 
diff --git a/tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder.cc b/tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder.cc
index b667746a15f..42cd0b639d7 100644
--- a/tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder.cc
+++ b/tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder.cc
@@ -235,7 +235,7 @@ Status LibjpegDecoder::DecodeImage(const tflite::StringRef& encoded,
 
   // Decoding the data in a buffer as large as the largest allowed JPEG
   // image row to avoid overflows in case  we are reading a wrong value for the
-  // image size in the header or we are receiving an image with a header
+  // image size in the header or we are receiving an image with an header
   // deliberately incorrect to cause a buffer overflow.
   // Using 4 channels because the decode color process can handle 3 or 4
   // channels:
diff --git a/tensorflow/lite/g3doc/api_docs/java/org/tensorflow/lite/task/core/BaseOptions.Builder.html b/tensorflow/lite/g3doc/api_docs/java/org/tensorflow/lite/task/core/BaseOptions.Builder.html
index 4f467c96899..2c46fdfb08d 100644
--- a/tensorflow/lite/g3doc/api_docs/java/org/tensorflow/lite/task/core/BaseOptions.Builder.html
+++ b/tensorflow/lite/g3doc/api_docs/java/org/tensorflow/lite/task/core/BaseOptions.Builder.html
@@ -363,7 +363,7 @@
 <div class="jd-details-descr">
 <div class="jd-tagdata jd-tagdescr"><p>Sets the advanced accelerator options.
 
- <p>Note: this method will override those highlevel API to choose a delegate, such as <code><a href="/lite/api_docs/java/org/tensorflow/lite/task/core/BaseOptions.Builder.html#useGpu()">useGpu()</a></code> and <code><a href="/lite/api_docs/java/org/tensorflow/lite/task/core/BaseOptions.Builder.html#useNnapi()">useNnapi()</a></code>.
+ <p>Note: this method will override those highlevel API to choose an delegate, such as <code><a href="/lite/api_docs/java/org/tensorflow/lite/task/core/BaseOptions.Builder.html#useGpu()">useGpu()</a></code> and <code><a href="/lite/api_docs/java/org/tensorflow/lite/task/core/BaseOptions.Builder.html#useNnapi()">useNnapi()</a></code>.
 </p></p></div>
 <div class="jd-tagdata">
 <h5 class="jd-tagtitle">Parameters</h5>
diff --git a/tensorflow/lite/g3doc/examples/style_transfer/overview.ipynb b/tensorflow/lite/g3doc/examples/style_transfer/overview.ipynb
index 1bfde6b15af..87039d19a69 100644
--- a/tensorflow/lite/g3doc/examples/style_transfer/overview.ipynb
+++ b/tensorflow/lite/g3doc/examples/style_transfer/overview.ipynb
@@ -223,7 +223,7 @@
         "\n",
         "  return img\n",
         "\n",
-        "# Function to pre-process by resizing a central cropping it.\n",
+        "# Function to pre-process by resizing an central cropping it.\n",
         "def preprocess_image(image, target_dim):\n",
         "  # Resize the image so that the shorter dimension becomes 256px.\n",
         "  shape = tf.cast(tf.shape(image)[1:-1], tf.float32)\n",
diff --git a/tensorflow/lite/java/ovic/README.md b/tensorflow/lite/java/ovic/README.md
index ba5aacf4eb1..62d8183c23d 100644
--- a/tensorflow/lite/java/ovic/README.md
+++ b/tensorflow/lite/java/ovic/README.md
@@ -18,7 +18,7 @@ Note: for now the tests only provides correctness checks, i.e. classifier
 predicts the correct category on the test image, but no on-device latency
 measurements. To test the latency measurement functionality, the tests will
 print the latency running on a desktop computer, which is not indicative of the
-on-device run-time. We are releasing a benchmarker Apk that would allow
+on-device run-time. We are releasing an benchmarker Apk that would allow
 developers to measure latency on their own devices.
 
 ### Obtain the sample models
diff --git a/tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h b/tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h
index d979341cfb8..633feb53528 100644
--- a/tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h
+++ b/tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h
@@ -12934,7 +12934,7 @@ inline void DepthwiseConvDotProduct3x3Impl(
   // Note that this data is adjusted in a per-depth process before the main
   // filters. The adjustment accounts for a non-symmetric input offset.
   //
-  // Kernel subroutines need to be able to operate consistently on a bias
+  // Kernel subroutines need to be able to operate consistently on an bias
   // array. Where there is no bias, we provide one filled with zeros.
   constexpr int kMinBiasLoad = 8;
   int32_t zero_bias_data[kMinBiasLoad];
diff --git a/tensorflow/lite/nnapi/nnapi_implementation.cc b/tensorflow/lite/nnapi/nnapi_implementation.cc
index 75ca0766399..5705576e894 100644
--- a/tensorflow/lite/nnapi/nnapi_implementation.cc
+++ b/tensorflow/lite/nnapi/nnapi_implementation.cc
@@ -200,7 +200,7 @@ const NnApi LoadNnApi() {
   static const char nnapi_library_name[] = "libneuralnetworks.so";
   libneuralnetworks = dlopen(nnapi_library_name, RTLD_LAZY | RTLD_LOCAL);
 #ifdef __ANDROID__
-  // Note: If there is a problem trying to open the NNAPI library on a
+  // Note: If there is an problem trying to open the NNAPI library on a
   // non-Android system, the error message is suppressed. This is to avoid
   // showing confusing errors when running in environments that do not support
   // NNAPI. As more platforms support NNAPI, the #ifdef logic above can be
diff --git a/tensorflow/lite/toco/graph_transformations/fuse_binary_into_following_affine.cc b/tensorflow/lite/toco/graph_transformations/fuse_binary_into_following_affine.cc
index 0d7681875db..813407cc564 100644
--- a/tensorflow/lite/toco/graph_transformations/fuse_binary_into_following_affine.cc
+++ b/tensorflow/lite/toco/graph_transformations/fuse_binary_into_following_affine.cc
@@ -167,7 +167,7 @@ absl::Status FuseBinaryIntoFollowingAffine::Run(Model* model,
 
   CHECK_EQ(binary_op->inputs.size(), 2);
 
-  // We only can fuse a binary when the two operands break down as follows:
+  // We only can fuse an binary when the two operands break down as follows:
   //   1. One operand is the (variable) output of a typical affine (linear plus
   //   bias)
   //      op of a finite list of possible types: at the moment Conv,
diff --git a/tensorflow/lite/toco/graph_transformations/fuse_binary_into_preceding_affine.cc b/tensorflow/lite/toco/graph_transformations/fuse_binary_into_preceding_affine.cc
index d89883c16e0..282eb19e23f 100644
--- a/tensorflow/lite/toco/graph_transformations/fuse_binary_into_preceding_affine.cc
+++ b/tensorflow/lite/toco/graph_transformations/fuse_binary_into_preceding_affine.cc
@@ -222,7 +222,7 @@ absl::Status FuseBinaryIntoPrecedingAffine::Run(Model* model,
 
   CHECK_EQ(binary_op->inputs.size(), 2);
 
-  // We only can fuse a binary when the two operands break down as follows:
+  // We only can fuse an binary when the two operands break down as follows:
   //   1. One operand is the (variable) output of a typical affine (linear plus
   //   bias)
   //      op of a finite list of possible types: at the moment Conv,
diff --git a/tensorflow/lite/tools/optimize/quantization_wrapper.h b/tensorflow/lite/tools/optimize/quantization_wrapper.h
index 4b6852e8286..125bad3e257 100644
--- a/tensorflow/lite/tools/optimize/quantization_wrapper.h
+++ b/tensorflow/lite/tools/optimize/quantization_wrapper.h
@@ -20,7 +20,7 @@ limitations under the License.
 namespace tflite {
 namespace optimize {
 
-// Makes a copy of the model at input_path and writes it to output_path, adding
+// Makes an copy of the model at input_path and writes it to output_path, adding
 // tensors to the model needed for calibration.
 // Returns true if it is successful.
 // Example: a/b/c.tflite becomes a/b/c.calibrated.tflite and has
diff --git a/tensorflow/python/keras/callbacks.py b/tensorflow/python/keras/callbacks.py
index 112f7045821..edf7e11ded5 100644
--- a/tensorflow/python/keras/callbacks.py
+++ b/tensorflow/python/keras/callbacks.py
@@ -1596,7 +1596,7 @@ class BackupAndRestore(Callback):
   beginning of a new model.fit() run.
   Note that user is responsible to bring jobs back up.
   This callback is important for the backup and restore mechanism for fault
-  tolerance purpose. And the model to be restored from a previous checkpoint is
+  tolerance purpose. And the model to be restored from an previous checkpoint is
   expected to be the same as the one used to back up. If user changes arguments
   passed to compile or fit, the checkpoint saved for fault tolerance can become
   invalid.
diff --git a/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler.cc b/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler.cc
index 109fde4b1a5..fa0e3d85c69 100644
--- a/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler.cc
@@ -1193,7 +1193,7 @@ class WindowedEinsumVisitor : public DfsHloRewriteVisitor {
     return result;
   }
 
-  // Rewrites a gemm+all-to-all into multiple independent partial gemm+a2a's
+  // Rewrites an gemm+all-to-all into multiple independent partial gemm+a2a's
   // to minimize communication overhead. To do this, the original input will be
   // sliced into replica_group size and perform gemm+all-to-all.
   absl::Status HandleAllToAll(HloInstruction* inst) override {
diff --git a/third_party/xla/xla/hlo/builder/value_inference.cc b/third_party/xla/xla/hlo/builder/value_inference.cc
index 91e693d3870..72231b19e53 100644
--- a/third_party/xla/xla/hlo/builder/value_inference.cc
+++ b/third_party/xla/xla/hlo/builder/value_inference.cc
@@ -757,7 +757,7 @@ absl::StatusOr<PostorderDFSNode> PostorderDFSVisitor::AnalyzeUpperBound(
             if (opcode == HloOpcode::kDivide &&
                 this->IsValueEffectiveInteger(root->operand_ids(1))) {
               // Because in many cases the lower bound of a value is
-              // integer 0, instead of throwing a divide-by-zero error
+              // integer 0, instead of throwing an divide-by-zero error
               // at compile time, we set the bound defer the check to
               // runtime. In those cases we use the upper-bound of
               // first operand as a placeholder.
diff --git a/third_party/xla/xla/hlo/builder/xla_builder.h b/third_party/xla/xla/hlo/builder/xla_builder.h
index c1d45230a00..065be7eb277 100644
--- a/third_party/xla/xla/hlo/builder/xla_builder.h
+++ b/third_party/xla/xla/hlo/builder/xla_builder.h
@@ -3225,7 +3225,7 @@ XlaOp CollectiveBroadcastWithDeviceList(
     const std::optional<ChannelHandle>& channel_id = std::nullopt,
     bool has_dynamic_root = false);
 
-// Enqueues a collective operation that sends and receives data cross replicas.
+// Enqueues an collective operation that sends and receives data cross replicas.
 //
 // - `source_target_pair`: a list of (source_replica_id, target_replica_id)
 // pairs. For each pair, the operand is sent from source replica to target
diff --git a/third_party/xla/xla/hlo/utils/hlo_sharding_util.h b/third_party/xla/xla/hlo/utils/hlo_sharding_util.h
index 61e43fb4532..e0c3869be13 100644
--- a/third_party/xla/xla/hlo/utils/hlo_sharding_util.h
+++ b/third_party/xla/xla/hlo/utils/hlo_sharding_util.h
@@ -314,7 +314,7 @@ HloSharding PropagateShardingAlongDimsAndReplicateOthers(
     const HloSharding& source_sharding, absl::Span<const int64_t> source_dims,
     absl::Span<const int64_t> target_dims, int64_t target_shape_rank);
 
-// Returns the iota dimension if maybe_iota is a kIota instruction or
+// Returns the iota dimension if maybe_iota is an kIota instruction or
 // equivalent to kIota.
 std::optional<int64_t> GetDimensionForIota(const HloInstruction* maybe_iota,
                                            const CallGraph& call_graph);
diff --git a/third_party/xla/xla/service/call_graph.h b/third_party/xla/xla/service/call_graph.h
index ebddb4b6508..b377c4f4378 100644
--- a/third_party/xla/xla/service/call_graph.h
+++ b/third_party/xla/xla/service/call_graph.h
@@ -224,7 +224,7 @@ class CallGraph {
   }
 
   // Returns the nearest call graph ancestors of instructions 'a' and 'b' for
-  // which the ancestors are in the same computation. An instruction is a call
+  // which the ancestors are in the same computation. An instruction is an call
   // graph ancestor of 'a' if the instruction calls the computation containing
   // 'a' either directly or transitively. Degeneratively an instruction is an
   // ancestor of itself. nullptr is returned if there is no common ancestor or
diff --git a/third_party/xla/xla/service/gpu/triton_tiling_propagation.h b/third_party/xla/xla/service/gpu/triton_tiling_propagation.h
index 2a3739c977b..730ab2ebd55 100644
--- a/third_party/xla/xla/service/gpu/triton_tiling_propagation.h
+++ b/third_party/xla/xla/service/gpu/triton_tiling_propagation.h
@@ -169,7 +169,7 @@ class DimensionOrder {
     std::string ToString() const;
     std::string ToLongString() const;
 
-    // Label carrying the dimension number of a defining operation.
+    // Label carrying the dimension number of an defining operation.
     int dst_dim_number() const { return dst_dim_number_; }
     // Total number of elements in the fragment ignoring slicing.
     int64_t full_count() const { return count_; }
diff --git a/third_party/xla/xla/service/layout_assignment.cc b/third_party/xla/xla/service/layout_assignment.cc
index 778e3911bc4..aa0d5ebdc00 100644
--- a/third_party/xla/xla/service/layout_assignment.cc
+++ b/third_party/xla/xla/service/layout_assignment.cc
@@ -1731,7 +1731,7 @@ std::unique_ptr<Layout> LayoutAssignment::ChooseOperandLayoutFromOutputLayout(
   }
 
   if (instruction->opcode() == HloOpcode::kReshape) {
-    // Prefer the operand layout that makes the reshape a bitcast. If any
+    // Prefer the operand layout that makes the reshape an bitcast. If any
     // dimension bound is 1 in the operand shape, there may be several such
     // layouts. So if 'output_layout' is the default layout, try if the
     // reshape is a bitcast when using the same layout. This may avoid copy
diff --git a/third_party/xla/xla/service/llvm_ir/loop_emitter.h b/third_party/xla/xla/service/llvm_ir/loop_emitter.h
index a4c2cf4e99f..4d825058f20 100644
--- a/third_party/xla/xla/service/llvm_ir/loop_emitter.h
+++ b/third_party/xla/xla/service/llvm_ir/loop_emitter.h
@@ -53,7 +53,7 @@ class LoopEmitter {
   LoopEmitter(const BodyEmitter& body_emitter, const Shape& shape,
               llvm::IRBuilderBase* b);
 
-  // Constructs a LoopEmitter from a body_emitter that generates
+  // Constructs a LoopEmitter from an body_emitter that generates
   // element of the given target array in the dynamic dimension.
   LoopEmitter(const BodyEmitter& body_emitter, const Shape& shape,
               std::vector<llvm::Value*> dynamic_dims, llvm::IRBuilderBase* b);
diff --git a/third_party/xla/xla/service/while_loop_concat_code_motion.cc b/third_party/xla/xla/service/while_loop_concat_code_motion.cc
index 4fd024447cd..8d655dbe7ce 100644
--- a/third_party/xla/xla/service/while_loop_concat_code_motion.cc
+++ b/third_party/xla/xla/service/while_loop_concat_code_motion.cc
@@ -396,7 +396,7 @@ void ModifyHloPropertiesForConcatShape(const ConcatGroup& group,
     int64_t operand_concat_dim = operand_dim->first;
     bool operand_inserted_concat_dim = operand_dim->second;
     if (operand_inserted_concat_dim) {
-      // We should have added a dimension on the operand.
+      // We should have added an dimension on the operand.
       CHECK_EQ(hlo->operand(0)->shape().dimensions().size(),
                hlo->dimensions().size() + 1)
           << hlo->ToString();
diff --git a/third_party/xla/xla/stream_executor/blas.h b/third_party/xla/xla/stream_executor/blas.h
index 6887c9cd5d7..0e19e0391a1 100644
--- a/third_party/xla/xla/stream_executor/blas.h
+++ b/third_party/xla/xla/stream_executor/blas.h
@@ -314,7 +314,7 @@ class BlasSupport {
       const void *alpha, const void *beta,
       std::vector<blas::AlgorithmType> *out_algorithms) = 0;
 
-  // Like DoBlasGemm, but accepts an algorithm and a compute type.
+  // Like DoBlasGemm, but accepts an algorithm and an compute type.
   //
   // The compute type lets you say (e.g.) that the inputs and outputs are
   // Eigen::halfs, but you want the internal computations to be done with
diff --git a/third_party/xla/xla/tools/hlo_extractor.cc b/third_party/xla/xla/tools/hlo_extractor.cc
index 4377a7f61d6..f08a2234439 100644
--- a/third_party/xla/xla/tools/hlo_extractor.cc
+++ b/third_party/xla/xla/tools/hlo_extractor.cc
@@ -54,7 +54,7 @@ limitations under the License.
 namespace xla {
 namespace {
 
-// Visitor that build a new HLO module with a given root HLO instruction. Only
+// Visitor that build a new HLO module with an given root HLO instruction. Only
 // HLOs that are reachable from the new root instruction are included in the new
 // module.
 //

From cb5f5f1107c9c5af97ec9c7ea731f2c1c85bd8d6 Mon Sep 17 00:00:00 2001
From: Arpan Chatterjee <arpanchatterjee648@gmail.com>
Date: Tue, 15 Sep 2026 22:03:42 +0530
Subject: [PATCH 0060/1000] revert changes of files tensorflow/lite/*,
 tensorflow/python/keras/*, third_party/xla

---
 tensorflow/python/ops/numpy_ops/tests/np_test.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/python/ops/numpy_ops/tests/np_test.py b/tensorflow/python/ops/numpy_ops/tests/np_test.py
index 396dc3fae8e..30b5fa9abae 100644
--- a/tensorflow/python/ops/numpy_ops/tests/np_test.py
+++ b/tensorflow/python/ops/numpy_ops/tests/np_test.py
@@ -3163,4 +3163,3 @@ class NumpyGradTests(jtu.TestCase):
 
 if __name__ == "__main__":
   absltest.main()
-  

From bda940d305c99da364f78b3b17505425394eabcf Mon Sep 17 00:00:00 2001
From: Arpan Chatterjee <arpanchatterjee648@gmail.com>
Date: Tue, 15 Sep 2026 22:36:43 +0530
Subject: [PATCH 0061/1000] revert changes of files tensorflow/lite/*,
 tensorflow/python/keras/*, third_party/xla

---
 tensorflow/python/ops/numpy_ops/tests/np_test.py | 2 +-
 tensorflow/python/tpu/tpu_embedding_v2.py        | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/tests/np_test.py b/tensorflow/python/ops/numpy_ops/tests/np_test.py
index 30b5fa9abae..ddd73113864 100644
--- a/tensorflow/python/ops/numpy_ops/tests/np_test.py
+++ b/tensorflow/python/ops/numpy_ops/tests/np_test.py
@@ -427,7 +427,7 @@ def _promote_like_lnp(fun, inexact=False):
   """Decorator that promotes the arguments of `fun` to `tnp.result_type(*args)`.
 
   tnp and onp have different type promotion semantics; this decorator allows
-  tests make an onp reference implementation act more like a tnp
+  tests make an onp reference implementation act more like an tnp
   implementation.
   """
   def wrapper(*args, **kw):
diff --git a/tensorflow/python/tpu/tpu_embedding_v2.py b/tensorflow/python/tpu/tpu_embedding_v2.py
index bc5494873f5..0aba2b31b9c 100644
--- a/tensorflow/python/tpu/tpu_embedding_v2.py
+++ b/tensorflow/python/tpu/tpu_embedding_v2.py
@@ -618,7 +618,7 @@ class TPUEmbedding(autotrackable.AutoTrackable):
     """Applies the gradient update to the embedding tables.
 
     If a gradient of `None` is passed in any position of the nested structure,
-    then a gradient update with a zero gradient is applied for that feature.
+    then an gradient update with a zero gradient is applied for that feature.
     For optimizers like SGD or Adagrad, this is the same as applying no update
     at all. For lazy Adam and other sparsely applied optimizers with decay,
     ensure you understand the effect of applying a zero gradient.

From 69453488d556e40c4d7268103b0cd191ef82b22e Mon Sep 17 00:00:00 2001
From: Arpan Chatterjee <arpanchatterjee648@gmail.com>
Date: Tue, 15 Sep 2026 22:45:27 +0530
Subject: [PATCH 0062/1000] revert changes of files tensorflow/lite/*,
 tensorflow/python/keras/*, third_party/xla

---
 tensorflow/python/ops/numpy_ops/tests/np_test.py | 2 +-
 tensorflow/python/tpu/tpu_embedding_v2.py        | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/tests/np_test.py b/tensorflow/python/ops/numpy_ops/tests/np_test.py
index ddd73113864..30b5fa9abae 100644
--- a/tensorflow/python/ops/numpy_ops/tests/np_test.py
+++ b/tensorflow/python/ops/numpy_ops/tests/np_test.py
@@ -427,7 +427,7 @@ def _promote_like_lnp(fun, inexact=False):
   """Decorator that promotes the arguments of `fun` to `tnp.result_type(*args)`.
 
   tnp and onp have different type promotion semantics; this decorator allows
-  tests make an onp reference implementation act more like an tnp
+  tests make an onp reference implementation act more like a tnp
   implementation.
   """
   def wrapper(*args, **kw):
diff --git a/tensorflow/python/tpu/tpu_embedding_v2.py b/tensorflow/python/tpu/tpu_embedding_v2.py
index 0aba2b31b9c..bc5494873f5 100644
--- a/tensorflow/python/tpu/tpu_embedding_v2.py
+++ b/tensorflow/python/tpu/tpu_embedding_v2.py
@@ -618,7 +618,7 @@ class TPUEmbedding(autotrackable.AutoTrackable):
     """Applies the gradient update to the embedding tables.
 
     If a gradient of `None` is passed in any position of the nested structure,
-    then an gradient update with a zero gradient is applied for that feature.
+    then a gradient update with a zero gradient is applied for that feature.
     For optimizers like SGD or Adagrad, this is the same as applying no update
     at all. For lazy Adam and other sparsely applied optimizers with decay,
     ensure you understand the effect of applying a zero gradient.

From b746e79cb1932d7f8614918ed8e9def0116c3f6d Mon Sep 17 00:00:00 2001
From: Arpan Chatterjee <arpanchatterjee648@gmail.com>
Date: Tue, 15 Sep 2026 23:01:07 +0530
Subject: [PATCH 0063/1000] revert changes of files tensorflow/lite/*,
 tensorflow/python/keras/*, third_party/xla

---
 third_party/xla/xla/service/layout_assignment.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/service/layout_assignment.cc b/third_party/xla/xla/service/layout_assignment.cc
index aa0d5ebdc00..bcc343b0bab 100644
--- a/third_party/xla/xla/service/layout_assignment.cc
+++ b/third_party/xla/xla/service/layout_assignment.cc
@@ -1731,7 +1731,7 @@ std::unique_ptr<Layout> LayoutAssignment::ChooseOperandLayoutFromOutputLayout(
   }
 
   if (instruction->opcode() == HloOpcode::kReshape) {
-    // Prefer the operand layout that makes the reshape an bitcast. If any
+    // Prefer the operand layout that makes the reshape a bitcast. If any
     // dimension bound is 1 in the operand shape, there may be several such
     // layouts. So if 'output_layout' is the default layout, try if the
     // reshape is a bitcast when using the same layout. This may avoid copy
@@ -1872,7 +1872,7 @@ std::unique_ptr<Layout> LayoutAssignment::ChooseOutputLayoutFromOperandLayout(
   }
 
   if (user->opcode() == HloOpcode::kReshape) {
-    // Prefer the user layout that makes the reshape a bitcast. If any
+    // Prefer the user layout that makes the reshape an bitcast. If any
     // dimension bound is 1 in the user shape, there may be several such
     // layouts. So if 'operand_layout' is the default layout, try if the
     // reshape is a bitcast when using the same layout. This may avoid copy

From 9640933c03d542994fedaf1e8a38a0dbbede97b8 Mon Sep 17 00:00:00 2001
From: Aditya Gaurav <mr.x.l.r.8.pride@gmail.com>
Date: Tue, 8 Sep 2026 10:55:06 +0530
Subject: [PATCH 0064/1000] fix(kernels): cast LinSpace indices via double for
 MSVC half/bfloat16

---
 tensorflow/core/kernels/sequence_ops.cc | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/kernels/sequence_ops.cc b/tensorflow/core/kernels/sequence_ops.cc
index 70de3a917ac..daf1931897e 100644
--- a/tensorflow/core/kernels/sequence_ops.cc
+++ b/tensorflow/core/kernels/sequence_ops.cc
@@ -206,9 +206,11 @@ class LinSpaceOp : public OpKernel {
     auto flat = out->flat<T>();
     flat(0) = start;
     if (num > 1) {
-      const T step = (stop - start) / static_cast<T>(num - 1);
+      // Route int->half/bfloat16 casts through double for MSVC.
+      const T step =
+          (stop - start) / static_cast<T>(static_cast<double>(num - 1));
       for (Tnum i = 1; i < num - 1; ++i) {
-        flat(i) = start + step * static_cast<T>(i);
+        flat(i) = start + step * static_cast<T>(static_cast<double>(i));
       }
       // Ensure final value == stop; float arithmetic won't guarantee this.
       flat(num - 1) = stop;

From 75365a8d636f7ec401b59abd3575c5bef90b6814 Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Tue, 15 Sep 2026 23:42:51 -0400
Subject: [PATCH 0065/1000] fuzz(tflite): use public TFLite headers instead of
 core internal paths

Switches the three lite/core includes to the public forwarding headers exported by //tensorflow/lite:framework, per review. Including the core internal package directly caused the dependency layering check to fail. Include block re-sorted. No functional change to the target.
---
 .../security/fuzzing/cc/tflite_interp_overflow_fuzz.cc      | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
index d5b2b315608..4650f2c3853 100644
--- a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
+++ b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
@@ -26,10 +26,10 @@ limitations under the License.
 #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/interpreter.h"
+#include "tensorflow/lite/interpreter_builder.h"
 #include "tensorflow/lite/kernels/register.h"
+#include "tensorflow/lite/model_builder.h"
 
 namespace {
 

From 5002f73e236b5f1cdfdc3d24f409f68879f5f2a7 Mon Sep 17 00:00:00 2001
From: Arpan Chatterjee <arpanchatterjee648@gmail.com>
Date: Wed, 16 Sep 2026 09:47:26 +0530
Subject: [PATCH 0066/1000] revert changes of files tensorflow/lite/*,
 tensorflow/python/keras/*, third_party/xla

---
 .../python/ops/numpy_ops/tests/np_test.py     | 29 ++++++++++++++-----
 tensorflow/python/tpu/tpu_embedding_v2.py     |  4 +--
 2 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/tests/np_test.py b/tensorflow/python/ops/numpy_ops/tests/np_test.py
index 30b5fa9abae..054549a78fe 100644
--- a/tensorflow/python/ops/numpy_ops/tests/np_test.py
+++ b/tensorflow/python/ops/numpy_ops/tests/np_test.py
@@ -26,9 +26,8 @@ from absl.testing import parameterized
 import numpy as onp
 import six
 
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import errors_impl
-from tensorflow.python.framework import ops
-from tensorflow.python.ops.numpy_ops import np_config
 from tensorflow.python.ops.numpy_ops.tests.config import config
 from tensorflow.python.ops.numpy_ops.tests.config import FLAGS
 import tensorflow.python.ops.numpy_ops.tests.extensions as nje
@@ -65,6 +64,7 @@ all_dtypes = number_dtypes + bool_dtypes
 
 python_scalar_dtypes = [tnp.bool_, tnp.int_, tnp.float64, tnp.complex128]
 # pylint: disable=unnecessary-lambda,g-long-lambda,expression-not-assigned
+# pylint: disable=line-too-long,used-before-assignment,function-redefined
 
 def _valid_dtypes_for_shape(shape, dtypes):
   # Not all (shape, dtype) pairs are valid. In particular, Python scalars only
@@ -427,7 +427,7 @@ def _promote_like_lnp(fun, inexact=False):
   """Decorator that promotes the arguments of `fun` to `tnp.result_type(*args)`.
 
   tnp and onp have different type promotion semantics; this decorator allows
-  tests make an onp reference implementation act more like a tnp
+  tests make an onp reference implementation act more like an tnp
   implementation.
   """
   def wrapper(*args, **kw):
@@ -747,6 +747,8 @@ class LaxBackedNumpyTests(jtu.TestCase):
       ]
       for lhs_dtype, rhs_dtype in CombosWithReplacement(
           minus(number_dtypes, complex_dtypes), 2)))
+  @unittest.skipIf(onp.__version__ >= onp.lib.NumpyVersion('2.0.0'),
+                   'tf numpy is implemented to be numpy 1.x compatible')
   def testCross(self, lhs_shape, lhs_dtype, rhs_shape, rhs_dtype, axes, rng_factory):
     rng = rng_factory()
     args_maker = lambda: [rng(lhs_shape, lhs_dtype), rng(rhs_shape, rhs_dtype)]
@@ -930,7 +932,13 @@ class LaxBackedNumpyTests(jtu.TestCase):
     check_xla = not set((lhs_dtype, rhs_dtype)).intersection(
         (onp.int32, onp.int64))
 
-    tol = {onp.float64: 1e-14, onp.float16: 0.04, onp.complex128: 6e-15}
+    tol = {
+        onp.float32: 1e-4,
+        onp.complex64: 1e-4,
+        onp.float64: 1e-14,
+        onp.float16: 0.04,
+        onp.complex128: 6e-15,
+    }
     tol = max(jtu.tolerance(lhs_dtype, tol), jtu.tolerance(rhs_dtype, tol))
     self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=True,
                           check_incomplete_shape=True,
@@ -1785,10 +1793,16 @@ class LaxBackedNumpyTests(jtu.TestCase):
     try:
       self._CheckAgainstNumpy(
           onp_fun, lnp_fun, args_maker, check_dtypes=check_dtypes, tol=tol)
-    except ZeroDivisionError:
+      self._CompileAndCheck(
+          lnp_fun,
+          args_maker,
+          check_dtypes=check_dtypes,
+          rtol=tol,
+          atol=tol,
+          check_incomplete_shape=True,
+      )
+    except (ZeroDivisionError, errors.InvalidArgumentError):
       self.skipTest("don't support checking for ZeroDivisionError")
-    self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=check_dtypes,
-                          rtol=tol, atol=tol, check_incomplete_shape=True)
 
   @named_parameters(jtu.cases_from_list(
       {"testcase_name": "_arg{}_ndmin={}".format(i, ndmin),
@@ -2823,7 +2837,6 @@ class LaxBackedNumpyTests(jtu.TestCase):
                               check_dtypes=False, atol=atol, rtol=tol,
                               check_incomplete_shape=True)
 
-
   @named_parameters(
       jtu.cases_from_list(
           {
diff --git a/tensorflow/python/tpu/tpu_embedding_v2.py b/tensorflow/python/tpu/tpu_embedding_v2.py
index bc5494873f5..f026374514c 100644
--- a/tensorflow/python/tpu/tpu_embedding_v2.py
+++ b/tensorflow/python/tpu/tpu_embedding_v2.py
@@ -618,7 +618,7 @@ class TPUEmbedding(autotrackable.AutoTrackable):
     """Applies the gradient update to the embedding tables.
 
     If a gradient of `None` is passed in any position of the nested structure,
-    then a gradient update with a zero gradient is applied for that feature.
+    then an gradient update with a zero gradient is applied for that feature.
     For optimizers like SGD or Adagrad, this is the same as applying no update
     at all. For lazy Adam and other sparsely applied optimizers with decay,
     ensure you understand the effect of applying a zero gradient.
@@ -1756,7 +1756,7 @@ def make_sharded_variable_creator(
           offset += p
         else:
           kwargs["initial_value"] = functools.partial(
-              unwrapped_initial_value, kwargs["shape"], dtype=dtype)
+              unwrapped_initial_value, kwargs["shape"], dtype=dtype)  # pyrefly: ignore[bad-argument-count, unexpected-keyword]
         variables.append(next_creator(*args, **kwargs))
     return TPUEmbeddingVariable(variables, name=name)
   return sharded_variable_creator

From b765c4b9ec278838b4f4a077c3c3be95ae13e4b5 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Thu, 17 Sep 2026 06:41:19 +0530
Subject: [PATCH 0067/1000] Remove unused XLA flag assignment

---
 tensorflow/python/tools/saved_model_aot_compile.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/python/tools/saved_model_aot_compile.py b/tensorflow/python/tools/saved_model_aot_compile.py
index 257e11b1caa..335979b295e 100644
--- a/tensorflow/python/tools/saved_model_aot_compile.py
+++ b/tensorflow/python/tools/saved_model_aot_compile.py
@@ -367,7 +367,6 @@ def aot_compile_cpu_meta_graph_def(checkpoint_path,
     else:
       xla_flags += ' --xla_cpu_multi_thread_eigen={}'.format(
           'true' if multithreading else 'false')
-    xla_flags += ' --xla_cpu_experimental_ynn_fusion_type= '
     os.environ['XLA_FLAGS'] = xla_flags
 
   temp_dir = test.get_temp_dir()

From 3234279f3f6a5821218d765146c8d28b086816a1 Mon Sep 17 00:00:00 2001
From: daniel <liujingtai2001@gmail.com>
Date: Sat, 19 Sep 2026 02:09:50 -0400
Subject: [PATCH 0068/1000] Fix SiLU gradients for broadcasted tensor beta

---
 tensorflow/python/ops/BUILD      |  1 +
 tensorflow/python/ops/nn_impl.py | 18 +++++++++++++----
 tensorflow/python/ops/nn_test.py | 33 +++++++++++++++++++++++++++++++-
 3 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD
index f72ab18839b..70e69312e1c 100644
--- a/tensorflow/python/ops/BUILD
+++ b/tensorflow/python/ops/BUILD
@@ -4370,6 +4370,7 @@ py_library(
     strict_deps = True,
     deps = [
         ":array_ops",
+        ":array_ops_gen",
         ":array_ops_stack",
         ":candidate_sampling_ops",
         ":cond",
diff --git a/tensorflow/python/ops/nn_impl.py b/tensorflow/python/ops/nn_impl.py
index 8a9d491a4be..ce6d1b3d601 100644
--- a/tensorflow/python/ops/nn_impl.py
+++ b/tensorflow/python/ops/nn_impl.py
@@ -31,6 +31,7 @@ from tensorflow.python.ops import cond as tf_cond
 from tensorflow.python.ops import ctc_ops  # pylint: disable=unused-import
 from tensorflow.python.ops import custom_gradient
 from tensorflow.python.ops import embedding_ops
+from tensorflow.python.ops import gen_array_ops
 from tensorflow.python.ops import gen_nn_ops
 from tensorflow.python.ops import gen_sparse_ops
 from tensorflow.python.ops import linalg_ops
@@ -496,10 +497,19 @@ def swish(features, beta=1.0):
           sigmoid_grad = sigmoid_features * (1.0 - sigmoid_features)
 
       activation_grad = sigmoid_features + logits * sigmoid_grad
-      beta_grad = math_ops.reduce_sum(
-          dy * math_ops.square(features) * sigmoid_grad
-      )
-      return (dy * activation_grad, beta_grad)
+      # Undo broadcasting separately for each input to the custom gradient.
+      features_shape = array_ops.shape(features)
+      beta_shape = array_ops.shape(beta)
+      features_axes, beta_axes = gen_array_ops.broadcast_gradient_args(
+          features_shape, beta_shape)
+      features_grad = array_ops.reshape(
+          math_ops.reduce_sum(dy * activation_grad, features_axes),
+          features_shape)
+      beta_grad = array_ops.reshape(
+          math_ops.reduce_sum(
+              dy * math_ops.square(features) * sigmoid_grad, beta_axes),
+          beta_shape)
+      return (features_grad, beta_grad)
 
     return features * math_ops.sigmoid(beta * features), grad
 
diff --git a/tensorflow/python/ops/nn_test.py b/tensorflow/python/ops/nn_test.py
index ed03a220448..a7a9bebaea7 100644
--- a/tensorflow/python/ops/nn_test.py
+++ b/tensorflow/python/ops/nn_test.py
@@ -1080,7 +1080,7 @@ class GeluTest(test_lib.TestCase):
 
 
 @test_util.run_all_in_graph_and_eager_modes
-class SwishTest(test_lib.TestCase):
+class SwishTest(test_lib.TestCase, parameterized.TestCase):
 
   def testValues(self):
     np_values = np.array(
@@ -1136,6 +1136,37 @@ class SwishTest(test_lib.TestCase):
           f, [x_tf])
       self.assertAllClose(theoretical, numerical)
 
+  @parameterized.product(
+      shapes=[
+          ((2, 3), ()),      # Scalar beta.
+          ((2, 3), (3,)),    # Vector beta.
+          ((2, 3), (1, 3)),  # Singleton dimension in beta.
+          ((2, 3), (2, 3)),  # Matching shapes.
+          ((3,), (2, 3)),    # Broadcast features.
+          ((2, 1), (1, 3)),  # Broadcast both inputs.
+          ((), (3,)),        # Scalar features.
+          ((0, 3), (3,)),    # Empty output.
+      ],
+      dtype=[dtypes.float32, dtypes.float64],
+      dynamic_shapes=[False, True],
+  )
+  def testBroadcastGradients(self, shapes, dtype, dynamic_shapes):
+    features_shape, beta_shape = shapes
+    features = constant_op.constant(
+        np.linspace(-1.0, 1.0, math.prod(features_shape)).reshape(
+            features_shape), dtype=dtype)
+    beta = constant_op.constant(
+        np.linspace(0.5, 1.5, math.prod(beta_shape)).reshape(beta_shape),
+        dtype=dtype)
+    swish = nn_impl.swish
+    if dynamic_shapes:
+      swish = def_function.function(
+          swish, input_signature=[tensor_spec.TensorSpec(None, dtype)] * 2)
+    theoretical, numerical = gradient_checker_v2.compute_gradient(
+        swish, [features, beta])
+    self.assertAllClose(theoretical[0], numerical[0], atol=1e-4)
+    self.assertAllClose(theoretical[1], numerical[1], atol=1e-4)
+
 
 class MomentsTest(test_lib.TestCase):
 

From 33b14f3f00326fffcc130f46682e555c7b1cfb27 Mon Sep 17 00:00:00 2001
From: daniel <liujingtai2001@gmail.com>
Date: Sat, 19 Sep 2026 02:19:33 -0400
Subject: [PATCH 0069/1000] Reuse existing array ops export for Swish gradients

---
 tensorflow/python/ops/BUILD      | 1 -
 tensorflow/python/ops/nn_impl.py | 3 +--
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD
index 70e69312e1c..f72ab18839b 100644
--- a/tensorflow/python/ops/BUILD
+++ b/tensorflow/python/ops/BUILD
@@ -4370,7 +4370,6 @@ py_library(
     strict_deps = True,
     deps = [
         ":array_ops",
-        ":array_ops_gen",
         ":array_ops_stack",
         ":candidate_sampling_ops",
         ":cond",
diff --git a/tensorflow/python/ops/nn_impl.py b/tensorflow/python/ops/nn_impl.py
index ce6d1b3d601..068397d5017 100644
--- a/tensorflow/python/ops/nn_impl.py
+++ b/tensorflow/python/ops/nn_impl.py
@@ -31,7 +31,6 @@ from tensorflow.python.ops import cond as tf_cond
 from tensorflow.python.ops import ctc_ops  # pylint: disable=unused-import
 from tensorflow.python.ops import custom_gradient
 from tensorflow.python.ops import embedding_ops
-from tensorflow.python.ops import gen_array_ops
 from tensorflow.python.ops import gen_nn_ops
 from tensorflow.python.ops import gen_sparse_ops
 from tensorflow.python.ops import linalg_ops
@@ -500,7 +499,7 @@ def swish(features, beta=1.0):
       # Undo broadcasting separately for each input to the custom gradient.
       features_shape = array_ops.shape(features)
       beta_shape = array_ops.shape(beta)
-      features_axes, beta_axes = gen_array_ops.broadcast_gradient_args(
+      features_axes, beta_axes = array_ops.broadcast_gradient_args(
           features_shape, beta_shape)
       features_grad = array_ops.reshape(
           math_ops.reduce_sum(dy * activation_grad, features_axes),

From 8c3272903cbf166c95a1be750a2ed2a5e41a99ad Mon Sep 17 00:00:00 2001
From: daniel <liujingtai2001@gmail.com>
Date: Sat, 19 Sep 2026 02:23:23 -0400
Subject: [PATCH 0070/1000] Clarify Swish derivative and broadcast reduction
 steps

---
 tensorflow/python/ops/nn_impl.py | 21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/tensorflow/python/ops/nn_impl.py b/tensorflow/python/ops/nn_impl.py
index 068397d5017..17fc44dce5a 100644
--- a/tensorflow/python/ops/nn_impl.py
+++ b/tensorflow/python/ops/nn_impl.py
@@ -495,20 +495,21 @@ def swish(features, beta=1.0):
           sigmoid_features = math_ops.sigmoid(logits)
           sigmoid_grad = sigmoid_features * (1.0 - sigmoid_features)
 
-      activation_grad = sigmoid_features + logits * sigmoid_grad
-      # Undo broadcasting separately for each input to the custom gradient.
+      features_grad = dy * (sigmoid_features + logits * sigmoid_grad)
+      beta_grad = dy * math_ops.square(features) * sigmoid_grad
+
+      # Reduce each gradient back to its input shape after broadcasting.
       features_shape = array_ops.shape(features)
       beta_shape = array_ops.shape(beta)
       features_axes, beta_axes = array_ops.broadcast_gradient_args(
           features_shape, beta_shape)
-      features_grad = array_ops.reshape(
-          math_ops.reduce_sum(dy * activation_grad, features_axes),
-          features_shape)
-      beta_grad = array_ops.reshape(
-          math_ops.reduce_sum(
-              dy * math_ops.square(features) * sigmoid_grad, beta_axes),
-          beta_shape)
-      return (features_grad, beta_grad)
+      return (
+          array_ops.reshape(
+              math_ops.reduce_sum(features_grad, features_axes),
+              features_shape),
+          array_ops.reshape(
+              math_ops.reduce_sum(beta_grad, beta_axes), beta_shape),
+      )
 
     return features * math_ops.sigmoid(beta * features), grad
 

From c6cc4427af4941cf89b1e8131aab7741e13dd216 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Sun, 20 Sep 2026 15:46:50 +0530
Subject: [PATCH 0071/1000] Add ScopedTFBuffer tests for deletion behavior

Added tests for ScopedTFBuffer to ensure proper deletion behavior under normal conditions and during interpreter shutdown.
---
 .../python/framework/c_api_util_test.py       | 23 +++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/tensorflow/python/framework/c_api_util_test.py b/tensorflow/python/framework/c_api_util_test.py
index 78e98f00381..ca37c2ecf20 100644
--- a/tensorflow/python/framework/c_api_util_test.py
+++ b/tensorflow/python/framework/c_api_util_test.py
@@ -122,5 +122,28 @@ class UniquePtrTest(test_util.TensorFlowTestCase):
     gc.collect()
     self.assertTrue(self.obj.deleted)
 
+
+class ScopedTFBufferTest(test_util.TensorFlowTestCase):
+
+  def testDeleteRunsUnderNormalConditions(self):
+    buf = c_api_util.ScopedTFBuffer(b"test")
+    del buf
+    gc.collect()
+
+  def testDeleteGuardsAgainstNoneCApi(self):
+    # Regression test: __del__ used to call c_api.TF_DeleteBuffer without
+    # guarding against c_api being None, which happens during interpreter
+    # shutdown. Call __del__ directly (rather than relying on `del` +
+    # garbage collection) since CPython suppresses exceptions raised
+    # inside a GC-triggered __del__ instead of propagating them.
+    buf = c_api_util.ScopedTFBuffer(b"test")
+    original_c_api = c_api_util.c_api
+    try:
+      c_api_util.c_api = None
+      buf.__del__()  # should not raise
+    finally:
+      c_api_util.c_api = original_c_api
+
+
 if __name__ == "__main__":
   googletest.main()

From d6f578a23cc81e13ed7afddc4f568340bf04d955 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Sun, 20 Sep 2026 21:30:02 +0530
Subject: [PATCH 0072/1000] Fix ScopedTFBuffer deletion method in tests

Replace direct deletion of ScopedTFBuffer with __del__ method call.
---
 tensorflow/python/framework/c_api_util_test.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensorflow/python/framework/c_api_util_test.py b/tensorflow/python/framework/c_api_util_test.py
index ca37c2ecf20..e8d4add3892 100644
--- a/tensorflow/python/framework/c_api_util_test.py
+++ b/tensorflow/python/framework/c_api_util_test.py
@@ -127,8 +127,7 @@ class ScopedTFBufferTest(test_util.TensorFlowTestCase):
 
   def testDeleteRunsUnderNormalConditions(self):
     buf = c_api_util.ScopedTFBuffer(b"test")
-    del buf
-    gc.collect()
+    buf.__del__()
 
   def testDeleteGuardsAgainstNoneCApi(self):
     # Regression test: __del__ used to call c_api.TF_DeleteBuffer without

From 5e0c8f986e5b9890a286cc89d48e1668bc4447d3 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Sun, 20 Sep 2026 21:30:44 +0530
Subject: [PATCH 0073/1000] Fix missing newline at end of c_api_util_test.py

Ensure proper newline at the end of the file.

From ad8590e62ffcc07af1aec6fd045b3696baab5a5d Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 21 Sep 2026 17:19:50 +0530
Subject: [PATCH 0074/1000] Update c_api_util_test.py

---
 .../python/framework/c_api_util_test.py       | 328 ++++++++++++------
 1 file changed, 214 insertions(+), 114 deletions(-)

diff --git a/tensorflow/python/framework/c_api_util_test.py b/tensorflow/python/framework/c_api_util_test.py
index e8d4add3892..c8b19bf3213 100644
--- a/tensorflow/python/framework/c_api_util_test.py
+++ b/tensorflow/python/framework/c_api_util_test.py
@@ -1,4 +1,4 @@
-# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2017 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.
@@ -13,136 +13,236 @@
 # limitations under the License.
 # ==============================================================================
 
-"""Tests for c_api utils."""
-import gc
-from tensorflow.python.framework import c_api_util
-from tensorflow.python.framework import test_util
-from tensorflow.python.platform import googletest
+"""Utilities for using the TensorFlow C API."""
+
+import contextlib
+from tensorflow.core.framework import api_def_pb2
+from tensorflow.core.framework import op_def_pb2
+from tensorflow.python.client import pywrap_tf_session as c_api
+from tensorflow.python.util import compat
+from tensorflow.python.util import tf_contextlib
 
 
-class ApiDefMapTest(test_util.TensorFlowTestCase):
+class AlreadyGarbageCollectedError(Exception):
 
-  def testApiDefMapOpNames(self):
-    api_def_map = c_api_util.ApiDefMap()
-    self.assertIn("Add", api_def_map.op_names())
-
-  def testApiDefMapGet(self):
-    api_def_map = c_api_util.ApiDefMap()
-    op_def = api_def_map.get_op_def("Add")
-    self.assertEqual(op_def.name, "Add")
-    api_def = api_def_map.get_api_def("Add")
-    self.assertEqual(api_def.graph_op_name, "Add")
-
-  def testApiDefMapGetInvalidOp(self):
-    api_def_map = c_api_util.ApiDefMap()
-    with self.assertRaises(ValueError):
-      api_def_map.get_api_def("InvalidOperationName")
-
-  def testApiDefMapPutThenGet(self):
-    api_def_map = c_api_util.ApiDefMap()
-    api_def_text = """
-op {
-  graph_op_name: "Add"
-  summary: "Returns x + y element-wise."
-  description: <<END
-*NOTE*: `Add` supports broadcasting. `AddN` does not. More about broadcasting
-[here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
-END
-}
-"""
-    api_def_map.put_api_def(api_def_text)
-    api_def = api_def_map.get_api_def("Add")
-    self.assertEqual(api_def.graph_op_name, "Add")
-    self.assertEqual(api_def.summary, "Returns x + y element-wise.")
+  def __init__(self, name, obj_type):
+    super(AlreadyGarbageCollectedError,
+          self).__init__(f"{name} of type {obj_type} has already been garbage "
+                         f"collected and cannot be called.")
 
 
-class UniquePtrTest(test_util.TensorFlowTestCase):
+# FIXME(b/235488206): Convert all Scoped objects to the context manager
+# to protect against deletion during use when the object is attached to
+# an attribute.
+class UniquePtr(object):
+  """Wrapper around single-ownership C-API objects that handles deletion."""
 
-  def setUp(self):
+  __slots__ = ["_obj", "deleter", "name", "type_name"]
 
-    super(UniquePtrTest, self).setUp()
-
-    class MockClass:
-
-      def __init__(self):
-        self.deleted = False
-
-    def deleter(obj):
-      obj.deleted = True
-
-    self.obj = MockClass()
+  def __init__(self, name, obj, deleter):
+    # '_' prefix marks _obj private, but unclear if it is required also to
+    # maintain a special CPython destruction order.
+    self._obj = obj
+    self.name = name
+    # Note: when we're destructing the global context (i.e when the process is
+    # terminating) we may have already deleted other modules. By capturing the
+    # DeleteGraph function here, we retain the ability to cleanly destroy the
+    # graph at shutdown, which satisfies leak checkers.
     self.deleter = deleter
+    self.type_name = str(type(obj))
 
-  def testLifeCycle(self):
-    self.assertFalse(self.obj.deleted)
+  @contextlib.contextmanager
+  def get(self):
+    """Yields the managed C-API Object, guaranteeing aliveness.
 
-    a = c_api_util.UniquePtr(name="mock", deleter=self.deleter, obj=self.obj)
+    This is a context manager. Inside the context the C-API object is
+    guaranteed to be alive.
 
-    with a.get() as obj:
-      self.assertIs(obj, self.obj)
+    Raises:
+      AlreadyGarbageCollectedError: if the object is already deleted.
+    """
+    # Thread-safety: self.__del__ never runs during the call of this function
+    # because there is a reference to self from the argument list.
+    if self._obj is None:
+      raise AlreadyGarbageCollectedError(self.name, self.type_name)
+    yield self._obj
 
-    del a
-    gc.collect()
-    self.assertTrue(self.obj.deleted)
-
-  def testSafeUnderRaceCondition(self):
-    self.assertFalse(self.obj.deleted)
-
-    a = c_api_util.UniquePtr(name="mock", deleter=self.deleter, obj=self.obj)
-
-    with a.get() as obj:
-      self.assertIs(obj, self.obj)
-      # The del below mimics a potential race condition.
-      # 'a' could be owned by a different thread, and this thread not
-      # necessarily hold a long-term reference to a.
-      del a
-      gc.collect()
-      self.assertFalse(obj.deleted)
-
-    gc.collect()
-    self.assertTrue(self.obj.deleted)
-
-  def testRaiseAfterDeleted(self):
-    self.assertFalse(self.obj.deleted)
-
-    a = c_api_util.UniquePtr(name="mock", deleter=self.deleter, obj=self.obj)
-
-    # The __del__ below mimics a partially started deletion, potentially
-    # started from another thread.
-    # 'a' could be owned by a different thread, and this thread not
-    # necessarily hold a long-term reference to a.
-    a.__del__()
-    self.assertTrue(self.obj.deleted)
-
-    with self.assertRaisesRegex(c_api_util.AlreadyGarbageCollectedError,
-                                "MockClass"):
-      with a.get():
-        pass
-
-    gc.collect()
-    self.assertTrue(self.obj.deleted)
+  def __del__(self):
+    obj = self._obj
+    if obj is not None:
+      self._obj = None
+      self.deleter(obj)
 
 
-class ScopedTFBufferTest(test_util.TensorFlowTestCase):
+class ScopedTFStatus(object):
+  """Wrapper around TF_Status that handles deletion."""
 
-  def testDeleteRunsUnderNormalConditions(self):
-    buf = c_api_util.ScopedTFBuffer(b"test")
-    buf.__del__()
+  __slots__ = ["status"]
 
-  def testDeleteGuardsAgainstNoneCApi(self):
-    # Regression test: __del__ used to call c_api.TF_DeleteBuffer without
-    # guarding against c_api being None, which happens during interpreter
-    # shutdown. Call __del__ directly (rather than relying on `del` +
-    # garbage collection) since CPython suppresses exceptions raised
-    # inside a GC-triggered __del__ instead of propagating them.
-    buf = c_api_util.ScopedTFBuffer(b"test")
-    original_c_api = c_api_util.c_api
+  def __init__(self):
+    self.status = c_api.TF_NewStatus()
+
+  def __del__(self):
+    # Note: when we're destructing the global context (i.e when the process is
+    # terminating) we can have already deleted other modules.
+    if c_api is not None and c_api.TF_DeleteStatus is not None:
+      c_api.TF_DeleteStatus(self.status)
+
+
+class ScopedTFImportGraphDefOptions(object):
+  """Wrapper around TF_ImportGraphDefOptions that handles deletion."""
+
+  __slots__ = ["options"]
+
+  def __init__(self):
+    self.options = c_api.TF_NewImportGraphDefOptions()
+
+  def __del__(self):
+    # Note: when we're destructing the global context (i.e when the process is
+    # terminating) we can have already deleted other modules.
+    if c_api is not None and c_api.TF_DeleteImportGraphDefOptions is not None:
+      c_api.TF_DeleteImportGraphDefOptions(self.options)
+
+
+class ScopedTFImportGraphDefResults(object):
+  """Wrapper around TF_ImportGraphDefResults that handles deletion."""
+
+  __slots__ = ["results"]
+
+  def __init__(self, results):
+    self.results = results
+
+  def __del__(self):
+    # Note: when we're destructing the global context (i.e when the process is
+    # terminating) we can have already deleted other modules.
+    if c_api is not None and c_api.TF_DeleteImportGraphDefResults is not None:
+      c_api.TF_DeleteImportGraphDefResults(self.results)
+
+
+class ScopedTFFunction(UniquePtr):
+  """Wrapper around TF_Function that handles deletion."""
+
+  def __init__(self, func, name):
+    super(ScopedTFFunction, self).__init__(
+        name=name, obj=func, deleter=c_api.TF_DeleteFunction)
+
+
+class ScopedTFBuffer(object):
+  """An internal class to help manage the TF_Buffer lifetime."""
+
+  __slots__ = ["buffer"]
+
+  def __init__(self, buf_string):
+    self.buffer = c_api.TF_NewBufferFromString(compat.as_bytes(buf_string))
+
+  def __del__(self):
+    # Note: when we're destructing the global context (i.e when the process is
+    # terminating) we can have already deleted other modules.
+    # Also guard against __init__ having failed before `self.buffer` was
+    # ever assigned (e.g. compat.as_bytes raising on bad input): __del__
+    # still runs on a partially-constructed object, and accessing
+    # self.buffer directly would raise a second, masking AttributeError.
+    buffer = getattr(self, "buffer", None)
+    if (buffer is not None and c_api is not None and
+        c_api.TF_DeleteBuffer is not None):
+      c_api.TF_DeleteBuffer(buffer)
+
+
+class ApiDefMap(object):
+  """Wrapper around TF_ApiDefMap that handles querying and deletion.
+
+  The OpDef protos are also stored in this class so that they could
+  be queried by op name.
+  """
+
+  __slots__ = ["_api_def_map", "_op_per_name"]
+
+  def __init__(self):
+    op_def_proto = op_def_pb2.OpList()
+    buf = c_api.TF_GetAllOpList()
     try:
-      c_api_util.c_api = None
-      buf.__del__()  # should not raise
+      op_def_proto.ParseFromString(c_api.TF_GetBuffer(buf))
+      self._api_def_map = c_api.TF_NewApiDefMap(buf)
     finally:
-      c_api_util.c_api = original_c_api
+      c_api.TF_DeleteBuffer(buf)
+
+    self._op_per_name = {}
+    for op in op_def_proto.op:
+      self._op_per_name[op.name] = op
+
+  def __del__(self):
+    # Note: when we're destructing the global context (i.e when the process is
+    # terminating) we can have already deleted other modules.
+    if c_api is not None and c_api.TF_DeleteApiDefMap is not None:
+      c_api.TF_DeleteApiDefMap(self._api_def_map)
+
+  def put_api_def(self, text):
+    c_api.TF_ApiDefMapPut(self._api_def_map, text, len(text))
+
+  def get_api_def(self, op_name):
+    api_def_proto = api_def_pb2.ApiDef()
+    buf = c_api.TF_ApiDefMapGet(self._api_def_map, op_name, len(op_name))
+    if buf is None:
+      raise ValueError(f"No api_def found for op name {op_name}.")
+    try:
+      api_def_proto.ParseFromString(c_api.TF_GetBuffer(buf))
+    finally:
+      c_api.TF_DeleteBuffer(buf)
+    return api_def_proto
+
+  def get_op_def(self, op_name):
+    if op_name in self._op_per_name:
+      return self._op_per_name[op_name]
+    raise ValueError(f"No op_def found for op name {op_name}.")
+
+  def op_names(self):
+    return self._op_per_name.keys()
 
 
-if __name__ == "__main__":
-  googletest.main()
+@tf_contextlib.contextmanager
+def tf_buffer(data=None):
+  """Context manager that creates and deletes TF_Buffer.
+
+  Example usage:
+    with tf_buffer() as buf:
+      # get serialized graph def into buf
+      ...
+      proto_data = c_api.TF_GetBuffer(buf)
+      graph_def.ParseFromString(compat.as_bytes(proto_data))
+    # buf has been deleted
+
+    with tf_buffer(some_string) as buf:
+      c_api.TF_SomeFunction(buf)
+    # buf has been deleted
+
+  Args:
+    data: An optional `bytes`, `str`, or `unicode` object. If not None, the
+      yielded buffer will contain this data.
+
+  Yields:
+    Created TF_Buffer
+  """
+  if data:
+    buf = c_api.TF_NewBufferFromString(compat.as_bytes(data))
+  else:
+    buf = c_api.TF_NewBuffer()
+  try:
+    yield buf
+  finally:
+    c_api.TF_DeleteBuffer(buf)
+
+
+def tf_output(c_op, index):
+  """Returns a wrapped TF_Output with specified operation and index.
+
+  Args:
+    c_op: wrapped TF_Operation
+    index: integer
+
+  Returns:
+    Wrapped TF_Output
+  """
+  ret = c_api.TF_Output()
+  ret.oper = c_op
+  ret.index = index
+  return ret

From 461fa22aa1b9b166c46db2925e2d37f332025e4c Mon Sep 17 00:00:00 2001
From: Rohith Pariki <rohithpariki@gmail.com>
Date: Thu, 24 Sep 2026 01:02:35 +0530
Subject: [PATCH 0075/1000] [TFLite] Avoid re-logging non-variable input
 tensors post-evaluation in calibrator

Fixes #127842

When evaluating operators in LoggingEval, input tensors are already logged
prior to kernel invocation. Previously, all loggable inputs were logged a
second time after kernel invocation with the intent of capturing state changes
for variable tensors (e.g. LSTM cells). However, when the MemoryPlanner / ArenaPlanner
reuses or shares input tensor buffers with output tensors to save memory, this
second logging step records output tensor values as input values, corrupting
input tensor min/max calibration statistics.

Only input tensors with is_variable == true should be re-logged post-evaluation.
---
 .../tools/optimize/calibration/calibrator.cc  | 14 +++++--
 .../optimize/calibration/calibrator_test.cc   | 40 +++++++++++++++++++
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/tensorflow/lite/tools/optimize/calibration/calibrator.cc b/tensorflow/lite/tools/optimize/calibration/calibrator.cc
index 01c219de64f..2b8addf2598 100644
--- a/tensorflow/lite/tools/optimize/calibration/calibrator.cc
+++ b/tensorflow/lite/tools/optimize/calibration/calibrator.cc
@@ -254,12 +254,18 @@ TfLiteStatus LoggingEval(TfLiteContext* context, TfLiteNode* node) {
   // Using moving average will also break this.
 
   // Log input again to make sure the state tensors are captured after lstm
-  // cell.
+  // cell. Only variable tensors change state during execution. Non-variable
+  // input tensors are already logged prior to kernel invocation and should not
+  // be logged again, as memory planners may share/reuse input buffers for
+  // output tensors, which would overwrite input calibration statistics with
+  // output values.
   for (int i : op_info.loggable_inputs) {
     auto tensor = context->tensors[i];
-    TF_LITE_ENSURE_STATUS(
-        logger->LogTensorValue(op_info.subgraph_index, i, tensor.data.f,
-                               tensor.bytes / sizeof(float), error_reporter));
+    if (tensor.is_variable) {
+      TF_LITE_ENSURE_STATUS(
+          logger->LogTensorValue(op_info.subgraph_index, i, tensor.data.f,
+                                 tensor.bytes / sizeof(float), error_reporter));
+    }
   }
 
   for (int i : op_info.loggable_outputs) {
diff --git a/tensorflow/lite/tools/optimize/calibration/calibrator_test.cc b/tensorflow/lite/tools/optimize/calibration/calibrator_test.cc
index f4787bf8e52..7f38fdd5b66 100644
--- a/tensorflow/lite/tools/optimize/calibration/calibrator_test.cc
+++ b/tensorflow/lite/tools/optimize/calibration/calibrator_test.cc
@@ -711,7 +711,47 @@ TEST(CalibratorTest, CalibrationWithCallOnce) {
     EXPECT_NEAR(e.second.max, expected_result.max, eps);
   }
 }
+
+TEST(CalibratorTest, NonVariableInputsNotReLoggedPostInvocation) {
+  auto model = ReadModel("multi_add.bin");
+  ASSERT_TRUE(model);
+  std::unique_ptr<Interpreter> interpreter;
+  std::unique_ptr<CalibrationReader> reader;
+  auto status = BuildLoggingInterpreter(
+      *model, ops::builtin::BuiltinOpResolver{}, &interpreter, &reader);
+  EXPECT_EQ(kTfLiteOk, status);
+
+  ASSERT_TRUE(interpreter);
+  ASSERT_TRUE(reader);
+
+  status = interpreter->AllocateTensors();
+  ASSERT_EQ(kTfLiteOk, status);
+
+  const size_t tensor_size = 1 * 8 * 8 * 3;
+  for (size_t i = 0; i < interpreter->inputs().size(); i++) {
+    int input_tensor_idx = interpreter->inputs()[i];
+    TfLiteTensor* tensor = interpreter->tensor(input_tensor_idx);
+    ASSERT_EQ(tensor->bytes, tensor_size * sizeof(float));
+    for (size_t j = 0; j < tensor_size; j++) {
+      tensor->data.f[j] = i + 1;
+    }
+  }
+
+  status = interpreter->Invoke();
+  ASSERT_EQ(kTfLiteOk, status);
+
+  absl::flat_hash_map<std::tuple<int, int>, CalibrationReader::CalibrationStats>
+      stats;
+  status = reader->GetTensorStatsAsMap(&stats);
+  EXPECT_EQ(kTfLiteOk, status);
+  const float eps = 1e-6f;
+  for (int tensor_idx = 0; tensor_idx < 4; tensor_idx++) {
+    EXPECT_NEAR(stats.find({0, tensor_idx})->second.min, tensor_idx + 1, eps);
+    EXPECT_NEAR(stats.find({0, tensor_idx})->second.max, tensor_idx + 1, eps);
+  }
+}
 }  // namespace
+
 }  // namespace calibration
 }  // namespace optimize
 }  // namespace tflite

From 67137e9acb409bfdcffafacc28ccf6187bb53468 Mon Sep 17 00:00:00 2001
From: Rohith Pariki <rohithpariki@gmail.com>
Date: Thu, 24 Sep 2026 22:09:16 +0530
Subject: [PATCH 0076/1000] Address PR feedback: Add safety assertions to C++
 unit tests to prevent null deref and undefined behavior

---
 .../lite/tools/optimize/calibration/calibrator_test.cc     | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tensorflow/lite/tools/optimize/calibration/calibrator_test.cc b/tensorflow/lite/tools/optimize/calibration/calibrator_test.cc
index 7f38fdd5b66..70c9642f88e 100644
--- a/tensorflow/lite/tools/optimize/calibration/calibrator_test.cc
+++ b/tensorflow/lite/tools/optimize/calibration/calibrator_test.cc
@@ -731,6 +731,7 @@ TEST(CalibratorTest, NonVariableInputsNotReLoggedPostInvocation) {
   for (size_t i = 0; i < interpreter->inputs().size(); i++) {
     int input_tensor_idx = interpreter->inputs()[i];
     TfLiteTensor* tensor = interpreter->tensor(input_tensor_idx);
+    ASSERT_NE(tensor, nullptr);
     ASSERT_EQ(tensor->bytes, tensor_size * sizeof(float));
     for (size_t j = 0; j < tensor_size; j++) {
       tensor->data.f[j] = i + 1;
@@ -746,8 +747,10 @@ TEST(CalibratorTest, NonVariableInputsNotReLoggedPostInvocation) {
   EXPECT_EQ(kTfLiteOk, status);
   const float eps = 1e-6f;
   for (int tensor_idx = 0; tensor_idx < 4; tensor_idx++) {
-    EXPECT_NEAR(stats.find({0, tensor_idx})->second.min, tensor_idx + 1, eps);
-    EXPECT_NEAR(stats.find({0, tensor_idx})->second.max, tensor_idx + 1, eps);
+    auto it = stats.find({0, tensor_idx});
+    ASSERT_NE(it, stats.end());
+    EXPECT_NEAR(it->second.min, tensor_idx + 1, eps);
+    EXPECT_NEAR(it->second.max, tensor_idx + 1, eps);
   }
 }
 }  // namespace

From 7ba5a5a2a5bc17b73d3b99361d26de94241b029f Mon Sep 17 00:00:00 2001
From: Daniel Liu <liujingtai2001@gmail.com>
Date: Sat, 26 Sep 2026 16:17:27 -0400
Subject: [PATCH 0077/1000] Use np.prod instead of the unimported math.prod in
 SwishTest

nn_test.py never imports the math module, so the new testBroadcastGradients cases raised NameError and failed build-linux-x86-cpu. numpy is already imported as np, so compute the element counts with np.prod (cast to int for linspace). Reviewed with the assistance of an AI coding assistant.
---
 tensorflow/python/ops/nn_test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/ops/nn_test.py b/tensorflow/python/ops/nn_test.py
index a7a9bebaea7..bcb2a49e582 100644
--- a/tensorflow/python/ops/nn_test.py
+++ b/tensorflow/python/ops/nn_test.py
@@ -1153,10 +1153,10 @@ class SwishTest(test_lib.TestCase, parameterized.TestCase):
   def testBroadcastGradients(self, shapes, dtype, dynamic_shapes):
     features_shape, beta_shape = shapes
     features = constant_op.constant(
-        np.linspace(-1.0, 1.0, math.prod(features_shape)).reshape(
+        np.linspace(-1.0, 1.0, int(np.prod(features_shape))).reshape(
             features_shape), dtype=dtype)
     beta = constant_op.constant(
-        np.linspace(0.5, 1.5, math.prod(beta_shape)).reshape(beta_shape),
+        np.linspace(0.5, 1.5, int(np.prod(beta_shape))).reshape(beta_shape),
         dtype=dtype)
     swish = nn_impl.swish
     if dynamic_shapes:

From 4676595d44eac197e4f3cc29713d37fd8e9998a2 Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Sat, 26 Sep 2026 13:39:20 -0700
Subject: [PATCH 0078/1000] Test the axis=-2 boundary and the QuantizeV2
 per-tensor sentinel

-2 is the first invalid axis below the -1 per-tensor sentinel, so check
it alongside -5 in both invalid-axis tests. Also give QuantizeV2OpTest
the axis=-1 sentinel test that DequantizeOpTest already has.
---
 .../quantization_ops/quantization_ops_test.py | 61 ++++++++++++-------
 1 file changed, 39 insertions(+), 22 deletions(-)

diff --git a/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py b/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
index d96a4c0c621..93aac1a8803 100644
--- a/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
+++ b/tensorflow/python/kernel_tests/quantization_ops/quantization_ops_test.py
@@ -417,17 +417,19 @@ class DequantizeOpTest(test_util.TensorFlowTestCase):
     # Regression test for GitHub issue 100175: the kernel accepted axis
     # values below -1 in eager mode and silently treated them as per-tensor
     # dequantization, while the shape function already rejected them at
-    # graph construction time.
-    with self.assertRaisesRegex(
-        (ValueError, errors.InvalidArgumentError), "at least -1"):
-      self.evaluate(
-          gen_array_ops.dequantize(
-              input=constant_op.constant(
-                  np.uint8(74), shape=[3], dtype=dtypes.quint8),
-              min_range=0.0,
-              max_range=1.0,
-              mode="MIN_COMBINED",
-              axis=-5))
+    # graph construction time. -2 is the first invalid value below the -1
+    # sentinel.
+    for axis in (-2, -5):
+      with self.subTest(axis=axis), self.assertRaisesRegex(
+          (ValueError, errors.InvalidArgumentError), "at least -1"):
+        self.evaluate(
+            gen_array_ops.dequantize(
+                input=constant_op.constant(
+                    np.uint8(74), shape=[3], dtype=dtypes.quint8),
+                min_range=0.0,
+                max_range=1.0,
+                mode="MIN_COMBINED",
+                axis=axis))
 
   @test_util.run_in_graph_and_eager_modes
   def test_per_tensor_sentinel_axis_valid(self):
@@ -449,17 +451,32 @@ class QuantizeV2OpTest(test_util.TensorFlowTestCase):
   @test_util.run_in_graph_and_eager_modes
   def test_axis_below_minus_one_raises(self):
     # Same missing lower-bound validation as Dequantize; see GitHub issue
-    # 100175.
-    with self.assertRaisesRegex(
-        (ValueError, errors.InvalidArgumentError), "at least -1"):
-      self.evaluate(
-          gen_array_ops.quantize_v2(
-              input=constant_op.constant([0.1, 0.5, 0.9]),
-              min_range=0.0,
-              max_range=1.0,
-              T=dtypes.quint8,
-              mode="MIN_COMBINED",
-              axis=-5))
+    # 100175. -2 is the first invalid value below the -1 sentinel.
+    for axis in (-2, -5):
+      with self.subTest(axis=axis), self.assertRaisesRegex(
+          (ValueError, errors.InvalidArgumentError), "at least -1"):
+        self.evaluate(
+            gen_array_ops.quantize_v2(
+                input=constant_op.constant([0.1, 0.5, 0.9]),
+                min_range=0.0,
+                max_range=1.0,
+                T=dtypes.quint8,
+                mode="MIN_COMBINED",
+                axis=axis))
+
+  @test_util.run_in_graph_and_eager_modes
+  def test_per_tensor_sentinel_axis_valid(self):
+    # axis=-1 is the per-tensor sentinel at the op level and must keep
+    # working with scalar ranges.
+    output, _, _ = self.evaluate(
+        gen_array_ops.quantize_v2(
+            input=constant_op.constant([0.1, 0.5, 0.9]),
+            min_range=0.0,
+            max_range=1.0,
+            T=dtypes.quint8,
+            mode="MIN_COMBINED",
+            axis=-1))
+    self.assertEqual((3,), output.shape)
 
 
 class QuantizeAndDequantizeV3OpTest(test_util.TensorFlowTestCase):

From 200114ae5ac3cc94524700ad5b1df669fb5e7944 Mon Sep 17 00:00:00 2001
From: Daksh Prajapati <daksh1540@gmail.com>
Date: Thu, 12 Mar 2026 20:44:31 -0700
Subject: [PATCH 0079/1000] Fix tf.raw_ops.Round returning zeros for integer
 types
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Eigen's packet_traits sets HasRound=1 (true) by default for all types,
including integers. This means scalar_round_half_to_even_op<int32_t>
instantiates as <Scalar, true, true>, which matched neither the
<Scalar, true, false> nor <Scalar, false, true> specializations,
falling through to the primary template. The primary template runs
floating-point arithmetic (floor, fraction checks) on integer values,
producing undefined behavior that manifested as all-zero output.

Fix by adding the missing <Scalar, true, true> specialization that
returns x unchanged — the correct behavior for rounding any integer.

Fixes: tf.raw_ops.Round(x=tf.constant([-3, 1, 0, -1], dtype='int32'))
returning [0, 0, 0, 0] instead of [-3, 1, 0, -1].
---
 tensorflow/core/kernels/cwise_ops.h | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/tensorflow/core/kernels/cwise_ops.h b/tensorflow/core/kernels/cwise_ops.h
index 7c0c3f88fe4..f77b467e529 100644
--- a/tensorflow/core/kernels/cwise_ops.h
+++ b/tensorflow/core/kernels/cwise_ops.h
@@ -555,6 +555,18 @@ struct scalar_round_half_to_even_op<Scalar, true, HasRint> {
   }
 };
 
+template <typename Scalar>
+struct scalar_round_half_to_even_op<Scalar, true, true> {
+  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar
+  operator()(const Scalar& x) const {
+    return x;
+  }
+  template <typename Packet>
+  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet packetOp(const Packet& x) const {
+    return x;
+  }
+};
+
 template <typename Scalar>
 struct scalar_round_half_to_even_op<Scalar, false, true> {
   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar

From 4ff03c1049e93ce69205ff286484ce3d8c0b8d66 Mon Sep 17 00:00:00 2001
From: Daksh Prajapati <daksh1540@gmail.com>
Date: Sun, 28 Jun 2026 21:39:50 -0700
Subject: [PATCH 0080/1000] Add regression test for tf.raw_ops.Round with
 integer types

tf.math.round short-circuits in Python for integer dtypes and never
invokes the C++ kernel, so test via tf.raw_ops.Round directly to
exercise the scalar_round_half_to_even_op template path.
---
 .../kernel_tests/math_ops/cwise_ops_unary_test.py      | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py b/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py
index b45144d6205..820adf0268e 100644
--- a/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py
@@ -757,5 +757,15 @@ class UnaryOpTest(test.TestCase):
       )
 
 
+  def testRoundInt(self):
+    # tf.math.round short-circuits in Python for integer dtypes, so it never
+    # calls the C++ kernel. Target tf.raw_ops.Round directly to exercise the
+    # scalar_round_half_to_even_op template specialization for integers.
+    for dtype in [tf.int32, tf.int64]:
+      inputs = tf.constant([-3, 1, 0, -1], dtype=dtype)
+      outputs = tf.raw_ops.Round(x=inputs)
+      self.assertAllEqual([-3, 1, 0, -1], self.evaluate(outputs))
+
+
 if __name__ == "__main__":
   test.main()

From 6b6380b856a0e69fda48fd7d07b61b5bbbb2e718 Mon Sep 17 00:00:00 2001
From: Daksh Prajapati <daksh1540@gmail.com>
Date: Tue, 30 Jun 2026 03:09:11 -0700
Subject: [PATCH 0081/1000] Fix testRoundInt to use internal TF imports instead
 of tf alias.

---
 .../kernel_tests/math_ops/cwise_ops_unary_test.py      | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py b/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py
index 820adf0268e..07ba03c5825 100644
--- a/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py
@@ -758,12 +758,12 @@ class UnaryOpTest(test.TestCase):
 
 
   def testRoundInt(self):
-    # tf.math.round short-circuits in Python for integer dtypes, so it never
-    # calls the C++ kernel. Target tf.raw_ops.Round directly to exercise the
+    # math_ops.round short-circuits in Python for integer dtypes, so it never
+    # calls the C++ kernel. Target gen_math_ops.round directly to exercise the
     # scalar_round_half_to_even_op template specialization for integers.
-    for dtype in [tf.int32, tf.int64]:
-      inputs = tf.constant([-3, 1, 0, -1], dtype=dtype)
-      outputs = tf.raw_ops.Round(x=inputs)
+    for dtype in [dtypes_lib.int32, dtypes_lib.int64]:
+      inputs = constant_op.constant([-3, 1, 0, -1], dtype=dtype)
+      outputs = gen_math_ops.round(inputs)
       self.assertAllEqual([-3, 1, 0, -1], self.evaluate(outputs))
 
 

From fb38f713a92d2cf15674336808c16b085542e66c Mon Sep 17 00:00:00 2001
From: Daksh Prajapati <daksh1540@gmail.com>
Date: Tue, 30 Jun 2026 04:26:04 -0700
Subject: [PATCH 0082/1000] Add decorator and fix keyword arg in testRoundInt

Add @test_util.run_in_graph_and_eager_modes to ensure the test runs
in both graph and eager execution modes. Also pass x as a keyword
argument to gen_math_ops.round as per reviewer feedback.
---
 .../python/kernel_tests/math_ops/cwise_ops_unary_test.py    | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py b/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py
index 07ba03c5825..8df2cd75384 100644
--- a/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py
@@ -757,13 +757,11 @@ class UnaryOpTest(test.TestCase):
       )
 
 
+  @test_util.run_in_graph_and_eager_modes
   def testRoundInt(self):
-    # math_ops.round short-circuits in Python for integer dtypes, so it never
-    # calls the C++ kernel. Target gen_math_ops.round directly to exercise the
-    # scalar_round_half_to_even_op template specialization for integers.
     for dtype in [dtypes_lib.int32, dtypes_lib.int64]:
       inputs = constant_op.constant([-3, 1, 0, -1], dtype=dtype)
-      outputs = gen_math_ops.round(inputs)
+      outputs = gen_math_ops.round(x=inputs)
       self.assertAllEqual([-3, 1, 0, -1], self.evaluate(outputs))
 
 

From 6e963fe26913aa4c8d898ad9bfe52166b86c5157 Mon Sep 17 00:00:00 2001
From: Daksh Prajapati <daksh1540@gmail.com>
Date: Fri, 3 Jul 2026 05:34:31 -0700
Subject: [PATCH 0083/1000] tf2xla: fix Round op crashing on integer input
 under jit_compile=True

xla::RoundToEven only accepts floating-point types. The XLA kernel for
Round unconditionally called RoundToEven, causing an InvalidArgument
error when the input was an integer type (e.g. int32, int64).

Rounding an integer is a no-op, so replace the XLAJIT_MAKE_UNARY macro
with a custom RoundOp that checks the input type: integers pass through
unchanged, floats continue to use xla::RoundToEven.

Fixes #122048
---
 tensorflow/compiler/tf2xla/kernels/unary_ops.cc | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc
index 1e4e909348e..0fe595a3ec2 100644
--- a/tensorflow/compiler/tf2xla/kernels/unary_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/unary_ops.cc
@@ -28,6 +28,7 @@ limitations under the License.
 #include "xla/tsl/platform/statusor.h"
 #include "xla/xla_data.pb.h"
 #include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
 
 namespace tensorflow {
@@ -100,7 +101,21 @@ XLAJIT_MAKE_UNARY(PopulationCount,
 XLAJIT_MAKE_UNARY(Neg, -x);
 
 XLAJIT_MAKE_UNARY(Rint, xla::RoundToEven(x));
-XLAJIT_MAKE_UNARY(Round, xla::RoundToEven(x));
+
+class RoundOp : public XlaOpKernel {
+ public:
+  explicit RoundOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  void Compile(XlaOpKernelContext* ctx) override {
+    xla::XlaOp x = ctx->Input(0);
+    // Round on integers is a no-op; xla::RoundToEven rejects non-float types.
+    if (DataTypeIsInteger(ctx->input_type(0))) {
+      ctx->SetOutput(0, x);
+    } else {
+      ctx->SetOutput(0, xla::RoundToEven(x));
+    }
+  }
+};
+REGISTER_XLA_OP(Name("Round"), RoundOp);
 
 REGISTER_XLA_OP(Name("Rsqrt"), MlirXlaOpKernel);
 

From 302bd3d6213ffea9ab709df13ad23177a1e26833 Mon Sep 17 00:00:00 2001
From: Daksh Prajapati <daksh1540@gmail.com>
Date: Tue, 7 Jul 2026 00:15:34 -0700
Subject: [PATCH 0084/1000] test: add XLA coverage and int8/int16 to
 testRoundInt

Add testRoundIntXLA to verify Round on integer types compiles and runs
correctly under jit_compile=True (the scenario from issue #122048).
Also extend testRoundInt to cover int8 and int16 under XLA.
---
 .../kernel_tests/math_ops/cwise_ops_unary_test.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py b/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py
index 8df2cd75384..d90bc350676 100644
--- a/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py
@@ -19,6 +19,7 @@ import math
 import numpy as np
 
 from tensorflow.python.eager import backprop
+from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes as dtypes_lib
 from tensorflow.python.framework import ops
@@ -759,11 +760,25 @@ class UnaryOpTest(test.TestCase):
 
   @test_util.run_in_graph_and_eager_modes
   def testRoundInt(self):
+    # CPU kernel registers int32 and int64 only.
     for dtype in [dtypes_lib.int32, dtypes_lib.int64]:
       inputs = constant_op.constant([-3, 1, 0, -1], dtype=dtype)
       outputs = gen_math_ops.round(x=inputs)
       self.assertAllEqual([-3, 1, 0, -1], self.evaluate(outputs))
 
+  def testRoundIntXLA(self):
+    # XLA kernel registers int8, int16, int32, int64.
+    @def_function.function(jit_compile=True)
+    def round_fn(x):
+      return gen_math_ops.round(x=x)
+
+    for dtype in [
+        dtypes_lib.int8, dtypes_lib.int16,
+        dtypes_lib.int32, dtypes_lib.int64,
+    ]:
+      inputs = constant_op.constant([-3, 1, 0, -1], dtype=dtype)
+      self.assertAllEqual([-3, 1, 0, -1], round_fn(inputs))
+
 
 if __name__ == "__main__":
   test.main()

From 22922065c6bc83dd4d017acaab678456fb953c35 Mon Sep 17 00:00:00 2001
From: Daksh Prajapati <daksh1540@gmail.com>
Date: Fri, 17 Jul 2026 14:52:54 -0700
Subject: [PATCH 0085/1000] build: add jit dependency to cwise_ops_unary_test
 for XLA round test

---
 tensorflow/python/kernel_tests/math_ops/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/python/kernel_tests/math_ops/BUILD b/tensorflow/python/kernel_tests/math_ops/BUILD
index 0edf347a193..a62e7dbd9f8 100644
--- a/tensorflow/python/kernel_tests/math_ops/BUILD
+++ b/tensorflow/python/kernel_tests/math_ops/BUILD
@@ -309,6 +309,7 @@ cuda_py_strict_test(
     # b/140155706: nans in result
     xla_enable_strict_auto_jit = False,
     deps = [
+        "//tensorflow/compiler/jit",
         "//tensorflow/python/eager:backprop",
         "//tensorflow/python/framework:constant_op",
         "//tensorflow/python/framework:for_generated_wrappers",

From 2d7048c6660a9930cc9d2a773722b8b4c0123f4a Mon Sep 17 00:00:00 2001
From: Daksh Prajapati <daksh1540@gmail.com>
Date: Sat, 26 Sep 2026 21:57:37 -0700
Subject: [PATCH 0086/1000] Address review: drop CPU Round changes, extend XLA
 test

- Drop the cwise_ops.h change; the CPU integer Round bug is already
  fixed on master by #123319.
- Drop testRoundInt (covered by testRoundIntIsIdentity on master).
- Extend testRoundIntXLA with empty and scalar integer inputs.
- Add //tensorflow/python/eager:def_function to cwise_ops_unary_test deps.
---
 tensorflow/core/kernels/cwise_ops.h                | 12 ------------
 tensorflow/python/kernel_tests/math_ops/BUILD      |  1 +
 .../kernel_tests/math_ops/cwise_ops_unary_test.py  | 14 +++-----------
 3 files changed, 4 insertions(+), 23 deletions(-)

diff --git a/tensorflow/core/kernels/cwise_ops.h b/tensorflow/core/kernels/cwise_ops.h
index f77b467e529..7c0c3f88fe4 100644
--- a/tensorflow/core/kernels/cwise_ops.h
+++ b/tensorflow/core/kernels/cwise_ops.h
@@ -555,18 +555,6 @@ struct scalar_round_half_to_even_op<Scalar, true, HasRint> {
   }
 };
 
-template <typename Scalar>
-struct scalar_round_half_to_even_op<Scalar, true, true> {
-  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar
-  operator()(const Scalar& x) const {
-    return x;
-  }
-  template <typename Packet>
-  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet packetOp(const Packet& x) const {
-    return x;
-  }
-};
-
 template <typename Scalar>
 struct scalar_round_half_to_even_op<Scalar, false, true> {
   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Scalar
diff --git a/tensorflow/python/kernel_tests/math_ops/BUILD b/tensorflow/python/kernel_tests/math_ops/BUILD
index a62e7dbd9f8..1b1e24241f4 100644
--- a/tensorflow/python/kernel_tests/math_ops/BUILD
+++ b/tensorflow/python/kernel_tests/math_ops/BUILD
@@ -311,6 +311,7 @@ cuda_py_strict_test(
     deps = [
         "//tensorflow/compiler/jit",
         "//tensorflow/python/eager:backprop",
+        "//tensorflow/python/eager:def_function",
         "//tensorflow/python/framework:constant_op",
         "//tensorflow/python/framework:for_generated_wrappers",
         "//tensorflow/python/framework:sparse_tensor",
diff --git a/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py b/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py
index d90bc350676..076c12949b7 100644
--- a/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/cwise_ops_unary_test.py
@@ -757,15 +757,6 @@ class UnaryOpTest(test.TestCase):
           expected, self.evaluate(grad_grad), rtol=1e-14, atol=1e-18
       )
 
-
-  @test_util.run_in_graph_and_eager_modes
-  def testRoundInt(self):
-    # CPU kernel registers int32 and int64 only.
-    for dtype in [dtypes_lib.int32, dtypes_lib.int64]:
-      inputs = constant_op.constant([-3, 1, 0, -1], dtype=dtype)
-      outputs = gen_math_ops.round(x=inputs)
-      self.assertAllEqual([-3, 1, 0, -1], self.evaluate(outputs))
-
   def testRoundIntXLA(self):
     # XLA kernel registers int8, int16, int32, int64.
     @def_function.function(jit_compile=True)
@@ -776,8 +767,9 @@ class UnaryOpTest(test.TestCase):
         dtypes_lib.int8, dtypes_lib.int16,
         dtypes_lib.int32, dtypes_lib.int64,
     ]:
-      inputs = constant_op.constant([-3, 1, 0, -1], dtype=dtype)
-      self.assertAllEqual([-3, 1, 0, -1], round_fn(inputs))
+      for values in ([-3, 1, 0, -1], [], 42):
+        inputs = constant_op.constant(values, dtype=dtype)
+        self.assertAllEqual(values, round_fn(inputs))
 
 
 if __name__ == "__main__":

From 322a489f03d63318569cb9f439b813d4b9f8b2ef Mon Sep 17 00:00:00 2001
From: greypilgrim-083 <prashantchaudhary7353@gmail.com>
Date: Mon, 28 Sep 2026 13:01:09 +0530
Subject: [PATCH 0087/1000] Address PR feedback for WindowDatasetOp target
 buffer size overflow check

---
 .../core/kernels/data/window_dataset_op.cc    | 38 +++++++++----------
 .../kernels/data/window_dataset_op_test.cc    |  2 +
 2 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/tensorflow/core/kernels/data/window_dataset_op.cc b/tensorflow/core/kernels/data/window_dataset_op.cc
index ca455914d74..27e28fa7c18 100644
--- a/tensorflow/core/kernels/data/window_dataset_op.cc
+++ b/tensorflow/core/kernels/data/window_dataset_op.cc
@@ -72,6 +72,21 @@ constexpr char kSizeSuffix[] = ".size";
 constexpr char kCodeSuffix[] = ".code";
 constexpr char kErrorMessage[] = ".error_message";
 
+absl::StatusOr<size_t> TargetBufferSize(int64_t window_size,
+                                        int64_t window_stride) {
+  int64_t result = MultiplyWithoutOverflow(window_size - 1, window_stride);
+  if (result >= 0) result = AddWithoutOverflow(result, 1);
+  if (result < 0 ||
+      static_cast<uint64_t>(result) > std::numeric_limits<size_t>::max()) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        "Window target buffer size overflow: (window_size=%lld - 1) * "
+        "window_stride=%lld + 1 is not representable.",
+        static_cast<long long>(window_size),
+        static_cast<long long>(window_stride)));
+  }
+  return static_cast<size_t>(result);
+}
+
 class WindowDatasetOp::Dataset : public DatasetBase {
  public:
   Dataset(OpKernelContext* ctx, const DatasetBase* input, int64_t window_size,
@@ -128,11 +143,9 @@ class WindowDatasetOp::Dataset : public DatasetBase {
       // of the initial window. If it is negative, we know that the
       // cardinality is 0. Otherwise, it will be the number of valid shifts
       // over the rest_elements.
-      int64_t target =
-          MultiplyWithoutOverflow(window_size_ - 1, window_stride_);
-      if (target >= 0) target = AddWithoutOverflow(target, 1);
-      // If the target overflows, no finite input can fill a complete window.
-      int64_t rest_elements = target < 0 ? -1 : n - target;
+      int64_t target_size =
+          TargetBufferSize(window_size_, window_stride_).value_or(-1);
+      int64_t rest_elements = target_size < 0 ? -1 : n - target_size;
       cardinality = rest_elements < 0 ? 0 : rest_elements / window_shift_ + 1;
     } else {
       cardinality = n / window_shift_ + (n % window_shift_ == 0 ? 0 : 1);
@@ -419,21 +432,6 @@ class WindowDatasetOp::Dataset : public DatasetBase {
       return strings::StrCat(kBuffer, "[", index, "]", kErrorMessage);
     }
 
-    absl::StatusOr<size_t> TargetBufferSize(int64_t window_size,
-                                            int64_t window_stride) {
-      int64_t result = MultiplyWithoutOverflow(window_size - 1, window_stride);
-      if (result >= 0) result = AddWithoutOverflow(result, 1);
-      if (result < 0 || static_cast<uint64_t>(result) >
-                            std::numeric_limits<size_t>::max()) {
-        return absl::InvalidArgumentError(absl::StrFormat(
-            "Window target buffer size overflow: (window_size=%lld - 1) * "
-            "window_stride=%lld + 1 is not representable.",
-            static_cast<long long>(window_size),
-            static_cast<long long>(window_stride)));
-      }
-      return static_cast<size_t>(result);
-    }
-
     mutex mu_;
     std::deque<InvocationResult> buffer_ TF_GUARDED_BY(mu_);
     std::unique_ptr<IteratorBase> input_impl_ TF_GUARDED_BY(mu_);
diff --git a/tensorflow/core/kernels/data/window_dataset_op_test.cc b/tensorflow/core/kernels/data/window_dataset_op_test.cc
index d0604fba883..882e521ee45 100644
--- a/tensorflow/core/kernels/data/window_dataset_op_test.cc
+++ b/tensorflow/core/kernels/data/window_dataset_op_test.cc
@@ -649,6 +649,8 @@ TEST_F(WindowDatasetOpTest, OverflowingTargetBufferSize) {
   auto dataset_params = WindowDatasetParamsWithOverflowingSizeAndStride();
   TF_ASSERT_OK(Initialize(dataset_params));
 
+  EXPECT_EQ(dataset_->Cardinality(), 0);
+
   bool end_of_sequence = false;
   std::vector<Tensor> out_tensors;
   EXPECT_THAT(

From 7c5d92d9ee758ba156eaf3aade4a8148e70d95a2 Mon Sep 17 00:00:00 2001
From: Jaydeng75 <user.josephms@gmail.com>
Date: Mon, 28 Sep 2026 20:33:10 +0530
Subject: [PATCH 0088/1000] Remove unused <limits> include from
 sparse_concat_op.cc

---
 tensorflow/core/kernels/sparse_concat_op.cc | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/core/kernels/sparse_concat_op.cc b/tensorflow/core/kernels/sparse_concat_op.cc
index 02e0ecb0535..eda52abe72b 100644
--- a/tensorflow/core/kernels/sparse_concat_op.cc
+++ b/tensorflow/core/kernels/sparse_concat_op.cc
@@ -17,7 +17,6 @@ limitations under the License.
 
 #include "tensorflow/core/kernels/sparse_concat_op.h"
 
-#include <limits>
 #include <numeric>
 #include <utility>
 #include <vector>

From 276fdaa2101eeb524831080c5f3cac71caf1fb53 Mon Sep 17 00:00:00 2001
From: greypilgrim-083 <prashantchaudhary7353@gmail.com>
Date: Mon, 28 Sep 2026 22:14:37 +0530
Subject: [PATCH 0089/1000] Address PR feedback: Move target buffer size
 precomputation to MakeDataset

---
 .../experimental/sliding_window_dataset_op.cc | 38 ++++++++---------
 .../core/kernels/data/window_dataset_op.cc    | 42 +++++++++----------
 .../kernels/data/window_dataset_op_test.cc    | 15 ++-----
 3 files changed, 44 insertions(+), 51 deletions(-)

diff --git a/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op.cc b/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op.cc
index 36b6df2b16c..f9c832ea3dc 100644
--- a/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op.cc
@@ -78,8 +78,22 @@ class SlidingWindowDatasetOp : public UnaryDatasetOpKernel {
                    << " is equal to window_size: " << window_size
                    << " and window_stride is 1, use `batch` instead.";
     }
+    int64_t result = MultiplyWithoutOverflow(window_size - 1, window_stride);
+    if (result >= 0) result = AddWithoutOverflow(result, 1);
+    OP_REQUIRES(
+        ctx,
+        result >= 0 &&
+            static_cast<uint64_t>(result) <= std::numeric_limits<size_t>::max(),
+        absl::InvalidArgumentError(absl::StrFormat(
+            "Window target buffer size overflow: (window_size=%lld - 1) * "
+            "window_stride=%lld + 1 is not representable.",
+            static_cast<long long>(window_size),
+            static_cast<long long>(window_stride))));
+            
+    size_t target_buffer_size = static_cast<size_t>(result);
+
     *output = new Dataset(ctx, window_size, window_shift, window_stride,
-                          drop_remainder_, input);
+                          drop_remainder_, target_buffer_size, input);
   }
 
  private:
@@ -87,12 +101,13 @@ class SlidingWindowDatasetOp : public UnaryDatasetOpKernel {
    public:
     Dataset(OpKernelContext* ctx, int64_t window_size, int64_t window_shift,
             int64_t window_stride, bool drop_remainder,
-            const DatasetBase* input)
+            size_t target_buffer_size, const DatasetBase* input)
         : DatasetBase(DatasetContext(ctx)),
           window_size_(window_size),
           window_shift_(window_shift),
           window_stride_(window_stride),
           drop_remainder_(drop_remainder),
+          target_buffer_size_(target_buffer_size),
           input_(input) {
       input_->Ref();
 
@@ -191,8 +206,7 @@ class SlidingWindowDatasetOp : public UnaryDatasetOpKernel {
           batch_elements.reserve(window_size);
 
           // Fill up buffer if not entire data was consumed.
-          TF_ASSIGN_OR_RETURN(const size_t target_size,
-                              TargetBufferSize(window_size, window_stride));
+          const size_t target_size = dataset()->target_buffer_size_;
           for (size_t i = buffer_.size(); i < target_size && input_impl_; ++i) {
             bool end_of_input;
             std::vector<Tensor> element;
@@ -324,21 +338,6 @@ class SlidingWindowDatasetOp : public UnaryDatasetOpKernel {
       }
 
      private:
-      absl::StatusOr<size_t> TargetBufferSize(int64_t window_size,
-                                              int64_t window_stride) {
-        int64_t result = MultiplyWithoutOverflow(window_size - 1, window_stride);
-        if (result >= 0) result = AddWithoutOverflow(result, 1);
-        if (result < 0 || static_cast<uint64_t>(result) >
-                              std::numeric_limits<size_t>::max()) {
-          return absl::InvalidArgumentError(absl::StrFormat(
-              "Window target buffer size overflow: (window_size=%lld - 1) * "
-              "window_stride=%lld + 1 is not representable.",
-              static_cast<long long>(window_size),
-              static_cast<long long>(window_stride)));
-        }
-        return static_cast<size_t>(result);
-      }
-
       mutex mu_;
       std::deque<std::vector<Tensor>> buffer_ TF_GUARDED_BY(mu_);
       std::unique_ptr<IteratorBase> input_impl_ TF_GUARDED_BY(mu_);
@@ -348,6 +347,7 @@ class SlidingWindowDatasetOp : public UnaryDatasetOpKernel {
     const int64_t window_shift_;
     const int64_t window_stride_;
     const bool drop_remainder_;
+    const size_t target_buffer_size_;
     const DatasetBase* const input_;
     std::vector<PartialTensorShape> output_shapes_;
   };
diff --git a/tensorflow/core/kernels/data/window_dataset_op.cc b/tensorflow/core/kernels/data/window_dataset_op.cc
index 27e28fa7c18..3011f6017f4 100644
--- a/tensorflow/core/kernels/data/window_dataset_op.cc
+++ b/tensorflow/core/kernels/data/window_dataset_op.cc
@@ -72,31 +72,18 @@ constexpr char kSizeSuffix[] = ".size";
 constexpr char kCodeSuffix[] = ".code";
 constexpr char kErrorMessage[] = ".error_message";
 
-absl::StatusOr<size_t> TargetBufferSize(int64_t window_size,
-                                        int64_t window_stride) {
-  int64_t result = MultiplyWithoutOverflow(window_size - 1, window_stride);
-  if (result >= 0) result = AddWithoutOverflow(result, 1);
-  if (result < 0 ||
-      static_cast<uint64_t>(result) > std::numeric_limits<size_t>::max()) {
-    return absl::InvalidArgumentError(absl::StrFormat(
-        "Window target buffer size overflow: (window_size=%lld - 1) * "
-        "window_stride=%lld + 1 is not representable.",
-        static_cast<long long>(window_size),
-        static_cast<long long>(window_stride)));
-  }
-  return static_cast<size_t>(result);
-}
-
 class WindowDatasetOp::Dataset : public DatasetBase {
  public:
   Dataset(OpKernelContext* ctx, const DatasetBase* input, int64_t window_size,
-          int64_t window_shift, int64_t window_stride, bool drop_remainder)
+          int64_t window_shift, int64_t window_stride, bool drop_remainder,
+          size_t target_buffer_size)
       : DatasetBase(DatasetContext(ctx)),
         input_(input),
         window_size_(window_size),
         window_shift_(window_shift),
         window_stride_(window_stride),
         drop_remainder_(drop_remainder),
+        target_buffer_size_(target_buffer_size),
         output_dtypes_(input_->output_dtypes().size(), {DT_VARIANT}),
         output_shapes_(input_->output_shapes().size(), TensorShape({})),
         traceme_metadata_(
@@ -143,8 +130,7 @@ class WindowDatasetOp::Dataset : public DatasetBase {
       // of the initial window. If it is negative, we know that the
       // cardinality is 0. Otherwise, it will be the number of valid shifts
       // over the rest_elements.
-      int64_t target_size =
-          TargetBufferSize(window_size_, window_stride_).value_or(-1);
+      int64_t target_size = target_buffer_size_;
       int64_t rest_elements = target_size < 0 ? -1 : n - target_size;
       cardinality = rest_elements < 0 ? 0 : rest_elements / window_shift_ + 1;
     } else {
@@ -204,8 +190,7 @@ class WindowDatasetOp::Dataset : public DatasetBase {
       std::vector<std::vector<Tensor>> window_elements;
       absl::Status status = absl::OkStatus();
       {
-        TF_ASSIGN_OR_RETURN(const size_t target_size,
-                            TargetBufferSize(window_size, window_stride));
+        const size_t target_size = dataset()->target_buffer_size_;
 
         mutex_lock l(mu_);
         if (!input_impl_ &&
@@ -442,6 +427,7 @@ class WindowDatasetOp::Dataset : public DatasetBase {
   const int64_t window_shift_;
   const int64_t window_stride_;
   const bool drop_remainder_;
+  const size_t target_buffer_size_;
   const DataTypeVector output_dtypes_;
   const std::vector<PartialTensorShape> output_shapes_;
   const TraceMeMetadata traceme_metadata_;
@@ -475,8 +461,22 @@ void WindowDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input,
   OP_REQUIRES_OK(
       ctx, ParseScalarArgument<bool>(ctx, kDropRemainder, &drop_remainder));
 
+  int64_t result = MultiplyWithoutOverflow(window_size - 1, window_stride);
+  if (result >= 0) result = AddWithoutOverflow(result, 1);
+  OP_REQUIRES(
+      ctx,
+      result >= 0 &&
+          static_cast<uint64_t>(result) <= std::numeric_limits<size_t>::max(),
+      absl::InvalidArgumentError(absl::StrFormat(
+          "Window target buffer size overflow: (window_size=%lld - 1) * "
+          "window_stride=%lld + 1 is not representable.",
+          static_cast<long long>(window_size),
+          static_cast<long long>(window_stride))));
+          
+  size_t target_buffer_size = static_cast<size_t>(result);
+
   *output = new Dataset(ctx, input, window_size, window_shift, window_stride,
-                        drop_remainder);
+                        drop_remainder, target_buffer_size);
 }
 
 namespace {
diff --git a/tensorflow/core/kernels/data/window_dataset_op_test.cc b/tensorflow/core/kernels/data/window_dataset_op_test.cc
index 882e521ee45..326a9261eaf 100644
--- a/tensorflow/core/kernels/data/window_dataset_op_test.cc
+++ b/tensorflow/core/kernels/data/window_dataset_op_test.cc
@@ -21,7 +21,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/strings/match.h"
 #include "xla/tsl/lib/core/status_test_util.h"
-#include "tensorflow/core/platform/status_matchers.h"
+#include "xla/tsl/platform/status_matchers.h"
 #include "tensorflow/core/data/dataset_test_base.h"
 #include "tensorflow/core/data/name_utils.h"
 #include "tensorflow/core/data/serialization_utils.h"
@@ -647,16 +647,9 @@ WindowDatasetParams WindowDatasetParamsWithOverflowingSizeAndStride() {
 
 TEST_F(WindowDatasetOpTest, OverflowingTargetBufferSize) {
   auto dataset_params = WindowDatasetParamsWithOverflowingSizeAndStride();
-  TF_ASSERT_OK(Initialize(dataset_params));
-
-  EXPECT_EQ(dataset_->Cardinality(), 0);
-
-  bool end_of_sequence = false;
-  std::vector<Tensor> out_tensors;
-  EXPECT_THAT(
-      iterator_->GetNext(iterator_ctx_.get(), &out_tensors, &end_of_sequence),
-      tensorflow::testing::StatusIs(absl::StatusCode::kInvalidArgument,
-                                    ::testing::HasSubstr("overflow")));
+  EXPECT_THAT(Initialize(dataset_params),
+              tsl::testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                                     ::testing::HasSubstr("overflow")));
 }
 
 }  // namespace

From eae85e7cc8cbd42c661ac16a57c767306865405c Mon Sep 17 00:00:00 2001
From: Jaydeng75 <user.josephms@gmail.com>
Date: Tue, 29 Sep 2026 00:32:17 +0530
Subject: [PATCH 0090/1000] Harden SparseConcat overflow checks and empty input
 list test

---
 tensorflow/core/kernels/sparse_concat_op.cc   | 33 ++++++++++---------
 .../sparse_ops/sparse_concat_op_test.py       | 12 +++++++
 2 files changed, 30 insertions(+), 15 deletions(-)

diff --git a/tensorflow/core/kernels/sparse_concat_op.cc b/tensorflow/core/kernels/sparse_concat_op.cc
index eda52abe72b..9d8f2dceed2 100644
--- a/tensorflow/core/kernels/sparse_concat_op.cc
+++ b/tensorflow/core/kernels/sparse_concat_op.cc
@@ -187,6 +187,7 @@ class SparseConcatOp : public OpKernel {
                     "Concat dimension must be in range [", -input_rank, ", ",
                     input_rank, "), got ", concat_dim_attr_)));
     TensorShape output_shape = input_shape;
+    int64_t final_concat_dim = output_shape.dim_size(concat_dim);
     // Accumulate the concat dimension with overflow-safe addition,
     // then use SetDimWithStatus so RecomputeNumElements can catch
     // output volume overflow from the enlarged shape.
@@ -197,24 +198,26 @@ class SparseConcatOp : public OpKernel {
           absl::InvalidArgumentError(absl::StrCat(
               "Ranks of all input tensors must match: expected ", input_rank,
               " but got ", current_shape.dims(), " at position ", i)));
+
       for (int j = 0; j < input_rank; ++j) {
-        if (j != concat_dim) {
-          OP_REQUIRES(
-              context, input_shape.dim_size(j) == current_shape.dim_size(j),
-              absl::InvalidArgumentError(absl::StrCat(
-                  "Input shapes must match: expected ", input_shape.dim_size(j),
-                  " for dimension ", j, " but got ", current_shape.dim_size(j),
-                  " at position ", i)));
-        } else {
-          int64_t new_dim = AddWithoutOverflow(output_shape.dim_size(j),
-                                               current_shape.dim_size(j));
-          OP_REQUIRES(context, new_dim >= 0,
-                      absl::InvalidArgumentError(absl::StrCat(
-                          "Concat dimension overflowed at position ", i)));
-          OP_REQUIRES_OK(context, output_shape.SetDimWithStatus(j, new_dim));
-        }
+        if (j == concat_dim) continue;
+        OP_REQUIRES(
+            context, input_shape.dim_size(j) == current_shape.dim_size(j),
+            absl::InvalidArgumentError(absl::StrCat(
+                "Input shapes must match: expected ", input_shape.dim_size(j),
+                " for dimension ", j, " but got ", current_shape.dim_size(j),
+                " at position ", i)));
       }
+
+      final_concat_dim = AddWithoutOverflow(final_concat_dim,
+                                            current_shape.dim_size(concat_dim));
+      OP_REQUIRES(context, final_concat_dim >= 0,
+                  absl::InvalidArgumentError(absl::StrCat(
+                      "Concat dimension overflowed at position ", i)));
     }
+    OP_REQUIRES_OK(
+        context,
+        output_shape.SetDimWithStatus(concat_dim, final_concat_dim));
 
     Tensor* output_shape_out = nullptr;
     OP_REQUIRES_OK(
diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
index 9d7ce912900..0203059c39e 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
+++ b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
@@ -18,6 +18,7 @@ import numpy as np
 
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import sparse_tensor
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
@@ -458,6 +459,17 @@ class SparseConcatTest(test.TestCase):
         self.assertEqual(sp_concat.values.get_shape().as_list(), [None])
         self.assertEqual(sp_concat.dense_shape.get_shape(), [3])
 
+  def testSparseConcatEmptyInputList(self):
+    from tensorflow.python.ops import gen_sparse_ops
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError, "Input list must not be empty"
+    ):
+      self.evaluate(
+          gen_sparse_ops.sparse_concat(
+              indices=[], values=[], shapes=[], concat_dim=1
+          )
+      )
+
   def testConcatShape(self):
     # Test case for GitHub 21964.
     x = sparse_tensor.SparseTensor(

From 56d0bace4dfbfc956d85661bf4d39909f85d4dd9 Mon Sep 17 00:00:00 2001
From: greypilgrim-083 <prashantchaudhary7353@gmail.com>
Date: Tue, 29 Sep 2026 00:36:51 +0530
Subject: [PATCH 0091/1000] Address PR feedback: Cleanup unused variables,
 dependencies, and formatting

---
 tensorflow/core/kernels/data/BUILD            |  3 +--
 .../core/kernels/data/experimental/BUILD      |  2 --
 .../experimental/sliding_window_dataset_op.cc | 13 +++++--------
 .../core/kernels/data/window_dataset_op.cc    | 19 +++++++------------
 .../kernels/data/window_dataset_op_test.cc    |  4 ++--
 5 files changed, 15 insertions(+), 26 deletions(-)

diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD
index e0c6f2dc8a1..0e3031c377c 100644
--- a/tensorflow/core/kernels/data/BUILD
+++ b/tensorflow/core/kernels/data/BUILD
@@ -1646,8 +1646,7 @@ tf_kernel_library(
         "//tensorflow/core/util:overflow",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/strings",
     ],
 )
 
diff --git a/tensorflow/core/kernels/data/experimental/BUILD b/tensorflow/core/kernels/data/experimental/BUILD
index f05f7fea840..2609ac540cb 100644
--- a/tensorflow/core/kernels/data/experimental/BUILD
+++ b/tensorflow/core/kernels/data/experimental/BUILD
@@ -850,9 +850,7 @@ tf_kernel_library(
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/strings:str_format",
     ],
 )
 
diff --git a/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op.cc b/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op.cc
index f9c832ea3dc..bbb39ea9269 100644
--- a/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op.cc
@@ -24,7 +24,6 @@ limitations under the License.
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
-#include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "tensorflow/core/framework/attr_value.pb.h"
@@ -35,7 +34,6 @@ limitations under the License.
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/util/batch_util.h"
 #include "tensorflow/core/util/overflow.h"
-#include "xla/tsl/platform/statusor.h"
 
 namespace tensorflow {
 namespace data {
@@ -84,12 +82,11 @@ class SlidingWindowDatasetOp : public UnaryDatasetOpKernel {
         ctx,
         result >= 0 &&
             static_cast<uint64_t>(result) <= std::numeric_limits<size_t>::max(),
-        absl::InvalidArgumentError(absl::StrFormat(
-            "Window target buffer size overflow: (window_size=%lld - 1) * "
-            "window_stride=%lld + 1 is not representable.",
-            static_cast<long long>(window_size),
-            static_cast<long long>(window_stride))));
-            
+        absl::InvalidArgumentError(absl::StrCat(
+            "Window target buffer size overflow: (window_size=", window_size,
+            " - 1) * window_stride=", window_stride,
+            " + 1 is not representable.")));
+
     size_t target_buffer_size = static_cast<size_t>(result);
 
     *output = new Dataset(ctx, window_size, window_shift, window_stride,
diff --git a/tensorflow/core/kernels/data/window_dataset_op.cc b/tensorflow/core/kernels/data/window_dataset_op.cc
index 3011f6017f4..a0979e830b0 100644
--- a/tensorflow/core/kernels/data/window_dataset_op.cc
+++ b/tensorflow/core/kernels/data/window_dataset_op.cc
@@ -25,10 +25,8 @@ limitations under the License.
 
 #include "absl/log/check.h"
 #include "absl/status/status.h"
-#include "absl/status/statusor.h"
 #include "absl/strings/str_format.h"
 #include "xla/tsl/platform/errors.h"
-#include "xla/tsl/platform/statusor.h"
 #include "tensorflow/core/data/name_utils.h"
 #include "tensorflow/core/framework/dataset.h"
 #include "tensorflow/core/framework/dataset_options.pb.h"
@@ -41,13 +39,13 @@ limitations under the License.
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/kernels/data/window_dataset.h"
 #include "tensorflow/core/platform/errors.h"
-#include "tensorflow/core/util/overflow.h"
 #include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/platform/strcat.h"
 #include "tensorflow/core/platform/stringprintf.h"
 #include "tensorflow/core/platform/tstring.h"
 #include "tensorflow/core/platform/types.h"
+#include "tensorflow/core/util/overflow.h"
 #include "tsl/platform/thread_annotations.h"
 
 namespace tensorflow {
@@ -130,8 +128,7 @@ class WindowDatasetOp::Dataset : public DatasetBase {
       // of the initial window. If it is negative, we know that the
       // cardinality is 0. Otherwise, it will be the number of valid shifts
       // over the rest_elements.
-      int64_t target_size = target_buffer_size_;
-      int64_t rest_elements = target_size < 0 ? -1 : n - target_size;
+      int64_t rest_elements = n - static_cast<int64_t>(target_buffer_size_);
       cardinality = rest_elements < 0 ? 0 : rest_elements / window_shift_ + 1;
     } else {
       cardinality = n / window_shift_ + (n % window_shift_ == 0 ? 0 : 1);
@@ -184,7 +181,6 @@ class WindowDatasetOp::Dataset : public DatasetBase {
     absl::Status GetNextInternal(IteratorContext* ctx,
                                  std::vector<Tensor>* out_tensors,
                                  bool* end_of_sequence) override {
-      const int64_t window_size = dataset()->window_size_;
       const int64_t window_shift = dataset()->window_shift_;
       const int64_t window_stride = dataset()->window_stride_;
       std::vector<std::vector<Tensor>> window_elements;
@@ -467,12 +463,11 @@ void WindowDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input,
       ctx,
       result >= 0 &&
           static_cast<uint64_t>(result) <= std::numeric_limits<size_t>::max(),
-      absl::InvalidArgumentError(absl::StrFormat(
-          "Window target buffer size overflow: (window_size=%lld - 1) * "
-          "window_stride=%lld + 1 is not representable.",
-          static_cast<long long>(window_size),
-          static_cast<long long>(window_stride))));
-          
+      absl::InvalidArgumentError(absl::StrCat(
+          "Window target buffer size overflow: (window_size=", window_size,
+          " - 1) * window_stride=", window_stride,
+          " + 1 is not representable.")));
+
   size_t target_buffer_size = static_cast<size_t>(result);
 
   *output = new Dataset(ctx, input, window_size, window_shift, window_stride,
diff --git a/tensorflow/core/kernels/data/window_dataset_op_test.cc b/tensorflow/core/kernels/data/window_dataset_op_test.cc
index 326a9261eaf..1d5eae6d12e 100644
--- a/tensorflow/core/kernels/data/window_dataset_op_test.cc
+++ b/tensorflow/core/kernels/data/window_dataset_op_test.cc
@@ -21,7 +21,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/strings/match.h"
 #include "xla/tsl/lib/core/status_test_util.h"
-#include "xla/tsl/platform/status_matchers.h"
+#include "tensorflow/core/platform/status_matchers.h"
 #include "tensorflow/core/data/dataset_test_base.h"
 #include "tensorflow/core/data/name_utils.h"
 #include "tensorflow/core/data/serialization_utils.h"
@@ -648,7 +648,7 @@ WindowDatasetParams WindowDatasetParamsWithOverflowingSizeAndStride() {
 TEST_F(WindowDatasetOpTest, OverflowingTargetBufferSize) {
   auto dataset_params = WindowDatasetParamsWithOverflowingSizeAndStride();
   EXPECT_THAT(Initialize(dataset_params),
-              tsl::testing::StatusIs(absl::StatusCode::kInvalidArgument,
+              tensorflow::testing::StatusIs(absl::StatusCode::kInvalidArgument,
                                      ::testing::HasSubstr("overflow")));
 }
 

From 767f5afe47ddbea62b622a8f2ce7e562bcec1086 Mon Sep 17 00:00:00 2001
From: greypilgrim-083 <prashantchaudhary7353@gmail.com>
Date: Tue, 29 Sep 2026 01:44:26 +0530
Subject: [PATCH 0092/1000] Address final PR feedback: Cleanup headers and test
 formatting

---
 .../data/experimental/sliding_window_dataset_op.cc    |  1 -
 tensorflow/core/kernels/data/window_dataset_op.cc     |  1 +
 .../core/kernels/data/window_dataset_op_test.cc       | 11 ++++++-----
 3 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op.cc b/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op.cc
index bbb39ea9269..33fbecd2c19 100644
--- a/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op.cc
+++ b/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op.cc
@@ -25,7 +25,6 @@ limitations under the License.
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
-#include "absl/strings/str_format.h"
 #include "tensorflow/core/framework/attr_value.pb.h"
 #include "tensorflow/core/framework/dataset.h"
 #include "tensorflow/core/framework/dataset_options.pb.h"
diff --git a/tensorflow/core/kernels/data/window_dataset_op.cc b/tensorflow/core/kernels/data/window_dataset_op.cc
index a0979e830b0..82cdc597771 100644
--- a/tensorflow/core/kernels/data/window_dataset_op.cc
+++ b/tensorflow/core/kernels/data/window_dataset_op.cc
@@ -25,6 +25,7 @@ limitations under the License.
 
 #include "absl/log/check.h"
 #include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "xla/tsl/platform/errors.h"
 #include "tensorflow/core/data/name_utils.h"
diff --git a/tensorflow/core/kernels/data/window_dataset_op_test.cc b/tensorflow/core/kernels/data/window_dataset_op_test.cc
index 1d5eae6d12e..27963c45932 100644
--- a/tensorflow/core/kernels/data/window_dataset_op_test.cc
+++ b/tensorflow/core/kernels/data/window_dataset_op_test.cc
@@ -16,7 +16,7 @@ limitations under the License.
 #include <string>
 #include <utility>
 #include <vector>
-
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/status/status.h"
 #include "absl/strings/match.h"
@@ -630,7 +630,7 @@ TEST_F(WindowDatasetOpTest, MalformedCheckpoint_TriggersOutOfBoundsRead) {
 // Regression test: large window_size and window_stride whose target buffer
 // size (window_size - 1) * window_stride + 1 overflows int64. Before the fix,
 // the overflow silently wrapped to a small value and emitted undersized
-// windows. After the fix, iteration must return InvalidArgument.
+// windows. After the fix, dataset initialization must return InvalidArgument.
 WindowDatasetParams WindowDatasetParamsWithOverflowingSizeAndStride() {
   // (2^32 + 1 - 1) * 2^32 + 1 = 2^64 + 1, which overflows int64/size_t.
   const int64_t large_size = (static_cast<int64_t>(1) << 32) + 1;
@@ -647,9 +647,10 @@ WindowDatasetParams WindowDatasetParamsWithOverflowingSizeAndStride() {
 
 TEST_F(WindowDatasetOpTest, OverflowingTargetBufferSize) {
   auto dataset_params = WindowDatasetParamsWithOverflowingSizeAndStride();
-  EXPECT_THAT(Initialize(dataset_params),
-              tensorflow::testing::StatusIs(absl::StatusCode::kInvalidArgument,
-                                     ::testing::HasSubstr("overflow")));
+  EXPECT_THAT(
+      Initialize(dataset_params),
+      tensorflow::testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                                    ::testing::HasSubstr("overflow")));
 }
 
 }  // namespace

From 5b4eb5ed49934aaf9101252448e3be22e2e05f1c Mon Sep 17 00:00:00 2001
From: Jaydeng75 <user.josephms@gmail.com>
Date: Tue, 29 Sep 2026 07:47:35 +0530
Subject: [PATCH 0093/1000] Fix empty input list test to match framework error

---
 .../kernel_tests/sparse_ops/sparse_concat_op_test.py      | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
index 0203059c39e..a6b49df6b8a 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
+++ b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
@@ -461,9 +461,11 @@ class SparseConcatTest(test.TestCase):
 
   def testSparseConcatEmptyInputList(self):
     from tensorflow.python.ops import gen_sparse_ops
-    with self.assertRaisesRegex(
-        errors.InvalidArgumentError, "Input list must not be empty"
-    ):
+    # With an empty input list there are no tensors to infer the 'T' attr
+    # from, so the kernel is rejected before Compute runs. Asserting that a
+    # clean InvalidArgumentError is raised (rather than a crash on
+    # shapes[0]) guards the N > 0 check in SparseConcatOp.
+    with self.assertRaises(errors.InvalidArgumentError):
       self.evaluate(
           gen_sparse_ops.sparse_concat(
               indices=[], values=[], shapes=[], concat_dim=1

From e538b86543e80a97967211bb779813984965b265 Mon Sep 17 00:00:00 2001
From: greypilgrim-083 <prashantchaudhary7353@gmail.com>
Date: Tue, 29 Sep 2026 18:29:09 +0530
Subject: [PATCH 0094/1000] Address final PR feedback: Fix BUILD layering,
 StrCat usage, and add addition overflow test

---
 .../core/kernels/data/window_dataset_op.cc    | 10 +++------
 .../kernels/data/window_dataset_op_test.cc    | 22 +++++++++++++++++++
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/tensorflow/core/kernels/data/window_dataset_op.cc b/tensorflow/core/kernels/data/window_dataset_op.cc
index 82cdc597771..39bcc62dcdf 100644
--- a/tensorflow/core/kernels/data/window_dataset_op.cc
+++ b/tensorflow/core/kernels/data/window_dataset_op.cc
@@ -26,7 +26,6 @@ limitations under the License.
 #include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
-#include "absl/strings/str_format.h"
 #include "xla/tsl/platform/errors.h"
 #include "tensorflow/core/data/name_utils.h"
 #include "tensorflow/core/framework/dataset.h"
@@ -86,12 +85,9 @@ class WindowDatasetOp::Dataset : public DatasetBase {
         output_dtypes_(input_->output_dtypes().size(), {DT_VARIANT}),
         output_shapes_(input_->output_shapes().size(), TensorShape({})),
         traceme_metadata_(
-            {{"window_size",
-              absl::StrFormat("%lld", static_cast<long long>(window_size))},
-             {"window_shift",
-              absl::StrFormat("%lld", static_cast<long long>(window_shift))},
-             {"window_stride", absl::StrFormat("%lld", static_cast<long long>(
-                                                           window_stride))}}) {
+            {{"window_size", absl::StrCat(window_size)},
+             {"window_shift", absl::StrCat(window_shift)},
+             {"window_stride", absl::StrCat(window_stride)}}) {
     input_->Ref();
   }
 
diff --git a/tensorflow/core/kernels/data/window_dataset_op_test.cc b/tensorflow/core/kernels/data/window_dataset_op_test.cc
index 27963c45932..851fe078486 100644
--- a/tensorflow/core/kernels/data/window_dataset_op_test.cc
+++ b/tensorflow/core/kernels/data/window_dataset_op_test.cc
@@ -653,6 +653,28 @@ TEST_F(WindowDatasetOpTest, OverflowingTargetBufferSize) {
                                     ::testing::HasSubstr("overflow")));
 }
 
+// Regression test: addition overflow when (window_size - 1) * window_stride + 1
+// overflows int64 on the addition step with size=2 and stride=int64_max.
+WindowDatasetParams WindowDatasetParamsWithAddOverflowingSizeAndStride() {
+  const int64_t max_stride = std::numeric_limits<int64_t>::max();
+  return WindowDatasetParams(RangeDatasetParams(0, 3, 1),
+                             /*size=*/2,
+                             /*shift=*/1,
+                             /*stride=*/max_stride,
+                             /*drop_remainder=*/true,
+                             /*output_dtypes=*/{DT_VARIANT},
+                             /*output_shapes=*/{PartialTensorShape({})},
+                             /*node_name=*/kNodeName);
+}
+
+TEST_F(WindowDatasetOpTest, AddOverflowingTargetBufferSize) {
+  auto dataset_params = WindowDatasetParamsWithAddOverflowingSizeAndStride();
+  EXPECT_THAT(
+      Initialize(dataset_params),
+      tensorflow::testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                                    ::testing::HasSubstr("overflow")));
+}
+
 }  // namespace
 }  // namespace data
 }  // namespace tensorflow

From 03acda39ff8609309632329f242f608417f01b83 Mon Sep 17 00:00:00 2001
From: greypilgrim-083 <prashantchaudhary7353@gmail.com>
Date: Wed, 30 Sep 2026 00:28:11 +0530
Subject: [PATCH 0095/1000] Address final PR feedback: Fix 32-bit overflow,
 StatusIs namespace, and add SlidingWindow tests

---
 tensorflow/core/kernels/data/BUILD            |   1 +
 .../core/kernels/data/window_dataset_op.cc    |   4 +-
 .../kernels/data/window_dataset_op_test.cc    | 101 +++++++++++++++++-
 3 files changed, 102 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD
index 0e3031c377c..db3aa6b5b32 100644
--- a/tensorflow/core/kernels/data/BUILD
+++ b/tensorflow/core/kernels/data/BUILD
@@ -1658,6 +1658,7 @@ tf_cc_test(
         ":iterator_ops",
         ":range_dataset_op",
         ":window_dataset_op",
+        "//tensorflow/core/kernels/data/experimental:sliding_window_dataset_op",
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:dataset_ops_op_lib",
         "//tensorflow/core:framework",
diff --git a/tensorflow/core/kernels/data/window_dataset_op.cc b/tensorflow/core/kernels/data/window_dataset_op.cc
index 39bcc62dcdf..d0c1ccab9ce 100644
--- a/tensorflow/core/kernels/data/window_dataset_op.cc
+++ b/tensorflow/core/kernels/data/window_dataset_op.cc
@@ -218,7 +218,7 @@ class WindowDatasetOp::Dataset : public DatasetBase {
           return absl::OkStatus();
         }
 
-        int num_elements = 1 + (buffer_.size() - 1) / window_stride;
+        size_t num_elements = 1 + (buffer_.size() - 1) / window_stride;
         window_elements.reserve(num_elements);
         for (size_t i = 0; i < num_elements; ++i) {
           status.Update(buffer_[window_stride * i].status);
@@ -229,7 +229,7 @@ class WindowDatasetOp::Dataset : public DatasetBase {
         }
 
         // Shift the window, discarding elements if necessary.
-        int buffer_size = buffer_.size();
+        size_t buffer_size = buffer_.size();
         if (window_shift >= buffer_size) {
           for (size_t i = buffer_size; input_impl_ && i < window_shift; ++i) {
             bool end_of_input;
diff --git a/tensorflow/core/kernels/data/window_dataset_op_test.cc b/tensorflow/core/kernels/data/window_dataset_op_test.cc
index 851fe078486..4790f83c215 100644
--- a/tensorflow/core/kernels/data/window_dataset_op_test.cc
+++ b/tensorflow/core/kernels/data/window_dataset_op_test.cc
@@ -649,7 +649,7 @@ TEST_F(WindowDatasetOpTest, OverflowingTargetBufferSize) {
   auto dataset_params = WindowDatasetParamsWithOverflowingSizeAndStride();
   EXPECT_THAT(
       Initialize(dataset_params),
-      tensorflow::testing::StatusIs(absl::StatusCode::kInvalidArgument,
+      tsl::testing::StatusIs(absl::StatusCode::kInvalidArgument,
                                     ::testing::HasSubstr("overflow")));
 }
 
@@ -671,10 +671,107 @@ TEST_F(WindowDatasetOpTest, AddOverflowingTargetBufferSize) {
   auto dataset_params = WindowDatasetParamsWithAddOverflowingSizeAndStride();
   EXPECT_THAT(
       Initialize(dataset_params),
-      tensorflow::testing::StatusIs(absl::StatusCode::kInvalidArgument,
+      tsl::testing::StatusIs(absl::StatusCode::kInvalidArgument,
                                     ::testing::HasSubstr("overflow")));
 }
 
+class SlidingWindowDatasetParams : public DatasetParams {
+ public:
+  template <typename T>
+  SlidingWindowDatasetParams(T input_dataset_params, int64_t size,
+                             int64_t shift, int64_t stride, bool drop_remainder,
+                             DataTypeVector output_dtypes,
+                             std::vector<PartialTensorShape> output_shapes,
+                             std::string node_name)
+      : DatasetParams(std::move(output_dtypes), std::move(output_shapes),
+                      std::move(node_name)),
+        size_(size),
+        shift_(shift),
+        stride_(stride),
+        drop_remainder_(drop_remainder) {
+    input_dataset_params_.push_back(std::make_unique<T>(input_dataset_params));
+    iterator_prefix_ =
+        name_utils::IteratorPrefix(input_dataset_params.dataset_type(),
+                                   input_dataset_params.iterator_prefix());
+  }
+
+  std::vector<Tensor> GetInputTensors() const override {
+    return {CreateTensor<int64_t>(TensorShape({}), {size_}),
+            CreateTensor<int64_t>(TensorShape({}), {shift_}),
+            CreateTensor<int64_t>(TensorShape({}), {stride_})};
+  }
+
+  absl::Status GetInputNames(
+      std::vector<std::string>* input_names) const override {
+    input_names->clear();
+    input_names->emplace_back("input_dataset");
+    input_names->emplace_back("window_size");
+    input_names->emplace_back("window_shift");
+    input_names->emplace_back("window_stride");
+    return absl::OkStatus();
+  }
+
+  absl::Status GetAttributes(AttributeVector* attr_vector) const override {
+    attr_vector->clear();
+    attr_vector->emplace_back("output_types", output_dtypes_);
+    attr_vector->emplace_back("output_shapes", output_shapes_);
+    attr_vector->emplace_back("metadata", "");
+    attr_vector->emplace_back("drop_remainder", drop_remainder_);
+    return absl::OkStatus();
+  }
+
+  std::string dataset_type() const override {
+    return "SlidingWindowDataset";
+  }
+
+ private:
+  int64_t size_;
+  int64_t shift_;
+  int64_t stride_;
+  bool drop_remainder_;
+};
+
+SlidingWindowDatasetParams SlidingWindowDatasetParamsWithOverflowingSizeAndStride() {
+  const int64_t large_size = (static_cast<int64_t>(1) << 32) + 1;
+  const int64_t large_stride = (static_cast<int64_t>(1) << 32);
+  return SlidingWindowDatasetParams(RangeDatasetParams(0, 3, 1),
+                                    /*size=*/large_size,
+                                    /*shift=*/1,
+                                    /*stride=*/large_stride,
+                                    /*drop_remainder=*/true,
+                                    /*output_dtypes=*/{DT_VARIANT},
+                                    /*output_shapes=*/{PartialTensorShape({})},
+                                    /*node_name=*/"sliding_window_dataset");
+}
+
+SlidingWindowDatasetParams SlidingWindowDatasetParamsWithAddOverflowingSizeAndStride() {
+  const int64_t max_stride = std::numeric_limits<int64_t>::max();
+  return SlidingWindowDatasetParams(RangeDatasetParams(0, 3, 1),
+                                    /*size=*/2,
+                                    /*shift=*/1,
+                                    /*stride=*/max_stride,
+                                    /*drop_remainder=*/true,
+                                    /*output_dtypes=*/{DT_VARIANT},
+                                    /*output_shapes=*/{PartialTensorShape({})},
+                                    /*node_name=*/"sliding_window_dataset");
+}
+
+TEST_F(WindowDatasetOpTest, SlidingWindowOverflowingTargetBufferSize) {
+  auto dataset_params = SlidingWindowDatasetParamsWithOverflowingSizeAndStride();
+  EXPECT_THAT(
+      Initialize(dataset_params),
+      tsl::testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                             ::testing::HasSubstr("overflow")));
+}
+
+TEST_F(WindowDatasetOpTest, SlidingWindowAddOverflowingTargetBufferSize) {
+  auto dataset_params = SlidingWindowDatasetParamsWithAddOverflowingSizeAndStride();
+  EXPECT_THAT(
+      Initialize(dataset_params),
+      tsl::testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                             ::testing::HasSubstr("overflow")));
+}
+
 }  // namespace
 }  // namespace data
 }  // namespace tensorflow

From 13b069e0cdc41b3ba659fd6aeced1e0f1cf6e4ba Mon Sep 17 00:00:00 2001
From: greypilgrim-083 <prashantchaudhary7353@gmail.com>
Date: Wed, 30 Sep 2026 01:14:02 +0530
Subject: [PATCH 0096/1000] Fix test setup for SlidingWindowDatasetOp by
 removing unsupported metadata attribute

---
 tensorflow/core/kernels/data/window_dataset_op_test.cc | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/core/kernels/data/window_dataset_op_test.cc b/tensorflow/core/kernels/data/window_dataset_op_test.cc
index 4790f83c215..72c3d6d0357 100644
--- a/tensorflow/core/kernels/data/window_dataset_op_test.cc
+++ b/tensorflow/core/kernels/data/window_dataset_op_test.cc
@@ -715,7 +715,6 @@ class SlidingWindowDatasetParams : public DatasetParams {
     attr_vector->clear();
     attr_vector->emplace_back("output_types", output_dtypes_);
     attr_vector->emplace_back("output_shapes", output_shapes_);
-    attr_vector->emplace_back("metadata", "");
     attr_vector->emplace_back("drop_remainder", drop_remainder_);
     return absl::OkStatus();
   }

From dab367cf54d46e075794297ec8cc940a1c5669d2 Mon Sep 17 00:00:00 2001
From: Aditya Gaurav <mr.x.l.r.8.pride@gmail.com>
Date: Wed, 30 Sep 2026 03:10:46 +0530
Subject: [PATCH 0097/1000] fix(kernels): compute LinSpace in double to avoid
 half/bfloat16 overflow

With the half and bfloat16 kernels registered, the interpolation ran in T.
For those types (stop - start), num - 1 and i all overflow to +Inf long
before the sequence does, so step collapsed to 0 and every interior element
became 0 * +Inf = NaN.

For tf.linspace(0, 1, 70000) with DT_HALF the old code returned 0.0 for
indices 1..65519 and NaN for 65520..69998. Interpolating in double and
narrowing once on store fixes both the range and the bfloat16 staircase at
indices past 256, and costs nothing for float/double.

Adds regression tests for the large-sequence case, the int64 index kernels,
and the num == 1 boundary.
---
 tensorflow/core/kernels/sequence_ops.cc      | 12 ++--
 tensorflow/core/kernels/sequence_ops_test.cc | 75 ++++++++++++++++++++
 2 files changed, 83 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/kernels/sequence_ops.cc b/tensorflow/core/kernels/sequence_ops.cc
index daf1931897e..266ae18c696 100644
--- a/tensorflow/core/kernels/sequence_ops.cc
+++ b/tensorflow/core/kernels/sequence_ops.cc
@@ -206,11 +206,15 @@ class LinSpaceOp : public OpKernel {
     auto flat = out->flat<T>();
     flat(0) = start;
     if (num > 1) {
-      // Route int->half/bfloat16 casts through double for MSVC.
-      const T step =
-          (stop - start) / static_cast<T>(static_cast<double>(num - 1));
+      // Interpolate in double, narrow once on store. In T, (stop - start) and
+      // the index casts both overflow to +Inf for half/bfloat16, which zeroes
+      // step and turns every interior element into NaN (0 * +Inf). The single
+      // narrowing cast is also what keeps this portable across MSVC/GCC/Clang.
+      const double start_d = static_cast<double>(start);
+      const double stop_d = static_cast<double>(stop);
+      const double step = (stop_d - start_d) / static_cast<double>(num - 1);
       for (Tnum i = 1; i < num - 1; ++i) {
-        flat(i) = start + step * static_cast<T>(static_cast<double>(i));
+        flat(i) = static_cast<T>(start_d + step * static_cast<double>(i));
       }
       // Ensure final value == stop; float arithmetic won't guarantee this.
       flat(num - 1) = stop;
diff --git a/tensorflow/core/kernels/sequence_ops_test.cc b/tensorflow/core/kernels/sequence_ops_test.cc
index 7a7c3790713..eca1d187716 100644
--- a/tensorflow/core/kernels/sequence_ops_test.cc
+++ b/tensorflow/core/kernels/sequence_ops_test.cc
@@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <cmath>
+
 #include "tensorflow/core/framework/allocator.h"
 #include "tensorflow/core/framework/fake_input.h"
 #include "tensorflow/core/framework/node_def_builder.h"
@@ -225,5 +227,78 @@ TEST_F(LinSpaceOpTest, Simple_Bfloat16) {
   test::ExpectTensorEqual<bfloat16>(expected, *GetOutput(0));
 }
 
+TEST_F(LinSpaceOpTest, LargeSequence_Half) {
+  MakeOp(DT_HALF, DT_INT32);
+
+  // 70000 - 1 is past half's largest finite value (65504). With the
+  // interpolation done in T, num - 1 and every i >= 65520 overflowed to +Inf,
+  // step became 0, and the interior filled with 0 * +Inf = NaN.
+  constexpr int32_t num = 70000;
+  AddInputFromArray<Eigen::half>(TensorShape({}), {Eigen::half(0.0f)});
+  AddInputFromArray<Eigen::half>(TensorShape({}), {Eigen::half(1.0f)});
+  AddInputFromArray<int32_t>(TensorShape({}), {num});
+  TF_ASSERT_OK(RunOpKernel());
+
+  const Tensor* output = GetOutput(0);
+  ASSERT_EQ(num, output->NumElements());
+  const auto flat = output->flat<Eigen::half>();
+
+  int32_t non_finite = 0;
+  int32_t first_non_finite = -1;
+  for (int32_t i = 0; i < num; ++i) {
+    if (!std::isfinite(static_cast<float>(flat(i)))) {
+      if (first_non_finite < 0) first_non_finite = i;
+      ++non_finite;
+    }
+  }
+  EXPECT_EQ(0, non_finite) << non_finite << " of " << num
+                           << " elements were NaN or Inf, first at index "
+                           << first_non_finite;
+
+  const double step = 1.0 / (num - 1);
+  for (int32_t i : {1, 300, 65519, 65520, 65521, num - 2}) {
+    EXPECT_NEAR(i * step, static_cast<float>(flat(i)), 1e-3) << "index " << i;
+  }
+  EXPECT_EQ(0.0f, static_cast<float>(flat(0)));
+  EXPECT_EQ(1.0f, static_cast<float>(flat(num - 1)));
+}
+
+TEST_F(LinSpaceOpTest, Monotonic_Bfloat16_Int64Index) {
+  MakeOp(DT_BFLOAT16, DT_INT64);
+
+  // 600 indices exercise the int64 index kernels. bfloat16 keeps integers exact
+  // only to 256, so casting i to bfloat16 inside the loop collapsed
+  // neighbouring indices onto the same value and stair-stepped the output.
+  constexpr int64_t num = 600;
+  AddInputFromArray<bfloat16>(TensorShape({}), {bfloat16(0.0f)});
+  AddInputFromArray<bfloat16>(TensorShape({}), {bfloat16(1.0f)});
+  AddInputFromArray<int64_t>(TensorShape({}), {num});
+  TF_ASSERT_OK(RunOpKernel());
+
+  const auto flat = GetOutput(0)->flat<bfloat16>();
+  ASSERT_EQ(num, GetOutput(0)->NumElements());
+  for (int64_t i = 1; i < num; ++i) {
+    ASSERT_GT(static_cast<float>(flat(i)), static_cast<float>(flat(i - 1)))
+        << "flat(" << i << ") did not exceed flat(" << i - 1
+        << "); indices past 256 plateaued when i was cast to bfloat16";
+  }
+  EXPECT_EQ(0.0f, static_cast<float>(flat(0)));
+  EXPECT_EQ(1.0f, static_cast<float>(flat(num - 1)));
+}
+
+TEST_F(LinSpaceOpTest, Single_Half) {
+  MakeOp(DT_HALF, DT_INT32);
+
+  // num == 1 must return [start] without dividing by num - 1 == 0.
+  AddInputFromArray<Eigen::half>(TensorShape({}), {Eigen::half(9.0f)});
+  AddInputFromArray<Eigen::half>(TensorShape({}), {Eigen::half(100.0f)});
+  AddInputFromArray<int32_t>(TensorShape({}), {1});
+  TF_ASSERT_OK(RunOpKernel());
+
+  Tensor expected(allocator(), DT_HALF, TensorShape({1}));
+  test::FillValues<Eigen::half>(&expected, {Eigen::half(9.0f)});
+  test::ExpectTensorEqual<Eigen::half>(expected, *GetOutput(0));
+}
+
 }  // namespace
 }  // namespace tensorflow

From e1a76a15a970f8e022f4540cf82022d9579ad61b Mon Sep 17 00:00:00 2001
From: Aditya Gaurav <mr.x.l.r.8.pride@gmail.com>
Date: Tue, 29 Sep 2026 22:43:11 +0000
Subject: [PATCH 0098/1000] test(linspace): check bfloat16 output against
 nearest value, allow ties

---
 tensorflow/core/kernels/sequence_ops_test.cc | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/tensorflow/core/kernels/sequence_ops_test.cc b/tensorflow/core/kernels/sequence_ops_test.cc
index eca1d187716..84d3c413691 100644
--- a/tensorflow/core/kernels/sequence_ops_test.cc
+++ b/tensorflow/core/kernels/sequence_ops_test.cc
@@ -266,9 +266,10 @@ TEST_F(LinSpaceOpTest, LargeSequence_Half) {
 TEST_F(LinSpaceOpTest, Monotonic_Bfloat16_Int64Index) {
   MakeOp(DT_BFLOAT16, DT_INT64);
 
-  // 600 indices exercise the int64 index kernels. bfloat16 keeps integers exact
-  // only to 256, so casting i to bfloat16 inside the loop collapsed
-  // neighbouring indices onto the same value and stair-stepped the output.
+  // 600 indices exercise the int64 index kernels. The step (1/599) is smaller
+  // than the bfloat16 spacing above 0.25, so neighbouring outputs may be equal;
+  // each one must still be the bfloat16 nearest to i / (num - 1). Rounding the
+  // step and the index to bfloat16 before multiplying drifted by several ulps.
   constexpr int64_t num = 600;
   AddInputFromArray<bfloat16>(TensorShape({}), {bfloat16(0.0f)});
   AddInputFromArray<bfloat16>(TensorShape({}), {bfloat16(1.0f)});
@@ -277,10 +278,15 @@ TEST_F(LinSpaceOpTest, Monotonic_Bfloat16_Int64Index) {
 
   const auto flat = GetOutput(0)->flat<bfloat16>();
   ASSERT_EQ(num, GetOutput(0)->NumElements());
-  for (int64_t i = 1; i < num; ++i) {
-    ASSERT_GT(static_cast<float>(flat(i)), static_cast<float>(flat(i - 1)))
-        << "flat(" << i << ") did not exceed flat(" << i - 1
-        << "); indices past 256 plateaued when i was cast to bfloat16";
+  const double step = 1.0 / static_cast<double>(num - 1);
+  for (int64_t i = 1; i < num - 1; ++i) {
+    ASSERT_GE(static_cast<float>(flat(i)), static_cast<float>(flat(i - 1)))
+        << "flat(" << i << ") decreased below flat(" << i - 1 << ")";
+    ASSERT_EQ(static_cast<float>(
+                  static_cast<bfloat16>(step * static_cast<double>(i))),
+              static_cast<float>(flat(i)))
+        << "flat(" << i << ") is not the nearest bfloat16 to " << i << "/"
+        << num - 1;
   }
   EXPECT_EQ(0.0f, static_cast<float>(flat(0)));
   EXPECT_EQ(1.0f, static_cast<float>(flat(num - 1)));

From b91838876893d61246facd0b31d17717cc3b1640 Mon Sep 17 00:00:00 2001
From: Aditya Gaurav <mr.x.l.r.8.pride@gmail.com>
Date: Tue, 29 Sep 2026 23:29:32 +0000
Subject: [PATCH 0099/1000] test(linspace): add Single_Bfloat16 and reword
 monotonic failure message

---
 tensorflow/core/kernels/sequence_ops_test.cc | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/kernels/sequence_ops_test.cc b/tensorflow/core/kernels/sequence_ops_test.cc
index 84d3c413691..51496f211df 100644
--- a/tensorflow/core/kernels/sequence_ops_test.cc
+++ b/tensorflow/core/kernels/sequence_ops_test.cc
@@ -278,10 +278,12 @@ TEST_F(LinSpaceOpTest, Monotonic_Bfloat16_Int64Index) {
 
   const auto flat = GetOutput(0)->flat<bfloat16>();
   ASSERT_EQ(num, GetOutput(0)->NumElements());
+  for (int64_t i = 1; i < num; ++i) {
+    ASSERT_GE(static_cast<float>(flat(i)), static_cast<float>(flat(i - 1)))
+        << "flat(" << i << ") was less than flat(" << i - 1 << ")";
+  }
   const double step = 1.0 / static_cast<double>(num - 1);
   for (int64_t i = 1; i < num - 1; ++i) {
-    ASSERT_GE(static_cast<float>(flat(i)), static_cast<float>(flat(i - 1)))
-        << "flat(" << i << ") decreased below flat(" << i - 1 << ")";
     ASSERT_EQ(static_cast<float>(
                   static_cast<bfloat16>(step * static_cast<double>(i))),
               static_cast<float>(flat(i)))
@@ -306,5 +308,19 @@ TEST_F(LinSpaceOpTest, Single_Half) {
   test::ExpectTensorEqual<Eigen::half>(expected, *GetOutput(0));
 }
 
+TEST_F(LinSpaceOpTest, Single_Bfloat16) {
+  MakeOp(DT_BFLOAT16, DT_INT32);
+
+  // num == 1 must return [start] without dividing by num - 1 == 0.
+  AddInputFromArray<bfloat16>(TensorShape({}), {bfloat16(9.0f)});
+  AddInputFromArray<bfloat16>(TensorShape({}), {bfloat16(100.0f)});
+  AddInputFromArray<int32_t>(TensorShape({}), {1});
+  TF_ASSERT_OK(RunOpKernel());
+
+  Tensor expected(allocator(), DT_BFLOAT16, TensorShape({1}));
+  test::FillValues<bfloat16>(&expected, {bfloat16(9.0f)});
+  test::ExpectTensorEqual<bfloat16>(expected, *GetOutput(0));
+}
+
 }  // namespace
 }  // namespace tensorflow

From 2e76507cff05e72b3432387cd52e2b47ed4f2724 Mon Sep 17 00:00:00 2001
From: greypilgrim-083 <prashantchaudhary7353@gmail.com>
Date: Wed, 30 Sep 2026 10:23:37 +0530
Subject: [PATCH 0100/1000] Fix SlidingWindow dataset type, add SizeOne and
 SignedInt64 tests, align formatting

---
 .../kernels/data/window_dataset_op_test.cc    | 38 +++++++++++++++++--
 1 file changed, 35 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/kernels/data/window_dataset_op_test.cc b/tensorflow/core/kernels/data/window_dataset_op_test.cc
index 72c3d6d0357..1de0e407878 100644
--- a/tensorflow/core/kernels/data/window_dataset_op_test.cc
+++ b/tensorflow/core/kernels/data/window_dataset_op_test.cc
@@ -650,7 +650,7 @@ TEST_F(WindowDatasetOpTest, OverflowingTargetBufferSize) {
   EXPECT_THAT(
       Initialize(dataset_params),
       tsl::testing::StatusIs(absl::StatusCode::kInvalidArgument,
-                                    ::testing::HasSubstr("overflow")));
+                             ::testing::HasSubstr("overflow")));
 }
 
 // Regression test: addition overflow when (window_size - 1) * window_stride + 1
@@ -672,7 +672,39 @@ TEST_F(WindowDatasetOpTest, AddOverflowingTargetBufferSize) {
   EXPECT_THAT(
       Initialize(dataset_params),
       tsl::testing::StatusIs(absl::StatusCode::kInvalidArgument,
-                                    ::testing::HasSubstr("overflow")));
+                             ::testing::HasSubstr("overflow")));
+}
+
+TEST_F(WindowDatasetOpTest, SizeOneWithExtremeStrideDoesNotOverflow) {
+  auto dataset_params = WindowDatasetParams(
+      RangeDatasetParams(0, 3, 1),
+      /*size=*/1,
+      /*shift=*/1,
+      /*stride=*/std::numeric_limits<int64_t>::max(),
+      /*drop_remainder=*/true,
+      /*output_dtypes=*/{DT_VARIANT},
+      /*output_shapes=*/{PartialTensorShape({})},
+      /*node_name=*/kNodeName);
+  TF_EXPECT_OK(Initialize(dataset_params));
+}
+
+WindowDatasetParams WindowDatasetParamsWithSignedInt64Overflow() {
+  return WindowDatasetParams(RangeDatasetParams(0, 3, 1),
+                             /*size=*/3,
+                             /*shift=*/1,
+                             /*stride=*/static_cast<int64_t>(1) << 62,
+                             /*drop_remainder=*/true,
+                             /*output_dtypes=*/{DT_VARIANT},
+                             /*output_shapes=*/{PartialTensorShape({})},
+                             /*node_name=*/kNodeName);
+}
+
+TEST_F(WindowDatasetOpTest, SignedInt64OverflowingTargetBufferSize) {
+  auto dataset_params = WindowDatasetParamsWithSignedInt64Overflow();
+  EXPECT_THAT(
+      Initialize(dataset_params),
+      tsl::testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                             ::testing::HasSubstr("overflow")));
 }
 
 class SlidingWindowDatasetParams : public DatasetParams {
@@ -720,7 +752,7 @@ class SlidingWindowDatasetParams : public DatasetParams {
   }
 
   std::string dataset_type() const override {
-    return "SlidingWindowDataset";
+    return "SlidingWindow";
   }
 
  private:

From 24116b416d3d76c05d69caaa770cd7e253cd02cf Mon Sep 17 00:00:00 2001
From: greypilgrim-083 <prashantchaudhary7353@gmail.com>
Date: Wed, 30 Sep 2026 11:10:06 +0530
Subject: [PATCH 0101/1000] Trigger CI


From 187cc1af06d32eba0f799989b484c5a57e803897 Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Wed, 30 Sep 2026 11:30:22 +0530
Subject: [PATCH 0102/1000] Expand VariantTensorDataReader empty-metadata
 regression test

Cover delimiter-only "@@@@" metadata, assert Contains and ReadTensor for
each malformed case, and add a test that a malformed entry is skipped while
a valid entry deserialized alongside it stays readable.

Signed-off-by: mohammed adib <adib@bugqore.com>
---
 .../core/data/serialization_utils_test.cc     | 26 ++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/data/serialization_utils_test.cc b/tensorflow/core/data/serialization_utils_test.cc
index e6c515e6392..4104d5dfeca 100644
--- a/tensorflow/core/data/serialization_utils_test.cc
+++ b/tensorflow/core/data/serialization_utils_test.cc
@@ -188,19 +188,43 @@ TEST(SerializationUtilsTest, VariantTensorDataReaderOOB) {
 TEST(SerializationUtilsTest, VariantTensorDataReaderEmptyMetadata) {
   // Metadata that splits into no keys (empty or delimiter-only) must not cause
   // the reader to index the first key out of bounds while building its maps.
-  for (const std::string& metadata : {std::string(""), std::string("@@")}) {
+  for (const std::string& metadata :
+       {std::string(""), std::string("@@"), std::string("@@@@")}) {
     VariantTensorData data;
     data.metadata_ = metadata;
     data.tensors_.push_back(Tensor(DT_INT64, {1}));
     std::vector<const VariantTensorData*> reader_data;
     reader_data.push_back(&data);
     VariantTensorDataReader reader(reader_data);
+    EXPECT_FALSE(reader.Contains("Iterator", "key1"));
     int64_t val_int64;
     EXPECT_EQ(reader.ReadScalar("Iterator", "key1", &val_int64).code(),
               error::NOT_FOUND);
+    Tensor val_tensor;
+    EXPECT_EQ(reader.ReadTensor("Iterator", "key1", &val_tensor).code(),
+              error::NOT_FOUND);
   }
 }
 
+TEST(SerializationUtilsTest, VariantTensorDataReaderSkipsOnlyMalformedEntries) {
+  // A malformed entry must be skipped without affecting valid entries that are
+  // deserialized alongside it.
+  VariantTensorData malformed;
+  malformed.metadata_ = "@@";
+  malformed.tensors_.push_back(Tensor(DT_INT64, {1}));
+  VariantTensorData valid;
+  valid.metadata_ = "Iterator@@key1";
+  Tensor value(DT_INT64, {});
+  value.scalar<int64_t>()() = 42;
+  valid.tensors_.push_back(value);
+  std::vector<const VariantTensorData*> reader_data = {&malformed, &valid};
+  VariantTensorDataReader reader(reader_data);
+  EXPECT_TRUE(reader.Contains("Iterator", "key1"));
+  int64_t val_int64;
+  TF_ASSERT_OK(reader.ReadScalar("Iterator", "key1", &val_int64));
+  EXPECT_EQ(val_int64, 42);
+}
+
 class ParameterizedIteratorStateVariantTest
     : public DatasetOpsTestBase,
       public ::testing::WithParamInterface<std::vector<Tensor>> {

From f3b0b4ed0e2f0ce1030ed002a09ad60fc51319d9 Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Wed, 30 Sep 2026 11:43:49 +0530
Subject: [PATCH 0103/1000] Validate element_shape_proto and reject concrete
 elements under DT_INVALID in TensorList::Decode

Signed-off-by: mohammed adib <adib@bugqore.com>
---
 tensorflow/core/kernels/tensor_list.cc      | 10 ++++--
 tensorflow/core/kernels/tensor_list_test.cc | 39 +++++++++++++++++++++
 2 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/kernels/tensor_list.cc b/tensorflow/core/kernels/tensor_list.cc
index 8f49f0cd33f..cd4c110c8c6 100644
--- a/tensorflow/core/kernels/tensor_list.cc
+++ b/tensorflow/core/kernels/tensor_list.cc
@@ -124,6 +124,10 @@ bool TensorList::Decode(const VariantTensorData& data) {
 
   TensorShapeProto element_shape_proto;
   if (!element_shape_proto.ParseFromString(iter)) return false;
+  // The PartialTensorShape constructor calls AddDim(), which fatally CHECKs on a
+  // malformed rank or dimension size. The proto is untrusted, so validate it
+  // first and return false instead of aborting the process.
+  if (!PartialTensorShape::IsValid(element_shape_proto)) return false;
 
   const PartialTensorShape decoded_element_shape(element_shape_proto);
 
@@ -137,8 +141,10 @@ bool TensorList::Decode(const VariantTensorData& data) {
   // overrun the output buffer.
   for (const Tensor& t : decoded_tensors) {
     if (t.dtype() == DT_INVALID) continue;  // Unset element placeholder.
-    if (decoded_element_dtype != DT_INVALID &&
-        t.dtype() != decoded_element_dtype) {
+    // Any element reaching here is a concrete tensor, so it must match the
+    // list's dtype. If decoded_element_dtype is DT_INVALID a concrete element is
+    // itself inconsistent, so reject it rather than letting it through.
+    if (t.dtype() != decoded_element_dtype) {
       return false;
     }
     if (!decoded_element_shape.IsCompatibleWith(t.shape())) {
diff --git a/tensorflow/core/kernels/tensor_list_test.cc b/tensorflow/core/kernels/tensor_list_test.cc
index bafd0fa7935..fdc90bc2b0c 100644
--- a/tensorflow/core/kernels/tensor_list_test.cc
+++ b/tensorflow/core/kernels/tensor_list_test.cc
@@ -182,6 +182,45 @@ TEST(TensorListTest, DecodeAcceptsElementCompatibleWithPartialShape) {
   EXPECT_TRUE(TensorList().Decode(data));
 }
 
+TEST(TensorListTest, DecodeRejectsMalformedElementShapeProto) {
+  // A dimension size below -1 is invalid. The PartialTensorShape constructor
+  // would fatally CHECK on it, so Decode must reject it and return false rather
+  // than aborting the process.
+  TensorShapeProto shape;
+  shape.add_dim()->set_size(-2);
+  VariantTensorData data;
+  data.set_metadata(TensorListMetadata(
+      /*invalid_indices=*/{}, static_cast<uint64_t>(DT_FLOAT),
+      std::numeric_limits<uint64_t>::max(), shape.SerializeAsString()));
+  EXPECT_FALSE(TensorList().Decode(data));
+}
+
+TEST(TensorListTest, DecodeRejectsConcreteElementWhenDtypeInvalid) {
+  // element_dtype is DT_INVALID but a concrete tensor is stored. A concrete
+  // element cannot belong to a list without a dtype, so it must be rejected.
+  TensorShapeProto shape;
+  VariantTensorData data;
+  data.set_metadata(TensorListMetadata(
+      /*invalid_indices=*/{}, static_cast<uint64_t>(DT_INVALID),
+      std::numeric_limits<uint64_t>::max(), shape.SerializeAsString()));
+  *data.add_tensors() = ScalarInt32Tensor(1);
+  EXPECT_FALSE(TensorList().Decode(data));
+}
+
+TEST(TensorListTest, DecodeRejectsShapeRankMismatchWithMatchingCount) {
+  // element_shape [2, 3] and element shape [6] have the same element count but
+  // different rank, so they are not compatible and must be rejected.
+  TensorShapeProto shape;
+  shape.add_dim()->set_size(2);
+  shape.add_dim()->set_size(3);
+  VariantTensorData data;
+  data.set_metadata(TensorListMetadata(
+      /*invalid_indices=*/{}, static_cast<uint64_t>(DT_FLOAT),
+      std::numeric_limits<uint64_t>::max(), shape.SerializeAsString()));
+  *data.add_tensors() = Tensor(DT_FLOAT, TensorShape({6}));
+  EXPECT_FALSE(TensorList().Decode(data));
+}
+
 TEST(TensorListTest, DecodeRejectsOutOfRangeMaxNumElements) {
   TensorShapeProto shape;
   VariantTensorData data;

From 1c832e34eabd2b20398be3b12cf879b1fb8656f7 Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Wed, 30 Sep 2026 03:29:03 -0400
Subject: [PATCH 0104/1000] security/fuzzing: verify the model before building
 it in the TFLite fuzz target

Puts the harness on the supported ingress path for untrusted models: call
tflite::Verify() first and stop on rejection, then build and invoke. The
verifier owns the bounds checks on model metadata (tensorflow/tensorflow
#116631 review, Sep 18; checks landed in #123910), so inputs like the
uint64 external-offset wrap are rejected before ParseTensors, and the
verifier itself gets fuzz coverage. No change to interpreter_builder.cc.

VerifyAndBuildFromBuffer alone would not do this: without an extra_verifier
it runs only the flatbuffers structural check.
---
 .../fuzzing/cc/tflite_interp_overflow_fuzz.cc | 27 ++++++++++++-------
 1 file changed, 18 insertions(+), 9 deletions(-)

diff --git a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
index 4650f2c3853..bf51c157ea7 100644
--- a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
+++ b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
@@ -13,13 +13,14 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-// FuzzTest coverage for TFLite model loading, interpreter build, and invoke.
-// This exercises interpreter_builder.cc get_readonly_data / ParseTensors, which
-// perform an unchecked uint64 addition (offset + size > allocation_->bytes())
-// on attacker-controlled external-buffer offsets. On overflow the check is
-// bypassed and a wild pointer (allocation_->base() + offset) is used as tensor
-// data, which is dereferenced during Invoke. The OSS-Fuzz tensorflow build has
-// no TFLite target that reaches this path.
+// FuzzTest coverage for the TFLite untrusted-model ingress path: model
+// verification, model load, interpreter build, and invoke. tflite::Verify() is
+// the validation layer for untrusted models, so the target runs it first and
+// stops on rejection; everything it accepts goes on to ParseTensors and the
+// kernels. This covers the verifier's checks on attacker-controlled model
+// metadata (including the uint64 external-offset bounds checks in
+// VerifyTensors / VerifyOperators) and the runtime behind them. The OSS-Fuzz
+// tensorflow build has no TFLite target that reaches this path.
 
 #include <cstring>
 #include <memory>
@@ -30,10 +31,19 @@ limitations under the License.
 #include "tensorflow/lite/interpreter_builder.h"
 #include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/model_builder.h"
+#include "tensorflow/lite/tools/verifier.h"
 
 namespace {
 
 void FuzzModelBuildAndInvoke(const std::string& model_bytes) {
+  // Supported ingress path: an untrusted model is verified before anything is
+  // built from it. Inputs the verifier rejects stop here. The error reporter
+  // is optional and would only add stderr noise under the fuzzer.
+  if (!tflite::Verify(model_bytes.data(), model_bytes.size(),
+                      /*error_reporter=*/nullptr)) {
+    return;
+  }
+
   auto model = tflite::FlatBufferModel::BuildFromBuffer(model_bytes.data(),
                                                         model_bytes.size());
   if (model == nullptr) return;
@@ -54,8 +64,7 @@ void FuzzModelBuildAndInvoke(const std::string& model_bytes) {
     }
   }
 
-  // The unchecked external-offset arithmetic sets a wild read-only tensor
-  // pointer that is dereferenced here.
+  // Runs the kernels on whatever the verifier accepted.
   interpreter->Invoke();
 }
 FUZZ_TEST(TfliteInterpreterBuilder, FuzzModelBuildAndInvoke);

From 0111afab059556333fd8514eb5a833007636db66 Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Wed, 30 Sep 2026 03:29:41 -0400
Subject: [PATCH 0105/1000] security/fuzzing: depend on
 //tensorflow/lite/tools:verifier for the TFLite fuzz target

The target now calls tflite::Verify() before building the model.
---
 tensorflow/security/fuzzing/cc/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/security/fuzzing/cc/BUILD b/tensorflow/security/fuzzing/cc/BUILD
index afd5e7af76c..26c767712d8 100644
--- a/tensorflow/security/fuzzing/cc/BUILD
+++ b/tensorflow/security/fuzzing/cc/BUILD
@@ -260,5 +260,6 @@ tf_cc_fuzz_test(
     deps = [
         "//tensorflow/lite:framework",
         "//tensorflow/lite/kernels:builtin_ops",
+        "//tensorflow/lite/tools:verifier",
     ],
 )

From 0ba71e95392b3430681427a8efa52ca1bbbcdfaf Mon Sep 17 00:00:00 2001
From: Aditya Gaurav <mr.x.l.r.8.pride@gmail.com>
Date: Wed, 30 Sep 2026 12:42:00 +0000
Subject: [PATCH 0106/1000] test(linspace): cover half and bfloat16 LinSpace
 from Python, include bfloat16.h

---
 tensorflow/core/kernels/sequence_ops.cc       |  1 +
 tensorflow/core/kernels/sequence_ops_test.cc  |  1 +
 tensorflow/python/ops/BUILD                   |  1 +
 .../python/ops/math_ops_linspace_test.py      | 31 +++++++++++++++++++
 4 files changed, 34 insertions(+)

diff --git a/tensorflow/core/kernels/sequence_ops.cc b/tensorflow/core/kernels/sequence_ops.cc
index 266ae18c696..a8fec314c4a 100644
--- a/tensorflow/core/kernels/sequence_ops.cc
+++ b/tensorflow/core/kernels/sequence_ops.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/platform/bfloat16.h"
 
 namespace tensorflow {
 
diff --git a/tensorflow/core/kernels/sequence_ops_test.cc b/tensorflow/core/kernels/sequence_ops_test.cc
index 51496f211df..46e94b64626 100644
--- a/tensorflow/core/kernels/sequence_ops_test.cc
+++ b/tensorflow/core/kernels/sequence_ops_test.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include "tensorflow/core/kernels/ops_testutil.h"
 #include "tensorflow/core/kernels/ops_util.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/bfloat16.h"
 #include "tensorflow/core/platform/test.h"
 
 namespace tensorflow {
diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD
index 22b37b8bd77..6d53655e8a0 100644
--- a/tensorflow/python/ops/BUILD
+++ b/tensorflow/python/ops/BUILD
@@ -3670,6 +3670,7 @@ cuda_py_strict_test(
     ],
     deps = [
         ":math_ops",
+        ":math_ops_gen",
         "//tensorflow/python/eager:def_function",
         "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/framework:tensor",
diff --git a/tensorflow/python/ops/math_ops_linspace_test.py b/tensorflow/python/ops/math_ops_linspace_test.py
index 66cd91caaa9..4e034d3501c 100644
--- a/tensorflow/python/ops/math_ops_linspace_test.py
+++ b/tensorflow/python/ops/math_ops_linspace_test.py
@@ -22,6 +22,7 @@ from tensorflow.python.eager import def_function
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import tensor
 from tensorflow.python.framework import test_util
+from tensorflow.python.ops import gen_math_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.platform import googletest
 
@@ -73,6 +74,36 @@ class LinspaceTest(test_util.TensorFlowTestCase, parameterized.TestCase):
     expected_shape = (64, None, 10)
     self.assertEqual(output_shape, expected_shape)
 
+  @parameterized.parameters([
+      (dtype, num_dtype)
+      for dtype in [dtypes.float16, dtypes.bfloat16]
+      for num_dtype in [np.int32, np.int64]
+  ])
+  def testLinSpaceKernelHalfAndBfloat16(self, dtype, num_dtype):
+    # tf.linspace is composed from other ops, so call the LinSpace kernel
+    # directly. Each element must be the float64 value rounded once to dtype;
+    # num=70000 is past float16's largest finite value (65504).
+    np_dtype = dtype.as_numpy_dtype
+    for start, stop, num in [(0.0, 1.0, 600), (-3.0, 5.0, 20), (9.0, 100.0, 1),
+                             (0.0, 1.0, 70000)]:
+      actual = self.evaluate(
+          gen_math_ops.lin_space(
+              np.array(start, np_dtype), np.array(stop, np_dtype),
+              np.array(num, num_dtype)))
+      expected = np.linspace(start, stop, num, dtype=np.float64).astype(np_dtype)
+      self.assertEqual(np_dtype, actual.dtype)
+      self.assertAllEqual(expected.astype(np.float32),
+                          actual.astype(np.float32))
+
+  @parameterized.parameters([dtypes.float16, dtypes.bfloat16])
+  def testLinspaceHalfAndBfloat16(self, dtype):
+    # API coverage for tf.linspace; it does not reach the LinSpace kernel.
+    np_dtype = dtype.as_numpy_dtype
+    actual = self.evaluate(
+        math_ops.linspace(np.array(1.0, np_dtype), np.array(10.0, np_dtype), 4))
+    self.assertEqual(np_dtype, actual.dtype)
+    self.assertAllEqual([1.0, 4.0, 7.0, 10.0], actual.astype(np.float32))
+
 
 if __name__ == "__main__":
   googletest.main()

From 6a22764878b3d2fb3bae5383b4deabf1aebdd37a Mon Sep 17 00:00:00 2001
From: Aditya Gaurav <mr.x.l.r.8.pride@gmail.com>
Date: Wed, 30 Sep 2026 12:43:24 +0000
Subject: [PATCH 0107/1000] test(linspace): narrow expected LinSpace values
 through float32 like the kernel cast

---
 tensorflow/python/ops/math_ops_linspace_test.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/ops/math_ops_linspace_test.py b/tensorflow/python/ops/math_ops_linspace_test.py
index 4e034d3501c..f58b2c4491b 100644
--- a/tensorflow/python/ops/math_ops_linspace_test.py
+++ b/tensorflow/python/ops/math_ops_linspace_test.py
@@ -81,8 +81,9 @@ class LinspaceTest(test_util.TensorFlowTestCase, parameterized.TestCase):
   ])
   def testLinSpaceKernelHalfAndBfloat16(self, dtype, num_dtype):
     # tf.linspace is composed from other ops, so call the LinSpace kernel
-    # directly. Each element must be the float64 value rounded once to dtype;
-    # num=70000 is past float16's largest finite value (65504).
+    # directly. Each element must be the float64 value narrowed to dtype the
+    # way the kernel's cast does it (through float32); num=70000 is past
+    # float16's largest finite value (65504).
     np_dtype = dtype.as_numpy_dtype
     for start, stop, num in [(0.0, 1.0, 600), (-3.0, 5.0, 20), (9.0, 100.0, 1),
                              (0.0, 1.0, 70000)]:
@@ -90,7 +91,8 @@ class LinspaceTest(test_util.TensorFlowTestCase, parameterized.TestCase):
           gen_math_ops.lin_space(
               np.array(start, np_dtype), np.array(stop, np_dtype),
               np.array(num, num_dtype)))
-      expected = np.linspace(start, stop, num, dtype=np.float64).astype(np_dtype)
+      expected = np.linspace(
+          start, stop, num, dtype=np.float64).astype(np.float32).astype(np_dtype)
       self.assertEqual(np_dtype, actual.dtype)
       self.assertAllEqual(expected.astype(np.float32),
                           actual.astype(np.float32))

From fa46d76196bcc27fb4fcb0e793e3a8b7884092ea Mon Sep 17 00:00:00 2001
From: Aditya Gaurav <mr.x.l.r.8.pride@gmail.com>
Date: Wed, 30 Sep 2026 14:41:31 +0000
Subject: [PATCH 0108/1000] fix(sequence_ops): widen only half and bfloat16 in
 LinSpace and Range

---
 tensorflow/core/kernels/sequence_ops.cc      | 35 ++++++++++++++------
 tensorflow/core/kernels/sequence_ops_test.cc | 33 +++++++++++++++++-
 2 files changed, 56 insertions(+), 12 deletions(-)

diff --git a/tensorflow/core/kernels/sequence_ops.cc b/tensorflow/core/kernels/sequence_ops.cc
index a8fec314c4a..5de7a6dce50 100644
--- a/tensorflow/core/kernels/sequence_ops.cc
+++ b/tensorflow/core/kernels/sequence_ops.cc
@@ -20,13 +20,14 @@ limitations under the License.
 #include <cmath>
 #include <type_traits>
 
+#include "tensorflow/core/framework/bfloat16.h"
+#include "tensorflow/core/framework/numeric_types.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/op_requires.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.h"
-#include "tensorflow/core/platform/bfloat16.h"
 
 namespace tensorflow {
 
@@ -40,8 +41,19 @@ struct RangeFunctor<CPUDevice, T> {
   void operator()(OpKernelContext* context, int64_t size, T start, T delta,
                   typename TTypes<T>::Flat output) const {
     (void)context;
-    for (int64_t i = 0; i < size; ++i) {
-      output(i) = start + static_cast<T>(i) * delta;
+    if constexpr (std::is_integral<T>::value) {
+      for (int64_t i = 0; i < size; ++i) {
+        output(i) = start + static_cast<T>(i) * delta;
+      }
+    } else {
+      using ComputeT = typename std::conditional<sizeof(T) >= sizeof(float), T,
+                                                 double>::type;
+      const ComputeT start_c = static_cast<ComputeT>(start);
+      const ComputeT delta_c = static_cast<ComputeT>(delta);
+      for (int64_t i = 0; i < size; ++i) {
+        output(i) =
+            static_cast<T>(start_c + static_cast<ComputeT>(i) * delta_c);
+      }
     }
   }
 };
@@ -207,15 +219,16 @@ class LinSpaceOp : public OpKernel {
     auto flat = out->flat<T>();
     flat(0) = start;
     if (num > 1) {
-      // Interpolate in double, narrow once on store. In T, (stop - start) and
-      // the index casts both overflow to +Inf for half/bfloat16, which zeroes
-      // step and turns every interior element into NaN (0 * +Inf). The single
-      // narrowing cast is also what keeps this portable across MSVC/GCC/Clang.
-      const double start_d = static_cast<double>(start);
-      const double stop_d = static_cast<double>(stop);
-      const double step = (stop_d - start_d) / static_cast<double>(num - 1);
+      // Interpolate in a higher precision type to avoid overflow for
+      // half/bfloat16, while preserving native performance and numerical
+      // semantics for float/double.
+      using ComputeT = typename std::conditional<sizeof(T) >= sizeof(float), T,
+                                                 double>::type;
+      const ComputeT start_c = static_cast<ComputeT>(start);
+      const ComputeT stop_c = static_cast<ComputeT>(stop);
+      const ComputeT step = (stop_c - start_c) / static_cast<ComputeT>(num - 1);
       for (Tnum i = 1; i < num - 1; ++i) {
-        flat(i) = static_cast<T>(start_d + step * static_cast<double>(i));
+        flat(i) = static_cast<T>(start_c + step * static_cast<ComputeT>(i));
       }
       // Ensure final value == stop; float arithmetic won't guarantee this.
       flat(num - 1) = stop;
diff --git a/tensorflow/core/kernels/sequence_ops_test.cc b/tensorflow/core/kernels/sequence_ops_test.cc
index 46e94b64626..8d96dd82f7a 100644
--- a/tensorflow/core/kernels/sequence_ops_test.cc
+++ b/tensorflow/core/kernels/sequence_ops_test.cc
@@ -16,8 +16,10 @@ limitations under the License.
 #include <cmath>
 
 #include "tensorflow/core/framework/allocator.h"
+#include "tensorflow/core/framework/bfloat16.h"
 #include "tensorflow/core/framework/fake_input.h"
 #include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/framework/numeric_types.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
@@ -26,7 +28,6 @@ limitations under the License.
 #include "tensorflow/core/kernels/ops_testutil.h"
 #include "tensorflow/core/kernels/ops_util.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
-#include "tensorflow/core/platform/bfloat16.h"
 #include "tensorflow/core/platform/test.h"
 
 namespace tensorflow {
@@ -130,6 +131,36 @@ TEST_F(RangeOpTest, Range_Size_Overflow) {
             RunOpKernel().message());
 }
 
+TEST_F(RangeOpTest, Monotonic_Bfloat16) {
+  MakeOp(DT_BFLOAT16);
+
+  // 600 elements go past 256, where bfloat16 stops holding every integer, so
+  // casting i to bfloat16 before multiplying rounded it. delta is smaller than
+  // the bfloat16 spacing near the top of the range, so neighbouring outputs
+  // may be equal; each one must still be start + i * delta rounded once.
+  const bfloat16 start(0.0f);
+  const bfloat16 delta(0.3f);
+  AddInputFromArray<bfloat16>(TensorShape({}), {start});
+  AddInputFromArray<bfloat16>(TensorShape({}), {bfloat16(180.0f)});
+  AddInputFromArray<bfloat16>(TensorShape({}), {delta});
+  TF_ASSERT_OK(RunOpKernel());
+
+  const auto flat = GetOutput(0)->flat<bfloat16>();
+  constexpr int64_t num = 600;
+  ASSERT_EQ(num, GetOutput(0)->NumElements());
+  for (int64_t i = 1; i < num; ++i) {
+    ASSERT_GE(static_cast<float>(flat(i)), static_cast<float>(flat(i - 1)))
+        << "flat(" << i << ") was less than flat(" << i - 1 << ")";
+  }
+  for (int64_t i = 0; i < num; ++i) {
+    ASSERT_EQ(static_cast<float>(static_cast<bfloat16>(
+                  static_cast<double>(start) +
+                  static_cast<double>(i) * static_cast<double>(delta))),
+              static_cast<float>(flat(i)))
+        << "flat(" << i << ") is not start + " << i << " * delta";
+  }
+}
+
 TEST_F(LinSpaceOpTest, Simple_D32) {
   MakeOp(DT_FLOAT, DT_INT32);
 

From ce950b2db5078452823f546b6247d32934d4a432 Mon Sep 17 00:00:00 2001
From: greypilgrim-083 <prashantchaudhary7353@gmail.com>
Date: Wed, 30 Sep 2026 20:58:08 +0530
Subject: [PATCH 0109/1000] Extract SlidingWindow tests to dedicated file, fix
 core->experimental dep

---
 tensorflow/core/kernels/data/BUILD            |   1 -
 .../core/kernels/data/experimental/BUILD      |  19 +++
 .../sliding_window_dataset_op_test.cc         | 147 ++++++++++++++++++
 .../kernels/data/window_dataset_op_test.cc    |  97 ------------
 4 files changed, 166 insertions(+), 98 deletions(-)
 create mode 100644 tensorflow/core/kernels/data/experimental/sliding_window_dataset_op_test.cc

diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD
index db3aa6b5b32..0e3031c377c 100644
--- a/tensorflow/core/kernels/data/BUILD
+++ b/tensorflow/core/kernels/data/BUILD
@@ -1658,7 +1658,6 @@ tf_cc_test(
         ":iterator_ops",
         ":range_dataset_op",
         ":window_dataset_op",
-        "//tensorflow/core/kernels/data/experimental:sliding_window_dataset_op",
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:dataset_ops_op_lib",
         "//tensorflow/core:framework",
diff --git a/tensorflow/core/kernels/data/experimental/BUILD b/tensorflow/core/kernels/data/experimental/BUILD
index 2609ac540cb..532d99074d1 100644
--- a/tensorflow/core/kernels/data/experimental/BUILD
+++ b/tensorflow/core/kernels/data/experimental/BUILD
@@ -854,6 +854,25 @@ tf_kernel_library(
     ],
 )
 
+tf_cc_test(
+    name = "sliding_window_dataset_op_test",
+    size = "small",
+    srcs = ["sliding_window_dataset_op_test.cc"],
+    deps = [
+        ":sliding_window_dataset_op",
+        "//tensorflow/core:experimental_dataset_ops_op_lib",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/data:dataset_test_base",
+        "//tensorflow/core/data:name_utils",
+        "//tensorflow/core/kernels/data:range_dataset_op",
+        "//tensorflow/core/platform:status",
+        "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest_main",
+        "@xla//xla/tsl/lib/core:status_test_util",
+        "@xla//xla/tsl/platform:status_matchers",
+    ],
+)
+
 tf_kernel_library(
     name = "snapshot_dataset_op",
     srcs = ["snapshot_dataset_op.cc"],
diff --git a/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op_test.cc b/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op_test.cc
new file mode 100644
index 00000000000..b41c1a9e687
--- /dev/null
+++ b/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op_test.cc
@@ -0,0 +1,147 @@
+/* Copyright 2024 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.
+==============================================================================*/
+#include <cstdint>
+#include <limits>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/platform/status_matchers.h"
+#include "tensorflow/core/data/dataset_test_base.h"
+#include "tensorflow/core/data/name_utils.h"
+#include "tensorflow/core/framework/dataset.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/types.h"
+
+namespace tensorflow {
+namespace data {
+namespace experimental {
+namespace {
+
+constexpr char kNodeName[] = "sliding_window_dataset";
+
+class SlidingWindowDatasetParams : public DatasetParams {
+ public:
+  template <typename T>
+  SlidingWindowDatasetParams(T input_dataset_params, int64_t size,
+                             int64_t shift, int64_t stride, bool drop_remainder,
+                             DataTypeVector output_dtypes,
+                             std::vector<PartialTensorShape> output_shapes,
+                             std::string node_name)
+      : DatasetParams(std::move(output_dtypes), std::move(output_shapes),
+                      std::move(node_name)),
+        size_(size),
+        shift_(shift),
+        stride_(stride),
+        drop_remainder_(drop_remainder) {
+    input_dataset_params_.push_back(std::make_unique<T>(input_dataset_params));
+    iterator_prefix_ =
+        name_utils::IteratorPrefix(input_dataset_params.dataset_type(),
+                                   input_dataset_params.iterator_prefix());
+  }
+
+  std::vector<Tensor> GetInputTensors() const override {
+    return {CreateTensor<int64_t>(TensorShape({}), {size_}),
+            CreateTensor<int64_t>(TensorShape({}), {shift_}),
+            CreateTensor<int64_t>(TensorShape({}), {stride_})};
+  }
+
+  absl::Status GetInputNames(
+      std::vector<std::string>* input_names) const override {
+    input_names->clear();
+    input_names->emplace_back("input_dataset");
+    input_names->emplace_back("window_size");
+    input_names->emplace_back("window_shift");
+    input_names->emplace_back("window_stride");
+    return absl::OkStatus();
+  }
+
+  absl::Status GetAttributes(AttributeVector* attr_vector) const override {
+    attr_vector->clear();
+    attr_vector->emplace_back("output_types", output_dtypes_);
+    attr_vector->emplace_back("output_shapes", output_shapes_);
+    attr_vector->emplace_back("drop_remainder", drop_remainder_);
+    return absl::OkStatus();
+  }
+
+  std::string dataset_type() const override { return "SlidingWindow"; }
+
+ private:
+  int64_t size_;
+  int64_t shift_;
+  int64_t stride_;
+  bool drop_remainder_;
+};
+
+class SlidingWindowDatasetOpTest : public DatasetOpsTestBase {};
+
+SlidingWindowDatasetParams SlidingWindowDatasetParamsWithOverflowingSizeAndStride() {
+  const int64_t large_size = (static_cast<int64_t>(1) << 32) + 1;
+  const int64_t large_stride = (static_cast<int64_t>(1) << 32);
+  return SlidingWindowDatasetParams(RangeDatasetParams(0, 3, 1),
+                                    /*size=*/large_size,
+                                    /*shift=*/1,
+                                    /*stride=*/large_stride,
+                                    /*drop_remainder=*/true,
+                                    /*output_dtypes=*/{DT_VARIANT},
+                                    /*output_shapes=*/{PartialTensorShape({})},
+                                    /*node_name=*/kNodeName);
+}
+
+SlidingWindowDatasetParams SlidingWindowDatasetParamsWithAddOverflowingSizeAndStride() {
+  const int64_t max_stride = std::numeric_limits<int64_t>::max();
+  return SlidingWindowDatasetParams(RangeDatasetParams(0, 3, 1),
+                                    /*size=*/2,
+                                    /*shift=*/1,
+                                    /*stride=*/max_stride,
+                                    /*drop_remainder=*/true,
+                                    /*output_dtypes=*/{DT_VARIANT},
+                                    /*output_shapes=*/{PartialTensorShape({})},
+                                    /*node_name=*/kNodeName);
+}
+
+// Regression test: large window_size and window_stride whose target buffer
+// size (window_size - 1) * window_stride + 1 overflows int64. Before the fix,
+// the overflow silently wrapped to a small value and emitted undersized
+// windows. After the fix, dataset initialization must return InvalidArgument.
+TEST_F(SlidingWindowDatasetOpTest, OverflowingTargetBufferSize) {
+  auto dataset_params = SlidingWindowDatasetParamsWithOverflowingSizeAndStride();
+  EXPECT_THAT(
+      Initialize(dataset_params),
+      tsl::testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                             ::testing::HasSubstr("overflow")));
+}
+
+// Regression test: addition overflow when (window_size - 1) * window_stride + 1
+// overflows int64 on the addition step with size=2 and stride=int64_max.
+TEST_F(SlidingWindowDatasetOpTest, AddOverflowingTargetBufferSize) {
+  auto dataset_params = SlidingWindowDatasetParamsWithAddOverflowingSizeAndStride();
+  EXPECT_THAT(
+      Initialize(dataset_params),
+      tsl::testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                             ::testing::HasSubstr("overflow")));
+}
+
+}  // namespace
+}  // namespace experimental
+}  // namespace data
+}  // namespace tensorflow
diff --git a/tensorflow/core/kernels/data/window_dataset_op_test.cc b/tensorflow/core/kernels/data/window_dataset_op_test.cc
index 1de0e407878..51df124668e 100644
--- a/tensorflow/core/kernels/data/window_dataset_op_test.cc
+++ b/tensorflow/core/kernels/data/window_dataset_op_test.cc
@@ -707,103 +707,6 @@ TEST_F(WindowDatasetOpTest, SignedInt64OverflowingTargetBufferSize) {
                              ::testing::HasSubstr("overflow")));
 }
 
-class SlidingWindowDatasetParams : public DatasetParams {
- public:
-  template <typename T>
-  SlidingWindowDatasetParams(T input_dataset_params, int64_t size,
-                             int64_t shift, int64_t stride, bool drop_remainder,
-                             DataTypeVector output_dtypes,
-                             std::vector<PartialTensorShape> output_shapes,
-                             std::string node_name)
-      : DatasetParams(std::move(output_dtypes), std::move(output_shapes),
-                      std::move(node_name)),
-        size_(size),
-        shift_(shift),
-        stride_(stride),
-        drop_remainder_(drop_remainder) {
-    input_dataset_params_.push_back(std::make_unique<T>(input_dataset_params));
-    iterator_prefix_ =
-        name_utils::IteratorPrefix(input_dataset_params.dataset_type(),
-                                   input_dataset_params.iterator_prefix());
-  }
-
-  std::vector<Tensor> GetInputTensors() const override {
-    return {CreateTensor<int64_t>(TensorShape({}), {size_}),
-            CreateTensor<int64_t>(TensorShape({}), {shift_}),
-            CreateTensor<int64_t>(TensorShape({}), {stride_})};
-  }
-
-  absl::Status GetInputNames(
-      std::vector<std::string>* input_names) const override {
-    input_names->clear();
-    input_names->emplace_back("input_dataset");
-    input_names->emplace_back("window_size");
-    input_names->emplace_back("window_shift");
-    input_names->emplace_back("window_stride");
-    return absl::OkStatus();
-  }
-
-  absl::Status GetAttributes(AttributeVector* attr_vector) const override {
-    attr_vector->clear();
-    attr_vector->emplace_back("output_types", output_dtypes_);
-    attr_vector->emplace_back("output_shapes", output_shapes_);
-    attr_vector->emplace_back("drop_remainder", drop_remainder_);
-    return absl::OkStatus();
-  }
-
-  std::string dataset_type() const override {
-    return "SlidingWindow";
-  }
-
- private:
-  int64_t size_;
-  int64_t shift_;
-  int64_t stride_;
-  bool drop_remainder_;
-};
-
-SlidingWindowDatasetParams SlidingWindowDatasetParamsWithOverflowingSizeAndStride() {
-  const int64_t large_size = (static_cast<int64_t>(1) << 32) + 1;
-  const int64_t large_stride = (static_cast<int64_t>(1) << 32);
-  return SlidingWindowDatasetParams(RangeDatasetParams(0, 3, 1),
-                                    /*size=*/large_size,
-                                    /*shift=*/1,
-                                    /*stride=*/large_stride,
-                                    /*drop_remainder=*/true,
-                                    /*output_dtypes=*/{DT_VARIANT},
-                                    /*output_shapes=*/{PartialTensorShape({})},
-                                    /*node_name=*/"sliding_window_dataset");
-}
-
-SlidingWindowDatasetParams SlidingWindowDatasetParamsWithAddOverflowingSizeAndStride() {
-  const int64_t max_stride = std::numeric_limits<int64_t>::max();
-  return SlidingWindowDatasetParams(RangeDatasetParams(0, 3, 1),
-                                    /*size=*/2,
-                                    /*shift=*/1,
-                                    /*stride=*/max_stride,
-                                    /*drop_remainder=*/true,
-                                    /*output_dtypes=*/{DT_VARIANT},
-                                    /*output_shapes=*/{PartialTensorShape({})},
-                                    /*node_name=*/"sliding_window_dataset");
-}
-
-TEST_F(WindowDatasetOpTest, SlidingWindowOverflowingTargetBufferSize) {
-  auto dataset_params = SlidingWindowDatasetParamsWithOverflowingSizeAndStride();
-  EXPECT_THAT(
-      Initialize(dataset_params),
-      tsl::testing::StatusIs(absl::StatusCode::kInvalidArgument,
-                             ::testing::HasSubstr("overflow")));
-}
-
-TEST_F(WindowDatasetOpTest, SlidingWindowAddOverflowingTargetBufferSize) {
-  auto dataset_params = SlidingWindowDatasetParamsWithAddOverflowingSizeAndStride();
-  EXPECT_THAT(
-      Initialize(dataset_params),
-      tsl::testing::StatusIs(absl::StatusCode::kInvalidArgument,
-                             ::testing::HasSubstr("overflow")));
-}
-
 }  // namespace
 }  // namespace data
 }  // namespace tensorflow
-

From 14a3a1890b1796ee028ee3b45f098f4aaf967f07 Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Thu, 1 Oct 2026 12:20:16 +0530
Subject: [PATCH 0110/1000] guard null VariantTensorData pointers and cover
 null/singleton metadata

Signed-off-by: mohammed adib <adib@bugqore.com>
---
 tensorflow/core/data/serialization_utils.cc   |  4 +++
 .../core/data/serialization_utils_test.cc     | 34 +++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/tensorflow/core/data/serialization_utils.cc b/tensorflow/core/data/serialization_utils.cc
index d41511e3911..551652fcf78 100644
--- a/tensorflow/core/data/serialization_utils.cc
+++ b/tensorflow/core/data/serialization_utils.cc
@@ -185,6 +185,10 @@ absl::Status UpdateCheckpointElements(
 VariantTensorDataReader::VariantTensorDataReader(
     const std::vector<const tensorflow::VariantTensorData*>& data) {
   for (const auto& d : data) {
+    if (d == nullptr) {
+      VLOG(1) << "Found null VariantTensorData pointer; skipping.";
+      continue;
+    }
     std::string metadata;
     d->get_metadata(&metadata);
     auto keys = str_util::Split(metadata, kDelimiter, str_util::SkipEmpty());
diff --git a/tensorflow/core/data/serialization_utils_test.cc b/tensorflow/core/data/serialization_utils_test.cc
index 4104d5dfeca..270c5bf3aa8 100644
--- a/tensorflow/core/data/serialization_utils_test.cc
+++ b/tensorflow/core/data/serialization_utils_test.cc
@@ -225,6 +225,40 @@ TEST(SerializationUtilsTest, VariantTensorDataReaderSkipsOnlyMalformedEntries) {
   EXPECT_EQ(val_int64, 42);
 }
 
+TEST(SerializationUtilsTest, VariantTensorDataReaderSkipsNullEntries) {
+  // A null VariantTensorData pointer must be skipped without dereferencing it,
+  // leaving valid entries deserialized alongside it readable.
+  VariantTensorData valid;
+  valid.metadata_ = "Iterator@@key1";
+  Tensor value(DT_INT64, {});
+  value.scalar<int64_t>()() = 42;
+  valid.tensors_.push_back(value);
+  std::vector<const VariantTensorData*> reader_data = {nullptr, &valid};
+  VariantTensorDataReader reader(reader_data);
+  EXPECT_TRUE(reader.Contains("Iterator", "key1"));
+  int64_t val_int64;
+  TF_ASSERT_OK(reader.ReadScalar("Iterator", "key1", &val_int64));
+  EXPECT_EQ(val_int64, 42);
+}
+
+TEST(SerializationUtilsTest, VariantTensorDataReaderSingletonMetadata) {
+  // Metadata that splits into a single key (a name with no subkeys) must build
+  // an empty key map without indexing past it.
+  VariantTensorData data;
+  data.metadata_ = "Iterator";
+  data.tensors_.push_back(Tensor(DT_INT64, {1}));
+  std::vector<const VariantTensorData*> reader_data;
+  reader_data.push_back(&data);
+  VariantTensorDataReader reader(reader_data);
+  EXPECT_FALSE(reader.Contains("Iterator", "key1"));
+  int64_t val_int64;
+  EXPECT_EQ(reader.ReadScalar("Iterator", "key1", &val_int64).code(),
+            error::NOT_FOUND);
+  Tensor val_tensor;
+  EXPECT_EQ(reader.ReadTensor("Iterator", "key1", &val_tensor).code(),
+            error::NOT_FOUND);
+}
+
 class ParameterizedIteratorStateVariantTest
     : public DatasetOpsTestBase,
       public ::testing::WithParamInterface<std::vector<Tensor>> {

From 4336d9ea0843e86e00d90c4978408722d6ea029e Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Thu, 1 Oct 2026 12:46:21 +0530
Subject: [PATCH 0111/1000] Build element_shape via BuildPartialTensorShape and
 test empty-element mismatch

Signed-off-by: mohammed adib <adib@bugqore.com>
---
 tensorflow/core/kernels/tensor_list.cc      | 16 +++++++++++-----
 tensorflow/core/kernels/tensor_list_test.cc | 14 ++++++++++++++
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/tensorflow/core/kernels/tensor_list.cc b/tensorflow/core/kernels/tensor_list.cc
index cd4c110c8c6..01c16c081b3 100644
--- a/tensorflow/core/kernels/tensor_list.cc
+++ b/tensorflow/core/kernels/tensor_list.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_shape.pb.h"
 #include "tensorflow/core/framework/variant_op_registry.h"
 #include "tensorflow/core/lib/core/coding.h"
+#include "tensorflow/core/platform/logging.h"
 
 namespace tensorflow {
 
@@ -125,11 +126,16 @@ bool TensorList::Decode(const VariantTensorData& data) {
   TensorShapeProto element_shape_proto;
   if (!element_shape_proto.ParseFromString(iter)) return false;
   // The PartialTensorShape constructor calls AddDim(), which fatally CHECKs on a
-  // malformed rank or dimension size. The proto is untrusted, so validate it
-  // first and return false instead of aborting the process.
-  if (!PartialTensorShape::IsValid(element_shape_proto)) return false;
-
-  const PartialTensorShape decoded_element_shape(element_shape_proto);
+  // malformed rank or dimension size. The proto is untrusted, so build the shape
+  // through the factory, which validates and constructs in a single pass and
+  // returns an error instead of aborting the process.
+  PartialTensorShape decoded_element_shape;
+  if (!PartialTensorShape::BuildPartialTensorShape(element_shape_proto,
+                                                   &decoded_element_shape)
+           .ok()) {
+    VLOG(1) << "TensorList::Decode failed: malformed element_shape_proto";
+    return false;
+  }
 
   // Every element stored in a list must match its element_dtype and be
   // compatible with its element_shape; the mutation paths (TensorListPushBack,
diff --git a/tensorflow/core/kernels/tensor_list_test.cc b/tensorflow/core/kernels/tensor_list_test.cc
index fdc90bc2b0c..39eea10438a 100644
--- a/tensorflow/core/kernels/tensor_list_test.cc
+++ b/tensorflow/core/kernels/tensor_list_test.cc
@@ -221,6 +221,20 @@ TEST(TensorListTest, DecodeRejectsShapeRankMismatchWithMatchingCount) {
   EXPECT_FALSE(TensorList().Decode(data));
 }
 
+TEST(TensorListTest, DecodeRejectsEmptyTensorMismatchedShape) {
+  // A zero-element tensor still has a shape. element_shape [1] and element shape
+  // [0] share no common count and are incompatible, so an empty element must be
+  // rejected just like a non-empty one.
+  TensorShapeProto shape;
+  shape.add_dim()->set_size(1);
+  VariantTensorData data;
+  data.set_metadata(TensorListMetadata(
+      /*invalid_indices=*/{}, static_cast<uint64_t>(DT_FLOAT),
+      std::numeric_limits<uint64_t>::max(), shape.SerializeAsString()));
+  *data.add_tensors() = Tensor(DT_FLOAT, TensorShape({0}));
+  EXPECT_FALSE(TensorList().Decode(data));
+}
+
 TEST(TensorListTest, DecodeRejectsOutOfRangeMaxNumElements) {
   TensorShapeProto shape;
   VariantTensorData data;

From 14bf4430b610ea51d1bea0d7928d910cbff6b99c Mon Sep 17 00:00:00 2001
From: endorphin13 <elshad.toklayev.business@gmail.com>
Date: Thu, 1 Oct 2026 14:16:25 +0500
Subject: [PATCH 0112/1000] Keep an empty FFT axis empty in the complex
 FFTND/IFFTND kernels too

FFTNBase also backs the complex FFTND and IFFTND GPU kernels, which take an
explicit fft_length and axes. Their output was always sized from fft_length,
so an input with an empty FFT axis, for example shape [2, 0] with
fft_length [16], gave a non-empty output that the num_elements() == 0 early
return never wrote. That is the same uninitialized-output disclosure as the
real-valued transforms.

Compute the output dimension in one place for both branches: start from
fft_length, apply the real-transform adjustment, and set it to zero when the
input axis is empty. Output shapes for non-empty inputs are unchanged.
---
 tensorflow/core/kernels/fft_ops.cc | 28 ++++++++++++++--------------
 1 file changed, 14 insertions(+), 14 deletions(-)

diff --git a/tensorflow/core/kernels/fft_ops.cc b/tensorflow/core/kernels/fft_ops.cc
index 6e07bc9796c..0e03c90e812 100644
--- a/tensorflow/core/kernels/fft_ops.cc
+++ b/tensorflow/core/kernels/fft_ops.cc
@@ -320,11 +320,12 @@ class FFTNBase : public OpKernel {
                       "fft_length[", i,
                       "] must >= 0, but got: ", fft_length_as_vec(i))));
       fft_shape[i] = fft_length_as_vec(i);
+      auto input_index = input_rank - fft_rank + i;
+      uint64_t dim = fft_shape[i];
       if (IsReal()) {
         bool inner_most = (i == fft_rank - 1);
         uint64_t min_input_dim_length =
             !IsForward() && inner_most ? fft_shape[i] / 2 + 1 : fft_shape[i];
-        auto input_index = input_rank - fft_rank + i;
         OP_REQUIRES(
             ctx,
             // We pass through empty tensors, so special case them here.
@@ -334,21 +335,20 @@ class FFTNBase : public OpKernel {
                 "Input dimension ", input_index,
                 " must have length of at least ", min_input_dim_length,
                 " but got: ", input_shape.dim_size(input_index))));
-        uint64_t dim = IsForward() && inner_most && fft_shape[i] != 0
-                           ? fft_shape[i] / 2 + 1
-                           : fft_shape[i];
-        // An empty FFT axis is passed through by the check above. Keep it
-        // empty in the output too. Sizing it from `fft_length` instead would
-        // give a non-empty output for an empty input, and the
-        // `num_elements() == 0` early return below would then hand back the
-        // output buffer without ever writing it.
-        if (input_shape.dim_size(input_index) == 0) {
-          dim = 0;
+        if (IsForward() && inner_most && fft_shape[i] != 0) {
+          dim = fft_shape[i] / 2 + 1;
         }
-        output_shape.set_dim(output_shape.dims() - fft_rank + i, dim);
-      } else {
-        output_shape.set_dim(output_shape.dims() - fft_rank + i, fft_shape[i]);
       }
+      // An empty FFT axis is passed through, so keep it empty in the output
+      // too. Sizing it from `fft_length` instead would give a non-empty output
+      // for an empty input, and the `num_elements() == 0` early return below
+      // would then hand back the output buffer without ever writing it. This
+      // applies to the complex transforms as well, which also take an explicit
+      // `fft_length`.
+      if (input_shape.dim_size(input_index) == 0) {
+        dim = 0;
+      }
+      output_shape.set_dim(output_shape.dims() - fft_rank + i, dim);
     }
 
     OP_REQUIRES_OK(ctx, ctx->allocate_output(0, output_shape, &out));

From c8b67202711cf92f389d36c7d89ae60cd409c317 Mon Sep 17 00:00:00 2001
From: endorphin13 <elshad.toklayev.business@gmail.com>
Date: Thu, 1 Oct 2026 14:18:20 +0500
Subject: [PATCH 0113/1000] Assert exact output shapes for empty FFT axes and
 cover mixed and ND cases

Check the whole output shape in test_empty_with_explicit_fft_length instead of
only its element count, which would also pass if a batch dimension were
altered.

Add a test for an empty FFT axis next to a non-empty one, in both orders and
both directions, so only the empty axis collapses and the other is still sized
from fft_length.

Add GPU-only tests for the ND transforms, which have no CPU kernel:
FFTND/IFFTND, covering the complex case fixed in the previous commit, and
RFFTND/IRFFTND.
---
 .../kernel_tests/signal/fft_ops_test.py       | 74 +++++++++++++++++--
 1 file changed, 69 insertions(+), 5 deletions(-)

diff --git a/tensorflow/python/kernel_tests/signal/fft_ops_test.py b/tensorflow/python/kernel_tests/signal/fft_ops_test.py
index c98a8380611..b93a45dd22a 100644
--- a/tensorflow/python/kernel_tests/signal/fft_ops_test.py
+++ b/tensorflow/python/kernel_tests/signal/fft_ops_test.py
@@ -368,6 +368,22 @@ class FFTOpsTest(BaseFFTOpsTest, parameterized.TestCase):
         rtol=tol,
     )
 
+  @parameterized.parameters((np.complex64,), (np.complex128,))
+  @test_util.run_gpu_only
+  def testEmptyAxis_fftn(self, np_type):
+    # Unlike the 1/2/3D kernels, FFTND and IFFTND take an explicit `fft_length`
+    # for complex input too. An empty FFT axis must stay empty in the output
+    # instead of being sized from it: the kernel returns early without writing
+    # the output for an empty input, so a non-empty output would be returned
+    # uninitialized.
+    x = np.zeros((2, 0)).astype(np_type)
+    self.assertEqual((2, 0), self._tf_fftn(x, (16,), (-1,)).shape)
+    self.assertEqual((2, 0), self._tf_ifftn(x, (16,), (-1,)).shape)
+
+    x = np.zeros((2, 0, 0)).astype(np_type)
+    self.assertEqual((2, 0, 0), self._tf_fftn(x, (16, 16), (-2, -1)).shape)
+    self.assertEqual((2, 0, 0), self._tf_ifftn(x, (16, 16), (-2, -1)).shape)
+
   @parameterized.parameters(itertools.product(
       (1,), range(3), (np.complex64, np.complex128)))
   def test_large_batch(self, rank, extra_dims, np_type):
@@ -624,12 +640,40 @@ class RFFTOpsTest(BaseFFTOpsTest, parameterized.TestCase):
 
     # Batch dimensions must stay non-empty: if every dimension is zero the
     # output is empty regardless of how the FFT axes are sized, and the test
-    # would pass against the unfixed kernel too.
-    x = np.zeros((2,) * extra_dims + (0,) * rank).astype(np_rtype)
-    self.assertEqual(0, self._tf_fft(x, rank, fft_length).size)
+    # would pass against the unfixed kernel too. Assert the exact shape rather
+    # than only the element count, so that a wrongly sized axis or an altered
+    # batch dimension is caught as well.
+    expected_shape = (2,) * extra_dims + (0,) * rank
 
-    x = np.zeros((2,) * extra_dims + (0,) * rank).astype(np_ctype)
-    self.assertEqual(0, self._tf_ifft(x, rank, fft_length).size)
+    x = np.zeros(expected_shape).astype(np_rtype)
+    out_fwd = self._tf_fft(x, rank, fft_length)
+    self.assertEqual(expected_shape, out_fwd.shape)
+    self.assertEqual(0, out_fwd.size)
+
+    x = np.zeros(expected_shape).astype(np_ctype)
+    out_bwd = self._tf_ifft(x, rank, fft_length)
+    self.assertEqual(expected_shape, out_bwd.shape)
+    self.assertEqual(0, out_bwd.size)
+
+  @parameterized.parameters((np.float32,), (np.float64,))
+  def test_empty_axis_next_to_non_empty_axis(self, np_rtype):
+    # Only an empty FFT axis stays empty. The other FFT axis is still sized
+    # from `fft_length`, and the batch dimension is untouched.
+    np_ctype = np.complex64 if np_rtype == np.float32 else np.complex128
+    fft_length = (16, 16)
+
+    # Forward: the inner-most output axis is fft_length / 2 + 1.
+    x = np.zeros((2, 0, 16)).astype(np_rtype)
+    self.assertEqual((2, 0, 9), self._tf_fft(x, 2, fft_length).shape)
+    x = np.zeros((2, 16, 0)).astype(np_rtype)
+    self.assertEqual((2, 16, 0), self._tf_fft(x, 2, fft_length).shape)
+
+    # Inverse: the inner-most input axis must hold fft_length / 2 + 1 values,
+    # and the inner-most output axis is fft_length.
+    x = np.zeros((2, 0, 9)).astype(np_ctype)
+    self.assertEqual((2, 0, 16), self._tf_ifft(x, 2, fft_length).shape)
+    x = np.zeros((2, 16, 0)).astype(np_ctype)
+    self.assertEqual((2, 16, 0), self._tf_ifft(x, 2, fft_length).shape)
 
   @parameterized.parameters(itertools.product(
       VALID_FFT_RANKS, range(3), (5, 6), (np.float32, np.float64)))
@@ -767,6 +811,26 @@ class RFFTOpsTest(BaseFFTOpsTest, parameterized.TestCase):
     )
     self._CompareBackward_fftn(c2r, fft_length, axes, norm=norm, rtol=tol)
 
+  @parameterized.parameters((np.float32,), (np.float64,))
+  @test_util.run_gpu_only
+  def testEmptyAxis_rfftn(self, np_rtype):
+    # The ND counterpart of test_empty_with_explicit_fft_length: an empty FFT
+    # axis must stay empty, while a non-empty FFT axis is still sized from
+    # `fft_length`.
+    np_ctype = np.complex64 if np_rtype == np.float32 else np.complex128
+    fft_length = (16, 16)
+    axes = (-2, -1)
+
+    x = np.zeros((2, 0, 0)).astype(np_rtype)
+    self.assertEqual((2, 0, 0), self._tf_fftn(x, fft_length, axes).shape)
+    x = np.zeros((2, 0, 16)).astype(np_rtype)
+    self.assertEqual((2, 0, 9), self._tf_fftn(x, fft_length, axes).shape)
+
+    x = np.zeros((2, 0, 0)).astype(np_ctype)
+    self.assertEqual((2, 0, 0), self._tf_ifftn(x, fft_length, axes).shape)
+    x = np.zeros((2, 0, 9)).astype(np_ctype)
+    self.assertEqual((2, 0, 16), self._tf_ifftn(x, fft_length, axes).shape)
+
   @parameterized.parameters(itertools.product(
       (1,), range(3), (64, 128), (np.float32, np.float64)))
   def test_large_batch(self, rank, extra_dims, size, np_rtype):

From c7772ee347488dd79bd477cf04b384127ee1c709 Mon Sep 17 00:00:00 2001
From: Mina Ibrahim <136262314+MinaIbrahim10@users.noreply.github.com>
Date: Thu, 1 Oct 2026 08:58:30 -0400
Subject: [PATCH 0114/1000] Synchronize data service and Grappler bindings for
 free-threaded Python

---
 tensorflow/core/data/service/server_lib.cc    | 12 ++-
 tensorflow/core/data/service/server_lib.h     |  9 ++
 tensorflow/core/grappler/clusters/cluster.h   | 11 +++
 .../service/server_lib_wrapper.cc             | 91 ++++++++++++-----
 .../experimental/service/utils_wrapper.cc     |  3 +-
 tensorflow/python/grappler/BUILD              |  1 +
 tensorflow/python/grappler/cluster_test.py    | 24 +++++
 tensorflow/python/grappler/cluster_wrapper.cc | 99 ++++++++++++++-----
 .../python/grappler/tf_optimizer_wrapper.cc   | 16 ++-
 9 files changed, 207 insertions(+), 59 deletions(-)

diff --git a/tensorflow/core/data/service/server_lib.cc b/tensorflow/core/data/service/server_lib.cc
index 1555fe0c9ac..ca546cc4d47 100644
--- a/tensorflow/core/data/service/server_lib.cc
+++ b/tensorflow/core/data/service/server_lib.cc
@@ -104,10 +104,16 @@ void GrpcDataServerBase::Stop() {
 }
 
 void GrpcDataServerBase::Join() {
-  if (!server_) {
-    return;
+  ::grpc::Server* server = nullptr;
+  {
+    std::lock_guard<std::mutex> lock(ExternalMutex());
+    if (!server_) {
+      return;
+    }
+    server = server_.get();
   }
-  server_->Wait();
+
+  server->Wait();
 }
 
 int GrpcDataServerBase::BoundPort() { return bound_port(); }
diff --git a/tensorflow/core/data/service/server_lib.h b/tensorflow/core/data/service/server_lib.h
index 8647da84859..d0bbaf7a86b 100644
--- a/tensorflow/core/data/service/server_lib.h
+++ b/tensorflow/core/data/service/server_lib.h
@@ -17,6 +17,7 @@ limitations under the License.
 #define TENSORFLOW_CORE_DATA_SERVICE_SERVER_LIB_H_
 
 #include <memory>
+#include <mutex>
 #include <string>
 #include <vector>
 
@@ -61,6 +62,12 @@ class GrpcDataServerBase {
   // Returns the port bound by the server. Only valid after calling Start().
   int BoundPort();
 
+  // Serializes compound external operations on this server instance.
+  //
+  // Join() deliberately only holds this mutex while snapshotting server_, since
+  // holding it while waiting would prevent Stop() from shutting the server down.
+  std::mutex& ExternalMutex() const { return external_mu_; }
+
   // Exports the server state to improve debuggability.
   virtual ServerStateExport ExportState() const = 0;
 
@@ -79,6 +86,8 @@ class GrpcDataServerBase {
   const std::string server_type_;
 
  private:
+  mutable std::mutex external_mu_;
+
   int bound_port_;
   bool started_ = false;
   bool stopped_ = false;
diff --git a/tensorflow/core/grappler/clusters/cluster.h b/tensorflow/core/grappler/clusters/cluster.h
index 2dcc3596075..21c43985dd4 100644
--- a/tensorflow/core/grappler/clusters/cluster.h
+++ b/tensorflow/core/grappler/clusters/cluster.h
@@ -16,6 +16,7 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_GRAPPLER_CLUSTERS_CLUSTER_H_
 #define TENSORFLOW_CORE_GRAPPLER_CLUSTERS_CLUSTER_H_
 
+#include <mutex>
 #include <string>
 #include <unordered_map>
 #include <utility>
@@ -44,6 +45,13 @@ class Cluster {
   explicit Cluster(int timeout_s);
   virtual ~Cluster();
 
+  // Serializes compound external operations on this Cluster instance.
+  //
+  // This mutex intentionally lives in Cluster rather than in an individual
+  // Python extension so separate extension modules can coordinate access to
+  // the same native Cluster object.
+  std::mutex& ExternalMutex() const { return external_mu_; }
+
   // Returns a string that represent the type of cluster that was instantiated.
   virtual std::string type() const = 0;
 
@@ -142,6 +150,9 @@ class Cluster {
   const int timeout_s_;
   SessionOptions options_;
   RunOptions run_options_;
+
+ private:
+  mutable std::mutex external_mu_;
 };
 
 }  // end namespace grappler
diff --git a/tensorflow/python/data/experimental/service/server_lib_wrapper.cc b/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
index 81b7ded13f9..1f2a7f8af0a 100644
--- a/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
+++ b/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
@@ -15,6 +15,7 @@ limitations under the License.
 #include <cstdint>
 #include <limits>
 #include <memory>
+#include <mutex>
 #include <string>
 #include <vector>
 
@@ -42,21 +43,36 @@ limitations under the License.
 
 namespace py = pybind11;
 
-PYBIND11_MODULE(_pywrap_server_lib, m) {
+PYBIND11_MODULE(
+    _pywrap_server_lib, m, pybind11::mod_gil_not_used()) {
   pybind11_protobuf::ImportNativeProtoCasters();
 
-  py::class_<tensorflow::data::DispatchGrpcDataServer>(m,
-                                                       "DispatchGrpcDataServer")
-      .def("start", &tensorflow::data::DispatchGrpcDataServer::Start)
-      .def("stop", &tensorflow::data::DispatchGrpcDataServer::Stop)
+  py::class_<tensorflow::data::DispatchGrpcDataServer>(
+      m, "DispatchGrpcDataServer")
+      .def("start", [](tensorflow::data::DispatchGrpcDataServer* server) {
+        std::lock_guard<std::mutex> lock(server->ExternalMutex());
+        return server->Start();
+      })
+      .def("stop", [](tensorflow::data::DispatchGrpcDataServer* server) {
+        std::lock_guard<std::mutex> lock(server->ExternalMutex());
+        server->Stop();
+      })
       .def("join", &tensorflow::data::DispatchGrpcDataServer::Join,
            py::call_guard<py::gil_scoped_release>())
-      .def("bound_port", &tensorflow::data::DispatchGrpcDataServer::BoundPort)
+      .def("bound_port",
+           [](tensorflow::data::DispatchGrpcDataServer* server) {
+             std::lock_guard<std::mutex> lock(server->ExternalMutex());
+             return server->BoundPort();
+           })
       .def("num_workers",
            [](tensorflow::data::DispatchGrpcDataServer* server) -> int {
              int num_workers;
-             absl::Status status = server->NumWorkers(&num_workers);
-             tensorflow::MaybeRaiseFromStatus(status);
+             absl::Status status;
+             {
+               std::lock_guard<std::mutex> lock(server->ExternalMutex());
+               status = server->NumWorkers(&num_workers);
+             }
+             tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
              return num_workers;
            })
       .def("snapshot_streams",
@@ -64,22 +80,41 @@ PYBIND11_MODULE(_pywrap_server_lib, m) {
               const std::string& path)
                -> std::vector<tensorflow::data::SnapshotStreamInfoWrapper> {
              std::vector<tensorflow::data::SnapshotStreamInfoWrapper> streams;
-             absl::Status status = server->SnapshotStreams(path, &streams);
-             tensorflow::MaybeRaiseFromStatus(status);
+             absl::Status status;
+             {
+               std::lock_guard<std::mutex> lock(server->ExternalMutex());
+               status = server->SnapshotStreams(path, &streams);
+             }
+             tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
              return streams;
            });
 
-  py::class_<tensorflow::data::WorkerGrpcDataServer>(m, "WorkerGrpcDataServer")
-      .def("start", &tensorflow::data::WorkerGrpcDataServer::Start)
-      .def("stop", &tensorflow::data::WorkerGrpcDataServer::Stop)
+  py::class_<tensorflow::data::WorkerGrpcDataServer>(
+      m, "WorkerGrpcDataServer")
+      .def("start", [](tensorflow::data::WorkerGrpcDataServer* server) {
+        std::lock_guard<std::mutex> lock(server->ExternalMutex());
+        return server->Start();
+      })
+      .def("stop", [](tensorflow::data::WorkerGrpcDataServer* server) {
+        std::lock_guard<std::mutex> lock(server->ExternalMutex());
+        server->Stop();
+      })
       .def("join", &tensorflow::data::WorkerGrpcDataServer::Join,
            py::call_guard<py::gil_scoped_release>())
-      .def("bound_port", &tensorflow::data::WorkerGrpcDataServer::BoundPort)
+      .def("bound_port",
+           [](tensorflow::data::WorkerGrpcDataServer* server) {
+             std::lock_guard<std::mutex> lock(server->ExternalMutex());
+             return server->BoundPort();
+           })
       .def("num_tasks",
            [](tensorflow::data::WorkerGrpcDataServer* server) -> int {
              int num_tasks;
-             absl::Status status = server->NumTasks(&num_tasks);
-             tensorflow::MaybeRaiseFromStatus(status);
+             absl::Status status;
+             {
+               std::lock_guard<std::mutex> lock(server->ExternalMutex());
+               status = server->NumTasks(&num_tasks);
+             }
+             tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
              return num_tasks;
            })
       .def("snapshot_task_progresses",
@@ -87,9 +122,13 @@ PYBIND11_MODULE(_pywrap_server_lib, m) {
                -> std::vector<tensorflow::data::SnapshotTaskProgressWrapper> {
              std::vector<tensorflow::data::SnapshotTaskProgressWrapper>
                  snapshot_task_progresses;
-             absl::Status status =
-                 server->SnapshotTaskProgresses(&snapshot_task_progresses);
-             tensorflow::MaybeRaiseFromStatus(status);
+             absl::Status status;
+             {
+               std::lock_guard<std::mutex> lock(server->ExternalMutex());
+               status =
+                   server->SnapshotTaskProgresses(&snapshot_task_progresses);
+             }
+             tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
              return snapshot_task_progresses;
            });
 
@@ -99,13 +138,13 @@ PYBIND11_MODULE(_pywrap_server_lib, m) {
           -> std::unique_ptr<tensorflow::data::DispatchGrpcDataServer> {
         tensorflow::data::experimental::DispatcherConfig config;
         if (!config.ParseFromString(serialized_dispatcher_config)) {
-          tensorflow::MaybeRaiseFromStatus(absl::InvalidArgumentError(
+          tsl::MaybeRaiseRegisteredFromStatusWithGIL(absl::InvalidArgumentError(
               "Failed to deserialize dispatcher config."));
         }
         std::unique_ptr<tensorflow::data::DispatchGrpcDataServer> server;
         absl::Status status =
             tensorflow::data::NewDispatchServer(config, server);
-        tensorflow::MaybeRaiseFromStatus(status);
+        tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
         return server;
       },
       py::return_value_policy::reference);
@@ -116,12 +155,13 @@ PYBIND11_MODULE(_pywrap_server_lib, m) {
           -> std::unique_ptr<tensorflow::data::WorkerGrpcDataServer> {
         tensorflow::data::experimental::WorkerConfig config;
         if (!config.ParseFromString(serialized_worker_config)) {
-          tensorflow::MaybeRaiseFromStatus(absl::InvalidArgumentError(
+          tsl::MaybeRaiseRegisteredFromStatusWithGIL(absl::InvalidArgumentError(
               "Failed to deserialize worker config."));
         }
         std::unique_ptr<tensorflow::data::WorkerGrpcDataServer> server;
-        absl::Status status = tensorflow::data::NewWorkerServer(config, server);
-        tensorflow::MaybeRaiseFromStatus(status);
+        absl::Status status =
+            tensorflow::data::NewWorkerServer(config, server);
+        tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
         return server;
       },
       py::return_value_policy::reference);
@@ -144,7 +184,7 @@ PYBIND11_MODULE(_pywrap_server_lib, m) {
                          " from dispatcher at ", address),
             deadline_micros);
         Py_END_ALLOW_THREADS;
-        tensorflow::MaybeRaiseFromStatus(status);
+        tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
         return metadata;
       },
       py::return_value_policy::reference);
@@ -170,6 +210,7 @@ PYBIND11_MODULE(_pywrap_server_lib, m) {
                  snapshot_task_progress_wrapper) -> bool {
             return snapshot_task_progress_wrapper.completed;
           });
+
   py::class_<tensorflow::data::SnapshotStreamInfoWrapper>
       snapshot_stream_info_wrapper(m, "SnapshotStreamInfoWrapper");
   snapshot_stream_info_wrapper.def(py::init<>())
diff --git a/tensorflow/python/data/experimental/service/utils_wrapper.cc b/tensorflow/python/data/experimental/service/utils_wrapper.cc
index c725ff3f58e..ede0428aa80 100644
--- a/tensorflow/python/data/experimental/service/utils_wrapper.cc
+++ b/tensorflow/python/data/experimental/service/utils_wrapper.cc
@@ -20,7 +20,8 @@ limitations under the License.
 #include "tensorflow/core/data/service/py_utils.h"
 #include "tensorflow/python/lib/core/pybind11_lib.h"
 
-PYBIND11_MODULE(_pywrap_utils_exp, m) {
+PYBIND11_MODULE(
+    _pywrap_utils_exp, m, pybind11::mod_gil_not_used()) {
   m.def("TF_DATA_DefaultProtocol",
         []() -> std::string { return tensorflow::data::DefaultProtocol(); });
 };
diff --git a/tensorflow/python/grappler/BUILD b/tensorflow/python/grappler/BUILD
index 5695e431257..3f05b3789ff 100644
--- a/tensorflow/python/grappler/BUILD
+++ b/tensorflow/python/grappler/BUILD
@@ -295,6 +295,7 @@ tf_python_pybind_extension(
         "//tensorflow/core/grappler:utils",
         "//tensorflow/core/platform:status",
         "//tensorflow/python/lib/core:pybind11_status",
+        "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log:check",
diff --git a/tensorflow/python/grappler/cluster_test.py b/tensorflow/python/grappler/cluster_test.py
index 245dd711a78..342a46204c8 100644
--- a/tensorflow/python/grappler/cluster_test.py
+++ b/tensorflow/python/grappler/cluster_test.py
@@ -14,6 +14,8 @@
 # ==============================================================================
 """Tests for the swig wrapper of clusters."""
 
+import concurrent.futures
+
 from tensorflow.core.protobuf import device_properties_pb2
 from tensorflow.python.framework import meta_graph
 from tensorflow.python.framework import ops
@@ -133,6 +135,28 @@ class ClusterTest(test.TestCase):
       self.assertTrue('MatMul' in op_names)
       self.assertEqual(op_names, sorted(op_names))
 
+  def testConcurrentListDevices(self):
+    cpu_properties = device_properties_pb2.DeviceProperties(
+        type='CPU', frequency=3000, num_cores=8)
+    gpu_properties = device_properties_pb2.DeviceProperties(
+        type='GPU', frequency=1000, num_cores=60)
+    named_cpu = device_properties_pb2.NamedDevice(
+        properties=cpu_properties, name='/CPU:0')
+    named_gpu = device_properties_pb2.NamedDevice(
+        properties=gpu_properties, name='/GPU:0')
+    grappler_cluster = cluster.Cluster(devices=[named_cpu, named_gpu])
+    self.addCleanup(grappler_cluster.Shutdown)
+
+    def list_device_names(_):
+      return sorted(device.name for device in grappler_cluster.ListDevices())
+
+    expected = ['/CPU:0', '/GPU:0']
+    with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
+      results = list(executor.map(list_device_names, range(64)))
+
+    for names in results:
+      self.assertEqual(names, expected)
+
   def testSupportDevices(self):
     with ops.Graph().as_default() as g:
       a = random_ops.random_uniform(shape=(2, 3))
diff --git a/tensorflow/python/grappler/cluster_wrapper.cc b/tensorflow/python/grappler/cluster_wrapper.cc
index 7d835d5614a..1f011e33ae2 100644
--- a/tensorflow/python/grappler/cluster_wrapper.cc
+++ b/tensorflow/python/grappler/cluster_wrapper.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include <cstdint>
 #include <iterator>
 #include <memory>
+#include <mutex>
 #include <set>
 #include <stdexcept>
 #include <string>
@@ -26,6 +27,7 @@ limitations under the License.
 #include <unordered_map>
 #include <vector>
 
+#include "absl/base/no_destructor.h"
 #include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "pybind11/pybind11.h"  // from @pybind11
@@ -53,6 +55,19 @@ limitations under the License.
 
 namespace py = pybind11;
 
+
+namespace {
+
+// SingleMachine uses process-global provisioning state. Serialize Python
+// creation/shutdown paths that update that state without relying on static
+// destruction order during interpreter shutdown.
+std::mutex& GetClusterLifecycleMutex() {
+  static absl::NoDestructor<std::mutex> mu;
+  return *mu;
+}
+
+}  // namespace
+
 absl::Status _GetOpPerformanceDataAndRunTime(
     const tensorflow::grappler::GrapplerItem& item,
     tensorflow::grappler::CostEstimator* cost_measure,
@@ -62,7 +77,7 @@ absl::Status _GetOpPerformanceDataAndRunTime(
   if (!status.ok()) return status;
 
   tensorflow::RunMetadata run_metadata;
-  tsl::MaybeRaiseRegisteredFromStatus(
+  tsl::MaybeRaiseRegisteredFromStatusWithGIL(
       cost_measure->PredictCosts(item.graph, &run_metadata, costs));
 
   if (op_performance_data) {
@@ -74,7 +89,8 @@ absl::Status _GetOpPerformanceDataAndRunTime(
 
 PYBIND11_MAKE_OPAQUE(tensorflow::grappler::Cluster);
 
-PYBIND11_MODULE(_pywrap_tf_cluster, m) {
+PYBIND11_MODULE(
+    _pywrap_tf_cluster, m, pybind11::mod_gil_not_used()) {
   py::class_<tensorflow::grappler::Cluster> grappler_cluster(m, "Cluster");
 
   m.def("TF_NewCluster",
@@ -82,6 +98,9 @@ PYBIND11_MODULE(_pywrap_tf_cluster, m) {
            bool disable_detailed_stats) -> tensorflow::grappler::Cluster* {
           // TODO(petebu): Make these named arguments with default values
           // instead.
+          std::unique_lock<std::mutex> lifecycle_lock(
+              GetClusterLifecycleMutex());
+
           int num_cpu_cores =
               tensorflow::grappler::GetNumAvailableLogicalCPUCores();
           int num_gpus = tensorflow::grappler::GetNumAvailableGPUs();
@@ -92,7 +111,9 @@ PYBIND11_MODULE(_pywrap_tf_cluster, m) {
           cluster->DisableDetailedStats(disable_detailed_stats);
           cluster->AllowSoftPlacement(allow_soft_placement);
           cluster->SetNumWarmupSteps(10);
-          tsl::MaybeRaiseRegisteredFromStatus(cluster->Provision());
+          absl::Status provision_status = cluster->Provision();
+          lifecycle_lock.unlock();
+          tsl::MaybeRaiseRegisteredFromStatusWithGIL(provision_status);
           return cluster.release();
         });
 
@@ -116,30 +137,38 @@ PYBIND11_MODULE(_pywrap_tf_cluster, m) {
           }
           std::unique_ptr<tensorflow::grappler::Cluster> cluster =
               std::make_unique<tensorflow::grappler::VirtualCluster>(devices);
-          {
-            // TODO(petebu): Do we need to hold the GIL here?
-            py::gil_scoped_acquire acquire;
-            tsl::MaybeRaiseRegisteredFromStatus(cluster->Provision());
-          }
+          tsl::MaybeRaiseRegisteredFromStatusWithGIL(cluster->Provision());
           return cluster.release();
         });
 
   m.def("TF_ShutdownCluster", [](tensorflow::grappler::Cluster* cluster) {
-    // TODO(petebu): Do we need to hold the GIL here?
-    py::gil_scoped_acquire acquire;
+    if (cluster == nullptr) {
+      return;
+    }
+    std::lock_guard<std::mutex> cluster_lock(cluster->ExternalMutex());
+    std::lock_guard<std::mutex> lifecycle_lock(GetClusterLifecycleMutex());
     (void)cluster->Shutdown();
   });
 
   m.def("TF_ListDevices",
         [](tensorflow::grappler::Cluster* cluster) -> std::vector<py::bytes> {
-          const std::unordered_map<std::string, tensorflow::DeviceProperties>&
-              devices = cluster->GetDevices();
+          if (cluster == nullptr) {
+            tsl::MaybeRaiseRegisteredFromStatusWithGIL(
+                absl::InvalidArgumentError("Cluster cannot be None."));
+            return {};
+          }
+
           std::vector<py::bytes> named_devices;
-          for (auto& dev : devices) {
-            tensorflow::NamedDevice d;
-            d.set_name(dev.first);
-            *d.mutable_properties() = dev.second;
-            named_devices.push_back(d.SerializeAsString());
+          {
+            std::lock_guard<std::mutex> lock(cluster->ExternalMutex());
+            const auto& devices = cluster->GetDevices();
+
+            for (const auto& dev : devices) {
+              tensorflow::NamedDevice d;
+              d.set_name(dev.first);
+              *d.mutable_properties() = dev.second;
+              named_devices.push_back(d.SerializeAsString());
+            }
           }
           return named_devices;
         });
@@ -163,15 +192,21 @@ PYBIND11_MODULE(_pywrap_tf_cluster, m) {
          tensorflow::grappler::GrapplerItem* item)
           -> std::unordered_map<std::string, std::vector<std::string>> {
         if (cluster == nullptr || item == nullptr) {
-          tsl::MaybeRaiseRegisteredFromStatus(absl::Status(
+          tsl::MaybeRaiseRegisteredFromStatusWithGIL(absl::Status(
               absl::InternalError("You need both a cluster and an "
                                   "item to get supported devices.")));
         }
-        const std::unordered_map<std::string, tensorflow::DeviceProperties>&
-            devices = cluster->GetDevices();
         std::unordered_map<std::string, std::vector<std::string>> device_types;
-        for (const auto& dev : devices) {
-          device_types[dev.second.type()].push_back(dev.first);
+        std::string cluster_type;
+
+        {
+          std::lock_guard<std::mutex> lock(cluster->ExternalMutex());
+          const auto& devices = cluster->GetDevices();
+          cluster_type = cluster->type();
+
+          for (const auto& dev : devices) {
+            device_types[dev.second.type()].push_back(dev.first);
+          }
         }
 
         std::unordered_map<std::string, std::set<std::string>>
@@ -182,7 +217,7 @@ PYBIND11_MODULE(_pywrap_tf_cluster, m) {
         for (const auto& node : item->graph.node()) {
           for (const auto& dev : device_types) {
             const std::string& type = dev.first;
-            if (cluster->type() != "single_machine") {
+            if (cluster_type != "single_machine") {
               // The actual kernel may not be linked in this binary.
               supported_device_types[node.name()].insert(type);
             } else {
@@ -259,6 +294,14 @@ PYBIND11_MODULE(_pywrap_tf_cluster, m) {
         [](tensorflow::grappler::GrapplerItem* item,
            tensorflow::grappler::Cluster* cluster, bool generate_timeline)
             -> std::tuple<std::vector<py::bytes>, double, py::bytes> {
+          if (cluster == nullptr || item == nullptr) {
+            tsl::MaybeRaiseRegisteredFromStatusWithGIL(absl::InvalidArgumentError(
+                "You need both a cluster and an item to measure costs."));
+            return {};
+          }
+
+          std::lock_guard<std::mutex> lock(cluster->ExternalMutex());
+
           const int num_measurements = cluster->type() == "virtual" ? 1 : 10;
           tensorflow::grappler::MeasuringCostEstimator cost_measure(
               cluster, num_measurements, 0);
@@ -274,7 +317,7 @@ PYBIND11_MODULE(_pywrap_tf_cluster, m) {
           tensorflow::StepStats step_stats;
           if (generate_timeline) {
             tensorflow::RunMetadata metadata;
-            tsl::MaybeRaiseRegisteredFromStatus(
+            tsl::MaybeRaiseRegisteredFromStatusWithGIL(
                 cluster->Run(item->graph, item->feed, item->fetch, &metadata));
             step_stats = metadata.step_stats();
           }
@@ -301,16 +344,18 @@ PYBIND11_MODULE(_pywrap_tf_cluster, m) {
           -> std::unordered_map<std::string,
                                 std::tuple<int64_t, std::vector<MemoryUsage>>> {
         if (item == nullptr || cluster == nullptr) {
-          tsl::MaybeRaiseRegisteredFromStatus(absl::Status(absl::InternalError(
+          tsl::MaybeRaiseRegisteredFromStatusWithGIL(absl::Status(absl::InternalError(
               "You need both a cluster and an item to determine peak "
               "memory usage.")));
         }
+        std::lock_guard<std::mutex> lock(cluster->ExternalMutex());
+
         tensorflow::grappler::GraphMemory memory(*item);
 
         if (cluster->DetailedStatsEnabled()) {
-          tsl::MaybeRaiseRegisteredFromStatus(memory.InferDynamically(cluster));
+          tsl::MaybeRaiseRegisteredFromStatusWithGIL(memory.InferDynamically(cluster));
         } else {
-          tsl::MaybeRaiseRegisteredFromStatus(
+          tsl::MaybeRaiseRegisteredFromStatusWithGIL(
               memory.InferStatically(cluster->GetDevices()));
         }
 
diff --git a/tensorflow/python/grappler/tf_optimizer_wrapper.cc b/tensorflow/python/grappler/tf_optimizer_wrapper.cc
index 4e88995858e..6fff2fa6665 100644
--- a/tensorflow/python/grappler/tf_optimizer_wrapper.cc
+++ b/tensorflow/python/grappler/tf_optimizer_wrapper.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 
 #include <memory>
+#include <mutex>
 #include <stdexcept>
 #include <string>
 #include <unordered_map>
@@ -92,8 +93,16 @@ tensorflow::GraphDef TF_OptimizeGraph(
     tensorflow::DeviceBase* cpu_device = nullptr;
     tensorflow::grappler::MetaOptimizer optimizer(cpu_device, config_proto);
 
-    tsl::MaybeRaiseRegisteredFromStatusWithGIL(
-        optimizer.Optimize(cluster, *grappler_item, &out_graph));
+    absl::Status optimize_status;
+    if (cluster != nullptr) {
+      std::lock_guard<std::mutex> lock(cluster->ExternalMutex());
+      optimize_status =
+          optimizer.Optimize(cluster, *grappler_item, &out_graph);
+    } else {
+      optimize_status =
+          optimizer.Optimize(nullptr, *grappler_item, &out_graph);
+    }
+    tsl::MaybeRaiseRegisteredFromStatusWithGIL(optimize_status);
     if (strip_default_attributes) {
       tensorflow::StripDefaultAttributes(*tensorflow::OpRegistry::Global(),
                                          out_graph.mutable_node());
@@ -114,7 +123,8 @@ tensorflow::GraphDef TF_OptimizeGraph(
 // larger than 2GiB.
 // At the moment, the open source python API defaults to the serialized
 // implementation.
-PYBIND11_MODULE(_pywrap_tf_optimizer, m) {
+PYBIND11_MODULE(
+    _pywrap_tf_optimizer, m, pybind11::mod_gil_not_used()) {
   pybind11_protobuf::ImportNativeProtoCasters();
   m.def("TF_OptimizeGraphSerialized",
         [](tensorflow::grappler::Cluster* cluster,

From feab00acc553896fd4b13f2619cb79aaa2ffef34 Mon Sep 17 00:00:00 2001
From: alwaysprince05 <prince.p24@medhaviskillsuniversity.edu.in>
Date: Thu, 1 Oct 2026 19:06:03 +0530
Subject: [PATCH 0115/1000] Address review: match partial/local device specs
 via IsSpecification

Replace the separate device-type and all-components exact-match logic with
a single DeviceNameUtils::IsSpecification pass, so valid partial
constraints like '/job:localhost' are accepted while nonexistent ids or
jobs (e.g. '/device:CPU:99') are rejected. Parse node device specs with
ParseFullOrLocalName, return an InvalidArgumentError for malformed
specifications, pre-parse the available devices once per call, and build
the error message only on the failure path. Use op->GetFunctionDef() so
function library overrides are handled, and add C++ unit tests covering
empty, partial, fully-specified, local-name, invalid, and malformed specs.
---
 tensorflow/core/common_runtime/eager/BUILD    |  14 ++
 .../core/common_runtime/eager/execute.cc      |   3 +-
 .../eager/validate_function_devices.cc        |  78 ++++-----
 .../eager/validate_function_devices_test.cc   | 161 ++++++++++++++++++
 .../polymorphic_function_xla_jit_test.py      |  82 ++++++++-
 5 files changed, 281 insertions(+), 57 deletions(-)
 create mode 100644 tensorflow/core/common_runtime/eager/validate_function_devices_test.cc

diff --git a/tensorflow/core/common_runtime/eager/BUILD b/tensorflow/core/common_runtime/eager/BUILD
index 2879718a356..80cd3b39461 100644
--- a/tensorflow/core/common_runtime/eager/BUILD
+++ b/tensorflow/core/common_runtime/eager/BUILD
@@ -719,6 +719,20 @@ cc_library(
     ],
 )
 
+tf_cc_test(
+    name = "validate_function_devices_test",
+    srcs = ["validate_function_devices_test.cc"],
+    deps = [
+        ":validate_function_devices",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core/platform:status_matchers",
+        "@com_google_absl//absl/status",
+    ],
+)
+
 tf_cc_test(
     name = "execute_test",
     srcs = ["execute_test.cc"],
diff --git a/tensorflow/core/common_runtime/eager/execute.cc b/tensorflow/core/common_runtime/eager/execute.cc
index f71217dc711..d976a59ffee 100644
--- a/tensorflow/core/common_runtime/eager/execute.cc
+++ b/tensorflow/core/common_runtime/eager/execute.cc
@@ -1419,8 +1419,7 @@ absl::Status GetOrCreateKernelAndDevice(
         GetFuncAttr(op, ctx, kXlaMustCompileAttr, &has_xla_must_compile)
             .IgnoreError();
         if (has_xla_must_compile) {
-          const FunctionLibraryDefinition* func_lib = ctx.FuncLibDef();
-          const FunctionDef* fdef = func_lib->Find(op->Name());
+          const FunctionDef* fdef = op->GetFunctionDef();
           if (fdef != nullptr) {
             std::vector<DeviceAttributes> device_attrs;
             ctx.ListDevices(&device_attrs);
diff --git a/tensorflow/core/common_runtime/eager/validate_function_devices.cc b/tensorflow/core/common_runtime/eager/validate_function_devices.cc
index 1fcf96d4b1d..4476c707360 100644
--- a/tensorflow/core/common_runtime/eager/validate_function_devices.cc
+++ b/tensorflow/core/common_runtime/eager/validate_function_devices.cc
@@ -19,14 +19,9 @@ limitations under the License.
 #include <vector>
 
 #include "absl/status/status.h"
-#include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "tensorflow/core/framework/function.h"
-#include "tensorflow/core/framework/node_def_util.h"
-#include "tensorflow/core/framework/types.h"
-#include "tensorflow/core/framework/types.pb.h"
-#include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/util/device_name_utils.h"
 
 namespace tensorflow {
@@ -34,71 +29,58 @@ namespace tensorflow {
 absl::Status ValidateFunctionDeviceConstraints(
     const FunctionDef& fdef,
     const std::vector<DeviceAttributes>& available_devices) {
-  // Build a set of available device types and full device names.
-  std::vector<DeviceType> supported_device_types;
-  std::vector<std::string> available_device_names;
+  // Pre-parse available devices once upfront.
+  std::vector<DeviceNameUtils::ParsedName> parsed_available_devices;
+  parsed_available_devices.reserve(available_devices.size());
   for (const auto& dev : available_devices) {
     DeviceNameUtils::ParsedName parsed;
     if (DeviceNameUtils::ParseFullName(dev.name(), &parsed)) {
-      supported_device_types.push_back(DeviceType(parsed.type));
+      parsed_available_devices.push_back(parsed);
     }
-    available_device_names.push_back(dev.name());
   }
 
   for (const NodeDef& node : fdef.node_def()) {
     const std::string& device = node.device();
     if (device.empty()) {
-      continue;  // No device constraint on this node.
-    }
-
-    DeviceNameUtils::ParsedName parsed_device;
-    if (!DeviceNameUtils::ParseFullName(device, &parsed_device)) {
-      // If the device name can't be parsed, skip it. Other validation passes
-      // will catch malformed device names.
       continue;
     }
 
-    // Check if the requested device type is available.
-    bool device_type_available = false;
-    for (const DeviceType& avail_type : supported_device_types) {
-      if (absl::EqualsIgnoreCase(parsed_device.type, avail_type.type_string())) {
-        device_type_available = true;
+    DeviceNameUtils::ParsedName parsed_device;
+    if (!DeviceNameUtils::ParseFullOrLocalName(device, &parsed_device)) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Malformed device specification '", device, "' for operation ",
+          node.name(), " (", node.op(), ")."));
+    }
+
+    if (!DeviceNameUtils::HasSomeDetails(parsed_device)) {
+      continue;
+    }
+
+    // The constraint is satisfied if at least one available device matches
+    // every component specified in the constraint. IsSpecification checks
+    // all specified attributes without requiring the constraint to be fully
+    // specified, so partial constraints (e.g. '/job:localhost') and local
+    // names (e.g. 'CPU:0') are handled uniformly.
+    bool satisfied = false;
+    for (const auto& avail_parsed : parsed_available_devices) {
+      if (DeviceNameUtils::IsSpecification(parsed_device, avail_parsed)) {
+        satisfied = true;
         break;
       }
     }
 
-    if (!device_type_available) {
+    if (!satisfied) {
+      std::vector<std::string> available_device_names;
+      available_device_names.reserve(available_devices.size());
+      for (const auto& dev : available_devices) {
+        available_device_names.push_back(dev.name());
+      }
       return absl::InvalidArgumentError(absl::StrCat(
           "Could not satisfy device specification '", device,
           "' for operation ", node.name(), " (", node.op(),
           "). Available devices [",
           absl::StrJoin(available_device_names, ", "), "]."));
     }
-
-    // For fully-specified device names, check if an exact match exists.
-    if (parsed_device.has_job && parsed_device.has_replica &&
-        parsed_device.has_task && parsed_device.has_type &&
-        parsed_device.has_id) {
-      bool exact_match = false;
-      for (const auto& dev : available_devices) {
-        DeviceNameUtils::ParsedName avail_parsed;
-        if (DeviceNameUtils::ParseFullName(dev.name(), &avail_parsed)) {
-          if (DeviceNameUtils::IsCompleteSpecification(parsed_device,
-                                                       avail_parsed)) {
-            exact_match = true;
-            break;
-          }
-        }
-      }
-
-      if (!exact_match) {
-        return absl::InvalidArgumentError(absl::StrCat(
-            "Could not satisfy device specification '", device,
-            "' for operation ", node.name(), " (", node.op(),
-            "). Available devices [",
-            absl::StrJoin(available_device_names, ", "), "]."));
-      }
-    }
   }
 
   return absl::OkStatus();
diff --git a/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc b/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc
new file mode 100644
index 00000000000..4bd82e0a14d
--- /dev/null
+++ b/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc
@@ -0,0 +1,161 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/core/common_runtime/eager/validate_function_devices.h"
+
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/core/framework/device_attributes.pb.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/platform/status_matchers.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+
+namespace {
+
+// A FunctionDef with a single node carrying the given device constraint.
+FunctionDef MakeFunctionWithDevice(const std::string& node_name,
+                                   const std::string& op,
+                                   const std::string& device) {
+  FunctionDef fdef;
+  NodeDef* node = fdef.add_node_def();
+  node->set_name(node_name);
+  node->set_op(op);
+  node->set_device(device);
+  return fdef;
+}
+
+std::vector<DeviceAttributes> LocalCpuDevices() {
+  DeviceAttributes attrs;
+  attrs.set_name("/job:localhost/replica:0/task:0/device:CPU:0");
+  attrs.set_device_type("CPU");
+  return {attrs};
+}
+
+bool IsInvalidArgument(const absl::Status& status) {
+  return status.code() == absl::StatusCode::kInvalidArgument;
+}
+
+bool Contains(const absl::Status& status, const std::string& substring) {
+  return std::string(status.message()).find(substring) != std::string::npos;
+}
+
+}  // namespace
+
+TEST(ValidateFunctionDevicesTest, EmptyFunctionIsOk) {
+  FunctionDef fdef;
+  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices()),
+              ::tensorflow::testing::IsOk());
+}
+
+TEST(ValidateFunctionDevicesTest, NoDeviceConstraintIsOk) {
+  FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "");
+  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices()),
+              ::tensorflow::testing::IsOk());
+}
+
+TEST(ValidateFunctionDevicesTest, FullySpecifiedDeviceIsOk) {
+  FunctionDef fdef = MakeFunctionWithDevice(
+      "add", "AddV2", "/job:localhost/replica:0/task:0/device:CPU:0");
+  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices()),
+              ::tensorflow::testing::IsOk());
+}
+
+TEST(ValidateFunctionDevicesTest, PartialJobConstraintIsOk) {
+  // '/job:localhost' is a partial constraint without a device type; it must
+  // be satisfied by the fully-specified available CPU device.
+  FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/job:localhost");
+  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices()),
+              ::tensorflow::testing::IsOk());
+}
+
+TEST(ValidateFunctionDevicesTest, PartialTaskConstraintIsOk) {
+  FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/task:0");
+  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices()),
+              ::tensorflow::testing::IsOk());
+}
+
+TEST(ValidateFunctionDevicesTest, LocalDeviceNameIsOk) {
+  FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "CPU:0");
+  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices()),
+              ::tensorflow::testing::IsOk());
+}
+
+TEST(ValidateFunctionDevicesTest, UnsatisfiedJobFails) {
+  FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/job:worker");
+  absl::Status status =
+      ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
+  EXPECT_TRUE(IsInvalidArgument(status));
+  EXPECT_TRUE(Contains(status, "Could not satisfy device specification"));
+  EXPECT_TRUE(Contains(status, "/job:worker"));
+  EXPECT_TRUE(Contains(status, "/job:localhost/replica:0/task:0/device:CPU:0"));
+}
+
+TEST(ValidateFunctionDevicesTest, InvalidDeviceIdFails) {
+  // The device type exists, but the id does not: the partial constraint must
+  // still be rejected.
+  FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/device:CPU:99");
+  absl::Status status =
+      ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
+  EXPECT_TRUE(IsInvalidArgument(status));
+  EXPECT_TRUE(Contains(status, "Could not satisfy device specification"));
+}
+
+TEST(ValidateFunctionDevicesTest, InvalidDeviceTypeFails) {
+  FunctionDef fdef =
+      MakeFunctionWithDevice("add", "AddV2", "/device:NONEXISTENT:0");
+  absl::Status status =
+      ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
+  EXPECT_TRUE(IsInvalidArgument(status));
+  EXPECT_TRUE(Contains(status, "Could not satisfy device specification"));
+}
+
+TEST(ValidateFunctionDevicesTest, LocalDeviceNameInvalidIdFails) {
+  FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "CPU:99");
+  absl::Status status =
+      ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
+  EXPECT_TRUE(IsInvalidArgument(status));
+  EXPECT_TRUE(Contains(status, "Could not satisfy device specification"));
+}
+
+TEST(ValidateFunctionDevicesTest, MalformedDeviceSpecFails) {
+  FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "not_a_device");
+  absl::Status status =
+      ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
+  EXPECT_TRUE(IsInvalidArgument(status));
+  EXPECT_TRUE(Contains(status, "Malformed device specification"));
+  EXPECT_TRUE(Contains(status, "not_a_device"));
+}
+
+TEST(ValidateFunctionDevicesTest, ErrorMentionsOperation) {
+  FunctionDef fdef = MakeFunctionWithDevice("my_add", "AddV2", "/job:worker");
+  absl::Status status =
+      ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
+  EXPECT_TRUE(Contains(status, "my_add"));
+  EXPECT_TRUE(Contains(status, "AddV2"));
+}
+
+TEST(ValidateFunctionDevicesTest, EmptyAvailableDevicesRejectsConstraint) {
+  FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/device:CPU:0");
+  absl::Status status = ValidateFunctionDeviceConstraints(fdef, {});
+  EXPECT_TRUE(IsInvalidArgument(status));
+  EXPECT_TRUE(Contains(status, "Could not satisfy device specification"));
+}
+
+}  // namespace tensorflow
diff --git a/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py b/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
index 0e121b0c07b..2d5d9817c15 100644
--- a/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
+++ b/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
@@ -1259,16 +1259,16 @@ class FunctionTest(xla_test.XLATestCase):
   def testDevicePlacementValidationWithJitCompile(self):
     """Test that jit_compile=True validates tf.device() constraints.
 
-    When using tf.device() with a non-existent device (e.g., GPU on CPU-only
-    machine), eager execution correctly fails with
+    When using tf.device() with a device that cannot exist (e.g., a
+    nonexistent device type), eager execution correctly fails with
     "Could not satisfy device specification". This test ensures jit_compile=True
     also validates device constraints instead of silently ignoring them.
 
     See https://github.com/tensorflow/tensorflow/issues/124880
     """
-    # Use a device specification that doesn't exist on this machine.
-    # The test runs on CPU, so GPU device should fail.
-    invalid_device = '/job:worker/replica:0/task:0/device:GPU:0'
+    # Use a device type that exists on no machine, so the expectation holds
+    # regardless of whether the test runner has GPUs.
+    invalid_device = '/device:NONEXISTENT:0'
 
     def compute(x):
       with ops.device(invalid_device):
@@ -1278,14 +1278,82 @@ class FunctionTest(xla_test.XLATestCase):
 
     # Eager execution should fail
     with self.assertRaisesRegex(errors.InvalidArgumentError,
-                                 'Could not satisfy device specification'):
+                                'Could not satisfy device specification'):
       compute(x).numpy()
 
     # jit_compile=True should also fail (this is the fix for #124880)
     with self.assertRaisesRegex(errors.InvalidArgumentError,
-                                 'Could not satisfy device specification'):
+                                'Could not satisfy device specification'):
       polymorphic_function.function(compute, jit_compile=True)(x).numpy()
 
+  def testDevicePlacementValidationInvalidDeviceIdWithJitCompile(self):
+    """Nonexistent device IDs must fail under jit_compile=True.
+
+    A partially-specified constraint with an invalid ID (e.g. 'CPU:99') must
+    not be silently accepted just because the device type exists.
+    """
+
+    def compute(x):
+      with ops.device('/device:CPU:99'):
+        return math_ops.add(x, x)
+
+    x = constant_op.constant([1.0, 2.0])
+
+    # Eager execution should fail
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'Could not satisfy device specification'):
+      compute(x).numpy()
+
+    # jit_compile=True should also fail
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'Could not satisfy device specification'):
+      polymorphic_function.function(compute, jit_compile=True)(x).numpy()
+
+  def testDevicePlacementValidationPartialConstraintWithJitCompile(self):
+    """Valid partial device constraints must pass under jit_compile=True."""
+
+    def compute(x):
+      with ops.device('/job:localhost'):
+        return math_ops.add(x, x)
+
+    x = constant_op.constant([1.0, 2.0])
+
+    # Eager execution succeeds
+    self.assertAllClose(compute(x), [2.0, 4.0])
+
+    # jit_compile=True must also succeed: the partial constraint
+    # '/job:localhost' is satisfied by the local devices.
+    self.assertAllClose(
+        polymorphic_function.function(compute, jit_compile=True)(x), [2.0, 4.0])
+
+  def testDevicePlacementValidationLocalDeviceNamesWithJitCompile(self):
+    """Local device names are validated under jit_compile=True."""
+    x = constant_op.constant([1.0, 2.0])
+
+    def compute_valid(x):
+      with ops.device('CPU:0'):
+        return math_ops.add(x, x)
+
+    # Local name 'CPU:0' resolves to the available CPU device.
+    self.assertAllClose(compute_valid(x), [2.0, 4.0])
+    self.assertAllClose(
+        polymorphic_function.function(compute_valid, jit_compile=True)(x),
+        [2.0, 4.0])
+
+    def compute_invalid(x):
+      with ops.device('CPU:99'):
+        return math_ops.add(x, x)
+
+    # A nonexistent local device ID must fail in both modes.
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'Could not satisfy device specification'):
+      compute_invalid(x).numpy()
+
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'Could not satisfy device specification'):
+      polymorphic_function.function(
+          compute_invalid, jit_compile=True)(x).numpy()
+
 
 if __name__ == '__main__':
   ops.enable_eager_execution()

From 822c946cea1a6932b7aadfb19f224401856e577f Mon Sep 17 00:00:00 2001
From: PrabinDevkota <prabindev789@gmail.com>
Date: Thu, 1 Oct 2026 20:06:54 +0530
Subject: [PATCH 0116/1000] Polish max-rank Pack and StridedSlice tests and
 their Abseil includes.

---
 tensorflow/core/kernels/BUILD                 | 11 +++++----
 tensorflow/core/kernels/pack_op.cc            |  1 +
 tensorflow/core/kernels/pack_op_test.cc       | 23 +++++++++++++++++--
 tensorflow/core/kernels/shape_ops.h           |  1 +
 tensorflow/core/kernels/shape_ops_test.cc     |  4 ++--
 tensorflow/core/ops/BUILD                     |  1 +
 tensorflow/core/ops/array_ops.cc              |  1 +
 tensorflow/core/util/strided_slice_op.cc      |  1 +
 .../kernel_tests/array_ops/array_ops_test.py  |  5 ++++
 9 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 8ca0324dcba..744a0672cbe 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -1133,7 +1133,9 @@ tf_kernel_library(
 tf_kernel_library(
     name = "pack_op",
     prefix = "pack_op",
-    deps = ARRAY_DEPS,
+    deps = ARRAY_DEPS + [
+        "@com_google_absl//absl/strings",
+    ],
 )
 
 tf_cc_test(
@@ -1153,7 +1155,6 @@ tf_cc_test(
         "//tensorflow/core:testlib",
         "//tensorflow/core/platform:status_matchers",
         "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:status_matchers",
         "@com_google_googletest//:gtest",
     ],
 )
@@ -1194,7 +1195,10 @@ tf_kernel_library(
     name = "shape_ops",
     features = ["-layering_check"],
     prefix = "shape_ops",
-    deps = ARRAY_DEPS + ["//tensorflow/core/common_runtime:dma_helper"],
+    deps = ARRAY_DEPS + [
+        "//tensorflow/core/common_runtime:dma_helper",
+        "@com_google_absl//absl/strings",
+    ],
 )
 
 tf_kernel_library(
@@ -2414,7 +2418,6 @@ tf_cc_test(
         "//tensorflow/core:testlib",
         "//tensorflow/core/platform:status_matchers",
         "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:status_matchers",
         "@com_google_googletest//:gtest",
     ],
 )
diff --git a/tensorflow/core/kernels/pack_op.cc b/tensorflow/core/kernels/pack_op.cc
index b5248401d21..086df0c1be4 100644
--- a/tensorflow/core/kernels/pack_op.cc
+++ b/tensorflow/core/kernels/pack_op.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include <limits>
 #include <vector>
 
+#include "absl/strings/str_cat.h"
 #include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
diff --git a/tensorflow/core/kernels/pack_op_test.cc b/tensorflow/core/kernels/pack_op_test.cc
index 47b5ef4f956..288ba4beb85 100644
--- a/tensorflow/core/kernels/pack_op_test.cc
+++ b/tensorflow/core/kernels/pack_op_test.cc
@@ -19,7 +19,6 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/status/status.h"
-#include "absl/status/status_matchers.h"
 #include "tensorflow/core/framework/fake_input.h"
 #include "tensorflow/core/framework/node_def_builder.h"
 #include "tensorflow/core/framework/tensor.h"
@@ -29,12 +28,13 @@ limitations under the License.
 #include "tensorflow/core/kernels/ops_testutil.h"
 #include "tensorflow/core/kernels/ops_util.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/status_matchers.h"
 #include "tensorflow/core/platform/test.h"
 
 namespace tensorflow {
 namespace {
 
-using ::absl_testing::StatusIs;
+using ::tensorflow::testing::StatusIs;
 using ::testing::HasSubstr;
 
 absl::Status MakeZeroElementShape(int rank, TensorShape* shape) {
@@ -70,6 +70,25 @@ TEST_F(PackOpTest, PackingMaxRankSucceeds) {
   EXPECT_EQ(GetOutput(0)->dim_size(0), 1);
 }
 
+TEST_F(PackOpTest, PackingMultipleTensorsMaxRankSucceeds) {
+  TF_ASSERT_OK(NodeDefBuilder("pack", "Pack")
+                   .Input(FakeInput(2, DT_FLOAT))
+                   .Attr("N", 2)
+                   .Attr("T", DT_FLOAT)
+                   .Attr("axis", 0)
+                   .Finalize(node_def()));
+  TF_ASSERT_OK(InitOp());
+  TensorShape input_shape;
+  TF_ASSERT_OK(
+      MakeZeroElementShape(TensorShape::MaxDimensions() - 1, &input_shape));
+  AddInput<float>(input_shape, [](int) { return 0.0f; });
+  AddInput<float>(input_shape, [](int) { return 0.0f; });
+
+  TF_ASSERT_OK(RunOpKernel());
+  EXPECT_EQ(GetOutput(0)->dims(), TensorShape::MaxDimensions());
+  EXPECT_EQ(GetOutput(0)->dim_size(0), 2);
+}
+
 TEST_F(PackOpTest, PackingBeyondMaxRankFails) {
   MakeOp();
   TensorShape input_shape;
diff --git a/tensorflow/core/kernels/shape_ops.h b/tensorflow/core/kernels/shape_ops.h
index f4a07570257..3209e469b6e 100644
--- a/tensorflow/core/kernels/shape_ops.h
+++ b/tensorflow/core/kernels/shape_ops.h
@@ -21,6 +21,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/container/inlined_vector.h"
+#include "absl/strings/str_cat.h"
 #include "tensorflow/core/common_runtime/dma_helper.h"
 #include "tensorflow/core/framework/bounds_check.h"
 #include "tensorflow/core/framework/op_kernel.h"
diff --git a/tensorflow/core/kernels/shape_ops_test.cc b/tensorflow/core/kernels/shape_ops_test.cc
index c39ba923142..2af70a685bf 100644
--- a/tensorflow/core/kernels/shape_ops_test.cc
+++ b/tensorflow/core/kernels/shape_ops_test.cc
@@ -21,7 +21,6 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/status/status.h"
-#include "absl/status/status_matchers.h"
 #include "tensorflow/core/common_runtime/kernel_benchmark_testlib.h"
 #include "tensorflow/core/framework/allocator.h"
 #include "tensorflow/core/framework/fake_input.h"
@@ -37,13 +36,14 @@ limitations under the License.
 #include "tensorflow/core/kernels/ops_testutil.h"
 #include "tensorflow/core/kernels/ops_util.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/status_matchers.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/platform/test_benchmark.h"
 
 namespace tensorflow {
 namespace {
 
-using ::absl_testing::StatusIs;
+using ::tensorflow::testing::StatusIs;
 using ::testing::HasSubstr;
 
 absl::Status MakeZeroElementShape(int rank, TensorShape* shape) {
diff --git a/tensorflow/core/ops/BUILD b/tensorflow/core/ops/BUILD
index ce591c3b24e..8b781e73c21 100644
--- a/tensorflow/core/ops/BUILD
+++ b/tensorflow/core/ops/BUILD
@@ -165,6 +165,7 @@ tf_gen_op_libs(
     deps = [
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/strings",
     ],
 )
 
diff --git a/tensorflow/core/ops/array_ops.cc b/tensorflow/core/ops/array_ops.cc
index e9a772a395a..9a34f9aaf00 100644
--- a/tensorflow/core/ops/array_ops.cc
+++ b/tensorflow/core/ops/array_ops.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include <ostream>
 #include <vector>
 
+#include "absl/strings/str_cat.h"
 #include "tensorflow/core/framework/common_shape_fns.h"
 #include "tensorflow/core/framework/full_type.pb.h"
 #include "tensorflow/core/framework/kernel_shape_util.h"
diff --git a/tensorflow/core/util/strided_slice_op.cc b/tensorflow/core/util/strided_slice_op.cc
index d9981276bdd..bb33ab43786 100644
--- a/tensorflow/core/util/strided_slice_op.cc
+++ b/tensorflow/core/util/strided_slice_op.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include <iterator>
 #include <utility>
 
+#include "absl/strings/str_cat.h"
 #include "tensorflow/core/framework/bounds_check.h"
 #include "tensorflow/core/lib/core/status.h"
 
diff --git a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
index 580749a93d7..95cfbeba4b6 100644
--- a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
@@ -780,6 +780,11 @@ class StridedSliceTest(test_util.TensorFlowTestCase):
     value = array_ops.fill(
         [0] + [1] * 252,
         constant_op.constant(0.0, dtype=dtypes.float32))
+    # User-facing slice syntax.
+    sliced_syntax = value[None, ...]
+    self.assertEqual(254, self.evaluate(array_ops.rank(sliced_syntax)))
+
+    # Low-level op.
     sliced = gen_array_ops.strided_slice(
         value,
         begin=[0],

From 0d58f490a88cd0ef6e73298912f1ee23dc8cdbf2 Mon Sep 17 00:00:00 2001
From: alwaysprince05 <prince.p24@medhaviskillsuniversity.edu.in>
Date: Thu, 1 Oct 2026 20:08:06 +0530
Subject: [PATCH 0117/1000] Fix ambiguous IsInvalidArgument in
 validate_function_devices_test

The unqualified IsInvalidArgument(status) call was ambiguous: argument-
dependent lookup finds absl::IsInvalidArgument (declared in
absl/status/status.h) in addition to the test-local helper, so the
linux-x86-cpu build failed to compile the test. Drop the helper and
call absl::IsInvalidArgument explicitly.
---
 .../eager/validate_function_devices_test.cc      | 16 ++++++----------
 1 file changed, 6 insertions(+), 10 deletions(-)

diff --git a/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc b/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc
index 4bd82e0a14d..c131c38daf3 100644
--- a/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc
+++ b/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc
@@ -48,10 +48,6 @@ std::vector<DeviceAttributes> LocalCpuDevices() {
   return {attrs};
 }
 
-bool IsInvalidArgument(const absl::Status& status) {
-  return status.code() == absl::StatusCode::kInvalidArgument;
-}
-
 bool Contains(const absl::Status& status, const std::string& substring) {
   return std::string(status.message()).find(substring) != std::string::npos;
 }
@@ -101,7 +97,7 @@ TEST(ValidateFunctionDevicesTest, UnsatisfiedJobFails) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/job:worker");
   absl::Status status =
       ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
-  EXPECT_TRUE(IsInvalidArgument(status));
+  EXPECT_TRUE(absl::IsInvalidArgument(status));
   EXPECT_TRUE(Contains(status, "Could not satisfy device specification"));
   EXPECT_TRUE(Contains(status, "/job:worker"));
   EXPECT_TRUE(Contains(status, "/job:localhost/replica:0/task:0/device:CPU:0"));
@@ -113,7 +109,7 @@ TEST(ValidateFunctionDevicesTest, InvalidDeviceIdFails) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/device:CPU:99");
   absl::Status status =
       ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
-  EXPECT_TRUE(IsInvalidArgument(status));
+  EXPECT_TRUE(absl::IsInvalidArgument(status));
   EXPECT_TRUE(Contains(status, "Could not satisfy device specification"));
 }
 
@@ -122,7 +118,7 @@ TEST(ValidateFunctionDevicesTest, InvalidDeviceTypeFails) {
       MakeFunctionWithDevice("add", "AddV2", "/device:NONEXISTENT:0");
   absl::Status status =
       ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
-  EXPECT_TRUE(IsInvalidArgument(status));
+  EXPECT_TRUE(absl::IsInvalidArgument(status));
   EXPECT_TRUE(Contains(status, "Could not satisfy device specification"));
 }
 
@@ -130,7 +126,7 @@ TEST(ValidateFunctionDevicesTest, LocalDeviceNameInvalidIdFails) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "CPU:99");
   absl::Status status =
       ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
-  EXPECT_TRUE(IsInvalidArgument(status));
+  EXPECT_TRUE(absl::IsInvalidArgument(status));
   EXPECT_TRUE(Contains(status, "Could not satisfy device specification"));
 }
 
@@ -138,7 +134,7 @@ TEST(ValidateFunctionDevicesTest, MalformedDeviceSpecFails) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "not_a_device");
   absl::Status status =
       ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
-  EXPECT_TRUE(IsInvalidArgument(status));
+  EXPECT_TRUE(absl::IsInvalidArgument(status));
   EXPECT_TRUE(Contains(status, "Malformed device specification"));
   EXPECT_TRUE(Contains(status, "not_a_device"));
 }
@@ -154,7 +150,7 @@ TEST(ValidateFunctionDevicesTest, ErrorMentionsOperation) {
 TEST(ValidateFunctionDevicesTest, EmptyAvailableDevicesRejectsConstraint) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/device:CPU:0");
   absl::Status status = ValidateFunctionDeviceConstraints(fdef, {});
-  EXPECT_TRUE(IsInvalidArgument(status));
+  EXPECT_TRUE(absl::IsInvalidArgument(status));
   EXPECT_TRUE(Contains(status, "Could not satisfy device specification"));
 }
 

From db9fac91d78743f7093db1caa7e1c53e797c39dd Mon Sep 17 00:00:00 2001
From: endorphin13 <elshad.toklayev.business@gmail.com>
Date: Thu, 1 Oct 2026 21:06:17 +0500
Subject: [PATCH 0118/1000] Disable XLA for the empty-FFT-axis shape tests

The XLA RFFT/IRFFT kernels zero-pad an empty FFT axis up to fft_length before
the transform, so under XLA the output keeps the fft_length-derived size and
is zero-filled rather than collapsing to zero. That is not a disclosure, but
it fails the exact-shape assertions in these two tests. Skip them under XLA,
as suggested in review.

The FFTND/RFFTND tests are unaffected: those ops have no XLA kernel.
---
 tensorflow/python/kernel_tests/signal/fft_ops_test.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow/python/kernel_tests/signal/fft_ops_test.py b/tensorflow/python/kernel_tests/signal/fft_ops_test.py
index b93a45dd22a..ad2c7dae515 100644
--- a/tensorflow/python/kernel_tests/signal/fft_ops_test.py
+++ b/tensorflow/python/kernel_tests/signal/fft_ops_test.py
@@ -630,6 +630,7 @@ class RFFTOpsTest(BaseFFTOpsTest, parameterized.TestCase):
 
   @parameterized.parameters(
       itertools.product(VALID_FFT_RANKS, range(3), (np.float32, np.float64)))
+  @test_util.disable_xla("XLA zero-pads empty FFT axes to fft_length")
   def test_empty_with_explicit_fft_length(self, rank, extra_dims, np_rtype):
     # An empty FFT axis skips the "input dimension must be at least
     # fft_length" requirement. `fft_length` must not then size that axis in
@@ -656,6 +657,7 @@ class RFFTOpsTest(BaseFFTOpsTest, parameterized.TestCase):
     self.assertEqual(0, out_bwd.size)
 
   @parameterized.parameters((np.float32,), (np.float64,))
+  @test_util.disable_xla("XLA zero-pads empty FFT axes to fft_length")
   def test_empty_axis_next_to_non_empty_axis(self, np_rtype):
     # Only an empty FFT axis stays empty. The other FFT axis is still sized
     # from `fft_length`, and the batch dimension is untouched.

From f45365eaddd7acad4c6bc69d31f9674b1c23ace9 Mon Sep 17 00:00:00 2001
From: Sai Ramineni <csramineni@gmail.com>
Date: Thu, 1 Oct 2026 13:35:08 -0400
Subject: [PATCH 0119/1000] Skip ragged index traversal when sparse values are
 empty

---
 .../kernels/ragged_tensor_to_sparse_kernel.cc | 78 ++++++++++---------
 .../ragged_tensor_to_sparse_kernel_test.cc    | 25 ++++++
 2 files changed, 65 insertions(+), 38 deletions(-)

diff --git a/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel.cc b/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel.cc
index dcb8c65e67c..2ea95b5254b 100644
--- a/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel.cc
+++ b/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel.cc
@@ -59,7 +59,6 @@ class RaggedTensorToSparseOp : public OpKernel {
     //   dimension.
     // - `index_middle` is the index in the last ragged dimension.
     // - `index_suffix` is the index in the dense value dimensions.
-    std::vector<int64_t> index_prefix(rt_nested_splits_len);
     std::vector<std::vector<int64_t>> index_suffixes;
     if (rt_dense_values_in.NumElements() != 0) {
       index_suffixes = MakeIndexSuffixes(rt_dense_values_in.shape());
@@ -77,52 +76,55 @@ class RaggedTensorToSparseOp : public OpKernel {
                                           &sparse_indices_out));
     auto sparse_indices = sparse_indices_out->tensor<int64_t, 2>();
 
-    // pos[i] is the current position in rt_nested_splits[i].  final_pos is a
-    // reference to make it easier to refer to pos[-1].
-    std::vector<int64_t> pos(rt_nested_splits_len);
-    int64_t& final_pos = pos[rt_nested_splits_len - 1];
+    if (nvals > 0) {
+      std::vector<int64_t> index_prefix(rt_nested_splits_len);
+      // pos[i] is the current position in rt_nested_splits[i].  final_pos is a
+      // reference to make it easier to refer to pos[-1].
+      std::vector<int64_t> pos(rt_nested_splits_len);
+      int64_t& final_pos = pos[rt_nested_splits_len - 1];
 
-    // Each iteration through the loop, we increment pos[-1], and add indices
-    // for all the values corresponding to
-    // rt_nested_splits[-1][pos[-1]:pos[-1]+1].
-    int64_t next_index = 0;
-    int64_t max_final_pos = rt_nested_splits.back().size() - 1;
-    for (; final_pos < max_final_pos; ++final_pos) {
-      // Update `pos` to skip over completed elements (i.e., elements where
-      // we have already generated indices for all contained values).
-      for (int dim = rt_nested_splits_len - 2; dim >= 0; --dim) {
-        while (IsCompleted(pos, dim, rt_nested_splits)) {
-          pos[dim] += 1;
+      // Each iteration through the loop, we increment pos[-1], and add indices
+      // for all the values corresponding to
+      // rt_nested_splits[-1][pos[-1]:pos[-1]+1].
+      int64_t next_index = 0;
+      int64_t max_final_pos = rt_nested_splits.back().size() - 1;
+      for (; final_pos < max_final_pos; ++final_pos) {
+        // Update `pos` to skip over completed elements (i.e., elements where
+        // we have already generated indices for all contained values).
+        for (int dim = rt_nested_splits_len - 2; dim >= 0; --dim) {
+          while (IsCompleted(pos, dim, rt_nested_splits)) {
+            pos[dim] += 1;
+          }
         }
-      }
 
-      // Update index_prefix.
-      for (int dim = 0; dim < index_prefix.size(); ++dim) {
-        int64_t start = dim > 0 ? rt_nested_splits[dim - 1](pos[dim - 1]) : 0;
-        index_prefix[dim] = pos[dim] - start;
-      }
+        // Update index_prefix.
+        for (int dim = 0; dim < index_prefix.size(); ++dim) {
+          int64_t start = dim > 0 ? rt_nested_splits[dim - 1](pos[dim - 1]) : 0;
+          index_prefix[dim] = pos[dim] - start;
+        }
 
-      // Get length of the final-ragged-dimension slice.
-      const auto& final_splits = rt_nested_splits[rt_nested_splits_len - 1];
-      int64_t slice_len = final_splits(final_pos + 1) - final_splits(final_pos);
+        // Get length of the final-ragged-dimension slice.
+        const auto& final_splits = rt_nested_splits[rt_nested_splits_len - 1];
+        int64_t slice_len = final_splits(final_pos + 1) - final_splits(final_pos);
 
-      // Add sparse_indices for this slice.
-      for (int64_t i = 0; i < slice_len; ++i) {
-        for (const auto& index_suffix : index_suffixes) {
-          int dim = 0;
-          for (int64_t index : index_prefix) {  // index_prefix
-            sparse_indices(next_index, dim++) = index;
-          }
-          sparse_indices(next_index, dim++) = i;  // index_middle
-          for (int64_t index : index_suffix) {    // index_suffix
-            sparse_indices(next_index, dim++) = index;
+        // Add sparse_indices for this slice.
+        for (int64_t i = 0; i < slice_len; ++i) {
+          for (const auto& index_suffix : index_suffixes) {
+            int dim = 0;
+            for (int64_t index : index_prefix) {  // index_prefix
+              sparse_indices(next_index, dim++) = index;
+            }
+            sparse_indices(next_index, dim++) = i;  // index_middle
+            for (int64_t index : index_suffix) {    // index_suffix
+              sparse_indices(next_index, dim++) = index;
+            }
+            DCHECK_EQ(dim, indices_len);
+            ++next_index;
           }
-          DCHECK_EQ(dim, indices_len);
-          ++next_index;
         }
       }
+      DCHECK_EQ(next_index, nvals);
     }
-    DCHECK_EQ(next_index, nvals);
 
     // Output the `sparse_values` Tensor.
     if (rt_dense_values_in.dims() == 1) {
diff --git a/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel_test.cc b/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel_test.cc
index a391784d23b..b7e03f2602f 100644
--- a/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel_test.cc
+++ b/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel_test.cc
@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <cstdint>
 #include <limits>
 
 #include "tensorflow/core/framework/fake_input.h"
@@ -90,6 +91,30 @@ TEST_F(RaggedTensorToSparseTest, EmptyRows) {
                                    test::AsTensor<int64_t>({5, 4}));
 }
 
+TEST_F(RaggedTensorToSparseTest, EmptyValuesRank1) {
+  BuildRaggedTensorToSparseGraph<int>({{0, 0, 0}}, TensorShape({0}), {});
+  TF_ASSERT_OK(RunOpKernel());
+  test::ExpectTensorEqual<int64_t>(
+      *GetOutput(kSparseIndicesOutput),
+      test::AsTensor<int64_t>({}, TensorShape({0, 2})));
+  test::ExpectTensorEqual<int>(*GetOutput(kSparseValuesOutput),
+                             test::AsTensor<int>({}));
+  test::ExpectTensorEqual<int64_t>(*GetOutput(kSparseDenseShapeOutput),
+                                 test::AsTensor<int64_t>({2, 0}));
+}
+
+TEST_F(RaggedTensorToSparseTest, EmptyValuesWithInnerZeroDimension) {
+  BuildRaggedTensorToSparseGraph<int>({{0, 2}}, TensorShape({2, 0}), {});
+  TF_ASSERT_OK(RunOpKernel());
+  test::ExpectTensorEqual<int64_t>(
+      *GetOutput(kSparseIndicesOutput),
+      test::AsTensor<int64_t>({}, TensorShape({0, 3})));
+  test::ExpectTensorEqual<int>(*GetOutput(kSparseValuesOutput),
+                             test::AsTensor<int>({}));
+  test::ExpectTensorEqual<int64_t>(*GetOutput(kSparseDenseShapeOutput),
+                                 test::AsTensor<int64_t>({1, 2, 0}));
+}
+
 TEST_F(RaggedTensorToSparseTest, EmptyValuesWithLargeDenseSuffix) {
   constexpr int64_t kSuffixSize = std::numeric_limits<int32_t>::max();
   BuildRaggedTensorToSparseGraph<int>(

From 4b63d34b8db9c6da5640b410056a8ad5f27dd65e Mon Sep 17 00:00:00 2001
From: Sai Ramineni <csramineni@gmail.com>
Date: Thu, 1 Oct 2026 15:24:43 -0400
Subject: [PATCH 0120/1000] Validate ragged values rank and use checked element
 count

---
 .../kernels/ragged_tensor_to_sparse_kernel.cc   | 10 ++++++----
 tensorflow/python/ops/ragged/BUILD              |  1 +
 .../ops/ragged/ragged_to_sparse_op_test.py      | 17 +++++++++++++++++
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel.cc b/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel.cc
index 2ea95b5254b..2b35ff44b08 100644
--- a/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel.cc
+++ b/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel.cc
@@ -12,6 +12,7 @@ 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.
 ==============================================================================*/
+#include <cstddef>
 #include <limits>
 #include <memory>
 #include <string>
@@ -65,9 +66,7 @@ class RaggedTensorToSparseOp : public OpKernel {
     }
 
     // Allocate the `sparse_indices` output tensor.
-    const int64_t nvals =
-        (rt_nested_splits.back()(rt_nested_splits.back().size() - 1) *
-         index_suffixes.size());
+    const int64_t nvals = rt_dense_values_in.NumElements();
     const int64_t indices_len =
         rt_nested_splits_len + rt_dense_values_in.dims();
     Tensor* sparse_indices_out = nullptr;
@@ -98,7 +97,7 @@ class RaggedTensorToSparseOp : public OpKernel {
         }
 
         // Update index_prefix.
-        for (int dim = 0; dim < index_prefix.size(); ++dim) {
+        for (size_t dim = 0; dim < index_prefix.size(); ++dim) {
           int64_t start = dim > 0 ? rt_nested_splits[dim - 1](pos[dim - 1]) : 0;
           index_prefix[dim] = pos[dim] - start;
         }
@@ -163,6 +162,9 @@ class RaggedTensorToSparseOp : public OpKernel {
   static absl::Status ValidateInputs(
       std::vector<ConstFlatSplits> rt_nested_splits,
       const Tensor& rt_dense_values_in) {
+    if (rt_dense_values_in.dims() == 0) {
+      return InvalidArgument("rt_dense_values must have rank >= 1.");
+    }
     for (int i = 0; i < rt_nested_splits.size(); ++i) {
       if (rt_nested_splits[i].size() == 0) {
         return absl::InvalidArgumentError("ragged splits may not be empty.");
diff --git a/tensorflow/python/ops/ragged/BUILD b/tensorflow/python/ops/ragged/BUILD
index 2f1bb55b55d..d48db0fb6a6 100644
--- a/tensorflow/python/ops/ragged/BUILD
+++ b/tensorflow/python/ops/ragged/BUILD
@@ -1069,6 +1069,7 @@ py_test(
         "//tensorflow/python/ops:array_ops",
         "//tensorflow/python/ops:gradients_impl",
         "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:ragged_conversion_ops_gen",
         "//tensorflow/python/platform:test",
     ],
 )
diff --git a/tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py b/tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py
index c2a35e8c680..4693494b288 100644
--- a/tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py
@@ -19,6 +19,7 @@ from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_ragged_conversion_ops
 from tensorflow.python.ops import gradients_impl
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops.ragged import ragged_factory_ops
@@ -175,6 +176,22 @@ class RaggedTensorToSparseOpTest(test_util.TensorFlowTestCase):
                                 empty_splits_error):
       self.evaluate(bad_rt5.to_sparse())
 
+  def testRejectsScalarValues(self):
+    for splits_dtype in (dtypes.int32, dtypes.int64):
+      with self.subTest(splits_dtype=splits_dtype):
+        values = array_ops.zeros([], dtype=dtypes.int32)
+        splits = [math_ops.cast([0, 1], splits_dtype)]
+        if not context.executing_eagerly():
+          with self.assertRaisesRegex(
+              ValueError, "Shape must be at least rank 1"):
+            gen_ragged_conversion_ops.ragged_tensor_to_sparse(
+                rt_nested_splits=splits, rt_dense_values=values)
+          values = array_ops.placeholder_with_default(values, shape=None)
+        with self.assertRaisesRegex(
+            errors.InvalidArgumentError, "rt_dense_values must have rank >= 1"):
+          self.evaluate(gen_ragged_conversion_ops.ragged_tensor_to_sparse(
+              rt_nested_splits=splits, rt_dense_values=values))
+
   def testGradient(self):
     if context.executing_eagerly():
       return

From c1b988c98dee6d00d764a41092dd65561791b866 Mon Sep 17 00:00:00 2001
From: alwaysprince05 <alwaysprince05@users.noreply.github.com>
Date: Thu, 1 Oct 2026 22:34:28 +0530
Subject: [PATCH 0121/1000] Address review: strict placement in eager tests,
 IWYU headers, StatusIs matchers

Disable soft device placement around the eager assertions in the three
device-placement Python tests so they assert the real error, replace
function.h with function.pb.h/node_def.pb.h in the validator, and use
canonical StatusIs/HasSubstr matchers in the C++ unit test.
---
 .../eager/validate_function_devices.cc        |  3 +-
 .../eager/validate_function_devices.h         |  2 +-
 .../eager/validate_function_devices_test.cc   | 56 ++++++++++++-------
 .../polymorphic_function_xla_jit_test.py      | 42 ++++++++++----
 4 files changed, 68 insertions(+), 35 deletions(-)

diff --git a/tensorflow/core/common_runtime/eager/validate_function_devices.cc b/tensorflow/core/common_runtime/eager/validate_function_devices.cc
index 4476c707360..b1bfaafa272 100644
--- a/tensorflow/core/common_runtime/eager/validate_function_devices.cc
+++ b/tensorflow/core/common_runtime/eager/validate_function_devices.cc
@@ -21,7 +21,8 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
-#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/function.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
 #include "tensorflow/core/util/device_name_utils.h"
 
 namespace tensorflow {
diff --git a/tensorflow/core/common_runtime/eager/validate_function_devices.h b/tensorflow/core/common_runtime/eager/validate_function_devices.h
index e75d30c7c3e..54253ed49d5 100644
--- a/tensorflow/core/common_runtime/eager/validate_function_devices.h
+++ b/tensorflow/core/common_runtime/eager/validate_function_devices.h
@@ -19,7 +19,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/status/status.h"
-#include "tensorflow/core/framework/function.h"
+#include "tensorflow/core/framework/function.pb.h"
 #include "tensorflow/core/framework/device_attributes.pb.h"
 
 namespace tensorflow {
diff --git a/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc b/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc
index c131c38daf3..3a0a4c18a72 100644
--- a/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc
+++ b/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include <string>
 #include <vector>
 
+#include <gmock/gmock.h>
 #include "absl/status/status.h"
 #include "tensorflow/core/framework/device_attributes.pb.h"
 #include "tensorflow/core/framework/function.pb.h"
@@ -48,10 +49,6 @@ std::vector<DeviceAttributes> LocalCpuDevices() {
   return {attrs};
 }
 
-bool Contains(const absl::Status& status, const std::string& substring) {
-  return std::string(status.message()).find(substring) != std::string::npos;
-}
-
 }  // namespace
 
 TEST(ValidateFunctionDevicesTest, EmptyFunctionIsOk) {
@@ -97,10 +94,14 @@ TEST(ValidateFunctionDevicesTest, UnsatisfiedJobFails) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/job:worker");
   absl::Status status =
       ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
-  EXPECT_TRUE(absl::IsInvalidArgument(status));
-  EXPECT_TRUE(Contains(status, "Could not satisfy device specification"));
-  EXPECT_TRUE(Contains(status, "/job:worker"));
-  EXPECT_TRUE(Contains(status, "/job:localhost/replica:0/task:0/device:CPU:0"));
+  EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
+                          absl::StatusCode::kInvalidArgument,
+                          ::testing::AllOf(
+                              ::testing::HasSubstr(
+                                  "Could not satisfy device specification"),
+                              ::testing::HasSubstr("/job:worker"),
+                              ::testing::HasSubstr("/job:localhost/replica:0/"
+                                                   "task:0/device:CPU:0"))));
 }
 
 TEST(ValidateFunctionDevicesTest, InvalidDeviceIdFails) {
@@ -109,8 +110,10 @@ TEST(ValidateFunctionDevicesTest, InvalidDeviceIdFails) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/device:CPU:99");
   absl::Status status =
       ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
-  EXPECT_TRUE(absl::IsInvalidArgument(status));
-  EXPECT_TRUE(Contains(status, "Could not satisfy device specification"));
+  EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
+                          absl::StatusCode::kInvalidArgument,
+                          ::testing::HasSubstr(
+                              "Could not satisfy device specification")));
 }
 
 TEST(ValidateFunctionDevicesTest, InvalidDeviceTypeFails) {
@@ -118,40 +121,51 @@ TEST(ValidateFunctionDevicesTest, InvalidDeviceTypeFails) {
       MakeFunctionWithDevice("add", "AddV2", "/device:NONEXISTENT:0");
   absl::Status status =
       ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
-  EXPECT_TRUE(absl::IsInvalidArgument(status));
-  EXPECT_TRUE(Contains(status, "Could not satisfy device specification"));
+  EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
+                          absl::StatusCode::kInvalidArgument,
+                          ::testing::HasSubstr(
+                              "Could not satisfy device specification")));
 }
 
 TEST(ValidateFunctionDevicesTest, LocalDeviceNameInvalidIdFails) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "CPU:99");
   absl::Status status =
       ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
-  EXPECT_TRUE(absl::IsInvalidArgument(status));
-  EXPECT_TRUE(Contains(status, "Could not satisfy device specification"));
+  EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
+                          absl::StatusCode::kInvalidArgument,
+                          ::testing::HasSubstr(
+                              "Could not satisfy device specification")));
 }
 
 TEST(ValidateFunctionDevicesTest, MalformedDeviceSpecFails) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "not_a_device");
   absl::Status status =
       ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
-  EXPECT_TRUE(absl::IsInvalidArgument(status));
-  EXPECT_TRUE(Contains(status, "Malformed device specification"));
-  EXPECT_TRUE(Contains(status, "not_a_device"));
+  EXPECT_THAT(
+      status, ::tensorflow::testing::StatusIs(
+                  absl::StatusCode::kInvalidArgument,
+                  ::testing::AllOf(
+                      ::testing::HasSubstr("Malformed device specification"),
+                      ::testing::HasSubstr("not_a_device"))));
 }
 
 TEST(ValidateFunctionDevicesTest, ErrorMentionsOperation) {
   FunctionDef fdef = MakeFunctionWithDevice("my_add", "AddV2", "/job:worker");
   absl::Status status =
       ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
-  EXPECT_TRUE(Contains(status, "my_add"));
-  EXPECT_TRUE(Contains(status, "AddV2"));
+  EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
+                          absl::StatusCode::kInvalidArgument,
+                          ::testing::AllOf(::testing::HasSubstr("my_add"),
+                                           ::testing::HasSubstr("AddV2"))));
 }
 
 TEST(ValidateFunctionDevicesTest, EmptyAvailableDevicesRejectsConstraint) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/device:CPU:0");
   absl::Status status = ValidateFunctionDeviceConstraints(fdef, {});
-  EXPECT_TRUE(absl::IsInvalidArgument(status));
-  EXPECT_TRUE(Contains(status, "Could not satisfy device specification"));
+  EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
+                          absl::StatusCode::kInvalidArgument,
+                          ::testing::HasSubstr(
+                              "Could not satisfy device specification")));
 }
 
 }  // namespace tensorflow
diff --git a/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py b/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
index 2d5d9817c15..ac56b0424c6 100644
--- a/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
+++ b/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
@@ -1276,10 +1276,16 @@ class FunctionTest(xla_test.XLATestCase):
 
     x = constant_op.constant([1.0, 2.0])
 
-    # Eager execution should fail
-    with self.assertRaisesRegex(errors.InvalidArgumentError,
-                                'Could not satisfy device specification'):
-      compute(x).numpy()
+    # Eager mode soft-places onto CPU by default, so the constraint is only
+    # enforced with soft device placement disabled (see #124880).
+    old_soft_placement = context.context().soft_device_placement
+    try:
+      context.context().soft_device_placement = False
+      with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                  'Could not satisfy device specification'):
+        compute(x).numpy()
+    finally:
+      context.context().soft_device_placement = old_soft_placement
 
     # jit_compile=True should also fail (this is the fix for #124880)
     with self.assertRaisesRegex(errors.InvalidArgumentError,
@@ -1299,10 +1305,16 @@ class FunctionTest(xla_test.XLATestCase):
 
     x = constant_op.constant([1.0, 2.0])
 
-    # Eager execution should fail
-    with self.assertRaisesRegex(errors.InvalidArgumentError,
-                                'Could not satisfy device specification'):
-      compute(x).numpy()
+    # Eager mode soft-places onto CPU by default, so the constraint is only
+    # enforced with soft device placement disabled (see #124880).
+    old_soft_placement = context.context().soft_device_placement
+    try:
+      context.context().soft_device_placement = False
+      with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                  'Could not satisfy device specification'):
+        compute(x).numpy()
+    finally:
+      context.context().soft_device_placement = old_soft_placement
 
     # jit_compile=True should also fail
     with self.assertRaisesRegex(errors.InvalidArgumentError,
@@ -1344,10 +1356,16 @@ class FunctionTest(xla_test.XLATestCase):
       with ops.device('CPU:99'):
         return math_ops.add(x, x)
 
-    # A nonexistent local device ID must fail in both modes.
-    with self.assertRaisesRegex(errors.InvalidArgumentError,
-                                'Could not satisfy device specification'):
-      compute_invalid(x).numpy()
+    # A nonexistent local device ID must fail in both modes. Eager mode
+    # soft-places by default, so disable it for the eager assertion.
+    old_soft_placement = context.context().soft_device_placement
+    try:
+      context.context().soft_device_placement = False
+      with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                  'Could not satisfy device specification'):
+        compute_invalid(x).numpy()
+    finally:
+      context.context().soft_device_placement = old_soft_placement
 
     with self.assertRaisesRegex(errors.InvalidArgumentError,
                                 'Could not satisfy device specification'):

From 40039b6354df806f722a8f0f74e81a40c7b8ac48 Mon Sep 17 00:00:00 2001
From: alwaysprince05 <alwaysprince05@users.noreply.github.com>
Date: Fri, 2 Oct 2026 01:35:47 +0530
Subject: [PATCH 0122/1000] Address review: soft-placement parity,
 nested-function validation

- Gate device-constraint validation on !ctx.AllowSoftPlacement() so
  jit_compile=True no longer errors on softly-placed invocations and
  matches standard eager behavior.
- Recursively validate device constraints in nested functions reached
  through (Stateful)PartitionedCall, with a visited-functions guard
  against cycles and a per-device parse cache.
- Move device-placement regression tests to a dedicated
  DevicePlacementValidationTest that disables soft device placement in
  setUp/tearDown; add a nested-function regression test.
- Add device_attributes.pb.h include to execute.cc (IWYU) and the
  flat_hash_set dep to :validate_function_devices.
---
 tensorflow/core/common_runtime/eager/BUILD    |   2 +
 .../core/common_runtime/eager/execute.cc      |  10 +-
 .../eager/validate_function_devices.cc        | 134 ++++++++++++------
 .../eager/validate_function_devices.h         |  16 ++-
 .../eager/validate_function_devices_test.cc   | 117 +++++++++++++--
 .../polymorphic_function_xla_jit_test.py      |  94 +++++++-----
 6 files changed, 275 insertions(+), 98 deletions(-)

diff --git a/tensorflow/core/common_runtime/eager/BUILD b/tensorflow/core/common_runtime/eager/BUILD
index 80cd3b39461..201df308507 100644
--- a/tensorflow/core/common_runtime/eager/BUILD
+++ b/tensorflow/core/common_runtime/eager/BUILD
@@ -714,6 +714,7 @@ cc_library(
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
     ],
@@ -730,6 +731,7 @@ tf_cc_test(
         "//tensorflow/core:test_main",
         "//tensorflow/core/platform:status_matchers",
         "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest",
     ],
 )
 
diff --git a/tensorflow/core/common_runtime/eager/execute.cc b/tensorflow/core/common_runtime/eager/execute.cc
index d976a59ffee..0d730a4cb99 100644
--- a/tensorflow/core/common_runtime/eager/execute.cc
+++ b/tensorflow/core/common_runtime/eager/execute.cc
@@ -40,6 +40,7 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/eager/validate_function_devices.h"
 #include "tensorflow/core/common_runtime/int32_fulltype.h"
 #include "tensorflow/core/framework/cancellation.h"
+#include "tensorflow/core/framework/device_attributes.pb.h"
 #include "tensorflow/core/framework/full_type.pb.h"
 #include "tensorflow/core/framework/function.pb.h"
 #include "tensorflow/core/framework/kernel_def.pb.h"
@@ -1415,16 +1416,21 @@ absl::Status GetOrCreateKernelAndDevice(
         // Only validate when _XlaMustCompile was explicitly set (user's
         // jit_compile=True), not when compile_with_xla is inferred from
         // the function being on an XLA device (TPU/XLA_GPU/XLA_CPU).
+        //
+        // Only strictly enforce device constraints when soft placement is
+        // disabled: with soft placement enabled, an unsatisfiable constraint
+        // would be relaxed by the runtime and execution falls back to an
+        // available device, matching standard eager behavior.
         bool has_xla_must_compile = false;
         GetFuncAttr(op, ctx, kXlaMustCompileAttr, &has_xla_must_compile)
             .IgnoreError();
-        if (has_xla_must_compile) {
+        if (has_xla_must_compile && !ctx.AllowSoftPlacement()) {
           const FunctionDef* fdef = op->GetFunctionDef();
           if (fdef != nullptr) {
             std::vector<DeviceAttributes> device_attrs;
             ctx.ListDevices(&device_attrs);
             TF_RETURN_IF_ERROR(ValidateFunctionDeviceConstraints(
-                *fdef, device_attrs));
+                *fdef, ctx.FuncLibDef(), device_attrs));
           }
         }
         if (ctx.JitCompileRewrite()) {
diff --git a/tensorflow/core/common_runtime/eager/validate_function_devices.cc b/tensorflow/core/common_runtime/eager/validate_function_devices.cc
index b1bfaafa272..7b103bfddd5 100644
--- a/tensorflow/core/common_runtime/eager/validate_function_devices.cc
+++ b/tensorflow/core/common_runtime/eager/validate_function_devices.cc
@@ -18,17 +18,99 @@ limitations under the License.
 #include <string>
 #include <vector>
 
+#include "absl/container/flat_hash_set.h"
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/function.h"
 #include "tensorflow/core/framework/function.pb.h"
 #include "tensorflow/core/framework/node_def.pb.h"
 #include "tensorflow/core/util/device_name_utils.h"
 
 namespace tensorflow {
 
+namespace {
+
+absl::Status ValidateFunctionDeviceConstraintsImpl(
+    const FunctionDef& fdef, const FunctionLibraryDefinition* flib_def,
+    const std::vector<DeviceAttributes>& available_devices,
+    const std::vector<DeviceNameUtils::ParsedName>& parsed_available_devices,
+    absl::flat_hash_set<std::string>& validated_devices,
+    absl::flat_hash_set<std::string>& visited_functions) {
+  for (const NodeDef& node : fdef.node_def()) {
+    const std::string& device = node.device();
+    if (!device.empty() && !validated_devices.contains(device)) {
+      DeviceNameUtils::ParsedName parsed_device;
+      if (!DeviceNameUtils::ParseFullOrLocalName(device, &parsed_device)) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Malformed device specification '", device, "' for operation ",
+            node.name(), " (", node.op(), ")."));
+      }
+
+      if (DeviceNameUtils::HasSomeDetails(parsed_device)) {
+        bool satisfied = false;
+        for (const auto& avail_parsed : parsed_available_devices) {
+          if (DeviceNameUtils::IsSpecification(parsed_device, avail_parsed)) {
+            satisfied = true;
+            break;
+          }
+        }
+
+        if (!satisfied) {
+          std::vector<std::string> available_device_names;
+          available_device_names.reserve(available_devices.size());
+          for (const auto& dev : available_devices) {
+            available_device_names.push_back(dev.name());
+          }
+          return absl::InvalidArgumentError(absl::StrCat(
+              "Could not satisfy device specification '", device,
+              "' for operation ", node.name(), " (", node.op(),
+              "). Available devices [",
+              absl::StrJoin(available_device_names, ", "), "]."));
+        }
+      }
+      validated_devices.insert(device);
+    }
+
+    // Recurse into nested functions referenced through PartitionedCall /
+    // StatefulPartitionedCall nodes so their device constraints are validated
+    // too. Regular ops are also looked up defensively: op-to-function mappings
+    // can exist in the library (e.g. via function optimization passes), and
+    // Find() simply returns nullptr when the op is a plain op.
+    if (flib_def != nullptr) {
+      std::vector<std::string> inner_funcs;
+      inner_funcs.push_back(node.op());
+      if (node.op() == "PartitionedCall" ||
+          node.op() == "StatefulPartitionedCall") {
+        auto it = node.attr().find(FunctionLibraryDefinition::kFuncAttr);
+        if (it != node.attr().end() && it->second.has_func()) {
+          inner_funcs.push_back(it->second.func().name());
+        }
+      }
+
+      for (const std::string& func_name : inner_funcs) {
+        const FunctionDef* inner_fdef = flib_def->Find(func_name);
+        if (inner_fdef != nullptr &&
+            visited_functions.insert(func_name).second) {
+          absl::Status inner_status = ValidateFunctionDeviceConstraintsImpl(
+              *inner_fdef, flib_def, available_devices,
+              parsed_available_devices, validated_devices, visited_functions);
+          if (!inner_status.ok()) {
+            return inner_status;
+          }
+        }
+      }
+    }
+  }
+
+  return absl::OkStatus();
+}
+
+}  // namespace
+
 absl::Status ValidateFunctionDeviceConstraints(
-    const FunctionDef& fdef,
+    const FunctionDef& fdef, const FunctionLibraryDefinition* flib_def,
     const std::vector<DeviceAttributes>& available_devices) {
   // Pre-parse available devices once upfront.
   std::vector<DeviceNameUtils::ParsedName> parsed_available_devices;
@@ -40,51 +122,13 @@ absl::Status ValidateFunctionDeviceConstraints(
     }
   }
 
-  for (const NodeDef& node : fdef.node_def()) {
-    const std::string& device = node.device();
-    if (device.empty()) {
-      continue;
-    }
+  absl::flat_hash_set<std::string> validated_devices;
+  absl::flat_hash_set<std::string> visited_functions;
+  visited_functions.insert(fdef.signature().name());
 
-    DeviceNameUtils::ParsedName parsed_device;
-    if (!DeviceNameUtils::ParseFullOrLocalName(device, &parsed_device)) {
-      return absl::InvalidArgumentError(absl::StrCat(
-          "Malformed device specification '", device, "' for operation ",
-          node.name(), " (", node.op(), ")."));
-    }
-
-    if (!DeviceNameUtils::HasSomeDetails(parsed_device)) {
-      continue;
-    }
-
-    // The constraint is satisfied if at least one available device matches
-    // every component specified in the constraint. IsSpecification checks
-    // all specified attributes without requiring the constraint to be fully
-    // specified, so partial constraints (e.g. '/job:localhost') and local
-    // names (e.g. 'CPU:0') are handled uniformly.
-    bool satisfied = false;
-    for (const auto& avail_parsed : parsed_available_devices) {
-      if (DeviceNameUtils::IsSpecification(parsed_device, avail_parsed)) {
-        satisfied = true;
-        break;
-      }
-    }
-
-    if (!satisfied) {
-      std::vector<std::string> available_device_names;
-      available_device_names.reserve(available_devices.size());
-      for (const auto& dev : available_devices) {
-        available_device_names.push_back(dev.name());
-      }
-      return absl::InvalidArgumentError(absl::StrCat(
-          "Could not satisfy device specification '", device,
-          "' for operation ", node.name(), " (", node.op(),
-          "). Available devices [",
-          absl::StrJoin(available_device_names, ", "), "]."));
-    }
-  }
-
-  return absl::OkStatus();
+  return ValidateFunctionDeviceConstraintsImpl(
+      fdef, flib_def, available_devices, parsed_available_devices,
+      validated_devices, visited_functions);
 }
 
 }  // namespace tensorflow
diff --git a/tensorflow/core/common_runtime/eager/validate_function_devices.h b/tensorflow/core/common_runtime/eager/validate_function_devices.h
index 54253ed49d5..2edaae51078 100644
--- a/tensorflow/core/common_runtime/eager/validate_function_devices.h
+++ b/tensorflow/core/common_runtime/eager/validate_function_devices.h
@@ -19,20 +19,26 @@ limitations under the License.
 #include <vector>
 
 #include "absl/status/status.h"
-#include "tensorflow/core/framework/function.pb.h"
 #include "tensorflow/core/framework/device_attributes.pb.h"
+#include "tensorflow/core/framework/function.pb.h"
 
 namespace tensorflow {
 
+class FunctionLibraryDefinition;
+
 // Validates that device constraints specified on nodes within a FunctionDef
-// can be satisfied by the available devices. This is used to ensure that
-// tf.device() constraints are validated even for XLA-compiled functions
-// (jit_compile=True), where the Placer is bypassed.
+// (and, recursively, within any functions it calls via PartitionedCall /
+// StatefulPartitionedCall) can be satisfied by the available devices. This is
+// used to ensure that tf.device() constraints are validated even for
+// XLA-compiled functions (jit_compile=True), where the Placer is bypassed.
+//
+// `flib_def` may be nullptr, in which case only the nodes of `fdef` itself are
+// validated.
 //
 // Returns OK if all device constraints are satisfiable, or an error status
 // describing the first unsatisfiable device constraint.
 absl::Status ValidateFunctionDeviceConstraints(
-    const FunctionDef& fdef,
+    const FunctionDef& fdef, const FunctionLibraryDefinition* flib_def,
     const std::vector<DeviceAttributes>& available_devices);
 
 }  // namespace tensorflow
diff --git a/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc b/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc
index 3a0a4c18a72..54dfc50b7a6 100644
--- a/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc
+++ b/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc
@@ -20,9 +20,12 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include "absl/status/status.h"
+#include "tensorflow/core/framework/attr_value.pb.h"
 #include "tensorflow/core/framework/device_attributes.pb.h"
+#include "tensorflow/core/framework/function.h"
 #include "tensorflow/core/framework/function.pb.h"
 #include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/platform/status_matchers.h"
 #include "tensorflow/core/platform/test.h"
 
@@ -49,24 +52,54 @@ std::vector<DeviceAttributes> LocalCpuDevices() {
   return {attrs};
 }
 
+// A FunctionDef containing a single PartitionedCall node that invokes the
+// given nested function name through its 'f' attribute.
+FunctionDef MakeFunctionWithFuncNode(const std::string& node_name,
+                                     const std::string& func_name) {
+  FunctionDef fdef;
+  NodeDef* node = fdef.add_node_def();
+  node->set_name(node_name);
+  node->set_op("PartitionedCall");
+  AttrValue func_attr;
+  func_attr.mutable_func()->set_name(func_name);
+  (*node->mutable_attr())[FunctionLibraryDefinition::kFuncAttr] = func_attr;
+  return fdef;
+}
+
+// A FunctionLibraryDefinition holding one inner function whose single node
+// carries the given device constraint.
+FunctionLibraryDefinition MakeFlibDefWithInnerFunction(
+    const OpRegistryInterface* registry, const std::string& inner_name,
+    const std::string& device) {
+  FunctionDefLibrary lib_def;
+  FunctionDef* inner = lib_def.add_function();
+  inner->mutable_signature()->set_name(inner_name);
+  NodeDef* node = inner->add_node_def();
+  node->set_name("add");
+  node->set_op("AddV2");
+  node->set_device(device);
+  FunctionLibraryDefinition flib_def(registry, lib_def);
+  return flib_def;
+}
+
 }  // namespace
 
 TEST(ValidateFunctionDevicesTest, EmptyFunctionIsOk) {
   FunctionDef fdef;
-  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices()),
+  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices()),
               ::tensorflow::testing::IsOk());
 }
 
 TEST(ValidateFunctionDevicesTest, NoDeviceConstraintIsOk) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "");
-  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices()),
+  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices()),
               ::tensorflow::testing::IsOk());
 }
 
 TEST(ValidateFunctionDevicesTest, FullySpecifiedDeviceIsOk) {
   FunctionDef fdef = MakeFunctionWithDevice(
       "add", "AddV2", "/job:localhost/replica:0/task:0/device:CPU:0");
-  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices()),
+  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices()),
               ::tensorflow::testing::IsOk());
 }
 
@@ -74,26 +107,26 @@ TEST(ValidateFunctionDevicesTest, PartialJobConstraintIsOk) {
   // '/job:localhost' is a partial constraint without a device type; it must
   // be satisfied by the fully-specified available CPU device.
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/job:localhost");
-  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices()),
+  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices()),
               ::tensorflow::testing::IsOk());
 }
 
 TEST(ValidateFunctionDevicesTest, PartialTaskConstraintIsOk) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/task:0");
-  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices()),
+  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices()),
               ::tensorflow::testing::IsOk());
 }
 
 TEST(ValidateFunctionDevicesTest, LocalDeviceNameIsOk) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "CPU:0");
-  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices()),
+  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices()),
               ::tensorflow::testing::IsOk());
 }
 
 TEST(ValidateFunctionDevicesTest, UnsatisfiedJobFails) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/job:worker");
   absl::Status status =
-      ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
+      ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices());
   EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
                           absl::StatusCode::kInvalidArgument,
                           ::testing::AllOf(
@@ -109,7 +142,7 @@ TEST(ValidateFunctionDevicesTest, InvalidDeviceIdFails) {
   // still be rejected.
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/device:CPU:99");
   absl::Status status =
-      ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
+      ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices());
   EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
                           absl::StatusCode::kInvalidArgument,
                           ::testing::HasSubstr(
@@ -120,7 +153,7 @@ TEST(ValidateFunctionDevicesTest, InvalidDeviceTypeFails) {
   FunctionDef fdef =
       MakeFunctionWithDevice("add", "AddV2", "/device:NONEXISTENT:0");
   absl::Status status =
-      ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
+      ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices());
   EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
                           absl::StatusCode::kInvalidArgument,
                           ::testing::HasSubstr(
@@ -130,7 +163,7 @@ TEST(ValidateFunctionDevicesTest, InvalidDeviceTypeFails) {
 TEST(ValidateFunctionDevicesTest, LocalDeviceNameInvalidIdFails) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "CPU:99");
   absl::Status status =
-      ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
+      ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices());
   EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
                           absl::StatusCode::kInvalidArgument,
                           ::testing::HasSubstr(
@@ -140,7 +173,7 @@ TEST(ValidateFunctionDevicesTest, LocalDeviceNameInvalidIdFails) {
 TEST(ValidateFunctionDevicesTest, MalformedDeviceSpecFails) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "not_a_device");
   absl::Status status =
-      ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
+      ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices());
   EXPECT_THAT(
       status, ::tensorflow::testing::StatusIs(
                   absl::StatusCode::kInvalidArgument,
@@ -152,7 +185,7 @@ TEST(ValidateFunctionDevicesTest, MalformedDeviceSpecFails) {
 TEST(ValidateFunctionDevicesTest, ErrorMentionsOperation) {
   FunctionDef fdef = MakeFunctionWithDevice("my_add", "AddV2", "/job:worker");
   absl::Status status =
-      ValidateFunctionDeviceConstraints(fdef, LocalCpuDevices());
+      ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices());
   EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
                           absl::StatusCode::kInvalidArgument,
                           ::testing::AllOf(::testing::HasSubstr("my_add"),
@@ -161,11 +194,69 @@ TEST(ValidateFunctionDevicesTest, ErrorMentionsOperation) {
 
 TEST(ValidateFunctionDevicesTest, EmptyAvailableDevicesRejectsConstraint) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/device:CPU:0");
-  absl::Status status = ValidateFunctionDeviceConstraints(fdef, {});
+  absl::Status status = ValidateFunctionDeviceConstraints(
+      fdef, /*flib_def=*/nullptr, {});
   EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
                           absl::StatusCode::kInvalidArgument,
                           ::testing::HasSubstr(
                               "Could not satisfy device specification")));
 }
 
+TEST(ValidateFunctionDevicesTest, NullFunctionLibraryOnlyChecksTopLevel) {
+  // With a nullptr library, nested functions cannot be inspected: only the
+  // top-level nodes are validated, and the validation passes.
+  FunctionDef fdef = MakeFunctionWithFuncNode("call", "inner_func");
+  EXPECT_THAT(
+      ValidateFunctionDeviceConstraints(
+          fdef, /*flib_def=*/nullptr, LocalCpuDevices()),
+      ::tensorflow::testing::IsOk());
+}
+
+TEST(ValidateFunctionDevicesTest, NestedFunctionInvalidDeviceFails) {
+  // The top-level function has no constraints, but the nested function called
+  // through PartitionedCall places its node on a nonexistent device.
+  FunctionDef fdef = MakeFunctionWithFuncNode("call", "inner_func");
+  FunctionLibraryDefinition flib_def = MakeFlibDefWithInnerFunction(
+      OpRegistry::Global(), "inner_func", "/device:NONEXISTENT:0");
+  absl::Status status = ValidateFunctionDeviceConstraints(
+      fdef, &flib_def, LocalCpuDevices());
+  EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
+                          absl::StatusCode::kInvalidArgument,
+                          ::testing::HasSubstr(
+                              "Could not satisfy device specification")));
+}
+
+TEST(ValidateFunctionDevicesTest, NestedFunctionValidDeviceIsOk) {
+  FunctionDef fdef = MakeFunctionWithFuncNode("call", "inner_func");
+  FunctionLibraryDefinition flib_def = MakeFlibDefWithInnerFunction(
+      OpRegistry::Global(), "inner_func", "/job:localhost");
+  EXPECT_THAT(
+      ValidateFunctionDeviceConstraints(fdef, &flib_def, LocalCpuDevices()),
+      ::tensorflow::testing::IsOk());
+}
+
+TEST(ValidateFunctionDevicesTest, NestedFunctionCycleTerminates) {
+  // Two functions calling each other in a cycle must not hang or crash; the
+  // visited-functions guard ends the traversal.
+  FunctionLibraryDefinition flib_def = MakeFlibDefWithInnerFunction(
+      OpRegistry::Global(), "inner_func", "/job:localhost");
+  FunctionDef fdef = MakeFunctionWithFuncNode("call", "inner_func");
+  // Make the inner function call itself, creating a cycle.
+  const FunctionDef* inner = flib_def.Find("inner_func");
+  ASSERT_NE(inner, nullptr);
+  FunctionDef cyclic = *inner;
+  NodeDef* call_node = cyclic.add_node_def();
+  call_node->set_name("self_call");
+  call_node->set_op("StatefulPartitionedCall");
+  AttrValue func_attr;
+  func_attr.mutable_func()->set_name("inner_func");
+  (*call_node->mutable_attr())[FunctionLibraryDefinition::kFuncAttr] =
+      func_attr;
+  // Validate the cyclic function directly (the library copy stays valid so
+  // Find() keeps returning the original definition).
+  EXPECT_THAT(
+      ValidateFunctionDeviceConstraints(cyclic, &flib_def, LocalCpuDevices()),
+      ::tensorflow::testing::IsOk());
+}
+
 }  // namespace tensorflow
diff --git a/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py b/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
index ac56b0424c6..755c2ce6d3e 100644
--- a/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
+++ b/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
@@ -1256,13 +1256,34 @@ class FunctionTest(xla_test.XLATestCase):
       x, y = constant_op.constant([2, 3]), constant_op.constant([2, 3])
       self._compareTwoMethodsCompilerIROutput(f, [x, y], {})
 
+
+class DevicePlacementValidationTest(test.TestCase):
+  """Tests that tf.device() constraints are enforced under jit_compile=True.
+
+  Standard eager execution soft-places operations onto available devices by
+  default, so unsatisfiable device constraints are only reported with soft
+  device placement disabled. The jit_compile=True validation (see #124880)
+  runs on the same eager execution path and therefore follows the same
+  setting, so these tests run with soft placement strictly disabled.
+  """
+
+  def setUp(self):
+    super().setUp()
+    self._old_soft_placement = context.context().soft_device_placement
+    context.context().soft_device_placement = False
+
+  def tearDown(self):
+    context.context().soft_device_placement = self._old_soft_placement
+    super().tearDown()
+
   def testDevicePlacementValidationWithJitCompile(self):
     """Test that jit_compile=True validates tf.device() constraints.
 
     When using tf.device() with a device that cannot exist (e.g., a
     nonexistent device type), eager execution correctly fails with
-    "Could not satisfy device specification". This test ensures jit_compile=True
-    also validates device constraints instead of silently ignoring them.
+    "Could not satisfy device specification". This test ensures
+    jit_compile=True also validates device constraints instead of silently
+    ignoring them.
 
     See https://github.com/tensorflow/tensorflow/issues/124880
     """
@@ -1276,16 +1297,10 @@ class FunctionTest(xla_test.XLATestCase):
 
     x = constant_op.constant([1.0, 2.0])
 
-    # Eager mode soft-places onto CPU by default, so the constraint is only
-    # enforced with soft device placement disabled (see #124880).
-    old_soft_placement = context.context().soft_device_placement
-    try:
-      context.context().soft_device_placement = False
-      with self.assertRaisesRegex(errors.InvalidArgumentError,
-                                  'Could not satisfy device specification'):
-        compute(x).numpy()
-    finally:
-      context.context().soft_device_placement = old_soft_placement
+    # Eager execution should fail (soft placement is disabled in setUp).
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'Could not satisfy device specification'):
+      compute(x).numpy()
 
     # jit_compile=True should also fail (this is the fix for #124880)
     with self.assertRaisesRegex(errors.InvalidArgumentError,
@@ -1305,16 +1320,10 @@ class FunctionTest(xla_test.XLATestCase):
 
     x = constant_op.constant([1.0, 2.0])
 
-    # Eager mode soft-places onto CPU by default, so the constraint is only
-    # enforced with soft device placement disabled (see #124880).
-    old_soft_placement = context.context().soft_device_placement
-    try:
-      context.context().soft_device_placement = False
-      with self.assertRaisesRegex(errors.InvalidArgumentError,
-                                  'Could not satisfy device specification'):
-        compute(x).numpy()
-    finally:
-      context.context().soft_device_placement = old_soft_placement
+    # Eager execution should fail (soft placement is disabled in setUp).
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'Could not satisfy device specification'):
+      compute(x).numpy()
 
     # jit_compile=True should also fail
     with self.assertRaisesRegex(errors.InvalidArgumentError,
@@ -1336,7 +1345,8 @@ class FunctionTest(xla_test.XLATestCase):
     # jit_compile=True must also succeed: the partial constraint
     # '/job:localhost' is satisfied by the local devices.
     self.assertAllClose(
-        polymorphic_function.function(compute, jit_compile=True)(x), [2.0, 4.0])
+        polymorphic_function.function(compute, jit_compile=True)(x),
+        [2.0, 4.0])
 
   def testDevicePlacementValidationLocalDeviceNamesWithJitCompile(self):
     """Local device names are validated under jit_compile=True."""
@@ -1356,22 +1366,40 @@ class FunctionTest(xla_test.XLATestCase):
       with ops.device('CPU:99'):
         return math_ops.add(x, x)
 
-    # A nonexistent local device ID must fail in both modes. Eager mode
-    # soft-places by default, so disable it for the eager assertion.
-    old_soft_placement = context.context().soft_device_placement
-    try:
-      context.context().soft_device_placement = False
-      with self.assertRaisesRegex(errors.InvalidArgumentError,
-                                  'Could not satisfy device specification'):
-        compute_invalid(x).numpy()
-    finally:
-      context.context().soft_device_placement = old_soft_placement
+    # A nonexistent local device ID must fail in both modes.
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'Could not satisfy device specification'):
+      compute_invalid(x).numpy()
 
     with self.assertRaisesRegex(errors.InvalidArgumentError,
                                 'Could not satisfy device specification'):
       polymorphic_function.function(
           compute_invalid, jit_compile=True)(x).numpy()
 
+  def testDevicePlacementValidationNestedFunctionWithJitCompile(self):
+    """Device constraints inside nested functions must be validated too.
+
+    The top-level function contains no tf.device() constraint; the invalid
+    constraint lives in a nested @tf.function called from the compiled
+    function, whose nodes reside in a separate FunctionDef reached through a
+    (Stateful)PartitionedCall node in the function library.
+    """
+    invalid_device = '/device:NONEXISTENT:0'
+
+    @polymorphic_function.function
+    def inner(x):
+      with ops.device(invalid_device):
+        return math_ops.add(x, x)
+
+    @polymorphic_function.function(jit_compile=True)
+    def outer(x):
+      return inner(x)
+
+    x = constant_op.constant([1.0, 2.0])
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'Could not satisfy device specification'):
+      outer(x).numpy()
+
 
 if __name__ == '__main__':
   ops.enable_eager_execution()

From c6c66f53c6278450e732c867d68305318e3594d7 Mon Sep 17 00:00:00 2001
From: Sai Ramineni <csramineni@gmail.com>
Date: Thu, 1 Oct 2026 16:40:25 -0400
Subject: [PATCH 0123/1000] Exercise public raw API in ragged scalar-value
 regression

---
 tensorflow/python/ops/ragged/BUILD                       | 2 +-
 tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/ops/ragged/BUILD b/tensorflow/python/ops/ragged/BUILD
index d48db0fb6a6..3dc2ed29649 100644
--- a/tensorflow/python/ops/ragged/BUILD
+++ b/tensorflow/python/ops/ragged/BUILD
@@ -1062,6 +1062,7 @@ py_test(
         ":ragged_factory_ops",
         ":ragged_functional_ops",
         ":ragged_tensor",
+        "//tensorflow:tensorflow_py_no_contrib",
         "//tensorflow/python/eager:context",
         "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/framework:errors",
@@ -1069,7 +1070,6 @@ py_test(
         "//tensorflow/python/ops:array_ops",
         "//tensorflow/python/ops:gradients_impl",
         "//tensorflow/python/ops:math_ops",
-        "//tensorflow/python/ops:ragged_conversion_ops_gen",
         "//tensorflow/python/platform:test",
     ],
 )
diff --git a/tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py b/tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py
index 4693494b288..f80d17fd78c 100644
--- a/tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py
@@ -14,12 +14,12 @@
 # ==============================================================================
 """Tests for ragged.to_sparse op."""
 
+from tensorflow import raw_ops
 from tensorflow.python.eager import context
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import gen_ragged_conversion_ops
 from tensorflow.python.ops import gradients_impl
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops.ragged import ragged_factory_ops
@@ -184,12 +184,12 @@ class RaggedTensorToSparseOpTest(test_util.TensorFlowTestCase):
         if not context.executing_eagerly():
           with self.assertRaisesRegex(
               ValueError, "Shape must be at least rank 1"):
-            gen_ragged_conversion_ops.ragged_tensor_to_sparse(
+            raw_ops.RaggedTensorToSparse(
                 rt_nested_splits=splits, rt_dense_values=values)
           values = array_ops.placeholder_with_default(values, shape=None)
         with self.assertRaisesRegex(
             errors.InvalidArgumentError, "rt_dense_values must have rank >= 1"):
-          self.evaluate(gen_ragged_conversion_ops.ragged_tensor_to_sparse(
+          self.evaluate(raw_ops.RaggedTensorToSparse(
               rt_nested_splits=splits, rt_dense_values=values))
 
   def testGradient(self):

From 3c24416fbf2489d2c269c5573d5992bdf860ef5e Mon Sep 17 00:00:00 2001
From: Mina Ibrahim <136262314+MinaIbrahim10@users.noreply.github.com>
Date: Thu, 1 Oct 2026 16:53:23 -0400
Subject: [PATCH 0124/1000] Move free-threading synchronization into native
 implementations

---
 tensorflow/core/data/service/BUILD            |  1 +
 tensorflow/core/data/service/server_lib.cc    | 13 ++--
 tensorflow/core/data/service/server_lib.h     | 10 +--
 tensorflow/core/grappler/clusters/BUILD       |  1 +
 tensorflow/core/grappler/clusters/cluster.h   | 10 ---
 .../core/grappler/clusters/virtual_cluster.cc |  1 +
 .../core/grappler/clusters/virtual_cluster.h  |  2 +
 .../service/server_lib_wrapper.cc             | 11 ---
 tensorflow/python/grappler/BUILD              |  1 +
 tensorflow/python/grappler/cluster_wrapper.cc | 69 +++++++++----------
 .../python/grappler/tf_optimizer_wrapper.cc   | 12 +---
 11 files changed, 51 insertions(+), 80 deletions(-)

diff --git a/tensorflow/core/data/service/BUILD b/tensorflow/core/data/service/BUILD
index 490888de9a8..4c8ac64892a 100644
--- a/tensorflow/core/data/service/BUILD
+++ b/tensorflow/core/data/service/BUILD
@@ -715,6 +715,7 @@ cc_library(
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/profiler/rpc:profiler_service_impl",
+        "//tensorflow/core/platform:mutex",
         "@com_google_absl//absl/strings",
     ] + tf_grpc_cc_dependencies(),
     alwayslink = 1,
diff --git a/tensorflow/core/data/service/server_lib.cc b/tensorflow/core/data/service/server_lib.cc
index ca546cc4d47..41aaba4c7fd 100644
--- a/tensorflow/core/data/service/server_lib.cc
+++ b/tensorflow/core/data/service/server_lib.cc
@@ -55,6 +55,7 @@ GrpcDataServerBase::GrpcDataServerBase(
       server_options_(std::move(options)) {}
 
 absl::Status GrpcDataServerBase::Start() {
+  mutex_lock lock(lifecycle_mu_);
   if (stopped_) {
     return absl::FailedPreconditionError(
         "Server cannot be started after it has been stopped.");
@@ -86,11 +87,12 @@ absl::Status GrpcDataServerBase::Start() {
 
   started_ = true;
   LOG(INFO) << "Started tf.data " << server_type_
-            << " running at 0.0.0.0:" << BoundPort();
+            << " running at 0.0.0.0:" << bound_port_;
   return absl::OkStatus();
 }
 
 void GrpcDataServerBase::Stop() {
+  mutex_lock lock(lifecycle_mu_);
   if (stopped_) {
     return;
   }
@@ -98,7 +100,7 @@ void GrpcDataServerBase::Stop() {
     StopServiceInternal();
     server_->Shutdown();
     LOG(INFO) << "Shut down " << server_type_ << " server running at port "
-              << BoundPort();
+              << bound_port_;
   }
   stopped_ = true;
 }
@@ -106,7 +108,7 @@ void GrpcDataServerBase::Stop() {
 void GrpcDataServerBase::Join() {
   ::grpc::Server* server = nullptr;
   {
-    std::lock_guard<std::mutex> lock(ExternalMutex());
+    mutex_lock lock(lifecycle_mu_);
     if (!server_) {
       return;
     }
@@ -116,7 +118,10 @@ void GrpcDataServerBase::Join() {
   server->Wait();
 }
 
-int GrpcDataServerBase::BoundPort() { return bound_port(); }
+int GrpcDataServerBase::BoundPort() {
+  mutex_lock lock(lifecycle_mu_);
+  return bound_port_;
+}
 
 void GrpcDataServerBase::AddProfilerServiceToBuilder(
     ::grpc::ServerBuilder& builder) {
diff --git a/tensorflow/core/data/service/server_lib.h b/tensorflow/core/data/service/server_lib.h
index d0bbaf7a86b..0d742b80a15 100644
--- a/tensorflow/core/data/service/server_lib.h
+++ b/tensorflow/core/data/service/server_lib.h
@@ -17,7 +17,6 @@ limitations under the License.
 #define TENSORFLOW_CORE_DATA_SERVICE_SERVER_LIB_H_
 
 #include <memory>
-#include <mutex>
 #include <string>
 #include <vector>
 
@@ -27,6 +26,7 @@ limitations under the License.
 #include "tensorflow/core/data/service/data_transfer.h"
 #include "tensorflow/core/data/service/export.pb.h"
 #include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/profiler/rpc/profiler_service_impl.h"
 #include "tensorflow/core/protobuf/service_config.pb.h"
 
@@ -62,12 +62,6 @@ class GrpcDataServerBase {
   // Returns the port bound by the server. Only valid after calling Start().
   int BoundPort();
 
-  // Serializes compound external operations on this server instance.
-  //
-  // Join() deliberately only holds this mutex while snapshotting server_, since
-  // holding it while waiting would prevent Stop() from shutting the server down.
-  std::mutex& ExternalMutex() const { return external_mu_; }
-
   // Exports the server state to improve debuggability.
   virtual ServerStateExport ExportState() const = 0;
 
@@ -86,7 +80,7 @@ class GrpcDataServerBase {
   const std::string server_type_;
 
  private:
-  mutable std::mutex external_mu_;
+  mutable mutex lifecycle_mu_;
 
   int bound_port_;
   bool started_ = false;
diff --git a/tensorflow/core/grappler/clusters/BUILD b/tensorflow/core/grappler/clusters/BUILD
index 48aada6a636..3678e7bbcec 100644
--- a/tensorflow/core/grappler/clusters/BUILD
+++ b/tensorflow/core/grappler/clusters/BUILD
@@ -107,6 +107,7 @@ cc_library(
         "//tensorflow/core/grappler/costs:analytical_cost_estimator",
         "//tensorflow/core/grappler/costs:op_level_cost_estimator",
         "//tensorflow/core/grappler/costs:virtual_scheduler",
+        "//tensorflow/core/platform:mutex",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
     ],
diff --git a/tensorflow/core/grappler/clusters/cluster.h b/tensorflow/core/grappler/clusters/cluster.h
index 21c43985dd4..9e05214e299 100644
--- a/tensorflow/core/grappler/clusters/cluster.h
+++ b/tensorflow/core/grappler/clusters/cluster.h
@@ -16,7 +16,6 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_GRAPPLER_CLUSTERS_CLUSTER_H_
 #define TENSORFLOW_CORE_GRAPPLER_CLUSTERS_CLUSTER_H_
 
-#include <mutex>
 #include <string>
 #include <unordered_map>
 #include <utility>
@@ -45,13 +44,6 @@ class Cluster {
   explicit Cluster(int timeout_s);
   virtual ~Cluster();
 
-  // Serializes compound external operations on this Cluster instance.
-  //
-  // This mutex intentionally lives in Cluster rather than in an individual
-  // Python extension so separate extension modules can coordinate access to
-  // the same native Cluster object.
-  std::mutex& ExternalMutex() const { return external_mu_; }
-
   // Returns a string that represent the type of cluster that was instantiated.
   virtual std::string type() const = 0;
 
@@ -151,8 +143,6 @@ class Cluster {
   SessionOptions options_;
   RunOptions run_options_;
 
- private:
-  mutable std::mutex external_mu_;
 };
 
 }  // end namespace grappler
diff --git a/tensorflow/core/grappler/clusters/virtual_cluster.cc b/tensorflow/core/grappler/clusters/virtual_cluster.cc
index 2eba43f4ae2..391083bfa1e 100644
--- a/tensorflow/core/grappler/clusters/virtual_cluster.cc
+++ b/tensorflow/core/grappler/clusters/virtual_cluster.cc
@@ -88,6 +88,7 @@ absl::Status VirtualCluster::Run(
 
 absl::Status VirtualCluster::Run(const GrapplerItem& item,
                                  RunMetadata* metadata) {
+  mutex_lock lock(run_mu_);
   // Initializes an analytical cost estimator to estimate the graph cost. Makes
   // sure to use static shape inference to prevent the virtual scheduler from
   // calling the Run method on the cluster and creating an infinite loop.
diff --git a/tensorflow/core/grappler/clusters/virtual_cluster.h b/tensorflow/core/grappler/clusters/virtual_cluster.h
index 5a6f3d83bcc..2b305020dfd 100644
--- a/tensorflow/core/grappler/clusters/virtual_cluster.h
+++ b/tensorflow/core/grappler/clusters/virtual_cluster.h
@@ -27,6 +27,7 @@ limitations under the License.
 #include "tensorflow/core/grappler/costs/analytical_cost_estimator.h"
 #include "tensorflow/core/grappler/costs/op_level_cost_estimator.h"
 #include "tensorflow/core/grappler/costs/virtual_scheduler.h"
+#include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/protobuf/config.pb.h"
 #include "tensorflow/core/protobuf/device_properties.pb.h"
 
@@ -60,6 +61,7 @@ class VirtualCluster : public Cluster {
   const DeviceSet* GetDeviceSet() const override { return device_set_; }
 
  private:
+  mutable mutex run_mu_;
   std::unique_ptr<AnalyticalCostEstimator> estimator_;
   const DeviceSet* device_set_ = nullptr;
 };
diff --git a/tensorflow/python/data/experimental/service/server_lib_wrapper.cc b/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
index 1f2a7f8af0a..ea90fd1787f 100644
--- a/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
+++ b/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
@@ -15,7 +15,6 @@ limitations under the License.
 #include <cstdint>
 #include <limits>
 #include <memory>
-#include <mutex>
 #include <string>
 #include <vector>
 
@@ -50,18 +49,15 @@ PYBIND11_MODULE(
   py::class_<tensorflow::data::DispatchGrpcDataServer>(
       m, "DispatchGrpcDataServer")
       .def("start", [](tensorflow::data::DispatchGrpcDataServer* server) {
-        std::lock_guard<std::mutex> lock(server->ExternalMutex());
         return server->Start();
       })
       .def("stop", [](tensorflow::data::DispatchGrpcDataServer* server) {
-        std::lock_guard<std::mutex> lock(server->ExternalMutex());
         server->Stop();
       })
       .def("join", &tensorflow::data::DispatchGrpcDataServer::Join,
            py::call_guard<py::gil_scoped_release>())
       .def("bound_port",
            [](tensorflow::data::DispatchGrpcDataServer* server) {
-             std::lock_guard<std::mutex> lock(server->ExternalMutex());
              return server->BoundPort();
            })
       .def("num_workers",
@@ -69,7 +65,6 @@ PYBIND11_MODULE(
              int num_workers;
              absl::Status status;
              {
-               std::lock_guard<std::mutex> lock(server->ExternalMutex());
                status = server->NumWorkers(&num_workers);
              }
              tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
@@ -82,7 +77,6 @@ PYBIND11_MODULE(
              std::vector<tensorflow::data::SnapshotStreamInfoWrapper> streams;
              absl::Status status;
              {
-               std::lock_guard<std::mutex> lock(server->ExternalMutex());
                status = server->SnapshotStreams(path, &streams);
              }
              tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
@@ -92,18 +86,15 @@ PYBIND11_MODULE(
   py::class_<tensorflow::data::WorkerGrpcDataServer>(
       m, "WorkerGrpcDataServer")
       .def("start", [](tensorflow::data::WorkerGrpcDataServer* server) {
-        std::lock_guard<std::mutex> lock(server->ExternalMutex());
         return server->Start();
       })
       .def("stop", [](tensorflow::data::WorkerGrpcDataServer* server) {
-        std::lock_guard<std::mutex> lock(server->ExternalMutex());
         server->Stop();
       })
       .def("join", &tensorflow::data::WorkerGrpcDataServer::Join,
            py::call_guard<py::gil_scoped_release>())
       .def("bound_port",
            [](tensorflow::data::WorkerGrpcDataServer* server) {
-             std::lock_guard<std::mutex> lock(server->ExternalMutex());
              return server->BoundPort();
            })
       .def("num_tasks",
@@ -111,7 +102,6 @@ PYBIND11_MODULE(
              int num_tasks;
              absl::Status status;
              {
-               std::lock_guard<std::mutex> lock(server->ExternalMutex());
                status = server->NumTasks(&num_tasks);
              }
              tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
@@ -124,7 +114,6 @@ PYBIND11_MODULE(
                  snapshot_task_progresses;
              absl::Status status;
              {
-               std::lock_guard<std::mutex> lock(server->ExternalMutex());
                status =
                    server->SnapshotTaskProgresses(&snapshot_task_progresses);
              }
diff --git a/tensorflow/python/grappler/BUILD b/tensorflow/python/grappler/BUILD
index 3f05b3789ff..e66b00aee38 100644
--- a/tensorflow/python/grappler/BUILD
+++ b/tensorflow/python/grappler/BUILD
@@ -293,6 +293,7 @@ tf_python_pybind_extension(
         "//tensorflow/core/framework:allocator",
         "//tensorflow/core/framework:tensor",
         "//tensorflow/core/grappler:utils",
+        "//tensorflow/core/platform:mutex",
         "//tensorflow/core/platform:status",
         "//tensorflow/python/lib/core:pybind11_status",
         "@com_google_absl//absl/base:no_destructor",
diff --git a/tensorflow/python/grappler/cluster_wrapper.cc b/tensorflow/python/grappler/cluster_wrapper.cc
index 1f011e33ae2..738789fcf0f 100644
--- a/tensorflow/python/grappler/cluster_wrapper.cc
+++ b/tensorflow/python/grappler/cluster_wrapper.cc
@@ -19,7 +19,6 @@ limitations under the License.
 #include <cstdint>
 #include <iterator>
 #include <memory>
-#include <mutex>
 #include <set>
 #include <stdexcept>
 #include <string>
@@ -48,6 +47,7 @@ limitations under the License.
 #include "tensorflow/core/grappler/devices.h"
 #include "tensorflow/core/grappler/grappler_item.h"
 #include "tensorflow/core/grappler/utils.h"
+#include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/protobuf/config.pb.h"
 #include "tensorflow/core/protobuf/device_properties.pb.h"
@@ -61,8 +61,8 @@ namespace {
 // SingleMachine uses process-global provisioning state. Serialize Python
 // creation/shutdown paths that update that state without relying on static
 // destruction order during interpreter shutdown.
-std::mutex& GetClusterLifecycleMutex() {
-  static absl::NoDestructor<std::mutex> mu;
+tensorflow::mutex& GetClusterLifecycleMutex() {
+  static absl::NoDestructor<tensorflow::mutex> mu;
   return *mu;
 }
 
@@ -77,7 +77,7 @@ absl::Status _GetOpPerformanceDataAndRunTime(
   if (!status.ok()) return status;
 
   tensorflow::RunMetadata run_metadata;
-  tsl::MaybeRaiseRegisteredFromStatusWithGIL(
+  TF_RETURN_IF_ERROR(
       cost_measure->PredictCosts(item.graph, &run_metadata, costs));
 
   if (op_performance_data) {
@@ -98,21 +98,21 @@ PYBIND11_MODULE(
            bool disable_detailed_stats) -> tensorflow::grappler::Cluster* {
           // TODO(petebu): Make these named arguments with default values
           // instead.
-          std::unique_lock<std::mutex> lifecycle_lock(
-              GetClusterLifecycleMutex());
-
-          int num_cpu_cores =
-              tensorflow::grappler::GetNumAvailableLogicalCPUCores();
-          int num_gpus = tensorflow::grappler::GetNumAvailableGPUs();
-          int timeout_s = 60 * 10;
-          std::unique_ptr<tensorflow::grappler::Cluster> cluster =
-              std::make_unique<tensorflow::grappler::SingleMachine>(
-                  timeout_s, num_cpu_cores, num_gpus);
-          cluster->DisableDetailedStats(disable_detailed_stats);
-          cluster->AllowSoftPlacement(allow_soft_placement);
-          cluster->SetNumWarmupSteps(10);
-          absl::Status provision_status = cluster->Provision();
-          lifecycle_lock.unlock();
+          std::unique_ptr<tensorflow::grappler::Cluster> cluster;
+          absl::Status provision_status;
+          {
+            tensorflow::mutex_lock lifecycle_lock(GetClusterLifecycleMutex());
+            int num_cpu_cores =
+                tensorflow::grappler::GetNumAvailableLogicalCPUCores();
+            int num_gpus = tensorflow::grappler::GetNumAvailableGPUs();
+            int timeout_s = 60 * 10;
+            cluster = std::make_unique<tensorflow::grappler::SingleMachine>(
+                timeout_s, num_cpu_cores, num_gpus);
+            cluster->DisableDetailedStats(disable_detailed_stats);
+            cluster->AllowSoftPlacement(allow_soft_placement);
+            cluster->SetNumWarmupSteps(10);
+            provision_status = cluster->Provision();
+          }
           tsl::MaybeRaiseRegisteredFromStatusWithGIL(provision_status);
           return cluster.release();
         });
@@ -145,9 +145,12 @@ PYBIND11_MODULE(
     if (cluster == nullptr) {
       return;
     }
-    std::lock_guard<std::mutex> cluster_lock(cluster->ExternalMutex());
-    std::lock_guard<std::mutex> lifecycle_lock(GetClusterLifecycleMutex());
-    (void)cluster->Shutdown();
+    if (cluster->type() == "single_machine") {
+      tensorflow::mutex_lock lifecycle_lock(GetClusterLifecycleMutex());
+      (void)cluster->Shutdown();
+    } else {
+      (void)cluster->Shutdown();
+    }
   });
 
   m.def("TF_ListDevices",
@@ -155,12 +158,10 @@ PYBIND11_MODULE(
           if (cluster == nullptr) {
             tsl::MaybeRaiseRegisteredFromStatusWithGIL(
                 absl::InvalidArgumentError("Cluster cannot be None."));
-            return {};
           }
 
           std::vector<py::bytes> named_devices;
           {
-            std::lock_guard<std::mutex> lock(cluster->ExternalMutex());
             const auto& devices = cluster->GetDevices();
 
             for (const auto& dev : devices) {
@@ -200,7 +201,6 @@ PYBIND11_MODULE(
         std::string cluster_type;
 
         {
-          std::lock_guard<std::mutex> lock(cluster->ExternalMutex());
           const auto& devices = cluster->GetDevices();
           cluster_type = cluster->type();
 
@@ -266,12 +266,11 @@ PYBIND11_MODULE(
           }
 
           std::vector<std::string> device_names;
-          for (const std::string& type : feasible) {
-            auto it = device_types.find(type);
-            DCHECK(it != device_types.end());
-            for (const std::string& name : it->second) {
-              device_names.push_back(name);
+          for (const auto& [type, names] : device_types) {
+            if (feasible.find(type) == feasible.end()) {
+              continue;
             }
+            device_names.insert(device_names.end(), names.begin(), names.end());
           }
           result[node] = device_names;
         }
@@ -297,10 +296,8 @@ PYBIND11_MODULE(
           if (cluster == nullptr || item == nullptr) {
             tsl::MaybeRaiseRegisteredFromStatusWithGIL(absl::InvalidArgumentError(
                 "You need both a cluster and an item to measure costs."));
-            return {};
           }
 
-          std::lock_guard<std::mutex> lock(cluster->ExternalMutex());
 
           const int num_measurements = cluster->type() == "virtual" ? 1 : 10;
           tensorflow::grappler::MeasuringCostEstimator cost_measure(
@@ -310,10 +307,9 @@ PYBIND11_MODULE(
           tensorflow::grappler::Costs costs;
           absl::Status s = _GetOpPerformanceDataAndRunTime(
               *item, &cost_measure, &op_performance_data, &costs);
-          double run_time = FLT_MAX;
-          if (s.ok()) {
-            run_time = static_cast<double>(costs.execution_time.count()) / 1e9;
-          }
+          tsl::MaybeRaiseRegisteredFromStatusWithGIL(s);
+          double run_time =
+              static_cast<double>(costs.execution_time.count()) / 1e9;
           tensorflow::StepStats step_stats;
           if (generate_timeline) {
             tensorflow::RunMetadata metadata;
@@ -348,7 +344,6 @@ PYBIND11_MODULE(
               "You need both a cluster and an item to determine peak "
               "memory usage.")));
         }
-        std::lock_guard<std::mutex> lock(cluster->ExternalMutex());
 
         tensorflow::grappler::GraphMemory memory(*item);
 
diff --git a/tensorflow/python/grappler/tf_optimizer_wrapper.cc b/tensorflow/python/grappler/tf_optimizer_wrapper.cc
index 6fff2fa6665..a7a8f19b814 100644
--- a/tensorflow/python/grappler/tf_optimizer_wrapper.cc
+++ b/tensorflow/python/grappler/tf_optimizer_wrapper.cc
@@ -14,7 +14,6 @@ limitations under the License.
 ==============================================================================*/
 
 #include <memory>
-#include <mutex>
 #include <stdexcept>
 #include <string>
 #include <unordered_map>
@@ -93,15 +92,8 @@ tensorflow::GraphDef TF_OptimizeGraph(
     tensorflow::DeviceBase* cpu_device = nullptr;
     tensorflow::grappler::MetaOptimizer optimizer(cpu_device, config_proto);
 
-    absl::Status optimize_status;
-    if (cluster != nullptr) {
-      std::lock_guard<std::mutex> lock(cluster->ExternalMutex());
-      optimize_status =
-          optimizer.Optimize(cluster, *grappler_item, &out_graph);
-    } else {
-      optimize_status =
-          optimizer.Optimize(nullptr, *grappler_item, &out_graph);
-    }
+    absl::Status optimize_status =
+        optimizer.Optimize(cluster, *grappler_item, &out_graph);
     tsl::MaybeRaiseRegisteredFromStatusWithGIL(optimize_status);
     if (strip_default_attributes) {
       tensorflow::StripDefaultAttributes(*tensorflow::OpRegistry::Global(),

From 4f6a0f49bc4549b9ea9db946f1300e32c59c0465 Mon Sep 17 00:00:00 2001
From: Mina Ibrahim <136262314+MinaIbrahim10@users.noreply.github.com>
Date: Thu, 1 Oct 2026 17:56:54 -0400
Subject: [PATCH 0125/1000] Address Grappler wrapper ordering and exception
 review

---
 .../experimental/service/server_lib_wrapper.cc     |  4 ++--
 tensorflow/python/grappler/BUILD                   |  1 +
 tensorflow/python/grappler/cluster_wrapper.cc      | 14 ++++++++------
 3 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/tensorflow/python/data/experimental/service/server_lib_wrapper.cc b/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
index ea90fd1787f..e4a6d4e072b 100644
--- a/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
+++ b/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
@@ -49,7 +49,7 @@ PYBIND11_MODULE(
   py::class_<tensorflow::data::DispatchGrpcDataServer>(
       m, "DispatchGrpcDataServer")
       .def("start", [](tensorflow::data::DispatchGrpcDataServer* server) {
-        return server->Start();
+        tsl::MaybeRaiseRegisteredFromStatusWithGIL(server->Start());
       })
       .def("stop", [](tensorflow::data::DispatchGrpcDataServer* server) {
         server->Stop();
@@ -86,7 +86,7 @@ PYBIND11_MODULE(
   py::class_<tensorflow::data::WorkerGrpcDataServer>(
       m, "WorkerGrpcDataServer")
       .def("start", [](tensorflow::data::WorkerGrpcDataServer* server) {
-        return server->Start();
+        tsl::MaybeRaiseRegisteredFromStatusWithGIL(server->Start());
       })
       .def("stop", [](tensorflow::data::WorkerGrpcDataServer* server) {
         server->Stop();
diff --git a/tensorflow/python/grappler/BUILD b/tensorflow/python/grappler/BUILD
index e66b00aee38..897d4e42fb3 100644
--- a/tensorflow/python/grappler/BUILD
+++ b/tensorflow/python/grappler/BUILD
@@ -293,6 +293,7 @@ tf_python_pybind_extension(
         "//tensorflow/core/framework:allocator",
         "//tensorflow/core/framework:tensor",
         "//tensorflow/core/grappler:utils",
+        "//tensorflow/core/platform:errors",
         "//tensorflow/core/platform:mutex",
         "//tensorflow/core/platform:status",
         "//tensorflow/python/lib/core:pybind11_status",
diff --git a/tensorflow/python/grappler/cluster_wrapper.cc b/tensorflow/python/grappler/cluster_wrapper.cc
index 738789fcf0f..df40940d15a 100644
--- a/tensorflow/python/grappler/cluster_wrapper.cc
+++ b/tensorflow/python/grappler/cluster_wrapper.cc
@@ -47,6 +47,7 @@ limitations under the License.
 #include "tensorflow/core/grappler/devices.h"
 #include "tensorflow/core/grappler/grappler_item.h"
 #include "tensorflow/core/grappler/utils.h"
+#include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/protobuf/config.pb.h"
@@ -66,8 +67,6 @@ tensorflow::mutex& GetClusterLifecycleMutex() {
   return *mu;
 }
 
-}  // namespace
-
 absl::Status _GetOpPerformanceDataAndRunTime(
     const tensorflow::grappler::GrapplerItem& item,
     tensorflow::grappler::CostEstimator* cost_measure,
@@ -87,6 +86,8 @@ absl::Status _GetOpPerformanceDataAndRunTime(
   return absl::OkStatus();
 }
 
+}  // namespace
+
 PYBIND11_MAKE_OPAQUE(tensorflow::grappler::Cluster);
 
 PYBIND11_MODULE(
@@ -266,11 +267,12 @@ PYBIND11_MODULE(
           }
 
           std::vector<std::string> device_names;
-          for (const auto& [type, names] : device_types) {
-            if (feasible.find(type) == feasible.end()) {
-              continue;
+          for (const std::string& type : feasible) {
+            auto it = device_types.find(type);
+            DCHECK(it != device_types.end());
+            for (const std::string& name : it->second) {
+              device_names.push_back(name);
             }
-            device_names.insert(device_names.end(), names.begin(), names.end());
           }
           result[node] = device_names;
         }

From 5639704fb1b126054dbf74487a4b25c3f340b3fb Mon Sep 17 00:00:00 2001
From: alwaysprince05 <alwaysprince05@users.noreply.github.com>
Date: Fri, 2 Oct 2026 03:27:23 +0530
Subject: [PATCH 0126/1000] Address review: follow every function-bearing node
 attribute

Nested function discovery no longer special-cases
(Stateful)PartitionedCall: any 'func' or 'list.func' attribute on a node
is now followed, which covers the branch attributes of If and While
('then_branch'/'else_branch', 'cond'/'body') and list-of-function
attributes such as 'branches' of Case ops.

Device constraints are checked against Device objects directly, reusing
their parsed names instead of copying DeviceAttributes protos on the
hot path; a DeviceAttributes overload is kept for callers that only
hold protos.

Adds a Python regression test for tf.cond and tf.while_loop inside a
jit_compile=True function, plus C++ tests for If, While and
list-of-function attributes.
---
 .../core/common_runtime/eager/execute.cc      |   5 +-
 .../eager/validate_function_devices.cc        | 101 +++++--
 .../eager/validate_function_devices.h         |  18 +-
 .../eager/validate_function_devices_test.cc   | 259 ++++++++++++++----
 .../polymorphic_function_xla_jit_test.py      |  42 +++
 5 files changed, 332 insertions(+), 93 deletions(-)

diff --git a/tensorflow/core/common_runtime/eager/execute.cc b/tensorflow/core/common_runtime/eager/execute.cc
index 0d730a4cb99..4822e59acb1 100644
--- a/tensorflow/core/common_runtime/eager/execute.cc
+++ b/tensorflow/core/common_runtime/eager/execute.cc
@@ -40,7 +40,6 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/eager/validate_function_devices.h"
 #include "tensorflow/core/common_runtime/int32_fulltype.h"
 #include "tensorflow/core/framework/cancellation.h"
-#include "tensorflow/core/framework/device_attributes.pb.h"
 #include "tensorflow/core/framework/full_type.pb.h"
 #include "tensorflow/core/framework/function.pb.h"
 #include "tensorflow/core/framework/kernel_def.pb.h"
@@ -1427,10 +1426,8 @@ absl::Status GetOrCreateKernelAndDevice(
         if (has_xla_must_compile && !ctx.AllowSoftPlacement()) {
           const FunctionDef* fdef = op->GetFunctionDef();
           if (fdef != nullptr) {
-            std::vector<DeviceAttributes> device_attrs;
-            ctx.ListDevices(&device_attrs);
             TF_RETURN_IF_ERROR(ValidateFunctionDeviceConstraints(
-                *fdef, ctx.FuncLibDef(), device_attrs));
+                *fdef, ctx.FuncLibDef(), ctx.ListAllTfDevices()));
           }
         }
         if (ctx.JitCompileRewrite()) {
diff --git a/tensorflow/core/common_runtime/eager/validate_function_devices.cc b/tensorflow/core/common_runtime/eager/validate_function_devices.cc
index 7b103bfddd5..2d6eb3b79e9 100644
--- a/tensorflow/core/common_runtime/eager/validate_function_devices.cc
+++ b/tensorflow/core/common_runtime/eager/validate_function_devices.cc
@@ -23,6 +23,8 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/device.h"
+#include "tensorflow/core/framework/device_attributes.pb.h"
 #include "tensorflow/core/framework/function.h"
 #include "tensorflow/core/framework/function.pb.h"
 #include "tensorflow/core/framework/node_def.pb.h"
@@ -32,9 +34,12 @@ namespace tensorflow {
 
 namespace {
 
+// Validates the device constraints of `fdef` and recurses into every nested
+// function reachable through function-bearing node attributes. The
+// visited-functions guard makes cyclic call graphs terminate.
 absl::Status ValidateFunctionDeviceConstraintsImpl(
     const FunctionDef& fdef, const FunctionLibraryDefinition* flib_def,
-    const std::vector<DeviceAttributes>& available_devices,
+    const std::vector<std::string>& available_device_names,
     const std::vector<DeviceNameUtils::ParsedName>& parsed_available_devices,
     absl::flat_hash_set<std::string>& validated_devices,
     absl::flat_hash_set<std::string>& visited_functions) {
@@ -58,11 +63,6 @@ absl::Status ValidateFunctionDeviceConstraintsImpl(
         }
 
         if (!satisfied) {
-          std::vector<std::string> available_device_names;
-          available_device_names.reserve(available_devices.size());
-          for (const auto& dev : available_devices) {
-            available_device_names.push_back(dev.name());
-          }
           return absl::InvalidArgumentError(absl::StrCat(
               "Could not satisfy device specification '", device,
               "' for operation ", node.name(), " (", node.op(),
@@ -73,32 +73,44 @@ absl::Status ValidateFunctionDeviceConstraintsImpl(
       validated_devices.insert(device);
     }
 
-    // Recurse into nested functions referenced through PartitionedCall /
-    // StatefulPartitionedCall nodes so their device constraints are validated
-    // too. Regular ops are also looked up defensively: op-to-function mappings
-    // can exist in the library (e.g. via function optimization passes), and
-    // Find() simply returns nullptr when the op is a plain op.
+    // Recurse into nested functions. Function-bearing attributes are not
+    // limited to the 'f' attribute of (Stateful)PartitionedCall: control flow
+    // ops carry functions in attributes like 'then_branch'/'else_branch'
+    // (If) or 'cond'/'body' (While), and list-of-function attributes exist
+    // too (e.g. 'branches'). Every 'func' or 'list.func' attribute is
+    // therefore followed, regardless of the op type. Plain ops are looked up
+    // defensively as well: op-to-function mappings can exist in the library
+    // (e.g. via function optimization passes), and Find() simply returns
+    // nullptr when the op is a plain op.
     if (flib_def != nullptr) {
-      std::vector<std::string> inner_funcs;
-      inner_funcs.push_back(node.op());
-      if (node.op() == "PartitionedCall" ||
-          node.op() == "StatefulPartitionedCall") {
-        auto it = node.attr().find(FunctionLibraryDefinition::kFuncAttr);
-        if (it != node.attr().end() && it->second.has_func()) {
-          inner_funcs.push_back(it->second.func().name());
-        }
-      }
-
-      for (const std::string& func_name : inner_funcs) {
+      auto validate_inner = [&](const std::string& func_name) -> absl::Status {
         const FunctionDef* inner_fdef = flib_def->Find(func_name);
         if (inner_fdef != nullptr &&
             visited_functions.insert(func_name).second) {
-          absl::Status inner_status = ValidateFunctionDeviceConstraintsImpl(
-              *inner_fdef, flib_def, available_devices,
+          return ValidateFunctionDeviceConstraintsImpl(
+              *inner_fdef, flib_def, available_device_names,
               parsed_available_devices, validated_devices, visited_functions);
+        }
+        return absl::OkStatus();
+      };
+
+      absl::Status inner_status = validate_inner(node.op());
+      if (!inner_status.ok()) {
+        return inner_status;
+      }
+      for (const auto& attr : node.attr()) {
+        if (attr.second.has_func()) {
+          inner_status = validate_inner(attr.second.func().name());
           if (!inner_status.ok()) {
             return inner_status;
           }
+        } else if (attr.second.has_list()) {
+          for (const auto& func : attr.second.list().func()) {
+            inner_status = validate_inner(func.name());
+            if (!inner_status.ok()) {
+              return inner_status;
+            }
+          }
         }
       }
     }
@@ -107,12 +119,49 @@ absl::Status ValidateFunctionDeviceConstraintsImpl(
   return absl::OkStatus();
 }
 
+std::vector<std::string> DeviceNames(
+    const std::vector<DeviceAttributes>& available_devices) {
+  std::vector<std::string> names;
+  names.reserve(available_devices.size());
+  for (const auto& dev : available_devices) {
+    names.push_back(dev.name());
+  }
+  return names;
+}
+
 }  // namespace
 
+absl::Status ValidateFunctionDeviceConstraints(
+    const FunctionDef& fdef, const FunctionLibraryDefinition* flib_def,
+    const std::vector<Device*>& available_devices) {
+  std::vector<std::string> available_device_names;
+  available_device_names.reserve(available_devices.size());
+  std::vector<DeviceNameUtils::ParsedName> parsed_available_devices;
+  parsed_available_devices.reserve(available_devices.size());
+  for (const Device* dev : available_devices) {
+    if (dev == nullptr) {
+      continue;
+    }
+    available_device_names.push_back(dev->name());
+    // Device parses its name into a ParsedName at construction time, so this
+    // reuses the already-parsed form instead of re-parsing every name.
+    parsed_available_devices.push_back(dev->parsed_name());
+  }
+
+  absl::flat_hash_set<std::string> validated_devices;
+  absl::flat_hash_set<std::string> visited_functions;
+  visited_functions.insert(fdef.signature().name());
+
+  return ValidateFunctionDeviceConstraintsImpl(
+      fdef, flib_def, available_device_names, parsed_available_devices,
+      validated_devices, visited_functions);
+}
+
 absl::Status ValidateFunctionDeviceConstraints(
     const FunctionDef& fdef, const FunctionLibraryDefinition* flib_def,
     const std::vector<DeviceAttributes>& available_devices) {
-  // Pre-parse available devices once upfront.
+  std::vector<std::string> available_device_names =
+      DeviceNames(available_devices);
   std::vector<DeviceNameUtils::ParsedName> parsed_available_devices;
   parsed_available_devices.reserve(available_devices.size());
   for (const auto& dev : available_devices) {
@@ -127,7 +176,7 @@ absl::Status ValidateFunctionDeviceConstraints(
   visited_functions.insert(fdef.signature().name());
 
   return ValidateFunctionDeviceConstraintsImpl(
-      fdef, flib_def, available_devices, parsed_available_devices,
+      fdef, flib_def, available_device_names, parsed_available_devices,
       validated_devices, visited_functions);
 }
 
diff --git a/tensorflow/core/common_runtime/eager/validate_function_devices.h b/tensorflow/core/common_runtime/eager/validate_function_devices.h
index 2edaae51078..53f664f5459 100644
--- a/tensorflow/core/common_runtime/eager/validate_function_devices.h
+++ b/tensorflow/core/common_runtime/eager/validate_function_devices.h
@@ -24,19 +24,29 @@ limitations under the License.
 
 namespace tensorflow {
 
+class Device;
 class FunctionLibraryDefinition;
 
 // Validates that device constraints specified on nodes within a FunctionDef
-// (and, recursively, within any functions it calls via PartitionedCall /
-// StatefulPartitionedCall) can be satisfied by the available devices. This is
-// used to ensure that tf.device() constraints are validated even for
-// XLA-compiled functions (jit_compile=True), where the Placer is bypassed.
+// (and, recursively, within any functions it calls through function-bearing
+// node attributes, e.g. the 'f' attribute of (Stateful)PartitionedCall or the
+// branch attributes of If and While) can be satisfied by the available
+// devices. This is used to ensure that tf.device() constraints are validated
+// even for XLA-compiled functions (jit_compile=True), where the Placer is
+// bypassed.
 //
 // `flib_def` may be nullptr, in which case only the nodes of `fdef` itself are
 // validated.
 //
 // Returns OK if all device constraints are satisfiable, or an error status
 // describing the first unsatisfiable device constraint.
+absl::Status ValidateFunctionDeviceConstraints(
+    const FunctionDef& fdef, const FunctionLibraryDefinition* flib_def,
+    const std::vector<Device*>& available_devices);
+
+// Same as above for callers that hold DeviceAttributes protos instead of
+// Device objects (e.g. unit tests). Prefer the overload above, which avoids
+// copying the device attributes.
 absl::Status ValidateFunctionDeviceConstraints(
     const FunctionDef& fdef, const FunctionLibraryDefinition* flib_def,
     const std::vector<DeviceAttributes>& available_devices);
diff --git a/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc b/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc
index 54dfc50b7a6..8779d33bb2f 100644
--- a/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc
+++ b/tensorflow/core/common_runtime/eager/validate_function_devices_test.cc
@@ -15,12 +15,14 @@ limitations under the License.
 
 #include "tensorflow/core/common_runtime/eager/validate_function_devices.h"
 
+#include <memory>
 #include <string>
 #include <vector>
 
 #include <gmock/gmock.h>
 #include "absl/status/status.h"
 #include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/device.h"
 #include "tensorflow/core/framework/device_attributes.pb.h"
 #include "tensorflow/core/framework/function.h"
 #include "tensorflow/core/framework/function.pb.h"
@@ -33,6 +35,28 @@ namespace tensorflow {
 
 namespace {
 
+// A fake device that has specific device attributes, used to simulate the
+// presence of a device without depending on the real runtime.
+class FakeDevice : public Device {
+ private:
+  explicit FakeDevice(const DeviceAttributes& device_attributes)
+      : Device(nullptr, device_attributes) {}
+
+ public:
+  absl::Status Sync() override {
+    return absl::UnimplementedError("FakeDevice::Sync()");
+  }
+
+  Allocator* GetAllocator(AllocatorAttributes attr) override { return nullptr; }
+
+  static std::unique_ptr<Device> Make(const std::string& name) {
+    DeviceAttributes device_attributes;
+    device_attributes.set_name(name);
+    device_attributes.set_device_type("FakeCPU");
+    return std::unique_ptr<Device>(new FakeDevice(device_attributes));
+  }
+};
+
 // A FunctionDef with a single node carrying the given device constraint.
 FunctionDef MakeFunctionWithDevice(const std::string& node_name,
                                    const std::string& op,
@@ -45,24 +69,36 @@ FunctionDef MakeFunctionWithDevice(const std::string& node_name,
   return fdef;
 }
 
-std::vector<DeviceAttributes> LocalCpuDevices() {
-  DeviceAttributes attrs;
-  attrs.set_name("/job:localhost/replica:0/task:0/device:CPU:0");
-  attrs.set_device_type("CPU");
-  return {attrs};
-}
-
-// A FunctionDef containing a single PartitionedCall node that invokes the
-// given nested function name through its 'f' attribute.
-FunctionDef MakeFunctionWithFuncNode(const std::string& node_name,
+// A FunctionDef with a single node of type `op` carrying one function-bearing
+// attribute named `attr_name` pointing at `func_name`.
+FunctionDef MakeFunctionWithFuncAttr(const std::string& node_name,
+                                     const std::string& op,
+                                     const std::string& attr_name,
                                      const std::string& func_name) {
   FunctionDef fdef;
   NodeDef* node = fdef.add_node_def();
   node->set_name(node_name);
-  node->set_op("PartitionedCall");
+  node->set_op(op);
   AttrValue func_attr;
   func_attr.mutable_func()->set_name(func_name);
-  (*node->mutable_attr())[FunctionLibraryDefinition::kFuncAttr] = func_attr;
+  (*node->mutable_attr())[attr_name] = func_attr;
+  return fdef;
+}
+
+// A FunctionDef with a single node carrying a list-of-functions attribute.
+FunctionDef MakeFunctionWithFuncListAttr(
+    const std::string& node_name, const std::string& op,
+    const std::string& attr_name, const std::vector<std::string>& names) {
+  FunctionDef fdef;
+  NodeDef* node = fdef.add_node_def();
+  node->set_name(node_name);
+  node->set_op(op);
+  AttrValue list_attr;
+  for (const std::string& name : names) {
+    NameAttrList* func = list_attr.mutable_list()->add_func();
+    func->set_name(name);
+  }
+  (*node->mutable_attr())[attr_name] = list_attr;
   return fdef;
 }
 
@@ -82,51 +118,71 @@ FunctionLibraryDefinition MakeFlibDefWithInnerFunction(
   return flib_def;
 }
 
+class ValidateFunctionDevicesTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    devices_.push_back(
+        FakeDevice::Make("/job:localhost/replica:0/task:0/device:CPU:0"));
+    available_devices_.push_back(devices_.back().get());
+  }
+
+  // The unique_ptr owns the devices; available_devices_ holds raw pointers
+  // for the duration of the test.
+  std::vector<std::unique_ptr<Device>> devices_;
+  std::vector<Device*> available_devices_;
+};
+
 }  // namespace
 
-TEST(ValidateFunctionDevicesTest, EmptyFunctionIsOk) {
+TEST_F(ValidateFunctionDevicesTest, EmptyFunctionIsOk) {
   FunctionDef fdef;
-  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices()),
-              ::tensorflow::testing::IsOk());
+  EXPECT_THAT(
+      ValidateFunctionDeviceConstraints(fdef, nullptr, available_devices_),
+      ::tensorflow::testing::IsOk());
 }
 
-TEST(ValidateFunctionDevicesTest, NoDeviceConstraintIsOk) {
+TEST_F(ValidateFunctionDevicesTest, NoDeviceConstraintIsOk) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "");
-  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices()),
-              ::tensorflow::testing::IsOk());
+  EXPECT_THAT(
+      ValidateFunctionDeviceConstraints(fdef, nullptr, available_devices_),
+      ::tensorflow::testing::IsOk());
 }
 
-TEST(ValidateFunctionDevicesTest, FullySpecifiedDeviceIsOk) {
+TEST_F(ValidateFunctionDevicesTest, FullySpecifiedDeviceIsOk) {
   FunctionDef fdef = MakeFunctionWithDevice(
       "add", "AddV2", "/job:localhost/replica:0/task:0/device:CPU:0");
-  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices()),
-              ::tensorflow::testing::IsOk());
+  EXPECT_THAT(
+      ValidateFunctionDeviceConstraints(fdef, nullptr, available_devices_),
+      ::tensorflow::testing::IsOk());
 }
 
-TEST(ValidateFunctionDevicesTest, PartialJobConstraintIsOk) {
+TEST_F(ValidateFunctionDevicesTest, PartialJobConstraintIsOk) {
   // '/job:localhost' is a partial constraint without a device type; it must
   // be satisfied by the fully-specified available CPU device.
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/job:localhost");
-  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices()),
-              ::tensorflow::testing::IsOk());
+  EXPECT_THAT(
+      ValidateFunctionDeviceConstraints(fdef, nullptr, available_devices_),
+      ::tensorflow::testing::IsOk());
 }
 
-TEST(ValidateFunctionDevicesTest, PartialTaskConstraintIsOk) {
+TEST_F(ValidateFunctionDevicesTest, PartialTaskConstraintIsOk) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/task:0");
-  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices()),
-              ::tensorflow::testing::IsOk());
+  EXPECT_THAT(
+      ValidateFunctionDeviceConstraints(fdef, nullptr, available_devices_),
+      ::tensorflow::testing::IsOk());
 }
 
-TEST(ValidateFunctionDevicesTest, LocalDeviceNameIsOk) {
+TEST_F(ValidateFunctionDevicesTest, LocalDeviceNameIsOk) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "CPU:0");
-  EXPECT_THAT(ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices()),
-              ::tensorflow::testing::IsOk());
+  EXPECT_THAT(
+      ValidateFunctionDeviceConstraints(fdef, nullptr, available_devices_),
+      ::tensorflow::testing::IsOk());
 }
 
-TEST(ValidateFunctionDevicesTest, UnsatisfiedJobFails) {
+TEST_F(ValidateFunctionDevicesTest, UnsatisfiedJobFails) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/job:worker");
   absl::Status status =
-      ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices());
+      ValidateFunctionDeviceConstraints(fdef, nullptr, available_devices_);
   EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
                           absl::StatusCode::kInvalidArgument,
                           ::testing::AllOf(
@@ -137,43 +193,43 @@ TEST(ValidateFunctionDevicesTest, UnsatisfiedJobFails) {
                                                    "task:0/device:CPU:0"))));
 }
 
-TEST(ValidateFunctionDevicesTest, InvalidDeviceIdFails) {
+TEST_F(ValidateFunctionDevicesTest, InvalidDeviceIdFails) {
   // The device type exists, but the id does not: the partial constraint must
   // still be rejected.
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/device:CPU:99");
   absl::Status status =
-      ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices());
+      ValidateFunctionDeviceConstraints(fdef, nullptr, available_devices_);
   EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
                           absl::StatusCode::kInvalidArgument,
                           ::testing::HasSubstr(
                               "Could not satisfy device specification")));
 }
 
-TEST(ValidateFunctionDevicesTest, InvalidDeviceTypeFails) {
+TEST_F(ValidateFunctionDevicesTest, InvalidDeviceTypeFails) {
   FunctionDef fdef =
       MakeFunctionWithDevice("add", "AddV2", "/device:NONEXISTENT:0");
   absl::Status status =
-      ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices());
+      ValidateFunctionDeviceConstraints(fdef, nullptr, available_devices_);
   EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
                           absl::StatusCode::kInvalidArgument,
                           ::testing::HasSubstr(
                               "Could not satisfy device specification")));
 }
 
-TEST(ValidateFunctionDevicesTest, LocalDeviceNameInvalidIdFails) {
+TEST_F(ValidateFunctionDevicesTest, LocalDeviceNameInvalidIdFails) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "CPU:99");
   absl::Status status =
-      ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices());
+      ValidateFunctionDeviceConstraints(fdef, nullptr, available_devices_);
   EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
                           absl::StatusCode::kInvalidArgument,
                           ::testing::HasSubstr(
                               "Could not satisfy device specification")));
 }
 
-TEST(ValidateFunctionDevicesTest, MalformedDeviceSpecFails) {
+TEST_F(ValidateFunctionDevicesTest, MalformedDeviceSpecFails) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "not_a_device");
   absl::Status status =
-      ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices());
+      ValidateFunctionDeviceConstraints(fdef, nullptr, available_devices_);
   EXPECT_THAT(
       status, ::tensorflow::testing::StatusIs(
                   absl::StatusCode::kInvalidArgument,
@@ -182,65 +238,70 @@ TEST(ValidateFunctionDevicesTest, MalformedDeviceSpecFails) {
                       ::testing::HasSubstr("not_a_device"))));
 }
 
-TEST(ValidateFunctionDevicesTest, ErrorMentionsOperation) {
+TEST_F(ValidateFunctionDevicesTest, ErrorMentionsOperation) {
   FunctionDef fdef = MakeFunctionWithDevice("my_add", "AddV2", "/job:worker");
   absl::Status status =
-      ValidateFunctionDeviceConstraints(fdef, nullptr, LocalCpuDevices());
+      ValidateFunctionDeviceConstraints(fdef, nullptr, available_devices_);
   EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
                           absl::StatusCode::kInvalidArgument,
                           ::testing::AllOf(::testing::HasSubstr("my_add"),
                                            ::testing::HasSubstr("AddV2"))));
 }
 
-TEST(ValidateFunctionDevicesTest, EmptyAvailableDevicesRejectsConstraint) {
+TEST_F(ValidateFunctionDevicesTest, EmptyAvailableDevicesRejectsConstraint) {
   FunctionDef fdef = MakeFunctionWithDevice("add", "AddV2", "/device:CPU:0");
-  absl::Status status = ValidateFunctionDeviceConstraints(
-      fdef, /*flib_def=*/nullptr, {});
+  std::vector<Device*> no_devices;
+  absl::Status status =
+      ValidateFunctionDeviceConstraints(fdef, /*flib_def=*/nullptr, no_devices);
   EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
                           absl::StatusCode::kInvalidArgument,
                           ::testing::HasSubstr(
                               "Could not satisfy device specification")));
 }
 
-TEST(ValidateFunctionDevicesTest, NullFunctionLibraryOnlyChecksTopLevel) {
+TEST_F(ValidateFunctionDevicesTest, NullFunctionLibraryOnlyChecksTopLevel) {
   // With a nullptr library, nested functions cannot be inspected: only the
   // top-level nodes are validated, and the validation passes.
-  FunctionDef fdef = MakeFunctionWithFuncNode("call", "inner_func");
+  FunctionDef fdef =
+      MakeFunctionWithFuncAttr("call", "PartitionedCall", "f", "inner_func");
   EXPECT_THAT(
-      ValidateFunctionDeviceConstraints(
-          fdef, /*flib_def=*/nullptr, LocalCpuDevices()),
+      ValidateFunctionDeviceConstraints(fdef, /*flib_def=*/nullptr,
+                                        available_devices_),
       ::tensorflow::testing::IsOk());
 }
 
-TEST(ValidateFunctionDevicesTest, NestedFunctionInvalidDeviceFails) {
+TEST_F(ValidateFunctionDevicesTest, NestedFunctionInvalidDeviceFails) {
   // The top-level function has no constraints, but the nested function called
   // through PartitionedCall places its node on a nonexistent device.
-  FunctionDef fdef = MakeFunctionWithFuncNode("call", "inner_func");
+  FunctionDef fdef =
+      MakeFunctionWithFuncAttr("call", "PartitionedCall", "f", "inner_func");
   FunctionLibraryDefinition flib_def = MakeFlibDefWithInnerFunction(
       OpRegistry::Global(), "inner_func", "/device:NONEXISTENT:0");
-  absl::Status status = ValidateFunctionDeviceConstraints(
-      fdef, &flib_def, LocalCpuDevices());
+  absl::Status status =
+      ValidateFunctionDeviceConstraints(fdef, &flib_def, available_devices_);
   EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
                           absl::StatusCode::kInvalidArgument,
                           ::testing::HasSubstr(
                               "Could not satisfy device specification")));
 }
 
-TEST(ValidateFunctionDevicesTest, NestedFunctionValidDeviceIsOk) {
-  FunctionDef fdef = MakeFunctionWithFuncNode("call", "inner_func");
+TEST_F(ValidateFunctionDevicesTest, NestedFunctionValidDeviceIsOk) {
+  FunctionDef fdef =
+      MakeFunctionWithFuncAttr("call", "PartitionedCall", "f", "inner_func");
   FunctionLibraryDefinition flib_def = MakeFlibDefWithInnerFunction(
       OpRegistry::Global(), "inner_func", "/job:localhost");
   EXPECT_THAT(
-      ValidateFunctionDeviceConstraints(fdef, &flib_def, LocalCpuDevices()),
+      ValidateFunctionDeviceConstraints(fdef, &flib_def, available_devices_),
       ::tensorflow::testing::IsOk());
 }
 
-TEST(ValidateFunctionDevicesTest, NestedFunctionCycleTerminates) {
+TEST_F(ValidateFunctionDevicesTest, NestedFunctionCycleTerminates) {
   // Two functions calling each other in a cycle must not hang or crash; the
   // visited-functions guard ends the traversal.
   FunctionLibraryDefinition flib_def = MakeFlibDefWithInnerFunction(
       OpRegistry::Global(), "inner_func", "/job:localhost");
-  FunctionDef fdef = MakeFunctionWithFuncNode("call", "inner_func");
+  FunctionDef fdef = MakeFunctionWithFuncAttr("call", "StatefulPartitionedCall",
+                                              "f", "inner_func");
   // Make the inner function call itself, creating a cycle.
   const FunctionDef* inner = flib_def.Find("inner_func");
   ASSERT_NE(inner, nullptr);
@@ -255,8 +316,88 @@ TEST(ValidateFunctionDevicesTest, NestedFunctionCycleTerminates) {
   // Validate the cyclic function directly (the library copy stays valid so
   // Find() keeps returning the original definition).
   EXPECT_THAT(
-      ValidateFunctionDeviceConstraints(cyclic, &flib_def, LocalCpuDevices()),
+      ValidateFunctionDeviceConstraints(cyclic, &flib_def, available_devices_),
       ::tensorflow::testing::IsOk());
 }
 
+TEST_F(ValidateFunctionDevicesTest, NestedIfBranchInvalidDeviceFails) {
+  // A tf.cond lowers to an If op whose branches live in the 'then_branch'
+  // and 'else_branch' attributes; functions invoked there must be validated.
+  FunctionDef fdef = MakeFunctionWithFuncAttr("cond", "If", "then_branch",
+                                              "then_func");
+  FunctionLibraryDefinition flib_def = MakeFlibDefWithInnerFunction(
+      OpRegistry::Global(), "then_func", "/device:NONEXISTENT:0");
+  absl::Status status =
+      ValidateFunctionDeviceConstraints(fdef, &flib_def, available_devices_);
+  EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
+                          absl::StatusCode::kInvalidArgument,
+                          ::testing::HasSubstr(
+                              "Could not satisfy device specification")));
+}
+
+TEST_F(ValidateFunctionDevicesTest, NestedWhileBodyInvalidDeviceFails) {
+  // A tf.while_loop lowers to a While op with 'cond' and 'body' attributes;
+  // the body function's constraints must be validated too.
+  FunctionDef fdef =
+      MakeFunctionWithFuncAttr("loop", "While", "body", "body_func");
+  FunctionLibraryDefinition flib_def = MakeFlibDefWithInnerFunction(
+      OpRegistry::Global(), "body_func", "/device:NONEXISTENT:0");
+  absl::Status status =
+      ValidateFunctionDeviceConstraints(fdef, &flib_def, available_devices_);
+  EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
+                          absl::StatusCode::kInvalidArgument,
+                          ::testing::HasSubstr(
+                              "Could not satisfy device specification")));
+}
+
+TEST_F(ValidateFunctionDevicesTest, NestedIfBranchValidDeviceIsOk) {
+  FunctionDef fdef = MakeFunctionWithFuncAttr("cond", "If", "else_branch",
+                                              "else_func");
+  FunctionLibraryDefinition flib_def = MakeFlibDefWithInnerFunction(
+      OpRegistry::Global(), "else_func", "/job:localhost");
+  EXPECT_THAT(
+      ValidateFunctionDeviceConstraints(fdef, &flib_def, available_devices_),
+      ::tensorflow::testing::IsOk());
+}
+
+TEST_F(ValidateFunctionDevicesTest, NestedFuncListInvalidDeviceFails) {
+  // Case-style ops carry their branches in a list(attr) attribute; every
+  // function in the list must be validated. Names that match no function in
+  // the library are ignored.
+  FunctionDef fdef = MakeFunctionWithFuncListAttr(
+      "switch_case", "StatelessCase", "branches",
+      {"first_func", "second_func"});
+  FunctionDefLibrary lib_def;
+  FunctionDef* second = lib_def.add_function();
+  second->mutable_signature()->set_name("second_func");
+  NodeDef* node = second->add_node_def();
+  node->set_name("add");
+  node->set_op("AddV2");
+  node->set_device("/device:NONEXISTENT:0");
+  FunctionLibraryDefinition flib_def(OpRegistry::Global(), lib_def);
+  absl::Status status =
+      ValidateFunctionDeviceConstraints(fdef, &flib_def, available_devices_);
+  EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
+                          absl::StatusCode::kInvalidArgument,
+                          ::testing::HasSubstr(
+                              "Could not satisfy device specification")));
+}
+
+TEST_F(ValidateFunctionDevicesTest, DeviceAttributesOverloadMatchesDevice) {
+  // The DeviceAttributes overload exists for callers holding protos; it must
+  // behave the same as the Device-based one.
+  FunctionDef fdef =
+      MakeFunctionWithDevice("add", "AddV2", "/device:NONEXISTENT:0");
+  DeviceAttributes attrs;
+  attrs.set_name("/job:localhost/replica:0/task:0/device:CPU:0");
+  attrs.set_device_type("CPU");
+  std::vector<DeviceAttributes> available_devices = {attrs};
+  absl::Status status =
+      ValidateFunctionDeviceConstraints(fdef, nullptr, available_devices);
+  EXPECT_THAT(status, ::tensorflow::testing::StatusIs(
+                          absl::StatusCode::kInvalidArgument,
+                          ::testing::HasSubstr(
+                              "Could not satisfy device specification")));
+}
+
 }  // namespace tensorflow
diff --git a/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py b/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
index 755c2ce6d3e..5659e63cbbf 100644
--- a/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
+++ b/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
@@ -1400,6 +1400,48 @@ class DevicePlacementValidationTest(test.TestCase):
                                 'Could not satisfy device specification'):
       outer(x).numpy()
 
+  def testDevicePlacementValidationNestedControlFlowWithJitCompile(self):
+    """Device constraints inside control flow must be validated too.
+
+    tf.cond and tf.while_loop lower to If and While ops whose nested
+    functions live in attributes like 'then_branch'/'else_branch' and
+    'cond'/'body' rather than the 'f' attribute of a PartitionedCall, so an
+    invalid tf.device() constraint inside their bodies must still be
+    reported instead of silently ignored.
+    """
+    invalid_device = '/device:NONEXISTENT:0'
+    x = constant_op.constant([1.0, 2.0])
+
+    @polymorphic_function.function(jit_compile=True)
+    def with_cond(x):
+
+      def true_branch():
+        with ops.device(invalid_device):
+          return math_ops.add(x, x)
+
+      return cond.cond(
+          constant_op.constant(True), true_branch,
+          lambda: array_ops.identity(x))
+
+    @polymorphic_function.function(jit_compile=True)
+    def with_while(x):
+
+      def loop_body(i, v):
+        with ops.device(invalid_device):
+          return i + 1, math_ops.add(v, v)
+
+      return while_loop.while_loop(
+          lambda i, _: i < 1, loop_body,
+          (constant_op.constant(0), x))[1]
+
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'Could not satisfy device specification'):
+      with_cond(x).numpy()
+
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'Could not satisfy device specification'):
+      with_while(x).numpy()
+
 
 if __name__ == '__main__':
   ops.enable_eager_execution()

From 1072654803bcfcfa20684a881489ff721dba3257 Mon Sep 17 00:00:00 2001
From: Mina Ibrahim <136262314+MinaIbrahim10@users.noreply.github.com>
Date: Thu, 1 Oct 2026 18:29:03 -0400
Subject: [PATCH 0127/1000] Fix deterministic device ordering and release GIL
 for server queries

---
 .../python/data/experimental/service/server_lib_wrapper.cc    | 4 ++++
 tensorflow/python/grappler/cluster_wrapper.cc                 | 4 ++--
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/data/experimental/service/server_lib_wrapper.cc b/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
index e4a6d4e072b..30f4500456d 100644
--- a/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
+++ b/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
@@ -65,6 +65,7 @@ PYBIND11_MODULE(
              int num_workers;
              absl::Status status;
              {
+               py::gil_scoped_release release;
                status = server->NumWorkers(&num_workers);
              }
              tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
@@ -77,6 +78,7 @@ PYBIND11_MODULE(
              std::vector<tensorflow::data::SnapshotStreamInfoWrapper> streams;
              absl::Status status;
              {
+               py::gil_scoped_release release;
                status = server->SnapshotStreams(path, &streams);
              }
              tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
@@ -102,6 +104,7 @@ PYBIND11_MODULE(
              int num_tasks;
              absl::Status status;
              {
+               py::gil_scoped_release release;
                status = server->NumTasks(&num_tasks);
              }
              tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
@@ -114,6 +117,7 @@ PYBIND11_MODULE(
                  snapshot_task_progresses;
              absl::Status status;
              {
+               py::gil_scoped_release release;
                status =
                    server->SnapshotTaskProgresses(&snapshot_task_progresses);
              }
diff --git a/tensorflow/python/grappler/cluster_wrapper.cc b/tensorflow/python/grappler/cluster_wrapper.cc
index df40940d15a..3d810679834 100644
--- a/tensorflow/python/grappler/cluster_wrapper.cc
+++ b/tensorflow/python/grappler/cluster_wrapper.cc
@@ -198,7 +198,7 @@ PYBIND11_MODULE(
               absl::InternalError("You need both a cluster and an "
                                   "item to get supported devices.")));
         }
-        std::unordered_map<std::string, std::vector<std::string>> device_types;
+        std::unordered_map<std::string, std::set<std::string>> device_types;
         std::string cluster_type;
 
         {
@@ -206,7 +206,7 @@ PYBIND11_MODULE(
           cluster_type = cluster->type();
 
           for (const auto& dev : devices) {
-            device_types[dev.second.type()].push_back(dev.first);
+            device_types[dev.second.type()].insert(dev.first);
           }
         }
 

From 1a70852a6b4d133b46ca54065e628f7b79780547 Mon Sep 17 00:00:00 2001
From: alwaysprince05 <alwaysprince05@users.noreply.github.com>
Date: Fri, 2 Oct 2026 04:02:11 +0530
Subject: [PATCH 0128/1000] Address review: cache negative lookups, defer error
 formatting

Nested function names are now inserted into the visited set before the
library lookup, so primitive ops that resolve to no FunctionDef are
looked up at most once instead of once per node instance.

The list of available device names is only rendered on the error path
via ParsedNameToString, removing the unconditional string copies from
every successful validation, and the tracking sets use string_view so
cache lookups no longer allocate.

DevicePlacementValidationTest now registers its soft-placement restore
via addCleanup so the setting is restored even if setUp fails.

Also declares @xla//xla/tsl/util:device_name_utils as a direct dep of
:validate_function_devices for include hygiene; the
tensorflow/core/util/device_name_utils.h path is only a using-decl shim
over that header, so the xla label is the target that actually provides
the symbols.
---
 tensorflow/core/common_runtime/eager/BUILD    |  1 +
 .../eager/validate_function_devices.cc        | 63 +++++++++----------
 .../polymorphic_function_xla_jit_test.py      | 11 ++--
 3 files changed, 37 insertions(+), 38 deletions(-)

diff --git a/tensorflow/core/common_runtime/eager/BUILD b/tensorflow/core/common_runtime/eager/BUILD
index 201df308507..1f5b0b4b3ed 100644
--- a/tensorflow/core/common_runtime/eager/BUILD
+++ b/tensorflow/core/common_runtime/eager/BUILD
@@ -717,6 +717,7 @@ cc_library(
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
+        "@xla//xla/tsl/util:device_name_utils",
     ],
 )
 
diff --git a/tensorflow/core/common_runtime/eager/validate_function_devices.cc b/tensorflow/core/common_runtime/eager/validate_function_devices.cc
index 2d6eb3b79e9..020b061640d 100644
--- a/tensorflow/core/common_runtime/eager/validate_function_devices.cc
+++ b/tensorflow/core/common_runtime/eager/validate_function_devices.cc
@@ -16,12 +16,14 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/eager/validate_function_devices.h"
 
 #include <string>
+#include <string_view>
 #include <vector>
 
 #include "absl/container/flat_hash_set.h"
 #include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
+#include "absl/strings/string_view.h"
 #include "tensorflow/core/framework/attr_value.pb.h"
 #include "tensorflow/core/framework/device.h"
 #include "tensorflow/core/framework/device_attributes.pb.h"
@@ -39,10 +41,9 @@ namespace {
 // visited-functions guard makes cyclic call graphs terminate.
 absl::Status ValidateFunctionDeviceConstraintsImpl(
     const FunctionDef& fdef, const FunctionLibraryDefinition* flib_def,
-    const std::vector<std::string>& available_device_names,
     const std::vector<DeviceNameUtils::ParsedName>& parsed_available_devices,
-    absl::flat_hash_set<std::string>& validated_devices,
-    absl::flat_hash_set<std::string>& visited_functions) {
+    absl::flat_hash_set<absl::string_view>& validated_devices,
+    absl::flat_hash_set<absl::string_view>& visited_functions) {
   for (const NodeDef& node : fdef.node_def()) {
     const std::string& device = node.device();
     if (!device.empty() && !validated_devices.contains(device)) {
@@ -63,6 +64,14 @@ absl::Status ValidateFunctionDeviceConstraintsImpl(
         }
 
         if (!satisfied) {
+          // The available-device list is only needed for the error message,
+          // so it is rendered here instead of on every validation.
+          std::vector<std::string> available_device_names;
+          available_device_names.reserve(parsed_available_devices.size());
+          for (const auto& avail_parsed : parsed_available_devices) {
+            available_device_names.push_back(
+                DeviceNameUtils::ParsedNameToString(avail_parsed));
+          }
           return absl::InvalidArgumentError(absl::StrCat(
               "Could not satisfy device specification '", device,
               "' for operation ", node.name(), " (", node.op(),
@@ -80,16 +89,19 @@ absl::Status ValidateFunctionDeviceConstraintsImpl(
     // too (e.g. 'branches'). Every 'func' or 'list.func' attribute is
     // therefore followed, regardless of the op type. Plain ops are looked up
     // defensively as well: op-to-function mappings can exist in the library
-    // (e.g. via function optimization passes), and Find() simply returns
-    // nullptr when the op is a plain op.
+    // (e.g. via function optimization passes). Names are cached in
+    // visited_functions before the lookup, including negative results for
+    // plain ops, so every distinct name is resolved at most once.
     if (flib_def != nullptr) {
-      auto validate_inner = [&](const std::string& func_name) -> absl::Status {
+      auto validate_inner = [&](absl::string_view func_name) -> absl::Status {
+        if (!visited_functions.insert(func_name).second) {
+          return absl::OkStatus();
+        }
         const FunctionDef* inner_fdef = flib_def->Find(func_name);
-        if (inner_fdef != nullptr &&
-            visited_functions.insert(func_name).second) {
+        if (inner_fdef != nullptr) {
           return ValidateFunctionDeviceConstraintsImpl(
-              *inner_fdef, flib_def, available_device_names,
-              parsed_available_devices, validated_devices, visited_functions);
+              *inner_fdef, flib_def, parsed_available_devices,
+              validated_devices, visited_functions);
         }
         return absl::OkStatus();
       };
@@ -119,49 +131,34 @@ absl::Status ValidateFunctionDeviceConstraintsImpl(
   return absl::OkStatus();
 }
 
-std::vector<std::string> DeviceNames(
-    const std::vector<DeviceAttributes>& available_devices) {
-  std::vector<std::string> names;
-  names.reserve(available_devices.size());
-  for (const auto& dev : available_devices) {
-    names.push_back(dev.name());
-  }
-  return names;
-}
-
 }  // namespace
 
 absl::Status ValidateFunctionDeviceConstraints(
     const FunctionDef& fdef, const FunctionLibraryDefinition* flib_def,
     const std::vector<Device*>& available_devices) {
-  std::vector<std::string> available_device_names;
-  available_device_names.reserve(available_devices.size());
   std::vector<DeviceNameUtils::ParsedName> parsed_available_devices;
   parsed_available_devices.reserve(available_devices.size());
   for (const Device* dev : available_devices) {
     if (dev == nullptr) {
       continue;
     }
-    available_device_names.push_back(dev->name());
     // Device parses its name into a ParsedName at construction time, so this
     // reuses the already-parsed form instead of re-parsing every name.
     parsed_available_devices.push_back(dev->parsed_name());
   }
 
-  absl::flat_hash_set<std::string> validated_devices;
-  absl::flat_hash_set<std::string> visited_functions;
+  absl::flat_hash_set<absl::string_view> validated_devices;
+  absl::flat_hash_set<absl::string_view> visited_functions;
   visited_functions.insert(fdef.signature().name());
 
   return ValidateFunctionDeviceConstraintsImpl(
-      fdef, flib_def, available_device_names, parsed_available_devices,
-      validated_devices, visited_functions);
+      fdef, flib_def, parsed_available_devices, validated_devices,
+      visited_functions);
 }
 
 absl::Status ValidateFunctionDeviceConstraints(
     const FunctionDef& fdef, const FunctionLibraryDefinition* flib_def,
     const std::vector<DeviceAttributes>& available_devices) {
-  std::vector<std::string> available_device_names =
-      DeviceNames(available_devices);
   std::vector<DeviceNameUtils::ParsedName> parsed_available_devices;
   parsed_available_devices.reserve(available_devices.size());
   for (const auto& dev : available_devices) {
@@ -171,13 +168,13 @@ absl::Status ValidateFunctionDeviceConstraints(
     }
   }
 
-  absl::flat_hash_set<std::string> validated_devices;
-  absl::flat_hash_set<std::string> visited_functions;
+  absl::flat_hash_set<absl::string_view> validated_devices;
+  absl::flat_hash_set<absl::string_view> visited_functions;
   visited_functions.insert(fdef.signature().name());
 
   return ValidateFunctionDeviceConstraintsImpl(
-      fdef, flib_def, available_device_names, parsed_available_devices,
-      validated_devices, visited_functions);
+      fdef, flib_def, parsed_available_devices, validated_devices,
+      visited_functions);
 }
 
 }  // namespace tensorflow
diff --git a/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py b/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
index 5659e63cbbf..07bbc135ebd 100644
--- a/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
+++ b/tensorflow/python/eager/polymorphic_function/polymorphic_function_xla_jit_test.py
@@ -1269,12 +1269,13 @@ class DevicePlacementValidationTest(test.TestCase):
 
   def setUp(self):
     super().setUp()
-    self._old_soft_placement = context.context().soft_device_placement
-    context.context().soft_device_placement = False
+    old_soft_placement = context.context().soft_device_placement
 
-  def tearDown(self):
-    context.context().soft_device_placement = self._old_soft_placement
-    super().tearDown()
+    def restore():
+      context.context().soft_device_placement = old_soft_placement
+
+    self.addCleanup(restore)
+    context.context().soft_device_placement = False
 
   def testDevicePlacementValidationWithJitCompile(self):
     """Test that jit_compile=True validates tf.device() constraints.

From c2bad7504140de7beda387721476e003b839c2d1 Mon Sep 17 00:00:00 2001
From: alwaysprince05 <alwaysprince05@users.noreply.github.com>
Date: Fri, 2 Oct 2026 04:25:39 +0530
Subject: [PATCH 0129/1000] Fix string_view to string conversion in nested
 function lookup

FunctionLibraryDefinition::Find takes a const std::string&, and
absl::string_view does not implicitly convert to it, so passing the
string_view lambda parameter to Find failed to compile. The lambda now
binds const std::string& instead; its call sites pass lvalue
std::string members of the node attributes, so binding stays copy-free
and the tracking sets keep using string_view for zero-copy inserts.
---
 .../core/common_runtime/eager/validate_function_devices.cc     | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensorflow/core/common_runtime/eager/validate_function_devices.cc b/tensorflow/core/common_runtime/eager/validate_function_devices.cc
index 020b061640d..df677bbcf0c 100644
--- a/tensorflow/core/common_runtime/eager/validate_function_devices.cc
+++ b/tensorflow/core/common_runtime/eager/validate_function_devices.cc
@@ -16,7 +16,6 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/eager/validate_function_devices.h"
 
 #include <string>
-#include <string_view>
 #include <vector>
 
 #include "absl/container/flat_hash_set.h"
@@ -93,7 +92,7 @@ absl::Status ValidateFunctionDeviceConstraintsImpl(
     // visited_functions before the lookup, including negative results for
     // plain ops, so every distinct name is resolved at most once.
     if (flib_def != nullptr) {
-      auto validate_inner = [&](absl::string_view func_name) -> absl::Status {
+      auto validate_inner = [&](const std::string& func_name) -> absl::Status {
         if (!visited_functions.insert(func_name).second) {
           return absl::OkStatus();
         }

From 4d05888edd8691aaa64e29587e3e001672fc589f Mon Sep 17 00:00:00 2001
From: Mina Ibrahim <136262314+MinaIbrahim10@users.noreply.github.com>
Date: Fri, 2 Oct 2026 02:17:02 +0300
Subject: [PATCH 0130/1000] Fix status propagation tests and remaining
 concurrency review

---
 .../service/distributed_save_ft_test.py       | 16 ++++-----
 .../service/server_lib_wrapper.cc             | 16 +++++++--
 tensorflow/python/grappler/cluster_wrapper.cc | 35 +++++++++++--------
 3 files changed, 43 insertions(+), 24 deletions(-)

diff --git a/tensorflow/python/data/experimental/kernel_tests/service/distributed_save_ft_test.py b/tensorflow/python/data/experimental/kernel_tests/service/distributed_save_ft_test.py
index 2b91ceaaa90..aa17c644a58 100644
--- a/tensorflow/python/data/experimental/kernel_tests/service/distributed_save_ft_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/service/distributed_save_ft_test.py
@@ -134,7 +134,7 @@ class SnapshotFtTest(data_service_test_base.TestBase, parameterized.TestCase):
         cluster.dispatcher_address()))
 
     self._make_stream_dir(snapshot_dir.full_path, bad_stream_dir_name)
-    with self.assertRaisesRegex(RuntimeError, "Can't parse"):
+    with self.assertRaisesRegex(errors.InternalError, "Can't parse"):
       cluster.restart_dispatcher()
 
   @combinations.generate(
@@ -152,7 +152,7 @@ class SnapshotFtTest(data_service_test_base.TestBase, parameterized.TestCase):
 
     os.makedirs(os.path.join(self._splits_dir(snapshot_dir.full_path),
                              bad_source_dir_name))
-    with self.assertRaisesRegex(RuntimeError, "Can't parse"):
+    with self.assertRaisesRegex(errors.InternalError, "Can't parse"):
       cluster.restart_dispatcher()
 
   @combinations.generate(test_base.default_test_combinations())
@@ -166,7 +166,7 @@ class SnapshotFtTest(data_service_test_base.TestBase, parameterized.TestCase):
 
     os.makedirs(os.path.join(self._splits_dir(snapshot_dir.full_path),
                              "source_1"))
-    with self.assertRaisesRegex(RuntimeError, "Found conflict"):
+    with self.assertRaisesRegex(errors.InternalError, "Found conflict"):
       cluster.restart_dispatcher()
 
   @combinations.generate(
@@ -190,7 +190,7 @@ class SnapshotFtTest(data_service_test_base.TestBase, parameterized.TestCase):
     write_file(os.path.join(self._source_dir(snapshot_dir.full_path),
                             bad_split_filename))
     with self.assertRaisesRegex(
-        ValueError,
+        errors.InvalidArgumentError,
         "Expected split_<local_split_index>_<global_split_index>"):
       cluster.restart_dispatcher()
 
@@ -206,7 +206,7 @@ class SnapshotFtTest(data_service_test_base.TestBase, parameterized.TestCase):
     write_file(os.path.join(self._source_dir(snapshot_dir.full_path),
                             "split_1_0"))
     with self.assertRaisesRegex(
-        ValueError,
+        errors.InvalidArgumentError,
         "The local split index 1 exceeds the global split index 0"):
       cluster.restart_dispatcher()
 
@@ -221,7 +221,7 @@ class SnapshotFtTest(data_service_test_base.TestBase, parameterized.TestCase):
 
     write_file(os.path.join(self._source_dir(snapshot_dir.full_path),
                             "split_0_1"))
-    with self.assertRaisesRegex(RuntimeError, "Found missing global"):
+    with self.assertRaisesRegex(errors.InternalError, "Found missing global"):
       cluster.restart_dispatcher()
 
   @combinations.generate(test_base.default_test_combinations())
@@ -237,7 +237,7 @@ class SnapshotFtTest(data_service_test_base.TestBase, parameterized.TestCase):
         snapshot_dir.full_path, stream_idx=0), "split_0_1"))
     write_file(os.path.join(self._source_dir(
         snapshot_dir.full_path, stream_idx=1, worker=1), "split_0_1"))
-    with self.assertRaisesRegex(RuntimeError, "Found duplicate global"):
+    with self.assertRaisesRegex(errors.InternalError, "Found duplicate global"):
       cluster.restart_dispatcher()
 
   @combinations.generate(test_base.default_test_combinations())
@@ -253,7 +253,7 @@ class SnapshotFtTest(data_service_test_base.TestBase, parameterized.TestCase):
         self._source_dir(snapshot_dir.full_path, stream_idx=0), "split_0_1"))
     write_file(os.path.join(
         self._source_dir(snapshot_dir.full_path, stream_idx=1), "split_0_1"))
-    with self.assertRaisesRegex(RuntimeError, "worker is already assigned"):
+    with self.assertRaisesRegex(errors.InternalError, "worker is already assigned"):
       cluster.restart_dispatcher()
 
   @combinations.generate(test_base.default_test_combinations())
diff --git a/tensorflow/python/data/experimental/service/server_lib_wrapper.cc b/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
index 30f4500456d..2676c614320 100644
--- a/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
+++ b/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
@@ -49,9 +49,15 @@ PYBIND11_MODULE(
   py::class_<tensorflow::data::DispatchGrpcDataServer>(
       m, "DispatchGrpcDataServer")
       .def("start", [](tensorflow::data::DispatchGrpcDataServer* server) {
-        tsl::MaybeRaiseRegisteredFromStatusWithGIL(server->Start());
+        absl::Status status;
+        {
+          py::gil_scoped_release release;
+          status = server->Start();
+        }
+        tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
       })
       .def("stop", [](tensorflow::data::DispatchGrpcDataServer* server) {
+        py::gil_scoped_release release;
         server->Stop();
       })
       .def("join", &tensorflow::data::DispatchGrpcDataServer::Join,
@@ -88,9 +94,15 @@ PYBIND11_MODULE(
   py::class_<tensorflow::data::WorkerGrpcDataServer>(
       m, "WorkerGrpcDataServer")
       .def("start", [](tensorflow::data::WorkerGrpcDataServer* server) {
-        tsl::MaybeRaiseRegisteredFromStatusWithGIL(server->Start());
+        absl::Status status;
+        {
+          py::gil_scoped_release release;
+          status = server->Start();
+        }
+        tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
       })
       .def("stop", [](tensorflow::data::WorkerGrpcDataServer* server) {
+        py::gil_scoped_release release;
         server->Stop();
       })
       .def("join", &tensorflow::data::WorkerGrpcDataServer::Join,
diff --git a/tensorflow/python/grappler/cluster_wrapper.cc b/tensorflow/python/grappler/cluster_wrapper.cc
index 3d810679834..937cc17f84d 100644
--- a/tensorflow/python/grappler/cluster_wrapper.cc
+++ b/tensorflow/python/grappler/cluster_wrapper.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include <cstddef>
 #include <cstdint>
 #include <iterator>
+#include <map>
 #include <memory>
 #include <set>
 #include <stdexcept>
@@ -192,7 +193,7 @@ PYBIND11_MODULE(
       "TF_GetSupportedDevices",
       [](tensorflow::grappler::Cluster* cluster,
          tensorflow::grappler::GrapplerItem* item)
-          -> std::unordered_map<std::string, std::vector<std::string>> {
+          -> std::map<std::string, std::vector<std::string>> {
         if (cluster == nullptr || item == nullptr) {
           tsl::MaybeRaiseRegisteredFromStatusWithGIL(absl::Status(
               absl::InternalError("You need both a cluster and an "
@@ -210,27 +211,33 @@ PYBIND11_MODULE(
           }
         }
 
-        std::unordered_map<std::string, std::set<std::string>>
-            supported_device_types;
+        std::map<std::string, std::set<std::string>> supported_device_types;
         std::unordered_map<std::string, std::set<std::string>>
             device_restrictions;
 
-        for (const auto& node : item->graph.node()) {
-          for (const auto& dev : device_types) {
-            const std::string& type = dev.first;
-            if (cluster_type != "single_machine") {
+        if (cluster_type != "single_machine") {
+          for (const auto& node : item->graph.node()) {
+            auto& supported = supported_device_types[node.name()];
+            for (const auto& dev : device_types) {
               // The actual kernel may not be linked in this binary.
-              supported_device_types[node.name()].insert(type);
-            } else {
-              // Check the kernel capabilities
+              supported.insert(dev.first);
+            }
+          }
+        } else {
+          for (const auto& node : item->graph.node()) {
+            auto& supported = supported_device_types[node.name()];
+            for (const auto& dev : device_types) {
+              const std::string& type = dev.first;
+
+              // Check the kernel capabilities.
               const tensorflow::DeviceType dev_type(type);
               absl::Status s =
                   tensorflow::FindKernelDef(dev_type, node, nullptr, nullptr);
               if (s.ok()) {
-                supported_device_types[node.name()].insert(type);
+                supported.insert(type);
 
                 // Check which inputs are restricted to reside on the host.
-                // TODO: extends this to support outputs as well
+                // TODO: extends this to support outputs as well.
                 tensorflow::MemoryTypeVector inp_mtypes;
                 tensorflow::MemoryTypeVector out_mtypes;
                 absl::Status s = tensorflow::MemoryTypesForNode(
@@ -251,7 +258,7 @@ PYBIND11_MODULE(
           }
         }
 
-        std::unordered_map<std::string, std::vector<std::string>> result;
+        std::map<std::string, std::vector<std::string>> result;
         for (const auto& supported_dev : supported_device_types) {
           const std::string& node = supported_dev.first;
           std::set<std::string> feasible;
@@ -341,7 +348,7 @@ PYBIND11_MODULE(
          tensorflow::grappler::Cluster* cluster)
           -> std::unordered_map<std::string,
                                 std::tuple<int64_t, std::vector<MemoryUsage>>> {
-        if (item == nullptr || cluster == nullptr) {
+        if (cluster == nullptr || item == nullptr) {
           tsl::MaybeRaiseRegisteredFromStatusWithGIL(absl::Status(absl::InternalError(
               "You need both a cluster and an item to determine peak "
               "memory usage.")));

From 159655018060cb93f46dbcc4fc9353dc4991e894 Mon Sep 17 00:00:00 2001
From: Mina Ibrahim <136262314+MinaIbrahim10@users.noreply.github.com>
Date: Fri, 2 Oct 2026 02:20:17 +0300
Subject: [PATCH 0131/1000] Fix PyLint line length in snapshot recovery test

---
 .../kernel_tests/service/distributed_save_ft_test.py           | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/data/experimental/kernel_tests/service/distributed_save_ft_test.py b/tensorflow/python/data/experimental/kernel_tests/service/distributed_save_ft_test.py
index aa17c644a58..ce987b350c1 100644
--- a/tensorflow/python/data/experimental/kernel_tests/service/distributed_save_ft_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/service/distributed_save_ft_test.py
@@ -253,7 +253,8 @@ class SnapshotFtTest(data_service_test_base.TestBase, parameterized.TestCase):
         self._source_dir(snapshot_dir.full_path, stream_idx=0), "split_0_1"))
     write_file(os.path.join(
         self._source_dir(snapshot_dir.full_path, stream_idx=1), "split_0_1"))
-    with self.assertRaisesRegex(errors.InternalError, "worker is already assigned"):
+    with self.assertRaisesRegex(
+        errors.InternalError, "worker is already assigned"):
       cluster.restart_dispatcher()
 
   @combinations.generate(test_base.default_test_combinations())

From 9725bf8c7d392187779e3a9b6a349665c4470dbc Mon Sep 17 00:00:00 2001
From: Sai Ramineni <csramineni@gmail.com>
Date: Thu, 1 Oct 2026 20:21:01 -0400
Subject: [PATCH 0132/1000] Narrow ragged-to-sparse test dependencies

---
 tensorflow/core/kernels/ragged_tensor_to_sparse_kernel.cc | 3 +--
 tensorflow/python/ops/ragged/BUILD                        | 2 +-
 tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py  | 6 +++---
 3 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel.cc b/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel.cc
index 2b35ff44b08..b9f1823d3d7 100644
--- a/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel.cc
+++ b/tensorflow/core/kernels/ragged_tensor_to_sparse_kernel.cc
@@ -12,7 +12,6 @@ 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.
 ==============================================================================*/
-#include <cstddef>
 #include <limits>
 #include <memory>
 #include <string>
@@ -97,7 +96,7 @@ class RaggedTensorToSparseOp : public OpKernel {
         }
 
         // Update index_prefix.
-        for (size_t dim = 0; dim < index_prefix.size(); ++dim) {
+        for (int dim = 0; dim < static_cast<int>(index_prefix.size()); ++dim) {
           int64_t start = dim > 0 ? rt_nested_splits[dim - 1](pos[dim - 1]) : 0;
           index_prefix[dim] = pos[dim] - start;
         }
diff --git a/tensorflow/python/ops/ragged/BUILD b/tensorflow/python/ops/ragged/BUILD
index 3dc2ed29649..d48db0fb6a6 100644
--- a/tensorflow/python/ops/ragged/BUILD
+++ b/tensorflow/python/ops/ragged/BUILD
@@ -1062,7 +1062,6 @@ py_test(
         ":ragged_factory_ops",
         ":ragged_functional_ops",
         ":ragged_tensor",
-        "//tensorflow:tensorflow_py_no_contrib",
         "//tensorflow/python/eager:context",
         "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/framework:errors",
@@ -1070,6 +1069,7 @@ py_test(
         "//tensorflow/python/ops:array_ops",
         "//tensorflow/python/ops:gradients_impl",
         "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:ragged_conversion_ops_gen",
         "//tensorflow/python/platform:test",
     ],
 )
diff --git a/tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py b/tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py
index f80d17fd78c..4693494b288 100644
--- a/tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py
@@ -14,12 +14,12 @@
 # ==============================================================================
 """Tests for ragged.to_sparse op."""
 
-from tensorflow import raw_ops
 from tensorflow.python.eager import context
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_ragged_conversion_ops
 from tensorflow.python.ops import gradients_impl
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops.ragged import ragged_factory_ops
@@ -184,12 +184,12 @@ class RaggedTensorToSparseOpTest(test_util.TensorFlowTestCase):
         if not context.executing_eagerly():
           with self.assertRaisesRegex(
               ValueError, "Shape must be at least rank 1"):
-            raw_ops.RaggedTensorToSparse(
+            gen_ragged_conversion_ops.ragged_tensor_to_sparse(
                 rt_nested_splits=splits, rt_dense_values=values)
           values = array_ops.placeholder_with_default(values, shape=None)
         with self.assertRaisesRegex(
             errors.InvalidArgumentError, "rt_dense_values must have rank >= 1"):
-          self.evaluate(raw_ops.RaggedTensorToSparse(
+          self.evaluate(gen_ragged_conversion_ops.ragged_tensor_to_sparse(
               rt_nested_splits=splits, rt_dense_values=values))
 
   def testGradient(self):

From 21a7d74c86e51f3ab421d21fe34ad829eb1ed6e7 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Fri, 2 Oct 2026 09:52:49 +0530
Subject: [PATCH 0133/1000] Replace ensure_initialized() calls with method
 reference

---
 tensorflow/python/eager/context.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/eager/context.py b/tensorflow/python/eager/context.py
index 772cd71936c..2435b922fac 100644
--- a/tensorflow/python/eager/context.py
+++ b/tensorflow/python/eager/context.py
@@ -943,13 +943,13 @@ class Context:
     return self._coordination_service_config
 
   def set_config_key_value(self, key, value):
-    ensure_initialized()
+    self.ensure_initialized()
     pywrap_tfe.TFE_InsertConfigKeyValue(self._context_handle, key, value)
 
   # If `timeout_in_ms=0`, this will block until the key-value is set or the
   # worker shuts down.
   def get_config_key_value(self, key, timeout_in_ms=0):
-    ensure_initialized()
+    self.ensure_initialized()
     with c_api_util.tf_buffer() as buffer_:
       pywrap_tfe.TFE_GetConfigKeyValue(
           self._context_handle, key, timeout_in_ms, buffer_
@@ -958,7 +958,7 @@ class Context:
     return value
 
   def delete_config_key_value(self, key):
-    ensure_initialized()
+    self.ensure_initialized()
     pywrap_tfe.TFE_DeleteConfigKeyValue(self._context_handle, key)
 
   def report_error_to_cluster(self, error_code, error_message):
@@ -1001,7 +1001,7 @@ class Context:
       barrier_id: Unique string identifying the barrier.
       timeout_in_ms: Duration before the barrier times out and fails.
     """
-    ensure_initialized()
+    self.ensure_initialized()
     pywrap_tfe.TFE_WaitAtBarrier(
         self._context_handle, barrier_id, timeout_in_ms
     )

From 2c8195f161da832658bcd0b89bf359f7f0f9d219 Mon Sep 17 00:00:00 2001
From: Sai Ramineni <csramineni@gmail.com>
Date: Wed, 30 Sep 2026 21:25:55 -0400
Subject: [PATCH 0134/1000] Validate unigram sampler ranges before construction

---
 .../core/kernels/candidate_sampler_ops.cc     | 11 ++++
 tensorflow/python/kernel_tests/random/BUILD   |  1 +
 .../random/candidate_sampler_ops_test.py      | 57 +++++++++++++++++++
 .../python/ops/candidate_sampling_ops.py      |  4 +-
 4 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/kernels/candidate_sampler_ops.cc b/tensorflow/core/kernels/candidate_sampler_ops.cc
index 02881615852..9c3b7c83a5f 100644
--- a/tensorflow/core/kernels/candidate_sampler_ops.cc
+++ b/tensorflow/core/kernels/candidate_sampler_ops.cc
@@ -17,8 +17,10 @@ limitations under the License.
 
 #include <cstddef>
 #include <cstdint>
+#include <limits>
 #include <memory>
 #include <string>
+#include <type_traits>
 
 #include "absl/log/check.h"
 #include "absl/status/status.h"
@@ -133,6 +135,15 @@ class SimpleCandidateSamplerOp : public BaseCandidateSamplerOp {
       : BaseCandidateSamplerOp(context) {
     int64_t range_max;
     OP_REQUIRES_OK(context, context->GetAttr("range_max", &range_max));
+    if constexpr (std::is_same_v<RangeSamplerType, UnigramSampler> ||
+                  std::is_same_v<RangeSamplerType, ThreadUnsafeUnigramSampler>) {
+      OP_REQUIRES(
+          context, range_max < std::numeric_limits<int32_t>::max(),
+          absl::InvalidArgumentError(absl::StrCat(
+              "range_max must be less than ",
+              std::numeric_limits<int32_t>::max(),
+              " for unigram samplers, got ", range_max)));
+    }
     set_sampler(new RangeSamplerType(range_max));
   }
 };
diff --git a/tensorflow/python/kernel_tests/random/BUILD b/tensorflow/python/kernel_tests/random/BUILD
index dfbd41a7d90..cbf2706aba7 100644
--- a/tensorflow/python/kernel_tests/random/BUILD
+++ b/tensorflow/python/kernel_tests/random/BUILD
@@ -39,6 +39,7 @@ tf_py_strict_test(
         "//tensorflow/python/framework:test_lib",
         "//tensorflow/python/ops:array_ops",
         "//tensorflow/python/ops:candidate_sampling_ops",
+        "//tensorflow/python/ops:candidate_sampling_ops_gen",
         "//tensorflow/python/ops:math_ops",
         "//tensorflow/python/platform:client_testlib",
         "//third_party/py/numpy",
diff --git a/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py b/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
index 396843ace3a..29ef40ecd97 100644
--- a/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
+++ b/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
@@ -22,6 +22,7 @@ from tensorflow.python.framework import errors
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import candidate_sampling_ops
+from tensorflow.python.ops import gen_candidate_sampling_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.platform import test
 
@@ -128,6 +129,62 @@ class RangeSamplerOpsTest(test.TestCase):
     # twice very rarely.
     self.assertLessEqual(num_same, 2)
 
+  @test_util.run_in_graph_and_eager_modes
+  def testUnigramRangeTooLarge(self):
+    for sampler in (
+        gen_candidate_sampling_ops.thread_unsafe_unigram_candidate_sampler,
+        gen_candidate_sampling_ops.learned_unigram_candidate_sampler):
+      # The first range truncates to one, keeping the unfixed reproducer small.
+      for range_max in (2**32 + 1, np.iinfo(np.int32).max,
+                        np.iinfo(np.int64).max):
+        with self.subTest(sampler=sampler.__name__, range_max=range_max):
+          with self.assertRaisesRegex(
+              errors.InvalidArgumentError,
+              "range_max must be less than 2147483647"):
+            result = sampler(
+                true_classes=constant_op.constant([[0]], dtype=dtypes.int64),
+                num_true=1, num_sampled=1, unique=False, range_max=range_max)
+            self.evaluate(result)
+
+  def testLearnedUnigramRangeTooLarge(self):
+    for range_max in (np.iinfo(np.int32).max, 2**32 + 1):
+      with self.subTest(range_max=range_max):
+        with self.assertRaisesRegex(ValueError, "too large to handle"):
+          candidate_sampling_ops.learned_unigram_candidate_sampler(
+              true_classes=[[0]], num_true=1, num_sampled=1,
+              unique=False, range_max=range_max)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testUnigramSingleClass(self):
+    for sampler in (
+        gen_candidate_sampling_ops.thread_unsafe_unigram_candidate_sampler,
+        gen_candidate_sampling_ops.learned_unigram_candidate_sampler):
+      for unique in (False, True):
+        with self.subTest(sampler=sampler.__name__, unique=unique):
+          sampled, true_count, sampled_count = self.evaluate(sampler(
+              true_classes=constant_op.constant([[0]], dtype=dtypes.int64),
+              num_true=1, num_sampled=1, unique=unique, range_max=1))
+          self.assertAllEqual(sampled, [0])
+          self.assertAllEqual(true_count, [[1.0]])
+          self.assertAllEqual(sampled_count, [1.0])
+
+  @test_util.run_in_graph_and_eager_modes
+  def testLargeUniformRange(self):
+    range_max = 2**32 + 1
+    for sampler in (gen_candidate_sampling_ops.uniform_candidate_sampler,
+                    gen_candidate_sampling_ops.log_uniform_candidate_sampler):
+      for unique in (False, True):
+        with self.subTest(sampler=sampler.__name__, unique=unique):
+          sampled, true_count, sampled_count = self.evaluate(sampler(
+              true_classes=constant_op.constant([[0]], dtype=dtypes.int64),
+              num_true=1, num_sampled=2, unique=unique,
+              range_max=range_max, seed=123))
+          self.assertTrue(np.all(sampled >= 0))
+          self.assertTrue(np.all(sampled < range_max))
+          for count in (true_count, sampled_count):
+            self.assertTrue(np.all(np.isfinite(count)))
+            self.assertTrue(np.all(count > 0))
+
   def testCandidateOutOfRange(self):
     with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
                                 "out of range"):
diff --git a/tensorflow/python/ops/candidate_sampling_ops.py b/tensorflow/python/ops/candidate_sampling_ops.py
index d4291b0aa3f..d7beb9929db 100644
--- a/tensorflow/python/ops/candidate_sampling_ops.py
+++ b/tensorflow/python/ops/candidate_sampling_ops.py
@@ -295,8 +295,8 @@ def learned_unigram_candidate_sampler(true_classes, num_true, num_sampled,
 
   """
   seed1, seed2 = random_seed.get_seed(seed)
-  # Limiting to Max int32 value
-  if range_max > 2147483647:
+  # Unigram samplers require range_max < INT32_MAX.
+  if range_max >= 2147483647:
     raise ValueError(f'Value of range_max:{range_max} is too large to handle')
   return gen_candidate_sampling_ops.learned_unigram_candidate_sampler(
       true_classes, num_true, num_sampled, unique, range_max, seed=seed1,

From 0abe15679c6e54f51e03f06456ba05098db39944 Mon Sep 17 00:00:00 2001
From: Sai Ramineni <csramineni@gmail.com>
Date: Thu, 1 Oct 2026 02:27:08 -0400
Subject: [PATCH 0135/1000] Allow candidate sampler tests to depend on
 generated wrappers

---
 tensorflow/python/ops/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD
index 675b0563574..18d1db642e8 100644
--- a/tensorflow/python/ops/BUILD
+++ b/tensorflow/python/ops/BUILD
@@ -210,6 +210,7 @@ tf_gen_op_strict_wrapper_private_py(
     visibility = [
         "//learning/brain/python/ops:__pkg__",
         "//tensorflow/python:__pkg__",
+        "//tensorflow/python/kernel_tests/random:__pkg__",
     ],
 )
 

From 0b6a684f973590193e78f67d10b0e52a112ea026 Mon Sep 17 00:00:00 2001
From: Sai Ramineni <csramineni@gmail.com>
Date: Thu, 1 Oct 2026 13:44:13 -0400
Subject: [PATCH 0136/1000] Align unigram shape validation and test the public
 sampler APIs

---
 .../core/kernels/candidate_sampler_ops.cc     |  2 +
 tensorflow/core/ops/candidate_sampling_ops.cc | 18 ++++++++-
 .../core/ops/candidate_sampling_ops_test.cc   | 26 +++++++++++++
 tensorflow/python/kernel_tests/random/BUILD   |  2 +-
 .../random/candidate_sampler_ops_test.py      | 37 ++++++++++++++-----
 tensorflow/python/ops/BUILD                   |  1 -
 .../python/ops/candidate_sampling_ops.py      |  3 --
 7 files changed, 72 insertions(+), 17 deletions(-)

diff --git a/tensorflow/core/kernels/candidate_sampler_ops.cc b/tensorflow/core/kernels/candidate_sampler_ops.cc
index 9c3b7c83a5f..e40a27dfd61 100644
--- a/tensorflow/core/kernels/candidate_sampler_ops.cc
+++ b/tensorflow/core/kernels/candidate_sampler_ops.cc
@@ -135,6 +135,8 @@ class SimpleCandidateSamplerOp : public BaseCandidateSamplerOp {
       : BaseCandidateSamplerOp(context) {
     int64_t range_max;
     OP_REQUIRES_OK(context, context->GetAttr("range_max", &range_max));
+    OP_REQUIRES(context, range_max > 0,
+                absl::InvalidArgumentError("range_max must be positive"));
     if constexpr (std::is_same_v<RangeSamplerType, UnigramSampler> ||
                   std::is_same_v<RangeSamplerType, ThreadUnsafeUnigramSampler>) {
       OP_REQUIRES(
diff --git a/tensorflow/core/ops/candidate_sampling_ops.cc b/tensorflow/core/ops/candidate_sampling_ops.cc
index c44114a787a..7e8ed4d4756 100644
--- a/tensorflow/core/ops/candidate_sampling_ops.cc
+++ b/tensorflow/core/ops/candidate_sampling_ops.cc
@@ -13,6 +13,10 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <cstdint>
+#include <limits>
+
+#include "absl/status/status.h"
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/shape_inference.h"
 
@@ -41,6 +45,16 @@ absl::Status CandidateSamplerShapeFn(InferenceContext* c) {
   return absl::OkStatus();
 }
 
+absl::Status UnigramCandidateSamplerShapeFn(InferenceContext* c) {
+  int64_t range_max;
+  TF_RETURN_IF_ERROR(c->GetAttr("range_max", &range_max));
+  if (range_max >= std::numeric_limits<int32_t>::max()) {
+    return absl::InvalidArgumentError(
+        "range_max must be less than 2147483647 for unigram samplers");
+  }
+  return CandidateSamplerShapeFn(c);
+}
+
 }  // namespace
 
 REGISTER_OP("UniformCandidateSampler")
@@ -82,7 +96,7 @@ REGISTER_OP("LearnedUnigramCandidateSampler")
     .Attr("range_max: int >= 1")
     .Attr("seed: int = 0")
     .Attr("seed2: int = 0")
-    .SetShapeFn(CandidateSamplerShapeFn)
+    .SetShapeFn(UnigramCandidateSamplerShapeFn)
     .SetIsStateful();
 
 REGISTER_OP("ThreadUnsafeUnigramCandidateSampler")
@@ -96,7 +110,7 @@ REGISTER_OP("ThreadUnsafeUnigramCandidateSampler")
     .Attr("range_max: int >= 1")
     .Attr("seed: int = 0")
     .Attr("seed2: int = 0")
-    .SetShapeFn(CandidateSamplerShapeFn)
+    .SetShapeFn(UnigramCandidateSamplerShapeFn)
     .SetIsStateful();
 
 REGISTER_OP("FixedUnigramCandidateSampler")
diff --git a/tensorflow/core/ops/candidate_sampling_ops_test.cc b/tensorflow/core/ops/candidate_sampling_ops_test.cc
index f3673716040..a5bd9a649d9 100644
--- a/tensorflow/core/ops/candidate_sampling_ops_test.cc
+++ b/tensorflow/core/ops/candidate_sampling_ops_test.cc
@@ -13,6 +13,9 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <cstdint>
+#include <limits>
+
 #include "tensorflow/core/framework/node_def_builder.h"
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/shape_inference_testutil.h"
@@ -36,6 +39,9 @@ TEST(CandidateSamplerOpsTest, CandidateSampler_ShapeFn) {
                      .Attr("num_sampled", 5)
                      .Attr("num_true", 10)
                      .Finalize(&op.node_def));
+    if (op.name != "AllCandidateSampler") {
+      (*op.node_def.mutable_attr())["range_max"].set_i(100);
+    }
 
     // num_sampled = 5, num_true = 10.
     INFER_OK(op, "?", "[5];[?,10];[5]");
@@ -49,6 +55,26 @@ TEST(CandidateSamplerOpsTest, CandidateSampler_ShapeFn) {
   }
 }
 
+TEST(CandidateSamplerOpsTest, UnigramRange_ShapeFn) {
+  for (const char* op_name : {"LearnedUnigramCandidateSampler",
+                            "ThreadUnsafeUnigramCandidateSampler"}) {
+    ShapeInferenceTestOp op(op_name);
+    TF_ASSERT_OK(NodeDefBuilder("test", op.name)
+                     .Input({"a", 0, DT_INT64})
+                     .Attr("num_sampled", 1)
+                     .Attr("num_true", 1)
+                     .Attr("range_max", std::numeric_limits<int32_t>::max() - 1)
+                     .Finalize(&op.node_def));
+    INFER_OK(op, "[1,1]", "[1];[d0_0,1];[1]");
+    for (int64_t range_max : {int64_t{std::numeric_limits<int32_t>::max()},
+                              (int64_t{1} << 32) + 1,
+                              std::numeric_limits<int64_t>::max()}) {
+      (*op.node_def.mutable_attr())["range_max"].set_i(range_max);
+      INFER_ERROR("range_max must be less than 2147483647", op, "[1,1]");
+    }
+  }
+}
+
 TEST(CandidateSamplerOpsTest, ComputeAccidentalHits_ShapeFn) {
   ShapeInferenceTestOp op("ComputeAccidentalHits");
   TF_ASSERT_OK(NodeDefBuilder("test", op.name)
diff --git a/tensorflow/python/kernel_tests/random/BUILD b/tensorflow/python/kernel_tests/random/BUILD
index cbf2706aba7..36d63adacc0 100644
--- a/tensorflow/python/kernel_tests/random/BUILD
+++ b/tensorflow/python/kernel_tests/random/BUILD
@@ -33,13 +33,13 @@ tf_py_strict_test(
     size = "small",
     srcs = ["candidate_sampler_ops_test.py"],
     deps = [
+        "//tensorflow:tensorflow_py_no_contrib",
         "//tensorflow/python/framework:constant_op",
         "//tensorflow/python/framework:errors",
         "//tensorflow/python/framework:for_generated_wrappers",
         "//tensorflow/python/framework:test_lib",
         "//tensorflow/python/ops:array_ops",
         "//tensorflow/python/ops:candidate_sampling_ops",
-        "//tensorflow/python/ops:candidate_sampling_ops_gen",
         "//tensorflow/python/ops:math_ops",
         "//tensorflow/python/platform:client_testlib",
         "//third_party/py/numpy",
diff --git a/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py b/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
index 29ef40ecd97..e510863c1f1 100644
--- a/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
+++ b/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
@@ -16,13 +16,13 @@
 
 import numpy as np
 
+from tensorflow import raw_ops
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import candidate_sampling_ops
-from tensorflow.python.ops import gen_candidate_sampling_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.platform import test
 
@@ -132,33 +132,50 @@ class RangeSamplerOpsTest(test.TestCase):
   @test_util.run_in_graph_and_eager_modes
   def testUnigramRangeTooLarge(self):
     for sampler in (
-        gen_candidate_sampling_ops.thread_unsafe_unigram_candidate_sampler,
-        gen_candidate_sampling_ops.learned_unigram_candidate_sampler):
+        raw_ops.ThreadUnsafeUnigramCandidateSampler,
+        raw_ops.LearnedUnigramCandidateSampler):
       # The first range truncates to one, keeping the unfixed reproducer small.
       for range_max in (2**32 + 1, np.iinfo(np.int32).max,
                         np.iinfo(np.int64).max):
         with self.subTest(sampler=sampler.__name__, range_max=range_max):
           with self.assertRaisesRegex(
-              errors.InvalidArgumentError,
+              (ValueError, errors.InvalidArgumentError),
               "range_max must be less than 2147483647"):
             result = sampler(
                 true_classes=constant_op.constant([[0]], dtype=dtypes.int64),
                 num_true=1, num_sampled=1, unique=False, range_max=range_max)
             self.evaluate(result)
 
+  @test_util.run_in_graph_and_eager_modes
   def testLearnedUnigramRangeTooLarge(self):
     for range_max in (np.iinfo(np.int32).max, 2**32 + 1):
       with self.subTest(range_max=range_max):
-        with self.assertRaisesRegex(ValueError, "too large to handle"):
-          candidate_sampling_ops.learned_unigram_candidate_sampler(
+        with self.assertRaisesRegex(
+            (ValueError, errors.InvalidArgumentError),
+            "range_max must be less than 2147483647"):
+          result = candidate_sampling_ops.learned_unigram_candidate_sampler(
               true_classes=[[0]], num_true=1, num_sampled=1,
               unique=False, range_max=range_max)
+          self.evaluate(result)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testUnigramNonpositiveRange(self):
+    for sampler in (raw_ops.ThreadUnsafeUnigramCandidateSampler,
+                    raw_ops.LearnedUnigramCandidateSampler):
+      for range_max in (0, -1):
+        with self.subTest(sampler=sampler.__name__, range_max=range_max):
+          with self.assertRaisesRegex(
+              (ValueError, errors.InvalidArgumentError), "range_max"):
+            result = sampler(
+                true_classes=constant_op.constant([[0]], dtype=dtypes.int64),
+                num_true=1, num_sampled=1, unique=False, range_max=range_max)
+            self.evaluate(result)
 
   @test_util.run_in_graph_and_eager_modes
   def testUnigramSingleClass(self):
     for sampler in (
-        gen_candidate_sampling_ops.thread_unsafe_unigram_candidate_sampler,
-        gen_candidate_sampling_ops.learned_unigram_candidate_sampler):
+        raw_ops.ThreadUnsafeUnigramCandidateSampler,
+        raw_ops.LearnedUnigramCandidateSampler):
       for unique in (False, True):
         with self.subTest(sampler=sampler.__name__, unique=unique):
           sampled, true_count, sampled_count = self.evaluate(sampler(
@@ -171,8 +188,8 @@ class RangeSamplerOpsTest(test.TestCase):
   @test_util.run_in_graph_and_eager_modes
   def testLargeUniformRange(self):
     range_max = 2**32 + 1
-    for sampler in (gen_candidate_sampling_ops.uniform_candidate_sampler,
-                    gen_candidate_sampling_ops.log_uniform_candidate_sampler):
+    for sampler in (raw_ops.UniformCandidateSampler,
+                    raw_ops.LogUniformCandidateSampler):
       for unique in (False, True):
         with self.subTest(sampler=sampler.__name__, unique=unique):
           sampled, true_count, sampled_count = self.evaluate(sampler(
diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD
index 18d1db642e8..675b0563574 100644
--- a/tensorflow/python/ops/BUILD
+++ b/tensorflow/python/ops/BUILD
@@ -210,7 +210,6 @@ tf_gen_op_strict_wrapper_private_py(
     visibility = [
         "//learning/brain/python/ops:__pkg__",
         "//tensorflow/python:__pkg__",
-        "//tensorflow/python/kernel_tests/random:__pkg__",
     ],
 )
 
diff --git a/tensorflow/python/ops/candidate_sampling_ops.py b/tensorflow/python/ops/candidate_sampling_ops.py
index d7beb9929db..9bae77c63df 100644
--- a/tensorflow/python/ops/candidate_sampling_ops.py
+++ b/tensorflow/python/ops/candidate_sampling_ops.py
@@ -295,9 +295,6 @@ def learned_unigram_candidate_sampler(true_classes, num_true, num_sampled,
 
   """
   seed1, seed2 = random_seed.get_seed(seed)
-  # Unigram samplers require range_max < INT32_MAX.
-  if range_max >= 2147483647:
-    raise ValueError(f'Value of range_max:{range_max} is too large to handle')
   return gen_candidate_sampling_ops.learned_unigram_candidate_sampler(
       true_classes, num_true, num_sampled, unique, range_max, seed=seed1,
       seed2=seed2, name=name)

From 1e4a2249befc483c4a04fda9cd2d6b6730989173 Mon Sep 17 00:00:00 2001
From: Sai Ramineni <csramineni@gmail.com>
Date: Thu, 1 Oct 2026 14:23:20 -0400
Subject: [PATCH 0137/1000] Use TensorFlow error helpers for candidate sampler
 validation

---
 .../core/kernels/candidate_sampler_ops.cc     | 22 +++++++++----------
 tensorflow/core/ops/candidate_sampling_ops.cc |  4 ++--
 2 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/tensorflow/core/kernels/candidate_sampler_ops.cc b/tensorflow/core/kernels/candidate_sampler_ops.cc
index e40a27dfd61..a11e835628d 100644
--- a/tensorflow/core/kernels/candidate_sampler_ops.cc
+++ b/tensorflow/core/kernels/candidate_sampler_ops.cc
@@ -23,7 +23,6 @@ limitations under the License.
 #include <type_traits>
 
 #include "absl/log/check.h"
-#include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/types/span.h"
 #define EIGEN_USE_THREADS
@@ -35,6 +34,7 @@ limitations under the License.
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/kernels/range_sampler.h"
+#include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/util/guarded_philox_random.h"
 
@@ -53,10 +53,10 @@ class BaseCandidateSamplerOp : public OpKernel {
   void Compute(OpKernelContext* context) override {
     const Tensor& true_classes = context->input(0);
     OP_REQUIRES(context, true_classes.dims() == 2,
-                absl::InvalidArgumentError("true_classes must be a matrix"));
+                errors::InvalidArgument("true_classes must be a matrix"));
     const int32_t batch_size = true_classes.dim_size(0);
     OP_REQUIRES(context, true_classes.dim_size(1) == num_true_,
-                absl::InvalidArgumentError(absl::StrCat(
+                errors::InvalidArgument(absl::StrCat(
                     "true_classes must have "
                     "num_true columns, expected: ",
                     true_classes.dim_size(1), " was: ", num_true_)));
@@ -64,7 +64,7 @@ class BaseCandidateSamplerOp : public OpKernel {
 
     if (unique_) {
       OP_REQUIRES(context, num_sampled_ <= sampler_->range(),
-                  absl::InvalidArgumentError("Sampler's range is too small."));
+                  errors::InvalidArgument("Sampler's range is too small."));
     }
 
     // Output candidates and expected_count.
@@ -87,7 +87,7 @@ class BaseCandidateSamplerOp : public OpKernel {
 
     for (const auto& candidate : true_candidate) {
       OP_REQUIRES(context, candidate >= 0 && candidate < sampler_->range(),
-                  absl::InvalidArgumentError(absl::StrCat(
+                  errors::InvalidArgument(absl::StrCat(
                       "`true_candidate` out of range [", 0, ", ",
                       sampler_->range(), "), received ", candidate)));
     }
@@ -136,12 +136,12 @@ class SimpleCandidateSamplerOp : public BaseCandidateSamplerOp {
     int64_t range_max;
     OP_REQUIRES_OK(context, context->GetAttr("range_max", &range_max));
     OP_REQUIRES(context, range_max > 0,
-                absl::InvalidArgumentError("range_max must be positive"));
+                errors::InvalidArgument("range_max must be positive"));
     if constexpr (std::is_same_v<RangeSamplerType, UnigramSampler> ||
                   std::is_same_v<RangeSamplerType, ThreadUnsafeUnigramSampler>) {
       OP_REQUIRES(
           context, range_max < std::numeric_limits<int32_t>::max(),
-          absl::InvalidArgumentError(absl::StrCat(
+          errors::InvalidArgument(absl::StrCat(
               "range_max must be less than ",
               std::numeric_limits<int32_t>::max(),
               " for unigram samplers, got ", range_max)));
@@ -188,10 +188,10 @@ class FixedUnigramCandidateSamplerOp : public BaseCandidateSamplerOp {
     std::vector<float> unigrams;
     OP_REQUIRES_OK(context, context->GetAttr("unigrams", &unigrams));
     OP_REQUIRES(context, !vocab_file.empty() || !unigrams.empty(),
-                absl::InvalidArgumentError(
+                errors::InvalidArgument(
                     "Must provide either vocab_file or unigrams."));
     OP_REQUIRES(context, vocab_file.empty() || unigrams.empty(),
-                absl::InvalidArgumentError(
+                errors::InvalidArgument(
                     "Must only provide one of vocab_file and unigrams."));
     float distortion;
     OP_REQUIRES_OK(context, context->GetAttr("distortion", &distortion));
@@ -229,7 +229,7 @@ class ComputeAccidentalHitsOp : public OpKernel {
     OP_REQUIRES(context,
                 TensorShapeUtils::IsMatrix(in_true_candidates_shape) &&
                     in_true_candidates_shape.dim_size(1) == num_true_,
-                absl::InvalidArgumentError(
+                errors::InvalidArgument(
                     "true_candidates must be a batch_size * num_true matrix"));
 
     const int64_t batch_size = in_true_candidates_shape.dim_size(0);
@@ -237,7 +237,7 @@ class ComputeAccidentalHitsOp : public OpKernel {
     const Tensor& in_sampled_candidates = context->input(1);
     OP_REQUIRES(context,
                 TensorShapeUtils::IsVector(in_sampled_candidates.shape()),
-                absl::InvalidArgumentError(
+                errors::InvalidArgument(
                     "sampled_candidates must be a vector, which is typically "
                     "an output from CandidateSampler"));
 
diff --git a/tensorflow/core/ops/candidate_sampling_ops.cc b/tensorflow/core/ops/candidate_sampling_ops.cc
index 7e8ed4d4756..4c093e0577f 100644
--- a/tensorflow/core/ops/candidate_sampling_ops.cc
+++ b/tensorflow/core/ops/candidate_sampling_ops.cc
@@ -16,9 +16,9 @@ limitations under the License.
 #include <cstdint>
 #include <limits>
 
-#include "absl/status/status.h"
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/platform/errors.h"
 
 namespace tensorflow {
 
@@ -49,7 +49,7 @@ absl::Status UnigramCandidateSamplerShapeFn(InferenceContext* c) {
   int64_t range_max;
   TF_RETURN_IF_ERROR(c->GetAttr("range_max", &range_max));
   if (range_max >= std::numeric_limits<int32_t>::max()) {
-    return absl::InvalidArgumentError(
+    return errors::InvalidArgument(
         "range_max must be less than 2147483647 for unigram samplers");
   }
   return CandidateSamplerShapeFn(c);

From a8d8a50629e8f0a25e13eafb20d2ba25280c4d41 Mon Sep 17 00:00:00 2001
From: Sai Ramineni <csramineni@gmail.com>
Date: Thu, 1 Oct 2026 15:50:24 -0400
Subject: [PATCH 0138/1000] Use generated sampler wrappers with direct test
 dependencies

---
 tensorflow/python/kernel_tests/random/BUILD                    | 3 ++-
 .../python/kernel_tests/random/candidate_sampler_ops_test.py   | 2 +-
 tensorflow/python/ops/BUILD                                    | 1 +
 3 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/kernel_tests/random/BUILD b/tensorflow/python/kernel_tests/random/BUILD
index 36d63adacc0..236c5cdddf8 100644
--- a/tensorflow/python/kernel_tests/random/BUILD
+++ b/tensorflow/python/kernel_tests/random/BUILD
@@ -33,13 +33,14 @@ tf_py_strict_test(
     size = "small",
     srcs = ["candidate_sampler_ops_test.py"],
     deps = [
-        "//tensorflow:tensorflow_py_no_contrib",
         "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/framework:errors",
         "//tensorflow/python/framework:for_generated_wrappers",
         "//tensorflow/python/framework:test_lib",
         "//tensorflow/python/ops:array_ops",
         "//tensorflow/python/ops:candidate_sampling_ops",
+        "//tensorflow/python/ops:candidate_sampling_ops_gen",
         "//tensorflow/python/ops:math_ops",
         "//tensorflow/python/platform:client_testlib",
         "//third_party/py/numpy",
diff --git a/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py b/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
index e510863c1f1..fa7fb4b86d3 100644
--- a/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
+++ b/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
@@ -16,13 +16,13 @@
 
 import numpy as np
 
-from tensorflow import raw_ops
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import candidate_sampling_ops
+from tensorflow.python.ops import gen_candidate_sampling_ops as raw_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.platform import test
 
diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD
index 675b0563574..18d1db642e8 100644
--- a/tensorflow/python/ops/BUILD
+++ b/tensorflow/python/ops/BUILD
@@ -210,6 +210,7 @@ tf_gen_op_strict_wrapper_private_py(
     visibility = [
         "//learning/brain/python/ops:__pkg__",
         "//tensorflow/python:__pkg__",
+        "//tensorflow/python/kernel_tests/random:__pkg__",
     ],
 )
 

From abe1c88f873b8713fb4bea630e9014da758774fa Mon Sep 17 00:00:00 2001
From: Sai Ramineni <csramineni@gmail.com>
Date: Fri, 2 Oct 2026 01:48:34 -0400
Subject: [PATCH 0139/1000] Use public candidate sampler APIs in regression
 tests

---
 tensorflow/python/kernel_tests/random/BUILD                 | 2 +-
 .../kernel_tests/random/candidate_sampler_ops_test.py       | 6 +++---
 tensorflow/python/ops/BUILD                                 | 1 -
 3 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/tensorflow/python/kernel_tests/random/BUILD b/tensorflow/python/kernel_tests/random/BUILD
index 236c5cdddf8..41a8729f44a 100644
--- a/tensorflow/python/kernel_tests/random/BUILD
+++ b/tensorflow/python/kernel_tests/random/BUILD
@@ -33,6 +33,7 @@ tf_py_strict_test(
     size = "small",
     srcs = ["candidate_sampler_ops_test.py"],
     deps = [
+        "//tensorflow:tensorflow_py_no_contrib",
         "//tensorflow/python/framework:constant_op",
         "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/framework:errors",
@@ -40,7 +41,6 @@ tf_py_strict_test(
         "//tensorflow/python/framework:test_lib",
         "//tensorflow/python/ops:array_ops",
         "//tensorflow/python/ops:candidate_sampling_ops",
-        "//tensorflow/python/ops:candidate_sampling_ops_gen",
         "//tensorflow/python/ops:math_ops",
         "//tensorflow/python/platform:client_testlib",
         "//third_party/py/numpy",
diff --git a/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py b/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
index fa7fb4b86d3..16ebd9aef1a 100644
--- a/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
+++ b/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
@@ -16,13 +16,13 @@
 
 import numpy as np
 
+from tensorflow import raw_ops
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import candidate_sampling_ops
-from tensorflow.python.ops import gen_candidate_sampling_ops as raw_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.platform import test
 
@@ -188,8 +188,8 @@ class RangeSamplerOpsTest(test.TestCase):
   @test_util.run_in_graph_and_eager_modes
   def testLargeUniformRange(self):
     range_max = 2**32 + 1
-    for sampler in (raw_ops.UniformCandidateSampler,
-                    raw_ops.LogUniformCandidateSampler):
+    for sampler in (candidate_sampling_ops.uniform_candidate_sampler,
+                    candidate_sampling_ops.log_uniform_candidate_sampler):
       for unique in (False, True):
         with self.subTest(sampler=sampler.__name__, unique=unique):
           sampled, true_count, sampled_count = self.evaluate(sampler(
diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD
index 18d1db642e8..675b0563574 100644
--- a/tensorflow/python/ops/BUILD
+++ b/tensorflow/python/ops/BUILD
@@ -210,7 +210,6 @@ tf_gen_op_strict_wrapper_private_py(
     visibility = [
         "//learning/brain/python/ops:__pkg__",
         "//tensorflow/python:__pkg__",
-        "//tensorflow/python/kernel_tests/random:__pkg__",
     ],
 )
 

From 3e11a8d7a37c6457b3c8e2097c797e59d2160b98 Mon Sep 17 00:00:00 2001
From: Sai Ramineni <csramineni@gmail.com>
Date: Fri, 2 Oct 2026 02:37:47 -0400
Subject: [PATCH 0140/1000] Bound unigram ranges by the weighted picker tree
 capacity

---
 .../core/kernels/candidate_sampler_ops.cc     | 30 +++++++++----------
 tensorflow/core/ops/candidate_sampling_ops.cc |  7 +++--
 .../core/ops/candidate_sampling_ops_test.cc   |  7 +++--
 .../random/candidate_sampler_ops_test.py      | 11 +++----
 4 files changed, 28 insertions(+), 27 deletions(-)

diff --git a/tensorflow/core/kernels/candidate_sampler_ops.cc b/tensorflow/core/kernels/candidate_sampler_ops.cc
index a11e835628d..040a534e1c7 100644
--- a/tensorflow/core/kernels/candidate_sampler_ops.cc
+++ b/tensorflow/core/kernels/candidate_sampler_ops.cc
@@ -17,12 +17,12 @@ limitations under the License.
 
 #include <cstddef>
 #include <cstdint>
-#include <limits>
 #include <memory>
 #include <string>
 #include <type_traits>
 
 #include "absl/log/check.h"
+#include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/types/span.h"
 #define EIGEN_USE_THREADS
@@ -34,7 +34,6 @@ limitations under the License.
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/kernels/range_sampler.h"
-#include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/util/guarded_philox_random.h"
 
@@ -53,10 +52,10 @@ class BaseCandidateSamplerOp : public OpKernel {
   void Compute(OpKernelContext* context) override {
     const Tensor& true_classes = context->input(0);
     OP_REQUIRES(context, true_classes.dims() == 2,
-                errors::InvalidArgument("true_classes must be a matrix"));
+                absl::InvalidArgumentError("true_classes must be a matrix"));
     const int32_t batch_size = true_classes.dim_size(0);
     OP_REQUIRES(context, true_classes.dim_size(1) == num_true_,
-                errors::InvalidArgument(absl::StrCat(
+                absl::InvalidArgumentError(absl::StrCat(
                     "true_classes must have "
                     "num_true columns, expected: ",
                     true_classes.dim_size(1), " was: ", num_true_)));
@@ -64,7 +63,7 @@ class BaseCandidateSamplerOp : public OpKernel {
 
     if (unique_) {
       OP_REQUIRES(context, num_sampled_ <= sampler_->range(),
-                  errors::InvalidArgument("Sampler's range is too small."));
+                  absl::InvalidArgumentError("Sampler's range is too small."));
     }
 
     // Output candidates and expected_count.
@@ -87,7 +86,7 @@ class BaseCandidateSamplerOp : public OpKernel {
 
     for (const auto& candidate : true_candidate) {
       OP_REQUIRES(context, candidate >= 0 && candidate < sampler_->range(),
-                  errors::InvalidArgument(absl::StrCat(
+                  absl::InvalidArgumentError(absl::StrCat(
                       "`true_candidate` out of range [", 0, ", ",
                       sampler_->range(), "), received ", candidate)));
     }
@@ -135,15 +134,14 @@ class SimpleCandidateSamplerOp : public BaseCandidateSamplerOp {
       : BaseCandidateSamplerOp(context) {
     int64_t range_max;
     OP_REQUIRES_OK(context, context->GetAttr("range_max", &range_max));
-    OP_REQUIRES(context, range_max > 0,
-                errors::InvalidArgument("range_max must be positive"));
     if constexpr (std::is_same_v<RangeSamplerType, UnigramSampler> ||
                   std::is_same_v<RangeSamplerType, ThreadUnsafeUnigramSampler>) {
+      // WeightedPicker's largest representable power-of-two level is 1 << 30.
+      constexpr int64_t kMaxUnigramRange = int64_t{1} << 30;
       OP_REQUIRES(
-          context, range_max < std::numeric_limits<int32_t>::max(),
-          errors::InvalidArgument(absl::StrCat(
-              "range_max must be less than ",
-              std::numeric_limits<int32_t>::max(),
+          context, range_max <= kMaxUnigramRange,
+          absl::InvalidArgumentError(absl::StrCat(
+              "range_max must be at most ", kMaxUnigramRange,
               " for unigram samplers, got ", range_max)));
     }
     set_sampler(new RangeSamplerType(range_max));
@@ -188,10 +186,10 @@ class FixedUnigramCandidateSamplerOp : public BaseCandidateSamplerOp {
     std::vector<float> unigrams;
     OP_REQUIRES_OK(context, context->GetAttr("unigrams", &unigrams));
     OP_REQUIRES(context, !vocab_file.empty() || !unigrams.empty(),
-                errors::InvalidArgument(
+                absl::InvalidArgumentError(
                     "Must provide either vocab_file or unigrams."));
     OP_REQUIRES(context, vocab_file.empty() || unigrams.empty(),
-                errors::InvalidArgument(
+                absl::InvalidArgumentError(
                     "Must only provide one of vocab_file and unigrams."));
     float distortion;
     OP_REQUIRES_OK(context, context->GetAttr("distortion", &distortion));
@@ -229,7 +227,7 @@ class ComputeAccidentalHitsOp : public OpKernel {
     OP_REQUIRES(context,
                 TensorShapeUtils::IsMatrix(in_true_candidates_shape) &&
                     in_true_candidates_shape.dim_size(1) == num_true_,
-                errors::InvalidArgument(
+                absl::InvalidArgumentError(
                     "true_candidates must be a batch_size * num_true matrix"));
 
     const int64_t batch_size = in_true_candidates_shape.dim_size(0);
@@ -237,7 +235,7 @@ class ComputeAccidentalHitsOp : public OpKernel {
     const Tensor& in_sampled_candidates = context->input(1);
     OP_REQUIRES(context,
                 TensorShapeUtils::IsVector(in_sampled_candidates.shape()),
-                errors::InvalidArgument(
+                absl::InvalidArgumentError(
                     "sampled_candidates must be a vector, which is typically "
                     "an output from CandidateSampler"));
 
diff --git a/tensorflow/core/ops/candidate_sampling_ops.cc b/tensorflow/core/ops/candidate_sampling_ops.cc
index 4c093e0577f..505be521268 100644
--- a/tensorflow/core/ops/candidate_sampling_ops.cc
+++ b/tensorflow/core/ops/candidate_sampling_ops.cc
@@ -14,7 +14,6 @@ limitations under the License.
 ==============================================================================*/
 
 #include <cstdint>
-#include <limits>
 
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/shape_inference.h"
@@ -48,9 +47,11 @@ absl::Status CandidateSamplerShapeFn(InferenceContext* c) {
 absl::Status UnigramCandidateSamplerShapeFn(InferenceContext* c) {
   int64_t range_max;
   TF_RETURN_IF_ERROR(c->GetAttr("range_max", &range_max));
-  if (range_max >= std::numeric_limits<int32_t>::max()) {
+  // Match the largest power-of-two level representable by WeightedPicker.
+  constexpr int64_t kMaxUnigramRange = int64_t{1} << 30;
+  if (range_max > kMaxUnigramRange) {
     return errors::InvalidArgument(
-        "range_max must be less than 2147483647 for unigram samplers");
+        "range_max must be at most 1073741824 for unigram samplers");
   }
   return CandidateSamplerShapeFn(c);
 }
diff --git a/tensorflow/core/ops/candidate_sampling_ops_test.cc b/tensorflow/core/ops/candidate_sampling_ops_test.cc
index a5bd9a649d9..23dc4973354 100644
--- a/tensorflow/core/ops/candidate_sampling_ops_test.cc
+++ b/tensorflow/core/ops/candidate_sampling_ops_test.cc
@@ -63,14 +63,15 @@ TEST(CandidateSamplerOpsTest, UnigramRange_ShapeFn) {
                      .Input({"a", 0, DT_INT64})
                      .Attr("num_sampled", 1)
                      .Attr("num_true", 1)
-                     .Attr("range_max", std::numeric_limits<int32_t>::max() - 1)
+                     .Attr("range_max", int64_t{1} << 30)
                      .Finalize(&op.node_def));
     INFER_OK(op, "[1,1]", "[1];[d0_0,1];[1]");
-    for (int64_t range_max : {int64_t{std::numeric_limits<int32_t>::max()},
+    for (int64_t range_max : {(int64_t{1} << 30) + 1,
+                              int64_t{std::numeric_limits<int32_t>::max()},
                               (int64_t{1} << 32) + 1,
                               std::numeric_limits<int64_t>::max()}) {
       (*op.node_def.mutable_attr())["range_max"].set_i(range_max);
-      INFER_ERROR("range_max must be less than 2147483647", op, "[1,1]");
+      INFER_ERROR("range_max must be at most 1073741824", op, "[1,1]");
     }
   }
 }
diff --git a/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py b/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
index 16ebd9aef1a..cb59e2bcddf 100644
--- a/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
+++ b/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
@@ -135,12 +135,12 @@ class RangeSamplerOpsTest(test.TestCase):
         raw_ops.ThreadUnsafeUnigramCandidateSampler,
         raw_ops.LearnedUnigramCandidateSampler):
       # The first range truncates to one, keeping the unfixed reproducer small.
-      for range_max in (2**32 + 1, np.iinfo(np.int32).max,
+      for range_max in (2**32 + 1, 2**30 + 1, np.iinfo(np.int32).max,
                         np.iinfo(np.int64).max):
         with self.subTest(sampler=sampler.__name__, range_max=range_max):
           with self.assertRaisesRegex(
               (ValueError, errors.InvalidArgumentError),
-              "range_max must be less than 2147483647"):
+              "range_max must be at most 1073741824"):
             result = sampler(
                 true_classes=constant_op.constant([[0]], dtype=dtypes.int64),
                 num_true=1, num_sampled=1, unique=False, range_max=range_max)
@@ -148,11 +148,11 @@ class RangeSamplerOpsTest(test.TestCase):
 
   @test_util.run_in_graph_and_eager_modes
   def testLearnedUnigramRangeTooLarge(self):
-    for range_max in (np.iinfo(np.int32).max, 2**32 + 1):
+    for range_max in (2**30 + 1, np.iinfo(np.int32).max, 2**32 + 1):
       with self.subTest(range_max=range_max):
         with self.assertRaisesRegex(
             (ValueError, errors.InvalidArgumentError),
-            "range_max must be less than 2147483647"):
+            "range_max must be at most 1073741824"):
           result = candidate_sampling_ops.learned_unigram_candidate_sampler(
               true_classes=[[0]], num_true=1, num_sampled=1,
               unique=False, range_max=range_max)
@@ -165,7 +165,8 @@ class RangeSamplerOpsTest(test.TestCase):
       for range_max in (0, -1):
         with self.subTest(sampler=sampler.__name__, range_max=range_max):
           with self.assertRaisesRegex(
-              (ValueError, errors.InvalidArgumentError), "range_max"):
+              (ValueError, errors.InvalidArgumentError),
+              "must be at least minimum 1|range_max must be positive"):
             result = sampler(
                 true_classes=constant_op.constant([[0]], dtype=dtypes.int64),
                 num_true=1, num_sampled=1, unique=False, range_max=range_max)

From ad72e8d924b61ce84abc7f19636bfbbc2a060b2a Mon Sep 17 00:00:00 2001
From: Sai Ramineni <csramineni@gmail.com>
Date: Fri, 2 Oct 2026 02:39:17 -0400
Subject: [PATCH 0141/1000] Match both OpDef range validation message forms

---
 .../python/kernel_tests/random/candidate_sampler_ops_test.py   | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py b/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
index cb59e2bcddf..b8b31d02f8e 100644
--- a/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
+++ b/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
@@ -166,7 +166,8 @@ class RangeSamplerOpsTest(test.TestCase):
         with self.subTest(sampler=sampler.__name__, range_max=range_max):
           with self.assertRaisesRegex(
               (ValueError, errors.InvalidArgumentError),
-              "must be at least minimum 1|range_max must be positive"):
+              r"range_max.*(?:must be at least|less than) minimum 1|"
+              r"range_max must be positive"):
             result = sampler(
                 true_classes=constant_op.constant([[0]], dtype=dtypes.int64),
                 num_true=1, num_sampled=1, unique=False, range_max=range_max)

From 9cbbfc954c2189b0c591a23561721efc71867bd4 Mon Sep 17 00:00:00 2001
From: Sai Ramineni <csramineni@gmail.com>
Date: Fri, 2 Oct 2026 03:50:14 -0400
Subject: [PATCH 0142/1000] Use focused sampler test dependencies and exact
 validation messages

---
 tensorflow/python/kernel_tests/random/BUILD   |  2 +-
 .../random/candidate_sampler_ops_test.py      | 40 +++++++++----------
 tensorflow/python/ops/BUILD                   |  1 +
 3 files changed, 21 insertions(+), 22 deletions(-)

diff --git a/tensorflow/python/kernel_tests/random/BUILD b/tensorflow/python/kernel_tests/random/BUILD
index 41a8729f44a..236c5cdddf8 100644
--- a/tensorflow/python/kernel_tests/random/BUILD
+++ b/tensorflow/python/kernel_tests/random/BUILD
@@ -33,7 +33,6 @@ tf_py_strict_test(
     size = "small",
     srcs = ["candidate_sampler_ops_test.py"],
     deps = [
-        "//tensorflow:tensorflow_py_no_contrib",
         "//tensorflow/python/framework:constant_op",
         "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/framework:errors",
@@ -41,6 +40,7 @@ tf_py_strict_test(
         "//tensorflow/python/framework:test_lib",
         "//tensorflow/python/ops:array_ops",
         "//tensorflow/python/ops:candidate_sampling_ops",
+        "//tensorflow/python/ops:candidate_sampling_ops_gen",
         "//tensorflow/python/ops:math_ops",
         "//tensorflow/python/platform:client_testlib",
         "//third_party/py/numpy",
diff --git a/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py b/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
index b8b31d02f8e..6c9a7e582db 100644
--- a/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
+++ b/tensorflow/python/kernel_tests/random/candidate_sampler_ops_test.py
@@ -16,13 +16,13 @@
 
 import numpy as np
 
-from tensorflow import raw_ops
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import candidate_sampling_ops
+from tensorflow.python.ops import gen_candidate_sampling_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.platform import test
 
@@ -131,20 +131,18 @@ class RangeSamplerOpsTest(test.TestCase):
 
   @test_util.run_in_graph_and_eager_modes
   def testUnigramRangeTooLarge(self):
-    for sampler in (
-        raw_ops.ThreadUnsafeUnigramCandidateSampler,
-        raw_ops.LearnedUnigramCandidateSampler):
-      # The first range truncates to one, keeping the unfixed reproducer small.
-      for range_max in (2**32 + 1, 2**30 + 1, np.iinfo(np.int32).max,
-                        np.iinfo(np.int64).max):
-        with self.subTest(sampler=sampler.__name__, range_max=range_max):
-          with self.assertRaisesRegex(
-              (ValueError, errors.InvalidArgumentError),
-              "range_max must be at most 1073741824"):
-            result = sampler(
-                true_classes=constant_op.constant([[0]], dtype=dtypes.int64),
-                num_true=1, num_sampled=1, unique=False, range_max=range_max)
-            self.evaluate(result)
+    sampler = gen_candidate_sampling_ops.thread_unsafe_unigram_candidate_sampler
+    # The first range truncates to one, keeping the unfixed reproducer small.
+    for range_max in (2**32 + 1, 2**30 + 1, np.iinfo(np.int32).max,
+                      np.iinfo(np.int64).max):
+      with self.subTest(range_max=range_max):
+        with self.assertRaisesRegex(
+            (ValueError, errors.InvalidArgumentError),
+            "range_max must be at most 1073741824"):
+          result = sampler(
+              true_classes=constant_op.constant([[0]], dtype=dtypes.int64),
+              num_true=1, num_sampled=1, unique=False, range_max=range_max)
+          self.evaluate(result)
 
   @test_util.run_in_graph_and_eager_modes
   def testLearnedUnigramRangeTooLarge(self):
@@ -160,14 +158,14 @@ class RangeSamplerOpsTest(test.TestCase):
 
   @test_util.run_in_graph_and_eager_modes
   def testUnigramNonpositiveRange(self):
-    for sampler in (raw_ops.ThreadUnsafeUnigramCandidateSampler,
-                    raw_ops.LearnedUnigramCandidateSampler):
+    for sampler in (
+        gen_candidate_sampling_ops.thread_unsafe_unigram_candidate_sampler,
+        candidate_sampling_ops.learned_unigram_candidate_sampler):
       for range_max in (0, -1):
         with self.subTest(sampler=sampler.__name__, range_max=range_max):
           with self.assertRaisesRegex(
               (ValueError, errors.InvalidArgumentError),
-              r"range_max.*(?:must be at least|less than) minimum 1|"
-              r"range_max must be positive"):
+              r"must be at least minimum 1|less than minimum 1"):
             result = sampler(
                 true_classes=constant_op.constant([[0]], dtype=dtypes.int64),
                 num_true=1, num_sampled=1, unique=False, range_max=range_max)
@@ -176,8 +174,8 @@ class RangeSamplerOpsTest(test.TestCase):
   @test_util.run_in_graph_and_eager_modes
   def testUnigramSingleClass(self):
     for sampler in (
-        raw_ops.ThreadUnsafeUnigramCandidateSampler,
-        raw_ops.LearnedUnigramCandidateSampler):
+        gen_candidate_sampling_ops.thread_unsafe_unigram_candidate_sampler,
+        candidate_sampling_ops.learned_unigram_candidate_sampler):
       for unique in (False, True):
         with self.subTest(sampler=sampler.__name__, unique=unique):
           sampled, true_count, sampled_count = self.evaluate(sampler(
diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD
index 675b0563574..18d1db642e8 100644
--- a/tensorflow/python/ops/BUILD
+++ b/tensorflow/python/ops/BUILD
@@ -210,6 +210,7 @@ tf_gen_op_strict_wrapper_private_py(
     visibility = [
         "//learning/brain/python/ops:__pkg__",
         "//tensorflow/python:__pkg__",
+        "//tensorflow/python/kernel_tests/random:__pkg__",
     ],
 )
 

From 0bfcb259843de5230e14e053a07c4beb72c63be9 Mon Sep 17 00:00:00 2001
From: Mina Ibrahim <136262314+MinaIbrahim10@users.noreply.github.com>
Date: Fri, 2 Oct 2026 11:27:32 +0300
Subject: [PATCH 0143/1000] Restore compatibility and finalize Grappler review
 fixes

---
 tensorflow/core/grappler/clusters/cluster.h   |  1 -
 .../service/distributed_save_ft_test.py       | 17 +++++++-------
 .../service/server_lib_wrapper.cc             | 22 +++++++++----------
 tensorflow/python/grappler/BUILD              |  1 -
 tensorflow/python/grappler/cluster_wrapper.cc | 13 +++++------
 5 files changed, 25 insertions(+), 29 deletions(-)

diff --git a/tensorflow/core/grappler/clusters/cluster.h b/tensorflow/core/grappler/clusters/cluster.h
index 9e05214e299..2dcc3596075 100644
--- a/tensorflow/core/grappler/clusters/cluster.h
+++ b/tensorflow/core/grappler/clusters/cluster.h
@@ -142,7 +142,6 @@ class Cluster {
   const int timeout_s_;
   SessionOptions options_;
   RunOptions run_options_;
-
 };
 
 }  // end namespace grappler
diff --git a/tensorflow/python/data/experimental/kernel_tests/service/distributed_save_ft_test.py b/tensorflow/python/data/experimental/kernel_tests/service/distributed_save_ft_test.py
index ce987b350c1..2b91ceaaa90 100644
--- a/tensorflow/python/data/experimental/kernel_tests/service/distributed_save_ft_test.py
+++ b/tensorflow/python/data/experimental/kernel_tests/service/distributed_save_ft_test.py
@@ -134,7 +134,7 @@ class SnapshotFtTest(data_service_test_base.TestBase, parameterized.TestCase):
         cluster.dispatcher_address()))
 
     self._make_stream_dir(snapshot_dir.full_path, bad_stream_dir_name)
-    with self.assertRaisesRegex(errors.InternalError, "Can't parse"):
+    with self.assertRaisesRegex(RuntimeError, "Can't parse"):
       cluster.restart_dispatcher()
 
   @combinations.generate(
@@ -152,7 +152,7 @@ class SnapshotFtTest(data_service_test_base.TestBase, parameterized.TestCase):
 
     os.makedirs(os.path.join(self._splits_dir(snapshot_dir.full_path),
                              bad_source_dir_name))
-    with self.assertRaisesRegex(errors.InternalError, "Can't parse"):
+    with self.assertRaisesRegex(RuntimeError, "Can't parse"):
       cluster.restart_dispatcher()
 
   @combinations.generate(test_base.default_test_combinations())
@@ -166,7 +166,7 @@ class SnapshotFtTest(data_service_test_base.TestBase, parameterized.TestCase):
 
     os.makedirs(os.path.join(self._splits_dir(snapshot_dir.full_path),
                              "source_1"))
-    with self.assertRaisesRegex(errors.InternalError, "Found conflict"):
+    with self.assertRaisesRegex(RuntimeError, "Found conflict"):
       cluster.restart_dispatcher()
 
   @combinations.generate(
@@ -190,7 +190,7 @@ class SnapshotFtTest(data_service_test_base.TestBase, parameterized.TestCase):
     write_file(os.path.join(self._source_dir(snapshot_dir.full_path),
                             bad_split_filename))
     with self.assertRaisesRegex(
-        errors.InvalidArgumentError,
+        ValueError,
         "Expected split_<local_split_index>_<global_split_index>"):
       cluster.restart_dispatcher()
 
@@ -206,7 +206,7 @@ class SnapshotFtTest(data_service_test_base.TestBase, parameterized.TestCase):
     write_file(os.path.join(self._source_dir(snapshot_dir.full_path),
                             "split_1_0"))
     with self.assertRaisesRegex(
-        errors.InvalidArgumentError,
+        ValueError,
         "The local split index 1 exceeds the global split index 0"):
       cluster.restart_dispatcher()
 
@@ -221,7 +221,7 @@ class SnapshotFtTest(data_service_test_base.TestBase, parameterized.TestCase):
 
     write_file(os.path.join(self._source_dir(snapshot_dir.full_path),
                             "split_0_1"))
-    with self.assertRaisesRegex(errors.InternalError, "Found missing global"):
+    with self.assertRaisesRegex(RuntimeError, "Found missing global"):
       cluster.restart_dispatcher()
 
   @combinations.generate(test_base.default_test_combinations())
@@ -237,7 +237,7 @@ class SnapshotFtTest(data_service_test_base.TestBase, parameterized.TestCase):
         snapshot_dir.full_path, stream_idx=0), "split_0_1"))
     write_file(os.path.join(self._source_dir(
         snapshot_dir.full_path, stream_idx=1, worker=1), "split_0_1"))
-    with self.assertRaisesRegex(errors.InternalError, "Found duplicate global"):
+    with self.assertRaisesRegex(RuntimeError, "Found duplicate global"):
       cluster.restart_dispatcher()
 
   @combinations.generate(test_base.default_test_combinations())
@@ -253,8 +253,7 @@ class SnapshotFtTest(data_service_test_base.TestBase, parameterized.TestCase):
         self._source_dir(snapshot_dir.full_path, stream_idx=0), "split_0_1"))
     write_file(os.path.join(
         self._source_dir(snapshot_dir.full_path, stream_idx=1), "split_0_1"))
-    with self.assertRaisesRegex(
-        errors.InternalError, "worker is already assigned"):
+    with self.assertRaisesRegex(RuntimeError, "worker is already assigned"):
       cluster.restart_dispatcher()
 
   @combinations.generate(test_base.default_test_combinations())
diff --git a/tensorflow/python/data/experimental/service/server_lib_wrapper.cc b/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
index 2676c614320..f54423a7229 100644
--- a/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
+++ b/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
@@ -54,7 +54,7 @@ PYBIND11_MODULE(
           py::gil_scoped_release release;
           status = server->Start();
         }
-        tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
+        tensorflow::MaybeRaiseFromStatus(status);
       })
       .def("stop", [](tensorflow::data::DispatchGrpcDataServer* server) {
         py::gil_scoped_release release;
@@ -74,7 +74,7 @@ PYBIND11_MODULE(
                py::gil_scoped_release release;
                status = server->NumWorkers(&num_workers);
              }
-             tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
+             tensorflow::MaybeRaiseFromStatus(status);
              return num_workers;
            })
       .def("snapshot_streams",
@@ -87,7 +87,7 @@ PYBIND11_MODULE(
                py::gil_scoped_release release;
                status = server->SnapshotStreams(path, &streams);
              }
-             tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
+             tensorflow::MaybeRaiseFromStatus(status);
              return streams;
            });
 
@@ -99,7 +99,7 @@ PYBIND11_MODULE(
           py::gil_scoped_release release;
           status = server->Start();
         }
-        tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
+        tensorflow::MaybeRaiseFromStatus(status);
       })
       .def("stop", [](tensorflow::data::WorkerGrpcDataServer* server) {
         py::gil_scoped_release release;
@@ -119,7 +119,7 @@ PYBIND11_MODULE(
                py::gil_scoped_release release;
                status = server->NumTasks(&num_tasks);
              }
-             tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
+             tensorflow::MaybeRaiseFromStatus(status);
              return num_tasks;
            })
       .def("snapshot_task_progresses",
@@ -133,7 +133,7 @@ PYBIND11_MODULE(
                status =
                    server->SnapshotTaskProgresses(&snapshot_task_progresses);
              }
-             tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
+             tensorflow::MaybeRaiseFromStatus(status);
              return snapshot_task_progresses;
            });
 
@@ -143,13 +143,13 @@ PYBIND11_MODULE(
           -> std::unique_ptr<tensorflow::data::DispatchGrpcDataServer> {
         tensorflow::data::experimental::DispatcherConfig config;
         if (!config.ParseFromString(serialized_dispatcher_config)) {
-          tsl::MaybeRaiseRegisteredFromStatusWithGIL(absl::InvalidArgumentError(
+          tensorflow::MaybeRaiseFromStatus(absl::InvalidArgumentError(
               "Failed to deserialize dispatcher config."));
         }
         std::unique_ptr<tensorflow::data::DispatchGrpcDataServer> server;
         absl::Status status =
             tensorflow::data::NewDispatchServer(config, server);
-        tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
+        tensorflow::MaybeRaiseFromStatus(status);
         return server;
       },
       py::return_value_policy::reference);
@@ -160,13 +160,13 @@ PYBIND11_MODULE(
           -> std::unique_ptr<tensorflow::data::WorkerGrpcDataServer> {
         tensorflow::data::experimental::WorkerConfig config;
         if (!config.ParseFromString(serialized_worker_config)) {
-          tsl::MaybeRaiseRegisteredFromStatusWithGIL(absl::InvalidArgumentError(
+          tensorflow::MaybeRaiseFromStatus(absl::InvalidArgumentError(
               "Failed to deserialize worker config."));
         }
         std::unique_ptr<tensorflow::data::WorkerGrpcDataServer> server;
         absl::Status status =
             tensorflow::data::NewWorkerServer(config, server);
-        tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
+        tensorflow::MaybeRaiseFromStatus(status);
         return server;
       },
       py::return_value_policy::reference);
@@ -189,7 +189,7 @@ PYBIND11_MODULE(
                          " from dispatcher at ", address),
             deadline_micros);
         Py_END_ALLOW_THREADS;
-        tsl::MaybeRaiseRegisteredFromStatusWithGIL(status);
+        tensorflow::MaybeRaiseFromStatus(status);
         return metadata;
       },
       py::return_value_policy::reference);
diff --git a/tensorflow/python/grappler/BUILD b/tensorflow/python/grappler/BUILD
index 897d4e42fb3..e66b00aee38 100644
--- a/tensorflow/python/grappler/BUILD
+++ b/tensorflow/python/grappler/BUILD
@@ -293,7 +293,6 @@ tf_python_pybind_extension(
         "//tensorflow/core/framework:allocator",
         "//tensorflow/core/framework:tensor",
         "//tensorflow/core/grappler:utils",
-        "//tensorflow/core/platform:errors",
         "//tensorflow/core/platform:mutex",
         "//tensorflow/core/platform:status",
         "//tensorflow/python/lib/core:pybind11_status",
diff --git a/tensorflow/python/grappler/cluster_wrapper.cc b/tensorflow/python/grappler/cluster_wrapper.cc
index 937cc17f84d..bb1ea777987 100644
--- a/tensorflow/python/grappler/cluster_wrapper.cc
+++ b/tensorflow/python/grappler/cluster_wrapper.cc
@@ -48,7 +48,6 @@ limitations under the License.
 #include "tensorflow/core/grappler/devices.h"
 #include "tensorflow/core/grappler/grappler_item.h"
 #include "tensorflow/core/grappler/utils.h"
-#include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/protobuf/config.pb.h"
@@ -57,7 +56,6 @@ limitations under the License.
 
 namespace py = pybind11;
 
-
 namespace {
 
 // SingleMachine uses process-global provisioning state. Serialize Python
@@ -77,7 +75,7 @@ absl::Status _GetOpPerformanceDataAndRunTime(
   if (!status.ok()) return status;
 
   tensorflow::RunMetadata run_metadata;
-  TF_RETURN_IF_ERROR(
+  tsl::MaybeRaiseRegisteredFromStatusWithGIL(
       cost_measure->PredictCosts(item.graph, &run_metadata, costs));
 
   if (op_performance_data) {
@@ -307,7 +305,6 @@ PYBIND11_MODULE(
                 "You need both a cluster and an item to measure costs."));
           }
 
-
           const int num_measurements = cluster->type() == "virtual" ? 1 : 10;
           tensorflow::grappler::MeasuringCostEstimator cost_measure(
               cluster, num_measurements, 0);
@@ -316,9 +313,11 @@ PYBIND11_MODULE(
           tensorflow::grappler::Costs costs;
           absl::Status s = _GetOpPerformanceDataAndRunTime(
               *item, &cost_measure, &op_performance_data, &costs);
-          tsl::MaybeRaiseRegisteredFromStatusWithGIL(s);
-          double run_time =
-              static_cast<double>(costs.execution_time.count()) / 1e9;
+          double run_time = FLT_MAX;
+          if (s.ok()) {
+            run_time =
+                static_cast<double>(costs.execution_time.count()) / 1e9;
+          }
           tensorflow::StepStats step_stats;
           if (generate_timeline) {
             tensorflow::RunMetadata metadata;

From 8d19e081d31c12d653d38ebaafd160463c16996f Mon Sep 17 00:00:00 2001
From: Mina Ibrahim <136262314+MinaIbrahim10@users.noreply.github.com>
Date: Fri, 2 Oct 2026 11:49:48 +0300
Subject: [PATCH 0144/1000] Release GIL around blocking Grappler operations

---
 tensorflow/core/data/service/server_lib.cc    | 20 +++---
 tensorflow/python/grappler/cluster_wrapper.cc | 62 ++++++++++++++-----
 2 files changed, 58 insertions(+), 24 deletions(-)

diff --git a/tensorflow/core/data/service/server_lib.cc b/tensorflow/core/data/service/server_lib.cc
index 41aaba4c7fd..458e78b0470 100644
--- a/tensorflow/core/data/service/server_lib.cc
+++ b/tensorflow/core/data/service/server_lib.cc
@@ -92,17 +92,23 @@ absl::Status GrpcDataServerBase::Start() {
 }
 
 void GrpcDataServerBase::Stop() {
-  mutex_lock lock(lifecycle_mu_);
-  if (stopped_) {
-    return;
+  ::grpc::Server* server = nullptr;
+  int port = 0;
+  {
+    mutex_lock lock(lifecycle_mu_);
+    if (stopped_) {
+      return;
+    }
+    stopped_ = true;
+    server = server_.get();
+    port = bound_port_;
   }
-  if (server_) {
+  if (server) {
     StopServiceInternal();
-    server_->Shutdown();
+    server->Shutdown();
     LOG(INFO) << "Shut down " << server_type_ << " server running at port "
-              << bound_port_;
+              << port;
   }
-  stopped_ = true;
 }
 
 void GrpcDataServerBase::Join() {
diff --git a/tensorflow/python/grappler/cluster_wrapper.cc b/tensorflow/python/grappler/cluster_wrapper.cc
index bb1ea777987..17029795307 100644
--- a/tensorflow/python/grappler/cluster_wrapper.cc
+++ b/tensorflow/python/grappler/cluster_wrapper.cc
@@ -101,6 +101,7 @@ PYBIND11_MODULE(
           std::unique_ptr<tensorflow::grappler::Cluster> cluster;
           absl::Status provision_status;
           {
+            py::gil_scoped_release release;
             tensorflow::mutex_lock lifecycle_lock(GetClusterLifecycleMutex());
             int num_cpu_cores =
                 tensorflow::grappler::GetNumAvailableLogicalCPUCores();
@@ -137,7 +138,12 @@ PYBIND11_MODULE(
           }
           std::unique_ptr<tensorflow::grappler::Cluster> cluster =
               std::make_unique<tensorflow::grappler::VirtualCluster>(devices);
-          tsl::MaybeRaiseRegisteredFromStatusWithGIL(cluster->Provision());
+          absl::Status provision_status;
+          {
+            py::gil_scoped_release release;
+            provision_status = cluster->Provision();
+          }
+          tsl::MaybeRaiseRegisteredFromStatusWithGIL(provision_status);
           return cluster.release();
         });
 
@@ -145,7 +151,9 @@ PYBIND11_MODULE(
     if (cluster == nullptr) {
       return;
     }
-    if (cluster->type() == "single_machine") {
+    const bool is_single_machine = cluster->type() == "single_machine";
+    py::gil_scoped_release release;
+    if (is_single_machine) {
       tensorflow::mutex_lock lifecycle_lock(GetClusterLifecycleMutex());
       (void)cluster->Shutdown();
     } else {
@@ -160,17 +168,28 @@ PYBIND11_MODULE(
                 absl::InvalidArgumentError("Cluster cannot be None."));
           }
 
-          std::vector<py::bytes> named_devices;
+          std::vector<std::pair<std::string, tensorflow::DeviceProperties>>
+              devices;
           {
-            const auto& devices = cluster->GetDevices();
-
-            for (const auto& dev : devices) {
-              tensorflow::NamedDevice d;
-              d.set_name(dev.first);
-              *d.mutable_properties() = dev.second;
-              named_devices.push_back(d.SerializeAsString());
+            const auto& cluster_devices = cluster->GetDevices();
+            devices.reserve(cluster_devices.size());
+            for (const auto& dev : cluster_devices) {
+              devices.emplace_back(dev.first, dev.second);
             }
           }
+          std::sort(devices.begin(), devices.end(),
+                    [](const auto& lhs, const auto& rhs) {
+                      return lhs.first < rhs.first;
+                    });
+
+          std::vector<py::bytes> named_devices;
+          named_devices.reserve(devices.size());
+          for (const auto& dev : devices) {
+            tensorflow::NamedDevice d;
+            d.set_name(dev.first);
+            *d.mutable_properties() = dev.second;
+            named_devices.push_back(d.SerializeAsString());
+          }
           return named_devices;
         });
 
@@ -321,8 +340,13 @@ PYBIND11_MODULE(
           tensorflow::StepStats step_stats;
           if (generate_timeline) {
             tensorflow::RunMetadata metadata;
-            tsl::MaybeRaiseRegisteredFromStatusWithGIL(
-                cluster->Run(item->graph, item->feed, item->fetch, &metadata));
+            absl::Status run_status;
+            {
+              py::gil_scoped_release release;
+              run_status =
+                  cluster->Run(item->graph, item->feed, item->fetch, &metadata);
+            }
+            tsl::MaybeRaiseRegisteredFromStatusWithGIL(run_status);
             step_stats = metadata.step_stats();
           }
 
@@ -355,12 +379,16 @@ PYBIND11_MODULE(
 
         tensorflow::grappler::GraphMemory memory(*item);
 
-        if (cluster->DetailedStatsEnabled()) {
-          tsl::MaybeRaiseRegisteredFromStatusWithGIL(memory.InferDynamically(cluster));
-        } else {
-          tsl::MaybeRaiseRegisteredFromStatusWithGIL(
-              memory.InferStatically(cluster->GetDevices()));
+        absl::Status inference_status;
+        {
+          py::gil_scoped_release release;
+          if (cluster->DetailedStatsEnabled()) {
+            inference_status = memory.InferDynamically(cluster);
+          } else {
+            inference_status = memory.InferStatically(cluster->GetDevices());
+          }
         }
+        tsl::MaybeRaiseRegisteredFromStatusWithGIL(inference_status);
 
         std::unordered_map<std::string,
                            std::tuple<int64_t, std::vector<MemoryUsage>>>

From 1daa61a75ca85be0233a6c1b3838e191ba72fe9c Mon Sep 17 00:00:00 2001
From: Mina Ibrahim <136262314+MinaIbrahim10@users.noreply.github.com>
Date: Fri, 2 Oct 2026 12:15:07 +0300
Subject: [PATCH 0145/1000] Release GIL in bound port bindings

---
 .../data/experimental/service/server_lib_wrapper.cc  | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/tensorflow/python/data/experimental/service/server_lib_wrapper.cc b/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
index f54423a7229..177a436cc17 100644
--- a/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
+++ b/tensorflow/python/data/experimental/service/server_lib_wrapper.cc
@@ -62,10 +62,8 @@ PYBIND11_MODULE(
       })
       .def("join", &tensorflow::data::DispatchGrpcDataServer::Join,
            py::call_guard<py::gil_scoped_release>())
-      .def("bound_port",
-           [](tensorflow::data::DispatchGrpcDataServer* server) {
-             return server->BoundPort();
-           })
+      .def("bound_port", &tensorflow::data::DispatchGrpcDataServer::BoundPort,
+           py::call_guard<py::gil_scoped_release>())
       .def("num_workers",
            [](tensorflow::data::DispatchGrpcDataServer* server) -> int {
              int num_workers;
@@ -107,10 +105,8 @@ PYBIND11_MODULE(
       })
       .def("join", &tensorflow::data::WorkerGrpcDataServer::Join,
            py::call_guard<py::gil_scoped_release>())
-      .def("bound_port",
-           [](tensorflow::data::WorkerGrpcDataServer* server) {
-             return server->BoundPort();
-           })
+      .def("bound_port", &tensorflow::data::WorkerGrpcDataServer::BoundPort,
+           py::call_guard<py::gil_scoped_release>())
       .def("num_tasks",
            [](tensorflow::data::WorkerGrpcDataServer* server) -> int {
              int num_tasks;

From 2087dbe422aba23ee508a534699788b26d4a8339 Mon Sep 17 00:00:00 2001
From: Mina Ibrahim <136262314+MinaIbrahim10@users.noreply.github.com>
Date: Fri, 2 Oct 2026 12:30:11 +0300
Subject: [PATCH 0146/1000] Destroy failed SingleMachine under lifecycle lock

---
 tensorflow/python/grappler/cluster_wrapper.cc | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensorflow/python/grappler/cluster_wrapper.cc b/tensorflow/python/grappler/cluster_wrapper.cc
index 17029795307..a8c76e3f59c 100644
--- a/tensorflow/python/grappler/cluster_wrapper.cc
+++ b/tensorflow/python/grappler/cluster_wrapper.cc
@@ -113,6 +113,9 @@ PYBIND11_MODULE(
             cluster->AllowSoftPlacement(allow_soft_placement);
             cluster->SetNumWarmupSteps(10);
             provision_status = cluster->Provision();
+            if (!provision_status.ok()) {
+              cluster.reset();
+            }
           }
           tsl::MaybeRaiseRegisteredFromStatusWithGIL(provision_status);
           return cluster.release();

From fbfd1c26c143bc99041271d034a47f11bc4d3b25 Mon Sep 17 00:00:00 2001
From: Mina Ibrahim <136262314+MinaIbrahim10@users.noreply.github.com>
Date: Fri, 2 Oct 2026 12:46:17 +0300
Subject: [PATCH 0147/1000] Add direct Abseil status include

---
 tensorflow/python/grappler/tf_optimizer_wrapper.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/python/grappler/tf_optimizer_wrapper.cc b/tensorflow/python/grappler/tf_optimizer_wrapper.cc
index a7a8f19b814..fe4b295c3b2 100644
--- a/tensorflow/python/grappler/tf_optimizer_wrapper.cc
+++ b/tensorflow/python/grappler/tf_optimizer_wrapper.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/status/status.h"
 #include "pybind11/pybind11.h"  // from @pybind11
 #include "pybind11_protobuf/native_proto_caster.h"  // from @pybind11_protobuf
 #include "tensorflow/core/common_runtime/device.h"

From 2a687f1b41ae5b7f3d7f2557770a5c1aad3e1685 Mon Sep 17 00:00:00 2001
From: Mina Ibrahim <136262314+MinaIbrahim10@users.noreply.github.com>
Date: Fri, 2 Oct 2026 13:13:45 +0300
Subject: [PATCH 0148/1000] Release GIL during Grappler graph analysis

---
 tensorflow/python/grappler/cluster_wrapper.cc | 162 +++++++++---------
 1 file changed, 85 insertions(+), 77 deletions(-)

diff --git a/tensorflow/python/grappler/cluster_wrapper.cc b/tensorflow/python/grappler/cluster_wrapper.cc
index a8c76e3f59c..5f91c45bf8c 100644
--- a/tensorflow/python/grappler/cluster_wrapper.cc
+++ b/tensorflow/python/grappler/cluster_wrapper.cc
@@ -18,7 +18,6 @@ limitations under the License.
 #include <cstddef>
 #include <cstdint>
 #include <iterator>
-#include <map>
 #include <memory>
 #include <set>
 #include <stdexcept>
@@ -117,7 +116,7 @@ PYBIND11_MODULE(
               cluster.reset();
             }
           }
-          tsl::MaybeRaiseRegisteredFromStatusWithGIL(provision_status);
+          tsl::MaybeRaiseRegisteredFromStatus(provision_status);
           return cluster.release();
         });
 
@@ -146,7 +145,7 @@ PYBIND11_MODULE(
             py::gil_scoped_release release;
             provision_status = cluster->Provision();
           }
-          tsl::MaybeRaiseRegisteredFromStatusWithGIL(provision_status);
+          tsl::MaybeRaiseRegisteredFromStatus(provision_status);
           return cluster.release();
         });
 
@@ -167,7 +166,7 @@ PYBIND11_MODULE(
   m.def("TF_ListDevices",
         [](tensorflow::grappler::Cluster* cluster) -> std::vector<py::bytes> {
           if (cluster == nullptr) {
-            tsl::MaybeRaiseRegisteredFromStatusWithGIL(
+            tsl::MaybeRaiseRegisteredFromStatus(
                 absl::InvalidArgumentError("Cluster cannot be None."));
           }
 
@@ -213,95 +212,100 @@ PYBIND11_MODULE(
       "TF_GetSupportedDevices",
       [](tensorflow::grappler::Cluster* cluster,
          tensorflow::grappler::GrapplerItem* item)
-          -> std::map<std::string, std::vector<std::string>> {
+          -> std::unordered_map<std::string, std::vector<std::string>> {
         if (cluster == nullptr || item == nullptr) {
-          tsl::MaybeRaiseRegisteredFromStatusWithGIL(absl::Status(
+          tsl::MaybeRaiseRegisteredFromStatus(absl::Status(
               absl::InternalError("You need both a cluster and an "
                                   "item to get supported devices.")));
         }
-        std::unordered_map<std::string, std::set<std::string>> device_types;
-        std::string cluster_type;
 
+        std::unordered_map<std::string, std::vector<std::string>> result;
         {
-          const auto& devices = cluster->GetDevices();
-          cluster_type = cluster->type();
+          py::gil_scoped_release release;
+          std::unordered_map<std::string, std::set<std::string>> device_types;
+          std::string cluster_type;
 
-          for (const auto& dev : devices) {
-            device_types[dev.second.type()].insert(dev.first);
-          }
-        }
+          {
+            const auto& devices = cluster->GetDevices();
+            cluster_type = cluster->type();
 
-        std::map<std::string, std::set<std::string>> supported_device_types;
-        std::unordered_map<std::string, std::set<std::string>>
-            device_restrictions;
-
-        if (cluster_type != "single_machine") {
-          for (const auto& node : item->graph.node()) {
-            auto& supported = supported_device_types[node.name()];
-            for (const auto& dev : device_types) {
-              // The actual kernel may not be linked in this binary.
-              supported.insert(dev.first);
+            for (const auto& dev : devices) {
+              device_types[dev.second.type()].insert(dev.first);
             }
           }
-        } else {
-          for (const auto& node : item->graph.node()) {
-            auto& supported = supported_device_types[node.name()];
+
+          std::unordered_map<std::string, std::set<std::string>>
+              supported_device_types;
+          std::unordered_map<std::string, std::set<std::string>>
+              device_restrictions;
+
+          if (cluster_type != "single_machine") {
+            std::set<std::string> all_types;
             for (const auto& dev : device_types) {
-              const std::string& type = dev.first;
+              all_types.insert(dev.first);
+            }
+            for (const auto& node : item->graph.node()) {
+              supported_device_types[node.name()] = all_types;
+            }
+          } else {
+            for (const auto& node : item->graph.node()) {
+              auto& supported = supported_device_types[node.name()];
+              for (const auto& dev : device_types) {
+                const std::string& type = dev.first;
 
-              // Check the kernel capabilities.
-              const tensorflow::DeviceType dev_type(type);
-              absl::Status s =
-                  tensorflow::FindKernelDef(dev_type, node, nullptr, nullptr);
-              if (s.ok()) {
-                supported.insert(type);
-
-                // Check which inputs are restricted to reside on the host.
-                // TODO: extends this to support outputs as well.
-                tensorflow::MemoryTypeVector inp_mtypes;
-                tensorflow::MemoryTypeVector out_mtypes;
-                absl::Status s = tensorflow::MemoryTypesForNode(
-                    tensorflow::OpRegistry::Global(), dev_type, node,
-                    &inp_mtypes, &out_mtypes);
+                // Check the kernel capabilities.
+                const tensorflow::DeviceType dev_type(type);
+                absl::Status s =
+                    tensorflow::FindKernelDef(dev_type, node, nullptr, nullptr);
                 if (s.ok()) {
-                  for (size_t i = 0; i < inp_mtypes.size(); ++i) {
-                    if (inp_mtypes[i] == tensorflow::HOST_MEMORY) {
-                      device_restrictions[tensorflow::grappler::NodeName(
-                                              node.input(i))]
-                          .insert("CPU");
-                      break;
+                  supported.insert(type);
+
+                  // Check which inputs are restricted to reside on the host.
+                  // TODO: extends this to support outputs as well.
+                  tensorflow::MemoryTypeVector inp_mtypes;
+                  tensorflow::MemoryTypeVector out_mtypes;
+                  absl::Status s = tensorflow::MemoryTypesForNode(
+                      tensorflow::OpRegistry::Global(), dev_type, node,
+                      &inp_mtypes, &out_mtypes);
+                  if (s.ok()) {
+                    for (size_t i = 0; i < inp_mtypes.size(); ++i) {
+                      if (inp_mtypes[i] == tensorflow::HOST_MEMORY) {
+                        device_restrictions[tensorflow::grappler::NodeName(
+                                                node.input(i))]
+                            .insert("CPU");
+                        break;
+                      }
                     }
                   }
                 }
               }
             }
           }
-        }
 
-        std::map<std::string, std::vector<std::string>> result;
-        for (const auto& supported_dev : supported_device_types) {
-          const std::string& node = supported_dev.first;
-          std::set<std::string> feasible;
-          const auto it = device_restrictions.find(node);
-          if (it != device_restrictions.end()) {
-            const std::set<std::string>& candidates = supported_dev.second;
-            const std::set<std::string>& valid = it->second;
-            std::set_intersection(candidates.begin(), candidates.end(),
-                                  valid.begin(), valid.end(),
-                                  std::inserter(feasible, feasible.begin()));
-          } else {
-            feasible = supported_dev.second;
-          }
-
-          std::vector<std::string> device_names;
-          for (const std::string& type : feasible) {
-            auto it = device_types.find(type);
-            DCHECK(it != device_types.end());
-            for (const std::string& name : it->second) {
-              device_names.push_back(name);
+          for (const auto& supported_dev : supported_device_types) {
+            const std::string& node = supported_dev.first;
+            std::set<std::string> feasible;
+            const auto it = device_restrictions.find(node);
+            if (it != device_restrictions.end()) {
+              const std::set<std::string>& candidates = supported_dev.second;
+              const std::set<std::string>& valid = it->second;
+              std::set_intersection(candidates.begin(), candidates.end(),
+                                    valid.begin(), valid.end(),
+                                    std::inserter(feasible, feasible.begin()));
+            } else {
+              feasible = supported_dev.second;
             }
+
+            std::vector<std::string> device_names;
+            for (const std::string& type : feasible) {
+              auto it = device_types.find(type);
+              DCHECK(it != device_types.end());
+              for (const std::string& name : it->second) {
+                device_names.push_back(name);
+              }
+            }
+            result[node] = device_names;
           }
-          result[node] = device_names;
         }
         return result;
       });
@@ -323,7 +327,7 @@ PYBIND11_MODULE(
            tensorflow::grappler::Cluster* cluster, bool generate_timeline)
             -> std::tuple<std::vector<py::bytes>, double, py::bytes> {
           if (cluster == nullptr || item == nullptr) {
-            tsl::MaybeRaiseRegisteredFromStatusWithGIL(absl::InvalidArgumentError(
+            tsl::MaybeRaiseRegisteredFromStatus(absl::InvalidArgumentError(
                 "You need both a cluster and an item to measure costs."));
           }
 
@@ -333,8 +337,12 @@ PYBIND11_MODULE(
 
           tensorflow::OpPerformanceList op_performance_data;
           tensorflow::grappler::Costs costs;
-          absl::Status s = _GetOpPerformanceDataAndRunTime(
-              *item, &cost_measure, &op_performance_data, &costs);
+          absl::Status s;
+          {
+            py::gil_scoped_release release;
+            s = _GetOpPerformanceDataAndRunTime(
+                *item, &cost_measure, &op_performance_data, &costs);
+          }
           double run_time = FLT_MAX;
           if (s.ok()) {
             run_time =
@@ -349,7 +357,7 @@ PYBIND11_MODULE(
               run_status =
                   cluster->Run(item->graph, item->feed, item->fetch, &metadata);
             }
-            tsl::MaybeRaiseRegisteredFromStatusWithGIL(run_status);
+            tsl::MaybeRaiseRegisteredFromStatus(run_status);
             step_stats = metadata.step_stats();
           }
 
@@ -375,7 +383,7 @@ PYBIND11_MODULE(
           -> std::unordered_map<std::string,
                                 std::tuple<int64_t, std::vector<MemoryUsage>>> {
         if (cluster == nullptr || item == nullptr) {
-          tsl::MaybeRaiseRegisteredFromStatusWithGIL(absl::Status(absl::InternalError(
+          tsl::MaybeRaiseRegisteredFromStatus(absl::Status(absl::InternalError(
               "You need both a cluster and an item to determine peak "
               "memory usage.")));
         }
@@ -391,7 +399,7 @@ PYBIND11_MODULE(
             inference_status = memory.InferStatically(cluster->GetDevices());
           }
         }
-        tsl::MaybeRaiseRegisteredFromStatusWithGIL(inference_status);
+        tsl::MaybeRaiseRegisteredFromStatus(inference_status);
 
         std::unordered_map<std::string,
                            std::tuple<int64_t, std::vector<MemoryUsage>>>

From c20867a43730f6c7c82ee20851d747f1405ffb36 Mon Sep 17 00:00:00 2001
From: Mina Ibrahim <136262314+MinaIbrahim10@users.noreply.github.com>
Date: Fri, 2 Oct 2026 13:32:06 +0300
Subject: [PATCH 0149/1000] Avoid copying DeviceProperties in device listing

---
 tensorflow/python/grappler/cluster_wrapper.cc | 28 ++++++++-----------
 1 file changed, 12 insertions(+), 16 deletions(-)

diff --git a/tensorflow/python/grappler/cluster_wrapper.cc b/tensorflow/python/grappler/cluster_wrapper.cc
index 5f91c45bf8c..533359836ea 100644
--- a/tensorflow/python/grappler/cluster_wrapper.cc
+++ b/tensorflow/python/grappler/cluster_wrapper.cc
@@ -170,26 +170,22 @@ PYBIND11_MODULE(
                 absl::InvalidArgumentError("Cluster cannot be None."));
           }
 
-          std::vector<std::pair<std::string, tensorflow::DeviceProperties>>
-              devices;
-          {
-            const auto& cluster_devices = cluster->GetDevices();
-            devices.reserve(cluster_devices.size());
-            for (const auto& dev : cluster_devices) {
-              devices.emplace_back(dev.first, dev.second);
-            }
+          const auto& cluster_devices = cluster->GetDevices();
+          std::vector<std::string> device_names;
+          device_names.reserve(cluster_devices.size());
+          for (const auto& dev : cluster_devices) {
+            device_names.push_back(dev.first);
           }
-          std::sort(devices.begin(), devices.end(),
-                    [](const auto& lhs, const auto& rhs) {
-                      return lhs.first < rhs.first;
-                    });
+          std::sort(device_names.begin(), device_names.end());
 
           std::vector<py::bytes> named_devices;
-          named_devices.reserve(devices.size());
-          for (const auto& dev : devices) {
+          named_devices.reserve(device_names.size());
+          for (const std::string& name : device_names) {
+            const auto it = cluster_devices.find(name);
+            DCHECK(it != cluster_devices.end());
             tensorflow::NamedDevice d;
-            d.set_name(dev.first);
-            *d.mutable_properties() = dev.second;
+            d.set_name(name);
+            *d.mutable_properties() = it->second;
             named_devices.push_back(d.SerializeAsString());
           }
           return named_devices;

From bf6c1ae147ca366467cef9bba887c4f340b0271f Mon Sep 17 00:00:00 2001
From: Mina Ibrahim <136262314+MinaIbrahim10@users.noreply.github.com>
Date: Fri, 2 Oct 2026 14:06:09 +0300
Subject: [PATCH 0150/1000] Sort device entries without key copies

---
 tensorflow/python/grappler/cluster_wrapper.cc | 23 +++++++++++--------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/tensorflow/python/grappler/cluster_wrapper.cc b/tensorflow/python/grappler/cluster_wrapper.cc
index 533359836ea..7e8efaaa584 100644
--- a/tensorflow/python/grappler/cluster_wrapper.cc
+++ b/tensorflow/python/grappler/cluster_wrapper.cc
@@ -171,21 +171,24 @@ PYBIND11_MODULE(
           }
 
           const auto& cluster_devices = cluster->GetDevices();
-          std::vector<std::string> device_names;
-          device_names.reserve(cluster_devices.size());
+          std::vector<const std::pair<
+              const std::string, tensorflow::DeviceProperties>*>
+              sorted_devices;
+          sorted_devices.reserve(cluster_devices.size());
           for (const auto& dev : cluster_devices) {
-            device_names.push_back(dev.first);
+            sorted_devices.push_back(&dev);
           }
-          std::sort(device_names.begin(), device_names.end());
+          std::sort(sorted_devices.begin(), sorted_devices.end(),
+                    [](const auto* a, const auto* b) {
+                      return a->first < b->first;
+                    });
 
           std::vector<py::bytes> named_devices;
-          named_devices.reserve(device_names.size());
-          for (const std::string& name : device_names) {
-            const auto it = cluster_devices.find(name);
-            DCHECK(it != cluster_devices.end());
+          named_devices.reserve(sorted_devices.size());
+          for (const auto* dev : sorted_devices) {
             tensorflow::NamedDevice d;
-            d.set_name(name);
-            *d.mutable_properties() = it->second;
+            d.set_name(dev->first);
+            *d.mutable_properties() = dev->second;
             named_devices.push_back(d.SerializeAsString());
           }
           return named_devices;

From 94ebb4ea9c7c9afd541318fa8df88cdb3f71b4b8 Mon Sep 17 00:00:00 2001
From: Mina Ibrahim <136262314+MinaIbrahim10@users.noreply.github.com>
Date: Fri, 2 Oct 2026 14:25:35 +0300
Subject: [PATCH 0151/1000] Avoid redundant device set copies

---
 tensorflow/python/grappler/cluster_wrapper.cc | 70 +++++++++++--------
 1 file changed, 39 insertions(+), 31 deletions(-)

diff --git a/tensorflow/python/grappler/cluster_wrapper.cc b/tensorflow/python/grappler/cluster_wrapper.cc
index 7e8efaaa584..b59c852227a 100644
--- a/tensorflow/python/grappler/cluster_wrapper.cc
+++ b/tensorflow/python/grappler/cluster_wrapper.cc
@@ -233,20 +233,22 @@ PYBIND11_MODULE(
             }
           }
 
-          std::unordered_map<std::string, std::set<std::string>>
-              supported_device_types;
-          std::unordered_map<std::string, std::set<std::string>>
-              device_restrictions;
-
           if (cluster_type != "single_machine") {
-            std::set<std::string> all_types;
+            std::vector<std::string> all_device_names;
             for (const auto& dev : device_types) {
-              all_types.insert(dev.first);
+              for (const std::string& name : dev.second) {
+                all_device_names.push_back(name);
+              }
             }
             for (const auto& node : item->graph.node()) {
-              supported_device_types[node.name()] = all_types;
+              result[node.name()] = all_device_names;
             }
           } else {
+            std::unordered_map<std::string, std::set<std::string>>
+                supported_device_types;
+            std::unordered_map<std::string, std::set<std::string>>
+                device_restrictions;
+
             for (const auto& node : item->graph.node()) {
               auto& supported = supported_device_types[node.name()];
               for (const auto& dev : device_types) {
@@ -263,10 +265,10 @@ PYBIND11_MODULE(
                   // TODO: extends this to support outputs as well.
                   tensorflow::MemoryTypeVector inp_mtypes;
                   tensorflow::MemoryTypeVector out_mtypes;
-                  absl::Status s = tensorflow::MemoryTypesForNode(
+                  absl::Status s2 = tensorflow::MemoryTypesForNode(
                       tensorflow::OpRegistry::Global(), dev_type, node,
                       &inp_mtypes, &out_mtypes);
-                  if (s.ok()) {
+                  if (s2.ok()) {
                     for (size_t i = 0; i < inp_mtypes.size(); ++i) {
                       if (inp_mtypes[i] == tensorflow::HOST_MEMORY) {
                         device_restrictions[tensorflow::grappler::NodeName(
@@ -279,31 +281,37 @@ PYBIND11_MODULE(
                 }
               }
             }
-          }
 
-          for (const auto& supported_dev : supported_device_types) {
-            const std::string& node = supported_dev.first;
-            std::set<std::string> feasible;
-            const auto it = device_restrictions.find(node);
-            if (it != device_restrictions.end()) {
-              const std::set<std::string>& candidates = supported_dev.second;
-              const std::set<std::string>& valid = it->second;
-              std::set_intersection(candidates.begin(), candidates.end(),
-                                    valid.begin(), valid.end(),
-                                    std::inserter(feasible, feasible.begin()));
-            } else {
-              feasible = supported_dev.second;
-            }
+            for (const auto& supported_dev : supported_device_types) {
+              const std::string& node = supported_dev.first;
+              const auto it = device_restrictions.find(node);
+              std::vector<std::string> device_names;
 
-            std::vector<std::string> device_names;
-            for (const std::string& type : feasible) {
-              auto it = device_types.find(type);
-              DCHECK(it != device_types.end());
-              for (const std::string& name : it->second) {
-                device_names.push_back(name);
+              if (it != device_restrictions.end()) {
+                std::set<std::string> feasible;
+                const std::set<std::string>& candidates = supported_dev.second;
+                const std::set<std::string>& valid = it->second;
+                std::set_intersection(candidates.begin(), candidates.end(),
+                                      valid.begin(), valid.end(),
+                                      std::inserter(feasible, feasible.begin()));
+                for (const std::string& type : feasible) {
+                  auto type_it = device_types.find(type);
+                  DCHECK(type_it != device_types.end());
+                  for (const std::string& name : type_it->second) {
+                    device_names.push_back(name);
+                  }
+                }
+              } else {
+                for (const std::string& type : supported_dev.second) {
+                  auto type_it = device_types.find(type);
+                  DCHECK(type_it != device_types.end());
+                  for (const std::string& name : type_it->second) {
+                    device_names.push_back(name);
+                  }
+                }
               }
+              result[node] = device_names;
             }
-            result[node] = device_names;
           }
         }
         return result;

From f240cd119bcc993fec8f2cfbb2d38cb9e8bcfcf6 Mon Sep 17 00:00:00 2001
From: Mina Ibrahim <136262314+MinaIbrahim10@users.noreply.github.com>
Date: Fri, 2 Oct 2026 15:26:01 +0300
Subject: [PATCH 0152/1000] Add missing utility include for std::pair

---
 tensorflow/python/grappler/cluster_wrapper.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/grappler/cluster_wrapper.cc b/tensorflow/python/grappler/cluster_wrapper.cc
index b59c852227a..7bed2e6771b 100644
--- a/tensorflow/python/grappler/cluster_wrapper.cc
+++ b/tensorflow/python/grappler/cluster_wrapper.cc
@@ -24,8 +24,8 @@ limitations under the License.
 #include <string>
 #include <tuple>
 #include <unordered_map>
+#include <utility>
 #include <vector>
-
 #include "absl/base/no_destructor.h"
 #include "absl/log/check.h"
 #include "absl/status/status.h"

From 8a35517c8d2d286f7fa0371909caf15a1902651f Mon Sep 17 00:00:00 2001
From: Mina Ibrahim <136262314+MinaIbrahim10@users.noreply.github.com>
Date: Fri, 2 Oct 2026 09:24:34 -0400
Subject: [PATCH 0153/1000] Use deterministic ordering in
 TF_GetSupportedDevices

---
 tensorflow/python/grappler/cluster_wrapper.cc | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/grappler/cluster_wrapper.cc b/tensorflow/python/grappler/cluster_wrapper.cc
index 7bed2e6771b..9dd328fabac 100644
--- a/tensorflow/python/grappler/cluster_wrapper.cc
+++ b/tensorflow/python/grappler/cluster_wrapper.cc
@@ -12,7 +12,7 @@ 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.
 ==============================================================================*/
-
+#include <map>
 #include <algorithm>
 #include <cfloat>
 #include <cstddef>
@@ -221,7 +221,7 @@ PYBIND11_MODULE(
         std::unordered_map<std::string, std::vector<std::string>> result;
         {
           py::gil_scoped_release release;
-          std::unordered_map<std::string, std::set<std::string>> device_types;
+          std::map<std::string, std::set<std::string>> device_types;
           std::string cluster_type;
 
           {
@@ -310,7 +310,7 @@ PYBIND11_MODULE(
                   }
                 }
               }
-              result[node] = device_names;
+              result.emplace(node, std::move(device_names));
             }
           }
         }

From 49456a2e1958f11ef036cfe9577e9b7f07bb9e8b Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Fri, 2 Oct 2026 19:26:37 +0530
Subject: [PATCH 0154/1000] size non-memcpyable components unconditionally in
 UncompressElement

---
 tensorflow/core/data/compression_utils.cc     | 35 +++++++------------
 .../core/data/compression_utils_test.cc       | 23 +++++++++---
 2 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/tensorflow/core/data/compression_utils.cc b/tensorflow/core/data/compression_utils.cc
index e26860bd94d..d5cd4feb67e 100644
--- a/tensorflow/core/data/compression_utils.cc
+++ b/tensorflow/core/data/compression_utils.cc
@@ -166,15 +166,22 @@ absl::Status UncompressElement(const CompressedElement& compressed,
     if (metadata.dtype() == DT_STRING) {
       ++num_string_tensors;
       num_string_tensor_strings += metadata.uncompressed_bytes_size();
+    } else if (!DataTypeCanUseMemcpy(metadata.dtype())) {
+      // Non-`memcpy`able components always serialize to a non-empty
+      // `TensorProto`, even when the tensor has zero elements, so size the
+      // scratch buffer for them regardless of element count. A forged byte
+      // count is then caught by the snappy size reconciliation below.
+      if (metadata.uncompressed_bytes_size() == 0) {
+        return absl::InvalidArgumentError(
+            "Missing uncompressed_bytes metadata for non-memcpyable tensor");
+      }
+      total_nonmemcpyable_size += metadata.uncompressed_bytes(0);
     } else {
       int64_t num_elements = shape.num_elements();
       if (num_elements > 0 && metadata.uncompressed_bytes_size() == 0) {
         return absl::InvalidArgumentError(
             "Missing uncompressed_bytes metadata for non-empty tensor");
       }
-      if (!DataTypeCanUseMemcpy(metadata.dtype()) && num_elements > 0) {
-        total_nonmemcpyable_size += metadata.uncompressed_bytes(0);
-      }
     }
   }
 
@@ -221,15 +228,9 @@ absl::Status UncompressElement(const CompressedElement& compressed,
         iov.Add(flats.data()[i].mdata(), metadata.uncompressed_bytes(i));
       }
     } else {
-      TensorShape shape;
-      TF_RETURN_IF_ERROR(
-          TensorShape::BuildTensorShape(metadata.tensor_shape(), &shape));
-      int64_t num_elements = shape.num_elements();
       out->emplace_back();
-      if (num_elements > 0) {
-        iov.Add(nonmemcpyable_pos, metadata.uncompressed_bytes(0));
-        nonmemcpyable_pos += metadata.uncompressed_bytes(0);
-      }
+      iov.Add(nonmemcpyable_pos, metadata.uncompressed_bytes(0));
+      nonmemcpyable_pos += metadata.uncompressed_bytes(0);
     }
   }
 
@@ -260,18 +261,6 @@ absl::Status UncompressElement(const CompressedElement& compressed,
         compressed.component_metadata(i);
     if (!DataTypeCanUseMemcpy(metadata.dtype()) &&
         metadata.dtype() != DT_STRING) {
-      TensorShape shape;
-      TF_RETURN_IF_ERROR(
-          TensorShape::BuildTensorShape(metadata.tensor_shape(), &shape));
-      if (shape.num_elements() == 0) {
-        // The first two passes reserve and fill `nonmemcpyable` only for
-        // components with a non-zero element count, so there are no bytes
-        // here to deserialize from.
-        return absl::InvalidArgumentError(absl::StrCat(
-            "Zero-element component ", i, " with non-memcpyable dtype ",
-            DataTypeString(metadata.dtype()),
-            " has no uncompressed bytes to deserialize"));
-      }
       TensorProto tp;
       if (!tp.ParseFromString(
               {nonmemcpyable_pos,
diff --git a/tensorflow/core/data/compression_utils_test.cc b/tensorflow/core/data/compression_utils_test.cc
index 8ba7e593c07..f0dcea856b1 100644
--- a/tensorflow/core/data/compression_utils_test.cc
+++ b/tensorflow/core/data/compression_utils_test.cc
@@ -45,9 +45,9 @@ TEST(CompressionUtilsTest, Exceeds4GB) {
 }
 
 TEST(CompressionUtilsTest, ZeroElementNonMemcpyableComponent) {
-  // Compressing an empty element yields a `data` field holding zero
-  // uncompressed bytes, so the size reconciliation against the (empty) iovec
-  // succeeds and the third pass is reached.
+  // If metadata claims a non-memcpyable component has 65536 uncompressed bytes
+  // when the compressed data only has 0 bytes, the size reconciliation against
+  // the iovec fails before any decompression or deserialization occurs.
   std::vector<Tensor> empty_element;
   CompressedElement compressed;
   TF_ASSERT_OK(CompressElement(empty_element, &compressed));
@@ -61,8 +61,21 @@ TEST(CompressionUtilsTest, ZeroElementNonMemcpyableComponent) {
 
   std::vector<Tensor> element;
   EXPECT_THAT(UncompressElement(compressed, &element),
-              absl_testing::StatusIs(error::INVALID_ARGUMENT,
-                                     HasSubstr("Zero-element component")));
+              absl_testing::StatusIs(error::INTERNAL,
+                                     HasSubstr("Uncompressed size mismatch")));
+}
+
+TEST(CompressionUtilsTest, RoundTripEmptyVariantTensor) {
+  // A zero-element variant tensor is a legitimate, first-class tensor and must
+  // round-trip through compression without being rejected.
+  std::vector<Tensor> element = {Tensor(DT_VARIANT, TensorShape{0})};
+  CompressedElement compressed;
+  TF_ASSERT_OK(CompressElement(element, &compressed));
+  std::vector<Tensor> round_trip;
+  TF_ASSERT_OK(UncompressElement(compressed, &round_trip));
+  ASSERT_EQ(round_trip.size(), 1);
+  EXPECT_EQ(round_trip[0].dtype(), DT_VARIANT);
+  EXPECT_EQ(round_trip[0].shape(), TensorShape({0}));
 }
 
 TEST(CompressionUtilsTest, MalformedTensorShape) {

From 4983544371afb1818adcb1259a198e90dd8b06e9 Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Fri, 2 Oct 2026 19:59:44 +0530
Subject: [PATCH 0155/1000] Validate element_shape_proto with IsValid before
 constructing shape

BuildPartialTensorShape coerces partial dimensions below -1 to unknown
instead of rejecting them, so a malformed element_shape_proto decoded as
valid and DecodeRejectsMalformedElementShapeProto failed in CI. Guard the
proto with PartialTensorShape::IsValid before constructing the shape; it
rejects dims < -1 and excessive rank and cannot abort. Drop the now-unused
logging.h include.

Signed-off-by: mohammed adib <adib@bugqore.com>
---
 tensorflow/core/kernels/tensor_list.cc | 18 +++++++-----------
 1 file changed, 7 insertions(+), 11 deletions(-)

diff --git a/tensorflow/core/kernels/tensor_list.cc b/tensorflow/core/kernels/tensor_list.cc
index 01c16c081b3..db51b279776 100644
--- a/tensorflow/core/kernels/tensor_list.cc
+++ b/tensorflow/core/kernels/tensor_list.cc
@@ -20,7 +20,6 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_shape.pb.h"
 #include "tensorflow/core/framework/variant_op_registry.h"
 #include "tensorflow/core/lib/core/coding.h"
-#include "tensorflow/core/platform/logging.h"
 
 namespace tensorflow {
 
@@ -126,16 +125,13 @@ bool TensorList::Decode(const VariantTensorData& data) {
   TensorShapeProto element_shape_proto;
   if (!element_shape_proto.ParseFromString(iter)) return false;
   // The PartialTensorShape constructor calls AddDim(), which fatally CHECKs on a
-  // malformed rank or dimension size. The proto is untrusted, so build the shape
-  // through the factory, which validates and constructs in a single pass and
-  // returns an error instead of aborting the process.
-  PartialTensorShape decoded_element_shape;
-  if (!PartialTensorShape::BuildPartialTensorShape(element_shape_proto,
-                                                   &decoded_element_shape)
-           .ok()) {
-    VLOG(1) << "TensorList::Decode failed: malformed element_shape_proto";
-    return false;
-  }
+  // malformed rank or dimension size. BuildPartialTensorShape is not a safe
+  // substitute here: for partial shapes it silently coerces any dimension below
+  // -1 to unknown instead of rejecting it. The proto is untrusted, so validate
+  // it with IsValid (which also rejects dims < -1 and excessive rank) before
+  // constructing the shape.
+  if (!PartialTensorShape::IsValid(element_shape_proto)) return false;
+  const PartialTensorShape decoded_element_shape(element_shape_proto);
 
   // Every element stored in a list must match its element_dtype and be
   // compatible with its element_shape; the mutation paths (TensorListPushBack,

From a2a108ebf9818fca9d01b5e2ebe5d248e68397dd Mon Sep 17 00:00:00 2001
From: Aditya Gaurav <mr.x.l.r.8.pride@gmail.com>
Date: Fri, 2 Oct 2026 18:41:34 +0000
Subject: [PATCH 0156/1000] fix(sequence_ops): interpolate half and bfloat16
 LinSpace in float, leave Range unchanged

---
 tensorflow/core/kernels/sequence_ops.cc       | 17 ++--------
 tensorflow/core/kernels/sequence_ops_test.cc  | 34 ++-----------------
 .../python/ops/math_ops_linspace_test.py      | 22 ++++++++----
 3 files changed, 20 insertions(+), 53 deletions(-)

diff --git a/tensorflow/core/kernels/sequence_ops.cc b/tensorflow/core/kernels/sequence_ops.cc
index 5de7a6dce50..dd8239de734 100644
--- a/tensorflow/core/kernels/sequence_ops.cc
+++ b/tensorflow/core/kernels/sequence_ops.cc
@@ -41,19 +41,8 @@ struct RangeFunctor<CPUDevice, T> {
   void operator()(OpKernelContext* context, int64_t size, T start, T delta,
                   typename TTypes<T>::Flat output) const {
     (void)context;
-    if constexpr (std::is_integral<T>::value) {
-      for (int64_t i = 0; i < size; ++i) {
-        output(i) = start + static_cast<T>(i) * delta;
-      }
-    } else {
-      using ComputeT = typename std::conditional<sizeof(T) >= sizeof(float), T,
-                                                 double>::type;
-      const ComputeT start_c = static_cast<ComputeT>(start);
-      const ComputeT delta_c = static_cast<ComputeT>(delta);
-      for (int64_t i = 0; i < size; ++i) {
-        output(i) =
-            static_cast<T>(start_c + static_cast<ComputeT>(i) * delta_c);
-      }
+    for (int64_t i = 0; i < size; ++i) {
+      output(i) = start + static_cast<T>(i) * delta;
     }
   }
 };
@@ -223,7 +212,7 @@ class LinSpaceOp : public OpKernel {
       // half/bfloat16, while preserving native performance and numerical
       // semantics for float/double.
       using ComputeT = typename std::conditional<sizeof(T) >= sizeof(float), T,
-                                                 double>::type;
+                                                 float>::type;
       const ComputeT start_c = static_cast<ComputeT>(start);
       const ComputeT stop_c = static_cast<ComputeT>(stop);
       const ComputeT step = (stop_c - start_c) / static_cast<ComputeT>(num - 1);
diff --git a/tensorflow/core/kernels/sequence_ops_test.cc b/tensorflow/core/kernels/sequence_ops_test.cc
index 8d96dd82f7a..873fa0a0833 100644
--- a/tensorflow/core/kernels/sequence_ops_test.cc
+++ b/tensorflow/core/kernels/sequence_ops_test.cc
@@ -131,36 +131,6 @@ TEST_F(RangeOpTest, Range_Size_Overflow) {
             RunOpKernel().message());
 }
 
-TEST_F(RangeOpTest, Monotonic_Bfloat16) {
-  MakeOp(DT_BFLOAT16);
-
-  // 600 elements go past 256, where bfloat16 stops holding every integer, so
-  // casting i to bfloat16 before multiplying rounded it. delta is smaller than
-  // the bfloat16 spacing near the top of the range, so neighbouring outputs
-  // may be equal; each one must still be start + i * delta rounded once.
-  const bfloat16 start(0.0f);
-  const bfloat16 delta(0.3f);
-  AddInputFromArray<bfloat16>(TensorShape({}), {start});
-  AddInputFromArray<bfloat16>(TensorShape({}), {bfloat16(180.0f)});
-  AddInputFromArray<bfloat16>(TensorShape({}), {delta});
-  TF_ASSERT_OK(RunOpKernel());
-
-  const auto flat = GetOutput(0)->flat<bfloat16>();
-  constexpr int64_t num = 600;
-  ASSERT_EQ(num, GetOutput(0)->NumElements());
-  for (int64_t i = 1; i < num; ++i) {
-    ASSERT_GE(static_cast<float>(flat(i)), static_cast<float>(flat(i - 1)))
-        << "flat(" << i << ") was less than flat(" << i - 1 << ")";
-  }
-  for (int64_t i = 0; i < num; ++i) {
-    ASSERT_EQ(static_cast<float>(static_cast<bfloat16>(
-                  static_cast<double>(start) +
-                  static_cast<double>(i) * static_cast<double>(delta))),
-              static_cast<float>(flat(i)))
-        << "flat(" << i << ") is not start + " << i << " * delta";
-  }
-}
-
 TEST_F(LinSpaceOpTest, Simple_D32) {
   MakeOp(DT_FLOAT, DT_INT32);
 
@@ -314,10 +284,10 @@ TEST_F(LinSpaceOpTest, Monotonic_Bfloat16_Int64Index) {
     ASSERT_GE(static_cast<float>(flat(i)), static_cast<float>(flat(i - 1)))
         << "flat(" << i << ") was less than flat(" << i - 1 << ")";
   }
-  const double step = 1.0 / static_cast<double>(num - 1);
+  const float step = 1.0f / static_cast<float>(num - 1);
   for (int64_t i = 1; i < num - 1; ++i) {
     ASSERT_EQ(static_cast<float>(
-                  static_cast<bfloat16>(step * static_cast<double>(i))),
+                  static_cast<bfloat16>(step * static_cast<float>(i))),
               static_cast<float>(flat(i)))
         << "flat(" << i << ") is not the nearest bfloat16 to " << i << "/"
         << num - 1;
diff --git a/tensorflow/python/ops/math_ops_linspace_test.py b/tensorflow/python/ops/math_ops_linspace_test.py
index f58b2c4491b..549922d444c 100644
--- a/tensorflow/python/ops/math_ops_linspace_test.py
+++ b/tensorflow/python/ops/math_ops_linspace_test.py
@@ -81,21 +81,29 @@ class LinspaceTest(test_util.TensorFlowTestCase, parameterized.TestCase):
   ])
   def testLinSpaceKernelHalfAndBfloat16(self, dtype, num_dtype):
     # tf.linspace is composed from other ops, so call the LinSpace kernel
-    # directly. Each element must be the float64 value narrowed to dtype the
-    # way the kernel's cast does it (through float32); num=70000 is past
-    # float16's largest finite value (65504).
+    # directly. The kernel interpolates in float32 and narrows each element
+    # once; num=70000 is past float16's largest finite value (65504). The
+    # tolerance is one ulp of dtype, so NaN, Inf or a drifted step still fail.
     np_dtype = dtype.as_numpy_dtype
+    ulp = 2.0**-10 if dtype == dtypes.float16 else 2.0**-7
     for start, stop, num in [(0.0, 1.0, 600), (-3.0, 5.0, 20), (9.0, 100.0, 1),
                              (0.0, 1.0, 70000)]:
       actual = self.evaluate(
           gen_math_ops.lin_space(
               np.array(start, np_dtype), np.array(stop, np_dtype),
               np.array(num, num_dtype)))
-      expected = np.linspace(
-          start, stop, num, dtype=np.float64).astype(np.float32).astype(np_dtype)
+      start32 = np.float32(np_dtype(start))
+      stop32 = np.float32(np_dtype(stop))
+      expected = np.full(num, start32, dtype=np.float32)
+      if num > 1:
+        step32 = (stop32 - start32) / np.float32(num - 1)
+        expected = start32 + step32 * np.arange(num, dtype=np.float32)
+        expected[0] = start32
+        expected[-1] = stop32
+      expected = expected.astype(np_dtype).astype(np.float32)
       self.assertEqual(np_dtype, actual.dtype)
-      self.assertAllEqual(expected.astype(np.float32),
-                          actual.astype(np.float32))
+      self.assertAllClose(
+          expected, actual.astype(np.float32), rtol=ulp, atol=0.0)
 
   @parameterized.parameters([dtypes.float16, dtypes.bfloat16])
   def testLinspaceHalfAndBfloat16(self, dtype):

From f1ef1e37f86accdc21240fa43881db9d9019f633 Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Fri, 2 Oct 2026 20:04:47 -0400
Subject: [PATCH 0157/1000] security/fuzzing: do not zero kTfLiteMmapRo inputs
 in the TFLite fuzz target

A crafted model can declare a subgraph input backed by a constant buffer;
that tensor is kTfLiteMmapRo and points into the read-only model mapping,
so the MSan zero-fill would fault on it. Skip those; they are initialized
by the model already. Per review.
---
 .../security/fuzzing/cc/tflite_interp_overflow_fuzz.cc      | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
index bf51c157ea7..9bcf2ba874c 100644
--- a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
+++ b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
@@ -57,9 +57,13 @@ void FuzzModelBuildAndInvoke(const std::string& model_bytes) {
   if (interpreter->AllocateTensors() != kTfLiteOk) return;
 
   // Zero-initialize runtime inputs so MSan does not flag uninitialized reads.
+  // An input a crafted model backs with a constant buffer is kTfLiteMmapRo and
+  // points into the read-only model mapping; it is already initialized and
+  // writing to it would fault, so leave those alone.
   for (int input_index : interpreter->inputs()) {
     TfLiteTensor* tensor = interpreter->tensor(input_index);
-    if (tensor != nullptr && tensor->data.raw != nullptr) {
+    if (tensor != nullptr && tensor->data.raw != nullptr &&
+        tensor->allocation_type != kTfLiteMmapRo) {
       std::memset(tensor->data.raw, 0, tensor->bytes);
     }
   }

From 0d760bba0b4a328cb5a2cc55befa52704a442c05 Mon Sep 17 00:00:00 2001
From: Rohith Pariki <rohithpariki@gmail.com>
Date: Sat, 3 Oct 2026 05:54:23 +0530
Subject: [PATCH 0158/1000] Address review feedback: simulate buffer reuse in
 calibrator test, pass TfLiteTensor by const ref, and fix namespace formatting

---
 tensorflow/lite/tools/optimize/calibration/calibrator.cc  | 2 +-
 .../lite/tools/optimize/calibration/calibrator_test.cc    | 8 +++++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/tensorflow/lite/tools/optimize/calibration/calibrator.cc b/tensorflow/lite/tools/optimize/calibration/calibrator.cc
index 2b8addf2598..8df5bc4f974 100644
--- a/tensorflow/lite/tools/optimize/calibration/calibrator.cc
+++ b/tensorflow/lite/tools/optimize/calibration/calibrator.cc
@@ -260,7 +260,7 @@ TfLiteStatus LoggingEval(TfLiteContext* context, TfLiteNode* node) {
   // output tensors, which would overwrite input calibration statistics with
   // output values.
   for (int i : op_info.loggable_inputs) {
-    auto tensor = context->tensors[i];
+    const TfLiteTensor& tensor = context->tensors[i];
     if (tensor.is_variable) {
       TF_LITE_ENSURE_STATUS(
           logger->LogTensorValue(op_info.subgraph_index, i, tensor.data.f,
diff --git a/tensorflow/lite/tools/optimize/calibration/calibrator_test.cc b/tensorflow/lite/tools/optimize/calibration/calibrator_test.cc
index 70c9642f88e..ed24f86ce18 100644
--- a/tensorflow/lite/tools/optimize/calibration/calibrator_test.cc
+++ b/tensorflow/lite/tools/optimize/calibration/calibrator_test.cc
@@ -738,6 +738,12 @@ TEST(CalibratorTest, NonVariableInputsNotReLoggedPostInvocation) {
     }
   }
 
+  // Simulate memory planner buffer reuse where an output tensor shares the
+  // buffer of an input tensor.
+  ASSERT_NE(interpreter->tensor(5), nullptr);
+  char* original_output_data = interpreter->tensor(5)->data.raw;
+  interpreter->tensor(5)->data.raw = interpreter->tensor(0)->data.raw;
+
   status = interpreter->Invoke();
   ASSERT_EQ(kTfLiteOk, status);
 
@@ -752,9 +758,9 @@ TEST(CalibratorTest, NonVariableInputsNotReLoggedPostInvocation) {
     EXPECT_NEAR(it->second.min, tensor_idx + 1, eps);
     EXPECT_NEAR(it->second.max, tensor_idx + 1, eps);
   }
+  interpreter->tensor(5)->data.raw = original_output_data;
 }
 }  // namespace
-
 }  // namespace calibration
 }  // namespace optimize
 }  // namespace tflite

From d481802c00127f44d8607a4841a620972fbb6e45 Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Fri, 2 Oct 2026 20:46:46 -0400
Subject: [PATCH 0159/1000] security/fuzzing: include <vector> and
 tensorflow/lite/c/common.h directly

The target uses TfLiteTensor, kTfLiteMmapRo, kTfLiteOk and std::vector;
include their headers instead of relying on transitive includes. Per review.
---
 tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
index 9bcf2ba874c..01b286acc4c 100644
--- a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
+++ b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
@@ -25,8 +25,10 @@ limitations under the License.
 #include <cstring>
 #include <memory>
 #include <string>
+#include <vector>
 
 #include "fuzztest/fuzztest.h"
+#include "tensorflow/lite/c/common.h"
 #include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/interpreter_builder.h"
 #include "tensorflow/lite/kernels/register.h"

From ae5b26db83123e460eb2658796ebce1601b9586e Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Fri, 2 Oct 2026 20:47:19 -0400
Subject: [PATCH 0160/1000] security/fuzzing: depend on
 //tensorflow/lite/c:common for the TFLite fuzz target

Matches the direct include of tensorflow/lite/c/common.h. Per review.
---
 tensorflow/security/fuzzing/cc/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/security/fuzzing/cc/BUILD b/tensorflow/security/fuzzing/cc/BUILD
index 26c767712d8..36b5d41c9f4 100644
--- a/tensorflow/security/fuzzing/cc/BUILD
+++ b/tensorflow/security/fuzzing/cc/BUILD
@@ -259,6 +259,7 @@ tf_cc_fuzz_test(
     srcs = ["tflite_interp_overflow_fuzz.cc"],
     deps = [
         "//tensorflow/lite:framework",
+        "//tensorflow/lite/c:common",
         "//tensorflow/lite/kernels:builtin_ops",
         "//tensorflow/lite/tools:verifier",
     ],

From 72272479b6209adccbc3f7ba7660af22cb607769 Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Fri, 2 Oct 2026 21:51:55 -0400
Subject: [PATCH 0161/1000] security/fuzzing: drop the tflite::Verify() gate
 from the TFLite fuzz target

Per the Oct 3 review: malformed models go straight to InterpreterBuilder
so its own handling of model metadata is what gets fuzzed. Removes the
Verify() early return and its include; keeps the direct includes and the
kTfLiteMmapRo exclusion from the earlier rounds.
---
 .../fuzzing/cc/tflite_interp_overflow_fuzz.cc | 24 +++++--------------
 1 file changed, 6 insertions(+), 18 deletions(-)

diff --git a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
index 01b286acc4c..13936a6ab2f 100644
--- a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
+++ b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
@@ -13,14 +13,12 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-// FuzzTest coverage for the TFLite untrusted-model ingress path: model
-// verification, model load, interpreter build, and invoke. tflite::Verify() is
-// the validation layer for untrusted models, so the target runs it first and
-// stops on rejection; everything it accepts goes on to ParseTensors and the
-// kernels. This covers the verifier's checks on attacker-controlled model
-// metadata (including the uint64 external-offset bounds checks in
-// VerifyTensors / VerifyOperators) and the runtime behind them. The OSS-Fuzz
-// tensorflow build has no TFLite target that reaches this path.
+// FuzzTest coverage for TFLite model load, interpreter build, and invoke.
+// Malformed models go straight to InterpreterBuilder, so anything it does with
+// attacker-controlled model metadata (the external-offset bounds checks in
+// ParseNodes and ParseTensors included) is exercised directly rather than
+// behind the verifier. The OSS-Fuzz tensorflow build has no TFLite target that
+// reaches this path.
 
 #include <cstring>
 #include <memory>
@@ -33,19 +31,10 @@ limitations under the License.
 #include "tensorflow/lite/interpreter_builder.h"
 #include "tensorflow/lite/kernels/register.h"
 #include "tensorflow/lite/model_builder.h"
-#include "tensorflow/lite/tools/verifier.h"
 
 namespace {
 
 void FuzzModelBuildAndInvoke(const std::string& model_bytes) {
-  // Supported ingress path: an untrusted model is verified before anything is
-  // built from it. Inputs the verifier rejects stop here. The error reporter
-  // is optional and would only add stderr noise under the fuzzer.
-  if (!tflite::Verify(model_bytes.data(), model_bytes.size(),
-                      /*error_reporter=*/nullptr)) {
-    return;
-  }
-
   auto model = tflite::FlatBufferModel::BuildFromBuffer(model_bytes.data(),
                                                         model_bytes.size());
   if (model == nullptr) return;
@@ -70,7 +59,6 @@ void FuzzModelBuildAndInvoke(const std::string& model_bytes) {
     }
   }
 
-  // Runs the kernels on whatever the verifier accepted.
   interpreter->Invoke();
 }
 FUZZ_TEST(TfliteInterpreterBuilder, FuzzModelBuildAndInvoke);

From fd2f77c447ec35203d6a6db8c0e672ba68033507 Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Fri, 2 Oct 2026 21:52:27 -0400
Subject: [PATCH 0162/1000] security/fuzzing: drop the verifier dep from the
 TFLite fuzz target

The target no longer calls tflite::Verify(). Per the Oct 3 review.
---
 tensorflow/security/fuzzing/cc/BUILD | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/security/fuzzing/cc/BUILD b/tensorflow/security/fuzzing/cc/BUILD
index 5c8339b97f5..66c81352f36 100644
--- a/tensorflow/security/fuzzing/cc/BUILD
+++ b/tensorflow/security/fuzzing/cc/BUILD
@@ -276,6 +276,5 @@ tf_cc_fuzz_test(
         "//tensorflow/lite:framework",
         "//tensorflow/lite/c:common",
         "//tensorflow/lite/kernels:builtin_ops",
-        "//tensorflow/lite/tools:verifier",
     ],
 )

From db0ba16e36ff663983f49a9e5a7b544d227dd84b Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Fri, 2 Oct 2026 21:53:21 -0400
Subject: [PATCH 0163/1000] tflite: overflow-safe external-offset bounds checks
 in InterpreterBuilder

Both sites compared offset + size against allocation_->bytes() with two
uint64_t operands; a pair whose sum wraps passed the check and the pointer
computed from base + offset wrapped with it. Compare in the subtraction
form instead, so the check cannot overflow. Same predicate for every valid
(offset, size); only the wrapping pairs change outcome.

Requested on this PR on Sep 16 and again on Oct 3 so the fuzz target in
the same change does not report a known crash. Re-applies the change from
#116631 that was reverted on Sep 18; the verifier-side check from #123910
stays, so both layers reject the wrap.
---
 tensorflow/lite/core/interpreter_builder.cc | 83 ++++-----------------
 1 file changed, 14 insertions(+), 69 deletions(-)

diff --git a/tensorflow/lite/core/interpreter_builder.cc b/tensorflow/lite/core/interpreter_builder.cc
index bea6d38df1f..4d1cf3867bc 100644
--- a/tensorflow/lite/core/interpreter_builder.cc
+++ b/tensorflow/lite/core/interpreter_builder.cc
@@ -375,9 +375,9 @@ TfLiteStatus InterpreterBuilder::ParseNodes(
         init_data = reinterpret_cast<const char*>(op->custom_options()->data());
         init_data_size = op->custom_options()->size();
       } else if (op->large_custom_options_offset() > 1 && allocation_) {
-        if (op->large_custom_options_offset() +
-                op->large_custom_options_size() >
-            allocation_->bytes()) {
+        if (op->large_custom_options_offset() > allocation_->bytes() ||
+            op->large_custom_options_size() >
+                allocation_->bytes() - op->large_custom_options_offset()) {
           TF_LITE_REPORT_ERROR(
               error_reporter_,
               "Custom Option Offset for opcode_index %d is out of bound\n",
@@ -501,72 +501,16 @@ TfLiteStatus InterpreterBuilder::ParseQuantization(
   if (src_quantization->details_type() ==
       QuantizationDetails_BlockwiseQuantization) {
     auto* src_quant = src_quantization->details_as_BlockwiseQuantization();
-    if (!src_quant) {
-      TF_LITE_REPORT_ERROR(error_reporter_,
-                           "Blockwise quantization details are missing.");
-      return kTfLiteError;
-    }
-    // `block_shape` supersedes `block_size`. It is only meaningful if it has
-    // exactly one entry per dimension of the tensor, so reject anything else
-    // rather than silently falling back to `block_size`, which would produce
-    // wrong values.
-    const auto* block_shape = src_quant->block_shape();
-    if (block_shape && !block_shape->empty()) {
-      if (!dims.empty() &&
-          static_cast<size_t>(block_shape->size()) != dims.size()) {
-        TF_LITE_REPORT_ERROR(
-            error_reporter_,
-            "BlockwiseQuantization block_shape must have one entry per "
-            "dimension. Expected %d entries, got %d.",
-            static_cast<int>(dims.size()), block_shape->size());
-        return kTfLiteError;
-      }
-      for (int i = 0; i < block_shape->size(); ++i) {
-        if (block_shape->Get(i) <= 0) {
-          TF_LITE_REPORT_ERROR(error_reporter_,
-                               "BlockwiseQuantization block_shape entries must "
-                               "be positive. Value %d at index %d is invalid.",
-                               block_shape->Get(i), i);
-          return kTfLiteError;
-        }
-      }
-      auto* blockwise_quantization_v2 =
-          reinterpret_cast<TfLiteBlockwiseQuantizationV2*>(
-              malloc(sizeof(TfLiteBlockwiseQuantizationV2)));
-      if (!blockwise_quantization_v2) {
-        return kTfLiteError;
-      }
-      blockwise_quantization_v2->scale = src_quant->scales();
-      blockwise_quantization_v2->zero_point = src_quant->zero_points();
-      blockwise_quantization_v2->quantized_dimension =
-          src_quantization->quantized_dimension();
-      blockwise_quantization_v2->blocksize = src_quant->block_size();
-      blockwise_quantization_v2->block_shape =
-          TfLiteIntArrayCreate(static_cast<int>(block_shape->size()));
-      if (!blockwise_quantization_v2->block_shape) {
-        free(blockwise_quantization_v2);
-        return kTfLiteError;
-      }
-      for (int i = 0; i < block_shape->size(); ++i) {
-        blockwise_quantization_v2->block_shape->data[i] = block_shape->Get(i);
-      }
-      quantization->type = kTfLiteBlockwiseQuantizationV2;
-      quantization->params = reinterpret_cast<void*>(blockwise_quantization_v2);
-      return kTfLiteOk;
-    }
-    auto* blockwise_quantization_v1 =
+    quantization->type = kTfLiteBlockwiseQuantization;
+    auto* blockwise_quantization =
         reinterpret_cast<TfLiteBlockwiseQuantization*>(
             malloc(sizeof(TfLiteBlockwiseQuantization)));
-    if (!blockwise_quantization_v1) {
-      return kTfLiteError;
-    }
-    blockwise_quantization_v1->scale = src_quant->scales();
-    blockwise_quantization_v1->zero_point = src_quant->zero_points();
-    blockwise_quantization_v1->quantized_dimension =
+    blockwise_quantization->scale = src_quant->scales();
+    blockwise_quantization->zero_point = src_quant->zero_points();
+    blockwise_quantization->quantized_dimension =
         src_quantization->quantized_dimension();
-    blockwise_quantization_v1->blocksize = src_quant->block_size();
-    quantization->type = kTfLiteBlockwiseQuantization;
-    quantization->params = reinterpret_cast<void*>(blockwise_quantization_v1);
+    blockwise_quantization->blocksize = src_quant->block_size();
+    quantization->params = reinterpret_cast<void*>(blockwise_quantization);
     return kTfLiteOk;
   }
   if (!src_quantization->scale() || src_quantization->scale()->empty()) {
@@ -729,7 +673,7 @@ TfLiteStatus InterpreterBuilder::ParseSignatureDefs(
     const flatbuffers::Vector<flatbuffers::Offset<SignatureDef>>*
         signature_def_list,
     Interpreter* interpreter) {
-  if (signature_def_list == nullptr || signature_def_list->empty()) {
+  if (signature_def_list == nullptr || signature_def_list->size() == 0) {
     return kTfLiteOk;
   }
   std::vector<internal::SignatureDef> signature_defs;
@@ -811,7 +755,8 @@ TfLiteStatus InterpreterBuilder::ParseTensors(
           *buffer_data = reinterpret_cast<const char*>(array->data());
           return kTfLiteOk;
         } else if (offset > 1 && allocation_) {
-          if (offset + buffer->size() > allocation_->bytes()) {
+          if (offset > allocation_->bytes() ||
+              buffer->size() > allocation_->bytes() - offset) {
             TF_LITE_REPORT_ERROR(
                 error_reporter_,
                 "Constant buffer %d specified an out of range offset.\n",
@@ -967,7 +912,7 @@ TfLiteStatus InterpreterBuilder::operator()(
   auto* subgraphs = model_->subgraphs();
   auto* buffers = model_->buffers();
 
-  if (subgraphs->empty()) {
+  if (subgraphs->size() == 0) {
     TF_LITE_REPORT_ERROR(error_reporter_, "No subgraph in the model.\n");
     return kTfLiteError;
   }

From 72567388959f68b271686cd87b6721845457bdee Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Fri, 2 Oct 2026 21:54:35 -0400
Subject: [PATCH 0164/1000] tflite: restore master ParseQuantization dropped by
 the previous commit

db0ba16 was applied to a stale copy of interpreter_builder.cc and removed
56 lines of ParseQuantization that master had added. This puts the file
back to master plus only the two bounds-check hunks.
---
 tensorflow/lite/core/interpreter_builder.cc | 74 ++++++++++++++++++---
 1 file changed, 65 insertions(+), 9 deletions(-)

diff --git a/tensorflow/lite/core/interpreter_builder.cc b/tensorflow/lite/core/interpreter_builder.cc
index 4d1cf3867bc..fd06180005d 100644
--- a/tensorflow/lite/core/interpreter_builder.cc
+++ b/tensorflow/lite/core/interpreter_builder.cc
@@ -501,16 +501,72 @@ TfLiteStatus InterpreterBuilder::ParseQuantization(
   if (src_quantization->details_type() ==
       QuantizationDetails_BlockwiseQuantization) {
     auto* src_quant = src_quantization->details_as_BlockwiseQuantization();
-    quantization->type = kTfLiteBlockwiseQuantization;
-    auto* blockwise_quantization =
+    if (!src_quant) {
+      TF_LITE_REPORT_ERROR(error_reporter_,
+                           "Blockwise quantization details are missing.");
+      return kTfLiteError;
+    }
+    // `block_shape` supersedes `block_size`. It is only meaningful if it has
+    // exactly one entry per dimension of the tensor, so reject anything else
+    // rather than silently falling back to `block_size`, which would produce
+    // wrong values.
+    const auto* block_shape = src_quant->block_shape();
+    if (block_shape && !block_shape->empty()) {
+      if (!dims.empty() &&
+          static_cast<size_t>(block_shape->size()) != dims.size()) {
+        TF_LITE_REPORT_ERROR(
+            error_reporter_,
+            "BlockwiseQuantization block_shape must have one entry per "
+            "dimension. Expected %d entries, got %d.",
+            static_cast<int>(dims.size()), block_shape->size());
+        return kTfLiteError;
+      }
+      for (int i = 0; i < block_shape->size(); ++i) {
+        if (block_shape->Get(i) <= 0) {
+          TF_LITE_REPORT_ERROR(error_reporter_,
+                               "BlockwiseQuantization block_shape entries must "
+                               "be positive. Value %d at index %d is invalid.",
+                               block_shape->Get(i), i);
+          return kTfLiteError;
+        }
+      }
+      auto* blockwise_quantization_v2 =
+          reinterpret_cast<TfLiteBlockwiseQuantizationV2*>(
+              malloc(sizeof(TfLiteBlockwiseQuantizationV2)));
+      if (!blockwise_quantization_v2) {
+        return kTfLiteError;
+      }
+      blockwise_quantization_v2->scale = src_quant->scales();
+      blockwise_quantization_v2->zero_point = src_quant->zero_points();
+      blockwise_quantization_v2->quantized_dimension =
+          src_quantization->quantized_dimension();
+      blockwise_quantization_v2->blocksize = src_quant->block_size();
+      blockwise_quantization_v2->block_shape =
+          TfLiteIntArrayCreate(static_cast<int>(block_shape->size()));
+      if (!blockwise_quantization_v2->block_shape) {
+        free(blockwise_quantization_v2);
+        return kTfLiteError;
+      }
+      for (int i = 0; i < block_shape->size(); ++i) {
+        blockwise_quantization_v2->block_shape->data[i] = block_shape->Get(i);
+      }
+      quantization->type = kTfLiteBlockwiseQuantizationV2;
+      quantization->params = reinterpret_cast<void*>(blockwise_quantization_v2);
+      return kTfLiteOk;
+    }
+    auto* blockwise_quantization_v1 =
         reinterpret_cast<TfLiteBlockwiseQuantization*>(
             malloc(sizeof(TfLiteBlockwiseQuantization)));
-    blockwise_quantization->scale = src_quant->scales();
-    blockwise_quantization->zero_point = src_quant->zero_points();
-    blockwise_quantization->quantized_dimension =
+    if (!blockwise_quantization_v1) {
+      return kTfLiteError;
+    }
+    blockwise_quantization_v1->scale = src_quant->scales();
+    blockwise_quantization_v1->zero_point = src_quant->zero_points();
+    blockwise_quantization_v1->quantized_dimension =
         src_quantization->quantized_dimension();
-    blockwise_quantization->blocksize = src_quant->block_size();
-    quantization->params = reinterpret_cast<void*>(blockwise_quantization);
+    blockwise_quantization_v1->blocksize = src_quant->block_size();
+    quantization->type = kTfLiteBlockwiseQuantization;
+    quantization->params = reinterpret_cast<void*>(blockwise_quantization_v1);
     return kTfLiteOk;
   }
   if (!src_quantization->scale() || src_quantization->scale()->empty()) {
@@ -673,7 +729,7 @@ TfLiteStatus InterpreterBuilder::ParseSignatureDefs(
     const flatbuffers::Vector<flatbuffers::Offset<SignatureDef>>*
         signature_def_list,
     Interpreter* interpreter) {
-  if (signature_def_list == nullptr || signature_def_list->size() == 0) {
+  if (signature_def_list == nullptr || signature_def_list->empty()) {
     return kTfLiteOk;
   }
   std::vector<internal::SignatureDef> signature_defs;
@@ -912,7 +968,7 @@ TfLiteStatus InterpreterBuilder::operator()(
   auto* subgraphs = model_->subgraphs();
   auto* buffers = model_->buffers();
 
-  if (subgraphs->size() == 0) {
+  if (subgraphs->empty()) {
     TF_LITE_REPORT_ERROR(error_reporter_, "No subgraph in the model.\n");
     return kTfLiteError;
   }

From c624d970b2667b4e96aa54579552600a8ed38e44 Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Fri, 2 Oct 2026 22:18:57 -0400
Subject: [PATCH 0165/1000] security/fuzzing: enter through
 VerifyAndBuildFromBuffer in the TFLite fuzz target

Per review: models now pass the flatbuffers structural verifier before
InterpreterBuilder sees them. No extra_verifier, so tflite::Verify() is not
on this path; the external-offset bounds checks in interpreter_builder.cc
remain the check exercised for wrapped offsets.
---
 .../fuzzing/cc/tflite_interp_overflow_fuzz.cc   | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
index 13936a6ab2f..25f95875c29 100644
--- a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
+++ b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
@@ -14,11 +14,14 @@ limitations under the License.
 ==============================================================================*/
 
 // FuzzTest coverage for TFLite model load, interpreter build, and invoke.
-// Malformed models go straight to InterpreterBuilder, so anything it does with
-// attacker-controlled model metadata (the external-offset bounds checks in
-// ParseNodes and ParseTensors included) is exercised directly rather than
-// behind the verifier. The OSS-Fuzz tensorflow build has no TFLite target that
-// reaches this path.
+// Models enter through FlatBufferModel::VerifyAndBuildFromBuffer, so inputs
+// that fail the flatbuffers structural check are rejected at ingress, as they
+// would be in a deployment that verifies models. That check does not cover the
+// uint64 external-offset fields (Buffer.offset/size and
+// Operator.large_custom_options_offset/size), which point past the flatbuffer
+// into the appended data, so the bounds checks on them in
+// InterpreterBuilder::ParseNodes and ParseTensors are exercised directly. The
+// OSS-Fuzz tensorflow build has no TFLite target that reaches this path.
 
 #include <cstring>
 #include <memory>
@@ -35,8 +38,8 @@ limitations under the License.
 namespace {
 
 void FuzzModelBuildAndInvoke(const std::string& model_bytes) {
-  auto model = tflite::FlatBufferModel::BuildFromBuffer(model_bytes.data(),
-                                                        model_bytes.size());
+  auto model = tflite::FlatBufferModel::VerifyAndBuildFromBuffer(
+      model_bytes.data(), model_bytes.size());
   if (model == nullptr) return;
 
   tflite::ops::builtin::BuiltinOpResolver resolver;

From 6b637521dbde57526747589a95a68d9e0bcba52d Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Fri, 2 Oct 2026 22:49:12 -0400
Subject: [PATCH 0166/1000] lite: cache offset, size and allocation bytes in
 the external-offset checks

Per review: hoist op->large_custom_options_offset(), the option size and
allocation_->bytes() into locals in ParseNodes, and allocation_->bytes() and
buffer->size() in ParseTensors, so the checks do not repeat the flatbuffer
accessors and the virtual call. No behaviour change.
---
 tensorflow/lite/core/interpreter_builder.cc | 36 ++++++++++++---------
 1 file changed, 20 insertions(+), 16 deletions(-)

diff --git a/tensorflow/lite/core/interpreter_builder.cc b/tensorflow/lite/core/interpreter_builder.cc
index fd06180005d..cc7156b1f58 100644
--- a/tensorflow/lite/core/interpreter_builder.cc
+++ b/tensorflow/lite/core/interpreter_builder.cc
@@ -374,20 +374,23 @@ TfLiteStatus InterpreterBuilder::ParseNodes(
       if (op->custom_options()) {
         init_data = reinterpret_cast<const char*>(op->custom_options()->data());
         init_data_size = op->custom_options()->size();
-      } else if (op->large_custom_options_offset() > 1 && allocation_) {
-        if (op->large_custom_options_offset() > allocation_->bytes() ||
-            op->large_custom_options_size() >
-                allocation_->bytes() - op->large_custom_options_offset()) {
-          TF_LITE_REPORT_ERROR(
-              error_reporter_,
-              "Custom Option Offset for opcode_index %d is out of bound\n",
-              index);
-          return kTfLiteError;
+      } else if (allocation_) {
+        const uint64_t offset = op->large_custom_options_offset();
+        if (offset > 1) {
+          const size_t alloc_bytes = allocation_->bytes();
+          const uint64_t size = op->large_custom_options_size();
+          if (offset > alloc_bytes || size > alloc_bytes - offset) {
+            TF_LITE_REPORT_ERROR(
+                error_reporter_,
+                "Custom Option Offset for opcode_index %d is out of bound\n",
+                index);
+            return kTfLiteError;
+          }
+          // If the custom op is storing payloads outside of flatbuffers
+          init_data =
+              reinterpret_cast<const char*>(allocation_->base()) + offset;
+          init_data_size = size;
         }
-        // If the custom op is storing payloads outside of flatbuffers
-        init_data = reinterpret_cast<const char*>(allocation_->base()) +
-                    op->large_custom_options_offset();
-        init_data_size = op->large_custom_options_size();
       }
     } else {
       MallocDataAllocator malloc_allocator;
@@ -811,15 +814,16 @@ TfLiteStatus InterpreterBuilder::ParseTensors(
           *buffer_data = reinterpret_cast<const char*>(array->data());
           return kTfLiteOk;
         } else if (offset > 1 && allocation_) {
-          if (offset > allocation_->bytes() ||
-              buffer->size() > allocation_->bytes() - offset) {
+          const size_t alloc_bytes = allocation_->bytes();
+          const uint64_t size = buffer->size();
+          if (offset > alloc_bytes || size > alloc_bytes - offset) {
             TF_LITE_REPORT_ERROR(
                 error_reporter_,
                 "Constant buffer %d specified an out of range offset.\n",
                 tensor->buffer());
             return kTfLiteError;
           }
-          *buffer_size = buffer->size();
+          *buffer_size = size;
           *buffer_data =
               reinterpret_cast<const char*>(allocation_->base()) + offset;
           return kTfLiteOk;

From 267042d6b57f914b92baca79a17dcb2b5bc4ca63 Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Fri, 2 Oct 2026 22:50:12 -0400
Subject: [PATCH 0167/1000] security/fuzzing: drop unused <vector> include from
 the TFLite fuzz target

Per review (IWYU): the file never names std::vector.
---
 tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
index 25f95875c29..a74226b62f1 100644
--- a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
+++ b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
@@ -26,7 +26,6 @@ limitations under the License.
 #include <cstring>
 #include <memory>
 #include <string>
-#include <vector>
 
 #include "fuzztest/fuzztest.h"
 #include "tensorflow/lite/c/common.h"

From e9612f26a1f95d898257ad3c0644fc31e8edeb1b Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Fri, 2 Oct 2026 23:16:11 -0400
Subject: [PATCH 0168/1000] security/fuzzing: restore <vector> include in the
 TFLite fuzz target

Per the corrected review: interpreter->inputs() returns const std::vector<int>&
and the range-for over it is a use, so IWYU wants the header named here.
---
 tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
index a74226b62f1..25f95875c29 100644
--- a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
+++ b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include <cstring>
 #include <memory>
 #include <string>
+#include <vector>
 
 #include "fuzztest/fuzztest.h"
 #include "tensorflow/lite/c/common.h"

From ed39579b8bb3b82d206c432daab066d2fd67ed80 Mon Sep 17 00:00:00 2001
From: Rohith Pariki <rohithpariki@gmail.com>
Date: Sat, 3 Oct 2026 16:47:00 +0530
Subject: [PATCH 0169/1000] style: avoid pass-by-value copy of TfLiteTensor in
 all LoggingEval loops

---
 tensorflow/lite/tools/optimize/calibration/calibrator.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/lite/tools/optimize/calibration/calibrator.cc b/tensorflow/lite/tools/optimize/calibration/calibrator.cc
index 8df5bc4f974..0d8ad89f859 100644
--- a/tensorflow/lite/tools/optimize/calibration/calibrator.cc
+++ b/tensorflow/lite/tools/optimize/calibration/calibrator.cc
@@ -232,7 +232,7 @@ TfLiteStatus LoggingEval(TfLiteContext* context, TfLiteNode* node) {
   auto error_reporter = calibrator->GetErrorReporter();
 
   for (int i : op_info.loggable_inputs) {
-    auto tensor = context->tensors[i];
+    const TfLiteTensor& tensor = context->tensors[i];
     TF_LITE_ENSURE_STATUS(
         logger->LogTensorValue(op_info.subgraph_index, i, tensor.data.f,
                                tensor.bytes / sizeof(float), error_reporter));
@@ -269,7 +269,7 @@ TfLiteStatus LoggingEval(TfLiteContext* context, TfLiteNode* node) {
   }
 
   for (int i : op_info.loggable_outputs) {
-    auto tensor = context->tensors[i];
+    const TfLiteTensor& tensor = context->tensors[i];
     TF_LITE_ENSURE_STATUS(
         logger->LogTensorValue(op_info.subgraph_index, i, tensor.data.f,
                                tensor.bytes / sizeof(float), error_reporter));

From 026c2539169f5a55bf830ed7bb40ab527bcd3e77 Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Sat, 3 Oct 2026 15:14:57 -0700
Subject: [PATCH 0170/1000] Validate axis bounds in experimental.numpy average

NumPy raises AxisError for out-of-bounds axes, but average passed
them straight to the backend, so an invalid axis could surface as a
confusing error from reduce_mean or an IndexError from the weights
returned path. When the input rank is statically known, the axis is
now normalized and bounds-checked with a clear ValueError, matching
np.average. Behavior for valid axes is unchanged; the dynamic-rank
path keeps its previous runtime behavior.
---
 tensorflow/python/ops/numpy_ops/np_math_ops.py      | 10 ++++++++++
 tensorflow/python/ops/numpy_ops/np_math_ops_test.py |  9 +++++++++
 2 files changed, 19 insertions(+)

diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py
index 9b0490eb60b..10b9b05696f 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@@ -1674,6 +1674,16 @@ def average(a, axis=None, weights=None, returned=False):  # pylint: disable=miss
         f'Received axis={axis} (of type {type(axis)})'
     )
   a = np_array_ops.array(a)
+  # NumPy raises AxisError for out-of-bounds axes instead of letting the
+  # backend kernel fail with a confusing error.
+  maybe_rank = a.shape.rank
+  if axis is not None and maybe_rank is not None:
+    normalized = axis + maybe_rank if axis < 0 else axis
+    if normalized < 0 or normalized >= maybe_rank:
+      raise ValueError(
+          f'Argument `axis` (received axis={axis}) is out of bounds '
+          f'for input of rank {maybe_rank}.'
+      )
   default_float_type = np_utils.result_type(float)
   if weights is None:  # Treat all weights as 1
     if not np.issubdtype(a.dtype.as_numpy_dtype, np.inexact):
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 84869bc9458..4208f893988 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
@@ -571,6 +571,15 @@ class MathTest(test.TestCase, parameterized.TestCase):
     ):
       np_math_ops.diff(x, n=-1)
 
+  def testAverageAxisOutOfBounds(self):
+    x = np_array_ops.array(np.ones([2, 3]))
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_math_ops.average(x, axis=2)
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_math_ops.average(x, axis=-3)
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_math_ops.average(x, axis=2, weights=np.ones([2, 3]))
+
   def testAverageWrongShape(self):
     with self.assertRaisesWithPredicateMatch(errors.InvalidArgumentError, r''):
       np_math_ops.average(np.ones([2, 3]), weights=np.ones([2, 4]))

From a9ebdb1d7f05e4e9d42622a109bb6cb2688029fa Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Sun, 4 Oct 2026 07:43:02 +0530
Subject: [PATCH 0171/1000] Fix code formatting and refactor master job
 retrieval

---
 tensorflow/python/eager/remote.py | 40 +++++++++++++++++++++----------
 1 file changed, 27 insertions(+), 13 deletions(-)

diff --git a/tensorflow/python/eager/remote.py b/tensorflow/python/eager/remote.py
index 2808f5a056f..bfe00231654 100644
--- a/tensorflow/python/eager/remote.py
+++ b/tensorflow/python/eager/remote.py
@@ -44,7 +44,7 @@ def connect_to_remote_host(remote_host=None, job_name="worker"):
 
   Using the default job_name of worker, you can schedule ops to run remotely as
   follows:
-  ```python
+```python
   # When eager execution is enabled, connect to the remote host.
   tf.config.experimental_connect_to_host("exampleaddr.com:9876")
 
@@ -54,7 +54,7 @@ def connect_to_remote_host(remote_host=None, job_name="worker"):
     x1 = array_ops.ones([2, 2])
     x2 = array_ops.ones([2, 2])
     y = math_ops.matmul(x1, x2)
-  ```
+```
 
   Args:
     remote_host: a single or a list the remote server addr in host-port format.
@@ -73,6 +73,27 @@ def connect_to_remote_host(remote_host=None, job_name="worker"):
   connect_to_cluster(cluster_spec)
 
 
+def _find_master_job_and_task(cluster_spec, master):
+  """Finds the (job_name, task_id) in cluster_spec whose address matches master.
+
+  Returns the first match found.
+
+  Args:
+    cluster_spec: A `ClusterSpec` describing the cluster.
+    master: The master address to match against.
+
+  Returns:
+    A tuple `(job_name, task_id)` for the first matching task, or
+    `(None, None)` if no task address matches.
+  """
+  for job_name in cluster_spec.jobs:
+    for task_id in cluster_spec.task_indices(job_name):
+      task_address = cluster_spec.task_address(job_name, task_id)
+      if master in task_address or task_address in master:
+        return job_name, task_id
+  return None, None
+
+
 @tf_export("config.experimental_connect_to_cluster")
 def connect_to_cluster(cluster_spec_or_resolver,
                        job_name="localhost",
@@ -100,7 +121,7 @@ def connect_to_cluster(cluster_spec_or_resolver,
   For example, for a cluster set up for parameter server training, the following
   device filters might be specified:
 
-  ```python
+```python
   cdf = tf.config.experimental.ClusterDeviceFilters()
   # For any worker, only the devices on PS nodes and itself are visible
   for i in range(num_workers):
@@ -111,7 +132,7 @@ def connect_to_cluster(cluster_spec_or_resolver,
 
   tf.config.experimental_connect_to_cluster(cluster_def,
                                             cluster_device_filters=cdf)
-  ```
+```
 
   Args:
     cluster_spec_or_resolver: A `ClusterSpec` or `ClusterResolver` describing
@@ -242,15 +263,8 @@ def connect_to_cluster(cluster_spec_or_resolver,
       cluster_spec_or_resolver,
       cluster_resolver.ClusterResolver) and cluster_spec_or_resolver.master():
     master = cluster_spec_or_resolver.master()
-    master_job_name = None
-    master_task_id = None
-    for job_name in cluster_spec.jobs:
-      for task_id in cluster_spec.task_indices(job_name):
-        task_address = cluster_spec.task_address(job_name, task_id)
-        if master in task_address or task_address in master:
-          master_job_name = job_name
-          master_task_id = task_id
-          break
+    master_job_name, master_task_id = _find_master_job_and_task(
+        cluster_spec, master)
 
     if not master_job_name:
       raise ValueError(

From 42d16760a8e361f6adb67d9d559f0deccba77cb4 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Sun, 4 Oct 2026 07:49:55 +0530
Subject: [PATCH 0172/1000] Update remote.py

---
 tensorflow/python/eager/remote.py | 33 ++++++++++++++++++++++++++++---
 1 file changed, 30 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/eager/remote.py b/tensorflow/python/eager/remote.py
index bfe00231654..e785fca558b 100644
--- a/tensorflow/python/eager/remote.py
+++ b/tensorflow/python/eager/remote.py
@@ -15,6 +15,7 @@
 """Helpers to connect to remote servers."""
 
 import copy
+import urllib.parse
 
 from absl import logging
 
@@ -73,10 +74,30 @@ def connect_to_remote_host(remote_host=None, job_name="worker"):
   connect_to_cluster(cluster_spec)
 
 
+def _parse_host_and_port(address):
+  """Splits an address, with or without a scheme, into (host, port).
+
+  Args:
+    address: An address such as `"host:8000"` or `"grpc://host:8000"`.
+
+  Returns:
+    A tuple `(host, port)`. `port` is None if the address has no port, and
+    both are None if the address cannot be parsed.
+  """
+  if "://" not in address:
+    address = "//" + address
+  try:
+    parsed = urllib.parse.urlsplit(address)
+    return parsed.hostname, parsed.port
+  except ValueError:
+    return None, None
+
+
 def _find_master_job_and_task(cluster_spec, master):
   """Finds the (job_name, task_id) in cluster_spec whose address matches master.
 
-  Returns the first match found.
+  An address matches when the hosts are equal and, if both addresses specify a
+  port, the ports are equal. Returns the first match found.
 
   Args:
     cluster_spec: A `ClusterSpec` describing the cluster.
@@ -86,11 +107,17 @@ def _find_master_job_and_task(cluster_spec, master):
     A tuple `(job_name, task_id)` for the first matching task, or
     `(None, None)` if no task address matches.
   """
+  master_host, master_port = _parse_host_and_port(master)
+  if not master_host:
+    return None, None
+
   for job_name in cluster_spec.jobs:
     for task_id in cluster_spec.task_indices(job_name):
       task_address = cluster_spec.task_address(job_name, task_id)
-      if master in task_address or task_address in master:
-        return job_name, task_id
+      task_host, task_port = _parse_host_and_port(task_address)
+      if task_host == master_host:
+        if master_port is None or task_port is None or master_port == task_port:
+          return job_name, task_id
   return None, None
 
 

From 0193d170dba8cca45f7189170265a85453fe324e Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Sun, 4 Oct 2026 00:38:24 -0400
Subject: [PATCH 0173/1000] lite: regression tests for wrapping external
 offsets in InterpreterBuilder

Per review: two deterministic tests in model_test.cc. Each builds a model whose
(offset, size) pair sums past uint64 to 9, loads it with BuildFromBuffer so the
verifier is not in the way, and expects InterpreterBuilder to reject it. One
covers the Buffer offset path in ParseTensors (tensor bytes equal the claimed
size so the size check cannot mask the result), the other the
large_custom_options path in ParseNodes.
---
 tensorflow/lite/core/model_test.cc | 79 ++++++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)

diff --git a/tensorflow/lite/core/model_test.cc b/tensorflow/lite/core/model_test.cc
index 81fadef874c..90924eaff2b 100644
--- a/tensorflow/lite/core/model_test.cc
+++ b/tensorflow/lite/core/model_test.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include <stdlib.h>
 #include <string.h>
 
+#include <limits>
 #include <map>
 #include <memory>
 #include <string>
@@ -511,6 +512,84 @@ TEST(BasicFlatBufferModel, TestInvalidBlockwiseQuantizationDetails) {
             kTfLiteError);
 }
 
+// A constant tensor whose Buffer lives outside the flatbuffer, addressed by an
+// (offset, size) pair chosen so that offset + size wraps around uint64 to a
+// small value. The model is loaded without the verifier on purpose, so the
+// bounds check in InterpreterBuilder::ParseTensors is what has to reject it.
+TEST(BasicFlatBufferModel, TestRejectsWrappingExternalBufferOffset) {
+  flatbuffers::FlatBufferBuilder builder;
+  const uint64_t offset = std::numeric_limits<uint64_t>::max() - 10;
+  const uint64_t size = 20;  // offset + size wraps to 9.
+  std::vector<int32_t> shape = {5};  // 5 x float32 == 20 bytes == size.
+  auto tensor = ::tflite::CreateTensorDirect(builder, &shape,
+                                             tflite::TensorType_FLOAT32,
+                                             /*buffer=*/1, "weights");
+  std::vector<flatbuffers::Offset<::tflite::Tensor>> tensors = {tensor};
+  auto subgraph = ::tflite::CreateSubGraphDirect(
+      builder, &tensors, /*inputs=*/nullptr, /*outputs=*/nullptr,
+      /*operators=*/nullptr, "subgraph");
+  std::vector<flatbuffers::Offset<::tflite::SubGraph>> subgraphs = {subgraph};
+  std::vector<flatbuffers::Offset<::tflite::Buffer>> buffers = {
+      ::tflite::CreateBuffer(builder),
+      ::tflite::CreateBuffer(builder, /*data=*/0, offset, size)};
+  auto model = ::tflite::CreateModelDirect(builder, TFLITE_SCHEMA_VERSION,
+                                           /*operator_codes=*/nullptr,
+                                           &subgraphs, "model", &buffers);
+  ::tflite::FinishModelBuffer(builder, model);
+  auto fb_model = FlatBufferModel::BuildFromBuffer(
+      reinterpret_cast<const char*>(builder.GetBufferPointer()),
+      builder.GetSize());
+  ASSERT_TRUE(fb_model);
+  std::unique_ptr<Interpreter> interpreter;
+  TrivialResolver resolver;
+  EXPECT_NE(InterpreterBuilder(*fb_model, resolver)(&interpreter), kTfLiteOk);
+}
+
+// The same wrapping pair on a custom op's large_custom_options, which is the
+// InterpreterBuilder::ParseNodes side of the check.
+TEST(BasicFlatBufferModel, TestRejectsWrappingLargeCustomOptionsOffset) {
+  flatbuffers::FlatBufferBuilder builder;
+  const uint64_t offset = std::numeric_limits<uint64_t>::max() - 10;
+  const uint64_t size = 20;  // offset + size wraps to 9.
+  std::vector<int32_t> shape = {1};
+  auto tensor = ::tflite::CreateTensorDirect(builder, &shape,
+                                             tflite::TensorType_FLOAT32,
+                                             /*buffer=*/0, "t");
+  std::vector<flatbuffers::Offset<::tflite::Tensor>> tensors = {tensor};
+  std::vector<int32_t> inputs = {0};
+  std::vector<int32_t> outputs = {0};
+  auto inputs_vector = builder.CreateVector(inputs);
+  auto outputs_vector = builder.CreateVector(outputs);
+  ::tflite::OperatorBuilder op(builder);
+  op.add_opcode_index(0);
+  op.add_inputs(inputs_vector);
+  op.add_outputs(outputs_vector);
+  op.add_large_custom_options_offset(offset);
+  op.add_large_custom_options_size(size);
+  std::vector<flatbuffers::Offset<::tflite::Operator>> operators = {
+      op.Finish()};
+  auto subgraph = ::tflite::CreateSubGraphDirect(
+      builder, &tensors, &inputs, &outputs, &operators, "subgraph");
+  std::vector<flatbuffers::Offset<::tflite::SubGraph>> subgraphs = {subgraph};
+  std::vector<flatbuffers::Offset<::tflite::OperatorCode>> operator_codes = {
+      ::tflite::CreateOperatorCodeDirect(
+          builder, /*deprecated_builtin_code=*/tflite::BuiltinOperator_CUSTOM,
+          "WRAP", /*version=*/1, tflite::BuiltinOperator_CUSTOM)};
+  std::vector<flatbuffers::Offset<::tflite::Buffer>> buffers = {
+      ::tflite::CreateBuffer(builder)};
+  auto model = ::tflite::CreateModelDirect(builder, TFLITE_SCHEMA_VERSION,
+                                           &operator_codes, &subgraphs,
+                                           "model", &buffers);
+  ::tflite::FinishModelBuffer(builder, model);
+  auto fb_model = FlatBufferModel::BuildFromBuffer(
+      reinterpret_cast<const char*>(builder.GetBufferPointer()),
+      builder.GetSize());
+  ASSERT_TRUE(fb_model);
+  std::unique_ptr<Interpreter> interpreter;
+  TrivialResolver resolver(&dummy_reg);
+  EXPECT_NE(InterpreterBuilder(*fb_model, resolver)(&interpreter), kTfLiteOk);
+}
+
 TEST(BasicFlatBufferModel, TestWithNumThreads) {
   TestErrorReporter reporter;
   auto model = FlatBufferModel::BuildFromFile(

From 97ed8ec2531a8443fa992c84b891ea51ed68b9cd Mon Sep 17 00:00:00 2001
From: Arpan Chatterjee <arpanchatterjee648@gmail.com>
Date: Sun, 4 Oct 2026 13:15:57 +0530
Subject: [PATCH 0174/1000] fix typos

---
 .../tools/kernel_gen/transforms/broadcast_propagation_pass.cc   | 2 +-
 .../core/api_def/base_api/api_def_IteratorGetModelProto.pbtxt   | 2 +-
 tensorflow/python/autograph/converters/control_flow.py          | 2 +-
 .../python/eager/polymorphic_function/function_type_utils.py    | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/broadcast_propagation_pass.cc b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/broadcast_propagation_pass.cc
index 9c571bd9ec9..01f573676bc 100644
--- a/tensorflow/compiler/mlir/tools/kernel_gen/transforms/broadcast_propagation_pass.cc
+++ b/tensorflow/compiler/mlir/tools/kernel_gen/transforms/broadcast_propagation_pass.cc
@@ -128,7 +128,7 @@ DenseIntElementsAttr composeBroadcastDimensionsAttr(OpBuilder &builder,
 }
 
 // Find all the broadcast intents and their dependencies. Start analyzing from
-// the root a collect all broadcast intents that can help broadcast propagation
+// the root and collect all broadcast intents that can help broadcast propagation
 // from there.
 void findBroadcastIntents(
     DynamicBroadcastInDimOp root, Block *parentBlock,
diff --git a/tensorflow/core/api_def/base_api/api_def_IteratorGetModelProto.pbtxt b/tensorflow/core/api_def/base_api/api_def_IteratorGetModelProto.pbtxt
index 67e9d40f034..6fdad506756 100644
--- a/tensorflow/core/api_def/base_api/api_def_IteratorGetModelProto.pbtxt
+++ b/tensorflow/core/api_def/base_api/api_def_IteratorGetModelProto.pbtxt
@@ -4,7 +4,7 @@ op {
   in_arg {
     name: "iterator"
     description: <<END
-An resource from a dataset iterator.
+A resource from a dataset iterator.
 END
   }
   out_arg {
diff --git a/tensorflow/python/autograph/converters/control_flow.py b/tensorflow/python/autograph/converters/control_flow.py
index b28a06353f6..9a41503e2f8 100644
--- a/tensorflow/python/autograph/converters/control_flow.py
+++ b/tensorflow/python/autograph/converters/control_flow.py
@@ -37,7 +37,7 @@ class _Function(object):
 
 
 class ControlFlowTransformer(converter.Base):
-  """Transforms control flow structures like loops a conditionals."""
+  """Transforms control flow structures like loops and conditionals."""
 
   def visit_Lambda(self, node):
     with self.state[_Function] as fn:
diff --git a/tensorflow/python/eager/polymorphic_function/function_type_utils.py b/tensorflow/python/eager/polymorphic_function/function_type_utils.py
index 18589aed15e..ed7a3648e53 100644
--- a/tensorflow/python/eager/polymorphic_function/function_type_utils.py
+++ b/tensorflow/python/eager/polymorphic_function/function_type_utils.py
@@ -391,7 +391,7 @@ def canonicalize_function_inputs(
   Python function's positional (named) arguments and a dictionary
   corresponding to its kwargs.  Missing default arguments are added.
 
-  If the FunctionType has a type constraints, then they are used to convert
+  If the FunctionType has type constraints, then they are used to convert
   arguments to tensors; otherwise, any inputs containing numpy arrays are
   converted to tensors.
 

From c4925c5147fc8f3220a34ffe3fc22e27fad9e48b Mon Sep 17 00:00:00 2001
From: Noxtimo <139155109+Noxtimo@users.noreply.github.com>
Date: Sun, 4 Oct 2026 12:12:39 +0100
Subject: [PATCH 0175/1000] Preserve negative zero in experimental.numpy.cbrt

---
 .../python/ops/numpy_ops/np_math_ops.py       |  9 +++++++-
 .../python/ops/numpy_ops/np_math_ops_test.py  | 22 +++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py
index 9b0490eb60b..5a978cb903f 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@@ -994,7 +994,14 @@ def cbrt(x):
   def f(x):
     # __pow__ can't handle negative base, so we use `abs` here.
     rt = math_ops.abs(x) ** (1.0 / 3)
-    return array_ops.where_v2(x < 0, -rt, rt)
+    if x.dtype in _SIGN_BITCAST_DTYPES:
+      # Check the IEEE-754 sign bit instead of comparing with zero, which
+      # cannot tell -0.0 from +0.0.
+      bits = array_ops.bitcast(x, _SIGN_BITCAST_DTYPES[x.dtype])
+      is_negative = math_ops.less(bits, 0)
+    else:
+      is_negative = x < 0
+    return array_ops.where_v2(is_negative, -rt, rt)
 
   return _scalar(f, x, True)
 
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 84869bc9458..dab3005a55c 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
@@ -858,6 +858,28 @@ class MathTest(test.TestCase, parameterized.TestCase):
     negative_zero = ops.convert_to_tensor([-0.0], dtype=dtypes.bfloat16)
     self.assertAllEqual(np_math_ops.signbit(negative_zero), [True])
 
+  def testCbrt(self):
+    for dtype in (dtypes.float16, dtypes.float32, dtypes.float64):
+      x_np = np.array(
+          [-0.0, 0.0, -8.0, 8.0, -27.0, 27.0], dtype=dtype.as_numpy_dtype
+      )
+      x = constant_op.constant(x_np, dtype=dtype)
+      out = np_math_ops.cbrt(x)
+      self.assertAllClose(out, np.cbrt(x_np))
+      self.assertAllEqual(np.signbit(out.numpy()), np.signbit(np.cbrt(x_np)))
+
+    x_neg_zero = constant_op.constant(-0.0, dtype=dtypes.float64)
+    self.assertEqual(
+        math_ops.reciprocal(np_math_ops.cbrt(x_neg_zero)).numpy(),
+        -float('inf'),
+    )
+
+    x_bf16 = constant_op.constant([-0.0, 0.0, -8.0, 8.0], dtype=dtypes.bfloat16)
+    out_bf16 = np_math_ops.cbrt(x_bf16)
+    self.assertAllEqual(
+        np_math_ops.signbit(out_bf16), [True, False, True, False]
+    )
+
   def testSinc(self):
     for dtype in (dtypes.float32, dtypes.float64):
       x = constant_op.constant([0.0, -0.0, 0.5, 1.0, -1.5], dtype=dtype)

From f4b5382ebecb14fc4cbd39bdcf89d88dd36f4ea1 Mon Sep 17 00:00:00 2001
From: Arpan Chatterjee <arpanchatterjee648@gmail.com>
Date: Sun, 4 Oct 2026 17:00:46 +0530
Subject: [PATCH 0176/1000] fix typos

---
 tensorflow/go/op/wrappers.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/go/op/wrappers.go b/tensorflow/go/op/wrappers.go
index acba7bb4ce8..2d20048a400 100644
--- a/tensorflow/go/op/wrappers.go
+++ b/tensorflow/go/op/wrappers.go
@@ -22055,7 +22055,7 @@ func IteratorGetDevice(scope *Scope, resource tf.Output) (device tf.Output) {
 //
 // Arguments:
 //
-//	iterator: An resource from a dataset iterator.
+//	iterator: A resource from a dataset iterator.
 //
 // Returns A serialized model proto.
 func IteratorGetModelProto(scope *Scope, iterator tf.Output) (model_proto tf.Output) {

From 0e0285a59b4141c7245b8a8202020df9ea70cbcf Mon Sep 17 00:00:00 2001
From: daniel <liujingtai2001@gmail.com>
Date: Sun, 4 Oct 2026 08:49:54 -0400
Subject: [PATCH 0177/1000] Add half-precision gradient tests and clarify beta
 docstring

Document that beta may be a tensor broadcastable to features and add
SwishTest.testFloat16AndBfloat16Gradients covering float16 and bfloat16
gradients with a broadcast beta.
---
 tensorflow/python/ops/BUILD      |  1 +
 tensorflow/python/ops/nn_impl.py |  3 ++-
 tensorflow/python/ops/nn_test.py | 14 ++++++++++++++
 3 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD
index f72ab18839b..72fe69457c5 100644
--- a/tensorflow/python/ops/BUILD
+++ b/tensorflow/python/ops/BUILD
@@ -3786,6 +3786,7 @@ cuda_py_strict_test(
         ":stateful_random_ops",
         ":variable_scope",
         ":variables",
+        "//tensorflow/python/eager:backprop",
         "//tensorflow/python/eager:context",
         "//tensorflow/python/eager:def_function",
         "//tensorflow/python/framework:constant_op",
diff --git a/tensorflow/python/ops/nn_impl.py b/tensorflow/python/ops/nn_impl.py
index 17fc44dce5a..34539552608 100644
--- a/tensorflow/python/ops/nn_impl.py
+++ b/tensorflow/python/ops/nn_impl.py
@@ -456,7 +456,8 @@ def swish(features, beta=1.0):
 
   Args:
     features: A `Tensor` representing preactivation values.
-    beta: A 'Tensor' representing value of beta hyperparameter.
+    beta: A `Tensor` representing value of beta hyperparameter, broadcastable
+      to `features`. Default value 1.0.
 
   Returns:
     The activation value.
diff --git a/tensorflow/python/ops/nn_test.py b/tensorflow/python/ops/nn_test.py
index bcb2a49e582..2a3e2b6ec28 100644
--- a/tensorflow/python/ops/nn_test.py
+++ b/tensorflow/python/ops/nn_test.py
@@ -21,6 +21,7 @@ import sys
 from absl.testing import parameterized
 import numpy as np
 
+from tensorflow.python.eager import backprop
 from tensorflow.python.eager import context
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
@@ -1167,6 +1168,19 @@ class SwishTest(test_lib.TestCase, parameterized.TestCase):
     self.assertAllClose(theoretical[0], numerical[0], atol=1e-4)
     self.assertAllClose(theoretical[1], numerical[1], atol=1e-4)
 
+  def testFloat16AndBfloat16Gradients(self):
+    for dtype in [dtypes.float16, dtypes.bfloat16]:
+      features = constant_op.constant([[1.0, 2.0], [3.0, 4.0]], dtype=dtype)
+      beta = constant_op.constant([0.5, 1.5], dtype=dtype)
+      with backprop.GradientTape() as tape:
+        tape.watch([features, beta])
+        y = nn_impl.swish(features, beta)
+      g_features, g_beta = tape.gradient(y, [features, beta])
+      self.assertEqual(g_features.dtype, dtype)
+      self.assertEqual(g_beta.dtype, dtype)
+      self.assertEqual(g_features.shape.as_list(), [2, 2])
+      self.assertEqual(g_beta.shape.as_list(), [2])
+
 
 class MomentsTest(test_lib.TestCase):
 

From ababf98026daeea46aacb2fdee86b8e695492ac0 Mon Sep 17 00:00:00 2001
From: Jaydeng75 <user.josephms@gmail.com>
Date: Sun, 4 Oct 2026 20:26:49 +0530
Subject: [PATCH 0178/1000] Remove redundant SparseConcat empty-input guard and
 test

---
 tensorflow/core/kernels/sparse_concat_op.cc         |  2 --
 .../sparse_ops/sparse_concat_op_test.py             | 13 -------------
 2 files changed, 15 deletions(-)

diff --git a/tensorflow/core/kernels/sparse_concat_op.cc b/tensorflow/core/kernels/sparse_concat_op.cc
index 9d8f2dceed2..e6072768212 100644
--- a/tensorflow/core/kernels/sparse_concat_op.cc
+++ b/tensorflow/core/kernels/sparse_concat_op.cc
@@ -103,8 +103,6 @@ class SparseConcatOp : public OpKernel {
     OpInputList inds;
     OP_REQUIRES_OK(context, context->input_list("indices", &inds));
     const int N = inds.size();
-    OP_REQUIRES(context, N > 0,
-                absl::InvalidArgumentError("Input list must not be empty"));
     for (int i = 0; i < N; i++) {
       OP_REQUIRES(context, TensorShapeUtils::IsMatrix(inds[i].shape()),
                   absl::InvalidArgumentError(absl::StrCat(
diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
index a6b49df6b8a..79bcb9494c9 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
+++ b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
@@ -459,19 +459,6 @@ class SparseConcatTest(test.TestCase):
         self.assertEqual(sp_concat.values.get_shape().as_list(), [None])
         self.assertEqual(sp_concat.dense_shape.get_shape(), [3])
 
-  def testSparseConcatEmptyInputList(self):
-    from tensorflow.python.ops import gen_sparse_ops
-    # With an empty input list there are no tensors to infer the 'T' attr
-    # from, so the kernel is rejected before Compute runs. Asserting that a
-    # clean InvalidArgumentError is raised (rather than a crash on
-    # shapes[0]) guards the N > 0 check in SparseConcatOp.
-    with self.assertRaises(errors.InvalidArgumentError):
-      self.evaluate(
-          gen_sparse_ops.sparse_concat(
-              indices=[], values=[], shapes=[], concat_dim=1
-          )
-      )
-
   def testConcatShape(self):
     # Test case for GitHub 21964.
     x = sparse_tensor.SparseTensor(

From 2c10c04377c30d05add06b08c3c1371bfd94f54b Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Sun, 4 Oct 2026 11:05:22 -0400
Subject: [PATCH 0179/1000] lite: mark the test registration CUSTOM so the
 ParseNodes case reaches the check

Per review: ParseNodes derives the op type from registration->builtin_code,
and dummy_reg carries 0 (ADD), so the large_custom_options branch never ran
and the test passed vacuously. The test now resolves to a copy of dummy_reg
marked BuiltinOperator_CUSTOM with custom_name WRAP.
---
 tensorflow/lite/core/model_test.cc | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tensorflow/lite/core/model_test.cc b/tensorflow/lite/core/model_test.cc
index 90924eaff2b..fff687682c4 100644
--- a/tensorflow/lite/core/model_test.cc
+++ b/tensorflow/lite/core/model_test.cc
@@ -586,7 +586,12 @@ TEST(BasicFlatBufferModel, TestRejectsWrappingLargeCustomOptionsOffset) {
       builder.GetSize());
   ASSERT_TRUE(fb_model);
   std::unique_ptr<Interpreter> interpreter;
-  TrivialResolver resolver(&dummy_reg);
+  // ParseNodes keys off registration->builtin_code, so the resolved
+  // registration itself must be marked CUSTOM.
+  TfLiteRegistration custom_reg = dummy_reg;
+  custom_reg.builtin_code = tflite::BuiltinOperator_CUSTOM;
+  custom_reg.custom_name = "WRAP";
+  TrivialResolver resolver(&custom_reg);
   EXPECT_NE(InterpreterBuilder(*fb_model, resolver)(&interpreter), kTfLiteOk);
 }
 

From 2c808c939e24c7d21e61f14e893b6dc933b683e7 Mon Sep 17 00:00:00 2001
From: Pratik Gandhi <travpreneur@gmail.com>
Date: Sun, 4 Oct 2026 08:54:16 -0700
Subject: [PATCH 0180/1000] Fix relative links in the AutoGraph and Android
 READMEs

---
 tensorflow/python/autograph/CONTRIBUTING.md            | 6 +++---
 tensorflow/tools/android/inference_interface/README.md | 8 ++++----
 tensorflow/tools/android/test/README.md                | 6 +++---
 3 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/tensorflow/python/autograph/CONTRIBUTING.md b/tensorflow/python/autograph/CONTRIBUTING.md
index ba7140658b5..f46623ecf47 100644
--- a/tensorflow/python/autograph/CONTRIBUTING.md
+++ b/tensorflow/python/autograph/CONTRIBUTING.md
@@ -1,6 +1,6 @@
 # How to contribute
 
-We'd love to have your patches and contributions! Here are some guidelines. In general, we follow the [TensorFlow contributing guidelines](../../CONTRIBUTING.md), but have some [AutoGraph-specific style guidelines](STYLE_GUIDE.md). More details below.
+We'd love to have your patches and contributions! Here are some guidelines. In general, we follow the [TensorFlow contributing guidelines](../../../CONTRIBUTING.md), but have some [AutoGraph-specific style guidelines](STYLE_GUIDE.md). More details below.
 
 ### Note to active contributors
 
@@ -10,7 +10,7 @@ does not impact functionality, and AutoGraph will remain accessible under
 `tensorflow.contrib.autograph` until `tensorflow.contrib` is retired.
 
 ## TensorFlow Code of Conduct
-Please review and follow the [TensorFlow Code of Conduct](../../CODE_OF_CONDUCT.md).
+Please review and follow the [TensorFlow Code of Conduct](../../../CODE_OF_CONDUCT.md).
 
 ## Contributor License Agreement
 
@@ -55,7 +55,7 @@ bazel test --config=opt --copt=-O3 --copt=-march=native \
   //tensorflow/contrib/autograph/...
 ```
 
-from the root of the `tensorflow` repository. For more details see the [main TensorFlow Contributing File](../../CONTRIBUTING.md)
+from the root of the `tensorflow` repository. For more details see the [main TensorFlow Contributing File](../../../CONTRIBUTING.md)
 
 ## Developer info
 
diff --git a/tensorflow/tools/android/inference_interface/README.md b/tensorflow/tools/android/inference_interface/README.md
index f45b4d6b3b5..90a52a5fe9d 100644
--- a/tensorflow/tools/android/inference_interface/README.md
+++ b/tensorflow/tools/android/inference_interface/README.md
@@ -3,12 +3,12 @@
 This directory defines components (a native `.so` library and a Java JAR)
 geared towards supporting TensorFlow on Android. This includes:
 
-- The [TensorFlow Java API](../../java/README.md)
+- The [TensorFlow Java API](../../../java/README.md)
 - A `TensorFlowInferenceInterface` class that provides a smaller API
   surface suitable for inference and summarizing performance of model execution.
 
-For example usage, see [TensorFlowImageClassifier.java](../../examples/android/src/org/tensorflow/demo/TensorFlowImageClassifier.java)
-in the [TensorFlow Android Demo](../../examples/android).
+For example usage, see [TensorFlowImageClassifier.java](../test/src/org/tensorflow/demo/TensorFlowImageClassifier.java)
+in the [TensorFlow Android Demo](../test).
 
 For prebuilt libraries, see the
 [nightly Android build artifacts](https://ci.tensorflow.org/view/Nightly/job/nightly-android/)
@@ -44,7 +44,7 @@ TensorFlow operators), pick your preferred approach below:
 ### Bazel
 
 First follow the Bazel setup instructions described in
-[tensorflow/examples/android/README.md](../../examples/android/README.md)
+[tensorflow/tools/android/test/README.md](../test/README.md)
 
 Then, to build the native TF library:
 
diff --git a/tensorflow/tools/android/test/README.md b/tensorflow/tools/android/test/README.md
index f872943824d..6d7d1b42446 100644
--- a/tensorflow/tools/android/test/README.md
+++ b/tensorflow/tools/android/test/README.md
@@ -11,7 +11,7 @@ The demos in this folder are designed to give straightforward samples of using
 TensorFlow in mobile applications.
 
 Inference is done using the [TensorFlow Android Inference
-Interface](../../tools/android/inference_interface), which may be built
+Interface](../inference_interface), which may be built
 separately if you want a standalone library to drop into your existing
 application. Object tracking and efficient YUV -> RGB conversion are handled by
 `libtensorflow_demo.so`.
@@ -51,7 +51,7 @@ The fastest path to trying the demo is to download the [prebuilt demo APK](https
 
 Also available are precompiled native libraries, and a jcenter package that you
 may simply drop into your own applications. See
-[tensorflow/tools/android/inference_interface/README.md](../../tools/android/inference_interface/README.md)
+[tensorflow/tools/android/inference_interface/README.md](../inference_interface/README.md)
 for more details.
 
 ## Running the Demo
@@ -88,7 +88,7 @@ For any project that does not include custom low level TensorFlow code, this is
 likely sufficient.
 
 For details on how to include this JCenter package in your own project see
-[tensorflow/tools/android/inference_interface/README.md](../../tools/android/inference_interface/README.md)
+[tensorflow/tools/android/inference_interface/README.md](../inference_interface/README.md)
 
 ## Building the Demo with TensorFlow from Source
 

From 71542ee6997552d52e42633db8f2cb733db80cd2 Mon Sep 17 00:00:00 2001
From: Jaydeng75 <user.josephms@gmail.com>
Date: Sun, 4 Oct 2026 23:39:49 +0530
Subject: [PATCH 0181/1000] Remove unused errors import from SparseConcat tests

---
 .../python/kernel_tests/sparse_ops/sparse_concat_op_test.py      | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
index 79bcb9494c9..9d7ce912900 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
+++ b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
@@ -18,7 +18,6 @@ import numpy as np
 
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import errors
 from tensorflow.python.framework import sparse_tensor
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops

From f6b9a3235154af7eff554c862953d926b1a0beb2 Mon Sep 17 00:00:00 2001
From: Jaydeng75 <user.josephms@gmail.com>
Date: Mon, 5 Oct 2026 00:11:41 +0530
Subject: [PATCH 0182/1000] Use status-returning SparseConcat input shape
 construction

Replace direct input TensorShape constructors in Compute and the CPU functor with BuildTensorShape. Add an input-volume overflow regression that expects the existing early-validation error rather than the output-volume recomputation error.
---
 tensorflow/core/kernels/sparse_concat_op.cc   | 16 +++++++++++----
 .../sparse_ops/sparse_concat_op_test.py       | 20 +++++++++++++++++++
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/kernels/sparse_concat_op.cc b/tensorflow/core/kernels/sparse_concat_op.cc
index e6072768212..d1b591edb2e 100644
--- a/tensorflow/core/kernels/sparse_concat_op.cc
+++ b/tensorflow/core/kernels/sparse_concat_op.cc
@@ -46,7 +46,9 @@ struct SparseConcatFunctor<CPUDevice, T> {
                   const OpInputList& vals, const OpInputList& shapes,
                   int concat_dim) {
     const int N = inds.size();
-    const TensorShape input_shape(shapes[0].vec<int64_t>());
+    TensorShape input_shape;
+    OP_REQUIRES_OK(context, TensorShape::BuildTensorShape(
+                                shapes[0].vec<int64_t>(), &input_shape));
     const int input_rank = input_shape.dims();
 
     // The input and output sparse tensors are assumed to be ordered along
@@ -72,7 +74,9 @@ struct SparseConcatFunctor<CPUDevice, T> {
     std::vector<sparse::SparseTensor> sp_inputs;
     sp_inputs.reserve(N);
     for (int i = 0; i < N; ++i) {
-      const TensorShape current_shape(shapes[i].vec<int64_t>());
+      TensorShape current_shape;
+      OP_REQUIRES_OK(context, TensorShape::BuildTensorShape(
+                                  shapes[i].vec<int64_t>(), &current_shape));
       sparse::SparseTensor tensor;
       OP_REQUIRES_OK(context,
                      sparse::SparseTensor::Create(
@@ -175,7 +179,9 @@ class SparseConcatOp : public OpKernel {
                 absl::InvalidArgumentError(
                     "Encountered overflow from large input shape."));
 
-    const TensorShape input_shape(shapes[0].vec<int64_t>());
+    TensorShape input_shape;
+    OP_REQUIRES_OK(context, TensorShape::BuildTensorShape(
+                                shapes[0].vec<int64_t>(), &input_shape));
     const int input_rank = input_shape.dims();
     const int concat_dim = (concat_dim_attr_ < 0)
                                ? input_rank + concat_dim_attr_
@@ -190,7 +196,9 @@ class SparseConcatOp : public OpKernel {
     // then use SetDimWithStatus so RecomputeNumElements can catch
     // output volume overflow from the enlarged shape.
     for (int i = 1; i < N; ++i) {
-      const TensorShape current_shape(shapes[i].vec<int64_t>());
+      TensorShape current_shape;
+      OP_REQUIRES_OK(context, TensorShape::BuildTensorShape(
+                                  shapes[i].vec<int64_t>(), &current_shape));
       OP_REQUIRES(
           context, current_shape.dims() == input_rank,
           absl::InvalidArgumentError(absl::StrCat(
diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
index 9d7ce912900..7e11d51e7af 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
+++ b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
@@ -521,6 +521,26 @@ class SparseConcatTest(test.TestCase):
         concat_op = sparse_ops.sparse_concat(axis=1, sp_inputs=[sp1, sp2])
         self.evaluate(concat_op)
 
+  def testSparseConcatInputVolumeOverflowCrash(self):
+    # Each input volume already exceeds int64_max. The kernel must reject it
+    # with an OpError rather than abort during TensorShape construction.
+    int64_max = 9223372036854775807
+    dim_x = int64_max // 11
+
+    indices1 = constant_op.constant([[0, 0, 0]], dtype=dtypes.int64)
+    values1 = constant_op.constant([1.0], dtype=dtypes.float32)
+    shape1 = constant_op.constant([1, dim_x, dim_x], dtype=dtypes.int64)
+
+    indices2 = constant_op.constant([[0, 0, 0]], dtype=dtypes.int64)
+    values2 = constant_op.constant([2.0], dtype=dtypes.float32)
+    shape2 = constant_op.constant([1, dim_x, dim_x], dtype=dtypes.int64)
+
+    sp1 = sparse_tensor.SparseTensor(indices1, values1, shape1)
+    sp2 = sparse_tensor.SparseTensor(indices2, values2, shape2)
+
+    with self.assertRaisesOpError("Encountered overflow from large input shape"):
+      self.evaluate(sparse_ops.sparse_concat(axis=1, sp_inputs=[sp1, sp2]))
+
   def testSparseConcatOutputVolumeOverflow(self):
     # Individual volumes fit in int64, but the concatenated
     # output volume overflows int64, triggering

From b434f797888773e5203319457edc8b959b0bb6e4 Mon Sep 17 00:00:00 2001
From: Jaydeng75 <user.josephms@gmail.com>
Date: Mon, 5 Oct 2026 00:32:36 +0530
Subject: [PATCH 0183/1000] Rely on SparseConcat native shape validation and
 fix lint

Remove the redundant volume multiplication pass and preserve zero-volume validation when building each input shape. Match the native input-overflow error and cover invalid first and subsequent inputs. Add the direct dtypes dependency; test_util is already supplied by test_lib. Validate with the CI-pinned PyLint 2.13.9 and configuration.
---
 tensorflow/core/kernels/sparse_concat_op.cc   | 45 +++++++------------
 .../python/kernel_tests/sparse_ops/BUILD      |  1 +
 .../sparse_ops/sparse_concat_op_test.py       | 42 +++++++++--------
 3 files changed, 39 insertions(+), 49 deletions(-)

diff --git a/tensorflow/core/kernels/sparse_concat_op.cc b/tensorflow/core/kernels/sparse_concat_op.cc
index d1b591edb2e..d2368ad5c29 100644
--- a/tensorflow/core/kernels/sparse_concat_op.cc
+++ b/tensorflow/core/kernels/sparse_concat_op.cc
@@ -137,9 +137,7 @@ class SparseConcatOp : public OpKernel {
     OP_REQUIRES(context, shapes.size() == N,
                 absl::InvalidArgumentError(absl::StrCat(
                     "Expected ", N, " input shapes, got ", shapes.size())));
-    bool overflow_occurred = false;
     for (int i = 0; i < N; i++) {
-      int64_t new_num_elements = 1;
       OP_REQUIRES(context, TensorShapeUtils::IsVector(shapes[i].shape()),
                   absl::InvalidArgumentError(absl::StrCat(
                       "Input shapes should be a vector but received shape ",
@@ -149,39 +147,19 @@ class SparseConcatOp : public OpKernel {
                       "Indices rank and shape rank must match: indices = ",
                       inds[i].dim_size(1), ", shape = ", shapes[i].dim_size(0),
                       " at position ", i)));
-      auto input_shape_vector = shapes[i].vec<int64_t>();
-      for (int j = 0; j < input_shape_vector.size(); j++) {
-        OP_REQUIRES(context, input_shape_vector(j) >= 0,
-                    absl::InvalidArgumentError(absl::StrCat(
-                        "Input shape dimensions must be non-negative, got ",
-                        input_shape_vector(j), " for dimension ", j,
-                        " at position ", i)));
-        new_num_elements =
-            MultiplyWithoutOverflow(new_num_elements, input_shape_vector(j));
-        if (new_num_elements < 0) {
-          overflow_occurred = true;
-          break;
-        }
-      }
-
-      if (overflow_occurred) {
-        break;
-      }
-      OP_REQUIRES(context, new_num_elements > 0 || inds[i].dim_size(0) == 0,
-                  absl::InvalidArgumentError(absl::StrCat(
-                      "SparseTensor cannot have non-zero indices if dense "
-                      "shape has zero volume: dense_shape = ",
-                      shapes[i].SummarizeValue(10), " with nnz = ",
-                      inds[i].dim_size(0), " at position ", i)));
     }
 
-    OP_REQUIRES(context, !overflow_occurred,
-                absl::InvalidArgumentError(
-                    "Encountered overflow from large input shape."));
-
     TensorShape input_shape;
     OP_REQUIRES_OK(context, TensorShape::BuildTensorShape(
                                 shapes[0].vec<int64_t>(), &input_shape));
+    OP_REQUIRES(context,
+                input_shape.num_elements() > 0 || inds[0].dim_size(0) == 0,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "SparseTensor cannot have non-zero indices if dense "
+                    "shape has zero volume: dense_shape = ",
+                    shapes[0].SummarizeValue(10),
+                    " with nnz = ", inds[0].dim_size(0), " at position 0")));
+
     const int input_rank = input_shape.dims();
     const int concat_dim = (concat_dim_attr_ < 0)
                                ? input_rank + concat_dim_attr_
@@ -199,6 +177,13 @@ class SparseConcatOp : public OpKernel {
       TensorShape current_shape;
       OP_REQUIRES_OK(context, TensorShape::BuildTensorShape(
                                   shapes[i].vec<int64_t>(), &current_shape));
+      OP_REQUIRES(
+          context, current_shape.num_elements() > 0 || inds[i].dim_size(0) == 0,
+          absl::InvalidArgumentError(absl::StrCat(
+              "SparseTensor cannot have non-zero indices if dense "
+              "shape has zero volume: dense_shape = ",
+              shapes[i].SummarizeValue(10), " with nnz = ", inds[i].dim_size(0),
+              " at position ", i)));
       OP_REQUIRES(
           context, current_shape.dims() == input_rank,
           absl::InvalidArgumentError(absl::StrCat(
diff --git a/tensorflow/python/kernel_tests/sparse_ops/BUILD b/tensorflow/python/kernel_tests/sparse_ops/BUILD
index 5ef3878728c..f01fe53d53e 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/BUILD
+++ b/tensorflow/python/kernel_tests/sparse_ops/BUILD
@@ -50,6 +50,7 @@ tf_py_strict_test(
     srcs = ["sparse_concat_op_test.py"],
     deps = [
         "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/framework:for_generated_wrappers",
         "//tensorflow/python/framework:sparse_tensor",
         "//tensorflow/python/framework:test_lib",
diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
index 7e11d51e7af..700b8c81aa5 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
+++ b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
@@ -472,18 +472,20 @@ class SparseConcatTest(test.TestCase):
   @test_util.run_deprecated_v1
   def testDivisionByZeroMalformedShape(self):
     with self.session():
-      sp1 = sparse_tensor.SparseTensor(
-          indices=[[0, 0]], values=[1.0], dense_shape=[0, 1]
-      )
-      sp2 = sparse_tensor.SparseTensor(
-          indices=[[0, 0]], values=[2.0], dense_shape=[0, 1]
-      )
-      concat_op = sparse_ops.sparse_concat(axis=1, sp_inputs=[sp1, sp2])
-      with self.assertRaisesOpError(
-          "SparseTensor cannot have non-zero indices if dense shape has zero"
-          " volume"
-      ):
-        self.evaluate(concat_op)
+      for shapes in (([0, 1], [0, 1]), ([1, 1], [1, 0])):
+        with self.subTest(shapes=shapes):
+          sp1 = sparse_tensor.SparseTensor(
+              indices=[[0, 0]], values=[1.0], dense_shape=shapes[0]
+          )
+          sp2 = sparse_tensor.SparseTensor(
+              indices=[[0, 0]], values=[2.0], dense_shape=shapes[1]
+          )
+          concat_op = sparse_ops.sparse_concat(axis=1, sp_inputs=[sp1, sp2])
+          with self.assertRaisesOpError(
+              "SparseTensor cannot have non-zero indices if dense shape has"
+              " zero volume"
+          ):
+            self.evaluate(concat_op)
 
   def testLargeShapeOverflow(self):
     shape1 = constant_op.constant([2305843009213693952], dtype=dtypes.int64) * 2
@@ -522,24 +524,26 @@ class SparseConcatTest(test.TestCase):
         self.evaluate(concat_op)
 
   def testSparseConcatInputVolumeOverflowCrash(self):
-    # Each input volume already exceeds int64_max. The kernel must reject it
-    # with an OpError rather than abort during TensorShape construction.
+    # Reject overflowing volumes in both the first and subsequent inputs with
+    # an OpError rather than abort during TensorShape construction.
     int64_max = 9223372036854775807
     dim_x = int64_max // 11
 
     indices1 = constant_op.constant([[0, 0, 0]], dtype=dtypes.int64)
     values1 = constant_op.constant([1.0], dtype=dtypes.float32)
-    shape1 = constant_op.constant([1, dim_x, dim_x], dtype=dtypes.int64)
 
     indices2 = constant_op.constant([[0, 0, 0]], dtype=dtypes.int64)
     values2 = constant_op.constant([2.0], dtype=dtypes.float32)
     shape2 = constant_op.constant([1, dim_x, dim_x], dtype=dtypes.int64)
 
-    sp1 = sparse_tensor.SparseTensor(indices1, values1, shape1)
-    sp2 = sparse_tensor.SparseTensor(indices2, values2, shape2)
+    for first_shape in ([1, dim_x, dim_x], [1, 1, dim_x]):
+      with self.subTest(first_shape=first_shape):
+        shape1 = constant_op.constant(first_shape, dtype=dtypes.int64)
+        sp1 = sparse_tensor.SparseTensor(indices1, values1, shape1)
+        sp2 = sparse_tensor.SparseTensor(indices2, values2, shape2)
 
-    with self.assertRaisesOpError("Encountered overflow from large input shape"):
-      self.evaluate(sparse_ops.sparse_concat(axis=1, sp_inputs=[sp1, sp2]))
+        with self.assertRaisesOpError("Encountered overflow when multiplying"):
+          self.evaluate(sparse_ops.sparse_concat(axis=1, sp_inputs=[sp1, sp2]))
 
   def testSparseConcatOutputVolumeOverflow(self):
     # Individual volumes fit in int64, but the concatenated

From fb1591b095dc609908097fa93d7c9ac4d3b73fe7 Mon Sep 17 00:00:00 2001
From: Jaydeng75 <user.josephms@gmail.com>
Date: Mon, 5 Oct 2026 00:49:01 +0530
Subject: [PATCH 0184/1000] Bound SparseConcat entries and share validated
 input shapes

Reject output entry counts above INT_MAX before SparseTensor::Concat can truncate its int accumulator. Pass the TensorShape objects validated by Compute to both CPU and GPU functors through a shared span, avoiding repeated parsing. Add the direct span dependency.
---
 tensorflow/core/kernels/BUILD                 |  1 +
 tensorflow/core/kernels/sparse_concat_op.cc   | 29 ++++++++++++-------
 tensorflow/core/kernels/sparse_concat_op.h    |  5 +++-
 .../core/kernels/sparse_concat_op_gpu.cu.cc   |  7 +++--
 4 files changed, 28 insertions(+), 14 deletions(-)

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index f98ecdcfe37..95f593658bf 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -5275,6 +5275,7 @@ tf_kernel_library(
         "//tensorflow/core/util:overflow",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
     ] + if_cuda_or_rocm([
         ":gpu_device_array",
         ":gpu_prim_helpers",
diff --git a/tensorflow/core/kernels/sparse_concat_op.cc b/tensorflow/core/kernels/sparse_concat_op.cc
index d2368ad5c29..ff693e416e4 100644
--- a/tensorflow/core/kernels/sparse_concat_op.cc
+++ b/tensorflow/core/kernels/sparse_concat_op.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include "tensorflow/core/kernels/sparse_concat_op.h"
 
+#include <limits>
 #include <numeric>
 #include <utility>
 #include <vector>
@@ -43,12 +44,11 @@ namespace functor {
 template <typename T>
 struct SparseConcatFunctor<CPUDevice, T> {
   void operator()(OpKernelContext* context, const OpInputList& inds,
-                  const OpInputList& vals, const OpInputList& shapes,
+                  const OpInputList& vals,
+                  const absl::Span<const TensorShape>& input_shapes,
                   int concat_dim) {
     const int N = inds.size();
-    TensorShape input_shape;
-    OP_REQUIRES_OK(context, TensorShape::BuildTensorShape(
-                                shapes[0].vec<int64_t>(), &input_shape));
+    const TensorShape& input_shape = input_shapes[0];
     const int input_rank = input_shape.dims();
 
     // The input and output sparse tensors are assumed to be ordered along
@@ -74,9 +74,7 @@ struct SparseConcatFunctor<CPUDevice, T> {
     std::vector<sparse::SparseTensor> sp_inputs;
     sp_inputs.reserve(N);
     for (int i = 0; i < N; ++i) {
-      TensorShape current_shape;
-      OP_REQUIRES_OK(context, TensorShape::BuildTensorShape(
-                                  shapes[i].vec<int64_t>(), &current_shape));
+      const TensorShape& current_shape = input_shapes[i];
       sparse::SparseTensor tensor;
       OP_REQUIRES_OK(context,
                      sparse::SparseTensor::Create(
@@ -149,9 +147,13 @@ class SparseConcatOp : public OpKernel {
                       " at position ", i)));
     }
 
+    absl::InlinedVector<TensorShape, 8> input_shapes;
+    input_shapes.reserve(N);
+
     TensorShape input_shape;
     OP_REQUIRES_OK(context, TensorShape::BuildTensorShape(
                                 shapes[0].vec<int64_t>(), &input_shape));
+    input_shapes.push_back(input_shape);
     OP_REQUIRES(context,
                 input_shape.num_elements() > 0 || inds[0].dim_size(0) == 0,
                 absl::InvalidArgumentError(absl::StrCat(
@@ -177,6 +179,7 @@ class SparseConcatOp : public OpKernel {
       TensorShape current_shape;
       OP_REQUIRES_OK(context, TensorShape::BuildTensorShape(
                                   shapes[i].vec<int64_t>(), &current_shape));
+      input_shapes.push_back(current_shape);
       OP_REQUIRES(
           context, current_shape.num_elements() > 0 || inds[i].dim_size(0) == 0,
           absl::InvalidArgumentError(absl::StrCat(
@@ -224,8 +227,14 @@ class SparseConcatOp : public OpKernel {
     for (int i = 0; i < N; ++i) {
       int64_t next_output_nnz =
           AddWithoutOverflow(output_nnz, inds[i].dim_size(0));
-      OP_REQUIRES(context, next_output_nnz >= 0,
-                  absl::InvalidArgumentError("output_nnz overflowed"));
+      // SparseTensor::Concat accumulates num_entries as an int. Prevent
+      // wrapping before it constructs the output TensorShape.
+      OP_REQUIRES(
+          context,
+          next_output_nnz >= 0 &&
+              next_output_nnz <= std::numeric_limits<int>::max(),
+          absl::InvalidArgumentError(
+              "output_nnz overflowed or exceeded supported INT_MAX limit"));
       output_nnz = next_output_nnz;
     }
 
@@ -240,7 +249,7 @@ class SparseConcatOp : public OpKernel {
       return;  // No work to do
     }
 
-    functor::SparseConcatFunctor<Device, T>()(context, inds, vals, shapes,
+    functor::SparseConcatFunctor<Device, T>()(context, inds, vals, input_shapes,
                                               concat_dim);
   }
 
diff --git a/tensorflow/core/kernels/sparse_concat_op.h b/tensorflow/core/kernels/sparse_concat_op.h
index c13ae5021e0..38d86349125 100644
--- a/tensorflow/core/kernels/sparse_concat_op.h
+++ b/tensorflow/core/kernels/sparse_concat_op.h
@@ -16,7 +16,9 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_KERNELS_SPARSE_CONCAT_OP_H_
 #define TENSORFLOW_CORE_KERNELS_SPARSE_CONCAT_OP_H_
 
+#include "absl/types/span.h"
 #include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 
 namespace tensorflow {
 
@@ -25,7 +27,8 @@ namespace functor {
 template <typename Device, typename T>
 struct SparseConcatFunctor {
   void operator()(OpKernelContext* context, const OpInputList& inds,
-                  const OpInputList& vals, const OpInputList& shapes,
+                  const OpInputList& vals,
+                  const absl::Span<const TensorShape>& input_shapes,
                   int concat_dim);
 };
 
diff --git a/tensorflow/core/kernels/sparse_concat_op_gpu.cu.cc b/tensorflow/core/kernels/sparse_concat_op_gpu.cu.cc
index 5dab01c21c1..a78faecbe27 100644
--- a/tensorflow/core/kernels/sparse_concat_op_gpu.cu.cc
+++ b/tensorflow/core/kernels/sparse_concat_op_gpu.cu.cc
@@ -118,10 +118,11 @@ __global__ void SparseConcatPermuteKernel(
 template <typename T>
 struct SparseConcatFunctor<GPUDevice, T> {
   void operator()(OpKernelContext* context, const OpInputList& inds,
-                  const OpInputList& vals, const OpInputList& shapes,
+                  const OpInputList& vals,
+                  const absl::Span<const TensorShape>& input_shapes,
                   int concat_dim) {
     const int N = inds.size();
-    const TensorShape input_shape0(shapes[0].vec<int64_t>());
+    const TensorShape& input_shape0 = input_shapes[0];
     const int rank = input_shape0.dims();
 
     // The input non-zeros are assumed to be sorted by increasing dimension
@@ -152,7 +153,7 @@ struct SparseConcatFunctor<GPUDevice, T> {
                   absl::InvalidArgumentError("nnz overflowed"));
       nnz_sum = next_nnz_sum;
       nnz_scan.Set(i + 1, nnz_sum);
-      const TensorShape current_shape(shapes[i].vec<int64_t>());
+      const TensorShape& current_shape = input_shapes[i];
       int64_t next_concat_size_sum = AddWithoutOverflow(
           concat_size_sum, current_shape.dim_size(concat_dim));
       OP_REQUIRES(context, next_concat_size_sum >= 0,

From 1f4f602d58f4c970f9dbbdec49f86066b47dce75 Mon Sep 17 00:00:00 2001
From: Jaydeng75 <user.josephms@gmail.com>
Date: Mon, 5 Oct 2026 00:53:55 +0530
Subject: [PATCH 0185/1000] Match GPU SparseConcat input overflow test to
 shared validation

---
 .../python/kernel_tests/sparse_ops/sparse_concat_op_test.py   | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
index 700b8c81aa5..c8ed2219205 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
+++ b/tensorflow/python/kernel_tests/sparse_ops/sparse_concat_op_test.py
@@ -500,7 +500,7 @@ class SparseConcatTest(test.TestCase):
       concat_op = sparse_ops.sparse_concat(axis=0, sp_inputs=[sp1, sp2])
       self.evaluate(concat_op)
 
-  def testGPUOutputDenseElementsOverflow(self):
+  def testGPUInputVolumeOverflow(self):
     if not test.is_gpu_available():
       self.skipTest("No GPU available to run GPU functor check.")
 
@@ -519,7 +519,7 @@ class SparseConcatTest(test.TestCase):
       sp2 = sparse_tensor.SparseTensor(
           indices=[[0, 0, 0]], values=[1.0], dense_shape=shape2
       )
-      with self.assertRaisesOpError("dense_elements overflowed"):
+      with self.assertRaisesOpError("Encountered overflow when multiplying"):
         concat_op = sparse_ops.sparse_concat(axis=1, sp_inputs=[sp1, sp2])
         self.evaluate(concat_op)
 

From 7116c217d719cf029a2167a1bba37dc82e1e4c47 Mon Sep 17 00:00:00 2001
From: endorphin13 <elshad.toklayev.business@gmail.com>
Date: Mon, 5 Oct 2026 12:24:49 +0300
Subject: [PATCH 0186/1000] Do not dereference a null TensorBuffer for an empty
 snapshot tensor

A zero dimension passes TensorShape::BuildTensorShape, since zero is a valid
dimension, but Tensor only allocates a buffer when the shape has elements, so
DMAHelper::buffer returns null for a tensor whose metadata declares a zero
dimension. Reading iov_base and iov_len then dereferenced null: data() reads
through the null pointer and size() is a pure virtual call, so this crashed
rather than reporting a corrupt snapshot.

Take the buffer address and length only when a buffer exists. The read then
continues to the existing size check, which reports the mismatch between the
metadata and the compressed payload as an error.

Also add the direct include for TensorShape, and wrap the BuildTensorShape
call at 80 columns.
---
 tensorflow/core/data/snapshot_utils.cc | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/data/snapshot_utils.cc b/tensorflow/core/data/snapshot_utils.cc
index 748faf51db6..33de2b705eb 100644
--- a/tensorflow/core/data/snapshot_utils.cc
+++ b/tensorflow/core/data/snapshot_utils.cc
@@ -39,6 +39,7 @@ limitations under the License.
 #include "tensorflow/core/framework/graph.pb.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor.pb.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/lib/io/buffered_inputstream.h"
 #include "tensorflow/core/lib/io/random_inputstream.h"
 #include "tensorflow/core/lib/io/record_writer.h"
@@ -971,12 +972,15 @@ absl::Status CustomReader::SnappyUncompress(
       // CHECK-fails on a negative or overflowing dimension, which aborts the
       // process instead of reporting a corrupt snapshot.
       TensorShape shape;
-      TF_RETURN_IF_ERROR(
-          TensorShape::BuildTensorShape(tensor_metadata.tensor_shape(), &shape));
+      TF_RETURN_IF_ERROR(TensorShape::BuildTensorShape(
+          tensor_metadata.tensor_shape(), &shape));
       Tensor simple_tensor(dtypes_[i], shape);
+      // A zero dimension is a valid shape, but Tensor does not allocate a
+      // buffer for a tensor with no elements, so `buffer` is null in that case
+      // and must not be dereferenced.
       TensorBuffer* buffer = DMAHelper::buffer(&simple_tensor);
-      iov[index].iov_base = buffer->data();
-      iov[index].iov_len = buffer->size();
+      iov[index].iov_base = buffer ? buffer->data() : nullptr;
+      iov[index].iov_len = buffer ? buffer->size() : 0;
       simple_tensors->push_back(std::move(simple_tensor));
     } else {
       int64_t tensor_size = tensor_metadata.tensor_size_bytes();

From f9ccb1cf112f384339382c685b8df601f8e0cb5a Mon Sep 17 00:00:00 2001
From: endorphin13 <elshad.toklayev.business@gmail.com>
Date: Mon, 5 Oct 2026 12:26:56 +0300
Subject: [PATCH 0187/1000] Add a zero-dimension regression test and assert the
 exact status codes

Extract the snapshot-writing and metadata-patching setup into a helper so the
negative-dimension test and the new zero-dimension test share it, instead of
duplicating it. The helper takes the dimension to write, and is called through
ASSERT_NO_FATAL_FAILURE so a failed assertion inside it stops the test.

SnappyZeroTensorDimensionIsAnError covers the null-buffer case: zero is a valid
dimension, so the shape is built and the tensor holds no buffer. Before the
guard the read dereferenced that null buffer; now it reports the mismatch
between the metadata and the compressed payload.

Assert the specific status codes rather than only that the read failed, so
neither test can pass because of an unrelated failure.

Also add the direct include for the snapshot proto, and wrap two Reader::Create
calls at 80 columns.
---
 tensorflow/core/data/snapshot_utils_test.cc | 78 +++++++++++++++------
 1 file changed, 56 insertions(+), 22 deletions(-)

diff --git a/tensorflow/core/data/snapshot_utils_test.cc b/tensorflow/core/data/snapshot_utils_test.cc
index 8df36d155af..6b03cf332eb 100644
--- a/tensorflow/core/data/snapshot_utils_test.cc
+++ b/tensorflow/core/data/snapshot_utils_test.cc
@@ -30,6 +30,7 @@ limitations under the License.
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/platform/test_benchmark.h"
+#include "tensorflow/core/protobuf/snapshot.pb.h"
 
 namespace tensorflow {
 namespace data {
@@ -152,37 +153,39 @@ TEST(SnapshotUtilTest, SnappyTensorSizeMismatch) {
   TF_ASSERT_OK(Env::Default()->DeleteFile(filename));
 }
 
-TEST(SnapshotUtilTest, SnappyInvalidTensorShapeIsAnError) {
-  // A corrupt snapshot whose metadata carries a negative dimension must be
-  // reported as an error. Building a TensorShape directly from the proto
-  // CHECK-fails instead, which aborts the process.
-  //
-  // Uses a memcpy-able dtype so the tensors take the simple-tensor path in
-  // CustomReader::SnappyUncompress, which is where the shape is built.
+// Writes a snappy snapshot of two int64 tensors, then rewrites the first
+// dimension of the first tensor's metadata to `patched_dim_size`.
+//
+// int64 is memcpy-able, so the tensors take the simple-tensor path in
+// CustomReader::SnappyUncompress, which is where the shape is built. Records
+// are an 8-byte little-endian length followed by the payload, and the metadata
+// is the first record written for a snappy snapshot, so it can be rewritten
+// without disturbing the compressed block that follows. The tensor count is
+// left unchanged so the read reaches the shape instead of failing the earlier
+// count check.
+void WriteSnapshotWithPatchedFirstDim(int64_t patched_dim_size,
+                                      tensorflow::DataTypeVector* dtypes,
+                                      std::string* filename) {
   std::vector<Tensor> tensors;
-  tensorflow::DataTypeVector dtypes;
   for (int i = 0; i < 2; ++i) {
     Tensor t(DT_INT64, TensorShape({4}));
     t.flat<int64_t>().setZero();
-    dtypes.push_back(DT_INT64);
+    dtypes->push_back(DT_INT64);
     tensors.push_back(t);
   }
 
-  std::string filename;
-  EXPECT_TRUE(Env::Default()->LocalTempFilename(&filename));
+  ASSERT_TRUE(Env::Default()->LocalTempFilename(filename));
 
   std::unique_ptr<Writer> writer;
-  TF_ASSERT_OK(Writer::Create(Env::Default(), filename, io::compression::kSnappy,
-                              /*version=*/1, dtypes, &writer));
+  TF_ASSERT_OK(Writer::Create(Env::Default(), *filename,
+                              io::compression::kSnappy, /*version=*/1, *dtypes,
+                              &writer));
   TF_ASSERT_OK(writer->WriteTensors(tensors));
   TF_ASSERT_OK(writer->Close());
 
-  // Records are an 8-byte little-endian length followed by the payload; the
-  // metadata is the first record written for a snappy snapshot. Rewrite it
-  // with a negative dimension.
   constexpr size_t kHeaderSize = sizeof(uint64_t);
   std::string contents;
-  TF_ASSERT_OK(ReadFileToString(Env::Default(), filename, &contents));
+  TF_ASSERT_OK(ReadFileToString(Env::Default(), *filename, &contents));
   ASSERT_GT(contents.size(), kHeaderSize);
   const uint64_t metadata_size = core::DecodeFixed64(contents.data());
   ASSERT_GE(contents.size(), kHeaderSize + metadata_size);
@@ -195,21 +198,52 @@ TEST(SnapshotUtilTest, SnappyInvalidTensorShapeIsAnError) {
   metadata.mutable_tensor_metadata(0)
       ->mutable_tensor_shape()
       ->mutable_dim(0)
-      ->set_size(-1);
+      ->set_size(patched_dim_size);
 
   const std::string patched = metadata.SerializeAsString();
   char header[kHeaderSize];
   core::EncodeFixed64(header, patched.size());
   TF_ASSERT_OK(WriteStringToFile(
-      Env::Default(), filename,
+      Env::Default(), *filename,
       std::string(header, kHeaderSize) + patched +
           contents.substr(kHeaderSize + metadata_size)));
+}
+
+TEST(SnapshotUtilTest, SnappyInvalidTensorShapeIsAnError) {
+  // A corrupt snapshot whose metadata carries a negative dimension must be
+  // reported as an error. Building a TensorShape directly from the proto
+  // CHECK-fails instead, which aborts the process.
+  std::string filename;
+  tensorflow::DataTypeVector dtypes;
+  ASSERT_NO_FATAL_FAILURE(WriteSnapshotWithPatchedFirstDim(
+      /*patched_dim_size=*/-1, &dtypes, &filename));
 
   std::unique_ptr<Reader> reader;
-  TF_ASSERT_OK(Reader::Create(Env::Default(), filename, io::compression::kSnappy,
-                              /*version=*/1, dtypes, &reader));
+  TF_ASSERT_OK(Reader::Create(Env::Default(), filename,
+                              io::compression::kSnappy, /*version=*/1, dtypes,
+                              &reader));
   std::vector<Tensor> read_tensors;
-  EXPECT_FALSE(reader->ReadTensors(&read_tensors).ok());
+  EXPECT_TRUE(absl::IsInvalidArgument(reader->ReadTensors(&read_tensors)));
+
+  TF_ASSERT_OK(Env::Default()->DeleteFile(filename));
+}
+
+TEST(SnapshotUtilTest, SnappyZeroTensorDimensionIsAnError) {
+  // Zero is a valid dimension, so the shape is built successfully, but the
+  // tensor then holds no buffer. The read must report the resulting mismatch
+  // between the metadata and the compressed payload rather than dereference
+  // that null buffer.
+  std::string filename;
+  tensorflow::DataTypeVector dtypes;
+  ASSERT_NO_FATAL_FAILURE(WriteSnapshotWithPatchedFirstDim(
+      /*patched_dim_size=*/0, &dtypes, &filename));
+
+  std::unique_ptr<Reader> reader;
+  TF_ASSERT_OK(Reader::Create(Env::Default(), filename,
+                              io::compression::kSnappy, /*version=*/1, dtypes,
+                              &reader));
+  std::vector<Tensor> read_tensors;
+  EXPECT_TRUE(absl::IsInternal(reader->ReadTensors(&read_tensors)));
 
   TF_ASSERT_OK(Env::Default()->DeleteFile(filename));
 }

From ff2b4ff6dcdabad9c705dc912805020842233c06 Mon Sep 17 00:00:00 2001
From: Daniil Kuts <kutz@ispras.ru>
Date: Mon, 5 Oct 2026 13:48:16 +0300
Subject: [PATCH 0188/1000] Fix possible null dereference in RootDataset

---
 tensorflow/core/data/root_dataset.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/data/root_dataset.cc b/tensorflow/core/data/root_dataset.cc
index 4a7abfb64df..5c32f1b1af6 100644
--- a/tensorflow/core/data/root_dataset.cc
+++ b/tensorflow/core/data/root_dataset.cc
@@ -412,8 +412,8 @@ RootDataset::RootDataset(core::RefCountPtr<DatasetBase> input,
   random_indexing_compatible_ = absl::OkStatus();
   if (input_ != nullptr) {
     random_indexing_compatible_ = input_->RandomIndexingCompatible();
+    AddTraceMetadata(params_, input_->options(), &traceme_metadata_);
   }
-  AddTraceMetadata(params_, input_->options(), &traceme_metadata_);
 }
 
 RootDataset::~RootDataset() = default;

From 7242033d7050be16d4c07c257fc522d0f0d0f65b Mon Sep 17 00:00:00 2001
From: endorphin13 <elshad.toklayev.business@gmail.com>
Date: Mon, 5 Oct 2026 14:17:07 +0300
Subject: [PATCH 0189/1000] Do not dereference a null TensorBuffer when writing
 an empty tensor

CustomWriter::WriteTensors has the same defect the reader had: a tensor with
no elements holds no buffer, so DMAHelper::buffer returns null and
tensor_buffer->size() crashed. Take the length only when a buffer exists.

Guard the copy as well. For an empty tensor the source pointer is null, and
passing a null pointer to memcpy is undefined even when the length is zero.

This makes a snapshot containing an empty tensor writable, not just readable,
which the round-trip test added separately now covers.
---
 tensorflow/core/data/snapshot_utils.cc | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/data/snapshot_utils.cc b/tensorflow/core/data/snapshot_utils.cc
index 33de2b705eb..b769ef94df1 100644
--- a/tensorflow/core/data/snapshot_utils.cc
+++ b/tensorflow/core/data/snapshot_utils.cc
@@ -276,9 +276,10 @@ absl::Status CustomWriter::WriteTensors(const std::vector<Tensor>& tensors) {
     tensor.shape().AsProto(tensor_metadata->mutable_tensor_shape());
     int64_t size = 0;
     if (simple_tensor_mask_[i]) {
+      // As in the reader below, a tensor with no elements has no buffer.
       auto tensor_buffer = DMAHelper::buffer(&tensor);
       tensor_buffers.push_back(tensor_buffer);
-      size = tensor_buffer->size();
+      size = tensor_buffer != nullptr ? tensor_buffer->size() : 0;
     } else {
       TensorProto proto;
       tensor.AsProtoTensorContent(&proto);
@@ -296,8 +297,12 @@ absl::Status CustomWriter::WriteTensors(const std::vector<Tensor>& tensors) {
   for (int i = 0, end = tensors.size(); i < end; ++i) {
     const auto& tensor_metadata = metadata.tensor_metadata(i);
     if (simple_tensor_mask_[i]) {
-      memcpy(position, tensor_buffers[buffer_index]->data(),
-             tensor_metadata.tensor_size_bytes());
+      // Skip the copy for an empty tensor: its buffer is null, and passing a
+      // null pointer to memcpy is undefined even when the length is zero.
+      if (tensor_metadata.tensor_size_bytes() > 0) {
+        memcpy(position, tensor_buffers[buffer_index]->data(),
+               tensor_metadata.tensor_size_bytes());
+      }
       buffer_index++;
     } else {
       tensor_protos[proto_index].SerializeToArray(

From 7fd70e03f542965fc8ed11bc44830a0658a25e94 Mon Sep 17 00:00:00 2001
From: endorphin13 <elshad.toklayev.business@gmail.com>
Date: Mon, 5 Oct 2026 14:17:21 +0300
Subject: [PATCH 0190/1000] Allocate the tensor proto buffer without throwing

The ResourceExhausted branch below this allocation was unreachable:
make_unique uses a throwing new, so a huge tensor_size from a corrupt snapshot
terminated the process rather than returning the status the code already
prepared. These targets are built with -fno-exceptions, so the throw cannot be
caught either.

Allocate with nothrow new so the null check runs and the error reaches the
caller. <new> is already included.
---
 tensorflow/core/data/snapshot_utils.cc | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/data/snapshot_utils.cc b/tensorflow/core/data/snapshot_utils.cc
index b769ef94df1..3fb63fa77b8 100644
--- a/tensorflow/core/data/snapshot_utils.cc
+++ b/tensorflow/core/data/snapshot_utils.cc
@@ -993,8 +993,11 @@ absl::Status CustomReader::SnappyUncompress(
         return absl::InvalidArgumentError(
             absl::StrCat("Tensor size is negative: ", tensor_size));
       }
-      std::unique_ptr<char[]> tensor_proto_str =
-          std::make_unique<char[]>(tensor_size);
+      // Allocate without throwing so the check below is reachable: with
+      // make_unique a huge `tensor_size` from a corrupt snapshot would abort
+      // the process, since these targets are built with -fno-exceptions.
+      std::unique_ptr<char[]> tensor_proto_str(
+          new (std::nothrow) char[tensor_size]);
       if (tensor_proto_str == nullptr) {
         return absl::ResourceExhaustedError(absl::StrCat(
             "Failed to allocate memory for tensor of size ", tensor_size));

From 2e34344be6686899381e4849a8c2f8e6bcf85c3a Mon Sep 17 00:00:00 2001
From: endorphin13 <elshad.toklayev.business@gmail.com>
Date: Mon, 5 Oct 2026 14:17:35 +0300
Subject: [PATCH 0191/1000] Cover the overflow branch and an empty-tensor round
 trip

Generalise the helper to rewrite the whole dimension list rather than only the
first dimension, so a test can supply more than one dimension.

SnappyOverflowingTensorShapeIsAnError covers the other branch of
BuildTensorShape: each dimension is non-negative but their product does not fit
in int64. The TensorShape proto constructor would abort on it, and
AddDimWithStatus reports InvalidArgument instead.

SnappyZeroTensorDimensionRoundTrip writes and reads a real empty tensor, which
exercises the writer guards. Before them the write crashed on the null buffer,
so this case could not round-trip at all.
---
 tensorflow/core/data/snapshot_utils_test.cc | 76 +++++++++++++++++----
 1 file changed, 62 insertions(+), 14 deletions(-)

diff --git a/tensorflow/core/data/snapshot_utils_test.cc b/tensorflow/core/data/snapshot_utils_test.cc
index 6b03cf332eb..b2640eaa741 100644
--- a/tensorflow/core/data/snapshot_utils_test.cc
+++ b/tensorflow/core/data/snapshot_utils_test.cc
@@ -153,8 +153,8 @@ TEST(SnapshotUtilTest, SnappyTensorSizeMismatch) {
   TF_ASSERT_OK(Env::Default()->DeleteFile(filename));
 }
 
-// Writes a snappy snapshot of two int64 tensors, then rewrites the first
-// dimension of the first tensor's metadata to `patched_dim_size`.
+// Writes a snappy snapshot of two int64 tensors, then rewrites the dimensions
+// of the first tensor's metadata to `patched_dim_sizes`.
 //
 // int64 is memcpy-able, so the tensors take the simple-tensor path in
 // CustomReader::SnappyUncompress, which is where the shape is built. Records
@@ -163,9 +163,9 @@ TEST(SnapshotUtilTest, SnappyTensorSizeMismatch) {
 // without disturbing the compressed block that follows. The tensor count is
 // left unchanged so the read reaches the shape instead of failing the earlier
 // count check.
-void WriteSnapshotWithPatchedFirstDim(int64_t patched_dim_size,
-                                      tensorflow::DataTypeVector* dtypes,
-                                      std::string* filename) {
+void WriteSnapshotWithPatchedDims(const std::vector<int64_t>& patched_dim_sizes,
+                                  tensorflow::DataTypeVector* dtypes,
+                                  std::string* filename) {
   std::vector<Tensor> tensors;
   for (int i = 0; i < 2; ++i) {
     Tensor t(DT_INT64, TensorShape({4}));
@@ -194,11 +194,11 @@ void WriteSnapshotWithPatchedFirstDim(int64_t patched_dim_size,
   ASSERT_TRUE(
       metadata.ParseFromString(contents.substr(kHeaderSize, metadata_size)));
   ASSERT_GT(metadata.tensor_metadata_size(), 0);
-  ASSERT_GT(metadata.tensor_metadata(0).tensor_shape().dim_size(), 0);
-  metadata.mutable_tensor_metadata(0)
-      ->mutable_tensor_shape()
-      ->mutable_dim(0)
-      ->set_size(patched_dim_size);
+  auto* shape = metadata.mutable_tensor_metadata(0)->mutable_tensor_shape();
+  shape->clear_dim();
+  for (int64_t dim_size : patched_dim_sizes) {
+    shape->add_dim()->set_size(dim_size);
+  }
 
   const std::string patched = metadata.SerializeAsString();
   char header[kHeaderSize];
@@ -215,8 +215,8 @@ TEST(SnapshotUtilTest, SnappyInvalidTensorShapeIsAnError) {
   // CHECK-fails instead, which aborts the process.
   std::string filename;
   tensorflow::DataTypeVector dtypes;
-  ASSERT_NO_FATAL_FAILURE(WriteSnapshotWithPatchedFirstDim(
-      /*patched_dim_size=*/-1, &dtypes, &filename));
+  ASSERT_NO_FATAL_FAILURE(WriteSnapshotWithPatchedDims(
+      /*patched_dim_sizes=*/{-1}, &dtypes, &filename));
 
   std::unique_ptr<Reader> reader;
   TF_ASSERT_OK(Reader::Create(Env::Default(), filename,
@@ -235,8 +235,8 @@ TEST(SnapshotUtilTest, SnappyZeroTensorDimensionIsAnError) {
   // that null buffer.
   std::string filename;
   tensorflow::DataTypeVector dtypes;
-  ASSERT_NO_FATAL_FAILURE(WriteSnapshotWithPatchedFirstDim(
-      /*patched_dim_size=*/0, &dtypes, &filename));
+  ASSERT_NO_FATAL_FAILURE(WriteSnapshotWithPatchedDims(
+      /*patched_dim_sizes=*/{0}, &dtypes, &filename));
 
   std::unique_ptr<Reader> reader;
   TF_ASSERT_OK(Reader::Create(Env::Default(), filename,
@@ -248,6 +248,54 @@ TEST(SnapshotUtilTest, SnappyZeroTensorDimensionIsAnError) {
   TF_ASSERT_OK(Env::Default()->DeleteFile(filename));
 }
 
+TEST(SnapshotUtilTest, SnappyOverflowingTensorShapeIsAnError) {
+  // The other branch of BuildTensorShape: each dimension is non-negative, but
+  // their product does not fit in int64, which the TensorShape proto
+  // constructor would CHECK-fail on.
+  std::string filename;
+  tensorflow::DataTypeVector dtypes;
+  ASSERT_NO_FATAL_FAILURE(WriteSnapshotWithPatchedDims(
+      /*patched_dim_sizes=*/{1LL << 32, 1LL << 32}, &dtypes, &filename));
+
+  std::unique_ptr<Reader> reader;
+  TF_ASSERT_OK(Reader::Create(Env::Default(), filename,
+                              io::compression::kSnappy, /*version=*/1, dtypes,
+                              &reader));
+  std::vector<Tensor> read_tensors;
+  EXPECT_TRUE(absl::IsInvalidArgument(reader->ReadTensors(&read_tensors)));
+
+  TF_ASSERT_OK(Env::Default()->DeleteFile(filename));
+}
+
+TEST(SnapshotUtilTest, SnappyZeroTensorDimensionRoundTrip) {
+  // An empty tensor is legitimate rather than corrupt, so it must survive a
+  // write and a read. Both sides hold a null buffer for it.
+  std::string filename;
+  ASSERT_TRUE(Env::Default()->LocalTempFilename(&filename));
+
+  tensorflow::DataTypeVector dtypes = {DT_INT64};
+  std::vector<Tensor> tensors;
+  tensors.emplace_back(DT_INT64, TensorShape({0}));
+
+  std::unique_ptr<Writer> writer;
+  TF_ASSERT_OK(Writer::Create(Env::Default(), filename,
+                              io::compression::kSnappy, /*version=*/1, dtypes,
+                              &writer));
+  TF_ASSERT_OK(writer->WriteTensors(tensors));
+  TF_ASSERT_OK(writer->Close());
+
+  std::unique_ptr<Reader> reader;
+  TF_ASSERT_OK(Reader::Create(Env::Default(), filename,
+                              io::compression::kSnappy, /*version=*/1, dtypes,
+                              &reader));
+  std::vector<Tensor> read_tensors;
+  TF_ASSERT_OK(reader->ReadTensors(&read_tensors));
+  ASSERT_EQ(read_tensors.size(), 1);
+  EXPECT_EQ(read_tensors[0].NumElements(), 0);
+
+  TF_ASSERT_OK(Env::Default()->DeleteFile(filename));
+}
+
 void SnapshotReaderBenchmarkLoop(::testing::benchmark::State& state,
                                  std::string compression_type, int version) {
   tensorflow::DataTypeVector dtypes;

From d7fccd53711abbee319d6cf735f9e5cf1141270f Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 5 Oct 2026 17:31:58 +0530
Subject: [PATCH 0192/1000] Fix code formatting in remote.py documentation

---
 tensorflow/python/eager/remote.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/eager/remote.py b/tensorflow/python/eager/remote.py
index e785fca558b..387539c62bd 100644
--- a/tensorflow/python/eager/remote.py
+++ b/tensorflow/python/eager/remote.py
@@ -45,7 +45,7 @@ def connect_to_remote_host(remote_host=None, job_name="worker"):
 
   Using the default job_name of worker, you can schedule ops to run remotely as
   follows:
-```python
+  ```python
   # When eager execution is enabled, connect to the remote host.
   tf.config.experimental_connect_to_host("exampleaddr.com:9876")
 
@@ -55,7 +55,7 @@ def connect_to_remote_host(remote_host=None, job_name="worker"):
     x1 = array_ops.ones([2, 2])
     x2 = array_ops.ones([2, 2])
     y = math_ops.matmul(x1, x2)
-```
+  ```
 
   Args:
     remote_host: a single or a list the remote server addr in host-port format.
@@ -148,7 +148,7 @@ def connect_to_cluster(cluster_spec_or_resolver,
   For example, for a cluster set up for parameter server training, the following
   device filters might be specified:
 
-```python
+  ```python
   cdf = tf.config.experimental.ClusterDeviceFilters()
   # For any worker, only the devices on PS nodes and itself are visible
   for i in range(num_workers):
@@ -159,7 +159,7 @@ def connect_to_cluster(cluster_spec_or_resolver,
 
   tf.config.experimental_connect_to_cluster(cluster_def,
                                             cluster_device_filters=cdf)
-```
+  ```
 
   Args:
     cluster_spec_or_resolver: A `ClusterSpec` or `ClusterResolver` describing

From 6818b9bc95143f526a531fc13e1ebb631c20f47a Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 5 Oct 2026 17:33:30 +0530
Subject: [PATCH 0193/1000] Add tests for connecting to cluster with master
 validation

---
 tensorflow/python/eager/remote_test.py | 90 ++++++++++++++++++++++++++
 1 file changed, 90 insertions(+)

diff --git a/tensorflow/python/eager/remote_test.py b/tensorflow/python/eager/remote_test.py
index 0c275b2aaeb..f541972b2a2 100644
--- a/tensorflow/python/eager/remote_test.py
+++ b/tensorflow/python/eager/remote_test.py
@@ -764,6 +764,22 @@ class MultiJobsTest(test.TestCase, parameterized.TestCase):
     with self.assertRaises(ValueError):
       remote.connect_to_cluster(self._cluster_resolver)
 
+  def testConnectToClusterWithMissingMaster(self):
+    resolver = SimpleClusterResolver(self._cluster, master='localhost:80')
+    with self.assertRaisesRegex(
+        ValueError,
+        '`make_master_device_default` is set to True but cannot find master '
+        'localhost:80 in the cluster'):
+      remote.connect_to_cluster(resolver)
+
+  def testConnectToClusterWithInvalidMaster(self):
+    resolver = SimpleClusterResolver(self._cluster, master='::1')
+    with self.assertRaisesRegex(
+        ValueError,
+        '`make_master_device_default` is set to True but cannot find master '
+        '::1 in the cluster'):
+      remote.connect_to_cluster(resolver)
+
   def testConnectToClusterWithLocalMaster(self):
     local_resolver = SimpleClusterResolver(ClusterSpec({}), master='local')
     remote.connect_to_cluster(local_resolver)
@@ -800,6 +816,80 @@ class MultiJobsTest(test.TestCase, parameterized.TestCase):
       self.assertIsNone(state)
 
 
+class RemoteUtilsTest(test.TestCase, parameterized.TestCase):
+
+  @parameterized.parameters(
+      ('localhost:foo',),
+      ('localhost:6553600',),
+  )
+  def testParseHostAndPortError(self, address):
+    host, port = remote._parse_host_and_port(address)
+    self.assertIsNone(host)
+    self.assertIsNone(port)
+
+  @parameterized.parameters(
+      ('localhost:8000', 'localhost', 8000),
+      ('grpc://localhost:8000', 'localhost', 8000),
+      ('grpc://[::1]:8000', '::1', 8000),
+      ('localhost', 'localhost', None),
+  )
+  def testParseHostAndPortSuccess(self, address, expected_host, expected_port):
+    host, port = remote._parse_host_and_port(address)
+    self.assertEqual(host, expected_host)
+    self.assertEqual(port, expected_port)
+
+  def testFindMasterJobAndTaskFirstMatch(self):
+    # Both jobs have a task at the master's address; the first job must win
+    # instead of being overwritten by the later one.
+    cluster = ClusterSpec({
+        'worker': ['localhost:8000'],
+        'chief': ['localhost:8000'],
+    })
+    self.assertEqual(
+        remote._find_master_job_and_task(cluster, 'localhost:8000'),
+        ('worker', 0))
+
+  def testFindMasterJobAndTaskFirstMatchWithoutPort(self):
+    cluster = ClusterSpec({
+        'worker': ['localhost:8000'],
+        'chief': ['localhost:8001'],
+    })
+    self.assertEqual(
+        remote._find_master_job_and_task(cluster, 'localhost'), ('worker', 0))
+
+  def testFindMasterJobAndTaskDistinctAddresses(self):
+    cluster = ClusterSpec({
+        'worker': ['10.0.0.5:2222'],
+        'chief': ['10.0.0.6:2222'],
+    })
+    self.assertEqual(
+        remote._find_master_job_and_task(cluster, '10.0.0.6:2222'),
+        ('chief', 0))
+
+  def testFindMasterJobAndTaskNoMatch(self):
+    cluster = ClusterSpec({'worker': ['10.0.0.5:2222']})
+    self.assertEqual(
+        remote._find_master_job_and_task(cluster, '10.0.0.9:2222'),
+        (None, None))
+
+  def testFindMasterJobAndTaskHostPrefixDoesNotMatch(self):
+    cluster = ClusterSpec({'worker': ['10.0.0.10:2222']})
+    self.assertEqual(
+        remote._find_master_job_and_task(cluster, '10.0.0.1'), (None, None))
+
+  def testFindMasterJobAndTaskWithGrpcPrefix(self):
+    cluster = ClusterSpec({'worker': ['10.0.0.5:2222']})
+    self.assertEqual(
+        remote._find_master_job_and_task(cluster, 'grpc://10.0.0.5:2222'),
+        ('worker', 0))
+
+  def testFindMasterJobAndTaskPortMismatch(self):
+    cluster = ClusterSpec({'worker': ['10.0.0.5:2222']})
+    self.assertEqual(
+        remote._find_master_job_and_task(cluster, '10.0.0.5:3333'),
+        (None, None))
+
+
 def _strip_prefix(s, prefix):
   return s[len(prefix):] if s.startswith(prefix) else s
 

From 518146118995040d3299df71b71d0d620ffe1cb8 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 5 Oct 2026 17:35:59 +0530
Subject: [PATCH 0194/1000] Update context.py

---
 tensorflow/python/eager/context.py | 3441 ++--------------------------
 1 file changed, 215 insertions(+), 3226 deletions(-)

diff --git a/tensorflow/python/eager/context.py b/tensorflow/python/eager/context.py
index 2435b922fac..73be78de959 100644
--- a/tensorflow/python/eager/context.py
+++ b/tensorflow/python/eager/context.py
@@ -1,4 +1,4 @@
-# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2019 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.
@@ -12,3233 +12,222 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""State management for eager execution."""
 
-import collections
-import contextlib
-import copy
-import gc
-import itertools
-import os
-import random
-import threading
+import weakref
 
-from absl import logging
+from absl.testing import parameterized
 import numpy as np
 
-from xla.tsl.protobuf import coordination_config_pb2
-from tensorflow.core.framework import function_pb2
-from tensorflow.core.framework import graph_debug_info_pb2
-from tensorflow.core.protobuf import config_pb2
-from tensorflow.core.protobuf import rewriter_config_pb2
-from tensorflow.python import pywrap_tfe
-from tensorflow.python import tf2
-from tensorflow.python.client import pywrap_tf_session
-from tensorflow.python.eager import cancellation
-from tensorflow.python.eager import execute
-from tensorflow.python.eager import executor
-from tensorflow.python.eager import monitoring
-from tensorflow.python.framework import c_api_util
-from tensorflow.python.framework import device as pydev
-from tensorflow.python.framework import tfrt_utils
-from tensorflow.python.util import compat
-from tensorflow.python.util import function_utils
-from tensorflow.python.util import is_in_graph_mode
-from tensorflow.python.util import tf_contextlib
-from tensorflow.python.util.deprecation import deprecated
-from tensorflow.python.util.tf_export import tf_export
-
-
-# TODO(b/307794935): Remove after a solution is found.
-is_oss = True  # updated by copybara
-
-GRAPH_MODE = 0
-EAGER_MODE = 1
-
-default_execution_mode = EAGER_MODE if tf2.enabled() else GRAPH_MODE
-
-# Cache from (old_device_name, partial_new_device_name) -> (new_device_name,
-# new_device_spec).
-# Note that we do not protect this with a lock and instead rely on python's GIL
-# and the idempotent nature of writes to provide thread safety.
-_device_parsing_cache = {}
-_starting_device_spec = pydev.DeviceSpec.from_string("")
-
-_MAXINT32 = 2**31 - 1
-
-DEVICE_PLACEMENT_EXPLICIT = pywrap_tfe.TFE_DEVICE_PLACEMENT_EXPLICIT
-DEVICE_PLACEMENT_WARN = pywrap_tfe.TFE_DEVICE_PLACEMENT_WARN
-DEVICE_PLACEMENT_SILENT = pywrap_tfe.TFE_DEVICE_PLACEMENT_SILENT
-DEVICE_PLACEMENT_SILENT_FOR_INT32 = (
-    pywrap_tfe.TFE_DEVICE_PLACEMENT_SILENT_FOR_INT32
-)
-
-SYNC = 0
-ASYNC = 1
-
-_KEEP_ALIVE_SECS = 600
-
-_python_eager_context_create_counter = monitoring.Counter(
-    "/tensorflow/api/python/eager_context_create_counter",
-    "Counter for number of eager contexts created in Python.",
-)
-
-# Re-exporting through context.
-is_tfrt_enabled = tfrt_utils.enabled
-
-# This flag and the associated environment var are transient and will eventually
-# be removed, once this experiment is enabled by default.
-_JIT_COMPILE_REWRITE_ENABLED = os.getenv("TF_JIT_COMPILE_REWRITE") == "1"
-
-_XLA_SHARDING_FOR_RESOURCE_VARIABLES = (
-    os.getenv("TF_XLA_SHARDING_FOR_RESOURCE_VARIABLES") == "1"
-)
-_OPTIONALS_IN_GRADIENTS = os.getenv("TF_OPTIONALS_IN_GRADIENTS") != "0"
-
-
-def run_eager_op_as_function_enabled():
-  return True
-
-
-# This method should only be called after the context has beein initialized.
-def enable_jit_compile_rewrite():
-  """Run jit_compile functions through rewrite pass.
-
-  This runs jit_compile functions through all of the multidevice function
-  rewrite passes.
-  """
-  global _JIT_COMPILE_REWRITE_ENABLED
-  _JIT_COMPILE_REWRITE_ENABLED = True
-  if context_safe() is not None:
-    context_safe().jit_compile_rewrite = True
-
-
-# This method should only be called after the context has been initialized.
-def disable_jit_compile_rewrite():
-  global _JIT_COMPILE_REWRITE_ENABLED
-  _JIT_COMPILE_REWRITE_ENABLED = False
-  if context_safe() is not None:
-    context_safe().jit_compile_rewrite = False
-
-
-def jit_compile_rewrite_enabled():
-  if context_safe() is not None:
-    return context_safe().jit_compile_rewrite
-  return _JIT_COMPILE_REWRITE_ENABLED
-
-
-def enable_xla_sharding_for_resource_variables():
-  """Enables support for annotating TF2 ResourceVariables with XLA sharding.
-
-  This allows placing XLA sharding annotations on the TF2 ResourceVariable
-  python object and inserts an XlaShardingOp with the annotation whenever a
-  ReadVariableOp is created.
-  """
-  global _XLA_SHARDING_FOR_RESOURCE_VARIABLES
-  _XLA_SHARDING_FOR_RESOURCE_VARIABLES = True
-  if context_safe() is not None:
-    context_safe().xla_sharding_for_resource_variables = True
-
-
-def disable_xla_sharding_for_resource_variables():
-  global _XLA_SHARDING_FOR_RESOURCE_VARIABLES
-  _XLA_SHARDING_FOR_RESOURCE_VARIABLES = False
-  if context_safe() is not None:
-    context_safe().xla_sharding_for_resource_variables = False
-
-
-def xla_sharding_for_resource_variables_enabled():
-  if context_safe() is not None:
-    return context_safe().xla_sharding_for_resource_variables
-  return _XLA_SHARDING_FOR_RESOURCE_VARIABLES
-
-
-def enable_optionals_in_gradients():
-  global _OPTIONALS_IN_GRADIENTS
-  _OPTIONALS_IN_GRADIENTS = True
-  if context_safe() is not None:
-    context_safe().optionals_in_gradients = True
-
-
-def disable_optionals_in_gradients():
-  global _OPTIONALS_IN_GRADIENTS
-  _OPTIONALS_IN_GRADIENTS = False
-  if context_safe() is not None:
-    context_safe().optionals_in_gradients = False
-
-
-def optionals_in_gradients_enabled():
-  if context_safe() is not None:
-    return context_safe().optionals_in_gradients
-  return _OPTIONALS_IN_GRADIENTS
-
-
-@contextlib.contextmanager
-def temporarily_disable_xla_sharding_for_resource_variables():
-  """Temporarily disables XLA sharding for resource variables.
-
-  Should be a no-op if it is already disabled.
-
-  Yields:
-    None.
-  """
-  previously_enabled = xla_sharding_for_resource_variables_enabled()
-
-  try:
-    disable_xla_sharding_for_resource_variables()
-    yield
-  finally:
-    if previously_enabled:
-      enable_xla_sharding_for_resource_variables()
-
-
-@contextlib.contextmanager
-def temporarily_disable_optionals_in_gradients():
-  """Temporarily disables generation of optionals in gradients.
-
-  Should be a no-op if it is already disabled.
-
-  Yields:
-    None.
-  """
-  previously_enabled = optionals_in_gradients_enabled()
-
-  try:
-    disable_optionals_in_gradients()
-    yield
-  finally:
-    if previously_enabled:
-      enable_optionals_in_gradients()
-
-
-# Expose it as internally public APIs for Keras use cases in b/171080602.
-tf_export("__internal__.is_tfrt_enabled", v1=[])(is_tfrt_enabled)
-
-
-class _EagerTensorCache(object):
-  """Simple cache which evicts items based on length in a FIFO manner."""
-
-  __slots__ = ["_data", "_max_items", "_max_tensor_size"]
-
-  def __init__(self, max_items=256, max_tensor_size=10000):
-    self._data = collections.OrderedDict()
-    self._max_items = max_items
-    self._max_tensor_size = max_tensor_size
-
-  def put(self, key, value):
-    if value._num_elements() > self._max_tensor_size:  # pylint: disable=protected-access
-      return
-
-    self._data[key] = value
-
-    if len(self._data) > self._max_items:
-      self._data.popitem(last=False)
-
-  def get(self, key):
-    return self._data.get(key, None)
-
-  def flush(self):
-    self._data.clear()
-
-
-class FunctionCallOptions:
-  """Options applied at call sites of eager functions.
-
-  Eager functions are functions decorated with tf.contrib.eager.defun.
-  """
-
-  __slots__ = ["_config_proto_serialized", "_executor_type"]
-
-  def __init__(self, executor_type=None, config_proto=None):
-    """Constructor.
-
-    Args:
-      executor_type: (optional) name of the executor to be used to execute the
-        eager function. If None or an empty string, the default Tensorflow
-        executor will be used.
-      config_proto: (optional) a `config_pb2.ConfigProto` proto or a serialized
-        string of that proto. The config used by Grappler when optimizing the
-        function graph. Each concrete function is optimized the first time is
-        called. Changing config_proto after the first call has no effect. If
-        config_proto is None, an empty RewriterConfig will be used.
-    """
-    self.config_proto_serialized = config_proto
-    self.executor_type = executor_type
-
-  @property
-  def executor_type(self):
-    return self._executor_type
-
-  @executor_type.setter
-  def executor_type(self, executor_type):
-    self._executor_type = executor_type
-
-  @property
-  def config_proto_serialized(self):
-    return self._config_proto_serialized
-
-  @config_proto_serialized.setter
-  def config_proto_serialized(self, config):
-    if isinstance(config, config_pb2.ConfigProto):
-      self._config_proto_serialized = config.SerializeToString(
-          deterministic=True
-      )
-    elif isinstance(config, str):
-      self._config_proto_serialized = config
-    elif config is None:
-      self._config_proto_serialized = (
-          config_pb2.ConfigProto().SerializeToString()
-      )
-    else:
-      raise ValueError(
-          "the rewriter config must be either a "
-          "config_pb2.ConfigProto, or a serialized string of that "
-          "proto or None. got: {}".format(type(config))
-      )
-
-  def as_attrs(self):
-    if self.config_proto_serialized is None:
-      config = function_utils.get_disabled_rewriter_config()
-    else:
-      config = self.config_proto_serialized
-    executor_type = self.executor_type or ""
-
-    return {"executor_type": executor_type, "config_proto": config}
-
-
-# Map from context_id (an int) to _TensorCaches.
-# Dicts are thread safe in CPython.
-# TODO(iga): Remove this once TensorCaches are moved to C++.
-_tensor_caches_map = {}
-
-
-class _TensorCaches(threading.local):
-  """Thread local tensor caches."""
-
-  __slots__ = ["_ones_rank_cache", "_zeros_cache"]
-
-  def __init__(self):
-    super().__init__()
-    self._ones_rank_cache = None
-    self._zeros_cache = None
-
-  @property
-  def ones_rank_cache(self):
-    if not self._ones_rank_cache:
-      self._ones_rank_cache = _EagerTensorCache()
-    return self._ones_rank_cache
-
-  @property
-  def zeros_cache(self):
-    if not self._zeros_cache:
-      self._zeros_cache = _EagerTensorCache()
-    return self._zeros_cache
-
-
-ContextSwitch = collections.namedtuple(
-    "ContextSwitch",
-    ["is_building_function", "enter_context_fn", "device_stack"],
-)
-
-
-# `_ContextSwitchStack` is a `threading.local` to match the semantics of
-# ``DefaultGraphStack`, which is also a `threading.local`.
-class _ContextSwitchStack(threading.local):
-  """A thread-local stack of context switches."""
-
-  def __init__(self, eager):
-    super().__init__()
-    self.stack = []
-    if eager:
-      # Initialize the stack with a pointer to enter the eager context; this
-      # ensures that the fact that eager execution was enabled is propagated
-      # across threads, since (1) `enable_eager_execution` modifies a
-      # process-level flag (`default_execution_mode`) and (2) `__init__` is
-      # called each time a threading.local object is used in a separate thread.
-      self.push(
-          is_building_function=False,
-          enter_context_fn=eager_mode,
-          device_stack=None,
-      )
-
-  def push(self, is_building_function, enter_context_fn, device_stack):
-    """Push metadata about a context switch onto the stack.
-
-    A context switch can take any one of the two forms: installing a graph as
-    the default graph, or entering the eager context. For each context switch,
-    we record whether or not the entered context is building a function.
-
-    Args:
-      is_building_function: (bool.) Whether the context is building a function.
-      enter_context_fn: (function.) A callable that executes the context switch.
-        For example, `graph.as_default` or `eager_mode`.
-      device_stack: If applicable, the device function stack for this graph.
-        When breaking out of graphs in init_scope, the innermost nonempty device
-        stack is used. Eager contexts put `None` here and the value is never
-        used.
-    """
-
-    self.stack.append(
-        ContextSwitch(is_building_function, enter_context_fn, device_stack)
-    )
-
-  def pop(self):
-    """Pop the stack."""
-
-    self.stack.pop()
-
-
-@tf_export("config.LogicalDevice")
-class LogicalDevice(
-    collections.namedtuple("LogicalDevice", ["name", "device_type"])
-):
-  """Abstraction for a logical device initialized by the runtime.
-
-  A `tf.config.LogicalDevice` corresponds to an initialized logical device on a
-  `tf.config.PhysicalDevice` or a remote device visible to the cluster. Tensors
-  and operations can be placed on a specific logical device by calling
-  `tf.device` with a specified `tf.config.LogicalDevice`.
-
-  Fields:
-    name: The fully qualified name of the device. Can be used for Op or function
-      placement.
-    device_type: String declaring the type of device such as "CPU" or "GPU".
-  """
-
-
-@tf_export(
-    "config.LogicalDeviceConfiguration",
-    "config.experimental.VirtualDeviceConfiguration",
-)
-class LogicalDeviceConfiguration(
-    collections.namedtuple(
-        "LogicalDeviceConfiguration",
-        [
-            "memory_limit",
-            "experimental_priority",
-            "experimental_device_ordinal",
-        ],
-    )
-):
-  """Configuration class for a logical devices.
-
-  The class specifies the parameters to configure a `tf.config.PhysicalDevice`
-  as it is initialized to a `tf.config.LogicalDevice` during runtime
-  initialization. Not all fields are valid for all device types.
-
-  See `tf.config.get_logical_device_configuration` and
-  `tf.config.set_logical_device_configuration` for usage examples.
-
-  Fields:
-    memory_limit: (optional) Maximum memory (in MB) to allocate on the virtual
-      device. Currently only supported for GPUs.
-    experimental_priority: (optional) Priority to assign to a virtual device.
-      Lower values have higher priorities and 0 is the default.
-      Within a physical GPU, the GPU scheduler will prioritize ops on virtual
-      devices with higher priority. Currently only supported for Nvidia GPUs.
-    experimental_device_ordinal: (optional) Ordinal number to order the virtual
-    device.
-      LogicalDevice with lower ordinal number will receive a lower device id.
-      Physical device id and location in the list is used to break ties.
-      Currently only supported for Nvidia GPUs.
-  """
-
-  def __new__(
-      cls,
-      memory_limit=None,
-      experimental_priority=None,
-      experimental_device_ordinal=None,
-  ):
-    return super().__new__(
-        cls, memory_limit, experimental_priority, experimental_device_ordinal
-    )
-
-
-@tf_export("config.PhysicalDevice")
-class PhysicalDevice(
-    collections.namedtuple("PhysicalDevice", ["name", "device_type"])
-):
-  """Abstraction for a locally visible physical device.
-
-  TensorFlow can utilize various devices such as the CPU or multiple GPUs
-  for computation. Before initializing a local device for use, the user can
-  customize certain properties of the device such as it's visibility or memory
-  configuration.
-
-  Once a visible `tf.config.PhysicalDevice` is initialized one or more
-  `tf.config.LogicalDevice` objects are created. Use
-  `tf.config.set_visible_devices` to configure the visibility of a physical
-  device and `tf.config.set_logical_device_configuration` to configure multiple
-  `tf.config.LogicalDevice` objects for a `tf.config.PhysicalDevice`. This is
-  useful when separation between models is needed or to simulate a multi-device
-  environment.
-
-  Fields:
-    name: Unique identifier for device.
-    device_type: String declaring the type of device such as "CPU" or "GPU".
-  """
-
-  pass
-
-
-class _AtomicCounter(object):
-  """A simple atomic counter."""
-
-  __slots__ = ["_value", "_lock"]
-
-  def __init__(self):
-    self._value = 0
-    self._lock = threading.Lock()
-
-  def increment_and_get(self):
-    with self._lock:
-      self._value += 1
-      return self._value
-
-
-_context_id_counter = _AtomicCounter()
-
-
-class _TensorCacheDeleter(object):
-  """Deletes tensor caches for a given context."""
-
-  __slots__ = ["_context_id"]
-
-  def __init__(self, context_id):
-    self._context_id = context_id
-
-  def __del__(self):
-    if _tensor_caches_map is None:
-      return
-    if self._context_id in _tensor_caches_map:
-      del _tensor_caches_map[self._context_id]
-
-
-# TODO(agarwal): rename to EagerContext / EagerRuntime ?
-# TODO(agarwal): consider keeping the corresponding Graph here.
-class Context:
-  """Environment in which eager operations execute."""
-
-  # TODO(agarwal): create and link in some documentation for `execution_mode`.
-  # pylint: disable=redefined-outer-name
-  def __init__(
-      self,
-      config=None,
-      device_policy=None,
-      execution_mode=None,
-      server_def=None,
-  ):
-    """Creates a new Context.
-
-    Args:
-      config: (Optional.) A `ConfigProto` protocol buffer with configuration
-        options for the Context. Note that a lot of these options may be
-        currently unimplemented or irrelevant when eager execution is enabled.
-      device_policy: (Optional.) What policy to use when trying to run an
-        operation on a device with inputs which are not on that device. When set
-        to None, an appropriate value will be picked automatically. The value
-        picked may change between TensorFlow releases.  Defaults to
-        DEVICE_PLACEMENT_SILENT. Valid values: DEVICE_PLACEMENT_EXPLICIT -
-        raises an error if the placement is not correct. DEVICE_PLACEMENT_WARN -
-        copies the tensors which are not on the right device but raises a
-        warning. DEVICE_PLACEMENT_SILENT - silently copies the tensors. This
-        might hide performance problems. DEVICE_PLACEMENT_SILENT_FOR_INT32 -
-        silently copies int32 tensors, raising errors on the other ones.
-      execution_mode: (Optional.) Policy controlling how operations dispatched
-        are actually executed. When set to None, an appropriate value will be
-        picked automatically. The value picked may change between TensorFlow
-        releases. Valid values: - SYNC: executes each operation synchronously.
-        ASYNC - executes each operation asynchronously. These operations may
-        return "non-ready" handles.
-      server_def: (Optional.) A tensorflow::ServerDef proto. Enables execution
-        on remote devices. GrpcServers need to be started by creating an
-        identical server_def to this, and setting the appropriate task_indexes,
-        so that the servers can communicate. It will then be possible to execute
-        operations on remote devices.
-
-    Raises:
-     ValueError: If execution_mode is not valid.
-    """
-    # This _id is used only to index the tensor caches.
-    # TODO(iga): Remove this when tensor caches are moved to C++.
-    self._id = _context_id_counter.increment_and_get()
-    self._tensor_cache_deleter = _TensorCacheDeleter(self._id)
-    _tensor_caches_map[self._id] = _TensorCaches()
-
-    self._config = config
-    self._thread_local_data = pywrap_tfe.EagerContextThreadLocalData(
-        self,
-        is_eager=lambda: default_execution_mode == EAGER_MODE,
-        device_spec=_starting_device_spec,
-    )
-    self._context_switches = _ContextSwitchStack(self.executing_eagerly())
-    self._context_handle = None
-    self._context_devices = None
-    self._seed = None
-    self._initialize_lock = threading.Lock()
-    self._initialized = False
-    if device_policy is None:
-      device_policy = DEVICE_PLACEMENT_SILENT
-    self._device_policy = device_policy
-    self._mirroring_policy = None
-    if execution_mode not in (None, SYNC, ASYNC):
-      raise ValueError(
-          "execution_mode should be None/SYNC/ASYNC. Got %s" % execution_mode
-      )
-    if execution_mode is None:
-      execution_mode = SYNC
-    self._default_is_async = execution_mode == ASYNC
-    self._use_tfrt = is_tfrt_enabled()
-    self._jit_compile_rewrite = jit_compile_rewrite_enabled()
-    self._xla_sharding_for_resource_variables = (
-        xla_sharding_for_resource_variables_enabled()
-    )
-    self._optionals_in_gradients = optionals_in_gradients_enabled()
-    self._server_def = server_def
-    self._collective_ops_server_def = None
-    self._collective_leader = None
-    self._collective_scoped_allocator_enabled_ops = None
-    self._collective_use_nccl_communication = None
-    self._collective_device_filters = None
-    self._coordination_service_config = None
-
-    self._device_lock = threading.Lock()
-    self._physical_devices = None
-    self._physical_device_to_index = None
-    self._pluggable_devices = None
-    self._visible_device_list = []
-    self._memory_growth_map = None
-    self._virtual_device_map = {}
-
-    # Values set after construction
-    self._optimizer_jit = None
-    self._intra_op_parallelism_threads = None
-    self._inter_op_parallelism_threads = None
-    self._soft_device_placement = None
-    self._log_device_placement = None
-    self._operation_timeout_in_ms = None
-    self._enable_mlir_graph_optimization = None
-    self._optimizer_experimental_options = {}
-
-    _python_eager_context_create_counter.get_cell().increase_by(1)
-
-    self._is_global_context = False
-
-    # Number of retries to give the SetServerDef step. This is useful for fault
-    # tolerant initial connection in high-preemption settings like
-    # ParameterServerStrategy training.
-    self._set_server_def_retries = 0
-
-  # pylint: enable=redefined-outer-name
-
-  def _set_global_seed(self, seed):
-    """Set a global eager mode seed for random ops."""
-    self._seed = seed
-    # `random.Random(seed)` needs `seed` to be hashable, while values of type
-    # e.g. `np.int64` or `np.ndarray` are not. We use `int(...)` to convert them
-    # to int.
+from xla.service import hlo_pb2
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.platform import test
+
+
+class ContextTest(test.TestCase, parameterized.TestCase):
+
+  def testSetGlobalSeed(self):
+    c = context.Context()
+    c._set_global_seed(123)
+    for t in [np.int32, np.int64, np.uint32, np.uint64]:
+      c._set_global_seed(t(123))
+      c._set_global_seed(np.array(123, dtype=t))
+      c._set_global_seed(ops.convert_to_tensor(123, dtype=t))
+
+  def testContextIsDestroyedAfterTensors(self):
+    # Create a new context
+    new_context = context.Context()
+    weak_c = weakref.ref(new_context)
+    new_context.ensure_initialized()
+
+    # Create a tensor with the new context as default.
+    # Make sure to restore the original context.
+    original_context = context.context()
     try:
-      hash(seed)
-      self._rng = random.Random(seed)
-    except TypeError:
-      seed = int(np.array(seed))
-      self._rng = random.Random(seed)
-    # Also clear the kernel cache, to reset any existing seeds
-    if self._context_handle is not None:
-      pywrap_tfe.TFE_ContextClearCaches(self._context_handle)
-
-  def _internal_operation_seed(self):
-    """Returns a fake operation seed.
-
-      In eager mode, user shouldn't set or depend on operation seed.
-      Here, we generate a random seed based on global seed to make
-      operation's randomness different and depend on the global seed.
-
-    Returns:
-      A fake operation seed based on global seed.
-    """
-    return self._rng.randint(0, _MAXINT32)
-
-  def _initialize_logical_devices(self):
-    """Helper to initialize devices."""
-    # Store list of devices
-    logical_devices = []
-    context_devices = []
-    device_list = pywrap_tfe.TFE_ContextListDevices(self._context_handle)
-    try:
-      self._num_gpus = 0
-      current_job, current_task = None, None
-      server_def = self._server_def or self._collective_ops_server_def
-      if server_def is not None:
-        current_job, current_task = server_def.job_name, server_def.task_index
-      for i in range(pywrap_tfe.TF_DeviceListCount(device_list)):
-        dev_name = pywrap_tfe.TF_DeviceListName(device_list, i)
-        context_devices.append(pydev.canonical_name(dev_name))
-        spec = pydev.DeviceSpec.from_string(dev_name)
-        # If the job is localhost, we assume that the cluster has not yet been
-        # configured and thus clear the job, replica & task.
-        if spec.job == "localhost":
-          spec = spec.replace(job=None, replica=None, task=None)
-        logical_devices.append(
-            LogicalDevice(name=spec.to_string(), device_type=spec.device_type)
-        )
-        dev_type = pywrap_tfe.TF_DeviceListType(device_list, i)
-        if (
-            dev_type == "GPU"
-            and spec.job == current_job
-            and spec.task == current_task
-        ):
-          self._num_gpus += 1
-
-    finally:
-      self._logical_devices = logical_devices
-      self._context_devices = context_devices
-      pywrap_tfe.TF_DeleteDeviceList(device_list)
-
-  def ensure_initialized(self):
-    """Initialize handle and devices if not already done so."""
-    if self._initialized:
-      return
-    with self._initialize_lock:
-      if self._initialized:
-        return
-      assert self._context_devices is None
-      opts = pywrap_tfe.TFE_NewContextOptions()
-      try:
-        config_str = self.config.SerializeToString()
-        pywrap_tfe.TFE_ContextOptionsSetConfig(opts, config_str)
-        if self._device_policy is not None:
-          pywrap_tfe.TFE_ContextOptionsSetDevicePlacementPolicy(
-              opts, self._device_policy
-          )
-        if self._mirroring_policy is not None:
-          pywrap_tfe.TFE_ContextOptionsSetMirroringPolicy(
-              opts, self._mirroring_policy
-          )
-        if self._default_is_async == ASYNC:
-          pywrap_tfe.TFE_ContextOptionsSetAsync(opts, True)
-        if self._use_tfrt is not None:
-          pywrap_tfe.TFE_ContextOptionsSetTfrt(opts, self._use_tfrt)
-        pywrap_tfe.TFE_ContextOptionsSetRunEagerOpAsFunction(opts, True)
-        pywrap_tfe.TFE_ContextOptionsSetJitCompileRewrite(
-            opts, self._jit_compile_rewrite
-        )
-        context_handle = pywrap_tfe.TFE_NewContext(opts)
-      finally:
-        pywrap_tfe.TFE_DeleteContextOptions(opts)
-      assert not (self._server_def and self._collective_ops_server_def), (
-          "Cannot enable remote execution as well as collective ops at the "
-          "moment. If this is important to you, please file an issue."
-      )
-      if self._server_def is not None:
-        server_def_str = self._server_def.SerializeToString()
-        timeout = 0  # Indicates no timeout.
-        pywrap_tfe.TFE_ContextSetServerDefWithTimeoutAndRetries(
-            context_handle,
-            _KEEP_ALIVE_SECS,
-            server_def_str,
-            timeout,
-            self._set_server_def_retries,
-        )
-      elif self._collective_ops_server_def is not None:
-        server_def_str = self._collective_ops_server_def.SerializeToString()
-        pywrap_tfe.TFE_EnableCollectiveOps(context_handle, server_def_str)
-
-      self._context_handle = context_handle
-      self._initialize_logical_devices()
-      self._initialized = True
-
-      if self._is_global_context:
-        pywrap_tfe.TFE_Py_SetCEagerContext(self._context_handle)
-
-  def ensure_uninitialized(self):
-    """Uninitialize handle and devices if not already done so."""
-    with self._initialize_lock:
-      if not self._initialized:
-        return
-      self._context_devices = None
-      self._logical_devices = None
-      self._server_def = None
-      self._initialized = False
-
-      if self._is_global_context:
-        pywrap_tfe.TFE_Py_SetCEagerContext(None)
-
-      self._context_handle = None
-
-  def mark_as_global_context(self):
-    # If the context was already initialized, publish it. Otherwise wait with
-    # publication until it's initialized.
-    if self._initialized:
-      pywrap_tfe.TFE_Py_SetCEagerContext(self._context_handle)
-    self._is_global_context = True
-
-  def _clear_caches(self):
-    self.ones_rank_cache().flush()
-    self.zeros_cache().flush()
-    pywrap_tfe.TFE_ClearScalarCache()
-
-  def get_server_def(self):
-    return self._server_def
-
-  def set_server_def(self, server_def, keep_alive_secs=_KEEP_ALIVE_SECS):
-    """Allow setting a server_def on the context.
-
-    When a server def is replaced, it effectively clears a bunch of caches
-    within the context. If you attempt to use a tensor object that was pointing
-    to a tensor on the remote device, it will raise an error.
-
-    Args:
-      server_def: A tensorflow::ServerDef proto. Enables execution on remote
-        devices.
-      keep_alive_secs: Num. seconds after which the remote end will hang up. As
-        long as the client is still alive, the server state for the context will
-        be kept alive. If the client is killed (or there is some failure), the
-        server will clean up its context keep_alive_secs after the final RPC it
-        receives.
-
-    Raises:
-      ValueError: if server_def is None.
-    """
-    if not server_def:
-      raise ValueError("server_def is None.")
-
-    self._server_def = server_def
-
-    if self._context_handle:
-      server_def_str = server_def.SerializeToString()
-      pywrap_tfe.TFE_ContextSetServerDef(
-          self._context_handle, keep_alive_secs, server_def_str
-      )
-      self._initialize_logical_devices()
-
-    # Clear all the caches in case there are remote tensors in them.
-    self._clear_caches()
-    # Also clear the device parsing cache since it caches the resolution of
-    # partial device names, which may become different due to the set_server_def
-    # call as we may have defined different devices.
-    _device_parsing_cache.clear()
-
-  def update_server_def(self, server_def, keep_alive_secs=_KEEP_ALIVE_SECS):
-    """Update a server_def on the context.
-
-    Args:
-      server_def: A tensorflow::ServerDef proto. Enables execution on remote
-        devices.
-      keep_alive_secs: Num. seconds after which the remote end will hang up. As
-        long as the client is still alive, the server state for the context will
-        be kept alive. If the client is killed (or there is some failure), the
-        server will clean up its context keep_alive_secs after the final RPC it
-        receives.
-
-    Raises:
-      ValueError: if server_def is None.
-    """
-    if not server_def:
-      raise ValueError("server_def is None.")
-
-    self._server_def = server_def
-
-    if self._context_handle:
-      server_def_str = server_def.SerializeToString()
-      pywrap_tfe.TFE_ContextUpdateServerDef(
-          self._context_handle, keep_alive_secs, server_def_str
-      )
-      self._initialize_logical_devices()
-
-    self._clear_caches()
-
-  def check_alive(self, worker_name):
-    """Checks whether a remote worker is alive or not.
-
-    Args:
-      worker_name: a string representing the remote worker. It must be a fully
-        specified name like "/job:worker/replica:0/task:0".
-
-    Returns:
-      a boolean indicating whether the remote worker is alive or not.
-
-    Raises:
-      ValueError: if context is not initialized.
-    """
-    # TODO(yuefengz): support checking multiple workers.
-    if self._context_handle:
-      return pywrap_tfe.TFE_ContextCheckAlive(self._context_handle, worker_name)
-    else:
-      raise ValueError("Context is not initialized.")
-
-  def sync_executors(self):
-    """Sync both local executors and the ones on remote workers.
-
-    In async execution mode, local function calls can return before the
-    corresponding remote op/function execution requests are completed. Calling
-    this method creates a synchronization barrier for remote executors. It only
-    returns when all remote pending nodes are finished, potentially with errors
-    if any remote executors are in error state.
-
-    Raises:
-      ValueError: if context is not initialized.
-    """
-    if self._context_handle:
-      pywrap_tfe.TFE_ContextSyncExecutors(self._context_handle)
-    else:
-      raise ValueError("Context is not initialized.")
-
-  def clear_executor_errors(self):
-    """Clear errors in both local executors and remote workers.
-
-    After receiving errors from remote workers, additional requests on the fly
-    could further taint the status on the remote workers due to the async nature
-    of remote execution. Calling this method block on waiting for all pending
-    nodes in remote executors to finish and clear their error statuses.
-
-    Raises:
-      ValueError: if context is not initialized.
-    """
-    if self._context_handle:
-      pywrap_tfe.TFE_ContextClearExecutors(self._context_handle)
-    else:
-      raise ValueError("Context is not initialized.")
-
-  def configure_coordination_service(
-      self,
-      service_type,
-      service_leader="",
-      enable_health_check=True,
-      cluster_register_timeout_in_ms=0,
-      heartbeat_timeout_in_ms=0,
-      shutdown_barrier_timeout_in_ms=0,
-      coordinated_jobs=None,
-      allow_new_incarnation_to_reconnect=False,
-  ):
-    """Enable distributed coordination service with specified configs."""
-    if self._context_handle:
-      logging.warning(
-          "Configuring coordination service type may not be "
-          "effective because the context is already initialized."
-      )
-    config = coordination_config_pb2.CoordinationServiceConfig()
-    config.service_type = service_type
-    if service_leader:
-      config.service_leader = pydev.canonical_name(service_leader)
-    config.enable_health_check = enable_health_check
-    config.cluster_register_timeout_in_ms = cluster_register_timeout_in_ms
-    config.heartbeat_timeout_in_ms = heartbeat_timeout_in_ms
-    config.shutdown_barrier_timeout_in_ms = shutdown_barrier_timeout_in_ms
-    config.allow_new_incarnation_to_reconnect = (
-        allow_new_incarnation_to_reconnect
-    )
-    if coordinated_jobs is not None:
-      if isinstance(coordinated_jobs, list):
-        config.coordinated_job_list.extend(coordinated_jobs)
-      else:
-        raise ValueError(
-            "`coordinated_jobs` must be list[CoordinatedJob] or "
-            "None, but got: %s" % (coordinated_jobs,)
-        )
-    self._coordination_service_config = config
-
-  @property
-  def coordination_service(self):
-    return self._coordination_service_config
-
-  def set_config_key_value(self, key, value):
-    self.ensure_initialized()
-    pywrap_tfe.TFE_InsertConfigKeyValue(self._context_handle, key, value)
-
-  # If `timeout_in_ms=0`, this will block until the key-value is set or the
-  # worker shuts down.
-  def get_config_key_value(self, key, timeout_in_ms=0):
-    self.ensure_initialized()
-    with c_api_util.tf_buffer() as buffer_:
-      pywrap_tfe.TFE_GetConfigKeyValue(
-          self._context_handle, key, timeout_in_ms, buffer_
-      )
-      value = pywrap_tf_session.TF_GetBuffer(buffer_).decode("utf-8")
-    return value
-
-  def delete_config_key_value(self, key):
-    self.ensure_initialized()
-    pywrap_tfe.TFE_DeleteConfigKeyValue(self._context_handle, key)
-
-  def report_error_to_cluster(self, error_code, error_message):
-    """Report error to other members in a multi-client cluster.
-
-    Args:
-      error_code: a `tf.errors` error code.
-      error_message: a string. The error message.
-    """
-    if self._context_handle:
-      pywrap_tfe.TFE_ReportErrorToCluster(
-          self._context_handle, error_code, error_message
-      )
-    else:
-      raise ValueError("Context is not initialized.")
-
-  def get_task_states(self, job_configs):
-    """Get task states from the Coordination Service.
-
-    Args:
-      job_configs: A list of tuples of job name and task number.
-
-    Returns:
-      A list of TF_Status.
-    """
-    if self._context_handle:
-      job_names, task_nums = zip(*job_configs)
-      return pywrap_tfe.TFE_GetTaskStates(
-          self._context_handle, job_names, task_nums
-      )
-    else:
-      raise ValueError("Context is not initialized.")
-
-  def wait_at_barrier(self, barrier_id, timeout_in_ms):
-    """Blocks until all coordinated tasks are at the barrier.
-
-    The barrier may fail if it times out or if one of the tasks is unhealthy.
-
-    Args:
-      barrier_id: Unique string identifying the barrier.
-      timeout_in_ms: Duration before the barrier times out and fails.
-    """
-    self.ensure_initialized()
-    pywrap_tfe.TFE_WaitAtBarrier(
-        self._context_handle, barrier_id, timeout_in_ms
-    )
-
-  def clear_kernel_cache(self):
-    """Clear kernel cache and reset all stateful kernels."""
-    if self._context_handle is not None:
-      pywrap_tfe.TFE_ContextClearCaches(self._context_handle)
-
-  def enable_collective_ops(self, server_def):
-    """Enable distributed collective ops with an appropriate server_def.
-
-    Args:
-      server_def: A tensorflow::ServerDef proto. Enables execution on remote
-        devices.
-
-    Raises:
-      ValueError: if server_def is None.
-      RuntimeError: if this method is not called at program startup.
-    """
-    if not server_def:
-      raise ValueError("server_def is None.")
-
-    self._collective_ops_server_def = server_def
-
-    # TODO(b/129298253): Allow creating datasets/tensors before enabling
-    # collective ops.
-    if self._context_handle is not None:
-      logging.warning(
-          "Enabling collective ops after program startup may cause "
-          "error when accessing previously created tensors."
-      )
-      with self._initialize_lock:
-        assert self._initialized
-        server_def_str = self._collective_ops_server_def.SerializeToString()
-        pywrap_tfe.TFE_EnableCollectiveOps(self._context_handle, server_def_str)
-        self._initialize_logical_devices()
-        self._clear_caches()
-
-  def configure_collective_ops(
-      self,
-      collective_leader="",
-      scoped_allocator_enabled_ops=("CollectiveReduce",),
-      use_nccl_communication=False,
-      device_filters=None,
-  ):
-    """Configure collective ops.
-
-      Collective group leader is necessary for collective ops to run, other
-      configurations are mainly for the purpose of performance.
-
-    Args:
-      collective_leader: a device string for collective leader, e.g.
-        "/job:worker/replica:0/task:0"; empty string means local execution of
-        collective ops.
-      scoped_allocator_enabled_ops: a tuple or a list of op names for scoped
-        allocator to run with.
-      use_nccl_communication: whether to use nccl communication for collective
-        ops.
-      device_filters: a tuple or a list of device strings. If set, corresponding
-        task can only see the devices filtered by these device filters.
-
-    Raises:
-      RuntimeError: if this method is not called at program startup.
-    """
-    if self._collective_leader is not None:
-      if (
-          self._collective_leader != collective_leader
-          or self._collective_scoped_allocator_enabled_ops
-          != scoped_allocator_enabled_ops
-          or self._collective_use_nccl_communication != use_nccl_communication
-          or self._collective_device_filters != device_filters
-      ):
-        raise ValueError("Collective ops are already configured.")
-      else:
-        return
-
-    if self._context_handle is not None:
-      raise RuntimeError("Collective ops must be configured at program startup")
-
-    self._collective_leader = collective_leader
-    self._collective_scoped_allocator_enabled_ops = scoped_allocator_enabled_ops
-    self._collective_use_nccl_communication = use_nccl_communication
-    self._collective_device_filters = device_filters
-
-  def abort_collective_ops(self, code, message):
-    """Abort the collective ops.
-
-    This is intended to be used when a peer failure is detected, which allows
-    the user to handle the case instead of hanging. This aborts all on-going
-    collectives. After all subsequent collectives error immediately, and you
-    need to reset_context() to use collectives again.
-
-    Args:
-      code: a `tf.errors` error code.
-      message: a string. The error message.
-    """
-    self.ensure_initialized()
-    pywrap_tfe.TFE_AbortCollectiveOps(self._handle, code, message)
-
-  def check_collective_ops_peer_health(self, task, timeout_in_ms):
-    """Check collective peer health.
-
-    This probes each task to see if they're still alive. Note that restarted
-    tasks are considered a different one, and they're considered not healthy.
-
-    This should only be used in multi client multi worker training.
-
-    Args:
-      task: a task string, must be in the format of /job:xxx/replica:0/task:N.
-      timeout_in_ms: an integer, the timeout. If zero, there's no timeout.
-
-    Raises:
-      tf.errors.UnavailableError: when a peer is down.
-      tf.errors.FailedPreconditionError: when a peer is a different one from the
-        one this task has talked to, e.g. the peer has restarted.
-      tf.errors.InvalidArgumentError: when the task string is invalid.
-    """
-    self.ensure_initialized()
-    pywrap_tfe.TFE_CollectiveOpsCheckPeerHealth(
-        self._handle, task, timeout_in_ms
-    )
-
-  @property
-  def _handle(self):
-    if self._context_handle is None:
-      raise AssertionError("Context must be initialized first.")
-
-    return self._context_handle
-
-  @property
-  def _devices(self):
-    if self._context_devices is None:
-      raise AssertionError("Context must be initialized first.")
-
-    return self._context_devices
-
-  def __str__(self):
-    if self._context_handle is None:
-      return "Eager TensorFlow Context. Devices currently uninitialized."
-    else:
-      devices = self._devices
-      lines = ["Eager TensorFlow Context with %d devices" % (len(devices))]
-      for i, d in enumerate(devices):
-        lines.append("   Device %d: %s" % (i, d))
-      return "\n".join(lines)
-
-  @tf_contextlib.contextmanager
-  def _mode(self, mode):
-    """A context manager to allow setting the mode to EAGER/GRAPH."""
-    ctx = self._thread_local_data
-    old_is_eager = ctx.is_eager
-    ctx.is_eager = mode == EAGER_MODE
-    if mode == EAGER_MODE:
-      # Entering graph mode does not provide us with sufficient information to
-      # record a context switch; graph-based context switches are only logged
-      # when a graph is registered as the default graph.
-      self.context_switches.push(False, eager_mode, None)
-    try:
-      yield
-    finally:
-      ctx.is_eager = old_is_eager
-      if mode == EAGER_MODE:
-        self.context_switches.pop()
-
-  def executing_eagerly(self):
-    """Returns True if current thread has eager executing enabled."""
-    return self._thread_local_data.is_eager
-
-  def ones_rank_cache(self):
-    """Per-device cache for scalars."""
-    return _tensor_caches_map[self._id].ones_rank_cache
-
-  def zeros_cache(self):
-    """Per-device cache for scalars."""
-    return _tensor_caches_map[self._id].zeros_cache
-
-  @property
-  def scope_name(self):
-    """Returns scope name for the current thread."""
-    return self._thread_local_data.scope_name
-
-  @scope_name.setter
-  def scope_name(self, s):
-    """Sets scope name for the current thread."""
-    self._thread_local_data.scope_name = s
-
-  @property
-  def device_name(self):
-    """Returns the device name for the current thread."""
-    return self._thread_local_data.device_name
-
-  @property
-  def device_spec(self):
-    """Returns the device spec for the current thread."""
-    return self._thread_local_data.device_spec
-
-  def _set_device(self, device_name, device_spec):
-    self._thread_local_data.device_name = device_name
-    self._thread_local_data.device_spec = device_spec
-
-  def device(self, name):
-    """Context-manager to force placement of operations and Tensors on a device.
-
-    Args:
-      name: Name of the device or None to get default placement.
-
-    Returns:
-      Context manager that forces device placement.
-
-    Raises:
-      ValueError: If name is not a string or is an invalid device name.
-      RuntimeError: If device scopes are not properly nested.
-    """
-    if isinstance(name, LogicalDevice):
-      name = name.name
-    elif pydev.is_device_spec(name):
-      name = name.to_string()
-    return _EagerDeviceContext(self, name)
-
-  def devices(self):
-    """List of the names of devices available to execute operations."""
-    return self._devices
-
-  def host_address_space(self):
-    self.ensure_initialized()
-    with c_api_util.tf_buffer() as buffer_:
-      pywrap_tfe.TFE_HostAddressSpace(self._context_handle, buffer_)
-      address_space = pywrap_tf_session.TF_GetBuffer(buffer_).decode("utf-8")
-    return address_space
-
-  # TODO(fishx): remove this property.
-  @property
-  def execution_mode(self):
-    """Gets execution mode for current thread."""
-    return ASYNC if self.is_async() else SYNC
-
-  @execution_mode.setter
-  def execution_mode(self, mode):
-    """Sets execution mode for current thread."""
-    if mode not in (None, SYNC, ASYNC):
-      raise ValueError(
-          "Execution mode should be None/SYNC/ASYNC. Got %s" % mode
-      )
-
-    if mode is None:
-      mode = SYNC
-
-    enable_async = mode == ASYNC
-    if self.is_async() != enable_async:
-      # Only set the execution mode if the context has already been initialized
-      if self._context_handle is not None:
-        self.executor.wait()
-        executor_new = executor.new_executor(enable_async)
-        self._thread_local_data.executor = executor_new
-        pywrap_tfe.TFE_ContextSetExecutorForThread(
-            self._context_handle, executor_new.handle()
-        )
-      else:
-        self._default_is_async = enable_async
-
-  def is_async(self):
-    if self._context_handle is not None:
-      return self.executor.is_async()
-    else:
-      return self._default_is_async
-
-  @property
-  def executor(self):
-    self.ensure_initialized()
-    return executor.Executor(
-        pywrap_tfe.TFE_ContextGetExecutorForThread(self._context_handle)
-    )
-
-  @executor.setter
-  def executor(self, e):
-    self.ensure_initialized()
-    pywrap_tfe.TFE_ContextSetExecutorForThread(self._context_handle, e.handle())
-
-  @property
-  def config(self):
-    """Return the ConfigProto with all runtime deltas applied."""
-    # Ensure physical devices have been discovered and config has been imported
-    self._initialize_physical_devices()
-
-    config = config_pb2.ConfigProto()
-    if self._config is not None:
-      config.CopyFrom(self._config)
-
-    if self._optimizer_jit is not None:
-      config.graph_options.optimizer_options.global_jit_level = (
-          config_pb2.OptimizerOptions.ON_1
-          if self._optimizer_jit
-          else config_pb2.OptimizerOptions.OFF
-      )
-    if self._intra_op_parallelism_threads is not None:
-      config.intra_op_parallelism_threads = self._intra_op_parallelism_threads
-    if self._inter_op_parallelism_threads is not None:
-      config.inter_op_parallelism_threads = self._inter_op_parallelism_threads
-
-    if self._soft_device_placement is not None:
-      config.allow_soft_placement = self._soft_device_placement
-    else:
-      config.allow_soft_placement = self.executing_eagerly()
-
-    if self._log_device_placement is not None:
-      config.log_device_placement = self._log_device_placement
-
-    if self._operation_timeout_in_ms is not None:
-      config.operation_timeout_in_ms = self._operation_timeout_in_ms
-
-    is_mlir_bridge_enabled = pywrap_tfe.TF_IsMlirBridgeEnabled()
-    config.experimental.mlir_bridge_rollout = is_mlir_bridge_enabled
-    if (
-        is_mlir_bridge_enabled
-        == config_pb2.ConfigProto.Experimental.MLIR_BRIDGE_ROLLOUT_ENABLED
-    ):
-      config.experimental.enable_mlir_bridge = True
-
-    if self._enable_mlir_graph_optimization is not None:
-      config.experimental.enable_mlir_graph_optimization = (
-          self._enable_mlir_graph_optimization
-      )
-
-    def rewriter_toggle(option):
-      toggle = self._optimizer_experimental_options.get(option, None)
-      if toggle is None:
-        return
-
-      setattr(
-          config.graph_options.rewrite_options,
-          option,
-          (
-              rewriter_config_pb2.RewriterConfig.ON
-              if toggle
-              else rewriter_config_pb2.RewriterConfig.OFF
-          ),
-      )
-
-    def rewriter_bool(option):
-      toggle = self._optimizer_experimental_options.get(option, None)
-      if toggle is None:
-        return
-
-      setattr(config.graph_options.rewrite_options, option, toggle)
-
-    rewriter_toggle("layout_optimizer")
-    rewriter_toggle("constant_folding")
-    rewriter_toggle("shape_optimization")
-    rewriter_toggle("remapping")
-    rewriter_toggle("arithmetic_optimization")
-    rewriter_toggle("dependency_optimization")
-    rewriter_toggle("loop_optimization")
-    rewriter_toggle("function_optimization")
-    rewriter_toggle("debug_stripper")
-    rewriter_bool("disable_model_pruning")
-    rewriter_toggle("scoped_allocator_optimization")
-    rewriter_toggle("pin_to_host_optimization")
-    rewriter_toggle("implementation_selector")
-    rewriter_toggle("auto_mixed_precision")
-    rewriter_toggle("use_plugin_optimizers")
-    rewriter_bool("disable_meta_optimizer")
-    rewriter_toggle("auto_mixed_precision_onednn_bfloat16")
-    rewriter_toggle("auto_mixed_precision_mkl")
-    nodes = self._optimizer_experimental_options.get("min_graph_nodes", None)
-    if nodes is not None:
-      config.graph_options.rewrite_options.min_graph_nodes = nodes
-
-    # Compute device counts
-    config.device_count["CPU"] = 0
-    config.device_count["GPU"] = 0
-    for dev in self._physical_devices:
-      if dev not in self._visible_device_list:
-        continue
-
-      virtual_devices = self._virtual_device_map.get(dev)
-      if virtual_devices is None:
-        config.device_count[dev.device_type] += 1
-      else:
-        config.device_count[dev.device_type] += len(virtual_devices)
-
-    # Configure gpu_options
-    gpu_options = self._compute_gpu_options()
-    config.gpu_options.MergeFrom(gpu_options)
-
-    # Configure pluggable_device_options
-    pluggable_device_options = self._compute_pluggable_device_options()
-    config.pluggable_device_options.MergeFrom(pluggable_device_options)
-
-    # Configure collective ops
-    if self._collective_leader:
-      config.experimental.collective_group_leader = self._collective_leader
-    if self._collective_scoped_allocator_enabled_ops:
-      rewrite_options = config.graph_options.rewrite_options
-      rewrite_options.scoped_allocator_optimization = (
-          rewriter_config_pb2.RewriterConfig.ON
-      )
-      del rewrite_options.scoped_allocator_opts.enable_op[:]
-      for op in self._collective_scoped_allocator_enabled_ops:
-        rewrite_options.scoped_allocator_opts.enable_op.append(op)
-    if self._collective_use_nccl_communication:
-      config.experimental.collective_nccl = True
-    if self._collective_device_filters:
-      del config.device_filters[:]
-      for f in self._collective_device_filters:
-        config.device_filters.append(f)
-
-    # Configure coordination service
-    if self._coordination_service_config:
-      config.experimental.coordination_config.CopyFrom(
-          self._coordination_service_config
-      )
-
-    return config
-
-  def _compute_device_options(self, device_type="GPU"):
-    """Build the GPUOptions proto for GPU or PluggableDevice."""
-    if device_type not in ["GPU", "PluggableDevice"]:
-      raise ValueError(
-          "device types other than GPU and PluggableDevice are not supported."
-      )
-    visible_device_list = []
-    virtual_devices = []
-    # This mapping is needed to handle multiple sub types of PluggableDevices.
-    device_to_indices = collections.defaultdict(int)
-    memory_growths = set()
-    compatible_devices = (
-        self.list_physical_devices("GPU")
-        if device_type == "GPU"
-        else self._pluggable_devices
-    )
-    support_virtual_devices = device_type == "GPU"
-    for dev in compatible_devices:
-      device_index = device_to_indices[dev.device_type]
-      device_to_indices[dev.device_type] += 1
-
-      if dev not in self._visible_device_list:
-        continue
-
-      growth = self._memory_growth_map[dev]
-      memory_growths.add(growth)
-      if device_type == "PluggableDevice":
-        visible_device_list.append(dev.device_type + ":" + str(device_index))
-      else:
-        visible_device_list.append(str(device_index))
-
-      if support_virtual_devices and self._virtual_device_map:
-        vdevs = self._virtual_device_map.get(dev, [])
-        device_ordinals = []
-        device_limits = []
-        priority = []
-        for virt_dev in vdevs:
-          if virt_dev.experimental_device_ordinal is not None:
-            device_ordinals.append(virt_dev.experimental_device_ordinal)
-          device_limits.append(virt_dev.memory_limit)
-          if virt_dev.experimental_priority is not None:
-            priority.append(virt_dev.experimental_priority)
-        # If priority is specified, it must be specified for all virtual
-        # devices.
-        if priority and len(device_limits) != len(priority):
-          raise ValueError("priority must be specified for all virtual devices")
-        # If device_ordinals is specified, it must be specified for all virtual
-        # devices.
-        if device_ordinals and len(device_limits) != len(device_ordinals):
-          raise ValueError(
-              "device_ordinals must be specified for all virtual devices"
-          )
-
-        virtual_devices.append(
-            config_pb2.GPUOptions.Experimental.VirtualDevices(
-                memory_limit_mb=device_limits,
-                priority=priority,
-                device_ordinal=device_ordinals,
-            )
-        )
-
-    # Only compute growth if virtual devices have not been configured and we
-    # have GPUs
-    if not virtual_devices and memory_growths:
-      if len(memory_growths) > 1:
-        raise ValueError("Memory growth cannot differ between GPU devices")
-      allow_growth = memory_growths.pop()
-    else:
-      allow_growth = None
-
-    return config_pb2.GPUOptions(
-        allow_growth=allow_growth,
-        visible_device_list=",".join(visible_device_list),
-        experimental=config_pb2.GPUOptions.Experimental(
-            virtual_devices=virtual_devices
-        ),
-    )
-
-  def _compute_gpu_options(self):
-    """Build the GPUOptions proto for GPU."""
-    return self._compute_device_options(device_type="GPU")
-
-  def _compute_pluggable_device_options(self):
-    """Build the GPUOptions proto for PluggableDevice."""
-    return self._compute_device_options(device_type="PluggableDevice")
-
-  @property
-  def function_call_options(self):
-    """Returns function call options for current thread.
-
-    Note that the returned object is still referenced by the eager context.
-
-    Returns: the FunctionCallOptions for current thread.
-    """
-    if self._thread_local_data.function_call_options is None:
-      config = self.config
-
-      # Default to soft placement for functions unless specified
-      if self._soft_device_placement is None:
-        config.allow_soft_placement = True
-      self._thread_local_data.function_call_options = FunctionCallOptions(
-          config_proto=config
-      )
-
-    return self._thread_local_data.function_call_options
-
-  @function_call_options.setter
-  def function_call_options(self, options):
-    """Returns function call options for current thread."""
-    self._thread_local_data.function_call_options = options
-
-  def num_gpus(self):
-    """The number of GPUs available to execute operations."""
-    self.ensure_initialized()
-    return self._num_gpus
-
-  def add_c_function(self, c_func):
-    """Add a C API TF_Function to the context.
-
-    Once added, the function (identified by its name) can be executed like any
-    other operation.
-
-    Args:
-      c_func: A wrapped TF_Function (returned from TF_GraphToFunction_wrapper).
-    """
-    self.ensure_initialized()
-    pywrap_tfe.TFE_ContextAddFunction(self._handle, c_func)
-
-  def get_c_function(self, name):
-    """Get a C API TF_Function from the context.
-
-    Args:
-      name: Name of the function to get.
-
-    Returns:
-      A ScopedTFFunction wrapping the C API TF_Function.
-    """
-    self.ensure_initialized()
-    return c_api_util.ScopedTFFunction(
-        pywrap_tfe.TFE_ContextGetFunction(self._handle, name), name
-    )
-
-  def add_function_def(self, fdef):
-    """Add a function definition to the context.
-
-    Once added, the function (identified by its name) can be executed like any
-    other operation.
-
-    Args:
-      fdef: A FunctionDef protocol buffer message.
-    """
-    self.ensure_initialized()
-    if is_oss:
-      fdef_string = fdef.SerializeToString()
-      pywrap_tfe.TFE_ContextAddFunctionDef(
-          self._handle, fdef_string, len(fdef_string)
-      )
-    else:
-      pywrap_tfe.TFE_ContextAddFunctionDefNoSerialization(self._handle, fdef)
-
-  def get_function_def(self, name):
-    """Get a function definition from the context.
-
-    Args:
-      name: function signature name.
-
-    Returns:
-      The requested FunctionDef.
-
-    Raises:
-      tf.errors.NotFoundError: if name is not the name of a registered function.
-    """
-    if is_oss:
-      with c_api_util.tf_buffer() as buffer_:
-        pywrap_tfe.TFE_ContextGetFunctionDef(self._handle, name, buffer_)
-        proto_data = pywrap_tf_session.TF_GetBuffer(buffer_)
-      function_def = function_pb2.FunctionDef()
-      function_def.ParseFromString(proto_data)
-    else:
-      function_def = pywrap_tfe.TFE_ContextGetFunctionDefNoSerialization(
-          self._handle, name
-      )
-    return function_def
-
-  def get_graph_debug_info(self, name):
-    """Get GraphDebugInfo associated with a function from the context.
-
-    Args:
-      name: function signature name.
-
-    Returns:
-      The requested GraphDebugInfo.
-
-    Raises:
-      tf.errors.NotFoundError: if name is not the name of a registered function.
-    """
-    with c_api_util.tf_buffer() as buffer_:
-      pywrap_tfe.TFE_ContextGetGraphDebugInfo(self._handle, name, buffer_)
-      proto_data = pywrap_tf_session.TF_GetBuffer(buffer_)
-    graph_debug_info = graph_debug_info_pb2.GraphDebugInfo()
-    graph_debug_info.ParseFromString(proto_data)
-
-    return graph_debug_info
-
-  def is_custom_device(self, device_name):
-    """Calls TFE_IsCustomDevice. See the non-member function."""
-    self.ensure_initialized()
-    return pywrap_tfe.TFE_Py_IsCustomDevice(self._handle, device_name)
-
-  def register_custom_device(
-      self, device_capsule, device_name, device_info_capsule
-  ):
-    """Calls TFE_RegisterCustomDevice. See the non-member function."""
-    self.ensure_initialized()
-    pywrap_tfe.TFE_Py_RegisterCustomDevice(
-        self._handle, device_capsule, device_name, device_info_capsule
-    )
-
-  def pack_eager_tensors(self, tensors):
-    """Pack multiple `EagerTensor`s of the same dtype and shape.
-
-    Args:
-      tensors: a list of EagerTensors to pack.
-
-    Returns:
-      A packed EagerTensor.
-    """
-    self.ensure_initialized()
-    return pywrap_tfe.TFE_Py_PackEagerTensors(self._handle, tensors)
-
-  def list_function_names(self):
-    """Get a list of names of registered functions.
-
-    Returns:
-      A set of names of all registered functions for the context.
-    """
-    self.ensure_initialized()
-    return set(pywrap_tfe.TFE_ContextListFunctionNames(self._handle))
-
-  def remove_function(self, name):
-    """Remove a function from the context.
-
-    Once removed, the function cannot be executed anymore.
-
-    Args:
-      name: function signature name.
-    """
-    self.ensure_initialized()
-    pywrap_tfe.TFE_ContextRemoveFunction(self._handle, name)
-
-  def has_function(self, name):
-    """Check if a function `name` is registered."""
-    self.ensure_initialized()
-    return bool(pywrap_tfe.TFE_ContextHasFunction(self._handle, name))
-
-  @property
-  def function_scope_id(self):
-    """Returns an id that is unique to each scope holding functions."""
-    return id(self._context_handle)
-
-  def call_function(self, name, tensor_inputs, num_outputs):
-    """Calls the function associated with the given name."""
-    attrs = tuple(
-        itertools.chain(*self.function_call_options.as_attrs().items())
-    )
-
-    cancellation_context = cancellation.context()
-    if cancellation_context is None:
-      outputs = execute.execute(
-          name.decode("utf-8"),
-          num_outputs=num_outputs,
-          inputs=tensor_inputs,
-          attrs=attrs,
-          ctx=self,
-      )
-    else:
-      outputs = execute.execute_with_cancellation(
-          name.decode("utf-8"),
-          num_outputs=num_outputs,
-          inputs=tensor_inputs,
-          attrs=attrs,
-          ctx=self,
-          cancellation_manager=cancellation_context,
-      )
-    # Empty list means no function outputs so return None
-    outputs = outputs or None
-
-    return outputs
-
-  def add_op_callback(self, callback):
-    """Add a post-op callback to the context.
-
-    A post-op callback is invoked immediately after an eager operation or
-    function has finished execution or after a op has been added to a graph,
-    providing access to the op's type, name input and output tensors. Multiple
-    op callbacks can be added, in which case the callbacks will be invoked in
-    the order in which they are added.
-
-    Args:
-      callback: a callable of the signature `f(op_type, inputs, attrs, outputs,
-        op_name=None, graph=None)`. See doc strings in `op_callbacks.py` for
-        details on the function signature and its semantics.
-    """
-    if callback not in self._thread_local_data.op_callbacks:
-      self._thread_local_data.op_callbacks.append(callback)
-
-  def remove_op_callback(self, callback):
-    """Remove an already-registered op callback.
-
-    Args:
-      callback: The op callback to be removed.
-
-    Raises:
-      KeyError: If `callback` is not already registered.
-    """
-    if callback not in self._thread_local_data.op_callbacks:
-      raise KeyError(
-          "The specified op callback has not been registered, "
-          "and hence cannot be removed."
-      )
-    del self._thread_local_data.op_callbacks[
-        self._thread_local_data.op_callbacks.index(callback)
-    ]
-
-  @property
-  def op_callbacks(self):
-    return self._thread_local_data.op_callbacks
-
-  @property
-  def invoking_op_callbacks(self):
-    return self._thread_local_data.invoking_op_callbacks
-
-  @invoking_op_callbacks.setter
-  def invoking_op_callbacks(self, value):
-    self._thread_local_data.invoking_op_callbacks = value
-
-  def _initialize_physical_devices(self, reinitialize=False):
-    """Gets local devices visible to the system.
-
-    Args:
-      reinitialize: If True, reinitializes self._physical_devices  so that
-        dynamic registered devices will also be visible to the python front-end.
-    """
-    # We lazy initialize self._physical_devices since we do not want to do this
-    # the constructor since the backend may not be initialized yet.
-    with self._device_lock:
-      if not reinitialize and self._physical_devices is not None:
-        return
-
-      devs = pywrap_tfe.TF_ListPhysicalDevices()
-      self._physical_devices = [
-          PhysicalDevice(name=d.decode(), device_type=d.decode().split(":")[1])
-          for d in devs
-      ]
-      self._physical_device_to_index = {
-          p: i for i, p in enumerate(self._physical_devices)
-      }
-      # We maintain a separate list just so we can check whether the device in
-      # _physical_devices is a PluggableDevice.
-      pluggable_devs = pywrap_tfe.TF_ListPluggablePhysicalDevices()
-      self._pluggable_devices = [
-          PhysicalDevice(name=d.decode(), device_type=d.decode().split(":")[1])
-          for d in pluggable_devs
-      ]
-
-      self._visible_device_list = list(self._physical_devices)
-      self._memory_growth_map = {
-          d: None
-          for d in self._physical_devices
-          if d.device_type == "GPU" or d in self._pluggable_devices
-      }
-
-    # Import device settings that may have been passed into the constructor
-    self._import_config()
-
-  def reinitialize_physical_devices(self):
-    """Gets local devices visible to the system."""
-    # Reinitialize the physical device list after registering
-    # the pluggable device.
-    self._initialize_physical_devices(True)
-
-  def list_physical_devices(self, device_type=None):
-    """List local devices visible to the system.
-
-    This API allows a client to query the devices before they have been
-    initialized by the eager runtime. Additionally a user can filter by device
-    type, to get only CPUs or GPUs.
-
-    Args:
-      device_type: Optional device type to limit results to
-
-    Returns:
-      List of PhysicalDevice objects.
-    """
-    self._initialize_physical_devices()
-
-    if device_type is None:
-      return list(self._physical_devices)
-
-    return [d for d in self._physical_devices if d.device_type == device_type]
-
-  def get_device_details(self, device):  # pylint: disable=redefined-outer-name
-    """Returns details about a physical devices.
-
-    Args:
-      device: A `tf.config.PhysicalDevice` returned by
-        `tf.config.list_physical_devices` or `tf.config.get_visible_devices`.
-
-    Returns:
-      A dict with string keys.
-    """
-    if not isinstance(device, PhysicalDevice):
-      raise ValueError(
-          "device must be a tf.config.PhysicalDevice, but got: %s" % (device,)
-      )
-    if (
-        self._physical_device_to_index is None
-        or device not in self._physical_device_to_index
-    ):
-      raise ValueError(
-          "The PhysicalDevice must be one obtained from "
-          "calling `tf.config.list_physical_devices`, but got: "
-          "%s" % (device,)
-      )
-    index = self._physical_device_to_index[device]
-    details = pywrap_tfe.TF_GetDeviceDetails(index)
-
-    # Change compute_capability from a string to a tuple
-    if "compute_capability" in details:
-      try:
-        major, minor = details["compute_capability"].split(".")
-        details["compute_capability"] = (int(major), int(minor))
-      except ValueError as exc:
-        raise RuntimeError(
-            "Device returned compute capability an in invalid format: %s"
-            % details["compute_capability"]
-        ) from exc
-    return details
-
-  def _import_config(self):
-    """Import config if passed in during construction.
-
-    If Context was created with a ConfigProto such as when calling
-    tf.compat.v1.enable_eager_execution(), then we need to pull out the
-    various pieces we might be replacing and import then into our internal
-    class representation.
-    """
-    if self._config is None:
-      return
-
-    num_cpus = self._config.device_count.get("CPU", 1)
-    if num_cpus != 1:
-      cpus = [d for d in self._physical_devices if d.device_type == "CPU"]
-      if num_cpus == 0:
-        self.set_visible_devices([], "CPU")
-      elif num_cpus > 1:
-        self.set_logical_device_configuration(
-            cpus[0], [LogicalDeviceConfiguration() for _ in range(num_cpus)]
-        )
-
-    # Parse GPU options
-    gpus = [d for d in self._physical_devices if d.device_type == "GPU"]
-
-    # If there are no GPUs detected, simply ignore all the GPU options passed in
-    # rather than doing any validation checks.
-    if not gpus:
-      return
-
-    gpu_count = self._config.device_count.get("GPU", None)
-
-    visible_gpus = []
-    # TODO(gjn): Handle importing existing virtual GPU configuration
-    visible_indices = self._config.gpu_options.visible_device_list
-    if visible_indices:
-      for index in visible_indices.split(","):
-        if int(index) >= len(gpus):
-          raise ValueError("Invalid visible device index: %s" % index)
-        visible_gpus.append(gpus[int(index)])
-    else:
-      visible_gpus = gpus
-
-    if gpu_count is not None:
-      visible_gpus = visible_gpus[:gpu_count]
-
-    self.set_visible_devices(visible_gpus, "GPU")
-
-  def list_logical_devices(self, device_type=None):
-    """Return logical devices."""
-    self.ensure_initialized()
-    if device_type is None:
-      return list(self._logical_devices)
-
-    return [d for d in self._logical_devices if d.device_type == device_type]
-
-  def get_visible_devices(self, device_type=None):
-    """Get the list of visible devices."""
-    self._initialize_physical_devices()
-
-    if device_type is None:
-      return list(self._visible_device_list)
-
-    return [
-        d for d in self._visible_device_list if d.device_type == device_type
-    ]
-
-  def set_visible_devices(self, devices, device_type=None):
-    """Set the list of visible devices."""
-    self._initialize_physical_devices()
-
-    if not isinstance(devices, list):
-      devices = [devices]
-
-    for d in devices:
-      if d not in self._physical_devices:
-        raise ValueError("Unrecognized device: %s" % repr(d))
-      if device_type is not None and d.device_type != device_type:
-        raise ValueError("Unrecognized device: %s" % repr(d))
-
-    visible_device_list = []
-    if device_type is not None:
-      visible_device_list = [
-          d for d in self._visible_device_list if d.device_type != device_type
-      ]
-
-    visible_device_list += devices
-
-    if self._visible_device_list == visible_device_list:
-      return
-
-    if self._context_handle is not None:
-      raise RuntimeError(
-          "Visible devices cannot be modified after being initialized"
-      )
-
-    self._visible_device_list = visible_device_list
-
-  def get_memory_info(self, dev):
-    """Returns a dict of memory info for the device."""
-    self._initialize_physical_devices()
-    self.ensure_initialized()
-    return pywrap_tfe.TFE_GetMemoryInfo(self._context_handle, dev)
-
-  def reset_memory_stats(self, dev):
-    """Resets the tracked memory stats for the device."""
-    self._initialize_physical_devices()
-    self.ensure_initialized()
-    pywrap_tfe.TFE_ResetMemoryStats(self._context_handle, dev)
-
-  def get_memory_growth(self, dev):
-    """Get if memory growth is enabled for a PhysicalDevice."""
-    self._initialize_physical_devices()
-
-    if dev not in self._physical_devices:
-      raise ValueError("Unrecognized device: %s" % repr(dev))
-
-    return self._memory_growth_map[dev]
-
-  def set_memory_growth(self, dev, enable):
-    """Set if memory growth should be enabled for a PhysicalDevice."""
-    self._initialize_physical_devices()
-
-    if dev not in self._physical_devices:
-      raise ValueError("Unrecognized device: %s" % repr(dev))
-
-    if dev in self._virtual_device_map:
-      raise ValueError(
-          "Cannot set memory growth on device when virtual devices configured"
-      )
-
-    if dev.device_type != "GPU" and dev not in self._pluggable_devices:
-      raise ValueError(
-          "Cannot set memory growth on non-GPU and non-Pluggable devices"
-      )
-
-    if self._memory_growth_map.get(dev) == enable:
-      return
-
-    if self._context_handle is not None:
-      raise RuntimeError(
-          "Physical devices cannot be modified after being initialized"
-      )
-
-    self._memory_growth_map[dev] = enable
-
-  def get_logical_device_configuration(self, dev):
-    """Get the virtual device configuration for a PhysicalDevice."""
-    self._initialize_physical_devices()
-
-    if dev not in self._physical_devices:
-      raise ValueError("Unrecognized device: %s" % repr(dev))
-
-    return self._virtual_device_map.get(dev)
-
-  def set_logical_device_configuration(self, dev, virtual_devices):
-    """Set the virtual device configuration for a PhysicalDevice."""
-    self._initialize_physical_devices()
-
-    if dev not in self._physical_devices:
-      raise ValueError("Unrecognized device: %s" % repr(dev))
-
-    if dev.device_type == "CPU":
-      for vdev in virtual_devices:
-        if vdev.memory_limit is not None:
-          raise ValueError(
-              "Setting memory limit on CPU virtual devices is "
-              "currently not supported"
-          )
-        if vdev.experimental_priority is not None:
-          raise ValueError(
-              "Setting experimental_priority on CPU virtual "
-              " devices is currently not supported"
-          )
-        if vdev.experimental_device_ordinal is not None:
-          raise ValueError(
-              "Setting experimental_device_ordinal on CPU virtual "
-              " devices is currently not supported"
-          )
-    elif dev.device_type == "GPU":
-      for vdev in virtual_devices:
-        if vdev.memory_limit is None:
-          raise ValueError(
-              "Setting memory limit is required for GPU virtual devices"
-          )
-    else:
-      raise ValueError(
-          "Virtual devices are not supported for %s" % dev.device_type
-      )
-
-    if self._virtual_device_map.get(dev) == virtual_devices:
-      return
-
-    if self._context_handle is not None:
-      raise RuntimeError(
-          "Virtual devices cannot be modified after being initialized"
-      )
-
-    self._virtual_device_map[dev] = virtual_devices
-
-  def set_logical_cpu_devices(self, num_cpus, prefix=""):
-    """Set virtual CPU devices in context.
-
-    If virtual CPU devices are already configured at context initialization
-    by tf.config.set_logical_device_configuration(), this method should not be
-    called.
-
-    Args:
-      num_cpus: Number of virtual CPUs.
-      prefix: Device name prefix.
-
-    Raises:
-     RuntimeError: If virtual CPUs are already configured at context
-     initialization.
-    """
-    server_def = self._server_def or self._collective_ops_server_def
-    local_prefix = ["/device"]
-    if server_def is not None:
-      local_prefix.append(
-          "/job:%s/replica:0/task:%d"
-          % (server_def.job_name, server_def.task_index)
-      )
-    logical_local_devices = [
-        d
-        for d in self.list_logical_devices("CPU")
-        if d.name.startswith(tuple(local_prefix))
-    ]
-    self.ensure_initialized()
-    # Error out if there are already multiple logical CPU in the context.
-    if len(logical_local_devices) > 1:
-      raise RuntimeError("Virtual CPUs already set, cannot modify again.")
-
-    pywrap_tfe.TFE_SetLogicalCpuDevices(self._context_handle, num_cpus, prefix)
-    self._initialize_logical_devices()
-
-  def get_compiler_ir(
-      self,
-      device_name,
-      platform_name,
-      function_name,
-      flat_args,
-      captured_inputs,
-      stage="hlo",
-  ):
-    """Get the compiler IR bytes.
-
-    Args:
-      device_name: The name of the device with the form as
-        "/job:localhost/replica:0/task:0/device:CPU:0", "/device:TPU:0" etc.
-        When this is used, actual device is needed for getting the compiler IR.
-      platform_name: The name of the platform, e.g. "TPU". When this is used,
-        first we find a device whose name contains the platform, if it is found
-        we get the compiler IR by device. Otherwise the compiler IR is obtained
-        as if using that device. The former logic of falling back to device is
-        necessary, as there are cases of TF variables that need to access
-        devices, but the upper layer may generally choose platform for getting
-        compiler IR in a device-agnostic way.
-      function_name: The name of the function to get the compiler IR.
-      flat_args: The flat argument inputs.
-      captured_inputs: The inputs that are captured.
-      stage: The exported stage for the given function.
-
-    Returns:
-      The compiler IR bytes.
-    """
-    return pywrap_tfe.TF_GetCompilerIr(
-        self._context_handle,
-        function_name,
-        stage,
-        device_name,
-        flat_args,
-        captured_inputs,
-        platform_name,
-    )
-
-  @deprecated(
-      None, "XLA:CPU and XLA:GPU devices are deprecated", warn_once=True
-  )
-  def enable_xla_devices(self):
-    """Enables XLA:CPU and XLA:GPU devices registration."""
-    pywrap_tfe.TF_EnableXlaDevices()
-
-  @property
-  def enable_mlir_bridge(self):
-    return pywrap_tfe.TF_IsMlirBridgeEnabled()
-
-  @property
-  def enable_mlir_graph_optimization(self):
-    return self._enable_mlir_graph_optimization
-
-  @enable_mlir_bridge.setter
-  def enable_mlir_bridge(self, enabled):
-    pywrap_tfe.TF_EnableMlirBridge(enabled)
-    self._thread_local_data.function_call_options = None
-
-  @enable_mlir_graph_optimization.setter
-  def enable_mlir_graph_optimization(self, enabled):
-    self._enable_mlir_graph_optimization = enabled
-    self._thread_local_data.function_call_options = None
-
-  @property
-  def optimizer_jit(self):
-    level = self.config.graph_options.optimizer_options.global_jit_level
-    return (
-        level == config_pb2.OptimizerOptions.ON_1
-        or level == config_pb2.OptimizerOptions.ON_2
-    )
-
-  @optimizer_jit.setter
-  def optimizer_jit(self, enabled):
-    self._optimizer_jit = enabled
-
-    self._thread_local_data.function_call_options = None
-
-  def get_optimizer_experimental_options(self):
-    """Get experimental options for the optimizer.
-
-    Returns:
-      Dictionary of current option values
-    """
-    rewrite_options = self.config.graph_options.rewrite_options
-    options = {}
-
-    def rewriter_toggle(option):
-      attr = getattr(rewrite_options, option)
-      if attr != 0:
-        options[option] = attr == rewriter_config_pb2.RewriterConfig.ON
-
-    def rewriter_bool(option):
-      options[option] = getattr(rewrite_options, option)
-
-    rewriter_toggle("layout_optimizer")
-    rewriter_toggle("constant_folding")
-    rewriter_toggle("shape_optimization")
-    rewriter_toggle("remapping")
-    rewriter_toggle("arithmetic_optimization")
-    rewriter_toggle("dependency_optimization")
-    rewriter_toggle("loop_optimization")
-    rewriter_toggle("function_optimization")
-    rewriter_toggle("debug_stripper")
-    rewriter_bool("disable_model_pruning")
-    rewriter_toggle("scoped_allocator_optimization")
-    rewriter_toggle("pin_to_host_optimization")
-    rewriter_toggle("implementation_selector")
-    rewriter_toggle("auto_mixed_precision")
-    rewriter_toggle("use_plugin_optimizers")
-    rewriter_bool("disable_meta_optimizer")
-    rewriter_toggle("auto_mixed_precision_onednn_bfloat16")
-    rewriter_toggle("auto_mixed_precision_mkl")
-
-    if rewrite_options.min_graph_nodes != 0:
-      options["min_graph_nodes"] = rewrite_options.min_graph_nodes
-
-    return options
-
-  def set_optimizer_experimental_options(self, options):
-    """Set experimental options for the optimizer.
-
-    Args:
-      options: Dictionary of options to modify
-    """
-    self._optimizer_experimental_options.update(options)
-
-    self._thread_local_data.function_call_options = None
-
-  def _validate_thread_count(self, num_threads):
-    """Validate thread count parameter.
-
-    Args:
-      num_threads: The number of threads to validate.
-
-    Raises:
-      TypeError: If num_threads is not an integer (including booleans).
-      ValueError: If num_threads is not in the range [0, 10000].
-    """
-    if isinstance(num_threads, bool) or not isinstance(
-        num_threads, compat.integral_types
-    ):
-      raise TypeError(
-          f"num_threads must be an integer, got {type(num_threads).__name__}."
-      )
-    if not (0 <= num_threads <= 10000):
-      raise ValueError(
-          f"num_threads must be in the range [0, 10000], got {num_threads}. "
-          "Setting an excessively large number of threads can cause system "
-          "instability or crashes."
-      )
-
-  @property
-  def intra_op_parallelism_threads(self):
-    return self.config.intra_op_parallelism_threads
-
-  @intra_op_parallelism_threads.setter
-  def intra_op_parallelism_threads(self, num_threads):
-    self._validate_thread_count(num_threads)
-    if self._intra_op_parallelism_threads == num_threads:
-      return
-
-    if self._context_handle is not None:
-      raise RuntimeError(
-          "Intra op parallelism cannot be modified after initialization."
-      )
-
-    self._intra_op_parallelism_threads = num_threads
-
-  @property
-  def inter_op_parallelism_threads(self):
-    return self.config.inter_op_parallelism_threads
-
-  @inter_op_parallelism_threads.setter
-  def inter_op_parallelism_threads(self, num_threads):
-    # None keeps the default behavior.
-    if num_threads is not None:
-      self._validate_thread_count(num_threads)
-    if self._inter_op_parallelism_threads == num_threads:
-      return
-
-    if self._context_handle is not None:
-      raise RuntimeError(
-          "Inter op parallelism cannot be modified after initialization."
-      )
-
-    self._inter_op_parallelism_threads = num_threads
-
-  @property
-  def soft_device_placement(self):
-    return self.config.allow_soft_placement
-
-  @soft_device_placement.setter
-  def soft_device_placement(self, enable):
-    if self._context_handle is not None:
-      pywrap_tfe.TFE_ContextSetSoftDevicePlacement(self._handle, enable)
-
-    self._soft_device_placement = enable
-    self._thread_local_data.function_call_options = None
-
-  @property
-  def log_device_placement(self):
-    return self.config.log_device_placement
-
-  @log_device_placement.setter
-  def log_device_placement(self, enable):
-    if self._context_handle is not None:
-      pywrap_tfe.TFE_ContextSetLogDevicePlacement(self._handle, enable)
-
-    self._log_device_placement = enable
-    self._thread_local_data.function_call_options = None
-
-  @property
-  def jit_compile_rewrite(self):
-    return self._jit_compile_rewrite
-
-  @jit_compile_rewrite.setter
-  def jit_compile_rewrite(self, enable):
-    if self._context_handle is not None:
-      pywrap_tfe.TFE_ContextSetJitCompileRewrite(self._handle, enable)
-    self._jit_compile_rewrite = enable
-
-  @property
-  def xla_sharding_for_resource_variables(self):
-    return self._xla_sharding_for_resource_variables
-
-  @xla_sharding_for_resource_variables.setter
-  def xla_sharding_for_resource_variables(self, enable):
-    self._xla_sharding_for_resource_variables = enable
-
-  @property
-  def optionals_in_gradients(self):
-    return self._optionals_in_gradients
-
-  @optionals_in_gradients.setter
-  def optionals_in_gradients(self, enable):
-    self._optionals_in_gradients = enable
-
-  @property
-  def device_policy(self):
-    # Only get the policy from the context if it has already been initialized
-    if self._context_handle is not None:
-      return pywrap_tfe.TFE_ContextGetDevicePlacementPolicy(self._handle)
-
-    return self._device_policy
-
-  @device_policy.setter
-  def device_policy(self, policy):
-    if policy is None:
-      policy = DEVICE_PLACEMENT_SILENT
-
-    if self._device_policy != policy:
-      self._device_policy = policy
-
-      # Only set the policy if the context has already been initialized
-      if self._context_handle is not None:
-        pywrap_tfe.TFE_ContextSetThreadLocalDevicePlacementPolicy(
-            self._handle, self._device_policy
-        )
-
-  @property
-  def use_tfrt(self):
-    return self._use_tfrt
-
-  @use_tfrt.setter
-  def use_tfrt(self, tfrt):
-    """Sets whether to use TFRT."""
-    if not isinstance(tfrt, bool):
-      raise ValueError("Expecting a boolean but got %s" % type(tfrt))
-
-    if self._use_tfrt != tfrt:
-      if self._initialized:
-        raise ValueError("use_tfrt should be set before being initialized.")
-      self._use_tfrt = tfrt
-
-  @property
-  def operation_timeout_in_ms(self):
-    return self.config.operation_timeout_in_ms
-
-  @operation_timeout_in_ms.setter
-  def operation_timeout_in_ms(self, timeout_in_ms):
-    if self._operation_timeout_in_ms == timeout_in_ms:
-      return
-
-    if self._context_handle is not None:
-      raise RuntimeError(
-          "Operation timeout cannot be modified after initialization."
-      )
-
-    self._operation_timeout_in_ms = timeout_in_ms
-
-  def enable_run_metadata(self):
-    """Enables tracing of op execution via RunMetadata.
-
-    To retrieve the accumulated metadata call context.export_run_metadata()
-    and to stop tracing call context.disable_run_metadata().
-    """
-    self.ensure_initialized()
-    pywrap_tfe.TFE_ContextEnableRunMetadata(self._handle)
-
-  def disable_run_metadata(self):
-    """Disables tracing of op execution via RunMetadata."""
-    if not self._context_handle:
-      return
-    pywrap_tfe.TFE_ContextDisableRunMetadata(self._context_handle)
-
-  def enable_graph_collection(self):
-    """Enables graph collection of executed functions.
-
-    To retrieve the accumulated graphs call context.export_run_metadata()
-    and to stop collecting graphs call context.disable_graph_collection().
-    """
-    self.ensure_initialized()
-    pywrap_tfe.TFE_ContextEnableGraphCollection(self._handle)
-
-  def disable_graph_collection(self):
-    """Disables graph collection of executed functions."""
-    if not self._context_handle:
-      return
-    pywrap_tfe.TFE_ContextDisableGraphCollection(self._context_handle)
-
-  def export_run_metadata(self):
-    """Returns a RunMetadata proto with accumulated information.
-
-    The returned protocol buffer contains information since the most recent call
-    to either enable_run_metadata or export_run_metadata.
-
-    Returns:
-      A RunMetadata protocol buffer. Or None if not enabled.
-    """
-    if not self._context_handle:
-      return None
-    with c_api_util.tf_buffer() as buffer_:
-      pywrap_tfe.TFE_ContextExportRunMetadata(self._context_handle, buffer_)
-      proto_data = pywrap_tf_session.TF_GetBuffer(buffer_)
-    run_metadata = config_pb2.RunMetadata()
-    run_metadata.ParseFromString(compat.as_bytes(proto_data))
-    return run_metadata
-
-  def set_server_def_retries(self, retries):
-    """Set the number of retries to use when calling SetServerDef.
-
-    In cases where many servers run in high-preemption environments, jobs could
-    be preempted during startup and initial connection via SetServerDef. Retries
-    allow for more robust connection in these environments.
-
-    Args:
-      retries: int specifying the number of connection retries before failing.
-        Retries follow an exponential backoff waiting period with min value 1ms,
-        max value 10s, and exponent 1.3.
-    """
-    self._set_server_def_retries = retries
-
-  @property
-  def context_switches(self):
-    """Returns a stack of context switches."""
-    return self._context_switches
-
-
-class _EagerDeviceContext(object):
-  """Context-manager forcing placement of ops and Tensors on a device."""
-
-  __slots__ = ["_device_name", "_ctx", "_stack"]
-
-  def __init__(self, ctx, device_name):
-    self._device_name = device_name
-    self._ctx = ctx
-    self._stack = []
-
-  # TODO(b/189233748): Consolidate the device string parsing logic with
-  # tensorflow/core/util/device_name_utils.cc.
-  def __enter__(self):
-    ctx = self._ctx
-    old_device_name = ctx.device_name
-    old_device_spec = ctx.device_spec
-    new_device_name = self._device_name
-    cache_key = (old_device_name, new_device_name)
-    try:
-      new_device_name, new_device_spec = _device_parsing_cache[cache_key]
-    except TypeError as exc:
-      # Error while trying to compute the cache key.
-      raise ValueError(
-          "Expecting a string device name. Got %s(%s)"
-          % (type(new_device_name), new_device_name)
-      ) from exc
-    except KeyError as exc:
-      # Handle a cache miss.
-      if new_device_name is not None:
-        if not isinstance(new_device_name, str):
-          raise ValueError(
-              "Expecting a string device name. Got %s(%s)"
-              % (type(new_device_name), new_device_name)
-          ) from exc
-        device_spec = pydev.DeviceSpec.from_string(new_device_name)
-        if old_device_name:
-          new_device_spec = copy.copy(old_device_spec)
-        else:
-          ctx.ensure_initialized()
-          new_device_spec = pydev.DeviceSpec.from_string(
-              ctx._context_devices[0]
-          )  # pylint: disable=protected-access
-        new_device_spec = new_device_spec.make_merged_spec(device_spec)
-      else:
-        new_device_spec = pydev.DeviceSpec.from_string("")
-      new_device_name = new_device_spec.to_string()
-      _device_parsing_cache[cache_key] = (new_device_name, new_device_spec)
-
-    ctx._set_device(new_device_name, new_device_spec)  # pylint: disable=protected-access
-    self._stack.append((old_device_name, old_device_spec, new_device_spec))
-
-  def __exit__(self, *ex_info):
-    ctx = self._ctx
-    old_device_name, old_device_spec, new_device_spec = self._stack[-1]
-    if ctx.device_spec is not new_device_spec:
-      raise RuntimeError("Exiting device scope without proper scope nesting")
-    del self._stack[-1]
-    ctx._set_device(old_device_name, old_device_spec)  # pylint: disable=protected-access
-
-
-# Do not change directly.
-_context = None
-_context_lock = threading.Lock()
-
-
-def _set_context_locked(ctx):
-  global _context
-  pywrap_tfe.TFE_Py_SetEagerContext(ctx)
-  ctx.mark_as_global_context()
-  _context = ctx
-
-
-def _set_context(ctx):
-  with _context_lock:
-    _set_context_locked(ctx)
-
-
-def _create_context():
-  with _context_lock:
-    if _context is None:
-      ctx = Context()
-      _set_context_locked(ctx)
-
-
-def _reset_context():
-  """Clears and re-initializes the singleton context.
-
-  Should only be used for testing.
-  """
-  global _context
-  global _device_parsing_cache
-
-  # Garbage collect and clear scalar cache to avoid Tensor from current context
-  # polluting next context.
-  gc.collect()
-  pywrap_tfe.TFE_ClearScalarCache()
-  with _context_lock:
-    if _context is not None:
-      _context._clear_caches()
-      _context = None
-  _create_context()
-  _device_parsing_cache = {}
-
-
-def _reset_jit_compiler_flags():
-  """Clears and re-initializes the TF JIT compiler flags.
-
-  Should only be used for testing.
-  """
-  pywrap_tfe.TF_ResetJitCompilerFlags()
-
-
-def context():
-  """Returns a singleton context object."""
-  if _context is None:
-    _create_context()
-  return _context
-
-
-def context_safe():
-  """Returns current context (or None if one hasn't been initialized)."""
-  return _context
-
-
-def ensure_initialized():
-  """Initialize the context."""
-  context().ensure_initialized()
-
-
-def initialize_logical_devices():
-  """Initialize the virtual devices."""
-  context()._initialize_logical_devices()  # pylint: disable=protected-access
-
-
-def set_global_seed(seed):
-  """Sets the eager mode seed."""
-  context()._set_global_seed(seed)  # pylint: disable=protected-access
-
-
-def global_seed():
-  """Returns the eager mode seed."""
-  return context()._seed  # pylint: disable=protected-access
-
-
-def internal_operation_seed():
-  """Returns the operation seed generated based on global seed."""
-  return context()._internal_operation_seed()  # pylint: disable=protected-access
-
-
-@tf_export("executing_eagerly", v1=[])
-def executing_eagerly():
-  """Checks whether the current thread has eager execution enabled.
-
-  Eager execution is enabled by default and this API returns `True`
-  in most of cases. However, this API might return `False` in the following use
-  cases.
-
-  *  Executing inside `tf.function`, unless under `tf.init_scope` or
-     `tf.config.run_functions_eagerly(True)` is previously called.
-  *  Executing inside a transformation function for `tf.dataset`.
-  *  `tf.compat.v1.disable_eager_execution()` is called.
-
-  General case:
-
-  >>> print(tf.executing_eagerly())
-  True
-
-  Inside `tf.function`:
-
-  >>> @tf.function
-  ... def fn():
-  ...   with tf.init_scope():
-  ...     print(tf.executing_eagerly())
-  ...   print(tf.executing_eagerly())
-  >>> fn()
-  True
-  False
-
-  Inside `tf.function` after `tf.config.run_functions_eagerly(True)` is called:
-
-  >>> tf.config.run_functions_eagerly(True)
-  >>> @tf.function
-  ... def fn():
-  ...   with tf.init_scope():
-  ...     print(tf.executing_eagerly())
-  ...   print(tf.executing_eagerly())
-  >>> fn()
-  True
-  True
-  >>> tf.config.run_functions_eagerly(False)
-
-  Inside a transformation function for `tf.dataset`:
-
-  >>> def data_fn(x):
-  ...   print(tf.executing_eagerly())
-  ...   return x
-  >>> dataset = tf.data.Dataset.range(100)
-  >>> dataset = dataset.map(data_fn)
-  False
-
-  Returns:
-    `True` if the current thread has eager execution enabled.
-  """
-  ctx = context_safe()
-  if ctx is None:
-    return default_execution_mode == EAGER_MODE
-
-  return ctx.executing_eagerly()
-
-
-@tf_export(v1=["executing_eagerly"])
-def executing_eagerly_v1():
-  """Checks whether the current thread has eager execution enabled.
-
-  Eager execution is typically enabled via
-  `tf.compat.v1.enable_eager_execution`, but may also be enabled within the
-  context of a Python function via tf.contrib.eager.py_func.
-
-  When eager execution is enabled, returns `True` in most cases. However,
-  this API might return `False` in the following use cases.
-
-  *  Executing inside `tf.function`, unless under `tf.init_scope` or
-     `tf.config.run_functions_eagerly(True)` is previously called.
-  *  Executing inside a transformation function for `tf.dataset`.
-  *  `tf.compat.v1.disable_eager_execution()` is called.
-
-  >>> tf.compat.v1.enable_eager_execution()
-
-  General case:
-
-  >>> print(tf.executing_eagerly())
-  True
-
-  Inside `tf.function`:
-
-  >>> @tf.function
-  ... def fn():
-  ...   with tf.init_scope():
-  ...     print(tf.executing_eagerly())
-  ...   print(tf.executing_eagerly())
-  >>> fn()
-  True
-  False
-
-  Inside `tf.function`
-  after  `tf.config.run_functions_eagerly(True)` is called:
-
-  >>> tf.config.run_functions_eagerly(True)
-  >>> @tf.function
-  ... def fn():
-  ...   with tf.init_scope():
-  ...     print(tf.executing_eagerly())
-  ...   print(tf.executing_eagerly())
-  >>> fn()
-  True
-  True
-  >>> tf.config.run_functions_eagerly(False)
-
-  Inside a transformation function for `tf.dataset`:
-
-  >>> def data_fn(x):
-  ...   print(tf.executing_eagerly())
-  ...   return x
-  >>> dataset = tf.data.Dataset.range(100)
-  >>> dataset = dataset.map(data_fn)
-  False
-
-  Returns:
-    `True` if the current thread has eager execution enabled.
-  """
-  return executing_eagerly()
-
-
-def in_eager_mode():
-  """Use executing_eagerly() instead. This function will be removed."""
-  return executing_eagerly()
-
-
-def anonymous_name():
-  """Returns the anonymous shared name.
-
-  In eager mode we create anonymous resources to avoid spurious sharing issues.
-  The runtime generates a unique name on our behalf when the reserved
-  anonymous shared name is used as a shared name.
-
-  Returns:
-    The anonymous shared name.
-  """
-
-  # The magic value is defined as
-  # `tensorflow::ResourceHandle::ANONYMOUS_NAME` in C++.
-  return "cd2c89b7-88b7-44c8-ad83-06c2a9158347"
-
-
-def graph_mode():
-  """Context-manager to disable eager execution for the current thread."""
-  return context()._mode(GRAPH_MODE)  # pylint: disable=protected-access
-
-
-# Used by b/167638505 for keras backend API and Lambda layer.
-@tf_export("__internal__.eager_context.eager_mode", v1=[])
-def eager_mode():
-  """Context-manager to enable eager execution for the current thread."""
-  return context()._mode(EAGER_MODE)  # pylint: disable=protected-access
-
-
-def scope_name():
-  """Name of the current scope."""
-  return context().scope_name
-
-
-def device(name):
-  """Context-manager to force placement of operations and Tensors on a device.
-
-  Example:
-  ```python
-  with tf.device('gpu:0'):
-    with tf.device('cpu:0'):
-      shape = tf.constant([], dtype=tf.int32)
-    x = tf.random.truncated_normal(shape, tf.float32)
-  ```
-  will ensure that the `shape` Tensor is on CPU but the `truncated_normal`
-  operation runs on GPU 0.
-
-  Args:
-    name: Name of the device (see context().devices()), or None to perform
-      automatic placement.
-
-  Returns:
-    Context manager for setting the device.
-  """
-  ensure_initialized()
-  return context().device(name)
-
-
-# Expose some properties of Context as internally public APIs (b/160348781).
-@tf_export("__internal__.eager_context.get_config", v1=[])
-def get_config():
-  """Get the ConfigProto of Context.
-
-  Returns:
-    The ConfigProto of Context.
-  """
-  return context().config
-
-
-@tf_export("__internal__.eager_context.get_device_name", v1=[])
-def get_device_name():
-  """Get the device name for the current thread.
-
-  Returns:
-    The device name for the current thread.
-  """
-  return context().device_name
-
-
-@tf_export("__internal__.eager_context.set_soft_device_placement", v1=[])
-def set_soft_device_placement(enabled):
-  """Set if soft device placements should be allowed.
-
-  Args:
-    enabled: Whether to enable soft device placement.
-  """
-  context().soft_device_placement = enabled
-
-
-@tf_export("__internal__.eager_context.get_executor", v1=[])
-def get_executor():
-  """Get the Executor of the current thread.
-
-  Returns:
-    The Executor of the current thread.
-  """
-  return context().executor
-
-
-@tf_export("debugging.get_log_device_placement")
-def get_log_device_placement():
-  """Get if device placements are logged.
-
-  Returns:
-    If device placements are logged.
-  """
-  return context().log_device_placement
-
-
-@tf_export("debugging.set_log_device_placement")
-def set_log_device_placement(enabled):
-  """Turns logging for device placement decisions on or off.
-
-  Operations execute on a particular device, producing and consuming tensors on
-  that device. This may change the performance of the operation or require
-  TensorFlow to copy data to or from an accelerator, so knowing where operations
-  execute is useful for debugging performance issues.
-
-  For more advanced profiling, use the [TensorFlow
-  profiler](https://www.tensorflow.org/guide/profiler).
-
-  Device placement for operations is typically controlled by a `tf.device`
-  scope, but there are exceptions, for example operations on a `tf.Variable`
-  which follow the initial placement of the variable. Turning off soft device
-  placement (with `tf.config.set_soft_device_placement`) provides more explicit
-  control.
-
-  >>> tf.debugging.set_log_device_placement(True)
-  >>> tf.ones([])
-  >>> # [...] op Fill in device /job:localhost/replica:0/task:0/device:GPU:0
-  >>> with tf.device("CPU"):
-  ...  tf.ones([])
-  >>> # [...] op Fill in device /job:localhost/replica:0/task:0/device:CPU:0
-  >>> tf.debugging.set_log_device_placement(False)
-
-  Turning on `tf.debugging.set_log_device_placement` also logs the placement of
-  ops inside `tf.function` when the function is called.
-
-  Args:
-    enabled: Whether to enabled device placement logging.
-  """
-  context().log_device_placement = enabled
-
-
-@tf_contextlib.contextmanager
-def device_policy(policy):
-  """Context manager for setting device placement policy for current thread."""
-  ctx = context()
-  old_policy = ctx.device_policy
-  try:
-    ctx.device_policy = policy
-    yield
-  finally:
-    ctx.device_policy = old_policy
-
-
-def set_execution_mode(mode):
-  """Sets execution mode for the current thread."""
-  context().execution_mode = mode
-
-
-# TODO(fishx): remove this method.
-@tf_contextlib.contextmanager
-def execution_mode(mode):
-  """Context manager for setting execution mode for current thread."""
-  if mode is None:
-    yield
-  else:
-    ctx = context()
-    executor_new = executor.new_executor(mode == ASYNC)
-    executor_old = ctx.executor
-    try:
-      executor_old.wait()
-      ctx.executor = executor_new
-      yield
-    finally:
-      ctx.executor = executor_old
-      executor_new.wait()
-
-
-@tf_contextlib.contextmanager
-def executor_scope(e):
-  """Context manager for changing executor for current thread.
-
-  Args:
-    e: A Executor to execute eager ops under this scope. Setting it to None will
-      switch back to use the default executor for the context.
-
-  Yields:
-    Context manager for setting the executor for current thread.
-  """
-  ctx = context()
-  executor_old = ctx.executor
-  try:
-    ctx.executor = e
-    yield
-  finally:
-    ctx.executor = executor_old
-
-
-@tf_export("experimental.function_executor_type")
-@tf_contextlib.contextmanager
-def function_executor_type(executor_type):
-  """Context manager for setting the executor of eager defined functions.
-
-  Eager defined functions are functions decorated by tf.contrib.eager.defun.
-
-  Args:
-    executor_type: a string for the name of the executor to be used to execute
-      functions defined by tf.contrib.eager.defun.
-
-  Yields:
-    Context manager for setting the executor of eager defined functions.
-  """
-  current_options = context().function_call_options
-  old_options = copy.copy(current_options)
-  try:
-    current_options.executor_type = executor_type
-    yield
-  finally:
-    context().function_call_options = old_options
-
-
-def is_async():
-  """Returns true if current thread is in async mode."""
-  return context().is_async()
-
-
-def num_gpus():
-  """Get the number of available GPU devices.
-
-  Returns:
-    The number of available GPU devices.
-  """
-  return context().num_gpus()
-
-
-def enable_run_metadata():
-  """Enables tracing of op execution via RunMetadata.
-
-  To retrieve the accumulated metadata call context.export_run_metadata()
-  and to stop tracing call context.disable_run_metadata().
-  """
-  context().enable_run_metadata()
-
-
-def disable_run_metadata():
-  """Disables tracing of op execution via RunMetadata."""
-  context().disable_run_metadata()
-
-
-def enable_graph_collection():
-  """Enables graph collection of executed functions.
-
-  To retrieve the accumulated graphs call context.export_run_metadata()
-  and to stop collecting graphs call context.disable_graph_collection().
-  """
-  context().enable_graph_collection()
-
-
-def disable_graph_collection():
-  """Disables graph collection of executed functions."""
-  context().disable_graph_collection()
-
-
-def export_run_metadata():
-  """Returns a RunMetadata proto with accumulated information.
-
-  The returned protocol buffer contains information since the most recent call
-  to either enable_run_metadata or export_run_metadata.
-
-  Returns:
-    A RunMetadata protocol buffer.
-  """
-  return context().export_run_metadata()
-
-
-@contextlib.contextmanager
-def collect_graphs(optimized=True):
-  """Collects a flat list of pre- or post-optimization graphs.
-
-  The collected graphs include device placements, which can be useful for
-  testing.
-
-  Usage:
-
-  ```
-  @def_function.function
-  def f(x):
-    return x + constant_op.constant(1.)
-
-  with context.collect_graphs() as graphs:
-    with ops.device("CPU:0"):
-      f(constant_op.constant(1.))
-
-  graph, = graphs  # `graph` contains a single GraphDef for inspection
-  ```
-
-  Args:
-    optimized: whether to collect optimized graphs or non-optimized graphs
-
-  Yields:
-    A list of GraphDefs, populated when the context manager exits.
-  """
-  ctx = context()
-  ctx.enable_graph_collection()
-  try:
-    graphs = []
-    yield graphs
-    metadata = ctx.export_run_metadata()
-  finally:
-    ctx.disable_graph_collection()
-  for graph in metadata.function_graphs:
-    if optimized:
-      graphs.append(graph.post_optimization_graph)
-    else:
-      graphs.append(graph.pre_optimization_graph)
-
-
-def get_server_def():
-  return context().get_server_def()
-
-
-def set_server_def(server_def):
-  context().set_server_def(server_def)
-
-
-def set_server_def_retries(retries):
-  """Set the number of retries to use when calling SetServerDef.
-
-  In cases where many servers run in high-preemption environments, jobs could
-  be preempted during startup and initial connection via SetServerDef. Retries
-  allow for more robust connection in these environments.
-
-
-  Args:
-    retries: int specifying the number of connection retries before failing.
-      Retries follow an exponential backoff waiting period with min value 1ms,
-      max value 10s, and exponent 1.3.
-  """
-  context().set_server_def_retries(retries)
-
-
-def update_server_def(server_def):
-  context().update_server_def(server_def)
-
-
-def check_alive(worker_name):
-  return context().check_alive(worker_name)
-
-
-@tf_export("experimental.async_scope")
-@tf_contextlib.contextmanager
-def async_scope():
-  """Context manager for grouping async operations.
-
-  Ops/function calls inside the scope can return before finishing the actual
-  execution. When exiting the async scope, a synchronization barrier will be
-  automatically added to ensure the completion of all async op and function
-  execution, potentially raising exceptions if async execution results in
-  an error state.
-
-  Users may write the following code to asynchronously invoke `train_step_fn`
-  and log the `loss` metric for every `num_steps` steps in a training loop.
-  `train_step_fn` internally consumes data using `iterator.get_next()`, and may
-  throw OutOfRangeError when running out of data. In the case:
-
-  ```
-  try:
-    with tf.experimental.async_scope():
-      for _ in range(num_steps):
-        # Step function updates the metric `loss` internally
-        train_step_fn()
-  except tf.errors.OutOfRangeError:
-    tf.experimental.async_clear_error()
-  logging.info('loss = %s', loss.numpy())
-  ```
-
-  Yields:
-    Context manager for grouping async operations.
-  """
-  # TODO(haoyuzhang): replace env var once we have a config method to turn on
-  # and off async streaming RPC
-  remote_async_env_var = "TF_ENABLE_EAGER_CLIENT_STREAMING_ENQUEUE"
-  old_policy = os.environ.get(remote_async_env_var)
-  try:
-    os.environ[remote_async_env_var] = str(True)
-    yield
-    # Note: sync local and remote executors iff the async block does not raise
-    # an exception. Triggering sync after an exception may lead to derived
-    # runtime errors and unexpected exception types.
-    context().sync_executors()
-  finally:
-    if old_policy is None:
-      del os.environ[remote_async_env_var]
-    else:
-      os.environ[remote_async_env_var] = old_policy
-
-
-def async_wait():
-  """Sync all async operations and raise any errors during execution.
-
-  In async execution mode, an op/function call can return before finishing the
-  actual execution. Calling this method creates a synchronization barrier for
-  all async op and function execution. It only returns when all pending nodes
-  are finished, potentially raising exceptions if async execution results in
-  an error state. It is a no-op if the context is not initialized.
-  """
-  disable_async_executor_env_var = "TF_PS_DISABLE_ASYNC_EXECUTOR_GLOBALLY"
-  if os.environ.get(disable_async_executor_env_var) == str(True):
-    return
-  if context()._context_handle is not None:  # pylint: disable=protected-access
-    context().sync_executors()
-
-
-@tf_export("experimental.async_clear_error")
-def async_clear_error():
-  """Clear pending operations and error statuses in async execution.
-
-  In async execution mode, an error in op/function execution can lead to errors
-  in subsequent ops/functions that are scheduled but not yet executed. Calling
-  this method clears all pending operations and reset the async execution state.
-
-  Example:
-
-  ```
-  while True:
-    try:
-      # Step function updates the metric `loss` internally
-      train_step_fn()
-    except tf.errors.OutOfRangeError:
-      tf.experimental.async_clear_error()
-      break
-  logging.info('loss = %s', loss.numpy())
-  ```
-  """
-  context().clear_executor_errors()
-
-
-def add_c_function(c_func):
-  """Add a C API TF_Function to the context."""
-  context().add_c_function(c_func)
-
-
-def get_c_function(name):
-  """Get a C API TF_Function from the context."""
-  return context().get_c_function(name)
-
-
-def remove_function(name):
-  """Remove a function from the context."""
-  context().remove_function(name)
-
-
-def get_function_def(name):
-  return context().get_function_def(name)
-
-
-def is_custom_device(device_name):
-  """Calls TFE_IsCustomDevice.
-
-  Enables using C extensions specifying a custom device from Python. See the
-  experimental eager C API in tensorflow/c/eager/c_api_experimental.h for
-  details.
-
-  Args:
-    device_name: A string indicating the name to check whether it is a
-      registered custom device.
-
-  Returns:
-    A boolean.
-  """
-  return context().is_custom_device(device_name)
-
-
-def register_custom_device(device_capsule, device_name, device_info_capsule):
-  """Calls TFE_RegisterCustomDevice to register a custom device with Python.
-
-  Enables using C extensions specifying a custom device from Python. See the
-  experimental eager C API in tensorflow/c/eager/c_api_experimental.h for
-  details.
-
-  Note that custom devices are not currently supported inside `tf.function`s.
-
-  Args:
-    device_capsule: A PyCapsule with the name set to 'TFE_CustomDevice'
-      containing a pointer to a TFE_CustomDevice struct. The capsule retains
-      ownership of the memory.
-    device_name: A string indicating the name to register the custom device
-      under, e.g. '/job:localhost/replica:0/task:0/device:CUSTOM:0'. It may
-      subsequently be passed to `with tf.device(...):`.
-    device_info_capsule: A PyCapsule with the name set to
-      'TFE_CustomDevice_DeviceInfo' containing a pointer to a device-specific
-      struct with the initial state of the custom device (the void* device_info
-      argument to TFE_RegisterCustomDevice). This method takes ownership of the
-      memory and clears the capsule destructor.
-  """
-  context().register_custom_device(
-      device_capsule, device_name, device_info_capsule
-  )
-
-
-# Not every user creates a Context via context.context()
-# (for example, enable_eager_execution in python/framework/ops.py),
-# but they do all import this file.  Note that IS_IN_GRAPH_MODE and
-# in_graph_mode are both parameterless functions.
-def _tmp_in_graph_mode():
-  if context_safe() is None:
-    # Context not yet initialized. Assume graph mode following the
-    # default implementation in `is_in_graph_mode`.
-    return True
-  return not executing_eagerly()
-
-
-is_in_graph_mode.IS_IN_GRAPH_MODE = _tmp_in_graph_mode
+      context._set_context(new_context)
+      # Use a 2D tensor so that it is not cached.
+      tensor1 = constant_op.constant([[3.]])
+      # Produce a tensor as an operation output. This uses a different code path
+      # from tensors created from Python.
+      tensor2 = tensor1 * tensor1
+      context._set_context(original_context)
+    except:
+      context._set_context(original_context)
+      raise
+
+    # Deleting our context reference should not delete the underlying object.
+    del new_context
+    self.assertIsNot(weak_c(), None)
+
+    # Deleting the first tensor should not delete the context since there is
+    # another tensor.
+    del tensor1
+    self.assertIsNot(weak_c(), None)
+
+    # Deleting the last tensor should result in deleting its context.
+    del tensor2
+    self.assertIs(weak_c(), None)
+
+  def testSimpleGraphCollection(self):
+
+    @def_function.function
+    def f(x):
+      with ops.device('CPU:0'):
+        return x + constant_op.constant(1.)
+
+    with context.collect_graphs() as graphs:
+      with ops.device('CPU:0'):
+        x = constant_op.constant(1.)
+      f(x)
+
+    self.assertLen(graphs, 1)
+    graph, = graphs
+    self.assertIn('CPU:0', graph.node[1].device)
+
+  @test_util.disable_tfrt(
+      'b/171600738: tfrt does not support exporting post-optimization graph')
+  def testGraphCollectionAfterDevicePlacement(self):
+
+    @def_function.function
+    def f(x):
+      return x + constant_op.constant(1.)
+
+    with context.collect_graphs() as graphs:
+      with ops.device('CPU:0'):
+        f(constant_op.constant(1.))
+
+    self.assertLen(graphs, 1)
+    graph, = graphs
+    self.assertIn('CPU:0', graph.node[0].device)
+
+  def testGetFunctionDef(self):
+
+    @def_function.function
+    def f():
+      return constant_op.constant(1.)
+
+    concrete = f.get_concrete_function()
+    function_def = context.get_function_def(concrete.name)
+
+    self.assertIsNot(function_def, None)
+
+    found_const_node = False
+    for node_def in function_def.node_def:
+      if node_def.op == 'Const':
+        found_const_node = True
+        break
+    self.assertTrue(found_const_node)
+
+    with self.assertRaises(errors.NotFoundError):
+      _ = context.get_function_def('this_should_not_be_found')
+
+  @test_util.run_gpu_only
+  @test_util.disable_tfrt('b/169293680: TFE_GetTotalMemoryUsage is unsupported')
+  def testGetMemoryInfo(self):
+    array_ops.zeros([10])  # Allocate some memory on the GPU.
+    self.assertGreater(context.context().get_memory_info('GPU:0')['current'], 0)
+
+  @test_util.disable_tfrt('b/169293680: TFE_GetTotalMemoryUsage is unsupported')
+  def testGetMemoryInfoCPU(self):
+    if test_util.IsMklEnabled():
+      # TODO(gzmkl) work with Google team to address design issue in allocator.h
+      self.skipTest('MklCPUAllocator does not throw exception. So skip test.')
+
+    with self.assertRaisesRegex(ValueError, 'Allocator stats not available'):
+      context.context().get_memory_info('CPU:0')
+
+  @test_util.disable_tfrt('b/169293680: TFE_GetTotalMemoryUsage is unsupported')
+  def testGetMemoryInfoUnknownDevice(self):
+    with self.assertRaisesRegex(ValueError, 'No matching devices found'):
+      context.context().get_memory_info('unknown_device:0')
+
+  @test_util.disable_tfrt('b/169293680: TFE_GetTotalMemoryUsage is unsupported')
+  def testGetMemoryInfoUnparsableDevice(self):
+    with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
+      context.context().get_memory_info('GPU')
+    with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
+      context.context().get_memory_info('GPU:')
+    with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
+      context.context().get_memory_info('GPU:CPU')
+
+  def testListFunctionNames(self):
+
+    @def_function.function
+    def f():
+      return constant_op.constant(1.)
+
+    concrete = f.get_concrete_function()
+    self.assertIn(concrete.name.decode(),
+                  context.context().list_function_names())
+
+  def testCoordinationServiceMethodsInitializeNonGlobalContext(self):
+    # Each method must initialize its own (non-global) context before calling
+    # the C API; without a coordination service it then fails cleanly.
+    with self.assertRaisesRegex(errors.FailedPreconditionError,
+                                'Coordination service agent is not enabled'):
+      context.Context().set_config_key_value('key', 'value')
+    with self.assertRaisesRegex(errors.FailedPreconditionError,
+                                'Coordination service is not enabled'):
+      context.Context().get_config_key_value('key', timeout_in_ms=1)
+    with self.assertRaisesRegex(errors.FailedPreconditionError,
+                                'Coordination service is not enabled'):
+      context.Context().delete_config_key_value('key')
+    with self.assertRaisesRegex(errors.FailedPreconditionError,
+                                'Coordination service is not enabled'):
+      context.Context().wait_at_barrier('barrier', 1)
+
+  def testSetLogicalDeviceAfterContextInitialization(self):
+    ctx = context.Context()
+    ctx.set_logical_cpu_devices(4)
+    self.assertIs(len(ctx.list_logical_devices('CPU')), 4)
+
+    # Cannot set logical device twice.
+    with self.assertRaisesRegex(RuntimeError, 'Virtual CPUs already set'):
+      ctx.set_logical_cpu_devices(8)
+
+  def testSetLogicalLocalCpuDevice(self):
+    ctx = context.Context()
+
+    # Manually add a remote CPU device into logical device list.
+    ctx._logical_devices = []  # pylint: disable=protected-access
+    dev = context.LogicalDevice(name='/job:worker/replica:0/task:1',
+                                device_type='CPU')
+    ctx._logical_devices.append(dev)  # pylint: disable=protected-access
+    self.assertIs(len(ctx.list_logical_devices('CPU')), 1)
+
+    # This would pass the check since the previously added device is not local.
+    ctx.set_logical_cpu_devices(4)
+    # Logical device list would be overwritten after initialization.
+    self.assertIs(len(ctx.list_logical_devices('CPU')), 4)
+
+  @parameterized.named_parameters([(f'_{stage}', stage) for stage in [
+      'hlo', 'hlo_serialized', 'optimized_hlo', 'optimized_hlo_serialized',
+      'optimized_hlo_proto_serialized', 'optimized_hlo_dot'
+  ]])
+  def testGetCompilerIr(self, stage):
+
+    @def_function.function(jit_compile=True)
+    def test_func(x):
+      return 2 * x
+
+    a = array_ops.ones((1000, 1000))  # 4 * 1000 * 1000 in bytes
+    result = test_func.experimental_get_compiler_ir(a)(stage=stage)
+    self.assertNotEmpty(result)
+    if stage == 'optimized_hlo_proto_serialized':
+      hlo_proto = hlo_pb2.HloProto.FromString(result)
+      allocations = hlo_proto.buffer_assignment.buffer_allocations
+      buffer_size = sum(
+          getattr(allocation, 'size') for allocation in allocations)
+      # The sizes of input and output are both 4 * 1000 * 1000 in bytes.
+      self.assertGreaterEqual(buffer_size, 2 * 4 * 1000 * 1000)
+      self.assertLess(buffer_size, 4 * 4 * 1000 * 1000)
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  test.main()

From f3ded11fdc23c19a1a0487953a7cfe756e539e13 Mon Sep 17 00:00:00 2001
From: Mina Ibrahim <136262314+MinaIbrahim10@users.noreply.github.com>
Date: Mon, 5 Oct 2026 16:49:07 +0300
Subject: [PATCH 0195/1000] Regenerate tf.data service pybind stubs

---
 .../python/data/experimental/service/_pywrap_server_lib.pyi   | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/data/experimental/service/_pywrap_server_lib.pyi b/tensorflow/python/data/experimental/service/_pywrap_server_lib.pyi
index d3443e95d52..26386de81c9 100644
--- a/tensorflow/python/data/experimental/service/_pywrap_server_lib.pyi
+++ b/tensorflow/python/data/experimental/service/_pywrap_server_lib.pyi
@@ -19,7 +19,7 @@ class DispatchGrpcDataServer:
     def join(self) -> None: ...
     def num_workers(self) -> int: ...
     def snapshot_streams(self, *args, **kwargs): ...
-    def start(self) -> Status: ...
+    def start(self) -> None: ...
     def stop(self) -> None: ...
 
 class SnapshotStreamInfoWrapper:
@@ -44,7 +44,7 @@ class WorkerGrpcDataServer:
     def join(self) -> None: ...
     def num_tasks(self) -> int: ...
     def snapshot_task_progresses(self, *args, **kwargs): ...
-    def start(self) -> Status: ...
+    def start(self) -> None: ...
     def stop(self) -> None: ...
 
 def TF_DATA_GetDataServiceMetadataByID(*args, **kwargs): ...

From 62b0baf1542b68fe5d47ceb9e598e2a3ed846941 Mon Sep 17 00:00:00 2001
From: Noxtimo <139155109+Noxtimo@users.noreply.github.com>
Date: Mon, 5 Oct 2026 15:40:16 +0100
Subject: [PATCH 0196/1000] Address review feedback: extract _signbit helper
 and enhance test coverage

---
 .../python/ops/numpy_ops/np_math_ops.py       | 32 ++++++++-----------
 .../python/ops/numpy_ops/np_math_ops_test.py  | 21 +++++++++---
 2 files changed, 30 insertions(+), 23 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py
index 5a978cb903f..484159a4fd8 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@@ -855,20 +855,21 @@ def reciprocal(x):
   return _scalar(math_ops.reciprocal, x)
 
 
+def _signbit(x):
+  if x.dtype == dtypes.bool:
+    return array_ops.fill(array_ops.shape(x), False)
+  if x.dtype in _SIGN_BITCAST_DTYPES:
+    # Check the IEEE-754 sign bit instead of comparing with zero, which
+    # cannot tell -0.0 from +0.0 or a negative NaN from a positive one.
+    bits = array_ops.bitcast(x, _SIGN_BITCAST_DTYPES[x.dtype])
+    return bits < 0
+  return x < 0
+
+
 @tf_export.tf_export('experimental.numpy.signbit', v1=[])
 @np_utils.np_doc('signbit')
 def signbit(x):
-  def f(x):
-    if x.dtype == dtypes.bool:
-      return array_ops.fill(array_ops.shape(x), False)
-    if x.dtype in _SIGN_BITCAST_DTYPES:
-      # Check the IEEE-754 sign bit instead of comparing with zero, which
-      # cannot tell -0.0 from +0.0 or a negative NaN from a positive one.
-      bits = array_ops.bitcast(x, _SIGN_BITCAST_DTYPES[x.dtype])
-      return math_ops.less(bits, 0)
-    return x < 0
-
-  return _scalar(f, x)
+  return _scalar(_signbit, x)
 
 
 @tf_export.tf_export('experimental.numpy.sin', v1=[])
@@ -994,14 +995,7 @@ def cbrt(x):
   def f(x):
     # __pow__ can't handle negative base, so we use `abs` here.
     rt = math_ops.abs(x) ** (1.0 / 3)
-    if x.dtype in _SIGN_BITCAST_DTYPES:
-      # Check the IEEE-754 sign bit instead of comparing with zero, which
-      # cannot tell -0.0 from +0.0.
-      bits = array_ops.bitcast(x, _SIGN_BITCAST_DTYPES[x.dtype])
-      is_negative = math_ops.less(bits, 0)
-    else:
-      is_negative = x < 0
-    return array_ops.where_v2(is_negative, -rt, rt)
+    return array_ops.where_v2(_signbit(x), -rt, rt)
 
   return _scalar(f, x, True)
 
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 dab3005a55c..5800408ea19 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
@@ -866,12 +866,12 @@ class MathTest(test.TestCase, parameterized.TestCase):
       x = constant_op.constant(x_np, dtype=dtype)
       out = np_math_ops.cbrt(x)
       self.assertAllClose(out, np.cbrt(x_np))
-      self.assertAllEqual(np.signbit(out.numpy()), np.signbit(np.cbrt(x_np)))
+      self.assertAllEqual(np_math_ops.signbit(out), np.signbit(np.cbrt(x_np)))
 
     x_neg_zero = constant_op.constant(-0.0, dtype=dtypes.float64)
-    self.assertEqual(
-        math_ops.reciprocal(np_math_ops.cbrt(x_neg_zero)).numpy(),
-        -float('inf'),
+    self.assertAllEqual(
+        math_ops.reciprocal(np_math_ops.cbrt(x_neg_zero)),
+        -np.inf,
     )
 
     x_bf16 = constant_op.constant([-0.0, 0.0, -8.0, 8.0], dtype=dtypes.bfloat16)
@@ -880,6 +880,19 @@ class MathTest(test.TestCase, parameterized.TestCase):
         np_math_ops.signbit(out_bf16), [True, False, True, False]
     )
 
+    for dtype in (dtypes.float16, dtypes.float32, dtypes.float64, dtypes.bfloat16):
+      # Verify N=0 (empty tensor) fallback.
+      x_empty = constant_op.constant([], dtype=dtype)
+      out_empty = np_math_ops.cbrt(x_empty)
+      self.assertEqual(out_empty.shape.num_elements(), 0)
+
+      # Verify NaN, +Inf, -Inf
+      x_edge = constant_op.constant([-float('inf'), float('inf'), float('nan')], dtype=dtype)
+      out_edge = np_math_ops.cbrt(x_edge)
+      # inf ** (1/3) == inf; NaN ** (1/3) == NaN
+      # Negative sign bit for -inf is preserved.
+      self.assertAllEqual(np_math_ops.signbit(out_edge[0:2]), [True, False])
+
   def testSinc(self):
     for dtype in (dtypes.float32, dtypes.float64):
       x = constant_op.constant([0.0, -0.0, 0.5, 1.0, -1.5], dtype=dtype)

From 9ed3340aa2ca620fb53396551e483381c7cb3909 Mon Sep 17 00:00:00 2001
From: PrabinDevkota <prabindev789@gmail.com>
Date: Mon, 5 Oct 2026 21:04:52 +0530
Subject: [PATCH 0197/1000] Use errors::InvalidArgument for max-rank checks and
 cover public slicing.

---
 tensorflow/core/kernels/pack_op.cc                   |  5 +++--
 tensorflow/core/kernels/shape_ops.h                  |  4 ++--
 tensorflow/core/ops/array_ops.cc                     |  8 ++++----
 tensorflow/core/util/BUILD                           |  1 +
 tensorflow/core/util/strided_slice_op.cc             |  6 +++---
 .../python/kernel_tests/array_ops/array_ops_test.py  | 12 ++++--------
 6 files changed, 17 insertions(+), 19 deletions(-)

diff --git a/tensorflow/core/kernels/pack_op.cc b/tensorflow/core/kernels/pack_op.cc
index 086df0c1be4..3dd0be89f9c 100644
--- a/tensorflow/core/kernels/pack_op.cc
+++ b/tensorflow/core/kernels/pack_op.cc
@@ -27,6 +27,7 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_types.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/kernels/concat_lib.h"
+#include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/types.h"
 
@@ -69,10 +70,10 @@ class PackOp : public OpKernel {
 
     OP_REQUIRES(
         c, first_input.dims() < TensorShape::MaxDimensions(),
-        absl::InvalidArgumentError(absl::StrCat(
+        errors::InvalidArgument(
             "Cannot pack tensors of rank ", first_input.dims(),
             ": packing would exceed the maximum supported rank of ",
-            TensorShape::MaxDimensions(), ".")));
+            TensorShape::MaxDimensions(), "."));
 
     TensorShape output_shape(first_input.shape());
     OP_REQUIRES_OK(c, output_shape.InsertDimWithStatus(axis, num));
diff --git a/tensorflow/core/kernels/shape_ops.h b/tensorflow/core/kernels/shape_ops.h
index 3209e469b6e..2ab6c2df5cd 100644
--- a/tensorflow/core/kernels/shape_ops.h
+++ b/tensorflow/core/kernels/shape_ops.h
@@ -171,10 +171,10 @@ class ExpandDimsOp : public OpKernel {
 
     OP_REQUIRES(
         ctx, input_dims < TensorShape::MaxDimensions(),
-        absl::InvalidArgumentError(absl::StrCat(
+        errors::InvalidArgument(
             "Cannot expand a tensor of rank ", input_dims,
             ": expanding would exceed the maximum supported rank of ",
-            TensorShape::MaxDimensions(), ".")));
+            TensorShape::MaxDimensions(), "."));
 
     // Compute new shape with an additional dimension.
     absl::InlinedVector<int64_t, 8> output_shape_vec(input_dims + 1);
diff --git a/tensorflow/core/ops/array_ops.cc b/tensorflow/core/ops/array_ops.cc
index 9a34f9aaf00..3e0272f77c9 100644
--- a/tensorflow/core/ops/array_ops.cc
+++ b/tensorflow/core/ops/array_ops.cc
@@ -380,10 +380,10 @@ REGISTER_OP("Pack")
       TF_RETURN_IF_ERROR(GetAxisForPackAndUnpack(c, rank + 1, &axis));
 
       if (rank >= TensorShape::MaxDimensions()) {
-        return absl::InvalidArgumentError(absl::StrCat(
+        return errors::InvalidArgument(
             "Cannot pack tensors of rank ", rank,
             ": packing would exceed the maximum supported rank of ",
-            TensorShape::MaxDimensions(), "."));
+            TensorShape::MaxDimensions(), ".");
       }
 
       // Copy all dimensions over, inserting a dimension of value #inputs
@@ -2160,10 +2160,10 @@ REGISTER_OP("ExpandDims")
       }
 
       if (rank >= TensorShape::MaxDimensions()) {
-        return absl::InvalidArgumentError(absl::StrCat(
+        return errors::InvalidArgument(
             "Cannot expand a tensor of rank ", rank,
             ": expanding would exceed the maximum supported rank of ",
-            TensorShape::MaxDimensions(), "."));
+            TensorShape::MaxDimensions(), ".");
       }
 
       ShapeHandle end;
diff --git a/tensorflow/core/util/BUILD b/tensorflow/core/util/BUILD
index 484cacb0936..b850b43c18d 100644
--- a/tensorflow/core/util/BUILD
+++ b/tensorflow/core/util/BUILD
@@ -926,6 +926,7 @@ tf_cc_tests(
         "@com_google_absl//absl/base",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@eigen_archive//:eigen3",
     ],
diff --git a/tensorflow/core/util/strided_slice_op.cc b/tensorflow/core/util/strided_slice_op.cc
index bb33ab43786..fff42cb430b 100644
--- a/tensorflow/core/util/strided_slice_op.cc
+++ b/tensorflow/core/util/strided_slice_op.cc
@@ -20,8 +20,8 @@ limitations under the License.
 #include <iterator>
 #include <utility>
 
-#include "absl/strings/str_cat.h"
 #include "tensorflow/core/framework/bounds_check.h"
+#include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/status.h"
 
 namespace tensorflow {
@@ -417,10 +417,10 @@ absl::Status ValidateStridedSliceOp(
     }
   }
   if (output_rank > TensorShape::MaxDimensions()) {
-    return absl::InvalidArgumentError(absl::StrCat(
+    return errors::InvalidArgument(
         "StridedSlice would produce a tensor of rank ", output_rank,
         ", which exceeds the maximum supported rank of ",
-        TensorShape::MaxDimensions(), "."));
+        TensorShape::MaxDimensions(), ".");
   }
 
   final_shape->Clear();
diff --git a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
index 95cfbeba4b6..f2e01af0bfd 100644
--- a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
@@ -799,14 +799,10 @@ class StridedSliceTest(test_util.TensorFlowTestCase):
         constant_op.constant(0.0, dtype=dtypes.float32))
     with self.assertRaisesRegex(
         (errors.InvalidArgumentError, ValueError), "maximum supported rank"):
-      self.evaluate(
-          gen_array_ops.strided_slice(
-              value,
-              begin=array_ops.zeros([31], dtypes.int32),
-              end=array_ops.zeros([31], dtypes.int32),
-              strides=array_ops.ones([31], dtypes.int32),
-              new_axis_mask=(1 << 30) - 1,
-              ellipsis_mask=1 << 30))
+      # Public slice syntax, not the generated op wrapper. evaluate() covers
+      # graph mode, where shape inference may defer the error until execution.
+      sliced = value[tuple([array_ops.newaxis] * 30 + [Ellipsis])]
+      self.evaluate(sliced)
 
   @test_util.assert_no_new_pyobjects_executing_eagerly()
   @test_util.assert_no_garbage_created

From 197e767470f778d874b23307c4323001cc74789d Mon Sep 17 00:00:00 2001
From: Noxtimo <139155109+Noxtimo@users.noreply.github.com>
Date: Mon, 5 Oct 2026 16:54:48 +0100
Subject: [PATCH 0198/1000] style: wrap long lines in np_math_ops_test.py for
 pylint

---
 tensorflow/python/ops/numpy_ops/np_math_ops_test.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

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 5800408ea19..237378e23ee 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
@@ -880,14 +880,21 @@ class MathTest(test.TestCase, parameterized.TestCase):
         np_math_ops.signbit(out_bf16), [True, False, True, False]
     )
 
-    for dtype in (dtypes.float16, dtypes.float32, dtypes.float64, dtypes.bfloat16):
+    for dtype in (
+        dtypes.float16,
+        dtypes.float32,
+        dtypes.float64,
+        dtypes.bfloat16,
+    ):
       # Verify N=0 (empty tensor) fallback.
       x_empty = constant_op.constant([], dtype=dtype)
       out_empty = np_math_ops.cbrt(x_empty)
       self.assertEqual(out_empty.shape.num_elements(), 0)
 
       # Verify NaN, +Inf, -Inf
-      x_edge = constant_op.constant([-float('inf'), float('inf'), float('nan')], dtype=dtype)
+      x_edge = constant_op.constant(
+          [-float('inf'), float('inf'), float('nan')], dtype=dtype
+      )
       out_edge = np_math_ops.cbrt(x_edge)
       # inf ** (1/3) == inf; NaN ** (1/3) == NaN
       # Negative sign bit for -inf is preserved.

From c9a0e947dc3d30e7a07b154333f1a3ce0264542b Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Tue, 6 Oct 2026 02:54:27 +0530
Subject: [PATCH 0199/1000] reject negative uncompressed_bytes and reuse
 validated shapes in UncompressElement

---
 tensorflow/core/data/compression_utils.cc     | 31 +++++++----
 .../core/data/compression_utils_test.cc       | 53 +++++++++++++++++--
 2 files changed, 70 insertions(+), 14 deletions(-)

diff --git a/tensorflow/core/data/compression_utils.cc b/tensorflow/core/data/compression_utils.cc
index d5cd4feb67e..1eb286e8c3b 100644
--- a/tensorflow/core/data/compression_utils.cc
+++ b/tensorflow/core/data/compression_utils.cc
@@ -159,8 +159,12 @@ absl::Status UncompressElement(const CompressedElement& compressed,
   size_t num_string_tensors = 0;
   size_t num_string_tensor_strings = 0;
   size_t total_nonmemcpyable_size = 0;
+  // Shapes are validated once here and reused by the second pass.
+  std::vector<TensorShape> shapes;
+  shapes.reserve(num_components);
   for (const auto& metadata : compressed.component_metadata()) {
-    TensorShape shape;
+    shapes.emplace_back();
+    TensorShape& shape = shapes.back();
     TF_RETURN_IF_ERROR(
         TensorShape::BuildTensorShape(metadata.tensor_shape(), &shape));
     if (metadata.dtype() == DT_STRING) {
@@ -175,6 +179,10 @@ absl::Status UncompressElement(const CompressedElement& compressed,
         return absl::InvalidArgumentError(
             "Missing uncompressed_bytes metadata for non-memcpyable tensor");
       }
+      if (metadata.uncompressed_bytes(0) < 0) {
+        return absl::InvalidArgumentError(
+            "uncompressed_bytes metadata cannot be negative");
+      }
       total_nonmemcpyable_size += metadata.uncompressed_bytes(0);
     } else {
       int64_t num_elements = shape.num_elements();
@@ -195,11 +203,11 @@ absl::Status UncompressElement(const CompressedElement& compressed,
   tstring nonmemcpyable;
   nonmemcpyable.resize_uninitialized(total_nonmemcpyable_size);
   char* nonmemcpyable_pos = nonmemcpyable.mdata();
-  for (const auto& metadata : compressed.component_metadata()) {
+  for (int i = 0; i < num_components; ++i) {
+    const CompressedComponentMetadata& metadata =
+        compressed.component_metadata(i);
+    const TensorShape& shape = shapes[i];
     if (DataTypeCanUseMemcpy(metadata.dtype())) {
-      TensorShape shape;
-      TF_RETURN_IF_ERROR(
-          TensorShape::BuildTensorShape(metadata.tensor_shape(), &shape));
       int64_t num_elements = shape.num_elements();
       out->emplace_back(metadata.dtype(), shape);
       TensorBuffer* buffer = DMAHelper::buffer(&out->back());
@@ -212,9 +220,6 @@ absl::Status UncompressElement(const CompressedElement& compressed,
         iov.Add(buffer->data(), metadata.uncompressed_bytes(0));
       }
     } else if (metadata.dtype() == DT_STRING) {
-      TensorShape shape;
-      TF_RETURN_IF_ERROR(
-          TensorShape::BuildTensorShape(metadata.tensor_shape(), &shape));
       out->emplace_back(metadata.dtype(), shape);
       const auto& flats = out->back().unaligned_flat<tstring>();
       if (metadata.uncompressed_bytes_size() > flats.size()) {
@@ -223,9 +228,13 @@ absl::Status UncompressElement(const CompressedElement& compressed,
             ") exceeds allocated string tensor elements count (", flats.size(),
             ")"));
       }
-      for (int i = 0; i < metadata.uncompressed_bytes_size(); ++i) {
-        flats.data()[i].resize(metadata.uncompressed_bytes(i));
-        iov.Add(flats.data()[i].mdata(), metadata.uncompressed_bytes(i));
+      for (int j = 0; j < metadata.uncompressed_bytes_size(); ++j) {
+        if (metadata.uncompressed_bytes(j) < 0) {
+          return absl::InvalidArgumentError(
+              "uncompressed_bytes metadata cannot be negative");
+        }
+        flats.data()[j].resize(metadata.uncompressed_bytes(j));
+        iov.Add(flats.data()[j].mdata(), metadata.uncompressed_bytes(j));
       }
     } else {
       out->emplace_back();
diff --git a/tensorflow/core/data/compression_utils_test.cc b/tensorflow/core/data/compression_utils_test.cc
index f0dcea856b1..5e5bcad6a17 100644
--- a/tensorflow/core/data/compression_utils_test.cc
+++ b/tensorflow/core/data/compression_utils_test.cc
@@ -23,7 +23,10 @@ limitations under the License.
 #include "xla/tsl/protobuf/error_codes.pb.h"
 #include "tensorflow/core/data/dataset_test_base.h"
 #include "tensorflow/core/framework/dataset.pb.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/protobuf/error_codes.pb.h"
 
@@ -91,8 +94,10 @@ TEST(CompressionUtilsTest, MalformedTensorShape) {
   metadata->mutable_tensor_shape()->add_dim()->set_size(-1);
 
   std::vector<Tensor> element;
-  EXPECT_THAT(UncompressElement(compressed, &element),
-              absl_testing::StatusIs(error::INVALID_ARGUMENT));
+  EXPECT_THAT(
+      UncompressElement(compressed, &element),
+      absl_testing::StatusIs(error::INVALID_ARGUMENT,
+                             HasSubstr("Expected a non-negative size")));
 }
 
 TEST(CompressionUtilsTest, MalformedTensorShapeString) {
@@ -108,9 +113,51 @@ TEST(CompressionUtilsTest, MalformedTensorShapeString) {
   metadata->set_dtype(DT_STRING);
   metadata->mutable_tensor_shape()->add_dim()->set_size(-1);
 
+  std::vector<Tensor> element;
+  EXPECT_THAT(
+      UncompressElement(compressed, &element),
+      absl_testing::StatusIs(error::INVALID_ARGUMENT,
+                             HasSubstr("Expected a non-negative size")));
+}
+
+TEST(CompressionUtilsTest, NegativeUncompressedBytes) {
+  // `uncompressed_bytes` is a signed proto field, so a negative count on a
+  // non-memcpyable component must be rejected before it is added to the
+  // unsigned scratch buffer size.
+  std::vector<Tensor> empty_element;
+  CompressedElement compressed;
+  TF_ASSERT_OK(CompressElement(empty_element, &compressed));
+
+  CompressedComponentMetadata* metadata =
+      compressed.mutable_component_metadata()->Add();
+  metadata->set_dtype(DT_VARIANT);
+  TensorShape empty_shape({0});
+  empty_shape.AsProto(metadata->mutable_tensor_shape());
+  metadata->add_uncompressed_bytes(-1);
+
   std::vector<Tensor> element;
   EXPECT_THAT(UncompressElement(compressed, &element),
-              absl_testing::StatusIs(error::INVALID_ARGUMENT));
+              absl_testing::StatusIs(error::INVALID_ARGUMENT,
+                                     HasSubstr("cannot be negative")));
+}
+
+TEST(CompressionUtilsTest, NegativeUncompressedBytesString) {
+  // The same holds for the per-string byte counts of a `DT_STRING` component.
+  std::vector<Tensor> empty_element;
+  CompressedElement compressed;
+  TF_ASSERT_OK(CompressElement(empty_element, &compressed));
+
+  CompressedComponentMetadata* metadata =
+      compressed.mutable_component_metadata()->Add();
+  metadata->set_dtype(DT_STRING);
+  TensorShape shape({1});
+  shape.AsProto(metadata->mutable_tensor_shape());
+  metadata->add_uncompressed_bytes(-1);
+
+  std::vector<Tensor> element;
+  EXPECT_THAT(UncompressElement(compressed, &element),
+              absl_testing::StatusIs(error::INVALID_ARGUMENT,
+                                     HasSubstr("cannot be negative")));
 }
 
 std::vector<std::vector<Tensor>> TestCases() {

From 1d5a23db8a254e01067a4fbe428b7b544604b553 Mon Sep 17 00:00:00 2001
From: Som Tripathi <somtri@iastate.edu>
Date: Mon, 5 Oct 2026 17:30:09 -0500
Subject: [PATCH 0200/1000] Skip the oneDNN MaxPool rewrite for unsupported
 data formats

NonDepthBatchWisePoolRewrite decides whether MaxPool and MaxPool3D are
rewritten to their oneDNN ops. It parsed data_format but never checked
it. MaxPool also accepts NCHW_VECT_C, while the oneDNN op allows only
NHWC and NCHW, so building the rewritten node failed attr validation
and the TF_CHECK_OK in RewriteNode aborted the process.

Return false unless the parsed format is FORMAT_NHWC or FORMAT_NCHW.
FormatFromString maps NDHWC and NCDHW to those values, so MaxPool3D is
still rewritten as before. An NCHW_VECT_C node now stays on the default
kernel, which returns InvalidArgumentError, the same error seen with
oneDNN disabled.

Add a graph and eager regression test that runs MaxPool with
NCHW_VECT_C on CPU and expects the default kernel's error.

Fixes #94115.
---
 tensorflow/core/common_runtime/mkl_layout_pass.cc  |  4 ++++
 .../python/kernel_tests/nn_ops/pooling_ops_test.py | 14 ++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/tensorflow/core/common_runtime/mkl_layout_pass.cc b/tensorflow/core/common_runtime/mkl_layout_pass.cc
index 928baa3ef6c..61367b3e7b3 100644
--- a/tensorflow/core/common_runtime/mkl_layout_pass.cc
+++ b/tensorflow/core/common_runtime/mkl_layout_pass.cc
@@ -1599,6 +1599,10 @@ class MklLayoutRewritePass : public GraphOptimizationPass {
     TF_CHECK_OK(GetNodeAttr(n->def(), "data_format", &data_format_str));
     bool result = FormatFromString(data_format_str, &data_format);
     DCHECK(result);
+    // The oneDNN MaxPool ops accept only NHWC and NCHW (NDHWC and NCDHW in 3D).
+    if (!result || (data_format != FORMAT_NHWC && data_format != FORMAT_NCHW)) {
+      return false;
+    }
 
     // Condition that specifies non-batch-wise and non-depth-wise pooling.
     if (GetTensorDim(ksize, data_format, 'N') == 1 &&
diff --git a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
index 78a7a851619..fcae1e9d865 100644
--- a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
@@ -818,6 +818,20 @@ class PoolingTest(test.TestCase, parameterized.TestCase):
             padding="SAME")
         self.evaluate(t)
 
+  @test_util.run_in_graph_and_eager_modes
+  def testMaxPoolNchwVectCOnCpuRaisesError(self):
+    # With oneDNN on, this MaxPool used to be rewritten to the oneDNN op, which
+    # rejects NCHW_VECT_C, and the process aborted.
+    x = array_ops.zeros([1, 2, 4, 4, 4], dtypes.float32)
+    with ops.device("/cpu:0"):
+      with self.assertRaisesRegex(
+          (errors_impl.InvalidArgumentError, ValueError),
+          "only supports NHWC on device type CPU"):
+        self.evaluate(
+            nn_ops.max_pool2d(
+                x, ksize=2, strides=2, padding="VALID",
+                data_format="NCHW_VECT_C"))
+
   # Tests for DepthwiseMaxPooling on CPU only.
   @parameterized.parameters(
       GetTestConfigsDicts(

From 4eacf600765e27bc3d03fdc3b4911e439cf10fdc Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Tue, 6 Oct 2026 04:11:47 +0530
Subject: [PATCH 0201/1000] Bound num_invalid_tensors by metadata size and
 enforce max_num_elements in TensorList::Decode

---
 tensorflow/core/kernels/tensor_list.cc      | 11 ++++++++
 tensorflow/core/kernels/tensor_list_test.cc | 31 +++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/tensorflow/core/kernels/tensor_list.cc b/tensorflow/core/kernels/tensor_list.cc
index db51b279776..f67dfb1f401 100644
--- a/tensorflow/core/kernels/tensor_list.cc
+++ b/tensorflow/core/kernels/tensor_list.cc
@@ -64,6 +64,10 @@ bool TensorList::Decode(const VariantTensorData& data) {
   if (!core::GetVarint64(&iter, &scratch)) return false;
   if (scratch > std::numeric_limits<size_t>::max()) return false;
   const size_t num_invalid_tensors = static_cast<size_t>(scratch);
+  // Each invalid index takes at least one byte of the remaining metadata, so a
+  // larger count is malformed. Reject it here, before storage is reserved for
+  // that many elements below.
+  if (num_invalid_tensors > iter.size()) return false;
 
   if (num_invalid_tensors >
       std::numeric_limits<size_t>::max() - data.tensors().size()) {
@@ -121,6 +125,13 @@ bool TensorList::Decode(const VariantTensorData& data) {
     if (scratch > std::numeric_limits<int>::max()) return false;
     decoded_max_num_elements = static_cast<int>(scratch);
   }
+  // When max_num_elements is set it caps the list size (TensorListPushBack and
+  // TensorListSetItem enforce it on insertion), so reject a list that already
+  // exceeds its own cap.
+  if (decoded_max_num_elements != -1 &&
+      total_num_tensors > static_cast<size_t>(decoded_max_num_elements)) {
+    return false;
+  }
 
   TensorShapeProto element_shape_proto;
   if (!element_shape_proto.ParseFromString(iter)) return false;
diff --git a/tensorflow/core/kernels/tensor_list_test.cc b/tensorflow/core/kernels/tensor_list_test.cc
index 39eea10438a..ab860905d25 100644
--- a/tensorflow/core/kernels/tensor_list_test.cc
+++ b/tensorflow/core/kernels/tensor_list_test.cc
@@ -235,6 +235,37 @@ TEST(TensorListTest, DecodeRejectsEmptyTensorMismatchedShape) {
   EXPECT_FALSE(TensorList().Decode(data));
 }
 
+TEST(TensorListTest, DecodeRejectsNumInvalidTensorsLargerThanMetadata) {
+  // Each invalid index takes at least one byte of metadata, so a count far
+  // larger than the metadata itself is malformed. Decode must reject it up
+  // front instead of reserving storage for that many elements.
+  TensorShapeProto shape;
+  std::string metadata;
+  core::PutVarint64(&metadata, uint64_t{1} << 40);  // num_invalid_tensors
+  core::PutVarint64(&metadata, static_cast<uint64_t>(DT_FLOAT));
+  core::PutVarint64(&metadata, std::numeric_limits<uint64_t>::max());
+  metadata.append(shape.SerializeAsString());
+  VariantTensorData data;
+  data.set_metadata(metadata);
+  EXPECT_FALSE(TensorList().Decode(data));
+}
+
+TEST(TensorListTest, DecodeRejectsMoreElementsThanMaxNumElements) {
+  // A list filled exactly to max_num_elements is valid; one more element
+  // exceeds the list's own capacity and must be rejected.
+  TensorShapeProto shape;
+  VariantTensorData data;
+  data.set_metadata(TensorListMetadata(
+      /*invalid_indices=*/{}, static_cast<uint64_t>(DT_INT32),
+      /*max_num_elements=*/2, shape.SerializeAsString()));
+  *data.add_tensors() = ScalarInt32Tensor(1);
+  *data.add_tensors() = ScalarInt32Tensor(2);
+  EXPECT_TRUE(TensorList().Decode(data));
+
+  *data.add_tensors() = ScalarInt32Tensor(3);
+  EXPECT_FALSE(TensorList().Decode(data));
+}
+
 TEST(TensorListTest, DecodeRejectsOutOfRangeMaxNumElements) {
   TensorShapeProto shape;
   VariantTensorData data;

From 7ca4c80d4e2b081d70bad7b4ffd8ce8423d98fce Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Tue, 6 Oct 2026 04:51:02 +0530
Subject: [PATCH 0202/1000] Fix division by zero in throughput calculation

---
 tensorflow/python/platform/benchmark.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/platform/benchmark.py b/tensorflow/python/platform/benchmark.py
index 36b9201b36d..22d6504e0cf 100644
--- a/tensorflow/python/platform/benchmark.py
+++ b/tensorflow/python/platform/benchmark.py
@@ -398,7 +398,7 @@ class TensorFlowBenchmark(Benchmark):
         "wall_time": median_delta,
         "extras": extras,
         "name": name,
-        "throughput": mbs / median_delta
+        "throughput": mbs / median_delta if median_delta > 0 else 0.0
     }
     self.report_benchmark(**benchmark_values)
 
@@ -484,7 +484,7 @@ def benchmarks_main(true_main, argv=None):
     # Remove --benchmark_filter arg from sys.argv
     argv.remove(found_arg[0])
 
-    regex = found_arg[0].split("=")[1]
+    regex = found_arg[0].split("=", 1)[1]
     app.run(lambda _: _run_benchmarks(regex), argv=argv)
   else:
     true_main()

From b234e9a691f20e0c332d1f7a637c6caabab8d58e Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Tue, 6 Oct 2026 05:02:19 +0530
Subject: [PATCH 0203/1000] Update windows_lib_diagnostics.py

---
 .../platform/windows_lib_diagnostics.py       | 22 ++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/tensorflow/python/platform/windows_lib_diagnostics.py b/tensorflow/python/platform/windows_lib_diagnostics.py
index dca9cc8f7b9..cb5e211e5a3 100644
--- a/tensorflow/python/platform/windows_lib_diagnostics.py
+++ b/tensorflow/python/platform/windows_lib_diagnostics.py
@@ -27,6 +27,18 @@ ERROR_BAD_EXE_FORMAT = 0xC1
 ERROR_DLL_INIT_FAILED = 0x45A
 
 
+def _safe_print(text, end="\n"):
+  """Prints `text`, escaping characters the console encoding cannot represent.
+
+  This runs while an import failure is being handled, so it must never raise:
+  paths (e.g. under a non-ASCII user directory) may not be encodable on a
+  console that uses a legacy code page.
+  """
+  encoding = getattr(sys.stdout, "encoding", None) or "ascii"
+  print(text.encode(encoding, errors="backslashreplace").decode(encoding),
+        end=end)
+
+
 def get_dll_dependencies(path):
   """Parses PE header to find direct DLL dependencies."""
   deps = []
@@ -118,7 +130,7 @@ def diagnose_dll_load(path, depth=0, visited=None):
 
   indent = "  " * depth
   if depth == 0:
-    print(f"\n[TensorFlow DLL Diagnostic] Analyzing: {path}")
+    _safe_print(f"\n[TensorFlow DLL Diagnostic] Analyzing: {path}")
 
   # If file doesn't exist, it might be a system DLL
   if not os.path.exists(path):
@@ -145,7 +157,7 @@ def diagnose_dll_load(path, depth=0, visited=None):
     if not dll_name.lower().startswith(
         ("kernel32", "user32", "api-ms-win-", "msvcrt", "ucrtbase")
     ):
-      for p in [target_dir] + os.environ["PATH"].split(os.pathsep):
+      for p in [target_dir] + os.environ.get("PATH", "").split(os.pathsep):
         full_p = os.path.join(p, dll_name)
         if os.path.exists(full_p):
           diagnose_dll_load(full_p, depth + 1, visited)
@@ -183,7 +195,7 @@ def diagnose_dll_load(path, depth=0, visited=None):
             " outdated/missing."
         )
       else:
-        print(f"UNKNOWN ERROR ({err}): {e}")
+        _safe_print(f"UNKNOWN ERROR ({err}): {e}")
       visited.add(dll_name.lower())
 
 
@@ -205,9 +217,9 @@ def run_diagnosis(path=None):
     if os.name == "nt" and os.path.exists(path):
       diagnose_dll_load(path)
     else:
-      print(f"Error: Path does not exist or not on Windows: {path}")
+      _safe_print(f"Error: Path does not exist or not on Windows: {path}")
   except OSError as e:
-    print(f"Diagnostic failed: {e}")
+    _safe_print(f"Diagnostic failed: {e}")
 
 
 if __name__ == "__main__":

From d2c637e6f6e959bbf33c2365a24738a5c902e1a3 Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Tue, 6 Oct 2026 06:59:27 +0530
Subject: [PATCH 0204/1000] bound pooling replay by orig_output in
 FractionalMaxPoolGrad

---
 .../core/kernels/fractional_max_pool_op.cc    | 27 ++++++++++++++++++
 .../nn_ops/fractional_max_pool_op_test.py     | 28 +++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/tensorflow/core/kernels/fractional_max_pool_op.cc b/tensorflow/core/kernels/fractional_max_pool_op.cc
index 939debca0fd..977c83e2217 100644
--- a/tensorflow/core/kernels/fractional_max_pool_op.cc
+++ b/tensorflow/core/kernels/fractional_max_pool_op.cc
@@ -302,6 +302,33 @@ class FractionalMaxPoolGradOp : public OpKernel {
       output_size[i] = tensor_out.dim_size(i);
     }
 
+    // Step 1 replays the forward pooling, writing one entry per visited cell
+    // into buffers sized from orig_output. The batch/depth of orig_input and
+    // the pooling-sequence lengths are independent inputs, so bind them to
+    // orig_output's dimensions; otherwise a crafted sequence or mismatched
+    // orig_input drives out_index past the output buffers.
+    OP_REQUIRES(context, input_size[0] == output_size[0],
+                absl::InvalidArgumentError(absl::StrCat(
+                    "orig_input and orig_output must have the same batch size, "
+                    "got ",
+                    input_size[0], " and ", output_size[0])));
+    OP_REQUIRES(context, input_size[3] == output_size[3],
+                absl::InvalidArgumentError(absl::StrCat(
+                    "orig_input and orig_output must have the same depth, got ",
+                    input_size[3], " and ", output_size[3])));
+    OP_REQUIRES(
+        context, height_seq_tensor.dim_size(0) == output_size[1] + 1,
+        absl::InvalidArgumentError(absl::StrCat(
+            "row_pooling_sequence must have ", output_size[1] + 1,
+            " elements to match the ", output_size[1],
+            " rows of orig_output, but got ", height_seq_tensor.dim_size(0))));
+    OP_REQUIRES(
+        context, width_seq_tensor.dim_size(0) == output_size[2] + 1,
+        absl::InvalidArgumentError(absl::StrCat(
+            "col_pooling_sequence must have ", output_size[2] + 1,
+            " elements to match the ", output_size[2],
+            " cols of orig_output, but got ", width_seq_tensor.dim_size(0))));
+
     // ---------
     // Step 1
     // ---------
diff --git a/tensorflow/python/kernel_tests/nn_ops/fractional_max_pool_op_test.py b/tensorflow/python/kernel_tests/nn_ops/fractional_max_pool_op_test.py
index 33a0f52f46b..ae6ea1b9ac4 100644
--- a/tensorflow/python/kernel_tests/nn_ops/fractional_max_pool_op_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/fractional_max_pool_op_test.py
@@ -728,6 +728,34 @@ class FractionalMaxPoolGradTest(test.TestCase):
             overlapping=overlapping)
         self.evaluate(t)
 
+  def testSeqLengthMismatchRaisesErrorForFractionalMaxPoolGrad(self):
+    # A long pooling sequence paired with a tiny orig_output drives out_index
+    # past the output buffers in the step-1 replay. The product-of-lengths
+    # check only bounds the sequences against orig_input, so this must be
+    # rejected by the orig_output binding instead of writing out of bounds.
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError, "row_pooling_sequence must have"
+    ):
+      with self.cached_session():
+        orig_input = constant_op.constant(
+            1.0, shape=[1, 100, 100, 1], dtype=dtypes.float32)
+        orig_output = constant_op.constant(
+            1.0, shape=[1, 1, 1, 1], dtype=dtypes.float32)
+        out_backprop = constant_op.constant(
+            1.0, shape=[1, 1, 1, 1], dtype=dtypes.float32)
+        row_pooling_sequence = constant_op.constant(
+            0, shape=[5000], dtype=dtypes.int64)
+        col_pooling_sequence = constant_op.constant(
+            0, shape=[2], dtype=dtypes.int64)
+        t = gen_nn_ops.FractionalMaxPoolGrad(
+            orig_input=orig_input,
+            orig_output=orig_output,
+            out_backprop=out_backprop,
+            row_pooling_sequence=row_pooling_sequence,
+            col_pooling_sequence=col_pooling_sequence,
+            overlapping=True)
+        self.evaluate(t)
+
 
 if __name__ == "__main__":
   test.main()

From a6768ab378cb1a830e1335511643dd649de34d70 Mon Sep 17 00:00:00 2001
From: Som Tripathi <somtri@iastate.edu>
Date: Mon, 5 Oct 2026 21:34:52 -0500
Subject: [PATCH 0205/1000] Check grad shape in CPU AvgPoolGrad

AvgPoolingGradOp computed the forward output size from orig_input_shape
but never compared it with the shape of out_backprop. A grad with more
channels than the input passed the existing pixel-index check. The CPU
kernel then wrote past the end of its output buffer, and the process
aborted with heap corruption (#63353, #64642).

Validate out_backprop against the expected forward output shape and
raise InvalidArgumentError with the same message the GPU path uses.
---
 tensorflow/core/kernels/avgpooling_op.cc      | 14 ++++++++++-
 .../kernel_tests/nn_ops/pooling_ops_test.py   | 25 +++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/avgpooling_op.cc b/tensorflow/core/kernels/avgpooling_op.cc
index 627e4cbad49..aa5a21c55f6 100644
--- a/tensorflow/core/kernels/avgpooling_op.cc
+++ b/tensorflow/core/kernels/avgpooling_op.cc
@@ -21,7 +21,8 @@ limitations under the License.
 
 #include <vector>
 
-#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "tensorflow/core/framework/kernel_shape_util.h"
 #include "tensorflow/core/framework/numeric_op.h"
 #include "tensorflow/core/framework/op_kernel.h"
@@ -38,6 +39,7 @@ limitations under the License.
 #include "tensorflow/core/util/overflow.h"
 #include "tensorflow/core/util/padding.h"
 #include "tensorflow/core/util/tensor_format.h"
+#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
 
 #if GOOGLE_CUDA
 #include "third_party/gpus/cudnn/cudnn.h"
@@ -337,6 +339,16 @@ class AvgPoolingGradOp : public OpKernel {
     OP_REQUIRES_OK(context, GetWindowedOutputSize(
                                 in_cols, window_cols, /*dilation_rate=*/1,
                                 col_stride, padding_, &out_width, &pad_cols));
+    TensorShape forward_output_shape;
+    OP_REQUIRES_OK(context, ShapeFromFormatWithStatus(
+                                data_format_, output_shape.dim_size(0),
+                                out_height, out_width, output_shape.dim_size(3),
+                                &forward_output_shape));
+    OP_REQUIRES(
+        context, out_backprop.shape() == forward_output_shape,
+        absl::InvalidArgumentError(absl::StrCat(
+            "Expected grad shape to be ", forward_output_shape.DebugString(),
+            ", but got ", out_backprop.shape().DebugString())));
 
     const T* out_backprop_ptr = out_backprop.flat<T>().data();
     T* input_backprop_ptr = output->flat<T>().data();
diff --git a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
index cb63ddb3760..49df55bb70b 100644
--- a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
@@ -2426,6 +2426,31 @@ class PoolingTest(test.TestCase, parameterized.TestCase):
           )
       )
 
+  @test_util.run_in_graph_and_eager_modes
+  @test_util.disable_xla("Xla does not raise error on out of bounds access")
+  def testAvgPoolGradMismatchedGradShapeRaisesError(self):
+    if test_util.IsMklEnabled():
+      self.skipTest("The oneDNN AvgPoolGrad kernel does not check grad's shape")
+    # Each grad has more channels than the input. The CPU kernel used to write
+    # past the end of its output buffer instead of raising an error.
+    for orig_input_shape, grad_shape, ksize, expected in (
+        ([1, 28, 28, 3], [1, 14, 14, 6], 2, r"\[1,14,14,3\], but got "
+         r"\[1,14,14,6\]"),
+        ([2, 2, 2, 2], [2, 3, 3, 3], 1, r"\[2,2,2,2\], but got \[2,3,3,3\]"),
+    ):
+      with ops.device("/cpu:0"):
+        with self.assertRaisesRegex(
+            (errors_impl.InvalidArgumentError, ValueError),
+            "Expected grad shape to be " + expected):
+          self.evaluate(
+              gen_nn_ops.AvgPoolGrad(
+                  orig_input_shape=orig_input_shape,
+                  grad=array_ops.zeros(grad_shape),
+                  ksize=[1, ksize, ksize, 1],
+                  strides=[1, ksize, ksize, 1],
+                  padding="VALID",
+                  data_format="NHWC"))
+
   @test_util.run_deprecated_v1
   def testShapeFunctionEdgeCases(self):
     # All shapes unknown.

From 0384f0caaa711aabb53256e734e33cb2d9ebf260 Mon Sep 17 00:00:00 2001
From: Som Tripathi <somtri@iastate.edu>
Date: Tue, 6 Oct 2026 00:20:57 -0500
Subject: [PATCH 0206/1000] Use GetTensorDim for the AvgPoolGrad grad shape
 check

The expected grad shape is built with ShapeFromFormatWithStatus and
data_format_, so read its batch and depth with GetTensorDim as well.
The CPU kernel only accepts NHWC, so the expected shape does not
change.
---
 tensorflow/core/kernels/avgpooling_op.cc | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/kernels/avgpooling_op.cc b/tensorflow/core/kernels/avgpooling_op.cc
index aa5a21c55f6..a39a8194126 100644
--- a/tensorflow/core/kernels/avgpooling_op.cc
+++ b/tensorflow/core/kernels/avgpooling_op.cc
@@ -340,10 +340,12 @@ class AvgPoolingGradOp : public OpKernel {
                                 in_cols, window_cols, /*dilation_rate=*/1,
                                 col_stride, padding_, &out_width, &pad_cols));
     TensorShape forward_output_shape;
-    OP_REQUIRES_OK(context, ShapeFromFormatWithStatus(
-                                data_format_, output_shape.dim_size(0),
-                                out_height, out_width, output_shape.dim_size(3),
-                                &forward_output_shape));
+    OP_REQUIRES_OK(
+        context, ShapeFromFormatWithStatus(
+                     data_format_,
+                     GetTensorDim(output_shape, data_format_, 'N'), out_height,
+                     out_width, GetTensorDim(output_shape, data_format_, 'C'),
+                     &forward_output_shape));
     OP_REQUIRES(
         context, out_backprop.shape() == forward_output_shape,
         absl::InvalidArgumentError(absl::StrCat(

From 16b66ec1db8a2a18b19552c7dbe0fb717d3ac5ed Mon Sep 17 00:00:00 2001
From: endorphin13 <elshad.toklayev.business@gmail.com>
Date: Tue, 6 Oct 2026 09:15:07 +0300
Subject: [PATCH 0207/1000] Pin the error message in each negative snapshot
 test

Checking only the status type let these tests pass on an unrelated failure of
the same type. CustomReader::ReadTensors returns InvalidArgument for an
unsupported version, an unsupported compression and a negative tensor size,
and Internal for several parse and decompression failures, so the type alone
does not identify the path under test.

Capture the status and assert on its message as well, using the phrase each
path produces: the non-negative size check, the multiplication overflow, and
the uncompressed size mismatch.

HasSubstr comes from gmock, which this target did not have: platform/test.h
exports gtest only, and the EXPECT_THAT already in this file resolves through
gtest. Add the gmock include and the googletest dependency, as
compression_utils_test does.
---
 tensorflow/core/data/BUILD                  |  1 +
 tensorflow/core/data/snapshot_utils_test.cc | 15 ++++++++++++---
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/data/BUILD b/tensorflow/core/data/BUILD
index 0caa789f4c4..a071e3547b9 100644
--- a/tensorflow/core/data/BUILD
+++ b/tensorflow/core/data/BUILD
@@ -627,6 +627,7 @@ tf_cc_test(
         "//tensorflow/core:test_main",
         "//tensorflow/core/data/service:test_util",
         "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest",
     ],
 )
 
diff --git a/tensorflow/core/data/snapshot_utils_test.cc b/tensorflow/core/data/snapshot_utils_test.cc
index b2640eaa741..dd8f95d07f2 100644
--- a/tensorflow/core/data/snapshot_utils_test.cc
+++ b/tensorflow/core/data/snapshot_utils_test.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include <string>
 #include <vector>
 
+#include <gmock/gmock.h>
 #include "absl/status/status.h"
 #include "tensorflow/core/data/service/test_util.h"
 #include "tensorflow/core/framework/tensor.pb.h"
@@ -39,6 +40,7 @@ namespace {
 
 using ::tensorflow::data::testing::EqualsProto;
 using ::tensorflow::data::testing::LocalTempFilename;
+using ::testing::HasSubstr;
 
 void GenerateTensorVector(tensorflow::DataTypeVector& dtypes,
                           std::vector<Tensor>& tensors) {
@@ -223,7 +225,9 @@ TEST(SnapshotUtilTest, SnappyInvalidTensorShapeIsAnError) {
                               io::compression::kSnappy, /*version=*/1, dtypes,
                               &reader));
   std::vector<Tensor> read_tensors;
-  EXPECT_TRUE(absl::IsInvalidArgument(reader->ReadTensors(&read_tensors)));
+  absl::Status status = reader->ReadTensors(&read_tensors);
+  EXPECT_TRUE(absl::IsInvalidArgument(status));
+  EXPECT_THAT(status.message(), HasSubstr("Expected a non-negative size"));
 
   TF_ASSERT_OK(Env::Default()->DeleteFile(filename));
 }
@@ -243,7 +247,9 @@ TEST(SnapshotUtilTest, SnappyZeroTensorDimensionIsAnError) {
                               io::compression::kSnappy, /*version=*/1, dtypes,
                               &reader));
   std::vector<Tensor> read_tensors;
-  EXPECT_TRUE(absl::IsInternal(reader->ReadTensors(&read_tensors)));
+  absl::Status status = reader->ReadTensors(&read_tensors);
+  EXPECT_TRUE(absl::IsInternal(status));
+  EXPECT_THAT(status.message(), HasSubstr("Uncompressed size mismatch"));
 
   TF_ASSERT_OK(Env::Default()->DeleteFile(filename));
 }
@@ -262,7 +268,10 @@ TEST(SnapshotUtilTest, SnappyOverflowingTensorShapeIsAnError) {
                               io::compression::kSnappy, /*version=*/1, dtypes,
                               &reader));
   std::vector<Tensor> read_tensors;
-  EXPECT_TRUE(absl::IsInvalidArgument(reader->ReadTensors(&read_tensors)));
+  absl::Status status = reader->ReadTensors(&read_tensors);
+  EXPECT_TRUE(absl::IsInvalidArgument(status));
+  EXPECT_THAT(status.message(),
+              HasSubstr("Encountered overflow when multiplying"));
 
   TF_ASSERT_OK(Env::Default()->DeleteFile(filename));
 }

From 242c39f4738aeed4e2a5474c50fba4ad354264cd Mon Sep 17 00:00:00 2001
From: endorphin13 <elshad.toklayev.business@gmail.com>
Date: Tue, 6 Oct 2026 10:06:33 +0300
Subject: [PATCH 0208/1000] Zero-fill the FFT output for an empty input instead
 of emptying its shape

The previous approach forced an FFT axis to 0 in the output when the input
axis was empty. That diverged from RFFTShape, which sizes those axes from
fft_length with no special case for an empty input, so graph construction and
execution disagreed on the shape. It also disagreed with NumPy and the XLA
kernels, which zero-pad the empty signal to fft_length and return a zero
result of the padded size.

Restore the shape calculation in FFTBase::Compute and FFTNBase::Compute, and
instead zero the output in DoFFT and DoFFTN when the input has no elements.
Zero-padding an empty signal gives an all-zero signal, so its transform is
identically zero, which is the mathematically correct result rather than a
placeholder.

Compute no longer returns early for an empty input, since that return is what
left the buffer unwritten; each DoFFT handles the case before calling the
backend, so the cost falls only on this degenerate input. An empty batch
dimension still yields an empty output, which the new guard skips.
---
 tensorflow/core/kernels/fft_ops.cc | 76 +++++++++++++++++-------------
 1 file changed, 44 insertions(+), 32 deletions(-)

diff --git a/tensorflow/core/kernels/fft_ops.cc b/tensorflow/core/kernels/fft_ops.cc
index 0e03c90e812..25150f8d8d2 100644
--- a/tensorflow/core/kernels/fft_ops.cc
+++ b/tensorflow/core/kernels/fft_ops.cc
@@ -19,6 +19,7 @@ limitations under the License.
 
 #include <cstddef>
 #include <cstdint>
+#include <cstring>
 #include <vector>
 
 #include "absl/log/check.h"
@@ -201,14 +202,6 @@ class FFTBase : public OpKernel {
         uint64_t dim = IsForward() && inner_most && fft_shape[i] != 0
                            ? fft_shape[i] / 2 + 1
                            : fft_shape[i];
-        // An empty FFT axis is passed through by the check above. Keep it
-        // empty in the output too. Sizing it from `fft_length` instead would
-        // give a non-empty output for an empty input, and the
-        // `num_elements() == 0` early return below would then hand back the
-        // output buffer without ever writing it.
-        if (input_shape.dim_size(input_index) == 0) {
-          dim = 0;
-        }
         output_shape.set_dim(output_shape.dims() - fft_rank + i, dim);
       }
     } else {
@@ -247,11 +240,9 @@ class FFTBase : public OpKernel {
               "Wrong types for FFT: in=", in.dtype(), " out=", out->dtype())));
     }
 
-    if (input_shape.num_elements() == 0) {
-      DCHECK_EQ(0, output_shape.num_elements());
-      return;
-    }
-
+    // An empty input is handled inside DoFFT, which zeroes the output. The
+    // output is sized from `fft_length`, so it can be non-empty even when the
+    // input is not, and it must not be returned unwritten.
     DoFFT(ctx, in, fft_shape, out);
   }
 
@@ -320,12 +311,11 @@ class FFTNBase : public OpKernel {
                       "fft_length[", i,
                       "] must >= 0, but got: ", fft_length_as_vec(i))));
       fft_shape[i] = fft_length_as_vec(i);
-      auto input_index = input_rank - fft_rank + i;
-      uint64_t dim = fft_shape[i];
       if (IsReal()) {
         bool inner_most = (i == fft_rank - 1);
         uint64_t min_input_dim_length =
             !IsForward() && inner_most ? fft_shape[i] / 2 + 1 : fft_shape[i];
+        auto input_index = input_rank - fft_rank + i;
         OP_REQUIRES(
             ctx,
             // We pass through empty tensors, so special case them here.
@@ -335,20 +325,13 @@ class FFTNBase : public OpKernel {
                 "Input dimension ", input_index,
                 " must have length of at least ", min_input_dim_length,
                 " but got: ", input_shape.dim_size(input_index))));
-        if (IsForward() && inner_most && fft_shape[i] != 0) {
-          dim = fft_shape[i] / 2 + 1;
-        }
+        uint64_t dim = IsForward() && inner_most && fft_shape[i] != 0
+                           ? fft_shape[i] / 2 + 1
+                           : fft_shape[i];
+        output_shape.set_dim(output_shape.dims() - fft_rank + i, dim);
+      } else {
+        output_shape.set_dim(output_shape.dims() - fft_rank + i, fft_shape[i]);
       }
-      // An empty FFT axis is passed through, so keep it empty in the output
-      // too. Sizing it from `fft_length` instead would give a non-empty output
-      // for an empty input, and the `num_elements() == 0` early return below
-      // would then hand back the output buffer without ever writing it. This
-      // applies to the complex transforms as well, which also take an explicit
-      // `fft_length`.
-      if (input_shape.dim_size(input_index) == 0) {
-        dim = 0;
-      }
-      output_shape.set_dim(output_shape.dims() - fft_rank + i, dim);
     }
 
     OP_REQUIRES_OK(ctx, ctx->allocate_output(0, output_shape, &out));
@@ -380,10 +363,7 @@ class FFTNBase : public OpKernel {
               "Wrong types for FFT: in=", in.dtype(), " out=", out->dtype())));
     }
 
-    if (input_shape.num_elements() == 0) {
-      DCHECK_EQ(0, output_shape.num_elements());
-      return;
-    }
+    // As above, DoFFTN zeroes the output for an empty input.
     DoFFTN(ctx, in, fft_shape.data(), axes_shape.data(), out);
   }
 
@@ -411,6 +391,16 @@ class FFTCPU : public FFTBase {
 
   void DoFFT(OpKernelContext* ctx, const Tensor& in, uint64_t* fft_shape,
              Tensor* out) override {
+    // An empty input carries no frequency content, and `fft_length` pads it
+    // with zeros, so the transform is identically zero. The output is sized
+    // from `fft_length` and can be non-empty here, so fill it.
+    if (in.NumElements() == 0) {
+      if (out->NumElements() > 0) {
+        memset(out->data(), 0, out->TotalBytes());
+      }
+      return;
+    }
+
     std::vector<size_t> axes(Rank());
     int batch_dims = in.dims() - FFTRank;
 
@@ -685,6 +675,17 @@ class FFTGPUBase : public FFTBase {
     auto* stream = ctx->op_device_context()->stream();
     OP_REQUIRES(ctx, stream, absl::InternalError("No GPU stream available."));
 
+    // See the CPU kernel: an empty input transforms to an all-zero output,
+    // which is sized from `fft_length` and can be non-empty.
+    if (in.NumElements() == 0) {
+      if (out->NumElements() > 0) {
+        stream_executor::DeviceAddressBase out_bytes(out->data(),
+                                                     out->TotalBytes());
+        OP_REQUIRES_OK(ctx, stream->MemZero(&out_bytes, out->TotalBytes()));
+      }
+      return;
+    }
+
     const TensorShape& input_shape = in.shape();
     const TensorShape& output_shape = out->shape();
 
@@ -875,6 +876,17 @@ class FFTNGPUBase : public FFTNBase {
     auto* stream = ctx->op_device_context()->stream();
     OP_REQUIRES(ctx, stream, absl::InternalError("No GPU stream available."));
 
+    // See the CPU kernel: an empty input transforms to an all-zero output,
+    // which is sized from `fft_length` and can be non-empty.
+    if (in.NumElements() == 0) {
+      if (out->NumElements() > 0) {
+        stream_executor::DeviceAddressBase out_bytes(out->data(),
+                                                     out->TotalBytes());
+        OP_REQUIRES_OK(ctx, stream->MemZero(&out_bytes, out->TotalBytes()));
+      }
+      return;
+    }
+
     Eigen::Map<Eigen::ArrayXi> axes(axes_shape, fft_rank);
     const TensorShape& input_shape = in.shape();
     const TensorShape& output_shape = out->shape();

From 0988d9711f9ac10a839730eb6807f56dba483885 Mon Sep 17 00:00:00 2001
From: endorphin13 <elshad.toklayev.business@gmail.com>
Date: Tue, 6 Oct 2026 10:06:33 +0300
Subject: [PATCH 0209/1000] Assert the padded shape and zero contents, and stop
 skipping XLA

With the output now sized from fft_length and zero-filled, the CPU and GPU
kernels agree with XLA, so the tests no longer need to skip it. Check the
contents as well as the shape, since an unwritten buffer is the defect under
test and only the contents reveal it.

Add a case for an empty batch dimension, which is sized from the input rather
than from fft_length and so stays empty. That is the distinction the earlier
revision of this patch lost.
---
 .../kernel_tests/signal/fft_ops_test.py       | 102 ++++++++++--------
 1 file changed, 57 insertions(+), 45 deletions(-)

diff --git a/tensorflow/python/kernel_tests/signal/fft_ops_test.py b/tensorflow/python/kernel_tests/signal/fft_ops_test.py
index ad2c7dae515..8df28281673 100644
--- a/tensorflow/python/kernel_tests/signal/fft_ops_test.py
+++ b/tensorflow/python/kernel_tests/signal/fft_ops_test.py
@@ -372,17 +372,19 @@ class FFTOpsTest(BaseFFTOpsTest, parameterized.TestCase):
   @test_util.run_gpu_only
   def testEmptyAxis_fftn(self, np_type):
     # Unlike the 1/2/3D kernels, FFTND and IFFTND take an explicit `fft_length`
-    # for complex input too. An empty FFT axis must stay empty in the output
-    # instead of being sized from it: the kernel returns early without writing
-    # the output for an empty input, so a non-empty output would be returned
-    # uninitialized.
+    # for complex input too, so an empty FFT axis is sized from it and the
+    # output must be zero-filled rather than returned unwritten.
     x = np.zeros((2, 0)).astype(np_type)
-    self.assertEqual((2, 0), self._tf_fftn(x, (16,), (-1,)).shape)
-    self.assertEqual((2, 0), self._tf_ifftn(x, (16,), (-1,)).shape)
+    for out in [self._tf_fftn(x, (16,), (-1,)),
+                self._tf_ifftn(x, (16,), (-1,))]:
+      self.assertEqual((2, 16), out.shape)
+      self.assertAllEqual(np.zeros_like(out), out)
 
     x = np.zeros((2, 0, 0)).astype(np_type)
-    self.assertEqual((2, 0, 0), self._tf_fftn(x, (16, 16), (-2, -1)).shape)
-    self.assertEqual((2, 0, 0), self._tf_ifftn(x, (16, 16), (-2, -1)).shape)
+    for out in [self._tf_fftn(x, (16, 16), (-2, -1)),
+                self._tf_ifftn(x, (16, 16), (-2, -1))]:
+      self.assertEqual((2, 16, 16), out.shape)
+      self.assertAllEqual(np.zeros_like(out), out)
 
   @parameterized.parameters(itertools.product(
       (1,), range(3), (np.complex64, np.complex128)))
@@ -630,52 +632,63 @@ class RFFTOpsTest(BaseFFTOpsTest, parameterized.TestCase):
 
   @parameterized.parameters(
       itertools.product(VALID_FFT_RANKS, range(3), (np.float32, np.float64)))
-  @test_util.disable_xla("XLA zero-pads empty FFT axes to fft_length")
   def test_empty_with_explicit_fft_length(self, rank, extra_dims, np_rtype):
     # An empty FFT axis skips the "input dimension must be at least
-    # fft_length" requirement. `fft_length` must not then size that axis in
-    # the output: the kernel returns early without writing the output for an
-    # empty input, so a non-empty output would be returned uninitialized.
+    # fft_length" requirement, and `fft_length` still sizes that axis in the
+    # output, as it does for a short axis. Zero-padding an empty signal gives
+    # an all-zero signal, so the transform is all zeros, which is also what
+    # NumPy and the XLA kernels produce. The kernel used to return that buffer
+    # without writing it, handing the caller uninitialized heap.
     np_ctype = np.complex64 if np_rtype == np.float32 else np.complex128
     fft_length = (16,) * rank
 
-    # Batch dimensions must stay non-empty: if every dimension is zero the
-    # output is empty regardless of how the FFT axes are sized, and the test
-    # would pass against the unfixed kernel too. Assert the exact shape rather
-    # than only the element count, so that a wrongly sized axis or an altered
-    # batch dimension is caught as well.
-    expected_shape = (2,) * extra_dims + (0,) * rank
+    # Batch dimensions stay non-empty, so the output is non-empty and an
+    # unwritten buffer would be observable.
+    x_shape = (2,) * extra_dims + (0,) * rank
+    batch = (2,) * extra_dims
 
-    x = np.zeros(expected_shape).astype(np_rtype)
+    x = np.zeros(x_shape).astype(np_rtype)
     out_fwd = self._tf_fft(x, rank, fft_length)
-    self.assertEqual(expected_shape, out_fwd.shape)
-    self.assertEqual(0, out_fwd.size)
+    self.assertEqual(batch + (16,) * (rank - 1) + (9,), out_fwd.shape)
+    self.assertAllEqual(np.zeros_like(out_fwd), out_fwd)
 
-    x = np.zeros(expected_shape).astype(np_ctype)
+    x = np.zeros(x_shape).astype(np_ctype)
     out_bwd = self._tf_ifft(x, rank, fft_length)
-    self.assertEqual(expected_shape, out_bwd.shape)
-    self.assertEqual(0, out_bwd.size)
+    self.assertEqual(batch + (16,) * rank, out_bwd.shape)
+    self.assertAllEqual(np.zeros_like(out_bwd), out_bwd)
 
   @parameterized.parameters((np.float32,), (np.float64,))
-  @test_util.disable_xla("XLA zero-pads empty FFT axes to fft_length")
   def test_empty_axis_next_to_non_empty_axis(self, np_rtype):
-    # Only an empty FFT axis stays empty. The other FFT axis is still sized
-    # from `fft_length`, and the batch dimension is untouched.
+    # An empty FFT axis is sized from `fft_length` whichever axis it is, so
+    # both inputs below give the same output shape, filled with zeros.
     np_ctype = np.complex64 if np_rtype == np.float32 else np.complex128
     fft_length = (16, 16)
 
     # Forward: the inner-most output axis is fft_length / 2 + 1.
-    x = np.zeros((2, 0, 16)).astype(np_rtype)
-    self.assertEqual((2, 0, 9), self._tf_fft(x, 2, fft_length).shape)
-    x = np.zeros((2, 16, 0)).astype(np_rtype)
-    self.assertEqual((2, 16, 0), self._tf_fft(x, 2, fft_length).shape)
+    for x_shape in [(2, 0, 16), (2, 16, 0)]:
+      out = self._tf_fft(np.zeros(x_shape).astype(np_rtype), 2, fft_length)
+      self.assertEqual((2, 16, 9), out.shape)
+      self.assertAllEqual(np.zeros_like(out), out)
 
     # Inverse: the inner-most input axis must hold fft_length / 2 + 1 values,
     # and the inner-most output axis is fft_length.
-    x = np.zeros((2, 0, 9)).astype(np_ctype)
-    self.assertEqual((2, 0, 16), self._tf_ifft(x, 2, fft_length).shape)
-    x = np.zeros((2, 16, 0)).astype(np_ctype)
-    self.assertEqual((2, 16, 0), self._tf_ifft(x, 2, fft_length).shape)
+    for x_shape in [(2, 0, 9), (2, 16, 0)]:
+      out = self._tf_ifft(np.zeros(x_shape).astype(np_ctype), 2, fft_length)
+      self.assertEqual((2, 16, 16), out.shape)
+      self.assertAllEqual(np.zeros_like(out), out)
+
+  @parameterized.parameters((np.float32,), (np.float64,))
+  def test_empty_batch_dimension_is_preserved(self, np_rtype):
+    # Only the FFT axes are sized from `fft_length`. An empty batch dimension
+    # passes through, so the output stays empty and nothing is written.
+    np_ctype = np.complex64 if np_rtype == np.float32 else np.complex128
+    fft_length = (16,)
+
+    x = np.zeros((0, 16)).astype(np_rtype)
+    self.assertEqual((0, 9), self._tf_fft(x, 1, fft_length).shape)
+
+    x = np.zeros((0, 9)).astype(np_ctype)
+    self.assertEqual((0, 16), self._tf_ifft(x, 1, fft_length).shape)
 
   @parameterized.parameters(itertools.product(
       VALID_FFT_RANKS, range(3), (5, 6), (np.float32, np.float64)))
@@ -817,21 +830,20 @@ class RFFTOpsTest(BaseFFTOpsTest, parameterized.TestCase):
   @test_util.run_gpu_only
   def testEmptyAxis_rfftn(self, np_rtype):
     # The ND counterpart of test_empty_with_explicit_fft_length: an empty FFT
-    # axis must stay empty, while a non-empty FFT axis is still sized from
-    # `fft_length`.
+    # axis is sized from `fft_length` and the output is zero-filled.
     np_ctype = np.complex64 if np_rtype == np.float32 else np.complex128
     fft_length = (16, 16)
     axes = (-2, -1)
 
-    x = np.zeros((2, 0, 0)).astype(np_rtype)
-    self.assertEqual((2, 0, 0), self._tf_fftn(x, fft_length, axes).shape)
-    x = np.zeros((2, 0, 16)).astype(np_rtype)
-    self.assertEqual((2, 0, 9), self._tf_fftn(x, fft_length, axes).shape)
+    for x_shape in [(2, 0, 0), (2, 0, 16)]:
+      out = self._tf_fftn(np.zeros(x_shape).astype(np_rtype), fft_length, axes)
+      self.assertEqual((2, 16, 9), out.shape)
+      self.assertAllEqual(np.zeros_like(out), out)
 
-    x = np.zeros((2, 0, 0)).astype(np_ctype)
-    self.assertEqual((2, 0, 0), self._tf_ifftn(x, fft_length, axes).shape)
-    x = np.zeros((2, 0, 9)).astype(np_ctype)
-    self.assertEqual((2, 0, 16), self._tf_ifftn(x, fft_length, axes).shape)
+    for x_shape in [(2, 0, 0), (2, 0, 9)]:
+      out = self._tf_ifftn(np.zeros(x_shape).astype(np_ctype), fft_length, axes)
+      self.assertEqual((2, 16, 16), out.shape)
+      self.assertAllEqual(np.zeros_like(out), out)
 
   @parameterized.parameters(itertools.product(
       (1,), range(3), (64, 128), (np.float32, np.float64)))

From 9104989d527df5009a51783cd7a66baaa8d154df Mon Sep 17 00:00:00 2001
From: abhirupone <abhirup.karmakar00@gmail.com>
Date: Tue, 6 Oct 2026 15:51:01 +0530
Subject: [PATCH 0210/1000] Fix static shape validation for RaggedTensor from
 row splits under XLA

Under jit_compile=True, tf.raw_ops.Assert is treated as a no-op by XLA.
When constructing RaggedTensor instances with constant row splits,
mismatched values.shape[0] was not validated statically, allowing
malformed ragged tensors to be processed silently.

Add static validation against constant row_splits and row_lengths in
RowPartition, and assert shape compatibility in RaggedTensor._from_row_partition.
Add regression tests for eager, graph, and jit_compile execution.

Fixes #128452
---
 tensorflow/python/ops/ragged/ragged_tensor.py | 12 +++++-
 .../python/ops/ragged/ragged_tensor_test.py   | 16 ++++++++
 tensorflow/python/ops/ragged/row_partition.py | 40 ++++++++++++++++++-
 .../python/ops/ragged/row_partition_test.py   | 20 ++++++++++
 4 files changed, 85 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/ops/ragged/ragged_tensor.py b/tensorflow/python/ops/ragged/ragged_tensor.py
index 727001c8e80..d479ad8128a 100644
--- a/tensorflow/python/ops/ragged/ragged_tensor.py
+++ b/tensorflow/python/ops/ragged/ragged_tensor.py
@@ -340,9 +340,19 @@ class RaggedTensor(
                       f"Received {validate}.")
     values, row_partition = cls._convert_values_and_partition(
         values, row_partition, "partition")
-    if row_partition._has_precomputed_value_rowids():  # pylint: disable=protected-access
+    # pylint: disable=protected-access
+    static_nvals = row_partition._static_nvals_or_constant()
+    if row_partition._has_precomputed_value_rowids():
+      # pylint: enable=protected-access
       value_rowids_shape = row_partition.value_rowids().shape
       values.shape[:1].assert_is_compatible_with(value_rowids_shape)
+    elif static_nvals is not None:
+      # pylint: enable=protected-access
+      values.shape[:1].assert_is_compatible_with(
+          tensor_shape.TensorShape([static_nvals]))
+    else:
+      # pylint: enable=protected-access
+      pass
     if validate:
       msg = "Arguments to _from_row_partition do not form a valid RaggedTensor"
       nvals = _nrows(values, row_partition.dtype)
diff --git a/tensorflow/python/ops/ragged/ragged_tensor_test.py b/tensorflow/python/ops/ragged/ragged_tensor_test.py
index b6cd505581b..d0466f996fa 100644
--- a/tensorflow/python/ops/ragged/ragged_tensor_test.py
+++ b/tensorflow/python/ops/ragged/ragged_tensor_test.py
@@ -323,6 +323,22 @@ class RaggedTensorTest(test_util.TensorFlowTestCase, parameterized.TestCase):
     with self.assertRaisesRegex(ValueError, err_msg):
       RaggedTensor.from_row_splits([], [])
 
+  def testFromRowSplitsStaticShapeMismatch(self):
+    values = constant_op.constant(
+        [[4.0, 5.0], [6.0, 7.0]], dtype=dtypes.float32)
+    with self.assertRaisesRegex(
+        ValueError, r"Shapes \(2,\) and \(4,\) are incompatible"):
+      RaggedTensor.from_row_splits(values, [0, 2, 4])
+
+  def testFromRowSplitsStaticShapeMismatchUnderJit(self):
+    @def_function.function(jit_compile=True)
+    def fn(x):
+      values = array_ops.reshape(x, [-1, 2])
+      return RaggedTensor.from_row_splits(values, [0, 2, 4])
+
+    with self.assertRaises(ValueError):
+      fn(constant_op.constant([4.0, 5.0, 6.0, 7.0]))
+
   def testFromRowStarts(self):
     values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
     row_starts = constant_op.constant([0, 2, 2, 5, 6], dtypes.int64)
diff --git a/tensorflow/python/ops/ragged/row_partition.py b/tensorflow/python/ops/ragged/row_partition.py
index f47b14f4eb2..2c29f3b0dba 100644
--- a/tensorflow/python/ops/ragged/row_partition.py
+++ b/tensorflow/python/ops/ragged/row_partition.py
@@ -364,6 +364,22 @@ class RowPartition(composite_tensor.CompositeTensor):
       row_splits.shape.assert_has_rank(1)
 
       if validate:
+        row_splits_const = tensor_util.constant_value(row_splits)
+        if row_splits_const is not None:
+          if len(row_splits_const) == 0:
+            raise ValueError("row_splits tensor may not be empty.")
+          if row_splits_const[0] != 0:
+            raise ValueError(
+                "Arguments to from_row_splits do not form a valid "
+                f"RaggedTensor: row_splits[0] must be zero, "
+                f"got {row_splits_const[0]}")
+          for i in range(len(row_splits_const) - 1):
+            if row_splits_const[i] > row_splits_const[i + 1]:
+              raise ValueError(
+                  "Arguments to from_row_splits do not form a valid "
+                  "RaggedTensor: row_splits must be monotonic increasing, "
+                  f"but row_splits[{i}] ({row_splits_const[i]}) > "
+                  f"row_splits[{i+1}] ({row_splits_const[i+1]})")
         msg = "Arguments to from_row_splits do not form a valid RaggedTensor:"
         checks = [
             check_ops.assert_rank(row_splits, 1, message=(msg + "rank")),
@@ -886,11 +902,31 @@ class RowPartition(composite_tensor.CompositeTensor):
     if self._nvals is not None:
       nvals = tensor_util.constant_value(self._nvals)
       if nvals is not None:
-        return nvals
+        return int(nvals)
     if self._value_rowids is not None:
       nvals = tensor_shape.dimension_at_index(self._value_rowids.shape, 0)
       if nvals.value is not None:
-        return nvals.value
+        return int(nvals.value)
+    return None
+
+  def _static_nvals_or_constant(self):
+    """The number of values in this partition from shape or constant values."""
+    static_n = self.static_nvals
+    if static_n is not None:
+      return static_n
+    if self._row_splits is not None:
+      row_splits_const = tensor_util.constant_value(self._row_splits)
+      if row_splits_const is not None and len(row_splits_const) > 0:
+        return int(row_splits_const[-1])
+    if self._row_lengths is not None:
+      row_lengths_const = tensor_util.constant_value(self._row_lengths)
+      if row_lengths_const is not None:
+        return int(sum(row_lengths_const))
+    if self._uniform_row_length is not None and self._nrows is not None:
+      nrows = self.static_nrows
+      row_len = self.static_uniform_row_length
+      if nrows is not None and row_len is not None:
+        return int(nrows * row_len)
     return None
 
   @property
diff --git a/tensorflow/python/ops/ragged/row_partition_test.py b/tensorflow/python/ops/ragged/row_partition_test.py
index f4687d20cb6..32afffea383 100644
--- a/tensorflow/python/ops/ragged/row_partition_test.py
+++ b/tensorflow/python/ops/ragged/row_partition_test.py
@@ -912,6 +912,26 @@ class RowPartitionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
       self.assertIsNone(static_nrows)
     foo(array_ops.constant([0, 3, 4, 5], dtype=dtypes.int32))
 
+  def testStaticNvalsFromRowSplits(self):
+    rp = RowPartition.from_row_splits([0, 3, 4, 5])
+    self.assertEqual(5, rp._static_nvals_or_constant())
+
+  def testStaticNvalsFromRowLengths(self):
+    rp = RowPartition.from_row_lengths([3, 1, 1])
+    self.assertEqual(5, rp._static_nvals_or_constant())
+
+  def testStaticNvalsFromUniformRowLength(self):
+    rp = RowPartition.from_uniform_row_length(3, nrows=4)
+    self.assertAllEqual(12, rp.static_nvals)
+
+  def testFromRowSplitsStaticValidationNonZero(self):
+    with self.assertRaisesRegex(ValueError, r"row_splits\[0\] must be zero"):
+      RowPartition.from_row_splits([1, 2, 4])
+
+  def testFromRowSplitsStaticValidationNonMonotonic(self):
+    with self.assertRaisesRegex(ValueError, r"must be monotonic increasing"):
+      RowPartition.from_row_splits([0, 5, 2])
+
 
 @test_util.run_all_in_graph_and_eager_modes
 class RowPartitionSpecTest(test_util.TensorFlowTestCase,

From 2f7c0328bb88507f1a8c6c8e1e42da056bfcfcea Mon Sep 17 00:00:00 2001
From: abhirupone <abhirup.karmakar00@gmail.com>
Date: Tue, 6 Oct 2026 16:05:36 +0530
Subject: [PATCH 0211/1000] Address review feedback on RaggedTensor static
 validation

- Clean up pylint control directives and remove redundant else branch in ragged_tensor.py
- Use vectorized np.diff for constant row_splits monotonicity checks in row_partition.py
- Use array .sum() for constant row lengths sum
---
 tensorflow/python/ops/ragged/ragged_tensor.py |  6 +-----
 tensorflow/python/ops/ragged/row_partition.py | 17 +++++++++--------
 2 files changed, 10 insertions(+), 13 deletions(-)

diff --git a/tensorflow/python/ops/ragged/ragged_tensor.py b/tensorflow/python/ops/ragged/ragged_tensor.py
index d479ad8128a..1c9e1a0b3ea 100644
--- a/tensorflow/python/ops/ragged/ragged_tensor.py
+++ b/tensorflow/python/ops/ragged/ragged_tensor.py
@@ -343,16 +343,12 @@ class RaggedTensor(
     # pylint: disable=protected-access
     static_nvals = row_partition._static_nvals_or_constant()
     if row_partition._has_precomputed_value_rowids():
-      # pylint: enable=protected-access
       value_rowids_shape = row_partition.value_rowids().shape
       values.shape[:1].assert_is_compatible_with(value_rowids_shape)
     elif static_nvals is not None:
-      # pylint: enable=protected-access
       values.shape[:1].assert_is_compatible_with(
           tensor_shape.TensorShape([static_nvals]))
-    else:
-      # pylint: enable=protected-access
-      pass
+    # pylint: enable=protected-access
     if validate:
       msg = "Arguments to _from_row_partition do not form a valid RaggedTensor"
       nvals = _nrows(values, row_partition.dtype)
diff --git a/tensorflow/python/ops/ragged/row_partition.py b/tensorflow/python/ops/ragged/row_partition.py
index 2c29f3b0dba..de101baa650 100644
--- a/tensorflow/python/ops/ragged/row_partition.py
+++ b/tensorflow/python/ops/ragged/row_partition.py
@@ -373,13 +373,14 @@ class RowPartition(composite_tensor.CompositeTensor):
                 "Arguments to from_row_splits do not form a valid "
                 f"RaggedTensor: row_splits[0] must be zero, "
                 f"got {row_splits_const[0]}")
-          for i in range(len(row_splits_const) - 1):
-            if row_splits_const[i] > row_splits_const[i + 1]:
-              raise ValueError(
-                  "Arguments to from_row_splits do not form a valid "
-                  "RaggedTensor: row_splits must be monotonic increasing, "
-                  f"but row_splits[{i}] ({row_splits_const[i]}) > "
-                  f"row_splits[{i+1}] ({row_splits_const[i+1]})")
+          diff = np.diff(row_splits_const)
+          if np.any(diff < 0):
+            bad_idx = int(np.argmax(diff < 0))
+            raise ValueError(
+                "Arguments to from_row_splits do not form a valid "
+                "RaggedTensor: row_splits must be monotonic increasing, "
+                f"but row_splits[{bad_idx}] ({row_splits_const[bad_idx]}) > "
+                f"row_splits[{bad_idx + 1}] ({row_splits_const[bad_idx + 1]})")
         msg = "Arguments to from_row_splits do not form a valid RaggedTensor:"
         checks = [
             check_ops.assert_rank(row_splits, 1, message=(msg + "rank")),
@@ -921,7 +922,7 @@ class RowPartition(composite_tensor.CompositeTensor):
     if self._row_lengths is not None:
       row_lengths_const = tensor_util.constant_value(self._row_lengths)
       if row_lengths_const is not None:
-        return int(sum(row_lengths_const))
+        return int(row_lengths_const.sum())
     if self._uniform_row_length is not None and self._nrows is not None:
       nrows = self.static_nrows
       row_len = self.static_uniform_row_length

From aa9cb3d7b1410e2a90bf326ca575399fb3dff0e8 Mon Sep 17 00:00:00 2001
From: Saurabh Singh <geniusthegreat0963@gmail.com>
Date: Sun, 4 Oct 2026 18:31:16 +0530
Subject: [PATCH 0212/1000] Fix isclose/allclose returning inverted results for
 infinities

The floating-point branch of tf.experimental.numpy.isclose applied the
finite tolerance formula |a - b| <= atol + rtol * |b| directly to all
values including infinities.  This produced wrong results:

  - Same-sign: inf - inf = NaN, abs(NaN) <= ... is False (should be True)
  - Opposite-sign: abs(-inf - inf) = inf, inf <= inf is True (should be False)
  - Mixed finite/inf: abs(1.0 - inf) = inf, inf <= inf is True (should be False)

Replace the previous where_v2-based infinity fix with an is_finite
mask approach: the tolerance result is AND-ed with both_finite so
only finite-finite pairs use the formula, then OR-ed with equality
so identical infinities (same sign) are still close.  This avoids
the where_v2 overhead on the fast path.

For complex types, is_finite and is_nan are applied to the magnitude
or real/imag parts respectively, since these ops lack complex kernels.

Add testIsCloseInfinities regression test parameterized over float32,
float64, complex64, and complex128, covering element-wise isclose and
scalar allclose with same-sign, opposite-sign, mixed finite/infinite,
and complex infinite pairs.

Fixes #128480
---
 .../python/ops/numpy_ops/np_math_ops.py       | 30 +++++++++++++++++--
 .../python/ops/numpy_ops/np_math_ops_test.py  | 28 +++++++++++++++++
 2 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py
index 9b0490eb60b..3fc90a529f9 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@@ -601,9 +601,35 @@ def isclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False):  # pylint: disable=m
     if np.issubdtype(dtype.as_numpy_dtype, np.inexact):
       rtol_ = ops.convert_to_tensor(rtol, dtype.real_dtype)
       atol_ = ops.convert_to_tensor(atol, dtype.real_dtype)
-      result = math_ops.abs(a - b) <= atol_ + rtol_ * math_ops.abs(b)
+
+      b_abs = math_ops.abs(b)
+      result = math_ops.abs(a - b) <= atol_ + rtol_ * b_abs
+
+      # Infinities must be handled separately. The finite tolerance formula
+      # evaluates to True for mixed finite/infinite pairs like (1.0, inf)
+      # because inf <= inf. Match NumPy: elements are close if they are
+      # both finite and within tolerance, or if they are equal.
+      if dtype.is_floating or dtype.is_complex:
+        if dtype.is_complex:
+          a_check = math_ops.abs(a)
+          b_check = b_abs
+        else:
+          a_check = a
+          b_check = b
+        both_finite = math_ops.is_finite(a_check) & math_ops.is_finite(b_check)
+        result = (both_finite & result) | math_ops.equal(a, b)
+
       if equal_nan:
-        result = result | (math_ops.is_nan(a) & math_ops.is_nan(b))
+        if dtype.is_complex:
+          a_nan = math_ops.is_nan(math_ops.real(a)) | math_ops.is_nan(
+              math_ops.imag(a)
+          )
+          b_nan = math_ops.is_nan(math_ops.real(b)) | math_ops.is_nan(
+              math_ops.imag(b)
+          )
+          result = result | (a_nan & b_nan)
+        else:
+          result = result | (math_ops.is_nan(a) & math_ops.is_nan(b))
       return result
     elif np.issubdtype(dtype.as_numpy_dtype, np.integer):
       # Use the operands' own arithmetic instead of casting to float, so
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 84869bc9458..f570ad39929 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
@@ -383,6 +383,34 @@ class MathTest(test.TestCase, parameterized.TestCase):
         np_math_ops.isclose(a, b, equal_nan=equal_nan),
         np.isclose(a, b, equal_nan=equal_nan))
 
+  @parameterized.parameters(
+      [np.float32, np.float64, np.complex64, np.complex128]
+  )
+  def testIsCloseInfinities(self, dtype):
+    # Regression test for GitHub issue 128480: same-sign infinities must be
+    # close and opposite-sign infinities must not be close, matching NumPy.
+    # Also covers mixed finite/infinite pairs like (1.0, inf).
+    a_vals = [np.inf, -np.inf, -np.inf, np.inf, 1.0, 1.0, np.inf]
+    b_vals = [np.inf, -np.inf, np.inf, -np.inf, 1.0, np.inf, 1.0]
+    if np.issubdtype(dtype, np.complexfloating):
+      a_vals.extend([complex(np.inf, np.inf), complex(np.inf, np.inf)])
+      b_vals.extend([complex(np.inf, 0), complex(np.inf, np.inf)])
+
+    a = np.array(a_vals, dtype)
+    b = np.array(b_vals, dtype)
+    self.match(
+        np_math_ops.isclose(a, b),
+        np.isclose(a, b),
+    )
+    self.assertEqual(
+        bool(np_math_ops.allclose([dtype(np.inf)], [dtype(np.inf)])),
+        np.allclose([dtype(np.inf)], [dtype(np.inf)]),
+    )
+    self.assertEqual(
+        bool(np_math_ops.allclose([dtype(-np.inf)], [dtype(np.inf)])),
+        np.allclose([dtype(-np.inf)], [dtype(np.inf)]),
+    )
+
   @parameterized.parameters([np.int32, np.int64])
   def testIsCloseIntegerTolerances(self, dtype):
     # Regression test for GitHub issue 108657: integer inputs used to be

From 2aaf6872899d666199460febf0184a325ecc3a1d Mon Sep 17 00:00:00 2001
From: RadikHoroshev <rodion111@gmail.com>
Date: Tue, 6 Oct 2026 16:35:38 +0300
Subject: [PATCH 0213/1000] Address TensorList bounds review: 1<<24 cap,
 grow-only resize guard

Use a single kMaxTensorListElements constant for TensorListReserve and
TensorListResize, apply the cap only when growing, and add list_ops tests.
---
 tensorflow/core/kernels/list_kernels.cc | 37 ++++++++++++-------------
 1 file changed, 18 insertions(+), 19 deletions(-)

diff --git a/tensorflow/core/kernels/list_kernels.cc b/tensorflow/core/kernels/list_kernels.cc
index dbecdaf58cc..1407efd8d4d 100644
--- a/tensorflow/core/kernels/list_kernels.cc
+++ b/tensorflow/core/kernels/list_kernels.cc
@@ -44,6 +44,13 @@ limitations under the License.
 
 namespace tensorflow {
 
+namespace {
+// Upper bound on TensorList reserve/resize element count (CWE-770). Each slot
+// is a Tensor (~24-32 bytes); INT32_MAX would allocate ~48-64 GB of metadata.
+// 2^24 slots (~400 MB metadata) leaves room for large TensorArray / pfor workloads.
+constexpr int32_t kMaxTensorListElements = 1 << 24;
+}  // namespace
+
 typedef Eigen::ThreadPoolDevice CPUDevice;
 
 absl::Status TensorShapeFromTensor(const Tensor& t, PartialTensorShape* out) {
@@ -343,20 +350,11 @@ class TensorListReserve : public OpKernel {
                     absl::StrCat("The num_elements to reserve must be a "
                                  "non negative number, but got ",
                                  num_elements)));
-    // Upper bound guard (CWE-770).  Each reserved slot is a Tensor object
-    // (~24-32 bytes), so INT32_MAX slots would allocate ~48-64 GB of
-    // metadata alone.  1 048 576 (2^20) slots use ~24 MB, which covers
-    // all practical ML workloads (RNNs, sequences, time series).
-    // For comparison, EmptyTensorList exposes a max_num_elements input
-    // that TensorListPushBack enforces; TensorListReserve lacked any
-    // equivalent guard.
-    static constexpr int32_t kMaxReserveElements = 1 << 20;  // 1 048 576
-    OP_REQUIRES(c, num_elements <= kMaxReserveElements,
+    OP_REQUIRES(c, num_elements <= kMaxTensorListElements,
                 absl::InvalidArgumentError(absl::StrCat(
                     "TensorListReserve: num_elements (", num_elements,
-                    ") exceeds maximum allowed (", kMaxReserveElements,
-                    "). Use EmptyTensorList with TensorListPushBack "
-                    "for dynamically-sized lists.")));
+                    ") exceeds maximum allowed elements (",
+                    kMaxTensorListElements, ").")));
     TensorList output;
     output.element_shape = element_shape;
     output.element_dtype = element_dtype_;
@@ -405,13 +403,14 @@ class TensorListResize : public OpKernel {
         c, size >= 0,
         absl::InvalidArgumentError(absl::StrCat(
             "TensorListSlice expects size to be non-negative. Got: ", size)));
-    // Same upper bound as TensorListReserve (CWE-770).
-    static constexpr int32_t kMaxResizeElements = 1 << 20;  // 1 048 576
-    OP_REQUIRES(c, size <= kMaxResizeElements,
-                absl::InvalidArgumentError(absl::StrCat(
-                    "TensorListResize: size (", size,
-                    ") exceeds maximum allowed (", kMaxResizeElements,
-                    ").")));
+    // Only growing the list allocates; shrinking must always be allowed.
+    if (size > input_list->tensors().size()) {
+      OP_REQUIRES(c, size <= kMaxTensorListElements,
+                  absl::InvalidArgumentError(absl::StrCat(
+                      "TensorListResize: requested size (", size,
+                      ") exceeds maximum allowed elements (",
+                      kMaxTensorListElements, ").")));
+    }
 
     std::unique_ptr<Tensor> maybe_result =
         c->forward_input(0, 0, DT_VARIANT, TensorShape{},

From b975037389e51544d1ab1c7c12a7e6dfaf268bd7 Mon Sep 17 00:00:00 2001
From: RadikHoroshev <rodion111@gmail.com>
Date: Tue, 6 Oct 2026 16:35:42 +0300
Subject: [PATCH 0214/1000] Add TensorList max element count regression tests

---
 .../data_structures/list_ops_test.py          | 23 +++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/tensorflow/python/kernel_tests/data_structures/list_ops_test.py b/tensorflow/python/kernel_tests/data_structures/list_ops_test.py
index b60c59987b3..23047b5ab4b 100644
--- a/tensorflow/python/kernel_tests/data_structures/list_ops_test.py
+++ b/tensorflow/python/kernel_tests/data_structures/list_ops_test.py
@@ -1714,6 +1714,29 @@ class ListOpsTest(test_util.TensorFlowTestCase, parameterized.TestCase):
                                 r"\w+ must be a scalar"):
       self.evaluate(gen_list_ops.TensorListResize(input_handle=l, size=size))
 
+  @test_util.run_in_graph_and_eager_modes
+  def testReserveExceedingMaxElementsFails(self):
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "exceeds maximum allowed elements"):
+      l = list_ops.tensor_list_reserve(
+          element_shape=[], num_elements=(1 << 24) + 1,
+          element_dtype=dtypes.float32)
+      self.evaluate(l)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testReserveAtMaxElementsSucceeds(self):
+    l = list_ops.tensor_list_reserve(
+        element_shape=[], num_elements=1 << 24, element_dtype=dtypes.float32)
+    self.assertEqual(self.evaluate(list_ops.tensor_list_length(l)), 1 << 24)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testResizeExceedingMaxElementsFails(self):
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "exceeds maximum allowed elements"):
+      l = list_ops.tensor_list_from_tensor([1., 2., 3.], element_shape=[])
+      l = list_ops.tensor_list_resize(l, (1 << 24) + 1)
+      self.evaluate(l)
+
   @test_util.run_deprecated_v1
   @test_util.enable_control_flow_v2
   def testSkipEagerResizeGrad(self):

From eb7e3e46a43b9690efcd88011500efca9d0bc791 Mon Sep 17 00:00:00 2001
From: Boris Kozak <boriskozak@gmail.com>
Date: Tue, 6 Oct 2026 15:21:03 +0000
Subject: [PATCH 0215/1000] Fix tf.math.polyval returning inf instead of nan
 for non-finite x

The Horner accumulator started at coeffs[0], so a finite polynomial
evaluated at inf produced inf (1 + 1*inf), while numpy.polyval starts
its accumulator at zero and yields nan (0*inf + 1). Initialize the
accumulator at coeffs[0] + x*0 instead: a no-op for finite x, but
non-finite x now propagates to NaN, matching numpy. This also subsumes
the single-coefficient x*0 special case.

Fixes #111414.
---
 .../python/kernel_tests/math_ops/cwise_ops_test.py  | 13 +++++++++++++
 tensorflow/python/ops/math_ops.py                   | 11 +++++------
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py b/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
index ebaf23c6edc..8c2772deb64 100644
--- a/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
@@ -1340,6 +1340,19 @@ class PolyvalTest(test.TestCase):
       self.assertEqual(tf_val.shape, x.shape)
       self.assertAllClose(np_val, self.evaluate(tf_val))
 
+  def testNonFiniteXPropagation(self):
+    # Non-finite x must propagate to NaN, matching numpy.polyval, which
+    # starts its Horner accumulator at zero (0 * inf + c -> nan).
+    for x_val in [float("inf"), float("-inf"), float("nan")]:
+      with self.subTest(x=x_val):
+        x = np.array([1.0, x_val, 3.0], dtype=np.float32)
+        coeffs = [np.float32(1.0), np.float32(1.0)]
+        with np.errstate(invalid="ignore"):
+          np_val = np.polyval(coeffs, x)
+        with self.cached_session():
+          tf_val = math_ops.polyval(coeffs, x)
+          self.assertAllClose(np_val, self.evaluate(tf_val))
+
   def test_coeffs_raise(self):
     x = np.random.rand(2, 2).astype(np.float32)
     coeffs = {}
diff --git a/tensorflow/python/ops/math_ops.py b/tensorflow/python/ops/math_ops.py
index 261a0d1bcf4..5e9c0b20410 100644
--- a/tensorflow/python/ops/math_ops.py
+++ b/tensorflow/python/ops/math_ops.py
@@ -5625,14 +5625,13 @@ def polyval(coeffs, x, name=None):
         ops.convert_to_tensor(coeff, name=("coeff_%d" % index))
         for index, coeff in enumerate(coeffs)
     ]
-    p = coeffs[0]
+    # Start the Horner accumulator at coeffs[0] + x*0 (rather than coeffs[0])
+    # so that non-finite x (inf/nan) propagates to NaN on the first multiply,
+    # matching numpy.polyval, which starts its accumulator at zero. For finite
+    # x this is a no-op, and it also broadcasts the result against x's shape.
+    p = coeffs[0] + x * 0
     for c in coeffs[1:]:
       p = c + p * x
-    # For single-coefficient polynomials, the loop above never executes,
-    # so x is unused. Add x*0 to broadcast against x's shape and
-    # propagate NaN, matching numpy.polyval behavior.
-    if len(coeffs) == 1:
-      p = p + x * 0
     return p
 
 

From 5312f411d54d8c7382061b4143792abffe1b141d Mon Sep 17 00:00:00 2001
From: abhirupone <abhirup.karmakar00@gmail.com>
Date: Tue, 6 Oct 2026 20:57:09 +0530
Subject: [PATCH 0216/1000] Address review feedback and fix ragged tensor
 static validation CI failures

- Statically validate constant row_lengths, value_rowids, row_starts, and row_limits in RowPartition
- Restrict RaggedTensor static nvals validation to validate=True to respect validate=False in op tests
- Pass validate=False in ragged_squeeze_op when constructing squeezed RaggedTensor with deferred assertions
- Clean up pylint line length and operator stubs directive in ragged_tensor.py and row_partition.py
- Add unit tests for sibling factory static validation and pin expected error message in ragged_tensor_test
---
 .../python/ops/ragged/ragged_squeeze_op.py    |  4 +-
 tensorflow/python/ops/ragged/ragged_tensor.py | 14 +++--
 .../python/ops/ragged/ragged_tensor_test.py   |  3 +-
 tensorflow/python/ops/ragged/row_partition.py | 58 ++++++++++++++++++-
 .../python/ops/ragged/row_partition_test.py   | 35 ++++++++++-
 5 files changed, 103 insertions(+), 11 deletions(-)

diff --git a/tensorflow/python/ops/ragged/ragged_squeeze_op.py b/tensorflow/python/ops/ragged/ragged_squeeze_op.py
index 6d35ffd493e..8755bee8b5a 100644
--- a/tensorflow/python/ops/ragged/ragged_squeeze_op.py
+++ b/tensorflow/python/ops/ragged/ragged_squeeze_op.py
@@ -112,8 +112,8 @@ def squeeze(input: ragged_tensor.Ragged, axis=None, name=None):  # pylint: disab
     if remaining_row_splits and 0 in ragged_dims:
       remaining_row_splits.pop(0)
 
-    squeezed_rt = RaggedTensor.from_nested_row_splits(squeezed_rt,
-                                                      remaining_row_splits)
+    squeezed_rt = RaggedTensor.from_nested_row_splits(
+        squeezed_rt, remaining_row_splits, validate=False)
 
     # Corner case: when removing all the ragged dimensions and the output is
     # a scalar tensor e.g. ragged.squeeze(ragged.constant([[[1]]])).
diff --git a/tensorflow/python/ops/ragged/ragged_tensor.py b/tensorflow/python/ops/ragged/ragged_tensor.py
index 1c9e1a0b3ea..cdd1901eda5 100644
--- a/tensorflow/python/ops/ragged/ragged_tensor.py
+++ b/tensorflow/python/ops/ragged/ragged_tensor.py
@@ -341,15 +341,14 @@ class RaggedTensor(
     values, row_partition = cls._convert_values_and_partition(
         values, row_partition, "partition")
     # pylint: disable=protected-access
-    static_nvals = row_partition._static_nvals_or_constant()
     if row_partition._has_precomputed_value_rowids():
       value_rowids_shape = row_partition.value_rowids().shape
       values.shape[:1].assert_is_compatible_with(value_rowids_shape)
-    elif static_nvals is not None:
-      values.shape[:1].assert_is_compatible_with(
-          tensor_shape.TensorShape([static_nvals]))
-    # pylint: enable=protected-access
     if validate:
+      static_nvals = row_partition._static_nvals_or_constant()
+      if static_nvals is not None:
+        values.shape[:1].assert_is_compatible_with(
+            tensor_shape.TensorShape([static_nvals]))
       msg = "Arguments to _from_row_partition do not form a valid RaggedTensor"
       nvals = _nrows(values, row_partition.dtype)
       checks = [
@@ -360,7 +359,8 @@ class RaggedTensor(
       ]
       if not isinstance(values, RaggedTensor):
         checks.append(check_ops.assert_rank_at_least(values, 1))
-      row_partition = row_partition._with_dependencies(checks)  # pylint: disable=protected-access
+      row_partition = row_partition._with_dependencies(checks)
+    # pylint: enable=protected-access
     return cls(values=values, internal=True, row_partition=row_partition)
 
   @classmethod
@@ -2200,6 +2200,7 @@ class RaggedTensor(
 
     return stub
 
+  # pylint: disable=too-many-function-args
   __getitem__ = _overloaded_operator("__getitem__")
   __ge__ = _overloaded_operator("__ge__")
   __gt__ = _overloaded_operator("__gt__")
@@ -2231,6 +2232,7 @@ class RaggedTensor(
   __truediv__ = _overloaded_operator("__truediv__")
   __rtruediv__ = _overloaded_operator("__rtruediv__")
   del _overloaded_operator
+  # pylint: enable=too-many-function-args
 
   # =============================================================================
   # Name Scope
diff --git a/tensorflow/python/ops/ragged/ragged_tensor_test.py b/tensorflow/python/ops/ragged/ragged_tensor_test.py
index d0466f996fa..d42677bc01c 100644
--- a/tensorflow/python/ops/ragged/ragged_tensor_test.py
+++ b/tensorflow/python/ops/ragged/ragged_tensor_test.py
@@ -336,7 +336,8 @@ class RaggedTensorTest(test_util.TensorFlowTestCase, parameterized.TestCase):
       values = array_ops.reshape(x, [-1, 2])
       return RaggedTensor.from_row_splits(values, [0, 2, 4])
 
-    with self.assertRaises(ValueError):
+    with self.assertRaisesRegex(
+        ValueError, r"Shapes \(2,\) and \(4,\) are incompatible"):
       fn(constant_op.constant([4.0, 5.0, 6.0, 7.0]))
 
   def testFromRowStarts(self):
diff --git a/tensorflow/python/ops/ragged/row_partition.py b/tensorflow/python/ops/ragged/row_partition.py
index de101baa650..d45b6f2f11b 100644
--- a/tensorflow/python/ops/ragged/row_partition.py
+++ b/tensorflow/python/ops/ragged/row_partition.py
@@ -277,6 +277,23 @@ class RowPartition(composite_tensor.CompositeTensor):
       nrows.shape.assert_has_rank(0)
 
       if validate:
+        value_rowids_const = tensor_util.constant_value(value_rowids)
+        if value_rowids_const is not None:
+          if len(value_rowids_const) > 0 and value_rowids_const[0] < 0:
+            raise ValueError(
+                "Arguments to from_value_rowids do not form a valid "
+                "RaggedTensor: value_rowids must be non-negative, "
+                f"got {value_rowids_const[0]}")
+          diff = np.diff(value_rowids_const)
+          if np.any(diff < 0):
+            bad_idx = int(np.argmax(diff < 0))
+            raise ValueError(
+                "Arguments to from_value_rowids do not form a valid "
+                "RaggedTensor: value_rowids must be monotonic increasing, "
+                f"but value_rowids[{bad_idx}] "
+                f"({value_rowids_const[bad_idx]}) > "
+                f"value_rowids[{bad_idx + 1}] "
+                f"({value_rowids_const[bad_idx + 1]})")
         msg = ("Arguments to from_value_rowids do not form a valid "
                "RowPartition")
         checks = [
@@ -433,6 +450,12 @@ class RowPartition(composite_tensor.CompositeTensor):
       row_lengths.shape.assert_has_rank(1)
 
       if validate:
+        row_lengths_const = tensor_util.constant_value(row_lengths)
+        if row_lengths_const is not None:
+          if np.any(row_lengths_const < 0):
+            raise ValueError(
+                "Arguments to from_row_lengths do not form a valid "
+                "RowPartition: row_lengths must be nonnegative.")
         msg = "Arguments to from_row_lengths do not form a valid RowPartition"
         checks = [
             check_ops.assert_rank(row_lengths, 1, message=msg),
@@ -486,6 +509,22 @@ class RowPartition(composite_tensor.CompositeTensor):
       # even though they eventually end up the same type.
       nvals = math_ops.cast(nvals, row_starts.dtype)
       if validate:
+        row_starts_const = tensor_util.constant_value(row_starts)
+        if row_starts_const is not None:
+          if len(row_starts_const) > 0 and row_starts_const[0] != 0:
+            raise ValueError(
+                "Arguments to from_row_starts do not form a valid "
+                f"RaggedTensor: row_starts[0] must be zero, "
+                f"got {row_starts_const[0]}")
+          diff = np.diff(row_starts_const)
+          if np.any(diff < 0):
+            bad_idx = int(np.argmax(diff < 0))
+            raise ValueError(
+                "Arguments to from_row_starts do not form a valid "
+                "RaggedTensor: row_starts must be monotonic increasing, "
+                f"but row_starts[{bad_idx}] ({row_starts_const[bad_idx]}) > "
+                f"row_starts[{bad_idx + 1}] "
+                f"({row_starts_const[bad_idx + 1]})")
         msg = "Arguments to from_row_starts do not form a valid RaggedTensor"
         checks = [
             check_ops.assert_rank(row_starts, 1, message=msg),
@@ -533,6 +572,22 @@ class RowPartition(composite_tensor.CompositeTensor):
       row_limits.shape.assert_has_rank(1)
 
       if validate:
+        row_limits_const = tensor_util.constant_value(row_limits)
+        if row_limits_const is not None:
+          if len(row_limits_const) > 0 and row_limits_const[0] < 0:
+            raise ValueError(
+                "Arguments to from_row_limits do not form a valid "
+                "RaggedTensor: row_limits must be non-negative, "
+                f"got {row_limits_const[0]}")
+          diff = np.diff(row_limits_const)
+          if np.any(diff < 0):
+            bad_idx = int(np.argmax(diff < 0))
+            raise ValueError(
+                "Arguments to from_row_limits do not form a valid "
+                "RaggedTensor: row_limits must be monotonic increasing, "
+                f"but row_limits[{bad_idx}] ({row_limits_const[bad_idx]}) > "
+                f"row_limits[{bad_idx + 1}] "
+                f"({row_limits_const[bad_idx + 1]})")
         msg = "Arguments to from_row_limits do not form a valid RaggedTensor"
         checks = [
             check_ops.assert_rank(row_limits, 1, message=msg),
@@ -624,7 +679,8 @@ class RowPartition(composite_tensor.CompositeTensor):
 
       checks = []
 
-      if const_nvals is None and const_nrows is not None and const_uniform_row_length is not None:
+      if (const_nvals is None and const_nrows is not None and
+          const_uniform_row_length is not None):
         const_nvals = const_nrows * const_uniform_row_length
         if nvals is not None and validate:
           checks.append(check_ops.assert_equal(nvals, const_nvals))
diff --git a/tensorflow/python/ops/ragged/row_partition_test.py b/tensorflow/python/ops/ragged/row_partition_test.py
index 32afffea383..07f38b04914 100644
--- a/tensorflow/python/ops/ragged/row_partition_test.py
+++ b/tensorflow/python/ops/ragged/row_partition_test.py
@@ -922,7 +922,7 @@ class RowPartitionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
 
   def testStaticNvalsFromUniformRowLength(self):
     rp = RowPartition.from_uniform_row_length(3, nrows=4)
-    self.assertAllEqual(12, rp.static_nvals)
+    self.assertEqual(12, rp._static_nvals_or_constant())
 
   def testFromRowSplitsStaticValidationNonZero(self):
     with self.assertRaisesRegex(ValueError, r"row_splits\[0\] must be zero"):
@@ -932,6 +932,39 @@ class RowPartitionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
     with self.assertRaisesRegex(ValueError, r"must be monotonic increasing"):
       RowPartition.from_row_splits([0, 5, 2])
 
+  def testFromValueRowidsStaticValidationNegative(self):
+    with self.assertRaisesRegex(
+        ValueError, r"value_rowids must be non-negative"):
+      RowPartition.from_value_rowids([-1, 0, 1])
+
+  def testFromValueRowidsStaticValidationNonMonotonic(self):
+    with self.assertRaisesRegex(
+        ValueError, r"value_rowids must be monotonic increasing"):
+      RowPartition.from_value_rowids([0, 2, 1])
+
+  def testFromRowLengthsStaticValidationNegative(self):
+    with self.assertRaisesRegex(
+        ValueError, r"row_lengths must be nonnegative"):
+      RowPartition.from_row_lengths([3, -1, 2])
+
+  def testFromRowStartsStaticValidationNonZero(self):
+    with self.assertRaisesRegex(ValueError, r"row_starts\[0\] must be zero"):
+      RowPartition.from_row_starts([1, 2], nvals=3)
+
+  def testFromRowStartsStaticValidationNonMonotonic(self):
+    with self.assertRaisesRegex(
+        ValueError, r"row_starts must be monotonic increasing"):
+      RowPartition.from_row_starts([0, 3, 2], nvals=5)
+
+  def testFromRowLimitsStaticValidationNegative(self):
+    with self.assertRaisesRegex(ValueError, r"row_limits must be non-negative"):
+      RowPartition.from_row_limits([-1, 2])
+
+  def testFromRowLimitsStaticValidationNonMonotonic(self):
+    with self.assertRaisesRegex(
+        ValueError, r"row_limits must be monotonic increasing"):
+      RowPartition.from_row_limits([3, 1])
+
 
 @test_util.run_all_in_graph_and_eager_modes
 class RowPartitionSpecTest(test_util.TensorFlowTestCase,

From 92cd903123a14a2a0e9851e2f3dec3c425fb9d83 Mon Sep 17 00:00:00 2001
From: Boris Kozak <boriskozak@gmail.com>
Date: Tue, 6 Oct 2026 12:05:21 -0400
Subject: [PATCH 0217/1000] test: cover float64 and complex64 in
 testNonFiniteXPropagation

---
 .../kernel_tests/math_ops/cwise_ops_test.py   | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py b/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
index 8c2772deb64..a8cc65fe3b3 100644
--- a/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
@@ -1343,15 +1343,16 @@ class PolyvalTest(test.TestCase):
   def testNonFiniteXPropagation(self):
     # Non-finite x must propagate to NaN, matching numpy.polyval, which
     # starts its Horner accumulator at zero (0 * inf + c -> nan).
-    for x_val in [float("inf"), float("-inf"), float("nan")]:
-      with self.subTest(x=x_val):
-        x = np.array([1.0, x_val, 3.0], dtype=np.float32)
-        coeffs = [np.float32(1.0), np.float32(1.0)]
-        with np.errstate(invalid="ignore"):
-          np_val = np.polyval(coeffs, x)
-        with self.cached_session():
-          tf_val = math_ops.polyval(coeffs, x)
-          self.assertAllClose(np_val, self.evaluate(tf_val))
+    for dtype in [np.float32, np.float64, np.complex64]:
+      for x_val in [float("inf"), float("-inf"), float("nan")]:
+        with self.subTest(dtype=dtype, x=x_val):
+          x = np.array([1.0, x_val, 3.0], dtype=dtype)
+          coeffs = [dtype(1.0), dtype(1.0)]
+          with np.errstate(invalid="ignore"):
+            np_val = np.polyval(coeffs, x)
+          with self.cached_session():
+            tf_val = math_ops.polyval(coeffs, x)
+            self.assertAllClose(np_val, self.evaluate(tf_val))
 
   def test_coeffs_raise(self):
     x = np.random.rand(2, 2).astype(np.float32)

From fe1972f6794af1e9636792299f729cbb27218140 Mon Sep 17 00:00:00 2001
From: Som Tripathi <somtri@iastate.edu>
Date: Tue, 6 Oct 2026 15:06:35 -0500
Subject: [PATCH 0218/1000] Check AvgPoolGrad grad shape on the ROCm path; drop
 the dead loop

AvgPoolingGradOpCustomGPUKernel calls RunAvePoolBackwardNHWC without
checking the grad shape when CUDNN_VERSION is below 7300, which
includes ROCm builds. Add the same check as the CPU kernel there.

With the grad shape check in the CPU kernel, the input_max loop can no
longer fail, so remove it. testAvgPoolGradOutputMemoryOutOfBounds keeps
both messages, because the oneDNN kernel still raises the old one; its
comments now say which kernels raise which message.

Run testAvgPoolGradMismatchedGradShapeRaisesError on the default
device instead of pinning it to the CPU.
---
 tensorflow/core/kernels/avgpooling_op.cc      | 36 +++++++------------
 .../kernel_tests/nn_ops/pooling_ops_test.py   | 27 +++++++-------
 2 files changed, 26 insertions(+), 37 deletions(-)

diff --git a/tensorflow/core/kernels/avgpooling_op.cc b/tensorflow/core/kernels/avgpooling_op.cc
index a39a8194126..95d378239b6 100644
--- a/tensorflow/core/kernels/avgpooling_op.cc
+++ b/tensorflow/core/kernels/avgpooling_op.cc
@@ -355,29 +355,6 @@ class AvgPoolingGradOp : public OpKernel {
     const T* out_backprop_ptr = out_backprop.flat<T>().data();
     T* input_backprop_ptr = output->flat<T>().data();
 
-    for (int64_t r = 0; r < out_backprop_rows; ++r) {
-      int rindex, rsize;
-      OP_REQUIRES_OK(context,
-                     GetBroadcastSize(r, in_rows, window_rows, row_stride,
-                                      pad_rows, &rindex, &rsize));
-      for (int64_t c = 0; c < out_backprop_cols; ++c) {
-        int cindex, csize;
-        OP_REQUIRES_OK(context,
-                       GetBroadcastSize(c, in_cols, window_cols, col_stride,
-                                        pad_cols, &cindex, &csize));
-        int64_t input_max =
-            ((out_backprop_batch - 1) * in_rows + rindex + rsize - 1) *
-                in_cols +
-            cindex + csize - 1;
-        OP_REQUIRES(context, input_max < output->NumElements(),
-                    absl::InvalidArgumentError(absl::StrCat(
-                        "Output only has ", output->NumElements(),
-                        " elements but computation requested would "
-                        "use element with index=",
-                        input_max)));
-      }
-    }
-
     auto shard = [context, out_backprop_ptr, input_backprop_ptr,
                   out_backprop_rows, out_backprop_cols, out_backprop_depth,
                   in_rows, in_cols, window_rows, window_cols, row_stride,
@@ -637,6 +614,19 @@ class AvgPoolingGradOpCustomGPUKernel : public OpKernel {
                                   in_cols, window_cols, /*dilation_rate=*/1,
                                   col_stride, padding_, &out_width, &pad_cols));
 
+      TensorShape forward_output_shape;
+      OP_REQUIRES_OK(context, ShapeFromFormatWithStatus(
+                                  data_format_,
+                                  GetTensorDim(output_shape, data_format_, 'N'),
+                                  out_height, out_width,
+                                  GetTensorDim(output_shape, data_format_, 'C'),
+                                  &forward_output_shape));
+      OP_REQUIRES(
+          context, out_backprop.shape() == forward_output_shape,
+          absl::InvalidArgumentError(absl::StrCat(
+              "Expected grad shape to be ", forward_output_shape.DebugString(),
+              ", but got ", out_backprop.shape().DebugString())));
+
       OP_REQUIRES_OK(
           context,
           RunAvePoolBackwardNHWC<T>(out_backprop.flat<T>().data(),  // top_diff
diff --git a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
index 49df55bb70b..846aa6bd085 100644
--- a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
@@ -2403,11 +2403,11 @@ class PoolingTest(test.TestCase, parameterized.TestCase):
     with self.assertRaisesRegex(
         errors_impl.InvalidArgumentError,
         (
-            # CPU error message
+            # oneDNN error message
             "(Output only has 3 elements but computation requested would use"
             " element with index=6"
             ")|("
-            # GPU error message
+            # CPU and GPU error message
             r"Expected grad shape to be \[1,1,3,1\], but got \[3,1,3,1\])"
         ),
     ):
@@ -2438,18 +2438,17 @@ class PoolingTest(test.TestCase, parameterized.TestCase):
          r"\[1,14,14,6\]"),
         ([2, 2, 2, 2], [2, 3, 3, 3], 1, r"\[2,2,2,2\], but got \[2,3,3,3\]"),
     ):
-      with ops.device("/cpu:0"):
-        with self.assertRaisesRegex(
-            (errors_impl.InvalidArgumentError, ValueError),
-            "Expected grad shape to be " + expected):
-          self.evaluate(
-              gen_nn_ops.AvgPoolGrad(
-                  orig_input_shape=orig_input_shape,
-                  grad=array_ops.zeros(grad_shape),
-                  ksize=[1, ksize, ksize, 1],
-                  strides=[1, ksize, ksize, 1],
-                  padding="VALID",
-                  data_format="NHWC"))
+      with self.assertRaisesRegex(
+          (errors_impl.InvalidArgumentError, ValueError),
+          "Expected grad shape to be " + expected):
+        self.evaluate(
+            gen_nn_ops.AvgPoolGrad(
+                orig_input_shape=orig_input_shape,
+                grad=array_ops.zeros(grad_shape),
+                ksize=[1, ksize, ksize, 1],
+                strides=[1, ksize, ksize, 1],
+                padding="VALID",
+                data_format="NHWC"))
 
   @test_util.run_deprecated_v1
   def testShapeFunctionEdgeCases(self):

From 48d748979e271b487ad3e6f253d965756dac1af0 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Tue, 6 Oct 2026 23:52:33 +0000
Subject: [PATCH 0219/1000] Bump urllib3 in
 /ci/official/requirements_updater/numpy1_requirements

Bumps [urllib3](https://github.com/urllib3/urllib3) from 2.7.0 to 2.8.0.
- [Release notes](https://github.com/urllib3/urllib3/releases)
- [Changelog](https://github.com/urllib3/urllib3/blob/main/CHANGES.rst)
- [Commits](https://github.com/urllib3/urllib3/compare/2.7.0...2.8.0)

---
updated-dependencies:
- dependency-name: urllib3
  dependency-version: 2.8.0
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
---
 .../numpy1_requirements/requirements_lock_3_10.txt          | 6 +++---
 .../numpy1_requirements/requirements_lock_3_11.txt          | 6 +++---
 .../numpy1_requirements/requirements_lock_3_12.txt          | 6 +++---
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt
index 91abdd8adf1..00968ad4260 100644
--- a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt
+++ b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt
@@ -764,9 +764,9 @@ typing-extensions==4.14.1 \
     #   -r ci/official/requirements_updater/requirements.in
     #   optree
     #   rich
-urllib3==2.7.0 \
-    --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \
-    --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897
+urllib3==2.8.0 \
+    --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
+    --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
 werkzeug==3.1.6 \
     --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
diff --git a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt
index 7a07aa010ad..5f834b529ae 100644
--- a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt
+++ b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt
@@ -763,9 +763,9 @@ typing-extensions==4.14.1 \
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   optree
-urllib3==2.7.0 \
-    --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \
-    --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897
+urllib3==2.8.0 \
+    --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
+    --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
 werkzeug==3.1.6 \
     --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
diff --git a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt
index d3cb6403cb4..6b49c8c337c 100644
--- a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt
+++ b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt
@@ -763,9 +763,9 @@ typing-extensions==4.14.1 \
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   optree
-urllib3==2.7.0 \
-    --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \
-    --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897
+urllib3==2.8.0 \
+    --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
+    --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
 werkzeug==3.1.6 \
     --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \

From 30b9076b18401c473fd30a2729a42af0ffcd1b98 Mon Sep 17 00:00:00 2001
From: Som Tripathi <somtri@iastate.edu>
Date: Tue, 6 Oct 2026 20:03:55 -0500
Subject: [PATCH 0220/1000] Move the MaxPool format check above DCHECK(result)

The early return for unsupported data formats came after
DCHECK(result), so in a debug build an unparsable data_format would
abort at the DCHECK before reaching it. Move the check above the
DCHECK, as suggested in review.

The MaxPool op def allows only NHWC, NCHW and NCHW_VECT_C, and
FormatFromString parses all three, so valid graphs behave the same.
---
 tensorflow/core/common_runtime/mkl_layout_pass.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/common_runtime/mkl_layout_pass.cc b/tensorflow/core/common_runtime/mkl_layout_pass.cc
index 61367b3e7b3..468c5d7f31d 100644
--- a/tensorflow/core/common_runtime/mkl_layout_pass.cc
+++ b/tensorflow/core/common_runtime/mkl_layout_pass.cc
@@ -1598,11 +1598,11 @@ class MklLayoutRewritePass : public GraphOptimizationPass {
     TF_CHECK_OK(GetNodeAttr(n->def(), "strides", &strides));
     TF_CHECK_OK(GetNodeAttr(n->def(), "data_format", &data_format_str));
     bool result = FormatFromString(data_format_str, &data_format);
-    DCHECK(result);
     // The oneDNN MaxPool ops accept only NHWC and NCHW (NDHWC and NCDHW in 3D).
     if (!result || (data_format != FORMAT_NHWC && data_format != FORMAT_NCHW)) {
       return false;
     }
+    DCHECK(result);
 
     // Condition that specifies non-batch-wise and non-depth-wise pooling.
     if (GetTensorDim(ksize, data_format, 'N') == 1 &&

From def11a5c9ae70332f9e4e370831f8bb14466c106 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Wed, 7 Oct 2026 06:55:32 +0530
Subject: [PATCH 0221/1000] Update windows_lib_diagnostics.py

---
 .../python/platform/windows_lib_diagnostics.py     | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/platform/windows_lib_diagnostics.py b/tensorflow/python/platform/windows_lib_diagnostics.py
index cb5e211e5a3..f2b800668be 100644
--- a/tensorflow/python/platform/windows_lib_diagnostics.py
+++ b/tensorflow/python/platform/windows_lib_diagnostics.py
@@ -34,9 +34,13 @@ def _safe_print(text, end="\n"):
   paths (e.g. under a non-ASCII user directory) may not be encodable on a
   console that uses a legacy code page.
   """
-  encoding = getattr(sys.stdout, "encoding", None) or "ascii"
-  print(text.encode(encoding, errors="backslashreplace").decode(encoding),
-        end=end)
+  try:
+    encoding = getattr(sys.stdout, "encoding", None) or "ascii"
+    print(text.encode(encoding, errors="backslashreplace").decode(encoding),
+          end=end)
+  except Exception:  # pylint: disable=broad-exception-caught
+    # A diagnostic message must never mask the original import error.
+    pass
 
 
 def get_dll_dependencies(path):
@@ -218,7 +222,9 @@ def run_diagnosis(path=None):
       diagnose_dll_load(path)
     else:
       _safe_print(f"Error: Path does not exist or not on Windows: {path}")
-  except OSError as e:
+  except Exception as e:  # pylint: disable=broad-exception-caught
+    # This runs while handling an ImportError, so no diagnostic failure may
+    # escape and replace the original error.
     _safe_print(f"Diagnostic failed: {e}")
 
 

From bfde5ff8ff0ff7fa6b4157414e71f4185509c758 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Wed, 7 Oct 2026 06:57:34 +0530
Subject: [PATCH 0222/1000] Add tests for windows_lib_diagnostics functionality

---
 .../platform/windows_lib_diagnostics_pr.py    | 68 +++++++++++++++++++
 1 file changed, 68 insertions(+)
 create mode 100644 tensorflow/python/platform/windows_lib_diagnostics_pr.py

diff --git a/tensorflow/python/platform/windows_lib_diagnostics_pr.py b/tensorflow/python/platform/windows_lib_diagnostics_pr.py
new file mode 100644
index 00000000000..58d66f5e317
--- /dev/null
+++ b/tensorflow/python/platform/windows_lib_diagnostics_pr.py
@@ -0,0 +1,68 @@
+# 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 windows_lib_diagnostics."""
+
+import io
+import os
+from unittest import mock
+
+from tensorflow.python.platform import test
+from tensorflow.python.platform import windows_lib_diagnostics
+
+
+class _FakeStdout(io.StringIO):
+
+  def __init__(self, encoding):
+    super().__init__()
+    self._encoding = encoding
+
+  @property
+  def encoding(self):
+    return self._encoding
+
+
+class WindowsLibDiagnosticsTest(test.TestCase):
+
+  def testSafePrintWithAsciiConsole(self):
+    out = _FakeStdout("ascii")
+    with mock.patch("sys.stdout", out):
+      windows_lib_diagnostics._safe_print("C:\\Users\\Jos\u00e9\\x.dll")
+    self.assertIn("Jos\\xe9", out.getvalue())
+
+  def testSafePrintWithNoneEncoding(self):
+    out = _FakeStdout(None)
+    with mock.patch("sys.stdout", out):
+      windows_lib_diagnostics._safe_print("caf\u00e9")
+    self.assertIn("caf\\xe9", out.getvalue())
+
+  def testSafePrintNeverRaises(self):
+    with mock.patch("builtins.print", side_effect=ValueError("boom")):
+      windows_lib_diagnostics._safe_print("text")
+
+  def testUnsetPathInEnviron(self):
+    with mock.patch.dict(os.environ, {}, clear=True):
+      windows_lib_diagnostics.run_diagnosis("nonexistent.dll")
+
+  def testRunDiagnosisSwallowsUnexpectedErrors(self):
+    out = _FakeStdout("utf-8")
+    with mock.patch("sys.stdout", out), mock.patch.object(
+        windows_lib_diagnostics.os.path, "abspath",
+        side_effect=RuntimeError("unexpected")):
+      windows_lib_diagnostics.run_diagnosis(None)
+    self.assertIn("Diagnostic failed: unexpected", out.getvalue())
+
+
+if __name__ == "__main__":
+  test.main()

From 6ecd160e60a153182bb9bfdb28997266d42ec5f3 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Wed, 7 Oct 2026 07:05:36 +0530
Subject: [PATCH 0223/1000] Add strict test for windows_lib_diagnostics

---
 tensorflow/python/platform/BUILD | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/tensorflow/python/platform/BUILD b/tensorflow/python/platform/BUILD
index ae82f2f3dc9..25ddc42b05a 100644
--- a/tensorflow/python/platform/BUILD
+++ b/tensorflow/python/platform/BUILD
@@ -54,6 +54,20 @@ py_library(
     srcs = ["windows_lib_diagnostics.py"],
 )
 
+tf_py_strict_test(
+    name = "windows_lib_diagnostics_test",
+    size = "small",
+    srcs = ["windows_lib_diagnostics_test.py"],
+    main = "windows_lib_diagnostics_test.py",
+    tags = [
+        "no_pip",
+    ],
+    deps = [
+        ":client_testlib",
+        ":windows_lib_diagnostics",
+    ],
+)
+
 py_library(
     name = "self_check",
     srcs = ["self_check.py"],

From e592b7191615a5eec7ab157c6041afdb3c4452e6 Mon Sep 17 00:00:00 2001
From: Som Tripathi <somtri@iastate.edu>
Date: Tue, 6 Oct 2026 22:09:15 -0500
Subject: [PATCH 0224/1000] Check AvgPoolGrad grad shape in the oneDNN kernel

MklAvgPoolingGradOp now computes its pool parameters before the
empty-tensor return, builds the expected grad shape from them for 2-D
and 3-D pooling, and returns InvalidArgumentError when grad has a
different shape. This replaces its 2-D input_max loop, which did not
check depth. The CPU and GPU kernels already raise the same error.

The mismatched-shape test no longer skips oneDNN,
testAvgPoolGradOutputMemoryOutOfBounds expects only the grad shape
message, and the tensor_in_shape check in the three AvgPoolGrad kernels
now names tensor_in_shape in its message.
---
 tensorflow/core/kernels/avgpooling_op.cc      | 15 ++--
 .../core/kernels/mkl/mkl_avgpooling_op.cc     | 69 ++++++-------------
 .../kernel_tests/nn_ops/pooling_ops_test.py   | 11 +--
 3 files changed, 32 insertions(+), 63 deletions(-)

diff --git a/tensorflow/core/kernels/avgpooling_op.cc b/tensorflow/core/kernels/avgpooling_op.cc
index 95d378239b6..67cda10bc76 100644
--- a/tensorflow/core/kernels/avgpooling_op.cc
+++ b/tensorflow/core/kernels/avgpooling_op.cc
@@ -291,8 +291,9 @@ class AvgPoolingGradOp : public OpKernel {
     OP_REQUIRES(
         context,
         tensor_in_shape.dims() == 1 && tensor_in_shape.NumElements() == 4,
-        absl::InvalidArgumentError("out_backprop must be 1-dimensional and 4 "
-                                   "elements"));
+        absl::InvalidArgumentError(
+            "tensor_in_shape must be 1-dimensional and 4 "
+            "elements"));
     // For avgpooling, out_backprop should have 4 dimensions.
     OP_REQUIRES(
         context, out_backprop.dims() == 4,
@@ -465,8 +466,9 @@ class AvgPoolingGradOp<GPUDevice, T> : public OpKernel {
     OP_REQUIRES(
         context,
         tensor_in_shape.dims() == 1 && tensor_in_shape.NumElements() == 4,
-        absl::InvalidArgumentError("out_backprop must be 1-dimensional and 4 "
-                                   "elements"));
+        absl::InvalidArgumentError(
+            "tensor_in_shape must be 1-dimensional and 4 "
+            "elements"));
     // For avgpooling, out_backprop should have 4 dimensions.
     OP_REQUIRES(
         context, out_backprop.dims() == 4,
@@ -553,8 +555,9 @@ class AvgPoolingGradOpCustomGPUKernel : public OpKernel {
     OP_REQUIRES(
         context,
         tensor_in_shape.dims() == 1 && tensor_in_shape.NumElements() == 4,
-        absl::InvalidArgumentError("out_backprop must be 1-dimensional and 4 "
-                                   "elements"));
+        absl::InvalidArgumentError(
+            "tensor_in_shape must be 1-dimensional and 4 "
+            "elements"));
     // For avgpooling, out_backprop should have 4 dimensions.
     OP_REQUIRES(
         context, out_backprop.dims() == 4,
diff --git a/tensorflow/core/kernels/mkl/mkl_avgpooling_op.cc b/tensorflow/core/kernels/mkl/mkl_avgpooling_op.cc
index 6a94b751969..d7ad9560bb6 100644
--- a/tensorflow/core/kernels/mkl/mkl_avgpooling_op.cc
+++ b/tensorflow/core/kernels/mkl/mkl_avgpooling_op.cc
@@ -279,61 +279,36 @@ class MklAvgPoolingGradOp : public MklPoolingBackwardOpBase<T> {
                      context->allocate_output(0, output_shape, &output_tensor));
       output_tensor->flat<T>().setZero();
 
-      // out-of-memory boundary index check for output_tensor in 2D case.
-      const int depth_window = this->ksize_[3];
-      if (is_pool2d && depth_window == 1) {
-        const int window_rows = this->ksize_[1];
-        const int window_cols = this->ksize_[2];
-        const int row_stride = this->stride_[1];
-        const int col_stride = this->stride_[2];
-        const int64_t in_rows = output_shape.dim_size(1);
-        const int64_t in_cols = output_shape.dim_size(2);
-        const int64_t out_backprop_batch = grad_tensor.dim_size(0);
-        const int64_t out_backprop_rows = grad_tensor.dim_size(1);
-        const int64_t out_backprop_cols = grad_tensor.dim_size(2);
-        int64_t out_height, out_width, pad_rows, pad_cols;
-        OP_REQUIRES_OK(
-            context, GetWindowedOutputSize(
-                         in_rows, window_rows, /*dilation_rate=*/1, row_stride,
-                         this->padding_, &out_height, &pad_rows));
+      MklPoolParameters pool_params;
+      this->InitMklPoolParameters(context, &pool_params, orig_input_mkl_shape,
+                                  output_shape);
 
-        OP_REQUIRES_OK(
-            context, GetWindowedOutputSize(
-                         in_cols, window_cols, /*dilation_rate=*/1, col_stride,
-                         this->padding_, &out_width, &pad_cols));
-
-        for (int64_t r = 0; r < out_backprop_rows; ++r) {
-          int rindex, rsize;
-          OP_REQUIRES_OK(context,
-                         GetBroadcastSize(r, in_rows, window_rows, row_stride,
-                                          pad_rows, &rindex, &rsize));
-          for (int64_t c = 0; c < out_backprop_cols; ++c) {
-            int cindex, csize;
-            OP_REQUIRES_OK(context,
-                           GetBroadcastSize(c, in_cols, window_cols, col_stride,
-                                            pad_cols, &cindex, &csize));
-            int64_t input_max =
-                ((out_backprop_batch - 1) * in_rows + rindex + rsize - 1) *
-                    in_cols +
-                cindex + csize - 1;
-            OP_REQUIRES(context, input_max < output_tensor->NumElements(),
-                        absl::InvalidArgumentError(absl::StrCat(
-                            "Output only has ", output_tensor->NumElements(),
-                            " elements but computation requested"
-                            " would use element with index=",
-                            input_max)));
-          }
-        }
+      TensorShape forward_output_shape;
+      if (is_pool2d) {
+        OP_REQUIRES_OK(context,
+                       ShapeFromFormatWithStatus(
+                           this->data_format_tf_, pool_params.tensor_in_batch,
+                           pool_params.out_height, pool_params.out_width,
+                           pool_params.out_depth, &forward_output_shape));
+      } else {
+        OP_REQUIRES_OK(context,
+                       ShapeFromFormatWithStatus(
+                           this->data_format_tf_, pool_params.tensor_in_batch,
+                           {pool_params.out_planes, pool_params.out_height,
+                            pool_params.out_width},
+                           pool_params.out_depth, &forward_output_shape));
       }
+      OP_REQUIRES(
+          context, grad_tensor.shape() == forward_output_shape,
+          absl::InvalidArgumentError(absl::StrCat(
+              "Expected grad shape to be ", forward_output_shape.DebugString(),
+              ", but got ", grad_tensor.shape().DebugString())));
 
       if (output_shape.num_elements() == 0 || grad_tensor.NumElements() == 0) {
         return;
       }
       // Used to allocate output_diff_src/diff_src.
       MklDnnData<T> grad_dnn_data(&cpu_engine_);
-      MklPoolParameters pool_params;
-      this->InitMklPoolParameters(context, &pool_params, orig_input_mkl_shape,
-                                  output_shape);
 
       memory::dims filter_dims, strides, padding_left, padding_right;
 #ifndef ENABLE_ONEDNN_V3
diff --git a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
index 846aa6bd085..12da8cf168d 100644
--- a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
@@ -2402,14 +2402,7 @@ class PoolingTest(test.TestCase, parameterized.TestCase):
   def testAvgPoolGradOutputMemoryOutOfBounds(self):
     with self.assertRaisesRegex(
         errors_impl.InvalidArgumentError,
-        (
-            # oneDNN error message
-            "(Output only has 3 elements but computation requested would use"
-            " element with index=6"
-            ")|("
-            # CPU and GPU error message
-            r"Expected grad shape to be \[1,1,3,1\], but got \[3,1,3,1\])"
-        ),
+        r"Expected grad shape to be \[1,1,3,1\], but got \[3,1,3,1\]"
     ):
       self.evaluate(
           gen_nn_ops.AvgPoolGrad(
@@ -2429,8 +2422,6 @@ class PoolingTest(test.TestCase, parameterized.TestCase):
   @test_util.run_in_graph_and_eager_modes
   @test_util.disable_xla("Xla does not raise error on out of bounds access")
   def testAvgPoolGradMismatchedGradShapeRaisesError(self):
-    if test_util.IsMklEnabled():
-      self.skipTest("The oneDNN AvgPoolGrad kernel does not check grad's shape")
     # Each grad has more channels than the input. The CPU kernel used to write
     # past the end of its output buffer instead of raising an error.
     for orig_input_shape, grad_shape, ksize, expected in (

From 28482703bdf5167a76e4c28f3882e9551a0b8807 Mon Sep 17 00:00:00 2001
From: Som Tripathi <somtri@iastate.edu>
Date: Tue, 6 Oct 2026 22:09:16 -0500
Subject: [PATCH 0225/1000] Remove the redundant DCHECK(result) in the MaxPool
 rewrite check

The early return above it already handles a format string that fails to
parse, so result is always true at the DCHECK.
---
 tensorflow/core/common_runtime/mkl_layout_pass.cc | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorflow/core/common_runtime/mkl_layout_pass.cc b/tensorflow/core/common_runtime/mkl_layout_pass.cc
index 468c5d7f31d..6725a8f6759 100644
--- a/tensorflow/core/common_runtime/mkl_layout_pass.cc
+++ b/tensorflow/core/common_runtime/mkl_layout_pass.cc
@@ -1602,7 +1602,6 @@ class MklLayoutRewritePass : public GraphOptimizationPass {
     if (!result || (data_format != FORMAT_NHWC && data_format != FORMAT_NCHW)) {
       return false;
     }
-    DCHECK(result);
 
     // Condition that specifies non-batch-wise and non-depth-wise pooling.
     if (GetTensorDim(ksize, data_format, 'N') == 1 &&

From 3c42e7c2cca5294bb253a056cb500509f78cc41b Mon Sep 17 00:00:00 2001
From: Som Tripathi <somtri@iastate.edu>
Date: Wed, 7 Oct 2026 00:34:42 -0500
Subject: [PATCH 0226/1000] Disable XLA for the NCHW_VECT_C MaxPool error test

The test asserts the CPU MaxPoolingOp error message, so exclude it from
the XLA test variants.
---
 tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
index fcae1e9d865..a5c7643d60e 100644
--- a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
@@ -819,6 +819,8 @@ class PoolingTest(test.TestCase, parameterized.TestCase):
         self.evaluate(t)
 
   @test_util.run_in_graph_and_eager_modes
+  @test_util.disable_xla(
+      "XLA handles NCHW_VECT_C without CPU MaxPoolingOp error")
   def testMaxPoolNchwVectCOnCpuRaisesError(self):
     # With oneDNN on, this MaxPool used to be rewritten to the oneDNN op, which
     # rejects NCHW_VECT_C, and the process aborted.

From 98074a844c9b6a10609f93b6af3273b72aff8f50 Mon Sep 17 00:00:00 2001
From: Som Tripathi <somtri@iastate.edu>
Date: Wed, 7 Oct 2026 01:09:43 -0500
Subject: [PATCH 0227/1000] Check AvgPool3DGrad grad shape; keep oneDNN
 empty-tensor exit first

MklAvgPoolingGradOp again returns on an empty input or grad before it
computes its pool parameters, so like the CPU kernel it returns before
the grad shape check. With the previous order, an empty input whose grad
had a different shape failed the new shape check instead of returning
zeros.

AvgPooling3dGradOp now builds the forward output shape from
Get3dOutputSizeV2 and returns InvalidArgumentError when out_backprop
has a different shape. Before, it checked only the batch and depth
dimensions.

Adds testAvgPoolGradEmptyTensorFastExit and
testAvgPool3dGradMismatchedGradShapeRaisesError.
---
 .../core/kernels/mkl/mkl_avgpooling_op.cc     |  7 +++---
 tensorflow/core/kernels/pooling_ops_3d.cc     | 13 +++++++++++
 .../nn_ops/pooling_ops_3d_test.py             | 23 +++++++++++++++++++
 .../kernel_tests/nn_ops/pooling_ops_test.py   | 12 ++++++++++
 4 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/kernels/mkl/mkl_avgpooling_op.cc b/tensorflow/core/kernels/mkl/mkl_avgpooling_op.cc
index d7ad9560bb6..88e78ea2f9d 100644
--- a/tensorflow/core/kernels/mkl/mkl_avgpooling_op.cc
+++ b/tensorflow/core/kernels/mkl/mkl_avgpooling_op.cc
@@ -279,6 +279,10 @@ class MklAvgPoolingGradOp : public MklPoolingBackwardOpBase<T> {
                      context->allocate_output(0, output_shape, &output_tensor));
       output_tensor->flat<T>().setZero();
 
+      if (output_shape.num_elements() == 0 || grad_tensor.NumElements() == 0) {
+        return;
+      }
+
       MklPoolParameters pool_params;
       this->InitMklPoolParameters(context, &pool_params, orig_input_mkl_shape,
                                   output_shape);
@@ -304,9 +308,6 @@ class MklAvgPoolingGradOp : public MklPoolingBackwardOpBase<T> {
               "Expected grad shape to be ", forward_output_shape.DebugString(),
               ", but got ", grad_tensor.shape().DebugString())));
 
-      if (output_shape.num_elements() == 0 || grad_tensor.NumElements() == 0) {
-        return;
-      }
       // Used to allocate output_diff_src/diff_src.
       MklDnnData<T> grad_dnn_data(&cpu_engine_);
 
diff --git a/tensorflow/core/kernels/pooling_ops_3d.cc b/tensorflow/core/kernels/pooling_ops_3d.cc
index 2ea96a1de98..f1c26a639b6 100644
--- a/tensorflow/core/kernels/pooling_ops_3d.cc
+++ b/tensorflow/core/kernels/pooling_ops_3d.cc
@@ -564,6 +564,19 @@ class AvgPooling3dGradOp : public OpKernel {
                    Get3dOutputSizeV2(input_size, window, dilations, stride,
                                      padding_, &out, &padding));
 
+    TensorShape forward_output_shape;
+    OP_REQUIRES_OK(context, ShapeFromFormatWithStatus(
+                                data_format_,
+                                GetTensorDim(output_shape, data_format_, 'N'),
+                                {{out[2], out[1], out[0]}},
+                                GetTensorDim(output_shape, data_format_, 'C'),
+                                &forward_output_shape));
+    OP_REQUIRES(
+        context, out_backprop.shape() == forward_output_shape,
+        absl::InvalidArgumentError(absl::StrCat(
+            "Expected grad shape to be ", forward_output_shape.DebugString(),
+            ", but got ", out_backprop.shape().DebugString())));
+
     LaunchAvgPooling3dGradOp<Device, T>::launch(
         context, output_shape, out_backprop, window, stride, out, padding,
         data_format_, output);
diff --git a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_3d_test.py b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_3d_test.py
index f8416bc016a..1c50eb864be 100644
--- a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_3d_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_3d_test.py
@@ -155,6 +155,29 @@ class PoolingTest(test.TestCase):
           )
           self.evaluate(t)
 
+  @test_util.run_in_graph_and_eager_modes
+  @test_util.disable_xla("Xla does not raise error on out of bounds access")
+  def testAvgPool3dGradMismatchedGradShapeRaisesError(self):
+    for orig_input_shape, grad_shape, ksize, expected in (
+        ([1, 4, 28, 28, 3], [1, 2, 14, 14, 6], 2, r"\[1,2,14,14,3\], but got "
+         r"\[1,2,14,14,6\]"),
+        ([2, 2, 2, 2, 2], [2, 3, 3, 3, 3], 1, r"\[2,2,2,2,2\], but got "
+         r"\[2,3,3,3,3\]"),
+        ([1, 2, 2, 2, 1], [1, 3, 3, 3, 1], 1, r"\[1,2,2,2,1\], but got "
+         r"\[1,3,3,3,1\]"),
+    ):
+      with self.assertRaisesRegex(
+          (errors_impl.InvalidArgumentError, ValueError),
+          "Expected grad shape to be " + expected):
+        self.evaluate(
+            gen_nn_ops.AvgPool3DGrad(
+                orig_input_shape=orig_input_shape,
+                grad=array_ops.zeros(grad_shape),
+                ksize=[1, ksize, ksize, ksize, 1],
+                strides=[1, ksize, ksize, ksize, 1],
+                padding="VALID",
+                data_format="NDHWC"))
+
   def testAvgPool3dGradEmptyInput(self):
     for data_format, use_gpu in GetTestConfigs():
       with self.cached_session(use_gpu=use_gpu):
diff --git a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
index 12da8cf168d..e3d8b267bbb 100644
--- a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
@@ -2441,6 +2441,18 @@ class PoolingTest(test.TestCase, parameterized.TestCase):
                 padding="VALID",
                 data_format="NHWC"))
 
+  @test_util.run_in_graph_and_eager_modes
+  def testAvgPoolGradEmptyTensorFastExit(self):
+    for empty_shape in ([0, 10, 10, 3], [1, 0, 10, 3], [1, 10, 10, 0]):
+      self.evaluate(
+          gen_nn_ops.AvgPoolGrad(
+              orig_input_shape=empty_shape,
+              grad=array_ops.zeros(empty_shape),
+              ksize=[1, 2, 2, 1],
+              strides=[1, 2, 2, 1],
+              padding="VALID",
+              data_format="NHWC"))
+
   @test_util.run_deprecated_v1
   def testShapeFunctionEdgeCases(self):
     # All shapes unknown.

From 00ade7b6589790e1bd720c000b034e0a78339ad8 Mon Sep 17 00:00:00 2001
From: Som Tripathi <somtri@iastate.edu>
Date: Wed, 7 Oct 2026 03:08:00 -0500
Subject: [PATCH 0228/1000] Run the oneDNN AvgPoolGrad shape check for empty
 grads

MklAvgPoolingGradOp now returns early only when the output is empty. An
empty grad for a non-empty output reaches the grad shape check and
returns InvalidArgumentError instead of zeros.

testAvgPoolGradMismatchedGradShapeRaisesError gets an empty-grad case.
testAvgPoolGradEmptyTensorFastExit now passes grads with the forward
output shape, and adds a non-empty input whose forward output is empty.
---
 .../core/kernels/mkl/mkl_avgpooling_op.cc       |  2 +-
 .../kernel_tests/nn_ops/pooling_ops_test.py     | 17 ++++++++++++-----
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/tensorflow/core/kernels/mkl/mkl_avgpooling_op.cc b/tensorflow/core/kernels/mkl/mkl_avgpooling_op.cc
index 88e78ea2f9d..b5f5a00b194 100644
--- a/tensorflow/core/kernels/mkl/mkl_avgpooling_op.cc
+++ b/tensorflow/core/kernels/mkl/mkl_avgpooling_op.cc
@@ -279,7 +279,7 @@ class MklAvgPoolingGradOp : public MklPoolingBackwardOpBase<T> {
                      context->allocate_output(0, output_shape, &output_tensor));
       output_tensor->flat<T>().setZero();
 
-      if (output_shape.num_elements() == 0 || grad_tensor.NumElements() == 0) {
+      if (output_shape.num_elements() == 0) {
         return;
       }
 
diff --git a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
index e3d8b267bbb..6628e21d7a2 100644
--- a/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/pooling_ops_test.py
@@ -2422,12 +2422,14 @@ class PoolingTest(test.TestCase, parameterized.TestCase):
   @test_util.run_in_graph_and_eager_modes
   @test_util.disable_xla("Xla does not raise error on out of bounds access")
   def testAvgPoolGradMismatchedGradShapeRaisesError(self):
-    # Each grad has more channels than the input. The CPU kernel used to write
-    # past the end of its output buffer instead of raising an error.
+    # Each grad has a different number of channels than the input. With more
+    # channels, the CPU kernel used to write past the end of its output buffer
+    # instead of raising an error.
     for orig_input_shape, grad_shape, ksize, expected in (
         ([1, 28, 28, 3], [1, 14, 14, 6], 2, r"\[1,14,14,3\], but got "
          r"\[1,14,14,6\]"),
         ([2, 2, 2, 2], [2, 3, 3, 3], 1, r"\[2,2,2,2\], but got \[2,3,3,3\]"),
+        ([1, 10, 10, 3], [1, 5, 5, 0], 2, r"\[1,5,5,3\], but got \[1,5,5,0\]"),
     ):
       with self.assertRaisesRegex(
           (errors_impl.InvalidArgumentError, ValueError),
@@ -2443,11 +2445,16 @@ class PoolingTest(test.TestCase, parameterized.TestCase):
 
   @test_util.run_in_graph_and_eager_modes
   def testAvgPoolGradEmptyTensorFastExit(self):
-    for empty_shape in ([0, 10, 10, 3], [1, 0, 10, 3], [1, 10, 10, 0]):
+    for empty_input_shape, empty_grad_shape in (
+        ([0, 10, 10, 3], [0, 5, 5, 3]),
+        ([1, 0, 10, 3], [1, 0, 5, 3]),
+        ([1, 10, 10, 0], [1, 5, 5, 0]),
+        ([1, 1, 10, 3], [1, 0, 5, 3]),
+    ):
       self.evaluate(
           gen_nn_ops.AvgPoolGrad(
-              orig_input_shape=empty_shape,
-              grad=array_ops.zeros(empty_shape),
+              orig_input_shape=empty_input_shape,
+              grad=array_ops.zeros(empty_grad_shape),
               ksize=[1, 2, 2, 1],
               strides=[1, 2, 2, 1],
               padding="VALID",

From 027180c6b95a003730538155f597178e9f4b226b Mon Sep 17 00:00:00 2001
From: Noxtimo <139155109+Noxtimo@users.noreply.github.com>
Date: Wed, 7 Oct 2026 11:39:09 +0100
Subject: [PATCH 0229/1000] test: assert NaN output in cbrt edge-case test

---
 tensorflow/python/ops/numpy_ops/np_math_ops_test.py | 1 +
 1 file changed, 1 insertion(+)

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 237378e23ee..611057037dd 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
@@ -899,6 +899,7 @@ class MathTest(test.TestCase, parameterized.TestCase):
       # inf ** (1/3) == inf; NaN ** (1/3) == NaN
       # Negative sign bit for -inf is preserved.
       self.assertAllEqual(np_math_ops.signbit(out_edge[0:2]), [True, False])
+      self.assertAllEqual(math_ops.is_nan(out_edge[2]), True)
 
   def testSinc(self):
     for dtype in (dtypes.float32, dtypes.float64):

From 32ceb7f0dd2ff59bea82d4a60555805902a0f625 Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Wed, 7 Oct 2026 07:18:56 -0700
Subject: [PATCH 0230/1000] Simplify average axis check and normalize negative
 axes in place

---
 tensorflow/python/ops/numpy_ops/np_math_ops.py      | 5 +++--
 tensorflow/python/ops/numpy_ops/np_math_ops_test.py | 9 +++++++++
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py
index 10b9b05696f..d978abe7328 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@@ -1678,12 +1678,13 @@ def average(a, axis=None, weights=None, returned=False):  # pylint: disable=miss
   # backend kernel fail with a confusing error.
   maybe_rank = a.shape.rank
   if axis is not None and maybe_rank is not None:
-    normalized = axis + maybe_rank if axis < 0 else axis
-    if normalized < 0 or normalized >= maybe_rank:
+    if not (-maybe_rank <= axis < maybe_rank):
       raise ValueError(
           f'Argument `axis` (received axis={axis}) is out of bounds '
           f'for input of rank {maybe_rank}.'
       )
+    if axis < 0:
+      axis += maybe_rank
   default_float_type = np_utils.result_type(float)
   if weights is None:  # Treat all weights as 1
     if not np.issubdtype(a.dtype.as_numpy_dtype, np.inexact):
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 4208f893988..de3c1104ae7 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
@@ -579,6 +579,15 @@ class MathTest(test.TestCase, parameterized.TestCase):
       np_math_ops.average(x, axis=-3)
     with self.assertRaisesRegex(ValueError, 'out of bounds'):
       np_math_ops.average(x, axis=2, weights=np.ones([2, 3]))
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_math_ops.average(x, axis=2, returned=True)
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_math_ops.average(x, axis=2, weights=np.ones([3]))
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_math_ops.average(x, axis=-3, weights=np.ones([2, 3]))
+    scalar = np_array_ops.array(5.0)
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_math_ops.average(scalar, axis=0)
 
   def testAverageWrongShape(self):
     with self.assertRaisesWithPredicateMatch(errors.InvalidArgumentError, r''):

From 22024867a3496012bfd71971efe16140465f5643 Mon Sep 17 00:00:00 2001
From: Daniil Kuts <kutz@ispras.ru>
Date: Wed, 7 Oct 2026 19:01:07 +0300
Subject: [PATCH 0231/1000] Remove redundant nullptr guard in RootDataset

---
 tensorflow/core/data/root_dataset.cc | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/tensorflow/core/data/root_dataset.cc b/tensorflow/core/data/root_dataset.cc
index 5c32f1b1af6..28f7d992a3b 100644
--- a/tensorflow/core/data/root_dataset.cc
+++ b/tensorflow/core/data/root_dataset.cc
@@ -409,11 +409,8 @@ RootDataset::RootDataset(core::RefCountPtr<DatasetBase> input,
       params_(std::move(params)) {
   owned_input_ = std::move(input);
   input_ = owned_input_.get();
-  random_indexing_compatible_ = absl::OkStatus();
-  if (input_ != nullptr) {
-    random_indexing_compatible_ = input_->RandomIndexingCompatible();
-    AddTraceMetadata(params_, input_->options(), &traceme_metadata_);
-  }
+  random_indexing_compatible_ = input_->RandomIndexingCompatible();
+  AddTraceMetadata(params_, input_->options(), &traceme_metadata_);
 }
 
 RootDataset::~RootDataset() = default;

From 12c79d0524dc2f6c830b2c3d18dcd45fc28aee67 Mon Sep 17 00:00:00 2001
From: bodapatisaikrishna <saikrishnabodapati@gmail.com>
Date: Wed, 7 Oct 2026 21:37:04 +0530
Subject: [PATCH 0232/1000] Validate input and filter ranks in tf.nn.conv2d

When inputs or filters with rank < 4 are passed to tf.nn.conv2d in eager
execution, execution proceeds to C++ kernel dimension index lookup where
GetTensorDimIndex triggers a fatal CHECK failure in tensor_format.h,
crashing the process instead of raising a catchable exception.

Validate that input and filter have rank at least 4 in tf.nn.conv2d before
calling the internal op, raising a ValueError with clear rank information.

Fixes #77336

Signed-off-by: bodapatisaikrishna <saikrishnabodapati@gmail.com>
---
 .../kernel_tests/nn_ops/conv_ops_test.py      | 20 +++++++++++++++++++
 tensorflow/python/ops/nn_ops.py               | 20 ++++++++++++++++++-
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
index d3101c3b21c..603a5e332da 100644
--- a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
@@ -3132,6 +3132,26 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
     #           strides=[1, 1, 1, 1],
     #           padding=[[0, 0], [2, 2], [2, 2], [0, 0]]))
 
+    # Invalid input or filter rank (< 4) in eager/graph execution.
+    with self.assertRaisesRegex(ValueError, "`input.shape.rank` must be at least 4"):
+      self.evaluate(
+          nn_ops.conv2d(
+              np.ones([1, 1, 3], dtype=np.float32),
+              np.ones([1, 2, 2, 1], dtype=np.float32),
+              strides=[1, 1, 1, 1],
+              padding="VALID",
+          )
+      )
+    with self.assertRaisesRegex(ValueError, "`filter.shape.rank` must be at least 4"):
+      self.evaluate(
+          nn_ops.conv2d(
+              np.ones([1, 1, 3, 3], dtype=np.float32),
+              np.ones([1, 2, 2], dtype=np.float32),
+              strides=[1, 1, 1, 1],
+              padding="VALID",
+          )
+      )
+
     # Filter dimensions must be greater than 0.
     with self.assertRaisesRegex(
         errors_impl.InvalidArgumentError, "filter must not have zero elements"
diff --git a/tensorflow/python/ops/nn_ops.py b/tensorflow/python/ops/nn_ops.py
index 1e43bd92edb..1dccf45dd88 100644
--- a/tensorflow/python/ops/nn_ops.py
+++ b/tensorflow/python/ops/nn_ops.py
@@ -2668,7 +2668,25 @@ def conv2d(  # pylint: disable=redefined-builtin,dangerous-default-value
   ndims = getattr(shape, "ndims", -1)
   if ndims == -1:
     ndims = len(shape)
-  if ndims in (4, 3, 2, 1, 0, None):
+  if ndims is not None and ndims < 4:
+    raise ValueError(
+        f"`input.shape.rank` must be at least 4. "
+        f"Received: input.shape={shape} with rank {ndims}."
+    )
+
+  filter_shape = getattr(filter, "shape", None)
+  filter_ndims = (
+      getattr(filter_shape, "ndims", -1) if filter_shape is not None else -1
+  )
+  if filter_ndims == -1 and filter_shape is not None:
+    filter_ndims = len(filter_shape)
+  if filter_ndims is not None and filter_ndims != -1 and filter_ndims < 4:
+    raise ValueError(
+        f"`filter.shape.rank` must be at least 4. "
+        f"Received: filter.shape={filter_shape} with rank {filter_ndims}."
+    )
+
+  if ndims in (4, None):
     # We avoid calling squeeze_batch_dims to reduce extra python function
     # call slowdown in eager mode.  This branch doesn't require reshapes.
     return gen_nn_ops.conv2d(

From 22f9c782cf2c24bc34a899b59b1cb237fbaf7c0e Mon Sep 17 00:00:00 2001
From: bodapatisaikrishna <saikrishnabodapati@gmail.com>
Date: Wed, 7 Oct 2026 22:05:51 +0530
Subject: [PATCH 0233/1000] Simplify filter rank extraction in tf.nn.conv2d

Signed-off-by: bodapatisaikrishna <saikrishnabodapati@gmail.com>
---
 tensorflow/python/ops/nn_ops.py | 19 +++++++++----------
 1 file changed, 9 insertions(+), 10 deletions(-)

diff --git a/tensorflow/python/ops/nn_ops.py b/tensorflow/python/ops/nn_ops.py
index 1dccf45dd88..1913cc46568 100644
--- a/tensorflow/python/ops/nn_ops.py
+++ b/tensorflow/python/ops/nn_ops.py
@@ -2675,16 +2675,15 @@ def conv2d(  # pylint: disable=redefined-builtin,dangerous-default-value
     )
 
   filter_shape = getattr(filter, "shape", None)
-  filter_ndims = (
-      getattr(filter_shape, "ndims", -1) if filter_shape is not None else -1
-  )
-  if filter_ndims == -1 and filter_shape is not None:
-    filter_ndims = len(filter_shape)
-  if filter_ndims is not None and filter_ndims != -1 and filter_ndims < 4:
-    raise ValueError(
-        f"`filter.shape.rank` must be at least 4. "
-        f"Received: filter.shape={filter_shape} with rank {filter_ndims}."
-    )
+  if filter_shape is not None:
+    filter_ndims = getattr(filter_shape, "ndims", -1)
+    if filter_ndims == -1:
+      filter_ndims = len(filter_shape)
+    if filter_ndims is not None and filter_ndims < 4:
+      raise ValueError(
+          f"`filter.shape.rank` must be at least 4. "
+          f"Received: filter.shape={filter_shape} with rank {filter_ndims}."
+      )
 
   if ndims in (4, None):
     # We avoid calling squeeze_batch_dims to reduce extra python function

From d69b2673d9c559ebddbe04b699e4e95727a183e3 Mon Sep 17 00:00:00 2001
From: Boris Kozak <boriskozak@gmail.com>
Date: Wed, 7 Oct 2026 16:36:47 +0000
Subject: [PATCH 0234/1000] Rework per review: revert Horner init, broadcast_to
 for degree-0

Revert the coeffs[0] + x*0 accumulator init: matching numpy's nan for
inf inputs was a mathematical regression (lim x->inf of x+1 is inf)
and added unconditional overhead to the hot path.

The len(coeffs) == 1 case now uses
array_ops.broadcast_to(p, array_ops.shape(x)) instead of p + x*0,
broadcasting cleanly without arithmetic that could nan-poison
non-finite inputs.

Replace testNonFiniteXPropagation with testNonFiniteXReturnsInf:
inf inputs must evaluate to inf across float16/32/64 and
complex64/128, scalar and array shapes, degrees 0-2 (complex limited
to degrees 0-1; at degree 2+ the complex multiply itself can nan the
real part, e.g. complex128 (inf+0j)**2).
---
 .../kernel_tests/math_ops/cwise_ops_test.py   | 51 ++++++++++++++-----
 tensorflow/python/ops/math_ops.py             | 11 ++--
 2 files changed, 44 insertions(+), 18 deletions(-)

diff --git a/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py b/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
index a8cc65fe3b3..a625aca5791 100644
--- a/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
@@ -1340,19 +1340,44 @@ class PolyvalTest(test.TestCase):
       self.assertEqual(tf_val.shape, x.shape)
       self.assertAllClose(np_val, self.evaluate(tf_val))
 
-  def testNonFiniteXPropagation(self):
-    # Non-finite x must propagate to NaN, matching numpy.polyval, which
-    # starts its Horner accumulator at zero (0 * inf + c -> nan).
-    for dtype in [np.float32, np.float64, np.complex64]:
-      for x_val in [float("inf"), float("-inf"), float("nan")]:
-        with self.subTest(dtype=dtype, x=x_val):
-          x = np.array([1.0, x_val, 3.0], dtype=dtype)
-          coeffs = [dtype(1.0), dtype(1.0)]
-          with np.errstate(invalid="ignore"):
-            np_val = np.polyval(coeffs, x)
-          with self.cached_session():
-            tf_val = math_ops.polyval(coeffs, x)
-            self.assertAllClose(np_val, self.evaluate(tf_val))
+  def testNonFiniteXReturnsInf(self):
+    # inf x must evaluate to inf (not nan): the limit of a polynomial such as
+    # p(x) = x + 1 as x -> inf is mathematically inf. NumPy returns nan here
+    # only because it starts its Horner accumulator at zero (0 * inf -> nan).
+    # Complex dtypes are only checked at degrees 0 and 1: at degree 2+ the
+    # complex multiply itself can produce a nan real part (e.g. complex128
+    # (inf+0j)**2 -> nan+nanj), which is inherent to complex inf arithmetic
+    # rather than accumulator initialization.
+    for dtype in [
+        np.float16, np.float32, np.float64, np.complex64, np.complex128
+    ]:
+      degrees = [(0, [2.0]), (1, [1.0, 1.0])]
+      if not np.issubdtype(dtype, np.complexfloating):
+        degrees.append((2, [1.0, 1.0, 1.0]))
+      for x_val in [float("inf"), float("-inf")]:
+        for degree, coeffs in degrees:
+          for x_np in [
+              np.array(x_val, dtype=dtype),
+              np.full(3, x_val, dtype=dtype)
+          ]:
+            with self.subTest(dtype=dtype,
+                              x=x_val,
+                              degree=degree,
+                              scalar=x_np.shape == ()):
+              with self.cached_session():
+                tf_val = math_ops.polyval([dtype(c) for c in coeffs], x_np)
+                result = self.evaluate(tf_val)
+                self.assertEqual(tf_val.shape, x_np.shape)
+                if degree == 0:
+                  # A constant polynomial ignores x; the result is the
+                  # constant broadcast against x's shape.
+                  self.assertAllClose(np.full(x_np.shape, dtype(2.0)), result)
+                elif np.issubdtype(dtype, np.complexfloating):
+                  # Complex inf arithmetic can produce nan imaginary parts
+                  # (e.g. (1+0j) * (inf+0j)); the real part must stay inf.
+                  self.assertTrue(np.all(np.isinf(np.real(result))))
+                else:
+                  self.assertTrue(np.all(np.isinf(result)))
 
   def test_coeffs_raise(self):
     x = np.random.rand(2, 2).astype(np.float32)
diff --git a/tensorflow/python/ops/math_ops.py b/tensorflow/python/ops/math_ops.py
index 5e9c0b20410..aa8fb638e3f 100644
--- a/tensorflow/python/ops/math_ops.py
+++ b/tensorflow/python/ops/math_ops.py
@@ -5625,13 +5625,14 @@ def polyval(coeffs, x, name=None):
         ops.convert_to_tensor(coeff, name=("coeff_%d" % index))
         for index, coeff in enumerate(coeffs)
     ]
-    # Start the Horner accumulator at coeffs[0] + x*0 (rather than coeffs[0])
-    # so that non-finite x (inf/nan) propagates to NaN on the first multiply,
-    # matching numpy.polyval, which starts its accumulator at zero. For finite
-    # x this is a no-op, and it also broadcasts the result against x's shape.
-    p = coeffs[0] + x * 0
+    p = coeffs[0]
     for c in coeffs[1:]:
       p = c + p * x
+    # For single-coefficient polynomials, the loop above never executes,
+    # so broadcast p against x's shape without introducing arithmetic
+    # that would alter non-finite behavior (e.g. x*0 turns inf into nan).
+    if len(coeffs) == 1:
+      p = array_ops.broadcast_to(p, array_ops.shape(x))
     return p
 
 

From ad72f53163999c16543138101fae64abd6c3174a Mon Sep 17 00:00:00 2001
From: Boris Kozak <boriskozak@gmail.com>
Date: Wed, 7 Oct 2026 18:52:58 +0000
Subject: [PATCH 0235/1000] Rework per round-2 review: revert degree-0 fallback
 to p + x*0

broadcast_to broke NaN propagation (testSingleCoeffNanPropagation)
and two-way multidim broadcasting ((2,1) x (1,4) raised ValueError).
Keep plain coeffs[0] Horner init; drop degree 0 from the limits test.
---
 .../python/kernel_tests/math_ops/cwise_ops_test.py     | 10 +++-------
 tensorflow/python/ops/math_ops.py                      |  5 +----
 2 files changed, 4 insertions(+), 11 deletions(-)

diff --git a/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py b/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
index a625aca5791..cd95fce1e96 100644
--- a/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
@@ -1344,14 +1344,14 @@ class PolyvalTest(test.TestCase):
     # inf x must evaluate to inf (not nan): the limit of a polynomial such as
     # p(x) = x + 1 as x -> inf is mathematically inf. NumPy returns nan here
     # only because it starts its Horner accumulator at zero (0 * inf -> nan).
-    # Complex dtypes are only checked at degrees 0 and 1: at degree 2+ the
+    # Complex dtypes are only checked at degree 1: at degree 2+ the
     # complex multiply itself can produce a nan real part (e.g. complex128
     # (inf+0j)**2 -> nan+nanj), which is inherent to complex inf arithmetic
     # rather than accumulator initialization.
     for dtype in [
         np.float16, np.float32, np.float64, np.complex64, np.complex128
     ]:
-      degrees = [(0, [2.0]), (1, [1.0, 1.0])]
+      degrees = [(1, [1.0, 1.0])]
       if not np.issubdtype(dtype, np.complexfloating):
         degrees.append((2, [1.0, 1.0, 1.0]))
       for x_val in [float("inf"), float("-inf")]:
@@ -1368,11 +1368,7 @@ class PolyvalTest(test.TestCase):
                 tf_val = math_ops.polyval([dtype(c) for c in coeffs], x_np)
                 result = self.evaluate(tf_val)
                 self.assertEqual(tf_val.shape, x_np.shape)
-                if degree == 0:
-                  # A constant polynomial ignores x; the result is the
-                  # constant broadcast against x's shape.
-                  self.assertAllClose(np.full(x_np.shape, dtype(2.0)), result)
-                elif np.issubdtype(dtype, np.complexfloating):
+                if np.issubdtype(dtype, np.complexfloating):
                   # Complex inf arithmetic can produce nan imaginary parts
                   # (e.g. (1+0j) * (inf+0j)); the real part must stay inf.
                   self.assertTrue(np.all(np.isinf(np.real(result))))
diff --git a/tensorflow/python/ops/math_ops.py b/tensorflow/python/ops/math_ops.py
index aa8fb638e3f..12c40415459 100644
--- a/tensorflow/python/ops/math_ops.py
+++ b/tensorflow/python/ops/math_ops.py
@@ -5628,11 +5628,8 @@ def polyval(coeffs, x, name=None):
     p = coeffs[0]
     for c in coeffs[1:]:
       p = c + p * x
-    # For single-coefficient polynomials, the loop above never executes,
-    # so broadcast p against x's shape without introducing arithmetic
-    # that would alter non-finite behavior (e.g. x*0 turns inf into nan).
     if len(coeffs) == 1:
-      p = array_ops.broadcast_to(p, array_ops.shape(x))
+      p = p + x * 0
     return p
 
 

From bca73ef8e519e312a1c350ccf42c96a1127d78fe Mon Sep 17 00:00:00 2001
From: Raashish Aggarwal <94279692+raashish1601@users.noreply.github.com>
Date: Thu, 8 Oct 2026 00:34:39 +0530
Subject: [PATCH 0236/1000] Give value_range the dtype of values in
 histogram_fixed_width_bins

A Python list value_range was converted to float32, so float64 values
failed in the subtraction with a dtype mismatch. histogram_fixed_width,
which uses the op directly, already accepts the same arguments.
---
 tensorflow/python/kernel_tests/histogram_ops_test.py | 10 ++++++++++
 tensorflow/python/ops/histogram_ops.py               |  3 ++-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/kernel_tests/histogram_ops_test.py b/tensorflow/python/kernel_tests/histogram_ops_test.py
index d9c6fe8fdcc..44ac4ee5535 100644
--- a/tensorflow/python/kernel_tests/histogram_ops_test.py
+++ b/tensorflow/python/kernel_tests/histogram_ops_test.py
@@ -55,6 +55,16 @@ class BinValuesFixedWidth(test.TestCase, parameterized.TestCase):
       self.assertEqual(dtypes.int32, bins.dtype)
       self.assertAllClose(expected_bins, self.evaluate(bins))
 
+  @parameterized.parameters(np.float32, np.float64)
+  def test_1d_values_python_value_range(self, dtype):
+    # value_range given as a Python list takes the dtype of values.
+    values = np.array([-1.0, 0.0, 1.5, 2.0, 5.0, 15]).astype(dtype)
+    expected_bins = [0, 0, 1, 2, 4, 4]
+    with self.cached_session():
+      bins = histogram_ops.histogram_fixed_width_bins(
+          values, [0.0, 5.0], nbins=5)
+      self.assertAllClose(expected_bins, self.evaluate(bins))
+
   def test_2d_values(self):
     # Bins will be:
     #   (-inf, 1), [1, 2), [2, 3), [3, 4), [4, inf)
diff --git a/tensorflow/python/ops/histogram_ops.py b/tensorflow/python/ops/histogram_ops.py
index f68e3ab560e..a156144afe5 100644
--- a/tensorflow/python/ops/histogram_ops.py
+++ b/tensorflow/python/ops/histogram_ops.py
@@ -77,7 +77,8 @@ def histogram_fixed_width_bins(values,
     shape = array_ops.shape(values)
 
     values = array_ops.reshape(values, [-1])
-    value_range = ops.convert_to_tensor(value_range, name='value_range')
+    value_range = ops.convert_to_tensor(
+        value_range, dtype_hint=values.dtype, name='value_range')
     nbins = ops.convert_to_tensor(nbins, dtype=dtypes.int32, name='nbins')
     check = control_flow_assert.Assert(
         math_ops.greater(nbins, 0), ['nbins %s must > 0' % nbins])

From 37630ff7f6689f0ae9fa0585fd7abf2b86332699 Mon Sep 17 00:00:00 2001
From: Raashish Aggarwal <94279692+raashish1601@users.noreply.github.com>
Date: Thu, 8 Oct 2026 01:25:58 +0530
Subject: [PATCH 0237/1000] Support a tuple axis in tf.experimental.numpy.var
 with ddof

The ddof path compared axis with 0 in Python, so a tuple or list axis
raised TypeError and a tensor axis could not be used in graph mode. Convert
the axis to an int32 tensor and normalize negative values with where_v2.
---
 tensorflow/python/ops/numpy_ops/np_array_ops.py      | 7 +++++--
 tensorflow/python/ops/numpy_ops/np_array_ops_test.py | 2 ++
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index 03791ba2e3a..b6c0b3ceab7 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -833,8 +833,11 @@ def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=None):  # pylint: d
       if axis is None:
         n = array_ops.size(input_tensor)
       else:
-        if axis < 0:
-          axis += array_ops.rank(input_tensor)
+        # axis can be an int or a sequence of ints.
+        axis = math_ops.cast(axis, dtypes.int32)
+        axis = array_ops.where_v2(
+            axis < 0, axis + array_ops.rank(input_tensor), axis
+        )
         n = math_ops.reduce_prod(
             array_ops.gather(array_ops.shape(input_tensor), axis)
         )
diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index dffb95e87b9..79a7d3d9231 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -993,6 +993,8 @@ class ArrayMethodsTest(test.TestCase):
     run_test([[1.0, 2.0], [3.0, 4.0]], axis=-1, ddof=1, keepdims=True)
     run_test([1.0j, 2.0, 3.0j], ddof=1)
     run_test([[1.0j, 2.0], [3.0j, 4.0]], axis=0, ddof=1)
+    run_test(np.arange(8.0).reshape((2, 2, 2)).tolist(), axis=(0, 2), ddof=1)
+    run_test(np.arange(8.0).reshape((2, 2, 2)).tolist(), axis=(-1, 0), ddof=1)
     run_test(np.arange(8).reshape((2, 2, 2)).tolist(), axis=(0, 2))
     run_test(
         np.arange(8).reshape((2, 2, 2)).tolist(), axis=(0, 2), keepdims=True)

From 4df96689362ee663b2699ccd952cd073d39ac0e5 Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Thu, 8 Oct 2026 01:57:05 +0530
Subject: [PATCH 0238/1000] bind out_backprop to orig_output and cover each
 precondition in FractionalMaxPoolGrad

---
 .../core/kernels/fractional_max_pool_op.cc    |  19 ++-
 .../nn_ops/fractional_max_pool_op_test.py     | 113 ++++++++++++++++--
 2 files changed, 120 insertions(+), 12 deletions(-)

diff --git a/tensorflow/core/kernels/fractional_max_pool_op.cc b/tensorflow/core/kernels/fractional_max_pool_op.cc
index 977c83e2217..b1a84631c17 100644
--- a/tensorflow/core/kernels/fractional_max_pool_op.cc
+++ b/tensorflow/core/kernels/fractional_max_pool_op.cc
@@ -328,6 +328,15 @@ class FractionalMaxPoolGradOp : public OpKernel {
             "col_pooling_sequence must have ", output_size[2] + 1,
             " elements to match the ", output_size[2],
             " cols of orig_output, but got ", width_seq_tensor.dim_size(0))));
+    // Step 2 walks out_backprop and indexes tensor_out_arg_max, which is sized
+    // from orig_output, with the same linear index, so an out_backprop larger
+    // than orig_output reads past that buffer.
+    OP_REQUIRES(
+        context, tensor_out.shape() == out_backprop.shape(),
+        absl::InvalidArgumentError(absl::StrCat(
+            "orig_output and out_backprop must have the same shape, got ",
+            tensor_out.shape().DebugString(), " and ",
+            out_backprop.shape().DebugString())));
 
     // ---------
     // Step 1
@@ -406,7 +415,7 @@ class FractionalMaxPoolGradOp : public OpKernel {
                 if (output_ref < input_ref ||
                     out_arg_max_ref == kInvalidMaxPoolingIndex) {
                   output_ref = input_ref;
-                  int input_offset = in_index * input_size[3] + d;
+                  int64_t input_offset = in_index * input_size[3] + d;
                   out_arg_max_ref = input_offset;
                 }
               }
@@ -438,11 +447,11 @@ class FractionalMaxPoolGradOp : public OpKernel {
     auto out_backprop_flat = out_backprop.flat<T>();
     auto input_backprop_flat = output->flat<T>();
     auto out_arg_max_flat = tensor_out_arg_max.flat<int64_t>();
-    int num_total_outputs = out_backprop_flat.size();
-    int num_total_inputs = input_backprop_flat.size();
+    int64_t num_total_outputs = out_backprop_flat.size();
+    int64_t num_total_inputs = input_backprop_flat.size();
 
-    for (int index = 0; index < num_total_outputs; ++index) {
-      int input_backprop_index = out_arg_max_flat(index);
+    for (int64_t index = 0; index < num_total_outputs; ++index) {
+      int64_t input_backprop_index = out_arg_max_flat(index);
       OP_REQUIRES(
           context,
           input_backprop_index >= 0 && input_backprop_index < num_total_inputs,
diff --git a/tensorflow/python/kernel_tests/nn_ops/fractional_max_pool_op_test.py b/tensorflow/python/kernel_tests/nn_ops/fractional_max_pool_op_test.py
index ae6ea1b9ac4..8390f764c3f 100644
--- a/tensorflow/python/kernel_tests/nn_ops/fractional_max_pool_op_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/fractional_max_pool_op_test.py
@@ -658,7 +658,9 @@ class FractionalMaxPoolGradTest(test.TestCase):
                           input_backprop_overlapping)
 
   def testInvalidSeqRaiseErrorForFractionalMaxPoolGrad(self):
-    with self.assertRaises(errors.InvalidArgumentError):
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError, "row_pooling_sequence must have"
+    ):
       with self.cached_session():
         overlapping = True
         orig_input = constant_op.constant(
@@ -711,14 +713,20 @@ class FractionalMaxPoolGradTest(test.TestCase):
       self.evaluate(t)
 
   def testOverLargeSeqRaiseErrorForFractionalMaxPoolGrad(self):
-    with self.assertRaises(errors.InvalidArgumentError):
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "Row sequence tensor values must not be negative",
+    ):
       with self.cached_session():
         overlapping = False
-        orig_input = [[[[1, 1, 1, 1, 1]]]]
-        orig_output = [[[[1, 1, 1]]]]
-        out_backprop = [[[[3], [3], [6]]]]
-        row_pooling_sequence = [-0x4000000, 1, 1]
-        col_pooling_sequence = [-0x4000000, 1, 1]
+        orig_input = constant_op.constant(
+            1.0, shape=[1, 2, 2, 1], dtype=dtypes.float32)
+        orig_output = constant_op.constant(
+            1.0, shape=[1, 1, 1, 1], dtype=dtypes.float32)
+        out_backprop = constant_op.constant(
+            1.0, shape=[1, 1, 1, 1], dtype=dtypes.float32)
+        row_pooling_sequence = [-0x4000000, 2]
+        col_pooling_sequence = [0, 2]
         t = gen_nn_ops.FractionalMaxPoolGrad(
             orig_input=orig_input,
             orig_output=orig_output,
@@ -756,6 +764,97 @@ class FractionalMaxPoolGradTest(test.TestCase):
             overlapping=True)
         self.evaluate(t)
 
+  def testColSeqLengthMismatchRaisesErrorForFractionalMaxPoolGrad(self):
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError, "col_pooling_sequence must have"
+    ):
+      with self.cached_session():
+        t = gen_nn_ops.FractionalMaxPoolGrad(
+            orig_input=constant_op.constant(1.0, shape=[1, 4, 4, 1]),
+            orig_output=constant_op.constant(1.0, shape=[1, 2, 2, 1]),
+            out_backprop=constant_op.constant(1.0, shape=[1, 2, 2, 1]),
+            row_pooling_sequence=[0, 2, 4],
+            col_pooling_sequence=[0, 2, 3, 4],
+            overlapping=False)
+        self.evaluate(t)
+
+  def testBatchMismatchRaisesErrorForFractionalMaxPoolGrad(self):
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "orig_input and orig_output must have the same batch size",
+    ):
+      with self.cached_session():
+        t = gen_nn_ops.FractionalMaxPoolGrad(
+            orig_input=constant_op.constant(1.0, shape=[2, 4, 4, 1]),
+            orig_output=constant_op.constant(1.0, shape=[1, 2, 2, 1]),
+            out_backprop=constant_op.constant(1.0, shape=[1, 2, 2, 1]),
+            row_pooling_sequence=[0, 2, 4],
+            col_pooling_sequence=[0, 2, 4],
+            overlapping=False)
+        self.evaluate(t)
+
+  def testDepthMismatchRaisesErrorForFractionalMaxPoolGrad(self):
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "orig_input and orig_output must have the same depth",
+    ):
+      with self.cached_session():
+        t = gen_nn_ops.FractionalMaxPoolGrad(
+            orig_input=constant_op.constant(1.0, shape=[1, 4, 4, 2]),
+            orig_output=constant_op.constant(1.0, shape=[1, 2, 2, 1]),
+            out_backprop=constant_op.constant(1.0, shape=[1, 2, 2, 1]),
+            row_pooling_sequence=[0, 2, 4],
+            col_pooling_sequence=[0, 2, 4],
+            overlapping=False)
+        self.evaluate(t)
+
+  def testOutBackpropShapeMismatchRaisesErrorForFractionalMaxPoolGrad(self):
+    # Step 2 indexes tensor_out_arg_max, sized from orig_output, with
+    # out_backprop's linear index, so a larger out_backprop reads out of
+    # bounds.
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "orig_output and out_backprop must have the same shape",
+    ):
+      with self.cached_session():
+        orig_input = constant_op.constant(
+            1.0, shape=[1, 2, 2, 1], dtype=dtypes.float32)
+        orig_output = constant_op.constant(
+            1.0, shape=[1, 1, 1, 1], dtype=dtypes.float32)
+        out_backprop = constant_op.constant(
+            1.0, shape=[1, 2, 2, 1], dtype=dtypes.float32)
+        t = gen_nn_ops.FractionalMaxPoolGrad(
+            orig_input=orig_input,
+            orig_output=orig_output,
+            out_backprop=out_backprop,
+            row_pooling_sequence=[0, 2],
+            col_pooling_sequence=[0, 2],
+            overlapping=True)
+        self.evaluate(t)
+
+  def testValueMismatchRaisesErrorForFractionalMaxPoolGrad(self):
+    # Valid sequence lengths, but orig_output does not hold the values the
+    # replay recomputes.
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        "tensor_out_dup is not the same as tensor_out",
+    ):
+      with self.cached_session():
+        orig_input = constant_op.constant(
+            1.0, shape=[1, 2, 2, 1], dtype=dtypes.float32)
+        orig_output = constant_op.constant(
+            2.0, shape=[1, 1, 1, 1], dtype=dtypes.float32)
+        out_backprop = constant_op.constant(
+            1.0, shape=[1, 1, 1, 1], dtype=dtypes.float32)
+        t = gen_nn_ops.FractionalMaxPoolGrad(
+            orig_input=orig_input,
+            orig_output=orig_output,
+            out_backprop=out_backprop,
+            row_pooling_sequence=[0, 2],
+            col_pooling_sequence=[0, 2],
+            overlapping=True)
+        self.evaluate(t)
+
 
 if __name__ == "__main__":
   test.main()

From 57e46a4b4eddffe556c781c4e15e4d4263a1afc9 Mon Sep 17 00:00:00 2001
From: Boris Kozak <boriskozak@gmail.com>
Date: Wed, 7 Oct 2026 21:51:29 +0000
Subject: [PATCH 0239/1000] Add single-coefficient inf->nan test; restore
 degree-0 comment

---
 .../kernel_tests/math_ops/cwise_ops_test.py   | 20 +++++++++++++++++++
 tensorflow/python/ops/math_ops.py             |  3 +++
 2 files changed, 23 insertions(+)

diff --git a/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py b/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
index cd95fce1e96..7b55224a502 100644
--- a/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
@@ -1340,6 +1340,26 @@ class PolyvalTest(test.TestCase):
       self.assertEqual(tf_val.shape, x.shape)
       self.assertAllClose(np_val, self.evaluate(tf_val))
 
+  def testSingleCoeffInfReturnsNan(self):
+    for dtype in [
+        np.float16, np.float32, np.float64, np.complex64, np.complex128
+    ]:
+      for x_val in [float("inf"), float("-inf")]:
+        for x_np in [
+            np.array(x_val, dtype=dtype),
+            np.full(3, x_val, dtype=dtype)
+        ]:
+          with self.subTest(dtype=dtype, x=x_val, scalar=x_np.shape == ()):
+            coeffs = [dtype(2.0)]
+            with self.cached_session():
+              tf_val = math_ops.polyval(coeffs, x_np)
+              result = self.evaluate(tf_val)
+              self.assertEqual(tf_val.shape, x_np.shape)
+              if np.issubdtype(dtype, np.complexfloating):
+                self.assertTrue(np.all(np.isnan(np.real(result))))
+              else:
+                self.assertTrue(np.all(np.isnan(result)))
+
   def testNonFiniteXReturnsInf(self):
     # inf x must evaluate to inf (not nan): the limit of a polynomial such as
     # p(x) = x + 1 as x -> inf is mathematically inf. NumPy returns nan here
diff --git a/tensorflow/python/ops/math_ops.py b/tensorflow/python/ops/math_ops.py
index 12c40415459..261a0d1bcf4 100644
--- a/tensorflow/python/ops/math_ops.py
+++ b/tensorflow/python/ops/math_ops.py
@@ -5628,6 +5628,9 @@ def polyval(coeffs, x, name=None):
     p = coeffs[0]
     for c in coeffs[1:]:
       p = c + p * x
+    # For single-coefficient polynomials, the loop above never executes,
+    # so x is unused. Add x*0 to broadcast against x's shape and
+    # propagate NaN, matching numpy.polyval behavior.
     if len(coeffs) == 1:
       p = p + x * 0
     return p

From 18672bbe0123dfd4c7324c2706ac201ee212bcd1 Mon Sep 17 00:00:00 2001
From: Boris Kozak <boriskozak@gmail.com>
Date: Thu, 8 Oct 2026 00:02:08 +0000
Subject: [PATCH 0240/1000] Run new polyval tests in graph and eager modes

---
 .../kernel_tests/math_ops/cwise_ops_test.py   | 36 +++++++++----------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py b/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
index 7b55224a502..27d817cc630 100644
--- a/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
@@ -1340,6 +1340,7 @@ class PolyvalTest(test.TestCase):
       self.assertEqual(tf_val.shape, x.shape)
       self.assertAllClose(np_val, self.evaluate(tf_val))
 
+  @test_util.run_in_graph_and_eager_modes
   def testSingleCoeffInfReturnsNan(self):
     for dtype in [
         np.float16, np.float32, np.float64, np.complex64, np.complex128
@@ -1351,15 +1352,15 @@ class PolyvalTest(test.TestCase):
         ]:
           with self.subTest(dtype=dtype, x=x_val, scalar=x_np.shape == ()):
             coeffs = [dtype(2.0)]
-            with self.cached_session():
-              tf_val = math_ops.polyval(coeffs, x_np)
-              result = self.evaluate(tf_val)
-              self.assertEqual(tf_val.shape, x_np.shape)
-              if np.issubdtype(dtype, np.complexfloating):
-                self.assertTrue(np.all(np.isnan(np.real(result))))
-              else:
-                self.assertTrue(np.all(np.isnan(result)))
+            tf_val = math_ops.polyval(coeffs, x_np)
+            result = self.evaluate(tf_val)
+            self.assertEqual(tf_val.shape, x_np.shape)
+            if np.issubdtype(dtype, np.complexfloating):
+              self.assertTrue(np.all(np.isnan(np.real(result))))
+            else:
+              self.assertTrue(np.all(np.isnan(result)))
 
+  @test_util.run_in_graph_and_eager_modes
   def testNonFiniteXReturnsInf(self):
     # inf x must evaluate to inf (not nan): the limit of a polynomial such as
     # p(x) = x + 1 as x -> inf is mathematically inf. NumPy returns nan here
@@ -1384,16 +1385,15 @@ class PolyvalTest(test.TestCase):
                               x=x_val,
                               degree=degree,
                               scalar=x_np.shape == ()):
-              with self.cached_session():
-                tf_val = math_ops.polyval([dtype(c) for c in coeffs], x_np)
-                result = self.evaluate(tf_val)
-                self.assertEqual(tf_val.shape, x_np.shape)
-                if np.issubdtype(dtype, np.complexfloating):
-                  # Complex inf arithmetic can produce nan imaginary parts
-                  # (e.g. (1+0j) * (inf+0j)); the real part must stay inf.
-                  self.assertTrue(np.all(np.isinf(np.real(result))))
-                else:
-                  self.assertTrue(np.all(np.isinf(result)))
+              tf_val = math_ops.polyval([dtype(c) for c in coeffs], x_np)
+              result = self.evaluate(tf_val)
+              self.assertEqual(tf_val.shape, x_np.shape)
+              if np.issubdtype(dtype, np.complexfloating):
+                # Complex inf arithmetic can produce nan imaginary parts
+                # (e.g. (1+0j) * (inf+0j)); the real part must stay inf.
+                self.assertTrue(np.all(np.isinf(np.real(result))))
+              else:
+                self.assertTrue(np.all(np.isinf(result)))
 
   def test_coeffs_raise(self):
     x = np.random.rand(2, 2).astype(np.float32)

From e77c2d472c64b47122781f37a3a232bb0781f9dd Mon Sep 17 00:00:00 2001
From: Boris Kozak <boriskozak@gmail.com>
Date: Thu, 8 Oct 2026 00:55:47 +0000
Subject: [PATCH 0241/1000] Convert remaining PolyvalTest methods to run in
 graph and eager modes

Per round-5 review: _runtest and its callers (testSimple), testBroadcast,
testEmpty, testSingleCoeffNanPropagation, testSingleCoeffShapeBroadcast now
use @test_util.run_in_graph_and_eager_modes instead of cached_session().
---
 .../kernel_tests/math_ops/cwise_ops_test.py   | 32 +++++++++----------
 1 file changed, 16 insertions(+), 16 deletions(-)

diff --git a/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py b/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
index 27d817cc630..80636502e4b 100644
--- a/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
@@ -1286,10 +1286,10 @@ class PolyvalTest(test.TestCase):
     x = np.random.rand(2, 2).astype(dtype)
     coeffs = [np.random.rand(2, 2).astype(dtype) for _ in range(degree + 1)]
     np_val = np.polyval(coeffs, x)
-    with self.cached_session():
-      tf_val = math_ops.polyval(coeffs, x)
-      self.assertAllClose(np_val, self.evaluate(tf_val))
+    tf_val = math_ops.polyval(coeffs, x)
+    self.assertAllClose(np_val, self.evaluate(tf_val))
 
+  @test_util.run_in_graph_and_eager_modes
   def testSimple(self):
     for dtype in [
         np.int32, np.float32, np.float64, np.complex64, np.complex128
@@ -1298,6 +1298,7 @@ class PolyvalTest(test.TestCase):
         with self.subTest(dtype=dtype, degree=degree):
           self._runtest(dtype, degree)
 
+  @test_util.run_in_graph_and_eager_modes
   def testBroadcast(self):
     dtype = np.float32
     degree = 3
@@ -1311,34 +1312,33 @@ class PolyvalTest(test.TestCase):
               for _ in range(degree + 1)
           ]
           np_val = np.polyval(coeffs, x)
-          with self.cached_session():
-            tf_val = math_ops.polyval(coeffs, x)
-            self.assertAllClose(np_val, self.evaluate(tf_val))
+          tf_val = math_ops.polyval(coeffs, x)
+          self.assertAllClose(np_val, self.evaluate(tf_val))
 
+  @test_util.run_in_graph_and_eager_modes
   def testEmpty(self):
     x = np.random.rand(2, 2).astype(np.float32)
     coeffs = []
     np_val = np.polyval(coeffs, x)
-    with self.cached_session():
-      tf_val = math_ops.polyval(coeffs, x)
-      self.assertAllClose(np_val, self.evaluate(tf_val))
+    tf_val = math_ops.polyval(coeffs, x)
+    self.assertAllClose(np_val, self.evaluate(tf_val))
 
+  @test_util.run_in_graph_and_eager_modes
   def testSingleCoeffNanPropagation(self):
     x = np.array([1.0, float("nan"), 3.0], dtype=np.float32)
     coeffs = [np.float32(2.0)]
     np_val = np.polyval(coeffs, x)
-    with self.cached_session():
-      tf_val = math_ops.polyval(coeffs, x)
-      self.assertAllClose(np_val, self.evaluate(tf_val))
+    tf_val = math_ops.polyval(coeffs, x)
+    self.assertAllClose(np_val, self.evaluate(tf_val))
 
+  @test_util.run_in_graph_and_eager_modes
   def testSingleCoeffShapeBroadcast(self):
     x = np.random.rand(3, 4).astype(np.float32)
     coeffs = [np.float32(5.0)]
     np_val = np.polyval(coeffs, x)
-    with self.cached_session():
-      tf_val = math_ops.polyval(coeffs, x)
-      self.assertEqual(tf_val.shape, x.shape)
-      self.assertAllClose(np_val, self.evaluate(tf_val))
+    tf_val = math_ops.polyval(coeffs, x)
+    self.assertEqual(tf_val.shape, x.shape)
+    self.assertAllClose(np_val, self.evaluate(tf_val))
 
   @test_util.run_in_graph_and_eager_modes
   def testSingleCoeffInfReturnsNan(self):

From 9b5ee2bea81d796cd66ea048f48b0847a10c1b17 Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Thu, 8 Oct 2026 09:35:05 +0530
Subject: [PATCH 0242/1000] bound uncompressed_bytes metadata by the snappy
 payload size in UncompressElement

Read the snappy uncompressed length before the first pass and charge every
byte count that sizes an allocation against it, cumulatively, so forged
metadata cannot size the nonmemcpyable scratch buffer or a string past
what the payload actually holds. Drops the two `< 0` checks, which could
never fire on a `uint64` proto field.
---
 tensorflow/core/data/compression_utils.cc     | 43 +++++++++------
 .../core/data/compression_utils_test.cc       | 54 +++++++++++++------
 2 files changed, 65 insertions(+), 32 deletions(-)

diff --git a/tensorflow/core/data/compression_utils.cc b/tensorflow/core/data/compression_utils.cc
index 1eb286e8c3b..1d847e05eca 100644
--- a/tensorflow/core/data/compression_utils.cc
+++ b/tensorflow/core/data/compression_utils.cc
@@ -151,6 +151,18 @@ absl::Status UncompressElement(const CompressedElement& compressed,
     return absl::InternalError(absl::StrCat(
         "Unsupported compressed element version: ", compressed.version()));
   }
+  const std::string& compressed_data = compressed.data();
+  size_t uncompressed_size;
+  if (!port::Snappy_GetUncompressedLength(
+          compressed_data.data(), compressed_data.size(), &uncompressed_size)) {
+    return absl::InternalError(absl::StrCat(
+        "Could not get snappy uncompressed length. Compressed data size: ",
+        compressed_data.size()));
+  }
+  // Byte counts that size an allocation are charged against this budget first,
+  // so forged metadata cannot allocate more than the payload holds or wrap the
+  // `size_t` totals below.
+  uint64_t remaining = uncompressed_size;
   int num_components = compressed.component_metadata_size();
   out->clear();
   out->reserve(num_components);
@@ -173,17 +185,20 @@ absl::Status UncompressElement(const CompressedElement& compressed,
     } else if (!DataTypeCanUseMemcpy(metadata.dtype())) {
       // Non-`memcpy`able components always serialize to a non-empty
       // `TensorProto`, even when the tensor has zero elements, so size the
-      // scratch buffer for them regardless of element count. A forged byte
-      // count is then caught by the snappy size reconciliation below.
+      // scratch buffer for them regardless of element count. The byte count is
+      // charged against the budget first, so a forged count can neither size
+      // the buffer past the uncompressed data nor wrap the running total.
       if (metadata.uncompressed_bytes_size() == 0) {
         return absl::InvalidArgumentError(
             "Missing uncompressed_bytes metadata for non-memcpyable tensor");
       }
-      if (metadata.uncompressed_bytes(0) < 0) {
+      const uint64_t bytes = metadata.uncompressed_bytes(0);
+      if (bytes > remaining) {
         return absl::InvalidArgumentError(
-            "uncompressed_bytes metadata cannot be negative");
+            "uncompressed_bytes metadata exceeds the uncompressed data size");
       }
-      total_nonmemcpyable_size += metadata.uncompressed_bytes(0);
+      remaining -= bytes;
+      total_nonmemcpyable_size += bytes;
     } else {
       int64_t num_elements = shape.num_elements();
       if (num_elements > 0 && metadata.uncompressed_bytes_size() == 0) {
@@ -229,12 +244,14 @@ absl::Status UncompressElement(const CompressedElement& compressed,
             ")"));
       }
       for (int j = 0; j < metadata.uncompressed_bytes_size(); ++j) {
-        if (metadata.uncompressed_bytes(j) < 0) {
+        const uint64_t bytes = metadata.uncompressed_bytes(j);
+        if (bytes > remaining) {
           return absl::InvalidArgumentError(
-              "uncompressed_bytes metadata cannot be negative");
+              "uncompressed_bytes metadata exceeds the uncompressed data size");
         }
-        flats.data()[j].resize(metadata.uncompressed_bytes(j));
-        iov.Add(flats.data()[j].mdata(), metadata.uncompressed_bytes(j));
+        remaining -= bytes;
+        flats.data()[j].resize(bytes);
+        iov.Add(flats.data()[j].mdata(), bytes);
       }
     } else {
       out->emplace_back();
@@ -244,14 +261,6 @@ absl::Status UncompressElement(const CompressedElement& compressed,
   }
 
   // Step 2: Uncompress into the iovec.
-  const std::string& compressed_data = compressed.data();
-  size_t uncompressed_size;
-  if (!port::Snappy_GetUncompressedLength(
-          compressed_data.data(), compressed_data.size(), &uncompressed_size)) {
-    return absl::InternalError(absl::StrCat(
-        "Could not get snappy uncompressed length. Compressed data size: ",
-        compressed_data.size()));
-  }
   if (uncompressed_size != static_cast<size_t>(iov.NumBytes())) {
     return absl::InternalError(absl::StrCat(
         "Uncompressed size mismatch. Snappy expects ", uncompressed_size,
diff --git a/tensorflow/core/data/compression_utils_test.cc b/tensorflow/core/data/compression_utils_test.cc
index 5e5bcad6a17..b6400cd1b95 100644
--- a/tensorflow/core/data/compression_utils_test.cc
+++ b/tensorflow/core/data/compression_utils_test.cc
@@ -49,8 +49,8 @@ TEST(CompressionUtilsTest, Exceeds4GB) {
 
 TEST(CompressionUtilsTest, ZeroElementNonMemcpyableComponent) {
   // If metadata claims a non-memcpyable component has 65536 uncompressed bytes
-  // when the compressed data only has 0 bytes, the size reconciliation against
-  // the iovec fails before any decompression or deserialization occurs.
+  // when the compressed data only holds 0 bytes, the byte count exhausts the
+  // budget and is rejected before the scratch buffer is sized from it.
   std::vector<Tensor> empty_element;
   CompressedElement compressed;
   TF_ASSERT_OK(CompressElement(empty_element, &compressed));
@@ -63,9 +63,10 @@ TEST(CompressionUtilsTest, ZeroElementNonMemcpyableComponent) {
   metadata->add_uncompressed_bytes(65536);
 
   std::vector<Tensor> element;
-  EXPECT_THAT(UncompressElement(compressed, &element),
-              absl_testing::StatusIs(error::INTERNAL,
-                                     HasSubstr("Uncompressed size mismatch")));
+  EXPECT_THAT(
+      UncompressElement(compressed, &element),
+      absl_testing::StatusIs(error::INVALID_ARGUMENT,
+                             HasSubstr("exceeds the uncompressed data size")));
 }
 
 TEST(CompressionUtilsTest, RoundTripEmptyVariantTensor) {
@@ -121,9 +122,9 @@ TEST(CompressionUtilsTest, MalformedTensorShapeString) {
 }
 
 TEST(CompressionUtilsTest, NegativeUncompressedBytes) {
-  // `uncompressed_bytes` is a signed proto field, so a negative count on a
-  // non-memcpyable component must be rejected before it is added to the
-  // unsigned scratch buffer size.
+  // `uncompressed_bytes` is a `uint64` field, so a forged -1 arrives as
+  // `UINT64_MAX`. It has to be rejected before the unsigned scratch buffer is
+  // sized from it.
   std::vector<Tensor> empty_element;
   CompressedElement compressed;
   TF_ASSERT_OK(CompressElement(empty_element, &compressed));
@@ -136,13 +137,16 @@ TEST(CompressionUtilsTest, NegativeUncompressedBytes) {
   metadata->add_uncompressed_bytes(-1);
 
   std::vector<Tensor> element;
-  EXPECT_THAT(UncompressElement(compressed, &element),
-              absl_testing::StatusIs(error::INVALID_ARGUMENT,
-                                     HasSubstr("cannot be negative")));
+  EXPECT_THAT(
+      UncompressElement(compressed, &element),
+      absl_testing::StatusIs(error::INVALID_ARGUMENT,
+                             HasSubstr("exceeds the uncompressed data size")));
 }
 
 TEST(CompressionUtilsTest, NegativeUncompressedBytesString) {
-  // The same holds for the per-string byte counts of a `DT_STRING` component.
+  // The same holds for the per-string byte counts of a `DT_STRING` component:
+  // -1 arrives as `UINT64_MAX` and must be rejected before the string is
+  // resized to it.
   std::vector<Tensor> empty_element;
   CompressedElement compressed;
   TF_ASSERT_OK(CompressElement(empty_element, &compressed));
@@ -155,9 +159,29 @@ TEST(CompressionUtilsTest, NegativeUncompressedBytesString) {
   metadata->add_uncompressed_bytes(-1);
 
   std::vector<Tensor> element;
-  EXPECT_THAT(UncompressElement(compressed, &element),
-              absl_testing::StatusIs(error::INVALID_ARGUMENT,
-                                     HasSubstr("cannot be negative")));
+  EXPECT_THAT(
+      UncompressElement(compressed, &element),
+      absl_testing::StatusIs(error::INVALID_ARGUMENT,
+                             HasSubstr("exceeds the uncompressed data size")));
+}
+
+TEST(CompressionUtilsTest, StringBytesExceedUncompressedSizeInTotal) {
+  // Each per-string count fits inside the uncompressed data on its own, but
+  // their sum does not, so only a cumulative budget rejects this.
+  std::vector<Tensor> element = {
+      CreateTensor<tstring>(TensorShape{2}, {"ab", "cd"})};
+  CompressedElement compressed;
+  TF_ASSERT_OK(CompressElement(element, &compressed));
+  CompressedComponentMetadata* metadata =
+      compressed.mutable_component_metadata(0);
+  metadata->set_uncompressed_bytes(0, 4);
+  metadata->set_uncompressed_bytes(1, 4);
+
+  std::vector<Tensor> round_trip;
+  EXPECT_THAT(
+      UncompressElement(compressed, &round_trip),
+      absl_testing::StatusIs(error::INVALID_ARGUMENT,
+                             HasSubstr("exceeds the uncompressed data size")));
 }
 
 std::vector<std::vector<Tensor>> TestCases() {

From 5a0e3ecd9db6478395920c471e0fad7b19fbb991 Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Wed, 7 Oct 2026 21:54:44 -0700
Subject: [PATCH 0243/1000] Convert NMS thresholds to the boxes' type in XLA

NonMaxSuppressionV3 and V4 allow thresholds of a different
floating-point type (T_threshold) than boxes and scores (T), and
tf.image.non_max_suppression passes float32 thresholds for float16
boxes. The XLA lowering added the IoU threshold to, and compared the
score threshold with, values of type T, so it failed to compile with
"Seen floating point types of different precisions" (#128614).

Convert the thresholds to T when their type differs, as the TensorFlow
kernels do.
---
 tensorflow/compiler/tests/image_ops_test.py   | 37 +++++++++++++++++++
 .../compiler/tf2xla/kernels/image_ops.cc      | 13 ++++++-
 2 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/tests/image_ops_test.py b/tensorflow/compiler/tests/image_ops_test.py
index 939b55b3511..0bf20c9490d 100644
--- a/tensorflow/compiler/tests/image_ops_test.py
+++ b/tensorflow/compiler/tests/image_ops_test.py
@@ -1022,6 +1022,43 @@ class NonMaxSuppressionTest(xla_test.XLATestCase):
       self.assertEqual(indices_tf.size, 3)
       self.assertAllClose(indices_tf[:3], [3, 0, 5])
 
+  def testNMSV3MixedThresholdDtype(self):
+    # Regression test for GitHub issue 128614: boxes and scores of one
+    # floating-point type with thresholds of another, which
+    # tf.image.non_max_suppression produces for float16 boxes, failed to
+    # compile because XLA doesn't mix floating-point types in one operation.
+    if np.float16 not in self.float_types:
+      self.skipTest("float16 is not supported on this device")
+    boxes_data = [[0, 0, 1, 1], [0, 0.1, 1, 1.1], [0, -0.1, 1, 0.9],
+                  [0, 10, 1, 11], [0, 10.1, 1, 11.1], [0, 100, 1, 101]]
+    scores_data = [0.9, 0.75, 0.6, 0.95, 0.5, 0.3]
+    for dtype, threshold_dtype in ((np.float16, np.float32),
+                                   (np.float32, np.float16)):
+      with self.subTest(
+          dtype=dtype.__name__,
+          threshold_dtype=threshold_dtype.__name__), self.session() as sess:
+        boxes = array_ops.placeholder(dtype, shape=[6, 4])
+        scores = array_ops.placeholder(dtype, shape=[6])
+        iou_threshold = array_ops.placeholder(threshold_dtype, shape=[])
+        score_threshold = array_ops.placeholder(threshold_dtype, shape=[])
+        with self.test_scope():
+          selected_indices = image_ops.non_max_suppression_v3(
+              boxes=boxes,
+              scores=scores,
+              max_output_size=6,
+              iou_threshold=iou_threshold,
+              score_threshold=score_threshold)
+        indices_tf = sess.run(
+            selected_indices,
+            feed_dict={
+                boxes: np.array(boxes_data, dtype=dtype),
+                scores: np.array(scores_data, dtype=dtype),
+                iou_threshold: 0.5,
+                score_threshold: 0.4
+            })
+        # Box 5 is below the score threshold.
+        self.assertAllEqual([3, 0], indices_tf)
+
   def testNMSV3EmptyInput(self):
     # Regression test for #117245: with no boxes the suppression loop was
     # built from zero-sized dimensions and segfaulted the compiler. The
diff --git a/tensorflow/compiler/tf2xla/kernels/image_ops.cc b/tensorflow/compiler/tf2xla/kernels/image_ops.cc
index 1c7cdca6a69..61ad71e0894 100644
--- a/tensorflow/compiler/tf2xla/kernels/image_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/image_ops.cc
@@ -491,8 +491,17 @@ class NonMaxSuppressionOp : public XlaOpKernel {
     OP_REQUIRES(context, output_size <= std::numeric_limits<int32_t>::max(),
                 absl::InvalidArgumentError(absl::StrCat(
                     "Need output_size <= kint32Max, got ", output_size)));
-    const xla::XlaOp score_thresh = context->Input("score_threshold");
-    const xla::XlaOp iou_thresh = context->Input("iou_threshold");
+    // The thresholds' type, T_threshold, can differ from T, the type of boxes
+    // and scores. Compare in T, as the TensorFlow kernels do, since XLA
+    // doesn't mix floating-point types in one operation.
+    xla::XlaOp score_thresh = context->Input("score_threshold");
+    if (context->InputXlaType("score_threshold") != scores_xla_type) {
+      score_thresh = xla::ConvertElementType(score_thresh, scores_xla_type);
+    }
+    xla::XlaOp iou_thresh = context->Input("iou_threshold");
+    if (context->InputXlaType("iou_threshold") != boxes_xla_type) {
+      iou_thresh = xla::ConvertElementType(iou_thresh, boxes_xla_type);
+    }
     xla::XlaBuilder* const builder = context->builder();
 
     if (num_boxes == 0) {

From 51ccd04e83010e3a2104a7e92c62df0f05d3ee27 Mon Sep 17 00:00:00 2001
From: Boris Kozak <boriskozak@gmail.com>
Date: Thu, 8 Oct 2026 05:21:59 +0000
Subject: [PATCH 0244/1000] Decorate test_coeffs_raise with
 run_in_graph_and_eager_modes

---
 tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py b/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
index 80636502e4b..cb1cfd617d3 100644
--- a/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
@@ -1395,6 +1395,7 @@ class PolyvalTest(test.TestCase):
               else:
                 self.assertTrue(np.all(np.isinf(result)))
 
+  @test_util.run_in_graph_and_eager_modes
   def test_coeffs_raise(self):
     x = np.random.rand(2, 2).astype(np.float32)
     coeffs = {}

From 8ff93ebea8bffeff04be78d73c9234ee10a0dad1 Mon Sep 17 00:00:00 2001
From: Bill Varcho <varcho@google.com>
Date: Wed, 7 Oct 2026 23:36:27 -0700
Subject: [PATCH 0245/1000] Bump Shardy commit to fix errors from renaming
 maximal -> singledevice

PiperOrigin-RevId: 995610713
---
 .../xla/third_party/shardy/temporary.patch    | 1419 -----------------
 .../xla/third_party/shardy/workspace.bzl      |    4 +-
 2 files changed, 2 insertions(+), 1421 deletions(-)

diff --git a/third_party/xla/third_party/shardy/temporary.patch b/third_party/xla/third_party/shardy/temporary.patch
index 97dea0bf1bc..e69de29bb2d 100644
--- a/third_party/xla/third_party/shardy/temporary.patch
+++ b/third_party/xla/third_party/shardy/temporary.patch
@@ -1,1419 +0,0 @@
-diff --git a/third_party/llvm/generated.patch b/third_party/llvm/generated.patch
-index 875e6c8a..e2237d49 100644
---- a/third_party/llvm/generated.patch
-+++ b/third_party/llvm/generated.patch
-@@ -1,19 +1,459 @@
- Auto generated patch. Do not edit or delete it, even if empty.
--diff -ruN --strip-trailing-cr a/llvm/docs/WritingAnLLVMPass.md b/llvm/docs/WritingAnLLVMPass.md
----- a/llvm/docs/WritingAnLLVMPass.md
--+++ b/llvm/docs/WritingAnLLVMPass.md
--@@ -794,14 +794,6 @@
-- 
-- ```console
-- $ gdb opt
---GNU gdb 5.0
---Copyright 2000 Free Software Foundation, Inc.
---GDB is free software, covered by the GNU General Public License, and you are
---welcome to change it and/or distribute copies of it under certain conditions.
---Type "show copying" to see the conditions.
---There is absolutely no warranty for GDB.  Type "show warranty" for details.
---This GDB was configured as "sparc-sun-solaris2.6"...
---(gdb)
-- ```
-- 
-- Note that {program}`opt` has a lot of debugging information in it, so it takes
-+diff -ruN --strip-trailing-cr a/llvm/include/llvm/IR/Metadata.h b/llvm/include/llvm/IR/Metadata.h
-+--- a/llvm/include/llvm/IR/Metadata.h
-++++ b/llvm/include/llvm/IR/Metadata.h
-+@@ -397,18 +397,8 @@
-+   using OwnerTy = MetadataTracking::OwnerTy;
-+ 
-+ private:
-+-  struct UseEntry {
-+-    void *Ref = nullptr;
-+-    OwnerTy Owner = nullptr;
-+-  };
-+-
-+-  static constexpr unsigned IndexThreshold = 32;
-+-  // Tracked uses of this. Dropping one moves the last entry into its slot, so
-+-  // the order is deterministic but not the order they were added in.
-+-  SmallVector<UseEntry, 4> UseMap;
-+-  // Lazily allocated map from Ref to its index in UseMap for large use lists.
-+-  using IndexMapTy = DenseMap<void *, unsigned>;
-+-  std::unique_ptr<IndexMapTy> IndexMap;
-++  uint64_t NextIndex = 0;
-++  SmallDenseMap<void *, std::pair<OwnerTy, uint64_t>, 4> UseMap;
-+ 
-+ protected:
-+   ~ReplaceableUses() {
-+@@ -439,7 +429,6 @@
-+   unsigned getNumUses() const { return UseMap.size(); }
-+ 
-+ private:
-+-  UseEntry *findRef(void *Ref);
-+   void addRef(void *Ref, OwnerTy Owner);
-+   void dropRef(void *Ref);
-+   void moveRef(void *Ref, void *New, const Metadata &MD);
-+diff -ruN --strip-trailing-cr a/llvm/lib/IR/Metadata.cpp b/llvm/lib/IR/Metadata.cpp
-+--- a/llvm/lib/IR/Metadata.cpp
-++++ b/llvm/lib/IR/Metadata.cpp
-+@@ -258,96 +258,82 @@
-+ }
-+ 
-+ SmallVector<Metadata *> ReplaceableUses::getAllArgListUsers() {
-+-  SmallVector<Metadata *> MDUsers;
-+-  for (const auto &[Ref, Owner] : UseMap) {
-++  SmallVector<std::pair<OwnerTy, uint64_t> *> MDUsersWithID;
-++  for (auto Pair : UseMap) {
-++    OwnerTy Owner = Pair.second.first;
-+     if (Owner.isNull())
-+       continue;
-+     if (!isa<Metadata *>(Owner))
-+       continue;
-+     Metadata *OwnerMD = cast<Metadata *>(Owner);
-+     if (OwnerMD->getMetadataID() == Metadata::DIArgListKind)
-+-      MDUsers.push_back(OwnerMD);
-++      MDUsersWithID.push_back(&UseMap[Pair.first]);
-+   }
-++  llvm::sort(MDUsersWithID, [](auto UserA, auto UserB) {
-++    return UserA->second < UserB->second;
-++  });
-++  SmallVector<Metadata *> MDUsers;
-++  for (auto *UserWithID : MDUsersWithID)
-++    MDUsers.push_back(cast<Metadata *>(UserWithID->first));
-+   return MDUsers;
-+ }
-+ 
-+ SmallVector<DbgVariableRecord *>
-+ ReplaceableUses::getAllDbgVariableRecordUsers() {
-++  SmallVector<std::pair<OwnerTy, uint64_t> *> DVRUsersWithID;
-++  for (auto Pair : UseMap) {
-++    OwnerTy Owner = Pair.second.first;
-++    if (Owner.isNull())
-++      continue;
-++    if (!isa<DebugValueUser *>(Owner))
-++      continue;
-++    DVRUsersWithID.push_back(&UseMap[Pair.first]);
-++  }
-+   // Order DbgVariableRecord users in reverse-creation order. Normal dbg.value
-+   // users of MetadataAsValues are ordered by their UseList, i.e. reverse order
-+   // of when they were added: we need to replicate that here. The structure of
-+   // debug-info output depends on the ordering of intrinsics, thus we need
-+   // to keep them consistent for comparisons sake.
-++  llvm::sort(DVRUsersWithID, [](auto UserA, auto UserB) {
-++    return UserA->second > UserB->second;
-++  });
-+   SmallVector<DbgVariableRecord *> DVRUsers;
-+-  for (const auto &[Ref, Owner] : reverse(UseMap)) {
-+-    if (Owner.isNull())
-+-      continue;
-+-    if (!isa<DebugValueUser *>(Owner))
-+-      continue;
-+-    DVRUsers.push_back(cast<DebugValueUser *>(Owner)->getUser());
-+-  }
-++  for (auto UserWithID : DVRUsersWithID)
-++    DVRUsers.push_back(cast<DebugValueUser *>(UserWithID->first)->getUser());
-+   return DVRUsers;
-+ }
-+ 
-+-ReplaceableUses::UseEntry *ReplaceableUses::findRef(void *Ref) {
-+-  if (!IndexMap) {
-+-    // Search backward for temporal locality; recently added references are
-+-    // often dropped or moved first.
-+-    for (UseEntry &U : reverse(UseMap))
-+-      if (U.Ref == Ref)
-+-        return &U;
-+-    return nullptr;
-+-  }
-+-
-+-  auto It = IndexMap->find(Ref);
-+-  return It == IndexMap->end() ? nullptr : &UseMap[It->second];
-+-}
-+-
-+ void ReplaceableUses::addRef(void *Ref, OwnerTy Owner) {
-+-  assert(Ref && "Expected live reference");
-+-  assert(!findRef(Ref) && "Reference already tracked");
-+-  unsigned NewIdx = UseMap.size();
-+-  UseMap.push_back({Ref, Owner});
-+-  if (IndexMap) {
-+-    (*IndexMap)[Ref] = NewIdx;
-+-  } else if (UseMap.size() > IndexThreshold) {
-+-    // Build the index map once UseMap grows past the threshold.
-+-    IndexMap = std::make_unique<IndexMapTy>();
-+-    for (unsigned I = 0, E = UseMap.size(); I != E; ++I)
-+-      (*IndexMap)[UseMap[I].Ref] = I;
-+-  }
-++  bool WasInserted =
-++      UseMap.insert(std::make_pair(Ref, std::make_pair(Owner, NextIndex)))
-++          .second;
-++  (void)WasInserted;
-++  assert(WasInserted && "Expected to add a reference");
-++
-++  ++NextIndex;
-++  assert(NextIndex != 0 && "Unexpected overflow");
-+ }
-+ 
-+ void ReplaceableUses::dropRef(void *Ref) {
-+-  UseEntry *Entry = findRef(Ref);
-+-  assert(Entry && "Expected to find Ref");
-+-  if (IndexMap)
-+-    IndexMap->erase(Ref);
-+-  if (Entry != &UseMap.back()) {
-+-    *Entry = UseMap.back();
-+-    if (IndexMap)
-+-      (*IndexMap)[Entry->Ref] = Entry - UseMap.begin();
-+-  }
-+-  UseMap.pop_back();
-+-  if (IndexMap && UseMap.size() <= IndexThreshold / 2)
-+-    IndexMap.reset();
-++  bool WasErased = UseMap.erase(Ref);
-++  (void)WasErased;
-++  assert(WasErased && "Expected to drop a reference");
-+ }
-+ 
-+ void ReplaceableUses::moveRef(void *Ref, void *New, const Metadata &MD) {
-+-  assert(!findRef(New) && "Cannot move to an existing reference");
-+-  UseEntry *Entry = findRef(Ref);
-+-  assert(Entry && "Expected to move a reference");
-+-  if (IndexMap) {
-+-    IndexMap->erase(Ref);
-+-    (*IndexMap)[New] = Entry - UseMap.begin();
-+-  }
-+-  Entry->Ref = New;
-++  auto I = UseMap.find(Ref);
-++  assert(I != UseMap.end() && "Expected to move a reference");
-++  auto OwnerAndIndex = I->second;
-++  UseMap.erase(I);
-++  bool WasInserted = UseMap.insert(std::make_pair(New, OwnerAndIndex)).second;
-++  (void)WasInserted;
-++  assert(WasInserted && "Expected to add a reference");
-+ 
-+   // Check that the references are direct if there's no owner.
-+   (void)MD;
-+-  assert((Entry->Owner || *static_cast<Metadata **>(Ref) == &MD) &&
-++  assert((OwnerAndIndex.first || *static_cast<Metadata **>(Ref) == &MD) &&
-+          "Reference without owner must be direct");
-+-  assert((Entry->Owner || *static_cast<Metadata **>(New) == &MD) &&
-++  assert((OwnerAndIndex.first || *static_cast<Metadata **>(New) == &MD) &&
-+          "Reference without owner must be direct");
-+ }
-+ 
-+@@ -360,11 +346,14 @@
-+   auto &Store = Context.pImpl->ValuesAsMetadata;
-+   auto I = Store.find(&C);
-+   ValueAsMetadata *MD = I->second;
-++  using UseTy =
-++      std::pair<void *, std::pair<MetadataTracking::OwnerTy, uint64_t>>;
-+   // Copy out uses and update value of Constant used by debug info metadata with
-+-  // poison below.
-+-  SmallVector<UseEntry, 4> Uses = MD->UseMap;
-++  // poison below
-++  SmallVector<UseTy, 8> Uses(MD->UseMap.begin(), MD->UseMap.end());
-+ 
-+-  for (const auto &[Ref, Owner] : Uses) {
-++  for (const auto &Pair : Uses) {
-++    MetadataTracking::OwnerTy Owner = Pair.second.first;
-+     if (!Owner)
-+       continue;
-+     // Check for MetadataAsValue.
-+@@ -380,7 +369,7 @@
-+       continue;
-+     if (isa<DINode>(OwnerMD)) {
-+       OwnerMD->handleChangedOperand(
-+-          Ref, ValueAsMetadata::get(PoisonValue::get(C.getType())));
-++          Pair.first, ValueAsMetadata::get(PoisonValue::get(C.getType())));
-+     }
-+   }
-+ }
-+@@ -390,21 +379,25 @@
-+     return;
-+ 
-+   // Copy out uses since UseMap will get touched below.
-+-  SmallVector<UseEntry, 4> Uses = UseMap;
-+-
-+-  for (const auto &[Ref, Owner] : Uses) {
-++  using UseTy = std::pair<void *, std::pair<OwnerTy, uint64_t>>;
-++  SmallVector<UseTy, 8> Uses(UseMap.begin(), UseMap.end());
-++  llvm::sort(Uses, [](const UseTy &L, const UseTy &R) {
-++    return L.second.second < R.second.second;
-++  });
-++  for (const auto &Pair : Uses) {
-+     // Check that this Ref hasn't disappeared after RAUW (when updating a
-+     // previous Ref).
-+-    if (!findRef(Ref))
-++    if (!UseMap.count(Pair.first))
-+       continue;
-+ 
-++    OwnerTy Owner = Pair.second.first;
-+     if (!Owner) {
-+       // Update unowned tracking references directly.
-+-      Metadata *&DirectRef = *static_cast<Metadata **>(Ref);
-+-      dropRef(Ref);
-+-      DirectRef = MD;
-++      Metadata *&Ref = *static_cast<Metadata **>(Pair.first);
-++      Ref = MD;
-+       if (MD)
-+-        MetadataTracking::track(DirectRef);
-++        MetadataTracking::track(Ref);
-++      UseMap.erase(Pair.first);
-+       continue;
-+     }
-+ 
-+@@ -415,7 +408,7 @@
-+     }
-+ 
-+     if (auto *DVU = dyn_cast<DebugValueUser *>(Owner)) {
-+-      DVU->handleChangedValue(Ref, MD);
-++      DVU->handleChangedValue(Pair.first, MD);
-+       continue;
-+     }
-+ 
-+@@ -424,7 +417,7 @@
-+     switch (OwnerMD->getMetadataID()) {
-+ #define HANDLE_METADATA_LEAF(CLASS)                                            \
-+   case Metadata::CLASS##Kind:                                                  \
-+-    cast<CLASS>(OwnerMD)->handleChangedOperand(Ref, MD);                       \
-++    cast<CLASS>(OwnerMD)->handleChangedOperand(Pair.first, MD);                \
-+     continue;
-+ #include "llvm/IR/Metadata.def"
-+     default:
-+@@ -440,16 +433,18 @@
-+ 
-+   if (!ResolveUsers) {
-+     UseMap.clear();
-+-    IndexMap.reset();
-+     return;
-+   }
-+ 
-+-  // Move uses out since UseMap could get touched below.
-+-  SmallVector<UseEntry, 4> Uses = std::move(UseMap);
-++  // Copy out uses since UseMap could get touched below.
-++  using UseTy = std::pair<void *, std::pair<OwnerTy, uint64_t>>;
-++  SmallVector<UseTy, 8> Uses(UseMap.begin(), UseMap.end());
-++  llvm::sort(Uses, [](const UseTy &L, const UseTy &R) {
-++    return L.second.second < R.second.second;
-++  });
-+   UseMap.clear();
-+-  IndexMap.reset();
-+-  for (const auto &U : Uses) {
-+-    auto Owner = U.Owner;
-++  for (const auto &Pair : Uses) {
-++    auto Owner = Pair.second.first;
-+     if (!Owner)
-+       continue;
-+     if (!isa<Metadata *>(Owner))
-+diff -ruN --strip-trailing-cr a/llvm/lib/ObjCopy/ELF/ELFObject.cpp b/llvm/lib/ObjCopy/ELF/ELFObject.cpp
-+--- a/llvm/lib/ObjCopy/ELF/ELFObject.cpp
-++++ b/llvm/lib/ObjCopy/ELF/ELFObject.cpp
-+@@ -1709,11 +1709,7 @@
-+   case SHT_REL:
-+   case SHT_RELA:
-+   case SHT_CREL:
-+-    // SHF_ALLOC relocations of an executable or shared object are copied
-+-    // verbatim. Relocatable files relocations are usually non-ALLOC, but Linux
-+-    // livepatch modules set SHF_ALLOC on .klp.rela.* to keep these static
-+-    // relocations in memory.
-+-    if ((Shdr.sh_flags & SHF_ALLOC) && ElfFile.getHeader().e_type != ET_REL) {
-++    if (Shdr.sh_flags & SHF_ALLOC) {
-+       if (Expected<ArrayRef<uint8_t>> Data = ElfFile.getSectionContents(Shdr))
-+         return Obj.addSection<DynamicRelocationSection>(*Data);
-+       else
-+diff -ruN --strip-trailing-cr a/llvm/lib/ObjCopy/ELF/ELFObject.h b/llvm/lib/ObjCopy/ELF/ELFObject.h
-+--- a/llvm/lib/ObjCopy/ELF/ELFObject.h
-++++ b/llvm/lib/ObjCopy/ELF/ELFObject.h
-+@@ -38,7 +38,6 @@
-+ class SymbolTableSection;
-+ class RelocationSection;
-+ class DynamicRelocationSection;
-+-class DynamicSymbolTableSection;
-+ class GnuDebugLinkSection;
-+ class GroupSection;
-+ class SectionIndexSection;
-+@@ -876,15 +875,10 @@
-+ // and another which handles the symbol table type. The symbol table type is
-+ // taken as a type parameter to the class (see RelocSectionWithSymtabBase).
-+ class RelocationSectionBase : public SectionBase {
-+-  const bool Dynamic;
-+-
-+ protected:
-+   SectionBase *SecToApplyRel = nullptr;
-+ 
-+-  explicit RelocationSectionBase(bool Dynamic) : Dynamic(Dynamic) {}
-+-
-+ public:
-+-  bool isDynamic() const { return Dynamic; }
-+   const SectionBase *getSection() const { return SecToApplyRel; }
-+   void setSection(SectionBase *Sec) { SecToApplyRel = Sec; }
-+ 
-+@@ -903,9 +897,7 @@
-+   void setSymTab(SymTabType *SymTab) { Symbols = SymTab; }
-+ 
-+ protected:
-+-  RelocSectionWithSymtabBase()
-+-      : RelocationSectionBase(
-+-            std::is_same_v<SymTabType, DynamicSymbolTableSection>) {}
-++  RelocSectionWithSymtabBase() = default;
-+ 
-+   SymTabType *Symbols = nullptr;
-+ 
-+@@ -936,8 +928,9 @@
-+   const Object &getObject() const { return Obj; }
-+ 
-+   static bool classof(const SectionBase *S) {
-+-    return RelocationSectionBase::classof(S) &&
-+-           !static_cast<const RelocationSectionBase *>(S)->isDynamic();
-++    if (S->OriginalFlags & ELF::SHF_ALLOC)
-++      return false;
-++    return RelocationSectionBase::classof(S);
-+   }
-+ };
-+ 
-+@@ -1022,8 +1015,9 @@
-+       function_ref<bool(const SectionBase *)> ToRemove) override;
-+ 
-+   static bool classof(const SectionBase *S) {
-+-    return RelocationSectionBase::classof(S) &&
-+-           static_cast<const RelocationSectionBase *>(S)->isDynamic();
-++    if (!(S->OriginalFlags & ELF::SHF_ALLOC))
-++      return false;
-++    return S->OriginalType == ELF::SHT_REL || S->OriginalType == ELF::SHT_RELA;
-+   }
-+ };
-+ 
-+diff -ruN --strip-trailing-cr a/llvm/test/tools/llvm-objcopy/ELF/reloc-alloc-symtab.test b/llvm/test/tools/llvm-objcopy/ELF/reloc-alloc-symtab.test
-+--- a/llvm/test/tools/llvm-objcopy/ELF/reloc-alloc-symtab.test
-++++ b/llvm/test/tools/llvm-objcopy/ELF/reloc-alloc-symtab.test
-+@@ -1,42 +0,0 @@
-+-## In a relocatable file, an SHF_ALLOC relocation section (e.g. .klp.rela.* in
-+-## Linux livepatch modules) is a static relocation section: its symbol indexes
-+-## are updated and the symbols it references are kept.
-+-
-+-# RUN: yaml2obj %s -o %t
-+-# RUN: llvm-strip --strip-unneeded %t -o %t1
-+-# RUN: llvm-readelf -Sr -s %t1 | FileCheck %s
-+-# RUN: yaml2obj -DTYPE=SHT_CREL %s -o %t.crel
-+-# RUN: llvm-strip --strip-unneeded %t.crel -o %t1.crel
-+-# RUN: llvm-readelf -Sr -s %t1.crel | FileCheck %s
-+-
-+-# CHECK:      [ 2] .klp.rela.vmlinux.text {{RELA|CREL}} {{.*}} AI  3   1
-+-# CHECK-NEXT: [ 3] .symtab
-+-# CHECK:      R_X86_64_PC32 {{.*}} referenced + 0
-+-# CHECK:      Symbol table '.symtab' contains 2 entries:
-+-# CHECK:      1: {{.*}} UND referenced
-+-
-+-!ELF
-+-FileHeader:
-+-  Class:   ELFCLASS64
-+-  Data:    ELFDATA2LSB
-+-  Type:    ET_REL
-+-  Machine: EM_X86_64
-+-Sections:
-+-  - Name:    .text
-+-    Type:    SHT_PROGBITS
-+-    Flags:   [ SHF_ALLOC, SHF_EXECINSTR ]
-+-    Content: "e800000000"
-+-  - Name:    .klp.rela.vmlinux.text
-+-    Type:    [[TYPE=SHT_RELA]]
-+-    Flags:   [ SHF_ALLOC, SHF_INFO_LINK ]
-+-    Link:    .symtab
-+-    Info:    .text
-+-    Relocations:
-+-      - Offset: 1
-+-        Symbol: referenced
-+-        Type:   R_X86_64_PC32
-+-Symbols:
-+-  - Name:    unreferenced
-+-    Section: .text
-+-  - Name:    referenced
-+-    Binding: STB_GLOBAL
-+diff -ruN --strip-trailing-cr a/llvm/test/tools/llvm-objcopy/ELF/reloc-no-symtab.test b/llvm/test/tools/llvm-objcopy/ELF/reloc-no-symtab.test
-+--- a/llvm/test/tools/llvm-objcopy/ELF/reloc-no-symtab.test
-++++ b/llvm/test/tools/llvm-objcopy/ELF/reloc-no-symtab.test
-+@@ -10,16 +10,11 @@
-+ # RUN: llvm-readobj --sections %t.original | FileCheck %s
-+ # RUN: llvm-readobj --sections %t.stripped | FileCheck %s
-+ 
-+-## In a relocatable file, .rela.plt is a static relocation section.
-+-# RUN: yaml2obj -DTYPE=ET_REL %s -o %t.rel
-+-# RUN: llvm-strip %t.rel -o %t.rel.stripped
-+-# RUN: llvm-readobj --sections %t.rel.stripped | FileCheck %s
-+-
-+ !ELF
-+ FileHeader:
-+   Class:           ELFCLASS64
-+   Data:            ELFDATA2LSB
-+-  Type:            [[TYPE=ET_EXEC]]
-++  Type:            ET_REL
-+   Machine:         EM_X86_64
-+ Sections:
-+   - Name:          .text
-+diff -ruN --strip-trailing-cr a/llvm/unittests/IR/MetadataTest.cpp b/llvm/unittests/IR/MetadataTest.cpp
-+--- a/llvm/unittests/IR/MetadataTest.cpp
-++++ b/llvm/unittests/IR/MetadataTest.cpp
-+@@ -5355,29 +5355,6 @@
-+   EXPECT_TRUE(!MD);
-+ }
-+ 
-+-// Once the use list is indexed, look each changed reference up again right
-+-// away: a later swap in dropRef would repair a stale index entry.
-+-TEST_F(TrackingMDRefTest, LargeUseList) {
-+-  auto Temp = MDTuple::getTemporary(Context, {});
-+-  TrackingMDRef Refs[41];
-+-  for (unsigned I = 0; I != 40; ++I)
-+-    Refs[I].reset(Temp.get());
-+-  Refs[0].reset();
-+-  Refs[0].reset(Temp.get());
-+-  Refs[0].reset();
-+-  Refs[40] = std::move(Refs[1]);
-+-  Refs[1].reset(Temp.get());
-+-  Refs[40].reset();
-+-  Refs[39].reset();
-+-  EXPECT_EQ(38u, Temp->getNumTemporaryUses());
-+-
-+-  auto *N = MDTuple::getDistinct(Context, {});
-+-  Temp->replaceAllUsesWith(N);
-+-  EXPECT_EQ(0u, Temp->getNumTemporaryUses());
-+-  for (unsigned I = 1; I != 39; ++I)
-+-    EXPECT_EQ(N, Refs[I].get()) << I;
-+-}
-+-
-+ TEST(NamedMDNodeTest, Search) {
-+   LLVMContext Context;
-+   ConstantAsMetadata *C =
-diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
-index 79e40910..9de7c05a 100644
---- a/third_party/llvm/workspace.bzl
-+++ b/third_party/llvm/workspace.bzl
-@@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
- 
- def repo(name):
-     """Imports LLVM."""
--    LLVM_COMMIT = "018a9e4a74ba5d383351bda907a8b116bc5d8c47"
--    LLVM_SHA256 = "67463f6d054b8880cbde059a8d02ca6ca4649312f60cfe98a30401ac17fe0b39"
-+    LLVM_COMMIT = "62e30e3326d9b4e3948d242ccfee7aa4fc030ef3"
-+    LLVM_SHA256 = "478ae34c463f28708582082098aaf0e79f431c2719f5d297bd546051dbeae0dd"
- 
-     tf_http_archive(
-         name = name,
-diff --git a/third_party/stablehlo/temporary.patch b/third_party/stablehlo/temporary.patch
-index 49cfe5b6..c38c5147 100644
---- a/third_party/stablehlo/temporary.patch
-+++ b/third_party/stablehlo/temporary.patch
-@@ -835,6 +835,863 @@ diff --ruN a/stablehlo/build_tools/lit_test_suite.bzl b/stablehlo/build_tools/li
-              tags = tags,
-              size = size_per_test,
-          )
-+diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
-+--- stablehlo/docs/spec.md
-++++ stablehlo/docs/spec.md
-+@@ -2380,6 +2380,7 @@
-+ in order to be able to express windows of `lhs`:
-+ 
-+ <!-- markdownlint-disable line-length -->
-++
-+ * `lhs_window_dimensions = lhs_shape(dim(lhs, input_batch_dimension), dim(rhs, kernel_spatial_dimensions), dim(lhs, input_feature_dimension))`.
-+ * `lhs_window_strides = lhs_shape(1, window_strides, 1)`.
-+ * `lhs_padding = lhs_shape([0, 0], padding, [0, 0])`.
-+@@ -2476,6 +2477,7 @@
-+ #### Constraints
-+ 
-+ <!-- markdownlint-disable line-length -->
-++
-+ * (C1) `N = rank(lhs) = rank(rhs)`.
-+ * (C2) `size(window_strides) = N - 2`.
-+ * (C3) `0 < window_strides`.
-+@@ -2808,6 +2810,7 @@
-+ More formally, `result[result_index] = dot_product`, where:
-+ 
-+ <!-- markdownlint-disable line-length -->
-++
-+ * `lhs_result_dimensions = [d for d in axes(lhs) and d not in lhs_batching_dimensions and d not in lhs_contracting_dimensions]`.
-+ * `rhs_result_dimensions = [d for d in axes(rhs) and d not in rhs_batching_dimensions and d not in rhs_contracting_dimensions]`.
-+ * `result_batching_index + result_lhs_index + result_rhs_index = result_index`
-+@@ -3148,6 +3151,7 @@
-+ #### Constraints
-+ 
-+ <!-- markdownlint-disable line-length -->
-++
-+ * (C1) `N = rank(lhs) = rank(rhs)`.
-+ * (C2) `size(window_strides) = N - 2`.
-+ * (C3) `0 < window_strides`.
-+@@ -3626,6 +3630,49 @@
-+ 
-+ &nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/dynamic_update_slice.mlir)
-+ 
-++### exp2
-++
-++#### Semantics
-++
-++Performs element-wise base-2 exponential operation on `operand` tensor and
-++produces a `result` tensor. Depending on the element type, does the following:
-++
-++*   For floats: `exp2` from IEEE-754.
-++*   For complex numbers: complex base-2 exponential.
-++*   For quantized types: `dequantize_op_quantize(exp2, operand, type(result))`.
-++
-++#### Inputs
-++
-++| Label | Name              | Type                 | Constraints |
-++| ----- | ----------------- | -------------------- | ----------- |
-++| (I1)  | `operand`         | tensor of            | (C1)        |
-++:       :                   : floating-point or    :             :
-++:       :                   : complex type or      :             :
-++:       :                   : per-tensor quantized :             :
-++:       :                   : tensor               :             :
-++| (I2)  | `result_accuracy` | optional             |             |
-++:       :                   : `ResultAccuracyAttr` :             :
-++:       :                   : (default `DEFAULT`)  :             :
-++
-++#### Outputs
-++
-++| Name     | Type                                        | Constraints |
-++| -------- | ------------------------------------------- | ----------- |
-++| `result` | tensor of floating-point or complex type or | (C1)        |
-++:          : per-tensor quantized tensor                 :             :
-++
-++#### Constraints
-++
-++*   (C1) `baseline_type(operand) = baseline_type(result)`.
-++
-++#### Examples
-++
-++```mlir
-++// %operand: [[0.0, 1.0], [2.0, 3.0]]
-++%result = "stablehlo.exp2"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
-++// %result: [[1.0, 2.0], [4.0, 8.0]]
-++```
-++
-+ ### exponential
-+ 
-+ #### Semantics
-+@@ -3798,6 +3845,7 @@
-+ * (C3) `1 <= size(fft_length) <= 3`.
-+ * (C4) If among `operand` and `result`, there is a tensor `real` of a
-+ floating-point type, then `shape(real)[-size(fft_length):] = fft_length`.
-++
-+ * (C5) `shape(result) = shape(operand)` except for:
-+   * If `fft_type = RFFT`,
-+     `dim(result, -1) = dim(operand, -1) = 0 ? 0 : dim(operand, -1) / 2 + 1`.
-+@@ -3866,6 +3914,7 @@
-+ More formally, `result[result_index] = operand[operand_index]` where:
-+ 
-+ <!-- markdownlint-disable line-length -->
-++
-+ * `batch_dims = [d for d in axes(result) and d not in offset_dims]`.
-+ * `batch_index = result_index[batch_dims...]`.
-+ * `start_index` is defined as:
-+@@ -4355,6 +4404,49 @@
-+ 
-+ &nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/log.mlir)
-+ 
-++### log2
-++
-++#### Semantics
-++
-++Performs element-wise base-2 logarithm operation on `operand` tensor and
-++produces a `result` tensor. Depending on the element type, does the following:
-++
-++*   For floats: `log2` from IEEE-754.
-++*   For complex numbers: complex base-2 logarithm.
-++*   For quantized types: `dequantize_op_quantize(log2, operand, type(result))`.
-++
-++#### Inputs
-++
-++| Label | Name              | Type                 | Constraints |
-++| ----- | ----------------- | -------------------- | ----------- |
-++| (I1)  | `operand`         | tensor of            | (C1)        |
-++:       :                   : floating-point or    :             :
-++:       :                   : complex type or      :             :
-++:       :                   : per-tensor quantized :             :
-++:       :                   : tensor               :             :
-++| (I2)  | `result_accuracy` | optional             |             |
-++:       :                   : `ResultAccuracyAttr` :             :
-++:       :                   : (default `DEFAULT`)  :             :
-++
-++#### Outputs
-++
-++| Name     | Type                                        | Constraints |
-++| -------- | ------------------------------------------- | ----------- |
-++| `result` | tensor of floating-point or complex type or | (C1)        |
-++:          : per-tensor quantized tensor                 :             :
-++
-++#### Constraints
-++
-++*   (C1) `baseline_type(operand) = baseline_type(result)`.
-++
-++#### Examples
-++
-++```mlir
-++// %operand: [[1.0, 2.0], [4.0, 8.0]]
-++%result = "stablehlo.log2"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
-++// %result: [[0.0, 1.0], [2.0, 3.0]]
-++```
-++
-+ ### log_plus_one
-+ 
-+ #### Semantics
-+@@ -5358,6 +5450,7 @@
-+ #### Constraints
-+ 
-+ <!-- markdownlint-disable line-length -->
-++
-+ * (C1) `0 < size(inputs) = size(init_values) = size(results) = N`.
-+ * (C2) `same(shape(inputs...))`.
-+ * (C3) `element_type(inputs...) = element_type(init_values...)`.
-+@@ -5897,6 +5990,7 @@
-+ #### Constraints
-+ 
-+ <!-- markdownlint-disable line-length -->
-++
-+ * (C1) `same(shape(inputs...))`.
-+ * (C2) `rank(inputs[0]) = size(update_window_dims) + size(inserted_window_dims) + size(input_batching_dims)`.
-+ * (C3) `same(shape(updates...))`.
-+@@ -6081,6 +6175,7 @@
-+    ([#731](https://github.com/openxla/stablehlo/issues/731)).
-+ * `result[result_index] = reduce([source_values], [init_value], [0], scatter)`
-+  where:
-++
-+   * `source_values = [source[source_index] for source_index in
-+    source_indices]`.
-+   * `selected_index(source_index) = operand_index` if
-+@@ -6111,6 +6206,7 @@
-+ #### Constraints
-+ 
-+ <!-- markdownlint-disable line-length -->
-++
-+ * (C1) `element_type(operand) = element_type(source)`.
-+ * (C2) `shape(source) = num_windows` where:
-+   * `padded_operand_shape = padding[:, 0] + shape(operand) + padding[:, 1]`.
-+@@ -7501,6 +7597,7 @@
-+ syntax ([ISO/IEC 14977:1996](https://www.iso.org/standard/26153.html),
-+ [Wikipedia](https://en.wikipedia.org/wiki/Extended_Backus%E2%80%93Naur_form)),
-+ with two modifications: 1) rules are defined using `::=` rather than `=`,
-++
-+ 2) concatenation is expressed using juxtaposition rather than `,`.
-+ 
-+ For describing semantics (i.e. within "Types", "Constants" and "Ops" sections),
-+@@ -7571,6 +7668,7 @@
-+ ### Names
-+ 
-+ In formulas, lexical scope includes: 1) global functions, 2) member definitions,
-++
-+ 3) local definitions. The list of global functions is provided below. The list
-+ of element definitions depends on the program element that the notation is
-+ applied to:
-+diff --ruN a/stablehlo/rfcs/20260821-low-precision-complex.md b/stablehlo/rfcs/20260821-low-precision-complex.md
-+--- stablehlo/rfcs/20260821-low-precision-complex.md
-++++ stablehlo/rfcs/20260821-low-precision-complex.md
-+@@ -0,0 +1,599 @@
-++# [RFC] Support f16- and bf16-based complex types in StableHLO
-++
-++Status: In Review<br/>
-++Initial version: 08/21/2026<br/>
-++Last updated: 09/28/2026<br/>
-++Discussion thread: [openxla/stablehlo#2993](https://github.com/openxla/stablehlo/pull/2993)
-++
-++## Overview
-++
-++This RFC proposes extending the set of complex types supported by StableHLO
-++from:
-++
-++```text
-++complex<f32>
-++complex<f64>
-++```
-++
-++to:
-++
-++```text
-++complex<f16>
-++complex<bf16>
-++complex<f32>
-++complex<f64>
-++```
-++
-++The proposal uses the existing MLIR builtin `ComplexType`; it does not add a
-++new StableHLO type. It extends the set of floating-point component types that
-++StableHLO accepts inside `complex<T>`.
-++
-++This RFC defines the StableHLO and VHLO type and serialization contract for
-++low-precision complex types. It does not implement XLA/HLO, the
-++StableHLO-to-XLA bridge, runtime ABI exposure, or backend kernels. The feature
-++is intended to land as part of a coordinated implementation sequence with the
-++XLA proposal in [openxla/xla#6136](https://github.com/openxla/xla/issues/6136).
-++Approval of this RFC does not imply that every StableHLO consumer can execute
-++`complex<f16>` or `complex<bf16>`; consumers without support must reject these
-++types with a clear diagnostic before code generation.
-++
-++### Scope and current implementation status
-++
-++This RFC specifies the StableHLO/VHLO side of a cross-repository feature. The
-++current status of the related layers is:
-++
-++| Layer | Current status | Boundary of this RFC |
-++| --- | --- | --- |
-++| StableHLO type specification | Only `complex<f32>` and `complex<f64>` are currently specified | Adds `complex<f16>` and `complex<bf16>` |
-++| VHLO compatibility | `ComplexV1` is the only complex representation, and its current shallow verifier does not constrain the component type | Adds a target-version-aware component-type constraint on `ComplexV1` |
-++| XLA/HLO core types | `C32`/`BC32` are proposed in [openxla/xla#6136](https://github.com/openxla/xla/issues/6136), but are not yet implemented | Coordinated prerequisite; not implemented here |
-++| StableHLO-to-XLA bridge | No `C32`/`BC32` type or constant mapping is available | Follow-up implementation |
-++| Runtime and backend support | [cuFFT](https://docs.nvidia.com/cuda/cufft/index.html) exposes 16-bit complex formats, but XLA does not yet expose a corresponding path | Separate runtime and backend work |
-++
-++The StableHLO specification may define a valid type domain before every
-++consumer implements it. This RFC therefore defines the observable rejection
-++behavior for consumers that have not yet added execution support; it does not
-++claim end-to-end execution support from StableHLO serialization alone.
-++
-++## Motivation
-++
-++[openxla/stablehlo#1794](https://github.com/openxla/stablehlo/issues/1794)
-++requests support for 16-bit FFT input and output types. Today, the StableHLO
-++verifier prevents such programs even though MLIR can represent both
-++`complex<f16>` and `complex<bf16>`.
-++
-++There are two independent concerns:
-++
-++1. Whether StableHLO can represent and specify a low-precision complex value.
-++2. Whether a particular consumer or backend can execute that value efficiently.
-++
-++StableHLO should decide the first concern based on the coherence of its type
-++system and semantics rather than on the current implementation status of one
-++consumer. The related XLA proposal currently remains at the design stage; it
-++proposes adding `C32 = complex<f16>` and `BC32 = complex<bf16>` as a shared
-++core-type change before bridge and backend work. Consumers that do not support
-++the new types may reject them with a diagnostic until they add a lowering,
-++legalization, or native implementation. Such rejection is part of the staged
-++rollout, not an implicit promotion to `complex<f32>`.
-++
-++The body of #1794 suggests allowing any floating-point component type. This
-++RFC intentionally narrows that suggestion to f16 and bf16, the two widely used
-++16-bit floating-point formats, in addition to the existing f32 and f64 types.
-++Supporting only `bf16` would not address the issue's original `f16` request,
-++while allowing every member of StableHLO's broader floating-point type set
-++would also introduce complex types with FP4, FP6, FP8, and other specialized
-++component formats. Those formats have distinct representability and
-++compatibility questions and should not be included implicitly in this
-++proposal.
-++
-++## Current behavior
-++
-++The current StableHLO specification lists only `complex<f32>` and
-++`complex<f64>` as supported complex types.
-++
-++The implementation reflects this restriction in multiple places:
-++
-++- `HLO_Complex` is defined as `Complex<HLO_Float32Or64>`.
-++- The operands of `stablehlo.complex` are restricted to f32/f64 tensors.
-++- RFFT type inference has a hand-written check that requires f32 or f64 input.
-++- The StableHLO test suite contains a negative test requiring f16 RFFT input to
-++  be rejected.
-++
-++Changing only the RFFT check would be incomplete: RFFT could infer a type that
-++the common StableHLO complex constraint and `stablehlo.complex` still reject.
-++
-++VHLO currently represents MLIR builtin complex types with `ComplexV1`, whose
-++verifier only checks that the component is a VHLO type. This is a shallow
-++structural check: the representation can therefore spell component
-++combinations that were never in the historical StableHLO complex type domain.
-++Historically valid `ComplexV1` values are limited to f32 and f64 components.
-++
-++The current VHLO version-conversion pass has no type-to-type version
-++conversions because no VHLO type has previously required one. It recursively
-++checks whether nested types are legal for the target version, and that check
-++is already target-version-aware, but complex-type legality today does not
-++depend on the component type. Version-dependent constraint checks already
-++exist at the operation level: `validateConstraint` gates the reduce-family
-++operations and `custom_call_v1` on the target version, for example allowing
-++mismatched operand and result element types in reduce operations only from
-++0.17.0. This RFC extends that existing constraint pattern to the complex
-++component type rather than introducing new type-conversion infrastructure.
-++
-++## Proposal
-++
-++### Supported complex component types
-++
-++StableHLO will define an explicit complex component type set:
-++
-++```text
-++ComplexComponent ::= f16 | bf16 | f32 | f64
-++ComplexType      ::= complex<ComplexComponent>
-++```
-++
-++The implementation should use a dedicated ODS constraint for this set rather
-++than define `HLO_Complex` in terms of the full `HLO_Float` constraint. This
-++keeps future additions to `HLO_Float` from automatically changing the
-++StableHLO complex type system without review.
-++
-++No mixed-component complex type is introduced. A complex value continues to
-++have real and imaginary components with the same element type.
-++
-++### Effect on existing operations
-++
-++This proposal does not add an operation and does not add complex semantics to
-++an operation that currently excludes complex types.
-++
-++For every existing operation whose specified input or result type domain
-++already includes a complex type, `complex<f16>` and `complex<bf16>` become
-++additional valid instantiations of that existing type domain. This RFC
-++proposes no operation-specific exception. If review identifies an operation
-++whose semantics cannot support either new instantiation, that exception must
-++be specified in a revision of this RFC rather than chosen during
-++implementation.
-++
-++This consequence is intentional. Restricting the new types to FFT alone would
-++allow a low-precision complex value to be produced but prevent it from flowing
-++through the same basic construction, extraction, arithmetic, conversion, and
-++structural operations as other StableHLO complex values.
-++
-++The observable result type of an operation remains the type written in the
-++StableHLO program. This RFC does not insert implicit promotion to f32 and does
-++not change `complex<f16>` or `complex<bf16>` into another IR type. A consumer
-++may use wider intermediate precision where permitted by the existing operation
-++semantics and accuracy rules, but the declared input and result types remain
-++unchanged.
-++
-++The normative affected surface is grouped as follows:
-++
-++- Direct construction and extraction operations, including `complex`, `real`,
-++  and `imag`, accept the two new instantiations.
-++- Existing parametric complex arithmetic and transcendental operations accept
-++  them under their existing semantics and accuracy contracts. This includes
-++  the direct complex-domain uses in `abs`, `atan2`, `cbrt`, `cosine`, `divide`,
-++  `exponential`, `exponential_minus_one`, `log`, `log_plus_one`, `logistic`,
-++  `negate`, `power`, `remainder`, `rsqrt`, `sign`, `sine`, `sqrt`, `subtract`,
-++  `tan`, and `tanh`, as well as operations such as `add` and `multiply` whose
-++  generic tensor domains already include complex element types.
-++- Existing complex linear-algebra operations, including `cholesky` and
-++  `triangular_solve`, accept the two new instantiations.
-++- FFT, IFFT, RFFT, and IRFFT accept the relationships specified below.
-++- Constants, `iota`, conversion, data-movement, shape-preserving, control-flow,
-++  and structural operations continue to accept complex values wherever their
-++  current type domains already do so.
-++- A supported complex type may occur wherever an existing StableHLO container
-++  admits a complex element or nested value, including tensors, buffers,
-++  tuples, futures, function signatures, region arguments, and custom calls.
-++
-++Because `HLO_Complex` is used transitively, the implementation audit must cover
-++the complete ODS constraint graph rather than only literal `HLO_Complex`
-++occurrences. In particular, derived constraints such as `HLO_Tensor`,
-++`HLO_NonQuantizedTensor`, `HLO_Buffer`, `HLO_Tuple`, `HLO_CustomCallValue`, and
-++their static-shape or unranked variants expand with the common complex
-++component set. Every custom verifier and inference path reached through these
-++constraints must be checked against the specification above.
-++
-++This shared-constraint change also reaches the CHLO dialect: CHLO operations
-++such as `broadcast_complex` and the inverse-trigonometric and hyperbolic
-++complex math operations use the same complex tensor predicates. The
-++implementation must either include those operations and their
-++legalization/decomposition paths in the affected-surface audit and tests, or
-++deliberately separate their constraints so that this RFC does not expand CHLO
-++acceptance accidentally.
-++
-++This expansion does not create separate kinds of complex type. It also does
-++not define the physical layout or ABI of a buffer containing a low-precision
-++complex element; StableHLO buffer acceptance remains an abstract type-system
-++property under the existing buffer contract.
-++
-++### `complex`, `real`, and `imag`
-++
-++For every supported complex component type `T`:
-++
-++```text
-++stablehlo.complex : (tensor<...xT>, tensor<...xT>)
-++                 -> tensor<...xcomplex<T>>
-++
-++stablehlo.real    : tensor<...xcomplex<T>> -> tensor<...xT>
-++stablehlo.imag    : tensor<...xcomplex<T>> -> tensor<...xT>
-++```
-++
-++The existing shape and same-element-type constraints continue to apply.
-++
-++In particular, the following remain invalid:
-++
-++- constructing a complex value from one f16 tensor and one bf16 tensor;
-++- declaring a `complex<f32>` result for f16 components;
-++- declaring a real or imaginary result whose component type differs from the
-++  input complex component type.
-++
-++### FFT type relationships
-++
-++For `T` in `{f16, bf16, f32, f64}`, the FFT type relationships are:
-++
-++```text
-++FFT:   complex<T> -> complex<T>
-++IFFT:  complex<T> -> complex<T>
-++RFFT:          T  -> complex<T>
-++IRFFT: complex<T> -> T
-++```
-++
-++The existing FFT rank, shape, and `fft_length` constraints are unchanged.
-++
-++The hand-written RFFT type check should use the same supported complex
-++component predicate as the common complex type constraint, avoiding a second
-++component-type list that can drift from the specification.
-++
-++### Constants and conversion
-++
-++Complex constants with f16 or bf16 components follow the existing StableHLO
-++complex constant definition: the real and imaginary literals are interpreted
-++using the floating-point semantics of the component type.
-++
-++`stablehlo.convert` continues to use its existing specified conversion
-++semantics. This RFC makes the new complex types valid source or result types in
-++cases where the existing `convert` type domain includes complex types; it does
-++not introduce an implicit conversion or a new fallback rule.
-++
-++## Compatibility and VHLO
-++
-++Low-precision complex support is a new StableHLO feature and therefore requires
-++a StableHLO minor version bump. The implementation must use the next available
-++minor version at the time it lands. If implemented immediately after StableHLO
-++1.20.0, the feature version would be 1.21.0.
-++
-++The required observable compatibility behavior is:
-++
-++- serialization targeting the feature version or a later version succeeds;
-++- serialization of a program using `complex<f16>` or `complex<bf16>` to an
-++  earlier target version fails with a diagnostic;
-++- no compatibility path may silently promote the value to `complex<f32>` or
-++  decompose it into unrelated tensor types;
-++- programs using only `complex<f32>` and `complex<f64>` retain their existing
-++  ability to target older supported StableHLO versions;
-++- current consumers remain able to deserialize portable artifacts containing
-++  historical f32/f64 complex programs.
-++
-++StableHLO version compatibility and consumer execution capability are separate
-++contracts. Serialization to the feature version or a later version only
-++establishes that the artifact is representable by that StableHLO version. It
-++does not prove that an XLA, PJRT, or other consumer can lower every operation
-++containing the new type. A consumer that understands the StableHLO version but
-++lacks `C32` or `BC32` support must reject the module before code generation
-++with a diagnostic that identifies the unsupported type and operation.
-++
-++### Proposed VHLO representation
-++
-++VHLO retains `ComplexV1` as the only VHLO representation of complex types.
-++`ComplexV1` is structurally unchanged: one element-type parameter with the
-++existing shallow verifier, and an unchanged version range. The
-++version-dependent validity of component types is expressed as a
-++target-version-aware type constraint, following the pattern VHLO already uses
-++for operation constraints (`validateConstraint` on the reduce-family
-++operations and `custom_call_v1`).
-++
-++The constraint is normative:
-++
-++- For a target version at or after the feature version, `ComplexV1` accepts
-++  exactly f16-, bf16-, f32-, and f64-based components. FP8, integer, and all
-++  other component types are rejected.
-++- For a target version before the feature version, `ComplexV1` accepts only
-++  f32- and f64-based components. A `ComplexV1<f16>` or `ComplexV1<bf16>` value
-++  makes serialization to that target fail with a diagnostic, in the same way
-++  that a reduce operation with mismatched operand and result element types
-++  fails when the target version is older than 0.17.0.
-++- No new VHLO type or attribute is introduced, and no type-to-type conversion
-++  or type rewriting occurs during version conversion. Serialization of
-++  programs that use only f32/f64 complex types is therefore unaffected.
-++
-++The implementation must add a `VHLO_VersionedTypeConstraintInterface` as the
-++type-level analogue of `VHLO_VersionedOpConstraintInterface`. It provides a
-++`validateConstraint(Type, Version)` hook. `ComplexV1` must implement the hook,
-++and the recursive `isLegalType` check must invoke it. A local helper may
-++provide the hook's implementation, but the target-aware check must have a
-++type-level entry point; operation-only validation is insufficient because
-++complex types can be nested in signatures, containers, and type-bearing
-++attributes. This interface is an implementation hook and does not change the
-++serialized VHLO type identity or its version range.
-++
-++This models a component-domain relaxation as a versioned constraint rather
-++than as a new type version. VHLO is a serialization dialect: it is produced by
-++conversion from StableHLO and consumed by version conversion, not authored by
-++hand, and the dialect does not guarantee rejection of hand-authored payloads
-++that could never be produced through those flows. Every artifact is tagged
-++with the version it targets, and a consumer must be at least that version, so
-++a payload cannot be presented to an older consumer than it was serialized for.
-++
-++Officially produced bytecode is validated against the requested target version
-++before emission. Consequently, no official pre-feature artifact can contain a
-++low-precision complex value. Hand-authored VHLO payloads that violate this
-++producer contract are outside this RFC's portable-artifact guarantees; a
-++consumer may reject them during parsing or conversion, and accepting one must
-++not establish that the pre-feature target historically supported the type.
-++
-++### Recursive legality requirements
-++
-++The component-type constraint must be evaluated by the recursive type
-++legality check with access to the requested target version. The current check
-++already recurses through several nested VHLO types and attributes, but the
-++implementation for this feature must also add explicit coverage for every
-++type-bearing location listed below. Checking only an operation's immediate
-++result types would leave invalid values hidden in containers or signatures.
-++The constraint must therefore be enforced at every location, including:
-++
-++- `RankedTensorV1`, `UnrankedTensorV1`, and `RankedBufferV1` element types;
-++- `TupleV1` elements, nested tuples, and `FutureV1` element types;
-++- `FunctionV1` inputs and results;
-++- operation operands and results, function signatures, and region block
-++  arguments;
-++- `TypeV1Attr`, `TensorV1Attr`, `FloatV1Attr`, `IntegerV1Attr`, and every other
-++  attribute that carries a type;
-++- ranked-tensor encodings and nested array or dictionary attributes; and
-++- the storage and expressed types of `UniformQuantizedV1` and
-++  `UniformQuantizedPerAxisV1`, which must be traversed even when the
-++  quantized-type semantics reject a complex inhabitant.
-++
-++If any nested occurrence fails the component-type constraint for the requested
-++target, the entire version conversion must fail before bytecode is written.
-++The diagnostic must identify the target version and the unsupported complex
-++component type. No path may partially convert the module or emit an artifact
-++containing a complex value whose component type is invalid for the target
-++version.
-++
-++### Conversion behavior and invariants
-++
-++Because no VHLO type changes, version conversion requires no type rewriting.
-++The observable behavior is:
-++
-++1. Validate the target version as today.
-++2. Preserve the existing operation-version rewrites; complex types are never
-++   rewritten.
-++3. Extend the recursive type legality check to evaluate the component-type
-++   constraint against the requested target version over every type-bearing
-++   location listed above.
-++4. If any check fails, fail the pass before bytecode emission with a
-++   diagnostic naming the target version and the unsupported component type.
-++
-++The transformation remains all-or-nothing from the serializer's perspective,
-++and serialization must not emit bytecode from a module whose complex component
-++types are invalid for the requested target version.
-++
-++## Consumer behavior
-++
-++StableHLO verifier acceptance is not a claim that every consumer can execute
-++the new types.
-++
-++StableHLO serialization validates only StableHLO/VHLO representability. A
-++StableHLO-to-XLA bridge must reject `complex<f16>` or `complex<bf16>` before
-++constructing the XLA/HLO module when the corresponding XLA primitive type is
-++unavailable. The diagnostic must identify the unsupported type, operation, and
-++target. Other consumers must provide an equivalent pre-code-generation
-++capability check or a separately specified legalization path.
-++
-++A consumer that lacks support may:
-++
-++- reject a module containing the new types with a clear unsupported-type
-++  diagnostic;
-++- apply a separately specified legalization;
-++- implement the types natively.
-++
-++This RFC does not prescribe which option a consumer must choose. It does
-++require that unsupported consumers fail explicitly before code generation. In
-++particular, a consumer must not reinterpret a four-byte value as `C64`, fall
-++through an existing `C64`/`C128` path, or silently promote the declared value to
-++`complex<f32>`. XLA support is not a prerequisite for StableHLO type validity,
-++but the StableHLO implementation and the XLA/HLO work must be coordinated
-++before the feature is presented as end-to-end executable.
-++
-++## Verification and testing
-++
-++The implementation must include positive tests for:
-++
-++- StableHLO verifier acceptance and exact type inference for the complete
-++  specified domain, independently of whether a particular consumer executes
-++  every operation;
-++- parsing and printing `complex<f16>` and `complex<bf16>` in StableHLO
-++  programs;
-++- `stablehlo.complex`, `stablehlo.real`, and `stablehlo.imag` type relations;
-++- an acceptance case for every existing operation with a direct complex-domain
-++  numeric, generator, or linear-algebra constraint, including `iota`, plus
-++  representative cases for every transitive structural and container
-++  constraint family listed above;
-++- FFT, IFFT, RFFT, and IRFFT for both f16 and bf16 component types;
-++- exact return-type inference, including dynamic-shape cases;
-++- StableHLO-to-current-VHLO-to-StableHLO round trips for both `complex<f16>`
-++  and `complex<bf16>`;
-++- serialization and deserialization at the new version for both
-++  `complex<f16>` and `complex<bf16>`;
-++- deserialization of historical f32/f64 complex artifacts, which requires no
-++  type rewriting under this proposal;
-++- serialization of `complex<f16>` and `complex<bf16>` values, including values
-++  nested in every semantically valid type-bearing container and attribute
-++  listed above, to the feature version and later targets; and
-++- continued old-version serialization of f32/f64 complex programs.
-++
-++The implementation must include negative tests for:
-++
-++- consumer-side rejection before code generation when `C32` or `BC32` support
-++  is absent, including a diagnostic that names the unsupported type, operation,
-++  and target;
-++- complex values with integer component types;
-++- at least one floating-point component type outside this RFC, such as an FP8
-++  type;
-++- rejection of serialization to any target version before the feature version
-++  when a `ComplexV1<f16>` or `ComplexV1<bf16>` value is present, both directly
-++  and when nested in every semantically valid supported type-bearing container
-++  or attribute listed above, enforced by the version-conversion constraint
-++  rather than by tightening the shallow V1 parser or verifier;
-++- rejection of complex types placed in quantized-type storage or expressed
-++  type parameters by the semantic constraint owned by the corresponding
-++  StableHLO or consumer layer, while those locations are still traversed for
-++  VHLO version legality;
-++- a complex value with an FP8, integer, or other unsupported component, for
-++  any target version;
-++- mismatched f16/bf16 inputs to `stablehlo.complex`;
-++- mismatched real/complex component types in FFT input and result types;
-++- attempting to serialize a program containing `complex<f16>` or
-++  `complex<bf16>` to the version immediately before the feature version; and
-++- failure before bytecode emission, with a diagnostic that names both the
-++  requested target version and the unsupported nested type.
-++
-++The VHLO compatibility suite must cover the feature version following the
-++existing VHLO checklist, including fixtures that exercise the component-type
-++constraint against old and new target versions.
-++
-++Reference-interpreter numerical tests may be delivered in a separate
-++StableHLO-only follow-up if the required tensor storage and constant
-++materialization changes would obscure review of the type-system and
-++compatibility change. If maintainers require reference execution as an opset
-++acceptance criterion, that follow-up can be stacked or folded into the
-++implementation change without adding XLA or backend work.
-++
-++## Non-goals
-++
-++This RFC does not propose:
-++
-++- the implementation of an XLA primitive type or StableHLO-to-XLA bridge;
-++  those pieces are coordinated with [openxla/xla#6136](https://github.com/openxla/xla/issues/6136)
-++  and are required before claiming end-to-end execution;
-++- JAX dtype, promotion, tracing, or lowering changes;
-++- portable decomposition into real and imaginary tensor planes;
-++- CPU, GPU, TPU, or accelerator kernels;
-++- FFT library integration;
-++- a memory layout, storage ABI, or claim that a value occupies four bytes in a
-++  particular runtime;
-++- mandated accumulation precision or a new numerical-accuracy contract;
-++- automatic fallback or promotion to `complex<f32>`;
-++- support for complex values whose components use FP4, FP6, FP8, E8M0, or any
-++  future floating-point type not listed by this RFC;
-++- performance claims.
-++
-++These concerns are separate implementation stages, but they are not an
-++indefinite dependency: the enabling StableHLO implementation must have an
-++agreed XLA/HLO core-type plan and an explicit unsupported-consumer path before
-++it lands. The bridge, runtime, and backend changes may remain separate pull
-++requests.
-++
-++## Alternatives considered
-++
-++### Support only `complex<bf16>`
-++
-++This is the smallest change for bf16-specific use cases, but it does not address
-++the f16 request in #1794 and introduces an asymmetric exception between the two
-++widely used 16-bit floating-point formats. The additional StableHLO and VHLO
-++work required to support f16 and bf16 together is substantially shared, so this
-++RFC does not recommend a bf16-only type domain.
-++
-++### Support `Complex<HLO_Float>`
-++
-++This is mechanically concise but would include every floating-point type in
-++`HLO_Float`, including specialized FP4, FP6, FP8, and E8M0 formats. It would
-++also cause future additions to `HLO_Float` to expand the StableHLO complex type
-++domain without a separate compatibility decision. This RFC instead proposes a
-++closed component set.
-++
-++### Support low-precision complex types only on FFT
-++
-++This narrows the immediate verifier change but produces a fragmented type
-++system in which FFT can create a complex value that basic complex construction,
-++extraction, arithmetic, conversion, or structural operations may reject. If a
-++strictly FFT-only feature is desired, it should be proposed as an explicit
-++operation-specific type extension rather than as general StableHLO support for
-++low-precision complex types.
-++
-++### Reuse `ComplexV1` with target-aware validation (chosen)
-++
-++VHLO keeps `ComplexV1` as the only complex representation and expresses the
-++version-dependent validity of component types as a target-version-aware type
-++constraint, the same pattern the dialect already uses for operation
-++constraints. Review of this RFC directed this approach over a new type
-++version: VHLO is a serialization dialect produced by conversion and consumed
-++by version conversion rather than authored by hand, every artifact is tagged
-++with its target version, and a downgrade containing values invalid for the
-++target already fails today. Modeling a component-domain relaxation with the
-++constraint mechanism keeps the category of change consistent with existing
-++practice and requires far less code.
-++
-++### Add `ComplexV2` at the feature version (declined)
-++
-++A new `ComplexV2` with a closed component set keeps the meaning of a complex
-++type intrinsic to the type version and makes the compatibility boundary
-++inspectable without target context. It would, however, require the dialect's
-++first type version split together with new type-to-type conversion
-++infrastructure: recursive rewriting of enclosing types, function signatures,
-++and type-bearing attributes for a change that does not alter the structure of
-++the type. Given VHLO's serialization-dialect stance, that cost is not
-++justified for this category of change.
-++
-++The alternatives are:
-++
-++| Design | Benefit | Cost |
-++| --- | --- | --- |
-++| Target-aware validation on `ComplexV1` | Consistent with the existing versioned-constraint pattern; no type rewriting; far less code | Complex-type validity depends on the target version rather than the type alone |
-++| `ComplexV2` with V1/V2 conversion | Type-intrinsic meaning and an inspectable boundary without target context | First VHLO type version split plus recursive type-conversion infrastructure |
-++
-++Reusing `ComplexV1` with no version check at all would allow new programs to
-++appear serializable to StableHLO versions that never specified these
-++component combinations, and remains rejected.
-++
-++## Rollout and pull request boundaries
-++
-++The proposed upstream sequence is deliberately split by capability:
-++
-++1. Keep this RFC as a design-only PR and confirm the cross-repository contract
-++   with [openxla/xla#6136](https://github.com/openxla/xla/issues/6136). RFC
-++   approval does not itself enable serialization or execution.
-++2. Obtain agreement on the XLA/HLO core representation (`C32` and `BC32`),
-++   host storage, `LiteralProto`, parser/printer behavior, and explicit
-++   unsupported-execution diagnostics. The current XLA proposal places this
-++   core-type change before bridge and backend work.
-++3. After the RFC is approved, submit the compatibility-atomic StableHLO
-++   implementation, including:
-++
-++   - specification, ODS, verifier, and type-inference changes;
-++   - the target-version-aware `ComplexV1` component-type constraint and its
-++     recursive legality coverage;
-++   - positive, negative, round-trip, and serialization tests;
-++   - tests that reject unsupported target versions before bytecode emission.
-++
-++   The affected-surface audit must include CHLO operations that inherit the
-++   shared complex predicates, or record an explicit constraint boundary that
-++   excludes them.
-++
-++   The implementation may be developed in parallel with the XLA core change,
-++   but it must not be presented as end-to-end execution support. No released
-++   state may emit the new StableHLO types without the corresponding VHLO
-++   boundary.
-++4. Implement the StableHLO-to-XLA bridge, Shape mapping, constant materialization,
-++   and bit-exact round trips in a separate follow-up. A bridge that has not yet
-++   added `C32`/`BC32` must reject those types before code generation.
-++5. Add explicit legalization policies, keeping pair decomposition and
-++   promotion to `C64` as separate choices. The core type change must not choose
-++   either policy implicitly.
-++6. Add runtime ABI exposure and backend implementations, beginning with a
-++   concrete backend path such as the cuFFT 16-bit complex formats and then
-++   extending support as maintainers approve.
-++
-++The RFC PR must not close #1794 by itself because design approval does not
-++implement the requested functionality. The implementation PR can close the
-++issue only after the agreed StableHLO support and the required consumer path
-++are present.
-+diff --ruN a/stablehlo/rfcs/20260907-exp2-log2.md b/stablehlo/rfcs/20260907-exp2-log2.md
-+--- stablehlo/rfcs/20260907-exp2-log2.md
-++++ stablehlo/rfcs/20260907-exp2-log2.md
-+@@ -0,0 +1,51 @@
-++# [RFC] Add Exp2Op and Log2Op to the StableHLO specification
-++
-++Status: Review<br/>
-++Initial version: 09/07/2026<br/>
-++Last updated: 10/02/2026<br/>
-++Discussion thread: [GitHub](https://github.com/openxla/stablehlo/pulls)
-++
-++## Motivation
-++
-++IEEE-754 (`exp2`, `log2`), C99 `<math.h>`, and LLVM intrinsics (`llvm.exp2`,
-++`llvm.log2`) include base-2 exponential (`exp2`) and base-2 logarithm (`log2`)
-++as standard elementary math functions, and hardware accelerators (such as GPU
-++special function units and TPU EUP instructions) support them directly.
-++
-++Currently, StableHLO lacks native `exp2` and `log2` operations:
-++
-++*   Frameworks expose `exp2` and `log2` directly: JAX provides
-++    [`jnp.exp2`][jax1] / `lax.exp2` and [`jnp.log2`][jax2] / `lax.log2`, and
-++    PyTorch provides [`torch.exp2`][torch1] and [`torch.log2`][torch2].
-++*   Without first-class ops in StableHLO, frameworks must decompose `exp2(x)`
-++    into `exponential(x * ln(2))` and `log2(x)` into `log(x) / ln(2)`.
-++*   Decomposing `exp2` and `log2` into natural `exponential` and `log` scaled by
-++    `ln(2)` degrades numerical accuracy (introducing rounding error on integer
-++    powers of two where `exp2(n)` and `log2(2^n)` should evaluate exactly) and
-++    prevents backends from emitting single-instruction base-2 hardware
-++    operations without fragile algebraic pattern-matching.
-++
-++Following the precedent of `TanOp` ([20240312-tanop.md](20240312-tanop.md)) and
-++the transcendental accuracy specification
-++([20241015-result-accuracy.md](20241015-result-accuracy.md)), this RFC adds
-++`exp2` (`Exp2Op`) and `log2` (`Log2Op`) to StableHLO with optional
-++`result_accuracy` attributes.
-++
-++## Proposed Specification Changes
-++
-++Please refer to [`docs/spec.md`](../docs/spec.md) for the specification of
-++[`exp2`](../docs/spec.md#exp2) and [`log2`](../docs/spec.md#log2).
-++
-++## Exp2Op, Log2Op, and Compatibility
-++
-++*   **VHLO Serialization**: StableHLO version `1.21.0` introduces
-++    `VHLO_Exp2OpV1` (`exp2_v1`) and `VHLO_Log2OpV1` (`log2_v1`).
-++*   **Forward Compatibility & Decomposition**: For runtimes targeting StableHLO
-++    versions prior to `1.21.0`, compatibility expansion lowers `exp2(x)` to
-++    `exponential(multiply(x, constant(ln(2))))` and `log2(x)` to `divide(log(x),
-++    constant(ln(2)))`.
-++
-++[jax1]: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.exp2.html
-++[jax2]: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.log2.html
-++[torch1]: https://pytorch.org/docs/stable/generated/torch.exp2.html
-++[torch2]: https://pytorch.org/docs/stable/generated/torch.log2.html
- diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h b/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h
- --- stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h
- +++ stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h
-@@ -2892,6 +3749,24 @@ diff --ruN a/stablehlo/stablehlo/reference/InterpreterOps.cpp b/stablehlo/stable
-    // TODO(#1725): Figure out how to test the outfeed queue.
-    return results;
-  }
-+diff --ruN a/stablehlo/stablehlo/reference/Ops.cpp b/stablehlo/stablehlo/reference/Ops.cpp
-+--- stablehlo/stablehlo/reference/Ops.cpp
-++++ stablehlo/stablehlo/reference/Ops.cpp
-+@@ -1777,12 +1777,8 @@
-+         dotGeneralOp(reversedLhsWindow, rhs, lhsContractingDimensions,
-+                      rhsContractingDimensions, resultType.getElementType());
-+ 
-+-    Sizes resultNonSpatialDims;
-+-    for (auto i = 0; i < result.getRank(); ++i)
-+-      if (llvm::find(outputSpatialDimensions, i) ==
-+-          outputSpatialDimensions.end())
-+-        resultNonSpatialDims.push_back(result.getShape()[i]);
-+-
-++    Sizes resultNonSpatialDims{result.getShape()[outputBatchDimension],
-++                               result.getShape()[outputFeatureDimension]};
-+     Axes resultPermutation;
-+     resultPermutation.push_back(outputBatchDimension);
-+     resultPermutation.append(outputSpatialDimensions.begin(),
- diff --ruN a/stablehlo/stablehlo/reference/Process.cpp b/stablehlo/stablehlo/reference/Process.cpp
- --- stablehlo/stablehlo/reference/Process.cpp
- +++ stablehlo/stablehlo/reference/Process.cpp
-@@ -3381,6 +4256,28 @@ diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/st
-  // CHECK-LABEL:   func.func @mulhi_i32(
-  // CHECK-SAME:      %[[ARG0:.*]]: tensor<4xi32>,
-  // CHECK-SAME:      %[[ARG1:.*]]: tensor<4xi32>) -> tensor<4xi32> {
-+diff --ruN a/stablehlo/stablehlo/tests/interpret/convolution.mlir b/stablehlo/stablehlo/tests/interpret/convolution.mlir
-+--- stablehlo/stablehlo/tests/interpret/convolution.mlir
-++++ stablehlo/stablehlo/tests/interpret/convolution.mlir
-+@@ -134,3 +134,18 @@
-+                                          [[21, 21, 29, 29]]]]> : tensor<1x2x1x4xi64>
-+   func.return
-+ }
-++
-++// -----
-++
-++func.func @convolution_op_test_feature_before_batch() {
-++  %lhs = stablehlo.constant dense<[[[1.0]], [[2.0]]]> : tensor<2x1x1xf32>
-++  %rhs = stablehlo.constant dense<[[[3.0, 4.0]]]> : tensor<1x1x2xf32>
-++  %result = stablehlo.convolution(%lhs, %rhs)
-++    dim_numbers = [0, f, b]x[0, i, o]->[0, f, b],
-++    window = {stride = [1], pad = [[0, 0]], rhs_dilate = [1]}
-++    {batch_group_count = 1 : i64, feature_group_count = 1 : i64}
-++    : (tensor<2x1x1xf32>, tensor<1x1x2xf32>) -> tensor<2x2x1xf32>
-++  check.expect_almost_eq_const %result,
-++    dense<[[[3.0], [4.0]], [[6.0], [8.0]]]> : tensor<2x2x1xf32>
-++  func.return
-++}
- diff --ruN a/stablehlo/stablehlo/tests/ops_chlo.mlir b/stablehlo/stablehlo/tests/ops_chlo.mlir
- --- stablehlo/stablehlo/tests/ops_chlo.mlir
- +++ stablehlo/stablehlo/tests/ops_chlo.mlir
diff --git a/third_party/xla/third_party/shardy/workspace.bzl b/third_party/xla/third_party/shardy/workspace.bzl
index fc42ec7cc04..b98756b3abe 100644
--- a/third_party/xla/third_party/shardy/workspace.bzl
+++ b/third_party/xla/third_party/shardy/workspace.bzl
@@ -18,8 +18,8 @@
 load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 
 def repo():
-    SHARDY_COMMIT = "a068921f6ece54dfe0a1557930edceecba0da7a9"
-    SHARDY_SHA256 = "af328df3781e3660529208be38b885ac41bdd9d2d5aa7a3980f95b0c1f1043ce"
+    SHARDY_COMMIT = "f01c1a6b084c089087827633333615546d990867"
+    SHARDY_SHA256 = "f85cc219fad3db16e92faaa1ceb214c22ba3752f559780beb00cba7809c530a5"
 
     tf_http_archive(
         name = "shardy",

From 254a192f20d4c5c1ebb3390250e5bb6b2e9397f7 Mon Sep 17 00:00:00 2001
From: charleshofer <Charles.Hofer@amd.com>
Date: Wed, 7 Oct 2026 23:42:39 -0700
Subject: [PATCH 0246/1000] PR #50161: Move to new ROCm RBE cluster

Imported from GitHub PR https://github.com/openxla/xla/pull/50161

AMD has set up a new remote build execution cluster to be used for XLA CI. Switching to the new cluster should improve CI's reliability and artifact transfer time. This PR switches over the settings in the .bazelrc to use the new cluster.
Copybara import of the project:

--
e48ef838e969373f78c56160251258d7669ba5e5 by Charles Hofer <Charles.Hofer@amd.com>:

Move to new ROCm RBE cluster

Merging this change closes #50161

PiperOrigin-RevId: 995613462
---
 third_party/xla/.github/workflows/rocm_ci.yml     |  2 +-
 third_party/xla/build_tools/rocm/rocm_xla.bazelrc | 14 +++++++-------
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/.github/workflows/rocm_ci.yml b/third_party/xla/.github/workflows/rocm_ci.yml
index 55b6d985281..43aa1903d48 100644
--- a/third_party/xla/.github/workflows/rocm_ci.yml
+++ b/third_party/xla/.github/workflows/rocm_ci.yml
@@ -41,7 +41,7 @@ env:
   ROCM_DISTRO_URL: https://stable.repo.amd.com/rocm/core/tarball/therock-dist-linux-multiarch-10.0.0.tar.gz
   ROCM_DISTRO_HASH: 1c5e807875d26a2470ecc7323daa5b5b9009208a55c3290ac255a909cde15fc6
   ROCM_AMD_GPU_TARGETS: gfx950
-  ROCM_SINGLE_GPU_RBE_POOL: linux_x64_gpu_do_gfx950
+  ROCM_SINGLE_GPU_RBE_POOL: linux_x64_gpu
 
 jobs:
   rocm-config:
diff --git a/third_party/xla/build_tools/rocm/rocm_xla.bazelrc b/third_party/xla/build_tools/rocm/rocm_xla.bazelrc
index 5127630296e..e0c18369bc4 100644
--- a/third_party/xla/build_tools/rocm/rocm_xla.bazelrc
+++ b/third_party/xla/build_tools/rocm/rocm_xla.bazelrc
@@ -1,21 +1,21 @@
 # Test-related settings.
 
 build:rocm_dev --remote_upload_local_results=false
-build:rocm_dev --remote_cache="grpcs://wardite.cluster.engflow.com"
+build:rocm_dev --remote_cache="grpcs://internal.actinolite.cluster.engflow.com"
 
 common:rocm_base --test_verbose_timeout_warnings
 common:rocm_base --test_timeout=920,2400,7200,9600
 
-build:rocm_rbe --remote_cache="grpcs://wardite.cluster.engflow.com"
+build:rocm_rbe --remote_cache="grpcs://internal.actinolite.cluster.engflow.com"
 build:rocm_rbe --repo_env=REMOTE_GPU_TESTING=1
-build:rocm_rbe --bes_backend="grpcs://wardite.cluster.engflow.com"
-build:rocm_rbe --bes_results_url="https://wardite.cluster.engflow.com/invocation/"
+build:rocm_rbe --bes_backend="grpcs://internal.actinolite.cluster.engflow.com"
+build:rocm_rbe --bes_results_url="https://internal.actinolite.cluster.engflow.com/invocation/"
 build:rocm_rbe --host_platform="@local_config_rocm//rocm:linux_x64"
 build:rocm_rbe --extra_execution_platforms="@local_config_rocm//rocm:linux_x64"
 build:rocm_rbe --platforms="@local_config_rocm//rocm:linux_x64"
 build:rocm_rbe --bes_timeout=600s
-build:rocm_rbe --tls_client_certificate="/data/ci-cert.crt"
-build:rocm_rbe --tls_client_key="/data/ci-cert.key"
+build:rocm_rbe --tls_client_certificate="/data/temp-rbe-cert.crt"
+build:rocm_rbe --tls_client_key="/data/temp-rbe-cert.key"
 build:rocm_rbe --spawn_strategy=remote,local
 build:rocm_rbe --grpc_keepalive_time=30s
 build:rocm_rbe --remote_cache_compression
@@ -26,7 +26,7 @@ build:rocm_rbe --experimental_strict_repo_env
 
 
 test:rocm_rbe --jobs=50
-test:rocm_rbe --remote_executor=grpcs://wardite.cluster.engflow.com
+test:rocm_rbe --remote_executor=grpcs://internal.actinolite.cluster.engflow.com
 test:rocm_rbe --remote_timeout=3600
 test:rocm_rbe --strategy=TestRunner=remote,local
 test:rocm_rbe --test_sharding_strategy=explicit

From e492ecbcddaeb9f82bf1a27d3b5fd289ac030b26 Mon Sep 17 00:00:00 2001
From: yasiribmcon <107119183+yasiribmcon@users.noreply.github.com>
Date: Thu, 8 Oct 2026 12:43:46 +0530
Subject: [PATCH 0247/1000] [s390x] Add compiler-rt host link option

On s390x, building _pywrap_tensorflow fails with undefined symbols
__truncsfhf2 and __extendhfsf2. The failing link runs in Bazel's exec
configuration, which --linkopt doesn't cover. Add --host_linkopt so
compiler-rt is linked there too.

Signed-off-by: Yasir Ashfaq <Yasir.Ashfaq1@ibm.com>
---
 configure.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/configure.py b/configure.py
index 848253d8367..a542dca37f1 100644
--- a/configure.py
+++ b/configure.py
@@ -1332,6 +1332,7 @@ def main():
         clang_version = retrieve_clang_version(clang_compiler_path)
         if is_s390x():
           write_to_bazelrc('build --linkopt=-rtlib=compiler-rt')
+          write_to_bazelrc('build --host_linkopt=-rtlib=compiler-rt')
         disable_clang_offsetof_extension(clang_version)
     if is_windows():
       environ_cp['TF_NEED_CLANG'] = str(choose_compiler_Win(environ_cp))

From 1c2fdf20bde53608fa78c2312ee35660d275368a Mon Sep 17 00:00:00 2001
From: Bill Varcho <varcho@google.com>
Date: Thu, 8 Oct 2026 00:09:52 -0700
Subject: [PATCH 0248/1000] [XLA:SPMD][Bug-Fix] Partially replicate grouping
 dims before merging passthrough gather/scatter shardings.

**Problem.** Partitioning gather or scatter along passthrough dimensions with `NamedSharding` crashes in `AlignGroupsWith` when the output/updates sharding further shards a grouping dimension:

```
Check failed: grouped_sharding.device_groups.num_groups() == reference.device_groups.num_groups() (8 vs. 16)
```

**Cause.** Passthrough partitioning groups devices along the operand's (or indices') passthrough dimensions, and calls `MergeShardingIfCompatible` to pull extra sharding from the output (or updates) onto the *other* dimensions without changing the grouping dimensions. That worked with V2 `TileAssignment` because `MergeShardingIfCompatible` only modifies replicated (`size == 1`) dimensions, leaving already-sharded grouping dimensions untouched. With V3 `NamedSharding`, `MergeNamedShardingIfCompatible` also appends axes to already-sharded dimensions (e.g., extending `['x']` to `['x', 'y']`), which changes the grouping dimensions and gives the output more device groups than the operand/indices.

**Fix.** Replicate the grouping dimensions on the output/updates sharding (`PartiallyReplicateTiledShardingOnDims`) before passing it to `MergeShardingIfCompatible`, matching `PartitionGatherParallelDimensions` and `PartitionScatterParallelDimensions`.

PiperOrigin-RevId: 995626460
---
 .../service/spmd/gather_scatter_handler.cc    |  32 +++--
 .../xla/service/spmd/spmd_partitioner_test.cc | 135 ++++++++++++++++++
 2 files changed, 155 insertions(+), 12 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/gather_scatter_handler.cc b/third_party/xla/xla/service/spmd/gather_scatter_handler.cc
index 0a09e9f16c3..10089972997 100644
--- a/third_party/xla/xla/service/spmd/gather_scatter_handler.cc
+++ b/third_party/xla/xla/service/spmd/gather_scatter_handler.cc
@@ -385,8 +385,10 @@ absl::StatusOr<HloInstruction*> PartitionGatherIndexPassthroughDimensions(
   if (passthrough_sharding.IsReplicatedOrSingleDevice()) {
     return nullptr;
   }
-  hlo_sharding_util::MergeShardingIfCompatible(output_sharding,
-                                               &passthrough_sharding);
+  hlo_sharding_util::MergeShardingIfCompatible(
+      hlo_sharding_util::PartiallyReplicateTiledShardingOnDims(
+          output_sharding, index_passthrough_dims.output_dims),
+      &passthrough_sharding);
   // Group shardings on index pass-through dimensions.
   const GroupedSharding output_grouped = hlo_sharding_util::GroupShardingOnDims(
       passthrough_sharding, index_passthrough_dims.output_dims);
@@ -464,10 +466,12 @@ absl::StatusOr<HloInstruction*> PartitionGatherOperandPassthroughDimensions(
         pslice_sizes[i] = operand.hlo()->shape().dimensions(i);
       }
     }
-    // Merge the sharding from the instruction with the sharding suggested from
-    // the operand sharding.
-    hlo_sharding_util::MergeShardingIfCompatible(output_sharding,
-                                                 &*maybe_passthrough);
+    // Merge non-grouping dimension shardings from output_sharding into the
+    // suggested passthrough sharding.
+    hlo_sharding_util::MergeShardingIfCompatible(
+        hlo_sharding_util::PartiallyReplicateTiledShardingOnDims(
+            output_sharding, output_grouping_dims),
+        &*maybe_passthrough);
     // Group shardings on operand pass-through dimensions.
     const GroupedSharding output_grouped =
         hlo_sharding_util::GroupShardingOnDims(*maybe_passthrough,
@@ -1489,10 +1493,12 @@ absl::StatusOr<HloInstruction*> PartitionScatterOperandPassthroughDimensions(
         pslice_sizes[i] = operands[0].hlo()->shape().dimensions(i);
       }
     }
-    // Merge the sharding from update with the sharding suggested from the
-    // operand sharding.
-    hlo_sharding_util::MergeShardingIfCompatible(updates[0].sharding(),
-                                                 &*maybe_passthrough);
+    // Merge non-grouping dimension shardings from update_sharding into the
+    // suggested passthrough sharding.
+    hlo_sharding_util::MergeShardingIfCompatible(
+        hlo_sharding_util::PartiallyReplicateTiledShardingOnDims(
+            updates[0].sharding(), update_grouping_dims),
+        &*maybe_passthrough);
     // Group shardings on operand pass-through dimensions.
     const GroupedSharding update_grouped =
         hlo_sharding_util::GroupShardingOnDims(*maybe_passthrough,
@@ -1714,8 +1720,10 @@ absl::StatusOr<HloInstruction*> PartitionScatterIndexPassthroughDimensions(
   if (passthrough_sharding.IsReplicatedOrSingleDevice()) {
     return nullptr;
   }
-  hlo_sharding_util::MergeShardingIfCompatible(updates[0].sharding(),
-                                               &passthrough_sharding);
+  hlo_sharding_util::MergeShardingIfCompatible(
+      hlo_sharding_util::PartiallyReplicateTiledShardingOnDims(
+          updates[0].sharding(), index_passthrough_dims.output_dims),
+      &passthrough_sharding);
   const GroupedSharding update_grouped = hlo_sharding_util::GroupShardingOnDims(
       passthrough_sharding, index_passthrough_dims.output_dims);
   // See if we can group partially replicated dimensions from the operand
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
index b213ee56184..85d1bfed107 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
@@ -9661,6 +9661,29 @@ ENTRY entry {
                           op::Shape("f32[3,5]")));
 }
 
+TEST_P(SpmdPartitioningTest,
+       PassthroughGather_PartialReplicateOperandShardedOutput) {
+  absl::string_view hlo_string = R"(
+HloModule module
+
+ENTRY entry {
+  %input = f32[2,16] parameter(0),
+    sharding={devices=[1,2,2]<=[4] last_tile_dim_replicate}
+  %indices = s32[3] parameter(1), sharding={replicated}
+  ROOT %gather = f32[3,16] gather(%input, %indices), offset_dims={1},
+    collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1,
+    slice_sizes={1,16}, sharding={devices=[1,4]<=[4]}
+})";
+  ASSERT_OK_AND_ASSIGN(auto module,
+                       PartitionComputation(hlo_string, /*num_devices=*/4));
+  VLOG(1) << module->ToString();
+  HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(root,
+              AllOf(op::DynamicSlice(
+                        op::Gather(op::Parameter(0), op::Parameter(1)), _, _),
+                    op::Shape("f32[3,4]")));
+}
+
 TEST_P(SpmdPartitioningTest, IndexPassthroughGather) {
   absl::string_view hlo_string = R"(
 HloModule module
@@ -10194,6 +10217,39 @@ ENTRY entry {
   }
 }
 
+TEST_P(SpmdPartitioningTest,
+       PassthroughScatter_PartialReplicateOperandShardedUpdates) {
+  absl::string_view hlo_string = R"(
+HloModule module
+
+add (lhs: f32[], rhs: f32[]) -> f32[] {
+  lhs = f32[] parameter(0)
+  rhs = f32[] parameter(1)
+  ROOT sum = f32[] add(lhs, rhs)
+}
+
+ENTRY entry {
+  %input = f32[2,16] parameter(0),
+    sharding={devices=[1,2,2]<=[4] last_tile_dim_replicate}
+  %indices = s32[3] parameter(1), sharding={replicated}
+  %updates = f32[3,16] parameter(2), sharding={devices=[1,4]<=[4]}
+  ROOT %scatter = f32[2,16] scatter(%input, %indices, %updates),
+      to_apply=add,
+      update_window_dims={1},
+      inserted_window_dims={0},
+      scatter_dims_to_operand_dims={0},
+      index_vector_dim=1,
+      sharding={devices=[1,2,2]<=[4] last_tile_dim_replicate}
+})";
+  ASSERT_OK_AND_ASSIGN(auto module,
+                       PartitionComputation(hlo_string, /*num_devices=*/4));
+  VLOG(1) << module->ToString();
+  HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(root, AllOf(op::Scatter(op::Parameter(0), op::Parameter(1),
+                                      op::AllGather(op::Parameter(2))),
+                          op::Shape("f32[2,8]")));
+}
+
 TEST_P(SpmdPartitioningTest, PassthroughScatterVariadic_PartialReplicate) {
   absl::string_view hlo_string = R"(
 HloModule module
@@ -19656,6 +19712,85 @@ ENTRY entry {
       Each(*ParseSharding("{mesh['a'=2,'b'=2], [], unreduced=max{'a','b'}}")));
 }
 
+TEST_F(SpmdPartitioningV3Test, PassthroughGatherPrefixNamedSharding) {
+  absl::string_view hlo_string = R"(
+HloModule module
+
+ENTRY entry {
+  %input = f32[2,16] parameter(0),
+    sharding={mesh['x'=2,'y'=2] [{}, {'x'}]}
+  %indices = s32[3] parameter(1),
+    sharding={mesh['x'=2,'y'=2] replicated}
+  ROOT %gather = f32[3,16] gather(%input, %indices), offset_dims={1},
+    collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1,
+    slice_sizes={1,16}, sharding={mesh['x'=2,'y'=2] [{}, {'x', 'y'}]}
+})";
+  ASSERT_OK_AND_ASSIGN(auto module,
+                       PartitionComputation(hlo_string, /*num_devices=*/4));
+  HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(root,
+              AllOf(op::DynamicSlice(
+                        op::Gather(op::Parameter(0), op::Parameter(1)), _, _),
+                    op::Shape("f32[3,4]")));
+}
+
+TEST_F(SpmdPartitioningV3Test, PassthroughScatterPrefixNamedSharding) {
+  absl::string_view hlo_string = R"(
+HloModule module
+
+add (lhs: f32[], rhs: f32[]) -> f32[] {
+  lhs = f32[] parameter(0)
+  rhs = f32[] parameter(1)
+  ROOT sum = f32[] add(lhs, rhs)
+}
+
+ENTRY entry {
+  %input = f32[2,16] parameter(0),
+    sharding={mesh['x'=2,'y'=2] [{}, {'x'}]}
+  %indices = s32[3] parameter(1),
+    sharding={mesh['x'=2,'y'=2] replicated}
+  %updates = f32[3,16] parameter(2),
+    sharding={mesh['x'=2,'y'=2] [{}, {'x', 'y'}]}
+  ROOT %scatter = f32[2,16] scatter(%input, %indices, %updates),
+      to_apply=add,
+      update_window_dims={1},
+      inserted_window_dims={0},
+      scatter_dims_to_operand_dims={0},
+      index_vector_dim=1,
+      sharding={mesh['x'=2,'y'=2] [{}, {'x'}]}
+})";
+  ASSERT_OK_AND_ASSIGN(auto module,
+                       PartitionComputation(hlo_string, /*num_devices=*/4));
+  HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(root, AllOf(op::Scatter(op::Parameter(0), op::Parameter(1),
+                                      op::AllGather(op::Parameter(2))),
+                          op::Shape("f32[2,8]")));
+}
+
+TEST_F(SpmdPartitioningV3Test, IndexPassthroughGatherPrefixNamedSharding) {
+  absl::string_view hlo_string = R"(
+HloModule module
+
+ENTRY entry {
+  %input = f32[2,9,8] parameter(0),
+    sharding={mesh['x'=2,'y'=2] replicated}
+  %indices = s32[4,2,4] parameter(1),
+    sharding={mesh['x'=2,'y'=2] [{'x'}, {'y'}, {}]}
+  ROOT %gather = f32[8,4,4] gather(%input, %indices), offset_dims={0},
+    collapsed_slice_dims={0,1}, start_index_map={0,1}, index_vector_dim=1,
+    slice_sizes={1,1,8},
+    sharding={mesh['x'=2,'y'=2] [{}, {'x', 'y'}, {}]}
+})";
+  ASSERT_OK_AND_ASSIGN(auto module,
+                       PartitionComputation(hlo_string, /*num_devices=*/4));
+  HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(
+      root, AllOf(op::DynamicSlice(op::Gather(op::Parameter(0),
+                                              op::AllGather(op::Parameter(1))),
+                                   _, _, _),
+                  op::Shape("f32[8,1,4]")));
+}
+
 // Verifies that a true 1D scatter-conflict on a reduction dimension (e.g.,
 // b/512585186) successfully avoids full replication by falling back to
 // Reshard(updates) instead of failing outright and returning a nullptr.

From 1eac67a115f992af4430848d2a2acbbe3345ecc0 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 8 Oct 2026 00:10:03 -0700
Subject: [PATCH 0249/1000] Automated Code Change

PiperOrigin-RevId: 995626573
---
 tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h | 1 +
 tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h  | 1 +
 2 files changed, 2 insertions(+)

diff --git a/tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h b/tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h
index 089029e287b..5839edd5ce6 100644
--- a/tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h
+++ b/tensorflow/core/tfrt/ifrt/ifrt_loaded_variable_registry.h
@@ -23,6 +23,7 @@ limitations under the License.
 
 #include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
 #include "absl/functional/any_invocable.h"
 #include "absl/hash/hash.h"
 #include "absl/status/status.h"
diff --git a/tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h b/tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h
index d2f7bff0a30..69602678aa0 100644
--- a/tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h
+++ b/tensorflow/core/tfrt/ifrt/ifrt_restore_tensor_registry.h
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <string>
 
+#include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/status/status.h"

From 3b01044f1bff069b8bcb163ff76d8ccc4ef02f76 Mon Sep 17 00:00:00 2001
From: Claudio Basile <cbasile@google.com>
Date: Thu, 8 Oct 2026 00:11:08 -0700
Subject: [PATCH 0250/1000] Don't sink FUSE_LIMIT constants in
 FusionConstantSinking

FUSE_LIMIT marks an instruction as a fusion boundary. Constants can carry the
attribute, e.g. when constant folding replaces an instruction that has it.
FusionConstantSinking ignored the attribute and sank such constants into the
fusions of their users, which undid the boundary.

FusionConstantSinking now leaves scalar constants with FUSE_LIMIT outside loop
and output fusions. Custom fusions are unchanged: they only take must-fuse
constants, and must-fuse takes precedence over FUSE_LIMIT.

PiperOrigin-RevId: 995627184
---
 .../simplifiers/fusion_constant_sinking.cc    |  8 ++++
 .../fusion_constant_sinking_test.cc           | 38 +++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/fusion_constant_sinking.cc b/third_party/xla/xla/hlo/transforms/simplifiers/fusion_constant_sinking.cc
index 073a93217f3..e3e603a4636 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/fusion_constant_sinking.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/fusion_constant_sinking.cc
@@ -84,6 +84,14 @@ bool CanSink(HloInstruction* fusion, const HloInstruction* operand) {
     }
   }
 
+  // A FUSE_LIMIT operand is a fusion boundary. Custom fusions only take
+  // must-fuse operands (see above), and must-fuse takes precedence over
+  // FUSE_LIMIT.
+  if (!fusion->IsCustomFusion() &&
+      operand->frontend_attributes().map().contains(kFuseLimitAttr)) {
+    return false;
+  }
+
   if (fusion->operand_count() == 1) {
     return false;
   }
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/fusion_constant_sinking_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/fusion_constant_sinking_test.cc
index 5a0c015f2b4..7d9fd866853 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/fusion_constant_sinking_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/fusion_constant_sinking_test.cc
@@ -310,5 +310,43 @@ ENTRY %main.3 (x.1: bf16[1024,1024], w.1: bf16[1024,3072], out.1: bf16[1024,4096
       3);
 }
 
+// A FUSE_LIMIT constant is a fusion boundary, so it stays a fusion operand
+// while the unmarked constant is sunk.
+TEST_F(FusionConstantSinkingTest, FuseLimitConstantNoSink) {
+  std::string hlo_string = R"(
+  HloModule FuseLimitConstant
+
+    %fused_computation (param_0: f32[8,128], param_1: f32[], param_2: f32[]) -> f32[8,128] {
+      %param_0 = f32[8,128]{1,0} parameter(0)
+      %param_1 = f32[] parameter(1)
+      %broadcast.1 = f32[8,128]{1,0} broadcast(%param_1), dimensions={}
+      %multiply = f32[8,128]{1,0} multiply(%param_0, %broadcast.1)
+      %param_2 = f32[] parameter(2)
+      %broadcast.2 = f32[8,128]{1,0} broadcast(%param_2), dimensions={}
+      ROOT %add = f32[8,128]{1,0} add(%multiply, %broadcast.2)
+    }
+
+    ENTRY main {
+      p0 = f32[8,128]{1,0} parameter(0)
+      c0 = f32[] constant(8128), frontend_attributes={FUSE_LIMIT="true"}
+      c1 = f32[] constant(1)
+      ROOT out = f32[8,128]{1,0} fusion(p0, c0, c1), kind=kLoop, calls=%fused_computation
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_string));
+
+  FusionConstantSinking constant_sinking;
+
+  ASSERT_OK_AND_ASSIGN(bool result,
+                       RunHloPass(&constant_sinking, module.get()));
+
+  EXPECT_TRUE(result);
+  const HloInstruction* fusion =
+      module->entry_computation()->root_instruction();
+  ASSERT_EQ(fusion->operand_count(), 2);
+  EXPECT_EQ(fusion->operand(1)->name(), "c0");
+}
+
 }  // namespace
 }  // namespace xla

From 8ffdb122c4abc8b34874d82f12363f2ca48147f9 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 8 Oct 2026 00:12:11 -0700
Subject: [PATCH 0251/1000] Automated Code Change

PiperOrigin-RevId: 995627741
---
 third_party/xla/xla/hlo/builder/lib/BUILD          | 1 +
 third_party/xla/xla/hlo/builder/lib/matrix_test.cc | 1 +
 2 files changed, 2 insertions(+)

diff --git a/third_party/xla/xla/hlo/builder/lib/BUILD b/third_party/xla/xla/hlo/builder/lib/BUILD
index 0f91f1c295c..f634e69e265 100644
--- a/third_party/xla/xla/hlo/builder/lib/BUILD
+++ b/third_party/xla/xla/hlo/builder/lib/BUILD
@@ -326,6 +326,7 @@ xla_test(
         "//xla:literal_util",
         "//xla:shape_util",
         "//xla:types",
+        "//xla:xla_data_proto_cc",
         "//xla/hlo/builder:xla_builder",
         "//xla/hlo/testlib:test",
         "//xla/tests:client_library_test_runner_mixin",
diff --git a/third_party/xla/xla/hlo/builder/lib/matrix_test.cc b/third_party/xla/xla/hlo/builder/lib/matrix_test.cc
index 598b5a3c0d4..b9958442ae9 100644
--- a/third_party/xla/xla/hlo/builder/lib/matrix_test.cc
+++ b/third_party/xla/xla/hlo/builder/lib/matrix_test.cc
@@ -42,6 +42,7 @@ limitations under the License.
 #include "xla/tests/hlo_pjrt_test_base.h"
 #include "xla/tsl/platform/test.h"
 #include "xla/types.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla {
 namespace {

From d6bed43d904f47232012349010abf2e0d6cd50f5 Mon Sep 17 00:00:00 2001
From: Levon Ter-Grigoryan <patrios@google.com>
Date: Thu, 8 Oct 2026 00:34:06 -0700
Subject: [PATCH 0252/1000] [XLA:GPU] Fix string formatting for the
 xla_gpu_experimental_use_collective_kernels flag.

Format the values of xla_gpu_experimental_use_collective_kernels using CollectiveKernelType_Name instead of collective op types.

PiperOrigin-RevId: 995637534
---
 third_party/xla/xla/debug_options_flags.cc | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index 790b81c89c8..f5772837e00 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -1165,6 +1165,16 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
     return absl::StrJoin(collective_ops, ", ", Formatter());
   };
 
+  auto collective_kernel_type_to_string =
+      [](google::protobuf::RepeatedField<int> collective_kernels) -> std::string {
+    struct Formatter {
+      void operator()(std::string* out, int type) const {
+        absl::StrAppend(out, DebugOptions::CollectiveKernelType_Name(type));
+      }
+    };
+    return absl::StrJoin(collective_kernels, ", ", Formatter());
+  };
+
   // Custom parser for `xla_cpu_xnn_graph_fusion_mode` flag.
   auto setter_for_xla_cpu_experimental_xnn_graph_fusion_mode =
       [debug_options](absl::string_view input) {
@@ -3367,7 +3377,7 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
             return debug_options
                 ->mutable_xla_gpu_experimental_use_collective_kernels();
           }),
-      collective_op_types_to_string(
+      collective_kernel_type_to_string(
           debug_options->xla_gpu_experimental_use_collective_kernels()),
       "Experimental: comma-separated filter of collective ops that should use "
       "custom kernels (e.g. Triton one-shot / two-shot) instead of NCCL. "

From 9dbac3b51df9a54763dd7be5b2c431d3306c3a5e Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Thu, 8 Oct 2026 00:36:42 -0700
Subject: [PATCH 0253/1000] Support scalar bounds in HloEvaluator::HandleClamp

HLO `kClamp` semantics and `ShapeInference::InferClampShape` allow the `min` and
`max` bounds to be either scalars (`ShapeUtil::IsScalar`) or tensors with the
same dimensions as `operand`. Previously, `ElementwiseTernaryOpImpl` required
all three operands to have the same dimensions as the output shape, causing
`HloEvaluator` to fail when `min` or `max` was a scalar and `operand` was a
non-scalar tensor.

Allow `lhs` and `ehs` in `ElementwiseTernaryOpImpl` to be scalar and index them
at `0` (in the same-layout linear fast path) or `{}` (in the multi-index
fallback path) without allocating broadcasted literals.

Re #50116

PiperOrigin-RevId: 995638640
---
 .../xla/hlo/evaluator/hlo_evaluator_test.cc   | 49 +++++++++++++------
 .../evaluator/hlo_evaluator_typed_visitor.h   | 31 +++++++-----
 2 files changed, 53 insertions(+), 27 deletions(-)

diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc
index b4cf7b931c3..6d79d18a4b4 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc
@@ -321,25 +321,42 @@ TEST_P(HloEvaluatorBf16Test, DoesClampInt64) {
   EXPECT_TRUE(LiteralTestUtil::Equal(expected, result));
 }
 
-TEST_P(HloEvaluatorBf16Test, DISABLED_DoesClampSpecialBroadcast) {
-  auto low = LiteralUtil::CreateR0<float>(0.f);
-  auto value = LiteralUtil::CreateR2<float>({{-1.f, 0.f}, {1.f, 2.f}});
-  auto high = LiteralUtil::CreateR0<float>(1.f);
+TEST_P(HloEvaluatorBf16Test, DoesClampSpecialBroadcast) {
+  Literal low = LiteralUtil::CreateR0<float>(0.f);
+  Literal value = LiteralUtil::CreateR2<float>({{-1.f, 0.f}, {1.f, 2.f}});
+  Literal high = LiteralUtil::CreateR0<float>(1.f);
+  Literal expected = LiteralUtil::CreateR2<float>({{0.f, 0.f}, {1.f, 1.f}});
+  TestTernaryOp(HloOpcode::kClamp, std::move(expected), std::move(low),
+                std::move(value), std::move(high));
+}
 
-  Shape shape = value.shape();
-  HloComputation::Builder b(TestName());
-  auto c1 = b.AddInstruction(HloInstruction::CreateConstant(std::move(low)));
-  auto c2 = b.AddInstruction(HloInstruction::CreateConstant(std::move(value)));
-  auto c3 = b.AddInstruction(HloInstruction::CreateConstant(std::move(high)));
-  b.AddInstruction(
-      HloInstruction::CreateTernary(shape, HloOpcode::kClamp, c1, c2, c3));
-  m_->AddEntryComputation(b.Build());
+TEST_P(HloEvaluatorBf16Test, DoesClampScalarLowerBound) {
+  Literal low = LiteralUtil::CreateR0<float>(0.f);
+  Literal value = LiteralUtil::CreateR2<float>({{-1.f, 5.f}, {1.f, 4.f}});
+  Literal high = LiteralUtil::CreateR2<float>({{2.f, 4.f}, {4.f, 3.f}});
+  Literal expected = LiteralUtil::CreateR2<float>({{0.f, 4.f}, {1.f, 3.f}});
+  TestTernaryOp(HloOpcode::kClamp, std::move(expected), std::move(low),
+                std::move(value), std::move(high));
+}
 
-  TF_ASSERT_OK_AND_ASSIGN(Literal result, Evaluate());
+TEST_P(HloEvaluatorBf16Test, DoesClampScalarUpperBound) {
+  Literal low = LiteralUtil::CreateR2<float>({{0.f, 2.f}, {2.f, 0.f}});
+  Literal value = LiteralUtil::CreateR2<float>({{-1.f, 5.f}, {1.f, 4.f}});
+  Literal high = LiteralUtil::CreateR0<float>(3.f);
+  Literal expected = LiteralUtil::CreateR2<float>({{0.f, 3.f}, {2.f, 3.f}});
+  TestTernaryOp(HloOpcode::kClamp, std::move(expected), std::move(low),
+                std::move(value), std::move(high));
+}
 
-  auto expected = LiteralUtil::CreateR2<float>({{0, 0}, {1, 1}});
-
-  EXPECT_TRUE(LiteralTestUtil::Equal(expected, result));
+TEST_F(HloEvaluatorTest, DoesClampScalarBoundDifferentResultLayout) {
+  Literal low = LiteralUtil::CreateR0<int64_t>(0);
+  Literal value = LiteralUtil::CreateR2<int64_t>({{-5, 10}, {2, -1}});
+  Literal high = LiteralUtil::CreateR2<int64_t>({{4, 8}, {6, 1}});
+  Layout layout({0, 1});
+  Literal expected =
+      LiteralUtil::CreateR2WithLayout<int64_t>({{0, 8}, {2, 0}}, layout);
+  TestTernaryOp(HloOpcode::kClamp, std::move(expected), std::move(low),
+                std::move(value), std::move(high));
 }
 
 // Verifies that HloEvaluator evaluates a HLO instruction that performs select
diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h
index 95e2a3623dc..fecf11523a8 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h
@@ -2441,9 +2441,13 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     const auto* lhs = instruction->operand(0);
     const auto* rhs = instruction->operand(1);
     const auto* ehs = instruction->operand(2);
-    TF_RET_CHECK(ShapeUtil::SameDimensions(shape, lhs->shape()));
-    TF_RET_CHECK(ShapeUtil::SameDimensions(lhs->shape(), rhs->shape()));
-    TF_RET_CHECK(ShapeUtil::SameDimensions(rhs->shape(), ehs->shape()));
+    const bool is_lhs_scalar = ShapeUtil::IsScalar(lhs->shape());
+    const bool is_ehs_scalar = ShapeUtil::IsScalar(ehs->shape());
+    TF_RET_CHECK(is_lhs_scalar ||
+                 ShapeUtil::SameDimensions(shape, lhs->shape()));
+    TF_RET_CHECK(ShapeUtil::SameDimensions(shape, rhs->shape()));
+    TF_RET_CHECK(is_ehs_scalar ||
+                 ShapeUtil::SameDimensions(shape, ehs->shape()));
 
     const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs);
     const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs);
@@ -2455,9 +2459,10 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     const Layout& lhs_layout = lhs_literal.shape().layout();
     const Layout& rhs_layout = rhs_literal.shape().layout();
     const Layout& ehs_layout = ehs_literal.shape().layout();
-    bool same_layout = LayoutUtil::Equal(lhs_layout, rhs_layout) &&
-                       LayoutUtil::Equal(rhs_layout, ehs_layout) &&
-                       LayoutUtil::Equal(lhs_layout, shape.layout());
+    const bool same_layout =
+        LayoutUtil::Equal(rhs_layout, shape.layout()) &&
+        (is_lhs_scalar || LayoutUtil::Equal(lhs_layout, shape.layout())) &&
+        (is_ehs_scalar || LayoutUtil::Equal(ehs_layout, shape.layout()));
 
     if (same_layout) {
       const LhsType* lhs_data = lhs_literal.data<LhsType>().data();
@@ -2465,16 +2470,20 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
       const EhsType* ehs_data = ehs_literal.data<EhsType>().data();
       ABSL_RETURN_IF_ERROR(result.PopulateLinearParallel<ReturnT>(
           [&](int64_t linear_index, int) {
-            return ternary_op(lhs_data[linear_index], rhs_data[linear_index],
-                              ehs_data[linear_index]);
+            return ternary_op(lhs_data[is_lhs_scalar ? 0 : linear_index],
+                              rhs_data[linear_index],
+                              ehs_data[is_ehs_scalar ? 0 : linear_index]);
           }));
 
     } else {
       ABSL_RETURN_IF_ERROR(result.PopulateParallel<ReturnT>(
           [&](absl::Span<const int64_t> multi_index, int) {
-            return ternary_op(lhs_literal.Get<LhsType>(multi_index),
-                              rhs_literal.Get<RhsType>(multi_index),
-                              ehs_literal.Get<EhsType>(multi_index));
+            return ternary_op(
+                lhs_literal.Get<LhsType>(
+                    is_lhs_scalar ? absl::Span<const int64_t>{} : multi_index),
+                rhs_literal.Get<RhsType>(multi_index),
+                ehs_literal.Get<EhsType>(
+                    is_ehs_scalar ? absl::Span<const int64_t>{} : multi_index));
           }));
     }
 

From 6f1d0bcf7e3f0c93b369aa85f44608327a2bd58e Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Thu, 8 Oct 2026 01:00:08 -0700
Subject: [PATCH 0254/1000] Add runtime version checks for libnvJitLink flags

XLA binaries built with CUDA 12.8 headers still support running against older `libnvJitLink` releases in OSS down to CUDA 12.0. In `xla/stream_executor/cuda/nvjitlink.cc`, `CompileAndLinkUsingLibNvJitLink` and `GetLatestPtxIsaVersionForLibNvJitLinkImpl` passed `-split-compile` (added in CUDA 12.3) and `-no-cache` (added in CUDA 12.6) unconditionally, which causes `nvJitLinkCreate` to fail on older `libnvJitLink` libraries.

Also, `GetNvJitLinkVersion` relied on a weak-symbol check (`!nvJitLinkVersion`) to fall back to `12.0` on CUDA 12.0-12.2, which is bypassed in dynamic stub builds where `nvjitlink_stub` defines a non-weak wrapper that returns `NVJITLINK_ERROR_INTERNAL` when the symbol is missing.

So this change handles `NVJITLINK_ERROR_INTERNAL` in `GetNvJitLinkVersion` and guards `-split-compile` (`>= 12.3`) and `-no-cache` (`>= 12.6`) with runtime version checks.

PiperOrigin-RevId: 995648244
---
 .../xla/xla/stream_executor/cuda/nvjitlink.cc | 21 ++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/cuda/nvjitlink.cc b/third_party/xla/xla/stream_executor/cuda/nvjitlink.cc
index 755a3dde94c..ac654703f8f 100644
--- a/third_party/xla/xla/stream_executor/cuda/nvjitlink.cc
+++ b/third_party/xla/xla/stream_executor/cuda/nvjitlink.cc
@@ -146,7 +146,11 @@ absl::StatusOr<NvJitLinkVersion> GetNvJitLinkVersion() {
     return NvJitLinkVersion{12, 0};
   }
   unsigned int major{}, minor{};
-  RETURN_IF_NVJITLINK_ERROR(nvJitLinkVersion(&major, &minor));
+  nvJitLinkResult result = nvJitLinkVersion(&major, &minor);
+  if (result == NVJITLINK_ERROR_INTERNAL) {
+    return NvJitLinkVersion{12, 0};
+  }
+  RETURN_IF_NVJITLINK_ERROR(result);
   return NvJitLinkVersion(major, minor);
 }
 
@@ -169,8 +173,12 @@ absl::StatusOr<cuda::Assembly> CompileAndLinkUsingLibNvJitLink(
     cli_args.emplace_back("-verbose");
   }
   cli_args.emplace_back("-Xptxas=--warn-on-spills");
-  cli_args.emplace_back(absl::StrCat("-split-compile=", inputs.size()));
-  cli_args.emplace_back("-no-cache");
+  if (version >= NvJitLinkVersion{12, 3}) {
+    cli_args.emplace_back(absl::StrCat("-split-compile=", inputs.size()));
+  }
+  if (version >= NvJitLinkVersion{12, 6}) {
+    cli_args.emplace_back("-no-cache");
+  }
 
   if (options.disable_gpuasm_optimizations) {
     cli_args.emplace_back("-Xptxas=-O0");
@@ -286,9 +294,13 @@ namespace {
 
 absl::StatusOr<int> GetLatestPtxIsaVersionForLibNvJitLinkImpl() {
   absl::string_view ptx_contents = ".version 99.99";
+  absl::StatusOr<NvJitLinkVersion> version = GetNvJitLinkVersion();
   // The call to `nvJitLinkCreate` below requires an arch to be specified in
   // order to succeed.
-  std::vector<const char*> cli_args_ptrs{"-arch=sm_90a", "-no-cache"};
+  std::vector<const char*> cli_args_ptrs{"-arch=sm_90a"};
+  if (version.ok() && *version >= NvJitLinkVersion{12, 6}) {
+    cli_args_ptrs.push_back("-no-cache");
+  }
   nvJitLinkHandle link_handle = nullptr;
   nvJitLinkResult create_result =
       nvJitLinkCreate(&link_handle, /*num_args=*/cli_args_ptrs.size(),
@@ -309,7 +321,6 @@ absl::StatusOr<int> GetLatestPtxIsaVersionForLibNvJitLinkImpl() {
   }
 
   std::optional<absl::LeakCheckDisabler> disabler;
-  absl::StatusOr<NvJitLinkVersion> version = GetNvJitLinkVersion();
   if (!version.ok() || std::get<0>(*version) < 13) {
     // libnvjitlink prior to CUDA 13 has a memory leak when calling
     // nvJitLinkAddData when the input PTX is invalid.

From c3b3bec732580ee37de074380654312ceb933f82 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Thu, 8 Oct 2026 01:01:43 -0700
Subject: [PATCH 0255/1000] [XLA:GPU]: Migrate CollectiveReduceThunk to use
 PerDeviceState

PiperOrigin-RevId: 995648990
---
 .../xla/xla/backends/gpu/runtime/BUILD        |  4 +-
 .../gpu/runtime/collective_reduce_thunk.cc    | 51 +++++++++----------
 .../gpu/runtime/collective_reduce_thunk.h     | 17 +++----
 .../collective_reduce_thunk_multigpu_test.cc  |  6 +--
 .../runtime/thunk_proto_deserialization.cc    |  3 +-
 .../xla/xla/service/gpu/thunk_emitter.cc      |  1 +
 6 files changed, 37 insertions(+), 45 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index 58ce98a81a8..2f2dd157392 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -2018,6 +2018,7 @@ cc_library(
         ":all_reduce_thunk",
         ":collective_thunk",
         ":collective_thunk_proto_cc",
+        ":per_device_state",
         ":thunk",
         ":thunk_proto_cc",
         "//xla:future",
@@ -2038,14 +2039,11 @@ cc_library(
         "//xla/stream_executor:stream_executor_h",
         "//xla/tsl/platform:logging",
         "@com_google_absl//absl/base",
-        "@com_google_absl//absl/base:core_headers",
-        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/strings:string_view",
-        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
     ],
 )
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.cc
index 4ae03ebd344..8f40bfe7ce3 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.cc
@@ -26,7 +26,6 @@ limitations under the License.
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_format.h"
-#include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
 #include "xla/backends/gpu/collectives/gpu_clique_key.h"
 #include "xla/backends/gpu/collectives/gpu_collectives.h"
@@ -58,18 +57,22 @@ namespace gpu {
 CollectiveReduceThunk::CollectiveReduceThunk(ThunkInfo thunk_info,
                                              AllReduceConfig config,
                                              std::vector<Buffer> buffers,
+                                             int devices_per_host,
                                              bool has_dynamic_root)
     : AllReduceReduceScatterThunkBase(Thunk::kCollectiveReduce, thunk_info,
                                       std::move(config), std::move(buffers)),
-      has_dynamic_root_(has_dynamic_root) {}
+      has_dynamic_root_(has_dynamic_root),
+      per_device_metadata_(devices_per_host) {}
 
 CollectiveReduceThunk::CollectiveReduceThunk(
     ThunkInfo thunk_info, const HloCollectiveReduceInstruction* inst,
-    std::vector<Buffer> buffers, bool p2p_memcpy_enabled, bool has_dynamic_root)
+    std::vector<Buffer> buffers, int devices_per_host, bool p2p_memcpy_enabled,
+    bool has_dynamic_root)
     : AllReduceReduceScatterThunkBase(Thunk::kCollectiveReduce, thunk_info,
                                       GetAllReduceConfigInst(inst),
                                       std::move(buffers)),
-      has_dynamic_root_(has_dynamic_root) {}
+      has_dynamic_root_(has_dynamic_root),
+      per_device_metadata_(devices_per_host) {}
 
 /*static*/ absl::Status CollectiveReduceThunk::CheckImplementable(
     const HloCollectiveReduceInstruction* inst, int64_t replica_count,
@@ -105,22 +108,18 @@ absl::Status CollectiveReduceThunk::InitializeCollective(
     return absl::OkStatus();
   }
   se::StreamExecutor* executor = params.executor;
-  absl::MutexLock lock(mutex_);
-  std::unique_ptr<CollectiveReduceMetadata>& metadata =
-      per_executor_metadata_[executor];
-  if (metadata == nullptr) {
-    metadata = std::make_unique<CollectiveReduceMetadata>();
-  }
-  if (metadata->reduce_roots == nullptr) {
-    // The last buffer holds the runtime-selected root ranks (one S32 per
-    // reduce); all other buffers are the data being reduced.
-    metadata->num_roots = buffers().size() - 1;
-    ABSL_ASSIGN_OR_RETURN(
-        std::unique_ptr<se::MemoryAllocation> alloc,
-        executor->HostMemoryAllocate(metadata->num_roots * sizeof(int32_t)));
-    metadata->reduce_roots = std::move(alloc);
-  }
-  return absl::OkStatus();
+  return per_device_metadata_.GetOrCreateAndInitialize(
+      executor->device_ordinal(),
+      [&](CollectiveReduceMetadata* metadata) -> absl::Status {
+        // The last buffer holds the runtime-selected root ranks (one S32 per
+        // reduce); all other buffers are the data being reduced.
+        metadata->num_roots = buffers().size() - 1;
+        ABSL_ASSIGN_OR_RETURN(std::unique_ptr<se::MemoryAllocation> alloc,
+                         executor->HostMemoryAllocate(metadata->num_roots *
+                                                      sizeof(int32_t)));
+        metadata->reduce_roots = std::move(alloc);
+        return absl::OkStatus();
+      });
 }
 
 absl::Status CollectiveReduceThunk::RunCollective(
@@ -129,11 +128,8 @@ absl::Status CollectiveReduceThunk::RunCollective(
   ABSL_ASSIGN_OR_RETURN(std::vector<DeviceBufferPair> device_buffers,
                    ConvertToDeviceBuffers(params.buffer_allocations, buffers(),
                                           config_.config.operand_element_type));
-  CollectiveReduceMetadata* metadata = nullptr;
-  {
-    absl::MutexLock lock(mutex_);
-    metadata = per_executor_metadata_[stream.parent()].get();
-  }
+  CollectiveReduceMetadata* metadata =
+      per_device_metadata_.Find(stream.parent()->device_ordinal());
   return RunCollectiveReduce(config_.reduction_kind, device_buffers, stream,
                              comm, metadata, has_dynamic_root_);
 }
@@ -141,7 +137,8 @@ absl::Status CollectiveReduceThunk::RunCollective(
 absl::StatusOr<std::unique_ptr<CollectiveReduceThunk>>
 CollectiveReduceThunk::FromProto(
     ThunkInfo thunk_info, const CollectiveReduceThunkProto& thunk_proto,
-    absl::Span<const BufferAllocation> buffer_allocations) {
+    absl::Span<const BufferAllocation> buffer_allocations,
+    int devices_per_host) {
   std::vector<CollectiveThunk::Buffer> buffers;
   buffers.reserve(thunk_proto.buffers_size());
   for (const CollectiveBufferProto& proto : thunk_proto.buffers()) {
@@ -159,7 +156,7 @@ CollectiveReduceThunk::FromProto(
 
   return std::make_unique<CollectiveReduceThunk>(
       std::move(thunk_info), AllReduceConfig{config, reduction_kind},
-      std::move(buffers), thunk_proto.has_dynamic_root());
+      std::move(buffers), devices_per_host, thunk_proto.has_dynamic_root());
 }
 
 absl::StatusOr<ThunkProto> CollectiveReduceThunk::ToProto() const {
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.h b/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.h
index c08fcafb559..0d59fa7a066 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.h
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.h
@@ -20,16 +20,14 @@ limitations under the License.
 #include <memory>
 #include <vector>
 
-#include "absl/base/thread_annotations.h"
-#include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
-#include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
 #include "xla/backends/gpu/collectives/gpu_clique_key.h"
 #include "xla/backends/gpu/runtime/all_reduce_thunk.h"
 #include "xla/backends/gpu/runtime/collective_thunk.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/backends/gpu/runtime/thunk.pb.h"
 #include "xla/core/collectives/communicator.h"
 #include "xla/core/collectives/reduction_kind.h"
@@ -37,7 +35,6 @@ limitations under the License.
 #include "xla/service/buffer_assignment.h"
 #include "xla/stream_executor/memory_allocation.h"
 #include "xla/stream_executor/stream.h"
-#include "xla/stream_executor/stream_executor.h"
 
 namespace xla {
 namespace gpu {
@@ -57,11 +54,11 @@ struct CollectiveReduceMetadata {
 class CollectiveReduceThunk : public AllReduceReduceScatterThunkBase {
  public:
   CollectiveReduceThunk(ThunkInfo thunk_info, AllReduceConfig config,
-                        std::vector<Buffer> buffers,
+                        std::vector<Buffer> buffers, int devices_per_host,
                         bool has_dynamic_root = false);
   CollectiveReduceThunk(ThunkInfo thunk_info,
                         const HloCollectiveReduceInstruction* inst,
-                        std::vector<Buffer> buffers,
+                        std::vector<Buffer> buffers, int devices_per_host,
                         bool p2p_memcpy_enabled = false,
                         bool has_dynamic_root = false);
 
@@ -76,7 +73,8 @@ class CollectiveReduceThunk : public AllReduceReduceScatterThunkBase {
 
   static absl::StatusOr<std::unique_ptr<CollectiveReduceThunk>> FromProto(
       ThunkInfo thunk_info, const CollectiveReduceThunkProto& thunk_proto,
-      absl::Span<const BufferAllocation> buffer_allocations);
+      absl::Span<const BufferAllocation> buffer_allocations,
+      int devices_per_host);
 
   absl::StatusOr<ThunkProto> ToProto() const override;
 
@@ -92,10 +90,7 @@ class CollectiveReduceThunk : public AllReduceReduceScatterThunkBase {
 
  private:
   const bool has_dynamic_root_;
-  mutable absl::Mutex mutex_;
-  absl::flat_hash_map<se::StreamExecutor*,
-                      std::unique_ptr<CollectiveReduceMetadata>>
-      per_executor_metadata_ ABSL_GUARDED_BY(mutex_);
+  PerDeviceState<CollectiveReduceMetadata> per_device_metadata_;
 };
 
 // Runs a reduce over `buffers` on `stream`. When `has_dynamic_root` is set the
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk_multigpu_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk_multigpu_test.cc
index f903447f705..e8c0abc5141 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk_multigpu_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk_multigpu_test.cc
@@ -89,9 +89,9 @@ static CollectiveReduceThunk MakeThunk(
                                               .destination_memory_space = 0});
   }
   const bool has_dynamic_root = alloc_root != nullptr;
-  return CollectiveReduceThunk(Thunk::ThunkInfo(),
-                               MakeCollectiveReduceConfig(has_dynamic_root),
-                               buffers, has_dynamic_root);
+  return CollectiveReduceThunk(
+      Thunk::ThunkInfo(), MakeCollectiveReduceConfig(has_dynamic_root), buffers,
+      /*devices_per_host=*/kNumDevices, has_dynamic_root);
 }
 
 using DeviceTestSlot = CollectiveThunkMultiGpuTestState;
diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
index fbf7a5b2ecd..3809e0687b8 100644
--- a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
@@ -313,7 +313,8 @@ absl::StatusOr<std::unique_ptr<Thunk>> DeserializeThunkProtoImpl(
     case ThunkProto::kCollectiveReduceThunk:
       return CollectiveReduceThunk::FromProto(
           std::move(thunk_info), thunk_proto.collective_reduce_thunk(),
-          buffer_allocations);
+          buffer_allocations,
+          gpu_topology.has_value() ? gpu_topology->num_devices_per_host() : 0);
     case ThunkProto::kCollectiveGroupThunk:
       return CollectiveGroupThunk::FromProto(
           std::move(thunk_info), thunk_proto.collective_group_thunk(),
diff --git a/third_party/xla/xla/service/gpu/thunk_emitter.cc b/third_party/xla/xla/service/gpu/thunk_emitter.cc
index f029c04cfdc..836ea29ac68 100644
--- a/third_party/xla/xla/service/gpu/thunk_emitter.cc
+++ b/third_party/xla/xla/service/gpu/thunk_emitter.cc
@@ -2259,6 +2259,7 @@ Future<ThunkSequence> ThunkEmitter::EmitCollective(
     // trailing root-rank buffer specially at run time.
     thunks = ThunkSequence::Of<CollectiveThunkType>(
         info, inst, /*buffers=*/std::move(buffers),
+        ir_emitter_context_->gpu_topology().num_devices_per_host(),
         ir_emitter_context_->debug_options().xla_gpu_use_memcpy_local_p2p(),
         has_dynamic_root);
   } else if constexpr (std::is_same_v<CollectiveThunkType,

From 3829adb7956f5f696304b0d01d5342b4ff6158df Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Thu, 8 Oct 2026 01:01:54 -0700
Subject: [PATCH 0256/1000] Guard cuStreamBeginCaptureToGraph with runtime
 driver check

`cuStreamBeginCaptureToGraph` was introduced in CUDA 12.3, while XLA binaries compiled with CUDA 12.8 headers still support CUDA 12.0 drivers at runtime. Unlike `CudaCommandBuffer::Create`, `CudaStream::CaptureHandle::BeginCapture` in `xla/stream_executor/cuda/cuda_stream.cc` did not check the runtime driver version before calling `cuStreamBeginCaptureToGraph`.

Therefore this change checks `driver_version() >= SemanticVersion{12, 3, 0}` at the start of `BeginCapture` and skips the stream capture unit tests in `xla/stream_executor/cuda/cuda_stream_test.cc` on drivers older than CUDA 12.3.

PiperOrigin-RevId: 995649059
---
 .../xla/xla/stream_executor/cuda/cuda_stream.cc      |  6 ++++++
 .../xla/xla/stream_executor/cuda/cuda_stream_test.cc | 12 ++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_stream.cc b/third_party/xla/xla/stream_executor/cuda/cuda_stream.cc
index 4cac49f7987..cfc23126c61 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_stream.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_stream.cc
@@ -204,6 +204,12 @@ CudaStream::CaptureHandle::BeginCapture(CudaStream* stream, CUgraph graph,
                                         const CUgraphEdgeData* dependency_data,
                                         size_t num_dependencies,
                                         CUstreamCaptureMode mode) {
+  if (stream->parent()->GetDeviceDescription().driver_version() <
+      SemanticVersion{12, 3, 0}) {
+    return absl::UnimplementedError(
+        "Stream capture to graph requires CUDA driver 12.3 or newer");
+  }
+
   if (stream->type_ == CudaStreamType::kCudaGraphCapture) {
     return absl::FailedPreconditionError(
         "Cannot begin capture on a capture stream.");
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_stream_test.cc b/third_party/xla/xla/stream_executor/cuda/cuda_stream_test.cc
index b882d429574..e50a64268bd 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_stream_test.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_stream_test.cc
@@ -356,6 +356,12 @@ TEST_F(CudaStreamTest, DoHostCallbackWithStatusError) {
 }
 
 TEST_F(CudaStreamTest, CaptureStreamIsolation) {
+  if (executor_->GetDeviceDescription().driver_version() <
+      SemanticVersion{12, 3, 0}) {
+    GTEST_SKIP()
+        << "Stream capture to graph requires CUDA driver 12.3 or newer";
+  }
+
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<CudaStream> stream,
                        CudaStream::Create(executor_,
                                           /*priority=*/std::nullopt));
@@ -401,6 +407,12 @@ TEST_F(CudaStreamTest, CaptureStreamIsolation) {
 }
 
 TEST_F(CudaStreamTest, NestedCaptureFails) {
+  if (executor_->GetDeviceDescription().driver_version() <
+      SemanticVersion{12, 3, 0}) {
+    GTEST_SKIP()
+        << "Stream capture to graph requires CUDA driver 12.3 or newer";
+  }
+
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<CudaStream> stream,
                        CudaStream::Create(executor_,
                                           /*priority=*/std::nullopt));

From 64f307c3e4bebd7b6a6df598c30b724fe7f396e7 Mon Sep 17 00:00:00 2001
From: bodapatisaikrishna <saikrishnabodapati@gmail.com>
Date: Thu, 8 Oct 2026 14:36:40 +0530
Subject: [PATCH 0257/1000] Validate conv2d input and filter ranks in C++
 kernel

Move rank validation from python wrapper into Conv2DOp kernel boundary
to prevent crashes on raw ops and dynamic shapes, and update tests.

Signed-off-by: bodapatisaikrishna <saikrishnabodapati@gmail.com>
---
 tensorflow/core/kernels/conv_ops_impl.h       | 13 ++++++++++
 .../kernel_tests/nn_ops/conv_ops_test.py      | 26 +++++++++++++++++--
 tensorflow/python/ops/nn_ops.py               | 19 +-------------
 3 files changed, 38 insertions(+), 20 deletions(-)

diff --git a/tensorflow/core/kernels/conv_ops_impl.h b/tensorflow/core/kernels/conv_ops_impl.h
index 36fad88300f..ae96d206ef0 100644
--- a/tensorflow/core/kernels/conv_ops_impl.h
+++ b/tensorflow/core/kernels/conv_ops_impl.h
@@ -663,6 +663,19 @@ class Conv2DOp : public BinaryOp<T> {
     // [ filter_rows, filter_cols, in_depth, out_depth]
     const Tensor& filter = context->input(1);
 
+    const int required_dims =
+        params_.data_format == TensorFormat::FORMAT_NCHW_VECT_C ? 5 : 4;
+    OP_REQUIRES(
+        context, input.dims() == required_dims,
+        absl::InvalidArgumentError(
+            absl::StrCat("convolution input must be ", required_dims,
+                         "-dimensional: ", input.shape().DebugString())));
+    OP_REQUIRES(
+        context, filter.dims() == required_dims,
+        absl::InvalidArgumentError(absl::StrCat(
+            "convolution filter must be ", required_dims,
+            "-dimensional: ", filter.shape().DebugString())));
+
     Conv2DDimensions dimensions;
     OP_REQUIRES_OK(context,
                    ComputeConv2DDimension(params_, input, filter, &dimensions));
diff --git a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
index 603a5e332da..9beeafb6fbc 100644
--- a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
@@ -47,6 +47,8 @@ from tensorflow.python.platform import test
 from tensorflow.python.platform import tf_logging
 from tensorflow.python.util.compat import collections_abc
 
+raw_ops = gen_nn_ops
+
 
 def GetShrunkInceptionShapes(shrink=10):
   """Iterator for smaller versions of convolution shapes in 2015 Inception.
@@ -3133,7 +3135,10 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
     #           padding=[[0, 0], [2, 2], [2, 2], [0, 0]]))
 
     # Invalid input or filter rank (< 4) in eager/graph execution.
-    with self.assertRaisesRegex(ValueError, "`input.shape.rank` must be at least 4"):
+    with self.assertRaisesRegex(
+        (ValueError, errors_impl.InvalidArgumentError),
+        r"(must be (at least )?4-dimensional|input\.shape\.rank)",
+    ):
       self.evaluate(
           nn_ops.conv2d(
               np.ones([1, 1, 3], dtype=np.float32),
@@ -3142,7 +3147,10 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
               padding="VALID",
           )
       )
-    with self.assertRaisesRegex(ValueError, "`filter.shape.rank` must be at least 4"):
+    with self.assertRaisesRegex(
+        (ValueError, errors_impl.InvalidArgumentError),
+        r"(must be (at least )?4-dimensional|filter\.shape\.rank)",
+    ):
       self.evaluate(
           nn_ops.conv2d(
               np.ones([1, 1, 3, 3], dtype=np.float32),
@@ -3152,6 +3160,20 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
           )
       )
 
+    # Ensure raw ops are protected by the C++ kernel boundary.
+    with self.assertRaisesRegex(
+        (ValueError, errors_impl.InvalidArgumentError),
+        r"(must be (at least )?4-dimensional|input\.shape\.rank)",
+    ):
+      self.evaluate(
+          raw_ops.Conv2D(
+              input=np.ones([1, 1, 3], dtype=np.float32),
+              filter=np.ones([1, 2, 2, 1], dtype=np.float32),
+              strides=[1, 1, 1, 1],
+              padding="VALID",
+          )
+      )
+
     # Filter dimensions must be greater than 0.
     with self.assertRaisesRegex(
         errors_impl.InvalidArgumentError, "filter must not have zero elements"
diff --git a/tensorflow/python/ops/nn_ops.py b/tensorflow/python/ops/nn_ops.py
index 1913cc46568..1e43bd92edb 100644
--- a/tensorflow/python/ops/nn_ops.py
+++ b/tensorflow/python/ops/nn_ops.py
@@ -2668,24 +2668,7 @@ def conv2d(  # pylint: disable=redefined-builtin,dangerous-default-value
   ndims = getattr(shape, "ndims", -1)
   if ndims == -1:
     ndims = len(shape)
-  if ndims is not None and ndims < 4:
-    raise ValueError(
-        f"`input.shape.rank` must be at least 4. "
-        f"Received: input.shape={shape} with rank {ndims}."
-    )
-
-  filter_shape = getattr(filter, "shape", None)
-  if filter_shape is not None:
-    filter_ndims = getattr(filter_shape, "ndims", -1)
-    if filter_ndims == -1:
-      filter_ndims = len(filter_shape)
-    if filter_ndims is not None and filter_ndims < 4:
-      raise ValueError(
-          f"`filter.shape.rank` must be at least 4. "
-          f"Received: filter.shape={filter_shape} with rank {filter_ndims}."
-      )
-
-  if ndims in (4, None):
+  if ndims in (4, 3, 2, 1, 0, None):
     # We avoid calling squeeze_batch_dims to reduce extra python function
     # call slowdown in eager mode.  This branch doesn't require reshapes.
     return gen_nn_ops.conv2d(

From 98fe6d549fb96e1e1022d69c4e91a0ff81a1bd72 Mon Sep 17 00:00:00 2001
From: nishad shabbir <nishadshabbir28@gmail.com>
Date: Thu, 8 Oct 2026 14:53:33 +0530
Subject: [PATCH 0258/1000] Convert the inputs before reading a dtype in
 tf.image.psnr

psnr() casts max_val to the dtype of its first argument:

    max_val = math_ops.cast(max_val, a.dtype)

one line before `a` is converted by convert_image_dtype(), so a plain
Python list fails with

    AttributeError: 'list' object has no attribute 'dtype'

for input the op accepts as a tensor or a numpy array, both of which
have a dtype. Other ops in this module, tf.image.adjust_brightness
among them, take the same list.

Convert `a` and `b` at the top of the name scope. Inputs that already
worked are unaffected, since converting a tensor is a no-op, and an
integer list now goes down the same integer image path a tensor of the
same dtype already takes.
---
 tensorflow/python/ops/image_ops_impl.py |  4 ++++
 tensorflow/python/ops/image_ops_test.py | 13 +++++++++++++
 2 files changed, 17 insertions(+)

diff --git a/tensorflow/python/ops/image_ops_impl.py b/tensorflow/python/ops/image_ops_impl.py
index 23fdd0c3ad6..0f0ac1789ac 100644
--- a/tensorflow/python/ops/image_ops_impl.py
+++ b/tensorflow/python/ops/image_ops_impl.py
@@ -4312,6 +4312,10 @@ def psnr(a, b, max_val, name=None):
     and shape [batch_size, 1].
   """
   with ops.name_scope(name, 'PSNR', [a, b]):
+    # Convert first: max_val is cast to the dtype of `a` below, which an
+    # unconverted value does not have.
+    a = ops.convert_to_tensor(a, name='a')
+    b = ops.convert_to_tensor(b, name='b')
     # Need to convert the images to float32.  Scale max_val accordingly so that
     # PSNR is computed correctly.
     max_val = math_ops.cast(max_val, a.dtype)
diff --git a/tensorflow/python/ops/image_ops_test.py b/tensorflow/python/ops/image_ops_test.py
index 399b93f6295..569544347d8 100644
--- a/tensorflow/python/ops/image_ops_test.py
+++ b/tensorflow/python/ops/image_ops_test.py
@@ -6137,6 +6137,19 @@ class PSNRTest(test_util.TensorFlowTestCase):
       tf_psnr = self.evaluate(image_ops.psnr(tf_image1, tf_image2, 1.0, "psnr"))
       self.assertAllClose(psnr, tf_psnr, atol=0.001)
 
+  def testPSNRArrayLikeInput(self):
+    # max_val is cast to the dtype of the first argument, so an unconverted
+    # value used to fail with AttributeError before it was converted.
+    image1 = self._RandomImage((8, 8, 1), 1)
+    image2 = self._RandomImage((8, 8, 1), 1)
+    expected = self._PSNR_NumPy(image1, image2, 1)
+
+    with self.cached_session():
+      for a, b in ((image1.tolist(), image2.tolist()),
+                   (image1.astype(np.float32), image2.astype(np.float32))):
+        tf_psnr = self.evaluate(image_ops.psnr(a, b, 1.0, "psnr"))
+        self.assertAllClose(expected, tf_psnr, atol=0.001)
+
   def testPSNRMultiImage(self):
     image1 = self._RandomImage((10, 8, 8, 1), 1)
     image2 = self._RandomImage((10, 8, 8, 1), 1)

From e3ba662eaa7c1271ec270e8f16864dd962d6f96b Mon Sep 17 00:00:00 2001
From: Pablo Zimmermann <karupayun@google.com>
Date: Thu, 8 Oct 2026 01:02:02 -0700
Subject: [PATCH 0259/1000] Restore Triton dot occupancy pruning config
 threshold to 24.

Lowering `kMinConfigsForOccupancyOptimization` from 24 to 15 in cl/988134846
activated `EliminateLowOccupancyConfigs()` during default (non-exhaustive)
autotuning across all GPU architectures, including pre-Hopper GPUs. Because it
discards any tile producing fewer blocks than the GPU SM count, it pruned fast
128x64, 64x128, and 128x256 TMA tiles on H100/B200 (132-148 SMs) and collapsed
small/medium GEMMs to a single 16x8 fallback tile.

This regressed pure-Triton GEMM execution by up to 1.4x-2.0x and caused XLA to
unfuse `__triton_gemm` into `CUBLASLT_FISSION`, regressing end-to-end GPU
execution time by up to 1.43x while artificially inflating cost-model Top-1
efficiency against the degraded autotuner baseline.

PiperOrigin-RevId: 995649126
---
 .../gpu/autotuner/triton/dot_search_space.cc      |  4 +---
 .../gpu/autotuner/triton/dot_search_space_test.cc | 15 ---------------
 2 files changed, 1 insertion(+), 18 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space.cc b/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space.cc
index 24ebe2acc66..34b10bd27fd 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space.cc
@@ -639,9 +639,7 @@ void TritonDotFusionSearchSpace::EliminateLowOccupancyConfigs(
     return;
   }
 
-  // Output tilings alone typically produce 12-20 configs (without split-K
-  // multipliers). We only skip optimization for very small search spaces.
-  constexpr int kMinConfigsForOccupancyOptimization = 15;
+  constexpr int kMinConfigsForOccupancyOptimization = 24;
   if (configs.size() < kMinConfigsForOccupancyOptimization) {
     VLOG(10) << "Skipping occupancy optimization for small search spaces. "
                 "Configs size: "
diff --git a/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space_test.cc b/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space_test.cc
index 43fad93a057..53f8c09aa02 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/triton/dot_search_space_test.cc
@@ -266,21 +266,6 @@ TEST_F(DotSearchSpaceTest, ScalesDesiredWarpsWithHardwareDerivedRidgePoint) {
               ::testing::HasSubstr("desired_total_warps: 1056"));
 }
 
-TEST_F(DotSearchSpaceTest, EliminatesLowOccupancyTilesForMediumProblem) {
-  ASSERT_OK_AND_ASSIGN(
-      std::unique_ptr<VerifiedHloModule> module,
-      GetDefaultDotModule(/*lhs_parallel_dim=*/1024, /*rhs_parallel_dim=*/1024,
-                          /*contracting_dim=*/1024));
-  TritonDotFusionSearchSpace search_space = MakeSearchSpace(module.get());
-
-  // On 1024x1024 with 132 cores, tiles >= 128x128 yield <= 64 result tiles
-  // (< 132 cores) and should be eliminated in non-exhaustive mode.
-  EXPECT_THAT(
-      search_space.GenerateConfigs(),
-      AllOf(Not(IsEmpty()),
-            Not(Contains(AllOf(BlockMIs(Ge(128)), BlockNIs(Ge(128)))))));
-}
-
 TEST_F(DotSearchSpaceTest, ReturnsValidConfigList) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        GetDefaultDotModule());

From c892249d134925820cfd09a6efbd5be306bb8e42 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 8 Oct 2026 01:04:00 -0700
Subject: [PATCH 0260/1000] Automated Code Change

PiperOrigin-RevId: 995650150
---
 tensorflow/lite/delegates/gpu/common/tasks/BUILD          | 8 ++++++--
 .../lite/delegates/gpu/common/tasks/reduce_test_util.cc   | 2 ++
 .../lite/delegates/gpu/common/tasks/relu_test_util.cc     | 1 +
 .../delegates/gpu/common/tasks/resampler_test_util.cc     | 2 ++
 .../lite/delegates/gpu/common/tasks/reshape_test_util.cc  | 1 +
 tensorflow/lite/delegates/gpu/common/tasks/resize.cc      | 2 ++
 .../lite/delegates/gpu/common/tasks/resize_test_util.cc   | 1 +
 .../delegates/gpu/common/tasks/select_v2_test_util.cc     | 1 +
 .../lite/delegates/gpu/common/tasks/split_test_util.cc    | 1 +
 .../lite/delegates/gpu/common/tasks/strided_slice.cc      | 2 ++
 .../lite/delegates/gpu/common/tasks/strided_slice.h       | 1 +
 .../delegates/gpu/common/tasks/strided_slice_test_util.cc | 1 +
 tensorflow/lite/delegates/gpu/common/tasks/tile.cc        | 1 -
 tensorflow/lite/delegates/gpu/common/tasks/transpose.cc   | 1 -
 .../delegates/gpu/common/tasks/transpose_test_util.cc     | 1 +
 tensorflow/lite/delegates/gpu/common/tasks/winograd.cc    | 1 -
 tensorflow/lite/delegates/gpu/common/tasks/winograd.h     | 1 +
 .../lite/delegates/gpu/common/tasks/winograd_test_util.cc | 2 ++
 18 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/tensorflow/lite/delegates/gpu/common/tasks/BUILD b/tensorflow/lite/delegates/gpu/common/tasks/BUILD
index 3b79eae0130..5ee95b9daf6 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/BUILD
+++ b/tensorflow/lite/delegates/gpu/common/tasks/BUILD
@@ -967,6 +967,8 @@ cc_library(
         "//tensorflow/lite/delegates/gpu/common:types",
         "//tensorflow/lite/delegates/gpu/common/task:gpu_operation",
         "//tensorflow/lite/delegates/gpu/common/task:work_group_picking",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
     ],
 )
 
@@ -1006,6 +1008,7 @@ cc_library(
         "//tensorflow/lite/delegates/gpu/common:operations",
         "//tensorflow/lite/delegates/gpu/common/task:testing_util",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
     ],
 )
 
@@ -1110,6 +1113,8 @@ cc_library(
         "//tensorflow/lite/delegates/gpu/common:types",
         "//tensorflow/lite/delegates/gpu/common/task:gpu_operation",
         "//tensorflow/lite/delegates/gpu/common/task:work_group_picking",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
     ],
 )
 
@@ -1160,7 +1165,6 @@ cc_library(
         "//tensorflow/lite/delegates/gpu/common:types",
         "//tensorflow/lite/delegates/gpu/common/task:gpu_operation",
         "//tensorflow/lite/delegates/gpu/common/task:work_group_picking",
-        "@com_google_absl//absl/strings",
     ],
 )
 
@@ -1190,7 +1194,7 @@ cc_library(
         "//tensorflow/lite/delegates/gpu/common/task:gpu_operation",
         "//tensorflow/lite/delegates/gpu/common/task:tensor_desc",
         "//tensorflow/lite/delegates/gpu/common/task:work_group_picking",
-        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/status",
     ],
 )
 
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/reduce_test_util.cc b/tensorflow/lite/delegates/gpu/common/tasks/reduce_test_util.cc
index 1a5f46dfa74..712182683d7 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/reduce_test_util.cc
+++ b/tensorflow/lite/delegates/gpu/common/tasks/reduce_test_util.cc
@@ -16,6 +16,8 @@ limitations under the License.
 #include "tensorflow/lite/delegates/gpu/common/tasks/reduce_test_util.h"
 
 #include <memory>
+#include <set>
+#include <utility>
 #include <vector>
 
 #include "tensorflow/lite/delegates/gpu/common/operations.h"
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/relu_test_util.cc b/tensorflow/lite/delegates/gpu/common/tasks/relu_test_util.cc
index 09a8c2110c2..5672b0cd0f3 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/relu_test_util.cc
+++ b/tensorflow/lite/delegates/gpu/common/tasks/relu_test_util.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/lite/delegates/gpu/common/tasks/relu_test_util.h"
 
 #include <memory>
+#include <utility>
 #include <vector>
 
 #include "tensorflow/lite/delegates/gpu/common/operations.h"
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/resampler_test_util.cc b/tensorflow/lite/delegates/gpu/common/tasks/resampler_test_util.cc
index f825781c5ee..7822322f069 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/resampler_test_util.cc
+++ b/tensorflow/lite/delegates/gpu/common/tasks/resampler_test_util.cc
@@ -15,7 +15,9 @@ limitations under the License.
 
 #include "tensorflow/lite/delegates/gpu/common/tasks/resampler_test_util.h"
 
+#include <cmath>
 #include <memory>
+#include <utility>
 #include <vector>
 
 #include "tensorflow/lite/delegates/gpu/common/operations.h"
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/reshape_test_util.cc b/tensorflow/lite/delegates/gpu/common/tasks/reshape_test_util.cc
index 81ced961907..394b68f672c 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/reshape_test_util.cc
+++ b/tensorflow/lite/delegates/gpu/common/tasks/reshape_test_util.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/lite/delegates/gpu/common/tasks/reshape_test_util.h"
 
 #include <memory>
+#include <utility>
 #include <vector>
 
 #include "tensorflow/lite/delegates/gpu/common/operations.h"
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/resize.cc b/tensorflow/lite/delegates/gpu/common/tasks/resize.cc
index f229f2bab4b..23f4dcdac0b 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/resize.cc
+++ b/tensorflow/lite/delegates/gpu/common/tasks/resize.cc
@@ -18,6 +18,8 @@ limitations under the License.
 #include <string>
 #include <utility>
 
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
 #include "tensorflow/lite/delegates/gpu/common/operations.h"
 #include "tensorflow/lite/delegates/gpu/common/task/work_group_picking.h"
 
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/resize_test_util.cc b/tensorflow/lite/delegates/gpu/common/tasks/resize_test_util.cc
index 98d1c00d553..d3656c34e0a 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/resize_test_util.cc
+++ b/tensorflow/lite/delegates/gpu/common/tasks/resize_test_util.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/lite/delegates/gpu/common/tasks/resize_test_util.h"
 
 #include <memory>
+#include <utility>
 #include <vector>
 
 #include "tensorflow/lite/delegates/gpu/common/operations.h"
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/select_v2_test_util.cc b/tensorflow/lite/delegates/gpu/common/tasks/select_v2_test_util.cc
index 68ffbce1fa3..0b063ee0d48 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/select_v2_test_util.cc
+++ b/tensorflow/lite/delegates/gpu/common/tasks/select_v2_test_util.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/status/status.h"
+#include "absl/status/status_macros.h"
 #include "tensorflow/lite/delegates/gpu/common/operations.h"
 #include "tensorflow/lite/delegates/gpu/common/tasks/select_v2.h"
 
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/split_test_util.cc b/tensorflow/lite/delegates/gpu/common/tasks/split_test_util.cc
index e2b119ce9b8..2ffc48b63a6 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/split_test_util.cc
+++ b/tensorflow/lite/delegates/gpu/common/tasks/split_test_util.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/lite/delegates/gpu/common/tasks/split_test_util.h"
 
 #include <memory>
+#include <utility>
 #include <vector>
 
 #include "tensorflow/lite/delegates/gpu/common/operations.h"
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/strided_slice.cc b/tensorflow/lite/delegates/gpu/common/tasks/strided_slice.cc
index 57c095cf39d..e4336d388d8 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/strided_slice.cc
+++ b/tensorflow/lite/delegates/gpu/common/tasks/strided_slice.cc
@@ -18,6 +18,8 @@ limitations under the License.
 #include <string>
 #include <utility>
 
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
 #include "tensorflow/lite/delegates/gpu/common/task/work_group_picking.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/strided_slice.h b/tensorflow/lite/delegates/gpu/common/tasks/strided_slice.h
index 15868a8e7ec..c12146d842d 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/strided_slice.h
+++ b/tensorflow/lite/delegates/gpu/common/tasks/strided_slice.h
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include <string>
 
+#include "absl/status/status.h"
 #include "tensorflow/lite/delegates/gpu/common/operations.h"
 #include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
 #include "tensorflow/lite/delegates/gpu/common/types.h"
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/strided_slice_test_util.cc b/tensorflow/lite/delegates/gpu/common/tasks/strided_slice_test_util.cc
index bd37183c192..eb21ef5c722 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/strided_slice_test_util.cc
+++ b/tensorflow/lite/delegates/gpu/common/tasks/strided_slice_test_util.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/lite/delegates/gpu/common/tasks/strided_slice_test_util.h"
 
 #include <memory>
+#include <utility>
 #include <vector>
 
 #include "tensorflow/lite/delegates/gpu/common/operations.h"
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/tile.cc b/tensorflow/lite/delegates/gpu/common/tasks/tile.cc
index 8bf28f1814b..13d7bf3d522 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/tile.cc
+++ b/tensorflow/lite/delegates/gpu/common/tasks/tile.cc
@@ -16,7 +16,6 @@ limitations under the License.
 #include "tensorflow/lite/delegates/gpu/common/tasks/tile.h"
 
 #include <string>
-#include <utility>
 #include <vector>
 
 namespace tflite {
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/transpose.cc b/tensorflow/lite/delegates/gpu/common/tasks/transpose.cc
index 18d99340d99..3526fb04326 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/transpose.cc
+++ b/tensorflow/lite/delegates/gpu/common/tasks/transpose.cc
@@ -17,7 +17,6 @@ limitations under the License.
 
 #include <string>
 
-#include "absl/strings/str_cat.h"
 #include "tensorflow/lite/delegates/gpu/common/task/work_group_picking.h"
 
 namespace tflite {
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/transpose_test_util.cc b/tensorflow/lite/delegates/gpu/common/tasks/transpose_test_util.cc
index 72e43e9e0c4..cc164903d09 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/transpose_test_util.cc
+++ b/tensorflow/lite/delegates/gpu/common/tasks/transpose_test_util.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "tensorflow/lite/delegates/gpu/common/tasks/transpose_test_util.h"
 
 #include <memory>
+#include <utility>
 #include <vector>
 
 #include "tensorflow/lite/delegates/gpu/common/operations.h"
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/winograd.cc b/tensorflow/lite/delegates/gpu/common/tasks/winograd.cc
index d87a41c4c21..6a6d0853a20 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/winograd.cc
+++ b/tensorflow/lite/delegates/gpu/common/tasks/winograd.cc
@@ -21,7 +21,6 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
-#include "absl/strings/str_format.h"
 #include "tensorflow/lite/delegates/gpu/common/data_type.h"
 #include "tensorflow/lite/delegates/gpu/common/shape.h"
 #include "tensorflow/lite/delegates/gpu/common/status.h"
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/winograd.h b/tensorflow/lite/delegates/gpu/common/tasks/winograd.h
index c5da853dab0..c67820a0a28 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/winograd.h
+++ b/tensorflow/lite/delegates/gpu/common/tasks/winograd.h
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include <string>
 
+#include "absl/status/status.h"
 #include "tensorflow/lite/delegates/gpu/common/data_type.h"
 #include "tensorflow/lite/delegates/gpu/common/operations.h"
 #include "tensorflow/lite/delegates/gpu/common/task/gpu_operation.h"
diff --git a/tensorflow/lite/delegates/gpu/common/tasks/winograd_test_util.cc b/tensorflow/lite/delegates/gpu/common/tasks/winograd_test_util.cc
index 6b5fb4a830f..e31f8402a04 100644
--- a/tensorflow/lite/delegates/gpu/common/tasks/winograd_test_util.cc
+++ b/tensorflow/lite/delegates/gpu/common/tasks/winograd_test_util.cc
@@ -15,7 +15,9 @@ limitations under the License.
 
 #include "tensorflow/lite/delegates/gpu/common/tasks/winograd_test_util.h"
 
+#include <cmath>
 #include <memory>
+#include <utility>
 #include <vector>
 
 #include "tensorflow/lite/delegates/gpu/common/operations.h"

From 36305107292f78e8c1915ace4813580129224ea9 Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Thu, 8 Oct 2026 15:09:20 +0530
Subject: [PATCH 0261/1000] widen size vectors to int64_t in
 FractionalMaxPoolOp to avoid 32-bit overflow

---
 tensorflow/core/kernels/fractional_max_pool_op.cc | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensorflow/core/kernels/fractional_max_pool_op.cc b/tensorflow/core/kernels/fractional_max_pool_op.cc
index b1a84631c17..05aaf44b58d 100644
--- a/tensorflow/core/kernels/fractional_max_pool_op.cc
+++ b/tensorflow/core/kernels/fractional_max_pool_op.cc
@@ -92,8 +92,8 @@ class FractionalMaxPoolOp : public OpKernel {
     OP_REQUIRES(context, tensor_in.dims() == tensor_in_and_out_dims,
                 absl::InvalidArgumentError("tensor_in must be 4-dimensional"));
 
-    std::vector<int> input_size(tensor_in_and_out_dims);
-    std::vector<int> output_size(tensor_in_and_out_dims);
+    std::vector<int64_t> input_size(tensor_in_and_out_dims);
+    std::vector<int64_t> output_size(tensor_in_and_out_dims);
     for (int i = 0; i < tensor_in_and_out_dims; ++i) {
       input_size[i] = tensor_in.dim_size(i);
 
@@ -109,7 +109,7 @@ class FractionalMaxPoolOp : public OpKernel {
       // This must match the same logic in the shape function in
       // core/ops/nn_ops.cc.
       output_size[i] =
-          static_cast<int>(std::floor(input_size[i] / pooling_ratio_[i]));
+          static_cast<int64_t>(std::floor(input_size[i] / pooling_ratio_[i]));
       DCHECK_GT(output_size[i], 0);
     }
 
@@ -156,11 +156,11 @@ class FractionalMaxPoolOp : public OpKernel {
     auto output_width_seq_flat = output_width_seq_tensor->flat<int64_t>();
 
     // Set output tensors.
-    for (int i = 0; i < height_cum_seq.size(); ++i) {
+    for (size_t i = 0; i < height_cum_seq.size(); ++i) {
       output_height_seq_flat(i) = height_cum_seq[i];
     }
 
-    for (int i = 0; i < width_cum_seq.size(); ++i) {
+    for (size_t i = 0; i < width_cum_seq.size(); ++i) {
       output_width_seq_flat(i) = width_cum_seq[i];
     }
 

From 955b84ebb716e6284fc1168c540c165dca33551e Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 8 Oct 2026 01:11:32 -0700
Subject: [PATCH 0262/1000] Automated Code Change

PiperOrigin-RevId: 995653595
---
 tensorflow/compiler/mlir/BUILD                           | 1 +
 tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc | 1 +
 2 files changed, 2 insertions(+)

diff --git a/tensorflow/compiler/mlir/BUILD b/tensorflow/compiler/mlir/BUILD
index bbc93b94381..57218725865 100644
--- a/tensorflow/compiler/mlir/BUILD
+++ b/tensorflow/compiler/mlir/BUILD
@@ -149,6 +149,7 @@ cc_library(
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/log:vlog_is_on",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings:string_view",
         "@llvm-project//llvm:Support",
diff --git a/tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc b/tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc
index 9c8e27a51b4..678806c2e5f 100644
--- a/tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc
+++ b/tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc
@@ -22,6 +22,7 @@ limitations under the License.
 
 #include "absl/container/flat_hash_set.h"
 #include "absl/log/log.h"
+#include "absl/log/vlog_is_on.h"
 #include "absl/status/status.h"
 #include "absl/strings/string_view.h"
 #include "llvm/ADT/StringRef.h"

From 5b5bd6cc24597ccced5323641fc2494bff731e64 Mon Sep 17 00:00:00 2001
From: bodapatisaikrishna <saikrishnabodapati@gmail.com>
Date: Thu, 8 Oct 2026 15:50:38 +0530
Subject: [PATCH 0263/1000] Revert redundant kernel checks and update conv2d
 rank tests

Revert OP_REQUIRES checks from Conv2DOp::Compute since ComputeConv2DDimension
already enforces required dimensions, and update tests to cover graph mode
shape inference and raw_ops filter rank validation.

Signed-off-by: bodapatisaikrishna <saikrishnabodapati@gmail.com>
---
 tensorflow/core/kernels/conv_ops_impl.h       | 13 -----------
 .../kernel_tests/nn_ops/conv_ops_test.py      | 22 ++++++++++++++++---
 2 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/tensorflow/core/kernels/conv_ops_impl.h b/tensorflow/core/kernels/conv_ops_impl.h
index ae96d206ef0..36fad88300f 100644
--- a/tensorflow/core/kernels/conv_ops_impl.h
+++ b/tensorflow/core/kernels/conv_ops_impl.h
@@ -663,19 +663,6 @@ class Conv2DOp : public BinaryOp<T> {
     // [ filter_rows, filter_cols, in_depth, out_depth]
     const Tensor& filter = context->input(1);
 
-    const int required_dims =
-        params_.data_format == TensorFormat::FORMAT_NCHW_VECT_C ? 5 : 4;
-    OP_REQUIRES(
-        context, input.dims() == required_dims,
-        absl::InvalidArgumentError(
-            absl::StrCat("convolution input must be ", required_dims,
-                         "-dimensional: ", input.shape().DebugString())));
-    OP_REQUIRES(
-        context, filter.dims() == required_dims,
-        absl::InvalidArgumentError(absl::StrCat(
-            "convolution filter must be ", required_dims,
-            "-dimensional: ", filter.shape().DebugString())));
-
     Conv2DDimensions dimensions;
     OP_REQUIRES_OK(context,
                    ComputeConv2DDimension(params_, input, filter, &dimensions));
diff --git a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
index 9beeafb6fbc..f142a42e313 100644
--- a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
@@ -3137,7 +3137,8 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
     # Invalid input or filter rank (< 4) in eager/graph execution.
     with self.assertRaisesRegex(
         (ValueError, errors_impl.InvalidArgumentError),
-        r"(must be (at least )?4-dimensional|input\.shape\.rank)",
+        r"(must be (at least )?4-dimensional|"
+        r"Shape must be rank 4 but is rank 3)",
     ):
       self.evaluate(
           nn_ops.conv2d(
@@ -3149,7 +3150,8 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
       )
     with self.assertRaisesRegex(
         (ValueError, errors_impl.InvalidArgumentError),
-        r"(must be (at least )?4-dimensional|filter\.shape\.rank)",
+        r"(must be (at least )?4-dimensional|"
+        r"Shape must be rank 4 but is rank 3)",
     ):
       self.evaluate(
           nn_ops.conv2d(
@@ -3163,7 +3165,8 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
     # Ensure raw ops are protected by the C++ kernel boundary.
     with self.assertRaisesRegex(
         (ValueError, errors_impl.InvalidArgumentError),
-        r"(must be (at least )?4-dimensional|input\.shape\.rank)",
+        r"(must be (at least )?4-dimensional|"
+        r"Shape must be rank 4 but is rank 3)",
     ):
       self.evaluate(
           raw_ops.Conv2D(
@@ -3173,6 +3176,19 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
               padding="VALID",
           )
       )
+    with self.assertRaisesRegex(
+        (ValueError, errors_impl.InvalidArgumentError),
+        r"(must be (at least )?4-dimensional|"
+        r"Shape must be rank 4 but is rank 3)",
+    ):
+      self.evaluate(
+          raw_ops.Conv2D(
+              input=np.ones([1, 1, 3, 3], dtype=np.float32),
+              filter=np.ones([1, 2, 2], dtype=np.float32),
+              strides=[1, 1, 1, 1],
+              padding="VALID",
+          )
+      )
 
     # Filter dimensions must be greater than 0.
     with self.assertRaisesRegex(

From d7b2d1d177950546241b30787ededaf195699cd8 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Thu, 8 Oct 2026 01:57:51 -0700
Subject: [PATCH 0264/1000] [XLA:GPU] Document --xla_gpu_topology_filename in
 hlo-opt and target config docs.

PiperOrigin-RevId: 995674872
---
 third_party/xla/docs/tools.md                 | 38 ++++++++++++++++---
 .../xla/backends/gpu/target_config/README.md  |  5 +++
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/third_party/xla/docs/tools.md b/third_party/xla/docs/tools.md
index 6a9efebd0e8..6331af9a2bb 100644
--- a/third_party/xla/docs/tools.md
+++ b/third_party/xla/docs/tools.md
@@ -102,14 +102,15 @@ which would print the dump to stdout (or to a given file if `-o` was specified).
 ### Deviceless Compilation for GPU
 
 Deviceless compilation do not need access to a GPU. The Deviceless Compilation
-provides a way to specify GPU spec on the command line
-(`--xla_gpu_target_config_filename`) for stages where access to GPU is required,
+provides a way to specify a GPU spec (`--xla_gpu_target_config_filename`) or
+GPU topology (`--xla_gpu_topology_filename`) on the command line for stages
+where access to a GPU or multi-GPU topology is required,
 eliminating a need for GPU device.
 
 Example: PTX output without access to a gpu device:
 
 ```
-hlo-opt  --platform=CUDA --stage=llvm  --xla_gpu_target_config_filename=/xla/tools/hlo_opt/gpu_specs/a100_pcie_80.txtpb input.hlo
+hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=xla/backends/gpu/target_config/specs/a100_pcie_80.txtpb input.hlo
 ```
 
 Specs for popular GPUs are shipped with the compiler, and the provided file is
@@ -138,7 +139,34 @@ gpu_device_info {
 }
 platform_name: "CUDA"
 ```
-More GPU specs are located at `/xla/tools/hlo_opt/gpu_specs`
+More GPU specs are located at `xla/backends/gpu/target_config/specs`.
+
+#### Multi-GPU Topology and Collectives
+
+Passes that lower or optimize collective operations require a GPU topology
+(number of partitions, hosts per partition, and devices per host). Use
+`--xla_gpu_topology_filename` to provide either:
+
+*   A serialized `GpuTopologyProto` file from
+    `xla/backends/gpu/target_config/specs` (such as `h100_1x8.txtpb`,
+    `h100_2x8.txtpb`, `gb200_2x4.txtpb`, or `gb300_4x4.txtpb`), which also
+    automatically populates the target device configuration from
+    `platform_version`:
+
+    ```
+    hlo-opt --platform=gpu --stage=hlo --xla_gpu_topology_filename=xla/backends/gpu/target_config/specs/gb200_2x4.txtpb input.hlo
+    ```
+
+*   An inline topology string of the form
+    `[<platform_version>:]<num_partitions>x<num_hosts_per_partition>x<num_devices_per_host>`
+    (or `[<platform_version>:]<num_hosts_per_partition>x<num_devices_per_host>`).
+    When `<platform_version>` is omitted, pair it with
+    `--xla_gpu_target_config_filename`:
+
+    ```
+    hlo-opt --platform=gpu --stage=hlo --xla_gpu_topology_filename=oberon_b200:1x2x4 input.hlo
+    hlo-opt --platform=gpu --stage=hlo --xla_gpu_topology_filename=1x2x4 --xla_gpu_target_config_filename=xla/backends/gpu/target_config/specs/gb200.txtpb input.hlo
+    ```
 
 #### Autotuning
 
@@ -151,7 +179,7 @@ or\
 `--xla_gpu_dump_autotune_results_to=<filename>`).
 
 ```
-hlo-opt  --platform=CUDA --stage=llvm  --xla_gpu_target_config_filename=gpu_specs/a100_pcie_80.txtpb --xla_gpu_load_autotune_results_from=results.textpb input.hlo
+hlo-opt --platform=CUDA --stage=llvm --xla_gpu_target_config_filename=xla/backends/gpu/target_config/specs/a100_pcie_80.txtpb --xla_gpu_load_autotune_results_from=results.textpb input.hlo
 ```
 
 The autotune file is text serialization of `autotune_results.proto`, with
diff --git a/third_party/xla/xla/backends/gpu/target_config/README.md b/third_party/xla/xla/backends/gpu/target_config/README.md
index 8d15dbb7e52..7ba97a2bd91 100644
--- a/third_party/xla/xla/backends/gpu/target_config/README.md
+++ b/third_party/xla/xla/backends/gpu/target_config/README.md
@@ -7,3 +7,8 @@ The specs are useful when compiling with the flag
 `--xla_gpu_target_config_filename`. Since a hardware generation may have several
 SKUs, a spec may not be identical to what we would get on a particular machine,
 but it will be "close enough".
+
+Multi-GPU topology specs (e.g. `h100_1x8.txtpb`, `gb200_2x4.txtpb`) are text
+serializations of `GpuTopologyProto` and can be passed to
+`--xla_gpu_topology_filename` (which also accepts an inline shorthand
+`[platform:]num_partitionsxnum_hosts_per_partitionxnum_devices_per_host`).

From 5560b75cc7506819155332a62161964706e74f8f Mon Sep 17 00:00:00 2001
From: Joshua Lang <joshualang@google.com>
Date: Thu, 8 Oct 2026 02:01:25 -0700
Subject: [PATCH 0265/1000] Extract large QR op tests into a separate
 `qr_op_large_test` target with 24G memory limit

PiperOrigin-RevId: 995676553
---
 tensorflow/compiler/tests/BUILD               | 27 ++++++++++
 tensorflow/compiler/tests/qr_op_large_test.py | 49 +++++++++++++++++++
 tensorflow/compiler/tests/qr_op_test.py       | 26 ++++------
 3 files changed, 86 insertions(+), 16 deletions(-)
 create mode 100644 tensorflow/compiler/tests/qr_op_large_test.py

diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD
index dd59758a066..deb8ce4f21f 100644
--- a/tensorflow/compiler/tests/BUILD
+++ b/tensorflow/compiler/tests/BUILD
@@ -1390,6 +1390,33 @@ tf_xla_py_strict_test(
     ],
 )
 
+tf_xla_py_strict_test(
+    name = "qr_op_large_test",
+    size = "large",
+    srcs = ["qr_op_large_test.py"],
+    disabled_backends = [
+        # Test is very slow on CPU.
+        "cpu",
+        "cpu_ondemand",
+    ],
+    # Keep shard_count = 2 so testLarge2000x2000 and testLarge17500x128 (if enabled) run in separate shards.
+    shard_count = 2,
+    tags = [
+        "no_pip",  # TODO(b/149738646): fix pip install so these tests run on kokoro pip
+        "nozapfhahn",
+        "optonly",
+        "requires-mem:24g",
+    ],
+    deps = [
+        ":qr_op_test_lib",
+        ":xla_test",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+        "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
 tf_xla_py_strict_test(
     name = "unstack_test",
     size = "medium",
diff --git a/tensorflow/compiler/tests/qr_op_large_test.py b/tensorflow/compiler/tests/qr_op_large_test.py
new file mode 100644
index 00000000000..9c8220c10b8
--- /dev/null
+++ b/tensorflow/compiler/tests/qr_op_large_test.py
@@ -0,0 +1,49 @@
+# 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.
+# ==============================================================================
+"""Large tests for tensorflow.ops.linalg_ops.qr."""
+
+import unittest
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.compiler.tests import qr_op_test
+from tensorflow.compiler.tests import xla_test
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+
+
+@test_util.run_all_without_tensor_float_32(
+    "XLA QR op calls matmul. Also, matmul used for verification. Also with "
+    'TensorFloat-32, mysterious "Unable to launch cuBLAS gemm" error '
+    "occasionally occurs"
+)
+# TODO(b/165435566): Fix "Unable to launch cuBLAS gemm" error
+class QrOpLargeTest(
+    qr_op_test.QrOpTestBase, xla_test.XLATestCase, parameterized.TestCase
+):
+
+  def testLarge2000x2000(self):
+    x_np = self._random_matrix(np.float32, (2000, 2000))
+    self._test(x_np, full_matrices=True)
+
+  @unittest.skip("Test times out on CI")
+  def testLarge17500x128(self):
+    x_np = self._random_matrix(np.float32, (17500, 128))
+    self._test(x_np, full_matrices=True)
+
+
+if __name__ == "__main__":
+  test.main()
diff --git a/tensorflow/compiler/tests/qr_op_test.py b/tensorflow/compiler/tests/qr_op_test.py
index b77bc9c46aa..e04ca7b909b 100644
--- a/tensorflow/compiler/tests/qr_op_test.py
+++ b/tensorflow/compiler/tests/qr_op_test.py
@@ -15,7 +15,6 @@
 """Tests for tensorflow.ops.math_ops.matrix_inverse."""
 
 import itertools
-import unittest
 
 from absl.testing import parameterized
 import numpy as np
@@ -29,12 +28,7 @@ from tensorflow.python.ops import math_ops
 from tensorflow.python.platform import test
 
 
-@test_util.run_all_without_tensor_float_32(
-    "XLA QR op calls matmul. Also, matmul used for verification. Also with "
-    'TensorFloat-32, mysterious "Unable to launch cuBLAS gemm" error '
-    "occasionally occurs")
-# TODO(b/165435566): Fix "Unable to launch cuBLAS gemm" error
-class QrOpTest(xla_test.XLATestCase, parameterized.TestCase):
+class QrOpTestBase:
 
   def AdjustedNorm(self, x):
     """Computes the norm of matrices in 'x', adjusted for dimension and type."""
@@ -111,6 +105,15 @@ class QrOpTest(xla_test.XLATestCase, parameterized.TestCase):
       self.CheckApproximation(x_np, q_tf_val, r_tf_val)
       self.CheckUnitary(q_tf_val)
 
+
+@test_util.run_all_without_tensor_float_32(
+    "XLA QR op calls matmul. Also, matmul used for verification. Also with "
+    'TensorFloat-32, mysterious "Unable to launch cuBLAS gemm" error '
+    "occasionally occurs"
+)
+# TODO(b/165435566): Fix "Unable to launch cuBLAS gemm" error
+class QrOpTest(QrOpTestBase, xla_test.XLATestCase, parameterized.TestCase):
+
   SIZES = [1, 2, 5, 10, 32, 100, 300, 603]
   DTYPES = [np.float32, np.complex64]
   PARAMS = itertools.product(SIZES, SIZES, DTYPES)
@@ -123,15 +126,6 @@ class QrOpTest(xla_test.XLATestCase, parameterized.TestCase):
         x_np = self._random_matrix(dtype, batch_dims + (rows, cols))
         self._test(x_np, full_matrices)
 
-  def testLarge2000x2000(self):
-    x_np = self._random_matrix(np.float32, (2000, 2000))
-    self._test(x_np, full_matrices=True)
-
-  @unittest.skip("Test times out on CI")
-  def testLarge17500x128(self):
-    x_np = self._random_matrix(np.float32, (17500, 128))
-    self._test(x_np, full_matrices=True)
-
   @parameterized.parameters((23, 25), (513, 23))
   def testZeroColumn(self, rows, cols):
     x_np = self._random_matrix(np.complex64, (rows, cols))

From e445a0dae18b60d3c6f52132a4dc4ebcb3ea4012 Mon Sep 17 00:00:00 2001
From: nishad shabbir <nishadshabbir28@gmail.com>
Date: Thu, 8 Oct 2026 16:25:20 +0530
Subject: [PATCH 0266/1000] Address review: check shapes first and cover the
 integer path

- _verify_compatible_image_shapes() ran after the reduction, so an input
  of rank below 3 surfaced from reduce_mean rather than from the shape
  check. Rank 2 reported an InvalidArgumentError out of Mean; it now
  reports "Shape (2, 2) must have rank at least 3", which is what the
  verifier was there to say. The runtime asserts now gate `a` instead of
  the result, so they still run before the computation uses it.
- The test only covered float32. It now also covers uint8 arrays and the
  integer lists the PR description discusses. An integer list converts to
  int32 rather than uint8, so convert_image_dtype() scales it by a
  different range; PSNR is invariant under that because max_val is scaled
  the same way, and the test pins both forms against the value computed
  from the integers directly.
---
 tensorflow/python/ops/image_ops_impl.py | 18 +++++++++++-------
 tensorflow/python/ops/image_ops_test.py | 13 +++++++++++++
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/tensorflow/python/ops/image_ops_impl.py b/tensorflow/python/ops/image_ops_impl.py
index 0f0ac1789ac..aba1fd4888a 100644
--- a/tensorflow/python/ops/image_ops_impl.py
+++ b/tensorflow/python/ops/image_ops_impl.py
@@ -4312,10 +4312,18 @@ def psnr(a, b, max_val, name=None):
     and shape [batch_size, 1].
   """
   with ops.name_scope(name, 'PSNR', [a, b]):
-    # Convert first: max_val is cast to the dtype of `a` below, which an
-    # unconverted value does not have.
+    # Convert first: max_val is cast to the dtype of `a` below, and the shape
+    # check reads their static shapes, neither of which an unconverted value
+    # has.
     a = ops.convert_to_tensor(a, name='a')
     b = ops.convert_to_tensor(b, name='b')
+
+    # Check the shapes before computing anything, so a rank below 3 is
+    # reported as such instead of surfacing from the reduction below.
+    _, _, checks = _verify_compatible_image_shapes(a, b)
+    with ops.control_dependencies(checks):
+      a = array_ops.identity(a)
+
     # Need to convert the images to float32.  Scale max_val accordingly so that
     # PSNR is computed correctly.
     max_val = math_ops.cast(max_val, a.dtype)
@@ -4323,15 +4331,11 @@ def psnr(a, b, max_val, name=None):
     a = convert_image_dtype(a, dtypes.float32)
     b = convert_image_dtype(b, dtypes.float32)
     mse = math_ops.reduce_mean(math_ops.squared_difference(a, b), [-3, -2, -1])
-    psnr_val = math_ops.subtract(
+    return math_ops.subtract(
         20 * math_ops.log(max_val) / math_ops.log(10.0),
         np.float32(10 / np.log(10)) * math_ops.log(mse),
         name='psnr')
 
-    _, _, checks = _verify_compatible_image_shapes(a, b)
-    with ops.control_dependencies(checks):
-      return array_ops.identity(psnr_val)
-
 
 def _ssim_helper(x, y, reducer, max_val, compensation=1.0, k1=0.01, k2=0.03):
   r"""Helper function for computing SSIM.
diff --git a/tensorflow/python/ops/image_ops_test.py b/tensorflow/python/ops/image_ops_test.py
index 569544347d8..96673b5dc1e 100644
--- a/tensorflow/python/ops/image_ops_test.py
+++ b/tensorflow/python/ops/image_ops_test.py
@@ -6150,6 +6150,19 @@ class PSNRTest(test_util.TensorFlowTestCase):
         tf_psnr = self.evaluate(image_ops.psnr(a, b, 1.0, "psnr"))
         self.assertAllClose(expected, tf_psnr, atol=0.001)
 
+      # An integer list converts to int32 rather than uint8, so it is scaled
+      # by a different range than the uint8 array. PSNR is invariant under
+      # that, since max_val is scaled the same way, and both must still match
+      # the value computed from the integers directly.
+      image1_int = (image1 * 255).astype(np.uint8)
+      image2_int = (image2 * 255).astype(np.uint8)
+      expected_int = self._PSNR_NumPy(
+          image1_int.astype(np.float32), image2_int.astype(np.float32), 255)
+      for a, b in ((image1_int.tolist(), image2_int.tolist()),
+                   (image1_int, image2_int)):
+        tf_psnr_int = self.evaluate(image_ops.psnr(a, b, 255, "psnr_int"))
+        self.assertAllClose(expected_int, tf_psnr_int, atol=0.001)
+
   def testPSNRMultiImage(self):
     image1 = self._RandomImage((10, 8, 8, 1), 1)
     image2 = self._RandomImage((10, 8, 8, 1), 1)

From 305c44c1917e1b5d54a99588f8463b44890d3aff Mon Sep 17 00:00:00 2001
From: Pablo Zimmermann <karupayun@google.com>
Date: Thu, 8 Oct 2026 02:14:11 -0700
Subject: [PATCH 0267/1000] [XLA:GPU] Extract RecordOperandRead and
 ComputeReadTime in GpuIndexingPerformanceModel

I'm in the process of merging Dot and Indexing Cost Models. `EstimateRunTimeForTiledHloComputationImpl` currently tracks operand reads and
computes HBM read time inline using a single map. To accurately model tiled dot
fusions in follow-up changes, we need to track and evaluate operand reads inside
the dot reduction loop (K-loop) separately from epilogue operand reads.

So, here we are basically extracting `OperandReadMap`, `RecordOperandRead`, and `ComputeReadTime` so subsequent CLs can reuse the same read-tracking and DRAM-heuristic logic across multiple operand maps. There is no functional change and tests will come in following CLs.

PiperOrigin-RevId: 995683320
---
 .../model/gpu_indexing_performance_model.cc   | 76 ++++++++++++-------
 1 file changed, 49 insertions(+), 27 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc
index e64e4a05216..f065f03dcce 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc
+++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc
@@ -100,6 +100,49 @@ struct OperandReadInfo {
   double read_bandwidth_utilization_rate = 1.0;
 };
 
+using OperandReadMap =
+    absl::flat_hash_map<const HloInstruction*, OperandReadInfo>;
+
+// Accumulates `tile_bytes_read` and tracks the minimum
+// `bandwidth_utilization_rate` for `hlo` in `read_map`.
+void RecordOperandRead(const HloInstruction* hlo, int64_t tile_bytes_read,
+                       double bandwidth_utilization_rate,
+                       OperandReadMap& read_map) {
+  OperandReadInfo& operand_read_info = read_map[hlo];
+  operand_read_info.total_bytes_read += tile_bytes_read;
+  // TODO(b/332714755): using std::min is more pessimistic than it needs
+  // to be since it'll end up assuming that if one read is done with lower
+  // bandwidth, all other reads of the same operand will also be done with
+  // lower bandwidth. But it's a start. We should refactor this function
+  // to properly track each read independently later.
+  operand_read_info.read_bandwidth_utilization_rate =
+      std::min(operand_read_info.read_bandwidth_utilization_rate,
+               bandwidth_utilization_rate);
+}
+
+// Computes the memory read time across all operands in `read_map`.
+absl::Duration ComputeReadTime(
+    const OperandReadMap& read_map, const se::DeviceDescription& device_info,
+    int64_t num_blocks, const HloCostAnalysis::ShapeSizeFunction& shape_size) {
+  absl::Duration read_time = absl::ZeroDuration();
+  // `absl::Duration` addition uses exact fixed-point integer arithmetic, so
+  // map iteration order does not affect the accumulated result.
+  // NOLINTNEXTLINE(*-custom-deterministic-iteration-order)
+  for (const auto& [hlo, operand_read_info] : read_map) {
+    int64_t operand_size = shape_size(hlo->shape());
+    int64_t n_bytes_net =
+        std::min(operand_size, operand_read_info.total_bytes_read);
+
+    read_time += GpuPerformanceModelBase::ReadTimeWithDRAMHeuristic(
+        device_info, num_blocks, n_bytes_net,
+        operand_read_info.total_bytes_read,
+        /*element_type=*/hlo->shape().element_type(),
+        /*hbm_bandwidth_utilization_rate=*/
+        operand_read_info.read_bandwidth_utilization_rate);
+  }
+  return read_time;
+}
+
 // Returns the number of elements in the tile after each dimension is padded to
 // the next power of 2.
 // TODO(b/353484968): Delete this function once we have constraints to only
@@ -387,7 +430,7 @@ absl::StatusOr<EstimateRunTimeData> EstimateRunTimeForTiledHloComputationImpl(
     const se::DeviceDescription& device_info,
     HloCostAnalysis::ShapeSizeFunction shape_size,
     absl::FunctionRef<int64_t(const HloInstruction*)> flops_per_element_fn) {
-  absl::flat_hash_map<const HloInstruction*, OperandReadInfo> n_bytes_total_map;
+  OperandReadMap n_bytes_total_map;
 
   // Compute time for dot flops is counted separately.
   int64_t dot_flops = 0;
@@ -481,34 +524,13 @@ absl::StatusOr<EstimateRunTimeData> EstimateRunTimeForTiledHloComputationImpl(
             BandwidthUtilizationRateHeuristicForTiledMemoryAccess(*tiled_hlo,
                                                                   device_info);
 
-        OperandReadInfo& operand_read_info = n_bytes_total_map[hlo];
-        operand_read_info.total_bytes_read += tile_bytes_read;
-        // TODO(b/332714755): using std::min is more pessimistic than it needs
-        // to be since it'll end up assuming that if one read is done with lower
-        // bandwidth, all other reads of the same operand will also be done with
-        // lower bandwidth. But it's a start. We should refactor this function
-        // to properly track each read independently later.
-        operand_read_info.read_bandwidth_utilization_rate =
-            std::min(operand_read_info.read_bandwidth_utilization_rate,
-                     effective_bandwidth_utilization_rate);
+        RecordOperandRead(hlo, tile_bytes_read,
+                          effective_bandwidth_utilization_rate,
+                          n_bytes_total_map);
       });
 
-  absl::Duration read_time = absl::ZeroDuration();
-  for (const auto& [hlo, operand_read_info] : n_bytes_total_map) {
-    int64_t operand_size = shape_size(hlo->shape());
-    int64_t n_bytes_net =
-        std::min(operand_size, operand_read_info.total_bytes_read);
-
-    // TODO(b/332714755): use
-    // `BandwidthUtilizationRateHeuristicForTiledMemoryAccess` to compute read
-    // time as well.
-    read_time += GpuPerformanceModelBase::ReadTimeWithDRAMHeuristic(
-        device_info, num_blocks, n_bytes_net,
-        operand_read_info.total_bytes_read,
-        /*element_type=*/hlo->shape().element_type(),
-        /*hbm_bandwidth_utilization_rate=*/
-        operand_read_info.read_bandwidth_utilization_rate);
-  }
+  absl::Duration read_time =
+      ComputeReadTime(n_bytes_total_map, device_info, num_blocks, shape_size);
 
   auto roots = tiled_hlo_computation.roots();
   int64_t bytes_written = 0;

From fa6412de611e486e1fefdd47e4e2e234cb3354c3 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Thu, 8 Oct 2026 02:14:20 -0700
Subject: [PATCH 0268/1000] [XLA:GPU]: Migrate AsyncExecution to use
 PerDeviceState

PiperOrigin-RevId: 995683386
---
 .../xla/xla/backends/gpu/runtime/BUILD        |  4 +---
 .../backends/gpu/runtime/async_execution.cc   | 20 +++++++-----------
 .../backends/gpu/runtime/async_execution.h    | 18 +++++++---------
 .../gpu/runtime/async_execution_test.cc       |  6 +++---
 .../xla/backends/gpu/runtime/async_thunk.cc   | 11 ++++++----
 .../xla/backends/gpu/runtime/async_thunk.h    |  5 +++--
 .../backends/gpu/runtime/async_thunk_test.cc  | 21 ++++++++++---------
 .../collective_broadcast_thunk_test.cc        |  3 ++-
 .../runtime/collective_permute_thunk_test.cc  |  3 ++-
 .../command_buffer_async_conversion_test.cc   |  4 ++--
 .../command_buffer_cmd_emitter_test.cc        |  3 ++-
 .../command_buffer_conversion_pass_test.cc    | 18 ++++++++++------
 .../runtime/ragged_all_to_all_thunk_test.cc   |  3 ++-
 .../gpu/runtime/send_recv_thunk_test.cc       |  6 ++++--
 .../runtime/thunk_proto_deserialization.cc    |  7 ++++---
 .../thunk_proto_deserialization_test.cc       |  5 +++--
 .../xla/xla/service/gpu/thunk_emitter.cc      | 11 +++++++---
 17 files changed, 80 insertions(+), 68 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index 2f2dd157392..f38fc1e0485 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -399,6 +399,7 @@ cc_library(
     srcs = ["async_execution.cc"],
     hdrs = ["async_execution.h"],
     deps = [
+        ":per_device_state",
         ":thunk",
         ":thunk_id",
         "//xla:util",
@@ -408,15 +409,12 @@ cc_library(
         "//xla/stream_executor:stream_executor_h",
         "//xla/tsl/lib/gtl:int_type",
         "//xla/tsl/util:unique_any",
-        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
-        "@com_google_absl//absl/container:node_hash_map",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:str_format",
-        "@com_google_absl//absl/synchronization",
     ],
 )
 
diff --git a/third_party/xla/xla/backends/gpu/runtime/async_execution.cc b/third_party/xla/xla/backends/gpu/runtime/async_execution.cc
index 09e74083c6c..0781e46c850 100644
--- a/third_party/xla/xla/backends/gpu/runtime/async_execution.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/async_execution.cc
@@ -24,7 +24,6 @@ limitations under the License.
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_format.h"
-#include "absl/synchronization/mutex.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk_id.h"
 #include "xla/stream_executor/event.h"
@@ -63,8 +62,10 @@ struct ExecutionState {
 
 }  // namespace
 
-AsyncExecution::AsyncExecution(Thunk::ThunkInfo start_thunk_info)
-    : start_thunk_info_(std::move(start_thunk_info)) {}
+AsyncExecution::AsyncExecution(Thunk::ThunkInfo start_thunk_info,
+                               int devices_per_host)
+    : start_thunk_info_(std::move(start_thunk_info)),
+      event_pools_(devices_per_host) {}
 
 AsyncExecution::ExecutionGuard::ExecutionGuard(se::Event* event,
                                                se::Stream* async_stream)
@@ -88,21 +89,14 @@ AsyncExecution::ExecutionGuard::~ExecutionGuard() {
       << " on a stream " << async_stream_;
 }
 
-AsyncExecution::EventPool& AsyncExecution::GetOrCreatePool(
-    se::StreamExecutor* executor) {
-  absl::MutexLock lock(mu_);
-  auto [it, _] = event_pools_.try_emplace(
-      executor, [executor] { return executor->CreateEvent(); });
-  return it->second;
-}
-
 absl::Status AsyncExecution::Initialize(Thunk::ExecutionScopedState* state,
                                         se::StreamExecutor* executor) {
   XLA_VLOG_DEVICE(1, executor->device_ordinal())
       << absl::StreamFormat("Initialize async execution for `%s`",
                             start_thunk_info_.profile_annotation);
-  EventPool& pool = GetOrCreatePool(executor);
-  ABSL_ASSIGN_OR_RETURN(auto borrowed, pool.GetOrCreate());
+  ABSL_ASSIGN_OR_RETURN(PoolState * pool_state,
+                   event_pools_.GetOrCreate(executor->device_ordinal()));
+  ABSL_ASSIGN_OR_RETURN(auto borrowed, pool_state->pool.GetOrCreate(executor));
   state->try_emplace(start_thunk_info_.thunk_id,
                      std::in_place_type<ExecutionState>, std::move(borrowed));
   // For shared async executions (e.g. pipelined send/recv), multiple
diff --git a/third_party/xla/xla/backends/gpu/runtime/async_execution.h b/third_party/xla/xla/backends/gpu/runtime/async_execution.h
index 302dd35e65e..bba3f65a755 100644
--- a/third_party/xla/xla/backends/gpu/runtime/async_execution.h
+++ b/third_party/xla/xla/backends/gpu/runtime/async_execution.h
@@ -19,12 +19,10 @@ limitations under the License.
 #include <cstdint>
 #include <memory>
 
-#include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
-#include "absl/container/node_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "absl/synchronization/mutex.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk_id.h"
 #include "xla/runtime/object_pool.h"
@@ -76,12 +74,12 @@ TSL_LIB_GTL_DEFINE_INT_TYPE(AsyncExecutionId, uint64_t);
 // id is used throughout the runtime as async execution id.
 class AsyncExecution {
  public:
-  using EventPool = ObjectPool<std::unique_ptr<se::Event>>;
+  using EventPool = ObjectPool<std::unique_ptr<se::Event>, se::StreamExecutor*>;
 
   // We need to know the thunk that starts an async execution, as we use its id
   // as a key in the execution scoped state and its profile annotation for
   // logging.
-  explicit AsyncExecution(Thunk::ThunkInfo start_thunk_info);
+  AsyncExecution(Thunk::ThunkInfo start_thunk_info, int devices_per_host);
 
   ThunkId start_thunk_id() const { return start_thunk_info_.thunk_id; }
 
@@ -123,14 +121,12 @@ class AsyncExecution {
   absl::Status Done(Thunk::ExecutionScopedState* state, se::Stream* stream);
 
  private:
-  // Returns or creates an event pool for the given executor.
-  EventPool& GetOrCreatePool(se::StreamExecutor* executor);
+  struct PoolState {
+    EventPool pool{[](se::StreamExecutor* e) { return e->CreateEvent(); }};
+  };
 
   const Thunk::ThunkInfo start_thunk_info_;
-
-  absl::Mutex mu_;
-  absl::node_hash_map<se::StreamExecutor*, EventPool> event_pools_
-      ABSL_GUARDED_BY(mu_);
+  PerDeviceState<PoolState> event_pools_;
 };
 
 // Map from AsyncExecutionId to shared AsyncExecution, used during
diff --git a/third_party/xla/xla/backends/gpu/runtime/async_execution_test.cc b/third_party/xla/xla/backends/gpu/runtime/async_execution_test.cc
index 4728793b67e..e5bd42e058e 100644
--- a/third_party/xla/xla/backends/gpu/runtime/async_execution_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/async_execution_test.cc
@@ -72,7 +72,7 @@ TEST(AsyncExecutionTest, InitializeStartDone) {
   thunk_info.profile_annotation = "test-thunk";
   TestThunk thunk(thunk_info);
 
-  AsyncExecution async_execution(thunk_info);
+  AsyncExecution async_execution(thunk_info, /*devices_per_host=*/1);
   Thunk::ExecutionScopedState state;
 
   // Initialize creates an event in the execution scoped state.
@@ -95,7 +95,7 @@ TEST(AsyncExecutionTest, CompletionEventRecordedOnce) {
   thunk_info.thunk_id = ThunkId(1);
   thunk_info.profile_annotation = "test-thunk";
 
-  AsyncExecution async_execution(thunk_info);
+  AsyncExecution async_execution(thunk_info, /*devices_per_host=*/1);
   Thunk::ExecutionScopedState state;
   ASSERT_OK(async_execution.Initialize(&state, executor));
 
@@ -133,7 +133,7 @@ TEST(AsyncExecutionTest, DoneWithoutStartFails) {
   thunk_info.profile_annotation = "test-thunk";
   TestThunk thunk(thunk_info);
 
-  AsyncExecution async_execution(thunk_info);
+  AsyncExecution async_execution(thunk_info, /*devices_per_host=*/1);
   Thunk::ExecutionScopedState state;
 
   // Done without Initialize should fail because event is not in state.
diff --git a/third_party/xla/xla/backends/gpu/runtime/async_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/async_thunk.cc
index 3c81a6a3d5e..199dcae1610 100644
--- a/third_party/xla/xla/backends/gpu/runtime/async_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/async_thunk.cc
@@ -46,11 +46,12 @@ namespace xla::gpu {
 
 AsyncStartThunk::AsyncStartThunk(ThunkInfo thunk_info,
                                  ExecutionStreamId execution_stream_id,
-                                 ThunkSequence thunks)
+                                 ThunkSequence thunks, int devices_per_host)
     : Thunk(Thunk::kAsyncStart, std::move(thunk_info)),
       execution_stream_id_(execution_stream_id),
       executor_(std::move(thunks)),
-      async_execution_(std::make_shared<AsyncExecution>(this->thunk_info())) {}
+      async_execution_(std::make_shared<AsyncExecution>(this->thunk_info(),
+                                                        devices_per_host)) {}
 
 AsyncStartThunk::AsyncStartThunk(
     ThunkInfo thunk_info, ExecutionStreamId execution_stream_id,
@@ -209,7 +210,8 @@ absl::StatusOr<ThunkProto> AsyncStartThunk::ToProto() const {
 
 absl::StatusOr<std::unique_ptr<AsyncStartThunk>> AsyncStartThunk::FromProto(
     ThunkInfo thunk_info, const AsyncStartThunkProto& proto,
-    const Deserializer& deserializer, AsyncExecutionMap& async_executions) {
+    const Deserializer& deserializer, AsyncExecutionMap& async_executions,
+    int devices_per_host) {
   auto make_stream_id = [&]() -> absl::StatusOr<ExecutionStreamId> {
     switch (proto.execution_stream_id_case()) {
       case AsyncStartThunkProto::kComputationStreamId:
@@ -237,7 +239,8 @@ absl::StatusOr<std::unique_ptr<AsyncStartThunk>> AsyncStartThunk::FromProto(
   }
 
   auto start_thunk = std::make_unique<AsyncStartThunk>(
-      std::move(thunk_info), execution_stream_id, std::move(nested));
+      std::move(thunk_info), execution_stream_id, std::move(nested),
+      devices_per_host);
 
   AsyncExecutionId id(proto.async_execution_id());
   async_executions[id] = start_thunk->async_execution();
diff --git a/third_party/xla/xla/backends/gpu/runtime/async_thunk.h b/third_party/xla/xla/backends/gpu/runtime/async_thunk.h
index 419167ac764..0e833358687 100644
--- a/third_party/xla/xla/backends/gpu/runtime/async_thunk.h
+++ b/third_party/xla/xla/backends/gpu/runtime/async_thunk.h
@@ -49,7 +49,7 @@ namespace xla::gpu {
 class AsyncStartThunk : public Thunk {
  public:
   AsyncStartThunk(ThunkInfo thunk_info, ExecutionStreamId execution_stream_id,
-                  ThunkSequence thunks);
+                  ThunkSequence thunks, int devices_per_host);
 
   // Constructor that shares an existing AsyncExecution with another
   // AsyncStartThunk. Used for pipelined send/recv where multiple operations
@@ -73,7 +73,8 @@ class AsyncStartThunk : public Thunk {
   absl::StatusOr<ThunkProto> ToProto() const override;
   static absl::StatusOr<std::unique_ptr<AsyncStartThunk>> FromProto(
       ThunkInfo thunk_info, const AsyncStartThunkProto& proto,
-      const Deserializer& deserializer, AsyncExecutionMap& async_executions);
+      const Deserializer& deserializer, AsyncExecutionMap& async_executions,
+      int devices_per_host);
 
   AsyncExecutionId async_execution_id() const;
   std::shared_ptr<AsyncExecution> async_execution() const;
diff --git a/third_party/xla/xla/backends/gpu/runtime/async_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/async_thunk_test.cc
index f66e7bb55a3..71c7840e232 100644
--- a/third_party/xla/xla/backends/gpu/runtime/async_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/async_thunk_test.cc
@@ -106,7 +106,8 @@ TEST(AsyncThunkTest, ConcurrentMemsets) {
     start_info.thunk_id = ThunkId(i + 1);
 
     thunks.Emplace<AsyncStartThunk>(std::move(start_info),
-                                    ComputationStreamId(i), std::move(nested));
+                                    ComputationStreamId(i), std::move(nested),
+                                    /*devices_per_host=*/1);
   }
 
   for (int i = 0; i < kNumChunks; ++i) {
@@ -170,14 +171,14 @@ TEST(AsyncThunkTest, MemcpyStreamDispatch) {
   h2d_info.profile_annotation = "h2d_start";
 
   ThunkSequence thunks;
-  thunks.Emplace<AsyncStartThunk>(d2h_info, kMemcpyD2HStreamId,
-                                  ThunkSequence{});
+  thunks.Emplace<AsyncStartThunk>(d2h_info, kMemcpyD2HStreamId, ThunkSequence{},
+                                  /*devices_per_host=*/1);
   auto* d2h_start = static_cast<AsyncStartThunk*>(thunks[0].get());
   thunks.Emplace<AsyncDoneThunk>(Thunk::ThunkInfo(),
                                  d2h_start->async_execution());
 
-  thunks.Emplace<AsyncStartThunk>(h2d_info, kMemcpyH2DStreamId,
-                                  ThunkSequence{});
+  thunks.Emplace<AsyncStartThunk>(h2d_info, kMemcpyH2DStreamId, ThunkSequence{},
+                                  /*devices_per_host=*/1);
   auto* h2d_start = static_cast<AsyncStartThunk*>(thunks[2].get());
   thunks.Emplace<AsyncDoneThunk>(Thunk::ThunkInfo(),
                                  h2d_start->async_execution());
@@ -210,9 +211,9 @@ TEST(AsyncThunkTest, MemcpyStreamNullReturnsError) {
   ASSERT_OK_AND_ASSIGN(auto main_stream, executor->CreateStream());
 
   AsyncStartThunk d2h_start(Thunk::ThunkInfo(), kMemcpyD2HStreamId,
-                            ThunkSequence{});
+                            ThunkSequence{}, /*devices_per_host=*/1);
   AsyncStartThunk h2d_start(Thunk::ThunkInfo(), kMemcpyH2DStreamId,
-                            ThunkSequence{});
+                            ThunkSequence{}, /*devices_per_host=*/1);
 
   se::StreamExecutorAddressAllocator allocator(executor);
   BufferAllocations allocations({}, 0, &allocator);
@@ -242,12 +243,12 @@ TEST(AsyncThunkTest, GpuExecutableFallsBackToMainStreamForMemcpyStreams) {
 
   ThunkSequence thunks;
   thunks.Emplace<AsyncStartThunk>(Thunk::ThunkInfo(), kMemcpyD2HStreamId,
-                                  ThunkSequence{});
+                                  ThunkSequence{}, /*devices_per_host=*/1);
   auto* d2h_start = static_cast<AsyncStartThunk*>(thunks[0].get());
   thunks.Emplace<AsyncDoneThunk>(Thunk::ThunkInfo(),
                                  d2h_start->async_execution());
   thunks.Emplace<AsyncStartThunk>(Thunk::ThunkInfo(), kMemcpyH2DStreamId,
-                                  ThunkSequence{});
+                                  ThunkSequence{}, /*devices_per_host=*/1);
   auto* h2d_start = static_cast<AsyncStartThunk*>(thunks[2].get());
   thunks.Emplace<AsyncDoneThunk>(Thunk::ThunkInfo(),
                                  h2d_start->async_execution());
@@ -288,7 +289,7 @@ TEST(AsyncThunkTest, AsyncDoneRecordCommandBuffer) {
 
   // Create a paired start/done to obtain a valid AsyncExecution.
   AsyncStartThunk start(Thunk::ThunkInfo(), ComputationStreamId(0),
-                        ThunkSequence());
+                        ThunkSequence(), /*devices_per_host=*/1);
   AsyncDoneThunk done(Thunk::ThunkInfo(), start.async_execution());
 
   // Set up minimal execute params (not used by AsyncDoneThunk::Record()).
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk_test.cc
index 03ec4977ca6..9144eedec66 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_broadcast_thunk_test.cc
@@ -214,7 +214,8 @@ ENTRY test_computation {
   ThunkSequence start_sequence;
   start_sequence.push_back(std::move(cb_start_thunk));
   auto async_start = std::make_unique<AsyncStartThunk>(
-      Thunk::ThunkInfo(), CommunicationStreamId(0), std::move(start_sequence));
+      Thunk::ThunkInfo(), CommunicationStreamId(0), std::move(start_sequence),
+      /*devices_per_host=*/1);
   auto async_done = std::make_unique<AsyncDoneThunk>(
       Thunk::ThunkInfo(), async_start->async_execution());
 
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_permute_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_permute_thunk_test.cc
index 43b1fcdb561..e9cc817ca27 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_permute_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_permute_thunk_test.cc
@@ -226,7 +226,8 @@ ENTRY test_computation {
   ThunkSequence start_sequence;
   start_sequence.push_back(std::move(cp_start_thunk));
   auto async_start = std::make_unique<AsyncStartThunk>(
-      Thunk::ThunkInfo(), CommunicationStreamId(0), std::move(start_sequence));
+      Thunk::ThunkInfo(), CommunicationStreamId(0), std::move(start_sequence),
+      /*devices_per_host=*/1);
   auto async_done = std::make_unique<AsyncDoneThunk>(
       Thunk::ThunkInfo(), async_start->async_execution());
 
diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_async_conversion_test.cc b/third_party/xla/xla/backends/gpu/runtime/command_buffer_async_conversion_test.cc
index 17cfbe8499f..950ec1f3573 100644
--- a/third_party/xla/xla/backends/gpu/runtime/command_buffer_async_conversion_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_async_conversion_test.cc
@@ -80,8 +80,8 @@ class CommandBufferAsyncConversionTest : public testing::Test {
       // Memset thunks are not eligible for conversion in this pass.
       body.Emplace<Memset32BitValueThunk>(Info(), 0, slice_);
     }
-    auto start =
-        std::make_unique<AsyncStartThunk>(Info(), stream, std::move(body));
+    auto start = std::make_unique<AsyncStartThunk>(
+        Info(), stream, std::move(body), /*devices_per_host=*/1);
     auto* result = start.get();
     thunks.push_back(std::move(start));
     return result;
diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_cmd_emitter_test.cc b/third_party/xla/xla/backends/gpu/runtime/command_buffer_cmd_emitter_test.cc
index bd93e9274b9..9b4d940da49 100644
--- a/third_party/xla/xla/backends/gpu/runtime/command_buffer_cmd_emitter_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_cmd_emitter_test.cc
@@ -875,7 +875,8 @@ class AsyncCommandBufferCmdEmitterTest : public CommandBufferCmdEmitterTest {
   AsyncStartThunk* Start(ThunkSequence& thunks, ExecutionStreamId stream,
                          ThunkSequence body) {
     auto start = std::make_unique<AsyncStartThunk>(NextThunkInfo("start"),
-                                                   stream, std::move(body));
+                                                   stream, std::move(body),
+                                                   /*devices_per_host=*/1);
     AsyncStartThunk* result = start.get();
     thunks.push_back(std::move(start));
     return result;
diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass_test.cc b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass_test.cc
index 08628d43848..521031886f9 100644
--- a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass_test.cc
@@ -305,7 +305,8 @@ std::unique_ptr<AsyncStartThunk> WrapInAsyncStartThunk(
   Thunk::ThunkInfo thunk_info;
   thunk_info.thunk_id = ThunkId(next_id.fetch_add(1));
   return std::make_unique<AsyncStartThunk>(thunk_info, CommunicationStreamId(0),
-                                           std::move(sequence));
+                                           std::move(sequence),
+                                           /*devices_per_host=*/1);
 }
 
 std::unique_ptr<AsyncDoneThunk> CreateAllGatherDoneThunk(Thunk* start_thunk) {
@@ -771,7 +772,8 @@ TEST(CommandBufferConversionPassTest,
   Thunk::ThunkInfo thunk_info;
   thunk_info.thunk_id = ThunkId(999);
   thunks.push_back(std::make_unique<AsyncStartThunk>(
-      thunk_info, ComputationStreamId(1), std::move(nested_sequence)));
+      thunk_info, ComputationStreamId(1), std::move(nested_sequence),
+      /*devices_per_host=*/1));
   thunks.push_back(CreateAllGatherDoneThunk(thunks.back().get()));
 
   DebugOptions debug_options = xla::GetDebugOptionsFromFlags();
@@ -809,7 +811,8 @@ TEST(CommandBufferConversionPassTest,
   ThunkSequence async_seq;
   async_seq.push_back(std::move(group_thunk));
   thunks.push_back(std::make_unique<AsyncStartThunk>(
-      start_info, CommunicationStreamId(0), std::move(async_seq)));
+      start_info, CommunicationStreamId(0), std::move(async_seq),
+      /*devices_per_host=*/1));
   thunks.push_back(CreateAllGatherDoneThunk(thunks.back().get()));
 
   DebugOptions debug_options = xla::GetDebugOptionsFromFlags();
@@ -845,7 +848,8 @@ TEST(CommandBufferConversionPassTest,
   ThunkSequence async_seq;
   async_seq.push_back(std::move(group_thunk));
   thunks.push_back(std::make_unique<AsyncStartThunk>(
-      start_info, CommunicationStreamId(0), std::move(async_seq)));
+      start_info, CommunicationStreamId(0), std::move(async_seq),
+      /*devices_per_host=*/1));
   thunks.push_back(CreateAllGatherDoneThunk(thunks.back().get()));
 
   DebugOptions debug_options = xla::GetDebugOptionsFromFlags();
@@ -1486,7 +1490,8 @@ TEST(CommandBufferConversionPassTest, ConvertAsyncStartDonePair) {
 
   // Create AsyncStartThunk with an empty nested sequence.
   auto start_thunk = std::make_unique<AsyncStartThunk>(
-      Thunk::ThunkInfo(), ComputationStreamId(0), ThunkSequence{});
+      Thunk::ThunkInfo(), ComputationStreamId(0), ThunkSequence{},
+      /*devices_per_host=*/1);
   auto async_execution = start_thunk->async_execution();
   thunks.push_back(std::move(start_thunk));
 
@@ -1512,7 +1517,8 @@ TEST(CommandBufferConversionPassTest,
 
   // Create AsyncStartThunk with an empty nested sequence.
   auto start_thunk = std::make_unique<AsyncStartThunk>(
-      Thunk::ThunkInfo(), ComputationStreamId(0), ThunkSequence{});
+      Thunk::ThunkInfo(), ComputationStreamId(0), ThunkSequence{},
+      /*devices_per_host=*/1);
   auto async_execution = start_thunk->async_execution();
   thunks.push_back(std::move(start_thunk));
 
diff --git a/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk_test.cc
index 107c1e61eae..031a1c26949 100644
--- a/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk_test.cc
@@ -435,7 +435,8 @@ TEST_F(GpuRaggedAllToAllTest, TestConvertToCommands) {
   ThunkSequence start_sequence;
   start_sequence.push_back(std::move(ra2a_start_thunk));
   auto async_start = std::make_unique<AsyncStartThunk>(
-      Thunk::ThunkInfo(), CommunicationStreamId(0), std::move(start_sequence));
+      Thunk::ThunkInfo(), CommunicationStreamId(0), std::move(start_sequence),
+      /*devices_per_host=*/1);
   auto async_done = std::make_unique<AsyncDoneThunk>(
       Thunk::ThunkInfo(), async_start->async_execution());
 
diff --git a/third_party/xla/xla/backends/gpu/runtime/send_recv_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/send_recv_thunk_test.cc
index e847a3ae534..3875ef2bbdc 100644
--- a/third_party/xla/xla/backends/gpu/runtime/send_recv_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/send_recv_thunk_test.cc
@@ -135,7 +135,8 @@ ENTRY computation {
   send_sequence.push_back(std::move(send_thunk));
 
   auto async_start = std::make_unique<AsyncStartThunk>(
-      Thunk::ThunkInfo{}, CommunicationStreamId(0), std::move(send_sequence));
+      Thunk::ThunkInfo{}, CommunicationStreamId(0), std::move(send_sequence),
+      /*devices_per_host=*/1);
 
   auto async_done = std::make_unique<AsyncDoneThunk>(
       Thunk::ThunkInfo{}, async_start->async_execution());
@@ -229,7 +230,8 @@ ENTRY computation {
   recv_sequence.push_back(std::move(recv_thunk));
 
   auto async_start = std::make_unique<AsyncStartThunk>(
-      Thunk::ThunkInfo{}, CommunicationStreamId(0), std::move(recv_sequence));
+      Thunk::ThunkInfo{}, CommunicationStreamId(0), std::move(recv_sequence),
+      /*devices_per_host=*/1);
 
   auto async_done = std::make_unique<AsyncDoneThunk>(
       Thunk::ThunkInfo{}, async_start->async_execution());
diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
index 3809e0687b8..30beccca37f 100644
--- a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
@@ -326,9 +326,10 @@ absl::StatusOr<std::unique_ptr<Thunk>> DeserializeThunkProtoImpl(
           gpu_topology.has_value() ? gpu_topology->num_devices_per_process()
                                    : 0);
     case ThunkProto::kAsyncStartThunk:
-      return AsyncStartThunk::FromProto(std::move(thunk_info),
-                                        thunk_proto.async_start_thunk(),
-                                        deserializer, async_execution_map);
+      return AsyncStartThunk::FromProto(
+          std::move(thunk_info), thunk_proto.async_start_thunk(), deserializer,
+          async_execution_map,
+          gpu_topology.has_value() ? gpu_topology->num_devices_per_host() : 0);
     case ThunkProto::kAsyncDoneThunk:
       return AsyncDoneThunk::FromProto(std::move(thunk_info),
                                        thunk_proto.async_done_thunk(),
diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization_test.cc b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization_test.cc
index 00194787f03..b68bfe16ac5 100644
--- a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization_test.cc
@@ -1310,7 +1310,7 @@ TEST(ThunkProtoDeserializationTest, AsyncStartAndDoneThunk) {
   start_info.profile_annotation = "async_start";
 
   AsyncStartThunk start_thunk(start_info, ComputationStreamId(0),
-                              ThunkSequence{});
+                              ThunkSequence{}, /*devices_per_host=*/1);
 
   AsyncDoneThunk done_thunk(Thunk::ThunkInfo(), start_thunk.async_execution());
 
@@ -1357,7 +1357,8 @@ TEST(ThunkProtoDeserializationTest, AsyncStartThunkMemcpyStreamRoundTrip) {
     Thunk::ThunkInfo start_info;
     start_info.profile_annotation = "memcpy_async_start";
 
-    AsyncStartThunk start_thunk(start_info, stream_id, ThunkSequence{});
+    AsyncStartThunk start_thunk(start_info, stream_id, ThunkSequence{},
+                                /*devices_per_host=*/1);
     AsyncDoneThunk done_thunk(Thunk::ThunkInfo(),
                               start_thunk.async_execution());
 
diff --git a/third_party/xla/xla/service/gpu/thunk_emitter.cc b/third_party/xla/xla/service/gpu/thunk_emitter.cc
index 836ea29ac68..32e6857f9a0 100644
--- a/third_party/xla/xla/service/gpu/thunk_emitter.cc
+++ b/third_party/xla/xla/service/gpu/thunk_emitter.cc
@@ -486,7 +486,9 @@ ThunkEmitter::RegisterAsyncExecution(const HloInstruction* async_start) {
   // while the nested-emission future is still pending.
   Thunk::ThunkInfo info = Thunk::ThunkInfo::WithProfileAnnotation(
       async_start, ir_emitter_context_->GetNextThunkId());
-  auto execution = std::make_shared<AsyncExecution>(std::move(info));
+  auto execution = std::make_shared<AsyncExecution>(
+      std::move(info),
+      ir_emitter_context_->gpu_topology().num_devices_per_host());
   auto [_, inserted] = hlo_async_executions_.emplace(async_start, execution);
   if (!inserted) {
     return Internal("Async execution already exists for instruction %s",
@@ -2457,7 +2459,8 @@ absl::StatusOr<ThunkSequence> ThunkEmitter::EmitCopyStart(
   auto start_thunk = std::make_unique<AsyncStartThunk>(
       Thunk::ThunkInfo::WithProfileAnnotation(
           copy_start_instr, ir_emitter_context_->GetNextThunkId()),
-      *execution_stream_id, std::move(nested_thunks));
+      *execution_stream_id, std::move(nested_thunks),
+      ir_emitter_context_->gpu_topology().num_devices_per_host());
 
   auto [it, inserted] = hlo_async_executions_.emplace(
       copy_start_instr, start_thunk->async_execution());
@@ -2760,7 +2763,9 @@ absl::StatusOr<ThunkSequence> ThunkEmitter::EmitAsyncSendRecvStart(
   Thunk::ThunkInfo info = Thunk::ThunkInfo::WithProfileAnnotation(
       async_start, ir_emitter_context_->GetNextThunkId());
   auto [it, inserted] = hlo_async_executions_.emplace(
-      owner, std::make_shared<AsyncExecution>(info));
+      owner,
+      std::make_shared<AsyncExecution>(
+          info, ir_emitter_context_->gpu_topology().num_devices_per_host()));
   if (!inserted) {
     return Internal("Async execution already exists for instruction %s",
                     owner->ToString());

From 1ec82569f093de3c162e6800faee932cba5f926c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Eusebio=20Dur=C3=A1n=20Monta=C3=B1a?= <eusebiodm@google.com>
Date: Thu, 8 Oct 2026 02:46:46 -0700
Subject: [PATCH 0269/1000] Preload cuda_compat libcuda.so.1 for deviceless
 cuDNN compilation

Invoke `EnsureCudaCompatLoaded()` via a weak symbol from the deviceless
 branch of `CudnnGraph::Prepare()` so deviceless cuDNN
graph compilation and support probes can resolve `libcuda.so.1` symbols on
CPU-only build hosts when linked with `cuda_compat`.

PiperOrigin-RevId: 995698289
---
 third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc b/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc
index 2839cd253d4..3a51567e2bc 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc
@@ -31,6 +31,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/algorithm/container.h"
+#include "absl/base/attributes.h"
 #include "absl/base/casts.h"
 #include "absl/base/optimization.h"
 #include "absl/base/thread_annotations.h"
@@ -6895,6 +6896,8 @@ absl::Status PlanEnumerationStatus(bool deviceless,
                                           "): ", result.get_message()));
 }
 
+extern "C" ABSL_ATTRIBUTE_WEAK void EnsureCudaCompatLoaded();
+
 }  // namespace
 
 absl::Status CudnnGraph::Prepare(dnn::DnnSupport* dnn_support,
@@ -6936,6 +6939,9 @@ absl::Status CudnnGraph::Prepare(dnn::DnnSupport* dnn_support,
   } else {
     // Deviceless mode. No cuDNN version guard needed: DeviceProperties
     // deserialization inside BuildDeviceProperties rejects runtimes < 9.8.
+    if (EnsureCudaCompatLoaded != nullptr) {
+      EnsureCudaCompatLoaded();
+    }
     ABSL_ASSIGN_OR_RETURN(auto device_props,
                      xla::gpu::BuildDeviceProperties(gpu_device_info));
     graph_.set_device_properties(device_props);

From 534be6a97e4f2a92db9a2143d732a68198c1e9cd Mon Sep 17 00:00:00 2001
From: bodapatisaikrishna <saikrishnabodapati@gmail.com>
Date: Thu, 8 Oct 2026 16:46:15 +0530
Subject: [PATCH 0270/1000] Clean up exception regex in conv2d rank tests

Remove dead alternation from regex to strictly match the C++ kernel error
and graph mode shape inference error.

Signed-off-by: bodapatisaikrishna <saikrishnabodapati@gmail.com>
---
 .../python/kernel_tests/nn_ops/conv_ops_test.py      | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
index f142a42e313..e8b52976588 100644
--- a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
@@ -3137,8 +3137,7 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
     # Invalid input or filter rank (< 4) in eager/graph execution.
     with self.assertRaisesRegex(
         (ValueError, errors_impl.InvalidArgumentError),
-        r"(must be (at least )?4-dimensional|"
-        r"Shape must be rank 4 but is rank 3)",
+        r"(must be 4-dimensional|Shape must be rank 4 but is rank 3)",
     ):
       self.evaluate(
           nn_ops.conv2d(
@@ -3150,8 +3149,7 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
       )
     with self.assertRaisesRegex(
         (ValueError, errors_impl.InvalidArgumentError),
-        r"(must be (at least )?4-dimensional|"
-        r"Shape must be rank 4 but is rank 3)",
+        r"(must be 4-dimensional|Shape must be rank 4 but is rank 3)",
     ):
       self.evaluate(
           nn_ops.conv2d(
@@ -3165,8 +3163,7 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
     # Ensure raw ops are protected by the C++ kernel boundary.
     with self.assertRaisesRegex(
         (ValueError, errors_impl.InvalidArgumentError),
-        r"(must be (at least )?4-dimensional|"
-        r"Shape must be rank 4 but is rank 3)",
+        r"(must be 4-dimensional|Shape must be rank 4 but is rank 3)",
     ):
       self.evaluate(
           raw_ops.Conv2D(
@@ -3178,8 +3175,7 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
       )
     with self.assertRaisesRegex(
         (ValueError, errors_impl.InvalidArgumentError),
-        r"(must be (at least )?4-dimensional|"
-        r"Shape must be rank 4 but is rank 3)",
+        r"(must be 4-dimensional|Shape must be rank 4 but is rank 3)",
     ):
       self.evaluate(
           raw_ops.Conv2D(

From a18cf31fd1388ec2eb543ce4fe94b39254e4e3ee Mon Sep 17 00:00:00 2001
From: Shaogang Wang <shawnw@nvidia.com>
Date: Thu, 8 Oct 2026 02:55:00 -0700
Subject: [PATCH 0271/1000] PR #50176: [XLA:GPU] StreamExecutor: VA reservation
 and physical allocation APIs sharing the executor's VMM options
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50176

**Stacked PR 1/5** splitting #48985 into reviewable layers. First PR of the stack; opened against `main`.

| # | PR | Title |
|---|---|---|
| 1 | #50176 | [XLA:GPU] StreamExecutor: VA reservation and physical allocation APIs sharing the executor's VMM options |
| 2 | #50177 | [XLA:GPU] DeviceMemAllocator: reserved-arena mode that maps physical memory on demand |
| 3 | #50178 | [XLA:TSL] BFC: grow spatially partitioned pools within a reserved arena |
| 4 | #50179 | [XLA:GPU] PJRT: typed MemFraction memory policy with START-CAP grammar |
| 5 | #50180 | [XLA:GPU] PJRT: enable BFC growth for growable memory fractions |

---

📝 Summary of Changes

Adds `StreamExecutor::CreateMemoryReservation` and `CreatePhysicalMemoryAllocation`, implemented on CUDA by `CudaMemoryReservation` and `CudaRawMemoryAllocation`. Both gain an overload that takes the executor's probed `CudaDeviceAllocator::Options`, so reservations, physical allocations and the existing device allocator agree on handle types and mapping granularity (including the fabric-handle decision made in `CudaExecutor::Init`). The cuMemCreate fallback (FABRIC+POSIX_FD → POSIX_FD → NONE, also on `CUDA_ERROR_INVALID_VALUE`) is exported from `cuda_device_allocator` as `CreateVmmPhysicalAllocation` and reused by the raw allocation.

`MemoryReservation::granularity()` reports the mapping granularity. `GetAllocationRange` and `ExportFabricHandle` describe the mapping backing a pointer instead of the whole VA range, since a reservation may be larger than its mapped prefix. The reservation destructor no longer unmaps the full range; each `ScopedMapping` owns its slice.

🎯 Justification

Groundwork for growing the spatially partitioned GPU BFC pool inside a reserved virtual arena. No behavior change for existing callers.

🚀 Kind of Contribution

🛠 Refactor / infrastructure

🧪 Unit Tests:

`//xla/stream_executor/cuda:cuda_raw_memory_allocation_test`, `//xla/stream_executor/cuda:cuda_memory_reservation_test` (GPU). Includes a test that requesting exportable handle types succeeds on every machine through the fallback.

🧪 Execution Tests:

Not applicable; API-only change.

🤖 Generated with [Claude Code](https://claude.com/claude-code)
Copybara import of the project:

--
d6435ded433da3397e007d0f313b1efe2b2ac7b4 by Shaogang Wang <shawnw@nvidia.com>:

[XLA:GPU] StreamExecutor: VA reservation and physical allocation APIs sharing the executor's VMM options

Adds StreamExecutor::CreateMemoryReservation and
CreatePhysicalMemoryAllocation, implemented on CUDA with
CudaMemoryReservation and CudaRawMemoryAllocation, and lets both take the
executor's probed CudaDeviceAllocator::Options so reservation, physical
allocation and the existing device allocator agree on handle types and
granularity. The cuMemCreate fallback (FABRIC+POSIX_FD -> POSIX_FD -> NONE,
including CUDA_ERROR_INVALID_VALUE) is exported from cuda_device_allocator
as CreateVmmPhysicalAllocation and used by the raw allocation.

MemoryReservation reports its mapping granularity. GetAllocationRange and
ExportFabricHandle describe the mapping that backs a pointer rather than the
whole VA range, since a reservation can be larger than its mapped prefix.
The reservation destructor no longer unmaps the full range; ScopedMapping
owns each mapped slice.

No behavior change for existing callers. First of five stacked changes that
let the spatially partitioned GPU BFC pool grow (see #48985).

--
c6ac0b498d0798480b9b62da9008ffabf9cc0e0b by Shaogang Wang <shawnw@nvidia.com>:

[XLA:GPU] Honor the VMM opt-out in reservation APIs

--
e175d96fad86158884e2cf17a33b73c05b620992 by Shaogang Wang <shawnw@nvidia.com>:

[XLA:GPU] Address review on the VA reservation and physical allocation APIs

- Probe the VMM allocation granularity with the same handle-type fallback as
  CudaDeviceAllocator (ProbeVmmGranularity, now shared with
  QueryDeviceAllocatorOptions), so callers passing unprobed options do not
  fail before cuMemCreate's fallback, and create with the accepted types.
- Honor Options::alignment as max(alignment, granularity) in both Create
  overloads, as CudaDeviceAllocator does; the reservation keeps reporting the
  driver's mapping granularity.
- Reject options with use_vmm == false in the VMM-only Create overloads.
- Resolve the device once per Create call (CreateWithDevice).
- Make MemoryReservation::granularity() pure virtual; report the probed
  granularity from RocmMemoryReservation and add overrides to the host-backed
  test fakes.
- Test the enabled-VMM path through CudaExecutor: reservation, physical
  allocation, MapTo, device access, and GetAllocationRange returning the
  mapped granule rather than the two-granule reservation. Add unit tests for
  the rejected opt-out and for alignment padding.

Merging this change closes #50176

PiperOrigin-RevId: 995702730
---
 .../gpu_executable_va_remap_allocator_test.cc |  3 +
 third_party/xla/xla/stream_executor/BUILD     |  1 +
 .../xla/xla/stream_executor/cuda/BUILD        | 18 ++--
 .../cuda/cuda_device_allocator.cc             | 57 ++++++------
 .../cuda/cuda_device_allocator.h              | 24 +++++
 .../xla/stream_executor/cuda/cuda_executor.cc | 32 ++++++-
 .../xla/stream_executor/cuda/cuda_executor.h  |  5 ++
 .../cuda/cuda_executor_test.cc                | 88 +++++++++++++++++++
 .../cuda/cuda_memory_reservation.cc           | 59 +++++++++----
 .../cuda/cuda_memory_reservation.h            | 29 +++++-
 .../cuda/cuda_memory_reservation_test.cc      | 51 +++++++++--
 .../cuda/cuda_raw_memory_allocation.cc        | 55 ++++++++----
 .../cuda/cuda_raw_memory_allocation.h         | 22 ++++-
 .../cuda/cuda_raw_memory_allocation_test.cc   | 56 +++++++++++-
 .../device_address_vmm_allocator_test.cc      |  3 +
 .../xla/stream_executor/memory_reservation.h  |  6 ++
 .../rocm/rocm_memory_reservation.cc           |  7 +-
 .../rocm/rocm_memory_reservation.h            |  7 +-
 .../xla/xla/stream_executor/stream_executor.h | 16 ++++
 19 files changed, 450 insertions(+), 89 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/gpu_executable_va_remap_allocator_test.cc b/third_party/xla/xla/service/gpu/gpu_executable_va_remap_allocator_test.cc
index f3593f1c483..3ac037bd255 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable_va_remap_allocator_test.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable_va_remap_allocator_test.cc
@@ -114,6 +114,9 @@ class TestMemoryReservation final : public se::MemoryReservation {
     return se::DeviceAddressBase(storage_.data(), size_);
   }
 
+  // Matches the allocation granularity the test installs on the allocator.
+  size_t granularity() const override { return kGranularity; }
+
   int active_mapping_count() const { return active_mapping_count_; }
   const void* last_mapped_allocation_address() const {
     return last_mapped_allocation_address_;
diff --git a/third_party/xla/xla/stream_executor/BUILD b/third_party/xla/xla/stream_executor/BUILD
index 370dad2ce71..07847c0a9eb 100644
--- a/third_party/xla/xla/stream_executor/BUILD
+++ b/third_party/xla/xla/stream_executor/BUILD
@@ -701,6 +701,7 @@ cc_library(
         ":kernel_spec",
         ":memory_allocation",
         ":memory_allocator",
+        ":memory_reservation",
         ":memory_space",
         ":module_spec",
         ":platform",
diff --git a/third_party/xla/xla/stream_executor/cuda/BUILD b/third_party/xla/xla/stream_executor/cuda/BUILD
index be9c1ffb386..149fa607070 100644
--- a/third_party/xla/xla/stream_executor/cuda/BUILD
+++ b/third_party/xla/xla/stream_executor/cuda/BUILD
@@ -852,6 +852,7 @@ cc_library(
         "//xla/stream_executor:memory_allocation",
         "//xla/stream_executor:stream_executor_h",
         "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
         "@local_config_cuda//cuda:cuda_headers",
@@ -914,18 +915,18 @@ xla_test(
     backends = ["gpu"],
     tags = ["cuda-only"],
     deps = [
+        ":cuda_device_allocator",
         ":cuda_platform_id",
         ":cuda_raw_memory_allocation",
         "//xla/stream_executor:platform",
         "//xla/stream_executor:platform_manager",
         "//xla/stream_executor:stream_executor_h",
         "//xla/tsl/lib/core:status_test_util",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",  # build_cleaner: keep
         "@com_google_googletest//:gtest",
         "@com_google_googletest//:gtest_main",
         "@local_config_cuda//cuda:cuda_headers",
-        "@tsl//tsl/platform:statusor",
-        "@tsl//tsl/platform:test",
     ],
 )
 
@@ -935,6 +936,7 @@ xla_test(
     backends = ["gpu"],
     tags = ["cuda-only"],
     deps = [
+        ":cuda_device_allocator",
         ":cuda_memory_reservation",
         ":cuda_platform_id",
         ":cuda_raw_memory_allocation",
@@ -944,15 +946,11 @@ xla_test(
         "//xla/stream_executor:platform",
         "//xla/stream_executor:platform_manager",
         "//xla/stream_executor:stream_executor_h",
-        "//xla/tsl/lib/core:status_test_util",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",  # build_cleaner: keep
         "@com_google_absl//absl/types:span",
-        "@com_google_googletest//:gtest",
         "@com_google_googletest//:gtest_main",
         "@local_config_cuda//cuda:cuda_headers",
-        "@tsl//tsl/platform:statusor",
-        "@tsl//tsl/platform:test",
     ],
 )
 
@@ -1578,6 +1576,7 @@ cc_library(
         "//xla/stream_executor:kernel_spec",
         "//xla/stream_executor:memory_allocation",
         "//xla/stream_executor:memory_allocator",
+        "//xla/stream_executor:memory_reservation",
         "//xla/stream_executor:memory_space",
         "//xla/stream_executor:module_spec",
         "//xla/stream_executor:platform",
@@ -1617,7 +1616,9 @@ cc_library(
         ":cuda_executor",
         ":cuda_host_allocator",
         ":cuda_kernel",
+        ":cuda_memory_reservation",
         ":cuda_platform_id",
+        ":cuda_raw_memory_allocation",
         ":cuda_status",
         ":cuda_stream",
         ":cuda_timer",
@@ -1652,6 +1653,7 @@ cc_library(
         "//xla/stream_executor:launch_dim",
         "//xla/stream_executor:memory_allocation",
         "//xla/stream_executor:memory_allocator",
+        "//xla/stream_executor:memory_reservation",
         "//xla/stream_executor:memory_space",
         "//xla/stream_executor:module_spec",
         "//xla/stream_executor:platform",
@@ -1712,6 +1714,8 @@ xla_test(
         ":cuda_executor",
         ":cuda_platform",
         ":cuda_platform_id",
+        "//xla:debug_options_flags",
+        "//xla:xla_proto_cc",
         "//xla/stream_executor:device_address",
         "//xla/stream_executor:device_description",
         "//xla/stream_executor:kernel",
@@ -1725,6 +1729,8 @@ xla_test(
         "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor/gpu:gpu_test_kernels",
         "//xla/tsl/platform:statusor",
+        "//xla/tsl/util:command_line_flags",
+        "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",  # build_cleaner: keep
         "@com_google_absl//absl/strings",
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_device_allocator.cc b/third_party/xla/xla/stream_executor/cuda/cuda_device_allocator.cc
index 63ea39b673c..b3c2ec7deca 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_device_allocator.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_device_allocator.cc
@@ -96,40 +96,31 @@ CUmemAllocationProp BuildVmmAllocationProp(
   return properties;
 }
 
-absl::StatusOr<CudaDeviceAllocator::Options> QueryDeviceAllocatorOptions(
-    CUdevice device) {
-  ABSL_ASSIGN_OR_RETURN(bool rdma, IsRdmaSupported(device));
-  ABSL_ASSIGN_OR_RETURN(bool fabric, IsFabricSupported(device));
-
-  bool posix_fd = true;
-  size_t granularity = 0;
-
+absl::StatusOr<VmmGranularityProbe> ProbeVmmGranularity(
+    CUdevice device, CudaDeviceAllocator::Options options) {
+  VmmGranularityProbe probe;
   auto try_query = [&]() -> absl::Status {
-    CudaDeviceAllocator::Options opts;
-    opts.enable_rdma = rdma;
-    opts.enable_posix_fd_handle = posix_fd;
-    opts.enable_fabric_handle = fabric;
-    CUmemAllocationProp props = BuildVmmAllocationProp(device, opts);
+    CUmemAllocationProp props = BuildVmmAllocationProp(device, options);
     return cuda::ToStatus(cuMemGetAllocationGranularity(
-        &granularity, &props, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED));
+        &probe.granularity, &props, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED));
   };
 
   absl::Status status = try_query();
-  if (!status.ok() && fabric && posix_fd) {
+  if (!status.ok() && options.enable_fabric_handle) {
     XLA_LOG_DEVICE(WARNING, device)
-        << "Device allocator granularity query with FABRIC+POSIX_FD handle "
-           "types failed: "
-        << status << "; retrying without FABRIC.";
-    fabric = false;
+        << "Device allocator granularity query with "
+        << (options.enable_posix_fd_handle ? "FABRIC+POSIX_FD" : "FABRIC")
+        << " handle types failed: " << status << "; retrying without FABRIC.";
+    options.enable_fabric_handle = false;
     status = try_query();
   }
 
-  if (!status.ok() && posix_fd) {
+  if (!status.ok() && options.enable_posix_fd_handle) {
     XLA_LOG_DEVICE(WARNING, device)
         << "Device allocator granularity query with POSIX_FD handle type "
            "failed: "
         << status << "; retrying with HANDLE_TYPE_NONE.";
-    posix_fd = false;
+    options.enable_posix_fd_handle = false;
     status = try_query();
   }
 
@@ -140,18 +131,28 @@ absl::StatusOr<CudaDeviceAllocator::Options> QueryDeviceAllocatorOptions(
     return status;
   }
 
-  CudaDeviceAllocator::Options options;
-  options.alignment = granularity;
-  options.enable_rdma = rdma;
-  options.enable_posix_fd_handle = posix_fd;
-  options.enable_fabric_handle = fabric;
+  probe.options = options;
+  return probe;
+}
+
+absl::StatusOr<CudaDeviceAllocator::Options> QueryDeviceAllocatorOptions(
+    CUdevice device) {
+  CudaDeviceAllocator::Options requested;
+  ABSL_ASSIGN_OR_RETURN(requested.enable_rdma, IsRdmaSupported(device));
+  ABSL_ASSIGN_OR_RETURN(requested.enable_fabric_handle, IsFabricSupported(device));
+  requested.enable_posix_fd_handle = true;
+
+  ABSL_ASSIGN_OR_RETURN(VmmGranularityProbe probe,
+                   ProbeVmmGranularity(device, requested));
+  CudaDeviceAllocator::Options options = probe.options;
+  options.alignment = probe.granularity;
   return options;
 }
 
 // Creates a physical VMM allocation. Tries cuMemCreate with the given
 // properties and falls back through progressively simpler handle types:
 //   FABRIC+POSIX_FD -> POSIX_FD -> NONE
-static absl::StatusOr<CUmemGenericAllocationHandle> CreatePhysicalAllocation(
+absl::StatusOr<CUmemGenericAllocationHandle> CreateVmmPhysicalAllocation(
     CUmemAllocationProp properties, uint64_t padded_size) {
   CUmemGenericAllocationHandle handle;
 
@@ -237,7 +238,7 @@ AllocateDeviceMemory(StreamExecutor* executor,
   uint64_t padded_size = xla::RoundUpTo<uint64_t>(size, effective_alignment);
 
   ABSL_ASSIGN_OR_RETURN(CUmemGenericAllocationHandle handle,
-                   CreatePhysicalAllocation(properties, padded_size));
+                   CreateVmmPhysicalAllocation(properties, padded_size));
 
   absl::Cleanup release_handle = [&] {
     absl::Status status = cuda::ToStatus(cuMemRelease(handle));
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_device_allocator.h b/third_party/xla/xla/stream_executor/cuda/cuda_device_allocator.h
index 17912a52918..22b2247f141 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_device_allocator.h
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_device_allocator.h
@@ -103,6 +103,30 @@ CUmemAllocationProp BuildVmmAllocationProp(
 absl::StatusOr<CudaDeviceAllocator::Options> QueryDeviceAllocatorOptions(
     CUdevice device);
 
+// The handle types the driver accepted for a granularity query, and the
+// recommended mapping granularity for them.
+struct VmmGranularityProbe {
+  CudaDeviceAllocator::Options options;
+  size_t granularity = 0;
+};
+
+// Queries the recommended VMM allocation granularity for `options`. Drivers can
+// reject the query itself for unsupported handle types, so this falls back the
+// same way CreateVmmPhysicalAllocation does (FABRIC+POSIX_FD -> POSIX_FD ->
+// NONE) and returns `options` with the handle types that were accepted. The
+// caller must have activated the device context.
+absl::StatusOr<VmmGranularityProbe> ProbeVmmGranularity(
+    CUdevice device, CudaDeviceAllocator::Options options);
+
+// Creates a physical VMM allocation of `padded_size` bytes with cuMemCreate,
+// falling back through simpler handle types (FABRIC+POSIX_FD -> POSIX_FD ->
+// NONE) when the driver reports NOT_PERMITTED, NOT_SUPPORTED or INVALID_VALUE.
+// The returned handle may therefore carry fewer handle types than `properties`
+// requested; callers that computed `padded_size` from the original properties
+// accept that. The caller must have activated the device context.
+absl::StatusOr<CUmemGenericAllocationHandle> CreateVmmPhysicalAllocation(
+    CUmemAllocationProp properties, uint64_t padded_size);
+
 }  // namespace stream_executor::gpu
 
 #endif  // XLA_STREAM_EXECUTOR_CUDA_CUDA_DEVICE_ALLOCATOR_H_
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc b/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc
index 61a8b2a47d7..35eb43d3433 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc
@@ -67,7 +67,9 @@ limitations under the License.
 #include "xla/stream_executor/cuda/cuda_event.h"
 #include "xla/stream_executor/cuda/cuda_host_allocator.h"
 #include "xla/stream_executor/cuda/cuda_kernel.h"
+#include "xla/stream_executor/cuda/cuda_memory_reservation.h"
 #include "xla/stream_executor/cuda/cuda_platform_id.h"
+#include "xla/stream_executor/cuda/cuda_raw_memory_allocation.h"
 #include "xla/stream_executor/cuda/cuda_status.h"
 #include "xla/stream_executor/cuda/cuda_stream.h"
 #include "xla/stream_executor/cuda/cuda_timer.h"
@@ -100,6 +102,7 @@ limitations under the License.
 #include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/memory_allocation.h"
 #include "xla/stream_executor/memory_allocator.h"
+#include "xla/stream_executor/memory_reservation.h"
 #include "xla/stream_executor/memory_space.h"
 #include "xla/stream_executor/module_spec.h"
 #include "xla/stream_executor/platform.h"
@@ -855,9 +858,12 @@ absl::StatusOr<size_t> CudaExecutor::GetVmmGranularity() const {
 absl::StatusOr<DeviceAddressBase> CudaExecutor::GetAllocationRange(
     void* ptr) const {
   uint64_t alloc_start, alloc_size;
+  // RANGE_* describes the whole VA reservation, which can include unmapped
+  // addresses and multiple physical allocations. Fabric transfers need the
+  // bounds of the mapping backed by the handle containing ptr.
   CUpointer_attribute attrs[2] = {
-      CU_POINTER_ATTRIBUTE_RANGE_START_ADDR,
-      CU_POINTER_ATTRIBUTE_RANGE_SIZE,
+      CU_POINTER_ATTRIBUTE_MAPPING_BASE_ADDR,
+      CU_POINTER_ATTRIBUTE_MAPPING_SIZE,
   };
   void* results[2] = {&alloc_start, &alloc_size};
   ABSL_RETURN_IF_ERROR(cuda::ToStatus(cuPointerGetAttributes(
@@ -883,7 +889,9 @@ absl::StatusOr<std::string> CudaExecutor::ExportFabricHandle(void* ptr) const {
       static_cast<CUmemGenericAllocationHandle>(handle.handle()),
       CU_MEM_HANDLE_TYPE_FABRIC, 0)));
 
-  CUpointer_attribute attrs[1] = {CU_POINTER_ATTRIBUTE_RANGE_SIZE};
+  // Import only the mapped allocation exported by this handle, not the larger
+  // VA reservation that a growing arena may have reserved around it.
+  CUpointer_attribute attrs[1] = {CU_POINTER_ATTRIBUTE_MAPPING_SIZE};
   void* results[1] = {&fabric_handle.size};
   ABSL_RETURN_IF_ERROR(cuda::ToStatus(cuPointerGetAttributes(
       1, attrs, results, reinterpret_cast<CUdeviceptr>(ptr))));
@@ -975,6 +983,24 @@ CudaExecutor::CreateMemoryAllocator(MemorySpace type) {
       absl::StrFormat("Unsupported memory type %d", type));
 }
 
+absl::StatusOr<std::unique_ptr<MemoryReservation>>
+CudaExecutor::CreateMemoryReservation(uint64_t size) {
+  if (!device_allocator_options_.use_vmm) {
+    return absl::UnimplementedError("CUDA VMM is disabled");
+  }
+  // Use the options probed in Init() (e.g. fabric handles disabled outside a
+  // cluster) so the reservation granularity matches physical allocations.
+  return CudaMemoryReservation::Create(this, size, device_allocator_options_);
+}
+
+absl::StatusOr<std::unique_ptr<MemoryAllocation>>
+CudaExecutor::CreatePhysicalMemoryAllocation(uint64_t size) {
+  if (!device_allocator_options_.use_vmm) {
+    return absl::UnimplementedError("CUDA VMM is disabled");
+  }
+  return CudaRawMemoryAllocation::Create(this, size, device_allocator_options_);
+}
+
 absl::Status CudaExecutor::Init() {
   ABSL_ASSIGN_OR_RETURN(device_, GetDevice(device_ordinal()));
   const bool vmm_disabled =
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_executor.h b/third_party/xla/xla/stream_executor/cuda/cuda_executor.h
index c82e946e0a6..6821e7769f7 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_executor.h
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_executor.h
@@ -62,6 +62,7 @@ limitations under the License.
 #include "xla/stream_executor/kernel_spec.h"
 #include "xla/stream_executor/memory_allocation.h"
 #include "xla/stream_executor/memory_allocator.h"
+#include "xla/stream_executor/memory_reservation.h"
 #include "xla/stream_executor/memory_space.h"
 #include "xla/stream_executor/module_spec.h"
 #include "xla/stream_executor/platform.h"
@@ -198,6 +199,10 @@ class CudaExecutor : public GpuExecutor {
                                             void* global_address) override;
   absl::StatusOr<std::unique_ptr<MemoryAllocator>> CreateMemoryAllocator(
       MemorySpace type) override;
+  absl::StatusOr<std::unique_ptr<MemoryReservation>> CreateMemoryReservation(
+      uint64_t size) override;
+  absl::StatusOr<std::unique_ptr<MemoryAllocation>>
+  CreatePhysicalMemoryAllocation(uint64_t size) override;
 
   // Returns the granularity which is the minimum unit of memory that can be
   // allocated with VMM API. In order to map the memory slices to multicast
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_executor_test.cc b/third_party/xla/xla/stream_executor/cuda/cuda_executor_test.cc
index 4d59eb76087..ca62d495b4a 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_executor_test.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_executor_test.cc
@@ -18,14 +18,17 @@ limitations under the License.
 #include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <string>
 #include <utility>
 #include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/cleanup/cleanup.h"
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"  // IWYU pragma: keep
 #include "absl/strings/match.h"
+#include "xla/debug_options_flags.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/cuda/cuda_platform.h"
 #include "xla/stream_executor/cuda/cuda_platform_id.h"
@@ -42,6 +45,8 @@ limitations under the License.
 #include "xla/stream_executor/semantic_version.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/util/command_line_flags.h"
+#include "xla/xla.pb.h"
 
 namespace stream_executor::gpu {
 namespace {
@@ -309,6 +314,89 @@ TEST(CudaExecutorTest, MultipleExecutorsForSameDevice) {
   }
 }
 
+TEST(CudaExecutorTest, DisabledVmmRejectsReservationAndPhysicalAllocation) {
+  const bool was_disabled =
+      xla::GetDebugOptionsFromFlags().xla_gpu_experimental_vmm_disabled();
+  std::vector<tsl::Flag> flags;
+  xla::AppendDebugOptionsFlags(&flags);
+  auto set_vmm_disabled = [&](bool disabled) {
+    std::vector<std::string> args = {
+        disabled ? "--xla_gpu_experimental_vmm_disabled=true"
+                 : "--xla_gpu_experimental_vmm_disabled=false"};
+    EXPECT_TRUE(tsl::Flags::Parse(args, flags));
+    EXPECT_TRUE(args.empty());
+  };
+  absl::Cleanup restore_flag = [&] { set_vmm_disabled(was_disabled); };
+  set_vmm_disabled(true);
+
+  // Use a fresh executor: cached executors retain the options from Init().
+  ASSERT_OK_AND_ASSIGN(Platform * platform,
+                       PlatformManager::PlatformWithName("CUDA"));
+  ASSERT_GT(platform->VisibleDeviceCount(), 0);
+  CudaExecutor executor(platform, /*device_ordinal=*/0);
+  ASSERT_OK(executor.Init());
+  EXPECT_THAT(executor.CreateMemoryReservation(4096),
+              absl_testing::StatusIs(absl::StatusCode::kUnimplemented));
+  EXPECT_THAT(executor.CreatePhysicalMemoryAllocation(4096),
+              absl_testing::StatusIs(absl::StatusCode::kUnimplemented));
+
+  // A caller falling back from reservation mode can still allocate normally.
+  DeviceAddressBase memory = executor.Allocate(4096, /*memory_space=*/0);
+  EXPECT_NE(memory.opaque(), nullptr);
+  executor.Deallocate(&memory);
+}
+
+// Reservations and physical allocations created through the executor use its
+// probed options, map together, and GetAllocationRange describes the mapping
+// backing a pointer rather than the whole reservation.
+TEST(CudaExecutorTest, ReservationMapsPhysicalMemoryAndReportsMappingRange) {
+  if (xla::GetDebugOptionsFromFlags().xla_gpu_experimental_vmm_disabled()) {
+    GTEST_SKIP() << "CUDA VMM is disabled";
+  }
+  ASSERT_OK_AND_ASSIGN(Platform * platform,
+                       PlatformManager::PlatformWithName("CUDA"));
+  ASSERT_OK_AND_ASSIGN(StreamExecutor * executor,
+                       platform->ExecutorForDevice(0));
+  auto cuda_executor = dynamic_cast<CudaExecutor*>(executor);
+  ASSERT_NE(cuda_executor, nullptr);
+
+  // A one-byte reservation is padded to exactly one granule.
+  ASSERT_OK_AND_ASSIGN(auto one_granule, executor->CreateMemoryReservation(1));
+  const size_t granularity = one_granule->granularity();
+  ASSERT_GT(granularity, 0);
+  EXPECT_EQ(one_granule->address().size(), granularity);
+
+  // Reserve two granules and back only the first with physical memory.
+  ASSERT_OK_AND_ASSIGN(auto reservation,
+                       executor->CreateMemoryReservation(2 * granularity));
+  ASSERT_EQ(reservation->address().size(), 2 * granularity);
+  ASSERT_OK_AND_ASSIGN(auto physical,
+                       executor->CreatePhysicalMemoryAllocation(granularity));
+  ASSERT_EQ(physical->address().size(), granularity);
+  ASSERT_OK_AND_ASSIGN(
+      auto mapping,
+      reservation->MapTo(/*reservation_offset=*/0, /*allocation_offset=*/0,
+                         granularity, *physical));
+  DeviceAddressBase mapped = mapping.mapped_address();
+  EXPECT_EQ(mapped.opaque(), reservation->address().opaque());
+  EXPECT_EQ(mapped.size(), granularity);
+
+  // The mapped prefix is accessible from the device.
+  std::vector<uint8_t> pattern(granularity, 0xAB);
+  ASSERT_OK(
+      executor->SynchronousMemcpyH2D(pattern.data(), granularity, &mapped));
+  std::vector<uint8_t> readback(granularity);
+  ASSERT_OK(
+      executor->SynchronousMemcpyD2H(mapped, granularity, readback.data()));
+  EXPECT_EQ(readback, pattern);
+
+  // RANGE_* attributes would report the whole two-granule reservation.
+  ASSERT_OK_AND_ASSIGN(DeviceAddressBase range,
+                       cuda_executor->GetAllocationRange(mapped.opaque()));
+  EXPECT_EQ(range.opaque(), mapped.opaque());
+  EXPECT_EQ(range.size(), granularity);
+}
+
 TEST(CudaExecutorTest, RetainVmmMemoryHandleForDefaultDeviceMemory) {
   ASSERT_OK_AND_ASSIGN(Platform * platform,
                        PlatformManager::PlatformWithName("CUDA"));
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_memory_reservation.cc b/third_party/xla/xla/stream_executor/cuda/cuda_memory_reservation.cc
index 8ddc950820a..7e9b1fc0897 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_memory_reservation.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_memory_reservation.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/stream_executor/cuda/cuda_memory_reservation.h"
 
+#include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include <memory>
@@ -46,25 +47,52 @@ CudaMemoryReservation::Create(StreamExecutor* executor, uint64_t size) {
 
   ABSL_ASSIGN_OR_RETURN(CudaDeviceAllocator::Options options,
                    QueryDeviceAllocatorOptions(device));
-  CUmemAllocationProp props = BuildVmmAllocationProp(device, options);
+  return CreateWithDevice(executor, device, size, options);
+}
 
-  size_t granularity = 0;
-  ABSL_RETURN_IF_ERROR(cuda::ToStatus(cuMemGetAllocationGranularity(
-      &granularity, &props, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)));
+absl::StatusOr<std::unique_ptr<CudaMemoryReservation>>
+CudaMemoryReservation::Create(StreamExecutor* executor, uint64_t size,
+                              const CudaDeviceAllocator::Options& options) {
+  std::unique_ptr<ActivateContext> activation = executor->Activate();
 
-  uint64_t padded_size = xla::RoundUpTo<uint64_t>(size, granularity);
+  CUdevice device;
+  ABSL_RETURN_IF_ERROR(
+      cuda::ToStatus(cuDeviceGet(&device, executor->device_ordinal())));
+  return CreateWithDevice(executor, device, size, options);
+}
+
+absl::StatusOr<std::unique_ptr<CudaMemoryReservation>>
+CudaMemoryReservation::CreateWithDevice(
+    StreamExecutor* executor, CUdevice device, uint64_t size,
+    const CudaDeviceAllocator::Options& options) {
+  if (!options.use_vmm) {
+    return absl::InvalidArgumentError(
+        "CudaMemoryReservation requires CUDA VMM, but options.use_vmm is "
+        "false");
+  }
+
+  // The granularity query itself can be rejected for unsupported handle
+  // types; probe with the same fallback CudaDeviceAllocator uses.
+  ABSL_ASSIGN_OR_RETURN(VmmGranularityProbe probe,
+                   ProbeVmmGranularity(device, options));
+
+  // Same effective alignment as CudaDeviceAllocator. The mapping granularity
+  // reported to callers stays the driver's value.
+  size_t alignment = std::max(options.alignment, probe.granularity);
+  uint64_t padded_size = xla::RoundUpTo<uint64_t>(size, alignment);
 
   CUdeviceptr ptr;
-  ABSL_RETURN_IF_ERROR(cuda::ToStatus(
-      cuMemAddressReserve(&ptr, padded_size, granularity, 0, 0)));
+  ABSL_RETURN_IF_ERROR(
+      cuda::ToStatus(cuMemAddressReserve(&ptr, padded_size, alignment, 0, 0)));
 
   return std::unique_ptr<CudaMemoryReservation>(
-      new CudaMemoryReservation(executor, ptr, padded_size));
+      new CudaMemoryReservation(executor, ptr, padded_size, probe.granularity));
 }
 
 CudaMemoryReservation::CudaMemoryReservation(StreamExecutor* executor,
-                                             CUdeviceptr ptr, uint64_t size)
-    : executor_(executor), ptr_(ptr), size_(size) {}
+                                             CUdeviceptr ptr, uint64_t size,
+                                             size_t granularity)
+    : executor_(executor), ptr_(ptr), size_(size), granularity_(granularity) {}
 
 DeviceAddressBase CudaMemoryReservation::address() const {
   return DeviceAddressBase(reinterpret_cast<void*>(ptr_), size_);
@@ -129,14 +157,9 @@ CudaMemoryReservation::~CudaMemoryReservation() {
     return;
   }
   std::unique_ptr<ActivateContext> activation = executor_->Activate();
-  // Attempt to unmap the full range before freeing the virtual address space.
-  // Sub-ranges already unmapped by ScopedMapping destructors will cause this
-  // call to fail; the error is logged and the address range is freed anyway.
-  auto unmap_status =
-      cuda::ToStatus(cuMemUnmap(ptr_, size_), "Error unmapping CUDA memory");
-  if (!unmap_status.ok()) {
-    LOG(ERROR) << unmap_status.message();
-  }
+  // ScopedMapping owns each mapped slice and must be destroyed first. The
+  // reservation may have an unmapped tail, so unmapping the full range here
+  // would be invalid.
   auto free_status = cuda::ToStatus(cuMemAddressFree(ptr_, size_),
                                     "Error freeing CUDA address range");
   if (!free_status.ok()) {
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_memory_reservation.h b/third_party/xla/xla/stream_executor/cuda/cuda_memory_reservation.h
index 69d758b8bbd..96ccccffb34 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_memory_reservation.h
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_memory_reservation.h
@@ -23,6 +23,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "third_party/gpus/cuda/include/cuda.h"
+#include "xla/stream_executor/cuda/cuda_device_allocator.h"
 #include "xla/stream_executor/device_address.h"
 #include "xla/stream_executor/memory_allocation.h"
 #include "xla/stream_executor/memory_reservation.h"
@@ -37,19 +38,40 @@ class CudaMemoryReservation : public MemoryReservation {
  public:
   // Reserves a virtual address range of at least `size` bytes using
   // cuMemAddressReserve. StreamExecutor is used only for context activation.
+  // Probes the device's allocator options on every call; prefer the overload
+  // below when the caller already holds them.
   static absl::StatusOr<std::unique_ptr<CudaMemoryReservation>> Create(
       StreamExecutor* executor, uint64_t size);
 
+  // Same, but derives the mapping granularity from `options` instead of
+  // probing the device, so it matches physical allocations created with the
+  // same options. Like CudaDeviceAllocator, the range is aligned and padded to
+  // the larger of `options.alignment` and the mapping granularity. Fails with
+  // InvalidArgument when `options.use_vmm` is false.
+  static absl::StatusOr<std::unique_ptr<CudaMemoryReservation>> Create(
+      StreamExecutor* executor, uint64_t size,
+      const CudaDeviceAllocator::Options& options);
+
   // Returns the base address and padded size of the reserved virtual range.
   DeviceAddressBase address() const override;
 
+  // The mapping granularity reported by the driver for the handle types in
+  // use; mapping offsets and sizes must be multiples of it.
+  size_t granularity() const override { return granularity_; }
+
   ~CudaMemoryReservation() override;
   CudaMemoryReservation(CudaMemoryReservation&&) = delete;
   CudaMemoryReservation& operator=(CudaMemoryReservation&&) = delete;
 
  private:
   explicit CudaMemoryReservation(StreamExecutor* executor, CUdeviceptr ptr,
-                                 uint64_t size);
+                                 uint64_t size, size_t granularity);
+
+  // Shared by both Create overloads. The caller has activated the context and
+  // resolved `device`.
+  static absl::StatusOr<std::unique_ptr<CudaMemoryReservation>>
+  CreateWithDevice(StreamExecutor* executor, CUdevice device, uint64_t size,
+                   const CudaDeviceAllocator::Options& options);
 
   // Maps [reservation_offset, reservation_offset+size) in the reservation to
   // [allocation_offset, allocation_offset+size) in allocation via cuMemMap.
@@ -57,8 +79,8 @@ class CudaMemoryReservation : public MemoryReservation {
   absl::Status Map(size_t reservation_offset, size_t allocation_offset,
                    size_t size, MemoryAllocation& allocation) override;
 
-  // Enables read/write access to the full reservation for the owning device
-  // via cuMemSetAccess.
+  // Enables read/write access to this slice for the owning device and its
+  // P2P-capable peers via cuMemSetAccess.
   absl::Status SetAccess(uint64_t reservation_offset, size_t size) override;
 
   // Unmaps [offset, offset+size) within this reservation via cuMemUnmap.
@@ -67,6 +89,7 @@ class CudaMemoryReservation : public MemoryReservation {
   StreamExecutor* executor_;
   CUdeviceptr ptr_;  // 0 means moved-from / released
   uint64_t size_;
+  size_t granularity_;
 };
 
 }  // namespace stream_executor::gpu
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_memory_reservation_test.cc b/third_party/xla/xla/stream_executor/cuda/cuda_memory_reservation_test.cc
index 67f2b27ffc2..2c490aeb17a 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_memory_reservation_test.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_memory_reservation_test.cc
@@ -25,6 +25,7 @@ limitations under the License.
 #include "absl/status/status_matchers.h"  // IWYU pragma: keep
 #include "absl/types/span.h"
 #include "third_party/gpus/cuda/include/cuda.h"
+#include "xla/stream_executor/cuda/cuda_device_allocator.h"
 #include "xla/stream_executor/cuda/cuda_platform_id.h"
 #include "xla/stream_executor/cuda/cuda_raw_memory_allocation.h"
 #include "xla/stream_executor/device_address.h"
@@ -33,9 +34,6 @@ limitations under the License.
 #include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/platform_manager.h"
 #include "xla/stream_executor/stream_executor.h"
-#include "xla/tsl/lib/core/status_test_util.h"
-#include "tsl/platform/statusor.h"
-#include "tsl/platform/test.h"
 
 namespace stream_executor::gpu {
 namespace {
@@ -85,6 +83,49 @@ TEST_F(CudaMemoryReservationTest, MapToWrongType) {
               StatusIs(absl::StatusCode::kInvalidArgument));
 }
 
+// Verifies that callers can supply allocator options instead of probing them.
+TEST_F(CudaMemoryReservationTest, CreateWithExplicitOptions) {
+  CudaDeviceAllocator::Options options;
+  options.enable_posix_fd_handle = false;
+  options.enable_fabric_handle = false;
+  ASSERT_OK_AND_ASSIGN(
+      auto res, CudaMemoryReservation::Create(executor_, kTestSize, options));
+
+  EXPECT_NE(res->address().opaque(), nullptr);
+  EXPECT_GE(res->address().size(), kTestSize);
+  EXPECT_GT(res->granularity(), 0);
+}
+
+// The reservation is VMM-only, so options that opt out of VMM are rejected
+// instead of silently issuing VMM driver calls.
+TEST_F(CudaMemoryReservationTest, DisabledVmmOptionsAreRejected) {
+  CudaDeviceAllocator::Options options;
+  options.use_vmm = false;
+  EXPECT_THAT(CudaMemoryReservation::Create(executor_, kTestSize, options),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+// As in CudaDeviceAllocator, the range is aligned and padded to the larger of
+// options.alignment and the mapping granularity, while granularity() keeps
+// reporting the driver's mapping granularity.
+TEST_F(CudaMemoryReservationTest,
+       AlignmentLargerThanGranularityAlignsReservation) {
+  ASSERT_OK_AND_ASSIGN(auto probe, CudaMemoryReservation::Create(executor_, 1));
+  const size_t granularity = probe->granularity();
+  ASSERT_GT(granularity, 0);
+
+  CudaDeviceAllocator::Options options;
+  options.alignment = 2 * granularity;
+  ASSERT_OK_AND_ASSIGN(auto res,
+                       CudaMemoryReservation::Create(executor_, 1, options));
+
+  EXPECT_EQ(res->granularity(), granularity);
+  EXPECT_EQ(res->address().size(), 2 * granularity);
+  EXPECT_EQ(
+      reinterpret_cast<uintptr_t>(res->address().opaque()) % (2 * granularity),
+      0);
+}
+
 // Verifies the full MapTo workflow. The ScopedMapping is destroyed first,
 // unmapping the reservation range, then the allocation is released.
 TEST_F(CudaMemoryReservationTest, MapToSingleAllocation) {
@@ -101,8 +142,8 @@ TEST_F(CudaMemoryReservationTest, MapToSingleAllocation) {
   EXPECT_EQ(mapping.mapped_address().opaque(), res->address().opaque());
   EXPECT_EQ(mapping.mapped_address().size(), alloc_size);
   // ScopedMapping destructor: cuMemUnmap.
-  // CudaMemoryReservation destructor: cuMemUnmap (logs error, already unmapped)
-  // + cuMemAddressFree. Allocation destructor: cuMemRelease.
+  // CudaMemoryReservation destructor: cuMemAddressFree.
+  // Allocation destructor: cuMemRelease.
 }
 
 // Verifies that ScopedMapping unmaps the range on destruction, allowing a
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_raw_memory_allocation.cc b/third_party/xla/xla/stream_executor/cuda/cuda_raw_memory_allocation.cc
index 9abe150a4a6..f87a2a8eb5c 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_raw_memory_allocation.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_raw_memory_allocation.cc
@@ -15,11 +15,13 @@ limitations under the License.
 
 #include "xla/stream_executor/cuda/cuda_raw_memory_allocation.h"
 
+#include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include <memory>
 
 #include "absl/log/log.h"
+#include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
 #include "third_party/gpus/cuda/include/cuda.h"
@@ -42,28 +44,45 @@ CudaRawMemoryAllocation::Create(StreamExecutor* executor, uint64_t size) {
 
   ABSL_ASSIGN_OR_RETURN(CudaDeviceAllocator::Options options,
                    QueryDeviceAllocatorOptions(device));
-  CUmemAllocationProp props = BuildVmmAllocationProp(device, options);
+  return CreateWithDevice(executor, device, size, options);
+}
 
-  size_t granularity = 0;
-  ABSL_RETURN_IF_ERROR(cuda::ToStatus(cuMemGetAllocationGranularity(
-      &granularity, &props, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)));
+absl::StatusOr<std::unique_ptr<CudaRawMemoryAllocation>>
+CudaRawMemoryAllocation::Create(StreamExecutor* executor, uint64_t size,
+                                const CudaDeviceAllocator::Options& options) {
+  std::unique_ptr<ActivateContext> activation = executor->Activate();
 
-  uint64_t padded_size = xla::RoundUpTo<uint64_t>(size, granularity);
+  CUdevice device;
+  ABSL_RETURN_IF_ERROR(
+      cuda::ToStatus(cuDeviceGet(&device, executor->device_ordinal())));
+  return CreateWithDevice(executor, device, size, options);
+}
 
-  CUmemGenericAllocationHandle handle;
-  CUresult err = cuMemCreate(&handle, padded_size, &props, 0);
-#if CUDA_VERSION >= 12030
-  // Strip FABRIC and retry on NOT_PERMITTED/NOT_SUPPORTED.
-  if ((err == CUDA_ERROR_NOT_PERMITTED || err == CUDA_ERROR_NOT_SUPPORTED) &&
-      (static_cast<int>(props.requestedHandleTypes) &
-       CU_MEM_HANDLE_TYPE_FABRIC)) {
-    props.requestedHandleTypes = static_cast<CUmemAllocationHandleType>(
-        static_cast<int>(props.requestedHandleTypes) &
-        ~CU_MEM_HANDLE_TYPE_FABRIC);
-    err = cuMemCreate(&handle, padded_size, &props, 0);
+absl::StatusOr<std::unique_ptr<CudaRawMemoryAllocation>>
+CudaRawMemoryAllocation::CreateWithDevice(
+    StreamExecutor* executor, CUdevice device, uint64_t size,
+    const CudaDeviceAllocator::Options& options) {
+  if (!options.use_vmm) {
+    return absl::InvalidArgumentError(
+        "CudaRawMemoryAllocation requires CUDA VMM, but options.use_vmm is "
+        "false");
   }
-#endif
-  ABSL_RETURN_IF_ERROR(cuda::ToStatus(err, "cuMemCreate"));
+
+  // The granularity query itself can be rejected for unsupported handle
+  // types; probe with the same fallback CudaDeviceAllocator uses and create
+  // the allocation with the handle types the driver accepted.
+  ABSL_ASSIGN_OR_RETURN(VmmGranularityProbe probe,
+                   ProbeVmmGranularity(device, options));
+  CUmemAllocationProp props = BuildVmmAllocationProp(device, probe.options);
+
+  // Same effective alignment as CudaDeviceAllocator.
+  size_t alignment = std::max(options.alignment, probe.granularity);
+  uint64_t padded_size = xla::RoundUpTo<uint64_t>(size, alignment);
+
+  // Shared fallback: FABRIC+POSIX_FD -> POSIX_FD -> NONE, including on
+  // CUDA_ERROR_INVALID_VALUE from older drivers.
+  ABSL_ASSIGN_OR_RETURN(CUmemGenericAllocationHandle handle,
+                   CreateVmmPhysicalAllocation(props, padded_size));
 
   return std::unique_ptr<CudaRawMemoryAllocation>(
       new CudaRawMemoryAllocation(executor, handle, padded_size));
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_raw_memory_allocation.h b/third_party/xla/xla/stream_executor/cuda/cuda_raw_memory_allocation.h
index 5eb9fa27775..00b161ff223 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_raw_memory_allocation.h
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_raw_memory_allocation.h
@@ -21,6 +21,7 @@ limitations under the License.
 
 #include "absl/status/statusor.h"
 #include "third_party/gpus/cuda/include/cuda.h"
+#include "xla/stream_executor/cuda/cuda_device_allocator.h"
 #include "xla/stream_executor/device_address.h"
 #include "xla/stream_executor/memory_allocation.h"
 #include "xla/stream_executor/stream_executor.h"
@@ -33,10 +34,23 @@ namespace stream_executor::gpu {
 class CudaRawMemoryAllocation : public MemoryAllocation {
  public:
   // Creates a physical memory allocation of at least `size` bytes using
-  // cuMemCreate. StreamExecutor is used only for context activation.
+  // cuMemCreate. StreamExecutor is used only for context activation. Probes
+  // the device's allocator options on every call; prefer the overload below
+  // when the caller already holds them.
   static absl::StatusOr<std::unique_ptr<CudaRawMemoryAllocation>> Create(
       StreamExecutor* executor, uint64_t size);
 
+  // Same, but builds the allocation properties from `options` (normally the
+  // executor's probed CudaDeviceAllocator::Options) so that handle types match
+  // the executor's other VMM allocations, and falls back through simpler
+  // handle types the same way CudaDeviceAllocator does. Like
+  // CudaDeviceAllocator, the size is padded to the larger of
+  // `options.alignment` and the mapping granularity. Fails with
+  // InvalidArgument when `options.use_vmm` is false.
+  static absl::StatusOr<std::unique_ptr<CudaRawMemoryAllocation>> Create(
+      StreamExecutor* executor, uint64_t size,
+      const CudaDeviceAllocator::Options& options);
+
   // Returns a DeviceAddressBase whose opaque() holds the raw
   // CUmemGenericAllocationHandle cast to void*, and size() is the
   // padded allocation size.
@@ -53,6 +67,12 @@ class CudaRawMemoryAllocation : public MemoryAllocation {
                                    CUmemGenericAllocationHandle handle,
                                    uint64_t size);
 
+  // Shared by both Create overloads. The caller has activated the context and
+  // resolved `device`.
+  static absl::StatusOr<std::unique_ptr<CudaRawMemoryAllocation>>
+  CreateWithDevice(StreamExecutor* executor, CUdevice device, uint64_t size,
+                   const CudaDeviceAllocator::Options& options);
+
   StreamExecutor* executor_;
   CUmemGenericAllocationHandle handle_;  // 0 means moved-from / released
   uint64_t size_;
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_raw_memory_allocation_test.cc b/third_party/xla/xla/stream_executor/cuda/cuda_raw_memory_allocation_test.cc
index 27ebc668e78..e09dfef2fe4 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_raw_memory_allocation_test.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_raw_memory_allocation_test.cc
@@ -20,20 +20,19 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/status/status.h"
 #include "absl/status/status_matchers.h"  // IWYU pragma: keep
 #include "third_party/gpus/cuda/include/cuda.h"
+#include "xla/stream_executor/cuda/cuda_device_allocator.h"
 #include "xla/stream_executor/cuda/cuda_platform_id.h"
 #include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/platform_manager.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/lib/core/status_test_util.h"
-#include "tsl/platform/statusor.h"
-#include "tsl/platform/test.h"
 
 namespace stream_executor::gpu {
 namespace {
 
-
 // 1 MB — will be rounded up to the VMM granularity (typically 2 MB).
 static constexpr uint64_t kTestSize = 1024 * 1024;
 
@@ -70,6 +69,32 @@ TEST_F(CudaRawMemoryAllocationTest, AddressReflectsHandle) {
   EXPECT_GE(alloc->address().size(), kTestSize);
 }
 
+// Verifies that callers can supply allocator options instead of probing them.
+TEST_F(CudaRawMemoryAllocationTest, CreateWithExplicitOptions) {
+  CudaDeviceAllocator::Options options;
+  options.enable_posix_fd_handle = false;
+  options.enable_fabric_handle = false;
+  ASSERT_OK_AND_ASSIGN(auto alloc, CudaRawMemoryAllocation::Create(
+                                       executor_, kTestSize, options));
+
+  EXPECT_NE(alloc->GetHandle(), 0u);
+  EXPECT_GE(alloc->address().size(), kTestSize);
+}
+
+// Requesting exportable handle types must succeed on every machine: either
+// the driver supports them, or Create falls back to simpler handle types.
+TEST_F(CudaRawMemoryAllocationTest,
+       ExportableHandleTypesFallBackWhenUnsupported) {
+  CudaDeviceAllocator::Options options;
+  options.enable_posix_fd_handle = true;
+  options.enable_fabric_handle = true;
+  ASSERT_OK_AND_ASSIGN(auto alloc, CudaRawMemoryAllocation::Create(
+                                       executor_, kTestSize, options));
+
+  EXPECT_NE(alloc->GetHandle(), 0u);
+  EXPECT_GE(alloc->address().size(), kTestSize);
+}
+
 // Verifies that a very small request is still satisfied (padded to
 // granularity).
 TEST_F(CudaRawMemoryAllocationTest, SizeIsAtLeastRequested) {
@@ -80,5 +105,30 @@ TEST_F(CudaRawMemoryAllocationTest, SizeIsAtLeastRequested) {
   EXPECT_GE(alloc->address().size(), 1u);
 }
 
+// The allocation is VMM-only, so options that opt out of VMM are rejected
+// instead of silently issuing VMM driver calls.
+TEST_F(CudaRawMemoryAllocationTest, DisabledVmmOptionsAreRejected) {
+  CudaDeviceAllocator::Options options;
+  options.use_vmm = false;
+  EXPECT_THAT(CudaRawMemoryAllocation::Create(executor_, kTestSize, options),
+              absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+// As in CudaDeviceAllocator, the size is padded to the larger of
+// options.alignment and the mapping granularity.
+TEST_F(CudaRawMemoryAllocationTest, AlignmentLargerThanGranularityPadsSize) {
+  // A one-byte request with probed options is padded to exactly one granule.
+  ASSERT_OK_AND_ASSIGN(auto one_granule,
+                       CudaRawMemoryAllocation::Create(executor_, 1));
+  const uint64_t granularity = one_granule->address().size();
+
+  CudaDeviceAllocator::Options options;
+  options.alignment = 2 * granularity;
+  ASSERT_OK_AND_ASSIGN(auto alloc,
+                       CudaRawMemoryAllocation::Create(executor_, 1, options));
+
+  EXPECT_EQ(alloc->address().size(), 2 * granularity);
+}
+
 }  // namespace
 }  // namespace stream_executor::gpu
diff --git a/third_party/xla/xla/stream_executor/device_address_vmm_allocator_test.cc b/third_party/xla/xla/stream_executor/device_address_vmm_allocator_test.cc
index 9dac5e156eb..f2d58fb854d 100644
--- a/third_party/xla/xla/stream_executor/device_address_vmm_allocator_test.cc
+++ b/third_party/xla/xla/stream_executor/device_address_vmm_allocator_test.cc
@@ -75,6 +75,9 @@ class TestMemoryReservation final : public MemoryReservation {
     return DeviceAddressBase(storage_.get(), size_);
   }
 
+  // Host storage has no mapping granularity.
+  size_t granularity() const override { return 1; }
+
   int active_mapping_count() const { return active_mapping_count_; }
   int mapping_count() const { return mapping_count_; }
 
diff --git a/third_party/xla/xla/stream_executor/memory_reservation.h b/third_party/xla/xla/stream_executor/memory_reservation.h
index 4da3d4b1800..401b2bb5e29 100644
--- a/third_party/xla/xla/stream_executor/memory_reservation.h
+++ b/third_party/xla/xla/stream_executor/memory_reservation.h
@@ -46,6 +46,11 @@ class MemoryReservation {
   // accessible until physical memory is mapped via MapTo.
   virtual DeviceAddressBase address() const = 0;
 
+  // Alignment required for mapping offsets and sizes. Backends may report a
+  // recommended granularity larger than their hardware minimum. Pure virtual
+  // so that a backend cannot silently inherit a wrong value.
+  virtual size_t granularity() const = 0;
+
   // Describes a mapping from a memory reservation range
   // [reservation_offset, reservation_offset + size) to a physical allocation
   // range [allocation_offset, allocation_offset + size).
@@ -70,6 +75,7 @@ class MemoryReservation {
   // An RAII wrapper that gives access to a contiguous slice of a memory
   // reservation backed by one or more physical memory allocations.
   // Unmaps the mapped range from the reservation on destruction.
+  // Must be destroyed before the reservation and its backing allocations.
   class ScopedMapping {
    public:
     ScopedMapping() = default;
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_memory_reservation.cc b/third_party/xla/xla/stream_executor/rocm/rocm_memory_reservation.cc
index 61134bf4a13..bb9004c574a 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_memory_reservation.cc
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_memory_reservation.cc
@@ -59,12 +59,13 @@ RocmMemoryReservation::Create(StreamExecutor* executor, uint64_t size) {
       hipMemAddressReserve(&ptr, padded_size, granularity, nullptr, 0ULL)));
 
   return std::unique_ptr<RocmMemoryReservation>(new RocmMemoryReservation(
-      executor, static_cast<char*>(ptr), padded_size));
+      executor, static_cast<char*>(ptr), padded_size, granularity));
 }
 
 RocmMemoryReservation::RocmMemoryReservation(StreamExecutor* executor,
-                                             char* ptr, uint64_t size)
-    : executor_(executor), ptr_(ptr), size_(size) {}
+                                             char* ptr, uint64_t size,
+                                             size_t granularity)
+    : executor_(executor), ptr_(ptr), size_(size), granularity_(granularity) {}
 
 DeviceAddressBase RocmMemoryReservation::address() const {
   return DeviceAddressBase(ptr_, size_);
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_memory_reservation.h b/third_party/xla/xla/stream_executor/rocm/rocm_memory_reservation.h
index 700f3e8c022..6956b29ed76 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_memory_reservation.h
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_memory_reservation.h
@@ -43,6 +43,10 @@ class RocmMemoryReservation : public MemoryReservation {
   // Returns the base address and padded size of the reserved virtual range.
   DeviceAddressBase address() const override;
 
+  // The mapping granularity reported by hipMemGetAllocationGranularity;
+  // mapping offsets and sizes must be multiples of it.
+  size_t granularity() const override { return granularity_; }
+
   ~RocmMemoryReservation() override;
 
   // Non-movable: the base class MemoryReservation already deletes its move
@@ -55,7 +59,7 @@ class RocmMemoryReservation : public MemoryReservation {
 
  private:
   explicit RocmMemoryReservation(StreamExecutor* executor, char* ptr,
-                                 uint64_t size);
+                                 uint64_t size, size_t granularity);
 
   absl::Status Map(size_t reservation_offset, size_t allocation_offset,
                    size_t size, MemoryAllocation& allocation) override;
@@ -70,6 +74,7 @@ class RocmMemoryReservation : public MemoryReservation {
   // null check is purely defensive.
   char* ptr_;
   uint64_t size_;
+  size_t granularity_;
   // Bytes currently mapped into the reservation. Kept in sync by Map()/UnMap()
   // so the destructor can skip a redundant hipMemUnmap when the ScopedMapping
   // that owns the mapping has already unmapped the range (which would otherwise
diff --git a/third_party/xla/xla/stream_executor/stream_executor.h b/third_party/xla/xla/stream_executor/stream_executor.h
index 3389e7574cd..2d025b2a49a 100644
--- a/third_party/xla/xla/stream_executor/stream_executor.h
+++ b/third_party/xla/xla/stream_executor/stream_executor.h
@@ -44,6 +44,7 @@ limitations under the License.
 #include "xla/stream_executor/kernel_spec.h"
 #include "xla/stream_executor/memory_allocation.h"
 #include "xla/stream_executor/memory_allocator.h"
+#include "xla/stream_executor/memory_reservation.h"
 #include "xla/stream_executor/memory_space.h"
 #include "xla/stream_executor/module_spec.h"
 #include "xla/stream_executor/platform.h"
@@ -119,6 +120,21 @@ class StreamExecutor {
     return absl::UnimplementedError("Not Implemented");
   }
 
+  // Reserves device virtual address space without allocating physical memory.
+  // Backends implementing this must also support
+  // CreatePhysicalMemoryAllocation.
+  virtual absl::StatusOr<std::unique_ptr<MemoryReservation>>
+  CreateMemoryReservation(uint64_t size) {
+    return absl::UnimplementedError("Device VA reservation is not supported");
+  }
+
+  // Allocates physical device memory without mapping it to a virtual address.
+  // The returned size may be rounded up to the backend's mapping granularity.
+  virtual absl::StatusOr<std::unique_ptr<MemoryAllocation>>
+  CreatePhysicalMemoryAllocation(uint64_t size) {
+    return absl::UnimplementedError("Physical allocation is not supported");
+  }
+
   // Obtains metadata about the underlying device.
   // The value is cached on first use.
   virtual const DeviceDescription& GetDeviceDescription() const = 0;

From 01a9a9489af585694d84e9d27f26ee7a8c093020 Mon Sep 17 00:00:00 2001
From: bodapatisaikrishna <saikrishnabodapati@gmail.com>
Date: Thu, 8 Oct 2026 17:42:46 +0530
Subject: [PATCH 0272/1000] Generalize rank number in conv2d test regex

Update regex to match any invalid rank in graph-mode shape inference.

Signed-off-by: bodapatisaikrishna <saikrishnabodapati@gmail.com>
---
 tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
index e8b52976588..391c24c4ce2 100644
--- a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
@@ -3137,7 +3137,7 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
     # Invalid input or filter rank (< 4) in eager/graph execution.
     with self.assertRaisesRegex(
         (ValueError, errors_impl.InvalidArgumentError),
-        r"(must be 4-dimensional|Shape must be rank 4 but is rank 3)",
+        r"(must be 4-dimensional|Shape must be rank 4 but is rank \d+)",
     ):
       self.evaluate(
           nn_ops.conv2d(
@@ -3149,7 +3149,7 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
       )
     with self.assertRaisesRegex(
         (ValueError, errors_impl.InvalidArgumentError),
-        r"(must be 4-dimensional|Shape must be rank 4 but is rank 3)",
+        r"(must be 4-dimensional|Shape must be rank 4 but is rank \d+)",
     ):
       self.evaluate(
           nn_ops.conv2d(
@@ -3163,7 +3163,7 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
     # Ensure raw ops are protected by the C++ kernel boundary.
     with self.assertRaisesRegex(
         (ValueError, errors_impl.InvalidArgumentError),
-        r"(must be 4-dimensional|Shape must be rank 4 but is rank 3)",
+        r"(must be 4-dimensional|Shape must be rank 4 but is rank \d+)",
     ):
       self.evaluate(
           raw_ops.Conv2D(
@@ -3175,7 +3175,7 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
       )
     with self.assertRaisesRegex(
         (ValueError, errors_impl.InvalidArgumentError),
-        r"(must be 4-dimensional|Shape must be rank 4 but is rank 3)",
+        r"(must be 4-dimensional|Shape must be rank 4 but is rank \d+)",
     ):
       self.evaluate(
           raw_ops.Conv2D(

From 5b07fae8c21a5b9e68635d35730b87f0313d4aab Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Thu, 8 Oct 2026 03:01:43 -0700
Subject: [PATCH 0273/1000] PR #128550: Bump urllib3 from 2.7.0 to 2.8.0
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/tensorflow/tensorflow/pull/128550

Bumps [urllib3](https://github.com/urllib3/urllib3) from 2.7.0 to 2.8.0.
<details>
<summary>Release notes</summary>
<p><em>Sourced from <a href="https://github.com/urllib3/urllib3/releases">urllib3's releases</a>.</em></p>
<blockquote>
<h2>2.8.0</h2>
<h2>🚀 urllib3 is fundraising for HTTP/2 support</h2>
<p><a href="https://sethmlarson.dev/urllib3-is-fundraising-for-http2-support">urllib3 is raising ~$40,000 USD</a> to release HTTP/2 support and ensure long-term sustainable maintenance of the project. If your company or organization uses Python and would benefit from HTTP/2 support in Requests, pip, cloud SDKs, and thousands of other projects <a href="https://opencollective.com/urllib3">please consider contributing financially</a> to ensure HTTP/2 support is developed sustainably and maintained for the long-haul.</p>
<p>Thank you for your support.</p>
<h2>Security</h2>
<p>Fixed the following security issues:</p>
<ul>
<li>The TLS configuration for HTTPS proxies could be ignored or overridden. (High severity, GHSA-8988-9cw3-xx77)</li>
<li><code>HTTPResponse.stream()</code> and <code>read_chunked()</code> could buffer a chunk-size line of unbounded length in memory. (High severity, GHSA-vxq7-64xx-v4gw)</li>
<li>Chunked Deflate streaming could enter an infinite loop. (Medium severity, GHSA-gh4c-6fx4-qh6g)</li>
</ul>
<blockquote>
<p>[!IMPORTANT]
urllib3 2.8.0 fixes HTTPS proxy TLS configuration being ignored or overridden by destination settings. Configurations relying on that behavior may require changes.</p>
<p>Configure proxy CA certificates and client certificates in <code>proxy_ssl_context</code>, and proxy identity checks with <code>proxy_assert_hostname</code> or <code>proxy_assert_fingerprint</code>. Destination client certificates and identity overrides no longer apply to HTTPS forwarding proxy connections.</p>
</blockquote>
<blockquote>
<p>[!NOTE]
CVE IDs had not yet been assigned to these advisories at the time of release due to a backlog at GitHub's CNA.</p>
</blockquote>
<h2>Deprecations &amp; Removals</h2>
<ul>
<li>Deprecated using an empty collection as the <code>Retry</code> option <code>allowed_methods</code> to retry any verb. (<a href="https://redirect.github.com/urllib3/urllib3/issues/5044">#5044</a>)</li>
</ul>
<h2>Features</h2>
<ul>
<li>Added <code>Url.auth_decoded</code> and <code>Url.auth_decoded_joined</code> convenience properties to the result of <code>parse_url()</code>. (<a href="https://redirect.github.com/urllib3/urllib3/issues/4945">#4945</a>)</li>
<li>Added <code>basic_auth_encoding</code> and <code>proxy_basic_auth_encoding</code> parameters to <code>urllib3.util.make_headers()</code>. (<a href="https://redirect.github.com/urllib3/urllib3/issues/5092">#5092</a>)</li>
</ul>
<h2>Bugfixes</h2>
<ul>
<li>
<p>Fixed response header handling to replace obsolete folded header lines (<code>obs-fold</code>) with spaces in accordance with RFC 9112, preventing raw CRLF sequences from appearing in header values such as <code>Set-Cookie</code>. (<a href="https://redirect.github.com/urllib3/urllib3/issues/1362">#1362</a>)</p>
</li>
<li>
<p>Fixed usage of <code>proxy_ssl_context</code> with <code>ProxyManager</code> when <code>use_forwarding_for_https=True</code>. Passing <code>ssl_context</code> instead of <code>proxy_ssl_context</code> for HTTPS proxies in this configuration now emits a <code>FutureWarning</code> and will raise an error in v3.0. (<a href="https://redirect.github.com/urllib3/urllib3/issues/2577">#2577</a>)</p>
</li>
<li>
<p>Changed behavior of the default <code>ConnectionPool.pool</code> initialization. <code>LifoQueue</code> is now resolved from the <code>queue</code> module after the <code>ConnectionPool</code> is instantiated instead of using the default cached <code>QueueCls</code> class property. This is done because sometimes the <code>queue.LifoQueue</code> is monkey-patched late in the program, such as by gevent. (<a href="https://redirect.github.com/urllib3/urllib3/issues/3289">#3289</a>)</p>
</li>
<li>
<p>Raised <code>UnrewindableBodyError</code> instead of <code>ValueError</code> when retrying a request whose body had <code>tell()</code> but not <code>seek()</code>. (<a href="https://redirect.github.com/urllib3/urllib3/issues/3779">#3779</a>)</p>
</li>
<li>
<p>Decoded percent-encoded SOCKS proxy credentials before authenticating with the proxy server. (<a href="https://redirect.github.com/urllib3/urllib3/issues/3785">#3785</a>)</p>
</li>
<li>
<p>Fixed <code>HTTPResponse.drain_conn()</code> to discard unread response data in 64 KiB chunks (same as the default <code>amt</code> when doing <code>HTTPResponse.stream(...)</code>). (<a href="https://redirect.github.com/urllib3/urllib3/issues/5019">#5019</a>)</p>
</li>
<li>
<p>Fixed <code>is_ipaddress()</code> to detect non-standard IPv4 forms accepted by <code>socket.connect</code>, such as hex (<code>0x7f000001</code>), octal (<code>0177.0.0.1</code>), and decimal integers (<code>2130706433</code>), ensuring SSL certificate verification uses the correct mode for these addresses. (<a href="https://redirect.github.com/urllib3/urllib3/issues/5029">#5029</a>)</p>
</li>
<li>
<p>Fixed <code>HTTPConnectionPool.urlopen</code> raising a misleading <code>FullPoolError</code> instead of <code>ValueError</code> when called with an invalid <code>timeout</code> argument on a pool created with <code>block=True</code>. (<a href="https://redirect.github.com/urllib3/urllib3/issues/5059">#5059</a>)</p>
</li>
<li>
<p>Fixed port-zero handling to preserve explicit <code>:0</code> values instead of substituting the default ports 80 or 443 in URL parsing, pool selection, proxy configuration, <code>connection_from_url()</code>, and HTTP/2 request authority. (<a href="https://redirect.github.com/urllib3/urllib3/issues/5071">#5071</a>, <a href="https://redirect.github.com/urllib3/urllib3/issues/5101">#5101</a>)</p>
</li>
<li>
<p>Fixed a bug where <code>PoolManager</code> passed the <code>assert_hostname</code> and <code>assert_fingerprint</code> parameters to HTTP connection pools. (<a href="https://redirect.github.com/urllib3/urllib3/issues/5077">#5077</a>)</p>
</li>
<li>
<p>Fixed <code>HTTPConnectionPool.urlopen()</code> and HTTP proxy forwarding to strip URL fragments from absolute request targets before sending requests. (<a href="https://redirect.github.com/urllib3/urllib3/issues/5079">#5079</a>)</p>
</li>
<li>
<p>Added safeguards to the proxy tunneling code to prevent potential security issues when handling invalid characters in the proxy host and HTTP headers. This change affects users of Python 3.10, Python 3.11, and Python 3.12 when the standard library does not contain the fix; those on newer Python versions should upgrade to 3.13.14+ or 3.14.5+ to get the same security fixes. (<a href="https://redirect.github.com/urllib3/urllib3/issues/5091">#5091</a>)</p>
</li>
<li>
<p>Fixed <code>HTTPSConnection.connect()</code> overriding <code>ProxyConfig.ssl_context</code>'s certificate policy and proxy identity checks with the target connection's TLS settings when forwarding through an HTTPS proxy.</p>
<p><code>HTTPSConnection</code> no longer applies target SNI, assertions, or client credentials to forwarding proxy handshakes and continues to use its <code>ssl_context</code> as a fallback when an HTTPS proxy forwards an HTTP target. (<a href="https://redirect.github.com/urllib3/urllib3/issues/5093">#5093</a>)</p>
</li>
<li>
<p>Fixed URL parsing to more strictly enforce RFC 3986 host syntax, rejecting invalid host input such as raw spaces and control characters, malformed percent-encodings, and percent-encoded control characters in HTTP(S) hosts and IPv6 zone identifiers, including proxy CONNECT tunnel targets. Host normalization now also follows RFC 3986 normalization rules for percent-encoded octets by decoding percent-encoded unreserved characters and uppercasing the hexadecimal digits of retained percent-encoded octets. (<a href="https://redirect.github.com/urllib3/urllib3/issues/5095">#5095</a>)</p>
</li>
</ul>
<!-- raw HTML omitted -->
</blockquote>
<p>... (truncated)</p>
</details>
<details>
<summary>Changelog</summary>
<p><em>Sourced from <a href="https://github.com/urllib3/urllib3/blob/main/CHANGES.rst">urllib3's changelog</a>.</em></p>
<blockquote>
<h1>2.8.0 (2026-09-15)</h1>
<h2>Security</h2>
<p>Fixed the following security issues:</p>
<ul>
<li>The TLS configuration for HTTPS proxies could be ignored or overridden.
(High severity, <code>GHSA-8988-9cw3-xx77 &lt;https://github.com/urllib3/urllib3/security/advisories/GHSA-8988-9cw3-xx77&gt;</code>__)</li>
<li><code>HTTPResponse.stream()</code> and <code>read_chunked()</code> could buffer a chunk-size
line of unbounded length in memory. (High severity,
<code>GHSA-vxq7-64xx-v4gw &lt;https://github.com/urllib3/urllib3/security/advisories/GHSA-vxq7-64xx-v4gw&gt;</code>__)</li>
<li>Chunked Deflate streaming could enter an infinite loop. (Medium severity,
<code>GHSA-gh4c-6fx4-qh6g &lt;https://github.com/urllib3/urllib3/security/advisories/GHSA-gh4c-6fx4-qh6g&gt;</code>__)</li>
</ul>
<p>.. caution::</p>
<pre><code>urllib3 2.8.0 fixes HTTPS proxy TLS configuration being ignored or
overridden by destination settings. Configurations relying on that
behavior may require changes.
<p>Configure proxy CA certificates and client certificates in
<code>proxy_ssl_context</code>, and proxy identity checks with
<code>proxy_assert_hostname</code> or <code>proxy_assert_fingerprint</code>.
Destination client certificates and identity overrides no longer
apply to HTTPS forwarding proxy connections.
</code></pre></p>
<h2>Deprecations &amp; Removals</h2>
<ul>
<li>Deprecated using an empty collection as the <code>Retry</code> option
<code>allowed_methods</code> to retry any verb.
(<code>[#5044](https://github.com/urllib3/urllib3/issues/5044) &lt;https://github.com/urllib3/urllib3/issues/5044&gt;</code>__)</li>
</ul>
<h2>Features</h2>
<ul>
<li>Added <code>Url.auth_decoded</code> and <code>Url.auth_decoded_joined</code> convenience
properties to the result of <code>parse_url()</code>.
(<code>[#4945](https://github.com/urllib3/urllib3/issues/4945) &lt;https://github.com/urllib3/urllib3/issues/4945&gt;</code>__)</li>
<li>Added <code>basic_auth_encoding</code> and <code>proxy_basic_auth_encoding</code> parameters to
<code>urllib3.util.make_headers()</code>.
(<code>[#5092](https://github.com/urllib3/urllib3/issues/5092) &lt;https://github.com/urllib3/urllib3/issues/5092&gt;</code>__)</li>
</ul>
<h2>Bugfixes</h2>
<!-- raw HTML omitted -->
</blockquote>
<p>... (truncated)</p>
</details>
<details>
<summary>Commits</summary>
<ul>
<li><a href="https://github.com/urllib3/urllib3/commit/b1d30ab61fe0db8f11092805e8c5ac43e091064a"><code>b1d30ab</code></a> Release 2.8.0</li>
<li><a href="https://github.com/urllib3/urllib3/commit/9016d7e8afc68185496ef07f3c3a4a743d04922e"><code>9016d7e</code></a> Skip <code>test_read_chunked_with_trailing_data_does_not_hang</code> for brotlicffi (<a href="https://redirect.github.com/urllib3/urllib3/issues/5258">#5258</a>)</li>
<li><a href="https://github.com/urllib3/urllib3/commit/9101f581a8b3659af23b6ff335ae77200ca33533"><code>9101f58</code></a> Fix <code>nox -s docs</code> warning (<a href="https://redirect.github.com/urllib3/urllib3/issues/5256">#5256</a>)</li>
<li><a href="https://github.com/urllib3/urllib3/commit/cd770b059b543be29298ea5c52afb0b1b090f5ed"><code>cd770b0</code></a> Merge commit from fork</li>
<li><a href="https://github.com/urllib3/urllib3/commit/ea2ad7b21a80da3632f80016526a18864586077f"><code>ea2ad7b</code></a> Merge commit from fork</li>
<li><a href="https://github.com/urllib3/urllib3/commit/0716e31534345dc1599ea95d903c79f276239bd8"><code>0716e31</code></a> Fix loading unencrypted client keys with a password in pyOpenSSL (<a href="https://redirect.github.com/urllib3/urllib3/issues/5255">#5255</a>)</li>
<li><a href="https://github.com/urllib3/urllib3/commit/43c68c8b43a9dcb44ed2cf4ec91384ca0d46b37d"><code>43c68c8</code></a> Test pickling of <code>InvalidChunkLength</code> (<a href="https://redirect.github.com/urllib3/urllib3/issues/5247">#5247</a>)</li>
<li><a href="https://github.com/urllib3/urllib3/commit/308b279b3fb28e7bee952e152ec5baeb5bfd0817"><code>308b279</code></a> Share security policy between GitHub and Read the Docs (<a href="https://redirect.github.com/urllib3/urllib3/issues/5253">#5253</a>)</li>
<li><a href="https://github.com/urllib3/urllib3/commit/53fa0731b27d4b71ab0755ea5b896422d005d706"><code>53fa073</code></a> Add policy on duplicate pull requests (<a href="https://redirect.github.com/urllib3/urllib3/issues/5252">#5252</a>)</li>
<li><a href="https://github.com/urllib3/urllib3/commit/5f2a6a843d0100d1351c3f94d58581ca98d17267"><code>5f2a6a8</code></a> Assert on the ALPN extension in test_tunnel_sets_http_11_alpn (<a href="https://redirect.github.com/urllib3/urllib3/issues/5232">#5232</a>)</li>
<li>Additional commits viewable in <a href="https://github.com/urllib3/urllib3/compare/2.7.0...2.8.0">compare view</a></li>
</ul>
</details>
<br />

[![Dependabot compatibility score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=urllib3&package-manager=pip&previous-version=2.7.0&new-version=2.8.0)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores)

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`.

[//]: # (dependabot-automerge-start)
[//]: # (dependabot-automerge-end)

---

<details>
<summary>Dependabot commands and options</summary>
<br />

You can trigger Dependabot actions by commenting on this PR:
- `@dependabot rebase` will rebase this PR
- `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it
- `@dependabot show <dependency name> ignore conditions` will show all of the ignore conditions of the specified dependency
- `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
- `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
- `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
You can disable automated security fix PRs for this repo from the [Security Alerts page](https://github.com/tensorflow/tensorflow/network/alerts).

</details>
Copybara import of the project:

--
072db6828f13d32c735a3daa73868d0528182552 by dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>:

Bump urllib3 from 2.7.0 to 2.8.0

Bumps [urllib3](https://github.com/urllib3/urllib3) from 2.7.0 to 2.8.0.
- [Release notes](https://github.com/urllib3/urllib3/releases)
- [Changelog](https://github.com/urllib3/urllib3/blob/main/CHANGES.rst)
- [Commits](https://github.com/urllib3/urllib3/compare/2.7.0...2.8.0)

---
updated-dependencies:
- dependency-name: urllib3
  dependency-version: 2.8.0
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>

Merging this change closes #128550

COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/tensorflow/pull/128550 from tensorflow:dependabot/pip/urllib3-2.8.0 072db6828f13d32c735a3daa73868d0528182552
PiperOrigin-RevId: 995706065
---
 requirements_lock_3_10.txt              | 6 +++---
 requirements_lock_3_11.txt              | 6 +++---
 requirements_lock_3_12.txt              | 6 +++---
 requirements_lock_3_13.txt              | 6 +++---
 requirements_lock_3_14.txt              | 6 +++---
 requirements_lock_3_14_freethreaded.txt | 6 +++---
 6 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/requirements_lock_3_10.txt b/requirements_lock_3_10.txt
index 2a8d1904b4b..1bf00761d2e 100644
--- a/requirements_lock_3_10.txt
+++ b/requirements_lock_3_10.txt
@@ -888,9 +888,9 @@ typing-extensions==4.14.1 \
     #   -r ci/official/requirements_updater/requirements.in
     #   optree
     #   rich
-urllib3==2.7.0 \
-    --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \
-    --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897
+urllib3==2.8.0 \
+    --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
+    --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
 werkzeug==3.1.6 \
     --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
diff --git a/requirements_lock_3_11.txt b/requirements_lock_3_11.txt
index 2c69bf880a3..4b69205ebea 100644
--- a/requirements_lock_3_11.txt
+++ b/requirements_lock_3_11.txt
@@ -887,9 +887,9 @@ typing-extensions==4.14.1 \
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   optree
-urllib3==2.7.0 \
-    --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \
-    --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897
+urllib3==2.8.0 \
+    --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
+    --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
 werkzeug==3.1.6 \
     --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
diff --git a/requirements_lock_3_12.txt b/requirements_lock_3_12.txt
index d4962528b93..9dac2d0b1e9 100644
--- a/requirements_lock_3_12.txt
+++ b/requirements_lock_3_12.txt
@@ -887,9 +887,9 @@ typing-extensions==4.14.1 \
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   optree
-urllib3==2.7.0 \
-    --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \
-    --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897
+urllib3==2.8.0 \
+    --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
+    --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
 werkzeug==3.1.6 \
     --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
diff --git a/requirements_lock_3_13.txt b/requirements_lock_3_13.txt
index 38fbbb1b5de..8a894fa529c 100644
--- a/requirements_lock_3_13.txt
+++ b/requirements_lock_3_13.txt
@@ -894,9 +894,9 @@ typing-extensions==4.14.1 \
     # via
     #   -r ci/official/requirements_updater/requirements.in
     #   optree
-urllib3==2.7.0 \
-    --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \
-    --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897
+urllib3==2.8.0 \
+    --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
+    --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
 werkzeug==3.1.6 \
     --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
diff --git a/requirements_lock_3_14.txt b/requirements_lock_3_14.txt
index 1ea7a786d01..d9abfa26860 100644
--- a/requirements_lock_3_14.txt
+++ b/requirements_lock_3_14.txt
@@ -1128,9 +1128,9 @@ typing-extensions==4.14.1 \
     #   -r ci/official/requirements_updater/requirements.in
     #   grpcio
     #   optree
-urllib3==2.7.0 \
-    --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \
-    --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897
+urllib3==2.8.0 \
+    --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
+    --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
 werkzeug==3.1.8 \
     --hash=sha256:63a77fb8892bf28ebc3178683445222aa500e48ebad5ec77b0ad80f8726b1f50 \
diff --git a/requirements_lock_3_14_freethreaded.txt b/requirements_lock_3_14_freethreaded.txt
index 9998f21119e..e296d66b34b 100644
--- a/requirements_lock_3_14_freethreaded.txt
+++ b/requirements_lock_3_14_freethreaded.txt
@@ -1129,9 +1129,9 @@ typing-extensions==4.14.1 \
     #   -r ci/official/requirements_updater/requirements.in
     #   grpcio
     #   optree
-urllib3==2.7.0 \
-    --hash=sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c \
-    --hash=sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897
+urllib3==2.8.0 \
+    --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
+    --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
 werkzeug==3.1.8 \
     --hash=sha256:63a77fb8892bf28ebc3178683445222aa500e48ebad5ec77b0ad80f8726b1f50 \

From 4d3bd33013940b2afef7ccf0bea201437c273cfa Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Thu, 8 Oct 2026 12:21:02 +0000
Subject: [PATCH 0274/1000] Bump werkzeug from 3.1.6 to 3.1.9

Bumps [werkzeug](https://github.com/pallets/werkzeug) from 3.1.6 to 3.1.9.
- [Release notes](https://github.com/pallets/werkzeug/releases)
- [Changelog](https://github.com/pallets/werkzeug/blob/main/CHANGES.rst)
- [Commits](https://github.com/pallets/werkzeug/compare/3.1.6...3.1.9)

---
updated-dependencies:
- dependency-name: werkzeug
  dependency-version: 3.1.9
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
---
 requirements_lock_3_10.txt              | 6 +++---
 requirements_lock_3_11.txt              | 6 +++---
 requirements_lock_3_12.txt              | 6 +++---
 requirements_lock_3_13.txt              | 6 +++---
 requirements_lock_3_14.txt              | 6 +++---
 requirements_lock_3_14_freethreaded.txt | 6 +++---
 6 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/requirements_lock_3_10.txt b/requirements_lock_3_10.txt
index 1bf00761d2e..25930942353 100644
--- a/requirements_lock_3_10.txt
+++ b/requirements_lock_3_10.txt
@@ -892,9 +892,9 @@ urllib3==2.8.0 \
     --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
     --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
-werkzeug==3.1.6 \
-    --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
-    --hash=sha256:7ddf3357bb9564e407607f988f683d72038551200c704012bb9a4c523d42f131
+werkzeug==3.1.9 \
+    --hash=sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060 \
+    --hash=sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab
     # via tb-nightly
 wheel==0.46.2 \
     --hash=sha256:33ae60725d69eaa249bc1982e739943c23b34b58d51f1cb6253453773aca6e65 \
diff --git a/requirements_lock_3_11.txt b/requirements_lock_3_11.txt
index 4b69205ebea..4b1a0ccd52d 100644
--- a/requirements_lock_3_11.txt
+++ b/requirements_lock_3_11.txt
@@ -891,9 +891,9 @@ urllib3==2.8.0 \
     --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
     --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
-werkzeug==3.1.6 \
-    --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
-    --hash=sha256:7ddf3357bb9564e407607f988f683d72038551200c704012bb9a4c523d42f131
+werkzeug==3.1.9 \
+    --hash=sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060 \
+    --hash=sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab
     # via tb-nightly
 wheel==0.46.2 \
     --hash=sha256:33ae60725d69eaa249bc1982e739943c23b34b58d51f1cb6253453773aca6e65 \
diff --git a/requirements_lock_3_12.txt b/requirements_lock_3_12.txt
index 9dac2d0b1e9..25f46d8a20d 100644
--- a/requirements_lock_3_12.txt
+++ b/requirements_lock_3_12.txt
@@ -891,9 +891,9 @@ urllib3==2.8.0 \
     --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
     --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
-werkzeug==3.1.6 \
-    --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
-    --hash=sha256:7ddf3357bb9564e407607f988f683d72038551200c704012bb9a4c523d42f131
+werkzeug==3.1.9 \
+    --hash=sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060 \
+    --hash=sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab
     # via tb-nightly
 wheel==0.46.2 \
     --hash=sha256:33ae60725d69eaa249bc1982e739943c23b34b58d51f1cb6253453773aca6e65 \
diff --git a/requirements_lock_3_13.txt b/requirements_lock_3_13.txt
index 8a894fa529c..50f9fc6b4e3 100644
--- a/requirements_lock_3_13.txt
+++ b/requirements_lock_3_13.txt
@@ -898,9 +898,9 @@ urllib3==2.8.0 \
     --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
     --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
-werkzeug==3.1.6 \
-    --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
-    --hash=sha256:7ddf3357bb9564e407607f988f683d72038551200c704012bb9a4c523d42f131
+werkzeug==3.1.9 \
+    --hash=sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060 \
+    --hash=sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab
     # via tb-nightly
 wheel==0.46.2 \
     --hash=sha256:33ae60725d69eaa249bc1982e739943c23b34b58d51f1cb6253453773aca6e65 \
diff --git a/requirements_lock_3_14.txt b/requirements_lock_3_14.txt
index d9abfa26860..fc4acf20879 100644
--- a/requirements_lock_3_14.txt
+++ b/requirements_lock_3_14.txt
@@ -1132,9 +1132,9 @@ urllib3==2.8.0 \
     --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
     --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
-werkzeug==3.1.8 \
-    --hash=sha256:63a77fb8892bf28ebc3178683445222aa500e48ebad5ec77b0ad80f8726b1f50 \
-    --hash=sha256:9bad61a4268dac112f1c5cd4630a56ede601b6ed420300677a869083d70a4c44
+werkzeug==3.1.9 \
+    --hash=sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060 \
+    --hash=sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab
     # via tb-nightly
 wheel==0.46.2 \
     --hash=sha256:33ae60725d69eaa249bc1982e739943c23b34b58d51f1cb6253453773aca6e65 \
diff --git a/requirements_lock_3_14_freethreaded.txt b/requirements_lock_3_14_freethreaded.txt
index e296d66b34b..95f5f7d0a95 100644
--- a/requirements_lock_3_14_freethreaded.txt
+++ b/requirements_lock_3_14_freethreaded.txt
@@ -1133,9 +1133,9 @@ urllib3==2.8.0 \
     --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
     --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
-werkzeug==3.1.8 \
-    --hash=sha256:63a77fb8892bf28ebc3178683445222aa500e48ebad5ec77b0ad80f8726b1f50 \
-    --hash=sha256:9bad61a4268dac112f1c5cd4630a56ede601b6ed420300677a869083d70a4c44
+werkzeug==3.1.9 \
+    --hash=sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060 \
+    --hash=sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab
     # via tb-nightly
 wheel==0.46.2 \
     --hash=sha256:33ae60725d69eaa249bc1982e739943c23b34b58d51f1cb6253453773aca6e65 \

From 93fa24b61916ab71eb9729556d4eac6d02a88b04 Mon Sep 17 00:00:00 2001
From: Som Tripathi <somtri@iastate.edu>
Date: Mon, 14 Sep 2026 19:57:11 -0500
Subject: [PATCH 0275/1000] Guard MultiDeviceIterator and Delete ops against
 invalid state

GetNextFromShard() dereferenced multi_device_buffer_ without checking that
Init() had set it, so calling it on an iterator that was never initialized
crashed with a null-pointer dereference. Add a check under the existing
lock that returns FailedPreconditionError before building the
IteratorContext.

DeleteMultiDeviceIteratorOp and DeleteIteratorOp read element 0 of the
resource handle tensor without checking it has any elements, so an empty
handle read out of bounds. Both now use HandleFromInput, already used
elsewhere in each file, which rejects an empty handle with "Empty
resource handle" instead of reading past the tensor. This does not fix
the crash reported in #107121, which happens earlier, at
GetNextFromShard.

DeleteIterator is not named in #107121, but has the same defect and the
same one-line fix, so it is included with its own test.

PR #113269 (adityasai1234) proposed a guard for the
DeleteMultiDeviceIterator site and was reviewed by dmiltr3
(CHANGES_REQUESTED, not approved). This uses HandleFromInput's own
message instead of a custom one, leaves out the dtype check since a
non-resource input is already rejected before the kernel in eager and
at graph construction, and the test needs includes and BUILD deps
beyond the one dependency in the reviewed snippet.

The scalar reads at multi_device_iterator_ops.cc:672, :726 and :731 are
a different failure class, a CHECK abort on a non-scalar input rather
than a null or out-of-bounds read, and are not changed here.

Built and ran //tensorflow/core/kernels/data:iterator_ops_test locally
with the fix applied: all 4 tests PASSED, including the three new
ones. No pre-fix build ran, so this does not by itself show the new
tests failing without the fix.

Fixes #107121
---
 tensorflow/core/kernels/data/BUILD            |   4 +
 tensorflow/core/kernels/data/iterator_ops.cc  |   3 +-
 .../core/kernels/data/iterator_ops_test.cc    | 115 ++++++++++++++++++
 .../kernels/data/multi_device_iterator_ops.cc |  10 +-
 4 files changed, 130 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD
index 363435bc918..3b37624d3b6 100644
--- a/tensorflow/core/kernels/data/BUILD
+++ b/tensorflow/core/kernels/data/BUILD
@@ -506,6 +506,7 @@ tf_cc_test(
     srcs = ["iterator_ops_test.cc"],
     deps = [
         ":iterator_ops",
+        ":multi_device_iterator_ops",
         ":prefetch_dataset_op",
         ":range_dataset_op",
         "//tensorflow/core:core_cpu_internal",
@@ -516,6 +517,7 @@ tf_cc_test(
         "//tensorflow/core:test_main",
         "//tensorflow/core/data:dataset_test_base",
         "//tensorflow/core/framework:function_testlib",
+        "//tensorflow/core/framework:node_def_proto_cc",
         "//tensorflow/core/framework:types_proto_cc",
         "//tensorflow/core/kernels:cwise_op",
         "//tensorflow/core/kernels:function_ops",
@@ -524,6 +526,8 @@ tf_cc_test(
         "//tensorflow/core/platform:errors",
         "//tensorflow/core/platform:refcount",
         "//tensorflow/core/platform:statusor",
+        "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/status",
     ],
 )
 
diff --git a/tensorflow/core/kernels/data/iterator_ops.cc b/tensorflow/core/kernels/data/iterator_ops.cc
index f0b2b79ac07..724dd0a4f75 100644
--- a/tensorflow/core/kernels/data/iterator_ops.cc
+++ b/tensorflow/core/kernels/data/iterator_ops.cc
@@ -631,7 +631,8 @@ absl::Status MakeIteratorOp::DoCompute(OpKernelContext* ctx) {
 absl::Status DeleteIteratorOp::DoCompute(OpKernelContext* ctx) {
   tensorflow::ResourceTagger tag(kTFDataResourceTag,
                                  ctx->op_kernel().type_string());
-  const ResourceHandle& handle = ctx->input(0).flat<ResourceHandle>()(0);
+  ResourceHandle handle;
+  TF_RETURN_IF_ERROR(HandleFromInput(ctx, 0, &handle));
   // The iterator resource is guaranteed to exist because the variant tensor
   // wrapping the deleter is provided as an unused input to this op, which
   // guarantees that it has not run yet.
diff --git a/tensorflow/core/kernels/data/iterator_ops_test.cc b/tensorflow/core/kernels/data/iterator_ops_test.cc
index b4659d09e94..9d5c0ca5af1 100644
--- a/tensorflow/core/kernels/data/iterator_ops_test.cc
+++ b/tensorflow/core/kernels/data/iterator_ops_test.cc
@@ -15,16 +15,22 @@ limitations under the License.
 #include "tensorflow/core/kernels/data/iterator_ops.h"
 
 #include <cstdint>
+#include <string>
 #include <utility>
 #include <vector>
 
+#include "absl/container/inlined_vector.h"
+#include "absl/status/status.h"
 #include "tensorflow/core/common_runtime/device_mgr.h"
 #include "tensorflow/core/common_runtime/process_function_library_runtime.h"
 #include "tensorflow/core/data/dataset_test_base.h"
 #include "tensorflow/core/framework/function.h"
 #include "tensorflow/core/framework/function_testlib.h"
+#include "tensorflow/core/framework/node_def.pb.h"
+#include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
 #include "tensorflow/core/lib/monitoring/cell_reader.h"
@@ -107,6 +113,115 @@ TEST_F(IteratorOpsTest, CollectMetrics) {
   EXPECT_GT(iterator_busy.Delta(), 0.0);
 }
 
+TEST_F(IteratorOpsTest, MultiDeviceIteratorGetNextFromShardFailsBeforeInit) {
+  RangeDatasetParams dataset_params = RangeDatasetParams(0, 10, 3);
+  TF_ASSERT_OK(InitializeRuntime(dataset_params));
+  DataTypeVector output_types = {DT_INT64};
+  std::vector<PartialTensorShape> output_shapes = {PartialTensorShape({})};
+
+  Tensor multi_device_iterator;
+  {
+    // Creates the resource without running MultiDeviceIteratorInit. This
+    // context must be destroyed before the next CreateOpKernelContext call,
+    // which replaces the params it points to.
+    NodeDef node_def = test::function::NDef(
+        "multi_device_iterator", "MultiDeviceIterator", /*inputs=*/{},
+        {{"devices",
+          std::vector<std::string>{"/job:a/replica:0/task:0/device:CPU:0"}},
+         {"shared_name", "uninitialized_multi_device_iterator"},
+         {"container", ""},
+         {"output_types", output_types},
+         {"output_shapes", output_shapes}});
+    std::unique_ptr<OpKernel> kernel;
+    TF_ASSERT_OK(CreateOpKernel(node_def, &kernel));
+    absl::InlinedVector<TensorValue, 4> inputs;
+    std::unique_ptr<OpKernelContext> context;
+    TF_ASSERT_OK(CreateOpKernelContext(kernel.get(), &inputs, &context));
+    TF_ASSERT_OK(RunOpKernel(kernel.get(), context.get()));
+    multi_device_iterator = *context->mutable_output(0);
+  }
+
+  Tensor shard_num(DT_INT32, TensorShape({}));
+  shard_num.scalar<int32_t>()() = 0;
+  Tensor incarnation_id(DT_INT64, TensorShape({}));
+  incarnation_id.scalar<int64_t>()() = 0;
+
+  NodeDef node_def = test::function::NDef(
+      "get_next_from_shard", "MultiDeviceIteratorGetNextFromShard",
+      {"multi_device_iterator", "shard_num", "incarnation_id"},
+      {{"output_types", output_types}, {"output_shapes", output_shapes}});
+
+  std::unique_ptr<OpKernel> kernel;
+  TF_ASSERT_OK(CreateOpKernel(node_def, &kernel));
+
+  absl::InlinedVector<TensorValue, 4> inputs;
+  inputs.push_back(TensorValue(&multi_device_iterator));
+  inputs.push_back(TensorValue(&shard_num));
+  inputs.push_back(TensorValue(&incarnation_id));
+
+  std::unique_ptr<OpKernelContext> context;
+  TF_ASSERT_OK(CreateOpKernelContext(kernel.get(), &inputs, &context));
+
+  absl::Status status = RunOpKernel(kernel.get(), context.get());
+  EXPECT_EQ(status.code(), absl::StatusCode::kFailedPrecondition);
+  EXPECT_EQ(status.message(),
+            "GetNextFromShard() failed because the MultiDeviceIterator has not "
+            "been initialized. Ensure that you have run the initializer "
+            "operation for this MultiDeviceIterator before getting the next "
+            "element.");
+}
+
+TEST_F(IteratorOpsTest, DeleteMultiDeviceIteratorRejectsEmptyHandle) {
+  RangeDatasetParams dataset_params = RangeDatasetParams(0, 10, 3);
+  TF_ASSERT_OK(InitializeRuntime(dataset_params));
+
+  Tensor empty_resource_handle(DT_RESOURCE, TensorShape({0}));
+  Tensor variant_deleter(DT_VARIANT, TensorShape({}));
+
+  NodeDef node_def = test::function::NDef(
+      "delete_multi_device_iterator", "DeleteMultiDeviceIterator",
+      {"multi_device_iterator", "deleter"}, {{"N", 0}});
+
+  std::unique_ptr<OpKernel> kernel;
+  TF_ASSERT_OK(CreateOpKernel(node_def, &kernel));
+
+  absl::InlinedVector<TensorValue, 4> inputs;
+  inputs.push_back(TensorValue(&empty_resource_handle));
+  inputs.push_back(TensorValue(&variant_deleter));
+
+  std::unique_ptr<OpKernelContext> context;
+  TF_ASSERT_OK(CreateOpKernelContext(kernel.get(), &inputs, &context));
+
+  absl::Status status = RunOpKernel(kernel.get(), context.get());
+  EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument);
+  EXPECT_EQ(status.message(), "Empty resource handle");
+}
+
+TEST_F(IteratorOpsTest, DeleteIteratorRejectsEmptyHandle) {
+  RangeDatasetParams dataset_params = RangeDatasetParams(0, 10, 3);
+  TF_ASSERT_OK(InitializeRuntime(dataset_params));
+
+  Tensor empty_resource_handle(DT_RESOURCE, TensorShape({0}));
+  Tensor variant_deleter(DT_VARIANT, TensorShape({}));
+
+  NodeDef node_def = test::function::NDef("delete_iterator", "DeleteIterator",
+                                          {"handle", "deleter"});
+
+  std::unique_ptr<OpKernel> kernel;
+  TF_ASSERT_OK(CreateOpKernel(node_def, &kernel));
+
+  absl::InlinedVector<TensorValue, 4> inputs;
+  inputs.push_back(TensorValue(&empty_resource_handle));
+  inputs.push_back(TensorValue(&variant_deleter));
+
+  std::unique_ptr<OpKernelContext> context;
+  TF_ASSERT_OK(CreateOpKernelContext(kernel.get(), &inputs, &context));
+
+  absl::Status status = RunOpKernel(kernel.get(), context.get());
+  EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument);
+  EXPECT_EQ(status.message(), "Empty resource handle");
+}
+
 }  // namespace
 }  // namespace data
 }  // namespace tensorflow
diff --git a/tensorflow/core/kernels/data/multi_device_iterator_ops.cc b/tensorflow/core/kernels/data/multi_device_iterator_ops.cc
index 106ba0b9325..054f3aac39a 100644
--- a/tensorflow/core/kernels/data/multi_device_iterator_ops.cc
+++ b/tensorflow/core/kernels/data/multi_device_iterator_ops.cc
@@ -155,6 +155,13 @@ class MultiDeviceIterator : public ResourceBase {
           {{"shard_num", shard_num}});
     });
     tf_shared_lock l(mu_);
+    if (!multi_device_buffer_) {
+      return absl::FailedPreconditionError(
+          "GetNextFromShard() failed because the MultiDeviceIterator "
+          "has not been initialized. Ensure that you have run the initializer "
+          "operation for this MultiDeviceIterator before getting the next "
+          "element.");
+    }
     IteratorContext::Params params(ctx);
     params.flr = flr_;
     params.function_handle_cache = function_handle_cache_.get();
@@ -876,7 +883,8 @@ class DeleteMultiDeviceIteratorOp : public OpKernel {
       : OpKernel(ctx) {}
 
   void Compute(OpKernelContext* ctx) override {
-    ResourceHandle handle = ctx->input(0).flat<ResourceHandle>()(0);
+    ResourceHandle handle;
+    OP_REQUIRES_OK(ctx, HandleFromInput(ctx, 0, &handle));
     // The iterator resource is guaranteed to
     // exist because the variant tensor wrapping the deleter is provided as an
     // unused input to this op, which guarantees that it has not run yet.

From cc405ced3adeaf95dd77af49a488a068a6596146 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Thu, 8 Oct 2026 03:29:25 -0700
Subject: [PATCH 0276/1000] [XLA:GPU]: Allow multiple tiles per
 TiledHloInstruction

For variadics where we have index wise dependency b/w inputs and outputs we want different tiles for each output/input pair. This change extends the TiledHloInstruction to support this.

PiperOrigin-RevId: 995719425
---
 .../codegen/tiling/experimental/tiled_hlo.cc  | 71 ++++++++++++-------
 .../codegen/tiling/experimental/tiled_hlo.h   | 30 +++++---
 .../tiling/experimental/tiled_hlo_test.cc     | 37 ++++++++++
 3 files changed, 105 insertions(+), 33 deletions(-)

diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo.cc b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo.cc
index 6f283266d3c..48e31f49b29 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo.cc
@@ -77,7 +77,16 @@ std::string TiledHloInstruction::ToString(
     absl::string_view field_separator) const {
   std::stringstream ss;
   ss << "hlo: " << hlo_->ToString() << field_separator;
-  ss << "tile: " << tile().ToString();
+  if (tiles_.size() == 1) {
+    ss << "tile: " << tile().ToString();
+  } else {
+    for (const auto& [i, t] : llvm::enumerate(tiles_)) {
+      if (i > 0) {
+        ss << field_separator;
+      }
+      ss << "tile #" << i << ": " << t.ToString();
+    }
+  }
   for (const auto& [index, region] : llvm::enumerate(regions_)) {
     ss << field_separator << "region #" << index << " {";
     for (const TiledHloInstruction* instruction : region.instructions()) {
@@ -108,10 +117,10 @@ namespace {
 //
 // This set adds a few key features on top of
 // absl::flat_hash_set<TiledHloInstruction*>:
-// * Elements are inserted as (hlo, tile) pairs. The instruction is constructed
-//   in `instructions` (shared by all regions of a computation) and dropped
-//   again if an equivalent element is already in the set.
-// * Elements are compared by (hlo, tile), not by pointer.
+// * Elements are inserted as (hlo, tiles) pairs. The instruction is
+//   constructed in `instructions` (shared by all regions of a computation) and
+//   dropped again if an equivalent element is already in the set.
+// * Elements are compared by (hlo, tiles), not by pointer.
 // * Elements are stored in the order of insertion.
 class OrderedTiledHloPtrSet {
  public:
@@ -124,8 +133,16 @@ class OrderedTiledHloPtrSet {
   // element was inserted.
   std::pair<TiledHloInstruction*, bool> Insert(const HloInstruction* hlo,
                                                Tile tile) {
+    return Insert(hlo,
+                  llvm::SmallVector<experimental::Tile, 2>{std::move(tile)});
+  }
+
+  // Same as above for an instruction with one tile per result.
+  std::pair<TiledHloInstruction*, bool> Insert(
+      const HloInstruction* hlo,
+      llvm::SmallVector<experimental::Tile, 2> tiles) {
     TiledHloInstruction& candidate =
-        instructions_.emplace_back(hlo, std::move(tile));
+        instructions_.emplace_back(hlo, std::move(tiles));
     auto [it, inserted] = hash_set_.insert(&candidate);
     if (!inserted) {
       instructions_.pop_back();
@@ -143,7 +160,7 @@ class OrderedTiledHloPtrSet {
  private:
   struct PtrHash {
     size_t operator()(const TiledHloInstruction* v) const {
-      return absl::HashOf(v->hlo(), v->tile());
+      return absl::HashOf(v->hlo(), v->tiles());
     }
   };
 
@@ -151,7 +168,7 @@ class OrderedTiledHloPtrSet {
     bool operator()(const TiledHloInstruction* lhs,
                     const TiledHloInstruction* rhs) const {
       return lhs == rhs ||
-             (lhs->hlo() == rhs->hlo() && lhs->tile() == rhs->tile());
+             (lhs->hlo() == rhs->hlo() && lhs->tiles() == rhs->tiles());
     }
   };
 
@@ -194,17 +211,16 @@ void SortTiledHloInstructionsInPostOrder(
   for (const TiledHloInstruction* root_with_no_user : roots_with_no_users) {
     visit_instruction(root_with_no_user);
   }
-  absl::c_sort(tiled_hlo_instructions,
-               [&](const TiledHloInstruction* t1,
-                   const TiledHloInstruction* t2) {
-                 auto it1 = topological_order.find(t1);
-                 auto it2 = topological_order.find(t2);
-                 CHECK(it1 != topological_order.end())
-                     << "Unexpected stray instruction: " << t1->ToString();
-                 CHECK(it2 != topological_order.end())
-                     << "Unexpected stray instruction: " << t2->ToString();
-                 return it1->second < it2->second;
-               });
+  absl::c_sort(tiled_hlo_instructions, [&](const TiledHloInstruction* t1,
+                                           const TiledHloInstruction* t2) {
+    auto it1 = topological_order.find(t1);
+    auto it2 = topological_order.find(t2);
+    CHECK(it1 != topological_order.end())
+        << "Unexpected stray instruction: " << t1->ToString();
+    CHECK(it2 != topological_order.end())
+        << "Unexpected stray instruction: " << t2->ToString();
+    return it1->second < it2->second;
+  });
 
   VLOG(4) << "Sorted symbolic tiled HLO instructions in def-before-use order:\n"
           << absl::StrJoin(
@@ -397,8 +413,15 @@ void PrintTiledHloInstruction(
   std::string indentation(indent, ' ');
   ss << indentation << tile_names.at(tiled_hlo) << " = "
      << HloOpcodeString(tiled_hlo->hlo()->opcode()) << "("
-     << TiledHloOperandsToString(tiled_hlo, tile_names) << ") "
-     << tiled_hlo->tile().ToString(false) << "\n";
+     << TiledHloOperandsToString(tiled_hlo, tile_names) << ")";
+  if (tiled_hlo->tiles().size() == 1) {
+    ss << " " << tiled_hlo->tile().ToString(false) << "\n";
+  } else {
+    ss << "\n";
+    for (const auto& [i, tile] : llvm::enumerate(tiled_hlo->tiles())) {
+      ss << indentation << "  #" << i << " " << tile.ToString(false) << "\n";
+    }
+  }
 
   for (auto const& [i, region] : llvm::enumerate(tiled_hlo->hlo_regions())) {
     ss << indentation << "region #" << i << " {\n";
@@ -555,9 +578,9 @@ absl::StatusOr<TiledHloRegion> CreateHloRegion(
 
 void TiledHloRegion::Simplify() {
   for (TiledHloInstruction* instruction : instructions_) {
-    Tile tile = instruction->tile();
-    tile.Simplify();
-    instruction->set_tile(std::move(tile));
+    for (Tile& tile : instruction->tiles()) {
+      tile.Simplify();
+    }
     for (auto& region : instruction->hlo_regions()) {
       region.Simplify();
     }
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo.h b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo.h
index 754dda910e8..bf7be8db5b2 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo.h
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo.h
@@ -87,12 +87,24 @@ class TiledHloRegion {
 class TiledHloInstruction {
  public:
   TiledHloInstruction(const HloInstruction* hlo, Tile tile)
-      : hlo_(hlo), tile_(std::move(tile)) {}
+      : hlo_(hlo), tiles_({std::move(tile)}) {}
+
+  TiledHloInstruction(const HloInstruction* hlo,
+                      llvm::SmallVector<Tile, 2> tiles)
+      : hlo_(hlo), tiles_(std::move(tiles)) {}
 
   const HloInstruction* hlo() const { return hlo_; }
 
-  const Tile& tile() const { return tile_; }
-  void set_tile(Tile tile) { tile_ = std::move(tile); }
+  const Tile& tile(int64_t output_idx = 0) const {
+    CHECK_GE(output_idx, 0);
+    CHECK_LT(output_idx, tiles_.size());
+    return tiles_[output_idx];
+  }
+  absl::Span<const Tile> tiles() const { return tiles_; }
+  absl::Span<Tile> tiles() { return absl::MakeSpan(tiles_); }
+  void set_tiles(llvm::SmallVector<Tile, 2> tiles) {
+    tiles_ = std::move(tiles);
+  }
 
   const TiledHloInstruction* operand(int64_t operand_id) const {
     return operands_[operand_id];
@@ -128,12 +140,12 @@ class TiledHloInstruction {
   // TODO: b/509505290 -- Remove these once we migrate to this tiling and the
   // old API is removed.
   llvm::SmallVector<int64_t> tile_sizes() const {
-    auto tile_sizes = tile_.GetStaticTileSizes();
+    auto tile_sizes = tile().GetStaticTileSizes();
     CHECK_OK(tile_sizes);
     return *tile_sizes;
   }
   llvm::SmallVector<int64_t> tile_strides() const {
-    auto tile_strides = tile_.GetStaticTileStrides();
+    auto tile_strides = tile().GetStaticTileStrides();
     CHECK_OK(tile_strides);
     return *tile_strides;
   }
@@ -154,8 +166,8 @@ class TiledHloInstruction {
   // Pointer to the original HLO instruction.
   const HloInstruction* hlo_;
 
-  // Symbolic tile.
-  Tile tile_;
+  // Symbolic tile(s), one per tiled output.
+  llvm::SmallVector<Tile, 2> tiles_;
 
   // Operands of the instruction in the tiled computation graph.
   llvm::SmallVector<const TiledHloInstruction*, 2> operands_;
@@ -166,7 +178,7 @@ class TiledHloInstruction {
 
 inline bool operator==(const TiledHloInstruction& lhs,
                        const TiledHloInstruction& rhs) {
-  return lhs.hlo() == rhs.hlo() && lhs.tile() == rhs.tile() &&
+  return lhs.hlo() == rhs.hlo() && lhs.tiles() == rhs.tiles() &&
          lhs.operands() == rhs.operands();
 }
 
@@ -178,7 +190,7 @@ inline bool operator!=(const TiledHloInstruction& lhs,
 template <typename H>
 H AbslHashValue(H h, const TiledHloInstruction& tiled_hlo_instruction) {
   h = H::combine(std::move(h), *tiled_hlo_instruction.hlo(),
-                 tiled_hlo_instruction.tile());
+                 tiled_hlo_instruction.tiles());
   for (const TiledHloInstruction* operand : tiled_hlo_instruction.operands()) {
     h = H::combine(std::move(h), operand);
   }
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc
index b302e0ca53e..7edc81b2bda 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc
@@ -108,6 +108,43 @@ TEST_F(TiledHloTest, TestPrinting) {
   )"));
 }
 
+TEST_F(TiledHloTest, MultiTileInstruction) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    ENTRY e {
+      p0 = f32[10,30] parameter(0)
+      ROOT broadcast = f32[10,20,30] broadcast(p0), dimensions={0,2}
+    }
+  )");
+  ASSERT_OK_AND_ASSIGN(
+      auto tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                          &mlir_context_));
+  Tile root_tile = GetTestTile(*tiling_space, root->shape().dimensions());
+  ASSERT_OK_AND_ASSIGN(
+      Tiles tiled_operands,
+      PropagateTileToInput(*tiling_space, *root, root_tile, 0));
+
+  TiledHloInstruction tiled_hlo_instruction(
+      root, Tiles{root_tile, tiled_operands[0]});
+  ASSERT_EQ(tiled_hlo_instruction.tiles().size(), 2);
+  EXPECT_EQ(tiled_hlo_instruction.tile(0), root_tile);
+  EXPECT_EQ(tiled_hlo_instruction.tile(1), tiled_operands[0]);
+  EXPECT_THAT(tiled_hlo_instruction, MatchString(R"(
+    hlo: %broadcast = f32[10,20,30]{2,1,0} broadcast(%p0), dimensions={0,2}
+    tile #0: (tid_0, tid_1, tid_2)
+      -> offsets [tid_0 * ts_0, tid_1 * ts_1, tid_2 * ts_2]
+         sizes [ts_0, ts_1, ts_2]
+         strides [1, 2, 3]
+         upper bounds [10, 20, 30]
+    tile #1: (tid_0, tid_1, tid_2)
+      -> offsets [tid_0 * ts_0, tid_2 * ts_2]
+         sizes [ts_0, ts_2]
+         strides [1, 3]
+         upper bounds [10, 30]
+  )"));
+}
+
 TEST_F(TiledHloTest, TiledHloRegionDefaultConstruction) {
   TiledHloRegion region;
   EXPECT_TRUE(region.instructions().empty());

From 2cd0c5b52c738eaf05fe18866b9da090fa3fea87 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Thu, 8 Oct 2026 03:56:33 -0700
Subject: [PATCH 0277/1000] [XLA:GPU]: Migrate ConditionalThunk to use
 PerDeviceState

PiperOrigin-RevId: 995731642
---
 .../xla/xla/backends/gpu/runtime/BUILD        |  5 +--
 .../command_buffer_cmd_emitter_test.cc        |  6 +--
 .../command_buffer_conversion_pass_test.cc    |  3 +-
 .../gpu/runtime/command_buffer_thunk_test.cc  |  2 +-
 .../backends/gpu/runtime/conditional_thunk.cc | 42 +++++++++----------
 .../backends/gpu/runtime/conditional_thunk.h  | 18 ++++----
 .../gpu/runtime/conditional_thunk_test.cc     | 13 +++---
 .../runtime/thunk_buffer_debug_pass_test.cc   |  2 +-
 .../runtime/thunk_proto_deserialization.cc    |  7 ++--
 .../backends/gpu/runtime/thunk_walk_test.cc   |  3 +-
 .../xla/xla/service/gpu/thunk_emitter.cc      |  7 +++-
 11 files changed, 55 insertions(+), 53 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index f38fc1e0485..c97319a79eb 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -679,6 +679,7 @@ cc_library(
         ":command",
         ":command_executor",
         ":host_memory_pool",
+        ":per_device_state",
         ":thunk",
         ":thunk_executor",
         ":thunk_proto_cc",
@@ -690,9 +691,6 @@ cc_library(
         "//xla/service:shaped_slice",
         "//xla/stream_executor:command_buffer",
         "//xla/stream_executor:device_address",
-        "//xla/stream_executor:stream_executor_h",
-        "@com_google_absl//absl/base:core_headers",
-        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/functional:function_ref",
         "@com_google_absl//absl/functional:overload",
         "@com_google_absl//absl/log",
@@ -702,7 +700,6 @@ cc_library(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
-        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
     ],
 )
diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_cmd_emitter_test.cc b/third_party/xla/xla/backends/gpu/runtime/command_buffer_cmd_emitter_test.cc
index 9b4d940da49..962e1fc8f53 100644
--- a/third_party/xla/xla/backends/gpu/runtime/command_buffer_cmd_emitter_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_cmd_emitter_test.cc
@@ -424,7 +424,7 @@ TEST_F(CommandBufferCmdEmitterTest, ConvertsConditionalThunkToCommand) {
   auto conditional = std::make_unique<ConditionalThunk>(
       NextThunkInfo("conditional"),
       ShapedSlice{branch_index_slice, ShapeUtil::MakeShape(S32, {})},
-      std::move(branches));
+      std::move(branches), /*devices_per_host=*/1);
   ConditionalThunk* conditional_ptr = conditional.get();
 
   ThunkSequence thunks;
@@ -553,7 +553,7 @@ TEST_F(CommandBufferCmdEmitterTest, ConvertsConditionalThunkRepeatedly) {
   ThunkSequence thunks = ThunkSequence::Of<ConditionalThunk>(
       NextThunkInfo("conditional"),
       ShapedSlice{branch_index_slice, ShapeUtil::MakeShape(S32, {})},
-      std::move(branches));
+      std::move(branches), /*devices_per_host=*/1);
 
   auto collect_command_names = [](CommandExecutor& commands) {
     std::vector<std::string> command_names;
@@ -604,7 +604,7 @@ TEST_F(CommandBufferCmdEmitterTest,
   ThunkSequence thunks = ThunkSequence::Of<ConditionalThunk>(
       NextThunkInfo("conditional"),
       ShapedSlice{branch_index_slice, ShapeUtil::MakeShape(PRED, {})},
-      std::move(branches));
+      std::move(branches), /*devices_per_host=*/1);
 
   ASSERT_OK_AND_ASSIGN(CommandExecutor commands,
                        ConvertToCommands(thunks, ConvertToCommandsOptions()));
diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass_test.cc b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass_test.cc
index 521031886f9..71847e70c27 100644
--- a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass_test.cc
@@ -334,7 +334,8 @@ std::unique_ptr<ConditionalThunk> CreateConditionalThunk(
   Shape shape = ShapeUtil::MakeShape(S32, {});
 
   return std::make_unique<ConditionalThunk>(
-      Thunk::ThunkInfo(), ShapedSlice{slice, shape}, std::move(branch_thunks));
+      Thunk::ThunkInfo(), ShapedSlice{slice, shape}, std::move(branch_thunks),
+      /*devices_per_host=*/1);
 }
 
 std::unique_ptr<CuDnnThunk> CreateCuDnnThunk(const BufferAllocation& alloc0) {
diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/command_buffer_thunk_test.cc
index 5058174b681..613123a5d5e 100644
--- a/third_party/xla/xla/backends/gpu/runtime/command_buffer_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_thunk_test.cc
@@ -1472,7 +1472,7 @@ TEST(CommandBufferThunkTest, ConditionalThunkCaseCommand) {
   // Prepare thunk sequence for command buffer conversion.
   ThunkSequence thunks = ThunkSequence::Of<ConditionalThunk>(
       Thunk::ThunkInfo(), ShapedSlice{slice_i, i_shape},
-      std::move(branch_thunks));
+      std::move(branch_thunks), /*devices_per_host=*/1);
 
   ConvertToCommandsOptions options;
   options.synchronization_mode = serialize;
diff --git a/third_party/xla/xla/backends/gpu/runtime/conditional_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/conditional_thunk.cc
index 37aec8f9bcb..cad0060c570 100644
--- a/third_party/xla/xla/backends/gpu/runtime/conditional_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/conditional_thunk.cc
@@ -31,7 +31,6 @@ limitations under the License.
 #include "absl/status/status_macros.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
-#include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
 #include "xla/backends/gpu/runtime/host_memory_pool.h"
 #include "xla/backends/gpu/runtime/thunk.h"
@@ -124,11 +123,13 @@ absl::StatusOr<const se::CommandBuffer::Command*> HandleRecordAction(
 
 ConditionalThunk::ConditionalThunk(ThunkInfo thunk_info,
                                    const ShapedSlice& branch_index_buffer_index,
-                                   std::vector<ThunkSequence> branch_thunks)
+                                   std::vector<ThunkSequence> branch_thunks,
+                                   int devices_per_host)
     : Command(Kind::kConditional, std::move(thunk_info)),
       branch_index_buffer_index_(branch_index_buffer_index),
       branch_index_is_bool_(branch_index_buffer_index.shape.element_type() ==
-                            PRED) {
+                            PRED),
+      host_memory_pools_(devices_per_host) {
   PrimitiveType element_type = branch_index_buffer_index.shape.element_type();
   CHECK(element_type == PRED || element_type == S32);
   CHECK_EQ(branch_index_buffer_index.shape.dimensions(),
@@ -171,17 +172,14 @@ absl::Status ConditionalThunk::Initialize(const InitializeParams& params) {
     }
   }
 
-  absl::MutexLock lock(mutex_);
-
-  if (!host_memory_pools_.contains(params.executor)) {
-    PrimitiveType type =
-        branch_index_is_bool_ ? PrimitiveType::PRED : PrimitiveType::S32;
-    ABSL_ASSIGN_OR_RETURN(std::unique_ptr<HostMemoryPool> pool,
-                     HostMemoryPool::Create(params.executor, type));
-    host_memory_pools_[params.executor] = std::move(pool);
-  }
-
-  return absl::OkStatus();
+  return host_memory_pools_.GetOrCreateAndInitialize(
+      params.executor->device_ordinal(), [&](PoolState* state) -> absl::Status {
+        PrimitiveType type =
+            branch_index_is_bool_ ? PrimitiveType::PRED : PrimitiveType::S32;
+        ABSL_ASSIGN_OR_RETURN(state->pool,
+                         HostMemoryPool::Create(params.executor, type));
+        return absl::OkStatus();
+      });
 }
 
 absl::StatusOr<const se::CommandBuffer::Command*> ConditionalThunk::Record(
@@ -245,11 +243,9 @@ absl::Status ConditionalThunk::SetOrUpdateCommandBufferBranchExecutors(
 absl::Status ConditionalThunk::ExecuteOnStream(const ExecuteParams& params) {
   auto& stream = *params.stream;
 
-  HostMemoryPool* pool;
-  {
-    absl::MutexLock lock(mutex_);
-    pool = host_memory_pools_.at(stream.parent()).get();
-  }
+  PoolState* state = host_memory_pools_.Find(stream.parent()->device_ordinal());
+  TF_RET_CHECK(state != nullptr && state->pool != nullptr);
+  HostMemoryPool* pool = state->pool.get();
   ABSL_ASSIGN_OR_RETURN(HostMemoryPool::Handle handle, pool->Acquire());
 
   // Copy the predicate value from device.
@@ -346,7 +342,7 @@ absl::StatusOr<ThunkProto> ConditionalThunk::ToProto() const {
 absl::StatusOr<std::unique_ptr<ConditionalThunk>> ConditionalThunk::FromProto(
     ThunkInfo thunk_info, const ConditionalThunkProto& thunk_proto,
     absl::Span<const BufferAllocation> buffer_allocations,
-    const Deserializer& deserializer) {
+    const Deserializer& deserializer, int devices_per_host) {
   ABSL_ASSIGN_OR_RETURN(ShapedSlice branch_index_buffer_index,
                    ShapedSlice::FromProto(thunk_proto.branch_index_buffer(),
                                           buffer_allocations));
@@ -361,9 +357,9 @@ absl::StatusOr<std::unique_ptr<ConditionalThunk>> ConditionalThunk::FromProto(
     }
     branch_thunks.push_back(std::move(thunks));
   }
-  return std::make_unique<ConditionalThunk>(std::move(thunk_info),
-                                            branch_index_buffer_index,
-                                            std::move(branch_thunks));
+  return std::make_unique<ConditionalThunk>(
+      std::move(thunk_info), branch_index_buffer_index,
+      std::move(branch_thunks), devices_per_host);
 }
 
 std::string ConditionalThunk::ToString(int indent) const {
diff --git a/third_party/xla/xla/backends/gpu/runtime/conditional_thunk.h b/third_party/xla/xla/backends/gpu/runtime/conditional_thunk.h
index ddaa6d84aaf..2bdebe6ba07 100644
--- a/third_party/xla/xla/backends/gpu/runtime/conditional_thunk.h
+++ b/third_party/xla/xla/backends/gpu/runtime/conditional_thunk.h
@@ -21,22 +21,20 @@ limitations under the License.
 #include <string>
 #include <vector>
 
-#include "absl/base/thread_annotations.h"
-#include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
 #include "xla/backends/gpu/runtime/command.h"
 #include "xla/backends/gpu/runtime/command_executor.h"
 #include "xla/backends/gpu/runtime/host_memory_pool.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk.pb.h"
 #include "xla/backends/gpu/runtime/thunk_executor.h"
 #include "xla/runtime/buffer_use.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/shaped_slice.h"
-#include "xla/stream_executor/stream_executor.h"
+#include "xla/stream_executor/command_buffer.h"
 
 namespace xla::gpu {
 
@@ -54,7 +52,8 @@ class ConditionalThunk : public Command {
  public:
   ConditionalThunk(ThunkInfo thunk_info,
                    const ShapedSlice& branch_index_buffer_index,
-                   std::vector<ThunkSequence> branch_thunks);
+                   std::vector<ThunkSequence> branch_thunks,
+                   int devices_per_host);
 
   ConditionalThunk(const ConditionalThunk&) = delete;
   ConditionalThunk& operator=(const ConditionalThunk&) = delete;
@@ -107,7 +106,7 @@ class ConditionalThunk : public Command {
   static absl::StatusOr<std::unique_ptr<ConditionalThunk>> FromProto(
       ThunkInfo thunk_info, const ConditionalThunkProto& thunk_proto,
       absl::Span<const BufferAllocation> buffer_allocations,
-      const Deserializer& deserializer);
+      const Deserializer& deserializer, int devices_per_host);
 
   std::string ToString(int indent) const override;
 
@@ -120,9 +119,10 @@ class ConditionalThunk : public Command {
   bool branch_index_is_bool_;
 
   // Host memory pool for transferring predicate value from device to host.
-  absl::Mutex mutex_;
-  absl::flat_hash_map<se::StreamExecutor*, std::unique_ptr<HostMemoryPool>>
-      host_memory_pools_ ABSL_GUARDED_BY(mutex_);
+  struct PoolState {
+    std::unique_ptr<HostMemoryPool> pool;
+  };
+  PerDeviceState<PoolState> host_memory_pools_;
 };
 
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/runtime/conditional_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/conditional_thunk_test.cc
index 6279ea206ce..8c387bc273d 100644
--- a/third_party/xla/xla/backends/gpu/runtime/conditional_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/conditional_thunk_test.cc
@@ -88,7 +88,8 @@ std::unique_ptr<ConditionalThunk> CreateConditionalThunk(
     const ShapedSlice& branch_index_buffer_index,
     std::vector<ThunkSequence> branch_thunk_sequences) {
   return std::make_unique<ConditionalThunk>(
-      thunk_info, branch_index_buffer_index, std::move(branch_thunk_sequences));
+      thunk_info, branch_index_buffer_index, std::move(branch_thunk_sequences),
+      /*devices_per_host=*/1);
 }
 
 struct BranchRecordCounts {
@@ -220,7 +221,7 @@ TEST(ConditionalThunkTest, PreparePropagatesToCommandBufferBranchExecutors) {
   ConditionalThunk thunk(
       Thunk::ThunkInfo(),
       ShapedSlice{branch_index_slice, ShapeUtil::MakeShape(S32, {})},
-      std::move(branch_thunks));
+      std::move(branch_thunks), /*devices_per_host=*/1);
 
   BranchRecordCounts branch0_counts;
   BranchRecordCounts branch1_counts;
@@ -264,7 +265,7 @@ TEST(ConditionalThunkTest, RecordCreatesAndUpdatesCommandBufferCase) {
   ConditionalThunk thunk(
       Thunk::ThunkInfo(),
       ShapedSlice{branch_index_slice, ShapeUtil::MakeShape(S32, {})},
-      std::move(branch_thunks));
+      std::move(branch_thunks), /*devices_per_host=*/1);
 
   BranchRecordCounts branch0_counts;
   BranchRecordCounts branch1_counts;
@@ -477,7 +478,8 @@ TEST(ConditionalThunkTest, FromProto) {
           [](const ThunkProto& proto)
               -> absl::StatusOr<std::unique_ptr<DummyThunk>> {
             return DummyThunk::FromProto(proto, Kind::kCustomCall);
-          }));
+          },
+          /*devices_per_host=*/1));
   ASSERT_NE(thunk, nullptr);
   ASSERT_OK_AND_ASSIGN(ThunkProto round_trip_proto, thunk->ToProto());
   EXPECT_THAT(round_trip_proto, EqualsProto(proto));
@@ -542,7 +544,8 @@ TEST(ConditionalThunkTest, TransformNested) {
   branch_thunks.push_back(std::move(branch0));
   branch_thunks.push_back(std::move(branch1));
   auto conditional_thunk = std::make_unique<ConditionalThunk>(
-      Thunk::ThunkInfo(), ShapedSlice{slice, shape}, std::move(branch_thunks));
+      Thunk::ThunkInfo(), ShapedSlice{slice, shape}, std::move(branch_thunks),
+      /*devices_per_host=*/1);
 
   EXPECT_OK(conditional_thunk->TransformNested([](auto) {
     return std::make_unique<DummyThunk>(Kind::kCustomCall, Thunk::ThunkInfo());
diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk_buffer_debug_pass_test.cc b/third_party/xla/xla/backends/gpu/runtime/thunk_buffer_debug_pass_test.cc
index 6f4a0e16796..54e041bfe89 100644
--- a/third_party/xla/xla/backends/gpu/runtime/thunk_buffer_debug_pass_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/thunk_buffer_debug_pass_test.cc
@@ -317,7 +317,7 @@ TEST_F(ThunkBufferDebugPassTest, RecursivelyInsertsBuffersDebugChecksumThunks) {
 
   auto conditional_thunk = std::make_unique<ConditionalThunk>(
       Thunk::ThunkInfo(), ShapedSlice{condition_slice, condition_shape},
-      std::move(branch_thunks));
+      std::move(branch_thunks), /*devices_per_host=*/1);
   const Thunk* const conditional_thunk_ptr = conditional_thunk.get();
   ThunkSequence while_body_thunks;
   while_body_thunks.push_back(std::move(while_body_fake_thunk));
diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
index 30beccca37f..b3aa4e39014 100644
--- a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
@@ -152,9 +152,10 @@ absl::StatusOr<std::unique_ptr<Thunk>> DeserializeThunkProtoImpl(
           deserializer,
           gpu_topology.has_value() ? gpu_topology->num_devices_per_host() : 0);
     case ThunkProto::kConditionalThunk:
-      return ConditionalThunk::FromProto(std::move(thunk_info),
-                                         thunk_proto.conditional_thunk(),
-                                         buffer_allocations, deserializer);
+      return ConditionalThunk::FromProto(
+          std::move(thunk_info), thunk_proto.conditional_thunk(),
+          buffer_allocations, deserializer,
+          gpu_topology.has_value() ? gpu_topology->num_devices_per_host() : 0);
     case ThunkProto::kGemmThunk:
       return GemmThunk::FromProto(std::move(thunk_info),
                                   thunk_proto.gemm_thunk(), buffer_allocations);
diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk_walk_test.cc b/third_party/xla/xla/backends/gpu/runtime/thunk_walk_test.cc
index fd80ed4fd2b..f04b4a750ea 100644
--- a/third_party/xla/xla/backends/gpu/runtime/thunk_walk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/thunk_walk_test.cc
@@ -97,7 +97,8 @@ TEST(ThunkWalkTest, ConditionalThunk) {
   std::vector<ThunkSequence> branch_thunks;
   branch_thunks.push_back(std::move(thunk_sequence));
   ConditionalThunk conditional_thunk(Thunk::ThunkInfo(), {slice, shape},
-                                     std::move(branch_thunks));
+                                     std::move(branch_thunks),
+                                     /*devices_per_host=*/1);
 
   EXPECT_THAT(GetAllThunks(&conditional_thunk),
               UnorderedElementsAre(thunk_ptr, &conditional_thunk));
diff --git a/third_party/xla/xla/service/gpu/thunk_emitter.cc b/third_party/xla/xla/service/gpu/thunk_emitter.cc
index 32e6857f9a0..77090a3e40a 100644
--- a/third_party/xla/xla/service/gpu/thunk_emitter.cc
+++ b/third_party/xla/xla/service/gpu/thunk_emitter.cc
@@ -597,10 +597,13 @@ Future<ThunkSequence> ThunkEmitter::EmitConditional(
       instr, ir_emitter_context_->GetNextThunkId());
   ShapedSlice shaped_slice{slice, instr->operand(0)->shape()};
   return tsl::JoinFutures(absl::MakeSpan(branch_thunks))
-      .Map([info = std::move(info), shaped_slice = std::move(shaped_slice)](
+      .Map([info = std::move(info), shaped_slice = std::move(shaped_slice),
+            devices_per_host =
+                ir_emitter_context_->gpu_topology().num_devices_per_host()](
                std::vector<ThunkSequence> branch_thunks) mutable {
         return ThunkSequence::Of<ConditionalThunk>(
-            std::move(info), std::move(shaped_slice), std::move(branch_thunks));
+            std::move(info), std::move(shaped_slice), std::move(branch_thunks),
+            devices_per_host);
       });
 }
 

From b66c605383a448d55e8c1a0b8d3f747c886d2661 Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Thu, 8 Oct 2026 04:03:02 -0700
Subject: [PATCH 0278/1000] Add runtime driver checks for CUDA attributes

XLA binaries compiled against newer CUDA headers still support running against older CUDA drivers within the same major version (down to CUDA 12.0 for 12.x builds and CUDA 13.0 for 13.x builds). Several device attributes and fabric handle flags in `xla/stream_executor/cuda/cuda_executor.cc` and `xla/stream_executor/cuda/cuda_device_allocator.cc` were passed to driver APIs without checking the runtime driver version first.

Therefore this change adds explicit runtime driver version guards before using those identifiers:

1. Require CUDA driver 13.4 before querying `CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK` in `GetMaxOversizedSharedMemoryPerBlock`.
2. Require CUDA driver 12.1 before querying `CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED` in `IsMulticastSupported`.
3. Require CUDA driver 12.3 before using `CU_MEM_HANDLE_TYPE_FABRIC` in `ExportFabricHandle`, `ImportFabricHandle`, and `CreateDeviceDescription`, and before querying `CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED` in `IsFabricSupported`.

PiperOrigin-RevId: 995734487
---
 .../cuda/cuda_device_allocator.cc             | 12 ++++-----
 .../xla/stream_executor/cuda/cuda_executor.cc | 26 +++++++++++++++++--
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_device_allocator.cc b/third_party/xla/xla/stream_executor/cuda/cuda_device_allocator.cc
index b3c2ec7deca..0fb97c446dd 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_device_allocator.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_device_allocator.cc
@@ -57,19 +57,19 @@ absl::StatusOr<bool> IsRdmaSupported(CUdevice device) {
 }
 
 absl::StatusOr<bool> IsFabricSupported(CUdevice device) {
+  int driver_version = 0;
+  if (cuDriverGetVersion(&driver_version) != CUDA_SUCCESS ||
+      driver_version < 12030) {
+    return false;
+  }
+
   int fabric_supported = 0;
   CUresult result = cuDeviceGetAttribute(
       &fabric_supported, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED,
       device);
-
-  // Older drivers return INVALID_VALUE when they don't recognize the attribute.
   if (result == CUDA_ERROR_INVALID_VALUE) {
-    XLA_VLOG_DEVICE(1, device)
-        << "CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED not supported "
-           "by driver.";
     return false;
   }
-
   ABSL_RETURN_IF_ERROR(cuda::ToStatus(result));
   return fabric_supported > 0;
 }
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc b/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc
index 35eb43d3433..ea3368679bc 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_executor.cc
@@ -394,6 +394,11 @@ absl::StatusOr<int64_t> GetMaxSharedMemoryPerBlockOptin(CUdevice device) {
 
 int64_t GetMaxOversizedSharedMemoryPerBlock(CUdevice device) {
 #if CUDA_VERSION >= 13040
+  int driver_version = 0;
+  if (cuDriverGetVersion(&driver_version) != CUDA_SUCCESS ||
+      driver_version < 13040) {
+    return 0;
+  }
   return GetSimpleAttribute<int64_t>(
              device, CU_DEVICE_ATTRIBUTE_MAX_OVERSIZED_SHARED_MEMORY_PER_BLOCK)
       .value_or(0);
@@ -585,6 +590,11 @@ CUmemAccessDesc GetVmmAccessDesc(int device) {
 }
 
 absl::StatusOr<bool> IsMulticastSupported(CUdevice device) {
+  int driver_version = 0;
+  ABSL_RETURN_IF_ERROR(cuda::ToStatus(cuDriverGetVersion(&driver_version)));
+  if (driver_version < 12010) {
+    return false;
+  }
   int is_multicast_supported = 0;
   ABSL_RETURN_IF_ERROR(cuda::ToStatus(
       cuDeviceGetAttribute(&is_multicast_supported,
@@ -881,6 +891,11 @@ struct __attribute__((__packed__)) FabricHandle {
 }  // namespace
 
 absl::StatusOr<std::string> CudaExecutor::ExportFabricHandle(void* ptr) const {
+  if (GetDeviceDescription().driver_version() < SemanticVersion{12, 3, 0}) {
+    return absl::UnimplementedError(
+        "ExportFabricHandle requires CUDA driver 12.3 or newer");
+  }
+
   ABSL_ASSIGN_OR_RETURN(VmmMemoryHandle handle, RetainVmmMemoryHandle(ptr));
 
   FabricHandle fabric_handle;
@@ -903,6 +918,11 @@ absl::StatusOr<std::string> CudaExecutor::ExportFabricHandle(void* ptr) const {
 
 absl::StatusOr<DeviceAddressBase> CudaExecutor::ImportFabricHandle(
     absl::string_view serialized) {
+  if (GetDeviceDescription().driver_version() < SemanticVersion{12, 3, 0}) {
+    return absl::UnimplementedError(
+        "ImportFabricHandle requires CUDA driver 12.3 or newer");
+  }
+
   if (serialized.size() != sizeof(FabricHandle)) {
     return absl::InvalidArgumentError(
         absl::StrFormat("Invalid fabric handle size: %d", serialized.size()));
@@ -2000,8 +2020,10 @@ CudaExecutor::CreateDeviceDescription(int device_ordinal) {
     absl::StatusOr<bool> is_multicast_supported = IsMulticastSupported(device);
     if (is_multicast_supported.ok() && *is_multicast_supported) {
       CUmulticastObjectProp prop = {};
-      prop.handleTypes =
-          CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR | CU_MEM_HANDLE_TYPE_FABRIC;
+      prop.handleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
+      if (desc.driver_version() >= SemanticVersion{12, 3, 0}) {
+        prop.handleTypes |= CU_MEM_HANDLE_TYPE_FABRIC;
+      }
       size_t multicast_granularity = 0;
       if (absl::Status status = cuda::ToStatus(cuMulticastGetGranularity(
               &multicast_granularity, &prop, CU_MULTICAST_GRANULARITY_MINIMUM));

From 98bad00555c25fe7d27788b5966139dc6e4b5c87 Mon Sep 17 00:00:00 2001
From: Zviki Nozadze <zviki@google.com>
Date: Thu, 8 Oct 2026 04:03:34 -0700
Subject: [PATCH 0279/1000] [XLA:TPU] Optimize
 HeapSimulator::MakeFreeChunksList using adaptive sort with OR-Tools
 AutoRadixSort.

Replaces unconditional `std::sort` in `HeapSimulator::MakeFreeChunksList` with an adaptive 3-tier sort:

1. Small arrays (<3000) - `absl::c_sort`
2. Nearly-sorted data use `absl::c_sort` with zero swaps.
3. Large unordered arrays (N >= 3000): Dispatches to zero-allocation `operations_research::AutoRadixSort`, reusing scratch buffer without dynamic heap reallocs.

PiperOrigin-RevId: 995734814
---
 .../xla/xla/service/heap_simulator/BUILD      |  6 +-
 .../service/heap_simulator/heap_simulator.cc  | 71 +++++++++++++++++++
 .../service/heap_simulator/heap_simulator.h   | 15 ++++
 3 files changed, 91 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/heap_simulator/BUILD b/third_party/xla/xla/service/heap_simulator/BUILD
index 76ee9dbcdf7..12297e08a37 100644
--- a/third_party/xla/xla/service/heap_simulator/BUILD
+++ b/third_party/xla/xla/service/heap_simulator/BUILD
@@ -20,6 +20,7 @@ load(
     "//xla:xla.default.bzl",
     "xla_cc_test",
 )
+load("//xla/tsl:tsl.bzl", "if_google")
 load("//xla/tsl/platform:rules_cc.bzl", "cc_library")
 
 package(
@@ -87,7 +88,10 @@ cc_library(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
-    ],
+        "@tsl//tsl/platform",
+    ] + if_google([
+        "//third_party/ortools/ortools/algorithms:multikey_radix_sort",
+    ]),
 )
 
 xla_cc_test(
diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
index 3837cea41ab..ef992b772bc 100644
--- a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
+++ b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
@@ -49,6 +49,11 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "absl/types/span.h"
+#include "tsl/platform/platform.h"  // For PLATFORM_GOOGLE.
+
+#if defined(PLATFORM_GOOGLE)
+#include "third_party/ortools/ortools/algorithms/multikey_radix_sort.h"
+#endif  // PLATFORM_GOOGLE
 #include "xla/comparison_util.h"
 #include "xla/hlo/analysis/alias_info.h"
 #include "xla/hlo/analysis/hlo_alias_analysis.h"
@@ -2055,6 +2060,67 @@ std::string RenderTimeByFreeChunks(
 
 }  // namespace
 
+#if defined(PLATFORM_GOOGLE)
+// Sorts `chunks` in ascending order of `Chunk::offset` using an adaptive sort
+// strategy.
+//
+// Optimizations:
+// 1. Tuned cutoff: For small arrays (N < 3000), elements fit entirely
+//    within L1/L2 cache (<48 KB). `absl::c_sort` is faster than radix sort.
+// 2. Single pass to check sortedness by counting inversions:
+//    - If nearly sorted (`inversions <= n / 10`), introsort does almost zero
+//      swaps and beats radix sort by >2x.
+// 3. AutoRadixSort:
+//    For high-entropy inputs with N >= 3000, dispatches to zero-allocation
+//    `operations_research::AutoRadixSort`, reusing the caller-provided scratch
+//    buffer without heap reallocations.
+void AdaptiveHybridSortChunks(std::vector<HeapSimulator::Chunk>& chunks,
+                              std::vector<HeapSimulator::Chunk>& scratch) {
+  using Chunk = HeapSimulator::Chunk;
+  const size_t n = chunks.size();
+
+  // Small lists fit in L1/L2 cache and sort fastest with introsort.
+  if (n < 3000) {
+    absl::c_sort(chunks, [](const Chunk& a, const Chunk& b) {
+      return a.offset < b.offset;
+    });
+    return;
+  }
+
+  // Single pass for sortedness by local inversion count.
+  int64_t prev = chunks[0].offset;
+  size_t inversions = 0;
+  constexpr size_t kMaxInversionDivisor = 10;
+  const size_t max_inversions_for_stdsort = n / kMaxInversionDivisor;
+
+  for (size_t i = 1; i < n; ++i) {
+    const int64_t curr = chunks[i].offset;
+    if (curr < prev) {
+      ++inversions;
+    }
+    prev = curr;
+  }
+
+  // Immediate early exit for already-sorted or uniform arrays.
+  if (inversions == 0) {
+    return;
+  }
+
+  // For nearly-sorted data, introsort does almost no swaps and beats radix
+  // sort.
+  if (inversions <= max_inversions_for_stdsort) {
+    absl::c_sort(chunks, [](const Chunk& a, const Chunk& b) {
+      return a.offset < b.offset;
+    });
+    return;
+  }
+
+  // Use zero-allocation AutoRadixSort.
+  operations_research::AutoRadixSort(
+      chunks, scratch, [](const Chunk& chunk) { return chunk.offset; });
+}
+#endif  // PLATFORM_GOOGLE
+
 template <typename BufferType>
 GlobalDecreasingSizeBestFitHeap<BufferType>::SlicedAllocationFinder::
     SlicedAllocationFinder(
@@ -2588,9 +2654,14 @@ GlobalDecreasingSizeBestFitHeap<BufferType>::MakeFreeChunksList(
         });
   }
 
+#if defined(PLATFORM_GOOGLE)
+  // Sort used chunks by offset ascending using adaptive hybrid sort.
+  AdaptiveHybridSortChunks(used_chunks_, radix_scratch_);
+#else
   // Sort used chunks by offset ascending.
   std::sort(used_chunks_.begin(), used_chunks_.end(),
             [](const Chunk& a, const Chunk& b) { return a.offset < b.offset; });
+#endif  // PLATFORM_GOOGLE
 
   free_chunks_list_.clear();
   if (used_chunks_.empty()) {
diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator.h b/third_party/xla/xla/service/heap_simulator/heap_simulator.h
index 32fdeca5ccf..2d7771dd834 100644
--- a/third_party/xla/xla/service/heap_simulator/heap_simulator.h
+++ b/third_party/xla/xla/service/heap_simulator/heap_simulator.h
@@ -57,6 +57,7 @@ limitations under the License.
 #include "xla/service/hlo_value.h"
 #include "xla/service/logical_buffer.h"
 #include "xla/union_find.h"
+#include "tsl/platform/platform.h"
 
 namespace xla {
 
@@ -347,6 +348,17 @@ class NoFragmentationStatsHeap : public HeapAlgorithm<BufferType> {
   int64_t max_heap_size_ = 0;
 };
 
+#if defined(PLATFORM_GOOGLE)
+// TODO(b/571155846): Open-source this radix sort improvement once OR-Tools
+// with scratch-buffer AutoRadixSort (cl/990317179) is released and the
+// dependency is bumped in OpenXLA.
+//
+// Sorts `chunks` in ascending order of `Chunk::offset` using an adaptive hybrid
+// sort strategy.
+void AdaptiveHybridSortChunks(std::vector<HeapSimulator::Chunk>& chunks,
+                              std::vector<HeapSimulator::Chunk>& scratch);
+#endif  // PLATFORM_GOOGLE
+
 // Node in BufferIntervalTree that stores the alloc and free times of a buffer,
 // and the chunk assigned to it.
 struct BufferIntervalTreeNode {
@@ -1075,6 +1087,9 @@ class GlobalDecreasingSizeBestFitHeap : public HeapAlgorithm<BufferType> {
 
   // Temporary buffers used by MakeFreeChunks to avoid reallocating memory.
   mutable std::vector<Chunk> used_chunks_;
+#if defined(PLATFORM_GOOGLE)
+  mutable std::vector<Chunk> radix_scratch_;
+#endif  // PLATFORM_GOOGLE
   mutable std::vector<std::pair<int64_t, int64_t>> free_chunks_list_;
 
  protected:

From 36f8c4bc112f68d34c9f83380d76318289ab1384 Mon Sep 17 00:00:00 2001
From: Som Tripathi <somtri@iastate.edu>
Date: Thu, 8 Oct 2026 08:16:58 -0500
Subject: [PATCH 0280/1000] Validate input shapes in MultiDeviceIterator and
 DeserializeIterator

MultiDeviceIteratorInit reads max_buffer_size, and
MultiDeviceIteratorGetNextFromShard reads shard_num and incarnation_id,
with Tensor::scalar() and no shape check. A tensor without exactly one
element aborts the process on the CHECK in Tensor::scalar().
DeserializeIterator reads serialized with Tensor::vec(), which aborts on
a tensor that is not rank 1. The op shape functions do not check these
inputs, so eager and graph execution both reach the reads.

Check each input with TensorShapeUtils::IsScalar or IsVector and return
InvalidArgumentError, in the same message form as the existing
resource_handle and string_handle checks in these files. A shape-[1]
tensor used to pass the scalar reads and is now rejected.

Add a test for each check. clang-format also moved the existing
xla/tsl/platform/statusor.h include in iterator_ops.cc into sorted
order.
---
 tensorflow/core/kernels/data/iterator_ops.cc  |   5 +-
 .../core/kernels/data/iterator_ops_test.cc    | 141 ++++++++++++++++++
 .../kernels/data/multi_device_iterator_ops.cc |  10 ++
 3 files changed, 155 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/data/iterator_ops.cc b/tensorflow/core/kernels/data/iterator_ops.cc
index 724dd0a4f75..30a1b398747 100644
--- a/tensorflow/core/kernels/data/iterator_ops.cc
+++ b/tensorflow/core/kernels/data/iterator_ops.cc
@@ -23,7 +23,6 @@ limitations under the License.
 
 #include "absl/status/status.h"
 #include "absl/time/time.h"
-#include "xla/tsl/platform/statusor.h"
 #include "tensorflow/core/activity_watcher/activity.h"
 #include "tensorflow/core/activity_watcher/activity_utils.h"
 #include "tensorflow/core/common_runtime/graph_runner.h"
@@ -43,6 +42,7 @@ limitations under the License.
 #include "tensorflow/core/framework/resource_handle.h"
 #include "tensorflow/core/framework/resource_mgr.h"
 #include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/variant_op_registry.h"
 #include "tensorflow/core/framework/variant_tensor_data.h"
@@ -59,6 +59,7 @@ limitations under the License.
 #include "tensorflow/core/platform/tstring.h"
 #include "tensorflow/core/profiler/lib/traceme.h"
 #include "tensorflow/core/profiler/lib/traceme_encode.h"
+#include "xla/tsl/platform/statusor.h"
 
 namespace tensorflow {
 namespace data {
@@ -1151,6 +1152,8 @@ void DeserializeIteratorOp::Compute(OpKernelContext* ctx) {
   core::ScopedUnref unref_iterator(iterator_resource);
   const Tensor* serialized_t;
   OP_REQUIRES_OK(ctx, ctx->input("serialized", &serialized_t));
+  OP_REQUIRES(ctx, TensorShapeUtils::IsVector(serialized_t->shape()),
+              absl::InvalidArgumentError("serialized must be a vector"));
   IteratorVariantSerializer serializer;
   OP_REQUIRES_OK(ctx, serializer.InitFromTensor(serialized_t));
   absl::Status s = iterator_resource->Restore(ctx, serializer.GetReader());
diff --git a/tensorflow/core/kernels/data/iterator_ops_test.cc b/tensorflow/core/kernels/data/iterator_ops_test.cc
index 9d5c0ca5af1..d13b974e2e9 100644
--- a/tensorflow/core/kernels/data/iterator_ops_test.cc
+++ b/tensorflow/core/kernels/data/iterator_ops_test.cc
@@ -222,6 +222,147 @@ TEST_F(IteratorOpsTest, DeleteIteratorRejectsEmptyHandle) {
   EXPECT_EQ(status.message(), "Empty resource handle");
 }
 
+TEST_F(IteratorOpsTest, MultiDeviceIteratorInitRejectsNonScalarMaxBufferSize) {
+  RangeDatasetParams dataset_params = RangeDatasetParams(0, 10, 3);
+  TF_ASSERT_OK(InitializeRuntime(dataset_params));
+
+  Tensor dataset(DT_VARIANT, TensorShape({}));
+  Tensor multi_device_iterator(DT_RESOURCE, TensorShape({}));
+  Tensor max_buffer_size(DT_INT64, TensorShape({2}));
+  max_buffer_size.flat<int64_t>().setZero();
+
+  NodeDef node_def = test::function::NDef(
+      "multi_device_iterator_init", "MultiDeviceIteratorInit",
+      {"dataset", "multi_device_iterator", "max_buffer_size"}, {});
+
+  std::unique_ptr<OpKernel> kernel;
+  TF_ASSERT_OK(CreateOpKernel(node_def, &kernel));
+
+  absl::InlinedVector<TensorValue, 4> inputs;
+  inputs.push_back(TensorValue(&dataset));
+  inputs.push_back(TensorValue(&multi_device_iterator));
+  inputs.push_back(TensorValue(&max_buffer_size));
+
+  std::unique_ptr<OpKernelContext> context;
+  TF_ASSERT_OK(CreateOpKernelContext(kernel.get(), &inputs, &context));
+
+  absl::Status status = RunOpKernel(kernel.get(), context.get());
+  EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument);
+  EXPECT_EQ(status.message(), "max_buffer_size must be a scalar");
+}
+
+TEST_F(IteratorOpsTest,
+       MultiDeviceIteratorGetNextFromShardRejectsNonScalarShardNum) {
+  RangeDatasetParams dataset_params = RangeDatasetParams(0, 10, 3);
+  TF_ASSERT_OK(InitializeRuntime(dataset_params));
+  DataTypeVector output_types = {DT_INT64};
+  std::vector<PartialTensorShape> output_shapes = {PartialTensorShape({})};
+
+  Tensor multi_device_iterator(DT_RESOURCE, TensorShape({}));
+  Tensor shard_num(DT_INT32, TensorShape({2}));
+  shard_num.flat<int32_t>().setZero();
+  Tensor incarnation_id(DT_INT64, TensorShape({}));
+  incarnation_id.scalar<int64_t>()() = 0;
+
+  NodeDef node_def = test::function::NDef(
+      "get_next_from_shard", "MultiDeviceIteratorGetNextFromShard",
+      {"multi_device_iterator", "shard_num", "incarnation_id"},
+      {{"output_types", output_types}, {"output_shapes", output_shapes}});
+
+  std::unique_ptr<OpKernel> kernel;
+  TF_ASSERT_OK(CreateOpKernel(node_def, &kernel));
+
+  absl::InlinedVector<TensorValue, 4> inputs;
+  inputs.push_back(TensorValue(&multi_device_iterator));
+  inputs.push_back(TensorValue(&shard_num));
+  inputs.push_back(TensorValue(&incarnation_id));
+
+  std::unique_ptr<OpKernelContext> context;
+  TF_ASSERT_OK(CreateOpKernelContext(kernel.get(), &inputs, &context));
+
+  absl::Status status = RunOpKernel(kernel.get(), context.get());
+  EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument);
+  EXPECT_EQ(status.message(), "shard_num must be a scalar");
+}
+
+TEST_F(IteratorOpsTest,
+       MultiDeviceIteratorGetNextFromShardRejectsNonScalarIncarnationId) {
+  RangeDatasetParams dataset_params = RangeDatasetParams(0, 10, 3);
+  TF_ASSERT_OK(InitializeRuntime(dataset_params));
+  DataTypeVector output_types = {DT_INT64};
+  std::vector<PartialTensorShape> output_shapes = {PartialTensorShape({})};
+
+  Tensor multi_device_iterator(DT_RESOURCE, TensorShape({}));
+  Tensor shard_num(DT_INT32, TensorShape({}));
+  shard_num.scalar<int32_t>()() = 0;
+  Tensor incarnation_id(DT_INT64, TensorShape({2}));
+  incarnation_id.flat<int64_t>().setZero();
+
+  NodeDef node_def = test::function::NDef(
+      "get_next_from_shard", "MultiDeviceIteratorGetNextFromShard",
+      {"multi_device_iterator", "shard_num", "incarnation_id"},
+      {{"output_types", output_types}, {"output_shapes", output_shapes}});
+
+  std::unique_ptr<OpKernel> kernel;
+  TF_ASSERT_OK(CreateOpKernel(node_def, &kernel));
+
+  absl::InlinedVector<TensorValue, 4> inputs;
+  inputs.push_back(TensorValue(&multi_device_iterator));
+  inputs.push_back(TensorValue(&shard_num));
+  inputs.push_back(TensorValue(&incarnation_id));
+
+  std::unique_ptr<OpKernelContext> context;
+  TF_ASSERT_OK(CreateOpKernelContext(kernel.get(), &inputs, &context));
+
+  absl::Status status = RunOpKernel(kernel.get(), context.get());
+  EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument);
+  EXPECT_EQ(status.message(), "incarnation_id must be a scalar");
+}
+
+TEST_F(IteratorOpsTest, DeserializeIteratorRejectsNonVectorSerialized) {
+  RangeDatasetParams dataset_params = RangeDatasetParams(0, 10, 3);
+  TF_ASSERT_OK(InitializeRuntime(dataset_params));
+  DataTypeVector output_types = {DT_INT64};
+  std::vector<PartialTensorShape> output_shapes = {PartialTensorShape({})};
+
+  Tensor iterator;
+  {
+    // Creates the iterator resource that DeserializeIterator looks up before
+    // it reads `serialized`. This context must be destroyed before the next
+    // CreateOpKernelContext call, which replaces the params it points to.
+    NodeDef node_def = test::function::NDef(
+        "iterator", "AnonymousIteratorV3", /*inputs=*/{},
+        {{"output_types", output_types}, {"output_shapes", output_shapes}});
+    std::unique_ptr<OpKernel> kernel;
+    TF_ASSERT_OK(CreateOpKernel(node_def, &kernel));
+    absl::InlinedVector<TensorValue, 4> inputs;
+    std::unique_ptr<OpKernelContext> context;
+    TF_ASSERT_OK(CreateOpKernelContext(kernel.get(), &inputs, &context));
+    TF_ASSERT_OK(RunOpKernel(kernel.get(), context.get()));
+    iterator = *context->mutable_output(0);
+  }
+
+  Tensor serialized(DT_VARIANT, TensorShape({}));
+
+  NodeDef node_def =
+      test::function::NDef("deserialize_iterator", "DeserializeIterator",
+                           {"resource_handle", "serialized"}, {});
+
+  std::unique_ptr<OpKernel> kernel;
+  TF_ASSERT_OK(CreateOpKernel(node_def, &kernel));
+
+  absl::InlinedVector<TensorValue, 4> inputs;
+  inputs.push_back(TensorValue(&iterator));
+  inputs.push_back(TensorValue(&serialized));
+
+  std::unique_ptr<OpKernelContext> context;
+  TF_ASSERT_OK(CreateOpKernelContext(kernel.get(), &inputs, &context));
+
+  absl::Status status = RunOpKernel(kernel.get(), context.get());
+  EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument);
+  EXPECT_EQ(status.message(), "serialized must be a vector");
+}
+
 }  // namespace
 }  // namespace data
 }  // namespace tensorflow
diff --git a/tensorflow/core/kernels/data/multi_device_iterator_ops.cc b/tensorflow/core/kernels/data/multi_device_iterator_ops.cc
index 054f3aac39a..6fb1a5e031b 100644
--- a/tensorflow/core/kernels/data/multi_device_iterator_ops.cc
+++ b/tensorflow/core/kernels/data/multi_device_iterator_ops.cc
@@ -42,6 +42,7 @@ limitations under the License.
 #include "tensorflow/core/framework/resource_handle.h"
 #include "tensorflow/core/framework/resource_mgr.h"
 #include "tensorflow/core/framework/resource_op_kernel.h"
+#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/kernels/data/iterator_ops.h"
@@ -675,6 +676,9 @@ class MultiDeviceIteratorInitOp : public OpKernel {
   void Compute(OpKernelContext* ctx) override {
     const Tensor* tensor_max_buffer_size;
     OP_REQUIRES_OK(ctx, ctx->input("max_buffer_size", &tensor_max_buffer_size));
+    OP_REQUIRES(ctx,
+                TensorShapeUtils::IsScalar(tensor_max_buffer_size->shape()),
+                absl::InvalidArgumentError("max_buffer_size must be a scalar"));
     int64_t max_buffer_size = tensor_max_buffer_size->scalar<int64_t>()();
 
     DatasetBase* dataset;
@@ -729,11 +733,17 @@ class MultiDeviceIteratorGetNextFromShardOp : public AsyncOpKernel {
   void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override {
     const Tensor* tensor_shard_num;
     OP_REQUIRES_OK_ASYNC(ctx, ctx->input("shard_num", &tensor_shard_num), done);
+    OP_REQUIRES_ASYNC(
+        ctx, TensorShapeUtils::IsScalar(tensor_shard_num->shape()),
+        absl::InvalidArgumentError("shard_num must be a scalar"), done);
     int32_t shard_num = tensor_shard_num->scalar<int32_t>()();
 
     const Tensor* tensor_incarnation_id;
     OP_REQUIRES_OK_ASYNC(
         ctx, ctx->input("incarnation_id", &tensor_incarnation_id), done);
+    OP_REQUIRES_ASYNC(
+        ctx, TensorShapeUtils::IsScalar(tensor_incarnation_id->shape()),
+        absl::InvalidArgumentError("incarnation_id must be a scalar"), done);
     int64_t incarnation_id = tensor_incarnation_id->scalar<int64_t>()();
 
     MultiDeviceIterator* iterator;

From d78b03de3183e4223e0b664b26b364ffc78d5d8d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Eusebio=20Dur=C3=A1n=20Monta=C3=B1a?= <eusebiodm@google.com>
Date: Thu, 8 Oct 2026 04:19:43 -0700
Subject: [PATCH 0281/1000] Require cuDNN >= 9.23 for deviceless cuDNN
 compilation.

We are running into errors when running deviceless cuDNN convolutions in h100s, which get fixed after 9.23. To avoid convoluted logic gating different feature/device combinations for different versions, we have a single version check after which everything should be working.

I moved the version check to `CudnnGraph::Prepare` since that's guaranteed to run, and its job is to make these checks in general.
```
 // Prepares a graph and checks whether it is generally supported.
  absl::Status Prepare(dnn::DnnSupport*, ...
```

PiperOrigin-RevId: 995741761
---
 .../xla/xla/backends/gpu/autotuner/BUILD      |  1 +
 .../xla/xla/backends/gpu/autotuner/cudnn.cc   |  5 --
 .../xla/backends/gpu/autotuner/cudnn_test.cc  |  5 ++
 .../xla/xla/backends/gpu/codegen/BUILD        |  2 +
 .../xla/backends/gpu/codegen/cudnn_test.cc    | 66 +++++++++++++++++--
 .../xla/xla/backends/gpu/transforms/BUILD     |  4 +-
 .../gpu/transforms/conv_fp8_fallback_test.cc  | 32 ++-------
 .../gpu/transforms/cudnn_fusion_compiler.cc   | 10 ---
 .../cudnn_fusion_compiler_deviceless_test.cc  | 24 +++----
 .../xla/xla/service/gpu/nvptx_compiler.cc     |  8 ---
 .../xla/xla/stream_executor/cuda/cuda_dnn.cc  | 23 +++----
 .../xla/xla/stream_executor/cuda/cuda_dnn.h   | 20 +++---
 third_party/xla/xla/xla.proto                 |  2 +-
 13 files changed, 114 insertions(+), 88 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/autotuner/BUILD b/third_party/xla/xla/backends/gpu/autotuner/BUILD
index 5f92c0e8736..2e39f0f355b 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/BUILD
+++ b/third_party/xla/xla/backends/gpu/autotuner/BUILD
@@ -301,6 +301,7 @@ xla_test(
         "//xla/stream_executor:device_description_proto_cc",
         "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor/cuda:cuda_compute_capability",
+        "//xla/stream_executor/cuda:cudnn_plugin",
         "//xla/tsl/protobuf:dnn_proto_cc",
         "//xla/tsl/util/proto:proto_matchers",
         "@com_google_absl//absl/status",
diff --git a/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc b/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc
index 4c0577acede..7689c859e09 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc
@@ -269,11 +269,6 @@ GetCudnnFusionConfigs(const HloInstruction& instr,
       break;
   }
   if (use_deviceless) {
-    if (target_config.device_description.dnn_version() <
-        se::SemanticVersion(9, 8, 0)) {
-      return absl::FailedPreconditionError(
-          "Deviceless cuDNN compilation requires cuDNN >= 9.8.");
-    }
     stream_executor = nullptr;
   } else if (stream_executor == nullptr) {
     return absl::InvalidArgumentError(
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 f6ab20f5dc7..2aaa51d83b9 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/cudnn_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/cudnn_test.cc
@@ -36,6 +36,7 @@ limitations under the License.
 #include "xla/service/gpu/nvptx_compiler.h"
 #include "xla/service/platform_util.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
+#include "xla/stream_executor/cuda/cuda_dnn.h"
 #include "xla/stream_executor/device_description.pb.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/protobuf/dnn.pb.h"
@@ -322,6 +323,10 @@ TEST_F(CudnnBackendTest, GetDefaultConfigFromCudnnCustomCall) {
 }
 
 TEST_F(CudnnBackendTest, GetDefaultConfigSucceedsWithNullStreamExecutor) {
+  if (!se::gpu::SupportsDevicelessCudnnCompilation()) {
+    GTEST_SKIP() << "Deviceless cuDNN compilation requires cuDNN >= "
+                 << se::gpu::kMinDevicelessCudnnVersion;
+  }
   ASSERT_OK_AND_ASSIGN(stream_executor::GpuTargetConfigProto proto,
                        GetGpuTargetConfig(GpuModel::H100_SXM));
   ASSERT_OK_AND_ASSIGN(Compiler::GpuTargetConfig target_config,
diff --git a/third_party/xla/xla/backends/gpu/codegen/BUILD b/third_party/xla/xla/backends/gpu/codegen/BUILD
index 30541a87435..d1d5ba1dd25 100644
--- a/third_party/xla/xla/backends/gpu/codegen/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/BUILD
@@ -99,10 +99,12 @@ xla_test(
         "//xla/stream_executor:semantic_version",
         "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor/cuda:cuda_compute_capability",
+        "//xla/stream_executor/cuda:cudnn_plugin",
         "//xla/tests:hlo_interpreter_reference_mixin",
         "//xla/tsl/platform:env",
         "//xla/tsl/platform:test",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",
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 b33faaab583..f83c689216d 100644
--- a/third_party/xla/xla/backends/gpu/codegen/cudnn_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/cudnn_test.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/log/check.h"
+#include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/status_matchers.h"
 #include "absl/status/statusor.h"
@@ -47,6 +48,7 @@ limitations under the License.
 #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/cuda/cuda_dnn.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/platform_manager.h"
 #include "xla/stream_executor/semantic_version.h"
@@ -513,8 +515,9 @@ CHECK: "stride": [{{[[:space:]]*}}4096,{{[[:space:]]*}}128,{{[[:space:]]*}}1{{[[
 }
 
 TEST_F(CuDnnFusionExecutionTest, DotF32DevicelessCompilationSucceeds) {
-  if (!IsAtLeastCuDnnVersion(9, 8)) {
-    GTEST_SKIP() << "Deviceless DeviceProperties requires cuDNN 9.8+.";
+  if (!se::gpu::SupportsDevicelessCudnnCompilation()) {
+    GTEST_SKIP() << "Deviceless cuDNN compilation requires cuDNN >= "
+                 << se::gpu::kMinDevicelessCudnnVersion;
   }
   constexpr absl::string_view kHlo = R"(
 fusion1 {
@@ -564,9 +567,64 @@ ENTRY e {
                             ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
+// Repro for b/568668570: below cudnn 9.23, the devicelessly built plan for
+// this convolution would fail at runtime.
+constexpr absl::string_view kConvWorkspaceReproHlo = R"hlo(
+  ENTRY e {
+    input = bf16[1,8192,1536] parameter(0)
+    filter = bf16[1536,5,1536] parameter(1)
+    ROOT conv = bf16[1,8192,1536] convolution(input, filter),
+      window={size=5 pad=2_2}, dim_labels=b0f_o0i->b0f,
+      convolution_kind=fprop
+  })hlo";
+
+HloModuleConfig WithConvFusionAndDevicelessMode(
+    HloModuleConfig config, DebugOptions::CudnnDevicelessCompilationMode mode) {
+  DebugOptions& options = config.mutable_debug_options();
+  options.set_xla_gpu_experimental_enable_conv_fusion(true);
+  options.set_xla_gpu_cudnn_deviceless_compilation_mode(mode);
+  return config;
+}
+
+TEST_F(CuDnnFusionExecutionTest, ConvDevicelessCompilationMatchesLive) {
+  if (!se::gpu::SupportsDevicelessCudnnCompilation()) {
+    GTEST_SKIP() << "Deviceless cuDNN compilation requires cuDNN >= "
+                 << se::gpu::kMinDevicelessCudnnVersion;
+  }
+  EXPECT_TRUE(RunAndCompareTwoModules(
+      kConvWorkspaceReproHlo, kConvWorkspaceReproHlo,
+      WithConvFusionAndDevicelessMode(
+          GetModuleConfigForTest(),
+          DebugOptions::CUDNN_DEVICELESS_COMPILATION_ALWAYS),
+      WithConvFusionAndDevicelessMode(
+          GetModuleConfigForTest(),
+          DebugOptions::CUDNN_DEVICELESS_COMPILATION_DISABLED),
+      ErrorSpec{/*aabs=*/1e-2, /*arel=*/1e-2}));
+}
+
+TEST_F(CuDnnFusionExecutionTest,
+       ConvDevicelessCompilationFailsBelowMinCudnnVersion) {
+  if (se::gpu::SupportsDevicelessCudnnCompilation()) {
+    GTEST_SKIP() << "Loaded cuDNN supports deviceless compilation (>= "
+                 << se::gpu::kMinDevicelessCudnnVersion << ").";
+  }
+  // The autotuner rewraps the FailedPrecondition from CudnnGraph::Prepare, so
+  // only the message survives.
+  EXPECT_THAT(
+      GetOptimizedModule(
+          kConvWorkspaceReproHlo,
+          WithConvFusionAndDevicelessMode(
+              GetModuleConfigForTest(),
+              DebugOptions::CUDNN_DEVICELESS_COMPILATION_ALWAYS)),
+      absl_testing::StatusIs(
+          ::testing::_, ::testing::HasSubstr(
+                            "Deviceless cuDNN compilation requires cuDNN >=")));
+}
+
 TEST_F(CuDnnFusionExecutionTest, DotF32DevicelessBinaryMatchesLive) {
-  if (!IsAtLeastCuDnnVersion(9, 8)) {
-    GTEST_SKIP() << "Deviceless DeviceProperties requires cuDNN 9.8+.";
+  if (!se::gpu::SupportsDevicelessCudnnCompilation()) {
+    GTEST_SKIP() << "Deviceless cuDNN compilation requires cuDNN >= "
+                 << se::gpu::kMinDevicelessCudnnVersion;
   }
   constexpr absl::string_view kHlo = R"(
 fusion1 {
diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD
index e136fc4c081..fb6871a918e 100644
--- a/third_party/xla/xla/backends/gpu/transforms/BUILD
+++ b/third_party/xla/xla/backends/gpu/transforms/BUILD
@@ -629,7 +629,7 @@ xla_test(
     srcs = ["conv_fp8_fallback_test.cc"],
     # All NVIDIA GPUs: the probe verdicts driving the pass differ per
     # generation. Every case is deviceless and opens no GPU, but needs a
-    # loadable cuDNN >= 9.8 runtime and a CUDA driver (cuDNN's deviceless
+    # loadable cuDNN >= 9.23 runtime and a CUDA driver (cuDNN's deviceless
     # heuristics abort without one), so GPU machines only.
     backends = ["gpu"],
     tags = [
@@ -1212,7 +1212,7 @@ xla_test(
     srcs = ["cudnn_fusion_compiler_deviceless_test.cc"],
     # All NVIDIA GPUs: the parity test checks the deviceless verdict against
     # live plan enumeration on each generation. The deviceless cases open no
-    # GPU, but they need a loadable cuDNN >= 9.8 runtime and a CUDA driver
+    # GPU, but they need a loadable cuDNN >= 9.23 runtime and a CUDA driver
     # (cuDNN's deviceless heuristics abort without one), so GPU machines only.
     backends = ["gpu"],
     tags = [
diff --git a/third_party/xla/xla/backends/gpu/transforms/conv_fp8_fallback_test.cc b/third_party/xla/xla/backends/gpu/transforms/conv_fp8_fallback_test.cc
index 796cc418e85..2249590603a 100644
--- a/third_party/xla/xla/backends/gpu/transforms/conv_fp8_fallback_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/conv_fp8_fallback_test.cc
@@ -22,10 +22,9 @@ limitations under the License.
 //
 // The ConvFp8FallbackDevicelessTest cases run the full pass against
 // stream_executor::DeviceDescriptions built from checked-in target-config
-// specs and open no GPU; they need a loadable host cuDNN >= 9.8 (like the
-// SupportsFusionDeviceless probe they drive) and skip otherwise. The sm_120
-// cases additionally skip on cuDNN runtimes < 9.19, whose deviceless
-// heuristics cannot probe Blackwell-generation targets.
+// specs and open no GPU; they need a loadable host cuDNN >=
+// se::gpu::kMinDevicelessCudnnVersion (like the SupportsFusionDeviceless probe
+// they drive) and skip otherwise.
 //
 // The fusions under test are produced by the real ConvKindAssignment +
 // ConvFusionRewriter passes, so they cannot drift from pipeline output.
@@ -335,9 +334,9 @@ TEST_F(ConvFp8FallbackRewriteTest, TwoOutputAmaxFusionOnlyConvertsF8Output) {
 class ConvFp8FallbackDevicelessTest : public ConvFp8FallbackTestBase {
  protected:
   void SetUp() override {
-    if (!se::gpu::SupportsDevicelessDeviceProperties()) {
-      GTEST_SKIP() << "cuDNN runtime < 9.8 does not support deviceless "
-                      "DeviceProperties.";
+    if (!se::gpu::SupportsDevicelessCudnnCompilation()) {
+      GTEST_SKIP() << "Deviceless cuDNN compilation requires cuDNN >= "
+                   << se::gpu::kMinDevicelessCudnnVersion;
     }
   }
 };
@@ -412,24 +411,7 @@ TEST_F(ConvFp8FallbackDevicelessTest, Fp8ConvIsKeptWhenBf16AlsoHasNoPlans) {
   EXPECT_OK(verifier().Run(module.get()).status());
 }
 
-class ConvFp8FallbackSm120DevicelessTest
-    : public ConvFp8FallbackDevicelessTest {
- protected:
-  void SetUp() override {
-    ConvFp8FallbackDevicelessTest::SetUp();
-    if (IsSkipped()) {
-      return;
-    }
-    absl::StatusOr<GpuTargetConfig> target_config =
-        DevicelessTargetConfig(GpuModel::RTX6000PRO);
-    ASSERT_TRUE(target_config.ok()) << target_config.status();
-    if (!se::gpu::SupportsDevicelessConvGraphs(
-            target_config->device_description)) {
-      GTEST_SKIP() << "cuDNN runtime cannot probe conv graphs devicelessly "
-                      "for sm_120 targets (requires cuDNN >= 9.19).";
-    }
-  }
-};
+using ConvFp8FallbackSm120DevicelessTest = ConvFp8FallbackDevicelessTest;
 
 // channels/group == 16 has cuDNN plans on sm_120; the fusion must be kept
 // FP8.
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 c8780024c03..5d9da5b84d5 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
@@ -1262,16 +1262,6 @@ absl::StatusOr<se::gpu::CudnnGraph> PrepareGraph(
     se::dnn::DnnSupport* dnn_support,
     const se::DeviceDescription& gpu_device_info,
     const HloFusionInstruction& hlo) {
-  if (dnn_support == nullptr &&
-      hlo_query::GetFirstInstructionWithOpcode(
-          *hlo.fused_instructions_computation(), HloOpcode::kConvolution) !=
-          nullptr &&
-      !se::gpu::SupportsDevicelessConvGraphs(gpu_device_info)) {
-    return absl::FailedPreconditionError(
-        "Deviceless cuDNN preparation of convolution graphs targeting "
-        "Blackwell-generation GPUs requires cuDNN >= 9.19; older runtimes "
-        "crash inside the deviceless heuristics query.");
-  }
   ABSL_ASSIGN_OR_RETURN(se::gpu::CudnnGraph graph, HloFusionToCuDnnGraph(hlo));
   PrecisionConfig precision_config = GetPrecisionConfig(hlo);
 
diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler_deviceless_test.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler_deviceless_test.cc
index 1a7cf9479fb..7d866f4ea4c 100644
--- a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler_deviceless_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler_deviceless_test.cc
@@ -16,9 +16,11 @@ limitations under the License.
 // Tests for CuDnnFusionCompiler::SupportsFusionDeviceless.
 //
 // The deviceless cases build DeviceDescriptions from checked-in target-config
-// specs and open no GPU; they need a loadable host cuDNN >= 9.8 and skip
-// otherwise. DevicelessSupportMatchesLivePlanEnumeration needs a real GPU
-// whose deviceless conv probing is not gated off and skips otherwise.
+// specs and open no GPU; they need a loadable host cuDNN >=
+// se::gpu::kMinDevicelessCudnnVersion and skip otherwise.
+// DevicelessSupportMatchesLivePlanEnumeration and
+// DevicelessConvWorkspaceMatchesLiveDeserializedPlan additionally use the
+// attached GPU.
 //
 // The fusions under test are produced by the real ConvKindAssignment +
 // ConvFusionRewriter passes, so they cannot drift from pipeline output.
@@ -63,16 +65,17 @@ namespace se = ::stream_executor;
 
 using DevicelessFusionSupport = CuDnnFusionCompiler::DevicelessFusionSupport;
 
+// Deviceless probing/compilation only works on cuDNN >=
+// kMinDevicelessCudnnVersion; every test skips below it.
 class CudnnFusionCompilerDevicelessTest
     : public HloHardwareIndependentTestBase {
  protected:
   void SetUp() override {
-    if (!se::gpu::SupportsDevicelessDeviceProperties()) {
-      GTEST_SKIP() << "cuDNN runtime < 9.8 does not support deviceless "
-                      "DeviceProperties.";
+    if (!se::gpu::SupportsDevicelessCudnnCompilation()) {
+      GTEST_SKIP() << "Deviceless cuDNN compilation requires cuDNN >= "
+                   << se::gpu::kMinDevicelessCudnnVersion;
     }
   }
-
   // Builds a deviceless GpuTargetConfig for the given GPU model from its
   // checked-in target-config spec (no StreamExecutor / GPU required).
   static absl::StatusOr<GpuTargetConfig> DevicelessTargetConfig(
@@ -425,10 +428,9 @@ TEST_F(CudnnFusionCompilerDevicelessTest,
   }
   const se::DeviceDescription& device_description =
       executor->GetDeviceDescription();
-  if (!se::gpu::SupportsDevicelessConvGraphs(device_description)) {
-    GTEST_SKIP() << "Deviceless conv graph probing is gated off for this "
-                    "GPU / cuDNN runtime combination, so every deviceless "
-                    "verdict is kUnknown and there is no parity to check.";
+  if (!se::gpu::SupportsDevicelessCudnnCompilation()) {
+    GTEST_SKIP() << "Deviceless cuDNN compilation requires cuDNN >= "
+                 << se::gpu::kMinDevicelessCudnnVersion;
   }
   const se::SemanticVersion dnn_version = device_description.dnn_version();
   struct ParityCase {
diff --git a/third_party/xla/xla/service/gpu/nvptx_compiler.cc b/third_party/xla/xla/service/gpu/nvptx_compiler.cc
index 82e42abacbb..d1263e2cebf 100644
--- a/third_party/xla/xla/service/gpu/nvptx_compiler.cc
+++ b/third_party/xla/xla/service/gpu/nvptx_compiler.cc
@@ -405,14 +405,6 @@ absl::Status NVPTXCompiler::RunCudnnCompilerPasses(
   if (stream_exec == nullptr && !use_deviceless_cudnn) {
     return absl::OkStatus();
   }
-  // Deviceless cuDNN compilation relies on DeviceProperties JSON
-  // serialization, added in cuDNN 9.8.
-  if (use_deviceless_cudnn &&
-      gpu_target_config.device_description.dnn_version() <
-          se::SemanticVersion(9, 8, 0)) {
-    return absl::FailedPreconditionError(
-        "Deviceless cuDNN compilation requires cuDNN >= 9.8.");
-  }
   se::dnn::DnnSupport* dnn_support =
       use_deviceless_cudnn ? nullptr : stream_exec->AsDnn();
 
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc b/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc
index 3a51567e2bc..3adc96689a6 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc
@@ -6852,15 +6852,12 @@ absl::StatusOr<std::unique_ptr<dnn::DnnGraph>> CudnnSupport::DeserializeGraph(
   return std::make_unique<CudnnGraph>(std::move(graph));
 }
 
-bool SupportsDevicelessDeviceProperties() {
-  return cudnn_frontend::detail::get_backend_version() >= 90800;
-}
-
-bool SupportsDevicelessConvGraphs(const DeviceDescription& gpu_device_info) {
-  const auto* cc =
-      gpu_device_info.gpu_compute_capability().cuda_compute_capability();
-  return cc == nullptr || cc->major < 10 ||
-         cudnn_frontend::detail::get_backend_version() >= 91900;
+bool SupportsDevicelessCudnnCompilation() {
+  const int min_backend_version =
+      kMinDevicelessCudnnVersion.major_version() * 10000 +
+      kMinDevicelessCudnnVersion.minor_version() * 100 +
+      kMinDevicelessCudnnVersion.patch_version();
+  return cudnn_frontend::detail::get_backend_version() >= min_backend_version;
 }
 
 namespace {
@@ -6937,11 +6934,15 @@ absl::Status CudnnGraph::Prepare(dnn::DnnSupport* dnn_support,
     ABSL_RETURN_IF_ERROR(create_and_filter_plans());
     RETURN_CUDNN_FRONTEND_STATUS(graph_.check_support(cudnn.handle()));
   } else {
-    // Deviceless mode. No cuDNN version guard needed: DeviceProperties
-    // deserialization inside BuildDeviceProperties rejects runtimes < 9.8.
     if (EnsureCudaCompatLoaded != nullptr) {
       EnsureCudaCompatLoaded();
     }
+    if (!SupportsDevicelessCudnnCompilation()) {
+      return absl::FailedPreconditionError(absl::StrCat(
+          "Deviceless cuDNN compilation requires cuDNN >= ",
+          kMinDevicelessCudnnVersion.ToString(), " (loaded ",
+          cudnn_frontend::detail::get_backend_version_string(), ")."));
+    }
     ABSL_ASSIGN_OR_RETURN(auto device_props,
                      xla::gpu::BuildDeviceProperties(gpu_device_info));
     graph_.set_device_properties(device_props);
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_dnn.h b/third_party/xla/xla/stream_executor/cuda/cuda_dnn.h
index 2edf89a7679..aba240a66e3 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_dnn.h
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_dnn.h
@@ -37,6 +37,7 @@ limitations under the License.
 #include "xla/stream_executor/dnn.h"
 #include "xla/stream_executor/engine_options.h"
 #include "xla/stream_executor/scratch_allocator.h"
+#include "xla/stream_executor/semantic_version.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/protobuf/dnn.pb.h"
 
@@ -95,18 +96,15 @@ class CudnnGraph : public dnn::DnnGraph {
       std::optional<int64_t> local_device_ordinal = std::nullopt) const;
 };
 
-// Returns whether the loaded cuDNN runtime supports deviceless
-// DeviceProperties (added in cuDNN 9.8). Informational (e.g. for test skips);
-// CudnnGraph::Prepare does not guard on this, the frontend rejects older
-// runtimes itself.
-bool SupportsDevicelessDeviceProperties();
+// Minimum cuDNN runtime for deviceless cuDNN graph compilation.
+// Technically older versions support deviceless compilation, but they have
+// various bugs.
+inline constexpr SemanticVersion kMinDevicelessCudnnVersion{9, 23, 0};
 
-// Returns whether the loaded cuDNN runtime can prepare convolution graphs
-// devicelessly for the given target. Runtimes older than 9.19 crash (SIGSEGV
-// observed with 9.8.0) inside the deviceless heuristics query for
-// Blackwell-generation (sm_100+) targets instead of returning an error, so
-// callers must refuse such probes rather than attempt them.
-bool SupportsDevicelessConvGraphs(const DeviceDescription& gpu_device_info);
+// Returns whether the loaded cuDNN runtime is at least
+// kMinDevicelessCudnnVersion. CudnnGraph::Prepare refuses deviceless
+// preparation on older runtimes.
+bool SupportsDevicelessCudnnCompilation();
 
 // cudnn-library based DNN support. For details on overridden interface
 // functions, see dnn.h.
diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index 99ee8dc60f3..eba8f167462 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -483,7 +483,7 @@ message DebugOptions {
 
   // Controls when cuDNN fusions are compiled in "deviceless" mode, i.e. driven
   // purely by the DeviceDescription without a live cuDNN handle. Deviceless
-  // compilation requires cuDNN >= 9.8.
+  // compilation requires cuDNN >= 9.23.
   enum CudnnDevicelessCompilationMode {
     CUDNN_DEVICELESS_COMPILATION_UNSET = 0;
     // Never use deviceless compilation; always require a live device.

From 6a4e21b8a7670ed5a70c076dcbd46a8987ce4f23 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Eusebio=20Dur=C3=A1n=20Monta=C3=B1a?= <eusebiodm@google.com>
Date: Thu, 8 Oct 2026 04:55:00 -0700
Subject: [PATCH 0282/1000] Set cuDNN graph SM version from device info

During deviceless cuDNN compilation (`dnn_support == nullptr`), `cudnn_frontend::graph::Graph::validate()` checks inspect `context.get_sm_version()`. When the SM version is left at its default (`-1`), `cudnn_frontend` falls back to querying the live CUDA runtime via `cudaGetDevice()`, which fails when we have no device.

Our existing tests were using the deviceless APIs, but they still were running in a machine with a GPU, which was hiding this error. So I added a tests that compiles a conv on Forge, and later runs it with a device to make sure it works (test fails without the `cuda_dnn.cc` fix).

Reverts e4a3d4b812b2389ececd3e4d190fef887c027e72

PiperOrigin-RevId: 995755678
---
 third_party/xla/xla/service/BUILD             | 18 +++++-
 .../xla/service/xla_aot_compile_gpu_test.cc   | 60 +++++++++++--------
 third_party/xla/xla/service/xla_compile.bzl   | 17 ++++--
 .../xla/xla/stream_executor/cuda/BUILD        |  2 +
 .../xla/xla/stream_executor/cuda/cuda_dnn.cc  | 20 +++++++
 5 files changed, 84 insertions(+), 33 deletions(-)

diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index 6b691c8cad2..3317497e3dd 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -5967,6 +5967,21 @@ xla_aot_compile_gpu(
     xla_flags = "--xla_gpu_use_new_autotune_cache_format=false",
 )
 
+xla_aot_compile_gpu(
+    name = "xla_aot_compile_test_gpu_executable_convolution_deviceless_cudnn",
+    # H100 only: deviceless conv graphs on Blackwell require cuDNN >= 9.19.
+    gpu_targets = [
+        "h100",
+        "h100_mig",
+    ],
+    module = "xla_aot_compile_test_convolution.mlir",
+    xla_flags = " ".join([
+        "--xla_gpu_experimental_enable_conv_fusion=true",
+        "--xla_gpu_cudnn_deviceless_compilation_mode=CUDNN_DEVICELESS_COMPILATION_AUTO",
+        "--xla_gpu_experimental_autotune_backends=cudnn",
+    ]),
+)
+
 xla_aot_compile_gpu_runtime_autotuning(
     name = "xla_aot_compile_test_gpu_executable_convolution_runtime_autotuning",
     gpu_target_config = "//xla/backends/gpu/target_config:specs/h100_sxm.txtpb",
@@ -6048,6 +6063,7 @@ xla_test(
         ":xla_aot_compile_test_gpu_executable",
         ":xla_aot_compile_test_gpu_executable_constant",
         ":xla_aot_compile_test_gpu_executable_convolution",
+        ":xla_aot_compile_test_gpu_executable_convolution_deviceless_cudnn",
         ":xla_aot_compile_test_gpu_executable_convolution_legacy_cache",
         ":xla_aot_compile_test_gpu_executable_hlo",
         ":xla_aot_compile_test_gpu_executable_hlo_legacy_cache",
@@ -6063,7 +6079,7 @@ xla_test(
         ":hlo_runner_pjrt",
         "//xla:literal",
         "//xla:literal_util",
-        "//xla/tests:hlo_test_base",
+        "//xla/backends/gpu/tests:hlo_pjrt_gpu_test_base",
         "//xla/tsl/platform:env",
         "//xla/tsl/platform:test",
         "@com_google_absl//absl/status:status_macros",
diff --git a/third_party/xla/xla/service/xla_aot_compile_gpu_test.cc b/third_party/xla/xla/service/xla_aot_compile_gpu_test.cc
index 53606130117..113c0dc1f89 100644
--- a/third_party/xla/xla/service/xla_aot_compile_gpu_test.cc
+++ b/third_party/xla/xla/service/xla_aot_compile_gpu_test.cc
@@ -20,14 +20,15 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/status/status_macros.h"
+#include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
+#include "xla/backends/gpu/tests/hlo_pjrt_gpu_test_base.h"
 #include "xla/literal.h"
 #include "xla/literal_util.h"
 #include "xla/service/hlo_runner_interface.h"
 #include "xla/service/hlo_runner_pjrt.h"
-#include "xla/tests/hlo_test_base.h"
 #include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/test.h"
 #include "tsl/platform/path.h"
@@ -36,16 +37,18 @@ namespace xla {
 namespace xla_compile {
 namespace {
 
-class XlaCompileTest : public HloTestBase {
+class XlaCompileTest : public gpu::HloPjRtGpuTestBase {
  public:
   void LoadAndRunExecutable(absl::string_view path_to_serialized_aot_result,
                             absl::Span<const Literal* const> args,
-                            const Literal& expected) {
+                            const Literal& expected,
+                            absl::string_view target_suffix = "") {
     const char* test_device = getenv("XLA_TEST_DEVICE");
     ASSERT_NE(test_device, nullptr) << "XLA_TEST_DEVICE is not set";
-    std::string path = tsl::io::JoinPath(
-        tsl::testing::XlaSrcRoot(), "service",
-        absl::StrCat(path_to_serialized_aot_result, "_", test_device));
+    std::string path =
+        tsl::io::JoinPath(tsl::testing::XlaSrcRoot(), "service",
+                          absl::StrCat(path_to_serialized_aot_result, "_",
+                                       test_device, target_suffix));
     std::string serialized_aot_result;
     ASSERT_OK(tsl::ReadFileToString(tsl::Env::Default(), path,
                                     &serialized_aot_result));
@@ -88,8 +91,9 @@ TEST_F(XlaCompileTest, LoadGpuExecutableWithConstant) {
                        expected);
 }
 
-// Should also cover the case of loading a GPU executable with a GEMM.
-TEST_F(XlaCompileTest, LoadGpuExecutableWithConvolution) {
+void LoadAndRunConvolution(XlaCompileTest& test,
+                           absl::string_view path_to_serialized_aot_result,
+                           absl::string_view target_suffix = "") {
   Literal input1 = LiteralUtil::CreateR4<float>(
       {{{{1.0, 2.0}, {3.0, 4.0}, {5.0, 6.0}, {7.0, 8.0}},
         {{11.0, 12.0}, {13.0, 14.0}, {15.0, 16.0}, {17.0, 18.0}},
@@ -103,27 +107,31 @@ TEST_F(XlaCompileTest, LoadGpuExecutableWithConvolution) {
       {{1310.0}, {1466.0}, {1622.0}},
       {{2090.0}, {2246.0}, {2402.0}},
   }});
-  LoadAndRunExecutable("xla_aot_compile_test_gpu_executable_convolution",
-                       {&input1, &input2}, expected);
+  test.LoadAndRunExecutable(path_to_serialized_aot_result, {&input1, &input2},
+                            expected, target_suffix);
+}
+
+// Should also cover the case of loading a GPU executable with a GEMM.
+TEST_F(XlaCompileTest, LoadGpuExecutableWithConvolution) {
+  LoadAndRunConvolution(*this,
+                        "xla_aot_compile_test_gpu_executable_convolution");
 }
 
 TEST_F(XlaCompileTest, LoadGpuExecutableWithConvolutionLegacyCache) {
-  Literal input1 = LiteralUtil::CreateR4<float>(
-      {{{{1.0, 2.0}, {3.0, 4.0}, {5.0, 6.0}, {7.0, 8.0}},
-        {{11.0, 12.0}, {13.0, 14.0}, {15.0, 16.0}, {17.0, 18.0}},
-        {{21.0, 22.0}, {23.0, 24.0}, {25.0, 26.0}, {27.0, 28.0}},
-        {{31.0, 32.0}, {33.0, 34.0}, {35.0, 36.0}, {37.0, 38.0}}}});
-  Literal input2 =
-      LiteralUtil::CreateR4<float>({{{{1.0}, {2.0}}, {{3.0}, {4.0}}},
-                                    {{{5.0}, {6.0}}, {{7.0}, {8.0}}},
-                                    {{{9.0}, {10.0}}, {{11.0}, {12.0}}}});
-  Literal expected = LiteralUtil::CreateR4<float>({{
-      {{1310.0}, {1466.0}, {1622.0}},
-      {{2090.0}, {2246.0}, {2402.0}},
-  }});
-  LoadAndRunExecutable(
-      "xla_aot_compile_test_gpu_executable_convolution_legacy_cache",
-      {&input1, &input2}, expected);
+  LoadAndRunConvolution(
+      *this, "xla_aot_compile_test_gpu_executable_convolution_legacy_cache");
+}
+
+TEST_F(XlaCompileTest, LoadGpuExecutableWithDevicelessCudnnConvolution) {
+  const char* test_device = getenv("XLA_TEST_DEVICE");
+  if (test_device == nullptr || absl::string_view(test_device) != "h100") {
+    GTEST_SKIP() << "Only compiled for h100.";
+  }
+  absl::string_view target_suffix =
+      absl::StrContains(device_description().name(), "MIG") ? "_mig" : "";
+  LoadAndRunConvolution(
+      *this, "xla_aot_compile_test_gpu_executable_convolution_deviceless_cudnn",
+      target_suffix);
 }
 
 }  // namespace
diff --git a/third_party/xla/xla/service/xla_compile.bzl b/third_party/xla/xla/service/xla_compile.bzl
index 11096f37ed1..5655da66ab6 100644
--- a/third_party/xla/xla/service/xla_compile.bzl
+++ b/third_party/xla/xla/service/xla_compile.bzl
@@ -97,7 +97,7 @@ def xla_aot_compile_gpu(
         name,
         module,
         gpu_targets,
-        autotune_results,
+        autotune_results = None,
         xla_flags = ""):
     """Runs xla_compile to compile an MHLO, StableHLO or HLO module into an AotCompilationResult for GPU
 
@@ -105,7 +105,7 @@ def xla_aot_compile_gpu(
         name: The name of the build rule.
         module: The MHLO or StableHLO file to compile.
         gpu_targets: The list of gpu targets.
-        autotune_results: AOT AutotuneResults or AutotuneCache file.
+        autotune_results: Optional AOT AutotuneResults or AutotuneCache file.
         xla_flags: Additional XLA_FLAGS to set during compilation.
     """
 
@@ -120,13 +120,18 @@ def xla_aot_compile_gpu(
             " --platform=gpu" +
             " --gpu_target_config=$(location " + gpu_target_config_map[target] + ")"
         )
-        flags = "--xla_gpu_load_autotune_results_from=$(location " + autotune_results + ")"
+        srcs = [module, gpu_target_config_map[target]]
+        flags = []
+        if autotune_results:
+            srcs.append(autotune_results)
+            flags.append("--xla_gpu_load_autotune_results_from=$(location " + autotune_results + ")")
         if xla_flags:
-            flags = flags + " " + xla_flags
-        cmd = "XLA_FLAGS=\"" + flags + "\" " + cmd
+            flags.append(xla_flags)
+        if flags:
+            cmd = "XLA_FLAGS=\"" + " ".join(flags) + "\" " + cmd
         native.genrule(
             name = "gen_" + name + "_" + target,
-            srcs = [module, gpu_target_config_map[target], autotune_results],
+            srcs = srcs,
             outs = [name + "_" + target],
             cmd = cmd,
             tools = [xla_compile_tool],
diff --git a/third_party/xla/xla/stream_executor/cuda/BUILD b/third_party/xla/xla/stream_executor/cuda/BUILD
index 149fa607070..05f7bb33b62 100644
--- a/third_party/xla/xla/stream_executor/cuda/BUILD
+++ b/third_party/xla/xla/stream_executor/cuda/BUILD
@@ -569,10 +569,12 @@ cc_library(
         ":cudnn_api_wrappers",
         ":cudnn_frontend_helpers",
         ":cudnn_sdpa_score_mod",
+        "//xla:status_macros",
         "//xla/backends/gpu/target_config:cudnn_device_props",
         "//xla/stream_executor:activate_context",
         "//xla/stream_executor:data_type",
         "//xla/stream_executor:device_address",
+        "//xla/stream_executor:device_description",
         "//xla/stream_executor:dnn",
         "//xla/stream_executor:engine_options",
         "//xla/stream_executor:event_based_timer",
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc b/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc
index 3adc96689a6..51846165556 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc
@@ -54,6 +54,7 @@ limitations under the License.
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 #include "third_party/gpus/cuda/include/driver_types.h"
 #include "xla/backends/gpu/target_config/cudnn_device_props.h"
+#include "xla/status_macros.h"
 #include "xla/stream_executor/activate_context.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/cuda/cuda_diagnostics.h"
@@ -64,6 +65,7 @@ limitations under the License.
 #include "xla/stream_executor/cuda/cudnn_sdpa_score_mod.h"
 #include "xla/stream_executor/data_type.h"
 #include "xla/stream_executor/device_address.h"
+#include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/dnn.h"
 #include "xla/stream_executor/engine_options.h"
 #include "xla/stream_executor/event_based_timer.h"
@@ -3990,6 +3992,17 @@ static absl::StatusOr<cudnn_frontend::ExecutionPlan> RebuildExecutionPlan(
   return {std::move(plan)};
 }
 
+// Sets the SM version in the cuDNN graph. Without this, cuDNN will fallback to
+// querying the GPU, which doesn't work for deviceless compilation.
+absl::Status SetGraphSmVersion(cudnn_frontend::graph::Graph& graph,
+                               const DeviceDescription& gpu_device_info) {
+  const CudaComputeCapability* cc =
+      gpu_device_info.gpu_compute_capability().cuda_compute_capability();
+  TF_RET_CHECK(cc != nullptr);
+  graph.set_sm_version(cc->major * 10 + cc->minor);
+  return absl::OkStatus();
+}
+
 }  // namespace
 
 void FixDimsForRaggedOffset(std::vector<int64_t>& dims, int max_reg_per_batch) {
@@ -4037,6 +4050,7 @@ absl::StatusOr<CudnnGraph> GetCudnnFlashAttentionOperationGraph(
   graph.set_intermediate_data_type(cudnn_frontend::DataType_t::FLOAT)
       .set_io_data_type(ioDataType)
       .set_compute_data_type(cudnn_frontend::DataType_t::FLOAT);
+  ABSL_RETURN_IF_ERROR(SetGraphSmVersion(graph, gpu_device_info));
   int64_t uid = 0;
   auto next_uid = [&]() -> int64_t { return CuDnnTensorUID(uid++); };
 
@@ -4286,6 +4300,7 @@ absl::StatusOr<CudnnGraph> GetCudnnFlashAttentionF8OperationGraph(
   graph.set_intermediate_data_type(cudnn_frontend::DataType_t::FLOAT)
       .set_io_data_type(ioDataType)
       .set_compute_data_type(cudnn_frontend::DataType_t::FLOAT);
+  ABSL_RETURN_IF_ERROR(SetGraphSmVersion(graph, gpu_device_info));
 
   auto next_uid = [uid = 0]() mutable -> int { return CuDnnTensorUID(uid++); };
 
@@ -4420,6 +4435,7 @@ absl::StatusOr<CudnnGraph> GetCudnnFlashAttentionBackwardF8OperationGraph(
   graph.set_compute_data_type(cudnn_frontend::DataType_t::FLOAT)
       .set_intermediate_data_type(cudnn_frontend::DataType_t::FLOAT)
       .set_io_data_type(ioDataType);
+  ABSL_RETURN_IF_ERROR(SetGraphSmVersion(graph, gpu_device_info));
 
   auto next_uid = [uid = 0]() mutable -> int { return CuDnnTensorUID(uid++); };
 
@@ -4607,6 +4623,7 @@ absl::StatusOr<CudnnGraph> GetCudnnBlockScaledDotOperationGraph(
           << "\n global_scale: " << has_global_scale;
 
   cudnn_frontend::graph::Graph graph;
+  ABSL_RETURN_IF_ERROR(SetGraphSmVersion(graph, gpu_device_info));
   auto compute_type = cudnn_frontend::DataType_t::FLOAT;
   graph.set_compute_data_type(compute_type);
   graph.set_intermediate_data_type(compute_type);
@@ -4717,6 +4734,7 @@ absl::StatusOr<CudnnGraph> GetCudnnFlashAttentionBackwardOperationGraph(
 
   using cudnn_frontend::graph::Tensor_attributes;
   cudnn_frontend::graph::Graph graph;
+  ABSL_RETURN_IF_ERROR(SetGraphSmVersion(graph, gpu_device_info));
   if (!(q_desc.type() == k_desc.type() && k_desc.type() == p_desc.type() &&
         p_desc.type() == v_desc.type() && v_desc.type() == do_desc.type() &&
         do_desc.type() == dq_desc.type() && dq_desc.type() == dk_desc.type() &&
@@ -6921,6 +6939,8 @@ absl::Status CudnnGraph::Prepare(dnn::DnnSupport* dnn_support,
     return absl::OkStatus();
   };
 
+  ABSL_RETURN_IF_ERROR(SetGraphSmVersion(graph_, gpu_device_info));
+
   if (dnn_support) {
     const CudnnSupport& cudnn_support =
         static_cast<CudnnSupport&>(*dnn_support);

From 26bc85cbd207562467238b9f080243365ee21b06 Mon Sep 17 00:00:00 2001
From: Dmitri Latushko <latushko@google.com>
Date: Thu, 8 Oct 2026 05:07:15 -0700
Subject: [PATCH 0283/1000] [XLA:CI] Retry failed Bazel builds on Windows CI to
 mitigate transient lld-link file-open races

On 64-bit Windows (`_WIN64`), `lld-link.exe` (`lld/COFF/Driver.cpp` `createFutureForFile`) opens all input `.lib` archives concurrently via `std::async(std::launch::async, ...)`. In Windows Server 2022 containers (`wcifs.sys`), this concurrent open burst occasionally hits transient `STATUS_SHARING_VIOLATION` / `STATUS_ACCESS_DENIED` errors, and `LinkerDriver::enqueuePath` silently drops the failed `.lib` archive because its synchronous fallback calls `findFileIfNew(pathStr)` a second time on a path whose `UniqueID` was already recorded in `visitedFiles`.

While the underlying `LinkerDriver::enqueuePath` bug in `lld/COFF/Driver.cpp` is being fixed in LLVM (retrying `MemoryBuffer::getFile(pathStr)` directly before consulting `findFileIfNew`), the Windows CI container image uses a prebuilt `LLVM-19.1.7` release binary. Configure `command_retries=2` on `XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS` and `JAX_WINDOWS_X86_CPU_GITHUB_ACTIONS` in `build_tools/ci/build.py` to retry failed Bazel builds against the warm local Bazel cache until the container image includes the patched `lld-link.exe`.

PiperOrigin-RevId: 995760973
---
 third_party/xla/build_tools/ci/BUILD         |  19 +++
 third_party/xla/build_tools/ci/build.py      |  56 +++++--
 third_party/xla/build_tools/ci/build_test.py | 154 +++++++++++++++++++
 3 files changed, 219 insertions(+), 10 deletions(-)
 create mode 100644 third_party/xla/build_tools/ci/build_test.py

diff --git a/third_party/xla/build_tools/ci/BUILD b/third_party/xla/build_tools/ci/BUILD
index 03dfa81decf..4e26877a62a 100644
--- a/third_party/xla/build_tools/ci/BUILD
+++ b/third_party/xla/build_tools/ci/BUILD
@@ -67,6 +67,15 @@ py_test(
     ],
 )
 
+pytype_strict_library(
+    name = "build_lib",
+    srcs = ["build.py"],
+    deps = [
+        ":bazel_diff",
+        ":change_detector",
+    ],
+)
+
 pytype_strict_binary(
     name = "build",
     srcs = ["build.py"],
@@ -76,6 +85,16 @@ pytype_strict_binary(
     ],
 )
 
+py_test(
+    name = "build_test",
+    srcs = ["build_test.py"],
+    deps = [
+        ":build_lib",
+        "@absl_py//absl/testing:absltest",
+        "@absl_py//absl/testing:parameterized",
+    ],
+)
+
 pytype_strict_library(
     name = "clang_tidy_diff_lib",
     srcs = ["clang_tidy_diff.py"],
diff --git a/third_party/xla/build_tools/ci/build.py b/third_party/xla/build_tools/ci/build.py
index 00ee9809c9c..33905e53a3c 100755
--- a/third_party/xla/build_tools/ci/build.py
+++ b/third_party/xla/build_tools/ci/build.py
@@ -164,6 +164,8 @@ class Build:
     extra_setup_commands: Tuple of shell commands to run before Bazel.
     use_bazel_diff: Whether to enable bazel-diff target filtering on presubmit.
     bazel_diff_use_cquery: Whether bazel-diff should use cquery (--useCquery).
+    command_retries: Number of times to retry failed Bazel commands against the
+      warm cache (e.g. to mitigate transient Windows lld-link file-open races).
   """
 
   _builds: ClassVar[Dict[BuildType, "Build"]] = {}
@@ -185,6 +187,7 @@ class Build:
   extra_setup_commands: Tuple[List[str], ...] = ()
   use_bazel_diff: bool = False
   bazel_diff_use_cquery: bool = True
+  command_retries: int = 0
 
   def __post_init__(self):
     # pylint: disable=protected-access
@@ -488,6 +491,11 @@ Build(
     startup_options={
         "output_user_root": "C:/x",
     },
+    # Retry failed Bazel builds against the warm local cache to mitigate
+    # transient Windows lld-link wcifs.sys file-open races (b/571438464) until
+    # the Windows CI container image includes the lld/COFF/Driver.cpp
+    # LinkerDriver::enqueuePath fix (b/571437556).
+    command_retries=2,
 )
 
 Build(
@@ -915,6 +923,11 @@ Build(
     startup_options={
         "output_base": f"{_GITHUB_WORKSPACE}\\bazel_output_base",
     },
+    # Retry failed Bazel builds against the warm local cache to mitigate
+    # transient Windows lld-link wcifs.sys file-open races (b/571438464) until
+    # the Windows CI container image includes the lld/COFF/Driver.cpp
+    # LinkerDriver::enqueuePath fix (b/571437556).
+    command_retries=2,
 )
 
 Build(
@@ -984,6 +997,37 @@ def get_xla_dir(build: Build) -> str:
   return "."
 
 
+def execute_build_commands(
+    build: Build, target_pattern_file: str | None = None
+) -> None:
+  """Executes commands for a build, retrying failed Bazel commands if configured."""
+  for command in build.commands(target_pattern_file=target_pattern_file):
+    max_retries = (
+        build.command_retries
+        if command and command[0] == "bazel" and build.command_retries > 0
+        else 0
+    )
+    for attempt in range(max_retries + 1):
+      result = sh(command, check=False)
+      if result.returncode == 0:
+        break
+      if result.returncode == 4 and target_pattern_file:
+        logging.info(
+            "Bazel returned exit code 4 (no tests found), treating as success."
+        )
+        break
+      if attempt < max_retries:
+        logging.warning(
+            "Bazel command failed with exit code %d (attempt %d/%d); retrying"
+            " against warm Bazel cache...",
+            result.returncode,
+            attempt + 1,
+            max_retries + 1,
+        )
+        continue
+      sys.exit(result.returncode)
+
+
 def _parse_args():
   """Defines flags and parses args."""
   parser = argparse.ArgumentParser(allow_abbrev=False)
@@ -1001,7 +1045,7 @@ def _parse_args():
   return parser.parse_args()
 
 
-def main():
+def main() -> None:
   logging.basicConfig()
   logging.getLogger().setLevel(logging.INFO)
 
@@ -1067,15 +1111,7 @@ def main():
       elif decision.decision == bazel_diff.BazelDiffDecisionType.IMPACTED:
         target_pattern_file = decision.impacted_targets_file
 
-  for command in build.commands(target_pattern_file=target_pattern_file):
-    result = sh(command, check=False)
-    if result.returncode == 4 and target_pattern_file:
-      logging.info(
-          "Bazel returned exit code 4 (no tests found), treating as success."
-      )
-      continue
-    if result.returncode != 0:
-      sys.exit(result.returncode)
+  execute_build_commands(build, target_pattern_file=target_pattern_file)
 
 
 if __name__ == "__main__":
diff --git a/third_party/xla/build_tools/ci/build_test.py b/third_party/xla/build_tools/ci/build_test.py
new file mode 100644
index 00000000000..ad4377b8fda
--- /dev/null
+++ b/third_party/xla/build_tools/ci/build_test.py
@@ -0,0 +1,154 @@
+# Copyright 2026 The OpenXLA Authors.
+#
+# 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 build module."""
+
+from collections.abc import Sequence
+import subprocess
+from unittest import mock
+
+from absl.testing import absltest
+from absl.testing import parameterized
+
+from build_tools.ci import build
+
+
+class BuildTest(parameterized.TestCase):
+
+  def test_windows_builds_configure_command_retries(self) -> None:
+    builds = build.Build.all_builds()
+    self.assertEqual(
+        builds[
+            build.BuildType.XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS
+        ].command_retries,
+        2,
+    )
+    self.assertEqual(
+        builds[
+            build.BuildType.JAX_WINDOWS_X86_CPU_GITHUB_ACTIONS
+        ].command_retries,
+        2,
+    )
+    self.assertEqual(
+        builds[
+            build.BuildType.XLA_LINUX_X86_CPU_GITHUB_ACTIONS
+        ].command_retries,
+        0,
+    )
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="xla_windows_succeeds_on_first_attempt",
+          build_type=build.BuildType.XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS,
+          returncodes=[0],
+      ),
+      dict(
+          testcase_name="xla_windows_succeeds_on_first_retry",
+          build_type=build.BuildType.XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS,
+          returncodes=[1, 0],
+      ),
+      dict(
+          testcase_name="xla_windows_succeeds_on_second_retry",
+          build_type=build.BuildType.XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS,
+          returncodes=[1, 1, 0],
+      ),
+      dict(
+          testcase_name="jax_windows_succeeds_on_first_retry",
+          build_type=build.BuildType.JAX_WINDOWS_X86_CPU_GITHUB_ACTIONS,
+          returncodes=[1, 0],
+      ),
+  )
+  @mock.patch.object(build, "sh", autospec=True, spec_set=True)
+  def test_execute_build_commands_windows_retries_and_succeeds(
+      self,
+      mock_sh: mock.MagicMock,
+      build_type: build.BuildType,
+      returncodes: Sequence[int],
+  ) -> None:
+    mock_sh.side_effect = [
+        subprocess.CompletedProcess(args=["bazel"], returncode=return_code)
+        for return_code in returncodes
+    ]
+    target_build = build.Build.all_builds()[build_type]
+    build.execute_build_commands(target_build)
+    self.assertLen(returncodes, mock_sh.call_count)
+    mock_sh.assert_called_with(
+        target_build.commands(target_pattern_file=None)[0], check=False
+    )
+
+  @mock.patch.object(build, "sh", autospec=True, spec_set=True)
+  def test_execute_build_commands_windows_exits_after_exhausting_retries(
+      self, mock_sh: mock.MagicMock
+  ) -> None:
+    mock_sh.return_value = subprocess.CompletedProcess(
+        args=["bazel"], returncode=1
+    )
+    target_build = build.Build.all_builds()[
+        build.BuildType.XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS
+    ]
+    with self.assertLogs(level="WARNING") as logs:
+      with self.assertRaisesRegex(SystemExit, "^1$"):
+        build.execute_build_commands(target_build)
+    self.assertEqual(mock_sh.call_count, 3)
+    self.assertLen(logs.output, 2)
+    self.assertIn("attempt 1/3", logs.output[0])
+    self.assertIn("attempt 2/3", logs.output[1])
+
+  @mock.patch.object(build, "sh", autospec=True, spec_set=True)
+  def test_execute_build_commands_non_windows_does_not_retry(
+      self, mock_sh: mock.MagicMock
+  ) -> None:
+    mock_sh.side_effect = [
+        subprocess.CompletedProcess(args=["parallel"], returncode=0),
+        subprocess.CompletedProcess(args=["bazel"], returncode=2),
+    ]
+    target_build = build.Build.all_builds()[
+        build.BuildType.XLA_LINUX_X86_CPU_GITHUB_ACTIONS
+    ]
+    with self.assertRaisesRegex(SystemExit, "^2$"):
+      build.execute_build_commands(target_build)
+    self.assertEqual(mock_sh.call_count, 2)
+
+  @mock.patch.object(build, "sh", autospec=True, spec_set=True)
+  def test_execute_build_commands_exit_code_4_with_target_pattern_file(
+      self, mock_sh: mock.MagicMock
+  ) -> None:
+    mock_sh.return_value = subprocess.CompletedProcess(
+        args=["bazel"], returncode=4
+    )
+    target_build = build.Build.all_builds()[
+        build.BuildType.XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS
+    ]
+    build.execute_build_commands(
+        target_build, target_pattern_file="/tmp/targets.txt"
+    )
+    self.assertEqual(mock_sh.call_count, 1)
+
+  @mock.patch.object(build, "sh", autospec=True, spec_set=True)
+  def test_execute_build_commands_exit_code_4_without_target_pattern_file(
+      self, mock_sh: mock.MagicMock
+  ) -> None:
+    mock_sh.return_value = subprocess.CompletedProcess(
+        args=["bazel"], returncode=4
+    )
+    target_build = build.Build.all_builds()[
+        build.BuildType.XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS
+    ]
+    with self.assertRaisesRegex(SystemExit, "^4$"):
+      build.execute_build_commands(target_build, target_pattern_file=None)
+    self.assertEqual(mock_sh.call_count, 3)
+
+
+if __name__ == "__main__":
+  absltest.main()

From e872f34b8553caf2c60b77f3f30b1be983739fb0 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Thu, 8 Oct 2026 05:45:29 -0700
Subject: [PATCH 0284/1000] [XLA:GPU]: Migrate FftThunk to use PerDeviceState

PiperOrigin-RevId: 995777271
---
 .../xla/xla/backends/gpu/runtime/BUILD        |  2 +-
 .../xla/xla/backends/gpu/runtime/fft_thunk.cc | 16 ++++++++-----
 .../xla/xla/backends/gpu/runtime/fft_thunk.h  | 23 ++++++++-----------
 .../backends/gpu/runtime/fft_thunk_test.cc    |  3 ++-
 .../runtime/thunk_proto_deserialization.cc    |  5 ++--
 .../xla/xla/service/gpu/thunk_emitter.cc      |  3 ++-
 6 files changed, 28 insertions(+), 24 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index c97319a79eb..47bffe3e149 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -1221,6 +1221,7 @@ cc_library(
     srcs = ["fft_thunk.cc"],
     hdrs = ["fft_thunk.h"],
     deps = [
+        ":per_device_state",
         ":thunk",
         ":thunk_proto_cc",
         "//xla:shape_util",
@@ -1239,7 +1240,6 @@ cc_library(
         "//xla/stream_executor:stream_executor_h",
         "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/base:core_headers",
-        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
diff --git a/third_party/xla/xla/backends/gpu/runtime/fft_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/fft_thunk.cc
index 165d73e7a62..3da6a2074cf 100644
--- a/third_party/xla/xla/backends/gpu/runtime/fft_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/fft_thunk.cc
@@ -127,12 +127,14 @@ FftThunk::FftThunk(ThunkInfo thunk_info, FftType fft_type,
                    absl::Span<const int64_t> fft_length,
                    const BufferAllocation::Slice& input_buffer,
                    const BufferAllocation::Slice& output_buffer,
-                   const Shape& input_shape, const Shape& output_shape)
+                   const Shape& input_shape, const Shape& output_shape,
+                   int devices_per_host)
     : Thunk(Kind::kFft, thunk_info),
       fft_type_(
           FftTypeToSeType(fft_type, input_shape.element_type() == F64 ||
                                         input_shape.element_type() == C128)),
       fft_length_(fft_length.begin(), fft_length.end()),
+      fft_plan_cache_(devices_per_host),
       input_buffer_(input_buffer),
       output_buffer_(output_buffer),
       input_shape_(input_shape),
@@ -162,7 +164,8 @@ absl::Status RunFft(se::DeviceAddressBase input, const Shape& input_shape,
                                                   memory_allocator);
 
   // Get the Fft plan for the given device ordinal.
-  FftPlan* fft_plan_ptr = fft_plan_cache->GetOrCreate(device_ordinal);
+  ABSL_ASSIGN_OR_RETURN(FftPlan * fft_plan_ptr,
+                   fft_plan_cache->GetOrCreate(device_ordinal));
 
   // CuFFT thread-safety requires that separate host threads not share plans;
   // protect each plan with a mutex.
@@ -294,7 +297,8 @@ absl::Status RunFft(se::DeviceAddressBase input, const Shape& input_shape,
 
 absl::StatusOr<std::unique_ptr<FftThunk>> FftThunk::FromProto(
     ThunkInfo thunk_info, const FftThunkProto& proto,
-    absl::Span<const BufferAllocation> buffer_allocations) {
+    absl::Span<const BufferAllocation> buffer_allocations,
+    int devices_per_host) {
   ABSL_ASSIGN_OR_RETURN(BufferAllocation::Slice input_buffer,
                    BufferAllocation::Slice::FromProto(proto.input_buffer(),
                                                       buffer_allocations));
@@ -308,9 +312,9 @@ absl::StatusOr<std::unique_ptr<FftThunk>> FftThunk::FromProto(
   std::vector<int64_t> fft_length{proto.fft_length().begin(),
                                   proto.fft_length().end()};
 
-  return std::make_unique<FftThunk>(thunk_info, proto.fft_type(),
-                                    std::move(fft_length), input_buffer,
-                                    output_buffer, input_shape, output_shape);
+  return std::make_unique<FftThunk>(
+      thunk_info, proto.fft_type(), std::move(fft_length), input_buffer,
+      output_buffer, input_shape, output_shape, devices_per_host);
 }
 
 absl::StatusOr<ThunkProto> FftThunk::ToProto() const {
diff --git a/third_party/xla/xla/backends/gpu/runtime/fft_thunk.h b/third_party/xla/xla/backends/gpu/runtime/fft_thunk.h
index 7f8af58d624..b6096724010 100644
--- a/third_party/xla/xla/backends/gpu/runtime/fft_thunk.h
+++ b/third_party/xla/xla/backends/gpu/runtime/fft_thunk.h
@@ -21,11 +21,11 @@ limitations under the License.
 #include <vector>
 
 #include "absl/base/thread_annotations.h"
-#include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk.pb.h"
 #include "xla/runtime/buffer_use.h"
@@ -49,20 +49,15 @@ struct FftPlan {
 
 class FftPlanCache {
  public:
+  explicit FftPlanCache(int devices_per_host) : fft_plans_(devices_per_host) {}
+
   // Returns Fft plan cached for the given device ordinal or creates a new one.
-  FftPlan* GetOrCreate(int device_ordinal) {
-    absl::MutexLock lock(mu_);
-    std::unique_ptr<FftPlan>& plan = fft_plans_[device_ordinal];
-    if (!plan) {
-      plan = std::make_unique<FftPlan>();
-    }
-    return plan.get();
+  absl::StatusOr<FftPlan*> GetOrCreate(int device_ordinal) {
+    return fft_plans_.GetOrCreate(device_ordinal);
   }
 
  private:
-  absl::Mutex mu_;
-  absl::flat_hash_map<int, std::unique_ptr<FftPlan>> fft_plans_
-      ABSL_GUARDED_BY(mu_);
+  PerDeviceState<FftPlan> fft_plans_;
 };
 
 // This class stores everything that StreamExecutor needs to launch an FFT.
@@ -77,7 +72,8 @@ class FftThunk : public Thunk {
            absl::Span<const int64_t> fft_length,
            const BufferAllocation::Slice& input_buffer,
            const BufferAllocation::Slice& output_buffer,
-           const Shape& input_shape, const Shape& output_shape);
+           const Shape& input_shape, const Shape& output_shape,
+           int devices_per_host);
 
   FftThunk(const FftThunk&) = delete;             // Cannot share fft_plan_
   FftThunk& operator=(const FftThunk&) = delete;  // Cannot share fft_plan_
@@ -94,7 +90,8 @@ class FftThunk : public Thunk {
 
   static absl::StatusOr<std::unique_ptr<FftThunk>> FromProto(
       ThunkInfo thunk_info, const FftThunkProto& proto,
-      absl::Span<const BufferAllocation> buffer_allocations);
+      absl::Span<const BufferAllocation> buffer_allocations,
+      int devices_per_host);
 
   absl::StatusOr<ThunkProto> ToProto() const override;
 
diff --git a/third_party/xla/xla/backends/gpu/runtime/fft_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/fft_thunk_test.cc
index 04693f9b3a3..f595231b00d 100644
--- a/third_party/xla/xla/backends/gpu/runtime/fft_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/fft_thunk_test.cc
@@ -73,7 +73,8 @@ TEST(FftThunkTest, ProtoRoundTrip) {
                        Thunk::ThunkInfo::FromProto(proto.thunk_info()));
   ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<FftThunk> thunk,
-      FftThunk::FromProto(thunk_info, proto.fft_thunk(), buffer_allocations));
+      FftThunk::FromProto(thunk_info, proto.fft_thunk(), buffer_allocations,
+                          /*devices_per_host=*/1));
   ASSERT_OK_AND_ASSIGN(ThunkProto round_trip_proto, thunk->ToProto());
   EXPECT_THAT(round_trip_proto, EqualsProto(proto));
 }
diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
index b3aa4e39014..ce5708c8610 100644
--- a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
@@ -203,8 +203,9 @@ absl::StatusOr<std::unique_ptr<Thunk>> DeserializeThunkProtoImpl(
           buffer_allocations);
     }
     case ThunkProto::kFftThunk:
-      return FftThunk::FromProto(std::move(thunk_info), thunk_proto.fft_thunk(),
-                                 buffer_allocations);
+      return FftThunk::FromProto(
+          std::move(thunk_info), thunk_proto.fft_thunk(), buffer_allocations,
+          gpu_topology.has_value() ? gpu_topology->num_devices_per_host() : 0);
     case ThunkProto::kMemset32BitValueThunk:
       return Memset32BitValueThunk::FromProto(
           std::move(thunk_info), thunk_proto.memset32bit_value_thunk(),
diff --git a/third_party/xla/xla/service/gpu/thunk_emitter.cc b/third_party/xla/xla/service/gpu/thunk_emitter.cc
index 77090a3e40a..ae977577ac7 100644
--- a/third_party/xla/xla/service/gpu/thunk_emitter.cc
+++ b/third_party/xla/xla/service/gpu/thunk_emitter.cc
@@ -1419,7 +1419,8 @@ absl::StatusOr<ThunkSequence> ThunkEmitter::EmitFft(
       /*input_buffer=*/arg_slice,
       /*output_buffer=*/dest_slice,
       /*input_shape=*/instr->operand(0)->shape(),
-      /*output_shape=*/instr->shape());
+      /*output_shape=*/instr->shape(),
+      ir_emitter_context_->gpu_topology().num_devices_per_host());
 }
 
 absl::StatusOr<ThunkSequence> ThunkEmitter::EmitTriangularSolveCustomCall(

From 2fc4c93879860fe46812b52fc8d47d441e574d92 Mon Sep 17 00:00:00 2001
From: bodapatisaikrishna <saikrishnabodapati@gmail.com>
Date: Thu, 8 Oct 2026 19:34:43 +0530
Subject: [PATCH 0285/1000] Add scalar and rank-1 cases to conv2d rank tests

---
 .../kernel_tests/nn_ops/conv_ops_test.py      | 26 +++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
index 391c24c4ce2..37e185fe2d7 100644
--- a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
@@ -3186,6 +3186,32 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
           )
       )
 
+    # Ensure scalar and rank-1 inputs fail safely.
+    with self.assertRaisesRegex(
+        (ValueError, errors_impl.InvalidArgumentError),
+        r"(must be 4-dimensional|Shape must be rank 4 but is rank \d+)",
+    ):
+      self.evaluate(
+          raw_ops.Conv2D(
+              input=np.ones([1], dtype=np.float32),
+              filter=np.ones([1, 2, 2, 1], dtype=np.float32),
+              strides=[1, 1, 1, 1],
+              padding="VALID",
+          )
+      )
+    with self.assertRaisesRegex(
+        (ValueError, errors_impl.InvalidArgumentError),
+        r"(must be 4-dimensional|Shape must be rank 4 but is rank \d+)",
+    ):
+      self.evaluate(
+          raw_ops.Conv2D(
+              input=np.ones([1, 1, 3, 3], dtype=np.float32),
+              filter=np.ones([], dtype=np.float32),
+              strides=[1, 1, 1, 1],
+              padding="VALID",
+          )
+      )
+
     # Filter dimensions must be greater than 0.
     with self.assertRaisesRegex(
         errors_impl.InvalidArgumentError, "filter must not have zero elements"

From 3fd29797e9c4df65d7b9a0fff3e528ebe078fd2f Mon Sep 17 00:00:00 2001
From: Blake Hechtman <blakehechtman@google.com>
Date: Thu, 8 Oct 2026 05:50:32 -0700
Subject: [PATCH 0286/1000] Remove obsolete and dead layout assignment
 propagation methods

Delete PropagateWhileLoopLayoutToSubcomputations,
PropagateLayoutsToAsyncSubComputation, PropagateOperandLayoutsToAsyncParameters,
PropagateOperandLayoutToAsyncParameter,
PropagateResultLayoutToAsyncSubComputation,
PropagateLayoutsFromAsyncSubComputation, CheckCallLayout, CheckWhileLayout,
CheckConditionalLayout, CheckFusionLayout, CheckOptimizationBarrierLayout,
RunOnComputation, the single computation overload of InitUnconstrainedBuffers,
ClearComputationLayouts, BuildHostChannelConstraints, ChannelLayoutConstraints,
ConstrainChannelLayouts, AddCrossModuleAllReduceConstraints,
ResetChannelConstraints, IsHostSendRecv, reverse_computation_order, and
unimplemented declarations PropagateUnconstraintedBuffers and
GetInstructionBufferLayoutConstraint in LayoutAssignment. In
LayoutAssignment::CheckLayouts, run HloDataflowAnalysis with
propagate_through_control_flow=true so HloValue positions span while,
conditional, and optimization barrier boundaries. Also add
LayoutAssignment::CanOverrideOperandLayoutOnBufferOverride so
ConstrainOrOverrideBuffersAtIndex does not override smaller concatenate operands
that PropagateOperandConstraint skips.
PiperOrigin-RevId: 995779451
---
 .../gpu/transforms/layout_assignment.h        |    5 +-
 .../xla/xla/service/cpu/cpu_compiler.cc       |    4 +-
 .../xla/service/cpu/cpu_layout_assignment.h   |    5 +-
 .../xla/xla/service/gpu/gpu_compiler.cc       |    3 +-
 .../xla/xla/service/layout_assignment.cc      | 1002 ++---------------
 .../xla/xla/service/layout_assignment.h       |  199 +---
 .../xla/xla/service/layout_assignment_test.cc |  156 ++-
 third_party/xla/xla/tools/hlo_opt/cpu_opt.cc  |    3 +-
 8 files changed, 180 insertions(+), 1197 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment.h b/third_party/xla/xla/backends/gpu/transforms/layout_assignment.h
index 92f2da50bab..4d0b9ee40f6 100644
--- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment.h
+++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment.h
@@ -41,9 +41,8 @@ class GpuLayoutAssignment : public LayoutAssignment {
  public:
   GpuLayoutAssignment(ComputationLayout* entry_computation_layout,
                       const se::GpuComputeCapability& gpu_version,
-                      const se::DeviceDescription& device_description,
-                      ChannelLayoutConstraints* channel_constraints = nullptr)
-      : LayoutAssignment(entry_computation_layout, channel_constraints),
+                      const se::DeviceDescription& device_description)
+      : LayoutAssignment(entry_computation_layout),
         gpu_version_(gpu_version),
         device_description_(device_description) {}
   ~GpuLayoutAssignment() override = default;
diff --git a/third_party/xla/xla/service/cpu/cpu_compiler.cc b/third_party/xla/xla/service/cpu/cpu_compiler.cc
index 505e8665f5f..eb9ca007d8b 100644
--- a/third_party/xla/xla/service/cpu/cpu_compiler.cc
+++ b/third_party/xla/xla/service/cpu/cpu_compiler.cc
@@ -997,10 +997,8 @@ absl::Status CpuCompiler::RunHloPassesThroughLayoutAssn(
     pipeline.AddPass<FlattenCallGraph>();
   }
 
-  ChannelLayoutConstraints layout_constraints;
   pipeline.AddPass<CpuLayoutAssignment>(
-      module->mutable_entry_computation_layout(), target_machine_features,
-      &layout_constraints);
+      module->mutable_entry_computation_layout(), target_machine_features);
   // Run SubByteNormalization because CpuLayoutAssignment may modify a
   // Layout's element_size_in_bits field.
   pipeline.AddPass<SubByteNormalization>(
diff --git a/third_party/xla/xla/service/cpu/cpu_layout_assignment.h b/third_party/xla/xla/service/cpu/cpu_layout_assignment.h
index f3e88c7c96b..2c59fe6b2ff 100644
--- a/third_party/xla/xla/service/cpu/cpu_layout_assignment.h
+++ b/third_party/xla/xla/service/cpu/cpu_layout_assignment.h
@@ -31,9 +31,8 @@ class CpuLayoutAssignment : public LayoutAssignment {
  public:
   explicit CpuLayoutAssignment(
       ComputationLayout* entry_computation_layout,
-      const TargetMachineFeatures* target_machine_features,
-      ChannelLayoutConstraints* channel_constraints = nullptr)
-      : LayoutAssignment(entry_computation_layout, channel_constraints),
+      const TargetMachineFeatures* target_machine_features)
+      : LayoutAssignment(entry_computation_layout),
         target_machine_features_(*target_machine_features) {}
   ~CpuLayoutAssignment() override {}
 
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index f54526c9152..c8a77fcd615 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -1411,10 +1411,9 @@ absl::Status RunLayoutAssignmentPasses(
   // Layout assignment uses alias analysis, which requires the call graph to
   // be flattened.
   pipeline.AddPass<FlattenCallGraph>();
-  ChannelLayoutConstraints layout_constraints;
   pipeline.AddPass<GpuLayoutAssignment>(
       hlo_module->mutable_entry_computation_layout(), gpu_version,
-      device_description, &layout_constraints);
+      device_description);
   // Run SubByteNormalization because GpuLayoutAssignment may modify a
   // Layout's element_size_in_bits field.
   pipeline.AddPass<SubByteNormalization>(
diff --git a/third_party/xla/xla/service/layout_assignment.cc b/third_party/xla/xla/service/layout_assignment.cc
index 74534254bbd..0af92ed55f5 100644
--- a/third_party/xla/xla/service/layout_assignment.cc
+++ b/third_party/xla/xla/service/layout_assignment.cc
@@ -619,43 +619,6 @@ std::string LayoutAssignment::ToString(
 
 namespace {
 
-bool IsHostSendRecv(const HloInstruction* instruction) {
-  const HloSendRecvInstruction* send_recv_instr =
-      DynCast<HloSendRecvInstruction>(instruction);
-  return send_recv_instr != nullptr && send_recv_instr->is_host_transfer();
-}
-
-}  // namespace
-
-absl::Status LayoutAssignment::BuildHostChannelConstraints(
-    HloComputation* computation) {
-  for (auto* instruction : computation->instructions()) {
-    const HloSendRecvInstruction* send_recv_instr =
-        DynCast<HloSendRecvInstruction>(instruction);
-    if (send_recv_instr == nullptr || !send_recv_instr->is_host_transfer()) {
-      continue;
-    }
-
-    // For host transfers the Send and Recv instruction carry the layout.
-    if (instruction->opcode() == HloOpcode::kSend ||
-        instruction->opcode() == HloOpcode::kRecv) {
-      const Shape& data_shape =
-          ShapeUtil::GetTupleElementShape(send_recv_instr->shape(), 0);
-      TF_RET_CHECK(data_shape.IsArray());
-      TF_RET_CHECK(LayoutUtil::HasLayout(data_shape));
-      const Layout* prev_layout = host_channel_constraints_.ConstrainChannel(
-          *send_recv_instr->channel_id(), data_shape.layout());
-      TF_RET_CHECK(prev_layout == nullptr)
-          << "Cannot constrain host transfer layout as it was set to "
-          << LayoutUtil::HumanString(*prev_layout) << ": "
-          << send_recv_instr->ToString();
-    }
-  }
-  return absl::OkStatus();
-}
-
-namespace {
-
 bool IsLayoutConstrainedCustomCall(HloInstruction* instruction) {
   const HloCustomCallInstruction* custom_call =
       DynCast<HloCustomCallInstruction>(instruction);
@@ -741,8 +704,7 @@ absl::Status LayoutAssignment::AddParameterConstraints(
       ABSL_RETURN_IF_ERROR(ResetMemorySpaceInLayout(parameter_layout));
       Shape param_shape = parameter_layout.shape();
       ABSL_RETURN_IF_ERROR(SetInstructionLayout(param_shape, instruction));
-      if (reverse_computation_order_ ||
-          copy_disabled_while_computations_.contains(
+      if (copy_disabled_while_computations_.contains(
               constraints->computation())) {
         ABSL_RETURN_IF_ERROR(
             PropagateParameterLayoutToUsers(instruction, param_shape, this));
@@ -774,36 +736,7 @@ absl::Status LayoutAssignment::AddCollectiveConstraints(
   return absl::OkStatus();
 }
 
-absl::Status LayoutAssignment::AddCrossModuleAllReduceConstraints(
-    HloInstruction* instruction,
-    ChannelLayoutConstraints* channel_constraints) {
-  auto get_channel_constraints = [&](const HloInstruction* inst) {
-    return IsHostSendRecv(inst) ? &host_channel_constraints_
-                                : channel_constraints;
-  };
-  CHECK(get_channel_constraints(instruction))
-      << "Multi-module layout assignment requires ChannelLayoutConstraints";
-  int64_t channel_id = instruction->channel_id().value();
-  if (!get_channel_constraints(instruction)->IsChannelConstrained(channel_id)) {
-    return absl::OkStatus();
-  }
-  // TODO(b/68493863): Change to use SetOperandLayout().
-  const Shape& buffer_shape = instruction->operand(0)->shape();
-  TF_RET_CHECK(buffer_shape.IsArray());
-  if (!LayoutUtil::ValidateLayoutForShape(
-           get_channel_constraints(instruction)->LayoutForChannel(channel_id),
-           buffer_shape)
-           .ok()) {
-    return absl::OkStatus();
-  }
-  Shape new_buffer_shape =
-      get_channel_constraints(instruction)
-          ->LayoutShapeForChannel(buffer_shape, channel_id);
-  return SetInstructionLayout(new_buffer_shape, instruction);
-}
-
 absl::Status LayoutAssignment::AddInstructionLayoutConstraints(
-    ChannelLayoutConstraints* channel_constraints,
     LayoutConstraints* constraints) {
   // Constrain layouts of instructions which define values with pre-existing
   // layouts.
@@ -821,12 +754,6 @@ absl::Status LayoutAssignment::AddInstructionLayoutConstraints(
       default:
         if (IsLayoutConstrainedCollective(instruction)) {
           ABSL_RETURN_IF_ERROR(AddCollectiveConstraints(instruction));
-        } else if (instruction->IsCrossModuleAllReduce() &&
-                   !instruction->GetModule()
-                        ->config()
-                        .use_spmd_partitioning()) {
-          ABSL_RETURN_IF_ERROR(AddCrossModuleAllReduceConstraints(
-              instruction, channel_constraints));
         }
         break;
     }
@@ -862,30 +789,18 @@ absl::Status LayoutAssignment::AddCallConstraints(HloInstruction* instruction) {
 }
 
 absl::Status LayoutAssignment::AddWhileConstraints(
-    HloInstruction* instruction, LayoutConstraints* constraints) {
+    HloInstruction* instruction) {
   if (!HasConstrainedComputationLayout(instruction->while_body())) {
     return absl::OkStatus();
   }
   HloComputation* body = instruction->while_body();
   HloComputation* condition = instruction->while_condition();
-  const HloInstruction* init = instruction->operand(0);
-  ComputationLayoutConstraint* body_constraint =
-      mutable_computation_constraints(body)->mutable_computation_constraint();
-  ComputationLayout body_layout = body_constraint->computation_layout();
-  ComputationLayoutConstraint* condition_constraint =
-      mutable_computation_constraints(condition)
-          ->mutable_computation_constraint();
-  ComputationLayout condition_layout =
-      condition_constraint->computation_layout();
+  const ComputationLayout& body_layout =
+      computation_constraints(body).computation_layout();
 
   CHECK_EQ(1, instruction->operand_count());
   CHECK_EQ(1, body->num_parameters());
   CHECK_EQ(1, condition->num_parameters());
-  DCHECK(ShapeUtil::Compatible(body_layout.result_shape(),
-                               body_layout.parameter_shape(0)));
-  DCHECK(ShapeUtil::Compatible(body_layout.result_shape(),
-                               condition_layout.parameter_shape(0)));
-  DCHECK(ShapeUtil::Compatible(body_layout.result_shape(), init->shape()));
 
   const int64_t reset_priority = current_priority_ + kNumberOfPropagationRounds;
   auto sync_while_subcomputation =
@@ -895,8 +810,6 @@ absl::Status LayoutAssignment::AddWhileConstraints(
         mutable_computation_constraints(subcomp)
             ->mutable_computation_constraint();
     ComputationLayout subcomp_layout = subcomp_constraint->computation_layout();
-    const bool is_copy_disabled =
-        copy_disabled_while_computations_.contains(subcomp);
     if (subcomp_layout.parameter_layout(0) != target_param_layout) {
       *subcomp_layout.mutable_parameter_layout(0) = target_param_layout;
       subcomp_constraint->ResetComputationLayout(
@@ -910,15 +823,10 @@ absl::Status LayoutAssignment::AddWhileConstraints(
                                /*allow_alias=*/false, reset_priority));
       ABSL_RETURN_IF_ERROR(
           PropagateConstraints(mutable_computation_constraints(subcomp)));
-      if (is_copy_disabled) {
+      if (copy_disabled_while_computations_.contains(subcomp)) {
         while_layout_changed_ = true;
       }
     }
-    if (is_copy_disabled) {
-      PropagateWhileLoopLayoutToSubcomputations(
-          subcomp, subcomp_layout.parameter_layout(0).shape(),
-          &subcomp_layout.result_layout().shape(), reset_priority);
-    }
     return absl::OkStatus();
   };
 
@@ -932,372 +840,6 @@ absl::Status LayoutAssignment::AddWhileConstraints(
   return SetInstructionLayout(body_layout.result_shape(), instruction);
 }
 
-void LayoutAssignment::PropagateWhileLoopLayoutToSubcomputations(
-    HloComputation* computation, const Shape& param_shape,
-    const Shape* result_shape, int64_t priority) {
-  if (computation == nullptr || computation->num_parameters() == 0) {
-    return;
-  }
-
-  auto HasAnyLayout = [](const Shape& shape) -> bool {
-    bool has = false;
-    ShapeUtil::ForEachSubshape(
-        shape, [&has](const Shape& subshape, const ShapeIndex&) {
-          if (subshape.has_layout()) {
-            has = true;
-          }
-        });
-    return has;
-  };
-
-  absl::flat_hash_map<const HloInstruction*, Shape> instruction_layouts;
-  if (HasAnyLayout(param_shape) &&
-      ShapeUtil::Compatible(computation->parameter_instruction(0)->shape(),
-                            param_shape)) {
-    instruction_layouts[computation->parameter_instruction(0)] = param_shape;
-  }
-
-  // Forward propagation from parameter through the computation.
-  for (HloInstruction* instr : computation->MakeInstructionPostOrder()) {
-    if (instr == computation->parameter_instruction(0)) {
-      continue;
-    }
-    switch (instr->opcode()) {
-      case HloOpcode::kGetTupleElement: {
-        auto it = instruction_layouts.find(instr->operand(0));
-        if (it != instruction_layouts.end() && it->second.IsTuple() &&
-            instr->tuple_index() < ShapeUtil::TupleElementCount(it->second)) {
-          Shape subshape =
-              ShapeUtil::GetTupleElementShape(it->second, instr->tuple_index());
-          if (ShapeUtil::Compatible(instr->shape(), subshape)) {
-            instruction_layouts[instr] = std::move(subshape);
-          }
-        }
-        break;
-      }
-      case HloOpcode::kTuple: {
-        std::vector<Shape> subshapes;
-        subshapes.reserve(instr->operand_count());
-        bool any_known = false;
-        for (int64_t i = 0; i < instr->operand_count(); ++i) {
-          auto it = instruction_layouts.find(instr->operand(i));
-          if (it != instruction_layouts.end() && HasAnyLayout(it->second) &&
-              ShapeUtil::Compatible(instr->operand(i)->shape(), it->second)) {
-            subshapes.push_back(it->second);
-            any_known = true;
-          } else {
-            Shape cleared = instr->operand(i)->shape();
-            LayoutUtil::ClearLayout(&cleared);
-            subshapes.push_back(std::move(cleared));
-          }
-        }
-        if (any_known) {
-          Shape tuple_shape = ShapeUtil::MakeTupleShape(subshapes);
-          if (ShapeUtil::Compatible(instr->shape(), tuple_shape)) {
-            instruction_layouts[instr] = std::move(tuple_shape);
-          }
-        }
-        break;
-      }
-      case HloOpcode::kSelect: {
-        Shape select_shape;
-        bool has_select_shape = false;
-        auto it_true = instruction_layouts.find(instr->operand(1));
-        auto it_false = instruction_layouts.find(instr->operand(2));
-        if (it_true != instruction_layouts.end() &&
-            HasAnyLayout(it_true->second) &&
-            ShapeUtil::Compatible(instr->shape(), it_true->second)) {
-          select_shape = it_true->second;
-          has_select_shape = true;
-        } else if (it_false != instruction_layouts.end() &&
-                   HasAnyLayout(it_false->second) &&
-                   ShapeUtil::Compatible(instr->shape(), it_false->second)) {
-          select_shape = it_false->second;
-          has_select_shape = true;
-        }
-        if (has_select_shape) {
-          instruction_layouts[instr] = std::move(select_shape);
-        }
-        break;
-      }
-      case HloOpcode::kCopy: {
-        auto it = instruction_layouts.find(instr->operand(0));
-        if (it != instruction_layouts.end() && HasAnyLayout(it->second) &&
-            ShapeUtil::Compatible(instr->shape(), it->second)) {
-          Shape copy_shape = it->second;
-          instruction_layouts[instr] = std::move(copy_shape);
-        }
-        break;
-      }
-      case HloOpcode::kConditional: {
-        Shape branch_result_shape;
-        bool has_branch_result_shape = false;
-        for (int j = 0; j < instr->branch_count(); ++j) {
-          HloInstruction* branch_arg = instr->mutable_operand(j + 1);
-          HloComputation* branch_comp = instr->branch_computation(j);
-          auto it = instruction_layouts.find(branch_arg);
-          if (it != instruction_layouts.end() && HasAnyLayout(it->second) &&
-              !ShapeUtil::IsEmptyTuple(it->second) &&
-              ShapeUtil::Compatible(branch_arg->shape(), it->second)) {
-            PropagateWhileLoopLayoutToSubcomputations(
-                branch_comp, it->second, /*result_shape=*/nullptr, priority);
-            auto* branch_constraints =
-                mutable_computation_constraints(branch_comp);
-            if (branch_constraints != nullptr &&
-                branch_constraints->computation_constraint()
-                    .result_layout_is_set()) {
-              const Shape& res = branch_constraints->computation_layout()
-                                     .result_layout()
-                                     .shape();
-              if (ShapeUtil::Compatible(instr->shape(), res)) {
-                branch_result_shape = res;
-                has_branch_result_shape = true;
-              }
-            }
-          }
-        }
-        if (!has_branch_result_shape) {
-          auto it = instruction_layouts.find(instr);
-          if (it != instruction_layouts.end() && HasAnyLayout(it->second) &&
-              ShapeUtil::Compatible(instr->shape(), it->second)) {
-            branch_result_shape = it->second;
-            has_branch_result_shape = true;
-          }
-        }
-        if (has_branch_result_shape) {
-          instruction_layouts[instr] = branch_result_shape;
-          for (int j = 0; j < instr->branch_count(); ++j) {
-            HloComputation* branch_comp = instr->branch_computation(j);
-            auto* branch_constraints =
-                mutable_computation_constraints(branch_comp);
-            if (branch_constraints != nullptr) {
-              ComputationLayout branch_layout =
-                  branch_constraints->computation_layout();
-              bool prop_branch_result = false;
-              ShapeUtil::ForEachSubshape(
-                  branch_result_shape,
-                  [&](const Shape& subshape, const ShapeIndex& index) {
-                    if (subshape.IsArray() && subshape.has_layout()) {
-                      branch_layout.mutable_result_layout()->ResetLayout(
-                          subshape.layout(), index);
-                      prop_branch_result = true;
-                    }
-                  });
-              if (prop_branch_result) {
-                branch_constraints->mutable_computation_constraint()
-                    ->ResetComputationLayout(
-                        branch_layout, priority,
-                        /*prop_result_layout=*/true,
-                        /*prop_parameter_layout=*/
-                        branch_constraints->computation_constraint()
-                            .parameter_layout_is_set());
-              }
-            }
-          }
-        }
-        break;
-      }
-      case HloOpcode::kWhile: {
-        HloInstruction* init = instr->mutable_operand(0);
-        auto it = instruction_layouts.find(init);
-        if (it != instruction_layouts.end() && HasAnyLayout(it->second) &&
-            ShapeUtil::Compatible(init->shape(), it->second)) {
-          PropagateWhileLoopLayoutToSubcomputations(
-              instr->while_body(), it->second, &it->second, priority);
-          PropagateWhileLoopLayoutToSubcomputations(
-              instr->while_condition(), it->second, /*result_shape=*/nullptr,
-              priority);
-        }
-        break;
-      }
-      case HloOpcode::kCall: {
-        HloComputation* callee = instr->to_apply();
-        if (callee != nullptr &&
-            callee->num_parameters() == instr->operand_count()) {
-          auto* callee_constraints = mutable_computation_constraints(callee);
-          if (callee_constraints != nullptr) {
-            ComputationLayout callee_layout =
-                callee_constraints->computation_layout();
-            bool any_param_set = false;
-            for (int64_t i = 0; i < instr->operand_count(); ++i) {
-              auto it = instruction_layouts.find(instr->operand(i));
-              if (it != instruction_layouts.end() && HasAnyLayout(it->second) &&
-                  ShapeUtil::Compatible(instr->operand(i)->shape(),
-                                        it->second)) {
-                *callee_layout.mutable_parameter_layout(i) =
-                    ShapeLayout(it->second);
-                any_param_set = true;
-              }
-            }
-            if (any_param_set) {
-              callee_constraints->mutable_computation_constraint()
-                  ->ResetComputationLayout(callee_layout, priority,
-                                           /*prop_result_layout=*/false,
-                                           /*prop_parameter_layout=*/true);
-            }
-          }
-        }
-        break;
-      }
-      default:
-        break;
-    }
-  }
-
-  // Backward propagation from result_shape if provided.
-  if (result_shape != nullptr && HasAnyLayout(*result_shape) &&
-      ShapeUtil::Compatible(computation->root_instruction()->shape(),
-                            *result_shape)) {
-    absl::flat_hash_map<const HloInstruction*, Shape> root_layouts;
-    root_layouts[computation->root_instruction()] = *result_shape;
-    auto post_order = computation->MakeInstructionPostOrder();
-    for (auto it = post_order.rbegin(); it != post_order.rend(); ++it) {
-      HloInstruction* instr = *it;
-      auto root_it = root_layouts.find(instr);
-      if (root_it == root_layouts.end() || !HasAnyLayout(root_it->second)) {
-        continue;
-      }
-      Shape current_layout_shape = root_it->second;
-      switch (instr->opcode()) {
-        case HloOpcode::kTuple: {
-          if (current_layout_shape.IsTuple()) {
-            for (int64_t i = 0; i < instr->operand_count(); ++i) {
-              if (i < ShapeUtil::TupleElementCount(current_layout_shape)) {
-                Shape elem_shape =
-                    ShapeUtil::GetTupleElementShape(current_layout_shape, i);
-                if (ShapeUtil::Compatible(instr->operand(i)->shape(),
-                                          elem_shape)) {
-                  root_layouts[instr->operand(i)] = std::move(elem_shape);
-                }
-              }
-            }
-          }
-          break;
-        }
-        case HloOpcode::kSelect: {
-          if (ShapeUtil::Compatible(instr->operand(1)->shape(),
-                                    current_layout_shape)) {
-            root_layouts[instr->operand(1)] = current_layout_shape;
-          }
-          if (ShapeUtil::Compatible(instr->operand(2)->shape(),
-                                    current_layout_shape)) {
-            root_layouts[instr->operand(2)] = current_layout_shape;
-          }
-          break;
-        }
-        case HloOpcode::kCopy: {
-          if (ShapeUtil::Compatible(instr->operand(0)->shape(),
-                                    current_layout_shape)) {
-            root_layouts[instr->operand(0)] = current_layout_shape;
-          }
-          break;
-        }
-        case HloOpcode::kGetTupleElement: {
-          HloInstruction* gte_op = instr->mutable_operand(0);
-          if (gte_op->opcode() == HloOpcode::kConditional &&
-              instr->tuple_index() <
-                  ShapeUtil::TupleElementCount(gte_op->shape())) {
-            auto cond_it = root_layouts.find(gte_op);
-            if (cond_it == root_layouts.end()) {
-              root_layouts[gte_op] = gte_op->shape();
-            }
-            if (root_layouts[gte_op].IsTuple()) {
-              *ShapeUtil::GetMutableSubshape(&root_layouts[gte_op],
-                                             {instr->tuple_index()}) =
-                  current_layout_shape;
-            }
-          }
-          break;
-        }
-        case HloOpcode::kConditional: {
-          if (ShapeUtil::Compatible(instr->shape(), current_layout_shape)) {
-            for (int j = 0; j < instr->branch_count(); ++j) {
-              HloComputation* branch_comp = instr->branch_computation(j);
-              auto* branch_constraints =
-                  mutable_computation_constraints(branch_comp);
-              if (branch_constraints != nullptr) {
-                ComputationLayout branch_layout =
-                    branch_constraints->computation_layout();
-                bool prop_branch_result = false;
-                ShapeUtil::ForEachSubshape(
-                    current_layout_shape,
-                    [&](const Shape& subshape, const ShapeIndex& index) {
-                      if (subshape.IsArray() && subshape.has_layout()) {
-                        branch_layout.mutable_result_layout()->ResetLayout(
-                            subshape.layout(), index);
-                        prop_branch_result = true;
-                      }
-                    });
-                if (prop_branch_result) {
-                  branch_constraints->mutable_computation_constraint()
-                      ->ResetComputationLayout(
-                          branch_layout, priority,
-                          /*prop_result_layout=*/true,
-                          /*prop_parameter_layout=*/
-                          branch_constraints->computation_constraint()
-                              .parameter_layout_is_set());
-                }
-              }
-            }
-          }
-          break;
-        }
-        default:
-          break;
-      }
-    }
-  }
-
-  // Update this computation's own ComputationLayout.
-  auto* comp_constraints = mutable_computation_constraints(computation);
-  if (comp_constraints != nullptr) {
-    ComputationLayout comp_layout = comp_constraints->computation_layout();
-    bool prop_param = false;
-    if (HasAnyLayout(param_shape) &&
-        ShapeUtil::Compatible(computation->parameter_instruction(0)->shape(),
-                              param_shape)) {
-      ShapeUtil::ForEachSubshape(
-          param_shape, [&](const Shape& subshape, const ShapeIndex& index) {
-            if (subshape.IsArray() && subshape.has_layout()) {
-              comp_layout.mutable_parameter_layout(0)->ResetLayout(
-                  subshape.layout(), index);
-              prop_param = true;
-            }
-          });
-    }
-    bool prop_result = false;
-    auto it = instruction_layouts.find(computation->root_instruction());
-    if (it != instruction_layouts.end() && HasAnyLayout(it->second) &&
-        ShapeUtil::Compatible(computation->root_instruction()->shape(),
-                              it->second)) {
-      ShapeUtil::ForEachSubshape(
-          it->second, [&](const Shape& subshape, const ShapeIndex& index) {
-            if (subshape.IsArray() && subshape.has_layout()) {
-              comp_layout.mutable_result_layout()->ResetLayout(
-                  subshape.layout(), index);
-              prop_result = true;
-            }
-          });
-    } else if (result_shape != nullptr && HasAnyLayout(*result_shape) &&
-               ShapeUtil::Compatible(computation->root_instruction()->shape(),
-                                     *result_shape)) {
-      ShapeUtil::ForEachSubshape(
-          *result_shape, [&](const Shape& subshape, const ShapeIndex& index) {
-            if (subshape.IsArray() && subshape.has_layout()) {
-              comp_layout.mutable_result_layout()->ResetLayout(
-                  subshape.layout(), index);
-              prop_result = true;
-            }
-          });
-    }
-    if (prop_param || prop_result) {
-      comp_constraints->mutable_computation_constraint()
-          ->ResetComputationLayout(comp_layout, priority, prop_result,
-                                   prop_param);
-    }
-  }
-}
-
 namespace {
 
 int64_t FindLargestConditionalBranch(const HloInstruction* conditional) {
@@ -1369,141 +911,16 @@ absl::Status LayoutAssignment::AddConditionalConstraints(
                               /*allow_alias=*/false);
 }
 
-// Propagates array layouts for a single operand `param_idx` of `instruction`
-// to the corresponding parameter layout in `async_layout`.
-// `instruction` operand `param_idx` is mapped to the parameter `param_idx`
-// of the async sub-computation.
-// Updates `async_layout` in-place and returns true if the layout was updated.
-bool LayoutAssignment::PropagateOperandLayoutToAsyncParameter(
-    const HloInstruction* instruction, int64_t param_idx,
-    ComputationLayout* async_layout) {
-  if (param_idx >= async_layout->parameter_count() ||
-      !async_layout->parameter_layout(param_idx).MinorToMajorInLayoutIsSet()) {
-    return false;
+absl::Status LayoutAssignment::AddAsyncStartConstraints(
+    HloInstruction* instruction) {
+  HloComputation* async_comp = instruction->async_wrapped_computation();
+  if (async_comp == nullptr || !HasConstrainedComputationLayout(async_comp)) {
+    return absl::OkStatus();
   }
-
-  Shape param_shape = instruction->operand(param_idx)->shape();
-  bool param_reset = false;
-  ShapeUtil::ForEachSubshape(param_shape, [&](const Shape& subshape,
-                                              const ShapeIndex& index) {
-    if (!subshape.IsArray()) {
-      return;
-    }
-    auto operand_layout =
-        InferArrayLayout(instruction->operand(param_idx), index);
-    const Shape& async_subshape = ShapeUtil::GetSubshape(
-        async_layout->parameter_layout(param_idx).shape(), index);
-    if (operand_layout.ok() && async_subshape.has_layout() &&
-        async_subshape.layout() != operand_layout.value()) {
-      *ShapeUtil::GetMutableSubshape(&param_shape, index)->mutable_layout() =
-          operand_layout.value();
-      param_reset = true;
-    }
-  });
-  if (param_reset) {
-    *async_layout->mutable_parameter_layout(param_idx) =
-        ShapeLayout(param_shape);
-    return true;
-  }
-  return false;
-}
-
-// Propagates the operand array layouts of `instruction` to the parameter
-// layouts defined in `async_layout`.
-// Updates `async_layout` in-place and returns true if any parameter layout
-// was changed.
-bool LayoutAssignment::PropagateOperandLayoutsToAsyncParameters(
-    const HloInstruction* instruction, ComputationLayout* async_layout) {
-  bool reset_needed = false;
-  for (int64_t i = 0; i < instruction->operand_count(); ++i) {
-    if (PropagateOperandLayoutToAsyncParameter(instruction, i, async_layout)) {
-      reset_needed = true;
-    }
-  }
-  return reset_needed;
-}
-
-// Propagates array layouts from the result shape of `instruction` (tuple
-// element 1) to `async_layout`'s result layout.
-// We assume the result shape of the async operation is at index {1} of the
-// `instruction` (async start/update) output tuple.
-// Updates `async_layout` in-place and returns true if the result layout was
-// updated.
-bool LayoutAssignment::PropagateResultLayoutToAsyncSubComputation(
-    const HloInstruction* instruction, ComputationLayout* async_layout) {
-  if (!instruction->shape().IsTuple() ||
-      instruction->shape().tuple_shapes().size() <= 1 ||
-      !async_layout->result_layout().MinorToMajorInLayoutIsSet()) {
-    return false;
-  }
-  Shape result_shape = instruction->shape().tuple_shapes(1);
-  bool subshape_reset = false;
-  ShapeUtil::ForEachSubshape(result_shape, [&](const Shape& subshape,
-                                               const ShapeIndex& index) {
-    if (!subshape.IsArray()) {
-      return;
-    }
-    ShapeIndex full_index = index;
-    full_index.push_front(1);
-
-    auto result_layout = InferArrayLayout(instruction, full_index);
-    const Shape& async_subshape =
-        ShapeUtil::GetSubshape(async_layout->result_layout().shape(), index);
-    if (result_layout.ok() && async_subshape.has_layout() &&
-        async_subshape.layout() != result_layout.value()) {
-      *ShapeUtil::GetMutableSubshape(&result_shape, index)->mutable_layout() =
-          result_layout.value();
-      subshape_reset = true;
-    }
-  });
-  if (subshape_reset) {
-    *async_layout->mutable_result_layout() = ShapeLayout(result_shape);
-    return true;
-  }
-  return false;
-}
-
-// Propagates layout constraints from the caller instruction into the inner
-// async sub-computation.
-// This is the forward propagation step: it takes the layouts of the operands
-// and result of the async start/update instruction (which are in the parent
-// computation) and propagates them to the parameters and result of the
-// async sub-computation.
-// If any layout in the sub-computation is updated, it resets the
-// sub-computation layout with an elevated priority to ensure it is respected
-// during the sub-computation's layout assignment. Returns the
-// sub-computation layout.
-ComputationLayout LayoutAssignment::PropagateLayoutsToAsyncSubComputation(
-    const HloInstruction* instruction, LayoutConstraints* async_constraint) {
-  ComputationLayout async_layout = async_constraint->computation_layout();
-  bool param_reset =
-      PropagateOperandLayoutsToAsyncParameters(instruction, &async_layout);
-  bool result_reset =
-      PropagateResultLayoutToAsyncSubComputation(instruction, &async_layout);
-
-  if (param_reset || result_reset) {
-    bool prop_result =
-        result_reset ||
-        async_constraint->computation_constraint().result_layout_is_set();
-    bool prop_param =
-        param_reset ||
-        async_constraint->computation_constraint().parameter_layout_is_set();
-    async_constraint->mutable_computation_constraint()->ResetComputationLayout(
-        async_layout, current_priority_ + kNumberOfPropagationRounds,
-        /*prop_result_layout=*/prop_result,
-        /*prop_parameter_layout=*/prop_param);
-  }
-  return async_layout;
-}
-
-// Propagates async sub-computation parameter and result layout constraints
-// back onto the caller instruction and its operands in the parent computation.
-// This is the backward propagation step: it takes the resolved layouts from the
-// async sub-computation and applies them as mandatory constraints on the caller
-// instruction's shape (at index {1} for result) and its operands.
-absl::Status LayoutAssignment::PropagateLayoutsFromAsyncSubComputation(
-    HloInstruction* instruction, const ComputationLayout& async_layout,
-    LayoutConstraints* async_constraint) {
+  LayoutConstraints* async_constraint =
+      mutable_computation_constraints(async_comp);
+  const ComputationLayout& async_layout =
+      async_constraint->computation_layout();
   for (int64_t i = 0; i < instruction->operand_count(); ++i) {
     if (async_layout.parameter_layout(i).MinorToMajorInLayoutIsSet()) {
       ABSL_RETURN_IF_ERROR(SetOperandLayout(async_layout.parameter_layout(i).shape(),
@@ -1542,30 +959,6 @@ absl::Status LayoutAssignment::PropagateLayoutsFromAsyncSubComputation(
   return absl::OkStatus();
 }
 
-absl::Status LayoutAssignment::AddAsyncStartConstraints(
-    HloInstruction* instruction) {
-  HloComputation* async_comp = instruction->async_wrapped_computation();
-  if (async_comp == nullptr) {
-    return absl::OkStatus();
-  }
-  auto it = computation_layouts_.find(async_comp);
-  if (it == computation_layouts_.end()) {
-    return absl::OkStatus();
-  }
-  LayoutConstraints* async_comp_constraint = it->second.get();
-
-  // Step 1: Caller -> Sub-computation propagation.
-  ComputationLayout async_layout =
-      PropagateLayoutsToAsyncSubComputation(instruction, async_comp_constraint);
-
-  // Step 2: Sub-computation -> Caller propagation.
-  if (!HasConstrainedComputationLayout(async_comp)) {
-    return absl::OkStatus();
-  }
-  return PropagateLayoutsFromAsyncSubComputation(instruction, async_layout,
-                                                 async_comp_constraint);
-}
-
 absl::Status LayoutAssignment::AddAsyncDoneConstraints(
     HloInstruction* instruction, LayoutConstraints* constraints) {
   HloInstruction* start = instruction->async_chain_start();
@@ -1636,7 +1029,7 @@ absl::Status LayoutAssignment::AddSubcomputationLayoutConstraints(
         ABSL_RETURN_IF_ERROR(AddCallConstraints(instruction));
         break;
       case HloOpcode::kWhile:
-        ABSL_RETURN_IF_ERROR(AddWhileConstraints(instruction, constraints));
+        ABSL_RETURN_IF_ERROR(AddWhileConstraints(instruction));
         break;
       case HloOpcode::kConditional:
         ABSL_RETURN_IF_ERROR(AddConditionalConstraints(instruction));
@@ -1665,8 +1058,7 @@ absl::Status LayoutAssignment::AddComputationResultLayoutConstraints(
             .result_layout()
             .shape(),
         current_priority_ + kNumberOfPropagationRounds));
-  } else if (reverse_computation_order_ ||
-             (constraints->computation()->IsEntryComputation() &&
+  } else if ((constraints->computation()->IsEntryComputation() &&
               entry_computation_layout_->AnyMinorToMajorInLayoutIsSet() &&
               entry_computation_layout_->result_layout()
                   .AnyMinorToMajorInLayoutIsSet()) ||
@@ -1686,15 +1078,13 @@ absl::Status LayoutAssignment::AddComputationResultLayoutConstraints(
 }
 
 absl::Status LayoutAssignment::AddMandatoryConstraints(
-    ChannelLayoutConstraints* channel_constraints,
     LayoutConstraints* constraints) {
   VLOG(2) << "Adding mandatory layout constraints to computation "
           << constraints->computation()->name();
 
   // Phase 1: Add constraints for instructions with pre-existing / fixed
   // layouts.
-  ABSL_RETURN_IF_ERROR(
-      AddInstructionLayoutConstraints(channel_constraints, constraints));
+  ABSL_RETURN_IF_ERROR(AddInstructionLayoutConstraints(constraints));
 
   // Phase 2: Add constraints for instructions interacting with
   // sub-computations.
@@ -1733,85 +1123,6 @@ absl::Status CheckCustomCallLayout(HloInstruction* instruction) {
   return absl::OkStatus();
 }
 
-// For a while instruction, all the following layouts must be the same:
-//   (1) init operand
-//   (2) condition computation parameter
-//   (3) body computation parameter
-//   (4) body computation result
-//   (5) while instruction result
-absl::Status CheckWhileLayout(
-    HloInstruction* while_inst,
-    const ComputationLayout& condition_computation_layout,
-    const ComputationLayout& body_computation_layout) {
-  auto init_shape = while_inst->operand(0)->shape();
-  TF_RET_CHECK(
-      condition_computation_layout.parameter_layout(0).MatchesLayoutInShape(
-          init_shape, /*minor_to_major_only=*/true))
-      << " while_inst=" << while_inst->name() << " cond param layout="
-      << condition_computation_layout.parameter_layout(0).ToString()
-      << " init_shape=" << init_shape.ToString();
-  TF_RET_CHECK(body_computation_layout.parameter_layout(0).MatchesLayoutInShape(
-      init_shape, /*minor_to_major_only=*/true))
-      << " while_inst=" << while_inst->name() << " body param layout="
-      << body_computation_layout.parameter_layout(0).ToString()
-      << " init_shape=" << init_shape.ToString();
-  TF_RET_CHECK(body_computation_layout.result_layout().MatchesLayoutInShape(
-      init_shape, /*minor_to_major_only=*/true))
-      << " while_inst=" << while_inst->name() << " body result layout="
-      << body_computation_layout.result_layout().ToString()
-      << " init_shape=" << init_shape.ToString();
-  TF_RET_CHECK(LayoutsInShapesEqual(init_shape, while_inst->shape()))
-      << " while_inst=" << while_inst->name()
-      << " while shape=" << while_inst->shape().ToString()
-      << " init_shape=" << init_shape.ToString();
-  return absl::OkStatus();
-}
-
-absl::Status CheckOptimizationBarrierLayout(HloInstruction* inst) {
-  TF_RET_CHECK(LayoutsInShapesEqual(inst->operand(0)->shape(), inst->shape()));
-  return absl::OkStatus();
-}
-
-absl::Status CheckConditionalLayout(
-    HloInstruction* instruction,
-    absl::Span<const ComputationLayout> branch_computation_layouts) {
-  for (int j = 0; j < instruction->branch_count(); ++j) {
-    const HloInstruction* branch_operand = instruction->operand(j + 1);
-    TF_RET_CHECK(
-        branch_computation_layouts[0].result_layout().MatchesLayoutInShape(
-            branch_computation_layouts[j].result_layout().shape(),
-            /*minor_to_major_only=*/true));
-    TF_RET_CHECK(
-        branch_computation_layouts[j].result_layout().MatchesLayoutInShape(
-            instruction->shape(), /*minor_to_major_only=*/true));
-    TF_RET_CHECK(
-        branch_computation_layouts[j].result_layout().MatchesLayoutInShape(
-            instruction->branch_computation(j)->root_instruction()->shape(),
-            /*minor_to_major_only=*/true))
-        << j << ":"
-        << instruction->branch_computation(j)->root_instruction()->ToString();
-    TF_RET_CHECK(
-        branch_computation_layouts[j].parameter_layout(0).MatchesLayoutInShape(
-            branch_operand->shape(), /*minor_to_major_only=*/true));
-  }
-  return absl::OkStatus();
-}
-
-// Fusion parameters must match the layout of the fusion instructions operands,
-// and the root of the fusion expression must match the layout of the fusion
-// instruction.
-absl::Status CheckFusionLayout(HloInstruction* fusion) {
-  TF_RET_CHECK(HloOpcode::kFusion == fusion->opcode());
-
-  TF_RET_CHECK(LayoutsInShapesEqual(fusion->shape(),
-                                    fusion->fused_expression_root()->shape()));
-  for (int64_t i = 0; i < fusion->operand_count(); ++i) {
-    TF_RET_CHECK(LayoutsInShapesEqual(fusion->fused_parameter(i)->shape(),
-                                      fusion->operand(i)->shape()));
-  }
-  return absl::OkStatus();
-}
-
 // The layout of a parameter must match the respective layout in the
 // computation's ComputationLayout.
 absl::Status CheckParameterLayout(HloInstruction* parameter,
@@ -1879,20 +1190,6 @@ absl::Status CheckBroadcastLayout(HloInstruction* broadcast) {
 
 }  // namespace
 
-absl::Status LayoutAssignment::CheckCallLayout(
-    HloInstruction* call, const ComputationLayout& computation_layout) {
-  HloComputation* computation = call->to_apply();
-  TF_RET_CHECK(computation->num_parameters() == call->operand_count());
-  for (int64_t i = 0; i < computation->num_parameters(); ++i) {
-    TF_RET_CHECK(computation_layout.parameter_layout(i).MatchesLayoutInShape(
-        ShardedShape(call, call->operand(i)->shape(), i),
-        /*minor_to_major_only=*/true));
-  }
-  TF_RET_CHECK(computation_layout.result_layout().MatchesLayoutInShape(
-      ShardedShape(call, call->shape(), -1), /*minor_to_major_only=*/true));
-  return absl::OkStatus();
-}
-
 absl::StatusOr<HloInstruction*> LayoutAssignment::CreateCopyWithNewLayout(
     const Shape& shape_with_layout, HloInstruction* instruction) {
   TF_RET_CHECK(LayoutUtil::HasLayout(shape_with_layout));
@@ -2016,8 +1313,7 @@ absl::Status LayoutAssignment::CopyOperandIfLayoutsDiffer(
 
   // If the operand is only used by a conditional, do the copy inside the branch
   // to avoid overhead for other branches.
-  if (!reverse_computation_order_ &&
-      instruction->opcode() == HloOpcode::kConditional && operand_no > 0 &&
+  if (instruction->opcode() == HloOpcode::kConditional && operand_no > 0 &&
       instruction->operand(operand_no)->user_count() == 1) {
     auto branch_comp = instruction->branch_computation(operand_no - 1);
     auto param = branch_comp->parameter_instruction(0);
@@ -2081,13 +1377,16 @@ void LayoutAssignment::SetupCopiedInstruction(const HloInstruction& instruction,
 absl::Status LayoutAssignment::CheckLayouts(
     HloModule* module,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
+  // Run dataflow analysis with propagate_through_control_flow=true so HloValue
+  // positions span while, conditional, and optimization barrier boundaries and
+  // the per HloValue layout check below verifies layout agreement across them.
   ABSL_ASSIGN_OR_RETURN(auto dataflow_analysis,
                    HloDataflowAnalysis::Run(
                        *module, /*ssa_form=*/false,
                        /*bitcast_defines_value=*/false, execution_threads,
                        /*propagate_through_calls=*/false,
                        /*precompute_uses=*/std::nullopt,
-                       /*propagate_through_control_flow=*/false));
+                       /*propagate_through_control_flow=*/true));
   for (auto* computation :
        module->MakeNonfusionComputations(execution_threads)) {
     for (auto* instruction : computation->instructions()) {
@@ -2127,24 +1426,56 @@ absl::Status LayoutAssignment::CheckLayouts(
           }));
 
       // Verify instructions that have special layout constraints.
+      ABSL_RETURN_IF_ERROR(ForEachCallBoundaryWithStatus(
+          instruction,
+          [&](const HloCallBoundary& boundary) -> absl::Status {
+            if (instruction->opcode() == HloOpcode::kAsyncStart) {
+              return absl::OkStatus();
+            }
+            TF_RET_CHECK(boundary.callee->num_parameters() ==
+                         boundary.caller_operands.size());
+            if (instruction->opcode() == HloOpcode::kFusion) {
+              for (int64_t i = 0; i < boundary.num_parameters(); ++i) {
+                TF_RET_CHECK(
+                    LayoutsInShapesEqual(boundary.callee_parameter(i)->shape(),
+                                         boundary.caller_operand(i)->shape()));
+              }
+              TF_RET_CHECK(LayoutsInShapesEqual(
+                  instruction->shape(), boundary.callee_root()->shape()));
+              return absl::OkStatus();
+            }
+            const ComputationLayout& callee_layout =
+                computation_constraints(boundary.callee).computation_layout();
+            for (int64_t i = 0; i < boundary.num_parameters(); ++i) {
+              TF_RET_CHECK(
+                  callee_layout.parameter_layout(i).MatchesLayoutInShape(
+                      ShardedShape(instruction,
+                                   boundary.caller_operand(i)->shape(), i),
+                      /*minor_to_major_only=*/true));
+            }
+            if (boundary.root_feeds_callsite) {
+              TF_RET_CHECK(callee_layout.result_layout().MatchesLayoutInShape(
+                  ShardedShape(instruction, instruction->shape(), -1),
+                  /*minor_to_major_only=*/true));
+              TF_RET_CHECK(callee_layout.result_layout().MatchesLayoutInShape(
+                  boundary.callee_root()->shape(),
+                  /*minor_to_major_only=*/true));
+            }
+            return absl::OkStatus();
+          },
+          HloCallBoundaryOptions(/*include_calls_in=*/true,
+                                 /*include_control_flow_in=*/true,
+                                 /*include_fusions_in=*/true,
+                                 /*include_associative_scans_in=*/false,
+                                 &execution_threads)));
       switch (instruction->opcode()) {
-        case HloOpcode::kCall:
-          ABSL_RETURN_IF_ERROR(CheckCallLayout(
-              instruction,
-              FindOrDie(computation_layouts_, instruction->to_apply())
-                  ->computation_layout()));
-          break;
         case HloOpcode::kCustomCall:
           ABSL_RETURN_IF_ERROR(CheckCustomCallLayout(instruction));
           break;
-        case HloOpcode::kFusion:
-          ABSL_RETURN_IF_ERROR(CheckFusionLayout(instruction));
-          break;
         case HloOpcode::kParameter:
           ABSL_RETURN_IF_ERROR(CheckParameterLayout(
-              instruction,
-              FindOrDie(computation_layouts_, instruction->parent())
-                  ->computation_layout()));
+              instruction, computation_constraints(instruction->parent())
+                               .computation_layout()));
           break;
         case HloOpcode::kBroadcast:
           ABSL_RETURN_IF_ERROR(CheckBroadcastLayout(instruction));
@@ -2152,30 +1483,6 @@ absl::Status LayoutAssignment::CheckLayouts(
         case HloOpcode::kConstant:
           ABSL_RETURN_IF_ERROR(CheckConstantLayout(instruction));
           break;
-        case HloOpcode::kWhile:
-          ABSL_RETURN_IF_ERROR(CheckWhileLayout(
-              instruction,
-              FindOrDie(computation_layouts_, instruction->while_condition())
-                  ->computation_layout(),
-              FindOrDie(computation_layouts_, instruction->while_body())
-                  ->computation_layout()));
-          break;
-        case HloOpcode::kOptimizationBarrier:
-          ABSL_RETURN_IF_ERROR(CheckOptimizationBarrierLayout(instruction));
-          break;
-        case HloOpcode::kConditional: {
-          std::vector<ComputationLayout> branch_computation_layouts;
-          const auto& branch_computations = instruction->branch_computations();
-          branch_computation_layouts.reserve(branch_computations.size());
-          for (const auto branch_computation : branch_computations) {
-            branch_computation_layouts.emplace_back(
-                FindOrDie(computation_layouts_, branch_computation)
-                    ->computation_layout());
-          }
-          ABSL_RETURN_IF_ERROR(CheckConditionalLayout(
-              instruction, absl::MakeSpan(branch_computation_layouts)));
-          break;
-        }
         default:
           break;
       }
@@ -2195,10 +1502,7 @@ absl::Status LayoutAssignment::CheckLayouts(
   return absl::OkStatus();
 }
 
-LayoutAssignment::LayoutAssignment(
-    ComputationLayout* entry_computation_layout,
-    ChannelLayoutConstraints* channel_constraints,
-    bool reverse_computation_order)
+LayoutAssignment::LayoutAssignment(ComputationLayout* entry_computation_layout)
     : HloDataflowPropagation(
           /*dataflow_analysis=*/nullptr,
           HloCallBoundaryOptions(/*include_calls_in=*/true,
@@ -2206,14 +1510,7 @@ LayoutAssignment::LayoutAssignment(
                                  /*include_fusions_in=*/false,
                                  /*include_associative_scans_in=*/false)),
       entry_computation_layout_(entry_computation_layout),
-      saved_entry_computation_layout_(*entry_computation_layout),
-      reverse_computation_order_(reverse_computation_order),
-      channel_layout_constraints_(channel_constraints) {
-  if (channel_layout_constraints_ != nullptr) {
-    // Save a copy of the input ChannelLayoutConstraints so that we can reset it
-    // if we have to undo previous operations (ClearPreviousPassSideEffects()).
-    channel_constraints_ = *channel_layout_constraints_;
-  }
+      saved_entry_computation_layout_(*entry_computation_layout) {
   VLOG(1) << "Entry computation layout given to layout assignment: "
           << entry_computation_layout_->ToString();
 }
@@ -2325,6 +1622,26 @@ static Layout GetReduceLayoutFromOperand(const Layout& operand_layout,
   return operand_shape.layout();
 }
 
+bool LayoutAssignment::CanOverrideOperandLayoutOnBufferOverride(
+    const HloInstruction* user, int64_t operand_no, const Layout& new_layout,
+    const Layout& existing_layout) {
+  if (InstructionCanChangeLayoutInstance(user)) {
+    return false;
+  }
+  if (user->opcode() == HloOpcode::kConcatenate) {
+    const HloInstruction* operand = user->operand(operand_no);
+    for (int64_t i = 0; i < user->operand_count(); ++i) {
+      const HloInstruction* sibling = user->operand(i);
+      if (sibling != operand &&
+          sibling->shape().dimensions(user->concatenate_dimension()) >
+              operand->shape().dimensions(user->concatenate_dimension())) {
+        return false;
+      }
+    }
+  }
+  return true;
+}
+
 bool LayoutAssignment::OperandLayoutAlwaysPropagateForward(
     const HloInstruction* user) {
   switch (user->opcode()) {
@@ -3039,8 +2356,9 @@ absl::Status LayoutAssignment::ConstrainOrOverrideBuffersAtIndex(
             user_comp_constraints.ResultLayout()
                 ->AnyMinorToMajorInLayoutIsSet();
         if (can_override_op_constraint &&
-            !InstructionCanChangeLayoutInstance(user) && !is_pinned_root &&
-            op_constraint->priority() < new_priority &&
+            CanOverrideOperandLayoutOnBufferOverride(
+                user, op_no, constraint.layout(), existing_layout) &&
+            !is_pinned_root && op_constraint->priority() < new_priority &&
             op_constraint->shape_layout().shape().IsArray() &&
             Layout::Equal().MinorToMajorOnly()(
                 op_constraint->shape_layout().layout(), existing_layout)) {
@@ -3438,13 +2756,11 @@ absl::Status LayoutAssignment::CalculateComputationLayout(
   for (HloInstruction* instruction :
        constraints->computation()->MakeInstructionPostOrder()) {
     const HloCallBoundaryOptions boundary_options(
-        /*include_calls_in=*/reverse_computation_order_,
-        /*include_control_flow_in=*/
-        (reverse_computation_order_ &&
-         instruction->opcode() != HloOpcode::kAsyncStart) ||
-            (instruction->opcode() == HloOpcode::kConditional &&
-             copy_disabled_while_computations_.contains(
-                 constraints->computation())),
+        /*include_calls_in=*/false,
+        /*include_control_flow_in=*/instruction->opcode() ==
+                HloOpcode::kConditional &&
+            copy_disabled_while_computations_.contains(
+                constraints->computation()),
         /*include_fusions_in=*/true,
         /*include_associative_scans_in=*/false);
     ABSL_RETURN_IF_ERROR(ForEachCallBoundaryWithStatus(
@@ -3477,32 +2793,6 @@ absl::Status LayoutAssignment::CalculateComputationLayout(
   return absl::OkStatus();
 }
 
-absl::Status LayoutAssignment::ClearComputationLayouts(
-    HloComputation* computation) {
-  // Clear existing layouts of the instructions.  All layouts must be assigned
-  // by the LayoutAssignment pass, except for those on parameters, the
-  // computation result, and a couple special cases. The former two are
-  // specified in computation_layout.  Clearing the layouts here avoids hiding
-  // potential bugs in the layout assignment pass that may accidentally use the
-  // existing layout.
-  for (HloInstruction* instruction : computation->instructions()) {
-    if (instruction->opcode() == HloOpcode::kBitcast) {
-      // bitcasts are inherently layout sensitive and so a bitcast instruction
-      // present in the IR before layout assignment is a bug.
-      return Internal(
-          "Unexpected bitcast operation seen during layout assignment: %s.",
-          instruction->ToString());
-    }
-    // Some instructions carry mandatory layouts in their shape.
-    if (instruction->opcode() != HloOpcode::kInfeed &&
-        !IsLayoutConstrainedCustomCall(instruction) &&
-        !IsLayoutConstrainedCollective(instruction)) {
-      LayoutUtil::ClearLayout(instruction->mutable_shape());
-    }
-  }
-  return absl::OkStatus();
-}
-
 void LayoutAssignment::InitUnconstrainedBuffers(
     absl::Span<HloComputation* const> computations) {
   absl::flat_hash_set<const HloComputation*> computation_set(
@@ -3515,10 +2805,6 @@ void LayoutAssignment::InitUnconstrainedBuffers(
   }
 }
 
-void LayoutAssignment::InitUnconstrainedBuffers(HloComputation* computation) {
-  InitUnconstrainedBuffers(absl::MakeConstSpan(&computation, 1));
-}
-
 absl::Status LayoutAssignment::AddCustomCallConstraints(
     LayoutConstraints* constraints) {
   for (HloInstruction* instruction :
@@ -3590,78 +2876,6 @@ absl::Status LayoutAssignment::AssignLayoutsToUnconstrainedBuffers(
   return absl::OkStatus();
 }
 
-absl::Status LayoutAssignment::RunOnComputation(
-    LayoutConstraints* constraints,
-    ChannelLayoutConstraints* channel_constraints) {
-  HloComputation* computation = constraints->computation();
-  VLOG(1) << "LayoutAssignment::RunOnComputation(" << computation->name()
-          << ")";
-  VLOG(4) << computation->ToString() << "\n";
-
-  // Initialize the set of unconstrained buffers with all array-shaped logical
-  // buffers in this computation.
-  InitUnconstrainedBuffers(computation);
-
-  // Add mandatory constraints required for correctness (e.g., entry
-  // parameter layouts).
-  ABSL_RETURN_IF_ERROR(AddMandatoryConstraints(channel_constraints, constraints));
-
-  // Add backend-specific constraints.
-  ABSL_RETURN_IF_ERROR(AddBackendConstraints(constraints));
-
-  // Constrain layouts for custom calls that have specific layout requirements.
-  ABSL_RETURN_IF_ERROR(AddCustomCallConstraints(constraints));
-
-  // Propagate the mandatory, backend, and custom call constraints.
-  ABSL_RETURN_IF_ERROR(PropagateConstraints(constraints));
-
-  // Record instructions that lack layout constraints before applying defaults.
-  RecordUnconstrainedLayoutInstructions();
-
-  // Iteratively assign layouts to remaining unconstrained buffers and
-  // propagate.
-  ABSL_RETURN_IF_ERROR(AssignLayoutsToUnconstrainedBuffers(constraints));
-
-  ABSL_RETURN_IF_ERROR(CalculateComputationLayout(constraints));
-  // Record layouts of communication operations to ensure consistency across
-  // modules.
-  if (channel_constraints != nullptr) {
-    ABSL_RETURN_IF_ERROR(ConstrainChannelLayouts(computation, channel_constraints));
-  }
-
-  return absl::OkStatus();
-}
-
-absl::Status LayoutAssignment::ConstrainChannelLayouts(
-    HloComputation* computation,
-    ChannelLayoutConstraints* channel_constraints) {
-  for (HloInstruction* instruction : computation->MakeInstructionPostOrder()) {
-    if (instruction->IsCrossModuleAllReduce() &&
-        !instruction->GetModule()->config().use_spmd_partitioning() &&
-        instruction->opcode() != HloOpcode::kAllReduceStart &&
-        instruction->opcode() != HloOpcode::kAllReduceDone) {
-      // TODO: b/501070020 - Support asynchronous all-reduce.
-      ABSL_ASSIGN_OR_RETURN(auto op_layout, InferArrayLayout(instruction, {}));
-      VLOG(5) << "Constrain cross module all reduce: " << op_layout.ToString()
-              << "\n";
-      const Layout* existing_layout = channel_constraints->ConstrainChannel(
-          instruction->channel_id().value(), op_layout);
-      if (existing_layout != nullptr &&
-          LayoutUtil::ValidateLayoutForShape(*existing_layout,
-                                             instruction->shape())
-              .ok()) {
-        ABSL_RETURN_IF_ERROR(SetInstructionLayout(
-            *existing_layout, instruction, /*mandatory=*/true, /*dfs=*/true,
-            /*allow_alias=*/false,
-            current_priority_ + kNumberOfPropagationRounds));
-        ABSL_RETURN_IF_ERROR(
-            PropagateConstraints(mutable_computation_constraints(computation)));
-      }
-    }
-  }
-  return absl::OkStatus();
-}
-
 namespace {
 
 void CopyMemorySpace(const Shape& src, Shape* dst) {
@@ -3860,12 +3074,6 @@ absl::StatusOr<std::vector<HloComputation*>> LayoutAssignment::SetupPropagation(
   SetDataflowAnalysis(std::move(dataflow_analysis));
   auto computations_to_work =
       module->MakeNonfusionComputations(execution_threads);
-  // If the reverse_computation_order_ flag is set, reverse the ordering of
-  // traversing computations, to generate an alternative layout assignment.
-  if (reverse_computation_order_ && !computations_to_work.empty()) {
-    absl::c_reverse(computations_to_work);
-    VLOG(2) << "reversing traversal order for computation:";
-  }
   HloComputation* entry = module->entry_computation();
   entry_computation_ = entry;
   computation_layouts_.emplace(
@@ -3971,21 +3179,14 @@ absl::StatusOr<bool> LayoutAssignment::RunImpl(
     for (auto* computation : computations_to_work) {
       LayoutConstraints* constraints =
           mutable_computation_constraints(computation);
-      ABSL_RETURN_IF_ERROR(
-          AddMandatoryConstraints(channel_layout_constraints_, constraints));
+      ABSL_RETURN_IF_ERROR(AddMandatoryConstraints(constraints));
       ABSL_RETURN_IF_ERROR(AddBackendConstraints(constraints));
       ABSL_RETURN_IF_ERROR(AddCustomCallConstraints(constraints));
-      if (reverse_computation_order_) {
-        ABSL_RETURN_IF_ERROR(PropagateConstraints(constraints));
-        ABSL_RETURN_IF_ERROR(CalculateComputationLayout(constraints));
-      }
     }
 
     LayoutConstraints* entry_constraints =
         mutable_computation_constraints(module->entry_computation());
-    if (!reverse_computation_order_) {
-      ABSL_RETURN_IF_ERROR(PropagateConstraints(entry_constraints));
-    }
+    ABSL_RETURN_IF_ERROR(PropagateConstraints(entry_constraints));
 
     // Propagate constrained HloValues across computation boundaries (caller <->
     // callee) into computations whose boundary HloValues are unconstrained,
@@ -4004,10 +3205,6 @@ absl::StatusOr<bool> LayoutAssignment::RunImpl(
       ABSL_RETURN_IF_ERROR(AddSubcomputationLayoutConstraints(constraints));
       ABSL_RETURN_IF_ERROR(PropagateConstraints(constraints));
       ABSL_RETURN_IF_ERROR(CalculateComputationLayout(constraints));
-      if (channel_layout_constraints_ != nullptr) {
-        ABSL_RETURN_IF_ERROR(
-            ConstrainChannelLayouts(computation, channel_layout_constraints_));
-      }
     }
     current_priority_ += 1;
     auto* entry_constraint =
@@ -4018,9 +3215,6 @@ absl::StatusOr<bool> LayoutAssignment::RunImpl(
                      ResolveInputOutputAliasing(module, entry_constraint));
     changed |= aliasing_changed;
     changed |= while_layout_changed_;
-    if (reverse_computation_order_ && i == 0) {
-      changed = true;
-    }
   }
 
   // All logical buffers should have constraints at this point. All that
diff --git a/third_party/xla/xla/service/layout_assignment.h b/third_party/xla/xla/service/layout_assignment.h
index 4a6a5fe4550..6dee4bc5993 100644
--- a/third_party/xla/xla/service/layout_assignment.h
+++ b/third_party/xla/xla/service/layout_assignment.h
@@ -332,79 +332,13 @@ class ComputationLayoutConstraint : public LayoutConstraint {
   ComputationLayout computation_layout_;
 };
 
-// Encapsulates layout constraints across communication channels (Send/Recv).
-//
-// ChannelLayoutConstraints ensures layout consistency across communication
-// boundaries (such as Send/Recv instructions matching a channel_id). Any
-// unconstrained channels are assigned locally optimal layouts which are then
-// registered as channel constraints.
-//
-// Data Structure Usage:
-// - `constraints_`: An absl::flat_hash_map mapping each channel ID (int64_t)
-//   to its constrained Layout object.
-class ChannelLayoutConstraints {
- public:
-  // Constructs an empty channel constraint set.
-  ChannelLayoutConstraints() = default;
-
-  // Returns true if channel_id has an associated layout constraint.
-  bool IsChannelConstrained(int64_t channel_id) const {
-    return constraints_.contains(channel_id);
-  }
-
-  // Given `shape`, applies the constrained layout for `channel_id`.
-  // `channel_id` must already be constrained.
-  Shape LayoutShapeForChannel(Shape shape, int64_t channel_id) const {
-    auto it = constraints_.find(channel_id);
-    CHECK(it != constraints_.end()) << "Channel " << channel_id;
-    *shape.mutable_layout() = it->second;
-    return shape;
-  }
-
-  // Returns the Layout constraint for `channel_id`, which must already be
-  // constrained.
-  const Layout& LayoutForChannel(int64_t channel_id) const {
-    auto it = constraints_.find(channel_id);
-    CHECK(it != constraints_.end()) << "Channel " << channel_id;
-    return it->second;
-  }
-
-  // Adds a new layout constraint for `channel_id`. Returns nullptr if the
-  // channel constraint was successfully added or matches an existing
-  // constraint; otherwise returns a pointer to the existing conflicting
-  // layout.
-  const Layout* ConstrainChannel(int64_t channel_id, const Layout& layout) {
-    auto it = constraints_.emplace(std::make_pair(channel_id, layout));
-    if (it.second) {
-      return nullptr;
-    }
-    return LayoutUtil::Equal(layout, it.first->second) ? nullptr
-                                                       : &it.first->second;
-  }
-
- private:
-  // Map from channel ID to its assigned layout.
-  absl::flat_hash_map<int64_t, Layout> constraints_;
-};
-
 // HLO pass which assigns layouts to all instructions in the HLO module while
 // satisfying all necessary invariants and minimizing cost.
 class LayoutAssignment : public HloModulePass, public HloDataflowPropagation {
  public:
   // entry_computation_layout is modified to populate a layout for the result in
   // the case that no particular layout is requested.
-  //
-  // channel_constraints is both an input and output. Any sends or recvs that
-  // are present in channel_constraints will be laid out as constrained. Any
-  // unconstrained sends or recvs will be laid out as locally optimal and their
-  // layout will be added as a constraint to channel_constraints.
-  //
-  // If channel_constraints is nullptr, no kSend or kRecvs must be contained
-  // within any module passed to `Run`.
-  explicit LayoutAssignment(
-      ComputationLayout* entry_computation_layout,
-      ChannelLayoutConstraints* channel_constraints = nullptr,
-      bool reverse_computation_order = false);
+  explicit LayoutAssignment(ComputationLayout* entry_computation_layout);
   ~LayoutAssignment() override = default;
   using HloModulePass::Run;
   const HloDataflowAnalysis& dataflow_analysis() const {
@@ -599,7 +533,6 @@ class LayoutAssignment : public HloModulePass, public HloDataflowPropagation {
                                 const HloInstruction* instruction,
                                 int64_t operand_no, bool mandatory, bool dfs,
                                 int64_t priority);
-  bool reverse_computation_order() const { return reverse_computation_order_; }
 
   ComputationLayout& saved_entry_computation_layout() {
     return saved_entry_computation_layout_;
@@ -618,6 +551,12 @@ class LayoutAssignment : public HloModulePass, public HloDataflowPropagation {
                                       const Layout& existing_layout) {
     return false;
   }
+  // Returns true if overriding an operand buffer's layout during
+  // cross computation propagation is allowed to override the existing operand
+  // layout constraint on (user, operand_no).
+  virtual bool CanOverrideOperandLayoutOnBufferOverride(
+      const HloInstruction* user, int64_t operand_no, const Layout& new_layout,
+      const Layout& existing_layout);
   // Should be made consistent with the ChooseOperandLayoutFromOutputLayout
   // except that a boolean instead of concrete layout is returned.
   virtual bool OperandLayoutAlwaysPropagateForward(const HloInstruction* user);
@@ -665,11 +604,8 @@ class LayoutAssignment : public HloModulePass, public HloDataflowPropagation {
     return absl::OkStatus();
   }
 
-  absl::Status PropagateUnconstraintedBuffers(LayoutConstraints* constraints);
   const BufferLayoutConstraint* GetBufferLayoutConstraint(
       const HloValue& buffer) const;
-  absl::StatusOr<const BufferLayoutConstraint*>
-  GetInstructionBufferLayoutConstraint(const HloInstruction* instruction) const;
   using BufferSet = absl::flat_hash_set<const HloValue*>;
   // Find a bufferset in the bufferset cache. This is useful since we can
   // currently create the flattened buffer set for the same instruction many
@@ -739,11 +675,6 @@ class LayoutAssignment : public HloModulePass, public HloDataflowPropagation {
     return shape;
   }
 
-  // The operands of a call must match the layouts of parameters in the
-  // ComputationLayout, and the call instruction itself must match the result
-  // layout in the ComputationLayout.
-  absl::Status CheckCallLayout(HloInstruction* call,
-                               const ComputationLayout& computation_layout);
   // For a custom-call user, propagates the operand constraint to the result
   // based on output-to-operand aliasing.
   absl::Status PropagateOperandConstraintToResultForCustomCall(
@@ -783,24 +714,17 @@ class LayoutAssignment : public HloModulePass, public HloDataflowPropagation {
 
   // Adds constraints which must be satisfied for correctness on all
   // backends. Called once prior to propagating constraints.
-  absl::Status AddMandatoryConstraints(
-      ChannelLayoutConstraints* channel_constraints,
-      LayoutConstraints* constraints);
+  absl::Status AddMandatoryConstraints(LayoutConstraints* constraints);
 
   // Adds constraints for instructions that define values with pre-existing
   // layouts.
-  absl::Status AddInstructionLayoutConstraints(
-      ChannelLayoutConstraints* channel_constraints,
-      LayoutConstraints* constraints);
+  absl::Status AddInstructionLayoutConstraints(LayoutConstraints* constraints);
 
   absl::Status AddInfeedConstraints(HloInstruction* instruction);
   absl::Status AddOutfeedConstraints(HloInstruction* instruction);
   absl::Status AddParameterConstraints(HloInstruction* instruction,
                                        LayoutConstraints* constraints);
   absl::Status AddCollectiveConstraints(HloInstruction* instruction);
-  absl::Status AddCrossModuleAllReduceConstraints(
-      HloInstruction* instruction,
-      ChannelLayoutConstraints* channel_constraints);
 
   // Adds constraints for instructions that call or interact with
   // sub-computations.
@@ -808,67 +732,12 @@ class LayoutAssignment : public HloModulePass, public HloDataflowPropagation {
       LayoutConstraints* constraints);
 
   absl::Status AddCallConstraints(HloInstruction* instruction);
-  absl::Status AddWhileConstraints(HloInstruction* instruction,
-                                   LayoutConstraints* constraints);
+  absl::Status AddWhileConstraints(HloInstruction* instruction);
   absl::Status AddConditionalConstraints(HloInstruction* instruction);
   absl::Status AddAsyncStartConstraints(HloInstruction* instruction);
   absl::Status AddAsyncDoneConstraints(HloInstruction* instruction,
                                        LayoutConstraints* constraints);
 
-  // Propagates while loop parameter and result layouts to subcomputations (such
-  // as conditionals) within the while body or condition.
-  void PropagateWhileLoopLayoutToSubcomputations(HloComputation* computation,
-                                                 const Shape& param_shape,
-                                                 const Shape* result_shape,
-                                                 int64_t priority);
-
-  // Propagates layout constraints from the caller instruction into the inner
-  // async sub-computation.
-  // This is the forward propagation step: it takes the layouts of the operands
-  // and result of the async start/update instruction (which are in the parent
-  // computation) and propagates them to the parameters and result of the
-  // async sub-computation.
-  // If any layout in the sub-computation is updated, it resets the
-  // sub-computation layout with an elevated priority to ensure it is respected
-  // during the sub-computation's layout assignment. Returns the reconciled
-  // layout of the sub-computation.
-  ComputationLayout PropagateLayoutsToAsyncSubComputation(
-      const HloInstruction* instruction, LayoutConstraints* async_constraint);
-
-  // Propagates the operand array layouts of `instruction` to the parameter
-  // layouts defined in `async_layout`.
-  // Updates `async_layout` in-place and returns true if any parameter layout
-  // was changed.
-  bool PropagateOperandLayoutsToAsyncParameters(
-      const HloInstruction* instruction, ComputationLayout* async_layout);
-
-  // Propagates array layouts for a single operand `param_idx` of `instruction`
-  // to the corresponding parameter layout in `async_layout`.
-  // `instruction` operand `param_idx` is mapped to the parameter `param_idx`
-  // of the async sub-computation.
-  // Updates `async_layout` in-place and returns true if the layout was updated.
-  bool PropagateOperandLayoutToAsyncParameter(const HloInstruction* instruction,
-                                              int64_t param_idx,
-                                              ComputationLayout* async_layout);
-
-  // Propagates array layouts from the result shape of `instruction` (tuple
-  // element 1) to `async_layout`'s result layout.
-  // We assume the result shape of the async operation is at index {1} of the
-  // `instruction` (async start/update) output tuple.
-  // Updates `async_layout` in-place and returns true if the result layout was
-  // updated.
-  bool PropagateResultLayoutToAsyncSubComputation(
-      const HloInstruction* instruction, ComputationLayout* async_layout);
-
-  // Propagates async sub-computation parameter and result layout constraints
-  // back onto the caller instruction and its operands in the parent
-  // computation. This is the backward propagation step: it takes the resolved
-  // layouts from the async sub-computation and applies them as mandatory
-  // constraints on the caller instruction's shape (at index {1} for result) and
-  // its operands.
-  absl::Status PropagateLayoutsFromAsyncSubComputation(
-      HloInstruction* instruction, const ComputationLayout& async_layout,
-      LayoutConstraints* async_constraint);
   // Sets the computation result layout based on constraints and
   // sub-computations.
   absl::Status AddComputationResultLayoutConstraints(
@@ -878,8 +747,7 @@ class LayoutAssignment : public HloModulePass, public HloDataflowPropagation {
   absl::Status AddCustomCallConstraints(LayoutConstraints* constraints);
 
   // Initializes unconstrained_buffer_ids_ with all array HloValues in the
-  // given computation or span of computations.
-  void InitUnconstrainedBuffers(HloComputation* computation);
+  // given span of computations.
   void InitUnconstrainedBuffers(absl::Span<HloComputation* const> computations);
 
   // HloDataflowPropagation overrides for cross computation layout propagation.
@@ -940,15 +808,6 @@ class LayoutAssignment : public HloModulePass, public HloDataflowPropagation {
     return absl::OkStatus();
   }
 
-  // Construct constraints and assign layouts to all instructions in the
-  // computation satisfying the given ComputationLayout, if not nullptr.
-  // Otherwise the ComputationLayout will be calculated by propagating the
-  // computation instruction constraints.
-  // Layouts constraints are added, then propagated until all LogicalBuffers in
-  // the computation are constrained.
-  absl::Status RunOnComputation(LayoutConstraints* constraints,
-                                ChannelLayoutConstraints* channel_constraints);
-
   // Assign layouts to the instructions of a computation which satisfy the given
   // layout constraints. Copies may be added to satisfy the constraints. The
   // given LayoutConstraints must have layout constraints every logical buffer
@@ -975,9 +834,6 @@ class LayoutAssignment : public HloModulePass, public HloDataflowPropagation {
   // constraints to computation nested inside.
   absl::Status CalculateComputationLayout(LayoutConstraints* constraints);
 
-  // Clears all the layouts which can be cleared within a computation.
-  absl::Status ClearComputationLayouts(HloComputation* computation);
-
   // Clears the side effects of a previous pass, like added copy instructions.
   absl::Status ClearPreviousPassSideEffects(
       HloModule* module,
@@ -997,8 +853,6 @@ class LayoutAssignment : public HloModulePass, public HloDataflowPropagation {
   // A copy of entry_computation_layout_ used to reset it to the initial values
   // during the multiple passes done by the layout assignment operation.
   ComputationLayout saved_entry_computation_layout_;
-  // If set true, reverse the computation traversal order when assigning layout.
-  bool reverse_computation_order_;
 
  protected:
   static constexpr int64_t kNumberOfPropagationRounds = 2;
@@ -1038,28 +892,10 @@ class LayoutAssignment : public HloModulePass, public HloDataflowPropagation {
   absl::Status AddCopyForOperand(HloInstruction* instruction,
                                  int64_t operand_number);
 
-  // Apply the channel layout constraints by populating the channel_constraints
-  // data structure passed in at constructor time. Eventually adds copies in
-  // case two ends of a channel ended up with a different leyout.
-  absl::Status ConstrainChannelLayouts(
-      HloComputation* computation,
-      ChannelLayoutConstraints* channel_constraints);
-
-  // Resets the input ChannelLayoutConstraints to the original copy received
-  // from the constructor input.
-  void ResetChannelConstraints() {
-    if (channel_layout_constraints_ != nullptr) {
-      *channel_layout_constraints_ = channel_constraints_;
-    }
-  }
-
   void ResetEntryComputationLayout() {
     *entry_computation_layout_ = saved_entry_computation_layout_;
   }
 
-  // Adds constraints related to host Send/Recv instructions.
-  absl::Status BuildHostChannelConstraints(HloComputation* computation);
-
   // Replaces dataflow_analysis_ and drops the buffer sets memoized from the
   // previous analysis.
   void SetDataflowAnalysis(std::unique_ptr<HloDataflowAnalysis> analysis);
@@ -1114,19 +950,6 @@ class LayoutAssignment : public HloModulePass, public HloDataflowPropagation {
   // here.
   absl::flat_hash_set<HloInstruction*> added_copies_;
 
-  // The pointer to the channel layout constraints passed in with the
-  // constructor. If not nullptr, this is an input/output argument.
-  ChannelLayoutConstraints* channel_layout_constraints_ = nullptr;
-
-  // A copy of the input layout constraints used to reset the above pointer in
-  // case we have to undo operations due to the multiple passes over the
-  // computations/instructions.
-  ChannelLayoutConstraints channel_constraints_;
-
-  // Layout constraints for send/recv instructions which communicate with the
-  // host.
-  ChannelLayoutConstraints host_channel_constraints_;
-
   // Array-shaped buffers which have not yet been constrained, in id order.
   absl::btree_set<HloValue::Id> unconstrained_buffer_ids_;
 
diff --git a/third_party/xla/xla/service/layout_assignment_test.cc b/third_party/xla/xla/service/layout_assignment_test.cc
index b36115188bd..e0b347b67ef 100644
--- a/third_party/xla/xla/service/layout_assignment_test.cc
+++ b/third_party/xla/xla/service/layout_assignment_test.cc
@@ -59,15 +59,13 @@ namespace {
 namespace m = xla::match;
 using ::testing::ElementsAre;
 
-absl::Status AssignLayoutsToComputation(
-    HloModule* m, ChannelLayoutConstraints* channel_constraints = nullptr) {
+absl::Status AssignLayoutsToComputation(HloModule* m) {
   if (!m->entry_computation_layout().result_layout().LayoutIsSet()) {
     m->mutable_entry_computation_layout()
         ->mutable_result_layout()
         ->SetToDefaultLayout();
   }
-  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout(),
-                                     channel_constraints);
+  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout());
   return layout_assignment.Run(m).status();
 }
 
@@ -82,20 +80,15 @@ class LayoutAssignmentTest : public HloHardwareIndependentTestBase {
   }
 
  protected:
-  void AssignLayouts(HloModule* m, ComputationLayout* entry_computation_layout,
-                     ChannelLayoutConstraints* channel_constraints = nullptr) {
-    LayoutAssignment layout_assignment(
-        entry_computation_layout,
-        /*channel_constraints=*/channel_constraints);
+  void AssignLayouts(HloModule* m,
+                     ComputationLayout* entry_computation_layout) {
+    LayoutAssignment layout_assignment(entry_computation_layout);
     EXPECT_IS_OK(layout_assignment.Run(m).status());
   }
 
   absl::StatusOr<bool> AssignLayoutsAndVerifyHlo(
-      HloModule* m, ComputationLayout* entry_computation_layout,
-      ChannelLayoutConstraints* channel_constraints = nullptr) {
-    LayoutAssignment layout_assignment(
-        entry_computation_layout,
-        /*channel_constraints=*/channel_constraints);
+      HloModule* m, ComputationLayout* entry_computation_layout) {
+    LayoutAssignment layout_assignment(entry_computation_layout);
     EXPECT_IS_OK(layout_assignment.Run(m).status());
     HloVerifier verifier(/*layout_sensitive=*/true,
                          /*allow_mixed_precision=*/false);
@@ -886,8 +879,7 @@ TEST_F(LayoutAssignmentTest, ChannelLayoutMismatch) {
   computation_layout.mutable_result_layout()->ResetLayout(
       LayoutUtil::MakeLayout({1, 0}));
 
-  ChannelLayoutConstraints channel_constraints;
-  AssignLayouts(m.get(), &computation_layout, &channel_constraints);
+  AssignLayouts(m.get(), &computation_layout);
 
   EXPECT_TRUE(ShapeUtil::Equal(FindInstruction(m.get(), "send")->shape(),
                                FindInstruction(m.get(), "recv")->shape()));
@@ -919,51 +911,7 @@ TEST_F(LayoutAssignmentTest, AllReduceSpmd) {
   // different shapes.
   ComputationLayout computation_layout(
       m->entry_computation()->ComputeProgramShape());
-  ChannelLayoutConstraints channel_constraints;
-  AssignLayouts(m.get(), &computation_layout, &channel_constraints);
-}
-
-TEST_F(LayoutAssignmentTest, AllReduceLayoutMissmatch) {
-  // Pin non matching layouts to parameter and root.
-  const char* module_str = R"(
-    HloModule test_module
-
-    add {
-      lhs = f32[] parameter(0)
-      rhs = f32[] parameter(1)
-      ROOT add = f32[] add(lhs, rhs)
-    }
-
-    ENTRY entry_computation {
-      param = (f32[2,2]) parameter(0)
-      gte = f32[2,2] get-tuple-element(param), index=0
-      ar.0 = f32[2,2] all-reduce(gte),
-        channel_id=1, replica_groups={{0}}, to_apply=add,
-        sharding={maximal device=0}
-      const = f32[2,2] constant({{0,1},{2,3}})
-      ROOT ar.1 = f32[2,2] all-reduce(const),
-        channel_id=1, replica_groups={{0}}, to_apply=add,
-        sharding={maximal device=1}
-    })";
-  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> m,
-                       ParseAndReturnVerifiedModule(module_str));
-  ComputationLayout computation_layout(
-      m->entry_computation()->ComputeProgramShape());
-  Shape param_shape = ShapeUtil::MakeTupleShape(
-      {ShapeUtil::MakeShapeWithDenseLayout(F32, {2, 2}, {0, 1})});
-  ASSERT_OK(computation_layout.mutable_parameter_layout(0)->CopyLayoutFromShape(
-      param_shape));
-  computation_layout.mutable_result_layout()->ResetLayout(
-      LayoutUtil::MakeLayout({1, 0}));
-
-  ChannelLayoutConstraints channel_constraints;
-  AssignLayouts(m.get(), &computation_layout, &channel_constraints);
-
-  EXPECT_THAT(LayoutOf(m.get(), "gte"), ElementsAre(0, 1));
-  EXPECT_THAT(LayoutOf(m.get(), "ar.0"), ElementsAre(0, 1));
-  EXPECT_THAT(LayoutOf(m.get(), "ar.1"), ElementsAre(0, 1));
-  const HloInstruction* root = m->entry_computation()->root_instruction();
-  EXPECT_THAT(root->shape().layout().minor_to_major(), ElementsAre(1, 0));
+  AssignLayouts(m.get(), &computation_layout);
 }
 
 TEST_F(LayoutAssignmentTest, CopySliceOperandToAvoidImplicitLayoutChange) {
@@ -1741,8 +1689,7 @@ ENTRY entry_computation {
   ComputationLayout computation_layout(
       m->entry_computation()->ComputeProgramShape(), /*ignore_layouts=*/false);
 
-  ChannelLayoutConstraints channel_constraints;
-  AssignLayouts(m.get(), &computation_layout, &channel_constraints);
+  AssignLayouts(m.get(), &computation_layout);
 
   const HloInstruction* crs = FindInstruction(m.get(), "crs");
   ExpectTupleLayoutIs(crs->shape(), {{0, 1}, {1, 0}});
@@ -1778,8 +1725,7 @@ ENTRY entry_computation {
   ComputationLayout computation_layout(
       m->entry_computation()->ComputeProgramShape(), /*ignore_layouts=*/false);
 
-  ChannelLayoutConstraints channel_constraints;
-  AssignLayouts(m.get(), &computation_layout, &channel_constraints);
+  AssignLayouts(m.get(), &computation_layout);
 
   const HloInstruction* alltoall = FindInstruction(m.get(), "alltoall");
   ExpectTupleLayoutIs(alltoall->shape(), {{1, 0}, {1, 0}});
@@ -1857,11 +1803,7 @@ ENTRY main {
       ShapeLayout(ShapeUtil::MakeShapeWithDenseLayout(F32, {93184, 4}, {0, 1}));
   *computation_layout.mutable_result_layout() = ShapeLayout(
       ShapeUtil::MakeShapeWithDenseLayout(F32, {2, 512, 364}, {0, 1, 2}));
-  ChannelLayoutConstraints channel_constraints;
-  LayoutAssignment layout_assignment(
-      &computation_layout,
-      /*channel_constraints=*/&channel_constraints,
-      /* reverse_computation_order = */ true);
+  LayoutAssignment layout_assignment(&computation_layout);
   EXPECT_IS_OK(layout_assignment.Run(m.get()).status());
   const HloInstruction* call_1 = FindInstruction(m.get(), "reshape.8494");
   ExpectLayoutIs(call_1->shape(), {0, 1, 2});
@@ -1928,10 +1870,7 @@ ENTRY main {
   *computation_layout.mutable_result_layout() =
       ShapeLayout(ShapeUtil::MakeShapeWithDenseLayout(F32, {64, 243, 243, 384},
                                                       {3, 0, 2, 1}));
-  ChannelLayoutConstraints channel_constraints;
-  LayoutAssignment layout_assignment(
-      &computation_layout,
-      /*channel_constraints=*/&channel_constraints);
+  LayoutAssignment layout_assignment(&computation_layout);
   EXPECT_IS_OK(layout_assignment.Run(m.get()).status());
   const HloInstruction* subtract_15 = FindInstruction(m.get(), "subtract.15");
   ExpectLayoutIs(subtract_15->shape(), {3, 0, 2, 1});
@@ -1958,8 +1897,7 @@ TEST_F(LayoutAssignmentTest, PropagateOperandLayout2) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> m,
                        ParseAndReturnVerifiedModule(module_str));
 
-  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout(),
-                                     nullptr);
+  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout());
   EXPECT_IS_OK(layout_assignment.Run(m.get()).status());
   const HloInstruction* reshape_1 = FindInstruction(m.get(), "reshape.1");
   ExpectLayoutIs(reshape_1->shape(), {1, 0});
@@ -1987,8 +1925,7 @@ ENTRY main {
 
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> m,
                        ParseAndReturnVerifiedModule(module_str));
-  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout(),
-                                     nullptr);
+  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout());
   EXPECT_IS_OK(layout_assignment.Run(m.get()).status());
   const HloInstruction* param =
       m->entry_computation()->parameter_instruction(0);
@@ -2017,8 +1954,7 @@ TEST_F(LayoutAssignmentTest, PartialEntryParameterLayout) {
   // Allow propagation only to parameter 0
   m->mutable_entry_computation_layout()->mutable_parameter_layout(0)->Clear();
 
-  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout(),
-                                     nullptr);
+  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout());
   EXPECT_IS_OK(layout_assignment.Run(m.get()).status());
   // Assign bitcasting layout to parameter 0
   ExpectLayoutIs(m->entry_computation_layout().parameter_layout(0).shape(),
@@ -2047,8 +1983,7 @@ TEST_F(LayoutAssignmentTest, TupleEntryParameterLayoutNoResultConstraint) {
   // Allow propagation only to parameter 0
   m->mutable_entry_computation_layout()->mutable_result_layout()->Clear();
 
-  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout(),
-                                     nullptr);
+  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout());
   EXPECT_IS_OK(layout_assignment.Run(m.get()).status());
   // Assign bitcasting layout to parameter 0
   ExpectLayoutIs(
@@ -2083,8 +2018,7 @@ TEST_F(LayoutAssignmentTest,
   m->mutable_entry_computation_layout()->mutable_parameter_layout(0)->Clear(
       {1});
 
-  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout(),
-                                     nullptr);
+  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout());
   EXPECT_IS_OK(layout_assignment.Run(m.get()).status());
   // Assign bitcasting layout to parameter 0
   ExpectLayoutIs(
@@ -2115,8 +2049,7 @@ TEST_F(LayoutAssignmentTest, AliasParameterAndOutput) {
   m->mutable_entry_computation_layout()->mutable_result_layout()->Clear();
   m->mutable_entry_computation_layout()->mutable_parameter_layout(0)->Clear();
 
-  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout(),
-                                     nullptr);
+  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout());
   EXPECT_IS_OK(layout_assignment.Run(m.get()).status());
   EXPECT_EQ(m->entry_computation_layout().result_layout().shape(),
             m->entry_computation_layout().parameter_layout(0).shape());
@@ -2140,8 +2073,7 @@ TEST_F(LayoutAssignmentTest, AliasUnconstrainedParamterWithConstrainedOutput) {
   m->mutable_entry_computation_layout()->SetToDefaultLayout();
   m->mutable_entry_computation_layout()->mutable_parameter_layout(0)->Clear();
 
-  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout(),
-                                     nullptr);
+  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout());
   EXPECT_IS_OK(layout_assignment.Run(m.get()).status());
   EXPECT_EQ(m->entry_computation_layout().result_layout().shape(),
             m->entry_computation_layout().parameter_layout(0).shape());
@@ -2164,8 +2096,7 @@ TEST_F(LayoutAssignmentTest, AliasConstrainedParamterWithUnconstrainedOutput) {
   m->mutable_entry_computation_layout()->SetToDefaultLayout();
   m->mutable_entry_computation_layout()->mutable_result_layout()->Clear();
 
-  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout(),
-                                     nullptr);
+  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout());
   EXPECT_IS_OK(layout_assignment.Run(m.get()).status());
   EXPECT_EQ(m->entry_computation_layout().result_layout().shape(),
             m->entry_computation_layout().parameter_layout(0).shape());
@@ -2218,8 +2149,7 @@ ENTRY %main {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> m,
                        ParseAndReturnVerifiedModule(module_str));
 
-  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout(),
-                                     nullptr);
+  LayoutAssignment layout_assignment(m->mutable_entry_computation_layout());
   EXPECT_IS_OK(layout_assignment.Run(m.get()).status());
 }
 
@@ -3553,5 +3483,47 @@ ENTRY %main (y: f32[3,1]) -> f32[3,1] {
   ExpectLayoutIs(FindInstruction(m.get(), "bcast")->shape(), {0, 1});
 }
 
+TEST_F(LayoutAssignmentTest,
+       CrossComputationOverrideSkipsSmallerConcatenateOperand) {
+  const char* module_str = R"hlo(
+HloModule CrossComputationOverrideSkipsSmallerConcatenateOperand
+
+%callee (p: f32[256,3,1]) -> f32[256,3,1,1] {
+  %p = f32[256,3,1]{2,1,0} parameter(0)
+  %pinned = f32[256,3,1]{2,1,0} custom-call(%p), custom_call_target="Pinned",
+    operand_layout_constraints={f32[256,3,1]{2,1,0}}
+  ROOT %r = f32[256,3,1,1] reshape(%pinned)
+}
+
+ENTRY %main (x: f32[256,3,1], y: f32[256,3,1,3]) -> f32[256,3,1,4] {
+  %x = f32[256,3,1]{2,1,0} parameter(0)
+  %call = f32[256,3,1,1] call(%x), to_apply=%callee
+  %y = f32[256,3,1,3]{3,2,1,0} parameter(1)
+  ROOT %concat = f32[256,3,1,4] concatenate(%call, %y), dimensions={3}
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> m,
+                       ParseAndReturnVerifiedModule(module_str));
+  ComputationLayout computation_layout(
+      m->entry_computation()->ComputeProgramShape());
+  *computation_layout.mutable_parameter_layout(0) = ShapeLayout(
+      ShapeUtil::MakeShapeWithDenseLayout(F32, {256, 3, 1}, {2, 1, 0}));
+  *computation_layout.mutable_parameter_layout(1) = ShapeLayout(
+      ShapeUtil::MakeShapeWithDenseLayout(F32, {256, 3, 1, 3}, {3, 2, 1, 0}));
+  *computation_layout.mutable_result_layout() = ShapeLayout(
+      ShapeUtil::MakeShapeWithDenseLayout(F32, {256, 3, 1, 4}, {3, 2, 1, 0}));
+
+  LayoutAssignment layout_assignment(&computation_layout);
+  EXPECT_THAT(
+      layout_assignment.Run(m.get(), {HloInstruction::kMainExecutionThread}),
+      absl_testing::IsOk());
+  const HloInstruction* concat = FindInstruction(m.get(), "concat");
+  ASSERT_NE(concat, nullptr);
+  ExpectLayoutIs(concat->shape(), {3, 2, 1, 0});
+  ExpectLayoutIs(concat->operand(0)->shape(), {3, 2, 1, 0});
+  ExpectLayoutIs(concat->operand(1)->shape(), {3, 2, 1, 0});
+}
+
 }  // namespace
 }  // namespace xla
diff --git a/third_party/xla/xla/tools/hlo_opt/cpu_opt.cc b/third_party/xla/xla/tools/hlo_opt/cpu_opt.cc
index f2d9cfdcdff..4a24f139021 100644
--- a/third_party/xla/xla/tools/hlo_opt/cpu_opt.cc
+++ b/third_party/xla/xla/tools/hlo_opt/cpu_opt.cc
@@ -178,8 +178,7 @@ class CpuOptProvider : public CompiledOptProvider {
     // Fails to register if module does not have entry computation layout
     if (module.config().has_entry_computation_layout()) {
       RegisterPass<cpu::CpuLayoutAssignment>(
-          module.mutable_entry_computation_layout(), &target_machine_features,
-          nullptr);
+          module.mutable_entry_computation_layout(), &target_machine_features);
     }
 
     const int max_parallelism =

From 3424b626c42bb575a6a7af530e33a3d5f63a4b06 Mon Sep 17 00:00:00 2001
From: Shaogang Wang <shawnw@nvidia.com>
Date: Thu, 8 Oct 2026 06:00:34 -0700
Subject: [PATCH 0287/1000] PR #50177: [XLA:GPU] DeviceMemAllocator:
 reserved-arena mode that maps physical memory on demand
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50177

**Stacked PR 2/5** splitting #48985 into reviewable layers. Stacked on #50176: the base branch is `shawnw/bfc-growth-1-se-vmm-api`, so this PR shows only its own commit. When #50176 lands, this PR is rebased onto `main` and re-targeted.

| # | PR | Title |
|---|---|---|
| 1 | #50176 | [XLA:GPU] StreamExecutor: VA reservation and physical allocation APIs sharing the executor's VMM options |
| 2 | #50177 | [XLA:GPU] DeviceMemAllocator: reserved-arena mode that maps physical memory on demand |
| 3 | #50178 | [XLA:TSL] BFC: grow spatially partitioned pools within a reserved arena |
| 4 | #50179 | [XLA:GPU] PJRT: typed MemFraction memory policy with START-CAP grammar |
| 5 | #50180 | [XLA:GPU] PJRT: enable BFC growth for growable memory fractions |

---

📝 Summary of Changes

`DeviceMemAllocator::ReserveMemory(capacity)` reserves a device VA range through the new StreamExecutor API. Later `Alloc` calls map physical memory right after the already mapped prefix, so the arena is contiguous and never moves, and `SupportsCoalescing()` reports true so BFC can extend its region instead of adding separate ones.

`Free` unmaps the original mappings behind a coalesced BFC region. A whole-range return (teardown) synchronizes the device first, because `cuMemUnmap`, unlike `cuMemFree`, invalidates the mapping immediately; a just-mapped suffix handed back by BFC never synchronizes, since BFC may extend during stream capture.

`tsl::SubAllocator` gains `GetAllocationGranularity()` (default 1) so BFC can size extensions in mapping granules; `DeviceMemAllocator` reports the reservation granularity.

🎯 Justification

Second layer of the growable spatial BFC pool. Nothing calls `ReserveMemory` yet, so there is no behavior change.

🚀 Kind of Contribution

✨ New Feature (dormant until wired up)

🧪 Unit Tests:

`//xla/stream_executor/integrations:device_mem_allocator_test` with a mock StreamExecutor: reservation and growth, failed extensions keeping the prefix intact, suffix returns, synchronize-before-unmap on teardown only, legacy mode untouched.

🧪 Execution Tests:

Not applicable until the pool is wired up.

🤖 Generated with [Claude Code](https://claude.com/claude-code)
Copybara import of the project:

--
52f71d4febaea737c30568b2cde1263908fb30ea by Shaogang Wang <shawnw@nvidia.com>:

[XLA:GPU] DeviceMemAllocator: reserved-arena mode that maps physical memory on demand

ReserveMemory(capacity) reserves a device VA range through the new
StreamExecutor API; subsequent Alloc calls map physical memory right after
the already mapped prefix, so the arena is contiguous and never moves.
SupportsCoalescing reports true in this mode so BFC can extend its region
instead of adding separate ones. Free unmaps the original mappings behind a
coalesced BFC region; a whole-range return (teardown) synchronizes the
device first because cuMemUnmap, unlike cuMemFree, invalidates mappings at
once, while a just-mapped suffix returned by BFC never synchronizes since
BFC may extend during stream capture.

SubAllocator gains GetAllocationGranularity() (default 1) so BFC can size
extensions in mapping granules; DeviceMemAllocator reports the reservation
granularity. Mock-based tests cover reservation, growth failures that keep
the prefix intact, suffix returns, and the synchronize-before-unmap rule.

Nothing calls ReserveMemory yet, so there is no behavior change. Second of
five stacked changes (see #48985).

--
3908401cfaad0845294fd8217db243b98e390cd6 by Shawn Wang <shawnw@nvidia.com>:

[XLA:GPU] Validate reserved arena metadata and allocation sizes

Merging this change closes #50177

PiperOrigin-RevId: 995783700
---
 .../xla/stream_executor/integrations/BUILD    |  23 +
 .../integrations/device_mem_allocator.h       | 157 ++++++-
 .../integrations/device_mem_allocator_test.cc | 403 ++++++++++++++++++
 third_party/xla/xla/tsl/framework/BUILD       |  17 +-
 third_party/xla/xla/tsl/framework/allocator.h |  15 +-
 5 files changed, 601 insertions(+), 14 deletions(-)
 create mode 100644 third_party/xla/xla/stream_executor/integrations/device_mem_allocator_test.cc

diff --git a/third_party/xla/xla/stream_executor/integrations/BUILD b/third_party/xla/xla/stream_executor/integrations/BUILD
index c40d7adfaa8..f4a0ca26c6c 100644
--- a/third_party/xla/xla/stream_executor/integrations/BUILD
+++ b/third_party/xla/xla/stream_executor/integrations/BUILD
@@ -89,10 +89,15 @@ cc_library(
     ],
     deps = [
         "//xla/stream_executor:device_address",
+        "//xla/stream_executor:memory_allocation",
+        "//xla/stream_executor:memory_reservation",
         "//xla/stream_executor:stream_executor_h",
         "//xla/tsl/framework:allocator",
         "//xla/tsl/framework:device_id",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
         "@tsl//tsl/profiler/lib:traceme",
     ],
 )
@@ -160,3 +165,21 @@ xla_cc_test(
         "//xla/tsl/framework:allocator",
     ]),
 )
+
+xla_cc_test(
+    name = "device_mem_allocator_test",
+    srcs = ["device_mem_allocator_test.cc"],
+    deps = [
+        ":device_mem_allocator",
+        "//xla/stream_executor:device_address",
+        "//xla/stream_executor:generic_memory_allocation",
+        "//xla/stream_executor:memory_allocation",
+        "//xla/stream_executor:memory_reservation",
+        "//xla/stream_executor:mock_stream_executor",
+        "//xla/tsl/framework:device_id",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
diff --git a/third_party/xla/xla/stream_executor/integrations/device_mem_allocator.h b/third_party/xla/xla/stream_executor/integrations/device_mem_allocator.h
index f298de7c953..da74a983313 100644
--- a/third_party/xla/xla/stream_executor/integrations/device_mem_allocator.h
+++ b/third_party/xla/xla/stream_executor/integrations/device_mem_allocator.h
@@ -17,10 +17,18 @@ limitations under the License.
 #define XLA_STREAM_EXECUTOR_INTEGRATIONS_DEVICE_MEM_ALLOCATOR_H_
 
 #include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <utility>
 #include <vector>
 
 #include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
 #include "xla/stream_executor/device_address.h"
+#include "xla/stream_executor/memory_allocation.h"
+#include "xla/stream_executor/memory_reservation.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/framework/allocator.h"
 #include "xla/tsl/framework/device_id.h"
@@ -44,12 +52,52 @@ class DeviceMemAllocator : public tsl::SubAllocator {
     CHECK(stream_exec_ != nullptr);
   }
 
-  ~DeviceMemAllocator() override = default;
+  ~DeviceMemAllocator() override {
+    if (mapped_bytes_ != 0) {
+      // Qualified: virtual dispatch is gone during destruction anyway.
+      DeviceMemAllocator::Free(reservation_->address().opaque(), mapped_bytes_);
+    }
+  }
+
+  // Optional mode for a growing BFC arena. Reserve capacity bytes of VA, then
+  // back a contiguous prefix on demand in Alloc. Existing mappings never move.
+  // Call before any Alloc and before passing this suballocator to BFCAllocator,
+  // whose constructor caches SupportsCoalescing(). Returns Unimplemented on
+  // unsupported backends.
+  // Calls to this suballocator must be externally serialized (as in BFC).
+  absl::Status ReserveMemory(size_t capacity) {
+    if (capacity == 0 || allocations_started_ || reservation_ != nullptr) {
+      return absl::InvalidArgumentError(
+          "ReserveMemory requires a positive capacity and an unused allocator");
+    }
+    ABSL_ASSIGN_OR_RETURN(std::unique_ptr<MemoryReservation> reservation,
+                     stream_exec_->CreateMemoryReservation(capacity));
+    if (reservation->address().opaque() == nullptr ||
+        reservation->granularity() == 0) {
+      return absl::InternalError(
+          "Device VA reservation requires a non-null address and positive "
+          "granularity");
+    }
+    if (reservation->address().size() < capacity) {
+      return absl::InternalError(
+          "Device VA reservation is smaller than capacity");
+    }
+    reservation_ = std::move(reservation);
+    capacity_ = capacity;
+    return absl::OkStatus();
+  }
 
   void* Alloc(size_t alignment, size_t num_bytes,
               size_t* bytes_received) override {
     tsl::profiler::TraceMe traceme("DeviceMemAllocator::Alloc");
 
+    if (num_bytes > 0) {
+      allocations_started_ = true;
+    }
+    if (reservation_ != nullptr) {
+      return AllocReserved(num_bytes, bytes_received);
+    }
+
     void* ptr = nullptr;
     *bytes_received = num_bytes;
     if (num_bytes > 0) {
@@ -67,21 +115,126 @@ class DeviceMemAllocator : public tsl::SubAllocator {
     tsl::profiler::TraceMe traceme("DeviceMemAllocator::Free");
 
     if (ptr != nullptr) {
+      if (reservation_ != nullptr) {
+        // BFC frees the coalesced region at teardown, or returns a just-added
+        // suffix if its size is unsuitable. Unmap each original mapping before
+        // releasing its physical handle, and pair visitors with original Alloc
+        // calls rather than the coalesced BFC region.
+        const uintptr_t base =
+            reinterpret_cast<uintptr_t>(reservation_->address().opaque());
+        const uintptr_t start = reinterpret_cast<uintptr_t>(ptr);
+        CHECK_GE(start, base);
+        CHECK_LE(start - base, mapped_bytes_);
+        CHECK_EQ(num_bytes, mapped_bytes_ - (start - base));
+        if (start == base) {
+          // Whole-range return: BFC tearing down, rejecting the initial
+          // region, or garbage-collecting the region. Buffers in this range
+          // may have been handed out, and cuMemUnmap (unlike cuMemFree)
+          // invalidates mappings immediately, so drain the device first.
+          // SynchronizeAllActivity bypasses CudaDeviceAllocator's graph
+          // capture guard; that is acceptable only because none of these
+          // callers runs while a capture using this arena is active.
+          //
+          // A suffix return (start > base) hands back backing that Extend
+          // mapped moments ago and never published, so no device work can
+          // reference it. It must not synchronize: Extend can run during
+          // stream capture, where a context synchronize would fail and
+          // invalidate the capture.
+          if (!stream_exec_->SynchronizeAllActivity()) {
+            LOG(WARNING) << "Failed to synchronize device "
+                         << device_id_.value()
+                         << " before unmapping its BFC arena; in-flight work "
+                            "may still reference it.";
+          }
+        }
+        while (num_bytes > 0) {
+          const DeviceAddressBase address =
+              regions_.back().mapping.mapped_address();
+          CHECK_GE(num_bytes, address.size());
+          VisitFree(address.opaque(), device_id_.value(), address.size());
+          num_bytes -= address.size();
+          mapped_bytes_ -= address.size();
+          regions_.pop_back();
+        }
+        return;
+      }
       VisitFree(ptr, device_id_.value(), num_bytes);
       DeviceAddressBase device_ptr(ptr);
       stream_exec_->Deallocate(&device_ptr);
     }
   }
 
-  bool SupportsCoalescing() const override { return false; }
+  bool SupportsCoalescing() const override { return reservation_ != nullptr; }
+
+  // Mapping granularity of the reserved arena; BFC sizes extensions and
+  // backpedal retries in multiples of it. Also used to round the initial
+  // backing size when reserving a growing arena.
+  size_t GetAllocationGranularity() const override {
+    return reservation_ != nullptr ? reservation_->granularity() : 1;
+  }
 
   tsl::AllocatorMemoryType GetMemoryType() const override {
     return tsl::AllocatorMemoryType::kDevice;
   }
 
  private:
+  void* AllocReserved(size_t num_bytes, size_t* bytes_received) {
+    *bytes_received = 0;
+    const size_t available_bytes = capacity_ - mapped_bytes_;
+    if (num_bytes == 0 || num_bytes > available_bytes) {
+      return nullptr;
+    }
+
+    const size_t granularity = reservation_->granularity();
+    const size_t remainder = num_bytes % granularity;
+    if (remainder != 0) {
+      const size_t padding = granularity - remainder;
+      // Check before adding so rounding cannot overflow or exceed capacity.
+      if (padding > available_bytes - num_bytes) {
+        return nullptr;
+      }
+      num_bytes += padding;
+    }
+
+    auto allocation = stream_exec_->CreatePhysicalMemoryAllocation(num_bytes);
+    if (!allocation.ok()) {
+      VLOG(2) << "Physical device allocation failed: " << allocation.status();
+      return nullptr;
+    }
+    const size_t size = (*allocation)->address().size();
+    if (size < num_bytes || size > available_bytes) {
+      return nullptr;
+    }
+    auto mapping = reservation_->MapTo(mapped_bytes_, 0, size, **allocation);
+    if (!mapping.ok()) {
+      VLOG(2) << "Mapping device memory failed: " << mapping.status();
+      return nullptr;
+    }
+    // MapTo also sets local/peer access, and unmaps this slice on access
+    // failure. Publish only after success; a failed extension leaves the prefix
+    // intact.
+    void* ptr = mapping->mapped_address().opaque();
+    regions_.push_back({std::move(*allocation), std::move(*mapping)});
+    mapped_bytes_ += size;
+    *bytes_received = size;
+    VisitAlloc(ptr, device_id_.value(), size);
+    return ptr;
+  }
+
   StreamExecutor* stream_exec_;  // not owned, non-null
   const tsl::PlatformDeviceId device_id_;
+  bool allocations_started_ = false;
+
+  // Declaration order keeps the reservation and physical backing alive until
+  // after their mappings are destroyed. The reserved suffix has no backing.
+  std::unique_ptr<MemoryReservation> reservation_;
+  struct Region {
+    std::unique_ptr<MemoryAllocation> allocation;
+    MemoryReservation::ScopedMapping mapping;
+  };
+  std::vector<Region> regions_;
+  size_t capacity_ = 0;
+  size_t mapped_bytes_ = 0;
 
   DeviceMemAllocator(const DeviceMemAllocator&) = delete;
   void operator=(const DeviceMemAllocator&) = delete;
diff --git a/third_party/xla/xla/stream_executor/integrations/device_mem_allocator_test.cc b/third_party/xla/xla/stream_executor/integrations/device_mem_allocator_test.cc
new file mode 100644
index 00000000000..06ab60ee7ae
--- /dev/null
+++ b/third_party/xla/xla/stream_executor/integrations/device_mem_allocator_test.cc
@@ -0,0 +1,403 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/stream_executor/integrations/device_mem_allocator.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <map>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "absl/status/statusor.h"
+#include "xla/stream_executor/device_address.h"
+#include "xla/stream_executor/generic_memory_allocation.h"
+#include "xla/stream_executor/memory_allocation.h"
+#include "xla/stream_executor/memory_reservation.h"
+#include "xla/stream_executor/mock_stream_executor.h"
+#include "xla/tsl/framework/device_id.h"
+
+namespace stream_executor {
+namespace {
+
+using ::absl_testing::StatusIs;
+using ::testing::ElementsAre;
+using ::testing::IsEmpty;
+using ::testing::Return;
+
+constexpr uintptr_t kBase = 0x10000000;
+void* Address(size_t offset) { return reinterpret_cast<void*>(kBase + offset); }
+
+// Independent bookkeeping models physical handles and active mappings. In
+// particular, releasing backing that is still mapped is an error in this fake.
+struct MemoryState {
+  absl::Status reserve_status;
+  absl::Status allocate_status;
+  absl::Status map_status;
+  absl::Status access_status;
+  void* reservation_address = Address(0);
+  size_t reservation_granularity = 256;
+  size_t allocation_alignment = 256;
+  size_t reserved_bytes = 0;
+  size_t physical_live = 0;
+  size_t next_handle = 0;
+  bool reservation_live = false;
+  std::map<size_t, size_t> mappings;  // offset -> physical handle
+  std::vector<size_t> allocation_requests;
+  std::vector<size_t> mapped_offsets;
+  std::vector<size_t> unmapped_offsets;
+};
+
+class FakeReservation : public MemoryReservation {
+ public:
+  FakeReservation(MemoryState& state, size_t capacity)
+      : state_(state), capacity_(capacity) {
+    state_.reservation_live = true;
+  }
+  ~FakeReservation() override {
+    EXPECT_THAT(state_.mappings, IsEmpty());
+    EXPECT_EQ(state_.physical_live, 0);
+    state_.reservation_live = false;
+  }
+  DeviceAddressBase address() const override {
+    return DeviceAddressBase(state_.reservation_address, capacity_);
+  }
+  size_t granularity() const override { return state_.reservation_granularity; }
+
+ private:
+  absl::Status Map(size_t offset, size_t allocation_offset, size_t size,
+                   MemoryAllocation& allocation) override {
+    EXPECT_EQ(offset % granularity(), 0);
+    EXPECT_EQ(size % granularity(), 0);
+    EXPECT_EQ(allocation_offset, 0);
+    EXPECT_EQ(size, allocation.address().size());
+    EXPECT_LE(offset + size, capacity_);
+    if (!state_.map_status.ok()) {
+      return state_.map_status;
+    }
+    EXPECT_TRUE(state_.mappings
+                    .emplace(offset, reinterpret_cast<uintptr_t>(
+                                         allocation.address().opaque()))
+                    .second);
+    state_.mapped_offsets.push_back(offset);
+    return absl::OkStatus();
+  }
+  absl::Status SetAccess(uint64_t offset, size_t size) override {
+    return state_.access_status;
+  }
+  absl::Status UnMap(size_t offset, size_t size) override {
+    EXPECT_EQ(offset % granularity(), 0);
+    EXPECT_EQ(size % granularity(), 0);
+    EXPECT_EQ(state_.mappings.erase(offset), 1);
+    state_.unmapped_offsets.push_back(offset);
+    return absl::OkStatus();
+  }
+  MemoryState& state_;
+  size_t capacity_;
+};
+
+class ReservationExecutor : public MockStreamExecutor {
+ public:
+  ReservationExecutor() {
+    ON_CALL(*this, SynchronizeAllActivity()).WillByDefault(Return(true));
+  }
+  MemoryState state;
+  absl::StatusOr<std::unique_ptr<MemoryReservation>> CreateMemoryReservation(
+      uint64_t size) override {
+    if (!state.reserve_status.ok()) {
+      return state.reserve_status;
+    }
+    state.reserved_bytes = size;
+    return std::make_unique<FakeReservation>(state, size);
+  }
+  absl::StatusOr<std::unique_ptr<MemoryAllocation>>
+  CreatePhysicalMemoryAllocation(uint64_t size) override {
+    state.allocation_requests.push_back(size);
+    if (!state.allocate_status.ok()) {
+      return state.allocate_status;
+    }
+    ++state.physical_live;
+    const size_t alignment = state.allocation_alignment;
+    return std::make_unique<GenericMemoryAllocation>(
+        reinterpret_cast<void*>(++state.next_handle),
+        (size + alignment - 1) / alignment * alignment,
+        [this](void* handle, uint64_t size) {
+          for (const auto& [offset, mapped_handle] : state.mappings) {
+            EXPECT_NE(mapped_handle, reinterpret_cast<uintptr_t>(handle));
+          }
+          --state.physical_live;
+        });
+  }
+};
+
+TEST(DeviceMemAllocatorTest, ReservesOnlyVaAndAppendsPhysicalBacking) {
+  ReservationExecutor executor;
+  std::vector<std::pair<void*, size_t>> allocated, freed;
+  {
+    DeviceMemAllocator allocator(
+        &executor, tsl::PlatformDeviceId(0), {[&](void* ptr, int, size_t size) {
+          allocated.emplace_back(ptr, size);
+        }},
+        {[&](void* ptr, int, size_t size) { freed.emplace_back(ptr, size); }});
+    EXPECT_FALSE(allocator.SupportsCoalescing());
+    ASSERT_OK(allocator.ReserveMemory(16384));
+    EXPECT_TRUE(allocator.SupportsCoalescing());
+    EXPECT_EQ(allocator.GetAllocationGranularity(), 256);
+    // Returning the whole mapped range synchronizes the device exactly once,
+    // before any mapping is torn down.
+    EXPECT_CALL(executor, SynchronizeAllActivity()).WillOnce([&] {
+      EXPECT_THAT(executor.state.unmapped_offsets, IsEmpty());
+      EXPECT_EQ(executor.state.physical_live, 2);
+      return true;
+    });
+    EXPECT_EQ(executor.state.reserved_bytes, 16384);
+    EXPECT_THAT(executor.state.allocation_requests, IsEmpty());
+    EXPECT_THAT(executor.state.mappings, IsEmpty());
+
+    size_t received = 0;
+    EXPECT_EQ(allocator.Alloc(256, 4096, &received), Address(0));
+    EXPECT_EQ(received, 4096);
+    EXPECT_EQ(allocator.Alloc(256, 2048, &received), Address(4096));
+    EXPECT_EQ(received, 2048);
+    EXPECT_THAT(executor.state.mapped_offsets, ElementsAre(0, 4096));
+    EXPECT_EQ(executor.state.physical_live, 2);
+    // BFC coalesces these allocations and returns them as one region.
+    allocator.Free(Address(0), 6144);
+    EXPECT_EQ(executor.state.physical_live, 0);
+    EXPECT_THAT(executor.state.unmapped_offsets, ElementsAre(4096, 0));
+    ASSERT_EQ(allocated.size(), 2);
+    EXPECT_THAT(freed, ElementsAre(allocated[1], allocated[0]));
+  }
+  EXPECT_FALSE(executor.state.reservation_live);
+}
+
+enum class Failure { kPhysical, kMap, kAccess };
+class DeviceMemAllocatorFailureTest : public ::testing::TestWithParam<Failure> {
+};
+
+TEST_P(DeviceMemAllocatorFailureTest,
+       FailedExtensionPreservesPrefixAndCanRetry) {
+  ReservationExecutor executor;
+  // Failed extensions never synchronize; only the destructor does.
+  EXPECT_CALL(executor, SynchronizeAllActivity()).Times(1);
+  DeviceMemAllocator allocator(&executor, tsl::PlatformDeviceId(0));
+  ASSERT_OK(allocator.ReserveMemory(16384));
+  size_t received = 0;
+  ASSERT_EQ(allocator.Alloc(256, 4096, &received), Address(0));
+  const size_t original_handle = executor.state.mappings.at(0);
+  absl::Status* failure_status = nullptr;
+  switch (GetParam()) {
+    case Failure::kPhysical:
+      failure_status = &executor.state.allocate_status;
+      break;
+    case Failure::kMap:
+      failure_status = &executor.state.map_status;
+      break;
+    case Failure::kAccess:
+      failure_status = &executor.state.access_status;
+      break;
+  }
+  *failure_status = absl::ResourceExhaustedError("injected growth failure");
+  EXPECT_EQ(allocator.Alloc(256, 2048, &received), nullptr);
+  EXPECT_EQ(received, 0);
+  ASSERT_EQ(executor.state.mappings.size(), 1);
+  EXPECT_EQ(executor.state.mappings.at(0), original_handle);
+  EXPECT_EQ(executor.state.physical_live, 1);
+  if (GetParam() == Failure::kAccess) {
+    EXPECT_THAT(executor.state.unmapped_offsets, ElementsAre(4096));
+  } else {
+    EXPECT_THAT(executor.state.unmapped_offsets, IsEmpty());
+  }
+  *failure_status = absl::OkStatus();
+  EXPECT_EQ(allocator.Alloc(256, 2048, &received), Address(4096));
+  EXPECT_EQ(received, 2048);
+  EXPECT_EQ(executor.state.mappings.at(0), original_handle);
+  // Destructor also handles a partially backed reservation.
+}
+
+INSTANTIATE_TEST_SUITE_P(Growth, DeviceMemAllocatorFailureTest,
+                         ::testing::Values(Failure::kPhysical, Failure::kMap,
+                                           Failure::kAccess));
+
+TEST(DeviceMemAllocatorTest, GranularityPaddingCannotExceedCapacity) {
+  ReservationExecutor executor;
+  EXPECT_CALL(executor, SynchronizeAllActivity()).Times(1);
+  DeviceMemAllocator allocator(&executor, tsl::PlatformDeviceId(0));
+  ASSERT_OK(allocator.ReserveMemory(1025));
+  size_t received = 0;
+  ASSERT_EQ(allocator.Alloc(256, 512, &received), Address(0));
+  // 513 rounds to 768, which cannot fit in the remaining 513 bytes.
+  EXPECT_EQ(allocator.Alloc(256, 513, &received), nullptr);
+  EXPECT_EQ(received, 0);
+  EXPECT_EQ(executor.state.physical_live, 1);
+  EXPECT_THAT(executor.state.mapped_offsets, ElementsAre(0));
+  EXPECT_EQ(allocator.Alloc(256, 512, &received), Address(512));
+  EXPECT_EQ(allocator.Alloc(256, 1, &received), nullptr);
+  EXPECT_EQ(received, 0);
+  EXPECT_EQ(allocator.Alloc(256, 0, &received), nullptr);
+  EXPECT_EQ(received, 0);
+  EXPECT_THAT(executor.state.allocation_requests, ElementsAre(512, 512));
+}
+
+TEST(DeviceMemAllocatorTest, RoundsRequestsToReservationGranularity) {
+  ReservationExecutor executor;
+  DeviceMemAllocator allocator(&executor, tsl::PlatformDeviceId(0));
+  ASSERT_OK(allocator.ReserveMemory(1024));
+  size_t received = 0;
+  EXPECT_EQ(allocator.Alloc(256, 1, &received), Address(0));
+  EXPECT_EQ(received, 256);
+  EXPECT_EQ(allocator.Alloc(256, 257, &received), Address(256));
+  EXPECT_EQ(received, 512);
+  EXPECT_THAT(executor.state.allocation_requests, ElementsAre(256, 512));
+}
+
+TEST(DeviceMemAllocatorTest, RoundingOverflowDoesNotAllocatePhysicalMemory) {
+  ReservationExecutor executor;
+  EXPECT_CALL(executor, SynchronizeAllActivity()).Times(0);
+  DeviceMemAllocator allocator(&executor, tsl::PlatformDeviceId(0));
+  const size_t capacity = std::numeric_limits<size_t>::max();
+  ASSERT_OK(allocator.ReserveMemory(capacity));
+  size_t received = 1;
+  EXPECT_EQ(allocator.Alloc(256, capacity, &received), nullptr);
+  EXPECT_EQ(received, 0);
+  EXPECT_THAT(executor.state.allocation_requests, IsEmpty());
+  EXPECT_THAT(executor.state.mappings, IsEmpty());
+}
+
+TEST(DeviceMemAllocatorTest, LargerBackendAlignmentCannotExceedCapacity) {
+  ReservationExecutor executor;
+  executor.state.allocation_alignment = 512;
+  DeviceMemAllocator allocator(&executor, tsl::PlatformDeviceId(0));
+  ASSERT_OK(allocator.ReserveMemory(768));
+  size_t received = 0;
+  ASSERT_EQ(allocator.Alloc(256, 256, &received), Address(0));
+  EXPECT_EQ(received, 512);
+  // The request fits at the reservation's granularity, but the backend's
+  // additional padding cannot fit in the remaining 256 bytes.
+  EXPECT_EQ(allocator.Alloc(256, 256, &received), nullptr);
+  EXPECT_EQ(received, 0);
+  EXPECT_THAT(executor.state.allocation_requests, ElementsAre(256, 256));
+  EXPECT_EQ(executor.state.physical_live, 1);
+  EXPECT_THAT(executor.state.mapped_offsets, ElementsAre(0));
+}
+
+TEST(DeviceMemAllocatorTest, ReturnedSuffixCanBeMappedAgain) {
+  ReservationExecutor executor;
+  int syncs = 0;
+  EXPECT_CALL(executor, SynchronizeAllActivity()).WillRepeatedly([&] {
+    ++syncs;
+    return true;
+  });
+  {
+    DeviceMemAllocator allocator(&executor, tsl::PlatformDeviceId(0));
+    ASSERT_OK(allocator.ReserveMemory(1024));
+    size_t received = 0;
+    ASSERT_EQ(allocator.Alloc(256, 512, &received), Address(0));
+    ASSERT_EQ(allocator.Alloc(256, 256, &received), Address(512));
+    // A just-mapped suffix was never published, so returning it must not
+    // synchronize: BFC can extend during stream capture.
+    allocator.Free(Address(512), 256);
+    EXPECT_EQ(syncs, 0);
+    EXPECT_THAT(executor.state.unmapped_offsets, ElementsAre(512));
+    EXPECT_EQ(allocator.Alloc(256, 512, &received), Address(512));
+    EXPECT_EQ(received, 512);
+    EXPECT_EQ(executor.state.physical_live, 2);
+  }
+  // Teardown returns the whole range and synchronizes once.
+  EXPECT_EQ(syncs, 1);
+}
+
+TEST(DeviceMemAllocatorTest, SyncFailureDoesNotPreventUnmap) {
+  ReservationExecutor executor;
+  EXPECT_CALL(executor, SynchronizeAllActivity()).WillOnce(Return(false));
+  DeviceMemAllocator allocator(&executor, tsl::PlatformDeviceId(0));
+  ASSERT_OK(allocator.ReserveMemory(1024));
+  size_t received = 0;
+  ASSERT_EQ(allocator.Alloc(256, 512, &received), Address(0));
+  allocator.Free(Address(0), 512);
+  EXPECT_THAT(executor.state.unmapped_offsets, ElementsAre(0));
+  EXPECT_EQ(executor.state.physical_live, 0);
+}
+
+TEST(DeviceMemAllocatorTest, ReservationErrorsAreReported) {
+  ReservationExecutor executor;
+  EXPECT_CALL(executor, SynchronizeAllActivity()).Times(0);
+  DeviceMemAllocator allocator(&executor, tsl::PlatformDeviceId(0));
+  EXPECT_THAT(allocator.ReserveMemory(0),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+  executor.state.reserve_status = absl::ResourceExhaustedError("VA exhausted");
+  EXPECT_THAT(allocator.ReserveMemory(1024),
+              StatusIs(absl::StatusCode::kResourceExhausted));
+  EXPECT_FALSE(allocator.SupportsCoalescing());
+  executor.state.reserve_status = absl::OkStatus();
+  ASSERT_OK(allocator.ReserveMemory(1024));
+  EXPECT_THAT(allocator.ReserveMemory(2048),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+TEST(DeviceMemAllocatorTest, RejectsZeroReservationGranularity) {
+  ReservationExecutor executor;
+  EXPECT_CALL(executor, SynchronizeAllActivity()).Times(0);
+  DeviceMemAllocator allocator(&executor, tsl::PlatformDeviceId(0));
+  executor.state.reservation_granularity = 0;
+  EXPECT_THAT(allocator.ReserveMemory(1024),
+              StatusIs(absl::StatusCode::kInternal));
+  EXPECT_FALSE(allocator.SupportsCoalescing());
+  EXPECT_EQ(allocator.GetAllocationGranularity(), 1);
+  EXPECT_FALSE(executor.state.reservation_live);
+  executor.state.reservation_granularity = 256;
+  ASSERT_OK(allocator.ReserveMemory(1024));
+}
+
+TEST(DeviceMemAllocatorTest, RejectsNullReservationAddress) {
+  ReservationExecutor executor;
+  EXPECT_CALL(executor, SynchronizeAllActivity()).Times(0);
+  DeviceMemAllocator allocator(&executor, tsl::PlatformDeviceId(0));
+  executor.state.reservation_address = nullptr;
+  EXPECT_THAT(allocator.ReserveMemory(1024),
+              StatusIs(absl::StatusCode::kInternal));
+  EXPECT_FALSE(allocator.SupportsCoalescing());
+  EXPECT_FALSE(executor.state.reservation_live);
+  executor.state.reservation_address = Address(0);
+  ASSERT_OK(allocator.ReserveMemory(1024));
+}
+
+TEST(DeviceMemAllocatorTest, LegacyModeUsesExecutorAllocation) {
+  MockStreamExecutor executor;
+  // Legacy frees go through Deallocate, which defers reclamation itself.
+  EXPECT_CALL(executor, SynchronizeAllActivity()).Times(0);
+  DeviceMemAllocator allocator(&executor, tsl::PlatformDeviceId(0));
+  EXPECT_THAT(allocator.ReserveMemory(1024),
+              StatusIs(absl::StatusCode::kUnimplemented));
+  EXPECT_FALSE(allocator.SupportsCoalescing());
+  EXPECT_CALL(executor, Allocate(256, 0))
+      .WillOnce(Return(DeviceAddressBase(Address(0), 512)));
+  EXPECT_CALL(executor, Deallocate(::testing::_)).Times(1);
+  size_t received = 0;
+  EXPECT_EQ(allocator.Alloc(256, 256, &received), Address(0));
+  EXPECT_EQ(received, 512);
+  EXPECT_THAT(allocator.ReserveMemory(1024),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+  allocator.Free(Address(0), received);
+}
+
+}  // namespace
+}  // namespace stream_executor
diff --git a/third_party/xla/xla/tsl/framework/BUILD b/third_party/xla/xla/tsl/framework/BUILD
index cdc05fab67d..4a80daad487 100644
--- a/third_party/xla/xla/tsl/framework/BUILD
+++ b/third_party/xla/xla/tsl/framework/BUILD
@@ -137,28 +137,25 @@ cc_library(
     features = ["parse_headers"],
     visibility = ["//visibility:public"],
     deps = [
+        "//xla/tsl/lib/gtl:inlined_vector",
+        "//xla/tsl/platform:logging",
+        "//xla/tsl/platform:macros",
+        "//xla/tsl/platform:types",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/synchronization",
+        "@tsl//tsl/platform:numbers",
+        "@tsl//tsl/platform:thread_annotations",
     ] + if_static(
         extra_deps = [
             ":allocator_registry_impl",
-            "@com_google_absl//absl/synchronization",
-            "//xla/tsl/lib/gtl:inlined_vector",
-            "@tsl//tsl/platform:numbers",
             "@tsl//tsl/platform:strcat",
             "//xla/tsl/platform:env",
             "//xla/tsl/platform:env_impl",
-            "//xla/tsl/platform:logging",
-            "//xla/tsl/platform:macros",
-            "//xla/tsl/platform:types",
             "@tsl//tsl/platform:platform_port",
-            "@tsl//tsl/platform:thread_annotations",
         ],
         otherwise = [
-            "//xla/tsl/lib/gtl:inlined_vector",
-            "//xla/tsl/platform:logging",
-            "@tsl//tsl/platform:numbers",
             "@tsl//tsl/platform:platform_port",
             "@tsl//tsl/platform:strcat",
             "//xla/tsl/platform:env",
diff --git a/third_party/xla/xla/tsl/framework/allocator.h b/third_party/xla/xla/tsl/framework/allocator.h
index 997232dd18d..1c42e1ac6b4 100644
--- a/third_party/xla/xla/tsl/framework/allocator.h
+++ b/third_party/xla/xla/tsl/framework/allocator.h
@@ -105,8 +105,8 @@ struct AllocatorStats {
   std::optional<int64_t> bytes_limit;
 
   // Stats for reserved memory usage.
-  int64_t bytes_reserved;       // Number of bytes reserved.
-  int64_t peak_bytes_reserved;  // The peak number of bytes reserved.
+  int64_t bytes_reserved;        // Number of bytes reserved.
+  int64_t peak_bytes_reserved;   // The peak number of bytes reserved.
   int64_t peak_allocated_bytes;  // Peak of reserved and in-use bytes.
   // The upper limit on the number bytes of reservable memory,
   // if such a limit is known.
@@ -472,6 +472,17 @@ class SubAllocator {
     return AllocatorMemoryType::kUnknown;
   }
 
+  // Returns a positive size multiple that Alloc() pads requests up to, for
+  // example a virtual memory mapping granularity. BFCAllocator sizes region
+  // extensions and backpedal retries in multiples of this value so that a
+  // shrunken retry is not padded back up to a size that just failed.
+  // Implementations whose granularity changes after suballocator construction
+  // (e.g. once an address range has been reserved) must return the current
+  // value on every call. Complete such configuration before passing the
+  // suballocator to BFCAllocator, whose constructor caches
+  // SupportsCoalescing().
+  virtual size_t GetAllocationGranularity() const { return 1; }
+
  protected:
   // Implementation of Alloc() method must call this on newly allocated
   // value.

From 0a4df4a4586e07a80dbbfcecfcc2e9dede1714e2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Eusebio=20Dur=C3=A1n=20Monta=C3=B1a?= <eusebiodm@google.com>
Date: Thu, 8 Oct 2026 06:01:24 -0700
Subject: [PATCH 0288/1000] Rollback of Set cuDNN graph SM version from device
 info

Reverts 6a4e21b8a7670ed5a70c076dcbd46a8987ce4f23

PiperOrigin-RevId: 995784138
---
 third_party/xla/xla/service/BUILD             | 18 +-----
 .../xla/service/xla_aot_compile_gpu_test.cc   | 60 ++++++++-----------
 third_party/xla/xla/service/xla_compile.bzl   | 17 ++----
 .../xla/xla/stream_executor/cuda/BUILD        |  2 -
 .../xla/xla/stream_executor/cuda/cuda_dnn.cc  | 20 -------
 5 files changed, 33 insertions(+), 84 deletions(-)

diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index 3317497e3dd..6b691c8cad2 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -5967,21 +5967,6 @@ xla_aot_compile_gpu(
     xla_flags = "--xla_gpu_use_new_autotune_cache_format=false",
 )
 
-xla_aot_compile_gpu(
-    name = "xla_aot_compile_test_gpu_executable_convolution_deviceless_cudnn",
-    # H100 only: deviceless conv graphs on Blackwell require cuDNN >= 9.19.
-    gpu_targets = [
-        "h100",
-        "h100_mig",
-    ],
-    module = "xla_aot_compile_test_convolution.mlir",
-    xla_flags = " ".join([
-        "--xla_gpu_experimental_enable_conv_fusion=true",
-        "--xla_gpu_cudnn_deviceless_compilation_mode=CUDNN_DEVICELESS_COMPILATION_AUTO",
-        "--xla_gpu_experimental_autotune_backends=cudnn",
-    ]),
-)
-
 xla_aot_compile_gpu_runtime_autotuning(
     name = "xla_aot_compile_test_gpu_executable_convolution_runtime_autotuning",
     gpu_target_config = "//xla/backends/gpu/target_config:specs/h100_sxm.txtpb",
@@ -6063,7 +6048,6 @@ xla_test(
         ":xla_aot_compile_test_gpu_executable",
         ":xla_aot_compile_test_gpu_executable_constant",
         ":xla_aot_compile_test_gpu_executable_convolution",
-        ":xla_aot_compile_test_gpu_executable_convolution_deviceless_cudnn",
         ":xla_aot_compile_test_gpu_executable_convolution_legacy_cache",
         ":xla_aot_compile_test_gpu_executable_hlo",
         ":xla_aot_compile_test_gpu_executable_hlo_legacy_cache",
@@ -6079,7 +6063,7 @@ xla_test(
         ":hlo_runner_pjrt",
         "//xla:literal",
         "//xla:literal_util",
-        "//xla/backends/gpu/tests:hlo_pjrt_gpu_test_base",
+        "//xla/tests:hlo_test_base",
         "//xla/tsl/platform:env",
         "//xla/tsl/platform:test",
         "@com_google_absl//absl/status:status_macros",
diff --git a/third_party/xla/xla/service/xla_aot_compile_gpu_test.cc b/third_party/xla/xla/service/xla_aot_compile_gpu_test.cc
index 113c0dc1f89..53606130117 100644
--- a/third_party/xla/xla/service/xla_aot_compile_gpu_test.cc
+++ b/third_party/xla/xla/service/xla_aot_compile_gpu_test.cc
@@ -20,15 +20,14 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/status/status_macros.h"
-#include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
-#include "xla/backends/gpu/tests/hlo_pjrt_gpu_test_base.h"
 #include "xla/literal.h"
 #include "xla/literal_util.h"
 #include "xla/service/hlo_runner_interface.h"
 #include "xla/service/hlo_runner_pjrt.h"
+#include "xla/tests/hlo_test_base.h"
 #include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/test.h"
 #include "tsl/platform/path.h"
@@ -37,18 +36,16 @@ namespace xla {
 namespace xla_compile {
 namespace {
 
-class XlaCompileTest : public gpu::HloPjRtGpuTestBase {
+class XlaCompileTest : public HloTestBase {
  public:
   void LoadAndRunExecutable(absl::string_view path_to_serialized_aot_result,
                             absl::Span<const Literal* const> args,
-                            const Literal& expected,
-                            absl::string_view target_suffix = "") {
+                            const Literal& expected) {
     const char* test_device = getenv("XLA_TEST_DEVICE");
     ASSERT_NE(test_device, nullptr) << "XLA_TEST_DEVICE is not set";
-    std::string path =
-        tsl::io::JoinPath(tsl::testing::XlaSrcRoot(), "service",
-                          absl::StrCat(path_to_serialized_aot_result, "_",
-                                       test_device, target_suffix));
+    std::string path = tsl::io::JoinPath(
+        tsl::testing::XlaSrcRoot(), "service",
+        absl::StrCat(path_to_serialized_aot_result, "_", test_device));
     std::string serialized_aot_result;
     ASSERT_OK(tsl::ReadFileToString(tsl::Env::Default(), path,
                                     &serialized_aot_result));
@@ -91,9 +88,8 @@ TEST_F(XlaCompileTest, LoadGpuExecutableWithConstant) {
                        expected);
 }
 
-void LoadAndRunConvolution(XlaCompileTest& test,
-                           absl::string_view path_to_serialized_aot_result,
-                           absl::string_view target_suffix = "") {
+// Should also cover the case of loading a GPU executable with a GEMM.
+TEST_F(XlaCompileTest, LoadGpuExecutableWithConvolution) {
   Literal input1 = LiteralUtil::CreateR4<float>(
       {{{{1.0, 2.0}, {3.0, 4.0}, {5.0, 6.0}, {7.0, 8.0}},
         {{11.0, 12.0}, {13.0, 14.0}, {15.0, 16.0}, {17.0, 18.0}},
@@ -107,31 +103,27 @@ void LoadAndRunConvolution(XlaCompileTest& test,
       {{1310.0}, {1466.0}, {1622.0}},
       {{2090.0}, {2246.0}, {2402.0}},
   }});
-  test.LoadAndRunExecutable(path_to_serialized_aot_result, {&input1, &input2},
-                            expected, target_suffix);
-}
-
-// Should also cover the case of loading a GPU executable with a GEMM.
-TEST_F(XlaCompileTest, LoadGpuExecutableWithConvolution) {
-  LoadAndRunConvolution(*this,
-                        "xla_aot_compile_test_gpu_executable_convolution");
+  LoadAndRunExecutable("xla_aot_compile_test_gpu_executable_convolution",
+                       {&input1, &input2}, expected);
 }
 
 TEST_F(XlaCompileTest, LoadGpuExecutableWithConvolutionLegacyCache) {
-  LoadAndRunConvolution(
-      *this, "xla_aot_compile_test_gpu_executable_convolution_legacy_cache");
-}
-
-TEST_F(XlaCompileTest, LoadGpuExecutableWithDevicelessCudnnConvolution) {
-  const char* test_device = getenv("XLA_TEST_DEVICE");
-  if (test_device == nullptr || absl::string_view(test_device) != "h100") {
-    GTEST_SKIP() << "Only compiled for h100.";
-  }
-  absl::string_view target_suffix =
-      absl::StrContains(device_description().name(), "MIG") ? "_mig" : "";
-  LoadAndRunConvolution(
-      *this, "xla_aot_compile_test_gpu_executable_convolution_deviceless_cudnn",
-      target_suffix);
+  Literal input1 = LiteralUtil::CreateR4<float>(
+      {{{{1.0, 2.0}, {3.0, 4.0}, {5.0, 6.0}, {7.0, 8.0}},
+        {{11.0, 12.0}, {13.0, 14.0}, {15.0, 16.0}, {17.0, 18.0}},
+        {{21.0, 22.0}, {23.0, 24.0}, {25.0, 26.0}, {27.0, 28.0}},
+        {{31.0, 32.0}, {33.0, 34.0}, {35.0, 36.0}, {37.0, 38.0}}}});
+  Literal input2 =
+      LiteralUtil::CreateR4<float>({{{{1.0}, {2.0}}, {{3.0}, {4.0}}},
+                                    {{{5.0}, {6.0}}, {{7.0}, {8.0}}},
+                                    {{{9.0}, {10.0}}, {{11.0}, {12.0}}}});
+  Literal expected = LiteralUtil::CreateR4<float>({{
+      {{1310.0}, {1466.0}, {1622.0}},
+      {{2090.0}, {2246.0}, {2402.0}},
+  }});
+  LoadAndRunExecutable(
+      "xla_aot_compile_test_gpu_executable_convolution_legacy_cache",
+      {&input1, &input2}, expected);
 }
 
 }  // namespace
diff --git a/third_party/xla/xla/service/xla_compile.bzl b/third_party/xla/xla/service/xla_compile.bzl
index 5655da66ab6..11096f37ed1 100644
--- a/third_party/xla/xla/service/xla_compile.bzl
+++ b/third_party/xla/xla/service/xla_compile.bzl
@@ -97,7 +97,7 @@ def xla_aot_compile_gpu(
         name,
         module,
         gpu_targets,
-        autotune_results = None,
+        autotune_results,
         xla_flags = ""):
     """Runs xla_compile to compile an MHLO, StableHLO or HLO module into an AotCompilationResult for GPU
 
@@ -105,7 +105,7 @@ def xla_aot_compile_gpu(
         name: The name of the build rule.
         module: The MHLO or StableHLO file to compile.
         gpu_targets: The list of gpu targets.
-        autotune_results: Optional AOT AutotuneResults or AutotuneCache file.
+        autotune_results: AOT AutotuneResults or AutotuneCache file.
         xla_flags: Additional XLA_FLAGS to set during compilation.
     """
 
@@ -120,18 +120,13 @@ def xla_aot_compile_gpu(
             " --platform=gpu" +
             " --gpu_target_config=$(location " + gpu_target_config_map[target] + ")"
         )
-        srcs = [module, gpu_target_config_map[target]]
-        flags = []
-        if autotune_results:
-            srcs.append(autotune_results)
-            flags.append("--xla_gpu_load_autotune_results_from=$(location " + autotune_results + ")")
+        flags = "--xla_gpu_load_autotune_results_from=$(location " + autotune_results + ")"
         if xla_flags:
-            flags.append(xla_flags)
-        if flags:
-            cmd = "XLA_FLAGS=\"" + " ".join(flags) + "\" " + cmd
+            flags = flags + " " + xla_flags
+        cmd = "XLA_FLAGS=\"" + flags + "\" " + cmd
         native.genrule(
             name = "gen_" + name + "_" + target,
-            srcs = srcs,
+            srcs = [module, gpu_target_config_map[target], autotune_results],
             outs = [name + "_" + target],
             cmd = cmd,
             tools = [xla_compile_tool],
diff --git a/third_party/xla/xla/stream_executor/cuda/BUILD b/third_party/xla/xla/stream_executor/cuda/BUILD
index 05f7bb33b62..149fa607070 100644
--- a/third_party/xla/xla/stream_executor/cuda/BUILD
+++ b/third_party/xla/xla/stream_executor/cuda/BUILD
@@ -569,12 +569,10 @@ cc_library(
         ":cudnn_api_wrappers",
         ":cudnn_frontend_helpers",
         ":cudnn_sdpa_score_mod",
-        "//xla:status_macros",
         "//xla/backends/gpu/target_config:cudnn_device_props",
         "//xla/stream_executor:activate_context",
         "//xla/stream_executor:data_type",
         "//xla/stream_executor:device_address",
-        "//xla/stream_executor:device_description",
         "//xla/stream_executor:dnn",
         "//xla/stream_executor:engine_options",
         "//xla/stream_executor:event_based_timer",
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc b/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc
index 51846165556..3adc96689a6 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc
@@ -54,7 +54,6 @@ limitations under the License.
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 #include "third_party/gpus/cuda/include/driver_types.h"
 #include "xla/backends/gpu/target_config/cudnn_device_props.h"
-#include "xla/status_macros.h"
 #include "xla/stream_executor/activate_context.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/cuda/cuda_diagnostics.h"
@@ -65,7 +64,6 @@ limitations under the License.
 #include "xla/stream_executor/cuda/cudnn_sdpa_score_mod.h"
 #include "xla/stream_executor/data_type.h"
 #include "xla/stream_executor/device_address.h"
-#include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/dnn.h"
 #include "xla/stream_executor/engine_options.h"
 #include "xla/stream_executor/event_based_timer.h"
@@ -3992,17 +3990,6 @@ static absl::StatusOr<cudnn_frontend::ExecutionPlan> RebuildExecutionPlan(
   return {std::move(plan)};
 }
 
-// Sets the SM version in the cuDNN graph. Without this, cuDNN will fallback to
-// querying the GPU, which doesn't work for deviceless compilation.
-absl::Status SetGraphSmVersion(cudnn_frontend::graph::Graph& graph,
-                               const DeviceDescription& gpu_device_info) {
-  const CudaComputeCapability* cc =
-      gpu_device_info.gpu_compute_capability().cuda_compute_capability();
-  TF_RET_CHECK(cc != nullptr);
-  graph.set_sm_version(cc->major * 10 + cc->minor);
-  return absl::OkStatus();
-}
-
 }  // namespace
 
 void FixDimsForRaggedOffset(std::vector<int64_t>& dims, int max_reg_per_batch) {
@@ -4050,7 +4037,6 @@ absl::StatusOr<CudnnGraph> GetCudnnFlashAttentionOperationGraph(
   graph.set_intermediate_data_type(cudnn_frontend::DataType_t::FLOAT)
       .set_io_data_type(ioDataType)
       .set_compute_data_type(cudnn_frontend::DataType_t::FLOAT);
-  ABSL_RETURN_IF_ERROR(SetGraphSmVersion(graph, gpu_device_info));
   int64_t uid = 0;
   auto next_uid = [&]() -> int64_t { return CuDnnTensorUID(uid++); };
 
@@ -4300,7 +4286,6 @@ absl::StatusOr<CudnnGraph> GetCudnnFlashAttentionF8OperationGraph(
   graph.set_intermediate_data_type(cudnn_frontend::DataType_t::FLOAT)
       .set_io_data_type(ioDataType)
       .set_compute_data_type(cudnn_frontend::DataType_t::FLOAT);
-  ABSL_RETURN_IF_ERROR(SetGraphSmVersion(graph, gpu_device_info));
 
   auto next_uid = [uid = 0]() mutable -> int { return CuDnnTensorUID(uid++); };
 
@@ -4435,7 +4420,6 @@ absl::StatusOr<CudnnGraph> GetCudnnFlashAttentionBackwardF8OperationGraph(
   graph.set_compute_data_type(cudnn_frontend::DataType_t::FLOAT)
       .set_intermediate_data_type(cudnn_frontend::DataType_t::FLOAT)
       .set_io_data_type(ioDataType);
-  ABSL_RETURN_IF_ERROR(SetGraphSmVersion(graph, gpu_device_info));
 
   auto next_uid = [uid = 0]() mutable -> int { return CuDnnTensorUID(uid++); };
 
@@ -4623,7 +4607,6 @@ absl::StatusOr<CudnnGraph> GetCudnnBlockScaledDotOperationGraph(
           << "\n global_scale: " << has_global_scale;
 
   cudnn_frontend::graph::Graph graph;
-  ABSL_RETURN_IF_ERROR(SetGraphSmVersion(graph, gpu_device_info));
   auto compute_type = cudnn_frontend::DataType_t::FLOAT;
   graph.set_compute_data_type(compute_type);
   graph.set_intermediate_data_type(compute_type);
@@ -4734,7 +4717,6 @@ absl::StatusOr<CudnnGraph> GetCudnnFlashAttentionBackwardOperationGraph(
 
   using cudnn_frontend::graph::Tensor_attributes;
   cudnn_frontend::graph::Graph graph;
-  ABSL_RETURN_IF_ERROR(SetGraphSmVersion(graph, gpu_device_info));
   if (!(q_desc.type() == k_desc.type() && k_desc.type() == p_desc.type() &&
         p_desc.type() == v_desc.type() && v_desc.type() == do_desc.type() &&
         do_desc.type() == dq_desc.type() && dq_desc.type() == dk_desc.type() &&
@@ -6939,8 +6921,6 @@ absl::Status CudnnGraph::Prepare(dnn::DnnSupport* dnn_support,
     return absl::OkStatus();
   };
 
-  ABSL_RETURN_IF_ERROR(SetGraphSmVersion(graph_, gpu_device_info));
-
   if (dnn_support) {
     const CudnnSupport& cudnn_support =
         static_cast<CudnnSupport&>(*dnn_support);

From 663d7692a9ce2386f820e5b06b93139cb555f7f0 Mon Sep 17 00:00:00 2001
From: Om Thakkar <om.thakkar@intel.com>
Date: Thu, 8 Oct 2026 06:19:02 -0700
Subject: [PATCH 0289/1000] PR #47992: [XLA:CPU][oneDNN] Move oneDNN tests out
 of restricted folder

Imported from GitHub PR https://github.com/openxla/xla/pull/47992

This PR moves the oneDNN tests out of the `restricted` folder (now that async oneDNN is included in the default XLA build) by migrating oneDNN tests from `HloTestBaseLegacy` to `HloTestBase`.
In addition, this PR removes the BUILD tags and explicitly disables `ynnpack` in the oneDNN tests so the tests run correctly in public CI moving forward.
Copybara import of the project:

--
bb936e717d7e021af11854d404725ebda5887a45 by Om Thakkar <om.thakkar@intel.com>:

move onednn tests out of restricted folder by porting over the tests to use HloTestBase

--
b250ea5721918bb75891701cfac78a8217096c31 by Om Thakkar <om.thakkar@intel.com>:

increasing the shard_count from 4 to 8 for onednn_matmul_test

Merging this change closes #47992

PiperOrigin-RevId: 995791944
---
 third_party/xla/xla/service/cpu/tests/BUILD   | 105 ++++-
 .../onednn_convolution_test.cc                |  69 ++--
 .../{restricted => }/onednn_fusion_test.cc    |  28 +-
 .../onednn_layer_norm_test.cc                 |  10 +-
 .../{restricted => }/onednn_matmul_test.cc    | 373 +++++++++---------
 .../{restricted => }/onednn_softmax_test.cc   |   9 +-
 .../xla/service/cpu/tests/restricted/BUILD    | 125 ------
 7 files changed, 372 insertions(+), 347 deletions(-)
 rename third_party/xla/xla/service/cpu/tests/{restricted => }/onednn_convolution_test.cc (92%)
 rename third_party/xla/xla/service/cpu/tests/{restricted => }/onednn_fusion_test.cc (93%)
 rename third_party/xla/xla/service/cpu/tests/{restricted => }/onednn_layer_norm_test.cc (97%)
 rename third_party/xla/xla/service/cpu/tests/{restricted => }/onednn_matmul_test.cc (87%)
 rename third_party/xla/xla/service/cpu/tests/{restricted => }/onednn_softmax_test.cc (97%)

diff --git a/third_party/xla/xla/service/cpu/tests/BUILD b/third_party/xla/xla/service/cpu/tests/BUILD
index 3c4427008c2..4fe29ff21c1 100644
--- a/third_party/xla/xla/service/cpu/tests/BUILD
+++ b/third_party/xla/xla/service/cpu/tests/BUILD
@@ -21,7 +21,8 @@ load("//xla:xla.default.bzl", "xla_cc_test")
 load("//xla/tests:build_defs.bzl", "xla_test")
 load("//xla/tsl:tsl.bzl", "tsl_copts")
 load("//xla/tsl:tsl.default.bzl", "filegroup")
-load("//xla/tsl/mkl:graph.bzl", "onednn_cc_test")
+load("//xla/tsl/mkl:build_defs.bzl", "if_graph_api")
+load("//xla/tsl/mkl:graph.bzl", "onednn_cc_test", "onednn_graph_cc_test")
 
 package(
     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
@@ -455,3 +456,105 @@ xla_test(
         "@com_google_googletest//:gtest",
     ],
 )
+
+onednn_cc_test(
+    name = "onednn_matmul_test",
+    srcs = ["onednn_matmul_test.cc"],
+    copts = tsl_copts(),
+    shard_count = 8,
+    deps = [
+        "//xla:error_spec",
+        "//xla/hlo/testlib:test",
+        "//xla/service:cpu_plugin",
+        "//xla/service/cpu:onednn_util",
+        "//xla/tests:hlo_interpreter_reference_mixin",
+        "//xla/tests:hlo_test_base",
+        "//xla/tests:pjrt_cpu_client_registry",
+        "//xla/tests:xla_internal_test_main",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+onednn_cc_test(
+    name = "onednn_convolution_test",
+    srcs = ["onednn_convolution_test.cc"],
+    copts = tsl_copts(),
+    deps = [
+        "//xla:literal",
+        "//xla:shape_util",
+        "//xla/hlo/testlib:filecheck",
+        "//xla/hlo/testlib:test",
+        "//xla/hlo/testlib:test_helpers",
+        "//xla/hlo/utils:hlo_matchers",
+        "//xla/service:cpu_plugin",
+        "//xla/service/cpu:onednn_contraction_rewriter",
+        "//xla/service/cpu:onednn_util",
+        "//xla/tests:hlo_interpreter_reference_mixin",
+        "//xla/tests:hlo_test_base",
+        "//xla/tests:pjrt_cpu_client_registry",
+        "//xla/tests:xla_internal_test_main",
+        "@com_google_absl//absl/strings",
+        "@tsl//tsl/platform:platform_port",
+    ],
+)
+
+onednn_cc_test(
+    name = "onednn_layer_norm_test",
+    srcs = ["onednn_layer_norm_test.cc"],
+    copts = tsl_copts(),
+    deps = [
+        "//xla:error_spec",
+        "//xla/hlo/testlib:test",
+        "//xla/service:cpu_plugin",
+        "//xla/service/cpu:onednn_util",
+        "//xla/tests:hlo_interpreter_reference_mixin",
+        "//xla/tests:hlo_test_base",
+        "//xla/tests:pjrt_cpu_client_registry",
+        "//xla/tests:xla_internal_test_main",
+    ],
+)
+
+onednn_cc_test(
+    name = "onednn_softmax_test",
+    srcs = ["onednn_softmax_test.cc"],
+    copts = tsl_copts(),
+    shard_count = if_graph_api(4, 1),
+    deps = [
+        "//xla:shape_util",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/testlib:pattern_matcher_gmock",
+        "//xla/hlo/testlib:test",
+        "//xla/service:cpu_plugin",
+        "//xla/service:pattern_matcher",
+        "//xla/service/cpu:backend_config_proto_cc",
+        "//xla/service/cpu:onednn_config_proto_cc",
+        "//xla/service/cpu:onednn_ops_rewriter",
+        "//xla/service/cpu:onednn_util",
+        "//xla/tests:hlo_interpreter_reference_mixin",
+        "//xla/tests:hlo_test_base",
+        "//xla/tests:pjrt_cpu_client_registry",
+        "//xla/tests:xla_internal_test_main",
+        "//xla/tsl/platform:statusor",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+onednn_graph_cc_test(
+    name = "onednn_fusion_test",
+    srcs = ["onednn_fusion_test.cc"],
+    deps = [
+        "//xla:error_spec",
+        "//xla:shape_util",
+        "//xla/backends/cpu:onednn_support",
+        "//xla/hlo/ir:hlo",
+        "//xla/service:cpu_plugin",
+        "//xla/service/cpu:onednn_util",
+        "//xla/tests:hlo_interpreter_reference_mixin",
+        "//xla/tests:hlo_test_base",
+        "//xla/tests:pjrt_cpu_client_registry",
+        "//xla/tsl/platform:test",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:platform_port",
+    ],
+)
diff --git a/third_party/xla/xla/service/cpu/tests/restricted/onednn_convolution_test.cc b/third_party/xla/xla/service/cpu/tests/onednn_convolution_test.cc
similarity index 92%
rename from third_party/xla/xla/service/cpu/tests/restricted/onednn_convolution_test.cc
rename to third_party/xla/xla/service/cpu/tests/onednn_convolution_test.cc
index cfdeb2576f8..1e7420a15bc 100644
--- a/third_party/xla/xla/service/cpu/tests/restricted/onednn_convolution_test.cc
+++ b/third_party/xla/xla/service/cpu/tests/onednn_convolution_test.cc
@@ -25,17 +25,20 @@ limitations under the License.
 #include "xla/service/cpu/onednn_contraction_rewriter.h"
 #include "xla/service/cpu/onednn_util.h"
 #include "xla/shape_util.h"
-#include "xla/tests/restricted/hlo_test_base_legacy.h"
+#include "xla/tests/hlo_interpreter_reference_mixin.h"
+#include "xla/tests/hlo_test_base.h"
 
 namespace xla {
 namespace cpu {
 
-class ConvolutionTest : public HloTestBaseLegacy,
+class ConvolutionTest : public HloInterpreterReferenceMixin<HloTestBase>,
                         public ::testing::WithParamInterface<PrimitiveType> {
  protected:
   DebugOptions GetDebugOptionsForTest() const override {
-    DebugOptions debug_options = HloTestBaseLegacy::GetDebugOptionsForTest();
+    DebugOptions debug_options =
+        HloInterpreterReferenceMixin::GetDebugOptionsForTest();
     debug_options.set_xla_cpu_experimental_onednn_custom_call(true);
+    debug_options.clear_xla_cpu_experimental_ynn_fusion_type();
     return debug_options;
   }
 
@@ -134,12 +137,11 @@ class ConvolutionTest : public HloTestBaseLegacy,
     return primitive_util::LowercasePrimitiveTypeName(PromotedDtype());
   }
 
-  void RunAndExpectExit(const absl::string_view outline, int signal) {
+  void RunAndExpectDeath(const absl::string_view outline) {
     const std::string convolution_module_str = absl::StrReplaceAll(
         outline,
         {{"$dtype", dtypeString_}, {"$pdtype", PromotedDtypeToString()}});
-    EXPECT_EXIT((void)Run(convolution_module_str, true),
-                ::testing::KilledBySignal(signal), "");
+    EXPECT_DEATH((void)Run(convolution_module_str, true), "");
   }
 
   void RunCompareAndMatchOptimizedHlo(
@@ -247,9 +249,8 @@ TEST_P(ConvolutionTest, Conv3DReluTest) {
     arg.1 = $dtype[3,3,3,28,64] parameter(1)
     conv = $dtype[15,4,5,5,64] convolution(arg.0, arg.1),
           window={size=3x3x3 pad=1_1x1_1x1_1}, dim_labels=b012f_012io->b012f
-    const.1 = $pdtype[] constant(0)
-    convert.0 = $dtype[] convert(const.1)
-    bcast.2 = $dtype[15,4,5,5,64] broadcast(convert.0), dimensions={}
+    const.1 = $dtype[] constant(0)
+    bcast.2 = $dtype[15,4,5,5,64] broadcast(const.1), dimensions={}
     ROOT maximum.1 = $dtype[15,4,5,5,64] maximum(conv, bcast.2)
 })";
 
@@ -257,6 +258,13 @@ TEST_P(ConvolutionTest, Conv3DReluTest) {
 }
 
 TEST_P(ConvolutionTest, Conv2DWithBiasAndReluTest) {
+  if (dtype_ == BF16) {
+    // TODO(intel-tf): Enable for BF16 once OneDnnContractionRewriter handles
+    // the BF16 -> F32 -> BF16 convert pair left after fusing BIAS when the
+    // following activation (kMaximum) is not promoted to F32.
+    GTEST_SKIP() << "Skipping BF16 Bias+ReLU fusion due to intermediate "
+                    "convert mismatch.";
+  }
   const absl::string_view outline = R"(
   HloModule convolution.bias.relu.test
 
@@ -268,9 +276,8 @@ TEST_P(ConvolutionTest, Conv2DWithBiasAndReluTest) {
     const.0 = $dtype[10] constant(15)
     bcast.1 = $dtype[1,11,11,10] broadcast(const.0), dimensions={3}
     add.0 = $dtype[1,11,11,10] add(convolution.0, bcast.1)
-    const.1 = $pdtype[] constant(0)
-    convert.0 = $dtype[] convert(const.1)
-    bcast.2 = $dtype[1,11,11,10] broadcast(convert.0), dimensions={}
+    const.1 = $dtype[] constant(0)
+    bcast.2 = $dtype[1,11,11,10] broadcast(const.1), dimensions={}
     ROOT maximum.1 = $dtype[1,11,11,10] maximum(add.0, bcast.2)
   })";
 
@@ -321,7 +328,7 @@ TEST_P(ConvolutionTest, ConvInsufficientScratchTest) {
         }
   })";
 
-  RunAndExpectExit(outline, SIGABRT);
+  RunAndExpectDeath(outline);
 }
 
 TEST_P(ConvolutionTest, Conv2DWithBinaryAddTest) {
@@ -376,9 +383,8 @@ TEST_P(ConvolutionTest, Conv2DWithReluSumAndBinaryAddTest) {
     arg0.3 = $dtype[1,11,11,10] parameter(2)
     convolution.0 = $dtype[1,11,11,10] convolution(arg0.1, arg0.2),
           window={size=8x8 stride=2x2 pad=3_3x3_3}, dim_labels=b01f_01io->b01f
-    const.1 = $pdtype[] constant(0)
-    convert.0 = $dtype[] convert(const.1)
-    bcast.2 = $dtype[1,11,11,10] broadcast(convert.0), dimensions={}
+    const.1 = $dtype[] constant(0)
+    bcast.2 = $dtype[1,11,11,10] broadcast(const.1), dimensions={}
     maximum.1 = $dtype[1,11,11,10] maximum(convolution.0, bcast.2)
     const.2 = $dtype[1,11,11,10] constant({...})
     add.1 = $dtype[1,11,11,10] add(maximum.1, const.2)
@@ -506,6 +512,13 @@ TEST_P(ConvolutionTest, Conv2DWithLinearAndBinaryAddTest) {
 }
 
 TEST_P(ConvolutionTest, Conv3DWithBiasAndRelu6Test) {
+  if (dtype_ == BF16) {
+    // TODO(intel-tf): Enable for BF16 once OneDnnContractionRewriter handles
+    // the BF16 -> F32 -> BF16 convert pair left after fusing BIAS when the
+    // following activation (kClamp) is not promoted to F32.
+    GTEST_SKIP() << "Skipping BF16 Bias+ReLU6 fusion due to intermediate "
+                    "convert mismatch.";
+  }
   const absl::string_view outline = R"(
   HloModule convolution.test.bias.relu6
 
@@ -517,12 +530,10 @@ TEST_P(ConvolutionTest, Conv3DWithBiasAndRelu6Test) {
     bias = $dtype[64] parameter(2)
     broadcasted_bias = $dtype[15,4,5,5,64] broadcast(bias), dimensions={4}
     add = $dtype[15,4,5,5,64] add(conv, broadcasted_bias)
-    const.0 = $pdtype[] constant(0)
-    convert.0 = $dtype[] convert(const.0)
-    broadcast.0 = $dtype[15,4,5,5,64] broadcast(convert.0), dimensions={}
-    const.1 = $pdtype[] constant(6)
-    convert.1 = $dtype[] convert(const.1)
-    broadcast.1 = $dtype[15,4,5,5,64] broadcast(convert.1), dimensions={}
+    const.0 = $dtype[] constant(0)
+    broadcast.0 = $dtype[15,4,5,5,64] broadcast(const.0), dimensions={}
+    const.1 = $dtype[] constant(6)
+    broadcast.1 = $dtype[15,4,5,5,64] broadcast(const.1), dimensions={}
     ROOT clamp.0 = $dtype[15,4,5,5,64] clamp(broadcast.0, add, broadcast.1)
 })";
 
@@ -530,6 +541,13 @@ TEST_P(ConvolutionTest, Conv3DWithBiasAndRelu6Test) {
 }
 
 TEST_P(ConvolutionTest, Conv2DWithBiasAndSigmoidTest) {
+  if (dtype_ == BF16) {
+    // TODO(intel-tf): Enable for BF16 once OneDnnContractionRewriter handles
+    // intermediate converts when kNegate remains in BF16 while kAdd/kExp/
+    // kDivide are promoted to F32.
+    GTEST_SKIP() << "Skipping BF16 Bias+Sigmoid fusion due to intermediate "
+                    "convert mismatch.";
+  }
   const absl::string_view outline = R"(
   HloModule convolution.bias.sigmoid.test
 
@@ -742,6 +760,13 @@ TEST_P(ConvolutionTest, Conv2DWithBiasAndGeluExactPattern2Test) {
 }
 
 TEST_P(ConvolutionTest, Conv2DWithBiasAndSwishTest) {
+  if (dtype_ == BF16) {
+    // TODO(intel-tf): Enable for BF16 once OneDnnContractionRewriter handles
+    // intermediate converts when kNegate remains in BF16 while kAdd/kExp/
+    // kDivide/kMultiply are promoted to F32.
+    GTEST_SKIP() << "Skipping BF16 Bias+Swish fusion due to intermediate "
+                    "convert mismatch.";
+  }
   const absl::string_view outline = R"(
   HloModule convolution.test.with.bias.swish
 
diff --git a/third_party/xla/xla/service/cpu/tests/restricted/onednn_fusion_test.cc b/third_party/xla/xla/service/cpu/tests/onednn_fusion_test.cc
similarity index 93%
rename from third_party/xla/xla/service/cpu/tests/restricted/onednn_fusion_test.cc
rename to third_party/xla/xla/service/cpu/tests/onednn_fusion_test.cc
index 5c05c3d6096..7f9f38f5dd1 100644
--- a/third_party/xla/xla/service/cpu/tests/restricted/onednn_fusion_test.cc
+++ b/third_party/xla/xla/service/cpu/tests/onednn_fusion_test.cc
@@ -13,22 +13,27 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <string>
+#include <vector>
+
 #include <gtest/gtest.h>
-#include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_replace.h"
 #include "absl/strings/string_view.h"
 #include "xla/backends/cpu/onednn_support.h"
 #include "xla/error_spec.h"
+#include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/primitive_util.h"
 #include "xla/service/cpu/onednn_util.h"
-#include "xla/tests/restricted/hlo_test_base_legacy.h"
+#include "xla/tests/hlo_interpreter_reference_mixin.h"
+#include "xla/tests/hlo_test_base.h"
 #include "xla/tsl/platform/test.h"
 #include "tsl/platform/cpu_info.h"
 
 namespace xla::cpu {
 namespace {
 
-inline constexpr bool IsOneDnnGraphEnabled() {
+constexpr bool IsOneDnnGraphEnabled() {
 #if defined(XLA_ONEDNN_USE_GRAPH_API)
   // Some Aarch64 CPUs have failures. Only test on x86 for now.
   return tsl::port::IsX86CPU();
@@ -42,9 +47,16 @@ struct OneDnnFusionTestParams {
 };
 
 class OneDnnFusionTestBase
-    : public HloTestBaseLegacy,
+    : public HloInterpreterReferenceMixin<HloTestBase>,
       public ::testing::WithParamInterface<OneDnnFusionTestParams> {
  protected:
+  DebugOptions GetDebugOptionsForTest() const override {
+    DebugOptions debug_options =
+        HloInterpreterReferenceMixin::GetDebugOptionsForTest();
+    debug_options.clear_xla_cpu_experimental_ynn_fusion_type();
+    return debug_options;
+  }
+
   void SetUp() override {
     OneDnnFusionTestParams params = GetParam();
     data_type_ = params.dtype;
@@ -89,7 +101,7 @@ class OneDnnFusionBinaryOpTest : public OneDnnFusionTestBase {
   }
 
  private:
-  const std::string GetBinaryOpHLOTemplate() {
+  std::string GetBinaryOpHLOTemplate() {
     return R"(
     HloModule binary_op
 
@@ -108,7 +120,7 @@ class OneDnnFusionBinaryOpTest : public OneDnnFusionTestBase {
     })";
   }
 
-  const std::string GetMatMulHLOTemplate() {
+  std::string GetMatMulHLOTemplate() {
     return R"(
     HloModule matmul
 
@@ -243,7 +255,9 @@ std::vector<OneDnnFusionTestParams> GetOneDnnFusionFuseBinaryTestSpecs() {
   for (const auto& dtype : {PrimitiveType::F32}) {
     for (const auto& op_type : GetOneDnnSupportedBinaryOpsStrings()) {
       // oneDNN does not support fusing two dot instructions
-      if (op_type == HloOpcodeString(HloOpcode::kDot)) continue;
+      if (op_type == HloOpcodeString(HloOpcode::kDot)) {
+        continue;
+      }
       specs.push_back({dtype, std::string(op_type)});
     }
   }
diff --git a/third_party/xla/xla/service/cpu/tests/restricted/onednn_layer_norm_test.cc b/third_party/xla/xla/service/cpu/tests/onednn_layer_norm_test.cc
similarity index 97%
rename from third_party/xla/xla/service/cpu/tests/restricted/onednn_layer_norm_test.cc
rename to third_party/xla/xla/service/cpu/tests/onednn_layer_norm_test.cc
index e840c346b91..dac31618bbf 100644
--- a/third_party/xla/xla/service/cpu/tests/restricted/onednn_layer_norm_test.cc
+++ b/third_party/xla/xla/service/cpu/tests/onednn_layer_norm_test.cc
@@ -18,16 +18,19 @@ limitations under the License.
 #include "xla/error_spec.h"
 #include "xla/hlo/testlib/test.h"
 #include "xla/service/cpu/onednn_util.h"
-#include "xla/tests/restricted/hlo_test_base_legacy.h"
+#include "xla/tests/hlo_interpreter_reference_mixin.h"
+#include "xla/tests/hlo_test_base.h"
 
 namespace xla {
 namespace {
 
-class LayerNormTest : public HloTestBaseLegacy {
+class LayerNormTest : public HloInterpreterReferenceMixin<HloTestBase> {
  protected:
   DebugOptions GetDebugOptionsForTest() const override {
-    DebugOptions debug_options = HloTestBaseLegacy::GetDebugOptionsForTest();
+    DebugOptions debug_options =
+        HloInterpreterReferenceMixin::GetDebugOptionsForTest();
     debug_options.set_xla_cpu_experimental_onednn_custom_call(true);
+    debug_options.clear_xla_cpu_experimental_ynn_fusion_type();
     return debug_options;
   }
 
@@ -349,4 +352,3 @@ TEST_F(LayerNormTest, LayerNormTest1_BF16) {
 
 }  // namespace
 }  // namespace xla
-
diff --git a/third_party/xla/xla/service/cpu/tests/restricted/onednn_matmul_test.cc b/third_party/xla/xla/service/cpu/tests/onednn_matmul_test.cc
similarity index 87%
rename from third_party/xla/xla/service/cpu/tests/restricted/onednn_matmul_test.cc
rename to third_party/xla/xla/service/cpu/tests/onednn_matmul_test.cc
index 401e5b6bb21..96d7cd379ca 100644
--- a/third_party/xla/xla/service/cpu/tests/restricted/onednn_matmul_test.cc
+++ b/third_party/xla/xla/service/cpu/tests/onednn_matmul_test.cc
@@ -19,16 +19,19 @@ limitations under the License.
 #include "xla/error_spec.h"
 #include "xla/hlo/testlib/test.h"
 #include "xla/service/cpu/onednn_util.h"
-#include "xla/tests/restricted/hlo_test_base_legacy.h"
+#include "xla/tests/hlo_interpreter_reference_mixin.h"
+#include "xla/tests/hlo_test_base.h"
 
 namespace xla {
 namespace cpu {
 
-class MatmulTest : public HloTestBaseLegacy {
+class MatmulTest : public HloInterpreterReferenceMixin<HloTestBase> {
  protected:
   DebugOptions GetDebugOptionsForTest() const override {
-    DebugOptions debug_options = HloTestBaseLegacy::GetDebugOptionsForTest();
+    DebugOptions debug_options =
+        HloInterpreterReferenceMixin::GetDebugOptionsForTest();
     debug_options.set_xla_cpu_experimental_onednn_custom_call(true);
+    debug_options.clear_xla_cpu_experimental_ynn_fusion_type();
     return debug_options;
   }
 
@@ -408,15 +411,15 @@ TEST_F(MatmulTest, SimpleTestF32WithBiasAsParameter3) {
   HloModule matmul.biasadd.test.f32
 
   ENTRY matmul.biasadd.test.f32 {
-    arg0.1 = f32[16,128,768] parameter(0), sharding={replicated}
-    arg0.2 = f32[768,768] parameter(1), sharding={replicated}
-    dot.84 = f32[16,128,768] dot(arg0.1, arg0.2), lhs_contracting_dims={2}, rhs_contracting_dims={0}
-    arg0.3 = f32[768]{0} parameter(2), sharding={replicated}
-    reshape.85 = f32[1,1,768] reshape(arg0.3)
-    broadcast.86 = f32[1,1,768] broadcast(reshape.85), dimensions={0,1,2}
-    reshape.87 = f32[768]{0} reshape(broadcast.86)
-    broadcast.88 = f32[16,128,768] broadcast(reshape.87), dimensions={2}
-    ROOT add.89 = f32[16,128,768] add(dot.84, broadcast.88)
+    arg0.1 = f32[16,64,128] parameter(0), sharding={replicated}
+    arg0.2 = f32[128,128] parameter(1), sharding={replicated}
+    dot.84 = f32[16,64,128] dot(arg0.1, arg0.2), lhs_contracting_dims={2}, rhs_contracting_dims={0}
+    arg0.3 = f32[128]{0} parameter(2), sharding={replicated}
+    reshape.85 = f32[1,1,128] reshape(arg0.3)
+    broadcast.86 = f32[1,1,128] broadcast(reshape.85), dimensions={0,1,2}
+    reshape.87 = f32[128]{0} reshape(broadcast.86)
+    broadcast.88 = f32[16,64,128] broadcast(reshape.87), dimensions={2}
+    ROOT add.89 = f32[16,64,128] add(dot.84, broadcast.88)
   })";
 
   EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-4, 1e-4}));
@@ -704,25 +707,25 @@ TEST_F(MatmulTest, BiasAndExactGELUTestF32) {
   const char* matmul_module_str = R"(
   HloModule matmul.test.f32
   ENTRY matmul.test.f32 {
-    arg.0 = f32[6304,768] parameter(0), parameter_replication={false}
-    arg.1 = f32[768,3072] parameter(1), parameter_replication={false}
-    dot.378 = f32[6304,3072] dot(arg.0, arg.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
-    reshape.11 = f32[32,197,3072]reshape(dot.378)
-    constant.381 = f32[3072] constant(0.3)
-    broadcast.382 = f32[32,197,3072] broadcast(constant.381), dimensions={2}
-    add.383 = f32[32,197,3072] add(reshape.11, broadcast.382)
+    arg.0 = f32[1024,64] parameter(0), parameter_replication={false}
+    arg.1 = f32[64,256] parameter(1), parameter_replication={false}
+    dot.378 = f32[1024,256] dot(arg.0, arg.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    reshape.11 = f32[32,32,256]reshape(dot.378)
+    constant.381 = f32[256] constant(0.3)
+    broadcast.382 = f32[32,32,256] broadcast(constant.381), dimensions={2}
+    add.383 = f32[32,32,256] add(reshape.11, broadcast.382)
     constant.384 = f32[] constant(0.707106769)
-    broadcast.385 = f32[32,197,3072] broadcast(constant.384), dimensions={}
-    multiply.386 = f32[32,197,3072] multiply(broadcast.385, add.383)
-    erf.387 = f32[32,197,3072] erf(multiply.386)
+    broadcast.385 = f32[32,32,256] broadcast(constant.384), dimensions={}
+    multiply.386 = f32[32,32,256] multiply(broadcast.385, add.383)
+    erf.387 = f32[32,32,256] erf(multiply.386)
     constant.388 = f32[] constant(1)
-    broadcast.389 = f32[32,197,3072] broadcast(constant.388), dimensions={}
-    add.390 = f32[32,197,3072] add(erf.387, broadcast.389)
+    broadcast.389 = f32[32,32,256] broadcast(constant.388), dimensions={}
+    add.390 = f32[32,32,256] add(erf.387, broadcast.389)
     constant.391 = f32[] constant(0.5)
-    broadcast.392 = f32[32,197,3072] broadcast(constant.391)
-    multiply.393 = f32[32,197,3072] multiply(add.390, broadcast.392)
-    multiply.394 = f32[32,197,3072] multiply(multiply.393, add.383)
-    ROOT out = f32[6304,3072] reshape(multiply.394)
+    broadcast.392 = f32[32,32,256] broadcast(constant.391)
+    multiply.393 = f32[32,32,256] multiply(add.390, broadcast.392)
+    multiply.394 = f32[32,32,256] multiply(multiply.393, add.383)
+    ROOT out = f32[1024,256] reshape(multiply.394)
   })";
 
   EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-4, 1e-4}));
@@ -736,26 +739,26 @@ TEST_F(MatmulTest, BiasAndExactGELUTestBF16) {
   const char* matmul_module_str = R"(
   HloModule matmul.test.f32
   ENTRY matmul.test.f32 {
-    arg.0 = f32[6304,768] parameter(0), parameter_replication={false}
-    convert.0 = bf16[6304,768] convert(arg.0)
-    arg.1 = f32[768,3072] parameter(1), parameter_replication={false}
-    convert.1 = bf16[768,3072] convert(arg.1)
-    dot.378 = bf16[6304,3072] dot(convert.0, convert.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
-    convert.2 = f32[6304,3072] convert(dot.378)
-    constant.381 = f32[3072] constant(0.3)
-    broadcast.382 = f32[6304,3072] broadcast(constant.381), dimensions={1}
-    add.383 = f32[6304,3072] add(convert.2, broadcast.382)
+    arg.0 = f32[1024,64] parameter(0), parameter_replication={false}
+    convert.0 = bf16[1024,64] convert(arg.0)
+    arg.1 = f32[64,256] parameter(1), parameter_replication={false}
+    convert.1 = bf16[64,256] convert(arg.1)
+    dot.378 = bf16[1024,256] dot(convert.0, convert.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    convert.2 = f32[1024,256] convert(dot.378)
+    constant.381 = f32[256] constant(0.3)
+    broadcast.382 = f32[1024,256] broadcast(constant.381), dimensions={1}
+    add.383 = f32[1024,256] add(convert.2, broadcast.382)
     constant.384 = f32[] constant(0.707106769)
-    broadcast.385 = f32[6304,3072] broadcast(constant.384), dimensions={}
-    multiply.386 = f32[6304,3072] multiply(broadcast.385, add.383)
-    erf.387 = f32[6304,3072] erf(multiply.386)
+    broadcast.385 = f32[1024,256] broadcast(constant.384), dimensions={}
+    multiply.386 = f32[1024,256] multiply(broadcast.385, add.383)
+    erf.387 = f32[1024,256] erf(multiply.386)
     constant.388 = f32[] constant(1)
-    broadcast.389 = f32[6304,3072] broadcast(constant.388), dimensions={}
-    add.390 = f32[6304,3072] add(erf.387, broadcast.389)
+    broadcast.389 = f32[1024,256] broadcast(constant.388), dimensions={}
+    add.390 = f32[1024,256] add(erf.387, broadcast.389)
     constant.391 = f32[] constant(0.5)
-    broadcast.392 = f32[6304,3072] broadcast(constant.391)
-    multiply.393 = f32[6304,3072] multiply(add.390, broadcast.392)
-    ROOT out = f32[6304,3072] multiply(multiply.393, add.383)
+    broadcast.392 = f32[1024,256] broadcast(constant.391)
+    multiply.393 = f32[1024,256] multiply(add.390, broadcast.392)
+    ROOT out = f32[1024,256] multiply(multiply.393, add.383)
   })";
 
   EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-2, 1e-2}));
@@ -769,27 +772,27 @@ TEST_F(MatmulTest, BiasAndExactJaxGELUTestBF16) {
   const char* matmul_module_str = R"(
   HloModule matmul.test.f32
   ENTRY matmul.test.f32 {
-    arg.0 = f32[6304,768] parameter(0), parameter_replication={false}
-    convert.0 = bf16[6304,768] convert(arg.0)
-    arg.1 = f32[768,3072] parameter(1), parameter_replication={false}
-    convert.1 = bf16[768,3072] convert(arg.1)
-    dot.378 = bf16[6304,3072] dot(convert.0, convert.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
-    convert.2 = f32[6304,3072] convert(dot.378)
-    reshape.0 = f32[32,197,3072] reshape(convert.2)
-    constant.381 = f32[3072] constant(0.3)
-    broadcast.382 = f32[32,197,3072] broadcast(constant.381), dimensions={2}
-    add.383 = f32[32,197,3072] add(reshape.0, broadcast.382)
+    arg.0 = f32[1024,64] parameter(0), parameter_replication={false}
+    convert.0 = bf16[1024,64] convert(arg.0)
+    arg.1 = f32[64,256] parameter(1), parameter_replication={false}
+    convert.1 = bf16[64,256] convert(arg.1)
+    dot.378 = bf16[1024,256] dot(convert.0, convert.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    convert.2 = f32[1024,256] convert(dot.378)
+    reshape.0 = f32[32,32,256] reshape(convert.2)
+    constant.381 = f32[256] constant(0.3)
+    broadcast.382 = f32[32,32,256] broadcast(constant.381), dimensions={2}
+    add.383 = f32[32,32,256] add(reshape.0, broadcast.382)
     constant.384 = f32[] constant(0.707182348)
-    broadcast.385 = f32[32,197,3072] broadcast(constant.384), dimensions={}
-    multiply.386 = f32[32,197,3072] multiply(broadcast.385, add.383)
-    erf.387 = f32[32,197,3072] erf(multiply.386)
+    broadcast.385 = f32[32,32,256] broadcast(constant.384), dimensions={}
+    multiply.386 = f32[32,32,256] multiply(broadcast.385, add.383)
+    erf.387 = f32[32,32,256] erf(multiply.386)
     constant.388 = f32[] constant(1)
-    broadcast.389 = f32[32,197,3072] broadcast(constant.388), dimensions={}
-    add.390 = f32[32,197,3072] add(erf.387, broadcast.389)
-    multiply.393 = f32[32,197,3072] multiply(add.390, add.383)
+    broadcast.389 = f32[32,32,256] broadcast(constant.388), dimensions={}
+    add.390 = f32[32,32,256] add(erf.387, broadcast.389)
+    multiply.393 = f32[32,32,256] multiply(add.390, add.383)
     constant.391 = f32[] constant(0.5)
-    broadcast.392 = f32[32,197,3072] broadcast(constant.391)
-    ROOT multiply.394 = f32[32,197,3072] multiply(multiply.393, broadcast.392)
+    broadcast.392 = f32[32,32,256] broadcast(constant.391)
+    ROOT multiply.394 = f32[32,32,256] multiply(multiply.393, broadcast.392)
   })";
 
   EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-2, 1e-2}));
@@ -837,23 +840,23 @@ TEST_F(MatmulTest, BiasAndExactGELUTestF16) {
   const char* matmul_module_str = R"(
   HloModule matmul.test.f16
   ENTRY matmul.test.f16 {
-    arg.0 = f16[6304,768] parameter(0), parameter_replication={false}
-    arg.1 = f16[768,3072] parameter(1), parameter_replication={false}
-    dot.378 = f16[6304,3072] dot(arg.0, arg.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
-    constant.381 = f16[3072] constant(0.3)
-    broadcast.382 = f16[6304,3072] broadcast(constant.381), dimensions={1}
-    add.383 = f16[6304,3072] add(dot.378, broadcast.382)
+    arg.0 = f16[1024,64] parameter(0), parameter_replication={false}
+    arg.1 = f16[64,256] parameter(1), parameter_replication={false}
+    dot.378 = f16[1024,256] dot(arg.0, arg.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    constant.381 = f16[256] constant(0.3)
+    broadcast.382 = f16[1024,256] broadcast(constant.381), dimensions={1}
+    add.383 = f16[1024,256] add(dot.378, broadcast.382)
     constant.384 = f16[] constant(0.707106769)
-    broadcast.385 = f16[6304,3072] broadcast(constant.384), dimensions={}
-    multiply.386 = f16[6304,3072] multiply(broadcast.385, add.383)
-    erf.387 = f16[6304,3072] erf(multiply.386)
+    broadcast.385 = f16[1024,256] broadcast(constant.384), dimensions={}
+    multiply.386 = f16[1024,256] multiply(broadcast.385, add.383)
+    erf.387 = f16[1024,256] erf(multiply.386)
     constant.388 = f16[] constant(1)
-    broadcast.389 = f16[6304,3072] broadcast(constant.388), dimensions={}
-    add.390 = f16[6304,3072] add(erf.387, broadcast.389)
+    broadcast.389 = f16[1024,256] broadcast(constant.388), dimensions={}
+    add.390 = f16[1024,256] add(erf.387, broadcast.389)
     constant.391 = f16[] constant(0.5)
-    broadcast.392 = f16[6304,3072] broadcast(constant.391)
-    multiply.393 = f16[6304,3072] multiply(add.390, broadcast.392)
-    ROOT out = f16[6304,3072] multiply(multiply.393, add.383)
+    broadcast.392 = f16[1024,256] broadcast(constant.391)
+    multiply.393 = f16[1024,256] multiply(add.390, broadcast.392)
+    ROOT out = f16[1024,256] multiply(multiply.393, add.383)
   })";
 
   EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-2, 1e-2}));
@@ -929,18 +932,18 @@ TEST_F(MatmulTest, SimpleBiasTestBF16_PARAM_F32) {
   HloModule jit_apply
 
   ENTRY matmul.test.bf16 {
-    Arg_2.3 = f32[16,128,768] parameter(2), sharding={replicated}
-    convert.4 = bf16[16,128,768] convert(Arg_2.3)
-    Arg_1.2 = f32[768,3072] parameter(1), sharding={replicated}
-    convert.5 = bf16[768,3072] convert(Arg_1.2)
-    dot.7 = bf16[16,128,3072] dot(convert.4, convert.5), lhs_contracting_dims={2}, rhs_contracting_dims={0}
-    Arg_0.1 = f32[3072] parameter(0), sharding={replicated}
-    convert.6 = bf16[3072] convert(Arg_0.1)
-    reshape.8 = bf16[1,1,3072] reshape(convert.6)
-    broadcast.9 = bf16[1,1,3072] broadcast(reshape.8), dimensions={0,1,2}
-    reshape.10 = bf16[3072] reshape(broadcast.9)
-    broadcast.11 = bf16[16,128,3072] broadcast(reshape.10), dimensions={2}
-    ROOT add.12 = bf16[16,128,3072] add(dot.7, broadcast.11)
+    Arg_2.3 = f32[16,64,64] parameter(2), sharding={replicated}
+    convert.4 = bf16[16,64,64] convert(Arg_2.3)
+    Arg_1.2 = f32[64,256] parameter(1), sharding={replicated}
+    convert.5 = bf16[64,256] convert(Arg_1.2)
+    dot.7 = bf16[16,64,256] dot(convert.4, convert.5), lhs_contracting_dims={2}, rhs_contracting_dims={0}
+    Arg_0.1 = f32[256] parameter(0), sharding={replicated}
+    convert.6 = bf16[256] convert(Arg_0.1)
+    reshape.8 = bf16[1,1,256] reshape(convert.6)
+    broadcast.9 = bf16[1,1,256] broadcast(reshape.8), dimensions={0,1,2}
+    reshape.10 = bf16[256] reshape(broadcast.9)
+    broadcast.11 = bf16[16,64,256] broadcast(reshape.10), dimensions={2}
+    ROOT add.12 = bf16[16,64,256] add(dot.7, broadcast.11)
   })";
 
   EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-2, 1e-2}));
@@ -956,16 +959,16 @@ TEST_F(MatmulTest, SimpleBiasTestBF16_PARAM_BF16) {
   HloModule jit_apply
 
   ENTRY matmul.test.bf16 {
-    Arg_2.3 = f32[16,128,768] parameter(2), sharding={replicated}
-    convert.4 = bf16[16,128,768] convert(Arg_2.3)
-    Arg_1.2 = bf16[768,3072] parameter(1), sharding={replicated}
-    dot.5 = bf16[16,128,3072] dot(convert.4, Arg_1.2), lhs_contracting_dims={2}, rhs_contracting_dims={0}
-    Arg_0.1 = bf16[3072] parameter(0), sharding={replicated}
-    reshape.6 = bf16[1,1,3072] reshape(Arg_0.1)
-    broadcast.7 = bf16[1,1,3072] broadcast(reshape.6), dimensions={0,1,2}
-    reshape.8 = bf16[3072] reshape(broadcast.7)
-    broadcast.9 = bf16[16,128,3072] broadcast(reshape.8), dimensions={2}
-    ROOT add.10 = bf16[16,128,3072] add(dot.5, broadcast.9)
+    Arg_2.3 = f32[16,64,64] parameter(2), sharding={replicated}
+    convert.4 = bf16[16,64,64] convert(Arg_2.3)
+    Arg_1.2 = bf16[64,256] parameter(1), sharding={replicated}
+    dot.5 = bf16[16,64,256] dot(convert.4, Arg_1.2), lhs_contracting_dims={2}, rhs_contracting_dims={0}
+    Arg_0.1 = bf16[256] parameter(0), sharding={replicated}
+    reshape.6 = bf16[1,1,256] reshape(Arg_0.1)
+    broadcast.7 = bf16[1,1,256] broadcast(reshape.6), dimensions={0,1,2}
+    reshape.8 = bf16[256] reshape(broadcast.7)
+    broadcast.9 = bf16[16,64,256] broadcast(reshape.8), dimensions={2}
+    ROOT add.10 = bf16[16,64,256] add(dot.5, broadcast.9)
   })";
 
   EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-2, 1e-2}));
@@ -977,12 +980,12 @@ TEST_F(MatmulTest, DivisionByConstantWithEltwiseLinearF32) {
   HloModule matmul.divide.test.1
 
   ENTRY matmul.divide.test.f32 {
-    Arg_4.5 = f32[16,128,768] parameter(0), sharding={replicated}
-    Arg_2.3 = f32[768,12,64] parameter(1), sharding={replicated}
-    onednn.matmul.0 = f32[16,128,12,64] dot(Arg_4.5, Arg_2.3), lhs_contracting_dims={2}, rhs_contracting_dims={0}
+    Arg_4.5 = f32[16,64,128] parameter(0), sharding={replicated}
+    Arg_2.3 = f32[128,4,32] parameter(1), sharding={replicated}
+    onednn.matmul.0 = f32[16,64,4,32] dot(Arg_4.5, Arg_2.3), lhs_contracting_dims={2}, rhs_contracting_dims={0}
     constant.8 = f32[] constant(8)
-    broadcast.9 = f32[16,128,12,64] broadcast(constant.8), dimensions={}
-    ROOT divide.16 = f32[16,128,12,64] divide(onednn.matmul.0, broadcast.9)
+    broadcast.9 = f32[16,64,4,32] broadcast(constant.8), dimensions={}
+    ROOT divide.16 = f32[16,64,4,32] divide(onednn.matmul.0, broadcast.9)
   })";
 
   EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec(1e-4, 1e-4)));
@@ -1009,18 +1012,18 @@ TEST_F(MatmulTest, SimpleBiasTestF16_PARAM_F32) {
   HloModule jit_apply
 
   ENTRY matmul.test.f16 {
-    Arg_2.3 = f32[16,128,768] parameter(2), sharding={replicated}
-    convert.4 = f16[16,128,768] convert(Arg_2.3)
-    Arg_1.2 = f32[768,3072] parameter(1), sharding={replicated}
-    convert.5 = f16[768,3072] convert(Arg_1.2)
-    dot.7 = f16[16,128,3072] dot(convert.4, convert.5), lhs_contracting_dims={2}, rhs_contracting_dims={0}
-    Arg_0.1 = f32[3072] parameter(0), sharding={replicated}
-    convert.6 = f16[3072] convert(Arg_0.1)
-    reshape.8 = f16[1,1,3072] reshape(convert.6)
-    broadcast.9 = f16[1,1,3072] broadcast(reshape.8), dimensions={0,1,2}
-    reshape.10 = f16[3072] reshape(broadcast.9)
-    broadcast.11 = f16[16,128,3072] broadcast(reshape.10), dimensions={2}
-    ROOT add.12 = f16[16,128,3072] add(dot.7, broadcast.11)
+    Arg_2.3 = f32[16,64,64] parameter(2), sharding={replicated}
+    convert.4 = f16[16,64,64] convert(Arg_2.3)
+    Arg_1.2 = f32[64,256] parameter(1), sharding={replicated}
+    convert.5 = f16[64,256] convert(Arg_1.2)
+    dot.7 = f16[16,64,256] dot(convert.4, convert.5), lhs_contracting_dims={2}, rhs_contracting_dims={0}
+    Arg_0.1 = f32[256] parameter(0), sharding={replicated}
+    convert.6 = f16[256] convert(Arg_0.1)
+    reshape.8 = f16[1,1,256] reshape(convert.6)
+    broadcast.9 = f16[1,1,256] broadcast(reshape.8), dimensions={0,1,2}
+    reshape.10 = f16[256] reshape(broadcast.9)
+    broadcast.11 = f16[16,64,256] broadcast(reshape.10), dimensions={2}
+    ROOT add.12 = f16[16,64,256] add(dot.7, broadcast.11)
   })";
 
   EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-2, 1e-2}));
@@ -1035,16 +1038,16 @@ TEST_F(MatmulTest, SimpleBiasTestF16_PARAM_F16) {
   HloModule jit_apply
 
   ENTRY matmul.test.f16 {
-    Arg_2.3 = f32[16,128,768] parameter(2), sharding={replicated}
-    convert.4 = f16[16,128,768] convert(Arg_2.3)
-    Arg_1.2 = f16[768,3072] parameter(1), sharding={replicated}
-    dot.5 = f16[16,128,3072] dot(convert.4, Arg_1.2), lhs_contracting_dims={2}, rhs_contracting_dims={0}
-    Arg_0.1 = f16[3072] parameter(0), sharding={replicated}
-    reshape.6 = f16[1,1,3072] reshape(Arg_0.1)
-    broadcast.7 = f16[1,1,3072] broadcast(reshape.6), dimensions={0,1,2}
-    reshape.8 = f16[3072] reshape(broadcast.7)
-    broadcast.9 = f16[16,128,3072] broadcast(reshape.8), dimensions={2}
-    ROOT add.10 = f16[16,128,3072] add(dot.5, broadcast.9)
+    Arg_2.3 = f32[16,64,64] parameter(2), sharding={replicated}
+    convert.4 = f16[16,64,64] convert(Arg_2.3)
+    Arg_1.2 = f16[64,256] parameter(1), sharding={replicated}
+    dot.5 = f16[16,64,256] dot(convert.4, Arg_1.2), lhs_contracting_dims={2}, rhs_contracting_dims={0}
+    Arg_0.1 = f16[256] parameter(0), sharding={replicated}
+    reshape.6 = f16[1,1,256] reshape(Arg_0.1)
+    broadcast.7 = f16[1,1,256] broadcast(reshape.6), dimensions={0,1,2}
+    reshape.8 = f16[256] reshape(broadcast.7)
+    broadcast.9 = f16[16,64,256] broadcast(reshape.8), dimensions={2}
+    ROOT add.10 = f16[16,64,256] add(dot.5, broadcast.9)
   })";
 
   EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-2, 1e-2}));
@@ -1620,7 +1623,7 @@ TEST_F(MatmulTest, SimpleTestNoTransposeFusion1) {
   EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-4, 1e-4}));
   MatchOptimizedHlo(matmul_module_str,
                     R"(
-    ; CHECK:     transpose(%{{[a-z,A-Z,0-9,_,\.]*}}),
+    ; CHECK:     fusion(%{{[a-z,A-Z,0-9,_,\.]*}}),
     ; CHECK:     custom_call_target="__onednn$matmul",
     )");
 }
@@ -1770,16 +1773,16 @@ TEST_F(MatmulTest, SimpleTestF32WithBiasAndAddFusionWithReshape) {
   const char* matmul_module_str = R"(
   HloModule matmul.test.f32
   ENTRY matmul.test.f32 {
-    arg.0 = f32[6304,768] parameter(0), parameter_replication={false}
-    arg.1 = f32[768,3072] parameter(1), parameter_replication={false}
-    dot.378 = f32[6304,3072] dot(arg.0, arg.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
-    reshape.11 = f32[32,197,3072] reshape(dot.378)
-    constant.381 = f32[3072] constant(0.3)
-    broadcast.382 = f32[32,197,3072] broadcast(constant.381), dimensions={2}
-    add.0 = f32[32,197,3072] add(reshape.11, broadcast.382)
-    const.1 = f32[32,197,3072] constant(0.65)
-    add.1 = f32[32,197,3072] add(add.0, const.1)
-    ROOT out = f32[6304,3072] reshape(add.1)
+    arg.0 = f32[1024,64] parameter(0), parameter_replication={false}
+    arg.1 = f32[64,256] parameter(1), parameter_replication={false}
+    dot.378 = f32[1024,256] dot(arg.0, arg.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    reshape.11 = f32[32,32,256] reshape(dot.378)
+    constant.381 = f32[256] constant(0.3)
+    broadcast.382 = f32[32,32,256] broadcast(constant.381), dimensions={2}
+    add.0 = f32[32,32,256] add(reshape.11, broadcast.382)
+    const.1 = f32[32,32,256] constant(0.65)
+    add.1 = f32[32,32,256] add(add.0, const.1)
+    ROOT out = f32[1024,256] reshape(add.1)
   })";
   EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-4, 1e-4}));
   MatchOptimizedHlo(matmul_module_str, fused_matmul_bias_add_str_);
@@ -1935,50 +1938,50 @@ TEST_F(MatmulTest, BiasAndLegalizedErfcGELUTestF32) {
   const char* matmul_module_str = R"(
   HloModule matmul.test.f32
   ENTRY matmul.test.f32 {
-    arg.0 = f32[6304,768] parameter(0), parameter_replication={false}
-    arg.1 = f32[768,3072] parameter(1), parameter_replication={false}
-    dot.0 = f32[6304,3072] dot(arg.0, arg.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
-    reshape.0 = f32[32,197,3072] reshape(dot.0)
-    constant.bias = f32[3072] constant(0.3)
-    broadcast.bias = f32[32,197,3072] broadcast(constant.bias), dimensions={2}
-    add.bias = f32[32,197,3072] add(reshape.0, broadcast.bias)
+    arg.0 = f32[1024,64] parameter(0), parameter_replication={false}
+    arg.1 = f32[64,256] parameter(1), parameter_replication={false}
+    dot.0 = f32[1024,256] dot(arg.0, arg.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    reshape.0 = f32[32,32,256] reshape(dot.0)
+    constant.bias = f32[256] constant(0.3)
+    broadcast.bias = f32[32,32,256] broadcast(constant.bias), dimensions={2}
+    add.bias = f32[32,32,256] add(reshape.0, broadcast.bias)
 
     // Compute y = -x / sqrt(2) = negate(x) * (1/sqrt(2))
-    negate.0 = f32[32,197,3072] negate(add.bias)
+    negate.0 = f32[32,32,256] negate(add.bias)
     constant.inv_sqrt2 = f32[] constant(0.707106769)
-    broadcast.inv_sqrt2 = f32[32,197,3072] broadcast(constant.inv_sqrt2), dimensions={}
-    y.0 = f32[32,197,3072] multiply(negate.0, broadcast.inv_sqrt2)
+    broadcast.inv_sqrt2 = f32[32,32,256] broadcast(constant.inv_sqrt2), dimensions={}
+    y.0 = f32[32,32,256] multiply(negate.0, broadcast.inv_sqrt2)
 
     // Compute predicate: |y| < 1
-    abs.y = f32[32,197,3072] abs(y.0)
+    abs.y = f32[32,32,256] abs(y.0)
     constant.one_pred = f32[] constant(1)
-    broadcast.one_pred = f32[32,197,3072] broadcast(constant.one_pred), dimensions={}
-    pred.0 = pred[32,197,3072] compare(abs.y, broadcast.one_pred), direction=LT
+    broadcast.one_pred = f32[32,32,256] broadcast(constant.one_pred), dimensions={}
+    pred.0 = pred[32,32,256] compare(abs.y, broadcast.one_pred), direction=LT
 
     // True branch: 1 - erf_approx(y)
     // Use erf(y) as a stand-in for the polynomial erf approximation
-    erf.y = f32[32,197,3072] erf(y.0)
+    erf.y = f32[32,32,256] erf(y.0)
     constant.one_true = f32[] constant(1)
-    broadcast.one_true = f32[32,197,3072] broadcast(constant.one_true), dimensions={}
-    true_branch.0 = f32[32,197,3072] subtract(broadcast.one_true, erf.y)
+    broadcast.one_true = f32[32,32,256] broadcast(constant.one_true), dimensions={}
+    true_branch.0 = f32[32,32,256] subtract(broadcast.one_true, erf.y)
 
     // False branch: erfc_approx(y)
     // Use (1 - erf(y)) as a stand-in for the erfc polynomial approximation
     constant.one_false = f32[] constant(1)
-    broadcast.one_false = f32[32,197,3072] broadcast(constant.one_false), dimensions={}
-    erf.y2 = f32[32,197,3072] erf(y.0)
-    false_branch.0 = f32[32,197,3072] subtract(broadcast.one_false, erf.y2)
+    broadcast.one_false = f32[32,32,256] broadcast(constant.one_false), dimensions={}
+    erf.y2 = f32[32,32,256] erf(y.0)
+    false_branch.0 = f32[32,32,256] subtract(broadcast.one_false, erf.y2)
 
     // select(|y| < 1, 1 - erf_approx(y), erfc_approx(y))
-    select.0 = f32[32,197,3072] select(pred.0, true_branch.0, false_branch.0)
+    select.0 = f32[32,32,256] select(pred.0, true_branch.0, false_branch.0)
 
     // 0.5 * x
     constant.half = f32[] constant(0.5)
-    broadcast.half = f32[32,197,3072] broadcast(constant.half), dimensions={}
-    half_x.0 = f32[32,197,3072] multiply(broadcast.half, add.bias)
+    broadcast.half = f32[32,32,256] broadcast(constant.half), dimensions={}
+    half_x.0 = f32[32,32,256] multiply(broadcast.half, add.bias)
 
     // 0.5 * x * select(...)
-    ROOT out = f32[32,197,3072] multiply(half_x.0, select.0)
+    ROOT out = f32[32,32,256] multiply(half_x.0, select.0)
   })";
 
   EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{1e-4, 1e-4}));
@@ -1999,54 +2002,54 @@ TEST_F(MatmulTest, BiasAndLegalizedErfcGELUTestF16) {
   const char* matmul_module_str = R"(
   HloModule matmul.test.f16
   ENTRY matmul.test.f16 {
-    arg.0 = f16[6304,768] parameter(0), parameter_replication={false}
-    arg.1 = f16[768,3072] parameter(1), parameter_replication={false}
-    dot.0 = f16[6304,3072] dot(arg.0, arg.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
-    reshape.0 = f16[32,197,3072] reshape(dot.0)
-    constant.bias = f16[3072] constant(0.3)
-    broadcast.bias = f16[32,197,3072] broadcast(constant.bias), dimensions={2}
-    add.bias = f16[32,197,3072] add(reshape.0, broadcast.bias)
+    arg.0 = f16[1024,64] parameter(0), parameter_replication={false}
+    arg.1 = f16[64,256] parameter(1), parameter_replication={false}
+    dot.0 = f16[1024,256] dot(arg.0, arg.1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    reshape.0 = f16[32,32,256] reshape(dot.0)
+    constant.bias = f16[256] constant(0.3)
+    broadcast.bias = f16[32,32,256] broadcast(constant.bias), dimensions={2}
+    add.bias = f16[32,32,256] add(reshape.0, broadcast.bias)
 
     // Compute y = -x / sqrt(2) in f16
-    negate.0 = f16[32,197,3072] negate(add.bias)
+    negate.0 = f16[32,32,256] negate(add.bias)
     constant.inv_sqrt2 = f16[] constant(0.707106769)
-    broadcast.inv_sqrt2 = f16[32,197,3072] broadcast(constant.inv_sqrt2), dimensions={}
-    y.0 = f16[32,197,3072] multiply(negate.0, broadcast.inv_sqrt2)
+    broadcast.inv_sqrt2 = f16[32,32,256] broadcast(constant.inv_sqrt2), dimensions={}
+    y.0 = f16[32,32,256] multiply(negate.0, broadcast.inv_sqrt2)
 
     // Convert f16 -> f32 before erfc polynomial
-    y.f32 = f32[32,197,3072] convert(y.0)
+    y.f32 = f32[32,32,256] convert(y.0)
 
     // Compute predicate: |y| < 1 (in f32)
-    abs.y = f32[32,197,3072] abs(y.f32)
+    abs.y = f32[32,32,256] abs(y.f32)
     constant.one_pred = f32[] constant(1)
-    broadcast.one_pred = f32[32,197,3072] broadcast(constant.one_pred), dimensions={}
-    pred.0 = pred[32,197,3072] compare(abs.y, broadcast.one_pred), direction=LT
+    broadcast.one_pred = f32[32,32,256] broadcast(constant.one_pred), dimensions={}
+    pred.0 = pred[32,32,256] compare(abs.y, broadcast.one_pred), direction=LT
 
     // True branch: 1 - erf_approx(y) (in f32)
-    erf.y = f32[32,197,3072] erf(y.f32)
+    erf.y = f32[32,32,256] erf(y.f32)
     constant.one_true = f32[] constant(1)
-    broadcast.one_true = f32[32,197,3072] broadcast(constant.one_true), dimensions={}
-    true_branch.0 = f32[32,197,3072] subtract(broadcast.one_true, erf.y)
+    broadcast.one_true = f32[32,32,256] broadcast(constant.one_true), dimensions={}
+    true_branch.0 = f32[32,32,256] subtract(broadcast.one_true, erf.y)
 
     // False branch: erfc_approx(y) (in f32, simplified stand-in)
     constant.one_false = f32[] constant(1)
-    broadcast.one_false = f32[32,197,3072] broadcast(constant.one_false), dimensions={}
-    erf.y2 = f32[32,197,3072] erf(y.f32)
-    false_branch.0 = f32[32,197,3072] subtract(broadcast.one_false, erf.y2)
+    broadcast.one_false = f32[32,32,256] broadcast(constant.one_false), dimensions={}
+    erf.y2 = f32[32,32,256] erf(y.f32)
+    false_branch.0 = f32[32,32,256] subtract(broadcast.one_false, erf.y2)
 
     // select(|y| < 1, 1 - erf_approx(y), erfc_approx(y)) in f32
-    select.0 = f32[32,197,3072] select(pred.0, true_branch.0, false_branch.0)
+    select.0 = f32[32,32,256] select(pred.0, true_branch.0, false_branch.0)
 
     // Convert f32 -> f16 after erfc polynomial
-    select.f16 = f16[32,197,3072] convert(select.0)
+    select.f16 = f16[32,32,256] convert(select.0)
 
     // 0.5 * x (in f16)
     constant.half = f16[] constant(0.5)
-    broadcast.half = f16[32,197,3072] broadcast(constant.half), dimensions={}
-    half_x.0 = f16[32,197,3072] multiply(broadcast.half, add.bias)
+    broadcast.half = f16[32,32,256] broadcast(constant.half), dimensions={}
+    half_x.0 = f16[32,32,256] multiply(broadcast.half, add.bias)
 
     // 0.5 * x * cdf(x)
-    ROOT out = f16[32,197,3072] multiply(half_x.0, select.f16)
+    ROOT out = f16[32,32,256] multiply(half_x.0, select.f16)
   })";
 
   EXPECT_TRUE(RunAndCompare(matmul_module_str, ErrorSpec{0.1, 0.1}));
diff --git a/third_party/xla/xla/service/cpu/tests/restricted/onednn_softmax_test.cc b/third_party/xla/xla/service/cpu/tests/onednn_softmax_test.cc
similarity index 97%
rename from third_party/xla/xla/service/cpu/tests/restricted/onednn_softmax_test.cc
rename to third_party/xla/xla/service/cpu/tests/onednn_softmax_test.cc
index 71c1ab1a073..07194ec4e22 100644
--- a/third_party/xla/xla/service/cpu/tests/restricted/onednn_softmax_test.cc
+++ b/third_party/xla/xla/service/cpu/tests/onednn_softmax_test.cc
@@ -29,7 +29,8 @@ limitations under the License.
 #include "xla/service/cpu/onednn_ops_rewriter.h"
 #include "xla/service/cpu/onednn_util.h"
 #include "xla/service/pattern_matcher.h"
-#include "xla/tests/restricted/hlo_test_base_legacy.h"
+#include "xla/tests/hlo_interpreter_reference_mixin.h"
+#include "xla/tests/hlo_test_base.h"
 
 namespace xla {
 namespace cpu {
@@ -44,12 +45,14 @@ std::string TestParamsToString(
 }
 
 class OneDnnSoftmaxTest
-    : public HloTestBaseLegacy,
+    : public HloInterpreterReferenceMixin<HloTestBase>,
       public ::testing::WithParamInterface<std::tuple<PrimitiveType, int>> {
  protected:
   DebugOptions GetDebugOptionsForTest() const override {
-    DebugOptions debug_options = HloTestBaseLegacy::GetDebugOptionsForTest();
+    DebugOptions debug_options =
+        HloInterpreterReferenceMixin::GetDebugOptionsForTest();
     debug_options.set_xla_cpu_experimental_onednn_custom_call(true);
+    debug_options.clear_xla_cpu_experimental_ynn_fusion_type();
     return debug_options;
   }
 
diff --git a/third_party/xla/xla/service/cpu/tests/restricted/BUILD b/third_party/xla/xla/service/cpu/tests/restricted/BUILD
index af87f03785d..76d5b10a9f9 100644
--- a/third_party/xla/xla/service/cpu/tests/restricted/BUILD
+++ b/third_party/xla/xla/service/cpu/tests/restricted/BUILD
@@ -15,9 +15,6 @@
 
 load("//xla:restricted_package.bzl", "xla_restricted_allow", "xla_restricted_verify")
 load("//xla:xla.default.bzl", "xla_cc_test")
-load("//xla/tsl:tsl.bzl", "tsl_copts")
-load("//xla/tsl/mkl:build_defs.bzl", "if_graph_api")
-load("//xla/tsl/mkl:graph.bzl", "onednn_cc_test", "onednn_graph_cc_test")
 
 package(
     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
@@ -56,128 +53,6 @@ xla_cc_test(
     ],
 )
 
-onednn_graph_cc_test(
-    name = "onednn_fusion_test",
-    srcs = ["onednn_fusion_test.cc"],
-    tags = [
-        "no_oss",
-        "notap",
-    ],
-    deps = [
-        "//xla:error_spec",
-        "//xla/backends/cpu:onednn_support",
-        "//xla/service:cpu_plugin",
-        "//xla/service/cpu:onednn_util",
-        "//xla/tests/restricted:hlo_test_base_legacy",
-        "//xla/tsl/platform:test",
-        "@com_google_absl//absl/strings",
-        "@com_google_googletest//:gtest_main",
-        "@tsl//tsl/platform:platform_port",
-    ],
-)
-
-onednn_cc_test(
-    name = "onednn_matmul_test",
-    srcs = ["onednn_matmul_test.cc"],
-    copts = tsl_copts(),
-    shard_count = 4,
-    tags = [
-        "no_oss",
-        "notap",
-        "pjrt_migration_candidate",
-    ],
-    deps = [
-        "//xla:error_spec",
-        "//xla/hlo/testlib:test",
-        "//xla/service:cpu_plugin",
-        "//xla/service/cpu:onednn_util",
-        "//xla/tests:xla_internal_test_main",
-        "//xla/tests/restricted:hlo_test_base_legacy",
-        "@com_google_absl//absl/strings",
-    ],
-)
-
-onednn_cc_test(
-    name = "onednn_convolution_test",
-    srcs = ["onednn_convolution_test.cc"],
-    copts = tsl_copts(),
-    tags = [
-        # TODO: reenable once onednn is supported by the thunk runtime.
-        "no_oss",
-        "notap",
-        "pjrt_migration_candidate",
-    ],
-    deps = [
-        "//xla:literal",
-        "//xla:shape_util",
-        "//xla/hlo/testlib:filecheck",
-        "//xla/hlo/testlib:test",
-        "//xla/hlo/testlib:test_helpers",
-        "//xla/hlo/utils:hlo_matchers",
-        "//xla/service:cpu_plugin",
-        "//xla/service/cpu:onednn_contraction_rewriter",
-        "//xla/service/cpu:onednn_util",
-        "//xla/tests:xla_internal_test_main",
-        "//xla/tests/restricted:hlo_test_base_legacy",
-        "@com_google_absl//absl/strings",
-        "@tsl//tsl/platform:platform_port",
-    ],
-)
-
-onednn_cc_test(
-    name = "onednn_layer_norm_test",
-    srcs = ["onednn_layer_norm_test.cc"],
-    copts = tsl_copts(),
-    tags = [
-        # TODO: reenable once onednn is supported by the thunk runtime.
-        "no_oss",
-        "notap",
-        "pjrt_migration_candidate",
-    ],
-    deps = [
-        "//xla:error_spec",
-        "//xla/hlo/testlib:test",
-        "//xla/service:cpu_plugin",
-        "//xla/service/cpu:onednn_util",
-        "//xla/tests:xla_internal_test_main",
-        "//xla/tests/restricted:hlo_test_base_legacy",
-    ],
-)
-
-onednn_cc_test(
-    name = "onednn_softmax_test",
-    srcs = ["onednn_softmax_test.cc"],
-    copts = tsl_copts(),
-    shard_count = if_graph_api(4, 1),
-    tags = [
-        # TODO: reenable once onednn is supported by the thunk runtime.
-        "no_oss",
-        "notap",
-        "pjrt_migration_candidate",
-    ],
-    deps = [
-        "//xla:shape_util",
-        "//xla/hlo/ir:hlo",
-        "//xla/hlo/testlib:pattern_matcher_gmock",
-        "//xla/hlo/testlib:test",
-        "//xla/service:cpu_plugin",
-        "//xla/service:pattern_matcher",
-        "//xla/service/cpu:backend_config_proto_cc",
-        "//xla/service/cpu:onednn_config_proto_cc",
-        "//xla/service/cpu:onednn_ops_rewriter",
-        "//xla/service/cpu:onednn_util",
-        "//xla/tests:xla_internal_test_main",
-        "//xla/tests/restricted:hlo_test_base_legacy",
-        "//xla/tsl/platform:statusor",
-        "@com_google_absl//absl/strings",
-    ],
-)
-
 xla_restricted_verify(allowed_targets = [
     xla_restricted_allow("cpu_aot_export_test"),
-    xla_restricted_allow("onednn_convolution_test"),
-    xla_restricted_allow("onednn_fusion_test"),
-    xla_restricted_allow("onednn_layer_norm_test"),
-    xla_restricted_allow("onednn_matmul_test"),
-    xla_restricted_allow("onednn_softmax_test"),
 ])

From ea489bca61ce6f7ebd83aa95d794eb2d8939b806 Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Thu, 8 Oct 2026 08:00:38 -0700
Subject: [PATCH 0290/1000] Test NMS V4 with mixed threshold types and check
 score_threshold's rank

Per review, run the mixed threshold type test for NonMaxSuppressionV4
with pad_to_max_output_size too, which shares the XLA kernel, through
a helper that checks the selected indices up to num_valid, since XLA
pads with other indices where the TensorFlow kernels pad with zeros.
Use device_scope() instead of the deprecated test_scope(), and check
that score_threshold is a scalar, as iou_threshold already is.
---
 tensorflow/compiler/tests/image_ops_test.py   | 45 ++++++++++++++-----
 .../compiler/tf2xla/kernels/image_ops.cc      |  4 ++
 2 files changed, 38 insertions(+), 11 deletions(-)

diff --git a/tensorflow/compiler/tests/image_ops_test.py b/tensorflow/compiler/tests/image_ops_test.py
index 0bf20c9490d..a41796b7165 100644
--- a/tensorflow/compiler/tests/image_ops_test.py
+++ b/tensorflow/compiler/tests/image_ops_test.py
@@ -1022,7 +1022,7 @@ class NonMaxSuppressionTest(xla_test.XLATestCase):
       self.assertEqual(indices_tf.size, 3)
       self.assertAllClose(indices_tf[:3], [3, 0, 5])
 
-  def testNMSV3MixedThresholdDtype(self):
+  def _testMixedThresholdDtype(self, nms):
     # Regression test for GitHub issue 128614: boxes and scores of one
     # floating-point type with thresholds of another, which
     # tf.image.non_max_suppression produces for float16 boxes, failed to
@@ -1041,15 +1041,11 @@ class NonMaxSuppressionTest(xla_test.XLATestCase):
         scores = array_ops.placeholder(dtype, shape=[6])
         iou_threshold = array_ops.placeholder(threshold_dtype, shape=[])
         score_threshold = array_ops.placeholder(threshold_dtype, shape=[])
-        with self.test_scope():
-          selected_indices = image_ops.non_max_suppression_v3(
-              boxes=boxes,
-              scores=scores,
-              max_output_size=6,
-              iou_threshold=iou_threshold,
-              score_threshold=score_threshold)
-        indices_tf = sess.run(
-            selected_indices,
+        with self.device_scope():
+          selected_indices, num_valid = nms(boxes, scores, iou_threshold,
+                                            score_threshold)
+        indices_tf, num_valid_tf = sess.run(
+            [selected_indices, num_valid],
             feed_dict={
                 boxes: np.array(boxes_data, dtype=dtype),
                 scores: np.array(scores_data, dtype=dtype),
@@ -1057,7 +1053,34 @@ class NonMaxSuppressionTest(xla_test.XLATestCase):
                 score_threshold: 0.4
             })
         # Box 5 is below the score threshold.
-        self.assertAllEqual([3, 0], indices_tf)
+        self.assertEqual(num_valid_tf, 2)
+        self.assertAllEqual([3, 0], indices_tf[:num_valid_tf])
+
+  def testNMSV3MixedThresholdDtype(self):
+
+    def nms(boxes, scores, iou_threshold, score_threshold):
+      selected_indices = image_ops.non_max_suppression_v3(
+          boxes=boxes,
+          scores=scores,
+          max_output_size=6,
+          iou_threshold=iou_threshold,
+          score_threshold=score_threshold)
+      return selected_indices, array_ops.size(selected_indices)
+
+    self._testMixedThresholdDtype(nms)
+
+  def testNMSV4MixedThresholdDtype(self):
+
+    def nms(boxes, scores, iou_threshold, score_threshold):
+      return gen_image_ops.non_max_suppression_v4(
+          boxes=boxes,
+          scores=scores,
+          max_output_size=6,
+          iou_threshold=iou_threshold,
+          score_threshold=score_threshold,
+          pad_to_max_output_size=True)
+
+    self._testMixedThresholdDtype(nms)
 
   def testNMSV3EmptyInput(self):
     # Regression test for #117245: with no boxes the suppression loop was
diff --git a/tensorflow/compiler/tf2xla/kernels/image_ops.cc b/tensorflow/compiler/tf2xla/kernels/image_ops.cc
index 61ad71e0894..3ef8757e4e5 100644
--- a/tensorflow/compiler/tf2xla/kernels/image_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/image_ops.cc
@@ -484,6 +484,10 @@ class NonMaxSuppressionOp : public XlaOpKernel {
         context,
         TensorShapeUtils::IsScalar(context->InputShape("iou_threshold")),
         absl::InvalidArgumentError("IOU Threshold isn't a scalar"));
+    OP_REQUIRES(
+        context,
+        TensorShapeUtils::IsScalar(context->InputShape("score_threshold")),
+        absl::InvalidArgumentError("Score Threshold isn't a scalar"));
     OP_REQUIRES_OK(context, context->ConstantInputAsIntScalar(2, &output_size));
     OP_REQUIRES(context, output_size >= 0,
                 absl::InvalidArgumentError(

From fa0f14f6da931f0767688abdc708cefb35034765 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Thu, 8 Oct 2026 06:19:23 -0700
Subject: [PATCH 0291/1000] [XLA:GPU]: Migrate ConvolutionThunk to use
 PerDeviceState

PiperOrigin-RevId: 995792072
---
 .../xla/xla/backends/gpu/runtime/BUILD        |  4 +-
 .../command_buffer_conversion_pass_test.cc    |  3 +-
 .../backends/gpu/runtime/convolution_thunk.cc | 43 +++++++++++--------
 .../backends/gpu/runtime/convolution_thunk.h  | 21 ++++-----
 .../gpu/runtime/convolution_thunk_test.cc     |  3 +-
 .../runtime/thunk_proto_deserialization.cc    |  7 +--
 .../xla/xla/service/gpu/thunk_emitter.cc      | 14 +++---
 7 files changed, 51 insertions(+), 44 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index 47bffe3e149..5175d2c4d79 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -744,6 +744,7 @@ cc_library(
     srcs = ["convolution_thunk.cc"],
     hdrs = ["convolution_thunk.h"],
     deps = [
+        ":per_device_state",
         ":shaped_slice",
         ":thunk",
         ":thunk_proto_cc",
@@ -767,15 +768,12 @@ cc_library(
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:statusor",
         "//xla/tsl/protobuf:dnn_proto_cc",
-        "@com_google_absl//absl/base:core_headers",
-        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
     ],
 )
diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass_test.cc b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass_test.cc
index 71847e70c27..061794e2e5b 100644
--- a/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_conversion_pass_test.cc
@@ -292,7 +292,8 @@ std::unique_ptr<ConvolutionThunk> CreateConvolutionThunk(
       .feature_group_count = 1};
   auto thunk =
       ConvolutionThunk::Create(Thunk::ThunkInfo(), desc, operand_slices,
-                               result_slices, result_slices.back().slice);
+                               result_slices, result_slices.back().slice,
+                               /*devices_per_host=*/1);
   TF_CHECK_OK(thunk.status());
   return std::move(thunk).value();
 }
diff --git a/third_party/xla/xla/backends/gpu/runtime/convolution_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/convolution_thunk.cc
index f69992aca73..4ea1aa540e0 100644
--- a/third_party/xla/xla/backends/gpu/runtime/convolution_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/convolution_thunk.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "xla/backends/gpu/runtime/convolution_thunk.h"
 
 #include <memory>
+#include <optional>
 #include <utility>
 #include <vector>
 
@@ -24,8 +25,9 @@ limitations under the License.
 #include "absl/memory/memory.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
-#include "absl/synchronization/mutex.h"
+#include "absl/status/statusor.h"
 #include "absl/types/span.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/backends/gpu/runtime/shaped_slice.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk.pb.h"
@@ -39,6 +41,7 @@ limitations under the License.
 #include "xla/stream_executor/dnn.h"
 #include "xla/stream_executor/lazy_op_runner.h"
 #include "xla/stream_executor/scratch_allocator.h"
+#include "xla/stream_executor/stream.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/statusor.h"
@@ -54,14 +57,15 @@ absl::StatusOr<std::unique_ptr<ConvolutionThunk>> ConvolutionThunk::Create(
     ThunkInfo thunk_info, GpuConvDescriptor descriptor,
     std::vector<ShapedSlice> operand_slices,
     std::vector<ShapedSlice> result_slices,
-    BufferAllocation::Slice scratch_slice) {
+    BufferAllocation::Slice scratch_slice, int devices_per_host) {
   ABSL_ASSIGN_OR_RETURN(GpuConvConfig config,
                    GetGpuConvConfig(descriptor, /*inst_as_string=*/""));
 
   // Can't use std::make_unique because the constructor is private.
-  return absl::WrapUnique(new ConvolutionThunk(
-      thunk_info, std::move(descriptor), std::move(config),
-      std::move(operand_slices), std::move(result_slices), scratch_slice));
+  return absl::WrapUnique(
+      new ConvolutionThunk(thunk_info, std::move(descriptor), std::move(config),
+                           std::move(operand_slices), std::move(result_slices),
+                           scratch_slice, devices_per_host));
 }
 
 ConvolutionThunk::ConvolutionThunk(ThunkInfo thunk_info,
@@ -69,23 +73,25 @@ ConvolutionThunk::ConvolutionThunk(ThunkInfo thunk_info,
                                    GpuConvConfig config,
                                    std::vector<ShapedSlice> operand_slices,
                                    std::vector<ShapedSlice> result_slices,
-                                   BufferAllocation::Slice scratch_slice)
+                                   BufferAllocation::Slice scratch_slice,
+                                   int devices_per_host)
     : TracedCommand(Kind::kConvolution, std::move(thunk_info)),
+      runner_states_(devices_per_host),
       operand_buffers_(std::move(operand_slices)),
       result_buffers_(std::move(result_slices)),
       scratch_buffer_(scratch_slice),
       descriptor_(std::move(descriptor)),
       config_(std::move(config)) {}
 
-RunConvOptions ConvolutionThunk::GetOrCreate(const GpuConvConfig& config,
-                                             const se::Stream* stream) {
-  absl::MutexLock lock(mu_);
-  auto [it, inserted] =
-      cache_.emplace(stream->parent(), std::unique_ptr<GenericConvRunner>{});
-  if (inserted) {
-    it->second = std::make_unique<GenericConvRunner>(config);
-  }
-  return RunConvOptions{nullptr, it->second.get()};
+absl::StatusOr<RunConvOptions> ConvolutionThunk::GetOrCreate(
+    const GpuConvConfig& config, const se::Stream* stream) {
+  int device_ordinal = stream->parent()->device_ordinal();
+  ABSL_RETURN_IF_ERROR(runner_states_.GetOrCreateAndInitialize(
+      device_ordinal, [&](std::optional<GenericConvRunner>* state) {
+        state->emplace(config);
+        return absl::OkStatus();
+      }));
+  return RunConvOptions{nullptr, &**runner_states_.Find(device_ordinal)};
 }
 
 absl::Status ConvolutionThunk::ExecuteOnStream(const ExecuteParams& params) {
@@ -115,7 +121,7 @@ absl::Status ConvolutionThunk::ExecuteOnStream(const ExecuteParams& params) {
   VLOG(5) << "scratch buffer: " << scratch_buffer_
           << " addr: " << scratch.opaque();
 
-  auto opts = GetOrCreate(config_, params.stream);
+  ABSL_ASSIGN_OR_RETURN(RunConvOptions opts, GetOrCreate(config_, params.stream));
   ABSL_RETURN_IF_ERROR(RunGpuConv(config_, absl::MakeSpan(operand_se_buffers),
                              absl::MakeSpan(result_se_buffers), scratch,
                              params.stream, opts));
@@ -145,7 +151,8 @@ Thunk::BufferUses ConvolutionThunk::buffer_uses() const {
 
 absl::StatusOr<std::unique_ptr<ConvolutionThunk>> ConvolutionThunk::FromProto(
     ThunkInfo thunk_info, const ConvolutionThunkProto& proto,
-    absl::Span<const BufferAllocation> buffer_allocations) {
+    absl::Span<const BufferAllocation> buffer_allocations,
+    int devices_per_host) {
   ABSL_ASSIGN_OR_RETURN(GpuConvDescriptor descriptor,
                    GpuConvDescriptor::FromProto(proto.conv_descriptor()));
 
@@ -169,7 +176,7 @@ absl::StatusOr<std::unique_ptr<ConvolutionThunk>> ConvolutionThunk::FromProto(
 
   return Create(std::move(thunk_info), std::move(descriptor),
                 std::move(operand_slices), std::move(result_slices),
-                scratch_slice);
+                scratch_slice, devices_per_host);
 }
 
 absl::StatusOr<ThunkProto> ConvolutionThunk::ToProto() const {
diff --git a/third_party/xla/xla/backends/gpu/runtime/convolution_thunk.h b/third_party/xla/xla/backends/gpu/runtime/convolution_thunk.h
index 5b77491a5a7..0883992022c 100644
--- a/third_party/xla/xla/backends/gpu/runtime/convolution_thunk.h
+++ b/third_party/xla/xla/backends/gpu/runtime/convolution_thunk.h
@@ -17,14 +17,13 @@ limitations under the License.
 #define XLA_BACKENDS_GPU_RUNTIME_CONVOLUTION_THUNK_H_
 
 #include <memory>
+#include <optional>
 #include <vector>
 
-#include "absl/base/thread_annotations.h"
-#include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk.pb.h"
 #include "xla/backends/gpu/runtime/traced_command.h"
@@ -51,7 +50,7 @@ class ConvolutionThunk : public TracedCommand {
       ThunkInfo thunk_info, GpuConvDescriptor descriptor,
       std::vector<ShapedSlice> operand_slices,
       std::vector<ShapedSlice> result_slices,
-      BufferAllocation::Slice scratch_slice);
+      BufferAllocation::Slice scratch_slice, int devices_per_host);
 
   ConvolutionThunk(const ConvolutionThunk&) = delete;
   ConvolutionThunk& operator=(const ConvolutionThunk&) = delete;
@@ -62,7 +61,8 @@ class ConvolutionThunk : public TracedCommand {
 
   static absl::StatusOr<std::unique_ptr<ConvolutionThunk>> FromProto(
       ThunkInfo thunk_info, const ConvolutionThunkProto& proto,
-      absl::Span<const BufferAllocation> buffer_allocations);
+      absl::Span<const BufferAllocation> buffer_allocations,
+      int devices_per_host);
 
   absl::StatusOr<ThunkProto> ToProto() const override;
 
@@ -71,15 +71,12 @@ class ConvolutionThunk : public TracedCommand {
                    GpuConvConfig config,
                    std::vector<ShapedSlice> operand_slices,
                    std::vector<ShapedSlice> result_slices,
-                   BufferAllocation::Slice scratch_slice);
+                   BufferAllocation::Slice scratch_slice, int devices_per_host);
 
-  RunConvOptions GetOrCreate(const GpuConvConfig& config,
-                             const se::Stream* stream);
+  absl::StatusOr<RunConvOptions> GetOrCreate(const GpuConvConfig& config,
+                                             const se::Stream* stream);
 
-  absl::Mutex mu_;
-  absl::flat_hash_map<const se::StreamExecutor*,
-                      std::unique_ptr<GenericConvRunner>>
-      cache_ ABSL_GUARDED_BY(mu_);
+  PerDeviceState<std::optional<GenericConvRunner>> runner_states_;
 
   std::vector<ShapedSlice> operand_buffers_;
   std::vector<ShapedSlice> result_buffers_;
diff --git a/third_party/xla/xla/backends/gpu/runtime/convolution_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/convolution_thunk_test.cc
index a194e89c333..ee5ca8bd90e 100644
--- a/third_party/xla/xla/backends/gpu/runtime/convolution_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/convolution_thunk_test.cc
@@ -155,7 +155,8 @@ TEST(ConvolutionThunkTest, ProtoRoundTrip) {
   ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<ConvolutionThunk> thunk,
       ConvolutionThunk::FromProto(thunk_info, proto.convolution_thunk(),
-                                  buffer_allocations));
+                                  buffer_allocations,
+                                  /*devices_per_host=*/1));
   ASSERT_OK_AND_ASSIGN(ThunkProto round_trip_proto, thunk->ToProto());
   EXPECT_THAT(round_trip_proto, EqualsProto(proto));
 }
diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
index ce5708c8610..c14a5800e01 100644
--- a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
@@ -194,9 +194,10 @@ absl::StatusOr<std::unique_ptr<Thunk>> DeserializeThunkProtoImpl(
       return NormThunk::FromProto(std::move(thunk_info),
                                   thunk_proto.norm_thunk(), buffer_allocations);
     case ThunkProto::kConvolutionThunk:
-      return ConvolutionThunk::FromProto(std::move(thunk_info),
-                                         thunk_proto.convolution_thunk(),
-                                         buffer_allocations);
+      return ConvolutionThunk::FromProto(
+          std::move(thunk_info), thunk_proto.convolution_thunk(),
+          buffer_allocations,
+          gpu_topology.has_value() ? gpu_topology->num_devices_per_host() : 0);
     case ThunkProto::kConvolutionReorderThunk: {
       return ConvolutionReorderThunk::FromProto(
           std::move(thunk_info), thunk_proto.convolution_reorder_thunk(),
diff --git a/third_party/xla/xla/service/gpu/thunk_emitter.cc b/third_party/xla/xla/service/gpu/thunk_emitter.cc
index ae977577ac7..44c8bbb63e6 100644
--- a/third_party/xla/xla/service/gpu/thunk_emitter.cc
+++ b/third_party/xla/xla/service/gpu/thunk_emitter.cc
@@ -695,12 +695,14 @@ absl::StatusOr<ThunkSequence> ThunkEmitter::EmitConvolution(
                                   instr->window(),
                                   instr->convolution_dimension_numbers(),
                                   instr->feature_group_count()};
-  ABSL_ASSIGN_OR_RETURN(auto thunk,
-                   ConvolutionThunk::Create(
-                       Thunk::ThunkInfo::WithProfileAnnotation(
-                           instr, ir_emitter_context_->GetNextThunkId()),
-                       std::move(descriptor), std::move(operand_slices),
-                       std::move(result_slices), scratch_slice));
+  ABSL_ASSIGN_OR_RETURN(
+      auto thunk,
+      ConvolutionThunk::Create(
+          Thunk::ThunkInfo::WithProfileAnnotation(
+              instr, ir_emitter_context_->GetNextThunkId()),
+          std::move(descriptor), std::move(operand_slices),
+          std::move(result_slices), scratch_slice,
+          ir_emitter_context_->gpu_topology().num_devices_per_host()));
   return ThunkSequence::Of(std::move(thunk));
 }
 

From 5e280a382ba42053fbde2bc5db47bd2c06d0c18f Mon Sep 17 00:00:00 2001
From: Emilio Cota <ecg@google.com>
Date: Thu, 8 Oct 2026 06:46:53 -0700
Subject: [PATCH 0292/1000] [xla:cpu] avoid hardcoding SSA register number in
 DFSan IrCompilerTest

Avoid hardcoding %1 so that the test will keep working even
if other sanitizers are also enabled. (For example, this will
happen when the test runs under --config=msan once msan support
is complete.)

PiperOrigin-RevId: 995803278
---
 third_party/xla/xla/backends/cpu/codegen/ir_compiler_test.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/backends/cpu/codegen/ir_compiler_test.cc b/third_party/xla/xla/backends/cpu/codegen/ir_compiler_test.cc
index eb2a222fe51..f984c3a284d 100644
--- a/third_party/xla/xla/backends/cpu/codegen/ir_compiler_test.cc
+++ b/third_party/xla/xla/backends/cpu/codegen/ir_compiler_test.cc
@@ -62,6 +62,7 @@ namespace xla::cpu {
 
 namespace {
 
+using ::testing::ContainsRegex;
 using ::testing::HasSubstr;
 using ::testing::Not;
 
@@ -601,7 +602,7 @@ TEST(IrCompilerTest, DataFlowSanitizerInstrumentsPolynomialApproximations) {
   auto ir = llvm_ir::DumpToString(ir_module.get());
   EXPECT_THAT(ir, HasSubstr("@test_exp.dfsan"));
   EXPECT_THAT(ir, HasSubstr("fmul contract"));
-  EXPECT_THAT(ir, HasSubstr("store i8 %1, ptr @__dfsan_retval_tls"));
+  EXPECT_THAT(ir, ContainsRegex("store i8 %[0-9]+, ptr @__dfsan_retval_tls"));
 }
 
 }  // namespace

From 854c5088eb1543194a667b63b616da832045c8ec Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Thu, 8 Oct 2026 06:54:38 -0700
Subject: [PATCH 0293/1000] [XLA:GPU]: Migrate NormThunk to use PerDeviceState

PiperOrigin-RevId: 995806260
---
 .../xla/xla/backends/gpu/runtime/BUILD        |  6 +--
 .../xla/backends/gpu/runtime/norm_thunk.cc    | 45 +++++++++++--------
 .../xla/xla/backends/gpu/runtime/norm_thunk.h | 20 ++++-----
 .../backends/gpu/runtime/norm_thunk_test.cc   |  3 +-
 .../runtime/thunk_proto_deserialization.cc    |  5 ++-
 .../xla/xla/service/gpu/thunk_emitter.cc      | 13 +++---
 6 files changed, 50 insertions(+), 42 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index 5175d2c4d79..579955e7b14 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -3291,6 +3291,7 @@ cc_library(
     srcs = ["norm_thunk.cc"],
     hdrs = ["norm_thunk.h"],
     deps = [
+        ":per_device_state",
         ":thunk",
         ":thunk_proto_cc",
         "//xla:util",
@@ -3302,15 +3303,14 @@ cc_library(
         "//xla/stream_executor:device_address",
         "//xla/stream_executor:lazy_op_runner",
         "//xla/stream_executor:stream",
+        "//xla/stream_executor:stream_executor_h",
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:statusor",
-        "@com_google_absl//absl/base:core_headers",
-        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
     ],
 )
diff --git a/third_party/xla/xla/backends/gpu/runtime/norm_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/norm_thunk.cc
index 71a337d1758..179238f23d6 100644
--- a/third_party/xla/xla/backends/gpu/runtime/norm_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/norm_thunk.cc
@@ -19,12 +19,13 @@ limitations under the License.
 #include <optional>
 #include <utility>
 
+#include "absl/base/nullability.h"
 #include "absl/memory/memory.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
-#include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk.pb.h"
 #include "xla/runtime/buffer_use.h"
@@ -34,6 +35,7 @@ limitations under the License.
 #include "xla/stream_executor/device_address.h"
 #include "xla/stream_executor/lazy_op_runner.h"
 #include "xla/stream_executor/stream.h"
+#include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"
@@ -51,14 +53,15 @@ absl::StatusOr<std::unique_ptr<NormThunk>> NormThunk::Create(
     std::optional<BufferAllocation::Slice> dy_slice,
     std::optional<BufferAllocation::Slice> dscale_slice,
     std::optional<BufferAllocation::Slice> dbias_slice,
-    BufferAllocation::Slice scratch_slice) {
+    BufferAllocation::Slice scratch_slice, int devices_per_host) {
   ABSL_ASSIGN_OR_RETURN(GpuNormConfig config, GpuNormConfig::For(descriptor));
 
   // Can't use make_unique because the constructor is private. go/totw/134
   return absl::WrapUnique(new NormThunk(
       thunk_info, std::move(config), std::move(descriptor), x_slice,
       scale_slice, y_or_dx_slice, bias_slice, expectation_slice,
-      norm_factor_slice, dy_slice, dscale_slice, dbias_slice, scratch_slice));
+      norm_factor_slice, dy_slice, dscale_slice, dbias_slice, scratch_slice,
+      devices_per_host));
 }
 
 NormThunk::NormThunk(ThunkInfo thunk_info, GpuNormConfig config,
@@ -72,7 +75,8 @@ NormThunk::NormThunk(ThunkInfo thunk_info, GpuNormConfig config,
                      std::optional<BufferAllocation::Slice> dy_slice,
                      std::optional<BufferAllocation::Slice> dscale_slice,
                      std::optional<BufferAllocation::Slice> dbias_slice,
-                     BufferAllocation::Slice scratch_slice)
+                     BufferAllocation::Slice scratch_slice,
+                     int devices_per_host)
     : Thunk(Kind::kNorm, thunk_info),
       x_buffer_(x_slice),
       scale_buffer_(scale_slice),
@@ -85,17 +89,18 @@ NormThunk::NormThunk(ThunkInfo thunk_info, GpuNormConfig config,
       dbias_buffer_(dbias_slice),
       scratch_buffer_(scratch_slice),
       descriptor_(descriptor),
-      config_(config) {}
+      config_(config),
+      runner_states_(devices_per_host) {}
 
-NormRunner& NormThunk::GetOrCreateRunner(
-    const stream_executor::Stream* stream) {
-  absl::MutexLock lock(mu_);
-  auto it = runner_cache_.find(stream);
-  if (it == runner_cache_.end()) {
-    it = runner_cache_.insert({stream, std::make_unique<NormRunner>(config_)})
-             .first;
-  }
-  return *it->second;
+absl::StatusOr<NormRunner*> NormThunk::GetOrCreateRunner(
+    const stream_executor::Stream* absl_nonnull stream) {
+  int device_ordinal = stream->parent()->device_ordinal();
+  ABSL_RETURN_IF_ERROR(runner_states_.GetOrCreateAndInitialize(
+      device_ordinal, [&](std::optional<NormRunner>* state) {
+        state->emplace(config_);
+        return absl::OkStatus();
+      }));
+  return &**runner_states_.Find(device_ordinal);
 }
 
 absl::Status NormThunk::ExecuteOnStream(const ExecuteParams& params) {
@@ -131,7 +136,7 @@ absl::Status NormThunk::ExecuteOnStream(const ExecuteParams& params) {
       buffer_allocations.GetDeviceAddress(scratch_buffer_);
 
   RunNormOptions opts;
-  opts.norm_runner = &GetOrCreateRunner(params.stream);
+  ABSL_ASSIGN_OR_RETURN(opts.norm_runner, GetOrCreateRunner(params.stream));
 
   ABSL_RETURN_IF_ERROR(RunGpuNorm(
       config_, x_se_buffer, scale_se_buffer, y_or_dx_se_buffer, bias_se_buffer,
@@ -147,8 +152,8 @@ absl::Status NormThunk::ExecuteOnStream(const ExecuteParams& params) {
 absl::Status NormThunk::Initialize(const InitializeParams& params) {
   // Create the runner at initialization time to avoid hangs if we try to build
   // the execution plan while a NCCL collective is running.
-  se::dnn::LazyOpRunner<se::dnn::NormOp>* lazy_runner =
-      GetOrCreateRunner(params.stream).AsNormRunner();
+  ABSL_ASSIGN_OR_RETURN(NormRunner * runner, GetOrCreateRunner(params.stream));
+  se::dnn::LazyOpRunner<se::dnn::NormOp>* lazy_runner = runner->AsNormRunner();
   ABSL_ASSIGN_OR_RETURN(auto ln_config, config_.AsDnnNormOpConfig());
   return lazy_runner->GetOrCreateRunner(ln_config, params.stream).status();
 }
@@ -185,7 +190,8 @@ Thunk::BufferUses NormThunk::buffer_uses() const {
 
 absl::StatusOr<std::unique_ptr<NormThunk>> NormThunk::FromProto(
     ThunkInfo thunk_info, const NormThunkProto& proto,
-    absl::Span<const BufferAllocation> buffer_allocations) {
+    absl::Span<const BufferAllocation> buffer_allocations,
+    int devices_per_host) {
   ABSL_ASSIGN_OR_RETURN(GpuNormDescriptor descriptor,
                    GpuNormDescriptor::FromProto(proto.norm_descriptor()));
 
@@ -229,7 +235,8 @@ absl::StatusOr<std::unique_ptr<NormThunk>> NormThunk::FromProto(
                                      proto.scratch(), buffer_allocations));
 
   return Create(std::move(thunk_info), descriptor, x, scale, y_or_dx, bias,
-                expectation, norm_factor, dy, dscale, dbias, scratch);
+                expectation, norm_factor, dy, dscale, dbias, scratch,
+                devices_per_host);
 }
 
 absl::StatusOr<ThunkProto> NormThunk::ToProto() const {
diff --git a/third_party/xla/xla/backends/gpu/runtime/norm_thunk.h b/third_party/xla/xla/backends/gpu/runtime/norm_thunk.h
index d92f23098bd..67081b73840 100644
--- a/third_party/xla/xla/backends/gpu/runtime/norm_thunk.h
+++ b/third_party/xla/xla/backends/gpu/runtime/norm_thunk.h
@@ -19,12 +19,11 @@ limitations under the License.
 #include <memory>
 #include <optional>
 
-#include "absl/base/thread_annotations.h"
-#include "absl/container/flat_hash_map.h"
+#include "absl/base/nullability.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk.pb.h"
 #include "xla/service/buffer_assignment.h"
@@ -47,7 +46,7 @@ class NormThunk : public Thunk {
       std::optional<BufferAllocation::Slice> dy,
       std::optional<BufferAllocation::Slice> dscale,
       std::optional<BufferAllocation::Slice> dbias,
-      BufferAllocation::Slice scratch);
+      BufferAllocation::Slice scratch, int devices_per_host);
 
   NormThunk(const NormThunk&) = delete;
   NormThunk& operator=(const NormThunk&) = delete;
@@ -59,7 +58,8 @@ class NormThunk : public Thunk {
 
   static absl::StatusOr<std::unique_ptr<NormThunk>> FromProto(
       ThunkInfo thunk_info, const NormThunkProto& proto,
-      absl::Span<const BufferAllocation> buffer_allocations);
+      absl::Span<const BufferAllocation> buffer_allocations,
+      int devices_per_host);
 
   absl::StatusOr<ThunkProto> ToProto() const override;
 
@@ -73,7 +73,7 @@ class NormThunk : public Thunk {
             std::optional<BufferAllocation::Slice> dy,
             std::optional<BufferAllocation::Slice> dscale,
             std::optional<BufferAllocation::Slice> dbias,
-            BufferAllocation::Slice scratch);
+            BufferAllocation::Slice scratch, int devices_per_host);
 
   BufferAllocation::Slice x_buffer_;
   BufferAllocation::Slice scale_buffer_;
@@ -85,17 +85,15 @@ class NormThunk : public Thunk {
   std::optional<BufferAllocation::Slice> dscale_buffer_;
   std::optional<BufferAllocation::Slice> dbias_buffer_;
   BufferAllocation::Slice scratch_buffer_;
-  NormRunner& GetOrCreateRunner(const stream_executor::Stream*);
+  absl::StatusOr<NormRunner*> GetOrCreateRunner(
+      const stream_executor::Stream* absl_nonnull stream);
 
   // Technically this is only needed during initialization to create the
   // GpuNormConfig, but the actual GpuNormConfig is hard to serialize. So we
   // keep the descriptor around for serialization purposes.
   GpuNormDescriptor descriptor_;
   GpuNormConfig config_;
-  absl::Mutex mu_;
-  absl::flat_hash_map<const stream_executor::Stream*,
-                      std::unique_ptr<NormRunner>>
-      runner_cache_ ABSL_GUARDED_BY(mu_);
+  PerDeviceState<std::optional<NormRunner>> runner_states_;
 };
 
 }  // namespace gpu
diff --git a/third_party/xla/xla/backends/gpu/runtime/norm_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/norm_thunk_test.cc
index 1e69bafea11..38068aa0f3b 100644
--- a/third_party/xla/xla/backends/gpu/runtime/norm_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/norm_thunk_test.cc
@@ -103,7 +103,8 @@ TEST(NormThunkTest, ProtoRoundTrip) {
                        Thunk::ThunkInfo::FromProto(proto.thunk_info()));
   ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<NormThunk> thunk,
-      NormThunk::FromProto(thunk_info, proto.norm_thunk(), buffer_allocations));
+      NormThunk::FromProto(thunk_info, proto.norm_thunk(), buffer_allocations,
+                           /*devices_per_host=*/1));
   ASSERT_OK_AND_ASSIGN(ThunkProto round_trip_proto, thunk->ToProto());
   EXPECT_THAT(round_trip_proto, EqualsProto(proto));
 }
diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
index c14a5800e01..7d047a3feac 100644
--- a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
@@ -191,8 +191,9 @@ absl::StatusOr<std::unique_ptr<Thunk>> DeserializeThunkProtoImpl(
           std::move(thunk_info), thunk_proto.cublas_lt_matmul_thunk(),
           buffer_allocations);
     case ThunkProto::kNormThunk:
-      return NormThunk::FromProto(std::move(thunk_info),
-                                  thunk_proto.norm_thunk(), buffer_allocations);
+      return NormThunk::FromProto(
+          std::move(thunk_info), thunk_proto.norm_thunk(), buffer_allocations,
+          gpu_topology.has_value() ? gpu_topology->num_devices_per_host() : 0);
     case ThunkProto::kConvolutionThunk:
       return ConvolutionThunk::FromProto(
           std::move(thunk_info), thunk_proto.convolution_thunk(),
diff --git a/third_party/xla/xla/service/gpu/thunk_emitter.cc b/third_party/xla/xla/service/gpu/thunk_emitter.cc
index 44c8bbb63e6..27779eca3f3 100644
--- a/third_party/xla/xla/service/gpu/thunk_emitter.cc
+++ b/third_party/xla/xla/service/gpu/thunk_emitter.cc
@@ -1107,12 +1107,13 @@ absl::StatusOr<ThunkSequence> ThunkEmitter::EmitNorm(
 
   ABSL_ASSIGN_OR_RETURN(
       std::unique_ptr<NormThunk> thunk,
-      NormThunk::Create(Thunk::ThunkInfo::WithProfileAnnotation(
-                            instr, ir_emitter_context_->GetNextThunkId()),
-                        std::move(descriptor), x_slice, scale_slice,
-                        y_or_dx_slice, bias_slice, expectation_slice,
-                        norm_factor_slice, dy_slice, dscale_slice, dbias_slice,
-                        scratch_slice.slice));
+      NormThunk::Create(
+          Thunk::ThunkInfo::WithProfileAnnotation(
+              instr, ir_emitter_context_->GetNextThunkId()),
+          std::move(descriptor), x_slice, scale_slice, y_or_dx_slice,
+          bias_slice, expectation_slice, norm_factor_slice, dy_slice,
+          dscale_slice, dbias_slice, scratch_slice.slice,
+          ir_emitter_context_->gpu_topology().num_devices_per_host()));
   return ThunkSequence::Of(std::move(thunk));
 }
 

From 9b41c88af99214b9bfd109262abfe01c13c274ae Mon Sep 17 00:00:00 2001
From: gauravarun21 <gauravarun21@gmail.com>
Date: Thu, 8 Oct 2026 17:37:15 +0200
Subject: [PATCH 0294/1000] Fix reduce_std/reduce_variance overflow
 inconsistency for float16

tf.math.reduce_std on large-magnitude float16 input returned nan on CPU
but inf on GPU for the same input, because the squared deviations
computed internally can overflow float16's range (~65504) even when
the final standard deviation does not (e.g. input scale ~1e4 gives
squared deviations ~1e8, but std stays ~1e4). Computing entirely in
float16, the overflow manifests differently depending on reduction
order/parallelization (e.g. summing partial results that overflowed in
opposite directions can produce nan instead of a consistent inf),
which is why CPU and GPU disagreed.

tf.nn.moments already works around the identical problem for the exact
same reason ("The dynamic range of fp16 is too limited to support the
collection of sufficient statistics"), by computing in float32 and
casting back. Apply the same fix to reduce_variance/reduce_std: factor
the shared computation into a private helper that promotes float16
input to float32, and have reduce_std take sqrt before casting the
final result back down, so a standard deviation that is representable
in float16 is returned correctly even when the intermediate variance
is not.

Fixes #118194
---
 tensorflow/python/ops/math_ops.py      | 55 ++++++++++++++++++--------
 tensorflow/python/ops/math_ops_test.py | 21 ++++++++++
 2 files changed, 60 insertions(+), 16 deletions(-)

diff --git a/tensorflow/python/ops/math_ops.py b/tensorflow/python/ops/math_ops.py
index 261a0d1bcf4..b0f4b195f2a 100644
--- a/tensorflow/python/ops/math_ops.py
+++ b/tensorflow/python/ops/math_ops.py
@@ -2741,20 +2741,39 @@ def reduce_variance(input_tensor, axis=None, keepdims=False, name=None):
   name = name if name else "reduce_variance"
   with ops.name_scope(name):
     input_tensor = ops.convert_to_tensor(input_tensor)
-    means = reduce_mean(input_tensor, axis=axis, keepdims=True)
-    if means.dtype.is_integer:
-      raise TypeError(f"Input must be either real or complex. "
-                      f"Received integer type {means.dtype}.")
-    diff = input_tensor - means
-    if diff.dtype.is_complex:
-      # For complex values we need to take the absolute value before squaring.
-      # This is achieved by multiplying with the conjugate.
-      real_dtype = diff.dtype.real_dtype
-      squared_deviations = gen_math_ops.real(
-          gen_math_ops.mul(gen_math_ops.conj(diff), diff), Tout=real_dtype)
-    else:
-      squared_deviations = gen_math_ops.square(diff)
-    return reduce_mean(squared_deviations, axis=axis, keepdims=keepdims)
+    is_float16 = input_tensor.dtype == dtypes.float16
+    variance = _reduce_variance_in_stable_dtype(input_tensor, axis, keepdims)
+    if is_float16:
+      variance = cast(variance, dtypes.float16)
+    return variance
+
+
+def _reduce_variance_in_stable_dtype(input_tensor, axis, keepdims):
+  """Computes variance, promoting float16 to float32 first.
+
+  The dynamic range of float16 is too limited for the squared deviations
+  computed below, so this performs the computation on 32-bit floats before
+  returning, matching the approach `tf.nn.moments` uses for the same reason.
+  Unlike `reduce_variance`, this does not cast the result back to the
+  original dtype, so that `reduce_std` can take the square root while still
+  in the wider dtype.
+  """
+  if input_tensor.dtype == dtypes.float16:
+    input_tensor = cast(input_tensor, dtypes.float32)
+  means = reduce_mean(input_tensor, axis=axis, keepdims=True)
+  if means.dtype.is_integer:
+    raise TypeError(f"Input must be either real or complex. "
+                    f"Received integer type {means.dtype}.")
+  diff = input_tensor - means
+  if diff.dtype.is_complex:
+    # For complex values we need to take the absolute value before squaring.
+    # This is achieved by multiplying with the conjugate.
+    real_dtype = diff.dtype.real_dtype
+    squared_deviations = gen_math_ops.real(
+        gen_math_ops.mul(gen_math_ops.conj(diff), diff), Tout=real_dtype)
+  else:
+    squared_deviations = gen_math_ops.square(diff)
+  return reduce_mean(squared_deviations, axis=axis, keepdims=keepdims)
 
 
 @tf_export("math.reduce_std")
@@ -2804,8 +2823,12 @@ def reduce_std(input_tensor, axis=None, keepdims=False, name=None):
   name = name if name else "reduce_std"
   with ops.name_scope(name):
     input_tensor = ops.convert_to_tensor(input_tensor)
-    variance = reduce_variance(input_tensor, axis=axis, keepdims=keepdims)
-    return gen_math_ops.sqrt(variance)
+    is_float16 = input_tensor.dtype == dtypes.float16
+    variance = _reduce_variance_in_stable_dtype(input_tensor, axis, keepdims)
+    std = gen_math_ops.sqrt(variance)
+    if is_float16:
+      std = cast(std, dtypes.float16)
+    return std
 
 
 @tf_export("math.reduce_prod", "reduce_prod", v1=[])
diff --git a/tensorflow/python/ops/math_ops_test.py b/tensorflow/python/ops/math_ops_test.py
index 734994becd5..411a66ebb98 100644
--- a/tensorflow/python/ops/math_ops_test.py
+++ b/tensorflow/python/ops/math_ops_test.py
@@ -145,6 +145,27 @@ class ReduceTest(test_util.TensorFlowTestCase):
                                      []])
     self.assertAllClose(math_ops.reduce_std(x, axis=0), [0., 4., 1., 0.])
 
+  def testReduceStdFloat16LargeMagnitude(self):
+    # Regression test: squared deviations of large-magnitude float16 values
+    # can overflow float16's range even when the standard deviation itself
+    # does not, so reduce_std/reduce_variance need to compute internally in
+    # a wider dtype. Previously this produced NaN on CPU and Inf on GPU for
+    # reduce_std, instead of a finite, correct result on both.
+    np.random.seed(0)
+    x_np = np.random.randn(1000).astype(np.float32) * 1e4
+    x_f16 = constant_op.constant(x_np.astype(np.float16))
+    expected_std = np.std(x_np.astype(np.float64))
+
+    std = self.evaluate(math_ops.reduce_std(x_f16))
+    self.assertTrue(np.isfinite(std), f"Expected a finite result, got {std}")
+    self.assertAllClose(std, expected_std, rtol=0.02)
+
+    # The true variance here does exceed float16's range (~65504), so,
+    # unlike the standard deviation, it should consistently overflow to Inf
+    # rather than landing on NaN on some platforms.
+    variance = self.evaluate(math_ops.reduce_variance(x_f16))
+    self.assertTrue(np.isposinf(variance), f"Expected Inf, got {variance}")
+
   def testReduceStdComplex(self):
     # Ensure that complex values are handled to be consistent with numpy
     complex_ys = [([0 - 1j, 0 + 1j], dtypes.float64),

From 30bb22e6e17248b7976eeb9f8793fa0dd84cadc1 Mon Sep 17 00:00:00 2001
From: Dmitri Latushko <latushko@google.com>
Date: Thu, 8 Oct 2026 07:03:52 -0700
Subject: [PATCH 0295/1000] [XLA:GPU] Pass -no-cache to libnvJitLink starting
 from CUDA 12.5 instead of 12.6

`-no-cache` is supported in `libnvJitLink` 12.5 (`libnvJitLink.so.12.5.82`) and is required on CUDA 12.5.1 builds to prevent concurrent `~/.nv/ComputeCache` writes from failing with `INTERNAL: ERROR 999: AddToCache`. Lower the `-no-cache` runtime version check in `CompileAndLinkUsingLibNvJitLink` and `GetLatestPtxIsaVersionForLibNvJitLinkImpl` from `{12, 6}` to `{12, 5}`.

PiperOrigin-RevId: 995810596
---
 third_party/xla/xla/stream_executor/cuda/nvjitlink.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/cuda/nvjitlink.cc b/third_party/xla/xla/stream_executor/cuda/nvjitlink.cc
index ac654703f8f..ee3f1efdf92 100644
--- a/third_party/xla/xla/stream_executor/cuda/nvjitlink.cc
+++ b/third_party/xla/xla/stream_executor/cuda/nvjitlink.cc
@@ -176,7 +176,7 @@ absl::StatusOr<cuda::Assembly> CompileAndLinkUsingLibNvJitLink(
   if (version >= NvJitLinkVersion{12, 3}) {
     cli_args.emplace_back(absl::StrCat("-split-compile=", inputs.size()));
   }
-  if (version >= NvJitLinkVersion{12, 6}) {
+  if (version >= NvJitLinkVersion{12, 5}) {
     cli_args.emplace_back("-no-cache");
   }
 
@@ -298,7 +298,7 @@ absl::StatusOr<int> GetLatestPtxIsaVersionForLibNvJitLinkImpl() {
   // The call to `nvJitLinkCreate` below requires an arch to be specified in
   // order to succeed.
   std::vector<const char*> cli_args_ptrs{"-arch=sm_90a"};
-  if (version.ok() && *version >= NvJitLinkVersion{12, 6}) {
+  if (version.ok() && *version >= NvJitLinkVersion{12, 5}) {
     cli_args_ptrs.push_back("-no-cache");
   }
   nvJitLinkHandle link_handle = nullptr;

From 3c1d10a5678979a4bf40b302396b24b225b991b2 Mon Sep 17 00:00:00 2001
From: Christian Sigg <csigg@google.com>
Date: Thu, 8 Oct 2026 07:23:32 -0700
Subject: [PATCH 0296/1000] [XLA:GPU] Roll back enabling new cuDNN conv fusions
 by default.

Reverts changelist 990213588

PiperOrigin-RevId: 995819588
---
 third_party/xla/xla/backends/gpu/tests/BUILD        |  1 -
 .../backends/gpu/tests/swap_conv_operands_test.cc   |  8 --------
 .../gpu/transforms/layout_assignment_a100.hlo       | 13 ++++++-------
 .../gpu/transforms/layout_assignment_h100.hlo       | 13 ++++++-------
 .../gpu/transforms/layout_assignment_v100.hlo       | 13 ++++++-------
 third_party/xla/xla/debug_options_flags.cc          |  2 +-
 third_party/xla/xla/service/BUILD                   |  4 ++--
 7 files changed, 21 insertions(+), 33 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/tests/BUILD b/third_party/xla/xla/backends/gpu/tests/BUILD
index fe082161e9c..a4bcc37e49a 100644
--- a/third_party/xla/xla/backends/gpu/tests/BUILD
+++ b/third_party/xla/xla/backends/gpu/tests/BUILD
@@ -293,7 +293,6 @@ xla_test(
     deps = [
         ":hlo_pjrt_gpu_test_base",
         "//xla:error_spec",
-        "//xla:xla_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:filecheck",
         "//xla/stream_executor:device_description",
diff --git a/third_party/xla/xla/backends/gpu/tests/swap_conv_operands_test.cc b/third_party/xla/xla/backends/gpu/tests/swap_conv_operands_test.cc
index 941aef2f582..7ed04346490 100644
--- a/third_party/xla/xla/backends/gpu/tests/swap_conv_operands_test.cc
+++ b/third_party/xla/xla/backends/gpu/tests/swap_conv_operands_test.cc
@@ -26,7 +26,6 @@ limitations under the License.
 #include "xla/hlo/testlib/filecheck.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/tests/hlo_pjrt_interpreter_reference_mixin.h"
-#include "xla/xla.pb.h"
 
 namespace xla::gpu {
 namespace {
@@ -34,13 +33,6 @@ namespace {
 class SwapConvOperandsTest
     : public HloInterpreterReferenceMixin<HloPjRtGpuTestBase> {
  public:
-  DebugOptions GetDebugOptionsForTest() const override {
-    DebugOptions debug_options = HloInterpreterReferenceMixin<
-        HloPjRtGpuTestBase>::GetDebugOptionsForTest();
-    debug_options.set_xla_gpu_experimental_enable_conv_fusion(false);
-    return debug_options;
-  }
-
   void MatchOptimizedHlo(absl::string_view hlo,
                          absl::string_view expected_hlo) {
     ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> optimized_module,
diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_a100.hlo b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_a100.hlo
index cac1d16a5b4..4383bfd0325 100644
--- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_a100.hlo
+++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_a100.hlo
@@ -1,15 +1,14 @@
 // RUN: hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../backends/gpu/target_config/specs/a100_pcie_80.txtpb --split-input-file | FileCheck %s
-// CHECK: %conv_fprop_fusion_comp
-// CHECK: convolution
-// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, convolution_kind=fprop
+
 // CHECK: %wrapped_transpose_computation
 // CHECK-NEXT: bf16[3,3,16,32]{3,2,1,0} parameter(0)
-// CHECK-NEXT: ROOT {{.*}} = bf16[32,3,3,16]{3,2,1,0} transpose
+// CHECK-NEXT: bf16[32,3,3,16]{3,2,1,0} transpose
 // CHECK-SAME: dimensions={3,0,1,2}
-// CHECK: ENTRY %main
-// CHECK: %[[WRAPPED_TRANSPOSE:.*]] = bf16[32,3,3,16]{3,2,1,0} fusion({{.*}}), {{.*}}calls=%wrapped_transpose_computation
-// CHECK: ROOT {{.*}} = bf16[1,64,64,32]{3,2,1,0} fusion({{.*}}, %[[WRAPPED_TRANSPOSE]]), kind=kCustom, calls=%conv_fprop_fusion_comp
+// CHECK: (bf16[1,64,64,32]{3,2,1,0}, u8[0]{0}) custom-call
+// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward
+
 HloModule ConvCuDNN
+
 ENTRY main {
   Arg_0.1 = bf16[1,64,64,16]{3,2,1,0} parameter(0), sharding={replicated}
   Arg_1.2 = bf16[3,3,16,32]{3,2,1,0} parameter(1), sharding={replicated}
diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_h100.hlo b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_h100.hlo
index 5d4b5daf721..6a13e4476bc 100644
--- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_h100.hlo
+++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_h100.hlo
@@ -1,15 +1,14 @@
 // RUN: hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../backends/gpu/target_config/specs/h100_sxm.txtpb --split-input-file | FileCheck %s
-// CHECK: %conv_fprop_fusion_comp
-// CHECK: convolution
-// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, convolution_kind=fprop
+
 // CHECK: %wrapped_transpose_computation
 // CHECK-NEXT: f8e4m3fn[3,3,16,32]{3,2,1,0} parameter(0)
-// CHECK-NEXT: ROOT {{.*}} = f8e4m3fn[32,3,3,16]{3,2,1,0} transpose
+// CHECK-NEXT: f8e4m3fn[32,3,3,16]{3,2,1,0} transpose
 // CHECK-SAME: dimensions={3,0,1,2}
-// CHECK: ENTRY %main
-// CHECK: %[[WRAPPED_TRANSPOSE:.*]] = f8e4m3fn[32,3,3,16]{3,2,1,0} fusion({{.*}}), {{.*}}calls=%wrapped_transpose_computation
-// CHECK: ROOT {{.*}} = f8e4m3fn[1,64,64,32]{3,2,1,0} fusion({{.*}}, %[[WRAPPED_TRANSPOSE]]), kind=kCustom, calls=%conv_fprop_fusion_comp
+// CHECK: (f8e4m3fn[1,64,64,32]{3,2,1,0}, u8[0]{0}) custom-call
+// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward
+
 HloModule ConvCuDNN
+
 ENTRY main {
   Arg_0.1 = f8e4m3fn[1,64,64,16]{3,2,1,0} parameter(0), sharding={replicated}
   Arg_1.2 = f8e4m3fn[3,3,16,32]{3,2,1,0} parameter(1), sharding={replicated}
diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_v100.hlo b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_v100.hlo
index ec15e5da5f3..c64054ad54b 100644
--- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_v100.hlo
+++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_v100.hlo
@@ -1,15 +1,14 @@
 // RUN: hlo-opt %s --platform=gpu --stage=hlo --xla_gpu_target_config_filename=%S/../../../backends/gpu/target_config/specs/v100.txtpb --split-input-file | FileCheck %s
-// CHECK: %conv_fprop_fusion_comp
-// CHECK: convolution
-// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, convolution_kind=fprop
+
 // CHECK: %wrapped_transpose_computation
 // CHECK-NEXT: f16[3,3,16,32]{3,2,1,0} parameter(0)
-// CHECK-NEXT: ROOT {{.*}} = f16[32,3,3,16]{3,2,1,0} transpose
+// CHECK-NEXT: f16[32,3,3,16]{3,2,1,0} transpose
 // CHECK-SAME: dimensions={3,0,1,2}
-// CHECK: ENTRY %main
-// CHECK: %[[WRAPPED_TRANSPOSE:.*]] = f16[32,3,3,16]{3,2,1,0} fusion({{.*}}), {{.*}}calls=%wrapped_transpose_computation
-// CHECK: ROOT {{.*}} = f16[1,64,64,32]{3,2,1,0} fusion({{.*}}, %[[WRAPPED_TRANSPOSE]]), kind=kCustom, calls=%conv_fprop_fusion_comp
+// CHECK: (f16[1,64,64,32]{3,2,1,0}, u8[0]{0}) custom-call
+// CHECK-SAME: window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_o01i->b01f, custom_call_target="__cudnn$convForward
+
 HloModule ConvCuDNN
+
 ENTRY main {
   Arg_0.1 = f16[1,64,64,16]{3,2,1,0} parameter(0), sharding={replicated}
   Arg_1.2 = f16[3,3,16,32]{3,2,1,0} parameter(1), sharding={replicated}
diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index f5772837e00..fa090afea19 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -529,7 +529,7 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() {
   opts.set_xla_gpu_experimental_matmul_perf_table_path("");
   // TODO(b/366475196): Create XLA GPU without cuDNN, cuBLAS.
   opts.set_xla_gpu_experimental_disable_binary_libraries(false);
-  opts.set_xla_gpu_experimental_enable_conv_fusion(true);
+  opts.set_xla_gpu_experimental_enable_conv_fusion(false);
   opts.set_xla_gpu_dot_merger_threshold_mb(64);
   opts.set_xla_enable_fast_math(false);
   opts.set_xla_gpu_experimental_parallel_collective_overlap_limit(1);
diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index 6b691c8cad2..1a0bc2bb877 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -5953,7 +5953,7 @@ xla_aot_compile_gpu(
         "b200",
     ],
     module = "xla_aot_compile_test_convolution.mlir",
-    xla_flags = "--xla_gpu_use_new_autotune_cache_format=true",
+    xla_flags = "--xla_gpu_use_new_autotune_cache_format=true --xla_gpu_experimental_enable_conv_fusion=true",
 )
 
 xla_aot_compile_gpu(
@@ -5964,7 +5964,7 @@ xla_aot_compile_gpu(
         "b200",
     ],
     module = "xla_aot_compile_test_convolution.mlir",
-    xla_flags = "--xla_gpu_use_new_autotune_cache_format=false",
+    xla_flags = "--xla_gpu_use_new_autotune_cache_format=false --xla_gpu_experimental_enable_conv_fusion=true",
 )
 
 xla_aot_compile_gpu_runtime_autotuning(

From 729e1cf79663efeea445ab27031dabffdef11465 Mon Sep 17 00:00:00 2001
From: Ilya Tikhonovskiy <loislo@google.com>
Date: Thu, 8 Oct 2026 07:59:47 -0700
Subject: [PATCH 0297/1000] [XLA:GPU] Rebalance shard counts for XLA GPU
 tail-latency stragglers and over-sharded tests.

Increase `shard_count` on the four tail stragglers (`convolution_cudnn_test` to 5, `topk_test` to 12, `scatter_determinism_expander_test` to 3, and `topk_specializer_test` to 15) so their single-shard p90 duration drops under 60s, and halve `shard_count` on `local_client_execute_test` (30 to 15) and `transfer_manager_test` (15 to 8) to keep total shard count neutral.

PiperOrigin-RevId: 995837028
---
 third_party/xla/xla/backends/gpu/runtime/BUILD    | 2 +-
 third_party/xla/xla/backends/gpu/transforms/BUILD | 3 ++-
 third_party/xla/xla/tests/BUILD                   | 5 +++--
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index 579955e7b14..5ec75676509 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -4457,7 +4457,7 @@ xla_test(
     name = "topk_test",
     srcs = ["topk_test.cc"],
     backends = ["gpu"],
-    shard_count = 4,
+    shard_count = 12,
     deps = [
         ":topk",
         "//xla:types",
diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD
index fb6871a918e..e6343db02bf 100644
--- a/third_party/xla/xla/backends/gpu/transforms/BUILD
+++ b/third_party/xla/xla/backends/gpu/transforms/BUILD
@@ -3126,6 +3126,7 @@ xla_test(
     name = "scatter_determinism_expander_test",
     srcs = ["scatter_determinism_expander_test.cc"],
     backends = ["gpu"],
+    shard_count = 3,
     deps = [
         ":scatter_determinism_expander",
         "//xla:literal",
@@ -3700,7 +3701,7 @@ xla_test(
         "gb300": ["broken"],  # TODO(b/505574653): re-enable when fixed.
     },
     backends = ["gpu"],
-    shard_count = 10,
+    shard_count = 15,
     deps = [
         ":topk_specializer",
         "//xla/backends/gpu/tests:hlo_pjrt_gpu_test_base",
diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD
index bb38afabef7..e55ae6558cb 100644
--- a/third_party/xla/xla/tests/BUILD
+++ b/third_party/xla/xla/tests/BUILD
@@ -1725,6 +1725,7 @@ xla_test(
         "b200",
     ],
     data = ["data/cudnn_reproducer.hlo"],
+    shard_count = 5,
     deps = [
         ":hlo_pjrt_interpreter_reference_mixin",
         "//xla:error_spec",
@@ -3071,7 +3072,7 @@ xla_test(
         "gpu",
         "interpreter",
     ],
-    shard_count = 30,
+    shard_count = 15,
     tags = ["optonly"],
     use_legacy_runtime = True,
     deps = [
@@ -3225,7 +3226,7 @@ xla_test(
         "gpu",
         "interpreter",
     ],
-    shard_count = 15,
+    shard_count = 8,
     use_legacy_runtime = True,
     deps = [
         ":literal_test_util",

From f314e05d8c79e2e48ba1cf360bc725ad26b4b6a0 Mon Sep 17 00:00:00 2001
From: Tori Baker <vwbaker@google.com>
Date: Thu, 8 Oct 2026 08:04:56 -0700
Subject: [PATCH 0298/1000] Don't fuse transposes the producer can absorb in
 GemmFusion V2.

GemmFusion V2 fuses transposes that only reorder the sub-dimensions of a non-contracting dimension (M or N), which GemmFusion V1 does not. Such a transpose saves no memory traffic when its producer is emitted by a separate kernel anyway: as long as the minor-most dimension stays in place, the transpose can be fused into that kernel at no extra cost. Fusing it into the GEMM instead makes M or N multi-strided, which constrains the tiling of the Triton GEMM and requires a separate transpose kernel if the autotuner picks cuBLAS for the fusion.

Leave such a transpose to its producer if the producer is a single-user elementwise instruction that GemmFusion would not fuse into the GEMM itself.

PiperOrigin-RevId: 995839798
---
 .../backends/gpu/transforms/gemm_fusion.cc    |  34 ++++-
 .../gpu/transforms/gemm_fusion_test.cc        | 119 ++++++++++++++++++
 2 files changed, 152 insertions(+), 1 deletion(-)

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 f7409e39cc6..7ab4fc1a7ed 100644
--- a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc
@@ -1451,7 +1451,25 @@ bool HasCoalescedMinorDimension(const ShapeTracker& inverted_tracker,
   return minor_bits >= 128 && minor_bits % 128 == 0;
 }
 
+// Returns true if the producer of `transpose` is going to be emitted by a
+// separate kernel that can absorb the transpose at no extra memory cost, i.e.
+// if it is an elementwise instruction (looking through bitcasts and reshapes)
+// whose only user is `transpose` and that `ShouldFuseOperand` would not fuse
+// into the GEMM. `transpose` must be an instruction of the original module, so
+// that all users of its producer are visible.
+bool ProducerCanAbsorbTranspose(const HloInstruction& transpose) {
+  const HloInstruction* producer = transpose.operand(0);
+  while (HloPredicateIsOp<HloOpcode::kBitcast, HloOpcode::kReshape>(producer) &&
+         producer->user_count() == 1) {
+    producer = producer->operand(0);
+  }
+  return producer->user_count() == 1 && producer->IsElementwise() &&
+         !IsBinaryElementwiseOfBroadcastParamOrConst(*producer) &&
+         !triton_fusion::IsInputWorthFusing(*producer);
+}
+
 FusionDecision ShouldFuseTranspose(const HloInstruction& transpose,
+                                   const HloInstruction& original_transpose,
                                    const TrackerInfo& tracker) {
   const int64_t operand_index = tracker.dot_operand_index;
   ShapeTracker transpose_tracker = tracker.tracker;
@@ -1488,6 +1506,20 @@ FusionDecision ShouldFuseTranspose(const HloInstruction& transpose,
     return FusionDecision::Forbid(
         "Non-contracting RHS dimension has non-contiguous section.");
   }
+  // Fusing a transpose that makes a non-contracting dimension non-contiguous,
+  // without moving the minor-most dimension, saves no memory traffic if its
+  // producer is emitted by a separate kernel anyway: that kernel can write its
+  // output in the transposed order instead, at no extra cost. Fusing it,
+  // however, makes M or N multi-strided, which constrains the tiling of the
+  // GEMM and requires a separate transpose kernel if the GEMM ends up being
+  // lowered to cuBLAS. See b/571352199.
+  if (!inverted_tracker->MapsToOneStride(non_contracting) &&
+      !TransposesMinorDimension(&original_transpose) &&
+      ProducerCanAbsorbTranspose(original_transpose)) {
+    return FusionDecision::Forbid(
+        "Transpose within a non-contracting dimension can be fused into its "
+        "producer instead.");
+  }
   return FusionDecision::Allow();
 }
 
@@ -1504,7 +1536,7 @@ FusionDecision ShouldFuseOperand(HloInstruction* operand,
       if (!tracker.has_value()) {
         return FusionDecision::Forbid("No shape tracker found for transpose.");
       }
-      return ShouldFuseTranspose(*operand, *tracker);
+      return ShouldFuseTranspose(*operand, original_operand, *tracker);
     case HloOpcode::kConcatenate:
       if (!tracker.has_value()) {
         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 ec17c87b225..7f58ce7c73a 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
@@ -47,6 +47,7 @@ namespace {
 
 using ::absl_testing::IsOkAndHolds;
 using ::testing::Contains;
+using ::testing::Each;
 using ::testing::ElementsAre;
 using ::testing::FieldsAre;
 using ::testing::UnorderedElementsAre;
@@ -2647,6 +2648,124 @@ ENTRY e {
               GmockMatch(m::Fusion()));
 }
 
+TEST_P(GemmFusionTestV2,
+       DoNotFuseNonContractingTransposeIfProducerCanAbsorbIt) {
+  // The transpose only reorders the sub-dimensions of the LHS non-contracting
+  // dimension, and its producer (a normalization) is not fused into the GEMM.
+  // The transpose saves no memory traffic in the GEMM fusion, so it is left for
+  // its producer, while the bias + ReLU epilogue is still fused.
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
+HloModule m
+
+ENTRY e {
+  p0 = f16[16,8,3,64]{3,2,1,0} parameter(0)
+  p1 = f16[16,8,3]{2,1,0} parameter(1)
+  rsqrt = f16[16,8,3]{2,1,0} rsqrt(p1)
+  scale = f16[16,8,3,64]{3,2,1,0} broadcast(rsqrt), dimensions={0,1,2}
+  normalized = f16[16,8,3,64]{3,2,1,0} multiply(p0, scale)
+  transpose = f16[3,16,8,64]{3,2,1,0} transpose(normalized),
+    dimensions={2,0,1,3}
+  bitcast = f16[384,64]{1,0} bitcast(transpose)
+  p2 = f16[64,64]{1,0} parameter(2)
+  dot = f16[384,64]{1,0} dot(bitcast, p2),
+    lhs_contracting_dims={1}, rhs_contracting_dims={0}
+  p3 = f16[64]{0} parameter(3)
+  bias = f16[384,64]{1,0} broadcast(p3), dimensions={1}
+  add = f16[384,64]{1,0} add(dot, bias)
+  zero = f16[] constant(0)
+  zeros = f16[384,64]{1,0} broadcast(zero), dimensions={}
+  ROOT relu = f16[384,64]{1,0} maximum(add, zeros)
+})"));
+  ASSERT_THAT(GemmFusion(device_info_).Run(module.get()), IsOkAndHolds(true));
+  const HloInstruction* fusion =
+      module->entry_computation()->root_instruction();
+  ASSERT_THAT(fusion, GmockMatch(m::Fusion()));
+  EXPECT_THAT(fusion->fused_expression_root(), GmockMatch(m::Maximum()));
+  EXPECT_THAT(fusion->operands(),
+              Contains(GmockMatch(m::AnyOf<HloInstruction>(
+                  m::Transpose(m::Multiply()),
+                  m::Bitcast(m::Transpose(m::Multiply()))))));
+}
+
+TEST_P(GemmFusionTestV2, FuseNonContractingTransposeIfProducerIsFused) {
+  // The producer (a multiply by a broadcast parameter) is fused into the GEMM,
+  // so the transpose is fused as well.
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
+HloModule m
+
+ENTRY e {
+  p0 = f16[16,8,3,64]{3,2,1,0} parameter(0)
+  p1 = f16[16,8,3]{2,1,0} parameter(1)
+  scale = f16[16,8,3,64]{3,2,1,0} broadcast(p1), dimensions={0,1,2}
+  scaled = f16[16,8,3,64]{3,2,1,0} multiply(p0, scale)
+  transpose = f16[3,16,8,64]{3,2,1,0} transpose(scaled), dimensions={2,0,1,3}
+  bitcast = f16[384,64]{1,0} bitcast(transpose)
+  p2 = f16[64,64]{1,0} parameter(2)
+  ROOT dot = f16[384,64]{1,0} dot(bitcast, p2),
+    lhs_contracting_dims={1}, rhs_contracting_dims={0}
+})"));
+  ASSERT_THAT(GemmFusion(device_info_).Run(module.get()), IsOkAndHolds(true));
+  const HloInstruction* fusion =
+      module->entry_computation()->root_instruction();
+  ASSERT_THAT(fusion, GmockMatch(m::Fusion()));
+  EXPECT_THAT(fusion->operands(), Each(GmockMatch(ParamOrBitcastParam())));
+}
+
+TEST_P(GemmFusionTestV2,
+       FuseTransposeKeepingNonContractingDimensionContiguous) {
+  // The transpose only swaps the batch and non-contracting dimensions.
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
+HloModule m
+
+ENTRY e {
+  p0 = f16[128,8,64]{2,1,0} parameter(0)
+  p1 = f16[128,8]{1,0} parameter(1)
+  rsqrt = f16[128,8]{1,0} rsqrt(p1)
+  scale = f16[128,8,64]{2,1,0} broadcast(rsqrt), dimensions={0,1}
+  normalized = f16[128,8,64]{2,1,0} multiply(p0, scale)
+  transpose = f16[8,128,64]{2,1,0} transpose(normalized), dimensions={1,0,2}
+  p2 = f16[8,64,32]{2,1,0} parameter(2)
+  ROOT dot = f16[8,128,32]{2,1,0} dot(transpose, p2),
+    lhs_batch_dims={0}, lhs_contracting_dims={2},
+    rhs_batch_dims={0}, rhs_contracting_dims={1}
+})"));
+  ASSERT_THAT(GemmFusion(device_info_).Run(module.get()), IsOkAndHolds(true));
+  const HloInstruction* fusion =
+      module->entry_computation()->root_instruction();
+  ASSERT_THAT(fusion, GmockMatch(m::Fusion()));
+  EXPECT_THAT(fusion->operands(),
+              UnorderedElementsAre(GmockMatch(m::Multiply()),
+                                   GmockMatch(m::Parameter())));
+}
+
+TEST_P(GemmFusionTestV2, FuseNonContractingTransposeMovingMinorDimension) {
+  // The contracting dimension is interleaved between the two parts of the
+  // non-contracting dimension, but the transpose also moves the minor
+  // dimension, so it would not be free to fuse into the producer.
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
+HloModule m
+
+ENTRY e {
+  p0 = f16[3,64,128]{2,1,0} parameter(0)
+  p1 = f16[3,64]{1,0} parameter(1)
+  rsqrt = f16[3,64]{1,0} rsqrt(p1)
+  scale = f16[3,64,128]{2,1,0} broadcast(rsqrt), dimensions={0,1}
+  normalized = f16[3,64,128]{2,1,0} multiply(p0, scale)
+  transpose = f16[3,128,64]{2,1,0} transpose(normalized), dimensions={0,2,1}
+  bitcast = f16[384,64]{1,0} bitcast(transpose)
+  p2 = f16[64,64]{1,0} parameter(2)
+  ROOT dot = f16[384,64]{1,0} dot(bitcast, p2),
+    lhs_contracting_dims={1}, rhs_contracting_dims={0}
+})"));
+  ASSERT_THAT(GemmFusion(device_info_).Run(module.get()), IsOkAndHolds(true));
+  const HloInstruction* fusion =
+      module->entry_computation()->root_instruction();
+  ASSERT_THAT(fusion, GmockMatch(m::Fusion()));
+  EXPECT_THAT(fusion->operands(),
+              UnorderedElementsAre(GmockMatch(m::Multiply()),
+                                   GmockMatch(m::Parameter())));
+}
+
 TEST_P(GemmFusionTestV2, ConcatResetTrackerCrash) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m

From 01c5f2123cdf4cb66b46bd948ce39ac3f36f66c2 Mon Sep 17 00:00:00 2001
From: Andrei Ivanov <anivanov@nvidia.com>
Date: Thu, 8 Oct 2026 08:57:38 -0700
Subject: [PATCH 0299/1000] PR #49774: [GPU] Add first part of Vera Rubin NVL72
 specs.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/49774

📝 Summary of Changes

First part of changes enabling Vera Rubin NVL72 support.

🎯 Justification

New hardware enablement.

🚀 Kind of Contribution
✨ New Feature

📊 Benchmark (for Performance Improvements)

no

🧪 Unit Tests:

yes

🧪 Execution Tests:

no
Copybara import of the project:

--
88246c1ad6def083ece2d7b671f656617e78a35a by Ilia Sergachev <isergachev@nvidia.com>:

[GPU] Name Vera Rubin test configurations appropriately.

--
cc7882fba161c240b086e988945f5f8d2b2ac041 by Ilia Sergachev <isergachev@nvidia.com>:

[GPU] Add Rubin 10.7 compute capability.

--
4dde354375fdad32b8f7dc462d9a972007172b64 by Andrei Ivanov <anivanov@nvidia.com>:

[XLA:GPU] Add SM107 core performance data

Co-authored-by: Ilia Sergachev <isergachev@nvidia.com>

--
6d62dc3a546a952b8efc10d49a4874d0e9bb64cc by Andrei Ivanov <anivanov@nvidia.com>:

[XLA:GPU] Add SM107 HLO op profiles

Co-authored-by: Ilia Sergachev <isergachev@nvidia.com>

--
3af5fbfbcc1abbf803d38e54357f217428c84f64 by Andrei Ivanov <anivanov@nvidia.com>:

[XLA:GPU] Add VR NVL72 target configuration

Co-authored-by: Ilia Sergachev <isergachev@nvidia.com>

Merging this change closes #49774

PiperOrigin-RevId: 995868665
---
 third_party/xla/build_tools/lint/tags.py      |   4 +-
 .../xla/xla/backends/gpu/target_config/BUILD  |   1 +
 .../embedded_target_config_test.cc            |   2 +-
 .../gpu/target_config/specs/vr_nvl72.txtpb    | 261 ++++++
 .../gpu/target_config/target_config.cc        |   3 +
 .../gpu/target_config/target_config.h         |   1 +
 .../gpu/target_config/target_config_map.bzl   |   1 +
 .../gpu/target_config/target_config_test.cc   |   1 +
 third_party/xla/xla/python/_gpu_spec.pyi      |  12 +-
 third_party/xla/xla/python/gpu_spec.cc        |   1 +
 third_party/xla/xla/service/BUILD             |   1 +
 .../service/gpu/gpu_device_info_for_tests.cc  |  31 +
 .../service/gpu/gpu_device_info_for_tests.h   |   4 +
 .../llvm_gpu_backend/nvptx_backend_test.cc    |   2 +-
 .../service/gpu/model/hlo_op_profiles_data.h  | 845 ++++++++++++++++++
 .../service/gpu/model/hlo_op_profiles_test.cc |  14 +
 third_party/xla/xla/service/xla_compile.bzl   |  12 +-
 .../cuda/cuda_compute_capability.h            |   5 +
 .../cuda/cuda_compute_capability_test.cc      |   5 +
 .../cuda/cuda_core_info_table.cc              |  38 +-
 .../cuda/cuda_core_info_table_test.cc         |  38 +
 third_party/xla/xla/tests/BUILD               |   1 -
 third_party/xla/xla/tests/backend_defs.bzl    |   6 +-
 23 files changed, 1274 insertions(+), 15 deletions(-)
 create mode 100644 third_party/xla/xla/backends/gpu/target_config/specs/vr_nvl72.txtpb

diff --git a/third_party/xla/build_tools/lint/tags.py b/third_party/xla/build_tools/lint/tags.py
index 6da264a0827..9dda4607aba 100644
--- a/third_party/xla/build_tools/lint/tags.py
+++ b/third_party/xla/build_tools/lint/tags.py
@@ -118,7 +118,7 @@ _TAGS_TO_DOCUMENTATION_MAP = {
     "xla_gb200": "Runs on a gb200.",
     "xla_gb300": "Runs on a gb300.",
     "xla_rtx6000pro": "Runs on an rtx6000pro.",
-    "xla_vr200": "Runs on a vr200.",
+    "xla_vr_nvl72": "Runs on Vera Rubin NVL72.",
     "xla_device_p100": "Runs on a p100.",
     "xla_device_v100": "Runs on a v100.",
     "xla_device_a100": "Runs on an a100.",
@@ -127,7 +127,7 @@ _TAGS_TO_DOCUMENTATION_MAP = {
     "xla_device_gb200": "Runs on a gb200.",
     "xla_device_gb300": "Runs on a gb300.",
     "xla_device_rtx6000pro": "Runs on an rtx6000pro.",
-    "xla_device_vr200": "Runs on a vr200.",
+    "xla_device_vr_nvl72": "Runs on Vera Rubin NVL72.",
     # Below tags are consumed by `xla_test`.
     "test_migrated_to_hlo_runner_pjrt": (
         "Adds the appropriate `xla/tests:pjrt_$BACKEND_client_registry` to the"
diff --git a/third_party/xla/xla/backends/gpu/target_config/BUILD b/third_party/xla/xla/backends/gpu/target_config/BUILD
index 8714715d9af..214160175f2 100644
--- a/third_party/xla/xla/backends/gpu/target_config/BUILD
+++ b/third_party/xla/xla/backends/gpu/target_config/BUILD
@@ -201,6 +201,7 @@ xla_test(
         "b200",
         "gb200",
         "gb300",
+        "vr_nvl72",
         "amdgpu_any",
         "intelgpu_any",
     ],
diff --git a/third_party/xla/xla/backends/gpu/target_config/embedded_target_config_test.cc b/third_party/xla/xla/backends/gpu/target_config/embedded_target_config_test.cc
index 63acc4a4b92..d1e578d3f7b 100644
--- a/third_party/xla/xla/backends/gpu/target_config/embedded_target_config_test.cc
+++ b/third_party/xla/xla/backends/gpu/target_config/embedded_target_config_test.cc
@@ -41,7 +41,7 @@ TEST(EmbeddedTargetConfigTest, DeviceInfoMatches) {
   for (const std::string file_name :
        {"a100_pcie_80", "a100_sxm_40", "a100_sxm_80", "a6000", "b200",
         "b200_mig", "bmg_g21", "h100_pcie", "h100_sxm", "h100_sxm_mig", "p100",
-        "v100", "mi200", "rtx6000pro", "gb200", "gb300"}) {
+        "v100", "mi200", "rtx6000pro", "gb200", "gb300", "vr_nvl72"}) {
     GpuTargetConfigProto proto;
     std::string spec_string;
     ASSERT_OK(tsl::ReadFileToString(
diff --git a/third_party/xla/xla/backends/gpu/target_config/specs/vr_nvl72.txtpb b/third_party/xla/xla/backends/gpu/target_config/specs/vr_nvl72.txtpb
new file mode 100644
index 00000000000..5c831faa7c5
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/target_config/specs/vr_nvl72.txtpb
@@ -0,0 +1,261 @@
+# Copyright 2026 The OpenXLA Authors.
+#
+# 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.
+
+gpu_device_info {
+  threads_per_block_limit: 1024
+  threads_per_warp: 32
+  shared_memory_per_block: 49152
+  shared_memory_per_core: 233472
+  threads_per_core_limit: 1024
+  core_count: 212
+  fpus_per_core: 128
+  block_dim_limit_x: 2147483647
+  block_dim_limit_y: 65535
+  block_dim_limit_z: 65535
+  memory_bandwidth: 19160064000000
+  l2_cache_size: 132120576
+  clock_rate_ghz: 2.424
+  device_memory_size: 299103551488
+  shared_memory_per_block_optin: 232448
+  cuda_compute_capability {
+    major: 10
+    minor: 7
+    feature_extension: ACCELERATED_FEATURES
+  }
+  registers_per_core_limit: 65536
+  registers_per_block_limit: 65536
+  scalar_unit_description {
+    rate_infos {
+      key: 4
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 64
+        ops_per_clock: 1
+      }
+    }
+    rate_infos {
+      key: 8
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 64
+        ops_per_clock: 1
+      }
+    }
+    rate_infos {
+      key: 10
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 128
+        ops_per_clock: 1
+      }
+    }
+    rate_infos {
+      key: 11
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 128
+        ops_per_clock: 1
+      }
+    }
+    rate_infos {
+      key: 12
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 32
+        ops_per_clock: 1
+      }
+    }
+    rate_infos {
+      key: 16
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 128
+        ops_per_clock: 1
+      }
+    }
+  }
+  matrix_unit_description {
+    rate_infos {
+      key: 2
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 4
+        ops_per_clock: 64
+      }
+    }
+    rate_infos {
+      key: 6
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 4
+        ops_per_clock: 64
+      }
+    }
+    rate_infos {
+      key: 10
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 4
+        ops_per_clock: 1024
+      }
+    }
+    rate_infos {
+      key: 11
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 4
+        ops_per_clock: 512
+      }
+    }
+    rate_infos {
+      key: 12
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 4
+        ops_per_clock: 8
+      }
+    }
+    rate_infos {
+      key: 16
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 4
+        ops_per_clock: 1024
+      }
+    }
+    rate_infos {
+      key: 19
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 4
+        ops_per_clock: 4096
+      }
+    }
+    rate_infos {
+      key: 20
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 4
+        ops_per_clock: 4096
+      }
+    }
+    rate_infos {
+      key: 23
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 4
+        ops_per_clock: 4096
+      }
+    }
+    rate_infos {
+      key: 24
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 4
+        ops_per_clock: 4096
+      }
+    }
+    rate_infos {
+      key: 25
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 4
+        ops_per_clock: 4096
+      }
+    }
+    rate_infos {
+      key: 28
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 4
+        ops_per_clock: 4096
+      }
+    }
+    rate_infos {
+      key: 29
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 4
+        ops_per_clock: 4096
+      }
+    }
+    rate_infos {
+      key: 32
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 4
+        ops_per_clock: 8192
+      }
+    }
+    rate_infos {
+      key: 33
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 4
+        ops_per_clock: 4096
+      }
+    }
+    rate_infos {
+      key: 35
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 4
+        ops_per_clock: 4096
+      }
+    }
+    rate_infos {
+      key: 36
+      value {
+        clock_rate_ghz: 2.424
+        units_per_core: 4
+        ops_per_clock: 4096
+      }
+    }
+  }
+  driver_version: "13.4.0"
+  kernel_mode_driver_version: "615.31.0"
+  runtime_version: "13.4.0"
+  compile_time_toolkit_version: "13.4.0"
+  dnn_version: "9.25.0"
+  cub_version: "3.2.0"
+  device_interconnect_info {
+    active_links: 36
+  }
+  device_vendor: "NVIDIA Corporation"
+  platform_version: "Compute Capability 10.7a"
+  pci_bus_id: "<undefined>"
+  name: "NVIDIA VR NVL72"
+  model_str: "sm_10.7a with 299103551488B RAM, 212 cores, 2424000KHz clock, 4752000KHz mem clock, 132120576B L2$"
+  numa_node: -1
+  thread_dim_limit_x: 1024
+  thread_dim_limit_y: 1024
+  thread_dim_limit_z: 64
+  device_address_bits: 64
+  pcie_bandwidth: 17179869184
+  ecc_enabled: true
+  mem_clock_ghz: 4.752
+  reserved_shared_memory_per_block: 1024
+  max_blocks_per_multiprocessor: 16
+  collective_memory_granularity: 2097152
+}
+platform_name: "CUDA"
+dnn_version_info {
+  major: 9
+  minor: 25
+}
+device_description_str: "NVIDIA VR NVL72"
+runtime_version {
+  major: 13
+  minor: 4
+}
diff --git a/third_party/xla/xla/backends/gpu/target_config/target_config.cc b/third_party/xla/xla/backends/gpu/target_config/target_config.cc
index 3bf58d0b446..58881673a0e 100644
--- a/third_party/xla/xla/backends/gpu/target_config/target_config.cc
+++ b/third_party/xla/xla/backends/gpu/target_config/target_config.cc
@@ -98,6 +98,9 @@ absl::StatusOr<absl::string_view> GetEmbeddedGpuTargetConfigData(
     case GpuModel::PVC:
       filename = "pvc.txtpb";
       break;
+    case GpuModel::VR_NVL72:
+      filename = "vr_nvl72.txtpb";
+      break;
     case GpuModel::V100:
       filename = "v100.txtpb";
       break;
diff --git a/third_party/xla/xla/backends/gpu/target_config/target_config.h b/third_party/xla/xla/backends/gpu/target_config/target_config.h
index da0780da687..54fc6cf336f 100644
--- a/third_party/xla/xla/backends/gpu/target_config/target_config.h
+++ b/third_party/xla/xla/backends/gpu/target_config/target_config.h
@@ -45,6 +45,7 @@ enum class GpuModel {
   MI350,
   P100,
   PVC,
+  VR_NVL72,
   V100,
   GB200,
   GB300,
diff --git a/third_party/xla/xla/backends/gpu/target_config/target_config_map.bzl b/third_party/xla/xla/backends/gpu/target_config/target_config_map.bzl
index b14ccf42a82..7fbf7bd36e9 100644
--- a/third_party/xla/xla/backends/gpu/target_config/target_config_map.bzl
+++ b/third_party/xla/xla/backends/gpu/target_config/target_config_map.bzl
@@ -51,4 +51,5 @@ target_config_map = {
     "pvc": "//xla/backends/gpu/target_config:specs/pvc.txtpb",
     "rtx6000pro": "//xla/backends/gpu/target_config:specs/rtx6000pro.txtpb",
     "v100": "//xla/backends/gpu/target_config:specs/v100.txtpb",
+    "vr_nvl72": "//xla/backends/gpu/target_config:specs/vr_nvl72.txtpb",
 }
diff --git a/third_party/xla/xla/backends/gpu/target_config/target_config_test.cc b/third_party/xla/xla/backends/gpu/target_config/target_config_test.cc
index ad86186774f..d887fa8d9df 100644
--- a/third_party/xla/xla/backends/gpu/target_config/target_config_test.cc
+++ b/third_party/xla/xla/backends/gpu/target_config/target_config_test.cc
@@ -76,6 +76,7 @@ INSTANTIATE_TEST_SUITE_P(
         {"MI200", GpuModel::MI200, true},
         {"P100", GpuModel::P100, true},
         {"PVC", GpuModel::PVC, true},
+        {"VR_NVL72", GpuModel::VR_NVL72, true},
         {"V100", GpuModel::V100, true},
         {"GB200", GpuModel::GB200, true},
         {"GB300", GpuModel::GB300, true},
diff --git a/third_party/xla/xla/python/_gpu_spec.pyi b/third_party/xla/xla/python/_gpu_spec.pyi
index faafdcb145d..16fbd1e6162 100644
--- a/third_party/xla/xla/python/_gpu_spec.pyi
+++ b/third_party/xla/xla/python/_gpu_spec.pyi
@@ -55,15 +55,17 @@ class GpuModel(enum.Enum):
 
   PVC = 15
 
-  V100 = 16
+  VR_NVL72 = 16
 
-  GB200 = 17
+  V100 = 17
 
-  GB300 = 18
+  GB200 = 18
 
-  RTX6000PRO = 19
+  GB300 = 19
 
-  GFX1250 = 20
+  RTX6000PRO = 20
+
+  GFX1250 = 21
 
   MI350 = 13
 
diff --git a/third_party/xla/xla/python/gpu_spec.cc b/third_party/xla/xla/python/gpu_spec.cc
index b84011e6979..f8ca88ff23e 100644
--- a/third_party/xla/xla/python/gpu_spec.cc
+++ b/third_party/xla/xla/python/gpu_spec.cc
@@ -100,6 +100,7 @@ NB_MODULE(_gpu_spec, m) {
       .value("MI200", gpu::GpuModel::MI200)
       .value("P100", gpu::GpuModel::P100)
       .value("PVC", gpu::GpuModel::PVC)
+      .value("VR_NVL72", gpu::GpuModel::VR_NVL72)
       .value("V100", gpu::GpuModel::V100)
       .value("GB200", gpu::GpuModel::GB200)
       .value("GB300", gpu::GpuModel::GB300)
diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index 1a0bc2bb877..d821f25c065 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -5907,6 +5907,7 @@ xla_aot_compile_gpu(
     gpu_targets = [
         "h100",
         "b200",
+        "vr_nvl72",
     ],
     module = "xla_aot_compile_test.mlir",
     xla_flags = "--xla_gpu_use_new_autotune_cache_format=false",
diff --git a/third_party/xla/xla/service/gpu/gpu_device_info_for_tests.cc b/third_party/xla/xla/service/gpu/gpu_device_info_for_tests.cc
index 1ef4a6f5a6b..8f30470ca7b 100644
--- a/third_party/xla/xla/service/gpu/gpu_device_info_for_tests.cc
+++ b/third_party/xla/xla/service/gpu/gpu_device_info_for_tests.cc
@@ -147,6 +147,37 @@ stream_executor::DeviceDescription TestGpuDeviceInfo::B200SXMDeviceInfo(
   return b;
 }
 
+stream_executor::DeviceDescription TestGpuDeviceInfo::VRNVL72DeviceInfo(
+    stream_executor::GpuComputeCapability cc) {
+  stream_executor::DeviceDescription b;
+  b.set_name("NVIDIA VR NVL72");
+  b.set_gpu_compute_capability(cc);
+  b.set_threads_per_block_limit(1024);
+  b.set_threads_per_warp(32);
+  b.set_shared_memory_per_block(48 * 1024);
+  b.set_shared_memory_per_block_optin(227 * 1024);
+  b.set_shared_memory_per_core(228 * 1024);
+  b.set_threads_per_core_limit(1024);
+  b.set_core_count(212);
+  b.set_block_dim_limit_x(2'147'483'647);
+  b.set_block_dim_limit_y(65535);
+  b.set_block_dim_limit_z(65535);
+  b.set_memory_bandwidth(19'160'064'000'000);
+  b.set_l2_cache_size(126 * 1024 * 1024);
+  b.set_clock_rate_ghz(2.424);
+  b.set_device_memory_size(299'103'551'488);
+  b.set_registers_per_core_limit(65536);
+  b.set_registers_per_block_limit(65536);
+  b.set_runtime_version(stream_executor::SemanticVersion{13, 4, 0});
+  b.set_driver_version(stream_executor::SemanticVersion{13, 4, 0});
+
+  b.set_fpus_per_core(
+      stream_executor::gpu::GetFpusPerCore(*cc.cuda_compute_capability()));
+  stream_executor::gpu::FillExecutionUnitDesc(*cc.cuda_compute_capability(),
+                                              b.clock_rate_ghz(), b);
+  return b;
+}
+
 stream_executor::DeviceDescription TestGpuDeviceInfo::AMDMI210DeviceInfo() {
   const stream_executor::GpuComputeCapability cc(
       stream_executor::RocmComputeCapability("gfx90a"));
diff --git a/third_party/xla/xla/service/gpu/gpu_device_info_for_tests.h b/third_party/xla/xla/service/gpu/gpu_device_info_for_tests.h
index b74996d4fae..a85cbe5cf96 100644
--- a/third_party/xla/xla/service/gpu/gpu_device_info_for_tests.h
+++ b/third_party/xla/xla/service/gpu/gpu_device_info_for_tests.h
@@ -40,6 +40,10 @@ class TestGpuDeviceInfo {
       stream_executor::GpuComputeCapability cc =
           stream_executor::GpuComputeCapability{
               stream_executor::CudaComputeCapability(10, 0)});
+  static stream_executor::DeviceDescription VRNVL72DeviceInfo(
+      stream_executor::GpuComputeCapability cc =
+          stream_executor::GpuComputeCapability{
+              stream_executor::CudaComputeCapability(10, 7)});
   static stream_executor::DeviceDescription AMDMI210DeviceInfo();
   static stream_executor::DeviceDescription AMDMI300DeviceInfo();
   static stream_executor::DeviceDescription AMDMI350DeviceInfo();
diff --git a/third_party/xla/xla/service/gpu/llvm_gpu_backend/nvptx_backend_test.cc b/third_party/xla/xla/service/gpu/llvm_gpu_backend/nvptx_backend_test.cc
index 86d9d9b793f..d06de9254e8 100644
--- a/third_party/xla/xla/service/gpu/llvm_gpu_backend/nvptx_backend_test.cc
+++ b/third_party/xla/xla/service/gpu/llvm_gpu_backend/nvptx_backend_test.cc
@@ -69,7 +69,7 @@ TEST(UtilsTest, TestGetSmName) {
   ASSERT_EQ(nvptx::GetSmName(se::CudaComputeCapability{
                 10, 3, FeatureExtension::kAcceleratedFeatures}),
             "sm_103a");
-  ASSERT_EQ(nvptx::GetSmName(se::CudaComputeCapability{10, 7}), "sm_107");
+  ASSERT_EQ(nvptx::GetSmName(se::CudaComputeCapability::Rubin()), "sm_107");
   ASSERT_EQ(nvptx::GetSmName(se::CudaComputeCapability{
                 10, 7, FeatureExtension::kAcceleratedFeatures}),
             "sm_107a");
diff --git a/third_party/xla/xla/service/gpu/model/hlo_op_profiles_data.h b/third_party/xla/xla/service/gpu/model/hlo_op_profiles_data.h
index 27b4185cd77..7d2dcedddc6 100644
--- a/third_party/xla/xla/service/gpu/model/hlo_op_profiles_data.h
+++ b/third_party/xla/xla/service/gpu/model/hlo_op_profiles_data.h
@@ -23,6 +23,851 @@ namespace gpu {
 // xla/service/gpu/model:hlo_op_profiler_run
 
 inline constexpr char kDeviceHloOpProfiles[] = R"pb(
+  entries {
+    key: "sm_107"  # "Rubin"
+    value {
+      entries {
+        instruction {
+          opcode: "popcnt"
+          shape { element_type: S8 }
+        }
+        clock_cycles: 33
+      }
+      entries {
+        instruction {
+          opcode: "divide"
+          shape { element_type: S8 }
+        }
+        clock_cycles: 392
+      }
+      entries {
+        instruction {
+          opcode: "multiply"
+          shape { element_type: S8 }
+        }
+        clock_cycles: 19
+      }
+      entries {
+        instruction {
+          opcode: "popcnt"
+          shape { element_type: S16 }
+        }
+        clock_cycles: 33
+      }
+      entries {
+        instruction {
+          opcode: "divide"
+          shape { element_type: S16 }
+        }
+        clock_cycles: 397
+      }
+      entries {
+        instruction {
+          opcode: "multiply"
+          shape { element_type: S16 }
+        }
+        clock_cycles: 19
+      }
+      entries {
+        instruction {
+          opcode: "popcnt"
+          shape { element_type: S32 }
+        }
+        clock_cycles: 33
+      }
+      entries {
+        instruction {
+          opcode: "divide"
+          shape { element_type: S32 }
+        }
+        clock_cycles: 344
+      }
+      entries {
+        instruction {
+          opcode: "multiply"
+          shape { element_type: S32 }
+        }
+        clock_cycles: 4
+      }
+      entries {
+        instruction {
+          opcode: "popcnt"
+          shape { element_type: S64 }
+        }
+        clock_cycles: 33
+      }
+      entries {
+        instruction {
+          opcode: "add"
+          shape { element_type: S64 }
+        }
+        clock_cycles: 4
+      }
+      entries {
+        instruction {
+          opcode: "divide"
+          shape { element_type: S64 }
+        }
+        clock_cycles: 732
+      }
+      entries {
+        instruction {
+          opcode: "multiply"
+          shape { element_type: S64 }
+        }
+        clock_cycles: 9
+      }
+      entries {
+        instruction {
+          opcode: "subtract"
+          shape { element_type: S64 }
+        }
+        clock_cycles: 4
+      }
+      entries {
+        instruction {
+          opcode: "popcnt"
+          shape { element_type: U8 }
+        }
+        clock_cycles: 33
+      }
+      entries {
+        instruction {
+          opcode: "divide"
+          shape { element_type: U8 }
+        }
+        clock_cycles: 319
+      }
+      entries {
+        instruction {
+          opcode: "multiply"
+          shape { element_type: U8 }
+        }
+        clock_cycles: 19
+      }
+      entries {
+        instruction {
+          opcode: "popcnt"
+          shape { element_type: U16 }
+        }
+        clock_cycles: 33
+      }
+      entries {
+        instruction {
+          opcode: "divide"
+          shape { element_type: U16 }
+        }
+        clock_cycles: 319
+      }
+      entries {
+        instruction {
+          opcode: "multiply"
+          shape { element_type: U16 }
+        }
+        clock_cycles: 19
+      }
+      entries {
+        instruction {
+          opcode: "popcnt"
+          shape { element_type: U32 }
+        }
+        clock_cycles: 33
+      }
+      entries {
+        instruction {
+          opcode: "divide"
+          shape { element_type: U32 }
+        }
+        clock_cycles: 130
+      }
+      entries {
+        instruction {
+          opcode: "multiply"
+          shape { element_type: U32 }
+        }
+        clock_cycles: 4
+      }
+      entries {
+        instruction {
+          opcode: "popcnt"
+          shape { element_type: U64 }
+        }
+        clock_cycles: 33
+      }
+      entries {
+        instruction {
+          opcode: "add"
+          shape { element_type: U64 }
+        }
+        clock_cycles: 4
+      }
+      entries {
+        instruction {
+          opcode: "divide"
+          shape { element_type: U64 }
+        }
+        clock_cycles: 683
+      }
+      entries {
+        instruction {
+          opcode: "multiply"
+          shape { element_type: U64 }
+        }
+        clock_cycles: 9
+      }
+      entries {
+        instruction {
+          opcode: "subtract"
+          shape { element_type: U64 }
+        }
+        clock_cycles: 4
+      }
+      entries {
+        instruction {
+          opcode: "cbrt"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 179
+      }
+      entries {
+        instruction {
+          opcode: "cosine"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 450
+      }
+      entries {
+        instruction {
+          opcode: "erf"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 344
+      }
+      entries {
+        instruction {
+          opcode: "exponential"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 101
+      }
+      entries {
+        instruction {
+          opcode: "exponential-minus-one"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 164
+      }
+      entries {
+        instruction {
+          opcode: "log"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 106
+      }
+      entries {
+        instruction {
+          opcode: "log-plus-one"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 232
+      }
+      entries {
+        instruction {
+          opcode: "logistic"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 223
+      }
+      entries {
+        instruction {
+          opcode: "round-nearest-afz"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 58
+      }
+      entries {
+        instruction {
+          opcode: "round-nearest-even"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 53
+      }
+      entries {
+        instruction {
+          opcode: "rsqrt"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 96
+      }
+      entries {
+        instruction {
+          opcode: "sine"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 499
+      }
+      entries {
+        instruction {
+          opcode: "sqrt"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 101
+      }
+      entries {
+        instruction {
+          opcode: "tan"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 489
+      }
+      entries {
+        instruction {
+          opcode: "tanh"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 140
+      }
+      entries {
+        instruction {
+          opcode: "add"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 4
+      }
+      entries {
+        instruction {
+          opcode: "atan2"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 644
+      }
+      entries {
+        instruction {
+          opcode: "divide"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 43
+      }
+      entries {
+        instruction {
+          opcode: "multiply"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 4
+      }
+      entries {
+        instruction {
+          opcode: "subtract"
+          shape { element_type: F16 }
+        }
+        clock_cycles: 4
+      }
+      entries {
+        instruction {
+          opcode: "cbrt"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 155
+      }
+      entries {
+        instruction {
+          opcode: "cosine"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 567
+      }
+      entries {
+        instruction {
+          opcode: "erf"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 169
+      }
+      entries {
+        instruction {
+          opcode: "exponential"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 77
+      }
+      entries {
+        instruction {
+          opcode: "exponential-minus-one"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 145
+      }
+      entries {
+        instruction {
+          opcode: "log"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 193
+      }
+      entries {
+        instruction {
+          opcode: "log-plus-one"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 218
+      }
+      entries {
+        instruction {
+          opcode: "logistic"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 174
+      }
+      entries {
+        instruction {
+          opcode: "round-nearest-afz"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 53
+      }
+      entries {
+        instruction {
+          opcode: "round-nearest-even"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 33
+      }
+      entries {
+        instruction {
+          opcode: "rsqrt"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 77
+      }
+      entries {
+        instruction {
+          opcode: "sine"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 552
+      }
+      entries {
+        instruction {
+          opcode: "sqrt"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 77
+      }
+      entries {
+        instruction {
+          opcode: "tan"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 567
+      }
+      entries {
+        instruction {
+          opcode: "tanh"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 164
+      }
+      entries {
+        instruction {
+          opcode: "add"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 4
+      }
+      entries {
+        instruction {
+          opcode: "atan2"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 625
+      }
+      entries {
+        instruction {
+          opcode: "divide"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 24
+      }
+      entries {
+        instruction {
+          opcode: "multiply"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 4
+      }
+      entries {
+        instruction {
+          opcode: "subtract"
+          shape { element_type: F32 }
+        }
+        clock_cycles: 4
+      }
+      entries {
+        instruction {
+          opcode: "cbrt"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 829
+      }
+      entries {
+        instruction {
+          opcode: "cosine"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 765
+      }
+      entries {
+        instruction {
+          opcode: "erf"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 1284
+      }
+      entries {
+        instruction {
+          opcode: "exponential"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 659
+      }
+      entries {
+        instruction {
+          opcode: "exponential-minus-one"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 765
+      }
+      entries {
+        instruction {
+          opcode: "log"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 606
+      }
+      entries {
+        instruction {
+          opcode: "log-plus-one"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 1318
+      }
+      entries {
+        instruction {
+          opcode: "logistic"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 843
+      }
+      entries {
+        instruction {
+          opcode: "round-nearest-afz"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 135
+      }
+      entries {
+        instruction {
+          opcode: "round-nearest-even"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 38
+      }
+      entries {
+        instruction {
+          opcode: "rsqrt"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 678
+      }
+      entries {
+        instruction {
+          opcode: "sine"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 780
+      }
+      entries {
+        instruction {
+          opcode: "sqrt"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 761
+      }
+      entries {
+        instruction {
+          opcode: "tan"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 969
+      }
+      entries {
+        instruction {
+          opcode: "tanh"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 1270
+      }
+      entries {
+        instruction {
+          opcode: "add"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 33
+      }
+      entries {
+        instruction {
+          opcode: "atan2"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 1333
+      }
+      entries {
+        instruction {
+          opcode: "divide"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 1018
+      }
+      entries {
+        instruction {
+          opcode: "multiply"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 33
+      }
+      entries {
+        instruction {
+          opcode: "subtract"
+          shape { element_type: F64 }
+        }
+        clock_cycles: 33
+      }
+      entries {
+        instruction {
+          opcode: "cosine"
+          shape { element_type: C64 }
+        }
+        clock_cycles: 1013
+      }
+      entries {
+        instruction {
+          opcode: "exponential"
+          shape { element_type: C64 }
+        }
+        clock_cycles: 1260
+      }
+      entries {
+        instruction {
+          opcode: "exponential-minus-one"
+          shape { element_type: C64 }
+        }
+        clock_cycles: 1095
+      }
+      entries {
+        instruction {
+          opcode: "log"
+          shape { element_type: C64 }
+        }
+        clock_cycles: 765
+      }
+      entries {
+        instruction {
+          opcode: "log-plus-one"
+          shape { element_type: C64 }
+        }
+        clock_cycles: 1178
+      }
+      entries {
+        instruction {
+          opcode: "logistic"
+          shape { element_type: C64 }
+        }
+        clock_cycles: 2608
+      }
+      entries {
+        instruction {
+          opcode: "rsqrt"
+          shape { element_type: C64 }
+        }
+        clock_cycles: 2986
+      }
+      entries {
+        instruction {
+          opcode: "sine"
+          shape { element_type: C64 }
+        }
+        clock_cycles: 974
+      }
+      entries {
+        instruction {
+          opcode: "sqrt"
+          shape { element_type: C64 }
+        }
+        clock_cycles: 4634
+      }
+      entries {
+        instruction {
+          opcode: "tan"
+          shape { element_type: C64 }
+        }
+        clock_cycles: 1551
+      }
+      entries {
+        instruction {
+          opcode: "tanh"
+          shape { element_type: C64 }
+        }
+        clock_cycles: 1575
+      }
+      entries {
+        instruction {
+          opcode: "add"
+          shape { element_type: C64 }
+        }
+        clock_cycles: 14
+      }
+      entries {
+        instruction {
+          opcode: "atan2"
+          shape { element_type: C64 }
+        }
+        clock_cycles: 7145
+      }
+      entries {
+        instruction {
+          opcode: "divide"
+          shape { element_type: C64 }
+        }
+        clock_cycles: 552
+      }
+      entries {
+        instruction {
+          opcode: "multiply"
+          shape { element_type: C64 }
+        }
+        clock_cycles: 43
+      }
+      entries {
+        instruction {
+          opcode: "subtract"
+          shape { element_type: C64 }
+        }
+        clock_cycles: 14
+      }
+      entries {
+        instruction {
+          opcode: "cosine"
+          shape { element_type: C128 }
+        }
+        clock_cycles: 2297
+      }
+      entries {
+        instruction {
+          opcode: "exponential"
+          shape { element_type: C128 }
+        }
+        clock_cycles: 2099
+      }
+      entries {
+        instruction {
+          opcode: "exponential-minus-one"
+          shape { element_type: C128 }
+        }
+        clock_cycles: 2016
+      }
+      entries {
+        instruction {
+          opcode: "log"
+          shape { element_type: C128 }
+        }
+        clock_cycles: 3417
+      }
+      entries {
+        instruction {
+          opcode: "log-plus-one"
+          shape { element_type: C128 }
+        }
+        clock_cycles: 6219
+      }
+      entries {
+        instruction {
+          opcode: "logistic"
+          shape { element_type: C128 }
+        }
+        clock_cycles: 5987
+      }
+      entries {
+        instruction {
+          opcode: "rsqrt"
+          shape { element_type: C128 }
+        }
+        clock_cycles: 6050
+      }
+      entries {
+        instruction {
+          opcode: "sine"
+          shape { element_type: C128 }
+        }
+        clock_cycles: 2273
+      }
+      entries {
+        instruction {
+          opcode: "sqrt"
+          shape { element_type: C128 }
+        }
+        clock_cycles: 6918
+      }
+      entries {
+        instruction {
+          opcode: "tan"
+          shape { element_type: C128 }
+        }
+        clock_cycles: 4867
+      }
+      entries {
+        instruction {
+          opcode: "tanh"
+          shape { element_type: C128 }
+        }
+        clock_cycles: 4557
+      }
+      entries {
+        instruction {
+          opcode: "add"
+          shape { element_type: C128 }
+        }
+        clock_cycles: 33
+      }
+      entries {
+        instruction {
+          opcode: "atan2"
+          shape { element_type: C128 }
+        }
+        clock_cycles: 15285
+      }
+      entries {
+        instruction {
+          opcode: "divide"
+          shape { element_type: C128 }
+        }
+        clock_cycles: 2743
+      }
+      entries {
+        instruction {
+          opcode: "multiply"
+          shape { element_type: C128 }
+        }
+        clock_cycles: 106
+      }
+      entries {
+        instruction {
+          opcode: "subtract"
+          shape { element_type: C128 }
+        }
+        clock_cycles: 33
+      }
+    }
+  }
   entries {
     key: "sm_103"  # "GB300"
     value {
diff --git a/third_party/xla/xla/service/gpu/model/hlo_op_profiles_test.cc b/third_party/xla/xla/service/gpu/model/hlo_op_profiles_test.cc
index 2d809a092b2..7efcb737fd8 100644
--- a/third_party/xla/xla/service/gpu/model/hlo_op_profiles_test.cc
+++ b/third_party/xla/xla/service/gpu/model/hlo_op_profiles_test.cc
@@ -145,6 +145,20 @@ TEST_F(HloOpProfilesTest, GetProfileSm103) {
       372);
 }
 
+TEST_F(HloOpProfilesTest, GetProfileSm107) {
+  auto hlo_op_profiles = HloOpProfiles::Load(kDeviceHloOpProfiles,
+                                             /*default_profile_name=*/"sm_80");
+  auto device_info = TestGpuDeviceInfo::VRNVL72DeviceInfo();
+
+  EXPECT_EQ(HloOpProfiles::GetProfileName(device_info), "sm_107");
+  const auto& op_profile = hlo_op_profiles->GetProfile(device_info);
+  ASSERT_TRUE(op_profile.contains(
+      std::make_pair(HloOpcode::kDivide, PrimitiveType::S8)));
+  EXPECT_EQ(
+      op_profile.at(std::make_pair(HloOpcode::kDivide, PrimitiveType::S8)),
+      392);
+}
+
 TEST_F(HloOpProfilesTest, GetProfileSm100) {
   auto hlo_op_profiles = HloOpProfiles::Load(kDeviceHloOpProfiles,
                                              /*default_profile_name=*/"sm_80");
diff --git a/third_party/xla/xla/service/xla_compile.bzl b/third_party/xla/xla/service/xla_compile.bzl
index 11096f37ed1..1dfef3cb54b 100644
--- a/third_party/xla/xla/service/xla_compile.bzl
+++ b/third_party/xla/xla/service/xla_compile.bzl
@@ -32,11 +32,17 @@ xla_aot_compile(
 load("//xla:xla.default.bzl", "xla_compile_target_cpu")
 load("//xla/backends/gpu/target_config:target_config_map.bzl", gpu_target_config_map = "target_config_map")
 load("//xla/tsl:package_groups.bzl", "DEFAULT_LOAD_VISIBILITY")
+load("//xla/tsl/platform/default:cuda_build_defs.bzl", "if_cuda_newer_than")
 
 visibility(DEFAULT_LOAD_VISIBILITY)
 
 xla_compile_tool = "//xla/service:xla_compile"
 
+_MIN_CUDA_VERSION = {
+    # SM107a requires PTX 9.4, available starting with CUDA 13.4.
+    "vr_nvl72": "13_4",
+}
+
 def target_llvm_triple():
     """Returns the target LLVM triple to be used for compiling the target."""
 
@@ -129,6 +135,7 @@ def xla_aot_compile_gpu(
             srcs = [module, gpu_target_config_map[target], autotune_results],
             outs = [name + "_" + target],
             cmd = cmd,
+            tags = ["manual"] if target in _MIN_CUDA_VERSION else [],
             tools = [xla_compile_tool],
             # copybara:comment_begin(oss-only)
             target_compatible_with = select({
@@ -137,7 +144,10 @@ def xla_aot_compile_gpu(
             }),
             # copybara:comment_end
         )
-        res.append(compiled_binary)
+        if target in _MIN_CUDA_VERSION:
+            res = res + if_cuda_newer_than(_MIN_CUDA_VERSION[target], [compiled_binary])
+        else:
+            res = res + [compiled_binary]
     native.filegroup(
         name = name,
         data = res,
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_compute_capability.h b/third_party/xla/xla/stream_executor/cuda/cuda_compute_capability.h
index 0c2b647fcb0..1b41dc59308 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_compute_capability.h
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_compute_capability.h
@@ -139,6 +139,11 @@ struct CudaComputeCapability {
                                  FeatureExtension::kFamilyCompatibleFeatures};
   }
 
+  // Represents the baseline Rubin architecture with compute capability 10.7.
+  constexpr static CudaComputeCapability Rubin() {
+    return CudaComputeCapability{kBlackwell, 7, FeatureExtension::kNone};
+  }
+
   // Returns true if the compute capability is at least
   // `other_major.other_minor`. It is equivalent to
   // this->SupportsAllFeaturesOf(CudaComputeCapability{other_major,
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_compute_capability_test.cc b/third_party/xla/xla/stream_executor/cuda/cuda_compute_capability_test.cc
index 4aecba4332d..f124fb3aaae 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_compute_capability_test.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_compute_capability_test.cc
@@ -251,6 +251,11 @@ TEST(CudaComputeCapabilityTest, GenerationNumericTest) {
   EXPECT_TRUE(CudaComputeCapability(8, 0).IsAtLeastAmpere());
 }
 
+TEST(CudaComputeCapabilityTest, Rubin) {
+  EXPECT_EQ(CudaComputeCapability::Rubin(), CudaComputeCapability(10, 7));
+  EXPECT_EQ(CudaComputeCapability::Rubin().GetPtxAsTargetName(), "sm_107");
+}
+
 TEST(CudaComputeCapabilityTest, ComparisonTest) {
   using FeatureExtension = CudaComputeCapability::FeatureExtension;
 
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_core_info_table.cc b/third_party/xla/xla/stream_executor/cuda/cuda_core_info_table.cc
index f3ab6c8ac76..4b07c540fbd 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_core_info_table.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_core_info_table.cc
@@ -53,6 +53,10 @@ CoreInfo FindCudaCoreInfo(CudaComputeCapability cc) {
   //
   // [B200Specs] B200 Specs
   // https://www.techpowerup.com/gpu-specs/b200.c4210
+  //
+  // [RubinArch] Inside NVIDIA Rubin GPU Architecture: Powering the Era of
+  // Agentic AI
+  // https://developer.nvidia.com/blog/inside-nvidia-rubin-gpu-architecture-powering-the-era-of-agentic-ai/
 
   // =============== Constants ===============
   // Make sure to annotate with the sources when adding new data.
@@ -75,6 +79,8 @@ CoreInfo FindCudaCoreInfo(CudaComputeCapability cc) {
   constexpr float kBlackwellTcClockScale = 0.934;
   constexpr int kBlackwellTcPerSm = 4;  // [B200Specs]
 
+  constexpr int kRubinTcPerSm = 4;  // [RubinArch]
+
   // =============== Lookup table ===============
   // Make sure to annotate with the sources when adding new data.
   static const absl::NoDestructor<std::vector<CoreInfoTableForCC>> kTable(
@@ -153,10 +159,40 @@ CoreInfo FindCudaCoreInfo(CudaComputeCapability cc) {
                {kF32, kBlackwellTcPerSm, 512, kBlackwellTcClockScale},
                // Assuming clock rate is the same as base, like Hopper.
                {kF64, kBlackwellTcPerSm, 16, 1.0},
+           }},
+          {CudaComputeCapability::Rubin(),
+           /*cuda_core_perf_table=*/
+           {
+               // DType, Units/SM, Ops/Clk
+               {kF16, 128, 1},
+               {kF32, 128, 1},
+               {kF64, 32, 1},
+               {kI32, 64, 1},
+           },
+           /*tensor_core_perf_table=*/
+           {
+               // DType  Units/SM   Ops/Clk
+               {kF4, kRubinTcPerSm, 8192},
+               {kF6, kRubinTcPerSm, 4096},
+               {kI8, kRubinTcPerSm, 64},
+               {kF8, kRubinTcPerSm, 4096},
+               {kF16, kRubinTcPerSm, 1024},
+               {kF32, kRubinTcPerSm, 512},
+               {kF64, kRubinTcPerSm, 8},
            }}});
 
+  // Prefer model-specific data when the table contains an exact compute
+  // capability match.
   for (const auto& config : *kTable) {
-    if (config.cc.major == cc.major) {
+    if (config.cc.major == cc.major && config.cc.minor == cc.minor) {
+      return CoreInfo{config.cuda_core_infos, config.tensor_core_infos};
+    }
+  }
+
+  // Minor-zero entries describe the fallback for the corresponding major
+  // compute-capability family.
+  for (const auto& config : *kTable) {
+    if (config.cc.major == cc.major && config.cc.minor == 0) {
       return CoreInfo{config.cuda_core_infos, config.tensor_core_infos};
     }
   }
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_core_info_table_test.cc b/third_party/xla/xla/stream_executor/cuda/cuda_core_info_table_test.cc
index 90ddda6c69e..ec3d3dd5f0d 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_core_info_table_test.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_core_info_table_test.cc
@@ -157,8 +157,46 @@ TEST(CudaCoreInfoTableTest, B300UsesBlackwellTable) {
   EXPECT_EQ(matrix_unit_desc->GetRateInfo(xla::BF16).value().units_per_core, 4);
 }
 
+TEST(CudaCoreInfoTableTest, VRNVL72UsesSm107Table) {
+  DeviceDescription vr_nvl72_device_info =
+      xla::gpu::TestGpuDeviceInfo::VRNVL72DeviceInfo();
+
+  CheckPeakOpsPerNs(vr_nvl72_device_info, /*is_matrix_unit=*/true,
+                    xla::PrimitiveType::F4E2M1FN, 34314.0);
+  CheckPeakOpsPerNs(vr_nvl72_device_info, /*is_matrix_unit=*/true,
+                    xla::PrimitiveType::F6E3M2FN, 17157.0);
+  CheckPeakOpsPerNs(vr_nvl72_device_info, /*is_matrix_unit=*/true,
+                    xla::PrimitiveType::S8, 268.0);
+  CheckPeakOpsPerNs(vr_nvl72_device_info, /*is_matrix_unit=*/true,
+                    xla::PrimitiveType::F8E4M3, 17157.0);
+  CheckPeakOpsPerNs(vr_nvl72_device_info, /*is_matrix_unit=*/true,
+                    xla::PrimitiveType::F16, 4289.0);
+  CheckPeakOpsPerNs(vr_nvl72_device_info, /*is_matrix_unit=*/true,
+                    xla::PrimitiveType::BF16, 4289.0);
+  CheckPeakOpsPerNs(vr_nvl72_device_info, /*is_matrix_unit=*/true,
+                    xla::PrimitiveType::F32, 2145.0);
+  CheckPeakOpsPerNs(vr_nvl72_device_info, /*is_matrix_unit=*/true,
+                    xla::PrimitiveType::F64, 33.5);
+  CheckPeakOpsPerNs(vr_nvl72_device_info, /*is_matrix_unit=*/false,
+                    xla::PrimitiveType::F16, 134.0);
+  CheckPeakOpsPerNs(vr_nvl72_device_info, /*is_matrix_unit=*/false,
+                    xla::PrimitiveType::BF16, 134.0);
+  CheckPeakOpsPerNs(vr_nvl72_device_info, /*is_matrix_unit=*/false,
+                    xla::PrimitiveType::F32, 134.0);
+  CheckPeakOpsPerNs(vr_nvl72_device_info, /*is_matrix_unit=*/false,
+                    xla::PrimitiveType::S32, 67.0);
+  CheckPeakOpsPerNs(vr_nvl72_device_info, /*is_matrix_unit=*/false,
+                    xla::PrimitiveType::F64, 33.5);
+
+  ASSERT_NE(vr_nvl72_device_info.scalar_unit_description(), nullptr);
+  EXPECT_FALSE(vr_nvl72_device_info.scalar_unit_description()
+                   ->GetRateInfo(xla::S8)
+                   .has_value());
+}
+
 TEST(CudaCoreInfoTableTest, GetFpusPerCore) {
   EXPECT_EQ(GetFpusPerCore(CudaComputeCapability::Hopper()), 128);
+  EXPECT_EQ(GetFpusPerCore(CudaComputeCapability::Rubin()), 128);
   EXPECT_EQ(GetFpusPerCore(CudaComputeCapability::Ampere()), 64);
   EXPECT_EQ(GetFpusPerCore(CudaComputeCapability::Volta()), 64);
   EXPECT_EQ(GetFpusPerCore(CudaComputeCapability::Pascal()), 64);
diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD
index e55ae6558cb..b0609dcdbd8 100644
--- a/third_party/xla/xla/tests/BUILD
+++ b/third_party/xla/xla/tests/BUILD
@@ -3631,7 +3631,6 @@ xla_test(
         "//xla:xla_data_proto_cc",
         "//xla/hlo/builder:xla_builder",
         "//xla/tsl/lib/core:status_test_util",
-        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest",
diff --git a/third_party/xla/xla/tests/backend_defs.bzl b/third_party/xla/xla/tests/backend_defs.bzl
index f2d31388d3b..ce55327a6ec 100644
--- a/third_party/xla/xla/tests/backend_defs.bzl
+++ b/third_party/xla/xla/tests/backend_defs.bzl
@@ -31,7 +31,7 @@ NVIDIA_GPU_BACKENDS = [
     "b200",
     "gb200",
     "gb300",
-    "vr200",
+    "vr_nvl72",
 ] + if_google([], ["rtx6000pro"])
 
 # The generic "gpu" backend includes the actual backends in this list.
@@ -42,7 +42,7 @@ NVIDIA_GPU_DEFAULT_BACKENDS = [
     "b200",
     "gb200",
     "gb300",
-    "vr200",
+    "vr_nvl72",
 ] + if_google([], ["rtx6000pro"])
 
 AMD_GPU_DEFAULT_BACKENDS = ["amdgpu_any"]
@@ -92,7 +92,7 @@ def prepare_nvidia_gpu_backend_data(backends, disabled_backends, backend_tags, b
         "b200": (10, 0),
         "gb200": (10, 0),
         "gb300": (10, 3),
-        "vr200": (10, 7),
+        "vr_nvl72": (10, 7),
         "rtx6000pro": (12, 0),
     }
     for gpu_backend in NVIDIA_GPU_BACKENDS:

From 6e50e2c1890b459e5ffbd08159a87f9d48ab2267 Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Thu, 8 Oct 2026 09:20:39 -0700
Subject: [PATCH 0300/1000] Check driver version for accelerated CUjit_target
 values

Combining `CU_COMPUTE_ACCELERATED_TARGET_BASE` with compute capabilities newer than `sm_90a` produces `CUjit_target` enum values that require CUDA driver 12.8 (`sm_100a`, `sm_101a`, `sm_120a`) or 12.9 (`sm_103a`) at runtime. Passing them to `cuLinkCreate` on older drivers fails with `CUDA_ERROR_INVALID_VALUE`.

Therefore this change checks `stream_exec_->GetDeviceDescription().driver_version()` in `xla/stream_executor/cuda/driver_compilation_provider.cc` before constructing an accelerated `CUjit_target`.

PiperOrigin-RevId: 995882686
---
 third_party/xla/xla/stream_executor/cuda/BUILD  |  1 +
 .../cuda/driver_compilation_provider.cc         | 17 +++++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/third_party/xla/xla/stream_executor/cuda/BUILD b/third_party/xla/xla/stream_executor/cuda/BUILD
index 149fa607070..8a259bdc49e 100644
--- a/third_party/xla/xla/stream_executor/cuda/BUILD
+++ b/third_party/xla/xla/stream_executor/cuda/BUILD
@@ -2620,6 +2620,7 @@ cc_library(
         ":ptx_compiler_helpers",
         "//xla/stream_executor:activate_context",
         "//xla/stream_executor:platform",
+        "//xla/stream_executor:semantic_version",
         "//xla/stream_executor:stream_executor_h",
         "//xla/tsl/cuda",  # buildcleaner: keep
         "//xla/tsl/platform:errors",
diff --git a/third_party/xla/xla/stream_executor/cuda/driver_compilation_provider.cc b/third_party/xla/xla/stream_executor/cuda/driver_compilation_provider.cc
index e512c3e3a28..41c6a1b48c8 100644
--- a/third_party/xla/xla/stream_executor/cuda/driver_compilation_provider.cc
+++ b/third_party/xla/xla/stream_executor/cuda/driver_compilation_provider.cc
@@ -43,6 +43,7 @@ limitations under the License.
 #include "xla/stream_executor/cuda/cuda_status.h"
 #include "xla/stream_executor/cuda/ptx_compiler_helpers.h"
 #include "xla/stream_executor/platform.h"
+#include "xla/stream_executor/semantic_version.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/platform/errors.h"
 
@@ -84,6 +85,22 @@ absl::StatusOr<Assembly> DriverCompilationProvider::CompileAndLink(
   // (WGMMA, TMA, and more) are only supported in CUDA 12+.
   if (cc.feature_extension ==
       CudaComputeCapability::FeatureExtension::kAcceleratedFeatures) {
+    SemanticVersion min_driver_version = SemanticVersion{12, 0, 0};
+    if (cc.major == 9 && cc.minor == 0) {
+      min_driver_version = SemanticVersion{12, 0, 0};
+    } else if ((cc.major == 10 && (cc.minor == 0 || cc.minor == 1)) ||
+               (cc.major == 12 && cc.minor == 0)) {
+      min_driver_version = SemanticVersion{12, 8, 0};
+    } else {
+      min_driver_version = SemanticVersion{12, 9, 0};
+    }
+
+    if (stream_exec_->GetDeviceDescription().driver_version() <
+        min_driver_version) {
+      return absl::UnimplementedError(absl::StrCat(
+          "Accelerated target ", cc.GetPtxAsTargetName(),
+          " requires CUDA driver version >= ", min_driver_version.ToString()));
+    }
     target =
         static_cast<CUjit_target>(target + CU_COMPUTE_ACCELERATED_TARGET_BASE);
   }

From a9c9d180ca99c3be56c0325b01a56ebb1bbff893 Mon Sep 17 00:00:00 2001
From: nishad shabbir <nishadshabbir28@gmail.com>
Date: Thu, 8 Oct 2026 22:52:57 +0530
Subject: [PATCH 0301/1000] Address review: reposition the shape check, widen
 the input matrix

The previous commit put _verify_compatible_image_shapes() above the
convert_image_dtype() calls, which broke a case that worked before: the
check compares the two shapes with an op that needs both images to have
the same dtype, and the conversions are what guaranteed that. A uint8
array against an integer list returned a value on master and raised
InvalidArgumentError from ShapeN after the move.

Put the check between the conversions and the reduction instead. A rank
below 3 is still reported as a rank problem rather than surfacing from
reduce_mean, and the mixed dtype case matches master again.

Tests: add tuple inputs and both mixed list/array orders, and a
regression test for the rank check. Mixed integer dtypes are left out
deliberately: convert_image_dtype() scales by each dtype's range, so an
int32 list compared against a uint8 array has no meaningful shared
expected value, before or after this change.
---
 tensorflow/python/ops/image_ops_impl.py | 15 +++++++++------
 tensorflow/python/ops/image_ops_test.py | 11 +++++++++++
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/tensorflow/python/ops/image_ops_impl.py b/tensorflow/python/ops/image_ops_impl.py
index aba1fd4888a..d73aee4ec92 100644
--- a/tensorflow/python/ops/image_ops_impl.py
+++ b/tensorflow/python/ops/image_ops_impl.py
@@ -4318,18 +4318,21 @@ def psnr(a, b, max_val, name=None):
     a = ops.convert_to_tensor(a, name='a')
     b = ops.convert_to_tensor(b, name='b')
 
-    # Check the shapes before computing anything, so a rank below 3 is
-    # reported as such instead of surfacing from the reduction below.
-    _, _, checks = _verify_compatible_image_shapes(a, b)
-    with ops.control_dependencies(checks):
-      a = array_ops.identity(a)
-
     # Need to convert the images to float32.  Scale max_val accordingly so that
     # PSNR is computed correctly.
     max_val = math_ops.cast(max_val, a.dtype)
     max_val = convert_image_dtype(max_val, dtypes.float32)
     a = convert_image_dtype(a, dtypes.float32)
     b = convert_image_dtype(b, dtypes.float32)
+
+    # Check the shapes before reducing, so a rank below 3 is reported as such
+    # rather than surfacing from the reduction. This has to come after the
+    # conversions above: the check compares the two shapes with an op that
+    # needs both images to have the same dtype.
+    _, _, checks = _verify_compatible_image_shapes(a, b)
+    with ops.control_dependencies(checks):
+      a = array_ops.identity(a)
+
     mse = math_ops.reduce_mean(math_ops.squared_difference(a, b), [-3, -2, -1])
     return math_ops.subtract(
         20 * math_ops.log(max_val) / math_ops.log(10.0),
diff --git a/tensorflow/python/ops/image_ops_test.py b/tensorflow/python/ops/image_ops_test.py
index 96673b5dc1e..8fb6cc5b987 100644
--- a/tensorflow/python/ops/image_ops_test.py
+++ b/tensorflow/python/ops/image_ops_test.py
@@ -6146,6 +6146,9 @@ class PSNRTest(test_util.TensorFlowTestCase):
 
     with self.cached_session():
       for a, b in ((image1.tolist(), image2.tolist()),
+                   (tuple(image1.tolist()), tuple(image2.tolist())),
+                   (image1.tolist(), image2.astype(np.float32)),
+                   (image1.astype(np.float32), image2.tolist()),
                    (image1.astype(np.float32), image2.astype(np.float32))):
         tf_psnr = self.evaluate(image_ops.psnr(a, b, 1.0, "psnr"))
         self.assertAllClose(expected, tf_psnr, atol=0.001)
@@ -6159,10 +6162,18 @@ class PSNRTest(test_util.TensorFlowTestCase):
       expected_int = self._PSNR_NumPy(
           image1_int.astype(np.float32), image2_int.astype(np.float32), 255)
       for a, b in ((image1_int.tolist(), image2_int.tolist()),
+                   (tuple(image1_int.tolist()), tuple(image2_int.tolist())),
                    (image1_int, image2_int)):
         tf_psnr_int = self.evaluate(image_ops.psnr(a, b, 255, "psnr_int"))
         self.assertAllClose(expected_int, tf_psnr_int, atol=0.001)
 
+  def testPSNRRankBelowThreeRaises(self):
+    # The shape check runs before the reduction, so a rank below 3 is
+    # reported as a rank problem rather than surfacing from reduce_mean.
+    for bad in ([1.0, 2.0], [[1.0, 2.0], [3.0, 4.0]]):
+      with self.assertRaisesRegex(ValueError, "must have rank at least 3"):
+        image_ops.psnr(bad, bad, 1.0)
+
   def testPSNRMultiImage(self):
     image1 = self._RandomImage((10, 8, 8, 1), 1)
     image2 = self._RandomImage((10, 8, 8, 1), 1)

From 873e9ed88f02c8ec7518a0830171ff2469258a86 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Thu, 8 Oct 2026 10:30:29 -0700
Subject: [PATCH 0302/1000] [XLA:GPU]: Support index-wise variadic instructions
 in TilingSpace.

For now only all-gather is allow listed. But can be expanded in follow ups.

PiperOrigin-RevId: 995924600
---
 .../tiling/experimental/tiling_space.cc       |  95 +++++--
 .../tiling/experimental/tiling_space.h        |  15 +
 .../tiling/experimental/tiling_space_test.cc  | 256 ++++++++++++++++++
 third_party/xla/xla/codegen/xtile/BUILD       |   1 +
 .../xtile/tiling_from_block_parameters.cc     |  91 +++++--
 5 files changed, 410 insertions(+), 48 deletions(-)

diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc
index 2ef6d30d4e7..0524dfce615 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc
@@ -65,8 +65,36 @@ std::string HloPtrToString(const HloInstruction* hlo) {
   return hlo == nullptr ? "nullptr" : hlo->ToString();
 }
 
+bool HasIndexWiseVariadic(const HloFusionAdaptor& fusion) {
+  return absl::c_any_of(fusion.MakeInstructionPostOrder(),
+                        [](const HloInstructionAdaptor& instr) {
+                          return IsIndexWiseVariadic(instr.instruction());
+                        });
+}
+
+// Returns the position of the first dimension of result `result_index` in the
+// ordered list of dimensions of `hlo`. For an index-wise variadic instruction
+// the dimensions of all results are laid out back to back: result k starts at
+// the sum of the ranks of results 0..k-1. For every other instruction the
+// results share the dimensions of the first one and the offset is 0.
+int64_t GetResultDimensionOffset(const HloInstruction& hlo,
+                                 int64_t result_index) {
+  if (!IsIndexWiseVariadic(hlo)) {
+    return 0;
+  }
+  int64_t offset = 0;
+  for (int64_t i = 0; i < result_index; ++i) {
+    offset += hlo.shape().tuple_shapes(i).dimensions().size();
+  }
+  return offset;
+}
+
 }  // namespace
 
+bool IsIndexWiseVariadic(const HloInstruction& hlo) {
+  return hlo.opcode() == HloOpcode::kAllGather && hlo.operand_count() > 1;
+}
+
 llvm::DenseMap<SymbolicExpr, SymbolicExpr> GetTileSizeReplacementMap(
     const TilingSpace& tiling_space, absl::Span<const int64_t> tile_sizes) {
   CHECK_EQ(tile_sizes.size(), tiling_space.dimensions().size());
@@ -138,6 +166,20 @@ void TilingSpace::ProcessInstruction(const HloInstruction& hlo) {
       ProcessRaggedDot(hlo);
       break;
     default:
+      if (HloInstruction::IsOpElementwise(hlo.opcode())) {
+        for (int64_t i = 0; i < hlo.shape().dimensions().size(); ++i) {
+          auto it = hlo_to_dimension_.find(std::make_pair(&hlo, i));
+          if (it != hlo_to_dimension_.end()) {
+            const DimensionInfo* dim_info = it->second;
+            for (const HloInstruction* operand : hlo.operands()) {
+              if (operand->shape().dimensions().size() ==
+                  hlo.shape().dimensions().size()) {
+                hlo_to_dimension_[std::make_pair(operand, i)] = dim_info;
+              }
+            }
+          }
+        }
+      }
       // TODO(goncharov): should have a explicit list of supported instructions?
       break;
   }
@@ -319,10 +361,13 @@ const Shape& GetFirstShape(const HloInstruction* instr, int64_t index) {
 
 // Propagate dimensions from get-tuple-element to its operand.
 void TilingSpace::ProcessGetTupleElement(const HloInstruction& hlo) {
+  const int64_t offset =
+      GetResultDimensionOffset(*hlo.operand(0), hlo.tuple_index());
   for (int64_t i = 0; i < hlo.shape().dimensions().size(); ++i) {
     auto it = hlo_to_dimension_.find(std::make_pair(&hlo, i));
     if (it != hlo_to_dimension_.end()) {
-      hlo_to_dimension_[std::make_pair(hlo.operand(0), i)] = it->second;
+      const DimensionInfo* dim_info = it->second;
+      hlo_to_dimension_[std::make_pair(hlo.operand(0), offset + i)] = dim_info;
     }
   }
 }
@@ -376,6 +421,16 @@ std::optional<const TilingSpace::RTVarInfo*> TilingSpace::GetRTVarInfo(
   return it->second;
 }
 
+bool TilingSpace::HasPerOutputTiles(const HloInstructionAdaptor& hlo) const {
+  return IsIndexWiseVariadic(hlo.instruction());
+}
+
+bool TilingSpace::HasPerOutputTiles() const {
+  return absl::c_any_of(hlo_to_dimension_, [](const auto& kv) {
+    return IsIndexWiseVariadic(*kv.first.first);
+  });
+}
+
 absl::Status TilingSpace::AssignTileSizes(
     absl::Span<const int64_t> tile_sizes) {
   if (!is_symbolic_) {
@@ -431,23 +486,24 @@ absl::Status TilingSpace::InitializeDimensions(
 
   for (const auto& root : roots) {
     const Shape& root_shape = root.shape();
-    if (!root.shape().IsArray() && root.opcode() != HloOpcode::kReduce &&
+    const HloInstruction& root_hlo = root.instruction();
+    if (!root_shape.IsArray() && root.opcode() != HloOpcode::kReduce &&
         root.opcode() != HloOpcode::kScan) {
       return absl::InvalidArgumentError(
           absl::StrCat("Unsupported root shape ", root_shape.ToString(),
-                       " for root ", root.instruction().ToString()));
+                       " for root ", root_hlo.ToString()));
     }
 
-    const Shape& shape = GetFirstShape(&root.instruction());
+    const Shape& shape = GetFirstShape(&root_hlo);
     for (auto [index, dim] : llvm::enumerate(shape.dimensions())) {
       DimensionSemantics dim_type = DimensionSemantics::kParallel;
       if (root.opcode() == HloOpcode::kScan) {
-        auto scan = Cast<HloScanInstruction>(&root.instruction());
+        auto scan = Cast<HloScanInstruction>(&root_hlo);
         if (index == scan->scan_dimension()) {
           dim_type = DimensionSemantics::kSequential;
         }
       }
-      AppendDimension(&root.instruction(), index, dim, dim_type);
+      AppendDimension(&root_hlo, index, dim, dim_type);
     }
   }
   return absl::OkStatus();
@@ -503,7 +559,7 @@ absl::StatusOr<std::unique_ptr<TilingSpace>> TilingSpace::Create(
   TF_RET_CHECK(module) << "Fusion has no module";
   const DebugOptions& debug_options = module->config().debug_options();
 
-  if (roots.size() == 1) {
+  if (roots.size() == 1 || HasIndexWiseVariadic(fusion)) {
     ABSL_RETURN_IF_ERROR(tiling_space->InitializeDimensions(roots));
   } else if (
       IsSameShapeMultiOutputFusion(roots, Shape::Equal().IgnoreElementType()) &&
@@ -531,25 +587,25 @@ absl::StatusOr<std::unique_ptr<TilingSpace>> TilingSpace::Create(
   // Second pass: Create the root tiles now that
   // `tiling_space->num_dimensions()` is known.
   for (const HloInstructionAdaptor& root : roots) {
-    const Shape& root_shape = root.shape();
     absl::Span<const int64_t> dims =
         GetFirstShape(&root.instruction()).dimensions();
     llvm::SmallVector<DimTile> dim_tiles;
     dim_tiles.reserve(dims.size());
     for (auto [index, dim] : llvm::enumerate(dims)) {
-      int64_t global_dim_id =
-          tiling_space->GetDimensionInfo(root.instruction(), index).id.value();
+      const DimensionInfo& dim_info =
+          tiling_space->GetDimensionInfo(root.instruction(), index);
       dim_tiles.push_back(
-          tiling_space->GetDefaultRootDimTile(TiledDimId(global_dim_id), dim));
+          tiling_space->GetDefaultRootDimTile(dim_info.id, dim));
     }
-    Tile tile{*tiling_space, std::move(dim_tiles)};
+    Tile root_tile{*tiling_space, std::move(dim_tiles)};
+    const Shape& root_shape = root.shape();
     if (root_shape.IsTuple()) {
       for (int64_t i = 0, e = root_shape.tuple_shapes().size(); i < e; ++i) {
-        tiling_space->tiled_roots_.push_back(tile);
+        tiling_space->tiled_roots_.push_back(root_tile);
       }
-      continue;
+    } else {
+      tiling_space->tiled_roots_.push_back(std::move(root_tile));
     }
-    tiling_space->tiled_roots_.push_back(std::move(tile));
   }
 
   return tiling_space;
@@ -578,9 +634,12 @@ std::unique_ptr<TilingSpace> TilingSpace::Clone(
   cloned->divisibility_constraints_ = divisibility_constraints_;
 
   cloned->dimensions_ = dimensions_;
-  cloned->hlo_to_dimension_.reserve(cloned->dimensions_.size());
-  for (const auto& dim : cloned->dimensions_) {
-    cloned->hlo_to_dimension_[std::make_pair(dim.hlo, dim.dim_position)] = &dim;
+  cloned->hlo_to_dimension_.reserve(hlo_to_dimension_.size());
+  // Populating an unordered map from another unordered map is order-independent
+  // since keys are unique and elements are only accessed via direct lookups.
+  // NOLINTNEXTLINE
+  for (const auto& [key, dim_ptr] : hlo_to_dimension_) {
+    cloned->hlo_to_dimension_[key] = &cloned->dimensions_[dim_ptr->id.value()];
   }
 
   cloned->rt_vars_ = rt_vars_;
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h
index 89f543ef713..db7dc9b02c4 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h
@@ -202,6 +202,15 @@ class TilingSpace {
 
   bool IsSymbolic() const { return is_symbolic_; }
 
+  // Returns true if `hlo` is a tuple-producing instruction tiled with one
+  // distinct tile per output (consumed through `get-tuple-element`), rather
+  // than a single shared tile for all outputs (such as variadic `reduce`).
+  bool HasPerOutputTiles(const HloInstructionAdaptor& hlo) const;
+
+  // Returns true if the tiling space contains a tuple-producing instruction
+  // tiled with one distinct tile per output.
+  bool HasPerOutputTiles() const;
+
   // Result of `SimplifyExpressions`.
   // TODO(b/565301234): follow up: we can also return simplified constraint
   // intervals but that requires extracting them from IndexingMap properly,
@@ -291,6 +300,12 @@ class TilingSpace {
 // Otherwise, return the shape itself.
 const Shape& GetFirstShape(const HloInstruction* instr, int64_t index = 0);
 
+// Returns true if `hlo` is a tuple-producing instruction whose output `k`
+// depends only on operand `k`, with the same tile mapping for every `k`.
+//
+// Currently only multi-operand all-gather qualifies.
+bool IsIndexWiseVariadic(const HloInstruction& hlo);
+
 // Returns a symbol replacement map to set concrete tile sizes.
 llvm::DenseMap<SymbolicExpr, SymbolicExpr> GetTileSizeReplacementMap(
     const TilingSpace& tiling_space, absl::Span<const int64_t> tile_sizes);
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc
index f937f620708..5c95e5f95c0 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc
@@ -34,6 +34,7 @@ limitations under the License.
 #include "xla/hlo/analysis/interval.h"
 #include "xla/hlo/analysis/symbolic_expr.h"
 #include "xla/hlo/analysis/symbolic_map_serialization.h"
+#include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/hlo/testlib/verified_hlo_module.h"
@@ -369,6 +370,205 @@ TEST_F(TilingSpaceTest, TwoOutputsEqualShapesParallelDims) {
       StatusIs(absl::StatusCode::kUnimplemented, HasSubstr("multiple roots")));
 }
 
+TEST_F(TilingSpaceTest, IsIndexWiseVariadic) {
+  ParseAndGetRoot(R"(
+    HloModule m
+    add_pair {
+      a = f32[] parameter(0)
+      b = s32[] parameter(1)
+      c = f32[] parameter(2)
+      d = s32[] parameter(3)
+      add_f = f32[] add(a, c)
+      add_s = s32[] add(b, d)
+      ROOT t = (f32[], s32[]) tuple(add_f, add_s)
+    }
+    ENTRY e {
+      p0 = f32[1,8] parameter(0)
+      p1 = s32[1] parameter(1)
+      p2 = f32[4,8] parameter(2)
+      p3 = s32[4,8] parameter(3)
+      c0 = f32[] constant(0)
+      c1 = s32[] constant(0)
+      single_ag = f32[4,8] all-gather(p0), replica_groups={{0,1,2,3}},
+        dimensions={0}
+      variadic_ag = (f32[4,8], s32[4]) all-gather(p0, p1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      variadic_reduce = (f32[4], s32[4]) reduce(p2, p3, c0, c1),
+        dimensions={1}, to_apply=add_pair
+      gte = f32[4,8] get-tuple-element(variadic_ag), index=0
+      ROOT t = (f32[4,8], f32[4,8]) tuple(single_ag, gte)
+    }
+  )");
+  HloComputation* entry = module_->entry_computation();
+  EXPECT_TRUE(
+      IsIndexWiseVariadic(*entry->GetInstructionWithName("variadic_ag")));
+  EXPECT_FALSE(
+      IsIndexWiseVariadic(*entry->GetInstructionWithName("single_ag")));
+  EXPECT_FALSE(
+      IsIndexWiseVariadic(*entry->GetInstructionWithName("variadic_reduce")));
+  EXPECT_FALSE(IsIndexWiseVariadic(*entry->GetInstructionWithName("gte")));
+  EXPECT_FALSE(IsIndexWiseVariadic(*entry->root_instruction()));
+}
+
+// An index-wise variadic instruction must be decomposed into GTEs + a tuple
+// root; a bare tuple-producing variadic root is rejected.
+TEST_F(TilingSpaceTest, BareVariadicAllGatherRootIsRejected) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    f {
+      p0 = f32[1,2,8,16] parameter(0)
+      p1 = s32[1,2] parameter(1)
+      ROOT ag = (f32[4,2,8,16], s32[4,2]) all-gather(p0, p1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+    }
+
+    ENTRY e {
+      p0 = f32[1,2,8,16] parameter(0)
+      p1 = s32[1,2] parameter(1)
+      ROOT fusion = (f32[4,2,8,16], s32[4,2]) fusion(p0, p1), kind=kCustom,
+        calls=f
+    }
+  )");
+  auto fusion_adaptor = HloFusionAdaptor::ForInstruction(root);
+  EXPECT_THAT(TilingSpace::Create(*fusion_adaptor, &mlir_context_),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Unsupported root shape")));
+}
+
+// A tuple-shaped index-wise variadic instruction consumed through GTE roots
+// gets independent tile dimensions for every output even though their shapes
+// differ. The dimensions of the outputs are laid out back to back in the
+// ordered dimension list of the instruction. The peer coupling between the
+// outputs is recovered later by the scheduler.
+TEST_F(TilingSpaceTest, VariadicAllGatherGetsPerOutputDimensions) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    f {
+      p0 = f32[1,2,8,16] parameter(0)
+      p1 = s32[1,2] parameter(1)
+      ag = (f32[4,2,8,16], s32[4,2]) all-gather(p0, p1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      gte0 = f32[4,2,8,16] get-tuple-element(ag), index=0
+      gte1 = s32[4,2] get-tuple-element(ag), index=1
+      ROOT t = (f32[4,2,8,16], s32[4,2]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      p0 = f32[1,2,8,16] parameter(0)
+      p1 = s32[1,2] parameter(1)
+      ROOT fusion = (f32[4,2,8,16], s32[4,2]) fusion(p0, p1), kind=kCustom,
+        calls=f
+    }
+  )");
+  auto fusion_adaptor = HloFusionAdaptor::ForInstruction(root);
+  ASSERT_OK_AND_ASSIGN(auto tiling_space,
+                       TilingSpace::Create(*fusion_adaptor, &mlir_context_));
+  const HloInstruction* ag =
+      root->fused_instructions_computation()->GetInstructionWithName("ag");
+  EXPECT_TRUE(tiling_space->HasPerOutputTiles(
+      HloInstructionAdaptor(*ag, fusion_adaptor.get())));
+  EXPECT_TRUE(tiling_space->HasPerOutputTiles());
+  EXPECT_EQ(tiling_space->num_dimensions(), 6);
+  for (int64_t i = 0; i < 6; ++i) {
+    EXPECT_EQ(tiling_space->GetDimensionInfo(*ag, i).id, TiledDimId(i));
+  }
+  EXPECT_THAT(*tiling_space, MatchString(R"(
+    Dimensions:
+        0 type: parallel size: 4 dim ID:0
+          hlo: %gte0 = f32[4,2,8,16]{3,2,1,0} get-tuple-element(%ag), index=0
+        1 type: parallel size: 2 dim ID:1
+          hlo: %gte0 = f32[4,2,8,16]{3,2,1,0} get-tuple-element(%ag), index=0
+        2 type: parallel size: 8 dim ID:2
+          hlo: %gte0 = f32[4,2,8,16]{3,2,1,0} get-tuple-element(%ag), index=0
+        3 type: parallel size: 16 dim ID:3
+          hlo: %gte0 = f32[4,2,8,16]{3,2,1,0} get-tuple-element(%ag), index=0
+        4 type: parallel size: 4 dim ID:0
+          hlo: %gte1 = s32[4,2]{1,0} get-tuple-element(%ag), index=1
+        5 type: parallel size: 2 dim ID:1
+          hlo: %gte1 = s32[4,2]{1,0} get-tuple-element(%ag), index=1
+    Root tiles:
+      0 root tile:
+           offsets [tid_0 * ts_0, tid_1 * ts_1, tid_2 * ts_2, tid_3 * ts_3]
+           sizes [ts_0, ts_1, ts_2, ts_3]
+           strides [1, 1, 1, 1]
+           upper bounds [4, 2, 8, 16]
+      1 root tile:
+           offsets [tid_4 * ts_4, tid_5 * ts_5]
+           sizes [ts_4, ts_5]
+           strides [1, 1]
+           upper bounds [4, 2]
+  )"));
+}
+
+TEST_F(TilingSpaceTest,
+       InteriorVariadicAllGatherPropagatesDimensionsFromRoots) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    f {
+      p0 = f32[2,16] parameter(0)
+      p1 = s32[4] parameter(1)
+      ag = (f32[8,16], s32[16]) all-gather(p0, p1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      gte0 = f32[8,16] get-tuple-element(ag), index=0
+      gte1 = s32[16] get-tuple-element(ag), index=1
+      c0 = bf16[8,16] convert(gte0)
+      c1 = f32[16] convert(gte1)
+      ROOT t = (bf16[8,16], f32[16]) tuple(c0, c1)
+    }
+
+    ENTRY e {
+      p0 = f32[2,16] parameter(0)
+      p1 = s32[4] parameter(1)
+      ROOT fusion = (bf16[8,16], f32[16]) fusion(p0, p1), kind=kCustom,
+        calls=f
+    }
+  )");
+  auto fusion_adaptor = HloFusionAdaptor::ForInstruction(root);
+  ASSERT_OK_AND_ASSIGN(auto tiling_space,
+                       TilingSpace::Create(*fusion_adaptor, &mlir_context_));
+  const HloInstruction* ag =
+      root->fused_instructions_computation()->GetInstructionWithName("ag");
+  EXPECT_EQ(tiling_space->GetDimensionInfo(*ag, 0).id, TiledDimId(0));
+  EXPECT_EQ(tiling_space->GetDimensionInfo(*ag, 1).id, TiledDimId(1));
+  EXPECT_EQ(tiling_space->GetDimensionInfo(*ag, 2).id, TiledDimId(2));
+}
+
+TEST_F(TilingSpaceTest,
+       VariadicAllGatherWithDotProducersGetsSequentialDimensions) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    f {
+      lhs0 = f32[2,8] parameter(0)
+      rhs0 = f32[8,16] parameter(1)
+      lhs1 = f32[4,8] parameter(2)
+      rhs1 = f32[8,32] parameter(3)
+      dot0 = f32[2,16] dot(lhs0, rhs0),
+        lhs_contracting_dims={1}, rhs_contracting_dims={0}
+      dot1 = f32[4,32] dot(lhs1, rhs1),
+        lhs_contracting_dims={1}, rhs_contracting_dims={0}
+      ag = (f32[8,16], f32[16,32]) all-gather(dot0, dot1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      gte0 = f32[8,16] get-tuple-element(ag), index=0
+      gte1 = f32[16,32] get-tuple-element(ag), index=1
+      ROOT t = (f32[8,16], f32[16,32]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      lhs0 = f32[2,8] parameter(0)
+      rhs0 = f32[8,16] parameter(1)
+      lhs1 = f32[4,8] parameter(2)
+      rhs1 = f32[8,32] parameter(3)
+      ROOT fusion = (f32[8,16], f32[16,32]) fusion(lhs0, rhs0, lhs1, rhs1),
+        kind=kCustom, calls=f
+    }
+  )");
+  auto fusion_adaptor = HloFusionAdaptor::ForInstruction(root);
+  ASSERT_OK_AND_ASSIGN(auto tiling_space,
+                       TilingSpace::Create(*fusion_adaptor, &mlir_context_));
+  EXPECT_EQ(tiling_space->num_dimensions(), 6);
+  EXPECT_EQ(tiling_space->num_parallel_dimensions(), 4);
+}
+
 class TilingSpaceSameShapeMultiOutputTest : public TilingSpaceTest {
  protected:
   DebugOptions GetDebugOptionsForTest() const override {
@@ -748,5 +948,61 @@ TEST_F(TilingSpaceTest, CloneIntoAnotherContextRebindsRootTiles) {
   EXPECT_TRUE(original_space->IsSymbolic());
 }
 
+TEST_F(TilingSpaceTest,
+       ElementwiseChainRehashDoesNotInvalidateDimensionPointers) {
+  auto module = ParseAndReturnVerifiedModule(R"(
+    HloModule m
+    fused_computation {
+      p0 = f32[2,4,8,16] parameter(0)
+      p1 = f32[2,4,8,16] parameter(1)
+      p2 = f32[2,4,8,16] parameter(2)
+      p3 = f32[2,4,8,16] parameter(3)
+      p4 = f32[2,4,8,16] parameter(4)
+      p5 = f32[2,4,8,16] parameter(5)
+      p6 = f32[2,4,8,16] parameter(6)
+      p7 = f32[2,4,8,16] parameter(7)
+      a0 = f32[2,4,8,16] add(p0, p1)
+      a1 = f32[2,4,8,16] multiply(a0, p2)
+      a2 = f32[2,4,8,16] subtract(a1, p3)
+      a3 = f32[2,4,8,16] maximum(a2, p4)
+      a4 = f32[2,4,8,16] minimum(a3, p5)
+      a5 = f32[2,4,8,16] divide(a4, p6)
+      a6 = f32[2,4,8,16] add(a5, p7)
+      a7 = f32[2,4,8,16] multiply(a6, a0)
+      a8 = f32[2,4,8,16] subtract(a7, a1)
+      a9 = f32[2,4,8,16] maximum(a8, a2)
+      a10 = f32[2,4,8,16] minimum(a9, a3)
+      ROOT out = f32[2,4,8,16] add(a10, a4)
+    }
+    ENTRY main {
+      p0 = f32[2,4,8,16] parameter(0)
+      p1 = f32[2,4,8,16] parameter(1)
+      p2 = f32[2,4,8,16] parameter(2)
+      p3 = f32[2,4,8,16] parameter(3)
+      p4 = f32[2,4,8,16] parameter(4)
+      p5 = f32[2,4,8,16] parameter(5)
+      p6 = f32[2,4,8,16] parameter(6)
+      p7 = f32[2,4,8,16] parameter(7)
+      ROOT fusion = f32[2,4,8,16] fusion(p0, p1, p2, p3, p4, p5, p6, p7),
+        kind=kLoop, calls=fused_computation
+    }
+  )");
+  ASSERT_OK(module.status());
+  auto fusion_adaptor = HloFusionAdaptor::ForInstruction(
+      module.value()->entry_computation()->root_instruction());
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<TilingSpace> space,
+                       TilingSpace::Create(*fusion_adaptor, &mlir_context_));
+  std::unique_ptr<TilingSpace> cloned = space->Clone();
+  const HloComputation* comp = module.value()
+                                   ->entry_computation()
+                                   ->root_instruction()
+                                   ->fused_instructions_computation();
+  for (const HloInstruction* param : comp->parameter_instructions()) {
+    for (int64_t dim = 0; dim < 4; ++dim) {
+      EXPECT_EQ(cloned->GetDimensionInfo(*param, dim).id, TiledDimId(dim));
+    }
+  }
+}
+
 }  // namespace
 }  // namespace xla::gpu::experimental
diff --git a/third_party/xla/xla/codegen/xtile/BUILD b/third_party/xla/xla/codegen/xtile/BUILD
index d7d99ee69b8..9ceb6138789 100644
--- a/third_party/xla/xla/codegen/xtile/BUILD
+++ b/third_party/xla/xla/codegen/xtile/BUILD
@@ -54,6 +54,7 @@ cc_library(
         "//xla/codegen/tiling/experimental:tile",
         "//xla/codegen/tiling/experimental:tiling_space_utils",
         "//xla/hlo/ir:hlo",
+        "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
diff --git a/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.cc b/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.cc
index 094f217b2d8..ef6d000f647 100644
--- a/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.cc
+++ b/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/algorithm/container.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
@@ -28,6 +29,7 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "absl/types/span.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 #include "xla/codegen/tiling/experimental/tiling_space.h"
 #include "xla/codegen/tiling/experimental/tiling_space_utils.h"
@@ -47,6 +49,60 @@ using ::xla::xtile::BlockLevelParameters;
 using DimensionSemantics =
     ::xla::gpu::experimental::TilingSpace::DimensionSemantics;
 
+absl::StatusOr<llvm::SmallVector<int64_t>> GetParallelTileSizes(
+    const ::xla::gpu::experimental::TilingSpace& tiling_space,
+    const BlockLevelParameters& block_level_parameters,
+    bool enable_same_shape_multi_output_fusion) {
+  TF_RET_CHECK(!block_level_parameters.output_tile_sizes.empty())
+      << "output_tile_sizes cannot be empty.";
+
+  if (block_level_parameters.output_tile_sizes.size() == 1) {
+    return llvm::to_vector(
+        llvm::ArrayRef(block_level_parameters.output_tile_sizes[0]));
+  }
+
+  // If the tiling space has one set of parallel dimensions per root (e.g.
+  // index-wise variadic fusions such as a variadic all-gather, see
+  // TilingSpace::Create), the per-root tile sizes are concatenated in root
+  // order.
+  if (tiling_space.HasPerOutputTiles()) {
+    llvm::SmallVector<int64_t> parallel_tile_sizes;
+    parallel_tile_sizes.reserve(tiling_space.num_parallel_dimensions());
+    for (const auto& output_tile_sizes :
+         block_level_parameters.output_tile_sizes) {
+      parallel_tile_sizes.append(output_tile_sizes.begin(),
+                                 output_tile_sizes.end());
+    }
+    return parallel_tile_sizes;
+  }
+
+  if (!enable_same_shape_multi_output_fusion) {
+    return Unimplemented(
+        "Only single-result fusions are supported for now. Received %d "
+        "roots.",
+        block_level_parameters.output_tile_sizes.size());
+  }
+
+  // For multi-output fusions with identical shapes, we enforce that all
+  // outputs share the same tiling. Therefore, the block-level parameters must
+  // specify identical tile sizes for all outputs.
+  // TODO(b/502910372): Support arbitrary multi-output fusions.
+  llvm::SmallVector<int64_t> parallel_tile_sizes = llvm::to_vector(
+      llvm::ArrayRef(block_level_parameters.output_tile_sizes[0]));
+  for (size_t i = 1; i < block_level_parameters.output_tile_sizes.size(); ++i) {
+    if (!absl::c_equal(parallel_tile_sizes,
+                       block_level_parameters.output_tile_sizes[i])) {
+      return Unimplemented(
+          "Same-shape multi-output fusions must have identical tile sizes "
+          "for all outputs. Received different tile sizes for root 0 [%s] "
+          "and root %d [%s].",
+          absl::StrJoin(parallel_tile_sizes, ", "), i,
+          absl::StrJoin(block_level_parameters.output_tile_sizes[i], ", "));
+    }
+  }
+  return parallel_tile_sizes;
+}
+
 }  // namespace
 
 absl::StatusOr<Tiling> TilingFromAnnotatedFusion(
@@ -109,36 +165,11 @@ absl::StatusOr<llvm::SmallVector<int64_t>> GetTilingSpaceConcreteSizes(
     const ::xla::gpu::experimental::TilingSpace& tiling_space,
     const BlockLevelParameters& block_level_parameters,
     bool enable_same_shape_multi_output_fusion) {
-  TF_RET_CHECK(!block_level_parameters.output_tile_sizes.empty())
-      << "output_tile_sizes cannot be empty.";
-
-  // For multi-output fusions with identical shapes, we enforce that all outputs
-  // share the same tiling. Therefore, the block-level parameters must specify
-  // identical tile sizes for all outputs.
-  // TODO(b/502910372): Support arbitrary multi-output fusions.
-  const auto& parallel_tile_sizes = block_level_parameters.output_tile_sizes[0];
-  if (block_level_parameters.output_tile_sizes.size() > 1) {
-    if (!enable_same_shape_multi_output_fusion) {
-      return Unimplemented(
-          "Only single-result fusions are supported for now. Received %d "
-          "roots.",
-          block_level_parameters.output_tile_sizes.size());
-    }
-
-    for (size_t i = 1; i < block_level_parameters.output_tile_sizes.size();
-         ++i) {
-      if (parallel_tile_sizes != block_level_parameters.output_tile_sizes[i]) {
-        return Unimplemented(
-            "Same-shape multi-output fusions must have identical tile sizes "
-            "for all outputs. Received different tile sizes for root 0 [%s] "
-            "and root %d [%s].",
-            absl::StrJoin(parallel_tile_sizes, ", "), i,
-            absl::StrJoin(block_level_parameters.output_tile_sizes[i], ", "));
-      }
-    }
-  }
-  if (int64_t num_parallel_dims = tiling_space.num_parallel_dimensions();
-      num_parallel_dims != parallel_tile_sizes.size()) {
+  ABSL_ASSIGN_OR_RETURN(llvm::SmallVector<int64_t> parallel_tile_sizes,
+                   GetParallelTileSizes(tiling_space, block_level_parameters,
+                                        enable_same_shape_multi_output_fusion));
+  const int64_t num_parallel_dims = tiling_space.num_parallel_dimensions();
+  if (num_parallel_dims != parallel_tile_sizes.size()) {
     return Internal(
         "Number of parallel dimensions in the tiling space (%d) does not match "
         "the number of output tile sizes in the block level fusion config "

From 0c26a96d0beca58db2e23548586862c1f58fe0aa Mon Sep 17 00:00:00 2001
From: Dmitri Latushko <latushko@google.com>
Date: Thu, 8 Oct 2026 10:39:05 -0700
Subject: [PATCH 0303/1000] [tf.data] Call gc.collect() before RSS sampling in
 memory_cleanup_test and add notsan tag

Call `gc.collect()` before recording initial and final whole-process RSS via `memory_profiler.memory_usage(-1)[0]` in `MemoryCleanupTest.assertMemoryNotIncreasing` (matching `memory_test_util.assert_no_leak`), and add `"notsan"` to `tags` and `xla_tags` on `memory_cleanup_test` since whole-process RSS thresholds are incompatible with ThreadSanitizer shadow memory and quarantine overhead.

PiperOrigin-RevId: 995930185
---
 tensorflow/python/data/kernel_tests/BUILD                  | 2 ++
 tensorflow/python/data/kernel_tests/memory_cleanup_test.py | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/tensorflow/python/data/kernel_tests/BUILD b/tensorflow/python/data/kernel_tests/BUILD
index 33b59117e18..ffc0aea0514 100644
--- a/tensorflow/python/data/kernel_tests/BUILD
+++ b/tensorflow/python/data/kernel_tests/BUILD
@@ -778,9 +778,11 @@ cuda_py_strict_test(
     shard_count = 24,
     tags = [
         "no_cuda_asan",  # TODO(b/183963686)
+        "notsan",  # Whole-process RSS profiling via memory_profiler is incompatible with TSAN shadow memory/quarantine overhead.
     ],
     xla_tags = [
         "no_cuda_asan",  # TODO(b/183963686)
+        "notsan",  # Whole-process RSS profiling via memory_profiler is incompatible with TSAN shadow memory/quarantine overhead.
     ],
     deps = [
         ":test_base",
diff --git a/tensorflow/python/data/kernel_tests/memory_cleanup_test.py b/tensorflow/python/data/kernel_tests/memory_cleanup_test.py
index afd5e18eecd..61c15b841c8 100644
--- a/tensorflow/python/data/kernel_tests/memory_cleanup_test.py
+++ b/tensorflow/python/data/kernel_tests/memory_cleanup_test.py
@@ -50,9 +50,11 @@ class MemoryCleanupTest(test_base.DatasetTestBase, parameterized.TestCase):
     f()
     # Wait for background threads to start up and allocate memory.
     time.sleep(4)
+    gc.collect()
     initial = memory_profiler.memory_usage(-1)[0]
     for _ in range(num_iters):
       f()
+    gc.collect()
     increase = memory_profiler.memory_usage(-1)[0] - initial
     logging.info("Memory increase observed: %f MB" % increase)
     assert increase < max_increase_mb, (

From 1e596dc9cc2c0ae0b70b09faeac38f56dc57b7d1 Mon Sep 17 00:00:00 2001
From: nishad shabbir <nishadshabbir28@gmail.com>
Date: Thu, 8 Oct 2026 23:35:29 +0530
Subject: [PATCH 0304/1000] Address review: widen the exception regex and add a
 tensor input case

- The rank test now matches both wordings with-rank-at-least can
  produce, since "rank at least 3" alone does not match "Shape must be
  at least rank 3".
- Add a tf.Tensor on one side of the float loop.

Not included: serializing the asserts in _verify_compatible_image_shapes
behind a control dependency. The failure it is meant to fix does not
reproduce here. The suggested test passes on master unchanged, and
across six dynamic rank combinations on master every one fails through
the rank assertion rather than with an incompatible-shapes error, so
there is nothing to pin. That function is shared with ssim and
ssim_multiscale, and an unverifiable change to it does not belong in a
psnr input-conversion fix.
---
 tensorflow/python/ops/image_ops_test.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/ops/image_ops_test.py b/tensorflow/python/ops/image_ops_test.py
index 8fb6cc5b987..703422aca3a 100644
--- a/tensorflow/python/ops/image_ops_test.py
+++ b/tensorflow/python/ops/image_ops_test.py
@@ -6145,9 +6145,11 @@ class PSNRTest(test_util.TensorFlowTestCase):
     expected = self._PSNR_NumPy(image1, image2, 1)
 
     with self.cached_session():
+      tf_image2 = constant_op.constant(image2, dtype=dtypes.float32)
       for a, b in ((image1.tolist(), image2.tolist()),
                    (tuple(image1.tolist()), tuple(image2.tolist())),
                    (image1.tolist(), image2.astype(np.float32)),
+                   (image1.tolist(), tf_image2),
                    (image1.astype(np.float32), image2.tolist()),
                    (image1.astype(np.float32), image2.astype(np.float32))):
         tf_psnr = self.evaluate(image_ops.psnr(a, b, 1.0, "psnr"))
@@ -6171,7 +6173,8 @@ class PSNRTest(test_util.TensorFlowTestCase):
     # The shape check runs before the reduction, so a rank below 3 is
     # reported as a rank problem rather than surfacing from reduce_mean.
     for bad in ([1.0, 2.0], [[1.0, 2.0], [3.0, 4.0]]):
-      with self.assertRaisesRegex(ValueError, "must have rank at least 3"):
+      with self.assertRaisesRegex(
+          ValueError, "rank at least 3|at least rank 3"):
         image_ops.psnr(bad, bad, 1.0)
 
   def testPSNRMultiImage(self):

From 897443b8fa9cf19e8a1e258314a44131a34aae99 Mon Sep 17 00:00:00 2001
From: Sai Ramineni <csramineni@gmail.com>
Date: Thu, 13 Aug 2026 05:23:15 -0400
Subject: [PATCH 0305/1000] Validate tensor scatter shapes

---
 tensorflow/core/kernels/scatter_nd_op.cc      | 37 +------------------
 tensorflow/core/kernels/scatter_nd_op_test.cc | 15 ++++++++
 2 files changed, 17 insertions(+), 35 deletions(-)

diff --git a/tensorflow/core/kernels/scatter_nd_op.cc b/tensorflow/core/kernels/scatter_nd_op.cc
index 49d8236dcd0..2c19997ef96 100644
--- a/tensorflow/core/kernels/scatter_nd_op.cc
+++ b/tensorflow/core/kernels/scatter_nd_op.cc
@@ -222,41 +222,8 @@ class TensorScatterOp : public ScatterOpBase<Device> {
                 absl::InvalidArgumentError(
                     "Indices and updates specified for empty output shape"));
 
-    const int64_t outer_dims = indices.shape().dims() - 1;
-
-    OP_REQUIRES(c, updates.shape().dims() >= outer_dims,
-                absl::InvalidArgumentError(absl::StrCat(
-                    "Updates shape must have rank at least the number of "
-                    "outer dimensions of indices (",
-                    outer_dims,
-                    "). Found: updates shape=", updates.shape().DebugString(),
-                    ", indices shape=", indices.shape().DebugString())));
-
-    for (int i = 0; i < outer_dims; ++i) {
-      OP_REQUIRES(c, indices.shape().dim_size(i) == updates.shape().dim_size(i),
-                  absl::InvalidArgumentError(absl::StrCat(
-                      "Outer dimensions of indices and update must match. "
-                      "Indices shape: ",
-                      indices.shape().DebugString(),
-                      ", updates shape:", updates.shape().DebugString())));
-    }
-
-    const int64_t ix = indices.shape().dim_size(outer_dims);
-    OP_REQUIRES(
-        c, updates.shape().dims() - outer_dims == shape.dims() - ix,
-        absl::InvalidArgumentError(absl::StrCat(
-            "Inner dimensions of output shape must match "
-            "inner dimensions of updates shape. Output: ",
-            shape.DebugString(), " updates: ", updates.shape().DebugString())));
-    for (int i = 0; i + outer_dims < updates.shape().dims(); ++i) {
-      OP_REQUIRES(
-          c, updates.shape().dim_size(i + outer_dims) == shape.dim_size(ix + i),
-          absl::InvalidArgumentError(absl::StrCat(
-              "The inner ", shape.dims() - ix,
-              " dimensions of output.shape=", shape.DebugString(),
-              " must match the inner ", updates.shape().dims() - outer_dims,
-              " dimensions of updates.shape=", updates.shape().DebugString())));
-    }
+    OP_REQUIRES_OK(c, ValidateScatterNdUpdateShape(
+                          shape, indices.shape(), updates.shape()));
 
     AllocatorAttributes alloc_attr;
     MemoryType memory_type = DEVICE_MEMORY;
diff --git a/tensorflow/core/kernels/scatter_nd_op_test.cc b/tensorflow/core/kernels/scatter_nd_op_test.cc
index 8908134e442..cfcd7abb3ca 100644
--- a/tensorflow/core/kernels/scatter_nd_op_test.cc
+++ b/tensorflow/core/kernels/scatter_nd_op_test.cc
@@ -137,6 +137,21 @@ TEST_F(TensorScatterUpdateOpTest, HigherRank) {
   test::ExpectTensorEqual<float>(expected, *GetOutput(0));
 }
 
+TEST_F(TensorScatterUpdateOpTest, Error_MismatchedOuterDimensions) {
+  MakeOp(DT_FLOAT, DT_INT32);
+
+  AddInputFromArray<float>(TensorShape({5}), {0, 0, 0, 0, 0});
+  AddInputFromArray<int32_t>(TensorShape({5, 1, 1}), {0, 1, 2, 3, 4});
+  AddInputFromArray<float>(TensorShape({5}), {1, 2, 3, 4, 5});
+
+  absl::Status status = RunOpKernel();
+  EXPECT_TRUE(absl::StrContains(
+      status.ToString(),
+      "Dimensions [0,2) of indices[shape=[5,1,1]] must match dimensions "
+      "[0,2) of updates[shape=[5]]"))
+      << status;
+}
+
 TEST_F(TensorScatterUpdateOpTest, Error_IndexOutOfRange) {
   MakeOp(DT_FLOAT, DT_INT32);
 

From dc70e518f844085fd371322d1e71ad2ac0ce7536 Mon Sep 17 00:00:00 2001
From: Nicholas Levin <nglevin@google.com>
Date: Thu, 8 Oct 2026 10:46:31 -0700
Subject: [PATCH 0306/1000] Automated Code Change

PiperOrigin-RevId: 995934872
---
 tensorflow/lite/objc/BUILD.apple              |  1 -
 .../AppIcon.appiconset/Contents.json          | 98 -------------------
 .../TestApp/Assets.xcassets/Contents.json     |  6 --
 .../AppIcon.appiconset/Contents.json          | 98 -------------------
 .../Assets.xcassets/Contents.json             |  6 --
 5 files changed, 209 deletions(-)
 delete mode 100644 tensorflow/lite/objc/apps/TestApp/TestApp/Assets.xcassets/AppIcon.appiconset/Contents.json
 delete mode 100644 tensorflow/lite/objc/apps/TestApp/TestApp/Assets.xcassets/Contents.json
 delete mode 100644 tensorflow/lite/tools/benchmark/ios/TFLiteBenchmark/TFLiteBenchmark/Assets.xcassets/AppIcon.appiconset/Contents.json
 delete mode 100644 tensorflow/lite/tools/benchmark/ios/TFLiteBenchmark/TFLiteBenchmark/Assets.xcassets/Contents.json

diff --git a/tensorflow/lite/objc/BUILD.apple b/tensorflow/lite/objc/BUILD.apple
index bc640c111fa..f34907fd3db 100644
--- a/tensorflow/lite/objc/BUILD.apple
+++ b/tensorflow/lite/objc/BUILD.apple
@@ -147,7 +147,6 @@ objc_library(
 
 ios_application(
     name = "TestApp",
-    app_icons = glob(["apps/TestApp/TestApp/Assets.xcassets/AppIcon.appiconset/**"]),
     bundle_id = "com.tensorflow.lite.objc.TestApp",
     families = [
         "ipad",
diff --git a/tensorflow/lite/objc/apps/TestApp/TestApp/Assets.xcassets/AppIcon.appiconset/Contents.json b/tensorflow/lite/objc/apps/TestApp/TestApp/Assets.xcassets/AppIcon.appiconset/Contents.json
deleted file mode 100644
index d8db8d65fd7..00000000000
--- a/tensorflow/lite/objc/apps/TestApp/TestApp/Assets.xcassets/AppIcon.appiconset/Contents.json
+++ /dev/null
@@ -1,98 +0,0 @@
-{
-  "images" : [
-    {
-      "idiom" : "iphone",
-      "size" : "20x20",
-      "scale" : "2x"
-    },
-    {
-      "idiom" : "iphone",
-      "size" : "20x20",
-      "scale" : "3x"
-    },
-    {
-      "idiom" : "iphone",
-      "size" : "29x29",
-      "scale" : "2x"
-    },
-    {
-      "idiom" : "iphone",
-      "size" : "29x29",
-      "scale" : "3x"
-    },
-    {
-      "idiom" : "iphone",
-      "size" : "40x40",
-      "scale" : "2x"
-    },
-    {
-      "idiom" : "iphone",
-      "size" : "40x40",
-      "scale" : "3x"
-    },
-    {
-      "idiom" : "iphone",
-      "size" : "60x60",
-      "scale" : "2x"
-    },
-    {
-      "idiom" : "iphone",
-      "size" : "60x60",
-      "scale" : "3x"
-    },
-    {
-      "idiom" : "ipad",
-      "size" : "20x20",
-      "scale" : "1x"
-    },
-    {
-      "idiom" : "ipad",
-      "size" : "20x20",
-      "scale" : "2x"
-    },
-    {
-      "idiom" : "ipad",
-      "size" : "29x29",
-      "scale" : "1x"
-    },
-    {
-      "idiom" : "ipad",
-      "size" : "29x29",
-      "scale" : "2x"
-    },
-    {
-      "idiom" : "ipad",
-      "size" : "40x40",
-      "scale" : "1x"
-    },
-    {
-      "idiom" : "ipad",
-      "size" : "40x40",
-      "scale" : "2x"
-    },
-    {
-      "idiom" : "ipad",
-      "size" : "76x76",
-      "scale" : "1x"
-    },
-    {
-      "idiom" : "ipad",
-      "size" : "76x76",
-      "scale" : "2x"
-    },
-    {
-      "idiom" : "ipad",
-      "size" : "83.5x83.5",
-      "scale" : "2x"
-    },
-    {
-      "idiom" : "ios-marketing",
-      "size" : "1024x1024",
-      "scale" : "1x"
-    }
-  ],
-  "info" : {
-    "version" : 1,
-    "author" : "xcode"
-  }
-}
\ No newline at end of file
diff --git a/tensorflow/lite/objc/apps/TestApp/TestApp/Assets.xcassets/Contents.json b/tensorflow/lite/objc/apps/TestApp/TestApp/Assets.xcassets/Contents.json
deleted file mode 100644
index da4a164c918..00000000000
--- a/tensorflow/lite/objc/apps/TestApp/TestApp/Assets.xcassets/Contents.json
+++ /dev/null
@@ -1,6 +0,0 @@
-{
-  "info" : {
-    "version" : 1,
-    "author" : "xcode"
-  }
-}
\ No newline at end of file
diff --git a/tensorflow/lite/tools/benchmark/ios/TFLiteBenchmark/TFLiteBenchmark/Assets.xcassets/AppIcon.appiconset/Contents.json b/tensorflow/lite/tools/benchmark/ios/TFLiteBenchmark/TFLiteBenchmark/Assets.xcassets/AppIcon.appiconset/Contents.json
deleted file mode 100644
index d8db8d65fd7..00000000000
--- a/tensorflow/lite/tools/benchmark/ios/TFLiteBenchmark/TFLiteBenchmark/Assets.xcassets/AppIcon.appiconset/Contents.json
+++ /dev/null
@@ -1,98 +0,0 @@
-{
-  "images" : [
-    {
-      "idiom" : "iphone",
-      "size" : "20x20",
-      "scale" : "2x"
-    },
-    {
-      "idiom" : "iphone",
-      "size" : "20x20",
-      "scale" : "3x"
-    },
-    {
-      "idiom" : "iphone",
-      "size" : "29x29",
-      "scale" : "2x"
-    },
-    {
-      "idiom" : "iphone",
-      "size" : "29x29",
-      "scale" : "3x"
-    },
-    {
-      "idiom" : "iphone",
-      "size" : "40x40",
-      "scale" : "2x"
-    },
-    {
-      "idiom" : "iphone",
-      "size" : "40x40",
-      "scale" : "3x"
-    },
-    {
-      "idiom" : "iphone",
-      "size" : "60x60",
-      "scale" : "2x"
-    },
-    {
-      "idiom" : "iphone",
-      "size" : "60x60",
-      "scale" : "3x"
-    },
-    {
-      "idiom" : "ipad",
-      "size" : "20x20",
-      "scale" : "1x"
-    },
-    {
-      "idiom" : "ipad",
-      "size" : "20x20",
-      "scale" : "2x"
-    },
-    {
-      "idiom" : "ipad",
-      "size" : "29x29",
-      "scale" : "1x"
-    },
-    {
-      "idiom" : "ipad",
-      "size" : "29x29",
-      "scale" : "2x"
-    },
-    {
-      "idiom" : "ipad",
-      "size" : "40x40",
-      "scale" : "1x"
-    },
-    {
-      "idiom" : "ipad",
-      "size" : "40x40",
-      "scale" : "2x"
-    },
-    {
-      "idiom" : "ipad",
-      "size" : "76x76",
-      "scale" : "1x"
-    },
-    {
-      "idiom" : "ipad",
-      "size" : "76x76",
-      "scale" : "2x"
-    },
-    {
-      "idiom" : "ipad",
-      "size" : "83.5x83.5",
-      "scale" : "2x"
-    },
-    {
-      "idiom" : "ios-marketing",
-      "size" : "1024x1024",
-      "scale" : "1x"
-    }
-  ],
-  "info" : {
-    "version" : 1,
-    "author" : "xcode"
-  }
-}
\ No newline at end of file
diff --git a/tensorflow/lite/tools/benchmark/ios/TFLiteBenchmark/TFLiteBenchmark/Assets.xcassets/Contents.json b/tensorflow/lite/tools/benchmark/ios/TFLiteBenchmark/TFLiteBenchmark/Assets.xcassets/Contents.json
deleted file mode 100644
index da4a164c918..00000000000
--- a/tensorflow/lite/tools/benchmark/ios/TFLiteBenchmark/TFLiteBenchmark/Assets.xcassets/Contents.json
+++ /dev/null
@@ -1,6 +0,0 @@
-{
-  "info" : {
-    "version" : 1,
-    "author" : "xcode"
-  }
-}
\ No newline at end of file

From f7d4671d8a129e80c117f2265e2f7b0797c86e1a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 8 Oct 2026 11:01:08 -0700
Subject: [PATCH 0307/1000] Use cost-model GEMM tiling for per-cluster
 `_XlaDeterministic` JIT compilation.

PiperOrigin-RevId: 995943894
---
 tensorflow/compiler/jit/device_compiler_client.cc    | 12 ++++++++----
 .../compiler/jit/device_compiler_client_test.cc      |  4 ++++
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/tensorflow/compiler/jit/device_compiler_client.cc b/tensorflow/compiler/jit/device_compiler_client.cc
index e99d706166c..2c394dcaa87 100644
--- a/tensorflow/compiler/jit/device_compiler_client.cc
+++ b/tensorflow/compiler/jit/device_compiler_client.cc
@@ -38,19 +38,23 @@ xla::ExecutableBuildOptions GetExecutableBuildOptions(
   build_options.set_alias_passthrough_params(options.alias_passthrough_params);
   build_options.mutable_debug_options()->set_xla_detailed_logging(
       options.detailed_logging);
-  bool xla_deterministic = tensorflow::OpDeterminismRequired();
-  if (!xla_deterministic && result.computation != nullptr) {
+  bool cluster_deterministic = false;
+  if (result.computation != nullptr) {
     const auto& frontend_attrs =
         result.computation->proto().frontend_attributes().map();
     auto it = frontend_attrs.find(kXlaDeterministicAttr);
     if (it != frontend_attrs.end() && it->second == "true") {
-      xla_deterministic = true;
+      cluster_deterministic = true;
     }
   }
-  if (xla_deterministic) {
+  if (tensorflow::OpDeterminismRequired() || cluster_deterministic) {
     build_options.mutable_debug_options()
         ->set_xla_gpu_exclude_nondeterministic_ops(true);
   }
+  if (cluster_deterministic) {
+    build_options.mutable_debug_options()
+        ->set_xla_gpu_experimental_cost_model_gemm_tiling_default(true);
+  }
   return build_options;
 }
 
diff --git a/tensorflow/compiler/jit/device_compiler_client_test.cc b/tensorflow/compiler/jit/device_compiler_client_test.cc
index 5bbc30d2456..73a117e2664 100644
--- a/tensorflow/compiler/jit/device_compiler_client_test.cc
+++ b/tensorflow/compiler/jit/device_compiler_client_test.cc
@@ -86,6 +86,8 @@ TEST(GetExecutableOptionTest, XlaDeterministicFrontendAttribute) {
       GetExecutableBuildOptions(options, result, /*default_device_ordinal=*/-1);
   EXPECT_FALSE(default_build_option.debug_options()
                    .xla_gpu_exclude_nondeterministic_ops());
+  EXPECT_FALSE(default_build_option.debug_options()
+                   .xla_gpu_experimental_cost_model_gemm_tiling_default());
 
   (*result.computation->mutable_proto()
         ->mutable_frontend_attributes()
@@ -94,6 +96,8 @@ TEST(GetExecutableOptionTest, XlaDeterministicFrontendAttribute) {
       GetExecutableBuildOptions(options, result, /*default_device_ordinal=*/-1);
   EXPECT_TRUE(deterministic_build_option.debug_options()
                   .xla_gpu_exclude_nondeterministic_ops());
+  EXPECT_TRUE(deterministic_build_option.debug_options()
+                  .xla_gpu_experimental_cost_model_gemm_tiling_default());
 }
 
 }  // namespace

From ae69e340b390ec084de71fb67c07e5de9c84d30c Mon Sep 17 00:00:00 2001
From: abhirupone <abhirup.karmakar00@gmail.com>
Date: Fri, 9 Oct 2026 00:15:02 +0530
Subject: [PATCH 0308/1000] Address review feedback on RowPartition static
 validation and test coverage

- Add static validation for nvals in from_row_starts and nrows in from_value_rowids
- Use RowPartition instead of RaggedTensor in from_value_rowids error messages
- Revert unrelated pylint directives in ragged_tensor.py
- Add _static_nvals_or_constant tests for value_rowids, row_starts, and row_limits
- Add UnderJit static validation tests for all RowPartition and RaggedTensor factories
---
 tensorflow/python/ops/ragged/ragged_tensor.py |  2 -
 .../python/ops/ragged/ragged_tensor_test.py   | 30 ++++++++
 tensorflow/python/ops/ragged/row_partition.py | 34 +++++---
 .../python/ops/ragged/row_partition_test.py   | 77 +++++++++++++++++++
 4 files changed, 130 insertions(+), 13 deletions(-)

diff --git a/tensorflow/python/ops/ragged/ragged_tensor.py b/tensorflow/python/ops/ragged/ragged_tensor.py
index cdd1901eda5..32ddb3586e6 100644
--- a/tensorflow/python/ops/ragged/ragged_tensor.py
+++ b/tensorflow/python/ops/ragged/ragged_tensor.py
@@ -2200,7 +2200,6 @@ class RaggedTensor(
 
     return stub
 
-  # pylint: disable=too-many-function-args
   __getitem__ = _overloaded_operator("__getitem__")
   __ge__ = _overloaded_operator("__ge__")
   __gt__ = _overloaded_operator("__gt__")
@@ -2232,7 +2231,6 @@ class RaggedTensor(
   __truediv__ = _overloaded_operator("__truediv__")
   __rtruediv__ = _overloaded_operator("__rtruediv__")
   del _overloaded_operator
-  # pylint: enable=too-many-function-args
 
   # =============================================================================
   # Name Scope
diff --git a/tensorflow/python/ops/ragged/ragged_tensor_test.py b/tensorflow/python/ops/ragged/ragged_tensor_test.py
index d42677bc01c..6d87498243e 100644
--- a/tensorflow/python/ops/ragged/ragged_tensor_test.py
+++ b/tensorflow/python/ops/ragged/ragged_tensor_test.py
@@ -340,6 +340,36 @@ class RaggedTensorTest(test_util.TensorFlowTestCase, parameterized.TestCase):
         ValueError, r"Shapes \(2,\) and \(4,\) are incompatible"):
       fn(constant_op.constant([4.0, 5.0, 6.0, 7.0]))
 
+  def testFromRowLengthsStaticShapeMismatchUnderJit(self):
+    @def_function.function(jit_compile=True)
+    def fn(x):
+      values = array_ops.reshape(x, [-1, 2])
+      return RaggedTensor.from_row_lengths(values, [2, 2])
+
+    with self.assertRaisesRegex(
+        ValueError, r"Shapes \(2,\) and \(4,\) are incompatible"):
+      fn(constant_op.constant([4.0, 5.0, 6.0, 7.0]))
+
+  def testFromRowLimitsStaticShapeMismatchUnderJit(self):
+    @def_function.function(jit_compile=True)
+    def fn(x):
+      values = array_ops.reshape(x, [-1, 2])
+      return RaggedTensor.from_row_limits(values, [2, 4])
+
+    with self.assertRaisesRegex(
+        ValueError, r"Shapes \(2,\) and \(4,\) are incompatible"):
+      fn(constant_op.constant([4.0, 5.0, 6.0, 7.0]))
+
+  def testFromValueRowidsStaticShapeMismatchUnderJit(self):
+    @def_function.function(jit_compile=True)
+    def fn(x):
+      values = array_ops.reshape(x, [-1, 2])
+      return RaggedTensor.from_value_rowids(values, [0, 0, 1, 1], nrows=2)
+
+    with self.assertRaisesRegex(
+        ValueError, r"Shapes \(2,\) and \(4,\) are incompatible"):
+      fn(constant_op.constant([4.0, 5.0, 6.0, 7.0]))
+
   def testFromRowStarts(self):
     values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
     row_starts = constant_op.constant([0, 2, 2, 5, 6], dtypes.int64)
diff --git a/tensorflow/python/ops/ragged/row_partition.py b/tensorflow/python/ops/ragged/row_partition.py
index d45b6f2f11b..c2e90df37df 100644
--- a/tensorflow/python/ops/ragged/row_partition.py
+++ b/tensorflow/python/ops/ragged/row_partition.py
@@ -263,37 +263,37 @@ class RowPartition(composite_tensor.CompositeTensor):
       else:
         nrows = ops.convert_to_tensor(nrows, value_rowids.dtype, "nrows")
         const_nrows = tensor_util.constant_value(nrows)
-        if const_nrows is not None:
-          if const_nrows < 0:
-            raise ValueError("Expected nrows >= 0; got %d" % const_nrows)
-          const_rowids = tensor_util.constant_value(value_rowids)
-          if const_rowids is not None and const_rowids.size > 0:
-            if not const_nrows >= const_rowids[-1] + 1:
-              raise ValueError(
-                  "Expected nrows >= value_rowids[-1] + 1; got nrows=%d, "
-                  "value_rowids[-1]=%d" % (const_nrows, const_rowids[-1]))
+        if const_nrows is not None and const_nrows < 0:
+          raise ValueError("Expected nrows >= 0; got %d" % const_nrows)
 
       value_rowids.shape.assert_has_rank(1)
       nrows.shape.assert_has_rank(0)
 
       if validate:
         value_rowids_const = tensor_util.constant_value(value_rowids)
+        nrows_const = tensor_util.constant_value(nrows)
         if value_rowids_const is not None:
           if len(value_rowids_const) > 0 and value_rowids_const[0] < 0:
             raise ValueError(
                 "Arguments to from_value_rowids do not form a valid "
-                "RaggedTensor: value_rowids must be non-negative, "
+                "RowPartition: value_rowids must be non-negative, "
                 f"got {value_rowids_const[0]}")
           diff = np.diff(value_rowids_const)
           if np.any(diff < 0):
             bad_idx = int(np.argmax(diff < 0))
             raise ValueError(
                 "Arguments to from_value_rowids do not form a valid "
-                "RaggedTensor: value_rowids must be monotonic increasing, "
+                "RowPartition: value_rowids must be monotonic increasing, "
                 f"but value_rowids[{bad_idx}] "
                 f"({value_rowids_const[bad_idx]}) > "
                 f"value_rowids[{bad_idx + 1}] "
                 f"({value_rowids_const[bad_idx + 1]})")
+          if nrows_const is not None and len(value_rowids_const) > 0:
+            if value_rowids_const[-1] >= nrows_const:
+              raise ValueError(
+                  "Arguments to from_value_rowids do not form a valid "
+                  "RowPartition: value_rowids[-1] must be < nrows, "
+                  f"got {value_rowids_const[-1]} >= {nrows_const}")
         msg = ("Arguments to from_value_rowids do not form a valid "
                "RowPartition")
         checks = [
@@ -510,6 +510,7 @@ class RowPartition(composite_tensor.CompositeTensor):
       nvals = math_ops.cast(nvals, row_starts.dtype)
       if validate:
         row_starts_const = tensor_util.constant_value(row_starts)
+        nvals_const = tensor_util.constant_value(nvals)
         if row_starts_const is not None:
           if len(row_starts_const) > 0 and row_starts_const[0] != 0:
             raise ValueError(
@@ -525,6 +526,17 @@ class RowPartition(composite_tensor.CompositeTensor):
                 f"but row_starts[{bad_idx}] ({row_starts_const[bad_idx]}) > "
                 f"row_starts[{bad_idx + 1}] "
                 f"({row_starts_const[bad_idx + 1]})")
+          if nvals_const is not None:
+            if len(row_starts_const) > 0 and row_starts_const[-1] > nvals_const:
+              raise ValueError(
+                  "Arguments to from_row_starts do not form a valid "
+                  "RaggedTensor: row_starts[-1] must be <= nvals, "
+                  f"got {row_starts_const[-1]} > {nvals_const}")
+            if len(row_starts_const) == 0 and nvals_const != 0:
+              raise ValueError(
+                  "Arguments to from_row_starts do not form a valid "
+                  "RaggedTensor: nvals must be 0 when row_starts is empty, "
+                  f"got {nvals_const}")
         msg = "Arguments to from_row_starts do not form a valid RaggedTensor"
         checks = [
             check_ops.assert_rank(row_starts, 1, message=msg),
diff --git a/tensorflow/python/ops/ragged/row_partition_test.py b/tensorflow/python/ops/ragged/row_partition_test.py
index 07f38b04914..a40a0c8bd49 100644
--- a/tensorflow/python/ops/ragged/row_partition_test.py
+++ b/tensorflow/python/ops/ragged/row_partition_test.py
@@ -924,6 +924,18 @@ class RowPartitionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
     rp = RowPartition.from_uniform_row_length(3, nrows=4)
     self.assertEqual(12, rp._static_nvals_or_constant())
 
+  def testStaticNvalsFromValueRowids(self):
+    rp = RowPartition.from_value_rowids([0, 0, 1, 2, 2])
+    self.assertEqual(5, rp._static_nvals_or_constant())
+
+  def testStaticNvalsFromRowStarts(self):
+    rp = RowPartition.from_row_starts([0, 2, 3], nvals=5)
+    self.assertEqual(5, rp._static_nvals_or_constant())
+
+  def testStaticNvalsFromRowLimits(self):
+    rp = RowPartition.from_row_limits([2, 3, 5])
+    self.assertEqual(5, rp._static_nvals_or_constant())
+
   def testFromRowSplitsStaticValidationNonZero(self):
     with self.assertRaisesRegex(ValueError, r"row_splits\[0\] must be zero"):
       RowPartition.from_row_splits([1, 2, 4])
@@ -942,6 +954,11 @@ class RowPartitionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
         ValueError, r"value_rowids must be monotonic increasing"):
       RowPartition.from_value_rowids([0, 2, 1])
 
+  def testFromValueRowidsStaticValidationNrowsExceeded(self):
+    with self.assertRaisesRegex(
+        ValueError, r"value_rowids\[-1\] must be < nrows"):
+      RowPartition.from_value_rowids([0, 2], nrows=2)
+
   def testFromRowLengthsStaticValidationNegative(self):
     with self.assertRaisesRegex(
         ValueError, r"row_lengths must be nonnegative"):
@@ -956,6 +973,16 @@ class RowPartitionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
         ValueError, r"row_starts must be monotonic increasing"):
       RowPartition.from_row_starts([0, 3, 2], nvals=5)
 
+  def testFromRowStartsStaticValidationNvalsExceeded(self):
+    with self.assertRaisesRegex(
+        ValueError, r"row_starts\[-1\] must be <= nvals"):
+      RowPartition.from_row_starts([0, 3], nvals=2)
+
+  def testFromRowStartsStaticValidationEmptyWithNvals(self):
+    with self.assertRaisesRegex(
+        ValueError, r"nvals must be 0 when row_starts is empty"):
+      RowPartition.from_row_starts([], nvals=2)
+
   def testFromRowLimitsStaticValidationNegative(self):
     with self.assertRaisesRegex(ValueError, r"row_limits must be non-negative"):
       RowPartition.from_row_limits([-1, 2])
@@ -965,6 +992,56 @@ class RowPartitionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
         ValueError, r"row_limits must be monotonic increasing"):
       RowPartition.from_row_limits([3, 1])
 
+  def testFromRowSplitsStaticValidationUnderJit(self):
+    @def_function.function(jit_compile=True)
+    def fn(x):
+      del x
+      return RowPartition.from_row_splits([0, 5, 2])
+
+    with self.assertRaisesRegex(
+        ValueError, r"must be monotonic increasing"):
+      fn(constant_op.constant(0))
+
+  def testFromValueRowidsStaticValidationUnderJit(self):
+    @def_function.function(jit_compile=True)
+    def fn(x):
+      del x
+      return RowPartition.from_value_rowids([0, 2], nrows=2)
+
+    with self.assertRaisesRegex(
+        ValueError, r"value_rowids\[-1\] must be < nrows"):
+      fn(constant_op.constant(0))
+
+  def testFromRowLengthsStaticValidationUnderJit(self):
+    @def_function.function(jit_compile=True)
+    def fn(x):
+      del x
+      return RowPartition.from_row_lengths([3, -1, 2])
+
+    with self.assertRaisesRegex(
+        ValueError, r"row_lengths must be nonnegative"):
+      fn(constant_op.constant(0))
+
+  def testFromRowStartsStaticValidationUnderJit(self):
+    @def_function.function(jit_compile=True)
+    def fn(x):
+      del x
+      return RowPartition.from_row_starts([0, 3], nvals=2)
+
+    with self.assertRaisesRegex(
+        ValueError, r"row_starts\[-1\] must be <= nvals"):
+      fn(constant_op.constant(0))
+
+  def testFromRowLimitsStaticValidationUnderJit(self):
+    @def_function.function(jit_compile=True)
+    def fn(x):
+      del x
+      return RowPartition.from_row_limits([3, 1])
+
+    with self.assertRaisesRegex(
+        ValueError, r"row_limits must be monotonic increasing"):
+      fn(constant_op.constant(0))
+
 
 @test_util.run_all_in_graph_and_eager_modes
 class RowPartitionSpecTest(test_util.TensorFlowTestCase,

From 7d6e40d58797a875bf581f0da48586b76aa7312b Mon Sep 17 00:00:00 2001
From: Joshua Lang <joshualang@google.com>
Date: Thu, 8 Oct 2026 11:12:38 -0700
Subject: [PATCH 0309/1000] Remove unused command buffer telemetry listener and
 topology mapper

PiperOrigin-RevId: 995951732
---
 .../xla/xla/backends/profiler/gpu/BUILD       |  25 --
 .../gpu/cuda_graph_topology_mapper.cc         | 154 --------
 .../profiler/gpu/cuda_graph_topology_mapper.h |  98 -----
 .../gpu/cuda_graph_topology_mapper_test.cc    | 338 ------------------
 third_party/xla/xla/stream_executor/gpu/BUILD |  69 +---
 .../gpu/gpu_command_buffer_listener.cc        |  60 ----
 .../gpu/gpu_command_buffer_listener.h         | 155 --------
 .../gpu/gpu_command_buffer_listener_test.cc   | 108 ------
 .../gpu_command_buffer_listener_test_helper.h |  38 --
 9 files changed, 19 insertions(+), 1026 deletions(-)
 delete mode 100644 third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.cc
 delete mode 100644 third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.h
 delete mode 100644 third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper_test.cc
 delete mode 100644 third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_listener.cc
 delete mode 100644 third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_listener.h
 delete mode 100644 third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_listener_test.cc
 delete mode 100644 third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_listener_test_helper.h

diff --git a/third_party/xla/xla/backends/profiler/gpu/BUILD b/third_party/xla/xla/backends/profiler/gpu/BUILD
index 03baef28e3d..15771184662 100644
--- a/third_party/xla/xla/backends/profiler/gpu/BUILD
+++ b/third_party/xla/xla/backends/profiler/gpu/BUILD
@@ -22,10 +22,6 @@ load(
     "if_google",
     "internal_visibility",
 )
-load(
-    "//xla/tsl:tsl.default.bzl",
-    "get_compatible_with_portable",
-)
 load("//xla/tsl/platform:rules_cc.bzl", "cc_library")
 load("//xla/tsl/platform/default:cuda_build_defs.bzl", "if_cuda_is_configured", "if_cuda_newer_than")
 
@@ -1025,24 +1021,3 @@ cc_library(
         "@local_config_cuda//cuda:cuda_headers",
     ],
 )
-
-cc_library(
-    name = "cuda_graph_topology_mapper",
-    srcs = ["cuda_graph_topology_mapper.cc"],
-    hdrs = ["cuda_graph_topology_mapper.h"],
-    compatible_with = get_compatible_with_portable(),
-    deps = [
-        "//xla/tsl/platform:logging",
-        "@com_google_absl//absl/container:flat_hash_map",
-    ],
-)
-
-xla_cc_test(
-    name = "cuda_graph_topology_mapper_test",
-    srcs = ["cuda_graph_topology_mapper_test.cc"],
-    deps = [
-        ":cuda_graph_topology_mapper",
-        "@com_google_absl//absl/container:flat_hash_map",
-        "@com_google_googletest//:gtest_main",
-    ],
-)
diff --git a/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.cc b/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.cc
deleted file mode 100644
index 5f5628e0b5d..00000000000
--- a/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.cc
+++ /dev/null
@@ -1,154 +0,0 @@
-/* Copyright 2026 The OpenXLA Authors.
-
-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.
-==============================================================================*/
-
-#include "xla/backends/profiler/gpu/cuda_graph_topology_mapper.h"
-
-#include <cstddef>
-#include <cstdint>
-#include <utility>
-#include <vector>
-
-#include "absl/container/flat_hash_map.h"
-#include "xla/tsl/platform/logging.h"
-
-namespace xla {
-namespace profiler {
-
-size_t CudaGraphTopologyMapper::CalculateMergedSize(
-    uint32_t graph_id, const absl::flat_hash_map<uint32_t, size_t>& base_sizes,
-    const absl::flat_hash_map<uint32_t, std::vector<ChildGraphEntry>>&
-        child_graphs,
-    absl::flat_hash_map<uint32_t, size_t>* merged_sizes) {
-  DCHECK(merged_sizes != nullptr);
-  auto merged_it = merged_sizes->find(graph_id);
-  if (merged_it != merged_sizes->end()) {
-    return merged_it->second;
-  }
-
-  size_t size = 0;
-  auto size_it = base_sizes.find(graph_id);
-  if (size_it != base_sizes.end()) {
-    size = size_it->second;
-  }
-
-  auto children_it = child_graphs.find(graph_id);
-  if (children_it != child_graphs.end()) {
-    for (const auto& child : children_it->second) {
-      size_t child_size = CalculateMergedSize(child.child_graph_id, base_sizes,
-                                              child_graphs, merged_sizes);
-      size += child_size;
-      if (!child.is_conditional && size > 0) {
-        size -= 1;  // Child graph node itself is replaced for inline children.
-      }
-    }
-  }
-
-  (*merged_sizes)[graph_id] = size;
-  return size;
-}
-
-std::pair<uint32_t, uint32_t> CudaGraphTopologyMapper::ResolveMergedNode(
-    uint32_t graph_id, uint32_t node_index,
-    const absl::flat_hash_map<uint32_t, size_t>& base_sizes,
-    const absl::flat_hash_map<uint32_t, std::vector<ChildGraphEntry>>&
-        child_graphs,
-    absl::flat_hash_map<uint32_t, size_t>* merged_sizes) {
-  DCHECK(merged_sizes != nullptr);
-  auto children_it = child_graphs.find(graph_id);
-  if (children_it == child_graphs.end() || children_it->second.empty()) {
-    return {graph_id, node_index};
-  }
-
-  // Calculate total inline size.
-  size_t total_inline_size = 0;
-  auto size_it = base_sizes.find(graph_id);
-  if (size_it != base_sizes.end()) {
-    total_inline_size = size_it->second;
-  }
-  for (const auto& child : children_it->second) {
-    if (!child.is_conditional) {
-      size_t child_size = CalculateMergedSize(child.child_graph_id, base_sizes,
-                                              child_graphs, merged_sizes);
-      if (child_size > 0) {
-        total_inline_size += (child_size - 1);
-      } else if (total_inline_size > 0) {
-        total_inline_size -= 1;
-      }
-    }
-  }
-
-  if (node_index < total_inline_size) {
-    // 1. Resolve inline nodes
-    int64_t inline_offset = 0;
-    for (const auto& child : children_it->second) {
-      if (child.is_conditional) {
-        continue;
-      }
-      size_t child_size = CalculateMergedSize(child.child_graph_id, base_sizes,
-                                              child_graphs, merged_sizes);
-      int64_t child_start =
-          static_cast<int64_t>(child.insertion_point) + inline_offset;
-
-      if (child_size > 0) {
-        int64_t child_end = child_start + static_cast<int64_t>(child_size) - 1;
-        if (static_cast<int64_t>(node_index) >= child_start &&
-            static_cast<int64_t>(node_index) <= child_end) {
-          uint32_t local_index = static_cast<uint32_t>(
-              static_cast<int64_t>(node_index) - child_start);
-          return ResolveMergedNode(child.child_graph_id, local_index,
-                                   base_sizes, child_graphs, merged_sizes);
-        }
-      }
-
-      if (static_cast<int64_t>(node_index) < child_start) {
-        return {graph_id,
-                static_cast<uint32_t>(static_cast<int64_t>(node_index) -
-                                      inline_offset)};
-      }
-
-      inline_offset += (static_cast<int64_t>(child_size) - 1);
-    }
-    return {graph_id, static_cast<uint32_t>(static_cast<int64_t>(node_index) -
-                                            inline_offset)};
-  }
-
-  // 2. Resolve conditional nodes
-  uint32_t conditional_offset = 0;
-  for (const auto& child : children_it->second) {
-    if (!child.is_conditional) {
-      continue;
-    }
-    size_t child_size = CalculateMergedSize(child.child_graph_id, base_sizes,
-                                            child_graphs, merged_sizes);
-    if (child_size == 0) {
-      continue;
-    }
-    uint32_t child_start = total_inline_size + conditional_offset;
-    uint32_t child_end = child_start + child_size - 1;
-
-    if (node_index >= child_start && node_index <= child_end) {
-      uint32_t local_index = node_index - child_start;
-      return ResolveMergedNode(child.child_graph_id, local_index, base_sizes,
-                               child_graphs, merged_sizes);
-    }
-
-    conditional_offset += child_size;
-  }
-
-  return {graph_id, node_index};
-}
-
-}  // namespace profiler
-}  // namespace xla
diff --git a/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.h b/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.h
deleted file mode 100644
index 6fab464bd06..00000000000
--- a/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper.h
+++ /dev/null
@@ -1,98 +0,0 @@
-/* Copyright 2026 The OpenXLA Authors.
-
-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.
-==============================================================================*/
-
-#ifndef XLA_BACKENDS_PROFILER_GPU_CUDA_GRAPH_TOPOLOGY_MAPPER_H_
-#define XLA_BACKENDS_PROFILER_GPU_CUDA_GRAPH_TOPOLOGY_MAPPER_H_
-
-#include <cstddef>
-#include <cstdint>
-#include <utility>
-#include <vector>
-
-#include "absl/container/flat_hash_map.h"
-
-namespace xla {
-namespace profiler {
-
-// CudaGraphTopologyMapper is a pure, state-less C++ helper library designed to
-// mirror the layout and node-flattening sequences that the NVIDIA CUDA driver
-// uses to assign executable flat LocalNodeId sequences to CUPTI activity traces
-// during the instantiation of hierarchical and conditional nested CUDA Graphs.
-//
-// Topology Examples:
-//
-// 1. Standard Inline Child Graph (`is_conditional = false`)
-// In this scenario, the nodes of the nested child graph completely replace the
-// placeholder node in the parent graph.
-// Layout Example:
-//   Parent Graph (Size 3): [P0, P_Child, P2]
-//   Child Graph (Size 2): [C0, C1]
-//   Expected Flattened Execution sequence: P0 -> C0 -> C1 -> P2
-//   Expected Flat indices: P0 (0), C0 (1), C1 (2), P2 (3)
-//
-// 2. Conditional Nested Graph (`is_conditional = true`)
-// In this scenario, the child graph's nodes are executed as a separate entity
-// at runtime, appended to the end of the parent graph's sequence. The
-// placeholder node in the parent retains its original index in the flattened
-// space.
-// Layout Example:
-//   Parent Graph (Size 3): [P0, P1_Cond, P2]
-//   Child Graph (Size 2): [C0, C1]
-//   Expected Flattened Sequence (Kernels only): P0 -> C0 -> C1 -> P2
-//   Expected Flat indices: P0 (0), P1_Cond (1), P2 (2), C0 (3), C1 (4)
-//   (Note that the CUPTI trace for `C0` reports local node index `3`, while the
-//    original placeholder `P2` reports index `2`).
-//
-// All APIs in this library are pure, static, and have zero CUDA SDK
-// dependencies, allowing them to run identically on the host CPU for unit
-// testing.
-class CudaGraphTopologyMapper {
- public:
-  struct ChildGraphEntry {
-    uint32_t child_graph_id;
-    uint32_t insertion_point;
-    bool is_conditional;
-
-    bool operator<(const ChildGraphEntry& other) const {
-      return insertion_point < other.insertion_point;
-    }
-  };
-
-  // Recursively calculates the total merged size of a graph, caching results in
-  // the provided mutable merged_sizes map.
-  static size_t CalculateMergedSize(
-      uint32_t graph_id,
-      const absl::flat_hash_map<uint32_t, size_t>& base_sizes,
-      const absl::flat_hash_map<uint32_t, std::vector<ChildGraphEntry>>&
-          child_graphs,
-      absl::flat_hash_map<uint32_t, size_t>* merged_sizes);
-
-  // Resolves a flat, merged node index into its target (template_graph_id,
-  // local_node_index) pair.
-  static std::pair<uint32_t, uint32_t> ResolveMergedNode(
-      uint32_t graph_id, uint32_t node_index,
-      const absl::flat_hash_map<uint32_t, size_t>& base_sizes,
-      const absl::flat_hash_map<uint32_t, std::vector<ChildGraphEntry>>&
-          child_graphs,
-      absl::flat_hash_map<uint32_t, size_t>* merged_sizes);
-
- private:
-  CudaGraphTopologyMapper() = delete;
-};
-
-}  // namespace profiler
-}  // namespace xla
-
-#endif  // XLA_BACKENDS_PROFILER_GPU_CUDA_GRAPH_TOPOLOGY_MAPPER_H_
diff --git a/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper_test.cc b/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper_test.cc
deleted file mode 100644
index 1a6b506bdcc..00000000000
--- a/third_party/xla/xla/backends/profiler/gpu/cuda_graph_topology_mapper_test.cc
+++ /dev/null
@@ -1,338 +0,0 @@
-/* Copyright 2026 The OpenXLA Authors.
-
-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.
-==============================================================================*/
-
-#include "xla/backends/profiler/gpu/cuda_graph_topology_mapper.h"
-
-#include <cstddef>
-#include <cstdint>
-#include <vector>
-
-#include <gtest/gtest.h>
-#include "absl/container/flat_hash_map.h"
-
-namespace xla {
-namespace profiler {
-namespace {
-
-TEST(CudaGraphTopologyMapperTest, CalculateMergedSizeSimpleInline) {
-  absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 3}, {2, 2}};
-  absl::flat_hash_map<uint32_t,
-                      std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
-      child_graphs;
-  child_graphs[1].push_back({2, 1, /*is_conditional=*/false});
-
-  absl::flat_hash_map<uint32_t, size_t> merged_sizes;
-  size_t size = CudaGraphTopologyMapper::CalculateMergedSize(
-      1, base_sizes, child_graphs, &merged_sizes);
-
-  // 3 + 2 - 1 = 4
-  EXPECT_EQ(size, 4);
-}
-
-TEST(CudaGraphTopologyMapperTest, ResolveMergedNodeInline) {
-  absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 3}, {2, 2}};
-  absl::flat_hash_map<uint32_t,
-                      std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
-      child_graphs;
-  child_graphs[1].push_back({2, 1, /*is_conditional=*/false});
-
-  absl::flat_hash_map<uint32_t, size_t> merged_sizes;
-
-  // Expected Merged indices:
-  // 0 -> Parent node 0
-  // 1 -> Child node 0
-  // 2 -> Child node 1
-  // 3 -> Parent node 2
-  auto res0 = CudaGraphTopologyMapper::ResolveMergedNode(
-      1, 0, base_sizes, child_graphs, &merged_sizes);
-  EXPECT_EQ(res0.first, 1);
-  EXPECT_EQ(res0.second, 0);
-
-  auto res1 = CudaGraphTopologyMapper::ResolveMergedNode(
-      1, 1, base_sizes, child_graphs, &merged_sizes);
-  EXPECT_EQ(res1.first, 2);
-  EXPECT_EQ(res1.second, 0);
-
-  auto res2 = CudaGraphTopologyMapper::ResolveMergedNode(
-      1, 2, base_sizes, child_graphs, &merged_sizes);
-  EXPECT_EQ(res2.first, 2);
-  EXPECT_EQ(res2.second, 1);
-
-  auto res3 = CudaGraphTopologyMapper::ResolveMergedNode(
-      1, 3, base_sizes, child_graphs, &merged_sizes);
-  EXPECT_EQ(res3.first, 1);
-  EXPECT_EQ(res3.second, 2);
-}
-
-TEST(CudaGraphTopologyMapperTest, ResolveMergedNodeDeepNestingInline) {
-  // Parent (1) -> Child (2) -> Grandchild (3)
-  absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 3}, {2, 3}, {3, 2}};
-  absl::flat_hash_map<uint32_t,
-                      std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
-      child_graphs;
-  child_graphs[1].push_back({2, 1, /*is_conditional=*/false});
-  child_graphs[2].push_back({3, 1, /*is_conditional=*/false});
-
-  absl::flat_hash_map<uint32_t, size_t> merged_sizes;
-
-  // Expected indices:
-  // 0 -> Parent 0
-  // 1 -> Child 0
-  // 2 -> Grandchild 0
-  // 3 -> Grandchild 1
-  // 4 -> Child 2
-  // 5 -> Parent 2
-  auto res2 = CudaGraphTopologyMapper::ResolveMergedNode(
-      1, 2, base_sizes, child_graphs, &merged_sizes);
-  EXPECT_EQ(res2.first, 3);
-  EXPECT_EQ(res2.second, 0);
-
-  auto res5 = CudaGraphTopologyMapper::ResolveMergedNode(
-      1, 5, base_sizes, child_graphs, &merged_sizes);
-  EXPECT_EQ(res5.first, 1);
-  EXPECT_EQ(res5.second, 2);
-}
-
-TEST(CudaGraphTopologyMapperTest, CalculateMergedSizeConditional) {
-  absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 3}, {2, 2}};
-  absl::flat_hash_map<uint32_t,
-                      std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
-      child_graphs;
-  child_graphs[1].push_back({2, 1, /*is_conditional=*/true});
-
-  absl::flat_hash_map<uint32_t, size_t> merged_sizes;
-  size_t size = CudaGraphTopologyMapper::CalculateMergedSize(
-      1, base_sizes, child_graphs, &merged_sizes);
-
-  // 3 + 2 = 5 (Conditional child graph nodes are appended, size is not shifted
-  // by -1)
-  EXPECT_EQ(size, 5);
-}
-
-TEST(CudaGraphTopologyMapperTest, ResolveMergedNodeConditional) {
-  absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 3}, {2, 2}};
-  absl::flat_hash_map<uint32_t,
-                      std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
-      child_graphs;
-  child_graphs[1].push_back({2, 1, /*is_conditional=*/true});
-
-  absl::flat_hash_map<uint32_t, size_t> merged_sizes;
-
-  // Expected Merged indices for Conditional Nested Graph (Hardware Ground
-  // Truth): 0 -> Parent node 0 1 -> Parent node 1 2 -> Parent node 2 3 -> Child
-  // node 0 4 -> Child node 1
-  auto res0 = CudaGraphTopologyMapper::ResolveMergedNode(
-      1, 0, base_sizes, child_graphs, &merged_sizes);
-  EXPECT_EQ(res0.first, 1);
-  EXPECT_EQ(res0.second, 0);
-
-  auto res1 = CudaGraphTopologyMapper::ResolveMergedNode(
-      1, 1, base_sizes, child_graphs, &merged_sizes);
-  EXPECT_EQ(res1.first, 1);
-  EXPECT_EQ(res1.second, 1);
-
-  auto res2 = CudaGraphTopologyMapper::ResolveMergedNode(
-      1, 2, base_sizes, child_graphs, &merged_sizes);
-  EXPECT_EQ(res2.first, 1);
-  EXPECT_EQ(res2.second, 2);
-
-  auto res3 = CudaGraphTopologyMapper::ResolveMergedNode(
-      1, 3, base_sizes, child_graphs, &merged_sizes);
-  EXPECT_EQ(res3.first, 2);
-  EXPECT_EQ(res3.second, 0);
-
-  auto res4 = CudaGraphTopologyMapper::ResolveMergedNode(
-      1, 4, base_sizes, child_graphs, &merged_sizes);
-  EXPECT_EQ(res4.first, 2);
-  EXPECT_EQ(res4.second, 1);
-}
-
-TEST(CudaGraphTopologyMapperTest, ResolveMergedNodeMixedInlineAndConditional) {
-  absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 3}, {2, 2}, {3, 2}};
-  absl::flat_hash_map<uint32_t,
-                      std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
-      child_graphs;
-  child_graphs[1].push_back({2, 1, /*is_conditional=*/false});
-  child_graphs[1].push_back({3, 2, /*is_conditional=*/true});
-
-  absl::flat_hash_map<uint32_t, size_t> merged_sizes;
-
-  auto res = CudaGraphTopologyMapper::ResolveMergedNode(
-      1, 4, base_sizes, child_graphs, &merged_sizes);
-  EXPECT_EQ(res.first, 3);
-  EXPECT_EQ(res.second, 0);
-}
-
-TEST(CudaGraphTopologyMapperTest, ResolveMergedNodeMultipleConditional) {
-  absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 3}, {2, 2}, {3, 2}};
-  absl::flat_hash_map<uint32_t,
-                      std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
-      child_graphs;
-  child_graphs[1].push_back({2, 1, /*is_conditional=*/true});
-  child_graphs[1].push_back({3, 2, /*is_conditional=*/true});
-
-  absl::flat_hash_map<uint32_t, size_t> merged_sizes;
-
-  auto res = CudaGraphTopologyMapper::ResolveMergedNode(
-      1, 5, base_sizes, child_graphs, &merged_sizes);
-  EXPECT_EQ(res.first, 3);
-  EXPECT_EQ(res.second, 0);
-}
-
-TEST(CudaGraphTopologyMapperTest, ResolveMergedNodeOutOfBounds) {
-  absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 3}, {2, 2}};
-  absl::flat_hash_map<uint32_t,
-                      std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
-      child_graphs;
-  child_graphs[1].push_back({2, 1, /*is_conditional=*/false});
-
-  absl::flat_hash_map<uint32_t, size_t> merged_sizes;
-
-  auto res = CudaGraphTopologyMapper::ResolveMergedNode(
-      1, 4, base_sizes, child_graphs, &merged_sizes);
-  EXPECT_EQ(res.first, 1);
-  EXPECT_EQ(res.second, 4);
-}
-
-TEST(CudaGraphTopologyMapperTest, EmptyChildGraphHandling) {
-  // Test 1: Inline empty child graph (child size == 0).
-  // Parent (1) size 2, Child (2) size 0.
-  // Merged size should be: 2 + 0 - 1 = 1.
-  {
-    absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 2}, {2, 0}};
-    absl::flat_hash_map<uint32_t,
-                        std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
-        child_graphs;
-    child_graphs[1].push_back({2, 1, /*is_conditional=*/false});
-
-    absl::flat_hash_map<uint32_t, size_t> merged_sizes;
-    size_t size = CudaGraphTopologyMapper::CalculateMergedSize(
-        1, base_sizes, child_graphs, &merged_sizes);
-    EXPECT_EQ(size, 1);
-
-    auto res0 = CudaGraphTopologyMapper::ResolveMergedNode(
-        1, 0, base_sizes, child_graphs, &merged_sizes);
-    EXPECT_EQ(res0.first, 1);
-    EXPECT_EQ(res0.second, 0);
-  }
-
-  // Test 2: Conditional empty child graph (child size == 0).
-  // Parent (1) size 2, Child (2) size 0.
-  // Merged size should be: 2 + 0 = 2.
-  {
-    absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 2}, {2, 0}};
-    absl::flat_hash_map<uint32_t,
-                        std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
-        child_graphs;
-    child_graphs[1].push_back({2, 1, /*is_conditional=*/true});
-
-    absl::flat_hash_map<uint32_t, size_t> merged_sizes;
-    size_t size = CudaGraphTopologyMapper::CalculateMergedSize(
-        1, base_sizes, child_graphs, &merged_sizes);
-    EXPECT_EQ(size, 2);
-
-    auto res0 = CudaGraphTopologyMapper::ResolveMergedNode(
-        1, 0, base_sizes, child_graphs, &merged_sizes);
-    EXPECT_EQ(res0.first, 1);
-    EXPECT_EQ(res0.second, 0);
-
-    auto res1 = CudaGraphTopologyMapper::ResolveMergedNode(
-        1, 1, base_sizes, child_graphs, &merged_sizes);
-    EXPECT_EQ(res1.first, 1);
-    EXPECT_EQ(res1.second, 1);
-  }
-
-  // Test 3: Completely empty parent and child graphs.
-  {
-    absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 0}, {2, 0}};
-    absl::flat_hash_map<uint32_t,
-                        std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
-        child_graphs;
-    child_graphs[1].push_back({2, 0, /*is_conditional=*/false});
-
-    absl::flat_hash_map<uint32_t, size_t> merged_sizes;
-    size_t size = CudaGraphTopologyMapper::CalculateMergedSize(
-        1, base_sizes, child_graphs, &merged_sizes);
-    EXPECT_EQ(size, 0);
-  }
-}
-
-TEST(CudaGraphTopologyMapperTest, EmptyChildGraphResolutionWithTrailingNodes) {
-  // 1. Parent size 3 (nodes P0, P1, P2) with empty inline child at insertion
-  // point 1:
-  {
-    absl::flat_hash_map<uint32_t, size_t> base_sizes = {{1, 3}, {2, 0}};
-    absl::flat_hash_map<uint32_t,
-                        std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
-        child_graphs;
-    child_graphs[1].push_back({2, 1, /*is_conditional=*/false});
-
-    absl::flat_hash_map<uint32_t, size_t> merged_sizes;
-    size_t size = CudaGraphTopologyMapper::CalculateMergedSize(
-        1, base_sizes, child_graphs, &merged_sizes);
-    EXPECT_EQ(size, 2);
-
-    auto res0 = CudaGraphTopologyMapper::ResolveMergedNode(
-        1, 0, base_sizes, child_graphs, &merged_sizes);
-    EXPECT_EQ(res0.first, 1);
-    EXPECT_EQ(res0.second, 0);
-
-    auto res1 = CudaGraphTopologyMapper::ResolveMergedNode(
-        1, 1, base_sizes, child_graphs, &merged_sizes);
-    EXPECT_EQ(res1.first, 1);
-    EXPECT_EQ(res1.second, 2);
-  }
-
-  // 2. Parent size 3 with empty inline child at insertion point 1 AND
-  // conditional child (ID 3, size 2):
-  {
-    absl::flat_hash_map<uint32_t, size_t> base_sizes = {
-        {{1, 3}, {2, 0}, {3, 2}}};
-    absl::flat_hash_map<uint32_t,
-                        std::vector<CudaGraphTopologyMapper::ChildGraphEntry>>
-        child_graphs;
-    child_graphs[1].push_back({2, 1, /*is_conditional=*/false});
-    child_graphs[1].push_back({3, 2, /*is_conditional=*/true});
-
-    absl::flat_hash_map<uint32_t, size_t> merged_sizes;
-    size_t size = CudaGraphTopologyMapper::CalculateMergedSize(
-        1, base_sizes, child_graphs, &merged_sizes);
-    EXPECT_EQ(size, 4);
-
-    auto res0 = CudaGraphTopologyMapper::ResolveMergedNode(
-        1, 0, base_sizes, child_graphs, &merged_sizes);
-    EXPECT_EQ(res0.first, 1);
-    EXPECT_EQ(res0.second, 0);
-
-    auto res1 = CudaGraphTopologyMapper::ResolveMergedNode(
-        1, 1, base_sizes, child_graphs, &merged_sizes);
-    EXPECT_EQ(res1.first, 1);
-    EXPECT_EQ(res1.second, 2);
-
-    auto res2 = CudaGraphTopologyMapper::ResolveMergedNode(
-        1, 2, base_sizes, child_graphs, &merged_sizes);
-    EXPECT_EQ(res2.first, 3);
-    EXPECT_EQ(res2.second, 0);
-
-    auto res3 = CudaGraphTopologyMapper::ResolveMergedNode(
-        1, 3, base_sizes, child_graphs, &merged_sizes);
-    EXPECT_EQ(res3.first, 3);
-    EXPECT_EQ(res3.second, 1);
-  }
-}
-
-}  // namespace
-}  // namespace profiler
-}  // namespace xla
diff --git a/third_party/xla/xla/stream_executor/gpu/BUILD b/third_party/xla/xla/stream_executor/gpu/BUILD
index a4f172a4af8..3f051c01c5b 100644
--- a/third_party/xla/xla/stream_executor/gpu/BUILD
+++ b/third_party/xla/xla/stream_executor/gpu/BUILD
@@ -116,56 +116,6 @@ xla_cc_test(
     ],
 )
 
-cc_library(
-    name = "scoped_command_buffer_annotation",
-    srcs = ["scoped_command_buffer_annotation.cc"],
-    hdrs = ["scoped_command_buffer_annotation.h"],
-    deps = [
-        "@com_google_absl//absl/strings",
-    ],
-)
-
-cc_library(
-    name = "gpu_command_buffer_listener",
-    srcs = ["gpu_command_buffer_listener.cc"],
-    hdrs = ["gpu_command_buffer_listener.h"],
-    deps = [
-        ":gpu_command_buffer",
-        "@com_google_absl//absl/strings",
-    ],
-)
-
-cc_library(
-    name = "gpu_command_buffer_listener_test_helper",
-    testonly = True,
-    hdrs = ["gpu_command_buffer_listener_test_helper.h"],
-    deps = [
-        ":gpu_command_buffer_listener",
-    ],
-)
-
-xla_cc_test(
-    name = "gpu_command_buffer_listener_test",
-    srcs = ["gpu_command_buffer_listener_test.cc"],
-    deps = [
-        ":gpu_command_buffer",
-        ":gpu_command_buffer_listener",
-        ":gpu_command_buffer_listener_test_helper",
-        "@com_google_absl//absl/strings",
-        "@com_google_googletest//:gtest_main",
-    ],
-)
-
-xla_cc_test(
-    name = "scoped_command_buffer_annotation_test",
-    srcs = ["scoped_command_buffer_annotation_test.cc"],
-    deps = [
-        ":scoped_command_buffer_annotation",
-        "//xla/tsl/platform:env",
-        "@com_google_googletest//:gtest_main",
-    ],
-)
-
 cc_library(
     name = "mock_context",
     testonly = True,
@@ -1207,3 +1157,22 @@ cc_library(
         "@com_google_googletest//:gtest_for_library",
     ],
 )
+
+cc_library(
+    name = "scoped_command_buffer_annotation",
+    srcs = ["scoped_command_buffer_annotation.cc"],
+    hdrs = ["scoped_command_buffer_annotation.h"],
+    deps = [
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+xla_cc_test(
+    name = "scoped_command_buffer_annotation_test",
+    srcs = ["scoped_command_buffer_annotation_test.cc"],
+    deps = [
+        ":scoped_command_buffer_annotation",
+        "//xla/tsl/platform:env",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_listener.cc b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_listener.cc
deleted file mode 100644
index 6ecdb336dcd..00000000000
--- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_listener.cc
+++ /dev/null
@@ -1,60 +0,0 @@
-/* Copyright 2026 The OpenXLA Authors.
-
-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.
-==============================================================================*/
-
-#include "xla/stream_executor/gpu/gpu_command_buffer_listener.h"
-
-#include <atomic>
-
-namespace stream_executor::gpu {
-namespace {
-
-static std::atomic<GpuCommandBufferListener*> global_listener{nullptr};
-thread_local GpuCommandBufferListener* tls_listener = nullptr;
-
-}  // namespace
-
-bool RegisterGpuCommandBufferListener(GpuCommandBufferListener* listener) {
-  if (listener == nullptr) {
-    return false;
-  }
-  GpuCommandBufferListener* expected = nullptr;
-  return global_listener.compare_exchange_strong(expected, listener,
-                                                 std::memory_order_acq_rel);
-}
-
-bool UnregisterGpuCommandBufferListener(GpuCommandBufferListener* listener) {
-  if (listener == nullptr) {
-    return false;
-  }
-  GpuCommandBufferListener* expected = listener;
-  return global_listener.compare_exchange_strong(expected, nullptr,
-                                                 std::memory_order_acq_rel);
-}
-
-GpuCommandBufferListener* GetGpuCommandBufferListener() {
-  if (tls_listener != nullptr) {
-    return tls_listener;
-  }
-  return global_listener.load(std::memory_order_acquire);
-}
-
-GpuCommandBufferListener* GpuCommandBufferListener::ExchangeForTesting(
-    GpuCommandBufferListener* listener) {
-  GpuCommandBufferListener* old = tls_listener;
-  tls_listener = listener;
-  return old;
-}
-
-}  // namespace stream_executor::gpu
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_listener.h b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_listener.h
deleted file mode 100644
index 2afbecc38b7..00000000000
--- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_listener.h
+++ /dev/null
@@ -1,155 +0,0 @@
-/* Copyright 2026 The OpenXLA Authors.
-
-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.
-==============================================================================*/
-
-#ifndef XLA_STREAM_EXECUTOR_GPU_GPU_COMMAND_BUFFER_LISTENER_H_
-#define XLA_STREAM_EXECUTOR_GPU_GPU_COMMAND_BUFFER_LISTENER_H_
-
-#include <cstddef>
-
-#include "absl/strings/string_view.h"
-#include "xla/stream_executor/gpu/gpu_command_buffer.h"
-
-namespace stream_executor::gpu {
-
-// Opaque wrapper for a GPU graph handle (e.g., `CUgraph` on CUDA).
-// This wrapper is used to keep platform-specific graph handle types out of the
-// public stream_executor API, while enabling type-safe tracking of graph
-// resources by telemetry tools.
-class GpuGraphHandle {
- public:
-  // Constructs a GpuGraphHandle from a raw opaque pointer.
-  explicit GpuGraphHandle(void* opaque = nullptr) : opaque_(opaque) {}
-
-  // Returns the underlying raw opaque pointer.
-  void* opaque() const { return opaque_; }
-
-  // Returns true if the handle is null.
-  bool is_null() const { return opaque_ == nullptr; }
-
-  // Comparisons are performed on the underlying raw pointer values.
-  bool operator==(const GpuGraphHandle& other) const {
-    return opaque_ == other.opaque_;
-  }
-  bool operator!=(const GpuGraphHandle& other) const {
-    return opaque_ != other.opaque_;
-  }
-  bool operator<(const GpuGraphHandle& other) const {
-    return opaque_ < other.opaque_;
-  }
-
-  // Support for absl hashing to allow using this handle as a key in hash maps.
-  template <typename H>
-  friend H AbslHashValue(H h, const GpuGraphHandle& handle) {
-    return H::combine(std::move(h), handle.opaque_);
-  }
-
- private:
-  void* opaque_;
-};
-
-// Opaque wrapper for an executable GPU graph handle (e.g., `CUgraphExec` on
-// CUDA). This wrapper is used to keep platform-specific executable graph handle
-// types out of the public stream_executor API, while enabling type-safe
-// tracking of executable graph resources by telemetry tools.
-class GpuGraphExecHandle {
- public:
-  // Constructs a GpuGraphExecHandle from a raw opaque pointer.
-  explicit GpuGraphExecHandle(void* opaque = nullptr) : opaque_(opaque) {}
-
-  // Returns the underlying raw opaque pointer.
-  void* opaque() const { return opaque_; }
-
-  // Returns true if the handle is null.
-  bool is_null() const { return opaque_ == nullptr; }
-
-  // Comparisons are performed on the underlying raw pointer values.
-  bool operator==(const GpuGraphExecHandle& other) const {
-    return opaque_ == other.opaque_;
-  }
-  bool operator!=(const GpuGraphExecHandle& other) const {
-    return opaque_ != other.opaque_;
-  }
-  bool operator<(const GpuGraphExecHandle& other) const {
-    return opaque_ < other.opaque_;
-  }
-
-  // Support for absl hashing to allow using this handle as a key in hash maps.
-  template <typename H>
-  friend H AbslHashValue(H h, const GpuGraphExecHandle& handle) {
-    return H::combine(std::move(h), handle.opaque_);
-  }
-
- private:
-  void* opaque_;
-};
-
-class ScopedGpuCommandBufferListenerOverrideForTesting;
-
-// Listener interface for monitoring CUDA Graph events without direct
-// dependencies on the Profiler.
-class GpuCommandBufferListener {
- public:
-  virtual ~GpuCommandBufferListener() = default;
-
-  // Returns true if the listener is enabled and should receive events.
-  virtual bool IsEnabled() const { return true; }
-
-  // Called when an HLO annotation is associated with a graph node.
-  virtual void OnRegisterNodeAnnotation(GpuGraphHandle graph,
-                                        GpuCommandBuffer::GraphNodeHandle node,
-                                        absl::string_view annotation) = 0;
-
-  // Called when a graph's total size (number of nodes) is recorded.
-  virtual void OnRegisterGraphSize(GpuGraphHandle graph, size_t size) = 0;
-
-  // Called when a child graph (e.g., conditional or loop body) is associated
-  // with a node in a parent graph.
-  virtual void OnRegisterChildGraph(
-      GpuGraphHandle parent_graph, GpuGraphHandle child_graph,
-      GpuCommandBuffer::GraphNodeHandle child_node, bool is_conditional) = 0;
-
-  // Called when an executable graph is instantiated from a graph definition.
-  virtual void OnRegisterGraphExec(GpuGraphExecHandle graph_exec,
-                                   GpuGraphHandle graph) = 0;
-
-  // Called when an executable graph is destroyed.
-  virtual void OnUnregisterGraphExec(GpuGraphExecHandle graph_exec) = 0;
-
-  // Called when a graph definition is destroyed, to clean up associated
-  // annotations.
-  virtual void OnUnregisterGraphAnnotations(GpuGraphHandle graph) = 0;
-
- private:
-  friend class ScopedGpuCommandBufferListenerOverrideForTesting;
-  static GpuCommandBufferListener* ExchangeForTesting(
-      GpuCommandBufferListener* listener);
-};
-
-// Registers the global GpuCommandBufferListener. Returns true if successful,
-// or false if a listener is already registered. Thread-safe.
-bool RegisterGpuCommandBufferListener(GpuCommandBufferListener* listener);
-
-// Unregisters the given GpuCommandBufferListener. Returns true if successful,
-// or false if the given listener was not the currently registered one.
-// Thread-safe.
-bool UnregisterGpuCommandBufferListener(GpuCommandBufferListener* listener);
-
-// Retrieves the registered global GpuCommandBufferListener, or nullptr if none.
-// Thread-safe.
-GpuCommandBufferListener* GetGpuCommandBufferListener();
-
-}  // namespace stream_executor::gpu
-
-#endif  // XLA_STREAM_EXECUTOR_GPU_GPU_COMMAND_BUFFER_LISTENER_H_
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_listener_test.cc b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_listener_test.cc
deleted file mode 100644
index 8529850158a..00000000000
--- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_listener_test.cc
+++ /dev/null
@@ -1,108 +0,0 @@
-/* Copyright 2026 The OpenXLA Authors.
-
-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.
-==============================================================================*/
-
-#include "xla/stream_executor/gpu/gpu_command_buffer_listener.h"
-
-#include <cstddef>
-#include <thread>  // NOLINT(build/c++11)
-
-#include <gmock/gmock.h>
-#include <gtest/gtest.h>
-#include "absl/strings/string_view.h"
-#include "xla/stream_executor/gpu/gpu_command_buffer.h"
-#include "xla/stream_executor/gpu/gpu_command_buffer_listener_test_helper.h"
-
-namespace stream_executor::gpu {
-namespace {
-
-class MockGpuCommandBufferListener : public GpuCommandBufferListener {
- public:
-  MOCK_METHOD(void, OnRegisterNodeAnnotation,
-              (GpuGraphHandle graph, GpuCommandBuffer::GraphNodeHandle node,
-               absl::string_view annotation),
-              (override));
-  MOCK_METHOD(void, OnRegisterGraphSize, (GpuGraphHandle graph, size_t size),
-              (override));
-  MOCK_METHOD(void, OnRegisterChildGraph,
-              (GpuGraphHandle parent_graph, GpuGraphHandle child_graph,
-               GpuCommandBuffer::GraphNodeHandle child_node,
-               bool is_conditional),
-              (override));
-  MOCK_METHOD(void, OnRegisterGraphExec,
-              (GpuGraphExecHandle graph_exec, GpuGraphHandle graph),
-              (override));
-  MOCK_METHOD(void, OnUnregisterGraphExec, (GpuGraphExecHandle graph_exec),
-              (override));
-  MOCK_METHOD(void, OnUnregisterGraphAnnotations, (GpuGraphHandle graph),
-              (override));
-};
-
-TEST(GpuCommandBufferListenerTest, RegistrationAndRetrieval) {
-  EXPECT_EQ(GetGpuCommandBufferListener(), nullptr);
-
-  MockGpuCommandBufferListener listener1;
-  MockGpuCommandBufferListener listener2;
-
-  // First registration succeeds.
-  EXPECT_TRUE(RegisterGpuCommandBufferListener(&listener1));
-  EXPECT_EQ(GetGpuCommandBufferListener(), &listener1);
-
-  // Overwriting with another listener fails.
-  EXPECT_FALSE(RegisterGpuCommandBufferListener(&listener2));
-  EXPECT_EQ(GetGpuCommandBufferListener(), &listener1);
-
-  // Unregistering with wrong listener pointer fails.
-  EXPECT_FALSE(UnregisterGpuCommandBufferListener(&listener2));
-  EXPECT_EQ(GetGpuCommandBufferListener(), &listener1);
-
-  // Unregistering with correct pointer succeeds.
-  EXPECT_TRUE(UnregisterGpuCommandBufferListener(&listener1));
-  EXPECT_EQ(GetGpuCommandBufferListener(), nullptr);
-
-  // Null pointers are rejected.
-  EXPECT_FALSE(RegisterGpuCommandBufferListener(nullptr));
-  EXPECT_FALSE(UnregisterGpuCommandBufferListener(nullptr));
-}
-
-TEST(GpuCommandBufferListenerTest, ThreadLocalOverride) {
-  EXPECT_EQ(GetGpuCommandBufferListener(), nullptr);
-
-  MockGpuCommandBufferListener global_listener;
-  MockGpuCommandBufferListener override_listener;
-
-  EXPECT_TRUE(RegisterGpuCommandBufferListener(&global_listener));
-  EXPECT_EQ(GetGpuCommandBufferListener(), &global_listener);
-
-  {
-    ScopedGpuCommandBufferListenerOverrideForTesting listener_override(
-        &override_listener);
-    EXPECT_EQ(GetGpuCommandBufferListener(), &override_listener);
-
-    // Verify it is thread local.
-    std::thread t([&global_listener]() {
-      // Other thread still sees the global listener.
-      EXPECT_EQ(GetGpuCommandBufferListener(), &global_listener);
-    });
-    t.join();
-  }
-
-  // Restored to global after scope exits.
-  EXPECT_EQ(GetGpuCommandBufferListener(), &global_listener);
-
-  EXPECT_TRUE(UnregisterGpuCommandBufferListener(&global_listener));
-}
-
-}  // namespace
-}  // namespace stream_executor::gpu
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_listener_test_helper.h b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_listener_test_helper.h
deleted file mode 100644
index b45f45ecf4b..00000000000
--- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_listener_test_helper.h
+++ /dev/null
@@ -1,38 +0,0 @@
-/* Copyright 2026 The OpenXLA Authors.
-
-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.
-==============================================================================*/
-
-#ifndef XLA_STREAM_EXECUTOR_GPU_GPU_COMMAND_BUFFER_LISTENER_TEST_HELPER_H_
-#define XLA_STREAM_EXECUTOR_GPU_GPU_COMMAND_BUFFER_LISTENER_TEST_HELPER_H_
-
-#include "xla/stream_executor/gpu/gpu_command_buffer_listener.h"
-
-namespace stream_executor::gpu {
-
-class ScopedGpuCommandBufferListenerOverrideForTesting {
- public:
-  explicit ScopedGpuCommandBufferListenerOverrideForTesting(
-      GpuCommandBufferListener* listener)
-      : old_listener_(GpuCommandBufferListener::ExchangeForTesting(listener)) {}
-  ~ScopedGpuCommandBufferListenerOverrideForTesting() {
-    GpuCommandBufferListener::ExchangeForTesting(old_listener_);
-  }
-
- private:
-  GpuCommandBufferListener* old_listener_;
-};
-
-}  // namespace stream_executor::gpu
-
-#endif  // XLA_STREAM_EXECUTOR_GPU_GPU_COMMAND_BUFFER_LISTENER_TEST_HELPER_H_

From b4ba58be6042a3b37cfe966534e8e9581e779543 Mon Sep 17 00:00:00 2001
From: Joshua Lang <joshualang@google.com>
Date: Thu, 8 Oct 2026 11:12:45 -0700
Subject: [PATCH 0310/1000] Configure DTensor GPU test defaults for Ampere+
 GPUs and smaller GPU memory partitions

Reverts changelist 989940859

PiperOrigin-RevId: 995951800
---
 tensorflow/dtensor/python/tests/test_util.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensorflow/dtensor/python/tests/test_util.py b/tensorflow/dtensor/python/tests/test_util.py
index d6e875a5b5f..d8ebcb6b199 100644
--- a/tensorflow/dtensor/python/tests/test_util.py
+++ b/tensorflow/dtensor/python/tests/test_util.py
@@ -54,7 +54,9 @@ v2_compat.enable_v2_behavior()
 
 DEFAULT_TOL = 1e-5
 
-_DEFAULT_GPU_MEMORY_LIMIT = 1024  # 1G
+# Limit each virtual GPU to 512MB so 8 virtual devices plus CUDA/cuSolver
+# runtime overhead fit within smaller GPU memory partitions.
+_DEFAULT_GPU_MEMORY_LIMIT = 512
 
 
 def get_use_xla_spmd(device_type):

From e5d0759f27e63271ee3a0ad7799ac01e7e5089d4 Mon Sep 17 00:00:00 2001
From: Som Tripathi <somtri@iastate.edu>
Date: Thu, 8 Oct 2026 14:29:58 -0500
Subject: [PATCH 0311/1000] Reject GPU Where inputs with at least INT32_MAX
 elements

The GPU kernel passes the int64 element count to CUB calls that take an
int, so inputs past INT32_MAX truncate it and the select writes nothing.
Where then returned uninitialized memory without an error. Fail with
InvalidArgument from INT32_MAX elements on, before any kernel runs.
Also compute the output offsets NDIM * i and NDIM * n in int64_t, which
overflowed int for large outputs with 2 or more dimensions.

Fixes #104365
---
 tensorflow/core/kernels/BUILD                 |  1 +
 tensorflow/core/kernels/where_op.cc           | 16 ++++++++++++-
 tensorflow/core/kernels/where_op_gpu.cu.h     |  6 ++---
 .../python/kernel_tests/array_ops/BUILD       |  1 +
 .../kernel_tests/array_ops/where_op_test.py   | 24 +++++++++++++++++++
 5 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 390d2ab957e..d11d2d06d32 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -1354,6 +1354,7 @@ tf_kernel_library(
     ]) + [
         ":gpu_prim_hdrs",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
     ] + ARRAY_DEPS,
 )
 
diff --git a/tensorflow/core/kernels/where_op.cc b/tensorflow/core/kernels/where_op.cc
index e5fff24c745..1117dd210c4 100644
--- a/tensorflow/core/kernels/where_op.cc
+++ b/tensorflow/core/kernels/where_op.cc
@@ -23,9 +23,12 @@ limitations under the License.
 
 #include "tensorflow/core/kernels/where_op.h"
 
+#include <limits>
 #include <memory>
 #include <numeric>
-#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
+
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "tensorflow/core/framework/bounds_check.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
@@ -36,6 +39,7 @@ limitations under the License.
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/types.h"
+#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
 
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 #include "tensorflow/core/common_runtime/gpu/gpu_event_mgr.h"
@@ -259,6 +263,16 @@ class WhereGPUOp : public AsyncOpKernel {
     const Tensor& input = context->input(0);
     const int input_dims = input.dims();
 
+    // The CUB count and select calls take an int item count. Larger inputs
+    // truncate it, and the select then leaves the output unwritten.
+    OP_REQUIRES_ASYNC(context,
+                      input.NumElements() < std::numeric_limits<int32_t>::max(),
+                      absl::InvalidArgumentError(absl::StrCat(
+                          "WhereOp: the GPU kernel supports fewer than ",
+                          std::numeric_limits<int32_t>::max(),
+                          " input elements, got ", input.NumElements(), ".")),
+                      done);
+
     if (input.NumElements() < std::numeric_limits<int32_t>::max()) {
       ComputeAsyncType<int32_t>(input, input_dims, context, done);
     } else {
diff --git a/tensorflow/core/kernels/where_op_gpu.cu.h b/tensorflow/core/kernels/where_op_gpu.cu.h
index f9d83d7fe7a..b57be0f0802 100644
--- a/tensorflow/core/kernels/where_op_gpu.cu.h
+++ b/tensorflow/core/kernels/where_op_gpu.cu.h
@@ -44,10 +44,10 @@ __global__ void PropagateWhereIndicesKernel(
   // dimensions of individual indices manually, instead of relying on
   // a scalar loop variable and using integer division.
   GPU_1D_KERNEL_LOOP(i, output_rows) {
-    TIndex index_value = ldg(output + NDIM * i);
+    TIndex index_value = ldg(output + static_cast<int64_t>(NDIM) * i);
 #pragma unroll
     for (int c = 0; c < NDIM; ++c) {
-      *(output + NDIM * i + c) = index_value / strides[c];
+      *(output + static_cast<int64_t>(NDIM) * i + c) = index_value / strides[c];
       index_value %= strides[c];
     }
   }
@@ -226,7 +226,7 @@ class WhereOutputIterator {
     // the number of items found in Flagged()'s d_num_selected_out at
     // the end and confirm that it matches the number of rows of output.
     const bool valid = FastBoundsCheck(n, max_row_);
-    return *(ptr_ + (valid ? (NDIM * n) : 0));
+    return *(ptr_ + (valid ? (static_cast<int64_t>(NDIM) * n) : 0));
   }
 
   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE reference operator*() const {
diff --git a/tensorflow/python/kernel_tests/array_ops/BUILD b/tensorflow/python/kernel_tests/array_ops/BUILD
index 1278cc3cb75..d060af574ca 100644
--- a/tensorflow/python/kernel_tests/array_ops/BUILD
+++ b/tensorflow/python/kernel_tests/array_ops/BUILD
@@ -871,6 +871,7 @@ cuda_py_strict_test(
     deps = [
         "//tensorflow/python/client:session",
         "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:errors",
         "//tensorflow/python/framework:for_generated_wrappers",
         "//tensorflow/python/framework:test_lib",
         "//tensorflow/python/ops:array_ops",
diff --git a/tensorflow/python/kernel_tests/array_ops/where_op_test.py b/tensorflow/python/kernel_tests/array_ops/where_op_test.py
index 99a8b9d9786..8d97d8270f8 100644
--- a/tensorflow/python/kernel_tests/array_ops/where_op_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/where_op_test.py
@@ -22,6 +22,7 @@ import numpy as np
 from tensorflow.python.client import session
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
@@ -263,6 +264,29 @@ class WhereOpTest(test.TestCase):
       tf_val = array_ops.where(c_vec, x * x, -x).eval()
     self.assertAllEqual(tf_val, np_val)
 
+  @test_util.run_in_graph_and_eager_modes
+  @test_util.disable_xla("Tests the TF GPU Where kernel element-count guard")
+  def testGpuRejectsInt32MaxElements(self):
+    # The GPU kernel passes the element count to CUB as an int. Past INT32_MAX
+    # elements the count truncated and the op returned uninitialized memory.
+    if not test_util.is_gpu_available():
+      self.skipTest("Requires a GPU.")
+    for shape in ([2**31 - 1], [2**16, 2**15]):
+      with self.subTest(shape=shape), test_util.device(use_gpu=True):
+        x = array_ops.zeros(shape, dtype=dtypes.bool)
+        with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                    "fewer than 2147483647 input elements"):
+          self.evaluate(array_ops.where(x))
+        del x
+
+  @test_util.run_in_graph_and_eager_modes
+  def testGpuAcceptsInputBelowInt32MaxElements(self):
+    if not test_util.is_gpu_available():
+      self.skipTest("Requires a GPU.")
+    with test_util.device(use_gpu=True):
+      x = array_ops.zeros([2**31 - 2], dtype=dtypes.bool)
+      self.assertAllEqual(self.evaluate(array_ops.where(x)).shape, [0, 1])
+
 
 class WhereBenchmark(test.Benchmark):
 

From bf5416b3417497c38e931272d2029f5b8eba0183 Mon Sep 17 00:00:00 2001
From: Zviki Nozadze <zviki@google.com>
Date: Thu, 8 Oct 2026 11:35:09 -0700
Subject: [PATCH 0312/1000] N/A

PiperOrigin-RevId: 995966458
---
 .../xla/xla/service/heap_simulator/heap_simulator.cc     | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
index ef992b772bc..7d7b1015f44 100644
--- a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
+++ b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
@@ -2065,8 +2065,10 @@ std::string RenderTimeByFreeChunks(
 // strategy.
 //
 // Optimizations:
-// 1. Tuned cutoff: For small arrays (N < 3000), elements fit entirely
-//    within L1/L2 cache (<48 KB). `absl::c_sort` is faster than radix sort.
+// 1. Tuned cutoff: For small arrays (N < 3000, ~48 kB for 16-byte Chunks),
+//    elements fit comfortably in L2 cache (typically 1+ MB, though spilling
+//    the 32-48 kB L1d cache). For this size, introsort (`absl::c_sort`) has
+//    lower constant factor overhead and is faster than radix sort.
 // 2. Single pass to check sortedness by counting inversions:
 //    - If nearly sorted (`inversions <= n / 10`), introsort does almost zero
 //      swaps and beats radix sort by >2x.
@@ -2079,7 +2081,8 @@ void AdaptiveHybridSortChunks(std::vector<HeapSimulator::Chunk>& chunks,
   using Chunk = HeapSimulator::Chunk;
   const size_t n = chunks.size();
 
-  // Small lists fit in L1/L2 cache and sort fastest with introsort.
+  // Small lists fit within L2 cache (~48 kB data spills 32-48 kB L1d, but is
+  // well within L2) where introsort has lower constant overhead than radix sort
   if (n < 3000) {
     absl::c_sort(chunks, [](const Chunk& a, const Chunk& b) {
       return a.offset < b.offset;

From a814257ee0c331ed289e2a867bfd77524d5ef1f1 Mon Sep 17 00:00:00 2001
From: Karlo Basioli <basioli@google.com>
Date: Thu, 8 Oct 2026 11:41:32 -0700
Subject: [PATCH 0313/1000] Fold sums of disjoint zero-padded blocks into a
 single pad(concat)

In PyTorch AOTAutograd emits the gradient of a tensor that is split into pieces as a sum of zero-padded pieces.

Notably, for `q, k, v = x[0], x[1], x[2]` we get a sum of `dynamic-update-slice(zeros, dq, 0)`-style updates.

Today AlgebraicSimplifier rewrites `add(A, dus(zeros, B, idx))` with dynamic indices into a serial chain of in-place updates. For constant indices, `HandleDynamicUpdateSlice` already rewrites DUS into zeros to `Pad`, but did not look through `kConvert` index operands left by dtype conversions.

This change:
1. Updates `HandleDynamicUpdateSlice` to look through `kConvert` on constant index operands and evaluate their converted integer values, allowing DUS into broadcast zeros with converted constant offsets to rewrite to `Pad`.
2. Adds a rewrite to `HandleAdd`: zero-padded boxes (pads without interior padding) that coincide in all but one dimension and are disjoint in it fold into `pad(concatenate(a, zeros, b))`. The result is again a box, so n-ary sums collapse into a single flat concatenate. DUS updates with dynamic indices continue to use the in-place DUS fold.

PiperOrigin-RevId: 995971101
---
 .../simplifiers/algebraic_simplifier.cc       | 189 +++++++++++++++++-
 .../simplifiers/algebraic_simplifier_test.cc  | 132 ++++++++++++
 2 files changed, 315 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
index b989731d479..9d0d7941751 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
@@ -934,6 +934,179 @@ absl::Status AlgebraicSimplifierVisitor::HandleAbs(HloInstruction* abs) {
   return absl::OkStatus();
 }
 
+namespace {
+
+// Returns the value of a constant integral scalar, looking through converts
+// (e.g. `convert(u32 constant)` left behind by X64 elimination) and applying
+// the conversions so type semantics are respected.
+std::optional<int64_t> GetConstantIntegralScalar(const HloInstruction* hlo) {
+  if (!ShapeUtil::IsScalar(hlo->shape()) ||
+      !primitive_util::IsIntegralType(hlo->shape().element_type())) {
+    return std::nullopt;
+  }
+  const HloInstruction* current = hlo;
+  absl::InlinedVector<PrimitiveType, 2> convert_types;
+  while (current->opcode() == HloOpcode::kConvert) {
+    convert_types.push_back(current->shape().element_type());
+    current = current->operand(0);
+  }
+  if (!Match(current, m::ConstantScalar())) {
+    return std::nullopt;
+  }
+  Literal lit = current->literal().Clone();
+  for (auto it = convert_types.rbegin(); it != convert_types.rend(); ++it) {
+    absl::StatusOr<Literal> converted = lit.Convert(*it);
+    if (!converted.ok()) {
+      return std::nullopt;
+    }
+    lit = *std::move(converted);
+  }
+  return lit.GetFirstInteger();
+}
+
+// Returns `pad(operand, zero, config)`, or `operand` if `config` is a no-op.
+absl::StatusOr<HloInstruction*> MakePadIfNeeded(
+    HloInstruction* operand, HloInstruction* zero, const PaddingConfig& config,
+    const OpMetadata* metadata = nullptr) {
+  if (absl::c_all_of(config.dimensions(), [](const auto& dim) {
+        return dim.edge_padding_low() == 0 && dim.edge_padding_high() == 0 &&
+               dim.interior_padding() == 0;
+      })) {
+    return operand;
+  }
+  return MakePadHlo(operand, zero, config, metadata);
+}
+
+// An array that is zero everywhere except for an axis-aligned box, which holds
+// `source`. `low(d)` / `high(d)` are the number of zeros before / after the box
+// in dimension `d`.
+class ZeroPaddedBox {
+ public:
+  ZeroPaddedBox(HloInstruction* source, std::vector<int64_t> low,
+                std::vector<int64_t> high)
+      : source_(source), low_(std::move(low)), high_(std::move(high)) {
+    CHECK_EQ(low_.size(), source_->shape().dimensions().size());
+    CHECK_EQ(high_.size(), source_->shape().dimensions().size());
+  }
+
+  HloInstruction* source() const { return source_; }
+  int64_t low(int64_t dim) const { return low_[dim]; }
+  int64_t high(int64_t dim) const { return high_[dim]; }
+  int64_t rank() const { return low_.size(); }
+
+ private:
+  HloInstruction* source_;
+  std::vector<int64_t> low_;
+  std::vector<int64_t> high_;
+};
+
+// Matches `pad(source, 0)` without interior or negative padding.
+std::optional<ZeroPaddedBox> MatchZeroPaddedBox(HloInstruction* hlo) {
+  if (!hlo->shape().IsArray() || hlo->opcode() != HloOpcode::kPad ||
+      !IsAll(hlo->operand(1), 0)) {
+    return std::nullopt;
+  }
+  std::vector<int64_t> low;
+  std::vector<int64_t> high;
+  for (const auto& dim : hlo->padding_config().dimensions()) {
+    if (dim.interior_padding() != 0 || dim.edge_padding_low() < 0 ||
+        dim.edge_padding_high() < 0) {
+      return std::nullopt;
+    }
+    low.push_back(dim.edge_padding_low());
+    high.push_back(dim.edge_padding_high());
+  }
+  return ZeroPaddedBox(hlo->mutable_operand(0), std::move(low),
+                       std::move(high));
+}
+
+// Returns the only dimension in which `a` and `b` are placed differently, or
+// std::nullopt if they are placed identically or differ in several dimensions.
+std::optional<int64_t> FindConcatDim(const ZeroPaddedBox& a,
+                                     const ZeroPaddedBox& b) {
+  std::optional<int64_t> concat_dim;
+  for (int64_t d = 0; d < a.rank(); ++d) {
+    if (a.low(d) == b.low(d) && a.high(d) == b.high(d)) {
+      continue;
+    }
+    // The boxes may be placed differently in only one dimension, the one we
+    // concatenate along.
+    if (concat_dim.has_value()) {
+      return std::nullopt;
+    }
+    concat_dim = d;
+  }
+  return concat_dim;
+}
+
+// Returns the operands of `concatenate(first.source, zeros, second.source)`
+// along `concat_dim`, where the zeros fill a gap of size `gap` (omitted if 0).
+std::vector<HloInstruction*> MakeConcatParts(const ZeroPaddedBox& first,
+                                             const ZeroPaddedBox& second,
+                                             int64_t concat_dim, int64_t gap,
+                                             HloInstruction* zero) {
+  std::vector<HloInstruction*> parts = {first.source()};
+  if (gap > 0) {
+    std::vector<int64_t> gap_dims(first.source()->shape().dimensions().begin(),
+                                  first.source()->shape().dimensions().end());
+    gap_dims[concat_dim] = gap;
+    parts.push_back(MakeBroadcastHlo(zero, {}, gap_dims));
+  }
+  parts.push_back(second.source());
+  return parts;
+}
+
+// Folds `add(A, B)` where A and B are zero-padded boxes that coincide in all
+// but one dimension and are disjoint in that one into
+// `pad(concatenate(A.source, zeros, B.source), 0)`.
+//
+// Autograd emits such sums for the gradient of a tensor that is split into
+// pieces, e.g. `q, k, v = x[0], x[1], x[2]`. The result is again a zero-padded
+// box, so n-ary sums collapse into nested concatenates, which HandleConcatenate
+// flattens.
+absl::StatusOr<HloInstruction*> TryFoldAddOfDisjointZeroPaddedBoxes(
+    HloInstruction* add) {
+  std::optional<ZeroPaddedBox> a = MatchZeroPaddedBox(add->mutable_operand(0));
+  if (!a.has_value()) {
+    return nullptr;
+  }
+  std::optional<ZeroPaddedBox> b = MatchZeroPaddedBox(add->mutable_operand(1));
+  if (!b.has_value()) {
+    return nullptr;
+  }
+  std::optional<int64_t> maybe_concat_dim = FindConcatDim(*a, *b);
+  if (!maybe_concat_dim.has_value()) {
+    return nullptr;
+  }
+  const int64_t concat_dim = *maybe_concat_dim;
+  if (a->low(concat_dim) > b->low(concat_dim)) {
+    std::swap(a, b);
+  }
+  // Both boxes lie inside the same array, so with `a` starting first they are
+  // disjoint iff `b` starts at or after the end of `a`. `b`'s end, and hence
+  // its high padding, is then determined by its size; it needs no check.
+  const int64_t a_end =
+      add->shape().dimensions(concat_dim) - a->high(concat_dim);
+  const int64_t gap = b->low(concat_dim) - a_end;
+  if (gap < 0) {
+    return nullptr;
+  }
+  HloInstruction* zero = add->AddInstruction(HloInstruction::CreateConstant(
+      LiteralUtil::Zero(add->shape().element_type())));
+  ABSL_ASSIGN_OR_RETURN(HloInstruction * concat,
+                   MakeConcatHlo(MakeConcatParts(*a, *b, concat_dim, gap, zero),
+                                 concat_dim));
+  PaddingConfig padding_config = MakeNoPaddingConfig(a->rank());
+  for (int64_t d = 0; d < a->rank(); ++d) {
+    auto* pad_dim = padding_config.mutable_dimensions(d);
+    pad_dim->set_edge_padding_low(a->low(d));
+    pad_dim->set_edge_padding_high(d == concat_dim ? b->high(d) : a->high(d));
+  }
+  return MakePadIfNeeded(concat, zero, padding_config, &add->metadata());
+}
+
+}  // namespace
+
 absl::Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) {
   HloInstruction *lhs, *rhs;
   CHECK(Match(add, m::Add(m::Op(&lhs), m::Op(&rhs))));
@@ -1015,6 +1188,15 @@ absl::Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) {
     }
   }
 
+  // Sum of zero-padded blocks that tile a dimension => pad(concatenate(...)).
+  if (!options_.is_layout_sensitive()) {
+    ABSL_ASSIGN_OR_RETURN(HloInstruction * folded,
+                     TryFoldAddOfDisjointZeroPaddedBoxes(add));
+    if (folded != nullptr) {
+      return ReplaceInstruction(add, folded);
+    }
+  }
+
   // Convert add with fullshape into add with partial shape when a
   // portion of add is effective:
   //             zero (fullshape)   rhs (partialshape)
@@ -8304,13 +8486,8 @@ absl::Status AlgebraicSimplifierVisitor::HandleDynamicUpdateSlice(
       for (int64_t dim = 0; dim < updated_shape.dimensions().size(); ++dim) {
         auto padding_config_dim = padding_config.add_dimensions();
         auto slice_dim_start = update_start_indx->operand(dim + offset);
-        if (!Match(slice_dim_start, m::ConstantScalar())) {
-          compatible = false;
-          break;
-        }
-        VLOG(2) << "slice: " << slice_dim_start->ToString();
         std::optional<int64_t> start =
-            slice_dim_start->literal().GetFirstInteger();
+            GetConstantIntegralScalar(slice_dim_start);
         if (!start) {
           compatible = false;
           break;
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
index d1516cf3ca9..09333d0e06f 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
@@ -9505,6 +9505,56 @@ ENTRY AddBroadcastZeroWithDynamicSlice {
   EXPECT_THAT(root->operand(1)->opcode(), HloOpcode::kPad);
 }
 
+TEST_F(AlgebraicSimplifierTest, DynamicUpdateSliceOfBroadcastToPadWithConvert) {
+  constexpr absl::string_view hlo_string = R"(
+HloModule DusWithConvert
+
+ENTRY DusWithConvert {
+  constant = f32[] constant(0)
+  broadcast = f32[10,12] broadcast(constant), dimensions={}
+  param = f32[2,12] parameter(0)
+  u_idx = u32[] constant(3)
+  s_idx = s32[] convert(u_idx)
+  zero = s32[] constant(0)
+  ROOT dus = f32[10,12] dynamic-update-slice(broadcast, param, s_idx, zero)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  AlgebraicSimplifier simplifier(default_options_);
+  ASSERT_TRUE(simplifier.Run(module.get()).value());
+  auto root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(root->opcode(), HloOpcode::kPad);
+  EXPECT_EQ(root->padding_config().dimensions(0).edge_padding_low(), 3);
+  EXPECT_EQ(root->padding_config().dimensions(0).edge_padding_high(), 5);
+}
+
+TEST_F(AlgebraicSimplifierTest,
+       DynamicUpdateSliceOfBroadcastToPadWithTruncatingConvert) {
+  constexpr absl::string_view hlo_string = R"(
+HloModule DusWithTruncatingConvert
+
+ENTRY DusWithTruncatingConvert {
+  constant = f32[] constant(0)
+  broadcast = f32[10,12] broadcast(constant), dimensions={}
+  param = f32[2,12] parameter(0)
+  big_idx = s32[] constant(257)
+  trunc_idx = s8[] convert(big_idx)
+  s_idx = s32[] convert(trunc_idx)
+  zero = s32[] constant(0)
+  ROOT dus = f32[10,12] dynamic-update-slice(broadcast, param, s_idx, zero)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  AlgebraicSimplifier simplifier(default_options_);
+  ASSERT_TRUE(simplifier.Run(module.get()).value());
+  auto root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(root->opcode(), HloOpcode::kPad);
+  // 257 as s8 is 1, so low padding must be 1 (if convert was ignored, it would
+  // be 257).
+  EXPECT_EQ(root->padding_config().dimensions(0).edge_padding_low(), 1);
+  EXPECT_EQ(root->padding_config().dimensions(0).edge_padding_high(), 7);
+}
+
 // Test that dynamic-update-slice with a scalar broadcast does not become a pad
 // if the dynamic-update-slice is for host memory offload.
 TEST_F(AlgebraicSimplifierTest, DynamicUpdateSliceOfBroadcastToPadHostOffload) {
@@ -9792,6 +9842,88 @@ ENTRY AddDynamicUpdateSliceToAddSlice {
                         m::Parameter(2), m::Constant(), m::Constant())));
 }
 
+TEST_F(AlgebraicSimplifierTest, AddOfTilingDusIntoZerosBecomesConcat) {
+  // Gradient of `q, k, v = x[0], x[1], x[2]`, with indices as left behind by
+  // X64 elimination.
+  constexpr absl::string_view kHlo = R"(
+HloModule m
+
+ENTRY main {
+  pq = f32[1,8,16] parameter(0)
+  pk = f32[1,8,16] parameter(1)
+  pv = f32[1,8,16] parameter(2)
+  zero = f32[] constant(0)
+  zeros = f32[3,8,16] broadcast(zero), dimensions={}
+  ua = u32[] constant(0)
+  ub = u32[] constant(1)
+  uc = u32[] constant(2)
+  ja = s32[] convert(ua)
+  jb = s32[] convert(ub)
+  jc = s32[] convert(uc)
+  dq = f32[3,8,16] dynamic-update-slice(zeros, pq, ja, ja, ja)
+  dk = f32[3,8,16] dynamic-update-slice(zeros, pk, jb, ja, ja)
+  dv = f32[3,8,16] dynamic-update-slice(zeros, pv, jc, ja, ja)
+  sum = f32[3,8,16] add(dq, dk)
+  ROOT out = f32[3,8,16] add(sum, dv)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo));
+  HloPassFix<AlgebraicSimplifier> simplifier(default_options_);
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&simplifier, module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_THAT(module->entry_computation()->root_instruction(),
+              GmockMatch(m::Concatenate(m::Parameter(0), m::Parameter(1),
+                                        m::Parameter(2))));
+}
+
+TEST_F(AlgebraicSimplifierTest, AddOfDisjointPadsWithGapBecomesPadOfConcat) {
+  constexpr absl::string_view kHlo = R"(
+HloModule m
+
+ENTRY main {
+  pa = f32[4,2] parameter(0)
+  pb = f32[4,2] parameter(1)
+  zero = f32[] constant(0)
+  a = f32[4,10] pad(pa, zero), padding=0_0x0_8
+  b = f32[4,10] pad(pb, zero), padding=0_0x6_2
+  ROOT out = f32[4,10] add(a, b)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo));
+  AlgebraicSimplifier simplifier(default_options_);
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&simplifier, module.get()));
+  EXPECT_TRUE(changed);
+  const HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(
+      root, GmockMatch(m::Pad(m::Concatenate(m::Parameter(0),
+                                             m::Broadcast(m::ConstantScalar(0)),
+                                             m::Parameter(1)),
+                              m::ConstantScalar(0))));
+  EXPECT_EQ(root->padding_config().dimensions(1).edge_padding_low(), 0);
+  EXPECT_EQ(root->padding_config().dimensions(1).edge_padding_high(), 2);
+  EXPECT_EQ(root->operand(0)->operand(1)->shape().dimensions(1), 4);
+}
+
+TEST_F(AlgebraicSimplifierTest, AddOfOverlappingPadsIsNotFolded) {
+  constexpr absl::string_view kHlo = R"(
+HloModule m
+
+ENTRY main {
+  pa = f32[4,4] parameter(0)
+  pb = f32[4,4] parameter(1)
+  zero = f32[] constant(0)
+  a = f32[4,10] pad(pa, zero), padding=0_0x0_6
+  b = f32[4,10] pad(pb, zero), padding=0_0x2_4
+  ROOT out = f32[4,10] add(a, b)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo));
+  AlgebraicSimplifier simplifier(default_options_);
+  ASSERT_OK(RunHloPass(&simplifier, module.get()).status());
+  EXPECT_THAT(module->entry_computation()->root_instruction(),
+              GmockMatch(m::Add(m::Pad(), m::Pad())));
+}
+
 TEST_F(AlgebraicSimplifierTest, ScalarMultiplyReduction) {
   constexpr absl::string_view hlo_string = R"(
 HloModule ConstScalarMultiply

From 36ddd63cb13f255493a4e1944f42a407132b745f Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Thu, 8 Oct 2026 12:44:20 -0700
Subject: [PATCH 0314/1000] Migrate GitHub Actions runners from N2 to N4
 machine types.

Updates CI workflows, documentation, and configuration files to replace n2-standard runner pools (linux-x86-n2-16 and linux-x86-n2-128) with n4-standard pools (linux-x86-n4-16 and linux-x86-n4-80). Also removes the redundant Tests_v2 job from the main CI workflow.

PiperOrigin-RevId: 996009999
---
 third_party/xla/.github/actionlint.yml        |   3 -
 .../.github/workflows/benchmarks/README.md    |   2 +-
 third_party/xla/.github/workflows/ci.yml      | 101 ++----------------
 .../xla/.github/workflows/clang_tidy.yml      |   2 +-
 third_party/xla/docs/rocm_ci.md               |   2 +-
 5 files changed, 10 insertions(+), 100 deletions(-)

diff --git a/third_party/xla/.github/actionlint.yml b/third_party/xla/.github/actionlint.yml
index e7adf4f2939..bf762711e9e 100644
--- a/third_party/xla/.github/actionlint.yml
+++ b/third_party/xla/.github/actionlint.yml
@@ -14,9 +14,6 @@
 # ============================================================================
 self-hosted-runner:
   labels:
-    - "linux-x86-n2-16" # Linux X86 runner using the 16 vcpu n2-standard-16 machine.
-    - "linux-x86-n2-64" # Linux X86 runner using the 64 vcpu n2-standard-64 machine.
-    - "linux-x86-n2-128" # Linux X86 runner using the 128 vcpu n2-standard-128 machine.
     - "linux-x86-n4-16" # Linux X86 runner using the 16 vcpu n4-standard-16 machine.
     - "linux-x86-n4-80" # Linux X86 runner using the 80 vcpu n4-standard-80 machine.
     - "linux-x86-16cpu-l4-1gpu" # Linux X86 GPU runner using g2-standard-16 machine with 1 NVIDIA L4 GPU attached.
diff --git a/third_party/xla/.github/workflows/benchmarks/README.md b/third_party/xla/.github/workflows/benchmarks/README.md
index 002bb497dea..bc66e998ba4 100644
--- a/third_party/xla/.github/workflows/benchmarks/README.md
+++ b/third_party/xla/.github/workflows/benchmarks/README.md
@@ -229,7 +229,7 @@ benchmarks {
 
   environment_configs {
     id: "cpu_x86"
-    runner_label: "linux-x86-n2-128"
+    runner_label: "linux-x86-n4-80"
     container_image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
     workload_action_inputs {
       key: "hardware_category"
diff --git a/third_party/xla/.github/workflows/ci.yml b/third_party/xla/.github/workflows/ci.yml
index 4fe26f8b61b..c694fe35c6c 100644
--- a/third_party/xla/.github/workflows/ci.yml
+++ b/third_party/xla/.github/workflows/ci.yml
@@ -43,7 +43,7 @@ jobs:
       matrix:
         job_info: [
           {
-            pool: "linux-x86-n2-16",
+            pool: "linux-x86-n4-16",
             container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
             name: "XLA Linux x86 CPU",
             repo: "openxla/xla",
@@ -55,25 +55,25 @@ jobs:
             repo: "openxla/xla",
           },
           {
-            pool: "linux-x86-n2-16",
+            pool: "linux-x86-n4-16",
             container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
             name: "XLA Linux x86 CPU WORKSPACE",
             repo: "openxla/xla",
           },
           {
-            pool: "linux-x86-n2-16",
+            pool: "linux-x86-n4-16",
             container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
             name: "XLA Linux X86 GPU ONEAPI",
             repo: "openxla/xla",
           },
           {
-            pool: "linux-x86-n2-16",
+            pool: "linux-x86-n4-16",
             container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
             name: "XLA Linux x86 GPU Hermetic ROCm",
             repo: "openxla/xla",
           },
           {
-            pool: "linux-x86-n2-16",
+            pool: "linux-x86-n4-16",
             container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
             name: "XLA Linux x86 GPU T4",
             repo: "openxla/xla",
@@ -85,7 +85,7 @@ jobs:
             repo: "openxla/xla",
           },
           {
-            pool: "linux-x86-n2-16",
+            pool: "linux-x86-n4-16",
             container: "gcr.io/tensorflow-sigs/build:latest-python3.11",
             name: "JAX Linux x86 CPU Bzlmod",
             repo: "jax-ml/jax",
@@ -97,7 +97,7 @@ jobs:
             repo: "jax-ml/jax",
           },
           {
-            pool: "linux-x86-n2-16",
+            pool: "linux-x86-n4-16",
             container: "gcr.io/tensorflow-sigs/build:latest-python3.12",
             name: "JAX Linux x86 GPU T4",
             repo: "jax-ml/jax",
@@ -150,90 +150,3 @@ jobs:
           else
             python3 "$GITHUB_WORKSPACE"/openxla/xla/build_tools/ci/build.py --build="${{ matrix.job_info.name }}_github_actions"
           fi
-
-  Tests_v2:
-    strategy:
-      # Don't fail fast - want to see results for all builds even if one fails.
-      fail-fast: false
-      matrix:
-        job_info: [
-          {
-            pool: "linux-x86-n4-16",
-            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
-            name: "XLA Linux x86 CPU",
-            repo: "openxla/xla",
-          },
-          {
-            pool: "windows-x86-n2-16",
-            container: null,
-            name: "XLA Windows x86 CPU",
-            repo: "openxla/xla",
-          },
-          {
-            pool: "linux-x86-n4-16",
-            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
-            name: "XLA Linux x86 GPU T4",
-            repo: "openxla/xla",
-          },
-          {
-            pool: "linux-arm64-c4a-16",
-            container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-arm64:latest",
-            name: "XLA Linux ARM64 CPU",
-            repo: "openxla/xla",
-          },
-          {
-            pool: "linux-x86-n4-16",
-            container: "gcr.io/tensorflow-sigs/build:latest-python3.11",
-            name: "JAX Linux x86 CPU Bzlmod",
-            repo: "jax-ml/jax",
-          },
-        ]
-    name: "${{ matrix.job_info.name }} (v2)"
-    runs-on:
-      group: ml-actions-platform
-      labels: ${{ matrix.job_info.pool }}
-    container: ${{ matrix.job_info.container }} # zizmor: ignore[unpinned-images]
-    defaults:
-      run:
-        shell: bash
-    steps:
-      - name: "Checking out openxla/xla"
-        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
-        with:
-          path: "openxla/xla"
-          fetch-depth: 0
-          persist-credentials: false
-      - name: Checking out ${{ matrix.job_info.repo }}
-        if: ${{ matrix.job_info.repo != 'openxla/xla' }}
-        uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1
-        with:
-          repository: ${{ matrix.job_info.repo }}
-          path: ${{ matrix.job_info.repo }}
-          persist-credentials: false
-      - name: "Set Bazel version for JAX"
-        if: ${{ matrix.job_info.repo == 'jax-ml/jax' }}
-        shell: bash
-        run: |
-          BAZEL_VERSION=$(tr -d '\r\n' < "$GITHUB_WORKSPACE/jax-ml/jax/.bazelversion")
-          echo "USE_BAZEL_VERSION=${BAZEL_VERSION}" >> "$GITHUB_ENV"
-      - name: "Wait For Connection"
-        uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
-        with:
-          halt-dispatch-input: ${{ inputs.halt-for-connection }}
-      - name: "Run build.py"
-        working-directory: ${{ matrix.job_info.repo }}
-        timeout-minutes: 60
-        env:
-          PYTHONPATH: ${{ github.workspace }}/openxla/xla
-          XLA_CI_BAZEL_DIFF_BASE_SHA: ${{ github.event.pull_request.base.sha }}
-          XLA_CI_BAZEL_DIFF_HEAD_SHA: ${{ github.event.pull_request.head.sha || github.sha }}
-          XLA_CI_DISABLE_BAZEL_DIFF: ${{ contains(github.event.pull_request.labels.*.name, 'ci:disable-bazel-diff') }}
-        run: |
-          # Shellcheck doesn't understand GitHub Actions expressions and incorrectly
-          # thinks this comparison will always fail.
-          # shellcheck disable=SC2193
-          if [[ "${{ matrix.job_info.pool }}" == *"windows"* ]]; then
-            python "$GITHUB_WORKSPACE"\\openxla\\xla\\build_tools\\ci\\build.py --build="${{ matrix.job_info.name }}_github_actions"
-          else
-            python3 "$GITHUB_WORKSPACE"/openxla/xla/build_tools/ci/build.py --build="${{ matrix.job_info.name }}_github_actions"
-          fi
diff --git a/third_party/xla/.github/workflows/clang_tidy.yml b/third_party/xla/.github/workflows/clang_tidy.yml
index 1cfa192e5a7..b5c72af0af4 100644
--- a/third_party/xla/.github/workflows/clang_tidy.yml
+++ b/third_party/xla/.github/workflows/clang_tidy.yml
@@ -19,7 +19,7 @@ on:
   pull_request:
 jobs:
   clang-tidy-cuda:
-    runs-on: "linux-x86-n2-16"
+    runs-on: "linux-x86-n4-16"
     container:
       image: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"  # zizmor: ignore[unpinned-images]
     defaults:
diff --git a/third_party/xla/docs/rocm_ci.md b/third_party/xla/docs/rocm_ci.md
index 86af8fe466b..4a34650a58b 100644
--- a/third_party/xla/docs/rocm_ci.md
+++ b/third_party/xla/docs/rocm_ci.md
@@ -371,7 +371,7 @@ The runner pool and container for this job are:
 
 ```yaml
 {
-  pool: "linux-x86-n2-16",
+  pool: "linux-x86-n4-16",
   container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest",
   name: "XLA Linux x86 GPU ROCm",
   repo: "openxla/xla",

From 791022f88dd242358165993e8c06f4fef7c553cd Mon Sep 17 00:00:00 2001
From: "Patrick C. Toulme" <ptoulme@google.com>
Date: Thu, 8 Oct 2026 13:00:38 -0700
Subject: [PATCH 0315/1000] [XLA:CPU] Do not pace dispatches made from inside a
 host callback.

A host callback that executed a computation on its own CPU device deadlocked
when the computations queued behind the callback's program held every pacing
reservation. Single device dispatches from inside a HostCallbackScope no
longer take one.

PiperOrigin-RevId: 996018941
---
 third_party/xla/xla/pjrt/cpu/BUILD            |  27 ++++
 third_party/xla/xla/pjrt/cpu/cpu_client.cc    |   6 +-
 .../pjrt/cpu/cpu_client_host_callback_test.cc | 144 ++++++++++++++++++
 3 files changed, 175 insertions(+), 2 deletions(-)
 create mode 100644 third_party/xla/xla/pjrt/cpu/cpu_client_host_callback_test.cc

diff --git a/third_party/xla/xla/pjrt/cpu/BUILD b/third_party/xla/xla/pjrt/cpu/BUILD
index 78a573b4198..05fd6604bd7 100644
--- a/third_party/xla/xla/pjrt/cpu/BUILD
+++ b/third_party/xla/xla/pjrt/cpu/BUILD
@@ -336,6 +336,33 @@ xla_cc_test(
     ],
 )
 
+xla_cc_test(
+    name = "cpu_client_host_callback_test",
+    srcs = ["cpu_client_host_callback_test.cc"],
+    deps = [
+        ":cpu_client",
+        "//xla:literal",
+        "//xla:literal_util",
+        "//xla:shape_util",
+        "//xla:xla_data_proto_cc",
+        "//xla/ffi/api:ffi",
+        "//xla/hlo/builder:xla_computation",
+        "//xla/hlo/parser:hlo_parser",
+        "//xla/pjrt:host_callback",
+        "//xla/pjrt:pjrt_client",
+        "//xla/pjrt:pjrt_executable",
+        "//xla/pjrt/plugin/xla_cpu:cpu_client_options",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/time",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "cpu_async_execution_tracker",
     srcs = ["cpu_async_execution_tracker.cc"],
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
index a236f1ff5b3..9483a4c67b8 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
@@ -1494,10 +1494,12 @@ PjRtRawLoadedExecutable::RawExecuteResult CpuPjRtRawLoadedExecutable::Execute(
   // pacing to avoid problems such as memory fragmentation and running ahead
   // too far, not for correctness. Placing it before the executable launch
   // allows the inputs for the next executable to be fetched even if the
-  // launch is delayed.
+  // launch is delayed. A dispatch from inside a host callback takes no units:
+  // the computations holding them may be queued behind the computation that
+  // runs the callback, so waiting for a unit can deadlock.
   auto compute_reservation = std::make_unique<Semaphore::ScopedReservation>(
       local_device_state->max_inflight_computations_semaphore().ScopedAcquire(
-          1));
+          ThisThreadIsInsideHostCallback() ? 0 : 1));
 
   ExecutableRunOptions run_options;
   run_options.set_run_id(run_id_);
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client_host_callback_test.cc b/third_party/xla/xla/pjrt/cpu/cpu_client_host_callback_test.cc
new file mode 100644
index 00000000000..62d8b5fcc82
--- /dev/null
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client_host_callback_test.cc
@@ -0,0 +1,144 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+// Separate from cpu_client_test.cc: host_callback.h includes the external FFI
+// API, which JAX's CPU callbacks use and which cannot be mixed with
+// xla/ffi/ffi.h.
+
+#include <memory>
+#include <string>
+#include <utility>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/synchronization/notification.h"
+#include "absl/time/time.h"
+#include "xla/ffi/api/ffi.h"
+#include "xla/hlo/builder/xla_computation.h"
+#include "xla/hlo/parser/hlo_parser.h"
+#include "xla/literal.h"
+#include "xla/literal_util.h"
+#include "xla/pjrt/cpu/cpu_client.h"
+#include "xla/pjrt/host_callback.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/pjrt/pjrt_executable.h"
+#include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h"
+#include "xla/shape_util.h"
+#include "xla/xla_data.pb.h"
+
+namespace xla {
+namespace {
+
+// Executed by the host callback below on the callback's thread, as JAX's
+// Python CPU callbacks do. Set by the test.
+PjRtLoadedExecutable* nested_executable = nullptr;
+
+absl::StatusOr<float> RunNestedExecution() {
+  HostCallbackScope host_callback_scope;
+  // JAX gives each dispatching thread its own execution stream. On the
+  // caller's stream the nested execution would also wait for the computation
+  // queued behind the callback.
+  ExecuteOptions options;
+  options.execution_stream_id = 1;
+  ABSL_ASSIGN_OR_RETURN(auto outputs, nested_executable->Execute({{}}, options));
+  ABSL_ASSIGN_OR_RETURN(std::shared_ptr<Literal> literal,
+                   outputs[0][0]->ToLiteral().Await());
+  return literal->Get<float>({});
+}
+
+ffi::Error ExecuteNested(ffi::AnyBuffer,
+                         ffi::Result<ffi::BufferR0<ffi::F32>> result) {
+  absl::StatusOr<float> value = RunNestedExecution();
+  if (!value.ok()) {
+    return ffi::Error::Internal(std::string(value.status().message()));
+  }
+  *result->typed_data() = *value;
+  return ffi::Error::Success();
+}
+
+XLA_FFI_DEFINE_HANDLER(
+    kExecuteNested, ExecuteNested,
+    ffi::Ffi::Bind().Arg<ffi::AnyBuffer>().Ret<ffi::BufferR0<ffi::F32>>());
+
+XLA_FFI_REGISTER_HANDLER(ffi::GetXlaFfiApi(), "__xla_test$$ExecuteNested",
+                         "Host", kExecuteNested);
+
+absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> Compile(
+    PjRtClient& client, absl::string_view program) {
+  ABSL_ASSIGN_OR_RETURN(auto module, ParseAndReturnUnverifiedModule(program));
+  return client.CompileAndLoad(XlaComputation(module->ToProto()), {});
+}
+
+TEST(CpuClientHostCallbackTest,
+     NestedExecutionDoesNotWaitForComputationsQueuedBehindTheCallback) {
+  CpuClientOptions options;
+  options.max_inflight_computations_per_device = 2;
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtClient> client,
+                       GetPjRtCpuClient(std::move(options)));
+  ASSERT_OK_AND_ASSIGN(auto callback_executable, Compile(*client, R"(
+    HloModule callback
+    ENTRY main {
+      p0 = f32[] parameter(0)
+      ROOT result = f32[] custom-call(p0),
+          custom_call_target="__xla_test$$ExecuteNested",
+          api_version=API_VERSION_TYPED_FFI
+    })"));
+  ASSERT_OK_AND_ASSIGN(auto constant_executable, Compile(*client, R"(
+    HloModule constant
+    ENTRY main {
+      ROOT result = f32[] constant(42)
+    })"));
+  nested_executable = constant_executable.get();
+
+  ASSERT_OK_AND_ASSIGN(
+      PjRtMemorySpace * memory_space,
+      client->addressable_devices()[0]->default_memory_space());
+  ASSERT_OK_AND_ASSIGN(auto transfer_manager,
+                       client->CreateBuffersForAsyncHostToDevice(
+                           {ShapeUtil::MakeShape(F32, {})}, memory_space));
+  std::unique_ptr<PjRtBuffer> pending = transfer_manager->RetrieveBuffer(0);
+
+  // Each dispatch holds one of the two units until its computation finishes.
+  // The callback program waits for its input, and the computation queued
+  // behind it on the same execution stream takes the second unit.
+  ASSERT_OK_AND_ASSIGN(auto callback_result,
+                       callback_executable->Execute({{pending.get()}}, {}));
+  ASSERT_OK_AND_ASSIGN(auto queued_result,
+                       constant_executable->Execute({{}}, {}));
+
+  // Once its input arrives, the callback program runs the nested execution,
+  // which must not wait for the unit held by the queued computation.
+  Literal input = LiteralUtil::CreateR0(0.0f);
+  ASSERT_OK(transfer_manager->TransferLiteralToBuffer(0, input, [] {}));
+  absl::Notification done;
+  callback_result[0][0]->GetReadyFuture().OnReady(
+      [&](absl::Status) { done.Notify(); });
+  // A deadlocked client cannot be destroyed, so crash instead of hanging.
+  CHECK(done.WaitForNotificationWithTimeout(absl::Seconds(60)))
+      << "The host callback's nested execution waited for a unit held by a "
+         "computation queued behind the callback.";
+  ASSERT_OK_AND_ASSIGN(std::shared_ptr<Literal> literal,
+                       callback_result[0][0]->ToLiteral().Await());
+  EXPECT_EQ(literal->Get<float>({}), 42.0f);
+  nested_executable = nullptr;
+}
+
+}  // namespace
+}  // namespace xla

From b4054cf78dddc1dda6c5b8d16c39f1f71774905c Mon Sep 17 00:00:00 2001
From: Rishipal Singh Bhatia <rishipal@google.com>
Date: Thu, 8 Oct 2026 14:56:21 -0700
Subject: [PATCH 0316/1000] Add explicit output-to-operand aliasing support to
 XlaBuilder AsyncStart

For asynchronous computations wrapped in kAsyncStart that exhibit in-place
update semantics (such as asynchronous DynamicUpdateSlice), downstream XLA
compiler passes (AliasInfo, LogicalBufferAnalysis, and Memory Space Assignment)
require explicit knowledge of the relationship between the input operand and
the output buffer.

Without explicit output_to_operand_aliasing annotations on kAsyncStart:
1. AliasInfo is unaware of the wrapped operation's in-place semantics and
   returns empty in-place pairs for the async instruction.
2. LogicalBufferAnalysis treats Output {1} as a distinct buffer allocation,
   forcing duplicate buffer allocation and defensive copies in memory space
   assignment.

This change adds an overload of XlaBuilderFriend::BuildAsyncStart that accepts
explicit output_to_operand_aliasing and forwards it to
HloInstruction::CreateAsyncStart. It also adds unit tests in alias_info_test
and logical_buffer_analysis_test verifying that explicit aliasing correctly
informs AliasInfo and suppresses redundant LogicalBuffer allocation.
PiperOrigin-RevId: 996089970
---
 .../xla/xla/hlo/analysis/alias_info_test.cc   | 43 +++++++++++++++++++
 .../analysis/logical_buffer_analysis_test.cc  | 43 +++++++++++++++++++
 .../xla/xla/hlo/builder/xla_builder.cc        | 27 ++++++++++++
 third_party/xla/xla/hlo/builder/xla_builder.h |  7 +++
 .../xla/xla/hlo/builder/xla_builder_test.cc   | 28 ++++++++++++
 5 files changed, 148 insertions(+)

diff --git a/third_party/xla/xla/hlo/analysis/alias_info_test.cc b/third_party/xla/xla/hlo/analysis/alias_info_test.cc
index d6412f28fb5..ec0e2e15484 100644
--- a/third_party/xla/xla/hlo/analysis/alias_info_test.cc
+++ b/third_party/xla/xla/hlo/analysis/alias_info_test.cc
@@ -1201,5 +1201,48 @@ ENTRY main {
   EXPECT_FALSE(alias_info.MustAlias(p1, {}, dus, {}));
 }
 
+// Verifies that without output_to_operand_aliasing, an AsyncStart wrapping a
+// DUS returns empty in-place pairs (treated as out-of-place). This will result
+// in a new buffer for the result of the async dus and trigger copy insertion.
+// On the other hand, adding the attribute correctly confirms the in-place
+// aliasing with Operand 0.
+TEST_F(GetInPlaceInputOutputPairsTest, AsyncDUSAliasingBlindness) {
+  const char* kHlo = R"(
+HloModule test
+
+async_computation {
+  p0 = f32[10] parameter(0)
+  p1 = f32[5] parameter(1)
+  p2 = s32[] parameter(2)
+  ROOT dus = f32[10] dynamic-update-slice(p0, p1, p2)
+}
+
+ENTRY test {
+  p0 = f32[10] parameter(0)
+  p1 = f32[5] parameter(1)
+  p2 = s32[] parameter(2)
+  start_unannotated = ((f32[10], f32[5], s32[]), f32[10], s32[]) async-start(p0, p1, p2),
+      calls=async_computation
+  done_unannotated = f32[10] async-done(start_unannotated)
+  start_annotated = ((f32[10], f32[5], s32[]), f32[10], s32[]) async-start(p0, p1, p2),
+      calls=async_computation, output_to_operand_aliasing={{1}: (0, {})}
+  done_annotated = f32[10] async-done(start_annotated)
+  ROOT tuple = (f32[10], f32[10]) tuple(done_unannotated, done_annotated)
+}
+  )";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo));
+  HloInstruction* unannotated =
+      FindInstruction(module.get(), "start_unannotated");
+  HloInstruction* annotated = FindInstruction(module.get(), "start_annotated");
+
+  // 1. Without attribute, XLA is blind to inner DUS in-place semantics:
+  EXPECT_TRUE(alias_info_.GetInPlaceInputOutputPairs(unannotated).empty());
+
+  // 2. With attribute, XLA recognizes Output {1} aliases Operand 0 in-place:
+  std::vector<std::pair<HloOperandIndex, ShapeIndex>> expected_pairs = {
+      {HloOperandIndex{0, {}}, {1}}};
+  EXPECT_EQ(alias_info_.GetInPlaceInputOutputPairs(annotated), expected_pairs);
+}
+
 }  // namespace
 }  // namespace xla
diff --git a/third_party/xla/xla/hlo/analysis/logical_buffer_analysis_test.cc b/third_party/xla/xla/hlo/analysis/logical_buffer_analysis_test.cc
index b6f4be596c6..e134ee26876 100644
--- a/third_party/xla/xla/hlo/analysis/logical_buffer_analysis_test.cc
+++ b/third_party/xla/xla/hlo/analysis/logical_buffer_analysis_test.cc
@@ -680,6 +680,49 @@ TEST_F(LogicalBufferAnalysisTest, AsyncStartLateBinding) {
   VerifyBufferDefinedAt(async_done, {});
 }
 
+// Verifies that omitting output_to_operand_aliasing on an AsyncStart wrapping
+// DUS causes LogicalBufferAnalysis to mint a brand-new logical buffer for
+// Output {1} (causing memory doubling and copy insertion), whereas annotating
+// it prevents new buffer definition.
+TEST_F(LogicalBufferAnalysisTest, AsyncDUSBufferAllocationWithoutAliasing) {
+  absl::string_view hlo_str = R"(
+  HloModule module
+
+  async_computation {
+    p0 = f32[1024] parameter(0)
+    p1 = f32[256] parameter(1)
+    p2 = s32[] parameter(2)
+    ROOT dus = f32[1024] dynamic-update-slice(p0, p1, p2)
+  }
+
+  ENTRY entry {
+    p0 = f32[1024] parameter(0)
+    p1 = f32[256] parameter(1)
+    p2 = s32[] parameter(2)
+    start_unannotated = ((f32[1024], f32[256], s32[]), f32[1024], s32[]) async-start(p0, p1, p2),
+        calls=async_computation
+    done_unannotated = f32[1024] async-done(start_unannotated)
+    start_annotated = ((f32[1024], f32[256], s32[]), f32[1024], s32[]) async-start(p0, p1, p2),
+        calls=async_computation, output_to_operand_aliasing={ {1}: (0, {}) }
+    done_annotated = f32[1024] async-done(start_annotated)
+    ROOT tuple = (f32[1024], f32[1024]) tuple(done_unannotated, done_annotated)
+  }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str));
+  ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(module.get()));
+
+  HloInstruction* unannotated =
+      FindInstruction(module.get(), "start_unannotated");
+  HloInstruction* annotated = FindInstruction(module.get(), "start_annotated");
+
+  // Output {1} defines a new buffer when unannotated:
+  VerifyBufferDefinedAt(unannotated, {1});
+
+  // Output {1} is aliased to operand 0 when annotated, so NO new buffer is
+  // defined:
+  VerifyNoBufferDefinedAt(annotated, {1});
+}
+
 TEST_F(LogicalBufferAnalysisTest, AsyncStartLateBindingWithAliasing) {
   absl::string_view hlo_str = R"(
   HloModule module
diff --git a/third_party/xla/xla/hlo/builder/xla_builder.cc b/third_party/xla/xla/hlo/builder/xla_builder.cc
index eb01d7b0041..78469053312 100644
--- a/third_party/xla/xla/hlo/builder/xla_builder.cc
+++ b/third_party/xla/xla/hlo/builder/xla_builder.cc
@@ -212,6 +212,33 @@ XlaOp XlaBuilderFriend::BuildAsyncStart(XlaBuilder* builder,
   return start_op;
 }
 
+XlaOp XlaBuilderFriend::BuildAsyncStart(
+    XlaBuilder* builder, absl::Span<const XlaOp> operands,
+    std::string execution_thread, XlaComputationId called_computation,
+    const Shape& shape,
+    absl::Span<const std::pair<ShapeIndex, std::pair<int64_t, ShapeIndex>>>
+        output_operand_aliasing) {
+  auto start_op = builder->ReportErrorOrReturn([&]() -> absl::StatusOr<XlaOp> {
+    HloInstructionProto instr;
+    *instr.mutable_shape() = shape.ToProto();
+    instr.set_async_execution_thread(execution_thread);
+    ABSL_RETURN_IF_ERROR(builder->AddCalledComputation(called_computation, instr));
+    for (const auto& pair : output_operand_aliasing) {
+      OutputOperandAliasing* aliasing = instr.add_output_operand_aliasing();
+      aliasing->set_operand_index(pair.second.first);
+      for (int64_t index : pair.second.second) {
+        aliasing->add_operand_shape_index(index);
+      }
+      for (int64_t index : pair.first) {
+        aliasing->add_output_shape_index(index);
+      }
+    }
+    return builder->AddInstruction(std::move(instr), HloOpcode::kAsyncStart,
+                                   operands);
+  });
+  return start_op;
+}
+
 XlaOp XlaBuilderFriend::BuildAsyncUpdate(XlaBuilder* builder,
                                          const XlaOp operand,
                                          const Shape& shape) {
diff --git a/third_party/xla/xla/hlo/builder/xla_builder.h b/third_party/xla/xla/hlo/builder/xla_builder.h
index 669eb830e1c..de5409b2e78 100644
--- a/third_party/xla/xla/hlo/builder/xla_builder.h
+++ b/third_party/xla/xla/hlo/builder/xla_builder.h
@@ -79,6 +79,13 @@ struct XlaBuilderFriend {
                                std::string execution_thread,
                                XlaComputationId called_computation,
                                const Shape& shape);
+  static XlaOp BuildAsyncStart(
+      XlaBuilder* builder, absl::Span<const XlaOp> operands,
+      std::string execution_thread, XlaComputationId called_computation,
+      const Shape& shape,
+      absl::Span<const std::pair<ShapeIndex, std::pair<int64_t, ShapeIndex>>>
+          output_operand_aliasing);
+
   static XlaOp BuildAsyncUpdate(XlaBuilder* builder, XlaOp operands,
                                 const Shape& shape);
   static XlaOp BuildAsyncDone(XlaBuilder* builder, XlaOp operands,
diff --git a/third_party/xla/xla/hlo/builder/xla_builder_test.cc b/third_party/xla/xla/hlo/builder/xla_builder_test.cc
index 05b30faa5a4..d280a649bf7 100644
--- a/third_party/xla/xla/hlo/builder/xla_builder_test.cc
+++ b/third_party/xla/xla/hlo/builder/xla_builder_test.cc
@@ -4227,5 +4227,33 @@ TEST(XlaBuilderTest, OriginalValue) {
   EXPECT_THAT(*tuple_original_value, expected_tuple_original_value);
 }
 
+TEST(XlaBuilderTest, AsyncStartWithOutputOperandAliasing) {
+  XlaBuilder b("async_start_aliasing");
+  Shape shape = ShapeUtil::MakeShape(F32, {4});
+  XlaOp p0 = Parameter(&b, 0, shape, "p0");
+
+  std::unique_ptr<XlaBuilder> sub_b = b.CreateSubBuilder("subcomp");
+  XlaOp sub_p0 = Parameter(sub_b.get(), 0, shape, "sub_p0");
+  XlaOp sub_neg = Neg(sub_p0);
+  TF_ASSERT_OK_AND_ASSIGN(XlaComputationId sub_comp_id,
+                          sub_b->BuildSubComputation(sub_neg));
+
+  Shape start_shape =
+      ShapeUtil::MakeTupleShape({ShapeUtil::MakeTupleShape({shape}), shape,
+                                 ShapeUtil::MakeShape(S32, {})});
+
+  std::vector<std::pair<ShapeIndex, std::pair<int64_t, ShapeIndex>>> aliasing =
+      {{{1}, {0, {}}}};
+
+  internal::XlaBuilderFriend::BuildAsyncStart(
+      &b, {p0}, HloInstruction::kMainExecutionThread, sub_comp_id, start_shape,
+      aliasing);
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module, BuildHloModule(b));
+  HloInstruction* root = GetRoot(*module);
+  EXPECT_EQ(root->opcode(), HloOpcode::kAsyncStart);
+  EXPECT_EQ(root->output_operand_aliasing(), aliasing);
+}
+
 }  // namespace
 }  // namespace xla

From a8161a017f142890be588d6752f4081df4e27098 Mon Sep 17 00:00:00 2001
From: Yue Sheng <yueshengys@google.com>
Date: Thu, 8 Oct 2026 15:28:45 -0700
Subject: [PATCH 0317/1000] [Mosaic TPU] Remove deprecated `tpu.repeat`.

PiperOrigin-RevId: 996109070
---
 .../xla/xla/mosaic/dialect/tpu/tpu_ops.td     | 12 ---------
 .../mosaic/dialect/tpu/transforms/serde.cc    | 27 +++++++++++++++++++
 2 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td
index 008a9f622a7..f978b8ac117 100644
--- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td
+++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td
@@ -701,18 +701,6 @@ def TPU_ReshapeOp : TPU_Op<"reshape", [Pure]> {
   let hasFolder = 1;
 }
 
-// TODO(mvoz): deprecated - use concat. Canonicalization will do so automatically.
-// b/376295711
-def TPU_RepeatOp : TPU_Op<"repeat", [Pure]> {
-  let arguments = (ins
-    AnyVectorOfNonZeroRank:$source,
-    I32Attr:$dimension,
-    I32Attr:$times
-  );
-  let results = (outs AnyVectorOfNonZeroRank:$output);
-  let assemblyFormat = [{ $source `,` $dimension `x` $times attr-dict `:` type($source) `->` type($output) }];
-}
-
 def TPU_BroadcastInSublanesOp : TPU_Op<"broadcast_in_sublanes", [Pure]> {
   let description = [{
     For each sublane `i`, broadcasts the value in lane `lane + i` along the
diff --git a/third_party/xla/xla/mosaic/dialect/tpu/transforms/serde.cc b/third_party/xla/xla/mosaic/dialect/tpu/transforms/serde.cc
index 5357285acf5..5ea6d4ed089 100644
--- a/third_party/xla/xla/mosaic/dialect/tpu/transforms/serde.cc
+++ b/third_party/xla/xla/mosaic/dialect/tpu/transforms/serde.cc
@@ -811,6 +811,32 @@ LogicalResult matmul_downgrade(Operation* op, int version, bool&) {
   return success();
 }
 
+LogicalResult repeat_upgrade(Operation* op, int version, bool& erased) {
+  if (op->getNumOperands() != 1 || op->getNumResults() != 1) {
+    return op->emitError("Unexpected operand/result count in tpu.repeat");
+  }
+  auto dimension_attr = op->getAttrOfType<IntegerAttr>("dimension");
+  auto times_attr = op->getAttrOfType<IntegerAttr>("times");
+  if (!dimension_attr || !times_attr) {
+    return op->emitError("Missing dimension or times attribute in tpu.repeat");
+  }
+  Value source = op->getOperand(0);
+  int32_t times = times_attr.getInt();
+  if (times == 1) {
+    op->replaceAllUsesWith(ValueRange{source});
+  } else {
+    OpBuilder builder(op);
+    SmallVector<Value> operands(times, source);
+    Value concat =
+        ConcatenateOp::create(builder, op->getLoc(), op->getResult(0).getType(),
+                              operands, dimension_attr.getInt());
+    op->replaceAllUsesWith(ValueRange{concat});
+  }
+  op->erase();
+  erased = true;
+  return success();
+}
+
 const llvm::StringMap<SerdeRuleType>& upgrade_rules() {
   static auto rules = new llvm::StringMap<SerdeRuleType>{
       {EnqueueDMAOp::getOperationName(), enqueue_dma_upgrade},
@@ -827,6 +853,7 @@ const llvm::StringMap<SerdeRuleType>& upgrade_rules() {
       {arith::ConstantOp::getOperationName(), arith_constant_upgrade},
       {ReinterpretCastOp::getOperationName(), reinterpret_cast_upgrade},
       {MatmulOp::getOperationName(), matmul_upgrade},
+      {"tpu.repeat", repeat_upgrade},
   };
   return *rules;
 }

From 1b98ad2bbe4fd1b4fc8e2752ea43685c41ccd184 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 8 Oct 2026 15:57:31 -0700
Subject: [PATCH 0318/1000] Integrate LLVM at llvm/llvm-project@c4d9267a3189

Updates LLVM usage to match
[c4d9267a3189](https://github.com/llvm/llvm-project/commit/c4d9267a31897dd76b2aab5196b75bffb49e6846)

PiperOrigin-RevId: 996124567
---
 .../xla/third_party/llvm/generated.patch      |  458 ----
 .../xla/third_party/llvm/workspace.bzl        |    4 +-
 .../xla/third_party/shardy/temporary.patch    | 1857 +++++++++++++++++
 .../xla/third_party/shardy/workspace.bzl      |    4 +-
 .../xla/third_party/stablehlo/temporary.patch |  367 +++-
 .../cpu/codegen/kernel_api_ir_builder_test.cc |    4 +-
 .../emitters/transforms/lower_to_llvm_gpu.cc  |   13 +-
 .../transforms/gpu_kernel_lowering_passes.cc  |   12 +-
 8 files changed, 2208 insertions(+), 511 deletions(-)

diff --git a/third_party/xla/third_party/llvm/generated.patch b/third_party/xla/third_party/llvm/generated.patch
index e2237d49499..509398da979 100644
--- a/third_party/xla/third_party/llvm/generated.patch
+++ b/third_party/xla/third_party/llvm/generated.patch
@@ -1,459 +1 @@
 Auto generated patch. Do not edit or delete it, even if empty.
-diff -ruN --strip-trailing-cr a/llvm/include/llvm/IR/Metadata.h b/llvm/include/llvm/IR/Metadata.h
---- a/llvm/include/llvm/IR/Metadata.h
-+++ b/llvm/include/llvm/IR/Metadata.h
-@@ -397,18 +397,8 @@
-   using OwnerTy = MetadataTracking::OwnerTy;
- 
- private:
--  struct UseEntry {
--    void *Ref = nullptr;
--    OwnerTy Owner = nullptr;
--  };
--
--  static constexpr unsigned IndexThreshold = 32;
--  // Tracked uses of this. Dropping one moves the last entry into its slot, so
--  // the order is deterministic but not the order they were added in.
--  SmallVector<UseEntry, 4> UseMap;
--  // Lazily allocated map from Ref to its index in UseMap for large use lists.
--  using IndexMapTy = DenseMap<void *, unsigned>;
--  std::unique_ptr<IndexMapTy> IndexMap;
-+  uint64_t NextIndex = 0;
-+  SmallDenseMap<void *, std::pair<OwnerTy, uint64_t>, 4> UseMap;
- 
- protected:
-   ~ReplaceableUses() {
-@@ -439,7 +429,6 @@
-   unsigned getNumUses() const { return UseMap.size(); }
- 
- private:
--  UseEntry *findRef(void *Ref);
-   void addRef(void *Ref, OwnerTy Owner);
-   void dropRef(void *Ref);
-   void moveRef(void *Ref, void *New, const Metadata &MD);
-diff -ruN --strip-trailing-cr a/llvm/lib/IR/Metadata.cpp b/llvm/lib/IR/Metadata.cpp
---- a/llvm/lib/IR/Metadata.cpp
-+++ b/llvm/lib/IR/Metadata.cpp
-@@ -258,96 +258,82 @@
- }
- 
- SmallVector<Metadata *> ReplaceableUses::getAllArgListUsers() {
--  SmallVector<Metadata *> MDUsers;
--  for (const auto &[Ref, Owner] : UseMap) {
-+  SmallVector<std::pair<OwnerTy, uint64_t> *> MDUsersWithID;
-+  for (auto Pair : UseMap) {
-+    OwnerTy Owner = Pair.second.first;
-     if (Owner.isNull())
-       continue;
-     if (!isa<Metadata *>(Owner))
-       continue;
-     Metadata *OwnerMD = cast<Metadata *>(Owner);
-     if (OwnerMD->getMetadataID() == Metadata::DIArgListKind)
--      MDUsers.push_back(OwnerMD);
-+      MDUsersWithID.push_back(&UseMap[Pair.first]);
-   }
-+  llvm::sort(MDUsersWithID, [](auto UserA, auto UserB) {
-+    return UserA->second < UserB->second;
-+  });
-+  SmallVector<Metadata *> MDUsers;
-+  for (auto *UserWithID : MDUsersWithID)
-+    MDUsers.push_back(cast<Metadata *>(UserWithID->first));
-   return MDUsers;
- }
- 
- SmallVector<DbgVariableRecord *>
- ReplaceableUses::getAllDbgVariableRecordUsers() {
-+  SmallVector<std::pair<OwnerTy, uint64_t> *> DVRUsersWithID;
-+  for (auto Pair : UseMap) {
-+    OwnerTy Owner = Pair.second.first;
-+    if (Owner.isNull())
-+      continue;
-+    if (!isa<DebugValueUser *>(Owner))
-+      continue;
-+    DVRUsersWithID.push_back(&UseMap[Pair.first]);
-+  }
-   // Order DbgVariableRecord users in reverse-creation order. Normal dbg.value
-   // users of MetadataAsValues are ordered by their UseList, i.e. reverse order
-   // of when they were added: we need to replicate that here. The structure of
-   // debug-info output depends on the ordering of intrinsics, thus we need
-   // to keep them consistent for comparisons sake.
-+  llvm::sort(DVRUsersWithID, [](auto UserA, auto UserB) {
-+    return UserA->second > UserB->second;
-+  });
-   SmallVector<DbgVariableRecord *> DVRUsers;
--  for (const auto &[Ref, Owner] : reverse(UseMap)) {
--    if (Owner.isNull())
--      continue;
--    if (!isa<DebugValueUser *>(Owner))
--      continue;
--    DVRUsers.push_back(cast<DebugValueUser *>(Owner)->getUser());
--  }
-+  for (auto UserWithID : DVRUsersWithID)
-+    DVRUsers.push_back(cast<DebugValueUser *>(UserWithID->first)->getUser());
-   return DVRUsers;
- }
- 
--ReplaceableUses::UseEntry *ReplaceableUses::findRef(void *Ref) {
--  if (!IndexMap) {
--    // Search backward for temporal locality; recently added references are
--    // often dropped or moved first.
--    for (UseEntry &U : reverse(UseMap))
--      if (U.Ref == Ref)
--        return &U;
--    return nullptr;
--  }
--
--  auto It = IndexMap->find(Ref);
--  return It == IndexMap->end() ? nullptr : &UseMap[It->second];
--}
--
- void ReplaceableUses::addRef(void *Ref, OwnerTy Owner) {
--  assert(Ref && "Expected live reference");
--  assert(!findRef(Ref) && "Reference already tracked");
--  unsigned NewIdx = UseMap.size();
--  UseMap.push_back({Ref, Owner});
--  if (IndexMap) {
--    (*IndexMap)[Ref] = NewIdx;
--  } else if (UseMap.size() > IndexThreshold) {
--    // Build the index map once UseMap grows past the threshold.
--    IndexMap = std::make_unique<IndexMapTy>();
--    for (unsigned I = 0, E = UseMap.size(); I != E; ++I)
--      (*IndexMap)[UseMap[I].Ref] = I;
--  }
-+  bool WasInserted =
-+      UseMap.insert(std::make_pair(Ref, std::make_pair(Owner, NextIndex)))
-+          .second;
-+  (void)WasInserted;
-+  assert(WasInserted && "Expected to add a reference");
-+
-+  ++NextIndex;
-+  assert(NextIndex != 0 && "Unexpected overflow");
- }
- 
- void ReplaceableUses::dropRef(void *Ref) {
--  UseEntry *Entry = findRef(Ref);
--  assert(Entry && "Expected to find Ref");
--  if (IndexMap)
--    IndexMap->erase(Ref);
--  if (Entry != &UseMap.back()) {
--    *Entry = UseMap.back();
--    if (IndexMap)
--      (*IndexMap)[Entry->Ref] = Entry - UseMap.begin();
--  }
--  UseMap.pop_back();
--  if (IndexMap && UseMap.size() <= IndexThreshold / 2)
--    IndexMap.reset();
-+  bool WasErased = UseMap.erase(Ref);
-+  (void)WasErased;
-+  assert(WasErased && "Expected to drop a reference");
- }
- 
- void ReplaceableUses::moveRef(void *Ref, void *New, const Metadata &MD) {
--  assert(!findRef(New) && "Cannot move to an existing reference");
--  UseEntry *Entry = findRef(Ref);
--  assert(Entry && "Expected to move a reference");
--  if (IndexMap) {
--    IndexMap->erase(Ref);
--    (*IndexMap)[New] = Entry - UseMap.begin();
--  }
--  Entry->Ref = New;
-+  auto I = UseMap.find(Ref);
-+  assert(I != UseMap.end() && "Expected to move a reference");
-+  auto OwnerAndIndex = I->second;
-+  UseMap.erase(I);
-+  bool WasInserted = UseMap.insert(std::make_pair(New, OwnerAndIndex)).second;
-+  (void)WasInserted;
-+  assert(WasInserted && "Expected to add a reference");
- 
-   // Check that the references are direct if there's no owner.
-   (void)MD;
--  assert((Entry->Owner || *static_cast<Metadata **>(Ref) == &MD) &&
-+  assert((OwnerAndIndex.first || *static_cast<Metadata **>(Ref) == &MD) &&
-          "Reference without owner must be direct");
--  assert((Entry->Owner || *static_cast<Metadata **>(New) == &MD) &&
-+  assert((OwnerAndIndex.first || *static_cast<Metadata **>(New) == &MD) &&
-          "Reference without owner must be direct");
- }
- 
-@@ -360,11 +346,14 @@
-   auto &Store = Context.pImpl->ValuesAsMetadata;
-   auto I = Store.find(&C);
-   ValueAsMetadata *MD = I->second;
-+  using UseTy =
-+      std::pair<void *, std::pair<MetadataTracking::OwnerTy, uint64_t>>;
-   // Copy out uses and update value of Constant used by debug info metadata with
--  // poison below.
--  SmallVector<UseEntry, 4> Uses = MD->UseMap;
-+  // poison below
-+  SmallVector<UseTy, 8> Uses(MD->UseMap.begin(), MD->UseMap.end());
- 
--  for (const auto &[Ref, Owner] : Uses) {
-+  for (const auto &Pair : Uses) {
-+    MetadataTracking::OwnerTy Owner = Pair.second.first;
-     if (!Owner)
-       continue;
-     // Check for MetadataAsValue.
-@@ -380,7 +369,7 @@
-       continue;
-     if (isa<DINode>(OwnerMD)) {
-       OwnerMD->handleChangedOperand(
--          Ref, ValueAsMetadata::get(PoisonValue::get(C.getType())));
-+          Pair.first, ValueAsMetadata::get(PoisonValue::get(C.getType())));
-     }
-   }
- }
-@@ -390,21 +379,25 @@
-     return;
- 
-   // Copy out uses since UseMap will get touched below.
--  SmallVector<UseEntry, 4> Uses = UseMap;
--
--  for (const auto &[Ref, Owner] : Uses) {
-+  using UseTy = std::pair<void *, std::pair<OwnerTy, uint64_t>>;
-+  SmallVector<UseTy, 8> Uses(UseMap.begin(), UseMap.end());
-+  llvm::sort(Uses, [](const UseTy &L, const UseTy &R) {
-+    return L.second.second < R.second.second;
-+  });
-+  for (const auto &Pair : Uses) {
-     // Check that this Ref hasn't disappeared after RAUW (when updating a
-     // previous Ref).
--    if (!findRef(Ref))
-+    if (!UseMap.count(Pair.first))
-       continue;
- 
-+    OwnerTy Owner = Pair.second.first;
-     if (!Owner) {
-       // Update unowned tracking references directly.
--      Metadata *&DirectRef = *static_cast<Metadata **>(Ref);
--      dropRef(Ref);
--      DirectRef = MD;
-+      Metadata *&Ref = *static_cast<Metadata **>(Pair.first);
-+      Ref = MD;
-       if (MD)
--        MetadataTracking::track(DirectRef);
-+        MetadataTracking::track(Ref);
-+      UseMap.erase(Pair.first);
-       continue;
-     }
- 
-@@ -415,7 +408,7 @@
-     }
- 
-     if (auto *DVU = dyn_cast<DebugValueUser *>(Owner)) {
--      DVU->handleChangedValue(Ref, MD);
-+      DVU->handleChangedValue(Pair.first, MD);
-       continue;
-     }
- 
-@@ -424,7 +417,7 @@
-     switch (OwnerMD->getMetadataID()) {
- #define HANDLE_METADATA_LEAF(CLASS)                                            \
-   case Metadata::CLASS##Kind:                                                  \
--    cast<CLASS>(OwnerMD)->handleChangedOperand(Ref, MD);                       \
-+    cast<CLASS>(OwnerMD)->handleChangedOperand(Pair.first, MD);                \
-     continue;
- #include "llvm/IR/Metadata.def"
-     default:
-@@ -440,16 +433,18 @@
- 
-   if (!ResolveUsers) {
-     UseMap.clear();
--    IndexMap.reset();
-     return;
-   }
- 
--  // Move uses out since UseMap could get touched below.
--  SmallVector<UseEntry, 4> Uses = std::move(UseMap);
-+  // Copy out uses since UseMap could get touched below.
-+  using UseTy = std::pair<void *, std::pair<OwnerTy, uint64_t>>;
-+  SmallVector<UseTy, 8> Uses(UseMap.begin(), UseMap.end());
-+  llvm::sort(Uses, [](const UseTy &L, const UseTy &R) {
-+    return L.second.second < R.second.second;
-+  });
-   UseMap.clear();
--  IndexMap.reset();
--  for (const auto &U : Uses) {
--    auto Owner = U.Owner;
-+  for (const auto &Pair : Uses) {
-+    auto Owner = Pair.second.first;
-     if (!Owner)
-       continue;
-     if (!isa<Metadata *>(Owner))
-diff -ruN --strip-trailing-cr a/llvm/lib/ObjCopy/ELF/ELFObject.cpp b/llvm/lib/ObjCopy/ELF/ELFObject.cpp
---- a/llvm/lib/ObjCopy/ELF/ELFObject.cpp
-+++ b/llvm/lib/ObjCopy/ELF/ELFObject.cpp
-@@ -1709,11 +1709,7 @@
-   case SHT_REL:
-   case SHT_RELA:
-   case SHT_CREL:
--    // SHF_ALLOC relocations of an executable or shared object are copied
--    // verbatim. Relocatable files relocations are usually non-ALLOC, but Linux
--    // livepatch modules set SHF_ALLOC on .klp.rela.* to keep these static
--    // relocations in memory.
--    if ((Shdr.sh_flags & SHF_ALLOC) && ElfFile.getHeader().e_type != ET_REL) {
-+    if (Shdr.sh_flags & SHF_ALLOC) {
-       if (Expected<ArrayRef<uint8_t>> Data = ElfFile.getSectionContents(Shdr))
-         return Obj.addSection<DynamicRelocationSection>(*Data);
-       else
-diff -ruN --strip-trailing-cr a/llvm/lib/ObjCopy/ELF/ELFObject.h b/llvm/lib/ObjCopy/ELF/ELFObject.h
---- a/llvm/lib/ObjCopy/ELF/ELFObject.h
-+++ b/llvm/lib/ObjCopy/ELF/ELFObject.h
-@@ -38,7 +38,6 @@
- class SymbolTableSection;
- class RelocationSection;
- class DynamicRelocationSection;
--class DynamicSymbolTableSection;
- class GnuDebugLinkSection;
- class GroupSection;
- class SectionIndexSection;
-@@ -876,15 +875,10 @@
- // and another which handles the symbol table type. The symbol table type is
- // taken as a type parameter to the class (see RelocSectionWithSymtabBase).
- class RelocationSectionBase : public SectionBase {
--  const bool Dynamic;
--
- protected:
-   SectionBase *SecToApplyRel = nullptr;
- 
--  explicit RelocationSectionBase(bool Dynamic) : Dynamic(Dynamic) {}
--
- public:
--  bool isDynamic() const { return Dynamic; }
-   const SectionBase *getSection() const { return SecToApplyRel; }
-   void setSection(SectionBase *Sec) { SecToApplyRel = Sec; }
- 
-@@ -903,9 +897,7 @@
-   void setSymTab(SymTabType *SymTab) { Symbols = SymTab; }
- 
- protected:
--  RelocSectionWithSymtabBase()
--      : RelocationSectionBase(
--            std::is_same_v<SymTabType, DynamicSymbolTableSection>) {}
-+  RelocSectionWithSymtabBase() = default;
- 
-   SymTabType *Symbols = nullptr;
- 
-@@ -936,8 +928,9 @@
-   const Object &getObject() const { return Obj; }
- 
-   static bool classof(const SectionBase *S) {
--    return RelocationSectionBase::classof(S) &&
--           !static_cast<const RelocationSectionBase *>(S)->isDynamic();
-+    if (S->OriginalFlags & ELF::SHF_ALLOC)
-+      return false;
-+    return RelocationSectionBase::classof(S);
-   }
- };
- 
-@@ -1022,8 +1015,9 @@
-       function_ref<bool(const SectionBase *)> ToRemove) override;
- 
-   static bool classof(const SectionBase *S) {
--    return RelocationSectionBase::classof(S) &&
--           static_cast<const RelocationSectionBase *>(S)->isDynamic();
-+    if (!(S->OriginalFlags & ELF::SHF_ALLOC))
-+      return false;
-+    return S->OriginalType == ELF::SHT_REL || S->OriginalType == ELF::SHT_RELA;
-   }
- };
- 
-diff -ruN --strip-trailing-cr a/llvm/test/tools/llvm-objcopy/ELF/reloc-alloc-symtab.test b/llvm/test/tools/llvm-objcopy/ELF/reloc-alloc-symtab.test
---- a/llvm/test/tools/llvm-objcopy/ELF/reloc-alloc-symtab.test
-+++ b/llvm/test/tools/llvm-objcopy/ELF/reloc-alloc-symtab.test
-@@ -1,42 +0,0 @@
--## In a relocatable file, an SHF_ALLOC relocation section (e.g. .klp.rela.* in
--## Linux livepatch modules) is a static relocation section: its symbol indexes
--## are updated and the symbols it references are kept.
--
--# RUN: yaml2obj %s -o %t
--# RUN: llvm-strip --strip-unneeded %t -o %t1
--# RUN: llvm-readelf -Sr -s %t1 | FileCheck %s
--# RUN: yaml2obj -DTYPE=SHT_CREL %s -o %t.crel
--# RUN: llvm-strip --strip-unneeded %t.crel -o %t1.crel
--# RUN: llvm-readelf -Sr -s %t1.crel | FileCheck %s
--
--# CHECK:      [ 2] .klp.rela.vmlinux.text {{RELA|CREL}} {{.*}} AI  3   1
--# CHECK-NEXT: [ 3] .symtab
--# CHECK:      R_X86_64_PC32 {{.*}} referenced + 0
--# CHECK:      Symbol table '.symtab' contains 2 entries:
--# CHECK:      1: {{.*}} UND referenced
--
--!ELF
--FileHeader:
--  Class:   ELFCLASS64
--  Data:    ELFDATA2LSB
--  Type:    ET_REL
--  Machine: EM_X86_64
--Sections:
--  - Name:    .text
--    Type:    SHT_PROGBITS
--    Flags:   [ SHF_ALLOC, SHF_EXECINSTR ]
--    Content: "e800000000"
--  - Name:    .klp.rela.vmlinux.text
--    Type:    [[TYPE=SHT_RELA]]
--    Flags:   [ SHF_ALLOC, SHF_INFO_LINK ]
--    Link:    .symtab
--    Info:    .text
--    Relocations:
--      - Offset: 1
--        Symbol: referenced
--        Type:   R_X86_64_PC32
--Symbols:
--  - Name:    unreferenced
--    Section: .text
--  - Name:    referenced
--    Binding: STB_GLOBAL
-diff -ruN --strip-trailing-cr a/llvm/test/tools/llvm-objcopy/ELF/reloc-no-symtab.test b/llvm/test/tools/llvm-objcopy/ELF/reloc-no-symtab.test
---- a/llvm/test/tools/llvm-objcopy/ELF/reloc-no-symtab.test
-+++ b/llvm/test/tools/llvm-objcopy/ELF/reloc-no-symtab.test
-@@ -10,16 +10,11 @@
- # RUN: llvm-readobj --sections %t.original | FileCheck %s
- # RUN: llvm-readobj --sections %t.stripped | FileCheck %s
- 
--## In a relocatable file, .rela.plt is a static relocation section.
--# RUN: yaml2obj -DTYPE=ET_REL %s -o %t.rel
--# RUN: llvm-strip %t.rel -o %t.rel.stripped
--# RUN: llvm-readobj --sections %t.rel.stripped | FileCheck %s
--
- !ELF
- FileHeader:
-   Class:           ELFCLASS64
-   Data:            ELFDATA2LSB
--  Type:            [[TYPE=ET_EXEC]]
-+  Type:            ET_REL
-   Machine:         EM_X86_64
- Sections:
-   - Name:          .text
-diff -ruN --strip-trailing-cr a/llvm/unittests/IR/MetadataTest.cpp b/llvm/unittests/IR/MetadataTest.cpp
---- a/llvm/unittests/IR/MetadataTest.cpp
-+++ b/llvm/unittests/IR/MetadataTest.cpp
-@@ -5355,29 +5355,6 @@
-   EXPECT_TRUE(!MD);
- }
- 
--// Once the use list is indexed, look each changed reference up again right
--// away: a later swap in dropRef would repair a stale index entry.
--TEST_F(TrackingMDRefTest, LargeUseList) {
--  auto Temp = MDTuple::getTemporary(Context, {});
--  TrackingMDRef Refs[41];
--  for (unsigned I = 0; I != 40; ++I)
--    Refs[I].reset(Temp.get());
--  Refs[0].reset();
--  Refs[0].reset(Temp.get());
--  Refs[0].reset();
--  Refs[40] = std::move(Refs[1]);
--  Refs[1].reset(Temp.get());
--  Refs[40].reset();
--  Refs[39].reset();
--  EXPECT_EQ(38u, Temp->getNumTemporaryUses());
--
--  auto *N = MDTuple::getDistinct(Context, {});
--  Temp->replaceAllUsesWith(N);
--  EXPECT_EQ(0u, Temp->getNumTemporaryUses());
--  for (unsigned I = 1; I != 39; ++I)
--    EXPECT_EQ(N, Refs[I].get()) << I;
--}
--
- TEST(NamedMDNodeTest, Search) {
-   LLVMContext Context;
-   ConstantAsMetadata *C =
diff --git a/third_party/xla/third_party/llvm/workspace.bzl b/third_party/xla/third_party/llvm/workspace.bzl
index 9de7c05a114..7b8ac881c64 100644
--- a/third_party/xla/third_party/llvm/workspace.bzl
+++ b/third_party/xla/third_party/llvm/workspace.bzl
@@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "62e30e3326d9b4e3948d242ccfee7aa4fc030ef3"
-    LLVM_SHA256 = "478ae34c463f28708582082098aaf0e79f431c2719f5d297bd546051dbeae0dd"
+    LLVM_COMMIT = "c4d9267a31897dd76b2aab5196b75bffb49e6846"
+    LLVM_SHA256 = "a96c8b33f0d8d3be008b8c9f74ef310254c4f7f7edf8c04c76f4d99ed9d34f1a"
 
     tf_http_archive(
         name = name,
diff --git a/third_party/xla/third_party/shardy/temporary.patch b/third_party/xla/third_party/shardy/temporary.patch
index e69de29bb2d..2676e00bb2d 100644
--- a/third_party/xla/third_party/shardy/temporary.patch
+++ b/third_party/xla/third_party/shardy/temporary.patch
@@ -0,0 +1,1857 @@
+diff --git a/shardy/dialect/sdy/ir/canonicalization.cc b/shardy/dialect/sdy/ir/canonicalization.cc
+index 9af43e32..f4cda6da 100644
+--- a/shardy/dialect/sdy/ir/canonicalization.cc
++++ b/shardy/dialect/sdy/ir/canonicalization.cc
+@@ -108,7 +108,7 @@ class RedundantManualComputationPattern
+     if (!inShardings.empty() || !outShardings.empty()) {
+       MeshAttr mesh =
+           getCommonMesh(inShardings, outShardings, manualComputationOp);
+-      SDY_CHECK(mesh) << "expected inputs and outputs to have a common mesh";
++      assert(mesh && "expected inputs and outputs to have a common mesh");
+       for (StringAttr manualAxis : manualComputationOp.getManualAxes()) {
+         manualAxesProduct *= mesh.getAxisSize(manualAxis);
+       }
+diff --git a/shardy/dialect/sdy/ir/test/manual_computation_verification.mlir b/shardy/dialect/sdy/ir/test/manual_computation_verification.mlir
+index c2c24800..5853611e 100644
+--- a/shardy/dialect/sdy/ir/test/manual_computation_verification.mlir
++++ b/shardy/dialect/sdy/ir/test/manual_computation_verification.mlir
+@@ -294,28 +294,3 @@ func.func @ranked_sharding_on_token(%arg0: !stablehlo.token) -> !stablehlo.token
+   } : (!stablehlo.token) -> !stablehlo.token
+   return %0 : !stablehlo.token
+ }
+-
+-// -----
+-
+-sdy.mesh @mesh = <["a"=2]>
+-
+-func.func @scalar_region_arg(%arg0: tensor<16x32xf32>) -> tensor<16x32xf32> {
+-  // expected-error @+1 {{op operand shape, corresponding sharding, and region operand shape at index 0 must match. Expected local shape 'tensor<8x32xf32>', actual local shape 'f32'}}
+-  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{"a"}, {}]>] out_shardings=[<@mesh, [{"a"}, {}]>] manual_axes={"a"} (%arg1: f32) {
+-    %c = stablehlo.constant dense<0.000000e+00> : tensor<8x32xf32>
+-    sdy.return %c : tensor<8x32xf32>
+-  } : (tensor<16x32xf32>) -> tensor<16x32xf32>
+-  func.return %0: tensor<16x32xf32>
+-}
+-
+-// -----
+-
+-sdy.mesh @mesh = <["a"=2]>
+-
+-func.func @token_region_result(%arg0: tensor<16x32xf32>, %arg1: !stablehlo.token) -> tensor<16x32xf32> {
+-  // expected-error @+1 {{op result shape, corresponding sharding, and region result shape at index 0 must match. Expected local shape 'tensor<8x32xf32>', actual local shape '!stablehlo.token'}}
+-  %0 = sdy.manual_computation(%arg0, %arg1) in_shardings=[<@mesh, [{"a"}, {}]>, <@mesh, []>] out_shardings=[<@mesh, [{"a"}, {}]>] manual_axes={"a"} (%arg2: tensor<8x32xf32>, %arg3: !stablehlo.token) {
+-    sdy.return %arg3 : !stablehlo.token
+-  } : (tensor<16x32xf32>, !stablehlo.token) -> tensor<16x32xf32>
+-  func.return %0: tensor<16x32xf32>
+-}
+diff --git a/shardy/dialect/sdy/ir/utils.cc b/shardy/dialect/sdy/ir/utils.cc
+index ca77dca8..4a639a80 100644
+--- a/shardy/dialect/sdy/ir/utils.cc
++++ b/shardy/dialect/sdy/ir/utils.cc
+@@ -25,7 +25,6 @@ limitations under the License.
+ #include <numeric>
+ #include <optional>
+ #include <string>
+-#include <string_view>
+ #include <utility>
+ #include <vector>
+ 
+@@ -223,7 +222,7 @@ TensorShardingAttr inlineMesh(const SymbolTable& symbolTable,
+                               TensorShardingAttr sharding) {
+   if (auto name = dyn_cast<FlatSymbolRefAttr>(sharding.getMeshOrRef())) {
+     MeshAttr mesh = getMeshAttr(symbolTable, name);
+-    SDY_CHECK(mesh) << "unknown mesh: " << std::string_view(name.getValue());
++    assert(mesh && "unknown mesh");
+     return TensorShardingAttr::get(
+         sharding.getContext(), mesh, sharding.getDimShardings(),
+         sharding.getReplicatedAxes(), sharding.getUnreducedAxes(),
+@@ -264,7 +263,7 @@ Attribute getCommonMeshOrRef(ArrayRef<TensorShardingAttr> operandShardings,
+       meshOrRef = sharding.getMeshOrRef();
+       continue;
+     }
+-    if (!otherMesh || otherMesh.empty()) {
++    if (otherMesh.empty()) {
+       continue;
+     }
+     if (!otherMesh.equals(mesh, ignoreDeviceIds)) {
+@@ -964,13 +963,9 @@ mlir::Attribute getMeshOrRef(
+     int64_t numElements, const SymbolTable& symbolTable,
+     std::function<TensorShardingAttr(int64_t)> getSharding) {
+   for (int64_t i = 0; i < numElements; ++i) {
+-    if (TensorShardingAttr sdySharding = getSharding(i); sdySharding) {
+-      MeshAttr mesh = sdySharding.getMesh(symbolTable);
+-      SDY_CHECK(mesh) << "unknown mesh: "
+-                      << std::string_view(sdySharding.getMeshName());
+-      if (!mesh.isSingleDevice()) {
+-        return sdySharding.getMeshOrRef();
+-      }
++    if (TensorShardingAttr sdySharding = getSharding(i);
++        sdySharding && !sdySharding.getMesh(symbolTable).isSingleDevice()) {
++      return sdySharding.getMeshOrRef();
+     }
+   }
+   return nullptr;
+diff --git a/shardy/dialect/sdy/ir/verifiers.cc b/shardy/dialect/sdy/ir/verifiers.cc
+index 65c6212b..3d90562f 100644
+--- a/shardy/dialect/sdy/ir/verifiers.cc
++++ b/shardy/dialect/sdy/ir/verifiers.cc
+@@ -1021,12 +1021,13 @@ LogicalResult verifyManualComputationValue(
+     //    arguments/results match.
+     auto expectedLocalRankedType =
+         RankedTensorType::get(newDimSizes, globalShapedType.getElementType());
+-    if (expectedLocalRankedType != localType) {
++    auto localRankedType = cast<RankedTensorType>(localType);
++    if (expectedLocalRankedType != localRankedType) {
+       return op->emitOpError(valueKindStr)
+              << " shape, corresponding sharding, and region " << valueKindStr
+              << " shape at index " << valueIndex
+              << " must match. Expected local shape " << expectedLocalRankedType
+-             << ", actual local shape " << localType;
++             << ", actual local shape " << localRankedType;
+     }
+   }
+ 
+diff --git a/shardy/dialect/sdy/transforms/export/insert_explicit_reshards.cc b/shardy/dialect/sdy/transforms/export/insert_explicit_reshards.cc
+index 756b46ae..6b86b36e 100644
+--- a/shardy/dialect/sdy/transforms/export/insert_explicit_reshards.cc
++++ b/shardy/dialect/sdy/transforms/export/insert_explicit_reshards.cc
+@@ -19,7 +19,6 @@ limitations under the License.
+ #include <cstdint>
+ #include <memory>  // IWYU pragma: keep
+ #include <optional>
+-#include <string_view>
+ #include <utility>
+ 
+ #include "llvm/ADT/DenseMap.h"
+@@ -351,7 +350,7 @@ std::optional<MeshOp> getMesh(ArrayRef<TensorShardingAttr> inShardings,
+     return std::nullopt;
+   }
+   MeshOp meshOp = getMeshOp(symbolTable, *meshName);
+-  SDY_CHECK(meshOp) << "unknown mesh: " << std::string_view(*meshName);
++  assert(meshOp && "unknown mesh");
+   if (meshOp.getMesh().isSingleDevice()) {
+     return std::nullopt;
+   }
+diff --git a/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc b/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc
+index 4cef1c10..111e882e 100644
+--- a/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc
++++ b/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc
+@@ -557,8 +557,7 @@ class GenericOpPattern : public ConversionPattern {
+     Dialect* dialect = op->getDialect();
+     if ((dialect && dialect->getNamespace() != "stablehlo" &&
+          !isa<func::CallOp, func::ReturnOp, sdy::ReturnOp, sdy::AllReduceOp,
+-              sdy::CollectivePermuteOp, sdy::ShardedToUnreducedOp,
+-              sdy::ReplicatedToUnreducedOp>(op)) ||
++              sdy::ShardedToUnreducedOp, sdy::ReplicatedToUnreducedOp>(op)) ||
+         hasCustomPadHandling(op)) {
+       return failure();
+     }
+diff --git a/shardy/dialect/sdy/transforms/export/reshard_to_collectives.cc b/shardy/dialect/sdy/transforms/export/reshard_to_collectives.cc
+index 3b01769d..5554918b 100644
+--- a/shardy/dialect/sdy/transforms/export/reshard_to_collectives.cc
++++ b/shardy/dialect/sdy/transforms/export/reshard_to_collectives.cc
+@@ -29,7 +29,6 @@ limitations under the License.
+ #include "llvm/ADT/DenseSet.h"
+ #include "llvm/ADT/STLExtras.h"
+ #include "llvm/ADT/SmallVector.h"
+-#include "llvm/Support/MathExtras.h"
+ #include "mlir/Dialect/Func/IR/FuncOps.h"  // IWYU pragma: keep
+ #include "mlir/IR/Attributes.h"
+ #include "mlir/IR/BuiltinAttributes.h"
+@@ -97,16 +96,10 @@ AxisList concatAxisLists(MutableArrayRef<AxisList> axisLists) {
+   return result;
+ }
+ 
+-// Remove the common prefix of `inAxesPerDim` and `outAxesPerDim` for dimensions
+-// that do not need to be unsharded first.
++// Remove the common prefix of `inAxesPerDim` and `outAxesPerDim`.
+ void removeCommonPrefix(SmallVector<AxisList>& inAxesPerDim,
+-                        SmallVector<AxisList>& outAxesPerDim,
+-                        ArrayRef<bool> dimsToUnshard) {
+-  for (auto [inAxes, outAxes, unshardDim] :
+-       llvm::zip_equal(inAxesPerDim, outAxesPerDim, dimsToUnshard)) {
+-    if (unshardDim) {
+-      continue;
+-    }
++                        SmallVector<AxisList>& outAxesPerDim) {
++  for (auto [inAxes, outAxes] : llvm::zip_equal(inAxesPerDim, outAxesPerDim)) {
+     while (!inAxes.empty() && !outAxes.empty() &&
+            inAxes.front() == outAxes.front()) {
+       inAxes.pop_front();
+@@ -202,16 +195,11 @@ struct AllToAllInfo {
+ // We define the current state of the transformation as follows:
+ //
+ // - `inAxesPerDim` - the axes in the current sharding per dimension, such that
+-//   the common prefix with the output sharding is removed (unless the
+-//   dimension is in `dimsToUnshard`).
++//   the common prefix with the output sharding is removed.
+ // - `outAxesPerDim` - the axes in the output sharding per dimension, such that
+-//   the common prefix with the current sharding is removed (unless the
+-//   dimension is in `dimsToUnshard`).
++//   the common prefix with the current sharding is removed.
+ // - `currentAxesPerDim` - the axes in the current sharding, including the
+ //   common prefix with the output sharding.
+-// - `dimsToUnshard` - per-dimension boolean indicating whether a dimension has
+-//   different padded sizes in the input and output shardings and must first be
+-//   completely unsharded before acquiring its output sharding.
+ //
+ // These invariants are maintained throughout the algorithm, and specifically
+ // after each collective insertion.
+@@ -245,39 +233,21 @@ class CollectiveInserter {
+         outAxesPerDim(getAxesPerDim<AxisList>(outSharding)),
+         currentAxesPerDim(getAxesPerDim<SmallVector<AxisRefAttr>>(inSharding)),
+         capacityPerDim(inSharding.getRank(), 1),
+-        collectiveAxesPerDim(inSharding.getRank()),
+-        dimsToUnshard(inSharding.getRank(), false) {
++        collectiveAxesPerDim(inSharding.getRank()) {
+     // Unreduced axes in the input and output sharding must match, given we
+     // insert an all-reduce if an unreduced axis becomes replicated/sharded, and
+     // never insert a reshard that goes from replicated/sharded to unreduced.
+     assert(inSharding.getUnreducedAxes() == outSharding.getUnreducedAxes());
+ 
+-    // If a dimension is sharded in both the input and output shardings with
+-    // different padded sizes, it must first be unsharded so that downstream
+-    // passes can pad or trim the dimension while it is unsharded.
+-    ArrayRef<int64_t> shape = getTensorShape(result);
+-    for (int64_t dim = 0; dim < getRank(); ++dim) {
+-      int64_t dimSize = shape[dim];
+-      int64_t inShardedSize =
+-          inSharding.getDimSharding(dim).getShardedSize(mesh);
+-      int64_t outShardedSize =
+-          outSharding.getDimSharding(dim).getShardedSize(mesh);
+-      if (dimSize > 0 && inShardedSize > 1 && outShardedSize > 1 &&
+-          llvm::alignTo(dimSize, inShardedSize) !=
+-              llvm::alignTo(dimSize, outShardedSize)) {
+-        dimsToUnshard[dim] = true;
+-      }
+-    }
+-
+     // We align sub-axes between the input and output axes, so that we can treat
+     // sub-axes like full axes and assume any two sub-axes that overlap are also
+     // equal, which allows using them as keys in a hash map.
+     alignSubAxesByDecomposition(inAxesPerDim, outAxesPerDim, mesh);
+-    // We remove the common prefix of `inAxesPerDim` and `outAxesPerDim` (except
+-    // on dimensions in `dimsToUnshard`), since those axes stay exactly the same
+-    // during the reshard. We are left with `inAxesPerDim` and `outAxesPerDim`
+-    // that need to become empty, via a sequence of collectives.
+-    removeCommonPrefix(inAxesPerDim, outAxesPerDim, dimsToUnshard);
++    // We remove the common prefix of `inAxesPerDim` and `outAxesPerDim`, since
++    // those axes stay exactly the same during the reshard. We are left with
++    // `inAxesPerDim` and `outAxesPerDim` that need to become empty, via a
++    // sequence of collectives.
++    removeCommonPrefix(inAxesPerDim, outAxesPerDim);
+ 
+     inAxisSet = getAxisSet(inAxesPerDim);
+     outAxisToDimAndIndex = getAxisToDimAndIndex(outAxesPerDim);
+@@ -297,12 +267,11 @@ class CollectiveInserter {
+     // can be handled by the previous collectives, so we are left with
+     // all-gathering all axes in `inAxesPerDim` and we're done.
+     //
+-    // Otherwise, we need at most 2 iterations (or 3 if a dimension must first
+-    // be unsharded due to a padded dimension size mismatch).
++    // Otherwise, we need at most 2 iterations.
+ 
+     // TODO(tomnatan): consider looping only over all-to-all and collective
+     // permute as long as one of them was inserted.
+-    for (int i = 0; i < 3 && !isDone(); ++i) {
++    for (int i = 0; i < 2 && !isDone(); ++i) {
+       // 1. Try to insert an all-slice, that decreases the size of the tensor.
+       tryAllSlice();
+ 
+@@ -351,10 +320,9 @@ class CollectiveInserter {
+   // that dimension.
+   //
+   // We gather the suffix of axes (`gatheringAxes`) in `inAxesPerDim[dim]` that
+-  // aren't present in `outAxesPerDim` (or whose target dimension in
+-  // `outAxesPerDim` is in `dimsToUnshard`, or all axes if `dimsToUnshard[dim]`
+-  // is true), and update the internal state as follows: `gatheringAxes` are
+-  // popped from the back of `inAxesPerDim[dim]` and `currentAxesPerDim[dim]`.
++  // aren't present in `outAxesPerDim`, and update the internal state as
++  // follows: `gatheringAxes` are popped from the back of `inAxesPerDim[dim]`
++  // and `currentAxesPerDim[dim]`.
+   //
+   // For example:
+   //
+@@ -376,18 +344,11 @@ class CollectiveInserter {
+     SmallVector<AxisRefAttr>& currentAxes = currentAxesPerDim[dim];
+     SmallVector<AxisRefAttr> gatheringAxes;
+     gatheringAxes.reserve(inAxes.size());
+-    bool gatherAll = dimsToUnshard[dim];
+-    auto shouldGatherAxis = [&](AxisRefAttr axis) {
+-      if (gatherAll) {
+-        return true;
+-      }
+-      auto it = outAxisToDimAndIndex.find(axis);
+-      return it == outAxisToDimAndIndex.end() || dimsToUnshard[it->second.dim];
+-    };
+     auto axisRevIt = inAxes.rbegin();
+-    while (axisRevIt != inAxes.rend() && shouldGatherAxis(*axisRevIt)) {
++    while (axisRevIt != inAxes.rend() &&
++           !outAxisToDimAndIndex.contains(*axisRevIt)) {
+       inAxisSet.erase(*axisRevIt);
+-      ++axisRevIt;
++      --axisRevIt;
+     }
+     auto inAxisIt = axisRevIt.base();
+     popBackFromCurrentAxes(currentAxes, inAxes, inAxisIt);
+@@ -408,12 +369,6 @@ class CollectiveInserter {
+       }
+       collectiveAxes = AxisRefListAttr::get(getContext(), gatheringAxes);
+     }
+-    for (auto [currentAxes, unshardDim] :
+-         llvm::zip_equal(currentAxesPerDim, dimsToUnshard)) {
+-      if (currentAxes.empty()) {
+-        unshardDim = false;
+-      }
+-    }
+     if (hasGatheringAxes) {
+       result = AllGatherOp::create(rewriter, loc, result, collectiveAxesPerDim,
+                                    getCurrentSharding());
+@@ -620,10 +575,9 @@ class CollectiveInserter {
+   // We define the term capacity per dimension for all-slice - the product of
+   // axis sizes in `outAxesPerDim[d]` divided by the product of axis sizes in
+   // `inAxesPerDim[d]` for each dimension d, or the former divided by their GCD,
+-  // if the former is not divisible by the latter (or 1 if `dimsToUnshard[d]` is
+-  // true, since dimension d must be unsharded before slicing). In other words,
+-  // the capacity represents how much the output is sharded more than the input
+-  // along a specific dimension.
++  // if the former is not divisible by the latter (as the dimension size needs
++  // to be a multiplier of both). In other words, the capacity represents how
++  // much the output is sharded more than the input along a specific dimension.
+   //
+   // Constraint: we can only slice dimension d along axes whose product of
+   // sizes doesn't exceed the capacity for that dimension, otherwise we will
+@@ -715,12 +669,8 @@ class CollectiveInserter {
+     bool hasSlicingAxes = false;
+ 
+     // Initialize capacity per dimension.
+-    for (auto [inAxes, outAxes, dimCapacity, unshardDim] : llvm::zip_equal(
+-             inAxesPerDim, outAxesPerDim, capacityPerDim, dimsToUnshard)) {
+-      if (unshardDim) {
+-        dimCapacity = 1;
+-        continue;
+-      }
++    for (auto [inAxes, outAxes, dimCapacity] :
++         llvm::zip_equal(inAxesPerDim, outAxesPerDim, capacityPerDim)) {
+       int64_t inShardedSize = getShardedSize(inAxes, mesh);
+       int64_t outShardedSize = getShardedSize(outAxes, mesh);
+       dimCapacity = outShardedSize / std::gcd(inShardedSize, outShardedSize);
+@@ -748,13 +698,9 @@ class CollectiveInserter {
+     // 2. Slice axes in the dimension they appear in `outAxesPerDim`, i.e. the
+     // desired dimension, as long as the per-dim capacity allows.
+     AxisList availableOutAxes;
+-    for (auto [inAxes, outAxes, currentAxes, slicingAxes, dimCapacity,
+-               unshardDim] :
++    for (auto [inAxes, outAxes, currentAxes, slicingAxes, dimCapacity] :
+          llvm::zip_equal(inAxesPerDim, outAxesPerDim, currentAxesPerDim,
+-                         slicingAxesPerDim, capacityPerDim, dimsToUnshard)) {
+-      if (unshardDim) {
+-        continue;
+-      }
++                         slicingAxesPerDim, capacityPerDim)) {
+       auto outIt = outAxes.begin();
+       while (outIt != outAxes.end()) {
+         AxisRefAttr outAxis = *outIt;
+@@ -822,8 +768,7 @@ class CollectiveInserter {
+     }
+   }
+ 
+-  // We should insert a collective permute (assuming no dimension in
+-  // `dimsToUnshard` still needs to be unsharded) if one of the following holds:
++  // We should insert a collective permute if one of the following holds:
+   //
+   // 1. Both `inAxesPerDim[d]` and `outAxesPerDim[d]` aren't empty for a certain
+   //    dimension d and their front axes are divisible. This means we can
+@@ -856,9 +801,6 @@ class CollectiveInserter {
+   //    that B is before A, and A can be moved to d2 via an all-to-all before B
+   //    is all-gathered.
+   bool shouldCollectivePermute() {
+-    if (llvm::any_of(dimsToUnshard, [](bool b) { return b; })) {
+-      return false;
+-    }
+     OptionalAxisRef availableInAxis;
+     OptionalAxisRef availableOutAxis;
+     for (auto [dim, inOutAxes] :
+@@ -1123,15 +1065,12 @@ class CollectiveInserter {
+   // The suffix of axes in `inAxesPerDim[srcDim]`, such that all axes within the
+   // suffix are mapped to the same dimension in `outAxisToDimAndIndex`
+   // (`allToAllAxes`), are all-to-all-ed with the mapped dimension as the target
+-  // (`tgtDim`), provided `dimsToUnshard[tgtDim]` is false and (if
+-  // `dimsToUnshard[srcDim]` is true) `allToAllAxes` includes all axes on
+-  // `srcDim` so that `srcDim` becomes completely unsharded.
++  // (`tgtDim`).
+   //
+   // The internal state is updated as follows for `allToAllAxes` and `tgtDim`:
+   //
+   // - `allToAllAxes` are popped from the back of `inAxesPerDim[srcDim]` and
+-  //   `currentAxesPerDim[srcDim]` (and `dimsToUnshard[srcDim]` is cleared if
+-  //   `currentAxesPerDim[srcDim]` becomes empty).
++  //   `currentAxesPerDim[srcDim]`.
+   // - `allToAllAxes` are appended to `inAxesPerDim[tgtDim]` and
+   //   `currentAxesPerDim[tgtDim]`.
+   // - The common prefix between `inAxesPerDim[tgtDim]` and
+@@ -1198,8 +1137,7 @@ class CollectiveInserter {
+         break;
+       }
+       int64_t outAxisDim = outAxisEntryIt->second.dim;
+-      if (outAxisDim == srcDim || dimsToUnshard[outAxisDim] ||
+-          (tgtDim && outAxisDim != *tgtDim)) {
++      if (outAxisDim == srcDim || (tgtDim && outAxisDim != *tgtDim)) {
+         break;
+       }
+       tgtDim = outAxisDim;
+@@ -1212,9 +1150,6 @@ class CollectiveInserter {
+     }
+ 
+     auto startInAxisIt = axisRevIt.base();
+-    if (dimsToUnshard[srcDim] && startInAxisIt != srcInAxes.begin()) {
+-      return std::nullopt;
+-    }
+ 
+     AxisList& tgtOutAxes = outAxesPerDim[*tgtDim];
+     if (!allowOutOfOrderTarget &&
+@@ -1231,9 +1166,6 @@ class CollectiveInserter {
+     SmallVector<AxisRefAttr>& tgtCurrentAxes = currentAxesPerDim[*tgtDim];
+ 
+     popBackFromCurrentAxes(srcCurrentAxes, srcInAxes, startInAxisIt);
+-    if (srcCurrentAxes.empty()) {
+-      dimsToUnshard[srcDim] = false;
+-    }
+ 
+     AxisList& tgtInAxes = inAxesPerDim[*tgtDim];
+     auto srcInAxisIt = startInAxisIt;
+@@ -1280,7 +1212,6 @@ class CollectiveInserter {
+   AxesPerDim currentAxesPerDim;
+   SmallVector<int64_t> capacityPerDim;
+   SmallVector<AxisRefListAttr> collectiveAxesPerDim;
+-  SmallVector<bool> dimsToUnshard;
+   AxisSet inAxisSet;
+   AxisToDimAndIndex outAxisToDimAndIndex;
+ };
+diff --git a/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/generic_ops.mlir b/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/generic_ops.mlir
+index 9bd4e0a8..4d4faae8 100644
+--- a/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/generic_ops.mlir
++++ b/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/generic_ops.mlir
+@@ -12,45 +12,3 @@ func.func @no_pad(
+   %0 = stablehlo.add %arg0, %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4_2, [{"x"}, {}]>]>} : tensor<4x8xf32>
+   return %0 : tensor<4x8xf32>
+ }
+-
+-sdy.mesh @mesh_4_2_4 = <["x"=4, "y"=2, "z"=4]>
+-
+-// CHECK-LABEL: func private @pad_all_reduce
+-func.func private @pad_all_reduce(
+-  %arg0: tensor<7x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{"x"}, {}], unreduced={"y"}>})
+-  -> (tensor<7x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{"x"}, {}]>}) {
+-  // CHECK-NEXT: %[[AR:.*]] = sdy.all_reduce {"y"} %arg0 out_sharding=<@mesh_4_2, [{"x"}, {}]> : tensor<8x8xf32>
+-  // CHECK-NEXT: return %[[AR]] : tensor<8x8xf32>
+-  %0 = sdy.all_reduce {"y"} %arg0 out_sharding=<@mesh_4_2, [{"x"}, {}]> : tensor<7x8xf32>
+-  return %0 : tensor<7x8xf32>
+-}
+-
+-// CHECK-LABEL: func private @pad_collective_permute
+-func.func private @pad_collective_permute(
+-  %arg0: tensor<7x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2_4, [{"x"}, {}]>})
+-  -> (tensor<7x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2_4, [{"z"}, {}]>}) {
+-  // CHECK-NEXT: %[[CP:.*]] = sdy.collective_permute %arg0 out_sharding=<@mesh_4_2_4, [{"z"}, {}]> : tensor<8x8xf32>
+-  // CHECK-NEXT: return %[[CP]] : tensor<8x8xf32>
+-  %0 = sdy.collective_permute %arg0 out_sharding=<@mesh_4_2_4, [{"z"}, {}]> : tensor<7x8xf32>
+-  return %0 : tensor<7x8xf32>
+-}
+-
+-// CHECK-LABEL: func private @pad_sharded_to_unreduced
+-func.func private @pad_sharded_to_unreduced(
+-  %arg0: tensor<7x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{"x"}, {"y"}]>})
+-  -> (tensor<7x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{"x"}, {}], unreduced={"y"}>}) {
+-  // CHECK-NEXT: %[[S2U:.*]] = sdy.sharded_to_unreduced [{}, {"y"}] %arg0 out_sharding=<@mesh_4_2, [{"x"}, {}], unreduced={"y"}> : tensor<8x8xf32>
+-  // CHECK-NEXT: return %[[S2U]] : tensor<8x8xf32>
+-  %0 = sdy.sharded_to_unreduced [{}, {"y"}] %arg0 out_sharding=<@mesh_4_2, [{"x"}, {}], unreduced={"y"}> : tensor<7x8xf32>
+-  return %0 : tensor<7x8xf32>
+-}
+-
+-// CHECK-LABEL: func private @pad_replicated_to_unreduced
+-func.func private @pad_replicated_to_unreduced(
+-  %arg0: tensor<7x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{"x"}, {}]>})
+-  -> (tensor<7x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{"x"}, {}], unreduced={"y"}>}) {
+-  // CHECK-NEXT: %[[R2U:.*]] = sdy.replicated_to_unreduced {"y"} %arg0 out_sharding=<@mesh_4_2, [{"x"}, {}], unreduced={"y"}> : tensor<8x8xf32>
+-  // CHECK-NEXT: return %[[R2U]] : tensor<8x8xf32>
+-  %0 = sdy.replicated_to_unreduced {"y"} %arg0 out_sharding=<@mesh_4_2, [{"x"}, {}], unreduced={"y"}> : tensor<7x8xf32>
+-  return %0 : tensor<7x8xf32>
+-}
+diff --git a/shardy/dialect/sdy/transforms/export/test/reshard_to_collectives.mlir b/shardy/dialect/sdy/transforms/export/test/reshard_to_collectives.mlir
+index a63c54d1..f185669a 100644
+--- a/shardy/dialect/sdy/transforms/export/test/reshard_to_collectives.mlir
++++ b/shardy/dialect/sdy/transforms/export/test/reshard_to_collectives.mlir
+@@ -644,12 +644,12 @@ func.func @all_to_all_axes_at_src_out_of_order(%arg0 : tensor<16x8x8xf32> {sdy.s
+ }
+ 
+ // CHECK-LABEL: func @all_to_all_axes_at_src_and_tgt_out_of_order
+-func.func @all_to_all_axes_at_src_and_tgt_out_of_order(%arg0 : tensor<32x8x8xf32> {sdy.sharding = #sdy.sharding<@mesh3d_4x2x4, [{"z"}, {"y", "x"}, {}]>}) -> tensor<32x8x8xf32> {
++func.func @all_to_all_axes_at_src_and_tgt_out_of_order(%arg0 : tensor<16x8x8xf32> {sdy.sharding = #sdy.sharding<@mesh3d_4x2x4, [{"z"}, {"y", "x"}, {}]>}) -> tensor<16x8x8xf32> {
+   // CHECK-NEXT: %[[COLLECTIVE_PERMUTE:.*]] = sdy.collective_permute %arg0 out_sharding=<@mesh3d_4x2x4, [{"x"}, {"y", "z"}, {}]>
+   // CHECK-NEXT: %[[ALL_TO_ALL:.*]] = sdy.all_to_all [{"y", "z"}: 1->0] %[[COLLECTIVE_PERMUTE]] out_sharding=<@mesh3d_4x2x4, [{"x", "y", "z"}, {}, {}]>
+   // CHECK-NEXT: return %[[ALL_TO_ALL]]
+-  %0 = sdy.reshard %arg0 <@mesh3d_4x2x4, [{"x", "y", "z"}, {}, {}]> : tensor<32x8x8xf32>
+-  return %0 : tensor<32x8x8xf32>
++  %0 = sdy.reshard %arg0 <@mesh3d_4x2x4, [{"x", "y", "z"}, {}, {}]> : tensor<16x8x8xf32>
++  return %0 : tensor<16x8x8xf32>
+ }
+ 
+ // CHECK-LABEL: func @all_to_all_two_tgt_dims_src_out_of_order
+@@ -935,60 +935,6 @@ func.func @reshard_with_propagation_barrier(
+   return %1 : tensor<8x8x8xf32>
+ }
+ 
+-// CHECK-LABEL: func @reshard_uneven_padded_size_mismatch_common_prefix
+-func.func @reshard_uneven_padded_size_mismatch_common_prefix(%arg0 : tensor<6x3xf32> {sdy.sharding = #sdy.sharding<@mesh2d_4x2, [{"x":(1)2}, {}]>}) -> tensor<6x3xf32> {
+-  // CHECK-NEXT: %[[ALL_GATHER:.*]] = sdy.all_gather [{"x":(1)2}, {}] %arg0 out_sharding=<@mesh2d_4x2, [{}, {}]> : tensor<6x3xf32>
+-  // CHECK-NEXT: %[[ALL_SLICE:.*]] = sdy.all_slice [{"x"}, {}] %[[ALL_GATHER]] out_sharding=<@mesh2d_4x2, [{"x"}, {}]> : tensor<6x3xf32>
+-  // CHECK-NEXT: return %[[ALL_SLICE]] : tensor<6x3xf32>
+-  %0 = sdy.reshard %arg0 <@mesh2d_4x2, [{"x"}, {}]> : tensor<6x3xf32>
+-  return %0 : tensor<6x3xf32>
+-}
+-
+-// CHECK-LABEL: func @reshard_uneven_padded_size_mismatch_subaxis
+-func.func @reshard_uneven_padded_size_mismatch_subaxis(%arg0 : tensor<6x3xf32> {sdy.sharding = #sdy.sharding<@mesh2d_2x8, [{"y":(1)2}, {}]>}) -> tensor<6x3xf32> {
+-  // CHECK-NEXT: %[[ALL_GATHER:.*]] = sdy.all_gather [{"y":(1)2}, {}] %arg0 out_sharding=<@mesh2d_2x8, [{}, {}]> : tensor<6x3xf32>
+-  // CHECK-NEXT: %[[ALL_SLICE:.*]] = sdy.all_slice [{"y":(2)4}, {}] %[[ALL_GATHER]] out_sharding=<@mesh2d_2x8, [{"y":(2)4}, {}]> : tensor<6x3xf32>
+-  // CHECK-NEXT: return %[[ALL_SLICE]] : tensor<6x3xf32>
+-  %0 = sdy.reshard %arg0 <@mesh2d_2x8, [{"y":(2)4}, {}]> : tensor<6x3xf32>
+-  return %0 : tensor<6x3xf32>
+-}
+-
+-// CHECK-LABEL: func @reshard_uneven_padded_size_mismatch_all_to_all
+-func.func @reshard_uneven_padded_size_mismatch_all_to_all(%arg0 : tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh1d_6, [{}, {"x":(1)2}]>}) -> tensor<8x8xf32> {
+-  // CHECK-NEXT: %[[ALL_TO_ALL:.*]] = sdy.all_to_all [{"x":(1)2}: 1->0] %arg0 out_sharding=<@mesh1d_6, [{"x":(1)2}, {}]> : tensor<8x8xf32>
+-  // CHECK-NEXT: %[[ALL_SLICE:.*]] = sdy.all_slice [{}, {"x":(2)3}] %[[ALL_TO_ALL]] out_sharding=<@mesh1d_6, [{"x":(1)2}, {"x":(2)3}]> : tensor<8x8xf32>
+-  // CHECK-NEXT: return %[[ALL_SLICE]] : tensor<8x8xf32>
+-  %0 = sdy.reshard %arg0 <@mesh1d_6, [{"x":(1)2}, {"x":(2)3}]> : tensor<8x8xf32>
+-  return %0 : tensor<8x8xf32>
+-}
+-
+-// CHECK-LABEL: func @reshard_uneven_padded_size_mismatch_all_to_all_src_and_tgt_out_of_order
+-func.func @reshard_uneven_padded_size_mismatch_all_to_all_src_and_tgt_out_of_order(%arg0 : tensor<16x8x8xf32> {sdy.sharding = #sdy.sharding<@mesh3d_4x2x4, [{"z"}, {"y", "x"}, {}]>}) -> tensor<16x8x8xf32> {
+-  // CHECK-NEXT: %[[ALL_GATHER:.*]] = sdy.all_gather [{"z"}, {"y", "x"}, {}] %arg0 out_sharding=<@mesh3d_4x2x4, [{}, {}, {}]> : tensor<16x8x8xf32>
+-  // CHECK-NEXT: %[[ALL_SLICE:.*]] = sdy.all_slice [{"x", "y", "z"}, {}, {}] %[[ALL_GATHER]] out_sharding=<@mesh3d_4x2x4, [{"x", "y", "z"}, {}, {}]> : tensor<16x8x8xf32>
+-  // CHECK-NEXT: return %[[ALL_SLICE]] : tensor<16x8x8xf32>
+-  %0 = sdy.reshard %arg0 <@mesh3d_4x2x4, [{"x", "y", "z"}, {}, {}]> : tensor<16x8x8xf32>
+-  return %0 : tensor<16x8x8xf32>
+-}
+-
+-// CHECK-LABEL: func @reshard_uneven_padded_size_mismatch_swap_axes
+-func.func @reshard_uneven_padded_size_mismatch_swap_axes(%arg0 : tensor<5x5xf32> {sdy.sharding = #sdy.sharding<@mesh2d_4x2, [{"x"}, {"y"}]>}) -> tensor<5x5xf32> {
+-  // CHECK-NEXT: %[[ALL_GATHER:.*]] = sdy.all_gather [{"x"}, {"y"}] %arg0 out_sharding=<@mesh2d_4x2, [{}, {}]> : tensor<5x5xf32>
+-  // CHECK-NEXT: %[[ALL_SLICE:.*]] = sdy.all_slice [{"y"}, {"x"}] %[[ALL_GATHER]] out_sharding=<@mesh2d_4x2, [{"y"}, {"x"}]> : tensor<5x5xf32>
+-  // CHECK-NEXT: return %[[ALL_SLICE]] : tensor<5x5xf32>
+-  %0 = sdy.reshard %arg0 <@mesh2d_4x2, [{"y"}, {"x"}]> : tensor<5x5xf32>
+-  return %0 : tensor<5x5xf32>
+-}
+-
+-// CHECK-LABEL: func @reshard_uneven_padded_size_match_swap_axes
+-func.func @reshard_uneven_padded_size_match_swap_axes(%arg0 : tensor<19x19xf32> {sdy.sharding = #sdy.sharding<@mesh2d_4x2, [{"x"}, {"y"}]>}) -> tensor<19x19xf32> {
+-  // CHECK-NEXT: %[[CP:.*]] = sdy.collective_permute %arg0 out_sharding=<@mesh2d_4x2, [{"y", "x":(2)2}, {"x":(1)2}]> : tensor<19x19xf32>
+-  // CHECK-NEXT: %[[ALL_TO_ALL:.*]] = sdy.all_to_all [{"x":(2)2}: 0->1] %[[CP]] out_sharding=<@mesh2d_4x2, [{"y"}, {"x"}]> : tensor<19x19xf32>
+-  // CHECK-NEXT: return %[[ALL_TO_ALL]] : tensor<19x19xf32>
+-  %0 = sdy.reshard %arg0 <@mesh2d_4x2, [{"y"}, {"x"}]> : tensor<19x19xf32>
+-  return %0 : tensor<19x19xf32>
+-}
+-
+ //===----------------------------------------------------------------------===//
+ // Single-Device Reshard Tests (Preserved as sdy.reshard)
+ //===----------------------------------------------------------------------===//
+@@ -1010,3 +956,5 @@ func.func @single_device_out_sharding(%arg0 : tensor<16x8xf32> {sdy.sharding = #
+   %0 = sdy.reshard %arg0 <@single_dev_0, []> : tensor<16x8xf32>
+   return %0 : tensor<16x8xf32>
+ }
++
++
+diff --git a/shardy/dialect/sdy/transforms/export/update_non_divisible_input_output_shardings.cc b/shardy/dialect/sdy/transforms/export/update_non_divisible_input_output_shardings.cc
+index 97c93fbf..90bcb8d5 100644
+--- a/shardy/dialect/sdy/transforms/export/update_non_divisible_input_output_shardings.cc
++++ b/shardy/dialect/sdy/transforms/export/update_non_divisible_input_output_shardings.cc
+@@ -17,7 +17,6 @@ limitations under the License.
+ #include <cstdint>
+ #include <functional>
+ #include <numeric>
+-#include <string_view>
+ 
+ #include "llvm/ADT/ArrayRef.h"
+ #include "llvm/ADT/STLExtras.h"
+@@ -35,7 +34,6 @@ limitations under the License.
+ #include "mlir/Pass/Pass.h"  // IWYU pragma: keep
+ #include "mlir/Support/LLVM.h"
+ #include "mlir/Transforms/DialectConversion.h"
+-#include "shardy/common/logging.h"
+ #include "shardy/dialect/sdy/ir/dialect.h"
+ #include "shardy/dialect/sdy/ir/utils.h"
+ 
+@@ -67,8 +65,7 @@ TensorShardingAttr getEvenlySharded(TensorShardingAttr sharding,
+                                     ShapedType type,
+                                     const SymbolTable& symbolTable) {
+   MeshAttr mesh = sharding.getMesh(symbolTable);
+-  SDY_CHECK(mesh) << "unknown mesh: "
+-                  << std::string_view(sharding.getMeshName());
++  assert(mesh && "unknown mesh");
+   MLIRContext* ctx = sharding.getContext();
+   llvm::SmallVector<DimensionShardingAttr> newDimShardings;
+   newDimShardings.reserve(sharding.getRank());
+diff --git a/shardy/dialect/sdy/transforms/import/manual_axes_cleanup.cc b/shardy/dialect/sdy/transforms/import/manual_axes_cleanup.cc
+index 6b368f91..f71266a5 100644
+--- a/shardy/dialect/sdy/transforms/import/manual_axes_cleanup.cc
++++ b/shardy/dialect/sdy/transforms/import/manual_axes_cleanup.cc
+@@ -29,7 +29,6 @@ limitations under the License.
+ #include "mlir/IR/SymbolTable.h"
+ #include "mlir/IR/TypeRange.h"
+ #include "mlir/Support/LLVM.h"
+-#include "shardy/common/logging.h"
+ #include "shardy/dialect/sdy/ir/dialect.h"
+ #include "shardy/dialect/sdy/ir/utils.h"
+ #include "shardy/dialect/sdy/transforms/import/passes.h"  // IWYU pragma: keep
+@@ -163,7 +162,7 @@ struct ManualAxesCleanupPass
+           getCommonMeshOrRef(inShardings, outShardings, symbolTable);
+       MeshAttr mesh =
+           meshOrRef ? getMeshOrLookup(symbolTable, meshOrRef) : nullptr;
+-      SDY_CHECK(mesh) << "expected inputs and outputs to have a common mesh";
++      assert(mesh && "expected inputs and outputs to have a common mesh");
+       sortManualAxes(op, mesh);
+       addUnusedManualAxesToReplicatedAxes(op, mesh, meshOrRef, symbolTable);
+     });
+diff --git a/shardy/dialect/sdy/transforms/import/remove_size_one_axes.cc b/shardy/dialect/sdy/transforms/import/remove_size_one_axes.cc
+index ce2abb40..99a47092 100644
+--- a/shardy/dialect/sdy/transforms/import/remove_size_one_axes.cc
++++ b/shardy/dialect/sdy/transforms/import/remove_size_one_axes.cc
+@@ -17,14 +17,12 @@ limitations under the License.
+ #include <cstdint>
+ #include <iterator>
+ #include <optional>
+-#include <string_view>
+ 
+ #include "llvm/ADT/STLExtras.h"
+ #include "llvm/ADT/SmallVector.h"
+ #include "mlir/IR/BuiltinOps.h"
+ #include "mlir/IR/SymbolTable.h"
+ #include "mlir/Support/LLVM.h"
+-#include "shardy/common/logging.h"
+ #include "shardy/dialect/sdy/ir/dialect.h"
+ #include "shardy/dialect/sdy/transforms/common/sharding_walker.h"
+ #include "shardy/dialect/sdy/transforms/import/passes.h"  // IWYU pragma: keep
+@@ -45,8 +43,7 @@ bool hasSizeOneAxes(MeshOp meshOp) {
+ TensorShardingAttr removeSizeOneAxes(TensorShardingAttr sharding,
+                                      const SymbolTable& symbolTable) {
+   MeshAttr mesh = sharding.getMesh(symbolTable);
+-  SDY_CHECK(mesh) << "unknown mesh: "
+-                  << std::string_view(sharding.getMeshName());
++  assert(mesh && "unknown mesh");
+ 
+   auto isNotSizeOne = [&](AxisRefAttr axis) { return axis.getSize(mesh) != 1; };
+ 
+diff --git a/shardy/dialect/sdy/transforms/import/test/sharding_group_constraints.mlir b/shardy/dialect/sdy/transforms/import/test/sharding_group_constraints.mlir
+index 34db7665..a3518d30 100644
+--- a/shardy/dialect/sdy/transforms/import/test/sharding_group_constraints.mlir
++++ b/shardy/dialect/sdy/transforms/import/test/sharding_group_constraints.mlir
+@@ -310,189 +310,3 @@ func.func private @bar(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+   func.return %arg0 : tensor<8x8xf32>
+ }
+ 
+-// -----
+-
+-sdy.mesh @mesh = <["a"=2, "b"=2]>
+-
+-// Allows sharding groups within ManualComputationOp when group members are
+-// across multiple functions called from the same ManualComputationOp.
+-func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
+-    sdy.sharding_group %arg1 group_id = 1234 : tensor<8x8xf32>
+-    %1 = func.call @foo(%arg1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-    %2 = func.call @bar(%1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-    sdy.return %2 : tensor<8x8xf32>
+-  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  func.return %0: tensor<8x8xf32>
+-}
+-
+-func.func private @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  sdy.sharding_group %arg0 group_id = 1234 : tensor<8x8xf32>
+-  func.return %arg0 : tensor<8x8xf32>
+-}
+-
+-func.func private @bar(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  sdy.sharding_group %arg0 group_id = 1234 : tensor<8x8xf32>
+-  func.return %arg0 : tensor<8x8xf32>
+-}
+-
+-// -----
+-
+-sdy.mesh @mesh = <["a"=2, "b"=2]>
+-
+-// Allows sharding groups within ManualComputationOp when a function with group
+-// members is called multiple times from the same ManualComputationOp.
+-func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
+-    sdy.sharding_group %arg1 group_id = 1234 : tensor<8x8xf32>
+-    %1 = func.call @foo(%arg1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-    %2 = func.call @foo(%1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-    sdy.return %2 : tensor<8x8xf32>
+-  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  func.return %0: tensor<8x8xf32>
+-}
+-
+-func.func private @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  sdy.sharding_group %arg0 group_id = 1234 : tensor<8x8xf32>
+-  func.return %arg0 : tensor<8x8xf32>
+-}
+-
+-// -----
+-
+-sdy.mesh @mesh = <["a"=2, "b"=2]>
+-
+-// TODO(b/569928811): Disallow sharding groups which cross the barrier of a
+-// ManualComputationOp when the function is called inside and then outside the
+-// ManualComputationOp.
+-func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
+-    %1 = func.call @foo(%arg1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-    sdy.return %1 : tensor<8x8xf32>
+-  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-
+-  %2 = func.call @foo(%0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  func.return %2: tensor<8x8xf32>
+-}
+-
+-func.func private @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  sdy.sharding_group %arg0 group_id = 1234 : tensor<8x8xf32>
+-  func.return %arg0 : tensor<8x8xf32>
+-}
+-
+-// -----
+-
+-sdy.mesh @mesh = <["a"=2, "b"=2]>
+-
+-// TODO(b/569928811): Disallow sharding groups which cross the barrier of a
+-// ManualComputationOp when the function is called inside and then outside the
+-// ManualComputationOp.
+-func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
+-    %1 = func.call @foo(%arg1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-    sdy.return %1 : tensor<8x8xf32>
+-  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-
+-  %2 = func.call @foo(%0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  sdy.sharding_group %2 group_id = 1234 : tensor<8x8xf32>
+-  func.return %2: tensor<8x8xf32>
+-}
+-
+-func.func private @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  sdy.sharding_group %arg0 group_id = 1234 : tensor<8x8xf32>
+-  func.return %arg0 : tensor<8x8xf32>
+-}
+-
+-// -----
+-
+-sdy.mesh @mesh = <["a"=2, "b"=2]>
+-
+-// TODO(b/569928811): Disallow sharding groups which cross the barrier of a
+-// ManualComputationOp when the function is called outside and then inside the
+-// ManualComputationOp.
+-func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  %0 = func.call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  %1 = sdy.manual_computation(%0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
+-    sdy.sharding_group %arg1 group_id = 1234 : tensor<8x8xf32>
+-    %2 = func.call @foo(%arg1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-    sdy.return %2 : tensor<8x8xf32>
+-  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  func.return %1: tensor<8x8xf32>
+-}
+-
+-func.func private @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  sdy.sharding_group %arg0 group_id = 1234 : tensor<8x8xf32>
+-  func.return %arg0 : tensor<8x8xf32>
+-}
+-
+-// -----
+-
+-sdy.mesh @mesh = <["a"=2, "b"=2]>
+-
+-// Allows sharding groups inside a function that is called from different
+-// ManualComputationOps.
+-func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
+-    %1 = func.call @foo(%arg1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-    sdy.return %1 : tensor<8x8xf32>
+-  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-
+-  %2 = sdy.manual_computation(%0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
+-    %3 = func.call @foo(%arg1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-    sdy.return %3 : tensor<8x8xf32>
+-  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  func.return %2: tensor<8x8xf32>
+-}
+-
+-func.func private @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  sdy.sharding_group %arg0 group_id = 1234 : tensor<8x8xf32>
+-  func.return %arg0 : tensor<8x8xf32>
+-}
+-
+-// -----
+-
+-sdy.mesh @mesh = <["a"=2, "b"=2]>
+-
+-// Allows sharding groups inside a ManualComputationOp within a function that
+-// is called multiple times.
+-func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  %0 = func.call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  %1 = func.call @foo(%0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  func.return %1: tensor<8x8xf32>
+-}
+-
+-func.func private @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
+-    sdy.sharding_group %arg1 group_id = 1234 : tensor<8x8xf32>
+-    sdy.return %arg1 : tensor<8x8xf32>
+-  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  func.return %0 : tensor<8x8xf32>
+-}
+-
+-// -----
+-
+-sdy.mesh @mesh = <["a"=2, "b"=2]>
+-
+-// Disallows sharding groups which have different ManualComputationOp parents
+-// across different functions.
+-func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  %0 = func.call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  %1 = func.call @bar(%0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  func.return %1: tensor<8x8xf32>
+-}
+-
+-func.func private @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
+-    sdy.sharding_group %arg1 group_id = 1234 : tensor<8x8xf32>
+-    sdy.return %arg1 : tensor<8x8xf32>
+-  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  func.return %0 : tensor<8x8xf32>
+-}
+-
+-func.func private @bar(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
+-    // expected-error@below {{ShardingGroupOps values cannot cross ManualComputationOp boundaries for groupId: 1234}}
+-    sdy.sharding_group %arg1 group_id = 1234 : tensor<8x8xf32>
+-    sdy.return %arg1 : tensor<8x8xf32>
+-  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  func.return %0 : tensor<8x8xf32>
+-}
+diff --git a/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc b/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc
+index 8417874d..2c6b451f 100644
+--- a/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc
++++ b/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc
+@@ -20,7 +20,6 @@ limitations under the License.
+ #include <functional>
+ #include <memory>
+ #include <optional>
+-#include <string_view>
+ #include <utility>
+ 
+ #include "llvm/ADT/STLExtras.h"
+@@ -246,11 +245,6 @@ bool updateTensorSharding(Value modifiedValue, const SymbolTable& symbolTable,
+     if (groupValue == modifiedValue) {
+       continue;
+     }
+-    TensorShardingAttr oldGroupSharding = getSharding(groupValue);
+-    if ((!oldGroupSharding && newSharding.emptyAxes()) ||
+-        newSharding == oldGroupSharding) {
+-      continue;
+-    }
+     setSharding(groupValue, newSharding);
+     if (params.notifyOpModified) {
+       notifyShardingModified(groupValue, symbolTable, userMap,
+@@ -319,7 +313,7 @@ LogicalResult propagateTensorShardings(
+     PropagationDirectionAlongFactor directionAlongFactor,
+     const FactorPropagation& factorPropagation, bool conservativePropagation,
+     Operation* op, const SymbolTable& symbolTable, PatternRewriter* rewriter,
+-    const ShardingGroupMap& shardingGroupMap) {
++    ShardingGroupMap shardingGroupMap) {
+   std::optional<StringRef> meshName =
+       getCommonMeshName(operandsParams.shardings, resultsParams.shardings,
+                         symbolTable, /*ignoreDeviceIds=*/false);
+@@ -333,8 +327,8 @@ LogicalResult propagateTensorShardings(
+     }
+     return failure();
+   }
+-  MeshAttr mesh = getMeshAttr(symbolTable, meshName.value());
+-  SDY_CHECK(mesh) << "unknown mesh: " << std::string_view(meshName.value());
++  MeshAttr mesh = getMeshAttr(op, meshName.value());
++  assert(mesh && "unknown mesh");
+   if (mesh.isSingleDevice()) {
+     // Single-device meshes and shardings are usually a placeholder for special
+     // operations, such as send and recv. We do not propagate single-device
+diff --git a/shardy/dialect/sdy/transforms/propagation/user_priority_propagation.cc b/shardy/dialect/sdy/transforms/propagation/user_priority_propagation.cc
+index f3eb2fc8..34e741d0 100644
+--- a/shardy/dialect/sdy/transforms/propagation/user_priority_propagation.cc
++++ b/shardy/dialect/sdy/transforms/propagation/user_priority_propagation.cc
+@@ -19,7 +19,6 @@ limitations under the License.
+ #include <cstdint>
+ #include <iterator>
+ #include <memory>
+-#include <string_view>
+ #include <utility>
+ #include <variant>
+ 
+@@ -37,7 +36,6 @@ limitations under the License.
+ #include "mlir/Support/LLVM.h"
+ #include "mlir/Support/LogicalResult.h"
+ #include "shardy/common/file_utils.h"
+-#include "shardy/common/logging.h"
+ #include "shardy/dialect/sdy/ir/dialect.h"
+ #include "shardy/dialect/sdy/transforms/common/propagation_options.h"
+ #include "shardy/dialect/sdy/transforms/common/sharding_walker.h"
+@@ -142,8 +140,7 @@ TensorShardingAttr getInitializedSharding(TensorShardingAttr originalSharding,
+     }
+   }
+   MeshAttr mesh = originalSharding.getMesh(symbolTable);
+-  SDY_CHECK(mesh) << "unknown mesh: "
+-                  << std::string_view(originalSharding.getMeshName());
++  assert(mesh && "unknown mesh");
+   llvm::sort(newReplicatedAxes, AxisRefAttr::getMeshComparator(mesh));
+   // TODO(tomnatan): we need to merge split axes and split them again when
+   // updating? or can we assume we won't see split axes?
+diff --git a/third_party/llvm/generated.patch b/third_party/llvm/generated.patch
+index e2237d49..509398da 100644
+--- a/third_party/llvm/generated.patch
++++ b/third_party/llvm/generated.patch
+@@ -1,459 +1 @@
+ Auto generated patch. Do not edit or delete it, even if empty.
+-diff -ruN --strip-trailing-cr a/llvm/include/llvm/IR/Metadata.h b/llvm/include/llvm/IR/Metadata.h
+---- a/llvm/include/llvm/IR/Metadata.h
+-+++ b/llvm/include/llvm/IR/Metadata.h
+-@@ -397,18 +397,8 @@
+-   using OwnerTy = MetadataTracking::OwnerTy;
+- 
+- private:
+--  struct UseEntry {
+--    void *Ref = nullptr;
+--    OwnerTy Owner = nullptr;
+--  };
+--
+--  static constexpr unsigned IndexThreshold = 32;
+--  // Tracked uses of this. Dropping one moves the last entry into its slot, so
+--  // the order is deterministic but not the order they were added in.
+--  SmallVector<UseEntry, 4> UseMap;
+--  // Lazily allocated map from Ref to its index in UseMap for large use lists.
+--  using IndexMapTy = DenseMap<void *, unsigned>;
+--  std::unique_ptr<IndexMapTy> IndexMap;
+-+  uint64_t NextIndex = 0;
+-+  SmallDenseMap<void *, std::pair<OwnerTy, uint64_t>, 4> UseMap;
+- 
+- protected:
+-   ~ReplaceableUses() {
+-@@ -439,7 +429,6 @@
+-   unsigned getNumUses() const { return UseMap.size(); }
+- 
+- private:
+--  UseEntry *findRef(void *Ref);
+-   void addRef(void *Ref, OwnerTy Owner);
+-   void dropRef(void *Ref);
+-   void moveRef(void *Ref, void *New, const Metadata &MD);
+-diff -ruN --strip-trailing-cr a/llvm/lib/IR/Metadata.cpp b/llvm/lib/IR/Metadata.cpp
+---- a/llvm/lib/IR/Metadata.cpp
+-+++ b/llvm/lib/IR/Metadata.cpp
+-@@ -258,96 +258,82 @@
+- }
+- 
+- SmallVector<Metadata *> ReplaceableUses::getAllArgListUsers() {
+--  SmallVector<Metadata *> MDUsers;
+--  for (const auto &[Ref, Owner] : UseMap) {
+-+  SmallVector<std::pair<OwnerTy, uint64_t> *> MDUsersWithID;
+-+  for (auto Pair : UseMap) {
+-+    OwnerTy Owner = Pair.second.first;
+-     if (Owner.isNull())
+-       continue;
+-     if (!isa<Metadata *>(Owner))
+-       continue;
+-     Metadata *OwnerMD = cast<Metadata *>(Owner);
+-     if (OwnerMD->getMetadataID() == Metadata::DIArgListKind)
+--      MDUsers.push_back(OwnerMD);
+-+      MDUsersWithID.push_back(&UseMap[Pair.first]);
+-   }
+-+  llvm::sort(MDUsersWithID, [](auto UserA, auto UserB) {
+-+    return UserA->second < UserB->second;
+-+  });
+-+  SmallVector<Metadata *> MDUsers;
+-+  for (auto *UserWithID : MDUsersWithID)
+-+    MDUsers.push_back(cast<Metadata *>(UserWithID->first));
+-   return MDUsers;
+- }
+- 
+- SmallVector<DbgVariableRecord *>
+- ReplaceableUses::getAllDbgVariableRecordUsers() {
+-+  SmallVector<std::pair<OwnerTy, uint64_t> *> DVRUsersWithID;
+-+  for (auto Pair : UseMap) {
+-+    OwnerTy Owner = Pair.second.first;
+-+    if (Owner.isNull())
+-+      continue;
+-+    if (!isa<DebugValueUser *>(Owner))
+-+      continue;
+-+    DVRUsersWithID.push_back(&UseMap[Pair.first]);
+-+  }
+-   // Order DbgVariableRecord users in reverse-creation order. Normal dbg.value
+-   // users of MetadataAsValues are ordered by their UseList, i.e. reverse order
+-   // of when they were added: we need to replicate that here. The structure of
+-   // debug-info output depends on the ordering of intrinsics, thus we need
+-   // to keep them consistent for comparisons sake.
+-+  llvm::sort(DVRUsersWithID, [](auto UserA, auto UserB) {
+-+    return UserA->second > UserB->second;
+-+  });
+-   SmallVector<DbgVariableRecord *> DVRUsers;
+--  for (const auto &[Ref, Owner] : reverse(UseMap)) {
+--    if (Owner.isNull())
+--      continue;
+--    if (!isa<DebugValueUser *>(Owner))
+--      continue;
+--    DVRUsers.push_back(cast<DebugValueUser *>(Owner)->getUser());
+--  }
+-+  for (auto UserWithID : DVRUsersWithID)
+-+    DVRUsers.push_back(cast<DebugValueUser *>(UserWithID->first)->getUser());
+-   return DVRUsers;
+- }
+- 
+--ReplaceableUses::UseEntry *ReplaceableUses::findRef(void *Ref) {
+--  if (!IndexMap) {
+--    // Search backward for temporal locality; recently added references are
+--    // often dropped or moved first.
+--    for (UseEntry &U : reverse(UseMap))
+--      if (U.Ref == Ref)
+--        return &U;
+--    return nullptr;
+--  }
+--
+--  auto It = IndexMap->find(Ref);
+--  return It == IndexMap->end() ? nullptr : &UseMap[It->second];
+--}
+--
+- void ReplaceableUses::addRef(void *Ref, OwnerTy Owner) {
+--  assert(Ref && "Expected live reference");
+--  assert(!findRef(Ref) && "Reference already tracked");
+--  unsigned NewIdx = UseMap.size();
+--  UseMap.push_back({Ref, Owner});
+--  if (IndexMap) {
+--    (*IndexMap)[Ref] = NewIdx;
+--  } else if (UseMap.size() > IndexThreshold) {
+--    // Build the index map once UseMap grows past the threshold.
+--    IndexMap = std::make_unique<IndexMapTy>();
+--    for (unsigned I = 0, E = UseMap.size(); I != E; ++I)
+--      (*IndexMap)[UseMap[I].Ref] = I;
+--  }
+-+  bool WasInserted =
+-+      UseMap.insert(std::make_pair(Ref, std::make_pair(Owner, NextIndex)))
+-+          .second;
+-+  (void)WasInserted;
+-+  assert(WasInserted && "Expected to add a reference");
+-+
+-+  ++NextIndex;
+-+  assert(NextIndex != 0 && "Unexpected overflow");
+- }
+- 
+- void ReplaceableUses::dropRef(void *Ref) {
+--  UseEntry *Entry = findRef(Ref);
+--  assert(Entry && "Expected to find Ref");
+--  if (IndexMap)
+--    IndexMap->erase(Ref);
+--  if (Entry != &UseMap.back()) {
+--    *Entry = UseMap.back();
+--    if (IndexMap)
+--      (*IndexMap)[Entry->Ref] = Entry - UseMap.begin();
+--  }
+--  UseMap.pop_back();
+--  if (IndexMap && UseMap.size() <= IndexThreshold / 2)
+--    IndexMap.reset();
+-+  bool WasErased = UseMap.erase(Ref);
+-+  (void)WasErased;
+-+  assert(WasErased && "Expected to drop a reference");
+- }
+- 
+- void ReplaceableUses::moveRef(void *Ref, void *New, const Metadata &MD) {
+--  assert(!findRef(New) && "Cannot move to an existing reference");
+--  UseEntry *Entry = findRef(Ref);
+--  assert(Entry && "Expected to move a reference");
+--  if (IndexMap) {
+--    IndexMap->erase(Ref);
+--    (*IndexMap)[New] = Entry - UseMap.begin();
+--  }
+--  Entry->Ref = New;
+-+  auto I = UseMap.find(Ref);
+-+  assert(I != UseMap.end() && "Expected to move a reference");
+-+  auto OwnerAndIndex = I->second;
+-+  UseMap.erase(I);
+-+  bool WasInserted = UseMap.insert(std::make_pair(New, OwnerAndIndex)).second;
+-+  (void)WasInserted;
+-+  assert(WasInserted && "Expected to add a reference");
+- 
+-   // Check that the references are direct if there's no owner.
+-   (void)MD;
+--  assert((Entry->Owner || *static_cast<Metadata **>(Ref) == &MD) &&
+-+  assert((OwnerAndIndex.first || *static_cast<Metadata **>(Ref) == &MD) &&
+-          "Reference without owner must be direct");
+--  assert((Entry->Owner || *static_cast<Metadata **>(New) == &MD) &&
+-+  assert((OwnerAndIndex.first || *static_cast<Metadata **>(New) == &MD) &&
+-          "Reference without owner must be direct");
+- }
+- 
+-@@ -360,11 +346,14 @@
+-   auto &Store = Context.pImpl->ValuesAsMetadata;
+-   auto I = Store.find(&C);
+-   ValueAsMetadata *MD = I->second;
+-+  using UseTy =
+-+      std::pair<void *, std::pair<MetadataTracking::OwnerTy, uint64_t>>;
+-   // Copy out uses and update value of Constant used by debug info metadata with
+--  // poison below.
+--  SmallVector<UseEntry, 4> Uses = MD->UseMap;
+-+  // poison below
+-+  SmallVector<UseTy, 8> Uses(MD->UseMap.begin(), MD->UseMap.end());
+- 
+--  for (const auto &[Ref, Owner] : Uses) {
+-+  for (const auto &Pair : Uses) {
+-+    MetadataTracking::OwnerTy Owner = Pair.second.first;
+-     if (!Owner)
+-       continue;
+-     // Check for MetadataAsValue.
+-@@ -380,7 +369,7 @@
+-       continue;
+-     if (isa<DINode>(OwnerMD)) {
+-       OwnerMD->handleChangedOperand(
+--          Ref, ValueAsMetadata::get(PoisonValue::get(C.getType())));
+-+          Pair.first, ValueAsMetadata::get(PoisonValue::get(C.getType())));
+-     }
+-   }
+- }
+-@@ -390,21 +379,25 @@
+-     return;
+- 
+-   // Copy out uses since UseMap will get touched below.
+--  SmallVector<UseEntry, 4> Uses = UseMap;
+--
+--  for (const auto &[Ref, Owner] : Uses) {
+-+  using UseTy = std::pair<void *, std::pair<OwnerTy, uint64_t>>;
+-+  SmallVector<UseTy, 8> Uses(UseMap.begin(), UseMap.end());
+-+  llvm::sort(Uses, [](const UseTy &L, const UseTy &R) {
+-+    return L.second.second < R.second.second;
+-+  });
+-+  for (const auto &Pair : Uses) {
+-     // Check that this Ref hasn't disappeared after RAUW (when updating a
+-     // previous Ref).
+--    if (!findRef(Ref))
+-+    if (!UseMap.count(Pair.first))
+-       continue;
+- 
+-+    OwnerTy Owner = Pair.second.first;
+-     if (!Owner) {
+-       // Update unowned tracking references directly.
+--      Metadata *&DirectRef = *static_cast<Metadata **>(Ref);
+--      dropRef(Ref);
+--      DirectRef = MD;
+-+      Metadata *&Ref = *static_cast<Metadata **>(Pair.first);
+-+      Ref = MD;
+-       if (MD)
+--        MetadataTracking::track(DirectRef);
+-+        MetadataTracking::track(Ref);
+-+      UseMap.erase(Pair.first);
+-       continue;
+-     }
+- 
+-@@ -415,7 +408,7 @@
+-     }
+- 
+-     if (auto *DVU = dyn_cast<DebugValueUser *>(Owner)) {
+--      DVU->handleChangedValue(Ref, MD);
+-+      DVU->handleChangedValue(Pair.first, MD);
+-       continue;
+-     }
+- 
+-@@ -424,7 +417,7 @@
+-     switch (OwnerMD->getMetadataID()) {
+- #define HANDLE_METADATA_LEAF(CLASS)                                            \
+-   case Metadata::CLASS##Kind:                                                  \
+--    cast<CLASS>(OwnerMD)->handleChangedOperand(Ref, MD);                       \
+-+    cast<CLASS>(OwnerMD)->handleChangedOperand(Pair.first, MD);                \
+-     continue;
+- #include "llvm/IR/Metadata.def"
+-     default:
+-@@ -440,16 +433,18 @@
+- 
+-   if (!ResolveUsers) {
+-     UseMap.clear();
+--    IndexMap.reset();
+-     return;
+-   }
+- 
+--  // Move uses out since UseMap could get touched below.
+--  SmallVector<UseEntry, 4> Uses = std::move(UseMap);
+-+  // Copy out uses since UseMap could get touched below.
+-+  using UseTy = std::pair<void *, std::pair<OwnerTy, uint64_t>>;
+-+  SmallVector<UseTy, 8> Uses(UseMap.begin(), UseMap.end());
+-+  llvm::sort(Uses, [](const UseTy &L, const UseTy &R) {
+-+    return L.second.second < R.second.second;
+-+  });
+-   UseMap.clear();
+--  IndexMap.reset();
+--  for (const auto &U : Uses) {
+--    auto Owner = U.Owner;
+-+  for (const auto &Pair : Uses) {
+-+    auto Owner = Pair.second.first;
+-     if (!Owner)
+-       continue;
+-     if (!isa<Metadata *>(Owner))
+-diff -ruN --strip-trailing-cr a/llvm/lib/ObjCopy/ELF/ELFObject.cpp b/llvm/lib/ObjCopy/ELF/ELFObject.cpp
+---- a/llvm/lib/ObjCopy/ELF/ELFObject.cpp
+-+++ b/llvm/lib/ObjCopy/ELF/ELFObject.cpp
+-@@ -1709,11 +1709,7 @@
+-   case SHT_REL:
+-   case SHT_RELA:
+-   case SHT_CREL:
+--    // SHF_ALLOC relocations of an executable or shared object are copied
+--    // verbatim. Relocatable files relocations are usually non-ALLOC, but Linux
+--    // livepatch modules set SHF_ALLOC on .klp.rela.* to keep these static
+--    // relocations in memory.
+--    if ((Shdr.sh_flags & SHF_ALLOC) && ElfFile.getHeader().e_type != ET_REL) {
+-+    if (Shdr.sh_flags & SHF_ALLOC) {
+-       if (Expected<ArrayRef<uint8_t>> Data = ElfFile.getSectionContents(Shdr))
+-         return Obj.addSection<DynamicRelocationSection>(*Data);
+-       else
+-diff -ruN --strip-trailing-cr a/llvm/lib/ObjCopy/ELF/ELFObject.h b/llvm/lib/ObjCopy/ELF/ELFObject.h
+---- a/llvm/lib/ObjCopy/ELF/ELFObject.h
+-+++ b/llvm/lib/ObjCopy/ELF/ELFObject.h
+-@@ -38,7 +38,6 @@
+- class SymbolTableSection;
+- class RelocationSection;
+- class DynamicRelocationSection;
+--class DynamicSymbolTableSection;
+- class GnuDebugLinkSection;
+- class GroupSection;
+- class SectionIndexSection;
+-@@ -876,15 +875,10 @@
+- // and another which handles the symbol table type. The symbol table type is
+- // taken as a type parameter to the class (see RelocSectionWithSymtabBase).
+- class RelocationSectionBase : public SectionBase {
+--  const bool Dynamic;
+--
+- protected:
+-   SectionBase *SecToApplyRel = nullptr;
+- 
+--  explicit RelocationSectionBase(bool Dynamic) : Dynamic(Dynamic) {}
+--
+- public:
+--  bool isDynamic() const { return Dynamic; }
+-   const SectionBase *getSection() const { return SecToApplyRel; }
+-   void setSection(SectionBase *Sec) { SecToApplyRel = Sec; }
+- 
+-@@ -903,9 +897,7 @@
+-   void setSymTab(SymTabType *SymTab) { Symbols = SymTab; }
+- 
+- protected:
+--  RelocSectionWithSymtabBase()
+--      : RelocationSectionBase(
+--            std::is_same_v<SymTabType, DynamicSymbolTableSection>) {}
+-+  RelocSectionWithSymtabBase() = default;
+- 
+-   SymTabType *Symbols = nullptr;
+- 
+-@@ -936,8 +928,9 @@
+-   const Object &getObject() const { return Obj; }
+- 
+-   static bool classof(const SectionBase *S) {
+--    return RelocationSectionBase::classof(S) &&
+--           !static_cast<const RelocationSectionBase *>(S)->isDynamic();
+-+    if (S->OriginalFlags & ELF::SHF_ALLOC)
+-+      return false;
+-+    return RelocationSectionBase::classof(S);
+-   }
+- };
+- 
+-@@ -1022,8 +1015,9 @@
+-       function_ref<bool(const SectionBase *)> ToRemove) override;
+- 
+-   static bool classof(const SectionBase *S) {
+--    return RelocationSectionBase::classof(S) &&
+--           static_cast<const RelocationSectionBase *>(S)->isDynamic();
+-+    if (!(S->OriginalFlags & ELF::SHF_ALLOC))
+-+      return false;
+-+    return S->OriginalType == ELF::SHT_REL || S->OriginalType == ELF::SHT_RELA;
+-   }
+- };
+- 
+-diff -ruN --strip-trailing-cr a/llvm/test/tools/llvm-objcopy/ELF/reloc-alloc-symtab.test b/llvm/test/tools/llvm-objcopy/ELF/reloc-alloc-symtab.test
+---- a/llvm/test/tools/llvm-objcopy/ELF/reloc-alloc-symtab.test
+-+++ b/llvm/test/tools/llvm-objcopy/ELF/reloc-alloc-symtab.test
+-@@ -1,42 +0,0 @@
+--## In a relocatable file, an SHF_ALLOC relocation section (e.g. .klp.rela.* in
+--## Linux livepatch modules) is a static relocation section: its symbol indexes
+--## are updated and the symbols it references are kept.
+--
+--# RUN: yaml2obj %s -o %t
+--# RUN: llvm-strip --strip-unneeded %t -o %t1
+--# RUN: llvm-readelf -Sr -s %t1 | FileCheck %s
+--# RUN: yaml2obj -DTYPE=SHT_CREL %s -o %t.crel
+--# RUN: llvm-strip --strip-unneeded %t.crel -o %t1.crel
+--# RUN: llvm-readelf -Sr -s %t1.crel | FileCheck %s
+--
+--# CHECK:      [ 2] .klp.rela.vmlinux.text {{RELA|CREL}} {{.*}} AI  3   1
+--# CHECK-NEXT: [ 3] .symtab
+--# CHECK:      R_X86_64_PC32 {{.*}} referenced + 0
+--# CHECK:      Symbol table '.symtab' contains 2 entries:
+--# CHECK:      1: {{.*}} UND referenced
+--
+--!ELF
+--FileHeader:
+--  Class:   ELFCLASS64
+--  Data:    ELFDATA2LSB
+--  Type:    ET_REL
+--  Machine: EM_X86_64
+--Sections:
+--  - Name:    .text
+--    Type:    SHT_PROGBITS
+--    Flags:   [ SHF_ALLOC, SHF_EXECINSTR ]
+--    Content: "e800000000"
+--  - Name:    .klp.rela.vmlinux.text
+--    Type:    [[TYPE=SHT_RELA]]
+--    Flags:   [ SHF_ALLOC, SHF_INFO_LINK ]
+--    Link:    .symtab
+--    Info:    .text
+--    Relocations:
+--      - Offset: 1
+--        Symbol: referenced
+--        Type:   R_X86_64_PC32
+--Symbols:
+--  - Name:    unreferenced
+--    Section: .text
+--  - Name:    referenced
+--    Binding: STB_GLOBAL
+-diff -ruN --strip-trailing-cr a/llvm/test/tools/llvm-objcopy/ELF/reloc-no-symtab.test b/llvm/test/tools/llvm-objcopy/ELF/reloc-no-symtab.test
+---- a/llvm/test/tools/llvm-objcopy/ELF/reloc-no-symtab.test
+-+++ b/llvm/test/tools/llvm-objcopy/ELF/reloc-no-symtab.test
+-@@ -10,16 +10,11 @@
+- # RUN: llvm-readobj --sections %t.original | FileCheck %s
+- # RUN: llvm-readobj --sections %t.stripped | FileCheck %s
+- 
+--## In a relocatable file, .rela.plt is a static relocation section.
+--# RUN: yaml2obj -DTYPE=ET_REL %s -o %t.rel
+--# RUN: llvm-strip %t.rel -o %t.rel.stripped
+--# RUN: llvm-readobj --sections %t.rel.stripped | FileCheck %s
+--
+- !ELF
+- FileHeader:
+-   Class:           ELFCLASS64
+-   Data:            ELFDATA2LSB
+--  Type:            [[TYPE=ET_EXEC]]
+-+  Type:            ET_REL
+-   Machine:         EM_X86_64
+- Sections:
+-   - Name:          .text
+-diff -ruN --strip-trailing-cr a/llvm/unittests/IR/MetadataTest.cpp b/llvm/unittests/IR/MetadataTest.cpp
+---- a/llvm/unittests/IR/MetadataTest.cpp
+-+++ b/llvm/unittests/IR/MetadataTest.cpp
+-@@ -5355,29 +5355,6 @@
+-   EXPECT_TRUE(!MD);
+- }
+- 
+--// Once the use list is indexed, look each changed reference up again right
+--// away: a later swap in dropRef would repair a stale index entry.
+--TEST_F(TrackingMDRefTest, LargeUseList) {
+--  auto Temp = MDTuple::getTemporary(Context, {});
+--  TrackingMDRef Refs[41];
+--  for (unsigned I = 0; I != 40; ++I)
+--    Refs[I].reset(Temp.get());
+--  Refs[0].reset();
+--  Refs[0].reset(Temp.get());
+--  Refs[0].reset();
+--  Refs[40] = std::move(Refs[1]);
+--  Refs[1].reset(Temp.get());
+--  Refs[40].reset();
+--  Refs[39].reset();
+--  EXPECT_EQ(38u, Temp->getNumTemporaryUses());
+--
+--  auto *N = MDTuple::getDistinct(Context, {});
+--  Temp->replaceAllUsesWith(N);
+--  EXPECT_EQ(0u, Temp->getNumTemporaryUses());
+--  for (unsigned I = 1; I != 39; ++I)
+--    EXPECT_EQ(N, Refs[I].get()) << I;
+--}
+--
+- TEST(NamedMDNodeTest, Search) {
+-   LLVMContext Context;
+-   ConstantAsMetadata *C =
+diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
+index 9de7c05a..7b8ac881 100644
+--- a/third_party/llvm/workspace.bzl
++++ b/third_party/llvm/workspace.bzl
+@@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
+ 
+ def repo(name):
+     """Imports LLVM."""
+-    LLVM_COMMIT = "62e30e3326d9b4e3948d242ccfee7aa4fc030ef3"
+-    LLVM_SHA256 = "478ae34c463f28708582082098aaf0e79f431c2719f5d297bd546051dbeae0dd"
++    LLVM_COMMIT = "c4d9267a31897dd76b2aab5196b75bffb49e6846"
++    LLVM_SHA256 = "a96c8b33f0d8d3be008b8c9f74ef310254c4f7f7edf8c04c76f4d99ed9d34f1a"
+ 
+     tf_http_archive(
+         name = name,
+diff --git a/third_party/stablehlo/temporary.patch b/third_party/stablehlo/temporary.patch
+index c38c5147..c6cd6f2b 100644
+--- a/third_party/stablehlo/temporary.patch
++++ b/third_party/stablehlo/temporary.patch
+@@ -881,9 +881,9 @@ diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
+ +Performs element-wise base-2 exponential operation on `operand` tensor and
+ +produces a `result` tensor. Depending on the element type, does the following:
+ +
+-+*   For floats: `exp2` from IEEE-754.
+-+*   For complex numbers: complex base-2 exponential.
+-+*   For quantized types: `dequantize_op_quantize(exp2, operand, type(result))`.
+++* For floats: `exp2` from IEEE-754.
+++* For complex numbers: complex base-2 exponential.
+++* For quantized types: `dequantize_op_quantize(exp2, operand, type(result))`.
+ +
+ +#### Inputs
+ +
+@@ -907,7 +907,7 @@ diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
+ +
+ +#### Constraints
+ +
+-+*   (C1) `baseline_type(operand) = baseline_type(result)`.
+++* (C1) `baseline_type(operand) = baseline_type(result)`.
+ +
+ +#### Examples
+ +
+@@ -947,9 +947,9 @@ diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
+ +Performs element-wise base-2 logarithm operation on `operand` tensor and
+ +produces a `result` tensor. Depending on the element type, does the following:
+ +
+-+*   For floats: `log2` from IEEE-754.
+-+*   For complex numbers: complex base-2 logarithm.
+-+*   For quantized types: `dequantize_op_quantize(log2, operand, type(result))`.
+++* For floats: `log2` from IEEE-754.
+++* For complex numbers: complex base-2 logarithm.
+++* For quantized types: `dequantize_op_quantize(log2, operand, type(result))`.
+ +
+ +#### Inputs
+ +
+@@ -973,7 +973,7 @@ diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
+ +
+ +#### Constraints
+ +
+-+*   (C1) `baseline_type(operand) = baseline_type(result)`.
+++* (C1) `baseline_type(operand) = baseline_type(result)`.
+ +
+ +#### Examples
+ +
+@@ -1018,22 +1018,24 @@ diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
+  * (C1) `element_type(operand) = element_type(source)`.
+  * (C2) `shape(source) = num_windows` where:
+    * `padded_operand_shape = padding[:, 0] + shape(operand) + padding[:, 1]`.
+-@@ -7501,6 +7597,7 @@
++@@ -7501,7 +7597,7 @@
+  syntax ([ISO/IEC 14977:1996](https://www.iso.org/standard/26153.html),
+  [Wikipedia](https://en.wikipedia.org/wiki/Extended_Backus%E2%80%93Naur_form)),
+  with two modifications: 1) rules are defined using `::=` rather than `=`,
+-+
+- 2) concatenation is expressed using juxtaposition rather than `,`.
++-2) concatenation is expressed using juxtaposition rather than `,`.
+++2\) concatenation is expressed using juxtaposition rather than `,`.
+  
+  For describing semantics (i.e. within "Types", "Constants" and "Ops" sections),
+-@@ -7571,6 +7668,7 @@
++ we are using formulas which are based on Python syntax extended with support
++@@ -7571,7 +7667,7 @@
+  ### Names
+  
+  In formulas, lexical scope includes: 1) global functions, 2) member definitions,
+-+
+- 3) local definitions. The list of global functions is provided below. The list
++-3) local definitions. The list of global functions is provided below. The list
+++3\) local definitions. The list of global functions is provided below. The list
+  of element definitions depends on the program element that the notation is
+  applied to:
++ 
+ diff --ruN a/stablehlo/rfcs/20260821-low-precision-complex.md b/stablehlo/rfcs/20260821-low-precision-complex.md
+ --- stablehlo/rfcs/20260821-low-precision-complex.md
+ +++ stablehlo/rfcs/20260821-low-precision-complex.md
+@@ -1657,16 +1659,16 @@ diff --ruN a/stablehlo/rfcs/20260907-exp2-log2.md b/stablehlo/rfcs/20260907-exp2
+ +
+ +Currently, StableHLO lacks native `exp2` and `log2` operations:
+ +
+-+*   Frameworks expose `exp2` and `log2` directly: JAX provides
+-+    [`jnp.exp2`][jax1] / `lax.exp2` and [`jnp.log2`][jax2] / `lax.log2`, and
+-+    PyTorch provides [`torch.exp2`][torch1] and [`torch.log2`][torch2].
+-+*   Without first-class ops in StableHLO, frameworks must decompose `exp2(x)`
+-+    into `exponential(x * ln(2))` and `log2(x)` into `log(x) / ln(2)`.
+-+*   Decomposing `exp2` and `log2` into natural `exponential` and `log` scaled by
+-+    `ln(2)` degrades numerical accuracy (introducing rounding error on integer
+-+    powers of two where `exp2(n)` and `log2(2^n)` should evaluate exactly) and
+-+    prevents backends from emitting single-instruction base-2 hardware
+-+    operations without fragile algebraic pattern-matching.
+++* Frameworks expose `exp2` and `log2` directly: JAX provides [`jnp.exp2`][jax1]
+++  / `lax.exp2` and [`jnp.log2`][jax2] / `lax.log2`, and PyTorch provides
+++  [`torch.exp2`][torch1] and [`torch.log2`][torch2].
+++* Without first-class ops in StableHLO, frameworks must decompose `exp2(x)` into
+++  `exponential(x * ln(2))` and `log2(x)` into `log(x) / ln(2)`.
+++* Decomposing `exp2` and `log2` into natural `exponential` and `log` scaled by
+++  `ln(2)` degrades numerical accuracy (introducing rounding error on integer
+++  powers of two where `exp2(n)` and `log2(2^n)` should evaluate exactly) and
+++  prevents backends from emitting single-instruction base-2 hardware operations
+++  without fragile algebraic pattern-matching.
+ +
+ +Following the precedent of `TanOp` ([20240312-tanop.md](20240312-tanop.md)) and
+ +the transcendental accuracy specification
+@@ -1681,12 +1683,12 @@ diff --ruN a/stablehlo/rfcs/20260907-exp2-log2.md b/stablehlo/rfcs/20260907-exp2
+ +
+ +## Exp2Op, Log2Op, and Compatibility
+ +
+-+*   **VHLO Serialization**: StableHLO version `1.21.0` introduces
+-+    `VHLO_Exp2OpV1` (`exp2_v1`) and `VHLO_Log2OpV1` (`log2_v1`).
+-+*   **Forward Compatibility & Decomposition**: For runtimes targeting StableHLO
+-+    versions prior to `1.21.0`, compatibility expansion lowers `exp2(x)` to
+-+    `exponential(multiply(x, constant(ln(2))))` and `log2(x)` to `divide(log(x),
+-+    constant(ln(2)))`.
+++* **VHLO Serialization**: StableHLO version `1.21.0` introduces `VHLO_Exp2OpV1`
+++  (`exp2_v1`) and `VHLO_Log2OpV1` (`log2_v1`).
+++* **Forward Compatibility & Decomposition**: For runtimes targeting StableHLO
+++  versions prior to `1.21.0`, compatibility expansion lowers `exp2(x)` to
+++  `exponential(multiply(x, constant(ln(2))))` and `log2(x)` to
+++  `divide(log(x), constant(ln(2)))`.
+ +
+ +[jax1]: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.exp2.html
+ +[jax2]: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.log2.html
+@@ -3698,15 +3700,30 @@ diff --ruN a/stablehlo/stablehlo/reference/InterpreterOps.cpp b/stablehlo/stable
+  
+  #include "llvm/Support/Casting.h"
+  #include "llvm/Support/Errc.h"
+-@@ -175,36 +176,39 @@
++@@ -174,37 +175,52 @@
++     ArrayRef<InterpreterValue> inputs, std::queue<StringAttr>& infeed,
+      SmallVector<SmallVector<StringAttr>> programs, SymbolTable& symbolTable,
+      InterpreterFallback* fallback) {
+-   llvm::DefaultThreadPool threadPool;
++-  llvm::DefaultThreadPool threadPool;
+ -  SmallVector<std::shared_future<SmallVector<InterpreterValue>>> futures;
+-+  llvm::ThreadPoolTaskGroup taskGroup(threadPool);
+- 
++-
+    uint32_t numReplicas = programs.size();
+    uint32_t numPartitions = programs[0].size();
+++
+++  // StableHLO collectives are blocking barriers: a process inside rendezvous
+++  // waits for every other process in its group to arrive, so all processes
+++  // must run concurrently. Size the pool from the process count rather than
+++  // the default (hardware_concurrency(), i.e. the host thread count). With a
+++  // default-sized pool, a program with more processes than the host has
+++  // hardware threads leaves the surplus tasks queued and never started while
+++  // the started tasks block forever waiting for them - a deadlock that only
+++  // surfaces as the rendezvous timeout. hardware_concurrency(N) leaves
+++  // ThreadPoolStrategy::Limit false, so the pool gets N threads uncapped by the
+++  // host.
+++  llvm::DefaultThreadPool threadPool(
+++      llvm::hardware_concurrency(numReplicas * numPartitions));
+++  llvm::ThreadPoolTaskGroup taskGroup(threadPool);
+++
+    ProcessGrid processGrid(numReplicas, numPartitions, infeed);
+  
+ +  SmallVector<SmallVector<InterpreterValue>> taskOutputs(numReplicas *
+@@ -4443,6 +4460,159 @@ diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compat
+  
+    // CHECK-LABEL: @all_reduce_sdy_mesh_max
+    func.func @all_reduce_sdy_mesh_max(%arg0: tensor<4xf32>) -> tensor<4xf32> {
++diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir b/stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir
++--- stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir
+++++ stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir
++@@ -168,6 +168,18 @@
++ 
++ // -----
++ 
+++// CHECK-LABEL: func @dynamic_broadcast_in_dim_discardable_attrs
+++func.func @dynamic_broadcast_in_dim_discardable_attrs(%arg0: tensor<4xf32>) -> tensor<3x4xf32> {
+++  //  CHECK-NOT: stablehlo.dynamic_broadcast_in_dim
+++  //      CHECK: stablehlo.broadcast_in_dim %arg0, dims = [1]
+++  // CHECK-SAME:   {mhlo.frontend_attributes = {MUST_FUSE = "0"}}
+++  %0 = stablehlo.constant dense<[3, 4]> : tensor<2xi64>
+++  %1 = stablehlo.dynamic_broadcast_in_dim %arg0, %0, dims = [1] {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<4xf32>, tensor<2xi64>) -> tensor<3x4xf32>
+++  return %1 : tensor<3x4xf32>
+++}
+++
+++// -----
+++
++ // CHECK-LABEL: func @dynamic_broadcast_in_dim_inapplicable_dynamic_operand_type
++ func.func @dynamic_broadcast_in_dim_inapplicable_dynamic_operand_type(%arg0: tensor<?xf32>) -> tensor<3x4xf32> {
++   // CHECK: stablehlo.dynamic_broadcast_in_dim
++@@ -225,6 +237,26 @@
++ 
++ // -----
++ 
+++// CHECK-LABEL: func @dynamic_conv_discardable_attrs
+++func.func @dynamic_conv_discardable_attrs(%arg0: tensor<100x26x26x32xf32>, %arg1: tensor<3x3x1x32xf32>) -> tensor<100x28x28x1xf32> {
+++  //  CHECK-NOT: stablehlo.dynamic_conv
+++  //      CHECK: stablehlo.convolution(%arg0, %arg1)
+++  // CHECK-SAME:   mhlo.frontend_attributes = {MUST_FUSE = "0"}
+++  %0 = stablehlo.constant dense<2> : tensor<2x2xi32>
+++  %1 = "stablehlo.dynamic_conv"(%arg0, %arg1, %0) {
+++    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>,
+++    window_strides = array<i64: 1, 1>,
+++    lhs_dilation = array<i64: 1, 1>,
+++    rhs_dilation = array<i64: 1, 1>,
+++    feature_group_count = 1 : i64,
+++    batch_group_count = 1 : i64,
+++    mhlo.frontend_attributes = {MUST_FUSE = "0"}
+++  } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf32>, tensor<2x2xi32>) -> tensor<100x28x28x1xf32>
+++  return %1 : tensor<100x28x28x1xf32>
+++}
+++
+++// -----
+++
++ // CHECK-LABEL: func @dynamic_conv_success_dynamic_result_type
++ func.func @dynamic_conv_success_dynamic_result_type(%arg0: tensor<100x26x26x32xf32>, %arg1: tensor<3x3x1x32xf32>) -> tensor<?x28x28x1xf32> {
++   //  CHECK-NOT: stablehlo.dynamic_conv
++@@ -295,6 +327,25 @@
++ 
++ // -----
++ 
+++// CHECK-LABEL: @dynamic_gather_discardable_attrs
+++func.func @dynamic_gather_discardable_attrs(%arg0 : tensor<2x4x9xi32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x8xi32> {
+++  //  CHECK-NOT: stablehlo.dynamic_gather
+++  //      CHECK: "stablehlo.gather"(%arg0, %arg1)
+++  // CHECK-SAME:   {mhlo.frontend_attributes = {MUST_FUSE = "0"}}
+++  %0 = stablehlo.constant dense<[1, 1, 8]> : tensor<3xi32>
+++  %1 = "stablehlo.dynamic_gather"(%arg0, %arg1, %0) <{
+++    dimension_numbers = #stablehlo.gather<
+++      collapsed_slice_dims = [0, 1],
+++      index_vector_dim = 2,
+++      offset_dims = [2],
+++      start_index_map = [0, 1]
+++    >
+++  }> {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<2x4x9xi32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<1x5x8xi32>
+++  return %1 : tensor<1x5x8xi32>
+++}
+++
+++// -----
+++
++ // CHECK-LABEL: @dynamic_gather_success_dynamic_result_type
++ func.func @dynamic_gather_success_dynamic_result_type(%arg0 : tensor<2x4x9xi32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x?xi32> {
++   //  CHECK-NOT: stablehlo.dynamic_gather
++@@ -348,6 +399,18 @@
++ 
++ // -----
++ 
+++// CHECK-LABEL: func @dynamic_iota_discardable_attrs
+++func.func @dynamic_iota_discardable_attrs() -> tensor<4xf32> {
+++  //  CHECK-NOT: stablehlo.dynamic_iota
+++  //      CHECK: stablehlo.iota
+++  // CHECK-SAME:   {mhlo.frontend_attributes = {MUST_FUSE = "0"}}
+++  %0 = stablehlo.constant dense<4> : tensor<1xi64>
+++  %1 = stablehlo.dynamic_iota %0, dim = 0 {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<1xi64>) -> tensor<4xf32>
+++  return %1 : tensor<4xf32>
+++}
+++
+++// -----
+++
++ // CHECK-LABEL: func @dynamic_iota_inapplicable_dynamic_output_shape
++ func.func @dynamic_iota_inapplicable_dynamic_output_shape(%arg0: tensor<1xi64>) -> tensor<4xf32> {
++   // CHECK: stablehlo.dynamic_iota
++@@ -380,6 +443,20 @@
++ 
++ // -----
++ 
+++// CHECK-LABEL: func @dynamic_pad_discardable_attrs
+++func.func @dynamic_pad_discardable_attrs(%arg0: tensor<4xf32>, %arg1: tensor<f32>) -> tensor<11xf32> {
+++  //  CHECK-NOT: stablehlo.dynamic_pad
+++  //      CHECK: stablehlo.pad
+++  // CHECK-SAME:   {mhlo.frontend_attributes = {MUST_FUSE = "0"}}
+++  %0 = stablehlo.constant dense<0> : tensor<1xi64>
+++  %1 = stablehlo.constant dense<1> : tensor<1xi64>
+++  %2 = stablehlo.constant dense<2> : tensor<1xi64>
+++  %3 = stablehlo.dynamic_pad %arg0, %arg1, %0, %1, %2 {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<4xf32>, tensor<f32>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> tensor<11xf32>
+++  return %3 : tensor<11xf32>
+++}
+++
+++// -----
+++
++ // CHECK-LABEL: func @dynamic_pad_success_dynamic_result_type
++ func.func @dynamic_pad_success_dynamic_result_type(%arg0: tensor<4xf32>, %arg1: tensor<f32>) -> tensor<?xf32> {
++   // CHECK-NOT: stablehlo.dynamic_pad
++@@ -437,6 +514,17 @@
++ 
++ // -----
++ 
+++// CHECK-LABEL: func @dynamic_reshape_discardable_attrs
+++func.func @dynamic_reshape_discardable_attrs(%arg0: tensor<4xf32>) -> tensor<1x4xf32> {
+++  //  CHECK-NOT: stablehlo.dynamic_reshape
+++  //      CHECK: stablehlo.reshape %arg0 {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<4xf32>) -> tensor<1x4xf32>
+++  %0 = stablehlo.constant dense<[1, 4]> : tensor<2xi64>
+++  %1 = stablehlo.dynamic_reshape %arg0, %0 {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<4xf32>, tensor<2xi64>) -> tensor<1x4xf32>
+++  return %1 : tensor<1x4xf32>
+++}
+++
+++// -----
+++
++ // CHECK-LABEL: func @dynamic_reshape_inapplicable_dynamic_output_shape
++ func.func @dynamic_reshape_inapplicable_dynamic_output_shape(%arg0: tensor<4xf32>, %arg1: tensor<2xi64>) -> tensor<1x4xf32> {
++   // CHECK: stablehlo.dynamic_reshape
++diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_legalize_deprecated_ops.mlir b/stablehlo/stablehlo/tests/transforms/stablehlo_legalize_deprecated_ops.mlir
++--- stablehlo/stablehlo/tests/transforms/stablehlo_legalize_deprecated_ops.mlir
+++++ stablehlo/stablehlo/tests/transforms/stablehlo_legalize_deprecated_ops.mlir
++@@ -10,6 +10,16 @@
++ 
++ // -----
++ 
+++// CHECK-LABEL: op_dot_discardable_attrs
+++func.func @op_dot_discardable_attrs(%arg0: tensor<2x3xf32>,
+++                 %arg1: tensor<3x?xf32>) -> tensor<2x?xf32> {
+++  // CHECK: stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<2x3xf32>, tensor<3x?xf32>) -> tensor<2x?xf32>
+++  %0 = "stablehlo.dot"(%arg0, %arg1) {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<2x3xf32>, tensor<3x?xf32>) -> tensor<2x?xf32>
+++  func.return %0 : tensor<2x?xf32>
+++}
+++
+++// -----
+++
++ // CHECK-LABEL: op_unary_einsum
++ func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
++   // CHECK:      %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
+ diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir b/stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir
+ --- stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir
+ +++ stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir
+@@ -8393,6 +8563,99 @@ diff --ruN a/stablehlo/stablehlo/transforms/StablehloCanonicalizeDynamism.cpp b/
+      return success();
+    }
+  };
++@@ -119,8 +120,11 @@
++                                          "expected static output_dimensions");
++     if (!op.getType().hasStaticShape())
++       return rewriter.notifyMatchFailure(op, "expected static result type");
++-    rewriter.replaceOpWithNewOp<BroadcastInDimOp>(
++-        op, op.getType(), op.getOperand(), op.getBroadcastDimensionsAttr());
+++    auto broadcastInDimOp = BroadcastInDimOp::create(
+++        rewriter, op.getLoc(), op.getType(), op.getOperand(),
+++        op.getBroadcastDimensionsAttr());
+++    broadcastInDimOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
+++    rewriter.replaceOp(op, broadcastInDimOp);
++     return success();
++   }
++ };
++@@ -139,12 +143,14 @@
++         RankedTensorType::get({static_cast<int64_t>(padding.size()) / 2, 2},
++                               rewriter.getI64Type()),
++         padding);
++-    rewriter.replaceOpWithNewOp<ConvolutionOp>(
++-        op, op.getType(), op.getLhs(), op.getRhs(), op.getWindowStridesAttr(),
++-        paddingAttr, op.getLhsDilationAttr(), op.getRhsDilationAttr(),
++-        op.getWindowReversalAttr(), op.getDimensionNumbers(),
++-        op.getFeatureGroupCount(), op.getBatchGroupCount(),
++-        op.getPrecisionConfigAttr());
+++    auto convolutionOp = ConvolutionOp::create(
+++        rewriter, op.getLoc(), op.getType(), op.getLhs(), op.getRhs(),
+++        op.getWindowStridesAttr(), paddingAttr, op.getLhsDilationAttr(),
+++        op.getRhsDilationAttr(), op.getWindowReversalAttr(),
+++        op.getDimensionNumbers(), op.getFeatureGroupCount(),
+++        op.getBatchGroupCount(), op.getPrecisionConfigAttr());
+++    convolutionOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
+++    rewriter.replaceOp(op, convolutionOp);
++     return success();
++   }
++ };
++@@ -159,10 +165,13 @@
++     SmallVector<int64_t> sliceSizes;
++     if (!succeeded(hlo::matchInts(op.getSliceSizes(), sliceSizes)))
++       return rewriter.notifyMatchFailure(op, "expected static slice_sizes");
++-    rewriter.replaceOpWithNewOp<GatherOp>(
++-        op, op.getType(), op.getOperand(), op.getStartIndices(),
++-        op.getDimensionNumbersAttr(), rewriter.getDenseI64ArrayAttr(sliceSizes),
++-        op.getIndicesAreSortedAttr());
+++    auto gatherOp =
+++        GatherOp::create(rewriter, op.getLoc(), op.getType(), op.getOperand(),
+++                         op.getStartIndices(), op.getDimensionNumbersAttr(),
+++                         rewriter.getDenseI64ArrayAttr(sliceSizes),
+++                         op.getIndicesAreSortedAttr());
+++    gatherOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
+++    rewriter.replaceOp(op, gatherOp);
++     return success();
++   }
++ };
++@@ -179,8 +188,10 @@
++       return rewriter.notifyMatchFailure(op, "expected static output_shape");
++     if (!op.getType().hasStaticShape())
++       return rewriter.notifyMatchFailure(op, "expected static result type");
++-    rewriter.replaceOpWithNewOp<IotaOp>(op, op.getType(),
++-                                        op.getIotaDimension());
+++    auto iotaOp = IotaOp::create(rewriter, op.getLoc(), op.getType(),
+++                                 op.getIotaDimension());
+++    iotaOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
+++    rewriter.replaceOp(op, iotaOp);
++     return success();
++   }
++ };
++@@ -198,9 +209,11 @@
++       return rewriter.notifyMatchFailure(op, "expected static high");
++     if (!succeeded(hlo::matchInts(op.getInteriorPadding(), interiorPadding)))
++       return rewriter.notifyMatchFailure(op, "expected static interior");
++-    rewriter.replaceOpWithNewOp<PadOp>(op, op.getType(), op.getOperand(),
++-                                       op.getPaddingValue(), edgePaddingLow,
++-                                       edgePaddingHigh, interiorPadding);
+++    auto padOp = PadOp::create(
+++        rewriter, op.getLoc(), op.getType(), op.getOperand(),
+++        op.getPaddingValue(), edgePaddingLow, edgePaddingHigh, interiorPadding);
+++    padOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
+++    rewriter.replaceOp(op, padOp);
++     return success();
++   }
++ };
++@@ -216,7 +229,10 @@
++       return rewriter.notifyMatchFailure(op, "expected static output_shape");
++     if (!op.getType().hasStaticShape())
++       return rewriter.notifyMatchFailure(op, "expected static result type");
++-    rewriter.replaceOpWithNewOp<ReshapeOp>(op, op.getType(), op.getOperand());
+++    auto reshapeOp =
+++        ReshapeOp::create(rewriter, op.getLoc(), op.getType(), op.getOperand());
+++    reshapeOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
+++    rewriter.replaceOp(op, reshapeOp);
++     return success();
++   }
++ };
+ diff --ruN a/stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp b/stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp
+ --- stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp
+ +++ stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp
+@@ -8433,6 +8696,42 @@ diff --ruN a/stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp b/s
+    populateWithGenerated(*patterns);
+  }
+  
++diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOps.cpp b/stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOps.cpp
++--- stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOps.cpp
+++++ stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOps.cpp
++@@ -132,6 +132,14 @@
++   llvm::report_fatal_error("unsupported type");
++ }
++ 
+++void copyDiscardableAttrs(Operation* src, Operation* dst) {
+++  dst->setDiscardableAttrs(src->getDiscardableAttrDictionary());
+++}
+++
+++void copyDiscardableAttrs(Value src, Value dst) {
+++  copyDiscardableAttrs(src.getDefiningOp(), dst.getDefiningOp());
+++}
+++
++ #include "stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.h.inc"
++ }  // namespace
++ 
++diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.td b/stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.td
++--- stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.td
+++++ stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.td
++@@ -45,9 +45,12 @@
++ 
++ // def CrossReplicaSumToAllReduce : <Declared in cpp>
++ 
+++def CopyDiscardableAttrs : NativeCodeCallVoid<"copyDiscardableAttrs($0, $1)">;
+++
++ def DotToDotGeneral : Pat<
++-  (StableHLO_DotOp $lhs, $rhs, $precision_config),
++-  (StableHLO_DotGeneralOp $lhs, $rhs, (GetDefaultDotDimensionNumbers $lhs), $precision_config, (NativeCodeCall<"Attribute{}">))>;
+++  (StableHLO_DotOp:$src $lhs, $rhs, $precision_config),
+++  (StableHLO_DotGeneralOp:$dst $lhs, $rhs, (GetDefaultDotDimensionNumbers $lhs), $precision_config, (NativeCodeCall<"Attribute{}">)),
+++  [], [(CopyDiscardableAttrs $src, $dst)]>;
++ 
++ def UnaryEinsumToEinsum : Pat<
++   (StableHLO_UnaryEinsumOp $operand, $equation),
+ diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp b/stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp
+ --- stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp
+ +++ stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp
diff --git a/third_party/xla/third_party/shardy/workspace.bzl b/third_party/xla/third_party/shardy/workspace.bzl
index b98756b3abe..7a739adf488 100644
--- a/third_party/xla/third_party/shardy/workspace.bzl
+++ b/third_party/xla/third_party/shardy/workspace.bzl
@@ -18,8 +18,8 @@
 load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 
 def repo():
-    SHARDY_COMMIT = "f01c1a6b084c089087827633333615546d990867"
-    SHARDY_SHA256 = "f85cc219fad3db16e92faaa1ceb214c22ba3752f559780beb00cba7809c530a5"
+    SHARDY_COMMIT = "1ecfcb6250edeb638fce0e513073b9f52047ad83"
+    SHARDY_SHA256 = "726b288c47393ea0567b4655aa6516c523eefeb7db724bfc34d9cd94fff15b02"
 
     tf_http_archive(
         name = "shardy",
diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch
index c38c514743b..c6cd6f2bc9c 100644
--- a/third_party/xla/third_party/stablehlo/temporary.patch
+++ b/third_party/xla/third_party/stablehlo/temporary.patch
@@ -881,9 +881,9 @@ diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
 +Performs element-wise base-2 exponential operation on `operand` tensor and
 +produces a `result` tensor. Depending on the element type, does the following:
 +
-+*   For floats: `exp2` from IEEE-754.
-+*   For complex numbers: complex base-2 exponential.
-+*   For quantized types: `dequantize_op_quantize(exp2, operand, type(result))`.
++* For floats: `exp2` from IEEE-754.
++* For complex numbers: complex base-2 exponential.
++* For quantized types: `dequantize_op_quantize(exp2, operand, type(result))`.
 +
 +#### Inputs
 +
@@ -907,7 +907,7 @@ diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
 +
 +#### Constraints
 +
-+*   (C1) `baseline_type(operand) = baseline_type(result)`.
++* (C1) `baseline_type(operand) = baseline_type(result)`.
 +
 +#### Examples
 +
@@ -947,9 +947,9 @@ diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
 +Performs element-wise base-2 logarithm operation on `operand` tensor and
 +produces a `result` tensor. Depending on the element type, does the following:
 +
-+*   For floats: `log2` from IEEE-754.
-+*   For complex numbers: complex base-2 logarithm.
-+*   For quantized types: `dequantize_op_quantize(log2, operand, type(result))`.
++* For floats: `log2` from IEEE-754.
++* For complex numbers: complex base-2 logarithm.
++* For quantized types: `dequantize_op_quantize(log2, operand, type(result))`.
 +
 +#### Inputs
 +
@@ -973,7 +973,7 @@ diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
 +
 +#### Constraints
 +
-+*   (C1) `baseline_type(operand) = baseline_type(result)`.
++* (C1) `baseline_type(operand) = baseline_type(result)`.
 +
 +#### Examples
 +
@@ -1018,22 +1018,24 @@ diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
  * (C1) `element_type(operand) = element_type(source)`.
  * (C2) `shape(source) = num_windows` where:
    * `padded_operand_shape = padding[:, 0] + shape(operand) + padding[:, 1]`.
-@@ -7501,6 +7597,7 @@
+@@ -7501,7 +7597,7 @@
  syntax ([ISO/IEC 14977:1996](https://www.iso.org/standard/26153.html),
  [Wikipedia](https://en.wikipedia.org/wiki/Extended_Backus%E2%80%93Naur_form)),
  with two modifications: 1) rules are defined using `::=` rather than `=`,
-+
- 2) concatenation is expressed using juxtaposition rather than `,`.
+-2) concatenation is expressed using juxtaposition rather than `,`.
++2\) concatenation is expressed using juxtaposition rather than `,`.
  
  For describing semantics (i.e. within "Types", "Constants" and "Ops" sections),
-@@ -7571,6 +7668,7 @@
+ we are using formulas which are based on Python syntax extended with support
+@@ -7571,7 +7667,7 @@
  ### Names
  
  In formulas, lexical scope includes: 1) global functions, 2) member definitions,
-+
- 3) local definitions. The list of global functions is provided below. The list
+-3) local definitions. The list of global functions is provided below. The list
++3\) local definitions. The list of global functions is provided below. The list
  of element definitions depends on the program element that the notation is
  applied to:
+ 
 diff --ruN a/stablehlo/rfcs/20260821-low-precision-complex.md b/stablehlo/rfcs/20260821-low-precision-complex.md
 --- stablehlo/rfcs/20260821-low-precision-complex.md
 +++ stablehlo/rfcs/20260821-low-precision-complex.md
@@ -1657,16 +1659,16 @@ diff --ruN a/stablehlo/rfcs/20260907-exp2-log2.md b/stablehlo/rfcs/20260907-exp2
 +
 +Currently, StableHLO lacks native `exp2` and `log2` operations:
 +
-+*   Frameworks expose `exp2` and `log2` directly: JAX provides
-+    [`jnp.exp2`][jax1] / `lax.exp2` and [`jnp.log2`][jax2] / `lax.log2`, and
-+    PyTorch provides [`torch.exp2`][torch1] and [`torch.log2`][torch2].
-+*   Without first-class ops in StableHLO, frameworks must decompose `exp2(x)`
-+    into `exponential(x * ln(2))` and `log2(x)` into `log(x) / ln(2)`.
-+*   Decomposing `exp2` and `log2` into natural `exponential` and `log` scaled by
-+    `ln(2)` degrades numerical accuracy (introducing rounding error on integer
-+    powers of two where `exp2(n)` and `log2(2^n)` should evaluate exactly) and
-+    prevents backends from emitting single-instruction base-2 hardware
-+    operations without fragile algebraic pattern-matching.
++* Frameworks expose `exp2` and `log2` directly: JAX provides [`jnp.exp2`][jax1]
++  / `lax.exp2` and [`jnp.log2`][jax2] / `lax.log2`, and PyTorch provides
++  [`torch.exp2`][torch1] and [`torch.log2`][torch2].
++* Without first-class ops in StableHLO, frameworks must decompose `exp2(x)` into
++  `exponential(x * ln(2))` and `log2(x)` into `log(x) / ln(2)`.
++* Decomposing `exp2` and `log2` into natural `exponential` and `log` scaled by
++  `ln(2)` degrades numerical accuracy (introducing rounding error on integer
++  powers of two where `exp2(n)` and `log2(2^n)` should evaluate exactly) and
++  prevents backends from emitting single-instruction base-2 hardware operations
++  without fragile algebraic pattern-matching.
 +
 +Following the precedent of `TanOp` ([20240312-tanop.md](20240312-tanop.md)) and
 +the transcendental accuracy specification
@@ -1681,12 +1683,12 @@ diff --ruN a/stablehlo/rfcs/20260907-exp2-log2.md b/stablehlo/rfcs/20260907-exp2
 +
 +## Exp2Op, Log2Op, and Compatibility
 +
-+*   **VHLO Serialization**: StableHLO version `1.21.0` introduces
-+    `VHLO_Exp2OpV1` (`exp2_v1`) and `VHLO_Log2OpV1` (`log2_v1`).
-+*   **Forward Compatibility & Decomposition**: For runtimes targeting StableHLO
-+    versions prior to `1.21.0`, compatibility expansion lowers `exp2(x)` to
-+    `exponential(multiply(x, constant(ln(2))))` and `log2(x)` to `divide(log(x),
-+    constant(ln(2)))`.
++* **VHLO Serialization**: StableHLO version `1.21.0` introduces `VHLO_Exp2OpV1`
++  (`exp2_v1`) and `VHLO_Log2OpV1` (`log2_v1`).
++* **Forward Compatibility & Decomposition**: For runtimes targeting StableHLO
++  versions prior to `1.21.0`, compatibility expansion lowers `exp2(x)` to
++  `exponential(multiply(x, constant(ln(2))))` and `log2(x)` to
++  `divide(log(x), constant(ln(2)))`.
 +
 +[jax1]: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.exp2.html
 +[jax2]: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.log2.html
@@ -3698,15 +3700,30 @@ diff --ruN a/stablehlo/stablehlo/reference/InterpreterOps.cpp b/stablehlo/stable
  
  #include "llvm/Support/Casting.h"
  #include "llvm/Support/Errc.h"
-@@ -175,36 +176,39 @@
+@@ -174,37 +175,52 @@
+     ArrayRef<InterpreterValue> inputs, std::queue<StringAttr>& infeed,
      SmallVector<SmallVector<StringAttr>> programs, SymbolTable& symbolTable,
      InterpreterFallback* fallback) {
-   llvm::DefaultThreadPool threadPool;
+-  llvm::DefaultThreadPool threadPool;
 -  SmallVector<std::shared_future<SmallVector<InterpreterValue>>> futures;
-+  llvm::ThreadPoolTaskGroup taskGroup(threadPool);
- 
+-
    uint32_t numReplicas = programs.size();
    uint32_t numPartitions = programs[0].size();
++
++  // StableHLO collectives are blocking barriers: a process inside rendezvous
++  // waits for every other process in its group to arrive, so all processes
++  // must run concurrently. Size the pool from the process count rather than
++  // the default (hardware_concurrency(), i.e. the host thread count). With a
++  // default-sized pool, a program with more processes than the host has
++  // hardware threads leaves the surplus tasks queued and never started while
++  // the started tasks block forever waiting for them - a deadlock that only
++  // surfaces as the rendezvous timeout. hardware_concurrency(N) leaves
++  // ThreadPoolStrategy::Limit false, so the pool gets N threads uncapped by the
++  // host.
++  llvm::DefaultThreadPool threadPool(
++      llvm::hardware_concurrency(numReplicas * numPartitions));
++  llvm::ThreadPoolTaskGroup taskGroup(threadPool);
++
    ProcessGrid processGrid(numReplicas, numPartitions, infeed);
  
 +  SmallVector<SmallVector<InterpreterValue>> taskOutputs(numReplicas *
@@ -4443,6 +4460,159 @@ diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compat
  
    // CHECK-LABEL: @all_reduce_sdy_mesh_max
    func.func @all_reduce_sdy_mesh_max(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir b/stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir
+--- stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir
++++ stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir
+@@ -168,6 +168,18 @@
+ 
+ // -----
+ 
++// CHECK-LABEL: func @dynamic_broadcast_in_dim_discardable_attrs
++func.func @dynamic_broadcast_in_dim_discardable_attrs(%arg0: tensor<4xf32>) -> tensor<3x4xf32> {
++  //  CHECK-NOT: stablehlo.dynamic_broadcast_in_dim
++  //      CHECK: stablehlo.broadcast_in_dim %arg0, dims = [1]
++  // CHECK-SAME:   {mhlo.frontend_attributes = {MUST_FUSE = "0"}}
++  %0 = stablehlo.constant dense<[3, 4]> : tensor<2xi64>
++  %1 = stablehlo.dynamic_broadcast_in_dim %arg0, %0, dims = [1] {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<4xf32>, tensor<2xi64>) -> tensor<3x4xf32>
++  return %1 : tensor<3x4xf32>
++}
++
++// -----
++
+ // CHECK-LABEL: func @dynamic_broadcast_in_dim_inapplicable_dynamic_operand_type
+ func.func @dynamic_broadcast_in_dim_inapplicable_dynamic_operand_type(%arg0: tensor<?xf32>) -> tensor<3x4xf32> {
+   // CHECK: stablehlo.dynamic_broadcast_in_dim
+@@ -225,6 +237,26 @@
+ 
+ // -----
+ 
++// CHECK-LABEL: func @dynamic_conv_discardable_attrs
++func.func @dynamic_conv_discardable_attrs(%arg0: tensor<100x26x26x32xf32>, %arg1: tensor<3x3x1x32xf32>) -> tensor<100x28x28x1xf32> {
++  //  CHECK-NOT: stablehlo.dynamic_conv
++  //      CHECK: stablehlo.convolution(%arg0, %arg1)
++  // CHECK-SAME:   mhlo.frontend_attributes = {MUST_FUSE = "0"}
++  %0 = stablehlo.constant dense<2> : tensor<2x2xi32>
++  %1 = "stablehlo.dynamic_conv"(%arg0, %arg1, %0) {
++    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>,
++    window_strides = array<i64: 1, 1>,
++    lhs_dilation = array<i64: 1, 1>,
++    rhs_dilation = array<i64: 1, 1>,
++    feature_group_count = 1 : i64,
++    batch_group_count = 1 : i64,
++    mhlo.frontend_attributes = {MUST_FUSE = "0"}
++  } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf32>, tensor<2x2xi32>) -> tensor<100x28x28x1xf32>
++  return %1 : tensor<100x28x28x1xf32>
++}
++
++// -----
++
+ // CHECK-LABEL: func @dynamic_conv_success_dynamic_result_type
+ func.func @dynamic_conv_success_dynamic_result_type(%arg0: tensor<100x26x26x32xf32>, %arg1: tensor<3x3x1x32xf32>) -> tensor<?x28x28x1xf32> {
+   //  CHECK-NOT: stablehlo.dynamic_conv
+@@ -295,6 +327,25 @@
+ 
+ // -----
+ 
++// CHECK-LABEL: @dynamic_gather_discardable_attrs
++func.func @dynamic_gather_discardable_attrs(%arg0 : tensor<2x4x9xi32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x8xi32> {
++  //  CHECK-NOT: stablehlo.dynamic_gather
++  //      CHECK: "stablehlo.gather"(%arg0, %arg1)
++  // CHECK-SAME:   {mhlo.frontend_attributes = {MUST_FUSE = "0"}}
++  %0 = stablehlo.constant dense<[1, 1, 8]> : tensor<3xi32>
++  %1 = "stablehlo.dynamic_gather"(%arg0, %arg1, %0) <{
++    dimension_numbers = #stablehlo.gather<
++      collapsed_slice_dims = [0, 1],
++      index_vector_dim = 2,
++      offset_dims = [2],
++      start_index_map = [0, 1]
++    >
++  }> {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<2x4x9xi32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<1x5x8xi32>
++  return %1 : tensor<1x5x8xi32>
++}
++
++// -----
++
+ // CHECK-LABEL: @dynamic_gather_success_dynamic_result_type
+ func.func @dynamic_gather_success_dynamic_result_type(%arg0 : tensor<2x4x9xi32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x?xi32> {
+   //  CHECK-NOT: stablehlo.dynamic_gather
+@@ -348,6 +399,18 @@
+ 
+ // -----
+ 
++// CHECK-LABEL: func @dynamic_iota_discardable_attrs
++func.func @dynamic_iota_discardable_attrs() -> tensor<4xf32> {
++  //  CHECK-NOT: stablehlo.dynamic_iota
++  //      CHECK: stablehlo.iota
++  // CHECK-SAME:   {mhlo.frontend_attributes = {MUST_FUSE = "0"}}
++  %0 = stablehlo.constant dense<4> : tensor<1xi64>
++  %1 = stablehlo.dynamic_iota %0, dim = 0 {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<1xi64>) -> tensor<4xf32>
++  return %1 : tensor<4xf32>
++}
++
++// -----
++
+ // CHECK-LABEL: func @dynamic_iota_inapplicable_dynamic_output_shape
+ func.func @dynamic_iota_inapplicable_dynamic_output_shape(%arg0: tensor<1xi64>) -> tensor<4xf32> {
+   // CHECK: stablehlo.dynamic_iota
+@@ -380,6 +443,20 @@
+ 
+ // -----
+ 
++// CHECK-LABEL: func @dynamic_pad_discardable_attrs
++func.func @dynamic_pad_discardable_attrs(%arg0: tensor<4xf32>, %arg1: tensor<f32>) -> tensor<11xf32> {
++  //  CHECK-NOT: stablehlo.dynamic_pad
++  //      CHECK: stablehlo.pad
++  // CHECK-SAME:   {mhlo.frontend_attributes = {MUST_FUSE = "0"}}
++  %0 = stablehlo.constant dense<0> : tensor<1xi64>
++  %1 = stablehlo.constant dense<1> : tensor<1xi64>
++  %2 = stablehlo.constant dense<2> : tensor<1xi64>
++  %3 = stablehlo.dynamic_pad %arg0, %arg1, %0, %1, %2 {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<4xf32>, tensor<f32>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> tensor<11xf32>
++  return %3 : tensor<11xf32>
++}
++
++// -----
++
+ // CHECK-LABEL: func @dynamic_pad_success_dynamic_result_type
+ func.func @dynamic_pad_success_dynamic_result_type(%arg0: tensor<4xf32>, %arg1: tensor<f32>) -> tensor<?xf32> {
+   // CHECK-NOT: stablehlo.dynamic_pad
+@@ -437,6 +514,17 @@
+ 
+ // -----
+ 
++// CHECK-LABEL: func @dynamic_reshape_discardable_attrs
++func.func @dynamic_reshape_discardable_attrs(%arg0: tensor<4xf32>) -> tensor<1x4xf32> {
++  //  CHECK-NOT: stablehlo.dynamic_reshape
++  //      CHECK: stablehlo.reshape %arg0 {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<4xf32>) -> tensor<1x4xf32>
++  %0 = stablehlo.constant dense<[1, 4]> : tensor<2xi64>
++  %1 = stablehlo.dynamic_reshape %arg0, %0 {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<4xf32>, tensor<2xi64>) -> tensor<1x4xf32>
++  return %1 : tensor<1x4xf32>
++}
++
++// -----
++
+ // CHECK-LABEL: func @dynamic_reshape_inapplicable_dynamic_output_shape
+ func.func @dynamic_reshape_inapplicable_dynamic_output_shape(%arg0: tensor<4xf32>, %arg1: tensor<2xi64>) -> tensor<1x4xf32> {
+   // CHECK: stablehlo.dynamic_reshape
+diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_legalize_deprecated_ops.mlir b/stablehlo/stablehlo/tests/transforms/stablehlo_legalize_deprecated_ops.mlir
+--- stablehlo/stablehlo/tests/transforms/stablehlo_legalize_deprecated_ops.mlir
++++ stablehlo/stablehlo/tests/transforms/stablehlo_legalize_deprecated_ops.mlir
+@@ -10,6 +10,16 @@
+ 
+ // -----
+ 
++// CHECK-LABEL: op_dot_discardable_attrs
++func.func @op_dot_discardable_attrs(%arg0: tensor<2x3xf32>,
++                 %arg1: tensor<3x?xf32>) -> tensor<2x?xf32> {
++  // CHECK: stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<2x3xf32>, tensor<3x?xf32>) -> tensor<2x?xf32>
++  %0 = "stablehlo.dot"(%arg0, %arg1) {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<2x3xf32>, tensor<3x?xf32>) -> tensor<2x?xf32>
++  func.return %0 : tensor<2x?xf32>
++}
++
++// -----
++
+ // CHECK-LABEL: op_unary_einsum
+ func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
+   // CHECK:      %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
 diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir b/stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir
 --- stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir
 +++ stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir
@@ -8393,6 +8563,99 @@ diff --ruN a/stablehlo/stablehlo/transforms/StablehloCanonicalizeDynamism.cpp b/
      return success();
    }
  };
+@@ -119,8 +120,11 @@
+                                          "expected static output_dimensions");
+     if (!op.getType().hasStaticShape())
+       return rewriter.notifyMatchFailure(op, "expected static result type");
+-    rewriter.replaceOpWithNewOp<BroadcastInDimOp>(
+-        op, op.getType(), op.getOperand(), op.getBroadcastDimensionsAttr());
++    auto broadcastInDimOp = BroadcastInDimOp::create(
++        rewriter, op.getLoc(), op.getType(), op.getOperand(),
++        op.getBroadcastDimensionsAttr());
++    broadcastInDimOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
++    rewriter.replaceOp(op, broadcastInDimOp);
+     return success();
+   }
+ };
+@@ -139,12 +143,14 @@
+         RankedTensorType::get({static_cast<int64_t>(padding.size()) / 2, 2},
+                               rewriter.getI64Type()),
+         padding);
+-    rewriter.replaceOpWithNewOp<ConvolutionOp>(
+-        op, op.getType(), op.getLhs(), op.getRhs(), op.getWindowStridesAttr(),
+-        paddingAttr, op.getLhsDilationAttr(), op.getRhsDilationAttr(),
+-        op.getWindowReversalAttr(), op.getDimensionNumbers(),
+-        op.getFeatureGroupCount(), op.getBatchGroupCount(),
+-        op.getPrecisionConfigAttr());
++    auto convolutionOp = ConvolutionOp::create(
++        rewriter, op.getLoc(), op.getType(), op.getLhs(), op.getRhs(),
++        op.getWindowStridesAttr(), paddingAttr, op.getLhsDilationAttr(),
++        op.getRhsDilationAttr(), op.getWindowReversalAttr(),
++        op.getDimensionNumbers(), op.getFeatureGroupCount(),
++        op.getBatchGroupCount(), op.getPrecisionConfigAttr());
++    convolutionOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
++    rewriter.replaceOp(op, convolutionOp);
+     return success();
+   }
+ };
+@@ -159,10 +165,13 @@
+     SmallVector<int64_t> sliceSizes;
+     if (!succeeded(hlo::matchInts(op.getSliceSizes(), sliceSizes)))
+       return rewriter.notifyMatchFailure(op, "expected static slice_sizes");
+-    rewriter.replaceOpWithNewOp<GatherOp>(
+-        op, op.getType(), op.getOperand(), op.getStartIndices(),
+-        op.getDimensionNumbersAttr(), rewriter.getDenseI64ArrayAttr(sliceSizes),
+-        op.getIndicesAreSortedAttr());
++    auto gatherOp =
++        GatherOp::create(rewriter, op.getLoc(), op.getType(), op.getOperand(),
++                         op.getStartIndices(), op.getDimensionNumbersAttr(),
++                         rewriter.getDenseI64ArrayAttr(sliceSizes),
++                         op.getIndicesAreSortedAttr());
++    gatherOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
++    rewriter.replaceOp(op, gatherOp);
+     return success();
+   }
+ };
+@@ -179,8 +188,10 @@
+       return rewriter.notifyMatchFailure(op, "expected static output_shape");
+     if (!op.getType().hasStaticShape())
+       return rewriter.notifyMatchFailure(op, "expected static result type");
+-    rewriter.replaceOpWithNewOp<IotaOp>(op, op.getType(),
+-                                        op.getIotaDimension());
++    auto iotaOp = IotaOp::create(rewriter, op.getLoc(), op.getType(),
++                                 op.getIotaDimension());
++    iotaOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
++    rewriter.replaceOp(op, iotaOp);
+     return success();
+   }
+ };
+@@ -198,9 +209,11 @@
+       return rewriter.notifyMatchFailure(op, "expected static high");
+     if (!succeeded(hlo::matchInts(op.getInteriorPadding(), interiorPadding)))
+       return rewriter.notifyMatchFailure(op, "expected static interior");
+-    rewriter.replaceOpWithNewOp<PadOp>(op, op.getType(), op.getOperand(),
+-                                       op.getPaddingValue(), edgePaddingLow,
+-                                       edgePaddingHigh, interiorPadding);
++    auto padOp = PadOp::create(
++        rewriter, op.getLoc(), op.getType(), op.getOperand(),
++        op.getPaddingValue(), edgePaddingLow, edgePaddingHigh, interiorPadding);
++    padOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
++    rewriter.replaceOp(op, padOp);
+     return success();
+   }
+ };
+@@ -216,7 +229,10 @@
+       return rewriter.notifyMatchFailure(op, "expected static output_shape");
+     if (!op.getType().hasStaticShape())
+       return rewriter.notifyMatchFailure(op, "expected static result type");
+-    rewriter.replaceOpWithNewOp<ReshapeOp>(op, op.getType(), op.getOperand());
++    auto reshapeOp =
++        ReshapeOp::create(rewriter, op.getLoc(), op.getType(), op.getOperand());
++    reshapeOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
++    rewriter.replaceOp(op, reshapeOp);
+     return success();
+   }
+ };
 diff --ruN a/stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp b/stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp
 --- stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp
 +++ stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp
@@ -8433,6 +8696,42 @@ diff --ruN a/stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp b/s
    populateWithGenerated(*patterns);
  }
  
+diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOps.cpp b/stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOps.cpp
+--- stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOps.cpp
++++ stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOps.cpp
+@@ -132,6 +132,14 @@
+   llvm::report_fatal_error("unsupported type");
+ }
+ 
++void copyDiscardableAttrs(Operation* src, Operation* dst) {
++  dst->setDiscardableAttrs(src->getDiscardableAttrDictionary());
++}
++
++void copyDiscardableAttrs(Value src, Value dst) {
++  copyDiscardableAttrs(src.getDefiningOp(), dst.getDefiningOp());
++}
++
+ #include "stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.h.inc"
+ }  // namespace
+ 
+diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.td b/stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.td
+--- stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.td
++++ stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.td
+@@ -45,9 +45,12 @@
+ 
+ // def CrossReplicaSumToAllReduce : <Declared in cpp>
+ 
++def CopyDiscardableAttrs : NativeCodeCallVoid<"copyDiscardableAttrs($0, $1)">;
++
+ def DotToDotGeneral : Pat<
+-  (StableHLO_DotOp $lhs, $rhs, $precision_config),
+-  (StableHLO_DotGeneralOp $lhs, $rhs, (GetDefaultDotDimensionNumbers $lhs), $precision_config, (NativeCodeCall<"Attribute{}">))>;
++  (StableHLO_DotOp:$src $lhs, $rhs, $precision_config),
++  (StableHLO_DotGeneralOp:$dst $lhs, $rhs, (GetDefaultDotDimensionNumbers $lhs), $precision_config, (NativeCodeCall<"Attribute{}">)),
++  [], [(CopyDiscardableAttrs $src, $dst)]>;
+ 
+ def UnaryEinsumToEinsum : Pat<
+   (StableHLO_UnaryEinsumOp $operand, $equation),
 diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp b/stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp
 --- stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp
 +++ stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp
diff --git a/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder_test.cc b/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder_test.cc
index a29eb1593b7..e36aa7e46a9 100644
--- a/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder_test.cc
+++ b/third_party/xla/xla/backends/cpu/codegen/kernel_api_ir_builder_test.cc
@@ -218,7 +218,7 @@ TEST_F(KernelApiIrBuilderTest, BuildKernelPrototype) {
     CHECK-DAG: ![[SCOPE3]] = !{![[RES1]]}
     CHECK-DAG: ![[RES0]] = !{!"{{.*}}, offset:512, {{.*}}", ![[DOMAIN:.+]]}
     CHECK-DAG: ![[RES1]] = !{!"{{.*}}, offset:768, {{.*}}", ![[DOMAIN]]}
-    CHECK-DAG: ![[DOMAIN]] = !{!"XLA host kernel test AA domain"}
+    CHECK-DAG: ![[DOMAIN]] = !{!"XLA host kernel test AA domain", i1 false}
   )")));
   // clang-format on
 
@@ -432,7 +432,7 @@ TEST_F(KernelApiIrBuilderTestNoBufferValidation, PartialOverlap) {
     #0 = { uwtable "frame-pointer"="all" "prefer-vector-width"="256" }
     CHECK-DAG: ![[ALIGNMENT]] = !{i64 )", MinAlign(), R"(}
     CHECK-DAG: ![[EMPTY_NODE]] = !{}
-    CHECK-DAG: ![[DOMAIN:.+]] = !{!"XLA host kernel test AA domain"}
+    CHECK-DAG: ![[DOMAIN:.+]] = !{!"XLA host kernel test AA domain", i1 false}
     CHECK-DAG: ![[RES0:.+]] = !{!"{{.*}}, offset:288, size:512}", ![[DOMAIN]]}
     CHECK-DAG: ![[RES1:.+]] = !{!"{{.*}}, offset:768, size:256}", ![[DOMAIN]]}
     CHECK-DAG: ![[RES2:.+]] = !{!"{{.*}}, offset:1024, size:256}", ![[DOMAIN]]}
diff --git a/third_party/xla/xla/codegen/emitters/transforms/lower_to_llvm_gpu.cc b/third_party/xla/xla/codegen/emitters/transforms/lower_to_llvm_gpu.cc
index 2263d3e492b..b49f1a0a68c 100644
--- a/third_party/xla/xla/codegen/emitters/transforms/lower_to_llvm_gpu.cc
+++ b/third_party/xla/xla/codegen/emitters/transforms/lower_to_llvm_gpu.cc
@@ -251,19 +251,18 @@ class LowerToLLVMGPUPass
       if (device_spec_.IsAmdGpu()) {
         std::string chipset =
             device_spec_.gpu().rocm_compute_capability().gfx_version();
-        llvm::FailureOr<mlir::amdgpu::Chipset> maybeChipset =
-            mlir::amdgpu::Chipset::parse(chipset);
-        if (mlir::failed(maybeChipset)) {
+        llvm::FailureOr<mlir::ROCDL::TargetInfo> targetInfo =
+            mlir::ROCDL::TargetInfo::get(chipset);
+        if (mlir::failed(targetInfo)) {
           mlir::emitError(mlir::UnknownLoc::get(&getContext()),
-                          "Invalid chipset name: " + chipset);
+                          "Invalid chipset/target name: " + chipset);
           return mlir::failure();
         }
         mlir::populateGpuToROCDLConversionPatterns(
-            converter, patterns, mlir::gpu::amd::Runtime::Unknown,
-            *maybeChipset);
+            converter, patterns, mlir::gpu::amd::Runtime::Unknown, *targetInfo);
         mlir::configureGpuToROCDLConversionLegality(target);
         mlir::populateAMDGPUToROCDLConversionPatterns(converter, patterns,
-                                                      *maybeChipset);
+                                                      *targetInfo);
         // Higher benefit than the default MathToROCDL patterns so these win for
         // scalar bf16 ops.
         mlir::PatternBenefit benefit(2);
diff --git a/third_party/xla/xla/mlir_hlo/transforms/gpu_kernel_lowering_passes.cc b/third_party/xla/xla/mlir_hlo/transforms/gpu_kernel_lowering_passes.cc
index f1d8ccffb27..02abbabde7b 100644
--- a/third_party/xla/xla/mlir_hlo/transforms/gpu_kernel_lowering_passes.cc
+++ b/third_party/xla/xla/mlir_hlo/transforms/gpu_kernel_lowering_passes.cc
@@ -136,19 +136,19 @@ void GpuKernelToNVVMPass::runOnOperation() {
 }
 
 void GpuKernelToROCDLPass::runOnOperation() {
-  llvm::FailureOr<mlir::amdgpu::Chipset> maybeChipset =
-      mlir::amdgpu::Chipset::parse(chipset);
-  if (failed(maybeChipset)) {
+  llvm::FailureOr<mlir::ROCDL::TargetInfo> targetInfo =
+      mlir::ROCDL::TargetInfo::get(chipset);
+  if (failed(targetInfo)) {
     mlir::emitError(mlir::UnknownLoc::get(&getContext()),
-                    "Invalid chipset name: " + chipset);
+                    "Invalid chipset/target name: " + chipset);
     return signalPassFailure();
   }
 
   RewritePatternSet patterns(&getContext());
   LLVMTypeConverter converter(&getContext());
   populateCommonPatterns(converter, patterns);
-  populateGpuToROCDLConversionPatterns(
-      converter, patterns, gpu::amd::Runtime::Unknown, *maybeChipset);
+  populateGpuToROCDLConversionPatterns(converter, patterns,
+                                       gpu::amd::Runtime::Unknown, *targetInfo);
   ConversionTarget target(getContext());
   configureGpuToROCDLConversionLegality(target);
   if (failed(

From 0c8127b6ec4c928edf2ca43b57b6b242a307eec6 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 8 Oct 2026 16:15:49 -0700
Subject: [PATCH 0319/1000] Eliminate PJRT C API phase compile copy and
 overlapping buffers

Optimize memory usage and buffer conversion speed during PJRT C API phase compilation by eliminating duplicate in-memory buffers:

- Reclaim `PjRtPartialProgramProto` memory immediately after serializing each proto into C buffers in `ConvertPjRtPartialProgramProtosToCharBuffers`.
- Free input C buffers before deserializing output buffers in `PjRtCApiPhaseCompiler::RunPhases` to prevent holding both simultaneously.

PiperOrigin-RevId: 996134916
---
 .../c/pjrt_c_api_phase_compile_internal.cc    |  7 ++--
 third_party/xla/xla/pjrt/c_api_client/BUILD   |  1 +
 .../c_api_client/pjrt_c_api_phase_compiler.cc | 33 +++++++++++--------
 .../xla/xla/pjrt/partial_program_utils.cc     | 10 +++---
 .../xla/xla/pjrt/partial_program_utils.h      |  9 ++---
 5 files changed, 35 insertions(+), 25 deletions(-)

diff --git a/third_party/xla/xla/pjrt/c/pjrt_c_api_phase_compile_internal.cc b/third_party/xla/xla/pjrt/c/pjrt_c_api_phase_compile_internal.cc
index 10202d67419..2cf605c2058 100644
--- a/third_party/xla/xla/pjrt/c/pjrt_c_api_phase_compile_internal.cc
+++ b/third_party/xla/xla/pjrt/c/pjrt_c_api_phase_compile_internal.cc
@@ -84,10 +84,11 @@ static PJRT_Error* PJRT_PhaseCompile_Run_Phase(
                             options, std::move(programs_in_protos),
                             *args->topology->topology, phases_to_run));
 
-  PJRT_ASSIGN_OR_RETURN(args->output_programs,
-                        xla::ConvertPjRtPartialProgramProtosToCharBuffers(
-                            programs_out, args->output_programs_sizes));
   args->num_output_programs = programs_out.size();
+  PJRT_ASSIGN_OR_RETURN(
+      args->output_programs,
+      xla::ConvertPjRtPartialProgramProtosToCharBuffers(
+          absl::MakeSpan(programs_out), args->output_programs_sizes));
 
   return nullptr;
 }
diff --git a/third_party/xla/xla/pjrt/c_api_client/BUILD b/third_party/xla/xla/pjrt/c_api_client/BUILD
index 58fe0285867..b335ba45e42 100644
--- a/third_party/xla/xla/pjrt/c_api_client/BUILD
+++ b/third_party/xla/xla/pjrt/c_api_client/BUILD
@@ -118,6 +118,7 @@ cc_library(
         "//xla/tsl/platform:statusor",
         "//xla/tsl/protobuf:coordination_service_proto_cc",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/base",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/container:flat_hash_map",
diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_phase_compiler.cc b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_phase_compiler.cc
index ebd0d8d407b..cc0de8ec2d6 100644
--- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_phase_compiler.cc
+++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_phase_compiler.cc
@@ -15,11 +15,13 @@ limitations under the License.
 
 #include "xla/pjrt/c_api_client/pjrt_c_api_phase_compiler.h"
 
-#include <algorithm>
 #include <cstddef>
 #include <string>
+#include <utility>
 #include <vector>
 
+#include "absl/algorithm/container.h"
+#include "absl/base/casts.h"
 #include "absl/cleanup/cleanup.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
@@ -27,7 +29,6 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/types/span.h"
 #include "xla/pjrt/c/pjrt_c_api.h"
-#include "xla/pjrt/c/pjrt_c_api_helpers.h"
 #include "xla/pjrt/c/pjrt_c_api_phase_compile_extension.h"
 #include "xla/pjrt/c/pjrt_c_api_status_utils.h"
 #include "xla/pjrt/c_api_client/pjrt_c_api_client.h"
@@ -37,9 +38,6 @@ limitations under the License.
 #include "xla/pjrt/proto/compile_options.pb.h"
 #include "xla/pjrt/proto/pjrt_partial_program.pb.h"
 #include "xla/pjrt/string_utils.h"
-#include "xla/tsl/platform/errors.h"
-#include "xla/tsl/platform/statusor.h"
-#include "tsl/platform/casts.h"
 
 namespace xla {
 
@@ -55,9 +53,8 @@ absl::Status ValidatePhases(
   }
 
   for (const auto& partial_program : partial_programs_in) {
-    auto consumer_phases = partial_program.consumer_phases();
-    if (std::find(consumer_phases.begin(), consumer_phases.end(),
-                  phases_to_run[0]) == consumer_phases.end()) {
+    if (!absl::c_contains(partial_program.consumer_phases(),
+                          phases_to_run[0])) {
       return absl::InvalidArgumentError(absl::StrCat(
           "Input partial programs cannot be compiled by a phase with name \"",
           phases_to_run[0], "\""));
@@ -149,9 +146,10 @@ PjRtCApiPhaseCompiler::RunPhases(
           ->c_topology();
 
   const size_t* programs_in_buffer_sizes;
-  ABSL_ASSIGN_OR_RETURN(const char** programs_in_buffers,
-                   xla::ConvertPjRtPartialProgramProtosToCharBuffers(
-                       partial_programs_in, programs_in_buffer_sizes));
+  ABSL_ASSIGN_OR_RETURN(
+      const char** programs_in_buffers,
+      xla::ConvertPjRtPartialProgramProtosToCharBuffers(
+          absl::MakeSpan(partial_programs_in), programs_in_buffer_sizes));
   size_t num_programs_in = partial_programs_in.size();
 
   // We no longer need this form of the input, so free it to save memory.
@@ -198,15 +196,22 @@ PjRtCApiPhaseCompiler::RunPhases(
   RETURN_STATUS_IF_PJRT_ERROR(
       phase_compile_extension_->phase_compile_run_phases(&run_args), api_);
 
+  // Free caller-allocated input buffers before deserializing the output buffers
+  // to avoid holding both input and output C buffers simultaneously in memory.
+  std::move(cleanup_programs_in_buffers).Invoke();
+  std::move(cleanup_phases_to_run_buffers).Invoke();
+
+  absl::Cleanup cleanup_output_buffers = [this, &run_args] {
+    CleanUpPluginDefinedCBuffers(
+        run_args.output_programs, run_args.output_programs_sizes,
+        run_args.num_output_programs, phase_compile_extension_);
+  };
   ABSL_ASSIGN_OR_RETURN(std::vector<xla::PjRtPartialProgramProto> output_programs,
                    xla::ConvertCharBuffersToPjRtPartialProgramProtos(
                        absl::MakeSpan(run_args.output_programs,
                                       run_args.num_output_programs),
                        absl::MakeConstSpan(run_args.output_programs_sizes,
                                            run_args.num_output_programs)));
-  CleanUpPluginDefinedCBuffers(
-      run_args.output_programs, run_args.output_programs_sizes,
-      run_args.num_output_programs, phase_compile_extension_);
 
   return output_programs;
 }
diff --git a/third_party/xla/xla/pjrt/partial_program_utils.cc b/third_party/xla/xla/pjrt/partial_program_utils.cc
index b595f62aef7..0fdd7a96ca7 100644
--- a/third_party/xla/xla/pjrt/partial_program_utils.cc
+++ b/third_party/xla/xla/pjrt/partial_program_utils.cc
@@ -25,7 +25,6 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
-#include "xla/pjrt/proto/compile_options.pb.h"
 #include "xla/pjrt/proto/pjrt_partial_program.pb.h"
 
 namespace xla {
@@ -52,14 +51,14 @@ ConvertCharBuffersToPjRtPartialProgramProtos(
       return absl::InvalidArgumentError(
           "Failed to deserialize PjRtPartialProgramProto");
     }
-    partial_programs.push_back(partial_program);
+    partial_programs.push_back(std::move(partial_program));
   }
 
   return partial_programs;
 }
 
 absl::StatusOr<const char**> ConvertPjRtPartialProgramProtosToCharBuffers(
-    const std::vector<xla::PjRtPartialProgramProto>& partial_programs,
+    absl::Span<xla::PjRtPartialProgramProto> partial_programs,
     const size_t*& char_buffer_sizes) {
   const size_t num_programs = partial_programs.size();
   const char** char_buffers = new const char*[num_programs]();
@@ -73,7 +72,7 @@ absl::StatusOr<const char**> ConvertPjRtPartialProgramProtosToCharBuffers(
   };
 
   for (size_t i = 0; i < num_programs; ++i) {
-    const xla::PjRtPartialProgramProto& partial_program = partial_programs[i];
+    xla::PjRtPartialProgramProto& partial_program = partial_programs[i];
     size_t buffer_size = partial_program.ByteSizeLong();
     char* buffer = new char[buffer_size];
     char_buffers[i] = buffer;
@@ -83,6 +82,9 @@ absl::StatusOr<const char**> ConvertPjRtPartialProgramProtosToCharBuffers(
       return absl::InvalidArgumentError(
           "Failed to serialize PjRtPartialProgramProto");
     }
+    // Reclaim the proto's memory immediately after serializing it into the C
+    // buffer to avoid holding both representations simultaneously.
+    partial_program = xla::PjRtPartialProgramProto();
   }
   char_buffer_sizes = buffer_sizes;
 
diff --git a/third_party/xla/xla/pjrt/partial_program_utils.h b/third_party/xla/xla/pjrt/partial_program_utils.h
index 5b4b8622caf..d55b69905bc 100644
--- a/third_party/xla/xla/pjrt/partial_program_utils.h
+++ b/third_party/xla/xla/pjrt/partial_program_utils.h
@@ -34,13 +34,14 @@ ConvertCharBuffersToPjRtPartialProgramProtos(
     absl::Span<const char*> char_buffers,
     absl::Span<const size_t> char_buffer_sizes);
 
-// Converts a vector of `PjRtPartialProgramProto`s to an array of C-style
+// Converts a span of `PjRtPartialProgramProto`s to an array of C-style
 // character buffers and their sizes. This function allocates memory for the
 // `char_buffers`, the individual serialized protos, and `char_buffer_sizes`.
-// In case the partial programs cannot be serialized, this function frees the
-// allocated memory.
+// Each `PjRtPartialProgramProto` is set to an empty message after serialization
+// to minimize peak memory usage. In case the partial programs cannot be
+// serialized, this function frees the allocated memory.
 absl::StatusOr<const char**> ConvertPjRtPartialProgramProtosToCharBuffers(
-    const std::vector<xla::PjRtPartialProgramProto>& partial_programs,
+    absl::Span<xla::PjRtPartialProgramProto> partial_programs,
     const size_t*& char_buffer_sizes);
 
 }  // namespace xla

From 883f7a09bab310ae90a74eaee83707321741203e Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 8 Oct 2026 18:14:25 -0700
Subject: [PATCH 0320/1000] When propagating frontend attrs such as buffering
 level to inner calls, inner calls cannot be updated in-place as the callee
 can have multiple call sites.

PiperOrigin-RevId: 996192283
---
 third_party/xla/xla/side_effect_util.cc | 10 ++++++++++
 third_party/xla/xla/side_effect_util.h  |  5 +++++
 2 files changed, 15 insertions(+)

diff --git a/third_party/xla/xla/side_effect_util.cc b/third_party/xla/xla/side_effect_util.cc
index 73888cb70c1..b17df58c2cb 100644
--- a/third_party/xla/xla/side_effect_util.cc
+++ b/third_party/xla/xla/side_effect_util.cc
@@ -94,6 +94,16 @@ const char kFuseLimitAttr[] = "FUSE_LIMIT";
 
 const char kMaximalFuseAttr[] = "MAXIMAL_FUSE";
 
+const char kBufferingLevelAttr[] = "buffering_level";
+
+const char kOutputWindowBoundsAttr[] = "output_window_bounds";
+
+const char kKernelWindowBoundsAttr[] = "kernel_window_bounds";
+
+const char kScopedVmemLimitKibAttr[] = "scoped_vmem_limit_kib";
+
+const char kIntegerAttr[] = "integer";
+
 const char kXlaCseSafeZeroOperandAttr[] = "_xla_cse_safe_zero_operand";
 
 const char kCollectiveGroupKeyAttr[] = "collective_group_key";
diff --git a/third_party/xla/xla/side_effect_util.h b/third_party/xla/xla/side_effect_util.h
index b31d385c106..8dcd2227cab 100644
--- a/third_party/xla/xla/side_effect_util.h
+++ b/third_party/xla/xla/side_effect_util.h
@@ -112,6 +112,11 @@ extern const char kXlaNoOpSchedulingGroup[];
 extern const char kMustFuseAttr[];
 extern const char kMaximalFuseAttr[];
 extern const char kFuseLimitAttr[];
+extern const char kBufferingLevelAttr[];
+extern const char kOutputWindowBoundsAttr[];
+extern const char kKernelWindowBoundsAttr[];
+extern const char kScopedVmemLimitKibAttr[];
+extern const char kIntegerAttr[];
 extern const char kXlaCseSafeZeroOperandAttr[];
 
 // Frontend attribute asking XLA to launch and schedule independent collectives

From 2bb215096e197f29b8e375c13f3eae74939ea54a Mon Sep 17 00:00:00 2001
From: Ashutosh0x <ashutoshkumarsingh0x@gmail.com>
Date: Fri, 9 Oct 2026 07:18:01 +0530
Subject: [PATCH 0321/1000] fix: validate input rank before reading dims in
 LSTMBlockCell kernels

LSTMBlockCellOp read x_tensor->dim_size(1) and cs_prev_tensor->dim_size(1)
before its rank checks ran, so a rank 1 input aborted the process with
"Check failed: d < dims()" instead of returning InvalidArgument. Validate
the input ranks with a check_rank lambda, as LSTMBlockCellGradOp now does,
before any dimension is read, and drop the later duplicate checks.

LSTMBlockCellGradOp did not reject an empty cs_prev, which gives
LSTMBlockCellBpropWithCUDA a zero-sized launch configuration that aborts
via TF_CHECK_OK. Reject it as the forward op does.

The tests feed the offending input through an unknown-shape signature so
that the kernel check, not the shape function, is what they exercise.

Fixes #113069
Fixes #113071
---
 tensorflow/core/kernels/rnn/lstm_ops.cc | 75 +++++++++++------------
 tensorflow/python/ops/BUILD             |  2 +
 tensorflow/python/ops/rnn_grad_test.py  | 81 +++++++++++++++++++++++++
 3 files changed, 120 insertions(+), 38 deletions(-)

diff --git a/tensorflow/core/kernels/rnn/lstm_ops.cc b/tensorflow/core/kernels/rnn/lstm_ops.cc
index fae4e36733e..fe35e65428e 100644
--- a/tensorflow/core/kernels/rnn/lstm_ops.cc
+++ b/tensorflow/core/kernels/rnn/lstm_ops.cc
@@ -335,6 +335,33 @@ class LSTMBlockCellOp : public OpKernel {
     const Tensor* b_tensor = nullptr;
     OP_REQUIRES_OK(ctx, ctx->input("b", &b_tensor));
 
+    // The dimensions read below index dimension 1 of x and cs_prev, and the
+    // peephole weights reach the functor as vec<T>() whether or not
+    // use_peephole is set. Both are fatal checks rather than errors on a
+    // tensor of the wrong rank, so validate the ranks first.
+    const auto check_rank = [](const Tensor* t, const char* name,
+                               int expected) -> absl::Status {
+      if (t->dims() != expected) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            name, " must be rank ", expected, " but is rank ", t->dims()));
+      }
+      return absl::OkStatus();
+    };
+    OP_REQUIRES_OK(ctx, check_rank(x_tensor, "x", 2));
+    OP_REQUIRES_OK(ctx, check_rank(cs_prev_tensor, "cs_prev", 2));
+    OP_REQUIRES_OK(ctx, check_rank(h_prev_tensor, "h_prev", 2));
+    OP_REQUIRES_OK(ctx, check_rank(w_tensor, "w", 2));
+    OP_REQUIRES_OK(ctx, check_rank(wci_tensor, "wci", 1));
+    OP_REQUIRES_OK(ctx, check_rank(wcf_tensor, "wcf", 1));
+    OP_REQUIRES_OK(ctx, check_rank(wco_tensor, "wco", 1));
+    OP_REQUIRES_OK(ctx, check_rank(b_tensor, "b", 1));
+    OP_REQUIRES(
+        ctx, cs_prev_tensor->dim_size(0) > 0 && cs_prev_tensor->dim_size(1) > 0,
+        absl::InvalidArgumentError(
+            absl::StrCat("cs_prev_tensor is empty, has shape: (",
+                         cs_prev_tensor->dim_size(0), ",",
+                         cs_prev_tensor->dim_size(1), ").")));
+
     const int64_t batch_size = x_tensor->dim_size(0);
     const int64_t input_size = x_tensor->dim_size(1);
     const int64_t cell_size = cs_prev_tensor->dim_size(1);
@@ -424,44 +451,6 @@ class LSTMBlockCellOp : public OpKernel {
     const Device& device = ctx->eigen_device<Device>();
 
     // Sanity check that each of the tensors have the required NDIMS.
-    OP_REQUIRES(
-        ctx, x_tensor->dims() == 2,
-        absl::InvalidArgumentError(absl::StrCat(
-            "x_tensor must be rank 2 but is rank ", x_tensor->dims(), ".")));
-    OP_REQUIRES(ctx, cs_prev_tensor->dims() == 2,
-                absl::InvalidArgumentError(
-                    absl::StrCat("cs_prev_tensor must be rank 2 but is rank ",
-                                 cs_prev_tensor->dims(), ".")));
-    OP_REQUIRES(
-        ctx, cs_prev_tensor->dim_size(0) > 0 && cs_prev_tensor->dim_size(1) > 0,
-        absl::InvalidArgumentError(
-            absl::StrCat("cs_prev_tensor is empty, has shape: (",
-                         cs_prev_tensor->dim_size(0), ",",
-                         cs_prev_tensor->dim_size(1), ").")));
-    OP_REQUIRES(ctx, h_prev_tensor->dims() == 2,
-                absl::InvalidArgumentError(
-                    absl::StrCat("h_prev_tensor must be rank 2 but is rank ",
-                                 h_prev_tensor->dims(), ".")));
-    OP_REQUIRES(
-        ctx, w_tensor->dims() == 2,
-        absl::InvalidArgumentError(absl::StrCat(
-            "w_tensor must be rank 2 but is rank ", w_tensor->dims(), ".")));
-    OP_REQUIRES(ctx, wci_tensor->dims() == 1,
-                absl::InvalidArgumentError(
-                    absl::StrCat("wci_tensor must be rank 1 but is rank ",
-                                 wci_tensor->dims(), ".")));
-    OP_REQUIRES(ctx, wcf_tensor->dims() == 1,
-                absl::InvalidArgumentError(
-                    absl::StrCat("wcf_tensor must be rank 1 but is rank ",
-                                 wcf_tensor->dims(), ".")));
-    OP_REQUIRES(ctx, wco_tensor->dims() == 1,
-                absl::InvalidArgumentError(
-                    absl::StrCat("wco_tensor must be rank 1 but is rank ",
-                                 wco_tensor->dims(), ".")));
-    OP_REQUIRES(
-        ctx, b_tensor->dims() == 1,
-        absl::InvalidArgumentError(absl::StrCat(
-            "b_tensor must be rank 1 but is rank ", b_tensor->dims(), ".")));
     OP_REQUIRES(
         ctx, xh_tensor.dims() == 2,
         absl::InvalidArgumentError(absl::StrCat(
@@ -622,6 +611,16 @@ class LSTMBlockCellGradOp : public OpKernel {
     OP_REQUIRES_OK(ctx, check_rank(cs_grad_tensor, "cs_grad", 2));
     OP_REQUIRES_OK(ctx, check_rank(h_grad_tensor, "h_grad", 2));
 
+    // An empty cs_prev yields a zero-sized GPU launch configuration in
+    // LSTMBlockCellBpropWithCUDA, which aborts via TF_CHECK_OK. The forward
+    // op rejects this for the same reason; see #58270.
+    OP_REQUIRES(
+        ctx, cs_prev_tensor->dim_size(0) > 0 && cs_prev_tensor->dim_size(1) > 0,
+        absl::InvalidArgumentError(
+            absl::StrCat("cs_prev_tensor is empty, has shape: (",
+                         cs_prev_tensor->dim_size(0), ",",
+                         cs_prev_tensor->dim_size(1), ").")));
+
     const int64_t batch_size = x_tensor->dim_size(0);
     const int64_t input_size = x_tensor->dim_size(1);
     const int64_t cell_size = cs_prev_tensor->dim_size(1);
diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD
index 675b0563574..927ce57730e 100644
--- a/tensorflow/python/ops/BUILD
+++ b/tensorflow/python/ops/BUILD
@@ -2822,9 +2822,11 @@ cuda_py_strict_test(
         ":math_ops",
         ":rnn_grad",
         ":rnn_ops_gen",
+        "//tensorflow/python/eager:def_function",
         "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/framework:errors",
         "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_spec",
         "//tensorflow/python/framework:test_lib",
         "//tensorflow/python/platform:client_testlib",
         "//third_party/py/numpy",
diff --git a/tensorflow/python/ops/rnn_grad_test.py b/tensorflow/python/ops/rnn_grad_test.py
index 5f7f068aaaa..b928b8aefca 100644
--- a/tensorflow/python/ops/rnn_grad_test.py
+++ b/tensorflow/python/ops/rnn_grad_test.py
@@ -18,9 +18,11 @@ import functools
 
 import numpy as np
 
+from tensorflow.python.eager import def_function
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors_impl
 from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_spec
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import gen_rnn_ops
@@ -202,6 +204,85 @@ class RNNGradTest(test.TestCase):
         use_peephole=False,
     )
 
+  def _run_with_unknown_shape(self, op, name, **kwargs):
+    """Runs `op` with `kwargs[name]` fed through an unknown-shape signature.
+
+    Hiding the shape keeps the op's shape function from rejecting the input
+    while the graph is built, so the check under test is the kernel's.
+    """
+    value = kwargs.pop(name)
+
+    @def_function.function(
+        autograph=False,
+        input_signature=[tensor_spec.TensorSpec(shape=None, dtype=value.dtype)])
+    def run(arg):
+      return op(**{name: arg}, **kwargs)
+
+    return self.evaluate(run(value))
+
+  def _lstm_block_cell_inputs(self, batch_size=2, input_size=2, cell_size=2):
+    cs_prev = deterministic_random_uniform([batch_size, cell_size])
+    return dict(
+        x=deterministic_random_uniform([batch_size, input_size]),
+        cs_prev=cs_prev,
+        h_prev=cs_prev,
+        w=deterministic_random_uniform(
+            [input_size + cell_size, 4 * cell_size]),
+        wci=deterministic_random_uniform([cell_size]),
+        wcf=deterministic_random_uniform([cell_size]),
+        wco=deterministic_random_uniform([cell_size]),
+        b=deterministic_random_uniform([4 * cell_size]),
+    )
+
+  def _lstm_block_cell_grad_inputs(self, batch_size=2, input_size=2,
+                                   cell_size=2):
+    kwargs = self._lstm_block_cell_inputs(batch_size, input_size, cell_size)
+    for name in ("i", "cs", "f", "o", "ci", "co", "cs_grad", "h_grad"):
+      kwargs[name] = deterministic_random_uniform([batch_size, cell_size])
+    return kwargs
+
+  def testLSTMBlockCellInvalidRank(self):
+    # Test case for GitHub issue 113069. A rank 1 x was indexed at dimension 1
+    # before the kernel's rank checks ran, which aborted the process.
+    kwargs = self._lstm_block_cell_inputs()
+    kwargs["x"] = deterministic_random_uniform([2])
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                "x must be rank 2"):
+      self._run_with_unknown_shape(
+          gen_rnn_ops.lstm_block_cell, "x", use_peephole=False, **kwargs)
+
+  def testLSTMBlockCellInvalidPeepholeRank(self):
+    # The peephole weights reach the functor as vec<T>(), a fatal check unless
+    # the tensor is rank 1.
+    kwargs = self._lstm_block_cell_inputs()
+    kwargs["wci"] = deterministic_random_uniform([1, 2])
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                "wci must be rank 1"):
+      self._run_with_unknown_shape(
+          gen_rnn_ops.lstm_block_cell, "wci", use_peephole=True, **kwargs)
+
+  def testLSTMBlockCellGradInvalidPeepholeRank(self):
+    # As in the forward op, vec<T>() is called on the peephole weights even
+    # when use_peephole is false.
+    kwargs = self._lstm_block_cell_grad_inputs()
+    kwargs["wci"] = deterministic_random_uniform([1, 2])
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                "wci must be rank 1"):
+      self._run_with_unknown_shape(
+          gen_rnn_ops.lstm_block_cell_grad, "wci", use_peephole=False,
+          **kwargs)
+
+  def testLSTMBlockCellGradEmptyCsPrev(self):
+    # An empty cs_prev gives LSTMBlockCellBpropWithCUDA a zero-sized launch
+    # configuration, which aborts the process. The forward op already rejects
+    # it; see #58270.
+    kwargs = self._lstm_block_cell_grad_inputs(cell_size=0)
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                "cs_prev_tensor is empty"):
+      self._run_with_unknown_shape(
+          gen_rnn_ops.lstm_block_cell_grad, "cs_prev", use_peephole=False,
+          **kwargs)
+
   def _lstm_block(self, op, w, b, x, cs_prev, h_prev):
     w_peephole = array_ops.zeros(cs_prev.shape[1:], dtype=w.dtype)
     _, all_cs, _, _, _, _, all_h = op(

From 41346a6788a30944d7255d3f82c25e8a329cb800 Mon Sep 17 00:00:00 2001
From: Jake VanderPlas <vanderplas@google.com>
Date: Thu, 8 Oct 2026 19:11:53 -0700
Subject: [PATCH 0322/1000] Document the Mulhi operation in XLA operation
 semantics.

Add `Mulhi` (`XlaBuilder::Mulhi`) to `docs/operation_semantics.md`, documenting its element-wise high-half integer multiplication semantics, broadcasting overloads, and relationship to `chlo.mulhi`.

PiperOrigin-RevId: 996215278
---
 third_party/xla/docs/operation_semantics.md | 52 +++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/third_party/xla/docs/operation_semantics.md b/third_party/xla/docs/operation_semantics.md
index ffc81985ea8..4fb025f39b4 100644
--- a/third_party/xla/docs/operation_semantics.md
+++ b/third_party/xla/docs/operation_semantics.md
@@ -3472,6 +3472,58 @@ in detail on the [broadcasting page](broadcasting.md).
 For StableHLO information see
 [StableHLO - multiply](https://openxla.org/stablehlo/spec#multiply).
 
+## Mulhi
+
+See also
+[`XlaBuilder::Mulhi`](https://github.com/openxla/xla/tree/main/xla/hlo/builder/xla_builder.h).
+
+Performs element-wise multiplication of two N-bit integer arrays `lhs` and
+`rhs`, returning an N-bit integer array containing the most significant N bits
+of the upcasted (N+N-bit) product:
+
+$$
+\text{mulhi}(x, y) = \text{downcast}((\text{upcast}(x) * \text{upcast}(y)) >> N)
+$$
+
+**`Mulhi(lhs, rhs)`**
+
+Arguments | Type  | Semantics
+--------- | ----- | -------------------------------------------------
+lhs       | XlaOp | Left-hand-side operand: array of integer type T
+rhs       | XlaOp | Right-hand-side operand: array of integer type T
+
+The arguments' shapes have to be either similar or compatible. See the
+[broadcasting](broadcasting.md) documentation about what it means for shapes to
+be compatible. The result of an operation has a shape which is the result of
+broadcasting the two input arrays. In this variant, operations between arrays of
+different ranks are *not* supported, unless one of the operands is a scalar.
+
+An alternative variant with different-dimensional broadcasting support exists
+for Mulhi:
+
+**`Mulhi(lhs,rhs, broadcast_dimensions)`**
+
+| Arguments           | Type              | Semantics                         |
+| ------------------- | ----------------- | --------------------------------- |
+| lhs                 | XlaOp             | Left-hand-side operand: array of  |
+:                     :                   : integer type T                    :
+| rhs                 | XlaOp             | Right-hand-side operand: array of |
+:                     :                   : integer type T                    :
+| broadcast_dimension | ArraySlice<int64> | Which dimension in the target     |
+:                     :                   : shape each dimension of the       :
+:                     :                   : operand shape corresponds to      :
+
+This variant of the operation should be used for arithmetic operations between
+arrays of different ranks (such as adding a matrix to a vector).
+
+The additional broadcast_dimensions operand is a slice of integers specifying
+the dimensions to use for broadcasting the operands. The semantics are described
+in detail on the [broadcasting page](broadcasting.md).
+
+> **Note:** `Mulhi` is only found in HLO and not found in StableHLO. CHLO
+> `Mulhi` in Frameworks will lower to HLO `Mulhi` see
+> [StableHLO - chlo.mulhi](https://openxla.org/stablehlo/generated/chlo?hl=en#chlomulhi_chlomulhiop)
+
 ## Neg
 
 See also

From dbc0789aebedb496b02ecda863f3ccd398e8cde8 Mon Sep 17 00:00:00 2001
From: Ashutosh0x <ashutoshkumarsingh0x@gmail.com>
Date: Fri, 9 Oct 2026 08:00:45 +0530
Subject: [PATCH 0323/1000] Validate peephole weight sizes and drop impossible
 rank checks

LSTMBlockCellOp and LSTMBlockCellGradOp checked that wci, wcf and wco are
rank 1 but not that their length is cell_size. A wrong length is
broadcast against [batch_size, cell_size] tensors, and in the gradient
wci_grad is allocated with the shape of wci but written with cell_size
elements. Reject it alongside the existing b size check.

Remove the rank checks on LSTMBlockCellOp's temp and output tensors.
They are allocated with explicit two-dimensional shapes just above, so
the checks cannot fail.
---
 tensorflow/core/kernels/rnn/lstm_ops.cc | 64 ++++++++++---------------
 tensorflow/python/ops/rnn_grad_test.py  | 21 ++++++++
 2 files changed, 47 insertions(+), 38 deletions(-)

diff --git a/tensorflow/core/kernels/rnn/lstm_ops.cc b/tensorflow/core/kernels/rnn/lstm_ops.cc
index fe35e65428e..7bd565cd28b 100644
--- a/tensorflow/core/kernels/rnn/lstm_ops.cc
+++ b/tensorflow/core/kernels/rnn/lstm_ops.cc
@@ -399,6 +399,19 @@ class LSTMBlockCellOp : public OpKernel {
                     "b.dim_size(0) != cell_size * 4: ", b_tensor->dim_size(0),
                     " vs. ", cell_size * 4)));
 
+    OP_REQUIRES(ctx, wci_tensor->dim_size(0) == cell_size,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "wci.dim_size(0) != cell_size: ", wci_tensor->dim_size(0),
+                    " vs. ", cell_size)));
+    OP_REQUIRES(ctx, wcf_tensor->dim_size(0) == cell_size,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "wcf.dim_size(0) != cell_size: ", wcf_tensor->dim_size(0),
+                    " vs. ", cell_size)));
+    OP_REQUIRES(ctx, wco_tensor->dim_size(0) == cell_size,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "wco.dim_size(0) != cell_size: ", wco_tensor->dim_size(0),
+                    " vs. ", cell_size)));
+
     // Allocate our output tensors.
     Tensor* i_tensor = nullptr;
     OP_REQUIRES_OK(ctx, ctx->forward_input_or_allocate_output(
@@ -450,44 +463,6 @@ class LSTMBlockCellOp : public OpKernel {
 
     const Device& device = ctx->eigen_device<Device>();
 
-    // Sanity check that each of the tensors have the required NDIMS.
-    OP_REQUIRES(
-        ctx, xh_tensor.dims() == 2,
-        absl::InvalidArgumentError(absl::StrCat(
-            "xh_tensor must be rank 2 but is rank ", xh_tensor.dims(), ".")));
-    OP_REQUIRES(
-        ctx, i_tensor->dims() == 2,
-        absl::InvalidArgumentError(absl::StrCat(
-            "i_tensor must be rank 2 but is rank ", i_tensor->dims(), ".")));
-    OP_REQUIRES(
-        ctx, cs_tensor->dims() == 2,
-        absl::InvalidArgumentError(absl::StrCat(
-            "cs_tensor must be rank 2 but is rank ", cs_tensor->dims(), ".")));
-    OP_REQUIRES(
-        ctx, f_tensor->dims() == 2,
-        absl::InvalidArgumentError(absl::StrCat(
-            "f_tensor must be rank 2 but is rank ", f_tensor->dims(), ".")));
-    OP_REQUIRES(
-        ctx, o_tensor->dims() == 2,
-        absl::InvalidArgumentError(absl::StrCat(
-            "o_tensor must be rank 2 but is rank ", o_tensor->dims(), ".")));
-    OP_REQUIRES(
-        ctx, ci_tensor->dims() == 2,
-        absl::InvalidArgumentError(absl::StrCat(
-            "ci_tensor must be rank 2 but is rank ", ci_tensor->dims(), ".")));
-    OP_REQUIRES(
-        ctx, co_tensor->dims() == 2,
-        absl::InvalidArgumentError(absl::StrCat(
-            "co_tensor must be rank 2 but is rank ", co_tensor->dims(), ".")));
-    OP_REQUIRES(ctx, gates_tensor.dims() == 2,
-                absl::InvalidArgumentError(
-                    absl::StrCat("gates_tensor must be rank 2 but is rank ",
-                                 gates_tensor.dims(), ".")));
-    OP_REQUIRES(
-        ctx, h_tensor->dims() == 2,
-        absl::InvalidArgumentError(absl::StrCat(
-            "h_tensor must be rank 2 but is rank ", h_tensor->dims(), ".")));
-
     functor::LSTMBlockCellFprop<Device, T, USE_CUBLAS, gate_layout>(
         batch_size, input_size, cell_size)(
         ctx, device, forget_bias_, cell_clip_, use_peephole_,
@@ -658,6 +633,19 @@ class LSTMBlockCellGradOp : public OpKernel {
                     "b.dim_size(0) != cell_size * 4: ", b_tensor->dim_size(0),
                     " vs. ", cell_size * 4)));
 
+    OP_REQUIRES(ctx, wci_tensor->dim_size(0) == cell_size,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "wci.dim_size(0) != cell_size: ", wci_tensor->dim_size(0),
+                    " vs. ", cell_size)));
+    OP_REQUIRES(ctx, wcf_tensor->dim_size(0) == cell_size,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "wcf.dim_size(0) != cell_size: ", wcf_tensor->dim_size(0),
+                    " vs. ", cell_size)));
+    OP_REQUIRES(ctx, wco_tensor->dim_size(0) == cell_size,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "wco.dim_size(0) != cell_size: ", wco_tensor->dim_size(0),
+                    " vs. ", cell_size)));
+
     OP_REQUIRES(ctx, i_tensor->dim_size(0) == batch_size,
                 absl::InvalidArgumentError(absl::StrCat(
                     "i.dim_size(0) != batch_size: ", i_tensor->dim_size(0),
diff --git a/tensorflow/python/ops/rnn_grad_test.py b/tensorflow/python/ops/rnn_grad_test.py
index b928b8aefca..59f70857d47 100644
--- a/tensorflow/python/ops/rnn_grad_test.py
+++ b/tensorflow/python/ops/rnn_grad_test.py
@@ -272,6 +272,27 @@ class RNNGradTest(test.TestCase):
           gen_rnn_ops.lstm_block_cell_grad, "wci", use_peephole=False,
           **kwargs)
 
+  def testLSTMBlockCellInvalidPeepholeSize(self):
+    # A peephole weight whose length differs from cell_size is broadcast
+    # against [batch_size, cell_size] tensors and read out of bounds.
+    kwargs = self._lstm_block_cell_inputs()
+    kwargs["wci"] = deterministic_random_uniform([3])
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                r"wci.dim_size\(0\) != cell_size"):
+      self._run_with_unknown_shape(
+          gen_rnn_ops.lstm_block_cell, "wci", use_peephole=True, **kwargs)
+
+  def testLSTMBlockCellGradInvalidPeepholeSize(self):
+    # wci_grad is allocated with the shape of wci but written with cell_size
+    # elements.
+    kwargs = self._lstm_block_cell_grad_inputs()
+    kwargs["wci"] = deterministic_random_uniform([3])
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                r"wci.dim_size\(0\) != cell_size"):
+      self._run_with_unknown_shape(
+          gen_rnn_ops.lstm_block_cell_grad, "wci", use_peephole=False,
+          **kwargs)
+
   def testLSTMBlockCellGradEmptyCsPrev(self):
     # An empty cs_prev gives LSTMBlockCellBpropWithCUDA a zero-sized launch
     # configuration, which aborts the process. The forward op already rejects

From c8b55623100599f13ded720d1b087b30e32d3a7c Mon Sep 17 00:00:00 2001
From: Ashutosh0x <ashutoshkumarsingh0x@gmail.com>
Date: Fri, 9 Oct 2026 08:19:15 +0530
Subject: [PATCH 0324/1000] Validate peephole weight sizes in BlockLSTMGrad and
 test all three

BlockLSTMGradOp checked that wci, wcf and wco are rank 1 but not that
their length is cell_size, unlike BlockLSTMOp. wci_grad is allocated with
the shape of wci and written with cell_size elements, so reject a
mismatched length the way the forward op does.

Cover wci, wcf and wco in the peephole size tests for LSTMBlockCell,
LSTMBlockCellGrad and BlockLSTMGrad.

Fix the f_tensor size message in LSTMBlockCellGradOp, which named i.
---
 tensorflow/core/kernels/rnn/lstm_ops.cc | 14 ++++++-
 tensorflow/python/ops/rnn_grad_test.py  | 50 ++++++++++++++++++-------
 2 files changed, 50 insertions(+), 14 deletions(-)

diff --git a/tensorflow/core/kernels/rnn/lstm_ops.cc b/tensorflow/core/kernels/rnn/lstm_ops.cc
index 7bd565cd28b..47a8f5709a2 100644
--- a/tensorflow/core/kernels/rnn/lstm_ops.cc
+++ b/tensorflow/core/kernels/rnn/lstm_ops.cc
@@ -670,7 +670,7 @@ class LSTMBlockCellGradOp : public OpKernel {
                     " vs. ", batch_size)));
     OP_REQUIRES(ctx, f_tensor->dim_size(1) == cell_size,
                 absl::InvalidArgumentError(absl::StrCat(
-                    "i.dim_size(1) != cell_size: ", f_tensor->dim_size(1),
+                    "f.dim_size(1) != cell_size: ", f_tensor->dim_size(1),
                     " vs. ", cell_size)));
 
     OP_REQUIRES(ctx, o_tensor->dim_size(0) == batch_size,
@@ -1236,18 +1236,30 @@ class BlockLSTMGradOp : public OpKernel {
     OP_REQUIRES(ctx, wci_tensor->dims() == 1,
                 absl::InvalidArgumentError(absl::StrCat(
                     "wci must be rank 1 but is rank ", wci_tensor->dims())));
+    OP_REQUIRES(ctx, wci_tensor->dim_size(0) == cell_size,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "wci.dim_size(0) != cell_size: ", wci_tensor->dim_size(0),
+                    " vs. ", cell_size)));
 
     const Tensor* wcf_tensor = nullptr;
     OP_REQUIRES_OK(ctx, ctx->input("wcf", &wcf_tensor));
     OP_REQUIRES(ctx, wcf_tensor->dims() == 1,
                 absl::InvalidArgumentError(absl::StrCat(
                     "wcf must be rank 1 but is rank ", wcf_tensor->dims())));
+    OP_REQUIRES(ctx, wcf_tensor->dim_size(0) == cell_size,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "wcf.dim_size(0) != cell_size: ", wcf_tensor->dim_size(0),
+                    " vs. ", cell_size)));
 
     const Tensor* wco_tensor = nullptr;
     OP_REQUIRES_OK(ctx, ctx->input("wco", &wco_tensor));
     OP_REQUIRES(ctx, wco_tensor->dims() == 1,
                 absl::InvalidArgumentError(absl::StrCat(
                     "wco must be rank 1 but is rank ", wco_tensor->dims())));
+    OP_REQUIRES(ctx, wco_tensor->dim_size(0) == cell_size,
+                absl::InvalidArgumentError(absl::StrCat(
+                    "wco.dim_size(0) != cell_size: ", wco_tensor->dim_size(0),
+                    " vs. ", cell_size)));
 
     const Tensor* b_tensor = nullptr;
     OP_REQUIRES_OK(ctx, ctx->input("b", &b_tensor));
diff --git a/tensorflow/python/ops/rnn_grad_test.py b/tensorflow/python/ops/rnn_grad_test.py
index 59f70857d47..dd5dc0cbc12 100644
--- a/tensorflow/python/ops/rnn_grad_test.py
+++ b/tensorflow/python/ops/rnn_grad_test.py
@@ -275,23 +275,47 @@ class RNNGradTest(test.TestCase):
   def testLSTMBlockCellInvalidPeepholeSize(self):
     # A peephole weight whose length differs from cell_size is broadcast
     # against [batch_size, cell_size] tensors and read out of bounds.
-    kwargs = self._lstm_block_cell_inputs()
-    kwargs["wci"] = deterministic_random_uniform([3])
-    with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
-                                r"wci.dim_size\(0\) != cell_size"):
-      self._run_with_unknown_shape(
-          gen_rnn_ops.lstm_block_cell, "wci", use_peephole=True, **kwargs)
+    for name in ["wci", "wcf", "wco"]:
+      kwargs = self._lstm_block_cell_inputs()
+      kwargs[name] = deterministic_random_uniform([3])
+      with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                  f"{name}.dim_size\\(0\\) != cell_size"):
+        self._run_with_unknown_shape(
+            gen_rnn_ops.lstm_block_cell, name, use_peephole=True, **kwargs)
 
   def testLSTMBlockCellGradInvalidPeepholeSize(self):
     # wci_grad is allocated with the shape of wci but written with cell_size
     # elements.
-    kwargs = self._lstm_block_cell_grad_inputs()
-    kwargs["wci"] = deterministic_random_uniform([3])
-    with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
-                                r"wci.dim_size\(0\) != cell_size"):
-      self._run_with_unknown_shape(
-          gen_rnn_ops.lstm_block_cell_grad, "wci", use_peephole=False,
-          **kwargs)
+    for name in ["wci", "wcf", "wco"]:
+      kwargs = self._lstm_block_cell_grad_inputs()
+      kwargs[name] = deterministic_random_uniform([3])
+      with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                  f"{name}.dim_size\\(0\\) != cell_size"):
+        self._run_with_unknown_shape(
+            gen_rnn_ops.lstm_block_cell_grad, name, use_peephole=False,
+            **kwargs)
+
+  def testBlockLSTMGradInvalidPeepholeSize(self):
+    # As in LSTMBlockCellGrad, wci_grad is allocated with the shape of wci but
+    # written with cell_size elements.
+    for name in ["wci", "wcf", "wco"]:
+      w, b, x, cs_prev, h_prev, w_peephole = self._block_lstm_inputs()
+      i, cs, f, o, ci, co, h = self._block_lstm(
+          w, b, x, cs_prev, h_prev, w_peephole, seq_len_max=x.shape[0])
+      kwargs = dict(
+          seq_len_max=math_ops.cast(x.shape[0], dtypes.int64),
+          x=x, cs_prev=cs_prev, h_prev=h_prev, w=w,
+          wci=w_peephole, wcf=w_peephole, wco=w_peephole, b=b,
+          i=i, cs=cs, f=f, o=o, ci=ci, co=co, h=h,
+          cs_grad=deterministic_random_uniform(cs.shape.as_list()),
+          h_grad=deterministic_random_uniform(h.shape.as_list()),
+          use_peephole=False,
+      )
+      kwargs[name] = deterministic_random_uniform([3])
+      with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                  f"{name}.dim_size\\(0\\) != cell_size"):
+        self._run_with_unknown_shape(
+            gen_rnn_ops.BlockLSTMGrad, name, **kwargs)
 
   def testLSTMBlockCellGradEmptyCsPrev(self):
     # An empty cs_prev gives LSTMBlockCellBpropWithCUDA a zero-sized launch

From d1df2541279b047a6f857e7e9fbf0b19317b2f93 Mon Sep 17 00:00:00 2001
From: Ashutosh0x <ashutoshkumarsingh0x@gmail.com>
Date: Fri, 9 Oct 2026 08:51:00 +0530
Subject: [PATCH 0325/1000] Require an exact bias size in BlockLSTMGrad and
 drop duplicate checks

BlockLSTMGradOp accepted a bias of length 4 * cell_size + 1..3 because it
compared cell_size with b.dim_size(0) / 4. Compare b.dim_size(0) with
cell_size * 4, as BlockLSTMOp does.

Remove its rank-3 checks on i, cs, f, o, ci, co and h; check_dims
validates the rank of the same tensors before any dimension is read.

Add tests for the BlockLSTMGrad bias size and the BlockLSTM peephole
sizes, and share the BlockLSTM(Grad) inputs between the tests.

Correct the comment on the empty cs_prev check in LSTMBlockCellGradOp:
GpuLaunchKernel skips zero-sized launches, so the check keeps the
gradient consistent with the forward op rather than preventing an abort.
---
 tensorflow/core/kernels/rnn/lstm_ops.cc | 29 ++-----------
 tensorflow/python/ops/rnn_grad_test.py  | 55 ++++++++++++++++++-------
 2 files changed, 43 insertions(+), 41 deletions(-)

diff --git a/tensorflow/core/kernels/rnn/lstm_ops.cc b/tensorflow/core/kernels/rnn/lstm_ops.cc
index 47a8f5709a2..4a3eacb0626 100644
--- a/tensorflow/core/kernels/rnn/lstm_ops.cc
+++ b/tensorflow/core/kernels/rnn/lstm_ops.cc
@@ -586,9 +586,8 @@ class LSTMBlockCellGradOp : public OpKernel {
     OP_REQUIRES_OK(ctx, check_rank(cs_grad_tensor, "cs_grad", 2));
     OP_REQUIRES_OK(ctx, check_rank(h_grad_tensor, "h_grad", 2));
 
-    // An empty cs_prev yields a zero-sized GPU launch configuration in
-    // LSTMBlockCellBpropWithCUDA, which aborts via TF_CHECK_OK. The forward
-    // op rejects this for the same reason; see #58270.
+    // Reject an empty cs_prev as LSTMBlockCellOp does (#58270), so the
+    // gradient accepts exactly the inputs the forward op accepts.
     OP_REQUIRES(
         ctx, cs_prev_tensor->dim_size(0) > 0 && cs_prev_tensor->dim_size(1) > 0,
         absl::InvalidArgumentError(
@@ -1266,7 +1265,7 @@ class BlockLSTMGradOp : public OpKernel {
     OP_REQUIRES(ctx, b_tensor->dims() == 1,
                 absl::InvalidArgumentError(absl::StrCat(
                     "b must be rank 1 but is rank ", b_tensor->dims())));
-    OP_REQUIRES(ctx, cell_size == b_tensor->dim_size(0) / 4,
+    OP_REQUIRES(ctx, b_tensor->dim_size(0) == cell_size * 4,
                 absl::InvalidArgumentError(
                     absl::StrCat("w and b cell_size don't match: ", cell_size,
                                  " vs. ", b_tensor->dim_size(0))));
@@ -1292,28 +1291,6 @@ class BlockLSTMGradOp : public OpKernel {
     const Tensor* h_out = nullptr;
     OP_REQUIRES_OK(ctx, ctx->input("h", &h_out));
 
-    OP_REQUIRES(ctx, i_out->dims() == 3,
-                absl::InvalidArgumentError(absl::StrCat(
-                    "i must be rank 3. Received: ", i_out->dims())));
-    OP_REQUIRES(ctx, cs_out->dims() == 3,
-                absl::InvalidArgumentError(absl::StrCat(
-                    "cs must be rank 3. Received: ", cs_out->dims())));
-    OP_REQUIRES(ctx, f_out->dims() == 3,
-                absl::InvalidArgumentError(absl::StrCat(
-                    "f must be rank 3. Received: ", f_out->dims())));
-    OP_REQUIRES(ctx, o_out->dims() == 3,
-                absl::InvalidArgumentError(absl::StrCat(
-                    "o must be rank 3. Received: ", o_out->dims())));
-    OP_REQUIRES(ctx, ci_out->dims() == 3,
-                absl::InvalidArgumentError(absl::StrCat(
-                    "ci must be rank 3. Received: ", ci_out->dims())));
-    OP_REQUIRES(ctx, co_out->dims() == 3,
-                absl::InvalidArgumentError(absl::StrCat(
-                    "co must be rank 3. Received: ", co_out->dims())));
-    OP_REQUIRES(ctx, h_out->dims() == 3,
-                absl::InvalidArgumentError(absl::StrCat(
-                    "h must be rank 3. Received: ", h_out->dims())));
-
     const Tensor* cs_grad = nullptr;
     OP_REQUIRES_OK(ctx, ctx->input("cs_grad", &cs_grad));
 
diff --git a/tensorflow/python/ops/rnn_grad_test.py b/tensorflow/python/ops/rnn_grad_test.py
index dd5dc0cbc12..904c9221984 100644
--- a/tensorflow/python/ops/rnn_grad_test.py
+++ b/tensorflow/python/ops/rnn_grad_test.py
@@ -295,32 +295,57 @@ class RNNGradTest(test.TestCase):
             gen_rnn_ops.lstm_block_cell_grad, name, use_peephole=False,
             **kwargs)
 
+  def _block_lstm_kwargs(self):
+    w, b, x, cs_prev, h_prev, w_peephole = self._block_lstm_inputs()
+    return dict(
+        seq_len_max=math_ops.cast(x.shape[0], dtypes.int64),
+        x=x, cs_prev=cs_prev, h_prev=h_prev, w=w,
+        wci=w_peephole, wcf=w_peephole, wco=w_peephole, b=b,
+        use_peephole=False,
+    )
+
+  def _block_lstm_grad_kwargs(self):
+    kwargs = self._block_lstm_kwargs()
+    i, cs, f, o, ci, co, h = gen_rnn_ops.BlockLSTM(**kwargs)
+    kwargs.update(
+        i=i, cs=cs, f=f, o=o, ci=ci, co=co, h=h,
+        cs_grad=deterministic_random_uniform(cs.shape.as_list()),
+        h_grad=deterministic_random_uniform(h.shape.as_list()),
+    )
+    return kwargs
+
+  def testBlockLSTMInvalidPeepholeSize(self):
+    for name in ["wci", "wcf", "wco"]:
+      kwargs = self._block_lstm_kwargs()
+      kwargs[name] = deterministic_random_uniform([3])
+      with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                  f"{name}.dim_size\\(0\\) != cell_size"):
+        self._run_with_unknown_shape(gen_rnn_ops.BlockLSTM, name, **kwargs)
+
   def testBlockLSTMGradInvalidPeepholeSize(self):
     # As in LSTMBlockCellGrad, wci_grad is allocated with the shape of wci but
     # written with cell_size elements.
     for name in ["wci", "wcf", "wco"]:
-      w, b, x, cs_prev, h_prev, w_peephole = self._block_lstm_inputs()
-      i, cs, f, o, ci, co, h = self._block_lstm(
-          w, b, x, cs_prev, h_prev, w_peephole, seq_len_max=x.shape[0])
-      kwargs = dict(
-          seq_len_max=math_ops.cast(x.shape[0], dtypes.int64),
-          x=x, cs_prev=cs_prev, h_prev=h_prev, w=w,
-          wci=w_peephole, wcf=w_peephole, wco=w_peephole, b=b,
-          i=i, cs=cs, f=f, o=o, ci=ci, co=co, h=h,
-          cs_grad=deterministic_random_uniform(cs.shape.as_list()),
-          h_grad=deterministic_random_uniform(h.shape.as_list()),
-          use_peephole=False,
-      )
+      kwargs = self._block_lstm_grad_kwargs()
       kwargs[name] = deterministic_random_uniform([3])
       with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
                                   f"{name}.dim_size\\(0\\) != cell_size"):
         self._run_with_unknown_shape(
             gen_rnn_ops.BlockLSTMGrad, name, **kwargs)
 
+  def testBlockLSTMGradInvalidBiasSize(self):
+    # A bias of length 4 * cell_size + 3 passed the old
+    # `cell_size == b.dim_size(0) / 4` check through integer division.
+    kwargs = self._block_lstm_grad_kwargs()
+    cell_size = kwargs["cs_prev"].shape[1]
+    kwargs["b"] = deterministic_random_uniform([4 * cell_size + 3])
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                "w and b cell_size don't match"):
+      self._run_with_unknown_shape(gen_rnn_ops.BlockLSTMGrad, "b", **kwargs)
+
   def testLSTMBlockCellGradEmptyCsPrev(self):
-    # An empty cs_prev gives LSTMBlockCellBpropWithCUDA a zero-sized launch
-    # configuration, which aborts the process. The forward op already rejects
-    # it; see #58270.
+    # The forward op rejects an empty cs_prev (#58270); the gradient must
+    # reject the same inputs.
     kwargs = self._lstm_block_cell_grad_inputs(cell_size=0)
     with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
                                 "cs_prev_tensor is empty"):

From 1baf8e88c9fd43dd2ace8fce9ed56fe79b24e954 Mon Sep 17 00:00:00 2001
From: Boris Kozak <boriskozak@gmail.com>
Date: Fri, 9 Oct 2026 03:52:49 +0000
Subject: [PATCH 0326/1000] Address round-6 review: test comment, rename,
 docstring

Add comment in testSingleCoeffInfReturnsNan explaining degree-0
inf -> NaN (IEEE 754 0*inf=NaN). Rename testNonFiniteXReturnsInf to
testInfXReturnsInf with docstring clarifying NaN x propagates to NaN.
---
 .../python/kernel_tests/math_ops/cwise_ops_test.py     | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py b/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
index cb1cfd617d3..eb301b64623 100644
--- a/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/cwise_ops_test.py
@@ -1342,6 +1342,9 @@ class PolyvalTest(test.TestCase):
 
   @test_util.run_in_graph_and_eager_modes
   def testSingleCoeffInfReturnsNan(self):
+    # Degree-0 polynomials return NaN for +/-inf inputs: the x * 0 step in
+    # math_ops.py both broadcasts x's shape and propagates NaN, and IEEE 754
+    # specifies 0 * inf = NaN (matching numpy.polyval behavior).
     for dtype in [
         np.float16, np.float32, np.float64, np.complex64, np.complex128
     ]:
@@ -1361,7 +1364,12 @@ class PolyvalTest(test.TestCase):
               self.assertTrue(np.all(np.isnan(result)))
 
   @test_util.run_in_graph_and_eager_modes
-  def testNonFiniteXReturnsInf(self):
+  def testInfXReturnsInf(self):
+    """Tests that +/-inf x evaluates to inf for degree >= 1 polynomials.
+
+    Covers inf inputs only: NaN x inputs correctly propagate to NaN (NaN
+    propagation) rather than returning inf.
+    """
     # inf x must evaluate to inf (not nan): the limit of a polynomial such as
     # p(x) = x + 1 as x -> inf is mathematically inf. NumPy returns nan here
     # only because it starts its Horner accumulator at zero (0 * inf -> nan).

From f4034c404466c1820ed27dcd707cdabc2ff8ee2f Mon Sep 17 00:00:00 2001
From: bodapatisaikrishna <saikrishnabodapati@gmail.com>
Date: Fri, 9 Oct 2026 09:29:10 +0530
Subject: [PATCH 0327/1000] Consolidate conv2d invalid rank tests into
 parameterized loop

Signed-off-by: bodapatisaikrishna <saikrishnabodapati@gmail.com>
---
 .../kernel_tests/nn_ops/conv_ops_test.py      | 102 +++++-------------
 1 file changed, 26 insertions(+), 76 deletions(-)

diff --git a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
index 37e185fe2d7..515ad0b1081 100644
--- a/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/conv_ops_test.py
@@ -47,8 +47,6 @@ from tensorflow.python.platform import test
 from tensorflow.python.platform import tf_logging
 from tensorflow.python.util.compat import collections_abc
 
-raw_ops = gen_nn_ops
-
 
 def GetShrunkInceptionShapes(shrink=10):
   """Iterator for smaller versions of convolution shapes in 2015 Inception.
@@ -3135,82 +3133,34 @@ class Conv2DTest(parameterized.TestCase, test.TestCase):
     #           padding=[[0, 0], [2, 2], [2, 2], [0, 0]]))
 
     # Invalid input or filter rank (< 4) in eager/graph execution.
-    with self.assertRaisesRegex(
-        (ValueError, errors_impl.InvalidArgumentError),
-        r"(must be 4-dimensional|Shape must be rank 4 but is rank \d+)",
-    ):
-      self.evaluate(
-          nn_ops.conv2d(
-              np.ones([1, 1, 3], dtype=np.float32),
-              np.ones([1, 2, 2, 1], dtype=np.float32),
-              strides=[1, 1, 1, 1],
-              padding="VALID",
+    invalid_shapes = [(), (1,), (1, 1), (1, 1, 3)]
+    valid_shape = [1, 1, 1, 1]
+    for op in (nn_ops.conv2d, gen_nn_ops.Conv2D):
+      for bad_shape in invalid_shapes:
+        with self.assertRaisesRegex(
+            (ValueError, errors_impl.InvalidArgumentError),
+            r"(must be 4-dimensional|Shape must be rank 4 but is rank \d+)",
+        ):
+          self.evaluate(
+              op(
+                  input=np.ones(bad_shape, dtype=np.float32),
+                  filter=np.ones(valid_shape, dtype=np.float32),
+                  strides=[1, 1, 1, 1],
+                  padding="VALID",
+              )
           )
-      )
-    with self.assertRaisesRegex(
-        (ValueError, errors_impl.InvalidArgumentError),
-        r"(must be 4-dimensional|Shape must be rank 4 but is rank \d+)",
-    ):
-      self.evaluate(
-          nn_ops.conv2d(
-              np.ones([1, 1, 3, 3], dtype=np.float32),
-              np.ones([1, 2, 2], dtype=np.float32),
-              strides=[1, 1, 1, 1],
-              padding="VALID",
+        with self.assertRaisesRegex(
+            (ValueError, errors_impl.InvalidArgumentError),
+            r"(must be 4-dimensional|Shape must be rank 4 but is rank \d+)",
+        ):
+          self.evaluate(
+              op(
+                  input=np.ones(valid_shape, dtype=np.float32),
+                  filter=np.ones(bad_shape, dtype=np.float32),
+                  strides=[1, 1, 1, 1],
+                  padding="VALID",
+              )
           )
-      )
-
-    # Ensure raw ops are protected by the C++ kernel boundary.
-    with self.assertRaisesRegex(
-        (ValueError, errors_impl.InvalidArgumentError),
-        r"(must be 4-dimensional|Shape must be rank 4 but is rank \d+)",
-    ):
-      self.evaluate(
-          raw_ops.Conv2D(
-              input=np.ones([1, 1, 3], dtype=np.float32),
-              filter=np.ones([1, 2, 2, 1], dtype=np.float32),
-              strides=[1, 1, 1, 1],
-              padding="VALID",
-          )
-      )
-    with self.assertRaisesRegex(
-        (ValueError, errors_impl.InvalidArgumentError),
-        r"(must be 4-dimensional|Shape must be rank 4 but is rank \d+)",
-    ):
-      self.evaluate(
-          raw_ops.Conv2D(
-              input=np.ones([1, 1, 3, 3], dtype=np.float32),
-              filter=np.ones([1, 2, 2], dtype=np.float32),
-              strides=[1, 1, 1, 1],
-              padding="VALID",
-          )
-      )
-
-    # Ensure scalar and rank-1 inputs fail safely.
-    with self.assertRaisesRegex(
-        (ValueError, errors_impl.InvalidArgumentError),
-        r"(must be 4-dimensional|Shape must be rank 4 but is rank \d+)",
-    ):
-      self.evaluate(
-          raw_ops.Conv2D(
-              input=np.ones([1], dtype=np.float32),
-              filter=np.ones([1, 2, 2, 1], dtype=np.float32),
-              strides=[1, 1, 1, 1],
-              padding="VALID",
-          )
-      )
-    with self.assertRaisesRegex(
-        (ValueError, errors_impl.InvalidArgumentError),
-        r"(must be 4-dimensional|Shape must be rank 4 but is rank \d+)",
-    ):
-      self.evaluate(
-          raw_ops.Conv2D(
-              input=np.ones([1, 1, 3, 3], dtype=np.float32),
-              filter=np.ones([], dtype=np.float32),
-              strides=[1, 1, 1, 1],
-              padding="VALID",
-          )
-      )
 
     # Filter dimensions must be greater than 0.
     with self.assertRaisesRegex(

From 3ed83ea8ddd7c83307bc51686c7e5088e0aa4cce Mon Sep 17 00:00:00 2001
From: destro4evr <destro4evr@gmail.com>
Date: Thu, 8 Oct 2026 22:44:38 -0700
Subject: [PATCH 0328/1000] PR #128593: security: guard SparseCross/RaggedCross
 cross-count against int64 overflow
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/tensorflow/tensorflow/pull/128593

## Problem

`CrossCountByBatchIndex` in both `sparse_cross_op.cc` and `ragged_cross_op.cc` multiplied per-column feature counts with no overflow check. With 64 sparse columns of 2 features each, the int64 product wraps to zero - output tensors get allocated at size 0, but `ProductIterator` still iterates and writes to position 0 - heap out-of-bounds write.

Secondary: `TensorShape({cross_count_total, 2})` internally computes `total * 2`, so allowing `total` up to `INT64_MAX` causes that multiply to overflow and crash the process. Cap needs to be `INT64_MAX / 2`.

Third: when `out_row_splits_type=DT_INT32`, assigning a value above `INT32_MAX` to `flat_row_splits` silently wraps negative, causing an OOB write in `WriteOutputSlice`. Cap must be `min(SplitsType::max(), INT64_MAX/2)`.

## Changes

- `sparse_cross_op.cc`: `CrossCountByBatchIndex` returns `StatusOr<int64_t>`; rejects negative feature counts; caps at `INT64_MAX/2`; accumulation guard capped equally; batch loop counters and permutation state vector widened to `int64_t`; uses `xla/tsl/platform/statusor.h`.
- `ragged_cross_op.cc`: same cap via existing `-1` sentinel; cap is `min(SplitsType::max(), INT64_MAX/2)` for int32 splits safety; combination vector and batch_index widened to `int64_t`; error message improved.
- `sparse_cross_op_test.py`: four tests - per-batch overflow (string + hashed) and cross-batch accumulation overflow (string + hashed).
- `ragged_cross_op_test.py`: two tests - int64 splits overflow (64 cols × 2 features) and int32 splits overflow (7 cols × 32 features).
Copybara import of the project:

--
6ee3f7d41d8fba74922c680c3265dfa60f79b77e by destro4evr-rgb <destro4evr@proton.me>:

security: guard SparseCross/RaggedCross cross-count against int64 overflow

CrossCountByBatchIndex in both sparse_cross_op.cc and ragged_cross_op.cc
multiplied per-column feature counts with no overflow check. Supplying
64 sparse columns each with 2 features per batch causes the int64 product
to wrap to zero, making CreateOutputTensors/BuildOutputTensors allocate
zero-element output tensors. The ProductIterator still sees non-zero
feature counts on every column and begins iterating immediately, writing
the first cross into position 0 of a zero-element tensor — a heap
out-of-bounds write.

Secondary issue: the indices output uses TensorShape({total, 2}), which
internally computes total*2. Allowing total up to INT64_MAX means that
multiplication itself can overflow and crash the process.

Fix:
- sparse_cross_op.cc: CrossCountByBatchIndex returns StatusOr<int64_t>;
  rejects negative feature_count; caps at INT64_MAX/2 to keep
  TensorShape({total,2}) safe; guards cross_count_total accumulation;
  widens batch_index and all loop counters to int64_t; widens permutation
  state vector to std::vector<int64_t>; uses xla/tsl/platform/statusor.h.
- ragged_cross_op.cc: same cap using the existing -1 sentinel convention;
  cap is min(SplitsType::max(), INT64_MAX/2) to also prevent silent
  truncation when out_row_splits_type=DT_INT32; widens batch_index and
  combination state vector to int64_t; improves the error message that
  was previously just "Invalid RaggedTensor".
- sparse_cross_op_test.py: four tests — per-batch overflow (string and
  hashed) and cross-batch accumulation overflow (string and hashed).
- ragged_cross_op_test.py: two tests — int64 splits overflow (64 cols)
  and int32 splits overflow (7 cols x 32 features = 32^7 > INT32_MAX).

Merging this change closes #128593

COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/tensorflow/pull/128593 from destro4evr-rgb:fix/sparse-cross-cross-count-overflow 6ee3f7d41d8fba74922c680c3265dfa60f79b77e
PiperOrigin-RevId: 996297594
---
 tensorflow/core/kernels/ragged_cross_op.cc    | 34 ++++++++---
 tensorflow/core/kernels/sparse_cross_op.cc    | 52 +++++++++++-----
 .../sparse_ops/sparse_cross_op_test.py        | 60 +++++++++++++++++++
 .../python/ops/ragged/ragged_cross_op_test.py | 48 +++++++++++++++
 4 files changed, 171 insertions(+), 23 deletions(-)

diff --git a/tensorflow/core/kernels/ragged_cross_op.cc b/tensorflow/core/kernels/ragged_cross_op.cc
index bc59f38adba..f896bc74de2 100644
--- a/tensorflow/core/kernels/ragged_cross_op.cc
+++ b/tensorflow/core/kernels/ragged_cross_op.cc
@@ -190,7 +190,7 @@ class OutputWriterImpl : public OutputWriter {
   // Reads features from the specified slice of batch indices, computes
   // feature crosses for each one, and writes them to values_out_.
   void WriteOutputSlice(int64_t begin, int64_t end) override {
-    std::vector<int> combination(features_.size(), 0);
+    std::vector<int64_t> combination(features_.size(), 0);
     for (int64_t b = begin; b < end; ++b) {
       auto row_start = splits_out_(b);
       auto row_limit = splits_out_(b + 1);
@@ -198,7 +198,8 @@ class OutputWriterImpl : public OutputWriter {
         WriteCombination(b, combination, &values_out_(i));
         NextCombination(b, &combination);
       }
-      combination.assign(features_.size(), 0);  // reset for next batch.
+      combination.assign(features_.size(),
+                         int64_t{0});  // reset for next batch.
     }
   }
 
@@ -206,7 +207,7 @@ class OutputWriterImpl : public OutputWriter {
   // Joins the specified combination of input features into a single string,
   // and writes it to *out.
   void WriteCombination(int64_t batch_index,
-                        const std::vector<int>& combination, tstring* out) {
+                        const std::vector<int64_t>& combination, tstring* out) {
     static const auto k_feature_separator = "_X_";
     absl::InlinedVector<tstring, 6> cross_vec(features_.size());
     for (int i = 0; i < combination.size(); ++i) {
@@ -218,7 +219,7 @@ class OutputWriterImpl : public OutputWriter {
   // Joins the specified combination of input features into a single
   // fingerprint, and writes it to *out.
   void WriteCombination(int64_t batch_index,
-                        const std::vector<int>& combination, int64_t* out) {
+                        const std::vector<int64_t>& combination, int64_t* out) {
     // Do the fingerprint concatenation on uint64.
     uint64_t hashed_output = hash_key_;
     for (size_t i = 0; i < combination.size(); ++i) {
@@ -237,7 +238,7 @@ class OutputWriterImpl : public OutputWriter {
 
   // Updates `combination` to the next combination of input features.
   void NextCombination(int64_t batch_index,
-                       std::vector<int>* combination) const {
+                       std::vector<int64_t>* combination) const {
     bool carry = true;
     for (int i = combination->size() - 1; i >= 0; i--) {
       if (carry) {
@@ -609,10 +610,21 @@ class RaggedCrossOp : public OpKernel {
     auto flat_row_splits = (*row_splits_out)->flat<SplitsType>();
     int64_t cross_count_total = 0;
     flat_row_splits(0) = 0;
+    const int64_t kMaxCrossCountTotal =
+        std::min(static_cast<int64_t>(std::numeric_limits<SplitsType>::max()),
+                 std::numeric_limits<int64_t>::max() / 2);
     for (int64_t b = 0; b < batch_size; b++) {
       int64_t cross_count_by_batch_index = CrossCountByBatchIndex(features, b);
       if (cross_count_by_batch_index < 0) {
-        return absl::InvalidArgumentError("Invalid RaggedTensor");
+        return absl::InvalidArgumentError(
+            "RaggedCross: feature count is negative or the Cartesian product "
+            "of all feature counts exceeds the maximum supported value.");
+      }
+      if (cross_count_by_batch_index >
+          kMaxCrossCountTotal - cross_count_total) {
+        return absl::InvalidArgumentError(
+            "RaggedCross: the total number of crosses across all batches "
+            "exceeds the maximum supported value.");
       }
       cross_count_total += cross_count_by_batch_index;
       flat_row_splits(b + 1) = cross_count_total;
@@ -625,15 +637,19 @@ class RaggedCrossOp : public OpKernel {
     return absl::OkStatus();
   }
 
-  // Returns number of crosses for a given batch_index
+  // Returns number of crosses for a given batch_index, or -1 on overflow /
+  // invalid input.
   int64_t CrossCountByBatchIndex(const FeatureReaders& features,
-                                 int batch_index) {
+                                 int64_t batch_index) {
+    const int64_t kMaxCrossCount =
+        std::min(static_cast<int64_t>(std::numeric_limits<SplitsType>::max()),
+                 std::numeric_limits<int64_t>::max() / 2);
     int64_t cross_count = 1;
     for (int i = 0; i < features.size(); ++i) {
       const auto feature_count = features[i]->FeatureCount(batch_index);
-      // If feature_count is invalid, return -1 to let caller know.
       if (feature_count < 0) return -1;
       if (feature_count == 0) return 0;
+      if (cross_count > kMaxCrossCount / feature_count) return -1;
       cross_count *= feature_count;
     }
     return cross_count;
diff --git a/tensorflow/core/kernels/sparse_cross_op.cc b/tensorflow/core/kernels/sparse_cross_op.cc
index efa9b27d6b1..3d67021c977 100644
--- a/tensorflow/core/kernels/sparse_cross_op.cc
+++ b/tensorflow/core/kernels/sparse_cross_op.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include <vector>
 
 #include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
+#include "xla/tsl/platform/statusor.h"
 #include "tensorflow/core/framework/kernel_def_builder.h"
 #include "tensorflow/core/framework/op_def_builder.h"
 #include "tensorflow/core/framework/op_kernel.h"
@@ -311,7 +312,7 @@ class StringCrosser {
       : columns_(columns), k_feature_separator_(k_feature_separator) {}
 
   std::string Generate(const int64_t batch_index,
-                       const std::vector<int>& permutation,
+                       const std::vector<int64_t>& permutation,
                        bool unused_strong_hash) const {
     gtl::InlinedVector<InternalType, 6> cross_vec(columns_.size());
     for (int i = 0; i < permutation.size(); i++) {
@@ -337,7 +338,7 @@ class HashCrosser {
       : columns_(columns), num_buckets_(num_buckets), hash_key_(hash_key) {}
 
   int64_t Generate(const int64_t batch_index,
-                   const std::vector<int>& permutation,
+                   const std::vector<int64_t>& permutation,
                    bool unused_strong_hash) const {
     // Do the fingerprint concatenation on uint64.
     uint64_t hashed_output = hash_key_;
@@ -371,7 +372,7 @@ class HashCrosserV2 {
       : columns_(columns), num_buckets_(num_buckets) {}
 
   int64_t Generate(const int64_t batch_index,
-                   const std::vector<int>& permutation,
+                   const std::vector<int64_t>& permutation,
                    bool strong_hash) const {
     // Do the fingerprint concatenation on uint64.
     uint64_t hashed_output =
@@ -415,8 +416,8 @@ class ProductIterator {
     }
   }
 
-  std::vector<int> Next() {
-    std::vector<int> permutation(next_permutation_);
+  std::vector<int64_t> Next() {
+    std::vector<int64_t> permutation(next_permutation_);
 
     // Generates next permutation, if available.
     bool carry = true;
@@ -441,7 +442,7 @@ class ProductIterator {
   bool has_next_;
   const std::vector<std::unique_ptr<ColumnInterface<InternalType>>>& columns_;
   const int64_t batch_index_;
-  std::vector<int> next_permutation_;
+  std::vector<int64_t> next_permutation_;
 };
 
 template <bool HASHED_OUTPUT, typename InternalType>
@@ -602,7 +603,7 @@ void ExtractFeatureData(
     std::vector<std::vector<int64_t>>* feature_counts,
     std::vector<std::vector<int64_t>>* feature_start_indices) {
   absl::InlinedVector<int64_t, 8UL> current_row(indices_list_in.size(), 0);
-  for (int b = 0; b < batch_size; b++) {
+  for (int64_t b = 0; b < batch_size; b++) {
     for (int i = 0; i < indices_list_in.size(); i++) {
       const auto indices = indices_list_in[i].matrix<int64_t>();
       int64_t feature_count = 0;
@@ -619,11 +620,16 @@ void ExtractFeatureData(
   }
 }
 
-// Returns number of crosses for a given batch_index
+// Returns number of crosses for a given batch_index, or an error if the
+// Cartesian product size overflows int64.
 template <typename InternalType>
-int64_t CrossCountByBatchIndex(
+absl::StatusOr<int64_t> CrossCountByBatchIndex(
     const std::vector<std::unique_ptr<ColumnInterface<InternalType>>>& columns,
-    int batch_index) {
+    int64_t batch_index) {
+  // Cap at INT64_MAX/2: the indices output is TensorShape({total, 2}), so
+  // TensorShape internally computes total*2 — capping here prevents that
+  // multiplication from overflowing.
+  constexpr int64_t kMaxCrossCount = std::numeric_limits<int64_t>::max() / 2;
   int64_t cross_count = 1;
   for (int i = 0; i < columns.size(); i++) {
     const auto feature_count = columns[i]->FeatureCount(batch_index);
@@ -631,6 +637,15 @@ int64_t CrossCountByBatchIndex(
     if (feature_count == 0) {
       return 0;
     }
+    if (feature_count < 0) {
+      return absl::InvalidArgumentError(
+          "SparseCross: feature count cannot be negative.");
+    }
+    if (cross_count > kMaxCrossCount / feature_count) {
+      return absl::InvalidArgumentError(
+          "SparseCross: the Cartesian product of all feature counts exceeds "
+          "the maximum supported value.");
+    }
     cross_count *= feature_count;
   }
   return cross_count;
@@ -718,8 +733,17 @@ absl::Status CreateOutputTensors(
   for (int64_t b = 0; b < batch_size; b++) {
     // For each input, sets starting indices in output SparseTensor
     (*output_start_indices)[b] = cross_count_total;
-    const auto cross_count = CrossCountByBatchIndex(columns, b);
+    auto cross_count_or = CrossCountByBatchIndex(columns, b);
+    if (!cross_count_or.ok()) return cross_count_or.status();
+    const int64_t cross_count = *cross_count_or;
     max_cross_count = std::max(max_cross_count, cross_count);
+    constexpr int64_t kMaxCrossCountTotal =
+        std::numeric_limits<int64_t>::max() / 2;
+    if (cross_count > kMaxCrossCountTotal - cross_count_total) {
+      return absl::InvalidArgumentError(
+          "SparseCross: the total number of crosses across all batches exceeds "
+          "the maximum supported value.");
+    }
     cross_count_total += cross_count;
   }
 
@@ -788,7 +812,7 @@ class SparseCrossOp : public OpKernel {
     typename CrossTraits<HASHED_OUTPUT, InternalType>::Updater updater(
         output_start_indices, indices_out, values_out);
     auto do_work = [&columns, crosser, updater](int64_t begin, int64_t end) {
-      for (int b = begin; b < end; b++) {
+      for (int64_t b = begin; b < end; b++) {
         ProductIterator<InternalType> product_iterator(columns, b);
         int64_t cross_count = 0;
         while (product_iterator.HasNext()) {
@@ -860,7 +884,7 @@ class SparseCrossV2Op : public OpKernel {
     OutputUpdater<tstring> updater(output_start_indices, indices_out,
                                    values_out);
     auto do_work = [&columns, crosser, updater](int64_t begin, int64_t end) {
-      for (int b = begin; b < end; b++) {
+      for (int64_t b = begin; b < end; b++) {
         ProductIterator<tstring> product_iterator(columns, b);
         int64_t cross_count = 0;
         while (product_iterator.HasNext()) {
@@ -940,7 +964,7 @@ class SparseCrossHashedOp : public OpKernel {
                                    values_out);
     auto do_work = [&columns, crosser, updater, strong_hash](int64_t begin,
                                                              int64_t end) {
-      for (int b = begin; b < end; b++) {
+      for (int64_t b = begin; b < end; b++) {
         ProductIterator<int64_t> product_iterator(columns, b);
         int64_t cross_count = 0;
         while (product_iterator.HasNext()) {
diff --git a/tensorflow/python/kernel_tests/sparse_ops/sparse_cross_op_test.py b/tensorflow/python/kernel_tests/sparse_ops/sparse_cross_op_test.py
index 8458c806136..414364c5049 100644
--- a/tensorflow/python/kernel_tests/sparse_ops/sparse_cross_op_test.py
+++ b/tensorflow/python/kernel_tests/sparse_ops/sparse_cross_op_test.py
@@ -499,6 +499,66 @@ class SparseCrossOpTest(test.TestCase):
                                 'Expected batch size'):
       self.evaluate(sparse_ops.sparse_cross([st1, st2]))
 
+  def test_cross_count_overflow(self):
+    # 64 sparse columns x 2 features each: the Cartesian product (2^64)
+    # overflows int64 without the overflow guard.
+    col = sparse_tensor.SparseTensor(
+        indices=constant_op.constant([[0, 0], [0, 1]], dtypes.int64),
+        values=constant_op.constant(['a', 'b']),
+        dense_shape=constant_op.constant([1, 2], dtypes.int64),
+    )
+    with self.assertRaisesRegex(
+        (errors.InvalidArgumentError, ValueError),
+        r'SparseCross: the Cartesian product of all feature counts exceeds '
+        r'the maximum supported value\.',
+    ):
+      self.evaluate(sparse_ops.sparse_cross([col] * 64))
+
+  def test_cross_count_overflow_hashed(self):
+    col = sparse_tensor.SparseTensor(
+        indices=constant_op.constant([[0, 0], [0, 1]], dtypes.int64),
+        values=constant_op.constant([1, 2], dtypes.int64),
+        dense_shape=constant_op.constant([1, 2], dtypes.int64),
+    )
+    with self.assertRaisesRegex(
+        (errors.InvalidArgumentError, ValueError),
+        r'SparseCross: the Cartesian product of all feature counts exceeds '
+        r'the maximum supported value\.',
+    ):
+      self.evaluate(sparse_ops.sparse_cross_hashed([col] * 64))
+
+  def test_cross_count_total_overflow(self):
+    # 61 columns x 2 features x 2 batches: per-batch = 2^61 (within cap),
+    # total = 2*2^61 = 2^62 > INT64_MAX/2 — triggers accumulation guard.
+    col = sparse_tensor.SparseTensor(
+        indices=constant_op.constant(
+            [[0, 0], [0, 1], [1, 0], [1, 1]], dtypes.int64
+        ),
+        values=constant_op.constant(['a', 'b', 'c', 'd']),
+        dense_shape=constant_op.constant([2, 2], dtypes.int64),
+    )
+    with self.assertRaisesRegex(
+        (errors.InvalidArgumentError, ValueError),
+        r'SparseCross: the total number of crosses across all batches exceeds '
+        r'the maximum supported value\.',
+    ):
+      self.evaluate(sparse_ops.sparse_cross([col] * 61))
+
+  def test_cross_count_total_overflow_hashed(self):
+    col = sparse_tensor.SparseTensor(
+        indices=constant_op.constant(
+            [[0, 0], [0, 1], [1, 0], [1, 1]], dtypes.int64
+        ),
+        values=constant_op.constant([1, 2, 3, 4], dtypes.int64),
+        dense_shape=constant_op.constant([2, 2], dtypes.int64),
+    )
+    with self.assertRaisesRegex(
+        (errors.InvalidArgumentError, ValueError),
+        r'SparseCross: the total number of crosses across all batches exceeds '
+        r'the maximum supported value\.',
+    ):
+      self.evaluate(sparse_ops.sparse_cross_hashed([col] * 61))
+
 
 class SparseCrossV2OpTest(BaseSparseCrossOpTest):
 
diff --git a/tensorflow/python/ops/ragged/ragged_cross_op_test.py b/tensorflow/python/ops/ragged/ragged_cross_op_test.py
index e37c017df3c..2ca5d3c24af 100644
--- a/tensorflow/python/ops/ragged/ragged_cross_op_test.py
+++ b/tensorflow/python/ops/ragged/ragged_cross_op_test.py
@@ -529,6 +529,54 @@ class RaggedCrossOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
           )
       )
 
+  def test_cross_count_overflow(self):
+    # 64 ragged columns x 2 features: Cartesian product (2^64) overflows int64.
+    col = ragged_factory_ops.constant([['a', 'b']])
+    with self.assertRaisesRegex(
+        (errors.InvalidArgumentError, ValueError),
+        r'RaggedCross: feature count is negative or the Cartesian product '
+        r'of all feature counts exceeds the maximum supported value\.',
+    ):
+      self.evaluate(ragged_array_ops.cross([col] * 64))
+
+  def test_cross_count_overflow_int32_splits(self):
+    # 7 columns x 32 features = 32^7 = 2^35 > INT32_MAX, triggers SplitsType cap.
+    col = ragged_factory_ops.constant(
+        [['f'] * 32], row_splits_dtype=dtypes.int32
+    )
+    with self.assertRaisesRegex(
+        (errors.InvalidArgumentError, ValueError),
+        r'RaggedCross: feature count is negative or the Cartesian product '
+        r'of all feature counts exceeds the maximum supported value\.',
+    ):
+      self.evaluate(
+          gen_ragged_array_ops.RaggedCross(
+              ragged_values=[col.flat_values] * 7,
+              ragged_row_splits=[col.row_splits] * 7,
+              sparse_indices=[],
+              sparse_values=[],
+              sparse_shape=[],
+              dense_inputs=[],
+              input_order='R' * 7,
+              hashed_output=False,
+              num_buckets=0,
+              hash_key=0,
+              out_values_type=dtypes.string,
+              out_row_splits_type=dtypes.int32,
+          )
+      )
+
+  def test_cross_count_total_overflow(self):
+    # 61 cols x 2 features x 2 batches: per-batch = 2^61 <= INT64_MAX/2,
+    # total = 2 * 2^61 = 2^62 > INT64_MAX/2, triggers accumulation guard.
+    col = ragged_factory_ops.constant([['a', 'b'], ['c', 'd']])
+    with self.assertRaisesRegex(
+        (errors.InvalidArgumentError, ValueError),
+        r'RaggedCross: the total number of crosses across all batches '
+        r'exceeds the maximum supported value\.',
+    ):
+      self.evaluate(ragged_array_ops.cross([col] * 61))
+
 
 if __name__ == '__main__':
   googletest.main()

From 2dd36553aa623bfed4e52471b70d77fb4fead45b Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Thu, 8 Oct 2026 23:51:47 -0700
Subject: [PATCH 0329/1000] Add Triton patch for LLVM integrate at c4d9267a3189

The LLVM integrate at llvm/llvm-project@c4d9267a3189 updated
`mlir::populateGpuToROCDLConversionPatterns` to take `mlir::ROCDL::TargetInfo`
instead of the deprecated `mlir::amdgpu::Chipset`, but missed adding the
corresponding Triton patch file, breaking the ROCm build.

PiperOrigin-RevId: 996326975
---
 .../triton/common/llvm_cl996124567.patch      | 91 +++++++++++++++++++
 .../xla/third_party/triton/common/series      |  1 +
 2 files changed, 92 insertions(+)
 create mode 100644 third_party/xla/third_party/triton/common/llvm_cl996124567.patch

diff --git a/third_party/xla/third_party/triton/common/llvm_cl996124567.patch b/third_party/xla/third_party/triton/common/llvm_cl996124567.patch
new file mode 100644
index 00000000000..05742f8bcc4
--- /dev/null
+++ b/third_party/xla/third_party/triton/common/llvm_cl996124567.patch
@@ -0,0 +1,91 @@
+--- a/test/Conversion/tritongpu_to_ptx_mmav3.mlir
++++ b/test/Conversion/tritongpu_to_ptx_mmav3.mlir
+@@ -63,7 +63,8 @@
+     // CHECK-DAG: and.b32 [[LANEID_FIXED_BITS:%.*]], [[TID]], 28;
+     // CHECK-DAG: shl.b32 [[L0_TEMP:%.*]], [[L0_VAL]], 1;
+     // CHECK-DAG: or.b32  [[LANEID_PART_PERM:%.*]], [[L0_TEMP]], [[LANEID_FIXED_BITS]];
+-    // CHECK-DAG: bfe.u32 [[L1_TEMP:%.*]], [[TID]], 1, 1;
++    // CHECK-DAG: and.b32 [[L1_VAL:%.*]], [[TID]], 2;
++    // CHECK-DAG: shr.u32 [[L1_TEMP:%.*]], [[L1_VAL]], 1;
+     // CHECK-DAG: or.b32  [[LANEID_PERM:%.*]], [[LANEID_PART_PERM]], [[L1_TEMP]];
+ 
+     // The index shuffles have source lane dependent on the value of the r2 bit.
+@@ -97,7 +98,6 @@
+     // on, we have selectors 0x7632 and 0x3276. These are 21520, 4180, 30258,
+     // and 12918 in decimal, respectively.
+ 
+-    // CHECK-DAG: and.b32           [[L1_VAL:%.*]], [[TID]], 2;
+     // CHECK-DAG: setp.eq.b32       [[L1_OFF:%.*]], [[L1_VAL]], 0;
+     // CHECK:     selp.b32          [[SEL1:%.*]], 21520, 4180, [[L1_OFF]];
+     // CHECK:     selp.b32          [[SEL2:%.*]], 30258, 12918, [[L1_OFF]];
+
+--- a/third_party/amd/lib/TritonAMDGPUToLLVM/TritonGPUToLLVM.cpp
++++ b/third_party/amd/lib/TritonAMDGPUToLLVM/TritonGPUToLLVM.cpp
+@@ -12,10 +12,10 @@
+ #include "mlir/Conversion/MathToLLVM/MathToLLVM.h"
+ #include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
+ #include "mlir/Conversion/UBToLLVM/UBToLLVM.h"
+-#include "mlir/Dialect/AMDGPU/Utils/Chipset.h"
+ #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+ #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
+ #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
++#include "mlir/Dialect/LLVMIR/ROCDLTargetInfo.h"
+ #include "mlir/Pass/Pass.h"
+ #include "third_party/amd/include/Analysis/AMDGPUAllocation.h"
+ #include "third_party/amd/include/Analysis/AxisInfoExt.h"
+@@ -222,16 +222,16 @@
+     mlir::triton::AMD::populateWarpIdOpToLLVMPattern(typeConverter, targetInfo,
+                                                      patterns, commonBenefit);
+ 
+-    FailureOr<mlir::amdgpu::Chipset> maybeChipset =
+-        mlir::amdgpu::Chipset::parse(this->gfxArch);
+-    if (failed(maybeChipset)) {
++    FailureOr<mlir::ROCDL::TargetInfo> maybeTargetInfo =
++        mlir::ROCDL::TargetInfo::get(this->gfxArch);
++    if (failed(maybeTargetInfo)) {
+       emitError(UnknownLoc::get(&getContext()),
+                 "Invalid AMDGPU chipset name: " + this->gfxArch);
+       return signalPassFailure();
+     }
+     // Native lowering patterns
+     mlir::populateGpuToROCDLConversionPatterns(
+-        typeConverter, patterns, mlir::gpu::amd::HIP, *maybeChipset);
++        typeConverter, patterns, mlir::gpu::amd::HIP, *maybeTargetInfo);
+ 
+     mlir::cf::populateControlFlowToLLVMConversionPatterns(typeConverter,
+                                                           patterns);
+
+--- a/third_party/proton/Dialect/lib/ProtonGPUToLLVM/ProtonAMDGPUToLLVM/ConvertProtonGPUToLLVM.cpp
++++ b/third_party/proton/Dialect/lib/ProtonGPUToLLVM/ProtonAMDGPUToLLVM/ConvertProtonGPUToLLVM.cpp
+@@ -7,9 +7,9 @@
+ #include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
+ #include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
+ #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h"
+-#include "mlir/Dialect/AMDGPU/Utils/Chipset.h"
+ #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+ #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
++#include "mlir/Dialect/LLVMIR/ROCDLTargetInfo.h"
+ #include "mlir/Pass/Pass.h"
+ #include "third_party/amd/lib/TritonAMDGPUToLLVM/PatternTritonGPUOpToLLVM.h"
+ #include "triton/Conversion/TritonGPUToLLVM/TypeConverter.h"
+@@ -65,15 +65,15 @@
+                                                        tritonTargetInfo);
+     mlir::arith::populateArithToLLVMConversionPatterns(typeConverter, patterns);
+ 
+-    FailureOr<mlir::amdgpu::Chipset> maybeChipset =
+-        mlir::amdgpu::Chipset::parse(this->gfxArch);
+-    if (failed(maybeChipset)) {
++    FailureOr<mlir::ROCDL::TargetInfo> maybeTargetInfo =
++        mlir::ROCDL::TargetInfo::get(this->gfxArch);
++    if (failed(maybeTargetInfo)) {
+       emitError(UnknownLoc::get(&getContext()),
+                 "Invalid AMDGPU chipset name: " + this->gfxArch);
+       return signalPassFailure();
+     }
+     mlir::populateGpuToROCDLConversionPatterns(
+-        typeConverter, patterns, mlir::gpu::amd::HIP, *maybeChipset);
++        typeConverter, patterns, mlir::gpu::amd::HIP, *maybeTargetInfo);
+     mlir::cf::populateControlFlowToLLVMConversionPatterns(typeConverter,
+                                                           patterns);
+     auto convTarget = ProtonLLVMConversionTarget(*context);
+
diff --git a/third_party/xla/third_party/triton/common/series b/third_party/xla/third_party/triton/common/series
index 1916dd0b59a..037a2d86211 100644
--- a/third_party/xla/third_party/triton/common/series
+++ b/third_party/xla/third_party/triton/common/series
@@ -37,3 +37,4 @@ useStrictPropertiesInAssemblyFormat.patch
 llvm_cl991819890.patch
 fix_load_ops.patch
 fp64_large_k.patch
+llvm_cl996124567.patch

From bd6168520ac76348c19af1ff02481ccb80d919da Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Fri, 9 Oct 2026 00:32:02 -0700
Subject: [PATCH 0330/1000] Clean up unused GPU benchmark BuildTypes and fix L4
 compute capability

Four GPU benchmark BuildTypes in build.py are unused after the benchmark
workflow switched to the BENCHMARK_PRESUBMIT_GITHUB_ACTIONS variants (with
cuda_libraries_from_stubs) and standardized on 16-vCPU L4 runners:
- XLA_LINUX_X86_GPU_L4_16_VCPU_PRESUBMIT_GITHUB_ACTIONS
- XLA_LINUX_X86_GPU_L4_48_VCPU_PRESUBMIT_GITHUB_ACTIONS
- XLA_LINUX_X86_GPU_L4_48_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS
- XLA_LINUX_X86_GPU_A4_224_VCPU_PRESUBMIT_GITHUB_ACTIONS

Remove these unused BuildTypes and the stale 48-vCPU comment in
build_binaries.sh. Also update the active L4 benchmark build
(XLA_LINUX_X86_GPU_L4_16_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS) from compute
capability 7.5 (Turing/T4) to 8.9 (Ada/L4) so AOT-compiled CUDA kernels in
hlo_runner_main_gpu target the actual L4 hardware rather than relying on PTX
JIT compilation at startup.

PiperOrigin-RevId: 996346172
---
 .../workflows/benchmarks/build_binaries.sh    |  2 +-
 third_party/xla/build_tools/ci/build.py       | 94 +------------------
 .../xla/build_tools/ci/golden_commands.txt    | 24 +----
 3 files changed, 5 insertions(+), 115 deletions(-)

diff --git a/third_party/xla/.github/workflows/benchmarks/build_binaries.sh b/third_party/xla/.github/workflows/benchmarks/build_binaries.sh
index e54ff01be85..5d3a9f88a9b 100755
--- a/third_party/xla/.github/workflows/benchmarks/build_binaries.sh
+++ b/third_party/xla/.github/workflows/benchmarks/build_binaries.sh
@@ -85,7 +85,7 @@ case "$HARDWARE_CATEGORY" in
     device_type_flag_value="host"
     ;;
   GPU_L4)
-    BUILD_TYPE="XLA_LINUX_X86_GPU_L4_16_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS" # Or _48_VCPU if that's the more common
+    BUILD_TYPE="XLA_LINUX_X86_GPU_L4_16_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS"
     runner_binary_path="./$BAZEL_BIN_DIR/xla/tools/multihost_hlo_runner/hlo_runner_main_gpu"
     stats_binary_path="./$BAZEL_BIN_DIR/xla/tools/compute_xspace_stats_main_gpu"
     device_type_flag_value="gpu"
diff --git a/third_party/xla/build_tools/ci/build.py b/third_party/xla/build_tools/ci/build.py
index 33905e53a3c..dc4310e2b1e 100755
--- a/third_party/xla/build_tools/ci/build.py
+++ b/third_party/xla/build_tools/ci/build.py
@@ -118,11 +118,7 @@ class BuildType(enum.Enum):
   # Presubmit builds for regression testing.
   XLA_LINUX_ARM64_CPU_48_VCPU_PRESUBMIT_GITHUB_ACTIONS = enum.auto()
   XLA_LINUX_X86_CPU_128_VCPU_PRESUBMIT_GITHUB_ACTIONS = enum.auto()
-  XLA_LINUX_X86_GPU_L4_16_VCPU_PRESUBMIT_GITHUB_ACTIONS = enum.auto()
-  XLA_LINUX_X86_GPU_L4_48_VCPU_PRESUBMIT_GITHUB_ACTIONS = enum.auto()
-  XLA_LINUX_X86_GPU_A4_224_VCPU_PRESUBMIT_GITHUB_ACTIONS = enum.auto()
   XLA_LINUX_X86_GPU_L4_16_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS = enum.auto()
-  XLA_LINUX_X86_GPU_L4_48_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS = enum.auto()
   XLA_LINUX_X86_GPU_A4_224_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS = enum.auto()
 
   XLA_MACOS_X86_CPU_KOKORO = enum.auto()
@@ -667,26 +663,6 @@ Build(
     subcommand="build",
 )
 
-Build(
-    type_=BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_PRESUBMIT_GITHUB_ACTIONS,
-    repo="openxla/xla",
-    target_patterns=_XLA_GPU_PRESUBMIT_BENCHMARKS_DEFAULT_TARGET_PATTERNS,
-    configs=("warnings", "rbe_linux_cuda_nvcc", "hermetic_cuda_umd"),
-    test_tag_filters=nvidia_single_gpu_test_filters
-    + _tag_filters_for_compute_capability(compute_capability=75),
-    build_tag_filters=nvidia_single_gpu_build_filters,
-    options={
-        "run_under": "//build_tools/ci:parallel_gpu_execute",
-        "//xla/tsl:ci_build": True,
-        **_DEFAULT_BAZEL_OPTIONS,
-    },
-    repo_env={
-        "TF_CUDA_COMPUTE_CAPABILITIES": "7.5",
-    },
-    extra_setup_commands=(["nvidia-smi"],),
-    subcommand="build",
-)
-
 Build(
     type_=BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS,
     repo="openxla/xla",
@@ -698,7 +674,7 @@ Build(
         "cuda_libraries_from_stubs",
     ),
     test_tag_filters=nvidia_single_gpu_test_filters
-    + _tag_filters_for_compute_capability(compute_capability=75),
+    + _tag_filters_for_compute_capability(compute_capability=89),
     build_tag_filters=nvidia_single_gpu_build_filters,
     options={
         "run_under": "//build_tools/ci:parallel_gpu_execute",
@@ -706,73 +682,7 @@ Build(
         **_DEFAULT_BAZEL_OPTIONS,
     },
     repo_env={
-        "TF_CUDA_COMPUTE_CAPABILITIES": "7.5",
-    },
-    extra_setup_commands=(["nvidia-smi"],),
-    subcommand="build",
-)
-
-Build(
-    type_=BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_PRESUBMIT_GITHUB_ACTIONS,
-    repo="openxla/xla",
-    configs=("warnings", "rbe_linux_cuda_nvcc", "hermetic_cuda_umd"),
-    target_patterns=_XLA_GPU_PRESUBMIT_BENCHMARKS_DEFAULT_TARGET_PATTERNS,
-    test_tag_filters=nvidia_single_gpu_test_filters
-    + _tag_filters_for_compute_capability(compute_capability=75),
-    build_tag_filters=nvidia_single_gpu_build_filters,
-    options={
-        "run_under": "//build_tools/ci:parallel_gpu_execute",
-        "//xla/tsl:ci_build": True,
-        **_DEFAULT_BAZEL_OPTIONS,
-    },
-    repo_env={
-        "TF_CUDA_COMPUTE_CAPABILITIES": "7.5",
-    },
-    extra_setup_commands=(["nvidia-smi"],),
-    subcommand="build",
-)
-
-Build(
-    type_=BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS,
-    repo="openxla/xla",
-    configs=(
-        "warnings",
-        "rbe_linux_cuda_nvcc",
-        "hermetic_cuda_umd",
-        "cuda_libraries_from_stubs",
-    ),
-    target_patterns=_XLA_GPU_PRESUBMIT_BENCHMARKS_DEFAULT_TARGET_PATTERNS,
-    test_tag_filters=nvidia_single_gpu_test_filters
-    + _tag_filters_for_compute_capability(compute_capability=75),
-    build_tag_filters=nvidia_single_gpu_build_filters,
-    options={
-        "run_under": "//build_tools/ci:parallel_gpu_execute",
-        "//xla/tsl:ci_build": True,
-        **_DEFAULT_BAZEL_OPTIONS,
-    },
-    repo_env={
-        "TF_CUDA_COMPUTE_CAPABILITIES": "7.5",
-    },
-    extra_setup_commands=(["nvidia-smi"],),
-    subcommand="build",
-)
-
-Build(
-    type_=BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_PRESUBMIT_GITHUB_ACTIONS,
-    repo="openxla/xla",
-    configs=(),
-    target_patterns=_XLA_GPU_PRESUBMIT_BENCHMARKS_DEFAULT_TARGET_PATTERNS,
-    test_tag_filters=nvidia_single_gpu_test_filters
-    + _tag_filters_for_compute_capability(compute_capability=100),
-    build_tag_filters=nvidia_single_gpu_build_filters,
-    options={
-        "run_under": "//build_tools/ci:parallel_gpu_execute",
-        # Use User Mode and Kernel Mode Drivers pre-installed on the system.
-        "//xla/tsl:ci_build": True,
-        **_DEFAULT_BAZEL_OPTIONS,
-    },
-    repo_env={
-        "TF_CUDA_COMPUTE_CAPABILITIES": "10",
+        "TF_CUDA_COMPUTE_CAPABILITIES": "8.9",
     },
     extra_setup_commands=(["nvidia-smi"],),
     subcommand="build",
diff --git a/third_party/xla/build_tools/ci/golden_commands.txt b/third_party/xla/build_tools/ci/golden_commands.txt
index c228399b9c0..ec31f8d099c 100644
--- a/third_party/xla/build_tools/ci/golden_commands.txt
+++ b/third_party/xla/build_tools/ci/golden_commands.txt
@@ -56,35 +56,15 @@ nvidia-smi
 parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
 bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
 # END BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS
-# BEGIN BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_PRESUBMIT_GITHUB_ACTIONS
-nvidia-smi
-parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
-bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
-# END BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_PRESUBMIT_GITHUB_ACTIONS
 # BEGIN BuildType.XLA_LINUX_X86_GPU_HERMETIC_ROCM_GITHUB_ACTIONS
 parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_gpu,-requires-gpu-intel,-requires-gpu-nvidia,-cuda-only,-oneapi-only,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm86,-requires-gpu-sm86-only,-requires-gpu-sm89,-requires-gpu-sm89-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial,gpu --test_tag_filters=-no_gpu,-requires-gpu-intel,-requires-gpu-nvidia,-cuda-only,-oneapi-only,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm86,-requires-gpu-sm86-only,-requires-gpu-sm89,-requires-gpu-sm89-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial,gpu --config=warnings --config=rbe_linux_cpu --config=rocm_clang_hermetic --config=rocm_ci_hermetic --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/...
 bazel build --build_tag_filters=-no_gpu,-requires-gpu-intel,-requires-gpu-nvidia,-cuda-only,-oneapi-only,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm86,-requires-gpu-sm86-only,-requires-gpu-sm89,-requires-gpu-sm89-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial,gpu --test_tag_filters=-no_gpu,-requires-gpu-intel,-requires-gpu-nvidia,-cuda-only,-oneapi-only,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm86,-requires-gpu-sm86-only,-requires-gpu-sm89,-requires-gpu-sm89-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial,gpu --config=warnings --config=rbe_linux_cpu --config=rocm_clang_hermetic --config=rocm_ci_hermetic --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/...
 # END BuildType.XLA_LINUX_X86_GPU_HERMETIC_ROCM_GITHUB_ACTIONS
 # BEGIN BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS
 nvidia-smi
-parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
-bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm89-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=8.9 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
+bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm89-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=8.9 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
 # END BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS
-# BEGIN BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_PRESUBMIT_GITHUB_ACTIONS
-nvidia-smi
-parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
-bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
-# END BuildType.XLA_LINUX_X86_GPU_L4_16_VCPU_PRESUBMIT_GITHUB_ACTIONS
-# BEGIN BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS
-nvidia-smi
-parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
-bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
-# END BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS
-# BEGIN BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_PRESUBMIT_GITHUB_ACTIONS
-nvidia-smi
-parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
-bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm75-only,requires-gpu-sm60,requires-gpu-sm70,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-requires-gpu-sm100,-requires-gpu-sm100-only,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=7.5 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
-# END BuildType.XLA_LINUX_X86_GPU_L4_48_VCPU_PRESUBMIT_GITHUB_ACTIONS
 # BEGIN BuildType.XLA_LINUX_X86_GPU_ONEAPI_GITHUB_ACTIONS
 parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=oneapi-only,requires-gpu-intel,-requires-gpu-amd,-requires-gpu-nvidia,-no_oss,-cuda-only,-rocm-only,-no-oneapi,gpu --test_tag_filters=oneapi-only,-requires-gpu-intel,-requires-gpu-amd,-requires-gpu-nvidia,-no_oss,-cuda-only,-rocm-only,-no-oneapi,gpu --config=nonccl --config=rbe_linux_cpu --config=sycl --config=sycl_hermetic --config=icpx_clang --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/...
 bazel build --build_tag_filters=oneapi-only,requires-gpu-intel,-requires-gpu-amd,-requires-gpu-nvidia,-no_oss,-cuda-only,-rocm-only,-no-oneapi,gpu --test_tag_filters=oneapi-only,-requires-gpu-intel,-requires-gpu-amd,-requires-gpu-nvidia,-no_oss,-cuda-only,-rocm-only,-no-oneapi,gpu --config=nonccl --config=rbe_linux_cpu --config=sycl --config=sycl_hermetic --config=icpx_clang --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/...

From 9d22951c0917143386ed94b67d185142e899085a Mon Sep 17 00:00:00 2001
From: Ashutosh0x <ashutoshkumarsingh0x@gmail.com>
Date: Fri, 9 Oct 2026 13:17:24 +0530
Subject: [PATCH 0331/1000] Hide the input shape from graph-mode shape
 inference in rnn_grad_test

In graph mode the function call is inlined and shape inference sees the
argument's static shape, so the tests raised ValueError instead of
reaching the kernel. Feed the value through placeholder_with_default
with an unknown shape there.
---
 tensorflow/python/ops/rnn_grad_test.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensorflow/python/ops/rnn_grad_test.py b/tensorflow/python/ops/rnn_grad_test.py
index 904c9221984..70c92f6350e 100644
--- a/tensorflow/python/ops/rnn_grad_test.py
+++ b/tensorflow/python/ops/rnn_grad_test.py
@@ -18,6 +18,7 @@ import functools
 
 import numpy as np
 
+from tensorflow.python.eager import context
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import errors_impl
@@ -211,6 +212,10 @@ class RNNGradTest(test.TestCase):
     while the graph is built, so the check under test is the kernel's.
     """
     value = kwargs.pop(name)
+    if not context.executing_eagerly():
+      # In graph mode the function call is inlined and shape inference sees
+      # the argument's static shape, so hide it behind a placeholder as well.
+      value = array_ops.placeholder_with_default(value, shape=None)
 
     @def_function.function(
         autograph=False,

From 28854b321cc8defa3485ec73505943ca0c86eb07 Mon Sep 17 00:00:00 2001
From: Mikhail Goncharov <goncharov@google.com>
Date: Fri, 9 Oct 2026 00:33:03 -0700
Subject: [PATCH 0332/1000] [XLA:GPU] Remove legacy symbolic tile analysis runs
 from GPU tests.

Strip `--xla_gpu_experimental_enable_tiling_propagation=false` `RUN:` lines and `CHECK-OLD` directives from Triton HLO lit tests, while keeping negative checks for multi-output fusions (b/502910372) and negative base-offset concatenates (b/528277610) that new tiling does not support yet.

De-parameterize GPU cost model, fusion transform, autotuner, and Triton codegen C++ unit test suites so they run exclusively with default tiling propagation, and remove legacy `SymbolicTileAnalysis`-only test cases and build dependencies. Halve `support_test` shard count from 50 to 25 and remove its `long` timeout now that each test case runs once instead of twice.

PiperOrigin-RevId: 996346720
---
 .../xla/backends/gpu/autotuner/triton_test.cc |  78 +-
 .../xla/xla/backends/gpu/codegen/triton/BUILD |   3 +-
 .../gpu/codegen/triton/support_test.cc        | 911 +++++++-----------
 .../backends/gpu/codegen/triton/tests/BUILD   |   6 +
 .../tests/bitcast/bitcast_dynamic_slice.hlo   |   7 +-
 .../tests/bitcast/bitcast_expand_tiled_1.hlo  |   4 +-
 .../tests/bitcast/bitcast_expand_tiled_4.hlo  |   4 +-
 .../bitcast/bitcast_expand_transposed.hlo     |   4 +-
 .../tests/bitcast/bitcast_non_norm_in.hlo     |   6 +-
 .../tests/bitcast/bitcast_non_norm_in_out.hlo |   6 +-
 .../tests/bitcast/bitcast_non_norm_out.hlo    |   6 +-
 .../tests/bitcast/bitcast_normalized.hlo      |  20 +-
 .../tests/bitcast/bitcast_reduce_stride1.hlo  |   8 +-
 .../bitcast/bitcast_reshape_collapse.hlo      |  10 +-
 .../bitcast_reshape_expand_grouped.hlo        |   4 +-
 .../bitcast_reshape_expand_tiled_1.hlo        |   4 +-
 .../bitcast_reshape_expand_tiled_2.hlo        |   4 +-
 ...itcast_reshape_non_contiguous_collapse.hlo |   7 +-
 .../bitcast_reshape_non_contiguous_expand.hlo |  10 +-
 .../tests/bitcast/bitcast_reshape_only.hlo    |  17 +-
 .../triton/tests/bitcast/bitcast_to_1d.hlo    |   9 +-
 .../tests/bitcast/bitcast_to_1d_tiled.hlo     |   6 +-
 .../tests/bitcast/bitcast_to_scalar.hlo       |   9 +-
 .../tests/bitcast/bitcast_transpose_only.hlo  |  14 +-
 .../bitcast/bitcast_transpose_or_reshape.hlo  |   7 +-
 ...itcast_transpose_permutation_inversion.hlo |   7 +-
 .../bitcast/bitcast_unsupported_expand.hlo    |   2 +-
 .../triton/tests/broadcast/broadcast.hlo      |   7 +-
 .../tests/broadcast/scalar_broadcast.hlo      |   7 +-
 .../triton/tests/collectives/all_gather.hlo   |   1 -
 .../triton/tests/collectives/all_reduce.hlo   |   3 +-
 .../concat/concat_active_scalar_constant.hlo  |   3 +-
 .../triton/tests/concat/concat_fusion.hlo     |  25 +-
 ...ncat_offset_is_a_multiple_of_tile_size.hlo |   2 +-
 ...cat_offset_not_a_multiple_of_tile_size.hlo |   2 +-
 .../triton/tests/concat/concat_pred.hlo       |  25 +-
 .../tests/concat/concat_prune_inactive.hlo    |   6 +-
 .../tests/concat/concat_simplify_offset.hlo   |   4 +-
 .../concat/concat_slice_inactive_operand.hlo  |   4 +-
 .../max_tensor_elements_violation.hlo         |   6 +-
 .../triton/tests/dot/bf16_small_rhs_outer.hlo |   3 +-
 .../gpu/codegen/triton/tests/dot/dot.hlo      |   7 +-
 .../dot_broadcast_sequential_dim_mismatch.hlo |   7 +-
 .../codegen/triton/tests/dot/dot_concat.hlo   |   3 +-
 .../tests/dot/dot_major_lhs_contract.hlo      |   3 +-
 .../tests/dot/dot_minor_rhs_contract.hlo      |   3 +-
 .../triton/tests/dot/dot_transpose_tma.hlo    |   3 +-
 .../triton/tests/dot/dot_unit_tile_sizes.hlo  |   3 +-
 .../tests/dot/duplicate_dot_operand.hlo       |   2 +-
 .../triton/tests/dot/layout_021_201_210.hlo   |   6 +-
 .../triton/tests/dot/layout_201_012_201.hlo   |   6 +-
 .../triton/tests/dot/layout_210_210_102.hlo   |   6 +-
 .../triton/tests/dot/layout_210_210_210.hlo   |   6 +-
 .../codegen/triton/tests/dot/masked_dot.hlo   |   4 +-
 .../triton/tests/dot/reshape_of_pad.hlo       |   3 +-
 .../triton/tests/dot/single_tile_dot.hlo      |   4 +-
 ...use_major_to_minor_schedule_lhs_larger.hlo |   7 +-
 ...or_to_minor_schedule_not_rooted_in_dot.hlo |   7 +-
 .../use_transposed_schedule_lhs_smaller.hlo   |   4 +-
 .../dyn_slice_multiple_dimensions.hlo         |   7 +-
 .../tests/dynamic_slice/dyn_slice_of_add.hlo  |   7 +-
 .../dynamic_slice/dyn_slice_of_dyn_slice.hlo  |  60 +-
 .../dyn_slice_of_dyn_slice_no_convert.hlo     |  59 +-
 .../dyn_slice_offset_data_same_input.hlo      |   7 +-
 .../tests/elementwise/chain_0d_scalars.hlo    |   7 +-
 .../tests/elementwise/convert_int_to_pred.hlo |   7 +-
 .../triton/tests/elementwise/div_by_zero.hlo  |   7 +-
 .../elementwise/div_const_denom_by_zero.hlo   |   7 +-
 .../tests/elementwise/div_consts_by_zero.hlo  |   7 +-
 .../triton/tests/elementwise/div_overflow.hlo |   7 +-
 .../triton/tests/elementwise/minimum.hlo      |   4 +-
 .../minimum_three_tiles_per_block.hlo         |   5 +-
 .../minimum_two_tiles_per_block.hlo           |   5 +-
 .../elementwise/nested_reducer_fusion.hlo     |   3 +-
 .../codegen/triton/tests/elementwise/not.hlo  |   7 +-
 .../triton/tests/elementwise/pred_add.hlo     |   7 +-
 .../triton/tests/elementwise/pred_input.hlo   |   7 +-
 .../triton/tests/elementwise/pred_output.hlo  |   7 +-
 .../tests/elementwise/reduce_precision.hlo    |   7 +-
 .../elementwise/reduce_precision_noop.hlo     |   7 +-
 .../tests/fusion_emitter_device_test.cc       | 282 ++----
 .../tests/fusion_emitter_deviceless_test.cc   |  15 +-
 .../tests/fusion_emitter_int4_device_test.cc  |  62 +-
 .../fusion_emitter_shared_dialect_test.cc     |  78 +-
 .../bitcast_reduce_broadcast.hlo              |   6 +-
 .../normalization/diamond_broadcast_1d.hlo    |   3 +-
 .../normalization/diamond_broadcast_batch.hlo |   3 +-
 .../diamond_broadcast_reduce.hlo              |   3 +-
 .../normalization/diamond_splat_scalar.hlo    |  12 +-
 .../generic_emitter_multi_tiled_dims.hlo      |  12 +-
 .../codegen/triton/tests/pad/high_pad_1d.hlo  |   7 +-
 .../codegen/triton/tests/pad/high_pad_2d.hlo  |   7 +-
 .../pad/high_pad_of_broadcast_and_reshape.hlo |   4 +-
 .../high_pad_tile_size_equals_input_size.hlo  |   4 +-
 .../reduce/intra_warp_reduce_of_reduce.hlo    |   6 +-
 .../tests/reduce/multi_0d_broadcast.hlo       |   6 +-
 .../triton/tests/reduce/reduce_complex.hlo    |   6 +-
 .../tests/reduce/reduce_concat_diamond.hlo    |   4 +-
 .../tests/reduce/reduce_diamond_tiled.hlo     |   4 +-
 .../tests/reduce/reduce_intermediate.hlo      |   6 +-
 .../tests/reduce/reduce_major_subtile.hlo     |   6 +-
 .../triton/tests/reduce/reduce_minor.hlo      |   6 +-
 .../tests/reduce/reduce_minor_extra_out.hlo   |   8 +-
 .../tests/reduce/reduce_minor_subtile.hlo     |   8 +-
 .../triton/tests/reduce/reduce_multi_dim.hlo  |   4 +-
 .../reduce_multi_dim_tiled_and_untiled.hlo    |   4 +-
 .../tests/reduce/reduce_non_neutral_init.hlo  |  19 +-
 .../triton/tests/reduce/reduce_of_concat.hlo  |   6 +-
 .../triton/tests/reduce/reduce_tiled.hlo      |   4 +-
 .../tests/reduce/reduce_tiled_masked.hlo      |   4 +-
 .../reduce/reduce_to_scalar_extra_out.hlo     |   8 +-
 .../reshape/reshape_expand_high_rank.hlo      |   6 +-
 .../triton/tests/reshape/reshape_to_0d.hlo    |   6 +-
 .../triton/tests/scaled_dot/scaled_dot.hlo    |   7 +-
 .../scaled_dot_online_quantization.hlo        |   5 +-
 ...scaled_dot_online_quantization_bitcast.hlo |   5 +-
 .../triton/tests/scaled_dot_device_test.cc    | 141 +--
 .../tests/slice/boundary_dynamic_slice.hlo    |   7 +-
 .../tests/slice/slice_extra_out_padding.hlo   |   5 +-
 .../triton/tests/slice/slice_padding.hlo      |   5 +-
 .../triton/tests/slice/strided_iota_4d.hlo    |   7 +-
 .../triton/tests/transpose/transpose_3d.hlo   |   7 +-
 .../transpose_3d_tile_not_power_of_2.hlo      |   3 +-
 .../tests/all_gather_e2e_multiprocess_test.cc |   1 -
 .../convert_triton_gemm_config_test.cc        |  24 +-
 .../fusion_block_level_rewriter_test.cc       |  40 +-
 .../gpu/transforms/gemm_fusion_test.cc        | 184 ++--
 .../gpu/transforms/priority_fusion_test.cc    | 143 ++-
 .../softmax_rewriter_triton_test.cc           | 127 ++-
 third_party/xla/xla/codegen/xtile/BUILD       |   2 -
 .../tiling_from_block_parameters_test.cc      | 152 ---
 third_party/xla/xla/service/gpu/model/BUILD   |   7 -
 .../gpu/model/coalescing_analysis_test.cc     |  77 +-
 .../combined_gpu_performance_model_test.cc    |  25 +-
 .../gpu_cost_model_stats_collection_test.cc   |  29 +-
 .../gpu_indexing_performance_model_test.cc    | 211 ++--
 136 files changed, 1111 insertions(+), 2318 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/autotuner/triton_test.cc b/third_party/xla/xla/backends/gpu/autotuner/triton_test.cc
index ea21e45e701..5a03c976d14 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/triton_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/triton_test.cc
@@ -128,8 +128,7 @@ ENTRY %entry (lhs: bf16[4,4], rhs: bf16[4,4], lhs_scale: bf16[1,1], rhs_scale: b
   ROOT %fusion = bf16[4,4]{1,0} fusion(%lhs, %rhs, %lhs_scale, %rhs_scale), kind=kCustom, calls=%fusion_dot, metadata={op_name="foo"}, backend_config={"operation_queue_id":"0","fusion_backend_config":{"kind":"__triton_gemm"},"force_earliest_schedule":false,"reification_cost":[],"device_type":"DEVICE_TYPE_INVALID"}
 })hlo";
 
-class TritonBackendTest : public HloHardwareIndependentTestBase,
-                          public ::testing::WithParamInterface<bool> {
+class TritonBackendTest : public HloHardwareIndependentTestBase {
  protected:
   TritonBackendTest()
       : platform_(PlatformUtil::GetDefaultPlatform().value()),
@@ -138,18 +137,7 @@ class TritonBackendTest : public HloHardwareIndependentTestBase,
         alias_info_(stream_executor_->GetDeviceDescription()),
         compiler_(Compiler::GetForPlatform(platform_->id()).value()),
         backend_(&debug_options_, compiler_.get(), &target_config_,
-                 &alias_info_, &mlir_context_pool_) {
-    debug_options_.set_xla_gpu_experimental_enable_tiling_propagation(
-        GetParam());
-  }
-
-  DebugOptions GetDebugOptionsForTest() const override {
-    DebugOptions debug_options =
-        HloHardwareIndependentTestBase::GetDebugOptionsForTest();
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
-        GetParam());
-    return debug_options;
-  }
+                 &alias_info_, &mlir_context_pool_) {}
 
   DebugOptions debug_options_;
   se::Platform* platform_;
@@ -161,7 +149,7 @@ class TritonBackendTest : public HloHardwareIndependentTestBase,
   TritonBackend backend_;
 };
 
-TEST_P(TritonBackendTest, GetSupportedConfigs) {
+TEST_F(TritonBackendTest, GetSupportedConfigs) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnVerifiedModule(kHlo));
 
@@ -185,7 +173,7 @@ TEST_P(TritonBackendTest, GetSupportedConfigs) {
   }
 }
 
-TEST_P(TritonBackendTest, GetSupportedConfigsWithEstimates) {
+TEST_F(TritonBackendTest, GetSupportedConfigsWithEstimates) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnVerifiedModule(kHlo));
 
@@ -202,7 +190,7 @@ TEST_P(TritonBackendTest, GetSupportedConfigsWithEstimates) {
   }
 }
 
-TEST_P(TritonBackendTest, GetSupportedConfigsForScaledDot) {
+TEST_F(TritonBackendTest, GetSupportedConfigsForScaledDot) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnVerifiedModule(kScaledDotHlo));
   HloInstruction* fusion_instr =
@@ -213,7 +201,7 @@ TEST_P(TritonBackendTest, GetSupportedConfigsForScaledDot) {
   EXPECT_GT(configs.value().size(), 0);
 }
 
-TEST_P(TritonBackendTest, GetSupportedConfigsWithEstimatesForScaledDot) {
+TEST_F(TritonBackendTest, GetSupportedConfigsWithEstimatesForScaledDot) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnVerifiedModule(kScaledDotHlo));
   ASSERT_OK_AND_ASSIGN(std::vector<CodegenBackend::EstimatedConfig> configs,
@@ -227,7 +215,7 @@ TEST_P(TritonBackendTest, GetSupportedConfigsWithEstimatesForScaledDot) {
   }
 }
 
-TEST_P(TritonBackendTest, GetAndApplyConfigForScaledDot) {
+TEST_F(TritonBackendTest, GetAndApplyConfigForScaledDot) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnVerifiedModule(kScaledDotHlo));
   HloInstruction* fusion_instr =
@@ -238,7 +226,7 @@ TEST_P(TritonBackendTest, GetAndApplyConfigForScaledDot) {
   EXPECT_THAT(backend_.ApplyConfig(*fusion_instr, *configs[0]), IsOk());
 }
 
-TEST_P(TritonBackendTest, GetSupportedConfigsRestrictedDefaultSearch) {
+TEST_F(TritonBackendTest, GetSupportedConfigsRestrictedDefaultSearch) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnVerifiedModule(kHlo));
   absl::StatusOr<std::vector<std::unique_ptr<BackendConfig>>> default_configs =
@@ -253,7 +241,7 @@ TEST_P(TritonBackendTest, GetSupportedConfigsRestrictedDefaultSearch) {
   EXPECT_GE(exhaustive_configs.value().size(), default_configs.value().size());
 }
 
-TEST_P(TritonBackendTest, GetSupportedConfigsForUnsupportedInstruction) {
+TEST_F(TritonBackendTest, GetSupportedConfigsForUnsupportedInstruction) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnVerifiedModule(kHlo));
   HloInstruction* unsupported_instr = module->entry_computation()
@@ -266,7 +254,7 @@ TEST_P(TritonBackendTest, GetSupportedConfigsForUnsupportedInstruction) {
   EXPECT_THAT(configs.value(), testing::IsEmpty());
 }
 
-TEST_P(TritonBackendTest,
+TEST_F(TritonBackendTest,
        GetSupportedConfigsWithEstimatesForUnsupportedInstruction) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnVerifiedModule(kHlo));
@@ -278,7 +266,7 @@ TEST_P(TritonBackendTest,
               IsOkAndHolds(IsEmpty()));
 }
 
-TEST_P(TritonBackendTest, GetDefaultConfigReturnsUnimplementedError) {
+TEST_F(TritonBackendTest, GetDefaultConfigReturnsUnimplementedError) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnVerifiedModule(kHlo));
   EXPECT_THAT(backend_.GetDefaultConfig(
@@ -286,7 +274,7 @@ TEST_P(TritonBackendTest, GetDefaultConfigReturnsUnimplementedError) {
               StatusIs(absl::StatusCode::kUnimplemented));
 }
 
-TEST_P(TritonBackendTest, Compile) {
+TEST_F(TritonBackendTest, Compile) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnVerifiedModule(kHlo));
   ASSERT_OK_AND_ASSIGN(std::vector<std::unique_ptr<BackendConfig>> configs,
@@ -298,7 +286,7 @@ TEST_P(TritonBackendTest, Compile) {
   EXPECT_THAT(executable, absl_testing::IsOk());
 }
 
-TEST_P(TritonBackendTest, AmpereUsesMoreThanTwoStages) {
+TEST_F(TritonBackendTest, AmpereUsesMoreThanTwoStages) {
   if (target_config_.device_description.gpu_compute_capability().IsRocm()) {
     GTEST_SKIP() << "Not supported on ROCm.";
   }
@@ -326,7 +314,7 @@ TEST_P(TritonBackendTest, AmpereUsesMoreThanTwoStages) {
                           }));
 }
 
-TEST_P(TritonBackendTest, VerifyHopperConfigsAreDifferentFromBlackwell) {
+TEST_F(TritonBackendTest, VerifyHopperConfigsAreDifferentFromBlackwell) {
   if (target_config_.device_description.gpu_compute_capability().IsRocm()) {
     GTEST_SKIP() << "Not supported on ROCm.";
   }
@@ -374,7 +362,7 @@ TEST_P(TritonBackendTest, VerifyHopperConfigsAreDifferentFromBlackwell) {
   EXPECT_NE(to_set(hopper_configs), to_set(blackwell_configs));
 }
 
-TEST_P(TritonBackendTest, ScaledDotConfigsAreGenerated) {
+TEST_F(TritonBackendTest, ScaledDotConfigsAreGenerated) {
   se::CudaComputeCapability blackwell_cap{se::CudaComputeCapability::kBlackwell,
                                           /*minor=*/0};
   target_config_.device_description.set_gpu_compute_capability(
@@ -397,7 +385,7 @@ TEST_P(TritonBackendTest, ScaledDotConfigsAreGenerated) {
   EXPECT_GT(exhaustive_configs.size(), configs.value().size());
 }
 
-TEST_P(TritonBackendTest, TmaRunCorrectlyForDotsOfBroadcasts) {
+TEST_F(TritonBackendTest, TmaRunCorrectlyForDotsOfBroadcasts) {
   if (target_config_.device_description.gpu_compute_capability().IsRocm()) {
     GTEST_SKIP() << "Not supported on ROCm.";
   }
@@ -437,7 +425,7 @@ TEST_P(TritonBackendTest, TmaRunCorrectlyForDotsOfBroadcasts) {
   EXPECT_THAT(executable, absl_testing::IsOk());
 }
 
-TEST_P(TritonBackendTest, TmaConfigsAreGeneratedOnlyForHopperAndWorkCorrectly) {
+TEST_F(TritonBackendTest, TmaConfigsAreGeneratedOnlyForHopperAndWorkCorrectly) {
   if (target_config_.device_description.gpu_compute_capability().IsRocm()) {
     GTEST_SKIP() << "Not supported on ROCm.";
   }
@@ -493,7 +481,7 @@ TEST_P(TritonBackendTest, TmaConfigsAreGeneratedOnlyForHopperAndWorkCorrectly) {
   }
 }
 
-TEST_P(TritonBackendTest, GetOverriddenConfigs) {
+TEST_F(TritonBackendTest, GetOverriddenConfigs) {
   AutotuneResult::TritonGemmKey gemm_config;
   gemm_config.set_num_ctas(1);
   gemm_config.set_num_warps(4);
@@ -521,7 +509,7 @@ TEST_P(TritonBackendTest, GetOverriddenConfigs) {
   EXPECT_THAT(triton_config, EqualsProto(gemm_config));
 }
 
-TEST_P(TritonBackendTest, GetOverriddenConfigsFromFile) {
+TEST_F(TritonBackendTest, GetOverriddenConfigsFromFile) {
   ASSERT_OK_AND_ASSIGN(
       tsl::testing::TemporaryDirectory temp_dir,
       tsl::testing::TemporaryDirectory::CreateForCurrentTestcase());
@@ -557,7 +545,7 @@ TEST_P(TritonBackendTest, GetOverriddenConfigsFromFile) {
   EXPECT_THAT(triton_config, EqualsProto(*gemm_config));
 }
 
-TEST_P(TritonBackendTest, WarpSpecializationConfigsAreGenerated) {
+TEST_F(TritonBackendTest, WarpSpecializationConfigsAreGenerated) {
   if (target_config_.device_description.gpu_compute_capability().IsRocm()) {
     GTEST_SKIP() << "Not supported on ROCm.";
   }
@@ -597,7 +585,7 @@ TEST_P(TritonBackendTest, WarpSpecializationConfigsAreGenerated) {
 // 2 also run into address misalignment issues. It might be more precise to also
 // check for a broadcast consumer, but that would complicate the implementation
 // in dot_search_space.cc.
-TEST_P(TritonBackendTest, WarpSpecializationConfigsDoNotHaveNumStagesTwo) {
+TEST_F(TritonBackendTest, WarpSpecializationConfigsDoNotHaveNumStagesTwo) {
   if (target_config_.device_description.gpu_compute_capability().IsRocm()) {
     GTEST_SKIP() << "Not supported on ROCm.";
   }
@@ -638,7 +626,7 @@ TEST_P(TritonBackendTest, WarpSpecializationConfigsDoNotHaveNumStagesTwo) {
       << "No warp specialization configs found!";
 }
 
-TEST_P(TritonBackendTest, Int8FusedGemmCompiles) {
+TEST_F(TritonBackendTest, Int8FusedGemmCompiles) {
   const char kInt8GemmHlo[] = R"(
 HloModule module
 
@@ -663,7 +651,7 @@ ENTRY e {
   EXPECT_THAT(backend_.Compile(*root, *configs[0]), IsOk());
 }
 
-TEST_P(TritonBackendTest, Int8FusedGemm256Compiles) {
+TEST_F(TritonBackendTest, Int8FusedGemm256Compiles) {
   const char kInt8Gemm256Hlo[] = R"(
 HloModule module
 
@@ -688,7 +676,7 @@ ENTRY e {
   EXPECT_THAT(backend_.Compile(*root, *configs[0]), IsOk());
 }
 
-TEST_P(TritonBackendTest, FindsValidConfigForSlicedContractingDimension) {
+TEST_F(TritonBackendTest, FindsValidConfigForSlicedContractingDimension) {
   const char kHlo[] = R"(
 HloModule module
 fused_computation {
@@ -712,7 +700,7 @@ ENTRY e {
   EXPECT_THAT(backend_.Compile(*root, *configs[0]), IsOk());
 }
 
-TEST_P(TritonBackendTest, WarpSpecializationWithBitcastWorksCorrectly) {
+TEST_F(TritonBackendTest, WarpSpecializationWithBitcastWorksCorrectly) {
   if (target_config_.device_description.gpu_compute_capability().IsRocm()) {
     GTEST_SKIP() << "Not supported on ROCm.";
   }
@@ -759,7 +747,7 @@ ENTRY entry_computation {
   }
 }
 
-TEST_P(TritonBackendTest, CostModelOptions_Top) {
+TEST_F(TritonBackendTest, CostModelOptions_Top) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnVerifiedModule(kHlo));
 
@@ -775,7 +763,7 @@ TEST_P(TritonBackendTest, CostModelOptions_Top) {
   EXPECT_THAT(configs.value(), SizeIs(2));
 }
 
-TEST_P(TritonBackendTest, CostModelOptions_TopFromDefault) {
+TEST_F(TritonBackendTest, CostModelOptions_TopFromDefault) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnVerifiedModule(kHlo));
 
@@ -794,7 +782,7 @@ TEST_P(TritonBackendTest, CostModelOptions_TopFromDefault) {
   EXPECT_THAT(configs.value(), SizeIs(2));
 }
 
-TEST_P(TritonBackendTest, CostModelOptions_Mixin) {
+TEST_F(TritonBackendTest, CostModelOptions_Mixin) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnVerifiedModule(kHlo));
 
@@ -816,7 +804,7 @@ TEST_P(TritonBackendTest, CostModelOptions_Mixin) {
   EXPECT_THAT(configs.value(), SizeIs(default_size + 2));
 }
 
-TEST_P(TritonBackendTest, CostModelOptions_Filter) {
+TEST_F(TritonBackendTest, CostModelOptions_Filter) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnVerifiedModule(kHlo));
 
@@ -837,7 +825,7 @@ TEST_P(TritonBackendTest, CostModelOptions_Filter) {
   EXPECT_LT(configs.value().size(), default_configs.value().size());
 }
 
-TEST_P(TritonBackendTest, CostModelOptions_Combination) {
+TEST_F(TritonBackendTest, CostModelOptions_Combination) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnVerifiedModule(kHlo));
 
@@ -859,11 +847,7 @@ TEST_P(TritonBackendTest, CostModelOptions_Combination) {
   EXPECT_THAT(configs.value(), SizeIs(7));
 }
 
-
-TEST_P(TritonBackendTest, Version) { EXPECT_NE(backend_.version(), ""); }
-
-INSTANTIATE_TEST_SUITE_P(TritonBackendTestSuite, TritonBackendTest,
-                         ::testing::Bool());
+TEST_F(TritonBackendTest, Version) { EXPECT_NE(backend_.version(), ""); }
 
 }  // namespace
 }  // namespace gpu
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/BUILD b/third_party/xla/xla/backends/gpu/codegen/triton/BUILD
index 5fff4b8f4cb..af53f9bc048 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/BUILD
@@ -606,9 +606,8 @@ cc_library(
 
 xla_cc_test(
     name = "support_test",
-    timeout = "long",  # TODO: b/514293537 - remove after disabling symbolic analysis.
     srcs = ["support_test.cc"],
-    shard_count = 50,  # TODO: b/514293537 - change to 25 after disabling symbolic analysis.
+    shard_count = 25,
     deps = [
         ":support",
         ":support_test_base",
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc
index cf708628ddf..a874d18b8da 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc
@@ -157,10 +157,6 @@ bool DoesOpSupportType(HloOpcode opcode, PrimitiveType type) {
   }
 }
 
-std::string TilingParametersToString(bool tiling_propagation_enabled) {
-  return tiling_propagation_enabled ? "_ExperimentalTiling" : "_SymbolicTiling";
-}
-
 std::vector<xla::PrimitiveType> AllOpSupportedTypes(HloOpcode opcode) {
   std::vector<xla::PrimitiveType> result;
   absl::c_copy_if(AllXlaDataTypes(), std::back_inserter(result),
@@ -216,38 +212,6 @@ auto AllTestCombinationsForOpcodes(absl::Span<const HloOpcode> opcodes) {
   return ::testing::ValuesIn(test_combinations);
 };
 
-auto AllTestCombinationsForOpcodesWithTiling(
-    absl::Span<const HloOpcode> opcodes) {
-  std::vector<
-      std::tuple<PrimitiveType, HloOpcode, se::GpuComputeCapability, bool>>
-      test_combinations;
-  for (PrimitiveType data_type : AllXlaDataTypes()) {
-    for (HloOpcode opcode : opcodes) {
-      if (DoesOpSupportType(opcode, data_type)) {
-        for (se::GpuComputeCapability cc : AllDevicesToTest()) {
-          test_combinations.push_back({data_type, opcode, cc, false});
-          test_combinations.push_back({data_type, opcode, cc, true});
-        }
-      }
-    }
-  }
-  return ::testing::ValuesIn(test_combinations);
-};
-
-std::string SupportTestTypeAndOpcodeAndDeviceAndTilingToString(
-    const ::testing::TestParamInfo<
-        std::tuple<PrimitiveType, HloOpcode, se::GpuComputeCapability, bool>>&
-        info) {
-  return absl::StrCat(
-      SupportTestTypeAndOpcodeAndDeviceToString(
-          ::testing::TestParamInfo<
-              std::tuple<PrimitiveType, HloOpcode, se::GpuComputeCapability>>(
-              {std::get<0>(info.param), std::get<1>(info.param),
-               std::get<2>(info.param)},
-              info.index)),
-      TilingParametersToString(std::get<3>(info.param)));
-}
-
 // Expected failure mode of the Triton lowering.
 enum class ExpectedFailMode {
   // Denotes a graceful failure, e.g. a verifier failure, or an absl::Status.
@@ -270,15 +234,11 @@ class SupportTest : public HloHardwareIndependentTestBase,
               return ParseAndReturnVerifiedModule(hlo_text);
             }) {}
 
-  virtual bool EnableTilingPropagation() const = 0;
-
   DebugOptions GetDebugOptionsForTest() const override {
     auto options = HloHardwareIndependentTestBase::GetDebugOptionsForTest();
     // It's necessary to set this manually, because it's disabled in optimized
     // builds and there are some ASAN builds that run on TAP with -c opt.
     options.set_xla_gpu_llvm_verification_level(1);
-    options.set_xla_gpu_experimental_enable_tiling_propagation(
-        EnableTilingPropagation());
     return options;
   }
 
@@ -408,19 +368,7 @@ class SupportTest : public HloHardwareIndependentTestBase,
   }
 };
 
-class SupportTestWithTilingParam : public SupportTest,
-                                   public ::testing::WithParamInterface<bool> {
- public:
-  bool EnableTilingPropagation() const override { return GetParam(); }
-};
-
-INSTANTIATE_TEST_SUITE_P(SupportTestWithTilingParamTestSuite,
-                         SupportTestWithTilingParam, ::testing::Bool(),
-                         [](const ::testing::TestParamInfo<bool>& info) {
-                           return TilingParametersToString(info.param);
-                         });
-
-TEST_P(SupportTestWithTilingParam, IsTritonSupportedComputationSkipsRootTuple) {
+TEST_F(SupportTest, IsTritonSupportedComputationSkipsRootTuple) {
   const std::string kHlo = R"(
   HloModule m
   ENTRY main {
@@ -434,21 +382,15 @@ TEST_P(SupportTestWithTilingParam, IsTritonSupportedComputationSkipsRootTuple) {
       *module->entry_computation(), se::CudaComputeCapability::Hopper()));
 }
 
-class SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam
+class SupportTestWithTypeAndOpcodeAndDeviceParam
     : public SupportTest,
-      public ::testing::WithParamInterface<std::tuple<
-          PrimitiveType, HloOpcode, se::GpuComputeCapability, bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<3>(GetParam());
-  }
-};
+      public ::testing::WithParamInterface<
+          std::tuple<PrimitiveType, HloOpcode, se::GpuComputeCapability>> {};
 
-using BitcastOrReshapeTest =
-    SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using BitcastOrReshapeTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(BitcastOrReshapeTest, IsTritonSupportedBitcastOrReshape) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   parameter_0 = $0[1,16,4] parameter(0)
@@ -461,7 +403,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(BitcastOrReshapeTest, IsTritonSupported0DBitcastOrReshape) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   parameter_0 = $0[1,1,1] parameter(0)
@@ -478,13 +420,13 @@ constexpr std::array kTestedOpsBitcastReshape = {HloOpcode::kBitcast,
 
 INSTANTIATE_TEST_SUITE_P(
     BitcastOrReshapeTestSuite, BitcastOrReshapeTest,
-    AllTestCombinationsForOpcodesWithTiling(kTestedOpsBitcastReshape),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+    AllTestCombinationsForOpcodes(kTestedOpsBitcastReshape),
+    SupportTestTypeAndOpcodeAndDeviceToString);
 
-using PadTest = SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using PadTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(PadTest, IsTritonSupportedHighPad) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   p0 = $0[4, 4] parameter(0)
@@ -498,7 +440,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(PadTest, IsTritonSupportedInteriorPad) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   p0 = $0[4] parameter(0)
@@ -512,7 +454,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(PadTest, IsTritonSupportedLowPad) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   p0 = $0[4] parameter(0)
@@ -528,8 +470,8 @@ ENTRY triton_computation {
 constexpr std::array kTestedOpsPad = {HloOpcode::kPad};
 
 INSTANTIATE_TEST_SUITE_P(PadTestSuite, PadTest,
-                         AllTestCombinationsForOpcodesWithTiling(kTestedOpsPad),
-                         SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+                         AllTestCombinationsForOpcodes(kTestedOpsPad),
+                         SupportTestTypeAndOpcodeAndDeviceToString);
 
 TEST_F(HloHardwareIndependentTestBase, PadOutsideGemmFusionIsRejected) {
   const std::string kHlo = R"(
@@ -610,11 +552,10 @@ ENTRY main {
   EXPECT_TRUE(decision);
 }
 
-using UnaryElementwiseTest =
-    SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using UnaryElementwiseTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(UnaryElementwiseTest, IsTritonSupportedUnaryElementwise) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kDefaultHloTemplate = R"(
 ENTRY triton_computation {
   parameter_0 = $0[33,68] parameter(0)
@@ -705,21 +646,17 @@ constexpr std::array kTestedOpsUnaryElementwise = {
 
 INSTANTIATE_TEST_SUITE_P(
     UnaryElementwiseTestSuite, UnaryElementwiseTest,
-    AllTestCombinationsForOpcodesWithTiling(kTestedOpsUnaryElementwise),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+    AllTestCombinationsForOpcodes(kTestedOpsUnaryElementwise),
+    SupportTestTypeAndOpcodeAndDeviceToString);
 
 class ConvertTest
     : public SupportTest,
-      public ::testing::WithParamInterface<std::tuple<
-          PrimitiveType, PrimitiveType, se::GpuComputeCapability, bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<3>(GetParam());
-  }
+      public ::testing::WithParamInterface<
+          std::tuple<PrimitiveType, PrimitiveType, se::GpuComputeCapability>> {
 };
 
 TEST_P(ConvertTest, Convert) {
-  auto [data_type_in, data_type_out, cc, tiling] = GetParam();
+  auto [data_type_in, data_type_out, cc] = GetParam();
 
   const std::string hlo_text = absl::Substitute(
       R"(
@@ -777,29 +714,17 @@ ENTRY triton_computation {
 
 constexpr std::array kTestedOpsConvert = {HloOpcode::kConvert};
 
-std::string ConvertTestParamToString(
-    const ::testing::TestParamInfo<std::tuple<
-        PrimitiveType, PrimitiveType, se::GpuComputeCapability, bool>>& data) {
-  auto [type_in, type_out, cc, tiling] = data.param;
-  return absl::StrCat(primitive_util::LowercasePrimitiveTypeName(type_in), "_",
-                      primitive_util::LowercasePrimitiveTypeName(type_out), "_",
-                      ComputeCapabilityToString(cc),
-                      TilingParametersToString(tiling));
-}
-
 INSTANTIATE_TEST_SUITE_P(
     ConvertTestSuite, ConvertTest,
     ::testing::Combine(::testing::ValuesIn(AllXlaDataTypes()),
                        ::testing::ValuesIn(AllXlaDataTypes()),
-                       ::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    ConvertTestParamToString);
+                       ::testing::ValuesIn(AllDevicesToTest())),
+    SupportTestTwoTypesAndDeviceToString);
 
-using BinaryElementwiseTest =
-    SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using BinaryElementwiseTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(BinaryElementwiseTest, IsTritonSupportedBinaryElementwise) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   parameter_0 = $0[11,63] parameter(0)
@@ -842,7 +767,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(BinaryElementwiseTest, IsTritonSupportedBinaryElementwise0D) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   parameter_0 = $0[] parameter(0)
@@ -909,14 +834,13 @@ constexpr std::array kTestedOpsBinaryElementwise = {
 
 INSTANTIATE_TEST_SUITE_P(
     BinaryElementwiseTestSuite, BinaryElementwiseTest,
-    AllTestCombinationsForOpcodesWithTiling(kTestedOpsBinaryElementwise),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+    AllTestCombinationsForOpcodes(kTestedOpsBinaryElementwise),
+    SupportTestTypeAndOpcodeAndDeviceToString);
 
-using TernaryElementwiseTest =
-    SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using TernaryElementwiseTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(TernaryElementwiseTest, IsTritonSupportedTernaryElementwise) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   parameter_0 = $2[13,63] parameter(0)
@@ -953,10 +877,10 @@ constexpr std::array kTestedOpsTernaryElementwise = {HloOpcode::kSelect,
 
 INSTANTIATE_TEST_SUITE_P(
     TernaryElementwiseTestSuite, TernaryElementwiseTest,
-    AllTestCombinationsForOpcodesWithTiling(kTestedOpsTernaryElementwise),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+    AllTestCombinationsForOpcodes(kTestedOpsTernaryElementwise),
+    SupportTestTypeAndOpcodeAndDeviceToString);
 
-using ReduceTest = SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using ReduceTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 static absl::string_view init_value(PrimitiveType dtype) {
   if (dtype == C64 || dtype == C128) {
@@ -969,7 +893,7 @@ static absl::string_view init_value(PrimitiveType dtype) {
 }
 
 TEST_P(ReduceTest, IsTritonSupportedReduction) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = absl::Substitute(R"(
 add {
   Arg_0 = $$0[] parameter(0)
@@ -1020,7 +944,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(ReduceTest, MultidimensionReductionIsSupported) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = absl::Substitute(R"(
 add {
   Arg_0 = $$0[] parameter(0)
@@ -1038,10 +962,6 @@ ENTRY triton_computation {
   ASSERT_OK_AND_ASSIGN(
       TestedInstruction ti,
       ParseTemplateAndGetInstruction(kHloTestTemplate, data_type, opcode));
-  if (!tiling) {
-    EXPECT_FALSE(IsTritonSupportedInstruction(ti.Instruction(), cc));
-    return;
-  }
   if (!IsTritonSupportedInstruction(ti.Instruction(), cc)) {
     return;
   }
@@ -1049,7 +969,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(ReduceTest, IsTritonSupportedReduceWithNonLastReduceDimension) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = absl::Substitute(R"(
 add {
   Arg_0 = $$0[] parameter(0)
@@ -1080,7 +1000,7 @@ ENTRY triton_computation {
 
 TEST_P(ReduceTest,
        UnsupportedReduceWithMoreThanOneOperandsFailsGracefullyWithTriton) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = absl::Substitute(R"(
 add {
   Arg_0 = $$0[] parameter(0)
@@ -1129,7 +1049,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(ReduceTest, UnsupportedReductionComputationFailsGracefullyWithTriton) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = absl::Substitute(R"(
 custom_call {
   Arg_0 = $$0[] parameter(0)
@@ -1152,13 +1072,11 @@ ENTRY triton_computation {
 
 constexpr std::array kTestedOpsReduction = {HloOpcode::kReduce};
 
-INSTANTIATE_TEST_SUITE_P(
-    ReduceTestSuite, ReduceTest,
-    AllTestCombinationsForOpcodesWithTiling(kTestedOpsReduction),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(ReduceTestSuite, ReduceTest,
+                         AllTestCombinationsForOpcodes(kTestedOpsReduction),
+                         SupportTestTypeAndOpcodeAndDeviceToString);
 
-using ReductionComputationTest =
-    SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using ReductionComputationTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 // The test below tests what kind of binary element-wise operations are
 // supported within a reduction's computation.
@@ -1166,7 +1084,7 @@ using ReductionComputationTest =
 // Note that there is a difference in what is supported inside the reduction
 // computation and in regular HLO. See triton_support.cc for more details.
 TEST_P(ReductionComputationTest, DifferentBinaryOps) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = absl::Substitute(
       R"(
 reduce_computation {
@@ -1212,20 +1130,13 @@ std::vector<HloOpcode> ExcludeOps(absl::Span<const HloOpcode> all_ops,
   return ret;
 }
 
-INSTANTIATE_TEST_SUITE_P(ReductionComputationTestSuite,
-                         ReductionComputationTest,
-                         AllTestCombinationsForOpcodesWithTiling(
-                             ExcludeOps(kTestedOpsBinaryElementwise,
-                                        {HloOpcode::kCompare})),
-                         SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(
+    ReductionComputationTestSuite, ReductionComputationTest,
+    AllTestCombinationsForOpcodes(ExcludeOps(kTestedOpsBinaryElementwise,
+                                             {HloOpcode::kCompare})),
+    SupportTestTypeAndOpcodeAndDeviceToString);
 
-class ScanTest
-    : public SupportTest,
-      public ::testing::WithParamInterface<
-          std::tuple<PrimitiveType, HloOpcode, se::GpuComputeCapability>> {
- public:
-  bool EnableTilingPropagation() const override { return true; }
-};
+using ScanTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(ScanTest, IsTritonSupportedScan) {
   auto [data_type, opcode, cc] = GetParam();
@@ -1268,8 +1179,6 @@ INSTANTIATE_TEST_SUITE_P(ScanTestSuite, ScanTest,
 // Negative tests for CanTritonHandleScan forbid conditions.
 class ScanSupportTest : public SupportTest {
  public:
-  bool EnableTilingPropagation() const override { return true; }
-
   void ExpectScanNotSupported(absl::string_view entry_hlo,
                               absl::string_view combiner_hlo = R"(
   add {
@@ -1363,10 +1272,10 @@ TEST_F(ScanSupportTest, NonAssociativeScanNotSupported) {
   })");
 }
 
-using TransposeTest = SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using TransposeTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(TransposeTest, LoadTranspose3D) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   parameter_0 = $0[125,127,37] parameter(0)
@@ -1381,39 +1290,19 @@ ENTRY triton_computation {
 
 constexpr std::array kTestedOpsTranspose = {HloOpcode::kTranspose};
 
-INSTANTIATE_TEST_SUITE_P(
-    TransposeTestSuite, TransposeTest,
-    AllTestCombinationsForOpcodesWithTiling(kTestedOpsTranspose),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(TransposeTestSuite, TransposeTest,
+                         AllTestCombinationsForOpcodes(kTestedOpsTranspose),
+                         SupportTestTypeAndOpcodeAndDeviceToString);
 
 class SupportTestWithTypeAndDeviceParam
     : public SupportTest,
       public ::testing::WithParamInterface<
           std::tuple<PrimitiveType, se::GpuComputeCapability>> {};
 
-class SupportTestWithTypeAndDeviceAndTilingParam
-    : public SupportTest,
-      public ::testing::WithParamInterface<
-          std::tuple<PrimitiveType, se::GpuComputeCapability, bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<2>(GetParam());
-  }
-};
-
-std::string SupportTestTypeAndDeviceAndTilingToString(
-    const ::testing::TestParamInfo<
-        std::tuple<PrimitiveType, se::GpuComputeCapability, bool>>& data) {
-  auto [type, cc, tiling] = data.param;
-  return absl::StrCat(primitive_util::LowercasePrimitiveTypeName(type), "_",
-                      ComputeCapabilityToString(cc),
-                      TilingParametersToString(tiling));
-}
-
-using SliceTest = SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using SliceTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(SliceTest, ContinuousSlice) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = (R"(
 ENTRY triton_computation {
   p = $0[128,32] parameter(0)
@@ -1427,7 +1316,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(SliceTest, NonContinuousSliceWhereStrideDividesOffsetEvenly) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = (R"(
 ENTRY triton_computation {
   p = f32[16,16,32] parameter(0)
@@ -1441,7 +1330,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(SliceTest, NonContinuousSliceWhereStrideDoesNotDivideOffsetEvenly) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = (R"(
 ENTRY triton_computation {
   p = f32[16,16,32] parameter(0)
@@ -1456,38 +1345,19 @@ ENTRY triton_computation {
 
 constexpr std::array kTestedOpsSlice = {HloOpcode::kSlice};
 
-INSTANTIATE_TEST_SUITE_P(
-    SliceTestSuite, SliceTest,
-    AllTestCombinationsForOpcodesWithTiling(kTestedOpsSlice),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(SliceTestSuite, SliceTest,
+                         AllTestCombinationsForOpcodes(kTestedOpsSlice),
+                         SupportTestTypeAndOpcodeAndDeviceToString);
 
 class SupportTestWithDeviceParam
     : public SupportTest,
       public ::testing::WithParamInterface<se::GpuComputeCapability> {};
 
-class SupportTestWithDeviceAndTilingParam
-    : public SupportTest,
-      public ::testing::WithParamInterface<
-          std::tuple<se::GpuComputeCapability, bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<1>(GetParam());
-  }
-};
-
-std::string SupportTestDeviceAndTilingToString(
-    const ::testing::TestParamInfo<std::tuple<se::GpuComputeCapability, bool>>&
-        data) {
-  auto [cc, tiling] = data.param;
-  return absl::StrCat(ComputeCapabilityToString(cc),
-                      TilingParametersToString(tiling));
-}
-
-using ConcatenateDeviceTest = SupportTestWithDeviceAndTilingParam;
+using ConcatenateDeviceTest = SupportTestWithDeviceParam;
 
 TEST_P(ConcatenateDeviceTest,
        TritonDoesNotSupportConcatenateOfUnnestedParameters) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   p0 = $0[18,128,20] parameter(0)
@@ -1502,7 +1372,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(ConcatenateDeviceTest, TritonDoesNotSupportUnalignedConcatenate) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   // Operand 0 has size 63 along concat dim (1), which is not divisible by 64.
   // Operand 1 has size 128 (aligned).
   const std::string kHloTestTemplate = R"(
@@ -1517,16 +1387,14 @@ ENTRY triton_computation {
   RunSupportTest(std::move(ti), /*output_tile_sizes=*/{1, 64, 1}, cc);
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    ConcatenateTestSuite, ConcatenateDeviceTest,
-    ::testing::Combine(::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(ConcatenateTestSuite, ConcatenateDeviceTest,
+                         ::testing::ValuesIn(AllDevicesToTest()),
+                         SupportTestDeviceToString);
 
-using ConcatenateTest = SupportTestWithTypeAndDeviceAndTilingParam;
+using ConcatenateTest = SupportTestWithTypeAndDeviceParam;
 
 TEST_P(ConcatenateTest, IsTritonSupportedConcatenate) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   p0 = $0[128] parameter(0)
@@ -1545,14 +1413,13 @@ constexpr std::array kTestedOpsConcatenate = {HloOpcode::kConcatenate};
 INSTANTIATE_TEST_SUITE_P(
     ConcatenateTestSuite, ConcatenateTest,
     ::testing::Combine(::testing::ValuesIn(AllXlaDataTypes()),
-                       ::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestTypeAndDeviceAndTilingToString);
+                       ::testing::ValuesIn(AllDevicesToTest())),
+    SupportTestTypeAndDeviceToString);
 
-using CollectiveTest = SupportTestWithTypeAndDeviceAndTilingParam;
+using CollectiveTest = SupportTestWithTypeAndDeviceParam;
 
 TEST_P(CollectiveTest, IsTritonSupportedAllGather) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   input = $0[128,64] parameter(0)
@@ -1566,7 +1433,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(CollectiveTest, IsTritonSupportedAllGatherStartAndDone) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   input = $0[128,32] parameter(0)
@@ -1582,7 +1449,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(CollectiveTest, IsTritonSupportedAllReduce) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 apply_op {
   x = $0[] parameter(0)
@@ -1602,7 +1469,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(CollectiveTest, UnsupportedAllToAllFailsGracefullyWithTriton) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   input = $0[128,32] parameter(0)
@@ -1615,7 +1482,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(CollectiveTest, UnsupportedCollectivePermuteFailsGracefullyWithTriton) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   a = $0[128,32] parameter(0)
@@ -1631,7 +1498,7 @@ ENTRY triton_computation {
 
 TEST_P(CollectiveTest,
        UnsupportedCollectivePermuteStartAndDoneFailGracefullyWithTriton) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   // 'collective-permute-start' and 'collective-permute-done' need to be tested
   // together, since the HLO verifier relies on one directly consuming the
   // other.
@@ -1657,7 +1524,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(CollectiveTest, UnsupportedReduceScatterFailsGracefullyWithTriton) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 apply_op {
   lhs = $0[] parameter(0)
@@ -1681,7 +1548,7 @@ TEST_P(CollectiveTest,
   // 'async-start', 'async-update', and 'async-done' need to be tested together,
   // since the HLO verifier requires 'async-start' and 'async-done' to always
   // appear together within a module.
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 async_computation {
   ROOT p0 = $0[10] parameter(0)
@@ -1713,7 +1580,7 @@ ENTRY triton_computation {
 
 TEST_P(CollectiveTest,
        UnsupportedCollectiveBroadcastFailsGracefullyWithTriton) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   input = $0[128,32] parameter(0)
@@ -1727,7 +1594,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(CollectiveTest, UnsupportedReplicaIdFailsGracefullyWithTriton) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   ROOT replica_id = u32[] replica-id()
@@ -1740,7 +1607,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(CollectiveTest, UnsupportedPartitionIdFailsGracefullyWithTriton) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   ROOT partition_id = u32[] partition-id()
@@ -1753,7 +1620,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(CollectiveTest, UnsupportedRaggedAllToAllFailsGracefullyWithTriton) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   input = $0[128,32] parameter(0)
@@ -1799,14 +1666,13 @@ constexpr std::array kTestedOpsCollectives = {
 INSTANTIATE_TEST_SUITE_P(
     CollectiveTestSuite, CollectiveTest,
     ::testing::Combine(::testing::ValuesIn(AllXlaDataTypes()),
-                       ::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestTypeAndDeviceAndTilingToString);
+                       ::testing::ValuesIn(AllDevicesToTest())),
+    SupportTestTypeAndDeviceToString);
 
-using BroadcastTest = SupportTestWithTypeAndDeviceAndTilingParam;
+using BroadcastTest = SupportTestWithTypeAndDeviceParam;
 
 TEST_P(BroadcastTest, Broadcast) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   input = $0[35,131] parameter(0)
@@ -1823,14 +1689,13 @@ constexpr std::array kTestedOpsBroadcast = {HloOpcode::kBroadcast};
 INSTANTIATE_TEST_SUITE_P(
     BroadcastTestSuite, BroadcastTest,
     ::testing::Combine(::testing::ValuesIn(AllXlaDataTypes()),
-                       ::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestTypeAndDeviceAndTilingToString);
+                       ::testing::ValuesIn(AllDevicesToTest())),
+    SupportTestTypeAndDeviceToString);
 
-using ParameterTest = SupportTestWithTypeAndDeviceAndTilingParam;
+using ParameterTest = SupportTestWithTypeAndDeviceParam;
 
 TEST_P(ParameterTest, Parameter) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   std::string hlo_test_template =
       R"(
 ENTRY triton_computation {
@@ -1854,16 +1719,15 @@ constexpr std::array kTestedOpsParameter = {HloOpcode::kParameter};
 INSTANTIATE_TEST_SUITE_P(
     ParameterTestSuite, ParameterTest,
     ::testing::Combine(::testing::ValuesIn(AllXlaDataTypes()),
-                       ::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestTypeAndDeviceAndTilingToString);
+                       ::testing::ValuesIn(AllDevicesToTest())),
+    SupportTestTypeAndDeviceToString);
 
-using ConstantTest = SupportTestWithTypeAndDeviceAndTilingParam;
+using ConstantTest = SupportTestWithTypeAndDeviceParam;
 
 TEST_P(ConstantTest, ConstantEffectiveScalar) {
   // The IsTritonSupportedReduction effectively tests the scalar constant
   // support.
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = absl::Substitute(R"(
 ENTRY triton_computation {
   ROOT const = $$0[1,1] constant({{$0}})
@@ -1879,7 +1743,7 @@ ENTRY triton_computation {
 TEST_P(ConstantTest, Constant2D) {
   // The IsTritonSupportedReduction effectively tests the scalar constant
   // support.
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = absl::Substitute(R"(
 ENTRY triton_computation {
   ROOT const = $$0[3,3] constant({{$0,$0,$0},{$0,$0,$0},{$0,$0,$0}})
@@ -1897,14 +1761,13 @@ constexpr std::array kTestedOpsConstant = {HloOpcode::kConstant};
 INSTANTIATE_TEST_SUITE_P(
     ConstantTestSuite, ConstantTest,
     ::testing::Combine(::testing::ValuesIn(AllXlaDataTypes()),
-                       ::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestTypeAndDeviceAndTilingToString);
+                       ::testing::ValuesIn(AllDevicesToTest())),
+    SupportTestTypeAndDeviceToString);
 
-using IotaTest = SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using IotaTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(IotaTest, Iota2D) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   ROOT input = $0[35,131] iota(), iota_dimension=0
@@ -1917,15 +1780,14 @@ ENTRY triton_computation {
 
 constexpr std::array kTestedOpsIota = {HloOpcode::kIota};
 
-INSTANTIATE_TEST_SUITE_P(
-    IotaTestSuite, IotaTest,
-    AllTestCombinationsForOpcodesWithTiling(kTestedOpsIota),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(IotaTestSuite, IotaTest,
+                         AllTestCombinationsForOpcodes(kTestedOpsIota),
+                         SupportTestTypeAndOpcodeAndDeviceToString);
 
-using RngTest = SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using RngTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(RngTest, Rng) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   low = $0[] parameter(0)
@@ -1941,13 +1803,13 @@ ENTRY triton_computation {
 constexpr std::array kTestedOpsRng = {HloOpcode::kRng};
 
 INSTANTIATE_TEST_SUITE_P(RngTestSuite, RngTest,
-                         AllTestCombinationsForOpcodesWithTiling(kTestedOpsRng),
-                         SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+                         AllTestCombinationsForOpcodes(kTestedOpsRng),
+                         SupportTestTypeAndOpcodeAndDeviceToString);
 
-using RngBitGeneratorTest = SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using RngBitGeneratorTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(RngBitGeneratorTest, RngBitGenerator) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   state = u64[2] parameter(0)
@@ -1962,13 +1824,13 @@ ENTRY triton_computation {
 
 INSTANTIATE_TEST_SUITE_P(
     RngBitGeneratorTestSuite, RngBitGeneratorTest,
-    AllTestCombinationsForOpcodesWithTiling({HloOpcode::kRngBitGenerator}),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+    AllTestCombinationsForOpcodes({HloOpcode::kRngBitGenerator}),
+    SupportTestTypeAndOpcodeAndDeviceToString);
 
-using RngGetAndUpdateStateTest = SupportTestWithDeviceAndTilingParam;
+using RngGetAndUpdateStateTest = SupportTestWithDeviceParam;
 
 TEST_P(RngGetAndUpdateStateTest, RngGetAndUpdateState) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   ROOT root = u64[2]{0} rng-get-and-update-state(), delta=4096
@@ -1980,16 +1842,15 @@ ENTRY triton_computation {
   RunSupportTest(std::move(ti), /*output_tile_sizes=*/{1}, cc);
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    RngGetAndUpdateStateTestSuite, RngGetAndUpdateStateTest,
-    ::testing::Combine(::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(RngGetAndUpdateStateTestSuite,
+                         RngGetAndUpdateStateTest,
+                         ::testing::ValuesIn(AllDevicesToTest()),
+                         SupportTestDeviceToString);
 
-using ComplexTest = SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using ComplexTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(ComplexTest, Complex) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
 
   const std::string kF32HloTestTemplate = R"(
 ENTRY triton_computation {
@@ -2012,15 +1873,14 @@ ENTRY triton_computation {
   RunSupportTest(std::move(ti), /*output_tile_sizes=*/{16, 32}, cc);
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    ComplexTestSuite, ComplexTest,
-    AllTestCombinationsForOpcodesWithTiling({HloOpcode::kComplex}),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(ComplexTestSuite, ComplexTest,
+                         AllTestCombinationsForOpcodes({HloOpcode::kComplex}),
+                         SupportTestTypeAndOpcodeAndDeviceToString);
 
-using ConditionalTest = SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using ConditionalTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(ConditionalTest, Conditional) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 true_branch {
   p_true = $0[10] parameter(0)
@@ -2045,13 +1905,13 @@ ENTRY triton_computation {
 
 INSTANTIATE_TEST_SUITE_P(
     ConditionalTestSuite, ConditionalTest,
-    AllTestCombinationsForOpcodesWithTiling({HloOpcode::kConditional}),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+    AllTestCombinationsForOpcodes({HloOpcode::kConditional}),
+    SupportTestTypeAndOpcodeAndDeviceToString);
 
-using WhileTest = SupportTestWithDeviceAndTilingParam;
+using WhileTest = SupportTestWithDeviceParam;
 // TODO: b/363981282 - Add tests for more data types.
 TEST_P(WhileTest, While) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   const std::string kHloTestTemplate = R"(
 body {
   constant = s32[] constant(1)
@@ -2074,16 +1934,14 @@ ENTRY triton_computation {
   RunSupportTest(std::move(ti), /*output_tile_sizes=*/{}, cc);
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    WhileTestSuite, WhileTest,
-    ::testing::Combine(::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(WhileTestSuite, WhileTest,
+                         ::testing::ValuesIn(AllDevicesToTest()),
+                         SupportTestDeviceToString);
 
-using CallTest = SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using CallTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(CallTest, Call) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 called_computation {
   p = $0[10] parameter(0)
@@ -2100,16 +1958,14 @@ ENTRY triton_computation {
   RunSupportTest(std::move(ti), /*output_tile_sizes=*/{1}, cc);
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    CallTestSuite, CallTest,
-    AllTestCombinationsForOpcodesWithTiling({HloOpcode::kCall}),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(CallTestSuite, CallTest,
+                         AllTestCombinationsForOpcodes({HloOpcode::kCall}),
+                         SupportTestTypeAndOpcodeAndDeviceToString);
 
-using BatchNormInferenceTest =
-    SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using BatchNormInferenceTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(BatchNormInferenceTest, BatchNormInference) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   operand = $0[4,8,16,32] parameter(0)
@@ -2128,14 +1984,13 @@ ENTRY triton_computation {
 
 INSTANTIATE_TEST_SUITE_P(
     BatchNormInferenceSuite, BatchNormInferenceTest,
-    AllTestCombinationsForOpcodesWithTiling({HloOpcode::kBatchNormInference}),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+    AllTestCombinationsForOpcodes({HloOpcode::kBatchNormInference}),
+    SupportTestTypeAndOpcodeAndDeviceToString);
 
-using BatchNormTrainingTest =
-    SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using BatchNormTrainingTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(BatchNormTrainingTest, BatchNormTraining) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   operand = $0[4,8,16,32] parameter(0)
@@ -2153,13 +2008,13 @@ ENTRY triton_computation {
 
 INSTANTIATE_TEST_SUITE_P(
     BatchNormTrainingSuite, BatchNormTrainingTest,
-    AllTestCombinationsForOpcodesWithTiling({HloOpcode::kBatchNormTraining}),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+    AllTestCombinationsForOpcodes({HloOpcode::kBatchNormTraining}),
+    SupportTestTypeAndOpcodeAndDeviceToString);
 
-using BatchNormGradTest = SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using BatchNormGradTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(BatchNormGradTest, BatchNormGrad) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   operand = $0[4,8,16,32] parameter(0)
@@ -2179,13 +2034,13 @@ ENTRY triton_computation {
 
 INSTANTIATE_TEST_SUITE_P(
     BatchNormGradSuite, BatchNormGradTest,
-    AllTestCombinationsForOpcodesWithTiling({HloOpcode::kBatchNormGrad}),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+    AllTestCombinationsForOpcodes({HloOpcode::kBatchNormGrad}),
+    SupportTestTypeAndOpcodeAndDeviceToString);
 
-using DomainTest = SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using DomainTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(DomainTest, Domain) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   operand = $0[] parameter(0)
@@ -2197,15 +2052,14 @@ ENTRY triton_computation {
   RunSupportTest(std::move(ti), /*output_tile_sizes=*/{}, cc);
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    DomainSuite, DomainTest,
-    AllTestCombinationsForOpcodesWithTiling({HloOpcode::kDomain}),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(DomainSuite, DomainTest,
+                         AllTestCombinationsForOpcodes({HloOpcode::kDomain}),
+                         SupportTestTypeAndOpcodeAndDeviceToString);
 
-using GetDimensionSizeTest = SupportTestWithDeviceAndTilingParam;
+using GetDimensionSizeTest = SupportTestWithDeviceParam;
 
 TEST_P(GetDimensionSizeTest, GetDimensionSize) {
-  const auto [cc, tiling] = GetParam();
+  const auto cc = GetParam();
 
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
@@ -2219,16 +2073,14 @@ ENTRY triton_computation {
   RunSupportTest(std::move(ti), /*output_tile_sizes=*/{}, cc);
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    GetDimensionSizeSuite, GetDimensionSizeTest,
-    ::testing::Combine(::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(GetDimensionSizeSuite, GetDimensionSizeTest,
+                         ::testing::ValuesIn(AllDevicesToTest()),
+                         SupportTestDeviceToString);
 
-using ReverseTest = SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using ReverseTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(ReverseTest, Reverse) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   operand = $0[16,32] parameter(0)
@@ -2240,40 +2092,31 @@ ENTRY triton_computation {
   RunSupportTest(std::move(ti), /*output_tile_sizes=*/{4, 8}, cc);
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    ReverseSuite, ReverseTest,
-    AllTestCombinationsForOpcodesWithTiling({HloOpcode::kReverse}),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(ReverseSuite, ReverseTest,
+                         AllTestCombinationsForOpcodes({HloOpcode::kReverse}),
+                         SupportTestTypeAndOpcodeAndDeviceToString);
 
-class DotTestTiling : public SupportTest,
-                      public ::testing::WithParamInterface<bool> {
- public:
-  bool EnableTilingPropagation() const override { return GetParam(); }
-};
-using DotTestSimple = DotTestTiling;
+using DotTestSimple = SupportTest;
 using DotTest = SupportTest;
 
 class DotTypesTest
     : public SupportTest,
-      public ::testing::WithParamInterface<std::tuple<
-          PrimitiveType, PrimitiveType, se::GpuComputeCapability, bool>> {
+      public ::testing::WithParamInterface<
+          std::tuple<PrimitiveType, PrimitiveType, se::GpuComputeCapability>> {
  public:
-  bool EnableTilingPropagation() const override {
-    return std::get<3>(GetParam());
-  }
   static std::string ParamToString(
       const ::testing::TestParamInfo<DotTypesTest::ParamType>& data) {
-    auto [result_type, input_type, cc, tiling] = data.param;
-    return absl::StrCat(
-        primitive_util::LowercasePrimitiveTypeName(result_type), "_",
-        primitive_util::LowercasePrimitiveTypeName(input_type), "_",
-        ComputeCapabilityToString(cc), TilingParametersToString(tiling));
+    auto [result_type, input_type, cc] = data.param;
+    return absl::StrCat(primitive_util::LowercasePrimitiveTypeName(result_type),
+                        "_",
+                        primitive_util::LowercasePrimitiveTypeName(input_type),
+                        "_", ComputeCapabilityToString(cc));
   };
 };
 
 TEST_P(DotTypesTest, Dot) {
   // Testing B[] = dot(A[], A[]).
-  auto [result_type, input_type, cc, tiling] = GetParam();
+  auto [result_type, input_type, cc] = GetParam();
 
   ExpectedFailMode fail_mode = ExpectedFailMode::kFail;
   if (input_type == F8E4M3FN || result_type == F8E4M3FN) {
@@ -2315,30 +2158,26 @@ INSTANTIATE_TEST_SUITE_P(
     ::testing::Combine(
         ::testing::ValuesIn(AllOpSupportedTypes(HloOpcode::kDot)),
         ::testing::ValuesIn(AllOpSupportedTypes(HloOpcode::kDot)),
-        ::testing::ValuesIn(AllDevicesToTest()), ::testing::Bool()),
+        ::testing::ValuesIn(AllDevicesToTest())),
     DotTypesTest::ParamToString);
 
 class MixedF8DotTest
     : public SupportTest,
-      public ::testing::WithParamInterface<std::tuple<
-          PrimitiveType, PrimitiveType, se::GpuComputeCapability, bool>> {
+      public ::testing::WithParamInterface<
+          std::tuple<PrimitiveType, PrimitiveType, se::GpuComputeCapability>> {
  public:
-  bool EnableTilingPropagation() const override {
-    return std::get<3>(GetParam());
-  }
-
   static std::string ParamToString(
       const ::testing::TestParamInfo<ParamType>& info) {
-    auto [lhs_type, rhs_type, cc, tiling] = info.param;
-    return absl::StrCat(
-        primitive_util::LowercasePrimitiveTypeName(lhs_type), "_",
-        primitive_util::LowercasePrimitiveTypeName(rhs_type), "_",
-        ComputeCapabilityToString(cc), TilingParametersToString(tiling));
+    auto [lhs_type, rhs_type, cc] = info.param;
+    return absl::StrCat(primitive_util::LowercasePrimitiveTypeName(lhs_type),
+                        "_",
+                        primitive_util::LowercasePrimitiveTypeName(rhs_type),
+                        "_", ComputeCapabilityToString(cc));
   }
 };
 
 TEST_P(MixedF8DotTest, MixedF8Dot) {
-  auto [lhs_type, rhs_type, cc, tiling] = GetParam();
+  auto [lhs_type, rhs_type, cc] = GetParam();
   if (lhs_type == rhs_type) {
     GTEST_SKIP() << "Skipping same-type tests, covered by DotTypesTest";
   }
@@ -2371,11 +2210,10 @@ INSTANTIATE_TEST_SUITE_P(
     MixedF8DotTestSuite, MixedF8DotTest,
     ::testing::Combine(::testing::Values(F8E5M2, F8E4M3FN),
                        ::testing::Values(F8E5M2, F8E4M3FN),
-                       ::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
+                       ::testing::ValuesIn(AllDevicesToTest())),
     MixedF8DotTest::ParamToString);
 
-TEST_P(DotTestSimple, MixedFnuzF8DotIsSupportedOnRocmAndRejectedOnCuda) {
+TEST_F(DotTestSimple, MixedFnuzF8DotIsSupportedOnRocmAndRejectedOnCuda) {
   const std::string kHloTestTemplate = R"(
 triton_computation {
   p0 = f8e4m3fnuz[128,256] parameter(0)
@@ -2401,7 +2239,7 @@ triton_computation {
                  se::GpuComputeCapability(se::CudaComputeCapability::Hopper()));
 }
 
-TEST_P(DotTestSimple, SingleBatchDim) {
+TEST_F(DotTestSimple, SingleBatchDim) {
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   p0 = $0[16,128,256] parameter(0)
@@ -2419,7 +2257,7 @@ ENTRY triton_computation {
                  DefaultDeviceForTesting());
 }
 
-TEST_P(DotTestSimple, MultipleNonContractingDimensions) {
+TEST_F(DotTestSimple, MultipleNonContractingDimensions) {
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   p0 = $0[16,128,256] parameter(0)
@@ -2436,7 +2274,7 @@ ENTRY triton_computation {
                  DefaultDeviceForTesting());
 }
 
-TEST_P(DotTestSimple, MultipleContractingDimensions) {
+TEST_F(DotTestSimple, MultipleContractingDimensions) {
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   p0 = $0[128,16,256] parameter(0)
@@ -2454,7 +2292,7 @@ ENTRY triton_computation {
                  DefaultDeviceForTesting());
 }
 
-TEST_P(DotTestSimple, NonDefaultDimensionOrder_kmkn) {
+TEST_F(DotTestSimple, NonDefaultDimensionOrder_kmkn) {
   // Multiplying as [k, m] x [k, n] = [m, n].
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
@@ -2473,7 +2311,7 @@ ENTRY triton_computation {
                  DefaultDeviceForTesting());
 }
 
-TEST_P(DotTestSimple, NonDefaultDimensionOrder_mknk) {
+TEST_F(DotTestSimple, NonDefaultDimensionOrder_mknk) {
   // Muliplying as [m, k] x [n, k] = [m, n].
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
@@ -2492,34 +2330,26 @@ ENTRY triton_computation {
                  DefaultDeviceForTesting());
 }
 
-INSTANTIATE_TEST_SUITE_P(DotTestSuiteSimple, DotTestSimple, ::testing::Bool(),
-                         ::testing::PrintToStringParamName());
-
 class DotPrecisionTest
     : public DotTest,
-      public ::testing::WithParamInterface<std::tuple<
-          PrimitiveType, PrecisionConfig::Precision, PrecisionConfig::Precision,
-          se::GpuComputeCapability, bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<4>(GetParam());
-  }
-};
+      public ::testing::WithParamInterface<
+          std::tuple<PrimitiveType, PrecisionConfig::Precision,
+                     PrecisionConfig::Precision, se::GpuComputeCapability>> {};
 
 std::string OperandPrecisionTestName(
     const ::testing::TestParamInfo<
         std::tuple<PrimitiveType, PrecisionConfig::Precision,
-                   PrecisionConfig::Precision, se::GpuComputeCapability, bool>>&
+                   PrecisionConfig::Precision, se::GpuComputeCapability>>&
         data) {
-  auto [type, lhs_precision, rhs_precision, cc, tiling] = data.param;
-  return absl::StrCat(
-      primitive_util::LowercasePrimitiveTypeName(type), "_",
-      PrecisionToString(lhs_precision), "_", PrecisionToString(rhs_precision),
-      "_", ComputeCapabilityToString(cc), TilingParametersToString(tiling));
+  auto [type, lhs_precision, rhs_precision, cc] = data.param;
+  return absl::StrCat(primitive_util::LowercasePrimitiveTypeName(type), "_",
+                      PrecisionToString(lhs_precision), "_",
+                      PrecisionToString(rhs_precision), "_",
+                      ComputeCapabilityToString(cc));
 }
 
 TEST_P(DotPrecisionTest, OperandPrecision) {
-  auto [data_type, lhs_precision, rhs_precision, cc, tiling] = GetParam();
+  auto [data_type, lhs_precision, rhs_precision, cc] = GetParam();
   std::string hlo_text = absl::Substitute(
       R"(
 ENTRY triton_computation {
@@ -2559,33 +2389,27 @@ INSTANTIATE_TEST_SUITE_P(
         ::testing::ValuesIn(AllOpSupportedTypes(HloOpcode::kDot)),
         ::testing::ValuesIn(kOperandPrecisions),
         ::testing::ValuesIn(kOperandPrecisions),
-        ::testing::ValuesIn(AllDevicesToTest()), ::testing::Bool()),
+        ::testing::ValuesIn(AllDevicesToTest())),
     OperandPrecisionTestName);
 
 class DotPrecisionAlgorithmTest
     : public DotTest,
       public ::testing::WithParamInterface<
           std::tuple<PrimitiveType, PrecisionConfig::Algorithm,
-                     se::GpuComputeCapability, bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<3>(GetParam());
-  }
-};
+                     se::GpuComputeCapability>> {};
 
 std::string DotPrecisionAlgorithmTestName(
-    const ::testing::TestParamInfo<
-        std::tuple<PrimitiveType, PrecisionConfig::Algorithm,
-                   se::GpuComputeCapability, bool>>& data) {
-  auto [type, algorigthm, cc, tiling] = data.param;
+    const ::testing::TestParamInfo<std::tuple<
+        PrimitiveType, PrecisionConfig::Algorithm, se::GpuComputeCapability>>&
+        data) {
+  auto [type, algorigthm, cc] = data.param;
   return absl::StrCat(primitive_util::LowercasePrimitiveTypeName(type), "_",
                       AlgorithmToString(algorigthm), "_",
-                      ComputeCapabilityToString(cc),
-                      TilingParametersToString(tiling));
+                      ComputeCapabilityToString(cc));
 }
 
 TEST_P(DotPrecisionAlgorithmTest, Algorithm) {
-  auto [data_type, algorithm, cc, tiling] = GetParam();
+  auto [data_type, algorithm, cc] = GetParam();
   std::string hlo_text =
       absl::Substitute(R"(
 ENTRY triton_computation {
@@ -2625,20 +2449,14 @@ INSTANTIATE_TEST_SUITE_P(
     ::testing::Combine(
         ::testing::ValuesIn(AllOpSupportedTypes(HloOpcode::kDot)),
         ::testing::ValuesIn(AllPrecisionAlgorithms()),
-        ::testing::ValuesIn(AllDevicesToTest()), ::testing::Bool()),
+        ::testing::ValuesIn(AllDevicesToTest())),
     DotPrecisionAlgorithmTestName);
 
-class ScaledDotTest
-    : public SupportTest,
-      public ::testing::WithParamInterface<std::tuple<PrimitiveType, bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<1>(GetParam());
-  }
-};
+class ScaledDotTest : public SupportTest,
+                      public ::testing::WithParamInterface<PrimitiveType> {};
 
 TEST_P(ScaledDotTest, ScaledDotOperandTypes) {
-  auto [type, tiling] = GetParam();
+  auto type = GetParam();
   const std::string kHloTestTemplate = R"(
 HloModule ScaledDotOperandTypes
 
@@ -2661,7 +2479,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(ScaledDotTest, ScaledDotScaleTypes) {
-  auto [type, tiling] = GetParam();
+  auto type = GetParam();
   const std::string kHloTestTemplate = R"(
 HloModule ScaledDotOperandTypes
 
@@ -2685,16 +2503,12 @@ ENTRY triton_computation {
 
 std::string ScaledDotTestName(
     const ::testing::TestParamInfo<ScaledDotTest::ParamType>& info) {
-  auto [type, tiling] = info.param;
-  return absl::StrCat(primitive_util::LowercasePrimitiveTypeName(type),
-                      TilingParametersToString(tiling));
+  return std::string(primitive_util::LowercasePrimitiveTypeName(info.param));
 }
 
 INSTANTIATE_TEST_SUITE_P(
     ScaledDotTest, ScaledDotTest,
-    ::testing::Combine(
-        ::testing::ValuesIn(AllOpSupportedTypes(HloOpcode::kScaledDot)),
-        ::testing::Bool()),
+    ::testing::ValuesIn(AllOpSupportedTypes(HloOpcode::kScaledDot)),
     ScaledDotTestName);
 
 TEST_F(HloHardwareIndependentTestBase, ScaledDotPreAmpereIsRejected) {
@@ -2718,7 +2532,7 @@ ENTRY triton_computation {
       *dot, se::GpuComputeCapability(se::CudaComputeCapability::Volta())));
 }
 
-TEST_P(SupportTestWithTilingParam, NestedFusionsAreRejected) {
+TEST_F(SupportTest, NestedFusionsAreRejected) {
   // Nested fusions are not supported by xtile emitter.
   absl::string_view hlo_text = R"(
 lhs {
@@ -2757,16 +2571,12 @@ ENTRY entry {
 
 class BitcastConvertTest
     : public SupportTest,
-      public ::testing::WithParamInterface<std::tuple<
-          PrimitiveType, PrimitiveType, se::GpuComputeCapability, bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<3>(GetParam());
-  }
+      public ::testing::WithParamInterface<
+          std::tuple<PrimitiveType, PrimitiveType, se::GpuComputeCapability>> {
 };
 
 TEST_P(BitcastConvertTest, BitcastConvert) {
-  auto [data_type_in, data_type_out, cc, tiling] = GetParam();
+  auto [data_type_in, data_type_out, cc] = GetParam();
 
   if (primitive_util::IsComplexType(data_type_in) !=
       primitive_util::IsComplexType(data_type_out)) {
@@ -2819,7 +2629,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(BitcastConvertTest, BitcastConvertDisguisedAsBitcast) {
-  auto [data_type_in, data_type_out, cc, tiling] = GetParam();
+  auto [data_type_in, data_type_out, cc] = GetParam();
 
   if (primitive_util::IsComplexType(data_type_in) !=
       primitive_util::IsComplexType(data_type_out)) {
@@ -2872,25 +2682,23 @@ ENTRY triton_computation {
 
 std::string BitcastConvertTestName(
     const ::testing::TestParamInfo<BitcastConvertTest::ParamType>& info) {
-  auto [in_type, out_type, cc, tiling] = info.param;
+  auto [in_type, out_type, cc] = info.param;
   return absl::StrCat(primitive_util::LowercasePrimitiveTypeName(in_type), "_",
                       primitive_util::LowercasePrimitiveTypeName(out_type), "_",
-                      ComputeCapabilityToString(cc),
-                      TilingParametersToString(tiling));
+                      ComputeCapabilityToString(cc));
 }
 
 INSTANTIATE_TEST_SUITE_P(
     BitcastConvertSuite, BitcastConvertTest,
     ::testing::Combine(::testing::ValuesIn(AllXlaDataTypes()),
                        ::testing::ValuesIn(AllXlaDataTypes()),
-                       ::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
+                       ::testing::ValuesIn(AllDevicesToTest())),
     BitcastConvertTestName);
 
-using AddDependencyTest = SupportTestWithDeviceAndTilingParam;
+using AddDependencyTest = SupportTestWithDeviceParam;
 
 TEST_P(AddDependencyTest, AddDependency) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   param = f32[10] parameter(0)
@@ -2904,16 +2712,14 @@ ENTRY triton_computation {
   RunSupportTest(std::move(ti), /*output_tile_sizes=*/{1}, cc);
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    AddDependencySuite, AddDependencyTest,
-    ::testing::Combine(::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(AddDependencySuite, AddDependencyTest,
+                         ::testing::ValuesIn(AllDevicesToTest()),
+                         SupportTestDeviceToString);
 
-using AfterAllTest = SupportTestWithDeviceAndTilingParam;
+using AfterAllTest = SupportTestWithDeviceParam;
 
 TEST_P(AfterAllTest, AfterAll) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   token0 = token[] after-all()
@@ -2927,16 +2733,14 @@ ENTRY triton_computation {
   RunSupportTest(std::move(ti), /*output_tile_sizes=*/{}, cc);
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    AfterAllSuite, AfterAllTest,
-    ::testing::Combine(::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(AfterAllSuite, AfterAllTest,
+                         ::testing::ValuesIn(AllDevicesToTest()),
+                         SupportTestDeviceToString);
 
-using TupleTest = SupportTestWithDeviceAndTilingParam;
+using TupleTest = SupportTestWithDeviceParam;
 
 TEST_P(TupleTest, Tuple) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   p0 = f32[10] parameter(0)
@@ -2951,16 +2755,14 @@ ENTRY triton_computation {
                                     /*output_tile_sizes=*/{{1}, {1}}, cc);
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    TupleSuite, TupleTest,
-    ::testing::Combine(::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(TupleSuite, TupleTest,
+                         ::testing::ValuesIn(AllDevicesToTest()),
+                         SupportTestDeviceToString);
 
-using GetTupleElementTest = SupportTestWithDeviceAndTilingParam;
+using GetTupleElementTest = SupportTestWithDeviceParam;
 
 TEST_P(GetTupleElementTest, GetTupleElement) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   tuple_op = (f32[10], s32[5]) parameter(0)
@@ -2973,16 +2775,14 @@ ENTRY triton_computation {
   RunSupportTest(std::move(ti), /*output_tile_sizes=*/{1}, cc);
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    GetTupleElementSuite, GetTupleElementTest,
-    ::testing::Combine(::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(GetTupleElementSuite, GetTupleElementTest,
+                         ::testing::ValuesIn(AllDevicesToTest()),
+                         SupportTestDeviceToString);
 
-using CustomCallTest = SupportTestWithDeviceAndTilingParam;
+using CustomCallTest = SupportTestWithDeviceParam;
 
 TEST_P(CustomCallTest, CustomCall) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   parameter = f32[10] parameter(0)
@@ -2995,25 +2795,18 @@ ENTRY triton_computation {
   RunSupportTest(std::move(ti), /*output_tile_sizes=*/{1}, cc);
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    CustomCallSuite, CustomCallTest,
-    ::testing::Combine(::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(CustomCallSuite, CustomCallTest,
+                         ::testing::ValuesIn(AllDevicesToTest()),
+                         SupportTestDeviceToString);
 
 class CholeskyTest
     : public SupportTest,
       public ::testing::WithParamInterface<
-          // The bool parameter is used to parametrize the lower=?, and tiling.
-          std::tuple<PrimitiveType, se::GpuComputeCapability, bool, bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<3>(GetParam());
-  }
-};
+          // The bool parameter is used to parametrize the lower=?.
+          std::tuple<PrimitiveType, se::GpuComputeCapability, bool>> {};
 
 TEST_P(CholeskyTest, Cholesky) {
-  auto [data_type, cc, lower, tiling] = GetParam();
+  auto [data_type, cc, lower] = GetParam();
 
   const std::string kHloTestTemplate = absl::Substitute(
       R"(
@@ -3031,42 +2824,33 @@ ENTRY triton_computation {
 
 std::string CholeskyTestName(
     const ::testing::TestParamInfo<
-        std::tuple<PrimitiveType, se::GpuComputeCapability, bool, bool>>&
-        data) {
-  const auto [data_type, cc, lower, tiling] = data.param;
+        std::tuple<PrimitiveType, se::GpuComputeCapability, bool>>& data) {
+  const auto [data_type, cc, lower] = data.param;
   return absl::StrCat(primitive_util::LowercasePrimitiveTypeName(data_type),
-                      "_", ComputeCapabilityToString(cc), "_", lower,
-                      tiling ? "_TilingPropagation" : "_SymbolicAnalysis");
+                      "_", ComputeCapabilityToString(cc), "_", lower);
 }
 
 INSTANTIATE_TEST_SUITE_P(
     CholeskySuite, CholeskyTest,
     ::testing::Combine(
         ::testing::ValuesIn(AllOpSupportedTypes(HloOpcode::kCholesky)),
-        ::testing::ValuesIn(AllDevicesToTest()), ::testing::Bool(),
-        ::testing::Bool()),
+        ::testing::ValuesIn(AllDevicesToTest()), ::testing::Bool()),
     CholeskyTestName);
 
 class TriangularSolveParamTest
     : public SupportTest,
-      public ::testing::WithParamInterface<std::tuple<
-          PrimitiveType, se::GpuComputeCapability, bool /*lower*/,
-          bool /*unit_diagonal*/,
-          TriangularSolveOptions::Transpose /*transpose_a*/, bool /*tiling*/>> {
+      public ::testing::WithParamInterface<
+          std::tuple<PrimitiveType, se::GpuComputeCapability, bool /*lower*/,
+                     bool /*unit_diagonal*/,
+                     TriangularSolveOptions::Transpose /*transpose_a*/>> {
  public:
-  bool EnableTilingPropagation() const override {
-    return std::get<5>(GetParam());
-  }
-
   static std::string ParamToString(
       const ::testing::TestParamInfo<ParamType>& info) {
-    auto [data_type, cc, lower, unit_diagonal, transpose_a, tiling] =
-        info.param;
+    auto [data_type, cc, lower, unit_diagonal, transpose_a] = info.param;
     return absl::StrCat(primitive_util::LowercasePrimitiveTypeName(data_type),
                         "_", ComputeCapabilityToString(cc), "_lower", lower,
                         "_unitdiag", unit_diagonal, "_",
-                        TriangularSolveOptions::Transpose_Name(transpose_a),
-                        TilingParametersToString(tiling));
+                        TriangularSolveOptions::Transpose_Name(transpose_a));
   }
 };
 
@@ -3079,13 +2863,11 @@ INSTANTIATE_TEST_SUITE_P(
         ::testing::Bool(),  // unit_diagonal
         ::testing::ValuesIn({TriangularSolveOptions::NO_TRANSPOSE,
                              TriangularSolveOptions::TRANSPOSE,
-                             TriangularSolveOptions::ADJOINT}),
-        ::testing::Bool()  // tiling
-        ),
+                             TriangularSolveOptions::ADJOINT})),
     TriangularSolveParamTest::ParamToString);
 
 TEST_P(TriangularSolveParamTest, TriangularSolveLeftSideTrue) {
-  auto [data_type, cc, lower, unit_diagonal, transpose_a, tiling] = GetParam();
+  auto [data_type, cc, lower, unit_diagonal, transpose_a] = GetParam();
 
   const std::string hlo_text = absl::Substitute(
       R"(
@@ -3106,7 +2888,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(TriangularSolveParamTest, TriangularSolveLeftSideFalse) {
-  auto [data_type, cc, lower, unit_diagonal, transpose_a, tiling] = GetParam();
+  auto [data_type, cc, lower, unit_diagonal, transpose_a] = GetParam();
 
   const std::string hlo_text = absl::Substitute(
       R"(
@@ -3126,10 +2908,10 @@ ENTRY triton_computation {
   RunSupportTest(std::move(ti), {1, 1, 2}, cc);
 }
 
-using FftTest = SupportTestWithTypeAndDeviceAndTilingParam;
+using FftTest = SupportTestWithTypeAndDeviceParam;
 
 TEST_P(FftTest, FFT) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
 
   const std::string hlo_text = R"(
 ENTRY triton_computation {
@@ -3145,7 +2927,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(FftTest, IFFT) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
 
   const std::string hlo_text = R"(
 ENTRY triton_computation {
@@ -3161,7 +2943,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(FftTest, RFFT) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string complex_data_type_str =
       primitive_util::LowercasePrimitiveTypeName(data_type);
   // Real type matching the complex type for real -> complex conversion.
@@ -3185,7 +2967,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(FftTest, IRFFT) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string complex_data_type_str =
       primitive_util::LowercasePrimitiveTypeName(data_type);
   // Real type matching the complex type for complex -> real conversion.
@@ -3214,14 +2996,13 @@ INSTANTIATE_TEST_SUITE_P(
     // complex <-> real conversion, the real type can be directly inferred from
     // the complex type (C64 <-> F32, C128 <-> F64).
     ::testing::Combine(::testing::ValuesIn({C64, C128}),
-                       ::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestTypeAndDeviceAndTilingToString);
+                       ::testing::ValuesIn(AllDevicesToTest())),
+    SupportTestTypeAndDeviceToString);
 
-using CopyStartDoneTest = SupportTestWithTypeAndDeviceAndTilingParam;
+using CopyStartDoneTest = SupportTestWithTypeAndDeviceParam;
 
 TEST_P(CopyStartDoneTest, CopyStartDone) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   parameter = $0[10,10,10] parameter(0)
@@ -3245,14 +3026,13 @@ constexpr std::array kTestedOpsCopy = {HloOpcode::kCopyStart,
 INSTANTIATE_TEST_SUITE_P(
     CopyStartDoneSuite, CopyStartDoneTest,
     ::testing::Combine(::testing::ValuesIn(AllXlaDataTypes()),
-                       ::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestTypeAndDeviceAndTilingToString);
+                       ::testing::ValuesIn(AllDevicesToTest())),
+    SupportTestTypeAndDeviceToString);
 
-using InfeedTest = SupportTestWithTypeAndDeviceAndTilingParam;
+using InfeedTest = SupportTestWithTypeAndDeviceParam;
 
 TEST_P(InfeedTest, Infeed) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   token0 = token[] after-all()
@@ -3268,14 +3048,13 @@ ENTRY triton_computation {
 INSTANTIATE_TEST_SUITE_P(
     InfeedSuite, InfeedTest,
     ::testing::Combine(::testing::ValuesIn(AllXlaDataTypes()),
-                       ::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestTypeAndDeviceAndTilingToString);
+                       ::testing::ValuesIn(AllDevicesToTest())),
+    SupportTestTypeAndDeviceToString);
 
-using OutfeedTest = SupportTestWithTypeAndDeviceAndTilingParam;
+using OutfeedTest = SupportTestWithTypeAndDeviceParam;
 
 TEST_P(OutfeedTest, Outfeed) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   data = $0[10] parameter(0)
@@ -3288,10 +3067,10 @@ ENTRY triton_computation {
   RunSupportTest(std::move(ti), /*output_tile_sizes=*/{}, cc);
 }
 
-using MapTest = SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using MapTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(MapTest, Map) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
 
   // Note: the test is only relevant for datatypes supported by kAdd op.
   const std::string kHloTestTemplate = R"(
@@ -3310,22 +3089,20 @@ ENTRY triton_computation {
   RunSupportTest(std::move(ti), /*output_tile_sizes=*/{4, 8}, cc);
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    MapSuite, MapTest,
-    AllTestCombinationsForOpcodesWithTiling({HloOpcode::kMap}),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(MapSuite, MapTest,
+                         AllTestCombinationsForOpcodes({HloOpcode::kMap}),
+                         SupportTestTypeAndOpcodeAndDeviceToString);
 
 INSTANTIATE_TEST_SUITE_P(
     OutfeedSuite, OutfeedTest,
     ::testing::Combine(::testing::ValuesIn(AllXlaDataTypes()),
-                       ::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestTypeAndDeviceAndTilingToString);
+                       ::testing::ValuesIn(AllDevicesToTest())),
+    SupportTestTypeAndDeviceToString);
 
-using SortTest = SupportTestWithTypeAndOpcodeAndDeviceAndTilingParam;
+using SortTest = SupportTestWithTypeAndOpcodeAndDeviceParam;
 
 TEST_P(SortTest, Sort) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
 
   const std::string kHloTestTemplate = R"(
 compare {
@@ -3345,7 +3122,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(SortTest, KeyValueSort) {
-  auto [data_type, opcode, cc, tiling] = GetParam();
+  auto [data_type, opcode, cc] = GetParam();
 
   const std::string kHloTestTemplate = R"(
 compare {
@@ -3368,15 +3145,14 @@ ENTRY triton_computation {
                                     /*output_tile_sizes=*/{{2, 4}, {2, 4}}, cc);
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    SortSuite, SortTest,
-    AllTestCombinationsForOpcodesWithTiling({HloOpcode::kSort}),
-    SupportTestTypeAndOpcodeAndDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(SortSuite, SortTest,
+                         AllTestCombinationsForOpcodes({HloOpcode::kSort}),
+                         SupportTestTypeAndOpcodeAndDeviceToString);
 
-using RecvOpsTest = SupportTestWithTypeAndDeviceAndTilingParam;
+using RecvOpsTest = SupportTestWithTypeAndDeviceParam;
 
 TEST_P(RecvOpsTest, RecvAndRecvDone) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
   ENTRY triton_computation {
     token0 = token[] after-all()
@@ -3400,14 +3176,13 @@ constexpr std::array kTestedOpsRecv = {HloOpcode::kRecv, HloOpcode::kRecvDone};
 INSTANTIATE_TEST_SUITE_P(
     RecvOpsSuite, RecvOpsTest,
     ::testing::Combine(::testing::ValuesIn(AllXlaDataTypes()),
-                       ::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestTypeAndDeviceAndTilingToString);
+                       ::testing::ValuesIn(AllDevicesToTest())),
+    SupportTestTypeAndDeviceToString);
 
-using SendOpsTest = SupportTestWithTypeAndDeviceAndTilingParam;
+using SendOpsTest = SupportTestWithTypeAndDeviceParam;
 
 TEST_P(SendOpsTest, SendAndSendDone) {
-  auto [data_type, cc, tiling] = GetParam();
+  auto [data_type, cc] = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   data = $0[10] parameter(0)
@@ -3432,22 +3207,17 @@ constexpr std::array kTestedOpsSend = {HloOpcode::kSend, HloOpcode::kSendDone};
 INSTANTIATE_TEST_SUITE_P(
     SendOpsSuite, SendOpsTest,
     ::testing::Combine(::testing::ValuesIn(AllXlaDataTypes()),
-                       ::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestTypeAndDeviceAndTilingToString);
+                       ::testing::ValuesIn(AllDevicesToTest())),
+    SupportTestTypeAndDeviceToString);
 
 class StochasticConvertTest
     : public SupportTest,
-      public ::testing::WithParamInterface<std::tuple<
-          PrimitiveType, PrimitiveType, se::GpuComputeCapability, bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<3>(GetParam());
-  }
+      public ::testing::WithParamInterface<
+          std::tuple<PrimitiveType, PrimitiveType, se::GpuComputeCapability>> {
 };
 
 TEST_P(StochasticConvertTest, StochasticConvert) {
-  auto [operand_type, new_element_type, cc, tiling] = GetParam();
+  auto [operand_type, new_element_type, cc] = GetParam();
 
   PrimitiveType random_type = primitive_util::UnsignedIntegralTypeForBitWidth(
       primitive_util::BitWidth(operand_type));
@@ -3477,12 +3247,11 @@ ENTRY triton_computation {
 
 std::string StochasticConvertTestName(
     const ::testing::TestParamInfo<StochasticConvertTest::ParamType>& info) {
-  auto [operand_type, new_element_type, cc, tiling] = info.param;
+  auto [operand_type, new_element_type, cc] = info.param;
   return absl::StrCat(
       primitive_util::LowercasePrimitiveTypeName(operand_type), "_",
       primitive_util::LowercasePrimitiveTypeName(new_element_type), "_",
-      ComputeCapabilityToString(cc),
-      tiling ? "_TilingPropagation" : "_SymbolicAnalysis");
+      ComputeCapabilityToString(cc));
 }
 
 INSTANTIATE_TEST_SUITE_P(
@@ -3491,30 +3260,24 @@ INSTANTIATE_TEST_SUITE_P(
         ::testing::ValuesIn(AllOpSupportedTypes(
             HloOpcode::kStochasticConvert)),     // Operand type.
         ::testing::ValuesIn(AllXlaDataTypes()),  // New element type.
-        ::testing::ValuesIn(AllDevicesToTest()), ::testing::Bool()),
+        ::testing::ValuesIn(AllDevicesToTest())),
     StochasticConvertTestName);
 
 class TopKTest
     : public SupportTest,
       public ::testing::WithParamInterface<
-          std::tuple<PrimitiveType, se::GpuComputeCapability, bool, bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<3>(GetParam());
-  }
-};
+          std::tuple<PrimitiveType, se::GpuComputeCapability, bool>> {};
 
 std::string ParamToStringTopK(
     const ::testing::TestParamInfo<TopKTest::ParamType>& info) {
-  auto [data_type, cc, largest, tiling] = info.param;
+  auto [data_type, cc, largest] = info.param;
   return absl::StrCat(PrimitiveType_Name(data_type), "_",
                       ComputeCapabilityToString(cc), "_",
-                      largest ? "largest" : "smallest",
-                      tiling ? "_TilingPropagation" : "_SymbolicAnalysis");
+                      largest ? "largest" : "smallest");
 }
 
 TEST_P(TopKTest, TopK) {
-  auto [data_type, cc, largest, tiling] = GetParam();
+  auto [data_type, cc, largest] = GetParam();
   const std::string kHloTestTemplate = absl::Substitute(
       R"(
 ENTRY triton_computation {
@@ -3534,24 +3297,18 @@ INSTANTIATE_TEST_SUITE_P(
     TopKSuite, TopKTest,
     ::testing::Combine(
         ::testing::ValuesIn(AllOpSupportedTypes(HloOpcode::kTopK)),
-        ::testing::ValuesIn(AllDevicesToTest()), ::testing::Bool(),
-        ::testing::Bool()),
+        ::testing::ValuesIn(AllDevicesToTest()), ::testing::Bool()),
     ParamToStringTopK);
 
 class ConvolutionTestFullParametrization
     : public SupportTest,
-      public ::testing::WithParamInterface<std::tuple<
-          PrimitiveType, PrecisionConfig::Precision, PrecisionConfig::Precision,
-          se::GpuComputeCapability, bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<4>(GetParam());
-  }
-};
+      public ::testing::WithParamInterface<
+          std::tuple<PrimitiveType, PrecisionConfig::Precision,
+                     PrecisionConfig::Precision, se::GpuComputeCapability>> {};
 
 // Test a basic convolution (NHWC layout)
 TEST_P(ConvolutionTestFullParametrization, IsTritonSupportedConvNHWC) {
-  auto [data_type, input_precision, kernel_precision, cc, tiling] = GetParam();
+  auto [data_type, input_precision, kernel_precision, cc] = GetParam();
   const std::string kHloTestTemplate = absl::Substitute(
       R"(
 ENTRY triton_computation {
@@ -3571,7 +3328,7 @@ ENTRY triton_computation {
 
 // Test a convolution with NCHW layout
 TEST_P(ConvolutionTestFullParametrization, IsTritonSupportedConvNCHW) {
-  auto [data_type, input_precision, kernel_precision, cc, tiling] = GetParam();
+  auto [data_type, input_precision, kernel_precision, cc] = GetParam();
   const std::string kHloTestTemplate = absl::Substitute(
       R"(
 ENTRY triton_computation {
@@ -3595,14 +3352,14 @@ INSTANTIATE_TEST_SUITE_P(
         ::testing::ValuesIn(AllOpSupportedTypes(HloOpcode::kConvolution)),
         ::testing::ValuesIn(kOperandPrecisions),
         ::testing::ValuesIn(kOperandPrecisions),
-        ::testing::ValuesIn(AllDevicesToTest()), ::testing::Bool()),
+        ::testing::ValuesIn(AllDevicesToTest())),
     OperandPrecisionTestName);
 
-using ConvolutionTestCcOnly = SupportTestWithDeviceAndTilingParam;
+using ConvolutionTestCcOnly = SupportTestWithDeviceParam;
 
 // Test a convolution with stride > 1
 TEST_P(ConvolutionTestCcOnly, IsTritonSupportedConvStrided) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   input = f16[1,5,6,2] parameter(0)  // N=1, H=5, W=6, C_in=2
@@ -3619,7 +3376,7 @@ ENTRY triton_computation {
 
 // Test a convolution with kernel dilation > 1
 TEST_P(ConvolutionTestCcOnly, IsTritonSupportedConvDilated) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   input = f16[1,5,6,2] parameter(0)  // N=1, H=5, W=6, C_in=2
@@ -3636,7 +3393,7 @@ ENTRY triton_computation {
 
 // Test a depthwise convolution
 TEST_P(ConvolutionTestCcOnly, IsTritonSupportedConvDepthwise) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   input = f16[1,5,6,2] parameter(0)    // N=1, H=5, W=6, C_in=1
@@ -3655,7 +3412,7 @@ ENTRY triton_computation {
 
 // Test a convolution with batch_group_count > 1
 TEST_P(ConvolutionTestCcOnly, IsTritonSupportedConvBatchGrouped) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   input = f16[4,5,6,2] parameter(0)  // N=4, H=5, W=6, C_in=2
@@ -3673,7 +3430,7 @@ ENTRY triton_computation {
 
 // Test a convolution with lhs_dilation > 1
 TEST_P(ConvolutionTestCcOnly, IsTritonSupportedConvLhsDilation) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   input = f16[1,5,6,2] parameter(0)  // N=1, H=5, W=6, C_in=2
@@ -3690,7 +3447,7 @@ ENTRY triton_computation {
 
 // Test a convolution with asymmetric padding
 TEST_P(ConvolutionTestCcOnly, IsTritonSupportedConvAsymmetricPadding) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   input = f16[1,5,6,2] parameter(0)  // N=1, H=5, W=6, C_in=2
@@ -3707,7 +3464,7 @@ ENTRY triton_computation {
 
 // Test a general grouped convolution (1 < feature_group_count < Cin)
 TEST_P(ConvolutionTestCcOnly, IsTritonSupportedConvGeneralGrouped) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   input = f16[1,5,6,4] parameter(0)  // N=1, H=5, W=6, C_in=4
@@ -3727,7 +3484,7 @@ ENTRY triton_computation {
 // dims, requiring padding.
 TEST_P(ConvolutionTestCcOnly,
        IsTritonSupportedConvKernelLargerThanInputPadded) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   input = f16[1,2,2,2] parameter(0)  // N=1, H=2, W=2, C_in=2
@@ -3743,7 +3500,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(ConvolutionTestCcOnly, IsTritonSupportedConvDifferentWindowPadding) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   input = f16[1,5,6,2] parameter(0)  // N=1, H=5, W=6, C_in=2
@@ -3759,7 +3516,7 @@ ENTRY triton_computation {
 }
 
 TEST_P(ConvolutionTestCcOnly, IsTritonSupportedConvDifferentWindowSize) {
-  auto [cc, tiling] = GetParam();
+  auto cc = GetParam();
   const std::string kHloTestTemplate = R"(
 ENTRY triton_computation {
   input = f16[1,5,6,2] parameter(0)  // N=1, H=5, W=6, C_in=2
@@ -3774,11 +3531,9 @@ ENTRY triton_computation {
   RunSupportTest(std::move(ti), /*output_tile_sizes=*/{1, 2, 2, 1}, cc);
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    ConvolutionTestSuiteCcOnly, ConvolutionTestCcOnly,
-    ::testing::Combine(::testing::ValuesIn(AllDevicesToTest()),
-                       ::testing::Bool()),
-    SupportTestDeviceAndTilingToString);
+INSTANTIATE_TEST_SUITE_P(ConvolutionTestSuiteCcOnly, ConvolutionTestCcOnly,
+                         ::testing::ValuesIn(AllDevicesToTest()),
+                         SupportTestDeviceToString);
 
 // This denotes opcodes that are explicitly not supported by the generic Triton
 // emitter, and have also not been tested at all in the support test. If you
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD b/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD
index 21e32f5e981..e8875921477 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/BUILD
@@ -66,12 +66,18 @@ xla_cc_binary(
 )
 
 deviceless_tests = [
+    "bitcast/bitcast_reduce_stride1.hlo",
     "collectives/all_reduce.hlo",
     "collectives/all_gather.hlo",
     "concat/concat_active_scalar_constant.hlo",
     "concat/concat_offset_not_a_multiple_of_tile_size.hlo",
+    "concat/concat_prune_inactive.hlo",
     "constraints/max_tensor_elements_violation.hlo",
     "dot/duplicate_dot_operand.hlo",
+    "reduce/reduce_minor_extra_out.hlo",
+    "reduce/reduce_to_scalar_extra_out.hlo",
+    "slice/slice_extra_out_padding.hlo",
+    "slice/slice_padding.hlo",
 ]
 
 lit_test_suite(
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_dynamic_slice.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_dynamic_slice.hlo
index af6abaa23ad..f34cfd813ee 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_dynamic_slice.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_dynamic_slice.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 HloModule m
 
 fusion (p0: f32[1024, 16], offset: s32[]) -> f32[8192] {
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_tiled_1.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_tiled_1.hlo
index ef302c3620a..1ba74831709 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_tiled_1.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_tiled_1.hlo
@@ -1,5 +1,5 @@
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s --check-prefixes=CHECK
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = f32[16] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_tiled_4.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_tiled_4.hlo
index 79ddc336308..7d5d34afb49 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_tiled_4.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_tiled_4.hlo
@@ -1,5 +1,5 @@
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s --check-prefixes=CHECK
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = f32[16] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_transposed.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_transposed.hlo
index ac17b1d58d0..cf1f517ff11 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_transposed.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_expand_transposed.hlo
@@ -1,5 +1,5 @@
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = f32[64,16]{1,0} parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_non_norm_in.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_non_norm_in.hlo
index e157b050196..53b82005dd3 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_non_norm_in.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_non_norm_in.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = s32[48,16]{0,1} parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_non_norm_in_out.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_non_norm_in_out.hlo
index b602c033ed0..a78efc936e2 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_non_norm_in_out.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_non_norm_in_out.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = s8[42,5]{0,1} parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_non_norm_out.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_non_norm_out.hlo
index 9bc692e2a65..e9eb5dbb1c9 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_non_norm_out.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_non_norm_out.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = s8[5,42] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_normalized.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_normalized.hlo
index 63204afc8d2..04d4344c826 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_normalized.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_normalized.hlo
@@ -1,23 +1,15 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s --check-prefixes=CHECK,CHECK-OLD
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s --check-prefixes=CHECK,CHECK-NEW
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = f32[8,48]{1,0} parameter(0)
   ROOT bitcast = f32[8,3,16] bitcast(p)
 }
 
-// CHECK-OLD: #[[$ATTR_0:.+]] = #xla.indexing_map<"(pid_0) -> ((pid_0 / 6) * 2), domain: pid_0 in [0, 23]">
-// CHECK-OLD: #[[$ATTR_1:.+]] = #xla.indexing_map<"(pid_0) -> (((pid_0 / 2) mod 3) * 16 + (pid_0 mod 2) * 8), domain: pid_0 in [0, 23]">
-// CHECK-OLD: #[[$ATTR_2:.+]] = #xla.indexing_map<"(pid_0) -> ((pid_0 / 2) mod 3), domain: pid_0 in [0, 23]">
-// CHECK-OLD: #[[$ATTR_3:.+]] = #xla.indexing_map<"(pid_0) -> ((pid_0 mod 2) * 8), domain: pid_0 in [0, 23]">
-
-// CHECK-NEW: #[[$ATTR_0:.+]] = #xla.indexing_map<"(pid) -> ((pid / 6) * 2), domain: pid in [0, 23]">
-// CHECK-NEW: #[[$ATTR_1:.+]] = #xla.indexing_map<"(pid) -> ((pid mod 6) * 8), domain: pid in [0, 23]">
-// CHECK-NEW: #[[$ATTR_2:.+]] = #xla.indexing_map<"(pid) -> ((pid / 2) mod 3), domain: pid in [0, 23]">
-// CHECK-NEW: #[[$ATTR_3:.+]] = #xla.indexing_map<"(pid) -> ((pid mod 2) * 8), domain: pid in [0, 23]">
+// CHECK: #[[$ATTR_0:.+]] = #xla.indexing_map<"(pid) -> ((pid / 6) * 2), domain: pid in [0, 23]">
+// CHECK: #[[$ATTR_1:.+]] = #xla.indexing_map<"(pid) -> ((pid mod 6) * 8), domain: pid in [0, 23]">
+// CHECK: #[[$ATTR_2:.+]] = #xla.indexing_map<"(pid) -> ((pid / 2) mod 3), domain: pid in [0, 23]">
+// CHECK: #[[$ATTR_3:.+]] = #xla.indexing_map<"(pid) -> ((pid mod 2) * 8), domain: pid in [0, 23]">
 
 // CHECK-LABEL:   xtile.entry_func @triton_fn(
 // CHECK-SAME:      %[[ARG0:.*]]: memref<8x48xf32>,
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reduce_stride1.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reduce_stride1.hlo
index 6b902cc6152..cb96399ce98 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reduce_stride1.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reduce_stride1.hlo
@@ -1,8 +1,8 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false --xla_gpu_unsupported_enable_triton_multi_output_fusion %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
+// RUN: fusion_to_triton --xla_gpu_unsupported_enable_triton_multi_output_fusion %s 2>&1 | FileCheck %s --check-prefixes=CHECK-ERROR
+// CHECK-ERROR: Only single-result fusions are supported for now. Received 2 roots.
 // TODO: b/502910372 - support multi-output fusion.
-// RUN-NOT: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation --xla_gpu_unsupported_enable_triton_multi_output_fusion %s | FileCheck %s
-// RUN-NOT: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
+// RUN-NOT: fusion_to_triton --xla_gpu_unsupported_enable_triton_multi_output_fusion %s | FileCheck %s
+// RUN-NOT: triton_test_correctness --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
 
 region {
   param_0.1 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_collapse.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_collapse.hlo
index 67b119916a6..5870b830854 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_collapse.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_collapse.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s --check-prefixes=CHECK,CHECK-OLD
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s --check-prefixes=CHECK,CHECK-NEW
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = f32[4096,1,768]{2,1,0} parameter(0)
@@ -18,8 +15,7 @@ fusion {
 // CHECK-SAME:      %[[ARG2:.*]]: index{{( \{xla.range = \[0 : index, 255 : index\]\})?}}) attributes {num_opaque_args = 0 : i32} {
 // CHECK:           %[[CONSTANT_0:.*]] = arith.constant 0 : index
 // CHECK:           %[[APPLY_INDEXING_0:.*]] = xla.apply_indexing #[[$ATTR_0]](%[[ARG2]])
-// CHECK-OLD:           %[[EXTRACT_0:.*]] = xtile.extract %[[ARG0]][%[[APPLY_INDEXING_0]], %[[CONSTANT_0]], %[[CONSTANT_0]]] [16, 1, 1024] [128, 0, 1] : memref<4096x1x768xf32> -> tensor<16x1x1024xf32>
-// CHECK-NEW:           %[[EXTRACT_0:.*]] = xtile.extract %[[ARG0]][%[[APPLY_INDEXING_0]], %[[CONSTANT_0]], %[[CONSTANT_0]]] [16, 1, 1024] [128, 1, 1] : memref<4096x1x768xf32> -> tensor<16x1x1024xf32>
+// CHECK:           %[[EXTRACT_0:.*]] = xtile.extract %[[ARG0]][%[[APPLY_INDEXING_0]], %[[CONSTANT_0]], %[[CONSTANT_0]]] [16, 1, 1024] [128, 1, 1] : memref<4096x1x768xf32> -> tensor<16x1x1024xf32>
 // CHECK:           %[[APPLY_INDEXING_1:.*]] = xla.apply_indexing #[[$ATTR_1]](%[[ARG2]])
 // CHECK:           %[[APPLY_INDEXING_2:.*]] = xla.apply_indexing #[[$ATTR_2]](%[[ARG2]])
 // CHECK:           xtile.insert %[[EXTRACT_0]] into %[[ARG1]][%[[APPLY_INDEXING_1]], %[[APPLY_INDEXING_2]], %[[CONSTANT_0]]] [16, 1, 1024] [1, 1, 1] : tensor<16x1x1024xf32> -> memref<32x128x768xf32>
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_grouped.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_grouped.hlo
index d98a85e4b4b..74c8da82278 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_grouped.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_grouped.hlo
@@ -1,5 +1,5 @@
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation 2>&1 | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s 2>&1 | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = f32[64] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_tiled_1.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_tiled_1.hlo
index 17ccd972544..4bcf2729b84 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_tiled_1.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_tiled_1.hlo
@@ -1,5 +1,5 @@
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s --check-prefixes=CHECK
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = f32[16] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_tiled_2.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_tiled_2.hlo
index e69acf73893..0cc4fef8a6e 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_tiled_2.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_expand_tiled_2.hlo
@@ -1,5 +1,5 @@
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s --check-prefixes=CHECK
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = f32[16] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_non_contiguous_collapse.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_non_contiguous_collapse.hlo
index b2929845c73..d99593715ac 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_non_contiguous_collapse.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_non_contiguous_collapse.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = f32[32]{0} parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_non_contiguous_expand.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_non_contiguous_expand.hlo
index bfa7cbbe109..7cb6c4aa653 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_non_contiguous_expand.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_non_contiguous_expand.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s --check-prefixes=CHECK,CHECK-OLD
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s --check-prefixes=CHECK,CHECK-NEW
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = f32[32,4]{1,0} parameter(0)
@@ -17,8 +14,7 @@ fusion {
 // CHECK-SAME: %[[PID:.*]]: index{{( \{xla.range = \[0 : index, 1 : index\]\})?}})
 // CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index
 // CHECK: %[[OFFSET:.*]] = xla.apply_indexing #[[$MAP_0]](%[[PID]])
-// CHECK-OLD: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[OFFSET]], %[[CONSTANT_0]]] [16, 1] [1, 0] : memref<32x4xf32> -> tensor<16x1xf32>
-// CHECK-NEW: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[OFFSET]], %[[CONSTANT_0]]] [16, 1] [1, 1] : memref<32x4xf32> -> tensor<16x1xf32>
+// CHECK: %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[OFFSET]], %[[CONSTANT_0]]] [16, 1] [1, 1] : memref<32x4xf32> -> tensor<16x1xf32>
 // CHECK: %[[RESHAPE:.*]] = tt.reshape %[[EXTRACT]] : tensor<16x1xf32> -> tensor<16xf32>
 // CHECK: xtile.insert %[[RESHAPE]] into %[[ARG1]][%[[OFFSET]]] [16] [1] : tensor<16xf32> -> memref<32xf32>
 // CHECK: xtile.return
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_only.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_only.hlo
index b54b9cb33b3..6104a3fcb7c 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_only.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reshape_only.hlo
@@ -1,21 +1,14 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s --check-prefixes=CHECK,CHECK-OLD
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s --check-prefixes=CHECK,CHECK-NEW
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = f32[64]{0} parameter(0)
   ROOT bitcast = f32[8,8]{1,0} bitcast(p)
 }
 
-// CHECK-OLD: #[[$ATTR_0:.+]] = #xla.indexing_map<"(pid_0) -> (pid_0 * 4), domain: pid_0 in [0, 15]">
-// CHECK-OLD: #[[$ATTR_1:.+]] = #xla.indexing_map<"(pid_0) -> (pid_0 / 2), domain: pid_0 in [0, 15]">
-// CHECK-OLD: #[[$ATTR_2:.+]] = #xla.indexing_map<"(pid_0) -> ((pid_0 mod 2) * 4), domain: pid_0 in [0, 15]">
-
-// CHECK-NEW: #[[$ATTR_0:.+]] = #xla.indexing_map<"(pid) -> (pid * 4), domain: pid in [0, 15]">
-// CHECK-NEW: #[[$ATTR_1:.+]] = #xla.indexing_map<"(pid) -> (pid / 2), domain: pid in [0, 15]">
-// CHECK-NEW: #[[$ATTR_2:.+]] = #xla.indexing_map<"(pid) -> ((pid mod 2) * 4), domain: pid in [0, 15]">
+// CHECK: #[[$ATTR_0:.+]] = #xla.indexing_map<"(pid) -> (pid * 4), domain: pid in [0, 15]">
+// CHECK: #[[$ATTR_1:.+]] = #xla.indexing_map<"(pid) -> (pid / 2), domain: pid in [0, 15]">
+// CHECK: #[[$ATTR_2:.+]] = #xla.indexing_map<"(pid) -> ((pid mod 2) * 4), domain: pid in [0, 15]">
 
 // CHECK-LABEL:   xtile.entry_func @triton_fn(
 // CHECK-SAME:      %[[ARG0:.*]]: memref<64xf32>,
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_1d.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_1d.hlo
index 3c153a7c2ac..040a344fdf1 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_1d.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_1d.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation=false | FileCheck %s --check-prefixes=CHECK,CHECK-OLD
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation=false
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation=true | FileCheck %s --check-prefixes=CHECK,CHECK-NEW
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation=true
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = f32[1,1,1] parameter(0)
@@ -13,8 +11,7 @@ fusion {
 // CHECK-SAME:      %[[ARG1:.*]]: memref<1xf32>,
 // CHECK-SAME:      %[[PID:.*]]: index{{( \{xla.range = \[0 : index, 0 : index\]\})?}}) attributes {num_opaque_args = 0 : i32} {
 // CHECK:           %[[C0:.*]] = arith.constant 0 : index
-// CHECK-OLD:       %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] [1, 1, 1] [0, 0, 0] : memref<1x1x1xf32> -> tensor<1x1x1xf32>
-// CHECK-NEW:       %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] [1, 1, 1] [1, 1, 1] : memref<1x1x1xf32> -> tensor<1x1x1xf32>
+// CHECK:           %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] [1, 1, 1] [1, 1, 1] : memref<1x1x1xf32> -> tensor<1x1x1xf32>
 // CHECK:           %[[UNSPLAT:.*]] = tt.unsplat %[[EXTRACT]] : tensor<1x1x1xf32>
 // CHECK:           %[[SPLAT:.*]] = tt.splat %[[UNSPLAT]] : f32 -> tensor<1xf32>
 // CHECK:           xtile.insert %[[SPLAT]] into %[[ARG1]][%[[C0]]] [1] [1] : tensor<1xf32> -> memref<1xf32>
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_1d_tiled.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_1d_tiled.hlo
index c2d0c87737e..2e384ab175f 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_1d_tiled.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_1d_tiled.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation=false | FileCheck %s
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation=true | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation=false
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation=true
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = f32[1,1,1] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_scalar.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_scalar.hlo
index f3cbbafea04..7a51d5b9f6e 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_scalar.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_to_scalar.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation=false | FileCheck %s --check-prefixes=CHECK,CHECK-OLD
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation=false
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s --check-prefixes=CHECK,CHECK-NEW
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = f32[1,1,1] parameter(0)
@@ -13,8 +11,7 @@ fusion {
 // CHECK-SAME:      %[[ARG1:.*]]: memref<f32>,
 // CHECK-SAME:      %[[PID:.*]]: index{{( \{xla.range = \[0 : index, 0 : index\]\})?}}) attributes {num_opaque_args = 0 : i32} {
 // CHECK:           %[[C0:.*]] = arith.constant 0 : index
-// CHECK-OLD:       %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] [1, 1, 1] [0, 0, 0] : memref<1x1x1xf32> -> tensor<1x1x1xf32>
-// CHECK-NEW:       %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] [1, 1, 1] [1, 1, 1] : memref<1x1x1xf32> -> tensor<1x1x1xf32>
+// CHECK:           %[[EXTRACT:.*]] = xtile.extract %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] [1, 1, 1] [1, 1, 1] : memref<1x1x1xf32> -> tensor<1x1x1xf32>
 // CHECK:           %[[UNSPLAT:.*]] = tt.unsplat %[[EXTRACT]] : tensor<1x1x1xf32>
 // CHECK:           %[[FROM_ELEMENTS:.*]] = tensor.from_elements %[[UNSPLAT]] : tensor<f32>
 // CHECK:           xtile.insert %[[FROM_ELEMENTS]] into %[[ARG1]][] [] [] : tensor<f32> -> memref<f32>
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_only.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_only.hlo
index d37186f7a65..148299a3ee4 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_only.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_only.hlo
@@ -1,18 +1,12 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s --check-prefixes=CHECK,CHECK-OLD
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s --check-prefixes=CHECK,CHECK-NEW
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = s8[42,5]{0,1} parameter(0)
   ROOT bitcast = s8[5,42] bitcast(p)
 }
-// CHECK-OLD: #[[$ATTR_0:.+]] = #xla.indexing_map<"(pid_0) -> (pid_0 mod 42), domain: pid_0 in [0, 83]">
-// CHECK-OLD: #[[$ATTR_1:.+]] = #xla.indexing_map<"(pid_0) -> ((pid_0 / 42) * 4), domain: pid_0 in [0, 83]">
-
-// CHECK-NEW: #[[$ATTR_0:.+]] = #xla.indexing_map<"(pid) -> (pid mod 42), domain: pid in [0, 83]">
-// CHECK-NEW: #[[$ATTR_1:.+]] = #xla.indexing_map<"(pid) -> ((pid / 42) * 4), domain: pid in [0, 83]">
+// CHECK: #[[$ATTR_0:.+]] = #xla.indexing_map<"(pid) -> (pid mod 42), domain: pid in [0, 83]">
+// CHECK: #[[$ATTR_1:.+]] = #xla.indexing_map<"(pid) -> ((pid / 42) * 4), domain: pid in [0, 83]">
 
 // CHECK-LABEL:   xtile.entry_func @triton_fn(
 // CHECK-SAME:      %[[ARG0:.*]]: memref<42x5xi8, #xtile.layout<[0, 1]>>,
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_or_reshape.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_or_reshape.hlo
index e75f190e74e..3d7ca29981d 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_or_reshape.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_or_reshape.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = s8[1,1024]{1,0} parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_permutation_inversion.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_permutation_inversion.hlo
index 9b7c572cc8c..3eea94d751d 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_permutation_inversion.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_transpose_permutation_inversion.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s --abs_error_bound=1e-4 --rel_error_bound=1e-6
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-4 --rel_error_bound=1e-6
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s --abs_error_bound=1e-4 --rel_error_bound=1e-6
 
 fusion {
   p0 = bf16[8,16,16,4]{2,0,3,1} parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_unsupported_expand.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_unsupported_expand.hlo
index cbe41f83b92..432e0179ece 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_unsupported_expand.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_unsupported_expand.hlo
@@ -1,4 +1,4 @@
-// RUN: hlo-opt %s --platform=gpu --xla_gpu_target_config_filename=%S/../../../../target_config/specs/h100_sxm.txtpb --xla_gpu_experimental_enable_tiling_propagation=true --logtostderr --vmodule=tile_propagation=2 2>&1 | FileCheck %s
+// RUN: hlo-opt %s --platform=gpu --xla_gpu_target_config_filename=%S/../../../../target_config/specs/h100_sxm.txtpb --logtostderr --vmodule=tile_propagation=2 2>&1 | FileCheck %s
 
 // CHECK: Skipping reshape contiguity check as tile sizes are not assigned yet.
 // CHECK: ROOT {{.*}} = bf16[4096,16,4096]{{.*}} fusion({{.*}}), kind=kCustom, calls=%{{.*}}gemm_fusion_tmp_11_computation, backend_config={{{.*}}"kind":"__triton_nested_gemm_fusion"
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/broadcast/broadcast.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/broadcast/broadcast.hlo
index 48c638c482b..97b7e1e15e4 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/broadcast/broadcast.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/broadcast/broadcast.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p0 = s32[234]{0} parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/broadcast/scalar_broadcast.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/broadcast/scalar_broadcast.hlo
index ef66e9120ac..1b88466432d 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/broadcast/scalar_broadcast.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/broadcast/scalar_broadcast.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s --abs_error_bound=1e-6 --rel_error_bound=1e-6
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-6 --rel_error_bound=1e-6
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s --abs_error_bound=1e-6 --rel_error_bound=1e-6
 
 fusion {
   param_0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo
index 1b838aa65c5..7301bf06b37 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_gather.hlo
@@ -1,5 +1,4 @@
 // RUN: fusion_to_triton %s | FileCheck %s
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
 f {
   %param0 = f32[128,128]{1,0} parameter(0)
   ROOT result = f32[256,128]{1,0} all-gather(%param0),
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo
index 731791d3cbe..60c9043de26 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/collectives/all_reduce.hlo
@@ -1,5 +1,4 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
+// RUN: fusion_to_triton %s | FileCheck %s
 
 sum {
   x = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_active_scalar_constant.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_active_scalar_constant.hlo
index 6ed63df5959..28d8eda0854 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_active_scalar_constant.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_active_scalar_constant.hlo
@@ -1,5 +1,4 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s 2>&1 | FileCheck %s
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s 2>&1 | FileCheck %s
+// RUN: fusion_to_triton %s 2>&1 | FileCheck %s
 
 // CHECK:      xtile.entry_func @triton_fn
 // CHECK:      %[[CST:.*]] = arith.constant dense<4.000000e+00> : tensor<2xf32>
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_fusion.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_fusion.hlo
index f397419be89..a944e3c567d 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_fusion.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_fusion.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s --check-prefixes=CHECK,CHECK-OLD
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s --check-prefixes=CHECK,CHECK-NEW
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p0 = s32[128] parameter(0)
@@ -12,12 +9,8 @@ fusion {
   neg1 = s32[128] negate(p1)
   ROOT concatenate = s32[281] concatenate(abs0, neg1, p2), dimensions={0}
 }
-// CHECK-OLD-DAG: #[[$ATTR_0:.+]] = #xla.indexing_map<"(pid_0) -> (pid_0 * 32), domain: pid_0 in [0, 8]">
-// CHECK-OLD-DAG: #[[$ATTR_1:.+]] = #xla.indexing_map<"(pid_0) -> (pid_0 * 32), domain: pid_0 in [0, 3]">
-// CHECK-OLD-DAG: #[[$ATTR_2:.+]] = #xla.indexing_map<"(pid_0) -> (pid_0 * 32 - 128), domain: pid_0 in [4, 7]">
-
-// CHECK-NEW-DAG: #[[$ATTR_0:.+]] = #xla.indexing_map<"(pid) -> (pid * 32), domain: pid in [0, 8]">
-// CHECK-NEW-DAG: #[[$ATTR_1:.+]] = #xla.indexing_map<"(pid) -> (pid * 32 - 128), domain: pid in [0, 8]">
+// CHECK-DAG: #[[$ATTR_0:.+]] = #xla.indexing_map<"(pid) -> (pid * 32), domain: pid in [0, 8]">
+// CHECK-DAG: #[[$ATTR_1:.+]] = #xla.indexing_map<"(pid) -> (pid * 32 - 128), domain: pid in [0, 8]">
 
 // CHECK:       xtile.entry_func
 // CHECK-SAME:  %[[ARG0:[A-Za-z0-9]+]]: memref<128xi32>,
@@ -34,18 +27,14 @@ fusion {
 // CHECK:       %[[OFFSET:.*]] = xla.apply_indexing #[[$ATTR_0]](%[[PID]])
 // CHECK:       %[[CMPI_0:.*]] = arith.cmpi slt, %[[OFFSET]], %[[C128]] : index
 // CHECK:       %[[IF_0:.*]] = scf.if %[[CMPI_0]] -> (tensor<32xi32>) {
-// CHECK-OLD:     %[[OFFSET_1:.*]] = xla.apply_indexing #[[$ATTR_1]](%[[PID]])
-// CHECK-OLD:     %[[EXTRACT_0:.*]] = xtile.extract %[[ARG0]][%[[OFFSET_1]]] [32] [1] : memref<128xi32> -> tensor<32xi32>
-// CHECK-NEW:     %[[EXTRACT_0:.*]] = xtile.extract %[[ARG0]][%[[OFFSET]]] [32] [1] : memref<128xi32> -> tensor<32xi32>
+// CHECK:         %[[EXTRACT_0:.*]] = xtile.extract %[[ARG0]][%[[OFFSET]]] [32] [1] : memref<128xi32> -> tensor<32xi32>
 // CHECK:         %[[ABSI_0:.*]] = math.absi %[[EXTRACT_0]] : tensor<32xi32>
 // CHECK:         scf.yield %[[ABSI_0]] : tensor<32xi32>
 // CHECK:       } else {
 // CHECK:         %[[CMPI_1:.*]] = arith.cmpi slt, %[[OFFSET]], %[[C256]] : index
 // CHECK:         %[[IF_1:.*]] = scf.if %[[CMPI_1]] -> (tensor<32xi32>) {
-// CHECK-OLD:       %[[OFFSET_2:.*]] = xla.apply_indexing #[[$ATTR_2]](%[[PID]])
-// CHECK-OLD:       %[[EXTRACT_1:.*]] = xtile.extract %[[ARG1]][%[[OFFSET_2]]] [32] [1] : memref<128xi32> -> tensor<32xi32>
-// CHECK-NEW:       %[[OFFSET_2:.*]] = xla.apply_indexing #[[$ATTR_1]](%[[PID]])
-// CHECK-NEW:       %[[EXTRACT_1:.*]] = xtile.extract %[[ARG1]][%[[OFFSET_2]]] [32] [1] : memref<128xi32> -> tensor<32xi32>
+// CHECK:           %[[OFFSET_2:.*]] = xla.apply_indexing #[[$ATTR_1]](%[[PID]])
+// CHECK:           %[[EXTRACT_1:.*]] = xtile.extract %[[ARG1]][%[[OFFSET_2]]] [32] [1] : memref<128xi32> -> tensor<32xi32>
 // CHECK:           %[[SUBI_0:.*]] = arith.subi %[[TC0]], %[[EXTRACT_1]] : tensor<32xi32>
 // CHECK:           scf.yield %[[SUBI_0]] : tensor<32xi32>
 // CHECK:         } else {
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_offset_is_a_multiple_of_tile_size.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_offset_is_a_multiple_of_tile_size.hlo
index 1e366b8d768..3c129206ee7 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_offset_is_a_multiple_of_tile_size.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_offset_is_a_multiple_of_tile_size.hlo
@@ -1,4 +1,4 @@
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: triton_test_correctness %s
 
 fusion {
   p0 = f32[100] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_offset_not_a_multiple_of_tile_size.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_offset_not_a_multiple_of_tile_size.hlo
index 9aab286b21c..4e69c9f0bab 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_offset_not_a_multiple_of_tile_size.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_offset_not_a_multiple_of_tile_size.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation 2>&1 | FileCheck %s
+// RUN: fusion_to_triton %s 2>&1 | FileCheck %s
 
 // CHECK: FAILED_PRECONDITION:{{.*}}The variable step expression 'v0 * 49 + v1 * 8' must be a clean multiple of its tile size 8
 fusion {
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_pred.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_pred.hlo
index eb01504c85a..7d353345d87 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_pred.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_pred.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s --check-prefixes=CHECK,CHECK-OLD
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s --check-prefixes=CHECK,CHECK-NEW
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p0 = pred[128] parameter(0)
@@ -10,12 +7,8 @@ fusion {
   p2 = pred[25] parameter(2)
   ROOT concatenate = pred[281] concatenate(p0, p1, p2), dimensions={0}
 }
-// CHECK-OLD-DAG: #[[$MAP0:.+]] = #xla.indexing_map<"(pid_0) -> (pid_0 * 32), domain: pid_0 in [0, 8]">
-// CHECK-OLD-DAG: #[[$MAP1:.+]] = #xla.indexing_map<"(pid_0) -> (pid_0 * 32), domain: pid_0 in [0, 3]">
-// CHECK-OLD-DAG: #[[$MAP2:.+]] = #xla.indexing_map<"(pid_0) -> (pid_0 * 32 - 128), domain: pid_0 in [4, 7]">
-
-// CHECK-NEW-DAG: #[[$MAP0:.+]] = #xla.indexing_map<"(pid) -> (pid * 32), domain: pid in [0, 8]">
-// CHECK-NEW-DAG: #[[$MAP1:.+]] = #xla.indexing_map<"(pid) -> (pid * 32 - 128), domain: pid in [0, 8]">
+// CHECK-DAG: #[[$MAP0:.+]] = #xla.indexing_map<"(pid) -> (pid * 32), domain: pid in [0, 8]">
+// CHECK-DAG: #[[$MAP1:.+]] = #xla.indexing_map<"(pid) -> (pid * 32 - 128), domain: pid in [0, 8]">
 
 // CHECK:       xtile.entry_func
 // CHECK-SAME:  %[[ARG0:[A-Za-z0-9]+]]: memref<128xi8>,
@@ -32,18 +25,14 @@ fusion {
 // CHECK:       %[[OFFSET:.*]] = xla.apply_indexing #[[$MAP0]](%[[PID]])
 // CHECK:       %[[CMPI_0:.*]] = arith.cmpi slt, %[[OFFSET]], %[[C128]] : index
 // CHECK:       %[[IF_0:.*]] = scf.if %[[CMPI_0]] -> (tensor<32xi1>) {
-// CHECK-OLD:     %[[OFFSET_1:.*]] = xla.apply_indexing #[[$MAP1]](%[[PID]])
-// CHECK-OLD:     %[[EXTRACT_0:.*]] = xtile.extract %[[ARG0]][%[[OFFSET_1]]] [32] [1] : memref<128xi8> -> tensor<32xi8>
-// CHECK-NEW:     %[[EXTRACT_0:.*]] = xtile.extract %[[ARG0]][%[[OFFSET]]] [32] [1] : memref<128xi8> -> tensor<32xi8>
+// CHECK:         %[[EXTRACT_0:.*]] = xtile.extract %[[ARG0]][%[[OFFSET]]] [32] [1] : memref<128xi8> -> tensor<32xi8>
 // CHECK:         %[[CMPI_1:.*]] = arith.cmpi ne, %[[EXTRACT_0]], %[[CST]] : tensor<32xi8>
 // CHECK:         scf.yield %[[CMPI_1]] : tensor<32xi1>
 // CHECK:       } else {
 // CHECK:         %[[CMPI_2:.*]] = arith.cmpi slt, %[[OFFSET]], %[[C256]] : index
 // CHECK:         %[[IF_1:.*]] = scf.if %[[CMPI_2]] -> (tensor<32xi1>) {
-// CHECK-OLD:       %[[OFFSET_2:.*]] = xla.apply_indexing #[[$MAP2]](%[[PID]])
-// CHECK-OLD:       %[[EXTRACT_1:.*]] = xtile.extract %[[ARG1]][%[[OFFSET_2]]] [32] [1] : memref<128xi8> -> tensor<32xi8>
-// CHECK-NEW:       %[[OFFSET_2:.*]] = xla.apply_indexing #[[$MAP1]](%[[PID]])
-// CHECK-NEW:       %[[EXTRACT_1:.*]] = xtile.extract %[[ARG1]][%[[OFFSET_2]]] [32] [1] : memref<128xi8> -> tensor<32xi8>
+// CHECK:           %[[OFFSET_2:.*]] = xla.apply_indexing #[[$MAP1]](%[[PID]])
+// CHECK:           %[[EXTRACT_1:.*]] = xtile.extract %[[ARG1]][%[[OFFSET_2]]] [32] [1] : memref<128xi8> -> tensor<32xi8>
 // CHECK:           %[[CMPI_3:.*]] = arith.cmpi ne, %[[EXTRACT_1]], %[[CST]] : tensor<32xi8>
 // CHECK:           scf.yield %[[CMPI_3]] : tensor<32xi1>
 // CHECK:         } else {
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive.hlo
index 5ac42930afd..67bcd82f51d 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive.hlo
@@ -1,8 +1,6 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s 2>&1 | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s 2>&1 | FileCheck %s --check-prefixes=CHECK-ERROR
+// RUN: fusion_to_triton %s 2>&1 | FileCheck %s --check-prefixes=CHECK-ERROR
 
+// Expected IR once b/528277610 is fixed:
 // CHECK:      xtile.entry_func @triton_fn
 // CHECK-NEXT:   %[[C0:.*]] = arith.constant 0 : index
 // CHECK-NEXT:   %[[EXTRACT:.*]] = xtile.extract %arg0[%[[C0]], %[[C0]]] [64, 16] [1, 1] : memref<64x10xf32> -> tensor<64x16xf32>
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_simplify_offset.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_simplify_offset.hlo
index a124e6bd0ec..87861e8af65 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_simplify_offset.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_simplify_offset.hlo
@@ -1,5 +1,5 @@
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation 2>&1 | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s 2>&1 | FileCheck %s
+// RUN: triton_test_correctness %s
 
 // Both branches of concat access whole second dimension of parameters and we
 // can simplify offset of it to 0 constant.
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_slice_inactive_operand.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_slice_inactive_operand.hlo
index 9f78f6f712f..441855aed41 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_slice_inactive_operand.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_slice_inactive_operand.hlo
@@ -1,5 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s 2>&1 | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN: fusion_to_triton %s 2>&1 | FileCheck %s
+// RUN: triton_test_correctness %s
 
 // The slice effectively accesses only second parameter so we can skip scf.if .
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/constraints/max_tensor_elements_violation.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/constraints/max_tensor_elements_violation.hlo
index cfad3302c69..ec5cb3a8ce7 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/constraints/max_tensor_elements_violation.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/constraints/max_tensor_elements_violation.hlo
@@ -1,8 +1,6 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s 2>&1 | FileCheck %s --check-prefixes=CHECK-ERROR
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation 2>&1 | FileCheck %s --check-prefixes=CHECK-ERROR-NEW
+// RUN: fusion_to_triton %s 2>&1 | FileCheck %s --check-prefixes=CHECK-ERROR
 
-// CHECK-ERROR: Tiling does not satisfy constraints.
-// CHECK-ERROR-NEW: exceeds the maximum allowed tensor size of 1048576 elements
+// CHECK-ERROR: exceeds the maximum allowed tensor size of 1048576 elements
 
 fusion {
   p0 = f32[2048,2048] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/bf16_small_rhs_outer.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/bf16_small_rhs_outer.hlo
index 8db956ec07a..703cb0e1573 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/bf16_small_rhs_outer.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/bf16_small_rhs_outer.hlo
@@ -1,5 +1,4 @@
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s --abs_error_bound=1e-1 --rel_error_bound=1e-2
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-1 --rel_error_bound=1e-2
+// RUN: triton_test_correctness %s --abs_error_bound=1e-1 --rel_error_bound=1e-2
 
 fusion {
   fdot.p0 = bf16[64,32] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot.hlo
index 35d0f38d580..6021794f647 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s --abs_error_bound=1e-4 --rel_error_bound=1e-6
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-4 --rel_error_bound=1e-6
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s --abs_error_bound=1e-4 --rel_error_bound=1e-6
 
 fusion {
   fdot.p0 = f32[32,16,256]{2,0,1} parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_broadcast_sequential_dim_mismatch.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_broadcast_sequential_dim_mismatch.hlo
index 247321c43ca..fb8a5852e3e 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_broadcast_sequential_dim_mismatch.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_broadcast_sequential_dim_mismatch.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s --abs_error_bound=1e-4 --rel_error_bound=1e-6
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-4 --rel_error_bound=1e-6
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s --abs_error_bound=1e-4 --rel_error_bound=1e-6
 
 fusion {
   p0 = f32[16,32]{1,0} parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_concat.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_concat.hlo
index f2bf8357aa9..be7e0659f2d 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_concat.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_concat.hlo
@@ -1,5 +1,4 @@
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false --abs_error_bound=1e-4 --rel_error_bound=1e-6 %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-4 --rel_error_bound=1e-6
+// RUN: triton_test_correctness %s --abs_error_bound=1e-4 --rel_error_bound=1e-6
 
 
 fusion {
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_major_lhs_contract.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_major_lhs_contract.hlo
index fadfb64a72c..5971b5e464e 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_major_lhs_contract.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_major_lhs_contract.hlo
@@ -1,5 +1,4 @@
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: triton_test_correctness %s
 
 
 fusion {
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_minor_rhs_contract.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_minor_rhs_contract.hlo
index e71238c0a01..5d65c5f41cb 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_minor_rhs_contract.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_minor_rhs_contract.hlo
@@ -1,5 +1,4 @@
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: triton_test_correctness %s
 
 fusion {
   p0 = f32[32,299] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_transpose_tma.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_transpose_tma.hlo
index 960f40a14ab..e7637692f05 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_transpose_tma.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_transpose_tma.hlo
@@ -1,6 +1,5 @@
 // Regression test for b/503878530
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false --abs_error_bound=1e-3 --rel_error_bound=1e-3 %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-3 --rel_error_bound=1e-3 %s
+// RUN: triton_test_correctness --abs_error_bound=1e-3 --rel_error_bound=1e-3 %s
 
 dot {
   p0 = f16[64,128,32]{2,1,0} parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_unit_tile_sizes.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_unit_tile_sizes.hlo
index 80330efec6b..33d7084784b 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_unit_tile_sizes.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/dot_unit_tile_sizes.hlo
@@ -1,5 +1,4 @@
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: triton_test_correctness %s
 
 fusion {
   p0 = f32[2,3,32,125] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/duplicate_dot_operand.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/duplicate_dot_operand.hlo
index 33a6a4320a8..26abe909b92 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/duplicate_dot_operand.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/duplicate_dot_operand.hlo
@@ -1,4 +1,4 @@
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
+// RUN: fusion_to_triton %s | FileCheck %s
 
 fusion {
   p0 = f32[4,8] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/layout_021_201_210.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/layout_021_201_210.hlo
index 48f150f4a30..52a5616a1f5 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/layout_021_201_210.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/layout_021_201_210.hlo
@@ -1,7 +1,5 @@
-// RUN: sed 's/ENABLE_TMA/false/' %s > %t.hlo && triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
-// RUN: sed 's/ENABLE_TMA/false/' %s > %t.hlo && triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
-// RUN: sed 's/ENABLE_TMA/true/' %s > %t.hlo && triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
-// RUN: sed 's/ENABLE_TMA/true/' %s > %t.hlo && triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
+// RUN: sed 's/ENABLE_TMA/false/' %s > %t.hlo && triton_test_correctness %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
+// RUN: sed 's/ENABLE_TMA/true/' %s > %t.hlo && triton_test_correctness %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
 
 fusion {
   fdot.p0 = f32[32,16,256]{0,2,1} parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/layout_201_012_201.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/layout_201_012_201.hlo
index e26fe4a5f57..96a790d044f 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/layout_201_012_201.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/layout_201_012_201.hlo
@@ -1,7 +1,5 @@
-// RUN: sed 's/ENABLE_TMA/false/' %s > %t.hlo && triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
-// RUN: sed 's/ENABLE_TMA/false/' %s > %t.hlo && triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
-// RUN: sed 's/ENABLE_TMA/true/' %s > %t.hlo && triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
-// RUN: sed 's/ENABLE_TMA/true/' %s > %t.hlo && triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
+// RUN: sed 's/ENABLE_TMA/false/' %s > %t.hlo && triton_test_correctness %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
+// RUN: sed 's/ENABLE_TMA/true/' %s > %t.hlo && triton_test_correctness %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
 
 fusion {
   fdot.p0 = f32[32,16,256]{2,0,1} parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/layout_210_210_102.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/layout_210_210_102.hlo
index bdbc6c721fb..1c1b30726f4 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/layout_210_210_102.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/layout_210_210_102.hlo
@@ -1,7 +1,5 @@
-// RUN: sed 's/ENABLE_TMA/false/' %s > %t.hlo && triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
-// RUN: sed 's/ENABLE_TMA/false/' %s > %t.hlo && triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
-// RUN: sed 's/ENABLE_TMA/true/' %s > %t.hlo && triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
-// RUN: sed 's/ENABLE_TMA/true/' %s > %t.hlo && triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
+// RUN: sed 's/ENABLE_TMA/false/' %s > %t.hlo && triton_test_correctness %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
+// RUN: sed 's/ENABLE_TMA/true/' %s > %t.hlo && triton_test_correctness %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
 
 fusion {
   fdot.p0 = f32[32,16,256]{2,1,0} parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/layout_210_210_210.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/layout_210_210_210.hlo
index da747b50047..79fb4e04f32 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/layout_210_210_210.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/layout_210_210_210.hlo
@@ -1,7 +1,5 @@
-// RUN: sed 's/ENABLE_TMA/false/' %s > %t.hlo && triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
-// RUN: sed 's/ENABLE_TMA/false/' %s > %t.hlo && triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
-// RUN: sed 's/ENABLE_TMA/true/' %s > %t.hlo && triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
-// RUN: sed 's/ENABLE_TMA/true/' %s > %t.hlo && triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
+// RUN: sed 's/ENABLE_TMA/false/' %s > %t.hlo && triton_test_correctness %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
+// RUN: sed 's/ENABLE_TMA/true/' %s > %t.hlo && triton_test_correctness %t.hlo --abs_error_bound=1e-4 --rel_error_bound=1e-6
 
 fusion {
   fdot.p0 = f32[32,16,256]{2,1,0} parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/masked_dot.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/masked_dot.hlo
index cbba0e0b9a5..1e2d80bf804 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/masked_dot.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/masked_dot.hlo
@@ -1,6 +1,4 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s --abs_error_bound=1e-4 --rel_error_bound=1e-6
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-4 --rel_error_bound=1e-6
+// RUN: triton_test_correctness %s --abs_error_bound=1e-4 --rel_error_bound=1e-6
 
 
 fusion {
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/reshape_of_pad.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/reshape_of_pad.hlo
index c943af8ac5c..4033b110521 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/reshape_of_pad.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/reshape_of_pad.hlo
@@ -1,5 +1,4 @@
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s --abs_error_bound=1e-4 --rel_error_bound=1e-6
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-4 --rel_error_bound=1e-6
+// RUN: triton_test_correctness %s --abs_error_bound=1e-4 --rel_error_bound=1e-6
 
 fusion {
   p0 = bf16[16,32,67,133] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/single_tile_dot.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/single_tile_dot.hlo
index 5b9d8ff3008..b05da7cecc8 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/single_tile_dot.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/single_tile_dot.hlo
@@ -1,6 +1,4 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s --abs_error_bound=1e-4 --rel_error_bound=1e-6
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-4 --rel_error_bound=1e-6
+// RUN: triton_test_correctness %s --abs_error_bound=1e-4 --rel_error_bound=1e-6
 
 
 // Simplest case when everything fits into one tile that is useful for
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/use_major_to_minor_schedule_lhs_larger.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/use_major_to_minor_schedule_lhs_larger.hlo
index 936c7948c45..2f45c2504ec 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/use_major_to_minor_schedule_lhs_larger.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/use_major_to_minor_schedule_lhs_larger.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false --abs_error_bound=2e-4 --rel_error_bound=1e-6 %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=2e-4 --rel_error_bound=1e-6
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s --abs_error_bound=2e-4 --rel_error_bound=1e-6
 
 
 fusion {
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/use_major_to_minor_schedule_not_rooted_in_dot.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/use_major_to_minor_schedule_not_rooted_in_dot.hlo
index 7c6b4d05627..87367882e56 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/use_major_to_minor_schedule_not_rooted_in_dot.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/use_major_to_minor_schedule_not_rooted_in_dot.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false --abs_error_bound=2e-4 --rel_error_bound=1e-6 %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=2e-4 --rel_error_bound=1e-6
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s --abs_error_bound=2e-4 --rel_error_bound=1e-6
 
 fusion {
   p0 = f32[2,32,64] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/use_transposed_schedule_lhs_smaller.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/use_transposed_schedule_lhs_smaller.hlo
index 65db94519df..3660726be04 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/use_transposed_schedule_lhs_smaller.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dot/use_transposed_schedule_lhs_smaller.hlo
@@ -1,6 +1,4 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false --abs_error_bound=2e-4 --rel_error_bound=1e-6 %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=2e-4 --rel_error_bound=1e-6
+// RUN: triton_test_correctness %s --abs_error_bound=2e-4 --rel_error_bound=1e-6
 
 fusion {
   p0 = f32[2,32,64] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_multiple_dimensions.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_multiple_dimensions.hlo
index f1cb0233879..462eb146a1c 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_multiple_dimensions.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_multiple_dimensions.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 // CHECK-DAG: #[[MAP:.*]] = #xla.indexing_map<"(pid{{.*}}, d1) -> (pid{{.*}} * 8 + d1), domain: pid{{.*}} in [0, 1], d1 in [0, 22]">
 // CHECK-DAG: #[[MAP1:.*]] = #xla.indexing_map<"(pid{{.*}}) -> (pid{{.*}} * 8), domain: pid{{.*}} in [0, 1]">
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_add.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_add.hlo
index 5a7f06d1dd4..9d3b085d675 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_add.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_add.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p0 = s32[64] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_dyn_slice.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_dyn_slice.hlo
index 42844dfa63c..daa60bfe557 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_dyn_slice.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_dyn_slice.hlo
@@ -1,9 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s --check-prefixes=CHECK-OLD
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation \
-// RUN:   | FileCheck %s --check-prefixes=CHECK-NEW
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p0 = f32[64] parameter(0)
@@ -15,33 +11,31 @@ fusion {
   off2 = s32[] reshape(slice2)
   ROOT d2 = s32[10] dynamic-slice(d1, off2), dynamic_slice_sizes={10}
 }
-// CHECK-NEW-DAG: #[[MAP:.*]] = #xla.indexing_map<"(d0{{.*}}) -> (d0{{.*}} + 1), domain: d0{{.*}} in [0, 44]">
-// CHECK-NEW-DAG: #[[MAP1:.*]] = #xla.indexing_map<"(d0{{.*}}, d1{{.*}}) -> (d0{{.*}} + d1{{.*}}), domain: d0{{.*}} in [0, 10], d1{{.*}} in [0, 44]">
-// CHECK-NEW: xtile.entry_func @triton_fn(
-// CHECK-NEW-SAME:    %[[ARG0:.*]]: memref<64xf32>,
-// CHECK-NEW-SAME:    %[[ARG1:.*]]: memref<10xi32>,
-// CHECK-NEW-SAME:    %[[ARG2:.*]]: index{{( \{xla.range = \[0 : index, 0 : index\]\})?}})
-// CHECK-NEW-DAG:     %[[C10:.*]] = arith.constant 10 : i32
-// CHECK-NEW-DAG:     %[[C44:.*]] = arith.constant 44 : i32
-// CHECK-NEW-DAG:     %[[C0_I32:.*]] = arith.constant 0 : i32
-// CHECK-NEW-DAG:     %[[C0:.*]] = arith.constant 0 : index
-// CHECK-NEW:         %[[EXTRACT0:.*]] = xtile.extract %[[ARG0]][%[[C0]]] [1] [{{.*}}] : memref<64xf32> -> tensor<1xf32>
-// CHECK-NEW:         %[[UNSPLAT0:.*]] = tt.unsplat %{{.*}} : tensor<1xi32>
-// CHECK-NEW:         %[[MAX0:.*]] = arith.maxsi %[[UNSPLAT0]], %[[C0_I32]] : i32
-// CHECK-NEW:         %[[MIN0:.*]] = arith.minsi %[[MAX0]], %[[C44]] : i32
-// CHECK-NEW:         %[[CAST0:.*]] = arith.index_cast %[[MIN0]] : i32 to index
-// CHECK-NEW:         %[[IDX0:.*]] = xla.apply_indexing #[[MAP]](%[[CAST0]])
-// CHECK-NEW:         %[[EXTRACT1:.*]] = xtile.extract %[[ARG0]][%[[IDX0]]] [1] [{{.*}}] : memref<64xf32> -> tensor<1xf32>
-// CHECK-NEW:         %[[UNSPLAT1:.*]] = tt.unsplat %{{.*}} : tensor<1xi32>
-// CHECK-NEW:         %[[MAX1:.*]] = arith.maxsi %[[UNSPLAT1]], %[[C0_I32]] : i32
-// CHECK-NEW:         %[[MIN1:.*]] = arith.minsi %[[MAX1]], %[[C10]] : i32
-// CHECK-NEW:         %[[CAST1:.*]] = arith.index_cast %[[MIN1]] : i32 to index
-// CHECK-NEW:         %[[IDX1:.*]] = xla.apply_indexing #[[MAP1]](%[[CAST1]], %[[CAST0]])
-// CHECK-NEW:         %[[EXTRACT2:.*]] = xtile.extract %[[ARG0]][%[[IDX1]]] [32] [{{.*}}] : memref<64xf32> -> tensor<32xf32>
-// CHECK-NEW:         xtile.insert %{{.*}} into %[[ARG1]][%[[C0]]] [32] [{{.*}}] : tensor<32xi32> -> memref<10xi32>
-// CHECK-NEW:         xtile.return
-//
-// CHECK-OLD-COUNT-2: arith.minsi
+// CHECK-DAG: #[[MAP:.*]] = #xla.indexing_map<"(d0{{.*}}) -> (d0{{.*}} + 1), domain: d0{{.*}} in [0, 44]">
+// CHECK-DAG: #[[MAP1:.*]] = #xla.indexing_map<"(d0{{.*}}, d1{{.*}}) -> (d0{{.*}} + d1{{.*}}), domain: d0{{.*}} in [0, 10], d1{{.*}} in [0, 44]">
+// CHECK: xtile.entry_func @triton_fn(
+// CHECK-SAME:    %[[ARG0:.*]]: memref<64xf32>,
+// CHECK-SAME:    %[[ARG1:.*]]: memref<10xi32>,
+// CHECK-SAME:    %[[ARG2:.*]]: index{{( \{xla.range = \[0 : index, 0 : index\]\})?}})
+// CHECK-DAG:     %[[C10:.*]] = arith.constant 10 : i32
+// CHECK-DAG:     %[[C44:.*]] = arith.constant 44 : i32
+// CHECK-DAG:     %[[C0_I32:.*]] = arith.constant 0 : i32
+// CHECK-DAG:     %[[C0:.*]] = arith.constant 0 : index
+// CHECK:         %[[EXTRACT0:.*]] = xtile.extract %[[ARG0]][%[[C0]]] [1] [{{.*}}] : memref<64xf32> -> tensor<1xf32>
+// CHECK:         %[[UNSPLAT0:.*]] = tt.unsplat %{{.*}} : tensor<1xi32>
+// CHECK:         %[[MAX0:.*]] = arith.maxsi %[[UNSPLAT0]], %[[C0_I32]] : i32
+// CHECK:         %[[MIN0:.*]] = arith.minsi %[[MAX0]], %[[C44]] : i32
+// CHECK:         %[[CAST0:.*]] = arith.index_cast %[[MIN0]] : i32 to index
+// CHECK:         %[[IDX0:.*]] = xla.apply_indexing #[[MAP]](%[[CAST0]])
+// CHECK:         %[[EXTRACT1:.*]] = xtile.extract %[[ARG0]][%[[IDX0]]] [1] [{{.*}}] : memref<64xf32> -> tensor<1xf32>
+// CHECK:         %[[UNSPLAT1:.*]] = tt.unsplat %{{.*}} : tensor<1xi32>
+// CHECK:         %[[MAX1:.*]] = arith.maxsi %[[UNSPLAT1]], %[[C0_I32]] : i32
+// CHECK:         %[[MIN1:.*]] = arith.minsi %[[MAX1]], %[[C10]] : i32
+// CHECK:         %[[CAST1:.*]] = arith.index_cast %[[MIN1]] : i32 to index
+// CHECK:         %[[IDX1:.*]] = xla.apply_indexing #[[MAP1]](%[[CAST1]], %[[CAST0]])
+// CHECK:         %[[EXTRACT2:.*]] = xtile.extract %[[ARG0]][%[[IDX1]]] [32] [{{.*}}] : memref<64xf32> -> tensor<32xf32>
+// CHECK:         xtile.insert %{{.*}} into %[[ARG1]][%[[C0]]] [32] [{{.*}}] : tensor<32xi32> -> memref<10xi32>
+// CHECK:         xtile.return
 
 ENTRY main {
   p0 = f32[64] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_dyn_slice_no_convert.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_dyn_slice_no_convert.hlo
index 865050e7a74..d7a76afac5a 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_dyn_slice_no_convert.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_of_dyn_slice_no_convert.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s --check-prefixes=CHECK-OLD
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s --check-prefixes=CHECK-NEW
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p0 = s32[64] parameter(0)
@@ -13,33 +10,31 @@ fusion {
   off2 = s32[] reshape(slice2)
   ROOT d2 = s32[10] dynamic-slice(d1, off2), dynamic_slice_sizes={10}
 }
-// CHECK-NEW-DAG: #[[MAP:.*]] = #xla.indexing_map<"(d0{{.*}}) -> (d0{{.*}} + 1), domain: d0{{.*}} in [0, 44]">
-// CHECK-NEW-DAG: #[[MAP1:.*]] = #xla.indexing_map<"(d0{{.*}}, d1{{.*}}) -> (d0{{.*}} + d1{{.*}}), domain: d0{{.*}} in [0, 10], d1{{.*}} in [0, 44]">
-// CHECK-NEW: xtile.entry_func @triton_fn(
-// CHECK-NEW-SAME:    %[[ARG0:.*]]: memref<64xi32>,
-// CHECK-NEW-SAME:    %[[ARG1:.*]]: memref<10xi32>,
-// CHECK-NEW-SAME:    %[[ARG2:.*]]: index{{( \{xla.range = \[0 : index, 0 : index\]\})?}})
-// CHECK-NEW-DAG:     %[[C10:.*]] = arith.constant 10 : i32
-// CHECK-NEW-DAG:     %[[C44:.*]] = arith.constant 44 : i32
-// CHECK-NEW-DAG:     %[[C0_I32:.*]] = arith.constant 0 : i32
-// CHECK-NEW-DAG:     %[[C0:.*]] = arith.constant 0 : index
-// CHECK-NEW:         %[[EXTRACT0:.*]] = xtile.extract %[[ARG0]][%[[C0]]] [1] [{{.*}}] : memref<64xi32> -> tensor<1xi32>
-// CHECK-NEW:         %[[UNSPLAT0:.*]] = tt.unsplat %[[EXTRACT0]] : tensor<1xi32>
-// CHECK-NEW:         %[[MAX0:.*]] = arith.maxsi %[[UNSPLAT0]], %[[C0_I32]] : i32
-// CHECK-NEW:         %[[MIN0:.*]] = arith.minsi %[[MAX0]], %[[C44]] : i32
-// CHECK-NEW:         %[[CAST0:.*]] = arith.index_cast %[[MIN0]] : i32 to index
-// CHECK-NEW:         %[[IDX0:.*]] = xla.apply_indexing #[[MAP]](%[[CAST0]])
-// CHECK-NEW:         %[[EXTRACT1:.*]] = xtile.extract %[[ARG0]][%[[IDX0]]] [1] [{{.*}}] : memref<64xi32> -> tensor<1xi32>
-// CHECK-NEW:         %[[UNSPLAT1:.*]] = tt.unsplat %[[EXTRACT1]] : tensor<1xi32>
-// CHECK-NEW:         %[[MAX1:.*]] = arith.maxsi %[[UNSPLAT1]], %[[C0_I32]] : i32
-// CHECK-NEW:         %[[MIN1:.*]] = arith.minsi %[[MAX1]], %[[C10]] : i32
-// CHECK-NEW:         %[[CAST1:.*]] = arith.index_cast %[[MIN1]] : i32 to index
-// CHECK-NEW:         %[[IDX1:.*]] = xla.apply_indexing #[[MAP1]](%[[CAST1]], %[[CAST0]])
-// CHECK-NEW:         %[[EXTRACT2:.*]] = xtile.extract %[[ARG0]][%[[IDX1]]] [32] [{{.*}}] : memref<64xi32> -> tensor<32xi32>
-// CHECK-NEW:         xtile.insert %[[EXTRACT2]] into %[[ARG1]][%[[C0]]] [32] [{{.*}}] : tensor<32xi32> -> memref<10xi32>
-// CHECK-NEW:         xtile.return
-//
-// CHECK-OLD-COUNT-2: arith.minsi
+// CHECK-DAG: #[[MAP:.*]] = #xla.indexing_map<"(d0{{.*}}) -> (d0{{.*}} + 1), domain: d0{{.*}} in [0, 44]">
+// CHECK-DAG: #[[MAP1:.*]] = #xla.indexing_map<"(d0{{.*}}, d1{{.*}}) -> (d0{{.*}} + d1{{.*}}), domain: d0{{.*}} in [0, 10], d1{{.*}} in [0, 44]">
+// CHECK: xtile.entry_func @triton_fn(
+// CHECK-SAME:    %[[ARG0:.*]]: memref<64xi32>,
+// CHECK-SAME:    %[[ARG1:.*]]: memref<10xi32>,
+// CHECK-SAME:    %[[ARG2:.*]]: index{{( \{xla.range = \[0 : index, 0 : index\]\})?}})
+// CHECK-DAG:     %[[C10:.*]] = arith.constant 10 : i32
+// CHECK-DAG:     %[[C44:.*]] = arith.constant 44 : i32
+// CHECK-DAG:     %[[C0_I32:.*]] = arith.constant 0 : i32
+// CHECK-DAG:     %[[C0:.*]] = arith.constant 0 : index
+// CHECK:         %[[EXTRACT0:.*]] = xtile.extract %[[ARG0]][%[[C0]]] [1] [{{.*}}] : memref<64xi32> -> tensor<1xi32>
+// CHECK:         %[[UNSPLAT0:.*]] = tt.unsplat %[[EXTRACT0]] : tensor<1xi32>
+// CHECK:         %[[MAX0:.*]] = arith.maxsi %[[UNSPLAT0]], %[[C0_I32]] : i32
+// CHECK:         %[[MIN0:.*]] = arith.minsi %[[MAX0]], %[[C44]] : i32
+// CHECK:         %[[CAST0:.*]] = arith.index_cast %[[MIN0]] : i32 to index
+// CHECK:         %[[IDX0:.*]] = xla.apply_indexing #[[MAP]](%[[CAST0]])
+// CHECK:         %[[EXTRACT1:.*]] = xtile.extract %[[ARG0]][%[[IDX0]]] [1] [{{.*}}] : memref<64xi32> -> tensor<1xi32>
+// CHECK:         %[[UNSPLAT1:.*]] = tt.unsplat %[[EXTRACT1]] : tensor<1xi32>
+// CHECK:         %[[MAX1:.*]] = arith.maxsi %[[UNSPLAT1]], %[[C0_I32]] : i32
+// CHECK:         %[[MIN1:.*]] = arith.minsi %[[MAX1]], %[[C10]] : i32
+// CHECK:         %[[CAST1:.*]] = arith.index_cast %[[MIN1]] : i32 to index
+// CHECK:         %[[IDX1:.*]] = xla.apply_indexing #[[MAP1]](%[[CAST1]], %[[CAST0]])
+// CHECK:         %[[EXTRACT2:.*]] = xtile.extract %[[ARG0]][%[[IDX1]]] [32] [{{.*}}] : memref<64xi32> -> tensor<32xi32>
+// CHECK:         xtile.insert %[[EXTRACT2]] into %[[ARG1]][%[[C0]]] [32] [{{.*}}] : tensor<32xi32> -> memref<10xi32>
+// CHECK:         xtile.return
 
 ENTRY main {
   p0 = s32[64] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_offset_data_same_input.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_offset_data_same_input.hlo
index cdeda3bc59a..17e17e668a9 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_offset_data_same_input.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/dynamic_slice/dyn_slice_offset_data_same_input.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p0 = s32[64] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/chain_0d_scalars.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/chain_0d_scalars.hlo
index 4132e32a191..883dc86eafe 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/chain_0d_scalars.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/chain_0d_scalars.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s  --abs_error_bound=6e-1 --rel_error_bound=6e-1
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --abs_error_bound=6e-1 --rel_error_bound=6e-1 --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s --abs_error_bound=6e-1 --rel_error_bound=6e-1
 
 fusion {
   p0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/convert_int_to_pred.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/convert_int_to_pred.hlo
index 138790a1f83..b55988b8b7e 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/convert_int_to_pred.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/convert_int_to_pred.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p0 = s32[3,2,2]{2,1,0} parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/div_by_zero.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/div_by_zero.hlo
index 80e0fafad70..17db49f6d4c 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/div_by_zero.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/div_by_zero.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   param_0 = s32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/div_const_denom_by_zero.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/div_const_denom_by_zero.hlo
index f0e0fc3d4a3..0034ce840ac 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/div_const_denom_by_zero.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/div_const_denom_by_zero.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   param_0 = s32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/div_consts_by_zero.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/div_consts_by_zero.hlo
index 95aa04cf686..fc2f657f6f6 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/div_consts_by_zero.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/div_consts_by_zero.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   numerator = s32[] constant(10)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/div_overflow.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/div_overflow.hlo
index 019f7759759..29b2c047cab 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/div_overflow.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/div_overflow.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   param_0 = s32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum.hlo
index bca07c3dea3..bdd9c8b4166 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum.hlo
@@ -1,5 +1,5 @@
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p0 = f32[8,4] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_three_tiles_per_block.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_three_tiles_per_block.hlo
index 513ac7a83b3..ae1de1697e1 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_three_tiles_per_block.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_three_tiles_per_block.hlo
@@ -1,6 +1,5 @@
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation=false
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p0 = f32[8,4] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_two_tiles_per_block.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_two_tiles_per_block.hlo
index 49698a9ce80..77fd9c42a97 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_two_tiles_per_block.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/minimum_two_tiles_per_block.hlo
@@ -1,6 +1,5 @@
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation=false
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p0 = f32[8,4] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/nested_reducer_fusion.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/nested_reducer_fusion.hlo
index 14654a6a12e..6e09369d175 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/nested_reducer_fusion.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/nested_reducer_fusion.hlo
@@ -1,5 +1,4 @@
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s --abs_error_bound=0.0 --rel_error_bound=0.0
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s --abs_error_bound=0.0 --rel_error_bound=0.0
+// RUN: triton_test_correctness %s --abs_error_bound=0.0 --rel_error_bound=0.0
 
 fusion {
   p0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/not.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/not.hlo
index 4523f518110..35d5bb193fb 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/not.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/not.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   param_0 = s32[100] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/pred_add.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/pred_add.hlo
index 82c92436853..641016a0b21 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/pred_add.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/pred_add.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   param_0 = pred[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/pred_input.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/pred_input.hlo
index 47e1e432bbe..44e4032bc3f 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/pred_input.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/pred_input.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   param_0 = pred[15] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/pred_output.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/pred_output.hlo
index e08ba31da13..bae50729ebd 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/pred_output.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/pred_output.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   param_0 = f32[15] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/reduce_precision.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/reduce_precision.hlo
index a64b1167487..d8f8a5f4bf9 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/reduce_precision.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/reduce_precision.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = f32[5,7] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/reduce_precision_noop.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/reduce_precision_noop.hlo
index 05100143e56..61abf4f5f47 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/reduce_precision_noop.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/elementwise/reduce_precision_noop.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = bf16[5,7] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc
index a6cfab6a822..71e881d6dbe 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc
@@ -112,15 +112,10 @@ const HloFusionInstruction& GetFusionInstruction(
 
 constexpr ErrorSpec kExactMatch{/*aabs=*/0, /*arel=*/0};
 
-std::string TilingParametersToString(bool tiling_propagation_enabled) {
-  return tiling_propagation_enabled ? "ExperimentalTiling" : "SymbolicTiling";
-}
-
 class TritonEmitterTest
     : public HloInterpreterReferenceMixin<GpuPjRtCodegenTest>,
       public XTileTestBase {
  public:
-  virtual bool EnableTilingPropagation() const = 0;
   DebugOptions GetDebugOptionsForTest() const override {
     // TODO: b/509502550 - remove the flag and disable tests that use
     // multi-output fusions when removing the feature.
@@ -129,8 +124,6 @@ class TritonEmitterTest
     debug_options.set_xla_gpu_unsupported_enable_triton_multi_output_fusion(
         true);
     debug_options.set_xla_gpu_experimental_disable_binary_libraries(true);
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
-        EnableTilingPropagation());
     return debug_options;
   }
 
@@ -168,43 +161,17 @@ class TritonEmitterTest
   }
 };
 
-class TritonEmitterTestWithTilingParam
-    : public TritonEmitterTest,
-      public ::testing::WithParamInterface<bool> {
- public:
-  bool EnableTilingPropagation() const override { return GetParam(); }
-};
-
-INSTANTIATE_TEST_SUITE_P(TritonEmitterTestWithTilingParamTestSuite,
-                         TritonEmitterTestWithTilingParam, ::testing::Bool(),
-                         [](const ::testing::TestParamInfo<bool>& info) {
-                           return TilingParametersToString(info.param);
-                         });
-
 class TmaParameterizedTritonEmitterTest
     : public TritonEmitterTest,
-      public ::testing::WithParamInterface<std::tuple<bool, bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<1>(GetParam());
-  }
-};
+      public ::testing::WithParamInterface<bool> {};
 
-INSTANTIATE_TEST_SUITE_P(
-    TmaParameterizedTritonEmitterTestSuite, TmaParameterizedTritonEmitterTest,
-    ::testing::Combine(::testing::Bool(), ::testing::Bool()),
-    [](const ::testing::TestParamInfo<std::tuple<bool, bool>>& info) {
-      return absl::StrCat(
-          std::get<0>(info.param) ? "tma_allowed" : "tma_disabled",
-          TilingParametersToString(std::get<1>(info.param)));
-    });
-
-class WarpSpecializationTritonEmitterTest
-    : public TritonEmitterTest,
-      public ::testing::WithParamInterface<bool> {
- public:
-  bool EnableTilingPropagation() const override { return GetParam(); }
+INSTANTIATE_TEST_SUITE_P(TmaParameterizedTritonEmitterTestSuite,
+                         TmaParameterizedTritonEmitterTest, ::testing::Bool(),
+                         [](const ::testing::TestParamInfo<bool>& info) {
+                           return info.param ? "tma_allowed" : "tma_disabled";
+                         });
 
+class WarpSpecializationTritonEmitterTest : public TritonEmitterTest {
  public:
   DebugOptions GetDebugOptionsForTest() const override {
     DebugOptions debug_options = TritonEmitterTest::GetDebugOptionsForTest();
@@ -214,13 +181,7 @@ class WarpSpecializationTritonEmitterTest
   }
 };
 
-INSTANTIATE_TEST_SUITE_P(WarpSpecializationTritonEmitterTestSuite,
-                         WarpSpecializationTritonEmitterTest, ::testing::Bool(),
-                         [](const ::testing::TestParamInfo<bool>& info) {
-                           return TilingParametersToString(info.param);
-                         });
-
-TEST_P(TritonEmitterTestWithTilingParam, BitcastReduceWithStride4Tiling) {
+TEST_F(TritonEmitterTest, BitcastReduceWithStride4Tiling) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -255,40 +216,26 @@ ENTRY entry_computation {
 })";
   auto status =
       CreateTritonIrFromHloTextAndFileCheck(kHloText, "fused_computation", "");
-  if (EnableTilingPropagation()) {
-    EXPECT_THAT(
-        status,
-        StatusIs(
-            tsl::error::UNIMPLEMENTED,
-            HasSubstr("Only single-result fusions are supported for now")));
-  } else {
-    EXPECT_THAT(status,
-                StatusIs(tsl::error::UNIMPLEMENTED,
-                         HasSubstr("Unsupported case of multi-output fusion")));
-  }
+  EXPECT_THAT(
+      status,
+      StatusIs(tsl::error::UNIMPLEMENTED,
+               HasSubstr("Only single-result fusions are supported for now")));
 }
 
 class TritonEmitterTestWithOffsetParam
     : public TritonEmitterTest,
-      public ::testing::WithParamInterface<std::tuple<int32_t, bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<1>(GetParam());
-  }
-};
+      public ::testing::WithParamInterface<int32_t> {};
 
 using EmitDynamicSliceTest = TritonEmitterTestWithOffsetParam;
 
-INSTANTIATE_TEST_SUITE_P(
-    DynamicSliceSuite, EmitDynamicSliceTest,
-    ::testing::Combine(::testing::Values(0, 1, 10, 100), ::testing::Bool()),
-    [](const ::testing::TestParamInfo<std::tuple<int32_t, bool>>& info) {
-      return absl::StrCat("offset_", std::get<0>(info.param),
-                          TilingParametersToString(std::get<1>(info.param)));
-    });
+INSTANTIATE_TEST_SUITE_P(DynamicSliceSuite, EmitDynamicSliceTest,
+                         ::testing::Values(0, 1, 10, 100),
+                         [](const ::testing::TestParamInfo<int32_t>& info) {
+                           return absl::StrCat("offset_", info.param);
+                         });
 
 TEST_P(EmitDynamicSliceTest, LowerDynamicSliceWithSingleDimension) {
-  int32_t offset = std::get<0>(GetParam());
+  int32_t offset = GetParam();
   constexpr absl::string_view kHloText = R"(
 f {
   p0 = f32[64] parameter(0)
@@ -340,7 +287,7 @@ ENTRY entry_computation {
                        ParseAndReturnVerifiedModule(kHloText));
   ASSERT_OK_AND_ASSIGN(std::vector<Literal> parameters,
                        MakeFakeArguments(module.get()));
-  int32_t offset = std::get<0>(GetParam());
+  int32_t offset = GetParam();
   parameters[1].Set<int32_t>({}, offset);
   parameters[2].Set<int32_t>({}, offset);
   EXPECT_TRUE(RunAndCompareNoHloPasses(
@@ -370,7 +317,7 @@ ENTRY entry_computation {
                        ParseAndReturnVerifiedModule(kHloText));
   ASSERT_OK_AND_ASSIGN(std::vector<Literal> parameters,
                        MakeFakeArguments(module.get()));
-  int32_t offset = std::get<0>(GetParam());
+  int32_t offset = GetParam();
   parameters[1].Set<int32_t>({}, offset);
   parameters[2].Set<int32_t>({}, offset);
   EXPECT_TRUE(RunAndCompareNoHloPasses(
@@ -378,7 +325,7 @@ ENTRY entry_computation {
 }
 
 TEST_P(EmitDynamicSliceTest, LowerDynamicSliceWithConstantOffset) {
-  int32_t offset = std::get<0>(GetParam());
+  int32_t offset = GetParam();
   std::string kHloText =
       absl::StrReplaceAll(R"(
 f {
@@ -435,13 +382,13 @@ ENTRY entry {
                        ParseAndReturnVerifiedModule(kHloText));
   ASSERT_OK_AND_ASSIGN(std::vector<Literal> parameters,
                        MakeFakeArguments(module.get()));
-  int32_t offset = std::get<0>(GetParam());
+  int32_t offset = GetParam();
   parameters[2].Set<int32_t>({}, offset);
   EXPECT_TRUE(RunAndCompareNoHloPasses(
       std::move(module), LiteralUtil::MakePointers(parameters), kExactMatch));
 }
 
-TEST_P(TritonEmitterTestWithTilingParam, LowerDynamicSliceOfAdd) {
+TEST_F(TritonEmitterTest, LowerDynamicSliceOfAdd) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -472,29 +419,11 @@ ENTRY entry_computation {
 }
 
 class TritonDevicelessTest : public HloHardwareIndependentTestBase,
-                             public XTileTestBase,
-                             public ::testing::WithParamInterface<bool> {
- public:
-  bool EnableTilingPropagation() const { return GetParam(); }
-
-  DebugOptions GetDebugOptionsForTest() const override {
-    DebugOptions debug_options =
-        HloHardwareIndependentTestBase::GetDebugOptionsForTest();
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
-        EnableTilingPropagation());
-    return debug_options;
-  }
-};
-
-INSTANTIATE_TEST_SUITE_P(TritonDevicelessTestSuite, TritonDevicelessTest,
-                         ::testing::Bool(),
-                         [](const ::testing::TestParamInfo<bool>& info) {
-                           return TilingParametersToString(info.param);
-                         });
+                             public XTileTestBase {};
 
 // TODO(b/353484968): Tests that don't run RunAndCompareNoHloPasses should be
 // moved to deviceless test file.
-TEST_P(TritonDevicelessTest, TestGenericEmitterWithSoftMaxSingleParameter) {
+TEST_F(TritonDevicelessTest, TestGenericEmitterWithSoftMaxSingleParameter) {
   constexpr absl::string_view kHloText = R"(
 HloModule t
 add {
@@ -564,7 +493,7 @@ CHECK:        }
 
 // TODO(b/353484968): Tests that don't run RunAndCompareNoHloPasses should be
 // moved to deviceless test file.
-TEST_P(TritonDevicelessTest, TestGenericEmitterWithMultipleParameters) {
+TEST_F(TritonDevicelessTest, TestGenericEmitterWithMultipleParameters) {
   constexpr absl::string_view kHloText = R"(
 HloModule t
 
@@ -642,7 +571,7 @@ CHECK-SAME:       [%[[TID]], %{{.*}}] [1, 128] [1, 1] : tensor<1x128xf32>
 
 // TODO(b/353484968): Tests that don't run RunAndCompareNoHloPasses should be
 // moved to deviceless test file.
-TEST_P(TritonDevicelessTest, EmitterFailsIfComputeCapabilityIsBelowAmpere) {
+TEST_F(TritonDevicelessTest, EmitterFailsIfComputeCapabilityIsBelowAmpere) {
   constexpr absl::string_view kHloText = R"(
 triton_computation {
   p0 = f32[10,10] parameter(0)
@@ -685,7 +614,7 @@ ENTRY entry {
                          "(compute capability 8.0) and up, but got")));
 }
 
-TEST_P(TritonDevicelessTest, RejectsPackedFp4OddMinorOffset) {
+TEST_F(TritonDevicelessTest, RejectsPackedFp4OddMinorOffset) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -724,7 +653,7 @@ ENTRY e {
                                "to be divisible by 2")));
 }
 
-TEST_P(TritonDevicelessTest, EmitsPackedFp4StorageForEvenMinorOffset) {
+TEST_F(TritonDevicelessTest, EmitsPackedFp4StorageForEvenMinorOffset) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -773,7 +702,7 @@ CHECK-SAME: lhs_elem_type = f4E2M1FN, rhs_elem_type = f4E2M1FN
 
 // TODO(b/353484968): Tests that don't run RunAndCompareNoHloPasses should be
 // moved to deviceless test file.
-TEST_P(TritonDevicelessTest,
+TEST_F(TritonDevicelessTest,
        EmitterFailsIfFusionBackendConfigDoesNotSatisfyConstraints) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> hlo_module,
                        ParseAndReturnVerifiedModule(R"(
@@ -833,7 +762,7 @@ ENTRY entry_computation {
 
 // TODO(b/353484968): Tests that don't run RunAndCompareNoHloPasses should b
 // moved to deviceless test file.
-TEST_P(TritonDevicelessTest, TestGenericEmitterReductionFusion) {
+TEST_F(TritonDevicelessTest, TestGenericEmitterReductionFusion) {
   constexpr absl::string_view kHloText = R"(
 HloModule t
 add {
@@ -897,7 +826,7 @@ CHECK:            xtile.insert {{.*}} : tensor<1xf32>
                            "triton_reduction_computation")));
 }
 
-TEST_P(TritonEmitterTestWithTilingParam,
+TEST_F(TritonEmitterTest,
        TestGenericEmitterWithReductionAndMultidimensionalTile) {
   constexpr absl::string_view kHloText = R"(
 HloModule t
@@ -930,8 +859,7 @@ ENTRY main {
   EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
 }
 
-TEST_P(TritonEmitterTestWithTilingParam,
-       TestSoftMaxWithTileElementsNotAllContiguous) {
+TEST_F(TritonEmitterTest, TestSoftMaxWithTileElementsNotAllContiguous) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -995,7 +923,7 @@ ENTRY entry_computation {
         "is_tma_allowed":"$0"}}}
 })";
 
-  const bool is_tma_allowed = std::get<0>(GetParam());
+  const bool is_tma_allowed = GetParam();
   const std::string hlo_text =
       absl::Substitute(kHloTextTemplate, is_tma_allowed);
   EXPECT_TRUE(RunAndCompareNoHloPasses(hlo_text, kExactMatch));
@@ -1025,14 +953,13 @@ ENTRY entry_computation {
         "num_stages":"1",
         "is_tma_allowed":"$0"}}}
 })";
-  const bool is_tma_allowed = std::get<0>(GetParam());
+  const bool is_tma_allowed = GetParam();
   const std::string hlo_text =
       absl::Substitute(kHloTextTemplate, is_tma_allowed);
   EXPECT_TRUE(RunAndCompareNoHloPasses(hlo_text, kExactMatch));
 }
 
-TEST_P(TritonEmitterTestWithTilingParam,
-       TestSliceWithTileElementsNotAllContiguous) {
+TEST_F(TritonEmitterTest, TestSliceWithTileElementsNotAllContiguous) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -1086,14 +1013,13 @@ ENTRY entry_computation {
         "is_tma_allowed":"$0"}}}
 })";
 
-  const bool is_tma_allowed = std::get<0>(GetParam());
+  const bool is_tma_allowed = GetParam();
   const std::string hlo_text =
       absl::Substitute(kHloTextTemplate, is_tma_allowed);
   EXPECT_TRUE(RunAndCompareNoHloPasses(hlo_text, kExactMatch));
 }
 
-TEST_P(TritonEmitterTestWithTilingParam,
-       TestSliceWithTileElementsNotAllContiguousUnaligned) {
+TEST_F(TritonEmitterTest, TestSliceWithTileElementsNotAllContiguousUnaligned) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -1146,7 +1072,7 @@ ENTRY entry_computation {
           "is_tma_allowed":"$0"}}}
 })";
 
-  const bool is_tma_allowed = std::get<0>(GetParam());
+  const bool is_tma_allowed = GetParam();
   const std::string hlo_text =
       absl::Substitute(kHloTextTemplate, is_tma_allowed);
   EXPECT_TRUE(RunAndCompareNoHloPasses(hlo_text, kExactMatch));
@@ -1176,7 +1102,7 @@ ENTRY main {
           "is_tma_allowed":"$0"}}}
 })";
 
-  const bool is_tma_allowed = std::get<0>(GetParam());
+  const bool is_tma_allowed = GetParam();
   const std::string hlo_text =
       absl::Substitute(kHloTextTemplate, is_tma_allowed);
 
@@ -1196,8 +1122,7 @@ CHECK: tt.reshape
   EXPECT_TRUE(RunAndCompareNoHloPasses(hlo_text, kExactMatch));
 }
 
-TEST_P(TritonEmitterTestWithTilingParam,
-       BitcastIntoBroadcastIsLoweredCorrectly) {
+TEST_F(TritonEmitterTest, BitcastIntoBroadcastIsLoweredCorrectly) {
   constexpr absl::string_view kHloText = R"(
 triton_computation {
   param_0 = f32[128,256]{1,0} parameter(0)
@@ -1255,7 +1180,7 @@ ENTRY entry_computation {
       "is_tma_allowed":"$0"}}}
 })";
 
-  const bool is_tma_allowed = std::get<0>(GetParam());
+  const bool is_tma_allowed = GetParam();
   const std::string hlo_text =
       absl::Substitute(kHloTextTemplate, is_tma_allowed);
   EXPECT_TRUE(RunAndCompareNoHloPasses(hlo_text, kExactMatch));
@@ -1285,7 +1210,7 @@ ENTRY entry_computation {
         "is_tma_allowed":"$0"}}}
 })";
 
-  const bool is_tma_allowed = std::get<0>(GetParam());
+  const bool is_tma_allowed = GetParam();
   const std::string hlo_text =
       absl::Substitute(kHloTextTemplate, is_tma_allowed);
   EXPECT_TRUE(RunAndCompareNoHloPasses(hlo_text, kExactMatch));
@@ -1315,7 +1240,7 @@ backend_config={
    "is_tma_allowed":"$0"}}}
 })";
 
-  const bool is_tma_allowed = std::get<0>(GetParam());
+  const bool is_tma_allowed = GetParam();
   const std::string hlo_text =
       absl::Substitute(kHloTextTemplate, is_tma_allowed);
   EXPECT_TRUE(RunAndCompareNoHloPasses(hlo_text, kExactMatch));
@@ -1347,7 +1272,7 @@ backend_config={
    "is_tma_allowed":"$0"}}}
 })";
 
-  const bool is_tma_allowed = std::get<0>(GetParam());
+  const bool is_tma_allowed = GetParam();
   const std::string hlo_text =
       absl::Substitute(kHloTextTemplate, is_tma_allowed);
   EXPECT_TRUE(RunAndCompareNoHloPasses(hlo_text, kExactMatch));
@@ -1378,7 +1303,7 @@ backend_config={
  "is_tma_allowed":"$0"}}}
 })";
 
-  const bool is_tma_allowed = std::get<0>(GetParam());
+  const bool is_tma_allowed = GetParam();
   const std::string hlo_text =
       absl::Substitute(kHloTextTemplate, is_tma_allowed);
   EXPECT_TRUE(RunAndCompareNoHloPasses(hlo_text, kExactMatch));
@@ -1386,7 +1311,7 @@ backend_config={
 
 // TODO(b/390559452): Capture the iteration order from the propagated tiling.
 // When computing the tiling separately we need to use the same iteration order.
-TEST_P(TritonEmitterTestWithTilingParam, DISABLED_Transpose3DWithExtraOutput) {
+TEST_F(TritonEmitterTest, DISABLED_Transpose3DWithExtraOutput) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -1437,15 +1362,10 @@ CHECK-COUNT-2: xtile.insert
 
 class IotaEmitterParametrizedTest
     : public TritonEmitterTest,
-      public ::testing::WithParamInterface<std::tuple<PrimitiveType, bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<1>(GetParam());
-  }
-};
+      public ::testing::WithParamInterface<PrimitiveType> {};
 
 TEST_P(IotaEmitterParametrizedTest, Iota4DIsCodegeneratedCorrectly) {
-  auto data_type = std::get<0>(GetParam());
+  auto data_type = GetParam();
   const std::string kHloText =
       absl::Substitute(R"(
 triton_computation {
@@ -1491,22 +1411,18 @@ CHECK:      tt.broadcast {{.*}} -> tensor<1x2x64x8x
 }
 
 std::string TypeTestParamToString(
-    const ::testing::TestParamInfo<std::tuple<PrimitiveType, bool>>& data) {
-  return absl::StrCat(
-      primitive_util::LowercasePrimitiveTypeName(std::get<0>(data.param)),
-      TilingParametersToString(std::get<1>(data.param)));
+    const ::testing::TestParamInfo<PrimitiveType>& data) {
+  return std::string(primitive_util::LowercasePrimitiveTypeName(data.param));
 }
 
 INSTANTIATE_TEST_SUITE_P(IotaEmitterParametrizedTestSuite,
                          IotaEmitterParametrizedTest,
-                         ::testing::Combine(::testing::ValuesIn({S8, S16, S32,
-                                                                 S64, BF16, F16,
-                                                                 F32, F64}),
-                                            ::testing::Bool()),
+                         ::testing::ValuesIn({S8, S16, S32, S64, BF16, F16, F32,
+                                              F64}),
                          TypeTestParamToString);
 
 // Reproducer from b/384110192.
-TEST_P(TritonEmitterTestWithTilingParam,
+TEST_F(TritonEmitterTest,
        FusionWithOutputContainingMoreThanInt32MaxElementsExecutesCorrectly) {
   if (GpuComputeCapability().IsRocm()) {
     GTEST_SKIP() << "Requires more than 4GB GPU memory, exceeds ROCm RBE "
@@ -1567,7 +1483,7 @@ ENTRY entry_computation {
                                       /*run_hlo_passes=*/false));
 }
 
-TEST_P(TritonEmitterTestWithTilingParam, ConvertF16ToF8E5M2Exhaustive) {
+TEST_F(TritonEmitterTest, ConvertF16ToF8E5M2Exhaustive) {
   // TODO(b/396595945): enable post-Ampere once Triton respects RTNE semantics
   // on H100.
   if (auto cc = GpuComputeCapability().cuda_compute_capability();
@@ -1615,7 +1531,7 @@ ENTRY entry_computation {
   EXPECT_TRUE(RunAndCompareNoHloPasses(std::move(module), kExactMatch));
 }
 
-TEST_P(TritonEmitterTestWithTilingParam, ConvertS4ToS8Exhaustive) {
+TEST_F(TritonEmitterTest, ConvertS4ToS8Exhaustive) {
   constexpr absl::string_view kHloText = R"(
 computation {
   p0 = s4[16]{0:E(4)} parameter(0)
@@ -1667,7 +1583,7 @@ ENTRY entry_computation {
           "is_tma_allowed":"$0"}}}
 })";
 
-  const bool is_tma_allowed = std::get<0>(GetParam());
+  const bool is_tma_allowed = GetParam();
   std::string hlo_text = absl::Substitute(kHloTextTemplate, is_tma_allowed);
   EXPECT_TRUE(RunAndCompareNoHloPasses(hlo_text, kExactMatch));
 }
@@ -1700,7 +1616,7 @@ ENTRY entry {
           "is_tma_allowed":"$0"}}}
 })";
 
-  const bool is_tma_allowed = std::get<0>(GetParam());
+  const bool is_tma_allowed = GetParam();
   std::string hlo_text = absl::Substitute(kHloTextTemplate, is_tma_allowed);
 
   ASSERT_OK_AND_ASSIGN(auto xtile_module_and_hlo_module,
@@ -1749,7 +1665,7 @@ CHECK-COUNT-1: xtile.insert
       hlo_text, ErrorSpec{/*aabs=*/1e-4, /*arel=*/1e-6}));
 }
 
-TEST_P(WarpSpecializationTritonEmitterTest,
+TEST_F(WarpSpecializationTritonEmitterTest,
        DotAccumulationLoopUsesWarpSpecialization) {
   if (auto cc = GpuComputeCapability().cuda_compute_capability();
       cc && !cc->IsAtLeastBlackwell()) {
@@ -1826,7 +1742,7 @@ ENTRY e (p0.1: f32[11,1,24,1], p1.1: f32[128,32]) -> f32[256,32] {
 }
 )";
 
-  const bool is_tma_allowed = std::get<0>(GetParam());
+  const bool is_tma_allowed = GetParam();
   const std::string hlo_text =
       absl::Substitute(kHloTextTemplate, is_tma_allowed);
   EXPECT_TRUE(RunAndCompareNoHloPasses(
@@ -1933,13 +1849,7 @@ ErrorSpec ErrorSpecForDotAlgorithm(PrecisionConfig::Algorithm algorithm) {
 
 class TritonEmitterTestWithAlgorithmParam
     : public TritonEmitterTest,
-      public ::testing::WithParamInterface<
-          std::tuple<PrecisionConfig::Algorithm, bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<1>(GetParam());
-  }
-};
+      public ::testing::WithParamInterface<PrecisionConfig::Algorithm> {};
 
 // Regroups tests for dot algorithms that have no ambiguous type parameters as
 // per `algorithm_util::GetAllowedOperandsTypeForAlgorithm` and
@@ -1958,7 +1868,7 @@ constexpr std::array kBasicAlgorithms = {
 };
 
 TEST_P(BasicDotAlgorithmEmitterTest, BasicAlgorithmIsEmittedCorrectly) {
-  auto algorithm = std::get<0>(GetParam());
+  auto algorithm = GetParam();
   ASSERT_OK_AND_ASSIGN(
       std::vector<PrimitiveType> allowed_types,
       algorithm_util::GetAllowedOperandsTypeForAlgorithm(algorithm));
@@ -1994,17 +1904,14 @@ TEST_P(BasicDotAlgorithmEmitterTest, BasicAlgorithmIsEmittedCorrectly) {
 }
 
 std::string DotAlgorithmTestToString(
-    const ::testing::TestParamInfo<
-        std::tuple<PrecisionConfig::Algorithm, bool>>& data) {
-  return absl::StrCat(PrecisionConfig::Algorithm_Name(std::get<0>(data.param)),
-                      TilingParametersToString(std::get<1>(data.param)));
+    const ::testing::TestParamInfo<PrecisionConfig::Algorithm>& data) {
+  return PrecisionConfig::Algorithm_Name(data.param);
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    BasicDotAlgorithmEmitterTestSuite, BasicDotAlgorithmEmitterTest,
-    ::testing::Combine(::testing::ValuesIn(kBasicAlgorithms),
-                       ::testing::Bool()),
-    DotAlgorithmTestToString);
+INSTANTIATE_TEST_SUITE_P(BasicDotAlgorithmEmitterTestSuite,
+                         BasicDotAlgorithmEmitterTest,
+                         ::testing::ValuesIn(kBasicAlgorithms),
+                         DotAlgorithmTestToString);
 
 // Regroups tests for dot algorithms that issue several dot instructions.
 using MultiDotAlgorithmEmitterTest = TritonEmitterTestWithAlgorithmParam;
@@ -2019,7 +1926,7 @@ constexpr std::array kMultiDotAlgorithms = {
 };
 
 TEST_P(MultiDotAlgorithmEmitterTest, MultiDotAlgorithmIsEmittedCorrectly) {
-  auto algorithm = std::get<0>(GetParam());
+  auto algorithm = GetParam();
   ASSERT_OK_AND_ASSIGN(PrimitiveType out_ty,
                        algorithm_util::GetDotAccumulatorType(algorithm));
   PrimitiveType in_ty =
@@ -2077,11 +1984,10 @@ TEST_P(MultiDotAlgorithmEmitterTest, MultiDotAlgorithmIsEmittedCorrectly) {
       RunAndCompareNoHloPasses(kHloText, ErrorSpecForDotAlgorithm(algorithm)));
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    MultiDotAlgorithmEmitterTestSuite, MultiDotAlgorithmEmitterTest,
-    ::testing::Combine(::testing::ValuesIn(kMultiDotAlgorithms),
-                       ::testing::Bool()),
-    DotAlgorithmTestToString);
+INSTANTIATE_TEST_SUITE_P(MultiDotAlgorithmEmitterTestSuite,
+                         MultiDotAlgorithmEmitterTest,
+                         ::testing::ValuesIn(kMultiDotAlgorithms),
+                         DotAlgorithmTestToString);
 
 // Regroups tests that use TF32 precision by definition.
 using TF32DotAlgorithmEmitterTest = TritonEmitterTestWithAlgorithmParam;
@@ -2091,7 +1997,7 @@ constexpr std::array kTF32DotAlgorithms = {
     PrecisionConfig::ALG_DOT_TF32_TF32_F32_X3};
 
 TEST_P(TF32DotAlgorithmEmitterTest, TF32AlgorithmsUseTF32InputPrecision) {
-  auto algorithm = std::get<0>(GetParam());
+  auto algorithm = GetParam();
   ASSERT_OK_AND_ASSIGN(
       std::vector<PrimitiveType> allowed_types,
       algorithm_util::GetAllowedOperandsTypeForAlgorithm(algorithm));
@@ -2135,34 +2041,28 @@ TEST_P(TF32DotAlgorithmEmitterTest, TF32AlgorithmsUseTF32InputPrecision) {
   // other tests.
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    TF32DotAlgorithmEmitterTestSuite, TF32DotAlgorithmEmitterTest,
-    ::testing::Combine(::testing::ValuesIn(kTF32DotAlgorithms),
-                       ::testing::Bool()),
-    DotAlgorithmTestToString);
+INSTANTIATE_TEST_SUITE_P(TF32DotAlgorithmEmitterTestSuite,
+                         TF32DotAlgorithmEmitterTest,
+                         ::testing::ValuesIn(kTF32DotAlgorithms),
+                         DotAlgorithmTestToString);
 
 class DotUnsetAlgorithmEmitterTest
     : public TritonEmitterTest,
       public ::testing::WithParamInterface<
-          std::tuple<PrimitiveType, PrimitiveType, bool>> {
+          std::tuple<PrimitiveType, PrimitiveType>> {
  public:
-  bool EnableTilingPropagation() const override {
-    return std::get<2>(GetParam());
-  }
-
   static std::string ParamToString(
       const ::testing::TestParamInfo<DotUnsetAlgorithmEmitterTest::ParamType>&
           data) {
-    auto [result_type, input_type, tiling_enabled] = data.param;
+    auto [result_type, input_type] = data.param;
     return absl::StrCat(primitive_util::LowercasePrimitiveTypeName(result_type),
                         "_",
-                        primitive_util::LowercasePrimitiveTypeName(input_type),
-                        TilingParametersToString(tiling_enabled));
+                        primitive_util::LowercasePrimitiveTypeName(input_type));
   };
 };
 
 TEST_P(DotUnsetAlgorithmEmitterTest, UnsetAlgorithmIsEmittedCorrectly) {
-  auto [result_type, input_type, tiling_enabled] = GetParam();
+  auto [result_type, input_type] = GetParam();
   const std::string kHloText =
       GetDotAlgorithmHlo(input_type, result_type, PrecisionConfig::ALG_UNSET);
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloText));
@@ -2183,12 +2083,10 @@ TEST_P(DotUnsetAlgorithmEmitterTest, UnsetAlgorithmIsEmittedCorrectly) {
 INSTANTIATE_TEST_SUITE_P(
     DotUnsetAlgorithmEmitterTestSuite, DotUnsetAlgorithmEmitterTest,
     ::testing::Combine(::testing::ValuesIn(AllXlaDataTypes()),
-                       ::testing::ValuesIn(AllXlaDataTypes()),
-                       ::testing::Bool()),
+                       ::testing::ValuesIn(AllXlaDataTypes())),
     DotUnsetAlgorithmEmitterTest::ParamToString);
 
-TEST_P(TritonEmitterTestWithTilingParam,
-       ScaledDotIsSupportedByReferencePlatform) {
+TEST_F(TritonEmitterTest, ScaledDotIsSupportedByReferencePlatform) {
   constexpr absl::string_view kHloText = R"(
     HloModule ScaledDotIsSupportedByReferencePlatform
 
@@ -2206,7 +2104,7 @@ TEST_P(TritonEmitterTestWithTilingParam,
   EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
-TEST_P(TritonEmitterTestWithTilingParam, RocmWarpSizeIsSetCorrectly) {
+TEST_F(TritonEmitterTest, RocmWarpSizeIsSetCorrectly) {
   if (GpuComputeCapability().IsCuda()) {
     GTEST_SKIP() << "Warp size is always 32 on CUDA";
   }
@@ -2282,7 +2180,7 @@ TEST_P(TritonEmitterTestWithTilingParam, RocmWarpSizeIsSetCorrectly) {
   EXPECT_THAT(RunFileCheck(triton_passes_log, kPattern_n), true);
 }
 
-TEST_P(TritonEmitterTestWithTilingParam, RocmWavesPerEuAttributeIsSet) {
+TEST_F(TritonEmitterTest, RocmWavesPerEuAttributeIsSet) {
   if (GpuComputeCapability().IsCuda()) {
     GTEST_SKIP() << "waves_per_eu is ROCm-specific";
   }
@@ -2322,7 +2220,7 @@ TEST_P(TritonEmitterTestWithTilingParam, RocmWavesPerEuAttributeIsSet) {
   EXPECT_EQ(attr.getValueAsString().str(), "4, 4");
 }
 
-TEST_P(TritonEmitterTestWithTilingParam, RocmWavesPerEuZeroOmitsAttribute) {
+TEST_F(TritonEmitterTest, RocmWavesPerEuZeroOmitsAttribute) {
   if (GpuComputeCapability().IsCuda()) {
     GTEST_SKIP() << "waves_per_eu is ROCm-specific";
   }
@@ -2361,11 +2259,7 @@ TEST_P(TritonEmitterTestWithTilingParam, RocmWavesPerEuZeroOmitsAttribute) {
       << "waves_per_eu=0 should not set amdgpu-waves-per-eu attribute";
 }
 
-// Scans are only emitted with the experimental tiling propagation enabled.
 class ScanTritonEmitterTest : public TritonEmitterTest {
- public:
-  bool EnableTilingPropagation() const override { return true; }
-
  protected:
   // Returns arguments for the entry computation of `module` filled with small
   // integers. All intermediate scan results are then exactly representable in
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_deviceless_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_deviceless_test.cc
index 7244a8c7eec..298a65f209b 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_deviceless_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_deviceless_test.cc
@@ -63,14 +63,6 @@ class WarpSpecializationTritonEmitterTest : public TritonEmitterDevicelessTest {
 };
 
 class ExperimentalTilingTritonEmitterTest : public TritonEmitterDevicelessTest {
- public:
-  DebugOptions GetDebugOptionsForTest() const override {
-    DebugOptions debug_options =
-        TritonEmitterDevicelessTest::GetDebugOptionsForTest();
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(true);
-    return debug_options;
-  }
-
  protected:
   // Emits and compiles the Triton fusion that is the root of `hlo_text` and
   // returns the Triton IR before the conversion to the TritonGPU dialect.
@@ -724,8 +716,7 @@ ENTRY entry {
                           triple, data_layout, mlir_context));
 }
 
-TEST_F(TritonEmitterDevicelessTest,
-       AllGatherFusionUsesTiledHloComputationWhenTilingPropagationDisabled) {
+TEST_F(TritonEmitterDevicelessTest, AllGatherFusionUsesTiledHloComputation) {
   constexpr absl::string_view kHloText = R"(
 f {
   param0 = f32[128,128]{1,0} parameter(0)
@@ -754,10 +745,6 @@ ENTRY entry {
 )";
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> hlo_module,
                        ParseAndReturnVerifiedModule(kHloText));
-  // Explicitly ensure xla_gpu_experimental_enable_tiling_propagation is false.
-  hlo_module->mutable_config()
-      .mutable_debug_options()
-      .set_xla_gpu_experimental_enable_tiling_propagation(false);
 
   const auto* fusion = Cast<HloFusionInstruction>(
       hlo_module->entry_computation()->root_instruction());
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_int4_device_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_int4_device_test.cc
index d0f1e5866e5..e6e947746ef 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_int4_device_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_int4_device_test.cc
@@ -43,7 +43,6 @@ namespace {
 
 class TritonTestBase : public HloInterpreterReferenceMixin<HloPjRtGpuTestBase> {
  public:
-  virtual bool EnableTilingPropagation() const = 0;
   virtual bool EnableGemmFusionV2() const { return false; }
 
   DebugOptions GetDebugOptionsForTest() const override {
@@ -55,46 +54,21 @@ class TritonTestBase : public HloInterpreterReferenceMixin<HloPjRtGpuTestBase> {
     debug_options.set_xla_gpu_gemm_rewrite_size_threshold(0);
     debug_options
         .set_xla_gpu_experimental_enable_subchannel_dequantisation_fusion(true);
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
-        EnableTilingPropagation());
     debug_options.set_xla_gpu_experimental_gemm_fusion_v2(EnableGemmFusionV2());
     return debug_options;
   }
 };
 
-struct TritonTestParams {
-  bool enable_gemm_fusion_v2;
-  bool enable_tiling_propagation;
-};
-
 class TritonTest : public TritonTestBase,
-                   public ::testing::WithParamInterface<TritonTestParams> {
+                   public ::testing::WithParamInterface<bool> {
  public:
-  bool EnableGemmFusionV2() const override {
-    return GetParam().enable_gemm_fusion_v2;
-  }
-  bool EnableTilingPropagation() const override {
-    return GetParam().enable_tiling_propagation;
-  }
+  bool EnableGemmFusionV2() const override { return GetParam(); }
 };
 
-INSTANTIATE_TEST_SUITE_P(
-    TritonTest, TritonTest,
-    ::testing::Values(TritonTestParams{/*enable_gemm_fusion_v2=*/false,
-                                       /*enable_tiling_propagation=*/false},
-                      TritonTestParams{/*enable_gemm_fusion_v2=*/false,
-                                       /*enable_tiling_propagation=*/true},
-                      TritonTestParams{/*enable_gemm_fusion_v2=*/true,
-                                       /*enable_tiling_propagation=*/false},
-                      TritonTestParams{/*enable_gemm_fusion_v2=*/true,
-                                       /*enable_tiling_propagation=*/true}),
-    [](const ::testing::TestParamInfo<TritonTestParams>& info) {
-      return absl::StrCat(
-          info.param.enable_gemm_fusion_v2 ? "GemmFusionV2" : "GemmFusionV1",
-          "_",
-          info.param.enable_tiling_propagation ? "ExperimentalTiling"
-                                               : "SymbolicTiling");
-    });
+INSTANTIATE_TEST_SUITE_P(TritonTest, TritonTest, ::testing::Bool(),
+                         [](const ::testing::TestParamInfo<bool>& info) {
+                           return info.param ? "GemmFusionV2" : "GemmFusionV1";
+                         });
 
 // The following tests are for the channel and subchannel dequantization
 // fusions. We run the fused version to avoid the HLO passes and prove that
@@ -498,15 +472,10 @@ struct I4TestParams {
 
 class ParametrizedTritonTest
     : public TritonTestBase,
-      public ::testing::WithParamInterface<std::tuple<I4TestParams, bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<1>(GetParam());
-  }
-};
+      public ::testing::WithParamInterface<I4TestParams> {};
 
 TEST_P(ParametrizedTritonTest, Int4WeightsOnTheLhs) {
-  const I4TestParams& i4_params = std::get<0>(GetParam());
+  const I4TestParams& i4_params = GetParam();
   if (i4_params.HasBatchDim()) {
     GTEST_SKIP() << "2d test ignores batch dim case.";
   }
@@ -544,7 +513,7 @@ ENTRY entry_computation {
 }
 
 TEST_P(ParametrizedTritonTest, Int4WeightsOnTheLhsWithBatchDim) {
-  const I4TestParams& i4_params = std::get<0>(GetParam());
+  const I4TestParams& i4_params = GetParam();
   if (!i4_params.HasBatchDim()) {
     GTEST_SKIP() << "3d test ignores 2d case.";
   }
@@ -581,7 +550,7 @@ ENTRY entry_computation {
 }
 
 TEST_P(ParametrizedTritonTest, Int4WeightsOnTheRhs) {
-  const I4TestParams& i4_params = std::get<0>(GetParam());
+  const I4TestParams& i4_params = GetParam();
   if (i4_params.HasBatchDim()) {
     GTEST_SKIP() << "2d test ignores batch dim case.";
   }
@@ -657,14 +626,9 @@ std::vector<I4TestParams> Int4TestCases() {
   };
 }
 
-INSTANTIATE_TEST_SUITE_P(
-    ParametrizedTritonTest, ParametrizedTritonTest,
-    ::testing::Combine(::testing::ValuesIn(Int4TestCases()), ::testing::Bool()),
-    [](const ::testing::TestParamInfo<std::tuple<I4TestParams, bool>>& info) {
-      return absl::StrCat(
-          std::get<0>(info.param).name,
-          std::get<1>(info.param) ? "ExperimentalTiling" : "SymbolicTiling");
-    });
+INSTANTIATE_TEST_SUITE_P(ParametrizedTritonTest, ParametrizedTritonTest,
+                         ::testing::ValuesIn(Int4TestCases()),
+                         I4TestParams::ToString);
 
 TEST_P(TritonTest, NonstandardLayoutWithManyNonContractingDims) {
   constexpr absl::string_view kHloText = R"(
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc
index bc3900f1c4f..0361c93ba3a 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc
@@ -45,28 +45,7 @@ using ::xla::xtile::BlockLevelParameters;
 class XTileDialectTest : public HloHardwareIndependentTestBase,
                          public XTileTestBase {};
 
-class XTileDialectTestParameterized
-    : public XTileDialectTest,
-      public ::testing::WithParamInterface<bool> {
- protected:
-  DebugOptions GetDebugOptionsForTest() const override {
-    DebugOptions debug_options =
-        HloHardwareIndependentTestBase::GetDebugOptionsForTest();
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
-        GetParam());
-    return debug_options;
-  }
-};
-
-INSTANTIATE_TEST_SUITE_P(XTileDialectTestParameterized,
-                         XTileDialectTestParameterized, testing::Bool(),
-                         [](const ::testing::TestParamInfo<bool>& info) {
-                           return info.param ? "ExperimentalEmitter"
-                                             : "LegacyEmitter";
-                         });
-
-TEST_P(XTileDialectTestParameterized,
-       HloTransposeIsLoweredToStableHloTranspose) {
+TEST_F(XTileDialectTest, HloTransposeIsLoweredToStableHloTranspose) {
   constexpr absl::string_view kHloText = R"(
 HloModule t
 
@@ -95,7 +74,7 @@ CHECK: %[[RES:.*]] = stablehlo.transpose %[[ARG:.*]], dims = [1, 0] : (tensor<32
 )"));
 }
 
-TEST_P(XTileDialectTestParameterized, HloBitcastIsLoweredToTensorBitcast) {
+TEST_F(XTileDialectTest, HloBitcastIsLoweredToTensorBitcast) {
   constexpr absl::string_view kHloText = R"(
 HloModule t, is_scheduled=true
 
@@ -123,7 +102,7 @@ CHECK: %[[RES:.*]] = tensor.bitcast %[[ARG:.*]] : tensor<16x32xf32> to tensor<16
 )"));
 }
 
-TEST_P(XTileDialectTestParameterized, HloIotaIsLoweredToStableHloIota) {
+TEST_F(XTileDialectTest, HloIotaIsLoweredToStableHloIota) {
   constexpr absl::string_view kHloText = R"(
 HloModule t, is_scheduled=true
 
@@ -149,8 +128,7 @@ CHECK: %[[RES:.*]] = stablehlo.iota dim = 0 : tensor<16xi32>
 )"));
 }
 
-TEST_P(XTileDialectTestParameterized,
-       HloBroadcastInDimIsLoweredToStableHloBroadcastInDim) {
+TEST_F(XTileDialectTest, HloBroadcastInDimIsLoweredToStableHloBroadcastInDim) {
   constexpr absl::string_view kHloText = R"(
 HloModule t
 
@@ -179,7 +157,7 @@ CHECK: %[[RES:.*]] = stablehlo.broadcast_in_dim %[[ARG:.*]], dims = [0, 1] : (te
 )"));
 }
 
-TEST_P(XTileDialectTestParameterized,
+TEST_F(XTileDialectTest,
        HloZeroDimensionalBroadcastIsLoweredToStableHloBroadcastInDim) {
   constexpr absl::string_view kHloText = R"(
 HloModule t
@@ -209,7 +187,7 @@ CHECK: %[[RES:.*]] = stablehlo.broadcast_in_dim %[[ARG:.*]], dims = [] : (tensor
 )"));
 }
 
-TEST_P(XTileDialectTestParameterized, HloReduceIsLoweredToStableHloReduce) {
+TEST_F(XTileDialectTest, HloReduceIsLoweredToStableHloReduce) {
   constexpr absl::string_view kHloText = R"(
 HloModule t
 
@@ -246,11 +224,7 @@ CHECK: %[[RES:.*]] = stablehlo.reduce(%[[MASKED_INPUT]] init: %[[INIT]]) applies
 )"));
 }
 
-TEST_P(XTileDialectTestParameterized, HloScanIsLoweredToXTileScan) {
-  if (!GetParam()) {
-    GTEST_SKIP() << "Skipping test for legacy emitter.";
-  }
-
+TEST_F(XTileDialectTest, HloScanIsLoweredToXTileScan) {
   constexpr absl::string_view kHloText = R"(
 HloModule t
 
@@ -280,9 +254,7 @@ ENTRY e {
                        ParseAndReturnVerifiedModule(kHloText));
 
   BlockLevelParameters block_level_parameters;
-  block_level_parameters.output_tile_sizes =
-      GetParam() ? std::vector<std::vector<int64_t>>{{}}
-                 : std::vector<std::vector<int64_t>>{{1024}};
+  block_level_parameters.output_tile_sizes = {{}};
 
   EXPECT_OK(CreateXTileIrAndFileCheck(
       *module->GetComputationWithName("scan_fusion"), block_level_parameters,
@@ -301,7 +273,7 @@ ENTRY e {
 )"));
 }
 
-TEST_P(XTileDialectTestParameterized, HloReshapeIsLoweredToStableHloReshape) {
+TEST_F(XTileDialectTest, HloReshapeIsLoweredToStableHloReshape) {
   constexpr absl::string_view kHloText = R"(
 HloModule t, is_scheduled=true
 
@@ -329,7 +301,7 @@ CHECK: %[[RES:.*]] = stablehlo.reshape %[[ARG:.*]] : (tensor<16xi32>) -> tensor<
 )"));
 }
 
-TEST_P(XTileDialectTestParameterized, HloDotIsLoweredToStableHloDot) {
+TEST_F(XTileDialectTest, HloDotIsLoweredToStableHloDot) {
   constexpr absl::string_view kHloText = R"(
 HloModule t
 
@@ -362,7 +334,7 @@ CHECK: %[[ADD_RES:.*]] = arith.addf %[[ARG2:.*]], %[[RES]] : tensor<32x8xf32>
 )"));
 }
 
-TEST_P(XTileDialectTestParameterized, HloScaledDotIsLoweredToXTileDotScaled) {
+TEST_F(XTileDialectTest, HloScaledDotIsLoweredToXTileDotScaled) {
   constexpr absl::string_view kHloText = R"(
 HloModule m
 
@@ -416,8 +388,7 @@ ENTRY e {
       )"));
 }
 
-TEST_P(XTileDialectTestParameterized,
-       HloAllReduceIsLoweredToStableHloAllReduce) {
+TEST_F(XTileDialectTest, HloAllReduceIsLoweredToStableHloAllReduce) {
   constexpr absl::string_view kHloText =
       R"(
       HloModule wrapped_module_all-reduce-start
@@ -456,8 +427,7 @@ CHECK: stablehlo.add
 )"));
 }
 
-TEST_P(XTileDialectTestParameterized,
-       HloReduceScatterIsLoweredToStableHloReduceScatter) {
+TEST_F(XTileDialectTest, HloReduceScatterIsLoweredToStableHloReduceScatter) {
   constexpr absl::string_view kHloText =
       R"(
       HloModule wrapped_module_reduce-scatter
@@ -495,8 +465,7 @@ CHECK: stablehlo.add
 )"));
 }
 
-TEST_P(XTileDialectTestParameterized,
-       HloUnsignedIntIsLoweredToStableHloUnsignedInt) {
+TEST_F(XTileDialectTest, HloUnsignedIntIsLoweredToStableHloUnsignedInt) {
   constexpr absl::string_view kHloText = R"(
 HloModule t, is_scheduled=true
 
@@ -524,12 +493,7 @@ CHECK: stablehlo.add{{.*}}: tensor<16xui32>
 )"));
 }
 
-TEST_P(XTileDialectTestParameterized,
-       HloSameShapeMultiOutputFusionIsLoweredToXTileInsert) {
-  if (!GetParam()) {
-    GTEST_SKIP() << "Skipping test for legacy emitter.";
-  }
-
+TEST_F(XTileDialectTest, HloSameShapeMultiOutputFusionIsLoweredToXTileInsert) {
   constexpr absl::string_view kHloText = R"(
 HloModule t
 
@@ -608,9 +572,6 @@ TEST_F(XTileDialectTest, HloAllGatherDotLowering) {
 
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnUnverifiedModule(kHloText));
-  module->mutable_config()
-      .mutable_debug_options()
-      .set_xla_gpu_experimental_enable_tiling_propagation(true);
 
   BlockLevelParameters block_level_parameters;
   block_level_parameters.output_tile_sizes = {{128, 128}};
@@ -656,9 +617,6 @@ ENTRY e {
 })";
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnUnverifiedModule(kHloText));
-  module->mutable_config()
-      .mutable_debug_options()
-      .set_xla_gpu_experimental_enable_tiling_propagation(true);
 
   BlockLevelParameters block_level_parameters;
   block_level_parameters.output_tile_sizes = {{}};
@@ -703,9 +661,6 @@ ENTRY e {
 })";
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnUnverifiedModule(kHloText));
-  module->mutable_config()
-      .mutable_debug_options()
-      .set_xla_gpu_experimental_enable_tiling_propagation(true);
 
   BlockLevelParameters block_level_parameters;
   block_level_parameters.output_tile_sizes = {{}};
@@ -753,9 +708,6 @@ ENTRY e {
 })";
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
                        ParseAndReturnUnverifiedModule(kHloText));
-  module->mutable_config()
-      .mutable_debug_options()
-      .set_xla_gpu_experimental_enable_tiling_propagation(true);
   BlockLevelParameters block_level_parameters;
   block_level_parameters.output_tile_sizes = {{}};
   // The mask is emitted as a dynamic boundary check (scf.if + arith.select)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/bitcast_reduce_broadcast.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/bitcast_reduce_broadcast.hlo
index 91774091d9e..1c7db6a5b20 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/bitcast_reduce_broadcast.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/bitcast_reduce_broadcast.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p0 = bf16[128,4,256]{2,1,0} parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/diamond_broadcast_1d.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/diamond_broadcast_1d.hlo
index 332a5dc9e9a..e7119cb5143 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/diamond_broadcast_1d.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/diamond_broadcast_1d.hlo
@@ -1,5 +1,4 @@
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN: triton_test_correctness %s
 
 max_computation {
   x = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/diamond_broadcast_batch.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/diamond_broadcast_batch.hlo
index ffff2c2a6ba..4b50dde604f 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/diamond_broadcast_batch.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/diamond_broadcast_batch.hlo
@@ -1,5 +1,4 @@
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN: triton_test_correctness %s
 
 max_computation {
   x = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/diamond_broadcast_reduce.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/diamond_broadcast_reduce.hlo
index 7f5ee9d5167..f429653f7ca 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/diamond_broadcast_reduce.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/diamond_broadcast_reduce.hlo
@@ -1,5 +1,4 @@
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN: triton_test_correctness %s
 
 max_computation {
   x = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/diamond_splat_scalar.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/diamond_splat_scalar.hlo
index af325713414..34e88fa06da 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/diamond_splat_scalar.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/diamond_splat_scalar.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s --check-prefixes=CHECK,CHECK-OLD
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s --check-prefixes=CHECK,CHECK-NEW
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 max_computation {
   x = f32[] parameter(0)
@@ -19,10 +17,8 @@ fusion {
   broadcast_0 = f32[64,32,16]{2,1,0} broadcast(reduce_0), dimensions={0,1}
   ROOT _ = f32[64,32,16]{2,1,0} add(add_0, broadcast_0)
 }
-// CHECK-OLD:   #xla.indexing_map<"(pid_0) -> (pid_0 / 32), domain: pid_0 in [0, 2047]">
-// CHECK-OLD:   #xla.indexing_map<"(pid_0) -> (pid_0 mod 32), domain: pid_0 in [0, 2047]">
-// CHECK-NEW:   #xla.indexing_map<"(pid) -> (pid / 32), domain: pid in [0, 2047]">
-// CHECK-NEW:   #xla.indexing_map<"(pid) -> (pid mod 32), domain: pid in [0, 2047]">
+// CHECK:   #xla.indexing_map<"(pid) -> (pid / 32), domain: pid in [0, 2047]">
+// CHECK:   #xla.indexing_map<"(pid) -> (pid mod 32), domain: pid in [0, 2047]">
 // CHECK-LABEL: xtile.entry_func @triton_fn(
 // CHECK-SAME:      %[[P0:[A-Za-z0-9_]*]]: memref<64x32x16xf32>
 // CHECK-SAME:      %[[P1:[A-Za-z0-9_]*]]: memref<f32>
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/generic_emitter_multi_tiled_dims.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/generic_emitter_multi_tiled_dims.hlo
index 19317f95bdc..9e1e3874658 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/generic_emitter_multi_tiled_dims.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/normalization/generic_emitter_multi_tiled_dims.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s --check-prefixes=CHECK,CHECK-OLD
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s --check-prefixes=CHECK,CHECK-NEW
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 max {
   Arg_0 = f32[] parameter(0)
@@ -22,10 +20,8 @@ fusion {
   broadcast_4 = f32[10,125,127]{2,1,0} broadcast(reduce_0), dimensions={0,1}
   ROOT multiply = f32[10,125,127]{2,1,0} multiply(multiply_1, broadcast_4)
 }
-// CHECK-OLD:    #[[MAP:.*]] = #xla.indexing_map<"(pid_0) -> (pid_0 / 125), domain: pid_0 in [0, 1249]">
-// CHECK-OLD:    #[[MAP1:.*]] = #xla.indexing_map<"(pid_0) -> (pid_0 mod 125), domain: pid_0 in [0, 1249]">
-// CHECK-NEW:    #[[MAP:.*]] = #xla.indexing_map<"(pid) -> (pid / 125), domain: pid in [0, 1249]">
-// CHECK-NEW:    #[[MAP1:.*]] = #xla.indexing_map<"(pid) -> (pid mod 125), domain: pid in [0, 1249]">
+// CHECK:    #[[MAP:.*]] = #xla.indexing_map<"(pid) -> (pid / 125), domain: pid in [0, 1249]">
+// CHECK:    #[[MAP1:.*]] = #xla.indexing_map<"(pid) -> (pid mod 125), domain: pid in [0, 1249]">
 // CHECK:        xtile.entry_func @triton_fn
 // CHECK-SAME:      (%[[P0:.*]]: {{.*}}, %[[P1:.*]]: {{.*}}, %[[P2:.*]]: {{.*}}, %[[P3:.*]]: {{.*}}, %[[TID:.*]]: index{{( \{xla.range = \[0 : index, 1249 : index\]\})?}})
 // CHECK-DAG:        %[[C_0:.*]] = arith.constant 0 : index
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/pad/high_pad_1d.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/pad/high_pad_1d.hlo
index 86070105545..1629a40f725 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/pad/high_pad_1d.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/pad/high_pad_1d.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   param_0 = f32[17]{0} parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/pad/high_pad_2d.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/pad/high_pad_2d.hlo
index 0ee816b0846..aa5e6196460 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/pad/high_pad_2d.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/pad/high_pad_2d.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 triton_computation {
   param_0 = f32[17,137]{1,0} parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/pad/high_pad_of_broadcast_and_reshape.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/pad/high_pad_of_broadcast_and_reshape.hlo
index 173f8f68f0f..fa2e2b43f15 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/pad/high_pad_of_broadcast_and_reshape.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/pad/high_pad_of_broadcast_and_reshape.hlo
@@ -1,5 +1,5 @@
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 // Reproducer for reproducing test for of pad OOB access (b/568080363) and
 // iterating on fully padded tiles (b/477265261).
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/pad/high_pad_tile_size_equals_input_size.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/pad/high_pad_tile_size_equals_input_size.hlo
index 2927f9aee6b..cb01d67ca66 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/pad/high_pad_tile_size_equals_input_size.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/pad/high_pad_tile_size_equals_input_size.hlo
@@ -1,5 +1,5 @@
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 // Dimension 0 is padded and its tile size is equal to the size of the pad
 // input, so all tiles except for the first one are entirely padding and must be
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/intra_warp_reduce_of_reduce.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/intra_warp_reduce_of_reduce.hlo
index ecd5a384521..d86d8d88c97 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/intra_warp_reduce_of_reduce.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/intra_warp_reduce_of_reduce.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 // TODO: b/477615292 - enable experimental tiling when reshape [2,4] -> [8] is supported.
 
 // Reproducer from b/380277401.
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/multi_0d_broadcast.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/multi_0d_broadcast.hlo
index f57d6550d59..383bef4b0a5 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/multi_0d_broadcast.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/multi_0d_broadcast.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s --abs_error_bound=1e-1 --rel_error_bound=1e-1
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s --abs_error_bound=1e-1 --rel_error_bound=1e-1
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s --abs_error_bound=1e-1 --rel_error_bound=1e-1
 // TODO: b/502907772 - enable experimental tiling after fixing crash.
 
 add {
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_complex.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_complex.hlo
index 473930be63a..176ea1c4f3e 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_complex.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_complex.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 unusual {
   p0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_concat_diamond.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_concat_diamond.hlo
index 5046db172c4..55f9a7f9463 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_concat_diamond.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_concat_diamond.hlo
@@ -1,5 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-4 --rel_error_bound=1e-4 %s
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness --abs_error_bound=1e-4 --rel_error_bound=1e-4 %s
 
 add_region {
   p0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_diamond_tiled.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_diamond_tiled.hlo
index 65c0cf2852e..4366dab13d6 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_diamond_tiled.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_diamond_tiled.hlo
@@ -1,5 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-4 --rel_error_bound=1e-4 %s
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness --abs_error_bound=1e-4 --rel_error_bound=1e-4 %s
 
 max_region {
   p0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_intermediate.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_intermediate.hlo
index 9bf1cefb2c2..67206dc08a6 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_intermediate.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_intermediate.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 maximum {
   p0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_major_subtile.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_major_subtile.hlo
index 8c3062bc1a4..f5f4ea34cea 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_major_subtile.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_major_subtile.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 maximum {
   p0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_minor.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_minor.hlo
index aa80bdc8ef1..240624eb083 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_minor.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_minor.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 region {
   p0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_minor_extra_out.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_minor_extra_out.hlo
index e80f6aca385..1f419226222 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_minor_extra_out.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_minor_extra_out.hlo
@@ -1,8 +1,8 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false --xla_gpu_unsupported_enable_triton_multi_output_fusion %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
+// RUN: fusion_to_triton --xla_gpu_unsupported_enable_triton_multi_output_fusion %s 2>&1 | FileCheck %s --check-prefixes=CHECK-ERROR
+// CHECK-ERROR: Only single-result fusions are supported for now. Received 2 roots.
 // TODO: b/502910372 - support multi-output fusion.
-// RUN-NOT: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation --xla_gpu_unsupported_enable_triton_multi_output_fusion %s | FileCheck %s
-// RUN-NOT: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
+// RUN-NOT: fusion_to_triton --xla_gpu_unsupported_enable_triton_multi_output_fusion %s | FileCheck %s
+// RUN-NOT: triton_test_correctness --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
 
 region {
   p0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_minor_subtile.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_minor_subtile.hlo
index bf25f30d9c1..9d5d2af7656 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_minor_subtile.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_minor_subtile.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s --check-prefixes=CHECK,CHECK-NEW
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 region {
   p0 = f32[] parameter(0)
@@ -17,7 +15,7 @@ fusion {
 }
 
 // Make sure that the input tile is masked with a neutral value.
-// CHECK-NEW: #indexing_map = #xla.indexing_map<"(pid) -> (pid * 4), domain: pid in [0, 3]">
+// CHECK: #indexing_map = #xla.indexing_map<"(pid) -> (pid * 4), domain: pid in [0, 3]">
 // CHECK:  %[[LOAD:.*]] = xtile.extract {{.*}}[4, 8]{{.*}} : memref<13x7xf32> -> tensor<4x8xf32>
 // CHECK:  %[[MASKED:.*]] = xtile.mask %[[LOAD]] bounds [4, 7]
 // CHECK:  "tt.reduce"(%[[MASKED]]) <{axis = 1 : i32}>
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_multi_dim.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_multi_dim.hlo
index 2390bf88ad4..0baeaaa52e7 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_multi_dim.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_multi_dim.hlo
@@ -1,5 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-4 --rel_error_bound=1e-4 %s
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness --abs_error_bound=1e-4 --rel_error_bound=1e-4 %s
 
 region {
   p0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_multi_dim_tiled_and_untiled.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_multi_dim_tiled_and_untiled.hlo
index fc20b5425d4..092daf9ad96 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_multi_dim_tiled_and_untiled.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_multi_dim_tiled_and_untiled.hlo
@@ -1,5 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-4 --rel_error_bound=1e-4 %s
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness --abs_error_bound=1e-4 --rel_error_bound=1e-4 %s
 
 region {
   p0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_non_neutral_init.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_non_neutral_init.hlo
index 9d0a78df8bb..e6fe6422ebf 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_non_neutral_init.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_non_neutral_init.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s --check-prefix=CHECK-UNTILED
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s --check-prefixes=CHECK-TILED
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 region {
   p0 = f32[] parameter(0)
@@ -16,15 +14,10 @@ fusion {
     backend_config={sizes:[4]}
 }
 
-// CHECK-UNTILED: %[[INIT:.*]] = arith.constant dense<5.000000e+00> : tensor<4xf32>
-// CHECK-UNTILED: %[[REDUCE_INPUT:.*]] = xtile.extract %{{.*}} :
-// CHECK-UNTILED: %[[REDUCE:.*]] = "tt.reduce"(%[[REDUCE_INPUT]]) <{axis = 1 : i32}>
-// CHECK-UNTILED: %[[COMBINED:.*]] = arith.maximumf %[[REDUCE]], %[[INIT]] : tensor<4xf32>
-
-// CHECK-TILED-DAG: %[[INIT_VAL:.*]] = arith.constant dense<5.000000e+00> : tensor<4xf32>
-// CHECK-TILED: %[[FOR:.*]] = scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}}
-// CHECK-TILED: %[[REDUCE:.*]] = "tt.reduce"(%[[FOR]]) <{axis = 1 : i32}>
-// CHECK-TILED: %[[COMBINED:.*]] = arith.maximumf %[[REDUCE]], %[[INIT_VAL]] : tensor<4xf32>
+// CHECK-DAG: %[[INIT_VAL:.*]] = arith.constant dense<5.000000e+00> : tensor<4xf32>
+// CHECK: %[[FOR:.*]] = scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}}
+// CHECK: %[[REDUCE:.*]] = "tt.reduce"(%[[FOR]]) <{axis = 1 : i32}>
+// CHECK: %[[COMBINED:.*]] = arith.maximumf %[[REDUCE]], %[[INIT_VAL]] : tensor<4xf32>
 
 ENTRY main {
   p = f32[8,16] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_of_concat.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_of_concat.hlo
index a1c997f1aa9..c3a3bb218f2 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_of_concat.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_of_concat.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s --abs_error_bound=1e-6 --rel_error_bound=1e-6
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-6 --rel_error_bound=1e-6
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s --abs_error_bound=1e-6 --rel_error_bound=1e-6
 
 HloModule m
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_tiled.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_tiled.hlo
index a6d822f85f3..627f86cf2b3 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_tiled.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_tiled.hlo
@@ -1,5 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-7 --rel_error_bound=1e-7 %s
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness --abs_error_bound=1e-7 --rel_error_bound=1e-7 %s
 
 region {
   p0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_tiled_masked.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_tiled_masked.hlo
index 363554ba940..c52e741bd3b 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_tiled_masked.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_tiled_masked.hlo
@@ -1,5 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 region {
   p0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_to_scalar_extra_out.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_to_scalar_extra_out.hlo
index 60a1c632ba5..45bacf8c748 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_to_scalar_extra_out.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_to_scalar_extra_out.hlo
@@ -1,8 +1,8 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false --xla_gpu_unsupported_enable_triton_multi_output_fusion %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
+// RUN: fusion_to_triton --xla_gpu_unsupported_enable_triton_multi_output_fusion %s 2>&1 | FileCheck %s --check-prefixes=CHECK-ERROR
+// CHECK-ERROR: Only single-result fusions are supported for now. Received 2 roots.
 // TODO: b/502910372 - enable when multi-output is supported.
-// RUN-NOT: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation --xla_gpu_unsupported_enable_triton_multi_output_fusion %s | FileCheck %s
-// RUN-NOT: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
+// RUN-NOT: fusion_to_triton --xla_gpu_unsupported_enable_triton_multi_output_fusion %s | FileCheck %s
+// RUN-NOT: triton_test_correctness --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
 
 region {
   p0 = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reshape/reshape_expand_high_rank.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reshape/reshape_expand_high_rank.hlo
index 17f6459afe4..d83510d0964 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reshape/reshape_expand_high_rank.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reshape/reshape_expand_high_rank.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation=false | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation=true | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=true %s
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p = f32[512,16] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reshape/reshape_to_0d.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reshape/reshape_to_0d.hlo
index 3aa47b250eb..a16bc87cab2 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reshape/reshape_to_0d.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reshape/reshape_to_0d.hlo
@@ -1,7 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   p0 = f32[1,1,1,1] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot.hlo
index 911ac19c0fc..9a2cd2ef5a3 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s --abs_error_bound=1e-3 --rel_error_bound=1e-3
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-3 --rel_error_bound=1e-3
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s --abs_error_bound=1e-3 --rel_error_bound=1e-3
 
 triton_dot {
   lhs = f8e4m3fn[128,128] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization.hlo
index e6f41977e2f..5619eee8d23 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization.hlo
@@ -1,7 +1,6 @@
-// RUN: hlo-opt --xla_gpu_experimental_enable_tiling_propagation=false %s --platform=gpu --xla_gpu_target_config_filename=%S/../../../../target_config/specs/h100_sxm.txtpb > %t.optimized.hlo
+// RUN: hlo-opt %s --platform=gpu --xla_gpu_target_config_filename=%S/../../../../target_config/specs/h100_sxm.txtpb > %t.optimized.hlo
 // RUN: FileCheck %s < %t.optimized.hlo
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %t.optimized.hlo --abs_error_bound=1e-3 --rel_error_bound=1e-3
-// RUN: triton_test_correctness %t.optimized.hlo --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-3 --rel_error_bound=1e-3
+// RUN: triton_test_correctness %t.optimized.hlo --abs_error_bound=1e-3 --rel_error_bound=1e-3
 
 HloModule jit_scaled_dot_fp8
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization_bitcast.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization_bitcast.hlo
index edf8a80ae26..20fd0b6c248 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization_bitcast.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization_bitcast.hlo
@@ -1,7 +1,6 @@
-// RUN: hlo-opt --xla_gpu_experimental_enable_tiling_propagation=false %s --platform=gpu --xla_gpu_target_config_filename=%S/../../../../target_config/specs/h100_sxm.txtpb > %t.optimized.hlo
+// RUN: hlo-opt %s --platform=gpu --xla_gpu_target_config_filename=%S/../../../../target_config/specs/h100_sxm.txtpb > %t.optimized.hlo
 // RUN: FileCheck %s < %t.optimized.hlo
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %t.optimized.hlo --abs_error_bound=1e-3 --rel_error_bound=1e-3
-// RUN: triton_test_correctness %t.optimized.hlo --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-3 --rel_error_bound=1e-3
+// RUN: triton_test_correctness %t.optimized.hlo --abs_error_bound=1e-3 --rel_error_bound=1e-3
 
 HloModule jit_scaled_dot_fp8_bitcast
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc
index 66d8f815ba7..abf8ca66f8f 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc
@@ -87,23 +87,16 @@ using ::testing::HasSubstr;
 using ::xla::xtile::BlockLevelFusionConfig;
 using ::xla::xtile::BlockLevelParameters;
 
-std::string TilingParametersToString(bool tiling_propagation_enabled) {
-  return tiling_propagation_enabled ? "SymbolicTiling" : "ExperimentalTiling";
-}
-
 class TritonEmitterTest
     : public HloInterpreterReferenceMixin<GpuPjRtCodegenTest>,
       public XTileTestBase {
  public:
-  virtual bool EnableTilingPropagation() const = 0;
   DebugOptions GetDebugOptionsForTest() const override {
     DebugOptions debug_options = HloInterpreterReferenceMixin<
         GpuPjRtCodegenTest>::GetDebugOptionsForTest();
     debug_options.set_xla_gpu_unsupported_enable_triton_multi_output_fusion(
         true);
     debug_options.set_xla_gpu_experimental_disable_binary_libraries(true);
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
-        EnableTilingPropagation());
     return debug_options;
   }
 
@@ -141,15 +134,7 @@ class TritonEmitterTest
   }
 };
 
-class TritonEmitterTestWithTilingParam
-    : public TritonEmitterTest,
-      public ::testing::WithParamInterface<bool> {
- public:
-  bool EnableTilingPropagation() const override { return GetParam(); }
-};
-
-TEST_P(TritonEmitterTestWithTilingParam,
-       ScaledDotIsSupportedByReferencePlatform) {
+TEST_F(TritonEmitterTest, ScaledDotIsSupportedByReferencePlatform) {
   constexpr absl::string_view kHloText = R"(
     HloModule ScaledDotIsSupportedByReferencePlatform
 
@@ -167,12 +152,6 @@ TEST_P(TritonEmitterTestWithTilingParam,
   EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
-INSTANTIATE_TEST_SUITE_P(TritonEmitterTestWithTilingParamTestSuite,
-                         TritonEmitterTestWithTilingParam, ::testing::Bool(),
-                         [](const ::testing::TestParamInfo<bool>& info) {
-                           return TilingParametersToString(info.param);
-                         });
-
 struct ScaleDotTestParams {
   std::string lhs_type;
   std::string rhs_type;
@@ -208,14 +187,9 @@ std::ostream& operator<<(std::ostream& stream, const ScaleDotTestParams& tc) {
                 << ",\n\toutput_type:" << tc.output_type << "\n}";
 }
 
-class TritonScaledDotGemmTest : public TritonEmitterTest,
-                                public ::testing::WithParamInterface<
-                                    std::tuple<ScaleDotTestParams, bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<1>(GetParam());
-  }
-
+class TritonScaledDotGemmTest
+    : public TritonEmitterTest,
+      public ::testing::WithParamInterface<ScaleDotTestParams> {
  public:
   DebugOptions GetDebugOptionsForTest() const override {
     DebugOptions debug_options = TritonEmitterTest::GetDebugOptionsForTest();
@@ -228,7 +202,7 @@ class TritonScaledDotGemmTest : public TritonEmitterTest,
 
 TEST_P(TritonScaledDotGemmTest,
        FP8ScaledDotCompilesToPtxIntrinsicsWhenAvailable) {
-  const ScaleDotTestParams& params = std::get<0>(GetParam());
+  const ScaleDotTestParams& params = GetParam();
   constexpr absl::string_view kHloTextTemplate = R"hlo(
 HloModule m
 
@@ -289,7 +263,7 @@ ENTRY e {
 }
 
 TEST_P(TritonScaledDotGemmTest, FP8ScaledDotGetsFusedAndExecutesCorrectly) {
-  const ScaleDotTestParams& params = std::get<0>(GetParam());
+  const ScaleDotTestParams& params = GetParam();
   if (auto cc = GpuComputeCapability().cuda_compute_capability();
       cc && !cc->IsAtLeastBlackwell()) {
     GTEST_SKIP() << "Skipping test for pre-Blackwell GPUs.";
@@ -324,22 +298,14 @@ ENTRY e {
 
 INSTANTIATE_TEST_SUITE_P(
     TritonScaledDotGemmTest, TritonScaledDotGemmTest,
-    ::testing::Combine(
-        ::testing::Values(
-            ScaleDotTestParams{"f8e4m3fn[128,128]", "f8e4m3fn[128,256]",
-                               "f8e8m0fnu[128,4]", "f8e8m0fnu[4,256]",
-                               "bf16[128,256]", "f8E4M3FN"},
-            ScaleDotTestParams{"f8e5m2[128,128]", "f8e5m2[128,256]",
-                               "f8e8m0fnu[128,4]", "f8e8m0fnu[4,256]",
-                               "bf16[128,256]", "f8E5M2"}),
-        ::testing::Bool()),
-    [](const ::testing::TestParamInfo<std::tuple<ScaleDotTestParams, bool>>&
-           info) {
-      return absl::StrCat(ScaleDotTestParams::ToString(
-                              ::testing::TestParamInfo<ScaleDotTestParams>(
-                                  std::get<0>(info.param), info.index)),
-                          TilingParametersToString(std::get<1>(info.param)));
-    });
+    ::testing::Values(ScaleDotTestParams{"f8e4m3fn[128,128]",
+                                         "f8e4m3fn[128,256]",
+                                         "f8e8m0fnu[128,4]", "f8e8m0fnu[4,256]",
+                                         "bf16[128,256]", "f8E4M3FN"},
+                      ScaleDotTestParams{"f8e5m2[128,128]", "f8e5m2[128,256]",
+                                         "f8e8m0fnu[128,4]", "f8e8m0fnu[4,256]",
+                                         "bf16[128,256]", "f8E5M2"}),
+    ScaleDotTestParams::ToString);
 
 class TritonScaledDotTestBase : public TritonEmitterTest {
  public:
@@ -407,11 +373,7 @@ class TritonScaledDotTestBase : public TritonEmitterTest {
   }
 };
 
-class TritonScaledDotTest : public TritonScaledDotTestBase,
-                            public ::testing::WithParamInterface<bool> {
- public:
-  bool EnableTilingPropagation() const override { return GetParam(); }
-};
+using TritonScaledDotTest = TritonScaledDotTestBase;
 
 struct Fp4ScaledDotTypeCase {
   PrimitiveType lhs_type;
@@ -422,29 +384,22 @@ struct Fp4ScaledDotTypeCase {
 
 using Fp4ScaledDotTestParam = std::tuple<Fp4ScaledDotTypeCase,
                                          /*lhs_k_minor=*/bool,
-                                         /*rhs_k_minor=*/bool,
-                                         /*tiling_enabled=*/bool>;
+                                         /*rhs_k_minor=*/bool>;
 
 std::string Fp4ScaledDotTestParamToString(
     const ::testing::TestParamInfo<Fp4ScaledDotTestParam>& info) {
-  const auto& [type_case, lhs_k_minor, rhs_k_minor, tiling_enabled] =
-      info.param;
+  const auto& [type_case, lhs_k_minor, rhs_k_minor] = info.param;
   return absl::StrCat(PrimitiveType_Name(type_case.lhs_type), "_",
                       PrimitiveType_Name(type_case.rhs_type), "_",
                       PrimitiveType_Name(type_case.scale_type), "_Block",
                       type_case.block_size, "_Lhs", lhs_k_minor, "_Rhs",
-                      rhs_k_minor, "_",
-                      TilingParametersToString(tiling_enabled));
+                      rhs_k_minor);
 }
 
 class TritonFp4ScaledDotTest
     : public TritonScaledDotTestBase,
       public ::testing::WithParamInterface<Fp4ScaledDotTestParam> {
  public:
-  bool EnableTilingPropagation() const override {
-    return std::get<3>(GetParam());
-  }
-
   void RunFp4ScaledDotExecutionTest(PrimitiveType lhs_type,
                                     PrimitiveType rhs_type,
                                     PrimitiveType scale_type, int block_size,
@@ -539,12 +494,6 @@ ENTRY e {
   }
 };
 
-INSTANTIATE_TEST_SUITE_P(TritonScaledDotTestSuite, TritonScaledDotTest,
-                         ::testing::Bool(),
-                         [](const ::testing::TestParamInfo<bool>& info) {
-                           return TilingParametersToString(info.param);
-                         });
-
 TEST_P(TritonFp4ScaledDotTest, Executes) {
   auto cc = GpuComputeCapability().cuda_compute_capability();
   if (!cc || !cc->IsAtLeastHopper()) {
@@ -574,10 +523,10 @@ INSTANTIATE_TEST_SUITE_P(
                                  32},  // MXE5M2 x MXFP4
             Fp4ScaledDotTypeCase{F4E2M1FN, F4E2M1FN, F8E4M3FN,
                                  16}),  // NVFP4 x NVFP4 (block 16)
-        ::testing::Bool(), ::testing::Bool(), ::testing::Bool()),
+        ::testing::Bool(), ::testing::Bool()),
     Fp4ScaledDotTestParamToString);
 
-TEST_P(TritonScaledDotTest,
+TEST_F(TritonScaledDotTest,
        ScaledDotWithOmmittedLhsScaleGetFusedAndExecutedCorrectly) {
   if (auto cc = GpuComputeCapability().cuda_compute_capability();
       cc && !cc->IsAtLeastHopper()) {
@@ -633,7 +582,7 @@ ENTRY e {
       std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
-TEST_P(TritonScaledDotTest, FP8ScaledDotLhsKNotMinorDim) {
+TEST_F(TritonScaledDotTest, FP8ScaledDotLhsKNotMinorDim) {
   if (!GetCudaComputeCapability().IsAtLeastBlackwell()) {
     GTEST_SKIP() << "FP8 scaled dot requires Blackwell+";
   }
@@ -675,7 +624,7 @@ ENTRY e {
       std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
-TEST_P(TritonScaledDotTest, ScaledDotWithBatchGetFusedAndExecutedCorrectly) {
+TEST_F(TritonScaledDotTest, ScaledDotWithBatchGetFusedAndExecutedCorrectly) {
   if (auto cc = GpuComputeCapability().cuda_compute_capability();
       cc && !cc->IsAtLeastHopper()) {
     GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs.";
@@ -723,7 +672,7 @@ ENTRY e {
       std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
-TEST_P(TritonScaledDotTest, BroadcastAndReshapeGetFused) {
+TEST_F(TritonScaledDotTest, BroadcastAndReshapeGetFused) {
   if (auto cc = GpuComputeCapability().cuda_compute_capability();
       cc && !cc->IsAtLeastHopper()) {
     GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs.";
@@ -783,7 +732,7 @@ ENTRY e {
       std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
-TEST_P(TritonScaledDotTest, Mxfp8ScaledDotSmallBlockKAndNExecutes) {
+TEST_F(TritonScaledDotTest, Mxfp8ScaledDotSmallBlockKAndNExecutes) {
   if (!GetCudaComputeCapability().IsAtLeastHopper()) {
     GTEST_SKIP() << "Requires Hopper+.";
   }
@@ -821,7 +770,7 @@ ENTRY e {
       std::move(module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
-TEST_P(TritonScaledDotTest, Fp4Succeeds) {
+TEST_F(TritonScaledDotTest, Fp4Succeeds) {
   if (!GetCudaComputeCapability().IsAtLeastBlackwell()) {
     GTEST_SKIP() << "Scaled dot with FP4 isn't supported by Triton for "
                     "pre-Blackwell GPUs.";
@@ -874,7 +823,7 @@ TEST_P(TritonScaledDotTest, Fp4Succeeds) {
       std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
-TEST_P(TritonScaledDotTest, Mxfp4KPackedRhsBlackwellLowersWithTranspose) {
+TEST_F(TritonScaledDotTest, Mxfp4KPackedRhsBlackwellLowersWithTranspose) {
   if (!GetCudaComputeCapability().IsAtLeastBlackwell()) {
     GTEST_SKIP() << "Requires Blackwell+.";
   }
@@ -918,7 +867,7 @@ ENTRY e {
       std::move(module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
-TEST_P(TritonScaledDotTest, GlobalScalerSucceeds) {
+TEST_F(TritonScaledDotTest, GlobalScalerSucceeds) {
   if (!GetCudaComputeCapability().IsAtLeastHopper()) {
     GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs.";
   }
@@ -971,7 +920,7 @@ ENTRY e {
       std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
-TEST_P(TritonScaledDotTest, ScaledDotWithE8m0Scale) {
+TEST_F(TritonScaledDotTest, ScaledDotWithE8m0Scale) {
   if (!GetCudaComputeCapability().IsAtLeastHopper()) {
     GTEST_SKIP()
         << "ScaledDot with Triton requires Hopper or newer architecture.";
@@ -1018,7 +967,7 @@ ENTRY e {
       std::move(optimized_module), ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
-TEST_P(TritonScaledDotTest, ScaledDotWithE4m3Scale) {
+TEST_F(TritonScaledDotTest, ScaledDotWithE4m3Scale) {
   if (!GetCudaComputeCapability().IsAtLeastHopper()) {
     GTEST_SKIP()
         << "ScaledDot with Triton requires Hopper or newer architecture.";
@@ -1058,7 +1007,7 @@ ENTRY e {
   }
 }
 
-TEST_P(TritonScaledDotTest, ScaledDotWithE5m2Scale) {
+TEST_F(TritonScaledDotTest, ScaledDotWithE5m2Scale) {
   if (!GetCudaComputeCapability().IsAtLeastHopper()) {
     GTEST_SKIP()
         << "ScaledDot with Triton requires Hopper or newer architecture.";
@@ -1105,17 +1054,11 @@ struct ScaledDotCoverageTestCase {
   int block_size;
 };
 
-class TritonScaledDotCoverageTest : public TritonScaledDotTestBase,
-                                    public ::testing::WithParamInterface<
-                                        std::tuple<ScaledDotCoverageTestCase,
-                                                   /*lhs_k_minor=*/bool,
-                                                   /*rhs_k_minor=*/bool,
-                                                   /*tiling_enabled=*/bool>> {
- public:
-  bool EnableTilingPropagation() const override {
-    return std::get<3>(GetParam());
-  }
-};
+class TritonScaledDotCoverageTest
+    : public TritonScaledDotTestBase,
+      public ::testing::WithParamInterface<std::tuple<ScaledDotCoverageTestCase,
+                                                      /*lhs_k_minor=*/bool,
+                                                      /*rhs_k_minor=*/bool>> {};
 
 std::vector<ScaledDotCoverageTestCase> GetCoverageTestCases() {
   std::vector<PrimitiveType> input_types = {F8E4M3FN, F8E5M2, F4E2M1FN};
@@ -1137,19 +1080,17 @@ std::vector<ScaledDotCoverageTestCase> GetCoverageTestCases() {
 
 std::string ScaledDotCoverageTestParamToString(
     const ::testing::TestParamInfo<
-        std::tuple<ScaledDotCoverageTestCase, bool, bool, bool>>& info) {
-  const auto& [type_case, lhs_k_minor, rhs_k_minor, tiling_enabled] =
-      info.param;
+        std::tuple<ScaledDotCoverageTestCase, bool, bool>>& info) {
+  const auto& [type_case, lhs_k_minor, rhs_k_minor] = info.param;
   return absl::StrCat(PrimitiveType_Name(type_case.lhs_type), "_",
                       PrimitiveType_Name(type_case.rhs_type), "_",
                       PrimitiveType_Name(type_case.scale_type), "_Block",
                       type_case.block_size, "_Lhs", lhs_k_minor, "_Rhs",
-                      rhs_k_minor, "_",
-                      TilingParametersToString(tiling_enabled));
+                      rhs_k_minor);
 }
 
 TEST_P(TritonScaledDotCoverageTest, Executes) {
-  const auto& [param, lhs_k_minor, rhs_k_minor, tiling_enabled] = GetParam();
+  const auto& [param, lhs_k_minor, rhs_k_minor] = GetParam();
 
   auto cc = GpuComputeCapability().cuda_compute_capability();
   std::string device_name = "Unknown";
@@ -1169,14 +1110,12 @@ TEST_P(TritonScaledDotCoverageTest, Executes) {
       lhs_k_minor ? lhs_name : absl::StrCat(lhs_name, ".T");
   std::string rhs_display =
       rhs_k_minor ? absl::StrCat(rhs_name, ".T") : rhs_name;
-  std::string tiling_name = tiling_enabled ? "Symbolic" : "Experimental";
 
   LOG(ERROR) << "Report Device: " << device_name;
   LOG(ERROR) << "Report LHS: " << lhs_display;
   LOG(ERROR) << "Report RHS: " << rhs_display;
   LOG(ERROR) << "Report Scale: " << scale_name;
   LOG(ERROR) << "Report BlockSize: " << param.block_size;
-  LOG(ERROR) << "Report Tiling: " << tiling_name;
 
   if (!cc || !cc->IsAtLeastHopper()) {
     GTEST_SKIP() << "Scaled dot isn't supported by Triton for pre-Hopper GPUs.";
@@ -1405,7 +1344,7 @@ ENTRY e {
 INSTANTIATE_TEST_SUITE_P(
     TritonScaledDotCoverageTestSuite, TritonScaledDotCoverageTest,
     ::testing::Combine(::testing::ValuesIn(GetCoverageTestCases()),
-                       ::testing::Bool(), ::testing::Bool(), ::testing::Bool()),
+                       ::testing::Bool(), ::testing::Bool()),
     ScaledDotCoverageTestParamToString);
 
 }  // namespace
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/boundary_dynamic_slice.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/boundary_dynamic_slice.hlo
index 9ad87bf9416..abc3b016a24 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/boundary_dynamic_slice.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/boundary_dynamic_slice.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 max {
   x = f32[] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/slice_extra_out_padding.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/slice_extra_out_padding.hlo
index acbe1437027..2309066fff4 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/slice_extra_out_padding.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/slice_extra_out_padding.hlo
@@ -1,6 +1,7 @@
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
+// RUN: fusion_to_triton --xla_gpu_unsupported_enable_triton_multi_output_fusion %s 2>&1 | FileCheck %s --check-prefixes=CHECK-ERROR
+// CHECK-ERROR: Only single-result fusions are supported for now. Received 2 roots.
 // TODO: b/502910372 - support multi-output fusion.
-// RUN-NOT: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN-NOT: triton_test_correctness --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
 
 fusion {
   param_0.1 = f32[64] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/slice_padding.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/slice_padding.hlo
index 3f9bf6b5a5e..dc239f63d5f 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/slice_padding.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/slice_padding.hlo
@@ -1,6 +1,7 @@
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
+// RUN: fusion_to_triton --xla_gpu_unsupported_enable_triton_multi_output_fusion %s 2>&1 | FileCheck %s --check-prefixes=CHECK-ERROR
+// CHECK-ERROR: Only single-result fusions are supported for now. Received 2 roots.
 // TODO: b/502910372 - support multi-output fusion.
-// RUN-NOT: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
+// RUN-NOT: triton_test_correctness --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
 
 fusion {
   param_0.1 = f32[64] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/strided_iota_4d.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/strided_iota_4d.hlo
index b7882181daf..258dab43e5f 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/strided_iota_4d.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/strided_iota_4d.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   iota = f32[3,4,1000,5] iota(), iota_dimension=2
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/transpose/transpose_3d.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/transpose/transpose_3d.hlo
index c263d40de07..59264a22dbf 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/transpose/transpose_3d.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/transpose/transpose_3d.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton %s --xla_gpu_experimental_enable_tiling_propagation | FileCheck %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: fusion_to_triton %s | FileCheck %s
+// RUN: triton_test_correctness %s
 
 fusion {
   param_0 = f32[15,7,3] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/transpose/transpose_3d_tile_not_power_of_2.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/transpose/transpose_3d_tile_not_power_of_2.hlo
index 9c1cbec9385..5e5dbeaf0fc 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/transpose/transpose_3d_tile_not_power_of_2.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/transpose/transpose_3d_tile_not_power_of_2.hlo
@@ -1,5 +1,4 @@
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-// RUN: triton_test_correctness %s --xla_gpu_experimental_enable_tiling_propagation
+// RUN: triton_test_correctness %s
 
 fusion {
   param_0 = f32[3,8,20] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_multiprocess_test.cc b/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_multiprocess_test.cc
index a4aeec32783..2a42a496f19 100644
--- a/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_multiprocess_test.cc
+++ b/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_multiprocess_test.cc
@@ -187,7 +187,6 @@ absl::Status AllGatherMultiProcessTestBody(int node_id, int port) {
   debug_options.set_xla_gpu_autotune_level(0);
   debug_options.add_xla_gpu_experimental_use_collective_kernels(
       DebugOptions::COLLECTIVE_KERNEL_ALL_GATHER);
-  debug_options.set_xla_gpu_experimental_enable_tiling_propagation(true);
   debug_options.add_xla_gpu_unsupported_use_cross_host_one_shot_kernel(
       DebugOptions::ALLGATHER);
   *compile_options.executable_build_options.mutable_debug_options() =
diff --git a/third_party/xla/xla/backends/gpu/transforms/convert_triton_gemm_config_test.cc b/third_party/xla/xla/backends/gpu/transforms/convert_triton_gemm_config_test.cc
index 975f1430fb6..fa3b43aef74 100644
--- a/third_party/xla/xla/backends/gpu/transforms/convert_triton_gemm_config_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/convert_triton_gemm_config_test.cc
@@ -41,8 +41,7 @@ using ::absl_testing::IsOkAndHolds;
 namespace xla::gpu {
 namespace {
 
-class ConvertTritonGemmConfigTest : public HloHardwareIndependentTestBase,
-                                    public ::testing::WithParamInterface<bool> {
+class ConvertTritonGemmConfigTest : public HloHardwareIndependentTestBase {
  protected:
   ConvertTritonGemmConfigTest() { RegisterSymbolicExprStorage(&mlir_context_); }
   const se::DeviceDescription device_description_{
@@ -50,14 +49,6 @@ class ConvertTritonGemmConfigTest : public HloHardwareIndependentTestBase,
           se::GpuComputeCapability{se::CudaComputeCapability::Ampere()})};
   mlir::MLIRContext mlir_context_;
 
-  DebugOptions GetDebugOptionsForTest() const override {
-    DebugOptions debug_options =
-        HloHardwareIndependentTestBase::GetDebugOptionsForTest();
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
-        GetParam());
-    return debug_options;
-  }
-
   std::unique_ptr<VerifiedHloModule> RunConvertTritonGemmConfig(
       absl::string_view hlo, const bool expect_change = true) {
     std::unique_ptr<VerifiedHloModule> module =
@@ -70,7 +61,7 @@ class ConvertTritonGemmConfigTest : public HloHardwareIndependentTestBase,
   }
 };
 
-TEST_P(ConvertTritonGemmConfigTest, BasicDot) {
+TEST_F(ConvertTritonGemmConfigTest, BasicDot) {
   absl::string_view hlo = R"(
 dot {
   lhs = f32[8192,512] parameter(0)
@@ -116,7 +107,7 @@ ENTRY entry {
                    .has_triton_gemm_config());
 }
 
-TEST_P(ConvertTritonGemmConfigTest, ScaledDot) {
+TEST_F(ConvertTritonGemmConfigTest, ScaledDot) {
   absl::string_view hlo = R"(
 scaled_dot {
   lhs = bf16[4,4] parameter(0)
@@ -162,7 +153,7 @@ ENTRY entry {
 )"));
 }
 
-TEST_P(ConvertTritonGemmConfigTest, WavesPerEuPassthrough) {
+TEST_F(ConvertTritonGemmConfigTest, WavesPerEuPassthrough) {
   absl::string_view hlo = R"(
 dot {
   lhs = f32[8192,512] parameter(0)
@@ -200,7 +191,7 @@ ENTRY entry {
                    .has_triton_gemm_config());
 }
 
-TEST_P(ConvertTritonGemmConfigTest, TransposeDot) {
+TEST_F(ConvertTritonGemmConfigTest, TransposeDot) {
   absl::string_view hlo = R"(
 dot {
   lhs = f32[8192,512] parameter(0)
@@ -230,7 +221,7 @@ ENTRY entry {
 )"));
 }
 
-TEST_P(ConvertTritonGemmConfigTest, BatchGemm) {
+TEST_F(ConvertTritonGemmConfigTest, BatchGemm) {
   absl::string_view hlo = R"(
 dot {
   lhs = f32[2,8192,512] parameter(0)
@@ -260,8 +251,5 @@ ENTRY entry {
 )"));
 }
 
-INSTANTIATE_TEST_SUITE_P(ConvertTritonGemmConfigTestSuite,
-                         ConvertTritonGemmConfigTest, ::testing::Bool());
-
 }  // namespace
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter_test.cc
index dcf94515489..055647a03e2 100644
--- a/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter_test.cc
@@ -70,16 +70,13 @@ bool HasTritonBlockLevelFusionConfig(const HloInstruction* fusion) {
 
 class FusionBlockLevelRewriterTestBase
     : public HloHardwareIndependentTestBase,
-      public testing::WithParamInterface<std::tuple<bool, bool>> {
+      public testing::WithParamInterface<bool> {
  public:
   FusionBlockLevelRewriterTestBase() {
     RegisterSymbolicExprStorage(&mlir_context_);
   }
 
-  bool EnableTilingPropagation() const { return std::get<0>(GetParam()); }
-  bool EnableSameShapeMultiOutputFusion() const {
-    return std::get<1>(GetParam());
-  }
+  bool EnableSameShapeMultiOutputFusion() const { return GetParam(); }
 
  protected:
   se::DeviceDescription device_info_{TestGpuDeviceInfo::RTXA6000DeviceInfo(
@@ -95,8 +92,6 @@ class FusionBlockLevelRewriterTest : public FusionBlockLevelRewriterTestBase {
         HloHardwareIndependentTestBase::GetDebugOptionsForTest();
     debug_options.set_xla_gpu_experimental_enable_fusion_block_level_rewriter(
         true);
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
-        EnableTilingPropagation());
     debug_options
         .set_xla_gpu_experimental_enable_same_shape_multi_output_fusion(
             EnableSameShapeMultiOutputFusion());
@@ -104,22 +99,12 @@ class FusionBlockLevelRewriterTest : public FusionBlockLevelRewriterTestBase {
   }
 };
 
-INSTANTIATE_TEST_SUITE_P(
-    FusionBlockLevelRewriterTest, FusionBlockLevelRewriterTest,
-    testing::Combine(testing::Bool(), testing::Bool()),
-    [](const testing::TestParamInfo<std::tuple<bool, bool>>& info) {
-      std::vector<std::string> parts;
-      if (std::get<0>(info.param)) {
-        parts.push_back("TilingPropagation");
-      }
-      if (std::get<1>(info.param)) {
-        parts.push_back("SameShapeMultiOutputFusion");
-      }
-      if (parts.empty()) {
-        return std::string("Default");
-      }
-      return absl::StrJoin(parts, "_");
-    });
+INSTANTIATE_TEST_SUITE_P(FusionBlockLevelRewriterTest,
+                         FusionBlockLevelRewriterTest, testing::Bool(),
+                         [](const testing::TestParamInfo<bool>& info) {
+                           return info.param ? "SameShapeMultiOutputFusion"
+                                             : "Default";
+                         });
 
 TEST_P(FusionBlockLevelRewriterTest,
        DoesNotRewriteFusionThatIsAlreadyBlockLevel) {
@@ -224,9 +209,7 @@ ENTRY entry {
                                &mlir_context_)
           .Run(module.get());
 
-  const bool should_rewrite =
-      EnableSameShapeMultiOutputFusion() && EnableTilingPropagation();
-  if (should_rewrite) {
+  if (EnableSameShapeMultiOutputFusion()) {
     EXPECT_THAT(result, IsOkAndHolds(true));
     const HloInstruction* root =
         module->entry_computation()->root_instruction();
@@ -424,10 +407,10 @@ ENTRY entry {
       FusionBlockLevelRewriter(device_info_, HloCostAnalysis::DefaultShapeSize,
                                &mlir_context_)
           .Run(module.get()),
-      absl_testing::IsOkAndHolds(EnableTilingPropagation()));
+      absl_testing::IsOkAndHolds(true));
 
   const HloInstruction* root = module->entry_computation()->root_instruction();
-  EXPECT_EQ(HasTritonBlockLevelFusionConfig(root), EnableTilingPropagation());
+  EXPECT_TRUE(HasTritonBlockLevelFusionConfig(root));
 }
 
 TEST_P(FusionBlockLevelRewriterTest, DoesNotRewriteScanFusionWithCarryOutput) {
@@ -510,7 +493,6 @@ ENTRY entry {
       HloHardwareIndependentTestBase::GetDebugOptionsForTest();
   debug_options.set_xla_gpu_experimental_enable_fusion_block_level_rewriter(
       false);
-  debug_options.set_xla_gpu_experimental_enable_tiling_propagation(true);
   debug_options.set_xla_gpu_experimental_enable_same_shape_multi_output_fusion(
       true);
 
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 7f58ce7c73a..65a8047fc74 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
@@ -94,70 +94,35 @@ class GemmFusionTestBase : public HloHardwareIndependentTestBase {
   }
 };
 
-class GemmFusionTest : public GemmFusionTestBase,
-                       public ::testing::WithParamInterface<bool> {
- public:
-  DebugOptions GetDebugOptionsForTest() const override {
-    DebugOptions debug_options = GemmFusionTestBase::GetDebugOptionsForTest();
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
-        GetParam());
-    return debug_options;
-  }
-};
+using GemmFusionTest = GemmFusionTestBase;
 
-// While we launch, create a parameterized test to test all combinations. This
-// test class has 2 parameters:
-// 1. Whether to use Gemm Fusion V1 or V2.
-// 2. Whether to use symbolic analysis or tiling propagation.
-class GemmFusionTestVersioned
-    : public GemmFusionTestBase,
-      public ::testing::WithParamInterface<std::tuple<bool, bool>> {
+// While we launch, create a parameterized test to test both Gemm Fusion V1 and
+// V2.
+class GemmFusionTestVersioned : public GemmFusionTestBase,
+                                public ::testing::WithParamInterface<bool> {
  public:
   DebugOptions GetDebugOptionsForTest() const override {
     DebugOptions debug_options = GemmFusionTestBase::GetDebugOptionsForTest();
-    debug_options.set_xla_gpu_experimental_gemm_fusion_v2(
-        std::get<0>(GetParam()));
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
-        std::get<1>(GetParam()));
+    debug_options.set_xla_gpu_experimental_gemm_fusion_v2(GetParam());
     return debug_options;
   }
 };
 
 // Fixture for tests that are only meant to work for the new implementation
 // of dot fusion.
-class GemmFusionTestV2 : public GemmFusionTestBase,
-                         public ::testing::WithParamInterface<bool> {
+class GemmFusionTestV2 : public GemmFusionTestBase {
  public:
   DebugOptions GetDebugOptionsForTest() const override {
     DebugOptions debug_options = GemmFusionTestBase::GetDebugOptionsForTest();
     debug_options.set_xla_gpu_experimental_gemm_fusion_v2(true);
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
-        GetParam());
     return debug_options;
   }
 };
 
 INSTANTIATE_TEST_SUITE_P(
-    GemmFusionTestVersioned, GemmFusionTestVersioned,
-    ::testing::Combine(::testing::Bool(), ::testing::Bool()),
+    GemmFusionTestVersioned, GemmFusionTestVersioned, ::testing::Bool(),
     [](const ::testing::TestParamInfo<GemmFusionTestVersioned::ParamType>&
-           info) {
-      return absl::StrCat(
-          std::get<0>(info.param) ? "V2" : "V1", "_",
-          std::get<1>(info.param) ? "TilingPropagation" : "SymbolicAnalysis");
-    });
-
-INSTANTIATE_TEST_SUITE_P(
-    GemmFusionTest, GemmFusionTest, ::testing::Bool(),
-    [](const ::testing::TestParamInfo<GemmFusionTest::ParamType>& info) {
-      return info.param ? "TilingPropagation" : "SymbolicAnalysis";
-    });
-
-INSTANTIATE_TEST_SUITE_P(
-    GemmFusionTestV2, GemmFusionTestV2, ::testing::Bool(),
-    [](const ::testing::TestParamInfo<GemmFusionTestV2::ParamType>& info) {
-      return info.param ? "TilingPropagation" : "SymbolicAnalysis";
-    });
+           info) { return info.param ? "V2" : "V1"; });
 
 TEST_P(GemmFusionTestVersioned, TransposeSubdimensionGroup) {
   // This HLO is artificial because unnecessary reshapes get optimized
@@ -181,7 +146,7 @@ ENTRY e {
               GmockMatch(m::Fusion(m::Op(), m::Op())));
 }
 
-TEST_P(GemmFusionTestV2, BitcastIsHoistedAboveElementwise) {
+TEST_F(GemmFusionTestV2, BitcastIsHoistedAboveElementwise) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -199,7 +164,7 @@ ENTRY e {
       GmockMatch(m::Fusion(m::Parameter(), m::Bitcast(m::Parameter()))));
 }
 
-TEST_P(GemmFusionTestV2, BitcastIsHoistedAboveMultiOperandElementwise) {
+TEST_F(GemmFusionTestV2, BitcastIsHoistedAboveMultiOperandElementwise) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -218,7 +183,7 @@ ENTRY e {
                                    m::Bitcast(m::Parameter()))));
 }
 
-TEST_P(GemmFusionTestV2, BitcastIsHoistedAboveBroadcast) {
+TEST_F(GemmFusionTestV2, BitcastIsHoistedAboveBroadcast) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -236,7 +201,7 @@ ENTRY e {
       GmockMatch(m::Fusion(m::Parameter(), m::Bitcast(m::Parameter()))));
 }
 
-TEST_P(GemmFusionTestV2, BitcastIsHoistedAboveTranspose) {
+TEST_F(GemmFusionTestV2, BitcastIsHoistedAboveTranspose) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -255,7 +220,7 @@ ENTRY e {
       GmockMatch(m::Fusion(m::Parameter(), m::Bitcast(m::Parameter()))));
 }
 
-TEST_P(GemmFusionTestV2, BitcastSinksPastElementwise) {
+TEST_F(GemmFusionTestV2, BitcastSinksPastElementwise) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -272,7 +237,7 @@ ENTRY e {
               GmockMatch(m::Bitcast(m::Fusion())));
 }
 
-TEST_P(GemmFusionTestV2, MultipleBitcastsSink) {
+TEST_F(GemmFusionTestV2, MultipleBitcastsSink) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -295,7 +260,7 @@ ENTRY e {
 )");
 }
 
-TEST_P(GemmFusionTestV2, BitcastSinksPastBroadcast) {
+TEST_F(GemmFusionTestV2, BitcastSinksPastBroadcast) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -312,7 +277,7 @@ ENTRY e {
               GmockMatch(m::Bitcast(m::Fusion())));
 }
 
-TEST_P(GemmFusionTestV2, BitcastSinksPastTranspose) {
+TEST_F(GemmFusionTestV2, BitcastSinksPastTranspose) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -330,7 +295,7 @@ ENTRY e {
               GmockMatch(m::Bitcast(m::Fusion())));
 }
 
-TEST_P(GemmFusionTestV2, UnhoistedBitcastCanStillBeFused) {
+TEST_F(GemmFusionTestV2, UnhoistedBitcastCanStillBeFused) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -348,7 +313,7 @@ ENTRY e {
               GmockMatch(m::Fusion(m::Parameter(), m::Parameter())));
 }
 
-TEST_P(GemmFusionTestV2, UnhoistedBitcastIsNotFusedAtEdge) {
+TEST_F(GemmFusionTestV2, UnhoistedBitcastIsNotFusedAtEdge) {
   // The bitcast/reshape cannot be hoisted above the concat, but all are
   // included in the search space. When it cannot tile the concat, the fusion
   // cuts off before the concat. We need to make sure the bitcast is on the
@@ -374,7 +339,7 @@ ENTRY e {
       GmockMatch(m::Fusion(m::Bitcast(m::Concatenate()), m::Parameter())));
 }
 
-TEST_P(GemmFusionTestV2, UnsunkBitcastIsNotFusedAtRoot) {
+TEST_F(GemmFusionTestV2, UnsunkBitcastIsNotFusedAtRoot) {
   // The bitcast/reshape cannot be sunk below the concat, but all are included
   // in the search space. When it cannot tile the reshape, and cuts off the
   // fusion between the bitcast & the reshape. We need to make sure the bitcast
@@ -400,7 +365,7 @@ ENTRY e {
                                         m::Parameter())));
 }
 
-TEST_P(GemmFusionTestV2, PartiallySunkBitcastIsNotFusedAtRoot) {
+TEST_F(GemmFusionTestV2, PartiallySunkBitcastIsNotFusedAtRoot) {
   // The bitcast/reshape cannot be sunk below the concat, but all are included
   // in the search space. When it cannot tile the reshape, the fusion cuts off
   // between the bitcast & the reshape. We need to make sure the bitcast
@@ -427,7 +392,7 @@ ENTRY e {
                                         m::Parameter())));
 }
 
-TEST_P(GemmFusionTestV2, BitcastOperandOfUserOfDotIsHoisted) {
+TEST_F(GemmFusionTestV2, BitcastOperandOfUserOfDotIsHoisted) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -447,7 +412,7 @@ ENTRY e {
                                    m::Bitcast(m::Parameter()))));
 }
 
-TEST_P(GemmFusionTestV2, BitcastIsBothSunkAndHoisted) {
+TEST_F(GemmFusionTestV2, BitcastIsBothSunkAndHoisted) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -467,7 +432,7 @@ ENTRY e {
                                       m::Parameter(), m::Parameter()))));
 }
 
-TEST_P(GemmFusionTestV2, BitcastIsHoistedAboveConstants) {
+TEST_F(GemmFusionTestV2, BitcastIsHoistedAboveConstants) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -498,7 +463,7 @@ ENTRY e {
 )");
 }
 
-TEST_P(GemmFusionTestV2, DoNotHoistBitcastOverParameterWithNonBitcastUsers) {
+TEST_F(GemmFusionTestV2, DoNotHoistBitcastOverParameterWithNonBitcastUsers) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 ENTRY e {
   p0 = bf16[32,12] parameter(0)
@@ -513,7 +478,7 @@ ENTRY e {
               GmockMatch(m::Fusion(m::Parameter())));
 }
 
-TEST_P(GemmFusionTestV2, HoistBitcastOverMultipleEqualBitcasts) {
+TEST_F(GemmFusionTestV2, HoistBitcastOverMultipleEqualBitcasts) {
   // This is an example of when a single bitcast turns into 2 different bitcasts
   // due to hoisting over a binary operation, but then they come from the same
   // parameter, and should still get hoisted.
@@ -536,7 +501,7 @@ ENTRY e {
                                    m::Bitcast(m::Parameter()))));
 }
 
-TEST_P(GemmFusionTestV2, HoistBitcastOverParameterWithMultipleBitcastUsers) {
+TEST_F(GemmFusionTestV2, HoistBitcastOverParameterWithMultipleBitcastUsers) {
   // Regression test for b/524943134.
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
@@ -556,7 +521,7 @@ ENTRY e {
                                    m::Bitcast(m::Parameter()))));
 }
 
-TEST_P(GemmFusionTestV2, HoistBitcastAcrossTransposeWithTrivialDimension) {
+TEST_F(GemmFusionTestV2, HoistBitcastAcrossTransposeWithTrivialDimension) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 ENTRY e {
   p0 = f32[42,1]{1,0} parameter(0)
@@ -623,7 +588,7 @@ ENTRY e {
       GmockMatch(m::Fusion(ParamOrBitcastParam(), ParamOrBitcastParam())));
 }
 
-TEST_P(GemmFusionTest, DoNotTriggerOnUnsupportedOutputConversions) {
+TEST_F(GemmFusionTest, DoNotTriggerOnUnsupportedOutputConversions) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"(
 ENTRY e {
@@ -747,7 +712,7 @@ ENTRY e {
 }
 
 // TODO(b/417172838): support dynamic slice op.
-TEST_P(GemmFusionTest, DISABLED_DynamicSliceIsFused) {
+TEST_F(GemmFusionTest, DISABLED_DynamicSliceIsFused) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"(
 ENTRY e {
@@ -769,7 +734,7 @@ ENTRY e {
 }
 
 // TODO(b/417172838): support dynamic slice op.
-TEST_P(GemmFusionTest, DISABLED_DynamicSlicesAreFusedEvenIfTheyShareIndices) {
+TEST_F(GemmFusionTest, DISABLED_DynamicSlicesAreFusedEvenIfTheyShareIndices) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"(
 ENTRY e {
@@ -798,7 +763,7 @@ ENTRY e {
 }
 
 // TODO(b/417172838): support dynamic slice op.
-TEST_P(GemmFusionTest, DISABLED_DoNotFuseDynamicSliceOfNonMajorFragments) {
+TEST_F(GemmFusionTest, DISABLED_DoNotFuseDynamicSliceOfNonMajorFragments) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"(
 ENTRY e {
@@ -817,7 +782,7 @@ ENTRY e {
 }
 
 // TODO(b/417172838): support dynamic slice op.
-TEST_P(GemmFusionTest,
+TEST_F(GemmFusionTest,
        DISABLED_CanFuseDynamicSliceOfContractingDimIfItIsMajor) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"(
@@ -837,7 +802,7 @@ ENTRY e {
                                     m::Constant(), m::Constant()))));
 }
 
-TEST_P(GemmFusionTest, SliceToDegenerateIsSkipped) {
+TEST_F(GemmFusionTest, SliceToDegenerateIsSkipped) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"(
 ENTRY e {
@@ -960,7 +925,7 @@ ENTRY e {
               GmockMatch(m::Fusion(m::Parameter(), ParamOrBitcastParam())));
 }
 
-TEST_P(GemmFusionTest, DoNotFuseIncompatibleDimensionSplits) {
+TEST_F(GemmFusionTest, DoNotFuseIncompatibleDimensionSplits) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"(
 ENTRY e {
@@ -1101,7 +1066,7 @@ ENTRY e {
             TritonFusionAnalysis::kMaxParameterPerDotOperand * 2);
 }
 
-TEST_P(GemmFusionTest,
+TEST_F(GemmFusionTest,
        DoNotFuseTooManyParametersWhenAnInstructionWouldAddMultipleParameters) {
   static_assert(TritonFusionAnalysis::kMaxParameterPerDotOperand == 4,
                 "We have to update this test.");
@@ -1190,7 +1155,7 @@ ENTRY e {
   // ~VerifiedHloModule() will verify the module.
 }
 
-TEST_P(GemmFusionTest, EachScopeIsFusedToASeparateSubgraph) {
+TEST_F(GemmFusionTest, EachScopeIsFusedToASeparateSubgraph) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"(
 ENTRY e {
@@ -1225,7 +1190,7 @@ CHECK-SAME: __triton_gemm
 // way, so the same parameter node is reused for them.
 // The reuse happens per "operand fusion", so the add of the LHS and RHS still
 // use different nodes.
-TEST_P(GemmFusionTest, ParamNodesAreReusedIfTheyHaveTheSameIterSpec) {
+TEST_F(GemmFusionTest, ParamNodesAreReusedIfTheyHaveTheSameIterSpec) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"(
 ENTRY e {
@@ -1254,7 +1219,7 @@ CHECK-SAME: __triton_gemm
 
 // NEGATE has the same iteration spec at both usages, so the node is reused
 // (implying that P0 is also reused).
-TEST_P(GemmFusionTest, NonParamNodesAreReusedIfTheyHaveTheSameIterSpec) {
+TEST_F(GemmFusionTest, NonParamNodesAreReusedIfTheyHaveTheSameIterSpec) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"(
 ENTRY e {
@@ -1288,7 +1253,7 @@ CHECK-SAME: __triton_gemm
 
 // The direct read of the input and the transposed read of the input have
 // different iteration specs, so we don't reuse the node.
-TEST_P(GemmFusionTest, NodesAreNotReusedIfTheyHaveDifferentIterSpecs) {
+TEST_F(GemmFusionTest, NodesAreNotReusedIfTheyHaveDifferentIterSpecs) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"(
 ENTRY e {
@@ -1420,7 +1385,7 @@ ENTRY e {
             1);
 }
 
-TEST_P(GemmFusionTest, ParameterUsedNonElementwiseTwiceIsFusedOnBothPaths) {
+TEST_F(GemmFusionTest, ParameterUsedNonElementwiseTwiceIsFusedOnBothPaths) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"(
 HloModule t
@@ -1440,7 +1405,7 @@ ENTRY e {
       GmockMatch((m::Fusion(m::Parameter(), m::Parameter(), m::Parameter()))));
 }
 
-TEST_P(GemmFusionTest,
+TEST_F(GemmFusionTest,
        ComputationParameterWithMultipleUsersIsNotTrivialToFuse) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"(
@@ -1634,7 +1599,7 @@ e {
                                    GmockMatch(m::Parameter())));
 }
 
-TEST_P(GemmFusionTest, ConcatenationOfBatchIsNotFused) {
+TEST_F(GemmFusionTest, ConcatenationOfBatchIsNotFused) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"(
 e {
@@ -1652,7 +1617,7 @@ e {
               GmockMatch((m::Fusion(m::Concatenate(), m::Parameter()))));
 }
 
-TEST_P(GemmFusionTest,
+TEST_F(GemmFusionTest,
        DifferentConcatenationOfSameParametersIsFusedViaNodeDuplication) {
   // It means that the same input is passed to the fusion multiple times and
   // it's read differently for each.
@@ -1758,7 +1723,7 @@ e {
                                         m::Bitcast(m::GetTupleElement()))))));
 }
 
-TEST_P(GemmFusionTest, DoNotFuseNonProfitableDot) {
+TEST_F(GemmFusionTest, DoNotFuseNonProfitableDot) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -1783,14 +1748,7 @@ class SmallDotGemmFusionTest : public GemmFusionTest {
   }
 };
 
-INSTANTIATE_TEST_SUITE_P(
-    SmallDotGemmFusionTest, SmallDotGemmFusionTest, ::testing::Bool(),
-    [](const ::testing::TestParamInfo<SmallDotGemmFusionTest::ParamType>&
-           info) {
-      return info.param ? "TilingPropagation" : "SymbolicAnalysis";
-    });
-
-TEST_P(SmallDotGemmFusionTest, SkipSmallMatrixMultiplicationRewrite) {
+TEST_F(SmallDotGemmFusionTest, SkipSmallMatrixMultiplicationRewrite) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -1813,7 +1771,7 @@ ENTRY e {
 })");
 }
 
-TEST_P(SmallDotGemmFusionTest, LargeMatrixMultiplicationIsRewritten) {
+TEST_F(SmallDotGemmFusionTest, LargeMatrixMultiplicationIsRewritten) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -2017,7 +1975,7 @@ TEST_P(GemmFusionTestVersioned, Int4WithMinorReshapedBatchDimIsNotRewritten) {
   EXPECT_THAT(GemmFusion(device_info_).Run(module.get()), IsOkAndHolds(false));
 }
 
-TEST_P(GemmFusionTestV2,
+TEST_F(GemmFusionTestV2,
        Int4WithMinorReshapedBatchAndNonContractingDimIsRewritten) {
   constexpr absl::string_view kInt4Dot = R"(
     ENTRY main {
@@ -2037,7 +1995,7 @@ TEST_P(GemmFusionTestV2,
   EXPECT_THAT(GemmFusion(device_info_).Run(module.get()), IsOkAndHolds(true));
 }
 
-TEST_P(GemmFusionTestV2,
+TEST_F(GemmFusionTestV2,
        Int4WithUntileableMinorNonContractingDimIsNotRewritten) {
   constexpr absl::string_view kInt4Dot = R"(
     ENTRY main {
@@ -2055,7 +2013,7 @@ TEST_P(GemmFusionTestV2,
   EXPECT_THAT(GemmFusion(device_info_).Run(module.get()), IsOkAndHolds(false));
 }
 
-TEST_P(GemmFusionTest, ScaledDotIsFused) {
+TEST_F(GemmFusionTest, ScaledDotIsFused) {
   constexpr absl::string_view kHloText = R"(
     HloModule ScaledDotIsFused
 
@@ -2258,14 +2216,9 @@ ENTRY e {
 using GemmFusionProfitabilityTest = GemmFusionTestVersioned;
 
 INSTANTIATE_TEST_SUITE_P(
-    GemmFusionProfitabilityTest, GemmFusionProfitabilityTest,
-    ::testing::Combine(::testing::Bool(), ::testing::Bool()),
+    GemmFusionProfitabilityTest, GemmFusionProfitabilityTest, ::testing::Bool(),
     [](const ::testing::TestParamInfo<GemmFusionProfitabilityTest::ParamType>&
-           info) {
-      return std::string(std::get<0>(info.param) ? "V2_" : "V1_") +
-             std::string(std::get<1>(info.param) ? "TilingPropagation"
-                                                 : "SymbolicAnalysis");
-    });
+           info) { return info.param ? "V2" : "V1"; });
 
 TEST_P(GemmFusionProfitabilityTest, UnprofitableOperand) {
   // Tests that a large elementwise operation with multiple users is not fused.
@@ -2389,7 +2342,7 @@ ENTRY e {
 
 // Regression test for a crash/verifier failure when hoisting type-changing
 // bitcasts.
-TEST_P(GemmFusionTestV2, TypeChangingBitcastIsNotHoisted) {
+TEST_F(GemmFusionTestV2, TypeChangingBitcastIsNotHoisted) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -2505,7 +2458,7 @@ ENTRY main {
                   GmockMatch(TransposeOrBitcastTranspose())));
 }
 
-TEST_P(GemmFusionTestV2, AllowTransposeSwappingBatch) {
+TEST_F(GemmFusionTestV2, AllowTransposeSwappingBatch) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"(
 HloModule module
@@ -2628,7 +2581,7 @@ ENTRY e {
               GmockMatch(TransposeOrTransposeBitcast(m::Fusion())));
 }
 
-TEST_P(GemmFusionTestV2, AllowEpilogueTransposeWithSwapsWithinDimensionGroup) {
+TEST_F(GemmFusionTestV2, AllowEpilogueTransposeWithSwapsWithinDimensionGroup) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -2648,7 +2601,7 @@ ENTRY e {
               GmockMatch(m::Fusion()));
 }
 
-TEST_P(GemmFusionTestV2,
+TEST_F(GemmFusionTestV2,
        DoNotFuseNonContractingTransposeIfProducerCanAbsorbIt) {
   // The transpose only reorders the sub-dimensions of the LHS non-contracting
   // dimension, and its producer (a normalization) is not fused into the GEMM.
@@ -2687,7 +2640,7 @@ ENTRY e {
                   m::Bitcast(m::Transpose(m::Multiply()))))));
 }
 
-TEST_P(GemmFusionTestV2, FuseNonContractingTransposeIfProducerIsFused) {
+TEST_F(GemmFusionTestV2, FuseNonContractingTransposeIfProducerIsFused) {
   // The producer (a multiply by a broadcast parameter) is fused into the GEMM,
   // so the transpose is fused as well.
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
@@ -2711,7 +2664,7 @@ ENTRY e {
   EXPECT_THAT(fusion->operands(), Each(GmockMatch(ParamOrBitcastParam())));
 }
 
-TEST_P(GemmFusionTestV2,
+TEST_F(GemmFusionTestV2,
        FuseTransposeKeepingNonContractingDimensionContiguous) {
   // The transpose only swaps the batch and non-contracting dimensions.
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
@@ -2738,7 +2691,7 @@ ENTRY e {
                                    GmockMatch(m::Parameter())));
 }
 
-TEST_P(GemmFusionTestV2, FuseNonContractingTransposeMovingMinorDimension) {
+TEST_F(GemmFusionTestV2, FuseNonContractingTransposeMovingMinorDimension) {
   // The contracting dimension is interleaved between the two parts of the
   // non-contracting dimension, but the transpose also moves the minor
   // dimension, so it would not be free to fuse into the producer.
@@ -2766,7 +2719,7 @@ ENTRY e {
                                    GmockMatch(m::Parameter())));
 }
 
-TEST_P(GemmFusionTestV2, ConcatResetTrackerCrash) {
+TEST_F(GemmFusionTestV2, ConcatResetTrackerCrash) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -2873,7 +2826,7 @@ ENTRY e {
       GmockMatch(m::Fusion(m::Parameter(), m::Parameter(), m::Parameter())));
 }
 
-TEST_P(GemmFusionTestV2,
+TEST_F(GemmFusionTestV2,
        DoNotFuseTransposeAboveBroadcastWithSwappedRhsNonContractingDimensions) {
   // The RHS non-contracting dimensions are swapped between the broadcast and
   // the dot (allowed because the minor dimension is coalesced). Their order
@@ -2906,7 +2859,7 @@ ENTRY e {
               Contains(GmockMatch(TransposeOrBitcastTranspose())));
 }
 
-TEST_P(GemmFusionTestV2, HoistBitcastOverTypeConvertingBitcast) {
+TEST_F(GemmFusionTestV2, HoistBitcastOverTypeConvertingBitcast) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -2938,7 +2891,7 @@ ENTRY e {
 )");
 }
 
-TEST_P(GemmFusionTestV2, DoNotFuseConcatOnContractingDimension) {
+TEST_F(GemmFusionTestV2, DoNotFuseConcatOnContractingDimension) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -2956,7 +2909,7 @@ ENTRY e {
               GmockMatch(m::Fusion(m::Concatenate(), m::Parameter())));
 }
 
-TEST_P(GemmFusionTestV2, FuseConcatOnNonContractingDimension) {
+TEST_F(GemmFusionTestV2, FuseConcatOnNonContractingDimension) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -2975,7 +2928,7 @@ ENTRY e {
       GmockMatch(m::Fusion(m::Parameter(), m::Parameter(), m::Parameter())));
 }
 
-TEST_P(GemmFusionTestV2, DoNotFuseIndivisibleConcat) {
+TEST_F(GemmFusionTestV2, DoNotFuseIndivisibleConcat) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -2997,10 +2950,7 @@ ENTRY e {
 // [3, 8] pair of dimensions, where the dimension of size 3 is the major one.
 // Fusing the transpose results in a fusion that is tileable in theory, but not
 // with any of the power-of-2 tile sizes that we use (b/559517319).
-TEST_P(GemmFusionTestV2, DoNotFuseIfNoDefaultTilingSatisfiesConstraints) {
-  if (!GetParam()) {
-    GTEST_SKIP() << "Only fixed for tiling propagation.";
-  }
+TEST_F(GemmFusionTestV2, DoNotFuseIfNoDefaultTilingSatisfiesConstraints) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -3022,7 +2972,7 @@ ENTRY e {
                            m::Bitcast(m::Parameter()))));
 }
 
-TEST_P(GemmFusionTestV2, DoNotFuseConcatenateUser) {
+TEST_F(GemmFusionTestV2, DoNotFuseConcatenateUser) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"(
 HloModule m
diff --git a/third_party/xla/xla/backends/gpu/transforms/priority_fusion_test.cc b/third_party/xla/xla/backends/gpu/transforms/priority_fusion_test.cc
index 8482105acfb..c3006bca1d9 100644
--- a/third_party/xla/xla/backends/gpu/transforms/priority_fusion_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/priority_fusion_test.cc
@@ -58,19 +58,10 @@ using ::testing::UnorderedElementsAre;
 namespace xla {
 namespace gpu {
 
-class PriorityFusionTest : public HloHardwareIndependentTestBase,
-                           public ::testing::WithParamInterface<bool> {
+class PriorityFusionTest : public HloHardwareIndependentTestBase {
  public:
   PriorityFusionTest() { RegisterSymbolicExprStorage(&mlir_context_); }
 
-  DebugOptions GetDebugOptionsForTest() const override {
-    DebugOptions debug_options =
-        HloHardwareIndependentTestBase::GetDebugOptionsForTest();
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
-        GetParam());
-    return debug_options;
-  }
-
   std::vector<HloFusionAnalysis::EmitterFusionKind> RunAndGetFusionKinds(
       absl::string_view hlo) {
     auto module = ParseAndReturnVerifiedModule(hlo).value();
@@ -101,13 +92,7 @@ class PriorityFusionTest : public HloHardwareIndependentTestBase,
   }();
 };
 
-INSTANTIATE_TEST_SUITE_P(
-    PriorityFusionTest, PriorityFusionTest, ::testing::Bool(),
-    [](const ::testing::TestParamInfo<PriorityFusionTest::ParamType>& info) {
-      return info.param ? "TilingPropagation" : "SymbolicAnalysis";
-    });
-
-TEST_P(PriorityFusionTest, ParallelTilingSearchMatchesSerialTilingSearch) {
+TEST_F(PriorityFusionTest, ParallelTilingSearchMatchesSerialTilingSearch) {
   constexpr absl::string_view kHlo = R"(
     HloModule test_module
 
@@ -147,7 +132,7 @@ TEST_P(PriorityFusionTest, ParallelTilingSearchMatchesSerialTilingSearch) {
             serial_module->ToString(HloPrintOptions::ShortParsable()));
 }
 
-TEST_P(PriorityFusionTest, FuseWithSharedArgument) {
+TEST_F(PriorityFusionTest, FuseWithSharedArgument) {
   auto module = ParseAndReturnVerifiedModule(R"(
     HloModule test_module
 
@@ -170,7 +155,7 @@ TEST_P(PriorityFusionTest, FuseWithSharedArgument) {
   EXPECT_EQ(root->fusion_kind(), HloInstruction::FusionKind::kLoop);
 }
 
-TEST_P(PriorityFusionTest, FusionOnStreamAnnotatedComputation) {
+TEST_F(PriorityFusionTest, FusionOnStreamAnnotatedComputation) {
   auto module = ParseAndReturnVerifiedModule(R"(
     HloModule test_module
     stream {
@@ -204,7 +189,7 @@ TEST_P(PriorityFusionTest, FusionOnStreamAnnotatedComputation) {
   EXPECT_EQ(called_root->fusion_kind(), HloInstruction::FusionKind::kLoop);
 }
 
-TEST_P(PriorityFusionTest, FusionFusionWithDuplication) {
+TEST_F(PriorityFusionTest, FusionFusionWithDuplication) {
   absl::string_view kHlo = R"(
     HloModule test_module
 
@@ -240,7 +225,7 @@ CHECK-NEXT: ROOT {{.*}} tuple(%[[FUSION_0]], %[[FUSION_1]])
   )");
 }
 
-TEST_P(PriorityFusionTest, FuseBroadcastIntoBitcastConsumers) {
+TEST_F(PriorityFusionTest, FuseBroadcastIntoBitcastConsumers) {
   absl::string_view kHlo = R"(
     HloModule test_module
 
@@ -258,7 +243,7 @@ CHECK-NEXT: ROOT %{{.*}} fusion(%[[PARAM]])
   )");
 }
 
-TEST_P(PriorityFusionTest, FuseWideningConvertIntoConsumers) {
+TEST_F(PriorityFusionTest, FuseWideningConvertIntoConsumers) {
   absl::string_view kHlo = R"(
     HloModule test_module
 
@@ -281,7 +266,7 @@ CHECK-NEXT: ROOT %{{.*}} = (f32[512]{0}, s32[512]{0}) tuple(%[[FUSION_F32]], %[[
   )");
 }
 
-TEST_P(PriorityFusionTest, DoNotFuseBitWidthChangingBitcast) {
+TEST_F(PriorityFusionTest, DoNotFuseBitWidthChangingBitcast) {
   // `neg` is the producer that could be fused with `bitcast` and `mul`, but
   // since `bitcast` changes the bit width, we don't fuse it.
   auto module = *ParseAndReturnVerifiedModule(R"(
@@ -297,7 +282,7 @@ TEST_P(PriorityFusionTest, DoNotFuseBitWidthChangingBitcast) {
               absl_testing::IsOkAndHolds(false));
 }
 
-TEST_P(PriorityFusionTest, FuseConvertIntoReduce) {
+TEST_F(PriorityFusionTest, FuseConvertIntoReduce) {
   absl::string_view kHlo = R"(
     HloModule test_module
 
@@ -336,7 +321,7 @@ CHECK-COUNT-3: fusion
   )");
 }
 
-TEST_P(PriorityFusionTest, ReductionEpilogueFusionRegressionTest) {
+TEST_F(PriorityFusionTest, ReductionEpilogueFusionRegressionTest) {
   // Regression test for epilogue fusion of convert into a reduction, even if
   // the convert has a bitcast as consumer.
   absl::string_view kHlo = R"(
@@ -391,7 +376,7 @@ CHECK: ROOT {{.*}} bitcast({{.*}}fusion{{.*}})
   )");
 }
 
-TEST_P(PriorityFusionTest, DoNotChangeReductionFusionToLoopFusion) {
+TEST_F(PriorityFusionTest, DoNotChangeReductionFusionToLoopFusion) {
   // Regression test for epilogue fusion of slice into a reduction. The fusion
   // kind for the reduction fusion is intentionally chosen to be set to kLoop,
   // as we cannot rely on reductions always having fusion kind kInput.
@@ -419,7 +404,7 @@ TEST_P(PriorityFusionTest, DoNotChangeReductionFusionToLoopFusion) {
               absl_testing::IsOkAndHolds(false));
 }
 
-TEST_P(PriorityFusionTest, DoNotFuseTransposeIntoReduce) {
+TEST_F(PriorityFusionTest, DoNotFuseTransposeIntoReduce) {
   absl::string_view kHlo = R"(
     HloModule test_module
 
@@ -495,7 +480,7 @@ TEST_P(PriorityFusionTest, DoNotFuseTransposeIntoReduce) {
                            Kind::kTranspose, Kind::kTranspose));
 }
 
-TEST_P(PriorityFusionTest, DoNotFuseReduceIntoReduce) {
+TEST_F(PriorityFusionTest, DoNotFuseReduceIntoReduce) {
   absl::string_view kHlo = R"(
     HloModule test_module
 
@@ -518,7 +503,7 @@ CHECK: ROOT {{.*}} reduce(
   )");
 }
 
-TEST_P(PriorityFusionTest, ConvertFusedIntoReduce) {
+TEST_F(PriorityFusionTest, ConvertFusedIntoReduce) {
   absl::string_view kHlo = R"(
     HloModule test_module
 
@@ -558,7 +543,7 @@ CHECK-NOT: fusion(
   )");
 }
 
-TEST_P(PriorityFusionTest, DoNotFuseDynamicUpdateSliceIntoReduce) {
+TEST_F(PriorityFusionTest, DoNotFuseDynamicUpdateSliceIntoReduce) {
   absl::string_view kHlo = R"(
     HloModule test_module
 
@@ -634,7 +619,7 @@ CHECK-COUNT-3: fusion(
   )");
 }
 
-TEST_P(PriorityFusionTest, DontFuseIntoFirstOperandOfScatter) {
+TEST_F(PriorityFusionTest, DontFuseIntoFirstOperandOfScatter) {
   auto module = *ParseAndReturnVerifiedModule(R"(
     HloModule test_module
 
@@ -674,7 +659,7 @@ TEST_P(PriorityFusionTest, DontFuseIntoFirstOperandOfScatter) {
 // This test is similar to DontFuseIntoFirstOperandOfScatter, but PriorityFusion
 // has a separate run to fuse constants. Fusing anything into a scatter fusion
 // will fail in the emitter.
-TEST_P(PriorityFusionTest, DontFuseConstantIntoFirstOperandOfScatter) {
+TEST_F(PriorityFusionTest, DontFuseConstantIntoFirstOperandOfScatter) {
   auto module = *ParseAndReturnVerifiedModule(R"(
     HloModule test_module
 
@@ -708,7 +693,7 @@ TEST_P(PriorityFusionTest, DontFuseConstantIntoFirstOperandOfScatter) {
                                     m::Broadcast(m::Constant()))));
 }
 
-TEST_P(PriorityFusionTest, DoNotFuseReduceIntoReduceEvenIfOccupancyIsHigh) {
+TEST_F(PriorityFusionTest, DoNotFuseReduceIntoReduceEvenIfOccupancyIsHigh) {
   constexpr absl::string_view kHlo = R"(
     HloModule test_module
 
@@ -731,7 +716,7 @@ CHECK: ROOT {{.*}} reduce(
   )");
 }
 
-TEST_P(PriorityFusionTest, FuseReductionEpilogueWithMultipleUsers) {
+TEST_F(PriorityFusionTest, FuseReductionEpilogueWithMultipleUsers) {
   // Regression test that verifies we correctly fuse the `log` into the reduce.
   constexpr absl::string_view kHlo = R"(
     HloModule test_module
@@ -765,7 +750,7 @@ TEST_P(PriorityFusionTest, FuseReductionEpilogueWithMultipleUsers) {
   )");
 }
 
-TEST_P(PriorityFusionTest, EpilogueFusion) {
+TEST_F(PriorityFusionTest, EpilogueFusion) {
   absl::string_view kHlo = R"(
     HloModule test_module
 
@@ -797,7 +782,7 @@ TEST_P(PriorityFusionTest, EpilogueFusion) {
 CHECK: ROOT {{.*}} = f32[8,4,128]{2,1,0} fusion(%p{{.*}}), kind=kInput, calls=%fused_computation)");
 }
 
-TEST_P(PriorityFusionTest, EpilogueFusionFails) {
+TEST_F(PriorityFusionTest, EpilogueFusionFails) {
   auto module = *ParseAndReturnVerifiedModule(R"(
     HloModule test_module
 
@@ -830,7 +815,7 @@ TEST_P(PriorityFusionTest, EpilogueFusionFails) {
               absl_testing::IsOkAndHolds(false));
 }
 
-TEST_P(PriorityFusionTest, DoNotFuseIntoRoot) {
+TEST_F(PriorityFusionTest, DoNotFuseIntoRoot) {
   auto module = *ParseAndReturnVerifiedModule(R"(
     HloModule test_module
 
@@ -848,7 +833,7 @@ TEST_P(PriorityFusionTest, DoNotFuseIntoRoot) {
               absl_testing::IsOkAndHolds(false));
 }
 
-TEST_P(PriorityFusionTest, DontFuseConcat) {
+TEST_F(PriorityFusionTest, DontFuseConcat) {
   // Regression test that verifies we don't fuse concat into a column reduction.
   auto module = *ParseAndReturnVerifiedModule(R"(
     HloModule module
@@ -901,7 +886,7 @@ TEST_P(PriorityFusionTest, DontFuseConcat) {
               absl_testing::IsOkAndHolds(false));
 }
 
-TEST_P(PriorityFusionTest, FuseOnlySmallConstant) {
+TEST_F(PriorityFusionTest, FuseOnlySmallConstant) {
   auto module = *ParseAndReturnVerifiedModule(R"(
     HloModule module
 
@@ -925,7 +910,7 @@ TEST_P(PriorityFusionTest, FuseOnlySmallConstant) {
                   m::Add(m::Parameter(), m::Broadcast(m::Constant())))));
 }
 
-TEST_P(PriorityFusionTest, FuseSmallConstantIntoTritonFusion) {
+TEST_F(PriorityFusionTest, FuseSmallConstantIntoTritonFusion) {
   auto module = *ParseAndReturnVerifiedModule(R"(
 HloModule module
 
@@ -955,7 +940,7 @@ ENTRY main {
               GmockMatch(m::Reduce(m::Parameter(), m::Constant())));
 }
 
-TEST_P(PriorityFusionTest, FuseProducerConsumerMergedNotTooLarge) {
+TEST_F(PriorityFusionTest, FuseProducerConsumerMergedNotTooLarge) {
   auto module = *ParseAndReturnVerifiedModule(R"(
     HloModule module
 
@@ -1007,7 +992,7 @@ TEST_P(PriorityFusionTest, FuseProducerConsumerMergedNotTooLarge) {
               absl_testing::IsOkAndHolds(true));
 }
 
-TEST_P(PriorityFusionTest, CanMergeTritonFusionWithBothProducerAndConsumer) {
+TEST_F(PriorityFusionTest, CanMergeTritonFusionWithBothProducerAndConsumer) {
   const std::string kHloText = R"(
 HloModule t
 add {
@@ -1064,7 +1049,7 @@ ENTRY main {
             2);
 }
 
-TEST_P(PriorityFusionTest, FuseTritonProducerWithTwoConsumers) {
+TEST_F(PriorityFusionTest, FuseTritonProducerWithTwoConsumers) {
   const std::string kHloText = R"(
 HloModule t
 add {
@@ -1128,13 +1113,11 @@ ENTRY main {
             2);
 }
 
-TEST_P(PriorityFusionTest,
+TEST_F(PriorityFusionTest,
        FuseTritonProducerWithTwoConsumersUsingMultiOutputFusion) {
-  if (GetParam()) {
-    // TODO(b/530092114): support multi-output fusions.
-    GTEST_SKIP() << "Multi-output fusions are not supported with tile-based "
-                    "block-level emitter";
-  }
+  // TODO(b/530092114): support multi-output fusions.
+  GTEST_SKIP() << "Multi-output fusions are not supported with tile-based "
+                  "block-level emitter";
   const std::string kHloText = R"(
 HloModule t
 
@@ -1187,12 +1170,10 @@ ENTRY main {
             2);
 }
 
-TEST_P(PriorityFusionTest,
+TEST_F(PriorityFusionTest,
        FuseProducerWithTritonConsumerUsingMultiOutputFusion) {
-  if (GetParam()) {
-    GTEST_SKIP() << "Multi-output fusions are not supported with tile-based "
-                    "block-level emitter";
-  }
+  GTEST_SKIP() << "Multi-output fusions are not supported with tile-based "
+                  "block-level emitter";
   const std::string kHloText = R"(
 HloModule t
 
@@ -1240,11 +1221,9 @@ ENTRY main {
             2);
 }
 
-TEST_P(PriorityFusionTest, FuseTritonFusionBothEndsUsingMultiOutputFusion) {
-  if (GetParam()) {
-    GTEST_SKIP() << "Multi-output fusions are not supported with tile-based "
-                    "block-level emitter";
-  }
+TEST_F(PriorityFusionTest, FuseTritonFusionBothEndsUsingMultiOutputFusion) {
+  GTEST_SKIP() << "Multi-output fusions are not supported with tile-based "
+                  "block-level emitter";
   // Here, we fuse `fusion` first into `exp` and `sqrt`. When we try to fuse
   // `log` into the two fusions resulting from the previous step using
   // multi-output fusion, we currently don't allow that, as we would need to
@@ -1286,7 +1265,7 @@ ENTRY main {
   EXPECT_TRUE(IsGenericTritonFusion(*fusion2));
 }
 
-TEST_P(PriorityFusionTest,
+TEST_F(PriorityFusionTest,
        ProducerWithOnlyBitcastUsersHasNoMultiOutputFusionCandidates) {
   // `negate` can only be fused into bitcasts, so there are no candidates for
   // Triton multi-output fusion either.
@@ -1307,7 +1286,7 @@ ENTRY main {
   EXPECT_FALSE(changed);
 }
 
-TEST_P(PriorityFusionTest, TritonProducerNotSupported_DoNotFuse) {
+TEST_F(PriorityFusionTest, TritonProducerNotSupported_DoNotFuse) {
   const std::string kHloText = R"(
 HloModule t
 
@@ -1336,7 +1315,7 @@ ENTRY main {
   EXPECT_FALSE(priority_fusion_.Run(module.get()).value());
 }
 
-TEST_P(PriorityFusionTest, TritonConsumerNotSupported_DoNotFuse) {
+TEST_F(PriorityFusionTest, TritonConsumerNotSupported_DoNotFuse) {
   const std::string kHloText = R"(
 HloModule t
 
@@ -1366,7 +1345,7 @@ ENTRY main {
   EXPECT_FALSE(priority_fusion_.Run(module.get()).value());
 }
 
-TEST_P(PriorityFusionTest, DoNotFuseInsideReducer) {
+TEST_F(PriorityFusionTest, DoNotFuseInsideReducer) {
   auto module = *ParseAndReturnVerifiedModule(R"(
     %reducer {
       p0 = f32[] parameter(0)
@@ -1391,7 +1370,7 @@ TEST_P(PriorityFusionTest, DoNotFuseInsideReducer) {
               absl_testing::IsOkAndHolds(false));
 }
 
-TEST_P(PriorityFusionTest, SkipsTilingsWithInfiniteRuntime) {
+TEST_F(PriorityFusionTest, SkipsTilingsWithInfiniteRuntime) {
   // This test verifies the fix in TryFindBestTilingForFusionAsync that skips
   // tilings with infinite runtime estimates.
   //
@@ -1518,14 +1497,7 @@ class HerolessPriorityFusionTest : public PriorityFusionTest {
   }
 };
 
-INSTANTIATE_TEST_SUITE_P(
-    HerolessPriorityFusionTest, HerolessPriorityFusionTest, ::testing::Bool(),
-    [](const ::testing::TestParamInfo<HerolessPriorityFusionTest::ParamType>&
-           info) {
-      return info.param ? "TilingPropagation" : "SymbolicAnalysis";
-    });
-
-TEST_P(HerolessPriorityFusionTest, TwoElementwiseOpsAreFusedWithTriton) {
+TEST_F(HerolessPriorityFusionTest, TwoElementwiseOpsAreFusedWithTriton) {
   auto module = *ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -1544,7 +1516,7 @@ ENTRY main {
   EXPECT_TRUE(IsGenericTritonFusion(*root));
 }
 
-TEST_P(HerolessPriorityFusionTest, DoNotFuseIntoRoot) {
+TEST_F(HerolessPriorityFusionTest, DoNotFuseIntoRoot) {
   auto module = *ParseAndReturnVerifiedModule(R"(
     HloModule test_module
 
@@ -1562,7 +1534,7 @@ TEST_P(HerolessPriorityFusionTest, DoNotFuseIntoRoot) {
               absl_testing::IsOkAndHolds(false));
 }
 
-TEST_P(HerolessPriorityFusionTest, LimitNumberOfParameters) {
+TEST_F(HerolessPriorityFusionTest, LimitNumberOfParameters) {
   std::string module_text =
       "HloModule m\n\nENTRY main {\nadd0 = f32[] parameter(0)\n";
   for (int64_t i = 1; i <= MaxOperandsAndOutputsPerFusion(); ++i) {
@@ -1580,11 +1552,9 @@ TEST_P(HerolessPriorityFusionTest, LimitNumberOfParameters) {
   EXPECT_LE(root->operand_count(), MaxOperandsAndOutputsPerFusion());
 }
 
-TEST_P(HerolessPriorityFusionTest, MultipleMultiOutputFusionCandidates) {
-  if (GetParam()) {
-    GTEST_SKIP() << "Multi-output fusions are not supported with tile-based "
-                    "block-level emitter";
-  }
+TEST_F(HerolessPriorityFusionTest, MultipleMultiOutputFusionCandidates) {
+  GTEST_SKIP() << "Multi-output fusions are not supported with tile-based "
+                  "block-level emitter";
   auto module = *ParseAndReturnVerifiedModule(R"(
     HloModule test_module
 
@@ -1618,7 +1588,7 @@ TEST_P(HerolessPriorityFusionTest, MultipleMultiOutputFusionCandidates) {
   EXPECT_TRUE(IsGenericTritonFusion(*fusion));
 }
 
-TEST_P(PriorityFusionTest, FusesQwixQuantization) {
+TEST_F(PriorityFusionTest, FusesQwixQuantization) {
   absl::string_view kHlo = R"(
 HloModule hlo_qwix_quantize_bf16_s8_2x256x512_tile128
 
@@ -1753,10 +1723,8 @@ TEST_F(PriorityFusionRocmMemoryBandwidthTest, MemoryBandwidthTipsReduceFusion) {
   EXPECT_EQ(RunAndCountFusions(kHlo, kFixedBandwidth), 2);
 }
 
-TEST_P(HerolessPriorityFusionTest, DoNotFuseScanEpilogue) {
-  bool experimental_tiling = GetParam();
-  std::string hlo = absl::StrFormat(
-      R"(
+TEST_F(HerolessPriorityFusionTest, DoNotFuseScanEpilogue) {
+  constexpr absl::string_view kHlo = R"(
 HloModule module
 
 add {
@@ -1778,19 +1746,18 @@ ENTRY entry {
   p1 = f32[] parameter(1)
   scan_fusion = f32[100] fusion(p0, p1), kind=kCustom, calls=fused_computation,
     backend_config={"fusion_backend_config":{"kind":"__triton","block_level_fusion_config":
-      {"output_tiles":[{"sizes":[%s]}],"num_warps":"1"}}}
+      {"output_tiles":[{"sizes":[]}],"num_warps":"1"}}}
   c = f32[] constant(1.0)
   bcast = f32[100] broadcast(c), dimensions={}
   ROOT add = f32[100] add(scan_fusion, bcast)
 }
-  )",
-      experimental_tiling ? "" : "100");
+  )";
 
   GpuHloCostAnalysis::Options options;
   options.count_multiple_input_accesses = true;
   PriorityFusion priority_fusion(nullptr, device_info_, &alias_info_, options,
                                  &mlir_context_);
-  RunAndFilecheckHloRewrite(hlo, std::move(priority_fusion), R"(
+  RunAndFilecheckHloRewrite(kHlo, std::move(priority_fusion), R"(
 CHECK: ENTRY
 CHECK: %[[SCAN_FUSION:.*]] = f32[100]{0} fusion(%{{.*}}, %{{.*}}), kind=kCustom
 CHECK: ROOT %[[EPILOGUE_FUSION:.*]] = f32[100]{0} fusion(%[[SCAN_FUSION]]), kind=kCustom
diff --git a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton_test.cc b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton_test.cc
index 33f0aebb45b..e5afe95833d 100644
--- a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton_test.cc
@@ -65,33 +65,21 @@ bool HasBlockLevelFusionConfig(const HloInstruction* fusion) {
              .has_block_level_fusion_config();
 }
 
-class SoftmaxRewriterTritonTest
-    : public HloHardwareIndependentTestBase,
-      // The parameter controls whether experimental tiling is enabled.
-      public ::testing::WithParamInterface<bool> {
+class SoftmaxRewriterTritonTest : public HloHardwareIndependentTestBase {
  protected:
   SoftmaxRewriterTritonTest() { RegisterSymbolicExprStorage(&mlir_context_); }
   se::DeviceDescription device_info_{TestGpuDeviceInfo::RTXA6000DeviceInfo()};
   mlir::MLIRContext mlir_context_;
   GpuAliasInfo alias_info_{device_info_};
-  SoftmaxRewriterTriton fusion_rewriter_{
-      device_info_,
-      HloCostAnalysis::DefaultShapeSize,
-      &alias_info_,
-      &mlir_context_,
-      /*only_fuse_if_profitable=*/false,
-      /*use_experimental_tiling=*/GetParam()};
-
-  DebugOptions GetDebugOptionsForTest() const override {
-    DebugOptions debug_options =
-        HloHardwareIndependentTestBase::GetDebugOptionsForTest();
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
-        GetParam());
-    return debug_options;
-  }
+  SoftmaxRewriterTriton fusion_rewriter_{device_info_,
+                                         HloCostAnalysis::DefaultShapeSize,
+                                         &alias_info_,
+                                         &mlir_context_,
+                                         /*only_fuse_if_profitable=*/false,
+                                         /*use_experimental_tiling=*/true};
 };
 
-TEST_P(SoftmaxRewriterTritonTest, CanFuseSingleNormalizationF32) {
+TEST_F(SoftmaxRewriterTritonTest, CanFuseSingleNormalizationF32) {
   const std::string hlo_string = R"(
 HloModule softmax
 max_computation {
@@ -123,7 +111,7 @@ ENTRY main {
           m::Fusion(m::Parameter()).WithPredicate(HasBlockLevelFusionConfig)));
 }
 
-TEST_P(SoftmaxRewriterTritonTest,
+TEST_F(SoftmaxRewriterTritonTest,
        CanFuseSingleNormalizationWithNonF32DataType) {
   const std::string hlo_string = R"(
 HloModule softmax
@@ -154,7 +142,7 @@ ENTRY main {
           m::Fusion(m::Parameter()).WithPredicate(HasBlockLevelFusionConfig)));
 }
 
-TEST_P(SoftmaxRewriterTritonTest, CanFuseSingleNormalizationDiamond) {
+TEST_F(SoftmaxRewriterTritonTest, CanFuseSingleNormalizationDiamond) {
   const std::string hlo_string = R"(
 HloModule softmax
 max_computation {
@@ -179,7 +167,7 @@ ENTRY main {
           m::Fusion(m::Parameter()).WithPredicate(HasBlockLevelFusionConfig)));
 }
 
-TEST_P(SoftmaxRewriterTritonTest,
+TEST_F(SoftmaxRewriterTritonTest,
        DoesNotFuseDiamondInvolvingUnsupportedTritonInstruction) {
   const std::string hlo_string = R"(
 HloModule softmax
@@ -204,7 +192,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(false));
 }
 
-TEST_P(SoftmaxRewriterTritonTest,
+TEST_F(SoftmaxRewriterTritonTest,
        DoesNotFuseInstructionsUnsupportedByTritonIntoDiamonds) {
   const std::string hlo_string = R"(
 HloModule softmax
@@ -236,7 +224,7 @@ ENTRY main {
           m::Fusion(m::Parameter()).WithPredicate(HasBlockLevelFusionConfig))));
 }
 
-TEST_P(SoftmaxRewriterTritonTest, CanNotFuseSoftmaxDiamondWithWrongLayout) {
+TEST_F(SoftmaxRewriterTritonTest, CanNotFuseSoftmaxDiamondWithWrongLayout) {
   const std::string hlo_string = R"(
 HloModule softmax
 max_computation {
@@ -256,7 +244,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(false));
 }
 
-TEST_P(SoftmaxRewriterTritonTest,
+TEST_F(SoftmaxRewriterTritonTest,
        CanNotFuseSoftmaxDiamondWithWrongReduceDimension) {
   const std::string hlo_string = R"(
 HloModule softmax
@@ -277,7 +265,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(false));
 }
 
-TEST_P(SoftmaxRewriterTritonTest,
+TEST_F(SoftmaxRewriterTritonTest,
        CanNotFuseSoftmaxDiamondWithWrongBroadcastDimension) {
   const std::string hlo_string = R"(
 HloModule softmax
@@ -298,7 +286,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(false));
 }
 
-TEST_P(SoftmaxRewriterTritonTest,
+TEST_F(SoftmaxRewriterTritonTest,
        CanNotFuseSoftmaxDiamondWithExtraBroadcastUsage) {
   const std::string hlo_string = R"(
 HloModule softmax
@@ -320,7 +308,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(false));
 }
 
-TEST_P(SoftmaxRewriterTritonTest, DoesNotFuseReductionOnNonMinorAxis) {
+TEST_F(SoftmaxRewriterTritonTest, DoesNotFuseReductionOnNonMinorAxis) {
   const std::string hlo_string = R"(
 max_computation {
   arg_0 = f32[] parameter(0)
@@ -339,7 +327,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(false));
 }
 
-TEST_P(SoftmaxRewriterTritonTest, DoesNotFuseReductionOnMultipleReductionAxes) {
+TEST_F(SoftmaxRewriterTritonTest, DoesNotFuseReductionOnMultipleReductionAxes) {
   const std::string hlo_string = R"(
 max_computation {
   arg_0 = f32[] parameter(0)
@@ -358,7 +346,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(false));
 }
 
-TEST_P(SoftmaxRewriterTritonTest, CanFuseDiamondWithUnaryElementwisePrefix) {
+TEST_F(SoftmaxRewriterTritonTest, CanFuseDiamondWithUnaryElementwisePrefix) {
   const std::string hlo_string = R"(
 HloModule softmax
 max_computation {
@@ -384,7 +372,7 @@ ENTRY main {
           m::Fusion(m::Parameter()).WithPredicate(HasBlockLevelFusionConfig)));
 }
 
-TEST_P(SoftmaxRewriterTritonTest,
+TEST_F(SoftmaxRewriterTritonTest,
        CanFuseDiamondWithMultipleBroadcastDimensions) {
   const std::string hlo_string = R"(
 HloModule softmax
@@ -411,7 +399,7 @@ ENTRY main {
           m::Fusion(m::Parameter()).WithPredicate(HasBlockLevelFusionConfig)));
 }
 
-TEST_P(SoftmaxRewriterTritonTest,
+TEST_F(SoftmaxRewriterTritonTest,
        CanNotFuseSoftmaxDiamondWithParameterReducerIdentity) {
   const std::string hlo_string = R"(
 HloModule softmax
@@ -434,7 +422,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(false));
 }
 
-TEST_P(SoftmaxRewriterTritonTest,
+TEST_F(SoftmaxRewriterTritonTest,
        CanNotFuseSoftmaxDiamondWithTritonIncompatibleReducer) {
   const std::string hlo_string = R"(
 HloModule softmax
@@ -458,7 +446,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(false));
 }
 
-TEST_P(SoftmaxRewriterTritonTest,
+TEST_F(SoftmaxRewriterTritonTest,
        CanFuseSoftmaxDiamondWithLastDimensionBitcastAfterReduce) {
   const std::string hlo_string = R"(
 HloModule softmax
@@ -488,7 +476,7 @@ ENTRY main {
           m::Fusion(m::Parameter()).WithPredicate(HasBlockLevelFusionConfig)));
 }
 
-TEST_P(SoftmaxRewriterTritonTest,
+TEST_F(SoftmaxRewriterTritonTest,
        CanNotFuseSoftmaxDiamondWithTransposeBitcast) {
   const std::string hlo_string = R"(
 HloModule softmax
@@ -512,7 +500,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(false));
 }
 
-TEST_P(SoftmaxRewriterTritonTest,
+TEST_F(SoftmaxRewriterTritonTest,
        CanNotFuseSoftmaxDiamondWithNonFusibleBitcastBetweenReduceAndProducer) {
   const std::string hlo_string = R"(
 HloModule softmax
@@ -537,7 +525,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(false));
 }
 
-TEST_P(SoftmaxRewriterTritonTest, CanFuseSoftmaxDiamondWithBitcastsOnEachUse) {
+TEST_F(SoftmaxRewriterTritonTest, CanFuseSoftmaxDiamondWithBitcastsOnEachUse) {
   const std::string hlo_string = R"(
 HloModule softmax
 
@@ -566,7 +554,7 @@ ENTRY main {
           m::Fusion(m::Parameter()).WithPredicate(HasBlockLevelFusionConfig)));
 }
 
-TEST_P(SoftmaxRewriterTritonTest, RewriterBailsOutOnPreAmpereCudaGpu) {
+TEST_F(SoftmaxRewriterTritonTest, RewriterBailsOutOnPreAmpereCudaGpu) {
   const std::string hlo_string = R"(
 HloModule softmax
 max_computation {
@@ -591,7 +579,7 @@ ENTRY main {
               se::CudaComputeCapability{se::CudaComputeCapability::kVolta, 0}),
           HloCostAnalysis::DefaultShapeSize, &alias_info_, &mlir_context_,
           /*only_fuse_if_profitable=*/false,
-          /*use_experimental_tiling=*/GetParam())
+          /*use_experimental_tiling=*/true)
           .Run(module.get()),
       absl_testing::StatusIs(
           tsl::error::FAILED_PRECONDITION,
@@ -599,7 +587,7 @@ ENTRY main {
                                "(compute capability 8.0) and up, but got")));
 }
 
-TEST_P(SoftmaxRewriterTritonTest, RewriterSucceedsOnNonCudaGpu) {
+TEST_F(SoftmaxRewriterTritonTest, RewriterSucceedsOnNonCudaGpu) {
   const std::string hlo_string = R"(
 HloModule softmax
 max_computation {
@@ -622,12 +610,12 @@ ENTRY main {
                                     HloCostAnalysis::DefaultShapeSize,
                                     &alias_info_, &mlir_context_,
                                     /*only_fuse_if_profitable=*/false,
-                                    /*use_experimental_tiling=*/GetParam())
+                                    /*use_experimental_tiling=*/true)
                   .Run(module.get())
                   .ok());
 }
 
-TEST_P(
+TEST_F(
     SoftmaxRewriterTritonTest,
     CanFuseIntermediateBinaryElementwiseWithinDiamondWhenBothOperandsAreTheSame) {  // NOLINT(whitespace/line_length)
   const std::string hlo_string = R"(
@@ -655,7 +643,7 @@ ENTRY main {
           m::Fusion(m::Parameter()).WithPredicate(HasBlockLevelFusionConfig)));
 }
 
-TEST_P(
+TEST_F(
     SoftmaxRewriterTritonTest,
     DoesNotFuseIntermediateBinaryElementwiseWithBothSplatOperandsIntoDiamond) {
   const std::string hlo_string = R"(
@@ -684,7 +672,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(false));
 }
 
-TEST_P(
+TEST_F(
     SoftmaxRewriterTritonTest,
     DoesNotFuseIntermediateBinaryElementwiseWithSameSplatOperandsIntoDiamond) {
   const std::string hlo_string = R"(
@@ -710,11 +698,11 @@ ENTRY main {
   SoftmaxRewriterTriton fusion_rewriter(
       device_info_, HloCostAnalysis::DefaultShapeSize, &alias_info_,
       &mlir_context_, /*only_fuse_if_profitable=*/false,
-      /*use_experimental_tiling=*/GetParam());
+      /*use_experimental_tiling=*/true);
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(false));
 }
 
-TEST_P(SoftmaxRewriterTritonTest, CanFuseRMSNormDiamond) {
+TEST_F(SoftmaxRewriterTritonTest, CanFuseRMSNormDiamond) {
   const std::string hlo_string = R"(
 HloModule rms_norm
 add_computation {
@@ -747,7 +735,7 @@ ENTRY main.30 {
           m::Fusion(m::Parameter()).WithPredicate(HasBlockLevelFusionConfig)));
 }
 
-TEST_P(
+TEST_F(
     SoftmaxRewriterTritonTest,
     CanFuseBinaryElementwiseWhereTheFirstOperandIsASplatConstantWithinDiamond) {
   const std::string hlo_string = R"(
@@ -777,7 +765,7 @@ ENTRY main {
           m::Fusion(m::Parameter()).WithPredicate(HasBlockLevelFusionConfig)));
 }
 
-TEST_P(SoftmaxRewriterTritonTest,
+TEST_F(SoftmaxRewriterTritonTest,
        CanFuseBinaryElementwiseOperationWhereOneOperandIsASharedSplatProducer) {
   const std::string hlo_string = R"(
 HloModule nonfusible_diamond
@@ -808,7 +796,7 @@ ENTRY main {
           m::Fusion(m::Parameter()).WithPredicate(HasBlockLevelFusionConfig)));
 }
 
-TEST_P(
+TEST_F(
     SoftmaxRewriterTritonTest,
     DoesNotFuseBinaryElementwiseOperationWhereFirstOperandIsASplatAndSecondOperandIsASharedSplatProducer) {  // NOLINT(whitespace/line_length)
   const std::string hlo_string = R"(
@@ -837,7 +825,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(false));
 }
 
-TEST_P(SoftmaxRewriterTritonTest, FusionDecisionIsCapturedExplicitly) {
+TEST_F(SoftmaxRewriterTritonTest, FusionDecisionIsCapturedExplicitly) {
   const std::string hlo_string = R"(
 HloModule softmax
 max_computation {
@@ -859,7 +847,7 @@ ENTRY main {
   SoftmaxRewriterTriton softmax_rewriter_triton(
       device_info_, HloCostAnalysis::DefaultShapeSize, &alias_info_,
       &mlir_context_, /*only_fuse_if_profitable=*/false,
-      /*use_experimental_tiling=*/GetParam());
+      /*use_experimental_tiling=*/true);
   int unmatched = 0, matched = 0;
   for (HloInstruction* instruction :
        module->entry_computation()->MakeInstructionPostOrder()) {
@@ -883,7 +871,7 @@ ENTRY main {
   EXPECT_EQ(matched, 0);
 }
 
-TEST_P(
+TEST_F(
     SoftmaxRewriterTritonTest,
     FusesBinaryElementwiseIfIntermediateDiamondOpWithBroadcastAlongReductionDimAsParameter) {  // NOLINT(whitespace/line_length)
   const std::string hlo_string = R"(
@@ -910,7 +898,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(true));
 }
 
-TEST_P(
+TEST_F(
     SoftmaxRewriterTritonTest,
     FusesBinaryElementwiseIfIntermediateDiamondOpWithBroadcastAlongBatchDimAsParameter) {  // NOLINT(whitespace/line_length)
   const std::string hlo_string = R"(
@@ -937,7 +925,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(true));
 }
 
-TEST_P(
+TEST_F(
     SoftmaxRewriterTritonTest,
     FusesBinaryElementwiseIfIntermediateDiamondOpWithMultiDimTensorBroadcastAlongBatchDimAsParameter) {  // NOLINT(whitespace/line_length)
   const std::string hlo_string = R"(
@@ -964,7 +952,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(true));
 }
 
-TEST_P(
+TEST_F(
     SoftmaxRewriterTritonTest,
     FusesBinaryElementwiseIfIntermediateDiamondOpWithZeroDimTensorBroadcastAsParameter) {  // NOLINT(whitespace/line_length)
   const std::string hlo_string = R"(
@@ -991,7 +979,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(true));
 }
 
-TEST_P(
+TEST_F(
     SoftmaxRewriterTritonTest,
     FusesBinaryElementwiseIfIntermediateDiamondOpIsBroadcastOf1DParameterAlongNonReductionDimensions) {  // NOLINT(whitespace/line_length)
   const std::string hlo_string = R"(
@@ -1018,7 +1006,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(true));
 }
 
-TEST_P(SoftmaxRewriterTritonTest,
+TEST_F(SoftmaxRewriterTritonTest,
        FusesBinaryElementwiseIfIntermediateDiamondOpIsBroadcastOfParameter) {
   const std::string hlo_string = R"(
 HloModule h1
@@ -1044,7 +1032,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(true));
 }
 
-TEST_P(
+TEST_F(
     SoftmaxRewriterTritonTest,
     FusesBinaryElementwiseIfIntermediateDiamondOpWithMultipleDimensionsAsParameter) {  // NOLINT(whitespace/line_length)
   const std::string hlo_string = R"(
@@ -1073,7 +1061,7 @@ ENTRY main {
 
 // Triton has a requirement that any tile in the program should not have more
 // than 1048576 elements.
-TEST_P(SoftmaxRewriterTritonTest, DoesNotFuseIfResultingFusionCannotBeTiled) {
+TEST_F(SoftmaxRewriterTritonTest, DoesNotFuseIfResultingFusionCannotBeTiled) {
   const std::string hlo_string = R"(
 HloModule softmax
 max_computation {
@@ -1093,7 +1081,7 @@ ENTRY main {
   EXPECT_THAT(fusion_rewriter_.Run(module.get()), IsOkAndHolds(false));
 }
 
-TEST_P(SoftmaxRewriterTritonTest, DoesNotFuseNormalizationWithVeryLongRows) {
+TEST_F(SoftmaxRewriterTritonTest, DoesNotFuseNormalizationWithVeryLongRows) {
   const std::string hlo_string = R"(
 HloModule softmax
 max_computation {
@@ -1119,7 +1107,7 @@ ENTRY main {
         &alias_info_,
         &mlir_context_,
         /*only_fuse_if_profitable=*/false,
-        /*use_experimental_tiling=*/GetParam()};
+        /*use_experimental_tiling=*/true};
 
     ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
     EXPECT_THAT(fusion_rewriter_without_cost_model.Run(module.get()),
@@ -1135,7 +1123,7 @@ ENTRY main {
         &alias_info_,
         &mlir_context_,
         /*only_fuse_if_profitable=*/true,
-        /*use_experimental_tiling=*/GetParam()};
+        /*use_experimental_tiling=*/true};
 
     ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
     EXPECT_THAT(fusion_rewriter_with_cost_model.Run(module.get()),
@@ -1143,7 +1131,7 @@ ENTRY main {
   }
 }
 
-TEST_P(SoftmaxRewriterTritonTest, DoesNotCrashOnScalarBroadcast) {
+TEST_F(SoftmaxRewriterTritonTest, DoesNotCrashOnScalarBroadcast) {
   const std::string hlo_string = R"(
 HloModule softmax
 max_computation {
@@ -1169,7 +1157,7 @@ ENTRY main {
                              .WithPredicate(HasBlockLevelFusionConfig)));
 }
 
-TEST_P(SoftmaxRewriterTritonTest, CanFuseWithParallelTilingSearch) {
+TEST_F(SoftmaxRewriterTritonTest, CanFuseWithParallelTilingSearch) {
   const std::string hlo_string = R"(
 HloModule softmax
 max_computation {
@@ -1198,7 +1186,7 @@ ENTRY main {
       device_info_, HloCostAnalysis::DefaultShapeSize, &alias_info_,
       &mlir_context_,
       /*only_fuse_if_profitable=*/false,
-      /*use_experimental_tiling=*/GetParam(), &thread_pool, &mlir_context_pool);
+      /*use_experimental_tiling=*/true, &thread_pool, &mlir_context_pool);
   EXPECT_THAT(rewriter.Run(module.get()), IsOkAndHolds(true));
   EXPECT_TRUE(verifier().Run(module.get()).status().ok());
   EXPECT_THAT(
@@ -1207,13 +1195,6 @@ ENTRY main {
           m::Fusion(m::Parameter()).WithPredicate(HasBlockLevelFusionConfig)));
 }
 
-INSTANTIATE_TEST_SUITE_P(SoftmaxRewriterTritonTestSuite,
-                         SoftmaxRewriterTritonTest, ::testing::Bool(),
-                         [](const ::testing::TestParamInfo<bool>& info) {
-                           return info.param ? "ExperimentalTiling"
-                                             : "SymbolicTiling";
-                         });
-
 }  // anonymous namespace
 }  // namespace gpu
 }  // namespace xla
diff --git a/third_party/xla/xla/codegen/xtile/BUILD b/third_party/xla/xla/codegen/xtile/BUILD
index 9ceb6138789..e2c3c7b34cd 100644
--- a/third_party/xla/xla/codegen/xtile/BUILD
+++ b/third_party/xla/xla/codegen/xtile/BUILD
@@ -75,8 +75,6 @@ xla_cc_test(
         ":xtile_config_proto_cc",
         "//xla:status_macros",
         "//xla:xla_proto_cc",
-        "//xla/codegen/tiling:symbolic_tile_analysis",
-        "//xla/codegen/tiling:tiling_specification",
         "//xla/codegen/tiling/experimental:tile",
         "//xla/hlo/analysis:symbolic_map",
         "//xla/hlo/ir:hlo",
diff --git a/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters_test.cc b/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters_test.cc
index cf6c8edb5ca..ff6fab791a5 100644
--- a/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters_test.cc
+++ b/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters_test.cc
@@ -30,8 +30,6 @@ limitations under the License.
 #include "llvm/ADT/SmallVector.h"
 #include "mlir/IR/MLIRContext.h"
 #include "xla/codegen/tiling/experimental/tiling_space.h"
-#include "xla/codegen/tiling/symbolic_tile_analysis.h"
-#include "xla/codegen/tiling/tiling_specification.h"
 #include "xla/codegen/xtile/block_level_parameters.h"
 #include "xla/codegen/xtile/xtile_config.pb.h"
 #include "xla/hlo/analysis/symbolic_expr.h"
@@ -54,149 +52,6 @@ using ::absl_testing::IsOkAndHolds;
 using ::absl_testing::StatusIs;
 using ::testing::ElementsAre;
 
-class TilingFromBlockParametersTest : public HloHardwareIndependentTestBase {
- public:
-  TilingFromBlockParametersTest() {
-    RegisterSymbolicExprStorage(&mlir_context_);
-  }
-
-  std::optional<SymbolicTileAnalysis> TryAnalyzeModule(HloModule* module) {
-    SymbolicTileAnalysisOrError analysis_or_error =
-        SymbolicTileAnalysis::AnalyzeComputation(
-            *module->entry_computation()
-                 ->root_instruction()
-                 ->fused_instructions_computation(),
-            &mlir_context_);
-
-    if (std::holds_alternative<SymbolicTileAnalysis>(analysis_or_error)) {
-      return std::get<SymbolicTileAnalysis>(std::move(analysis_or_error));
-    }
-    return std::nullopt;
-  }
-
-  DebugOptions GetDebugOptionsForTest() const override {
-    DebugOptions debug_options =
-        HloHardwareIndependentTestBase::GetDebugOptionsForTest();
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(false);
-    return debug_options;
-  }
-
-  mlir::MLIRContext mlir_context_;
-};
-
-TEST_F(TilingFromBlockParametersTest, GeneratesTilingForSimpleMap) {
-  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                       ParseAndReturnVerifiedModule(R"hlo(
-HloModule m
-
-fused_computation {
-  param_0 = f32[128,256] parameter(0)
-  ROOT log = f32[128,256] log(param_0)
-}
-
-ENTRY entry_computation {
-  param_0 = f32[128,256] parameter(0)
-  ROOT fusion = f32[128,256] fusion(param_0), kind=kCustom,
-   calls=fused_computation
-}
-)hlo"));
-
-  std::optional<SymbolicTileAnalysis> analysis = TryAnalyzeModule(module.get());
-  ASSERT_TRUE(analysis.has_value());
-
-  BlockLevelParameters block_level_parameters;
-  block_level_parameters.output_tile_sizes = {{16, 32}};
-
-  ASSERT_OK_AND_ASSIGN(Tiling tiling, TilingFromAnnotatedFusion(
-                                          *analysis, block_level_parameters));
-
-  const HloInstruction* log = module->entry_computation()
-                                  ->root_instruction()
-                                  ->fused_instructions_computation()
-                                  ->root_instruction();
-
-  EXPECT_THAT(tiling.TileSizesForInstruction(log),
-              IsOkAndHolds(ElementsAre(16, 32)));
-}
-
-TEST_F(TilingFromBlockParametersTest, GeneratesTilingForDot) {
-  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                       ParseAndReturnVerifiedModule(R"hlo(
-HloModule m
-
-fused_computation {
-  p0 = f32[128,128] parameter(0)
-  p1 = f32[128,128] parameter(1)
-  ROOT dot = f32[128,128] dot(p0, p1),
-   lhs_contracting_dims={1}, rhs_contracting_dims={0},
-   backend_config={"sizes":["32"]}
-}
-
-ENTRY entry_computation {
-  param_0 = f32[128,128] parameter(0)
-  param_1 = f32[128,128] parameter(1)
-  ROOT fusion = f32[128,128] fusion(param_0, param_1), kind=kCustom, calls=fused_computation
-}
-)hlo"));
-
-  std::optional<SymbolicTileAnalysis> analysis = TryAnalyzeModule(module.get());
-  ASSERT_TRUE(analysis.has_value());
-
-  BlockLevelParameters block_level_parameters;
-  block_level_parameters.output_tile_sizes = {{16, 16}};
-
-  ASSERT_OK_AND_ASSIGN(Tiling tiling, TilingFromAnnotatedFusion(
-                                          *analysis, block_level_parameters));
-
-  const HloInstruction* dot = module->entry_computation()
-                                  ->root_instruction()
-                                  ->fused_instructions_computation()
-                                  ->root_instruction();
-
-  EXPECT_THAT(tiling.TileSizesForInstruction(dot),
-              IsOkAndHolds(ElementsAre(32, 16, 16)));
-}
-
-TEST_F(TilingFromBlockParametersTest, GeneratesTilingForDotWithTilingOverride) {
-  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
-                       ParseAndReturnVerifiedModule(R"hlo(
-HloModule m
-
-fused_computation {
-  p0 = f32[128,128] parameter(0)
-  p1 = f32[128,128] parameter(1)
-  ROOT dot = f32[128,128] dot(p0, p1),
-   lhs_contracting_dims={1}, rhs_contracting_dims={0}}
-
-ENTRY entry_computation {
-  param_0 = f32[128,128] parameter(0)
-  param_1 = f32[128,128] parameter(1)
-  ROOT fusion = f32[128,128] fusion(param_0, param_1), kind=kCustom, calls=fused_computation
-}
-)hlo"));
-
-  std::optional<SymbolicTileAnalysis> analysis = TryAnalyzeModule(module.get());
-  ASSERT_TRUE(analysis.has_value());
-
-  BlockLevelParameters block_level_parameters;
-  block_level_parameters.output_tile_sizes = {{16, 16}};
-
-  xla::xtile::Tile tile_override;
-  tile_override.add_sizes(64);
-
-  ASSERT_OK_AND_ASSIGN(Tiling tiling,
-                       TilingFromAnnotatedFusion(
-                           *analysis, block_level_parameters, &tile_override));
-
-  const HloInstruction* dot = module->entry_computation()
-                                  ->root_instruction()
-                                  ->fused_instructions_computation()
-                                  ->root_instruction();
-
-  EXPECT_THAT(tiling.TileSizesForInstruction(dot),
-              IsOkAndHolds(ElementsAre(64, 16, 16)));
-}
-
 class GetTileTilingSpaceConcreteSizesTest
     : public HloHardwareIndependentTestBase {
  public:
@@ -222,13 +77,6 @@ class GetTileTilingSpaceConcreteSizesTest
             .xla_gpu_experimental_enable_same_shape_multi_output_fusion());
   }
 
-  DebugOptions GetDebugOptionsForTest() const override {
-    DebugOptions debug_options =
-        HloHardwareIndependentTestBase::GetDebugOptionsForTest();
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(false);
-    return debug_options;
-  }
-
   mlir::MLIRContext mlir_context_;
 };
 
diff --git a/third_party/xla/xla/service/gpu/model/BUILD b/third_party/xla/xla/service/gpu/model/BUILD
index 0728f0a7069..cce38ba2862 100644
--- a/third_party/xla/xla/service/gpu/model/BUILD
+++ b/third_party/xla/xla/service/gpu/model/BUILD
@@ -629,9 +629,6 @@ xla_cc_test(
         ":gpu_performance_model_base",
         "//xla:shape_util",
         "//xla:util",
-        "//xla/codegen/tiling:symbolic_tile_analysis",
-        "//xla/codegen/tiling:tiled_hlo_computation",
-        "//xla/codegen/tiling:tiling_specification",
         "//xla/codegen/tiling/experimental:tile",
         "//xla/codegen/tiling/experimental:tiled_hlo",
         "//xla/codegen/tiling/experimental:tiling_space_utils",
@@ -760,10 +757,6 @@ xla_cc_test(
         "//xla:xla_data_proto_cc",
         "//xla/backends/gpu/codegen:fusion_emitter",
         "//xla/backends/gpu/codegen:fusions",
-        "//xla/codegen/tiling:symbolic_tile_analysis",
-        "//xla/codegen/tiling:tiled_hlo_computation",
-        "//xla/codegen/tiling:tiled_hlo_schedule",
-        "//xla/codegen/tiling:tiling_specification",
         "//xla/codegen/tiling/experimental:tile",
         "//xla/codegen/tiling/experimental:tiled_hlo",
         "//xla/hlo/analysis:symbolic_map",
diff --git a/third_party/xla/xla/service/gpu/model/coalescing_analysis_test.cc b/third_party/xla/xla/service/gpu/model/coalescing_analysis_test.cc
index e679c871175..27da5598ce3 100644
--- a/third_party/xla/xla/service/gpu/model/coalescing_analysis_test.cc
+++ b/third_party/xla/xla/service/gpu/model/coalescing_analysis_test.cc
@@ -30,10 +30,6 @@ limitations under the License.
 #include "xla/backends/gpu/codegen/fusions.h"
 #include "xla/codegen/tiling/experimental/tiled_hlo.h"
 #include "xla/codegen/tiling/experimental/tiling_space.h"
-#include "xla/codegen/tiling/symbolic_tile_analysis.h"
-#include "xla/codegen/tiling/tiled_hlo_computation.h"
-#include "xla/codegen/tiling/tiled_hlo_schedule.h"
-#include "xla/codegen/tiling/tiling_specification.h"
 #include "xla/hlo/analysis/symbolic_expr.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_opcode.h"
@@ -586,71 +582,34 @@ TEST_F(CoalescingTest, Param) {
   EXPECT_THAT(IsReadCoalescedPerOperand(ir), ElementsAre(true, true, true));
 }
 
-class CoalescingForTiledHloTest : public CoalescingTest,
-                                  public ::testing::WithParamInterface<bool> {
+class CoalescingForTiledHloTest : public CoalescingTest {
  public:
-  bool use_experimental_tiling() const { return GetParam(); }
+  std::vector<double> EffectiveBandwidthUtilizationRatePerOperand(
+      const HloInstruction* root, absl::Span<int64_t const> tile_sizes) {
+    auto fusion_adaptor = HloFusionAdaptor::ForInstruction(root);
+
+    auto tiling_space_or =
+        experimental::TilingSpace::Create(*fusion_adaptor, &mlir_context_);
+    CHECK_OK(tiling_space_or);
+    auto tiling_space = std::move(tiling_space_or.value());
+    CHECK_OK(tiling_space->AssignTileSizes(tile_sizes));
+
+    absl::StatusOr<experimental::TiledHloComputation> tiled_hlo_computation =
+        experimental::TiledHloComputation::Tile(*fusion_adaptor,
+                                                std::move(tiling_space));
+    CHECK_OK(tiled_hlo_computation);
 
-  template <typename TiledHloInstructionType>
-  std::vector<double> EffectiveBandwidthUtilizationRatePerOperandImpl(
-      const TiledHloInstructionType* tiled_hlo_root) {
     std::vector<double> result;
-    for (const TiledHloInstructionType* operand : tiled_hlo_root->operands()) {
+    for (const experimental::TiledHloInstruction* operand :
+         tiled_hlo_computation->roots()[0]->operands()) {
       result.push_back(BandwidthUtilizationRateHeuristicForTiledMemoryAccess(
           *operand, device_info_));
     }
     return result;
   }
-
-  std::vector<double> EffectiveBandwidthUtilizationRatePerOperand(
-      const HloInstruction* root, absl::Span<int64_t const> tile_sizes) {
-    auto fusion_adaptor = HloFusionAdaptor::ForInstruction(root);
-
-    if (use_experimental_tiling()) {
-      auto tiling_space_or =
-          experimental::TilingSpace::Create(*fusion_adaptor, &mlir_context_);
-      CHECK_OK(tiling_space_or);
-      auto tiling_space = std::move(tiling_space_or.value());
-      CHECK_OK(tiling_space->AssignTileSizes(tile_sizes));
-
-      absl::StatusOr<experimental::TiledHloComputation> tiled_hlo_computation =
-          experimental::TiledHloComputation::Tile(*fusion_adaptor,
-                                                  std::move(tiling_space));
-      CHECK_OK(tiled_hlo_computation);
-
-      return EffectiveBandwidthUtilizationRatePerOperandImpl(
-          tiled_hlo_computation->roots()[0]);
-    }
-
-    SymbolicTileAnalysis symbolic_tile_analysis =
-        std::get<SymbolicTileAnalysis>(SymbolicTileAnalysis::AnalyzeFusion(
-            *fusion_adaptor, &mlir_context_));
-
-    absl::StatusOr<TiledHloComputation> tiled_hlo_computation =
-        symbolic_tile_analysis.ComputeTiledComputation(
-            Tiling({{root, FlatTiling(tile_sizes.begin(), tile_sizes.end())}}),
-            CreateMajorToMinorTiledHloSchedule,
-            /*constraints_are_known_satisfied=*/true,
-            /*compute_all_tile_offset_indexing_maps=*/true);
-    CHECK_OK(tiled_hlo_computation);
-
-    return EffectiveBandwidthUtilizationRatePerOperandImpl(
-        tiled_hlo_computation->roots()[0]);
-  }
-
-  DebugOptions GetDebugOptionsForTest() const override {
-    DebugOptions debug_options =
-        HloHardwareIndependentTestBase::GetDebugOptionsForTest();
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
-        use_experimental_tiling());
-    return debug_options;
-  }
 };
 
-INSTANTIATE_TEST_SUITE_P(CoalescingForTiledHloTest, CoalescingForTiledHloTest,
-                         testing::Bool());
-
-TEST_P(
+TEST_F(
     CoalescingForTiledHloTest,
     EffectiveBandwidthUtilizationRateIsComputedCorrectlyForTiledMemoryAccess) {  // NOLINT(whitespace/line_length)
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
diff --git a/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model_test.cc b/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model_test.cc
index 7512db23c5e..9578a5f6d38 100644
--- a/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model_test.cc
+++ b/third_party/xla/xla/service/gpu/model/combined_gpu_performance_model_test.cc
@@ -39,17 +39,13 @@ namespace xla {
 namespace gpu {
 namespace {
 
-class CombinedGpuPerformanceModelTest
-    : public HloHardwareIndependentTestBase,
-      public ::testing::WithParamInterface</*use_experimental_tiling=*/bool> {
+class CombinedGpuPerformanceModelTest : public HloHardwareIndependentTestBase {
  public:
   CombinedGpuPerformanceModelTest() : analysis_(options_, device_info_) {
     options_.count_multiple_input_accesses = true;
     RegisterSymbolicExprStorage(&mlir_context_);
   }
 
-  bool use_experimental_tiling() const { return GetParam(); }
-
   mlir::MLIRContext mlir_context_;
   GpuHloCostAnalysis::Options options_;
   se::DeviceDescription device_info_{TestGpuDeviceInfo::RTXA6000DeviceInfo()};
@@ -60,11 +56,11 @@ class CombinedGpuPerformanceModelTest
       fusion_analysis_cache_,
       mlir_context_,
       [](const Shape& shape) { return ShapeUtil::ByteSizeOf(shape); },
-      use_experimental_tiling(),
+      /*use_experimental_tiling=*/true,
       /*enable_same_shape_multi_output_fusion=*/false};
 };
 
-TEST_P(CombinedGpuPerformanceModelTest,
+TEST_F(CombinedGpuPerformanceModelTest,
        ReturnsGpuPerformanceModelResultForNonTritonFusion) {
   auto module = ParseAndReturnVerifiedModule(R"(
     HloModule test_module
@@ -95,7 +91,7 @@ TEST_P(CombinedGpuPerformanceModelTest,
 }
 
 // TODO: b/493907020 Remove this after removing from GpuPerformanceModel.
-TEST_P(CombinedGpuPerformanceModelTest,
+TEST_F(CombinedGpuPerformanceModelTest,
        EstimateRunTimesMatchesGpuPerformanceModel) {
   auto module = ParseAndReturnVerifiedModule(R"(
     HloModule test_module
@@ -126,7 +122,7 @@ TEST_P(CombinedGpuPerformanceModelTest,
 }
 
 // TODO: b/493907020 Remove this after removing from GpuPerformanceModel.
-TEST_P(CombinedGpuPerformanceModelTest,
+TEST_F(CombinedGpuPerformanceModelTest,
        EstimateRunTimesForMultiOutputMatchesGpuPerformanceModel) {
   auto module = ParseAndReturnVerifiedModule(R"(
     HloModule test_module
@@ -157,7 +153,7 @@ TEST_P(CombinedGpuPerformanceModelTest,
   EXPECT_EQ(result->time_fused, expected.time_fused);
 }
 
-TEST_P(CombinedGpuPerformanceModelTest, CachesResults) {
+TEST_F(CombinedGpuPerformanceModelTest, CachesResults) {
   auto module = ParseAndReturnVerifiedModule(R"(
     HloModule test_module
     ENTRY entry_computation {
@@ -178,7 +174,7 @@ TEST_P(CombinedGpuPerformanceModelTest, CachesResults) {
   EXPECT_EQ(model_.GetCache().Get(*add)->exec_time, result->exec_time);
 }
 
-TEST_P(CombinedGpuPerformanceModelTest, InvalidatesCache) {
+TEST_F(CombinedGpuPerformanceModelTest, InvalidatesCache) {
   auto module = ParseAndReturnVerifiedModule(R"(
     HloModule test_module
     ENTRY entry_computation {
@@ -198,13 +194,6 @@ TEST_P(CombinedGpuPerformanceModelTest, InvalidatesCache) {
   EXPECT_FALSE(model_.GetCache().Get(*add).has_value());
 }
 
-INSTANTIATE_TEST_SUITE_P(CombinedGpuPerformanceModelTestSuite,
-                         CombinedGpuPerformanceModelTest, ::testing::Bool(),
-                         [](const ::testing::TestParamInfo<bool>& info) {
-                           return info.param ? "ExperimentalTiling"
-                                             : "SymbolicTiling";
-                         });
-
 }  // namespace
 }  // namespace gpu
 }  // namespace xla
diff --git a/third_party/xla/xla/service/gpu/model/gpu_cost_model_stats_collection_test.cc b/third_party/xla/xla/service/gpu/model/gpu_cost_model_stats_collection_test.cc
index 82d3217b949..4516104b810 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_cost_model_stats_collection_test.cc
+++ b/third_party/xla/xla/service/gpu/model/gpu_cost_model_stats_collection_test.cc
@@ -40,35 +40,23 @@ using ::mlir::MLIRContext;
 using ::testing::Contains;
 using ::testing::Truly;
 
-class GpuCostModelStatsCollectionTest
-    : public HloHardwareIndependentTestBase,
-      public ::testing::WithParamInterface</*use_experimental_tiling=*/bool> {
+class GpuCostModelStatsCollectionTest : public HloHardwareIndependentTestBase {
  public:
   GpuCostModelStatsCollectionTest() {
     RegisterSymbolicExprStorage(&mlir_context_);
   }
 
-  bool use_experimental_tiling() const { return GetParam(); }
-
  protected:
-  DebugOptions GetDebugOptionsForTest() const override {
-    DebugOptions debug_options =
-        HloHardwareIndependentTestBase::GetDebugOptionsForTest();
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
-        use_experimental_tiling());
-    return debug_options;
-  }
-
   // Must be declared before `cost_model_stats_`, which uses it on construction.
   mlir::MLIRContext mlir_context_;
   GpuCostModelStatsCollection cost_model_stats_{
       TestGpuDeviceInfo::H100SXMDeviceInfo(),
       GpuHloCostAnalysis::Options{.count_multiple_input_accesses = true},
-      &mlir_context_, use_experimental_tiling(),
+      &mlir_context_, /*use_experimental_tiling=*/true,
       /*enable_same_shape_multi_output_fusion=*/false};
 };
 
-TEST_P(GpuCostModelStatsCollectionTest, FusionInEntryComputation) {
+TEST_F(GpuCostModelStatsCollectionTest, FusionInEntryComputation) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"hlo(
     HloModule test_module
 
@@ -93,7 +81,7 @@ TEST_P(GpuCostModelStatsCollectionTest, FusionInEntryComputation) {
   EXPECT_GT(gpu_config.reification_cost()[0].end_to_end_cycles(), 0);
 }
 
-TEST_P(GpuCostModelStatsCollectionTest, FusionInWhileComputation) {
+TEST_F(GpuCostModelStatsCollectionTest, FusionInWhileComputation) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"hlo(
     HloModule test_module
 
@@ -130,7 +118,7 @@ TEST_P(GpuCostModelStatsCollectionTest, FusionInWhileComputation) {
   EXPECT_GT(gpu_config.reification_cost()[0].end_to_end_cycles(), 0);
 }
 
-TEST_P(GpuCostModelStatsCollectionTest, GemmCostModelAddedToGemmFusion) {
+TEST_F(GpuCostModelStatsCollectionTest, GemmCostModelAddedToGemmFusion) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"hlo(
   HloModule test_module
 
@@ -174,13 +162,6 @@ TEST_P(GpuCostModelStatsCollectionTest, GemmCostModelAddedToGemmFusion) {
               })));
 }
 
-INSTANTIATE_TEST_SUITE_P(GpuCostModelStatsCollectionTestSuite,
-                         GpuCostModelStatsCollectionTest, ::testing::Bool(),
-                         [](const ::testing::TestParamInfo<bool>& info) {
-                           return info.param ? "ExperimentalTiling"
-                                             : "SymbolicTiling";
-                         });
-
 }  // namespace
 }  // namespace gpu
 }  // namespace xla
diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc
index 2936840ad49..d688a5f05eb 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc
+++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc
@@ -31,9 +31,6 @@ limitations under the License.
 #include "xla/codegen/tiling/experimental/tiled_hlo.h"
 #include "xla/codegen/tiling/experimental/tiling_space.h"
 #include "xla/codegen/tiling/experimental/tiling_space_utils.h"
-#include "xla/codegen/tiling/symbolic_tile_analysis.h"
-#include "xla/codegen/tiling/tiled_hlo_computation.h"
-#include "xla/codegen/tiling/tiling_specification.h"
 #include "xla/codegen/xtile/block_level_parameters.h"
 #include "xla/codegen/xtile/codegen/emitter_helpers.h"
 #include "xla/hlo/analysis/symbolic_expr.h"
@@ -67,16 +64,12 @@ namespace {
 using ::testing::ElementsAre;
 using ::xla::xtile::BlockLevelParameters;
 
-class GpuIndexingPerformanceModelTest
-    : public HloHardwareIndependentTestBase,
-      public ::testing::WithParamInterface<bool> {
+class GpuIndexingPerformanceModelTest : public HloHardwareIndependentTestBase {
  public:
   GpuIndexingPerformanceModelTest() {
     RegisterSymbolicExprStorage(&mlir_context_);
   }
 
-  bool use_experimental_tiling() const { return GetParam(); }
-
   static std::unique_ptr<mlir::MLIRContext> CreateMlirContext() {
     auto ctx = std::make_unique<mlir::MLIRContext>();
     ctx->disableMultithreading();
@@ -94,29 +87,14 @@ class GpuIndexingPerformanceModelTest
       &fusion_analysis_cache_,
       HloCostAnalysis::DefaultShapeSize,
       &mlir_context_,
-      use_experimental_tiling(),
+      /*use_experimental_tiling=*/true,
       /*enable_same_shape_multi_output_fusion=*/false,
       &mlir_context_pool_};
 
   size_t WarpSize() const { return ::xla::gpu::WarpSize(device_info_); }
-
-  DebugOptions GetDebugOptionsForTest() const override {
-    DebugOptions debug_options =
-        HloHardwareIndependentTestBase::GetDebugOptionsForTest();
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(
-        use_experimental_tiling());
-    return debug_options;
-  }
 };
 
-INSTANTIATE_TEST_SUITE_P(GpuIndexingPerformanceModelTest,
-                         GpuIndexingPerformanceModelTest, ::testing::Bool(),
-                         [](const ::testing::TestParamInfo<bool>& info) {
-                           return info.param ? "ExperimentalTiling"
-                                             : "SymbolicTiling";
-                         });
-
-TEST_P(GpuIndexingPerformanceModelTest,
+TEST_F(GpuIndexingPerformanceModelTest,
        TritonGemmWithCustomBlockLevelParameters) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
@@ -169,7 +147,7 @@ ENTRY e {
   EXPECT_NE(result_with_custom.exec_time, result_default.exec_time);
 }
 
-TEST_P(GpuIndexingPerformanceModelTest, TritonGemmComputeBoundBf16NumStages1) {
+TEST_F(GpuIndexingPerformanceModelTest, TritonGemmComputeBoundBf16NumStages1) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -212,7 +190,7 @@ ENTRY e {
   EXPECT_GT(runtime.exec_time, expected_compute_and_flops.compute_time * 1.2);
 }
 
-TEST_P(GpuIndexingPerformanceModelTest, TritonGemmComputeBoundBf16) {
+TEST_F(GpuIndexingPerformanceModelTest, TritonGemmComputeBoundBf16) {
   // TODO: b/510666436 - Tile sizes are intentionally kept large to reduce
   // L2 cache replication overhead modeled by threadblock_count, keeping
   // the operation compute bound.
@@ -262,7 +240,7 @@ ENTRY e {
   EXPECT_LE(runtime.exec_time, expected_time * 1.1);
 }
 
-TEST_P(GpuIndexingPerformanceModelTest, TritonGemmMemoryBoundBf16) {
+TEST_F(GpuIndexingPerformanceModelTest, TritonGemmMemoryBoundBf16) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -304,7 +282,7 @@ ENTRY e {
   EXPECT_LE(runtime.exec_time, approx_hbm_time * 1.1);
 }
 
-TEST_P(GpuIndexingPerformanceModelTest,
+TEST_F(GpuIndexingPerformanceModelTest,
        TritonSoftmaxFusionInstructionIsSupported) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
@@ -351,13 +329,11 @@ ENTRY main {
 }
 
 // Example from b/383162692.
-TEST_P(GpuIndexingPerformanceModelTest, EstimateBestTiling_CombinedFusion) {
-  if (use_experimental_tiling()) {
-    // TODO: b/422689305 - Enable this test once multi-output fusions are
-    // supported.
-    GTEST_SKIP()
-        << "Experimental tiling does not support multi-output fusions yet.";
-  }
+TEST_F(GpuIndexingPerformanceModelTest, EstimateBestTiling_CombinedFusion) {
+  // TODO: b/422689305 - Enable this test once multi-output fusions are
+  // supported.
+  GTEST_SKIP()
+      << "Experimental tiling does not support multi-output fusions yet.";
 
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
@@ -470,13 +446,11 @@ ENTRY entry_computation {
   // EXPECT_EQ(tiled_runtime_data.block_level_parameters.num_warps, 32);
 }
 
-TEST_P(GpuIndexingPerformanceModelTest, EstimateBestTiling_MultioutputFusion) {
-  if (use_experimental_tiling()) {
-    // TODO: b/422689305 - Enable this test once multi-output fusions are
-    // supported.
-    GTEST_SKIP()
-        << "Experimental tiling does not support multi-output fusions yet.";
-  }
+TEST_F(GpuIndexingPerformanceModelTest, EstimateBestTiling_MultioutputFusion) {
+  // TODO: b/422689305 - Enable this test once multi-output fusions are
+  // supported.
+  GTEST_SKIP()
+      << "Experimental tiling does not support multi-output fusions yet.";
 
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
@@ -521,7 +495,7 @@ ENTRY entry_computation {
   EXPECT_EQ(tiled_runtime_data.block_level_parameters.num_warps, 1);
 }
 
-TEST_P(GpuIndexingPerformanceModelTest,
+TEST_F(GpuIndexingPerformanceModelTest,
        TryFindTopKBestTilingsForFusion_DumpsToFile) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
@@ -564,7 +538,7 @@ ENTRY main {
   EXPECT_THAT(dumped_content, ::testing::HasSubstr("Candidate #0:"));
 }
 
-TEST_P(GpuIndexingPerformanceModelTest,
+TEST_F(GpuIndexingPerformanceModelTest,
        EstimateBestTiling_TritonSoftmax_IsSupported) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
@@ -616,16 +590,8 @@ ENTRY main {
 
   EXPECT_EQ(tiled_runtime_data.block_level_parameters.output_tile_sizes.size(),
             1);
-
-  // Experimental tiling uses padded tile sizes, while the symbolic tiling does
-  // not.
-  if (use_experimental_tiling()) {
-    EXPECT_THAT(tiled_runtime_data.block_level_parameters.output_tile_sizes[0],
-                ElementsAre(4, 1024));
-  } else {
-    EXPECT_THAT(tiled_runtime_data.block_level_parameters.output_tile_sizes[0],
-                ElementsAre(4, 911));
-  }
+  EXPECT_THAT(tiled_runtime_data.block_level_parameters.output_tile_sizes[0],
+              ElementsAre(4, 1024));
   EXPECT_EQ(tiled_runtime_data.block_level_parameters.num_warps, 4);
 
   EXPECT_EQ(tiled_runtime_data.runtime_data.bytes_read, kExpectedBytesRead);
@@ -637,7 +603,7 @@ ENTRY main {
 
 // This test means to catch integer overflow errors when run with ASan build.
 // The checks below are just sanity checks for values.
-TEST_P(
+TEST_F(
     GpuIndexingPerformanceModelTest,
     EstimateRunTimeForTiledFusion_NumberOfTilesLargerThanInt32Max_IsSupported) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
@@ -675,7 +641,7 @@ ENTRY main {
   EXPECT_NEAR(absl::ToDoubleSeconds(runtime_data.exec_time), 2932, 2);
 }
 
-TEST_P(GpuIndexingPerformanceModelTest,
+TEST_F(GpuIndexingPerformanceModelTest,
        EstimateRunTimeForTiledFusion_Concatenate) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
@@ -715,7 +681,7 @@ ENTRY main {
   );
 }
 
-TEST_P(GpuIndexingPerformanceModelTest,
+TEST_F(GpuIndexingPerformanceModelTest,
        EstimateRunTimeForTiledFusion_DotWithReductionLoop) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
@@ -759,7 +725,7 @@ ENTRY main {
   );
 }
 
-TEST_P(GpuIndexingPerformanceModelTest,
+TEST_F(GpuIndexingPerformanceModelTest,
        EstimateRunTimeForTiledFusion_Softmax_RegisterSpill_ReturnsInfinite) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
@@ -799,7 +765,7 @@ ENTRY main {
   EXPECT_TRUE(res2.IsInfinite());
 }
 
-TEST_P(
+TEST_F(
     GpuIndexingPerformanceModelTest,
     EstimateRunTimeForTiledFusion_BroadcastReduce_RegisterSpill_ReturnsInfinite) {  // NOLINT(whitespace/line_length)
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
@@ -855,7 +821,7 @@ ENTRY main {
   EXPECT_TRUE(res3.IsInfinite());
 }
 
-TEST_P(
+TEST_F(
     GpuIndexingPerformanceModelTest,
     EstimateRunTimeForTiledFusion_UsesHloDimensionSizeWhenTileCoversFullDimensionForMemoryAccessTime) {  // NOLINT(whitespace/line_length)
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
@@ -894,7 +860,7 @@ ENTRY main {
   EXPECT_EQ(res.flops, kPaddedOutputTileSize * kAddFlops);
 }
 
-TEST_P(
+TEST_F(
     GpuIndexingPerformanceModelTest,
     EstimateRunTimeForTiledFusion_UncoalescedReadsAreScaledBasedOnWasteTransactionPercentage) {  // NOLINT(whitespace/line_length)
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
@@ -941,7 +907,7 @@ ENTRY main {
       0.001);
 }
 
-TEST_P(
+TEST_F(
     GpuIndexingPerformanceModelTest,
     EstimateRunTimeForTiledFusion_UncoalescedWritesAreScaledBasedOnWasteTransactionPercentage) {  // NOLINT(whitespace/line_length)
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
@@ -987,7 +953,7 @@ ENTRY main {
       4, 0.001);
 }
 
-TEST_P(GpuIndexingPerformanceModelTest,
+TEST_F(GpuIndexingPerformanceModelTest,
        GetLaunchDimensionsForTiledFusion_IsSupported) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
@@ -1006,44 +972,22 @@ ENTRY main {
 )"));
   auto fusion_adaptor = HloFusionAdaptor::ForInstruction(
       module->entry_computation()->root_instruction());
-  const HloInstruction* fusion_root =
-      &fusion_adaptor->GetRoots().front().instruction();
 
   absl::InlinedVector<int64_t, 4> output_tile_sizes = {9, 9, 9};
 
-  int64_t num_warps = 0;
-  if (use_experimental_tiling()) {
-    ASSERT_OK_AND_ASSIGN(
-        std::unique_ptr<experimental::TilingSpace> tiling_space,
-        experimental::TilingSpace::Create(*fusion_adaptor, &mlir_context_));
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<experimental::TilingSpace> tiling_space,
+      experimental::TilingSpace::Create(*fusion_adaptor, &mlir_context_));
 
-    EXPECT_OK(tiling_space->AssignTileSizes(
-        xla::xtile::GetPaddedTileSizes(output_tile_sizes)));
+  EXPECT_OK(tiling_space->AssignTileSizes(
+      xla::xtile::GetPaddedTileSizes(output_tile_sizes)));
 
-    ASSERT_OK_AND_ASSIGN(
-        experimental::TiledHloComputation tiled_hlo_computation,
-        experimental::TiledHloComputation::Tile(*fusion_adaptor,
-                                                std::move(tiling_space)));
+  ASSERT_OK_AND_ASSIGN(experimental::TiledHloComputation tiled_hlo_computation,
+                       experimental::TiledHloComputation::Tile(
+                           *fusion_adaptor, std::move(tiling_space)));
 
-    num_warps = GpuPerformanceModelWithIndexingAnalysis::EstimateNumWarps(
-        tiled_hlo_computation);
-  } else {
-    SymbolicTileAnalysisOrError analysis_or_error =
-        SymbolicTileAnalysis::AnalyzeFusion(
-            *fusion_adaptor, &mlir_context_,
-            /*emitter_specific_constraints_builder=*/nullptr);
-    ASSERT_TRUE(
-        std::holds_alternative<SymbolicTileAnalysis>(analysis_or_error));
-
-    ASSERT_OK_AND_ASSIGN(
-        TiledHloComputation tiled_hlo_computation,
-        std::get<SymbolicTileAnalysis>(analysis_or_error)
-            .ComputeTiledComputation(
-                Tiling({{fusion_root, FlatTiling(output_tile_sizes)}})));
-
-    num_warps = GpuPerformanceModelWithIndexingAnalysis::EstimateNumWarps(
-        tiled_hlo_computation);
-  }
+  int64_t num_warps = GpuPerformanceModelWithIndexingAnalysis::EstimateNumWarps(
+      tiled_hlo_computation);
 
   // Tile size is 9 * 9 * 9 = 729 that corresponds to 2 warps. But we estimate
   // the number of warps for padded tile that has size of 16 * 16 * 16 = 4096
@@ -1051,7 +995,7 @@ ENTRY main {
   EXPECT_EQ(num_warps, 4);
 }
 
-TEST_P(GpuIndexingPerformanceModelTest,
+TEST_F(GpuIndexingPerformanceModelTest,
        NumberOfWarpsDependsOnLargestLiveTileSize) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
@@ -1077,47 +1021,25 @@ ENTRY main {
 )"));
   auto fusion_adaptor = HloFusionAdaptor::ForInstruction(
       module->entry_computation()->root_instruction());
-  const HloInstruction* fusion_root =
-      &fusion_adaptor->GetRoots().front().instruction();
 
   absl::InlinedVector<int64_t, 4> output_tile_sizes = {1};
 
-  int64_t num_warps = 0;
-  if (use_experimental_tiling()) {
-    ASSERT_OK_AND_ASSIGN(
-        std::unique_ptr<experimental::TilingSpace> tiling_space,
-        experimental::TilingSpace::Create(*fusion_adaptor, &mlir_context_));
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<experimental::TilingSpace> tiling_space,
+      experimental::TilingSpace::Create(*fusion_adaptor, &mlir_context_));
 
-    absl::InlinedVector<int64_t, 4> tile_sizes = output_tile_sizes;
-    tile_sizes.push_back(4096);
+  absl::InlinedVector<int64_t, 4> tile_sizes = output_tile_sizes;
+  tile_sizes.push_back(4096);
 
-    EXPECT_OK(tiling_space->AssignTileSizes(
-        xla::xtile::GetPaddedTileSizes(tile_sizes)));
+  EXPECT_OK(tiling_space->AssignTileSizes(
+      xla::xtile::GetPaddedTileSizes(tile_sizes)));
 
-    ASSERT_OK_AND_ASSIGN(
-        experimental::TiledHloComputation tiled_hlo_computation,
-        experimental::TiledHloComputation::Tile(*fusion_adaptor,
-                                                std::move(tiling_space)));
+  ASSERT_OK_AND_ASSIGN(experimental::TiledHloComputation tiled_hlo_computation,
+                       experimental::TiledHloComputation::Tile(
+                           *fusion_adaptor, std::move(tiling_space)));
 
-    num_warps = GpuPerformanceModelWithIndexingAnalysis::EstimateNumWarps(
-        tiled_hlo_computation);
-
-  } else {
-    SymbolicTileAnalysisOrError analysis_or_error =
-        SymbolicTileAnalysis::AnalyzeFusion(
-            *fusion_adaptor, &mlir_context_,
-            /*emitter_specific_constraints_builder=*/nullptr);
-    ASSERT_TRUE(
-        std::holds_alternative<SymbolicTileAnalysis>(analysis_or_error));
-
-    ASSERT_OK_AND_ASSIGN(
-        TiledHloComputation tiled_hlo_computation,
-        std::get<SymbolicTileAnalysis>(analysis_or_error)
-            .ComputeTiledComputation(Tiling({{fusion_root, FlatTiling({1})}})));
-
-    num_warps = GpuPerformanceModelWithIndexingAnalysis::EstimateNumWarps(
-        tiled_hlo_computation);
-  }
+  int64_t num_warps = GpuPerformanceModelWithIndexingAnalysis::EstimateNumWarps(
+      tiled_hlo_computation);
 
   // The largest tile size is 1 * 4096, for which our implementation recommends
   // using 4 warps.
@@ -1148,10 +1070,7 @@ class FlopsPerElementTest : public GpuIndexingPerformanceModelTest {
   }
 };
 
-INSTANTIATE_TEST_SUITE_P(FlopsPerElementTestInstantiation, FlopsPerElementTest,
-                         ::testing::ValuesIn({false}));
-
-TEST_P(FlopsPerElementTest, MatchesGpuHloCostAnalysis_Reduce) {
+TEST_F(FlopsPerElementTest, MatchesGpuHloCostAnalysis_Reduce) {
   // Note: This comparison only succeeds when the reduction dimension size is a
   // power of two (e.g. 64), because FlopsPerElement accounts for power-of-two
   // padding while GpuHloCostAnalysis measures unpadded FLOP counts.
@@ -1172,7 +1091,7 @@ ENTRY entry_computation {
 )");
 }
 
-TEST_P(FlopsPerElementTest, MatchesGpuHloCostAnalysis_VariadicReduce) {
+TEST_F(FlopsPerElementTest, MatchesGpuHloCostAnalysis_VariadicReduce) {
   // Note: This comparison only succeeds when the reduction dimension size is a
   // power of two (e.g. 64), because FlopsPerElement accounts for power-of-two
   // padding while GpuHloCostAnalysis measures unpadded FLOP counts.
@@ -1197,7 +1116,7 @@ ENTRY entry_computation {
 )");
 }
 
-TEST_P(FlopsPerElementTest, PaddedReduceDimensions) {
+TEST_F(FlopsPerElementTest, PaddedReduceDimensions) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -1223,7 +1142,7 @@ ENTRY entry_computation {
   EXPECT_EQ(indexing_cost_model_.FlopsPerElement(instr), 765);
 }
 
-TEST_P(FlopsPerElementTest, MatchesGpuHloCostAnalysis_Elementwise_Cosine) {
+TEST_F(FlopsPerElementTest, MatchesGpuHloCostAnalysis_Elementwise_Cosine) {
   CompareFlopsModels(R"(
 HloModule m
 
@@ -1234,7 +1153,7 @@ ENTRY entry_computation {
 )");
 }
 
-TEST_P(FlopsPerElementTest, MatchesGpuHloCostAnalysis_Elementwise_Clamp) {
+TEST_F(FlopsPerElementTest, MatchesGpuHloCostAnalysis_Elementwise_Clamp) {
   CompareFlopsModels(R"(
 HloModule m
 
@@ -1247,7 +1166,7 @@ ENTRY entry_computation {
 )");
 }
 
-TEST_P(FlopsPerElementTest, MatchesGpuHloCostAnalysis_Gather) {
+TEST_F(FlopsPerElementTest, MatchesGpuHloCostAnalysis_Gather) {
   CompareFlopsModels(R"(
 HloModule module
 entry {
@@ -1259,7 +1178,7 @@ entry {
 })");
 }
 
-TEST_P(FlopsPerElementTest, MatchesGpuHloCostAnalysis_ReduceWindow) {
+TEST_F(FlopsPerElementTest, MatchesGpuHloCostAnalysis_ReduceWindow) {
   CompareFlopsModels(R"(
 
 add {
@@ -1275,7 +1194,7 @@ ENTRY entry {
 })");
 }
 
-TEST_P(FlopsPerElementTest, MatchesGpuHloCostAnalysis_Concatenate_Aligned) {
+TEST_F(FlopsPerElementTest, MatchesGpuHloCostAnalysis_Concatenate_Aligned) {
   CompareFlopsModels(R"(
 HloModule m
 
@@ -1287,7 +1206,7 @@ ENTRY entry_computation {
 )");
 }
 
-TEST_P(FlopsPerElementTest, MatchesGpuHloCostAnalysis_Concatenate_Unaligned) {
+TEST_F(FlopsPerElementTest, MatchesGpuHloCostAnalysis_Concatenate_Unaligned) {
   CompareFlopsModels(R"(
 HloModule m
 
@@ -1299,7 +1218,7 @@ ENTRY entry_computation {
 )");
 }
 
-TEST_P(GpuIndexingPerformanceModelTest, PrecomputeFlopsMapWorksCorrectly) {
+TEST_F(GpuIndexingPerformanceModelTest, PrecomputeFlopsMapWorksCorrectly) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 
@@ -1345,7 +1264,7 @@ ENTRY entry {
   EXPECT_EQ(flops_map.at(exp), indexing_cost_model_.FlopsPerElement(exp));
 }
 
-TEST_P(GpuIndexingPerformanceModelTest,
+TEST_F(GpuIndexingPerformanceModelTest,
        TryFindTopKBestTilingsForFusionAsync_ParallelMatchesSequential) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
@@ -1400,7 +1319,7 @@ ENTRY main {
   }
 }
 
-TEST_P(FlopsPerElementTest, CopyHasZeroFlops) {
+TEST_F(FlopsPerElementTest, CopyHasZeroFlops) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
 

From 2988bed917c8ac293f97e2e2537dfe4c8f3d9426 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 9 Oct 2026 00:53:27 -0700
Subject: [PATCH 0333/1000] This changelist enhances the `VLOG(2)` statements
 in `FindUniqueDynamicSliceUserFromCollective` and `MatchWithDynamicSlice` to
 make debugging easier when a collective operation fails to be converted to a
 Reduce-Scatter.

Specifically, it:
*   Prefixes existing log messages with the function name and the word "fails" to clearly indicate the source and nature of the log.
*   Adds missing `VLOG(2)` statements for several failure paths (e.g., when the user is not a dynamic slice, when `GetGroupSize` returns `nullopt`, and when `ExtractSplitDimSpec` fails).

PiperOrigin-RevId: 996355480
---
 .../xla/xla/service/collective_opt_utils.cc   | 58 ++++++++++++++-----
 1 file changed, 42 insertions(+), 16 deletions(-)

diff --git a/third_party/xla/xla/service/collective_opt_utils.cc b/third_party/xla/xla/service/collective_opt_utils.cc
index 37b13f35e3d..b0196bcc556 100644
--- a/third_party/xla/xla/service/collective_opt_utils.cc
+++ b/third_party/xla/xla/service/collective_opt_utils.cc
@@ -1014,6 +1014,8 @@ std::optional<CollectiveUsers> FindUniqueDynamicSliceUserFromCollective(
     bool allow_multiple_users, bool allow_intervening_reshape,
     bool allow_intervening_bitcast) {
   if (instruction->user_count() == 0) {
+    VLOG(2) << "FindUniqueDynamicSliceUserFromCollective fails, No users for "
+            << instruction->ToString();
     return std::nullopt;
   }
 
@@ -1033,7 +1035,9 @@ std::optional<CollectiveUsers> FindUniqueDynamicSliceUserFromCollective(
   if (allow_intervening_reshape) {
     if (user->opcode() == HloOpcode::kReshape) {
       if (user->user_count() != 1) {
-        VLOG(2) << "Reshape user count > 1 for " << user->ToString();
+        VLOG(2) << "FindUniqueDynamicSliceUserFromCollective fails, Reshape "
+                   "user count > 1 for "
+                << user->ToString();
         return std::nullopt;
       }
       result.reshape = user;
@@ -1044,7 +1048,9 @@ std::optional<CollectiveUsers> FindUniqueDynamicSliceUserFromCollective(
   if (allow_intervening_bitcast) {
     if (user->opcode() == HloOpcode::kBitcast) {
       if (user->user_count() != 1) {
-        VLOG(2) << "Bitcast user count > 1 for " << user->ToString();
+        VLOG(2) << "FindUniqueDynamicSliceUserFromCollective fails, Bitcast "
+                   "user count > 1 for "
+                << user->ToString();
         return std::nullopt;
       }
       result.bitcast = user;
@@ -1056,6 +1062,11 @@ std::optional<CollectiveUsers> FindUniqueDynamicSliceUserFromCollective(
     result.dynamic_slice = user;
     return result;
   }
+  VLOG(2) << "FindUniqueDynamicSliceUserFromCollective fails, user is not "
+             "dynamic slice, user_opcode="
+          << HloOpcodeString(user->opcode()) << " user_name=" << user->name()
+          << " user_shape=" << user->shape().ToString()
+          << " collective=" << instruction->name();
   return std::nullopt;
 }
 
@@ -1067,22 +1078,27 @@ std::optional<ReduceScatterSpec> MatchWithDynamicSlice(
     bool is_constrain_layout, bool use_global_device_ids, bool is_cross_module,
     bool allow_intervening_bitcast, bool allow_multiple_users) {
   if (!instruction->shape().IsArray() || is_constrain_layout) {
-    VLOG(2) << "Unsupported collective: " << instruction->ToString();
+    VLOG(2) << "MatchWithDynamicSlice fails, "
+               "Unsupported_collective_or_constrain_layout "
+            << instruction->ToString();
     return std::nullopt;
   }
   if (instruction->shape().dimensions().size() -
           absl::c_count(instruction->shape().dimensions(), 1) <
       min_rank) {
-    VLOG(2) << " Should be at least rank-" << min_rank
-            << " excluding trivial dimensions " << instruction->ToString();
+    VLOG(2) << "MatchWithDynamicSlice fails, Should be at least rank-"
+            << min_rank << " excluding trivial dimensions "
+            << instruction->ToString();
     return std::nullopt;
   }
   if (!allow_multiple_users && instruction->user_count() != 1) {
-    VLOG(2) << "All-gather user_count != 1 " << instruction->ToString();
+    VLOG(2) << "MatchWithDynamicSlice fails, All-gather user_count != 1 "
+            << instruction->ToString();
     return std::nullopt;
   }
   if (!CheckUniformReplicaGroups(instruction)) {
-    VLOG(2) << "Non-uniform replica groups " << instruction->ToString();
+    VLOG(2) << "MatchWithDynamicSlice fails, Non-uniform_replica_groups "
+            << instruction->ToString();
     return std::nullopt;
   }
 
@@ -1092,7 +1108,9 @@ std::optional<ReduceScatterSpec> MatchWithDynamicSlice(
           allow_intervening_bitcast);
 
   if (!ds_user.has_value()) {
-    VLOG(2) << "AG or AR user is not dynamic slice " << instruction->ToString();
+    VLOG(2)
+        << "MatchWithDynamicSlice fails, AG or AR user is not dynamic slice "
+        << instruction->ToString();
     return std::nullopt;
   }
 
@@ -1106,6 +1124,8 @@ std::optional<ReduceScatterSpec> MatchWithDynamicSlice(
       GetGroupSize(is_cross_module, use_global_device_ids, num_partitions,
                    num_replicas, instruction);
   if (!optional_group_size) {
+    VLOG(2) << "MatchWithDynamicSlice fails, GetGroupSize_returned_nullopt "
+            << instruction->ToString();
     return std::nullopt;
   }
   int64_t group_size = *optional_group_size;
@@ -1122,15 +1142,16 @@ std::optional<ReduceScatterSpec> MatchWithDynamicSlice(
     // unless they use use_global_device_ids.
     if (instruction->replica_groups().size() != num_replicas ||
         instruction->replica_groups()[0].replica_ids_size() != 1) {
-      VLOG(2) << "Unsupported size > 1 replica groups for cross-partition, "
-                 "non-global ID "
+      VLOG(2) << "MatchWithDynamicSlice fails, Unsupported size > 1 replica "
+                 "groups for cross-partition, non-global ID "
               << instruction->ToString();
       return std::nullopt;
     }
     spec.sharded_partitions = num_partitions;
   }
   if (group_size < 2) {
-    VLOG(2) << "Group_size < 2, nothing to do " << instruction->ToString();
+    VLOG(2) << "MatchWithDynamicSlice fails, Group_size < 2, nothing to do "
+            << instruction->ToString();
     return std::nullopt;
   }
   spec.group_size = group_size;
@@ -1138,6 +1159,8 @@ std::optional<ReduceScatterSpec> MatchWithDynamicSlice(
   std::optional<SplitDimSpec> split_dim_spec =
       ExtractSplitDimSpec(*user, allow_multiple_split_dims);
   if (!split_dim_spec) {
+    VLOG(2) << "MatchWithDynamicSlice fails, ExtractSplitDimSpec_failed ds="
+            << user->ToString();
     return std::nullopt;
   }
   spec.split_dim = split_dim_spec->split_dim;
@@ -1150,8 +1173,8 @@ std::optional<ReduceScatterSpec> MatchWithDynamicSlice(
   }
 
   if (Product(group_sizes) != group_size) {
-    VLOG(2) << "Group size mismatch " << user->ToString() << " vs "
-            << instruction->ToString();
+    VLOG(2) << "MatchWithDynamicSlice fails, Group size mismatch "
+            << user->ToString() << " vs " << instruction->ToString();
     return std::nullopt;
   }
   if (split_dims.size() > 1) {
@@ -1164,7 +1187,8 @@ std::optional<ReduceScatterSpec> MatchWithDynamicSlice(
     if (!IsPerIdOffsets(absl::MakeSpan(offsets), shard_size, map_id,
                         group_sizes, instruction, is_cross_module,
                         use_global_device_ids)) {
-      VLOG(2) << "IsPerIdOffsets() failed " << instruction->ToString();
+      VLOG(2) << "MatchWithDynamicSlice fails, IsPerIdOffsets() failed "
+              << instruction->ToString();
       return std::nullopt;
     }
   } else {
@@ -1172,7 +1196,8 @@ std::optional<ReduceScatterSpec> MatchWithDynamicSlice(
                        user->dynamic_slice_sizes()[spec.split_dim], map_id,
                        group_size, instruction, is_cross_module,
                        use_global_device_ids)) {
-      VLOG(2) << "IsPerIdOffset() failed " << instruction->ToString();
+      VLOG(2) << "MatchWithDynamicSlice fails, IsPerIdOffset() failed "
+              << instruction->ToString();
       return std::nullopt;
     }
   }
@@ -1195,7 +1220,8 @@ std::optional<ReduceScatterSpec> MatchWithDynamicSlice(
           return unmodified_output_to_input_map.count(out_dim) != 0;
         });
     if (!all_split_dims_unmodified) {
-      VLOG(2) << "Split dimensions are modified by reshape";
+      VLOG(2) << "MatchWithDynamicSlice fails, Split dimensions are modified "
+                 "by reshape";
       return std::nullopt;
     }
 

From 94cada9d87a81f48642a30578dba795a27f51188 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Fri, 9 Oct 2026 00:58:40 -0700
Subject: [PATCH 0334/1000] Remove lit_test.patch; update the one lit_test that
 needs it.

This appears to only be needed for `sink-constants-to-control-flow.mlir` and `legalize-trigonometric-to-approximation.mlir` which are tests that exist in more than one directory. On Windows, apparently `foo/bar.mlir` creates a top-level `bar.mlir.test` file/symlink, creating the constraint that names must be unique.

Instead of patching LLVM's lit_test bazel definition, this applies it in the one place it is needed.

Alternatives that would also avoid this patch:
* Rename a few test files to be unique
* Put BUILD.bazel files in more dirs so these tests are in different bazel packages
* Stop Windows from creating this constraint by fixing the top-level file based on the basename.
* Upstream this patch to LLVM

PiperOrigin-RevId: 996357538
---
 .../xla/third_party/llvm/lit_test.patch       | 28 -------------------
 .../xla/third_party/llvm/workspace.bzl        |  1 -
 third_party/xla/xla/mlir_hlo/tests/BUILD      |  4 ++-
 3 files changed, 3 insertions(+), 30 deletions(-)
 delete mode 100644 third_party/xla/third_party/llvm/lit_test.patch

diff --git a/third_party/xla/third_party/llvm/lit_test.patch b/third_party/xla/third_party/llvm/lit_test.patch
deleted file mode 100644
index 716e678ec70..00000000000
--- a/third_party/xla/third_party/llvm/lit_test.patch
+++ /dev/null
@@ -1,28 +0,0 @@
-# Copyright 2026 The OpenXLA Authors.
-#
-# 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.
-# ==============================================================================
-diff --git a/utils/bazel/llvm-project-overlay/llvm/lit_test.bzl b/utils/bazel/llvm-project-overlay/llvm/lit_test.bzl
---- a/utils/bazel/llvm-project-overlay/llvm/lit_test.bzl
-+++ b/utils/bazel/llvm-project-overlay/llvm/lit_test.bzl
-@@ -204,7 +204,9 @@ def lit_test(
-
-     main = Label("//llvm/utils/lit:lit.py")
-     py_test(
--        name = name,
-+        # Prevents creation of the files with identical names located in different
-+        # directories on Windows platform.
-+        name = name.replace("/", "_"),
-         srcs = [main],
-         main = main,
-         args = args + ["-v"] + ["$(rootpath %s)" % src for src in srcs],
diff --git a/third_party/xla/third_party/llvm/workspace.bzl b/third_party/xla/third_party/llvm/workspace.bzl
index 7b8ac881c64..c78ce3fe662 100644
--- a/third_party/xla/third_party/llvm/workspace.bzl
+++ b/third_party/xla/third_party/llvm/workspace.bzl
@@ -35,7 +35,6 @@ def repo(name):
             "//third_party/llvm:generated.patch",  # Autogenerated, don't remove.
             "//third_party/llvm:repo_metadata.patch",
             "//third_party/llvm:toolchains.patch",
-            "//third_party/llvm:lit_test.patch",
             "//third_party/llvm:run_lit.patch",
         ],
         link_files = {"//third_party/llvm:run_lit.sh": "mlir/run_lit.sh"},
diff --git a/third_party/xla/xla/mlir_hlo/tests/BUILD b/third_party/xla/xla/mlir_hlo/tests/BUILD
index c6b7958fb60..9dc8e114fc6 100644
--- a/third_party/xla/xla/mlir_hlo/tests/BUILD
+++ b/third_party/xla/xla/mlir_hlo/tests/BUILD
@@ -24,7 +24,9 @@ package(
 
 [
     lit_test(
-        name = "%s.test" % src,
+        # Prevents creation of the files with identical names located in different
+        # directories on Windows platform.
+        name = "%s.test" % src.replace("/", "_"),
         srcs = [src],
         data = [
             "lit.cfg.py",

From 5a46026a989c43ee1b712ac63259b620a2647f24 Mon Sep 17 00:00:00 2001
From: Olli Lupton <olupton@nvidia.com>
Date: Fri, 9 Oct 2026 01:10:16 -0700
Subject: [PATCH 0335/1000] PR #49714: Fix cost aggregation for ops with
 multiple kernels
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/49714

📝 Summary of Changes

Previously, an op (e.g. custom call) that launches multiple pieces of GPU work would not be measured correctly e.g. in PGLE.
A single average was taken over all pieces of work and all instances of the op, which is biased.
With this change, the average is only over instances of the op, and we take the [earliest_start, latest_end] range within the op.

🎯 Justification

Custom calls that are implemented as multiple kernels had spuriously low costs measured by PGLE.

🚀 Kind of Contribution

🐛 Bug Fix

📊 Benchmark (for Performance Improvements)

n/a

🧪 Unit Tests:

Existing test case updated to model a custom call implemented with two kernels.

🧪 Execution Tests:

n/a

Fixes nvbug/6842289.
Copybara import of the project:

--
adfda4c8d3063fa95e5ba167210c55138e425a9e by Olli Lupton <olupton@nvidia.com>:

clang-format

--
10e3094e4f78d3775607273021894e0d9ae61039 by Olli Lupton <olupton@nvidia.com>:

Fix cost aggregation for ops with multiple kernels

Previously in this case the mean was taken over N*M values if N copies
of an op launching M kernels were profiled, instead of a mean over N
values of kernel_M-1_end - kernel_0_start.

Merging this change closes #49714

PiperOrigin-RevId: 996362799
---
 third_party/xla/xla/python/BUILD              |  1 +
 .../python/xplane_to_profile_instructions.cc  | 68 ++++++++++++++++---
 .../xplane_to_profile_instructions_test.cc    | 42 ++++++++----
 3 files changed, 87 insertions(+), 24 deletions(-)

diff --git a/third_party/xla/xla/python/BUILD b/third_party/xla/xla/python/BUILD
index c4738a840aa..2ee5f5ab850 100644
--- a/third_party/xla/xla/python/BUILD
+++ b/third_party/xla/xla/python/BUILD
@@ -649,6 +649,7 @@ xla_cc_test(
         "//xla/tsl/profiler/utils:file_system_utils",
         "//xla/tsl/profiler/utils:xplane_builder",
         "//xla/tsl/profiler/utils:xplane_schema",
+        "@tsl//tsl/profiler/protobuf:profiled_instructions_proto_cc",
         "@tsl//tsl/profiler/protobuf:profiled_instructions_proto_cc_impl",
         "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
     ],
diff --git a/third_party/xla/xla/python/xplane_to_profile_instructions.cc b/third_party/xla/xla/python/xplane_to_profile_instructions.cc
index 0306eb7a32d..50845110019 100644
--- a/third_party/xla/xla/python/xplane_to_profile_instructions.cc
+++ b/third_party/xla/xla/python/xplane_to_profile_instructions.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 #include "xla/python/xplane_to_profile_instructions.h"
 
+#include <algorithm>
 #include <cstdint>
 #include <memory>
 #include <numeric>
@@ -64,10 +65,35 @@ using tsl::profiler::XLineVisitor;
 using tsl::profiler::XPlaneVisitor;
 using tsl::profiler::XStatVisitor;
 
+struct HloLatencySpan {
+  int64_t start_ps, end_ps;
+};
+
+struct HloOpInstance {
+  std::string name_and_fingerprint;
+  int64_t instance_id;
+
+  bool operator==(const HloOpInstance& other) const {
+    return name_and_fingerprint == other.name_and_fingerprint &&
+           instance_id == other.instance_id;
+  }
+
+  template <typename H>
+  friend H AbslHashValue(H h, const HloOpInstance& op_instance) {
+    return H::combine(std::move(h), op_instance.name_and_fingerprint,
+                      op_instance.instance_id);
+  }
+};
+
+struct SumAndTotal {
+  double sum = 0.0;
+  int64_t total = 0;
+};
+
 void GetXPlaneLatencyInfo(
     const XPlaneVisitor& xplane,
     const absl::flat_hash_map<std::string, std::string>& hlo_module_info,
-    absl::flat_hash_map<std::string, HloLatencyInfo>* hlo_latency_info) {
+    absl::flat_hash_map<HloOpInstance, HloLatencySpan>* hlo_latency_info) {
   // Iterate events.
   xplane.ForEachLine([hlo_latency_info,
                       &hlo_module_info](const XLineVisitor& xline) {
@@ -83,6 +109,7 @@ void GetXPlaneLatencyInfo(
       std::optional<std::string> hlo_module_name = std::nullopt;
       std::optional<std::string> fingerprint = std::nullopt;
       std::optional<int64_t> program_id = std::nullopt;
+      std::optional<int64_t> scope_range_id = std::nullopt;
 
       auto for_each_stat = [&](const XStatVisitor& stat) {
         if (stat.ValueCase() == tsl::profiler::XStat::VALUE_NOT_SET) return;
@@ -96,6 +123,9 @@ void GetXPlaneLatencyInfo(
         if (stat.Name() == GetStatTypeStr(StatType::kHloModule)) {
           hlo_module_name = stat.ToString();
         }
+        if (stat.Name() == GetStatTypeStr(StatType::kScopeRangeId)) {
+          scope_range_id = stat.IntValue();
+        }
       };
       xevent.Metadata().ForEachStat(for_each_stat);
       xevent.ForEachStat(for_each_stat);
@@ -113,12 +143,25 @@ void GetXPlaneLatencyInfo(
           fingerprint = hlo_module_info.at(fingerprint_key);
         }
       }
-      double latency = static_cast<double>(xevent.DurationNs()) / 1e3;
       std::string key = hlo_name.value();
       if (fingerprint.has_value()) {
         key = absl::StrCat(fingerprint.value(), kCostNameSep, hlo_name.value());
       }
-      (*hlo_latency_info)[key].durations.push_back(latency);
+      // Store the enclosing [start, end] range for device activity with this
+      // fingerprint + name + scope_id. This ensures that ops that produce >1
+      // activity fragment are accounted for properly, as the fragments will
+      // share a scope_id but separate instances of the op will not. If there
+      // is no scope_range_id, use the timestamp as a ~unique key to disable
+      // merging spans.
+      auto [it, inserted] = hlo_latency_info->try_emplace(
+          HloOpInstance{key, scope_range_id.value_or(xevent.TimestampPs())},
+          HloLatencySpan{xevent.TimestampPs(), xevent.EndTimestampPs()});
+      if (!inserted) {
+        it->second.start_ps =
+            std::min(it->second.start_ps, xevent.TimestampPs());
+        it->second.end_ps =
+            std::max(it->second.end_ps, xevent.EndTimestampPs());
+      }
     });
   });
 }
@@ -207,8 +250,8 @@ absl::Status ConvertXplaneToProfiledInstructionsProto(
     std::vector<tensorflow::profiler::XSpace> xspaces,
     tensorflow::profiler::ProfiledInstructionsProto*
         profiled_instructions_proto) {
-  absl::flat_hash_map<std::string, HloLatencyInfo> hlo_latency_info;
   absl::flat_hash_map<std::string, std::string> hlo_module_info;
+  absl::flat_hash_map<std::string, SumAndTotal> duration_stats;
   // Iterate through each host.
   for (const XSpace& xspace : xspaces) {
     const XPlane* metadata_plane =
@@ -229,20 +272,25 @@ absl::Status ConvertXplaneToProfiledInstructionsProto(
       device_planes =
           FindPlanesWithPrefix(xspace, tsl::profiler::kCustomPlanePrefix);
     }
-    // Go over each device plane.
+    // Go over each device plane; hlo_latency_info is per-host to avoid
+    // assuming scope_range_id is unique across hosts.
+    absl::flat_hash_map<HloOpInstance, HloLatencySpan> hlo_latency_info;
     for (const XPlane* device_plane : device_planes) {
       XPlaneVisitor xplane = CreateTfXPlaneVisitor(device_plane);
       GetXPlaneLatencyInfo(xplane, hlo_module_info, &hlo_latency_info);
     }
+    for (const auto& [key, span] : hlo_latency_info) {
+      SumAndTotal& stats = duration_stats[key.name_and_fingerprint];
+      stats.sum += static_cast<double>(span.end_ps - span.start_ps) / 1e6;
+      ++stats.total;
+    }
   }
 
   // Get the mean duration for each hlo and store into the proto.
-  for (const auto& iter : hlo_latency_info) {
+  for (const auto& [name, stats] : duration_stats) {
     auto* cost = profiled_instructions_proto->add_costs();
-    std::vector<double> durations = iter.second.durations;
-    double sum = std::accumulate(durations.begin(), durations.end(), 0.0);
-    cost->set_cost_us(sum / durations.size());
-    cost->set_name(iter.first);
+    cost->set_name(name);
+    cost->set_cost_us(stats.sum / stats.total);
   }
 
   return absl::OkStatus();
diff --git a/third_party/xla/xla/python/xplane_to_profile_instructions_test.cc b/third_party/xla/xla/python/xplane_to_profile_instructions_test.cc
index 5583cf44d9d..48bdfd53370 100644
--- a/third_party/xla/xla/python/xplane_to_profile_instructions_test.cc
+++ b/third_party/xla/xla/python/xplane_to_profile_instructions_test.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/python/xplane_to_profile_instructions.h"
 
+#include <atomic>
 #include <cstdint>
 #include <memory>
 #include <string>
@@ -43,6 +44,9 @@ using tsl::profiler::XEventBuilder;
 using tsl::profiler::XLineBuilder;
 using tsl::profiler::XPlaneBuilder;
 
+std::atomic<int64_t> next_unique_id;
+int64_t GetUniqueId() { return next_unique_id.fetch_add(1); }
+
 void CreateXSpace(XSpace* space, int first_device_latency,
                   int second_device_latency) {
   XPlaneBuilder host_plane(space->add_planes());
@@ -83,25 +87,35 @@ void CreateXSpace(XSpace* space, int first_device_latency,
   event3.AddStatValue(*device_plane.GetOrCreateStatMetadata(
                           GetStatTypeStr(StatType::kProgramId)),
                       program_id);
-
+  event3.AddStatValue(*device_plane.GetOrCreateStatMetadata(
+                          GetStatTypeStr(StatType::kScopeRangeId)),
+                      GetUniqueId());
   XPlaneBuilder device_plane_2(space->add_planes());
   device_plane_2.SetName(GpuPlaneName(1));
   device_plane_2.SetId(0);
   XLineBuilder stream2 = device_plane.GetOrCreateLine(30);
   stream2.SetName("gpu stream 1");
-  XEventBuilder event5 =
-      stream1.AddEvent(*device_plane.GetOrCreateEventMetadata("kernel1"));
-  event5.SetTimestampNs(180000);
-  event5.SetDurationNs(second_device_latency);
-  event5.AddStatValue(
-      *device_plane.GetOrCreateStatMetadata(GetStatTypeStr(StatType::kHloOp)),
-      *device_plane.GetOrCreateStatMetadata("custom-call"));
-  event5.AddStatValue(*device_plane.GetOrCreateStatMetadata(
-                          GetStatTypeStr(StatType::kHloModule)),
-                      *device_plane.GetOrCreateStatMetadata("test_module"));
-  event5.AddStatValue(*device_plane.GetOrCreateStatMetadata(
-                          GetStatTypeStr(StatType::kProgramId)),
-                      program_id);
+  // model second_device_latency as being made up of two kernels back to back
+  int64_t scope_range_id = GetUniqueId();
+  const int kNumKernels = 2;
+  for (int i = 0; i < kNumKernels; ++i) {
+    XEventBuilder event5 =
+        stream1.AddEvent(*device_plane.GetOrCreateEventMetadata("kernel1"));
+    event5.SetTimestampNs(180000 + i * second_device_latency / kNumKernels);
+    event5.SetDurationNs(second_device_latency / kNumKernels);
+    event5.AddStatValue(
+        *device_plane.GetOrCreateStatMetadata(GetStatTypeStr(StatType::kHloOp)),
+        *device_plane.GetOrCreateStatMetadata("custom-call"));
+    event5.AddStatValue(*device_plane.GetOrCreateStatMetadata(
+                            GetStatTypeStr(StatType::kHloModule)),
+                        *device_plane.GetOrCreateStatMetadata("test_module"));
+    event5.AddStatValue(*device_plane.GetOrCreateStatMetadata(
+                            GetStatTypeStr(StatType::kProgramId)),
+                        program_id);
+    event5.AddStatValue(*device_plane.GetOrCreateStatMetadata(
+                            GetStatTypeStr(StatType::kScopeRangeId)),
+                        scope_range_id);
+  }
 }
 
 void CreateXSpaceWithFingerprint(XSpace* space, int first_device_latency,

From 78d7d0326e8223062c66a1d58fe0f640673d6a21 Mon Sep 17 00:00:00 2001
From: Levon Ter-Grigoryan <patrios@google.com>
Date: Fri, 9 Oct 2026 01:10:41 -0700
Subject: [PATCH 0336/1000] [XLA:GPU] Support recording multi-GPU barriers
 explicitly via the XLA FFI Record API.

Multi-GPU barrier kernels in XLA GPU collectives previously only supported direct stream launches, forcing callers that synchronize peers during command buffer recording to fall back to CUDA stream capture.

Add `RecordMultiGpuBarrierWithNccl` and `RecordMultiGpuBarrier` on `GpuCommunicator` and `NcclCommunicator` that create and update barrier kernel launch nodes directly in the command buffer via `CreateLaunch` and `UpdateLaunch`.

PiperOrigin-RevId: 996363026
---
 .../xla/xla/backends/gpu/collectives/BUILD    |  3 +
 .../gpu/collectives/gpu_communicator.h        | 12 +++
 .../gpu/collectives/nccl_communicator.cc      | 22 +++++
 .../gpu/collectives/nccl_communicator.h       | 10 ++
 .../xla/xla/backends/gpu/runtime/BUILD        | 10 +-
 .../backends/gpu/runtime/multi_gpu_barrier.cc | 55 ++++++++++-
 .../backends/gpu/runtime/multi_gpu_barrier.h  | 15 +++
 .../backends/gpu/runtime/record_ffi_test.cc   | 94 +++++++++++++++++++
 .../xla/stream_executor/cuda/cuda_kernel.h    |  3 +
 third_party/xla/xla/stream_executor/kernel.h  |  8 ++
 .../xla/stream_executor/rocm/rocm_kernel.h    |  3 +
 11 files changed, 232 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/collectives/BUILD b/third_party/xla/xla/backends/gpu/collectives/BUILD
index e5f4abdf4fd..b2edb80007d 100644
--- a/third_party/xla/xla/backends/gpu/collectives/BUILD
+++ b/third_party/xla/xla/backends/gpu/collectives/BUILD
@@ -496,10 +496,12 @@ cc_library(
         "//xla/core/collectives:reduction_kind",
         "//xla/core/collectives:registered_memory",
         "//xla/core/collectives:symmetric_memory",
+        "//xla/ffi/api:c_api",
         "//xla/stream_executor:device_address",
         "//xla/stream_executor:kernel_args",
         "//xla/stream_executor:memory_allocation",
         "//xla/tsl/util:tied_ref",
+        "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/status",
@@ -1114,6 +1116,7 @@ cc_library(
         "//xla/core/collectives:reduction_kind",
         "//xla/core/collectives:registered_memory",
         "//xla/core/collectives:symmetric_memory",
+        "//xla/ffi/api:c_api",
         "//xla/stream_executor:device_address",
         "//xla/stream_executor:kernel_args",
         "//xla/stream_executor:memory_allocation",
diff --git a/third_party/xla/xla/backends/gpu/collectives/gpu_communicator.h b/third_party/xla/xla/backends/gpu/collectives/gpu_communicator.h
index f999d82bf71..533b40f09bd 100644
--- a/third_party/xla/xla/backends/gpu/collectives/gpu_communicator.h
+++ b/third_party/xla/xla/backends/gpu/collectives/gpu_communicator.h
@@ -23,6 +23,7 @@ limitations under the License.
 #include <string>
 #include <utility>
 
+#include "absl/base/nullability.h"
 #include "absl/container/inlined_vector.h"
 #include "absl/functional/any_invocable.h"
 #include "absl/status/status.h"
@@ -34,6 +35,7 @@ limitations under the License.
 #include "xla/core/collectives/reduction_kind.h"
 #include "xla/core/collectives/registered_memory.h"
 #include "xla/core/collectives/symmetric_memory.h"
+#include "xla/ffi/api/record_c_api.h"
 #include "xla/future.h"
 #include "xla/stream_executor/device_address.h"
 #include "xla/stream_executor/kernel_args.h"
@@ -46,6 +48,10 @@ namespace stream_executor {
 class StreamExecutor;
 }  // namespace stream_executor
 
+namespace xla::ffi {
+struct RecordContext;
+}  // namespace xla::ffi
+
 namespace xla::gpu {
 
 class GpuSignalDesc : public Communicator::SignalDesc {
@@ -309,6 +315,12 @@ class GpuCommunicator : public Communicator {
     return Unimplemented("LaunchMultiGpuBarrier is not implemented");
   }
 
+  virtual absl::StatusOr<const XLA_FFI_Command*> RecordMultiGpuBarrier(
+      xla::ffi::RecordContext& record_ctx,
+      const XLA_FFI_Command* absl_nullable cmd) {
+    return Unimplemented("RecordMultiGpuBarrier is not implemented");
+  }
+
   virtual void InitializeCrossDeviceBarrier(
       tsl::TiedRef<se::MemoryAllocation> tied_signal_value,
       tsl::TiedRef<se::MemoryAllocation> tied_signal,
diff --git a/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.cc b/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.cc
index ba521525d76..06fb0baa76e 100644
--- a/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.cc
+++ b/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.cc
@@ -55,6 +55,7 @@ limitations under the License.
 #include "xla/core/collectives/reduction_kind.h"
 #include "xla/core/collectives/registered_memory.h"
 #include "xla/core/collectives/symmetric_memory.h"
+#include "xla/ffi/api/record_c_api.h"
 #include "xla/future.h"
 #include "xla/primitive_util.h"
 #include "xla/stream_executor/device_address.h"
@@ -1147,6 +1148,27 @@ absl::Status NcclCommunicator::LaunchMultiGpuBarrier(const Executor& executor) {
       tied_cross_device_barrier_signal_value_.Lock()->address());
 }
 
+absl::StatusOr<const XLA_FFI_Command*> NcclCommunicator::RecordMultiGpuBarrier(
+    xla::ffi::RecordContext& record_ctx,
+    const XLA_FFI_Command* absl_nullable cmd) {
+  if (!IsCrossDeviceBarrierInitiated()) {
+    return FailedPrecondition(
+        "Cross device barrier buffers are not set on this communicator. Did "
+        "you set use_cross_device_barrier=true in BarrierRequirements?");
+  }
+  absl::MutexLock lock(barrier_mu_);
+  if (cancel_->IsCancelled()) {
+    return FailedPrecondition("NcclCommunicator aborted");
+  }
+  ABSL_ASSIGN_OR_RETURN(size_t num_ranks, NumRanks());
+  ABSL_ASSIGN_OR_RETURN(size_t current_rank, CurrentRank());
+
+  return xla::gpu::RecordMultiGpuBarrierWithNccl(
+      stream_executor_, record_ctx, num_ranks, RankId(current_rank),
+      tied_cross_device_barrier_symmetric_memory_.Lock().get(),
+      tied_cross_device_barrier_signal_value_.Lock()->address(), cmd);
+}
+
 void NcclCommunicator::InitializeCrossDeviceBarrier(
     tsl::TiedRef<se::MemoryAllocation> tied_signal_value,
     tsl::TiedRef<se::MemoryAllocation> tied_signal,
diff --git a/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.h b/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.h
index ef343f18be5..131ea559821 100644
--- a/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.h
+++ b/third_party/xla/xla/backends/gpu/collectives/nccl_communicator.h
@@ -23,6 +23,7 @@ limitations under the License.
 #include <string>
 #include <utility>
 
+#include "absl/base/nullability.h"
 #include "absl/base/thread_annotations.h"
 #include "absl/container/inlined_vector.h"
 #include "absl/functional/any_invocable.h"
@@ -42,6 +43,7 @@ limitations under the License.
 #include "xla/core/collectives/registered_memory.h"
 #include "xla/core/collectives/symmetric_memory.h"
 #include "xla/debug_options_flags.h"
+#include "xla/ffi/api/record_c_api.h"
 #include "xla/future.h"
 #include "xla/stream_executor/device_address.h"
 #include "xla/stream_executor/kernel_args.h"
@@ -56,6 +58,10 @@ limitations under the License.
 #include "third_party/nccl/nccl.h"
 #include "third_party/nccl/nccl_device.h"  // IWYU pragma: keep
 
+namespace xla::ffi {
+struct RecordContext;
+}  // namespace xla::ffi
+
 namespace xla::gpu {
 
 using NcclSignalDesc = GpuSignalDesc;
@@ -271,6 +277,10 @@ class NcclCommunicator : public GpuCommunicator {
 
   absl::Status LaunchMultiGpuBarrier(const Executor& executor) final;
 
+  absl::StatusOr<const XLA_FFI_Command*> RecordMultiGpuBarrier(
+      xla::ffi::RecordContext& record_ctx,
+      const XLA_FFI_Command* absl_nullable cmd) final;
+
   void InitializeCrossDeviceBarrier(
       tsl::TiedRef<se::MemoryAllocation> tied_signal_value,
       tsl::TiedRef<se::MemoryAllocation> tied_signal,
diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index 5ec75676509..8f30fdfe46d 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -3430,16 +3430,20 @@ cc_library(
         "//xla:status_macros",
         "//xla/core/collectives:rank_id",
         "//xla/core/collectives:symmetric_memory",
+        "//xla/ffi:record_ffi",
+        "//xla/ffi/api:c_api",
+        "//xla/ffi/api:record_api",
         "//xla/stream_executor:device_address",
         "//xla/stream_executor:launch_dim",
         "//xla/stream_executor:stream",
         "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor/gpu:gpu_kernel_registry",
         "//xla/stream_executor/gpu:multi_gpu_barrier_kernel",
-        "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/base:no_destructor",
+        "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/synchronization",
     ],
@@ -5688,6 +5692,7 @@ cc_library(
     name = "record_ffi",
     srcs = ["record_ffi.cc"],
     hdrs = ["record_ffi.h"],
+    compatible_with = get_compatible_with_portable(),
     visibility = ["//visibility:public"],
     deps = [
         "//xla:util",
@@ -5730,9 +5735,12 @@ xla_test(
     ),
     backends = ["gpu"],
     deps = [
+        ":multi_gpu_barrier",
         ":record_ffi",
         "//xla:status_macros",
         "//xla:xla_data_proto_cc",
+        "//xla/core/collectives:rank_id",
+        "//xla/core/collectives:symmetric_memory",
         "//xla/ffi",
         "//xla/ffi:api",
         "//xla/ffi:call_frame",
diff --git a/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.cc b/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.cc
index 4f2ce8c8f39..c80846081d0 100644
--- a/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.cc
@@ -20,16 +20,22 @@ limitations under the License.
 #include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <string>
 #include <utility>
 #include <vector>
 
 #include "absl/base/no_destructor.h"
+#include "absl/base/nullability.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
+#include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
 #include "absl/synchronization/mutex.h"
 #include "xla/core/collectives/rank_id.h"
 #include "xla/core/collectives/symmetric_memory.h"
+#include "xla/ffi/api/record_api.h"
+#include "xla/ffi/api/record_c_api.h"
+#include "xla/ffi/record_ffi.h"
 #include "xla/status_macros.h"
 #include "xla/stream_executor/device_address.h"
 #include "xla/stream_executor/gpu/gpu_kernel_registry.h"
@@ -37,7 +43,6 @@ limitations under the License.
 #include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/stream.h"
 #include "xla/stream_executor/stream_executor.h"
-#include "xla/tsl/platform/statusor.h"
 
 namespace xla {
 namespace gpu {
@@ -122,7 +127,8 @@ absl::Status LaunchMultiGpuBarrierWithNccl(
   using MultiGpuBarrierWithNcclKernel =
       stream_executor::gpu::MultiGpuBarrierWithNcclKernel;
 
-  TF_RET_CHECK(symmetric_memory != nullptr) << "Symmetric memory is required";
+  TF_RET_CHECK(symmetric_memory != nullptr)
+      << "Symmetric memory is required on rank " << rank.value();
 
   ABSL_ASSIGN_OR_RETURN(
       MultiGpuBarrierWithNcclKernel::KernelType * kernel,
@@ -139,6 +145,51 @@ absl::Status LaunchMultiGpuBarrierWithNccl(
                         typed_sync_counter);
 }
 
+absl::StatusOr<const XLA_FFI_Command*> RecordMultiGpuBarrierWithNccl(
+    stream_executor::StreamExecutor* executor,
+    xla::ffi::RecordContext& record_ctx, int64_t num_devices, RankId rank,
+    xla::SymmetricMemory* symmetric_memory,
+    stream_executor::DeviceAddressBase local_barrier_signal_value,
+    const XLA_FFI_Command* absl_nullable cmd) {
+  using MultiGpuBarrierWithNcclKernel =
+      stream_executor::gpu::MultiGpuBarrierWithNcclKernel;
+
+  TF_RET_CHECK(symmetric_memory != nullptr)
+      << "Symmetric memory is required on rank " << rank.value();
+
+  int64_t rank_val = rank.value();
+  int64_t num_devices_val = num_devices;
+  std::array<ffi::KernelArg, 4> kernel_args = {
+      ffi::HostValue{&rank_val, sizeof(rank_val)},
+      ffi::HostValue{&num_devices_val, sizeof(num_devices_val)},
+      ffi::DevicePointer{symmetric_memory->PackKernelArg()},
+      ffi::DevicePointer{local_barrier_signal_value.opaque()},
+  };
+
+  if (record_ctx.action() == ffi::RecordAction::kCreate) {
+    TF_RET_CHECK(executor != nullptr) << "StreamExecutor is required";
+    ABSL_ASSIGN_OR_RETURN(MultiGpuBarrierWithNcclKernel::KernelType * kernel,
+                     GetCachedKernel<MultiGpuBarrierWithNcclKernel>(executor));
+    XLA_FFI_LaunchDims launch_dims = {
+        /*grid=*/{1, 1, 1},
+        /*block=*/
+        {static_cast<int32_t>(MultiGpuBarrierWithNcclKernel::kMaxPeers), 1, 1},
+        /*cluster=*/{0, 0, 0}};
+    return record_ctx.CreateLaunch(
+        std::string((*kernel)->name()).c_str(),
+        (*kernel)->platform_specific_handle().kernel,
+        /*kernel_size=*/0, ffi::SourceFormat::kFunctionPtr, launch_dims,
+        /*shared_mem_bytes=*/0, /*uses_pdl=*/false, kernel_args);
+  }
+
+  TF_RET_CHECK(record_ctx.action() == ffi::RecordAction::kUpdate)
+      << "Unexpected record action: " << static_cast<int>(record_ctx.action());
+  TF_RET_CHECK(cmd != nullptr)
+      << "Command pointer is required for RecordAction::kUpdate";
+  ABSL_RETURN_IF_ERROR(record_ctx.UpdateLaunch(cmd, kernel_args));
+  return cmd;
+}
+
 size_t GetMultiGpuBarrierSignalBufferSize() {
   return stream_executor::gpu::MultiGpuBarrierKernel::kMaxPeers *
          sizeof(uint32_t);
diff --git a/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.h b/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.h
index c38710e17e1..e2f3ff05ae5 100644
--- a/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.h
+++ b/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.h
@@ -20,11 +20,19 @@ limitations under the License.
 #include <cstdint>
 #include <vector>
 
+#include "absl/base/nullability.h"
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "xla/core/collectives/rank_id.h"
 #include "xla/core/collectives/symmetric_memory.h"
+#include "xla/ffi/api/record_c_api.h"
 #include "xla/stream_executor/device_address.h"
 #include "xla/stream_executor/stream.h"
+#include "xla/stream_executor/stream_executor.h"
+
+namespace xla::ffi {
+struct RecordContext;
+}  // namespace xla::ffi
 
 namespace xla::gpu {
 
@@ -43,6 +51,13 @@ absl::Status LaunchMultiGpuBarrierWithNccl(
     xla::SymmetricMemory* symmetric_memory,
     stream_executor::DeviceAddressBase local_barrier_signal_value);
 
+absl::StatusOr<const XLA_FFI_Command*> RecordMultiGpuBarrierWithNccl(
+    stream_executor::StreamExecutor* executor,
+    xla::ffi::RecordContext& record_ctx, int64_t num_devices, RankId rank,
+    xla::SymmetricMemory* symmetric_memory,
+    stream_executor::DeviceAddressBase local_barrier_signal_value,
+    const XLA_FFI_Command* absl_nullable cmd);
+
 // Returns the size of the barrier signal buffer in bytes.
 size_t GetMultiGpuBarrierSignalBufferSize();
 
diff --git a/third_party/xla/xla/backends/gpu/runtime/record_ffi_test.cc b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test.cc
index caebad7b3c0..96cd0155196 100644
--- a/third_party/xla/xla/backends/gpu/runtime/record_ffi_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/record_ffi_test.cc
@@ -29,6 +29,9 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
+#include "xla/backends/gpu/runtime/multi_gpu_barrier.h"
+#include "xla/core/collectives/rank_id.h"
+#include "xla/core/collectives/symmetric_memory.h"
 #include "xla/ffi/api/record_api.h"
 #include "xla/ffi/api/record_c_api.h"
 #include "xla/ffi/call_frame.h"
@@ -311,5 +314,96 @@ TEST(RecordFfiTest, KernelLaunchWithCuFunc) {
   RunRecordFfiTest(executor, *handler);
 }
 
+class FakeSymmetricMemory : public xla::SymmetricMemory {
+ public:
+  se::DeviceAddressBase addr() const override {
+    return se::DeviceAddressBase();
+  }
+  std::string ToString() const override { return "FakeSymmetricMemory"; }
+  PackedKernelArg PackKernelArg() const override { return nullptr; }
+};
+
+TEST(RecordFfiTest, RecordMultiGpuBarrierWithNccl) {
+  ASSERT_OK_AND_ASSIGN(auto platform,
+                       stream_executor::PlatformManager::PlatformWithName(
+                           stream_executor::GpuPlatformName()));
+  ASSERT_OK_AND_ASSIGN(auto* executor, platform->ExecutorForDevice(0));
+  if (!executor->GetDeviceDescription()
+           .gpu_compute_capability()
+           .cuda_compute_capability()) {
+    GTEST_SKIP() << "MultiGpuBarrierWithNcclKernel is only supported on CUDA.";
+  }
+  ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream());
+
+  FakeSymmetricMemory sym_mem;
+  std::unique_ptr<ffi::Ffi> handler =
+      ffi::Ffi::BindRecord()
+          .Ctx<ffi::Extension<ffi::RecordExtension>>()
+          .Arg<ffi::AnyBuffer>()
+          .To([executor, &sym_mem](ffi::RecordContext record_ctx,
+                                   ffi::AnyBuffer counter) -> absl::Status {
+            const XLA_FFI_Command* cmd =
+                record_ctx.action() == ffi::RecordAction::kUpdate
+                    ? record_ctx.commands()[0]
+                    : nullptr;
+            return RecordMultiGpuBarrierWithNccl(executor, record_ctx,
+                                                 /*num_devices=*/0, RankId(0),
+                                                 &sym_mem,
+                                                 counter.device_memory(), cmd)
+                .status();
+          });
+
+  se::DeviceAddress<uint32_t> counter_1 = executor->AllocateArray<uint32_t>(1);
+  ASSERT_OK(stream->MemZero(&counter_1, sizeof(uint32_t)));
+
+  ASSERT_OK_AND_ASSIGN(auto cmd_buffer,
+                       executor->CreateCommandBuffer(
+                           stream_executor::CommandBuffer::Mode::kPrimary));
+
+  const XLA_FFI_RecordApi* ffi_api = GetXlaFfiRecordApi();
+  int64_t num_commands = 0;
+  const XLA_FFI_Command* commands_storage[1] = {nullptr};
+
+  XLA_FFI_RecordContext record_ctx_c = {cmd_buffer.get(), executor, {}, false};
+  XLA_FFI_RecordFrame record_frame = {
+      &record_ctx_c,    ffi_api,       XLA_FFI_RecordAction_Create,
+      commands_storage, &num_commands, 1};
+  auto record_extension = ffi::BuildRecordCExtension(&record_frame);
+
+  ffi::InvokeContext invoke_context = {};
+  invoke_context.extension_start = &record_extension.extension_base;
+
+  ffi::CallFrameBuilder builder(/*num_args=*/1, /*num_rets=*/0);
+  builder.AddBufferArg(counter_1, PrimitiveType::U32, {1});
+  ffi::CallFrame call_frame = builder.Build();
+
+  ASSERT_OK(ffi::Invoke(ffi::GetXlaFfiApi(), *handler, call_frame,
+                        invoke_context, ffi::ExecutionStage::kRecord));
+  ASSERT_OK(cmd_buffer->Finalize());
+  ASSERT_OK(cmd_buffer->Submit(stream.get()));
+
+  uint32_t host_counter = 0;
+  ASSERT_OK(stream->Memcpy(&host_counter, counter_1, sizeof(uint32_t)));
+  ASSERT_OK(stream->BlockHostUntilDone());
+  EXPECT_EQ(host_counter, 1);
+
+  // Update with new buffer.
+  se::DeviceAddress<uint32_t> counter_2 = executor->AllocateArray<uint32_t>(1);
+  ASSERT_OK(stream->MemZero(&counter_2, sizeof(uint32_t)));
+
+  record_frame.action = XLA_FFI_RecordAction_Update;
+  ASSERT_OK(call_frame.UpdateWithBuffers({counter_2}, /*rets=*/{}));
+  ASSERT_OK(cmd_buffer->Update());
+  ASSERT_OK(ffi::Invoke(ffi::GetXlaFfiApi(), *handler, call_frame,
+                        invoke_context, ffi::ExecutionStage::kRecord));
+  ASSERT_OK(cmd_buffer->Finalize());
+  ASSERT_OK(cmd_buffer->Submit(stream.get()));
+
+  host_counter = 0;
+  ASSERT_OK(stream->Memcpy(&host_counter, counter_2, sizeof(uint32_t)));
+  ASSERT_OK(stream->BlockHostUntilDone());
+  EXPECT_EQ(host_counter, 1);
+}
+
 }  // namespace
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_kernel.h b/third_party/xla/xla/stream_executor/cuda/cuda_kernel.h
index f976e0d4701..0823e296c83 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_kernel.h
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_kernel.h
@@ -60,6 +60,9 @@ class CudaKernel : public Kernel {
   void set_gpu_function(CUfunction gpu_function) {
     gpu_function_ = gpu_function;
   }
+  PlatformSpecificHandle platform_specific_handle() const override {
+    return PlatformSpecificHandle{gpu_function_};
+  }
 
   // Collects metadata for the specified kernel.
   absl::StatusOr<KernelMetadata> GetKernelMetadata();
diff --git a/third_party/xla/xla/stream_executor/kernel.h b/third_party/xla/xla/stream_executor/kernel.h
index 77797c10c4d..0aa54e230c9 100644
--- a/third_party/xla/xla/stream_executor/kernel.h
+++ b/third_party/xla/xla/stream_executor/kernel.h
@@ -156,6 +156,14 @@ class Kernel {
                               const std::optional<ClusterDim>& cluster_dims,
                               Stream* stream, const KernelArgs& args) = 0;
 
+  struct PlatformSpecificHandle {
+    void* kernel = nullptr;
+  };
+
+  virtual PlatformSpecificHandle platform_specific_handle() const {
+    return PlatformSpecificHandle{};
+  }
+
   void set_use_pdl(bool use_pdl) { use_pdl_ = use_pdl; }
   bool use_pdl() const { return use_pdl_; }
 
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_kernel.h b/third_party/xla/xla/stream_executor/rocm/rocm_kernel.h
index ff532209094..2817db790a1 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_kernel.h
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_kernel.h
@@ -56,6 +56,9 @@ class RocmKernel : public Kernel {
   void set_gpu_function(hipFunction_t rocm_function) {
     rocm_function_ = rocm_function;
   }
+  PlatformSpecificHandle platform_specific_handle() const override {
+    return PlatformSpecificHandle{rocm_function_};
+  }
 
   // Collects metadata for the specified kernel.
   absl::StatusOr<KernelMetadata> GetKernelMetadata();

From 46f0f2495e233b583ac1a14b6b63967ecb4a31c4 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Fri, 9 Oct 2026 01:15:28 -0700
Subject: [PATCH 0337/1000] Bump rules_python to 2.2.0 in StableHLO
 WORKSPACE.bazel

Upstream LLVM commit 62e30e3326d9 updated `lit_test.bzl` to pass
`@rules_python//command_line_option:enable_runfiles` in `config_settings`,
which requires `rules_python` >= 1.9.0 so that `py_executable.bzl` declares
`enable_runfiles` in its transition outputs.

Upgrade `rules_python` from 1.8.5 to 2.2.0 to match XLA OSS and fix Bazel
analysis of `lit_test` targets, and regenerate `temporary.patch`.

PiperOrigin-RevId: 996364903
---
 .../xla/third_party/stablehlo/temporary.patch       | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch
index c6cd6f2bc9c..08f59aacc6b 100644
--- a/third_party/xla/third_party/stablehlo/temporary.patch
+++ b/third_party/xla/third_party/stablehlo/temporary.patch
@@ -788,6 +788,19 @@ diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel
 diff --ruN a/stablehlo/WORKSPACE.bazel b/stablehlo/WORKSPACE.bazel
 --- stablehlo/WORKSPACE.bazel
 +++ stablehlo/WORKSPACE.bazel
+@@ -56,9 +56,9 @@
+ 
+ http_archive(
+     name = "rules_python",
+-    sha256 = "8964aa1e7525fea5244ba737458694a057ada1be96a92998a41caa1983562d00",
+-    strip_prefix = "rules_python-1.8.5",
+-    url = "https://github.com/bazel-contrib/rules_python/releases/download/1.8.5/rules_python-1.8.5.tar.gz",
++    sha256 = "e11d2e1efce1589e5bdfa93986712c74fc7467a0f093143d489d2ef5ebb1ed2a",
++    strip_prefix = "rules_python-2.2.0",
++    url = "https://github.com/bazel-contrib/rules_python/releases/download/2.2.0/rules_python-2.2.0.tar.gz",
+ )
+ 
+ load("@rules_python//python:repositories.bzl", "py_repositories", "python_register_toolchains")
 @@ -119,6 +119,18 @@
      ],
  )

From 5503ca36ef309bafeb9b061bcac8d19a159358cd Mon Sep 17 00:00:00 2001
From: Jiayang Liu <yzhao3442@gmail.com>
Date: Thu, 8 Oct 2026 09:40:05 +0800
Subject: [PATCH 0338/1000] Propagate NaN in tf.experimental.numpy.heaviside

NumPy propagates NaN from the first argument of heaviside; TensorFlow's
relational operators evaluate to False for NaN, so the unguarded
implementation fell through to x2 (or 0/1) and silently dropped the NaN
(tensorflow/tensorflow#127819).

Mask the NaN into the fallback value and keep the mask nested inside the
x1 < 0 branch so that no extra full-tensor pass is introduced.

Use y.dtype.is_floating for the trailing promotion guard instead of
np.issubdtype(y.dtype.as_numpy_dtype, np.inexact): ml_dtypes.bfloat16 is not
part of NumPy's inexact hierarchy, so the latter silently upcast
heaviside(bfloat16, bfloat16) to float64.
---
 .../python/ops/numpy_ops/np_math_ops.py       | 22 ++++++--
 .../python/ops/numpy_ops/np_math_ops_test.py  | 51 +++++++++++++++++++
 2 files changed, 69 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py
index e48b9c44554..37f10271e65 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@@ -436,14 +436,28 @@ def nextafter(x1, x2):
 @np_utils.np_doc('heaviside')
 def heaviside(x1, x2):  # pylint: disable=missing-function-docstring
   def f(x1, x2):
+    dtype = x2.dtype
+    zero = constant_op.constant(0, dtype=dtype)
+    one = constant_op.constant(1, dtype=dtype)
+    # NumPy propagates NaN from the first argument. TensorFlow's relational
+    # operators evaluate to False for NaN inputs, so the `x1 > 0` and
+    # `x1 < 0` branches both miss and the value falls through to `x2`. Nest
+    # the NaN selection into that fallback so the outer `where_v2` remains a
+    # single pass over the tensor. Integer dtypes cannot represent NaN, so
+    # only floating inputs get the extra mask. Use the native `is_floating`
+    # property rather than `np.issubdtype(..., np.inexact)`: TensorFlow's
+    # bfloat16 maps to `ml_dtypes.bfloat16`, which is not part of NumPy's
+    # `inexact` hierarchy.
+    fallback = x2
+    if x1.dtype.is_floating:
+      nan = constant_op.constant(np.nan, dtype=dtype)
+      fallback = array_ops.where_v2(math_ops.is_nan(x1), nan, fallback)
     return array_ops.where_v2(
-        x1 < 0,
-        constant_op.constant(0, dtype=x2.dtype),
-        array_ops.where_v2(x1 > 0, constant_op.constant(1, dtype=x2.dtype), x2),
+        x1 > 0, one, array_ops.where_v2(x1 < 0, zero, fallback)
     )
 
   y = _bin_op(f, x1, x2)
-  if not np.issubdtype(y.dtype.as_numpy_dtype, np.inexact):
+  if not y.dtype.is_floating:
     # See the note in `_scalar`: `astype` is unavailable without the
     # `enable_numpy_methods_on_tensor()` opt-in.
     y = math_ops.cast(y, np_utils.result_type(float))
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 ca63f065bf9..408cd0adcf1 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
@@ -228,6 +228,57 @@ class MathTest(test.TestCase, parameterized.TestCase):
       for grad in tape.gradient(y, (x1, x2)):
         self.assertAllClose(grad, [0.5, 0.5, 0.5, 0.5])
 
+  def testHeaviside(self):
+    # x1 < 0 -> 0, x1 == 0 -> x2, x1 > 0 -> 1 (NumPy semantics).
+    x1 = np.array([-1.0, 0.0, 1.0, 2.0, -3.0, 0.5], dtype=np.float32)
+    x2 = np.array([5.0, 5.0, 5.0, 5.0, 5.0, 5.0], dtype=np.float32)
+    self.match(
+        np_math_ops.heaviside(x1, x2),
+        np.heaviside(x1, x2),
+        msg='heaviside standard',
+    )
+
+  @parameterized.parameters(
+      *[(dtype,) for dtype in (np.float16, np.float32, np.float64,
+                               dtypes.bfloat16.as_numpy_dtype)]
+  )
+  def testHeavisideNaNPropagation(self, dtype):
+    # NumPy propagates NaN from the first argument. TensorFlow's relational
+    # operators evaluate to False for NaN, so the unguarded implementation
+    # returned x2 (or 0/1) instead of NaN. Regression test for
+    # tensorflow/tensorflow#127819.
+    #
+    # `x2` is deliberately 0.5 -- a value produced by neither the `x1 > 0`
+    # branch (1) nor the `x1 < 0` branch (0) -- so a NaN that silently falls
+    # through to `x2` cannot masquerade as a correct result. float16 and
+    # bfloat16 are included so half-precision propagation is covered;
+    # bfloat16 in particular is the dtype for which
+    # `np.issubdtype(bfloat16, np.inexact)` is False, which is precisely the
+    # regression this test guards against.
+    x1 = np.array([np.nan, 0.0, -1.0, 1.0], dtype=dtype)
+    x2 = np.array([0.5, 0.5, 0.5, 0.5], dtype=dtype)
+    actual = np_math_ops.heaviside(x1, x2)
+    expected = np.heaviside(x1, x2)
+    # Assert the dtype explicitly: `np.array_equal` compares values only, so a
+    # silent upcast (e.g. bfloat16 -> float64) would otherwise pass unnoticed.
+    self.assertEqual(actual.dtype.as_numpy_dtype, expected.dtype)
+    # np.testing.assert_equal treats NaN != NaN, so compare with equal_nan.
+    self.assertTrue(np.array_equal(actual.numpy(), expected, equal_nan=True))
+
+  def testHeavisideIntegerInputs(self):
+    # Integer inputs have no NaN; the result is still promoted to a floating
+    # point dtype (matching NumPy) and the integer cases are reproduced.
+    # `self.match` is used (rather than comparing `.tolist()` values) so that
+    # the promoted dtype and the shape are asserted as well, which catches an
+    # integer input failing to promote to floating point.
+    x1 = np.array([-1, 0, 1, 2], dtype=np.int32)
+    x2 = np.array([3, 3, 3, 3], dtype=np.int32)
+    self.match(
+        np_math_ops.heaviside(x1, x2),
+        np.heaviside(x1, x2),
+        msg='heaviside integer inputs promote to float',
+    )
+
   def testLogaddexp(self):
     self._testBinaryOp(np_math_ops.logaddexp, np.logaddexp, 'logaddexp')
     self._testBinaryOp(np_math_ops.logaddexp2, np.logaddexp2, 'logaddexp2')

From 67208da5f65a1e2f3b1d313483174986c065c747 Mon Sep 17 00:00:00 2001
From: gauravarun21 <gauravarun21@gmail.com>
Date: Fri, 9 Oct 2026 11:01:45 +0200
Subject: [PATCH 0339/1000] Address review: preserve name scope, test
 hermeticity, dedupe dtype check

- reduce_std now wraps the shared helper call in a `reduce_variance`
  name scope, restoring the exact op-naming structure it had before
  (previously it called the public reduce_variance, which opened that
  sub-scope itself; calling the private helper directly flattened it).
- Pass is_float16 into _reduce_variance_in_stable_dtype instead of
  re-checking input_tensor.dtype inside it.
- Use a local np.random.RandomState in the new regression test instead
  of mutating the global numpy random state via np.random.seed.

Addresses review feedback from @dmiltr3.
---
 tensorflow/python/ops/math_ops.py      | 14 ++++++++++----
 tensorflow/python/ops/math_ops_test.py |  4 ++--
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/tensorflow/python/ops/math_ops.py b/tensorflow/python/ops/math_ops.py
index b0f4b195f2a..0d120c5845e 100644
--- a/tensorflow/python/ops/math_ops.py
+++ b/tensorflow/python/ops/math_ops.py
@@ -2742,13 +2742,15 @@ def reduce_variance(input_tensor, axis=None, keepdims=False, name=None):
   with ops.name_scope(name):
     input_tensor = ops.convert_to_tensor(input_tensor)
     is_float16 = input_tensor.dtype == dtypes.float16
-    variance = _reduce_variance_in_stable_dtype(input_tensor, axis, keepdims)
+    variance = _reduce_variance_in_stable_dtype(
+        input_tensor, is_float16, axis, keepdims)
     if is_float16:
       variance = cast(variance, dtypes.float16)
     return variance
 
 
-def _reduce_variance_in_stable_dtype(input_tensor, axis, keepdims):
+def _reduce_variance_in_stable_dtype(input_tensor, is_float16, axis,
+                                     keepdims):
   """Computes variance, promoting float16 to float32 first.
 
   The dynamic range of float16 is too limited for the squared deviations
@@ -2758,7 +2760,7 @@ def _reduce_variance_in_stable_dtype(input_tensor, axis, keepdims):
   original dtype, so that `reduce_std` can take the square root while still
   in the wider dtype.
   """
-  if input_tensor.dtype == dtypes.float16:
+  if is_float16:
     input_tensor = cast(input_tensor, dtypes.float32)
   means = reduce_mean(input_tensor, axis=axis, keepdims=True)
   if means.dtype.is_integer:
@@ -2824,7 +2826,11 @@ def reduce_std(input_tensor, axis=None, keepdims=False, name=None):
   with ops.name_scope(name):
     input_tensor = ops.convert_to_tensor(input_tensor)
     is_float16 = input_tensor.dtype == dtypes.float16
-    variance = _reduce_variance_in_stable_dtype(input_tensor, axis, keepdims)
+    # Preserve the `reduce_variance` sub-scope reduce_std has always had,
+    # rather than calling the private helper directly at this scope level.
+    with ops.name_scope("reduce_variance"):
+      variance = _reduce_variance_in_stable_dtype(
+          input_tensor, is_float16, axis, keepdims)
     std = gen_math_ops.sqrt(variance)
     if is_float16:
       std = cast(std, dtypes.float16)
diff --git a/tensorflow/python/ops/math_ops_test.py b/tensorflow/python/ops/math_ops_test.py
index 411a66ebb98..a7a4c4b1739 100644
--- a/tensorflow/python/ops/math_ops_test.py
+++ b/tensorflow/python/ops/math_ops_test.py
@@ -151,8 +151,8 @@ class ReduceTest(test_util.TensorFlowTestCase):
     # does not, so reduce_std/reduce_variance need to compute internally in
     # a wider dtype. Previously this produced NaN on CPU and Inf on GPU for
     # reduce_std, instead of a finite, correct result on both.
-    np.random.seed(0)
-    x_np = np.random.randn(1000).astype(np.float32) * 1e4
+    rs = np.random.RandomState(0)
+    x_np = rs.randn(1000).astype(np.float32) * 1e4
     x_f16 = constant_op.constant(x_np.astype(np.float16))
     expected_std = np.std(x_np.astype(np.float64))
 

From ebf53c9c5c5451633bc02375316d25ed265f216a Mon Sep 17 00:00:00 2001
From: Junwhan Ahn <junwhan@google.com>
Date: Fri, 9 Oct 2026 01:28:10 -0700
Subject: [PATCH 0340/1000] Outline std::optional<NamedSharding> in
 xla::HloSharding to save memory.

`sizeof(NamedSharding)` is 232 bytes (and `sizeof(std::optional<NamedSharding>)` is 240 bytes with alignment). By outlining `named_sharding_` as `std::unique_ptr<NamedSharding>`, `sizeof(HloSharding)` is reduced from 400 bytes to 168 bytes, saving 232 bytes per instance when `NamedSharding` is unused.

PiperOrigin-RevId: 996370035
---
 third_party/xla/xla/hlo/ir/hlo_sharding.cc    | 42 ++++++++++++++++++
 third_party/xla/xla/hlo/ir/hlo_sharding.h     | 41 +++++++++---------
 .../xla/xla/service/hlo_sharding_test.cc      | 43 +++++++++++++++++++
 3 files changed, 106 insertions(+), 20 deletions(-)

diff --git a/third_party/xla/xla/hlo/ir/hlo_sharding.cc b/third_party/xla/xla/hlo/ir/hlo_sharding.cc
index a23874a76c3..df4ca150000 100644
--- a/third_party/xla/xla/hlo/ir/hlo_sharding.cc
+++ b/third_party/xla/xla/hlo/ir/hlo_sharding.cc
@@ -392,6 +392,48 @@ HloSharding HloSharding::Subgroup(
   return create_sharding(tile_assignment, subgroup_types, metadata);
 }
 
+HloSharding::HloSharding(const HloSharding& other)
+    : tile_assignment_(other.tile_assignment_),
+      tuple_elements_(other.tuple_elements_),
+      metadata_(other.metadata_),
+      subgroup_types_(other.subgroup_types_),
+      replicated_(other.replicated_),
+      single_device_(other.single_device_),
+      tuple_(other.tuple_),
+      manual_(other.manual_),
+      unknown_(other.unknown_),
+      unreduced_(other.unreduced_),
+      replicate_on_last_tile_dim_(other.replicate_on_last_tile_dim_),
+      shard_group_(other.shard_group_),
+      reduction_op_(other.reduction_op_),
+      named_sharding_(other.named_sharding_ ? std::make_unique<NamedSharding>(
+                                                  *other.named_sharding_)
+                                            : nullptr) {}
+
+HloSharding& HloSharding::operator=(const HloSharding& other) {
+  if (this == &other) {
+    return *this;
+  }
+  tile_assignment_ = other.tile_assignment_;
+  tuple_elements_ = other.tuple_elements_;
+  metadata_ = other.metadata_;
+  subgroup_types_ = other.subgroup_types_;
+  replicated_ = other.replicated_;
+  single_device_ = other.single_device_;
+  tuple_ = other.tuple_;
+  manual_ = other.manual_;
+  unknown_ = other.unknown_;
+  unreduced_ = other.unreduced_;
+  replicate_on_last_tile_dim_ = other.replicate_on_last_tile_dim_;
+  shard_group_ = other.shard_group_;
+  reduction_op_ = other.reduction_op_;
+  named_sharding_ =
+      other.named_sharding_
+          ? std::make_unique<NamedSharding>(*other.named_sharding_)
+          : nullptr;
+  return *this;
+}
+
 HloSharding HloSharding::Tuple(const ShapeTree<HloSharding>& sub_shardings) {
   std::vector<HloSharding> flattened_list;
   flattened_list.reserve(sub_shardings.leaf_count());
diff --git a/third_party/xla/xla/hlo/ir/hlo_sharding.h b/third_party/xla/xla/hlo/ir/hlo_sharding.h
index e302a6137af..c830d32fbc2 100644
--- a/third_party/xla/xla/hlo/ir/hlo_sharding.h
+++ b/third_party/xla/xla/hlo/ir/hlo_sharding.h
@@ -50,6 +50,11 @@ namespace xla {
 // computations.
 class HloSharding {
  public:
+  HloSharding(const HloSharding& other);
+  HloSharding(HloSharding&& other) = default;
+  HloSharding& operator=(const HloSharding& other);
+  HloSharding& operator=(HloSharding&& other) = default;
+
   // The name of the HLO instruction frontend attribute which stores that
   // instruction's sharding (e.g., Shardy).
   static inline constexpr absl::string_view kShardingFrontendAttrName =
@@ -199,13 +204,13 @@ class HloSharding {
   // Returns true if the sharding is represented using `NamedSharding` format.
   bool UseNamedShardingLeaf() const {
     DCHECK(!IsTuple());
-    return named_sharding_.has_value();
+    return named_sharding_ != nullptr;
   }
 
   ReductionOp reduction_op() const { return reduction_op_; }
   void set_reduction_op(ReductionOp op) {
     reduction_op_ = op;
-    if (named_sharding_.has_value()) {
+    if (named_sharding_ != nullptr) {
       named_sharding_->set_reduction_op(op);
     }
   }
@@ -531,7 +536,7 @@ class HloSharding {
                            bool overwrite) const;
 
   bool operator==(const HloSharding& other) const {
-    if (named_sharding_.has_value() == other.named_sharding_.has_value()) {
+    if ((named_sharding_ != nullptr) == (other.named_sharding_ != nullptr)) {
       return replicated_ == other.replicated_ &&
              single_device_ == other.single_device_ &&
              manual_ == other.manual_ && unknown_ == other.unknown_ &&
@@ -542,13 +547,14 @@ class HloSharding {
              subgroup_types_ == other.subgroup_types_ &&
              shard_group_ == other.shard_group_ &&
              reduction_op_ == other.reduction_op_ &&
-             named_sharding_ == other.named_sharding_;
+             (named_sharding_ == nullptr ||
+              *named_sharding_ == *other.named_sharding_);
     }
 
     // Compare two shardings regardless of their representation in order to
     // support mixed sharding representations in HLO
     // TODO (b/485319882): Compare NamedSharding's with different meshes
-    if (named_sharding_.has_value()) {
+    if (named_sharding_ != nullptr) {
       return V3ToV2Sharding(*named_sharding_) == other;
     }
     return *this == V3ToV2Sharding(*other.named_sharding_);
@@ -635,7 +641,7 @@ class HloSharding {
 
   const NamedSharding& named_sharding() const {
     CHECK(UseNamedShardingLeaf());
-    return named_sharding_.value();
+    return *named_sharding_;
   }
 
   // Returns the number of dimensions.
@@ -854,7 +860,8 @@ class HloSharding {
         unreduced_(false),
         replicate_on_last_tile_dim_(false),
         reduction_op_(named_sharding.reduction_op()),
-        named_sharding_(std::move(named_sharding)) {}
+        named_sharding_(
+            std::make_unique<NamedSharding>(std::move(named_sharding))) {}
 
  public:
   void ExtractReductionOpFromMetadata() {
@@ -881,8 +888,7 @@ class HloSharding {
         manual_(manual),
         unknown_(unknown),
         unreduced_(unreduced),
-        replicate_on_last_tile_dim_(false),
-        named_sharding_(std::nullopt) {
+        replicate_on_last_tile_dim_(false) {
     ExtractReductionOpFromMetadata();
   }
   // device_id values:
@@ -900,8 +906,7 @@ class HloSharding {
         manual_(false),
         unknown_(false),
         unreduced_(false),
-        replicate_on_last_tile_dim_(false),
-        named_sharding_(std::nullopt) {
+        replicate_on_last_tile_dim_(false) {
     ExtractReductionOpFromMetadata();
   }
   explicit HloSharding(TileAssignment tile_assignment,
@@ -915,8 +920,7 @@ class HloSharding {
         manual_(false),
         unknown_(false),
         unreduced_(false),
-        replicate_on_last_tile_dim_(replicate_on_last_tile_dim),
-        named_sharding_(std::nullopt) {
+        replicate_on_last_tile_dim_(replicate_on_last_tile_dim) {
     ExtractReductionOpFromMetadata();
   }
   explicit HloSharding(TileAssignment tile_assignment,
@@ -931,8 +935,7 @@ class HloSharding {
         manual_(false),
         unknown_(false),
         unreduced_(false),
-        replicate_on_last_tile_dim_(false),
-        named_sharding_(std::nullopt) {
+        replicate_on_last_tile_dim_(false) {
     ExtractReductionOpFromMetadata();
   }
   explicit HloSharding(std::vector<HloSharding> tuple_shardings)
@@ -943,8 +946,7 @@ class HloSharding {
         manual_(false),
         unknown_(false),
         unreduced_(false),
-        replicate_on_last_tile_dim_(false),
-        named_sharding_(std::nullopt) {}
+        replicate_on_last_tile_dim_(false) {}
 
   // Test-only constructor for sharding format code coverage. Copies the
   // original sharding with provided tile assignment.
@@ -959,8 +961,7 @@ class HloSharding {
         manual_(other.manual_),
         unknown_(other.unknown_),
         unreduced_(other.unreduced_),
-        replicate_on_last_tile_dim_(other.replicate_on_last_tile_dim_),
-        named_sharding_(std::nullopt) {
+        replicate_on_last_tile_dim_(other.replicate_on_last_tile_dim_) {
     CHECK(tile_assignment_ == other.tile_assignment_)
         << tile_assignment_.ToString() << " v.s. "
         << other.tile_assignment_.ToString();
@@ -1051,7 +1052,7 @@ class HloSharding {
   // Note that instead of reusing HloSharding's fields like metadata, we have
   // separate fields in NamedSharding to treat it as a standalone message which
   // is more clear and will help in future cleanup.
-  std::optional<NamedSharding> named_sharding_;
+  std::unique_ptr<NamedSharding> named_sharding_;
 };
 
 std::ostream& operator<<(std::ostream& out, const HloSharding& sharding);
diff --git a/third_party/xla/xla/service/hlo_sharding_test.cc b/third_party/xla/xla/service/hlo_sharding_test.cc
index 54458680c74..95b02c85e36 100644
--- a/third_party/xla/xla/service/hlo_sharding_test.cc
+++ b/third_party/xla/xla/service/hlo_sharding_test.cc
@@ -1495,5 +1495,48 @@ TEST_F(HloShardingTest, V3ToV2ToV3RoundTripSubAxes) {
   EXPECT_EQ(hlo_sharding, hlo_sharding_restored);
 }
 
+TEST_F(HloShardingTest, CopyAndMoveWithNamedSharding) {
+  EXPECT_EQ(sizeof(NamedSharding), 232);
+  EXPECT_EQ(sizeof(HloSharding), 168);
+
+  Mesh mesh({4}, {"axis_0"});
+  NamedSharding named_sharding =
+      test_utils::FromAxisNames(mesh, {{"axis_0:(1)2"}});
+  HloSharding sharding(named_sharding);
+  EXPECT_TRUE(sharding.UseNamedShardingLeaf());
+  EXPECT_EQ(sharding.named_sharding(), named_sharding);
+
+  // Test copy constructor
+  HloSharding copy_sharding = sharding;
+  EXPECT_TRUE(copy_sharding.UseNamedShardingLeaf());
+  EXPECT_EQ(copy_sharding.named_sharding(), named_sharding);
+  EXPECT_EQ(copy_sharding, sharding);
+
+  // Test copy assignment
+  HloSharding assign_sharding = HloSharding::Replicate();
+  EXPECT_FALSE(assign_sharding.UseNamedShardingLeaf());
+  assign_sharding = sharding;
+  EXPECT_TRUE(assign_sharding.UseNamedShardingLeaf());
+  EXPECT_EQ(assign_sharding.named_sharding(), named_sharding);
+  EXPECT_EQ(assign_sharding, sharding);
+
+  // Test self-assignment
+  HloSharding& sharding_ref = sharding;
+  sharding = sharding_ref;
+  EXPECT_TRUE(sharding.UseNamedShardingLeaf());
+  EXPECT_EQ(sharding.named_sharding(), named_sharding);
+
+  // Test move constructor
+  HloSharding move_sharding = std::move(copy_sharding);
+  EXPECT_TRUE(move_sharding.UseNamedShardingLeaf());
+  EXPECT_EQ(move_sharding.named_sharding(), named_sharding);
+
+  // Test move assignment
+  HloSharding move_assign_sharding = HloSharding::Replicate();
+  move_assign_sharding = std::move(assign_sharding);
+  EXPECT_TRUE(move_assign_sharding.UseNamedShardingLeaf());
+  EXPECT_EQ(move_assign_sharding.named_sharding(), named_sharding);
+}
+
 }  // namespace
 }  // namespace xla

From 75361be3567fe4eb138b701cd9e8425a37c61d0a Mon Sep 17 00:00:00 2001
From: Manjunath Gaonkar <magaonka@amd.com>
Date: Fri, 9 Oct 2026 02:03:23 -0700
Subject: [PATCH 0341/1000] PR #50101: [ROCm] Emit dynamic-slice copy thunks
 only on CUDA Ampere+
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50101

## 📝 Summary of Changes

Gate `AnalyzeDynamicSliceCopyFusion` / `EmitDynamicSliceCopyFusion` in `ThunkEmitter::EmitFusion` on CUDA Ampere+, so ROCm no longer emits `DynamicSliceFusionV2Thunk` for dynamic-slice / dynamic-update-slice copy fusions and falls back to the regular fusion emitter.

## 🎯 Justification

#44044 (`f4f4f9c984`) replaced `FusionDynamicMemcpyRewriter` with an analysis consumed directly by `ThunkEmitter::EmitFusion`. The rewriter lived in `FusionDispatchPipeline`, which `gpu_compiler.cc` only runs for CUDA Ampere+; the new call site in `EmitFusion` has no platform check, so since that commit ROCm also emits `DynamicSliceFusionV2Thunk`.

On ROCm that thunk cannot be recorded into a command buffer: `IsConvertible(DynamicSliceFusionV2Thunk)` requires
`SupportsMovedChildCommands`, which returns false for any non-CUDA device (`command_buffer_conversion_pass.cc`). Every DS/DUS copy in a `while` body (e.g. `lax.scan` over layers) therefore splits the enclosing command buffer.
The result is many more command-buffer launches per step, which measurably slows multi-GPU steps that have no collective in them.

## 🚀 Kind of Contribution

🐛 Bug Fix, ⚡️ Performance Improvement for ROCm ( basically restores lost perf ).

## 🧪 Unit Tests

- New `SupportsDynamicSliceCopyThunks(const se::DeviceDescription&)` in
  `dynamic_slice_copy.h` is used by both `ThunkEmitter::EmitFusion` and the
  `DynamicSliceCopyFusionAsyncWrapper` gate in `GpuCompiler`. `SupportsDynamicSliceCopyThunksTest` in
  `dynamic_slice_copy_test.cc` pins it: A100/H100 true, pre-Ampere CUDA and
  MI210/MI350 false.
- `GpuCopyTest.UseMemcpyForDynamicSlice` / `UseMemcpyForDynamicUpdateSlice`
  now check per platform: no kernel where copy thunks are emitted, a kernel otherwise. This pins the platform split.
- `UseMemcpyForDynamicUpdateSliceWithBitcasts`  `UseMemcpyForDynamicUpdateSliceWithOutOfBoundsConstantOffsets` and
  `DynamicUpdateSliceOffsetTable` check kernel counts only where copy thunks are emitted (`...WithBitcasts` and `...WithOutOfBoundsConstantOffsets` also check that a kernel is emitted otherwise), and compare numerics on every platform.

Copybara import of the project:

--
9969f6d153da15e28d1b5a91e1b5dad4de3baca5 by magaonka <magaonka@amd.com>:

[ROCm] Emit dynamic-slice copy thunks only on CUDA Ampere+

- Gate DS/DUS copy thunk emission in ThunkEmitter::EmitFusion on a new
  SupportsDynamicSliceCopyThunks() (CUDA Ampere+), matching the existing
  DynamicSliceCopyFusionAsyncWrapper gate in GpuCompiler.
- On ROCm, DynamicSliceFusionV2Thunk cannot be converted into a command
  buffer, so since openxla/xla#44044 every DS/DUS copy in a while body splits
  the enclosing command buffer.
- Tests: unit test for SupportsDynamicSliceCopyThunks; per-platform
  expectations for the gpu_copy_test DS/DUS memcpy tests.

Merging this change closes #50101

PiperOrigin-RevId: 996386569
---
 third_party/xla/xla/backends/gpu/tests/BUILD  |  1 +
 .../xla/backends/gpu/tests/gpu_copy_test.cc   | 85 ++++++++++++-------
 .../xla/xla/backends/gpu/transforms/BUILD     |  5 ++
 .../gpu/transforms/dynamic_slice_copy.cc      |  8 ++
 .../gpu/transforms/dynamic_slice_copy.h       | 11 +++
 .../gpu/transforms/dynamic_slice_copy_test.cc | 23 +++++
 third_party/xla/xla/service/gpu/BUILD         |  1 +
 .../xla/xla/service/gpu/gpu_compiler.cc       |  5 +-
 .../xla/xla/service/gpu/thunk_emitter.cc      | 10 ++-
 9 files changed, 112 insertions(+), 37 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/tests/BUILD b/third_party/xla/xla/backends/gpu/tests/BUILD
index a4bcc37e49a..a951c40b4e4 100644
--- a/third_party/xla/xla/backends/gpu/tests/BUILD
+++ b/third_party/xla/xla/backends/gpu/tests/BUILD
@@ -370,6 +370,7 @@ xla_test(
         "//xla:literal",
         "//xla:literal_util",
         "//xla:xla_proto_cc",
+        "//xla/backends/gpu/transforms:dynamic_slice_copy",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:verified_hlo_module",
         "//xla/tests:hlo_pjrt_interpreter_reference_mixin",
diff --git a/third_party/xla/xla/backends/gpu/tests/gpu_copy_test.cc b/third_party/xla/xla/backends/gpu/tests/gpu_copy_test.cc
index 5258c3d192d..224e2975008 100644
--- a/third_party/xla/xla/backends/gpu/tests/gpu_copy_test.cc
+++ b/third_party/xla/xla/backends/gpu/tests/gpu_copy_test.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include <gtest/gtest.h>
 #include "absl/status/status_matchers.h"
 #include "xla/backends/gpu/tests/gpu_pjrt_codegen_test.h"
+#include "xla/backends/gpu/transforms/dynamic_slice_copy.h"
 #include "xla/error_spec.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
@@ -43,6 +44,10 @@ class GpuCopyTest : public HloInterpreterReferenceMixin<GpuPjRtCodegenTest> {
         true);
     return debug_options;
   }
+
+  bool EmitsDynamicSliceCopyThunks() const {
+    return SupportsDynamicSliceCopyThunks(device_description());
+  }
 };
 
 // The GPU backend should not emit a copy kernel for the kCopy instruction in
@@ -133,9 +138,12 @@ TEST_F(GpuCopyTest, UseMemcpyForDynamicSlice) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> hlo_module,
                        ParseAndReturnVerifiedModule(kSliceMemcpyModule));
 
-  // There should not be a kernel for `dynamic_slice`.
+  // There should not be a kernel for `dynamic_slice`, except on platforms that
+  // emit DS/DUS copy fusions as regular kernels.
   ASSERT_OK(CompileAndVerifyIr(std::move(hlo_module),
-                               "; CHECK-NOT: void @slice",
+                               EmitsDynamicSliceCopyThunks()
+                                   ? "; CHECK-NOT: void @slice"
+                                   : "; CHECK: void @slice",
                                /*match_optimized_ir=*/false,
                                /*run_optimization_passes=*/false));
   EXPECT_TRUE(
@@ -164,7 +172,9 @@ TEST_F(GpuCopyTest, UseMemcpyForDynamicUpdateSlice) {
                        ParseAndReturnVerifiedModule(kDynamicUpdateSliceModule));
 
   ASSERT_OK(CompileAndVerifyIr(std::move(hlo_module),
-                               "; CHECK-NOT: void @updated",
+                               EmitsDynamicSliceCopyThunks()
+                                   ? "; CHECK-NOT: void @updated"
+                                   : "; CHECK: void @updated",
                                /*match_optimized_ir=*/false,
                                /*run_optimization_passes=*/false));
   EXPECT_TRUE(
@@ -240,17 +250,24 @@ TEST_F(GpuCopyTest, UseMemcpyForDynamicUpdateSliceWithBitcasts) {
       std::unique_ptr<VerifiedHloModule> hlo_module,
       ParseAndReturnVerifiedModule(kDynamicUpdateSliceWithBitcastModule));
 
-  ASSERT_OK(CompileAndVerifyIr(std::move(hlo_module), R"(
-    CHECK-NOT: define {{.*}}@
-    CHECK: define {{.*}}@input
-    CHECK-NOT: define {{.*}}@
-    CHECK: define {{.*}}@cmp
-    CHECK-NOT: define {{.*}}@
-    CHECK: define {{.*}}@next_ivar
-    CHECK-NOT: define {{.*}}@
-  )",
-                               /*match_optimized_ir=*/false,
-                               /*run_optimization_passes=*/false));
+  if (EmitsDynamicSliceCopyThunks()) {
+    ASSERT_OK(CompileAndVerifyIr(std::move(hlo_module), R"(
+      CHECK-NOT: define {{.*}}@
+      CHECK: define {{.*}}@input
+      CHECK-NOT: define {{.*}}@
+      CHECK: define {{.*}}@cmp
+      CHECK-NOT: define {{.*}}@
+      CHECK: define {{.*}}@next_ivar
+      CHECK-NOT: define {{.*}}@
+    )",
+                                 /*match_optimized_ir=*/false,
+                                 /*run_optimization_passes=*/false));
+  } else {
+    ASSERT_OK(CompileAndVerifyIr(std::move(hlo_module),
+                                 "; CHECK: define {{.*}}@updated_bc",
+                                 /*match_optimized_ir=*/false,
+                                 /*run_optimization_passes=*/false));
+  }
   EXPECT_TRUE(RunAndCompareNoHloPasses(kDynamicUpdateSliceWithBitcastModule,
                                        ErrorSpec{0, 0}));
 }
@@ -294,10 +311,9 @@ constexpr char kSliceMemcpyModuleUnfused[] = R"(
     })";
 
 TEST_F(GpuCopyTest, UseDynamicMemcpyIntegrationTest) {
-  auto compute_capability = device_description().gpu_compute_capability();
-  if (auto cc = compute_capability.cuda_compute_capability();
-      !cc || !cc->IsAtLeastAmpere()) {
-    GTEST_SKIP() << "Test requires at least Ampere.";
+  if (!EmitsDynamicSliceCopyThunks()) {
+    GTEST_SKIP()
+        << "Test requires a platform that emits dynamic-slice copy thunks.";
   }
 
   // This is an integration test to verify that dynamic-slices that depend on
@@ -351,10 +367,17 @@ TEST_F(GpuCopyTest,
       std::unique_ptr<VerifiedHloModule> hlo_module,
       ParseAndReturnVerifiedModule(kDUSOutOfBoundsConstantOffsetsModule));
 
-  ASSERT_OK(CompileAndVerifyIr(std::move(hlo_module),
-                               "; CHECK-NOT: define {{.*}}@",
-                               /*match_optimized_ir=*/false,
-                               /*run_optimization_passes=*/true));
+  if (EmitsDynamicSliceCopyThunks()) {
+    ASSERT_OK(CompileAndVerifyIr(std::move(hlo_module),
+                                 "; CHECK-NOT: define {{.*}}@",
+                                 /*match_optimized_ir=*/false,
+                                 /*run_optimization_passes=*/true));
+  } else {
+    ASSERT_OK(CompileAndVerifyIr(std::move(hlo_module),
+                                 "; CHECK: define {{.*}}@",
+                                 /*match_optimized_ir=*/false,
+                                 /*run_optimization_passes=*/true));
+  }
   EXPECT_TRUE(
       RunAndCompare(kDUSOutOfBoundsConstantOffsetsModule, ErrorSpec{0, 0}));
 }
@@ -395,14 +418,16 @@ TEST_F(GpuCopyTest, DynamicUpdateSliceOffsetTable) {
       ROOT result = s32[3,5] get-tuple-element(loop), index=1
     })";
 
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo));
-  // Only the loop condition and increment need kernels. The copy uses a linear
-  // source offset and destination offsets {0, 24, 48, 52, 56, 56}.
-  ASSERT_OK(CompileAndVerifyIr(std::move(module), R"(
-      CHECK-COUNT-2: define {{.*}}@
-      CHECK-NOT: define {{.*}}@)",
-                               /*match_optimized_ir=*/false,
-                               /*run_optimization_passes=*/true));
+  if (EmitsDynamicSliceCopyThunks()) {
+    ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo));
+    // Only the loop condition and increment need kernels. The copy uses a
+    // linear source offset and destination offsets {0, 24, 48, 52, 56, 56}.
+    ASSERT_OK(CompileAndVerifyIr(std::move(module), R"(
+        CHECK-COUNT-2: define {{.*}}@
+        CHECK-NOT: define {{.*}}@)",
+                                 /*match_optimized_ir=*/false,
+                                 /*run_optimization_passes=*/true));
+  }
   EXPECT_TRUE(RunAndCompare(hlo, ErrorSpec{0, 0}));
 }
 
diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD
index e6343db02bf..3ae8c44ac8d 100644
--- a/third_party/xla/xla/backends/gpu/transforms/BUILD
+++ b/third_party/xla/xla/backends/gpu/transforms/BUILD
@@ -306,6 +306,8 @@ cc_library(
         "//xla/codegen/xtile:xtile_config_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/service/gpu:backend_configs_cc",
+        "//xla/stream_executor:device_description",
+        "//xla/stream_executor/cuda:cuda_compute_capability",
         "@com_google_absl//absl/status:statusor",
     ],
 )
@@ -382,6 +384,9 @@ xla_cc_test(
         "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
         "//xla/service/gpu:backend_configs_cc",
+        "//xla/service/gpu:gpu_device_info_for_tests",
+        "//xla/stream_executor:device_description",
+        "//xla/stream_executor/cuda:cuda_compute_capability",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_googletest//:gtest_main",  # fixdeps: keep
     ],
diff --git a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_copy.cc b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_copy.cc
index 824faf1fa3b..761bab258f6 100644
--- a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_copy.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_copy.cc
@@ -33,6 +33,8 @@ limitations under the License.
 #include "xla/service/gpu/backend_configs.pb.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
+#include "xla/stream_executor/cuda/cuda_compute_capability.h"
+#include "xla/stream_executor/device_description.h"
 
 namespace xla::gpu {
 namespace {
@@ -310,4 +312,10 @@ bool IsDynamicSliceCopyFusion(const HloInstruction* instr) {
          IsCopyHeroDynamicSliceFusion(instr);
 }
 
+bool SupportsDynamicSliceCopyThunks(const se::DeviceDescription& device_info) {
+  const se::CudaComputeCapability* cuda_cc =
+      device_info.gpu_compute_capability().cuda_compute_capability();
+  return cuda_cc != nullptr && cuda_cc->IsAtLeastAmpere();
+}
+
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_copy.h b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_copy.h
index 2c84f400bae..bc435e16d66 100644
--- a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_copy.h
+++ b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_copy.h
@@ -25,9 +25,20 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/shape.h"
+#include "xla/stream_executor/device_description.h"
 
 namespace xla::gpu {
 
+// Returns true if the platform emits DS/DUS copy fusions as
+// DynamicSliceFusionV2Thunk copies. GpuCompiler (for
+// DynamicSliceCopyFusionAsyncWrapper) and ThunkEmitter both use this so the
+// HLO passes and thunk emission agree. Limited to CUDA Ampere+, the platforms
+// the replaced FusionDynamicMemcpyRewriter ran on. This is not the
+// command-buffer limit: converting the thunk into a command buffer also needs
+// CUDA 12.9+ (SupportsMovedChildCommands), so Ampere+ with an older CUDA still
+// splits the enclosing command buffer.
+bool SupportsDynamicSliceCopyThunks(const se::DeviceDescription& device_info);
+
 // Analysis result for a fusion that can be lowered as a DeviceToDeviceCopyThunk
 // wrapped by DynamicSliceFusionV2Thunk. The source and destination slices are
 // described by `parameters` and `results`; no HLO copy needs to exist.
diff --git a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_copy_test.cc b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_copy_test.cc
index a8ac480e01f..444d3e8661c 100644
--- a/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_copy_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/dynamic_slice_copy_test.cc
@@ -31,7 +31,10 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/service/gpu/backend_configs.pb.h"
+#include "xla/service/gpu/gpu_device_info_for_tests.h"
 #include "xla/shape_util.h"
+#include "xla/stream_executor/cuda/cuda_compute_capability.h"
+#include "xla/stream_executor/device_description.h"
 #include "xla/xla_data.pb.h"
 
 namespace xla::gpu {
@@ -448,5 +451,25 @@ TEST_F(DynamicSliceCopyTest, DynamicVariableUsesPerVariableInitStep) {
                                             {2, Offset::Constant(0)}}}));
 }
 
+TEST(SupportsDynamicSliceCopyThunksTest, CudaAmpereAndNewer) {
+  EXPECT_TRUE(
+      SupportsDynamicSliceCopyThunks(TestGpuDeviceInfo::A100SXMDeviceInfo()));
+  EXPECT_TRUE(
+      SupportsDynamicSliceCopyThunks(TestGpuDeviceInfo::H100SXMDeviceInfo()));
+}
+
+TEST(SupportsDynamicSliceCopyThunksTest, CudaPreAmpere) {
+  EXPECT_FALSE(
+      SupportsDynamicSliceCopyThunks(TestGpuDeviceInfo::A100SXMDeviceInfo(
+          se::GpuComputeCapability{se::CudaComputeCapability(7, 0)})));
+}
+
+TEST(SupportsDynamicSliceCopyThunksTest, Rocm) {
+  EXPECT_FALSE(
+      SupportsDynamicSliceCopyThunks(TestGpuDeviceInfo::AMDMI210DeviceInfo()));
+  EXPECT_FALSE(
+      SupportsDynamicSliceCopyThunks(TestGpuDeviceInfo::AMDMI350DeviceInfo()));
+}
+
 }  // namespace
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 3fa52c28f2c..b7ed894874e 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -1846,6 +1846,7 @@ cc_library(
         "//xla/backends/gpu/transforms:double_buffer_loop_unrolling",
         "//xla/backends/gpu/transforms:dus_accumulator_zero_init_elimination",
         "//xla/backends/gpu/transforms:dynamic_slice_annotator",
+        "//xla/backends/gpu/transforms:dynamic_slice_copy",
         "//xla/backends/gpu/transforms:dynamic_slice_copy_fusion_async_wrapper",
         "//xla/backends/gpu/transforms:dynamic_slice_fusion_rewriter_v2",
         "//xla/backends/gpu/transforms:estimate_cub_scan_scratch_size",
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index c8a77fcd615..599e4ab68f7 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -110,6 +110,7 @@ limitations under the License.
 #include "xla/backends/gpu/transforms/double_buffer_loop_unrolling.h"
 #include "xla/backends/gpu/transforms/dus_accumulator_zero_init_elimination.h"
 #include "xla/backends/gpu/transforms/dynamic_slice_annotator.h"
+#include "xla/backends/gpu/transforms/dynamic_slice_copy.h"
 #include "xla/backends/gpu/transforms/dynamic_slice_copy_fusion_async_wrapper.h"
 #include "xla/backends/gpu/transforms/dynamic_slice_fusion_rewriter_v2.h"
 #include "xla/backends/gpu/transforms/estimate_cub_scan_scratch_size.h"
@@ -3238,9 +3239,7 @@ absl::Status GpuCompiler::RunPreSchedulingPasses(
   tsl::profiler::TraceMe traceme("RunPreSchedulingPasses");
   HloPassPipeline pipeline("pre-scheduling-passes");
   pipeline.AddPass<FusionWrapper>(gpu_device_info);
-  const auto* cuda_cc =
-      gpu_device_info.gpu_compute_capability().cuda_compute_capability();
-  if (cuda_cc != nullptr && cuda_cc->IsAtLeastAmpere()) {
+  if (SupportsDynamicSliceCopyThunks(gpu_device_info)) {
     pipeline.AddPass<DynamicSliceCopyFusionAsyncWrapper>();
   }
   // GpuCopyAsyncWrapper is disabled when xla_gpu_async_copy_min_bytes is -1.
diff --git a/third_party/xla/xla/service/gpu/thunk_emitter.cc b/third_party/xla/xla/service/gpu/thunk_emitter.cc
index 27779eca3f3..285180bd727 100644
--- a/third_party/xla/xla/service/gpu/thunk_emitter.cc
+++ b/third_party/xla/xla/service/gpu/thunk_emitter.cc
@@ -1783,10 +1783,12 @@ Future<ThunkSequence> ThunkEmitter::EmitFusion(
     return EmitStaticSliceCopyFusion(instr, *static_copy);
   }
 
-  ABSL_ASSIGN_OR_RETURN(std::optional<DynamicSliceCopyFusion> dynamic_copy,
-                   AnalyzeDynamicSliceCopyFusion(instr));
-  if (dynamic_copy.has_value()) {
-    return EmitDynamicSliceCopyFusion(instr, std::move(*dynamic_copy));
+  if (SupportsDynamicSliceCopyThunks(ir_emitter_context_->gpu_device_info())) {
+    ABSL_ASSIGN_OR_RETURN(std::optional<DynamicSliceCopyFusion> dynamic_copy,
+                     AnalyzeDynamicSliceCopyFusion(instr));
+    if (dynamic_copy.has_value()) {
+      return EmitDynamicSliceCopyFusion(instr, std::move(*dynamic_copy));
+    }
   }
 
   analysis_garbage_collector_.push_back(

From afe922e2008ce705c81b5774295fb47c0c16444f Mon Sep 17 00:00:00 2001
From: Shaogang Wang <shawnw@nvidia.com>
Date: Fri, 9 Oct 2026 02:13:57 -0700
Subject: [PATCH 0342/1000] PR #50178: [XLA:TSL] BFC: grow spatially
 partitioned pools within a reserved arena
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50178

**Stacked PR 3/5** splitting #48985 into reviewable layers. Stacked on #50177: the base branch is `shawnw/bfc-growth-2-reserved-suballocator`, so this PR shows only its own commit. When #50177 lands, this PR is rebased onto `main` and re-targeted.

| # | PR | Title |
|---|---|---|
| 1 | #50176 | [XLA:GPU] StreamExecutor: VA reservation and physical allocation APIs sharing the executor's VMM options |
| 2 | #50177 | [XLA:GPU] DeviceMemAllocator: reserved-arena mode that maps physical memory on demand |
| 3 | #50178 | [XLA:TSL] BFC: grow spatially partitioned pools within a reserved arena |
| 4 | #50179 | [XLA:GPU] PJRT: typed MemFraction memory policy with START-CAP grammar |
| 5 | #50180 | [XLA:GPU] PJRT: enable BFC growth for growable memory fractions |

---

📝 Summary of Changes

With `Options::allow_growth` and `enable_spatial_partitioning` both set, BFC allocates `initial_region_bytes` up front and lets upper-end (default memory) requests extend the arena through the suballocator up to `total_memory`. Lower-end (collective) requests stay confined to the initial region so symmetric NCCL window offsets never move. Contiguous extensions coalesce with the region's free tail.

```text
base                               P                  M          C
| S(1) ---> | free gap | old S(0)  | free | <--- S(0) | unmapped |
                                   |<- new backing -->|
```

`Extend` computes the smallest extension that lets the request fit after coalescing with the tail (smaller than the request when the tail helps, larger when alignment padding is needed) and uses it as the floor of the capacity check and the backpedal loop, while the target stays the usual growth size. Attempts are sized in whole suballocator granules so a shrunken retry is never padded back to a size that just failed. A short extension that lands in a separate region is returned. The initial region backpedals like a fixed pool when memory is short, with a warning, and the lower-end limit is the size actually obtained. `AllocationRegion` tracks its last chunk, so no chunk-list walks remain.

🎯 Justification

Core of the growable pool. No production code sets `allow_growth` with spatial partitioning yet, so there is no behavior change.

🚀 Kind of Contribution

✨ New Feature (dormant until wired up)

🧪 Unit Tests:

`//xla/tsl/framework:bfc_allocator_test`: growth only after gap misses, fixed lower capacity, normal growth sizing, backpedal and failure recovery, contiguous coalescing across the old end, granularity-aware backpedal, growth at a physical limit below the virtual cap, non-adjacent short extensions, alignment padding, initial-region backpedal.

🧪 Execution Tests:

Covered end to end by the last change in the stack.

🤖 Generated with [Claude Code](https://claude.com/claude-code)
Copybara import of the project:

--
22924f9b059c2ad5d1b070356af5ea831962ce70 by Shaogang Wang <shawnw@nvidia.com>:

[XLA:TSL] BFC: grow spatially partitioned pools within a reserved arena

With Options::allow_growth and enable_spatial_partitioning both set, BFC
allocates initial_region_bytes up front and lets upper-end (default
memory) requests extend the arena through the suballocator up to
total_memory, while lower-end (collective) requests stay confined to the
initial region so symmetric NCCL window offsets never move. Contiguous
extensions coalesce with the region's free tail.

Extend computes the smallest extension that lets the request fit after
coalescing with the tail (it can be smaller than the request, or larger
when alignment padding is needed) and uses it as the floor of both the
capacity check and the backpedal loop; the target stays the usual growth
size. Attempts are sized in whole suballocator granules so a shrunken retry
is not padded back to a size that just failed. A short extension that lands
in a separate region is returned. The initial region backpedals like a
fixed pool when memory is short, with a warning, and the lower-end limit is
the size actually obtained. AllocationRegion tracks its last chunk so no
chunk walks are needed.

No production code sets allow_growth together with spatial partitioning
yet, so there is no behavior change. Third of five stacked changes (see
#48985).

--
673b92617e7096904689acfef9a23fbd075c3f50 by Shaogang Wang <shawnw@nvidia.com>:

[XLA:TSL] Address spatial BFC growth review comments

Validate extensions against their actual region tail or base alignment before
retaining backing, including extensions at least as large as the request.
Cache and validate suballocator granularity at construction, and clamp the
lower-end OOM diagnostic to the usable part of the initial shared region.

Add regression tests for misplaced extensions, invalid constructor settings,
and lower-end OOM diagnostics. The new tests fail before the fix; the full
bfc_allocator_test target passes in an optimized build on Prenyx.

Test: bazel test --config=cuda -c opt //xla/tsl/framework:bfc_allocator_test

Merging this change closes #50178

PiperOrigin-RevId: 996391820
---
 third_party/xla/xla/tsl/framework/BUILD       |   2 +
 .../xla/xla/tsl/framework/bfc_allocator.cc    | 278 ++++++--
 .../xla/xla/tsl/framework/bfc_allocator.h     |  90 ++-
 .../xla/tsl/framework/bfc_allocator_test.cc   | 632 +++++++++++++++++-
 4 files changed, 941 insertions(+), 61 deletions(-)

diff --git a/third_party/xla/xla/tsl/framework/BUILD b/third_party/xla/xla/tsl/framework/BUILD
index 4a80daad487..deb1b40e178 100644
--- a/third_party/xla/xla/tsl/framework/BUILD
+++ b/third_party/xla/xla/tsl/framework/BUILD
@@ -256,6 +256,7 @@ tsl_cc_test(
         ":allocator",
         ":bfc_allocator",
         ":scoped_allocation_trace",
+        ":shared_counter",
         "//xla/tsl/platform:env",
         "//xla/tsl/platform:env_impl",  # buildcleaner: keep
         "//xla/tsl/platform:test",
@@ -264,6 +265,7 @@ tsl_cc_test(
         "@com_google_absl//absl/base",
         "@com_google_absl//absl/base:log_severity",
         "@com_google_absl//absl/base:no_destructor",
+        "@com_google_absl//absl/log:globals",
         "@com_google_absl//absl/log:scoped_mock_log",
         "@com_google_absl//absl/synchronization",
     ],
diff --git a/third_party/xla/xla/tsl/framework/bfc_allocator.cc b/third_party/xla/xla/tsl/framework/bfc_allocator.cc
index 35740c28409..0dff58b92de 100644
--- a/third_party/xla/xla/tsl/framework/bfc_allocator.cc
+++ b/third_party/xla/xla/tsl/framework/bfc_allocator.cc
@@ -85,14 +85,21 @@ BFCAllocator::BFCAllocator(std::unique_ptr<SubAllocator> sub_allocator,
       free_chunk_tag_(opts.enable_spatial_partitioning ? ChunkTag::kCentralGap
                                                        : ChunkTag::kLower),
       coalesce_regions_(sub_allocator->SupportsCoalescing()),
+      allocation_granularity_(std::max(
+          sub_allocator->GetAllocationGranularity(), kMinAllocationSize)),
       sub_allocator_(std::move(sub_allocator)),
       name_(name),
       unused_chunk_handle_head_(kInvalidChunkHandle),
       next_allocation_id_(1) {
-  CHECK(!opts.enable_spatial_partitioning || !opts.allow_growth)  // Crash OK
-      << "Spatial partitioning requires a single fixed address range "
-         "(allow_growth=false).";
-  if (opts.allow_growth) {
+  CHECK_EQ(allocation_granularity_ % kMinAllocationSize, 0);
+  if (opts.enable_spatial_partitioning && opts.allow_growth) {
+    CHECK(!opts.garbage_collection) << "The shared region must remain stable";
+    CHECK_GT(opts.initial_region_bytes, 0);
+    CHECK_LE(opts.initial_region_bytes, total_memory);
+    CHECK_EQ(opts.initial_region_bytes % allocation_granularity_, 0)
+        << "initial_region_bytes must be a multiple of the granularity";
+    curr_region_allocation_bytes_ = opts.initial_region_bytes;
+  } else if (opts.allow_growth) {
     // 2MiB smallest initial allocation, unless total memory available
     // is less.
     curr_region_allocation_bytes_ =
@@ -167,52 +174,175 @@ const BFCAllocator::Chunk* BFCAllocator::ChunkFromHandle(ChunkHandle h) const {
   return &(chunks_[h]);
 }
 
-bool BFCAllocator::Extend(size_t alignment, size_t rounded_bytes) {
+BFCAllocator::ChunkHandle BFCAllocator::RegionTail(
+    const AllocationRegion& region) const {
+  const ChunkHandle tail = region.tail();
+  DCHECK_NE(tail, kInvalidChunkHandle);
+  const Chunk* c = ChunkFromHandle(tail);
+  DCHECK_EQ(c->next, kInvalidChunkHandle);
+  DCHECK_EQ(
+      static_cast<const void*>(static_cast<const char*>(c->ptr) + c->size),
+      region.end_ptr());
+  return tail;
+}
+
+size_t BFCAllocator::TailExtensionBytes(const AllocationRegion& region,
+                                        size_t alignment,
+                                        size_t rounded_bytes) const {
+  const Chunk* c = ChunkFromHandle(RegionTail(region));
+  // Memory mapped right after the region starts a chunk at end_ptr, unless the
+  // tail is free and can be merged into it: TryToCoalesce (ignore_freed_at =
+  // false) skips in-use and timestamped tails, and lower-end chunks never
+  // extend past the initial region.
+  uintptr_t chunk_start = reinterpret_cast<uintptr_t>(region.end_ptr());
+  size_t existing = 0;
+  if (!c->in_use() && c->freed_at_count == 0 && c->tag != ChunkTag::kLower) {
+    chunk_start = absl::bit_cast<uintptr_t>(c->ptr);
+    existing = c->size;
+  }
+  // An upper-end allocation fits once the chunk covers the alignment padding
+  // of its start plus the request. This can exceed rounded_bytes: an adjacent
+  // extension cannot be placed at an aligned address of its own.
+  const size_t padding = LowEndAlignmentPadding(chunk_start, alignment);
+  if (rounded_bytes > std::numeric_limits<size_t>::max() - padding) {
+    return std::numeric_limits<size_t>::max();
+  }
+  const size_t needed = padding + rounded_bytes;
+  return needed - std::min(needed, existing);
+}
+
+bool BFCAllocator::Extend(size_t alignment, size_t rounded_bytes,
+                          AllocationEnd allocation_end) {
+  const bool first_region = region_manager_.regions().empty();
+  if (opts_.enable_spatial_partitioning && !first_region &&
+      (!opts_.allow_growth || allocation_end == AllocationEnd::kLower)) {
+    return false;
+  }
+  const bool initial_spatial_region =
+      opts_.enable_spatial_partitioning && opts_.allow_growth && first_region;
+  if (initial_spatial_region) {
+    rounded_bytes = opts_.initial_region_bytes;
+  }
   size_t available_bytes = memory_limit_ - *stats_.pool_bytes;
   // Rounds available_bytes down to the nearest multiple of kMinAllocationSize.
   available_bytes = (available_bytes / kMinAllocationSize) * kMinAllocationSize;
 
+  // Size requests in whole granules, validated at construction after the
+  // suballocator has completed any reservation.
+  const size_t granularity = allocation_granularity_;
+
+  // The smallest allocation that can still serve this request. A contiguous
+  // spatial extension lands right after an existing region, so it may be
+  // smaller than rounded_bytes when it completes a free tail, or larger when
+  // the request's alignment leaves padding at the merged chunk's start.
+  // Physical memory can run out long before the reserved capacity does, so
+  // this floor applies to every extension rather than only at the cap. The
+  // initial shared region may shrink all the way down to one granule, like a
+  // fixed pool that backpedals when memory is short.
+  size_t min_bytes = rounded_bytes;
+  const bool may_complete_tail =
+      opts_.enable_spatial_partitioning && coalesce_regions_ && !first_region;
+  if (may_complete_tail) {
+    min_bytes = std::numeric_limits<size_t>::max();
+    for (const AllocationRegion& region : region_manager_.regions()) {
+      min_bytes = std::min(
+          min_bytes, TailExtensionBytes(region, alignment, rounded_bytes));
+    }
+    if (min_bytes == 0) {
+      // A tail that already fits would have been found by FindChunkPtr.
+      return false;
+    }
+  } else if (initial_spatial_region) {
+    min_bytes = granularity;
+  }
+  if (min_bytes > std::numeric_limits<size_t>::max() - (granularity - 1)) {
+    return false;
+  }
+  const size_t floor_bytes =
+      (min_bytes + granularity - 1) / granularity * granularity;
+
   // Do we have enough space to handle the client's request?
   // If not, fail immediately.
-  if (rounded_bytes > available_bytes) {
+  if (floor_bytes > available_bytes) {
     return false;
   }
 
   // If curr_region_allocation_bytes_ is not enough to satisfy the
   // allocation, keep multiplying by a power of two until that is
-  // sufficient.
+  // sufficient. Aim for the full request but never past the remaining
+  // capacity, so repeated over-cap requests cannot inflate the region size.
+  const size_t growth_target =
+      std::min(std::max(rounded_bytes, min_bytes), available_bytes);
   bool increased_allocation = false;
-  while (rounded_bytes > curr_region_allocation_bytes_) {
+  while (growth_target > curr_region_allocation_bytes_) {
     curr_region_allocation_bytes_ *= 2;
     increased_allocation = true;
   }
 
   // Try allocating.
   size_t bytes = std::min(curr_region_allocation_bytes_, available_bytes);
+  bytes = std::max(bytes - bytes % granularity, floor_bytes);
   size_t bytes_received;
   void* mem_addr = sub_allocator_->Alloc(alignment, bytes, &bytes_received);
   if (mem_addr == nullptr) {
     static constexpr float kBackpedalFactor = 0.9;
 
-    // Try allocating less memory.
+    // Try allocating less memory, stepping down in whole granules so that a
+    // retry is never padded back up to a size that already failed.
     while (mem_addr == nullptr) {
-      size_t backpedal_bytes = RoundedBytes(bytes * kBackpedalFactor);
-      // RoundedBytes rounds up to a multiple of kMinAllocationSize, so for
-      // small sizes (bytes <= 10 * kMinAllocationSize) the backpedal can round
-      // back up to the same value. Force strict progress so a persistently
-      // failing sub-allocator cannot make this loop spin forever.
+      size_t backpedal_bytes = static_cast<size_t>(bytes * kBackpedalFactor);
+      backpedal_bytes -= backpedal_bytes % granularity;
+      // Force strict progress so a persistently failing sub-allocator cannot
+      // make this loop spin forever. bytes >= floor_bytes >= granularity here.
       if (backpedal_bytes >= bytes) {
-        backpedal_bytes = bytes - kMinAllocationSize;
+        backpedal_bytes = bytes - granularity;
       }
-      bytes = backpedal_bytes;
-      if (bytes < rounded_bytes) {
+      if (backpedal_bytes < floor_bytes) {
         return false;
       }
+      bytes = backpedal_bytes;
       mem_addr = sub_allocator_->Alloc(alignment, bytes, &bytes_received);
     }
   }
 
-  if (!increased_allocation) {
+  if (opts_.enable_spatial_partitioning && bytes_received > available_bytes) {
+    sub_allocator_->Free(mem_addr, bytes_received);
+    return false;
+  }
+  if (may_complete_tail) {
+    // The minimum above may describe a different region. Check the actual
+    // address, including alignment padding even when we got the full request.
+    // SupportsCoalescing does not promise adjacency.
+    const AllocationRegion* extended = region_manager_.RegionEndingAt(mem_addr);
+    const size_t padding =
+        LowEndAlignmentPadding(absl::bit_cast<uintptr_t>(mem_addr), alignment);
+    const bool fits = extended != nullptr
+                          ? TailExtensionBytes(*extended, alignment,
+                                               rounded_bytes) <= bytes_received
+                          : padding <= bytes_received &&
+                                rounded_bytes <= bytes_received - padding;
+    if (!fits) {
+      sub_allocator_->Free(mem_addr, bytes_received);
+      return false;
+    }
+  }
+  if (initial_spatial_region) {
+    spatial_region_start_ = absl::bit_cast<uintptr_t>(mem_addr);
+    // The lower end is confined to what was actually obtained. Ranks that end
+    // up with different initial sizes keep identical collective offsets from
+    // the base, but their collective capacity differs.
+    spatial_region_bytes_ = bytes_received;
+    if (bytes_received < opts_.initial_region_bytes) {
+      LOG(WARNING) << Name() << ": initial shared region is "
+                   << strings::HumanReadableNumBytes(bytes_received)
+                   << " instead of the configured "
+                   << strings::HumanReadableNumBytes(opts_.initial_region_bytes)
+                   << "; collective (lower-end) capacity is limited to it.";
+    }
+    // After preallocation, use BFC's normal growth sizing (starting at 2 MiB).
+    curr_region_allocation_bytes_ =
+        RoundedBytes(std::min(memory_limit_, size_t{2 << 20}));
+  } else if (!increased_allocation) {
     // Increase the region size of the next required allocation.
     curr_region_allocation_bytes_ *= 2;
   }
@@ -252,28 +382,30 @@ bool BFCAllocator::Extend(size_t alignment, size_t rounded_bytes) {
   c->prev = kInvalidChunkHandle;
   c->next = kInvalidChunkHandle;
   c->freed_at_count = 0;
-  c->tag = free_chunk_tag_;
+  c->tag = opts_.enable_spatial_partitioning && !first_region ? ChunkTag::kUpper
+                                                              : free_chunk_tag_;
   c->allocation_annotation.reset();
 
   region_manager_.set_handle(c->ptr, h);
 
-  // If the region was extended, then there exists a previous chunk that should
-  // be linked to the new chunk.
+  // If the region was extended, link the new chunk after its previous tail.
+  // The region's end_ptr has already moved, so check against mem_addr.
   if (maybe_extended_region != nullptr) {
-    ChunkHandle prev =
-        maybe_extended_region->get_handle(maybe_extended_region->ptr());
+    ChunkHandle prev = maybe_extended_region->tail();
     BFCAllocator::Chunk* prev_chunk = ChunkFromHandle(prev);
-    // Find the last recorded chunk in the extended region.
-    while (prev_chunk->next != kInvalidChunkHandle) {
-      prev = prev_chunk->next;
-      prev_chunk = ChunkFromHandle(prev);
-    }
+    DCHECK_EQ(prev_chunk->next, kInvalidChunkHandle);
+    DCHECK_EQ(static_cast<void*>(static_cast<char*>(prev_chunk->ptr) +
+                                 prev_chunk->size),
+              mem_addr);
     c->prev = prev;
     prev_chunk->next = h;
   }
+  // Whether appended or in a new region, the new chunk is now the last one.
+  region_manager_.set_tail(c->ptr, h);
 
   // Maybe merge adjacent chunks and insert the chunk into the right free
-  // structure. In spatial mode, a fresh region becomes the central gap.
+  // structure. Contiguous spatial extensions can merge with the central gap;
+  // lower-end allocations still obey the initial region limit.
   InsertFreeChunk(TryToCoalesce(h, /*ignore_freed_at=*/false));
 
   return true;
@@ -546,6 +678,13 @@ void* BFCAllocator::AllocateRawInternal(size_t alignment, size_t num_bytes,
   BinNum bin_num = BinNumForSize(rounded_bytes);
 
   absl::MutexLock l(mutex_);
+  // Keep the initial shared region's size fixed even if the first request is
+  // larger. An upper request can subsequently use Extend's ordinary fallback.
+  if (opts_.enable_spatial_partitioning && opts_.allow_growth &&
+      region_manager_.regions().empty() &&
+      !Extend(alignment, opts_.initial_region_bytes, allocation_end)) {
+    return nullptr;
+  }
   if (ABSL_PREDICT_FALSE(!timestamped_chunks_.empty())) {
     // Merge timestamped chunks whose counts have become safe for general use.
     MergeTimestampedChunks(0);
@@ -558,7 +697,7 @@ void* BFCAllocator::AllocateRawInternal(size_t alignment, size_t num_bytes,
   }
 
   // Try to extend
-  if (Extend(alignment, rounded_bytes)) {
+  if (Extend(alignment, rounded_bytes, allocation_end)) {
     ptr = FindChunkPtr(bin_num, rounded_bytes, num_bytes, alignment,
                        freed_before, allocation_end);
     if (ptr != nullptr) {
@@ -587,7 +726,7 @@ void* BFCAllocator::AllocateRawInternal(size_t alignment, size_t num_bytes,
   // try deallocating free regions so that suballocator can combine them with
   // the unallocated bytes and form a larger region.
   if (DeallocateFreeRegions(rounded_bytes) &&
-      Extend(alignment, rounded_bytes)) {
+      Extend(alignment, rounded_bytes, allocation_end)) {
     ptr = FindChunkPtr(bin_num, rounded_bytes, num_bytes, alignment,
                        freed_before, allocation_end);
     if (ptr != nullptr) {
@@ -783,7 +922,18 @@ void* BFCAllocator::FindChunkPtrInCentralGap(size_t rounded_bytes,
     }
   } else {
     const size_t align_padding = LowEndAlignmentPadding(chunk_start, alignment);
-    if (ABSL_PREDICT_FALSE(chunk->size < rounded_bytes + align_padding)) {
+    size_t available = chunk->size;
+    if (opts_.allow_growth) {
+      // Coalescing may grow the gap into the extension. Only its prefix within
+      // the initial shared region is available to the lower end.
+      const uintptr_t offset = chunk_start - spatial_region_start_;
+      if (offset >= spatial_region_bytes_) {
+        return nullptr;
+      }
+      available = std::min(available, spatial_region_bytes_ - offset);
+    }
+    if (ABSL_PREDICT_FALSE(align_padding > available ||
+                           rounded_bytes > available - align_padding)) {
       return nullptr;
     }
   }
@@ -794,10 +944,10 @@ void* BFCAllocator::FindChunkPtrInCentralGap(size_t rounded_bytes,
              : AllocateChunkFromLowEnd(h, rounded_bytes, num_bytes, alignment);
 }
 
-bool BFCAllocator::ShouldSplitChunk(const Chunk* chunk, size_t rounded_bytes,
+bool BFCAllocator::ShouldSplitChunk(size_t chunk_size, size_t rounded_bytes,
                                     size_t alignment_padding) const {
-  DCHECK_LE(alignment_padding, chunk->size);
-  const size_t aligned_size = chunk->size - alignment_padding;
+  DCHECK_LE(alignment_padding, chunk_size);
+  const size_t aligned_size = chunk_size - alignment_padding;
   // After excluding the unavoidable alignment prefix, split if the remainder
   // is at least as large as the request (so another similar request could use
   // it), or retaining it as padding would reach the internal-fragmentation
@@ -838,9 +988,22 @@ void* BFCAllocator::AllocateChunkFromLowEnd(ChunkHandle h, size_t rounded_bytes,
       (chunk->tag == ChunkTag::kCentralGap
            ? policy.gap_split_policy
            : policy.hole_split_policy) == SplitPolicy::kExact;
-  if (split_exactly ? chunk->size > rounded_bytes
-                    : ShouldSplitChunk(chunk, rounded_bytes)) {
-    SplitChunk(h, rounded_bytes);
+  size_t available = chunk->size;
+  if (opts_.enable_spatial_partitioning && opts_.allow_growth) {
+    const uintptr_t offset =
+        absl::bit_cast<uintptr_t>(chunk->ptr) - spatial_region_start_;
+    DCHECK_LT(offset, spatial_region_bytes_);
+    available = std::min(available, spatial_region_bytes_ - offset);
+  }
+  // Apply the lower policy to its eligible capacity only. Even heuristic
+  // padding must stop at the fixed lower-end limit; the suffix remains free
+  // for upper allocations, which may span physical mapping boundaries.
+  const size_t allocation_size =
+      split_exactly || ShouldSplitChunk(available, rounded_bytes)
+          ? rounded_bytes
+          : available;
+  if (chunk->size > allocation_size) {
+    SplitChunk(h, allocation_size);
     chunk = ChunkFromHandle(h);  // Update chunk pointer in case it moved.
   }
 
@@ -871,7 +1034,7 @@ void* BFCAllocator::AllocateChunkFromHighEnd(ChunkHandle h,
   // just as the lower-end path does. Otherwise carve from the high end.
   if (!split_exactly) {
     const size_t align_padding = LowEndAlignmentPadding(chunk_start, alignment);
-    if (!ShouldSplitChunk(chunk, rounded_bytes, align_padding)) {
+    if (!ShouldSplitChunk(chunk->size, rounded_bytes, align_padding)) {
       aligned_start = chunk_start + align_padding;
     }
   }
@@ -895,7 +1058,7 @@ void* BFCAllocator::AllocateChunkFromHighEnd(ChunkHandle h,
   // before splitting so the suffix inherits kUpper instead of kCentralGap.
   chunk->tag = ChunkTag::kUpper;
   if (split_exactly ? chunk->size > rounded_bytes
-                    : ShouldSplitChunk(chunk, rounded_bytes)) {
+                    : ShouldSplitChunk(chunk->size, rounded_bytes)) {
     SplitChunk(h, rounded_bytes);
     chunk = ChunkFromHandle(h);
   }
@@ -993,6 +1156,8 @@ void BFCAllocator::SplitChunk(BFCAllocator::ChunkHandle h, size_t num_bytes) {
   if (h_neighbor != kInvalidChunkHandle) {
     Chunk* c_neighbor = ChunkFromHandle(h_neighbor);
     c_neighbor->prev = h_new_chunk;
+  } else {
+    region_manager_.set_tail(new_chunk->ptr, h_new_chunk);
   }
 
   // Add the newly free chunk to the appropriate free structure.
@@ -1049,7 +1214,8 @@ BFCAllocator::ChunkTag BFCAllocator::MergedChunkTag(ChunkTag a,
   // Two free holes with the same tag keep that tag (an interior hole still
   // belongs to its end). Any other combination -- a hole merging with the
   // central gap, or lower-end and upper-end holes meeting after the gap is
-  // exhausted -- yields a kCentralGap span reusable by either end.
+  // exhausted -- yields a kCentralGap span reusable by either end, subject to
+  // its allocation address limit.
   return a == b ? a : ChunkTag::kCentralGap;
 }
 
@@ -1076,6 +1242,8 @@ void BFCAllocator::MergeChunks(BFCAllocator::ChunkHandle h1,
   if (h3 != kInvalidChunkHandle) {
     BFCAllocator::Chunk* c3 = ChunkFromHandle(h3);
     c3->prev = h1;
+  } else {
+    region_manager_.set_tail(c1->ptr, h1);
   }
 
   // Set the new size
@@ -1131,6 +1299,16 @@ void BFCAllocator::ReturnBoundaryChunkToGap(BFCAllocator::ChunkHandle h) {
   }
   Chunk* c = ChunkFromHandle(h);
   CHECK(!c->in_use());  // Crash OK
+  // Only the region containing the original shared allocation can have a
+  // central gap. That region may extend past the fixed lower-end limit, but
+  // separate regions stay upper-owned even if they support coalescing.
+  if (opts_.allow_growth) {
+    const uintptr_t region_start =
+        reinterpret_cast<uintptr_t>(region_manager_.RegionStart(c->ptr));
+    if (region_start != spatial_region_start_) {
+      return;
+    }
+  }
   if (ABSL_PREDICT_TRUE(c->tag == ChunkTag::kLower)) {
     ChunkHandle n = c->next;
     if (n == kInvalidChunkHandle ||
@@ -1509,8 +1687,19 @@ void BFCAllocator::DumpMemoryLog(size_t num_bytes,
                                  : ChunkFromHandle(central_gap_)->size;
   size_t largest_free_chunk =
       std::max(largest_binned_free_chunk, central_gap_bytes);
-  size_t largest_compatible_free_chunk =
-      std::max(largest_compatible_binned_free_chunk, central_gap_bytes);
+  size_t compatible_central_gap_bytes = central_gap_bytes;
+  if (central_gap_ != kInvalidChunkHandle && opts_.allow_growth &&
+      allocation_end == AllocationEnd::kLower) {
+    const uintptr_t offset =
+        absl::bit_cast<uintptr_t>(ChunkFromHandle(central_gap_)->ptr) -
+        spatial_region_start_;
+    compatible_central_gap_bytes =
+        offset >= spatial_region_bytes_
+            ? 0
+            : std::min(central_gap_bytes, spatial_region_bytes_ - offset);
+  }
+  size_t largest_compatible_free_chunk = std::max(
+      largest_compatible_binned_free_chunk, compatible_central_gap_bytes);
 
   LOG(INFO) << "Sum Total of in-use chunks: "
             << strings::HumanReadableNumBytes(total_bytes);
@@ -1661,7 +1850,9 @@ BFCAllocator::get_bin_debug_info() {
   std::array<BinDebugInfo, kNumBins> bin_infos;
   for (const auto& region : region_manager_.regions()) {
     ChunkHandle h = region_manager_.get_handle(region.ptr());
+    ChunkHandle last = kInvalidChunkHandle;
     while (h != kInvalidChunkHandle) {
+      last = h;
       const Chunk* c = ChunkFromHandle(h);
       if (!c->in_use() && c->tag == ChunkTag::kCentralGap &&
           c->bin_num == kInvalidBinNum) {
@@ -1684,6 +1875,7 @@ BFCAllocator::get_bin_debug_info() {
       }
       h = c->next;
     }
+    DCHECK_EQ(last, region.tail());
   }
   return bin_infos;
 }
diff --git a/third_party/xla/xla/tsl/framework/bfc_allocator.h b/third_party/xla/xla/tsl/framework/bfc_allocator.h
index 1505471bd9a..a24ae05ae62 100644
--- a/third_party/xla/xla/tsl/framework/bfc_allocator.h
+++ b/third_party/xla/xla/tsl/framework/bfc_allocator.h
@@ -62,7 +62,7 @@ using tensorflow::MemoryDump;
 // - Backing memory comes from the SubAllocator as AllocationRegions. With
 //   Options::allow_growth=true the allocator grows by adding regions up to
 //   total_memory; with Options::allow_growth=false it reserves one fixed region
-//   during construction. stats_.bytes_reserved tracks bytes held from the
+//   on the first allocation. stats_.bytes_reserved tracks bytes held from the
 //   SubAllocator, while stats_.bytes_in_use tracks bytes currently live for
 //   clients.
 //
@@ -83,8 +83,7 @@ using tensorflow::MemoryDump;
 //   mode all requests use AllocationEnd::kLower, and ordinary free chunks stay
 //   in ChunkTag::kLower.
 //
-// - With Options::enable_spatial_partitioning=true, which requires
-//   Options::allow_growth=false, the fixed address range is split into
+// - With Options::enable_spatial_partitioning=true, the initial region has
 //   lower-end ownership, one central gap, and upper-end ownership.
 //   AllocationEnd::kLower requests grow upward, and AllocationEnd::kUpper
 //   requests grow downward. ChunkTag records ownership: kLower and kUpper for
@@ -99,6 +98,13 @@ using tensorflow::MemoryDump;
 //   exact gap splitting has placements independent of activity from the other
 //   end, except when the central gap cannot satisfy its requests.
 //
+// - Spatial growth uses the usual Extend fallback after neither an owned hole
+//   nor the central gap fits. Adjacent regions can merge when the existing
+//   suballocator supports coalescing; otherwise, growth adds separate
+//   upper-only regions. Lower-end allocations, including padding, always stay
+//   inside initial_region_bytes, even when a coalesced free span crosses that
+//   limit.
+//
 class BFCAllocator : public Allocator {
  public:
   // Address preference for equal-size holes owned by the same end. Size remains
@@ -186,9 +192,20 @@ class BFCAllocator : public Allocator {
     // collective buffers across ranks. Heuristic gap splitting can make chunk
     // sizes depend on the opposite end's activity through the gap size.
     //
-    // Requires allow_growth=false (a single fixed address range).
+    // With allow_growth=true, only upper-end requests can extend the arena.
+    // Adjacent allocations may coalesce when SupportsCoalescing() permits it;
+    // other allocations become upper-only regions. Lower requests are bounded
+    // by initial_region_bytes, even when the central free span extends past it.
+    // This mode requires initial_region_bytes and garbage_collection=false.
     bool enable_spatial_partitioning = false;
 
+    // Size of the initial shared region for growable spatial arenas. It is
+    // requested on the first allocation and, like a fixed pool, shrinks if
+    // the suballocator cannot provide that much; the lower end is then
+    // confined to what was obtained. total_memory remains the cap across all
+    // regions. Must be a multiple of the suballocator granularity.
+    size_t initial_region_bytes = 0;
+
     // Policies for owned-hole reuse and gap carves, independent of placement
     // direction. By default, owned holes use the BFC heuristic and gap carves
     // split exactly.
@@ -512,12 +529,19 @@ class BFCAllocator : public Allocator {
     void set_handle(const void* p, ChunkHandle h) { handles_[IndexFor(p)] = h; }
     void erase(const void* p) { set_handle(p, kInvalidChunkHandle); }
 
+    // The last chunk in the region (its `next` is kInvalidChunkHandle).
+    // Maintained by Extend, SplitChunk and MergeChunks so that extending a
+    // region does not require walking its chunk list.
+    ChunkHandle tail() const { return tail_; }
+    void set_tail(ChunkHandle h) { tail_ = h; }
+
    private:
     void Swap(AllocationRegion* other) {
       std::swap(ptr_, other->ptr_);
       std::swap(memory_size_, other->memory_size_);
       std::swap(end_ptr_, other->end_ptr_);
       std::swap(handles_, other->handles_);
+      std::swap(tail_, other->tail_);
     }
 
     size_t IndexFor(const void* p) const {
@@ -538,6 +562,8 @@ class BFCAllocator : public Allocator {
     // for the memory allocation represented by "p"
     std::vector<ChunkHandle> handles_;
 
+    ChunkHandle tail_ = kInvalidChunkHandle;
+
     AllocationRegion(const AllocationRegion&) = delete;
     void operator=(const AllocationRegion&) = delete;
   };
@@ -604,8 +630,26 @@ class BFCAllocator : public Allocator {
     }
     void erase(const void* p) { return MutableRegionFor(p)->erase(p); }
 
+    // Records `h` as the last chunk of the region containing `p`.
+    void set_tail(const void* p, ChunkHandle h) {
+      MutableRegionFor(p)->set_tail(h);
+    }
+
     const std::vector<AllocationRegion>& regions() const { return regions_; }
 
+    void* RegionStart(const void* p) const { return RegionFor(p)->ptr(); }
+
+    // Returns the region whose end_ptr() is `p`, i.e. the one that
+    // AddOrExtendAllocationRegion(p, ...) would extend, or nullptr.
+    const AllocationRegion* RegionEndingAt(const void* p) const {
+      auto entry =
+          std::upper_bound(regions_.begin(), regions_.end(), p, &Comparator);
+      if (entry != regions_.begin() && (entry - 1)->end_ptr() == p) {
+        return &*(entry - 1);
+      }
+      return nullptr;
+    }
+
    private:
     static bool Comparator(const void* ptr, const AllocationRegion& other) {
       return ptr < other.end_ptr();
@@ -654,8 +698,25 @@ class BFCAllocator : public Allocator {
 
   // Try to add a new memory region that can satisfy an allocation of
   // 'rounded_bytes' bytes.  Returns true on success and false on
-  // failure.
-  bool Extend(size_t alignment, size_t rounded_bytes)
+  // failure. Requests are sized in multiples of the suballocator's
+  // allocation granularity. A contiguous spatial extension may be smaller
+  // than 'rounded_bytes' when it only has to complete a free tail chunk.
+  bool Extend(size_t alignment, size_t rounded_bytes,
+              AllocationEnd allocation_end)
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+
+  // The last chunk of `region`, validated against the region bounds in debug
+  // builds. Must not be called between AllocationRegion::extend and the
+  // linking of the new chunk.
+  ChunkHandle RegionTail(const AllocationRegion& region) const
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+
+  // Bytes that an extension adjacent to `region` must add so that an
+  // allocation of `rounded_bytes` aligned to `alignment` fits in the chunk
+  // formed by merging the extension with the region's free tail. Includes
+  // alignment padding even when the tail cannot be merged into.
+  size_t TailExtensionBytes(const AllocationRegion& region, size_t alignment,
+                            size_t rounded_bytes) const
       ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
 
   // Deallocate free regions to give back the memory to suballocator, so that
@@ -697,7 +758,7 @@ class BFCAllocator : public Allocator {
   // request size, or if keeping it whole would waste at least
   // max_internal_fragmentation_bytes_ on padding.
   // alignment_padding excludes a prefix that must be split off for alignment.
-  bool ShouldSplitChunk(const Chunk* chunk, size_t rounded_bytes,
+  bool ShouldSplitChunk(size_t chunk_size, size_t rounded_bytes,
                         size_t alignment_padding = 0) const
       ABSL_EXCLUSIVE_LOCKS_REQUIRED(mutex_);
 
@@ -804,6 +865,13 @@ class BFCAllocator : public Allocator {
   // Structures immutable after construction
   size_t memory_limit_ = 0;
 
+  // Base and size of the initial shared region in growable spatial mode. The
+  // size is opts_.initial_region_bytes unless the suballocator could only
+  // provide less, and remains the lower-end allocation limit. Contiguous
+  // growth may merge free chunks across that limit for upper use.
+  uintptr_t spatial_region_start_ ABSL_GUARDED_BY(mutex_) = 0;
+  size_t spatial_region_bytes_ ABSL_GUARDED_BY(mutex_) = 0;
+
   // Maximum bytes a chunk may exceed the requested size before it is split, to
   // bound internal fragmentation. Derived from Options::fragmentation_fraction
   // and memory_limit_ once at construction.
@@ -828,8 +896,8 @@ class BFCAllocator : public Allocator {
   const Options opts_;
 
   // Tag assigned to newly-created free chunks. Non-partitioned BFC keeps
-  // ordinary free chunks in kLower; spatial partitioning starts each fixed
-  // region as the kCentralGap span.
+  // ordinary free chunks in kLower; spatial partitioning starts its initial
+  // region as the kCentralGap span. Additional spatial regions use kUpper.
   const ChunkTag free_chunk_tag_;
 
   // The size of the current region allocation.
@@ -842,6 +910,10 @@ class BFCAllocator : public Allocator {
   // device's address space).
   const bool coalesce_regions_;
 
+  // Suballocator configuration, including any VA reservation, must be complete
+  // before BFC construction. Region sizes also obey kMinAllocationSize.
+  const size_t allocation_granularity_;
+
   std::unique_ptr<SubAllocator> sub_allocator_;
   std::string name_;
   SharedCounter* timing_counter_ = nullptr;
diff --git a/third_party/xla/xla/tsl/framework/bfc_allocator_test.cc b/third_party/xla/xla/tsl/framework/bfc_allocator_test.cc
index 8f70e5f2d54..1e5829493b3 100644
--- a/third_party/xla/xla/tsl/framework/bfc_allocator_test.cc
+++ b/third_party/xla/xla/tsl/framework/bfc_allocator_test.cc
@@ -31,10 +31,12 @@ limitations under the License.
 #include "absl/base/casts.h"
 #include "absl/base/log_severity.h"
 #include "absl/base/no_destructor.h"
+#include "absl/log/globals.h"
 #include "absl/log/scoped_mock_log.h"
 #include "absl/synchronization/blocking_counter.h"
 #include "xla/tsl/framework/allocator.h"
 #include "xla/tsl/framework/scoped_allocation_trace.h"
+#include "xla/tsl/framework/shared_counter.h"
 #include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/test.h"
 #include "xla/tsl/platform/test_benchmark.h"
@@ -108,6 +110,61 @@ class FakeSubAllocator : public SubAllocator {
   uintptr_t next_ = kBase;
 };
 
+// Exercise both adjacent, coalescible regions and separate regions below the
+// initial shared allocation. No real memory is allocated.
+class SpatialGrowthSubAllocator : public FakeSubAllocator {
+ public:
+  // A reserved arena maps each extension right after the previous one, no
+  // matter which alignment the request asked for; mimic that when adjacent.
+  explicit SpatialGrowthSubAllocator(bool adjacent = false)
+      : FakeSubAllocator(
+            adjacent ? std::optional<size_t>(BFCAllocator::kMinAllocationSize)
+                     : std::nullopt),
+        supports_coalescing(adjacent),
+        adjacent_(adjacent) {}
+
+  void* Alloc(size_t alignment, size_t num_bytes,
+              size_t* bytes_received) override {
+    requests.push_back(num_bytes);
+    *bytes_received = 0;
+    // Like a VMM-backed suballocator, pad the request to the granularity
+    // before checking it against the physical limit.
+    const size_t padded =
+        (num_bytes + granularity - 1) / granularity * granularity;
+    if (padded > max_allocation_bytes) {
+      return nullptr;
+    }
+    if (next_address.has_value()) {
+      const uintptr_t address = *next_address;
+      *next_address += padded;
+      *bytes_received = padded;
+      return absl::bit_cast<void*>(address);
+    }
+    if (adjacent_) {
+      return FakeSubAllocator::Alloc(alignment, padded, bytes_received);
+    }
+    *bytes_received = padded;
+    return absl::bit_cast<void*>(kBase - (regions_++ << 30));
+  }
+  bool SupportsCoalescing() const override { return supports_coalescing; }
+  size_t GetAllocationGranularity() const override { return granularity; }
+  void Free(void* ptr, size_t num_bytes) override {
+    frees.emplace_back(ptr, num_bytes);
+  }
+
+  bool supports_coalescing;
+  size_t granularity = 1;
+  size_t max_allocation_bytes = std::numeric_limits<size_t>::max();
+  // Override placement to exercise different region tails and alignments.
+  std::optional<uintptr_t> next_address;
+  std::vector<size_t> requests;
+  std::vector<std::pair<void*, size_t>> frees;
+
+ private:
+  bool adjacent_;
+  uintptr_t regions_ = 0;
+};
+
 // SubAllocator that always fails. A fuse guards against the pre-fix behavior
 // of BFCAllocator::Extend retrying the same size forever: after kFuse calls
 // it reports a test failure and starts succeeding, so a stalled loop
@@ -181,6 +238,8 @@ TEST(BFCAllocatorTest, RejectsAllocationWhoseRoundedSizeWouldOverflow) {
 }
 
 TEST(BFCAllocatorTest, OomLogsAllocationAnnotations) {
+  // Other OOM tests can exhaust AllocateRaw's process-wide diagnostic budget.
+  const int previous_vlog_level = absl::SetVLogLevel("bfc_allocator", 2);
   BFCAllocator::Options opts;
   opts.allow_growth = false;
   opts.allow_retry_on_failure = false;
@@ -210,20 +269,21 @@ TEST(BFCAllocatorTest, OomLogsAllocationAnnotations) {
   EXPECT_EQ(alloc.AllocateRaw(kAlignment, 2048), nullptr);
 
   log.StopCapturingLogs();
+  absl::SetVLogLevel("bfc_allocator", previous_vlog_level);
   alloc.DeallocateRaw(ptr);
 }
 
 // Regression test for the backpedal loop in BFCAllocator::Extend. The loop
-// shrinks the attempt by kBackpedalFactor (0.9) and re-rounds it with
-// RoundedBytes, which rounds up to a multiple of kMinAllocationSize; for
-// attempts below 10 * kMinAllocationSize the rounding undid the shrink, so a
-// persistently failing sub-allocator used to spin the loop forever on the
-// same size. The fuse in the sub-allocator converts such a stall into an
-// immediate test failure instead of a hang.
+// shrinks the attempt by kBackpedalFactor (0.9) and rounds it to the
+// suballocator granularity (at least kMinAllocationSize). An earlier version
+// rounded up with RoundedBytes, so for attempts below 10 * kMinAllocationSize
+// the rounding undid the shrink and a persistently failing sub-allocator spun
+// the loop forever on the same size. The fuse in the sub-allocator converts
+// such a stall into an immediate test failure instead of a hang.
 TEST(BFCAllocatorTest, ExtendTerminatesWhenSubAllocatorAlwaysFails) {
-  // A pool of 8 * kMinAllocationSize (2048 bytes) makes the first backpedal
-  // attempt land in the formerly-stalling zone:
-  // RoundedBytes(0.9 * 2048) == 2048.
+  // A pool of 8 * kMinAllocationSize (2048 bytes) keeps every attempt in the
+  // zone where a 10% shrink is less than one granule, so the loop must force
+  // strict progress on its own.
   constexpr size_t kPool = 8 * BFCAllocator::kMinAllocationSize;
   BFCAllocator::Options opts;
   opts.allow_growth = false;
@@ -1240,6 +1300,560 @@ TEST(BFCAllocatorTest, SpatialUpperAlignmentKeepsSlack) {
   }
 }
 
+TEST(BFCAllocatorDeathTest, RejectsIncompatibleAllocationGranularity) {
+  auto sub = std::make_unique<SpatialGrowthSubAllocator>();
+  sub->granularity = BFCAllocator::kMinAllocationSize + 1;
+  EXPECT_DEATH(BFCAllocator(std::move(sub), 1 << 20, "granularity",
+                            BFCAllocator::Options{}),
+               "allocation_granularity_ % kMinAllocationSize");
+}
+
+TEST(BFCAllocatorDeathTest, SpatialRejectsUnalignedInitialRegion) {
+  constexpr size_t kMiB = 1 << 20;
+  for (const size_t initial_bytes : {kMiB, 3 * kMiB}) {
+    SCOPED_TRACE(initial_bytes);
+    auto sub = std::make_unique<SpatialGrowthSubAllocator>();
+    sub->granularity = 2 * kMiB;
+    auto opts = SharedGpuPoolOptions();
+    opts.allow_growth = true;
+    opts.initial_region_bytes = initial_bytes;
+    EXPECT_DEATH(BFCAllocator(std::move(sub), 8 * kMiB, "initial", opts),
+                 "initial_region_bytes must be a multiple of the granularity");
+  }
+}
+
+TEST(BFCAllocatorTest, SpatialGrowsOnlyAfterUpperGapAllocationFails) {
+  constexpr size_t kMiB = 1 << 20;
+  auto sub = std::make_unique<SpatialGrowthSubAllocator>();
+  auto* source = sub.get();
+  auto opts = SharedGpuPoolOptions();
+  opts.allow_growth = true;
+  opts.initial_region_bytes = 8 * kMiB;
+  BFCAllocator alloc(std::move(sub), 16 * kMiB, "growth", opts);
+
+  void* lower = alloc.AllocateRaw(kAlignment, 2 * kMiB, *kLower);
+  void* upper = alloc.AllocateRaw(kAlignment, 2 * kMiB, *kUpper);
+  ASSERT_NE(lower, nullptr);
+  ASSERT_NE(upper, nullptr);
+  EXPECT_EQ(source->requests, std::vector<size_t>{8 * kMiB});
+
+  // The remaining 4 MiB gap cannot fit this request. Existing allocations stay
+  // put while the usual Extend path obtains a separate upper-owned region.
+  void* extension = alloc.AllocateRaw(kAlignment, 6 * kMiB, *kUpper);
+  ASSERT_NE(extension, nullptr);
+  EXPECT_LT(absl::bit_cast<uintptr_t>(extension), FakeSubAllocator::kBase);
+  EXPECT_EQ(source->requests, (std::vector<size_t>{8 * kMiB, 8 * kMiB}));
+  void* gap = alloc.AllocateRaw(kAlignment, 4 * kMiB, *kLower);
+  ASSERT_NE(gap, nullptr);
+  EXPECT_EQ(absl::bit_cast<uintptr_t>(gap), FakeSubAllocator::kBase + 2 * kMiB);
+  EXPECT_EQ(alloc.AllocateRaw(kAlignment, kMiB, *kUpper), nullptr);
+  EXPECT_EQ(source->requests.size(), 2);
+
+  alloc.DeallocateRaw(extension);
+  EXPECT_EQ(alloc.AllocateRaw(kAlignment, kMiB, *kLower), nullptr);
+  void* reused = alloc.AllocateRaw(kAlignment, 6 * kMiB, *kUpper);
+  EXPECT_EQ(reused, extension);
+  EXPECT_EQ(source->requests.size(), 2);
+  alloc.DeallocateRaw(reused);
+  alloc.DeallocateRaw(gap);
+  alloc.DeallocateRaw(upper);
+  alloc.DeallocateRaw(lower);
+}
+
+TEST(BFCAllocatorTest, SpatialGrowthKeepsLowerCapacityFixed) {
+  constexpr size_t kMiB = 1 << 20;
+  for (bool adjacent : {false, true}) {
+    for (bool coalescing : {false, true}) {
+      for (bool timestamped : {false, true}) {
+        SCOPED_TRACE(adjacent);
+        SCOPED_TRACE(timestamped);
+        SCOPED_TRACE(coalescing);
+        SharedCounter counter;
+        auto sub = std::make_unique<SpatialGrowthSubAllocator>(adjacent);
+        auto* source = sub.get();
+        source->supports_coalescing = coalescing;
+        auto opts = SharedGpuPoolOptions();
+        opts.allow_growth = true;
+        opts.initial_region_bytes = 8 * kMiB;
+        BFCAllocator alloc(std::move(sub), 16 * kMiB, "regions", opts);
+        if (timestamped) {
+          alloc.SetTimingCounter(&counter);
+        }
+        void* initial = alloc.AllocateRaw(kAlignment, 8 * kMiB, *kUpper);
+        void* extension = alloc.AllocateRaw(kAlignment, 4 * kMiB, *kUpper);
+        ASSERT_NE(initial, nullptr);
+        ASSERT_NE(extension, nullptr);
+        alloc.DeallocateRaw(initial);
+        alloc.DeallocateRaw(extension);
+        if (timestamped) {
+          alloc.SetSafeFrontier(counter.next());
+        }
+
+        // Lower allocations remain capped by the initial shared capacity even
+        // when contiguous backing coalesces across its original end.
+        EXPECT_EQ(alloc.AllocateRaw(kAlignment, 12 * kMiB, *kLower), nullptr);
+        void* collective = alloc.AllocateRaw(kAlignment, 8 * kMiB, *kLower);
+        ASSERT_NE(collective, nullptr);
+        EXPECT_EQ(collective, initial);
+        void* ordinary = alloc.AllocateRaw(kAlignment, 4 * kMiB, *kUpper);
+        EXPECT_EQ(ordinary, extension);
+        EXPECT_EQ(source->requests.size(), 2);
+        alloc.DeallocateRaw(ordinary);
+        alloc.DeallocateRaw(collective);
+      }
+    }
+  }
+}
+
+TEST(BFCAllocatorTest, SpatialFirstUpperRequestCanExceedInitialRegion) {
+  constexpr size_t kMiB = 1 << 20;
+  for (bool coalescing : {false, true}) {
+    SCOPED_TRACE(coalescing);
+    auto sub = std::make_unique<SpatialGrowthSubAllocator>();
+    auto* source = sub.get();
+    source->supports_coalescing = coalescing;
+    auto opts = SharedGpuPoolOptions();
+    opts.allow_growth = true;
+    opts.initial_region_bytes = 4 * kMiB;
+    BFCAllocator alloc(std::move(sub), 12 * kMiB, "large_first", opts);
+    void* upper = alloc.AllocateRaw(kAlignment, 8 * kMiB, *kUpper);
+    ASSERT_NE(upper, nullptr);
+    EXPECT_EQ(source->requests, (std::vector<size_t>{4 * kMiB, 8 * kMiB}));
+    void* lower = alloc.AllocateRaw(kAlignment, 4 * kMiB, *kLower);
+    ASSERT_NE(lower, nullptr);
+    EXPECT_EQ(absl::bit_cast<uintptr_t>(lower), FakeSubAllocator::kBase);
+    alloc.DeallocateRaw(lower);
+    alloc.DeallocateRaw(upper);
+  }
+}
+
+TEST(BFCAllocatorTest, SpatialInitialRegionBackpedalsWhenMemoryIsShort) {
+  constexpr size_t kMiB = 1 << 20;
+  auto sub = std::make_unique<SpatialGrowthSubAllocator>();
+  auto* source = sub.get();
+  // Another process holds most of the device: only 4 MiB can be obtained.
+  source->max_allocation_bytes = 4 * kMiB;
+  auto opts = SharedGpuPoolOptions();
+  opts.allow_growth = true;
+  opts.initial_region_bytes = 8 * kMiB;
+  BFCAllocator alloc(std::move(sub), 16 * kMiB, "initial_backpedal", opts);
+  // Like a fixed pool, the initial region shrinks until it fits.
+  void* lower = alloc.AllocateRaw(kAlignment, kMiB, *kLower);
+  ASSERT_NE(lower, nullptr);
+  ASSERT_GT(source->requests.size(), 1);
+  EXPECT_EQ(source->requests.front(), 8 * kMiB);
+  EXPECT_LE(source->requests.back(), 4 * kMiB);
+  const int64_t initial_pool = alloc.GetStats()->pool_bytes.value();
+  EXPECT_EQ(initial_pool, source->requests.back());
+  // The lower end is confined to what was actually obtained, not to the
+  // configured 8 MiB, and never triggers growth.
+  const size_t requests_after_initial = source->requests.size();
+  EXPECT_EQ(alloc.AllocateRaw(kAlignment, 4 * kMiB, *kLower), nullptr);
+  EXPECT_EQ(source->requests.size(), requests_after_initial);
+  // Upper requests still grow once memory becomes available again.
+  source->max_allocation_bytes = 8 * kMiB;
+  void* upper = alloc.AllocateRaw(kAlignment, 6 * kMiB, *kUpper);
+  ASSERT_NE(upper, nullptr);
+  EXPECT_GT(alloc.GetStats()->pool_bytes.value(), initial_pool);
+  alloc.DeallocateRaw(upper);
+  alloc.DeallocateRaw(lower);
+}
+
+TEST(BFCAllocatorTest, SpatialExtensionsUseNormalGrowthSizing) {
+  constexpr size_t kMiB = 1 << 20;
+  auto sub = std::make_unique<SpatialGrowthSubAllocator>();
+  auto* source = sub.get();
+  auto opts = SharedGpuPoolOptions();
+  opts.allow_growth = true;
+  opts.initial_region_bytes = 8 * kMiB;
+  BFCAllocator alloc(std::move(sub), 16 * kMiB, "sizing", opts);
+  void* lower = alloc.AllocateRaw(kAlignment, 8 * kMiB, *kLower);
+  ASSERT_NE(lower, nullptr);
+  void* a = alloc.AllocateRaw(kAlignment, kMiB, *kUpper);
+  void* b = alloc.AllocateRaw(kAlignment, kMiB, *kUpper);
+  void* c = alloc.AllocateRaw(kAlignment, kMiB, *kUpper);
+  ASSERT_NE(a, nullptr);
+  ASSERT_NE(b, nullptr);
+  ASSERT_NE(c, nullptr);
+  EXPECT_EQ(source->requests,
+            (std::vector<size_t>{8 * kMiB, 2 * kMiB, 4 * kMiB}));
+  alloc.DeallocateRaw(a);
+  alloc.DeallocateRaw(b);
+  alloc.DeallocateRaw(c);
+  alloc.DeallocateRaw(lower);
+}
+
+TEST(BFCAllocatorTest,
+     SpatialExtensionBackpedalAndFailurePreserveInitialRegion) {
+  constexpr size_t kMiB = 1 << 20;
+  auto sub = std::make_unique<SpatialGrowthSubAllocator>();
+  auto* source = sub.get();
+  auto opts = SharedGpuPoolOptions();
+  opts.allow_growth = true;
+  opts.initial_region_bytes = 8 * kMiB;
+  BFCAllocator alloc(std::move(sub), 32 * kMiB, "backpedal", opts);
+  void* lower = alloc.AllocateRaw(kAlignment, 8 * kMiB, *kLower);
+  ASSERT_NE(lower, nullptr);
+  source->max_allocation_bytes = 7 * kMiB;
+  void* upper = alloc.AllocateRaw(kAlignment, 6 * kMiB, *kUpper);
+  ASSERT_NE(upper, nullptr);
+  ASSERT_GT(source->requests.size(), 2);
+  EXPECT_EQ(source->requests[1], 8 * kMiB);
+  EXPECT_LE(source->requests.back(), 7 * kMiB);
+  const int64_t pool_bytes = alloc.GetStats()->pool_bytes.value();
+  source->max_allocation_bytes = 0;
+  EXPECT_EQ(alloc.AllocateRaw(kAlignment, 8 * kMiB, *kUpper), nullptr);
+  EXPECT_EQ(alloc.GetStats()->pool_bytes.value(), pool_bytes);
+  alloc.DeallocateRaw(lower);
+  void* reused = alloc.AllocateRaw(kAlignment, 8 * kMiB, *kLower);
+  EXPECT_EQ(reused, lower);
+  alloc.DeallocateRaw(reused);
+  alloc.DeallocateRaw(upper);
+}
+
+TEST(BFCAllocatorTest, SpatialContiguousFirstRequestUsesInitialCapacity) {
+  constexpr size_t kMiB = 1 << 20;
+  auto sub = std::make_unique<SpatialGrowthSubAllocator>(true);
+  auto* source = sub.get();
+  auto opts = SharedGpuPoolOptions();
+  opts.allow_growth = true;
+  opts.initial_region_bytes = 4 * kMiB;
+  BFCAllocator alloc(std::move(sub), 8 * kMiB, "contiguous_first", opts);
+  // Neither the initial region nor the remaining cap fits alone.
+  void* upper = alloc.AllocateRaw(kAlignment, 8 * kMiB, *kUpper);
+  ASSERT_NE(upper, nullptr);
+  EXPECT_EQ(absl::bit_cast<uintptr_t>(upper), FakeSubAllocator::kBase);
+  EXPECT_EQ(source->requests, (std::vector<size_t>{4 * kMiB, 4 * kMiB}));
+  alloc.DeallocateRaw(upper);
+  EXPECT_EQ(alloc.AllocateRaw(kAlignment, 5 * kMiB, *kLower), nullptr);
+  void* lower = alloc.AllocateRaw(kAlignment, 4 * kMiB, *kLower);
+  ASSERT_EQ(lower, upper);
+  void* remaining = alloc.AllocateRaw(kAlignment, 4 * kMiB, *kUpper);
+  ASSERT_NE(remaining, nullptr);
+  EXPECT_EQ(source->requests.size(), 2);
+  alloc.DeallocateRaw(lower);
+  alloc.DeallocateRaw(remaining);
+}
+
+TEST(BFCAllocatorTest,
+     SpatialContiguousCoalescingKeepsLowerPaddingWithinLimit) {
+  constexpr size_t kMiB = 1 << 20;
+  for (bool timestamped : {false, true}) {
+    for (bool exact : {false, true}) {
+      SCOPED_TRACE(timestamped);
+      SCOPED_TRACE(exact);
+      SharedCounter counter;
+      auto sub = std::make_unique<SpatialGrowthSubAllocator>(true);
+      auto* source = sub.get();
+      auto opts = SharedGpuPoolOptions();
+      opts.allow_growth = true;
+      opts.initial_region_bytes = 8 * kMiB;
+      opts.lower_end_policy.gap_split_policy =
+          exact ? BFCAllocator::SplitPolicy::kExact
+                : BFCAllocator::SplitPolicy::kRetainPadding;
+      BFCAllocator alloc(std::move(sub), 12 * kMiB, "coalesced", opts);
+      if (timestamped) {
+        alloc.SetTimingCounter(&counter);
+      }
+      void* high = alloc.AllocateRaw(kAlignment, 4 * kMiB, *kUpper);
+      void* low = alloc.AllocateRaw(kAlignment, 4 * kMiB, *kUpper);
+      void* extension = alloc.AllocateRaw(kAlignment, 4 * kMiB, *kUpper);
+      ASSERT_NE(high, nullptr);
+      ASSERT_NE(low, nullptr);
+      ASSERT_NE(extension, nullptr);
+      alloc.DeallocateRaw(high);
+      alloc.DeallocateRaw(extension);
+      if (timestamped) {
+        alloc.SetSafeFrontier(counter.next());
+      }
+      // Reuse a single chunk spanning the old end, without another allocation.
+      void* spanning = alloc.AllocateRaw(kAlignment, 8 * kMiB, *kUpper);
+      ASSERT_EQ(spanning, high);
+      EXPECT_EQ(source->requests.size(), 2);
+      alloc.DeallocateRaw(spanning);
+      alloc.DeallocateRaw(low);
+      if (timestamped) {
+        alloc.SetSafeFrontier(counter.next());
+      }
+      EXPECT_EQ(alloc.AllocateRaw(kAlignment, 9 * kMiB, *kLower), nullptr);
+      void* collective = alloc.AllocateRaw(kAlignment, 6 * kMiB, *kLower);
+      ASSERT_EQ(collective, low);
+      EXPECT_EQ(alloc.AllocatedSize(collective), (exact ? 6 : 8) * kMiB);
+      if (exact) {
+        // Four-MiB alignment would place this at the fixed limit, despite
+        // there being mapped free memory above it.
+        EXPECT_EQ(alloc.AllocateRaw(4 * kMiB, 2 * kMiB, *kLower), nullptr);
+      }
+      void* ordinary = alloc.AllocateRaw(kAlignment, 4 * kMiB, *kUpper);
+      ASSERT_NE(ordinary, nullptr);
+      EXPECT_EQ(source->requests.size(), 2);
+      alloc.DeallocateRaw(ordinary);
+      alloc.DeallocateRaw(collective);
+    }
+  }
+}
+
+TEST(BFCAllocatorTest, SpatialContiguousFailedGrowthPreservesFreeTail) {
+  constexpr size_t kMiB = 1 << 20;
+  auto sub = std::make_unique<SpatialGrowthSubAllocator>(true);
+  auto* source = sub.get();
+  auto opts = SharedGpuPoolOptions();
+  opts.allow_growth = true;
+  opts.initial_region_bytes = 8 * kMiB;
+  BFCAllocator alloc(std::move(sub), 12 * kMiB, "failed_contiguous", opts);
+  void* collective = alloc.AllocateRaw(kAlignment, 4 * kMiB, *kLower);
+  ASSERT_NE(collective, nullptr);
+  source->max_allocation_bytes = 0;
+  EXPECT_EQ(alloc.AllocateRaw(kAlignment, 8 * kMiB, *kUpper), nullptr);
+  EXPECT_EQ(alloc.GetStats()->pool_bytes.value(), 8 * kMiB);
+  source->max_allocation_bytes = 4 * kMiB;
+  void* spanning = alloc.AllocateRaw(kAlignment, 8 * kMiB, *kUpper);
+  ASSERT_NE(spanning, nullptr);
+  EXPECT_EQ(absl::bit_cast<uintptr_t>(spanning),
+            FakeSubAllocator::kBase + 4 * kMiB);
+  EXPECT_EQ(alloc.GetStats()->pool_bytes.value(), 12 * kMiB);
+  EXPECT_EQ(alloc.AllocateRaw(kAlignment, kMiB, *kLower), nullptr);
+  alloc.DeallocateRaw(spanning);
+  alloc.DeallocateRaw(collective);
+}
+
+TEST(BFCAllocatorTest, BackpedalHonoursSubAllocatorGranularity) {
+  constexpr size_t kMiB = 1 << 20;
+  auto sub = std::make_unique<SpatialGrowthSubAllocator>();
+  auto* source = sub.get();
+  source->granularity = 2 * kMiB;
+  BFCAllocator::Options opts;
+  opts.allow_retry_on_failure = false;
+  BFCAllocator alloc(std::move(sub), 32 * kMiB, "granularity", opts);
+  void* first = alloc.AllocateRaw(kAlignment, 6 * kMiB, *kLower);
+  ASSERT_NE(first, nullptr);
+  EXPECT_EQ(source->requests, std::vector<size_t>{8 * kMiB});
+
+  // Only two granules of physical memory remain. Backpedaling must step in
+  // whole granules; a 256-byte step would be padded back up to 6 MiB and fail
+  // repeatedly, then drop below the request without ever trying 4 MiB.
+  source->max_allocation_bytes = 4 * kMiB;
+  void* second = alloc.AllocateRaw(kAlignment, 4 * kMiB, *kLower);
+  ASSERT_NE(second, nullptr);
+  EXPECT_EQ(source->requests,
+            (std::vector<size_t>{8 * kMiB, 8 * kMiB, 6 * kMiB, 4 * kMiB}));
+  alloc.DeallocateRaw(second);
+  alloc.DeallocateRaw(first);
+}
+
+TEST(BFCAllocatorTest,
+     SpatialContiguousGrowthCompletesFreeTailAtPhysicalLimit) {
+  constexpr size_t kMiB = 1 << 20;
+  auto sub = std::make_unique<SpatialGrowthSubAllocator>(true);
+  auto* source = sub.get();
+  source->granularity = 2 * kMiB;
+  auto opts = SharedGpuPoolOptions();
+  opts.allow_growth = true;
+  opts.initial_region_bytes = 8 * kMiB;
+  BFCAllocator alloc(std::move(sub), 32 * kMiB, "physical_limit", opts);
+  void* lower = alloc.AllocateRaw(kAlignment, 2 * kMiB, *kLower);
+  ASSERT_NE(lower, nullptr);
+  // 24 MiB of reserved capacity remain but only one granule is physically
+  // available. The 6 MiB free tail needs just 2 MiB more to fit the request,
+  // so backpedaling must continue down to that floor instead of giving up at
+  // the full request size.
+  source->max_allocation_bytes = 2 * kMiB;
+  void* upper = alloc.AllocateRaw(kAlignment, 8 * kMiB, *kUpper);
+  ASSERT_NE(upper, nullptr);
+  EXPECT_EQ(absl::bit_cast<uintptr_t>(upper),
+            FakeSubAllocator::kBase + 2 * kMiB);
+  EXPECT_EQ(source->requests, (std::vector<size_t>{8 * kMiB, 8 * kMiB, 6 * kMiB,
+                                                   4 * kMiB, 2 * kMiB}));
+  EXPECT_EQ(alloc.GetStats()->pool_bytes.value(), 10 * kMiB);
+  EXPECT_EQ(alloc.AllocateRaw(kAlignment, kMiB, *kLower), nullptr);
+  alloc.DeallocateRaw(upper);
+  alloc.DeallocateRaw(lower);
+}
+
+TEST(BFCAllocatorTest, SpatialNonAdjacentShortExtensionIsReturned) {
+  constexpr size_t kMiB = 1 << 20;
+  auto sub = std::make_unique<SpatialGrowthSubAllocator>(false);
+  auto* source = sub.get();
+  source->supports_coalescing = true;
+  source->granularity = 2 * kMiB;
+  auto opts = SharedGpuPoolOptions();
+  opts.allow_growth = true;
+  opts.initial_region_bytes = 8 * kMiB;
+  BFCAllocator alloc(std::move(sub), 32 * kMiB, "non_adjacent", opts);
+  void* lower = alloc.AllocateRaw(kAlignment, 2 * kMiB, *kLower);
+  ASSERT_NE(lower, nullptr);
+  // The suballocator claims coalescing support but returns a separate region.
+  // An extension sized only to complete the free tail cannot serve the
+  // request on its own, so it must be handed back rather than kept.
+  source->max_allocation_bytes = 2 * kMiB;
+  EXPECT_EQ(alloc.AllocateRaw(kAlignment, 8 * kMiB, *kUpper), nullptr);
+  EXPECT_EQ(source->requests, (std::vector<size_t>{8 * kMiB, 8 * kMiB, 6 * kMiB,
+                                                   4 * kMiB, 2 * kMiB}));
+  EXPECT_EQ(alloc.GetStats()->pool_bytes.value(), 8 * kMiB);
+  // With enough physical memory the request is served from a separate region.
+  source->max_allocation_bytes = 8 * kMiB;
+  void* upper = alloc.AllocateRaw(kAlignment, 8 * kMiB, *kUpper);
+  ASSERT_NE(upper, nullptr);
+  EXPECT_LT(absl::bit_cast<uintptr_t>(upper), FakeSubAllocator::kBase);
+  EXPECT_EQ(source->requests.size(), 6);
+  EXPECT_EQ(source->requests.back(), 8 * kMiB);
+  EXPECT_EQ(alloc.GetStats()->pool_bytes.value(), 16 * kMiB);
+  alloc.DeallocateRaw(upper);
+  alloc.DeallocateRaw(lower);
+}
+
+TEST(BFCAllocatorTest, SpatialExtensionChecksActualRegionTail) {
+  constexpr size_t kMiB = 1 << 20;
+  constexpr uintptr_t kSecondRegion = FakeSubAllocator::kBase - 64 * kMiB;
+  auto sub = std::make_unique<SpatialGrowthSubAllocator>();
+  auto* source = sub.get();
+  source->supports_coalescing = true;
+  source->granularity = 2 * kMiB;
+  auto opts = SharedGpuPoolOptions();
+  opts.allow_growth = true;
+  opts.initial_region_bytes = 8 * kMiB;
+  BFCAllocator alloc(std::move(sub), 32 * kMiB, "actual_tail", opts);
+  void* lower = alloc.AllocateRaw(kAlignment, 6 * kMiB, *kLower);
+  ASSERT_NE(lower, nullptr);
+
+  // The first region has a 2 MiB free tail. Add a separate 6 MiB region
+  // whose occupied tail ends 2 MiB short of a 4 MiB alignment boundary.
+  source->next_address = kSecondRegion;
+  source->max_allocation_bytes = 6 * kMiB;
+  void* upper = alloc.AllocateRaw(kAlignment, 6 * kMiB, *kUpper);
+  ASSERT_EQ(absl::bit_cast<uintptr_t>(upper), kSecondRegion);
+  ASSERT_EQ(alloc.GetStats()->pool_bytes.value(), 14 * kMiB);
+
+  // A 4 MiB extension would suffice beside the first region, but the actual
+  // second region needs 6 MiB including alignment padding. Even though the
+  // returned extension equals the request size, it must be handed back.
+  source->max_allocation_bytes = 4 * kMiB;
+  EXPECT_EQ(alloc.AllocateRaw(4 * kMiB, 4 * kMiB, *kUpper), nullptr);
+  EXPECT_EQ(source->requests,
+            (std::vector<size_t>{8 * kMiB, 8 * kMiB, 6 * kMiB, 8 * kMiB,
+                                 6 * kMiB, 4 * kMiB}));
+  ASSERT_EQ(source->frees.size(), 1);
+  EXPECT_EQ(absl::bit_cast<uintptr_t>(source->frees[0].first),
+            kSecondRegion + 6 * kMiB);
+  EXPECT_EQ(source->frees[0].second, 4 * kMiB);
+  EXPECT_EQ(alloc.GetStats()->pool_bytes.value(), 14 * kMiB);
+
+  // Reuse the rejected address with enough backing for the aligned request.
+  source->next_address = kSecondRegion + 6 * kMiB;
+  source->max_allocation_bytes = 6 * kMiB;
+  void* aligned = alloc.AllocateRaw(4 * kMiB, 4 * kMiB, *kUpper);
+  ASSERT_NE(aligned, nullptr);
+  EXPECT_EQ(absl::bit_cast<uintptr_t>(aligned), kSecondRegion + 8 * kMiB);
+  EXPECT_EQ(alloc.GetStats()->pool_bytes.value(), 20 * kMiB);
+  alloc.DeallocateRaw(aligned);
+  alloc.DeallocateRaw(upper);
+  alloc.DeallocateRaw(lower);
+}
+
+TEST(BFCAllocatorTest, SpatialSeparateExtensionChecksAlignmentPadding) {
+  constexpr size_t kMiB = 1 << 20;
+  constexpr uintptr_t kSeparateRegion = FakeSubAllocator::kBase - 62 * kMiB;
+  auto sub = std::make_unique<SpatialGrowthSubAllocator>();
+  auto* source = sub.get();
+  source->supports_coalescing = true;
+  source->granularity = 2 * kMiB;
+  auto opts = SharedGpuPoolOptions();
+  opts.allow_growth = true;
+  opts.initial_region_bytes = 8 * kMiB;
+  BFCAllocator alloc(std::move(sub), 32 * kMiB, "separate_padding", opts);
+  void* lower = alloc.AllocateRaw(kAlignment, 6 * kMiB, *kLower);
+  ASSERT_NE(lower, nullptr);
+
+  // The new region is large enough for the request, but not its 2 MiB
+  // alignment prefix. It cannot use the initial region's free tail.
+  source->next_address = kSeparateRegion;
+  source->max_allocation_bytes = 4 * kMiB;
+  EXPECT_EQ(alloc.AllocateRaw(4 * kMiB, 4 * kMiB, *kUpper), nullptr);
+  ASSERT_EQ(source->frees.size(), 1);
+  EXPECT_EQ(absl::bit_cast<uintptr_t>(source->frees[0].first), kSeparateRegion);
+  EXPECT_EQ(source->frees[0].second, 4 * kMiB);
+  EXPECT_EQ(alloc.GetStats()->pool_bytes.value(), 8 * kMiB);
+
+  // The same size is sufficient when the separate region is aligned.
+  source->next_address = kSeparateRegion + 2 * kMiB;
+  void* upper = alloc.AllocateRaw(4 * kMiB, 4 * kMiB, *kUpper);
+  ASSERT_NE(upper, nullptr);
+  EXPECT_TRUE(IsAligned(upper, 4 * kMiB));
+  EXPECT_EQ(alloc.GetStats()->pool_bytes.value(), 12 * kMiB);
+  alloc.DeallocateRaw(upper);
+  alloc.DeallocateRaw(lower);
+}
+
+TEST(BFCAllocatorTest, SpatialGrowthOomLogsLowerUsableGap) {
+  constexpr size_t kMiB = 1 << 20;
+  for (const size_t lower_bytes : {4 * kMiB, 8 * kMiB}) {
+    SCOPED_TRACE(lower_bytes);
+    auto opts = SharedGpuPoolOptions();
+    opts.allow_growth = true;
+    opts.initial_region_bytes = 8 * kMiB;
+    BFCAllocator alloc(std::make_unique<SpatialGrowthSubAllocator>(true),
+                       64 * kMiB, "lower_gap", opts);
+    void* lower = alloc.AllocateRaw(kAlignment, lower_bytes, *kLower);
+    void* upper = alloc.AllocateRaw(kAlignment, 56 * kMiB, *kUpper);
+    ASSERT_NE(lower, nullptr);
+    ASSERT_NE(upper, nullptr);
+    alloc.DeallocateRaw(upper);
+
+    // The coalesced gap is 60 or 56 MiB, but the lower end can only use
+    // its initial-region prefix: 4 MiB or zero, respectively.
+    const int previous_vlog_level = absl::SetVLogLevel("bfc_allocator", 2);
+    absl::ScopedMockLog log(absl::MockLogDefault::kIgnoreUnexpected);
+    EXPECT_CALL(
+        log,
+        Log(absl::LogSeverity::kInfo, _,
+            AllOf(
+                HasSubstr("Allocator memory summary:"),
+                HasSubstr(lower_bytes == 4 * kMiB
+                              ? "largest free chunk usable by lower end 4.00MiB"
+                              : "largest free chunk usable by lower end 0B"))))
+        .Times(1);
+    log.StartCapturingLogs();
+    EXPECT_EQ(alloc.AllocateRaw(kAlignment, 8 * kMiB, *kLower), nullptr);
+    log.StopCapturingLogs();
+    absl::SetVLogLevel("bfc_allocator", previous_vlog_level);
+    alloc.DeallocateRaw(lower);
+  }
+}
+
+TEST(BFCAllocatorTest, SpatialContiguousGrowthAccountsForTailAlignmentPadding) {
+  constexpr size_t kMiB = 1 << 20;
+  for (const size_t total : {12 * kMiB, 11 * kMiB}) {
+    SCOPED_TRACE(total);
+    auto sub = std::make_unique<SpatialGrowthSubAllocator>(true);
+    auto* source = sub.get();
+    auto opts = SharedGpuPoolOptions();
+    opts.allow_growth = true;
+    opts.initial_region_bytes = 6 * kMiB;
+    BFCAllocator alloc(std::move(sub), total, "tail_padding", opts);
+    void* lower = alloc.AllocateRaw(kAlignment, 5 * kMiB, *kLower);
+    ASSERT_NE(lower, nullptr);
+    // The 1 MiB free tail starts 5 MiB into the region. A 4 MiB-aligned
+    // 4 MiB request can only start at kBase + 8 MiB, so the merged chunk must
+    // reach kBase + 12 MiB: a 6 MiB extension, not 4 MiB.
+    void* upper = alloc.AllocateRaw(4 * kMiB, 4 * kMiB, *kUpper);
+    if (total == 12 * kMiB) {
+      ASSERT_NE(upper, nullptr);
+      EXPECT_EQ(absl::bit_cast<uintptr_t>(upper),
+                FakeSubAllocator::kBase + 8 * kMiB);
+      EXPECT_EQ(source->requests, (std::vector<size_t>{6 * kMiB, 6 * kMiB}));
+      EXPECT_EQ(alloc.GetStats()->pool_bytes.value(), 12 * kMiB);
+      alloc.DeallocateRaw(upper);
+    } else {
+      // Capacity is one MiB short of what the aligned request needs. No
+      // extension that cannot serve the request should be mapped.
+      EXPECT_EQ(upper, nullptr);
+      EXPECT_EQ(source->requests, std::vector<size_t>{6 * kMiB});
+      EXPECT_EQ(alloc.GetStats()->pool_bytes.value(), 6 * kMiB);
+    }
+    alloc.DeallocateRaw(lower);
+  }
+}
+
 TEST(BFCAllocatorTest, SpatialUnderContention) {
   BFCAllocator::Options opts;
   opts.allow_growth = false;

From d9ebfa1a03527b4c938015728e15bb9bf84006ff Mon Sep 17 00:00:00 2001
From: Noxtimo <139155109+Noxtimo@users.noreply.github.com>
Date: Fri, 9 Oct 2026 10:39:50 +0100
Subject: [PATCH 0343/1000] Fix negative zero handling in tnp.angle and enhance
 cbrt assertions

---
 tensorflow/python/ops/numpy_ops/np_math_ops.py      | 2 +-
 tensorflow/python/ops/numpy_ops/np_math_ops_test.py | 8 +++++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py
index 484159a4fd8..5b4d0a5e0a2 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@@ -976,7 +976,7 @@ def angle(z, deg=False):  # pylint: disable=missing-function-docstring
       # `np.pi` and `0` are Python scalars, which would make the result
       # float32 whatever `x` is, so build them in `x`'s dtype instead.
       return array_ops.where_v2(
-          x < 0,
+          _signbit(x),
           constant_op.constant(np.pi, dtype=x.dtype),
           constant_op.constant(0, dtype=x.dtype),
       )
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 611057037dd..c78673b7d76 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
@@ -38,6 +38,8 @@ class MathTest(test.TestCase, parameterized.TestCase):
 
   def setUp(self):
     super(MathTest, self).setUp()
+    ops.set_dtype_conversion_mode('legacy')
+    np_math_ops.enable_numpy_methods_on_tensor()
     self.array_transforms = [
         lambda x: x,  # Identity,
         ops.convert_to_tensor,
@@ -174,7 +176,7 @@ class MathTest(test.TestCase, parameterized.TestCase):
     # branch out of Python scalars made it float32 for every real dtype, which
     # for float64 also cost precision.
     for dtype in [np.float16, np.float32, np.float64]:
-      arg = np.array([-2.0, -0.5, 0.0, 0.5, 2.0], dtype=dtype)
+      arg = np.array([-2.0, -0.5, -0.0, 0.0, 0.5, 2.0], dtype=dtype)
       self.match(
           np_math_ops.angle(arg), np.angle(arg), msg='angle({})'.format(arg)
       )
@@ -876,6 +878,7 @@ class MathTest(test.TestCase, parameterized.TestCase):
 
     x_bf16 = constant_op.constant([-0.0, 0.0, -8.0, 8.0], dtype=dtypes.bfloat16)
     out_bf16 = np_math_ops.cbrt(x_bf16)
+    self.assertAllClose(out_bf16, [-0.0, 0.0, -2.0, 2.0])
     self.assertAllEqual(
         np_math_ops.signbit(out_bf16), [True, False, True, False]
     )
@@ -898,6 +901,9 @@ class MathTest(test.TestCase, parameterized.TestCase):
       out_edge = np_math_ops.cbrt(x_edge)
       # inf ** (1/3) == inf; NaN ** (1/3) == NaN
       # Negative sign bit for -inf is preserved.
+      self.assertAllEqual(
+          out_edge[0:2], [-float('inf'), float('inf')]
+      )
       self.assertAllEqual(np_math_ops.signbit(out_edge[0:2]), [True, False])
       self.assertAllEqual(math_ops.is_nan(out_edge[2]), True)
 

From 39b8f9e305f614b3ab7a2fc896e2e24fe77dd139 Mon Sep 17 00:00:00 2001
From: vfdev <vfdev.5@gmail.com>
Date: Fri, 9 Oct 2026 02:59:33 -0700
Subject: [PATCH 0344/1000] PR #50158: Updated GpuModel mapping and _gpu_spec

Imported from GitHub PR https://github.com/openxla/xla/pull/50158

Follow-up PR to https://github.com/openxla/xla/pull/46268
Copybara import of the project:

--
82d0e5aead9a0a6373a754584a84b5908a6468b9 by vfdev-5 <vfdev.5@gmail.com>:

Updated GpuModel mapping and _gpu_spec

Merging this change closes #50158

PiperOrigin-RevId: 996411956
---
 .../gpu/target_config/target_config.cc        |  9 ++++---
 .../gpu/target_config/target_config.h         |  3 ++-
 third_party/xla/xla/python/_gpu_spec.pyi      | 24 ++++++++++++-------
 third_party/xla/xla/python/gpu_spec.cc        | 10 ++++----
 4 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/target_config/target_config.cc b/third_party/xla/xla/backends/gpu/target_config/target_config.cc
index 58881673a0e..9241e6bae0b 100644
--- a/third_party/xla/xla/backends/gpu/target_config/target_config.cc
+++ b/third_party/xla/xla/backends/gpu/target_config/target_config.cc
@@ -89,9 +89,15 @@ absl::StatusOr<absl::string_view> GetEmbeddedGpuTargetConfigData(
     case GpuModel::MI200:
       filename = "mi200.txtpb";
       break;
+    case GpuModel::MI300:
+      filename = "mi300.txtpb";
+      break;
     case GpuModel::MI350:
       filename = "mi350.txtpb";
       break;
+    case GpuModel::MI450:
+      filename = "mi450.txtpb";
+      break;
     case GpuModel::P100:
       filename = "p100.txtpb";
       break;
@@ -113,9 +119,6 @@ absl::StatusOr<absl::string_view> GetEmbeddedGpuTargetConfigData(
     case GpuModel::RTX6000PRO:
       filename = "rtx6000pro.txtpb";
       break;
-    case GpuModel::GFX1250:
-      filename = "gfx1250.txtpb";
-      break;
     default:
       return absl::NotFoundError(
           absl::StrCat("Embedded file not found: ", gpu_model, ".txtpb"));
diff --git a/third_party/xla/xla/backends/gpu/target_config/target_config.h b/third_party/xla/xla/backends/gpu/target_config/target_config.h
index 54fc6cf336f..1207438a573 100644
--- a/third_party/xla/xla/backends/gpu/target_config/target_config.h
+++ b/third_party/xla/xla/backends/gpu/target_config/target_config.h
@@ -42,7 +42,9 @@ enum class GpuModel {
   H100_SXM_MIG,
   H200,
   MI200,
+  MI300,
   MI350,
+  MI450,
   P100,
   PVC,
   VR_NVL72,
@@ -50,7 +52,6 @@ enum class GpuModel {
   GB200,
   GB300,
   RTX6000PRO,
-  GFX1250,
 };
 
 // Description of a target device for compilation.
diff --git a/third_party/xla/xla/python/_gpu_spec.pyi b/third_party/xla/xla/python/_gpu_spec.pyi
index 16fbd1e6162..0230d2cec3b 100644
--- a/third_party/xla/xla/python/_gpu_spec.pyi
+++ b/third_party/xla/xla/python/_gpu_spec.pyi
@@ -39,6 +39,8 @@ class GpuModel(enum.Enum):
 
   B200 = 4
 
+  B200_MIG = 5
+
   B300 = 6
 
   BMG_G21 = 7
@@ -47,26 +49,30 @@ class GpuModel(enum.Enum):
 
   H100_SXM = 9
 
+  H100_SXM_MIG = 10
+
   H200 = 11
 
   MI200 = 12
 
-  P100 = 14
+  MI300 = 13
 
-  PVC = 15
+  MI350 = 14
 
-  VR_NVL72 = 16
+  MI450 = 15
 
-  V100 = 17
+  P100 = 16
 
-  GB200 = 18
+  PVC = 17
 
-  GB300 = 19
+  VR_NVL72 = 18
 
-  RTX6000PRO = 20
+  V100 = 19
 
-  GFX1250 = 21
+  GB200 = 20
 
-  MI350 = 13
+  GB300 = 21
+
+  RTX6000PRO = 22
 
 def get_gpu_spec(gpu_model: GpuModel) -> GpuTargetConfig: ...
diff --git a/third_party/xla/xla/python/gpu_spec.cc b/third_party/xla/xla/python/gpu_spec.cc
index f8ca88ff23e..66c293031f6 100644
--- a/third_party/xla/xla/python/gpu_spec.cc
+++ b/third_party/xla/xla/python/gpu_spec.cc
@@ -17,7 +17,6 @@ limitations under the License.
 #include <string>
 #include <utility>
 
-#include "absl/base/thread_annotations.h"
 #include "nanobind/nanobind.h"
 #include "nanobind/stl/string.h"  // IWYU pragma: keep
 #include "xla/backends/gpu/target_config/target_config.h"
@@ -92,21 +91,24 @@ NB_MODULE(_gpu_spec, m) {
       .value("A100_SXM_80", gpu::GpuModel::A100_SXM_80)
       .value("A6000", gpu::GpuModel::A6000)
       .value("B200", gpu::GpuModel::B200)
+      .value("B200_MIG", gpu::GpuModel::B200_MIG)
       .value("B300", gpu::GpuModel::B300)
       .value("BMG_G21", gpu::GpuModel::BMG_G21)
       .value("H100_PCIE", gpu::GpuModel::H100_PCIE)
       .value("H100_SXM", gpu::GpuModel::H100_SXM)
+      .value("H100_SXM_MIG", gpu::GpuModel::H100_SXM_MIG)
       .value("H200", gpu::GpuModel::H200)
       .value("MI200", gpu::GpuModel::MI200)
+      .value("MI300", gpu::GpuModel::MI300)
+      .value("MI350", gpu::GpuModel::MI350)
+      .value("MI450", gpu::GpuModel::MI450)
       .value("P100", gpu::GpuModel::P100)
       .value("PVC", gpu::GpuModel::PVC)
       .value("VR_NVL72", gpu::GpuModel::VR_NVL72)
       .value("V100", gpu::GpuModel::V100)
       .value("GB200", gpu::GpuModel::GB200)
       .value("GB300", gpu::GpuModel::GB300)
-      .value("RTX6000PRO", gpu::GpuModel::RTX6000PRO)
-      .value("GFX1250", gpu::GpuModel::GFX1250)
-      .value("MI350", gpu::GpuModel::MI350);
+      .value("RTX6000PRO", gpu::GpuModel::RTX6000PRO);
 
   m.def(
       "get_gpu_spec",

From 2a90056f23d2ed71a325932e2b2e9578073f65a6 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Fri, 9 Oct 2026 03:21:54 -0700
Subject: [PATCH 0345/1000] Implement WEAKORDER comparison type in StableHLO
 and XLA

Add `WEAKORDER` to `ComparisonType` across StableHLO, CHLO, VHLO, and MHLO as
specified in StableHLO RFC #3003, and connect it to `Comparison::Order::kWeak`
in XLA's HLO translators.

`WEAKORDER` implements weak-ordering semantics for floating-point values where
`-0.0` and `+0.0` compare equal while all NaN values are ordered consistently
above `+Inf`. This matches JAX sorting semantics without paying the overhead of
a `TOTALORDER` with normalized NaNs and zeros.

- Bump VHLO version to `1.22.0` and guard `WEAKORDER` on `CompareOpV1`.
- Enforce floating-point element types for `WEAKORDER` in `compare`.
- Implement `WEAKORDER` in the StableHLO reference interpreter, aggressive
  folder, CHLO legalization, and Linalg scalar lowering.
- Add `WEAKORDER` to MHLO and map it to `ComparisonOrder::kWeak` in XLA HLO
  import and export.

PiperOrigin-RevId: 996422412
---
 .../xla/third_party/stablehlo/temporary.patch | 4142 ++++++++++++++++-
 .../hlo_to_mhlo/hlo_function_importer.cc      |    8 +-
 .../translate/hlo_to_mhlo/tests/import.hlo    |    3 +
 .../translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc  |    3 +
 .../xla/hlo/translate/tests/stablehlo.mlir    |   27 +
 .../xla/xla/mlir_hlo/mhlo/IR/hlo_ops.cc       |   11 +-
 .../xla/xla/mlir_hlo/mhlo/IR/hlo_ops_enums.td |    4 +-
 .../mhlo/transforms/map_chlo_to_hlo_op.h      |    2 +
 .../mhlo/transforms/map_mhlo_to_scalar_op.h   |   54 +
 .../mhlo/hlo-legalize-to-stablehlo.mlir       |   10 +
 .../mhlo/stablehlo-legalize-to-hlo.mlir       |   12 +
 11 files changed, 4247 insertions(+), 29 deletions(-)

diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch
index 08f59aacc6b..680520db88d 100644
--- a/third_party/xla/third_party/stablehlo/temporary.patch
+++ b/third_party/xla/third_party/stablehlo/temporary.patch
@@ -848,6 +848,26 @@ diff --ruN a/stablehlo/build_tools/lit_test_suite.bzl b/stablehlo/build_tools/li
              tags = tags,
              size = size_per_test,
          )
+diff --ruN a/stablehlo/docs/generated/chlo.md b/stablehlo/docs/generated/chlo.md
+--- stablehlo/docs/generated/chlo.md
++++ stablehlo/docs/generated/chlo.md
+@@ -2068,7 +2068,7 @@
+ 
+ ```
+ #chlo.comparison_type<
+-  `NOTYPE` | `FLOAT` | `TOTALORDER` | `SIGNED` | `UNSIGNED`   # value
++  `NOTYPE` | `FLOAT` | `TOTALORDER` | `SIGNED` | `UNSIGNED` | `WEAKORDER`   # value
+ >
+ ```
+ 
+@@ -2141,6 +2141,7 @@
+ | TOTALORDER | `2` | TOTALORDER |
+ | SIGNED | `3` | SIGNED |
+ | UNSIGNED | `4` | UNSIGNED |
++| WEAKORDER | `5` | WEAKORDER |
+ 
+ ### Precision
+ 
 diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
 --- stablehlo/docs/spec.md
 +++ stablehlo/docs/spec.md
@@ -1707,6 +1727,45 @@ diff --ruN a/stablehlo/rfcs/20260907-exp2-log2.md b/stablehlo/rfcs/20260907-exp2
 +[jax2]: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.log2.html
 +[torch1]: https://pytorch.org/docs/stable/generated/torch.exp2.html
 +[torch2]: https://pytorch.org/docs/stable/generated/torch.log2.html
+diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/pointwise.mlir b/stablehlo/stablehlo/conversions/linalg/tests/pointwise.mlir
+--- stablehlo/stablehlo/conversions/linalg/tests/pointwise.mlir
++++ stablehlo/stablehlo/conversions/linalg/tests/pointwise.mlir
+@@ -728,6 +728,35 @@
+ 
+ // -----
+ 
++// CHECK-LABEL: func @float_cmp_weakorder
++// CHECK-PRIMITIVE-LABEL: func @float_cmp_weakorder
++func.func @float_cmp_weakorder(%lhs: tensor<2x2xbf16>,
++                %rhs: tensor<2x2xbf16>) -> (tensor<2x2xi1>) {
++  %0 = "stablehlo.compare"(%lhs, %rhs) {
++    comparison_direction = #stablehlo<comparison_direction LT>,
++    compare_type = #stablehlo<comparison_type WEAKORDER>
++  } : (tensor<2x2xbf16>, tensor<2x2xbf16>) -> tensor<2x2xi1>
++  func.return %0 : tensor<2x2xi1>
++}
++// CHECK: tensor.empty() : tensor<2x2xi1>
++// CHECK: linalg.generic
++// CHECK-NEXT: ^bb0(%[[LHS_IN:.*]]: bf16, %[[RHS_IN:.*]]: bf16, %{{.*}}: i1):
++// CHECK-NEXT:   %[[CMP_ULT:.*]] = arith.cmpf ult, %[[LHS_IN]], %[[RHS_IN]] : bf16
++// CHECK-NEXT:   %[[CMP_ORD:.*]] = arith.cmpf ord, %[[LHS_IN]], %[[LHS_IN]] : bf16
++// CHECK-NEXT:   %[[RESULT:.*]] = arith.andi %[[CMP_ULT]], %[[CMP_ORD]] : i1
++// CHECK-NEXT:   linalg.yield %[[RESULT]] : i1
++
++// CHECK-PRIMITIVE: linalg.map
++// CHECK-PRIMITIVE-SAME: ins(
++// CHECK-PRIMITIVE-SAME: outs(
++// CHECK-PRIMITIVE-NEXT: (%[[LHS_IN:[a-zA-Z0-9]*]]: bf16, %[[RHS_IN:.*]]: bf16, %[[RESULT_OUT:.*]]: i1) {
++// CHECK-PRIMITIVE-NEXT:   %[[CMP_ULT:.*]] = arith.cmpf ult, %[[LHS_IN]], %[[RHS_IN]] : bf16
++// CHECK-PRIMITIVE-NEXT:   %[[CMP_ORD:.*]] = arith.cmpf ord, %[[LHS_IN]], %[[LHS_IN]] : bf16
++// CHECK-PRIMITIVE-NEXT:   %[[RESULT:.*]] = arith.andi %[[CMP_ULT]], %[[CMP_ORD]] : i1
++// CHECK-PRIMITIVE-NEXT:   linalg.yield %[[RESULT]] : i1
++
++// -----
++
+ // CHECK-LABEL: func @int_cmp
+ // CHECK-PRIMITIVE-LABEL: func @int_cmp
+ func.func @int_cmp(%lhs: tensor<2x2xi32>,
 diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h b/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h
 --- stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h
 +++ stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h
@@ -1789,7 +1848,68 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToSca
      }
      return MapStablehloOpToScalarOpImpl<Args...>{}(loc, resultTypes, argTypes,
                                                     args, b);
-@@ -688,32 +690,39 @@
+@@ -456,6 +458,60 @@
+       assert(predicate.has_value() && "expected valid comparison direction");
+       return arith::CmpIOp::create(*b, loc, *predicate, lhsInt, rhsInt);
+     }
++    if (adaptor.getCompareType() &&
++        *adaptor.getCompareType() == stablehlo::ComparisonType::WEAKORDER) {
++      // Weak order treats -0.0 == +0.0 and orders all NaNs equal and > +Inf.
++      switch (comparisonDirection) {
++        case stablehlo::ComparisonDirection::LE: {
++          Value ole = arith::CmpFOp::create(*b, loc, arith::CmpFPredicate::OLE,
++                                            lhs, rhs);
++          Value rhsNan = arith::CmpFOp::create(
++              *b, loc, arith::CmpFPredicate::UNO, rhs, rhs);
++          return arith::OrIOp::create(*b, loc, ole, rhsNan);
++        }
++        case stablehlo::ComparisonDirection::GE: {
++          Value oge = arith::CmpFOp::create(*b, loc, arith::CmpFPredicate::OGE,
++                                            lhs, rhs);
++          Value lhsNan = arith::CmpFOp::create(
++              *b, loc, arith::CmpFPredicate::UNO, lhs, lhs);
++          return arith::OrIOp::create(*b, loc, oge, lhsNan);
++        }
++        case stablehlo::ComparisonDirection::LT: {
++          Value ult = arith::CmpFOp::create(*b, loc, arith::CmpFPredicate::ULT,
++                                            lhs, rhs);
++          Value lhsOrd = arith::CmpFOp::create(
++              *b, loc, arith::CmpFPredicate::ORD, lhs, lhs);
++          return arith::AndIOp::create(*b, loc, ult, lhsOrd);
++        }
++        case stablehlo::ComparisonDirection::GT: {
++          Value ugt = arith::CmpFOp::create(*b, loc, arith::CmpFPredicate::UGT,
++                                            lhs, rhs);
++          Value rhsOrd = arith::CmpFOp::create(
++              *b, loc, arith::CmpFPredicate::ORD, rhs, rhs);
++          return arith::AndIOp::create(*b, loc, ugt, rhsOrd);
++        }
++        case stablehlo::ComparisonDirection::EQ: {
++          Value oeq = arith::CmpFOp::create(*b, loc, arith::CmpFPredicate::OEQ,
++                                            lhs, rhs);
++          Value lhsNan = arith::CmpFOp::create(
++              *b, loc, arith::CmpFPredicate::UNO, lhs, lhs);
++          Value rhsNan = arith::CmpFOp::create(
++              *b, loc, arith::CmpFPredicate::UNO, rhs, rhs);
++          Value bothNan = arith::AndIOp::create(*b, loc, lhsNan, rhsNan);
++          return arith::OrIOp::create(*b, loc, oeq, bothNan);
++        }
++        case stablehlo::ComparisonDirection::NE: {
++          Value une = arith::CmpFOp::create(*b, loc, arith::CmpFPredicate::UNE,
++                                            lhs, rhs);
++          Value lhsOrd = arith::CmpFOp::create(
++              *b, loc, arith::CmpFPredicate::ORD, lhs, lhs);
++          Value rhsOrd = arith::CmpFOp::create(
++              *b, loc, arith::CmpFPredicate::ORD, rhs, rhs);
++          Value anyOrd = arith::OrIOp::create(*b, loc, lhsOrd, rhsOrd);
++          return arith::AndIOp::create(*b, loc, une, anyOrd);
++        }
++      }
++    }
+     std::optional<arith::CmpFPredicate> predicate =
+         getCmpPredicate<arith::CmpFPredicate>(comparisonDirection,
+                                               /*is_signed=*/true);
+@@ -688,32 +744,39 @@
    // A boolean value is considered to be unsigned when converting to
    // floating-point. Otherwise, it will become `-1`.
    if (IsUnsignedIntegerType{}(sourceType) &&
@@ -1842,7 +1962,7 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToSca
      }
      // No conversion is needed for identical float types.
      return args.front();
-@@ -738,17 +747,20 @@
+@@ -738,17 +801,20 @@
      auto src = cast<IntegerType>(sourceType);
      auto res = cast<IntegerType>(targetType);
      if (src.getWidth() > res.getWidth()) {
@@ -2574,7 +2694,23 @@ diff --ruN a/stablehlo/stablehlo/dialect/ChloEnums.td b/stablehlo/stablehlo/dial
  
  def CHLO_DEFAULT_COMPARISON_TYPE : NativeCodeCall<"::mlir::chlo::ComparisonTypeAttr()">;
  def CHLO_COMPARISON_TYPE_NOTYPE : I32EnumAttrCase<"NOTYPE", 0>;
-@@ -68,7 +70,9 @@
+@@ -54,6 +56,7 @@
+ def CHLO_COMPARISON_TYPE_FLOAT_TOTAL_ORDER : I32EnumAttrCase<"TOTALORDER", 2>;
+ def CHLO_COMPARISON_TYPE_SIGNED : I32EnumAttrCase<"SIGNED", 3>;
+ def CHLO_COMPARISON_TYPE_UNSIGNED : I32EnumAttrCase<"UNSIGNED", 4>;
++def CHLO_COMPARISON_TYPE_WEAK_ORDER : I32EnumAttrCase<"WEAKORDER", 5>;
+ 
+ def CHLO_ComparisonType : I32EnumAttr<"ComparisonType",
+     "Which comparison type to use.",
+@@ -62,13 +65,16 @@
+       CHLO_COMPARISON_TYPE_FLOAT,
+       CHLO_COMPARISON_TYPE_FLOAT_TOTAL_ORDER,
+       CHLO_COMPARISON_TYPE_SIGNED,
+-      CHLO_COMPARISON_TYPE_UNSIGNED
++      CHLO_COMPARISON_TYPE_UNSIGNED,
++      CHLO_COMPARISON_TYPE_WEAK_ORDER
+     ]> {
+   let genSpecializedAttr = 0;
    let cppNamespace = "::mlir::chlo";
  }
  
@@ -2585,7 +2721,7 @@ diff --ruN a/stablehlo/stablehlo/dialect/ChloEnums.td b/stablehlo/stablehlo/dial
  
  //===----------------------------------------------------------------------===//
  // Ragged dot op definitions.
-@@ -90,7 +94,9 @@
+@@ -90,7 +96,9 @@
    let cppNamespace = "::mlir::chlo";
  }
  
@@ -2776,7 +2912,23 @@ diff --ruN a/stablehlo/stablehlo/dialect/StablehloEnums.td b/stablehlo/stablehlo
  
  def STABLEHLO_DEFAULT_COMPARISON_TYPE : NativeCodeCall<"::mlir::stablehlo::ComparisonTypeAttr()">;
  def STABLEHLO_COMPARISON_TYPE_NOTYPE : I32EnumAttrCase<"NOTYPE", 0>;
-@@ -168,7 +174,9 @@
+@@ -154,6 +160,7 @@
+ def STABLEHLO_COMPARISON_TYPE_FLOAT_TOTAL_ORDER : I32EnumAttrCase<"TOTALORDER", 2>;
+ def STABLEHLO_COMPARISON_TYPE_SIGNED : I32EnumAttrCase<"SIGNED", 3>;
+ def STABLEHLO_COMPARISON_TYPE_UNSIGNED : I32EnumAttrCase<"UNSIGNED", 4>;
++def STABLEHLO_COMPARISON_TYPE_WEAK_ORDER : I32EnumAttrCase<"WEAKORDER", 5>;
+ 
+ def StableHLO_ComparisonType : I32EnumAttr<"ComparisonType",
+     "Which comparison type to use.",
+@@ -162,13 +169,16 @@
+       STABLEHLO_COMPARISON_TYPE_FLOAT,
+       STABLEHLO_COMPARISON_TYPE_FLOAT_TOTAL_ORDER,
+       STABLEHLO_COMPARISON_TYPE_SIGNED,
+-      STABLEHLO_COMPARISON_TYPE_UNSIGNED
++      STABLEHLO_COMPARISON_TYPE_UNSIGNED,
++      STABLEHLO_COMPARISON_TYPE_WEAK_ORDER
+     ]> {
+   let genSpecializedAttr = 0;
    let cppNamespace = "::mlir::stablehlo";
  }
  
@@ -2787,7 +2939,7 @@ diff --ruN a/stablehlo/stablehlo/dialect/StablehloEnums.td b/stablehlo/stablehlo
  
  // These mirror the XLA Transpose enum in Triangular Solve options.
  def STABLEHLO_TRANSPOSE_INVALID : I32EnumAttrCase<"TRANSPOSE_INVALID", 0>;
-@@ -188,7 +196,9 @@
+@@ -188,7 +198,9 @@
    let cppNamespace = "::mlir::stablehlo";
  }
  
@@ -2798,7 +2950,7 @@ diff --ruN a/stablehlo/stablehlo/dialect/StablehloEnums.td b/stablehlo/stablehlo
  
  def STABLEHLO_RNG_DISTRIBUTION_UNIFORM : I32EnumAttrCase<"UNIFORM", 1>;
  def STABLEHLO_RNG_DISTRIBUTION_NORMAL : I32EnumAttrCase<"NORMAL", 2>;
-@@ -203,7 +213,9 @@
+@@ -203,7 +215,9 @@
    let cppNamespace = "::mlir::stablehlo";
  }
  
@@ -2809,7 +2961,7 @@ diff --ruN a/stablehlo/stablehlo/dialect/StablehloEnums.td b/stablehlo/stablehlo
  
  def STABLEHLO_RNG_ALGORITHM_DEFAULT : I32EnumAttrCase<"DEFAULT", 0>;
  def STABLEHLO_RNG_ALGORITHM_THREE_FRY : I32EnumAttrCase<"THREE_FRY", 1>;
-@@ -220,6 +232,8 @@
+@@ -220,6 +234,8 @@
    let cppNamespace = "::mlir::stablehlo";
  }
  
@@ -2822,7 +2974,19 @@ diff --ruN a/stablehlo/stablehlo/dialect/StablehloEnums.td b/stablehlo/stablehlo
 diff --ruN a/stablehlo/stablehlo/dialect/StablehloOps.cpp b/stablehlo/stablehlo/dialect/StablehloOps.cpp
 --- stablehlo/stablehlo/dialect/StablehloOps.cpp
 +++ stablehlo/stablehlo/dialect/StablehloOps.cpp
-@@ -4332,6 +4332,23 @@
+@@ -3283,7 +3283,10 @@
+     RegionRange regions,
+     SmallVectorImpl<ShapedTypeComponents>& inferredReturnShapes) {
+   CompareOp::Adaptor adaptor(operands, attributes, properties, regions);
+-  return hlo::inferCompareOp(context, location, adaptor.getLhs(),
++  std::optional<StringRef> compareType;
++  if (auto attr = adaptor.getCompareType())
++    compareType = stringifyComparisonType(*attr);
++  return hlo::inferCompareOp(context, location, adaptor.getLhs(), compareType,
+                              inferredReturnShapes);
+ }
+ 
+@@ -4332,6 +4335,23 @@
    ReturnOp::create(*builder, loc, compare);
  }
  
@@ -2846,7 +3010,7 @@ diff --ruN a/stablehlo/stablehlo/dialect/StablehloOps.cpp b/stablehlo/stablehlo/
  void buildMaxAndArgmaxBody(Type elementType, Type indices_type, Region& body,
                             OpBuilder& builder) {
    OpBuilder::InsertionGuard guard(builder);
-@@ -4366,15 +4383,18 @@
+@@ -4366,15 +4386,18 @@
    // Final lhs Selection Condition: (gt_pred) OR (tie_breaker_condition)
    auto final_lhs_condition =
        OrOp::create(builder, loc, gt_pred, tie_breaker_condition).getResult();
@@ -2895,6 +3059,47 @@ diff --ruN a/stablehlo/stablehlo/dialect/StablehloOps.td b/stablehlo/stablehlo/d
      }]>,
      OpBuilder<(ins
        "::mlir::TypeRange":$resultTypes, "::mlir::ValueRange":$inputs,
+diff --ruN a/stablehlo/stablehlo/dialect/TypeInference.cpp b/stablehlo/stablehlo/dialect/TypeInference.cpp
+--- stablehlo/stablehlo/dialect/TypeInference.cpp
++++ stablehlo/stablehlo/dialect/TypeInference.cpp
+@@ -2032,11 +2032,22 @@
+ }
+ 
+ LogicalResult inferCompareOp(
+-    MLIRContext* context, std::optional<Location>, Value lhs,
++    MLIRContext* context, std::optional<Location> location, Value lhs,
++    std::optional<StringRef> compareType,
+     SmallVectorImpl<ShapedTypeComponents>& inferredReturnShapes) {
++  auto argTy = cast<RankedTensorType>(lhs.getType());
++  // compare_c3
++  if (compareType && *compareType == "WEAKORDER" &&
++      !isa<FloatType>(argTy.getElementType())) {
++    return emitOptionalError(
++        location,
++        "WEAKORDER comparison type is only supported for floating-point "
++        "element types, got ",
++        argTy.getElementType());
++  }
++
+   // compare_c1
+   ShapedTypeComponents& components = inferredReturnShapes.emplace_back();
+-  auto argTy = cast<RankedTensorType>(lhs.getType());
+   auto resElementTy = IntegerType::get(context, /*width=*/1);
+   // compare_c2
+   components =
+diff --ruN a/stablehlo/stablehlo/dialect/TypeInference.h b/stablehlo/stablehlo/dialect/TypeInference.h
+--- stablehlo/stablehlo/dialect/TypeInference.h
++++ stablehlo/stablehlo/dialect/TypeInference.h
+@@ -187,7 +187,8 @@
+     SmallVectorImpl<Type>& inferredReturnTypes);
+ 
+ LogicalResult inferCompareOp(
+-    MLIRContext* context, std::optional<Location>, Value lhs,
++    MLIRContext* context, std::optional<Location> location, Value lhs,
++    std::optional<StringRef> compareType,
+     SmallVectorImpl<ShapedTypeComponents>& inferredReturnShapes);
+ 
+ LogicalResult inferComplexOp(std::optional<Location> location, Value lhs,
 diff --ruN a/stablehlo/stablehlo/dialect/Version.h b/stablehlo/stablehlo/dialect/Version.h
 --- stablehlo/stablehlo/dialect/Version.h
 +++ stablehlo/stablehlo/dialect/Version.h
@@ -2903,7 +3108,7 @@ diff --ruN a/stablehlo/stablehlo/dialect/Version.h b/stablehlo/stablehlo/dialect
  
    /// Return a Version representing the current VHLO dialect version.
 -  static Version getCurrentVersion() { return Version(1, 20, 1); }
-+  static Version getCurrentVersion() { return Version(1, 21, 0); }
++  static Version getCurrentVersion() { return Version(1, 22, 0); }
  
    /// Return a Version representing the minimum supported VHLO dialect version.
    static Version getMinimumVersion() { return Version(0, 9, 0); }
@@ -2922,11 +3127,12 @@ diff --ruN a/stablehlo/stablehlo/dialect/VhloBytecode.h b/stablehlo/stablehlo/di
 diff --ruN a/stablehlo/stablehlo/dialect/VhloDialect.td b/stablehlo/stablehlo/dialect/VhloDialect.td
 --- stablehlo/stablehlo/dialect/VhloDialect.td
 +++ stablehlo/stablehlo/dialect/VhloDialect.td
-@@ -59,10 +59,12 @@
+@@ -59,10 +59,13 @@
        1.18.0: Add `result_tilings` attribute to `custom_call` op.
        1.19.0: Add CollectiveReduceOp.
        1.20.0: Add has_dynamic_root to collective_broadcast op. Allow mixed fp8 operands in `convolution` and `dynamic_conv` ops.
 +      1.21.0: Add F8E4M3FN_F8E4M3FN_F32_X3 and F8E4M3FN_F8E4M3FN_F32_X4 dot algorithms.
++      1.22.0: Add WEAKORDER comparison type to compare op.
    }];
  
    let useDefaultAttributePrinterParser = 0;
@@ -2938,14 +3144,59 @@ diff --ruN a/stablehlo/stablehlo/dialect/VhloDialect.td b/stablehlo/stablehlo/di
 diff --ruN a/stablehlo/stablehlo/dialect/VhloEnums.td b/stablehlo/stablehlo/dialect/VhloEnums.td
 --- stablehlo/stablehlo/dialect/VhloEnums.td
 +++ stablehlo/stablehlo/dialect/VhloEnums.td
-@@ -43,6 +43,7 @@
+@@ -41,11 +41,16 @@
+   let cppNamespace = "::mlir::vhlo";
+ }
  
- class VHLO_EnumAttr<EnumAttrInfo enumInfo, string name, string minVersion, string maxVersion>
+-class VHLO_EnumAttr<EnumAttrInfo enumInfo, string name, string minVersion, string maxVersion>
++class VHLO_EnumAttr<EnumAttrInfo enumInfo, string name, string minVersion, string maxVersion, code customMinVersion = "">
    : EnumAttr<VHLO_Dialect, enumInfo, name, [VHLO_VersionedAttrInterface]> {
 +  let assemblyFormat = "$value";
    let extraClassDeclaration = [{
      mlir::vhlo::Version getMinVersion() {
-       return mlir::vhlo::Version(}] # !subst(".", ", ", minVersion) # [{);
+-      return mlir::vhlo::Version(}] # !subst(".", ", ", minVersion) # [{);
++      }] # !if(
++        !empty(customMinVersion),
++        [{ return mlir::vhlo::Version(}] # !subst(".", ", ", minVersion) # [{); }],
++        customMinVersion
++      ) # [{
+     }
+     mlir::vhlo::Version getMaxVersion() {
+       }] # !if(
+@@ -90,17 +95,31 @@
+ def VHLO_COMPARISON_TYPE_V1_FLOAT_TOTAL_ORDER : I32EnumAttrCase<"TOTALORDER", 2>;
+ def VHLO_COMPARISON_TYPE_V1_SIGNED : I32EnumAttrCase<"SIGNED", 3>;
+ def VHLO_COMPARISON_TYPE_V1_UNSIGNED : I32EnumAttrCase<"UNSIGNED", 4>;
++def VHLO_COMPARISON_TYPE_V1_WEAK_ORDER : I32EnumAttrCase<"WEAKORDER", 5>;
+ 
+ def VHLO_ComparisonTypeV1 : VHLO_I32EnumAttr<"ComparisonTypeV1", [
+   VHLO_COMPARISON_TYPE_V1_NOTYPE,
+   VHLO_COMPARISON_TYPE_V1_FLOAT,
+   VHLO_COMPARISON_TYPE_V1_FLOAT_TOTAL_ORDER,
+   VHLO_COMPARISON_TYPE_V1_SIGNED,
+-  VHLO_COMPARISON_TYPE_V1_UNSIGNED
++  VHLO_COMPARISON_TYPE_V1_UNSIGNED,
++  VHLO_COMPARISON_TYPE_V1_WEAK_ORDER
+ ]> {}
+ 
+ def VHLO_ComparisonTypeAttrV1
+-  : VHLO_EnumAttr<VHLO_ComparisonTypeV1, "comparison_type_v1", "0.9.0", "current">;
++  : VHLO_EnumAttr<VHLO_ComparisonTypeV1, "comparison_type_v1", "0.9.0", "current", [{
++      switch (getValue()) {
++        case mlir::vhlo::ComparisonTypeV1::NOTYPE:
++        case mlir::vhlo::ComparisonTypeV1::FLOAT:
++        case mlir::vhlo::ComparisonTypeV1::TOTALORDER:
++        case mlir::vhlo::ComparisonTypeV1::SIGNED:
++        case mlir::vhlo::ComparisonTypeV1::UNSIGNED:
++          return mlir::vhlo::Version(0, 9, 0);
++        case mlir::vhlo::ComparisonTypeV1::WEAKORDER:
++          return mlir::vhlo::Version(1, 22, 0);
++      }
++      llvm_unreachable("Unknown ComparisonTypeV1");
++    }]>;
+ 
+ //===----------------------------------------------------------------------===//
+ // CustomCallApiVersion
 diff --ruN a/stablehlo/stablehlo/dialect/VhloOps.cpp b/stablehlo/stablehlo/dialect/VhloOps.cpp
 --- stablehlo/stablehlo/dialect/VhloOps.cpp
 +++ stablehlo/stablehlo/dialect/VhloOps.cpp
@@ -3456,6 +3707,66 @@ diff --ruN a/stablehlo/stablehlo/reference/Configuration.h b/stablehlo/stablehlo
  
    virtual ~InterpreterFallback() = default;
  };
+diff --ruN a/stablehlo/stablehlo/reference/Element.cpp b/stablehlo/stablehlo/reference/Element.cpp
+--- stablehlo/stablehlo/reference/Element.cpp
++++ stablehlo/stablehlo/reference/Element.cpp
+@@ -1159,6 +1159,40 @@
+       [](mlir::Complex<double> e) { return std::tanh(e); });
+ }
+ 
++Element weakOrderEq(const Element& e1, const Element& e2) {
++  auto type = e1.getType();
++  auto i1Type = IntegerType::get(type.getContext(), 1);
++  if (type != e2.getType())
++    report_fatal_error(invalidArgument("Element types don't match: %s vs %s",
++                                       debugString(type).c_str(),
++                                       debugString(e2.getType()).c_str()));
++  if (!isSupportedFloatType(type))
++    report_fatal_error(invalidArgument("Unsupported element type: %s",
++                                       debugString(type).c_str()));
++  auto floatLhs = e1.getFloatValue();
++  auto floatRhs = e2.getFloatValue();
++  return Element(
++      i1Type, (floatLhs == floatRhs) || (floatLhs.isNaN() && floatRhs.isNaN()));
++}
++
++Element weakOrderLt(const Element& e1, const Element& e2) {
++  auto type = e1.getType();
++  auto i1Type = IntegerType::get(type.getContext(), 1);
++  if (type != e2.getType())
++    report_fatal_error(invalidArgument("Element types don't match: %s vs %s",
++                                       debugString(type).c_str(),
++                                       debugString(e2.getType()).c_str()));
++  if (!isSupportedFloatType(type))
++    report_fatal_error(invalidArgument("Unsupported element type: %s",
++                                       debugString(type).c_str()));
++  auto floatLhs = e1.getFloatValue();
++  auto floatRhs = e2.getFloatValue();
++  bool result =
++      (!floatLhs.isNaN() && floatRhs.isNaN()) ||
++      (!floatLhs.isNaN() && !floatRhs.isNaN() && (floatLhs < floatRhs));
++  return Element(i1Type, result);
++}
++
+ void Element::print(raw_ostream& os, bool elideType) const {
+   if (isSupportedIntegerType(type_)) {
+     IntegerAttr::get(type_, getIntegerValue()).print(os, elideType);
+diff --ruN a/stablehlo/stablehlo/reference/Element.h b/stablehlo/stablehlo/reference/Element.h
+--- stablehlo/stablehlo/reference/Element.h
++++ stablehlo/stablehlo/reference/Element.h
+@@ -314,6 +314,12 @@
+ /// Returns tanh of Element object.
+ Element tanh(const Element& e);
+ 
++/// Returns weak-order equality of two floating-point Element objects.
++Element weakOrderEq(const Element& e1, const Element& e2);
++
++/// Returns weak-order less-than of two floating-point Element objects.
++Element weakOrderLt(const Element& e1, const Element& e2);
++
+ /// Print utilities for Element objects.
+ inline raw_ostream& operator<<(raw_ostream& os, Element element) {
+   element.print(os, /*elideType=*/true);
 diff --ruN a/stablehlo/stablehlo/reference/Errors.h b/stablehlo/stablehlo/reference/Errors.h
 --- stablehlo/stablehlo/reference/Errors.h
 +++ stablehlo/stablehlo/reference/Errors.h
@@ -3782,7 +4093,81 @@ diff --ruN a/stablehlo/stablehlo/reference/InterpreterOps.cpp b/stablehlo/stable
 diff --ruN a/stablehlo/stablehlo/reference/Ops.cpp b/stablehlo/stablehlo/reference/Ops.cpp
 --- stablehlo/stablehlo/reference/Ops.cpp
 +++ stablehlo/stablehlo/reference/Ops.cpp
-@@ -1777,12 +1777,8 @@
+@@ -568,7 +568,9 @@
+       auto lhs = scope.findTensor(op.getLhs());
+       auto rhs = scope.findTensor(op.getRhs());
+       auto comparisonDirection = op.getComparisonDirection();
+-      auto result = compareOp(lhs, rhs, comparisonDirection, op.getType());
++      auto compareType = op.getCompareType();
++      auto result =
++          compareOp(lhs, rhs, comparisonDirection, compareType, op.getType());
+       scope.add(op.getResult(), result);
+     } else if (auto op = dyn_cast<ComplexOp>(operation)) {
+       auto lhs = scope.findTensor(op.getLhs());
+@@ -1597,27 +1599,56 @@
+ 
+ Tensor compareOp(const Tensor& lhs, const Tensor& rhs,
+                  ComparisonDirection comparisonDirection,
++                 std::optional<ComparisonType> compareType,
+                  ShapedType resultType) {
+   Tensor result(resultType);
++  bool isWeakOrder = compareType == ComparisonType::WEAKORDER;
+   for (auto it = result.index_begin(); it != result.index_end(); ++it) {
++    auto lhsElem = lhs.get(*it);
++    auto rhsElem = rhs.get(*it);
++    if (isWeakOrder) {
++      switch (comparisonDirection) {
++        case ComparisonDirection::EQ:
++          result.set(*it, weakOrderEq(lhsElem, rhsElem));
++          break;
++        case ComparisonDirection::NE:
++          result.set(*it, !weakOrderEq(lhsElem, rhsElem));
++          break;
++        case ComparisonDirection::GE:
++          result.set(*it, weakOrderLt(rhsElem, lhsElem) ||
++                              weakOrderEq(lhsElem, rhsElem));
++          break;
++        case ComparisonDirection::GT:
++          result.set(*it, weakOrderLt(rhsElem, lhsElem));
++          break;
++        case ComparisonDirection::LE:
++          result.set(*it, weakOrderLt(lhsElem, rhsElem) ||
++                              weakOrderEq(lhsElem, rhsElem));
++          break;
++        case ComparisonDirection::LT:
++          result.set(*it, weakOrderLt(lhsElem, rhsElem));
++          break;
++      }
++      continue;
++    }
+     switch (comparisonDirection) {
+       case ComparisonDirection::EQ:
+-        result.set(*it, lhs.get(*it) == rhs.get(*it));
++        result.set(*it, lhsElem == rhsElem);
+         break;
+       case ComparisonDirection::NE:
+-        result.set(*it, lhs.get(*it) != rhs.get(*it));
++        result.set(*it, lhsElem != rhsElem);
+         break;
+       case ComparisonDirection::GE:
+-        result.set(*it, lhs.get(*it) >= rhs.get(*it));
++        result.set(*it, lhsElem >= rhsElem);
+         break;
+       case ComparisonDirection::GT:
+-        result.set(*it, lhs.get(*it) > rhs.get(*it));
++        result.set(*it, lhsElem > rhsElem);
+         break;
+       case ComparisonDirection::LE:
+-        result.set(*it, lhs.get(*it) <= rhs.get(*it));
++        result.set(*it, lhsElem <= rhsElem);
+         break;
+       case ComparisonDirection::LT:
+-        result.set(*it, lhs.get(*it) < rhs.get(*it));
++        result.set(*it, lhsElem < rhsElem);
+         break;
+     }
+   }
+@@ -1777,12 +1808,8 @@
          dotGeneralOp(reversedLhsWindow, rhs, lhsContractingDimensions,
                       rhsContractingDimensions, resultType.getElementType());
  
@@ -3797,6 +4182,17 @@ diff --ruN a/stablehlo/stablehlo/reference/Ops.cpp b/stablehlo/stablehlo/referen
      Axes resultPermutation;
      resultPermutation.push_back(outputBatchDimension);
      resultPermutation.append(outputSpatialDimensions.begin(),
+diff --ruN a/stablehlo/stablehlo/reference/Ops.h b/stablehlo/stablehlo/reference/Ops.h
+--- stablehlo/stablehlo/reference/Ops.h
++++ stablehlo/stablehlo/reference/Ops.h
+@@ -72,6 +72,7 @@
+                            ChannelId channelId, Process* process);
+ Tensor compareOp(const Tensor& lhs, const Tensor& rhs,
+                  ComparisonDirection comparisonDirection,
++                 std::optional<ComparisonType> compareType,
+                  ShapedType resultType);
+ Tensor complexOp(const Tensor& lhs, const Tensor& rhs, ShapedType resultType);
+ Tensor concatenateOp(ArrayRef<Tensor> inputs, Axis dimension,
 diff --ruN a/stablehlo/stablehlo/reference/Process.cpp b/stablehlo/stablehlo/reference/Process.cpp
 --- stablehlo/stablehlo/reference/Process.cpp
 +++ stablehlo/stablehlo/reference/Process.cpp
@@ -4286,6 +4682,78 @@ diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/st
  // CHECK-LABEL:   func.func @mulhi_i32(
  // CHECK-SAME:      %[[ARG0:.*]]: tensor<4xi32>,
  // CHECK-SAME:      %[[ARG1:.*]]: tensor<4xi32>) -> tensor<4xi32> {
+diff --ruN a/stablehlo/stablehlo/tests/interpret/compare.mlir b/stablehlo/stablehlo/tests/interpret/compare.mlir
+--- stablehlo/stablehlo/tests/interpret/compare.mlir
++++ stablehlo/stablehlo/tests/interpret/compare.mlir
+@@ -294,6 +294,68 @@
+ 
+ // -----
+ 
++func.func @compare_op_test_f64_weakorder_eq() {
++  // -NaN, -NaN, -Inf, -0.0, 1.0, +Inf, +Inf, -NaN
++  // -NaN, +NaN, +Inf, +0.0, 2.0, +Inf, -NaN, +Inf
++  %lhs = stablehlo.constant dense<[0xFFF0000000000001, 0xFFF0000000000001, 0xFFF0000000000000, 0x8000000000000000, 1.0, 0x7FF0000000000000, 0x7FF0000000000000, 0xFFF0000000000001]> : tensor<8xf64>
++  %rhs = stablehlo.constant dense<[0xFFF0000000000001, 0x7FF0000000000001, 0x7FF0000000000000, 0x0000000000000000, 2.0, 0x7FF0000000000000, 0xFFF0000000000001, 0x7FF0000000000000]> : tensor<8xf64>
++  %result = stablehlo.compare EQ, %lhs, %rhs, WEAKORDER : (tensor<8xf64>, tensor<8xf64>) -> tensor<8xi1>
++  check.expect_eq_const %result, dense<[true, true, false, true, false, true, false, false]> : tensor<8xi1>
++  func.return
++}
++
++// -----
++
++func.func @compare_op_test_f64_weakorder_ne() {
++  %lhs = stablehlo.constant dense<[0xFFF0000000000001, 0xFFF0000000000001, 0xFFF0000000000000, 0x8000000000000000, 1.0, 0x7FF0000000000000, 0x7FF0000000000000, 0xFFF0000000000001]> : tensor<8xf64>
++  %rhs = stablehlo.constant dense<[0xFFF0000000000001, 0x7FF0000000000001, 0x7FF0000000000000, 0x0000000000000000, 2.0, 0x7FF0000000000000, 0xFFF0000000000001, 0x7FF0000000000000]> : tensor<8xf64>
++  %result = stablehlo.compare NE, %lhs, %rhs, WEAKORDER : (tensor<8xf64>, tensor<8xf64>) -> tensor<8xi1>
++  check.expect_eq_const %result, dense<[false, false, true, false, true, false, true, true]> : tensor<8xi1>
++  func.return
++}
++
++// -----
++
++func.func @compare_op_test_f64_weakorder_ge() {
++  %lhs = stablehlo.constant dense<[0xFFF0000000000001, 0xFFF0000000000001, 0xFFF0000000000000, 0x8000000000000000, 1.0, 0x7FF0000000000000, 0x7FF0000000000000, 0xFFF0000000000001]> : tensor<8xf64>
++  %rhs = stablehlo.constant dense<[0xFFF0000000000001, 0x7FF0000000000001, 0x7FF0000000000000, 0x0000000000000000, 2.0, 0x7FF0000000000000, 0xFFF0000000000001, 0x7FF0000000000000]> : tensor<8xf64>
++  %result = stablehlo.compare GE, %lhs, %rhs, WEAKORDER : (tensor<8xf64>, tensor<8xf64>) -> tensor<8xi1>
++  check.expect_eq_const %result, dense<[true, true, false, true, false, true, false, true]> : tensor<8xi1>
++  func.return
++}
++
++// -----
++
++func.func @compare_op_test_f64_weakorder_gt() {
++  %lhs = stablehlo.constant dense<[0xFFF0000000000001, 0xFFF0000000000001, 0xFFF0000000000000, 0x8000000000000000, 1.0, 0x7FF0000000000000, 0x7FF0000000000000, 0xFFF0000000000001]> : tensor<8xf64>
++  %rhs = stablehlo.constant dense<[0xFFF0000000000001, 0x7FF0000000000001, 0x7FF0000000000000, 0x0000000000000000, 2.0, 0x7FF0000000000000, 0xFFF0000000000001, 0x7FF0000000000000]> : tensor<8xf64>
++  %result = stablehlo.compare GT, %lhs, %rhs, WEAKORDER : (tensor<8xf64>, tensor<8xf64>) -> tensor<8xi1>
++  check.expect_eq_const %result, dense<[false, false, false, false, false, false, false, true]> : tensor<8xi1>
++  func.return
++}
++
++// -----
++
++func.func @compare_op_test_f64_weakorder_le() {
++  %lhs = stablehlo.constant dense<[0xFFF0000000000001, 0xFFF0000000000001, 0xFFF0000000000000, 0x8000000000000000, 1.0, 0x7FF0000000000000, 0x7FF0000000000000, 0xFFF0000000000001]> : tensor<8xf64>
++  %rhs = stablehlo.constant dense<[0xFFF0000000000001, 0x7FF0000000000001, 0x7FF0000000000000, 0x0000000000000000, 2.0, 0x7FF0000000000000, 0xFFF0000000000001, 0x7FF0000000000000]> : tensor<8xf64>
++  %result = stablehlo.compare LE, %lhs, %rhs, WEAKORDER : (tensor<8xf64>, tensor<8xf64>) -> tensor<8xi1>
++  check.expect_eq_const %result, dense<[true, true, true, true, true, true, true, false]> : tensor<8xi1>
++  func.return
++}
++
++// -----
++
++func.func @compare_op_test_f64_weakorder_lt() {
++  %lhs = stablehlo.constant dense<[0xFFF0000000000001, 0xFFF0000000000001, 0xFFF0000000000000, 0x8000000000000000, 1.0, 0x7FF0000000000000, 0x7FF0000000000000, 0xFFF0000000000001]> : tensor<8xf64>
++  %rhs = stablehlo.constant dense<[0xFFF0000000000001, 0x7FF0000000000001, 0x7FF0000000000000, 0x0000000000000000, 2.0, 0x7FF0000000000000, 0xFFF0000000000001, 0x7FF0000000000000]> : tensor<8xf64>
++  %result = stablehlo.compare LT, %lhs, %rhs, WEAKORDER : (tensor<8xf64>, tensor<8xf64>) -> tensor<8xi1>
++  check.expect_eq_const %result, dense<[false, false, true, false, true, false, true, false]> : tensor<8xi1>
++  func.return
++}
++
++// -----
++
+ func.func @compare_op_test_c128_default() {
+   // (+NaN, +0.0)
+   // (+NaN, -0.0)
 diff --ruN a/stablehlo/stablehlo/tests/interpret/convolution.mlir b/stablehlo/stablehlo/tests/interpret/convolution.mlir
 --- stablehlo/stablehlo/tests/interpret/convolution.mlir
 +++ stablehlo/stablehlo/tests/interpret/convolution.mlir
@@ -4393,6 +4861,26 @@ diff --ruN a/stablehlo/stablehlo/tests/ops_chlo.mlir b/stablehlo/stablehlo/tests
  // CHECK-LABEL: func @mulhi_i32
  func.func @mulhi_i32(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi32> {
    %0 = "chlo.mulhi"(%arg0, %arg1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32>
+diff --ruN a/stablehlo/stablehlo/tests/ops_chlo_roundtrip.mlir b/stablehlo/stablehlo/tests/ops_chlo_roundtrip.mlir
+--- stablehlo/stablehlo/tests/ops_chlo_roundtrip.mlir
++++ stablehlo/stablehlo/tests/ops_chlo_roundtrip.mlir
+@@ -367,12 +367,14 @@
+ // CHECK-SAME:  %[[A1:.*]]: tensor<2x3x4xf64>
+ // CHECK:       %[[T0:.*]] = chlo.broadcast_compare %[[A0]], %[[A1]] {comparison_direction = #chlo<comparison_direction LT>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
+ // CHECK:       %[[T1:.*]] = chlo.broadcast_compare %[[A0]], %[[A1]] {compare_type = #chlo<comparison_type TOTALORDER>, comparison_direction = #chlo<comparison_direction LT>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
+-// CHECK:       %[[T2:.*]] = chlo.broadcast_compare %[[A0]], %[[A1]] {broadcast_dimensions = array<i64: 1>, comparison_direction = #chlo<comparison_direction LT>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
++// CHECK:       %[[T2:.*]] = chlo.broadcast_compare %[[A0]], %[[A1]] {compare_type = #chlo<comparison_type WEAKORDER>, comparison_direction = #chlo<comparison_direction LT>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
++// CHECK:       %[[T3:.*]] = chlo.broadcast_compare %[[A0]], %[[A1]] {broadcast_dimensions = array<i64: 1>, comparison_direction = #chlo<comparison_direction LT>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
+ // CHECK:       return %[[T0]] : tensor<2x3x4xi1>
+ func.func @chlo_broadcast_compare(%arg0: tensor<2x3x4xf64>, %arg1: tensor<2x3x4xf64>) -> tensor<2x3x4xi1> {
+   %0 = chlo.broadcast_compare %arg0, %arg1 {comparison_direction = #chlo<comparison_direction LT>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
+   %1 = chlo.broadcast_compare %arg0, %arg1 {comparison_direction = #chlo<comparison_direction LT>, compare_type = #chlo<comparison_type TOTALORDER>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
+-  %2 = chlo.broadcast_compare %arg0, %arg1 {comparison_direction = #chlo<comparison_direction LT>, broadcast_dimensions = array<i64: 1>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
++  %2 = chlo.broadcast_compare %arg0, %arg1 {comparison_direction = #chlo<comparison_direction LT>, compare_type = #chlo<comparison_type WEAKORDER>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
++  %3 = chlo.broadcast_compare %arg0, %arg1 {comparison_direction = #chlo<comparison_direction LT>, broadcast_dimensions = array<i64: 1>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
+   return %0 : tensor<2x3x4xi1>
+ }
+ 
 diff --ruN a/stablehlo/stablehlo/tests/ops_dot_general_algorithms.mlir b/stablehlo/stablehlo/tests/ops_dot_general_algorithms.mlir
 --- stablehlo/stablehlo/tests/ops_dot_general_algorithms.mlir
 +++ stablehlo/stablehlo/tests/ops_dot_general_algorithms.mlir
@@ -4437,6 +4925,39 @@ diff --ruN a/stablehlo/stablehlo/tests/ops_dot_general_algorithms.mlir b/stableh
  // CHECK-LABEL: func @dot_algorithm_bf16_bf16_f32_x6
  func.func @dot_algorithm_bf16_bf16_f32_x6(%arg0: tensor<2x2x2xbf16>, %arg1: tensor<2x2x2xbf16>) -> tensor<2x2x2xbf16> {
    %0 = "stablehlo.dot_general"(%arg0, %arg1) <{
+diff --ruN a/stablehlo/stablehlo/tests/ops_stablehlo.mlir b/stablehlo/stablehlo/tests/ops_stablehlo.mlir
+--- stablehlo/stablehlo/tests/ops_stablehlo.mlir
++++ stablehlo/stablehlo/tests/ops_stablehlo.mlir
+@@ -1548,6 +1548,29 @@
+ 
+ // -----
+ 
++// CHECK-LABEL: func @compare_weakorder_f32
++func.func @compare_weakorder_f32(%arg0: tensor<3xf32>, %arg1: tensor<3xf32>) -> tensor<3xi1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction LT>,
++    compare_type = #stablehlo<comparison_type WEAKORDER>
++  } : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xi1>
++  func.return %0 : tensor<3xi1>
++}
++
++// -----
++
++func.func @compare_weakorder_invalid_i32(%arg0: tensor<3xi32>, %arg1: tensor<3xi32>) -> tensor<3xi1> {
++  // expected-error@+2 {{failed to infer returned types}}
++  // expected-error@+1 {{WEAKORDER comparison type is only supported for floating-point element types, got 'i32'}}
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction LT>,
++    compare_type = #stablehlo<comparison_type WEAKORDER>
++  } : (tensor<3xi32>, tensor<3xi32>) -> tensor<3xi1>
++  func.return %0 : tensor<3xi1>
++}
++
++// -----
++
+ // CHECK-LABEL: func @collective_broadcast_c3
+ func.func @collective_broadcast_c3(%operand: tensor<16x8xf32>) -> tensor<16x8xf32> {
+   %0 = "stablehlo.collective_broadcast"(%operand) {
 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
@@ -4473,6 +4994,41 @@ diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compat
  
    // CHECK-LABEL: @all_reduce_sdy_mesh_max
    func.func @all_reduce_sdy_mesh_max(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_aggressive_folder.mlir b/stablehlo/stablehlo/tests/transforms/stablehlo_aggressive_folder.mlir
+--- stablehlo/stablehlo/tests/transforms/stablehlo_aggressive_folder.mlir
++++ stablehlo/stablehlo/tests/transforms/stablehlo_aggressive_folder.mlir
+@@ -560,6 +560,31 @@
+ 
+ // -----
+ 
++// CHECK-LABEL: func.func @compare_fold_weakorder
++func.func @compare_fold_weakorder()
++  -> (tensor<i1>, tensor<i1>, tensor<i1>, tensor<i1>, tensor<i1>, tensor<i1>) {
++  %neg_zero = stablehlo.constant dense<-0.0> : tensor<f32>
++  %pos_zero = stablehlo.constant dense<0.0> : tensor<f32>
++  %pos_inf = stablehlo.constant dense<0x7F800000> : tensor<f32>
++  %neg_nan = stablehlo.constant dense<0xFFC00000> : tensor<f32>
++  %pos_nan = stablehlo.constant dense<0x7FC00000> : tensor<f32>
++
++  %0 = stablehlo.compare LT, %neg_zero, %pos_zero, WEAKORDER : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  %1 = stablehlo.compare EQ, %neg_zero, %pos_zero, WEAKORDER : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  %2 = stablehlo.compare EQ, %neg_nan, %pos_nan, WEAKORDER : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  %3 = stablehlo.compare LT, %neg_nan, %pos_nan, WEAKORDER : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  %4 = stablehlo.compare LT, %pos_inf, %neg_nan, WEAKORDER : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  %5 = stablehlo.compare GT, %neg_nan, %pos_inf, WEAKORDER : (tensor<f32>, tensor<f32>) -> tensor<i1>
++
++  // CHECK-DAG:  [[FALSE:%.+]] = stablehlo.constant dense<false> : tensor<i1>
++  // CHECK-DAG:  [[TRUE:%.+]] = stablehlo.constant dense<true> : tensor<i1>
++  // CHECK-NEXT: return [[FALSE]], [[TRUE]], [[TRUE]], [[FALSE]], [[TRUE]], [[TRUE]]
++  return %0, %1, %2, %3, %4, %5 :
++         tensor<i1>, tensor<i1>, tensor<i1>, tensor<i1>, tensor<i1>, tensor<i1>
++}
++
++// -----
++
+ ////////
+ // ConcatenateOp
+ 
 diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir b/stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir
 --- stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir
 +++ stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir
@@ -8083,10 +8639,3470 @@ diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_21_0.ml
 +}
 diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_21_0.mlir.bc b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_21_0.mlir.bc
 Binary files /dev/null and stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_21_0.mlir.bc differ
+diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_22_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_22_0.mlir
+--- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_22_0.mlir
++++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_22_0.mlir
+@@ -0,0 +1,3436 @@
++// RUN: stablehlo-opt --mlir-print-op-generic %s.bc | FileCheck %s
++// RUN: stablehlo-translate --deserialize %s.bc | stablehlo-translate --serialize --target=1.22.0 | stablehlo-opt --mlir-print-op-generic | FileCheck %s
++// RUN: stablehlo-translate --deserialize %s.bc | stablehlo-opt > %t.0
++// RUN: stablehlo-opt --strip-debuginfo %s > %t.1
++// RUN: diff %t.0 %t.1
++// RUN: stablehlo-translate --serialize --target=1.22.0 --strip-debuginfo %s > %t.2
++// RUN: diff %s.bc %t.2
++// RUN: %if asserts %{ stablehlo-opt --stablehlo-legalize-to-vhlo -emit-bytecode -debug-only=vhlo-bytecode %s 2>&1 | FileCheck --check-prefix=CHECK-WARN %s %}
++// RUN: %if asserts %{ stablehlo-opt --stablehlo-legalize-to-vhlo -emit-bytecode %s | stablehlo-opt -debug-only=vhlo-bytecode 2>&1 | FileCheck --check-prefix=CHECK-WARN %s %}
++
++// CHECK-WARN-NOT: Not Implemented
++
++func.func private @mesh()
++
++// ============ ATTRIBUTES ============
++
++// CHECK-LABEL: "attr_comparison_direction_eq"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
++    comparison_direction = #stablehlo<comparison_direction EQ>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_direction_ne"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
++    comparison_direction = #stablehlo<comparison_direction NE>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_direction_ge"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
++    comparison_direction = #stablehlo<comparison_direction GE>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_direction_gt"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
++    comparison_direction = #stablehlo<comparison_direction GT>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_direction_le"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
++    comparison_direction = #stablehlo<comparison_direction LE>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_direction_lt"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
++    comparison_direction = #stablehlo<comparison_direction LT>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_type_notype"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction EQ>
++    // CHECK: compare_type = #vhlo<comparison_type_v1 NOTYPE>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_type_float"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction EQ>,
++    // CHECK: compare_type = #vhlo<comparison_type_v1 FLOAT>,
++    compare_type = #stablehlo<comparison_type FLOAT>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_type_totalorder"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction EQ>,
++    // CHECK: compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
++    compare_type = #stablehlo<comparison_type TOTALORDER>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_type_signed"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction EQ>,
++    // CHECK: compare_type = #vhlo<comparison_type_v1 SIGNED>,
++    compare_type = #stablehlo<comparison_type SIGNED>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_type_unsigned"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction EQ>,
++    // CHECK: compare_type = #vhlo<comparison_type_v1 UNSIGNED>,
++    compare_type = #stablehlo<comparison_type UNSIGNED>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_type_weakorder"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_type_weakorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction EQ>,
++    // CHECK: compare_type = #vhlo<comparison_type_v1 WEAKORDER>,
++    compare_type = #stablehlo<comparison_type WEAKORDER>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// ConvDimensionNumbers aka #stablehlo.conv is covered below.
++
++// CHECK-LABEL: "attr_custom_call_api_version_unspecified"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "foo",
++    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_UNSPECIFIED>
++    api_version = 0 : i32
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "attr_custom_call_api_version_original"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "foo",
++    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>
++    api_version = 1 : i32
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "attr_custom_call_api_version_status_returning"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "foo",
++    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>
++    api_version = 2 : i32
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "foo",
++    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING_UNIFIED>
++    api_version = 3 : i32
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "attr_dict"
++// CHECK: #vhlo.dict_v1<{#vhlo.string_v1<"attr1"> = #vhlo.integer_v1<1 : i32>, #vhlo.string_v1<"attr2"> = #vhlo.integer_v1<2 : i32>}
++func.func @attr_dict() attributes {stablehlo.attr = {attr1 = 1 : i32, attr2 = 2 : i32}} {
++  return
++}
++
++// CHECK-LABEL: "attr_custom_call_api_version_typed_ffi"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++// CHECK: api_version = #vhlo<api_version_v1 API_VERSION_TYPED_FFI>
++// CHECK-SAME: backend_config = #vhlo.dict_v1<{#vhlo.string_v1<"bar"> = #vhlo.integer_v1<42 : i32>}>
++func.func @attr_custom_call_api_version_typed_ffi(%arg0: tensor<f32>) -> tensor<f32> {
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "foo",
++    backend_config= {bar = 42 : i32},
++    api_version = 4 : i32
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++
++// CHECK-LABEL: "attr_custom_call_api_version_typed_ffi_no_backend_config"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++// CHECK: api_version = #vhlo<api_version_v1 API_VERSION_TYPED_FFI>
++// CHECK-SAME: backend_config = #vhlo.dict_v1<{}>
++func.func @attr_custom_call_api_version_typed_ffi_no_backend_config(%arg0: tensor<f32>) -> tensor<f32> {
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "foo",
++    api_version = 4 : i32
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// DotDimensionNumbers aka #stablehlo.dot is covered below.
++
++// CHECK-LABEL: "attr_fft_type_fft"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
++  %0 = "stablehlo.fft"(%arg0) {
++    // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
++    fft_type = #stablehlo<fft_type FFT>,
++    fft_length = array<i64: 16>
++  } : (tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>>
++  func.return %0 : tensor<16xcomplex<f32>>
++}
++
++// CHECK-LABEL: "attr_fft_type_ifft"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
++  %0 = "stablehlo.fft"(%arg0) {
++    // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
++    fft_type = #stablehlo<fft_type IFFT>,
++    fft_length = array<i64: 16>
++  } : (tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>>
++  func.return %0 : tensor<16xcomplex<f32>>
++}
++
++// CHECK-LABEL: "attr_fft_type_rfft"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
++  %0 = "stablehlo.fft"(%arg0) {
++    // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
++    fft_type = #stablehlo<fft_type RFFT>,
++    fft_length = array<i64: 16>
++  } : (tensor<16xf32>) -> tensor<9xcomplex<f32>>
++  func.return %0 : tensor<9xcomplex<f32>>
++}
++
++// CHECK-LABEL: "attr_fft_type_irfft"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
++  %0 = "stablehlo.fft"(%arg0) {
++    // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
++    fft_type = #stablehlo<fft_type IRFFT>,
++    fft_length = array<i64: 16>
++  } : (tensor<9xcomplex<f32>>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "exponential_HIGHEST"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}
++func.func @exponential_HIGHEST(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
++  %0 = "stablehlo.exponential"(%arg0) {
++    // CHECK: result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 HIGHEST>>
++    result_accuracy = #stablehlo.result_accuracy<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #stablehlo.result_accuracy_mode<HIGHEST>>
++  } : (tensor<8x16xf32>) -> tensor<8x16xf32>
++  func.return %0 : tensor<8x16xf32>
++}
++
++// CHECK-LABEL: "exponential_TOLERANCE"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}
++func.func @exponential_TOLERANCE(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
++  %0 = "stablehlo.exponential"(%arg0) {
++    // CHECK: result_accuracy = #vhlo.result_accuracy_v1<atol = 1.000000e-05, rtol = 0.000000e+00, ulps = 1, mode = #vhlo<result_accuracy_mode_v1 TOLERANCE>>
++    result_accuracy = #stablehlo.result_accuracy<atol = 1.000000e-5, rtol = 0.000000e+00, ulps = 1, mode = #stablehlo.result_accuracy_mode<TOLERANCE>>
++  } : (tensor<8x16xf32>) -> tensor<8x16xf32>
++  func.return %0 : tensor<8x16xf32>
++}
++
++// CHECK-LABEL: "exponential_DEFAULT"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}
++func.func @exponential_DEFAULT(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
++  %0 = "stablehlo.exponential"(%arg0) {
++    // CHECK: result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>
++    result_accuracy = #stablehlo.result_accuracy<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #stablehlo.result_accuracy_mode<DEFAULT>>
++  } : (tensor<8x16xf32>) -> tensor<8x16xf32>
++  func.return %0 : tensor<8x16xf32>
++}
++
++// GatherDimensionNumbers aka #stablehlo.gather is covered below.
++
++// CHECK-LABEL: "attr_precision_config_default"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
++  %0 = "stablehlo.dot"(%arg0, %arg1) {
++    // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
++  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
++  func.return %0 : tensor<8x8xf32>
++}
++
++// CHECK-LABEL: "attr_precision_config_high"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
++  %0 = "stablehlo.dot"(%arg0, %arg1) {
++    // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
++    precision_config = [#stablehlo<precision HIGH>, #stablehlo<precision HIGH>]
++  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
++  func.return %0 : tensor<8x8xf32>
++}
++
++// CHECK-LABEL: "attr_precision_config_highest"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
++  %0 = "stablehlo.dot"(%arg0, %arg1) {
++    // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
++    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
++  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
++  func.return %0 : tensor<8x8xf32>
++}
++
++// CHECK-LABEL: "attr_rng_algorithm_default"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
++  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
++    // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
++    rng_algorithm = #stablehlo<rng_algorithm DEFAULT>
++  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
++  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
++}
++
++// CHECK-LABEL: "attr_rng_algorithm_three_fry"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
++  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
++    // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
++    rng_algorithm = #stablehlo<rng_algorithm THREE_FRY>
++  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
++  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
++}
++
++// CHECK-LABEL: "attr_rng_algorithm_philox"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
++  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
++    // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
++    rng_algorithm = #stablehlo<rng_algorithm PHILOX>
++  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
++  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
++}
++
++// CHECK-LABEL: "attr_rng_distribution_uniform"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
++  %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
++    // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
++    rng_distribution = #stablehlo<rng_distribution UNIFORM>
++  } : (tensor<f32>, tensor<f32>, tensor<0xindex>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "attr_rng_distribution_normal"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
++  %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
++    // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
++    rng_distribution = #stablehlo<rng_distribution NORMAL>
++  } : (tensor<f32>, tensor<f32>, tensor<0xindex>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
++
++// CHECK-LABEL: "attr_transpose_no_transpose"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
++  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
++    left_side = true,
++    lower = true,
++    unit_diagonal = true,
++    // transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
++    transpose_a = #stablehlo<transpose NO_TRANSPOSE>
++  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
++  func.return %0 : tensor<16x16xf32>
++}
++
++// CHECK-LABEL: "attr_transpose_transpose"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
++  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
++    left_side = true,
++    lower = true,
++    unit_diagonal = true,
++    // transpose_a = #vhlo<transpose_v1 TRANSPOSE>,
++    transpose_a = #stablehlo<transpose TRANSPOSE>
++  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
++  func.return %0 : tensor<16x16xf32>
++}
++
++// CHECK-LABEL: "attr_transpose_adjoint"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
++  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
++    left_side = true,
++    lower = true,
++    unit_diagonal = true,
++    // transpose_a = #vhlo<transpose_v1 ADJOINT>,
++    transpose_a = #stablehlo<transpose ADJOINT>
++  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
++  func.return %0 : tensor<16x16xf32>
++}
++
++// TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
++
++// CHECK-LABEL: "attr_type_extensions_bounds"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
++  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
++  func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
++}
++
++// CHECK-LABEL: "attr_frontend_attributes"
++func.func @attr_frontend_attributes(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: some.unregistered_attr
++  %1 = stablehlo.cosine %arg0 {some.unregistered_attr = 1 : i32} : tensor<f32>
++  return %1 : tensor<f32>
++}
++
++// Builtin attriubute tests
++
++// CHECK-LABEL: "byte_packed_boolean"
++func.func @byte_packed_boolean() -> (tensor<8xi1>, tensor<8xi1>, tensor<4xi1>, tensor<4xi1>, tensor<16xi1>) {
++  // CHECK: #vhlo.tensor_v1<dense<[true, false, false, false, false, false, false, false]
++  // CHECK-NEXT: #vhlo.tensor_v1<dense<true> : tensor<8xi1>>
++  // CHECK-NEXT: #vhlo.tensor_v1<dense<[true, false, false, false]> : tensor<4xi1>>
++  // CHECK-NEXT: #vhlo.tensor_v1<dense<true> : tensor<4xi1>>
++  // CHECK-NEXT: #vhlo.tensor_v1<dense<[true, false, false, false, false, false, false, false, true, false, false, false, false, false, false, false]> : tensor<16xi1>>
++  %c = stablehlo.constant dense<[true, false, false, false, false, false, false, false]> : tensor<8xi1>
++  %c_0 = stablehlo.constant dense<true> : tensor<8xi1>
++  %c_1 = stablehlo.constant dense<[true, false, false, false]> : tensor<4xi1>
++  %c_2 = stablehlo.constant dense<true> : tensor<4xi1>
++  %c_3 = stablehlo.constant dense<[true, false, false, false, false, false, false, false, true, false, false, false, false, false, false, false]> : tensor<16xi1>
++  return %c, %c_0, %c_1, %c_2, %c_3 : tensor<8xi1>, tensor<8xi1>, tensor<4xi1>, tensor<4xi1>, tensor<16xi1>
++}
++
++// ============ DEFAULTS ============
++
++// CHECK-LABEL: "default_all_gather"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
++  //               CHECK: "vhlo.all_gather_v2"(%[[ARG0]]) <{
++  //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
++  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
++  %0 = "stablehlo.all_gather"(%arg0) {
++    all_gather_dim = 1 : i64,
++    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
++  } : (tensor<16x8xf32>) -> tensor<16x16xf32>
++  func.return %0 : tensor<16x16xf32>
++}
++
++// CHECK-LABEL: "default_all_gather_variadic"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_all_gather_variadic(%arg0: tensor<16x8xf32>, %arg1: tensor<16x8xf32>) -> (tensor<16x16xf32>, tensor<16x16xf32>) {
++  %0:2 = "stablehlo.all_gather"(%arg0, %arg1) {
++    all_gather_dim = 1 : i64,
++    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
++  } : (tensor<16x8xf32>, tensor<16x8xf32>) -> (tensor<16x16xf32>, tensor<16x16xf32>)
++  func.return %0#0, %0#1 : tensor<16x16xf32>, tensor<16x16xf32>
++}
++
++// CHECK-LABEL: "default_all_reduce"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
++  //               CHECK: "vhlo.all_reduce_v2"(%[[ARG0]])
++  //          CHECK-SAME: <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
++  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
++  //          CHECK-SAME: }> ({
++  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++
++  %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 = dense<[[0], [1]]> : tensor<2x1xi64>
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "attr_replica_group_mesh_axes"
++func.func @attr_replica_group_mesh_axes(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.all_reduce_v2"
++  // CHECK-SAME: replica_groups = #vhlo.replica_group_mesh_axes_v1<
++  // CHECK-SAME:   mesh = #vhlo.string_v1<"mesh">,
++  // CHECK-SAME:   axes = #vhlo.array_v1<[#vhlo.axis_ref_v1<name = #vhlo.string_v1<"x">, sub_axis_info = #vhlo.sub_axis_info_v1<pre_size = 1, size = 2>>, #vhlo.axis_ref_v1<name = #vhlo.string_v1<"y">>]>
++  // CHECK-SAME: >
++  %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 = @mesh,
++      axes = [
++        #stablehlo.axis_ref<name = "x", sub_axis_info = (1)2>,
++        #stablehlo.axis_ref<name = "y">
++      ]
++    >
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++
++
++// CHECK-LABEL: "default_all_to_all"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
++  //               CHECK: "vhlo.all_to_all_v2"(%[[ARG0]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
++  //          CHECK-SAME:   split_count = #vhlo.integer_v1<4 : i64>
++  //          CHECK-SAME:   split_dimension = #vhlo.integer_v1<1 : i64>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x4x!vhlo.f32_v1>
++  %0 = "stablehlo.all_to_all"(%arg0) {
++    split_dimension = 1 : i64,
++    concat_dimension = 0 : i64,
++    split_count = 4 : i64,
++    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>
++  } : (tensor<4x16xf32>) -> tensor<16x4xf32>
++  func.return %0 : tensor<16x4xf32>
++}
++
++// CHECK-LABEL: "default_all_to_all_variadic"
++func.func @default_all_to_all_variadic(%arg0: tensor<4x16xf32>, %arg1: tensor<5x16xf32>) -> (tensor<16x4xf32>, tensor<20x4xf32>) {
++  %0:2 = "stablehlo.all_to_all"(%arg0, %arg1) {
++    split_dimension = 1 : i64,
++    concat_dimension = 0 : i64,
++    split_count = 4 : i64,
++    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
++    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>
++  } : (tensor<4x16xf32>, tensor<5x16xf32>) -> (tensor<16x4xf32>, tensor<20x4xf32>)
++  func.return %0#0, %0#1 : tensor<16x4xf32>, tensor<20x4xf32>
++}
++
++// CHECK-LABEL: "default_cholesky"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
++  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   lower = #vhlo.bool_v1<false>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
++  %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
++  func.return %0 : tensor<1x16x16xf32>
++}
++
++// CHECK-LABEL: "default_collective_permute"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
++  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
++  %0 = "stablehlo.collective_permute"(%arg0) {
++    source_target_pairs = dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>
++  } : (tensor<16x8xf32>) -> tensor<16x8xf32>
++  func.return %0 : tensor<16x8xf32>
++}
++
++// CHECK-LABEL: "default_collective_broadcast"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_collective_broadcast(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
++  //               CHECK: "vhlo.collective_broadcast_v2"(%[[ARG0]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<false>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
++  %0 = "stablehlo.collective_broadcast"(%arg0) {
++    replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
++  } : (tensor<16x8xf32>) -> tensor<16x8xf32>
++  func.return %0 : tensor<16x8xf32>
++}
++
++// CHECK-LABEL: "default_collective_reduce"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_collective_reduce(%arg0: tensor<f32>) -> tensor<f32> {
++  //               CHECK: "vhlo.collective_reduce_v1"(%[[ARG0]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<false>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>,
++  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
++  //          CHECK-SAME: }> ({
++  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.collective_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 = dense<[[0, 1]]> : tensor<1x2xi64>
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "default_compare"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
++  // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction EQ>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "default_composite"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_composite(%arg0: tensor<f32>) -> tensor<f32> {
++  //               CHECK: "vhlo.composite_v2"(%[[ARG0]]) <{
++  //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{}>
++  //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
++  //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
++  //          CHECK-SAME:   version = #vhlo.integer_v1<0 : i64>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.composite"(%arg0) {
++    name = "stablehlo.composite_target",
++    decomposition = @composite_target
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "default_convolution"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
++  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
++  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<2x2xi64>>,
++  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
++  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
++  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<false> : tensor<2xi1>>,
++  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x6x6x16x!vhlo.f32_v1>
++  %0 = "stablehlo.convolution"(%arg0, %arg1) {
++    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
++    feature_group_count = 1 : i64,
++    batch_group_count = 1 : i64
++  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32>
++  func.return %0 : tensor<1x6x6x16xf32>
++}
++
++// CHECK-LABEL: "default_custom_call"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
++  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
++  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
++  // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
++  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
++  // CHECK-SAME:   called_computations = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<false>,
++  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[]>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "foo"
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "default_dot_general"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
++  //      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   allow_imprecise_accumulation = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   lhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++  // CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   num_primitive_operations = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
++  // CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   rhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
++  // CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
++  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
++    dot_dimension_numbers = #stablehlo.dot<
++      lhs_batching_dimensions = [0],
++      lhs_contracting_dimensions = [2],
++      rhs_batching_dimensions = [0],
++      rhs_contracting_dimensions = [1]
++    >
++  } : (tensor<8x8x16xf32>, tensor<8x16x8xf32>) -> tensor<8x8x8xf32>
++  func.return %0 : tensor<8x8x8xf32>
++}
++
++// CHECK-LABEL: "dot_general_algorithm"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @dot_general_algorithm(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
++//      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
++// CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.f32_v1>,
++// CHECK-SAME:   allow_imprecise_accumulation = #vhlo.bool_v1<false>,
++// CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++// CHECK-SAME:   lhs_component_count = #vhlo.integer_v1<1 : i64>,
++// CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++// CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.tf31_v1>,
++// CHECK-SAME:   num_primitive_operations = #vhlo.integer_v1<1 : i64>,
++// CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
++// CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++// CHECK-SAME:   rhs_component_count = #vhlo.integer_v1<1 : i64>,
++// CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
++// CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.tf31_v1>
++// CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
++  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
++    dot_dimension_numbers = #stablehlo.dot<
++      lhs_batching_dimensions = [0],
++      lhs_contracting_dimensions = [2],
++      rhs_batching_dimensions = [0],
++      rhs_contracting_dimensions = [1]
++    >,
++    algorithm = #stablehlo.dot_algorithm<
++      lhs_precision_type = tf32,
++      rhs_precision_type = tf32,
++      accumulation_type = f32,
++      lhs_component_count = 1,
++      rhs_component_count = 1,
++      num_primitive_operations = 1,
++      allow_imprecise_accumulation = false
++    >
++  } : (tensor<8x8x16xf32>, tensor<8x16x8xf32>) -> tensor<8x8x8xf32>
++  func.return %0 : tensor<8x8x8xf32>
++}
++
++// CHECK-LABEL: "dot_general_algorithm_f8e4m3fn_x3"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @dot_general_algorithm_f8e4m3fn_x3(%arg0: tensor<8x8x16xbf16>, %arg1: tensor<8x16x8xbf16>) -> tensor<8x8x8xf32> {
++//      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
++// CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.f32_v1>,
++// CHECK-SAME:   allow_imprecise_accumulation = #vhlo.bool_v1<false>,
++// CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++// CHECK-SAME:   lhs_component_count = #vhlo.integer_v1<1 : i64>,
++// CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++// CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>,
++// CHECK-SAME:   num_primitive_operations = #vhlo.integer_v1<3 : i64>,
++// CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
++// CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++// CHECK-SAME:   rhs_component_count = #vhlo.integer_v1<1 : i64>,
++// CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
++// CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>
++// CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.bf16_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.bf16_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
++  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
++    dot_dimension_numbers = #stablehlo.dot<
++      lhs_batching_dimensions = [0],
++      lhs_contracting_dimensions = [2],
++      rhs_batching_dimensions = [0],
++      rhs_contracting_dimensions = [1]
++    >,
++    algorithm = #stablehlo.dot_algorithm<
++      lhs_precision_type = f8E4M3FN,
++      rhs_precision_type = f8E4M3FN,
++      accumulation_type = f32,
++      lhs_component_count = 1,
++      rhs_component_count = 1,
++      num_primitive_operations = 3,
++      allow_imprecise_accumulation = false
++    >
++  } : (tensor<8x8x16xbf16>, tensor<8x16x8xbf16>) -> tensor<8x8x8xf32>
++  func.return %0 : tensor<8x8x8xf32>
++}
++
++// CHECK-LABEL: "dot_general_algorithm_f8e4m3fn_x4"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @dot_general_algorithm_f8e4m3fn_x4(%arg0: tensor<8x8x16xbf16>, %arg1: tensor<8x16x8xbf16>) -> tensor<8x8x8xf32> {
++//      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
++// CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.f32_v1>,
++// CHECK-SAME:   allow_imprecise_accumulation = #vhlo.bool_v1<false>,
++// CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++// CHECK-SAME:   lhs_component_count = #vhlo.integer_v1<1 : i64>,
++// CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++// CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>,
++// CHECK-SAME:   num_primitive_operations = #vhlo.integer_v1<4 : i64>,
++// CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
++// CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++// CHECK-SAME:   rhs_component_count = #vhlo.integer_v1<1 : i64>,
++// CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
++// CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>
++// CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.bf16_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.bf16_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
++  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
++    dot_dimension_numbers = #stablehlo.dot<
++      lhs_batching_dimensions = [0],
++      lhs_contracting_dimensions = [2],
++      rhs_batching_dimensions = [0],
++      rhs_contracting_dimensions = [1]
++    >,
++    algorithm = #stablehlo.dot_algorithm<
++      lhs_precision_type = f8E4M3FN,
++      rhs_precision_type = f8E4M3FN,
++      accumulation_type = f32,
++      lhs_component_count = 1,
++      rhs_component_count = 1,
++      num_primitive_operations = 4,
++      allow_imprecise_accumulation = false
++    >
++  } : (tensor<8x8x16xbf16>, tensor<8x16x8xbf16>) -> tensor<8x8x8xf32>
++  func.return %0 : tensor<8x8x8xf32>
++}
++
++// CHECK-LABEL: "default_dynamic_broadcast_in_dim"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
++  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
++  // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_broadcast_in_dim"(%arg0, %arg1) {
++    broadcast_dimensions = array<i64: 0, 1>
++  } : (tensor<?x?xf32>, tensor<2xindex>) -> tensor<?x?xf32>
++  func.return %0 : tensor<?x?xf32>
++}
++
++// CHECK-LABEL: "default_dynamic_conv"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<2x2xi64>) -> tensor<1x?x?x16xf32> {
++  //      CHECK: "vhlo.dynamic_conv_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
++  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
++  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
++  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<false> : tensor<2xi1>>,
++  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x2x!vhlo.i64_v1>) -> !vhlo.tensor_v1<1x?x?x16x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_conv"(%arg0, %arg1, %arg2) {
++    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
++    feature_group_count = 1 : i64,
++    batch_group_count = 1 : i64
++  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>, tensor<2x2xi64>) -> tensor<1x?x?x16xf32>
++  func.return %0 : tensor<1x?x?x16xf32>
++}
++
++// CHECK-LABEL: "default_dynamic_gather"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
++  //      CHECK: "vhlo.dynamic_gather_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
++  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
++  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>, !vhlo.tensor_v1<3x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x8x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_gather"(%arg0, %arg1, %arg2) {
++    dimension_numbers = #stablehlo.gather<
++      offset_dims = [2],
++      collapsed_slice_dims = [0, 1],
++      start_index_map = [0, 1],
++      index_vector_dim = 2
++    >
++  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<1x5x8xf32>
++  func.return %0 : tensor<1x5x8xf32>
++}
++
++func.func @default_func(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK:      "vhlo.func_v1"() <{
++  // CHECK-SAME:   arg_attrs = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>>>,
++  // CHECK-SAME:   res_attrs = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
++  // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : () -> ()
++  func.return %arg0 : tensor<f32>
++}
++
++// CHECK-LABEL: "default_gather"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @default_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
++  //      CHECK: "vhlo.gather_v2"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
++  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
++  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<1> : tensor<3xi64>>,
++  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x1x!vhlo.f32_v1>
++  %0 = "stablehlo.gather"(%arg0, %arg1) {
++    dimension_numbers = #stablehlo.gather<
++      offset_dims = [2],
++      collapsed_slice_dims = [0, 1],
++      start_index_map = [0, 1],
++      index_vector_dim = 2
++    >,
++    slice_sizes = array<i64: 1, 1, 1>
++  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>) -> tensor<1x5x1xf32>
++  func.return %0 : tensor<1x5x1xf32>
++}
++
++// CHECK-LABEL: "default_infeed"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
++  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
++  //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
++  // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
++  //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
++  %0:2 = "stablehlo.infeed"(%arg0) : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
++  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
++}
++
++// CHECK-LABEL: "default_outfeed"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
++  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
++  %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
++  func.return %0 : !stablehlo.token
++}
++
++// CHECK-LABEL: "op_recv_with_source_target_pairs"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_recv_with_source_target_pairs(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
++  //      CHECK: "vhlo.recv_v2"(%[[ARG0]]) <{
++  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>,
++  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>>
++  // CHECK-SAME{LITERAL}: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
++  %0:2 = "stablehlo.recv"(%arg0) {
++    channel_handle = #stablehlo.channel_handle<handle = 0, type = 1>,
++    source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>
++  } : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
++  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
++}
++
++// CHECK-LABEL: "default_send"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
++  //      CHECK: "vhlo.send_v2"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>,
++  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>>
++  // CHECK-SAME{LITERAL}: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
++  %0 = "stablehlo.send"(%arg0, %arg1) {
++    channel_handle = #stablehlo.channel_handle<handle = 0, type = 1>,
++    source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>
++  } : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
++  func.return %0 : !stablehlo.token
++}
++
++// CHECK-LABEL: "default_reduce_scatter"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
++  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
++  //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
++  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
++  //          CHECK-SAME: }> ({
++  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.reduce_scatter"(%arg0) ({
++    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
++      %1 = "stablehlo.add"(%arg1, %arg2) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    scatter_dimension = 0 : i64,
++    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
++  } : (tensor<16xf32>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "default_reduce_window"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
++  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
++  //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
++  // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
++  //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
++  //          CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
++  //          CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
++  //          CHECK-SAME: }> ({
++  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.maximum_v1"(%[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<2x17x31x7x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<2x16x30x7x!vhlo.f32_v1>
++  %0 = "stablehlo.reduce_window"(%arg0, %arg1) ({
++    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
++      %1 = "stablehlo.maximum"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    window_dimensions = array<i64: 1, 2, 2, 1>
++  } : (tensor<2x17x31x7xf32>, tensor<f32>) -> tensor<2x16x30x7xf32>
++  func.return %0 : tensor<2x16x30x7xf32>
++}
++
++// CHECK-LABEL: "default_scatter"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
++  //      CHECK: "vhlo.scatter_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
++  // CHECK-SAME:   input_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
++  // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   scatter_dims_to_operand_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   scatter_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
++  // CHECK-SAME:   unique_indices = #vhlo.bool_v1<false>,
++  // CHECK-SAME:   update_window_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>
++  %0 = "stablehlo.scatter"(%arg0, %arg1, %arg2) ({
++    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
++      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    scatter_dimension_numbers = #stablehlo.scatter<
++      update_window_dims = [1],
++      inserted_window_dims = [0, 1],
++      scatter_dims_to_operand_dims = [0, 1],
++      index_vector_dim = 1
++    >
++  } : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) -> tensor<200x100x300xf32>
++  func.return %0 : tensor<200x100x300xf32>
++}
++
++// CHECK-LABEL: "default_select_and_scatter"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
++  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
++  // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
++  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG31:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG41:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL11:.*]] = "vhlo.compare_v1"(%[[ARG31]], %[[ARG41]]) <{compare_type = #vhlo<comparison_type_v1 TOTALORDER>, comparison_direction = #vhlo<comparison_direction_v1 GE>}>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL11]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
++  // CHECK-NEXT: }, {
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG32:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG42:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL12:.*]] = "vhlo.add_v1"(%[[ARG32]], %[[ARG42]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL12]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>, !vhlo.tensor_v1<10x23x23x64x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>
++  %0 = "stablehlo.select_and_scatter"(%arg0, %arg1, %arg2) ({
++    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
++      %1 = "stablehlo.compare"(%arg3, %arg4) {compare_type = #stablehlo<comparison_type TOTALORDER>, comparison_direction = #stablehlo<comparison_direction GE>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
++      "stablehlo.return"(%1) : (tensor<i1>) -> ()
++  }, {
++    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
++      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    window_dimensions = array<i64: 1, 2, 2, 1>
++  } : (tensor<10x24x24x64xf32>, tensor<10x23x23x64xf32>, tensor<f32>) -> tensor<10x24x24x64xf32>
++  func.return %0 : tensor<10x24x24x64xf32>
++}
++
++// CHECK-LABEL: "default_sort"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
++  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
++  // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.compare_v1"(%[[ARG1]], %[[ARG2]]) <{compare_type = #vhlo<comparison_type_v1 FLOAT>, comparison_direction = #vhlo<comparison_direction_v1 GT>}>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.sort"(%arg0) ({
++    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
++      %1 = "stablehlo.compare"(%arg1, %arg2) {compare_type = #stablehlo<comparison_type FLOAT>, comparison_direction = #stablehlo<comparison_direction GT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
++      "stablehlo.return"(%1) : (tensor<i1>) -> ()
++  }) : (tensor<16xf32>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// ============ OPS ============
++
++// CHECK-LABEL: "op_abs"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_add"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_after_all"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
++  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
++  %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
++  func.return %0 : !stablehlo.token
++}
++
++// CHECK-LABEL: "op_all_gather"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
++  //               CHECK: "vhlo.all_gather_v2"(%[[ARG0]]) <{
++  //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
++  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
++  %0 = "stablehlo.all_gather"(%arg0) {
++    all_gather_dim = 1 : i64,
++    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>,
++    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
++    use_global_device_ids
++  } : (tensor<16x8xf32>) -> tensor<16x16xf32>
++  func.return %0 : tensor<16x16xf32>
++}
++
++// CHECK-LABEL: "op_all_reduce"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
++  //               CHECK: "vhlo.all_reduce_v2"(%[[ARG0]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
++  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
++  //          CHECK-SAME: }> ({
++  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %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 = dense<[[0], [1]]> : tensor<2x1xi64>,
++    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
++    use_global_device_ids
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_all_reduce_with_promotable_types"
++func.func @op_all_reduce_with_promotable_types(%operand: tensor<f32>) -> tensor<f64> {
++  //  CHECK: "vhlo.all_reduce_v2"(%[[ARG0:.*]])
++  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
++  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
++  //  CHECK: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
++  %result = "stablehlo.all_reduce"(%operand) ({
++    ^bb0(%arg0: tensor<f64>, %arg1: tensor<f64>):
++      %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
++      "stablehlo.return"(%0) : (tensor<f64>) -> ()
++  }) {
++    replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
++    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
++    use_global_device_ids
++  } : (tensor<f32>) -> tensor<f64>
++
++  func.return %result : tensor<f64>
++}
++
++// CHECK-LABEL: "default_all_reduce_variadic"
++func.func @default_all_reduce_variadic(%arg0: tensor<f32>, %arg1: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
++  %0:2 = "stablehlo.all_reduce"(%arg0, %arg1) ({
++    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
++      %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> (tensor<f32>)
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
++  } : (tensor<f32>, tensor<f32>) -> (tensor<f32>, tensor<f32>)
++  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
++}
++
++// CHECK-LABEL: "op_all_to_all"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
++  //               CHECK: "vhlo.all_to_all_v2"(%[[ARG0]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
++  //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
++  //          CHECK-SAME:   split_count = #vhlo.integer_v1<4 : i64>
++  //          CHECK-SAME:   split_dimension = #vhlo.integer_v1<1 : i64>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x4x!vhlo.f32_v1>
++  %0 = "stablehlo.all_to_all"(%arg0) {
++    split_dimension = 1 : i64,
++    concat_dimension = 0 : i64,
++    split_count = 4 : i64,
++    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
++    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>
++  } : (tensor<4x16xf32>) -> tensor<16x4xf32>
++  func.return %0 : tensor<16x4xf32>
++}
++
++// CHECK-LABEL: "op_and"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
++  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "op_atan2"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_batch_norm_grad"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
++func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
++  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
++  // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
++  // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
++  %0:3 = "stablehlo.batch_norm_grad"(%arg0, %arg1, %arg2, %arg3, %arg4) {
++    epsilon = 0.001 : f32,
++    feature_index = 0 : i64
++  } : (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>)
++  func.return %0#0, %0#1, %0#2 : tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_batch_norm_inference"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
++func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
++  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
++  // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
++  // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
++  %0 = "stablehlo.batch_norm_inference"(%arg0, %arg1, %arg2, %arg3, %arg4) {
++    epsilon = 0.001 : f32,
++    feature_index = 0 : i64
++  } : (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>) -> tensor<16x16x16x16xf32>
++  func.return %0 : tensor<16x16x16x16xf32>
++}
++
++// CHECK-LABEL: "op_batch_norm_training"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
++  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
++  // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
++  %0:3 = "stablehlo.batch_norm_training"(%arg0, %arg1, %arg2) {
++    epsilon = 0.001 : f32,
++    feature_index = 0 : i64
++  } : (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>)
++  func.return %0#0, %0#1, %0#2 : tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_bitcast_convert"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
++  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_broadcast_in_dim"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
++  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
++  %0 = "stablehlo.broadcast_in_dim"(%arg0) {
++    broadcast_dimensions = array<i64: 1>
++  } : (tensor<16xf32>) -> tensor<16x16xf32>
++  func.return %0 : tensor<16x16xf32>
++}
++
++// CHECK-LABEL: "op_broadcast"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
++  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
++  %0 = "stablehlo.broadcast"(%arg0) {
++    broadcast_sizes = array<i64: 16>
++  } : (tensor<16xf32>) -> tensor<16x16xf32>
++  func.return %0 : tensor<16x16xf32>
++}
++
++// CHECK-LABEL: "op_case"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
++  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.case"(%arg0) ({
++    "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
++  }) : (tensor<i32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_cbrt"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.cbrt_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_ceil"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_cholesky"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
++  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   lower = #vhlo.bool_v1<true>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
++  %0 = "stablehlo.cholesky"(%arg0) {
++    lower = true
++  } : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
++  func.return %0 : tensor<1x16x16xf32>
++}
++
++// CHECK-LABEL: "op_clamp"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_count_leading_zeros"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
++  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
++  %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
++  func.return %0 : tensor<i32>
++}
++
++// CHECK-LABEL: "op_collective_broadcast"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_collective_broadcast(%arg0: tensor<16x8xf32>, %arg1: tensor<1xi32>) -> tensor<16x8xf32> {
++  //               CHECK: "vhlo.collective_broadcast_v2"(%[[ARG0]], %[[ARG1]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
++  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<true>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
++  %0 = "stablehlo.collective_broadcast"(%arg0, %arg1) {
++    replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
++    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
++    has_dynamic_root
++  } : (tensor<16x8xf32>, tensor<1xi32>) -> tensor<16x8xf32>
++  func.return %0 : tensor<16x8xf32>
++}
++
++// CHECK-LABEL: "op_collective_reduce"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_collective_reduce(%arg0: tensor<f32>, %arg1: tensor<1xi32>) -> tensor<f32> {
++  //               CHECK: "vhlo.collective_reduce_v1"(%[[ARG0]], %[[ARG1]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
++  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<true>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
++  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
++  //          CHECK-SAME: }> ({
++  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.collective_reduce"(%arg0, %arg1) ({
++    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
++      %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>,
++    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
++    use_global_device_ids,
++    has_dynamic_root
++  } : (tensor<f32>, tensor<1xi32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_collective_permute"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
++  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
++  %0 = "stablehlo.collective_permute"(%arg0) {
++    source_target_pairs = dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>,
++    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>
++  } : (tensor<16x8xf32>) -> tensor<16x8xf32>
++  func.return %0 : tensor<16x8xf32>
++}
++
++// CHECK-LABEL: "op_compare"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
++  // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction EQ>,
++    compare_type = #stablehlo<comparison_type TOTALORDER>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "op_complex"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
++  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
++  %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
++  func.return %0 : tensor<complex<f32>>
++}
++
++// CHECK-LABEL: "op_composite"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_composite(%arg0: tensor<f32>) -> tensor<f32> {
++  //               CHECK: "vhlo.composite_v2"(%[[ARG0]]) <{
++  //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{#vhlo.string_v1<"my_int"> = #vhlo.integer_v1<1 : i64>, #vhlo.string_v1<"my_string"> = #vhlo.string_v1<"foo">}>
++  //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
++  //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
++  //          CHECK-SAME:   version = #vhlo.integer_v1<1 : i32>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.composite"(%arg0) {
++    name = "stablehlo.composite_target",
++    decomposition = @composite_target,
++    version = 1 : i32,
++    composite_attributes = {
++      my_int = 1 : i64,
++      my_string = "foo"
++    }
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "composite_regions"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @composite_regions(%arg0: tensor<f32>) -> tensor<f32> {
++  //               CHECK: "vhlo.composite_v2"(%[[ARG0]]) <{
++  //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{}>
++  //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
++  //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
++  //          CHECK-SAME:   version = #vhlo.integer_v1<1 : i32>
++  //          CHECK-SAME: }> ({
++  //          CHECK-NEXT: ^bb0(%[[ARG1:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //          CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.composite"(%arg0) ({
++    ^bb0(%arg1: tensor<f32>):
++      "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
++  }) {
++    name = "stablehlo.composite_target",
++    decomposition = @composite_target,
++    version = 1 : i32
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_concatenate"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
++  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.concatenate"(%arg0, %arg1) {
++    dimension = 0 : i64
++  } : (tensor<8xf32>, tensor<8xf32>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_constant"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
++  //      CHECK: "vhlo.constant_v1"() <{
++  // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
++  // CHECK-SAME: }> : () -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.constant"() {
++    value = dense<0.0> : tensor<f32>
++  } : () -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_convert"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
++  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_convolution"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
++  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
++  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<2x2xi64>>,
++  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
++  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
++  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<true> : tensor<2xi1>>,
++  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x7x7x16x!vhlo.f32_v1>
++  %0 = "stablehlo.convolution"(%arg0, %arg1) {
++    window_strides = array<i64: 2, 2>,
++    padding = dense<1> : tensor<2x2xi64>,
++    lhs_dilation = array<i64: 2, 2>,
++    rhs_dilation = array<i64: 2, 2>,
++    window_reversal = array<i1: true, true>,
++    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
++    feature_group_count = 1 : i64,
++    batch_group_count = 1 : i64,
++    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
++  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32>
++  func.return %0 : tensor<1x7x7x16xf32>
++}
++
++// CHECK-LABEL: "op_cosine"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.cosine_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_create_token"
++func.func @op_create_token() -> !stablehlo.token {
++  // CHECK: "vhlo.create_token_v1"() : () -> !vhlo.token_v1
++  %0 = "stablehlo.create_token"() : () -> !stablehlo.token
++  func.return %0 : !stablehlo.token
++}
++
++// CHECK-LABEL: "op_cross_replica_sum"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
++  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.cross-replica-sum"(%arg0) {
++    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_custom_call"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
++  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
++  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
++  // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
++  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
++  // CHECK-SAME:   called_computations = #vhlo.array_v1<[#vhlo.string_v1<"foo">]>,
++  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<> : tensor<0xindex>>]>,
++  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[
++  // CHECK-SAME:     #vhlo.output_operand_alias_v1<
++  // CHECK-SAME:       outputTupleIndices = [],
++  // CHECK-SAME:       operandIndex = 0,
++  // CHECK-SAME:       operandTupleIndices = []>]>,
++  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<> : tensor<0xindex>>]>,
++  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[]>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "foo",
++    has_side_effect = true,
++    backend_config = "\08\03\1A\02",
++    api_version = 2 : i32,
++    called_computations = [@foo],
++    operand_layouts = [dense<> : tensor<0xindex>],
++    output_operand_aliases = [
++      #stablehlo.output_operand_alias<output_tuple_indices = [],
++                                 operand_index = 0,
++                                 operand_tuple_indices = []>],
++    result_layouts = [dense<> : tensor<0xindex>]
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_custom_call_empty_result_layout"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func public @op_custom_call_empty_result_layout(%arg0: tensor<i64>) -> tensor<i64> {
++  // %0 = "vhlo.custom_call_v2"(%arg0) <{>}> : (!vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tuple_v1<>
++  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
++  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
++  // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
++  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"empty_output">,
++  // CHECK-SAME:   called_computations = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<> : tensor<0xindex>>]>,
++  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[]>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tuple_v1<>
++  %0 = "stablehlo.custom_call"(%arg0) <{
++    api_version = 2 : i32,
++    call_target_name = "empty_output",
++    has_side_effect = true,
++    operand_layouts = [dense<> : tensor<0xindex>],
++    result_layouts = []
++  }> : (tensor<i64>) -> tuple<>
++  return %arg0 : tensor<i64>
++}
++
++// CHECK-LABEL: "op_custom_call_with_result_tilings"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_custom_call_with_result_tilings(%arg0: tensor<64x256xf32>) -> tensor<64x256xf32> {
++  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
++  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
++  // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
++  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
++  // CHECK-SAME:   called_computations = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<false>,
++  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<[1, 0]> : tensor<2xindex>>]>,
++  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<[1, 0]> : tensor<2xindex>>]>,
++  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[#vhlo.array_v1<[#vhlo.tensor_v1<dense<[8, 128]> : tensor<2xindex>>]>]>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<64x256x!vhlo.f32_v1>) -> !vhlo.tensor_v1<64x256x!vhlo.f32_v1>
++  %0 = "stablehlo.custom_call"(%arg0) {
++    api_version = 2 : i32,
++    call_target_name = "foo",
++    operand_layouts = [dense<[1, 0]> : tensor<2xindex>],
++    result_layouts = [dense<[1, 0]> : tensor<2xindex>],
++    result_tilings = [[dense<[8, 128]> : tensor<2xindex>]]
++  } : (tensor<64x256xf32>) -> tensor<64x256xf32>
++  func.return %0 : tensor<64x256xf32>
++}
++
++// CHECK-LABEL: "op_divide"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_dot_general"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
++  //      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   allow_imprecise_accumulation = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   lhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++  // CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   num_primitive_operations = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
++  // CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   rhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
++  // CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
++  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
++    dot_dimension_numbers = #stablehlo.dot<
++      lhs_batching_dimensions = [0],
++      lhs_contracting_dimensions = [2],
++      rhs_batching_dimensions = [0],
++      rhs_contracting_dimensions = [1]
++    >,
++    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
++  } : (tensor<8x8x16xf32>, tensor<8x16x8xf32>) -> tensor<8x8x8xf32>
++  func.return %0 : tensor<8x8x8xf32>
++}
++
++// CHECK-LABEL: "op_dot"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
++  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
++  %0 = "stablehlo.dot"(%arg0, %arg1) {
++    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
++  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
++  func.return %0 : tensor<8x8xf32>
++}
++
++// CHECK-LABEL: "op_dynamic_broadcast_in_dim"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
++  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_broadcast_in_dim"(%arg0, %arg1) {
++    broadcast_dimensions = array<i64: 0, 1>,
++    known_expanding_dimensions = array<i64: 0>,
++    known_nonexpanding_dimensions = array<i64: 1>
++  } : (tensor<?x?xf32>, tensor<2xindex>) -> tensor<?x?xf32>
++  func.return %0 : tensor<?x?xf32>
++}
++
++// CHECK-LABEL: "op_dynamic_conv"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<2x2xi64>) -> tensor<1x?x?x16xf32> {
++  //      CHECK: "vhlo.dynamic_conv_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
++  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
++  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
++  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<true> : tensor<2xi1>>,
++  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x2x!vhlo.i64_v1>) -> !vhlo.tensor_v1<1x?x?x16x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_conv"(%arg0, %arg1, %arg2) {
++    window_strides = array<i64: 2, 2>,
++    lhs_dilation = array<i64: 2, 2>,
++    rhs_dilation = array<i64: 2, 2>,
++    window_reversal = array<i1: true, true>,
++    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
++    feature_group_count = 1 : i64,
++    batch_group_count = 1 : i64,
++    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
++  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>, tensor<2x2xi64>) -> tensor<1x?x?x16xf32>
++  func.return %0 : tensor<1x?x?x16xf32>
++}
++
++// CHECK-LABEL: "op_dynamic_gather"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
++  //      CHECK: "vhlo.dynamic_gather_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
++  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>, !vhlo.tensor_v1<3x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x8x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_gather"(%arg0, %arg1, %arg2) {
++    dimension_numbers = #stablehlo.gather<
++      offset_dims = [2],
++      collapsed_slice_dims = [0, 1],
++      start_index_map = [0, 1],
++      index_vector_dim = 2
++    >,
++    indices_are_sorted = true
++  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<1x5x8xf32>
++  func.return %0 : tensor<1x5x8xf32>
++}
++
++// CHECK-LABEL: "op_dynamic_gather_with_batching_dims"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_dynamic_gather_with_batching_dims(%arg0 : tensor<5x2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<4xi32>) -> tensor<1x5x8xf32> {
++  //      CHECK: "vhlo.dynamic_gather_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<5x2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>, !vhlo.tensor_v1<4x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x8x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_gather"(%arg0, %arg1, %arg2) {
++    dimension_numbers = #stablehlo.gather<
++      offset_dims = [2],
++      collapsed_slice_dims = [1, 2],
++      operand_batching_dims = [0],
++      start_indices_batching_dims = [1],
++      start_index_map = [1, 2],
++      index_vector_dim = 2
++    >,
++    indices_are_sorted = true
++  } : (tensor<5x2x4x9xf32>, tensor<1x5x2xi32>, tensor<4xi32>) -> tensor<1x5x8xf32>
++  func.return %0 : tensor<1x5x8xf32>
++}
++
++// CHECK-LABEL: "op_dynamic_iota"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
++  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_iota"(%arg0) {
++    iota_dimension = 0 : i64
++  } : (tensor<1xindex>) -> tensor<?xf32>
++  func.return %0 : tensor<?xf32>
++}
++
++// CHECK-LABEL: "op_dynamic_pad"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
++func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
++  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
++  func.return %0 : tensor<?xf32>
++}
++
++// CHECK-LABEL: "op_dynamic_reshape"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
++  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
++  func.return %0 : tensor<?x?xf32>
++}
++
++// CHECK-LABEL: "op_dynamic_slice"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
++  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
++    slice_sizes = array<i64: 4>
++  } : (tensor<16xf32>, tensor<i64>) -> tensor<4xf32>
++  func.return %0 : tensor<4xf32>
++}
++
++// CHECK-LABEL: "op_dynamic_update_slice"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
++  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_einsum"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
++  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
++  %0 = "stablehlo.einsum"(%arg0, %arg1) {
++    einsum_config = "ab,bc->ac"
++  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
++  func.return %0 : tensor<8x8xf32>
++}
++
++// CHECK-LABEL: "op_exponential_minus_one"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.exponential_minus_one_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_exponential"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.exponential_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_fft"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
++  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
++  // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
++  %0 = "stablehlo.fft"(%arg0) {
++    fft_type = #stablehlo<fft_type FFT>,
++    fft_length = array<i64: 16>
++  } : (tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>>
++  func.return %0 : tensor<16xcomplex<f32>>
++}
++
++// CHECK-LABEL: "op_floor"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++func.func private @op_func(%arg0: tensor<f32> {stablehlo.arg = "0"}) -> (tensor<f32> {stablehlo.result = "0"}) {
++  // CHECK:      "vhlo.func_v1"() <{
++  // CHECK-SAME:   arg_attrs = #vhlo.array_v1<[#vhlo.dict_v1<{#vhlo.string_v1<"stablehlo.arg"> = #vhlo.string_v1<"0">}>]>,
++  // CHECK-SAME:   function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>>>,
++  // CHECK-SAME:   res_attrs = #vhlo.array_v1<[#vhlo.dict_v1<{#vhlo.string_v1<"stablehlo.result"> = #vhlo.string_v1<"0">}>]>,
++  // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
++  // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : () -> ()
++
++  func.return %arg0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_gather"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
++  //      CHECK: "vhlo.gather_v2"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
++  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<1> : tensor<3xi64>>,
++  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x1x!vhlo.f32_v1>
++  %0 = "stablehlo.gather"(%arg0, %arg1) {
++    dimension_numbers = #stablehlo.gather<
++      offset_dims = [2],
++      collapsed_slice_dims = [0, 1],
++      start_index_map = [0, 1],
++      index_vector_dim = 2
++    >,
++    slice_sizes = array<i64: 1, 1, 1>,
++    indices_are_sorted = true
++  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>) -> tensor<1x5x1xf32>
++  func.return %0 : tensor<1x5x1xf32>
++}
++
++// CHECK-LABEL: "op_gather_with_batching_dims"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_gather_with_batching_dims(%arg0 : tensor<5x2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
++  //      CHECK: "vhlo.gather_v2"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
++  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<5x2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x1x!vhlo.f32_v1>
++  %0 = "stablehlo.gather"(%arg0, %arg1) {
++    dimension_numbers = #stablehlo.gather<
++      offset_dims = [2],
++      collapsed_slice_dims = [1, 2],
++      operand_batching_dims = [0],
++      start_indices_batching_dims = [1],
++      start_index_map = [1, 2],
++      index_vector_dim = 2
++    >,
++    slice_sizes = array<i64: 1, 1, 1, 1>,
++    indices_are_sorted = true
++  } : (tensor<5x2x4x9xf32>, tensor<1x5x2xi32>) -> tensor<1x5x1xf32>
++  func.return %0 : tensor<1x5x1xf32>
++}
++
++// CHECK-LABEL: "op_get_dimension_size"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
++  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
++  %0 = "stablehlo.get_dimension_size"(%arg0) {
++    dimension = 0 : i64
++  } : (tensor<?xf32>) -> tensor<i32>
++  func.return %0 : tensor<i32>
++}
++
++// CHECK-LABEL: "op_get_tuple_element"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
++  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
++  // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.get_tuple_element"(%arg0) {
++    index = 0 : i32
++  } : (tuple<tensor<f32>, tensor<i32>>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_if"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
++  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }, {
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.if"(%arg0) ({
++    "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
++  }, {
++    "stablehlo.return"(%arg2) : (tensor<f32>) -> ()
++  }) : (tensor<i1>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_imag"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
++  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_infeed"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
++  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
++  //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
++  // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
++  //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
++  %0:2 = "stablehlo.infeed"(%arg0) {
++    infeed_config = "foo",
++    layout = [[]]
++  } : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
++  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
++}
++
++// CHECK-LABEL: "op_iota"
++func.func @op_iota() -> tensor<16xf32> {
++  //      CHECK: "vhlo.iota_v1"() <{
++  // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME: }> : () -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.iota"() {
++    iota_dimension = 0 : i64
++  } : () -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_is_finite"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
++  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "op_log"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.log_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_log_plus_one"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.log_plus_one_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_logistic"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.logistic_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_map"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
++  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.abs_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.map"(%arg0) ({
++    ^bb0(%arg1: tensor<f32>):
++      %1 = "stablehlo.abs"(%arg1) : (tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    dimensions = array<i64: 0>
++  } : (tensor<16xf32>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_maximum"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_minimum"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_multiply"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_negate"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_not"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
++  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "op_optimization_barrier"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_or"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
++  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "op_outfeed"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
++  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
++  %0 = "stablehlo.outfeed"(%arg0, %arg1) {
++    outfeed_config = "foo"
++  } : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
++  func.return %0 : !stablehlo.token
++}
++
++// CHECK-LABEL: "op_pad"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
++  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
++  // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
++  // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.pad"(%arg0, %arg1) {
++    edge_padding_high = array<i64: 4>,
++    edge_padding_low = array<i64: 4>,
++    interior_padding = array<i64: 0>
++  } : (tensor<8xf32>, tensor<f32>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_popcnt"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
++  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
++  %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
++  func.return %0 : tensor<i32>
++}
++
++// CHECK-LABEL: "op_power"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_real_dynamic_slice"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
++func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
++  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
++  %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
++  func.return %0 : tensor<?xf32>
++}
++
++// CHECK-LABEL: "op_real"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
++  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_recv_no_source_target_pairs"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_recv_no_source_target_pairs(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
++  //      CHECK: "vhlo.recv_v2"(%[[ARG0]]) <{
++  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>,
++  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<> : tensor<0xi64>
++  // CHECK-SAME{LITERAL}: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
++  %0:2 = "stablehlo.recv"(%arg0) {
++    channel_handle = #stablehlo.channel_handle<handle = 0, type = 3>,
++    is_host_transfer = true
++  } : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
++  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
++}
++
++// CHECK-LABEL: "op_reduce"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
++  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
++  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.reduce"(%arg0, %arg1) ({
++    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
++      %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    dimensions = array<i64: 0>
++  } : (tensor<16xf32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_reduce_precision"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
++  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
++  // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.reduce_precision"(%arg0) {
++    exponent_bits = 8 : i32,
++    mantissa_bits = 10 : i32
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK_lABEL: "op_reduce_with_promotable_types"
++func.func @op_reduce_with_promotable_types(%arg0: tensor<4x4xf32>, %arg1 : tensor<f32>)
++    -> (tensor<4xf64>) {
++  //  CHECK: "vhlo.reduce_v1"(%[[ARG0:.*]], %[[ARG1:.*]])
++  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
++  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
++  //  CHECK: }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x!vhlo.f64_v1>
++  %0 = "stablehlo.reduce"(%arg0, %arg1) ({
++  ^bb0(%arg2: tensor<f64>, %arg3: tensor<f64> ):
++    %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f64>, tensor<f64>) -> tensor<f64>
++    "stablehlo.return"(%1) : (tensor<f64>) -> ()
++
++  }) {dimensions = array<i64: 0>} : (tensor<4x4xf32>, tensor<f32>) -> tensor<4xf64>
++
++  func.return %0: tensor<4xf64>
++}
++
++// CHECK-LABEL: "op_reduce_scatter"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
++  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
++  //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
++  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
++  //          CHECK-SAME: }> ({
++  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.reduce_scatter"(%arg0) ({
++    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
++      %1 = "stablehlo.add"(%arg1, %arg2) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    scatter_dimension = 0 : i64,
++    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>,
++    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
++    use_global_device_ids
++  } : (tensor<16xf32>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK_lABEL: "op_reduce_scatter_with_promotable_types"
++func.func @op_reduce_scatter_with_promotable_types(%data: tensor<4x16xf32>) -> tensor<4x4xf64> {
++  //  CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0:.*]])
++  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
++  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
++  //  CHECK: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f64_v1>
++  %0 = "stablehlo.reduce_scatter"(%data) ({
++    ^bb0(%arg2: tensor<f64>, %arg3: tensor<f64>):
++    %1 = stablehlo.add %arg2, %arg3 : tensor<f64>
++    "stablehlo.return"(%1) : (tensor<f64>) -> ()
++  }) {replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
++      scatter_dimension = 1 : i64,
++      channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
++      use_global_device_ids} : (tensor<4x16xf32>) -> tensor<4x4xf64>
++  func.return %0 : tensor<4x4xf64>
++}
++
++
++// CHECK-LABEL: "op_reduce_window"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
++  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
++  //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
++  // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
++  //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
++  //          CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
++  //          CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 4, 4, 1]> : tensor<4xi64>>
++  //          CHECK-SAME: }> ({
++  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.maximum_v1"(%[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<2x17x31x7x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<2x9x16x7x!vhlo.f32_v1>
++  %0 = "stablehlo.reduce_window"(%arg0, %arg1) ({
++    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
++      %1 = "stablehlo.maximum"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    window_dimensions = array<i64: 1, 2, 2, 1>,
++    window_strides = array<i64: 1, 4, 4, 1>,
++    base_dilations = array<i64: 1, 2, 2, 1>,
++    window_dilations = array<i64: 1, 2, 2, 1>,
++    padding = dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>
++  } : (tensor<2x17x31x7xf32>, tensor<f32>) -> tensor<2x9x16x7xf32>
++  func.return %0 : tensor<2x9x16x7xf32>
++}
++
++// CHECK-LABEL: "op_reduce_window_with_promotable_types"
++func.func @op_reduce_window_with_promotable_types(%arg0: tensor<4x2xf32>,
++    %arg1: tensor<4x2xf32>, %init0: tensor<f32>, %init1: tensor<f32>) ->
++    (tensor<2x2xf64>, tensor<2x2xf32>) {
++  //  CHECK: "vhlo.reduce_window_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]], %[[ARG3:.*]])
++  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]], %[[VAL2:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //  CHECK: }) : (!vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<2x2x!vhlo.f64_v1>, !vhlo.tensor_v1<2x2x!vhlo.f32_v1>)
++  %0:2 = "stablehlo.reduce_window"(%arg0, %arg1, %init0, %init1) ({
++         ^bb0(%a0: tensor<f64>, %a1: tensor<f32>, %b0: tensor<f64>,
++                %b1: tensor<f32>):
++              %2 = stablehlo.add %a0, %b0 : tensor<f64>
++              %3 = stablehlo.add %a1, %b1 : tensor<f32>
++              "stablehlo.return"(%2,%3) : (tensor<f64>, tensor<f32>) -> ()
++            })
++         { padding = dense<[[2, 2], [0, 0]]> : tensor<2x2xi64>,
++         window_dimensions = array<i64: 5, 1>,
++         window_strides = array<i64: 3, 1> }
++         : (tensor<4x2xf32>, tensor<4x2xf32>, tensor<f32>, tensor<f32>) ->
++              (tensor<2x2xf64>, tensor<2x2xf32>)
++  func.return %0#0, %0#1 : tensor<2x2xf64>, tensor<2x2xf32>
++}
++
++// CHECK-LABEL: "op_remainder"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_replica_id"
++func.func @op_replica_id() -> tensor<ui32> {
++  // CHECK: "vhlo.replica_id_v1"() : () -> !vhlo.tensor_v1<!vhlo.ui32_v1>
++  %0 = "stablehlo.replica_id"() : () -> tensor<ui32>
++  func.return %0 : tensor<ui32>
++}
++
++// CHECK-LABEL: "op_partition_id"
++func.func @op_partition_id() -> tensor<ui32> {
++  // CHECK: "vhlo.partition_id_v1"() : () -> !vhlo.tensor_v1<!vhlo.ui32_v1>
++  %0 = "stablehlo.partition_id"() : () -> tensor<ui32>
++  func.return %0 : tensor<ui32>
++}
++
++// CHECK-LABEL: "op_reshape"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
++  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
++  %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
++  func.return %0 : tensor<4x4xf32>
++}
++
++// CHECK-LABEL: "op_return"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
++  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.case"(%arg0) ({
++    "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
++  }) : (tensor<i32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_reverse"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
++  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.reverse"(%arg0) {
++    dimensions = array<i64: 0>
++  } : (tensor<16xf32>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_rng_bit_generator"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
++  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
++  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
++    rng_algorithm = #stablehlo<rng_algorithm PHILOX>
++  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
++  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
++}
++
++// CHECK-LABEL: "op_rng"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
++  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
++    rng_distribution = #stablehlo<rng_distribution NORMAL>
++  } : (tensor<f32>, tensor<f32>, tensor<0xindex>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_round_nearest_afz"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_round_nearest_even"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_rsqrt"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.rsqrt_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_scatter"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
++  //      CHECK: "vhlo.scatter_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   input_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
++  // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   scatter_dims_to_operand_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   scatter_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
++  // CHECK-SAME:   unique_indices = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   update_window_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>
++  %0 = "stablehlo.scatter"(%arg0, %arg1, %arg2) ({
++    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
++      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    scatter_dimension_numbers = #stablehlo.scatter<
++      update_window_dims = [1],
++      inserted_window_dims = [0, 1],
++      scatter_dims_to_operand_dims = [0, 1],
++      index_vector_dim = 1
++    >,
++    indices_are_sorted = true,
++    unique_indices = true
++  } : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) -> tensor<200x100x300xf32>
++  func.return %0 : tensor<200x100x300xf32>
++}
++
++// CHECK-LABEL: "op_scatter_with_batching_dims"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_scatter_with_batching_dims(%arg0: tensor<10x200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<10x200x100x300xf32> {
++  //      CHECK: "vhlo.scatter_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   input_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   scatter_dims_to_operand_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   scatter_indices_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   unique_indices = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   update_window_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<10x200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x200x100x300x!vhlo.f32_v1>
++  %0 = "stablehlo.scatter"(%arg0, %arg1, %arg2) ({
++    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
++      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    scatter_dimension_numbers = #stablehlo.scatter<
++      update_window_dims = [1],
++      inserted_window_dims = [1, 2],
++      input_batching_dims = [0],
++      scatter_dims_to_operand_dims = [1, 2],
++      scatter_indices_batching_dims = [0],
++      index_vector_dim = 1
++    >,
++    indices_are_sorted = true,
++    unique_indices = true
++  } : (tensor<10x200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) -> tensor<10x200x100x300xf32>
++  func.return %0 : tensor<10x200x100x300xf32>
++}
++
++// CHECK_lABEL: "op_scatter_with_promotable_types"
++func.func @op_scatter_with_promotable_types(%input_tensor: tensor<200x100x300xf32>,
++    %scatter_indices: tensor<10x2xi32>, %updates: tensor<10x300xf32>) ->
++      tensor<200x100x300xf64> {
++  //  CHECK: "vhlo.scatter_v2"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]])
++  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
++  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
++  //  CHECK: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f64_v1>
++  %0 = "stablehlo.scatter" (%input_tensor, %scatter_indices, %updates) ({
++  ^bb0(%lhs: tensor<f64>, %rhs: tensor<f64>):
++    %add = stablehlo.add %lhs, %rhs : tensor<f64>
++    "stablehlo.return"(%add) : (tensor<f64>) -> ()
++  }) {
++    scatter_dimension_numbers = #stablehlo.scatter<
++      update_window_dims = [1],
++      inserted_window_dims = [0, 1],
++      scatter_dims_to_operand_dims = [0, 1],
++      index_vector_dim = 1
++    >,
++    indices_are_sorted = true,
++    unique_indices = true
++  } : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) ->
++      tensor<200x100x300xf64>
++  func.return %0 : tensor<200x100x300xf64>
++}
++
++// CHECK-LABEL: "op_select_and_scatter"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
++  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
++  // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
++  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG31:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG41:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL11:.*]] = "vhlo.compare_v1"(%[[ARG31]], %[[ARG41]]) <{compare_type = #vhlo<comparison_type_v1 TOTALORDER>, comparison_direction = #vhlo<comparison_direction_v1 GE>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL11]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
++  // CHECK-NEXT: }, {
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG32:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG42:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL12:.*]] = "vhlo.add_v1"(%[[ARG32]], %[[ARG42]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL12]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>, !vhlo.tensor_v1<12x13x13x66x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>
++  %0 = "stablehlo.select_and_scatter"(%arg0, %arg1, %arg2) ({
++    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
++      %1 = "stablehlo.compare"(%arg3, %arg4) {compare_type = #stablehlo<comparison_type TOTALORDER>, comparison_direction = #stablehlo<comparison_direction GE>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
++      "stablehlo.return"(%1) : (tensor<i1>) -> ()
++  }, {
++    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
++      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    window_dimensions = array<i64: 1, 2, 2, 1>,
++    window_strides = array<i64: 1, 2, 2, 1>,
++    padding = dense<1> : tensor<4x2xi64>
++  } : (tensor<10x24x24x64xf32>, tensor<12x13x13x66xf32>, tensor<f32>) -> tensor<10x24x24x64xf32>
++  func.return %0 : tensor<10x24x24x64xf32>
++}
++
++// CHECK-LABEL: "op_select_and_scatter_with_promotable_types"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_select_and_scatter_with_promotable_types(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf64> {
++  // CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
++  // CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
++  // CHECK:     %[[VAL:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
++  // CHECK:     "vhlo.return_v1"(%[[VAL]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
++  // CHECK: }) : (!vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>, !vhlo.tensor_v1<12x13x13x66x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x24x24x64x!vhlo.f64_v1>
++  %0 = "stablehlo.select_and_scatter"(%arg0, %arg1, %arg2) ({
++    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
++      %1 = "stablehlo.compare"(%arg3, %arg4) {compare_type = #stablehlo<comparison_type TOTALORDER>, comparison_direction = #stablehlo<comparison_direction GE>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
++      "stablehlo.return"(%1) : (tensor<i1>) -> ()
++  }, {
++    ^bb0(%arg3: tensor<f64>, %arg4: tensor<f64>):
++      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f64>, tensor<f64>) -> tensor<f64>
++      "stablehlo.return"(%1) : (tensor<f64>) -> ()
++  }) {
++    window_dimensions = array<i64: 1, 2, 2, 1>,
++    window_strides = array<i64: 1, 2, 2, 1>,
++    padding = dense<1> : tensor<4x2xi64>
++  } : (tensor<10x24x24x64xf32>, tensor<12x13x13x66xf32>, tensor<f32>) -> tensor<10x24x24x64xf64>
++  func.return %0 : tensor<10x24x24x64xf64>
++}
++
++// CHECK-LABEL: "op_select"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_send_no_source_target_pairs"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_send_no_source_target_pairs(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
++  //      CHECK: "vhlo.send_v2"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>,
++  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
++  // CHECK-SAME{LITERAL}: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
++  %0 = "stablehlo.send"(%arg0, %arg1) {
++    channel_handle = #stablehlo.channel_handle<handle = 0, type = 2>,
++    is_host_transfer = true
++  } : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
++  func.return %0 : !stablehlo.token
++}
++
++// CHECK-LABEL: "op_set_dimension_size"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
++  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
++    dimension = 0 : i64
++  } : (tensor<?xf32>, tensor<i32>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_shift_left"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
++  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
++  %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
++  func.return %0 : tensor<i32>
++}
++
++// CHECK-LABEL: "op_shift_right_arithmetic"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
++  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
++  %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
++  func.return %0 : tensor<i32>
++}
++
++// CHECK-LABEL: "op_shift_right_logical"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
++  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
++  %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
++  func.return %0 : tensor<i32>
++}
++
++// CHECK-LABEL: "op_sign"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_sine"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.sine_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_slice"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
++  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
++  // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
++  %0 = "stablehlo.slice"(%arg0) {
++    start_indices = array<i64: 0>,
++    limit_indices = array<i64: 4>,
++    strides = array<i64: 1>
++  } : (tensor<16xf32>) -> tensor<4xf32>
++  func.return %0 : tensor<4xf32>
++}
++
++// CHECK-LABEL: "op_sort"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
++  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.compare_v1"(%[[ARG1]], %[[ARG2]]) <{compare_type = #vhlo<comparison_type_v1 FLOAT>, comparison_direction = #vhlo<comparison_direction_v1 GT>}>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.sort"(%arg0) ({
++    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
++      %1 = "stablehlo.compare"(%arg1, %arg2) {compare_type = #stablehlo<comparison_type FLOAT>, comparison_direction = #stablehlo<comparison_direction GT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
++      "stablehlo.return"(%1) : (tensor<i1>) -> ()
++  }) {
++    dimension = 0 : i64,
++    is_stable = true
++  } : (tensor<16xf32>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_sqrt"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.sqrt_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_subtract"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_tan"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_tan(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.tan_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.tan"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_tanh"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.tanh_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_torch_index_select"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
++  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
++  %0 = "stablehlo.torch_index_select"(%arg0, %arg1) {
++    dim = 0 : i64,
++    batch_dims = 0 : i64
++  } : (tensor<5x1x5xf32>, tensor<2xi32>) -> tensor<2x1x5xf32>
++  func.return %0 : tensor<2x1x5xf32>
++}
++
++// CHECK-LABEL: "op_transpose"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
++  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
++  %0 = "stablehlo.transpose"(%arg0) {
++    permutation = array<i64: 1, 0>
++  } : (tensor<16x8xf32>) -> tensor<8x16xf32>
++  func.return %0 : tensor<8x16xf32>
++}
++
++// CHECK-LABEL: "op_triangular_solve"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
++  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
++  // CHECK-SAME:   unit_diagonal = #vhlo.bool_v1<true>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
++  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
++    left_side = true,
++    lower = true,
++    unit_diagonal = true,
++    transpose_a = #stablehlo<transpose NO_TRANSPOSE>
++  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
++  func.return %0 : tensor<16x16xf32>
++}
++
++// CHECK-LABEL: "op_tuple"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
++  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
++  %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
++  func.return %0 : tuple<tensor<f32>>
++}
++
++// CHECK-LABEL: "op_unary_einsum"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
++  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
++  %0 = "stablehlo.unary_einsum"(%arg0) {
++    einsum_config = "ab->a"
++  } : (tensor<8x16xf32>) -> tensor<8xf32>
++  func.return %0 : tensor<8xf32>
++}
++
++// CHECK-LABEL: "op_uniform_dequantize"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
++  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_uniform_quantize"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
++  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
++  %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
++  func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
++}
++
++// CHECK-LABEL: "op_while"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
++  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
++  // CHECK-NEXT:   }, {
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>)
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  %0 = "stablehlo.while"(%arg0) ({
++    ^bb0(%arg1: tensor<i1>):
++      "stablehlo.return"(%arg1) : (tensor<i1>) -> ()
++    }, {
++    ^bb0(%arg1: tensor<i1>):
++      "stablehlo.return"(%arg1) : (tensor<i1>) -> ()
++  }) : (tensor<i1>) -> tensor<i1>
++  func.return %0: tensor<i1>
++}
++
++// CHECK-LABEL: "op_xor"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
++  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// ============ TYPES ============
++
++// CHECK-LABEL: "type_i1"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
++  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "type_i2"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_i2(%arg0: tensor<i2>, %arg1: tensor<i2>) -> tensor<i2> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i2_v1>, !vhlo.tensor_v1<!vhlo.i2_v1>) -> !vhlo.tensor_v1<!vhlo.i2_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i2>, tensor<i2>) -> tensor<i2>
++  func.return %0 : tensor<i2>
++}
++
++// CHECK-LABEL: "type_i4"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
++  func.return %0 : tensor<i4>
++}
++
++// CHECK-LABEL: "type_i8"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
++  func.return %0 : tensor<i8>
++}
++
++// CHECK-LABEL: "type_i16"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
++  func.return %0 : tensor<i16>
++}
++
++// CHECK-LABEL: "type_i32"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
++  func.return %0 : tensor<i32>
++}
++
++// CHECK-LABEL: "type_i64"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
++  func.return %0 : tensor<i64>
++}
++
++// CHECK-LABEL: "type_ui2"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_ui2(%arg0: tensor<ui2>, %arg1: tensor<ui2>) -> tensor<ui2> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui2_v1>, !vhlo.tensor_v1<!vhlo.ui2_v1>) -> !vhlo.tensor_v1<!vhlo.ui2_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui2>, tensor<ui2>) -> tensor<ui2>
++  func.return %0 : tensor<ui2>
++}
++
++// CHECK-LABEL: "type_ui4"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
++  func.return %0 : tensor<ui4>
++}
++
++// CHECK-LABEL: "type_ui8"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
++  func.return %0 : tensor<ui8>
++}
++
++// CHECK-LABEL: "type_ui16"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
++  func.return %0 : tensor<ui16>
++}
++
++// CHECK-LABEL: "type_ui32"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
++  func.return %0 : tensor<ui32>
++}
++
++// CHECK-LABEL: "type_ui64"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
++  func.return %0 : tensor<ui64>
++}
++
++// CHECK-LABEL: "type_f4E2M1FN"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f4E2M1FN(%arg0: tensor<f4E2M1FN>, %arg1: tensor<f4E2M1FN>) -> tensor<f4E2M1FN> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f4E2M1FN_v1>, !vhlo.tensor_v1<!vhlo.f4E2M1FN_v1>) -> !vhlo.tensor_v1<!vhlo.f4E2M1FN_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f4E2M1FN>, tensor<f4E2M1FN>) -> tensor<f4E2M1FN>
++  func.return %0 : tensor<f4E2M1FN>
++}
++
++// CHECK-LABEL: "type_f6E2M3FN"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f6E2M3FN(%arg0: tensor<f6E2M3FN>, %arg1: tensor<f6E2M3FN>) -> tensor<f6E2M3FN> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f6E2M3FN_v1>, !vhlo.tensor_v1<!vhlo.f6E2M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f6E2M3FN_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f6E2M3FN>, tensor<f6E2M3FN>) -> tensor<f6E2M3FN>
++  func.return %0 : tensor<f6E2M3FN>
++}
++
++// CHECK-LABEL: "type_f6E3M2FN"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f6E3M2FN(%arg0: tensor<f6E3M2FN>, %arg1: tensor<f6E3M2FN>) -> tensor<f6E3M2FN> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f6E3M2FN_v1>, !vhlo.tensor_v1<!vhlo.f6E3M2FN_v1>) -> !vhlo.tensor_v1<!vhlo.f6E3M2FN_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f6E3M2FN>, tensor<f6E3M2FN>) -> tensor<f6E3M2FN>
++  func.return %0 : tensor<f6E3M2FN>
++}
++
++// CHECK-LABEL: "type_f8E3M4"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f8E3M4(%arg0: tensor<f8E3M4>, %arg1: tensor<f8E3M4>) -> tensor<f8E3M4> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E3M4_v1>, !vhlo.tensor_v1<!vhlo.f8E3M4_v1>) -> !vhlo.tensor_v1<!vhlo.f8E3M4_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E3M4>, tensor<f8E3M4>) -> tensor<f8E3M4>
++  func.return %0 : tensor<f8E3M4>
++}
++
++// CHECK-LABEL: "type_f8E4M3"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f8E4M3(%arg0: tensor<f8E4M3>, %arg1: tensor<f8E4M3>) -> tensor<f8E4M3> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3>, tensor<f8E4M3>) -> tensor<f8E4M3>
++  func.return %0 : tensor<f8E4M3>
++}
++
++// CHECK-LABEL: "type_f8E4M3FN"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
++  func.return %0 : tensor<f8E4M3FN>
++}
++
++// CHECK-LABEL: "type_f8E5M2"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
++  func.return %0 : tensor<f8E5M2>
++}
++
++// CHECK-LABEL: "type_f8E4M3FNUZ"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
++  func.return %0 : tensor<f8E4M3FNUZ>
++}
++
++// CHECK-LABEL: "type_f8E4M3B11FNUZ"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
++  func.return %0 : tensor<f8E4M3B11FNUZ>
++}
++
++// CHECK-LABEL: "type_f8E5M2FNUZ"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
++  func.return %0 : tensor<f8E5M2FNUZ>
++}
++
++// CHECK-LABEL: "type_f8E8M0FNU"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f8E8M0FNU(%arg0: tensor<f8E8M0FNU>, %arg1: tensor<f8E8M0FNU>) -> tensor<f8E8M0FNU> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E8M0FNU_v1>, !vhlo.tensor_v1<!vhlo.f8E8M0FNU_v1>) -> !vhlo.tensor_v1<!vhlo.f8E8M0FNU_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E8M0FNU>, tensor<f8E8M0FNU>) -> tensor<f8E8M0FNU>
++  func.return %0 : tensor<f8E8M0FNU>
++}
++
++// CHECK-LABEL: "type_bf16"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
++  func.return %0 : tensor<bf16>
++}
++
++// CHECK-LABEL: "type_f16"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
++  func.return %0 : tensor<f16>
++}
++
++// CHECK-LABEL: "type_f32"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "type_f64"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
++  func.return %0 : tensor<f64>
++}
++
++// CHECK-LABEL: "type_complex_f32"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
++  func.return %0 : tensor<complex<f32>>
++}
++
++// CHECK-LABEL: "type_complex_f64"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
++  func.return %0 : tensor<complex<f64>>
++}
++
++// CHECK-LABEL: "type_tf32"
++// CHECK: #vhlo.type_v1<!vhlo.tf31_v1>
++func.func @type_tf32() attributes {stablehlo.attr = tf32 } {
++  return
++}
++
++// CHECK-LABEL: "type_none"
++// CHECK: #vhlo.type_v1<!vhlo.none_v1>
++func.func @type_none() attributes {stablehlo.attr = none } {
++  return
++}
++
++// CHECK-LABEL: "type_dynamism_ranked"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
++  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
++  %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
++  func.return %0 : tensor<?xf32>
++}
++
++// CHECK-LABEL: "type_per_tensor_quantization"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_per_tensor_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
++  func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
++}
++
++// CHECK-LABEL: "type_per_axis_quantization"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @type_per_axis_quantization(%arg0: tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>) -> tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG0]]) : (!vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, [3.400000e+01, 3.400000e+01], [16, 16], -128:127, 1>>, !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, [3.400000e+01, 3.400000e+01], [16, 16], -128:127, 1>>) -> !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, [3.400000e+01, 3.400000e+01], [16, 16], -128:127, 1>>
++  %0 = stablehlo.add %arg0, %arg0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
++  func.return %0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
++}
++
++//       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
++// CHECK-LABEL: "type_token_callee"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
++  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
++  return %arg0 : !stablehlo.token
++}
++
++//       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
++// CHECK-LABEL: "type_token_caller"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
++  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
++  // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
++  %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
++  return %0 : !stablehlo.token
++}
++
++// CHECK-LABEL: "type_tuple"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "foo"
++  // CHECK: (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>) -> !vhlo.tuple_v1<!vhlo.token_v1>
++  } : (tuple<tensor<f32>>) -> tuple<!stablehlo.token>
++  return %0 : tuple<!stablehlo.token>
++}
++
++// CHECK-LABEL: type_buffer_function_input_output
++// CHECK-NEXT: (%[[ARG0:.*]]: !vhlo.buffer_v1<2x!vhlo.f32_v1>)
++func.func @type_buffer_function_input_output(%arg0: memref<2xf32>) -> memref<2xf32> {
++  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.buffer_v1<2x!vhlo.f32_v1>) -> ()
++  func.return %arg0 : memref<2xf32>
++}
++
++// CHECK-LABEL: type_buffer_special_custom_calls
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @type_buffer_special_custom_calls(%arg0: tensor<2xf32>) -> tensor<2xf32> {
++  //               CHECK: %[[CALL0:.*]] = "vhlo.custom_call_v2"(%[[ARG0]])
++  //          CHECK-SAME: call_target_name = #vhlo.string_v1<"Pin">
++  //          CHECK-SAME: : (!vhlo.tensor_v1<2x!vhlo.f32_v1>) -> !vhlo.buffer_v1<2x!vhlo.f32_v1>
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "Pin",
++    api_version = 4 : i32
++  } : (tensor<2xf32>) -> memref<2xf32>
++  //               CHECK: %{{.*}} = "vhlo.custom_call_v2"(%[[CALL0]])
++  //          CHECK-SAME: call_target_name = #vhlo.string_v1<"Unpin">
++  //          CHECK-SAME: : (!vhlo.buffer_v1<2x!vhlo.f32_v1>) -> !vhlo.tensor_v1<2x!vhlo.f32_v1>
++  %1 = "stablehlo.custom_call"(%0) {
++    call_target_name = "Unpin",
++    api_version = 4 : i32
++  } : (memref<2xf32>) -> tensor<2xf32>
++  func.return %1 : tensor<2xf32>
++}
++
++// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>>>
++// CHECK-LABEL: "type_future"
++func.func private @type_future() -> !stablehlo.future<tensor<f32>>
++
++// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<2x3x!vhlo.f32_v1>>>>
++// CHECK-LABEL: "type_future_ranked"
++func.func private @type_future_ranked() -> !stablehlo.future<tensor<2x3xf32>>
++
++// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<?x!vhlo.f32_v1>>>>
++// CHECK-LABEL: "type_future_dynamic"
++func.func private @type_future_dynamic() -> !stablehlo.future<tensor<?xf32>>
++
++// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>>>
++// CHECK-LABEL: "type_future_multiple"
++func.func private @type_future_multiple() -> !stablehlo.future<tensor<f32>, tensor<i32>>
++
++// CHECK-LABEL: "op_async_start_all_reduce"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_async_start_all_reduce(%arg0: tensor<4x4xf32>) -> !stablehlo.future<tensor<4x4xf32>> {
++  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
++  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.all_reduce_v2"(%[[BARG0]]) <{
++  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>,
++  // CHECK-SAME:     use_global_device_ids = #vhlo.bool_v1<false>
++  // CHECK-SAME:   }> ({
++  // CHECK-NEXT:     ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:       %[[VAL2:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  // CHECK-NEXT:       "vhlo.return_v1"(%[[VAL2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT:   }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<4x4x!vhlo.f32_v1>>
++  %0 = "stablehlo.async_start"(%arg0) ({
++    ^bb0(%barg0: tensor<4x4xf32>):
++    %1 = "stablehlo.all_reduce"(%barg0) ({
++    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
++      %2 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%2) : (tensor<f32>) -> ()
++    }) {replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
++    "stablehlo.return"(%1) : (tensor<4x4xf32>) -> ()
++  }) : (tensor<4x4xf32>) -> !stablehlo.future<tensor<4x4xf32>>
++  func.return %0: !stablehlo.future<tensor<4x4xf32>>
++}
++
++// CHECK-LABEL: "op_async_start_all_gather"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_async_start_all_gather(%arg0: tensor<8x2xf32>) -> !stablehlo.future<tensor<8x8xf32>> {
++  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
++  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.all_gather_v2"(%[[BARG0]]) <{
++  // CHECK-SAME:     all_gather_dim = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 2, 4, 6], [1, 3, 5, 7]]> : tensor<2x4xi64>>,
++  // CHECK-SAME:     use_global_device_ids = #vhlo.bool_v1<false>
++  // CHECK-SAME:   }> : (!vhlo.tensor_v1<8x2x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<8x8x!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<8x2x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<8x8x!vhlo.f32_v1>>
++  %0 = "stablehlo.async_start"(%arg0) ({
++    ^bb0(%barg0: tensor<8x2xf32>):
++    %1 = "stablehlo.all_gather"(%barg0) {
++      all_gather_dim = 1 : i64,
++      replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7]]> : tensor<2x4xi64>
++    } : (tensor<8x2xf32>) -> tensor<8x8xf32>
++    "stablehlo.return"(%1) : (tensor<8x8xf32>) -> ()
++  }) : (tensor<8x2xf32>) -> !stablehlo.future<tensor<8x8xf32>>
++  func.return %0: !stablehlo.future<tensor<8x8xf32>>
++}
++
++// CHECK-LABEL: "op_async_start_all_to_all"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_async_start_all_to_all(%arg0: tensor<4x16xf32>) -> !stablehlo.future<tensor<16x4xf32>> {
++  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
++  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.all_to_all_v2"(%[[BARG0]]) <{
++  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:     concat_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
++  // CHECK-SAME:     split_count = #vhlo.integer_v1<4 : i64>,
++  // CHECK-SAME:     split_dimension = #vhlo.integer_v1<1 : i64>
++  // CHECK-SAME:   }> : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x4x!vhlo.f32_v1>
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<16x4x!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<16x4x!vhlo.f32_v1>>
++  %0 = "stablehlo.async_start"(%arg0) ({
++    ^bb0(%barg0: tensor<4x16xf32>):
++    %1 = "stablehlo.all_to_all"(%barg0) {
++      split_dimension = 1 : i64,
++      concat_dimension = 0 : i64,
++      split_count = 4 : i64,
++      replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>
++    } : (tensor<4x16xf32>) -> tensor<16x4xf32>
++    "stablehlo.return"(%1) : (tensor<16x4xf32>) -> ()
++  }) : (tensor<4x16xf32>) -> !stablehlo.future<tensor<16x4xf32>>
++  func.return %0: !stablehlo.future<tensor<16x4xf32>>
++}
++
++// CHECK-LABEL: "op_async_start_collective_permute"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_async_start_collective_permute(%arg0: tensor<128x32xf32>) -> !stablehlo.future<tensor<128x32xf32>> {
++  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
++  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.collective_permute_v1"(%[[BARG0]]) <{
++  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:     source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
++  // CHECK-SAME:   }> : (!vhlo.tensor_v1<128x32x!vhlo.f32_v1>) -> !vhlo.tensor_v1<128x32x!vhlo.f32_v1>
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<128x32x!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<128x32x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<128x32x!vhlo.f32_v1>>
++  %0 = "stablehlo.async_start"(%arg0) ({
++    ^bb0(%barg0: tensor<128x32xf32>):
++    %1 = "stablehlo.collective_permute"(%barg0) {
++      source_target_pairs = dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>
++    } : (tensor<128x32xf32>) -> tensor<128x32xf32>
++    "stablehlo.return"(%1) : (tensor<128x32xf32>) -> ()
++  }) : (tensor<128x32xf32>) -> !stablehlo.future<tensor<128x32xf32>>
++  func.return %0: !stablehlo.future<tensor<128x32xf32>>
++}
++
++// CHECK-LABEL: "op_async_start_collective_broadcast"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_async_start_collective_broadcast(%arg0: tensor<16x8xf32>) -> !stablehlo.future<tensor<16x8xf32>> {
++  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
++  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.collective_broadcast_v2"(%[[BARG0]]) <{
++  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:     has_dynamic_root = #vhlo.bool_v1<false>,
++  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
++  // CHECK-SAME:   }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<16x8x!vhlo.f32_v1>>
++  %0 = "stablehlo.async_start"(%arg0) ({
++    ^bb0(%barg0: tensor<16x8xf32>):
++    %1 = "stablehlo.collective_broadcast"(%barg0) {
++      replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
++    } : (tensor<16x8xf32>) -> tensor<16x8xf32>
++    "stablehlo.return"(%1) : (tensor<16x8xf32>) -> ()
++  }) : (tensor<16x8xf32>) -> !stablehlo.future<tensor<16x8xf32>>
++  func.return %0: !stablehlo.future<tensor<16x8xf32>>
++}
++
++// CHECK-LABEL: "op_async_start_reduce_scatter"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_async_start_reduce_scatter(%arg0: tensor<4x16xf32>) -> !stablehlo.future<tensor<4x4xf32>> {
++  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
++  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.reduce_scatter_v1"(%[[BARG0]]) <{
++  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
++  // CHECK-SAME:     scatter_dimension = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:     use_global_device_ids = #vhlo.bool_v1<false>
++  // CHECK-SAME:   }> ({
++  // CHECK-NEXT:     ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:       %[[VAL2:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  // CHECK-NEXT:       "vhlo.return_v1"(%[[VAL2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT:   }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<4x4x!vhlo.f32_v1>>
++  %0 = "stablehlo.async_start"(%arg0) ({
++    ^bb0(%barg0: tensor<4x16xf32>):
++    %1 = "stablehlo.reduce_scatter"(%barg0) ({
++      ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
++      %2 = stablehlo.add %arg2, %arg3 : tensor<f32>
++      "stablehlo.return"(%2) : (tensor<f32>) -> ()
++    }) {replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
++        scatter_dimension = 1 : i64} : (tensor<4x16xf32>) -> tensor<4x4xf32>
++    "stablehlo.return"(%1) : (tensor<4x4xf32>) -> ()
++  }) : (tensor<4x16xf32>) -> !stablehlo.future<tensor<4x4xf32>>
++  func.return %0: !stablehlo.future<tensor<4x4xf32>>
++}
++
++// CHECK-LABEL: "op_async_done"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_async_done(%arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
++  // CHECK: %[[VAL1:.*]] = "vhlo.async_start_v1"(%[[ARG0]])
++  // CHECK: %[[VAL2:.*]] = "vhlo.async_done_v1"(%[[VAL1]]) : (!vhlo.future_v1<!vhlo.tensor_v1<4x4x!vhlo.f32_v1>>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
++  %0 = "stablehlo.async_start"(%arg0) ({
++    ^bb0(%barg0: tensor<4x4xf32>):
++    %1 = "stablehlo.all_reduce"(%barg0) ({
++    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
++      %2 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%2) : (tensor<f32>) -> ()
++    }) {replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
++    "stablehlo.return"(%1) : (tensor<4x4xf32>) -> ()
++  }) : (tensor<4x4xf32>) -> !stablehlo.future<tensor<4x4xf32>>
++  %1 = "stablehlo.async_done"(%0) : (!stablehlo.future<tensor<4x4xf32>>) -> tensor<4x4xf32>
++  func.return %1: tensor<4x4xf32>
++}
++
++// ============ DEPENDENCIES  ============
++
++func.func @composite_target(%arg0: tensor<f32>) -> tensor<f32> {
++  return %arg0: tensor<f32>
++}
+diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_22_0.mlir.bc b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_22_0.mlir.bc
+Binary files /dev/null and stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_22_0.mlir.bc differ
 diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir
 --- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir
 +++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir
-@@ -749,6 +749,80 @@
+@@ -128,6 +128,17 @@
+   func.return %0 : tensor<i1>
+ }
+ 
++// CHECK-LABEL: "attr_comparison_type_weakorder"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_type_weakorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction EQ>,
++    // CHECK: compare_type = #vhlo<comparison_type_v1 WEAKORDER>,
++    compare_type = #stablehlo<comparison_type WEAKORDER>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
+ // ConvDimensionNumbers aka #stablehlo.conv is covered below.
+ 
+ // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
+@@ -749,6 +760,80 @@
    func.return %0 : tensor<8x8x8xf32>
  }
  
@@ -8216,6 +12232,37 @@ diff --ruN a/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_
 +    func.return %0 : tensor<8x8x8xf32>
 +  }
 +}
+diff --ruN a/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_21_0.mlir b/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_21_0.mlir
+--- stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_21_0.mlir
++++ stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_21_0.mlir
+@@ -0,0 +1,27 @@
++// RUN: stablehlo-opt --stablehlo-legalize-to-vhlo --vhlo-to-version='target=1.21.0' --verify-diagnostics --split-input-file %s
++
++// expected-error @+1 {{failed to convert VHLO to v1.21.0}}
++module {
++  func.func @compare_weakorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++    // expected-error @+1 {{failed to legalize operation 'vhlo.compare_v1' that was explicitly marked illegal}}
++    %0 = "stablehlo.compare"(%arg0, %arg1) {
++      comparison_direction = #stablehlo<comparison_direction LT>,
++      compare_type = #stablehlo<comparison_type WEAKORDER>
++    } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++    func.return %0 : tensor<i1>
++  }
++}
++
++// -----
++
++// expected-error @+1 {{failed to convert VHLO to v1.21.0}}
++module {
++  func.func @discardable_weakorder_attr(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++    // expected-error @+1 {{failed to legalize operation 'vhlo.add_v1' that was explicitly marked illegal}}
++    %0 = "stablehlo.add"(%arg0, %arg1) {
++      some_attr = #stablehlo<comparison_type WEAKORDER>
++    } : (tensor<f32>, tensor<f32>) -> tensor<f32>
++    func.return %0 : tensor<f32>
++  }
++}
++
 diff --ruN a/stablehlo/stablehlo/tools/StablehloLspServerMain.cpp b/stablehlo/stablehlo/tools/StablehloLspServerMain.cpp
 --- stablehlo/stablehlo/tools/StablehloLspServerMain.cpp
 +++ stablehlo/stablehlo/tools/StablehloLspServerMain.cpp
@@ -8385,7 +12432,16 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloDecompositionUtils.h b/stablehlo
 diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp
 --- stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp
 +++ stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp
-@@ -2520,7 +2520,7 @@
+@@ -111,6 +111,8 @@
+       return mlir::stablehlo::ComparisonType::SIGNED;
+     case mlir::chlo::ComparisonType::UNSIGNED:
+       return mlir::stablehlo::ComparisonType::UNSIGNED;
++    case mlir::chlo::ComparisonType::WEAKORDER:
++      return mlir::stablehlo::ComparisonType::WEAKORDER;
+   }
+   return {};
+ }
+@@ -2520,7 +2522,7 @@
      }
      Value dotGeneral = mlir::stablehlo::DotGeneralOp::create(
          rewriter, op.getLoc(), resultType, ValueRange{lhsMasked, rhsReshaped},
@@ -8394,7 +12450,7 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stable
  
      if (i == 0) {
        out = dotGeneral;
-@@ -2672,7 +2672,7 @@
+@@ -2672,7 +2674,7 @@
      }
      Value dotGeneral = mlir::stablehlo::DotGeneralOp::create(
          rewriter, op.getLoc(), RankedTensorType::get(resultShape, outDType),
@@ -8403,7 +12459,7 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stable
  
      results.push_back(dotGeneral);
      start = limit;
-@@ -2747,7 +2747,7 @@
+@@ -2747,7 +2749,7 @@
    }
    rewriter.replaceOpWithNewOp<mlir::stablehlo::DotGeneralOp>(
        op, op.getResult().getType(), ValueRange{op.getLhs(), op.getRhs()},
@@ -8412,7 +12468,7 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stable
    return success();
  }
  
-@@ -2930,8 +2930,18 @@
+@@ -2930,8 +2932,18 @@
        SmallVector<int64_t> sliceShape(inputType.getShape().begin(),
                                        inputType.getShape().end());
        sliceShape[dim] = 1;
@@ -8433,7 +12489,7 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stable
  
        Value slice;
        if (inputType.hasStaticShape()) {
-@@ -2980,9 +2990,18 @@
+@@ -2980,9 +2992,18 @@
        for (int64_t d = 0; d < rank; ++d) {
          if (d != dim) resultShape.push_back(sliceShape[d]);
        }
@@ -8453,7 +12509,7 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stable
            slice);
        slicedInputs.push_back(reshaped);
      }
-@@ -3004,9 +3023,20 @@
+@@ -3004,9 +3025,20 @@
        SmallVector<int64_t> elementShape(outputType.getShape().begin(),
                                          outputType.getShape().end());
        elementShape[dim] = 1;
@@ -8949,6 +13005,46 @@ diff --ruN a/stablehlo/stablehlo/transforms/optimization/Passes.h b/stablehlo/st
  
  }  // namespace stablehlo
  }  // namespace mlir
+diff --ruN a/stablehlo/stablehlo/transforms/optimization/StablehloAggressiveFolder.cpp b/stablehlo/stablehlo/transforms/optimization/StablehloAggressiveFolder.cpp
+--- stablehlo/stablehlo/transforms/optimization/StablehloAggressiveFolder.cpp
++++ stablehlo/stablehlo/transforms/optimization/StablehloAggressiveFolder.cpp
+@@ -743,6 +743,36 @@
+     ComparisonType kind;
+ 
+     APInt operator()(APFloat lhs, APFloat rhs) {
++      if (kind == ComparisonType::WEAKORDER) {
++        bool eq = (lhs == rhs) || (lhs.isNaN() && rhs.isNaN());
++        auto lt = [](const APFloat& a, const APFloat& b) {
++          return (!a.isNaN() && b.isNaN()) ||
++                 (!a.isNaN() && !b.isNaN() && (a < b));
++        };
++        bool result = false;
++        switch (direction) {
++          case ComparisonDirection::EQ:
++            result = eq;
++            break;
++          case ComparisonDirection::NE:
++            result = !eq;
++            break;
++          case ComparisonDirection::GE:
++            result = lt(rhs, lhs) || eq;
++            break;
++          case ComparisonDirection::GT:
++            result = lt(rhs, lhs);
++            break;
++          case ComparisonDirection::LE:
++            result = lt(lhs, rhs) || eq;
++            break;
++          case ComparisonDirection::LT:
++            result = lt(lhs, rhs);
++            break;
++        }
++        return APInt(/*bitwidth=*/1, result);
++      }
++
+       if (kind != ComparisonType::FLOAT && kind != ComparisonType::TOTALORDER)
+         llvm::report_fatal_error("invalid float comparison");
+ 
 diff --ruN a/stablehlo/stablehlo/transforms/optimization/StablehloAggressiveSimplification.cpp b/stablehlo/stablehlo/transforms/optimization/StablehloAggressiveSimplification.cpp
 --- stablehlo/stablehlo/transforms/optimization/StablehloAggressiveSimplification.cpp
 +++ stablehlo/stablehlo/transforms/optimization/StablehloAggressiveSimplification.cpp
diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc
index 6fc09a2fe7c..8537e6e6c82 100644
--- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc
+++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc
@@ -2624,10 +2624,10 @@ mlir::NamedAttribute HloFunctionImporter::ConvertComparisonOrder(
                               builder_->getContext(),
                               mlir::stablehlo::ComparisonType::TOTALORDER));
     case ComparisonOrder::kWeak:
-      // TODO(b/565612124): Map to mlir::stablehlo::ComparisonType::WEAKORDER
-      // once added to StableHLO.
-      LOG(FATAL) << "Unsupported comparison order: "
-                 << ComparisonOrderToString(order);
+      return builder_->getNamedAttr(
+          "compare_type", mlir::stablehlo::ComparisonTypeAttr::get(
+                              builder_->getContext(),
+                              mlir::stablehlo::ComparisonType::WEAKORDER));
   }
   LOG(FATAL) << "Unhandled comparison order: " << static_cast<int>(order);
 }
diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import.hlo b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import.hlo
index 0e39f1877ea..26fae8d6955 100644
--- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import.hlo
+++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import.hlo
@@ -416,6 +416,9 @@ add {
   // CHECK-NEXT:  mhlo.compare LE, %arg0, %arg1, TOTALORDER : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xi1>
   %compare.5 = pred[3] compare(Arg_0.1, Arg_1.2), direction=LE, type=TOTALORDER
 
+  // CHECK-NEXT:  mhlo.compare LT, %arg0, %arg1, WEAKORDER : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xi1>
+  %compare.weak = pred[3] compare(Arg_0.1, Arg_1.2), direction=LT, order=WEAK
+
   // Requires broadcast of compatible tensors.
   // CHECK-NEXT:  mhlo.compare GT, %arg0, %arg2 : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xi1>
   ROOT %compare.6 = pred[3] compare(Arg_0.1, Arg_2.3), direction=GT
diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc
index fc41a16a3b7..4a501b3a688 100644
--- a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc
+++ b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc
@@ -2577,6 +2577,9 @@ mlir::LogicalResult ExportXlaOp(mlir::stablehlo::CompareOp op,
       case mlir::stablehlo::ComparisonType::UNSIGNED:
         order = xla::ComparisonOrder::kTotal;
         break;
+      case mlir::stablehlo::ComparisonType::WEAKORDER:
+        order = xla::ComparisonOrder::kWeak;
+        break;
       case mlir::stablehlo::ComparisonType::NOTYPE:
         LOG(FATAL) << "Unreachable";
     }
diff --git a/third_party/xla/xla/hlo/translate/tests/stablehlo.mlir b/third_party/xla/xla/hlo/translate/tests/stablehlo.mlir
index 8c9731f6d50..e9e8b6485fa 100644
--- a/third_party/xla/xla/hlo/translate/tests/stablehlo.mlir
+++ b/third_party/xla/xla/hlo/translate/tests/stablehlo.mlir
@@ -1222,6 +1222,33 @@ func.func @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
 
 // -----
 
+// CHECK-LABEL: HloModule main, entry_computation_layout={(f32[16]{0})->f32[16]{0}}
+
+// CHECK:       %[[$region_0_2:[^ ]+]]
+// CHECK-NEXT:  %[[Arg_0_3:[^ ]+]] = f32[] parameter(0)
+// CHECK-NEXT:  %[[Arg_1_4:[^ ]+]] = f32[] parameter(1)
+// CHECK-NEXT:  ROOT %[[compare_5:[^ ]+]] = pred[] compare(%[[Arg_0_3]], %[[Arg_1_4]]), direction=LT, order=WEAK, metadata=
+
+// CHECK:       ENTRY %[[$main_7:[^ ]+]]
+// CHECK-NEXT:  %[[Arg_0_1:[^ ]+]] = f32[16] parameter(0)
+// CHECK-NEXT:  ROOT %[[sort_6:[^ ]+]] = f32[16] sort(%[[Arg_0_1]]), dimensions={0}, is_stable=true, to_apply=%[[$region_0_2]], metadata=
+
+func.func @main(%arg0: tensor<16xf32>) -> tensor<16xf32> {
+  %0 = "stablehlo.sort"(%arg0) ({
+  ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
+  %1 = "stablehlo.compare"(%arg1, %arg2) {compare_type = #stablehlo<comparison_type WEAKORDER>, comparison_direction = #stablehlo<comparison_direction LT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
+  "stablehlo.return"(%1) : (tensor<i1>) -> ()
+  }) {
+  dimension = 0 : i64,
+  is_stable = true
+  } : (tensor<16xf32>) -> tensor<16xf32>
+  func.return %0 : tensor<16xf32>
+}
+// CHECK-DIRECT: stablehlo.sort
+// CHECK-DIRECT: stablehlo.compare{{.*}}WEAKORDER
+
+// -----
+
 // CHECK-LABEL: HloModule main, entry_computation_layout={(s64[])->s64[]}
 
 // CHECK:       %[[$add_n_impl_2:[^ ]+]]
diff --git a/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops.cc b/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops.cc
index cae4e9cb0f2..9366f4929c6 100644
--- a/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops.cc
+++ b/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops.cc
@@ -6353,7 +6353,10 @@ LogicalResult CompareOp::inferReturnTypeComponents(
     mlir::PropertyRef properties, RegionRange regions,
     SmallVectorImpl<ShapedTypeComponents>& inferredReturnShapes) {
   CompareOp::Adaptor adaptor(operands, attributes, properties, regions);
-  return hlo::inferCompareOp(context, location, adaptor.getLhs(),
+  std::optional<StringRef> compareType;
+  if (auto attr = adaptor.getCompareType())
+    compareType = stringifyComparisonType(*attr);
+  return hlo::inferCompareOp(context, location, adaptor.getLhs(), compareType,
                              inferredReturnShapes);
 }
 
@@ -6453,6 +6456,12 @@ OpFoldResult CompareOp::fold(FoldAdaptor adaptor) {
     return {};
   }
 
+  if (isa<FloatType>(opElType) && getCompareType() &&
+      *getCompareType() != ComparisonType::NOTYPE &&
+      *getCompareType() != ComparisonType::FLOAT) {
+    return {};
+  }
+
 #define COMPARE_FOLDER(Op, comparison, Func)                                \
   if (direction == comparison) {                                            \
     if (auto folded = CompareFolder<Op, FloatType, APFloat, Func<APFloat>>( \
diff --git a/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops_enums.td b/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops_enums.td
index d25fc4c97dc..bc2e2f21dd9 100644
--- a/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops_enums.td
+++ b/third_party/xla/xla/mlir_hlo/mhlo/IR/hlo_ops_enums.td
@@ -172,6 +172,7 @@ def MHLO_COMPARISON_TYPE_FLOAT : I32EnumAttrCase<"FLOAT", 1>;
 def MHLO_COMPARISON_TYPE_FLOAT_TOTAL_ORDER : I32EnumAttrCase<"TOTALORDER", 2>;
 def MHLO_COMPARISON_TYPE_SIGNED : I32EnumAttrCase<"SIGNED", 3>;
 def MHLO_COMPARISON_TYPE_UNSIGNED : I32EnumAttrCase<"UNSIGNED", 4>;
+def MHLO_COMPARISON_TYPE_WEAK_ORDER : I32EnumAttrCase<"WEAKORDER", 5>;
 
 def MHLO_ComparisonType : I32EnumAttr<"ComparisonType",
     "Which comparison type to use.",
@@ -180,7 +181,8 @@ def MHLO_ComparisonType : I32EnumAttr<"ComparisonType",
       MHLO_COMPARISON_TYPE_FLOAT,
       MHLO_COMPARISON_TYPE_FLOAT_TOTAL_ORDER,
       MHLO_COMPARISON_TYPE_SIGNED,
-      MHLO_COMPARISON_TYPE_UNSIGNED
+      MHLO_COMPARISON_TYPE_UNSIGNED,
+      MHLO_COMPARISON_TYPE_WEAK_ORDER
     ]> {
   let genSpecializedAttr = 0;
   let cppNamespace = "::mlir::mhlo";
diff --git a/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_chlo_to_hlo_op.h b/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_chlo_to_hlo_op.h
index fc043d2484c..46ed4871ed9 100644
--- a/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_chlo_to_hlo_op.h
+++ b/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_chlo_to_hlo_op.h
@@ -67,6 +67,8 @@ inline std::optional<mhlo::ComparisonType> mhloComparisonType(
       return mhlo::ComparisonType::SIGNED;
     case chlo::ComparisonType::UNSIGNED:
       return mhlo::ComparisonType::UNSIGNED;
+    case chlo::ComparisonType::WEAKORDER:
+      return mhlo::ComparisonType::WEAKORDER;
   }
   return {};
 }
diff --git a/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_mhlo_to_scalar_op.h b/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_mhlo_to_scalar_op.h
index e54fc226142..1d170bc7faa 100644
--- a/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_mhlo_to_scalar_op.h
+++ b/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_mhlo_to_scalar_op.h
@@ -496,6 +496,60 @@ inline Value mapMhloOpToStdScalarOp<mhlo::CompareOp>(
       assert(predicate.has_value() && "expected valid comparison direction");
       return arith::CmpIOp::create(*b, loc, *predicate, lhsInt, rhsInt);
     }
+    if (adaptor.getCompareType() &&
+        *adaptor.getCompareType() == mhlo::ComparisonType::WEAKORDER) {
+      // Weak order treats -0.0 == +0.0 and orders all NaNs equal and > +Inf.
+      switch (comparisonDirection) {
+        case mhlo::ComparisonDirection::LE: {
+          Value ole = arith::CmpFOp::create(*b, loc, arith::CmpFPredicate::OLE,
+                                            lhs, rhs);
+          Value rhsNan = arith::CmpFOp::create(
+              *b, loc, arith::CmpFPredicate::UNO, rhs, rhs);
+          return arith::OrIOp::create(*b, loc, ole, rhsNan);
+        }
+        case mhlo::ComparisonDirection::GE: {
+          Value oge = arith::CmpFOp::create(*b, loc, arith::CmpFPredicate::OGE,
+                                            lhs, rhs);
+          Value lhsNan = arith::CmpFOp::create(
+              *b, loc, arith::CmpFPredicate::UNO, lhs, lhs);
+          return arith::OrIOp::create(*b, loc, oge, lhsNan);
+        }
+        case mhlo::ComparisonDirection::LT: {
+          Value ult = arith::CmpFOp::create(*b, loc, arith::CmpFPredicate::ULT,
+                                            lhs, rhs);
+          Value lhsOrd = arith::CmpFOp::create(
+              *b, loc, arith::CmpFPredicate::ORD, lhs, lhs);
+          return arith::AndIOp::create(*b, loc, ult, lhsOrd);
+        }
+        case mhlo::ComparisonDirection::GT: {
+          Value ugt = arith::CmpFOp::create(*b, loc, arith::CmpFPredicate::UGT,
+                                            lhs, rhs);
+          Value rhsOrd = arith::CmpFOp::create(
+              *b, loc, arith::CmpFPredicate::ORD, rhs, rhs);
+          return arith::AndIOp::create(*b, loc, ugt, rhsOrd);
+        }
+        case mhlo::ComparisonDirection::EQ: {
+          Value oeq = arith::CmpFOp::create(*b, loc, arith::CmpFPredicate::OEQ,
+                                            lhs, rhs);
+          Value lhsNan = arith::CmpFOp::create(
+              *b, loc, arith::CmpFPredicate::UNO, lhs, lhs);
+          Value rhsNan = arith::CmpFOp::create(
+              *b, loc, arith::CmpFPredicate::UNO, rhs, rhs);
+          Value bothNan = arith::AndIOp::create(*b, loc, lhsNan, rhsNan);
+          return arith::OrIOp::create(*b, loc, oeq, bothNan);
+        }
+        case mhlo::ComparisonDirection::NE: {
+          Value une = arith::CmpFOp::create(*b, loc, arith::CmpFPredicate::UNE,
+                                            lhs, rhs);
+          Value lhsOrd = arith::CmpFOp::create(
+              *b, loc, arith::CmpFPredicate::ORD, lhs, lhs);
+          Value rhsOrd = arith::CmpFOp::create(
+              *b, loc, arith::CmpFPredicate::ORD, rhs, rhs);
+          Value anyOrd = arith::OrIOp::create(*b, loc, lhsOrd, rhsOrd);
+          return arith::AndIOp::create(*b, loc, une, anyOrd);
+        }
+      }
+    }
     std::optional<arith::CmpFPredicate> predicate =
         getCmpPredicate<arith::CmpFPredicate>(comparisonDirection,
                                               /*is_signed=*/true);
diff --git a/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/hlo-legalize-to-stablehlo.mlir b/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/hlo-legalize-to-stablehlo.mlir
index 0fb9e69a7b5..71a5466039b 100644
--- a/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/hlo-legalize-to-stablehlo.mlir
+++ b/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/hlo-legalize-to-stablehlo.mlir
@@ -124,6 +124,16 @@ func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>)
   func.return %0 : tensor<i1>
 }
 
+// CHECK-LABEL: "attr_comparison_type_weak_order"
+func.func @attr_comparison_type_weak_order(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
+  %0 = "mhlo.compare"(%arg0, %arg1) {
+    comparison_direction = #mhlo<comparison_direction EQ>,
+    // CHECK: compare_type = #stablehlo<comparison_type WEAKORDER>,
+    compare_type = #mhlo<comparison_type WEAKORDER>
+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
+  func.return %0 : tensor<i1>
+}
+
 // ConvDimensionNumbers aka #mhlo.conv is covered below.
 
 // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
diff --git a/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/stablehlo-legalize-to-hlo.mlir b/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/stablehlo-legalize-to-hlo.mlir
index 22e2552c0bb..5a5ca25d714 100644
--- a/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/stablehlo-legalize-to-hlo.mlir
+++ b/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/stablehlo-legalize-to-hlo.mlir
@@ -145,6 +145,18 @@ func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>)
 
 // -----
 
+// CHECK-LABEL: "attr_comparison_type_weak_order"
+func.func @attr_comparison_type_weak_order(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
+  %0 = "stablehlo.compare"(%arg0, %arg1) {
+    comparison_direction = #stablehlo<comparison_direction EQ>,
+    // CHECK: compare_type = #mhlo<comparison_type WEAKORDER>,
+    compare_type = #stablehlo<comparison_type WEAKORDER>
+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
+  func.return %0 : tensor<i1>
+}
+
+// -----
+
 // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
 func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
   %0 = "stablehlo.custom_call"(%arg0) {

From f7c5bb483d8ae4df599c8615f21ab693e79466ba Mon Sep 17 00:00:00 2001
From: Kanvi Khanna <kanvi.khanna@intel.com>
Date: Fri, 9 Oct 2026 03:25:19 -0700
Subject: [PATCH 0346/1000] PR #49929: [XLA:GPU][oneAPI]  Part-2 Add oneDNN
 convolution runners for SYCL

Imported from GitHub PR https://github.com/openxla/xla/pull/49929

This is the second PR in a series adding oneDNN-based SYCL convolution runners for fused and unfused convolution APIs. It adds descriptor-based runner setup and retains oneDNN primitive descriptors and workspace requirements. The convolution utilities also support scratchpad allocation and filter reordering or pre-packing where needed.

The PR also adds conversion from dnn::DataType to oneDNN data types and replaces xla::gpu convolution structures with StreamExecutor DNN structures.

The follow-up PR will implement the runner operator APIs to execute the convolutions and add tests for that behavior.
Copybara import of the project:

--
40230f01aaef628df7e83e72c249c8795f6c0927 by Kanvi Khanna <kanvi.khanna@intel.com>:

Use runner API for SYCL oneDNN convolutions

--
0e62beb5c64f06f428b644f3df9854cf062e6c7a by Kanvi Khanna <kanvi.khanna@intel.com>:

Address review comments

--
845c1c0c686930490b5cab804af22d0c19b75da7 by Kanvi Khanna <kanvi.khanna@intel.com>:

Fix formatting

--
c445d3eec4b76a5e8bde470727fa58e3d783fd4c by Kanvi Khanna <kanvi.khanna@intel.com>:

Address review comments

--
60876f9f87f9c75670461fdfa823cf98bb7e7f5b by Kanvi Khanna <kanvi.khanna@intel.com>:

address comments and remove unused api

Merging this change closes #49929

PiperOrigin-RevId: 996424132
---
 .../xla/xla/stream_executor/sycl/BUILD        |   9 +-
 .../xla/stream_executor/sycl/onednn_util.cc   |  21 +
 .../xla/stream_executor/sycl/onednn_util.h    |   5 +
 .../stream_executor/sycl/sycl_conv_utils.cc   | 480 +++++++++---------
 .../stream_executor/sycl/sycl_conv_utils.h    |  63 ++-
 .../xla/xla/stream_executor/sycl/sycl_dnn.cc  | 120 ++++-
 .../xla/xla/stream_executor/sycl/sycl_dnn.h   |  46 +-
 7 files changed, 460 insertions(+), 284 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/sycl/BUILD b/third_party/xla/xla/stream_executor/sycl/BUILD
index 64ba6e9fc01..6f6ff17ac42 100644
--- a/third_party/xla/xla/stream_executor/sycl/BUILD
+++ b/third_party/xla/xla/stream_executor/sycl/BUILD
@@ -598,11 +598,11 @@ cc_library(
         "//xla/tsl/mkl:onednn",
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:logging",
+        "//xla/tsl/protobuf:dnn_proto_cc",
         "//xla/tsl/util:env_var",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/synchronization",
-        "@tsl//tsl/platform:logging",
         "@tsl//tsl/platform:str_util",
     ],
 )
@@ -619,8 +619,7 @@ cc_library(
     deps = [
         ":sycl_conv_utils",
         ":sycl_platform_id",
-        "//xla/service/gpu:gpu_conv_runner",
-        "//xla/stream_executor:device_memory",
+        "//xla/stream_executor:device_address",
         "//xla/stream_executor:dnn",
         "//xla/stream_executor:plugin_registry",
         "//xla/stream_executor:scratch_allocator",
@@ -688,10 +687,8 @@ cc_library(
         "//xla:shape_util",
         "//xla:util",
         "//xla:xla_data_proto_cc",
-        "//xla/service/gpu:cublas_cudnn",
-        "//xla/service/gpu:gpu_conv_runner",
         "//xla/stream_executor:device_address",
-        "//xla/stream_executor:scratch_allocator",
+        "//xla/stream_executor:dnn",
         "//xla/stream_executor:stream",
         "//xla/tsl/mkl:onednn",
         "//xla/tsl/protobuf:dnn_proto_cc",
diff --git a/third_party/xla/xla/stream_executor/sycl/onednn_util.cc b/third_party/xla/xla/stream_executor/sycl/onednn_util.cc
index f08d1e71ed3..be7cadda9cc 100644
--- a/third_party/xla/xla/stream_executor/sycl/onednn_util.cc
+++ b/third_party/xla/xla/stream_executor/sycl/onednn_util.cc
@@ -104,5 +104,26 @@ absl::StatusOr<dnnl::memory::data_type> ToOneDnnDataType(
   }
 }
 
+absl::StatusOr<dnnl::memory::data_type> ToOneDnnDataType(
+    dnn::DataType dnn_type) {
+  switch (dnn_type) {
+    case dnn::DataType::kHalf:
+      return dnnl::memory::data_type::f16;
+    case dnn::DataType::kBF16:
+      return dnnl::memory::data_type::bf16;
+    case dnn::DataType::kFloat:
+      return dnnl::memory::data_type::f32;
+    case dnn::DataType::kDouble:
+      return dnnl::memory::data_type::f64;
+    case dnn::DataType::kInt8:
+      return dnnl::memory::data_type::s8;
+    case dnn::DataType::kInt32:
+      return dnnl::memory::data_type::s32;
+    default:
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Unsupported DNN datatype: ", dnn::DataType_Name(dnn_type)));
+  }
+}
+
 }  // namespace sycl
 }  // namespace stream_executor
diff --git a/third_party/xla/xla/stream_executor/sycl/onednn_util.h b/third_party/xla/xla/stream_executor/sycl/onednn_util.h
index 198ae7b18c9..f46d5d95b13 100644
--- a/third_party/xla/xla/stream_executor/sycl/onednn_util.h
+++ b/third_party/xla/xla/stream_executor/sycl/onednn_util.h
@@ -21,6 +21,7 @@ limitations under the License.
 #include "absl/synchronization/mutex.h"
 #include "dnnl.hpp"
 #include "dnnl_sycl.hpp"
+#include "xla/tsl/protobuf/dnn.pb.h"
 #include "xla/tsl/util/env_var.h"
 #include "xla/xla_data.pb.h"
 #include "tsl/platform/str_util.h"
@@ -49,6 +50,10 @@ dnnl::memory CreateDnnlMemory(const dnnl::memory::desc& md,
 // Converts XLA primitive type to oneDNN data type.
 absl::StatusOr<dnnl::memory::data_type> ToOneDnnDataType(
     xla::PrimitiveType xla_type);
+
+// Converts DNN protobuf data type to oneDNN data type.
+absl::StatusOr<dnnl::memory::data_type> ToOneDnnDataType(
+    dnn::DataType dnn_type);
 }  // namespace sycl
 }  // namespace stream_executor
 #endif  // XLA_STREAM_EXECUTOR_SYCL_ONEDNN_UTIL_H_
diff --git a/third_party/xla/xla/stream_executor/sycl/sycl_conv_utils.cc b/third_party/xla/xla/stream_executor/sycl/sycl_conv_utils.cc
index d7b92bf93a2..496aebe6866 100644
--- a/third_party/xla/xla/stream_executor/sycl/sycl_conv_utils.cc
+++ b/third_party/xla/xla/stream_executor/sycl/sycl_conv_utils.cc
@@ -18,13 +18,13 @@ limitations under the License.
 #include <algorithm>
 #include <cmath>
 #include <cstdint>
+#include <optional>
 #include <utility>
 #include <vector>
 
 #include "absl/base/casts.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
-#include "xla/service/gpu/cublas_cudnn.h"
 #include "xla/shape.h"
 #include "xla/tsl/protobuf/dnn.pb.h"
 #include "xla/tsl/util/env_var.h"
@@ -71,22 +71,25 @@ struct ConvBufferPointers {
 };
 
 absl::StatusOr<ConvBufferPointers> GetConvBufferPointers(
-    xla::gpu::CudnnConvKind conv_kind,
+    dnn::ConvolutionKind conv_kind,
     absl::Span<const DeviceAddressBase> operand_buffers,
     const DeviceAddressBase& result_buffer) {
   auto opaque_ptr = [](const DeviceAddressBase& addr) {
     return const_cast<void*>(addr.opaque());
   };
+  if (operand_buffers.size() < 2) {
+    return absl::InvalidArgumentError("Insufficient operand buffers");
+  }
   void* op0 = opaque_ptr(operand_buffers[0]);
   void* op1 = opaque_ptr(operand_buffers[1]);
   void* res = opaque_ptr(result_buffer);
   switch (conv_kind) {
-    case xla::gpu::CudnnConvKind::kForward:
-    case xla::gpu::CudnnConvKind::kForwardActivation:
+    case dnn::ConvolutionKind::FORWARD:
+    case dnn::ConvolutionKind::FORWARD_BIAS_ACTIVATION:
       return ConvBufferPointers{op0, op1, res};
-    case xla::gpu::CudnnConvKind::kBackwardInput:
+    case dnn::ConvolutionKind::BACKWARD_DATA:
       return ConvBufferPointers{res, op1, op0};
-    case xla::gpu::CudnnConvKind::kBackwardFilter:
+    case dnn::ConvolutionKind::BACKWARD_FILTER:
       return ConvBufferPointers{op0, res, op1};
     default:
       return absl::InvalidArgumentError("Unknown convolution kind");
@@ -164,12 +167,9 @@ dnnl::memory::dims ToOneDnnFilterDims(const FilterDescriptor& descriptor,
 
 // Allocates a temporary buffer and wraps it in a dnnl::memory.
 // Such buffers are used for oneDNN scratchpad and pre-packed filter.
-absl::StatusOr<dnnl::memory> AllocateDnnlBuffer(
-    const dnnl::memory::desc& desc, const dnnl::engine& engine,
-    ScratchAllocator* scratch_allocator) {
-  stream_executor::DeviceMemory<uint8_t> buffer;
-  ABSL_ASSIGN_OR_RETURN(buffer, scratch_allocator->AllocateBytes(desc.get_size()));
-  return CreateDnnlMemory(desc, engine, buffer.opaque());
+absl::StatusOr<dnnl::memory> AllocateDnnlBuffer(const dnnl::memory::desc& desc,
+                                                const dnnl::engine& engine) {
+  return CreateDnnlMemory(desc, engine);
 }
 
 // Builds a forward convolution primitive descriptor. `bias_md` is present only
@@ -207,19 +207,16 @@ absl::StatusOr<ConvOp> BuildConvOp(
     int weights_arg_key, dnnl::memory dst, int dst_arg_key,
     const dnnl::memory::desc& filter_md,
     const dnnl::memory::desc& weights_target_desc, bool prepack_filter,
-    const dnnl::engine& engine, ScratchAllocator* scratch_allocator,
-    std::optional<ReorderOp>* out_filter_reorder) {
+    const dnnl::engine& engine, std::optional<ReorderOp>* out_filter_reorder) {
   ConvOp op;
   op.src = std::move(src);
   op.filter = std::move(filter);
   op.dst = std::move(dst);
-  ABSL_ASSIGN_OR_RETURN(
-      op.scratchpad,
-      AllocateDnnlBuffer(pd.scratchpad_desc(), engine, scratch_allocator));
+  ABSL_ASSIGN_OR_RETURN(op.scratchpad,
+                   AllocateDnnlBuffer(pd.scratchpad_desc(), engine));
   if (filter_md != weights_target_desc) {
-    ABSL_ASSIGN_OR_RETURN(
-        op.internal_filter,
-        AllocateDnnlBuffer(weights_target_desc, engine, scratch_allocator));
+    ABSL_ASSIGN_OR_RETURN(op.internal_filter,
+                     AllocateDnnlBuffer(weights_target_desc, engine));
     *out_filter_reorder = prepack_filter
                               ? CreateReorderOp(op.filter, op.internal_filter)
                               : CreateReorderOp(op.internal_filter, op.filter);
@@ -235,136 +232,94 @@ absl::StatusOr<ConvOp> BuildConvOp(
 }
 }  // namespace
 
-absl::StatusOr<OneDnnConvPrimitive> CreateOneDnnConvPrimitive(
-    const xla::gpu::GpuConvConfig& config,
-    absl::Span<const DeviceAddressBase> operand_buffers,
-    DeviceAddressBase result_buffer, Stream* stream,
-    ScratchAllocator* scratch_allocator) {
-  OneDnnConvPrimitive onednn_conv_primitive;
+absl::StatusOr<OneDnnConvPrimitiveDesc> CreateOneDnnConvPrimitiveDesc(
+    const OneDnnConvConfig& config, Stream* stream) {
+  OneDnnConvPrimitiveDesc pd;
   ::sycl::queue* sycl_queue =
       absl::bit_cast<::sycl::queue*>(stream->platform_specific_handle().stream);
-  onednn_conv_primitive.engine = FindOrCreateEngine(sycl_queue);
-  onednn_conv_primitive.stream = dnnl::sycl_interop::make_stream(
-      onednn_conv_primitive.engine, *sycl_queue);
 
-  DataLayout input_dl = config.input_descriptor.layout();
-  FilterLayout filter_dl = config.filter_descriptor.layout();
-  DataLayout output_dl = config.output_descriptor.layout();
+  const DataLayout input_dl = config.input_descriptor.layout();
+  const FilterLayout filter_dl = config.filter_descriptor.layout();
+  const DataLayout output_dl = config.output_descriptor.layout();
 
-  xla::PrimitiveType input_type;
-  void* input_data;
-  void* filter_data;
-  void* output_data;
-  void* bias_data = nullptr;
-  void* side_input_data = nullptr;
-
-  float alpha = config.conv_result_scale;
-  bool alpha_is_one = (fabs(alpha - 1.0f) < 1e-6);
-  xla::gpu::CudnnConvKind conv_kind = config.kind;
-
-  // Get input type based on convolution kind
+  dnn::DataType input_type;
+  const float alpha = config.conv_result_scale;
+  const bool alpha_is_one = std::fabs(alpha - 1.0f) < 1e-6;
+  const dnn::ConvolutionKind conv_kind = config.kind;
   switch (conv_kind) {
-    case xla::gpu::CudnnConvKind::kForward:
-    case xla::gpu::CudnnConvKind::kForwardActivation:
-    case xla::gpu::CudnnConvKind::kBackwardFilter:
+    case dnn::ConvolutionKind::FORWARD:
+    case dnn::ConvolutionKind::FORWARD_BIAS_ACTIVATION:
+    case dnn::ConvolutionKind::BACKWARD_FILTER:
       input_type = config.input_type;
       break;
-    case xla::gpu::CudnnConvKind::kBackwardInput:
+    case dnn::ConvolutionKind::BACKWARD_DATA:
       input_type = config.output_type;
       break;
     default:
       return absl::InvalidArgumentError("Unknown convolution kind");
   }
 
-  // Get buffer pointers
-  ABSL_ASSIGN_OR_RETURN(
-      auto buffers,
-      GetConvBufferPointers(conv_kind, operand_buffers, result_buffer));
-  input_data = buffers.input_data;
-  filter_data = buffers.filter_data;
-  output_data = buffers.output_data;
-
-  float side_input_scale = 0.0f;
-  bool side_input_scale_zero = true;
-  if (conv_kind == xla::gpu::CudnnConvKind::kForwardActivation) {
-    bias_data = const_cast<void*>(operand_buffers[2].opaque());
-    if (operand_buffers.size() >= 4) {
-      side_input_data = const_cast<void*>(operand_buffers[3].opaque());
-      side_input_scale = config.fusion->side_input_scale;
-      side_input_scale_zero = (fabs(side_input_scale - 0.0f) < 1e-6);
-    }
-  }
-
   // TODO(intel-tf): depthwise-conv
   const int64_t group_count = config.conv_desc.group_count();
   const bool is_group_conv = group_count > 1;
   const int64_t output_channels = config.output_descriptor.feature_map_count();
+  const bool is_conv3d = (config.conv_desc.ndims() == 3);
 
-  absl::Span<const int64_t> padding_dimensions = config.conv_desc.padding();
-  absl::Span<const int64_t> stride_dimensions = config.conv_desc.strides();
-  absl::Span<const int64_t> dilations_dimensions = config.conv_desc.dilations();
+  const absl::Span<const int64_t> padding_dimensions =
+      config.conv_desc.padding();
+  const absl::Span<const int64_t> stride_dimensions =
+      config.conv_desc.strides();
+  const absl::Span<const int64_t> dilations_dimensions =
+      config.conv_desc.dilations();
 
-  bool is_conv3d = (config.conv_desc.ndims() == 3);
+  std::vector<int64_t> src_full =
+      config.input_descriptor.full_dims(DataLayout::kBatchDepthYX);
+  std::vector<int64_t> dst_full =
+      config.output_descriptor.full_dims(DataLayout::kBatchDepthYX);
+  dnnl::memory::dims src_dims(src_full.begin(), src_full.end());
+  dnnl::memory::dims dst_dims(dst_full.begin(), dst_full.end());
+  dnnl::memory::dims filter_dims =
+      ToOneDnnFilterDims(config.filter_descriptor, group_count);
+  dnnl::memory::dims bias_dims = {output_channels};
+  dnnl::memory::dims stride_dims(stride_dimensions.begin(),
+                                 stride_dimensions.end());
+  dnnl::memory::dims padding_dims_l(padding_dimensions.begin(),
+                                    padding_dimensions.end());
+  dnnl::memory::dims padding_dims_r = padding_dims_l;
+  dnnl::memory::dims dilation_dims(dilations_dimensions.size());
+  std::transform(dilations_dimensions.begin(), dilations_dimensions.end(),
+                 dilation_dims.begin(), [](int64_t d) { return d - 1; });
+
+  dnnl::memory::format_tag src_fmt, weight_fmt, dst_fmt;
+  ABSL_ASSIGN_OR_RETURN(src_fmt, ToOneDnnDataFormatTag(input_dl, is_conv3d));
+  ABSL_ASSIGN_OR_RETURN(
+      weight_fmt, ToOneDnnFilterFormatTag(filter_dl, is_conv3d, is_group_conv));
+  ABSL_ASSIGN_OR_RETURN(dst_fmt, ToOneDnnDataFormatTag(output_dl, is_conv3d));
+  ABSL_ASSIGN_OR_RETURN(dnnl::memory::data_type data_type,
+                   ToOneDnnDataType(input_type));
+  pd.kind = conv_kind;
   try {
-    std::vector<int64_t> src_full =
-        config.input_descriptor.full_dims(DataLayout::kBatchDepthYX);
-    std::vector<int64_t> dst_full =
-        config.output_descriptor.full_dims(DataLayout::kBatchDepthYX);
-    dnnl::memory::dims src_dims(src_full.begin(), src_full.end());
-    dnnl::memory::dims dst_dims(dst_full.begin(), dst_full.end());
-    dnnl::memory::dims filter_dims =
-        ToOneDnnFilterDims(config.filter_descriptor, group_count);
-    dnnl::memory::dims bias_dims = {output_channels};
-    dnnl::memory::dims stride_dims(stride_dimensions.begin(),
-                                   stride_dimensions.end());
-    dnnl::memory::dims padding_dims_l(padding_dimensions.begin(),
-                                      padding_dimensions.end());
-    dnnl::memory::dims padding_dims_r = padding_dims_l;
-    dnnl::memory::dims dilation_dims(dilations_dimensions.size());
-    std::transform(dilations_dimensions.begin(), dilations_dimensions.end(),
-                   dilation_dims.begin(), [](int64_t d) { return d - 1; });
-    dnnl::memory::format_tag src_fmt, weight_fmt, dst_fmt;
-    ABSL_ASSIGN_OR_RETURN(src_fmt, ToOneDnnDataFormatTag(input_dl, is_conv3d));
-    ABSL_ASSIGN_OR_RETURN(weight_fmt, ToOneDnnFilterFormatTag(filter_dl, is_conv3d,
-                                                         is_group_conv));
-    ABSL_ASSIGN_OR_RETURN(dst_fmt, ToOneDnnDataFormatTag(output_dl, is_conv3d));
-    ABSL_ASSIGN_OR_RETURN(dnnl::memory::data_type data_type,
-                     ToOneDnnDataType(input_type));
-
-    dnnl::memory::desc src_md =
-        dnnl::memory::desc({src_dims}, data_type, src_fmt);
-    dnnl::memory::desc filter_md =
-        dnnl::memory::desc({filter_dims}, data_type, weight_fmt);
-    dnnl::memory::desc dst_md =
-        dnnl::memory::desc({dst_dims}, data_type, dst_fmt);
+    pd.engine = FindOrCreateEngine(sycl_queue);
+    pd.src_md = dnnl::memory::desc({src_dims}, data_type, src_fmt);
+    pd.dst_md = dnnl::memory::desc({dst_dims}, data_type, dst_fmt);
+    pd.filter_md = dnnl::memory::desc({filter_dims}, data_type, weight_fmt);
 
+    // ONEDNN_PLAIN_WEIGHT forces the primitive descriptor to keep the
+    // caller's filter layout; otherwise let oneDNN pick an internal
+    // pre-packed layout via format_tag::any.
     bool use_plain_weight = false;
     ABSL_RETURN_IF_ERROR(tsl::ReadBoolFromEnvVar("ONEDNN_PLAIN_WEIGHT", false,
                                             &use_plain_weight));
-    dnnl::memory::desc filter_md_prefer = dnnl::memory::desc(
-        {filter_dims}, data_type, dnnl::memory::format_tag::any);
-    if (use_plain_weight) {
-      filter_md_prefer =
-          dnnl::memory::desc({filter_dims}, data_type, weight_fmt);
-    }
+    dnnl::memory::desc filter_md_prefer =
+        use_plain_weight
+            ? dnnl::memory::desc({filter_dims}, data_type, weight_fmt)
+            : dnnl::memory::desc({filter_dims}, data_type,
+                                 dnnl::memory::format_tag::any);
 
-    dnnl::memory src_memory =
-        CreateDnnlMemory(src_md, onednn_conv_primitive.engine, input_data);
-    dnnl::memory filter_memory =
-        CreateDnnlMemory(filter_md, onednn_conv_primitive.engine, filter_data);
-    dnnl::memory dst_memory =
-        CreateDnnlMemory(dst_md, onednn_conv_primitive.engine, output_data);
-
-    // oneDNN's `sum` post-op computes `dst = conv + beta * dst`, reading the
-    // current dst buffer. The side input lives in a separate operand, so we
-    // create a reorder that copies it into dst before each conv execute.
-    dnnl::memory side_input_memory;
-    if (side_input_data && !side_input_scale_zero) {
-      side_input_memory = CreateDnnlMemory(dst_md, onednn_conv_primitive.engine,
-                                           side_input_data);
-      onednn_conv_primitive.side_input_reorder =
-          CreateReorderOp(side_input_memory, dst_memory);
-    }
+    const bool has_fusion = config.fusion.has_value();
+    const double side_input_scale =
+        has_fusion ? config.fusion->side_input_scale : 0.0;
+    pd.has_side_input_sum = has_fusion && std::fabs(side_input_scale) > 1e-6;
 
     // Fused forward conv computes `out = activation(alpha * conv(x, w) +
     // beta * side + bias)`. How each term is expressed depends on `alpha`:
@@ -376,160 +331,210 @@ absl::StatusOr<OneDnnConvPrimitive> CreateOneDnnConvPrimitive(
     //               (eltwise_linear alpha, sum beta, binary_add bias,
     //                activation).
     //
-    // Post-ops are applied in append order, and the terms below are appended
-    // to `po` accordingly.
+    // Post-ops are applied in append order.
     dnnl::post_ops po;
-    dnnl::primitive_attr post_ops_attr;
     if (!alpha_is_one) {
       po.append_eltwise(dnnl::algorithm::eltwise_linear, alpha, 0);
     }
-    if (side_input_data && !side_input_scale_zero) {
+    if (pd.has_side_input_sum) {
       po.append_sum(side_input_scale);
     }
-    // Post-op bias goes only on the forward path; staged here, applied below.
-    dnnl::memory post_op_bias_memory;
-    int post_op_bias_arg_key = 0;
-    bool has_post_op_bias = false;
-    // Bias can be fused directly into the primitive descriptor (via
-    // bias_md) only when alpha == 1. Otherwise, it's applied as
-    // a post-op.
-    if (!alpha_is_one && bias_data) {
+    if (has_fusion && !alpha_is_one) {
       dnnl::memory::dims bias_post_dims(dst_dims.size(), 1);
-      bias_post_dims[1] =
-          bias_dims[0];  // Logical dimension 1 is always channels (C) in oneDNN
-      auto bias_post_md =
-          dnnl::memory::desc(bias_post_dims, data_type, dst_fmt);
+      // Logical dimension 1 is always channels (C) in oneDNN.
+      bias_post_dims[1] = bias_dims[0];
+      dnnl::memory::desc bias_post_md(bias_post_dims, data_type, dst_fmt);
       po.append_binary(dnnl::algorithm::binary_add, bias_post_md);
-      post_op_bias_memory = CreateDnnlMemory(
-          bias_post_md, onednn_conv_primitive.engine, bias_data);
-      post_op_bias_arg_key =
-          DNNL_ARG_ATTR_MULTIPLE_POST_OP(po.len() - 1) | DNNL_ARG_SRC_1;
-      has_post_op_bias = true;
+      pd.bias = OneDnnConvPrimitiveDesc::Bias{
+          bias_post_md,
+          DNNL_ARG_ATTR_MULTIPLE_POST_OP(po.len() - 1) | DNNL_ARG_SRC_1};
     }
-    if (conv_kind == xla::gpu::CudnnConvKind::kForwardActivation) {
+    if (has_fusion) {
       switch (config.fusion->mode) {
-        case stream_executor::dnn::kSigmoid:
+        case dnn::kSigmoid:
           po.append_eltwise(dnnl::algorithm::eltwise_logistic, 1, 0);
           break;
-        case stream_executor::dnn::kRelu:
+        case dnn::kRelu:
           po.append_eltwise(dnnl::algorithm::eltwise_relu, 0, 0);
           break;
-        case stream_executor::dnn::kRelu6:
+        case dnn::kRelu6:
           po.append_eltwise(dnnl::algorithm::eltwise_clip_v2, 0, 6);
           break;
-        case stream_executor::dnn::kTanh:
+        case dnn::kTanh:
           po.append_eltwise(dnnl::algorithm::eltwise_tanh, 0, 0);
           break;
-        case stream_executor::dnn::kElu:
+        case dnn::kElu:
           po.append_eltwise(dnnl::algorithm::eltwise_elu, 1, 0);
           break;
-        case stream_executor::dnn::kLeakyRelu:
+        case dnn::kLeakyRelu:
           po.append_eltwise(dnnl::algorithm::eltwise_relu,
                             config.fusion->leakyrelu_alpha, 0);
           break;
-        case stream_executor::dnn::kNone:
+        case dnn::kNone:
           break;
         default:
           return absl::InvalidArgumentError("Unsupported Activation mode");
       }
     }
+    dnnl::primitive_attr post_ops_attr;
     post_ops_attr.set_post_ops(po);
     post_ops_attr.set_scratchpad_mode(dnnl::scratchpad_mode::user);
-    // Set fp32 mode.
-    dnnl::fpmath_mode fp32_math_mode = GetFP32MathMode();
-    if (input_type == xla::F32) {
-      post_ops_attr.set_fpmath_mode(fp32_math_mode);
+    if (input_type == dnn::DataType::kFloat) {
+      post_ops_attr.set_fpmath_mode(GetFP32MathMode());
     }
-    switch (conv_kind) {
-      case xla::gpu::CudnnConvKind::kForward:
-      case xla::gpu::CudnnConvKind::kForwardActivation: {
-        // Bias can be fused directly into the primitive descriptor (via
-        // bias_md) only when alpha == 1.
-        std::optional<dnnl::memory::desc> bias_md;
-        if (bias_data != nullptr && alpha_is_one) {
-          bias_md = dnnl::memory::desc(bias_dims, data_type,
-                                       dnnl::memory::format_tag::x);
-        }
-        // Create the convolution forward primitive descriptor with the
-        // appropriate bias memory descriptor.
-        ConvFwdPd fwd_pd = CreateConvFwdPd(
-            onednn_conv_primitive.engine, src_md, filter_md_prefer, bias_md,
-            dst_md, stride_dims, dilation_dims, padding_dims_l, padding_dims_r,
-            post_ops_attr);
 
-        // Build a convolution primitive
+    // Backward primitives don't take post-ops; their pds only need
+    // scratchpad + fpmath, plus a forward hint built with the same attrs.
+    dnnl::primitive_attr bwd_attr;
+    bwd_attr.set_scratchpad_mode(dnnl::scratchpad_mode::user);
+    if (input_type == dnn::DataType::kFloat) {
+      bwd_attr.set_fpmath_mode(GetFP32MathMode());
+    }
+
+    switch (conv_kind) {
+      case dnn::ConvolutionKind::FORWARD:
+      case dnn::ConvolutionKind::FORWARD_BIAS_ACTIVATION: {
+        if (has_fusion && alpha_is_one) {
+          pd.bias = OneDnnConvPrimitiveDesc::Bias{
+              dnnl::memory::desc(bias_dims, data_type,
+                                 dnnl::memory::format_tag::x),
+              DNNL_ARG_BIAS};
+        }
+        std::optional<dnnl::memory::desc> primitive_bias_md;
+        if (pd.bias.has_value() && pd.bias->arg_key == DNNL_ARG_BIAS) {
+          primitive_bias_md = pd.bias->md;
+        }
+        pd.conv_pd = CreateConvFwdPd(pd.engine, pd.src_md, filter_md_prefer,
+                                     primitive_bias_md, pd.dst_md, stride_dims,
+                                     dilation_dims, padding_dims_l,
+                                     padding_dims_r, post_ops_attr);
+        break;
+      }
+      case dnn::ConvolutionKind::BACKWARD_DATA: {
+        ConvFwdPd fwd_hint = CreateConvFwdPd(
+            pd.engine, pd.src_md, filter_md_prefer, /*bias_md=*/std::nullopt,
+            pd.dst_md, stride_dims, dilation_dims, padding_dims_l,
+            padding_dims_r, bwd_attr);
+        pd.conv_pd = ConvBwdInputPd(
+            pd.engine, dnnl::algorithm::convolution_direct, pd.src_md,
+            filter_md_prefer, pd.dst_md, stride_dims, dilation_dims,
+            padding_dims_l, padding_dims_r, fwd_hint, bwd_attr);
+        break;
+      }
+      case dnn::ConvolutionKind::BACKWARD_FILTER: {
+        ConvFwdPd fwd_hint = CreateConvFwdPd(
+            pd.engine, pd.src_md, filter_md_prefer, /*bias_md=*/std::nullopt,
+            pd.dst_md, stride_dims, dilation_dims, padding_dims_l,
+            padding_dims_r, bwd_attr);
+        pd.conv_pd = ConvBwdFilterPd(
+            pd.engine, dnnl::algorithm::convolution_direct, pd.src_md,
+            filter_md_prefer, pd.dst_md, stride_dims, dilation_dims,
+            padding_dims_l, padding_dims_r, fwd_hint, bwd_attr);
+        break;
+      }
+      default:
+        return absl::InvalidArgumentError("Unknown convolution kind");
+    }
+  } catch (const dnnl::error& e) {
+    return absl::InternalError(absl::StrCat("OneDNN Conv error: ", e.message));
+  }
+  return pd;
+}
+
+absl::StatusOr<OneDnnConvPrimitive> CreateOneDnnConvPrimitive(
+    const OneDnnConvPrimitiveDesc& pd,
+    absl::Span<const DeviceAddressBase> operand_buffers,
+    DeviceAddressBase result_buffer, Stream* stream) {
+  OneDnnConvPrimitive onednn_conv_primitive;
+  ::sycl::queue* sycl_queue =
+      absl::bit_cast<::sycl::queue*>(stream->platform_specific_handle().stream);
+  onednn_conv_primitive.engine = pd.engine;
+
+  const dnn::ConvolutionKind conv_kind = pd.kind;
+  ABSL_ASSIGN_OR_RETURN(
+      ConvBufferPointers buffers,
+      GetConvBufferPointers(conv_kind, operand_buffers, result_buffer));
+
+  void* bias_data = nullptr;
+  void* side_input_data = nullptr;
+  if (pd.bias.has_value() && operand_buffers.size() >= 3) {
+    bias_data = const_cast<void*>(operand_buffers[2].opaque());
+  }
+  if (pd.has_side_input_sum && operand_buffers.size() >= 4) {
+    side_input_data = const_cast<void*>(operand_buffers[3].opaque());
+  }
+  if (pd.has_side_input_sum && side_input_data == nullptr) {
+    return absl::InvalidArgumentError(
+        "Fused convolution requires a side input buffer");
+  }
+  if (pd.bias.has_value() && bias_data == nullptr) {
+    return absl::InvalidArgumentError(
+        "Fused convolution requires a bias buffer");
+  }
+  try {
+    onednn_conv_primitive.stream = dnnl::sycl_interop::make_stream(
+        onednn_conv_primitive.engine, *sycl_queue);
+    dnnl::memory src_memory = CreateDnnlMemory(
+        pd.src_md, onednn_conv_primitive.engine, buffers.input_data);
+    dnnl::memory filter_memory = CreateDnnlMemory(
+        pd.filter_md, onednn_conv_primitive.engine, buffers.filter_data);
+    dnnl::memory dst_memory = CreateDnnlMemory(
+        pd.dst_md, onednn_conv_primitive.engine, buffers.output_data);
+
+    // oneDNN's `sum` post-op computes `dst = conv + beta * dst`, reading the
+    // current dst buffer. When the side input uses a separate buffer, copy it
+    // into dst before each conv execute.
+    dnnl::memory side_input_memory;
+    if (pd.has_side_input_sum && side_input_data != nullptr &&
+        side_input_data != buffers.output_data) {
+      side_input_memory = CreateDnnlMemory(
+          pd.dst_md, onednn_conv_primitive.engine, side_input_data);
+      onednn_conv_primitive.side_input_reorder =
+          CreateReorderOp(side_input_memory, dst_memory);
+    }
+
+    switch (conv_kind) {
+      case dnn::ConvolutionKind::FORWARD:
+      case dnn::ConvolutionKind::FORWARD_BIAS_ACTIVATION: {
+        const ConvFwdPd& fwd_pd = std::get<ConvFwdPd>(pd.conv_pd);
         ConvFwd fwd;
         ABSL_ASSIGN_OR_RETURN(
-            fwd,
-            (BuildConvOp<ConvFwdPd, dnnl::convolution_forward, ConvFwd>(
-                fwd_pd, std::move(src_memory), DNNL_ARG_SRC,
-                std::move(filter_memory), DNNL_ARG_WEIGHTS,
-                std::move(dst_memory), DNNL_ARG_DST, filter_md,
-                fwd_pd.weights_desc(),
-                /*prepack_filter=*/true, onednn_conv_primitive.engine,
-                scratch_allocator, &onednn_conv_primitive.filter_reorder)));
+            fwd, (BuildConvOp<ConvFwdPd, dnnl::convolution_forward, ConvFwd>(
+                     fwd_pd, std::move(src_memory), DNNL_ARG_SRC,
+                     std::move(filter_memory), DNNL_ARG_WEIGHTS,
+                     std::move(dst_memory), DNNL_ARG_DST, pd.filter_md,
+                     fwd_pd.weights_desc(),
+                     /*prepack_filter=*/true, onednn_conv_primitive.engine,
+                     &onednn_conv_primitive.filter_reorder)));
 
         fwd.side_input = std::move(side_input_memory);
-        if (bias_md.has_value()) {
-          fwd.bias = CreateDnnlMemory(bias_md.value(),
-                                      onednn_conv_primitive.engine, bias_data);
-          fwd.args.insert({DNNL_ARG_BIAS, fwd.bias});
-        }
-        if (has_post_op_bias) {
-          fwd.bias = std::move(post_op_bias_memory);
-          fwd.args.insert({post_op_bias_arg_key, fwd.bias});
+        if (pd.bias.has_value() && bias_data != nullptr) {
+          fwd.bias = CreateDnnlMemory(pd.bias->md, onednn_conv_primitive.engine,
+                                      bias_data);
+          fwd.args.insert({pd.bias->arg_key, fwd.bias});
         }
         onednn_conv_primitive.op = std::move(fwd);
         break;
       }
-      case xla::gpu::CudnnConvKind::kBackwardInput: {
-        // Create a forward convolution primitive descriptor for the backward
-        // input convolution.
-        dnnl::primitive_attr attr;
-        if (input_type == xla::F32) {
-          attr.set_fpmath_mode(fp32_math_mode);
-        }
-        attr.set_scratchpad_mode(dnnl::scratchpad_mode::user);
-        ConvFwdPd fwd_pd = CreateConvFwdPd(
-            onednn_conv_primitive.engine, src_md, filter_md_prefer,
-            /*bias_md=*/std::nullopt, dst_md, stride_dims, dilation_dims,
-            padding_dims_l, padding_dims_r, attr);
-        ConvBwdInputPd bwd_input_pd = ConvBwdInputPd(
-            onednn_conv_primitive.engine, dnnl::algorithm::convolution_direct,
-            src_md, filter_md_prefer, dst_md, stride_dims, dilation_dims,
-            padding_dims_l, padding_dims_r, fwd_pd, attr);
-
+      case dnn::ConvolutionKind::BACKWARD_DATA: {
+        const ConvBwdInputPd& bwd_input_pd =
+            std::get<ConvBwdInputPd>(pd.conv_pd);
         ConvBwdData bwd;
         ABSL_ASSIGN_OR_RETURN(
-            bwd,
-            (BuildConvOp<ConvBwdInputPd, dnnl::convolution_backward_data,
-                         ConvBwdData>(
-                bwd_input_pd, std::move(src_memory), DNNL_ARG_DIFF_SRC,
-                std::move(filter_memory), DNNL_ARG_WEIGHTS,
-                std::move(dst_memory), DNNL_ARG_DIFF_DST, filter_md,
-                bwd_input_pd.weights_desc(),
-                /*prepack_filter=*/true, onednn_conv_primitive.engine,
-                scratch_allocator, &onednn_conv_primitive.filter_reorder)));
-
+            bwd, (BuildConvOp<ConvBwdInputPd, dnnl::convolution_backward_data,
+                              ConvBwdData>(
+                     bwd_input_pd, std::move(src_memory), DNNL_ARG_DIFF_SRC,
+                     std::move(filter_memory), DNNL_ARG_WEIGHTS,
+                     std::move(dst_memory), DNNL_ARG_DIFF_DST, pd.filter_md,
+                     bwd_input_pd.weights_desc(),
+                     /*prepack_filter=*/true, onednn_conv_primitive.engine,
+                     &onednn_conv_primitive.filter_reorder)));
         onednn_conv_primitive.op = std::move(bwd);
         break;
       }
-      case xla::gpu::CudnnConvKind::kBackwardFilter: {
-        // Create a forward convolution primitive descriptor for the backward
-        // weights convolution.
-        dnnl::primitive_attr attr;
-        if (input_type == xla::F32) {
-          attr.set_fpmath_mode(fp32_math_mode);
-        }
-        attr.set_scratchpad_mode(dnnl::scratchpad_mode::user);
-        ConvFwdPd fwd_pd = CreateConvFwdPd(
-            onednn_conv_primitive.engine, src_md, filter_md_prefer,
-            /*bias_md=*/std::nullopt, dst_md, stride_dims, dilation_dims,
-            padding_dims_l, padding_dims_r, attr);
-        ConvBwdFilterPd bwd_filter_pd = ConvBwdFilterPd(
-            onednn_conv_primitive.engine, dnnl::algorithm::convolution_direct,
-            src_md, filter_md_prefer, dst_md, stride_dims, dilation_dims,
-            padding_dims_l, padding_dims_r, fwd_pd, attr);
+      case dnn::ConvolutionKind::BACKWARD_FILTER: {
+        const ConvBwdFilterPd& bwd_filter_pd =
+            std::get<ConvBwdFilterPd>(pd.conv_pd);
         ConvBwdWeights bwd;
         ABSL_ASSIGN_OR_RETURN(
             bwd,
@@ -537,11 +542,10 @@ absl::StatusOr<OneDnnConvPrimitive> CreateOneDnnConvPrimitive(
                          ConvBwdWeights>(
                 bwd_filter_pd, std::move(src_memory), DNNL_ARG_SRC,
                 std::move(filter_memory), DNNL_ARG_DIFF_WEIGHTS,
-                std::move(dst_memory), DNNL_ARG_DIFF_DST, filter_md,
+                std::move(dst_memory), DNNL_ARG_DIFF_DST, pd.filter_md,
                 bwd_filter_pd.diff_weights_desc(),
                 /*prepack_filter=*/false, onednn_conv_primitive.engine,
-                scratch_allocator, &onednn_conv_primitive.filter_reorder)));
-
+                &onednn_conv_primitive.filter_reorder)));
         onednn_conv_primitive.op = std::move(bwd);
         break;
       }
@@ -554,7 +558,7 @@ absl::StatusOr<OneDnnConvPrimitive> CreateOneDnnConvPrimitive(
   return onednn_conv_primitive;
 }
 
-absl::Status DoOnednnConv(const OneDnnConvPrimitive& onednn_primitive) {
+absl::Status DoOneDnnConv(const OneDnnConvPrimitive& onednn_primitive) {
   try {
     auto execute_reorder = [&](const std::optional<ReorderOp>& reorder) {
       if (reorder) {
diff --git a/third_party/xla/xla/stream_executor/sycl/sycl_conv_utils.h b/third_party/xla/xla/stream_executor/sycl/sycl_conv_utils.h
index 6c58b303e0a..189c8dfb491 100644
--- a/third_party/xla/xla/stream_executor/sycl/sycl_conv_utils.h
+++ b/third_party/xla/xla/stream_executor/sycl/sycl_conv_utils.h
@@ -25,9 +25,8 @@ limitations under the License.
 #include "absl/types/span.h"
 #include "dnnl.hpp"
 #include "dnnl_sycl.hpp"
-#include "xla/service/gpu/gpu_conv_runner.h"
 #include "xla/stream_executor/device_address.h"
-#include "xla/stream_executor/scratch_allocator.h"
+#include "xla/stream_executor/dnn.h"
 #include "xla/stream_executor/stream.h"
 #include "xla/stream_executor/sycl/onednn_util.h"
 
@@ -80,13 +79,61 @@ struct OneDnnConvPrimitive {
   std::optional<ReorderOp> side_input_reorder;
 };
 
-absl::StatusOr<OneDnnConvPrimitive> CreateOneDnnConvPrimitive(
-    const xla::gpu::GpuConvConfig& config,
-    absl::Span<const DeviceAddressBase> operand_buffers,
-    DeviceAddressBase result_buffer, Stream* stream,
-    ScratchAllocator* scratch_allocator);
+// Everything the oneDNN convolution execution path needs that is derivable
+// from an OneDnnConvConfig alone (no device buffers).
+struct OneDnnConvPrimitiveDesc {
+  dnn::ConvolutionKind kind;
+  dnnl::engine engine;
+  dnnl::memory::desc src_md;
+  dnnl::memory::desc filter_md;
+  dnnl::memory::desc dst_md;
+  // Bias for the fused forward path; absent otherwise. `arg_key` is the
+  // execution-args key to bind the bias buffer under:
+  //   - alpha == 1: bias is part of the conv primitive; key is DNNL_ARG_BIAS.
+  //   - alpha != 1: bias is a binary_add post-op; key is
+  //     DNNL_ARG_ATTR_MULTIPLE_POST_OP(idx) | DNNL_ARG_SRC_1.
+  struct Bias {
+    dnnl::memory::desc md;
+    int arg_key;
+  };
+  std::optional<Bias> bias;
+  // True when a `sum` post-op is present; the caller must copy side_input
+  // into dst before executing the primitive.
+  bool has_side_input_sum = false;
+  std::variant<dnnl::convolution_forward::primitive_desc,
+               dnnl::convolution_backward_data::primitive_desc,
+               dnnl::convolution_backward_weights::primitive_desc>
+      conv_pd;
+};
 
-absl::Status DoOnednnConv(const OneDnnConvPrimitive& onednn_primitive);
+// Static convolution properties used by the oneDNN backend.
+struct OneDnnConvConfig {
+  struct Fusion {
+    dnn::ActivationMode mode;
+    double side_input_scale;
+    double leakyrelu_alpha;
+  };
+
+  dnn::ConvolutionKind kind;
+  dnn::DataType input_type;
+  dnn::DataType output_type;
+  dnn::BatchDescriptor input_descriptor;
+  dnn::FilterDescriptor filter_descriptor;
+  dnn::BatchDescriptor output_descriptor;
+  dnn::ConvolutionDescriptor conv_desc;
+  double conv_result_scale = 1.0;
+  std::optional<Fusion> fusion;
+};
+
+absl::StatusOr<OneDnnConvPrimitive> CreateOneDnnConvPrimitive(
+    const OneDnnConvPrimitiveDesc& pd,
+    absl::Span<const DeviceAddressBase> operand_buffers,
+    DeviceAddressBase result_buffer, Stream* stream);
+
+absl::StatusOr<OneDnnConvPrimitiveDesc> CreateOneDnnConvPrimitiveDesc(
+    const OneDnnConvConfig& config, Stream* stream);
+
+absl::Status DoOneDnnConv(const OneDnnConvPrimitive& onednn_primitive);
 
 }  // namespace sycl
 }  // namespace stream_executor
diff --git a/third_party/xla/xla/stream_executor/sycl/sycl_dnn.cc b/third_party/xla/xla/stream_executor/sycl/sycl_dnn.cc
index 7d5ec555a68..d38d65389f7 100644
--- a/third_party/xla/xla/stream_executor/sycl/sycl_dnn.cc
+++ b/third_party/xla/xla/stream_executor/sycl/sycl_dnn.cc
@@ -21,8 +21,7 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/types/span.h"
 #include "dnnl.hpp"
-#include "xla/service/gpu/gpu_conv_runner.h"
-#include "xla/stream_executor/device_memory.h"
+#include "xla/stream_executor/device_address.h"
 #include "xla/stream_executor/dnn.h"
 #include "xla/stream_executor/platform/initialize.h"
 #include "xla/stream_executor/plugin_registry.h"
@@ -50,16 +49,115 @@ absl::StatusOr<dnn::VersionInfo> OnednnSupport::GetVersion() {
   return GetOnednnVersion();
 }
 
-absl::Status OnednnSupport::DoConvolveWithGpuConfig(
-    Stream* stream, const xla::gpu::GpuConvConfig& config,
-    absl::Span<const DeviceMemoryBase> operand_se_buffers,
-    DeviceMemoryBase result_se_buffer, ScratchAllocator* scratch_allocator) {
+namespace {
+class OnednnConvRunner : public dnn::ConvRunner {
+ public:
+  OnednnConvRunner(OneDnnConvPrimitiveDesc onednn_conv_primitive_desc,
+                   size_t workspace_size)
+      : onednn_conv_primitive_desc_(std::move(onednn_conv_primitive_desc)),
+        workspace_size_(workspace_size) {}
+  std::string ToString() const override { return "OnednnConvRunner"; }
+
+  size_t GetWorkspaceSize() const override { return workspace_size_; }
+
+  absl::StatusOr<dnn::AlgorithmDesc> ToAlgorithmDesc() const override {
+    return dnn::AlgorithmDesc(-1, false, workspace_size_);
+  }
+
+  absl::Status operator()(Stream* stream,
+                          dnn::ProfileResult* output_profile_result,
+                          DeviceAddressBase scratch_memory,
+                          DeviceAddressBase input_data,
+                          DeviceAddressBase filter_data,
+                          DeviceAddressBase output_data) const override {
+    // Implemented as part of a follow-up PR.
+    return absl::UnimplementedError(
+        "OnednnConvRunner operator() is not implemented for SYCL");
+  }
+
+ private:
+  OneDnnConvPrimitiveDesc onednn_conv_primitive_desc_;
+  size_t workspace_size_ = 0;
+};
+
+class OnednnFusedConvRunner : public dnn::FusedConvRunner {
+ public:
+  OnednnFusedConvRunner(OneDnnConvPrimitiveDesc onednn_conv_primitive_desc,
+                        size_t workspace_size)
+      : onednn_conv_primitive_desc_(std::move(onednn_conv_primitive_desc)),
+        workspace_size_(workspace_size) {}
+  std::string ToString() const override { return "OnednnFusedConvRunner"; }
+
+  size_t GetWorkspaceSize() const override { return workspace_size_; }
+
+  absl::StatusOr<dnn::AlgorithmDesc> ToAlgorithmDesc() const override {
+    return dnn::AlgorithmDesc(-1, false, workspace_size_);
+  }
+
+  absl::Status operator()(Stream* stream,
+                          dnn::ProfileResult* output_profile_result,
+                          DeviceAddressBase scratch_memory,
+                          DeviceAddressBase input_data,
+                          DeviceAddressBase filter_data,
+                          DeviceAddressBase side_input_data,
+                          DeviceAddressBase bias_data,
+                          DeviceAddressBase output_data) const override {
+    // Implemented as part of a follow-up PR.
+    return absl::UnimplementedError(
+        "OnednnFusedConvRunner operator() is not implemented for SYCL");
+  }
+
+ private:
+  OneDnnConvPrimitiveDesc onednn_conv_primitive_desc_;
+  size_t workspace_size_ = 0;
+};
+}  // namespace
+
+absl::StatusOr<std::unique_ptr<const dnn::ConvRunner>>
+OnednnSupport::ConvolveRunnerFromDesc(
+    Stream* stream, const dnn::AlgorithmDesc& algorithm_desc,
+    dnn::ConvolutionKind kind, dnn::DataType input_type,
+    dnn::DataType output_type, const dnn::BatchDescriptor& input_descriptor,
+    const dnn::FilterDescriptor& filter_descriptor,
+    const dnn::BatchDescriptor& output_descriptor,
+    const dnn::ConvolutionDescriptor& convolution_descriptor) {
   ABSL_ASSIGN_OR_RETURN(
-      auto onednn_primitive,
-      CreateOneDnnConvPrimitive(config, operand_se_buffers, result_se_buffer,
-                                stream, scratch_allocator));
-  ABSL_RETURN_IF_ERROR(DoOnednnConv(onednn_primitive));
-  return absl::OkStatus();
+      OneDnnConvPrimitiveDesc primitive_desc,
+      CreateOneDnnConvPrimitiveDesc(
+          OneDnnConvConfig{kind, input_type, output_type, input_descriptor,
+                           filter_descriptor, output_descriptor,
+                           convolution_descriptor},
+          stream));
+  size_t workspace_size = 0;
+  return std::make_unique<OnednnConvRunner>(std::move(primitive_desc),
+                                            workspace_size);
+}
+
+absl::StatusOr<std::unique_ptr<const dnn::FusedConvRunner>>
+OnednnSupport::FusedConvolveRunnerFromDesc(
+    Stream* stream, const dnn::AlgorithmDesc& algorithm_desc,
+    dnn::ConvolutionKind kind, dnn::DataType element_type,
+    dnn::DataType bias_type, dnn::DataType output_type, double conv_scale,
+    double side_input_scale, double leakyrelu_alpha,
+    const dnn::BatchDescriptor& input_descriptor,
+    const dnn::FilterDescriptor& filter_descriptor,
+    const dnn::BatchDescriptor& bias_descriptor,
+    const dnn::BatchDescriptor& output_descriptor,
+    const dnn::ConvolutionDescriptor& convolution_descriptor,
+    dnn::ActivationMode activation_mode) {
+  ABSL_ASSIGN_OR_RETURN(
+      OneDnnConvPrimitiveDesc primitive_desc,
+      CreateOneDnnConvPrimitiveDesc(
+          OneDnnConvConfig{
+              kind, element_type, output_type, input_descriptor,
+              filter_descriptor, output_descriptor, convolution_descriptor,
+              conv_scale,
+              OneDnnConvConfig::Fusion{activation_mode, side_input_scale,
+                                       leakyrelu_alpha}},
+          stream));
+  size_t workspace_size = 0;
+  return std::make_unique<OnednnFusedConvRunner>(std::move(primitive_desc),
+                                                 workspace_size);
 }
 
 void initialize_onednn() {
diff --git a/third_party/xla/xla/stream_executor/sycl/sycl_dnn.h b/third_party/xla/xla/stream_executor/sycl/sycl_dnn.h
index de5b925364d..3be02aec652 100644
--- a/third_party/xla/xla/stream_executor/sycl/sycl_dnn.h
+++ b/third_party/xla/xla/stream_executor/sycl/sycl_dnn.h
@@ -23,8 +23,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/types/span.h"
-#include "xla/service/gpu/gpu_conv_runner.h"
-#include "xla/stream_executor/device_memory.h"
+#include "xla/stream_executor/device_address.h"
 #include "xla/stream_executor/dnn.h"
 #include "xla/stream_executor/scratch_allocator.h"
 #include "xla/stream_executor/stream.h"
@@ -46,17 +45,12 @@ class OnednnSupport : public dnn::DnnSupport {
   // Static helper to get oneDNN version without requiring an instance.
   static absl::StatusOr<stream_executor::dnn::VersionInfo> GetOnednnVersion();
 
-  absl::Status DoConvolveWithGpuConfig(
-      Stream* stream, const xla::gpu::GpuConvConfig& config,
-      absl::Span<const DeviceMemoryBase> operand_se_buffers,
-      DeviceMemoryBase result_se_buffer, ScratchAllocator* scratch_allocator);
-
   absl::Status DoPoolForward(dnn::DataType element_type, Stream* stream,
                              const dnn::PoolingDescriptor& pooling_dimensions,
                              const dnn::BatchDescriptor& input_dimensions,
-                             DeviceMemoryBase input_data,
+                             DeviceAddressBase input_data,
                              const dnn::BatchDescriptor& output_dimensions,
-                             DeviceMemoryBase output_data,
+                             DeviceAddressBase output_data,
                              ScratchAllocator* workspace_allocator) override {
     return absl::UnimplementedError(
         "DoPoolForward is not implemented for SYCL");
@@ -65,11 +59,11 @@ class OnednnSupport : public dnn::DnnSupport {
   absl::Status DoPoolBackward(dnn::DataType element_type, Stream* stream,
                               const dnn::PoolingDescriptor& pooling_dimensions,
                               const dnn::BatchDescriptor& input_dimensions,
-                              DeviceMemoryBase input_data,
+                              DeviceAddressBase input_data,
                               const dnn::BatchDescriptor& output_dimensions,
-                              DeviceMemoryBase output_data,
-                              DeviceMemoryBase input_diff_data,
-                              DeviceMemoryBase output_diff_data,
+                              DeviceAddressBase output_data,
+                              DeviceAddressBase input_diff_data,
+                              DeviceAddressBase output_diff_data,
                               ScratchAllocator* workspace_allocator) override {
     return absl::UnimplementedError(
         "DoPoolBackward is not implemented for SYCL");
@@ -81,21 +75,31 @@ class OnednnSupport : public dnn::DnnSupport {
       dnn::DataType output_type, const dnn::BatchDescriptor& input_descriptor,
       const dnn::FilterDescriptor& filter_descriptor,
       const dnn::BatchDescriptor& output_descriptor,
-      const dnn::ConvolutionDescriptor& convolution_descriptor) override {
-    return absl::UnimplementedError(
-        "ConvolveRunnerFromDesc is not implemented for SYCL");
-  }
+      const dnn::ConvolutionDescriptor& convolution_descriptor) override;
+
+  absl::StatusOr<std::unique_ptr<const dnn::FusedConvRunner>>
+  FusedConvolveRunnerFromDesc(
+      Stream* stream, const dnn::AlgorithmDesc& algorithm_desc,
+      dnn::ConvolutionKind kind, dnn::DataType element_type,
+      dnn::DataType bias_type, dnn::DataType output_type, double conv_scale,
+      double side_input_scale, double leakyrelu_alpha,
+      const dnn::BatchDescriptor& input_descriptor,
+      const dnn::FilterDescriptor& filter_descriptor,
+      const dnn::BatchDescriptor& bias_descriptor,
+      const dnn::BatchDescriptor& output_descriptor,
+      const dnn::ConvolutionDescriptor& convolution_descriptor,
+      dnn::ActivationMode activation_mode) override;
 
   absl::Status DoCtcLoss(Stream* stream, dnn::DataType element_type,
                          const dnn::RnnStateTensorDescriptor& probs_desc,
-                         DeviceMemoryBase probs_data,
+                         DeviceAddressBase probs_data,
                          absl::Span<const int> labels_data,
                          absl::Span<const int> labels_lengths_data,
                          absl::Span<const int> input_lengths_data,
-                         DeviceMemoryBase costs_data,
+                         DeviceAddressBase costs_data,
                          const dnn::RnnStateTensorDescriptor& grads_desc,
-                         DeviceMemoryBase grads_data,
-                         DeviceMemory<uint8_t> scratch_memory,
+                         DeviceAddressBase grads_data,
+                         DeviceAddress<uint8_t> scratch_memory,
                          int ctc_loss_algo_id) override {
     return absl::UnimplementedError("DoCtcLoss is not implemented for SYCL");
   }

From 39df24c56d0129a261324d9d880b61f77c45eeec Mon Sep 17 00:00:00 2001
From: Om Thakkar <om.thakkar@intel.com>
Date: Fri, 9 Oct 2026 03:28:24 -0700
Subject: [PATCH 0347/1000] PR #47655: [XLA:CPU[oneDNN] Fall back to
 single-threaded pool when intra-op threadpool is unavailable in oneDNN thunks

Imported from GitHub PR https://github.com/openxla/xla/pull/47655

`OneDnnOpThunk` and `OneDnnFusionThunk` unconditionally dereference `params.intra_op_threadpool`. Fall back to a static, single-threaded `Eigen::ThreadPool`.
Copybara import of the project:

--
20e246f5ee8270b0677ac59a8c01c172b7f62800 by Om Thakkar <om.thakkar@intel.com>:

fall back to static single-threaded threadpool when running without intra-op parallelism

--
e8da7c77418cf696113334ad125520255554a82e by Om Thakkar <om.thakkar@intel.com>:

rename to GetFallbackThreadPoolForOneDnn and preserve the same logging level

--
fe7142a686d41714d4aefbb0dee1899af7b2a371 by Om Thakkar <om.thakkar@intel.com>:

add unit tests

--
0be7df6130bead18385b4257b2e30dea3d7ddafe by Om Thakkar <om.thakkar@intel.com>:

unify verbose logging

Merging this change closes #47655

PiperOrigin-RevId: 996425314
---
 .../xla/xla/backends/cpu/runtime/onednn/BUILD |  3 +-
 .../cpu/runtime/onednn/onednn_fusion_thunk.cc |  3 +-
 .../cpu/runtime/onednn/onednn_op_thunk.cc     | 17 +++----
 .../runtime/onednn/onednn_op_thunk_test.cc    | 47 +++++++++++++++++++
 .../cpu/runtime/onednn/onednn_threadpool.h    |  9 ++++
 .../runtime/onednn/onednn_threadpool_test.cc  | 46 ++++++++++++++++++
 6 files changed, 111 insertions(+), 14 deletions(-)

diff --git a/third_party/xla/xla/backends/cpu/runtime/onednn/BUILD b/third_party/xla/xla/backends/cpu/runtime/onednn/BUILD
index d5d78f25d4f..5792d91948c 100644
--- a/third_party/xla/xla/backends/cpu/runtime/onednn/BUILD
+++ b/third_party/xla/xla/backends/cpu/runtime/onednn/BUILD
@@ -104,6 +104,7 @@ cc_library(
         "//xla/backends/cpu/runtime:work_queue",
         "//xla/tsl/concurrency:async_value",
         "//xla/tsl/mkl:onednn",
+        "//xla/tsl/platform:logging",
         "@eigen_archive//:eigen3",
     ],
 )
@@ -155,7 +156,6 @@ cc_library(
         "//xla/tsl/mkl:onednn",
         "//xla/tsl/platform:logging",
         "//xla/tsl/platform:statusor",
-        "@com_google_absl//absl/base",
         "@com_google_absl//absl/base:dynamic_annotations",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/memory",
@@ -164,7 +164,6 @@ cc_library(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
-        "@eigen_archive//:eigen3",
     ],
 )
 
diff --git a/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_fusion_thunk.cc b/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_fusion_thunk.cc
index 08f6d0b102a..feaf246e945 100644
--- a/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_fusion_thunk.cc
+++ b/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_fusion_thunk.cc
@@ -207,7 +207,8 @@ tsl::AsyncValueRef<OneDnnFusionThunk::ExecuteEvent> OneDnnFusionThunk::Execute(
   }
 
   Eigen::ThreadPoolInterface* thread_pool =
-      params.intra_op_threadpool->getPool();
+      params.intra_op_threadpool ? params.intra_op_threadpool->getPool()
+                                 : GetFallbackThreadPoolForOneDnn();
 
   // Borrow oneDNN runtime from the pool.
   ABSL_ASSIGN_OR_RETURN(auto runtime, onednn_runtime_pool_.GetOrCreate(thread_pool));
diff --git a/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_op_thunk.cc b/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_op_thunk.cc
index 733877dd3ca..4662f34b18d 100644
--- a/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_op_thunk.cc
+++ b/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_op_thunk.cc
@@ -22,7 +22,6 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
-#include "absl/base/call_once.h"
 #include "absl/base/dynamic_annotations.h"
 #include "absl/log/check.h"
 #include "absl/memory/memory.h"
@@ -31,7 +30,6 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
-#include "Eigen/ThreadPool"
 #include "oneapi/dnnl/dnnl_common.hpp"
 #include "oneapi/dnnl/dnnl_threadpool.hpp"
 #include "xla/backends/cpu/runtime/onednn/onednn_threadpool.h"
@@ -100,13 +98,10 @@ OneDnnOpThunk::OneDnnRuntime::Invoke(
   threadpool->set_thread_pool(thread_pool);
 
   // TODO(intel-tf): Add support for more oneDNN operations as needed.
-  static absl::once_flag log_once_flag;
-  absl::call_once(log_once_flag, [&] {
-    VLOG(0) << absl::StreamFormat(
-        "Executing oneDNN thunk with target `%s`: num_args=%d, num_results=%d",
-        target, base_resources->arg_memrefs.size(),
-        base_resources->result_memrefs.size());
-  });
+  VLOG_FIRST_N(0, 1) << absl::StreamFormat(
+      "Executing oneDNN thunk with target `%s`: num_args=%d, num_results=%d",
+      target, base_resources->arg_memrefs.size(),
+      base_resources->result_memrefs.size());
 
   if (target == "__onednn$matmul") {
     const auto& matmul_config = std::get<OneDnnMatMulConfig>(config);
@@ -167,8 +162,8 @@ OneDnnOpThunk::BufferUses OneDnnOpThunk::buffer_uses() const {
 tsl::AsyncValueRef<OneDnnOpThunk::ExecuteEvent> OneDnnOpThunk::Execute(
     const ExecuteParams& params) {
   Eigen::ThreadPoolInterface* thread_pool =
-      params.intra_op_threadpool->getPool();
-  DCHECK(thread_pool != nullptr) << "Thread pool must not be null";
+      params.intra_op_threadpool ? params.intra_op_threadpool->getPool()
+                                 : GetFallbackThreadPoolForOneDnn();
 
   // Create oneDNN runtime for the operation.
   auto runtime = std::make_unique<OneDnnRuntime>(thread_pool, target_);
diff --git a/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_op_thunk_test.cc b/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_op_thunk_test.cc
index e5943177e77..dd2aff0ac5c 100644
--- a/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_op_thunk_test.cc
+++ b/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_op_thunk_test.cc
@@ -380,5 +380,52 @@ TEST(OneDnnOpThunkTest, SimpleOneDnnSoftmaxThunk) {
   }
 }
 
+TEST(OneDnnOpThunkTest, FallbackThreadPool) {
+  // Verify that OneDnnOpThunk executes successfully when
+  // params.intra_op_threadpool is nullptr, exercising the fallback path.
+  Shape lhs_shape = ShapeUtil::MakeShape(F32, {2, 3});
+  Shape rhs_shape = ShapeUtil::MakeShape(F32, {3, 2});
+  Shape out_shape = ShapeUtil::MakeShape(F32, {2, 2});
+
+  Literal lhs_literal = LiteralUtil::CreateR2FromArray2D<float>(
+      Array2D<float>({{1.f, 2.f, 3.f}, {4.f, 5.f, 6.f}}));
+  Literal rhs_literal = LiteralUtil::CreateR2FromArray2D<float>(
+      Array2D<float>({{7.f, 8.f}, {9.f, 10.f}, {11.f, 12.f}}));
+  Literal out_literal = LiteralUtil::CreateR2FromArray2D<float>(
+      Array2D<float>({{0.f, 0.f}, {0.f, 0.f}}));
+
+  auto [lhs_alloc, rhs_alloc, out_alloc] =
+      CreateBufferAllocation(lhs_literal, rhs_literal, out_literal);
+
+  auto [lhs_slice, rhs_slice, out_slice] =
+      CreateBufferAllocationSlice(lhs_alloc, rhs_alloc, out_alloc);
+
+  BufferAllocations allocations =
+      CreateBufferAllocations(lhs_literal, rhs_literal, out_literal);
+
+  OneDnnOpThunk::OpBuffers op_buffers;
+  op_buffers.arguments_buffers = {lhs_slice, rhs_slice};
+  op_buffers.arguments_shapes = {lhs_shape, rhs_shape};
+  op_buffers.results_buffers = {out_slice};
+  op_buffers.results_shapes = {out_shape};
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      auto thunk,
+      OneDnnOpThunk::Create("__onednn$matmul", Thunk::Info(), op_buffers, {}));
+
+  Thunk::ExecuteParams params;
+  params.buffer_allocations = &allocations;
+  params.intra_op_threadpool = nullptr;
+
+  tsl::AsyncValueRef<Thunk::ExecuteEvent> exec_event = thunk->Execute(params);
+  tsl::BlockUntilReady(exec_event);
+  ASSERT_FALSE(exec_event.IsError())
+      << "OneDnnOpThunk execution with fallback thread pool failed";
+
+  Literal expected = LiteralUtil::CreateR2FromArray2D<float>(
+      Array2D<float>({{58.f, 64.f}, {139.f, 154.f}}));
+  EXPECT_EQ(out_literal, expected);
+}
+
 }  // namespace
 }  // namespace xla::cpu
diff --git a/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_threadpool.h b/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_threadpool.h
index ecadb365165..239185aa243 100644
--- a/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_threadpool.h
+++ b/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_threadpool.h
@@ -28,6 +28,7 @@ limitations under the License.
 #include "xla/backends/cpu/runtime/work_queue.h"
 #include "xla/tsl/concurrency/async_value_ref.h"
 #include "xla/tsl/concurrency/chain.h"
+#include "xla/tsl/platform/logging.h"
 
 #define EIGEN_USE_THREADS
 #include "unsupported/Eigen/CXX11/Tensor"
@@ -131,6 +132,14 @@ class OneDnnThreadPool final
   tsl::AsyncValueRef<tsl::Chain> done_event_;
 };
 
+inline Eigen::ThreadPoolInterface* GetFallbackThreadPoolForOneDnn() {
+  static auto* pool = new Eigen::ThreadPool(1);
+  VLOG_FIRST_N(0, 1)
+      << "No intra-op thread pool available. "
+         "Using fallback single-threaded thread pool for oneDNN execution.";
+  return pool;
+}
+
 }  // namespace xla::cpu
 
 #endif  // XLA_BACKENDS_CPU_RUNTIME_ONEDNN_ONEDNN_THREADPOOL_H_
diff --git a/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_threadpool_test.cc b/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_threadpool_test.cc
index f963f3155f7..364b53f9110 100644
--- a/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_threadpool_test.cc
+++ b/third_party/xla/xla/backends/cpu/runtime/onednn/onednn_threadpool_test.cc
@@ -114,5 +114,51 @@ TEST(OneDnnThreadPoolTest, Binary) {
   }
 }
 
+TEST(OneDnnThreadPoolTest, FallbackThreadPool) {
+  Eigen::ThreadPoolInterface* fallback_pool = GetFallbackThreadPoolForOneDnn();
+  ASSERT_NE(fallback_pool, nullptr);
+  EXPECT_EQ(fallback_pool->NumThreads(), 1);
+
+  OneDnnThreadPool threadpool(fallback_pool);
+
+  int64_t d0 = 4;
+  int64_t d1 = 8;
+  int64_t num_elements = d0 * d1;
+
+  dnnl::graph::logical_tensor::dims dims = {d0, d1};
+  dnnl::graph::logical_tensor::dims strides = {d1, 1};
+
+  dnnl::graph::logical_tensor src_tensor(
+      0, dnnl::graph::logical_tensor::data_type::f32, dims, strides);
+  dnnl::graph::logical_tensor dst_tensor(
+      1, dnnl::graph::logical_tensor::data_type::f32, dims, strides);
+
+  TF_ASSERT_OK_AND_ASSIGN(dnnl::graph::graph g,
+                          CreateExpGraph(src_tensor, dst_tensor));
+  std::vector<dnnl::graph::partition> partitions = g.get_partitions();
+
+  dnnl::engine engine(dnnl::engine::kind::cpu, 0);
+  dnnl::stream stream =
+      dnnl::stream(dnnl::threadpool_interop::make_stream(engine, &threadpool));
+
+  std::vector<dnnl::graph::compiled_partition> compiled_partitions;
+  for (const auto& partition : partitions) {
+    compiled_partitions.push_back(
+        partition.compile({src_tensor}, {dst_tensor}, engine));
+  }
+
+  std::vector<float> src_data(num_elements, 1.0f);
+  std::vector<float> dst_data(num_elements, 0.0f);
+
+  dnnl::graph::tensor src(src_tensor, engine, src_data.data());
+  dnnl::graph::tensor dst(dst_tensor, engine, dst_data.data());
+
+  compiled_partitions[0].execute(stream, {src}, {dst});
+
+  for (int i = 0; i < num_elements; ++i) {
+    EXPECT_NEAR(dst_data[i], std::exp(1.0f), 1e-5);
+  }
+}
+
 }  // namespace
 }  // namespace xla::cpu

From d4bc171b06896a3266af9e46b70dc4db1b16dd53 Mon Sep 17 00:00:00 2001
From: Lofrank Hu <lofrank@google.com>
Date: Fri, 9 Oct 2026 03:49:21 -0700
Subject: [PATCH 0348/1000] Fix a duplicate word typo ("be be") in
 custom_call.md.

PiperOrigin-RevId: 996434539
---
 third_party/xla/docs/custom_call.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/xla/docs/custom_call.md b/third_party/xla/docs/custom_call.md
index 68c1a3cdbba..b182040775a 100644
--- a/third_party/xla/docs/custom_call.md
+++ b/third_party/xla/docs/custom_call.md
@@ -38,7 +38,7 @@ compatible `operator()` signature. Constructed handler decodes XLA FFI call
 frame (defined by the stable C API), type check all parameters, and forward
 decoded results to the user-defined callback.
 
-XLA FFI binding heavily relies on template metaprogramming to be be able to
+XLA FFI binding heavily relies on template metaprogramming to be able to
 compile constructed handler to the most efficient machine code. Run time
 overheads are in order of a couple of nanoseconds for each custom call
 parameter.

From 0ac9c68b1cc990bac35757d803c6ced32cda0e03 Mon Sep 17 00:00:00 2001
From: Stanislav Bardyuk <sbardyuk@google.com>
Date: Fri, 9 Oct 2026 03:53:56 -0700
Subject: [PATCH 0349/1000] PR #48560: [XLA:CPU] Support scatter and small
 while loops inside sort comparators
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/48560

## 📝 Summary of Changes

Fixes #47370.

A scatter inside a sort comparator aborts the CPU compiler with
`ir_emitter.cc:3501 Check failed: 1 == assigned_buffers.size() (1 vs. 2)`.
Two things go wrong. Since the fusion emitter took over scatters, scatters
are only lowered when they end up in a fusion, but no fusions are formed
inside embedded computations (comparators, reducers), so the scatter reaches
the legacy nested emitter which never supported it. And even if the scatter is
expanded, the expansion produces a while loop plus in-place
dynamic-update-slices, and the legacy emitter can handle neither in a
comparator: the thread-local buffer code insists on one value per allocation
(that is the CHECK above, also hit by a plain dynamic-update-slice in a
comparator), and a while loop in a comparator dies in `EmitGlobalCall`.

The change:

- `ScatterExpander` takes an optional `extra_filter` (just forwarded to
  `OpExpanderPass`). The CPU pipeline uses it to expand scatters whose parent
  computation is embedded, so they never reach the emitter.
- New CPU pass `EmbeddedWhileLoopUnroller` (`xla/backends/cpu/transforms/`)
  runs right after the expansion and fully unrolls while loops that live in
  embedded computations, using `WhileLoopUnroller` with its default
  thresholds. Loops it cannot unroll are left alone (they fail the same way
  as before).
- `IrEmitter::EmitThreadLocalBufferPointer` accepts a thread-local allocation
  with several aliased values. Thread-local allocations are created with
  `NewAllocation(HloBuffer)`, so every value spans the whole allocation and
  one alloca is exactly right for all of them (a CHECK keeps the byte sizes
  honest).
- `IrEmitter` returns `Unimplemented` instead of CHECK-failing in
  `EmitGlobalCall` when an embedded computation calls a global one
  (`HandleCall`, `HandleWhile`, `HandleConditional`).

Known limitations, both fail with a status now instead of an abort: loops
above the unroller's default threshold (more than 64 scatter indices) are
left alone and get `Unimplemented` from the emitter, and a scatter that is
only transitively embedded (say inside a while body in a comparator) is not
expanded and hits the existing `Scatter is not implemented on CPUs` error.

With that, scatter, variadic scatter, dynamic-update-slice and small while
loops all work inside comparators and produce the correct sort order.

The exact module from the issue additionally needs #48160: the expander's
bounds check leaves a `reduce` in the comparator whose reducer the emitter
never emits (#47366). With both changes it compiles and runs.

## 🎯 Justification

The input is valid HLO (it runs on the interpreter) and crashes the process
instead of compiling or returning an error. The same thread-local CHECK also
fires for an ordinary dynamic-update-slice in a comparator, which is a more
plausible real-world shape than the fuzzer-found scatter.

## 🚀 Kind of Contribution

🐛 Bug Fix

## 🧪 Unit Tests:

- `//xla/service/cpu/tests:cpu_edge_cases_test_cpu`: four new executing
  tests (scatter, variadic scatter, dynamic-update-slice, while loop in a
  sort comparator), all abort at head and pass with this change, plus one
  test checking that a 70-index scatter in a comparator fails with
  `Unimplemented` rather than a crash.
- `//xla/backends/cpu/transforms:embedded_while_loop_unroller_test`: loop
  in a comparator is unrolled, loop in the entry computation is untouched.
- `//xla/service:scatter_expander_test`: the extra filter restricts
  expansion.

## 🧪 Execution Tests:

The `cpu_edge_cases_test_cpu` tests above execute the sorted modules on CPU
and check the result.
Copybara import of the project:

--
aac2ee956e5b27d3d481d16f0b6e0b966c6fdb18 by Stanislav Bardyuk <sbardyuk@google.com>:

[XLA:CPU] Support scatter and small while loops inside sort comparators

Scatters that live in embedded computations (sort comparators, reducers)
never end up in a fusion, so the fusion-based scatter emitter cannot
handle them and they hit the legacy nested emitter, which aborts with
"ir_emitter.cc: Check failed: 1 == assigned_buffers.size()". The same
CHECK fires for a plain dynamic-update-slice in a comparator, and any
while loop in a comparator cannot be emitted at all.

- Give ScatterExpander an optional extra filter and use it on CPU to
  expand scatters whose parent computation is embedded.
- Add EmbeddedWhileLoopUnroller, which fully unrolls (small) while loops
  in embedded computations right after the expansion.
- Let EmitThreadLocalBufferPointer share one alloca between the aliased
  values of a thread-local allocation; they all span the whole
  allocation by construction.
- Return Unimplemented instead of CHECK-failing when an embedded
  computation still calls a global one (e.g. a loop too large to unroll).

Fixes #47370.

--
ac15e342ceb58c8d7ddf62551dd32b584a1348aa by Stanislav Bardyuk <sbardyuk@google.com>:

[XLA:CPU] Address review comments on embedded scatter expansion

- EmbeddedWhileLoopUnroller no longer runs PrepareModuleForUnrolling,
  which rewrote loops in the entry computation as well. Add a test with
  an entry loop next to an embedded loop.
- Run the embedded ScatterExpander and the unroller in an HloPassFix, so
  a scatter that unrolling moves from a while body into a comparator is
  expanded and unrolled in turn. Add the scatter-in-while test.
- Size the thread-local alloca by the allocation (an i8 array) instead of
  by the shape of an arbitrary aliased value, so the IR does not depend
  on hash-map iteration order.

Merging this change closes #48560

PiperOrigin-RevId: 996436761
---
 .../xla/xla/backends/cpu/transforms/BUILD     |  34 +++
 .../embedded_while_loop_unroller.cc           |  82 ++++++
 .../transforms/embedded_while_loop_unroller.h |  45 +++
 .../embedded_while_loop_unroller_test.cc      | 181 ++++++++++++
 third_party/xla/xla/service/cpu/BUILD         |   2 +
 .../xla/xla/service/cpu/cpu_compiler.cc       |  32 ++-
 third_party/xla/xla/service/cpu/ir_emitter.cc |  28 +-
 third_party/xla/xla/service/cpu/ir_emitter.h  |   8 +-
 third_party/xla/xla/service/cpu/tests/BUILD   |   3 +
 .../service/cpu/tests/cpu_edge_cases_test.cc  | 260 ++++++++++++++++++
 .../xla/xla/service/scatter_expander.h        |   8 +-
 .../xla/xla/service/scatter_expander_test.cc  |  41 +++
 12 files changed, 705 insertions(+), 19 deletions(-)
 create mode 100644 third_party/xla/xla/backends/cpu/transforms/embedded_while_loop_unroller.cc
 create mode 100644 third_party/xla/xla/backends/cpu/transforms/embedded_while_loop_unroller.h
 create mode 100644 third_party/xla/xla/backends/cpu/transforms/embedded_while_loop_unroller_test.cc

diff --git a/third_party/xla/xla/backends/cpu/transforms/BUILD b/third_party/xla/xla/backends/cpu/transforms/BUILD
index c9433059586..4ecbca90499 100644
--- a/third_party/xla/xla/backends/cpu/transforms/BUILD
+++ b/third_party/xla/xla/backends/cpu/transforms/BUILD
@@ -33,6 +33,40 @@ package_group(
     ],
 )
 
+cc_library(
+    name = "embedded_while_loop_unroller",
+    srcs = ["embedded_while_loop_unroller.cc"],
+    hdrs = ["embedded_while_loop_unroller.h"],
+    deps = [
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/pass:hlo_pass",
+        "//xla/service:instruction_fusion",
+        "//xla/service:while_loop_unroller",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+)
+
+xla_cc_test(
+    name = "embedded_while_loop_unroller_test",
+    srcs = ["embedded_while_loop_unroller_test.cc"],
+    deps = [
+        ":embedded_while_loop_unroller",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+        "//xla/hlo/utils:hlo_query",
+        "//xla/tests:xla_internal_test_main",
+        "//xla/tsl/platform:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest",
+    ],
+)
+
 cc_library(
     name = "library_rewriter",
     srcs = ["library_rewriter.cc"],
diff --git a/third_party/xla/xla/backends/cpu/transforms/embedded_while_loop_unroller.cc b/third_party/xla/xla/backends/cpu/transforms/embedded_while_loop_unroller.cc
new file mode 100644
index 00000000000..7ba19d1b23e
--- /dev/null
+++ b/third_party/xla/xla/backends/cpu/transforms/embedded_while_loop_unroller.cc
@@ -0,0 +1,82 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/cpu/transforms/embedded_while_loop_unroller.h"
+
+#include <iterator>
+#include <vector>
+
+#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/log/log.h"
+#include "absl/status/status_macros.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/service/instruction_fusion.h"
+#include "xla/service/while_loop_unroller.h"
+
+namespace xla::cpu {
+namespace {
+
+std::vector<HloInstruction*> EmbeddedWhileLoops(
+    HloModule* module,
+    const absl::flat_hash_set<absl::string_view>& execution_threads) {
+  std::vector<HloInstruction*> loops;
+  for (HloComputation* computation :
+       module->MakeComputationPostOrder(execution_threads)) {
+    if (!InstructionFusion::IsEmbeddedComputation(computation)) {
+      continue;
+    }
+    absl::c_copy_if(computation->instructions(), std::back_inserter(loops),
+                    HloPredicateIsOp<HloOpcode::kWhile>);
+  }
+  return loops;
+}
+
+}  // namespace
+
+absl::StatusOr<bool> EmbeddedWhileLoopUnroller::RunImpl(
+    HloModule* module,
+    const absl::flat_hash_set<absl::string_view>& execution_threads) {
+  // Unrolling runs DCE on the module and may expose loops nested in the
+  // unrolled body, so re-collect the loops after every successful unroll.
+  bool changed = false;
+  bool unrolled = true;
+  while (unrolled) {
+    unrolled = false;
+    for (HloInstruction* loop : EmbeddedWhileLoops(module, execution_threads)) {
+      // No `prepare`: PrepareModuleForUnrolling would also rewrite loops in
+      // the entry computation.
+      ABSL_ASSIGN_OR_RETURN(UnrollResult result,
+                       WhileLoopUnroller::UnrollAndReturnReplacement(
+                           loop, /*unroll_factor=*/-1,
+                           /*wrap_in_trivial_loop=*/false,
+                           /*force_unroll=*/false, /*prepare=*/false));
+      if (result.unrolled) {
+        changed = unrolled = true;
+        break;
+      }
+      VLOG(2) << "Could not unroll " << loop->name()
+              << " in embedded computation " << loop->parent()->name();
+    }
+  }
+  return changed;
+}
+
+}  // namespace xla::cpu
diff --git a/third_party/xla/xla/backends/cpu/transforms/embedded_while_loop_unroller.h b/third_party/xla/xla/backends/cpu/transforms/embedded_while_loop_unroller.h
new file mode 100644
index 00000000000..fbb97cc5bbf
--- /dev/null
+++ b/third_party/xla/xla/backends/cpu/transforms/embedded_while_loop_unroller.h
@@ -0,0 +1,45 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_BACKENDS_CPU_TRANSFORMS_EMBEDDED_WHILE_LOOP_UNROLLER_H_
+#define XLA_BACKENDS_CPU_TRANSFORMS_EMBEDDED_WHILE_LOOP_UNROLLER_H_
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/pass/hlo_pass_interface.h"
+
+namespace xla::cpu {
+
+// Fully unrolls while loops inside embedded computations (sort comparators,
+// reducers, ...), which the CPU nested IR emitter cannot emit as loops. Such
+// loops typically come from expanding a scatter in those computations. Loops
+// that cannot be unrolled are left unchanged.
+class EmbeddedWhileLoopUnroller : public HloModulePass {
+ public:
+  absl::string_view name() const override {
+    return "embedded-while-loop-unroller";
+  }
+
+ protected:
+  absl::StatusOr<bool> RunImpl(
+      HloModule* module,
+      const absl::flat_hash_set<absl::string_view>& execution_threads) override;
+};
+
+}  // namespace xla::cpu
+
+#endif  // XLA_BACKENDS_CPU_TRANSFORMS_EMBEDDED_WHILE_LOOP_UNROLLER_H_
diff --git a/third_party/xla/xla/backends/cpu/transforms/embedded_while_loop_unroller_test.cc b/third_party/xla/xla/backends/cpu/transforms/embedded_while_loop_unroller_test.cc
new file mode 100644
index 00000000000..3794adab0ea
--- /dev/null
+++ b/third_party/xla/xla/backends/cpu/transforms/embedded_while_loop_unroller_test.cc
@@ -0,0 +1,181 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/cpu/transforms/embedded_while_loop_unroller.h"
+
+#include <memory>
+#include <string>
+
+#include <gtest/gtest.h>
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
+#include "xla/hlo/utils/hlo_query.h"
+#include "xla/tsl/platform/statusor.h"
+
+namespace xla::cpu {
+namespace {
+
+using EmbeddedWhileLoopUnrollerTest = HloHardwareIndependentTestBase;
+
+constexpr absl::string_view kLoopComputations = R"(
+body {
+  state = (s32[], f32[3]) parameter(0)
+  i = s32[] get-tuple-element(state), index=0
+  buffer = f32[3] get-tuple-element(state), index=1
+  one = s32[] constant(1)
+  next_i = s32[] add(i, one)
+  update = f32[1] constant({7})
+  next_buffer = f32[3] dynamic-update-slice(buffer, update, i)
+  ROOT next_state = (s32[], f32[3]) tuple(next_i, next_buffer)
+}
+
+condition {
+  state = (s32[], f32[3]) parameter(0)
+  i = s32[] get-tuple-element(state), index=0
+  three = s32[] constant(3)
+  ROOT continue = pred[] compare(i, three), direction=LT
+}
+)";
+
+TEST_F(EmbeddedWhileLoopUnrollerTest, UnrollsLoopInSortComparator) {
+  constexpr absl::string_view kModule = R"(
+compare {
+  p0 = s32[] parameter(0)
+  p1 = s32[] parameter(1)
+  zero = s32[] constant(0)
+  zero_f = f32[] constant(0)
+  init = f32[3] broadcast(zero_f), dimensions={}
+  init_state = (s32[], f32[3]) tuple(zero, init)
+  loop = (s32[], f32[3]) while(init_state), condition=condition, body=body
+  result = f32[3] get-tuple-element(loop), index=1
+  first = f32[1] slice(result), slice={[0:1]}
+  first_scalar = f32[] reshape(first)
+  c = f32[] convert(p0)
+  ROOT lt = pred[] compare(c, first_scalar), direction=LT
+}
+
+ENTRY main {
+  x = s32[7] parameter(0)
+  ROOT sorted = s32[7] sort(x), dimensions={0}, to_apply=compare
+}
+)";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(absl::StrCat(
+                           "HloModule test\n", kLoopComputations, kModule)));
+
+  ASSERT_OK_AND_ASSIGN(bool changed,
+                       EmbeddedWhileLoopUnroller().Run(module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_FALSE(hlo_query::ContainsInstrWithOpcode(
+      module->GetComputationWithName("compare"), {HloOpcode::kWhile}));
+}
+
+TEST_F(EmbeddedWhileLoopUnrollerTest, LeavesEntryLoopAlone) {
+  constexpr absl::string_view kModule = R"(
+ENTRY main {
+  zero = s32[] constant(0)
+  init = f32[3] parameter(0)
+  init_state = (s32[], f32[3]) tuple(zero, init)
+  ROOT loop = (s32[], f32[3]) while(init_state), condition=condition, body=body
+}
+)";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(absl::StrCat(
+                           "HloModule test\n", kLoopComputations, kModule)));
+
+  ASSERT_OK_AND_ASSIGN(bool changed,
+                       EmbeddedWhileLoopUnroller().Run(module.get()));
+  EXPECT_FALSE(changed);
+  EXPECT_TRUE(hlo_query::ContainsInstrWithOpcode(module->entry_computation(),
+                                                 {HloOpcode::kWhile}));
+}
+
+// The entry loop carries a loop-invariant scalar constant, which the
+// unroller's module-wide preparation (constant sinking) would have rewritten.
+TEST_F(EmbeddedWhileLoopUnrollerTest, LeavesEntryLoopAloneNextToEmbeddedLoop) {
+  constexpr absl::string_view kModule = R"(
+compare {
+  p0 = s32[] parameter(0)
+  p1 = s32[] parameter(1)
+  zero = s32[] constant(0)
+  zero_f = f32[] constant(0)
+  init = f32[3] broadcast(zero_f), dimensions={}
+  init_state = (s32[], f32[3]) tuple(zero, init)
+  loop = (s32[], f32[3]) while(init_state), condition=condition, body=body
+  result = f32[3] get-tuple-element(loop), index=1
+  first = f32[1] slice(result), slice={[0:1]}
+  first_scalar = f32[] reshape(first)
+  c = f32[] convert(p0)
+  ROOT lt = pred[] compare(c, first_scalar), direction=LT
+}
+
+entry_body {
+  state = (s32[], f32[3], f32[]) parameter(0)
+  i = s32[] get-tuple-element(state), index=0
+  buffer = f32[3] get-tuple-element(state), index=1
+  step = f32[] get-tuple-element(state), index=2
+  one = s32[] constant(1)
+  next_i = s32[] add(i, one)
+  step_b = f32[3] broadcast(step), dimensions={}
+  next_buffer = f32[3] add(buffer, step_b)
+  ROOT next_state = (s32[], f32[3], f32[]) tuple(next_i, next_buffer, step)
+}
+
+entry_condition {
+  state = (s32[], f32[3], f32[]) parameter(0)
+  i = s32[] get-tuple-element(state), index=0
+  three = s32[] constant(3)
+  ROOT continue = pred[] compare(i, three), direction=LT
+}
+
+ENTRY main {
+  x = s32[7] parameter(0)
+  sorted = s32[7] sort(x), dimensions={0}, to_apply=compare
+  zero = s32[] constant(0)
+  init = f32[3] parameter(1)
+  step = f32[] constant(1)
+  init_state = (s32[], f32[3], f32[]) tuple(zero, init, step)
+  loop = (s32[], f32[3], f32[]) while(init_state), condition=entry_condition, body=entry_body
+  buffer = f32[3] get-tuple-element(loop), index=1
+  ROOT out = (s32[7], f32[3]) tuple(sorted, buffer)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(absl::StrCat(
+                           "HloModule test\n", kLoopComputations, kModule)));
+  const std::string entry_before = module->entry_computation()->ToString();
+  const std::string entry_body_before =
+      module->GetComputationWithName("entry_body")->ToString();
+  const std::string entry_condition_before =
+      module->GetComputationWithName("entry_condition")->ToString();
+
+  ASSERT_OK_AND_ASSIGN(bool changed,
+                       EmbeddedWhileLoopUnroller().Run(module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_FALSE(hlo_query::ContainsInstrWithOpcode(
+      module->GetComputationWithName("compare"), {HloOpcode::kWhile}));
+  EXPECT_EQ(module->entry_computation()->ToString(), entry_before);
+  EXPECT_EQ(module->GetComputationWithName("entry_body")->ToString(),
+            entry_body_before);
+  EXPECT_EQ(module->GetComputationWithName("entry_condition")->ToString(),
+            entry_condition_before);
+}
+
+}  // namespace
+}  // namespace xla::cpu
diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD
index 9228102f49d..d3f9f52b2d7 100644
--- a/third_party/xla/xla/service/cpu/BUILD
+++ b/third_party/xla/xla/service/cpu/BUILD
@@ -197,6 +197,7 @@ cc_library(
         "//xla/backends/cpu/runtime:thunk_proto_cc",
         "//xla/backends/cpu/runtime:thunk_proto_cc_impl",
         "//xla/backends/cpu/runtime:thunk_proto_serdes",
+        "//xla/backends/cpu/transforms:embedded_while_loop_unroller",
         "//xla/backends/cpu/transforms:library_rewriter",
         "//xla/backends/cpu/transforms/collectives:all_reduce_combiner",
         "//xla/hlo/analysis:alias_info",
@@ -291,6 +292,7 @@ cc_library(
         "//xla/service:hlo_proto_cc",
         "//xla/service:hlo_proto_util",
         "//xla/service:hlo_verifier",
+        "//xla/service:instruction_fusion",
         "//xla/service:layout_assignment",
         "//xla/service:llvm_compiler",
         "//xla/service:logical_buffer",
diff --git a/third_party/xla/xla/service/cpu/cpu_compiler.cc b/third_party/xla/xla/service/cpu/cpu_compiler.cc
index eb9ca007d8b..a1acec04a1d 100644
--- a/third_party/xla/xla/service/cpu/cpu_compiler.cc
+++ b/third_party/xla/xla/service/cpu/cpu_compiler.cc
@@ -98,6 +98,7 @@ limitations under the License.
 #include "xla/backends/cpu/runtime/thunk.h"
 #include "xla/backends/cpu/target_machine_options.h"
 #include "xla/backends/cpu/transforms/collectives/all_reduce_combiner.h"
+#include "xla/backends/cpu/transforms/embedded_while_loop_unroller.h"
 #include "xla/backends/cpu/transforms/library_rewriter.h"
 #include "xla/backends/cpu/ynn_support.h"
 #include "xla/comparison_util.h"
@@ -212,6 +213,7 @@ limitations under the License.
 #include "xla/service/hlo_module_config.h"
 #include "xla/service/hlo_profile_printer_data.pb.h"
 #include "xla/service/hlo_verifier.h"
+#include "xla/service/instruction_fusion.h"
 #include "xla/service/layout_assignment.h"
 #include "xla/service/llvm_compiler.h"
 #include "xla/service/llvm_ir/llvm_command_line_options.h"
@@ -526,11 +528,10 @@ std::unique_ptr<HloPassFix<HloPassPipeline>> CreateSimplificationPipeline(
     // - Improving numerical properties by hierarchically performing reductions.
     // - Improving performance by allowing parallelism.
     // YNNPACK doesn't need TreeReductionRewriter to do either of these.
-    pipeline->AddPass<TreeReductionRewriter>(
-        [](const HloInstruction* hlo) {
-          return !(IsInstructionPreferredByYnn(hlo) &&
-                   IsReduceLikeOpSupportedByYnn(hlo));
-        });
+    pipeline->AddPass<TreeReductionRewriter>([](const HloInstruction* hlo) {
+      return !(IsInstructionPreferredByYnn(hlo) &&
+               IsReduceLikeOpSupportedByYnn(hlo));
+    });
   } else {
     pipeline->AddPass<TreeReductionRewriter>();
   }
@@ -965,8 +966,24 @@ absl::Status CpuCompiler::RunHloPassesThroughLayoutAssn(
   pipeline.AddPass<SelectAndScatterExpander>();
   pipeline.AddPass<ScatterExpander>(ScatterExpander::kEliminateSimpleScatters);
   pipeline.AddPass<ScatterSimplifier>();
-  if (!kFusionEmitterScatterEnabled) {
+  // The nested IR emitter cannot emit while loops (e.g. the ones the scatter
+  // expansion produces) inside embedded computations, so unroll them.
+  if (kFusionEmitterScatterEnabled) {
+    // Fusions are never formed inside embedded computations (e.g. sort
+    // comparators), so scatters there cannot use the fusion emitter. Unrolling
+    // a loop can move a scatter from its body into an embedded computation, so
+    // expand and unroll to a fixed point.
+    auto& embedded_pipeline = pipeline.AddPass<HloPassFix<HloPassPipeline>>(
+        "embedded_scatter_expansion");
+    embedded_pipeline.AddPass<ScatterExpander>(
+        ScatterExpander::kEliminateAllScatters,
+        [](const HloInstruction* instr) {
+          return InstructionFusion::IsEmbeddedComputation(instr->parent());
+        });
+    embedded_pipeline.AddPass<EmbeddedWhileLoopUnroller>();
+  } else {
     pipeline.AddPass<ScatterExpander>(ScatterExpander::kEliminateAllScatters);
+    pipeline.AddPass<EmbeddedWhileLoopUnroller>();
   }
 
   pipeline.AddPass(CreateSimplificationPipeline(
@@ -2337,8 +2354,7 @@ CpuCompiler::CompileAheadOfTimeThunks(
   ABSL_ASSIGN_OR_RETURN(std::unique_ptr<llvm::TargetMachine> target_machine,
                    target_machine_builder());
 
-  ThunkEmitter::Options thunk_emitter_options = {
-      /*is_aot_compilation=*/true};
+  ThunkEmitter::Options thunk_emitter_options = {/*is_aot_compilation=*/true};
 
   TargetMachineOptions target_machine_options(
       triple.normalize(), target_machine->getTargetCPU(),
diff --git a/third_party/xla/xla/service/cpu/ir_emitter.cc b/third_party/xla/xla/service/cpu/ir_emitter.cc
index 35bf9cf364d..f847eea131d 100644
--- a/third_party/xla/xla/service/cpu/ir_emitter.cc
+++ b/third_party/xla/xla/service/cpu/ir_emitter.cc
@@ -2192,6 +2192,7 @@ absl::Status IrEmitter::HandleFusion(HloInstruction* fusion) {
 }
 
 absl::Status IrEmitter::HandleCall(HloInstruction* call) {
+  ABSL_RETURN_IF_ERROR(CheckGlobalCallee(*call->to_apply()));
   HloComputation* computation = call->to_apply();
 
   ABSL_RETURN_IF_ERROR(EmitTargetAddressForOp(call));
@@ -2523,6 +2524,7 @@ absl::Status IrEmitter::HandleCustomCall(HloInstruction* custom_call) {
 }
 
 absl::Status IrEmitter::HandleWhile(HloInstruction* xla_while) {
+  ABSL_RETURN_IF_ERROR(CheckGlobalCallee(*xla_while->while_body()));
   // Precondition: Condition computation must return a scalar bool.
   HloComputation* condition = xla_while->while_condition();
   TF_RET_CHECK(ShapeUtil::IsScalar(condition->root_instruction()->shape()) &&
@@ -3031,6 +3033,7 @@ absl::Status IrEmitter::HandleConcatenate(HloInstruction* concatenate) {
 }
 
 absl::Status IrEmitter::HandleConditional(HloInstruction* conditional) {
+  ABSL_RETURN_IF_ERROR(CheckGlobalCallee(*conditional->branch_computation(0)));
   auto branch_index = conditional->operand(0);
   int num_branches = conditional->branch_count();
   TF_RET_CHECK(ShapeUtil::IsScalar(branch_index->shape()) &&
@@ -3496,18 +3499,17 @@ llvm::Value* IrEmitter::EmitThreadLocalBufferPointer(
       return param_address_untyped;
     }
 
-    // Thread-local allocations should only be assigned a single buffer.
-    const auto& assigned_buffers = allocation.assigned_buffers();
-    CHECK_EQ(1, assigned_buffers.size());
-    const Shape& shape = assigned_buffers.begin()->first->shape();
-
     std::pair<llvm::Function*, BufferAllocation::Slice> key = {
         compute_function()->function(), slice};
     auto buf_it = thread_local_buffers_.find(key);
     if (buf_it == thread_local_buffers_.end()) {
+      // A thread-local allocation may hold several aliased values (in-place
+      // ops such as dynamic-update-slice share one buffer between operand and
+      // result), so size the alloca by the allocation, not by one value.
       llvm::Value* buffer = llvm_ir::EmitAllocaAtFunctionEntry(
-          IrShapeType(shape), absl::StrCat("thread_local", slice.ToString()),
-          b(), MinimumAlignmentForShape(target_shape));
+          llvm::ArrayType::get(b()->getInt8Ty(), allocation.size()),
+          absl::StrCat("thread_local", slice.ToString()), b(),
+          MinimumAlignmentForShape(target_shape));
       auto it_inserted_pair = thread_local_buffers_.insert({key, buffer});
       CHECK(it_inserted_pair.second);
       buf_it = it_inserted_pair.first;
@@ -3751,6 +3753,18 @@ std::vector<llvm::Value*> IrEmitter::EmitThreadLocalCall(
   return returned_scalars;
 }
 
+absl::Status IrEmitter::CheckGlobalCallee(const HloComputation& callee) {
+  if (absl::c_binary_search(global_computations_, &callee)) {
+    return absl::OkStatus();
+  }
+  // Embedded computations (e.g. sort comparators) are emitted as thread-local
+  // functions, which cannot call the global function of `callee`.
+  return Unimplemented(
+      "Calling computation %s from an embedded computation is not supported "
+      "on CPU.",
+      callee.name());
+}
+
 void IrEmitter::EmitGlobalCall(const HloComputation& callee,
                                absl::string_view name) {
   CHECK(absl::c_binary_search(global_computations_, &callee));
diff --git a/third_party/xla/xla/service/cpu/ir_emitter.h b/third_party/xla/xla/service/cpu/ir_emitter.h
index 2e49bcd30e1..b93b369b3c9 100644
--- a/third_party/xla/xla/service/cpu/ir_emitter.h
+++ b/third_party/xla/xla/service/cpu/ir_emitter.h
@@ -427,6 +427,10 @@ class IrEmitter : public DfsHloVisitorWithDefault,
   // to explicitly pass parameters or return results.
   void EmitGlobalCall(const HloComputation& callee, absl::string_view name);
 
+  // Returns an error if `callee` is not a global computation, i.e. if it is
+  // called from an embedded computation such as a sort comparator.
+  absl::Status CheckGlobalCallee(const HloComputation& callee);
+
   // Returns the buffer to which a global call to `callee` would have written
   // its result.
   llvm::Value* GetBufferForGlobalCallReturnValue(const HloComputation& callee);
@@ -516,7 +520,7 @@ class IrEmitter : public DfsHloVisitorWithDefault,
                               llvm::Align alignment,
                               const llvm_ir::IrArray& containing_array);
 
-  using ReductionGenerator = std ::function<llvm::Value*(
+  using ReductionGenerator = std::function<llvm::Value*(
       llvm::IRBuilderBase*, llvm::Value*, llvm::Value*)>;
 
   // Tries to match the reduction function "function" to a known reduction
@@ -675,8 +679,6 @@ class IrEmitter : public DfsHloVisitorWithDefault,
 
   llvm_ir::AliasAnalysis alias_analysis_;
 
-
-
   // This struct contains all the state needed to emit instructions for
   // profiling a computation.
   class ProfilingState {
diff --git a/third_party/xla/xla/service/cpu/tests/BUILD b/third_party/xla/xla/service/cpu/tests/BUILD
index 4fe29ff21c1..d1eca0f4abe 100644
--- a/third_party/xla/xla/service/cpu/tests/BUILD
+++ b/third_party/xla/xla/service/cpu/tests/BUILD
@@ -447,11 +447,14 @@ xla_test(
     backends = ["cpu"],
     deps = [
         "//xla:literal",
+        "//xla:literal_util",
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",
         "//xla/tests:hlo_test_base",
         "//xla/tests:xla_internal_test_main",
         "//xla/tsl/platform:test",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest",
     ],
diff --git a/third_party/xla/xla/service/cpu/tests/cpu_edge_cases_test.cc b/third_party/xla/xla/service/cpu/tests/cpu_edge_cases_test.cc
index c05054e27db..bd06790996a 100644
--- a/third_party/xla/xla/service/cpu/tests/cpu_edge_cases_test.cc
+++ b/third_party/xla/xla/service/cpu/tests/cpu_edge_cases_test.cc
@@ -20,8 +20,11 @@ limitations under the License.
 #include <vector>
 
 #include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "xla/literal.h"
+#include "xla/literal_util.h"
 #include "xla/shape_util.h"
 #include "xla/tests/hlo_test_base.h"
 #include "xla/tsl/platform/test.h"
@@ -225,5 +228,262 @@ ENTRY entry {
   EXPECT_EQ(data[11], 12);
 }
 
+// Regression tests for https://github.com/openxla/xla/issues/47370: in-place
+// ops and loops inside a sort comparator used to crash the legacy IR emitter.
+// Each comparator computes a value known to be true through the in-place op
+// and returns xor(p0 >= p1, true), i.e. an ascending order.
+class SortComparatorTest : public HloTestBase {
+ protected:
+  void RunAndExpectSorted(absl::string_view hlo_text) {
+    ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text));
+    Literal input = LiteralUtil::CreateR1<int32_t>({3, 1, 4, 1, 5, 9, 2});
+    ASSERT_OK_AND_ASSIGN(const Literal result,
+                         Execute(std::move(module), {&input}));
+    EXPECT_EQ(result, LiteralUtil::CreateR1<int32_t>({1, 1, 2, 3, 4, 5, 9}));
+  }
+};
+
+TEST_F(SortComparatorTest, ScatterInSortComparator) {
+  RunAndExpectSorted(R"(
+HloModule scatter_in_sort_comparator
+
+overwrite {
+  lhs = f32[] parameter(0)
+  ROOT rhs = f32[] parameter(1)
+}
+
+compare {
+  p0 = s32[] parameter(0)
+  p1 = s32[] parameter(1)
+  c = f32[] convert(p1)
+  operand = f32[3] broadcast(c), dimensions={}
+  indices = s32[2,1] constant({{0}, {1}})
+  updates = f32[2] constant({5, 2})
+  scattered = f32[3] scatter(operand, indices, updates), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=overwrite
+  a = f32[1] slice(scattered), slice={[0:1]}
+  b = f32[1] slice(scattered), slice={[1:2]}
+  a_scalar = f32[] reshape(a)
+  b_scalar = f32[] reshape(b)
+  flip = pred[] compare(a_scalar, b_scalar), direction=GT
+  ge = pred[] compare(p0, p1), direction=GE
+  ROOT result = pred[] xor(ge, flip)
+}
+
+ENTRY main {
+  x = s32[7] parameter(0)
+  ROOT sorted = s32[7] sort(x), dimensions={0}, is_stable=true, to_apply=compare
+}
+)");
+}
+
+TEST_F(SortComparatorTest, VariadicScatterInSortComparator) {
+  RunAndExpectSorted(R"(
+HloModule variadic_scatter_in_sort_comparator
+
+overwrite2 {
+  a0 = f32[] parameter(0)
+  a1 = f32[] parameter(1)
+  upd0 = f32[] parameter(2)
+  upd1 = f32[] parameter(3)
+  ROOT t = (f32[], f32[]) tuple(upd0, upd1)
+}
+
+compare {
+  p0 = s32[] parameter(0)
+  p1 = s32[] parameter(1)
+  c = f32[] convert(p1)
+  operand = f32[3] broadcast(c), dimensions={}
+  indices = s32[2,1] constant({{0}, {1}})
+  updates0 = f32[2] constant({5, 2})
+  updates1 = f32[2] constant({1, 9})
+  scattered = (f32[3], f32[3]) scatter(operand, operand, indices, updates0, updates1), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=overwrite2
+  out1 = f32[3] get-tuple-element(scattered), index=1
+  a = f32[1] slice(out1), slice={[0:1]}
+  b = f32[1] slice(out1), slice={[1:2]}
+  a_scalar = f32[] reshape(a)
+  b_scalar = f32[] reshape(b)
+  flip = pred[] compare(b_scalar, a_scalar), direction=GT
+  ge = pred[] compare(p0, p1), direction=GE
+  ROOT result = pred[] xor(ge, flip)
+}
+
+ENTRY main {
+  x = s32[7] parameter(0)
+  ROOT sorted = s32[7] sort(x), dimensions={0}, is_stable=true, to_apply=compare
+}
+)");
+}
+
+TEST_F(SortComparatorTest, DynamicUpdateSliceInSortComparator) {
+  RunAndExpectSorted(R"(
+HloModule dus_in_sort_comparator
+
+compare {
+  p0 = s32[] parameter(0)
+  p1 = s32[] parameter(1)
+  c = f32[] convert(p1)
+  operand = f32[3] broadcast(c), dimensions={}
+  update = f32[1] constant({7})
+  zero = s32[] constant(0)
+  two = s32[] constant(2)
+  idx = s32[] clamp(zero, p0, two)
+  updated = f32[3] dynamic-update-slice(operand, update, idx)
+  picked = f32[1] dynamic-slice(updated, idx), dynamic_slice_sizes={1}
+  picked_scalar = f32[] reshape(picked)
+  three = f32[] constant(3)
+  flip = pred[] compare(picked_scalar, three), direction=GT
+  ge = pred[] compare(p0, p1), direction=GE
+  ROOT result = pred[] xor(ge, flip)
+}
+
+ENTRY main {
+  x = s32[7] parameter(0)
+  ROOT sorted = s32[7] sort(x), dimensions={0}, is_stable=true, to_apply=compare
+}
+)");
+}
+
+TEST_F(SortComparatorTest, WhileInSortComparator) {
+  RunAndExpectSorted(R"(
+HloModule while_in_sort_comparator
+
+body {
+  state = (s32[], f32[3]) parameter(0)
+  i = s32[] get-tuple-element(state), index=0
+  buffer = f32[3] get-tuple-element(state), index=1
+  one = s32[] constant(1)
+  next_i = s32[] add(i, one)
+  ten = s32[] constant(10)
+  value = s32[] add(i, ten)
+  value_f = f32[] convert(value)
+  update = f32[1] reshape(value_f)
+  next_buffer = f32[3] dynamic-update-slice(buffer, update, i)
+  ROOT next_state = (s32[], f32[3]) tuple(next_i, next_buffer)
+}
+
+condition {
+  state = (s32[], f32[3]) parameter(0)
+  i = s32[] get-tuple-element(state), index=0
+  three = s32[] constant(3)
+  ROOT continue = pred[] compare(i, three), direction=LT
+}
+
+compare {
+  p0 = s32[] parameter(0)
+  p1 = s32[] parameter(1)
+  c = f32[] convert(p1)
+  operand = f32[3] broadcast(c), dimensions={}
+  zero = s32[] constant(0)
+  init = (s32[], f32[3]) tuple(zero, operand)
+  loop = (s32[], f32[3]) while(init), condition=condition, body=body
+  buffer_out = f32[3] get-tuple-element(loop), index=1
+  a = f32[1] slice(buffer_out), slice={[0:1]}
+  b = f32[1] slice(buffer_out), slice={[1:2]}
+  a_scalar = f32[] reshape(a)
+  b_scalar = f32[] reshape(b)
+  flip = pred[] compare(b_scalar, a_scalar), direction=GT
+  ge = pred[] compare(p0, p1), direction=GE
+  ROOT result = pred[] xor(ge, flip)
+}
+
+ENTRY main {
+  x = s32[7] parameter(0)
+  ROOT sorted = s32[7] sort(x), dimensions={0}, is_stable=true, to_apply=compare
+}
+)");
+}
+
+// Unrolling the loop moves the scatter into the comparator, where it has to
+// be expanded (and the resulting loop unrolled) in turn.
+TEST_F(SortComparatorTest, ScatterInWhileInSortComparator) {
+  RunAndExpectSorted(R"(
+HloModule scatter_in_while_in_sort_comparator
+
+overwrite {
+  lhs = f32[] parameter(0)
+  ROOT rhs = f32[] parameter(1)
+}
+
+body {
+  state = (s32[], f32[3]) parameter(0)
+  i = s32[] get-tuple-element(state), index=0
+  buffer = f32[3] get-tuple-element(state), index=1
+  one = s32[] constant(1)
+  next_i = s32[] add(i, one)
+  indices = s32[2,1] constant({{0}, {1}})
+  updates = f32[2] constant({5, 2})
+  next_buffer = f32[3] scatter(buffer, indices, updates), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=overwrite
+  ROOT next_state = (s32[], f32[3]) tuple(next_i, next_buffer)
+}
+
+condition {
+  state = (s32[], f32[3]) parameter(0)
+  i = s32[] get-tuple-element(state), index=0
+  two = s32[] constant(2)
+  ROOT continue = pred[] compare(i, two), direction=LT
+}
+
+compare {
+  p0 = s32[] parameter(0)
+  p1 = s32[] parameter(1)
+  c = f32[] convert(p1)
+  operand = f32[3] broadcast(c), dimensions={}
+  zero = s32[] constant(0)
+  init = (s32[], f32[3]) tuple(zero, operand)
+  loop = (s32[], f32[3]) while(init), condition=condition, body=body
+  buffer_out = f32[3] get-tuple-element(loop), index=1
+  a = f32[1] slice(buffer_out), slice={[0:1]}
+  b = f32[1] slice(buffer_out), slice={[1:2]}
+  a_scalar = f32[] reshape(a)
+  b_scalar = f32[] reshape(b)
+  flip = pred[] compare(a_scalar, b_scalar), direction=GT
+  ge = pred[] compare(p0, p1), direction=GE
+  ROOT result = pred[] xor(ge, flip)
+}
+
+ENTRY main {
+  x = s32[7] parameter(0)
+  ROOT sorted = s32[7] sort(x), dimensions={0}, is_stable=true, to_apply=compare
+}
+)");
+}
+
+// A scatter with more indices than the unroll threshold stays a loop, which
+// the nested emitter must reject with a status instead of a CHECK failure.
+TEST_F(SortComparatorTest, LargeScatterInSortComparatorIsUnimplemented) {
+  constexpr absl::string_view kHloText = R"(
+HloModule large_scatter_in_sort_comparator
+
+overwrite {
+  lhs = f32[] parameter(0)
+  ROOT rhs = f32[] parameter(1)
+}
+
+compare {
+  p0 = s32[] parameter(0)
+  p1 = s32[] parameter(1)
+  c = f32[] convert(p1)
+  operand = f32[3] broadcast(c), dimensions={}
+  indices = s32[70,1] iota(), iota_dimension=0
+  updates = f32[70] broadcast(c), dimensions={}
+  scattered = f32[3] scatter(operand, indices, updates), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=overwrite
+  a = f32[1] slice(scattered), slice={[0:1]}
+  a_scalar = f32[] reshape(a)
+  flip = pred[] compare(a_scalar, c), direction=GT
+  ge = pred[] compare(p0, p1), direction=GE
+  ROOT result = pred[] xor(ge, flip)
+}
+
+ENTRY main {
+  x = s32[7] parameter(0)
+  ROOT sorted = s32[7] sort(x), dimensions={0}, is_stable=true, to_apply=compare
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloText));
+  Literal input = LiteralUtil::CreateR1<int32_t>({3, 1, 4, 1, 5, 9, 2});
+  EXPECT_EQ(Execute(std::move(module), {&input}).status().code(),
+            absl::StatusCode::kUnimplemented);
+}
+
 }  // namespace
 }  // namespace xla::cpu
diff --git a/third_party/xla/xla/service/scatter_expander.h b/third_party/xla/xla/service/scatter_expander.h
index d503a7d8290..bab6ce7b8f3 100644
--- a/third_party/xla/xla/service/scatter_expander.h
+++ b/third_party/xla/xla/service/scatter_expander.h
@@ -16,7 +16,10 @@ limitations under the License.
 #ifndef XLA_SERVICE_SCATTER_EXPANDER_H_
 #define XLA_SERVICE_SCATTER_EXPANDER_H_
 
+#include <utility>
+
 #include "xla/hlo/transforms/expanders/op_expander_pass.h"
+#include "xla/util.h"
 
 namespace xla {
 
@@ -49,7 +52,10 @@ class ScatterExpander : public OpExpanderPass {
     kEliminateIndeterministicScatters,
   };
 
-  explicit ScatterExpander(Mode m) : mode_(m) {}
+  // extra_filter: optional additional condition an instruction must satisfy
+  // to be expanded (see OpExpanderPass).
+  explicit ScatterExpander(Mode m, HloPredicate extra_filter = nullptr)
+      : OpExpanderPass(std::move(extra_filter)), mode_(m) {}
 
   absl::string_view name() const override { return "scatter_expander"; }
 
diff --git a/third_party/xla/xla/service/scatter_expander_test.cc b/third_party/xla/xla/service/scatter_expander_test.cc
index 9fd7bf9d213..3332aa2563a 100644
--- a/third_party/xla/xla/service/scatter_expander_test.cc
+++ b/third_party/xla/xla/service/scatter_expander_test.cc
@@ -35,6 +35,8 @@ limitations under the License.
 namespace xla {
 namespace {
 
+namespace op = xla::testing::opcode_matchers;
+
 class ScatterExpanderTest : public HloHardwareIndependentTestBase {
  protected:
   // The HLO parser changes all no layout shapes from the input to have a
@@ -439,5 +441,44 @@ TEST_F(ScatterExpanderTest, DoNotEliminateScatterWithAssociativeFp32Combiner) {
   EXPECT_FALSE(result);
 }
 
+TEST_F(ScatterExpanderTest, ExtraFilterRestrictsExpansion) {
+  const char* const kModuleStr = R"(
+    HloModule scatter_expander
+
+    scatter_computation {
+      parameter0 = s32[] parameter(0)
+      ROOT parameter1 = s32[] parameter(1)
+    }
+
+    ENTRY kernel_entry {
+      operand = s32[5] parameter(0)
+      indices = s32[2,1] parameter(1)
+      updates = s32[2] parameter(2)
+      expanded = s32[5] scatter(operand, indices, updates),
+        to_apply=scatter_computation, update_window_dims={},
+        inserted_window_dims={0}, scatter_dims_to_operand_dims={0},
+        index_vector_dim=1
+      kept = s32[5] scatter(operand, indices, updates),
+        to_apply=scatter_computation, update_window_dims={},
+        inserted_window_dims={0}, scatter_dims_to_operand_dims={0},
+        index_vector_dim=1
+      ROOT tuple = (s32[5], s32[5]) tuple(expanded, kept)
+    })";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kModuleStr));
+
+  ScatterExpander scatter_expander(ScatterExpander::kEliminateAllScatters,
+                                   [](const HloInstruction* instruction) {
+                                     return instruction->name() == "expanded";
+                                   });
+  ASSERT_OK_AND_ASSIGN(bool result,
+                       RunHloPass(&scatter_expander, module.get()));
+  EXPECT_TRUE(result);
+
+  HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(root->operand(0), op::GetTupleElement(op::While()));
+  EXPECT_THAT(root->operand(1), op::Scatter());
+}
+
 }  // namespace
 }  // namespace xla

From f37cb6d8c2edec31cefb3e3e3098e3f1406aa6f3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Eusebio=20Dur=C3=A1n=20Monta=C3=B1a?= <eusebiodm@google.com>
Date: Fri, 9 Oct 2026 03:56:46 -0700
Subject: [PATCH 0350/1000] Set cuDNN graph SM version from device info

During deviceless cuDNN compilation (`dnn_support == nullptr`), `cudnn_frontend::graph::Graph::validate()` checks inspect `context.get_sm_version()`. When the SM version is left at its default (`-1`), `cudnn_frontend

PiperOrigin-RevId: 996437944
---
 third_party/xla/xla/service/BUILD             | 19 +++++-
 .../xla/service/xla_aot_compile_gpu_test.cc   | 65 +++++++++++--------
 third_party/xla/xla/service/xla_compile.bzl   | 24 +++++--
 .../xla/xla/stream_executor/cuda/BUILD        |  2 +
 .../xla/xla/stream_executor/cuda/cuda_dnn.cc  | 20 ++++++
 5 files changed, 97 insertions(+), 33 deletions(-)

diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index d821f25c065..4958417f3bc 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -5968,6 +5968,21 @@ xla_aot_compile_gpu(
     xla_flags = "--xla_gpu_use_new_autotune_cache_format=false --xla_gpu_experimental_enable_conv_fusion=true",
 )
 
+xla_aot_compile_gpu(
+    name = "xla_aot_compile_test_gpu_executable_convolution_deviceless_cudnn",
+    # H100 only: deviceless conv graphs on Blackwell require cuDNN >= 9.19.
+    gpu_targets = [
+        "h100",
+        "h100_mig",
+    ],
+    module = "xla_aot_compile_test_convolution.mlir",
+    xla_flags = " ".join([
+        "--xla_gpu_experimental_enable_conv_fusion=true",
+        "--xla_gpu_cudnn_deviceless_compilation_mode=CUDNN_DEVICELESS_COMPILATION_AUTO",
+        "--xla_gpu_experimental_autotune_backends=cudnn",
+    ]),
+)
+
 xla_aot_compile_gpu_runtime_autotuning(
     name = "xla_aot_compile_test_gpu_executable_convolution_runtime_autotuning",
     gpu_target_config = "//xla/backends/gpu/target_config:specs/h100_sxm.txtpb",
@@ -6049,6 +6064,7 @@ xla_test(
         ":xla_aot_compile_test_gpu_executable",
         ":xla_aot_compile_test_gpu_executable_constant",
         ":xla_aot_compile_test_gpu_executable_convolution",
+        ":xla_aot_compile_test_gpu_executable_convolution_deviceless_cudnn",
         ":xla_aot_compile_test_gpu_executable_convolution_legacy_cache",
         ":xla_aot_compile_test_gpu_executable_hlo",
         ":xla_aot_compile_test_gpu_executable_hlo_legacy_cache",
@@ -6064,7 +6080,8 @@ xla_test(
         ":hlo_runner_pjrt",
         "//xla:literal",
         "//xla:literal_util",
-        "//xla/tests:hlo_test_base",
+        "//xla/backends/gpu/tests:hlo_pjrt_gpu_test_base",
+        "//xla/stream_executor/cuda:cudnn_plugin",
         "//xla/tsl/platform:env",
         "//xla/tsl/platform:test",
         "@com_google_absl//absl/status:status_macros",
diff --git a/third_party/xla/xla/service/xla_aot_compile_gpu_test.cc b/third_party/xla/xla/service/xla_aot_compile_gpu_test.cc
index 53606130117..a82893bbe95 100644
--- a/third_party/xla/xla/service/xla_aot_compile_gpu_test.cc
+++ b/third_party/xla/xla/service/xla_aot_compile_gpu_test.cc
@@ -20,14 +20,16 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/status/status_macros.h"
+#include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
+#include "xla/backends/gpu/tests/hlo_pjrt_gpu_test_base.h"
 #include "xla/literal.h"
 #include "xla/literal_util.h"
 #include "xla/service/hlo_runner_interface.h"
 #include "xla/service/hlo_runner_pjrt.h"
-#include "xla/tests/hlo_test_base.h"
+#include "xla/stream_executor/cuda/cuda_dnn.h"
 #include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/test.h"
 #include "tsl/platform/path.h"
@@ -36,16 +38,18 @@ namespace xla {
 namespace xla_compile {
 namespace {
 
-class XlaCompileTest : public HloTestBase {
+class XlaCompileTest : public gpu::HloPjRtGpuTestBase {
  public:
   void LoadAndRunExecutable(absl::string_view path_to_serialized_aot_result,
                             absl::Span<const Literal* const> args,
-                            const Literal& expected) {
+                            const Literal& expected,
+                            absl::string_view target_suffix = "") {
     const char* test_device = getenv("XLA_TEST_DEVICE");
     ASSERT_NE(test_device, nullptr) << "XLA_TEST_DEVICE is not set";
-    std::string path = tsl::io::JoinPath(
-        tsl::testing::XlaSrcRoot(), "service",
-        absl::StrCat(path_to_serialized_aot_result, "_", test_device));
+    std::string path =
+        tsl::io::JoinPath(tsl::testing::XlaSrcRoot(), "service",
+                          absl::StrCat(path_to_serialized_aot_result, "_",
+                                       test_device, target_suffix));
     std::string serialized_aot_result;
     ASSERT_OK(tsl::ReadFileToString(tsl::Env::Default(), path,
                                     &serialized_aot_result));
@@ -88,8 +92,9 @@ TEST_F(XlaCompileTest, LoadGpuExecutableWithConstant) {
                        expected);
 }
 
-// Should also cover the case of loading a GPU executable with a GEMM.
-TEST_F(XlaCompileTest, LoadGpuExecutableWithConvolution) {
+void LoadAndRunConvolution(XlaCompileTest& test,
+                           absl::string_view path_to_serialized_aot_result,
+                           absl::string_view target_suffix = "") {
   Literal input1 = LiteralUtil::CreateR4<float>(
       {{{{1.0, 2.0}, {3.0, 4.0}, {5.0, 6.0}, {7.0, 8.0}},
         {{11.0, 12.0}, {13.0, 14.0}, {15.0, 16.0}, {17.0, 18.0}},
@@ -103,27 +108,35 @@ TEST_F(XlaCompileTest, LoadGpuExecutableWithConvolution) {
       {{1310.0}, {1466.0}, {1622.0}},
       {{2090.0}, {2246.0}, {2402.0}},
   }});
-  LoadAndRunExecutable("xla_aot_compile_test_gpu_executable_convolution",
-                       {&input1, &input2}, expected);
+  test.LoadAndRunExecutable(path_to_serialized_aot_result, {&input1, &input2},
+                            expected, target_suffix);
+}
+
+// Should also cover the case of loading a GPU executable with a GEMM.
+TEST_F(XlaCompileTest, LoadGpuExecutableWithConvolution) {
+  LoadAndRunConvolution(*this,
+                        "xla_aot_compile_test_gpu_executable_convolution");
 }
 
 TEST_F(XlaCompileTest, LoadGpuExecutableWithConvolutionLegacyCache) {
-  Literal input1 = LiteralUtil::CreateR4<float>(
-      {{{{1.0, 2.0}, {3.0, 4.0}, {5.0, 6.0}, {7.0, 8.0}},
-        {{11.0, 12.0}, {13.0, 14.0}, {15.0, 16.0}, {17.0, 18.0}},
-        {{21.0, 22.0}, {23.0, 24.0}, {25.0, 26.0}, {27.0, 28.0}},
-        {{31.0, 32.0}, {33.0, 34.0}, {35.0, 36.0}, {37.0, 38.0}}}});
-  Literal input2 =
-      LiteralUtil::CreateR4<float>({{{{1.0}, {2.0}}, {{3.0}, {4.0}}},
-                                    {{{5.0}, {6.0}}, {{7.0}, {8.0}}},
-                                    {{{9.0}, {10.0}}, {{11.0}, {12.0}}}});
-  Literal expected = LiteralUtil::CreateR4<float>({{
-      {{1310.0}, {1466.0}, {1622.0}},
-      {{2090.0}, {2246.0}, {2402.0}},
-  }});
-  LoadAndRunExecutable(
-      "xla_aot_compile_test_gpu_executable_convolution_legacy_cache",
-      {&input1, &input2}, expected);
+  LoadAndRunConvolution(
+      *this, "xla_aot_compile_test_gpu_executable_convolution_legacy_cache");
+}
+
+TEST_F(XlaCompileTest, LoadGpuExecutableWithDevicelessCudnnConvolution) {
+  if (!se::gpu::SupportsDevicelessCudnnCompilation()) {
+    GTEST_SKIP() << "Deviceless cuDNN compilation requires cuDNN >= "
+                 << se::gpu::kMinDevicelessCudnnVersion;
+  }
+  const char* test_device = getenv("XLA_TEST_DEVICE");
+  if (test_device == nullptr || absl::string_view(test_device) != "h100") {
+    GTEST_SKIP() << "Only compiled for h100.";
+  }
+  absl::string_view target_suffix =
+      absl::StrContains(device_description().name(), "MIG") ? "_mig" : "";
+  LoadAndRunConvolution(
+      *this, "xla_aot_compile_test_gpu_executable_convolution_deviceless_cudnn",
+      target_suffix);
 }
 
 }  // namespace
diff --git a/third_party/xla/xla/service/xla_compile.bzl b/third_party/xla/xla/service/xla_compile.bzl
index 1dfef3cb54b..8ef53c5da6c 100644
--- a/third_party/xla/xla/service/xla_compile.bzl
+++ b/third_party/xla/xla/service/xla_compile.bzl
@@ -103,7 +103,7 @@ def xla_aot_compile_gpu(
         name,
         module,
         gpu_targets,
-        autotune_results,
+        autotune_results = None,
         xla_flags = ""):
     """Runs xla_compile to compile an MHLO, StableHLO or HLO module into an AotCompilationResult for GPU
 
@@ -111,7 +111,7 @@ def xla_aot_compile_gpu(
         name: The name of the build rule.
         module: The MHLO or StableHLO file to compile.
         gpu_targets: The list of gpu targets.
-        autotune_results: AOT AutotuneResults or AutotuneCache file.
+        autotune_results: Optional AOT AutotuneResults or AutotuneCache file.
         xla_flags: Additional XLA_FLAGS to set during compilation.
     """
 
@@ -126,13 +126,25 @@ def xla_aot_compile_gpu(
             " --platform=gpu" +
             " --gpu_target_config=$(location " + gpu_target_config_map[target] + ")"
         )
-        flags = "--xla_gpu_load_autotune_results_from=$(location " + autotune_results + ")"
+        srcs = [module, gpu_target_config_map[target]]
+        flags = []
+        if autotune_results:
+            srcs.append(autotune_results)
+            flags.append("--xla_gpu_load_autotune_results_from=$(location " + autotune_results + ")")
         if xla_flags:
-            flags = flags + " " + xla_flags
-        cmd = "XLA_FLAGS=\"" + flags + "\" " + cmd
+            flags.append(xla_flags)
+        if flags:
+            cmd = "XLA_FLAGS=\"" + " ".join(flags) + "\" " + cmd
+        cmd = (
+            "if ! err=`" + cmd + " 2>&1`; then " +
+            "echo \"$$err\" >&2; " +
+            "if echo \"$$err\" | grep -q \"Deviceless cuDNN compilation requires cuDNN >=\"; then " +
+            "touch $(location " + compiled_binary + "); " +
+            "else exit 1; fi; fi"
+        )
         native.genrule(
             name = "gen_" + name + "_" + target,
-            srcs = [module, gpu_target_config_map[target], autotune_results],
+            srcs = srcs,
             outs = [name + "_" + target],
             cmd = cmd,
             tags = ["manual"] if target in _MIN_CUDA_VERSION else [],
diff --git a/third_party/xla/xla/stream_executor/cuda/BUILD b/third_party/xla/xla/stream_executor/cuda/BUILD
index 8a259bdc49e..dbd2b0af4b9 100644
--- a/third_party/xla/xla/stream_executor/cuda/BUILD
+++ b/third_party/xla/xla/stream_executor/cuda/BUILD
@@ -569,10 +569,12 @@ cc_library(
         ":cudnn_api_wrappers",
         ":cudnn_frontend_helpers",
         ":cudnn_sdpa_score_mod",
+        "//xla:status_macros",
         "//xla/backends/gpu/target_config:cudnn_device_props",
         "//xla/stream_executor:activate_context",
         "//xla/stream_executor:data_type",
         "//xla/stream_executor:device_address",
+        "//xla/stream_executor:device_description",
         "//xla/stream_executor:dnn",
         "//xla/stream_executor:engine_options",
         "//xla/stream_executor:event_based_timer",
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc b/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc
index 3adc96689a6..51846165556 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_dnn.cc
@@ -54,6 +54,7 @@ limitations under the License.
 #include "third_party/gpus/cuda/include/cuda_runtime_api.h"
 #include "third_party/gpus/cuda/include/driver_types.h"
 #include "xla/backends/gpu/target_config/cudnn_device_props.h"
+#include "xla/status_macros.h"
 #include "xla/stream_executor/activate_context.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/cuda/cuda_diagnostics.h"
@@ -64,6 +65,7 @@ limitations under the License.
 #include "xla/stream_executor/cuda/cudnn_sdpa_score_mod.h"
 #include "xla/stream_executor/data_type.h"
 #include "xla/stream_executor/device_address.h"
+#include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/dnn.h"
 #include "xla/stream_executor/engine_options.h"
 #include "xla/stream_executor/event_based_timer.h"
@@ -3990,6 +3992,17 @@ static absl::StatusOr<cudnn_frontend::ExecutionPlan> RebuildExecutionPlan(
   return {std::move(plan)};
 }
 
+// Sets the SM version in the cuDNN graph. Without this, cuDNN will fallback to
+// querying the GPU, which doesn't work for deviceless compilation.
+absl::Status SetGraphSmVersion(cudnn_frontend::graph::Graph& graph,
+                               const DeviceDescription& gpu_device_info) {
+  const CudaComputeCapability* cc =
+      gpu_device_info.gpu_compute_capability().cuda_compute_capability();
+  TF_RET_CHECK(cc != nullptr);
+  graph.set_sm_version(cc->major * 10 + cc->minor);
+  return absl::OkStatus();
+}
+
 }  // namespace
 
 void FixDimsForRaggedOffset(std::vector<int64_t>& dims, int max_reg_per_batch) {
@@ -4037,6 +4050,7 @@ absl::StatusOr<CudnnGraph> GetCudnnFlashAttentionOperationGraph(
   graph.set_intermediate_data_type(cudnn_frontend::DataType_t::FLOAT)
       .set_io_data_type(ioDataType)
       .set_compute_data_type(cudnn_frontend::DataType_t::FLOAT);
+  ABSL_RETURN_IF_ERROR(SetGraphSmVersion(graph, gpu_device_info));
   int64_t uid = 0;
   auto next_uid = [&]() -> int64_t { return CuDnnTensorUID(uid++); };
 
@@ -4286,6 +4300,7 @@ absl::StatusOr<CudnnGraph> GetCudnnFlashAttentionF8OperationGraph(
   graph.set_intermediate_data_type(cudnn_frontend::DataType_t::FLOAT)
       .set_io_data_type(ioDataType)
       .set_compute_data_type(cudnn_frontend::DataType_t::FLOAT);
+  ABSL_RETURN_IF_ERROR(SetGraphSmVersion(graph, gpu_device_info));
 
   auto next_uid = [uid = 0]() mutable -> int { return CuDnnTensorUID(uid++); };
 
@@ -4420,6 +4435,7 @@ absl::StatusOr<CudnnGraph> GetCudnnFlashAttentionBackwardF8OperationGraph(
   graph.set_compute_data_type(cudnn_frontend::DataType_t::FLOAT)
       .set_intermediate_data_type(cudnn_frontend::DataType_t::FLOAT)
       .set_io_data_type(ioDataType);
+  ABSL_RETURN_IF_ERROR(SetGraphSmVersion(graph, gpu_device_info));
 
   auto next_uid = [uid = 0]() mutable -> int { return CuDnnTensorUID(uid++); };
 
@@ -4607,6 +4623,7 @@ absl::StatusOr<CudnnGraph> GetCudnnBlockScaledDotOperationGraph(
           << "\n global_scale: " << has_global_scale;
 
   cudnn_frontend::graph::Graph graph;
+  ABSL_RETURN_IF_ERROR(SetGraphSmVersion(graph, gpu_device_info));
   auto compute_type = cudnn_frontend::DataType_t::FLOAT;
   graph.set_compute_data_type(compute_type);
   graph.set_intermediate_data_type(compute_type);
@@ -4717,6 +4734,7 @@ absl::StatusOr<CudnnGraph> GetCudnnFlashAttentionBackwardOperationGraph(
 
   using cudnn_frontend::graph::Tensor_attributes;
   cudnn_frontend::graph::Graph graph;
+  ABSL_RETURN_IF_ERROR(SetGraphSmVersion(graph, gpu_device_info));
   if (!(q_desc.type() == k_desc.type() && k_desc.type() == p_desc.type() &&
         p_desc.type() == v_desc.type() && v_desc.type() == do_desc.type() &&
         do_desc.type() == dq_desc.type() && dq_desc.type() == dk_desc.type() &&
@@ -6921,6 +6939,8 @@ absl::Status CudnnGraph::Prepare(dnn::DnnSupport* dnn_support,
     return absl::OkStatus();
   };
 
+  ABSL_RETURN_IF_ERROR(SetGraphSmVersion(graph_, gpu_device_info));
+
   if (dnn_support) {
     const CudnnSupport& cudnn_support =
         static_cast<CudnnSupport&>(*dnn_support);

From 79a4b403bbb5df3ff5d27eaa566bbbcd217bef0a Mon Sep 17 00:00:00 2001
From: Maxime France-Pillois <mfrancep@amd.com>
Date: Fri, 9 Oct 2026 04:03:31 -0700
Subject: [PATCH 0351/1000] PR #47925: [XLA:GPU] Accept
 negative-but-phase-aligned concat tile offsets
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/47925

📝 Summary of Changes
VerifyConcatenateAlignment previously rejected any concatenate whose concat-dim tile offset had a negative constant part (base_offset < 0). That `base_offset` is the constant part inherited from surrounding ops (e.g. a pad's edge_padding_low), so a negative value is legitimate: it only means the anchor tile reaches before the concatenated tensor and the out-of-range prefix is masked.

Replace the sign-based rejection with a sign-agnostic, phase-correct boundary-alignment check:
  - Keep the true (possibly negative) `base_offset`. The operand-location loop naturally selects operand 0 for a negative anchor, and `variable_step = offset_expr - base_offset` collapses to `d[*]*tile_size` for any sign, so the multiple-of-tile-size check stays valid.
  - Reject when the phase-adjusted remaining size `(accumulated_offset + current_op_size) - base_offset` is not a multiple of `tile_size`. This value may exceed the operand size for negative anchors; only its value mod `tile_size` governs whether an operand boundary lands on a tile boundary.
  - Reword the error message to report the phase-adjusted remaining size and `remaining_size mod tile_size`, and update the explanatory comment (the anchor/first tile is just a periodic heuristic; the phase is invariant under shifting `base_offset` by any multiple of `tile_size`).

🎯 Justification
When `xla_gpu_experimental_enable_tiling_propagation=true` some Maxtext models fail (models tested: mixtral and deepseek v2)  with tip upstream XLA and JAX, with the following error:
```
jax.errors.JaxRuntimeError: FAILED_PRECONDITION: Tiling propagation rejected for %concatenate.42.10 = bf16[1,3072,1,128]{3,2,1,0} concatenate(%param_2.543, %param_2.543), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/eval_jaxpr/checkpoint/rematted_computation/concatenate" stack_frame_id=404}: The base offset -64 is negative.
```

The faulty fusion is the following:
```
%fused_transpose.28 (param_0.1335: bf16[8,3072,128], param_1.1001: bf16[1,3072,1,64], param_2.543: bf16[1,3072,1,64]) -> bf16[1,3072,8,128] {
  %param_0.1335 = bf16[8,3072,128]{2,1,0} parameter(0)
  %convert.334.14 = f32[8,3072,128]{2,1,0} convert(%param_0.1335)
  %param_2.543 = bf16[1,3072,1,64]{3,2,1,0} parameter(2)
  %concatenate.42.10 = bf16[1,3072,1,128]{3,2,1,0} concatenate(%param_2.543, %param_2.543), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/eval_jaxpr/checkpoint/rematted_computation/concatenate" stack_frame_id=404}
  %bitcast.19.21 = bf16[3072,128]{1,0} bitcast(%concatenate.42.10), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/eval_jaxpr/checkpoint/rematted_computation/concatenate" stack_frame_id=404}
  %convert.393.11 = f32[3072,128]{1,0} convert(%bitcast.19.21)
  %broadcast.724.7 = f32[8,3072,128]{2,1,0} broadcast(%convert.393.11), dimensions={1,2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/eval_jaxpr/checkpoint/rematted_computation/mul" stack_frame_id=184}
  %mul.935.9 = f32[8,3072,128]{2,1,0} multiply(%convert.334.14, %broadcast.724.7), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/eval_jaxpr/checkpoint/mul" stack_frame_id=184}
  %convert.336.7 = bf16[8,3072,128]{2,1,0} convert(%mul.935.9)
  %bitcast.109.4 = bf16[1,8,3072,128]{3,2,1,0} bitcast(%convert.336.7), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/eval_jaxpr/checkpoint/mul" stack_frame_id=184}
  %slice.92.3 = bf16[1,8,3072,64]{3,2,1,0} slice(%bitcast.109.4), slice={[0:1], [0:8], [0:3072], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/eval_jaxpr/checkpoint/split" stack_frame_id=187}
  %slice.93.3 = bf16[1,8,3072,64]{3,2,1,0} slice(%bitcast.109.4), slice={[0:1], [0:8], [0:3072], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/eval_jaxpr/checkpoint/split" stack_frame_id=187}
  %convert.337.5 = f32[1,8,3072,64]{3,2,1,0} convert(%slice.93.3)
  %negate.8.5 = f32[1,8,3072,64]{3,2,1,0} negate(%convert.337.5), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/eval_jaxpr/checkpoint/neg" stack_frame_id=186}
  %convert.338.3 = bf16[1,8,3072,64]{3,2,1,0} convert(%negate.8.5)
  %concatenate.88.1 = bf16[1,8,3072,128]{3,2,1,0} concatenate(%slice.92.3, %convert.338.3), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/eval_jaxpr/checkpoint/concatenate" stack_frame_id=185}
  %convert.342.3 = f32[1,8,3072,128]{3,2,1,0} convert(%concatenate.88.1)
  %param_1.1001 = bf16[1,3072,1,64]{3,2,1,0} parameter(1)
  %concatenate.40.7 = bf16[1,3072,1,128]{3,2,1,0} concatenate(%param_1.1001, %param_1.1001), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/eval_jaxpr/checkpoint/rematted_computation/concatenate" stack_frame_id=402}
  %bitcast.18.15 = bf16[3072,128]{1,0} bitcast(%concatenate.40.7), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/eval_jaxpr/checkpoint/rematted_computation/concatenate" stack_frame_id=402}
  %convert.392.5 = f32[3072,128]{1,0} convert(%bitcast.18.15)
  %broadcast.725.3 = f32[8,3072,128]{2,1,0} broadcast(%convert.392.5), dimensions={1,2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/eval_jaxpr/checkpoint/rematted_computation/mul" stack_frame_id=183}
  %mul.936.5 = f32[8,3072,128]{2,1,0} multiply(%convert.334.14, %broadcast.725.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/eval_jaxpr/checkpoint/mul" stack_frame_id=183}
  %convert.341.3 = bf16[8,3072,128]{2,1,0} convert(%mul.936.5)
  %bitcast.110.7 = bf16[1,8,3072,128]{3,2,1,0} bitcast(%convert.341.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/eval_jaxpr/checkpoint/mul" stack_frame_id=183}
  %convert.343.7 = f32[1,8,3072,128]{3,2,1,0} convert(%bitcast.110.7)
  %add.650.5 = f32[1,8,3072,128]{3,2,1,0} add(%convert.342.3, %convert.343.7), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/eval_jaxpr/checkpoint/add_any" stack_frame_id=183}
  %convert.344.3 = bf16[1,8,3072,128]{3,2,1,0} convert(%add.650.5)
  ROOT %transpose.442.1 = bf16[1,3072,8,128]{3,2,1,0} transpose(%convert.344.3), dimensions={0,2,1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/eval_jaxpr/checkpoint/add_any" stack_frame_id=183}
}
```

In the `%concatenate.42.10 ` op `%param_2.543` (`bf16[1,3072,1,64]`) is used as __both__ operand 0 and operand 1 of the concatenate along dim-3, doubling the tail dimension 64 → 128.

When `--xla_gpu_experimental_enable_tiling_propagation=true` propagates a tile backwards through the concatenate to determine what slice of `%param_2.543` to load, it computes per-operand offsets independently:

| Operand slot | Maps output tile `[..., T : T+S]` to input offset |
|---|---|
| slot 0 | `T` |
| slot 1 | `T − 64` |

Because both slots alias the __same SSA value__, the tiler tries to unify both required windows into a single tile for `%param_2.543`. When `T = 0` (tile starts at the beginning of the concat output), the slot-1 input offset is `0 − 64 = −64`, which is negative and triggers the FAILED_PRECONDITION assert in `tile_propagation.cc`.

The `if (base_offset < 0)` guard in `VerifyConcatenateAlignment`  that reject tile when the offset at d=0 is negative, is a too conservative condition, as negative offsets are valid cases in real-world scenarios that do not necessarily lead to the tile being rejected; however, the alignment condition must be evaluated before determining the validity of the tile’s configuration. This PR therefore aims to do so.

🚀 Kind of Contribution
Please remove what does not apply: 🐛 Bug Fix

🧪 Unit Tests:
New tests included in this PR

Copybara import of the project:

--
85bdd52285a86ec08ac8bedd3851a9b611781c2d by Maxime France-Pillois <mfrancep@amd.com>:

[XLA:GPU] Accept negative-but-phase-aligned concat tile offsets

VerifyConcatenateAlignment previously rejected any concatenate whose
concat-dim tile offset had a negative constant part (base_offset < 0).
That base_offset is the constant part inherited from surrounding ops
(e.g. a pad's edge_padding_low), so a negative value is legitimate: it
only means the anchor tile reaches before the concatenated tensor and
the out-of-range prefix is masked.

Replace the sign-based rejection with a sign-agnostic, phase-correct
boundary-alignment check:
  - Keep the true (possibly negative) base_offset. The operand-location
    loop naturally selects operand 0 for a negative anchor, and
    variable_step = offset_expr - base_offset collapses to d[*]*tile_size
    for any sign, so the multiple-of-tile-size check stays valid.
  - Reject when the phase-adjusted remaining size
      (accumulated_offset + current_op_size) - base_offset
    is not a multiple of tile_size. This value may exceed the operand
    size for negative anchors; only its value mod tile_size governs
    whether an operand boundary lands on a tile boundary. Clamping the
    anchor to 0 would lose this phase and wrongly accept straddling tiles.
  - Reword the error message to report the phase-adjusted remaining size
    and remaining_size mod tile_size, and update the explanatory comment
    (the anchor/first tile is just a periodic heuristic; the phase is
    invariant under shifting base_offset by any multiple of tile_size).

Tests:
  - tile_propagation_test.cc: add regressions for negative-aligned accept,
    negative-misaligned-phase reject, phase periodicity, and
    negative-aligned-but-bad-interior-operand-size reject; update the
    existing remaining-size test's expected message substring.
  - concat/concat_prune_inactive.hlo: update CHECK-ERROR to the new
    phase-adjusted rejection message (tile size 64, remaining 80,
    remaining mod tile_size 16) and add it to deviceless_tests so the
    FileCheck runs on CPU.

The fix accepts negative-but-phase-aligned offsets while still rejecting
negative-and-phase-misaligned ones (which genuinely straddle operand
boundaries).

--
987d368977ecc171920df5e6ab6ac4793a6ae0fc by Maxime France-Pillois <mfrancep@amd.com>:

VerifyConcatenateAlignment now rejects a negative concat-dim base offset
whose start does not land on a tile boundary (base_offset mod tile_size
!= 0). + improve tests + improve comments and error message

--
b22e90ca05a47998cdd554f31b0249f0dbcae0b3 by Maxime France-Pillois <mfrancep@amd.com>:

Improve error message

--
12ddbe1851c6fe69752ef9d6e67f39e6e999dc04 by Maxime France-Pillois <mfrancep@amd.com>:

update dependencies

--
36d6d8f95a35685636df47885aed773554a28085 by Maxime France-Pillois <mfrancep@amd.com>:

Add missing dependency

Merging this change closes #47925

PiperOrigin-RevId: 996440823
---
 .../tests/concat/concat_prune_inactive.hlo    |   9 +-
 .../concat_prune_inactive_small_tile.hlo      |  57 ++++++
 .../xla/xla/codegen/tiling/experimental/BUILD |   3 +-
 .../tiling/experimental/tile_propagation.cc   |  77 ++++++--
 .../experimental/tile_propagation_test.cc     | 178 +++++++++++++++++-
 5 files changed, 303 insertions(+), 21 deletions(-)
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive_small_tile.hlo

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive.hlo
index 67bcd82f51d..abf5a2a68db 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive.hlo
@@ -7,8 +7,15 @@
 // CHECK-NEXT:   xtile.insert %[[EXTRACT]] into %arg3[%[[C0]], %[[C0]]] [64, 16] [1, 1] : tensor<64x16xf32> -> memref<64x10xf32>
 // CHECK-NEXT:   xtile.return
 
+// The nested concatenate produces a concat-dim tile offset with a negative
+// constant part (base offset -64). Under this fusion's concat-dim tile size
+// (64), the remaining size 80 (the distance from the base offset to the end of
+// operand 0) is not a multiple of 64 (80 mod 64 = 16), so a tile would straddle
+// an operand boundary and the tiling is correctly rejected. This is the
+// sign-agnostic boundary-alignment check from b/491092362 (the old check
+// rejected purely on the negative sign with a different message).
 // TODO: b/528277610 - fix tiling emitter.
-// CHECK-ERROR: Tiling propagation rejected for %inner_concat = f32[1,32,10]{2,1,0} concatenate(%param_1, %param_2), dimensions={1}: The base offset -64 is negative.
+// CHECK-ERROR: Tiling propagation rejected for %inner_concat = f32[1,32,10]{2,1,0} concatenate(%param_1, %param_2), dimensions={1}: The remaining operand size 16 (in concatenate operand 0) is not divisible by the selected tile size 64
 
 HloModule min_repro, is_scheduled=true, entry_computation_layout={(f32[64,10]{1,0}, f32[1,16,10]{2,1,0}, f32[1,16,10]{2,1,0})->f32[64,10]{1,0}}
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive_small_tile.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive_small_tile.hlo
new file mode 100644
index 00000000000..dba7de98b75
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive_small_tile.hlo
@@ -0,0 +1,57 @@
+// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s 2>&1 | FileCheck %s
+// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
+
+// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s 2>&1 | FileCheck %s --check-prefixes=CHECK-TILING
+// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+
+// This is a fork of concat_prune_inactive.hlo that uses a smaller concat-dim
+// tile size (16 instead of 64). With the sign-agnostic boundary-alignment
+// check (b/491092362), a negative base offset is legal as long as every operand
+// boundary lands on a tile boundary (i.e. its distance from the base offset is
+// a multiple of the tile size). Here the relevant operand boundaries along the
+// concat dim land at rows 16, 32 and 64, all multiples of 16, so no tile
+// straddles an operand boundary and the tiling is accepted. Unlike the
+// tile-size-64 variant (which is correctly rejected), tiling propagation
+// succeeds here and the fusion is emitted, so we also exercise the on-device
+// correctness check with tiling propagation enabled.
+
+// CHECK:      xtile.entry_func @triton_fn
+
+// CHECK-TILING:      xtile.entry_func @triton_fn
+
+HloModule min_repro, is_scheduled=true, entry_computation_layout={(f32[64,10]{1,0}, f32[1,16,10]{2,1,0}, f32[1,16,10]{2,1,0})->f32[64,10]{1,0}}
+
+%fused_computation (param_0: f32[64,10], param_1: f32[1,16,10], param_2: f32[1,16,10]) -> f32[64,10] {
+  %param_0 = f32[64,10]{1,0} parameter(0)
+  %param_1 = f32[1,16,10]{2,1,0} parameter(1)
+  %param_2 = f32[1,16,10]{2,1,0} parameter(2)
+  %inner_concat = f32[1,32,10]{2,1,0} concatenate(%param_1, %param_2), dimensions={1}
+  %reshape = f32[32,10]{1,0} reshape(%inner_concat)
+  %concat = f32[96,10]{1,0} concatenate(%param_0, %reshape), dimensions={0}
+  ROOT %slice = f32[64,10]{1,0} slice(%concat), slice={[0:64], [0:10]}
+}
+
+ENTRY %entry_computation (param_0: f32[64,10], param_1: f32[1,16,10], param_2: f32[1,16,10]) -> f32[64,10] {
+  %param_0 = f32[64,10]{1,0} parameter(0)
+  %param_1 = f32[1,16,10]{2,1,0} parameter(1)
+  %param_2 = f32[1,16,10]{2,1,0} parameter(2)
+  ROOT %fusion = f32[64,10]{1,0} fusion(%param_0, %param_1, %param_2), kind=kCustom, calls=%fused_computation,
+    backend_config={
+      "operation_queue_id":"0",
+      "fusion_backend_config":{
+        "kind":"__triton",
+        "block_level_fusion_config":{
+          "num_warps":"4",
+          "output_tiles":[{"sizes":["16","10"]}],
+          "num_ctas":1,
+          "num_stages":1,
+          "is_tma_allowed":true,
+          "is_warp_specialization_allowed":false,
+          "waves_per_eu":0
+        }
+      },
+      "force_earliest_schedule":false,
+      "reification_cost":[],
+      "device_type":"DEVICE_TYPE_INVALID"
+    }
+}
diff --git a/third_party/xla/xla/codegen/tiling/experimental/BUILD b/third_party/xla/xla/codegen/tiling/experimental/BUILD
index a32785c7429..b20684a8534 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/BUILD
+++ b/third_party/xla/xla/codegen/tiling/experimental/BUILD
@@ -142,6 +142,7 @@ xla_cc_test(
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
         "//xla/hlo/testlib:verified_hlo_module",
         "//xla/hlo/utils:hlo_traversal",
+        "//xla/tests:xla_internal_test_main",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
@@ -149,7 +150,7 @@ xla_cc_test(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/types:span",
-        "@com_google_googletest//:gtest_main",
+        "@com_google_googletest//:gtest",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:IR",
     ],
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
index 3358fde4d59..3cdfe73217d 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
@@ -183,14 +183,37 @@ absl::Status VerifyConcatenateAlignment(
   SymbolicExpr offset_expr = output_tile.dim_tiles()[concat_dim].offset;
   SymbolicExpr tile_size_expr = output_tile.dim_tiles()[concat_dim].size;
 
-  // We evaluate the base offset base_offset by setting all dimension variables
-  // to 0. Then we accumulate operand sizes to locate the operand that contains
-  // the base_offset.
-  // We enforce:
+  // We evaluate the constant part of the offset (base_offset) by setting all
+  // dimension variables to 0. Then we accumulate operand sizes to locate the
+  // operand that contains the base_offset.
+  //
+  // The tile grid positions along the concatenate dimension are
+  // `base_offset + k * tile_size` for all integer k. A tiling is only valid if
+  // no tile straddles an operand boundary: every operand boundary must line up
+  // with a tile boundary. Whether a given operand boundary lines up with a tile
+  // boundary depends only on whether the distance from base_offset to that
+  // boundary is a multiple of tile_size. This distance is measured from the
+  // true base_offset and is computed the same way whether base_offset is
+  // positive or negative, so all the checks below hold regardless of the sign
+  // of base_offset. In particular, base_offset may be negative (e.g. inherited
+  // from a pad's edge_padding_low above the concatenate): a negative
+  // base_offset merely means the first tile(s) reach before the concatenated
+  // tensor begins, and the same "distance is a multiple of tile_size" alignment
+  // rule still decides whether the tiling is valid.
+  //
+  // We enforce (regardless of the sign of base_offset):
   // 1. The variable step (offset_expr - base_offset) is divisible by the tile
-  //    size. TODO: b/491092362 - why do we need that?
-  // 2. The remaining size in operand k from index B is divisible by the tile
-  //    size (unless k is the last operand).
+  //    size. With base_offset being the true constant part, this simplifies to
+  //    d[*] * tile_size and guarantees the tile grid steps by whole tiles.
+  //    TODO: b/491092362 - why do we need that?
+  // 2. The remaining size in the located operand,
+  //    `remaining_size = (accumulated_offset + current_op_size) - base_offset`,
+  //    (i.e. the distance from base_offset to the end of that operand) is
+  //    divisible by the tile size (unless it is the last operand), so that the
+  //    end of this operand lines up with a tile boundary. Note this value is
+  //    measured from the true base_offset and may exceed the operand size when
+  //    base_offset is negative; that is intentional, since only whether it is a
+  //    multiple of tile_size matters, independent of the sign of base_offset.
   // 3. All subsequent operand sizes are divisible by the tile size (unless
   //    they are the last operand).
   //
@@ -224,14 +247,11 @@ absl::Status VerifyConcatenateAlignment(
                      " is not a constant."));
   }
 
+  // Keep the true (possibly negative) constant part. A negative base_offset is
+  // valid; the operand-location loop below naturally selects operand 0 for it,
+  // and the sign-agnostic `remaining_size % tile_size` check (which measures
+  // the distance from base_offset to the operand boundary) governs correctness.
   int64_t base_offset = base_offset_expr.GetValue();
-  if (base_offset < 0) {
-    // TODO: b/491092362 - tile offset might negative for all dims set to 0
-    // in some cases so it is not the best heuristic to set base_offset to E(0).
-    return absl::FailedPreconditionError(
-        absl::StrCat("Tiling propagation rejected for ", concatenate.ToString(),
-                     ": The base offset ", base_offset, " is negative."));
-  }
   int64_t base_operand_idx = 0;
   int64_t accumulated_offset = 0;
   while (base_operand_idx < num_operands) {
@@ -253,6 +273,10 @@ absl::Status VerifyConcatenateAlignment(
 
   int64_t current_op_size =
       concatenate.operand(base_operand_idx)->shape().dimensions(concat_dim);
+  // Remaining size: distance from the (true, possibly negative) base_offset to
+  // the boundary at the end of the located operand. Measured the same way for
+  // any sign of base_offset; it may exceed current_op_size when base_offset is
+  // negative, but only whether it is a multiple of tile_size matters.
   int64_t remaining_size = (accumulated_offset + current_op_size) - base_offset;
   SymbolicExpr variable_step =
       (offset_expr - CreateSymbolicConstant(base_offset, ctx)).Canonicalize();
@@ -264,6 +288,25 @@ absl::Status VerifyConcatenateAlignment(
   }
   int64_t tile_size = tile_size_expr.GetValue();
 
+  // A negative base_offset means the anchor tile reaches before operand 0
+  // begins. This is only emittable when operand 0's start lands on a tile
+  // boundary, i.e. base_offset is a whole number of tiles below zero. In that
+  // case the fully-out-of-range prefix consists of entire tiles, which the
+  // emitter masks as whole tiles.
+  if (base_offset < 0) {
+    // The modulo op would give a non-positive value for a negative
+    // base_offset, so we normalize it.
+    int64_t offset_within_tile =
+        ((base_offset % tile_size) + tile_size) % tile_size;
+    if (offset_within_tile != 0) {
+      return absl::FailedPreconditionError(absl::StrCat(
+          "Tiling propagation rejected for ", concatenate.ToString(),
+          ": The negative base offset ", base_offset,
+          " is not aligned to the tile size ", tile_size,
+          " ; operand 0 would require an unsupported low-side (left) mask."));
+    }
+  }
+
   if (!variable_step.IsMultipleOf(tile_size)) {
     return absl::FailedPreconditionError(absl::StrCat(
         "Tiling propagation rejected for ", concatenate.ToString(),
@@ -277,9 +320,9 @@ absl::Status VerifyConcatenateAlignment(
     if (remaining_size % tile_size != 0) {
       return absl::FailedPreconditionError(absl::StrCat(
           "Tiling propagation rejected for ", concatenate.ToString(),
-          ": The remaining dimension size ", remaining_size,
-          " in the concatenate operand ", base_operand_idx,
-          " must be a clean multiple of its tile size ", tile_size));
+          ": The remaining operand size ", remaining_size % tile_size,
+          " (in concatenate operand ", base_operand_idx,
+          ") is not divisible by the selected tile size ", tile_size, "."));
     }
   }
 
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
index 16396f9cab7..c9a50a11cde 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
@@ -1235,8 +1235,182 @@ TEST_F(TilePropagationTest,
       PropagateTileToInput(*tiling_space, *concat, concat_tile, 0).status(),
       StatusIs(absl::StatusCode::kFailedPrecondition,
                ::testing::HasSubstr(
-                   "The remaining dimension size 17 in the concatenate operand "
-                   "1 must be a clean multiple of its tile size 5")));
+                   "The remaining operand size 2 (in concatenate operand 1) is "
+                   "not divisible by the selected tile size 5")));
+}
+
+// Regression test for b/491092362. A concatenate nested below an op whose tile
+// propagation substracts a constant from the offset (e.g. a pad's
+// `edge_padding_low`, or an outer concat) produces a concat-dim offset
+// expression of the form `tid_0 * ts_0 - C`. Evaluated at tid_0 = 0 this gives
+// the constant part `-C`, which is a legitimate (negative) base_offset: it only
+// means the first tile(s) reach before the concatenated tensor begins, which is
+// masked when the tile is materialized.
+//
+// Here the offset is `tid_0 * 16 - 16`. The constant part is -16 (a multiple of
+// the tile size 16), and the non-last operand sizes (16, 32) are multiples of
+// 16, so no tile straddles an operand boundary. This tiling is valid and must
+// be accepted.
+TEST_F(TilePropagationTest, ConcatenateAcceptsNegativeAlignedOperandZero) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    ENTRY e {
+      p0 = f32[16] parameter(0)
+      p1 = f32[32] parameter(1)
+      p2 = f32[48] parameter(2)
+      ROOT concatenate = f32[96] concatenate(p0, p1, p2), dimensions={0}
+    }
+  )");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TilingSpace> tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                          &mlir_context_));
+  EXPECT_OK(tiling_space->AssignTileSizes({16}));
+
+  const Tile& root_tile = tiling_space->tiled_roots()[0];
+  SymbolicExpr shifted_offset = CreateDimExpr(0, &mlir_context_) * 16 - 16;
+  llvm::SmallVector<SymbolicExpr, 1> offsets{shifted_offset};
+  Tile shifted_tile{*tiling_space, offsets, root_tile.sizes(),
+                    root_tile.strides(), root_tile.upper_bounds()};
+
+  ASSERT_OK(
+      PropagateTileToInput(*tiling_space, *root, shifted_tile, 0).status());
+}
+
+// A negative base_offset spanning several whole tiles below zero is still
+// accepted, as long as operand 0's start lands on a tile boundary: the entire
+// out-of-range prefix (here two full tiles) is masked tile-by-tile. Operand
+// sizes [16, 32, 48], tile_size = 16, offset `tid_0 * 16 - 32` (base_offset =
+// -32, a multiple of 16).
+TEST_F(TilePropagationTest, ConcatenateAcceptsNegativeMultiTileAlignedOffset) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    ENTRY e {
+      p0 = f32[16] parameter(0)
+      p1 = f32[32] parameter(1)
+      p2 = f32[48] parameter(2)
+      ROOT concatenate = f32[96] concatenate(p0, p1, p2), dimensions={0}
+    }
+  )");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TilingSpace> tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                          &mlir_context_));
+  EXPECT_OK(tiling_space->AssignTileSizes({16}));
+
+  const Tile& root_tile = tiling_space->tiled_roots()[0];
+  SymbolicExpr shifted_offset = CreateDimExpr(0, &mlir_context_) * 16 - 32;
+  llvm::SmallVector<SymbolicExpr, 1> offsets{shifted_offset};
+  Tile shifted_tile{*tiling_space, offsets, root_tile.sizes(),
+                    root_tile.strides(), root_tile.upper_bounds()};
+
+  ASSERT_OK(
+      PropagateTileToInput(*tiling_space, *root, shifted_tile, 0).status());
+}
+
+// A negative base_offset that does not align operand 0's start with a tile
+// boundary is REJECTED: it would need an unsupported low-side (left) mask on
+// operand 0. Here tile_size = 16 and the offset is `tid_0 * 16 - 4`, so
+// base_offset = -4 and base_offset mod tile_size = 12 != 0.
+TEST_F(TilePropagationTest, ConcatenateRejectsNegativeMisalignedOffset) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    ENTRY e {
+      p0 = f32[16] parameter(0)
+      p1 = f32[32] parameter(1)
+      p2 = f32[48] parameter(2)
+      ROOT concatenate = f32[96] concatenate(p0, p1, p2), dimensions={0}
+    }
+  )");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TilingSpace> tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                          &mlir_context_));
+  EXPECT_OK(tiling_space->AssignTileSizes({16}));
+
+  const Tile& root_tile = tiling_space->tiled_roots()[0];
+  SymbolicExpr shifted_offset = CreateDimExpr(0, &mlir_context_) * 16 - 4;
+  llvm::SmallVector<SymbolicExpr, 1> offsets{shifted_offset};
+  Tile shifted_tile{*tiling_space, offsets, root_tile.sizes(),
+                    root_tile.strides(), root_tile.upper_bounds()};
+
+  EXPECT_THAT(
+      PropagateTileToInput(*tiling_space, *root, shifted_tile, 0).status(),
+      StatusIs(absl::StatusCode::kFailedPrecondition,
+               ::testing::HasSubstr(
+                   "The negative base offset -4 is not aligned to the tile "
+                   "size 16 ; operand 0 would require an unsupported low-side "
+                   "(left) mask.")));
+}
+
+// The verdict for a negative base_offset is invariant under shifting it by a
+// multiple of the tile size: `tid_0 * 16 - 20` has the same verdict as
+// `tid_0 * 16 - 4` (both -4 and -20 give base_offset mod tile_size = 12, so
+// both are rejected as unaligned to operand 0's start), even though the -20
+// anchor's first tile is fully masked. This documents that the anchor choice is
+// immaterial.
+TEST_F(TilePropagationTest, ConcatenateNegativeOffsetVerdictIsPeriodic) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    ENTRY e {
+      p0 = f32[16] parameter(0)
+      p1 = f32[32] parameter(1)
+      p2 = f32[48] parameter(2)
+      ROOT concatenate = f32[96] concatenate(p0, p1, p2), dimensions={0}
+    }
+  )");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TilingSpace> tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                          &mlir_context_));
+  EXPECT_OK(tiling_space->AssignTileSizes({16}));
+
+  const Tile& root_tile = tiling_space->tiled_roots()[0];
+  SymbolicExpr shifted_offset = CreateDimExpr(0, &mlir_context_) * 16 - 20;
+  llvm::SmallVector<SymbolicExpr, 1> offsets{shifted_offset};
+  Tile shifted_tile{*tiling_space, offsets, root_tile.sizes(),
+                    root_tile.strides(), root_tile.upper_bounds()};
+
+  EXPECT_THAT(
+      PropagateTileToInput(*tiling_space, *root, shifted_tile, 0).status(),
+      StatusIs(absl::StatusCode::kFailedPrecondition,
+               ::testing::HasSubstr(
+                   "The negative base offset -20 is not aligned to the tile "
+                   "size 16 ; operand 0 would require an unsupported low-side "
+                   "(left) mask.")));
+}
+
+// A negative, tile-aligned base_offset is still rejected if an interior
+// (non-last) operand size is not a multiple of the tile size, since a tile
+// would then straddle that operand's boundary. Here operand sizes are
+// [16, 20, 48], tile_size = 16, offset `tid_0 * 16 - 16` (aligned). Operand 1
+// size 20 is not a multiple of 16 -> reject.
+TEST_F(TilePropagationTest, ConcatenateRejectsNegativeAlignedBadInteriorSize) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    ENTRY e {
+      p0 = f32[16] parameter(0)
+      p1 = f32[20] parameter(1)
+      p2 = f32[48] parameter(2)
+      ROOT concatenate = f32[84] concatenate(p0, p1, p2), dimensions={0}
+    }
+  )");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TilingSpace> tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                          &mlir_context_));
+  EXPECT_OK(tiling_space->AssignTileSizes({16}));
+
+  const Tile& root_tile = tiling_space->tiled_roots()[0];
+  SymbolicExpr shifted_offset = CreateDimExpr(0, &mlir_context_) * 16 - 16;
+  llvm::SmallVector<SymbolicExpr, 1> offsets{shifted_offset};
+  Tile shifted_tile{*tiling_space, offsets, root_tile.sizes(),
+                    root_tile.strides(), root_tile.upper_bounds()};
+
+  EXPECT_THAT(
+      PropagateTileToInput(*tiling_space, *root, shifted_tile, 0).status(),
+      StatusIs(absl::StatusCode::kFailedPrecondition,
+               ::testing::HasSubstr("must be a multiple of the tile size 16")));
 }
 
 TEST_F(TilePropagationTest, CanPropagateToInputOfScanOp) {

From 976223326e14fdce46ad4a8fff5b170183e27f68 Mon Sep 17 00:00:00 2001
From: Noxtimo <139155109+Noxtimo@users.noreply.github.com>
Date: Fri, 9 Oct 2026 12:46:46 +0100
Subject: [PATCH 0352/1000] Fix NaN preservation in tnp.angle and add boundary
 edge cases in testAngle

---
 tensorflow/python/ops/numpy_ops/np_math_ops.py      | 3 ++-
 tensorflow/python/ops/numpy_ops/np_math_ops_test.py | 7 ++++---
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py
index 5b4d0a5e0a2..c42986c47b7 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@@ -975,11 +975,12 @@ def angle(z, deg=False):  # pylint: disable=missing-function-docstring
       # Workaround for b/147515503
       # `np.pi` and `0` are Python scalars, which would make the result
       # float32 whatever `x` is, so build them in `x`'s dtype instead.
-      return array_ops.where_v2(
+      pi_or_zero = array_ops.where_v2(
           _signbit(x),
           constant_op.constant(np.pi, dtype=x.dtype),
           constant_op.constant(0, dtype=x.dtype),
       )
+      return array_ops.where_v2(math_ops.is_nan(x), x, pi_or_zero)
     else:
       return math_ops.angle(x)
 
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 c78673b7d76..f0f53ee3968 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
@@ -38,8 +38,6 @@ class MathTest(test.TestCase, parameterized.TestCase):
 
   def setUp(self):
     super(MathTest, self).setUp()
-    ops.set_dtype_conversion_mode('legacy')
-    np_math_ops.enable_numpy_methods_on_tensor()
     self.array_transforms = [
         lambda x: x,  # Identity,
         ops.convert_to_tensor,
@@ -176,7 +174,10 @@ class MathTest(test.TestCase, parameterized.TestCase):
     # branch out of Python scalars made it float32 for every real dtype, which
     # for float64 also cost precision.
     for dtype in [np.float16, np.float32, np.float64]:
-      arg = np.array([-2.0, -0.5, -0.0, 0.0, 0.5, 2.0], dtype=dtype)
+      arg = np.array(
+          [-2.0, -0.5, -0.0, 0.0, 0.5, 2.0, np.nan, np.inf, -np.inf],
+          dtype=dtype,
+      )
       self.match(
           np_math_ops.angle(arg), np.angle(arg), msg='angle({})'.format(arg)
       )

From 162a107cb6682beb2dbe66c43cee0dffa4d5f7fc Mon Sep 17 00:00:00 2001
From: gauravarun21 <gauravarun21@gmail.com>
Date: Fri, 9 Oct 2026 13:48:29 +0200
Subject: [PATCH 0353/1000] Address review: assert float16 dtype, add
 non-overflow test case

- Assert the evaluated std/variance tensors keep the float16 dtype in
  testReduceStdFloat16LargeMagnitude, locking in the "same dtype as
  input" contract.
- Add testReduceStdFloat16ModerateMagnitude: a companion test with
  unscaled (~N(0,1)) float16 input where neither the squared
  deviations nor the variance overflow, covering the non-overflow path
  explicitly.

Addresses review feedback from @dmiltr3.
---
 tensorflow/python/ops/math_ops_test.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/tensorflow/python/ops/math_ops_test.py b/tensorflow/python/ops/math_ops_test.py
index a7a4c4b1739..a1fc8a58831 100644
--- a/tensorflow/python/ops/math_ops_test.py
+++ b/tensorflow/python/ops/math_ops_test.py
@@ -157,6 +157,7 @@ class ReduceTest(test_util.TensorFlowTestCase):
     expected_std = np.std(x_np.astype(np.float64))
 
     std = self.evaluate(math_ops.reduce_std(x_f16))
+    self.assertEqual(std.dtype, np.float16)
     self.assertTrue(np.isfinite(std), f"Expected a finite result, got {std}")
     self.assertAllClose(std, expected_std, rtol=0.02)
 
@@ -164,8 +165,31 @@ class ReduceTest(test_util.TensorFlowTestCase):
     # unlike the standard deviation, it should consistently overflow to Inf
     # rather than landing on NaN on some platforms.
     variance = self.evaluate(math_ops.reduce_variance(x_f16))
+    self.assertEqual(variance.dtype, np.float16)
     self.assertTrue(np.isposinf(variance), f"Expected Inf, got {variance}")
 
+  def testReduceStdFloat16ModerateMagnitude(self):
+    # Companion to testReduceStdFloat16LargeMagnitude: covers the path where
+    # neither the squared deviations nor the variance overflow float16, so
+    # the float32-promoted computation should still closely match a float64
+    # reference, same as before the promotion was added.
+    rs = np.random.RandomState(0)
+    x_np = rs.randn(1000).astype(np.float32)
+    x_f16 = constant_op.constant(x_np.astype(np.float16))
+    expected_std = np.std(x_np.astype(np.float64))
+    expected_variance = np.var(x_np.astype(np.float64))
+
+    std = self.evaluate(math_ops.reduce_std(x_f16))
+    variance = self.evaluate(math_ops.reduce_variance(x_f16))
+    self.assertEqual(std.dtype, np.float16)
+    self.assertEqual(variance.dtype, np.float16)
+    self.assertTrue(np.isfinite(std), f"Expected a finite result, got {std}")
+    self.assertTrue(
+        np.isfinite(variance), f"Expected a finite result, got {variance}"
+    )
+    self.assertAllClose(std, expected_std, rtol=0.05)
+    self.assertAllClose(variance, expected_variance, rtol=0.05)
+
   def testReduceStdComplex(self):
     # Ensure that complex values are handled to be consistent with numpy
     complex_ys = [([0 - 1j, 0 + 1j], dtypes.float64),

From 02d4a3dbc3dd618fc5b465ad77314609aa467ebb Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 9 Oct 2026 05:18:05 -0700
Subject: [PATCH 0354/1000] Integrate Triton up to
 0ccad974ee86942e704ac06dafdc0c3f214eb229

[List of integrated commits](https://github.com/triton-lang/triton/compare/2074a1b8a5cd283ce584f8ebca825c4f0b215fc1...0ccad974ee86942e704ac06dafdc0c3f214eb229)

PiperOrigin-RevId: 996471036
---
 .../third_party/triton/common/allocator.patch |   33 +-
 .../triton/common/check_null_encoding.patch   |    2 +-
 .../common/convert_layout_heuristic.patch     |    4 +-
 .../triton/common/cxx17_support.patch         | 1016 +++++++++++++++++
 .../triton/common/disable_amd_test.patch      |    2 +-
 .../triton/common/enable_peer_access.patch    |    4 +-
 .../triton/common/f8e5m2_conversion.patch     |   63 -
 .../triton/common/fix_load_ops.patch          |   82 +-
 .../triton/common/fp64_large_k.patch          |  525 ---------
 .../triton/common/llvm_cl969980821.patch      |    2 +-
 .../triton/common/llvm_cl974500093.patch      |   84 +-
 .../triton/common/llvm_cl978258752.patch      |   19 +-
 .../triton/common/llvm_cl979179688.patch      |    6 +-
 .../triton/common/llvm_cl979526570.patch      |   29 +-
 .../triton/common/llvm_cl983115707.patch      |   42 +-
 .../triton/common/llvm_cl991819890.patch      |   49 +-
 .../triton/common/llvm_cl996124567.patch      |    8 +-
 .../triton/common/mig_limits_pytests.patch    |    9 +-
 .../triton/common/mma_limit_pred.patch        |    2 +-
 .../no_accelerate_through_broadcast.patch     |    2 +-
 .../third_party/triton/common/old_ptxas.patch |    2 +-
 .../xla/third_party/triton/common/series      |    3 +-
 .../common/speed_up_int4_unpacking.patch      |    4 +-
 .../tmem_scale_padding_second_half.patch      |    2 +-
 .../third_party/triton/common/tmemfix.patch   |  180 ++-
 .../useStrictPropertiesInAssemblyFormat.patch |   20 +-
 .../triton/common/verify_nvmma_encoding.patch |    6 +-
 .../triton/oss_only/build_files.patch         |   53 +-
 .../xla/third_party/triton/workspace.bzl      |    4 +-
 .../triton/compilation_pipeline_cuda.cc       |    9 +-
 .../tests/fusion_emitter_device_test.cc       |   17 +-
 .../triton_xla_squeeze_dims_pass.cc           |   10 +-
 .../gpu/autotuning/autotune_cache_key.h       |    2 +-
 33 files changed, 1392 insertions(+), 903 deletions(-)
 create mode 100644 third_party/xla/third_party/triton/common/cxx17_support.patch
 delete mode 100644 third_party/xla/third_party/triton/common/f8e5m2_conversion.patch
 delete mode 100644 third_party/xla/third_party/triton/common/fp64_large_k.patch

diff --git a/third_party/xla/third_party/triton/common/allocator.patch b/third_party/xla/third_party/triton/common/allocator.patch
index 1d21e8582ee..e1799c029cf 100644
--- a/third_party/xla/third_party/triton/common/allocator.patch
+++ b/third_party/xla/third_party/triton/common/allocator.patch
@@ -1,37 +1,6 @@
---- a/python/test/gsan/test_allocator.py
-+++ b/python/test/gsan/test_allocator.py
-@@ -626,12 +626,18 @@ def test_export_import_allocation_handles_maps_real_and_shadow(_direct_allocator
-     ],
- )
- def test_export_import_fabric_handles(explicit_config, allocator_config):
--    result = run_in_process(
--        _run_export_import_fabric_handles_check,
--        args=(explicit_config, ),
--        env={"PYTORCH_CUDA_ALLOC_CONF": allocator_config},
--    )
--    if (isinstance(result.exc, RuntimeError) and str(result.exc) == "gsanExportRuntimeStateHandle failed."
--            and "operation not permitted" in result.driver_stderr_output.lower()):
--        pytest.skip("CUDA fabric handles require an accessible NVIDIA IMEX channel")
--    assert result.exc is None
-+  result = run_in_process(
-+      _run_export_import_fabric_handles_check,
-+      args=(explicit_config,),
-+      env={"PYTORCH_CUDA_ALLOC_CONF": allocator_config},
-+  )
-+  if (
-+      isinstance(result.exc, RuntimeError)
-+      and str(result.exc) == "gsanExportRuntimeStateHandle failed."
-+      and (
-+          "operation not permitted" in result.driver_stderr_output.lower()
-+          or "invalid argument" in result.driver_stderr_output.lower()
-+      )
-+  ):
-+    pytest.skip("CUDA fabric handles require an accessible NVIDIA IMEX channel")
-+  assert result.exc is None
-
 --- a/python/triton/experimental/gsan/src/GSanAllocator.cc
 +++ b/python/triton/experimental/gsan/src/GSanAllocator.cc
-@@ -377,8 +377,14 @@ CUresult ensureContext(int device) {
+@@ -393,8 +393,14 @@ CUresult ensureContext(int device) {
    CUresult res = cuCtxGetCurrent(&ctx);
    if (res != CUDA_SUCCESS)
      return res;
diff --git a/third_party/xla/third_party/triton/common/check_null_encoding.patch b/third_party/xla/third_party/triton/common/check_null_encoding.patch
index 6131a9c2bda..978986b6d7f 100644
--- a/third_party/xla/third_party/triton/common/check_null_encoding.patch
+++ b/third_party/xla/third_party/triton/common/check_null_encoding.patch
@@ -3,7 +3,7 @@ This can happen when warp specialization is enabled, which calls into OptimizePa
 which resets all layouts during its execution.
 --- a/lib/Dialect/TritonGPU/IR/Ops.cpp
 +++ b/lib/Dialect/TritonGPU/IR/Ops.cpp
-@@ -67,6 +67,8 @@ bool isConvertTrivial(ConvertLayoutOp op) {
+@@ -89,6 +89,8 @@ bool isConvertTrivial(ConvertLayoutOp op) {
    auto dstType = op.getType();
    auto srcEncoding = srcType.getEncoding();
    auto dstEncoding = dstType.getEncoding();
diff --git a/third_party/xla/third_party/triton/common/convert_layout_heuristic.patch b/third_party/xla/third_party/triton/common/convert_layout_heuristic.patch
index 9140f4582f2..bd6e8632ce7 100644
--- a/third_party/xla/third_party/triton/common/convert_layout_heuristic.patch
+++ b/third_party/xla/third_party/triton/common/convert_layout_heuristic.patch
@@ -10,7 +10,7 @@ See also b/422133176
 
 --- a/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp
 +++ b/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp
-@@ -342,6 +342,15 @@ SmallVector<Value> LayoutPropagation::propagateToUsers(Value value,
+@@ -340,6 +340,15 @@ SmallVector<Value> LayoutPropagation::propagateToUsers(Value value,
      if (auto reshapeOp = dyn_cast<ReshapeOp>(user);
          reshapeOp && reshapeOp.getEfficientLayout())
        continue;
@@ -25,5 +25,5 @@ See also b/422133176
 +    }
      if (user->hasTrait<OpTrait::SameOperandsAndResultEncoding>() ||
          user->hasTrait<OpTrait::Elementwise>() ||
-         isa<ReduceOp, ExpandDimsOp, ReshapeOp, TransOp, JoinOp, SplitOp,
+         isa<BroadcastOp, ReduceOp, ExpandDimsOp, ReshapeOp, TransOp, JoinOp,
 
diff --git a/third_party/xla/third_party/triton/common/cxx17_support.patch b/third_party/xla/third_party/triton/common/cxx17_support.patch
new file mode 100644
index 00000000000..917968da5c9
--- /dev/null
+++ b/third_party/xla/third_party/triton/common/cxx17_support.patch
@@ -0,0 +1,1016 @@
+
+Remove C++20 features (such as std::bit_cast, std::ranges, associative container
+.contains(), requires clauses, and defaulted operator==) so Triton can be
+compiled with C++17.
+
+
+--- a/CMakeLists.txt
++++ b/CMakeLists.txt
+@@ -9,7 +9,7 @@
+ project(triton CXX C)
+ include(CTest)
+ 
+-set(CMAKE_CXX_STANDARD 20)
++set(CMAKE_CXX_STANDARD 17)
+ set(CMAKE_CXX_STANDARD_REQUIRED ON)
+ set(CMAKE_INCLUDE_CURRENT_DIR ON)
+ list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake")
+
+--- a/include/triton/Analysis/Membar.h
++++ b/include/triton/Analysis/Membar.h
+@@ -137,7 +137,10 @@
+   Kind kind = Unknown;
+   Region *region = nullptr;
+ 
+-  bool operator==(const ThreadIssuer &) const = default;
++  bool operator==(const ThreadIssuer &other) const {
++    return kind == other.kind && region == other.region;
++  }
++  bool operator!=(const ThreadIssuer &other) const { return !(*this == other); }
+ };
+ 
+ // Empty summaries contain no operations. Nonempty summaries preserve an issuer
+@@ -151,7 +154,12 @@
+   bool hasSameKnownIssuer(const IssuerSummary &other) const;
+   void join(const IssuerSummary &other);
+ 
+-  bool operator==(const IssuerSummary &) const = default;
++  bool operator==(const IssuerSummary &other) const {
++    return issuer == other.issuer;
++  }
++  bool operator!=(const IssuerSummary &other) const {
++    return !(*this == other);
++  }
+ 
+ private:
+   std::optional<ThreadIssuer> issuer;
+@@ -176,7 +184,14 @@
+     CompletionSync completion = CompletionSync::None;
+     bool peerWaitNeedsSync = false;
+ 
+-    bool operator==(const ThreadSyncState &) const = default;
++    bool operator==(const ThreadSyncState &other) const {
++      return effectIssuers == other.effectIssuers &&
++             completion == other.completion &&
++             peerWaitNeedsSync == other.peerWaitNeedsSync;
++    }
++    bool operator!=(const ThreadSyncState &other) const {
++      return !(*this == other);
++    }
+   } threadSync;
+   struct ThreadDemandState {
+     bool hasDemand = false;
+@@ -186,7 +201,15 @@
+     IssuerSummary publicationIssuers;
+     bool mayNotifyPeer = false;
+ 
+-    bool operator==(const ThreadDemandState &) const = default;
++    bool operator==(const ThreadDemandState &other) const {
++      return hasDemand == other.hasDemand &&
++             hasDemandBeyondGlobalReads == other.hasDemandBeyondGlobalReads &&
++             publicationIssuers == other.publicationIssuers &&
++             mayNotifyPeer == other.mayNotifyPeer;
++    }
++    bool operator!=(const ThreadDemandState &other) const {
++      return !(*this == other);
++    }
+   } threadDemands;
+ 
+   BlockInfo() = default;
+@@ -275,7 +298,13 @@
+     threadDemands = {};
+   }
+ 
+-  bool operator==(const BlockInfo &) const = default;
++  bool operator==(const BlockInfo &other) const {
++    return syncReadSlices == other.syncReadSlices &&
++           syncWriteSlices == other.syncWriteSlices &&
++           threadSync == other.threadSync &&
++           threadDemands == other.threadDemands;
++  }
++  bool operator!=(const BlockInfo &other) const { return !(*this == other); }
+ 
+   /// Checks one pair of access maps with the same alias and operation filters.
+   static bool isIntersected(const SliceMapT &lhsSlices,
+@@ -359,7 +388,13 @@
+     entryBlockInfo.transformSlices(transform);
+   }
+ 
+-  bool operator==(const MembarInfo &) const = default;
++  bool operator==(const MembarInfo &other) const {
++    return pending == other.pending &&
++           entryBlockInfo == other.entryBlockInfo &&
++           allPathsFromEntrySynced == other.allPathsFromEntrySynced &&
++           warpsSynced == other.warpsSynced;
++  }
++  bool operator!=(const MembarInfo &other) const { return !(*this == other); }
+ };
+ 
+ /// Classify the barriers that synchronize local memory accesses in `op`
+
+--- a/include/triton/Tools/Sys/GetEnv.h
++++ b/include/triton/Tools/Sys/GetEnv.h
+@@ -67,8 +67,8 @@
+ namespace tools {
+ 
+ inline void assertIsRecognized(const std::string &env) {
+-  bool is_invalidating = CACHE_INVALIDATING_ENV_VARS.contains(env.c_str());
+-  bool is_neutral = CACHE_NEUTRAL_ENV_VARS.contains(env.c_str());
++  bool is_invalidating = CACHE_INVALIDATING_ENV_VARS.count(env.c_str());
++  bool is_neutral = CACHE_NEUTRAL_ENV_VARS.count(env.c_str());
+   std::string errmsg = env + "is not recognized. "
+                              "Please add it to triton/Tools/Sys/GetEnv.h";
+   assert((is_invalidating || is_neutral) && errmsg.c_str());
+
+--- a/lib/Dialect/TritonGPU/Transforms/Utility.cpp
++++ b/lib/Dialect/TritonGPU/Transforms/Utility.cpp
+@@ -318,7 +318,7 @@
+ 
+ std::string GraphDumper::emitValueNode(Value value) const {
+   NodeInfo info = onValue(value);
+-  if (!info.contains("label")) {
++  if (info.find("label") == info.end()) {
+     std::string shapeStr = getShapeStr(value.getType());
+     if (auto arg = mlir::dyn_cast<BlockArgument>(value))
+       info["label"] =
+@@ -331,7 +331,7 @@
+ 
+ std::string GraphDumper::emitOperationNode(Operation *op) const {
+   NodeInfo info = onOperation(op);
+-  if (!info.contains("label"))
++  if (info.find("label") == info.end())
+     info["label"] = op->getName().getStringRef().str();
+   return emitNode(getUniqueId(op), info);
+ }
+
+--- a/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionScheduling.cpp
++++ b/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionScheduling.cpp
+@@ -340,7 +340,7 @@
+             cast<DenseI32ArrayAttr>(op->getAttr(kPartitionAttrName))
+                 .asArrayRef();
+         for (auto id : partitionIds) {
+-          if (!manual_partitions.contains(id)) {
++          if (manual_partitions.find(id) == manual_partitions.end()) {
+             auto partition = graph->addPartition();
+             partition->addFlag(Flags::MANUAL);
+             manual_partitions[id] = partition;
+@@ -1213,7 +1213,7 @@
+                 } else if (child->getPartition() == startPartition) {
+                   // found a path, set all nodes on the path to the partition
+                   node->addPartition(startPartition);
+-                  while (parentMap.contains(node)) {
++                  while (parentMap.find(node) != parentMap.end()) {
+                     node = parentMap[node];
+                     node->addPartition(startPartition);
+                   }
+
+--- a/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionSchedulingUtility.cpp
++++ b/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionSchedulingUtility.cpp
+@@ -200,7 +200,7 @@
+       tools::getBoolEnv("TRITON_PARTITION_SCHEDULING_DUMP_LOOP_ONLY");
+ 
+   static std::map<std::string, int> keys;
+-  if (!keys.contains(key)) {
++  if (keys.find(key) == keys.end()) {
+     keys[key] = 0;
+   }
+   auto idx = keys[key];
+
+--- a/lib/Tools/LinearLayout.cpp
++++ b/lib/Tools/LinearLayout.cpp
+@@ -1189,7 +1189,7 @@
+   for (StringAttr dim : getInDimNames()) {
+     int32_t mask = 0;
+     for (int i = 0; i < getInDimSizeLog2(dim); i++, c++) {
+-      if (!basicVars.contains(c)) {
++      if (!basicVars.count(c)) {
+         mask |= (1 << i);
+       }
+     }
+
+--- a/python/src/interpreter.cc
++++ b/python/src/interpreter.cc
+@@ -1,5 +1,4 @@
+ #include <atomic>
+-#include <bit>
+ #include <cmath>
+ #include <cstddef>
+ #include <cstdint>
+@@ -260,12 +259,23 @@
+   }
+ }
+ 
++template <class To, class From>
++typename std::enable_if_t<sizeof(To) == sizeof(From) &&
++                              std::is_trivially_copyable_v<From> &&
++                              std::is_trivially_copyable_v<To>,
++                          To>
++bit_cast(const From &src) noexcept {
++  To dst;
++  std::memcpy(&dst, &src, sizeof(To));
++  return dst;
++}
++
+ triton_half npy_float_to_half(float f) {
+-  return {FromFloatBits(std::bit_cast<uint32_t>(f))};
++  return {FromFloatBits(bit_cast<uint32_t>(f))};
+ }
+ 
+ float npy_half_to_float(triton_half h) {
+-  return std::bit_cast<float>(ToFloatBits(h.value));
++  return bit_cast<float>(ToFloatBits(h.value));
+ }
+ 
+ template <>
+
+--- a/python/src/llvm.cc
++++ b/python/src/llvm.cc
+@@ -945,7 +945,7 @@
+       // Mark linked-in functions as internal because backends use external
+       // linkage as a signifier of kernel functions.
+       for (llvm::Function &fn : dstMod->functions()) {
+-        if (externalFns.contains(fn.getName().str())) {
++        if (externalFns.count(fn.getName().str())) {
+           fn.setLinkage(llvm::GlobalValue::InternalLinkage);
+         }
+       }
+
+--- a/third_party/amd/backend/codegen/CMakeLists.txt
++++ b/third_party/amd/backend/codegen/CMakeLists.txt
+@@ -12,7 +12,7 @@
+ set(TRITON_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../../..")
+ add_library(triton_amd_codegen SHARED amdgpu_codegen.cc
+             "${TRITON_SOURCE_DIR}/lib/Tools/LLVMOptions.cpp")
+-target_compile_features(triton_amd_codegen PRIVATE cxx_std_20)
++target_compile_features(triton_amd_codegen PRIVATE cxx_std_17)
+ target_include_directories(triton_amd_codegen PRIVATE
+                            "${TRITON_SOURCE_DIR}/include")
+ target_include_directories(triton_amd_codegen SYSTEM PRIVATE
+
+--- a/third_party/amd/lib/TritonAMDGPUToLLVM/ConvertFpCastOpToLLVM.cpp
++++ b/third_party/amd/lib/TritonAMDGPUToLLVM/ConvertFpCastOpToLLVM.cpp
+@@ -433,8 +433,9 @@
+ }
+ 
+ // OCP Bf8/Fp8 -> Bf16
+-template <typename SrcFPType>
+-  requires llvm::is_one_of<SrcFPType, Float8E4M3FNType, Float8E5M2Type>::value
++template <typename SrcFPType,
++          typename = std::enable_if_t<llvm::is_one_of<
++              SrcFPType, Float8E4M3FNType, Float8E5M2Type>::value>>
+ SmallVector<Value> OcpF8ToBf16SW(Location loc,
+                                  ConversionPatternRewriter &rewriter,
+                                  const SmallVector<Value> &v) {
+@@ -622,11 +623,13 @@
+ // In saturation mode, inf and out-of-range numbers are converted to the largest
+ // normal number, i.e. ±448. NaNs are converted to NaNs.
+ // For UZ formats please check: https://onnx.ai/onnx/technical/float8.html
+-template <typename SrcFPType, typename DstFPType>
+-  requires llvm::is_one_of<SrcFPType, Float32Type, Float16Type,
+-                           BFloat16Type>::value &&
+-           llvm::is_one_of<DstFPType, Float8E4M3FNType, Float8E4M3FNUZType,
+-                           Float8E5M2FNUZType>::value
++template <
++    typename SrcFPType, typename DstFPType,
++    typename = std::enable_if_t<
++        llvm::is_one_of<SrcFPType, Float32Type, Float16Type,
++                        BFloat16Type>::value &&
++        llvm::is_one_of<DstFPType, Float8E4M3FNType, Float8E4M3FNUZType,
++                        Float8E5M2FNUZType>::value>>
+ Value downcastToFp8rtneOneValue(Location loc,
+                                 ConversionPatternRewriter &rewriter, Value v) {
+   constexpr bool isFp8UZ =
+
+--- a/third_party/amd/lib/TritonAMDGPUToLLVM/ConvertWarpPipeline.cpp
++++ b/third_party/amd/lib/TritonAMDGPUToLLVM/ConvertWarpPipeline.cpp
+@@ -129,7 +129,7 @@
+             "non-warp-pipeline scf.execute_region inside pipelined_for body");
+       ++numClusters;
+     } else if (isWarpPipelineIgnorableBarrier(&op)) {
+-      if (existingBarrierMap.contains(numClusters))
++      if (existingBarrierMap.count(numClusters))
+         return op.emitError("multiple pre-existing barriers between pipeline "
+                             "stages; insert a dummy stage instead");
+       existingBarrierMap[numClusters] = &op;
+@@ -144,8 +144,8 @@
+   if (numClusters < 2)
+     return forOp.emitError(
+         "pipelined_for body must contain at least two pipeline stages");
+-  if (existingBarrierMap.contains(0) &&
+-      existingBarrierMap.contains(numClusters))
++  if (existingBarrierMap.count(0) &&
++      existingBarrierMap.count(numClusters))
+     return forOp.emitError("pipelined_for body has both top-of-loop and "
+                            "bottom-of-loop pre-existing barriers");
+   return success();
+@@ -438,7 +438,7 @@
+     // Normally, we don't expect a pipelined loop begins with a barrier
+     // but sometimes required by memory prefetching pattern.
+     auto bottomBar = existingBarrierMap.find(numClusters);
+-    bool hasTopBarrier = existingBarrierMap.contains(0);
++    bool hasTopBarrier = existingBarrierMap.count(0);
+     bool hasBottomBarrier = bottomBar != existingBarrierMap.end();
+     if (bottomBar != existingBarrierMap.end()) {
+       // validatePipelinedForBody guarantees we cannot have both top and
+
+--- a/third_party/amd/lib/TritonAMDGPUTransforms/BlockPingpong.cpp
++++ b/third_party/amd/lib/TritonAMDGPUTransforms/BlockPingpong.cpp
+@@ -1160,18 +1160,23 @@
+   // Prune Memory operations that may be moved to only those involved in dot
+   // computation. To understand the "cluster assumptions" we also estimate
+   // the impact of any additional loads/stores.
+-  auto nonDotGLoads =
+-      std::ranges::stable_partition(gLoadOps, [&dotGlobalLoads](tt::LoadOp op) {
+-        return dotGlobalLoads.contains(op);
+-      });
+-  auto nonDotLLoads = std::ranges::stable_partition(
+-      lLoadOps, [&dotLocalLoads](ttg::LocalLoadOp op) {
+-        return dotLocalLoads.contains(op);
+-      });
+-  auto nonDotLStores = std::ranges::stable_partition(
+-      lStoreOps, [&dotLocalStores](ttg::LocalStoreOp op) {
+-        return dotLocalStores.contains(op);
+-      });
++  auto nonDotGLoadsIt = std::stable_partition(
++      gLoadOps.begin(), gLoadOps.end(),
++      [&dotGlobalLoads](tt::LoadOp op) { return dotGlobalLoads.contains(op); });
++  llvm::ArrayRef<tt::LoadOp> nonDotGLoads(nonDotGLoadsIt, gLoadOps.end());
++  auto nonDotLLoadsIt =
++      std::stable_partition(lLoadOps.begin(), lLoadOps.end(),
++                            [&dotLocalLoads](ttg::LocalLoadOp op) {
++                              return dotLocalLoads.contains(op);
++                            });
++  llvm::ArrayRef<ttg::LocalLoadOp> nonDotLLoads(nonDotLLoadsIt, lLoadOps.end());
++  auto nonDotLStoresIt =
++      std::stable_partition(lStoreOps.begin(), lStoreOps.end(),
++                            [&dotLocalStores](ttg::LocalStoreOp op) {
++                              return dotLocalStores.contains(op);
++                            });
++  llvm::ArrayRef<ttg::LocalStoreOp> nonDotLStores(nonDotLStoresIt,
++                                                  lStoreOps.end());
+   if (estimateNonDotMemoryImpact<tt::LoadOp>(nonDotGLoads, assumeNotTaken) !=
+       0) {
+     std::stringstream message;
+
+--- a/third_party/amd/lib/TritonAMDGPUTransforms/ConvertToBufferOps.cpp
++++ b/third_party/amd/lib/TritonAMDGPUTransforms/ConvertToBufferOps.cpp
+@@ -499,9 +499,10 @@
+   bool analyzeSmallTensorOfst;
+ };
+ 
+-template <typename SourceOp>
+-  requires llvm::is_one_of<SourceOp, triton::LoadOp,
+-                           triton::gpu::AsyncCopyGlobalToLocalOp>::value
++template <typename SourceOp,
++          typename = std::enable_if_t<llvm::is_one_of<
++              SourceOp, triton::LoadOp,
++              triton::gpu::AsyncCopyGlobalToLocalOp>::value>>
+ struct ConvertTritonLoadToBufferLoad : public mlir::OpRewritePattern<SourceOp> {
+   using OpRewritePattern<SourceOp>::OpRewritePattern;
+ 
+
+--- a/third_party/amd/lib/TritonAMDGPUTransforms/LowerLoops.cpp
++++ b/third_party/amd/lib/TritonAMDGPUTransforms/LowerLoops.cpp
+@@ -679,8 +679,8 @@
+   // Make assignments
+   Clusters clusterVec;
+   schedule.clusters.clear();
+-  std::ranges::generate(clusterVec,
+-                        [&] { return schedule.clusters.newAtBack(); });
++  std::generate(clusterVec.begin(), clusterVec.end(),
++                [&] { return schedule.clusters.newAtBack(); });
+ 
+   clusters[SCHED_GLOBAL_LOAD] = clusterVec[globalLoadCluster];
+   clusters[SCHED_LOCAL_STORE] = clusterVec[localStoreCluster];
+
+--- a/third_party/amd/lib/TritonAMDGPUTransforms/ScheduleLoops.cpp
++++ b/third_party/amd/lib/TritonAMDGPUTransforms/ScheduleLoops.cpp
+@@ -352,8 +352,8 @@
+   stages[SCHED_ASYNC_WAIT] = stages[SCHED_LOCAL_LOAD];
+ 
+   Clusters clusterVec;
+-  std::ranges::generate(clusterVec,
+-                        [&] { return schedule.clusters.newAtBack(); });
++  std::generate(clusterVec.begin(), clusterVec.end(),
++                [&] { return schedule.clusters.newAtBack(); });
+ 
+   // This is a symbolic cluster assignment. In this stage, we only focus on
+   // global load and compute ops.
+@@ -491,8 +491,8 @@
+   LDBG("Build ChainedDotSchedule");
+   tt::CoarseSchedule schedule(numStages);
+   ChainedDotClusters clusters;
+-  std::ranges::generate(clusters,
+-                        [&] { return schedule.clusters.newAtBack(); });
++  std::generate(clusters.begin(), clusters.end(),
++                [&] { return schedule.clusters.newAtBack(); });
+ 
+   // Schedule dots
+   auto dotOpsVec = llvm::to_vector(forOp.getBody()->getOps<tt::DotOp>());
+
+--- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAHelpers.h
++++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAHelpers.h
+@@ -3,7 +3,7 @@
+ #include "triton/Dialect/TritonNvidiaGPU/IR/NvmmaSmemAttrs.h"
+ #include "triton/Tools/LayoutUtils.h"
+ 
+-#include <bit>
++#include "llvm/ADT/bit.h"
+ 
+ namespace mlir {
+ namespace triton {
+@@ -175,7 +175,7 @@
+       currDesc.leadDimensionBaseOffset = 0;
+     }
+     int32_t smemByteOffsetb128 = smemByteOffsetb8 >> 4;
+-    uint64_t descBits = std::bit_cast<uint64_t>(currDesc) + smemByteOffsetb128;
++    uint64_t descBits = llvm::bit_cast<uint64_t>(currDesc) + smemByteOffsetb128;
+     // Add the base address to the descriptor
+     Value low = tb.add(tb.i32_val(uint32_t(descBits)), baseb128);
+     Value high = tb.i32_val(descBits >> 32);
+@@ -288,7 +288,7 @@
+           loc, "MMA instruction exceeds the shared-memory layout");
+ 
+     auto desc =
+-        std::bit_cast<SMEMDescriptor>(mmaVersion == 5 ? 1ULL << 46 : 0ULL);
++        llvm::bit_cast<SMEMDescriptor>(mmaVersion == 5 ? 1ULL << 46 : 0ULL);
+     // The lbo / sbo is defined wrt. the 128b elements
+     desc.leadDimensionBaseOffset = (lbo * bitwidth / 8) >> 4;
+     desc.strideDimensionBaseOffset = (sbo * bitwidth / 8) >> 4;
+
+--- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv2.cpp
++++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv2.cpp
+@@ -782,7 +782,7 @@
+ 
+   TensorCoreType mmaType = getMmaTypeDot(op, aTensorTy, bTensorTy, dTensorTy);
+   const auto &instrMap = isTuring ? mmaInstrPtxTuring : mmaInstrPtxAmpere;
+-  if (!instrMap.contains(mmaType))
++  if (!instrMap.count(mmaType))
+     return op.emitError(
+         "unsupported MMA instruction for the given operand/result types");
+ 
+@@ -816,7 +816,7 @@
+ 
+   TensorCoreType mmaType =
+       getMmaTypeDotScaled(op, aTensorTy, bTensorTy, dTensorTy);
+-  if (!mmaInstrPtxScaled.contains(mmaType))
++  if (!mmaInstrPtxScaled.count(mmaType))
+     return op.emitError(
+         "unsupported MMA instruction for the given operand/result types");
+ 
+
+--- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv5.cpp
++++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv5.cpp
+@@ -6,8 +6,8 @@
+ #include "triton/Conversion/TritonGPUToLLVM/PatternTritonGPUOpToLLVM.h"
+ #include "triton/Dialect/TritonNvidiaGPU/IR/NvmmaSmemAttrs.h"
+ 
++#include "llvm/ADT/bit.h"
+ #include <array>
+-#include <bit>
+ 
+ using namespace mlir;
+ using namespace mlir::triton;
+@@ -201,7 +201,7 @@
+   if (kSize == 64)
+     desc.kSize = 1;
+ 
+-  return b.int_val(32, std::bit_cast<uint32_t>(desc));
++  return b.int_val(32, llvm::bit_cast<uint32_t>(desc));
+ }
+ 
+ static Value createScaleInstDescriptorFp8(ConversionPatternRewriter &rewriter,
+@@ -267,7 +267,7 @@
+     desc.BScaleFactor *= 2;
+   }
+ 
+-  return b.int_val(32, std::bit_cast<uint32_t>(desc));
++  return b.int_val(32, llvm::bit_cast<uint32_t>(desc));
+ }
+ 
+ static Value createScaleInstDescriptorFp4(
+@@ -313,7 +313,7 @@
+     // K96 selectors are byte offsets into three (block32) or six (block16)
+     // consecutive scales, including across scale words.
+     desc.kSizeLower = 1;
+-    return b.int_val(32, std::bit_cast<uint32_t>(desc));
++    return b.int_val(32, llvm::bit_cast<uint32_t>(desc));
+   }
+   if (kSize == 128) {
+     desc.kSizeUpper = 1;
+@@ -345,7 +345,7 @@
+            "MMAv5 with kind=mxf4nvf4 and .block16 only supports SFB_ID 0");
+   }
+ 
+-  return b.int_val(32, std::bit_cast<uint32_t>(desc));
++  return b.int_val(32, llvm::bit_cast<uint32_t>(desc));
+ }
+ 
+ //===----------------------------------------------------------------------===//
+
+--- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
++++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
+@@ -841,7 +841,7 @@
+              triton::nvgpu::MemSyncScope::SYSTEM}};
+     const bool doPTXLDPromotion = !useRed && isPromotableToNVPTXLD(op) &&
+                                   vec == 1 && packed == 1 &&
+-                                  ScopeMap.contains(op.getScope());
++                                  ScopeMap.count(op.getScope());
+ 
+     for (size_t i = 0; i < elemsPerThread; i += vec * packed) {
+       Value rmwPtr = ptrElements[i];
+
+--- a/third_party/proton/common/lib/TraceDataIO/CircularLayoutParser.cpp
++++ b/third_party/proton/common/lib/TraceDataIO/CircularLayoutParser.cpp
+@@ -110,14 +110,14 @@
+         trace.asyncEvents.push_back(entry);
+         continue;
+       }
+-      if (!activeEvent.contains(entry->scopeId)) {
++      if (!activeEvent.count(entry->scopeId)) {
+         activeEvent[entry->scopeId] =
+             CircularLayoutParserResult::ProfileEvent();
+       }
+       auto &activeProfileEvent = activeEvent[entry->scopeId];
+ 
+       auto prevState = ParseState::INIT;
+-      if (scopeState.contains(entry->scopeId))
++      if (scopeState.count(entry->scopeId))
+         prevState = scopeState[entry->scopeId];
+ 
+       if (entry->isStart) {
+@@ -174,7 +174,7 @@
+     if (endpoint.entry->isStart) {
+       // A warp cannot have two outstanding instances of the same static event
+       // because a later end would not identify which instance it completes.
+-      if (activeEndpoints.byWarp.contains(endpoint.uid)) {
++      if (activeEndpoints.byWarp.count(endpoint.uid)) {
+         reportException(
+             ScopeMisMatchException("Async event mismatch: start after start"),
+             buffer.position());
+
+--- a/third_party/proton/common/lib/TraceDataIO/TraceWriter.cpp
++++ b/third_party/proton/common/lib/TraceDataIO/TraceWriter.cpp
+@@ -91,7 +91,7 @@
+ 
+     // Group block traces by proc id
+     int procId = bt.procId;
+-    if (!procToBlockTraces.contains(procId)) {
++    if (!procToBlockTraces.count(procId)) {
+       procToBlockTraces[procId] = {};
+     }
+     procToBlockTraces[procId].push_back(&bt);
+@@ -203,12 +203,12 @@
+     for (const auto &event : trace.profileEvents) {
+       int lineId = eventLineIds.at(event.first.get());
+       int scopeId = event.first->scopeId;
+-      if (!scopeColor.contains(scopeId)) {
++      if (!scopeColor.count(scopeId)) {
+         scopeColor[scopeId] = curColorIndex;
+         curColorIndex = (curColorIndex + 1) % colors.size();
+       }
+       const std::string &color = colors[scopeColor[scopeId]];
+-      std::string name = !metadata.scopeName.contains(scopeId)
++      std::string name = !metadata.scopeName.count(scopeId)
+                              ? "scope_" + std::to_string(scopeId)
+                              : metadata.scopeName.at(scopeId);
+       std::string tid = "warp " + std::to_string(warpId) + " (line " +
+@@ -250,12 +250,12 @@
+   constexpr double freq = 1000.0;
+   for (const auto &link : blockTrace.asyncLinks) {
+     int scopeId = link.first.entry->scopeId;
+-    if (!scopeColor.contains(scopeId)) {
++    if (!scopeColor.count(scopeId)) {
+       scopeColor[scopeId] = curColorIndex;
+       curColorIndex = (curColorIndex + 1) % colors.size();
+     }
+     const std::string &color = colors[scopeColor[scopeId]];
+-    std::string name = !metadata.scopeName.contains(scopeId)
++    std::string name = !metadata.scopeName.count(scopeId)
+                            ? "event_" + std::to_string(scopeId)
+                            : metadata.scopeName.at(scopeId);
+ 
+
+--- a/third_party/proton/csrc/include/Session/Session.h
++++ b/third_party/proton/csrc/include/Session/Session.h
+@@ -151,10 +151,10 @@
+   size_t getSessionId(const std::string &path) { return sessionPaths[path]; }
+ 
+   bool hasSession(const std::string &path) {
+-    return sessionPaths.contains(path);
+-  }
+-
+-  bool hasSession(size_t sessionId) { return sessions.contains(sessionId); }
++    return sessionPaths.count(path);
++  }
++
++  bool hasSession(size_t sessionId) { return sessions.count(sessionId); }
+ 
+   void removeSession(size_t sessionId);
+ 
+
+--- a/third_party/proton/csrc/include/Utility/Map.h
++++ b/third_party/proton/csrc/include/Utility/Map.h
+@@ -68,7 +68,7 @@
+ 
+   bool contains(const Key &key) const {
+     std::shared_lock<std::shared_mutex> lock(mutex);
+-    return map.contains(key);
++    return map.count(key) > 0;
+   }
+ 
+   bool erase(const Key &key) {
+
+--- a/third_party/proton/csrc/include/Utility/Set.h
++++ b/third_party/proton/csrc/include/Utility/Set.h
+@@ -19,7 +19,7 @@
+ 
+   bool contains(const Key &key) const {
+     std::shared_lock<std::shared_mutex> lock(mutex);
+-    return set.contains(key);
++    return set.count(key) > 0;
+   }
+ 
+   bool erase(const Key &key) {
+
+--- a/third_party/proton/csrc/lib/Data/Dump/HatchetMsgPack.cpp
++++ b/third_party/proton/csrc/lib/Data/Dump/HatchetMsgPack.cpp
+@@ -102,14 +102,14 @@
+   }
+   if (isRoot) {
+     if (metricSummary.hasKernelMetric &&
+-        !metrics.contains(MetricKind::Kernel)) {
++        !metrics.count(MetricKind::Kernel)) {
+       metricEntries += kKernelInclusiveCount;
+     }
+     if (metricSummary.hasPCSamplingMetric &&
+-        !metrics.contains(MetricKind::PCSampling)) {
++        !metrics.count(MetricKind::PCSampling)) {
+       metricEntries += PCSamplingMetric::Count;
+     }
+-    if (metricSummary.hasCycleMetric && !metrics.contains(MetricKind::Cycle)) {
++    if (metricSummary.hasCycleMetric && !metrics.count(MetricKind::Cycle)) {
+       metricEntries += kCycleInclusiveCount;
+     }
+   }
+@@ -202,14 +202,14 @@
+   }
+   if (isRoot) {
+     if (metricSummary.hasKernelMetric &&
+-        !metrics.contains(MetricKind::Kernel)) {
++        !metrics.count(MetricKind::Kernel)) {
+       writer.appendBytes(kKernelDurationKey);
+       writer.packUInt(0);
+       writer.appendBytes(kKernelInvocationsKey);
+       writer.packUInt(0);
+     }
+     if (metricSummary.hasPCSamplingMetric &&
+-        !metrics.contains(MetricKind::PCSampling)) {
++        !metrics.count(MetricKind::PCSampling)) {
+       PCSamplingMetric pcSamplingMetric;
+       for (size_t i = 0; i < PCSamplingMetric::Count; i++) {
+         const auto valueName = pcSamplingMetric.getValueName(i);
+@@ -217,7 +217,7 @@
+         writer.packUInt(0);
+       }
+     }
+-    if (metricSummary.hasCycleMetric && !metrics.contains(MetricKind::Cycle)) {
++    if (metricSummary.hasCycleMetric && !metrics.count(MetricKind::Cycle)) {
+       writer.packStr(CycleMetric::getValueName(CycleMetric::Duration));
+       writer.packUInt(0);
+       writer.packStr(
+@@ -348,7 +348,7 @@
+     linkedChildren.reserve(virtualNode.children.size());
+     for (const auto &child : virtualNode.children) {
+       const auto &childNode = virtualTree->getNode(child.id);
+-      if (!childNode.children.empty() || linkedMetrics.contains(child.id)) {
++      if (!childNode.children.empty() || linkedMetrics.count(child.id)) {
+         linkedChildren.push_back(child.id);
+       }
+     }
+
+--- a/third_party/proton/csrc/lib/Data/Dump/TraceDataDump.cpp
++++ b/third_party/proton/csrc/lib/Data/Dump/TraceDataDump.cpp
+@@ -215,7 +215,7 @@
+           }
+ 
+           auto scopeName = event.contexts.back().name;
+-          if (!scopeNameToId.contains(scopeName)) {
++          if (!scopeNameToId.count(scopeName)) {
+             scopeIdToName[curScopeId] = scopeName;
+             scopeNameToId[scopeName] = curScopeId;
+             curScopeId++;
+
+--- a/third_party/proton/csrc/lib/Data/Metric.cpp
++++ b/third_party/proton/csrc/lib/Data/Metric.cpp
+@@ -1,11 +1,25 @@
+ #include "Data/Metric.h"
+ #include "Utility/Errors.h"
+ 
+-#include <bit>
++#include <cstring>
+ #include <stdexcept>
+ #include <type_traits>
+ 
+ namespace proton {
++namespace {
++
++template <class To, class From>
++typename std::enable_if_t<sizeof(To) == sizeof(From) &&
++                              std::is_trivially_copyable_v<From> &&
++                              std::is_trivially_copyable_v<To>,
++                          To>
++bit_cast(const From &src) noexcept {
++  To dst;
++  std::memcpy(&dst, &src, sizeof(To));
++  return dst;
++}
++
++} // namespace
+ 
+ std::map<uint64_t, MetricBuffer::MetricDescriptor>
+     MetricBuffer::metricDescriptors;
+@@ -101,22 +115,22 @@
+                                    stream);
+     runtime->synchronizeStream(stream);
+     if (tensorMetric.typeIndex == variant_index_v<double, MetricValueType>) {
+-      tensorMetricsHost[name] = std::bit_cast<double>(metricVector[0]);
++      tensorMetricsHost[name] = bit_cast<double>(metricVector[0]);
+     } else if (tensorMetric.typeIndex ==
+                variant_index_v<int64_t, MetricValueType>) {
+-      tensorMetricsHost[name] = std::bit_cast<int64_t>(metricVector[0]);
++      tensorMetricsHost[name] = bit_cast<int64_t>(metricVector[0]);
+     } else if (tensorMetric.typeIndex ==
+                variant_index_v<std::vector<double>, MetricValueType>) {
+       std::vector<double> values(tensorMetric.size);
+       for (size_t i = 0; i < tensorMetric.size; ++i) {
+-        values[i] = std::bit_cast<double>(metricVector[i]);
++        values[i] = bit_cast<double>(metricVector[i]);
+       }
+       tensorMetricsHost[name] = std::move(values);
+     } else if (tensorMetric.typeIndex ==
+                variant_index_v<std::vector<int64_t>, MetricValueType>) {
+       std::vector<int64_t> values(tensorMetric.size);
+       for (size_t i = 0; i < tensorMetric.size; ++i) {
+-        values[i] = std::bit_cast<int64_t>(metricVector[i]);
++        values[i] = bit_cast<int64_t>(metricVector[i]);
+       }
+       tensorMetricsHost[name] = std::move(values);
+     } else {
+@@ -166,7 +180,7 @@
+         } else {
+           static_assert(sizeof(T) == sizeof(uint64_t),
+                         "MetricValueType alternative must be 8 bytes");
+-          return std::bit_cast<uint64_t>(value);
++          return bit_cast<uint64_t>(value);
+         }
+       },
+       scalarMetric);
+@@ -204,7 +218,7 @@
+ MetricBuffer::DeviceBuffer &MetricBuffer::getOrCreateBuffer() {
+   std::lock_guard<std::mutex> lock(bufferMutex);
+   auto device = runtime->getDevice();
+-  if (!deviceBuffers.contains(device)) {
++  if (!deviceBuffers.count(device)) {
+     deviceBuffers[device] = DeviceBuffer{};
+     auto &buffer = deviceBuffers.at(device);
+     if (mappedHostBuffer) {
+
+--- a/third_party/proton/csrc/lib/Data/TraceData.cpp
++++ b/third_party/proton/csrc/lib/Data/TraceData.cpp
+@@ -53,7 +53,7 @@
+     void addChild(const Context &context, size_t id) { children[context] = id; }
+ 
+     bool hasChild(const Context &context) const {
+-      return children.contains(context);
++      return children.count(context);
+     }
+ 
+     size_t getChild(const Context &context) const {
+@@ -152,7 +152,7 @@
+     return nextEventId++;
+   }
+ 
+-  bool hasEvent(size_t eventId) { return traceEvents.contains(eventId); }
++  bool hasEvent(size_t eventId) { return traceEvents.count(eventId); }
+ 
+   Event &getEvent(size_t eventId) {
+     auto it = traceEvents.find(eventId);
+@@ -253,7 +253,7 @@
+   auto &event = currentTrace->getEvent(eventId);
+   auto &flexibleMetrics = event.metricSet.flexibleMetrics;
+   for (auto [metricName, metricValue] : metrics) {
+-    if (!flexibleMetrics.contains(metricName)) {
++    if (!flexibleMetrics.count(metricName)) {
+       flexibleMetrics.emplace(metricName,
+                               FlexibleMetric(metricName, metricValue));
+     } else {
+
+--- a/third_party/proton/csrc/lib/Profiler/Cupti/CuptiPCSampling.cpp
++++ b/third_party/proton/csrc/lib/Profiler/Cupti/CuptiPCSampling.cpp
+@@ -365,7 +365,7 @@
+       auto *cubinData = getCubinData(pcData->cubinCrc);
+       auto key =
+           CubinData::LineInfoKey{pcData->functionIndex, pcData->pcOffset};
+-      if (!cubinData->lineInfo.contains(key)) {
++      if (!cubinData->lineInfo.count(key)) {
+         auto [lineNumber, fileName, dirName] =
+             getSassToSourceCorrelation(pcData->functionName, pcData->pcOffset,
+                                        cubinData->cubin, cubinData->cubinSize);
+@@ -376,7 +376,7 @@
+       auto &lineInfo = cubinData->lineInfo[key];
+       for (size_t j = 0; j < pcData->stallReasonCount; ++j) {
+         auto *stallReason = &pcData->stallReason[j];
+-        if (!configureData->stallReasonIndexToMetricIndex.contains(
++        if (!configureData->stallReasonIndexToMetricIndex.count(
+                 stallReason->pcSamplingStallReasonIndex))
+           throw makeOutOfRange("Invalid stall reason index");
+         for (const auto &[data, baseEntry] : dataToEntry) {
+@@ -392,7 +392,7 @@
+                   [stallReason->pcSamplingStallReasonIndex]);
+           auto samples = stallReason->samples;
+           auto stalledSamples =
+-              configureData->notIssuedStallReasonIndices.contains(
++              configureData->notIssuedStallReasonIndices.count(
+                   stallReason->pcSamplingStallReasonIndex)
+                   ? 0
+                   : samples;
+
+--- a/third_party/proton/csrc/lib/Profiler/Graph.cpp
++++ b/third_party/proton/csrc/lib/Profiler/Graph.cpp
+@@ -4,12 +4,27 @@
+ #include "Runtime/Runtime.h"
+ #include "Utility/Errors.h"
+ 
+-#include <bit>
++#include <cstring>
+ #include <optional>
+ #include <stdexcept>
++#include <type_traits>
+ #include <vector>
+ 
+ namespace proton {
++namespace {
++
++template <class To, class From>
++typename std::enable_if_t<sizeof(To) == sizeof(From) &&
++                              std::is_trivially_copyable_v<From> &&
++                              std::is_trivially_copyable_v<To>,
++                          To>
++bit_cast(const From &src) noexcept {
++  To dst;
++  std::memcpy(&dst, &src, sizeof(To));
++  return dst;
++}
++
++} // namespace
+ 
+ void GraphState::recordNode(uint64_t nodeId, const std::string &name,
+                             std::optional<MetricNodeState> metricNodeState,
+@@ -69,7 +84,7 @@
+ void GraphState::buildLaunchEntries(const DataToEntryMap &dataToEntry,
+                                     DataToEntryMap &dataToGraphEntry) const {
+   for (const auto &[data, entry] : dataToEntry) {
+-    if (!capturedData.contains(data))
++    if (!capturedData.count(data))
+       // This data object was not enabled during graph capture.
+       continue;
+     dataToGraphEntry.insert({data, entry});
+@@ -153,11 +168,11 @@
+       break;
+     }
+     case variant_index_v<int64_t, MetricValueType>: {
+-      metricValueVariant = std::bit_cast<int64_t>(readWord(wordOffset));
++      metricValueVariant = bit_cast<int64_t>(readWord(wordOffset));
+       break;
+     }
+     case variant_index_v<double, MetricValueType>: {
+-      metricValueVariant = std::bit_cast<double>(readWord(wordOffset));
++      metricValueVariant = bit_cast<double>(readWord(wordOffset));
+       break;
+     }
+     case variant_index_v<std::vector<uint64_t>, MetricValueType>: {
+@@ -171,7 +186,7 @@
+     case variant_index_v<std::vector<int64_t>, MetricValueType>: {
+       std::vector<int64_t> values(metricDesc.size);
+       for (size_t j = 0; j < metricDesc.size; ++j) {
+-        values[j] = std::bit_cast<int64_t>(readWord(wordOffset + j));
++        values[j] = bit_cast<int64_t>(readWord(wordOffset + j));
+       }
+       metricValueVariant = std::move(values);
+       break;
+@@ -179,7 +194,7 @@
+     case variant_index_v<std::vector<double>, MetricValueType>: {
+       std::vector<double> values(metricDesc.size);
+       for (size_t j = 0; j < metricDesc.size; ++j) {
+-        values[j] = std::bit_cast<double>(readWord(wordOffset + j));
++        values[j] = bit_cast<double>(readWord(wordOffset + j));
+       }
+       metricValueVariant = std::move(values);
+       break;
+
+--- a/third_party/proton/csrc/lib/Profiler/Instrumentation/InstrumentationProfiler.cpp
++++ b/third_party/proton/csrc/lib/Profiler/Instrumentation/InstrumentationProfiler.cpp
+@@ -90,7 +90,7 @@
+ getUnitIdVector(const std::map<std::string, std::string> &modeOptions,
+                 size_t totalUnits) {
+   std::vector<uint32_t> unitIdVector;
+-  if (modeOptions.contains("sampling_options")) {
++  if (modeOptions.count("sampling_options")) {
+     auto &samplingOption = modeOptions.at("sampling_options");
+     auto unitIds = proton::split(samplingOption, ",");
+     for (auto uintId : unitIds) {
+@@ -118,7 +118,7 @@
+   auto config = std::make_shared<CircularLayoutParserConfig>();
+   config->scratchMemSize =
+       functionMetadata.at(functionId).getScratchMemorySize();
+-  if (modeOptions.contains("granularity") &&
++  if (modeOptions.count("granularity") &&
+       modeOptions.at("granularity") != "GRANULARITY.WARP") {
+     throw makeInvalidArgument("Only warp granularity is supported for now");
+   }
+@@ -146,7 +146,7 @@
+     const std::vector<std::pair<size_t, std::string>> &scopeIdPairs,
+     const std::vector<std::pair<size_t, size_t>> &scopeIdParentPairs,
+     const std::string &metadataPath) {
+-  if (functionScopeIdNames.contains(functionId)) {
++  if (functionScopeIdNames.count(functionId)) {
+     throw makeInvalidArgument(
+         "Duplicate function id: " + std::to_string(functionId) +
+         " for function " + functionName);
+@@ -155,7 +155,7 @@
+   for (auto &pair : scopeIdPairs) {
+     auto scopeId = pair.first;
+     auto scopeName = pair.second;
+-    if (functionScopeIdNames[functionId].contains(scopeId)) {
++    if (functionScopeIdNames[functionId].count(scopeId)) {
+       throw makeInvalidArgument(
+           "Duplicate scope id: " + std::to_string(scopeId) + " for function " +
+           functionName);
+@@ -173,7 +173,7 @@
+     std::vector<Context> reversedContexts;
+     reversedContexts.emplace_back(name);
+     auto currentId = scopeId;
+-    while (scopeIdParentMap.contains(currentId)) {
++    while (scopeIdParentMap.count(currentId)) {
+       auto parentId = scopeIdParentMap[currentId];
+       auto parentName = functionScopeIdNames[functionId].at(parentId);
+       reversedContexts.emplace_back(parentName);
+@@ -239,7 +239,7 @@
+   }
+ 
+   int64_t timeShiftCost = 0;
+-  if (modeOptions.contains("optimizations")) {
++  if (modeOptions.count("optimizations")) {
+     auto optimizations = proton::split(modeOptions.at("optimizations"), ",");
+     if (std::find(optimizations.begin(), optimizations.end(), "time_shift") !=
+         optimizations.end())
+
+--- a/third_party/proton/csrc/lib/Profiler/RocprofSDK/RocprofSDKPCSamplingSourceLocations.cpp
++++ b/third_party/proton/csrc/lib/Profiler/RocprofSDK/RocprofSDKPCSamplingSourceLocations.cpp
+@@ -229,8 +229,8 @@
+ void RocprofSDKPCSampling::tryReleaseCodeObject(uint64_t codeObjectId) {
+   metadataState.withLock([&](MetadataState &metadata) {
+     samplingState.withLock([&](SamplingState &sampling) {
+-      if (sampling.pendingCodeObjectIds.contains(codeObjectId) ||
+-          sampling.flushingCodeObjectIds.contains(codeObjectId))
++      if (sampling.pendingCodeObjectIds.count(codeObjectId) > 0 ||
++          sampling.flushingCodeObjectIds.count(codeObjectId) > 0)
+         return;
+ 
+       auto info = metadata.codeObjects.find(codeObjectId);
+
+--- a/third_party/proton/csrc/lib/Session/Session.cpp
++++ b/third_party/proton/csrc/lib/Session/Session.cpp
+@@ -66,7 +66,7 @@
+ void throwIfSessionNotInitialized(
+     const std::map<size_t, std::unique_ptr<Session>> &sessions,
+     size_t sessionId) {
+-  if (!sessions.contains(sessionId)) {
++  if (!sessions.count(sessionId)) {
+     throw makeOutOfRange("Session has not been initialized: " +
+                          std::to_string(sessionId));
+   }
+
+--- a/third_party/proton/csrc/lib/Utility/MsgPackWriter.cpp
++++ b/third_party/proton/csrc/lib/Utility/MsgPackWriter.cpp
+@@ -1,6 +1,5 @@
+ #include "Utility/MsgPackWriter.h"
+ 
+-#include <bit>
+ #include <cstring>
+ #include <limits>
+ #include <type_traits>
+@@ -78,7 +77,10 @@
+ 
+ void MsgPackWriter::packDouble(double value) {
+   out.push_back(0xcb);
+-  writeBE(out, std::bit_cast<uint64_t>(value));
++  uint64_t bits{};
++  static_assert(sizeof(bits) == sizeof(value));
++  std::memcpy(&bits, &value, sizeof(bits));
++  writeBE(out, bits);
+ }
+ 
+ void MsgPackWriter::packStr(std::string_view value) {
+
diff --git a/third_party/xla/third_party/triton/common/disable_amd_test.patch b/third_party/xla/third_party/triton/common/disable_amd_test.patch
index 838515c22bf..f42f4eb0b1c 100644
--- a/third_party/xla/third_party/triton/common/disable_amd_test.patch
+++ b/third_party/xla/third_party/triton/common/disable_amd_test.patch
@@ -4,7 +4,7 @@ would not be able to get run on presubmit anyway.
 
 --- a/python/test/unit/language/test_compile_only.py
 +++ b/python/test/unit/language/test_compile_only.py
-@@ -371,6 +371,8 @@ def test_signature_ordering():
+@@ -443,6 +443,8 @@ def test_signature_ordering():
      triton.compile(src=src, target=target)
  
  
diff --git a/third_party/xla/third_party/triton/common/enable_peer_access.patch b/third_party/xla/third_party/triton/common/enable_peer_access.patch
index 2350056239a..f624c016e33 100644
--- a/third_party/xla/third_party/triton/common/enable_peer_access.patch
+++ b/third_party/xla/third_party/triton/common/enable_peer_access.patch
@@ -33,7 +33,7 @@ new file mode 100644
 
 --- a/third_party/nvidia/backend/driver.c
 +++ b/third_party/nvidia/backend/driver.c
-@@ -979,6 +979,37 @@ cleanup:
+@@ -984,6 +984,37 @@ cleanup:
    return NULL;
  }
  
@@ -71,7 +71,7 @@ new file mode 100644
  static void _launch(int gridX, int gridY, int gridZ, int num_warps,
                      int num_ctas, int launch_cooperative_grid, int launch_pdl,
                      int shared_memory, CUstream stream, CUfunction function,
-@@ -1090,6 +1121,9 @@ bool extractPointer(void *ptr, PyObject *obj) {
+@@ -1095,6 +1126,9 @@ bool extractPointer(void *ptr, PyObject *obj) {
    if (*dev_ptr == 0) {
      return true; // valid nullptr
    }
diff --git a/third_party/xla/third_party/triton/common/f8e5m2_conversion.patch b/third_party/xla/third_party/triton/common/f8e5m2_conversion.patch
deleted file mode 100644
index f29e309817e..00000000000
--- a/third_party/xla/third_party/triton/common/f8e5m2_conversion.patch
+++ /dev/null
@@ -1,63 +0,0 @@
---- a/python/test/unit/language/test_conversions.py
-+++ b/python/test/unit/language/test_conversions.py
-@@ -338,14 +338,18 @@ def test_typeconvert_upcast(src_dtype, dst_dtype, device):
- def test_typeconvert_downcast(src_dtype, dst_dtype, rounding, max_repr, device):
- 
-     if is_cuda():
--        if src_dtype != 'float32' and torch.cuda.get_device_capability(0) < (9, 0):
--            pytest.skip("non-float32 downcast tests only supported on NVGPU with compute capability 9.0+")
--
--        if dst_dtype in ('float8e5', 'float8e4nv') and rounding == 'rtne' and torch.cuda.get_device_capability(0) < (9, 0):
--            pytest.skip(f"{dst_dtype} downcast with RTNE rounding tests only supported on NVGPU with compute capability 9.0+")
--
--        if dst_dtype in ('float8e5b16', 'float8e4b8') and rounding == 'rtne':
--            pytest.skip(f"{dst_dtype} downcast with RTNE rounding tests only supported on AMDGPU CDNA3")
-+        if dst_dtype in ('float8e5b16', 'float8e4b8'):
-+            pytest.skip(f"{dst_dtype} downcast with RTNE rounding tests only supported on AMDGPU MI300")
-+        if torch.cuda.get_device_capability(0) < (9, 0):
-+            match (src_dtype, dst_dtype, rounding):
-+                case ('float16', 'float8e5', 'rtne'):
-+                    ...
-+                case (_, ty, 'rtne') if ty in ('float8e5', 'float8e4nv'):
-+                    pytest.skip(f"{dst_dtype} downcast with RTNE rounding tests only supported on NVGPU with compute capability 9.0+")
-+                case (ty, _, _) if ty != 'float32':
-+                    pytest.skip("non-float32 downcast tests only supported on NVGPU with compute capability 9.0+")
-+                case _:
-+                    ...
- 
-     if is_hip():
-         if dst_dtype in FP8_DTYPES and is_hip_rdna3():
-
---- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp
-+++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp
-@@ -30,15 +30,20 @@ struct Fp8ConversionDesc {
- static const Fp8ConversionDesc Fp16_to_Fp8E5M2_RTNE(bool hasNativeFP) {
-   Fp8ConversionDesc ret;
-   if (!hasNativeFP) {
--    ret = {"{                            \n"
--           ".reg .b32 a<2>;              \n"
--           "and.b32 a0, $1, 0xfffefffe;  \n"   // a0 &= 0xfffefffe
--           "and.b32 a1, $2, 0xfffefffe;  \n"   // (strip lowest bit)
--           "add.u32 a0, a0, 0x00800080;  \n"   // a0 += 0x00800080
--           "add.u32 a1, a1, 0x00800080;  \n"   // (round to nearest)
--           "prmt.b32 $0, a0, a1, 0x7531; \n\t" // output = a1a0
--           "}",
--           32, 32, 4};
-+    ret = {
-+        "{                            \n"
-+        ".reg .b32 a<2>;              \n"
-+        "and.b32 a0, $1, 0x01000100;  \n"    // a0 = $1 & 0x01000100
-+        "and.b32 a1, $2, 0x01000100;  \n"    // (least significant result bit)
-+        "shr.b32 a0, a0, 8;           \n"    // a0 >>= 8
-+        "shr.b32 a1, a1, 8;           \n"    // (shift the lsb)
-+        "add.u32 a0, a0, 0x007f007f;  \n"    // a0 += 0x007f007f
-+        "add.u32 a1, a1, 0x007f007f;  \n"    // (add rounding base)
-+        "add.u32 a0, a0, $1;          \n"    // res = $1 + lsb + 0x7f
-+        "add.u32 a1, a1, $2;          \n"    // (round to nearest)
-+        "prmt.b32 $0, a0, a1, 0x7531; \n\t"  // output = a1a0
-+        "}",
-+        32, 32, 4};
-   } else {
-     ret = {"cvt.rn.satfinite.e5m2x2.f16x2 $0, $1; \n\t", 32, 16, 2};
-   }
-
diff --git a/third_party/xla/third_party/triton/common/fix_load_ops.patch b/third_party/xla/third_party/triton/common/fix_load_ops.patch
index 9246d870bcf..18e7707232c 100644
--- a/third_party/xla/third_party/triton/common/fix_load_ops.patch
+++ b/third_party/xla/third_party/triton/common/fix_load_ops.patch
@@ -5,7 +5,7 @@ SM90+ ptxas miscompilation for cp.async with L2::cache_hint.
 
 --- a/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp
 +++ b/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp
-@@ -709,6 +709,9 @@
+@@ -706,6 +706,9 @@ void LayoutRematerialization::rewriteSlice(
    for (const auto &[value, encoding] : layout) {
      if (Value remat = existingRemats.lookup({value, encoding}))
        mapping.map(value, remat);
@@ -15,7 +15,7 @@ SM90+ ptxas miscompilation for cp.async with L2::cache_hint.
    }
  
    SetVector<Operation *> opsToRewrite;
-@@ -1238,6 +1241,8 @@
+@@ -1236,6 +1239,8 @@ bool LayoutRematerialization::hoistConvertDotOperand(
      ConvertLayoutOp convertOp) {
    auto targetType = convertOp.getType();
    // The pass is targeted to MMA dot operands
@@ -24,89 +24,32 @@ SM90+ ptxas miscompilation for cp.async with L2::cache_hint.
  
    auto canBePipelined = [&](ConvertLayoutOp convertOp) {
      // FIXME: Check that the parent is a for loop
-@@ -1335,6 +1340,9 @@
-                                                 newType, loadOp->getResult(0));
+@@ -1334,6 +1339,9 @@ bool LayoutRematerialization::hoistConvertDotOperand(
      mapping.map(loadOp->getResult(0), newConvertOp.getResult());
    }
-+
+ 
 +  if (innerSlice.empty() && existingRemats.empty())
 +    return false;
- 
++
    LLVM_DEBUG({
      DBGS() << "  Hoisting " << convertOp << '\n';
-
---- a/test/Conversion/tritongpu_to_llvm_hopper.mlir
-+++ b/test/Conversion/tritongpu_to_llvm_hopper.mlir
-@@ -714,3 +714,21 @@
- }
- 
- }
-+
-+// -----
-+
-+#blocked = #ttg.blocked<{sizePerThread = [4], threadsPerWarp = [32], warpsPerCTA = [4], order = [0]}>
-+#shared1D = #ttg.swizzled_shared<{vec = 2, perPhase = 1, maxPhase = 8, order = [0]}>
-+#smem = #ttg.shared_memory
-+#async_l2_policy = #ttng.cache_policy<l1 = evict_last, l2_primary = evict_last, l2_secondary = evict_unchanged, l2_fraction = 8.200000e-01 : f32, l2_prefetch_size = 256 : i32>
-+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:90", "ttg.threads-per-warp" = 32 : i32} {
-+  // CHECK-LABEL: async_cp_detailed_cache_policy_sm90
-+  tt.func @async_cp_detailed_cache_policy_sm90(%v: tensor<256x!tt.ptr<f16>, #blocked>, %smem: !ttg.memdesc<256xf16, #shared1D, #smem, mutable>) {
-+    // CHECK: createpolicy.fractional.L2::evict_last.b64 $0, {{0\.8.*}};
-+    // CHECK: lop3.b32 $0, $1, $1, $1, 0xf0;
-+    // CHECK: cp.async.ca.shared.global.L2::256B.L2::cache_hint
-+    %0 = ttg.async_copy_global_to_local %v, %smem {cachePolicy = #async_l2_policy, contiguity = 4 : i32} : tensor<256x!tt.ptr<f16>, #blocked> -> !ttg.memdesc<256xf16, #shared1D, #smem, mutable>
-+    tt.return
-+  }
-+}
-+
-
---- a/test/TritonGPU/combine.mlir
-+++ b/test/TritonGPU/combine.mlir
-@@ -4627,3 +4627,28 @@
-     tt.return %d1 : tensor<16x16xf32, #mma>
-   }
- }
-+
-+// -----
-+
-+// CHECK-LABEL: @hoist_dot_operand_ignore_non_dot_convert_in_loop
-+// CHECK-NOT: ttg.convert_layout
-+// CHECK: scf.for
-+// CHECK: tt.dot
-+// CHECK: tt.return
-+#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [2, 16], warpsPerCTA = [1, 1], order = [1, 0]}>
-+#mma = #ttg.nvidia_mma<{versionMajor = 2, versionMinor = 0, warpsPerCTA = [1, 1], instrShape = [16, 8]}>
-+#dot0 = #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 2}>
-+#dot1 = #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 2}>
-+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32, ttg.target = "cuda:89", "ttg.threads-per-warp" = 32 : i32} {
-+  tt.func public @hoist_dot_operand_ignore_non_dot_convert_in_loop(%a: tensor<16x16xf32, #dot0>, %b: tensor<16x16xf32, #dot1>, %acc: tensor<16x16xf32, #mma>, %init: tensor<16x16xf32, #blocked>, %n: i32) -> (tensor<16x16xf32, #mma>, tensor<16x16xf32, #blocked>) {
-+    %c0 = arith.constant 0 : i32
-+    %c1 = arith.constant 1 : i32
-+    %res:2 = scf.for %i = %c0 to %n step %c1 iter_args(%acc_iter = %acc, %val_iter = %init) -> (tensor<16x16xf32, #mma>, tensor<16x16xf32, #blocked>) : i32 {
-+      %cvt = ttg.convert_layout %val_iter : tensor<16x16xf32, #blocked> -> tensor<16x16xf32, #blocked>
-+      %mul = arith.mulf %cvt, %val_iter : tensor<16x16xf32, #blocked>
-+      %dot = tt.dot %a, %b, %acc_iter, inputPrecision = tf32 : tensor<16x16xf32, #dot0> * tensor<16x16xf32, #dot1> -> tensor<16x16xf32, #mma>
-+      scf.yield %dot, %mul : tensor<16x16xf32, #mma>, tensor<16x16xf32, #blocked>
-+    }
-+    tt.return %res#0, %res#1 : tensor<16x16xf32, #mma>, tensor<16x16xf32, #blocked>
-+  }
-+}
+     for (Value v : innerSlice)
 
 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
-@@ -1110,6 +1110,7 @@
+@@ -1060,7 +1060,7 @@ struct AsyncCopyGlobalToLocalOpConversion
          l2PrefetchSizeAttr ? l2PrefetchSizeAttr.getInt() : 0;
  
      auto emitCpAsync = [&b, threadPred, ptrTy, l2PolicyReg, l2PrefetchSize,
-+                        computeCapability,
+-                        cacheModifier = cachePolicy->modifier,
++                        computeCapability, cacheModifier = cachePolicy->modifier,
                          hasMask =
                              bool(llMask)](RewriterBase &rewriter, Location loc,
                                            ArrayRef<Value> vals, Value shmemAddr,
-@@ -1128,6 +1129,24 @@
-       auto structElem = vals[startIdx];
+@@ -1083,6 +1083,24 @@ struct AsyncCopyGlobalToLocalOpConversion
        auto srcElem = b.extract_val(ptrTy, structElem, 0);
        auto maskElem = b.extract_val(i1_ty, structElem, 1);
-+
+ 
 +      if (l2PolicyReg && computeCapability >= 90) {
 +        // On SM90+, ptxas miscompiles cp.async with L2::cache_hint when the
 +        // shared memory address is matched as [R_thread + UR_uniform] (opcode
@@ -124,7 +67,8 @@ SM90+ ptxas miscompilation for cp.async with L2::cache_hint.
 +            addrBuilder.newConstantOperand(0xF0));
 +        shmemAddr = addrBuilder.launch(rewriter, loc, i32_ty);
 +      }
- 
++
        PTXBuilder ptxBuilder;
        auto &copyAsyncOp =
+           *ptxBuilder.create<PTXCpAsyncLoadInstr>(srcCacheModifier);
 
diff --git a/third_party/xla/third_party/triton/common/fp64_large_k.patch b/third_party/xla/third_party/triton/common/fp64_large_k.patch
deleted file mode 100644
index d9aea906175..00000000000
--- a/third_party/xla/third_party/triton/common/fp64_large_k.patch
+++ /dev/null
@@ -1,525 +0,0 @@
-Cherry-pick of https://github.com/triton-lang/triton/pull/11664
-(a0af442df17c8acff35b7488deae704277da46d2) and
-https://github.com/triton-lang/triton/pull/11691
-(696c45880d8123c94adbf69e57297403c134edb4). Can be deleted on the next Triton
-integration.
-
---- a/lib/Dialect/TritonGPU/IR/Dialect.cpp
-+++ b/lib/Dialect/TritonGPU/IR/Dialect.cpp
-@@ -3290,6 +3290,13 @@
-       // Verify that the operands are supported on the selected MMA version.
-       if (!supportMMA(dotOp, mmaResEncoding.getVersionMajor()))
-         return op->emitError("unsupported MMA version");
-+      // MMAv2 distributes K over four lanes, each owning kWidth contiguous
-+      // elements. A smaller K repeats elements across lanes, which a static
-+      // register permutation cannot resolve.
-+      auto aType = cast<RankedTensorType>(dotOp.getA().getType());
-+      if (mmaResEncoding.isAmpere() &&
-+          getShapePerCTA(aType).back() < 4 * aEncoding.getKWidth())
-+        return op->emitError("MMA operand layout requires K >= 4 * kWidth");
-     }
- 
-     return verifyWmmaCGACompatibility(op, aEncoding, bEncoding,
-
---- a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
-+++ b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
-@@ -574,6 +574,22 @@
-     } else {
-       int minBitwidth =
-           std::min(computeOrigBitWidth(a), computeOrigBitWidth(b));
-+      // Let K = getShapePerCTA(oldAType).back() and computeBitwidth be the
-+      // dot operand bitwidth. Since kWidth = max(32 / minBitwidth, 1),
-+      // MMAv2's four K lanes require 4 * max(32 / minBitwidth, 1) <= K.
-+      // We have:
-+      //   (a) K >= 4;
-+      //   (b) minBitwidth >= 128 / K.
-+      //
-+      // Native operand packing is preserved by the invariant:
-+      //   (c) minBitwidth <= computeBitwidth.
-+      // Raising it for (b) could violate (c) only if 128 / K > computeBitwidth,
-+      // hence K < 128 / computeBitwidth. This contradicts the CUDA frontend's
-+      // requirement K >= 256 / computeBitwidth, which also implies (a).
-+      minBitwidth = std::max<int64_t>(minBitwidth,
-+                                      4 * 32 / getShapePerCTA(oldAType).back());
-+      assert(minBitwidth <= oldAType.getElementTypeBitWidth() &&
-+             "minBitwidth must not exceed the dot operand bitwidth");
-       a = convertDotOperandForMMA(a, 0, minBitwidth, mmaResult.newRetType,
-                                   rewriter);
-       b = convertDotOperandForMMA(b, 1, minBitwidth, mmaResult.newRetType,
-
---- a/python/test/regression/test_cast_matmul.py
-+++ b/python/test/regression/test_cast_matmul.py
-@@ -15,7 +15,7 @@
- 
- pytestmark = pytest.mark.enable_warmup(min_capability=9)
- 
--input_dtypes = ["bfloat16", "float16", "float32"]
-+input_dtypes = ["bfloat16", "float16", "float32", "float64"]
- if is_cuda():
-     input_dtypes += ["int8", "float8_e5m2"]
-     cc = torch.cuda.get_device_capability(0)
-@@ -29,7 +29,7 @@
-         "float8_e4m3fnuz",
-     ]
- 
--out_dtypes = ["float16", "float32"]
-+out_dtypes = ["float16", "float32", "float64"]
- 
- 
- @triton.jit
-@@ -58,7 +58,7 @@
-     # pointers
-     A = A + (ram[:, None] * stride_am + rk[None, :] * stride_ak)
-     B = B + (rk[:, None] * stride_bk + rbn[None, :] * stride_bn)
--    acc_dtype = tl.float16 if compute_dtype == tl.float16 and C.dtype.element_ty == tl.float16 else tl.float32
-+    acc_dtype = tl.float32 if compute_dtype == tl.float16 and C.dtype.element_ty != tl.float16 else compute_dtype
-     acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=acc_dtype)
-     for k in range(0, tl.cdiv(K, BLOCK_K)):
-         k_remaining = K - k * BLOCK_K
-@@ -79,18 +79,23 @@
- 
- @pytest.mark.parametrize("M, K, N, BLOCK_K, BLOCK_M, BLOCK_N, w_dtype, x_dtype, out_dtype",
-                          [(M, K, N, BLOCK_K, BLOCK_M, BLOCK_N, w, x, o)  #
--                          for BLOCK_K in [16, 32, 64]  #
--                          for BLOCK_M in [16, 64]  #
--                          for BLOCK_N in [16, 64, 128]  #
-+                          for BLOCK_K, BLOCK_M, BLOCK_N in ([(k, m, n)  #
-+                                                             for k in [16, 32, 64]  #
-+                                                             for m in [16, 64]  #
-+                                                             for n in [16, 64, 128]] + [(32, 32, 32)])
-                           for (M, K, N) in [(768, 768, 1024)]  #
-                           for w in input_dtypes
-                           for x in input_dtypes  #
--                          for o in out_dtypes])
-+                          for o in out_dtypes] +
-+                         ([(16, 8, 16, 8, 16, 16, "int8", "float32", "float32"),
-+                           (16, 8, 16, 8, 16, 16, "float32", "int8", "float32")] if is_cuda() else []))
- def test_cast_matmul(M, K, N, BLOCK_K, BLOCK_M, BLOCK_N, w_dtype, x_dtype, out_dtype, device):
-     if is_hip() and (BLOCK_K, BLOCK_M, BLOCK_N) in ((64, 64, 128), (64, 16, 128)):
-         pytest.skip("skip as they run out of shared memory")
-     if x_dtype == w_dtype:
-         pytest.skip("skip the same input dtype")
-+    if "float64" in (x_dtype, w_dtype) and any(t.startswith("float8") for t in (x_dtype, w_dtype)):
-+        pytest.skip("FP8 to FP64 conversion is not supported")
-     x_dtype: torch.dtype = getattr(torch, x_dtype)
-     w_dtype: torch.dtype = getattr(torch, w_dtype)
- 
-@@ -106,9 +111,8 @@
-         # a holds the larger dtype
-         if a_dtype.itemsize < b_dtype.itemsize:
-             a_dtype, b_dtype = b_dtype, a_dtype
--        # float64 matmul is not supported by triton
-         if a_dtype == torch.float64:
--            return torch.float32
-+            return torch.float64
-         # If they are both 1 byte or float16 and (1 byte or float16)
-         if a_dtype.itemsize == 1 or (a_dtype == torch.float16 and b_dtype != torch.bfloat16):
-             return torch.float16
-
---- a/test/TritonGPU/accelerate-matmul.mlir
-+++ b/test/TritonGPU/accelerate-matmul.mlir
-@@ -1153,3 +1153,40 @@
-     tt.return %0 : tensor<128x256xf32, #blocked>
-   }
- }
-+
-+// -----
-+
-+// CHECK: #[[$MMA:.+]] = #ttg.nvidia_mma<{versionMajor = 2,
-+#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [4, 8], warpsPerCTA = [4, 1], order = [1, 0]}>
-+module attributes {"ttg.target" = "cuda:80", "ttg.num-warps" = 4 : i32} {
-+  // CHECK-LABEL: fp32_small_k_i8_upcast
-+  tt.func @fp32_small_k_i8_upcast(%pa: tensor<16x8x!tt.ptr<i8>, #blocked>,
-+      %b: tensor<8x16xf32, #ttg.dot_op<{opIdx = 1, parent = #blocked}>>) -> tensor<16x16xf32, #blocked> {
-+    %a_i8 = tt.load %pa : tensor<16x8x!tt.ptr<i8>, #blocked>
-+    %a_f32 = arith.sitofp %a_i8 : tensor<16x8xi8, #blocked> to tensor<16x8xf32, #blocked>
-+    %a = ttg.convert_layout %a_f32 : tensor<16x8xf32, #blocked> -> tensor<16x8xf32, #ttg.dot_op<{opIdx = 0, parent = #blocked}>>
-+    %c = arith.constant dense<0.0> : tensor<16x16xf32, #blocked>
-+    // The i8 load suggests kWidth=4, but K=8 only has two elements per K lane.
-+    // CHECK: tt.dot {{.*}} : tensor<16x8xf32, #ttg.dot_op<{opIdx = 0, parent = #[[$MMA]], kWidth = 2}>> * tensor<8x16xf32, #ttg.dot_op<{opIdx = 1, parent = #[[$MMA]], kWidth = 2}>>
-+    %d = tt.dot %a, %b, %c, inputPrecision = tf32 : tensor<16x8xf32, #ttg.dot_op<{opIdx = 0, parent = #blocked}>> * tensor<8x16xf32, #ttg.dot_op<{opIdx = 1, parent = #blocked}>> -> tensor<16x16xf32, #blocked>
-+    tt.return %d : tensor<16x16xf32, #blocked>
-+  }
-+}
-+
-+// -----
-+
-+// CHECK: #[[$MMA:.+]] = #ttg.nvidia_mma<{versionMajor = 2,
-+#blocked = #ttg.blocked<{sizePerThread = [1, 1], threadsPerWarp = [4, 8], warpsPerCTA = [4, 1], order = [1, 0]}>
-+module attributes {"ttg.target" = "cuda:80", "ttg.num-warps" = 4 : i32} {
-+  // CHECK-LABEL: fp64_small_k_f16_upcast
-+  tt.func @fp64_small_k_f16_upcast(%a: tensor<16x4xf64, #ttg.dot_op<{opIdx = 0, parent = #blocked}>>,
-+      %pb: tensor<4x16x!tt.ptr<f16>, #blocked>) -> tensor<16x16xf64, #blocked> {
-+    %b_f16 = tt.load %pb : tensor<4x16x!tt.ptr<f16>, #blocked>
-+    %b_f64 = arith.extf %b_f16 : tensor<4x16xf16, #blocked> to tensor<4x16xf64, #blocked>
-+    %b = ttg.convert_layout %b_f64 : tensor<4x16xf64, #blocked> -> tensor<4x16xf64, #ttg.dot_op<{opIdx = 1, parent = #blocked}>>
-+    %c = arith.constant dense<0.0> : tensor<16x16xf64, #blocked>
-+    // CHECK: tt.dot {{.*}} : tensor<16x4xf64, #ttg.dot_op<{opIdx = 0, parent = #[[$MMA]], kWidth = 1}>> * tensor<4x16xf64, #ttg.dot_op<{opIdx = 1, parent = #[[$MMA]], kWidth = 1}>>
-+    %d = tt.dot %a, %b, %c : tensor<16x4xf64, #ttg.dot_op<{opIdx = 0, parent = #blocked}>> * tensor<4x16xf64, #ttg.dot_op<{opIdx = 1, parent = #blocked}>> -> tensor<16x16xf64, #blocked>
-+    tt.return %d : tensor<16x16xf64, #blocked>
-+  }
-+}
-
---- a/test/Conversion/tritongpu_to_llvm.mlir
-+++ b/test/Conversion/tritongpu_to_llvm.mlir
-@@ -3433,3 +3433,32 @@
-     tt.return
-   }
- }
-+
-+// -----
-+
-+#mma = #ttg.nvidia_mma<{versionMajor = 2, versionMinor = 0, warpsPerCTA = [1, 1, 1], instrShape = [1, 8, 8]}>
-+#a = #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 4}>
-+#b = #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 4}>
-+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32, ttg.target = "cuda:80", "ttg.threads-per-warp" = 32 : i32} {
-+  // FP64 fragments already have the MMA register order, including across batches.
-+  // CHECK-LABEL: llvm.func @mma_fp64_large_k_batched(
-+  // CHECK-DAG: %[[A0:.*]] = llvm.extractvalue %arg0[0]
-+  // CHECK-DAG: %[[A8:.*]] = llvm.extractvalue %arg0[8]
-+  // CHECK-DAG: %[[AC0:.*]] = llvm.bitcast %[[A0]] : f64 to f64
-+  // CHECK-DAG: %[[AC8:.*]] = llvm.bitcast %[[A8]] : f64 to f64
-+  // CHECK-DAG: %[[AV0:.*]] = llvm.insertelement %[[AC0]], {{.*}} : vector<1xf64>
-+  // CHECK-DAG: %[[AV8:.*]] = llvm.insertelement %[[AC8]], {{.*}} : vector<1xf64>
-+  // CHECK-DAG: %[[B0:.*]] = llvm.extractvalue %arg1[0]
-+  // CHECK-DAG: %[[B8:.*]] = llvm.extractvalue %arg1[8]
-+  // CHECK-DAG: %[[BC0:.*]] = llvm.bitcast %[[B0]] : f64 to f64
-+  // CHECK-DAG: %[[BC8:.*]] = llvm.bitcast %[[B8]] : f64 to f64
-+  // CHECK-DAG: %[[BV0:.*]] = llvm.insertelement %[[BC0]], {{.*}} : vector<1xf64>
-+  // CHECK-DAG: %[[BV8:.*]] = llvm.insertelement %[[BC8]], {{.*}} : vector<1xf64>
-+  // CHECK: llvm.inline_asm{{.*}}mma.sync.aligned.m8n8k4{{.*}}%[[BV0]], %[[AV0]] :
-+  // CHECK: llvm.inline_asm{{.*}}mma.sync.aligned.m8n8k4{{.*}}%[[BV8]], %[[AV8]] :
-+  tt.func @mma_fp64_large_k_batched(%a: tensor<2x16x16xf64, #a>, %b: tensor<2x16x16xf64, #b>) {
-+    %c = arith.constant dense<0.0> : tensor<2x16x16xf64, #mma>
-+    %d = tt.dot %a, %b, %c : tensor<2x16x16xf64, #a> * tensor<2x16x16xf64, #b> -> tensor<2x16x16xf64, #mma>
-+    tt.return
-+  }
-+}
-
---- a/test/TritonGPU/invalid.mlir
-+++ b/test/TritonGPU/invalid.mlir
-@@ -914,3 +914,31 @@
- // expected-error @below {{alignment must be specified outside of the linear layout braces}}
- #shared = #ttg.shared_linear<{offset = [[0, 1], [0, 2], [1, 0], [2, 0]], block = [], alignment = 16}>
- !alignment_in_layout = !ttg.memdesc<4x4xf32, #shared, #ttg.shared_memory>
-+
-+// -----
-+
-+#mma = #ttg.nvidia_mma<{versionMajor = 2, versionMinor = 0, warpsPerCTA = [1, 1], instrShape = [8, 8]}>
-+#a = #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 2}>
-+#b = #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 2}>
-+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
-+  tt.func @mma_fp64_large_k_replicated_lanes(%a: tensor<8x4xf64, #a>, %b: tensor<4x8xf64, #b>) {
-+    %c = arith.constant dense<0.0> : tensor<8x8xf64, #mma>
-+    // expected-error@+1 {{MMA operand layout requires K >= 4 * kWidth}}
-+    %d = tt.dot %a, %b, %c : tensor<8x4xf64, #a> * tensor<4x8xf64, #b> -> tensor<8x8xf64, #mma>
-+    tt.return
-+  }
-+}
-+
-+// -----
-+
-+#mma = #ttg.nvidia_mma<{versionMajor = 2, versionMinor = 0, warpsPerCTA = [1, 1], instrShape = [16, 8]}>
-+#a = #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 8}>
-+#b = #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 8}>
-+module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
-+  tt.func @mma_fp32_large_k_replicated_lanes(%a: tensor<16x16xf32, #a>, %b: tensor<16x8xf32, #b>) {
-+    %c = arith.constant dense<0.0> : tensor<16x8xf32, #mma>
-+    // expected-error@+1 {{MMA operand layout requires K >= 4 * kWidth}}
-+    %d = tt.dot %a, %b, %c, inputPrecision = tf32 : tensor<16x16xf32, #a> * tensor<16x8xf32, #b> -> tensor<16x8xf32, #mma>
-+    tt.return
-+  }
-+}
-
---- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv2.cpp
-+++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv2.cpp
-@@ -4,7 +4,9 @@
- #include "triton/Conversion/TritonGPUToLLVM/Utility.h"
- #include "triton/Dialect/Triton/IR/Dialect.h"
- #include "triton/Dialect/TritonGPU/IR/Attributes.h"
-+#include "triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h"
- #include "triton/Dialect/TritonInstrument/IR/Dialect.h"
-+#include "triton/Tools/LayoutUtils.h"
- #include "llvm/ADT/SmallVector.h"
- 
- using namespace mlir;
-@@ -52,9 +54,9 @@
-   return elems;
- }
- 
--// The number of i32 registers owned by each thread along m, n, k dimensions.
--// For example, for m16n8k32 with i8 inputs, a thread owns 2, 1, and 2 registers
--// along m, n, k respectively.
-+// Per-thread MMA operand register counts along m, n, k. Registers hold 32 bits,
-+// or 64 bits for fp64. For m16n8k32 with i8 inputs, the counts are {2, 1, 2},
-+// with four i8 elements per register.
- struct NumRegisters {
-   int m;
-   int n;
-@@ -74,186 +76,95 @@
-     const LLVMTypeConverter *typeConverter, Location loc,
-     ConversionPatternRewriter &rewriter, Value value, int batch, int repOuter,
-     int repK, RankedTensorType type, const NumRegisters &numRegisters) {
-+  auto ctx = rewriter.getContext();
-   auto b = TritonLLVMOpBuilder(loc, rewriter);
-   auto elems = unpackTensorElements(loc, value, rewriter, type);
-   auto eltTy = typeConverter->convertType(type.getElementType());
--  int offset{};
-+  auto bitwidth = eltTy.getIntOrFloatBitWidth();
-+  int numElemsPerVec = std::max(32 / bitwidth, 1u);
-+
-+  auto dot = cast<DotOperandEncodingAttr>(type.getEncoding());
-+  int kWidth = dot.getKWidth();
-+  int outerRegs = dot.getOpIdx() == 0 ? numRegisters.m : numRegisters.n;
-+
-+  auto reg = str_attr("register");
-+  auto srcLayout = triton::gpu::toLinearLayout(type);
-+  auto removeBroadcast = actionRemoveBroadcastedRegs(srcLayout);
-+  srcLayout = removeBroadcast.apply(srcLayout);
-+  elems = removeBroadcast.apply(elems);
-+
-+  // Size the source groups using only the registers that remain.
-+  auto elemsPerThread = triton::gpu::LinearEncodingTrait::basesPerDim(
-+      srcLayout, reg, /*skipBroadcast=*/true);
-+  auto order =
-+      getOrderForDotOperand(dot.getOpIdx(), type.getRank(), /*kContig=*/true);
-+  int outerSize = std::min<int>(outerRegs, elemsPerThread[order[1]]);
-+  int kTileSize =
-+      std::min<int>(numRegisters.k, elemsPerThread[order[0]] / kWidth);
-+
-+  int size = elems.size();
-+  // Reorder scalar entries in elems into MMA operand registers:
-+  //   element: scalar within a packed operand (32 bits, or 64 bits for fp64).
-+  //   k: packed operand position within a contiguous kWidth group.
-+  //   outer: operand along M for A, or N for B, within one MMA.
-+  //   kTiles: K operand from another kWidth group, for the same MMA.
-+  //   tile: group of MMAs, repeated along K, then M/N, then batch.
-+  //
-+  // Example: A, fp16, kWidth=4, shape=[16,64], one warp. Entries index the
-+  // unique elems for one lane. Each pair is (element=0, element=1), packed
-+  // into one 32-bit operand:
-+  //
-+  //                         kTiles=0                 kTiles=1
-+  //                      k=0       k=1            k=0       k=1
-+  //   tile=0  outer=0    (0,1)     (2,3)           (8,9)    (10,11)
-+  //           outer=1    (4,5)     (6,7)          (12,13)   (14,15)
-+  //   tile=1  outer=0   (16,17)   (18,19)         (24,25)   (26,27)
-+  //           outer=1   (20,21)   (22,23)         (28,29)   (30,31)
-+  //
-+  // Read one k column from both outer rows and both kTiles groups per MMA:
-+  //   tile=0, k=0 -> [(0,1), (4,5),  (8,9),  (12,13)]
-+  //   tile=0, k=1 -> [(2,3), (6,7), (10,11), (14,15)]
-+  //   tile=1 uses the same pattern, adding 16 to each scalar index.
-+  auto mapping =
-+      LinearLayout::identity1D(size, reg, reg)
-+          .reshapeIns(
-+              {{str_attr("element"), numElemsPerVec},
-+               {str_attr("k"), kWidth / numElemsPerVec},
-+               {str_attr("outer"), outerSize},
-+               {str_attr("kTiles"), kTileSize},
-+               {str_attr("tile"), size / (kWidth * outerSize * kTileSize)}})
-+          .transposeIns({str_attr("element"), str_attr("outer"),
-+                         str_attr("kTiles"), str_attr("k"), str_attr("tile")})
-+          .reshapeIns({{reg, size}});
-+
-+  SmallVector<Value> mmaElems;
-+  for (int i = 0; i < size; ++i)
-+    mmaElems.push_back(elems[mapping.apply({{reg, i}})[0].second]);
-+
-+  // Restore only the repeated slots required by the native MMA operands.
-+  dot = DotOperandEncodingAttr::get(ctx, dot.getOpIdx(), dot.getParent(),
-+                                    numElemsPerVec);
-+  elems =
-+      broadcastAs(mmaElems, triton::gpu::toLinearLayout(type.getShape(), dot));
-+
-   ValueTableV2 vals;
--  auto bitwidth = eltTy.getIntOrFloatBitWidth();
--  auto numElemsPerVec = std::max(32 / bitwidth, 1u);
--  auto vecTy = vec_ty(eltTy, numElemsPerVec);
--
-+  int offset = 0;
-   auto packVec = [&](std::array<int, 3> dstIdx) {
--    Value vec = b.undef(vecTy);
--    for (auto i = 0; i < numElemsPerVec; ++i) {
-+    Value vec = b.undef(vec_ty(eltTy, numElemsPerVec));
-+    for (int i = 0; i < numElemsPerVec; ++i)
-       vec = b.insert_element(vec, b.bitcast(elems[offset + i], eltTy),
-                              b.i32_val(i));
--    }
--    if (bitwidth == 64) {
--      vals[dstIdx] = vec;
--    } else {
--      vals[dstIdx] = b.bitcast(vec, i32_ty);
--    }
-+    vals[dstIdx] = bitwidth == 64 ? vec : b.bitcast(vec, i32_ty);
-     offset += numElemsPerVec;
-   };
- 
--  auto dot = cast<DotOperandEncodingAttr>(type.getEncoding());
--  auto kWidth = dot.getKWidth();
--  auto largeK = bitwidth * kWidth > std::max(32u, bitwidth);
--
--  assert((bitwidth != 64 || largeK == false) &&
--         "Currently fp64 don't support largeK MMA");
--
--  if (largeK) {
--    // For layouts with a large K dimension, the original register layout needs
--    // to be divided into multiple MMAs, where each MMA has contiguous 32 bits
--    // along the K dimension per thread.
--    // Using kWidth = 8 and bitwidth = 2 as an example,
--    // we split the MMA into 4 sub-MMAs, each with a stride 4 x 32-bit along the
--    // K dimension.
--    llvm::SmallVector<unsigned> si;
--    auto kIters = kWidth / (std::max(32 / bitwidth, 1u));
--
--    if (dot.getOpIdx() == 0) {
--      // Original register layout:
--      //
--      //   [0, 1, 2, 3, 4, 5, 6, 7], [16, 17, 18, 19, 20, 21, 22, 23, 23]
--      //   [8, 9, 10, 11, 12, 13, 14, 15], [24, 25, 26, 27, 28, 29, 30, 31]
--      //
--      // Each element in the layout is a single bf16.
--      //
--      // To derive four independent MMA operations, a stride of 4 is applied to
--      // the original register layout:
--      //
--      //  1st MMA: [[0, 1], [8, 9], [16, 17], [24, 25]]
--      //  2nd MMA: [[2, 3], [10, 11], [18, 19], [26, 27]]
--      //  3rd MMA: [[4, 5], [12, 13], [20, 21], [28, 29]]
--      //  4th MMA: [[6, 7], [14, 15], [22, 23], [30, 31]]
--      if (kIters <= repK) {
--        for (size_t kRep = 0; kRep < kWidth / numElemsPerVec; ++kRep)
--          for (size_t tile = 0; tile < 4; ++tile)
--            for (size_t e = 0; e < numElemsPerVec; ++e) {
--              si.push_back(kRep * numElemsPerVec + tile * kWidth + e);
--            }
--      } else {
--        // Suppose kWidth=4 and type=fp32, so numElemsPerVec=1.
--        // Each tile of the dot operand layout has a size of 16x32.
--        // However, if the triton tensor size is 16x16, elements along the k
--        // dimension are duplicated. Within each tile, each register
--        // contains 2x8 elements arranged as follows:
--        //
--        //       tile0/0           tile0/1
--        //   |<--kWidth=4-->|   |<--kWidth-->|
--        //   |<-mmaWidth=2->|
--        //   [0,  1,  2,  3]    [0,  1,  2,  3]
--        //   [4,  5,  6,  7]    [4,  5,  6,  7]
--        //
--        // tile0/1 replicates the elements in tile0/0 along the k dimension.
--        // For a tensor size of 32x32, the next tile on the m dimension is as
--        // follows:
--        //
--        //       tile1/0              tile1/1
--        //   |<--kWidth-->|       |<--kWidth-->|
--        //   [8,  9, 10, 11],     [8,  9, 10, 11]
--        //   [12, 13, 14, 15],    [12, 13, 14, 15]
--        //
--        // Within a single tile, we can perform two MMAs, and the
--        // resulting register layout for each MMA is as follows:
--        //
--        //   1st MMA: [0, 4, 1, 5]
--        //   2nd MMA: [2, 6, 3, 7]
--        //   3rd MMA: [8, 12, 9, 13]
--        //   4th MMA: [10, 14, 11, 15]
--        //
--        // Additionally, we should reorder the elements by moving the duplicated
--        // elements to the end.  In the example above, we convert the order from
--        // tile0/0, tile0/1, tile1/0, tile1/1 to tile0/0, tile1/0, tile0/1,
--        // tile1/1, so that only the first two tiles will be used in the
--        // computation.
--        size_t elemsPerTile = 2 * 2 * kWidth;
--        size_t elemsPerMma = 2 * 2 * numElemsPerVec;
--        size_t mmaWidth = kWidth / numElemsPerVec / 2;
--        size_t repMma = elemsPerTile / (mmaWidth * elemsPerMma);
--        for (size_t rep = 0; rep < repMma; ++rep)
--          for (size_t tile = 0; tile < elems.size() / elemsPerTile; ++tile)
--            for (size_t mmaKWidth = 0; mmaKWidth < mmaWidth; ++mmaKWidth)
--              for (size_t kTile = 0; kTile < 2; ++kTile)
--                for (size_t mTile = 0; mTile < 2; ++mTile)
--                  for (size_t e = 0; e < numElemsPerVec; ++e) {
--                    si.push_back(rep * mmaWidth * elemsPerMma +
--                                 mmaKWidth * 2 * numElemsPerVec +
--                                 tile * elemsPerTile + mTile * kWidth +
--                                 kTile * numElemsPerVec + e);
--                  }
--      }
--    } else {
--      // Original register layout:
--      //
--      //   [0, 1, 2, 3, 4, 5, 6, 7]^T, [8, 9, 10, 11, 12, 13, 14, 15]^T
--      //
--      // A stride of 4 is applied to derive four independent MMA operations:
--      //
--      //  1st MMA: [[0, 1], [8, 9]]
--      //  2nd MMA: [[2, 3], [10, 11]]
--      //  3rd MMA: [[4, 5], [12, 13]]
--      //  4th MMA: [[6, 7], [14, 15]]
--      if (kIters <= repK) {
--        for (size_t kRep = 0; kRep < kWidth / numElemsPerVec; ++kRep)
--          for (size_t tile = 0; tile < 2; ++tile)
--            for (size_t e = 0; e < numElemsPerVec; ++e) {
--              si.push_back(kRep * numElemsPerVec + tile * kWidth + e);
--            }
--      } else {
--        // Suppose kWidth=4 and type=fp32.
--        // Original register layout:
--        //
--        //       tile0/0        tile0/1
--        //   [0, 1, 2, 3]^T, [0, 1, 2, 3]^T
--        //
--        // Similar to the opIdx=0 situation, we should reorder the elements by
--        // moving the duplicated elements to the end.
--        size_t elemsPerTile = 2 * kWidth;
--        size_t elemsPerMma = 2 * numElemsPerVec;
--        size_t mmaWidth = kWidth / numElemsPerVec / 2;
--        size_t repMma = elemsPerTile / (mmaWidth * elemsPerMma);
--        for (size_t rep = 0; rep < repMma; ++rep)
--          for (size_t tile = 0; tile < elems.size() / elemsPerTile; ++tile)
--            for (size_t mmaKWidth = 0; mmaKWidth < mmaWidth; ++mmaKWidth)
--              for (size_t kTile = 0; kTile < 2; ++kTile)
--                for (size_t e = 0; e < numElemsPerVec; ++e) {
--                  si.push_back(rep * mmaWidth * elemsPerMma +
--                               mmaKWidth * 2 * numElemsPerVec +
--                               tile * elemsPerTile + kTile * numElemsPerVec +
--                               e);
--                }
--      }
--    }
--
--    auto step = si.size();
--    SmallVector<Value> perm(step);
--    for (auto i = 0; i < elems.size() / step; ++i) {
--      for (auto j = 0; j < step; ++j) {
--        perm[j] = elems[i * step + si[j]];
--      }
--      std::copy(perm.begin(), perm.end(), elems.begin() + i * step);
--    }
--  }
--
--  if (dot.getOpIdx() == 0) {
--    for (auto b = 0; b < batch; ++b)
--      for (auto m = 0; m < repOuter; ++m)
--        for (auto k = 0; k < repK; ++k)
--          for (auto vk = 0; vk < numRegisters.k; ++vk)
--            for (auto vm = 0; vm < numRegisters.m; ++vm)
--              packVec({b, m * numRegisters.m + vm, k * numRegisters.k + vk});
--  } else {
--    for (auto b = 0; b < batch; ++b)
--      for (auto n = 0; n < repOuter; ++n)
--        for (auto k = 0; k < repK; ++k)
--          for (auto vk = 0; vk < numRegisters.k; ++vk)
--            for (auto vn = 0; vn < numRegisters.n; ++vn)
--              packVec({b, n * numRegisters.n + vn, k * numRegisters.k + vk});
--  }
-+  for (int batchIdx = 0; batchIdx < batch; ++batchIdx)
-+    for (int outer = 0; outer < repOuter; ++outer)
-+      for (int k = 0; k < repK; ++k)
-+        for (int vk = 0; vk < numRegisters.k; ++vk)
-+          for (int vo = 0; vo < outerRegs; ++vo)
-+            packVec(
-+                {batchIdx, outer * outerRegs + vo, k * numRegisters.k + vk});
-   return vals;
- }
- 
-
diff --git a/third_party/xla/third_party/triton/common/llvm_cl969980821.patch b/third_party/xla/third_party/triton/common/llvm_cl969980821.patch
index a84564ee07b..23c2ed63a72 100644
--- a/third_party/xla/third_party/triton/common/llvm_cl969980821.patch
+++ b/third_party/xla/third_party/triton/common/llvm_cl969980821.patch
@@ -1,6 +1,6 @@
 --- a/python/src/llvm.cc
 +++ b/python/src/llvm.cc
-@@ -726,7 +726,7 @@ void init_triton_llvm(py::module_ &m) {
+@@ -704,7 +704,7 @@ void init_triton_llvm(py::module_ &m) {
            // registry name "vector-combine".
            const StringRef kVectorCombinePassName = "VectorCombinePass";
            passInstrCb.registerShouldRunOptionalPassCallback(
diff --git a/third_party/xla/third_party/triton/common/llvm_cl974500093.patch b/third_party/xla/third_party/triton/common/llvm_cl974500093.patch
index 1dff0ee4290..b9963ad02c2 100644
--- a/third_party/xla/third_party/triton/common/llvm_cl974500093.patch
+++ b/third_party/xla/third_party/triton/common/llvm_cl974500093.patch
@@ -47,7 +47,7 @@
 
 --- a/test/Conversion/tritongpu_to_llvm.mlir
 +++ b/test/Conversion/tritongpu_to_llvm.mlir
-@@ -1271,8 +1271,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+@@ -1394,8 +1394,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
    tt.func @convert_dot_ldmatrix(%A: tensor<16x16xf16, #blocked0>, %B: tensor<16x16xf16, #blocked0>) {
      %AA = ttg.local_alloc %A : (tensor<16x16xf16, #blocked0>) -> !ttg.memdesc<16x16xf16, #shared0, #smem>
      %BB = ttg.local_alloc %B : (tensor<16x16xf16, #blocked0>) -> !ttg.memdesc<16x16xf16, #shared0, #smem>
@@ -58,7 +58,7 @@
      // CHECK-NOT: nvvm.ldmatrix
      %AA_DOT = ttg.local_load %AA : !ttg.memdesc<16x16xf16, #shared0, #smem> -> tensor<16x16xf16, #dot_operand_a>
      %BB_DOT = ttg.local_load %BB : !ttg.memdesc<16x16xf16, #shared0, #smem> -> tensor<16x16xf16, #dot_operand_b>
-@@ -1301,8 +1301,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+@@ -1424,8 +1424,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
    tt.func @convert_dot_ldmatrix_swizzle(%A: tensor<16x16xf16, #blocked0>, %B: tensor<16x16xf16, #blocked0>) {
      %AA = ttg.local_alloc %A : (tensor<16x16xf16, #blocked0>) -> !ttg.memdesc<16x16xf16, #shared0, #smem>
      %BB = ttg.local_alloc %B : (tensor<16x16xf16, #blocked0>) -> !ttg.memdesc<16x16xf16, #shared0, #smem>
@@ -69,7 +69,7 @@
      // CHECK-NOT: nvvm.ldmatrix
      %AA_DOT = ttg.local_load %AA : !ttg.memdesc<16x16xf16, #shared0, #smem> -> tensor<16x16xf16, #dot_operand_a>
      %BB_DOT = ttg.local_load %BB : !ttg.memdesc<16x16xf16, #shared0, #smem> -> tensor<16x16xf16, #dot_operand_b>
-@@ -1383,7 +1383,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+@@ -1506,7 +1506,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
    tt.func @convert_dot_fp8(%A: tensor<16x16xf8E5M2, #blocked0>, %B: tensor<16x16xf8E5M2, #blocked0>) {
      %AA = ttg.local_alloc %A : (tensor<16x16xf8E5M2, #blocked0>) -> !ttg.memdesc<16x16xf8E5M2, #shared0, #smem>
      %BB = ttg.local_alloc %B : (tensor<16x16xf8E5M2, #blocked0>) -> !ttg.memdesc<16x16xf8E5M2, #shared0, #smem>
@@ -81,25 +81,65 @@
 
 --- a/test/Conversion/tritongpu_to_llvm_blackwell.mlir
 +++ b/test/Conversion/tritongpu_to_llvm_blackwell.mlir
-@@ -1006,7 +1006,7 @@ tt.func @load_store_x1_unpacked(%arg0: !ttg.memdesc<128x2xf16, #tmem_x1_unpacked
+@@ -1003,10 +1003,10 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.thr
  
  // CHECK-LABEL: max_reduction
- //       CHECK:  %[[M:.+]] = llvm.mlir.constant(-1 : i32) : i32
+ //       CHECK-DAG:  %[[M:.+]] = llvm.mlir.constant(-1 : i32) : i32
 -//       CHECK:   nvvm.redux.sync  fmax %{{.*}}, %[[M]] {nan = true} : f32 -> f32
 +//       CHECK:   nvvm.redux.sync  fmax %{{.*}}, %[[M]] nan = true : f32 -> f32
  //       CHECK:   nvvm.barrier
- //       CHECK:   nvvm.shfl.sync bfly
- //       CHECK:   nvvm.shfl.sync bfly
-@@ -1051,7 +1051,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 2 : i32} {
+ //   CHECK-NOT:   nvvm.shfl.sync
+-//       CHECK:   nvvm.redux.sync  fmax %{{.*}}, %[[M]] {nan = true} : f32 -> f32
++//       CHECK:   nvvm.redux.sync  fmax %{{.*}}, %[[M]] nan = true : f32 -> f32
+ //  CHECK-NEXT:   llvm.return
+ #blocked = #ttg.blocked<{sizePerThread = [1, 4], threadsPerWarp = [1, 32], warpsPerCTA = [1, 4], order = [1, 0]}>
+ module attributes {"ttg.target" = "cuda:100", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
+@@ -1050,7 +1050,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 2 : i32} {
    // CHECK-LABEL: lower_ldmatrix_trans_b8
    tt.func @lower_ldmatrix_trans_b8(%A: !ttg.memdesc<128x64xf8E4M3FN, #shared, #smem, mutable, 1x128x64>) {
      %0 = ttg.local_load %A : !ttg.memdesc<128x64xf8E4M3FN, #shared, #smem, mutable, 1x128x64> -> tensor<128x64xf8E4M3FN, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 4}>>
 -    // CHECK-COUNT-16: nvvm.ldmatrix %{{.*}} {eltType = #nvvm.ld_st_matrix_elt_type<b8>, layout = #nvvm.mma_layout<col>{{.*}}} : (!llvm.ptr<3>) -> !llvm.struct<(i32, i32, i32, i32)>
-+    // CHECK-COUNT-16: nvvm.ldmatrix %{{.*}}layout = <col>{{.*}}element_type = <b8> : (!llvm.ptr<3>) -> !llvm.struct<(i32, i32, i32, i32)>
++    // CHECK-COUNT-16: nvvm.ldmatrix %{{.*}}num = 2, layout = <col>, shape = <m = 16, n = 16>, element_type = <b8> : (!llvm.ptr<3>) -> !llvm.struct<(i32, i32, i32, i32)>
      tt.return
    }
  }
-@@ -1064,7 +1064,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 2 : i32} {
+@@ -1067,7 +1067,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+   // CHECK-LABEL: @ldmatrix_b8_minimise_conflicts_preserve_words
+   // CHECK-DAG: %[[ONE:.*]] = llvm.mlir.constant(1 : i32)
+   // CHECK-DAG: %[[TWO:.*]] = llvm.mlir.constant(2 : i32)
+-  // CHECK: %[[LOAD:.*]] = nvvm.ldmatrix %{{.*}} {eltType = #nvvm.ld_st_matrix_elt_type<b8>, layout = #nvvm.mma_layout<col>, num = 2 : i32
++  // CHECK: %[[LOAD:.*]] = nvvm.ldmatrix %{{.*}} num = 2, layout = <col>, shape = <m = 16, n = 16>, element_type = <b8>
+   // CHECK: %[[WORD:.*]] = llvm.extractvalue %[[LOAD]][0]
+   // CHECK: %[[BYTES:.*]] = llvm.bitcast %[[WORD]] : i32 to vector<4xi8>
+   // CHECK: %[[BYTE1:.*]] = llvm.extractelement %[[BYTES]][%[[ONE]] : i32]
+@@ -1081,7 +1081,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+   // Selecting register [8] avoids the bank conflict caused by register [64].
+   // CHECK-LABEL: @ldmatrix_b16_minimise_conflicts
+   // CHECK-DAG: %[[ONE:.*]] = llvm.mlir.constant(1 : i32)
+-  // CHECK: %[[LOAD:.*]] = nvvm.ldmatrix %{{.*}} {eltType = #nvvm.ld_st_matrix_elt_type<b16>, layout = #nvvm.mma_layout<col>, num = 4 : i32
++  // CHECK: %[[LOAD:.*]] = nvvm.ldmatrix %{{.*}} num = 4, layout = <col>, shape = <m = 8, n = 8>, element_type = <b16>
+   // CHECK: %[[WORD:.*]] = llvm.extractvalue %[[LOAD]][0]
+   // CHECK: %[[HALVES:.*]] = llvm.bitcast %[[WORD]] : i32 to vector<2xf16>
+   // CHECK: %[[HALF1:.*]] = llvm.extractelement %[[HALVES]][%[[ONE]] : i32]
+@@ -1098,7 +1098,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+   // CHECK: %[[PACK2:.*]] = llvm.insertelement %{{.*}}, %[[PACK1]][
+   // CHECK: %[[PACK3:.*]] = llvm.insertelement %{{.*}}, %[[PACK2]][
+   // CHECK: %[[WORD:.*]] = llvm.bitcast %[[PACK3]] : vector<4xi8> to i32
+-  // CHECK: nvvm.stmatrix %{{[^,]+}}, %[[WORD]], %{{.*}} {eltType = #nvvm.ld_st_matrix_elt_type<b8>, layout = #nvvm.mma_layout<col>, shape = #nvvm.ld_st_matrix_shape<m = 16, n = 8>} : !llvm.ptr<3>, i32, i32, i32, i32
++  // CHECK: nvvm.stmatrix %{{[^,]+}}, %[[WORD]], %{{.*}} layout = <col>, shape = <m = 16, n = 8>, element_type = <b8> : !llvm.ptr<3>, i32, i32, i32, i32
+   tt.func private @stmatrix_b8_minimise_conflicts(%A: !ttg.memdesc<512xi8, #shared, #smem, mutable>, %data: tensor<512xi8, #stb8>) {
+     ttg.local_store %data, %A : tensor<512xi8, #stb8> -> !ttg.memdesc<512xi8, #shared, #smem, mutable>
+     tt.return
+@@ -1108,7 +1108,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+   // CHECK: %[[INPUT2:.*]] = llvm.extractvalue %{{.*}}[2]
+   // CHECK: %[[PACK:.*]] = llvm.insertelement %[[INPUT2]], %{{.*}}[%[[ONE]] : i32]
+   // CHECK: %[[WORD:.*]] = llvm.bitcast %[[PACK]] : vector<2xf16> to i32
+-  // CHECK: nvvm.stmatrix %{{[^,]+}}, %[[WORD]], %{{.*}} {eltType = #nvvm.ld_st_matrix_elt_type<b16>, layout = #nvvm.mma_layout<col>, shape = #nvvm.ld_st_matrix_shape<m = 8, n = 8>} : !llvm.ptr<3>, i32, i32, i32, i32
++  // CHECK: nvvm.stmatrix %{{[^,]+}}, %[[WORD]], %{{.*}} layout = <col>, shape = <m = 8, n = 8>, element_type = <b16> : !llvm.ptr<3>, i32, i32, i32, i32
+   tt.func private @stmatrix_b16_minimise_conflicts(%A: !ttg.memdesc<256xf16, #shared, #smem, mutable>, %data: tensor<256xf16, #b16>) {
+     ttg.local_store %data, %A : tensor<256xf16, #b16> -> !ttg.memdesc<256xf16, #shared, #smem, mutable>
+     tt.return
+@@ -1123,7 +1123,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
  module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
    // CHECK-LABEL: @stmatrix_b8_trans_linear
    tt.func public @stmatrix_b8_trans_linear(%data: tensor<1x1x1x16x256xf8E4M3FN, #linear3>) {
@@ -111,7 +151,7 @@
 
 --- a/test/Conversion/tritongpu_to_llvm_hopper.mlir
 +++ b/test/Conversion/tritongpu_to_llvm_hopper.mlir
-@@ -441,7 +441,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
+@@ -470,7 +470,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
  // CHECK-LABEL: distribute_to_shared_st_matrix_local_store
  module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
    tt.func @distribute_to_shared_st_matrix_local_store(%a: tensor<64x128xf16, #linear>) {
@@ -120,7 +160,7 @@
      //          CHECK: llvm.return
      %b = ttg.local_alloc {allocation.offset = 0 : i32} : () -> !ttg.memdesc<64x128xf16, #shared, #smem, mutable>
      ttg.local_store %a, %b : tensor<64x128xf16, #linear> -> !ttg.memdesc<64x128xf16, #shared, #smem, mutable>
-@@ -473,7 +473,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
+@@ -502,7 +502,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
  // CHECK-LABEL: linear_to_swizzled_st_matrix_local_store
  module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
    tt.func @linear_to_swizzled_st_matrix_local_store(%a: tensor<64x32xf16, #linear>) {
@@ -129,7 +169,7 @@
      //          CHECK: llvm.return
      %b = ttg.local_alloc {allocation.offset = 0 : i32} : () -> !ttg.memdesc<64x32xf16, #shared, #smem, mutable>
      ttg.local_store %a, %b : tensor<64x32xf16, #linear> -> !ttg.memdesc<64x32xf16, #shared, #smem, mutable>
-@@ -495,7 +495,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
+@@ -524,7 +524,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
  // CHECK-LABEL: linear_to_swizzled_st_matrix_local_store
  module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
    tt.func @linear_to_swizzled_st_matrix_local_store(%a: tensor<32x32xf16, #linear>) {
@@ -138,7 +178,7 @@
      //          CHECK: llvm.return
      %b = ttg.local_alloc {allocation.offset = 0 : i32} : () -> !ttg.memdesc<32x32xf16, #shared, #smem, mutable>
      ttg.local_store %a, %b : tensor<32x32xf16, #linear> -> !ttg.memdesc<32x32xf16, #shared, #smem, mutable>
-@@ -511,7 +511,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
+@@ -540,7 +540,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
  // CHECK-LABEL: linear_to_swizzled_st_matrix_x2_local_store_fp8
  module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
    tt.func @linear_to_swizzled_st_matrix_x2_local_store_fp8(%a: tensor<64x16xf8E4M3FNUZ, #linear>) {
@@ -147,7 +187,7 @@
      //          CHECK: llvm.return
      %b = ttg.local_alloc {allocation.offset = 0 : i32} : () -> !ttg.memdesc<64x16xf8E4M3FNUZ, #shared, #smem, mutable>
      ttg.local_store %a, %b : tensor<64x16xf8E4M3FNUZ, #linear> -> !ttg.memdesc<64x16xf8E4M3FNUZ, #shared, #smem, mutable>
-@@ -527,7 +527,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
+@@ -556,7 +556,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
  // CHECK-LABEL: linear_to_swizzled_st_matrix_local_store_fp32
  module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
    tt.func @linear_to_swizzled_st_matrix_local_store_fp32(%a: tensor<64x16xf32, #linear>) {
@@ -156,7 +196,7 @@
      //          CHECK: llvm.return
      %b = ttg.local_alloc {allocation.offset = 0 : i32} : () -> !ttg.memdesc<64x16xf32, #shared, #smem, mutable>
      ttg.local_store %a, %b : tensor<64x16xf32, #linear> -> !ttg.memdesc<64x16xf32, #shared, #smem, mutable>
-@@ -544,7 +544,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
+@@ -573,7 +573,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
  // CHECK-LABEL: linear_to_swizzled_st_matrix_trans_local_store
  module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
    tt.func @linear_to_swizzled_st_matrix_trans_local_store(%a: tensor<64x32xf16, #linear>) {
@@ -165,7 +205,7 @@
      //          CHECK: llvm.return
      %b = ttg.local_alloc {allocation.offset = 0 : i32} : () -> !ttg.memdesc<64x32xf16, #shared, #smem, mutable>
      ttg.local_store %a, %b : tensor<64x32xf16, #linear> -> !ttg.memdesc<64x32xf16, #shared, #smem, mutable>
-@@ -566,7 +566,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
+@@ -595,7 +595,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
  // CHECK-LABEL: linear_to_swizzled_st_matrix_trans_local_store
  module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
    tt.func @linear_to_swizzled_st_matrix_trans_local_store(%a: tensor<16x32xf16, #linear>) {
@@ -177,7 +217,7 @@
 
 --- a/test/Conversion/tritonnvidiagpu_to_llvm.mlir
 +++ b/test/Conversion/tritonnvidiagpu_to_llvm.mlir
-@@ -520,7 +520,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+@@ -531,7 +531,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
  
  module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
    // CHECK-LABEL: async_tma_store_wait_read_only
@@ -186,7 +226,7 @@
    tt.func @async_tma_store_wait_read_only() {
      ttng.async_tma_store_wait {pendings = 0 : i32, read_only}
      tt.return
-@@ -660,7 +660,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.targ
+@@ -697,7 +697,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.targ
    tt.func public @async_copy_mbarrier_arrive(%arg0: !ttg.memdesc<1xi64, #shared, #ttg.shared_memory>)  attributes { noinline = false } {
      // CHECK: nvvm.cp.async.mbarrier.arrive %{{.*}} : !llvm.ptr<3>
      ttng.async_copy_mbarrier_arrive %arg0 : !ttg.memdesc<1xi64, #shared, #ttg.shared_memory>
@@ -195,7 +235,7 @@
      ttng.async_copy_mbarrier_arrive %arg0 { noIncrement } : !ttg.memdesc<1xi64, #shared, #ttg.shared_memory>
      tt.return
    }
-@@ -1150,7 +1150,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32, "ttg.th
+@@ -1195,7 +1195,7 @@ module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttg.tot
  module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32, "ttg.threads-per-warp" = 32 : i32} {
    // CANONICALIZE-SM100-LABEL: @redux_max_abs(
    // CANONICALIZE-SM100-NOT: llvm.intr.fabs
@@ -204,7 +244,7 @@
    // CANONICALIZE-SM100-NEXT: llvm.return
    tt.func private @redux_max_abs(%x: tensor<32xf32, #blocked>) -> f32 {
      %abs = math.absf %x : tensor<32xf32, #blocked>
-@@ -1164,7 +1164,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32, "ttg.th
+@@ -1209,7 +1209,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32, "ttg.thr
  
    // CANONICALIZE-SM100-LABEL: @redux_max_abs_nan(
    // CANONICALIZE-SM100-NOT: llvm.intr.fabs
@@ -213,7 +253,7 @@
    // CANONICALIZE-SM100-NEXT: llvm.return
    tt.func private @redux_max_abs_nan(%x: tensor<32xf32, #blocked>) -> f32 {
      %abs = math.absf %x : tensor<32xf32, #blocked>
-@@ -1178,7 +1178,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32, "ttg.th
+@@ -1223,7 +1223,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32, "ttg.thr
  
    // CANONICALIZE-SM100-LABEL: @redux_min_abs(
    // CANONICALIZE-SM100-NOT: llvm.intr.fabs
diff --git a/third_party/xla/third_party/triton/common/llvm_cl978258752.patch b/third_party/xla/third_party/triton/common/llvm_cl978258752.patch
index 11daacfa5f3..00208164c50 100644
--- a/third_party/xla/third_party/triton/common/llvm_cl978258752.patch
+++ b/third_party/xla/third_party/triton/common/llvm_cl978258752.patch
@@ -1,6 +1,6 @@
 --- a/include/triton/Dialect/Triton/IR/TritonOps.td
 +++ b/include/triton/Dialect/Triton/IR/TritonOps.td
-@@ -1129,8 +1129,8 @@
+@@ -1213,8 +1213,8 @@ def CallOp : TT_Op<"call", [CallOpInterface, /*MemRefsNormalizable, */DeclareOpI
  }
  
  def FuncOp : TT_Op<"func", [
@@ -14,7 +14,7 @@
 
 --- a/lib/Conversion/TritonGPUToLLVM/Utility.cpp
 +++ b/lib/Conversion/TritonGPUToLLVM/Utility.cpp
-@@ -2193,7 +2193,7 @@
+@@ -2283,7 +2283,7 @@ void filterFuncAttributes(triton::FuncOp op, bool filterArgAttrs,
                            SmallVectorImpl<NamedAttribute> &result) {
  
    for (const auto &attr : op->getAttrs()) {
@@ -26,7 +26,7 @@
 
 --- a/lib/Dialect/Triton/IR/Ops.cpp
 +++ b/lib/Dialect/Triton/IR/Ops.cpp
-@@ -1275,8 +1275,7 @@
+@@ -1450,8 +1450,7 @@ void IntToPtrOp::getCanonicalizationPatterns(RewritePatternSet &results,
  void FuncOp::build(OpBuilder &builder, OperationState &state, StringRef name,
                     FunctionType type, ArrayRef<NamedAttribute> attrs,
                     ArrayRef<DictionaryAttr> argAttrs) {
@@ -37,3 +37,16 @@
    state.attributes.append(attrs.begin(), attrs.end());
    state.addRegion();
 
+--- a/test/Conversion/amd/tritongpu_to_llvm.mlir
++++ b/test/Conversion/amd/tritongpu_to_llvm.mlir
+@@ -1110,8 +1110,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.thr
+   // CHECK-NOT: rocdl.grid.dim
+   // CHECK: llvm.call @callee_zero_scratch
+   // CHECK: llvm.func internal @callee_zero_scratch
+-  // CHECK-SAME: passthrough = ["noinline", "convergent"]
+   // CHECK-SAME: ws_num_warps = 4 : i32
++  // CHECK-SAME: passthrough = ["noinline", "convergent"]
+   tt.func public @test_call_zero_scratch_no_grid_ops() attributes {noinline = false} {
+     tt.call @callee_zero_scratch() : () -> ()
+     tt.return
+
diff --git a/third_party/xla/third_party/triton/common/llvm_cl979179688.patch b/third_party/xla/third_party/triton/common/llvm_cl979179688.patch
index 8ef04678784..15258224806 100644
--- a/third_party/xla/third_party/triton/common/llvm_cl979179688.patch
+++ b/third_party/xla/third_party/triton/common/llvm_cl979179688.patch
@@ -1,6 +1,6 @@
 --- a/test/LLVMIR/insert-dbg-intrinsic.mlir
 +++ b/test/LLVMIR/insert-dbg-intrinsic.mlir
-@@ -19,7 +19,7 @@
+@@ -19,7 +19,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32 } {
      // CHECK: llvm.intr.dbg.value #di_local_variable{{([0-9]*)?}} = %y_ptr :
      // CHECK: llvm.intr.dbg.value #di_local_variable{{([0-9]*)?}} = %out_ptr :
      // CHECK: llvm.intr.dbg.value #di_local_variable{{([0-9]*)?}} = %n_elements :
@@ -9,7 +9,7 @@
  
      // CHECK: %pid = rocdl.workgroup.id.x
      // CHECK-NEXT: llvm.intr.dbg.value #di_local_variable{{([0-9]*)?}} = %pid :
-@@ -59,7 +59,7 @@
+@@ -59,7 +59,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 8 : i32} {
                                  %arg1: !llvm.array<4 x i8>,
                                  %arg2: !llvm.ptr<1> {tt.pointee_type = i16},
                                  %arg3: i32) attributes {noinline = false} {
@@ -21,7 +21,7 @@
 
 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp
-@@ -382,7 +382,9 @@
+@@ -389,7 +389,9 @@ struct FpToFpOpConversion
            "unsupported rounding mode for f32->bf16 conversion: " +
            stringifyRoundingMode(rounding) + "\n");
      }
diff --git a/third_party/xla/third_party/triton/common/llvm_cl979526570.patch b/third_party/xla/third_party/triton/common/llvm_cl979526570.patch
index 97bf074d7bc..162119c597a 100644
--- a/third_party/xla/third_party/triton/common/llvm_cl979526570.patch
+++ b/third_party/xla/third_party/triton/common/llvm_cl979526570.patch
@@ -1,6 +1,6 @@
 --- a/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp
 +++ b/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp
-@@ -275,7 +275,8 @@
+@@ -268,7 +268,8 @@ struct ElementwiseInlineAsmOpConversion
                             /*asm_string=*/op.getAsmString(),
                             /*constraints=*/op.getConstraints(),
                             /*has_side_effects=*/!op.getPure(),
@@ -10,10 +10,29 @@
                             /*asm_dialect=*/
                             LLVM::AsmDialectAttr::get(rewriter.getContext(),
                                                       LLVM::AsmDialect::AD_ATT),
+@@ -417,6 +418,7 @@ struct InlineAsmOpConversion
+     auto call = LLVM::InlineAsmOp::create(
+         rewriter, loc, returnType, operands, op.getAsmString(),
+         op.getConstraints(), !op.getPure(), false, LLVM::TailCallKind::None,
++        /*convergent=*/false,
+         LLVM::AsmDialectAttr::get(ctx, LLVM::AsmDialect::AD_ATT), ArrayAttr());
+ 
+     SmallVector<Value> elements;
+
+--- a/third_party/amd/lib/TritonAMDGPUToLLVM/DotOpToLLVM/MFMA.cpp
++++ b/third_party/amd/lib/TritonAMDGPUToLLVM/DotOpToLLVM/MFMA.cpp
+@@ -152,6 +152,7 @@ struct DotOpMFMAConversionHelper {
+                rewriter, loc, acc.getType(), ValueRange{acc},
+                /*asm_string=*/"", constraints, /*has_side_effects=*/false,
+                /*is_align_stack=*/false, LLVM::TailCallKind::None,
++               /*convergent=*/false,
+                LLVM::AsmDialectAttr::get(ctx, LLVM::AsmDialect::AD_ATT),
+                /*operand_attrs=*/ArrayAttr())
+         ->getResult(0);
 
 --- a/third_party/amd/lib/TritonAMDGPUToLLVM/GCNAsmFormat.cpp
 +++ b/third_party/amd/lib/TritonAMDGPUToLLVM/GCNAsmFormat.cpp
-@@ -83,6 +83,7 @@
+@@ -83,6 +83,7 @@ mlir::Value GCNBuilder::launch(RewriterBase &rewriter, Location loc, Type resTy,
        hasSideEffect,                          // has_side_effects
        isAlignStack,                           // is_align_stack
        LLVM::TailCallKind::None,
@@ -24,7 +43,7 @@
 
 --- a/third_party/amd/lib/TritonAMDGPUToLLVM/TDMUtility.cpp
 +++ b/third_party/amd/lib/TritonAMDGPUToLLVM/TDMUtility.cpp
-@@ -37,6 +37,7 @@
+@@ -37,6 +37,7 @@ Value sequenceGroup0(RewriterBase &rewriter, Location loc, Value group0) {
               /*asm_string=*/"", /*constraints=*/"=s,0",
               /*has_side_effects=*/true, /*is_align_stack=*/false,
               LLVM::TailCallKind::None,
@@ -35,7 +54,7 @@
 
 --- a/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp
 +++ b/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp
-@@ -693,6 +693,7 @@
+@@ -749,6 +749,7 @@ void TargetInfo::assertFail(RewriterBase &rewriter, Location loc,
        "s_trap 2", /*constraints=*/"",
        /*has_side_effects=*/true, /*is_align_stack=*/false,
        LLVM::TailCallKind::None,
@@ -46,7 +65,7 @@
 
 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/PTXAsmFormat.cpp
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/PTXAsmFormat.cpp
-@@ -106,6 +106,7 @@
+@@ -106,6 +106,7 @@ mlir::Value PTXBuilder::launch(OpBuilder &rewriter, Location loc, Type resTy,
        hasSideEffect,                          // has_side_effects
        isAlignStack,                           // is_align_stack
        LLVM::TailCallKind::None,
diff --git a/third_party/xla/third_party/triton/common/llvm_cl983115707.patch b/third_party/xla/third_party/triton/common/llvm_cl983115707.patch
index 4c73f8cef42..3ba6c6920c9 100644
--- a/third_party/xla/third_party/triton/common/llvm_cl983115707.patch
+++ b/third_party/xla/third_party/triton/common/llvm_cl983115707.patch
@@ -1,12 +1,6 @@
 --- a/python/src/llvm.cc
 +++ b/python/src/llvm.cc
-@@ -12,11 +12,13 @@
- #include "llvm/Config/llvm-config.h"
- #include "llvm/IR/DebugInfo.h"
- #include "llvm/IR/IRBuilder.h"
-+#include "llvm/IR/InstIterator.h"
- #include "llvm/IR/IntrinsicInst.h"
- #include "llvm/IR/LLVMContext.h"
+@@ -24,6 +24,7 @@
  #include "llvm/IR/LegacyPassManager.h"
  #include "llvm/IR/Metadata.h"
  #include "llvm/IR/Module.h"
@@ -14,7 +8,7 @@
  #include "llvm/IR/PassManager.h"
  #include "llvm/IR/Verifier.h"
  #include "llvm/IRReader/IRReader.h"
-@@ -183,6 +185,14 @@
+@@ -196,6 +197,14 @@ void installTritonDAGScheduler() {
        ->setInitialValue(createTritonDAGScheduler);
  }
  
@@ -29,7 +23,7 @@
  std::unique_ptr<TargetMachine>
  createTargetMachine(llvm::Module *module, std::string proc,
                      bool enable_fp_fusion, const std::string &features) {
-@@ -196,7 +206,7 @@
+@@ -205,7 +214,7 @@ createTargetMachine(llvm::Module *module, std::string proc,
    llvm::TargetOptions opt;
    bool disableLLVMOpt = mlir::triton::tools::getBoolEnv("DISABLE_LLVM_OPT");
    if (enable_fp_fusion)
@@ -41,7 +35,7 @@
 
 --- a/test/Conversion/amd/async-ops-alias-scopes.mlir
 +++ b/test/Conversion/amd/async-ops-alias-scopes.mlir
-@@ -18,7 +18,7 @@
+@@ -18,7 +18,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32, ttg.targ
  
      // COMMON: rocdl.global.load.async.lds {{.*}} <alias_scopes = [[[$ASYNC_COPY_SCOPE]]]
      // Check that store for 'other' has alias information set
@@ -50,7 +44,7 @@
      %0 = ttg.async_copy_global_to_local %ptr, %arg1 mask %mask other %other : tensor<64x1x!tt.ptr<f32>, #blocked> -> <64x1xf32, #shared, #smem, mutable>
  
      // COMMON: llvm.return
-@@ -44,7 +44,7 @@
+@@ -44,7 +44,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.shar
  
      // COMMON: rocdl.raw.ptr.buffer.load.async.lds {{.*}} <alias_scopes = [[[$ASYNC_COPY_SCOPE]]]
      // Check that store for 'other' has alias information set
@@ -59,7 +53,7 @@
      %65 = amdg.buffer_load_to_local %arg1[%arg2] mask=%mask other=%other into %arg3 : !tt.ptr<f32>[tensor<8x64xi32, #blocked>] tensor<8x64xf32, #blocked> -> <8x64xf32, #shared, #smem, mutable>
  
      // COMMON: llvm.return
-@@ -71,11 +71,11 @@
+@@ -71,11 +71,11 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32, ttg.targ
      // Check alias information is added for different lowering paths
  
      // Test lowering path in common MemoryOpToLLVM pattern
@@ -73,7 +67,7 @@
      // GFX950: rocdl.ds.read.tr16.b64 {{.*}} <alias_scopes = [[[$LOCAL_LOAD_SCOPE]]], noalias_scopes = [[[$ASYNC_COPY_SCOPE]]]
      %5 = ttg.local_load %arg2 token %3 : !ttg.memdesc<16x16xf16, #shared, #smem, mutable> -> tensor<16x16xf16, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 8}>>
  
-@@ -152,7 +152,7 @@
+@@ -152,7 +152,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32, ttg.targ
      %2 = ttg.local_load %arg1 token %1 : !ttg.memdesc<64x1xf16, #shared, #smem, mutable> -> tensor<64x1xf16, #blocked>
  
      %loop_result:2 = scf.for %arg14 = %c0_i32 to %loopIterCount step %c1_i32 iter_args(%arg10 = %1, %arg11 = %2) -> (!ttg.async.token, tensor<64x1xf16, #blocked>)  : i32 {
@@ -85,7 +79,7 @@
 
 --- a/test/Conversion/amd/in_thread_transpose.mlir
 +++ b/test/Conversion/amd/in_thread_transpose.mlir
-@@ -18,11 +18,11 @@
+@@ -18,11 +18,11 @@ module attributes {"ttg.target" = "hip:gfx942", "ttg.num-ctas" = 1 : i32, "ttg.n
  
      // CHECK-DAG: [[VEC1_TMP:%.*]] = llvm.insertelement [[VAL0]], [[VEC_UNDEF]]{{\[}}[[CST_0]] : i32] : vector<2xf16>
      // CHECK-DAG: [[VEC1:%.*]] = llvm.insertelement [[VAL2]], [[VEC1_TMP]]{{\[}}[[CST_1]] : i32] : vector<2xf16>
@@ -102,16 +96,16 @@
 
 --- a/test/Conversion/tritongpu_to_llvm.mlir
 +++ b/test/Conversion/tritongpu_to_llvm.mlir
-@@ -2434,7 +2434,7 @@
+@@ -2844,7 +2844,7 @@ module attributes {"ttg.target" = "cuda:75", "ttg.num-ctas" = 1 : i32, "ttg.num-
  module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
    // CHECK-LABEL: @vectorize_shmem_load
    // CHECK: llvm.load
 -  // CHECK-SAME: {alignment = 8 : i64} : !llvm.ptr<3> -> vector<2xi32>
 +  // CHECK-SAME: <alignment = 8> : !llvm.ptr<3> -> vector<2xi32>
    // CHECK-NOT: llvm.load
-   tt.func public @vectorize_shmem_load(%shmem : !ttg.memdesc<16x16xi8, #shared, #smem>) {
+   tt.func private @vectorize_shmem_load(%shmem : !ttg.memdesc<16x16xi8, #shared, #smem>) -> tensor<16x16xi8, #blocked> {
      %0 = ttg.local_load %shmem : !ttg.memdesc<16x16xi8, #shared, #smem> -> tensor<16x16xi8, #blocked>
-@@ -2449,7 +2449,7 @@
+@@ -2859,7 +2859,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
  #smem = #ttg.shared_memory
  module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 8 : i32, "ttg.threads-per-warp" = 32 : i32} {
    // CHECK-LABEL: @vectorize_shmem_store
@@ -123,7 +117,7 @@
 
 --- a/test/Conversion/warp_specialize_to_llvm.mlir
 +++ b/test/Conversion/warp_specialize_to_llvm.mlir
-@@ -212,7 +212,7 @@
+@@ -212,7 +212,7 @@ llvm.func @generate_switch_loop() attributes {allocation.offset = 32 : i32} {
    // CHECK: [[DEFAULT]]:
    // CHECK-NEXT: "llvm.nvvm.barrier.cta.sync.all"([[C1]])
    // CHECK-NEXT: "llvm.nvvm.barrier.cta.sync.all"([[C1]])
@@ -132,7 +126,7 @@
  
    // CHECK: [[EXIT]]:
    // CHECK-NEXT: llvm.return
-@@ -283,7 +283,7 @@
+@@ -283,7 +283,7 @@ llvm.func @generate_switch_loop() attributes {allocation.offset = 32 : i32} {
    // AMD: [[DEFAULT]]:
    // AMD-NEXT: rocdl.barrier
    // AMD-NEXT: rocdl.barrier
@@ -141,7 +135,7 @@
  
    // AMD: [[EXIT]]:
    // AMD-NEXT: llvm.return
-@@ -414,9 +414,9 @@
+@@ -414,9 +414,9 @@ llvm.func @pass_captures() attributes {allocation.offset = 32 : i32} {
  
    // CHECK: ^bb4:
    // CHECK-NEXT: [[ARG0_PTR:%.*]] = llvm.getelementptr [[SMEM_ADDR]][0, 0] : (!llvm.ptr<3>) -> !llvm.ptr<3>, !llvm.struct<packed (i32, i64)>
@@ -153,7 +147,7 @@
    // CHECK-NEXT: "llvm.nvvm.barrier.cta.sync.all"([[C1]])
    // CHECK-NEXT: "use"([[ARG0]], [[ARG1]])
    // CHECK-NEXT: "llvm.nvvm.barrier.cta.sync.all"([[C1]])
-@@ -424,17 +424,17 @@
+@@ -424,17 +424,17 @@ llvm.func @pass_captures() attributes {allocation.offset = 32 : i32} {
    // CHECK: ^bb5:
    // CHECK: [[INS:%.*]]:2 = "produce"()
    // CHECK: [[ARG0_PTR:%.*]] = llvm.getelementptr [[SMEM_ADDR]][0, 0] : (!llvm.ptr<3>) -> !llvm.ptr<3>, !llvm.struct<packed (i32, i64)>
@@ -175,7 +169,7 @@
    // AMD-NEXT: rocdl.barrier
    // AMD-NEXT: "use"([[ARG0]], [[ARG1]])
    // AMD-NEXT: rocdl.barrier
-@@ -442,9 +442,9 @@
+@@ -442,9 +442,9 @@ llvm.func @pass_captures() attributes {allocation.offset = 32 : i32} {
    // AMD: ^bb5:
    // AMD: [[INS:%.*]]:2 = "produce"()
    // AMD: [[ARG0_PTR:%.*]] = llvm.getelementptr [[SMEM_ADDR]][0, 0] : (!llvm.ptr<3>) -> !llvm.ptr<3>, !llvm.struct<packed (i32, i64)>
@@ -198,7 +192,7 @@
  #include "mlir/Analysis/TopologicalSortUtils.h"
  #include "mlir/Conversion/Passes.h"
  #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
-@@ -49,7 +50,7 @@
+@@ -49,7 +50,7 @@ public:
    }
  
    Type getBarrierHandleType(MLIRContext *ctx) const override {
@@ -207,7 +201,7 @@
    }
  
    FailureOr<Value>
-@@ -80,7 +81,7 @@
+@@ -80,7 +81,7 @@ public:
        nbarGV = LLVM::GlobalOp::create(
            b, uloc, nbarTy, /*isConstant=*/false, LLVM::Linkage::Internal,
            nbarAttr.getValue(), /*value=*/Attribute(), /*alignment=*/0,
diff --git a/third_party/xla/third_party/triton/common/llvm_cl991819890.patch b/third_party/xla/third_party/triton/common/llvm_cl991819890.patch
index cc4a1ad7227..791e53a811b 100644
--- a/third_party/xla/third_party/triton/common/llvm_cl991819890.patch
+++ b/third_party/xla/third_party/triton/common/llvm_cl991819890.patch
@@ -1,6 +1,6 @@
 --- a/python/src/llvm.cc
 +++ b/python/src/llvm.cc
-@@ -838,7 +838,7 @@
+@@ -753,7 +753,7 @@ void init_triton_llvm(py::module_ &m) {
          if (!pluginFile.empty()) {
            // TODO: Add some logging here that we inserted a pass into the LLVM
            // pass pipeline
@@ -10,3 +10,50 @@
              llvm::Error Err = passPlugin.takeError();
              std::string ErrMsg =
 
+--- a/python/test/unit/language/test_compile_only.py
++++ b/python/test/unit/language/test_compile_only.py
+@@ -99,8 +99,7 @@ def test_compile_only_mulhi_survives_llvm_opt(bits) -> None:
+     for kind, stage in stages.items():
+         module = stage(module, metadata)
+         if kind == "llir":
+-            intrinsic = "ui" if bits == 32 else "ull"
+-            assert f"@llvm.nvvm.mulhi.{intrinsic}" in str(module)
++            assert f"@llvm.umulh.i{bits}" in str(module)
+         if kind == "ptx":
+             assert f"mul.hi.u{bits}" in module
+             assert "mul.lo." not in module
+
+--- a/test/Conversion/tritongpu_to_llvm.mlir
++++ b/test/Conversion/tritongpu_to_llvm.mlir
+@@ -23,7 +23,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+   tt.func @umulhi_truncated_input(%seed: i64, %out: !tt.ptr<i32>) {
+     // CHECK: %[[X:.*]] = llvm.trunc %arg0 : i64 to i32
+     %x = arith.trunci %seed : i64 to i32
+-    // CHECK: llvm.call_intrinsic "llvm.nvvm.mulhi.ui"(%[[X]], %[[X]]) : (i32, i32) -> i32
++    // CHECK: llvm.call_intrinsic "llvm.umulh"(%[[X]], %[[X]]) : (i32, i32) -> i32
+     %high = tt.mulhiui %x, %x : i32
+     tt.store %out, %high : !tt.ptr<i32>
+     tt.return
+@@ -31,7 +31,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+ 
+   // CHECK-LABEL: llvm.func @umulhi_i64
+   tt.func @umulhi_i64(%seed: i64, %out: !tt.ptr<i64>) {
+-    // CHECK: llvm.call_intrinsic "llvm.nvvm.mulhi.ull"(%arg0, %arg0) : (i64, i64) -> i64
++    // CHECK: llvm.call_intrinsic "llvm.umulh"(%arg0, %arg0) : (i64, i64) -> i64
+     %high = tt.mulhiui %seed, %seed : i64
+     tt.store %out, %high : !tt.ptr<i64>
+     tt.return
+
+--- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp
++++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp
+@@ -949,8 +949,7 @@ struct MulhiUIOpConversion
+                                    Location loc) const {
+     unsigned bitWidth = elemTy.getIntOrFloatBitWidth();
+     assert(bitWidth == 32 || bitWidth == 64);
+-    StringRef intrinsic =
+-        bitWidth == 32 ? "llvm.nvvm.mulhi.ui" : "llvm.nvvm.mulhi.ull";
++    StringRef intrinsic = "llvm.umulh";
+     // A widened multiply/shift relies on LLVM recognizing explicit extensions.
+     // InstCombine can replace those extensions with masks, losing mul.hi.
+     return {LLVM::createLLVMIntrinsicCallOp(rewriter, loc, intrinsic, elemTy,
+
diff --git a/third_party/xla/third_party/triton/common/llvm_cl996124567.patch b/third_party/xla/third_party/triton/common/llvm_cl996124567.patch
index 05742f8bcc4..b19f3c07ff9 100644
--- a/third_party/xla/third_party/triton/common/llvm_cl996124567.patch
+++ b/third_party/xla/third_party/triton/common/llvm_cl996124567.patch
@@ -21,7 +21,7 @@
 
 --- a/third_party/amd/lib/TritonAMDGPUToLLVM/TritonGPUToLLVM.cpp
 +++ b/third_party/amd/lib/TritonAMDGPUToLLVM/TritonGPUToLLVM.cpp
-@@ -12,10 +12,10 @@
+@@ -10,10 +10,10 @@
  #include "mlir/Conversion/MathToLLVM/MathToLLVM.h"
  #include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"
  #include "mlir/Conversion/UBToLLVM/UBToLLVM.h"
@@ -33,15 +33,15 @@
  #include "mlir/Pass/Pass.h"
  #include "third_party/amd/include/Analysis/AMDGPUAllocation.h"
  #include "third_party/amd/include/Analysis/AxisInfoExt.h"
-@@ -222,16 +222,16 @@
+@@ -213,16 +213,16 @@
      mlir::triton::AMD::populateWarpIdOpToLLVMPattern(typeConverter, targetInfo,
                                                       patterns, commonBenefit);
  
 -    FailureOr<mlir::amdgpu::Chipset> maybeChipset =
--        mlir::amdgpu::Chipset::parse(this->gfxArch);
+-        mlir::amdgpu::Chipset::parse(targetInfo.getBaseArch());
 -    if (failed(maybeChipset)) {
 +    FailureOr<mlir::ROCDL::TargetInfo> maybeTargetInfo =
-+        mlir::ROCDL::TargetInfo::get(this->gfxArch);
++        mlir::ROCDL::TargetInfo::get(targetInfo.getBaseArch());
 +    if (failed(maybeTargetInfo)) {
        emitError(UnknownLoc::get(&getContext()),
                  "Invalid AMDGPU chipset name: " + this->gfxArch);
diff --git a/third_party/xla/third_party/triton/common/mig_limits_pytests.patch b/third_party/xla/third_party/triton/common/mig_limits_pytests.patch
index 92159dccab8..89f926edc5b 100644
--- a/third_party/xla/third_party/triton/common/mig_limits_pytests.patch
+++ b/third_party/xla/third_party/triton/common/mig_limits_pytests.patch
@@ -1,9 +1,10 @@
 --- a/python/test/conftest.py
 +++ b/python/test/conftest.py
-@@ -2,6 +2,18 @@ import pytest
+@@ -2,7 +2,18 @@ import pytest
  
  
  def pytest_configure(config):
+-    config.addinivalue_line("markers", "gsan_fine_granularity(reason): requires sub-16-byte GSan tracking")
 +    # In internal hermetic GPU test environments (e.g. a100, h100, b200 test targets), Forge enforces go/forge-accel-or-fail
 +    # to verify that GPU test binaries actively utilize the GPU. For pure AST/compiler tests (e.g. test_cache_determinism),
 +    # lazy CUDA initialization causes 0 GPU usage and fails the runner. We actively guarantee GPU utilization during pytest setup.
@@ -22,7 +23,7 @@
 
 --- a/python/test/unit/language/test_core.py
 +++ b/python/test/unit/language/test_core.py
-@@ -2137,7 +2137,13 @@ def test_atomic_cas(sem, num_ctas, dtype_str, device):
+@@ -2351,7 +2351,13 @@ def test_atomic_cas(sem, num_ctas, dtype_str, device):
  
  @pytest.mark.skipif(not is_cuda() or torch.cuda.get_device_capability()[0] < 9,
                      reason="num_ctas > 1 requires NVIDIA SM90+ (Hopper)")
@@ -37,7 +38,7 @@
  def test_scalar_atomic_cas_multicta_result(device):
  
      @triton.jit
-@@ -2153,8 +2159,8 @@ def test_scalar_atomic_cas_multicta_result(device):
+@@ -2367,8 +2373,8 @@ def test_scalar_atomic_cas_multicta_result(device):
      torch.testing.assert_close(output, torch.full_like(output, 7))
  
  
@@ -48,7 +49,7 @@
  @pytest.mark.parametrize("size", [1, 4, 16, 128, 512])
  def test_tensor_atomic_cas_multicta_result(size, device):
  
-@@ -2205,8 +2211,8 @@ def test_tensor_atomic_cas(sem, size, dtype_str, num_ctas, device):
+@@ -2419,8 +2425,8 @@ def test_tensor_atomic_cas(sem, size, dtype_str, num_ctas, device):
  
  
  @pytest.mark.interpreter
diff --git a/third_party/xla/third_party/triton/common/mma_limit_pred.patch b/third_party/xla/third_party/triton/common/mma_limit_pred.patch
index 66e9f8092b7..f92111236a3 100644
--- a/third_party/xla/third_party/triton/common/mma_limit_pred.patch
+++ b/third_party/xla/third_party/triton/common/mma_limit_pred.patch
@@ -3,7 +3,7 @@ fusions reported.
 
 --- a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
 +++ b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
-@@ -351,6 +351,16 @@ getWarpsPerTile(DotOpInterface dotOp, const ArrayRef<int64_t> shape,
+@@ -352,6 +352,16 @@ getWarpsPerTile(DotOpInterface dotOp, const ArrayRef<int64_t> shape,
  }
  
  static bool bwdFilter(Operation *op) {
diff --git a/third_party/xla/third_party/triton/common/no_accelerate_through_broadcast.patch b/third_party/xla/third_party/triton/common/no_accelerate_through_broadcast.patch
index 9ca761080eb..6413a49d2c9 100644
--- a/third_party/xla/third_party/triton/common/no_accelerate_through_broadcast.patch
+++ b/third_party/xla/third_party/triton/common/no_accelerate_through_broadcast.patch
@@ -2,7 +2,7 @@ Fix for b/405045790.
 
 --- a/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
 +++ b/lib/Dialect/TritonGPU/Transforms/AccelerateMatmul.cpp
-@@ -361,9 +361,12 @@ static bool bwdFilter(Operation *op) {
+@@ -362,9 +362,12 @@ static bool bwdFilter(Operation *op) {
        return false;
    }
  
diff --git a/third_party/xla/third_party/triton/common/old_ptxas.patch b/third_party/xla/third_party/triton/common/old_ptxas.patch
index 3a00162af2e..40bb19ea4da 100644
--- a/third_party/xla/third_party/triton/common/old_ptxas.patch
+++ b/third_party/xla/third_party/triton/common/old_ptxas.patch
@@ -34,7 +34,7 @@ On Blackwell architectures with PTX 8.8+, upstream Triton generates 256-bit glob
 
 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
-@@ -186,12 +186,9 @@ struct LoadStoreConversionBase {
+@@ -189,12 +189,9 @@ struct LoadStoreConversionBase {
      auto pointeeBitWidth = triton::getPointeeBitWidth(tensorTy);
      LDBG("getVectorSize contiguity = " << contiguity << " pointeeBitWidth = "
                                         << pointeeBitWidth);
diff --git a/third_party/xla/third_party/triton/common/series b/third_party/xla/third_party/triton/common/series
index 037a2d86211..126e281f788 100644
--- a/third_party/xla/third_party/triton/common/series
+++ b/third_party/xla/third_party/triton/common/series
@@ -3,7 +3,6 @@ mixed_precision_fix.patch
 mma_limit_pred.patch
 tc_disabled_kwidth_fix.patch
 enable_peer_access.patch
-f8e5m2_conversion.patch
 no_accelerate_through_broadcast.patch
 speed_up_int4_unpacking.patch
 verify_nvmma_encoding.patch
@@ -36,5 +35,5 @@ llvm_cl986804242.patch
 useStrictPropertiesInAssemblyFormat.patch
 llvm_cl991819890.patch
 fix_load_ops.patch
-fp64_large_k.patch
+cxx17_support.patch
 llvm_cl996124567.patch
diff --git a/third_party/xla/third_party/triton/common/speed_up_int4_unpacking.patch b/third_party/xla/third_party/triton/common/speed_up_int4_unpacking.patch
index 2fcc145dc93..2521bdcd4ea 100644
--- a/third_party/xla/third_party/triton/common/speed_up_int4_unpacking.patch
+++ b/third_party/xla/third_party/triton/common/speed_up_int4_unpacking.patch
@@ -27,7 +27,7 @@ an early stage on the XLA side with further optimizations from nVidia #24 and
 
 --- a/lib/Dialect/TritonGPU/Transforms/Utility.cpp
 +++ b/lib/Dialect/TritonGPU/Transforms/Utility.cpp
-@@ -1176,6 +1176,39 @@ swizzleDotOperandLike(RankedTensorType type, ttg::CGAEncodingAttr cgaLayout) {
+@@ -1196,6 +1196,39 @@ swizzleDotOperandLike(RankedTensorType type, ttg::CGAEncodingAttr cgaLayout) {
        type.getElementTypeBitWidth(), false);
  }
  
@@ -67,7 +67,7 @@ an early stage on the XLA side with further optimizations from nVidia #24 and
  // If all the transitive uses of the given value have are used by a convert to
  // the same dot operand encoding, return the shared encoding that needs to be
  // used to be compatible with users' layouts. If there are incompatible shared
-@@ -1210,14 +1243,21 @@ getSharedEncIfAllUsersAreDotEnc(Value val, bool &incompatible) {
+@@ -1230,14 +1263,21 @@ getSharedEncIfAllUsersAreDotEnc(Value val, bool &incompatible) {
        auto CGALayout = isa<ttg::LinearEncodingAttr>(dstTy.getEncoding())
                             ? ttg::getCGALayout(srcTy.getEncoding())
                             : ttg::getCGALayout(dstTy.getEncoding());
diff --git a/third_party/xla/third_party/triton/common/tmem_scale_padding_second_half.patch b/third_party/xla/third_party/triton/common/tmem_scale_padding_second_half.patch
index d65779f754a..54f4fb6251c 100644
--- a/third_party/xla/third_party/triton/common/tmem_scale_padding_second_half.patch
+++ b/third_party/xla/third_party/triton/common/tmem_scale_padding_second_half.patch
@@ -7,7 +7,7 @@ XLA.
 
 --- a/lib/Dialect/TritonNvidiaGPU/IR/Dialect.cpp
 +++ b/lib/Dialect/TritonNvidiaGPU/IR/Dialect.cpp
-@@ -143,6 +143,10 @@ TMemAllocation getTmemAllocSizes(MemDescType memDescType) {
+@@ -148,6 +148,10 @@ TMemAllocation getTmemAllocSizes(MemDescType memDescType) {
    if (ll.getBasis(kRow, llvm::Log2_32(16)) == ArrayRef{0, 0}) {
      nRow /= 2;
    }
diff --git a/third_party/xla/third_party/triton/common/tmemfix.patch b/third_party/xla/third_party/triton/common/tmemfix.patch
index dec09dbd808..75d0f89a0d4 100644
--- a/third_party/xla/third_party/triton/common/tmemfix.patch
+++ b/third_party/xla/third_party/triton/common/tmemfix.patch
@@ -3,7 +3,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
 
 --- a/lib/Analysis/Utility.cpp
 +++ b/lib/Analysis/Utility.cpp
-@@ -1142,7 +1142,7 @@ bool supportMMA(triton::DotOp op, int version) {
+@@ -1046,7 +1046,7 @@ bool supportMMA(triton::DotOp op, int version) {
      if (k < 256 / aElemTy.getIntOrFloatBitWidth())
        return false;
      if (!(retShapePerCTA[rank - 2] % 64 == 0 &&
@@ -12,7 +12,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
        return false;
      if (aElemTy.isF64() || bElemTy.isF64() ||
          retType.getElementType().isF64()) {
-@@ -1167,7 +1167,7 @@ bool supportMMA(triton::DotOp op, int version) {
+@@ -1071,7 +1071,7 @@ bool supportMMA(triton::DotOp op, int version) {
      if (rank == 3)
        return false;
      if (!(numWarps % 4 == 0 && retShapePerCTA[rank - 2] % 64 == 0 &&
@@ -24,7 +24,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
 
 --- a/lib/Dialect/TritonNvidiaGPU/IR/Dialect.cpp
 +++ b/lib/Dialect/TritonNvidiaGPU/IR/Dialect.cpp
-@@ -289,10 +289,7 @@ getDistributedLayoutForTmemLdSt(const LinearLayout &ll, TMemAccessAtom atom,
+@@ -294,10 +294,7 @@ getDistributedLayoutForTmemLdSt(const LinearLayout &ll, TMemAccessAtom atom,
  
    auto nColsTile = tile.getOutDimSize(rowColDims[1]);
    auto nColsLL = ll.getInDimSize(rowColDims[1]);
@@ -36,7 +36,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    auto kReg = StringAttr::get(ctx, "register");
    auto kLane = StringAttr::get(ctx, "lane");
    auto kWarp = StringAttr::get(ctx, "warp");
-@@ -309,7 +306,9 @@ getDistributedLayoutForTmemLdSt(const LinearLayout &ll, TMemAccessAtom atom,
+@@ -314,7 +311,9 @@ getDistributedLayoutForTmemLdSt(const LinearLayout &ll, TMemAccessAtom atom,
    // basis as that would disallow the resulting layout to be left-divisible by
    // the tile
    auto trivialBlock = LinearLayout::identity1D(1, kBlock, kBlock);
@@ -47,7 +47,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
                    .compose(ll)
                    .sublayout({kReg, kLane}, to_vector(ll.getOutDimNames()));
    if (instr32Rows) {
-@@ -317,7 +316,7 @@ getDistributedLayoutForTmemLdSt(const LinearLayout &ll, TMemAccessAtom atom,
+@@ -322,7 +321,7 @@ getDistributedLayoutForTmemLdSt(const LinearLayout &ll, TMemAccessAtom atom,
      // basis
      comp = comp.resizeInDim(kLane, comp.getInDimSize(kLane) / 2);
    }
@@ -56,7 +56,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
      return std::nullopt;
  
    // Fit the warp bases either tiling on the RHS or in row=16
-@@ -372,9 +371,8 @@ getDistributedLayoutForTmemLdSt(const LinearLayout &ll, TMemAccessAtom atom,
+@@ -377,9 +376,8 @@ getDistributedLayoutForTmemLdSt(const LinearLayout &ll, TMemAccessAtom atom,
    // Add CTAs as a trivial map
    auto nCTAs = ll.getInDimSize(kBlock);
    tile *= LinearLayout::identity1D(nCTAs, kBlock, kBlock);
@@ -68,7 +68,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    return ret;
  }
  
-@@ -397,6 +395,18 @@ DistributedEncodingTrait getDefaultLayoutForTmemLdSt(gpu::MemDescType memType,
+@@ -402,6 +400,18 @@ DistributedEncodingTrait getDefaultLayoutForTmemLdSt(gpu::MemDescType memType,
    auto *ctx = memType.getContext();
    auto layout = getDistributedLayoutForTmemLdSt(
        memType, TMemAccessAtom::I32x32b, numWarps);
@@ -87,12 +87,14 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    assert(layout);
    return LinearEncodingAttr::get(ctx, std::move(*layout));
  }
-@@ -470,6 +480,10 @@ getTmemCompatibleLayouts(Operation *op, RankedTensorType tensorType,
+@@ -475,6 +485,12 @@ getTmemCompatibleLayouts(Operation *op, RankedTensorType tensorType,
  bool isDistributedLayoutTMemCompatible(Operation *op,
                                         RankedTensorType tensorType,
                                         gpu::MemDescType memType) {
++  auto kReg = StringAttr::get(op->getContext(), "register");
 +  for (auto l : getTmemCompatibleLayouts(op, tensorType, memType)) {
-+    if (l == tensorType.getEncoding())
++    if (toLinearLayout(tensorType.getShape(), l).removeZeroBasesAlongDim(kReg) ==
++        toLinearLayout(tensorType).removeZeroBasesAlongDim(kReg))
 +      return true;
 +  }
    auto maxnreg = getContextualMaxNReg(op);
@@ -129,21 +131,21 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
      if (failed(info))
        return failure();
      if (bestContig > 1) {
-@@ -305,8 +306,9 @@ computeTMemLdStEncodingInfo(RankedTensorType regTy, MemDescType memTy,
-   cvt = LinearLayout(std::move(bases), cvt.getOutDims(),
-                      /*isSurjective=*/cvt.isSurjective());
- 
+@@ -313,8 +314,9 @@ computeTMemLdStEncodingInfo(RankedTensorType regTy, MemDescType memTy,
+   auto layout = computeTMemLdStLayout(regTy, memTy, emitError);
+   if (failed(layout))
+     return failure();
 +  bool isScales = isa<TensorMemoryScalesEncodingAttr>(memTy.getEncoding());
    int bitwidth = memTy.getElementTypeBitWidth();
--  return lowerTMemLdSt(cvt, maxnreg, bitwidth, emitError);
-+  return lowerTMemLdSt(cvt, maxnreg, bitwidth, isScales, emitError);
+-  return lowerTMemLdSt(*layout, maxnreg, bitwidth, emitError);
++  return lowerTMemLdSt(*layout, maxnreg, bitwidth, isScales, emitError);
  }
  
  bool supportsTMemLoadReduce(RankedTensorType regTy, MemDescType memTy,
 
 --- a/python/test/gluon/test_fpsan.py
 +++ b/python/test/gluon/test_fpsan.py
-@@ -1909,6 +1909,8 @@ _WMMA_DOT_CASES = [
+@@ -1933,6 +1933,8 @@ _WMMA_DOT_CASES = [
  
  
  def _native_mma_k(type_a: str) -> int:
@@ -155,7 +157,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
 
 --- a/python/test/unit/language/test_compile_only.py
 +++ b/python/test/unit/language/test_compile_only.py
-@@ -341,7 +341,7 @@ def test_compile_only_dot_mxfp() -> None:
+@@ -413,7 +413,7 @@ def test_compile_only_dot_mxfp() -> None:
      assert re.search(pattern, str(ttgir)), "The TTGIR does not match the expected pattern."
  
      ptx = k.asm["ptx"]
@@ -167,7 +169,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
 
 --- a/python/test/unit/language/test_core.py
 +++ b/python/test/unit/language/test_core.py
-@@ -3981,10 +3981,7 @@ def get_test_small_dots_cases():
+@@ -4388,10 +4388,7 @@ def get_test_small_dots_cases():
      if not is_cuda():
          return []
      return [(2, 4, 32, 1, False, False, 'None', 'ieee', 'float16', 'float32', 1, None),
@@ -179,7 +181,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
  
  
  @pytest.mark.interpreter
-@@ -4223,21 +4220,19 @@ def test_dot(M, N, K, num_warps, col_a, col_b, epilogue, input_precision, in_dty
+@@ -4630,21 +4627,19 @@ def test_dot(M, N, K, num_warps, col_a, col_b, epilogue, input_precision, in_dty
      # make sure ld/st are vectorized
      ptx = pgm.asm['ptx']
  
@@ -217,7 +219,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
 
 --- a/python/test/unit/language/test_matmul.py
 +++ b/python/test/unit/language/test_matmul.py
-@@ -408,6 +408,8 @@ def test_mxfp(BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES, nonKDim, NUM_WARPS, device)
+@@ -400,6 +400,8 @@ def test_mxfp(BLOCK_M, BLOCK_N, BLOCK_K, NUM_STAGES, nonKDim, NUM_WARPS, device)
              pytest.skip(f"CDNA4 does not support {BLOCK_K=} for scaled mfma {nonKDim=} variants")
          if (BLOCK_M == 256 or BLOCK_N == 256) and BLOCK_K == 256:
              pytest.skip("Config requires too much shared memory")
@@ -226,6 +228,15 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
  
      if not is_rubin() and BLOCK_N == 256 and BLOCK_K == 256:
          NUM_STAGES = min(NUM_STAGES, 2)
+@@ -1436,6 +1438,8 @@ def test_batched_mxfp(BATCH_SIZE, BLOCK_BATCH_SIZE, BLOCK_M, BLOCK_N, BLOCK_K, N
+     kernel_kwargs = {}
+     if is_hip():
+         kernel_kwargs["matrix_instr_nonkdim"] = nonKDim
++    if is_cuda() and BLOCK_BATCH_SIZE == 4 and BLOCK_M == 128 and BLOCK_N == 128 and NUM_STAGES == 1 and NUM_WARPS == 8:
++        kernel_kwargs["maxnreg"] = 255
+ 
+     out = batched_mxfp_matmul[grid](
+         a,
 
 --- a/python/tutorials/09-persistent-matmul.py
 +++ b/python/tutorials/09-persistent-matmul.py
@@ -241,7 +252,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
 
 --- a/test/Conversion/tritongpu_to_llvm_blackwell.mlir
 +++ b/test/Conversion/tritongpu_to_llvm_blackwell.mlir
-@@ -342,10 +342,10 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+@@ -341,10 +341,10 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
    // CHECK-LABEL: @tc_gen5_mma_scaled_a_tmem
    // CHECK: %[[TMEM_BASE:.+]] = llvm.ptrtoint %arg2 : !llvm.ptr<3> to i32
    // CHECK: %[[A_BASE:.+]] = llvm.ptrtoint %arg0 : !llvm.ptr<3> to i32
@@ -256,41 +267,36 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    tt.func @tc_gen5_mma_scaled_a_tmem(
        %a: !ttg.memdesc<128x256xf8E5M2, #tmem, #ttng.tensor_memory>,
        %b: !ttg.memdesc<256x128xf8E5M2, #shared, #ttg.shared_memory>,
-@@ -380,10 +380,10 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+@@ -381,8 +381,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+   // CHECK: %[[P0:.+]] = llvm.icmp "eq" %[[WID]], %[[C0]] : i32
    // CHECK: %[[P1:.+]] = llvm.and %{{.*}}, %[[P0]]  : i1
    // CHECK: llvm.cond_br %[[P1]]
-   // CHECK: %[[DESC0:.+]] = llvm.mlir.constant(144708608 : i32) : i32
 -  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]], %{{.+}}, %{{.+}}, %arg5
-+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]], %{{.+}}, %{{.+}}, %arg5
-   // CHECK: %[[TRUE:.+]] = llvm.mlir.constant(true) : i1
-   // CHECK: %[[DESC1:.+]] = llvm.mlir.constant(681579536 : i32) : i32
 -  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC1]], %{{.+}}, %{{.+}}, %[[TRUE]]
++  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]], %{{.+}}, %{{.+}}, %arg5
 +  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC1]], %{{.+}}, %{{.+}}, %[[TRUE]]
    tt.func @tc_gen5_mma_block_scale(%a: !ttg.memdesc<128x64xf8E4M3FN, #shared, #ttg.shared_memory>,
                         %b: !ttg.memdesc<32x128xi8, #shared1, #ttg.shared_memory>,
                         %c: !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>,
-@@ -414,13 +414,13 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
- module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
-   // CHECK-LABEL: @tc_gen5_mma_block_scale_fp4_a
-   // CHECK: %[[DESC0:.+]] = llvm.mlir.constant(144769664 : i32) : i32
+@@ -416,10 +416,10 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+   // CHECK-DAG: %[[DESC1:.+]] = llvm.mlir.constant(681640592 : i32) : i32
+   // CHECK-DAG: %[[DESC2:.+]] = llvm.mlir.constant(1218511520 : i32) : i32
+   // CHECK-DAG: %[[DESC3:.+]] = llvm.mlir.constant(1755382448 : i32) : i32
 -  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC0]]
-+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC0]]
-   // CHECK: %[[DESC1:.+]] = llvm.mlir.constant(681640592 : i32) : i32
 -  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC1]]
-+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC1]]
-   // CHECK: %[[DESC2:.+]] = llvm.mlir.constant(1218511520 : i32) : i32
 -  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC2]]
-+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC2]]
-   // CHECK: %[[DESC3:.+]] = llvm.mlir.constant(1755382448 : i32) : i32
 -  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC3]]
++  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC0]]
++  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC1]]
++  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC2]]
 +  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.scale_vec::1X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %{{.+}}, %{{.+}}, %{{.+}}, %[[DESC3]]
    tt.func @tc_gen5_mma_block_scale_fp4_a(%a: !ttg.memdesc<128x64xi8, #shared1, #ttg.shared_memory>,
                         %b: !ttg.memdesc<128x128xi8, #shared, #ttg.shared_memory>,
                         %c: !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>,
-@@ -583,8 +583,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+@@ -582,8 +582,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
    // CHECK-LABEL: @tc_gen5_mma_block_scale_nvfp4
+   // CHECK-DAG: %[[DESC0:.+]] = llvm.mlir.constant(138413184 : i32) : i32
    // CHECK: %[[TMEM_BASE:.+]] = llvm.ptrtoint %arg2 : !llvm.ptr<3> to i32
-   // CHECK: %[[DESC0:.+]] = llvm.mlir.constant(138413184 : i32) : i32
 -  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]]
 -  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]]
 +  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::4X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]]
@@ -298,19 +304,18 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    tt.func @tc_gen5_mma_block_scale_nvfp4(%a: !ttg.memdesc<128x64xi8, #shared, #ttg.shared_memory>,
                         %b: !ttg.memdesc<64x256xi8, #shared1, #ttg.shared_memory>,
                         %c: !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable>,
-@@ -616,9 +616,9 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
-   // CHECK-LABEL: @tc_gen5_mma_block_scale_mxfp4
+@@ -616,8 +616,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+   // CHECK-DAG: %[[DESC0:.+]] = llvm.mlir.constant(146801792 : i32) : i32
+   // CHECK-DAG: %[[DESC1:.+]] = llvm.mlir.constant(1220543648 : i32) : i32
    // CHECK-DAG: %[[TMEM_BASE:.+]] = llvm.ptrtoint %arg2 : !llvm.ptr<3> to i32
-   // CHECK: %[[DESC0:.+]] = llvm.mlir.constant(146801792 : i32) : i32
 -  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]]
-+  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4.block_scale.scale_vec::2X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]]
-   // CHECK: %[[DESC1:.+]] = llvm.mlir.constant(1220543648 : i32) : i32
 -  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32 [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC1]]
++  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4.block_scale.scale_vec::2X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC0]]
 +  // CHECK: @$7 tcgen05.mma.cta_group::1.kind::mxf4.block_scale.scale_vec::2X [ $0 + 0 ], $1, $2, $3, [ $4 + 0 ], [ $5 + 0 ], $6;", "r,l,l,r,r,r,b,b" %[[TMEM_BASE]], %{{.+}}, %{{.+}}, %[[DESC1]]
    tt.func @tc_gen5_mma_block_scale_mxfp4(%a: !ttg.memdesc<128x64xi8, #shared, #ttg.shared_memory>,
                         %b: !ttg.memdesc<64x256xi8, #shared1, #ttg.shared_memory>,
                         %c: !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable>,
-@@ -1178,7 +1178,7 @@ module attributes {"ttg.num-warps" = 8 : i32} {
+@@ -1237,7 +1237,7 @@ module attributes {"ttg.num-warps" = 8 : i32} {
  // CHECK-LABEL: @load_store_16x32bx1_broadcast
  tt.func private @load_store_16x32bx1_broadcast(%arg0: !ttg.memdesc<16x8xi8, #tmem_scales, #ttng.tensor_memory, mutable>, %arg1: tensor<16x8xi8, #linear>) {
    %true = arith.constant true
@@ -319,7 +324,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    ttng.tmem_store %arg1, %arg0, %true : tensor<16x8xi8, #linear> -> !ttg.memdesc<16x8xi8, #tmem_scales, #ttng.tensor_memory, mutable>
    tt.return
  }
-@@ -1193,7 +1193,7 @@ module attributes {"ttg.num-warps" = 4 : i32} {
+@@ -1252,7 +1252,7 @@ module attributes {"ttg.num-warps" = 4 : i32} {
  // CHECK-LABEL: @load_store_16x2_scales_uses_zero_second_half
  tt.func private @load_store_16x2_scales_uses_zero_second_half(%arg0: !ttg.memdesc<16x2xi8, #tmem_scales, #ttng.tensor_memory, mutable>, %arg1: tensor<16x2xi8, #linear>) {
    %true = arith.constant true
@@ -328,7 +333,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    // CHECK: llvm.return
    ttng.tmem_store %arg1, %arg0, %true : tensor<16x2xi8, #linear> -> !ttg.memdesc<16x2xi8, #tmem_scales, #ttng.tensor_memory, mutable>
    tt.return
-@@ -1257,22 +1257,22 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+@@ -1316,22 +1316,22 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
        %pred: i1) {
      // CHECK: %[[ACC_BASE:.+]] = llvm.ptrtoint %arg4 : !llvm.ptr<3> to i32
      // CHECK: %[[A_E4M3_BASE:.+]] = llvm.ptrtoint %arg0 : !llvm.ptr<3> to i32
@@ -359,7 +364,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
      // CHECK-SAME: %[[ACC_BASE]], %[[A_E5M2_BASE]]
      ttng.tc_gen5_mma_scaled %a_e4m3, %b_e5m2, %c, %scale_a, %scale_b, %useAcc, %pred lhs = e4m3 rhs = e5m2 :
         !ttg.memdesc<128x128xf8E4M3FN, #tmem_fp8_lhs, #ttng.tensor_memory>,
-@@ -1300,13 +1300,13 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+@@ -1359,13 +1359,13 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
        %pred: i1) {
      // CHECK-DAG: %[[ACC_BASE:.+]] = llvm.ptrtoint %arg2 : !llvm.ptr<3> to i32
      // CHECK-DAG: %[[A_BASE:.+]] = llvm.ptrtoint %arg0 : !llvm.ptr<3> to i32
@@ -377,7 +382,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
      // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]]
      ttng.tc_gen5_mma_scaled %a, %b, %c, %scale_a, %scale_b, %useAcc, %pred lhs = e2m1 rhs = e5m2 :
         !ttg.memdesc<128x64xi8, #tmem_fp4_padded_lhs, #ttng.tensor_memory>,
-@@ -1336,13 +1336,13 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+@@ -1395,13 +1395,13 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
        %pred: i1) {
      // CHECK-DAG: %[[ACC_BASE:.+]] = llvm.ptrtoint %arg2 : !llvm.ptr<3> to i32
      // CHECK-DAG: %[[A_BASE:.+]] = llvm.ptrtoint %arg0 : !llvm.ptr<3> to i32
@@ -395,69 +400,56 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
      // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]]
      ttng.tc_gen5_mma_scaled %a, %b, %c, %scale_a, %scale_b, %useAcc, %pred lhs = e2m1 rhs = e2m1 :
         !ttg.memdesc<128x128xi8, #tmem_fp4_dense_lhs, #ttng.tensor_memory>,
-@@ -1589,16 +1589,16 @@ module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttng.t
- #sb = #ttng.tensor_memory_scales_encoding<CGALayout = [[0, 0]]>
- module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttng.two-ctas" = true} {
-   // CHECK-LABEL: @tc_gen5_mma_mxfp4_k256
+@@ -1651,13 +1651,13 @@ module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttng.tw
+   // SM103-DAG: %[[K96_0:.+]] = llvm.mlir.constant(-1868561280 : i32)
+   // SM103-DAG: %[[K96_96:.+]] = llvm.mlir.constant(-257948496 : i32)
+   // SM103-DAG: %[[K64_192:.+]] = llvm.mlir.constant(1352664224 : i32)
 -  // SM100-COUNT-4: tcgen05.mma.cta_group::2.kind::mxf4.block_scale.block32
 +  // SM100-COUNT-4: tcgen05.mma.cta_group::2.kind::mxf4.block_scale.scale_vec::2X
    // SM100-NOT: tcgen05.mma
    // K96 at K=0 selects scale bytes 0..2; at K=96 it selects bytes 3..5.
-   // SM103: %[[K96_0:.+]] = llvm.mlir.constant(-1868561280 : i32)
 -  // SM103: tcgen05.mma.cta_group::2.kind::mxf4.block_scale.block32 {{.*}} %[[K96_0]],
-+  // SM103: tcgen05.mma.cta_group::2.kind::mxf4.block_scale.scale_vec::2X {{.*}} %[[K96_0]],
-   // SM103: %[[K96_96:.+]] = llvm.mlir.constant(-257948496 : i32)
 -  // SM103: tcgen05.mma.cta_group::2.kind::mxf4.block_scale.block32 {{.*}} %[[K96_96]],
++  // SM103: tcgen05.mma.cta_group::2.kind::mxf4.block_scale.scale_vec::2X {{.*}} %[[K96_0]],
 +  // SM103: tcgen05.mma.cta_group::2.kind::mxf4.block_scale.scale_vec::2X {{.*}} %[[K96_96]],
    // K64 at K=192 selects bytes 2..3 of the next scale word.
--  // SM103: %[[K64_192:.+]] = llvm.mlir.constant(1352664224 : i32)
 -  // SM103: tcgen05.mma.cta_group::2.kind::mxf4.block_scale.block32 {{.*}} %[[K64_192]],
-+  // SM103: %[[K64_192:.+]] = llvm.mlir.constant(278922368 : i32)
 +  // SM103: tcgen05.mma.cta_group::2.kind::mxf4.block_scale.scale_vec::2X {{.*}} %[[K64_192]],
    // SM103-NOT: tcgen05.mma
    // CHECK: llvm.return
    tt.func @tc_gen5_mma_mxfp4_k256(
-@@ -1617,25 +1617,25 @@ module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttng.t
-   }
- 
-   // CHECK-LABEL: @tc_gen5_mma_nvfp4_k512
+@@ -1681,20 +1681,20 @@ module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttng.tw
+   // SM103-DAG: %[[K64_192:.+]] = llvm.mlir.constant(272630912 : i32)
+   // SM103-DAG: %[[C16:.+]] = llvm.mlir.constant(16 : i32)
+   // SM103-DAG: %[[C32:.+]] = llvm.mlir.constant(32 : i32)
 -  // SM100-COUNT-8: tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.block16
 +  // SM100-COUNT-8: tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X
    // SM100-NOT: tcgen05.mma
    // SM103: %[[SCALE_A:.+]] = llvm.ptrtoint %arg3
    // SM103: %[[SCALE_B:.+]] = llvm.ptrtoint %arg4
-   // SM103: %[[K96_0:.+]] = llvm.mlir.constant(-1874852736 : i32)
 -  // SM103: tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.block16 {{.*}} %[[K96_0]],
-+  // SM103: tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X {{.*}} %[[K96_0]],
-   // SM103: %[[K96_96:.+]] = llvm.mlir.constant(-801110880 : i32)
 -  // SM103: tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.block16 {{.*}} %[[K96_96]],
-+  // SM103: tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X {{.*}} %[[K96_96]],
-   // SM103: %[[K64_192:.+]] = llvm.mlir.constant(272630912 : i32)
 -  // SM103: tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.block16 {{.*}} %[[K64_192]],
++  // SM103: tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X {{.*}} %[[K96_0]],
++  // SM103: tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X {{.*}} %[[K96_96]],
 +  // SM103: tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X {{.*}} %[[K64_192]],
    // The second sector begins at scale word 4: four columns per A word,
    // eight per B word for N=256.
--  // SM103: %[[C16:.+]] = llvm.mlir.constant(16 : i32)
--  // SM103: %[[A256:.+]] = llvm.add %[[SCALE_A]], %[[C16]]
--  // SM103: %[[C32:.+]] = llvm.mlir.constant(32 : i32)
--  // SM103: %[[B256:.+]] = llvm.add %[[SCALE_B]], %[[C32]]
+   // SM103: %[[A256:.+]] = llvm.add %[[SCALE_A]], %[[C16]]
+   // SM103: %[[B256:.+]] = llvm.add %[[SCALE_B]], %[[C32]]
 -  // SM103: tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.block16 {{.*}} %[[K96_0]], %[[A256]], %[[B256]],
 -  // SM103: tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.block16 {{.*}} %[[K96_96]],
 -  // SM103: tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.block16 {{.*}} %[[K64_192]],
-+  // SM103: %[[C12:.+]] = llvm.mlir.constant(12 : i32)
-+  // SM103: %[[A256:.+]] = llvm.add %[[SCALE_A]], %[[C12]]
-+  // SM103: %[[C24:.+]] = llvm.mlir.constant(24 : i32)
-+  // SM103: %[[B256:.+]] = llvm.add %[[SCALE_B]], %[[C24]]
 +  // SM103: tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X {{.*}} %[[K96_0]], %[[A256]], %[[B256]],
 +  // SM103: tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X {{.*}} %[[K96_96]],
 +  // SM103: tcgen05.mma.cta_group::2.kind::mxf4nvf4.block_scale.scale_vec::4X {{.*}} %[[K64_192]],
    // SM103-NOT: tcgen05.mma
    // CHECK: llvm.return
    tt.func @tc_gen5_mma_nvfp4_k512(
-@@ -1654,11 +1654,11 @@ module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttng.t
-   }
- 
-   // CHECK-LABEL: @tc_gen5_mma_mxfp4_k1024_subslice
+@@ -1716,11 +1716,11 @@ module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttng.tw
+   // SM103-DAG: %[[ABS_HIGH:.+]] = llvm.mlir.constant(1074806848 : i32)
+   // SM103-DAG: %[[ADDR_MASK:.+]] = llvm.mlir.constant(16383 : i64)
+   // SM103-DAG: %[[SHIFT:.+]] = llvm.mlir.constant(16 : i64)
 -  // SM100-COUNT-16: tcgen05.mma.cta_group::2.kind::mxf4.block_scale.block32
 +  // SM100-COUNT-16: tcgen05.mma.cta_group::2.kind::mxf4.block_scale.scale_vec::2X
    // SM100-NOT: tcgen05.mma
@@ -465,10 +457,10 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    // The high word sets absolute mode (bit 52) and matrix base offset zero.
 -  // SM103-COUNT-2: tcgen05.mma.cta_group::2.kind::mxf4.block_scale.block32
 +  // SM103-COUNT-2: tcgen05.mma.cta_group::2.kind::mxf4.block_scale.scale_vec::2X
-   // SM103: %[[ABS_HIGH:.+]] = llvm.mlir.constant(1074806848 : i32)
    // SM103: %[[ADESC:.+]] = llvm.bitcast
    // SM103: %[[ANEXT:.+]] = llvm.bitcast
-@@ -1673,8 +1673,8 @@ module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttng.t
+   // SM103: %[[AADDR:.+]] = llvm.and %[[ANEXT]], %[[ADDR_MASK]]
+@@ -1732,8 +1732,8 @@ module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttng.tw
    // SM103: %[[BADDR:.+]] = llvm.and %[[BNEXT]], %[[ADDR_MASK]]
    // SM103: %[[BLEAD:.+]] = llvm.shl %[[BADDR]], %[[SHIFT]]
    // SM103: %[[BABS:.+]] = llvm.or %[[BDESC]], %[[BLEAD]]
@@ -494,7 +486,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
 @@ -42,7 +42,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.targ
  module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:107"} {
    // CHECK-LABEL: @tc_gen5_mma_block_scale_nvfp4_sm107_m256
-   // CHECK: %[[DESC:.+]] = llvm.mlir.constant(136316040 : i32) : i32
+   // CHECK-DAG: %[[DESC:.+]] = llvm.mlir.constant(136316040 : i32) : i32
 -  // CHECK-COUNT-4: tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block32
 +  // CHECK-COUNT-4: tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::2X
    // CHECK-NOT: tcgen05.mma
@@ -509,7 +501,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    // CHECK-NOT: tcgen05.mma
    tt.func @tc_gen5_mma_scaled_fp4_padded_shared_linear_a(
        %a: !ttg.memdesc<256x64xi8, #shared_fp4_padded_sl, #smem, mutable>,
-@@ -275,13 +275,13 @@ module attributes {"ttg.num-ctas" = 8 : i32, "ttg.num-warps" = 4 : i32, ttg.targ
+@@ -333,13 +333,13 @@ module attributes {"ttg.num-ctas" = 8 : i32, "ttg.num-warps" = 4 : i32, ttg.targ
  module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:107"} {
    // CHECK-LABEL: @tc_gen5_mma_block_scale_ab_reuse
    // FP8 A takes twice the bytes of packed FP4 B, despite equal M/N tile sizes.
@@ -527,10 +519,10 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    // CHECK-SAME: "r,l,l,r,r,r,b,b" %{{[^,]+}}, %{{[^,]+}}, %{{[^,]+}}, %[[SC_ID]], %[[SC_A1]], %[[SC_B0]],
    // CHECK-NOT: tcgen05.mma
    tt.func @tc_gen5_mma_block_scale_ab_reuse(%a: !ttg.memdesc<256x64xf8E4M3FN, #shared, #ttg.shared_memory>,
-@@ -319,7 +319,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 8 : i32, ttg.targ
+@@ -377,7 +377,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 8 : i32, ttg.targ
+   // UE5M3 uses i8 block16 scales, so split-K lowering must keep the UE5M3
    // scale kind in the descriptor for both K=128 sub-instructions of a
    // logical BLOCK_K=256 tile.
-   // CHECK: %[[DESC:.+]] = llvm.mlir.constant(153093256 : i32) : i32
 -  // CHECK-COUNT-2: tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16
 +  // CHECK-COUNT-2: tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.scale_vec::4X
    // CHECK-NOT: tcgen05.mma
@@ -539,7 +531,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
 
 --- a/test/TritonGPU/accelerate-matmul.mlir
 +++ b/test/TritonGPU/accelerate-matmul.mlir
-@@ -1118,40 +1118,6 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.targ
+@@ -1166,40 +1166,6 @@ module attributes {"ttg.target" = "cuda:80", "ttg.num-warps" = 4 : i32} {
  
  // -----
  
@@ -583,7 +575,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
 
 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv5.cpp
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv5.cpp
-@@ -847,14 +847,14 @@ LogicalResult convertScaledDot(const LLVMTypeConverter &typeConverter,
+@@ -810,14 +810,14 @@ LogicalResult convertScaledDot(const LLVMTypeConverter &typeConverter,
    DotConversion dot;
    switch (mxfpInstKind) {
    case mxfpKind::mxf8f6f4:
@@ -602,14 +594,10 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
      break;
    }
  
-@@ -921,19 +921,23 @@ LogicalResult convertScaledDot(const LLVMTypeConverter &typeConverter,
-   dot.getInstOperands = [&](const DotConversion::InstDesc &desc, int m, int n,
+@@ -879,15 +879,18 @@ LogicalResult convertScaledDot(const LLVMTypeConverter &typeConverter,
                              int k) -> MMAInstOperands {
      auto [numRepM, numRepN, numRepK] = desc.repShape;
--    int numRepKWords = ceil<int>(blockK, dot.mmaSizeK * mmasPerScaleWord);
-+    int scaleFactorColsPerSet =
-+        getScaleFactorColsPerSet(mxfpInstKind, op, dot.mmaSizeK);
-+    int numRepKWords = ceil<int>(numRepK, 4 / scaleFactorColsPerSet);
+     int numRepKWords = ceil<int>(blockK, dot.mmaSizeK * mmasPerScaleWord);
 +    int aMultibuffers = cast<MemDescType>(op.getAScale().getType()).getRank() > 2 ? cast<MemDescType>(op.getAScale().getType()).getDimSize(0) : 1;
 +    int bMultibuffers = cast<MemDescType>(op.getBScale().getType()).getRank() > 2 ? cast<MemDescType>(op.getBScale().getType()).getDimSize(0) : 1;
      int numColPerScaleBlockA = ceil<int>(
@@ -623,13 +611,11 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
 +            .numCols / bMultibuffers,
          numRepN * numRepKWords);
 -    numColPerScaleBlockB = std::max(numColPerScaleBlockB, 2);
--    int baseK = dot.getKOffset(k) / dot.mmaSizeK;
--    int subWordIdx = baseK % mmasPerScaleWord;
--    int wordIdx = baseK / mmasPerScaleWord;
 +    if (desc.mmaSizeN >= 32)
 +      numColPerScaleBlockB = std::max(numColPerScaleBlockB, 2);
-+    int subWordIdx = k % (4 / scaleFactorColsPerSet);
-+    int wordIdx = k / (4 / scaleFactorColsPerSet);
+     int baseK = dot.getKOffset(k) / dot.mmaSizeK;
+     int subWordIdx = baseK % mmasPerScaleWord;
+     int wordIdx = baseK / mmasPerScaleWord;
      int scaleIdxA = linearizeScaleBlockIdx(op.getAScale(), m, wordIdx, numRepM,
                                             numRepKWords);
      int scaleIdxB = linearizeScaleBlockIdx(op.getBScale(), n, wordIdx, numRepN,
diff --git a/third_party/xla/third_party/triton/common/useStrictPropertiesInAssemblyFormat.patch b/third_party/xla/third_party/triton/common/useStrictPropertiesInAssemblyFormat.patch
index 19f2372b6a4..9d15e1d75bb 100644
--- a/third_party/xla/third_party/triton/common/useStrictPropertiesInAssemblyFormat.patch
+++ b/third_party/xla/third_party/triton/common/useStrictPropertiesInAssemblyFormat.patch
@@ -4,7 +4,7 @@ See b/567092262 for more details.
 
 --- a/include/triton/Dialect/Gluon/IR/GluonDialect.td
 +++ b/include/triton/Dialect/Gluon/IR/GluonDialect.td
-@@ -16,6 +16,7 @@
+@@ -16,6 +16,7 @@ def Gluon_Dialect : Dialect {
      "mlir::gpu::GPUDialect",
    ];
    let useDefaultAttributePrinterParser = 1;
@@ -15,7 +15,7 @@ See b/567092262 for more details.
 
 --- a/include/triton/Dialect/Triton/IR/TritonDialect.td
 +++ b/include/triton/Dialect/Triton/IR/TritonDialect.td
-@@ -52,6 +52,7 @@
+@@ -52,6 +52,7 @@ def Triton_Dialect : Dialect {
    let hasConstantMaterializer = 1;
    let useDefaultAttributePrinterParser = 1;
    let useDefaultTypePrinterParser = 1;
@@ -26,7 +26,7 @@ See b/567092262 for more details.
 
 --- a/include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td
 +++ b/include/triton/Dialect/TritonGPU/IR/TritonGPUDialect.td
-@@ -35,6 +35,7 @@
+@@ -35,6 +35,7 @@ def TritonGPU_Dialect : Dialect {
  
    let useDefaultTypePrinterParser = 1;
    let useDefaultAttributePrinterParser = 1;
@@ -37,7 +37,7 @@ See b/567092262 for more details.
 
 --- a/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentDialect.td
 +++ b/include/triton/Dialect/TritonInstrument/IR/TritonInstrumentDialect.td
-@@ -6,6 +6,7 @@
+@@ -6,6 +6,7 @@ include "mlir/IR/OpBase.td"
  def TritonInstrument_Dialect : Dialect {
    let name = "tti";
    let cppNamespace = "::mlir::triton::instrument";
@@ -48,7 +48,7 @@ See b/567092262 for more details.
 
 --- a/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUDialect.td
 +++ b/include/triton/Dialect/TritonNvidiaGPU/IR/TritonNvidiaGPUDialect.td
-@@ -43,6 +43,7 @@
+@@ -44,6 +44,7 @@ def TritonNvidiaGPU_Dialect : Dialect {
  
    let useDefaultAttributePrinterParser = 1;
    let useDefaultTypePrinterParser = 1;
@@ -59,7 +59,7 @@ See b/567092262 for more details.
 
 --- a/third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUDialect.td
 +++ b/third_party/amd/include/Dialect/TritonAMDGPU/IR/TritonAMDGPUDialect.td
-@@ -37,6 +37,7 @@
+@@ -37,6 +37,7 @@ def TritonAMDGPU_Dialect : Dialect {
    let dependentDialects = ["triton::TritonDialect"];
  
    let useDefaultAttributePrinterParser = 1;
@@ -70,7 +70,7 @@ See b/567092262 for more details.
 
 --- a/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUDialect.td
 +++ b/third_party/nvidia/include/Dialect/NVGPU/IR/NVGPUDialect.td
-@@ -35,6 +35,7 @@
+@@ -35,6 +35,7 @@ def NVGPU_Dialect : Dialect {
    let dependentDialects = [
      "mlir::LLVM::LLVMDialect"
    ];
@@ -81,7 +81,7 @@ See b/567092262 for more details.
 
 --- a/third_party/nvidia/include/Dialect/NVWS/IR/NVWSDialect.td
 +++ b/third_party/nvidia/include/Dialect/NVWS/IR/NVWSDialect.td
-@@ -39,6 +39,7 @@
+@@ -39,6 +39,7 @@ def NVWS_Dialect : Dialect {
  
    let useDefaultTypePrinterParser = 1;
    let useDefaultAttributePrinterParser = 1;
@@ -92,7 +92,7 @@ See b/567092262 for more details.
 
 --- a/third_party/proton/Dialect/include/Dialect/Proton/IR/ProtonDialect.td
 +++ b/third_party/proton/Dialect/include/Dialect/Proton/IR/ProtonDialect.td
-@@ -13,6 +13,7 @@
+@@ -13,6 +13,7 @@ def Proton_Dialect : Dialect {
    }];
  
    let dependentDialects = [];
@@ -103,7 +103,7 @@ See b/567092262 for more details.
 
 --- a/third_party/proton/Dialect/include/Dialect/ProtonGPU/IR/ProtonGPUDialect.td
 +++ b/third_party/proton/Dialect/include/Dialect/ProtonGPU/IR/ProtonGPUDialect.td
-@@ -22,6 +22,7 @@
+@@ -22,6 +22,7 @@ def ProtonGPU_Dialect : Dialect {
  
    let useDefaultTypePrinterParser = 1;
    let useDefaultAttributePrinterParser = 1;
diff --git a/third_party/xla/third_party/triton/common/verify_nvmma_encoding.patch b/third_party/xla/third_party/triton/common/verify_nvmma_encoding.patch
index 76f5c04c903..137f4d2d479 100644
--- a/third_party/xla/third_party/triton/common/verify_nvmma_encoding.patch
+++ b/third_party/xla/third_party/triton/common/verify_nvmma_encoding.patch
@@ -22,7 +22,7 @@ reshape outcomes.
  #include "triton/Tools/LayoutUtils.h"
  #include "llvm/ADT/STLExtras.h"
  #include "llvm/Support/Casting.h"
-@@ -576,7 +579,45 @@ LogicalResult MemDescReshapeOp::verify() {
+@@ -598,7 +601,45 @@ LogicalResult MemDescReshapeOp::verify() {
    return OpTrait::impl::verifyEquivalentMemDescType(expectedTy, dstType);
  }
  
@@ -69,7 +69,7 @@ reshape outcomes.
                                                     Attribute srcEnc,
                                                     ArrayRef<int64_t> dstShape,
                                                     Attribute &dstEnc) {
-@@ -597,6 +638,11 @@ static LogicalResult inferMemDescReshapeOpEncoding(ArrayRef<int64_t> srcShape,
+@@ -619,6 +660,11 @@ static LogicalResult inferMemDescReshapeOpEncoding(ArrayRef<int64_t> srcShape,
              ctx, mmaEncoding.getSwizzlingByteWidth(),
              mmaEncoding.getTransposed(), mmaEncoding.getElementBitWidth(),
              mmaEncoding.getFp4Padded(), CGALayout);
@@ -81,7 +81,7 @@ reshape outcomes.
          auto srcLL = toLinearLayout(srcShape, srcEnc);
          auto dstLL = toLinearLayout(dstShape, candidateEncoding);
          if (reshapeLayout(ctx, srcLL, dstShape) == dstLL) {
-@@ -636,8 +682,8 @@ LogicalResult MemDescReshapeOp::inferReturnTypes(
+@@ -658,8 +704,8 @@ LogicalResult MemDescReshapeOp::inferReturnTypes(
  
    Attribute dstEncoding;
    if (Attribute srcEnc = srcTy.getEncoding()) {
diff --git a/third_party/xla/third_party/triton/oss_only/build_files.patch b/third_party/xla/third_party/triton/oss_only/build_files.patch
index ee990271f61..c0f09c75e0f 100644
--- a/third_party/xla/third_party/triton/oss_only/build_files.patch
+++ b/third_party/xla/third_party/triton/oss_only/build_files.patch
@@ -5,7 +5,7 @@ diff --git b/BUILD b/BUILD
 new file mode 100644
 --- /dev/null
 +++ b/BUILD
-@@ -0,0 +1,1331 @@
+@@ -0,0 +1,1337 @@
 +# This package imports OpenAI's Triton (https://github.com/triton-lang/triton).
 +
 +# copybara:uncomment load("@rules:copybara.bzl", "copybara_config_test")
@@ -569,8 +569,10 @@ new file mode 100644
 +        "include/triton/Conversion/TritonGPUToLLVM/TargetInfoBase.h",  # Avoid circular dependency.
 +        "include/triton/Conversion/TritonGPUToLLVM/Utility.h",  # Avoid circular dependency.
 +        "lib/Dialect/TritonGPU/Transforms/DescriptorMemoryLayouts.cpp",  # Avoid circular dependency.
++        "lib/Dialect/TritonGPU/Transforms/TargetInfo.cpp",  # Avoid circular dependency.
 +        "lib/Dialect/TritonGPU/Transforms/Utility.cpp",  # Avoid circular dependency.
 +        "lib/Dialect/TritonNvidiaGPU/Transforms/ClusterBarrierMbarAllocator.cpp",  # Avoid circular dependency.
++        "lib/Dialect/TritonNvidiaGPU/Transforms/MBarrierUtilities.cpp",  # Avoid circular dependency.
 +        "lib/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.cpp",  # Avoid circular dependency.
 +    ],
 +    hdrs = glob([
@@ -584,6 +586,7 @@ new file mode 100644
 +        "include/triton/Analysis/*.h",
 +    ]) + [
 +        "include/triton/Dialect/TritonNvidiaGPU/Transforms/ClusterBarrierMbarAllocator.h",  # Avoid circular dependency.
++        "include/triton/Dialect/TritonNvidiaGPU/Transforms/MBarrierUtilities.h",  # Avoid circular dependency.
 +        "include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h",  # Avoid circular dependency.
 +        "include/triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h",  # Avoid circular dependency.
 +        "include/triton/Dialect/TritonGPU/Transforms/DescriptorMemoryLayouts.h",  # Avoid circular dependency.
@@ -773,6 +776,7 @@ new file mode 100644
 +        ],
 +        exclude = [
 +            "lib/Dialect/TritonGPU/Transforms/DescriptorMemoryLayouts.cpp",
++            "lib/Dialect/TritonGPU/Transforms/TargetInfo.cpp",
 +            "lib/Dialect/TritonGPU/Transforms/Utility.cpp",
 +        ],
 +    ) + [
@@ -836,6 +840,7 @@ new file mode 100644
 +        "@llvm-project//mlir:InferTypeOpInterface",
 +        "@llvm-project//mlir:LLVMDialect",
 +        "@llvm-project//mlir:LoopLikeInterface",
++        "@llvm-project//mlir:MathDialect",
 +        "@llvm-project//mlir:Pass",
 +        "@llvm-project//mlir:SCFDialect",
 +        "@llvm-project//mlir:SCFTransforms",
@@ -935,8 +940,8 @@ new file mode 100644
 +        # The file included into TritonDialects target which this target depends on.
 +        exclude = [
 +            "lib/Dialect/TritonNvidiaGPU/Transforms/ClusterBarrierMbarAllocator.cpp",
++            "lib/Dialect/TritonNvidiaGPU/Transforms/MBarrierUtilities.cpp",
 +            "lib/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.cpp",
-+            "lib/Dialect/TritonNvidiaGPU/Transforms/Utility.cpp",
 +        ],
 +    ),
 +    hdrs = glob(
@@ -947,9 +952,9 @@ new file mode 100644
 +        # The file included into TritonDialects target which this target depends on.
 +        exclude = [
 +            "include/triton/Dialect/TritonNvidiaGPU/Transforms/ClusterBarrierMbarAllocator.h",
++            "include/triton/Dialect/TritonNvidiaGPU/Transforms/MBarrierUtilities.h",
 +            "include/triton/Dialect/TritonNvidiaGPU/Transforms/Passes.h",
 +            "include/triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h",
-+            "include/triton/Dialect/TritonNvidiaGPU/Transforms/Utility.h",
 +        ],
 +    ),
 +    copts = select({
@@ -982,6 +987,7 @@ new file mode 100644
 +        "@llvm-project//mlir:ControlFlowInterfaces",
 +        "@llvm-project//mlir:FunctionInterfaces",
 +        "@llvm-project//mlir:IR",
++        "@llvm-project//mlir:LLVMDialect",
 +        "@llvm-project//mlir:LoopLikeInterface",
 +        "@llvm-project//mlir:Pass",
 +        "@llvm-project//mlir:SCFDialect",
@@ -1342,7 +1348,7 @@ diff --git b/python/BUILD b/python/BUILD
 new file mode 100644
 --- /dev/null
 +++ b/python/BUILD
-@@ -0,0 +1,215 @@
+@@ -0,0 +1,217 @@
 +# NOTE: Do not depend on any targets from this directory,
 +# but use //third_party/py/triton instead.
 +
@@ -1485,6 +1491,7 @@ new file mode 100644
 +        ":ir",
 +        ":passes",
 +        "@llvm-project//llvm:Analysis",
++        "@llvm-project//llvm:BitWriter",
 +        "@llvm-project//llvm:CodeGen",
 +        "@llvm-project//llvm:Core",
 +        "@llvm-project//llvm:IPO",
@@ -1539,6 +1546,7 @@ new file mode 100644
 +        "@triton//third_party/amd:TDMUtility",
 +        "@triton//third_party/amd:TritonAMDGPU",
 +        "@triton//third_party/amd:TritonAMDGPUTransforms",
++        "@triton//third_party/amd:TritonAMDMIRDAG",
 +        "@triton//third_party/nvidia:triton_nvidia",
 +        "@triton//third_party/proton:ProtonIR",
 +    ],
@@ -1563,7 +1571,7 @@ diff --git b/python/test/BUILD b/python/test/BUILD
 new file mode 100644
 --- /dev/null
 +++ b/python/test/BUILD
-@@ -0,0 +1,431 @@
+@@ -0,0 +1,441 @@
 +load("//third_party/py/pytest:pytest_defs.bzl", "pytest_multi_tests", "pytest_test")
 +
 +package(
@@ -1656,6 +1664,8 @@ new file mode 100644
 +    "unit/language/test_matmul.py",  # TODO(b/462585714):this test fails when run on B200 MIGs
 +    "gsan/test_gsan.py",  # TODO(b/495776594): these tests complain about misaligned access on B200.
 +    "gsan/test_gsan_failures.py",
++    "gsan/test_write_once.py",
++    "gsan/test_write_once_failures.py",
 +]
 +
 +test_suite(
@@ -1696,6 +1706,8 @@ new file mode 100644
 +    name = "ampere",
 +    size = "large",
 +    srcs = _conftest + [
++        "gsan/test_gsan.py",
++        "gsan/test_gsan_failures.py",
 +        "unit/language/test_core.py",
 +        "unit/language/test_mxfp.py",
 +    ],
@@ -1722,6 +1734,8 @@ new file mode 100644
 +    name = "hopper",
 +    size = "large",
 +    srcs = _conftest + [
++        "gsan/test_gsan.py",
++        "gsan/test_gsan_failures.py",
 +        "unit/language/test_core.py",
 +        "unit/language/test_mxfp.py",
 +    ],
@@ -1893,6 +1907,8 @@ new file mode 100644
 +    name = "blackwell",
 +    size = "large",
 +    srcs = _conftest + [
++        "gsan/test_gsan.py",
++        "gsan/test_gsan_failures.py",
 +        "unit/language/test_core.py",
 +        "unit/language/test_mxfp.py",
 +    ],
@@ -1918,6 +1934,8 @@ new file mode 100644
 +    name = "blackwell_full",
 +    size = "large",
 +    srcs = _conftest + [
++        "gsan/test_gsan.py",
++        "gsan/test_gsan_failures.py",
 +        "unit/language/test_core.py",
 +        "unit/language/test_mxfp.py",
 +    ],
@@ -2329,7 +2347,7 @@ diff --git b/third_party/amd/BUILD b/third_party/amd/BUILD
 new file mode 100644
 --- /dev/null
 +++ b/third_party/amd/BUILD
-@@ -0,0 +1,383 @@
+@@ -0,0 +1,403 @@
 +load("@rules_cc//cc:cc_library.bzl", "cc_library")
 +load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
 +
@@ -2398,6 +2416,7 @@ new file mode 100644
 +        "@llvm-project//mlir:InferTypeOpInterface",
 +        "@llvm-project//mlir:LLVMCommonConversion",
 +        "@llvm-project//mlir:LLVMDialect",
++        "@llvm-project//mlir:MathDialect",
 +        "@llvm-project//mlir:Pass",
 +        "@llvm-project//mlir:ROCDLDialect",
 +        "@llvm-project//mlir:SCFDialect",
@@ -2706,6 +2725,25 @@ new file mode 100644
 +    ],
 +)
 +
++cc_library(
++    name = "TritonAMDMIRDAG",
++    srcs = ["lib/Target/MIRDAG/DAGBuilder.cpp"],
++    hdrs = ["lib/Target/MIRDAG/DAGBuilder.h"],
++    copts = _no_unused_variable,
++    includes = [
++        "../..",
++        "lib/Target/MIRDAG",
++    ],
++    deps = [
++        "@llvm-project//llvm:Analysis",
++        "@llvm-project//llvm:CodeGen",
++        "@llvm-project//llvm:Core",
++        "@llvm-project//llvm:Support",
++        "@llvm-project//llvm:Target",
++        "@llvm-project//llvm:TargetParser",
++    ],
++)
++
 +filegroup(
 +    name = "all-files",
 +    srcs = glob(
@@ -2816,7 +2854,7 @@ diff --git b/third_party/nvidia/BUILD b/third_party/nvidia/BUILD
 new file mode 100644
 --- /dev/null
 +++ b/third_party/nvidia/BUILD
-@@ -0,0 +1,567 @@
+@@ -0,0 +1,568 @@
 +load("@rules_cc//cc:cc_library.bzl", "cc_library")
 +load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
 +
@@ -2946,6 +2984,7 @@ new file mode 100644
 +        ":NVGPUDialect",
 +        ":triton_conversion_triton_nvidia_gpu_to_llvm_passes_inc_gen",
 +        "@llvm-project//llvm:Support",
++        "@llvm-project//llvm:ir_headers",
 +        "@llvm-project//mlir:Analysis",
 +        "@llvm-project//mlir:ArithToLLVM",
 +        "@llvm-project//mlir:ArithTransforms",
diff --git a/third_party/xla/third_party/triton/workspace.bzl b/third_party/xla/third_party/triton/workspace.bzl
index 2b82ab2dc21..4ba5608b5fb 100644
--- a/third_party/xla/third_party/triton/workspace.bzl
+++ b/third_party/xla/third_party/triton/workspace.bzl
@@ -71,8 +71,8 @@ triton_archive = repository_rule(
 def repo():
     """Imports Triton."""
 
-    TRITON_COMMIT = "2074a1b8a5cd283ce584f8ebca825c4f0b215fc1"
-    TRITON_SHA256 = "3cf7fb4bf850a1e06dd96704a8407046c028129ebc9469a338856e30646d3f4b"
+    TRITON_COMMIT = "0ccad974ee86942e704ac06dafdc0c3f214eb229"
+    TRITON_SHA256 = "05f3bf327718825595400fc0d797f4f78ec470cef731e41eef07b9bb3ef9b3d6"
     triton_archive(
         name = "triton",
         sha256 = TRITON_SHA256,
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/compilation_pipeline_cuda.cc b/third_party/xla/xla/backends/gpu/codegen/triton/compilation_pipeline_cuda.cc
index 6637a0d6883..99df4d24ab2 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/compilation_pipeline_cuda.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/compilation_pipeline_cuda.cc
@@ -168,13 +168,18 @@ static void MakeLLIR(mlir::OpPassManager* pm,
   pm->addPass(
       mt::createAllocateSharedMemoryNvPass(cuda_cc_as_int, final_ptx_version));
   pm->addPass(ttng::createTritonTensorMemoryAllocationPass());
+  pm->addPass(ttng::createTritonGPUProxyFenceInsertion({cuda_cc_as_int}));
+  pm->addPass(
+      mt::createTritonNvidiaGPUMembar({cuda_cc_as_int, final_ptx_version}));
+  pm->addPass(ttng::createTritonNvidiaGPUTMemBarrierInsertionPass());
+  pm->addPass(ttng::createTritonNvidiaGPUOptimizeMBarrierArrivalsPass());
   // We could add a flag to XLA to optionally enable the following passes:
   // if "consan" in options.instrumentation_mode
   // pm->addPass(mt::instrument::createTritonInstrumentConcurrencySanitizer());
   // pm->addPass(mlir::triton::gluon::createGluonCanonicalize());
   // pm->addPass(mlir::createCSEPass());
-  pm->addPass(ttng::createTritonGPUProxyFenceInsertion({cuda_cc_as_int}));
-  pm->addPass(ttng::createTritonNvidiaGPUTMemBarrierInsertionPass());
+  pm->addPass(ttng::createTritonNvidiaGPUClusterBarrierMbarAllocatorPass());
+  pm->addPass(mt::gpu::createTritonGPUGlobalScratchAllocationPass());
   pm->addPass(
       mt::createConvertTritonGPUToLLVMPass(cuda_cc_as_int, final_ptx_version));
   pm->addPass(mt::createInitializeWSClusterBarriers(
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc
index 71e881d6dbe..29a1327e315 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc
@@ -1484,13 +1484,6 @@ ENTRY entry_computation {
 }
 
 TEST_F(TritonEmitterTest, ConvertF16ToF8E5M2Exhaustive) {
-  // TODO(b/396595945): enable post-Ampere once Triton respects RTNE semantics
-  // on H100.
-  if (auto cc = GpuComputeCapability().cuda_compute_capability();
-      cc && cc->IsAtLeastHopper()) {
-    GTEST_SKIP() << "Skipping tests above Ampere, Triton's conversion isn't "
-                    "always correct";
-  }
   if (GpuComputeCapability().IsRocm()) {
     GTEST_SKIP() << "Triton's F16 to F8E5M2 conversion doesn't preserve "
                     "infinities on ROCm";
@@ -1518,8 +1511,14 @@ ENTRY entry_computation {
 
   std::vector<Eigen::half> all_f16_values;
   for (int i = 0; i < 65536; i++) {
-    all_f16_values.push_back(
-        Eigen::numext::bit_cast<Eigen::half>(static_cast<uint16_t>(i)));
+    Eigen::half v =
+        Eigen::numext::bit_cast<Eigen::half>(static_cast<uint16_t>(i));
+    // TODO(b/396595945): Triton's F16 to F8E5M2 conversion uses satfinite
+    // semantics and saturates values that overflow to infinity.
+    if (Eigen::numext::isinf(static_cast<tsl::float8_e5m2>(v))) {
+      v = Eigen::half(0.0f);
+    }
+    all_f16_values.push_back(v);
   }
 
   std::string hlo_text =
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_squeeze_dims_pass.cc b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_squeeze_dims_pass.cc
index 864697ffab6..b0d41570bea 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_squeeze_dims_pass.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_squeeze_dims_pass.cc
@@ -130,13 +130,19 @@ RankedTensorType SqueezeTensorType(RankedTensorType type,
   if (encoding) {
     auto inferLayoutInterface =
         cast<DialectInferLayoutInterface>(&encoding.getDialect());
-    Attribute dst_encoding = dst_encoding_hint;
+    Attribute dst_encoding;
     [[maybe_unused]] LogicalResult result =
         inferLayoutInterface->inferReshapeOpEncoding(
             type.getShape(), encoding, shape, dst_encoding,
             /*allowReorder=*/false, /*loc=*/std::nullopt);
     CHECK(succeeded(result));
-    encoding = dst_encoding;
+    if (dst_encoding_hint &&
+        succeeded(inferLayoutInterface->verifyLayoutsAreEqual(
+            shape, dst_encoding, dst_encoding_hint, /*loc=*/std::nullopt))) {
+      encoding = dst_encoding_hint;
+    } else {
+      encoding = dst_encoding;
+    }
   }
   return RankedTensorType::get(shape, type.getElementType(), encoding);
 }
diff --git a/third_party/xla/xla/service/gpu/autotuning/autotune_cache_key.h b/third_party/xla/xla/service/gpu/autotuning/autotune_cache_key.h
index 3e225fa78b7..b9370cb6647 100644
--- a/third_party/xla/xla/service/gpu/autotuning/autotune_cache_key.h
+++ b/third_party/xla/xla/service/gpu/autotuning/autotune_cache_key.h
@@ -34,7 +34,7 @@ class AutotuneCacheKey {
   // changes that may affect the autotuning results.
   // To prevent accidental merges of concurrent increments, update the comment
   // to explain why the version is bumped.
-  static constexpr int kCurrentVersion = 54;  // Triton integration 1.39
+  static constexpr int kCurrentVersion = 55;  // Triton integration 1.40
 
   AutotuneCacheKey(const se::DeviceDescription& device_description,
                    const HloInstruction& instruction,

From 0e2d96d26c9621dbe1998ba740b5e2e5d14ffec5 Mon Sep 17 00:00:00 2001
From: Nikhita Hegde <nhegde@habana.ai>
Date: Fri, 9 Oct 2026 05:19:12 -0700
Subject: [PATCH 0355/1000] PR #47829: [XLA:GPU] Add a common helper for
 invoking oneMKL calls

Imported from GitHub PR https://github.com/openxla/xla/pull/47829

Adds common code for invoking oneMKL calls, to be used by the oneMKL-based
implementations such as FFT, solvers, and complex GEMMs in upcoming PRs.
Copybara import of the project:

--
214ddadc704aff0f3e592e33d88861e07feb7815 by nikhitah <nikhita.hegde@intel.com>:

[XLA:GPU] Add a common helper for invoking oneMKL calls

Common code for invoking onemkl calls for the onemkl based implementations such as FFT, solvers, Complex gemms in upcoming PRs.

--
08413c4cd041e5aae0b0fc02b675de7f15ecf205 by nikhitah <nikhita.hegde@intel.com>:

Fix review comments

--
b795847de623078f1f5bcdedaba774f9a69346eb by nikhitah <nikhita.hegde@intel.com>:

Add one  unit test

--
b2032d300c78089efea5599b834ead305d394938 by nikhitah <nikhita.hegde@intel.com>:

Fix review comments

--
d4e2e3f2188abef67cb77eb64947c8282e38d20e by ag.ramesh <ag.ramesh@intel.com>:

Fixed missing new line

--
8438ad88d9913a7855ba92cfbdcb7d30a1161dab by nikhitah <nikhita.hegde@intel.com>:

Use xla_test for onemkl_util_test instead of linking @oneapi//:libs

Merging this change closes #47829

PiperOrigin-RevId: 996471374
---
 .../xla/xla/stream_executor/sycl/BUILD        | 33 ++++++++
 .../xla/stream_executor/sycl/onemkl_util.h    | 81 +++++++++++++++++++
 .../stream_executor/sycl/onemkl_util_test.cc  | 52 ++++++++++++
 3 files changed, 166 insertions(+)
 create mode 100644 third_party/xla/xla/stream_executor/sycl/onemkl_util.h
 create mode 100644 third_party/xla/xla/stream_executor/sycl/onemkl_util_test.cc

diff --git a/third_party/xla/xla/stream_executor/sycl/BUILD b/third_party/xla/xla/stream_executor/sycl/BUILD
index 6f6ff17ac42..384804af986 100644
--- a/third_party/xla/xla/stream_executor/sycl/BUILD
+++ b/third_party/xla/xla/stream_executor/sycl/BUILD
@@ -584,6 +584,39 @@ cc_library(
     alwayslink = True,
 )
 
+cc_library(
+    name = "onemkl_util_headers",
+    hdrs = ["onemkl_util.h"],
+    tags = [
+        "gpu",
+        "oneapi-only",
+    ],
+    deps = [
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@local_config_sycl//sycl:sycl_headers",
+    ],
+)
+
+xla_test(
+    name = "onemkl_util_test",
+    srcs = ["onemkl_util_test.cc"],
+    backends = ["gpu"],
+    tags = [
+        "gpu",
+        "oneapi-only",
+    ],
+    deps = [
+        ":onemkl_util_headers",
+        "//xla/tsl/platform:test",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_googletest//:gtest_main",
+        "@local_config_sycl//sycl:sycl_headers",
+    ],
+)
+
 cc_library(
     name = "onednn_util",
     srcs = ["onednn_util.cc"],
diff --git a/third_party/xla/xla/stream_executor/sycl/onemkl_util.h b/third_party/xla/xla/stream_executor/sycl/onemkl_util.h
new file mode 100644
index 00000000000..d260854be19
--- /dev/null
+++ b/third_party/xla/xla/stream_executor/sycl/onemkl_util.h
@@ -0,0 +1,81 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_STREAM_EXECUTOR_SYCL_ONEMKL_UTIL_H_
+#define XLA_STREAM_EXECUTOR_SYCL_ONEMKL_UTIL_H_
+
+#include <exception>
+#include <type_traits>
+#include <utility>
+
+#include "oneapi/mkl.hpp"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+
+namespace stream_executor::sycl {
+
+// Helper functions to make oneMKL calls. For MKL calls that take a queue it
+// returns a status if failed or an event object if it succeeds. The caller
+// can then check for ok, and use the returend object as needed.
+// Example to call oneapi::mkl::blas::trsm
+//  absl::StatusOr<::sycl::event> result = ExecMklFunc([&] {
+//    return oneapi::mkl::blas::trsm(q, left_right, upper_lower, trans,
+//                                   unit_diag, m, n, alpha, a, lda, b, ldb);
+//  });
+//
+// if (result.ok()) {
+//   result->DoSomethingCool();
+// } else {
+//   LOG(ERROR) << result.status();
+// }
+// For MKL calls that that do not take a queue of their own -- e.g. the DFT
+// ones, which run on the queue their descriptor was committed to -- are
+// called the same way:
+//  absl::Status status = ExecMklFunc([&] {
+//    oneapi::mkl::dft::compute_forward(descriptor, input, output);
+//  });
+// It returns a plain absl::Status.
+
+// The type ExecMklFunc() returns for a oneMKL routine returning `R`.
+template <typename R>
+using MklResult = std::conditional_t<std::is_void_v<R>, absl::Status,
+                                     absl::StatusOr<std::decay_t<R>>>;
+
+template <typename Callable>
+MklResult<std::invoke_result_t<Callable>> ExecMklFunc(Callable&& mkl_func) {
+  using Result = std::invoke_result_t<Callable>;
+  try {
+    if constexpr (std::is_void_v<Result>) {
+      std::forward<Callable>(mkl_func)();
+      return absl::OkStatus();
+    } else {
+      return std::forward<Callable>(mkl_func)();
+    }
+  } catch (oneapi::mkl::exception const& e) {
+    return absl::InternalError(absl::StrCat("Mkl exception: ", e.what()));
+  } catch (const ::sycl::exception& e) {
+    return absl::InternalError(absl::StrCat(
+        "SYCL exception: ", e.what(), " [sycl_code=", e.code().value(), "]"));
+  } catch (const std::exception& e) {
+    return absl::InternalError(absl::StrCat("Exception: ", e.what()));
+  } catch (...) {
+    return absl::InternalError("Unknown (non-std::exception) thrown");
+  }
+}
+
+}  // namespace stream_executor::sycl
+
+#endif  // XLA_STREAM_EXECUTOR_SYCL_ONEMKL_UTIL_H_
diff --git a/third_party/xla/xla/stream_executor/sycl/onemkl_util_test.cc b/third_party/xla/xla/stream_executor/sycl/onemkl_util_test.cc
new file mode 100644
index 00000000000..f11e32749d7
--- /dev/null
+++ b/third_party/xla/xla/stream_executor/sycl/onemkl_util_test.cc
@@ -0,0 +1,52 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/stream_executor/sycl/onemkl_util.h"
+
+#include "oneapi/mkl.hpp"
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "xla/tsl/platform/test.h"
+
+namespace stream_executor::sycl {
+namespace {
+
+using ::absl_testing::IsOk;
+using ::absl_testing::IsOkAndHolds;
+using ::absl_testing::StatusIs;
+using ::testing::HasSubstr;
+
+TEST(ExecMklFuncTest, VoidCallReturnsOkStatus) {
+  bool called = false;
+  EXPECT_THAT(ExecMklFunc([&] { called = true; }), IsOk());
+  EXPECT_TRUE(called);
+}
+
+TEST(ExecMklFuncTest, ReturnValueIsForwarded) {
+  EXPECT_THAT(ExecMklFunc([] { return 42; }), IsOkAndHolds(42));
+}
+
+TEST(ExecMklFuncTest, MklExceptionBecomesInternalError) {
+  absl::Status status = ExecMklFunc(
+      [] { throw oneapi::mkl::exception("blas", "trsm", "lda too small"); });
+  EXPECT_THAT(status, StatusIs(absl::StatusCode::kInternal,
+                               HasSubstr("Mkl exception: oneapi::mkl::blas::"
+                                         "trsm: lda too small")));
+}
+
+}  // namespace
+}  // namespace stream_executor::sycl

From 5daf5f25b82bf37d136b7a60ee2afa4b586676f9 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 9 Oct 2026 05:25:50 -0700
Subject: [PATCH 0356/1000] Add chunked staging for host-device transfers in
 PJRT StreamExecutor.

In Confidential Computing mode, host-to-device and device-to-host transfers are staged through pinned host memory. Large transfers previously required allocating a single staging buffer matching the full transfer size, which could lead to host memory exhaustion or fragmentation.

This change splits all staged transfers into fixed-size chunks (defaulting to 16 MiB and configurable via XLA_PJRT_GPU_CC_STAGING_CHUNK_MB) and pipelines them across a bounded pool of up to 4 staging slots, overlapping host copies with device DMAs.

PiperOrigin-RevId: 996474035
---
 .../xla/pjrt/gpu/se_gpu_pjrt_client_test.cc   | 103 +++++++++++
 third_party/xla/xla/pjrt/se/BUILD             |   3 +
 .../pjrt/se/pjrt_stream_executor_client.cc    |  29 +++
 .../xla/pjrt/se/pjrt_stream_executor_client.h |  36 +++-
 third_party/xla/xla/pjrt/se/se_raw_buffer.cc  | 172 +++++++++++++++++-
 5 files changed, 338 insertions(+), 5 deletions(-)

diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
index acd040d5fea..50815923322 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
@@ -1656,6 +1656,32 @@ TEST(StreamExecutorGpuClientTest, ShouldStageHostToDeviceTransfersSetToTrue) {
       *literal, LiteralUtil::CreateR1<float>(std::vector<float>(1024, 1.0f))));
 }
 
+TEST(StreamExecutorGpuClientTest, ShouldStageLargeTransfersWhenChunking) {
+  GpuClientOptions options = GetTestGpuClientOptions();
+  options.should_stage_host_to_device_transfers = true;
+  TF_ASSERT_OK_AND_ASSIGN(auto client, GetStreamExecutorGpuClient(options));
+  auto* raw_client = absl::down_cast<PjRtStreamExecutorRawClient*>(
+      absl::down_cast<CommonPjRtClient*>(client.get())->raw_client());
+
+  // Only the pointer value is used; no 1 GiB host buffer is allocated.
+  std::vector<char> host(16);
+  constexpr int64_t kOneGiB = int64_t{1} << 30;
+  ASSERT_FALSE(raw_client->IsDmaMapped(host.data(), kOneGiB));
+
+  // Without chunking, large transfers are not staged (unchanged behavior).
+  EXPECT_FALSE(
+      raw_client->ShouldStageHostToDeviceTransfers(host.data(), kOneGiB));
+  EXPECT_FALSE(
+      raw_client->ShouldStageHostToDeviceTransfers(host.data(), kOneGiB + 1));
+
+  // With chunking, staging uses bounded chunk buffers, so any size is staged.
+  raw_client->SetStagingChunkingForTesting(/*force=*/true, /*chunk_size=*/4096);
+  EXPECT_TRUE(
+      raw_client->ShouldStageHostToDeviceTransfers(host.data(), kOneGiB));
+  EXPECT_TRUE(
+      raw_client->ShouldStageHostToDeviceTransfers(host.data(), kOneGiB + 1));
+}
+
 TEST(StreamExecutorGpuClientTest, ShouldStageHostToDeviceTransfersSetToFalse) {
   GpuClientOptions options_no_staging = GetTestGpuClientOptions();
   options_no_staging.should_stage_host_to_device_transfers = false;
@@ -1687,6 +1713,83 @@ TEST(StreamExecutorGpuClientTest, ShouldStageHostToDeviceTransfersSetToFalse) {
       *literal, LiteralUtil::CreateR1<float>(std::vector<float>(1024, 1.0f))));
 }
 
+constexpr int64_t kTestStagingChunkSize = 4 * 1024;
+
+// Returns `size` bytes whose content differs between any two chunks, so that
+// misplaced or reordered chunks are detected.
+std::vector<uint8_t> MakeStagingTestPattern(int64_t size) {
+  std::vector<uint8_t> pattern(size);
+  for (int64_t i = 0; i < size; ++i) {
+    pattern[i] = static_cast<uint8_t>((i * 131) ^ (i >> 8));
+  }
+  return pattern;
+}
+
+// Copies `transfer_size` bytes host->device->host at `offset` of a device
+// buffer with chunked host staging forced on, and checks the bytes survive.
+void ExpectChunkedStagingRoundTrip(int64_t offset, int64_t transfer_size) {
+  GpuClientOptions options = GetTestGpuClientOptions();
+  options.should_stage_host_to_device_transfers = true;
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtClient> pjrt_client,
+                          GetStreamExecutorGpuClient(options));
+  auto* client = absl::down_cast<CommonPjRtClient*>(pjrt_client.get());
+  auto* raw_client =
+      absl::down_cast<PjRtStreamExecutorRawClient*>(client->raw_client());
+  raw_client->SetStagingChunkingForTesting(/*force=*/true,
+                                           kTestStagingChunkSize);
+
+  std::vector<uint8_t> src = MakeStagingTestPattern(transfer_size);
+  ASSERT_TRUE(
+      raw_client->ShouldStageHostToDeviceTransfers(src.data(), transfer_size));
+  TF_ASSERT_OK_AND_ASSIGN(
+      PjRtMemorySpace * memory_space,
+      client->addressable_devices()[0]->default_memory_space());
+  TF_ASSERT_OK_AND_ASSIGN(
+      PjRtRawBufferRef raw_buffer,
+      client->AllocateRawBuffer(memory_space, offset + transfer_size,
+                                /*retry_on_oom=*/true,
+                                /*allocate_after=*/{}));
+
+  TF_ASSERT_OK(
+      raw_buffer->CopyRawHostToDevice(src.data(), offset, transfer_size)
+          .Await());
+  std::vector<uint8_t> dst(transfer_size, 0);
+  TF_ASSERT_OK(
+      raw_buffer->CopyRawDeviceToHost(dst.data(), offset, transfer_size)
+          .Await());
+  EXPECT_EQ(dst, src);
+}
+
+TEST(StreamExecutorGpuClientTest, ChunkedStagingZeroBytes) {
+  ExpectChunkedStagingRoundTrip(/*offset=*/0, 0);
+}
+
+TEST(StreamExecutorGpuClientTest, ChunkedStagingOneByte) {
+  ExpectChunkedStagingRoundTrip(/*offset=*/0, 1);
+}
+
+TEST(StreamExecutorGpuClientTest, ChunkedStagingSmallerThanChunk) {
+  ExpectChunkedStagingRoundTrip(/*offset=*/0, kTestStagingChunkSize / 2);
+}
+
+TEST(StreamExecutorGpuClientTest, ChunkedStagingEqualToChunk) {
+  ExpectChunkedStagingRoundTrip(/*offset=*/0, kTestStagingChunkSize);
+}
+
+TEST(StreamExecutorGpuClientTest, ChunkedStagingMultipleOfChunk) {
+  // More chunks than staging slots, so slots are reused.
+  ExpectChunkedStagingRoundTrip(/*offset=*/0, 8 * kTestStagingChunkSize);
+}
+
+TEST(StreamExecutorGpuClientTest, ChunkedStagingNonMultipleOfChunk) {
+  ExpectChunkedStagingRoundTrip(/*offset=*/0, 10 * kTestStagingChunkSize + 123);
+}
+
+TEST(StreamExecutorGpuClientTest, ChunkedStagingWithNonzeroOffset) {
+  ExpectChunkedStagingRoundTrip(/*offset=*/123,
+                                10 * kTestStagingChunkSize + 45);
+}
+
 TEST(StreamExecutorGpuClientTest, BufferFromHostBufferPinnedMemory) {
   TF_ASSERT_OK_AND_ASSIGN(
       auto client, GetStreamExecutorGpuClient(GetTestGpuClientOptions()));
diff --git a/third_party/xla/xla/pjrt/se/BUILD b/third_party/xla/xla/pjrt/se/BUILD
index 47d24622f1a..9356860796a 100644
--- a/third_party/xla/xla/pjrt/se/BUILD
+++ b/third_party/xla/xla/pjrt/se/BUILD
@@ -275,7 +275,9 @@ cc_library(
         "//xla/service/gpu:gpu_executable_run_options",
         "//xla/stream_executor:device_address",
         "//xla/stream_executor:device_address_allocator",
+        "//xla/stream_executor:event",
         "//xla/stream_executor:stream",
+        "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor/abi:executable_abi_version",
         "//xla/stream_executor/integrations:tf_allocator_adapter",
         "//xla/tsl/concurrency:async_value",
@@ -285,6 +287,7 @@ cc_library(
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:logging",
         "//xla/tsl/platform:statusor",
+        "//xla/tsl/util:env_var",
         "//xla/util/split_proto:split_executable_and_options_writer",
         "//xla/util/split_proto:split_proto_reader",
         "@com_google_absl//absl/algorithm:container",
diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc
index 65ec1b2e934..693fb557690 100644
--- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc
+++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc
@@ -71,6 +71,7 @@ limitations under the License.
 #include <cstring>
 #include <functional>
 #include <initializer_list>
+#include <limits>
 #include <memory>
 #include <optional>
 #include <string>
@@ -174,6 +175,7 @@ limitations under the License.
 #include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/statusor.h"
 #include "xla/tsl/platform/threadpool.h"
+#include "xla/tsl/util/env_var.h"
 #include "xla/util.h"
 #include "xla/util/split_proto/split_executable_and_options_writer.h"
 #include "xla/util/split_proto/split_proto_reader.h"
@@ -186,6 +188,30 @@ limitations under the License.
 
 namespace xla {
 
+namespace {
+
+constexpr int64_t kDefaultStagingChunkSizeMb = 16;
+
+// Returns the chunk size, in bytes, used for chunked staging of host<->device
+// transfers. Overridable via the XLA_PJRT_GPU_CC_STAGING_CHUNK_MB environment
+// variable.
+int64_t GetStagingChunkSizeFromEnv() {
+  constexpr absl::string_view kEnvVar = "XLA_PJRT_GPU_CC_STAGING_CHUNK_MB";
+  int64_t chunk_size_mb;
+  absl::Status status = tsl::ReadInt64FromEnvVar(
+      kEnvVar, kDefaultStagingChunkSizeMb, &chunk_size_mb);
+  if (!status.ok() || chunk_size_mb <= 0 ||
+      chunk_size_mb > (std::numeric_limits<int64_t>::max() >> 20)) {
+    LOG(WARNING) << "Ignoring invalid " << kEnvVar
+                 << " (must be a positive integer number of MiB); using the "
+                 << "default of " << kDefaultStagingChunkSizeMb << " MiB.";
+    chunk_size_mb = kDefaultStagingChunkSizeMb;
+  }
+  return chunk_size_mb << 20;
+}
+
+}  // namespace
+
 template <typename MemorySpaceKind>
 static bool IsMemorySpaceKind(const PjRtMemorySpace* memory_space) {
   return memory_space->kind_id() == MemorySpaceKind::kKindId;
@@ -297,6 +323,9 @@ PjRtStreamExecutorRawClient::PjRtStreamExecutorRawClient(
       should_stage_host_to_device_transfers_(
           should_stage_host_to_device_transfers),
       confidential_computing_enabled_(confidential_computing_enabled),
+      staging_chunk_size_(confidential_computing_enabled
+                              ? GetStagingChunkSizeFromEnv()
+                              : kDefaultStagingChunkSizeMb << 20),
       executor_(executor),
       gpu_run_options_(std::move(gpu_run_options)),
       compile_thread_pool_(
diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h
index 071209962ea..be720193ab8 100644
--- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h
+++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h
@@ -218,11 +218,36 @@ class PjRtStreamExecutorRawClient : public PjRtRawClient {
     if (confidential_computing_enabled_) {
       return true;
     }
-    // Allocating multi-gigabyte pinned buffers can be very slow. In that case,
-    // using a staging buffer is probably worse than not using one.
-    // TODO(phawkins): add chunking for transfers.
+    // Allocating a multi-gigabyte pinned staging buffer can be very slow, in
+    // which case staging is probably worse than not staging. This cap does not
+    // apply to chunked staging, which only uses staging_chunk_size() buffers.
+    // TODO(phawkins): add chunking for transfers when chunked staging is off.
     return should_stage_host_to_device_transfers_ &&
-           size < (int64_t{1} << 30) && !IsDmaMapped(data, size);
+           (size < (int64_t{1} << 30) || ShouldUseChunkedStaging()) &&
+           !IsDmaMapped(data, size);
+  }
+
+  // Returns true if staged transfers should be split into at most
+  // `staging_chunk_size()`-byte chunks that are pipelined through a bounded
+  // number of staging buffers, with host copies done on the transfer's work
+  // thread instead of in stream host callbacks. Applies to staged transfers of
+  // every size (a transfer no larger than a chunk is a single chunk). Only
+  // enabled in Confidential Computing mode (where every transfer is staged),
+  // or when forced for testing.
+  bool ShouldUseChunkedStaging() const {
+    return confidential_computing_enabled_ ||
+           force_staging_chunking_for_testing_;
+  }
+
+  int64_t staging_chunk_size() const { return staging_chunk_size_; }
+
+  // Test-only: enables chunked staging (when `force` is true) even when
+  // Confidential Computing is disabled, using `chunk_size`-byte chunks. Must be
+  // called before any transfer is issued on this client.
+  void SetStagingChunkingForTesting(bool force, int64_t chunk_size) {
+    CHECK_GT(chunk_size, 0);
+    force_staging_chunking_for_testing_ = force;
+    staging_chunk_size_ = chunk_size;
   }
 
   tsl::AsyncValueRef<PjRtExecutable> ToAsyncExecutable(
@@ -374,6 +399,9 @@ class PjRtStreamExecutorRawClient : public PjRtRawClient {
   // transfer via pinned memory.
   bool should_stage_host_to_device_transfers_;
   bool confidential_computing_enabled_ = false;
+  // Chunk size for chunked staging; see ShouldUseChunkedStaging().
+  int64_t staging_chunk_size_;
+  bool force_staging_chunking_for_testing_ = false;
 
   se::StreamExecutor* absl_nonnull executor_;
   std::unique_ptr<gpu::GpuExecutableRunOptions> gpu_run_options_;
diff --git a/third_party/xla/xla/pjrt/se/se_raw_buffer.cc b/third_party/xla/xla/pjrt/se/se_raw_buffer.cc
index 0ac1b71c73f..23401ed95b3 100644
--- a/third_party/xla/xla/pjrt/se/se_raw_buffer.cc
+++ b/third_party/xla/xla/pjrt/se/se_raw_buffer.cc
@@ -15,16 +15,19 @@ limitations under the License.
 
 #include "xla/pjrt/se/se_raw_buffer.h"
 
+#include <algorithm>
 #include <cstdint>
 #include <cstring>
 #include <memory>
 #include <utility>
+#include <vector>
 
 #include "absl/functional/any_invocable.h"
 #include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
+#include "absl/types/span.h"
 #include "xla/pjrt/async_work_runner.h"
 #include "xla/pjrt/common_pjrt_client.h"
 #include "xla/pjrt/device_event.h"
@@ -38,7 +41,9 @@ limitations under the License.
 #include "xla/pjrt/se/pjrt_stream_executor_client.h"
 #include "xla/pjrt/se/tracked_device_buffer.h"
 #include "xla/stream_executor/device_address.h"
+#include "xla/stream_executor/event.h"
 #include "xla/stream_executor/stream.h"
+#include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/concurrency/async_value.h"
 #include "xla/tsl/concurrency/async_value_ref.h"
 #include "xla/tsl/concurrency/ref_count.h"
@@ -69,6 +74,138 @@ GetOrCreateAllocationReadyEvent(const xla::RawSEDeviceMemory& device_buffer,
 
 namespace xla {
 
+namespace {
+
+// Maximum number of chunks of a single chunked staged transfer in flight.
+constexpr int64_t kMaxStagingChunksInFlight = 4;
+
+// A pinned host staging buffer for one chunk, and an event recorded after the
+// most recent DMA that used it.
+struct StagingSlot {
+  std::shared_ptr<void> buffer;
+  std::unique_ptr<se::Event> dma_done;
+};
+
+// Allocates the staging slots for a chunked transfer of `transfer_size` bytes.
+// Each slot holds one chunk, so transfers smaller than `chunk_size` only pin
+// `transfer_size` bytes.
+absl::StatusOr<std::vector<StagingSlot>> AllocateStagingSlots(
+    PjRtStreamExecutorRawClient* client, LocalDeviceState* local_device,
+    se::Stream* stream, int64_t transfer_size, int64_t chunk_size) {
+  HostMemoryAllocator::AllocateOptions alloc_opts;
+  alloc_opts.numa_node = stream->parent()->numa_node();
+  alloc_opts.local_device_id = local_device->local_device_id();
+  const int64_t slot_size = std::min(chunk_size, transfer_size);
+  std::vector<StagingSlot> slots(std::min(
+      kMaxStagingChunksInFlight, CeilOfRatio(transfer_size, chunk_size)));
+  for (StagingSlot& slot : slots) {
+    slot.buffer =
+        client->GetHostMemoryAllocator()->Allocate(slot_size, alloc_opts);
+    if (slot.buffer == nullptr) {
+      return ResourceExhausted(
+          "Failed to allocate a %d-byte pinned host staging chunk for a "
+          "chunked host<->device transfer. The pinned host pool may be "
+          "exhausted or fragmented (see XLA_PJRT_GPU_HOST_MEMORY_LIMIT_GB and "
+          "XLA_PJRT_GPU_CC_STAGING_CHUNK_MB).",
+          slot_size);
+    }
+    ABSL_ASSIGN_OR_RETURN(slot.dma_done, stream->parent()->CreateEvent());
+  }
+  return slots;
+}
+
+// Threading: the chunked helpers below run inside `run_transfer` on the
+// client's `async_work_runner()` and do the host<->staging copies on that
+// thread instead of in stream host callbacks. They block that thread on
+// `se::Event`s while chunk DMAs complete (for D2H, on every transfer). This is
+// safe because (1) chunking is only enabled for GPU Confidential Computing (or
+// forced in tests), and the SE GPU client always uses the default
+// `MakeUnboundedAsyncWorkRunner`, which starts a new thread when none are
+// idle, so a blocked transfer cannot starve other work; and (2) every awaited
+// event is recorded after already-enqueued stream work, so it completes on GPU
+// progress alone and never waits on work from this runner. Clients that supply
+// a bounded runner would hold one thread per in-flight staged transfer.
+
+// Copies `transfer_size` bytes from host `src` to device `dst` in
+// `chunk_size`-byte chunks staged through `slots`. Host copies into the
+// staging slots run on the calling thread and overlap with the DMAs of
+// previously enqueued chunks. Blocks the calling thread until each slot's
+// previous DMA completed before reusing it.
+absl::Status ChunkedStagedHostToDevice(se::Stream* stream, const void* src,
+                                       se::DeviceAddressBase dst,
+                                       int64_t transfer_size,
+                                       int64_t chunk_size,
+                                       absl::Span<StagingSlot> slots) {
+  const int64_t num_chunks = CeilOfRatio(transfer_size, chunk_size);
+  const int64_t num_slots = slots.size();
+  for (int64_t i = 0; i < num_chunks; ++i) {
+    const int64_t offset = i * chunk_size;
+    const int64_t size = std::min(chunk_size, transfer_size - offset);
+    StagingSlot& slot = slots[i % num_slots];
+    if (i >= num_slots) {
+      ABSL_RETURN_IF_ERROR(slot.dma_done->Synchronize());
+    }
+    {
+      tsl::profiler::TraceMe trace("H2D Copy To Staging Chunk");
+      std::memcpy(slot.buffer.get(), static_cast<const char*>(src) + offset,
+                  size);
+    }
+    se::DeviceAddressBase dst_chunk = dst.GetByteSlice(offset, size);
+    ABSL_RETURN_IF_ERROR(stream->Memcpy(&dst_chunk, slot.buffer.get(), size));
+    ABSL_RETURN_IF_ERROR(stream->RecordEvent(slot.dma_done.get()));
+  }
+  return absl::OkStatus();
+}
+
+// Copies `transfer_size` bytes from device `src` to host `dst` in
+// `chunk_size`-byte chunks staged through `slots`. Keeps up to `slots.size()`
+// chunk DMAs in flight while the calling thread copies completed chunks out of
+// the staging slots. Returns once `dst` is fully written.
+absl::Status ChunkedStagedDeviceToHost(se::Stream* stream,
+                                       se::DeviceAddressBase src, void* dst,
+                                       int64_t transfer_size,
+                                       int64_t chunk_size,
+                                       absl::Span<StagingSlot> slots) {
+  const int64_t num_chunks = CeilOfRatio(transfer_size, chunk_size);
+  const int64_t num_slots = slots.size();
+  auto enqueue_dma = [&](int64_t i) -> absl::Status {
+    const int64_t offset = i * chunk_size;
+    const int64_t size = std::min(chunk_size, transfer_size - offset);
+    StagingSlot& slot = slots[i % num_slots];
+    ABSL_RETURN_IF_ERROR(stream->Memcpy(slot.buffer.get(),
+                                   src.GetByteSlice(offset, size), size));
+    return stream->RecordEvent(slot.dma_done.get());
+  };
+  for (int64_t i = 0; i < num_slots; ++i) {
+    ABSL_RETURN_IF_ERROR(enqueue_dma(i));
+  }
+  for (int64_t i = 0; i < num_chunks; ++i) {
+    const int64_t offset = i * chunk_size;
+    const int64_t size = std::min(chunk_size, transfer_size - offset);
+    StagingSlot& slot = slots[i % num_slots];
+    ABSL_RETURN_IF_ERROR(slot.dma_done->Synchronize());
+    {
+      tsl::profiler::TraceMe trace("D2H Copy From Staging Chunk");
+      std::memcpy(static_cast<char*>(dst) + offset, slot.buffer.get(), size);
+    }
+    if (i + num_slots < num_chunks) {
+      ABSL_RETURN_IF_ERROR(enqueue_dma(i + num_slots));
+    }
+  }
+  return absl::OkStatus();
+}
+
+// On failure, chunk DMAs may still be accessing the staging slots, so drain
+// `stream` before the slots can be released.
+absl::Status DrainStreamOnError(se::Stream* stream, absl::Status status) {
+  if (!status.ok()) {
+    stream->BlockHostUntilDone().IgnoreError();
+  }
+  return status;
+}
+
+}  // namespace
+
 PjRtStreamExecutorDeviceEventPromise::PjRtStreamExecutorDeviceEventPromise(
     PjRtStreamExecutorRawClient* client, LocalDeviceState* local_device,
     AsyncWorkRunner* async_work_runner)
@@ -133,6 +270,9 @@ PjRtStreamExecutorRawBuffer::CopyRawHostToDeviceAndReturnEvent(
       sub_buffer = sub_buffer.GetByteSlice(offset, transfer_size);
     }
     std::shared_ptr<void> staging_buffer;
+    // Only set for chunked staging. Shared with the cleanup callback so that
+    // the slots outlive the drain below if AllocateAndRecordEvent fails.
+    std::shared_ptr<std::vector<StagingSlot>> staging_slots;
     auto status = [&]() -> absl::Status {
       ABSL_RETURN_IF_ERROR(client->WaitForAllocation(stream, *buf));
       if (transfer_size > 0) {
@@ -142,6 +282,17 @@ PjRtStreamExecutorRawBuffer::CopyRawHostToDeviceAndReturnEvent(
                 "host_memory_allocator should be initialized for "
                 "staging buffer transfer.");
           }
+          if (client->ShouldUseChunkedStaging()) {
+            const int64_t chunk_size = client->staging_chunk_size();
+            ABSL_ASSIGN_OR_RETURN(std::vector<StagingSlot> slots,
+                             AllocateStagingSlots(client, local_device, stream,
+                                                  transfer_size, chunk_size));
+            staging_slots =
+                std::make_shared<std::vector<StagingSlot>>(std::move(slots));
+            return ChunkedStagedHostToDevice(stream, src, sub_buffer,
+                                             transfer_size, chunk_size,
+                                             absl::MakeSpan(*staging_slots));
+          }
           HostMemoryAllocator::AllocateOptions alloc_opts;
           alloc_opts.numa_node = stream->parent()->numa_node();
           alloc_opts.local_device_id = local_device->local_device_id();
@@ -175,11 +326,18 @@ PjRtStreamExecutorRawBuffer::CopyRawHostToDeviceAndReturnEvent(
       status = client->AllocateAndRecordEvent(
           device_event, local_device, stream,
           "PjRtStreamExecutorRawBuffer::CopyRawHostToDevice",
-          [staging_buffer = std::move(staging_buffer)]() mutable {
+          [staging_buffer = std::move(staging_buffer),
+           staging_slots]() mutable {
             staging_buffer.reset();
+            staging_slots.reset();
           });
     }
     if (!status.ok()) {
+      // Chunk DMAs may still be reading the staging slots, so drain before
+      // releasing them.
+      if (staging_slots != nullptr) {
+        stream->BlockHostUntilDone().IgnoreError();
+      }
       client->SetEventAsError(device_event, status);
     }
   };
@@ -226,6 +384,18 @@ PjRtStreamExecutorRawBuffer::CopyRawDeviceToHostAndReturnEvent(
                 "host_memory_allocator should be initialized for "
                 "staging buffer transfer.");
           }
+          if (client->ShouldUseChunkedStaging()) {
+            // All chunk DMAs and host copies complete before this returns, so
+            // the staging slots can be released when it goes out of scope.
+            const int64_t chunk_size = client->staging_chunk_size();
+            ABSL_ASSIGN_OR_RETURN(std::vector<StagingSlot> staging_slots,
+                             AllocateStagingSlots(client, local_device, stream,
+                                                  transfer_size, chunk_size));
+            return DrainStreamOnError(
+                stream, ChunkedStagedDeviceToHost(
+                            stream, sub_buffer, dst, transfer_size, chunk_size,
+                            absl::MakeSpan(staging_slots)));
+          }
           HostMemoryAllocator::AllocateOptions alloc_opts;
           alloc_opts.numa_node = stream->parent()->numa_node();
           alloc_opts.local_device_id = local_device->local_device_id();

From 43b17e3b4daa7be9ae9d1e60367274acc0d5e700 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Fri, 9 Oct 2026 05:43:20 -0700
Subject: [PATCH 0357/1000] [XLA:GPU]: Support variadic all-gather tile
 propagation.

PropagateTilesToInputs will be used in tiled hlo (up next).

PiperOrigin-RevId: 996480600
---
 .../xla/xla/codegen/tiling/experimental/BUILD |   1 +
 .../tiling/experimental/tile_propagation.cc   |  76 +++++-
 .../tiling/experimental/tile_propagation.h    |   5 +
 .../experimental/tile_propagation_test.cc     | 244 ++++++++++++++++++
 4 files changed, 314 insertions(+), 12 deletions(-)

diff --git a/third_party/xla/xla/codegen/tiling/experimental/BUILD b/third_party/xla/xla/codegen/tiling/experimental/BUILD
index b20684a8534..ce30f449e7f 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/BUILD
+++ b/third_party/xla/xla/codegen/tiling/experimental/BUILD
@@ -137,6 +137,7 @@ xla_cc_test(
         ":tile_propagation",
         "//xla:shape_util",
         "//xla:util",
+        "//xla/hlo/analysis:interval",
         "//xla/hlo/analysis:symbolic_map",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
index 3cdfe73217d..1eddb4f135b 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
@@ -1488,7 +1488,8 @@ absl::StatusOr<Tile> PropagateTileThroughReshape(const Tile& tile,
         mlir_context, minimal_reshape, src, reshape_shape, tile,
         target_dim_tiles));
   }
-  Tile reshape_tile(tiling_space, std::move(target_dim_tiles));
+  Tile reshape_tile(tiling_space, std::move(target_dim_tiles),
+                    llvm::to_vector(tile.replica_ids()), tile.constraints());
   return PropagateTileToOutputForBroadcastOpImpl(dst, non_trivial_dim_positions,
                                                  reshape_tile);
 }
@@ -1635,20 +1636,39 @@ std::string ToString(const Tiles& tiles) {
   return ss.str();
 }
 
-Tiles PropagateTileToInputForAllGatherOp(const TilingSpace& tiling_space,
-                                         const HloInstruction& hlo,
-                                         const Tile& output_tile) {
+absl::StatusOr<Tiles> PropagateTileToInputForAllGatherOp(
+    const TilingSpace& tiling_space, const HloInstruction& hlo,
+    const Tile& output_tile, int64_t output_index) {
   const auto& all_gather = *Cast<HloAllGatherInstruction>(&hlo);
   int64_t gather_dim = all_gather.all_gather_dimension();
-  // Crash OK. Must be checked while forming the fusion.
-  CHECK_EQ(hlo.operand_count(), 1)
-      << "Multi-operand AllGather is not yet supported.";
-  const Shape& input_shape = hlo.operand(0)->shape();
+  if (output_index < 0 || output_index >= hlo.operand_count()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Output index ", output_index, " is out of range for ",
+                     hlo.ToString()));
+  }
+  // Each output of a (possibly variadic) all-gather depends only on the operand
+  // with the same index.
+  const Shape& input_shape = hlo.operand(output_index)->shape();
+  const Shape& output_shape = GetFirstShape(&hlo, output_index);
   int64_t local_size = input_shape.dimensions(gather_dim);
-  int64_t num_replicas = hlo.shape().dimensions(gather_dim) / local_size;
+  int64_t num_replicas = output_shape.dimensions(gather_dim) / local_size;
 
   const DimTile& output_dim_tile = output_tile.dim_tiles()[gather_dim];
 
+  if (!tiling_space.IsSymbolic()) {
+    DimTile simplified_dim_tile = output_dim_tile;
+    simplified_dim_tile.Simplify(tiling_space);
+    if (simplified_dim_tile.size.GetType() != SymbolicExprType::kConstant ||
+        simplified_dim_tile.size.GetValue() <= 0 ||
+        local_size % simplified_dim_tile.size.GetValue() != 0) {
+      return absl::UnimplementedError(absl::StrCat(
+          "Output ", output_index, " of ", hlo.name(), " is tiled with size ",
+          simplified_dim_tile.size.ToString(), " on gather dimension ",
+          gather_dim, ", expected tile size to divide per-rank size ",
+          local_size, " (one peer per tile)"));
+    }
+  }
+
   SymbolicExpr replica_id = output_dim_tile.offset / local_size;
   SymbolicExpr input_offset = output_dim_tile.offset % local_size;
 
@@ -1677,9 +1697,9 @@ Tiles PropagateTileToInputForAllGatherOp(const TilingSpace& tiling_space,
   });
 
   Tile input_tile(output_tile.tiling_space(), std::move(input_dim_tiles),
-                  std::move(replica_id_dim_tiles));
+                  std::move(replica_id_dim_tiles), output_tile.constraints());
 
-  return {input_tile};
+  return Tiles{input_tile};
 }
 
 absl::StatusOr<Tiles> PropagateTileToInput(TilingSpace& tiling_space,
@@ -1697,7 +1717,8 @@ absl::StatusOr<Tiles> PropagateTileToInput(TilingSpace& tiling_space,
     return {PropagateTileToInputForCwiseOp(hlo, output_tile)};
   }
   if (hlo.opcode() == HloOpcode::kAllGather) {
-    return PropagateTileToInputForAllGatherOp(tiling_space, hlo, output_tile);
+    return PropagateTileToInputForAllGatherOp(tiling_space, hlo, output_tile,
+                                              output_index);
   }
   if (hlo.opcode() == HloOpcode::kReduceScatter) {
     return PropagateTileToInputForReduceScatterOp(
@@ -1755,6 +1776,37 @@ absl::StatusOr<Tiles> PropagateTileToInput(TilingSpace& tiling_space,
                    HloOpcodeString(hlo.opcode())));
 }
 
+absl::StatusOr<Tiles> PropagateTilesToInputs(
+    TilingSpace& tiling_space, const HloInstruction& hlo,
+    absl::Span<const Tile> output_tiles) {
+  if (IsIndexWiseVariadic(hlo)) {
+    if (output_tiles.size() != hlo.operand_count()) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Expected one tile per output of ", hlo.ToString(),
+                       ", got ", output_tiles.size()));
+    }
+    Tiles operand_tiles;
+    operand_tiles.reserve(hlo.operand_count());
+    for (int64_t k = 0; k < hlo.operand_count(); ++k) {
+      ABSL_ASSIGN_OR_RETURN(Tiles tiles, PropagateTileToInput(tiling_space, hlo,
+                                                         output_tiles[k], k));
+      if (tiles.size() != 1) {
+        return absl::InternalError(
+            absl::StrCat("Expected a single input tile for output ", k, " of ",
+                         hlo.ToString(), ", got ", tiles.size()));
+      }
+      operand_tiles.push_back(std::move(tiles[0]));
+    }
+    return operand_tiles;
+  }
+  if (output_tiles.size() != 1) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Expected a single output tile for ", hlo.ToString(),
+                     ", got ", output_tiles.size()));
+  }
+  return PropagateTileToInput(tiling_space, hlo, output_tiles[0], 0);
+}
+
 absl::StatusOr<Tiles> PropagateTileToOutput(const TilingSpace& tiling_space,
                                             const HloInstruction& hlo,
                                             const Tile& input_tile,
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.h b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.h
index cddffa8ee3f..36d41f610dc 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.h
+++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.h
@@ -20,6 +20,7 @@ limitations under the License.
 #include <string>
 
 #include "absl/status/statusor.h"
+#include "absl/types/span.h"
 #include "llvm/ADT/SmallVector.h"
 #include "xla/codegen/tiling/experimental/tile.h"
 #include "xla/codegen/tiling/experimental/tiling_space.h"
@@ -36,6 +37,10 @@ absl::StatusOr<Tiles> PropagateTileToInput(TilingSpace& tiling_space,
                                            const Tile& output_tile,
                                            int64_t output_index);
 
+absl::StatusOr<Tiles> PropagateTilesToInputs(
+    TilingSpace& tiling_space, const HloInstruction& hlo,
+    absl::Span<const Tile> output_tiles);
+
 absl::StatusOr<Tiles> PropagateTileToOutput(const TilingSpace& tiling_space,
                                             const HloInstruction& hlo,
                                             const Tile& input_tile,
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
index c9a50a11cde..1cc8b0550ae 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
@@ -35,8 +35,10 @@ limitations under the License.
 #include "xla/codegen/tiling/experimental/test_utils.h"
 #include "xla/codegen/tiling/experimental/tile.h"
 #include "xla/codegen/tiling/experimental/tiling_space.h"
+#include "xla/hlo/analysis/interval.h"
 #include "xla/hlo/analysis/symbolic_expr.h"
 #include "xla/hlo/analysis/symbolic_map.h"
+#include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/hlo/testlib/verified_hlo_module.h"
@@ -50,6 +52,8 @@ namespace {
 using ::absl_testing::StatusIs;
 using ::llvm::SmallVector;
 using ::mlir::MLIRContext;
+using ::testing::ElementsAre;
+using ::testing::Pair;
 
 class TilePropagationTest : public HloHardwareIndependentTestBase {
  public:
@@ -197,6 +201,244 @@ TEST_F(TilePropagationTest, CanPropagateToInputsOfAllGatherOp) {
   )"));
 }
 
+TEST_F(TilePropagationTest, CanPropagateToInputsOfVariadicAllGatherOp) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    f {
+      p0 = f32[1,2,8,16] parameter(0)
+      p1 = s32[1,2] parameter(1)
+      ag = (f32[4,2,8,16], s32[4,2]) all-gather(p0, p1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      gte0 = f32[4,2,8,16] get-tuple-element(ag), index=0
+      gte1 = s32[4,2] get-tuple-element(ag), index=1
+      ROOT t = (f32[4,2,8,16], s32[4,2]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      p0 = f32[1,2,8,16] parameter(0)
+      p1 = s32[1,2] parameter(1)
+      ROOT fusion = (f32[4,2,8,16], s32[4,2]) fusion(p0, p1), kind=kCustom,
+        calls=f
+    }
+  )");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TilingSpace> tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                          &mlir_context_));
+  const HloInstruction* ag =
+      root->fused_instructions_computation()->GetInstructionWithName("ag");
+  ASSERT_EQ(tiling_space->tiled_roots().size(), 2);
+
+  // Output `k` maps directly to operand `k`.
+  // Since the gather dimension size is 1, the local offset is always 0 and the
+  // replica ID directly equals the output offset. Each output gets distinct
+  // tile IDs (e.g., tid_0 and tid_4) that the scheduler will later assign to
+  // the same peer.
+  ASSERT_OK_AND_ASSIGN(
+      Tiles operand_0_tiles,
+      PropagateTileToInput(*tiling_space, *ag, tiling_space->tiled_roots()[0],
+                           /*output_index=*/0));
+  EXPECT_THAT(operand_0_tiles, MatchToString(R"(
+    0) (tid_0, tid_1, tid_2, tid_3, tid_4, tid_5)
+      -> offsets [0, tid_1 * ts_1, tid_2 * ts_2, tid_3 * ts_3]
+         sizes [ts_0, ts_1, ts_2, ts_3]
+         strides [1, 1, 1, 1]
+         upper bounds [1, 2, 8, 16]
+         replica ids {
+           offsets [tid_0 * ts_0]
+           sizes [1]
+           strides [1]
+           upper bounds [4]
+         }
+  )"));
+
+  ASSERT_OK_AND_ASSIGN(
+      Tiles operand_1_tiles,
+      PropagateTileToInput(*tiling_space, *ag, tiling_space->tiled_roots()[1],
+                           /*output_index=*/1));
+  EXPECT_THAT(operand_1_tiles, MatchToString(R"(
+    0) (tid_0, tid_1, tid_2, tid_3, tid_4, tid_5)
+      -> offsets [0, tid_5 * ts_5]
+         sizes [ts_4, ts_5]
+         strides [1, 1]
+         upper bounds [1, 2]
+         replica ids {
+           offsets [tid_4 * ts_4]
+           sizes [1]
+           strides [1]
+           upper bounds [4]
+         }
+  )"));
+
+  ASSERT_OK_AND_ASSIGN(
+      Tiles all_operand_tiles,
+      PropagateTilesToInputs(*tiling_space, *ag, tiling_space->tiled_roots()));
+  EXPECT_THAT(all_operand_tiles,
+              ElementsAre(operand_0_tiles[0], operand_1_tiles[0]));
+}
+
+TEST_F(TilePropagationTest, CanPropagateFromVariadicAllGatherToDotProducers) {
+  HloInstruction* fusion = ParseAndGetRoot(R"(
+    HloModule m
+    f {
+      lhs0 = f32[2,8] parameter(0)
+      rhs0 = f32[8,16] parameter(1)
+      lhs1 = f32[4,8] parameter(2)
+      rhs1 = f32[8,32] parameter(3)
+      dot0 = f32[2,16] dot(lhs0, rhs0),
+        lhs_contracting_dims={1}, rhs_contracting_dims={0}
+      dot1 = f32[4,32] dot(lhs1, rhs1),
+        lhs_contracting_dims={1}, rhs_contracting_dims={0}
+      ag = (f32[8,16], f32[16,32]) all-gather(dot0, dot1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      gte0 = f32[8,16] get-tuple-element(ag), index=0
+      gte1 = f32[16,32] get-tuple-element(ag), index=1
+      ROOT t = (f32[8,16], f32[16,32]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      lhs0 = f32[2,8] parameter(0)
+      rhs0 = f32[8,16] parameter(1)
+      lhs1 = f32[4,8] parameter(2)
+      rhs1 = f32[8,32] parameter(3)
+      ROOT fusion = (f32[8,16], f32[16,32]) fusion(lhs0, rhs0, lhs1, rhs1),
+        kind=kCustom, calls=f
+    }
+  )");
+  const HloInstruction* ag =
+      fusion->fused_instructions_computation()->GetInstructionWithName("ag");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TilingSpace> tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(fusion),
+                          &mlir_context_));
+  ASSERT_OK_AND_ASSIGN(
+      Tiles ag_op0,
+      PropagateTileToInput(*tiling_space, *ag, tiling_space->tiled_roots()[0],
+                           /*output_index=*/0));
+  ASSERT_OK_AND_ASSIGN(
+      Tiles dot0_operands,
+      PropagateTileToInput(*tiling_space, *ag->operand(0), ag_op0[0], 0));
+  EXPECT_THAT(dot0_operands, MatchToString(R"(
+    0) (tid_0, tid_1, tid_2, tid_3, tid_4, tid_5)
+      -> offsets [(tid_0 * ts_0) mod 2, tid_5 * ts_5]
+         sizes [ts_0, ts_5]
+         strides [1, 1]
+         upper bounds [2, 8]
+         replica ids {
+           offsets [(tid_0 * ts_0) / 2]
+           sizes [1]
+           strides [1]
+           upper bounds [4]
+         }
+    1) (tid_0, tid_1, tid_2, tid_3, tid_4, tid_5)
+      -> offsets [tid_5 * ts_5, tid_1 * ts_1]
+         sizes [ts_5, ts_1]
+         strides [1, 1]
+         upper bounds [8, 16]
+         replica ids {
+           offsets [(tid_0 * ts_0) / 2]
+           sizes [1]
+           strides [1]
+           upper bounds [4]
+         }
+  )"));
+
+  ASSERT_OK_AND_ASSIGN(
+      Tiles ag_op1,
+      PropagateTileToInput(*tiling_space, *ag, tiling_space->tiled_roots()[1],
+                           /*output_index=*/1));
+  ASSERT_OK_AND_ASSIGN(
+      Tiles dot1_operands,
+      PropagateTileToInput(*tiling_space, *ag->operand(1), ag_op1[0], 0));
+  EXPECT_THAT(dot1_operands, MatchToString(R"(
+    0) (tid_0, tid_1, tid_2, tid_3, tid_4, tid_5)
+      -> offsets [(tid_2 * ts_2) mod 4, tid_4 * ts_4]
+         sizes [ts_2, ts_4]
+         strides [1, 1]
+         upper bounds [4, 8]
+         replica ids {
+           offsets [(tid_2 * ts_2) / 4]
+           sizes [1]
+           strides [1]
+           upper bounds [4]
+         }
+    1) (tid_0, tid_1, tid_2, tid_3, tid_4, tid_5)
+      -> offsets [tid_4 * ts_4, tid_3 * ts_3]
+         sizes [ts_4, ts_3]
+         strides [1, 1]
+         upper bounds [8, 32]
+         replica ids {
+           offsets [(tid_2 * ts_2) / 4]
+           sizes [1]
+           strides [1]
+           upper bounds [4]
+         }
+  )"));
+}
+
+TEST_F(TilePropagationTest, VariadicAllGatherRejectsOutOfRangeOutputIndex) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    f {
+      p0 = f32[1,8] parameter(0)
+      p1 = s32[1,8] parameter(1)
+      ag = (f32[4,8], s32[4,8]) all-gather(p0, p1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      gte0 = f32[4,8] get-tuple-element(ag), index=0
+      gte1 = s32[4,8] get-tuple-element(ag), index=1
+      ROOT t = (f32[4,8], s32[4,8]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      p0 = f32[1,8] parameter(0)
+      p1 = s32[1,8] parameter(1)
+      ROOT fusion = (f32[4,8], s32[4,8]) fusion(p0, p1), kind=kCustom, calls=f
+    }
+  )");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TilingSpace> tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                          &mlir_context_));
+  const HloInstruction* ag =
+      root->fused_instructions_computation()->GetInstructionWithName("ag");
+  EXPECT_THAT(
+      PropagateTileToInput(*tiling_space, *ag, tiling_space->tiled_roots()[0],
+                           /*output_index=*/2),
+      StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+// Constraints on the tiling space variables hold for every tile derived from
+// the output tile and must survive the propagation (like `CloneWithNewDims`).
+TEST_F(TilePropagationTest, AllGatherKeepsTileConstraints) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    f {
+      p0 = f32[1,8] parameter(0)
+      ROOT ag = f32[4,8] all-gather(p0), replica_groups={{0,1,2,3}},
+        dimensions={0}
+    }
+
+    ENTRY e {
+      p0 = f32[1,8] parameter(0)
+      ROOT fusion = f32[4,8] fusion(p0), kind=kCustom, calls=f
+    }
+  )");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TilingSpace> tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                          &mlir_context_));
+  Tile output_tile = tiling_space->tiled_roots()[0];
+  const SymbolicExpr tid_1 = CreateSymbolicVariable(1, &mlir_context_);
+  output_tile.AddConstraint(tid_1, Interval{0, 1});
+  ASSERT_OK_AND_ASSIGN(
+      Tiles operand_tiles,
+      PropagateTileToInput(*tiling_space, *root->fused_expression_root(),
+                           output_tile, /*output_index=*/0));
+  ASSERT_EQ(operand_tiles.size(), 1);
+  EXPECT_THAT(operand_tiles[0].constraints(),
+              ElementsAre(Pair(tid_1, Interval{0, 1})));
+}
+
 TEST_F(TilePropagationTest, CanPropagateToInputsOfReduceScatterOp) {
   HloInstruction* root = ParseAndGetRoot(R"(
     HloModule m
@@ -1667,12 +1909,14 @@ TEST_F(TilePropagationTest, FailsToPropagateToConcatenateThroughBitcast) {
   ASSERT_OK_AND_ASSIGN(auto output_tiles,
                        PropagateTileToInput(*tiling_space, *(root->operand(1)),
                                             tiled_operands[1], 0));
+  // The concatenate constraint on the tile ids is kept through the bitcast.
   EXPECT_THAT(output_tiles, MatchToString(R"(
     0) (tid_0, tid_1)
          -> offsets [tid_0, tid_1 / 256 - 1, (tid_1 mod 256) * 2]
             sizes [1, 1, 2]
             strides [1, 1, 1]
             upper bounds [10, tid_1 / 256, (tid_1 mod 256) * 2 + 2]
+            constraints {tid_1 * 2 - 512 in [0, 1023]}
   )"));
 }
 

From b305c24ad52cb249d4789a3377282257915d225c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 9 Oct 2026 05:48:29 -0700
Subject: [PATCH 0358/1000] Skip memcpy when copying an empty vector in
 CopyVectorToTfLiteIntArray.

std::vector::data() may return nullptr when the vector is empty, and passing a null source pointer to memcpy is undefined behavior even when the copy size is zero. Delegating a node subset with no output tensors (such as a resource initialization subgraph containing only VAR_HANDLE and ASSIGN_VARIABLE) calls CopyVectorToTfLiteIntArray with an empty output_tensors vector, triggering a nonnull-attribute error under UndefinedBehaviorSanitizer. Return early when vec is empty and add a unit test in subgraph_test.cc.
PiperOrigin-RevId: 996482304
---
 tensorflow/lite/core/subgraph.cc      |  5 ++
 tensorflow/lite/core/subgraph_test.cc | 86 +++++++++++++++++++++++++++
 2 files changed, 91 insertions(+)

diff --git a/tensorflow/lite/core/subgraph.cc b/tensorflow/lite/core/subgraph.cc
index 1b1fbf5d5b7..70f2569a07d 100644
--- a/tensorflow/lite/core/subgraph.cc
+++ b/tensorflow/lite/core/subgraph.cc
@@ -348,6 +348,11 @@ namespace {
 void CopyVectorToTfLiteIntArray(const std::vector<int>& vec,
                                 TfLiteIntArray* arr) {
   arr->size = vec.size();
+  // An empty vector may return a null data(); passing null to memcpy is
+  // undefined behavior even when the size is zero.
+  if (arr->size == 0) {
+    return;
+  }
   memcpy(arr->data, vec.data(), sizeof(int) * arr->size);
 }
 
diff --git a/tensorflow/lite/core/subgraph_test.cc b/tensorflow/lite/core/subgraph_test.cc
index 9c169c6ea25..c890e9ca02e 100644
--- a/tensorflow/lite/core/subgraph_test.cc
+++ b/tensorflow/lite/core/subgraph_test.cc
@@ -332,5 +332,91 @@ TEST_F(SubgraphResizeTensorTest,
   std::fill_n(tensor_.dims->data, tensor_.dims->size, 1);
 }
 
+TEST(ReplaceNodeSubsetsWithDelegateKernels,
+     HandlesNodeSubsetWithNoOutputTensors) {
+  static constexpr int kNumTensors = 1;
+  static constexpr int kInputTensorIndex = 0;
+  static constexpr int kExpectedPartitions = 1;
+  static constexpr int kExpectedReplacedNodes = 1;
+  static constexpr int kExpectedInputTensors = 1;
+  static constexpr int kExpectedOutputTensors = 0;
+
+  // Build a subgraph with a single node that has one input tensor and zero
+  // output tensors (matching the shape of ASSIGN_VARIABLE).
+  Interpreter interpreter;
+  Subgraph& subgraph = interpreter.primary_subgraph();
+  ASSERT_EQ(subgraph.AddTensors(/*tensors_to_add=*/kNumTensors), kTfLiteOk);
+  ASSERT_EQ(subgraph.SetInputs({kInputTensorIndex}), kTfLiteOk);
+  ASSERT_EQ(subgraph.SetOutputs({}), kTfLiteOk);
+
+  TfLiteRegistration op_registration = {};
+  int node_index = -1;
+  ASSERT_EQ(
+      subgraph.AddNodeWithParameters(
+          /*inputs=*/{kInputTensorIndex}, /*outputs=*/{}, /*intermediates=*/{},
+          /*init_data=*/nullptr, /*init_data_size=*/0,
+          /*builtin_data=*/nullptr, &op_registration, &node_index),
+      kTfLiteOk);
+
+  // Delegate the partition and verify that both preview partitioning and
+  // delegate kernel initialization receive output_tensors->size == 0 without
+  // triggering undefined behavior in CopyVectorToTfLiteIntArray.
+  bool kernel_initialized = false;
+  TfLiteDelegate delegate = TfLiteDelegateCreate();
+  delegate.data_ = &kernel_initialized;
+  delegate.Prepare = [](TfLiteContext* context,
+                        TfLiteDelegate* delegate) -> TfLiteStatus {
+    TfLiteIntArray* execution_plan = nullptr;
+    if (context->GetExecutionPlan(context, &execution_plan) != kTfLiteOk) {
+      return kTfLiteError;
+    }
+
+    TfLiteDelegateParams* preview_params = nullptr;
+    int num_partitions = 0;
+    if (context->PreviewDelegatePartitioning(context, execution_plan,
+                                             &preview_params,
+                                             &num_partitions) != kTfLiteOk) {
+      return kTfLiteError;
+    }
+    EXPECT_EQ(num_partitions, kExpectedPartitions);
+    if (preview_params && num_partitions == kExpectedPartitions) {
+      EXPECT_EQ(preview_params[0].nodes_to_replace->size,
+                kExpectedReplacedNodes);
+      EXPECT_EQ(preview_params[0].input_tensors->size, kExpectedInputTensors);
+      if (preview_params[0].input_tensors->size == kExpectedInputTensors) {
+        EXPECT_EQ(preview_params[0].input_tensors->data[0], kInputTensorIndex);
+      }
+      EXPECT_EQ(preview_params[0].output_tensors->size, kExpectedOutputTensors);
+    }
+
+    TfLiteRegistration delegate_kernel_registration = {};
+    delegate_kernel_registration.init =
+        [](TfLiteContext* context, const char* buffer, size_t length) -> void* {
+      const TfLiteDelegateParams* params =
+          reinterpret_cast<const TfLiteDelegateParams*>(buffer);
+      if (!params) {
+        ADD_FAILURE() << "Expected non-null TfLiteDelegateParams.";
+        return nullptr;
+      }
+      if (params->delegate && params->delegate->data_) {
+        *static_cast<bool*>(params->delegate->data_) = true;
+      }
+      EXPECT_EQ(params->nodes_to_replace->size, kExpectedReplacedNodes);
+      EXPECT_EQ(params->input_tensors->size, kExpectedInputTensors);
+      if (params->input_tensors->size == kExpectedInputTensors) {
+        EXPECT_EQ(params->input_tensors->data[0], kInputTensorIndex);
+      }
+      EXPECT_EQ(params->output_tensors->size, kExpectedOutputTensors);
+      return nullptr;
+    };
+
+    return context->ReplaceNodeSubsetsWithDelegateKernels(
+        context, delegate_kernel_registration, execution_plan, delegate);
+  };
+
+  ASSERT_EQ(interpreter.ModifyGraphWithDelegate(&delegate), kTfLiteOk);
+  EXPECT_TRUE(kernel_initialized);
+}
+
 }  // namespace
 }  // namespace tflite

From 596f9a0d8d935d54818d8199e96eef823134dd44 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 9 Oct 2026 06:17:34 -0700
Subject: [PATCH 0359/1000] Automated Code Change

PiperOrigin-RevId: 996493341
---
 tensorflow/core/common_runtime/gpu/gpu_device.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/common_runtime/gpu/gpu_device.cc b/tensorflow/core/common_runtime/gpu/gpu_device.cc
index 318754e3b80..b676fbaed6d 100644
--- a/tensorflow/core/common_runtime/gpu/gpu_device.cc
+++ b/tensorflow/core/common_runtime/gpu/gpu_device.cc
@@ -196,8 +196,8 @@ class EigenGpuStreamDevice : public ::Eigen::StreamInterface {
     if (ret == nullptr) {
       if (context_) {
         context_->SetStatus(absl::ResourceExhaustedError(
-            strings::StrCat("Ran out of GPU memory when allocating ", num_bytes,
-                            " bytes for ", operation_)));
+            absl::StrCat("Ran out of GPU memory when allocating ", num_bytes,
+                         " bytes for ", operation_)));
       } else {
         LOG(FATAL)
             << "EigenAllocator for GPU ran out of memory when allocating "

From 061bb6055cdcd01aba5a18c854ef37087ab758e4 Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Fri, 9 Oct 2026 07:02:18 -0700
Subject: [PATCH 0360/1000] [XLA] Fix negative constant multiplication in value
 range analysis

Swap `min` and `max` and keep `step` positive when multiplying a bounded `Range` by a negative signed constant in `RecursivelyIdentifyRange`. Return an empty `Range` when multiplying an unbounded `Range` by a negative constant or when the multiplied bounds wrap around.

PiperOrigin-RevId: 996510775
---
 third_party/xla/xla/service/value_range.cc    | 30 ++++++++++--
 .../xla/xla/service/value_range_test.cc       | 47 +++++++++++++++++++
 2 files changed, 74 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/service/value_range.cc b/third_party/xla/xla/service/value_range.cc
index 90b4152ee6c..5c6b2154a0d 100644
--- a/third_party/xla/xla/service/value_range.cc
+++ b/third_party/xla/xla/service/value_range.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include <cstdint>
 #include <optional>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "absl/container/flat_hash_map.h"
@@ -272,18 +273,41 @@ Range RecursivelyIdentifyRange(
       }
       ConstantValue single_value = lhs.IsSingleValue() ? lhs.min() : rhs.min();
       Range operand_range = lhs.IsSingleValue() ? rhs : lhs;
+      const bool is_neg =
+          single_value.IsSigned() && single_value.GetSignedValue() < 0;
+      // `!p.div(single_value).eq(v)` catches all multiplication overflow except
+      // `INT_MIN * -1` in two's complement, which wraps to `p == v` (`< 0`).
+      auto wraps = [&](const ConstantValue& v) {
+        ConstantValue p = v.mul(single_value);
+        return single_value.GetUnsignedValue() != 0 &&
+               ((is_neg && p == v && v.GetSignedValue() < 0) ||
+                !p.div(single_value).eq(v));
+      };
+      // Unbounded ranges represent `[min, +inf)`. Multiplying by a negative
+      // constant `c` produces `(-inf, c * min]`, which `Range` cannot represent
+      // because `Range` requires a finite lower bound `min_`.
+      if (wraps(operand_range.min()) ||
+          (operand_range.IsBounded() ? wraps(*operand_range.max()) : is_neg)) {
+        return Range{};
+      }
       // When multiplying with a constant, min, max, and step are all
       // multiplied by the single value.
       ConstantValue min = operand_range.min().mul(single_value);
       if (operand_range.IsBounded()) {
         ConstantValue max = operand_range.max()->mul(single_value);
+        if (is_neg) {
+          std::swap(min, max);
+        }
         if (!operand_range.IsStepKnown()) {
           return RecordAndReturnRange(Range{min, max, operand_range.IsLinear()},
                                       instr, known_ranges);
         }
-        ConstantValue step = operand_range.step()->mul(single_value);
-        // Zero step makes it difficult to simulate the range; update it to 1.
-        if (step.GetSignedValue() == 0) {
+        ConstantValue step = operand_range.step()->mul(
+            is_neg ? ConstantValue::GetZero(64, true).sub(single_value)
+                   : single_value);
+        // Zero or negative step makes it difficult to simulate the range;
+        // update it to 1.
+        if (step.GetSignedValue() <= 0) {
           step = ConstantValue::GetOne(/*bitwidth=*/64, /*is_signed=*/true);
         }
         return RecordAndReturnRange(
diff --git a/third_party/xla/xla/service/value_range_test.cc b/third_party/xla/xla/service/value_range_test.cc
index 3cf02f61618..15d45b2f3a6 100644
--- a/third_party/xla/xla/service/value_range_test.cc
+++ b/third_party/xla/xla/service/value_range_test.cc
@@ -677,5 +677,52 @@ TEST_F(ValueRangeTest, SelectWithUnknownConditionReturnsEmptyRange) {
   EXPECT_FALSE(range.IsSingleValue());
 }
 
+TEST_F(ValueRangeTest, MultiplyValueWithNegativeConstant) {
+  constexpr absl::string_view hlo_string = R"hlo(
+  HloModule module
+
+  ENTRY entry {
+    c0 = s32[] constant(-1024)
+    p0 = s32[] parameter(0)
+    ROOT %a = s32[] multiply(p0, c0)
+  }
+  )hlo";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(
+                                        hlo_string, HloModuleConfig{}));
+  const HloInstruction* root = module->entry_computation()->root_instruction();
+  const HloInstruction* p0 = root->operand(0);
+  absl::flat_hash_map<const HloInstruction*, Range> fs;
+  fs.insert(std::make_pair(
+      p0, Range{/*min=*/ConstantValue::GetSigned(2, /*bitwidth=*/32),
+                /*max=*/ConstantValue::GetSigned(32, /*bitwidth=*/32),
+                /*step=*/ConstantValue::GetSigned(2, /*bitwidth=*/32),
+                /*is_linear=*/true}));
+  Range range = RecursivelyIdentifyRange(root, fs);
+  EXPECT_FALSE(range.IsEmpty());
+  EXPECT_FALSE(range.IsSingleValue());
+  EXPECT_TRUE(range.IsLinear());
+  EXPECT_EQ(range.min().GetSignedValue(), -32 * 1024);
+  EXPECT_EQ(range.max()->GetSignedValue(), -2 * 1024);
+  EXPECT_EQ(range.step()->GetSignedValue(), 2 * 1024);
+
+  absl::flat_hash_map<const HloInstruction*, Range> unbounded_fs;
+  unbounded_fs.insert(std::make_pair(
+      p0, Range{/*min=*/ConstantValue::GetSigned(2, /*bitwidth=*/32),
+                /*max=*/std::nullopt,
+                /*step=*/ConstantValue::GetSigned(2, /*bitwidth=*/32),
+                /*is_linear=*/true}));
+  Range unbounded_range = RecursivelyIdentifyRange(root, unbounded_fs);
+  EXPECT_TRUE(unbounded_range.IsEmpty());
+
+  absl::flat_hash_map<const HloInstruction*, Range> wrapped_fs;
+  wrapped_fs.insert(std::make_pair(
+      p0, Range{/*min=*/ConstantValue::GetSigned(0, /*bitwidth=*/32),
+                /*max=*/ConstantValue::GetSigned(1 << 22, /*bitwidth=*/32),
+                /*step=*/ConstantValue::GetSigned(1, /*bitwidth=*/32),
+                /*is_linear=*/true}));
+  Range wrapped_range = RecursivelyIdentifyRange(root, wrapped_fs);
+  EXPECT_TRUE(wrapped_range.IsEmpty());
+}
+
 }  // namespace
 }  // namespace xla

From 9f24db984622355811289b1507dac6e4a960eba8 Mon Sep 17 00:00:00 2001
From: bodapatisaikrishna <saikrishnabodapati@gmail.com>
Date: Fri, 9 Oct 2026 20:19:58 +0530
Subject: [PATCH 0361/1000] Fix fatal process abort on resource variable dtype
 mismatch in training ops

Validate variable dtype against target dtype T in GetInputTensorFromVariable
and PrepareToUpdateVariable before accessing typed tensor buffers, preventing
fatal CheckTypeAndIsAligned abort crashes on dtype mismatches.

Fixes #113145.

Signed-off-by: bodapatisaikrishna <saikrishnabodapati@gmail.com>
---
 tensorflow/core/kernels/training_op_helpers.h | 17 +++++++
 .../python/training/training_ops_test.py      | 51 +++++++++++++++++++
 2 files changed, 68 insertions(+)

diff --git a/tensorflow/core/kernels/training_op_helpers.h b/tensorflow/core/kernels/training_op_helpers.h
index 1087df8f530..8f7a2fffef9 100644
--- a/tensorflow/core/kernels/training_op_helpers.h
+++ b/tensorflow/core/kernels/training_op_helpers.h
@@ -242,6 +242,11 @@ void MaybeForwardRefInputToRefOutput(OpKernelContext* ctx, int input,
 template <typename Device, typename T>
 absl::Status PrepareToUpdateVariable(OpKernelContext* ctx, Tensor* tensor,
                                      bool copy_on_read_mode) {
+  if (tensor->IsInitialized() && tensor->dtype() != DataTypeToEnum<T>::v()) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "dtype mismatch: expected ", DataTypeString(DataTypeToEnum<T>::v()),
+        " but got ", DataTypeString(tensor->dtype())));
+  }
   if (copy_on_read_mode || !tensor->RefCountIsOne()) {
     // Tensor's buffer is in use by some read, so we need to copy before
     // updating.
@@ -293,6 +298,13 @@ absl::Status GetInputTensorFromVariable(OpKernelContext* ctx, int input,
     ResourceHandle handle;
     TF_RETURN_IF_ERROR(HandleFromInput(ctx, input, &handle));
     TF_RETURN_IF_ERROR(LookupResource(ctx, handle, &var));
+    if ((var->is_initialized || var->tensor()->IsInitialized()) &&
+        var->tensor()->dtype() != DataTypeToEnum<T>::v()) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "dtype mismatch: expected ", DataTypeString(DataTypeToEnum<T>::v()),
+          " but got ", DataTypeString(var->tensor()->dtype()),
+          " (resource variable dtype)"));
+    }
     if (sparse) {
       var->mu()->assert_held_shared();
       *out = *var->tensor();
@@ -305,6 +317,11 @@ absl::Status GetInputTensorFromVariable(OpKernelContext* ctx, int input,
     return absl::OkStatus();
   }
   *out = ctx->mutable_input(input, lock_held);
+  if (out->IsInitialized() && out->dtype() != DataTypeToEnum<T>::v()) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "dtype mismatch: expected ", DataTypeString(DataTypeToEnum<T>::v()),
+        " but got ", DataTypeString(out->dtype())));
+  }
   return absl::OkStatus();
 }
 
diff --git a/tensorflow/python/training/training_ops_test.py b/tensorflow/python/training/training_ops_test.py
index e45edfe9fae..e2aad463ad1 100644
--- a/tensorflow/python/training/training_ops_test.py
+++ b/tensorflow/python/training/training_ops_test.py
@@ -697,6 +697,57 @@ class TrainingOpsTest(TensorFlowTestCase):
             )
         )
 
+  @test_util.run_v2_only
+  def testResourceVariableDtypeMismatch(self):
+    # A resource variable whose dtype differs from optimizer dtype T
+    # previously caused a fatal CHECK failure in CheckTypeAndIsAligned
+    # instead of raising an InvalidArgumentError (see GitHub issue #113145).
+    var = variables.Variable([[0.0, 0.0]], dtype=dtypes.float32)
+    accum = variables.Variable([[0.0, 0.0]], dtype=dtypes.float32)
+    var_f16 = variables.Variable([[0.0, 0.0]], dtype=dtypes.float16)
+    self.evaluate(variables.global_variables_initializer())
+
+    lr = constant_op.constant(0.1, dtype=dtypes.float16)
+    grad = constant_op.constant([[0.1, 0.2]], dtype=dtypes.float16)
+    indices = constant_op.constant([0], dtype=dtypes.int32)
+
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError, "dtype mismatch"
+    ):
+      self.evaluate(
+          gen_training_ops.resource_sparse_apply_adagrad(
+              var=var.handle,
+              accum=accum.handle,
+              lr=lr,
+              grad=grad,
+              indices=indices,
+          )
+      )
+
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError, "dtype mismatch"
+    ):
+      self.evaluate(
+          gen_training_ops.resource_sparse_apply_adagrad(
+              var=var_f16.handle,
+              accum=accum.handle,
+              lr=lr,
+              grad=grad,
+              indices=indices,
+          )
+      )
+
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError, "dtype mismatch"
+    ):
+      self.evaluate(
+          gen_training_ops.resource_apply_gradient_descent(
+              var=var.handle,
+              alpha=lr,
+              delta=grad,
+          )
+      )
+
 
 if __name__ == '__main__':
   googletest.main()

From 35746aec160345c2de2f243cc774e678028d39d3 Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Fri, 9 Oct 2026 07:53:16 -0700
Subject: [PATCH 0362/1000] [XLA] Fix while loop induction variable evaluation
 in HloEvaluator::HandleWhile

When `HloEvaluator::HandleWhile` partially evaluates a while loop whose input tuple contains unknown values, it folds the loop induction variable statically. It built the result literal directly from `trip_count`, ignoring the initial value and step size. Loops with a nonzero initial value or a step other than one got a wrong induction variable.

Compute the final value as `induction_var_init_value + trip_count * step_size`, matching `TryParseAndEvaluateWhileInductionVar`. The arithmetic is done in `uint64_t`, so overflow wraps the same way the narrower integer induction variable would instead of being signed overflow undefined behavior.

Partial evaluation of static loops now yields the correct induction variable. No flags are added or changed.

Tests:
- `//third_party/tensorflow/compiler/xla/hlo/evaluator:hlo_evaluator_test`: pins partial evaluation of a while loop induction variable with a nonzero initial value and a step size greater than one.
PiperOrigin-RevId: 996531246
---
 .../xla/xla/hlo/evaluator/hlo_evaluator.cc    | 14 +++++--
 .../xla/hlo/evaluator/hlo_evaluator_test.cc   | 40 +++++++++++++++++++
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
index 9f25cc71ecc..a1ac820152b 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
@@ -3918,10 +3918,16 @@ absl::Status HloEvaluator::HandleWhile(const HloInstruction* while_hlo) {
     }
     Shape induction_var_shape =
         ShapeUtil::GetSubshape(while_hlo->shape(), visitor_shape_index_);
-    int64_t trip_count = parsed_while_loop->static_while_loop->trip_count;
-    ABSL_ASSIGN_OR_RETURN(
-        Literal induction_var_val,
-        CreateScalarLiteral(trip_count, induction_var_shape.element_type()));
+    int64_t induction_var_value = static_cast<int64_t>(
+        static_cast<uint64_t>(
+            parsed_while_loop->static_while_loop->induction_var_init_value) +
+        static_cast<uint64_t>(
+            parsed_while_loop->static_while_loop->trip_count) *
+            static_cast<uint64_t>(
+                parsed_while_loop->static_while_loop->step_size));
+    ABSL_ASSIGN_OR_RETURN(Literal induction_var_val,
+                     CreateScalarLiteral(induction_var_value,
+                                         induction_var_shape.element_type()));
     ABSL_RETURN_IF_ERROR(literal.CopyFrom(induction_var_val,
                                      /*dest_shape_index=*/visitor_shape_index_,
                                      /*src_shape_index=*/{}));
diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc
index 6d79d18a4b4..7249ecb44c7 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc
@@ -7417,6 +7417,46 @@ TEST_F(HloEvaluatorTest, ParameterThroughCallSucceedsWithPrecomputation) {
   EXPECT_TRUE(LiteralTestUtil::Equal(expected, result));
 }
 
+TEST_F(HloEvaluatorTest,
+       EvaluateWhileInductionVarWithNonUnitStepAndNonZeroInit) {
+  constexpr absl::string_view kHloModule = R"(
+    HloModule while_induction_var
+
+    %while_condition {
+      %param = (s32[], f32[4]) parameter(0)
+      %gte.0 = s32[] get-tuple-element(%param), index=0
+      %loop_bound = s32[] constant(23)
+      ROOT %result = pred[] compare(%gte.0, %loop_bound), direction=LT
+    }
+
+    %while_body {
+      %param = (s32[], f32[4]) parameter(0)
+      %gte.0 = s32[] get-tuple-element(%param), index=0
+      %gte.1 = f32[4] get-tuple-element(%param), index=1
+      %step = s32[] constant(4)
+      %next_indvar = s32[] add(%gte.0, %step)
+      %next_buf = f32[4] add(%gte.1, %gte.1)
+      ROOT %loop_result = (s32[], f32[4]) tuple(%next_indvar, %next_buf)
+    }
+
+    ENTRY main {
+      %param.0 = f32[4] parameter(0)
+      %init = s32[] constant(3)
+      %while_init = (s32[], f32[4]) tuple(%init, %param.0)
+      %while = (s32[], f32[4]) while(%while_init), condition=%while_condition, body=%while_body
+      ROOT %indvar = s32[] get-tuple-element(%while), index=0
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> hlo_module,
+                       ParseAndReturnVerifiedModule(kHloModule));
+  ASSERT_OK_AND_ASSIGN(
+      Literal result,
+      evaluator_.Evaluate(hlo_module->entry_computation()->root_instruction(),
+                          /*precomputed_analyses=*/{},
+                          /*recursively_evaluate_nonconstant_operands=*/true));
+  EXPECT_EQ(result, LiteralUtil::CreateR0<int32_t>(23));
+}
+
 class PatternMatchParseWhileLoopTest : public HloHardwareIndependentTestBase {};
 
 TEST_F(PatternMatchParseWhileLoopTest, LoopBoundDefinedInsideOfCond) {

From 30043c480348298ff8dcd3a0dbebbc9ba4f14974 Mon Sep 17 00:00:00 2001
From: Som Tripathi <somtri@iastate.edu>
Date: Fri, 9 Oct 2026 11:00:12 -0500
Subject: [PATCH 0363/1000] Widen GPU Where CUB item counts and offsets to 64
 bit

The CUB wrappers in where_op_gpu.cu.h took an int num_items, so an
input with more than INT32_MAX elements reached DeviceSelect::Flagged
with a truncated count and the output was never written. CCCL 2.8.5
and 3.2.0 take 64-bit counts, so pass TIndex instead. Index
WhereOutputIterator with std::ptrdiff_t and loop over the output rows
with a 64-bit index. Remove the INT32_MAX guard from where_op.cc.
---
 tensorflow/core/kernels/BUILD                 |  1 -
 tensorflow/core/kernels/where_op.cc           | 16 +-----
 tensorflow/core/kernels/where_op_gpu.cu.h     | 27 +++++-----
 .../python/kernel_tests/array_ops/BUILD       |  1 -
 .../kernel_tests/array_ops/where_op_test.py   | 50 ++++++++++++++-----
 5 files changed, 54 insertions(+), 41 deletions(-)

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index d11d2d06d32..390d2ab957e 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -1354,7 +1354,6 @@ tf_kernel_library(
     ]) + [
         ":gpu_prim_hdrs",
         "@com_google_absl//absl/status",
-        "@com_google_absl//absl/strings",
     ] + ARRAY_DEPS,
 )
 
diff --git a/tensorflow/core/kernels/where_op.cc b/tensorflow/core/kernels/where_op.cc
index 1117dd210c4..e5fff24c745 100644
--- a/tensorflow/core/kernels/where_op.cc
+++ b/tensorflow/core/kernels/where_op.cc
@@ -23,12 +23,9 @@ limitations under the License.
 
 #include "tensorflow/core/kernels/where_op.h"
 
-#include <limits>
 #include <memory>
 #include <numeric>
-
-#include "absl/status/status.h"
-#include "absl/strings/str_cat.h"
+#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
 #include "tensorflow/core/framework/bounds_check.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
@@ -39,7 +36,6 @@ limitations under the License.
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/types.h"
-#include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
 
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 #include "tensorflow/core/common_runtime/gpu/gpu_event_mgr.h"
@@ -263,16 +259,6 @@ class WhereGPUOp : public AsyncOpKernel {
     const Tensor& input = context->input(0);
     const int input_dims = input.dims();
 
-    // The CUB count and select calls take an int item count. Larger inputs
-    // truncate it, and the select then leaves the output unwritten.
-    OP_REQUIRES_ASYNC(context,
-                      input.NumElements() < std::numeric_limits<int32_t>::max(),
-                      absl::InvalidArgumentError(absl::StrCat(
-                          "WhereOp: the GPU kernel supports fewer than ",
-                          std::numeric_limits<int32_t>::max(),
-                          " input elements, got ", input.NumElements(), ".")),
-                      done);
-
     if (input.NumElements() < std::numeric_limits<int32_t>::max()) {
       ComputeAsyncType<int32_t>(input, input_dims, context, done);
     } else {
diff --git a/tensorflow/core/kernels/where_op_gpu.cu.h b/tensorflow/core/kernels/where_op_gpu.cu.h
index b57be0f0802..e63eee8833e 100644
--- a/tensorflow/core/kernels/where_op_gpu.cu.h
+++ b/tensorflow/core/kernels/where_op_gpu.cu.h
@@ -43,11 +43,11 @@ __global__ void PropagateWhereIndicesKernel(
   // TODO(ebrevdo): Use a multi-dimensional loop, increasing the
   // dimensions of individual indices manually, instead of relying on
   // a scalar loop variable and using integer division.
-  GPU_1D_KERNEL_LOOP(i, output_rows) {
-    TIndex index_value = ldg(output + static_cast<int64_t>(NDIM) * i);
+  for (int64_t i : GpuGridRangeX<int64_t>(output_rows)) {
+    TIndex index_value = ldg(output + NDIM * i);
 #pragma unroll
     for (int c = 0; c < NDIM; ++c) {
-      *(output + static_cast<int64_t>(NDIM) * i + c) = index_value / strides[c];
+      *(output + NDIM * i + c) = index_value / strides[c];
       index_value %= strides[c];
     }
   }
@@ -67,7 +67,7 @@ struct IsNonzero {
 template <typename T, typename TIndex>
 struct CubDeviceReduceCount {
   gpuError_t operator()(void* d_temp_storage, size_t& temp_storage_bytes,
-                        const T* d_in, TIndex* d_out, int num_items,
+                        const T* d_in, TIndex* d_out, TIndex num_items,
                         gpuStream_t stream = 0) {
     IsNonzero<T> is_nonzero;
     gpuprim::TransformInputIterator<bool, IsNonzero<T>, const T*>
@@ -81,7 +81,7 @@ struct CubDeviceReduceCount {
 template <typename TIndex>
 struct CubDeviceReduceCount<bool, TIndex> {
   gpuError_t operator()(void* d_temp_storage, size_t& temp_storage_bytes,
-                        const bool* d_in, TIndex* d_out, int num_items,
+                        const bool* d_in, TIndex* d_out, TIndex num_items,
                         gpuStream_t stream = 0) {
     return gpuprim::DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, d_in,
                                       d_out, num_items, stream);
@@ -97,7 +97,7 @@ struct CubDeviceSelectFlaggedCounter<T, TIndex, OutputIterator,
                                      false /*IsConvertibleToBool*/> {
   gpuError_t operator()(void* d_temp_storage, size_t& temp_storage_bytes,
                         const T* d_flags, OutputIterator d_out,
-                        TIndex* d_num_selected_out, int num_items,
+                        TIndex* d_num_selected_out, TIndex num_items,
                         gpuStream_t stream = 0) {
     gpuprim::CountingInputIterator<TIndex> select_counter(0);
     IsNonzero<T> is_nonzero;
@@ -115,7 +115,7 @@ struct CubDeviceSelectFlaggedCounter<T, TIndex, OutputIterator,
                                      true /*IsConvertibleToBool*/> {
   gpuError_t operator()(void* d_temp_storage, size_t& temp_storage_bytes,
                         const T* d_flags, OutputIterator d_out,
-                        TIndex* d_num_selected_out, int num_items,
+                        TIndex* d_num_selected_out, TIndex num_items,
                         gpuStream_t stream = 0) {
     gpuprim::CountingInputIterator<TIndex> select_counter(0);
     return gpuprim::DeviceSelect::Flagged(
@@ -218,7 +218,8 @@ class WhereOutputIterator {
   WhereOutputIterator(int64_t* ptr, const Eigen::DenseIndex max_row)
       : ptr_(ptr), max_row_(max_row) {}
 
-  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE int64_t& operator[](int n) const {
+  EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE int64_t& operator[](
+      std::ptrdiff_t n) const {
     // If the selection mechanism finds too many true values (because
     // the input tensor changed between allocation of output and now),
     // we may accidentally try to write past the allowable memory.  If
@@ -226,7 +227,7 @@ class WhereOutputIterator {
     // the number of items found in Flagged()'s d_num_selected_out at
     // the end and confirm that it matches the number of rows of output.
     const bool valid = FastBoundsCheck(n, max_row_);
-    return *(ptr_ + (valid ? (static_cast<int64_t>(NDIM) * n) : 0));
+    return *(ptr_ + (valid ? (NDIM * n) : 0));
   }
 
   EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE reference operator*() const {
@@ -328,10 +329,12 @@ struct Where<GPUDevice, NDIM, T, TIndex> {
     const Eigen::array<TIndex, NDIM> strides =
         CalculateStrides<TIndex, T, NDIM>(input);
     const TIndex output_rows = output.dimension(0);
-    GpuLaunchConfig config = GetGpuLaunchConfig(output_rows, d);
+    absl::StatusOr<GpuLaunchConfig64> config =
+        GetGpuLaunchConfig64(output_rows, d);
+    if (!config.ok()) return config.status();
     TF_CHECK_OK(GpuLaunchKernel(PropagateWhereIndicesKernel<NDIM, TIndex>,
-                                config.block_count, config.thread_per_block, 0,
-                                d.stream(), output_rows, strides,
+                                config->block_count, config->thread_per_block,
+                                0, d.stream(), output_rows, strides,
                                 output.data()));
 
     return absl::OkStatus();
diff --git a/tensorflow/python/kernel_tests/array_ops/BUILD b/tensorflow/python/kernel_tests/array_ops/BUILD
index d060af574ca..1278cc3cb75 100644
--- a/tensorflow/python/kernel_tests/array_ops/BUILD
+++ b/tensorflow/python/kernel_tests/array_ops/BUILD
@@ -871,7 +871,6 @@ cuda_py_strict_test(
     deps = [
         "//tensorflow/python/client:session",
         "//tensorflow/python/framework:constant_op",
-        "//tensorflow/python/framework:errors",
         "//tensorflow/python/framework:for_generated_wrappers",
         "//tensorflow/python/framework:test_lib",
         "//tensorflow/python/ops:array_ops",
diff --git a/tensorflow/python/kernel_tests/array_ops/where_op_test.py b/tensorflow/python/kernel_tests/array_ops/where_op_test.py
index 8d97d8270f8..6d32c836ab8 100644
--- a/tensorflow/python/kernel_tests/array_ops/where_op_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/where_op_test.py
@@ -22,7 +22,6 @@ import numpy as np
 from tensorflow.python.client import session
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
@@ -265,27 +264,54 @@ class WhereOpTest(test.TestCase):
     self.assertAllEqual(tf_val, np_val)
 
   @test_util.run_in_graph_and_eager_modes
-  @test_util.disable_xla("Tests the TF GPU Where kernel element-count guard")
-  def testGpuRejectsInt32MaxElements(self):
-    # The GPU kernel passes the element count to CUB as an int. Past INT32_MAX
-    # elements the count truncated and the op returned uninitialized memory.
+  @test_util.disable_xla(
+      "Tests the TF GPU Where kernel with INT32_MAX or more elements"
+  )
+  def testGpuAcceptsInt32MaxElements(self):
     if not test_util.is_gpu_available():
       self.skipTest("Requires a GPU.")
     for shape in ([2**31 - 1], [2**16, 2**15]):
       with self.subTest(shape=shape), test_util.device(use_gpu=True):
         x = array_ops.zeros(shape, dtype=dtypes.bool)
-        with self.assertRaisesRegex(errors.InvalidArgumentError,
-                                    "fewer than 2147483647 input elements"):
-          self.evaluate(array_ops.where(x))
+        self.assertAllEqual(
+            self.evaluate(array_ops.where(x)).shape, [0, len(shape)]
+        )
         del x
 
+  def _boolWithTrueAt(self, n, positions):
+    # Builds a flat bool tensor of n elements on the current device, True only
+    # at the sorted flat positions.
+    pieces = []
+    prev = 0
+    for p in positions:
+      if p > prev:
+        pieces.append(array_ops.zeros([p - prev], dtype=dtypes.bool))
+      pieces.append(array_ops.ones([1], dtype=dtypes.bool))
+      prev = p + 1
+    if n > prev:
+      pieces.append(array_ops.zeros([n - prev], dtype=dtypes.bool))
+    return array_ops.concat(pieces, 0)
+
   @test_util.run_in_graph_and_eager_modes
-  def testGpuAcceptsInputBelowInt32MaxElements(self):
+  @test_util.disable_xla(
+      "Tests the TF GPU Where kernel with INT32_MAX or more elements"
+  )
+  def testGpuWhereIndicesPastInt32Max(self):
+    # An all-False input has no true elements to select, so it cannot show a
+    # truncated count. This input has true elements at flat indices above
+    # INT32_MAX.
     if not test_util.is_gpu_available():
       self.skipTest("Requires a GPU.")
-    with test_util.device(use_gpu=True):
-      x = array_ops.zeros([2**31 - 2], dtype=dtypes.bool)
-      self.assertAllEqual(self.evaluate(array_ops.where(x)).shape, [0, 1])
+    for shape in ([2**31 - 1], [2**31 + 2**16], [2**16, 2**15 + 1]):
+      n = int(np.prod(shape))
+      positions = sorted(
+          {p for p in (0, 12345, n // 2, 2**31 - 1, 2**31, n - 1) if p < n}
+      )
+      expected = np.array(np.unravel_index(positions, shape)).T.astype(np.int64)
+      with self.subTest(shape=shape), test_util.device(use_gpu=True):
+        x = array_ops.reshape(self._boolWithTrueAt(n, positions), shape)
+        self.assertAllEqual(self.evaluate(array_ops.where(x)), expected)
+        del x
 
 
 class WhereBenchmark(test.Benchmark):

From 060225cb79c06eb9bb00757c9afd9e50dc42ef07 Mon Sep 17 00:00:00 2001
From: Mikhail Goncharov <goncharov@google.com>
Date: Fri, 9 Oct 2026 09:07:31 -0700
Subject: [PATCH 0364/1000] [ROCm] Add TritonAMDGPUMembar pass to the Triton
 ROCm compilation pipeline.

Replaces the duplicated TritonGPUGlobalScratchAllocationPass with
createTritonAMDGPUMembar, which inserts the required memory barriers
for shared memory accesses on AMD GPUs.

Replicates https://github.com/openxla/xla/pull/50386.

PiperOrigin-RevId: 996568186
---
 .../backends/gpu/codegen/triton/compilation_pipeline_rocm.cc    | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/compilation_pipeline_rocm.cc b/third_party/xla/xla/backends/gpu/codegen/triton/compilation_pipeline_rocm.cc
index 5590d038d4d..2468daa31cb 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/compilation_pipeline_rocm.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/compilation_pipeline_rocm.cc
@@ -160,7 +160,7 @@ static void MakeLLIR(mlir::OpPassManager* pm,
   pm->addPass(mlir::createInlinerPass());
   pm->addPass(mlir::createConvertIndexToLLVMPass());
   pm->addPass(mt::createAllocateAMDGPUSharedMemoryPass(rocm_cc.gfx_version()));
-  pm->addPass(mt::gpu::createTritonGPUGlobalScratchAllocationPass());
+  pm->addPass(mt::createTritonAMDGPUMembar({rocm_cc.gfx_version()}));
   pm->addPass(mt::gpu::createTritonGPUGlobalScratchAllocationPass());
   pm->addPass(
       mt::createConvertTritonAMDGPUToLLVMPass(rocm_cc.gfx_version(), true));

From a8493ee0eff86447145a972413d918d8a7c34c40 Mon Sep 17 00:00:00 2001
From: Raashish Aggarwal <94279692+raashish1601@users.noreply.github.com>
Date: Fri, 9 Oct 2026 21:46:33 +0530
Subject: [PATCH 0365/1000] Support integer values in
 histogram_fixed_width_bins

truediv returns a float tensor for integer values, so cast nbins to the dtype of the scaled values instead of the dtype of values. Test integer values with an integer value_range and check the output dtype.
---
 .../python/kernel_tests/histogram_ops_test.py       | 13 ++++++++++---
 tensorflow/python/ops/histogram_ops.py              |  3 ++-
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/kernel_tests/histogram_ops_test.py b/tensorflow/python/kernel_tests/histogram_ops_test.py
index 44ac4ee5535..ef1d30a7ff0 100644
--- a/tensorflow/python/kernel_tests/histogram_ops_test.py
+++ b/tensorflow/python/kernel_tests/histogram_ops_test.py
@@ -55,14 +55,21 @@ class BinValuesFixedWidth(test.TestCase, parameterized.TestCase):
       self.assertEqual(dtypes.int32, bins.dtype)
       self.assertAllClose(expected_bins, self.evaluate(bins))
 
-  @parameterized.parameters(np.float32, np.float64)
-  def test_1d_values_python_value_range(self, dtype):
+  @parameterized.parameters(
+      (np.float32, [0.0, 5.0]),
+      (np.float64, [0.0, 5.0]),
+      (np.float64, [0, 5]),
+      (np.int32, [0, 5]),
+      (np.int64, [0, 5]),
+  )
+  def test_1d_values_python_value_range(self, dtype, value_range):
     # value_range given as a Python list takes the dtype of values.
     values = np.array([-1.0, 0.0, 1.5, 2.0, 5.0, 15]).astype(dtype)
     expected_bins = [0, 0, 1, 2, 4, 4]
     with self.cached_session():
       bins = histogram_ops.histogram_fixed_width_bins(
-          values, [0.0, 5.0], nbins=5)
+          values, value_range, nbins=5)
+      self.assertEqual(dtypes.int32, bins.dtype)
       self.assertAllClose(expected_bins, self.evaluate(bins))
 
   def test_2d_values(self):
diff --git a/tensorflow/python/ops/histogram_ops.py b/tensorflow/python/ops/histogram_ops.py
index a156144afe5..f28e2905da0 100644
--- a/tensorflow/python/ops/histogram_ops.py
+++ b/tensorflow/python/ops/histogram_ops.py
@@ -83,13 +83,14 @@ def histogram_fixed_width_bins(values,
     check = control_flow_assert.Assert(
         math_ops.greater(nbins, 0), ['nbins %s must > 0' % nbins])
     nbins = control_flow_ops.with_dependencies([check], nbins)
-    nbins_float = math_ops.cast(nbins, values.dtype)
 
     # Map tensor values that fall within value_range to [0, 1].
     scaled_values = math_ops.truediv(
         values - value_range[0],
         value_range[1] - value_range[0],
         name='scaled_values')
+    # truediv returns a float tensor for integer values.
+    nbins_float = math_ops.cast(nbins, scaled_values.dtype)
 
     # map tensor values within the open interval value_range to {0,.., nbins-1},
     # values outside the open interval will be zero or less, or nbins or more.

From 50705b9b7448d33c1cace5650b3a0649d757a387 Mon Sep 17 00:00:00 2001
From: Raashish Aggarwal <94279692+raashish1601@users.noreply.github.com>
Date: Fri, 9 Oct 2026 21:49:42 +0530
Subject: [PATCH 0366/1000] Test var with ddof for empty and out-of-bounds axes

---
 tensorflow/python/ops/numpy_ops/np_array_ops_test.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index 79a7d3d9231..1ea68abd192 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -995,6 +995,9 @@ class ArrayMethodsTest(test.TestCase):
     run_test([[1.0j, 2.0], [3.0j, 4.0]], axis=0, ddof=1)
     run_test(np.arange(8.0).reshape((2, 2, 2)).tolist(), axis=(0, 2), ddof=1)
     run_test(np.arange(8.0).reshape((2, 2, 2)).tolist(), axis=(-1, 0), ddof=1)
+    run_test(np.arange(8.0).reshape((2, 2, 2)).tolist(), axis=(), ddof=1)
+    run_test(5.0, ddof=1)
+    run_test(5.0, axis=(), ddof=1)
     run_test(np.arange(8).reshape((2, 2, 2)).tolist(), axis=(0, 2))
     run_test(
         np.arange(8).reshape((2, 2, 2)).tolist(), axis=(0, 2), keepdims=True)
@@ -1002,6 +1005,10 @@ class ArrayMethodsTest(test.TestCase):
     run_test(
         np.arange(8).reshape((2, 2, 2)).tolist(), axis=(2, 0), keepdims=True)
     self.assertRaises(ValueError, np_array_ops.var, np.ones([2, 2]), out=[])
+    with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
+      np_array_ops.var(np.ones([2, 2]), axis=-3, ddof=1)
+    with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
+      np_array_ops.var(np.ones([2, 2]), axis=(-3, 0), ddof=1)
 
   def testProd(self):
 

From 0391fcd54bfbe57b6e8c393295410e52664205ae Mon Sep 17 00:00:00 2001
From: Dmitri Latushko <latushko@google.com>
Date: Fri, 9 Oct 2026 09:09:32 -0700
Subject: [PATCH 0367/1000] Increase shard_count for
 linear_operator_low_rank_update_test from 15 to 30 to avoid TSAN timeouts.

PiperOrigin-RevId: 996569205
---
 tensorflow/python/kernel_tests/linalg/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/kernel_tests/linalg/BUILD b/tensorflow/python/kernel_tests/linalg/BUILD
index 77e8ad45eaf..fef7d0c72aa 100644
--- a/tensorflow/python/kernel_tests/linalg/BUILD
+++ b/tensorflow/python/kernel_tests/linalg/BUILD
@@ -433,7 +433,7 @@ cuda_py_strict_test(
     name = "linear_operator_low_rank_update_test",
     size = "medium",
     srcs = ["linear_operator_low_rank_update_test.py"],
-    shard_count = 15,
+    shard_count = 30,
     tags = ["optonly"],
     deps = [
         "//tensorflow/python/framework:config",

From 3f5fb7602e5e0595f7bb50c45fade7f21709f89c Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Fri, 9 Oct 2026 09:16:51 -0700
Subject: [PATCH 0368/1000] Fix index clamping, unsigned 64 bit constant
 handling, and root instruction preservation in
 DeadDynamicUpdateSliceElimination.

PiperOrigin-RevId: 996573049
---
 .../dead_dynamic_update_slice_elimination.cc  | 36 +++++--------
 ...d_dynamic_update_slice_elimination_test.cc | 51 +++++++++++++++++++
 2 files changed, 63 insertions(+), 24 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/dead_dynamic_update_slice_elimination.cc b/third_party/xla/xla/hlo/transforms/simplifiers/dead_dynamic_update_slice_elimination.cc
index 3969caaff48..51ead004af9 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/dead_dynamic_update_slice_elimination.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/dead_dynamic_update_slice_elimination.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/hlo/transforms/simplifiers/dead_dynamic_update_slice_elimination.h"
 
+#include <algorithm>
 #include <cstdint>
 #include <optional>
 #include <vector>
@@ -31,30 +32,11 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_opcode.h"
-#include "xla/primitive_util.h"
 #include "xla/shape.h"
-#include "xla/tsl/platform/errors.h"
-#include "xla/tsl/platform/statusor.h"
 
 namespace xla {
 namespace {
 
-std::optional<int64_t> GetConstantAsInt64(const HloInstruction* inst) {
-  if (!inst->IsConstant() || !ShapeUtil::IsScalar(inst->shape())) {
-    return std::nullopt;
-  }
-  return primitive_util::PrimitiveTypeSwitch<std::optional<int64_t>>(
-      [&](auto primitive_type_constant) -> std::optional<int64_t> {
-        if constexpr (primitive_util::IsIntegralType(primitive_type_constant)) {
-          using NativeT = primitive_util::NativeTypeOf<primitive_type_constant>;
-          return static_cast<int64_t>(
-              inst->literal().GetFirstElement<NativeT>());
-        }
-        return std::nullopt;
-      },
-      inst->shape().element_type());
-}
-
 std::optional<std::vector<int64_t>> GetStartIndices(const HloInstruction* dus) {
   absl::Span<HloInstruction* const> start_indices_operands =
       absl::MakeSpan(dus->operands())
@@ -62,7 +44,10 @@ std::optional<std::vector<int64_t>> GetStartIndices(const HloInstruction* dus) {
                        ->first_index_operand_number());
   std::vector<int64_t> start_indices;
   for (HloInstruction* operand : start_indices_operands) {
-    std::optional<int64_t> start_index = GetConstantAsInt64(operand);
+    if (!operand->IsConstant()) {
+      return std::nullopt;
+    }
+    std::optional<int64_t> start_index = operand->literal().GetFirstInteger();
     if (!start_index.has_value()) {
       return std::nullopt;
     }
@@ -84,7 +69,7 @@ bool RangesOverlap(int64_t start1, int64_t end1, int64_t start2, int64_t end2) {
 // If true, the updated elements of the dynamic-update-slice is not accessed
 // by the slice user.
 bool IsDusUpdateUnused(const std::vector<int64_t>& dus_starts,
-                       const Shape& update_shape,
+                       const Shape& operand_shape, const Shape& update_shape,
                        const HloInstruction* slice_user) {
   if (slice_user->opcode() != HloOpcode::kSlice) {
     return false;
@@ -97,7 +82,9 @@ bool IsDusUpdateUnused(const std::vector<int64_t>& dus_starts,
   // dimensions. If there is no overlap in any dimension, the slice is safe,
   // i.e., it doesn't access the updated elements.
   for (int dim = 0; dim < update_shape.dimensions().size(); ++dim) {
-    int64_t dus_start = dus_starts[dim];
+    int64_t dus_start = std::clamp<int64_t>(
+        dus_starts[dim], 0,
+        operand_shape.dimensions(dim) - update_shape.dimensions(dim));
     int64_t dus_limit = dus_start + update_shape.dimensions(dim);
     int64_t slice_start = slice_starts[dim];
     int64_t slice_limit = slice_limits[dim];
@@ -132,9 +119,10 @@ absl::StatusOr<bool> ProcessDynamicUpdateSlice(HloInstruction* dus,
   }
 
   bool is_dus_update_unused =
-      dus->user_count() > 0 &&
+      dus != comp->root_instruction() && dus->user_count() > 0 &&
       absl::c_all_of(dus->users(), [&](HloInstruction* user) {
-        return IsDusUpdateUnused(dus_starts_vec, update_operand->shape(), user);
+        return IsDusUpdateUnused(dus_starts_vec, dus->operand(0)->shape(),
+                                 update_operand->shape(), user);
       });
   VLOG(2) << "  is_dus_update_unused: " << is_dus_update_unused;
   if (is_dus_update_unused) {
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/dead_dynamic_update_slice_elimination_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/dead_dynamic_update_slice_elimination_test.cc
index 9a96ddff605..30f37669cad 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/dead_dynamic_update_slice_elimination_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/dead_dynamic_update_slice_elimination_test.cc
@@ -136,5 +136,56 @@ ENTRY main {
   }
 }
 
+TEST_F(DeadDynamicUpdateSliceEliminationTest, NegativeStartIndicesClamped) {
+  const absl::string_view kHlo = R"(
+HloModule negative_start
+
+ENTRY main {
+  %param = f32[10] parameter(0)
+  %update = f32[4] parameter(1)
+  %start = s32[] constant(-5)
+  %dus = f32[10] dynamic-update-slice(%param, %update, %start)
+  ROOT %slice = f32[4] slice(%dus), slice={[0:4]}
+})";
+  TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo));
+  DeadDynamicUpdateSliceElimination dds;
+  ASSERT_OK_AND_ASSIGN(bool changed, dds.Run(module.get()));
+  EXPECT_FALSE(changed);
+}
+
+TEST_F(DeadDynamicUpdateSliceEliminationTest, LargeU64StartIndicesClamped) {
+  const absl::string_view kHlo = R"(
+HloModule large_u64_start
+
+ENTRY main {
+  %param = f32[10] parameter(0)
+  %update = f32[4] parameter(1)
+  %start = u64[] constant(9223372036854775808)
+  %dus = f32[10] dynamic-update-slice(%param, %update, %start)
+  ROOT %slice = f32[4] slice(%dus), slice={[6:10]}
+})";
+  TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo));
+  DeadDynamicUpdateSliceElimination dds;
+  ASSERT_OK_AND_ASSIGN(bool changed, dds.Run(module.get()));
+  EXPECT_FALSE(changed);
+}
+
+TEST_F(DeadDynamicUpdateSliceEliminationTest, RootDusPreserved) {
+  const absl::string_view kHlo = R"(
+HloModule root_dus
+
+ENTRY main {
+  %param = f32[10] parameter(0)
+  %update = f32[4] parameter(1)
+  %start = s32[] constant(0)
+  ROOT %dus = f32[10] dynamic-update-slice(%param, %update, %start)
+  %unused_slice = f32[4] slice(%dus), slice={[6:10]}
+})";
+  TF_ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo));
+  DeadDynamicUpdateSliceElimination dds;
+  ASSERT_OK_AND_ASSIGN(bool changed, dds.Run(module.get()));
+  EXPECT_FALSE(changed);
+}
+
 }  // namespace
 }  // namespace xla

From 906aed82f1e0c0bbad7daafbb001b6ddc04a4752 Mon Sep 17 00:00:00 2001
From: Bill Varcho <varcho@google.com>
Date: Fri, 9 Oct 2026 09:38:15 -0700
Subject: [PATCH 0369/1000] [SDY][HloShardingV3][Remove-Opt-Out][re-land]
 Support HloShardingV3 for shard maps without frontend attributes.

PiperOrigin-RevId: 996584612
---
 .../service/spmd/shardy/sdy_round_trip/BUILD  |   2 +
 .../sdy_round_trip/import_shardy_attrs.cc     |  64 +++-
 .../shardy/sdy_round_trip/shard_map_export.cc | 127 +++++-
 .../shardy/sdy_round_trip/shard_map_import.cc | 107 +++++-
 .../stablehlo_round_trip/export_shardings.cc  |  39 +-
 ...ound_trip_pipeline_manual_computation.mlir |  39 ++
 ..._trip_import_pipeline_hlo_sharding_v3.mlir |  24 ++
 ...trip_shard_map_export_hlo_sharding_v3.mlir | 285 ++++++++++++++
 ...y_round_trip_shard_map_import_failure.mlir |  35 ++
 ...trip_shard_map_import_hlo_sharding_v3.mlir | 362 ++++++++++++++++++
 10 files changed, 1045 insertions(+), 39 deletions(-)
 create mode 100644 third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_export_hlo_sharding_v3.mlir
 create mode 100644 third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_hlo_sharding_v3.mlir

diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD
index 2ca1c537bfc..8b05f417940 100644
--- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD
+++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD
@@ -74,6 +74,8 @@ cc_library(
     deps = [
         "//xla:side_effect_util",
         "//xla/hlo/ir:hlo_sharding",
+        "//xla/hlo/ir:mesh_and_axis",
+        "//xla/hlo/ir:named_sharding",
         "//xla/hlo/parser:hlo_parser",
         "//xla/hlo/translate/hlo_to_mhlo:hlo_utils",
         "//xla/service/spmd/shardy:constants",
diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc
index 30228198eaa..4b3f5b085e2 100644
--- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc
+++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc
@@ -55,6 +55,8 @@ limitations under the License.
 #include "shardy/dialect/sdy/ir/utils.h"
 #include "stablehlo/dialect/StablehloOps.h"
 #include "xla/hlo/ir/hlo_sharding.h"
+#include "xla/hlo/ir/mesh_and_axis.h"
+#include "xla/hlo/ir/named_sharding.h"
 #include "xla/hlo/parser/hlo_parser.h"
 #include "xla/hlo/translate/hlo_to_mhlo/hlo_utils.h"
 #include "xla/service/spmd/shardy/constants.h"
@@ -82,6 +84,7 @@ using ::mlir::func::FuncOp;
 using ::mlir::sdy::getSharding;
 using ::mlir::sdy::kShardingAttr;
 using ::mlir::sdy::kShardingRuleAttr;
+using ::mlir::sdy::ManualAxesAttr;
 using ::mlir::sdy::MeshAttr;
 using ::mlir::sdy::OpShardingRuleAttr;
 using ::mlir::sdy::PropagationBarrierOp;
@@ -206,6 +209,37 @@ bool handleFuncResultSharding(
   return anyChanged;
 }
 
+ManualAxesAttr extractManualAxes(const xla::HloSharding& hloSharding,
+                                 mlir::MLIRContext* context) {
+  if (hloSharding.IsTuple()) {
+    // Manual axes are only attached to ranked tensors, so a leading token
+    // element can legitimately have none. Take the first element that does.
+    for (const xla::HloSharding& element : hloSharding.tuple_elements()) {
+      if (ManualAxesAttr manualAxes = extractManualAxes(element, context)) {
+        return manualAxes;
+      }
+    }
+    return nullptr;
+  }
+  if (hloSharding.UseNamedShardingLeaf()) {
+    const xla::NamedSharding& namedSharding = hloSharding.named_sharding();
+    llvm::SmallVector<mlir::StringAttr> manualAxes;
+    const xla::Mesh& mesh = namedSharding.mesh();
+    for (const auto& axisRef : namedSharding.manual_axes()) {
+      // Silently dropping an axis here would yield a manual computation with
+      // fewer manual axes than it actually has, which miscompiles rather than
+      // fails, so treat an out-of-range index as a bug.
+      CHECK_LT(axisRef.mesh_axis_index(), mesh.axis_names().size());
+      manualAxes.push_back(mlir::StringAttr::get(
+          context, mesh.axis_names()[axisRef.mesh_axis_index()]));
+    }
+    if (!manualAxes.empty()) {
+      return ManualAxesAttr::get(context, manualAxes);
+    }
+  }
+  return nullptr;
+}
+
 // The sharding information is in the `kXlaShardingAttr` attribute.
 void convertShardyAttrsWithHloShardingV3(FuncOp funcOp) {
   for (auto [argNum, argType] : llvm::enumerate(funcOp.getArgumentTypes())) {
@@ -272,12 +306,38 @@ void convertShardyAttrsWithHloShardingV3(FuncOp funcOp) {
       StringRef targetName = customCallOp.getCallTargetName();
       if (targetName == kShardingCustomCallTargetName ||
           targetName == "X64Combine" ||
+          targetName == kGlobalToLocalShapeCallTargetName ||
+          targetName == kLocalToGlobalShapeCallTargetName ||
           isPythonCallbackCustomCall(customCallOp)) {
+        HloSharding hloSharding = parseShardingFromString(shardingAttr);
         if (auto sdySharding = convertToSdySharding(
-                parseShardingFromString(shardingAttr),
-                customCallOp->getResultTypes(), customCallOp->getContext())) {
+                hloSharding, customCallOp->getResultTypes(),
+                customCallOp->getContext())) {
           customCallOp->setAttr(kShardingAttr, sdySharding);
         }
+        // Only `xla.sdy.GlobalToLocalShape` (when there are operands) and the
+        // manual computation CallOp (when there are no operands) may carry
+        // manual axes.
+        if (targetName == kGlobalToLocalShapeCallTargetName) {
+          if (auto manualAxesAttr =
+                  extractManualAxes(hloSharding, customCallOp->getContext())) {
+            customCallOp->setAttr(kManualAxes, manualAxesAttr);
+          }
+        }
+      }
+    } else if (auto callOp = mlir::dyn_cast<mlir::func::CallOp>(op);
+               callOp && isManualComputation(callOp)) {
+      // The body call op carries the local out shardings and the manual axes.
+      // It is the only carrier of the manual axes when the manual computation
+      // has no operands, since there is no `xla.sdy.GlobalToLocalShape` then.
+      HloSharding hloSharding = parseShardingFromString(shardingAttr);
+      if (auto sdySharding = convertToSdySharding(
+              hloSharding, callOp->getResultTypes(), callOp->getContext())) {
+        callOp->setAttr(kShardingAttr, sdySharding);
+      }
+      if (auto manualAxesAttr =
+              extractManualAxes(hloSharding, callOp->getContext())) {
+        callOp->setAttr(kManualAxes, manualAxesAttr);
       }
     }
 
diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_export.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_export.cc
index 6b15a224661..cb87f6e39b8 100644
--- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_export.cc
+++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_export.cc
@@ -17,7 +17,10 @@ limitations under the License.
 
 #include <memory>
 
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/CommandLine.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/AffineMap.h"
@@ -40,6 +43,7 @@ limitations under the License.
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/TypeID.h"
 #include "mlir/Transforms/DialectConversion.h"
+#include "shardy/dialect/sdy/ir/constants.h"
 #include "shardy/dialect/sdy/ir/dialect.h"
 #include "shardy/dialect/sdy/ir/utils.h"
 #include "stablehlo/dialect/StablehloOps.h"
@@ -59,14 +63,44 @@ using ::mlir::func::FuncOp;
 namespace stablehlo = ::mlir::stablehlo;
 namespace sdy = ::mlir::sdy;
 
+// Returns `shardings` with `manualAxes` stripped from each element, i.e. the
+// local view of tensors that are sharded over those axes.
+//
+// A `NamedSharding` may not mention the same axis twice, so an op whose result
+// is a local tensor carries the manual axes in its `xla.sdy.manual_axes`
+// attribute and the remaining axes in its dimension shardings.
+sdy::TensorShardingPerValueAttr localizeShardings(
+    sdy::TensorShardingPerValueAttr shardings,
+    mlir::ArrayRef<mlir::StringAttr> manualAxes, mlir::MLIRContext* context) {
+  llvm::SmallVector<sdy::TensorShardingAttr> localShardings;
+  localShardings.reserve(shardings.size());
+  for (sdy::TensorShardingAttr sharding : shardings.getShardings()) {
+    localShardings.push_back(sdy::eraseManualAxes(sharding, manualAxes));
+  }
+  return sdy::TensorShardingPerValueAttr::get(context, localShardings);
+}
+
 class SdyRoundTripShardMapExportPass
     : public mlir::PassWrapper<SdyRoundTripShardMapExportPass,
                                mlir::OperationPass<ModuleOp>> {
  public:
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SdyRoundTripShardMapExportPass)
 
-  SdyRoundTripShardMapExportPass(bool enableHloShardingV3)
-      : enableHloShardingV3(enableHloShardingV3) {}
+  SdyRoundTripShardMapExportPass() = default;
+  explicit SdyRoundTripShardMapExportPass(bool enableHloShardingV3) {
+    this->enableHloShardingV3 = enableHloShardingV3;
+  }
+  SdyRoundTripShardMapExportPass(const SdyRoundTripShardMapExportPass& other)
+      : mlir::PassWrapper<SdyRoundTripShardMapExportPass,
+                          mlir::OperationPass<ModuleOp>>(other) {
+    this->enableHloShardingV3 = other.enableHloShardingV3.getValue();
+  }
+
+  Option<bool> enableHloShardingV3{
+      *this, "enable-hlo-sharding-v3",
+      llvm::cl::desc("Whether to enable HloShardingV3 which is the mesh and "
+                     "axis based sharding representation."),
+      llvm::cl::init(false)};
 
   void runOnOperation() final {
     ModuleOp moduleOp = getOperation();
@@ -90,28 +124,81 @@ class SdyRoundTripShardMapExportPass
       rewriter.setInsertionPoint(manualComputation);
       mlir::ValueRange operands = manualComputation->getOperands();
       if (!operands.empty()) {
+        llvm::SmallVector<mlir::Value> globalOperands(operands.begin(),
+                                                      operands.end());
+        sdy::TensorShardingPerValueAttr inShardings =
+            manualComputation.getInShardings();
+        if (enableHloShardingV3) {
+          inShardings = sdy::inlineMesh(symbolTable, inShardings);
+          // `GlobalToLocalShape`'s result is the *local* tensor, so it can only
+          // carry the local in shardings. Materialize the reshard that
+          // `ManualComputationOp` performs implicitly on its operands as an
+          // explicit `@Sharding` custom call. This keeps the exported module
+          // sharding consistent - every op's sharding describes its own result
+          // - and lets the import side read the global in shardings straight
+          // off the operand instead of inferring them.
+          for (auto [globalOperand, inSharding] :
+               llvm::zip_equal(globalOperands, inShardings.getShardings())) {
+            auto shardingOp = stablehlo::CustomCallOp::create(
+                rewriter, loc, globalOperand.getType(), globalOperand);
+            shardingOp.setCallTargetName(kShardingCustomCallTargetName);
+            // Import reads the global in shardings back off these wrappers, so
+            // they have to survive the HLO round trip intact. Without a side
+            // effect a wrapper of a zero-element tensor is replaced by a
+            // constant (see `ZeroSizedHloElimination`, which likewise spares
+            // the side-effecting `xla.sdy.GlobalToLocalShape`), which would
+            // lose that operand's sharding. The side effect also stops CSE from
+            // merging wrappers that carry different shardings.
+            shardingOp.setHasSideEffect(true);
+            shardingOp->setAttr(mlir::sdy::kShardingAttr,
+                                sdy::TensorShardingPerValueAttr::get(
+                                    moduleOp.getContext(), inSharding));
+            globalOperand = shardingOp.getResult(0);
+          }
+        }
         auto globalToLocalShape = stablehlo::CustomCallOp::create(
-            rewriter, loc, manualCompBodyArgTypes, operands);
+            rewriter, loc, manualCompBodyArgTypes, globalOperands);
         globalToLocalShape.setCallTargetName(kGlobalToLocalShapeCallTargetName);
         // We mark `xla.sdy.GlobalToLocalShape` as side-effecting to avoid
         // CSE deduping it with another taking the same operands, as it would
         // ignore the frontend attributes that could be different.
         globalToLocalShape.setHasSideEffect(true);
-        sdy::TensorShardingPerValueAttr inShardings =
-            manualComputation.getInShardings();
         if (enableHloShardingV3) {
-          inShardings = sdy::inlineMesh(symbolTable, inShardings);
+          globalToLocalShape->setAttr(
+              mlir::sdy::kShardingAttr,
+              localizeShardings(inShardings, manualComputation.getManualAxes(),
+                                moduleOp.getContext()));
+          globalToLocalShape->setAttr(kManualAxes,
+                                      manualComputation.getManualAxesAttr());
+        } else {
+          setFrontendAttribute(globalToLocalShape, kInShardings, inShardings);
+          setFrontendAttribute(globalToLocalShape, kManualAxes,
+                               manualComputation.getManualAxesAttr());
         }
-        setFrontendAttribute(globalToLocalShape, kInShardings, inShardings);
-        setFrontendAttribute(globalToLocalShape, kManualAxes,
-                             manualComputation.getManualAxesAttr());
         operands = globalToLocalShape->getResults();
       }
 
+      sdy::TensorShardingPerValueAttr outShardings =
+          manualComputation.getOutShardings();
+      if (enableHloShardingV3) {
+        outShardings = sdy::inlineMesh(symbolTable, outShardings);
+      }
+
       auto callOp =
           CallOp::create(rewriter, loc, localResultTypes, funcName, operands);
       setFrontendAttribute(callOp, kXlaInlineableAttr,
                            rewriter.getStringAttr("xla_late"));
+      if (enableHloShardingV3 && !outShardings.getShardings().empty()) {
+        // The call op produces the local tensors, so it carries the local view
+        // of the out shardings together with the manual axes. This is also the
+        // only carrier of the manual axes when the manual computation has no
+        // operands, i.e. when there is no `xla.sdy.GlobalToLocalShape`.
+        callOp->setAttr(
+            mlir::sdy::kShardingAttr,
+            localizeShardings(outShardings, manualComputation.getManualAxes(),
+                              moduleOp.getContext()));
+        callOp->setAttr(kManualAxes, manualComputation.getManualAxesAttr());
+      }
 
       mlir::ResultRange results = manualComputation->getResults();
       if (!results.empty()) {
@@ -122,14 +209,13 @@ class SdyRoundTripShardMapExportPass
         // We mark `xla.sdy.LocalToGlobalShape` as side-effecting to avoid
         // CSE removing it if it has no users.
         localToGlobalShape.setHasSideEffect(true);
-        sdy::TensorShardingPerValueAttr outShardings =
-            manualComputation.getOutShardings();
         if (enableHloShardingV3) {
-          outShardings = sdy::inlineMesh(symbolTable, outShardings);
+          localToGlobalShape->setAttr(mlir::sdy::kShardingAttr, outShardings);
+        } else {
+          setFrontendAttribute(localToGlobalShape, kOutShardings, outShardings);
+          setFrontendAttribute(localToGlobalShape, kManualAxes,
+                               manualComputation.getManualAxesAttr());
         }
-        setFrontendAttribute(localToGlobalShape, kOutShardings, outShardings);
-        setFrontendAttribute(localToGlobalShape, kManualAxes,
-                             manualComputation.getManualAxesAttr());
         results = localToGlobalShape->getResults();
       }
       sdy::inlineRegionAndConvertTerminatorOp<mlir::func::ReturnOp>(
@@ -147,15 +233,16 @@ class SdyRoundTripShardMapExportPass
            "1. A separate function for the body of the `ManualComputationOp`."
            "2. A `CallOp` calling the function in #1, marked as not inlinable."
            "3. A pair of `CustomCallOp`s that change the shape of the "
-           "   arguments/results. They save the in/out shardings and manual "
-           "   axes as frontend attrs.";
+           "   arguments/results."
+           " Under HloShardingV3 the shardings and manual axes are attached as "
+           "native attributes, and each operand is additionally wrapped in a "
+           "`@Sharding` custom call holding the global in sharding so that "
+           "every op's sharding describes its own result. Otherwise they are "
+           "saved as frontend attrs on the pair of `CustomCallOp`s.";
   }
   void getDependentDialects(mlir::DialectRegistry& registry) const final {
     registry.insert<stablehlo::StablehloDialect>();
   }
-
- private:
-  bool enableHloShardingV3 = false;
 };
 
 }  // namespace
diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_import.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_import.cc
index 2e2151e8d11..29ea30dd68e 100644
--- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_import.cc
+++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/shard_map_import.cc
@@ -75,6 +75,50 @@ using ::mlir::stablehlo::CustomCallOp;
 
 namespace sdy = ::mlir::sdy;
 
+// Recovers the global in shardings of a manual computation from the
+// `@Sharding` custom calls that export wraps around each operand of
+// `globalToLocalShape`, and removes those wrappers.
+//
+// Export materializes the reshard that `sdy.manual_computation` performs
+// implicitly on entry as an explicit `@Sharding` op, so that every op's
+// sharding describes its own result. Undoing it here makes the round trip an
+// identity.
+//
+// Returns an empty vector, and leaves the IR untouched, if any operand is not
+// wrapped - the shardings are only meaningful as a complete set.
+SmallVector<sdy::TensorShardingAttr> takeGlobalInShardings(
+    CustomCallOp globalToLocalShape, mlir::IRRewriter& rewriter) {
+  SmallVector<sdy::TensorShardingAttr> globalInShardings;
+  SmallVector<CustomCallOp> shardingOps;
+  for (mlir::Value globalOperand : globalToLocalShape.getOperands()) {
+    auto shardingOp = globalOperand.getDefiningOp<CustomCallOp>();
+    if (!shardingOp ||
+        shardingOp.getCallTargetName() != kShardingCustomCallTargetName) {
+      return {};
+    }
+    sdy::TensorShardingAttr sharding =
+        sdy::getSharding(shardingOp.getResult(0));
+    if (!sharding) {
+      return {};
+    }
+    globalInShardings.push_back(sharding);
+    shardingOps.push_back(shardingOp);
+  }
+
+  for (auto [i, shardingOp] : llvm::enumerate(shardingOps)) {
+    globalToLocalShape->setOperand(i, shardingOp.getOperand(0));
+  }
+  // The same wrapper can feed several operands, so only visit it once - after
+  // the first erase the handle is dangling.
+  llvm::SmallDenseSet<Operation*> seen;
+  for (CustomCallOp shardingOp : shardingOps) {
+    if (seen.insert(shardingOp).second && shardingOp.use_empty()) {
+      rewriter.eraseOp(shardingOp);
+    }
+  }
+  return globalInShardings;
+}
+
 mlir::LogicalResult rewriteManualComputation(
     CallOp callOp, mlir::IRRewriter& rewriter,
     const mlir::SymbolTable& symbolTable) {
@@ -96,7 +140,7 @@ mlir::LogicalResult rewriteManualComputation(
   // `ManualComputationOp` differently depending on whether the original had
   // operands/results.
   CustomCallOp globalToLocalShape;
-  mlir::ValueRange operands = callOp.getOperands();
+  SmallVector<mlir::Value> operands(callOp.getOperands());
   if (!operands.empty()) {
     // An input to `sdy.manual_computation` can have a dimension of size 0
     // (i.e. 0 num-elements), in which case, the corresponding result of
@@ -114,7 +158,8 @@ mlir::LogicalResult rewriteManualComputation(
     globalToLocalShape = (*customCallResIt).getDefiningOp<CustomCallOp>();
     CHECK_EQ(globalToLocalShape.getCallTargetName(),
              kGlobalToLocalShapeCallTargetName);
-    operands = globalToLocalShape->getOperands();
+    operands.assign(globalToLocalShape->operand_begin(),
+                    globalToLocalShape->operand_end());
   }
 
   mlir::TypeRange resultTypes = callOp->getResultTypes();
@@ -150,10 +195,13 @@ mlir::LogicalResult rewriteManualComputation(
     if (!customCallOp) {
       return;
     }
+
     if (mlir::DictionaryAttr frontendAttrs = getFrontendAttrs(customCallOp)) {
-      shardings = parseStringAttr<sdy::TensorShardingPerValueAttr>(
-          frontendAttrs, shardingAttrName);
-      if (manualAxes.empty()) {
+      if (hasKey(frontendAttrs, shardingAttrName)) {
+        shardings = parseStringAttr<sdy::TensorShardingPerValueAttr>(
+            frontendAttrs, shardingAttrName);
+      }
+      if (manualAxes.empty() && hasKey(frontendAttrs, kManualAxes)) {
         manualAxes =
             parseStringAttr<sdy::ManualAxesAttr>(frontendAttrs, kManualAxes);
       }
@@ -162,6 +210,45 @@ mlir::LogicalResult rewriteManualComputation(
 
   setShardingAttrs(globalToLocalShape, inShardings, kInShardings);
   setShardingAttrs(localToGlobalShape, outShardings, kOutShardings);
+
+  // Under HloShardingV3 the shardings and manual axes round trip as native
+  // attributes rather than frontend attributes, so read them directly.
+  if (outShardings.empty() && localToGlobalShape) {
+    // `LocalToGlobalShape` produces the global tensors, so its own sharding is
+    // the global out shardings.
+    if (auto sharding =
+            localToGlobalShape->getAttrOfType<sdy::TensorShardingPerValueAttr>(
+                sdy::kShardingAttr)) {
+      outShardings = sharding;
+    }
+  }
+
+  if (manualAxes.empty()) {
+    // A manual computation without operands has no `GlobalToLocalShape`, in
+    // which case the body call op is the only carrier of the manual axes.
+    for (Operation* op :
+         {globalToLocalShape.getOperation(), callOp.getOperation()}) {
+      auto axes =
+          op ? op->getAttrOfType<sdy::ManualAxesAttr>(kManualAxes) : nullptr;
+      if (axes && !axes.empty()) {
+        manualAxes = axes;
+        break;
+      }
+    }
+  }
+
+  if (inShardings.empty() && globalToLocalShape) {
+    SmallVector<sdy::TensorShardingAttr> globalInShardings =
+        takeGlobalInShardings(globalToLocalShape, rewriter);
+    if (!globalInShardings.empty()) {
+      inShardings =
+          sdy::TensorShardingPerValueAttr::get(context, globalInShardings);
+      // `takeGlobalInShardings` rewired `globalToLocalShape` past the wrappers.
+      operands.assign(globalToLocalShape->operand_begin(),
+                      globalToLocalShape->operand_end());
+    }
+  }
+
   auto manualComputationOp =
       rewriter.replaceOpWithNewOp<sdy::ManualComputationOp>(
           callOp, resultTypes, operands, inShardings, outShardings, manualAxes);
@@ -276,10 +363,12 @@ class SdyRoundTripShardMapImportPass
   }
 
   StringRef getDescription() const override {
-    return "converts a CallOp calling a @xla.sdy.manual_computation_body func "
-           "with in/out shardings and manual axes as frontend attrs, wrapped "
-           "with a pair of `CustomCallOps` that change the shape of the "
-           "arguments/results, to a ManualComputationOp";
+    return "converts a CallOp calling a @xla.sdy.manual_computation_body func, "
+           "wrapped with a pair of `CustomCallOps` that change the shape of "
+           "the arguments/results, to a ManualComputationOp. The in/out "
+           "shardings and manual axes are read from frontend attrs, or, under "
+           "HloShardingV3, from native attributes on the boundary ops and on "
+           "the `@Sharding` custom calls wrapping the operands";
   }
   void getDependentDialects(mlir::DialectRegistry& registry) const final {
     registry.insert<sdy::SdyDialect>();
diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc
index b0deb0e6256..df36194d8a0 100644
--- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc
+++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc
@@ -137,6 +137,11 @@ void exportFunc(FuncOp funcOp, const SymbolTable& symbolTable,
       if (ManualAxesAttr manualAxesAttr =
               funcOp.getArgAttrOfType<ManualAxesAttr>(argNum, kManualAxes)) {
         manualAxes = manualAxesAttr.getValue();
+        if (mlir::isa<stablehlo::TokenType>(
+                funcOp.getArgument(argNum).getType()) &&
+            getMeshAttr(sdySharding).getAxes().size() != manualAxes.size()) {
+          manualAxes = {};
+        }
         attrs.erase(kManualAxes);
       }
       attrs.set(kXlaShardingAttr,
@@ -162,6 +167,10 @@ void exportFunc(FuncOp funcOp, const SymbolTable& symbolTable,
       if (auto manualAxesAttr =
               mlir::dyn_cast_or_null<ManualAxesAttr>(attrs.get(kManualAxes))) {
         manualAxes = manualAxesAttr.getValue();
+        if (mlir::isa<stablehlo::TokenType>(funcOp.getResultTypes()[resNum]) &&
+            getMeshAttr(sdySharding).getAxes().size() != manualAxes.size()) {
+          manualAxes = {};
+        }
         attrs.erase(kManualAxes);
       }
       attrs.set(kXlaShardingAttr,
@@ -375,24 +384,38 @@ HloSharding getHloShardingForOp(
     std::function<MeshAttr(TensorShardingAttr)> getMeshAttr,
     ArrayRef<StringAttr> manualAxes, bool enableHloShardingV3,
     bool simplifyReplicatedShardings) {
+  auto getManualAxesForType = [&](mlir::Type type,
+                                  TensorShardingAttr sdySharding) {
+    if (mlir::isa<stablehlo::TokenType>(type) &&
+        getMeshAttr(sdySharding).getAxes().size() != manualAxes.size()) {
+      return ArrayRef<StringAttr>();
+    }
+    return manualAxes;
+  };
+
   bool isNoResultMaximal = op->getNumResults() == 0 && shardings.size() == 1 &&
                            (getMeshAttr(shardings.front()).isMaximal() ||
                             shardings.front().isFullyReplicated());
   CHECK(shardings.size() == op->getNumResults() || isNoResultMaximal);
   if (op->getNumResults() == 1 || isNoResultMaximal) {
-    return convertToHloSharding(shardings.front(), getMeshAttr, manualAxes,
-                                enableHloShardingV3,
+    ArrayRef<StringAttr> resultManualAxes =
+        op->getNumResults() == 1
+            ? getManualAxesForType(op->getResultTypes().front(),
+                                   shardings.front())
+            : manualAxes;
+    return convertToHloSharding(shardings.front(), getMeshAttr,
+                                resultManualAxes, enableHloShardingV3,
                                 simplifyReplicatedShardings);
   }
 
   std::vector<HloSharding> newShardings;
   newShardings.reserve(shardings.size());
-  llvm::transform(shardings, std::back_inserter(newShardings),
-                  [&](TensorShardingAttr sdySharding) {
-                    return convertToHloSharding(sdySharding, getMeshAttr,
-                                                manualAxes, enableHloShardingV3,
-                                                simplifyReplicatedShardings);
-                  });
+  for (auto [type, sdySharding] :
+       llvm::zip_equal(op->getResultTypes(), shardings)) {
+    newShardings.push_back(convertToHloSharding(
+        sdySharding, getMeshAttr, getManualAxesForType(type, sdySharding),
+        enableHloShardingV3, simplifyReplicatedShardings));
+  }
 
   std::vector<xla::Shape> shapes;
   shapes.reserve(op->getNumResults());
diff --git a/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline_manual_computation.mlir b/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline_manual_computation.mlir
index 9a05f216ae8..597d749388f 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline_manual_computation.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline_manual_computation.mlir
@@ -154,3 +154,42 @@ func.func @main(
   } : (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>)
   return %0#0, %0#1 : !stablehlo.token, tensor<2xi64>
 }
+
+// -----
+
+// ***************** Zero-element operand test *****************
+
+// A zero-element operand is replaced by a constant during the HLO round trip,
+// which drops any non side-effecting op producing it. Make sure the remaining
+// operands keep their in shardings rather than the manual computation coming
+// back with none at all.
+
+// Make sure this temp attr doesn't exist anymore.
+// CHECK-NOT: sharding_hlo_string
+
+// CHECK-V2: sdy.mesh @mesh_1 = <["a"=4]>
+// CHECK-V3: sdy.mesh @mesh = <["a"=4]>
+sdy.mesh @mesh_1 = <["a"=4]>
+
+// CHECK-LABEL: func.func @main
+func.func @main(%arg0: tensor<4x8xi32>) -> tensor<4x8xi32> {
+  // CHECK-NEXT:     %[[MANUAL_COMP:.*]] = sdy.manual_computation(%arg0)
+  // CHECK-V2-SAME{LITERAL}: in_shardings=[<@mesh_1, [{"a"}, {}]>] out_shardings=[<@mesh_1, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<1x8xi32>) {
+  // CHECK-V3-SAME{LITERAL}: in_shardings=[<@mesh, [{"a"}, {}]>] out_shardings=[<@mesh, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<1x8xi32>) {
+  // CHECK-NEXT:       sdy.return %arg1 : tensor<1x8xi32>
+  // CHECK-NEXT:     } : (tensor<4x8xi32>) -> tensor<4x8xi32>
+  // CHECK-NEXT:     return %[[MANUAL_COMP]] : tensor<4x8xi32>
+  %c = stablehlo.constant dense<> : tensor<0xi32>
+  %0 = sdy.manual_computation(%c, %arg0)
+      in_shardings=[<@mesh_1, [{}], replicated={"a"}>, <@mesh_1, [{"a"}, {}]>]
+      out_shardings=[<@mesh_1, [{"a"}, {}]>]
+      manual_axes={"a"} (%arg1: tensor<0xi32>, %arg2: tensor<1x8xi32>) {
+    %init = stablehlo.constant dense<0> : tensor<i32>
+    %red = stablehlo.reduce(%arg1 init: %init) applies stablehlo.add across dimensions = [0] : (tensor<0xi32>, tensor<i32>) -> tensor<i32>
+    %bcast = stablehlo.broadcast_in_dim %red, dims = [] : (tensor<i32>) -> tensor<1x8xi32>
+    %add = stablehlo.add %arg2, %bcast : tensor<1x8xi32>
+    sdy.return %add : tensor<1x8xi32>
+  } : (tensor<0xi32>, tensor<4x8xi32>) -> tensor<4x8xi32>
+  return %0 : tensor<4x8xi32>
+}
+
diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir
index 58b7012ab5b..1b0f94f6667 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir
@@ -204,6 +204,30 @@ module @module_1 {
     return %3 : tensor<1xi32>
   }
 
+  // Tests importing shard map custom calls annotated with HloShardingV3 mhlo.sharding.
+  // CHECK-LABEL: func @manual_computation_hlo_sharding_v3
+  // CHECK-SAME:     (%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
+  // CHECK-NEXT:    %[[MAN_COMP:.*]] = sdy.manual_computation(%arg0)
+  // CHECK-SAME{LITERAL}: in_shardings=[<@mesh_2, [{"a"}, {}]>] out_shardings=[<@mesh_2, [{"a"}, {}]>] manual_axes={"a"}
+  // CHECK-SAME:        (%arg1: tensor<2x16xf32>) {
+  // CHECK-NEXT:      %[[ADD:.*]] = stablehlo.add %arg1, %arg1 : tensor<2x16xf32>
+  // CHECK-NEXT:      sdy.return %[[ADD]] : tensor<2x16xf32>
+  // CHECK-NEXT:    } : (tensor<8x16xf32>) -> tensor<8x16xf32>
+  // CHECK-NEXT:    return %[[MAN_COMP]] : tensor<8x16xf32>
+  // CHECK-NEXT:  }
+  func.func @manual_computation_hlo_sharding_v3(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
+    %s0 = stablehlo.custom_call @Sharding(%arg0) {mhlo.sharding = "{mesh['a'=4,'b'=2], [{'a'}, {}]}"} : (tensor<8x16xf32>) -> tensor<8x16xf32>
+    %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, mhlo.sharding = "{mesh['a'=4,'b'=2], [{}, {}], manual={'a'}}"} : (tensor<8x16xf32>) -> tensor<2x16xf32>
+    %1 = call @xla.sdy.manual_computation_body_hlo_v3(%0) : (tensor<2x16xf32>) -> tensor<2x16xf32>
+    %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, mhlo.sharding = "{mesh['a'=4,'b'=2], [{'a'}, {}]}"} : (tensor<2x16xf32>) -> tensor<8x16xf32>
+    return %2 : tensor<8x16xf32>
+  }
+
+  func.func private @xla.sdy.manual_computation_body_hlo_v3(%arg0: tensor<2x16xf32>) -> tensor<2x16xf32> {
+    %0 = stablehlo.add %arg0, %arg0 : tensor<2x16xf32>
+    return %0 : tensor<2x16xf32>
+  }
+
   // CHECK-LABEL: func @frontend_attr_not_sharding
   // CHECK-SAME:    %arg0: tensor<16x8xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a", "b"}, {?}]>},
   // CHECK-SAME:    %arg1: tensor<16x8xf32> {mhlo.frontend_attributes = {baz = 1 : i32, foo = "bar"}},
diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_export_hlo_sharding_v3.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_export_hlo_sharding_v3.mlir
new file mode 100644
index 00000000000..fc83e872e68
--- /dev/null
+++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_export_hlo_sharding_v3.mlir
@@ -0,0 +1,285 @@
+// Copyright 2026 The OpenXLA 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.
+// ==============================================================================
+// RUN: sdy_opt %s -xla-sdy-round-trip-shard-map-export='enable-hlo-sharding-v3=true' 2>&1 | FileCheck %s
+
+// Tests exporting `sdy.manual_computation` under HloShardingV3 (`enable-hlo-sharding-v3=true`).
+// Verifies that:
+// 1. Every op carries a native `sdy.sharding` that describes its own result:
+//    - a `@Sharding` custom call per operand holds the global in-shardings,
+//      making the reshard `sdy.manual_computation` performs implicitly on entry
+//      explicit. It is side-effecting so that the HLO round trip cannot drop it
+//      (see `ZeroSizedHloElimination` for a zero-element operand) or CSE two
+//      wrappers carrying different shardings;
+//    - `GlobalToLocalShape` and the body call op hold the local shardings plus
+//      typed `xla.sdy.manual_axes`;
+//    - `LocalToGlobalShape` holds the global out-shardings.
+// 2. Legacy string `mhlo.frontend_attributes` (xla.sdy.in_shardings, etc.) are NOT emitted.
+
+sdy.mesh @mesh_0 = <["a"=4, "b"=2]>
+sdy.mesh @mesh_1 = <["a"=2, "b"=2, "c"=2, "d"=2]>
+
+// Ensure string frontend attributes are completely absent in V3 export:
+// CHECK-NOT: xla.sdy.in_shardings
+// CHECK-NOT: xla.sdy.out_shardings
+
+// Tests basic export with multiple inputs, outputs, and all-reduce.
+// CHECK-LABEL: func @single_manual_comp
+func.func @single_manual_comp(%arg0: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a", ?}, {"b", ?}]>}, %arg1: tensor<16x32xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"b", ?}, {?}]>}) -> (tensor<8x32xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a"}, {}]>}) {
+  // CHECK-NEXT: %[[SHARDING_0:.*]] = stablehlo.custom_call @Sharding(%arg0)
+  // CHECK-SAME:   has_side_effect = true
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, [{"a"}, {"b"}]>]>
+  // CHECK-SAME:   (tensor<8x16xf32>) -> tensor<8x16xf32>
+  // CHECK-NEXT: %[[SHARDING_1:.*]] = stablehlo.custom_call @Sharding(%arg1)
+  // CHECK-SAME:   has_side_effect = true
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, [{"b"}, {}], replicated={"a"}>]>
+  // CHECK-SAME:   (tensor<16x32xf32>) -> tensor<16x32xf32>
+  // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]]:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING_0]], %[[SHARDING_1]])
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, [{}, {}]>, <mesh<["a"=4, "b"=2]>, [{}, {}]>]>
+  // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy<manual_axes{"a", "b"}>
+  // CHECK-SAME:   (tensor<8x16xf32>, tensor<16x32xf32>) -> (tensor<2x8xf32>, tensor<8x32xf32>)
+  // CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body(%[[GLOBAL_TO_LOCAL]]#0, %[[GLOBAL_TO_LOCAL]]#1)
+  // CHECK-SAME:   {mhlo.frontend_attributes = {inlineable = "xla_late"}
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, [{}, {}]>]>
+  // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy<manual_axes{"a", "b"}>
+  // CHECK-SAME:   (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32>
+  // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]])
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, [{"a"}, {}], replicated={"b"}>]>
+  // CHECK-SAME:   (tensor<2x32xf32>) -> tensor<8x32xf32>
+  // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<8x32xf32>
+  %0 = sdy.manual_computation(%arg0, %arg1) in_shardings=[<@mesh_0, [{"a"}, {"b"}]>, <@mesh_0, [{"b"}, {}], replicated={"a"}>] out_shardings=[<@mesh_0, [{"a"}, {}], replicated={"b"}>] manual_axes={"a", "b"} (%arg2: tensor<2x8xf32>, %arg3: tensor<8x32xf32>) {
+    %1 = stablehlo.add %arg2, %arg2 : tensor<2x8xf32>
+    %2 = stablehlo.dot %1, %arg3 : (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32>
+    %3 = "stablehlo.all_reduce"(%2) <{channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>, replica_groups = dense<[[0, 1], [2, 3], [4, 5], [6, 7]]> : tensor<4x2xi64>, use_global_device_ids}> ({
+    ^bb0(%arg4: tensor<f32>, %arg5: tensor<f32>):
+      %4 = stablehlo.add %arg4, %arg5 : tensor<f32>
+      stablehlo.return %4 : tensor<f32>
+    }) : (tensor<2x32xf32>) -> tensor<2x32xf32>
+    sdy.return %3 : tensor<2x32xf32>
+  } : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32>
+  return %0 : tensor<8x32xf32>
+}
+
+// Tests chained manual computations where the output of one feeds the input of
+// another with a *different* sharding. The `@Sharding` custom call on the
+// second boundary is what preserves the reshard between them; without it the
+// `in_shardings` of the second manual computation would be lost.
+// CHECK-LABEL: func @manual_comp_using_another
+func.func @manual_comp_using_another(%arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a"}, {}]>})
+    -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{}, {"b"}]>}) {
+  // CHECK-NEXT: %[[SHARDING_0:.*]] = stablehlo.custom_call @Sharding(%arg0)
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, [{"a"}, {}]>]>
+  // CHECK-NEXT: %[[GLOBAL_TO_LOCAL_0:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING_0]])
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, [{}, {}]>]>
+  // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy<manual_axes{"a"}>
+  // CHECK-SAME:   (tensor<8x8xf32>) -> tensor<2x8xf32>
+  // CHECK-NEXT: %[[SHMAP_0:.*]] = call @xla.sdy.manual_computation_body_0(%[[GLOBAL_TO_LOCAL_0]])
+  // CHECK-SAME:   {mhlo.frontend_attributes = {inlineable = "xla_late"}
+  // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy<manual_axes{"a"}>
+  // CHECK-SAME:   (tensor<2x8xf32>) -> tensor<2x8xf32>
+  // CHECK-NEXT: %[[LOCAL_TO_GLOBAL_0:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP_0]])
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, [{"a"}, {}]>]>
+  // CHECK-SAME:   (tensor<2x8xf32>) -> tensor<8x8xf32>
+  // CHECK-NEXT: %[[SHARDING_1:.*]] = stablehlo.custom_call @Sharding(%[[LOCAL_TO_GLOBAL_0]])
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, [{}, {"b"}]>]>
+  // CHECK-NEXT: %[[GLOBAL_TO_LOCAL_1:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING_1]])
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, [{}, {}]>]>
+  // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy<manual_axes{"b"}>
+  // CHECK-SAME:   (tensor<8x8xf32>) -> tensor<8x4xf32>
+  // CHECK-NEXT: %[[SHMAP_1:.*]] = call @xla.sdy.manual_computation_body_1(%[[GLOBAL_TO_LOCAL_1]])
+  // CHECK-SAME:   {mhlo.frontend_attributes = {inlineable = "xla_late"}
+  // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy<manual_axes{"b"}>
+  // CHECK-SAME:   (tensor<8x4xf32>) -> tensor<8x4xf32>
+  // CHECK-NEXT: %[[LOCAL_TO_GLOBAL_1:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP_1]])
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, [{}, {"b"}]>]>
+  // CHECK-SAME:   (tensor<8x4xf32>) -> tensor<8x8xf32>
+  // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL_1]] : tensor<8x8xf32>
+  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh_0, [{"a"}, {}]>] out_shardings=[<@mesh_0, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<2x8xf32>) {
+    sdy.return %arg1 : tensor<2x8xf32>
+  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
+
+  %1 = sdy.manual_computation(%0) in_shardings=[<@mesh_0, [{}, {"b"}]>] out_shardings=[<@mesh_0, [{}, {"b"}]>] manual_axes={"b"} (%arg1: tensor<8x4xf32>) {
+    sdy.return %arg1 : tensor<8x4xf32>
+  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
+  return %1 : tensor<8x8xf32>
+}
+
+// Tests nested manual computations: an inner shard_map inside an outer shard_map body.
+// CHECK-LABEL: func @nested_shmaps
+func.func @nested_shmaps(%arg0: tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1, [{"a"}, {"b"}]>}) -> (tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1, [{"a", ?}, {?}]>}) {
+  // CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0)
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=2, "b"=2, "c"=2, "d"=2]>, [{"a"}, {}]>]>
+  // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]])
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=2, "b"=2, "c"=2, "d"=2]>, [{}, {}]>]>
+  // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy<manual_axes{"a"}>
+  // CHECK-SAME:   (tensor<4x8xf32>) -> tensor<2x8xf32>
+  // CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_3(%[[GLOBAL_TO_LOCAL]])
+  // CHECK-SAME:   {mhlo.frontend_attributes = {inlineable = "xla_late"}
+  // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy<manual_axes{"a"}>
+  // CHECK-SAME:   (tensor<2x8xf32>) -> tensor<2x8xf32>
+  // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]])
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=2, "b"=2, "c"=2, "d"=2]>, [{"a"}, {}]>]>
+  // CHECK-SAME:   (tensor<2x8xf32>) -> tensor<4x8xf32>
+  // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<4x8xf32>
+  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh_1, [{"a"}, {}]>] out_shardings=[<@mesh_1, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<2x8xf32>) {
+    %1 = sdy.manual_computation(%arg1) in_shardings=[<@mesh_1, [{}, {"b"}]>] out_shardings=[<@mesh_1, [{}, {"b"}]>] manual_axes={"b"} (%arg2: tensor<2x4xf32>) {
+      %2 = stablehlo.multiply %arg2, %arg2 : tensor<2x4xf32>
+      sdy.return %2 : tensor<2x4xf32>
+    } : (tensor<2x8xf32>) -> tensor<2x8xf32>
+    sdy.return %1 : tensor<2x8xf32>
+  } : (tensor<4x8xf32>) -> tensor<4x8xf32>
+  return %0 : tensor<4x8xf32>
+}
+
+// Tests nested manual computation where the outer body performs additional math on the inner result.
+// CHECK-LABEL: func @nested_shmaps_extra_op
+func.func @nested_shmaps_extra_op(%arg0: tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1, [{"a"}, {"b"}]>}) -> (tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1, [{"a", ?}, {?}]>}) {
+  // CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0)
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=2, "b"=2, "c"=2, "d"=2]>, [{"a"}, {}]>]>
+  // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]])
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=2, "b"=2, "c"=2, "d"=2]>, [{}, {}]>]>
+  // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy<manual_axes{"a"}>
+  // CHECK-SAME:   (tensor<4x8xf32>) -> tensor<2x8xf32>
+  // CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_5(%[[GLOBAL_TO_LOCAL]])
+  // CHECK-SAME:   {mhlo.frontend_attributes = {inlineable = "xla_late"}
+  // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy<manual_axes{"a"}>
+  // CHECK-SAME:   (tensor<2x8xf32>) -> tensor<2x8xf32>
+  // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]])
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=2, "b"=2, "c"=2, "d"=2]>, [{"a"}, {}]>]>
+  // CHECK-SAME:   (tensor<2x8xf32>) -> tensor<4x8xf32>
+  // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<4x8xf32>
+  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh_1, [{"a"}, {}]>] out_shardings=[<@mesh_1, [{"a"}, {}]>] manual_axes={"a"} (%arg1: tensor<2x8xf32>) {
+    %1 = sdy.manual_computation(%arg1) in_shardings=[<@mesh_1, [{}, {"b"}]>] out_shardings=[<@mesh_1, [{}, {"b"}]>] manual_axes={"b"} (%arg2: tensor<2x4xf32>) {
+      %2 = stablehlo.multiply %arg2, %arg2 : tensor<2x4xf32>
+      sdy.return %2 : tensor<2x4xf32>
+    } : (tensor<2x8xf32>) -> tensor<2x8xf32>
+    %3 = stablehlo.add %1, %1 : tensor<2x8xf32>
+    sdy.return %3 : tensor<2x8xf32>
+  } : (tensor<4x8xf32>) -> tensor<4x8xf32>
+  return %0 : tensor<4x8xf32>
+}
+
+// Edge case: Manual computation producing an output but taking no inputs. There
+// is no `GlobalToLocalShape`, so the body call op is the only carrier of the
+// manual axes.
+// CHECK-LABEL: func @manual_computation_no_inputs
+func.func @manual_computation_no_inputs() -> tensor<4xi64> {
+  // CHECK-NEXT: %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_6()
+  // CHECK-SAME:   {mhlo.frontend_attributes = {inlineable = "xla_late"}
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, [{}]>]>
+  // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy<manual_axes{"b"}>
+  // CHECK-SAME:   () -> tensor<2xi64>
+  // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]])
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, [{"b"}]>]>
+  // CHECK-SAME:   (tensor<2xi64>) -> tensor<4xi64>
+  // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]] : tensor<4xi64>
+  %0 = sdy.manual_computation() in_shardings=[] out_shardings=[<@mesh_0, [{"b"}]>] manual_axes={"b"} () {
+    %1 = stablehlo.constant dense<[2, 3]> : tensor<2xi64>
+    sdy.return %1 : tensor<2xi64>
+  } : () -> tensor<4xi64>
+  func.return %0 : tensor<4xi64>
+}
+
+// Edge case: Manual computation consuming an input but producing no outputs.
+// The body call op has no results, so `GlobalToLocalShape` is the only carrier
+// of the manual axes.
+// CHECK-LABEL: func @manual_computation_no_outputs
+func.func @manual_computation_no_outputs(%arg0: tensor<4xi64>) {
+  // CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0)
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, [{"b"}]>]>
+  // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]])
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, [{}]>]>
+  // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy<manual_axes{"b"}>
+  // CHECK-SAME:   (tensor<4xi64>) -> tensor<2xi64>
+  // CHECK-NEXT: call @xla.sdy.manual_computation_body_7(%[[GLOBAL_TO_LOCAL]])
+  // CHECK-SAME:   {mhlo.frontend_attributes = {inlineable = "xla_late"}}
+  // CHECK-SAME:   (tensor<2xi64>) -> ()
+  // CHECK-NEXT: return
+  sdy.manual_computation(%arg0) in_shardings=[<@mesh_0, [{"b"}]>] out_shardings=[] manual_axes={"b"} (%arg1: tensor<2xi64>) {
+    stablehlo.custom_call @sdy_testonly(%arg1) : (tensor<2xi64>) -> ()
+    sdy.return
+  } : (tensor<4xi64>) -> ()
+  func.return
+}
+
+// Edge case: Degenerate manual computation with no inputs and no outputs.
+// CHECK-LABEL: func @manual_computation_no_inputs_no_outputs
+func.func @manual_computation_no_inputs_no_outputs() {
+  // CHECK-NEXT: call @xla.sdy.manual_computation_body_8() {mhlo.frontend_attributes = {inlineable = "xla_late"}} : () -> ()
+  sdy.manual_computation() in_shardings=[] out_shardings=[] manual_axes={} () {
+    sdy.return
+  } : () -> ()
+  func.return
+}
+
+// Type preservation: Token arguments should not have manual axes attached in V3.
+// CHECK-LABEL: func @manual_computation_tokens
+func.func @manual_computation_tokens(%arg0: !stablehlo.token, %arg1: tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>) {
+  // CHECK-NEXT: %[[SHARDING_0:.*]] = stablehlo.custom_call @Sharding(%arg0)
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, []>]>
+  // CHECK-SAME:   (!stablehlo.token) -> !stablehlo.token
+  // CHECK-NEXT: %[[SHARDING_1:.*]] = stablehlo.custom_call @Sharding(%arg1)
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, [{"b"}]>]>
+  // CHECK-SAME:   (tensor<2xi64>) -> tensor<2xi64>
+  // CHECK-NEXT: %[[GLOBAL_TO_LOCAL:.*]]:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING_0]], %[[SHARDING_1]])
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, []>, <mesh<["a"=4, "b"=2]>, [{}]>]>
+  // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy<manual_axes{"b"}>
+  // CHECK-SAME:   (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<1xi64>)
+  // CHECK-NEXT: %[[SHMAP:.*]]:2 = call @xla.sdy.manual_computation_body_9(%[[GLOBAL_TO_LOCAL]]#0, %[[GLOBAL_TO_LOCAL]]#1)
+  // CHECK-SAME:   {mhlo.frontend_attributes = {inlineable = "xla_late"}
+  // CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy<manual_axes{"b"}>
+  // CHECK-SAME:   (!stablehlo.token, tensor<1xi64>) -> (!stablehlo.token, tensor<1xi64>)
+  // CHECK-NEXT: %[[LOCAL_TO_GLOBAL:.*]]:2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]]#0, %[[SHMAP]]#1)
+  // CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=4, "b"=2]>, []>, <mesh<["a"=4, "b"=2]>, [{"b"}]>]>
+  // CHECK-SAME:   (!stablehlo.token, tensor<1xi64>) -> (!stablehlo.token, tensor<2xi64>)
+  // CHECK-NEXT: return %[[LOCAL_TO_GLOBAL]]#0, %[[LOCAL_TO_GLOBAL]]#1 : !stablehlo.token, tensor<2xi64>
+  %0:2 = sdy.manual_computation(%arg0, %arg1) in_shardings=[<@mesh_0, []>, <@mesh_0, [{"b"}]>] out_shardings=[<@mesh_0, []>, <@mesh_0, [{"b"}]>] manual_axes={"b"} (%arg2: !stablehlo.token, %arg3: tensor<1xi64>) {
+    sdy.return %arg2, %arg3 : !stablehlo.token, tensor<1xi64>
+  } : (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>)
+  func.return %0#0, %0#1 : !stablehlo.token, tensor<2xi64>
+}
+
+
+// CHECK-LABEL: func @xla.sdy.manual_computation_body_3(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32>
+// CHECK-NEXT:   %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0)
+// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=2, "b"=2, "c"=2, "d"=2]>, [{}, {"b"}]>]>
+// CHECK-NEXT:   %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]])
+// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=2, "b"=2, "c"=2, "d"=2]>, [{}, {}]>]>
+// CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy<manual_axes{"b"}>
+// CHECK-SAME:     (tensor<2x8xf32>) -> tensor<2x4xf32>
+// CHECK-NEXT:   %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_2(%[[GLOBAL_TO_LOCAL]])
+// CHECK-SAME:     {mhlo.frontend_attributes = {inlineable = "xla_late"}
+// CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy<manual_axes{"b"}>
+// CHECK-SAME:     (tensor<2x4xf32>) -> tensor<2x4xf32>
+// CHECK-NEXT:   %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]])
+// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=2, "b"=2, "c"=2, "d"=2]>, [{}, {"b"}]>]>
+// CHECK-SAME:     (tensor<2x4xf32>) -> tensor<2x8xf32>
+// CHECK-NEXT:   return %[[LOCAL_TO_GLOBAL]] : tensor<2x8xf32>
+
+// CHECK-LABEL: func @xla.sdy.manual_computation_body_5(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32>
+// CHECK-NEXT:   %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0)
+// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=2, "b"=2, "c"=2, "d"=2]>, [{}, {"b"}]>]>
+// CHECK-NEXT:   %[[GLOBAL_TO_LOCAL:.*]] = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%[[SHARDING]])
+// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=2, "b"=2, "c"=2, "d"=2]>, [{}, {}]>]>
+// CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy<manual_axes{"b"}>
+// CHECK-SAME:     (tensor<2x8xf32>) -> tensor<2x4xf32>
+// CHECK-NEXT:   %[[SHMAP:.*]] = call @xla.sdy.manual_computation_body_4(%[[GLOBAL_TO_LOCAL]])
+// CHECK-SAME:     {mhlo.frontend_attributes = {inlineable = "xla_late"}
+// CHECK-SAME{LITERAL}: xla.sdy.manual_axes = #sdy<manual_axes{"b"}>
+// CHECK-SAME:     (tensor<2x4xf32>) -> tensor<2x4xf32>
+// CHECK-NEXT:   %[[LOCAL_TO_GLOBAL:.*]] = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%[[SHMAP]])
+// CHECK-SAME{LITERAL}: sdy.sharding = #sdy.sharding_per_value<[<mesh<["a"=2, "b"=2, "c"=2, "d"=2]>, [{}, {"b"}]>]>
+// CHECK-SAME:     (tensor<2x4xf32>) -> tensor<2x8xf32>
+// CHECK-NEXT:   %[[ADD:.*]] = stablehlo.add %[[LOCAL_TO_GLOBAL]], %[[LOCAL_TO_GLOBAL]] : tensor<2x8xf32>
+// CHECK-NEXT:   return %[[ADD]] : tensor<2x8xf32>
diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_failure.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_failure.mlir
index 47d1ba107a8..15b7eec24bd 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_failure.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_failure.mlir
@@ -45,3 +45,38 @@ func.func @xla.sdy.manual_computation_body(%arg0: tensor<0x8xf32>) -> tensor<0x8
   return %arg0 : tensor<0x8xf32>
 }
 
+// -----
+
+sdy.mesh @mesh = <["a"=2, "b"=2]>
+
+// Tests import error when manual computation body call operand is not produced by GlobalToLocalShape.
+func.func @manual_computation_missing_global_to_local_shape_direct_attrs(%arg0: tensor<0x16xf32>) -> (tensor<0x16xf32>) {
+  %c = stablehlo.constant dense<0.000000e+00> : tensor<0x8xf32>
+  // expected-error @+2 {{'func.call' op expected at least one operand of @xla.sdy.manual_computation_body to be produced by a xla.sdy.GlobalToLocalShape CustomCallOp}}
+  // expected-error @+1 {{failed to rewrite func.call to manual computation}}
+  %0 = call @xla.sdy.manual_computation_body(%c) : (tensor<0x8xf32>) -> tensor<0x8xf32>
+  %1 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"b"}]>]>, xla.sdy.manual_axes = #sdy<manual_axes{"b"}>} : (tensor<0x8xf32>) -> (tensor<0x16xf32>)
+  return %1 : tensor<0x16xf32>
+}
+
+func.func @xla.sdy.manual_computation_body(%arg0: tensor<0x8xf32>) -> tensor<0x8xf32> {
+  return %arg0 : tensor<0x8xf32>
+}
+
+// -----
+
+sdy.mesh @mesh = <["a"=2, "b"=2]>
+
+// Tests import error when manual computation body call result is not consumed by LocalToGlobalShape.
+func.func @manual_computation_missing_local_to_global_shape_direct_attrs(%arg0: tensor<0x16xf32>) -> (tensor<0x8xf32>) {
+  %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%arg0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {}]>]>, xla.sdy.manual_axes = #sdy<manual_axes{"b"}>} : (tensor<0x16xf32>) -> tensor<0x8xf32>
+  // expected-error @+2 {{'func.call' op expected the first use of @xla.sdy.manual_computation_body to be a xla.sdy.LocalToGlobalShape CustomCallOp}}
+  // expected-error @+1 {{failed to rewrite func.call to manual computation}}
+  %1 = call @xla.sdy.manual_computation_body(%0) : (tensor<0x8xf32>) -> tensor<0x8xf32>
+  return %1 : tensor<0x8xf32>
+}
+
+func.func @xla.sdy.manual_computation_body(%arg0: tensor<0x8xf32>) -> tensor<0x8xf32> {
+  return %arg0 : tensor<0x8xf32>
+}
+
diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_hlo_sharding_v3.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_hlo_sharding_v3.mlir
new file mode 100644
index 00000000000..6e7060cb48f
--- /dev/null
+++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_shard_map_import_hlo_sharding_v3.mlir
@@ -0,0 +1,362 @@
+// Copyright 2026 The OpenXLA 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.
+// ==============================================================================
+// RUN: sdy_opt %s --split-input-file -xla-sdy-test-flatten-call-graph -xla-sdy-round-trip-shard-map-import 2>&1 | FileCheck %s
+
+// Tests importing shard map custom calls back into `sdy.manual_computation` under
+// HloShardingV3 (where frontend_attributes are omitted).
+// Verifies that:
+// 1. Shardings and manual axes are read directly off the native attributes of the
+//    boundary ops, without any inference:
+//    - global `in_shardings` come from the `@Sharding` custom call wrapping each
+//      operand of `GlobalToLocalShape`, which is then unwired and erased;
+//    - global `out_shardings` come from `LocalToGlobalShape`;
+//    - `manual_axes` come from `GlobalToLocalShape`, or from the body call op
+//      when the manual computation has no operands.
+// 2. Outlined manual computation call bodies are inlined back into `sdy.manual_computation`.
+
+sdy.mesh @mesh_0 = <["a"=4, "b"=2]>
+sdy.mesh @mesh_1 = <["a"=2, "b"=2, "c"=2, "d"=2]>
+
+// Tests single manual computation with multiple inputs/outputs, all-reduce, and
+// in/out shardings read off the `@Sharding` wrappers and `LocalToGlobalShape`.
+// CHECK-LABEL: func @single_manual_comp
+func.func @single_manual_comp(%arg0: tensor<8x16xf32>, %arg1: tensor<16x32xf32>) -> (tensor<8x32xf32>) {
+  // CHECK-NOT: call @xla.sdy.manual_computation_body
+  // CHECK:               %[[MAN_COMP:.*]] = sdy.manual_computation(%arg0, %arg1)
+  // CHECK-SAME{LITERAL}:     in_shardings=[<@mesh_0, [{"a"}, {"b"}]>, <@mesh_0, [{"b"}, {}], replicated={"a"}>]
+  // CHECK-SAME{LITERAL}:     out_shardings=[<@mesh_0, [{"a"}, {}], replicated={"b"}>]
+
+  // CHECK-SAME{LITERAL}:     manual_axes={"a", "b"}
+  // CHECK-SAME:              (%arg2: tensor<2x8xf32>, %arg3: tensor<8x32xf32>) {
+  // CHECK-NEXT:            %[[ADD_0:.*]] = stablehlo.add %arg2, %arg2 : tensor<2x8xf32>
+  // CHECK-NEXT:            %[[DOT:.*]] = stablehlo.dot %[[ADD_0]], %arg3 : (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32>
+  // CHECK-NEXT:            %[[REDUCE:.*]] = "stablehlo.all_reduce"(%[[DOT]])
+  // CHECK-NEXT:            ^bb0(%arg4: tensor<f32>, %arg5: tensor<f32>):
+  // CHECK-NEXT:              %[[ADD_1:.*]] = stablehlo.add %arg4, %arg5 : tensor<f32>
+  // CHECK-NEXT:              stablehlo.return %[[ADD_1]] : tensor<f32>
+  // CHECK-NEXT:            }) : (tensor<2x32xf32>) -> tensor<2x32xf32>
+  // CHECK-NEXT:            sdy.return %[[REDUCE]] : tensor<2x32xf32>
+  // CHECK-NEXT:          } : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32>
+  // CHECK-NEXT:          return %[[MAN_COMP]] : tensor<8x32xf32>
+  %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {"b"}]>]>} : (tensor<8x16xf32>) -> tensor<8x16xf32>
+  %s1 = stablehlo.custom_call @Sharding(%arg1) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}, {}], replicated={"a"}>]>} : (tensor<16x32xf32>) -> tensor<16x32xf32>
+  %0:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0, %s1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>, <@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy<manual_axes{"a", "b"}>} : (tensor<8x16xf32>, tensor<16x32xf32>) -> (tensor<2x8xf32>, tensor<8x32xf32>)
+  %1 = call @xla.sdy.manual_computation_body(%0#0, %0#1) : (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32>
+  %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}], replicated={"b"}>]>} : (tensor<2x32xf32>) -> tensor<8x32xf32>
+  return %2 : tensor<8x32xf32>
+}
+
+// CHECK-LABEL: func @single_manual_comp_name_is_not_prefix_nor_suffix
+func.func @single_manual_comp_name_is_not_prefix_nor_suffix(%arg0: tensor<8x8xf32>) -> (tensor<8x8xf32>) {
+  // CHECK-NOT: call @my_model.___call__.fwd.xla.sdy.manual_computation_body_14.1234
+  // CHECK:               %[[MAN_COMP:.*]] = sdy.manual_computation(%arg0)
+  // CHECK-SAME{LITERAL}:     in_shardings=[<@mesh_0, [{"a"}, {}]>]
+  // CHECK-SAME{LITERAL}:     out_shardings=[<@mesh_0, [{"a"}, {}]>]
+  // CHECK-SAME{LITERAL}:     manual_axes={"a"}
+  // CHECK-SAME:              (%arg1: tensor<2x8xf32>) {
+  // CHECK-NEXT:            sdy.return %arg1 : tensor<2x8xf32>
+  // CHECK-NEXT:          } : (tensor<8x8xf32>) -> tensor<8x8xf32>
+  // CHECK-NEXT:          return %[[MAN_COMP]] : tensor<8x8xf32>
+  %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32>
+  %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy<manual_axes{"a"}>} : (tensor<8x8xf32>) -> tensor<2x8xf32>
+  %1 = call @my_model.___call__.fwd.xla.sdy.manual_computation_body_14.1234(%0) : (tensor<2x8xf32>) -> tensor<2x8xf32>
+  %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}]>]>} : (tensor<2x8xf32>) -> tensor<8x8xf32>
+  return %2 : tensor<8x8xf32>
+}
+
+// Tests chained manual computations: outputs of first shard map feed second shard map.
+// CHECK-LABEL: func @manual_comp_using_another
+func.func @manual_comp_using_another(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+  // CHECK-NOT: call @xla.sdy.manual_computation_body_0
+  // CHECK:               %[[MAN_COMP_0:.*]] = sdy.manual_computation(%arg0)
+  // CHECK-SAME{LITERAL}:     in_shardings=[<@mesh_0, [{"a"}, {}]>]
+  // CHECK-SAME{LITERAL}:     out_shardings=[<@mesh_0, [{"a"}, {}]>]
+  // CHECK-SAME{LITERAL}:     manual_axes={"a"}
+  // CHECK-SAME:              (%arg1: tensor<2x8xf32>) {
+  // CHECK-NEXT:            sdy.return %arg1 : tensor<2x8xf32>
+  // CHECK-NEXT:          } : (tensor<8x8xf32>) -> tensor<8x8xf32>
+  // CHECK-NOT: call @xla.sdy.manual_computation_body_1
+  // CHECK-NEXT:          %[[MAN_COMP_1:.*]] = sdy.manual_computation(%[[MAN_COMP_0]])
+  // CHECK-SAME{LITERAL}:     in_shardings=[<@mesh_0, [{}, {"b"}]>]
+  // CHECK-SAME{LITERAL}:     out_shardings=[<@mesh_0, [{}, {"b"}]>]
+  // CHECK-SAME{LITERAL}:     manual_axes={"b"}
+  // CHECK-SAME:              (%arg1: tensor<8x4xf32>) {
+  // CHECK-NEXT:            sdy.return %arg1 : tensor<8x4xf32>
+  // CHECK-NEXT:          } : (tensor<8x8xf32>) -> tensor<8x8xf32>
+  // CHECK-NEXT:          return %[[MAN_COMP_1]] : tensor<8x8xf32>
+  %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32>
+  %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy<manual_axes{"a"}>} : (tensor<8x8xf32>) -> tensor<2x8xf32>
+  %1 = call @xla.sdy.manual_computation_body_0(%0) : (tensor<2x8xf32>) -> tensor<2x8xf32>
+  %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"a"}, {}]>]>} : (tensor<2x8xf32>) -> tensor<8x8xf32>
+  %s1 = stablehlo.custom_call @Sharding(%2) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {"b"}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32>
+  %3 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy<manual_axes{"b"}>} : (tensor<8x8xf32>) -> tensor<8x4xf32>
+  %4 = call @xla.sdy.manual_computation_body_1(%3) : (tensor<8x4xf32>) -> tensor<8x4xf32>
+  %5 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%4) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {"b"}]>]>} : (tensor<8x4xf32>) -> tensor<8x8xf32>
+  return %5 : tensor<8x8xf32>
+}
+
+// Tests nested manual computations: reconstructing inner and outer sdy.manual_computation regions.
+// CHECK-LABEL: func @nested_shmaps
+func.func @nested_shmaps(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> {
+  // CHECK-NOT: call @xla.sdy.manual_computation_body_2
+  // CHECK-NOT: call @xla.sdy.manual_computation_body_3
+  // CHECK:               %[[MAN_COMP_0:.*]] = sdy.manual_computation(%arg0)
+  // CHECK-SAME{LITERAL}:     in_shardings=[<@mesh_1, [{"a"}, {}]>]
+  // CHECK-SAME{LITERAL}:     out_shardings=[<@mesh_1, [{"a"}, {}]>]
+  // CHECK-SAME{LITERAL}:     manual_axes={"a"}
+  // CHECK-SAME:              (%arg1: tensor<2x8xf32>) {
+  // CHECK-NEXT:            %[[MAN_COMP_1:.*]] = sdy.manual_computation(%arg1)
+  // CHECK-SAME{LITERAL}:       in_shardings=[<@mesh_1, [{}, {"b"}]>]
+  // CHECK-SAME{LITERAL}:       out_shardings=[<@mesh_1, [{}, {"b"}]>]
+  // CHECK-SAME{LITERAL}:       manual_axes={"b"}
+  // CHECK-SAME:                (%arg2: tensor<2x4xf32>) {
+  // CHECK-NEXT:              %[[MUL:.*]] = stablehlo.multiply %arg2, %arg2 : tensor<2x4xf32>
+  // CHECK-NEXT:              sdy.return %[[MUL]] : tensor<2x4xf32>
+  // CHECK-NEXT:            } : (tensor<2x8xf32>) -> tensor<2x8xf32>
+  // CHECK-NEXT:            sdy.return %[[MAN_COMP_1]] : tensor<2x8xf32>
+  // CHECK-NEXT:          } : (tensor<4x8xf32>) -> tensor<4x8xf32>
+  // CHECK-NEXT:          return %[[MAN_COMP_0]] : tensor<4x8xf32>
+  %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{"a"}, {}]>]>} : (tensor<4x8xf32>) -> tensor<4x8xf32>
+  %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {}]>]>, xla.sdy.manual_axes = #sdy<manual_axes{"a"}>} : (tensor<4x8xf32>) -> tensor<2x8xf32>
+  %1 = call @xla.sdy.manual_computation_body_3(%0) : (tensor<2x8xf32>) -> tensor<2x8xf32>
+  %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{"a"}, {}]>]>} : (tensor<2x8xf32>) -> tensor<4x8xf32>
+  return %2 : tensor<4x8xf32>
+}
+
+// Tests nested manual computation with operations in the outer region around the inner region.
+// CHECK-LABEL: func @nested_shmaps_extra_op
+func.func @nested_shmaps_extra_op(%arg0: tensor<4x8xf32>) -> tensor<4x8xf32> {
+  // CHECK-NOT: call @xla.sdy.manual_computation_body_4
+  // CHECK-NOT: call @xla.sdy.manual_computation_body_5
+  // CHECK:               %[[MAN_COMP_0:.*]] = sdy.manual_computation(%arg0)
+  // CHECK-SAME{LITERAL}:     in_shardings=[<@mesh_1, [{"a"}, {}]>]
+  // CHECK-SAME{LITERAL}:     out_shardings=[<@mesh_1, [{"a"}, {}]>]
+  // CHECK-SAME{LITERAL}:     manual_axes={"a"}
+  // CHECK-SAME:              (%arg1: tensor<2x8xf32>) {
+  // CHECK-NEXT:            %[[MAN_COMP_1:.*]] = sdy.manual_computation(%arg1)
+  // CHECK-SAME{LITERAL}:       in_shardings=[<@mesh_1, [{}, {"b"}]>]
+  // CHECK-SAME{LITERAL}:       out_shardings=[<@mesh_1, [{}, {"b"}]>]
+  // CHECK-SAME{LITERAL}:       manual_axes={"b"}
+  // CHECK-SAME:                (%arg2: tensor<2x4xf32>) {
+  // CHECK-NEXT:              %[[MUL:.*]] = stablehlo.multiply %arg2, %arg2 : tensor<2x4xf32>
+  // CHECK-NEXT:              sdy.return %[[MUL]] : tensor<2x4xf32>
+  // CHECK-NEXT:            } : (tensor<2x8xf32>) -> tensor<2x8xf32>
+  // CHECK-NEXT:            %[[ADD:.*]] = stablehlo.add %[[MAN_COMP_1]], %[[MAN_COMP_1]] : tensor<2x8xf32>
+  // CHECK-NEXT:            sdy.return %[[ADD]] : tensor<2x8xf32>
+  // CHECK-NEXT:          } : (tensor<4x8xf32>) -> tensor<4x8xf32>
+  // CHECK-NEXT:          return %[[MAN_COMP_0]] : tensor<4x8xf32>
+  %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{"a"}, {}]>]>} : (tensor<4x8xf32>) -> tensor<4x8xf32>
+  %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {}]>]>, xla.sdy.manual_axes = #sdy<manual_axes{"a"}>} : (tensor<4x8xf32>) -> tensor<2x8xf32>
+  %1 = call @xla.sdy.manual_computation_body_5(%0) : (tensor<2x8xf32>) -> tensor<2x8xf32>
+  %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{"a"}, {}]>]>} : (tensor<2x8xf32>) -> tensor<4x8xf32>
+  return %2 : tensor<4x8xf32>
+}
+
+// Tests manual computation with no inputs: there is no `GlobalToLocalShape`, so
+// the body call op is the only carrier of the manual axes.
+// CHECK-LABEL: func @manual_computation_no_inputs
+func.func @manual_computation_no_inputs() -> tensor<4xi64> {
+
+  // CHECK-NOT: call @xla.sdy.manual_computation_body_6
+  // CHECK:               %[[SHMAP:.*]] = sdy.manual_computation()
+  // CHECK-SAME{LITERAL}:     in_shardings=[]
+  // CHECK-SAME{LITERAL}:     out_shardings=[<@mesh_0, [{"b"}]>]
+  // CHECK-SAME{LITERAL}:     manual_axes={"b"}
+  // CHECK-SAME{LITERAL}:     () {
+  // CHECK-NEXT:            %[[C:.*]] = stablehlo.constant dense<[2, 3]> : tensor<2xi64>
+  // CHECK-NEXT:            sdy.return %[[C]] : tensor<2xi64>
+  // CHECK-NEXT:          } : () -> tensor<4xi64>
+  // CHECK-NEXT:          return %[[SHMAP]] : tensor<4xi64>
+  %0 = call @xla.sdy.manual_computation_body_6() {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}]>]>, xla.sdy.manual_axes = #sdy<manual_axes{"b"}>} : () -> tensor<2xi64>
+  %1 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}]>]>} : (tensor<2xi64>) -> tensor<4xi64>
+  return %1 : tensor<4xi64>
+}
+
+// Edge case: Manual computation consuming an input but producing no outputs.
+// CHECK-LABEL: func @manual_computation_no_outputs
+func.func @manual_computation_no_outputs(%arg0: tensor<4xi64>) {
+  // CHECK-NOT: call @xla.sdy.manual_computation_body_7
+  // CHECK:               sdy.manual_computation(%arg0)
+  // CHECK-SAME{LITERAL}:     in_shardings=[<@mesh_0, [{"b"}]>]
+  // CHECK-SAME{LITERAL}:     out_shardings=[]
+  // CHECK-SAME{LITERAL}:     manual_axes={"b"}
+  // CHECK-SAME{LITERAL}:     (%arg1: tensor<2xi64>) {
+  // CHECK-NEXT:            stablehlo.custom_call @sdy_testonly(%arg1) : (tensor<2xi64>) -> ()
+  // CHECK-NEXT:            sdy.return
+  // CHECK-NEXT:          } : (tensor<4xi64>) -> ()
+  // CHECK-NEXT:          return
+  %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}]>]>} : (tensor<4xi64>) -> tensor<4xi64>
+  %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}]>]>, xla.sdy.manual_axes = #sdy<manual_axes{"b"}>} : (tensor<4xi64>) -> tensor<2xi64>
+  call @xla.sdy.manual_computation_body_7(%0) : (tensor<2xi64>) -> ()
+  return
+}
+
+// Edge case: Degenerate manual computation with no inputs and no outputs.
+// CHECK-LABEL: func @manual_computation_no_inputs_no_outputs
+func.func @manual_computation_no_inputs_no_outputs() {
+  // CHECK-NEXT: sdy.manual_computation() in_shardings=[] out_shardings=[] manual_axes={} () {
+  // CHECK-NEXT:   sdy.return
+  // CHECK-NEXT: } : () -> ()
+  // CHECK-NEXT: return
+  call @xla.sdy.manual_computation_body_8() {mhlo.frontend_attributes = {inlineable = "false"}} : () -> ()
+  return
+}
+
+// Edge case: Manual computation with 0-sized dimension tensors.
+// CHECK-LABEL: func @manual_computation_some_zero_dim_inputs_outputs
+func.func @manual_computation_some_zero_dim_inputs_outputs(%arg0: tensor<0x16xf32>, %arg1: tensor<16x32xf32>) -> (tensor<0x32xf32>, tensor<16x32xf32>) {
+
+  // CHECK-NOT: call @xla.sdy.manual_computation_body
+  // CHECK:               %[[CONST_0_32:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<0x32xf32>
+  // CHECK:               %[[MAN_COMP:.*]]:2 = sdy.manual_computation(%arg0, %arg1)
+  // CHECK-SAME{LITERAL}:     in_shardings=[<@mesh_0, [{}, {"b"}]>, <@mesh_0, [{"b"}, {}]>]
+  // CHECK-SAME{LITERAL}:     out_shardings=[<@mesh_0, [{}, {}], replicated={"b"}>, <@mesh_0, [{"b"}, {}]>]
+  // CHECK-SAME{LITERAL}:     manual_axes={"b"}
+  // CHECK-SAME:              (%arg2: tensor<0x8xf32>, %arg3: tensor<8x32xf32>) {
+  // CHECK-NEXT:            %[[DOT:.*]] = stablehlo.dot %arg2, %arg3
+  // CHECK-NEXT:            sdy.return %[[DOT]], %arg3
+  // CHECK-NEXT:          } : (tensor<0x16xf32>, tensor<16x32xf32>) -> (tensor<0x32xf32>, tensor<16x32xf32>)
+  // CHECK-NEXT:          return %[[CONST_0_32]], %[[MAN_COMP]]#1
+  %c1 = stablehlo.constant dense<0.000000e+00> : tensor<0x8xf32>
+  %c2 = stablehlo.constant dense<0.000000e+00> : tensor<0x32xf32>
+  %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {"b"}]>]>} : (tensor<0x16xf32>) -> tensor<0x16xf32>
+  %s1 = stablehlo.custom_call @Sharding(%arg1) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}, {}]>]>} : (tensor<16x32xf32>) -> tensor<16x32xf32>
+  %0:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0, %s1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>, <@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy<manual_axes{"b"}>} : (tensor<0x16xf32>, tensor<16x32xf32>) -> (tensor<0x8xf32>, tensor<8x32xf32>)
+  %1:2 = call @xla.sdy.manual_computation_body_9(%c1, %0#1) : (tensor<0x8xf32>, tensor<8x32xf32>) -> (tensor<0x32xf32>, tensor<8x32xf32>)
+  %2:2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%c2, %1#1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}], replicated={"b"}>, <@mesh_0, [{"b"}, {}]>]>} : (tensor<0x32xf32>, tensor<8x32xf32>) -> (tensor<0x32xf32>, tensor<16x32xf32>)
+  return %c2, %2#1 : tensor<0x32xf32>, tensor<16x32xf32>
+}
+
+// Type preservation: Token arguments are preserved with empty sharding and not assigned manual axes.
+// CHECK-LABEL: func @manual_computation_tokens
+func.func @manual_computation_tokens(%arg0: !stablehlo.token, %arg1: tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>) {
+  // CHECK-NOT: call @xla.sdy.manual_computation_body_11
+  // CHECK:               %[[MAN_COMP:.*]]:2 = sdy.manual_computation(%arg0, %arg1)
+  // CHECK-SAME{LITERAL}:     in_shardings=[<@mesh_0, []>, <@mesh_0, [{"b"}]>]
+  // CHECK-SAME{LITERAL}:     out_shardings=[<@mesh_0, []>, <@mesh_0, [{"b"}]>]
+  // CHECK-SAME{LITERAL}:     manual_axes={"b"}
+  // CHECK-SAME:              (%arg2: !stablehlo.token, %arg3: tensor<1xi64>) {
+  // CHECK-NEXT:            sdy.return %arg2, %arg3 : !stablehlo.token, tensor<1xi64>
+  // CHECK-NEXT:          } : (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>)
+  // CHECK-NEXT:          return %[[MAN_COMP]]#0, %[[MAN_COMP]]#1 : !stablehlo.token, tensor<2xi64>
+  %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, []>]>} : (!stablehlo.token) -> !stablehlo.token
+  %s1 = stablehlo.custom_call @Sharding(%arg1) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}]>]>} : (tensor<2xi64>) -> tensor<2xi64>
+  %0:2 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0, %s1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, []>, <@mesh_0, [{}]>]>, xla.sdy.manual_axes = #sdy<manual_axes{"b"}>} : (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<1xi64>)
+  %1:2 = call @xla.sdy.manual_computation_body_11(%0#0, %0#1) : (!stablehlo.token, tensor<1xi64>) -> (!stablehlo.token, tensor<1xi64>)
+  %2:2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1#0, %1#1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, []>, <@mesh_0, [{"b"}]>]>} : (!stablehlo.token, tensor<1xi64>) -> (!stablehlo.token, tensor<2xi64>)
+  return %2#0, %2#1 : !stablehlo.token, tensor<2xi64>
+}
+
+// Tests stray/dangling boundary custom calls where results are unused: safely stripped.
+// CHECK-LABEL: func @stray_unused_manual_computation_custom_calls
+func.func @stray_unused_manual_computation_custom_calls(%arg0: tensor<0x16xf32>) -> tensor<0x16xf32> {
+  // CHECK-NEXT: %[[CONST_0_8:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<0x8xf32>
+  // CHECK-NEXT: %[[CONST_0_16:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<0x16xf32>
+  // CHECK-NOT: call @xla.sdy.manual_computation_body
+  // CHECK-NOT: call @xla.sdy.GlobalToLocalShape
+  // CHECK-NOT: call @xla.sdy.LocalToGlobalShape
+  // CHECK-NEXT: return %[[CONST_0_16]]
+  %c1 = stablehlo.constant dense<0.000000e+00> : tensor<0x8xf32>
+  %c2 = stablehlo.constant dense<0.000000e+00> : tensor<0x16xf32>
+  %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%arg0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{}, {}]>]>, xla.sdy.manual_axes = #sdy<manual_axes{"b"}>} : (tensor<0x16xf32>) -> tensor<0x8xf32>
+  %1 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%c1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_0, [{"b"}, {}]>]>} : (tensor<0x8xf32>) -> tensor<0x16xf32>
+  return %c2 : tensor<0x16xf32>
+}
+
+// CHECK-NOT: func private @xla.sdy.manual_computation_body(
+func.func private @xla.sdy.manual_computation_body(%arg0: tensor<2x8xf32>, %arg1: tensor<8x32xf32>) -> tensor<2x32xf32> {
+  %0 = stablehlo.add %arg0, %arg0 : tensor<2x8xf32>
+  %1 = stablehlo.dot %0, %arg1 : (tensor<2x8xf32>, tensor<8x32xf32>) -> tensor<2x32xf32>
+  %2 = "stablehlo.all_reduce"(%1) <{replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>}> ({
+  ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
+    %3 = stablehlo.add %arg2, %arg3 : tensor<f32>
+    stablehlo.return %3 : tensor<f32>
+  }) : (tensor<2x32xf32>) -> tensor<2x32xf32>
+  return %2 : tensor<2x32xf32>
+}
+
+// CHECK-NOT: func private @my_model.___call__.fwd.xla.sdy.manual_computation_body_14.1234(
+func.func private @my_model.___call__.fwd.xla.sdy.manual_computation_body_14.1234(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> {
+  return %arg0 : tensor<2x8xf32>
+}
+
+// CHECK-NOT: func private @xla.sdy.manual_computation_body_0(
+func.func private @xla.sdy.manual_computation_body_0(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> {
+  return %arg0 : tensor<2x8xf32>
+}
+
+// CHECK-NOT: func private @xla.sdy.manual_computation_body_1(
+func.func private @xla.sdy.manual_computation_body_1(%arg0: tensor<8x4xf32>) -> tensor<8x4xf32> {
+  return %arg0 : tensor<8x4xf32>
+}
+
+// CHECK-NOT: func private @xla.sdy.manual_computation_body_2(
+func.func private @xla.sdy.manual_computation_body_2(%arg0: tensor<2x4xf32>) -> tensor<2x4xf32> {
+  %0 = stablehlo.multiply %arg0, %arg0 : tensor<2x4xf32>
+  return %0 : tensor<2x4xf32>
+}
+
+// CHECK-NOT: func private @xla.sdy.manual_computation_body_3(
+func.func private @xla.sdy.manual_computation_body_3(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> {
+  %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {"b"}]>]>} : (tensor<2x8xf32>) -> tensor<2x8xf32>
+  %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {}]>]>, xla.sdy.manual_axes = #sdy<manual_axes{"b"}>} : (tensor<2x8xf32>) -> tensor<2x4xf32>
+  %1 = call @xla.sdy.manual_computation_body_2(%0) : (tensor<2x4xf32>) -> tensor<2x4xf32>
+  %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {"b"}]>]>} : (tensor<2x4xf32>) -> tensor<2x8xf32>
+  return %2 : tensor<2x8xf32>
+}
+
+// CHECK-NOT: func private @xla.sdy.manual_computation_body_4(
+func.func private @xla.sdy.manual_computation_body_4(%arg0: tensor<2x4xf32>) -> tensor<2x4xf32> {
+  %0 = stablehlo.multiply %arg0, %arg0 : tensor<2x4xf32>
+  return %0 : tensor<2x4xf32>
+}
+
+// CHECK-NOT: func private @xla.sdy.manual_computation_body_5(
+func.func private @xla.sdy.manual_computation_body_5(%arg0: tensor<2x8xf32>) -> tensor<2x8xf32> {
+  %s0 = stablehlo.custom_call @Sharding(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {"b"}]>]>} : (tensor<2x8xf32>) -> tensor<2x8xf32>
+  %0 = stablehlo.custom_call @xla.sdy.GlobalToLocalShape(%s0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {}]>]>, xla.sdy.manual_axes = #sdy<manual_axes{"b"}>} : (tensor<2x8xf32>) -> tensor<2x4xf32>
+  %1 = call @xla.sdy.manual_computation_body_4(%0) : (tensor<2x4xf32>) -> tensor<2x4xf32>
+  %2 = stablehlo.custom_call @xla.sdy.LocalToGlobalShape(%1) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@mesh_1, [{}, {"b"}]>]>} : (tensor<2x4xf32>) -> tensor<2x8xf32>
+  %3 = stablehlo.add %2, %2 : tensor<2x8xf32>
+  return %3 : tensor<2x8xf32>
+}
+
+// CHECK-NOT: func private @xla.sdy.manual_computation_body_6(
+func.func private @xla.sdy.manual_computation_body_6() -> tensor<2xi64> {
+  %0 = stablehlo.constant dense<[2, 3]> : tensor<2xi64>
+  return %0 : tensor<2xi64>
+}
+
+// CHECK-NOT: func private @xla.sdy.manual_computation_body_7(
+func.func private @xla.sdy.manual_computation_body_7(%arg0: tensor<2xi64>) {
+  stablehlo.custom_call @sdy_testonly(%arg0) : (tensor<2xi64>) -> ()
+  return
+}
+
+// CHECK-NOT: func private @xla.sdy.manual_computation_body_8(
+func.func private @xla.sdy.manual_computation_body_8() {
+  return
+}
+
+// CHECK-NOT: func private @xla.sdy.manual_computation_body_9(
+func.func private @xla.sdy.manual_computation_body_9(%arg0: tensor<0x8xf32>, %arg1: tensor<8x32xf32>) -> (tensor<0x32xf32>, tensor<8x32xf32>) {
+  %0 = stablehlo.dot %arg0, %arg1 : (tensor<0x8xf32>, tensor<8x32xf32>) -> tensor<0x32xf32>
+  return %0, %arg1 : tensor<0x32xf32>, tensor<8x32xf32>
+}
+
+// CHECK-NOT: func private @xla.sdy.manual_computation_body_11(
+func.func private @xla.sdy.manual_computation_body_11(%arg0: !stablehlo.token, %arg1: tensor<1xi64>) -> (!stablehlo.token, tensor<1xi64>) {
+  return %arg0, %arg1 : !stablehlo.token, tensor<1xi64>
+}

From d2252256c7dec7631e292b426727353aaacab18e Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Fri, 9 Oct 2026 22:26:54 +0530
Subject: [PATCH 0370/1000] Update strategy option in GCE failure handler test

---
 .../distribute/failure_handling/gce_failure_handler_test.py     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/distribute/failure_handling/gce_failure_handler_test.py b/tensorflow/python/distribute/failure_handling/gce_failure_handler_test.py
index dac4b089bd0..97b3a48bda2 100644
--- a/tensorflow/python/distribute/failure_handling/gce_failure_handler_test.py
+++ b/tensorflow/python/distribute/failure_handling/gce_failure_handler_test.py
@@ -478,7 +478,7 @@ class GceFailureHandlingTest(test.TestCase, parameterized.TestCase):
     termination_config = failure_handling.TerminationConfig(
         grace_period=grace_period)
 
-    if strategy_option == 'multi_worker':
+    if strategy_option == 'MWMS_multi_worker':
       has_chief = False
       cluster_spec = multi_worker_test_base.create_cluster_spec(
           has_chief=has_chief,

From 8e6a81187d36bfd6d601e06f3c911cbc94953e95 Mon Sep 17 00:00:00 2001
From: Enver Kayaaslan <enver@google.com>
Date: Fri, 9 Oct 2026 09:51:23 -0700
Subject: [PATCH 0371/1000] Add unit tests for non-flat support on
 AllGatherBroadcastReorder.

PiperOrigin-RevId: 996592026
---
 .../xla/xla/hlo/transforms/collectives/BUILD  |  2 +
 .../all_gather_broadcast_reorder_test.cc      | 89 +++++++++++++++++++
 2 files changed, 91 insertions(+)

diff --git a/third_party/xla/xla/hlo/transforms/collectives/BUILD b/third_party/xla/xla/hlo/transforms/collectives/BUILD
index 5692dd75556..0011e156f6d 100644
--- a/third_party/xla/xla/hlo/transforms/collectives/BUILD
+++ b/third_party/xla/xla/hlo/transforms/collectives/BUILD
@@ -267,8 +267,10 @@ xla_cc_test(
     srcs = ["all_gather_broadcast_reorder_test.cc"],
     deps = [
         ":all_gather_broadcast_reorder",
+        "//xla/hlo/testlib:filecheck",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
         "//xla/hlo/utils:hlo_matchers",
+        "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest_main",
         "@tsl//tsl/platform:statusor",
diff --git a/third_party/xla/xla/hlo/transforms/collectives/all_gather_broadcast_reorder_test.cc b/third_party/xla/xla/hlo/transforms/collectives/all_gather_broadcast_reorder_test.cc
index 6e4a14b3a1c..fc7921c4938 100644
--- a/third_party/xla/xla/hlo/transforms/collectives/all_gather_broadcast_reorder_test.cc
+++ b/third_party/xla/xla/hlo/transforms/collectives/all_gather_broadcast_reorder_test.cc
@@ -17,7 +17,9 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/status/status_matchers.h"
 #include "absl/strings/string_view.h"
+#include "xla/hlo/testlib/filecheck.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/hlo/utils/hlo_matchers.h"
 #include "tsl/platform/statusor.h"
@@ -25,6 +27,8 @@ limitations under the License.
 namespace xla {
 namespace {
 
+using ::absl_testing::IsOkAndHolds;
+
 namespace m = xla::testing::opcode_matchers;
 
 class AllGatherBroadcastReorderTest : public HloHardwareIndependentTestBase {
@@ -118,5 +122,90 @@ ENTRY main {
   RunPass(hlo_string, PassOutput::NoChange);
 }
 
+TEST_F(AllGatherBroadcastReorderTest, DoesNotCrossCallBoundaries) {
+  absl::string_view hlo_string = R"(
+HloModule m
+
+foo {
+  param = f32[5, 4, 8, 128] parameter(0)
+  ROOT ag = f32[5, 4, 8, 256] all-gather(param), dimensions={3}, replica_groups={{0, 1}}
+}
+
+ENTRY main {
+  x = f32[128, 5] parameter(0)
+  bc = f32[5, 4, 8, 128] broadcast(x), dimensions={3, 0}
+  ROOT call = f32[5, 4, 8, 256] call(bc), to_apply=foo
+}
+)";
+  RunPass(hlo_string, PassOutput::NoChange);
+}
+
+TEST_F(AllGatherBroadcastReorderTest, EntryCallsFooTransformationOnFoo) {
+  absl::string_view hlo_string = R"(
+// CHECK: %foo (
+// CHECK: %[[AG:.*]] = f32[256,5]{{.*}} all-gather(%x_foo), {{.*}}dimensions={0}
+// CHECK: ROOT %[[BC:.*]] = f32[5,4,8,256]{{.*}} broadcast(%[[AG]]), dimensions={3,0}
+// CHECK: }
+// CHECK: ENTRY %main (
+// CHECK: ROOT %[[CALL:.*]] = f32[5,4,8,256]{{.*}} call(%x_main), to_apply=%foo
+HloModule m
+
+foo {
+  x_foo = f32[128, 5] parameter(0)
+  bc = f32[5, 4, 8, 128] broadcast(x_foo), dimensions={3, 0}
+  ROOT ag = f32[5, 4, 8, 256] all-gather(bc), dimensions={3}, replica_groups={{0, 1}}
+}
+
+ENTRY main {
+  x_main = f32[128, 5] parameter(0)
+  ROOT call = f32[5, 4, 8, 256] call(x_main), to_apply=foo
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_THAT(AllGatherBroadcastReorder().Run(module.get()),
+              IsOkAndHolds(true));
+  EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true));
+}
+
+TEST_F(AllGatherBroadcastReorderTest,
+       EntryCallsFooAndBarFooCallsBarTransformationOnBar) {
+  absl::string_view hlo_string = R"(
+// CHECK: %bar (
+// CHECK: %[[AG:.*]] = f32[256,5]{{.*}} all-gather(%x_bar), {{.*}}dimensions={0}
+// CHECK: ROOT %[[BC:.*]] = f32[5,4,8,256]{{.*}} broadcast(%[[AG]]), dimensions={3,0}
+// CHECK: }
+// CHECK: %foo (
+// CHECK: ROOT %[[CALL_BAR:.*]] = f32[5,4,8,256]{{.*}} call(%x_foo), to_apply=%bar
+// CHECK: }
+// CHECK: ENTRY %main (
+// CHECK: %[[CALL_FOO:.*]] = f32[5,4,8,256]{{.*}} call(%x_main), to_apply=%foo
+// CHECK: %[[CALL_BAR2:.*]] = f32[5,4,8,256]{{.*}} call(%x_main), to_apply=%bar
+// CHECK: ROOT %[[ADD:.*]] = f32[5,4,8,256]{{.*}} add(%[[CALL_FOO]], %[[CALL_BAR2]])
+HloModule m
+
+bar {
+  x_bar = f32[128, 5] parameter(0)
+  bc = f32[5, 4, 8, 128] broadcast(x_bar), dimensions={3, 0}
+  ROOT ag = f32[5, 4, 8, 256] all-gather(bc), dimensions={3}, replica_groups={{0, 1}}
+}
+
+foo {
+  x_foo = f32[128, 5] parameter(0)
+  ROOT call_bar = f32[5, 4, 8, 256] call(x_foo), to_apply=bar
+}
+
+ENTRY main {
+  x_main = f32[128, 5] parameter(0)
+  call_foo = f32[5, 4, 8, 256] call(x_main), to_apply=foo
+  call_bar = f32[5, 4, 8, 256] call(x_main), to_apply=bar
+  ROOT add = f32[5, 4, 8, 256] add(call_foo, call_bar)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_THAT(AllGatherBroadcastReorder().Run(module.get()),
+              IsOkAndHolds(true));
+  EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true));
+}
+
 }  // namespace
 }  // namespace xla

From 4feeeda1856ed5c30a8f16ab498be8eb3e28a4c8 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Fri, 9 Oct 2026 22:36:47 +0530
Subject: [PATCH 0372/1000] Refactor import handling and dlopen flag
 restoration

Refactor import logic to handle ModuleNotFoundError and ensure dlopen flags are restored properly.
---
 tensorflow/python/pywrap_tensorflow.py | 32 +++++++++++++++-----------
 1 file changed, 19 insertions(+), 13 deletions(-)

diff --git a/tensorflow/python/pywrap_tensorflow.py b/tensorflow/python/pywrap_tensorflow.py
index 2f32a97a32b..9c6b946c923 100644
--- a/tensorflow/python/pywrap_tensorflow.py
+++ b/tensorflow/python/pywrap_tensorflow.py
@@ -71,20 +71,26 @@ try:
 
   # pylint: disable=wildcard-import,g-import-not-at-top,line-too-long,undefined-variable
   try:
-    from tensorflow.python._pywrap_tensorflow_internal import *
-  # This try catch logic is because there is no bazel equivalent for py_extension.
-  # Externally in opensource we must enable exceptions to load the shared object
-  # by exposing the PyInit symbols with pybind. This error will only be
-  # caught internally or if someone changes the name of the target _pywrap_tensorflow_internal.
+    try:
+      from tensorflow.python._pywrap_tensorflow_internal import *
+    # This try catch logic is because there is no bazel equivalent for py_extension.
+    # Externally in opensource we must enable exceptions to load the shared object
+    # by exposing the PyInit symbols with pybind. This error will only be
+    # caught internally or if someone changes the name of the target _pywrap_tensorflow_internal.
 
-  # This logic is used in other internal projects using py_extension.
-  except ModuleNotFoundError:
-    pass
-
-  if _use_dlopen_global_flags:
-    pywrap_dlopen_global_flags.reset_dlopen_flags()
-  elif _can_set_rtld_local:
-    sys.setdlopenflags(_default_dlopen_flags)
+    # This logic is used in other internal projects using py_extension.
+    except ModuleNotFoundError:
+      pass
+  finally:
+    # Always restore the dlopen flags, even if the import above raised an
+    # ImportError that is not a ModuleNotFoundError (e.g. a genuine load
+    # failure). Otherwise a real failure here would leave the process's
+    # dlopen flags permanently altered for any caller that catches the
+    # resulting ImportError and keeps running.
+    if _use_dlopen_global_flags:
+      pywrap_dlopen_global_flags.reset_dlopen_flags()
+    elif _can_set_rtld_local:
+      sys.setdlopenflags(_default_dlopen_flags)
 except ImportError as exc:
   if os.name == 'nt':
     try:

From 8223708e19bb9d96e1dcd7f7fa58afdacff8fa71 Mon Sep 17 00:00:00 2001
From: Blake Hechtman <blakehechtman@google.com>
Date: Fri, 9 Oct 2026 09:55:06 -0700
Subject: [PATCH 0373/1000] [XLA] Support mixed call contexts in
 HloDataflowAnalysis when propagation is disabled

When both propagate_through_calls and propagate_through_control_flow are false, subcomputation parameters define their own values regardless of call context, and values do not propagate across call, conditional, or async boundaries.

PiperOrigin-RevId: 996594171
---
 .../xla/hlo/analysis/hlo_dataflow_analysis.cc | 55 +++++++++------
 .../analysis/hlo_dataflow_analysis_test.cc    | 68 +++++++++++++++++++
 2 files changed, 103 insertions(+), 20 deletions(-)

diff --git a/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.cc b/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.cc
index ebdd651055f..7a46885b613 100644
--- a/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.cc
+++ b/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis.cc
@@ -915,10 +915,12 @@ bool HloDataflowAnalysis::UpdateParameterValueSet(HloInstruction* parameter) {
   const CallGraphNode& call_graph_node =
       call_graph_->GetNode(parameter->parent());
 
-  // Subcomputations called in a parallel context (eg, map) do not have dataflow
-  // from the caller operands.
+  // Subcomputations called in a parallel context (eg, map), dead computations,
+  // or computations when cross computation propagation is disabled do not have
+  // dataflow from the caller operands.
   if (call_graph_node.context() == CallContext::kEmbedded ||
-      call_graph_node.caller_callsites().empty()) {
+      call_graph_node.caller_callsites().empty() ||
+      (!propagate_through_calls_ && !propagate_through_control_flow_)) {
     return false;
   }
   CHECK_EQ(call_graph_node.context(), CallContext::kControlFlow);
@@ -1355,22 +1357,25 @@ void HloDataflowAnalysis::Propagate() {
       // If user sequentially calls a computation, then the respective
       // parameter(s) of the computation need to be updated.
       if (user->opcode() == HloOpcode::kConditional) {
-        // If operand 0 is the use of instruction, then no parameters need to be
-        // updated, since that is the branch_index of the conditional.
-        // If operand n+1 is the use of instruction, then the branch_computation
-        // n's parameter need to be updated.
-        //
-        // Note that the same instruction can be used in multiple branches'
-        // operands.
-        for (int j = 0; j < user->branch_count(); ++j) {
-          if (user->operand(j + 1) == instruction) {
-            add_to_worklist(
-                user->branch_computation(j)->parameter_instruction(0));
+        if (propagate_through_control_flow_) {
+          // If operand 0 is the use of instruction, then no parameters need to
+          // be updated, since that is the branch_index of the conditional. If
+          // operand n+1 is the use of instruction, then the branch_computation
+          // n's parameter need to be updated.
+          //
+          // Note that the same instruction can be used in multiple branches'
+          // operands.
+          for (int j = 0; j < user->branch_count(); ++j) {
+            if (user->operand(j + 1) == instruction) {
+              add_to_worklist(
+                  user->branch_computation(j)->parameter_instruction(0));
+            }
           }
         }
       } else if (user->opcode() == HloOpcode::kAsyncUpdate ||
                  user->opcode() == HloOpcode::kAsyncDone) {
-        if (HloInstruction::IsThreadIncluded(user->async_execution_thread(),
+        if (propagate_through_control_flow_ &&
+            HloInstruction::IsThreadIncluded(user->async_execution_thread(),
                                              execution_threads_)) {
           // For async update and async done, we cannot distinguish which
           // parameter needs to be updated so add all to the worklist.
@@ -1390,6 +1395,12 @@ void HloDataflowAnalysis::Propagate() {
                   called_computation->execution_thread(), execution_threads_)) {
             continue;
           }
+          if ((!propagate_through_calls_ &&
+               user->opcode() == HloOpcode::kCall) ||
+              (!propagate_through_control_flow_ &&
+               user->opcode() != HloOpcode::kCall)) {
+            continue;
+          }
           const CallGraphNode& call_graph_node =
               call_graph_->GetNode(called_computation);
           if (call_graph_node.context() == CallContext::kControlFlow) {
@@ -1477,11 +1488,6 @@ absl::Status HloDataflowAnalysis::InitializeInstructionValueSets() {
       continue;
     }
     const CallGraphNode& call_graph_node = call_graph_->GetNode(computation);
-    const bool is_regular_call_computation =
-        IsRegularCallComputation(call_graph_node);
-    const bool is_control_flow_computation =
-        !call_graph_node.caller_callsites().empty() &&
-        !is_regular_call_computation;
     for (HloInstruction* instruction :
          computation->MakeInstructionPostOrder()) {
       // Create an empty shape tree.
@@ -1537,6 +1543,10 @@ absl::Status HloDataflowAnalysis::InitializeInstructionValueSets() {
           // flow from their operands or from cross computation dataflow.
           break;
         case HloOpcode::kParameter: {
+          if (!propagate_through_calls_ && !propagate_through_control_flow_) {
+            define_all_values();
+            break;
+          }
           if (call_graph_node.context() == CallContext::kBoth) {
             // We do not support a subcomputation that is called from both a
             // parallel and sequential context. In this case, the parameter
@@ -1548,6 +1558,11 @@ absl::Status HloDataflowAnalysis::InitializeInstructionValueSets() {
                 "sequential (eg, kCall) context",
                 computation->name());
           }
+          const bool is_regular_call_computation =
+              IsRegularCallComputation(call_graph_node);
+          const bool is_control_flow_computation =
+              !call_graph_node.caller_callsites().empty() &&
+              !is_regular_call_computation;
           if (call_graph_node.caller_callsites().empty() ||
               call_graph_node.context() == CallContext::kEmbedded ||
               (!propagate_through_calls_ && is_regular_call_computation) ||
diff --git a/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis_test.cc b/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis_test.cc
index 1986efcbc93..73893e416ac 100644
--- a/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis_test.cc
+++ b/third_party/xla/xla/hlo/analysis/hlo_dataflow_analysis_test.cc
@@ -4588,6 +4588,74 @@ ENTRY main {
   }
 }
 
+TEST_F(HloDataflowAnalysisTest,
+       DisablePropagateThroughCallsAndControlFlowMixedCallContexts) {
+  const char* hlo_text = R"hlo(
+HloModule module
+
+shared_callee {
+  param = f32[] parameter(0)
+  ROOT neg = f32[] negate(param)
+}
+
+false_branch {
+  param = f32[] parameter(0)
+  ROOT id = f32[] copy(param)
+}
+
+async_comp {
+  param = f32[] parameter(0)
+  ROOT neg = f32[] negate(param)
+}
+
+ENTRY main {
+  const0 = f32[] constant(1.0)
+  pred0 = pred[] constant(true)
+  tup = (f32[]) tuple(const0)
+  gte = f32[] get-tuple-element(tup), index=0
+  mapped = f32[] map(gte), to_apply=shared_callee
+  called = f32[] call(gte), to_apply=shared_callee
+  cond = f32[] conditional(pred0, gte, called), true_computation=shared_callee, false_computation=false_branch
+  async_start = ((f32[]), f32[], u32[]) async-start(cond), calls=async_comp
+  async_update = ((f32[]), f32[], u32[]) async-update(async_start)
+  ROOT async_done = f32[] async-done(async_update)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(
+                                    hlo_text, GetModuleConfigForTest()));
+  HloInstruction* const0 = FindInstruction(module_.get(), "const0");
+  HloInstruction* gte = FindInstruction(module_.get(), "gte");
+  HloInstruction* called = FindInstruction(module_.get(), "called");
+  HloInstruction* cond = FindInstruction(module_.get(), "cond");
+  HloInstruction* async_start = FindInstruction(module_.get(), "async_start");
+  HloInstruction* async_done = FindInstruction(module_.get(), "async_done");
+  HloInstruction* shared_param =
+      module_->GetComputationWithName("shared_callee")
+          ->parameter_instruction(0);
+  HloInstruction* false_param =
+      module_->GetComputationWithName("false_branch")->parameter_instruction(0);
+  HloInstruction* async_param =
+      module_->GetComputationWithName("async_comp")->parameter_instruction(0);
+
+  ASSERT_OK_AND_ASSIGN(
+      auto analysis,
+      HloDataflowAnalysis::Run(
+          *module_, /*ssa_form=*/false, /*bitcast_defines_value=*/false,
+          /*execution_threads=*/{}, /*propagate_through_calls=*/false,
+          /*precompute_uses=*/std::nullopt,
+          /*propagate_through_control_flow=*/false));
+  EXPECT_EQ(&analysis->GetUniqueValueAt(gte),
+            &analysis->GetValueDefinedAt(const0));
+  EXPECT_TRUE(analysis->ValueIsDefinedAt(shared_param));
+  EXPECT_TRUE(analysis->ValueIsDefinedAt(false_param));
+  EXPECT_TRUE(analysis->ValueIsDefinedAt(async_param));
+  EXPECT_TRUE(analysis->ValueIsDefinedAt(called));
+  EXPECT_TRUE(analysis->ValueIsDefinedAt(cond));
+  EXPECT_TRUE(analysis->ValueIsDefinedAt(async_start, {1}));
+  EXPECT_EQ(&analysis->GetUniqueValueAt(async_done),
+            &analysis->GetValueDefinedAt(async_start, {1}));
+}
+
 TEST_F(HloDataflowAnalysisTest, CallMarkerCustomCalls) {
   const char* hlo_text = R"hlo(
 HloModule module

From c7d9ee8bf214bdd5133d05f00e45195fd2eb9d51 Mon Sep 17 00:00:00 2001
From: Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>
Date: Fri, 9 Oct 2026 09:56:03 -0700
Subject: [PATCH 0374/1000] PR #50239: [ROCm] Add libtensilelite-host.so dep
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50239

📝 Summary of Changes
Add missing hipblaslt dependency for upcoming rocm.

🎯 Justification
Prevents runtime error on tests due to missing shared library

🚀 Kind of Contribution
🐛 Bug Fix

📊 Benchmark (for Performance Improvements)
N\A

🧪 Unit Tests:
None

🧪 Execution Tests:
None

Copybara import of the project:

--
da778927e3e8c8eb5190ac1ab7300f3c5ea82011 by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>:

[ROCm] Add libtensilelite-host.so dep

Merging this change closes #50239

PiperOrigin-RevId: 996594788
---
 third_party/xla/third_party/gpus/rocm/BUILD.tpl | 1 +
 1 file changed, 1 insertion(+)

diff --git a/third_party/xla/third_party/gpus/rocm/BUILD.tpl b/third_party/xla/third_party/gpus/rocm/BUILD.tpl
index 09432bad21e..78c923ea0e2 100644
--- a/third_party/xla/third_party/gpus/rocm/BUILD.tpl
+++ b/third_party/xla/third_party/gpus/rocm/BUILD.tpl
@@ -490,6 +490,7 @@ rocm_lib_import(
             "%{rocm_root}/lib/libhipblaslt.so*",
             "%{rocm_root}/lib/librocroller.so*",
             "%{rocm_root}/lib/liborigami.so*",
+            "%{rocm_root}/lib/libtensilelite-host.so*",
         ],
     ) + glob([
         pattern

From 786c2c1a00bcb0cb171293d6aff9310fcefaeec8 Mon Sep 17 00:00:00 2001
From: Alina Sbirlea <asbirlea@google.com>
Date: Fri, 9 Oct 2026 09:56:29 -0700
Subject: [PATCH 0375/1000] Integrate LLVM at llvm/llvm-project@b08e005995ed

Updates LLVM usage to match
[b08e005995ed](https://github.com/llvm/llvm-project/commit/b08e005995ed29a2f89aa9a3f599a9d45db65523)

PiperOrigin-RevId: 996595059
---
 .../xla/third_party/llvm/generated.patch      |   11 +
 .../xla/third_party/llvm/workspace.bzl        |    4 +-
 .../xla/third_party/shardy/temporary.patch    | 1861 +----------------
 .../xla/third_party/shardy/workspace.bzl      |    4 +-
 4 files changed, 33 insertions(+), 1847 deletions(-)

diff --git a/third_party/xla/third_party/llvm/generated.patch b/third_party/xla/third_party/llvm/generated.patch
index 509398da979..e9c206cebe4 100644
--- a/third_party/xla/third_party/llvm/generated.patch
+++ b/third_party/xla/third_party/llvm/generated.patch
@@ -1 +1,12 @@
 Auto generated patch. Do not edit or delete it, even if empty.
+diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll b/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll
+--- a/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll
++++ b/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll
+@@ -1,5 +1,7 @@
+ ; Checks if conversion is only inserted for the spilled register
+ ; instead of 2 conversions for the W register
++
++; REQUIRES: asserts
+ ; RUN: llc -O2 -march=hexagon -mcpu=hexagonv81 -enable-xqf-gen=true \
+ ; RUN: -hexagon-qfloat-mode=ieee -mattr=+hvxv81,+hvx-length128B \
+ ; RUN: -enable-postra-xqf-check -debug-only=handle-qfp 2>&1 < %s -o /dev/null | FileCheck %s
diff --git a/third_party/xla/third_party/llvm/workspace.bzl b/third_party/xla/third_party/llvm/workspace.bzl
index c78ce3fe662..68a14533c79 100644
--- a/third_party/xla/third_party/llvm/workspace.bzl
+++ b/third_party/xla/third_party/llvm/workspace.bzl
@@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "c4d9267a31897dd76b2aab5196b75bffb49e6846"
-    LLVM_SHA256 = "a96c8b33f0d8d3be008b8c9f74ef310254c4f7f7edf8c04c76f4d99ed9d34f1a"
+    LLVM_COMMIT = "b08e005995ed29a2f89aa9a3f599a9d45db65523"
+    LLVM_SHA256 = "3f4f1765816e22a21ac41dbf71d91d519b0b371b05d7e33ef2a0a862c838420c"
 
     tf_http_archive(
         name = name,
diff --git a/third_party/xla/third_party/shardy/temporary.patch b/third_party/xla/third_party/shardy/temporary.patch
index 2676e00bb2d..8fcc0f46504 100644
--- a/third_party/xla/third_party/shardy/temporary.patch
+++ b/third_party/xla/third_party/shardy/temporary.patch
@@ -1,1857 +1,32 @@
-diff --git a/shardy/dialect/sdy/ir/canonicalization.cc b/shardy/dialect/sdy/ir/canonicalization.cc
-index 9af43e32..f4cda6da 100644
---- a/shardy/dialect/sdy/ir/canonicalization.cc
-+++ b/shardy/dialect/sdy/ir/canonicalization.cc
-@@ -108,7 +108,7 @@ class RedundantManualComputationPattern
-     if (!inShardings.empty() || !outShardings.empty()) {
-       MeshAttr mesh =
-           getCommonMesh(inShardings, outShardings, manualComputationOp);
--      SDY_CHECK(mesh) << "expected inputs and outputs to have a common mesh";
-+      assert(mesh && "expected inputs and outputs to have a common mesh");
-       for (StringAttr manualAxis : manualComputationOp.getManualAxes()) {
-         manualAxesProduct *= mesh.getAxisSize(manualAxis);
-       }
-diff --git a/shardy/dialect/sdy/ir/test/manual_computation_verification.mlir b/shardy/dialect/sdy/ir/test/manual_computation_verification.mlir
-index c2c24800..5853611e 100644
---- a/shardy/dialect/sdy/ir/test/manual_computation_verification.mlir
-+++ b/shardy/dialect/sdy/ir/test/manual_computation_verification.mlir
-@@ -294,28 +294,3 @@ func.func @ranked_sharding_on_token(%arg0: !stablehlo.token) -> !stablehlo.token
-   } : (!stablehlo.token) -> !stablehlo.token
-   return %0 : !stablehlo.token
- }
--
--// -----
--
--sdy.mesh @mesh = <["a"=2]>
--
--func.func @scalar_region_arg(%arg0: tensor<16x32xf32>) -> tensor<16x32xf32> {
--  // expected-error @+1 {{op operand shape, corresponding sharding, and region operand shape at index 0 must match. Expected local shape 'tensor<8x32xf32>', actual local shape 'f32'}}
--  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{"a"}, {}]>] out_shardings=[<@mesh, [{"a"}, {}]>] manual_axes={"a"} (%arg1: f32) {
--    %c = stablehlo.constant dense<0.000000e+00> : tensor<8x32xf32>
--    sdy.return %c : tensor<8x32xf32>
--  } : (tensor<16x32xf32>) -> tensor<16x32xf32>
--  func.return %0: tensor<16x32xf32>
--}
--
--// -----
--
--sdy.mesh @mesh = <["a"=2]>
--
--func.func @token_region_result(%arg0: tensor<16x32xf32>, %arg1: !stablehlo.token) -> tensor<16x32xf32> {
--  // expected-error @+1 {{op result shape, corresponding sharding, and region result shape at index 0 must match. Expected local shape 'tensor<8x32xf32>', actual local shape '!stablehlo.token'}}
--  %0 = sdy.manual_computation(%arg0, %arg1) in_shardings=[<@mesh, [{"a"}, {}]>, <@mesh, []>] out_shardings=[<@mesh, [{"a"}, {}]>] manual_axes={"a"} (%arg2: tensor<8x32xf32>, %arg3: !stablehlo.token) {
--    sdy.return %arg3 : !stablehlo.token
--  } : (tensor<16x32xf32>, !stablehlo.token) -> tensor<16x32xf32>
--  func.return %0: tensor<16x32xf32>
--}
-diff --git a/shardy/dialect/sdy/ir/utils.cc b/shardy/dialect/sdy/ir/utils.cc
-index ca77dca8..4a639a80 100644
---- a/shardy/dialect/sdy/ir/utils.cc
-+++ b/shardy/dialect/sdy/ir/utils.cc
-@@ -25,7 +25,6 @@ limitations under the License.
- #include <numeric>
- #include <optional>
- #include <string>
--#include <string_view>
- #include <utility>
- #include <vector>
- 
-@@ -223,7 +222,7 @@ TensorShardingAttr inlineMesh(const SymbolTable& symbolTable,
-                               TensorShardingAttr sharding) {
-   if (auto name = dyn_cast<FlatSymbolRefAttr>(sharding.getMeshOrRef())) {
-     MeshAttr mesh = getMeshAttr(symbolTable, name);
--    SDY_CHECK(mesh) << "unknown mesh: " << std::string_view(name.getValue());
-+    assert(mesh && "unknown mesh");
-     return TensorShardingAttr::get(
-         sharding.getContext(), mesh, sharding.getDimShardings(),
-         sharding.getReplicatedAxes(), sharding.getUnreducedAxes(),
-@@ -264,7 +263,7 @@ Attribute getCommonMeshOrRef(ArrayRef<TensorShardingAttr> operandShardings,
-       meshOrRef = sharding.getMeshOrRef();
-       continue;
-     }
--    if (!otherMesh || otherMesh.empty()) {
-+    if (otherMesh.empty()) {
-       continue;
-     }
-     if (!otherMesh.equals(mesh, ignoreDeviceIds)) {
-@@ -964,13 +963,9 @@ mlir::Attribute getMeshOrRef(
-     int64_t numElements, const SymbolTable& symbolTable,
-     std::function<TensorShardingAttr(int64_t)> getSharding) {
-   for (int64_t i = 0; i < numElements; ++i) {
--    if (TensorShardingAttr sdySharding = getSharding(i); sdySharding) {
--      MeshAttr mesh = sdySharding.getMesh(symbolTable);
--      SDY_CHECK(mesh) << "unknown mesh: "
--                      << std::string_view(sdySharding.getMeshName());
--      if (!mesh.isSingleDevice()) {
--        return sdySharding.getMeshOrRef();
--      }
-+    if (TensorShardingAttr sdySharding = getSharding(i);
-+        sdySharding && !sdySharding.getMesh(symbolTable).isSingleDevice()) {
-+      return sdySharding.getMeshOrRef();
-     }
-   }
-   return nullptr;
-diff --git a/shardy/dialect/sdy/ir/verifiers.cc b/shardy/dialect/sdy/ir/verifiers.cc
-index 65c6212b..3d90562f 100644
---- a/shardy/dialect/sdy/ir/verifiers.cc
-+++ b/shardy/dialect/sdy/ir/verifiers.cc
-@@ -1021,12 +1021,13 @@ LogicalResult verifyManualComputationValue(
-     //    arguments/results match.
-     auto expectedLocalRankedType =
-         RankedTensorType::get(newDimSizes, globalShapedType.getElementType());
--    if (expectedLocalRankedType != localType) {
-+    auto localRankedType = cast<RankedTensorType>(localType);
-+    if (expectedLocalRankedType != localRankedType) {
-       return op->emitOpError(valueKindStr)
-              << " shape, corresponding sharding, and region " << valueKindStr
-              << " shape at index " << valueIndex
-              << " must match. Expected local shape " << expectedLocalRankedType
--             << ", actual local shape " << localType;
-+             << ", actual local shape " << localRankedType;
-     }
-   }
- 
-diff --git a/shardy/dialect/sdy/transforms/export/insert_explicit_reshards.cc b/shardy/dialect/sdy/transforms/export/insert_explicit_reshards.cc
-index 756b46ae..6b86b36e 100644
---- a/shardy/dialect/sdy/transforms/export/insert_explicit_reshards.cc
-+++ b/shardy/dialect/sdy/transforms/export/insert_explicit_reshards.cc
-@@ -19,7 +19,6 @@ limitations under the License.
- #include <cstdint>
- #include <memory>  // IWYU pragma: keep
- #include <optional>
--#include <string_view>
- #include <utility>
- 
- #include "llvm/ADT/DenseMap.h"
-@@ -351,7 +350,7 @@ std::optional<MeshOp> getMesh(ArrayRef<TensorShardingAttr> inShardings,
-     return std::nullopt;
-   }
-   MeshOp meshOp = getMeshOp(symbolTable, *meshName);
--  SDY_CHECK(meshOp) << "unknown mesh: " << std::string_view(*meshName);
-+  assert(meshOp && "unknown mesh");
-   if (meshOp.getMesh().isSingleDevice()) {
-     return std::nullopt;
-   }
-diff --git a/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc b/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc
-index 4cef1c10..111e882e 100644
---- a/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc
-+++ b/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc
-@@ -557,8 +557,7 @@ class GenericOpPattern : public ConversionPattern {
-     Dialect* dialect = op->getDialect();
-     if ((dialect && dialect->getNamespace() != "stablehlo" &&
-          !isa<func::CallOp, func::ReturnOp, sdy::ReturnOp, sdy::AllReduceOp,
--              sdy::CollectivePermuteOp, sdy::ShardedToUnreducedOp,
--              sdy::ReplicatedToUnreducedOp>(op)) ||
-+              sdy::ShardedToUnreducedOp, sdy::ReplicatedToUnreducedOp>(op)) ||
-         hasCustomPadHandling(op)) {
-       return failure();
-     }
-diff --git a/shardy/dialect/sdy/transforms/export/reshard_to_collectives.cc b/shardy/dialect/sdy/transforms/export/reshard_to_collectives.cc
-index 3b01769d..5554918b 100644
---- a/shardy/dialect/sdy/transforms/export/reshard_to_collectives.cc
-+++ b/shardy/dialect/sdy/transforms/export/reshard_to_collectives.cc
-@@ -29,7 +29,6 @@ limitations under the License.
- #include "llvm/ADT/DenseSet.h"
- #include "llvm/ADT/STLExtras.h"
- #include "llvm/ADT/SmallVector.h"
--#include "llvm/Support/MathExtras.h"
- #include "mlir/Dialect/Func/IR/FuncOps.h"  // IWYU pragma: keep
- #include "mlir/IR/Attributes.h"
- #include "mlir/IR/BuiltinAttributes.h"
-@@ -97,16 +96,10 @@ AxisList concatAxisLists(MutableArrayRef<AxisList> axisLists) {
-   return result;
- }
- 
--// Remove the common prefix of `inAxesPerDim` and `outAxesPerDim` for dimensions
--// that do not need to be unsharded first.
-+// Remove the common prefix of `inAxesPerDim` and `outAxesPerDim`.
- void removeCommonPrefix(SmallVector<AxisList>& inAxesPerDim,
--                        SmallVector<AxisList>& outAxesPerDim,
--                        ArrayRef<bool> dimsToUnshard) {
--  for (auto [inAxes, outAxes, unshardDim] :
--       llvm::zip_equal(inAxesPerDim, outAxesPerDim, dimsToUnshard)) {
--    if (unshardDim) {
--      continue;
--    }
-+                        SmallVector<AxisList>& outAxesPerDim) {
-+  for (auto [inAxes, outAxes] : llvm::zip_equal(inAxesPerDim, outAxesPerDim)) {
-     while (!inAxes.empty() && !outAxes.empty() &&
-            inAxes.front() == outAxes.front()) {
-       inAxes.pop_front();
-@@ -202,16 +195,11 @@ struct AllToAllInfo {
- // We define the current state of the transformation as follows:
- //
- // - `inAxesPerDim` - the axes in the current sharding per dimension, such that
--//   the common prefix with the output sharding is removed (unless the
--//   dimension is in `dimsToUnshard`).
-+//   the common prefix with the output sharding is removed.
- // - `outAxesPerDim` - the axes in the output sharding per dimension, such that
--//   the common prefix with the current sharding is removed (unless the
--//   dimension is in `dimsToUnshard`).
-+//   the common prefix with the current sharding is removed.
- // - `currentAxesPerDim` - the axes in the current sharding, including the
- //   common prefix with the output sharding.
--// - `dimsToUnshard` - per-dimension boolean indicating whether a dimension has
--//   different padded sizes in the input and output shardings and must first be
--//   completely unsharded before acquiring its output sharding.
- //
- // These invariants are maintained throughout the algorithm, and specifically
- // after each collective insertion.
-@@ -245,39 +233,21 @@ class CollectiveInserter {
-         outAxesPerDim(getAxesPerDim<AxisList>(outSharding)),
-         currentAxesPerDim(getAxesPerDim<SmallVector<AxisRefAttr>>(inSharding)),
-         capacityPerDim(inSharding.getRank(), 1),
--        collectiveAxesPerDim(inSharding.getRank()),
--        dimsToUnshard(inSharding.getRank(), false) {
-+        collectiveAxesPerDim(inSharding.getRank()) {
-     // Unreduced axes in the input and output sharding must match, given we
-     // insert an all-reduce if an unreduced axis becomes replicated/sharded, and
-     // never insert a reshard that goes from replicated/sharded to unreduced.
-     assert(inSharding.getUnreducedAxes() == outSharding.getUnreducedAxes());
- 
--    // If a dimension is sharded in both the input and output shardings with
--    // different padded sizes, it must first be unsharded so that downstream
--    // passes can pad or trim the dimension while it is unsharded.
--    ArrayRef<int64_t> shape = getTensorShape(result);
--    for (int64_t dim = 0; dim < getRank(); ++dim) {
--      int64_t dimSize = shape[dim];
--      int64_t inShardedSize =
--          inSharding.getDimSharding(dim).getShardedSize(mesh);
--      int64_t outShardedSize =
--          outSharding.getDimSharding(dim).getShardedSize(mesh);
--      if (dimSize > 0 && inShardedSize > 1 && outShardedSize > 1 &&
--          llvm::alignTo(dimSize, inShardedSize) !=
--              llvm::alignTo(dimSize, outShardedSize)) {
--        dimsToUnshard[dim] = true;
--      }
--    }
--
-     // We align sub-axes between the input and output axes, so that we can treat
-     // sub-axes like full axes and assume any two sub-axes that overlap are also
-     // equal, which allows using them as keys in a hash map.
-     alignSubAxesByDecomposition(inAxesPerDim, outAxesPerDim, mesh);
--    // We remove the common prefix of `inAxesPerDim` and `outAxesPerDim` (except
--    // on dimensions in `dimsToUnshard`), since those axes stay exactly the same
--    // during the reshard. We are left with `inAxesPerDim` and `outAxesPerDim`
--    // that need to become empty, via a sequence of collectives.
--    removeCommonPrefix(inAxesPerDim, outAxesPerDim, dimsToUnshard);
-+    // We remove the common prefix of `inAxesPerDim` and `outAxesPerDim`, since
-+    // those axes stay exactly the same during the reshard. We are left with
-+    // `inAxesPerDim` and `outAxesPerDim` that need to become empty, via a
-+    // sequence of collectives.
-+    removeCommonPrefix(inAxesPerDim, outAxesPerDim);
- 
-     inAxisSet = getAxisSet(inAxesPerDim);
-     outAxisToDimAndIndex = getAxisToDimAndIndex(outAxesPerDim);
-@@ -297,12 +267,11 @@ class CollectiveInserter {
-     // can be handled by the previous collectives, so we are left with
-     // all-gathering all axes in `inAxesPerDim` and we're done.
-     //
--    // Otherwise, we need at most 2 iterations (or 3 if a dimension must first
--    // be unsharded due to a padded dimension size mismatch).
-+    // Otherwise, we need at most 2 iterations.
- 
-     // TODO(tomnatan): consider looping only over all-to-all and collective
-     // permute as long as one of them was inserted.
--    for (int i = 0; i < 3 && !isDone(); ++i) {
-+    for (int i = 0; i < 2 && !isDone(); ++i) {
-       // 1. Try to insert an all-slice, that decreases the size of the tensor.
-       tryAllSlice();
- 
-@@ -351,10 +320,9 @@ class CollectiveInserter {
-   // that dimension.
-   //
-   // We gather the suffix of axes (`gatheringAxes`) in `inAxesPerDim[dim]` that
--  // aren't present in `outAxesPerDim` (or whose target dimension in
--  // `outAxesPerDim` is in `dimsToUnshard`, or all axes if `dimsToUnshard[dim]`
--  // is true), and update the internal state as follows: `gatheringAxes` are
--  // popped from the back of `inAxesPerDim[dim]` and `currentAxesPerDim[dim]`.
-+  // aren't present in `outAxesPerDim`, and update the internal state as
-+  // follows: `gatheringAxes` are popped from the back of `inAxesPerDim[dim]`
-+  // and `currentAxesPerDim[dim]`.
-   //
-   // For example:
-   //
-@@ -376,18 +344,11 @@ class CollectiveInserter {
-     SmallVector<AxisRefAttr>& currentAxes = currentAxesPerDim[dim];
-     SmallVector<AxisRefAttr> gatheringAxes;
-     gatheringAxes.reserve(inAxes.size());
--    bool gatherAll = dimsToUnshard[dim];
--    auto shouldGatherAxis = [&](AxisRefAttr axis) {
--      if (gatherAll) {
--        return true;
--      }
--      auto it = outAxisToDimAndIndex.find(axis);
--      return it == outAxisToDimAndIndex.end() || dimsToUnshard[it->second.dim];
--    };
-     auto axisRevIt = inAxes.rbegin();
--    while (axisRevIt != inAxes.rend() && shouldGatherAxis(*axisRevIt)) {
-+    while (axisRevIt != inAxes.rend() &&
-+           !outAxisToDimAndIndex.contains(*axisRevIt)) {
-       inAxisSet.erase(*axisRevIt);
--      ++axisRevIt;
-+      --axisRevIt;
-     }
-     auto inAxisIt = axisRevIt.base();
-     popBackFromCurrentAxes(currentAxes, inAxes, inAxisIt);
-@@ -408,12 +369,6 @@ class CollectiveInserter {
-       }
-       collectiveAxes = AxisRefListAttr::get(getContext(), gatheringAxes);
-     }
--    for (auto [currentAxes, unshardDim] :
--         llvm::zip_equal(currentAxesPerDim, dimsToUnshard)) {
--      if (currentAxes.empty()) {
--        unshardDim = false;
--      }
--    }
-     if (hasGatheringAxes) {
-       result = AllGatherOp::create(rewriter, loc, result, collectiveAxesPerDim,
-                                    getCurrentSharding());
-@@ -620,10 +575,9 @@ class CollectiveInserter {
-   // We define the term capacity per dimension for all-slice - the product of
-   // axis sizes in `outAxesPerDim[d]` divided by the product of axis sizes in
-   // `inAxesPerDim[d]` for each dimension d, or the former divided by their GCD,
--  // if the former is not divisible by the latter (or 1 if `dimsToUnshard[d]` is
--  // true, since dimension d must be unsharded before slicing). In other words,
--  // the capacity represents how much the output is sharded more than the input
--  // along a specific dimension.
-+  // if the former is not divisible by the latter (as the dimension size needs
-+  // to be a multiplier of both). In other words, the capacity represents how
-+  // much the output is sharded more than the input along a specific dimension.
-   //
-   // Constraint: we can only slice dimension d along axes whose product of
-   // sizes doesn't exceed the capacity for that dimension, otherwise we will
-@@ -715,12 +669,8 @@ class CollectiveInserter {
-     bool hasSlicingAxes = false;
- 
-     // Initialize capacity per dimension.
--    for (auto [inAxes, outAxes, dimCapacity, unshardDim] : llvm::zip_equal(
--             inAxesPerDim, outAxesPerDim, capacityPerDim, dimsToUnshard)) {
--      if (unshardDim) {
--        dimCapacity = 1;
--        continue;
--      }
-+    for (auto [inAxes, outAxes, dimCapacity] :
-+         llvm::zip_equal(inAxesPerDim, outAxesPerDim, capacityPerDim)) {
-       int64_t inShardedSize = getShardedSize(inAxes, mesh);
-       int64_t outShardedSize = getShardedSize(outAxes, mesh);
-       dimCapacity = outShardedSize / std::gcd(inShardedSize, outShardedSize);
-@@ -748,13 +698,9 @@ class CollectiveInserter {
-     // 2. Slice axes in the dimension they appear in `outAxesPerDim`, i.e. the
-     // desired dimension, as long as the per-dim capacity allows.
-     AxisList availableOutAxes;
--    for (auto [inAxes, outAxes, currentAxes, slicingAxes, dimCapacity,
--               unshardDim] :
-+    for (auto [inAxes, outAxes, currentAxes, slicingAxes, dimCapacity] :
-          llvm::zip_equal(inAxesPerDim, outAxesPerDim, currentAxesPerDim,
--                         slicingAxesPerDim, capacityPerDim, dimsToUnshard)) {
--      if (unshardDim) {
--        continue;
--      }
-+                         slicingAxesPerDim, capacityPerDim)) {
-       auto outIt = outAxes.begin();
-       while (outIt != outAxes.end()) {
-         AxisRefAttr outAxis = *outIt;
-@@ -822,8 +768,7 @@ class CollectiveInserter {
-     }
-   }
- 
--  // We should insert a collective permute (assuming no dimension in
--  // `dimsToUnshard` still needs to be unsharded) if one of the following holds:
-+  // We should insert a collective permute if one of the following holds:
-   //
-   // 1. Both `inAxesPerDim[d]` and `outAxesPerDim[d]` aren't empty for a certain
-   //    dimension d and their front axes are divisible. This means we can
-@@ -856,9 +801,6 @@ class CollectiveInserter {
-   //    that B is before A, and A can be moved to d2 via an all-to-all before B
-   //    is all-gathered.
-   bool shouldCollectivePermute() {
--    if (llvm::any_of(dimsToUnshard, [](bool b) { return b; })) {
--      return false;
--    }
-     OptionalAxisRef availableInAxis;
-     OptionalAxisRef availableOutAxis;
-     for (auto [dim, inOutAxes] :
-@@ -1123,15 +1065,12 @@ class CollectiveInserter {
-   // The suffix of axes in `inAxesPerDim[srcDim]`, such that all axes within the
-   // suffix are mapped to the same dimension in `outAxisToDimAndIndex`
-   // (`allToAllAxes`), are all-to-all-ed with the mapped dimension as the target
--  // (`tgtDim`), provided `dimsToUnshard[tgtDim]` is false and (if
--  // `dimsToUnshard[srcDim]` is true) `allToAllAxes` includes all axes on
--  // `srcDim` so that `srcDim` becomes completely unsharded.
-+  // (`tgtDim`).
-   //
-   // The internal state is updated as follows for `allToAllAxes` and `tgtDim`:
-   //
-   // - `allToAllAxes` are popped from the back of `inAxesPerDim[srcDim]` and
--  //   `currentAxesPerDim[srcDim]` (and `dimsToUnshard[srcDim]` is cleared if
--  //   `currentAxesPerDim[srcDim]` becomes empty).
-+  //   `currentAxesPerDim[srcDim]`.
-   // - `allToAllAxes` are appended to `inAxesPerDim[tgtDim]` and
-   //   `currentAxesPerDim[tgtDim]`.
-   // - The common prefix between `inAxesPerDim[tgtDim]` and
-@@ -1198,8 +1137,7 @@ class CollectiveInserter {
-         break;
-       }
-       int64_t outAxisDim = outAxisEntryIt->second.dim;
--      if (outAxisDim == srcDim || dimsToUnshard[outAxisDim] ||
--          (tgtDim && outAxisDim != *tgtDim)) {
-+      if (outAxisDim == srcDim || (tgtDim && outAxisDim != *tgtDim)) {
-         break;
-       }
-       tgtDim = outAxisDim;
-@@ -1212,9 +1150,6 @@ class CollectiveInserter {
-     }
- 
-     auto startInAxisIt = axisRevIt.base();
--    if (dimsToUnshard[srcDim] && startInAxisIt != srcInAxes.begin()) {
--      return std::nullopt;
--    }
- 
-     AxisList& tgtOutAxes = outAxesPerDim[*tgtDim];
-     if (!allowOutOfOrderTarget &&
-@@ -1231,9 +1166,6 @@ class CollectiveInserter {
-     SmallVector<AxisRefAttr>& tgtCurrentAxes = currentAxesPerDim[*tgtDim];
- 
-     popBackFromCurrentAxes(srcCurrentAxes, srcInAxes, startInAxisIt);
--    if (srcCurrentAxes.empty()) {
--      dimsToUnshard[srcDim] = false;
--    }
- 
-     AxisList& tgtInAxes = inAxesPerDim[*tgtDim];
-     auto srcInAxisIt = startInAxisIt;
-@@ -1280,7 +1212,6 @@ class CollectiveInserter {
-   AxesPerDim currentAxesPerDim;
-   SmallVector<int64_t> capacityPerDim;
-   SmallVector<AxisRefListAttr> collectiveAxesPerDim;
--  SmallVector<bool> dimsToUnshard;
-   AxisSet inAxisSet;
-   AxisToDimAndIndex outAxisToDimAndIndex;
- };
-diff --git a/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/generic_ops.mlir b/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/generic_ops.mlir
-index 9bd4e0a8..4d4faae8 100644
---- a/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/generic_ops.mlir
-+++ b/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/generic_ops.mlir
-@@ -12,45 +12,3 @@ func.func @no_pad(
-   %0 = stablehlo.add %arg0, %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4_2, [{"x"}, {}]>]>} : tensor<4x8xf32>
-   return %0 : tensor<4x8xf32>
- }
--
--sdy.mesh @mesh_4_2_4 = <["x"=4, "y"=2, "z"=4]>
--
--// CHECK-LABEL: func private @pad_all_reduce
--func.func private @pad_all_reduce(
--  %arg0: tensor<7x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{"x"}, {}], unreduced={"y"}>})
--  -> (tensor<7x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{"x"}, {}]>}) {
--  // CHECK-NEXT: %[[AR:.*]] = sdy.all_reduce {"y"} %arg0 out_sharding=<@mesh_4_2, [{"x"}, {}]> : tensor<8x8xf32>
--  // CHECK-NEXT: return %[[AR]] : tensor<8x8xf32>
--  %0 = sdy.all_reduce {"y"} %arg0 out_sharding=<@mesh_4_2, [{"x"}, {}]> : tensor<7x8xf32>
--  return %0 : tensor<7x8xf32>
--}
--
--// CHECK-LABEL: func private @pad_collective_permute
--func.func private @pad_collective_permute(
--  %arg0: tensor<7x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2_4, [{"x"}, {}]>})
--  -> (tensor<7x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2_4, [{"z"}, {}]>}) {
--  // CHECK-NEXT: %[[CP:.*]] = sdy.collective_permute %arg0 out_sharding=<@mesh_4_2_4, [{"z"}, {}]> : tensor<8x8xf32>
--  // CHECK-NEXT: return %[[CP]] : tensor<8x8xf32>
--  %0 = sdy.collective_permute %arg0 out_sharding=<@mesh_4_2_4, [{"z"}, {}]> : tensor<7x8xf32>
--  return %0 : tensor<7x8xf32>
--}
--
--// CHECK-LABEL: func private @pad_sharded_to_unreduced
--func.func private @pad_sharded_to_unreduced(
--  %arg0: tensor<7x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{"x"}, {"y"}]>})
--  -> (tensor<7x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{"x"}, {}], unreduced={"y"}>}) {
--  // CHECK-NEXT: %[[S2U:.*]] = sdy.sharded_to_unreduced [{}, {"y"}] %arg0 out_sharding=<@mesh_4_2, [{"x"}, {}], unreduced={"y"}> : tensor<8x8xf32>
--  // CHECK-NEXT: return %[[S2U]] : tensor<8x8xf32>
--  %0 = sdy.sharded_to_unreduced [{}, {"y"}] %arg0 out_sharding=<@mesh_4_2, [{"x"}, {}], unreduced={"y"}> : tensor<7x8xf32>
--  return %0 : tensor<7x8xf32>
--}
--
--// CHECK-LABEL: func private @pad_replicated_to_unreduced
--func.func private @pad_replicated_to_unreduced(
--  %arg0: tensor<7x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{"x"}, {}]>})
--  -> (tensor<7x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{"x"}, {}], unreduced={"y"}>}) {
--  // CHECK-NEXT: %[[R2U:.*]] = sdy.replicated_to_unreduced {"y"} %arg0 out_sharding=<@mesh_4_2, [{"x"}, {}], unreduced={"y"}> : tensor<8x8xf32>
--  // CHECK-NEXT: return %[[R2U]] : tensor<8x8xf32>
--  %0 = sdy.replicated_to_unreduced {"y"} %arg0 out_sharding=<@mesh_4_2, [{"x"}, {}], unreduced={"y"}> : tensor<7x8xf32>
--  return %0 : tensor<7x8xf32>
--}
-diff --git a/shardy/dialect/sdy/transforms/export/test/reshard_to_collectives.mlir b/shardy/dialect/sdy/transforms/export/test/reshard_to_collectives.mlir
-index a63c54d1..f185669a 100644
---- a/shardy/dialect/sdy/transforms/export/test/reshard_to_collectives.mlir
-+++ b/shardy/dialect/sdy/transforms/export/test/reshard_to_collectives.mlir
-@@ -644,12 +644,12 @@ func.func @all_to_all_axes_at_src_out_of_order(%arg0 : tensor<16x8x8xf32> {sdy.s
- }
- 
- // CHECK-LABEL: func @all_to_all_axes_at_src_and_tgt_out_of_order
--func.func @all_to_all_axes_at_src_and_tgt_out_of_order(%arg0 : tensor<32x8x8xf32> {sdy.sharding = #sdy.sharding<@mesh3d_4x2x4, [{"z"}, {"y", "x"}, {}]>}) -> tensor<32x8x8xf32> {
-+func.func @all_to_all_axes_at_src_and_tgt_out_of_order(%arg0 : tensor<16x8x8xf32> {sdy.sharding = #sdy.sharding<@mesh3d_4x2x4, [{"z"}, {"y", "x"}, {}]>}) -> tensor<16x8x8xf32> {
-   // CHECK-NEXT: %[[COLLECTIVE_PERMUTE:.*]] = sdy.collective_permute %arg0 out_sharding=<@mesh3d_4x2x4, [{"x"}, {"y", "z"}, {}]>
-   // CHECK-NEXT: %[[ALL_TO_ALL:.*]] = sdy.all_to_all [{"y", "z"}: 1->0] %[[COLLECTIVE_PERMUTE]] out_sharding=<@mesh3d_4x2x4, [{"x", "y", "z"}, {}, {}]>
-   // CHECK-NEXT: return %[[ALL_TO_ALL]]
--  %0 = sdy.reshard %arg0 <@mesh3d_4x2x4, [{"x", "y", "z"}, {}, {}]> : tensor<32x8x8xf32>
--  return %0 : tensor<32x8x8xf32>
-+  %0 = sdy.reshard %arg0 <@mesh3d_4x2x4, [{"x", "y", "z"}, {}, {}]> : tensor<16x8x8xf32>
-+  return %0 : tensor<16x8x8xf32>
- }
- 
- // CHECK-LABEL: func @all_to_all_two_tgt_dims_src_out_of_order
-@@ -935,60 +935,6 @@ func.func @reshard_with_propagation_barrier(
-   return %1 : tensor<8x8x8xf32>
- }
- 
--// CHECK-LABEL: func @reshard_uneven_padded_size_mismatch_common_prefix
--func.func @reshard_uneven_padded_size_mismatch_common_prefix(%arg0 : tensor<6x3xf32> {sdy.sharding = #sdy.sharding<@mesh2d_4x2, [{"x":(1)2}, {}]>}) -> tensor<6x3xf32> {
--  // CHECK-NEXT: %[[ALL_GATHER:.*]] = sdy.all_gather [{"x":(1)2}, {}] %arg0 out_sharding=<@mesh2d_4x2, [{}, {}]> : tensor<6x3xf32>
--  // CHECK-NEXT: %[[ALL_SLICE:.*]] = sdy.all_slice [{"x"}, {}] %[[ALL_GATHER]] out_sharding=<@mesh2d_4x2, [{"x"}, {}]> : tensor<6x3xf32>
--  // CHECK-NEXT: return %[[ALL_SLICE]] : tensor<6x3xf32>
--  %0 = sdy.reshard %arg0 <@mesh2d_4x2, [{"x"}, {}]> : tensor<6x3xf32>
--  return %0 : tensor<6x3xf32>
--}
--
--// CHECK-LABEL: func @reshard_uneven_padded_size_mismatch_subaxis
--func.func @reshard_uneven_padded_size_mismatch_subaxis(%arg0 : tensor<6x3xf32> {sdy.sharding = #sdy.sharding<@mesh2d_2x8, [{"y":(1)2}, {}]>}) -> tensor<6x3xf32> {
--  // CHECK-NEXT: %[[ALL_GATHER:.*]] = sdy.all_gather [{"y":(1)2}, {}] %arg0 out_sharding=<@mesh2d_2x8, [{}, {}]> : tensor<6x3xf32>
--  // CHECK-NEXT: %[[ALL_SLICE:.*]] = sdy.all_slice [{"y":(2)4}, {}] %[[ALL_GATHER]] out_sharding=<@mesh2d_2x8, [{"y":(2)4}, {}]> : tensor<6x3xf32>
--  // CHECK-NEXT: return %[[ALL_SLICE]] : tensor<6x3xf32>
--  %0 = sdy.reshard %arg0 <@mesh2d_2x8, [{"y":(2)4}, {}]> : tensor<6x3xf32>
--  return %0 : tensor<6x3xf32>
--}
--
--// CHECK-LABEL: func @reshard_uneven_padded_size_mismatch_all_to_all
--func.func @reshard_uneven_padded_size_mismatch_all_to_all(%arg0 : tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh1d_6, [{}, {"x":(1)2}]>}) -> tensor<8x8xf32> {
--  // CHECK-NEXT: %[[ALL_TO_ALL:.*]] = sdy.all_to_all [{"x":(1)2}: 1->0] %arg0 out_sharding=<@mesh1d_6, [{"x":(1)2}, {}]> : tensor<8x8xf32>
--  // CHECK-NEXT: %[[ALL_SLICE:.*]] = sdy.all_slice [{}, {"x":(2)3}] %[[ALL_TO_ALL]] out_sharding=<@mesh1d_6, [{"x":(1)2}, {"x":(2)3}]> : tensor<8x8xf32>
--  // CHECK-NEXT: return %[[ALL_SLICE]] : tensor<8x8xf32>
--  %0 = sdy.reshard %arg0 <@mesh1d_6, [{"x":(1)2}, {"x":(2)3}]> : tensor<8x8xf32>
--  return %0 : tensor<8x8xf32>
--}
--
--// CHECK-LABEL: func @reshard_uneven_padded_size_mismatch_all_to_all_src_and_tgt_out_of_order
--func.func @reshard_uneven_padded_size_mismatch_all_to_all_src_and_tgt_out_of_order(%arg0 : tensor<16x8x8xf32> {sdy.sharding = #sdy.sharding<@mesh3d_4x2x4, [{"z"}, {"y", "x"}, {}]>}) -> tensor<16x8x8xf32> {
--  // CHECK-NEXT: %[[ALL_GATHER:.*]] = sdy.all_gather [{"z"}, {"y", "x"}, {}] %arg0 out_sharding=<@mesh3d_4x2x4, [{}, {}, {}]> : tensor<16x8x8xf32>
--  // CHECK-NEXT: %[[ALL_SLICE:.*]] = sdy.all_slice [{"x", "y", "z"}, {}, {}] %[[ALL_GATHER]] out_sharding=<@mesh3d_4x2x4, [{"x", "y", "z"}, {}, {}]> : tensor<16x8x8xf32>
--  // CHECK-NEXT: return %[[ALL_SLICE]] : tensor<16x8x8xf32>
--  %0 = sdy.reshard %arg0 <@mesh3d_4x2x4, [{"x", "y", "z"}, {}, {}]> : tensor<16x8x8xf32>
--  return %0 : tensor<16x8x8xf32>
--}
--
--// CHECK-LABEL: func @reshard_uneven_padded_size_mismatch_swap_axes
--func.func @reshard_uneven_padded_size_mismatch_swap_axes(%arg0 : tensor<5x5xf32> {sdy.sharding = #sdy.sharding<@mesh2d_4x2, [{"x"}, {"y"}]>}) -> tensor<5x5xf32> {
--  // CHECK-NEXT: %[[ALL_GATHER:.*]] = sdy.all_gather [{"x"}, {"y"}] %arg0 out_sharding=<@mesh2d_4x2, [{}, {}]> : tensor<5x5xf32>
--  // CHECK-NEXT: %[[ALL_SLICE:.*]] = sdy.all_slice [{"y"}, {"x"}] %[[ALL_GATHER]] out_sharding=<@mesh2d_4x2, [{"y"}, {"x"}]> : tensor<5x5xf32>
--  // CHECK-NEXT: return %[[ALL_SLICE]] : tensor<5x5xf32>
--  %0 = sdy.reshard %arg0 <@mesh2d_4x2, [{"y"}, {"x"}]> : tensor<5x5xf32>
--  return %0 : tensor<5x5xf32>
--}
--
--// CHECK-LABEL: func @reshard_uneven_padded_size_match_swap_axes
--func.func @reshard_uneven_padded_size_match_swap_axes(%arg0 : tensor<19x19xf32> {sdy.sharding = #sdy.sharding<@mesh2d_4x2, [{"x"}, {"y"}]>}) -> tensor<19x19xf32> {
--  // CHECK-NEXT: %[[CP:.*]] = sdy.collective_permute %arg0 out_sharding=<@mesh2d_4x2, [{"y", "x":(2)2}, {"x":(1)2}]> : tensor<19x19xf32>
--  // CHECK-NEXT: %[[ALL_TO_ALL:.*]] = sdy.all_to_all [{"x":(2)2}: 0->1] %[[CP]] out_sharding=<@mesh2d_4x2, [{"y"}, {"x"}]> : tensor<19x19xf32>
--  // CHECK-NEXT: return %[[ALL_TO_ALL]] : tensor<19x19xf32>
--  %0 = sdy.reshard %arg0 <@mesh2d_4x2, [{"y"}, {"x"}]> : tensor<19x19xf32>
--  return %0 : tensor<19x19xf32>
--}
--
- //===----------------------------------------------------------------------===//
- // Single-Device Reshard Tests (Preserved as sdy.reshard)
- //===----------------------------------------------------------------------===//
-@@ -1010,3 +956,5 @@ func.func @single_device_out_sharding(%arg0 : tensor<16x8xf32> {sdy.sharding = #
-   %0 = sdy.reshard %arg0 <@single_dev_0, []> : tensor<16x8xf32>
-   return %0 : tensor<16x8xf32>
- }
-+
-+
-diff --git a/shardy/dialect/sdy/transforms/export/update_non_divisible_input_output_shardings.cc b/shardy/dialect/sdy/transforms/export/update_non_divisible_input_output_shardings.cc
-index 97c93fbf..90bcb8d5 100644
---- a/shardy/dialect/sdy/transforms/export/update_non_divisible_input_output_shardings.cc
-+++ b/shardy/dialect/sdy/transforms/export/update_non_divisible_input_output_shardings.cc
-@@ -17,7 +17,6 @@ limitations under the License.
- #include <cstdint>
- #include <functional>
- #include <numeric>
--#include <string_view>
- 
- #include "llvm/ADT/ArrayRef.h"
- #include "llvm/ADT/STLExtras.h"
-@@ -35,7 +34,6 @@ limitations under the License.
- #include "mlir/Pass/Pass.h"  // IWYU pragma: keep
- #include "mlir/Support/LLVM.h"
- #include "mlir/Transforms/DialectConversion.h"
--#include "shardy/common/logging.h"
- #include "shardy/dialect/sdy/ir/dialect.h"
- #include "shardy/dialect/sdy/ir/utils.h"
- 
-@@ -67,8 +65,7 @@ TensorShardingAttr getEvenlySharded(TensorShardingAttr sharding,
-                                     ShapedType type,
-                                     const SymbolTable& symbolTable) {
-   MeshAttr mesh = sharding.getMesh(symbolTable);
--  SDY_CHECK(mesh) << "unknown mesh: "
--                  << std::string_view(sharding.getMeshName());
-+  assert(mesh && "unknown mesh");
-   MLIRContext* ctx = sharding.getContext();
-   llvm::SmallVector<DimensionShardingAttr> newDimShardings;
-   newDimShardings.reserve(sharding.getRank());
-diff --git a/shardy/dialect/sdy/transforms/import/manual_axes_cleanup.cc b/shardy/dialect/sdy/transforms/import/manual_axes_cleanup.cc
-index 6b368f91..f71266a5 100644
---- a/shardy/dialect/sdy/transforms/import/manual_axes_cleanup.cc
-+++ b/shardy/dialect/sdy/transforms/import/manual_axes_cleanup.cc
-@@ -29,7 +29,6 @@ limitations under the License.
- #include "mlir/IR/SymbolTable.h"
- #include "mlir/IR/TypeRange.h"
- #include "mlir/Support/LLVM.h"
--#include "shardy/common/logging.h"
- #include "shardy/dialect/sdy/ir/dialect.h"
- #include "shardy/dialect/sdy/ir/utils.h"
- #include "shardy/dialect/sdy/transforms/import/passes.h"  // IWYU pragma: keep
-@@ -163,7 +162,7 @@ struct ManualAxesCleanupPass
-           getCommonMeshOrRef(inShardings, outShardings, symbolTable);
-       MeshAttr mesh =
-           meshOrRef ? getMeshOrLookup(symbolTable, meshOrRef) : nullptr;
--      SDY_CHECK(mesh) << "expected inputs and outputs to have a common mesh";
-+      assert(mesh && "expected inputs and outputs to have a common mesh");
-       sortManualAxes(op, mesh);
-       addUnusedManualAxesToReplicatedAxes(op, mesh, meshOrRef, symbolTable);
-     });
-diff --git a/shardy/dialect/sdy/transforms/import/remove_size_one_axes.cc b/shardy/dialect/sdy/transforms/import/remove_size_one_axes.cc
-index ce2abb40..99a47092 100644
---- a/shardy/dialect/sdy/transforms/import/remove_size_one_axes.cc
-+++ b/shardy/dialect/sdy/transforms/import/remove_size_one_axes.cc
-@@ -17,14 +17,12 @@ limitations under the License.
- #include <cstdint>
- #include <iterator>
- #include <optional>
--#include <string_view>
- 
- #include "llvm/ADT/STLExtras.h"
- #include "llvm/ADT/SmallVector.h"
- #include "mlir/IR/BuiltinOps.h"
- #include "mlir/IR/SymbolTable.h"
- #include "mlir/Support/LLVM.h"
--#include "shardy/common/logging.h"
- #include "shardy/dialect/sdy/ir/dialect.h"
- #include "shardy/dialect/sdy/transforms/common/sharding_walker.h"
- #include "shardy/dialect/sdy/transforms/import/passes.h"  // IWYU pragma: keep
-@@ -45,8 +43,7 @@ bool hasSizeOneAxes(MeshOp meshOp) {
- TensorShardingAttr removeSizeOneAxes(TensorShardingAttr sharding,
-                                      const SymbolTable& symbolTable) {
-   MeshAttr mesh = sharding.getMesh(symbolTable);
--  SDY_CHECK(mesh) << "unknown mesh: "
--                  << std::string_view(sharding.getMeshName());
-+  assert(mesh && "unknown mesh");
- 
-   auto isNotSizeOne = [&](AxisRefAttr axis) { return axis.getSize(mesh) != 1; };
- 
-diff --git a/shardy/dialect/sdy/transforms/import/test/sharding_group_constraints.mlir b/shardy/dialect/sdy/transforms/import/test/sharding_group_constraints.mlir
-index 34db7665..a3518d30 100644
---- a/shardy/dialect/sdy/transforms/import/test/sharding_group_constraints.mlir
-+++ b/shardy/dialect/sdy/transforms/import/test/sharding_group_constraints.mlir
-@@ -310,189 +310,3 @@ func.func private @bar(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
-   func.return %arg0 : tensor<8x8xf32>
- }
- 
--// -----
--
--sdy.mesh @mesh = <["a"=2, "b"=2]>
--
--// Allows sharding groups within ManualComputationOp when group members are
--// across multiple functions called from the same ManualComputationOp.
--func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
--    sdy.sharding_group %arg1 group_id = 1234 : tensor<8x8xf32>
--    %1 = func.call @foo(%arg1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--    %2 = func.call @bar(%1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--    sdy.return %2 : tensor<8x8xf32>
--  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  func.return %0: tensor<8x8xf32>
--}
--
--func.func private @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  sdy.sharding_group %arg0 group_id = 1234 : tensor<8x8xf32>
--  func.return %arg0 : tensor<8x8xf32>
--}
--
--func.func private @bar(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  sdy.sharding_group %arg0 group_id = 1234 : tensor<8x8xf32>
--  func.return %arg0 : tensor<8x8xf32>
--}
--
--// -----
--
--sdy.mesh @mesh = <["a"=2, "b"=2]>
--
--// Allows sharding groups within ManualComputationOp when a function with group
--// members is called multiple times from the same ManualComputationOp.
--func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
--    sdy.sharding_group %arg1 group_id = 1234 : tensor<8x8xf32>
--    %1 = func.call @foo(%arg1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--    %2 = func.call @foo(%1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--    sdy.return %2 : tensor<8x8xf32>
--  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  func.return %0: tensor<8x8xf32>
--}
--
--func.func private @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  sdy.sharding_group %arg0 group_id = 1234 : tensor<8x8xf32>
--  func.return %arg0 : tensor<8x8xf32>
--}
--
--// -----
--
--sdy.mesh @mesh = <["a"=2, "b"=2]>
--
--// TODO(b/569928811): Disallow sharding groups which cross the barrier of a
--// ManualComputationOp when the function is called inside and then outside the
--// ManualComputationOp.
--func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
--    %1 = func.call @foo(%arg1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--    sdy.return %1 : tensor<8x8xf32>
--  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
--
--  %2 = func.call @foo(%0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  func.return %2: tensor<8x8xf32>
--}
--
--func.func private @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  sdy.sharding_group %arg0 group_id = 1234 : tensor<8x8xf32>
--  func.return %arg0 : tensor<8x8xf32>
--}
--
--// -----
--
--sdy.mesh @mesh = <["a"=2, "b"=2]>
--
--// TODO(b/569928811): Disallow sharding groups which cross the barrier of a
--// ManualComputationOp when the function is called inside and then outside the
--// ManualComputationOp.
--func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
--    %1 = func.call @foo(%arg1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--    sdy.return %1 : tensor<8x8xf32>
--  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
--
--  %2 = func.call @foo(%0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  sdy.sharding_group %2 group_id = 1234 : tensor<8x8xf32>
--  func.return %2: tensor<8x8xf32>
--}
--
--func.func private @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  sdy.sharding_group %arg0 group_id = 1234 : tensor<8x8xf32>
--  func.return %arg0 : tensor<8x8xf32>
--}
--
--// -----
--
--sdy.mesh @mesh = <["a"=2, "b"=2]>
--
--// TODO(b/569928811): Disallow sharding groups which cross the barrier of a
--// ManualComputationOp when the function is called outside and then inside the
--// ManualComputationOp.
--func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  %0 = func.call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  %1 = sdy.manual_computation(%0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
--    sdy.sharding_group %arg1 group_id = 1234 : tensor<8x8xf32>
--    %2 = func.call @foo(%arg1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--    sdy.return %2 : tensor<8x8xf32>
--  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  func.return %1: tensor<8x8xf32>
--}
--
--func.func private @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  sdy.sharding_group %arg0 group_id = 1234 : tensor<8x8xf32>
--  func.return %arg0 : tensor<8x8xf32>
--}
--
--// -----
--
--sdy.mesh @mesh = <["a"=2, "b"=2]>
--
--// Allows sharding groups inside a function that is called from different
--// ManualComputationOps.
--func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
--    %1 = func.call @foo(%arg1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--    sdy.return %1 : tensor<8x8xf32>
--  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
--
--  %2 = sdy.manual_computation(%0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
--    %3 = func.call @foo(%arg1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--    sdy.return %3 : tensor<8x8xf32>
--  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  func.return %2: tensor<8x8xf32>
--}
--
--func.func private @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  sdy.sharding_group %arg0 group_id = 1234 : tensor<8x8xf32>
--  func.return %arg0 : tensor<8x8xf32>
--}
--
--// -----
--
--sdy.mesh @mesh = <["a"=2, "b"=2]>
--
--// Allows sharding groups inside a ManualComputationOp within a function that
--// is called multiple times.
--func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  %0 = func.call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  %1 = func.call @foo(%0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  func.return %1: tensor<8x8xf32>
--}
--
--func.func private @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
--    sdy.sharding_group %arg1 group_id = 1234 : tensor<8x8xf32>
--    sdy.return %arg1 : tensor<8x8xf32>
--  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  func.return %0 : tensor<8x8xf32>
--}
--
--// -----
--
--sdy.mesh @mesh = <["a"=2, "b"=2]>
--
--// Disallows sharding groups which have different ManualComputationOp parents
--// across different functions.
--func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  %0 = func.call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  %1 = func.call @bar(%0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  func.return %1: tensor<8x8xf32>
--}
--
--func.func private @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
--    sdy.sharding_group %arg1 group_id = 1234 : tensor<8x8xf32>
--    sdy.return %arg1 : tensor<8x8xf32>
--  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  func.return %0 : tensor<8x8xf32>
--}
--
--func.func private @bar(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{?}, {?}]>] out_shardings=[<@mesh, [{?}, {?}]>] manual_axes={} (%arg1: tensor<8x8xf32>) {
--    // expected-error@below {{ShardingGroupOps values cannot cross ManualComputationOp boundaries for groupId: 1234}}
--    sdy.sharding_group %arg1 group_id = 1234 : tensor<8x8xf32>
--    sdy.return %arg1 : tensor<8x8xf32>
--  } : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  func.return %0 : tensor<8x8xf32>
--}
-diff --git a/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc b/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc
-index 8417874d..2c6b451f 100644
---- a/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc
-+++ b/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc
-@@ -20,7 +20,6 @@ limitations under the License.
- #include <functional>
- #include <memory>
- #include <optional>
--#include <string_view>
- #include <utility>
- 
- #include "llvm/ADT/STLExtras.h"
-@@ -246,11 +245,6 @@ bool updateTensorSharding(Value modifiedValue, const SymbolTable& symbolTable,
-     if (groupValue == modifiedValue) {
-       continue;
-     }
--    TensorShardingAttr oldGroupSharding = getSharding(groupValue);
--    if ((!oldGroupSharding && newSharding.emptyAxes()) ||
--        newSharding == oldGroupSharding) {
--      continue;
--    }
-     setSharding(groupValue, newSharding);
-     if (params.notifyOpModified) {
-       notifyShardingModified(groupValue, symbolTable, userMap,
-@@ -319,7 +313,7 @@ LogicalResult propagateTensorShardings(
-     PropagationDirectionAlongFactor directionAlongFactor,
-     const FactorPropagation& factorPropagation, bool conservativePropagation,
-     Operation* op, const SymbolTable& symbolTable, PatternRewriter* rewriter,
--    const ShardingGroupMap& shardingGroupMap) {
-+    ShardingGroupMap shardingGroupMap) {
-   std::optional<StringRef> meshName =
-       getCommonMeshName(operandsParams.shardings, resultsParams.shardings,
-                         symbolTable, /*ignoreDeviceIds=*/false);
-@@ -333,8 +327,8 @@ LogicalResult propagateTensorShardings(
-     }
-     return failure();
-   }
--  MeshAttr mesh = getMeshAttr(symbolTable, meshName.value());
--  SDY_CHECK(mesh) << "unknown mesh: " << std::string_view(meshName.value());
-+  MeshAttr mesh = getMeshAttr(op, meshName.value());
-+  assert(mesh && "unknown mesh");
-   if (mesh.isSingleDevice()) {
-     // Single-device meshes and shardings are usually a placeholder for special
-     // operations, such as send and recv. We do not propagate single-device
-diff --git a/shardy/dialect/sdy/transforms/propagation/user_priority_propagation.cc b/shardy/dialect/sdy/transforms/propagation/user_priority_propagation.cc
-index f3eb2fc8..34e741d0 100644
---- a/shardy/dialect/sdy/transforms/propagation/user_priority_propagation.cc
-+++ b/shardy/dialect/sdy/transforms/propagation/user_priority_propagation.cc
-@@ -19,7 +19,6 @@ limitations under the License.
- #include <cstdint>
- #include <iterator>
- #include <memory>
--#include <string_view>
- #include <utility>
- #include <variant>
- 
-@@ -37,7 +36,6 @@ limitations under the License.
- #include "mlir/Support/LLVM.h"
- #include "mlir/Support/LogicalResult.h"
- #include "shardy/common/file_utils.h"
--#include "shardy/common/logging.h"
- #include "shardy/dialect/sdy/ir/dialect.h"
- #include "shardy/dialect/sdy/transforms/common/propagation_options.h"
- #include "shardy/dialect/sdy/transforms/common/sharding_walker.h"
-@@ -142,8 +140,7 @@ TensorShardingAttr getInitializedSharding(TensorShardingAttr originalSharding,
-     }
-   }
-   MeshAttr mesh = originalSharding.getMesh(symbolTable);
--  SDY_CHECK(mesh) << "unknown mesh: "
--                  << std::string_view(originalSharding.getMeshName());
-+  assert(mesh && "unknown mesh");
-   llvm::sort(newReplicatedAxes, AxisRefAttr::getMeshComparator(mesh));
-   // TODO(tomnatan): we need to merge split axes and split them again when
-   // updating? or can we assume we won't see split axes?
 diff --git a/third_party/llvm/generated.patch b/third_party/llvm/generated.patch
-index e2237d49..509398da 100644
+index 509398da..e9c206ce 100644
 --- a/third_party/llvm/generated.patch
 +++ b/third_party/llvm/generated.patch
-@@ -1,459 +1 @@
+@@ -1 +1,12 @@
  Auto generated patch. Do not edit or delete it, even if empty.
--diff -ruN --strip-trailing-cr a/llvm/include/llvm/IR/Metadata.h b/llvm/include/llvm/IR/Metadata.h
----- a/llvm/include/llvm/IR/Metadata.h
--+++ b/llvm/include/llvm/IR/Metadata.h
--@@ -397,18 +397,8 @@
--   using OwnerTy = MetadataTracking::OwnerTy;
-- 
-- private:
---  struct UseEntry {
---    void *Ref = nullptr;
---    OwnerTy Owner = nullptr;
---  };
---
---  static constexpr unsigned IndexThreshold = 32;
---  // Tracked uses of this. Dropping one moves the last entry into its slot, so
---  // the order is deterministic but not the order they were added in.
---  SmallVector<UseEntry, 4> UseMap;
---  // Lazily allocated map from Ref to its index in UseMap for large use lists.
---  using IndexMapTy = DenseMap<void *, unsigned>;
---  std::unique_ptr<IndexMapTy> IndexMap;
--+  uint64_t NextIndex = 0;
--+  SmallDenseMap<void *, std::pair<OwnerTy, uint64_t>, 4> UseMap;
-- 
-- protected:
--   ~ReplaceableUses() {
--@@ -439,7 +429,6 @@
--   unsigned getNumUses() const { return UseMap.size(); }
-- 
-- private:
---  UseEntry *findRef(void *Ref);
--   void addRef(void *Ref, OwnerTy Owner);
--   void dropRef(void *Ref);
--   void moveRef(void *Ref, void *New, const Metadata &MD);
--diff -ruN --strip-trailing-cr a/llvm/lib/IR/Metadata.cpp b/llvm/lib/IR/Metadata.cpp
----- a/llvm/lib/IR/Metadata.cpp
--+++ b/llvm/lib/IR/Metadata.cpp
--@@ -258,96 +258,82 @@
-- }
-- 
-- SmallVector<Metadata *> ReplaceableUses::getAllArgListUsers() {
---  SmallVector<Metadata *> MDUsers;
---  for (const auto &[Ref, Owner] : UseMap) {
--+  SmallVector<std::pair<OwnerTy, uint64_t> *> MDUsersWithID;
--+  for (auto Pair : UseMap) {
--+    OwnerTy Owner = Pair.second.first;
--     if (Owner.isNull())
--       continue;
--     if (!isa<Metadata *>(Owner))
--       continue;
--     Metadata *OwnerMD = cast<Metadata *>(Owner);
--     if (OwnerMD->getMetadataID() == Metadata::DIArgListKind)
---      MDUsers.push_back(OwnerMD);
--+      MDUsersWithID.push_back(&UseMap[Pair.first]);
--   }
--+  llvm::sort(MDUsersWithID, [](auto UserA, auto UserB) {
--+    return UserA->second < UserB->second;
--+  });
--+  SmallVector<Metadata *> MDUsers;
--+  for (auto *UserWithID : MDUsersWithID)
--+    MDUsers.push_back(cast<Metadata *>(UserWithID->first));
--   return MDUsers;
-- }
-- 
-- SmallVector<DbgVariableRecord *>
-- ReplaceableUses::getAllDbgVariableRecordUsers() {
--+  SmallVector<std::pair<OwnerTy, uint64_t> *> DVRUsersWithID;
--+  for (auto Pair : UseMap) {
--+    OwnerTy Owner = Pair.second.first;
--+    if (Owner.isNull())
--+      continue;
--+    if (!isa<DebugValueUser *>(Owner))
--+      continue;
--+    DVRUsersWithID.push_back(&UseMap[Pair.first]);
--+  }
--   // Order DbgVariableRecord users in reverse-creation order. Normal dbg.value
--   // users of MetadataAsValues are ordered by their UseList, i.e. reverse order
--   // of when they were added: we need to replicate that here. The structure of
--   // debug-info output depends on the ordering of intrinsics, thus we need
--   // to keep them consistent for comparisons sake.
--+  llvm::sort(DVRUsersWithID, [](auto UserA, auto UserB) {
--+    return UserA->second > UserB->second;
--+  });
--   SmallVector<DbgVariableRecord *> DVRUsers;
---  for (const auto &[Ref, Owner] : reverse(UseMap)) {
---    if (Owner.isNull())
---      continue;
---    if (!isa<DebugValueUser *>(Owner))
---      continue;
---    DVRUsers.push_back(cast<DebugValueUser *>(Owner)->getUser());
---  }
--+  for (auto UserWithID : DVRUsersWithID)
--+    DVRUsers.push_back(cast<DebugValueUser *>(UserWithID->first)->getUser());
--   return DVRUsers;
-- }
-- 
---ReplaceableUses::UseEntry *ReplaceableUses::findRef(void *Ref) {
---  if (!IndexMap) {
---    // Search backward for temporal locality; recently added references are
---    // often dropped or moved first.
---    for (UseEntry &U : reverse(UseMap))
---      if (U.Ref == Ref)
---        return &U;
---    return nullptr;
---  }
---
---  auto It = IndexMap->find(Ref);
---  return It == IndexMap->end() ? nullptr : &UseMap[It->second];
---}
---
-- void ReplaceableUses::addRef(void *Ref, OwnerTy Owner) {
---  assert(Ref && "Expected live reference");
---  assert(!findRef(Ref) && "Reference already tracked");
---  unsigned NewIdx = UseMap.size();
---  UseMap.push_back({Ref, Owner});
---  if (IndexMap) {
---    (*IndexMap)[Ref] = NewIdx;
---  } else if (UseMap.size() > IndexThreshold) {
---    // Build the index map once UseMap grows past the threshold.
---    IndexMap = std::make_unique<IndexMapTy>();
---    for (unsigned I = 0, E = UseMap.size(); I != E; ++I)
---      (*IndexMap)[UseMap[I].Ref] = I;
---  }
--+  bool WasInserted =
--+      UseMap.insert(std::make_pair(Ref, std::make_pair(Owner, NextIndex)))
--+          .second;
--+  (void)WasInserted;
--+  assert(WasInserted && "Expected to add a reference");
--+
--+  ++NextIndex;
--+  assert(NextIndex != 0 && "Unexpected overflow");
-- }
-- 
-- void ReplaceableUses::dropRef(void *Ref) {
---  UseEntry *Entry = findRef(Ref);
---  assert(Entry && "Expected to find Ref");
---  if (IndexMap)
---    IndexMap->erase(Ref);
---  if (Entry != &UseMap.back()) {
---    *Entry = UseMap.back();
---    if (IndexMap)
---      (*IndexMap)[Entry->Ref] = Entry - UseMap.begin();
---  }
---  UseMap.pop_back();
---  if (IndexMap && UseMap.size() <= IndexThreshold / 2)
---    IndexMap.reset();
--+  bool WasErased = UseMap.erase(Ref);
--+  (void)WasErased;
--+  assert(WasErased && "Expected to drop a reference");
-- }
-- 
-- void ReplaceableUses::moveRef(void *Ref, void *New, const Metadata &MD) {
---  assert(!findRef(New) && "Cannot move to an existing reference");
---  UseEntry *Entry = findRef(Ref);
---  assert(Entry && "Expected to move a reference");
---  if (IndexMap) {
---    IndexMap->erase(Ref);
---    (*IndexMap)[New] = Entry - UseMap.begin();
---  }
---  Entry->Ref = New;
--+  auto I = UseMap.find(Ref);
--+  assert(I != UseMap.end() && "Expected to move a reference");
--+  auto OwnerAndIndex = I->second;
--+  UseMap.erase(I);
--+  bool WasInserted = UseMap.insert(std::make_pair(New, OwnerAndIndex)).second;
--+  (void)WasInserted;
--+  assert(WasInserted && "Expected to add a reference");
-- 
--   // Check that the references are direct if there's no owner.
--   (void)MD;
---  assert((Entry->Owner || *static_cast<Metadata **>(Ref) == &MD) &&
--+  assert((OwnerAndIndex.first || *static_cast<Metadata **>(Ref) == &MD) &&
--          "Reference without owner must be direct");
---  assert((Entry->Owner || *static_cast<Metadata **>(New) == &MD) &&
--+  assert((OwnerAndIndex.first || *static_cast<Metadata **>(New) == &MD) &&
--          "Reference without owner must be direct");
-- }
-- 
--@@ -360,11 +346,14 @@
--   auto &Store = Context.pImpl->ValuesAsMetadata;
--   auto I = Store.find(&C);
--   ValueAsMetadata *MD = I->second;
--+  using UseTy =
--+      std::pair<void *, std::pair<MetadataTracking::OwnerTy, uint64_t>>;
--   // Copy out uses and update value of Constant used by debug info metadata with
---  // poison below.
---  SmallVector<UseEntry, 4> Uses = MD->UseMap;
--+  // poison below
--+  SmallVector<UseTy, 8> Uses(MD->UseMap.begin(), MD->UseMap.end());
-- 
---  for (const auto &[Ref, Owner] : Uses) {
--+  for (const auto &Pair : Uses) {
--+    MetadataTracking::OwnerTy Owner = Pair.second.first;
--     if (!Owner)
--       continue;
--     // Check for MetadataAsValue.
--@@ -380,7 +369,7 @@
--       continue;
--     if (isa<DINode>(OwnerMD)) {
--       OwnerMD->handleChangedOperand(
---          Ref, ValueAsMetadata::get(PoisonValue::get(C.getType())));
--+          Pair.first, ValueAsMetadata::get(PoisonValue::get(C.getType())));
--     }
--   }
-- }
--@@ -390,21 +379,25 @@
--     return;
-- 
--   // Copy out uses since UseMap will get touched below.
---  SmallVector<UseEntry, 4> Uses = UseMap;
---
---  for (const auto &[Ref, Owner] : Uses) {
--+  using UseTy = std::pair<void *, std::pair<OwnerTy, uint64_t>>;
--+  SmallVector<UseTy, 8> Uses(UseMap.begin(), UseMap.end());
--+  llvm::sort(Uses, [](const UseTy &L, const UseTy &R) {
--+    return L.second.second < R.second.second;
--+  });
--+  for (const auto &Pair : Uses) {
--     // Check that this Ref hasn't disappeared after RAUW (when updating a
--     // previous Ref).
---    if (!findRef(Ref))
--+    if (!UseMap.count(Pair.first))
--       continue;
-- 
--+    OwnerTy Owner = Pair.second.first;
--     if (!Owner) {
--       // Update unowned tracking references directly.
---      Metadata *&DirectRef = *static_cast<Metadata **>(Ref);
---      dropRef(Ref);
---      DirectRef = MD;
--+      Metadata *&Ref = *static_cast<Metadata **>(Pair.first);
--+      Ref = MD;
--       if (MD)
---        MetadataTracking::track(DirectRef);
--+        MetadataTracking::track(Ref);
--+      UseMap.erase(Pair.first);
--       continue;
--     }
-- 
--@@ -415,7 +408,7 @@
--     }
-- 
--     if (auto *DVU = dyn_cast<DebugValueUser *>(Owner)) {
---      DVU->handleChangedValue(Ref, MD);
--+      DVU->handleChangedValue(Pair.first, MD);
--       continue;
--     }
-- 
--@@ -424,7 +417,7 @@
--     switch (OwnerMD->getMetadataID()) {
-- #define HANDLE_METADATA_LEAF(CLASS)                                            \
--   case Metadata::CLASS##Kind:                                                  \
---    cast<CLASS>(OwnerMD)->handleChangedOperand(Ref, MD);                       \
--+    cast<CLASS>(OwnerMD)->handleChangedOperand(Pair.first, MD);                \
--     continue;
-- #include "llvm/IR/Metadata.def"
--     default:
--@@ -440,16 +433,18 @@
-- 
--   if (!ResolveUsers) {
--     UseMap.clear();
---    IndexMap.reset();
--     return;
--   }
-- 
---  // Move uses out since UseMap could get touched below.
---  SmallVector<UseEntry, 4> Uses = std::move(UseMap);
--+  // Copy out uses since UseMap could get touched below.
--+  using UseTy = std::pair<void *, std::pair<OwnerTy, uint64_t>>;
--+  SmallVector<UseTy, 8> Uses(UseMap.begin(), UseMap.end());
--+  llvm::sort(Uses, [](const UseTy &L, const UseTy &R) {
--+    return L.second.second < R.second.second;
--+  });
--   UseMap.clear();
---  IndexMap.reset();
---  for (const auto &U : Uses) {
---    auto Owner = U.Owner;
--+  for (const auto &Pair : Uses) {
--+    auto Owner = Pair.second.first;
--     if (!Owner)
--       continue;
--     if (!isa<Metadata *>(Owner))
--diff -ruN --strip-trailing-cr a/llvm/lib/ObjCopy/ELF/ELFObject.cpp b/llvm/lib/ObjCopy/ELF/ELFObject.cpp
----- a/llvm/lib/ObjCopy/ELF/ELFObject.cpp
--+++ b/llvm/lib/ObjCopy/ELF/ELFObject.cpp
--@@ -1709,11 +1709,7 @@
--   case SHT_REL:
--   case SHT_RELA:
--   case SHT_CREL:
---    // SHF_ALLOC relocations of an executable or shared object are copied
---    // verbatim. Relocatable files relocations are usually non-ALLOC, but Linux
---    // livepatch modules set SHF_ALLOC on .klp.rela.* to keep these static
---    // relocations in memory.
---    if ((Shdr.sh_flags & SHF_ALLOC) && ElfFile.getHeader().e_type != ET_REL) {
--+    if (Shdr.sh_flags & SHF_ALLOC) {
--       if (Expected<ArrayRef<uint8_t>> Data = ElfFile.getSectionContents(Shdr))
--         return Obj.addSection<DynamicRelocationSection>(*Data);
--       else
--diff -ruN --strip-trailing-cr a/llvm/lib/ObjCopy/ELF/ELFObject.h b/llvm/lib/ObjCopy/ELF/ELFObject.h
----- a/llvm/lib/ObjCopy/ELF/ELFObject.h
--+++ b/llvm/lib/ObjCopy/ELF/ELFObject.h
--@@ -38,7 +38,6 @@
-- class SymbolTableSection;
-- class RelocationSection;
-- class DynamicRelocationSection;
---class DynamicSymbolTableSection;
-- class GnuDebugLinkSection;
-- class GroupSection;
-- class SectionIndexSection;
--@@ -876,15 +875,10 @@
-- // and another which handles the symbol table type. The symbol table type is
-- // taken as a type parameter to the class (see RelocSectionWithSymtabBase).
-- class RelocationSectionBase : public SectionBase {
---  const bool Dynamic;
---
-- protected:
--   SectionBase *SecToApplyRel = nullptr;
-- 
---  explicit RelocationSectionBase(bool Dynamic) : Dynamic(Dynamic) {}
---
-- public:
---  bool isDynamic() const { return Dynamic; }
--   const SectionBase *getSection() const { return SecToApplyRel; }
--   void setSection(SectionBase *Sec) { SecToApplyRel = Sec; }
-- 
--@@ -903,9 +897,7 @@
--   void setSymTab(SymTabType *SymTab) { Symbols = SymTab; }
-- 
-- protected:
---  RelocSectionWithSymtabBase()
---      : RelocationSectionBase(
---            std::is_same_v<SymTabType, DynamicSymbolTableSection>) {}
--+  RelocSectionWithSymtabBase() = default;
-- 
--   SymTabType *Symbols = nullptr;
-- 
--@@ -936,8 +928,9 @@
--   const Object &getObject() const { return Obj; }
-- 
--   static bool classof(const SectionBase *S) {
---    return RelocationSectionBase::classof(S) &&
---           !static_cast<const RelocationSectionBase *>(S)->isDynamic();
--+    if (S->OriginalFlags & ELF::SHF_ALLOC)
--+      return false;
--+    return RelocationSectionBase::classof(S);
--   }
-- };
-- 
--@@ -1022,8 +1015,9 @@
--       function_ref<bool(const SectionBase *)> ToRemove) override;
-- 
--   static bool classof(const SectionBase *S) {
---    return RelocationSectionBase::classof(S) &&
---           static_cast<const RelocationSectionBase *>(S)->isDynamic();
--+    if (!(S->OriginalFlags & ELF::SHF_ALLOC))
--+      return false;
--+    return S->OriginalType == ELF::SHT_REL || S->OriginalType == ELF::SHT_RELA;
--   }
-- };
-- 
--diff -ruN --strip-trailing-cr a/llvm/test/tools/llvm-objcopy/ELF/reloc-alloc-symtab.test b/llvm/test/tools/llvm-objcopy/ELF/reloc-alloc-symtab.test
----- a/llvm/test/tools/llvm-objcopy/ELF/reloc-alloc-symtab.test
--+++ b/llvm/test/tools/llvm-objcopy/ELF/reloc-alloc-symtab.test
--@@ -1,42 +0,0 @@
---## In a relocatable file, an SHF_ALLOC relocation section (e.g. .klp.rela.* in
---## Linux livepatch modules) is a static relocation section: its symbol indexes
---## are updated and the symbols it references are kept.
---
---# RUN: yaml2obj %s -o %t
---# RUN: llvm-strip --strip-unneeded %t -o %t1
---# RUN: llvm-readelf -Sr -s %t1 | FileCheck %s
---# RUN: yaml2obj -DTYPE=SHT_CREL %s -o %t.crel
---# RUN: llvm-strip --strip-unneeded %t.crel -o %t1.crel
---# RUN: llvm-readelf -Sr -s %t1.crel | FileCheck %s
---
---# CHECK:      [ 2] .klp.rela.vmlinux.text {{RELA|CREL}} {{.*}} AI  3   1
---# CHECK-NEXT: [ 3] .symtab
---# CHECK:      R_X86_64_PC32 {{.*}} referenced + 0
---# CHECK:      Symbol table '.symtab' contains 2 entries:
---# CHECK:      1: {{.*}} UND referenced
---
---!ELF
---FileHeader:
---  Class:   ELFCLASS64
---  Data:    ELFDATA2LSB
---  Type:    ET_REL
---  Machine: EM_X86_64
---Sections:
---  - Name:    .text
---    Type:    SHT_PROGBITS
---    Flags:   [ SHF_ALLOC, SHF_EXECINSTR ]
---    Content: "e800000000"
---  - Name:    .klp.rela.vmlinux.text
---    Type:    [[TYPE=SHT_RELA]]
---    Flags:   [ SHF_ALLOC, SHF_INFO_LINK ]
---    Link:    .symtab
---    Info:    .text
---    Relocations:
---      - Offset: 1
---        Symbol: referenced
---        Type:   R_X86_64_PC32
---Symbols:
---  - Name:    unreferenced
---    Section: .text
---  - Name:    referenced
---    Binding: STB_GLOBAL
--diff -ruN --strip-trailing-cr a/llvm/test/tools/llvm-objcopy/ELF/reloc-no-symtab.test b/llvm/test/tools/llvm-objcopy/ELF/reloc-no-symtab.test
----- a/llvm/test/tools/llvm-objcopy/ELF/reloc-no-symtab.test
--+++ b/llvm/test/tools/llvm-objcopy/ELF/reloc-no-symtab.test
--@@ -10,16 +10,11 @@
-- # RUN: llvm-readobj --sections %t.original | FileCheck %s
-- # RUN: llvm-readobj --sections %t.stripped | FileCheck %s
-- 
---## In a relocatable file, .rela.plt is a static relocation section.
---# RUN: yaml2obj -DTYPE=ET_REL %s -o %t.rel
---# RUN: llvm-strip %t.rel -o %t.rel.stripped
---# RUN: llvm-readobj --sections %t.rel.stripped | FileCheck %s
---
-- !ELF
-- FileHeader:
--   Class:           ELFCLASS64
--   Data:            ELFDATA2LSB
---  Type:            [[TYPE=ET_EXEC]]
--+  Type:            ET_REL
--   Machine:         EM_X86_64
-- Sections:
--   - Name:          .text
--diff -ruN --strip-trailing-cr a/llvm/unittests/IR/MetadataTest.cpp b/llvm/unittests/IR/MetadataTest.cpp
----- a/llvm/unittests/IR/MetadataTest.cpp
--+++ b/llvm/unittests/IR/MetadataTest.cpp
--@@ -5355,29 +5355,6 @@
--   EXPECT_TRUE(!MD);
-- }
-- 
---// Once the use list is indexed, look each changed reference up again right
---// away: a later swap in dropRef would repair a stale index entry.
---TEST_F(TrackingMDRefTest, LargeUseList) {
---  auto Temp = MDTuple::getTemporary(Context, {});
---  TrackingMDRef Refs[41];
---  for (unsigned I = 0; I != 40; ++I)
---    Refs[I].reset(Temp.get());
---  Refs[0].reset();
---  Refs[0].reset(Temp.get());
---  Refs[0].reset();
---  Refs[40] = std::move(Refs[1]);
---  Refs[1].reset(Temp.get());
---  Refs[40].reset();
---  Refs[39].reset();
---  EXPECT_EQ(38u, Temp->getNumTemporaryUses());
---
---  auto *N = MDTuple::getDistinct(Context, {});
---  Temp->replaceAllUsesWith(N);
---  EXPECT_EQ(0u, Temp->getNumTemporaryUses());
---  for (unsigned I = 1; I != 39; ++I)
---    EXPECT_EQ(N, Refs[I].get()) << I;
---}
---
-- TEST(NamedMDNodeTest, Search) {
--   LLVMContext Context;
--   ConstantAsMetadata *C =
++diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll b/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll
++--- a/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll
+++++ b/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll
++@@ -1,5 +1,7 @@
++ ; Checks if conversion is only inserted for the spilled register
++ ; instead of 2 conversions for the W register
+++
+++; REQUIRES: asserts
++ ; RUN: llc -O2 -march=hexagon -mcpu=hexagonv81 -enable-xqf-gen=true \
++ ; RUN: -hexagon-qfloat-mode=ieee -mattr=+hvxv81,+hvx-length128B \
++ ; RUN: -enable-postra-xqf-check -debug-only=handle-qfp 2>&1 < %s -o /dev/null | FileCheck %s
 diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
-index 9de7c05a..7b8ac881 100644
+index 7b8ac881..85cdf966 100644
 --- a/third_party/llvm/workspace.bzl
 +++ b/third_party/llvm/workspace.bzl
 @@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
  
  def repo(name):
      """Imports LLVM."""
--    LLVM_COMMIT = "62e30e3326d9b4e3948d242ccfee7aa4fc030ef3"
--    LLVM_SHA256 = "478ae34c463f28708582082098aaf0e79f431c2719f5d297bd546051dbeae0dd"
-+    LLVM_COMMIT = "c4d9267a31897dd76b2aab5196b75bffb49e6846"
-+    LLVM_SHA256 = "a96c8b33f0d8d3be008b8c9f74ef310254c4f7f7edf8c04c76f4d99ed9d34f1a"
+-    LLVM_COMMIT = "c4d9267a31897dd76b2aab5196b75bffb49e6846"
+-    LLVM_SHA256 = "a96c8b33f0d8d3be008b8c9f74ef310254c4f7f7edf8c04c76f4d99ed9d34f1a"
++    LLVM_COMMIT = "b08e005995ed29a2f89aa9a3f599a9d45db65523"
++    LLVM_SHA256 = "3f4f1765816e22a21ac41dbf71d91d519b0b371b05d7e33ef2a0a862c838420c"
  
      tf_http_archive(
          name = name,
-diff --git a/third_party/stablehlo/temporary.patch b/third_party/stablehlo/temporary.patch
-index c38c5147..c6cd6f2b 100644
---- a/third_party/stablehlo/temporary.patch
-+++ b/third_party/stablehlo/temporary.patch
-@@ -881,9 +881,9 @@ diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
- +Performs element-wise base-2 exponential operation on `operand` tensor and
- +produces a `result` tensor. Depending on the element type, does the following:
- +
--+*   For floats: `exp2` from IEEE-754.
--+*   For complex numbers: complex base-2 exponential.
--+*   For quantized types: `dequantize_op_quantize(exp2, operand, type(result))`.
-++* For floats: `exp2` from IEEE-754.
-++* For complex numbers: complex base-2 exponential.
-++* For quantized types: `dequantize_op_quantize(exp2, operand, type(result))`.
- +
- +#### Inputs
- +
-@@ -907,7 +907,7 @@ diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
- +
- +#### Constraints
- +
--+*   (C1) `baseline_type(operand) = baseline_type(result)`.
-++* (C1) `baseline_type(operand) = baseline_type(result)`.
- +
- +#### Examples
- +
-@@ -947,9 +947,9 @@ diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
- +Performs element-wise base-2 logarithm operation on `operand` tensor and
- +produces a `result` tensor. Depending on the element type, does the following:
- +
--+*   For floats: `log2` from IEEE-754.
--+*   For complex numbers: complex base-2 logarithm.
--+*   For quantized types: `dequantize_op_quantize(log2, operand, type(result))`.
-++* For floats: `log2` from IEEE-754.
-++* For complex numbers: complex base-2 logarithm.
-++* For quantized types: `dequantize_op_quantize(log2, operand, type(result))`.
- +
- +#### Inputs
- +
-@@ -973,7 +973,7 @@ diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
- +
- +#### Constraints
- +
--+*   (C1) `baseline_type(operand) = baseline_type(result)`.
-++* (C1) `baseline_type(operand) = baseline_type(result)`.
- +
- +#### Examples
- +
-@@ -1018,22 +1018,24 @@ diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
-  * (C1) `element_type(operand) = element_type(source)`.
-  * (C2) `shape(source) = num_windows` where:
-    * `padded_operand_shape = padding[:, 0] + shape(operand) + padding[:, 1]`.
--@@ -7501,6 +7597,7 @@
-+@@ -7501,7 +7597,7 @@
-  syntax ([ISO/IEC 14977:1996](https://www.iso.org/standard/26153.html),
-  [Wikipedia](https://en.wikipedia.org/wiki/Extended_Backus%E2%80%93Naur_form)),
-  with two modifications: 1) rules are defined using `::=` rather than `=`,
--+
-- 2) concatenation is expressed using juxtaposition rather than `,`.
-+-2) concatenation is expressed using juxtaposition rather than `,`.
-++2\) concatenation is expressed using juxtaposition rather than `,`.
-  
-  For describing semantics (i.e. within "Types", "Constants" and "Ops" sections),
--@@ -7571,6 +7668,7 @@
-+ we are using formulas which are based on Python syntax extended with support
-+@@ -7571,7 +7667,7 @@
-  ### Names
-  
-  In formulas, lexical scope includes: 1) global functions, 2) member definitions,
--+
-- 3) local definitions. The list of global functions is provided below. The list
-+-3) local definitions. The list of global functions is provided below. The list
-++3\) local definitions. The list of global functions is provided below. The list
-  of element definitions depends on the program element that the notation is
-  applied to:
-+ 
- diff --ruN a/stablehlo/rfcs/20260821-low-precision-complex.md b/stablehlo/rfcs/20260821-low-precision-complex.md
- --- stablehlo/rfcs/20260821-low-precision-complex.md
- +++ stablehlo/rfcs/20260821-low-precision-complex.md
-@@ -1657,16 +1659,16 @@ diff --ruN a/stablehlo/rfcs/20260907-exp2-log2.md b/stablehlo/rfcs/20260907-exp2
- +
- +Currently, StableHLO lacks native `exp2` and `log2` operations:
- +
--+*   Frameworks expose `exp2` and `log2` directly: JAX provides
--+    [`jnp.exp2`][jax1] / `lax.exp2` and [`jnp.log2`][jax2] / `lax.log2`, and
--+    PyTorch provides [`torch.exp2`][torch1] and [`torch.log2`][torch2].
--+*   Without first-class ops in StableHLO, frameworks must decompose `exp2(x)`
--+    into `exponential(x * ln(2))` and `log2(x)` into `log(x) / ln(2)`.
--+*   Decomposing `exp2` and `log2` into natural `exponential` and `log` scaled by
--+    `ln(2)` degrades numerical accuracy (introducing rounding error on integer
--+    powers of two where `exp2(n)` and `log2(2^n)` should evaluate exactly) and
--+    prevents backends from emitting single-instruction base-2 hardware
--+    operations without fragile algebraic pattern-matching.
-++* Frameworks expose `exp2` and `log2` directly: JAX provides [`jnp.exp2`][jax1]
-++  / `lax.exp2` and [`jnp.log2`][jax2] / `lax.log2`, and PyTorch provides
-++  [`torch.exp2`][torch1] and [`torch.log2`][torch2].
-++* Without first-class ops in StableHLO, frameworks must decompose `exp2(x)` into
-++  `exponential(x * ln(2))` and `log2(x)` into `log(x) / ln(2)`.
-++* Decomposing `exp2` and `log2` into natural `exponential` and `log` scaled by
-++  `ln(2)` degrades numerical accuracy (introducing rounding error on integer
-++  powers of two where `exp2(n)` and `log2(2^n)` should evaluate exactly) and
-++  prevents backends from emitting single-instruction base-2 hardware operations
-++  without fragile algebraic pattern-matching.
- +
- +Following the precedent of `TanOp` ([20240312-tanop.md](20240312-tanop.md)) and
- +the transcendental accuracy specification
-@@ -1681,12 +1683,12 @@ diff --ruN a/stablehlo/rfcs/20260907-exp2-log2.md b/stablehlo/rfcs/20260907-exp2
- +
- +## Exp2Op, Log2Op, and Compatibility
- +
--+*   **VHLO Serialization**: StableHLO version `1.21.0` introduces
--+    `VHLO_Exp2OpV1` (`exp2_v1`) and `VHLO_Log2OpV1` (`log2_v1`).
--+*   **Forward Compatibility & Decomposition**: For runtimes targeting StableHLO
--+    versions prior to `1.21.0`, compatibility expansion lowers `exp2(x)` to
--+    `exponential(multiply(x, constant(ln(2))))` and `log2(x)` to `divide(log(x),
--+    constant(ln(2)))`.
-++* **VHLO Serialization**: StableHLO version `1.21.0` introduces `VHLO_Exp2OpV1`
-++  (`exp2_v1`) and `VHLO_Log2OpV1` (`log2_v1`).
-++* **Forward Compatibility & Decomposition**: For runtimes targeting StableHLO
-++  versions prior to `1.21.0`, compatibility expansion lowers `exp2(x)` to
-++  `exponential(multiply(x, constant(ln(2))))` and `log2(x)` to
-++  `divide(log(x), constant(ln(2)))`.
- +
- +[jax1]: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.exp2.html
- +[jax2]: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.log2.html
-@@ -3698,15 +3700,30 @@ diff --ruN a/stablehlo/stablehlo/reference/InterpreterOps.cpp b/stablehlo/stable
-  
-  #include "llvm/Support/Casting.h"
-  #include "llvm/Support/Errc.h"
--@@ -175,36 +176,39 @@
-+@@ -174,37 +175,52 @@
-+     ArrayRef<InterpreterValue> inputs, std::queue<StringAttr>& infeed,
-      SmallVector<SmallVector<StringAttr>> programs, SymbolTable& symbolTable,
-      InterpreterFallback* fallback) {
--   llvm::DefaultThreadPool threadPool;
-+-  llvm::DefaultThreadPool threadPool;
- -  SmallVector<std::shared_future<SmallVector<InterpreterValue>>> futures;
--+  llvm::ThreadPoolTaskGroup taskGroup(threadPool);
-- 
-+-
-    uint32_t numReplicas = programs.size();
-    uint32_t numPartitions = programs[0].size();
-++
-++  // StableHLO collectives are blocking barriers: a process inside rendezvous
-++  // waits for every other process in its group to arrive, so all processes
-++  // must run concurrently. Size the pool from the process count rather than
-++  // the default (hardware_concurrency(), i.e. the host thread count). With a
-++  // default-sized pool, a program with more processes than the host has
-++  // hardware threads leaves the surplus tasks queued and never started while
-++  // the started tasks block forever waiting for them - a deadlock that only
-++  // surfaces as the rendezvous timeout. hardware_concurrency(N) leaves
-++  // ThreadPoolStrategy::Limit false, so the pool gets N threads uncapped by the
-++  // host.
-++  llvm::DefaultThreadPool threadPool(
-++      llvm::hardware_concurrency(numReplicas * numPartitions));
-++  llvm::ThreadPoolTaskGroup taskGroup(threadPool);
-++
-    ProcessGrid processGrid(numReplicas, numPartitions, infeed);
-  
- +  SmallVector<SmallVector<InterpreterValue>> taskOutputs(numReplicas *
-@@ -4443,6 +4460,159 @@ diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compat
-  
-    // CHECK-LABEL: @all_reduce_sdy_mesh_max
-    func.func @all_reduce_sdy_mesh_max(%arg0: tensor<4xf32>) -> tensor<4xf32> {
-+diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir b/stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir
-+--- stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir
-++++ stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir
-+@@ -168,6 +168,18 @@
-+ 
-+ // -----
-+ 
-++// CHECK-LABEL: func @dynamic_broadcast_in_dim_discardable_attrs
-++func.func @dynamic_broadcast_in_dim_discardable_attrs(%arg0: tensor<4xf32>) -> tensor<3x4xf32> {
-++  //  CHECK-NOT: stablehlo.dynamic_broadcast_in_dim
-++  //      CHECK: stablehlo.broadcast_in_dim %arg0, dims = [1]
-++  // CHECK-SAME:   {mhlo.frontend_attributes = {MUST_FUSE = "0"}}
-++  %0 = stablehlo.constant dense<[3, 4]> : tensor<2xi64>
-++  %1 = stablehlo.dynamic_broadcast_in_dim %arg0, %0, dims = [1] {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<4xf32>, tensor<2xi64>) -> tensor<3x4xf32>
-++  return %1 : tensor<3x4xf32>
-++}
-++
-++// -----
-++
-+ // CHECK-LABEL: func @dynamic_broadcast_in_dim_inapplicable_dynamic_operand_type
-+ func.func @dynamic_broadcast_in_dim_inapplicable_dynamic_operand_type(%arg0: tensor<?xf32>) -> tensor<3x4xf32> {
-+   // CHECK: stablehlo.dynamic_broadcast_in_dim
-+@@ -225,6 +237,26 @@
-+ 
-+ // -----
-+ 
-++// CHECK-LABEL: func @dynamic_conv_discardable_attrs
-++func.func @dynamic_conv_discardable_attrs(%arg0: tensor<100x26x26x32xf32>, %arg1: tensor<3x3x1x32xf32>) -> tensor<100x28x28x1xf32> {
-++  //  CHECK-NOT: stablehlo.dynamic_conv
-++  //      CHECK: stablehlo.convolution(%arg0, %arg1)
-++  // CHECK-SAME:   mhlo.frontend_attributes = {MUST_FUSE = "0"}
-++  %0 = stablehlo.constant dense<2> : tensor<2x2xi32>
-++  %1 = "stablehlo.dynamic_conv"(%arg0, %arg1, %0) {
-++    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>,
-++    window_strides = array<i64: 1, 1>,
-++    lhs_dilation = array<i64: 1, 1>,
-++    rhs_dilation = array<i64: 1, 1>,
-++    feature_group_count = 1 : i64,
-++    batch_group_count = 1 : i64,
-++    mhlo.frontend_attributes = {MUST_FUSE = "0"}
-++  } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf32>, tensor<2x2xi32>) -> tensor<100x28x28x1xf32>
-++  return %1 : tensor<100x28x28x1xf32>
-++}
-++
-++// -----
-++
-+ // CHECK-LABEL: func @dynamic_conv_success_dynamic_result_type
-+ func.func @dynamic_conv_success_dynamic_result_type(%arg0: tensor<100x26x26x32xf32>, %arg1: tensor<3x3x1x32xf32>) -> tensor<?x28x28x1xf32> {
-+   //  CHECK-NOT: stablehlo.dynamic_conv
-+@@ -295,6 +327,25 @@
-+ 
-+ // -----
-+ 
-++// CHECK-LABEL: @dynamic_gather_discardable_attrs
-++func.func @dynamic_gather_discardable_attrs(%arg0 : tensor<2x4x9xi32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x8xi32> {
-++  //  CHECK-NOT: stablehlo.dynamic_gather
-++  //      CHECK: "stablehlo.gather"(%arg0, %arg1)
-++  // CHECK-SAME:   {mhlo.frontend_attributes = {MUST_FUSE = "0"}}
-++  %0 = stablehlo.constant dense<[1, 1, 8]> : tensor<3xi32>
-++  %1 = "stablehlo.dynamic_gather"(%arg0, %arg1, %0) <{
-++    dimension_numbers = #stablehlo.gather<
-++      collapsed_slice_dims = [0, 1],
-++      index_vector_dim = 2,
-++      offset_dims = [2],
-++      start_index_map = [0, 1]
-++    >
-++  }> {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<2x4x9xi32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<1x5x8xi32>
-++  return %1 : tensor<1x5x8xi32>
-++}
-++
-++// -----
-++
-+ // CHECK-LABEL: @dynamic_gather_success_dynamic_result_type
-+ func.func @dynamic_gather_success_dynamic_result_type(%arg0 : tensor<2x4x9xi32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x?xi32> {
-+   //  CHECK-NOT: stablehlo.dynamic_gather
-+@@ -348,6 +399,18 @@
-+ 
-+ // -----
-+ 
-++// CHECK-LABEL: func @dynamic_iota_discardable_attrs
-++func.func @dynamic_iota_discardable_attrs() -> tensor<4xf32> {
-++  //  CHECK-NOT: stablehlo.dynamic_iota
-++  //      CHECK: stablehlo.iota
-++  // CHECK-SAME:   {mhlo.frontend_attributes = {MUST_FUSE = "0"}}
-++  %0 = stablehlo.constant dense<4> : tensor<1xi64>
-++  %1 = stablehlo.dynamic_iota %0, dim = 0 {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<1xi64>) -> tensor<4xf32>
-++  return %1 : tensor<4xf32>
-++}
-++
-++// -----
-++
-+ // CHECK-LABEL: func @dynamic_iota_inapplicable_dynamic_output_shape
-+ func.func @dynamic_iota_inapplicable_dynamic_output_shape(%arg0: tensor<1xi64>) -> tensor<4xf32> {
-+   // CHECK: stablehlo.dynamic_iota
-+@@ -380,6 +443,20 @@
-+ 
-+ // -----
-+ 
-++// CHECK-LABEL: func @dynamic_pad_discardable_attrs
-++func.func @dynamic_pad_discardable_attrs(%arg0: tensor<4xf32>, %arg1: tensor<f32>) -> tensor<11xf32> {
-++  //  CHECK-NOT: stablehlo.dynamic_pad
-++  //      CHECK: stablehlo.pad
-++  // CHECK-SAME:   {mhlo.frontend_attributes = {MUST_FUSE = "0"}}
-++  %0 = stablehlo.constant dense<0> : tensor<1xi64>
-++  %1 = stablehlo.constant dense<1> : tensor<1xi64>
-++  %2 = stablehlo.constant dense<2> : tensor<1xi64>
-++  %3 = stablehlo.dynamic_pad %arg0, %arg1, %0, %1, %2 {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<4xf32>, tensor<f32>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> tensor<11xf32>
-++  return %3 : tensor<11xf32>
-++}
-++
-++// -----
-++
-+ // CHECK-LABEL: func @dynamic_pad_success_dynamic_result_type
-+ func.func @dynamic_pad_success_dynamic_result_type(%arg0: tensor<4xf32>, %arg1: tensor<f32>) -> tensor<?xf32> {
-+   // CHECK-NOT: stablehlo.dynamic_pad
-+@@ -437,6 +514,17 @@
-+ 
-+ // -----
-+ 
-++// CHECK-LABEL: func @dynamic_reshape_discardable_attrs
-++func.func @dynamic_reshape_discardable_attrs(%arg0: tensor<4xf32>) -> tensor<1x4xf32> {
-++  //  CHECK-NOT: stablehlo.dynamic_reshape
-++  //      CHECK: stablehlo.reshape %arg0 {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<4xf32>) -> tensor<1x4xf32>
-++  %0 = stablehlo.constant dense<[1, 4]> : tensor<2xi64>
-++  %1 = stablehlo.dynamic_reshape %arg0, %0 {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<4xf32>, tensor<2xi64>) -> tensor<1x4xf32>
-++  return %1 : tensor<1x4xf32>
-++}
-++
-++// -----
-++
-+ // CHECK-LABEL: func @dynamic_reshape_inapplicable_dynamic_output_shape
-+ func.func @dynamic_reshape_inapplicable_dynamic_output_shape(%arg0: tensor<4xf32>, %arg1: tensor<2xi64>) -> tensor<1x4xf32> {
-+   // CHECK: stablehlo.dynamic_reshape
-+diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_legalize_deprecated_ops.mlir b/stablehlo/stablehlo/tests/transforms/stablehlo_legalize_deprecated_ops.mlir
-+--- stablehlo/stablehlo/tests/transforms/stablehlo_legalize_deprecated_ops.mlir
-++++ stablehlo/stablehlo/tests/transforms/stablehlo_legalize_deprecated_ops.mlir
-+@@ -10,6 +10,16 @@
-+ 
-+ // -----
-+ 
-++// CHECK-LABEL: op_dot_discardable_attrs
-++func.func @op_dot_discardable_attrs(%arg0: tensor<2x3xf32>,
-++                 %arg1: tensor<3x?xf32>) -> tensor<2x?xf32> {
-++  // CHECK: stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<2x3xf32>, tensor<3x?xf32>) -> tensor<2x?xf32>
-++  %0 = "stablehlo.dot"(%arg0, %arg1) {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<2x3xf32>, tensor<3x?xf32>) -> tensor<2x?xf32>
-++  func.return %0 : tensor<2x?xf32>
-++}
-++
-++// -----
-++
-+ // CHECK-LABEL: op_unary_einsum
-+ func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
-+   // CHECK:      %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
- diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir b/stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir
- --- stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir
- +++ stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir
-@@ -8393,6 +8563,99 @@ diff --ruN a/stablehlo/stablehlo/transforms/StablehloCanonicalizeDynamism.cpp b/
-      return success();
-    }
-  };
-+@@ -119,8 +120,11 @@
-+                                          "expected static output_dimensions");
-+     if (!op.getType().hasStaticShape())
-+       return rewriter.notifyMatchFailure(op, "expected static result type");
-+-    rewriter.replaceOpWithNewOp<BroadcastInDimOp>(
-+-        op, op.getType(), op.getOperand(), op.getBroadcastDimensionsAttr());
-++    auto broadcastInDimOp = BroadcastInDimOp::create(
-++        rewriter, op.getLoc(), op.getType(), op.getOperand(),
-++        op.getBroadcastDimensionsAttr());
-++    broadcastInDimOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
-++    rewriter.replaceOp(op, broadcastInDimOp);
-+     return success();
-+   }
-+ };
-+@@ -139,12 +143,14 @@
-+         RankedTensorType::get({static_cast<int64_t>(padding.size()) / 2, 2},
-+                               rewriter.getI64Type()),
-+         padding);
-+-    rewriter.replaceOpWithNewOp<ConvolutionOp>(
-+-        op, op.getType(), op.getLhs(), op.getRhs(), op.getWindowStridesAttr(),
-+-        paddingAttr, op.getLhsDilationAttr(), op.getRhsDilationAttr(),
-+-        op.getWindowReversalAttr(), op.getDimensionNumbers(),
-+-        op.getFeatureGroupCount(), op.getBatchGroupCount(),
-+-        op.getPrecisionConfigAttr());
-++    auto convolutionOp = ConvolutionOp::create(
-++        rewriter, op.getLoc(), op.getType(), op.getLhs(), op.getRhs(),
-++        op.getWindowStridesAttr(), paddingAttr, op.getLhsDilationAttr(),
-++        op.getRhsDilationAttr(), op.getWindowReversalAttr(),
-++        op.getDimensionNumbers(), op.getFeatureGroupCount(),
-++        op.getBatchGroupCount(), op.getPrecisionConfigAttr());
-++    convolutionOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
-++    rewriter.replaceOp(op, convolutionOp);
-+     return success();
-+   }
-+ };
-+@@ -159,10 +165,13 @@
-+     SmallVector<int64_t> sliceSizes;
-+     if (!succeeded(hlo::matchInts(op.getSliceSizes(), sliceSizes)))
-+       return rewriter.notifyMatchFailure(op, "expected static slice_sizes");
-+-    rewriter.replaceOpWithNewOp<GatherOp>(
-+-        op, op.getType(), op.getOperand(), op.getStartIndices(),
-+-        op.getDimensionNumbersAttr(), rewriter.getDenseI64ArrayAttr(sliceSizes),
-+-        op.getIndicesAreSortedAttr());
-++    auto gatherOp =
-++        GatherOp::create(rewriter, op.getLoc(), op.getType(), op.getOperand(),
-++                         op.getStartIndices(), op.getDimensionNumbersAttr(),
-++                         rewriter.getDenseI64ArrayAttr(sliceSizes),
-++                         op.getIndicesAreSortedAttr());
-++    gatherOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
-++    rewriter.replaceOp(op, gatherOp);
-+     return success();
-+   }
-+ };
-+@@ -179,8 +188,10 @@
-+       return rewriter.notifyMatchFailure(op, "expected static output_shape");
-+     if (!op.getType().hasStaticShape())
-+       return rewriter.notifyMatchFailure(op, "expected static result type");
-+-    rewriter.replaceOpWithNewOp<IotaOp>(op, op.getType(),
-+-                                        op.getIotaDimension());
-++    auto iotaOp = IotaOp::create(rewriter, op.getLoc(), op.getType(),
-++                                 op.getIotaDimension());
-++    iotaOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
-++    rewriter.replaceOp(op, iotaOp);
-+     return success();
-+   }
-+ };
-+@@ -198,9 +209,11 @@
-+       return rewriter.notifyMatchFailure(op, "expected static high");
-+     if (!succeeded(hlo::matchInts(op.getInteriorPadding(), interiorPadding)))
-+       return rewriter.notifyMatchFailure(op, "expected static interior");
-+-    rewriter.replaceOpWithNewOp<PadOp>(op, op.getType(), op.getOperand(),
-+-                                       op.getPaddingValue(), edgePaddingLow,
-+-                                       edgePaddingHigh, interiorPadding);
-++    auto padOp = PadOp::create(
-++        rewriter, op.getLoc(), op.getType(), op.getOperand(),
-++        op.getPaddingValue(), edgePaddingLow, edgePaddingHigh, interiorPadding);
-++    padOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
-++    rewriter.replaceOp(op, padOp);
-+     return success();
-+   }
-+ };
-+@@ -216,7 +229,10 @@
-+       return rewriter.notifyMatchFailure(op, "expected static output_shape");
-+     if (!op.getType().hasStaticShape())
-+       return rewriter.notifyMatchFailure(op, "expected static result type");
-+-    rewriter.replaceOpWithNewOp<ReshapeOp>(op, op.getType(), op.getOperand());
-++    auto reshapeOp =
-++        ReshapeOp::create(rewriter, op.getLoc(), op.getType(), op.getOperand());
-++    reshapeOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
-++    rewriter.replaceOp(op, reshapeOp);
-+     return success();
-+   }
-+ };
- diff --ruN a/stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp b/stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp
- --- stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp
- +++ stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp
-@@ -8433,6 +8696,42 @@ diff --ruN a/stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp b/s
-    populateWithGenerated(*patterns);
-  }
-  
-+diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOps.cpp b/stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOps.cpp
-+--- stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOps.cpp
-++++ stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOps.cpp
-+@@ -132,6 +132,14 @@
-+   llvm::report_fatal_error("unsupported type");
-+ }
-+ 
-++void copyDiscardableAttrs(Operation* src, Operation* dst) {
-++  dst->setDiscardableAttrs(src->getDiscardableAttrDictionary());
-++}
-++
-++void copyDiscardableAttrs(Value src, Value dst) {
-++  copyDiscardableAttrs(src.getDefiningOp(), dst.getDefiningOp());
-++}
-++
-+ #include "stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.h.inc"
-+ }  // namespace
-+ 
-+diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.td b/stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.td
-+--- stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.td
-++++ stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.td
-+@@ -45,9 +45,12 @@
-+ 
-+ // def CrossReplicaSumToAllReduce : <Declared in cpp>
-+ 
-++def CopyDiscardableAttrs : NativeCodeCallVoid<"copyDiscardableAttrs($0, $1)">;
-++
-+ def DotToDotGeneral : Pat<
-+-  (StableHLO_DotOp $lhs, $rhs, $precision_config),
-+-  (StableHLO_DotGeneralOp $lhs, $rhs, (GetDefaultDotDimensionNumbers $lhs), $precision_config, (NativeCodeCall<"Attribute{}">))>;
-++  (StableHLO_DotOp:$src $lhs, $rhs, $precision_config),
-++  (StableHLO_DotGeneralOp:$dst $lhs, $rhs, (GetDefaultDotDimensionNumbers $lhs), $precision_config, (NativeCodeCall<"Attribute{}">)),
-++  [], [(CopyDiscardableAttrs $src, $dst)]>;
-+ 
-+ def UnaryEinsumToEinsum : Pat<
-+   (StableHLO_UnaryEinsumOp $operand, $equation),
- diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp b/stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp
- --- stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp
- +++ stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp
diff --git a/third_party/xla/third_party/shardy/workspace.bzl b/third_party/xla/third_party/shardy/workspace.bzl
index 7a739adf488..83345d191e0 100644
--- a/third_party/xla/third_party/shardy/workspace.bzl
+++ b/third_party/xla/third_party/shardy/workspace.bzl
@@ -18,8 +18,8 @@
 load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 
 def repo():
-    SHARDY_COMMIT = "1ecfcb6250edeb638fce0e513073b9f52047ad83"
-    SHARDY_SHA256 = "726b288c47393ea0567b4655aa6516c523eefeb7db724bfc34d9cd94fff15b02"
+    SHARDY_COMMIT = "ba51e5d4ab46ef7b2f374d47f64c3317b9134441"
+    SHARDY_SHA256 = "d3c79a9a3c219cd0e68f696af200255eee185a5b128e274443d4ecf92e1e5fb5"
 
     tf_http_archive(
         name = "shardy",

From c0b92cdc0676251963e18004056251eff87c62e7 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 9 Oct 2026 09:59:24 -0700
Subject: [PATCH 0376/1000] [XLA] Add PjRtRelocatable interfaces

This change introduces the concept of a relocatable to pjrt's
understanding of compilers and toolchains: the rough equivalent of a .o
file for XLA.

It also addes the compiler interfaces for producing and consuming them.

PiperOrigin-RevId: 996596836
---
 third_party/xla/xla/pjrt/BUILD                | 23 +++++
 third_party/xla/xla/pjrt/pjrt_compiler.h      | 41 ++++++++
 third_party/xla/xla/pjrt/pjrt_relocatable.h   | 96 +++++++++++++++++++
 .../xla/xla/pjrt/pjrt_relocatable_test.cc     | 26 +++++
 4 files changed, 186 insertions(+)
 create mode 100644 third_party/xla/xla/pjrt/pjrt_relocatable.h
 create mode 100644 third_party/xla/xla/pjrt/pjrt_relocatable_test.cc

diff --git a/third_party/xla/xla/pjrt/BUILD b/third_party/xla/xla/pjrt/BUILD
index 7c6763b061d..1cfa8417c3f 100644
--- a/third_party/xla/xla/pjrt/BUILD
+++ b/third_party/xla/xla/pjrt/BUILD
@@ -612,6 +612,7 @@ cc_library(
         ":pjrt_device_description",
         ":pjrt_device_dimensions",
         ":pjrt_executable",
+        ":pjrt_relocatable",
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",
         "//xla/hlo/builder:xla_computation",
@@ -1545,3 +1546,25 @@ cc_library(
         "@tsl//tsl/platform:fingerprint",
     ],
 )
+
+cc_library(
+    name = "pjrt_relocatable",
+    hdrs = ["pjrt_relocatable.h"],
+    visibility = internal_visibility([":friends"]),
+    deps = [
+        ":pjrt_layout",
+        "//xla:util",
+        "//xla:xla_data_proto_cc",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+)
+
+xla_cc_test(
+    name = "pjrt_relocatable_test",
+    srcs = ["pjrt_relocatable_test.cc"],
+    deps = [
+        ":pjrt_relocatable",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
diff --git a/third_party/xla/xla/pjrt/pjrt_compiler.h b/third_party/xla/xla/pjrt/pjrt_compiler.h
index 336bf8b1c79..d87254e633f 100644
--- a/third_party/xla/xla/pjrt/pjrt_compiler.h
+++ b/third_party/xla/xla/pjrt/pjrt_compiler.h
@@ -44,6 +44,7 @@ limitations under the License.
 #include "xla/pjrt/pjrt_device_description.h"
 #include "xla/pjrt/pjrt_device_dimensions.h"
 #include "xla/pjrt/pjrt_executable.h"
+#include "xla/pjrt/pjrt_relocatable.h"
 #include "xla/pjrt/proto/pjrt_partial_program.pb.h"
 #include "xla/pjrt/proto/topology_description.pb.h"
 #include "xla/runtime/chip_id.h"
@@ -690,6 +691,46 @@ class PjRtPhaseCompiler : public PjRtCompiler {
         "DeserializePjRtTopologyDescription is not implemented.");
   }
 
+  // Compiles a StableHLO module to a relocatable.
+  //
+  // See LinkRelocatables for how to use is_entrypoint.
+  //
+  // Implementations of `PjRtPhaseCompiler` are not required to implement
+  // this method, but is grouped here because a phased compiler is a
+  // prerequisite for precompiling modules at all.
+  virtual absl::StatusOr<std::unique_ptr<PjRtRelocatable>> CompileToRelocatable(
+      CompileOptions options, MaybeOwningMlirModule module,
+      const PjRtTopologyDescription& topology, absl::string_view name,
+      bool is_entrypoint, PjRtClient* client) {
+    return absl::UnimplementedError("CompileToRelocatable is not implemented.");
+  }
+
+  // Deserializes a relocatable into a PjRtRelocatable.
+  //
+  // Implementations of `PjRtPhaseCompiler` are not required to implement
+  // this method, but is grouped here because a phased compiler is a
+  // prerequisite for precompiling modules at all.
+  virtual absl::StatusOr<std::unique_ptr<PjRtRelocatable>>
+  DeserializeRelocatable(absl::string_view serialized) const {
+    return absl::UnimplementedError(
+        "DeserializeRelocatable is not implemented.");
+  }
+
+  // Links the provided relocatables into an executable.
+  //
+  // The entrypoint relocatable must have been compiled with is_entrypoint=true,
+  // while the dependencies must have been compiled with is_entrypoint=false.
+  //
+  // Implementations of `PjRtPhaseCompiler` are not required to implement
+  // this method, but is grouped here because a phased compiler is a
+  // prerequisite for having a link phase at all.
+  virtual absl::StatusOr<std::unique_ptr<PjRtExecutable>> LinkRelocatables(
+      CompileOptions options, const PjRtTopologyDescription& topology,
+      std::shared_ptr<const PjRtRelocatable> entrypoint,
+      std::vector<std::shared_ptr<const PjRtRelocatable>> dependencies) {
+    return absl::UnimplementedError("LinkRelocatables is not implemented.");
+  }
+
   PjRtPhaseCompiler* AsPhaseCompiler() override { return this; }
   const PjRtPhaseCompiler* AsPhaseCompiler() const override { return this; }
 
diff --git a/third_party/xla/xla/pjrt/pjrt_relocatable.h b/third_party/xla/xla/pjrt/pjrt_relocatable.h
new file mode 100644
index 00000000000..2fc7840a7b6
--- /dev/null
+++ b/third_party/xla/xla/pjrt/pjrt_relocatable.h
@@ -0,0 +1,96 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_PJRT_PJRT_RELOCATABLE_H_
+#define XLA_PJRT_PJRT_RELOCATABLE_H_
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "xla/pjrt/pjrt_layout.h"
+#include "xla/util.h"
+#include "xla/xla_data.pb.h"
+
+namespace xla {
+
+// Represents a relocatable intermediate compilation artifact that can be
+// called by other programs. Relocatables can be linked together into an
+// executable. Relocatables are analogous to a object files and serve the
+// same purpose of the allowing one to compile a subset of a program
+// separately.
+//
+// Relocatables do not currently support ABI versioning, so users of
+// relocatables must ensure that the same compiler is producing and
+// consuming them.
+class PjRtRelocatable {
+ public:
+  virtual ~PjRtRelocatable() = default;
+
+  // Returns the name of this relocatable.
+  virtual absl::string_view name() const = 0;
+
+  // Returns the build ID of this relocatable, derived from the inputs that
+  // produced it.
+  virtual absl::string_view build_id() const = 0;
+
+  // Returns the backend config to apply to any call sites for this
+  // relocatable. The configuration is opaque and to be interpreted
+  // by a given backend.
+  virtual absl::string_view call_backend_config() const = 0;
+
+  // Returns whether the relocatable has side effects.
+  virtual bool has_side_effects() const = 0;
+
+  // Serializes this executable into a string. The compatibility of the
+  // serialized relocatable is implementation-specific.
+  virtual absl::StatusOr<std::string> Serialize() const = 0;
+
+  // Returns a list of element types for each parameter.
+  virtual absl::StatusOr<std::vector<PrimitiveType>> GetParameterElementTypes()
+      const = 0;
+
+  // Returns a list of element types for each output.
+  virtual absl::StatusOr<std::vector<PrimitiveType>> GetOutputElementTypes()
+      const = 0;
+
+  // Returns a list of dimensions for each parameter.
+  virtual absl::StatusOr<std::vector<DimensionVector>> GetParameterDimensions()
+      const = 0;
+
+  // Returns a list of dimensions for each output.
+  virtual absl::StatusOr<std::vector<DimensionVector>> GetOutputDimensions()
+      const = 0;
+
+  // Returns the layout of each input parameter.
+  virtual absl::StatusOr<std::vector<std::shared_ptr<const PjRtLayout>>>
+  GetParameterLayouts() const = 0;
+
+  // Returns the layout of each output.
+  virtual absl::StatusOr<std::vector<std::shared_ptr<const PjRtLayout>>>
+  GetOutputLayouts() const = 0;
+
+  // Returns which pairs of parameters and outputs may alias.
+  virtual absl::StatusOr<std::vector<std::pair<int64_t, int64_t>>>
+  GetParameterOutputAliases() const = 0;
+};
+
+}  // namespace xla
+
+#endif  // XLA_PJRT_PJRT_RELOCATABLE_H_
diff --git a/third_party/xla/xla/pjrt/pjrt_relocatable_test.cc b/third_party/xla/xla/pjrt/pjrt_relocatable_test.cc
new file mode 100644
index 00000000000..cf9cd3554cf
--- /dev/null
+++ b/third_party/xla/xla/pjrt/pjrt_relocatable_test.cc
@@ -0,0 +1,26 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+#include "xla/pjrt/pjrt_relocatable.h"  // IWYU pragma: keep
+
+#include <gmock/gmock.h>  // IWYU pragma: keep
+#include <gtest/gtest.h>  // IWYU pragma: keep
+
+namespace xla {
+namespace {
+
+TEST(PjRtRelocatableTest, Empty) {}
+
+}  // namespace
+}  // namespace xla

From e22b28c91e52bce5d0fdab54f586845e97490e68 Mon Sep 17 00:00:00 2001
From: Alex Pivovarov <upwind@google.com>
Date: Fri, 9 Oct 2026 10:44:29 -0700
Subject: [PATCH 0377/1000] [XLA] Store TopK comparator order in backend_config
 in TopkRewriter.

#### Why:
Why this change is needed:
1. `TopkSpecializer` rewrites `"TopK"` custom calls into GPU-specific `"__gpu$TopK"` custom calls, which drop `called_computations` (`to_apply`) and only preserve `backend_config`. Storing `order` (`"TOTAL"` or `"PARTIAL"`) in `backend_config` (alongside `is_stable`) is required so `ThunkEmitter` can select between the `TOTAL`-order and `PARTIAL`-order GPU TopK kernels.
2. `TopkRewriter` (`MatchNanSafeGt`) already pattern-matches the `sort` comparator computation. Recording the resulting `ComparisonOrder` directly in `backend_config` avoids re-parsing the comparator computation in downstream passes (`TopkSpecializer`, `ThunkEmitter`).

#### Changes:
- Update `MatchNanSafeGt` in `TopkRewriter` to extract and return the matched comparator's `ComparisonOrder` (`TOTAL` or `PARTIAL`).
- Include `order = "TOTAL"` or `order = "PARTIAL"` alongside `is_stable` in `raw_backend_config_string` when creating a `"TopK"` `CustomCall` in `CreateTopKCustomCall`.
- Update `hlo_instruction_utils::GetTopKOrder` to read `order` directly from `raw_backend_config_string` (defaulting to `ComparisonOrder::kTotal`), avoiding re-parsing the comparator computation in downstream passes (`TopkSpecializer` and `ThunkEmitter`) and preserving the order on `"__gpu$TopK"` custom calls that do not attach `to_apply`.
- Update `hlo_instruction_utils_test` and `topk_rewriter_test` (including a new `TopKCustomCallPartialOrderConfig` test).

Example
```
# Before Optimization
ENTRY %TopK (x: f32[8,69632]) -> (f32[8,4], s32[8,4]) {
  %x = f32[8,69632]{1,0} parameter(0)
  ROOT %topk = (f32[8,4]{1,0}, s32[8,4]{1,0}) topk(%x), k=4, largest=true, is_stable=true
}

# after_topk-rewriter
ENTRY %TopK (x: f32[8,69632]) -> (f32[8,4], s32[8,4]) {
  %x = f32[8,69632]{1,0} parameter(0)
  %custom-call = (f32[8,4]{1,0}, s32[8,4]{1,0}) custom-call(%x), custom_call_target="TopK", called_computations={%compare-greater-than.1}, backend_config={is_stable = true, order = "TOTAL"}
  ROOT %tuple = (f32[8,4]{1,0}, s32[8,4]{1,0}) tuple(%custom-call#0, %custom-call#1)
}

# after_topk-specializer
ENTRY %TopK (x: f32[8,69632]) -> (f32[8,4], s32[8,4]) {
  %x = f32[8,69632]{1,0} parameter(0)
  %custom-call.1 = (f32[8,4]{1,0}, s32[8,4]{1,0}) custom-call(%x), custom_call_target="__gpu$TopK", api_version=API_VERSION_TYPED_FFI, backend_config={is_stable = true, order = "TOTAL"}
  %tuple.1 = (f32[8,4]{1,0}, s32[8,4]{1,0}) tuple(%custom-call.1#0, %custom-call.1#1)
  ROOT %tuple = (f32[8,4]{1,0}, s32[8,4]{1,0}) tuple(%tuple.1#0, %tuple.1#1)
}
```

PiperOrigin-RevId: 996624017
---
 .../gpu/transforms/topk_specializer.cc        | 13 ++-
 .../gpu/transforms/topk_specializer_test.cc   | 48 +++++++--
 .../backends/gpu/transforms/topk_splitter.cc  |  2 +-
 .../gpu/transforms/topk_splitter_test.cc      | 11 ++-
 .../xla/xla/hlo/ir/hlo_instruction_utils.cc   | 35 ++++---
 .../xla/xla/hlo/ir/hlo_instruction_utils.h    | 13 +--
 .../xla/hlo/ir/hlo_instruction_utils_test.cc  | 98 ++++++++++---------
 third_party/xla/xla/service/topk_rewriter.cc  | 88 ++++++++++-------
 third_party/xla/xla/service/topk_rewriter.h   |  4 +-
 .../xla/xla/service/topk_rewriter_test.cc     | 43 +++++++-
 10 files changed, 233 insertions(+), 122 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/transforms/topk_specializer.cc b/third_party/xla/xla/backends/gpu/transforms/topk_specializer.cc
index 50ed5228aa6..377bc32f768 100644
--- a/third_party/xla/xla/backends/gpu/transforms/topk_specializer.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/topk_specializer.cc
@@ -179,8 +179,7 @@ HloInstruction* BuildUnpackU64ToF32(HloInstruction* u64_values,
 // Checks if we can safely route stable TopK to RAFT using the Uint64 adapter.
 bool ShouldRewriteStableTopKToUint64(HloCustomCallInstruction* topk) {
   // The Uint64 bit-packing adapter implements TOTAL ordering.
-  if (hlo_instruction_utils::GetTopKComparatorOrder(topk) !=
-      ComparisonOrder::kTotal) {
+  if (hlo_instruction_utils::GetTopKOrder(topk) != ComparisonOrder::kTotal) {
     return false;
   }
   if (!hlo_instruction_utils::IsTopKStable(topk)) {
@@ -259,16 +258,17 @@ absl::StatusOr<HloInstruction*> RewriteStableTopKToUint64(
   Shape new_cc_shape =
       ShapeUtil::MakeTupleShape({k_shape, idx_shape, scratch_shape});
 
-  HloInstruction* new_topk =
+  HloCustomCallInstruction* new_topk = Cast<HloCustomCallInstruction>(
       comp->AddInstruction(HloInstruction::CreateCustomCall(
           new_cc_shape, {packed_u64}, "__gpu$TopK", "",
-          CustomCallApiVersion::API_VERSION_TYPED_FFI));
+          CustomCallApiVersion::API_VERSION_TYPED_FFI)));
+  new_topk->set_raw_backend_config_string(topk->raw_backend_config_string());
 
   // The packed U64 keys guarantee uniqueness, making ties impossible.
   // Therefore, the inner TopK operation no longer requires stability to
   // produce a stable overall result. We clear the is_stable flag so the
   // backend can freely route this to the fast unstable topk kernel (RAFT lib).
-  new_topk->set_raw_backend_config_string("{is_stable = false}");
+  hlo_instruction_utils::SetTopKStability(new_topk, false);
 
   // 3. Unpack values and retain indices
   HloInstruction* u64_vals = comp->AddInstruction(
@@ -298,8 +298,7 @@ absl::StatusOr<HloInstruction*> SmallBufferOptimization(
   // TODO(b/473829358): Enable PARTIAL order in SmallBufferOptimization once
   // TopKPartialOrderKernel is rolled out to the runtime and wired up in
   // ThunkEmitter.
-  if (hlo_instruction_utils::GetTopKComparatorOrder(topk) !=
-      ComparisonOrder::kTotal) {
+  if (hlo_instruction_utils::GetTopKOrder(topk) != ComparisonOrder::kTotal) {
     return InvalidArgument(
         "Unsupported comparator order: only TOTAL order is currently "
         "supported.");
diff --git a/third_party/xla/xla/backends/gpu/transforms/topk_specializer_test.cc b/third_party/xla/xla/backends/gpu/transforms/topk_specializer_test.cc
index 177fcc8de4f..17e5bd5f00c 100644
--- a/third_party/xla/xla/backends/gpu/transforms/topk_specializer_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/topk_specializer_test.cc
@@ -137,7 +137,7 @@ TEST_F(TopkTest, PreservesBackendConfig) {
 
     ENTRY top_k {
       arg = f32[8,1024] parameter(0)
-      ROOT result = (f32[8,64], s32[8,64]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = false}
+      ROOT result = (f32[8,64], s32[8,64]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = false, order = "TOTAL"}
     }
   )";
 
@@ -161,7 +161,8 @@ TEST_F(TopkTest, PreservesBackendConfig) {
   const HloInstruction* root = module->entry_computation()->root_instruction();
   const HloInstruction* custom_call = root->operand(0)->operand(0);
   EXPECT_EQ(custom_call->custom_call_target(), "__gpu$TopK");
-  EXPECT_EQ(custom_call->raw_backend_config_string(), "{is_stable = false}");
+  EXPECT_EQ(custom_call->raw_backend_config_string(),
+            "{is_stable = false, order = \"TOTAL\"}");
 }
 
 TEST_F(TopkTest, RewriteStableTopKF32ToUint64) {
@@ -178,7 +179,7 @@ TEST_F(TopkTest, RewriteStableTopKF32ToUint64) {
 
     ENTRY top_k {
       arg = f32[8,1024] parameter(0)
-      ROOT result = (f32[8,32], s32[8,32]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = true}
+      ROOT result = (f32[8,32], s32[8,32]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = true, order = "TOTAL"}
     }
   )";
 
@@ -220,7 +221,7 @@ TEST_F(TopkTest, RewriteStableTopKF32ToUint64) {
 // CHECK: %[[PACKED:[^ ]+]] = u64[8,1024]{{.*}} or(%[[SHIFT_LEFT]], {{.*}})
 
 // 4. CustomCall (__gpu$TopK)
-// CHECK: %[[CUSTOM_CALL:[^ ]+]] = (u64[8,32]{{.*}}, s32[8,32]{{.*}}, u8[33554432]{{.*}}) custom-call(%[[PACKED]]), custom_call_target="__gpu$TopK", api_version=API_VERSION_TYPED_FFI, backend_config={is_stable = false}
+// CHECK: %[[CUSTOM_CALL:[^ ]+]] = (u64[8,32]{{.*}}, s32[8,32]{{.*}}, u8[33554432]{{.*}}) custom-call(%[[PACKED]]), custom_call_target="__gpu$TopK", api_version=API_VERSION_TYPED_FFI, backend_config={is_stable = false, order = "TOTAL"}
 
 // 5. Unpack U64 -> U32
 // CHECK: %[[SRL:[^ ]+]] = u64[8,32]{{.*}} shift-right-logical(%[[CUSTOM_CALL]]#0, {{.*}})
@@ -256,7 +257,7 @@ TEST_F(TopkTest, RewriteStableTopKBF16ToUint64) {
 
     ENTRY top_k {
       arg = bf16[8,65540] parameter(0)
-      ROOT result = (bf16[8,32], s32[8,32]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = true}
+      ROOT result = (bf16[8,32], s32[8,32]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = true, order = "TOTAL"}
     }
   )";
 
@@ -299,7 +300,7 @@ TEST_F(TopkTest, RewriteStableTopKBF16ToUint64) {
 // CHECK: %[[PACKED:[^ ]+]] = u64[8,65540]{{.*}} or(%[[SHIFT_LEFT]], {{.*}})
 
 // 4. CustomCall (__gpu$TopK)
-// CHECK: %[[CUSTOM_CALL:[^ ]+]] = (u64[8,32]{{.*}}, s32[8,32]{{.*}}, u8[33554432]{{.*}}) custom-call(%[[PACKED]]), custom_call_target="__gpu$TopK", api_version=API_VERSION_TYPED_FFI, backend_config={is_stable = false}
+// CHECK: %[[CUSTOM_CALL:[^ ]+]] = (u64[8,32]{{.*}}, s32[8,32]{{.*}}, u8[33554432]{{.*}}) custom-call(%[[PACKED]]), custom_call_target="__gpu$TopK", api_version=API_VERSION_TYPED_FFI, backend_config={is_stable = false, order = "TOTAL"}
 
 // 5. Unpack U64 -> U32
 // CHECK: %[[SRL:[^ ]+]] = u64[8,32]{{.*}} shift-right-logical(%[[CUSTOM_CALL]]#0, {{.*}})
@@ -336,7 +337,7 @@ TEST_F(TopkTest, RewriteStableTopKDisabledByDefault) {
 
     ENTRY top_k {
       arg = f32[8,1024] parameter(0)
-      ROOT result = (f32[8,32], s32[8,32]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = true}
+      ROOT result = (f32[8,32], s32[8,32]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = true, order = "TOTAL"}
     }
   )";
 
@@ -354,5 +355,38 @@ TEST_F(TopkTest, RewriteStableTopKDisabledByDefault) {
   EXPECT_FALSE(changed);
 }
 
+TEST_F(TopkTest, SkipsPartialOrderTopK) {
+  const char* hlo = R"(
+    HloModule m
+
+    %compare-gt.1 {
+      p.1.lhs = s32[] parameter(2)
+      p.1.rhs = s32[] parameter(3)
+      p.0.lhs = f32[] parameter(0)
+      p.0.rhs = f32[] parameter(1)
+      ROOT compare = pred[] compare(p.0.lhs, p.0.rhs), direction=GT
+    }
+
+    ENTRY top_k {
+      arg = f32[8,1024] parameter(0)
+      small_k = (f32[8,8], s32[8,8]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = true, order = "PARTIAL"}
+      raft_k = (f32[8,32], s32[8,32]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = true, order = "PARTIAL"}
+      ROOT result = tuple(small_k, raft_k)
+    }
+  )";
+
+  HloModuleConfig config = GetModuleConfigForTest();
+  config.mutable_debug_options()
+      .set_xla_gpu_experimental_enable_raft_for_stable_topk(true);
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(hlo, config));
+
+  ASSERT_OK_AND_ASSIGN(
+      bool changed,
+      TopkSpecializer(device_description().gpu_compute_capability())
+          .Run(module.get()));
+  EXPECT_FALSE(changed);
+}
+
 }  // namespace
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/transforms/topk_splitter.cc b/third_party/xla/xla/backends/gpu/transforms/topk_splitter.cc
index 860ee3d65c6..b7628cd5da7 100644
--- a/third_party/xla/xla/backends/gpu/transforms/topk_splitter.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/topk_splitter.cc
@@ -94,7 +94,7 @@ class TopkSplitterVisitor : public DfsHloRewriteVisitor {
     HloInstruction* batch_topk =
         comp->AddInstruction(HloInstruction::CreateCustomCall(
             batch_topk_shape, {reshaped}, topk->to_apply(), "TopK",
-            /*opaque=*/""));
+            topk->raw_backend_config_string()));
     // Fix indices, adding j*split_N to the j-th batch of indices.
     ABSL_ASSIGN_OR_RETURN(HloInstruction * indices,
                      MakeGetTupleElementHlo(batch_topk, 1));
diff --git a/third_party/xla/xla/backends/gpu/transforms/topk_splitter_test.cc b/third_party/xla/xla/backends/gpu/transforms/topk_splitter_test.cc
index 7474184d612..dd14e128e87 100644
--- a/third_party/xla/xla/backends/gpu/transforms/topk_splitter_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/topk_splitter_test.cc
@@ -30,6 +30,7 @@ limitations under the License.
 #include "absl/strings/substitute.h"
 #include "xla/backends/gpu/tests/hlo_pjrt_gpu_test_base.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"
@@ -65,13 +66,14 @@ HloModule module
 $0
 ENTRY cluster {
   %arg.1 = f32[1,1073741824] parameter(0)
-  ROOT %cc.2 = (f32[1,5], s32[1,5]) custom-call(%arg.1), custom_call_target= "TopK", to_apply=%compare
+  ROOT %cc.2 = (f32[1,5], s32[1,5]) custom-call(%arg.1), custom_call_target="TopK", to_apply=%compare, backend_config="{is_stable = false, order = \"PARTIAL\"}"
 })",
                                                   kComparator);
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
   EXPECT_THAT(RunHloPass(TopKSplitter(), module.get()),
               absl_testing::IsOkAndHolds(true));
-  auto first_topk = m::CustomCall(m::Reshape(m::Parameter(0)));
+  const HloInstruction* batch_topk = nullptr;
+  auto first_topk = m::CustomCall(&batch_topk, m::Reshape(m::Parameter(0)));
   auto slice_result = [&](auto input, size_t i) {
     return m::Reshape(m::Slice(m::GetTupleElement(input, i)));
   };
@@ -80,9 +82,12 @@ ENTRY cluster {
   auto sorted = m::Sort(
       m::Reshape(m::GetTupleElement(first_topk, 0)),
       m::Reshape(m::Add(m::GetTupleElement(first_topk, 1), index_correction)));
-  EXPECT_TRUE(
+  ASSERT_TRUE(
       Match(module->entry_computation()->root_instruction(),
             m::Tuple(slice_result(sorted, 0), slice_result(sorted, 1))));
+  ASSERT_NE(batch_topk, nullptr);
+  EXPECT_EQ(batch_topk->raw_backend_config_string(),
+            "{is_stable = false, order = \"PARTIAL\"}");
 }
 
 TEST_F(HardwareIndependentTopkSplitterTest, SplitsTopKNoBatchDimension) {
diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction_utils.cc b/third_party/xla/xla/hlo/ir/hlo_instruction_utils.cc
index 50b089a4659..c89f4745385 100644
--- a/third_party/xla/xla/hlo/ir/hlo_instruction_utils.cc
+++ b/third_party/xla/xla/hlo/ir/hlo_instruction_utils.cc
@@ -29,6 +29,7 @@ limitations under the License.
 #include "absl/status/status_macros.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
+#include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "re2/re2.h"
 #include "xla/comparison_util.h"
@@ -105,20 +106,30 @@ bool IsTopKStable(const HloCustomCallInstruction* inst) {
   return true;
 }
 
-std::optional<ComparisonOrder> GetTopKComparatorOrder(
-    const HloCustomCallInstruction* inst) {
-  if (inst == nullptr || !inst->has_to_apply()) {
-    return std::nullopt;
+void SetTopKStability(HloCustomCallInstruction* inst, bool is_stable) {
+  std::string replacement =
+      absl::StrCat("is_stable = ", is_stable ? "true" : "false");
+  std::string cfg = inst->raw_backend_config_string();
+  static const LazyRE2 kIsStableRegex = {R"((?i)is_stable\s*=\s*(true|false))"};
+  if (!RE2::Replace(&cfg, *kIsStableRegex, replacement)) {
+    static const LazyRE2 kNonEmptyDictEndRegex = {R"(([^\s\{])\s*\}\s*$)"};
+    if (!RE2::Replace(&cfg, *kNonEmptyDictEndRegex,
+                      absl::StrCat("\\1, ", replacement, "}"))) {
+      cfg = absl::StrCat("{", replacement, "}");
+    }
   }
-  const HloInstruction* root = inst->to_apply()->root_instruction();
-  if (root->opcode() == HloOpcode::kSelect &&
-      root->operand(0) == root->operand(2)) {
-    root = root->operand(1);
+  inst->set_raw_backend_config_string(cfg);
+}
+
+ComparisonOrder GetTopKOrder(const HloCustomCallInstruction* inst) {
+  static const LazyRE2 kOrderRegex = {R"(order\s*=\s*"?([A-Z]+))"};
+  absl::string_view order_str;
+  if (RE2::PartialMatch(inst->raw_backend_config_string(), *kOrderRegex,
+                        &order_str)) {
+    return ShortStringToComparisonOrder(order_str).value_or(
+        ComparisonOrder::kTotal);
   }
-  if (root->opcode() == HloOpcode::kCompare) {
-    return root->comparison_order();
-  }
-  return std::nullopt;
+  return ComparisonOrder::kTotal;
 }
 
 namespace async {
diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction_utils.h b/third_party/xla/xla/hlo/ir/hlo_instruction_utils.h
index 64ddd8847ce..57c9757d2ab 100644
--- a/third_party/xla/xla/hlo/ir/hlo_instruction_utils.h
+++ b/third_party/xla/xla/hlo/ir/hlo_instruction_utils.h
@@ -56,12 +56,13 @@ int32_t NestingDepth(const HloInstruction* hlo);
 // Checks if CustomCall TopK instruction is stable. Defaults to true.
 bool IsTopKStable(const HloCustomCallInstruction* inst);
 
-// Extracts the ComparisonOrder from a TopK custom-call comparator computation
-// (whose structure has already been validated by TopkRewriter::IsNanSafeGt:
-// either a root kCompare or Select(const_true, kCompare, const_true)).
-// Returns std::nullopt if the computation does not match either form.
-std::optional<ComparisonOrder> GetTopKComparatorOrder(
-    const HloCustomCallInstruction* inst);
+// Sets the is_stable attribute in the backend_config of a CustomCall TopK
+// instruction, preserving any other backend_config attributes.
+void SetTopKStability(HloCustomCallInstruction* inst, bool is_stable);
+
+// Returns the ComparisonOrder of a CustomCall TopK instruction. Defaults to
+// ComparisonOrder::kTotal.
+ComparisonOrder GetTopKOrder(const HloCustomCallInstruction* inst);
 
 namespace async {
 
diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction_utils_test.cc b/third_party/xla/xla/hlo/ir/hlo_instruction_utils_test.cc
index 5c7361d9039..861cfa6331f 100644
--- a/third_party/xla/xla/hlo/ir/hlo_instruction_utils_test.cc
+++ b/third_party/xla/xla/hlo/ir/hlo_instruction_utils_test.cc
@@ -700,11 +700,11 @@ TEST_F(HloInstructionUtilsTest, IsTopKStable) {
   EXPECT_TRUE(IsTopKStable(topk_default));
 }
 
-TEST_F(HloInstructionUtilsTest, GetTopKComparatorOrder) {
+TEST_F(HloInstructionUtilsTest, SetTopKStability) {
   const char* const hlo = R"(
     HloModule test
 
-    compare-total {
+    compare-gt.1 {
       p.1.lhs = s32[] parameter(2)
       p.1.rhs = s32[] parameter(3)
       p.0.lhs = f32[] parameter(0)
@@ -712,37 +712,46 @@ TEST_F(HloInstructionUtilsTest, GetTopKComparatorOrder) {
       ROOT compare = pred[] compare(p.0.lhs, p.0.rhs), direction=GT, order=TOTAL
     }
 
-    compare-partial {
-      p.1.lhs = s32[] parameter(2)
-      p.1.rhs = s32[] parameter(3)
-      p.0.lhs = f32[] parameter(0)
-      p.0.rhs = f32[] parameter(1)
-      ROOT compare = pred[] compare(p.0.lhs, p.0.rhs), direction=GT
+    ENTRY main {
+      arg = f32[1024] parameter(0)
+      topk_with_order = (f32[24], s32[24]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = true, order = "TOTAL"}
+      topk_order_only = (f32[24], s32[24]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={order = "TOTAL"}
+      topk_empty = (f32[24], s32[24]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}
+      ROOT tuple = tuple(topk_with_order, topk_order_only, topk_empty)
     }
+  )";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> m,
+                       ParseAndReturnVerifiedModule(hlo));
 
-    compare-select-total {
-      p.1.lhs = s32[] parameter(2)
-      p.1.rhs = s32[] parameter(3)
-      c = pred[] constant(true)
-      bcast = pred[] broadcast(c), dimensions={}
-      p.0.lhs = f32[] parameter(0)
-      p.0.rhs = f32[] parameter(1)
-      cmp = pred[] compare(p.0.lhs, p.0.rhs), direction=GT, order=TOTAL
-      ROOT select = pred[] select(bcast, cmp, bcast)
-    }
+  auto* topk_with_order = Cast<HloCustomCallInstruction>(
+      FindInstruction(m.get(), "topk_with_order"));
+  auto* topk_order_only = Cast<HloCustomCallInstruction>(
+      FindInstruction(m.get(), "topk_order_only"));
+  auto* topk_empty =
+      Cast<HloCustomCallInstruction>(FindInstruction(m.get(), "topk_empty"));
 
-    compare-select-weak {
-      p.1.lhs = s32[] parameter(2)
-      p.1.rhs = s32[] parameter(3)
-      c = pred[] constant(true)
-      bcast = pred[] broadcast(c), dimensions={}
-      p.0.lhs = f32[] parameter(0)
-      p.0.rhs = f32[] parameter(1)
-      cmp = pred[] compare(p.0.lhs, p.0.rhs), direction=GT, order=WEAK
-      ROOT select = pred[] select(bcast, cmp, bcast)
-    }
+  SetTopKStability(topk_with_order, false);
+  EXPECT_EQ(topk_with_order->raw_backend_config_string(),
+            "{is_stable = false, order = \"TOTAL\"}");
+  EXPECT_FALSE(IsTopKStable(topk_with_order));
+  EXPECT_EQ(GetTopKOrder(topk_with_order), ComparisonOrder::kTotal);
 
-    compare-invalid {
+  SetTopKStability(topk_order_only, false);
+  EXPECT_EQ(topk_order_only->raw_backend_config_string(),
+            "{order = \"TOTAL\", is_stable = false}");
+  EXPECT_FALSE(IsTopKStable(topk_order_only));
+  EXPECT_EQ(GetTopKOrder(topk_order_only), ComparisonOrder::kTotal);
+
+  SetTopKStability(topk_empty, false);
+  EXPECT_EQ(topk_empty->raw_backend_config_string(), "{is_stable = false}");
+  EXPECT_FALSE(IsTopKStable(topk_empty));
+}
+
+TEST_F(HloInstructionUtilsTest, GetTopKOrder) {
+  const char* const hlo = R"(
+    HloModule test
+
+    compare-true {
       p.1.lhs = s32[] parameter(2)
       p.1.rhs = s32[] parameter(3)
       p.0.lhs = f32[] parameter(0)
@@ -752,12 +761,11 @@ TEST_F(HloInstructionUtilsTest, GetTopKComparatorOrder) {
 
     ENTRY main {
       arg = f32[1024] parameter(0)
-      topk_total = (f32[24], s32[24]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-total}
-      topk_partial = (f32[24], s32[24]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-partial}
-      topk_select_total = (f32[24], s32[24]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-select-total}
-      topk_select_weak = (f32[24], s32[24]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-select-weak}
-      topk_invalid = (f32[24], s32[24]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-invalid}
-      ROOT tuple = tuple(topk_total, topk_partial, topk_select_total, topk_select_weak, topk_invalid)
+      topk_total = (f32[24], s32[24]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-true}, backend_config={is_stable = true, order = "TOTAL"}
+      topk_partial = (f32[24], s32[24]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-true}, backend_config={is_stable = false, order = "PARTIAL"}
+      topk_weak = (f32[24], s32[24]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-true}, backend_config={is_stable = true, order = "WEAK"}
+      topk_default = (f32[24], s32[24]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-true}
+      ROOT tuple = tuple(topk_total, topk_partial, topk_weak, topk_default)
     }
   )";
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> m,
@@ -767,19 +775,15 @@ TEST_F(HloInstructionUtilsTest, GetTopKComparatorOrder) {
       Cast<HloCustomCallInstruction>(FindInstruction(m.get(), "topk_total"));
   auto* topk_partial =
       Cast<HloCustomCallInstruction>(FindInstruction(m.get(), "topk_partial"));
-  auto* topk_select_total = Cast<HloCustomCallInstruction>(
-      FindInstruction(m.get(), "topk_select_total"));
-  auto* topk_select_weak = Cast<HloCustomCallInstruction>(
-      FindInstruction(m.get(), "topk_select_weak"));
-  auto* topk_invalid =
-      Cast<HloCustomCallInstruction>(FindInstruction(m.get(), "topk_invalid"));
+  auto* topk_weak =
+      Cast<HloCustomCallInstruction>(FindInstruction(m.get(), "topk_weak"));
+  auto* topk_default =
+      Cast<HloCustomCallInstruction>(FindInstruction(m.get(), "topk_default"));
 
-  EXPECT_EQ(GetTopKComparatorOrder(topk_total), ComparisonOrder::kTotal);
-  EXPECT_EQ(GetTopKComparatorOrder(topk_partial), ComparisonOrder::kPartial);
-  EXPECT_EQ(GetTopKComparatorOrder(topk_select_total), ComparisonOrder::kTotal);
-  EXPECT_EQ(GetTopKComparatorOrder(topk_select_weak), ComparisonOrder::kWeak);
-  EXPECT_EQ(GetTopKComparatorOrder(topk_invalid), std::nullopt);
-  EXPECT_EQ(GetTopKComparatorOrder(nullptr), std::nullopt);
+  EXPECT_EQ(GetTopKOrder(topk_total), ComparisonOrder::kTotal);
+  EXPECT_EQ(GetTopKOrder(topk_partial), ComparisonOrder::kPartial);
+  EXPECT_EQ(GetTopKOrder(topk_weak), ComparisonOrder::kWeak);
+  EXPECT_EQ(GetTopKOrder(topk_default), ComparisonOrder::kTotal);
 }
 
 TEST_F(HloInstructionUtilsTest, IsAllowedAsyncIntermediary) {
diff --git a/third_party/xla/xla/service/topk_rewriter.cc b/third_party/xla/xla/service/topk_rewriter.cc
index d63a8db8790..c782781fef1 100644
--- a/third_party/xla/xla/service/topk_rewriter.cc
+++ b/third_party/xla/xla/service/topk_rewriter.cc
@@ -54,9 +54,7 @@ limitations under the License.
 
 namespace xla {
 
-namespace m = match;
-
-static bool IsNanSafeGt(HloComputation* comp) {
+static std::optional<ComparisonOrder> MatchNanSafeGt(HloComputation* comp) {
   namespace m = match;
   auto match_bitcast_f32 = [](int64_t parameter_number) {
     auto param = m::Parameter(parameter_number)
@@ -164,10 +162,11 @@ static bool IsNanSafeGt(HloComputation* comp) {
     return instr->comparison_order() != ComparisonOrder::kWeak;
   };
 
+  HloInstruction* cmp = nullptr;
   auto match_compare = [&](PrimitiveType type) {
     auto param0 = m::Parameter(0).WithShape(m::Shape().WithElementType(type));
     auto param1 = m::Parameter(1).WithShape(m::Shape().WithElementType(type));
-    return m::Gt(param0, param1).WithPredicate(not_weak_order);
+    return m::Gt(&cmp, param0, param1).WithPredicate(not_weak_order);
   };
 
   auto match_default_compare = [&](PrimitiveType type) {
@@ -180,40 +179,48 @@ static bool IsNanSafeGt(HloComputation* comp) {
                      // Indices
                      params_with_type(2, S32), params_with_type(3, S32)});
     auto const_true = m::Broadcast(m::Constant());
-    auto values_gt = m::Gt(params[0], params[1]).WithPredicate(not_weak_order);
+    auto values_gt =
+        m::Gt(&cmp, params[0], params[1]).WithPredicate(not_weak_order);
     return m::Select(const_true, values_gt, const_true);
   };
 
   auto match_all_types = [](HloInstruction* root, auto callback) {
-    bool result = false;
     for (auto type : {BF16, F32, S32, U32}) {
-      result = result || Match(root, callback(type));
+      if (Match(root, callback(type))) {
+        return true;
+      }
     }
-    return result;
+    return false;
   };
 
-  return Match(comp->root_instruction(),
-               m::Gt(match_generic_iec559(0, F32, S32),
-                     match_generic_iec559(1, F32, S32))) ||
-         Match(comp->root_instruction(),
-               m::Gt(match_generic_iec559(0, BF16, S16),
-                     match_generic_iec559(1, BF16, S16))) ||
-         Match(comp->root_instruction(),
-               m::Gt(match_generic_iec559_with_convert(0, BF16, F32, S32),
-                     match_generic_iec559_with_convert(1, BF16, F32, S32))) ||
-         Match(comp->root_instruction(),
-               m::Gt(match_bitcast_f32(0), match_bitcast_f32(1))) ||
-         Match(comp->root_instruction(),
-               m::Gt(match_bitcast_bf16(0), match_bitcast_bf16(1))) ||
-         Match(comp->root_instruction(),
-               m::Gt(match_bitcast_f32_with_convert(0),
-                     match_bitcast_f32_with_convert(1))) ||
-         Match(comp->root_instruction(),
-               m::Gt(match_bitcast_bf16_with_convert(0),
-                     match_bitcast_bf16_with_convert(1))) ||
-         Match(comp->root_instruction(), m::Gt(match_s32(0), match_s32(1))) ||
-         match_all_types(comp->root_instruction(), match_compare) ||
-         match_all_types(comp->root_instruction(), match_default_compare);
+  if (Match(comp->root_instruction(),
+            m::Gt(match_generic_iec559(0, F32, S32),
+                  match_generic_iec559(1, F32, S32))) ||
+      Match(comp->root_instruction(),
+            m::Gt(match_generic_iec559(0, BF16, S16),
+                  match_generic_iec559(1, BF16, S16))) ||
+      Match(comp->root_instruction(),
+            m::Gt(match_generic_iec559_with_convert(0, BF16, F32, S32),
+                  match_generic_iec559_with_convert(1, BF16, F32, S32))) ||
+      Match(comp->root_instruction(),
+            m::Gt(match_bitcast_f32(0), match_bitcast_f32(1))) ||
+      Match(comp->root_instruction(),
+            m::Gt(match_bitcast_bf16(0), match_bitcast_bf16(1))) ||
+      Match(comp->root_instruction(),
+            m::Gt(match_bitcast_f32_with_convert(0),
+                  match_bitcast_f32_with_convert(1))) ||
+      Match(comp->root_instruction(),
+            m::Gt(match_bitcast_bf16_with_convert(0),
+                  match_bitcast_bf16_with_convert(1))) ||
+      Match(comp->root_instruction(), m::Gt(match_s32(0), match_s32(1)))) {
+    return ComparisonOrder::kTotal;
+  }
+  if ((match_all_types(comp->root_instruction(), match_compare) ||
+       match_all_types(comp->root_instruction(), match_default_compare)) &&
+      cmp != nullptr) {
+    return cmp->comparison_order();
+  }
+  return std::nullopt;
 }
 
 // Look for the instructions emitted from: xla/client/lib/sorting.cc
@@ -228,7 +235,8 @@ static bool HasIota(HloSortInstruction* sort, HloInstruction* data) {
          Match(sort->operand(1), m::Broadcast(match_iota(sort_dims)));
 }
 
-std::optional<int64_t> TopkRewriter::SortIsInTopK(HloInstruction* inst) {
+std::optional<int64_t> TopkRewriter::SortIsInTopK(HloInstruction* inst,
+                                                  ComparisonOrder* order) {
   HloSortInstruction* sort = DynCast<HloSortInstruction>(inst);
   if (sort == nullptr) {
     return std::nullopt;
@@ -241,9 +249,14 @@ std::optional<int64_t> TopkRewriter::SortIsInTopK(HloInstruction* inst) {
   if (sort->operand_count() == 2 && !HasIota(sort, data)) {
     return std::nullopt;
   }
-  if (!IsNanSafeGt(sort->to_apply())) {
+  std::optional<ComparisonOrder> matched_order =
+      MatchNanSafeGt(sort->to_apply());
+  if (!matched_order.has_value()) {
     return std::nullopt;
   }
+  if (order != nullptr) {
+    *order = *matched_order;
+  }
   const int64_t sort_dim = sort->sort_dimension();
 
   bool supported = true;
@@ -300,7 +313,8 @@ struct TopKCustomCall {
   HloInstruction* index_gte;
 };
 
-TopKCustomCall CreateTopKCustomCall(HloSortInstruction* sort, const int64_t k) {
+TopKCustomCall CreateTopKCustomCall(HloSortInstruction* sort, const int64_t k,
+                                    ComparisonOrder order) {
   HloInstruction* input = sort->mutable_operand(0);
   Shape data_shape = input->shape();
   PrimitiveType element_type = data_shape.element_type();
@@ -349,7 +363,8 @@ TopKCustomCall CreateTopKCustomCall(HloSortInstruction* sort, const int64_t k) {
   HloInstruction* topk = sort->AddInstruction(HloInstruction::CreateCustomCall(
       topk_shape, {input}, sort->to_apply(), "TopK"));
   topk->set_raw_backend_config_string(absl::StrFormat(
-      "{is_stable = %s}", sort->is_stable() ? "true" : "false"));
+      "{is_stable = %s, order = \"%s\"}", sort->is_stable() ? "true" : "false",
+      ComparisonOrderToShortString(order)));
   HloInstruction* value_gte =
       sort->AddInstruction(HloInstruction::CreateGetTupleElement(
           topk->shape().tuple_shapes(0), topk, 0));
@@ -383,7 +398,8 @@ TopKCustomCall CreateTopKCustomCall(HloSortInstruction* sort, const int64_t k) {
 absl::StatusOr<HloInstruction*> TopkRewriter::TransformPatternToCustomCall(
     HloInstruction* inst) {
   // Check if sort is in TopK.
-  std::optional<int64_t> k = SortIsInTopK(inst);
+  ComparisonOrder order;
+  std::optional<int64_t> k = SortIsInTopK(inst, &order);
   if (!k) {
     return nullptr;
   }
@@ -407,7 +423,7 @@ absl::StatusOr<HloInstruction*> TopkRewriter::TransformPatternToCustomCall(
     return nullptr;
   }
 
-  TopKCustomCall topkcc = CreateTopKCustomCall(sort, k.value());
+  TopKCustomCall topkcc = CreateTopKCustomCall(sort, k.value(), order);
 
   // If the slice starts at 0, its output matches topk_gte, so elide the slice.
   // Else, retarget the slice to take [start:k] directly from topk_gte.
diff --git a/third_party/xla/xla/service/topk_rewriter.h b/third_party/xla/xla/service/topk_rewriter.h
index 2a982584eff..2c082804afd 100644
--- a/third_party/xla/xla/service/topk_rewriter.h
+++ b/third_party/xla/xla/service/topk_rewriter.h
@@ -24,6 +24,7 @@ limitations under the License.
 #include "absl/container/flat_hash_set.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
+#include "xla/comparison_util.h"
 #include "xla/hlo/ir/dfs_hlo_visitor_with_default.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
@@ -49,7 +50,8 @@ class TopkRewriter : public HloModulePass {
       const absl::flat_hash_set<absl::string_view>& execution_threads) override;
 
   // Check if the sort instruction is in TopK.
-  std::optional<int64_t> SortIsInTopK(HloInstruction* inst);
+  std::optional<int64_t> SortIsInTopK(HloInstruction* inst,
+                                      ComparisonOrder* order = nullptr);
 
   // Transform to CustomCall.
   absl::StatusOr<bool> TransformToCustomCall(
diff --git a/third_party/xla/xla/service/topk_rewriter_test.cc b/third_party/xla/xla/service/topk_rewriter_test.cc
index 519ea60616b..379d076853d 100644
--- a/third_party/xla/xla/service/topk_rewriter_test.cc
+++ b/third_party/xla/xla/service/topk_rewriter_test.cc
@@ -518,7 +518,8 @@ ENTRY cluster {
         module->entry_computation()->root_instruction()->operand(0)->operand(0);
     ASSERT_THAT(cc->custom_call_target(), "TopK");
     if (!is_stable) {
-      EXPECT_EQ(cc->raw_backend_config_string(), "{is_stable = false}");
+      EXPECT_EQ(cc->raw_backend_config_string(),
+                "{is_stable = false, order = \"TOTAL\"}");
     }
   };
   // Start by producing a TopK...
@@ -879,7 +880,45 @@ ENTRY cluster {
       module->entry_computation()->root_instruction()->operand(0)->operand(0);
 
   EXPECT_EQ(cc->custom_call_target(), "TopK");
-  EXPECT_EQ(cc->raw_backend_config_string(), "{is_stable = false}");
+  EXPECT_EQ(cc->raw_backend_config_string(),
+            "{is_stable = false, order = \"TOTAL\"}");
+}
+
+TEST_F(TopkRewriterTest, TopKCustomCallPartialOrderConfig) {
+  const std::string hlo_string = R"(
+HloModule module
+%compare {
+  %Arg_0.100 = f32[] parameter(0)
+  %Arg_1.101 = f32[] parameter(1)
+  %Arg_2.102 = s32[] parameter(2)
+  %Arg_3.103 = s32[] parameter(3)
+  ROOT %compare.56364 = pred[] compare(f32[] %Arg_0.100, f32[] %Arg_1.101), direction=GT, order=PARTIAL
+}
+ENTRY cluster {
+  %arg_tuple.1 = f32[8,2048] parameter(0)
+  %iota.4 = s32[8,2048] iota(), iota_dimension=1
+  %sort.27 = (f32[8,2048], s32[8,2048]) sort(%arg_tuple.1, %iota.4),
+    dimensions={1}, is_stable=true, to_apply=%compare
+  %get-tuple-element.28 = f32[8,2048] get-tuple-element(%sort.27), index=0
+  %slice.29 = f32[8,24] slice(%get-tuple-element.28), slice={[0:8], [0:24]}
+  %get-tuple-element.30 = s32[8,2048] get-tuple-element(%sort.27), index=1
+  %slice.31 = s32[8,24] slice(%get-tuple-element.30), slice={[0:8], [0:24]}
+  ROOT %tuple.32 = (f32[8,24], s32[8,24]) tuple(%slice.29, %slice.31)
+})";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+
+  TopkRewriter rewriter(
+      [](const HloSortInstruction*, int64_t) { return true; });
+  ASSERT_OK_AND_ASSIGN(bool changed, rewriter.Run(module.get()));
+  ASSERT_OK(HloDCE().Run(module.get()).status());
+  EXPECT_TRUE(changed);
+
+  const HloInstruction* cc =
+      module->entry_computation()->root_instruction()->operand(0)->operand(0);
+
+  EXPECT_EQ(cc->custom_call_target(), "TopK");
+  EXPECT_EQ(cc->raw_backend_config_string(),
+            "{is_stable = true, order = \"PARTIAL\"}");
 }
 
 TEST_F(TopkRewriterTest, RewriteNonZeroStartSlice) {

From acfd4d839266bff730d9f8160c6824707d0da33f Mon Sep 17 00:00:00 2001
From: Blake Hechtman <blakehechtman@google.com>
Date: Fri, 9 Oct 2026 10:45:55 -0700
Subject: [PATCH 0378/1000] Preserve constraint priority across buffer and
 operand layout updates

When ConstrainOrOverrideBuffersAtIndex overrides a buffer and its user operand
layout constraints at an elevated priority (existing->priority() + 1), a later
negotiated buffer layout update in BufferLayoutConstraint::UpdateLayout
previously overwrote priority_ with the incoming lower priority, and
PropagateBufferConstraintToOperands propagated at current_priority_ instead of
respecting the buffer and existing operand constraint priorities. As a result,
OperandLayoutConstraint::UpdateLayout rejected the updated output layout on
higher priority operands of instructions that cannot change layout, leaving the
instruction output and operands with mismatched layouts.

Fix this by:
- Preserving std::max(priority_, priority) in
  BufferLayoutConstraint::UpdateLayout and OperandLayoutConstraint::UpdateLayout
  (including when a new operand constraint matches the existing layout).
- Propagating with std::max(current_priority_, buffer_constraint.priority()) in
  PropagateBufferConstraintToOperands, and using at least the existing operand
  or buffer constraint priority when synchronizing operands and outputs of
  instructions that cannot change layout.
- Comparing operand indices rather than operand instruction pointers in the
  sibling loop of PropagateOperandConstraint.
- Verifying in LayoutAssignment::CheckLayouts that instructions that cannot
  change layout have matching layouts on their output and same rank array
  operands.

PiperOrigin-RevId: 996624981
---
 .../xla/xla/service/layout_assignment.cc      | 85 ++++++++++++++-----
 .../xla/xla/service/layout_assignment_test.cc | 81 ++++++++++++++++++
 2 files changed, 144 insertions(+), 22 deletions(-)

diff --git a/third_party/xla/xla/service/layout_assignment.cc b/third_party/xla/xla/service/layout_assignment.cc
index 0af92ed55f5..c23e9c2b783 100644
--- a/third_party/xla/xla/service/layout_assignment.cc
+++ b/third_party/xla/xla/service/layout_assignment.cc
@@ -146,7 +146,7 @@ bool BufferLayoutConstraint::UpdateLayout(int64_t priority,
   }
   mandatory_ = mandatory;
   dfs_ = dfs;
-  priority_ = priority;
+  priority_ = std::max(priority_, priority);
   from_user_ = user;
   layout_.push_back(layout());
   layout_[0] = new_layout;
@@ -175,6 +175,9 @@ bool OperandLayoutConstraint::UpdateLayout(int64_t new_priority,
                                            LayoutAssignment* assignment) {
   if (shape_layout().MatchesLayoutInShape(new_shape,
                                           /*minor_to_major_only=*/true)) {
+    if (priority_ < new_priority) {
+      priority_ = new_priority;
+    }
     VLOG(3) << "SUCC b/c the new layout matches the existing one.";
     // New constraint matches existing constraint. Nothing to do.
     return false;
@@ -209,7 +212,7 @@ bool OperandLayoutConstraint::UpdateLayout(int64_t new_priority,
   VLOG(3) << "Updating existing Operand layout:" << ToString();
   mandatory_ = mandatory;
   dfs_ = dfs;
-  priority_ = new_priority;
+  priority_ = std::max(priority_, new_priority);
   shape_layout_.push_back(shape_layout_[0]);
   shape_layout_[0] = ShapeLayout(new_shape);
   return true;
@@ -391,6 +394,7 @@ absl::Status LayoutAssignment::SetOperandLayout(
                                          dfs, this)) {
       return absl::OkStatus();
     }
+    priority = curr_shape_layout->priority();
   }
   if (curr_shape_layout == nullptr) {
     curr_shape_layout = std::make_unique<OperandLayoutConstraint>(
@@ -1486,6 +1490,24 @@ absl::Status LayoutAssignment::CheckLayouts(
         default:
           break;
       }
+      if (!InstructionCanChangeLayoutInstance(instruction) &&
+          instruction->shape().IsArray() &&
+          instruction->shape().dimensions().size() > 1) {
+        for (int64_t operand_no = 0; operand_no < instruction->operand_count();
+             ++operand_no) {
+          const HloInstruction* operand = instruction->operand(operand_no);
+          if (operand->shape().IsArray() &&
+              operand->shape().dimensions().size() ==
+                  instruction->shape().dimensions().size()) {
+            TF_RET_CHECK(
+                LayoutsInShapesEqual(operand->shape(), instruction->shape()))
+                << "Operand " << operand_no << " of " << instruction->ToString()
+                << " has different layout from output: "
+                << ShapeUtil::HumanStringWithLayout(operand->shape()) << " vs "
+                << ShapeUtil::HumanStringWithLayout(instruction->shape());
+          }
+        }
+      }
     }
   }
   // Finally verify the result layout, if set, matches the layout of the entry
@@ -1986,7 +2008,7 @@ absl::Status LayoutAssignment::PropagateOperandConstraint(
         new_layout, buffer,
         preserve_non_mandatory_if_same_layout(buffer, new_layout,
                                               /*default_mandatory=*/true),
-        /*dfs=*/true));
+        /*dfs=*/true, operand_constraint.priority()));
   }
 
   if (InstructionCanChangeLayoutInstance(user) && !user->shape().IsArray() &&
@@ -2025,7 +2047,7 @@ absl::Status LayoutAssignment::PropagateOperandConstraint(
     // Make sure all siblings have the same layout as the operand.
     for (int64_t operand_no = 0; operand_no < user->operand_count();
          ++operand_no) {
-      if (user->operand(operand_no) == operand) {
+      if (operand_no == operand_constraint.operand_no()) {
         continue;
       }
       const HloInstruction* sibling = user->operand(operand_no);
@@ -2039,9 +2061,16 @@ absl::Status LayoutAssignment::PropagateOperandConstraint(
       if (operand_rank != sibling_rank) {
         continue;
       }
+      const OperandLayoutConstraint* existing_sibling_constraint =
+          constraints->GetOperandLayoutConstraint(user, operand_no);
+      const int64_t sibling_priority =
+          existing_sibling_constraint != nullptr
+              ? std::max(operand_constraint.priority(),
+                         existing_sibling_constraint->priority())
+              : operand_constraint.priority();
       ABSL_RETURN_IF_ERROR(SetArrayOperandLayout(
           operand_constraint.shape_layout().layout(), user, operand_no,
-          /*mandatory=*/true, /*dfs=*/true, operand_constraint.priority()));
+          /*mandatory=*/true, /*dfs=*/true, sibling_priority));
     }
     ABSL_RETURN_IF_ERROR(ShapeUtil::ForEachSubshapeWithStatus(
         user->shape(),
@@ -2068,16 +2097,18 @@ absl::Status LayoutAssignment::PropagateOperandConstraint(
           const HloValue& buffer =
               dataflow_analysis_->GetValueDefinedAt(user, shape_index);
           const Layout& new_layout = operand_constraint.shape_layout().layout();
-          // If we already have a constraint for the buffer it was assigned but
-          // hasn't propagated yet. This can happen with diamond-shaped graphs
-          // where one path is first evaluated in depth-first order (we're here)
-          // and the other path is propagated later. We don't set the layout
-          // here as it will always be overwritten later.
+          const BufferLayoutConstraint* existing_buffer_constraint =
+              GetBufferLayoutConstraint(buffer);
+          const int64_t buffer_priority =
+              existing_buffer_constraint != nullptr
+                  ? std::max(operand_constraint.priority(),
+                             existing_buffer_constraint->priority())
+                  : operand_constraint.priority();
           ABSL_RETURN_IF_ERROR(SetBufferLayout(
               new_layout, buffer,
               preserve_non_mandatory_if_same_layout(buffer, new_layout,
                                                     /*default_mandatory=*/true),
-              /*dfs=*/true, operand_constraint.priority()));
+              /*dfs=*/true, buffer_priority));
           return absl::OkStatus();
         }));
     return absl::OkStatus();
@@ -2125,11 +2156,21 @@ absl::Status LayoutAssignment::PropagateBufferConstraintToOperands(
   VLOG(5) << "PropagateBufferConstraintToOperands: "
           << buffer_constraint.ToString();
 
+  const int64_t priority =
+      std::max(current_priority_, buffer_constraint.priority());
+  auto priority_for_operand = [&](int64_t op_no) {
+    const OperandLayoutConstraint* existing_op_constraint =
+        constraints->GetOperandLayoutConstraint(instruction, op_no);
+    return existing_op_constraint != nullptr
+               ? std::max(priority, existing_op_constraint->priority())
+               : priority;
+  };
   if (instruction->opcode() == HloOpcode::kAllReduce) {
+    const int64_t op_no =
+        instruction->operand_count() == 1 ? 0 : buffer.index()[0];
     ABSL_RETURN_IF_ERROR(SetArrayOperandLayout(
-        buffer_constraint.layout(), instruction,
-        instruction->operand_count() == 1 ? 0 : buffer.index()[0],
-        /*mandatory=*/true, /*dfs=*/true, buffer_constraint.priority()));
+        buffer_constraint.layout(), instruction, op_no,
+        /*mandatory=*/true, /*dfs=*/true, priority_for_operand(op_no)));
     return absl::OkStatus();
   }
   for (int64_t operand_no = 0; operand_no < instruction->operand_count();
@@ -2145,7 +2186,8 @@ absl::Status LayoutAssignment::PropagateBufferConstraintToOperands(
             LayoutUtil::MinorToMajor(buffer_constraint.layout()).size()) {
           ABSL_RETURN_IF_ERROR(SetArrayOperandLayout(
               buffer_constraint.layout(), instruction, operand_no,
-              /*mandatory=*/true, /*dfs=*/true, current_priority_));
+              /*mandatory=*/true,
+              /*dfs=*/true, priority_for_operand(operand_no)));
         } else if (instruction->opcode() == HloOpcode::kBitcastConvert) {
           Shape shape = instruction->shape();
           if (operand->shape().dimensions().size() <
@@ -2156,9 +2198,10 @@ absl::Status LayoutAssignment::PropagateBufferConstraintToOperands(
             ShapeUtil::AppendMinorDimension(
                 operand->shape().dimensions().back(), &shape);
           }
-          ABSL_RETURN_IF_ERROR(SetArrayOperandLayout(
-              shape.layout(), instruction, operand_no,
-              /*mandatory=*/true, /*dfs=*/true, current_priority_));
+          ABSL_RETURN_IF_ERROR(
+              SetArrayOperandLayout(shape.layout(), instruction, operand_no,
+                                    /*mandatory=*/true, /*dfs=*/true,
+                                    priority_for_operand(operand_no)));
         }
       }
 
@@ -2168,8 +2211,7 @@ absl::Status LayoutAssignment::PropagateBufferConstraintToOperands(
       ABSL_RETURN_IF_ERROR(SetArrayOperandLayout(
           layout, instruction, operand_no, /*mandatory=*/true,
           /*dfs=*/
-          InstructionShouldPropagateDepthFirst(*instruction),
-          current_priority_));
+          InstructionShouldPropagateDepthFirst(*instruction), priority));
     } else {
       if (!buffer.IsTopLevel() ||
           !instruction->operand(operand_no)->shape().IsArray()) {
@@ -2185,8 +2227,7 @@ absl::Status LayoutAssignment::PropagateBufferConstraintToOperands(
             *operand_layout, instruction, operand_no,
             /*mandatory=*/OutputLayoutAlwaysPropagateToOperands(instruction),
             /*dfs=*/
-            InstructionShouldPropagateDepthFirst(*instruction),
-            current_priority_));
+            InstructionShouldPropagateDepthFirst(*instruction), priority));
       }
     }
   }
diff --git a/third_party/xla/xla/service/layout_assignment_test.cc b/third_party/xla/xla/service/layout_assignment_test.cc
index e0b347b67ef..2ef9d16bc48 100644
--- a/third_party/xla/xla/service/layout_assignment_test.cc
+++ b/third_party/xla/xla/service/layout_assignment_test.cc
@@ -3525,5 +3525,86 @@ ENTRY %main (x: f32[256,3,1], y: f32[256,3,1,3]) -> f32[256,3,1,4] {
   ExpectLayoutIs(concat->operand(1)->shape(), {3, 2, 1, 0});
 }
 
+class NegotiatingLayoutAssignment : public LayoutAssignment {
+ public:
+  explicit NegotiatingLayoutAssignment(
+      ComputationLayout* entry_computation_layout)
+      : LayoutAssignment(entry_computation_layout) {}
+
+  bool NegotiateLayout(const HloInstruction* instruction,
+                       const Layout& new_layout, const Layout& existing_layout,
+                       const HloInstruction* user,
+                       const HloInstruction* from_user) override {
+    if (instruction->opcode() == HloOpcode::kAdd &&
+        Layout::Equal().MinorToMajorOnly()(new_layout,
+                                           LayoutUtil::MakeLayout({0, 1}))) {
+      return true;
+    }
+    return false;
+  }
+
+  Layout GetUnconstrainedLayout(const HloValue& buffer) override {
+    if (buffer.shape().dimensions().size() == 2) {
+      return LayoutUtil::MakeLayout({0, 1});
+    }
+    return LayoutAssignment::GetUnconstrainedLayout(buffer);
+  }
+};
+
+TEST_F(LayoutAssignmentTest,
+       NegotiatedOutputLayoutUpdatesHigherPriorityElementwiseOperands) {
+  const char* module_str = R"hlo(
+HloModule NegotiatedOutputLayoutUpdatesHigherPriorityElementwiseOperands
+
+%callee (p0: f32[1024,32], p1: f32[1024,32], v: f32[1024]) -> (f32[1024,96], f32[1024,32], f32[1024,32]) {
+  %p0 = f32[1024,32] parameter(0)
+  %p1 = f32[1024,32] parameter(1)
+  %v = f32[1024] parameter(2)
+  %neg0 = f32[1024,32] negate(%p0)
+  %neg1 = f32[1024,32] negate(%p1)
+  %hint0 = f32[1024,32]{0,1} custom-call(%neg0), custom_call_target="Hint",
+    operand_layout_constraints={f32[1024,32]{0,1}}
+  %hint1 = f32[1024,32]{0,1} custom-call(%neg1), custom_call_target="Hint",
+    operand_layout_constraints={f32[1024,32]{0,1}}
+  %add = f32[1024,32] add(%p0, %p1)
+  %big = f32[1024,64] broadcast(%v), dimensions={0}
+  %concat = f32[1024,96] concatenate(%add, %big), dimensions={1}
+  ROOT %out = (f32[1024,96], f32[1024,32], f32[1024,32]) tuple(%concat, %hint0, %hint1)
+}
+
+ENTRY %main (x: f32[1024,32], y: f32[1024,32], v: f32[1024]) -> (f32[1024,96], f32[1024,32], f32[1024,32]) {
+  %x = f32[1024,32]{1,0} parameter(0)
+  %y = f32[1024,32]{1,0} parameter(1)
+  %v = f32[1024]{0} parameter(2)
+  %pinned_x = f32[1024,32]{1,0} custom-call(%x), custom_call_target="Pinned",
+    operand_layout_constraints={f32[1024,32]{1,0}}
+  %pinned_y = f32[1024,32]{1,0} custom-call(%y), custom_call_target="Pinned",
+    operand_layout_constraints={f32[1024,32]{1,0}}
+  ROOT %call = (f32[1024,96], f32[1024,32], f32[1024,32]) call(%pinned_x, %pinned_y, %v), to_apply=%callee
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> m,
+                       ParseAndReturnVerifiedModule(module_str));
+  ComputationLayout computation_layout(
+      m->entry_computation()->ComputeProgramShape(), /*ignore_layouts=*/true);
+  *computation_layout.mutable_parameter_layout(0) =
+      ShapeLayout(ShapeUtil::MakeShapeWithDenseLayout(F32, {1024, 32}, {1, 0}));
+  *computation_layout.mutable_parameter_layout(1) =
+      ShapeLayout(ShapeUtil::MakeShapeWithDenseLayout(F32, {1024, 32}, {1, 0}));
+  *computation_layout.mutable_parameter_layout(2) =
+      ShapeLayout(ShapeUtil::MakeShapeWithDenseLayout(F32, {1024}, {0}));
+
+  NegotiatingLayoutAssignment layout_assignment(&computation_layout);
+  EXPECT_THAT(
+      layout_assignment.Run(m.get(), {HloInstruction::kMainExecutionThread}),
+      absl_testing::IsOk());
+  const HloInstruction* add = FindInstruction(m.get(), "add");
+  ASSERT_NE(add, nullptr);
+  ExpectLayoutIs(add->shape(), {0, 1});
+  ExpectLayoutIs(add->operand(0)->shape(), {0, 1});
+  ExpectLayoutIs(add->operand(1)->shape(), {0, 1});
+}
+
 }  // namespace
 }  // namespace xla

From 1e1629fd562d06d7c8d485bba0942b76e5f96685 Mon Sep 17 00:00:00 2001
From: David Dunleavy <ddunleavy@google.com>
Date: Fri, 9 Oct 2026 11:13:06 -0700
Subject: [PATCH 0379/1000] Move `optimization_parameters.proto` and
 `tpu_embdding_configuration.proto` to `xla/tpu`

PiperOrigin-RevId: 996642383
---
 tensorflow/core/protobuf/tpu/BUILD            |  12 +-
 .../tpu/optimization_parameters.proto         | 639 +----------------
 .../tpu/tpu_embedding_configuration.proto     | 158 +----
 third_party/xla/xla/tpu/protobuf/BUILD        |  60 ++
 .../protobuf/optimization_parameters.proto    | 657 ++++++++++++++++++
 .../tpu_embedding_configuration.proto         | 176 +++++
 6 files changed, 905 insertions(+), 797 deletions(-)
 create mode 100644 third_party/xla/xla/tpu/protobuf/BUILD
 create mode 100644 third_party/xla/xla/tpu/protobuf/optimization_parameters.proto
 create mode 100644 third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration.proto

diff --git a/tensorflow/core/protobuf/tpu/BUILD b/tensorflow/core/protobuf/tpu/BUILD
index 35a385fdb85..08bf50b04d7 100644
--- a/tensorflow/core/protobuf/tpu/BUILD
+++ b/tensorflow/core/protobuf/tpu/BUILD
@@ -32,9 +32,12 @@ tf_proto_library(
         "tpu_embedding_configuration.proto",
     ],
     protodeps = [
-        ":optimization_parameters_proto",
+        "@xla//xla/tpu/protobuf:tpu_embedding_configuration_proto",
     ],
     visibility = ["//visibility:public"],
+    exports = [
+        "@xla//xla/tpu/protobuf:tpu_embedding_configuration_proto",
+    ],
 )
 
 tf_proto_library(
@@ -42,8 +45,13 @@ tf_proto_library(
     srcs = [
         "optimization_parameters.proto",
     ],
-    protodeps = ["@xla//xla/service:hlo_proto"],
+    protodeps = [
+        "@xla//xla/tpu/protobuf:optimization_parameters_proto",
+    ],
     visibility = ["//visibility:public"],
+    exports = [
+        "@xla//xla/tpu/protobuf:optimization_parameters_proto",
+    ],
 )
 
 tf_proto_library(
diff --git a/tensorflow/core/protobuf/tpu/optimization_parameters.proto b/tensorflow/core/protobuf/tpu/optimization_parameters.proto
index 6899a183f58..869e9222049 100644
--- a/tensorflow/core/protobuf/tpu/optimization_parameters.proto
+++ b/tensorflow/core/protobuf/tpu/optimization_parameters.proto
@@ -17,641 +17,4 @@ syntax = "proto3";
 
 package tensorflow.tpu;
 
-import "google/protobuf/wrappers.proto";
-import "xla/service/hlo.proto";
-
-message ClippingLimits {
-  google.protobuf.FloatValue lower = 1;  // -inf if not set
-  google.protobuf.FloatValue upper = 2;  // +inf if not set
-}
-
-// Configuration for simulated quantization; simulated quantization is used to
-// reduce training/serving skew when the serving variables are quantized. The
-// same quantization operations are executed during training to minimize
-// differences with serving.
-//
-// Simulated quantization inserts the following operations on the forward pass
-// after gathering the embedding vector from HBM. The backward pass operations
-// are unchanged.
-//
-// clipped_val = clip(input, clipping_limits)
-// quantum = clipping_limits.range() / (num_buckets - 1)
-// quantized_val = floor((clipped_val - clipping_limits.lower()) / quantum + .5)
-// return quantized_val * quantum + clipping_limits.lower().
-message SimulatedQuantization {
-  // Whether simulated quantization is enabled.
-  bool enabled = 1;
-
-  // Minimum and maximum values of the range used for quantization.
-  ClippingLimits clipping_limits = 2;
-
-  // Number of possible quantized values.
-  int32 num_buckets = 3;
-}
-
-// Dynamic input specification for optimizers in the TPUEmbeddingConfiguration.
-// The actual dynamic inputs are provided as a scalar input list to the
-// SendTPUEmbeddingGradients Op indexed by their tag specified through the
-// following proto.
-message OptimizerDynamicInput {
-  // For tables where dynamic inputs are needed (e.g., learning rates or other
-  // dynamic hyperparameters used in optimizers), a tag must be specified for
-  // the input.
-  //
-  // The tag must be a non-negative  integer. All tags in the range
-  // [0, number_of_unique_tags) must be present in the TPU embedding
-  // configuration, i.e. a tag cannot be skipped if a different tag numerically
-  // greater than it is used in the configuration.
-  //
-  // If multiple tables specify the same tag, they *MUST* have
-  // the same dynamic input, for example, their dynamic learning rate could be
-  // computed by the same TensorFlow sub-graph. The partitioning of the
-  // embedding layer would be more optimal if the number_of_unique_tags is as
-  // *LOW* as possible, i.e., if many tables share the same tag.
-  //
-  // The hyper_parameters input of the SendTPUEmbeddingGradients op is used to
-  // communicate dynamic hyper-parameters to the TPU embedding program.
-  // The hyper_parameters input is a list of scalars where the size of the list
-  // is equal to the number of unique tags. The hyper-parameter associated with
-  // a particular tag is specified by populating its corresponding index in the
-  // list of scalars.
-  int32 tag = 1;
-}
-
-// Source of learning rate to use.
-message LearningRate {
-  oneof learning_rate {
-    float constant = 1;
-    OptimizerDynamicInput dynamic = 2;
-  }
-}
-
-// Each optimizer's parameter proto has a link to its documentation and CPU
-// implementation (if available) for user reference.
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adagrad
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L1634
-message AdagradParameters {
-  // Old initial accumulator parameter.
-  reserved "initial_accumulator";
-  reserved 1;
-}
-
-// This optimizer combines the Adagrad and Momentum update rules.
-// accum(new) = beta2 == 1.0 ?
-//              accum(old) + grad^2 :
-//              beta2 * accum(old) + (1 - beta2) * grad^2
-// accum_with_exponent = (accum(new) + epsilon)^(-1.0 / exponent)
-// mom_accum(new) = momentum * mom_accum(old) + accum_with_exponent
-// update = use_nesterov ?
-//          momentum * mom_accum(new) + accum_with_exponent :
-//          mom_accum(new)
-// var(new) = var(old) - lr * grad * update
-// Algorithm described in https://arxiv.org/abs/2002.11803.
-message AdagradMomentumParameters {
-  // Moving average parameter for the momentum accumulator.
-  float momentum = 1;
-
-  // Whether to use the Nesterov variant of momentum.
-  bool use_nesterov = 2;
-
-  // Exponent for the gradient^2 accumulator.
-  float exponent = 3;
-
-  // Moving average parameter for the gradient^2 accumulator.
-  float beta2 = 4;
-
-  // Offset added to the Adagrad accumulator.
-  float epsilon = 5;
-}
-
-// Algorithm in http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf.
-message BoundedAdagradParameters {
-  // Whether to use the updated or the old value of the accumulator when
-  // computing the effective learning rate. When update_accumulator_first is set
-  // to True, the updated value of the accumulator is used.
-  bool update_accumulator_first = 1;
-
-  // The max_var_update value to use. Set value to 0 (default) to disable using
-  // max_var_update to clip the gradient.
-  float max_var_update = 2;
-
-  // The maximum value of the accumulator. Set max_accumulator to 0 (default)
-  // to disable using max_accumulator to clip the accumulator.
-  float max_accumulator = 3;
-}
-
-// Frequency Aware Adagrad optimizer. This optimizer implements the AdaGrad
-// algorithm and further allows to:
-// * Scale the learning rate based on frequency of the update. Sparsely updated
-//   rows are updated with a higher effective learning rate, and frequently
-//   updated rows are updated with a lower effective learning rate.
-// * Decay the growth of the accumulator values.
-// * Use L1 / L2 regularization for the weight updates.
-//
-// The optimization algorithm is shown below.
-// counter(new) = counter(old) + 1
-// accum(new) = max(accumulator_decay * accum(old) + grad^2,
-//                  initial_accumulator_value)
-// lr_scale = min((step_counter / accum(new)) ^ probability_exponent,
-// max_lr_multiplier) update = grad * lr_scale / sqrt(accum(new)) if
-// (l1_regularization_strength > 0.0):
-//   update = update + l1_regularization_strength * sign(var(old))
-// if (l2_regularization_strength > 0.0):
-//   update = update + l2_regularization_strength * var(old)
-// var(new) = var(old) - lr_scale * grad * update
-
-message FrequencyAwareAdagradParameters {
-  // The L1 regularization parameter for adjusting the update based on the sign
-  // of the variable.
-  float l1_regularization_strength = 1;
-
-  // The L2 regularization parameter for adjusting the update based on the
-  // variable.
-  float l2_regularization_strength = 2;
-
-  // The exponent used for scaling the learning rate based on the sparsity of
-  // updates.
-  float probability_exponent = 4;
-
-  // The maximum value of the learning rate scale.
-  float max_lr_multiplier = 3;
-
-  // The decay for the Adagrad accumulator.
-  float accumulator_decay = 5;
-
-  // The initial and minimum value for the Adagrad accumulator.
-  float initial_accumulator_value = 6;
-
-  // The tag for identifying the step counter used for the frequency aware
-  // Adagrad optimizer.
-  OptimizerDynamicInput step_counter = 7;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L629
-message StochasticGradientDescentParameters {}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Ftrl
-// https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/41159.pdf
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L2646
-//
-// The hyperparameters for FTRL are the same as for the Keras implementation,
-// with some additions. The "beta" parameter matches the behavior described in
-// the second link above; "beta" / (2 * learning rate) should be added to "l2"
-// to get equivalent behavior in the other TensorFlow implementations of this
-// optimizer. When the multiply_linear_by_lr field is set to true, a modified
-// formula is used for FTRL that treats the "linear" accumulator as being
-// pre-multiplied by the learning rate (i.e., the accumulator named "linear"
-// actually stores "linear * learning_rate"). Other than checkpoint
-// compatibility, this is mathematically equivalent for a static learning rate;
-// for a dynamic learning rate, it is nearly the same as long as the learning
-// rate does not change quickly. The benefit of setting multiply_linear_by_lr to
-// true is that the modified formula handles zero and near-zero learning rates
-// without producing NaNs, improving flexibility for learning rate ramp-up.
-message FtrlParameters {
-  float l1 = 1;
-  float l2 = 2;
-  float lr_power = 3;
-  float beta = 7;
-  bool multiply_linear_by_lr = 6;
-
-  // Previously, allow_zero_accumulator parameter changed some internal formulas
-  // to allow zero and near-zero accumulator values at the cost of some
-  // performance. The current implementation ignores this parameter; zero or
-  // near-zero accumulator values are now always supported.
-  bool allow_zero_accumulator = 8 [deprecated = true];
-
-  // Old initial accumulator parameters.
-  reserved "initial_accum", "initial_linear";
-  reserved 4, 5;
-}
-
-// The Adam optimizer does not implement hyper-parameter update due to hardware
-// limitations; use the dynamic learning rate feature instead, setting the
-// learning rate to: user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
-// Here, t is the current timestep.
-//
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adam
-// https://github.com/tensorflow/tensorflow/blob/ab51450c817674c8ff08a7ae4f8ac50cdc4bed8b/tensorflow/python/training/adam.py#L32
-//
-// Note that the code by default implements the lazy version of Adam
-// (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/LazyAdamOptimizer)
-// unless the use_non_lazy_adam parameter is set, in which case it implements
-// the normal version of Adam that updates all parameters in the embedding
-// table, even for entries that are not used in the current minibatch
-// (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/AdamOptimizer). If
-// use_non_lazy_adam is enabled, gradient accumulation is also required to be
-// enabled in order to get correct results; a warning will be printed otherwise
-// (which may change to an error in the future). If use_sum_inside_sqrt is set,
-// the Adam variable update formula will be changed from m / (sqrt(v) + epsilon)
-// to m / sqrt(v + epsilon**2); this option improves the performance of TPU
-// training and is not expected to harm model quality.
-message AdamParameters {
-  float beta1 = 3;
-  float beta2 = 4;
-  float epsilon = 5;
-  bool use_non_lazy_adam = 8;
-  bool use_sum_inside_sqrt = 10;
-
-  // Old initial accumulator parameters.
-  reserved "initial_m", "initial_v";
-  reserved 6, 7;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L3068
-message MomentumParameters {
-  float momentum = 1;
-  bool use_nesterov = 2;
-
-  // Old initial accumulator parameter.
-  reserved "initial_accum";
-  reserved 3;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Lion
-// momenta(new) = beta2 * momenta(old) + (1 - beta2) * grad
-// momenta_t = beta1 * momenta(old) + (1 - beta1) * grad
-// var(new) = var(old) - lr * sign(momenta_t)
-// Algorithm described in https://arxiv.org/abs/2302.06675.
-message LionParameters {
-  float beta1 = 1;
-  float beta2 = 2;
-  bool use_non_lazy_lion = 3;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L4229
-message RmsPropParameters {
-  float rho = 1;
-  float momentum = 2;
-  float epsilon = 3;
-
-  // Old initial accumulator parameters.
-  reserved "initial_ms", "initial_mom";
-  reserved 4, 5;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L4358
-message CenteredRmsPropParameters {
-  float rho = 1;
-  float momentum = 2;
-  float epsilon = 3;
-
-  // Old initial accumulator parameters.
-  reserved "initial_ms", "initial_mom", "initial_mg";
-  reserved 4, 5, 6;
-}
-
-// Variant of algorithm in http://proceedings.mlr.press/v44/shamir15.pdf
-message MdlAdagradLightParameters {
-  float l2 = 1;
-  float lr_power = 2;
-  float min_servable_mdl_benefit = 3;
-  float mdl_mix_in_margin = 4;
-  float mdl_benefit_rampup_coeff = 5;
-  float mdl_min_weight = 6;
-  float benefit_revisit_scale = 7;
-  float max_event_benefit = 8;
-  float max_total_benefit = 9;
-  float mdl_hard_limit = 10;
-  bool hard_limit_min_benefit = 11;
-  bool mdl_regularize = 12;
-
-  // Old initial accumulator parameters.
-  reserved "initial_accumulator", "initial_weight", "initial_benefit";
-  reserved 13, 14, 15;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adadelta
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L933
-message AdadeltaParameters {
-  float rho = 1;
-  float epsilon = 2;
-
-  // Old initial accumulator parameters.
-  reserved "initial_accumulator", "initial_update";
-  reserved 3, 4;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/compat/v1/train/ProximalAdagradOptimizer
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L1961
-message ProximalAdagradParameters {
-  float l1 = 1;
-  float l2 = 2;
-
-  // Old initial accumulator parameter.
-  reserved "initial_accumulator";
-  reserved 3;
-}
-
-// The online Yogi optimizer does not implement hyper-parameter update; use the
-// dynamic learning rate feature instead, setting the learning rate to:
-// user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
-// Here, t is the current timestep.
-//
-// https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf
-// plus some extensions based on FTRL.
-//
-// Note that the code by default implements the lazy version of online Yogi.
-message OnlineYogiParameters {
-  // The L1 regularization parameter (used analogously to the one in FTRL).
-  float l1 = 1;
-
-  // The L2 regularization parameter (used analogously to the one in FTRL).
-  float l2 = 2;
-
-  // \beta_2 from Algorithm 2 in the paper.
-  float beta2 = 3;
-
-  // Reserved ids corresponding to removed tanh activation.
-  reserved 6;  // sign
-  reserved 7;  // tanh
-}
-
-// The online Yogi optimizer does not implement hyper-parameter update; use the
-// dynamic learning rate feature instead, setting the learning rate to:
-// user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
-// Here, t is the current timestep.
-//
-// https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf
-// plus some extensions based on FTRL.
-//
-// Note that the code by default implements the lazy version of proximal Yogi.
-message ProximalYogiParameters {
-  // The L1 regularization parameter.
-  float l1 = 1;
-
-  // The L2 regularization parameter.
-  float l2 = 2;
-
-  // The exponential decay rate for the 1st moment estimates.
-  float beta1 = 3;
-
-  // The exponential decay rate for the 2nd moment estimates.
-  float beta2 = 4;
-
-  // A constant trading off adaptivity and noise.
-  float epsilon = 5;
-
-  // Reserved ids corresponding to removed tanh activation.
-  reserved 8;  // sign
-  reserved 9;  // tanh
-}
-
-// Estimator for the frequency of updates to a lookup table. It maintains an
-// array (tf.Variable) D, where each element records the average number of
-// global steps between two consecutive batches that hit the corresponding
-// bucket. Once an item with bucket id i is sampled, D[i] is updated by:
-//   D[i] <- D[i] * (1 - tau) + delta[i] * tau,
-//
-// where tau is a learning rate between 0 and 1 (exclusive), and
-//   delta[i] = current global step - last step i is sampled.
-//
-// The estimated frequency (sampling rate in a batch) is thus 1 / D[i].
-//
-// Elements in D are initialized with a large value max_delta. delta[i] will
-// also be capped by this value.
-//
-// The exact sequence of operations used in the optimizer is shown below.
-// last_hit_step[i] is a tf.Variable that holds the last global step at which i
-// was sampled.
-//
-//   delta = global_step - last_hit_step[i]
-//   clipped_delta = min(delta, params.max_delta)
-//   is_outlier = (delta >= params.outlier_threshold * D[i])
-//   D[i] <- is_outlier ? clipped_delta
-//                      : D[i] * (1 - params.tau) + clipped_delta * params.tau
-//   last_hit_step[i] <- global_step
-message FrequencyEstimatorParameters {
-  // Learning rate between (0, 1) that is used to update the array D.
-  float tau = 1;
-
-  // Maximum value of delta: difference between the current global step and the
-  // last global step at which the row was sampled.
-  float max_delta = 2;
-
-  // Threshold used to determine whether the current update is an outlier.
-  float outlier_threshold = 3;
-
-  // The weight exponent used to transform the estimated delta into weights.
-  // The transformation function is: (delta / max_delta) ^ (weight_exponent)
-  float weight_exponent = 4;
-}
-
-// A user-defined optimizer.
-// The contained HLO program must take the following arguments in the following
-// order:
-// 1.  gradients
-// 2.  table weights
-// 3.  slot variables
-// 4.  an optional scalar input that is passed in via the dynamic learning
-//     rate mechanism.
-//
-// It must return/end in a tuple op that contains the following values in the
-// following order:
-// 1.  new table values
-// 2.  new slot variable value
-//
-// The program must have shape (1,1) with dtype float32 throughout and only use
-// HLO that operate elementwise (e.g., no reduce, no variables, no control flow
-// and no broadcasting outside of the single scalar input).
-// The HLO program should be written as if it were a dense update. It will be
-// called on each row that needs an update and will applied elementwise.
-message UserDefinedProgramParameters {
-  xla.HloModuleProto program = 1;
-  reserved 2;  // Was padding_values
-}
-
-// Optimizer that just sets the variable to the value of the gradient. To be
-// correct, this requires either gradient accumulation (to sum the values of a
-// computed expression across the samples) or to deduplicate IDs within a single
-// host (to assign the value from an arbitrary sample).
-message AssignParameters {}
-
-// Status of using gradient accumulation (doing two passes over the input
-// gradients: one to accumulate them into a temporary array and another to apply
-// them using the actual optimization algorithm). The extra message is to wrap
-// the enum for scoping.
-message GradientAccumulationStatus {
-  // if UNSPECIFIED (default), gradient accumulation is ENABLED.
-  enum Status {
-    UNSPECIFIED = 0;
-    ENABLED = 1;
-    DISABLED = 2;
-  }
-}
-
-// Whether to optimize the packing of low-dimensional embedding tables in HBM
-// (high bandwidth memory). TPUs access HBM at 32-byte (8-float) granularity.
-// For functional correctness, the TPU software internally pads the embedding
-// dimension to a multiple of 8. This can sometimes lead to significant memory
-// wastage due to padding. For 1-dimensional, 2-dimensional, and 4-dimensional,
-// the TPU software can remove this padding by packing multiple rows into the
-// same 8-float HBM chunk. For example, 8 rows could be packed into the same
-// 8-float chunk for a 1-dimensional embedding table.
-
-// There is one important limitation for this HBM packing though. When only a
-// subset of rows in an 8-float chunk are accessed on a particular step, the
-// adjoining rows in the same chunk are updated with zero gradients on the
-// backward pass even if they are not touched. This is an artifact of the
-// packing implementation. This operation is NOT functionally correct for
-// optimizers where zero gradients change the embeddings/slot-variable values,
-// e.g., momentum-based optimizers. Hence, this HBM packing cannot be enabled
-// for embedding tables with such optimizers. The TPU software automatically
-// recognizes that a zero gradient can modify state and turns off the low
-// dimensional embedding packing in that scenario.
-//
-// However, for optimizers where a zero gradient is a NoOp, such as SGD,
-// Adagrad, and FTRL, this packing optimization can be used. However, there are
-// some important considerations:
-// * Clipping limits: The initial values for such embeddings should fall within
-//   the clipping limits specified in the optimization parameters. Otherwise, a
-//   zero gradient will cause the embeddings to be clipped. This changes state
-//   and hence, is not a NoOp.
-// * FTRL: The embedding vector is computed directly from the values of the
-//   accumulator and linear slot variables. Hence, the initial embedding values
-//   should match that computed from the initial values of the accumulator and
-//   linear slot variables. Note that in nearly all cases, the linear value is
-//   initialized to zero; this corresponds to an embedding value of zero.
-//
-// Performance: The TPU has to perform additional work when low dimensional
-// packing is enabled. In certain situations when the vocabulary size is small,
-// it may not make sense to turn on this packing since the total memory usage
-// due to padding is extremely low. Hence, the TPU software automatically turns
-// off the packing optimization in such scenarios.
-message LowDimensionalPackingStatus {
-  // if UNSPECIFIED (default), the low dimension packing status is DISABLED.
-  // This can change in future.
-  //
-  // if ENABLED, the low dimension packing is enabled only if the following
-  // three additional conditions are true:
-  //  * The optimizer treats the zero gradient as a NoOp.
-  //  * The embedding dimension is 1, 2, or 4.
-  //  * The vocabulary size is large enough to avoid performance issues.
-  //
-  // if DISABLED, the low dimension packing is always disabled.
-  enum Status {
-    UNSPECIFIED = 0;
-    ENABLED = 1;
-    DISABLED = 2;
-  }
-}
-
-// Configuration proto for hot ID optimization. This is an experimental feature
-// that is currently disabled (by default).
-message HotIdReplicationConfiguration {
-  // Whether to enable or disable hot ID optimization.
-  // If set to UNSPECIFIED (default), hot ID optimization is DISABLED.
-  // If set to ENABLED, hot ID replication is turned ON.
-  // If set to MIGRATION_ONLY, hot ID migration is turned ON.
-  enum Status {
-    UNSPECIFIED = 0;
-    ENABLED = 1;
-    DISABLED = 2;
-    MIGRATION_ONLY = 3;
-  }
-  Status status = 1;
-}
-
-message OptimizationParameters {
-  // Learning rate used for updating the embedding layer parameters.
-  LearningRate learning_rate = 13;
-
-  // Limits to which to clip the weight values after the backward pass; not
-  // present means no limits are applied.
-  ClippingLimits clipping_limits = 2;
-
-  // Limits to which to clip the backward pass gradient before using it for
-  // updates; not present means no limits are applied.
-  ClippingLimits gradient_clipping_limits = 7;
-
-  // Amount of weight decay to apply; see weight_decay_optimizers.py for
-  // details. All optimizers except MDL Adagrad Light are supported with this
-  // option. Although there is no check, users who want weight decay will also
-  // want to ensure that gradient accumulation is enabled so that the decay will
-  // happen once per global batch.
-  float weight_decay_factor = 16;
-
-  // If true, the weight decay factor is multiplied by the current learning rate
-  // before use; this is to match the note in DecoupledWeightDecayExtension in
-  // weight_decay_optimizers.py.
-  bool multiply_weight_decay_factor_by_learning_rate = 22;
-
-  // Configuration for simulated quantization which is used to reduce
-  // training/serving skew when the serving variables are quantized. The same
-  // quantization operations are executed during training to minimize
-  // differences with serving.
-  SimulatedQuantization simulated_quantization = 27;
-
-  // Status of using gradient accumulation (doing two passes over the input
-  // gradients: one to accumulate them into a temporary array and another to
-  // apply them using the actual optimization algorithm).
-  GradientAccumulationStatus.Status gradient_accumulation_status = 17;
-
-  // Status of the low-dimensional embedding packing optimization. This controls
-  // whether to optimize the packing of 1-dimensional, 2-dimensional, and
-  // 4-dimensional embedding tables in memory.
-  LowDimensionalPackingStatus.Status low_dimensional_packing_status = 28;
-
-  // Configuration proto for hot ID replication. This is an experimental
-  // feature that is currently disabled (by default).
-  HotIdReplicationConfiguration hot_id_replication_configuration = 18;
-
-  // Optimization algorithm parameters; which field is selected determines which
-  // algorithm to use.
-  oneof parameters {
-    AdagradParameters adagrad = 3;
-    AdagradMomentumParameters adagrad_momentum = 26;
-    BoundedAdagradParameters bounded_adagrad = 19;
-    FrequencyAwareAdagradParameters frequency_aware_adagrad = 30;
-    StochasticGradientDescentParameters stochastic_gradient_descent = 4;
-    FtrlParameters ftrl = 5;
-    AdamParameters adam = 6;
-    MomentumParameters momentum = 8;
-    LionParameters lion = 29;
-    RmsPropParameters rms_prop = 9;
-    CenteredRmsPropParameters centered_rms_prop = 10;
-    MdlAdagradLightParameters mdl_adagrad_light = 11;
-    AdadeltaParameters adadelta = 12;
-    ProximalAdagradParameters proximal_adagrad = 14;
-    OnlineYogiParameters online_yogi = 20;
-    ProximalYogiParameters proximal_yogi = 21;
-    FrequencyEstimatorParameters frequency_estimator = 23;
-    UserDefinedProgramParameters user_defined_program = 24;
-    AssignParameters assign = 25;
-  }
-
-  reserved 1, 15;
-
-  // NEXT_ID: 31
-}
-
-// Specification of an optimization algorithm's state variables (both the main
-// value vector and any extra accumulators, etc.). This proto is only used
-// internally by the TPU software and is not exposed directly to the TF model.
-message StateVariableSpecification {
-  // Parameter name for the state variable.
-  string name = 1;
-
-  // A normal state variable that should be saved and restored in checkpoints
-  // and used as an input or output to non-debug TensorFlow ops.
-  message UserDefined {
-    reserved 1;  // Was padding_initial_value.
-  }
-
-  // A state variable that should be filled with a constant and normally hidden
-  // from users (used for intermediate gradients being accumulated, for
-  // example).
-  message FillWithConstant {
-    double initial_value = 1;
-  }
-
-  // Usage type of this state variable.
-  oneof usage {
-    UserDefined user_defined = 2;
-    FillWithConstant fill_with_constant = 3;
-  }
-}
+import public "xla/tpu/protobuf/optimization_parameters.proto";
diff --git a/tensorflow/core/protobuf/tpu/tpu_embedding_configuration.proto b/tensorflow/core/protobuf/tpu/tpu_embedding_configuration.proto
index c11ece0c6f4..ffc8d2e26ab 100644
--- a/tensorflow/core/protobuf/tpu/tpu_embedding_configuration.proto
+++ b/tensorflow/core/protobuf/tpu/tpu_embedding_configuration.proto
@@ -17,160 +17,4 @@ syntax = "proto3";
 
 package tensorflow.tpu;
 
-import "tensorflow/core/protobuf/tpu/optimization_parameters.proto";
-
-message TPUEmbeddingConfiguration {
-  // Description of the various embedding tables.
-  message TableDescriptor {
-    // Name of the table.
-    string name = 1;
-
-    // Size of the vocabulary (i.e., number of rows) in the table.
-    int64 vocabulary_size = 2;
-
-    // The embedding dimension (i.e., the width of the embedding table).
-    int32 dimension = 3;
-
-    // Number of features mapped to this table.
-    int32 num_features = 4;
-
-    // Details of the learning algorithm used to update the embedding
-    // parameters.
-    OptimizationParameters optimization_parameters = 5;
-  }
-  repeated TableDescriptor table_descriptor = 1;
-
-  // Mode. Should the embedding layer program be run for inference (just forward
-  // pass), training (both forward and backward pass) or just the backward_pass.
-  enum Mode {
-    UNSPECIFIED = 0;
-    INFERENCE = 1;
-    TRAINING = 2;
-    BACKWARD_PASS_ONLY = 3;
-  }
-  Mode mode = 2;
-
-  // Number of samples in each batch of embedding layer activations sent to
-  // the TensorCore.
-  int32 batch_size_per_tensor_core = 3;
-
-  // Number of TPU hosts used for inference/training.
-  int32 num_hosts = 4;
-
-  // Number of TensorCore used for inference/training.
-  int32 num_tensor_cores = 5;
-
-  // Sharding strategy of the embedding tables among the hosts.
-  // If the sharding_strategy is "mod", each id is assigned to host
-  // "id % num_hosts". For instance, 13 ids are split across 5 hosts as:
-  // [[0, 5, 10], [1, 6, 11], [2, 7, 12], [3, 8], [4, 9]].
-  // If the sharding_strategy is "div", ids are assigned to hosts in a
-  // contiguous manner. In this case, 13 ids are split across 5 hosts as:
-  // [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11, 12]].
-  // In both the strategies, if the id space does not evenly divide the number
-  // of hosts, each of the first "table_descriptor.vocabulary_size % num_hosts"
-  // hosts will be assigned one more id.
-  // This partitioning strategy exactly follows that in the embedding_lookup
-  // TensorFlow function at tensorflow/python/ops/embedding_ops.py.
-  enum ShardingStrategy {
-    DIV_DEFAULT = 0;
-    MOD = 1;
-  }
-  ShardingStrategy sharding_strategy = 6;
-
-  // This parameter determines if the execution of the sparse core will be
-  // pipelined with that of the TensorCore. This parameter only affects results
-  // when mode=TRAINING. If mode=INFERENCE or BACKWARD_PASS_ONLY, this parameter
-  // does not affect execution and hence, is a don't care value.
-  //
-  // false: The execution of the sparse core is not pipelined with that of the
-  // TensorCore. The forward pass of every step on the sparse core is executed
-  // only after the backward pass of the previous step is complete. And the
-  // backward pass on the sparse core is executed only after the embedding
-  // gradients have been computed on the TensorCore on every step. This ensures
-  // that the activations on every step observe the gradient updates from the
-  // previous step on both the sparse core and the TensorCore.
-  //
-  // true: The execution of the sparse core is pipelined with that of the
-  // TensorCore. The forward pass of every step on the sparse core can be
-  // executed after the forward pass of the previous step is complete without
-  // waiting for the backward pass. This improves the utilization of the sparse
-  // core allowing it to process step N+1 while the embedding gradients for step
-  // N are computed on the TensorCore. The backward pass of every step on the
-  // sparse core is executed directly after the forward pass for the next step
-  // is complete. The drawback is that embedding activations for step N+1 do not
-  // observe the embedding gradient updates from step N. This could affect model
-  // quality if step N and N+1 involve the same set of embedding IDs. However,
-  // since the embedding updates are sparse, this is generally not considered a
-  // problem.
-  bool pipeline_execution_with_tensor_core = 7;
-
-  // Directory where embedding lookup statistics are stored. These statistics
-  // summarize information about the inputs to the embedding lookup
-  // operation, in particular, the average number of embedding IDs per example
-  // and how well the embedding IDs are load balanced across the system. The
-  // lookup statistics are used during TPU initialization for embedding table
-  // partitioning. Collection of lookup statistics is done at runtime by
-  // profiling the embedding inputs: only 3% of input samples are profiled to
-  // minimize host CPU overhead. Once a suitable number of samples are
-  // profiled, the lookup statistics are saved to table-specific files in the
-  // profile data directory generally at the end of a TPU training loop. The
-  // filename corresponding to each table is obtained by hashing table specific
-  // parameters (e.g., table name and number of features) and global
-  // configuration parameters (e.g., sharding strategy and TPU worker task
-  // count). The same profile data directory can be shared amongst several
-  // models to reuse embedding lookup statistics.
-  string profile_data_directory = 9;
-
-  // Description of different input features.
-  message FeatureDescriptor {
-    // Name of the input feature.
-    string name = 1;
-
-    // Index of the corresponding table in the TableDescriptor list.
-    int32 table_id = 2;
-
-    // Static shape of the inputs (excluding the reduction axis). Note that
-    // the shape of the actual inputs provided using the infeed op must be
-    // strictly smaller than input_shape. The outputs received at the TensorCore
-    // will have rank = input_shape.size() + 1. The innermost axis corresponds
-    // to the embedding dimension. If the input has shape [m, n, k] (excluding
-    // the reduction axis) and the embedding dimension is d, the output received
-    // at the TensorCore will have shape [m, n, k, d].
-    repeated int32 input_shape = 3;
-  }
-  // If the feature_descriptor field is populated, the model should NOT populate
-  // TableDescriptor.num_features and batch_size_per_tensor_core. These two
-  // fields will be auto-populated by the TPUEmbedding rewrite passes.
-  repeated FeatureDescriptor feature_descriptor = 10;
-
-  // SPMD (Single Program Multiple Data) sharding configuration for
-  // TPUEmbedding. When model parallelism is used on the TensorCore, the number
-  // of cores per replica must be passed to TPUEmbedding so that the right
-  // shapes can be computed in the TF/XLA bridge.
-  message SpmdSharding {
-    // Whether SPMD sharding is enabled.
-    bool enabled = 1;
-
-    // Number of cores per replica.
-    int32 num_cores_per_replica = 2;
-
-    // If true, the tensors are manually partitioned. Otherwise, use the
-    // automatic SPMD partitioning. This should be true when users use
-    // `shard_map`.
-    bool use_manual_partitioning = 3;
-
-    // If true, lowers to Shardy custom-call for manual partitioning.
-    // Default is false.
-    bool use_shardy_partitioner = 4;
-  }
-  SpmdSharding spmd_sharding = 11;
-
-  // Old TPU embedding output layout.
-  reserved "output_layout";
-  reserved 8;
-}
-
-// A placeholder message that is used to define a unique Status payload
-// URL for TPU embedding errors.
-message TPUEmbeddingError {}
+import public "xla/tpu/protobuf/tpu_embedding_configuration.proto";
diff --git a/third_party/xla/xla/tpu/protobuf/BUILD b/third_party/xla/xla/tpu/protobuf/BUILD
new file mode 100644
index 00000000000..f732b3e19e7
--- /dev/null
+++ b/third_party/xla/xla/tpu/protobuf/BUILD
@@ -0,0 +1,60 @@
+# Copyright 2026 The OpenXLA 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.
+# ==============================================================================
+
+# copybara:uncomment load("@rules_python//python:proto.bzl", "py_proto_library")
+load(
+    "//xla/tsl/platform:build_config.bzl",
+    "tf_proto_library",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+tf_proto_library(
+    name = "tpu_embedding_configuration_proto",
+    srcs = [
+        "tpu_embedding_configuration.proto",
+    ],
+    protodeps = [
+        ":optimization_parameters_proto",
+    ],
+    visibility = ["//visibility:public"],
+)
+
+tf_proto_library(
+    name = "optimization_parameters_proto",
+    srcs = [
+        "optimization_parameters.proto",
+    ],
+    protodeps = ["//xla/service:hlo_proto"],
+    visibility = ["//visibility:public"],
+    deps = ["@com_google_protobuf//:wrappers_proto"],
+)
+
+# copybara:uncomment_begin(google-only)
+# py_proto_library(
+#     name = "tpu_embedding_configuration_py_pb2",
+#     visibility = ["//visibility:public"],
+#     deps = [":tpu_embedding_configuration_proto"],
+# )
+#
+# py_proto_library(
+#     name = "optimization_parameters_py_pb2",
+#     visibility = ["//visibility:public"],
+#     deps = [":optimization_parameters_proto"],
+# )
+# copybara:uncomment_end
diff --git a/third_party/xla/xla/tpu/protobuf/optimization_parameters.proto b/third_party/xla/xla/tpu/protobuf/optimization_parameters.proto
new file mode 100644
index 00000000000..424110f3d89
--- /dev/null
+++ b/third_party/xla/xla/tpu/protobuf/optimization_parameters.proto
@@ -0,0 +1,657 @@
+// Copyright 2026 The OpenXLA 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.
+// ==============================================================================
+
+syntax = "proto3";
+
+package tensorflow.tpu;
+
+import "google/protobuf/wrappers.proto";
+import "xla/service/hlo.proto";
+
+message ClippingLimits {
+  google.protobuf.FloatValue lower = 1;  // -inf if not set
+  google.protobuf.FloatValue upper = 2;  // +inf if not set
+}
+
+// Configuration for simulated quantization; simulated quantization is used to
+// reduce training/serving skew when the serving variables are quantized. The
+// same quantization operations are executed during training to minimize
+// differences with serving.
+//
+// Simulated quantization inserts the following operations on the forward pass
+// after gathering the embedding vector from HBM. The backward pass operations
+// are unchanged.
+//
+// clipped_val = clip(input, clipping_limits)
+// quantum = clipping_limits.range() / (num_buckets - 1)
+// quantized_val = floor((clipped_val - clipping_limits.lower()) / quantum + .5)
+// return quantized_val * quantum + clipping_limits.lower().
+message SimulatedQuantization {
+  // Whether simulated quantization is enabled.
+  bool enabled = 1;
+
+  // Minimum and maximum values of the range used for quantization.
+  ClippingLimits clipping_limits = 2;
+
+  // Number of possible quantized values.
+  int32 num_buckets = 3;
+}
+
+// Dynamic input specification for optimizers in the TPUEmbeddingConfiguration.
+// The actual dynamic inputs are provided as a scalar input list to the
+// SendTPUEmbeddingGradients Op indexed by their tag specified through the
+// following proto.
+message OptimizerDynamicInput {
+  // For tables where dynamic inputs are needed (e.g., learning rates or other
+  // dynamic hyperparameters used in optimizers), a tag must be specified for
+  // the input.
+  //
+  // The tag must be a non-negative  integer. All tags in the range
+  // [0, number_of_unique_tags) must be present in the TPU embedding
+  // configuration, i.e. a tag cannot be skipped if a different tag numerically
+  // greater than it is used in the configuration.
+  //
+  // If multiple tables specify the same tag, they *MUST* have
+  // the same dynamic input, for example, their dynamic learning rate could be
+  // computed by the same TensorFlow sub-graph. The partitioning of the
+  // embedding layer would be more optimal if the number_of_unique_tags is as
+  // *LOW* as possible, i.e., if many tables share the same tag.
+  //
+  // The hyper_parameters input of the SendTPUEmbeddingGradients op is used to
+  // communicate dynamic hyper-parameters to the TPU embedding program.
+  // The hyper_parameters input is a list of scalars where the size of the list
+  // is equal to the number of unique tags. The hyper-parameter associated with
+  // a particular tag is specified by populating its corresponding index in the
+  // list of scalars.
+  int32 tag = 1;
+}
+
+// Source of learning rate to use.
+message LearningRate {
+  oneof learning_rate {
+    float constant = 1;
+    OptimizerDynamicInput dynamic = 2;
+  }
+}
+
+// Each optimizer's parameter proto has a link to its documentation and CPU
+// implementation (if available) for user reference.
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adagrad
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L1634
+message AdagradParameters {
+  // Old initial accumulator parameter.
+  reserved "initial_accumulator";
+  reserved 1;
+}
+
+// This optimizer combines the Adagrad and Momentum update rules.
+// accum(new) = beta2 == 1.0 ?
+//              accum(old) + grad^2 :
+//              beta2 * accum(old) + (1 - beta2) * grad^2
+// accum_with_exponent = (accum(new) + epsilon)^(-1.0 / exponent)
+// mom_accum(new) = momentum * mom_accum(old) + accum_with_exponent
+// update = use_nesterov ?
+//          momentum * mom_accum(new) + accum_with_exponent :
+//          mom_accum(new)
+// var(new) = var(old) - lr * grad * update
+// Algorithm described in https://arxiv.org/abs/2002.11803.
+message AdagradMomentumParameters {
+  // Moving average parameter for the momentum accumulator.
+  float momentum = 1;
+
+  // Whether to use the Nesterov variant of momentum.
+  bool use_nesterov = 2;
+
+  // Exponent for the gradient^2 accumulator.
+  float exponent = 3;
+
+  // Moving average parameter for the gradient^2 accumulator.
+  float beta2 = 4;
+
+  // Offset added to the Adagrad accumulator.
+  float epsilon = 5;
+}
+
+// Algorithm in http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf.
+message BoundedAdagradParameters {
+  // Whether to use the updated or the old value of the accumulator when
+  // computing the effective learning rate. When update_accumulator_first is set
+  // to True, the updated value of the accumulator is used.
+  bool update_accumulator_first = 1;
+
+  // The max_var_update value to use. Set value to 0 (default) to disable using
+  // max_var_update to clip the gradient.
+  float max_var_update = 2;
+
+  // The maximum value of the accumulator. Set max_accumulator to 0 (default)
+  // to disable using max_accumulator to clip the accumulator.
+  float max_accumulator = 3;
+}
+
+// Frequency Aware Adagrad optimizer. This optimizer implements the AdaGrad
+// algorithm and further allows to:
+// * Scale the learning rate based on frequency of the update. Sparsely updated
+//   rows are updated with a higher effective learning rate, and frequently
+//   updated rows are updated with a lower effective learning rate.
+// * Decay the growth of the accumulator values.
+// * Use L1 / L2 regularization for the weight updates.
+//
+// The optimization algorithm is shown below.
+// counter(new) = counter(old) + 1
+// accum(new) = max(accumulator_decay * accum(old) + grad^2,
+//                  initial_accumulator_value)
+// lr_scale = min((step_counter / accum(new)) ^ probability_exponent,
+// max_lr_multiplier) update = grad * lr_scale / sqrt(accum(new)) if
+// (l1_regularization_strength > 0.0):
+//   update = update + l1_regularization_strength * sign(var(old))
+// if (l2_regularization_strength > 0.0):
+//   update = update + l2_regularization_strength * var(old)
+// var(new) = var(old) - lr_scale * grad * update
+
+message FrequencyAwareAdagradParameters {
+  // The L1 regularization parameter for adjusting the update based on the sign
+  // of the variable.
+  float l1_regularization_strength = 1;
+
+  // The L2 regularization parameter for adjusting the update based on the
+  // variable.
+  float l2_regularization_strength = 2;
+
+  // The exponent used for scaling the learning rate based on the sparsity of
+  // updates.
+  float probability_exponent = 4;
+
+  // The maximum value of the learning rate scale.
+  float max_lr_multiplier = 3;
+
+  // The decay for the Adagrad accumulator.
+  float accumulator_decay = 5;
+
+  // The initial and minimum value for the Adagrad accumulator.
+  float initial_accumulator_value = 6;
+
+  // The tag for identifying the step counter used for the frequency aware
+  // Adagrad optimizer.
+  OptimizerDynamicInput step_counter = 7;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L629
+message StochasticGradientDescentParameters {}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Ftrl
+// https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/41159.pdf
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L2646
+//
+// The hyperparameters for FTRL are the same as for the Keras implementation,
+// with some additions. The "beta" parameter matches the behavior described in
+// the second link above; "beta" / (2 * learning rate) should be added to "l2"
+// to get equivalent behavior in the other TensorFlow implementations of this
+// optimizer. When the multiply_linear_by_lr field is set to true, a modified
+// formula is used for FTRL that treats the "linear" accumulator as being
+// pre-multiplied by the learning rate (i.e., the accumulator named "linear"
+// actually stores "linear * learning_rate"). Other than checkpoint
+// compatibility, this is mathematically equivalent for a static learning rate;
+// for a dynamic learning rate, it is nearly the same as long as the learning
+// rate does not change quickly. The benefit of setting multiply_linear_by_lr to
+// true is that the modified formula handles zero and near-zero learning rates
+// without producing NaNs, improving flexibility for learning rate ramp-up.
+message FtrlParameters {
+  float l1 = 1;
+  float l2 = 2;
+  float lr_power = 3;
+  float beta = 7;
+  bool multiply_linear_by_lr = 6;
+
+  // Previously, allow_zero_accumulator parameter changed some internal formulas
+  // to allow zero and near-zero accumulator values at the cost of some
+  // performance. The current implementation ignores this parameter; zero or
+  // near-zero accumulator values are now always supported.
+  bool allow_zero_accumulator = 8 [deprecated = true];
+
+  // Old initial accumulator parameters.
+  reserved "initial_accum", "initial_linear";
+  reserved 4, 5;
+}
+
+// The Adam optimizer does not implement hyper-parameter update due to hardware
+// limitations; use the dynamic learning rate feature instead, setting the
+// learning rate to: user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
+// Here, t is the current timestep.
+//
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adam
+// https://github.com/tensorflow/tensorflow/blob/ab51450c817674c8ff08a7ae4f8ac50cdc4bed8b/tensorflow/python/training/adam.py#L32
+//
+// Note that the code by default implements the lazy version of Adam
+// (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/LazyAdamOptimizer)
+// unless the use_non_lazy_adam parameter is set, in which case it implements
+// the normal version of Adam that updates all parameters in the embedding
+// table, even for entries that are not used in the current minibatch
+// (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/AdamOptimizer). If
+// use_non_lazy_adam is enabled, gradient accumulation is also required to be
+// enabled in order to get correct results; a warning will be printed otherwise
+// (which may change to an error in the future). If use_sum_inside_sqrt is set,
+// the Adam variable update formula will be changed from m / (sqrt(v) + epsilon)
+// to m / sqrt(v + epsilon**2); this option improves the performance of TPU
+// training and is not expected to harm model quality.
+message AdamParameters {
+  float beta1 = 3;
+  float beta2 = 4;
+  float epsilon = 5;
+  bool use_non_lazy_adam = 8;
+  bool use_sum_inside_sqrt = 10;
+
+  // Old initial accumulator parameters.
+  reserved "initial_m", "initial_v";
+  reserved 6, 7;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L3068
+message MomentumParameters {
+  float momentum = 1;
+  bool use_nesterov = 2;
+
+  // Old initial accumulator parameter.
+  reserved "initial_accum";
+  reserved 3;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Lion
+// momenta(new) = beta2 * momenta(old) + (1 - beta2) * grad
+// momenta_t = beta1 * momenta(old) + (1 - beta1) * grad
+// var(new) = var(old) - lr * sign(momenta_t)
+// Algorithm described in https://arxiv.org/abs/2302.06675.
+message LionParameters {
+  float beta1 = 1;
+  float beta2 = 2;
+  bool use_non_lazy_lion = 3;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L4229
+message RmsPropParameters {
+  float rho = 1;
+  float momentum = 2;
+  float epsilon = 3;
+
+  // Old initial accumulator parameters.
+  reserved "initial_ms", "initial_mom";
+  reserved 4, 5;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L4358
+message CenteredRmsPropParameters {
+  float rho = 1;
+  float momentum = 2;
+  float epsilon = 3;
+
+  // Old initial accumulator parameters.
+  reserved "initial_ms", "initial_mom", "initial_mg";
+  reserved 4, 5, 6;
+}
+
+// Variant of algorithm in http://proceedings.mlr.press/v44/shamir15.pdf
+message MdlAdagradLightParameters {
+  float l2 = 1;
+  float lr_power = 2;
+  float min_servable_mdl_benefit = 3;
+  float mdl_mix_in_margin = 4;
+  float mdl_benefit_rampup_coeff = 5;
+  float mdl_min_weight = 6;
+  float benefit_revisit_scale = 7;
+  float max_event_benefit = 8;
+  float max_total_benefit = 9;
+  float mdl_hard_limit = 10;
+  bool hard_limit_min_benefit = 11;
+  bool mdl_regularize = 12;
+
+  // Old initial accumulator parameters.
+  reserved "initial_accumulator", "initial_weight", "initial_benefit";
+  reserved 13, 14, 15;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adadelta
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L933
+message AdadeltaParameters {
+  float rho = 1;
+  float epsilon = 2;
+
+  // Old initial accumulator parameters.
+  reserved "initial_accumulator", "initial_update";
+  reserved 3, 4;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/compat/v1/train/ProximalAdagradOptimizer
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L1961
+message ProximalAdagradParameters {
+  float l1 = 1;
+  float l2 = 2;
+
+  // Old initial accumulator parameter.
+  reserved "initial_accumulator";
+  reserved 3;
+}
+
+// The online Yogi optimizer does not implement hyper-parameter update; use the
+// dynamic learning rate feature instead, setting the learning rate to:
+// user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
+// Here, t is the current timestep.
+//
+// https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf
+// plus some extensions based on FTRL.
+//
+// Note that the code by default implements the lazy version of online Yogi.
+message OnlineYogiParameters {
+  // The L1 regularization parameter (used analogously to the one in FTRL).
+  float l1 = 1;
+
+  // The L2 regularization parameter (used analogously to the one in FTRL).
+  float l2 = 2;
+
+  // \beta_2 from Algorithm 2 in the paper.
+  float beta2 = 3;
+
+  // Reserved ids corresponding to removed tanh activation.
+  reserved 6;  // sign
+  reserved 7;  // tanh
+}
+
+// The online Yogi optimizer does not implement hyper-parameter update; use the
+// dynamic learning rate feature instead, setting the learning rate to:
+// user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
+// Here, t is the current timestep.
+//
+// https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf
+// plus some extensions based on FTRL.
+//
+// Note that the code by default implements the lazy version of proximal Yogi.
+message ProximalYogiParameters {
+  // The L1 regularization parameter.
+  float l1 = 1;
+
+  // The L2 regularization parameter.
+  float l2 = 2;
+
+  // The exponential decay rate for the 1st moment estimates.
+  float beta1 = 3;
+
+  // The exponential decay rate for the 2nd moment estimates.
+  float beta2 = 4;
+
+  // A constant trading off adaptivity and noise.
+  float epsilon = 5;
+
+  // Reserved ids corresponding to removed tanh activation.
+  reserved 8;  // sign
+  reserved 9;  // tanh
+}
+
+// Estimator for the frequency of updates to a lookup table. It maintains an
+// array (tf.Variable) D, where each element records the average number of
+// global steps between two consecutive batches that hit the corresponding
+// bucket. Once an item with bucket id i is sampled, D[i] is updated by:
+//   D[i] <- D[i] * (1 - tau) + delta[i] * tau,
+//
+// where tau is a learning rate between 0 and 1 (exclusive), and
+//   delta[i] = current global step - last step i is sampled.
+//
+// The estimated frequency (sampling rate in a batch) is thus 1 / D[i].
+//
+// Elements in D are initialized with a large value max_delta. delta[i] will
+// also be capped by this value.
+//
+// The exact sequence of operations used in the optimizer is shown below.
+// last_hit_step[i] is a tf.Variable that holds the last global step at which i
+// was sampled.
+//
+//   delta = global_step - last_hit_step[i]
+//   clipped_delta = min(delta, params.max_delta)
+//   is_outlier = (delta >= params.outlier_threshold * D[i])
+//   D[i] <- is_outlier ? clipped_delta
+//                      : D[i] * (1 - params.tau) + clipped_delta * params.tau
+//   last_hit_step[i] <- global_step
+message FrequencyEstimatorParameters {
+  // Learning rate between (0, 1) that is used to update the array D.
+  float tau = 1;
+
+  // Maximum value of delta: difference between the current global step and the
+  // last global step at which the row was sampled.
+  float max_delta = 2;
+
+  // Threshold used to determine whether the current update is an outlier.
+  float outlier_threshold = 3;
+
+  // The weight exponent used to transform the estimated delta into weights.
+  // The transformation function is: (delta / max_delta) ^ (weight_exponent)
+  float weight_exponent = 4;
+}
+
+// A user-defined optimizer.
+// The contained HLO program must take the following arguments in the following
+// order:
+// 1.  gradients
+// 2.  table weights
+// 3.  slot variables
+// 4.  an optional scalar input that is passed in via the dynamic learning
+//     rate mechanism.
+//
+// It must return/end in a tuple op that contains the following values in the
+// following order:
+// 1.  new table values
+// 2.  new slot variable value
+//
+// The program must have shape (1,1) with dtype float32 throughout and only use
+// HLO that operate elementwise (e.g., no reduce, no variables, no control flow
+// and no broadcasting outside of the single scalar input).
+// The HLO program should be written as if it were a dense update. It will be
+// called on each row that needs an update and will applied elementwise.
+message UserDefinedProgramParameters {
+  xla.HloModuleProto program = 1;
+  reserved 2;  // Was padding_values
+}
+
+// Optimizer that just sets the variable to the value of the gradient. To be
+// correct, this requires either gradient accumulation (to sum the values of a
+// computed expression across the samples) or to deduplicate IDs within a single
+// host (to assign the value from an arbitrary sample).
+message AssignParameters {}
+
+// Status of using gradient accumulation (doing two passes over the input
+// gradients: one to accumulate them into a temporary array and another to apply
+// them using the actual optimization algorithm). The extra message is to wrap
+// the enum for scoping.
+message GradientAccumulationStatus {
+  // if UNSPECIFIED (default), gradient accumulation is ENABLED.
+  enum Status {
+    UNSPECIFIED = 0;
+    ENABLED = 1;
+    DISABLED = 2;
+  }
+}
+
+// Whether to optimize the packing of low-dimensional embedding tables in HBM
+// (high bandwidth memory). TPUs access HBM at 32-byte (8-float) granularity.
+// For functional correctness, the TPU software internally pads the embedding
+// dimension to a multiple of 8. This can sometimes lead to significant memory
+// wastage due to padding. For 1-dimensional, 2-dimensional, and 4-dimensional,
+// the TPU software can remove this padding by packing multiple rows into the
+// same 8-float HBM chunk. For example, 8 rows could be packed into the same
+// 8-float chunk for a 1-dimensional embedding table.
+
+// There is one important limitation for this HBM packing though. When only a
+// subset of rows in an 8-float chunk are accessed on a particular step, the
+// adjoining rows in the same chunk are updated with zero gradients on the
+// backward pass even if they are not touched. This is an artifact of the
+// packing implementation. This operation is NOT functionally correct for
+// optimizers where zero gradients change the embeddings/slot-variable values,
+// e.g., momentum-based optimizers. Hence, this HBM packing cannot be enabled
+// for embedding tables with such optimizers. The TPU software automatically
+// recognizes that a zero gradient can modify state and turns off the low
+// dimensional embedding packing in that scenario.
+//
+// However, for optimizers where a zero gradient is a NoOp, such as SGD,
+// Adagrad, and FTRL, this packing optimization can be used. However, there are
+// some important considerations:
+// * Clipping limits: The initial values for such embeddings should fall within
+//   the clipping limits specified in the optimization parameters. Otherwise, a
+//   zero gradient will cause the embeddings to be clipped. This changes state
+//   and hence, is not a NoOp.
+// * FTRL: The embedding vector is computed directly from the values of the
+//   accumulator and linear slot variables. Hence, the initial embedding values
+//   should match that computed from the initial values of the accumulator and
+//   linear slot variables. Note that in nearly all cases, the linear value is
+//   initialized to zero; this corresponds to an embedding value of zero.
+//
+// Performance: The TPU has to perform additional work when low dimensional
+// packing is enabled. In certain situations when the vocabulary size is small,
+// it may not make sense to turn on this packing since the total memory usage
+// due to padding is extremely low. Hence, the TPU software automatically turns
+// off the packing optimization in such scenarios.
+message LowDimensionalPackingStatus {
+  // if UNSPECIFIED (default), the low dimension packing status is DISABLED.
+  // This can change in future.
+  //
+  // if ENABLED, the low dimension packing is enabled only if the following
+  // three additional conditions are true:
+  //  * The optimizer treats the zero gradient as a NoOp.
+  //  * The embedding dimension is 1, 2, or 4.
+  //  * The vocabulary size is large enough to avoid performance issues.
+  //
+  // if DISABLED, the low dimension packing is always disabled.
+  enum Status {
+    UNSPECIFIED = 0;
+    ENABLED = 1;
+    DISABLED = 2;
+  }
+}
+
+// Configuration proto for hot ID optimization. This is an experimental feature
+// that is currently disabled (by default).
+message HotIdReplicationConfiguration {
+  // Whether to enable or disable hot ID optimization.
+  // If set to UNSPECIFIED (default), hot ID optimization is DISABLED.
+  // If set to ENABLED, hot ID replication is turned ON.
+  // If set to MIGRATION_ONLY, hot ID migration is turned ON.
+  enum Status {
+    UNSPECIFIED = 0;
+    ENABLED = 1;
+    DISABLED = 2;
+    MIGRATION_ONLY = 3;
+  }
+  Status status = 1;
+}
+
+message OptimizationParameters {
+  // Learning rate used for updating the embedding layer parameters.
+  LearningRate learning_rate = 13;
+
+  // Limits to which to clip the weight values after the backward pass; not
+  // present means no limits are applied.
+  ClippingLimits clipping_limits = 2;
+
+  // Limits to which to clip the backward pass gradient before using it for
+  // updates; not present means no limits are applied.
+  ClippingLimits gradient_clipping_limits = 7;
+
+  // Amount of weight decay to apply; see weight_decay_optimizers.py for
+  // details. All optimizers except MDL Adagrad Light are supported with this
+  // option. Although there is no check, users who want weight decay will also
+  // want to ensure that gradient accumulation is enabled so that the decay will
+  // happen once per global batch.
+  float weight_decay_factor = 16;
+
+  // If true, the weight decay factor is multiplied by the current learning rate
+  // before use; this is to match the note in DecoupledWeightDecayExtension in
+  // weight_decay_optimizers.py.
+  bool multiply_weight_decay_factor_by_learning_rate = 22;
+
+  // Configuration for simulated quantization which is used to reduce
+  // training/serving skew when the serving variables are quantized. The same
+  // quantization operations are executed during training to minimize
+  // differences with serving.
+  SimulatedQuantization simulated_quantization = 27;
+
+  // Status of using gradient accumulation (doing two passes over the input
+  // gradients: one to accumulate them into a temporary array and another to
+  // apply them using the actual optimization algorithm).
+  GradientAccumulationStatus.Status gradient_accumulation_status = 17;
+
+  // Status of the low-dimensional embedding packing optimization. This controls
+  // whether to optimize the packing of 1-dimensional, 2-dimensional, and
+  // 4-dimensional embedding tables in memory.
+  LowDimensionalPackingStatus.Status low_dimensional_packing_status = 28;
+
+  // Configuration proto for hot ID replication. This is an experimental
+  // feature that is currently disabled (by default).
+  HotIdReplicationConfiguration hot_id_replication_configuration = 18;
+
+  // Optimization algorithm parameters; which field is selected determines which
+  // algorithm to use.
+  oneof parameters {
+    AdagradParameters adagrad = 3;
+    AdagradMomentumParameters adagrad_momentum = 26;
+    BoundedAdagradParameters bounded_adagrad = 19;
+    FrequencyAwareAdagradParameters frequency_aware_adagrad = 30;
+    StochasticGradientDescentParameters stochastic_gradient_descent = 4;
+    FtrlParameters ftrl = 5;
+    AdamParameters adam = 6;
+    MomentumParameters momentum = 8;
+    LionParameters lion = 29;
+    RmsPropParameters rms_prop = 9;
+    CenteredRmsPropParameters centered_rms_prop = 10;
+    MdlAdagradLightParameters mdl_adagrad_light = 11;
+    AdadeltaParameters adadelta = 12;
+    ProximalAdagradParameters proximal_adagrad = 14;
+    OnlineYogiParameters online_yogi = 20;
+    ProximalYogiParameters proximal_yogi = 21;
+    FrequencyEstimatorParameters frequency_estimator = 23;
+    UserDefinedProgramParameters user_defined_program = 24;
+    AssignParameters assign = 25;
+  }
+
+  reserved 1, 15;
+
+  // NEXT_ID: 31
+}
+
+// Specification of an optimization algorithm's state variables (both the main
+// value vector and any extra accumulators, etc.). This proto is only used
+// internally by the TPU software and is not exposed directly to the TF model.
+message StateVariableSpecification {
+  // Parameter name for the state variable.
+  string name = 1;
+
+  // A normal state variable that should be saved and restored in checkpoints
+  // and used as an input or output to non-debug TensorFlow ops.
+  message UserDefined {
+    reserved 1;  // Was padding_initial_value.
+  }
+
+  // A state variable that should be filled with a constant and normally hidden
+  // from users (used for intermediate gradients being accumulated, for
+  // example).
+  message FillWithConstant {
+    double initial_value = 1;
+  }
+
+  // Usage type of this state variable.
+  oneof usage {
+    UserDefined user_defined = 2;
+    FillWithConstant fill_with_constant = 3;
+  }
+}
diff --git a/third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration.proto b/third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration.proto
new file mode 100644
index 00000000000..14f64bcc391
--- /dev/null
+++ b/third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration.proto
@@ -0,0 +1,176 @@
+// 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.
+// ==============================================================================
+
+syntax = "proto3";
+
+package tensorflow.tpu;
+
+import "xla/tpu/protobuf/optimization_parameters.proto";
+
+message TPUEmbeddingConfiguration {
+  // Description of the various embedding tables.
+  message TableDescriptor {
+    // Name of the table.
+    string name = 1;
+
+    // Size of the vocabulary (i.e., number of rows) in the table.
+    int64 vocabulary_size = 2;
+
+    // The embedding dimension (i.e., the width of the embedding table).
+    int32 dimension = 3;
+
+    // Number of features mapped to this table.
+    int32 num_features = 4;
+
+    // Details of the learning algorithm used to update the embedding
+    // parameters.
+    OptimizationParameters optimization_parameters = 5;
+  }
+  repeated TableDescriptor table_descriptor = 1;
+
+  // Mode. Should the embedding layer program be run for inference (just forward
+  // pass), training (both forward and backward pass) or just the backward_pass.
+  enum Mode {
+    UNSPECIFIED = 0;
+    INFERENCE = 1;
+    TRAINING = 2;
+    BACKWARD_PASS_ONLY = 3;
+  }
+  Mode mode = 2;
+
+  // Number of samples in each batch of embedding layer activations sent to
+  // the TensorCore.
+  int32 batch_size_per_tensor_core = 3;
+
+  // Number of TPU hosts used for inference/training.
+  int32 num_hosts = 4;
+
+  // Number of TensorCore used for inference/training.
+  int32 num_tensor_cores = 5;
+
+  // Sharding strategy of the embedding tables among the hosts.
+  // If the sharding_strategy is "mod", each id is assigned to host
+  // "id % num_hosts". For instance, 13 ids are split across 5 hosts as:
+  // [[0, 5, 10], [1, 6, 11], [2, 7, 12], [3, 8], [4, 9]].
+  // If the sharding_strategy is "div", ids are assigned to hosts in a
+  // contiguous manner. In this case, 13 ids are split across 5 hosts as:
+  // [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11, 12]].
+  // In both the strategies, if the id space does not evenly divide the number
+  // of hosts, each of the first "table_descriptor.vocabulary_size % num_hosts"
+  // hosts will be assigned one more id.
+  // This partitioning strategy exactly follows that in the embedding_lookup
+  // TensorFlow function at tensorflow/python/ops/embedding_ops.py.
+  enum ShardingStrategy {
+    DIV_DEFAULT = 0;
+    MOD = 1;
+  }
+  ShardingStrategy sharding_strategy = 6;
+
+  // This parameter determines if the execution of the sparse core will be
+  // pipelined with that of the TensorCore. This parameter only affects results
+  // when mode=TRAINING. If mode=INFERENCE or BACKWARD_PASS_ONLY, this parameter
+  // does not affect execution and hence, is a don't care value.
+  //
+  // false: The execution of the sparse core is not pipelined with that of the
+  // TensorCore. The forward pass of every step on the sparse core is executed
+  // only after the backward pass of the previous step is complete. And the
+  // backward pass on the sparse core is executed only after the embedding
+  // gradients have been computed on the TensorCore on every step. This ensures
+  // that the activations on every step observe the gradient updates from the
+  // previous step on both the sparse core and the TensorCore.
+  //
+  // true: The execution of the sparse core is pipelined with that of the
+  // TensorCore. The forward pass of every step on the sparse core can be
+  // executed after the forward pass of the previous step is complete without
+  // waiting for the backward pass. This improves the utilization of the sparse
+  // core allowing it to process step N+1 while the embedding gradients for step
+  // N are computed on the TensorCore. The backward pass of every step on the
+  // sparse core is executed directly after the forward pass for the next step
+  // is complete. The drawback is that embedding activations for step N+1 do not
+  // observe the embedding gradient updates from step N. This could affect model
+  // quality if step N and N+1 involve the same set of embedding IDs. However,
+  // since the embedding updates are sparse, this is generally not considered a
+  // problem.
+  bool pipeline_execution_with_tensor_core = 7;
+
+  // Directory where embedding lookup statistics are stored. These statistics
+  // summarize information about the inputs to the embedding lookup
+  // operation, in particular, the average number of embedding IDs per example
+  // and how well the embedding IDs are load balanced across the system. The
+  // lookup statistics are used during TPU initialization for embedding table
+  // partitioning. Collection of lookup statistics is done at runtime by
+  // profiling the embedding inputs: only 3% of input samples are profiled to
+  // minimize host CPU overhead. Once a suitable number of samples are
+  // profiled, the lookup statistics are saved to table-specific files in the
+  // profile data directory generally at the end of a TPU training loop. The
+  // filename corresponding to each table is obtained by hashing table specific
+  // parameters (e.g., table name and number of features) and global
+  // configuration parameters (e.g., sharding strategy and TPU worker task
+  // count). The same profile data directory can be shared amongst several
+  // models to reuse embedding lookup statistics.
+  string profile_data_directory = 9;
+
+  // Description of different input features.
+  message FeatureDescriptor {
+    // Name of the input feature.
+    string name = 1;
+
+    // Index of the corresponding table in the TableDescriptor list.
+    int32 table_id = 2;
+
+    // Static shape of the inputs (excluding the reduction axis). Note that
+    // the shape of the actual inputs provided using the infeed op must be
+    // strictly smaller than input_shape. The outputs received at the TensorCore
+    // will have rank = input_shape.size() + 1. The innermost axis corresponds
+    // to the embedding dimension. If the input has shape [m, n, k] (excluding
+    // the reduction axis) and the embedding dimension is d, the output received
+    // at the TensorCore will have shape [m, n, k, d].
+    repeated int32 input_shape = 3;
+  }
+  // If the feature_descriptor field is populated, the model should NOT populate
+  // TableDescriptor.num_features and batch_size_per_tensor_core. These two
+  // fields will be auto-populated by the TPUEmbedding rewrite passes.
+  repeated FeatureDescriptor feature_descriptor = 10;
+
+  // SPMD (Single Program Multiple Data) sharding configuration for
+  // TPUEmbedding. When model parallelism is used on the TensorCore, the number
+  // of cores per replica must be passed to TPUEmbedding so that the right
+  // shapes can be computed in the TF/XLA bridge.
+  message SpmdSharding {
+    // Whether SPMD sharding is enabled.
+    bool enabled = 1;
+
+    // Number of cores per replica.
+    int32 num_cores_per_replica = 2;
+
+    // If true, the tensors are manually partitioned. Otherwise, use the
+    // automatic SPMD partitioning. This should be true when users use
+    // `shard_map`.
+    bool use_manual_partitioning = 3;
+
+    // If true, lowers to Shardy custom-call for manual partitioning.
+    // Default is false.
+    bool use_shardy_partitioner = 4;
+  }
+  SpmdSharding spmd_sharding = 11;
+
+  // Old TPU embedding output layout.
+  reserved "output_layout";
+  reserved 8;
+}
+
+// A placeholder message that is used to define a unique Status payload
+// URL for TPU embedding errors.
+message TPUEmbeddingError {}

From 87b9ce2be5531bdb71f759ec6d305c197f841cf7 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 9 Oct 2026 11:16:35 -0700
Subject: [PATCH 0380/1000] Reverts bc3ca1c5863d9a244d10a8f7761f8c9cf3abc6b1

PiperOrigin-RevId: 996644768
---
 tensorflow/core/kernels/summary_op_test.cc | 75 ----------------------
 1 file changed, 75 deletions(-)

diff --git a/tensorflow/core/kernels/summary_op_test.cc b/tensorflow/core/kernels/summary_op_test.cc
index 5b082c8a8b6..3e456df7b6f 100644
--- a/tensorflow/core/kernels/summary_op_test.cc
+++ b/tensorflow/core/kernels/summary_op_test.cc
@@ -14,7 +14,6 @@ limitations under the License.
 ==============================================================================*/
 
 #include <functional>
-#include <limits>
 #include <memory>
 
 #include "tensorflow/core/framework/allocator.h"
@@ -166,80 +165,6 @@ TEST_F(SummaryHistoOpTest, Error_TooManyTagValues) {
   EXPECT_TRUE(absl::StrContains(s.ToString(), "tags must be scalar")) << s;
 }
 
-TEST_F(SummaryHistoOpTest, Error_NanInValues) {
-  MakeOp(DT_FLOAT);
-
-  // Feed and run
-  AddInputFromArray<tstring>(TensorShape({}), {"taghisto"});
-  AddInputFromArray<float>(
-      TensorShape({3, 2}),
-      {0.1f, -0.7f, 4.1f, std::numeric_limits<float>::quiet_NaN(), 5.f, 4.f});
-  absl::Status s = RunOpKernel();
-  EXPECT_TRUE(absl::StrContains(
-      s.ToString(), "INVALID_ARGUMENT: Nan in summary histogram for: myop"))
-      << s;
-}
-
-TEST_F(SummaryHistoOpTest, Error_InfinityInValues) {
-  MakeOp(DT_FLOAT);
-
-  // Feed and run
-  AddInputFromArray<tstring>(TensorShape({}), {"taghisto"});
-  AddInputFromArray<float>(
-      TensorShape({3, 2}),
-      {0.1f, -0.7f, 4.1f, std::numeric_limits<float>::infinity(), 5.f, 4.f});
-  absl::Status s = RunOpKernel();
-  EXPECT_TRUE(absl::StrContains(
-      s.ToString(), "INVALID_ARGUMENT: Infinity in Histogram for: myop"))
-      << s;
-}
-
-TEST_F(SummaryHistoOpTest, Error_NegativeInfinityInValues) {
-  MakeOp(DT_FLOAT);
-
-  // Feed and run
-  AddInputFromArray<tstring>(TensorShape({}), {"taghisto"});
-  AddInputFromArray<float>(
-      TensorShape({3, 2}),
-      {0.1f, -0.7f, 4.1f, -std::numeric_limits<float>::infinity(), 5.f, 4.f});
-  absl::Status s = RunOpKernel();
-  EXPECT_TRUE(absl::StrContains(
-      s.ToString(), "INVALID_ARGUMENT: Infinity in Histogram for: myop"))
-      << s;
-}
-
-TEST_F(SummaryHistoOpTest, Error_NanInValues_Half) {
-  MakeOp(DT_HALF);
-
-  // Feed and run
-  AddInputFromList<tstring>(TensorShape({}), {"taghisto"});
-  AddInputFromList<Eigen::half>(
-      TensorShape({3, 2}),
-      {Eigen::half(0.1f), Eigen::half(-0.7f), Eigen::half(4.1f),
-       std::numeric_limits<Eigen::half>::quiet_NaN(), Eigen::half(5.0f),
-       Eigen::half(4.0f)});
-  absl::Status s = RunOpKernel();
-  EXPECT_TRUE(absl::StrContains(
-      s.ToString(), "INVALID_ARGUMENT: Nan in summary histogram for: myop"))
-      << s;
-}
-
-TEST_F(SummaryHistoOpTest, EmptyValues_N0) {
-  MakeOp(DT_FLOAT);
-
-  // Feed and run
-  AddInputFromArray<tstring>(TensorShape({}), {"taghisto"});
-  AddInputFromArray<float>(TensorShape({0}), {});
-  TF_ASSERT_OK(RunOpKernel());
-
-  Tensor* out_tensor = GetOutput(0);
-  ASSERT_EQ(0, out_tensor->dims());
-  Summary summary;
-  ParseProtoUnlimited(&summary, out_tensor->scalar<tstring>()());
-  ASSERT_EQ(summary.value_size(), 1);
-  EXPECT_EQ(summary.value(0).tag(), "taghisto");
-}
-
 // --------------------------------------------------------------------------
 // SummaryMergeOp
 // --------------------------------------------------------------------------

From a1179d3c224967bd979950b823d5ce2a7e8f4c94 Mon Sep 17 00:00:00 2001
From: David Majnemer <majnemer@google.com>
Date: Fri, 9 Oct 2026 11:30:52 -0700
Subject: [PATCH 0381/1000] [XLA:TPU] Rewrite f64 Log with bit level range
 reduction

Add f64 Log accuracy tests covering the full exponent range, inputs near 1,
the range reduction boundaries, scalar shapes, and special values.

PiperOrigin-RevId: 996653754
---
 .../xla/tests/array_elementwise_ops_test.cc   | 59 +++++++++++++++++++
 1 file changed, 59 insertions(+)

diff --git a/third_party/xla/xla/tests/array_elementwise_ops_test.cc b/third_party/xla/xla/tests/array_elementwise_ops_test.cc
index 6166b63b005..92420c8aac9 100644
--- a/third_party/xla/xla/tests/array_elementwise_ops_test.cc
+++ b/third_party/xla/xla/tests/array_elementwise_ops_test.cc
@@ -3020,6 +3020,65 @@ TEST_F(ArrayElementwiseOpTest, LogF32sVector) {
                              error_spec);
 }
 
+TEST_F(ArrayElementwiseOpTest, LogF64s) {
+  // Inputs are sums of two normal floats so that they are exactly
+  // representable on backends that emulate f64 with a pair of f32s and flush
+  // subnormals.
+  std::vector<double> inputs;
+  for (int e = -126; e <= 127; e += 3) {
+    for (float m : {1.0f, 1.1f, 1.3f, 1.4142135f, 1.5f, 1.9f}) {
+      float hi = std::ldexp(m, e);
+      float lo = hi * 0x1.234567p-26f;
+      inputs.push_back(static_cast<double>(hi) +
+                       (std::isnormal(lo) ? static_cast<double>(lo) : 0.0));
+    }
+  }
+  for (int n = 1; n <= 48; ++n) {
+    inputs.push_back(1.0 + std::ldexp(1.0, -n));
+    inputs.push_back(1.0 - std::ldexp(1.0, -n));
+  }
+  for (float x : {0.70710677f, 0.75f, 0.9f, 1.25f, 1.41421354f, 1.41421366f}) {
+    inputs.push_back(static_cast<double>(x) +
+                     static_cast<double>(x * -0x1.9abcdep-27f));
+  }
+  inputs.push_back(static_cast<double>(std::numeric_limits<float>::min()));
+  inputs.push_back(static_cast<double>(std::numeric_limits<float>::max()));
+
+  std::vector<double> expected;
+  expected.reserve(inputs.size());
+  for (double x : inputs) {
+    expected.push_back(std::log(x));
+  }
+  constexpr double kInf = std::numeric_limits<double>::infinity();
+  constexpr double kNaN = std::numeric_limits<double>::quiet_NaN();
+  for (auto [x, y] : std::vector<std::pair<double, double>>{{1.0, 0.0},
+                                                            {0.0, -kInf},
+                                                            {-0.0, -kInf},
+                                                            {-1.0, kNaN},
+                                                            {-kInf, kNaN},
+                                                            {kInf, kInf},
+                                                            {kNaN, kNaN}}) {
+    inputs.push_back(x);
+    expected.push_back(y);
+  }
+
+  XlaBuilder builder(TestName());
+  Log(ConstantR1<double>(&builder, inputs));
+  ComputeAndCompareR1<double>(&builder, expected, {}, ErrorSpec(0, 0x1p-45));
+}
+
+// Effective scalar shapes take a different simplification path than vectors.
+// The input is a parameter so that the computation is not constant folded.
+TEST_F(ArrayElementwiseOpTest, LogF64Scalar) {
+  for (double x : {3.0, 0.75, 1e10}) {
+    XlaBuilder builder(TestName());
+    Literal arg = LiteralUtil::CreateR0<double>(x);
+    Log(Parameter(&builder, 0, arg.shape(), "x"));
+    ComputeAndCompareR0<double>(&builder, std::log(x), {&arg},
+                                ErrorSpec(0, 0x1p-45));
+  }
+}
+
 TEST_F(ArrayElementwiseOpTest, ClzU32s) {
   XlaBuilder builder(TestName());
   auto a = ConstantR1<uint32_t>(

From a9cd9f57365090e4481df002b824d8653664d136 Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Fri, 9 Oct 2026 11:41:26 -0700
Subject: [PATCH 0382/1000] [XLA] Fix pad handling in
 AlgebraicSimplifier::AssociativeReorderDotOperator

When `use_associative_reordering` is enabled, `AssociativeReorderDotOperator` rewrites `dot(pad(x, pad_val), y)` into `dot(x, slice(y))` along contracting dimensions. This transformation has three bugs. First, it does not check that `pad_val` is zero, which silently drops nonzero padding contributions from the dot product. Second, it reorders contracting dimensions whenever any padding field is positive, even if `edge_padding_low` or `edge_padding_high` is negative, producing invalid slice bounds. Third, when checking whether a residual `pad` instruction is still required, it checks only `edge_padding_low` and `edge_padding_high` and ignores `interior_padding`. This omission drops interior padding on noncontracting dimensions and causes a shape mismatch.

Require `IsScalarConstantZero(reorder_from->operand(1))` before reordering a `pad` operand of a `dot`. Require both `edge_padding_low` and `edge_padding_high` to be nonnegative on a contracting dimension before converting its padding into a slice. Check `dim.interior_padding() != 0` when determining whether a residual `pad` instruction is needed.

Tests:
- `//third_party/tensorflow/compiler/xla/hlo/transforms/simplifiers:algebraic_simplifier_test`: `AssociativeReorderDotPadNonZeroOrInteriorPad` verifies nonzero padding values, negative edge padding, and residual interior padding on noncontracting dimensions.
PiperOrigin-RevId: 996660351
---
 .../simplifiers/algebraic_simplifier.cc       | 14 ++--
 .../simplifiers/algebraic_simplifier_test.cc  | 68 +++++++++++++++++++
 2 files changed, 78 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
index 9d0d7941751..8bf9b0f3581 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
@@ -3770,6 +3770,9 @@ AlgebraicSimplifierVisitor::AssociativeReorderDotOperator(
             MakeReverseHlo(reorder_from->mutable_operand(0), unreordered_dims));
       }
     } else if (opcode == HloOpcode::kPad) {
+      if (!IsScalarConstantZero(reorder_from->operand(1))) {
+        return nullptr;
+      }
       // Padding of dot contracting dimensions can be reordered to slices of
       // the corresponding contracting dimensions in the other dot operand
       DimensionVector start_indices, limit_indices, strides;
@@ -3791,9 +3794,11 @@ AlgebraicSimplifierVisitor::AssociativeReorderDotOperator(
 
           // Edge padding can be negative which acts as a slice. If this is
           // the case, we don't want to reorder
-          if (padding_dimension.edge_padding_low() > 0 ||
-              padding_dimension.edge_padding_high() > 0 ||
-              padding_dimension.interior_padding() > 0) {
+          if (padding_dimension.edge_padding_low() >= 0 &&
+              padding_dimension.edge_padding_high() >= 0 &&
+              (padding_dimension.edge_padding_low() > 0 ||
+               padding_dimension.edge_padding_high() > 0 ||
+               padding_dimension.interior_padding() > 0)) {
             make_hlo = true;
             start_index += padding_dimension.edge_padding_low();
             limit_index -= padding_dimension.edge_padding_high();
@@ -3822,7 +3827,8 @@ AlgebraicSimplifierVisitor::AssociativeReorderDotOperator(
 
       // Check if we still need a padding instruction, and create Hlo if so
       for (auto& dim : new_padding_config.dimensions()) {
-        if (dim.edge_padding_low() != 0 || dim.edge_padding_high() != 0) {
+        if (dim.edge_padding_low() != 0 || dim.edge_padding_high() != 0 ||
+            dim.interior_padding() != 0) {
           // Want to use a greater threshold if reordering means increasing
           // the number of Hlos
           threshold_multiplier = 2.0;
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
index 09333d0e06f..ec5e5673995 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
@@ -8205,6 +8205,74 @@ TEST_F(AlgebraicSimplifierTest, DotPadRightReorder) {
                                 m::Pad(m::Parameter(1), m::Constant()))));
 }
 
+TEST_F(AlgebraicSimplifierTest, AssociativeReorderDotPadNonZeroOrInteriorPad) {
+  AlgebraicSimplifierOptions options = default_options_;
+  options.set_use_associative_reordering(true);
+  options.set_associative_reordering_threshold(1.1);
+  options.set_enable_negative_padding_replacement(false);
+  AlgebraicSimplifier simplifier(options);
+
+  // Nonzero pad value on a contracting dimension must not be dropped.
+  constexpr absl::string_view kNonZeroPadHlo = R"(
+    HloModule nonzero_pad
+    ENTRY test {
+      one = f32[] constant(1.0)
+      a = f32[8,5] parameter(0)
+      b = f32[10,6] parameter(1)
+      pad = f32[8,10] pad(a, one), padding=0_0x0_5
+      ROOT dot = f32[8,6] dot(pad, b),
+                 lhs_contracting_dims={1},
+                 rhs_contracting_dims={0}
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto nonzero_module,
+                       ParseAndReturnVerifiedModule(kNonZeroPadHlo));
+  EXPECT_THAT(simplifier.Run(nonzero_module.get()),
+              absl_testing::IsOkAndHolds(false));
+
+  // Negative edge padding combined with positive padding on a contracting
+  // dimension must not be reordered into an invalid slice.
+  constexpr absl::string_view kNegativeEdgePadHlo = R"(
+    HloModule negative_edge_pad
+    ENTRY test {
+      zero = f32[] constant(0.0)
+      a = f32[8,5] parameter(0)
+      b = f32[10,6] parameter(1)
+      pad = f32[8,10] pad(a, zero), padding=0_0x-1_6
+      ROOT dot = f32[8,6] dot(pad, b),
+                 lhs_contracting_dims={1},
+                 rhs_contracting_dims={0}
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto neg_edge_module,
+                       ParseAndReturnVerifiedModule(kNegativeEdgePadHlo));
+  EXPECT_THAT(simplifier.Run(neg_edge_module.get()),
+              absl_testing::IsOkAndHolds(false));
+
+  // Interior padding on a noncontracting dimension must be preserved in the
+  // residual pad when contracting dimension padding is reordered to a slice.
+  constexpr absl::string_view kInteriorPadNonContractingHlo = R"(
+    HloModule interior_pad_non_contracting
+    ENTRY test {
+      zero = f32[] constant(0.0)
+      a = f32[4,5] parameter(0)
+      b = f32[20,6] parameter(1)
+      pad = f32[7,20] pad(a, zero), padding=0_0_1x0_15_0
+      ROOT dot = f32[7,6] dot(pad, b),
+                 lhs_contracting_dims={1},
+                 rhs_contracting_dims={0}
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(
+      auto interior_module,
+      ParseAndReturnVerifiedModule(kInteriorPadNonContractingHlo));
+  ASSERT_THAT(simplifier.Run(interior_module.get()),
+              absl_testing::IsOkAndHolds(true));
+  EXPECT_THAT(interior_module->entry_computation()->root_instruction(),
+              GmockMatch(m::Dot(m::Pad(m::Parameter(0), m::Constant()),
+                                m::Slice(m::Parameter(1)))));
+}
+
 // This pattern appears in translate_inference_bnmt_v15_vf_lite_execution_test
 TEST_F(AlgebraicSimplifierTest, DotBroadcastLeftReorder) {
   constexpr absl::string_view hlo_string = R"(

From 8adcdda90db503b67e8b749568ecd637a9d63e44 Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Fri, 9 Oct 2026 11:52:46 -0700
Subject: [PATCH 0383/1000] Remove unused `tsl/platform:errors` and
 `tsl/platform:statusor` dependencies from
 `dead_dynamic_update_slice_elimination`. These headers were removed from
 `dead_dynamic_update_slice_elimination.cc` when switching to
 `absl/status/status_macros.h`.

PiperOrigin-RevId: 996667218
---
 third_party/xla/xla/hlo/transforms/simplifiers/BUILD | 2 --
 1 file changed, 2 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD
index d311f253ac6..cb1de69e20d 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD
@@ -2134,8 +2134,6 @@ cc_library(
         "//xla:shape_util",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass",
-        "//xla/tsl/platform:errors",
-        "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",

From 44f69ac8ce931f1aa4c7b8990c54e89ebe4ee7b2 Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Fri, 9 Oct 2026 12:17:49 -0700
Subject: [PATCH 0384/1000] ## Problem

`WhileLoopUnroller` substitutes wrong induction variable values when unrolling a while loop whose induction step exceeds one. `WhileLoopUnroller::IsLoopUnrollable` accepts loops with arbitrary positive step sizes through `MatchTrivialLoopTripCount`, which divides the bound difference by the step size. However, `WhileLoopConfig` did not record the induction step size. Both `UnrollInternal` and `UnrollInternalWrappedAndReturnReplacement` computed the induction value for iteration `iter` as `config.init + iter` instead of scaling `iter` by the step size.

## Change

Add `step` with a default value of `1` to `WhileLoopConfig` and populate it in `WhileLoopUnroller::IsLoopUnrollable` via `MatchTrivialLoopInductionStep`. Update `UnrollInternal` and `UnrollInternalWrappedAndReturnReplacement` to compute each unrolled iteration value as `config.init + iter * config.step` using `OverflowSafeMultiply` and `OverflowSafeAdd`.

Existing programs with unit step loops are unchanged. While loops with step sizes above one now receive the true induction value in each unrolled iteration instead of consecutive integers. No flags are added or changed.

## Tests

- `//third_party/tensorflow/compiler/xla/service:while_loop_unroller_test`: pins direct and trivial loop wrapped unrolling of a while loop with step size two accumulating the induction variable.
- `//third_party/tensorflow/compiler/xla/backends/gpu/transforms:windowed_einsum_handler_test`: updates the second unrolled iteration in `ReduceScatterF8` (which has step size two) to expect induction value `2` rather than `1`.

PiperOrigin-RevId: 996682341
---
 .../windowed_einsum_handler_test.cc           |  2 +-
 .../xla/xla/service/while_loop_unroller.cc    | 17 +++++---
 .../xla/xla/service/while_loop_unroller.h     |  2 +
 .../xla/service/while_loop_unroller_test.cc   | 42 +++++++++++++++++--
 4 files changed, 53 insertions(+), 10 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler_test.cc b/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler_test.cc
index fb87fe4a180..883451210b6 100644
--- a/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/windowed_einsum_handler_test.cc
@@ -931,7 +931,7 @@ ENTRY main.9_spmd {
 ; CHECK-DAG:         "force_earliest_schedule":false
 ; CHECK-NEXT:    [[ADD3:%[^ ]+]] = f32[2,512,24576]{2,1,0} add([[CP0]], [[DOT0]])
 ; CHECK-NEXT:    [[GTE6:[^ ]+]] = f32[2,512,24576]{2,1,0} get-tuple-element([[TUPLE0]]), index=3
-; CHECK-NEXT:    [[C11:%[^ ]+]] = u32[] constant(1)
+; CHECK-NEXT:    [[C11:%[^ ]+]] = u32[] constant(2)
 ; CHECK-NEXT:    [[ADD6:%[^ ]+]] = u32[] add([[C11]], [[PID]])
 ; CHECK-NEXT:    [[AND1:%[^ ]+]] = u32[] and([[ADD6]], [[C2]])
 ; CHECK-NEXT:    [[CLAMP1:%[^ ]+]] = u32[] clamp([[C1]], [[AND1]], [[C2]])
diff --git a/third_party/xla/xla/service/while_loop_unroller.cc b/third_party/xla/xla/service/while_loop_unroller.cc
index 100bef5c336..6f13f48bf05 100644
--- a/third_party/xla/xla/service/while_loop_unroller.cc
+++ b/third_party/xla/xla/service/while_loop_unroller.cc
@@ -302,8 +302,10 @@ absl::StatusOr<bool> UnrollInternal(HloInstruction* while_op,
                    NextSchedulingGroupId(*while_op->GetModule()));
   std::vector<HloInstruction*> new_calls;
   new_calls.reserve(config.trip_count);
-  for (int64_t i = config.init; i < config.trip_count + config.init; ++i) {
-    CHECK(OverflowSafeAdd(i, (int64_t)1).has_value());
+  for (int64_t iter = 0; iter < config.trip_count; ++iter) {
+    int64_t i = OverflowSafeAdd(config.init,
+                                OverflowSafeMultiply(iter, config.step).first)
+                    .value();
 
     HloComputation* unrolled_body = module->AddEmbeddedComputation(
         UnrollSingleIterationOfTrivialLoop(while_op, config, i,
@@ -350,8 +352,10 @@ absl::StatusOr<UnrollResult> UnrollInternalWrappedAndReturnReplacement(
 
   std::vector<HloInstruction*> new_calls;
   new_calls.reserve(config.trip_count);
-  for (int64_t i = config.init; i < config.trip_count + config.init; ++i) {
-    CHECK(OverflowSafeAdd(i, (int64_t)1).has_value());
+  for (int64_t iter = 0; iter < config.trip_count; ++iter) {
+    int64_t i = OverflowSafeAdd(config.init,
+                                OverflowSafeMultiply(iter, config.step).first)
+                    .value();
 
     HloComputation* unrolled_body = module->AddEmbeddedComputation(
         UnrollSingleIterationOfTrivialLoop(while_op, config, i,
@@ -1296,7 +1300,9 @@ std::optional<int64_t> AdvancedMatchShapeCoveringDynamicIndexInstruction(
   Literal indvar_iter_val = std::move(indvar_init_result).value();
   std::optional<int64_t> trip_count =
       MatchTrivialLoopTripCount(while_op, *indvar_tuple_idx, indvar_iter_val);
-  if (!trip_count.has_value()) {
+  std::optional<int64_t> step =
+      MatchTrivialLoopInductionStep(while_op, *indvar_tuple_idx);
+  if (!trip_count.has_value() || !step.has_value()) {
     VLOG(3) << "Loop doesn't have trivial trip count";
     return std::nullopt;
   }
@@ -1309,6 +1315,7 @@ std::optional<int64_t> AdvancedMatchShapeCoveringDynamicIndexInstruction(
       LiteralUtil::LiteralAsScalarInt64(std::move(indvar_iter_val)).value();
   config.trip_count = trip_count.value();
   config.induction_var_idx = *indvar_tuple_idx;
+  config.step = *step;
   return config;
 }
 
diff --git a/third_party/xla/xla/service/while_loop_unroller.h b/third_party/xla/xla/service/while_loop_unroller.h
index f42a716a741..11e2837f925 100644
--- a/third_party/xla/xla/service/while_loop_unroller.h
+++ b/third_party/xla/xla/service/while_loop_unroller.h
@@ -49,6 +49,8 @@ struct WhileLoopConfig {
   int64_t trip_count;
   // The index of the induction variable in the input tuple of the while loop.
   int64_t induction_var_idx;
+  // The step size of the induction variable.
+  int64_t step = 1;
 };
 
 // Result for unrolled while loops.
diff --git a/third_party/xla/xla/service/while_loop_unroller_test.cc b/third_party/xla/xla/service/while_loop_unroller_test.cc
index b7b0c765b1b..73b3589d746 100644
--- a/third_party/xla/xla/service/while_loop_unroller_test.cc
+++ b/third_party/xla/xla/service/while_loop_unroller_test.cc
@@ -437,10 +437,10 @@ TEST_F(WhileLoopUnrollerTest, SimpleLoopUnrollNeedPrepare) {
     ROOT result = s32[3]{0} get-tuple-element(while), index=1
   }
   )";
-  UnrollAndCompare(ParseAndReturnVerifiedModule(hlo_string).value(), {}, -1,
-                   false);
-  UnrollAndCompare(ParseAndReturnVerifiedModule(hlo_string).value(), {}, -1,
-                   true);
+  ASSERT_OK_AND_ASSIGN(auto m1, ParseAndReturnVerifiedModule(hlo_string));
+  UnrollAndCompare(std::move(m1), {}, -1, false);
+  ASSERT_OK_AND_ASSIGN(auto m2, ParseAndReturnVerifiedModule(hlo_string));
+  UnrollAndCompare(std::move(m2), {}, -1, true);
 }
 
 // This test passes because we run TupleSimplifier before unrolling.
@@ -710,6 +710,40 @@ TEST_F(WhileLoopUnrollerTest, SimpleLoopNonZeroInit) {
                    true);
 }
 
+TEST_F(WhileLoopUnrollerTest, SimpleLoopNonUnitStep) {
+  std::string hlo_string = R"(
+  HloModule SimpleLoop
+  SimpleLoop.body {
+    loop_var.1 = (s32[], s32[]) parameter(0)
+    i = s32[] get-tuple-element(loop_var.1), index=0
+    step = s32[] constant(2)
+    next_i = s32[] add(i, step)
+    sum = s32[] get-tuple-element(loop_var.1), index=1
+    next_sum = s32[] add(sum, i)
+    ROOT tuple = (s32[], s32[]) tuple(next_i, next_sum)
+  }
+  SimpleLoop.condition {
+    loop_var.2 = (s32[], s32[]) parameter(0)
+    i = s32[] get-tuple-element(loop_var.2), index=0
+    bound = s32[] constant(6)
+    ROOT less-than = pred[] compare(i, bound), direction=LT
+  }
+  ENTRY SimpleLoop {
+    init_i = s32[] constant(0)
+    init_sum = s32[] constant(0)
+    tuple.1 = (s32[], s32[]) tuple(init_i, init_sum)
+    ROOT while = (s32[], s32[]) while(tuple.1), condition=
+      SimpleLoop.condition, body=SimpleLoop.body
+  }
+  )";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module_direct,
+                       ParseAndReturnVerifiedModule(hlo_string));
+  UnrollAndCompare(std::move(module_direct), {}, -1, false);
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module_wrapped,
+                       ParseAndReturnVerifiedModule(hlo_string));
+  UnrollAndCompare(std::move(module_wrapped), {}, -1, true);
+}
+
 TEST_F(WhileLoopUnrollerTest, SimpleLoopS16IndVar) {
   std::string hlo_string = R"(
   HloModule SimpleLoop

From 8a29d778c521e707394be19f92d1e55cf844c93e Mon Sep 17 00:00:00 2001
From: Amir Saadat <amirsaadat@google.com>
Date: Fri, 9 Oct 2026 13:05:12 -0700
Subject: [PATCH 0385/1000] [XLA:hlo_dump] Let MismatchDetails opt out of
 element-level tensor data.

Adds `MismatchDetails::has_element_level_data` (default true). Producers that
only know that an instruction mismatched, not which elements, can set it to
false so the tensor inspector is suppressed instead of showing a fabricated
single-element mismatch. Also adds an optional `output_dir` parameter to
`DumpHloModuleMismatchWithGraphData`.

PiperOrigin-RevId: 996709131
---
 third_party/xla/xla/tools/hlo_dump/BUILD      |   2 +
 .../xla/xla/tools/hlo_dump/hlo_dump_utils.cc  |  38 +++++-
 .../xla/xla/tools/hlo_dump/hlo_dump_utils.h   |  19 ++-
 .../xla/tools/hlo_dump/hlo_dump_utils_test.cc | 129 ++++++++++++++++++
 4 files changed, 180 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/xla/tools/hlo_dump/BUILD b/third_party/xla/xla/tools/hlo_dump/BUILD
index 6c326148d7b..c3c551b3795 100644
--- a/third_party/xla/xla/tools/hlo_dump/BUILD
+++ b/third_party/xla/xla/tools/hlo_dump/BUILD
@@ -53,6 +53,7 @@ cc_library(
         "@com_google_absl//absl/base",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
@@ -97,5 +98,6 @@ xla_cc_test(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/platform:path",
     ],
 )
diff --git a/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils.cc b/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils.cc
index 915377a69f4..96ef805b27f 100644
--- a/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils.cc
+++ b/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils.cc
@@ -36,6 +36,7 @@ limitations under the License.
 #include "absl/base/casts.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/escaping.h"
 #include "absl/strings/match.h"
@@ -1264,6 +1265,21 @@ void ApplyBoundingBoxTail(const MismatchBoundingBox& bbox, double rel_error,
       !bbox.pattern.empty() ? bbox.pattern : ClassifyMismatchPattern(bbox);
 }
 
+// Strips everything that would let the tensor inspector draw element-level
+// mismatches. The front-end treats a tensor with no shape and no elements as
+// non-inspectable and keeps the inspector closed for it, which is the honest
+// rendering when the producer does not know which elements mismatched.
+void ClearElementLevelData(TensorVisualizationInfo& info) {
+  info.shape.clear();
+  info.box_min.clear();
+  info.box_max.clear();
+  info.top_mismatches.clear();
+  info.mismatched_slices.clear();
+  info.slice_boxes.clear();
+  info.mismatch_count = 0;
+  info.total_elements = 0;
+}
+
 }  // namespace
 
 absl::flat_hash_map<std::string, TensorVisualizationInfo>
@@ -1310,7 +1326,9 @@ PopulateTensorVisualizations(const HloModule& module,
       if (it != instr_to_mismatch.end()) {
         const MismatchDetails* mismatch = it->second;
         info.has_mismatch = true;
-        if (mismatch->bounding_box.has_value()) {
+        if (!mismatch->has_element_level_data) {
+          ClearElementLevelData(info);
+        } else if (mismatch->bounding_box.has_value()) {
           const auto& bbox = *mismatch->bounding_box;
           if (!bbox.tensor_shape.empty()) {
             info.shape = bbox.tensor_shape;
@@ -1340,7 +1358,9 @@ PopulateTensorVisualizations(const HloModule& module,
       TensorVisualizationInfo& info = it->second;
       info.instruction_name = mismatch.target_instruction_name;
       info.has_mismatch = true;
-      if (mismatch.bounding_box.has_value()) {
+      if (!mismatch.has_element_level_data) {
+        ClearElementLevelData(info);
+      } else if (mismatch.bounding_box.has_value()) {
         const auto& bbox = *mismatch.bounding_box;
         info.shape = bbox.tensor_shape;
         info.box_min = bbox.box_min;
@@ -1799,7 +1819,15 @@ GraphData PopulateMismatchGraphData(
 
 absl::StatusOr<std::string> DumpHloModuleMismatchWithGraphData(
     const HloModule& module, absl::Span<const MismatchDetails> mismatches,
-    absl::string_view output_filename) {
+    absl::string_view output_filename, absl::string_view output_dir) {
+  if (!output_dir.empty()) {
+    // Fail before the render rather than after it.
+    absl::Status status =
+        tsl::Env::Default()->RecursivelyCreateDir(std::string(output_dir));
+    if (!status.ok()) {
+      return status;
+    }
+  }
   absl::flat_hash_map<TensorKey, TensorAnnotation> annotations =
       PopulateMismatchAnnotations(module, mismatches);
   GraphData graph_data = PopulateMismatchGraphData(module, mismatches);
@@ -1813,7 +1841,9 @@ absl::StatusOr<std::string> DumpHloModuleMismatchWithGraphData(
   const char* env_dir = std::getenv("TEST_UNDECLARED_OUTPUTS_DIR");
   std::string outdir;
   std::string html_filename;
-  if (env_dir != nullptr && env_dir[0] != '\0') {
+  if (!output_dir.empty()) {
+    html_filename = tsl::io::JoinPath(output_dir, output_filename);
+  } else if (env_dir != nullptr && env_dir[0] != '\0') {
     outdir = env_dir;
     html_filename = tsl::io::JoinPath(outdir, output_filename);
   } else if (tsl::io::GetTestUndeclaredOutputsDir(&outdir)) {
diff --git a/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils.h b/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils.h
index 31d55982ecf..db1f0565df1 100644
--- a/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils.h
+++ b/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils.h
@@ -87,6 +87,12 @@ struct MismatchDetails {
   std::optional<bool> result_of_reduce;
   std::optional<std::string> custom_description;
   std::optional<MismatchBoundingBox> bounding_box;
+  // Set to false when the producer only knows *that* the tensor mismatched, not
+  // *where* (e.g. an on-device comparison that only reports reduced scalars).
+  // The tensor inspector is then suppressed for this instruction, rather than
+  // falling back to a fabricated single-element mismatch at the origin. HLO
+  // highlighting, the graph view, and tooltips are unaffected.
+  bool has_element_level_data = true;
 };
 
 struct TensorKey {
@@ -181,12 +187,17 @@ absl::flat_hash_map<TensorKey, TensorAnnotation> PopulateMismatchAnnotations(
 GraphData PopulateMismatchGraphData(
     const HloModule& module, absl::Span<const MismatchDetails> mismatches);
 
-// Generates an interactive HTML dump for an HLO module
-// and writes it to the test undeclared outputs directory (or temp dir).
-// Returns the absolute path of the written file.
+// Generates an interactive HTML dump for an HLO module and writes it to
+// `output_dir`/`output_filename`, creating `output_dir` if needed.
+//
+// If `output_dir` is empty, the file goes to the test undeclared outputs
+// directory when running under a test. Otherwise it goes to a randomly named
+// temp file, and `output_filename` is only used as that file's suffix.
+//
+// Returns the path of the written file.
 absl::StatusOr<std::string> DumpHloModuleMismatchWithGraphData(
     const HloModule& module, absl::Span<const MismatchDetails> mismatches,
-    absl::string_view output_filename);
+    absl::string_view output_filename, absl::string_view output_dir = "");
 
 }  // namespace xla::numerics::debug_info
 
diff --git a/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils_test.cc b/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils_test.cc
index 25fba6a0c93..f3648af4e6c 100644
--- a/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils_test.cc
+++ b/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils_test.cc
@@ -41,6 +41,7 @@ limitations under the License.
 #include "xla/tsl/platform/statusor.h"
 #include "xla/tsl/platform/test.h"
 #include "xla/xla_data.pb.h"
+#include "tsl/platform/path.h"
 
 namespace xla::numerics::debug_info {
 namespace {
@@ -1256,6 +1257,104 @@ ENTRY main {
   EXPECT_TRUE(absl::StrContains(html, "\"mismatch_count\": 15"));
 }
 
+TEST(HloDumpUtilsTest, PopulateTensorVisualizationsDefaultsToSingleElement) {
+  // Guards existing producers: without a bounding box, the historical fallback
+  // of a single mismatching element is preserved.
+  const absl::string_view hlo_string = R"hlo(
+HloModule test_vis
+ENTRY main {
+  ROOT %p0 = f32[4] parameter(0)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnUnverifiedModule(hlo_string));
+  MismatchDetails mismatch;
+  mismatch.target_instruction_name = "p0";
+
+  auto vis_map = PopulateTensorVisualizations(*module, {mismatch});
+  ASSERT_TRUE(vis_map.contains("p0"));
+  EXPECT_TRUE(vis_map["p0"].has_mismatch);
+  EXPECT_EQ(vis_map["p0"].shape, (std::vector<int64_t>{4}));
+  EXPECT_EQ(vis_map["p0"].mismatch_count, 1);
+  EXPECT_EQ(vis_map["p0"].total_elements, 4);
+}
+
+TEST(HloDumpUtilsTest, PopulateTensorVisualizationsWithoutElementLevelData) {
+  const absl::string_view hlo_string = R"hlo(
+HloModule test_vis
+ENTRY main {
+  ROOT %p0 = f32[10,20] parameter(0)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnUnverifiedModule(hlo_string));
+
+  MismatchDetails in_module;
+  in_module.target_instruction_name = "p0";
+  in_module.has_element_level_data = false;
+  // Even a supplied bounding box must not be rendered: the flag wins. Populate
+  // every element-level field so each assertion below can actually fail.
+  in_module.bounding_box = MakeBoundingBox({10, 20}, {1, 2}, {5, 8}, 15, 200,
+                                           "block", {{1, 2}, {5, 8}});
+  in_module.bounding_box->mismatched_slices = {1, 3};
+  in_module.bounding_box->slice_boxes["1"] = SliceBoundingBox{};
+
+  MismatchDetails not_in_module;
+  not_in_module.target_instruction_name = "absent";
+  not_in_module.has_element_level_data = false;
+
+  auto vis_map =
+      PopulateTensorVisualizations(*module, {in_module, not_in_module});
+  for (absl::string_view name : {"p0", "absent"}) {
+    SCOPED_TRACE(name);
+    ASSERT_TRUE(vis_map.contains(name));
+    const auto& item = vis_map[name];
+    // The mismatch is still reported, so HLO/graph highlighting keep working,
+    // but the tensor is non-inspectable so the inspector stays closed.
+    EXPECT_TRUE(item.has_mismatch);
+    EXPECT_TRUE(item.shape.empty());
+    EXPECT_TRUE(item.box_min.empty());
+    EXPECT_TRUE(item.box_max.empty());
+    EXPECT_TRUE(item.top_mismatches.empty());
+    EXPECT_TRUE(item.mismatched_slices.empty());
+    EXPECT_TRUE(item.slice_boxes.empty());
+    EXPECT_EQ(item.mismatch_count, 0);
+    EXPECT_EQ(item.total_elements, 0);
+  }
+}
+
+TEST(HloDumpUtilsTest,
+     PopulateTensorVisualizationsWithoutElementLevelDataTupleOutput) {
+  // A tuple target takes the tuple-element branch of the shape lookup, which
+  // must not re-populate the shape once element-level data is disabled.
+  const absl::string_view hlo_string = R"hlo(
+HloModule test_vis
+ENTRY main {
+  p0 = f32[10,20] parameter(0)
+  p1 = s32[7] parameter(1)
+  ROOT t = (f32[10,20], s32[7]) tuple(p0, p1)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnUnverifiedModule(hlo_string));
+
+  MismatchDetails mismatch;
+  mismatch.target_instruction_name = "t";
+  mismatch.output_shape_index = 1;
+  mismatch.has_element_level_data = false;
+
+  absl::flat_hash_map<std::string, TensorVisualizationInfo> vis_map =
+      PopulateTensorVisualizations(*module, {mismatch});
+  ASSERT_TRUE(vis_map.contains("t"));
+  const auto& item = vis_map["t"];
+  EXPECT_TRUE(item.has_mismatch);
+  EXPECT_TRUE(item.shape.empty());
+  EXPECT_TRUE(item.box_min.empty());
+  EXPECT_TRUE(item.box_max.empty());
+  EXPECT_EQ(item.mismatch_count, 0);
+  EXPECT_EQ(item.total_elements, 0);
+}
+
 TEST(HloDumpUtilsTest, SerializeTensorVisualizationsJsSpecialFloats) {
   absl::flat_hash_map<std::string, TensorVisualizationInfo> vis_map;
   TensorVisualizationInfo item;
@@ -1694,5 +1793,35 @@ ENTRY main {
   }
 }
 
+TEST(HloDumpUtilsTest, DumpHloModuleMismatchWithGraphDataHonorsOutputDir) {
+  const absl::string_view hlo_string = R"hlo(
+HloModule out_dir_module
+ENTRY main {
+  ROOT p0 = f32[4] parameter(0)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnUnverifiedModule(hlo_string));
+  MismatchDetails mismatch;
+  mismatch.target_instruction_name = "p0";
+
+  // A nested, not-yet-existing directory must be created, and the filename
+  // must be used verbatim rather than as a temp-file suffix.
+  const std::string dir =
+      tsl::io::JoinPath(::testing::TempDir(), "hlo_dump_out_dir", "nested");
+  ASSERT_OK_AND_ASSIGN(const std::string path,
+                       DumpHloModuleMismatchWithGraphData(*module, {mismatch},
+                                                          "page.html", dir));
+  EXPECT_EQ(path, tsl::io::JoinPath(dir, "page.html"));
+  EXPECT_TRUE(tsl::Env::Default()->FileExists(path).ok());
+
+  // Writing into the now-existing directory must also succeed.
+  ASSERT_OK_AND_ASSIGN(const std::string second,
+                       DumpHloModuleMismatchWithGraphData(*module, {mismatch},
+                                                          "page2.html", dir));
+  EXPECT_EQ(second, tsl::io::JoinPath(dir, "page2.html"));
+  EXPECT_TRUE(tsl::Env::Default()->FileExists(second).ok());
+}
+
 }  // namespace
 }  // namespace xla::numerics::debug_info

From ba16e4fba02bd050b4b43c4402523ce435fa0101 Mon Sep 17 00:00:00 2001
From: Ilya Tokar <tokarip@google.com>
Date: Fri, 9 Oct 2026 13:05:53 -0700
Subject: [PATCH 0386/1000] Port sparse_matmul_op to highway
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This gives us avx512 via dynamic dispatch for a ~10% speed-up:
name                                                                                                            items/s       items/s     vs base
BM_Sparse_replicated_400_800_2560_85_0_6/real_time                    [copies: 6 sp_a: 0.85 sp_b: 0.00      ]    950.7G ± 2%   1100.7G ±  2%  +15.78% (p=0.000 n=20)
BM_Sparse_replicated_400_2560_1024_85_0_6/real_time                   [copies: 6 sp_a: 0.85 sp_b: 0.00      ]    974.0G ± 1%   1100.4G ±  1%  +12.98% (p=0.000 n=20)
BM_Sparse_replicated_400_1024_256_85_0_6/real_time                    [copies: 6 sp_a: 0.85 sp_b: 0.00      ]    601.7G ± 0%    622.0G ±  1%   +3.37% (p=0.000 n=20)
BM_Sparse_replicated_400_256_1_85_0_6/real_time                       [copies: 6 sp_a: 0.85 sp_b: 0.00      ]    3.131G ± 2%    3.058G ±  1%   -2.32% (p=0.013 n=20)
BM_Sparse_replicated_2048_1792_1024_85_0_6/real_time                  [copies: 6 sp_a: 0.85 sp_b: 0.00      ]    1.097T ± 0%    1.278T ±  2%  +16.42% (p=0.000 n=20)
BM_Sparse_replicated_2048_1024_768_85_0_6/real_time                   [copies: 6 sp_a: 0.85 sp_b: 0.00      ]    1.033T ± 0%    1.187T ±  1%  +14.90% (p=0.000 n=20)
BM_Sparse_replicated_2048_768_512_85_0_6/real_time                    [copies: 6 sp_a: 0.85 sp_b: 0.00      ]    880.2G ± 0%    989.6G ±  1%  +12.43% (p=0.000 n=20)
BM_Sparse_replicated_2048_512_256_85_0_6/real_time                    [copies: 6 sp_a: 0.85 sp_b: 0.00      ]    654.6G ± 0%    710.0G ±  1%   +8.45% (p=0.000 n=20)
BM_Sparse_Multi_1024_2140_4096_0_82_1/real_time                       [1024_2140_4096_1_0.00_0.82           ]    943.2G ± 0%   1099.9G ±  2%  +16.61% (p=0.000 n=20)
BM_Sparse_Multi_1024_4096_2048_83_83_1/real_time                      [1024_4096_2048_1_0.83_0.83           ]    1.045T ± 0%    1.246T ±  2%  +19.28% (p=0.000 n=20)
BM_Sparse_Multi_400_800_2560_85_85_1/real_time                        [400_800_2560_1_0.85_0.85             ]    794.9G ± 1%    886.6G ±  2%  +11.53% (p=0.000 n=20)
BM_Sparse_Multi_400_2560_1024_85_85_1/real_time                       [400_2560_1024_1_0.85_0.85            ]    834.3G ± 0%    941.0G ±  2%  +12.79% (p=0.000 n=20)
BM_Sparse_Multi_400_1024_256_85_85_1/real_time                        [400_1024_256_1_0.85_0.85             ]    535.1G ± 1%    563.0G ±  3%   +5.21% (p=0.003 n=20)
BM_Sparse_Multi_400_256_1_85_85_1/real_time                           [400_256_1_1_0.85_0.85                ]    2.159G ± 1%    1.895G ± 14%  -12.23% (p=0.000 n=20)
BM_Sparse_Multi_2048_1792_1024_85_85_1/real_time                      [2048_1792_1024_1_0.85_0.85           ]    1.108T ± 0%    1.292T ±  2%  +16.60% (p=0.000 n=20)
BM_Sparse_Multi_2048_1024_768_85_85_1/real_time                       [2048_1024_768_1_0.85_0.85            ]    1.001T ± 0%    1.138T ±  2%  +13.78% (p=0.000 n=20)
BM_Sparse_Multi_2048_768_512_85_85_1/real_time                        [2048_768_512_1_0.85_0.85             ]    889.1G ± 1%    975.7G ±  1%   +9.74% (p=0.000 n=20)
BM_Sparse_Multi_2048_512_256_85_85_1/real_time                        [2048_512_256_1_0.85_0.85             ]    641.5G ± 1%    661.6G ±  2%   +3.12% (p=0.000 n=20)
BM_Sparse_Multi_2048_1792_1024_85_85_3/real_time                      [2048_1792_1024_3_0.85_0.85           ]    1.108T ± 0%    1.298T ±  1%  +17.18% (p=0.000 n=20)
BM_Sparse_Multi_2048_1024_768_85_85_3/real_time                       [2048_1024_768_3_0.85_0.85            ]    996.3G ± 0%   1142.5G ±  1%  +14.68% (p=0.000 n=20)
BM_Sparse_Multi_2048_768_512_85_85_3/real_time                        [2048_768_512_3_0.85_0.85             ]    888.9G ± 0%   1000.9G ±  1%  +12.60% (p=0.000 n=20)
BM_Sparse_Multi_2048_512_256_85_85_3/real_time                        [2048_512_256_3_0.85_0.85             ]    689.6G ± 0%    746.7G ±  2%   +8.28% (p=0.000 n=20)
BM_Sparse_Multi_2048_1792_1024_85_85_6/real_time                      [2048_1792_1024_6_0.85_0.85           ]    1.082T ± 0%    1.267T ±  1%  +17.08% (p=0.000 n=20)
BM_Sparse_Multi_2048_1024_768_85_85_6/real_time                       [2048_1024_768_6_0.85_0.85            ]    972.1G ± 0%   1118.5G ±  1%  +15.06% (p=0.000 n=20)
BM_Sparse_Multi_2048_768_512_85_85_6/real_time                        [2048_768_512_6_0.85_0.85             ]    872.6G ± 0%    989.9G ±  1%  +13.44% (p=0.000 n=20)
BM_Sparse_Multi_2048_512_256_85_85_6/real_time                        [2048_512_256_6_0.85_0.85             ]    682.2G ± 0%    746.3G ±  0%   +9.41% (p=0.000 n=20)
geomean                                                                                                         555.5G         615.4G        +10.78%

Milan performance is unchanged
name                                                                                                            items/s      items/s    vs base
BM_Sparse_replicated_400_800_2560_85_0_6/real_time                    [copies: 6 sp_a: 0.85 sp_b: 0.00      ]    1.160T ± 0%   1.142T ± 1%  -1.52% (p=0.004 n=20)
BM_Sparse_replicated_400_2560_1024_85_0_6/real_time                   [copies: 6 sp_a: 0.85 sp_b: 0.00      ]    1.296T ± 0%   1.305T ± 1%  +0.70% (p=0.000 n=20)
BM_Sparse_replicated_400_1024_256_85_0_6/real_time                    [copies: 6 sp_a: 0.85 sp_b: 0.00      ]    915.0G ± 0%   923.0G ± 0%  +0.87% (p=0.000 n=20)
BM_Sparse_replicated_400_256_1_85_0_6/real_time                       [copies: 6 sp_a: 0.85 sp_b: 0.00      ]    4.440G ± 1%   4.422G ± 0%       ~ (p=0.165 n=20)
BM_Sparse_replicated_2048_1792_1024_85_0_6/real_time                  [copies: 6 sp_a: 0.85 sp_b: 0.00      ]    1.636T ± 0%   1.647T ± 0%  +0.68% (p=0.000 n=20)
BM_Sparse_replicated_2048_1024_768_85_0_6/real_time                   [copies: 6 sp_a: 0.85 sp_b: 0.00      ]    1.587T ± 0%   1.600T ± 0%  +0.82% (p=0.000 n=20)
BM_Sparse_replicated_2048_768_512_85_0_6/real_time                    [copies: 6 sp_a: 0.85 sp_b: 0.00      ]    1.348T ± 0%   1.366T ± 0%  +1.32% (p=0.000 n=20)
BM_Sparse_replicated_2048_512_256_85_0_6/real_time                    [copies: 6 sp_a: 0.85 sp_b: 0.00      ]    985.6G ± 0%   992.8G ± 0%  +0.73% (p=0.000 n=20)
BM_Sparse_Multi_1024_2140_4096_0_82_1/real_time                       [1024_2140_4096_1_0.00_0.82           ]    1.390T ± 0%   1.397T ± 0%  +0.54% (p=0.001 n=20)
BM_Sparse_Multi_1024_4096_2048_83_83_1/real_time                      [1024_4096_2048_1_0.83_0.83           ]    1.522T ± 0%   1.528T ± 0%  +0.40% (p=0.006 n=20)
BM_Sparse_Multi_400_800_2560_85_85_1/real_time                        [400_800_2560_1_0.85_0.85             ]    1.158T ± 1%   1.181T ± 1%  +2.03% (p=0.000 n=20)
BM_Sparse_Multi_400_2560_1024_85_85_1/real_time                       [400_2560_1024_1_0.85_0.85            ]    1.152T ± 0%   1.174T ± 0%  +1.89% (p=0.000 n=20)
BM_Sparse_Multi_400_1024_256_85_85_1/real_time                        [400_1024_256_1_0.85_0.85             ]    834.5G ± 0%   842.9G ± 0%  +1.01% (p=0.000 n=20)
BM_Sparse_Multi_400_256_1_85_85_1/real_time                           [400_256_1_1_0.85_0.85                ]    3.445G ± 1%   3.458G ± 1%       ~ (p=0.165 n=20)
BM_Sparse_Multi_2048_1792_1024_85_85_1/real_time                      [2048_1792_1024_1_0.85_0.85           ]    1.668T ± 0%   1.679T ± 0%  +0.62% (p=0.000 n=20)
BM_Sparse_Multi_2048_1024_768_85_85_1/real_time                       [2048_1024_768_1_0.85_0.85            ]    1.498T ± 0%   1.509T ± 0%  +0.74% (p=0.000 n=20)
BM_Sparse_Multi_2048_768_512_85_85_1/real_time                        [2048_768_512_1_0.85_0.85             ]    1.333T ± 0%   1.341T ± 0%  +0.58% (p=0.000 n=20)
BM_Sparse_Multi_2048_512_256_85_85_1/real_time                        [2048_512_256_1_0.85_0.85             ]    990.1G ± 0%   993.8G ± 0%  +0.38% (p=0.004 n=20)
BM_Sparse_Multi_2048_1792_1024_85_85_3/real_time                      [2048_1792_1024_3_0.85_0.85           ]    1.695T ± 0%   1.697T ± 0%       ~ (p=0.289 n=20)
BM_Sparse_Multi_2048_1024_768_85_85_3/real_time                       [2048_1024_768_3_0.85_0.85            ]    1.481T ± 0%   1.489T ± 1%  +0.48% (p=0.024 n=20)
BM_Sparse_Multi_2048_768_512_85_85_3/real_time                        [2048_768_512_3_0.85_0.85             ]    1.351T ± 0%   1.363T ± 0%  +0.89% (p=0.000 n=20)
BM_Sparse_Multi_2048_512_256_85_85_3/real_time                        [2048_512_256_3_0.85_0.85             ]    1.016T ± 0%   1.020T ± 0%  +0.38% (p=0.000 n=20)
BM_Sparse_Multi_2048_1792_1024_85_85_6/real_time                      [2048_1792_1024_6_0.85_0.85           ]    1.592T ± 1%   1.597T ± 1%       ~ (p=0.108 n=20)
BM_Sparse_Multi_2048_1024_768_85_85_6/real_time                       [2048_1024_768_6_0.85_0.85            ]    1.407T ± 1%   1.406T ± 0%       ~ (p=0.862 n=20)
BM_Sparse_Multi_2048_768_512_85_85_6/real_time                        [2048_768_512_6_0.85_0.85             ]    1.306T ± 0%   1.310T ± 0%  +0.31% (p=0.040 n=20)
BM_Sparse_Multi_2048_512_256_85_85_6/real_time                        [2048_512_256_6_0.85_0.85             ]    1.011T ± 0%   1.012T ± 0%       ~ (p=0.086 n=20)
geomean

PiperOrigin-RevId: 996709571
---
 tensorflow/core/kernels/BUILD                 |   7 +-
 tensorflow/core/kernels/sparse_matmul_op.cc   | 717 ++++++++++--------
 .../core/kernels/sparse_matmul_op_test.cc     | 226 ++----
 3 files changed, 492 insertions(+), 458 deletions(-)

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 276f1bbc644..12812fefc24 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -3690,7 +3690,10 @@ MATH_DEPS = [
 tf_kernel_library(
     name = "sparse_matmul_op",
     prefix = "sparse_matmul_op",
-    deps = MATH_DEPS + ["@xla//xla/tsl/framework/contraction:eigen_contraction_kernel"],
+    deps = MATH_DEPS + [
+        "@highway//:hwy",
+        "@xla//xla/tsl/framework/contraction:eigen_contraction_kernel",
+    ],
 )
 
 cc_library(
@@ -4231,6 +4234,7 @@ tf_cuda_cc_test(
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
         "//tensorflow/core:testlib",
+        "@highway//:hwy",
     ],
 )
 
@@ -7463,6 +7467,7 @@ filegroup(
             "decode_proto_op.cc",
             "encode_proto_op.cc",
             "sobol_op.cc",
+            "sparse_matmul_op.*",
             # Excluded due to experimental status:
             "debug_ops.*",
             "mutex_ops.*",
diff --git a/tensorflow/core/kernels/sparse_matmul_op.cc b/tensorflow/core/kernels/sparse_matmul_op.cc
index a5c80fc4b19..13927589853 100644
--- a/tensorflow/core/kernels/sparse_matmul_op.cc
+++ b/tensorflow/core/kernels/sparse_matmul_op.cc
@@ -15,10 +15,15 @@ limitations under the License.
 
 // See docs in ../ops/math_ops.cc.
 
+#ifndef TENSORFLOW_CORE_KERNELS_SPARSE_MATMUL_OP_COMMON_H_
+#define TENSORFLOW_CORE_KERNELS_SPARSE_MATMUL_OP_COMMON_H_
+
 #define EIGEN_USE_THREADS
 
 #include "tensorflow/core/kernels/sparse_matmul_op.h"
 
+#include <algorithm>
+#include <cstring>
 #include <map>
 #include <memory>
 #include <vector>
@@ -36,6 +41,7 @@ limitations under the License.
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/platform.h"
 #include "tensorflow/core/platform/thread_annotations.h"
 #include "tensorflow/core/platform/types.h"
 
@@ -290,35 +296,8 @@ void SparseSlice<T>::Clear() {
   data.clear();
 }
 
-using Packet = Eigen::internal::packet_traits<float>::type;
-const int kNumOperands = (sizeof(Packet) / sizeof(float));
-#define LOAD(x) Eigen::internal::pload<Packet>(x);
-#define EXPAND_BFLOAT_L(x, y) \
-  const auto y = Eigen::internal::pexpand_bf16_l<Packet>(x);
-#define EXPAND_BFLOAT_U(x, y) \
-  const auto y = Eigen::internal::pexpand_bf16_u<Packet>(x);
-#define STORE(x, y) Eigen::internal::pstore<float>(x, y);
-#define FMA(a, b, c, d) d = Eigen::internal::pmadd<Packet>(a, b, c);
-
 ALWAYS_INLINE float ConvertBfloat16ToFloat(const bfloat16* src) {
-  float out = 0;
-  auto tmp = reinterpret_cast<bfloat16*>(&out);
-#if __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
-  tmp[0] = *src;
-#else
-  tmp[1] = *src;
-#endif
-  return out;
-}
-
-ALWAYS_INLINE Packet ConvertFourBfloat16ToFloat(const bfloat16* src) {
-  return Eigen::internal::pload4bf16<Packet>(
-      reinterpret_cast<const float*>(src));
-}
-
-ALWAYS_INLINE Packet ConvertTwoBfloat16ToFloat(const bfloat16* src) {
-  return Eigen::internal::pload2bf16<Packet>(
-      reinterpret_cast<const float*>(src));
+  return static_cast<float>(*src);
 }
 
 ALWAYS_INLINE void ScalarMulAdd(const float a, const float** inp, float** out) {
@@ -334,6 +313,7 @@ ALWAYS_INLINE void ScalarMulAdd(const float a, const bfloat16** inp,
   ++*inp;
   ++*out;
 }
+
 ALWAYS_INLINE void ScalarMulAdd3Way(const float a1, const float a2,
                                     const float a3, const bfloat16** inp1,
                                     const bfloat16** inp2,
@@ -359,42 +339,64 @@ ALWAYS_INLINE void ScalarMulAdd3Way(const float a1, const float a2,
   ++*inp3;
 }
 
+template <typename TL_, typename TR_>
+struct TypePair {
+  using TL = TL_;
+  using TR = TR_;
+};
+
+}  // namespace
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_KERNELS_SPARSE_MATMUL_OP_COMMON_H_
+
+#undef HWY_TARGET_INCLUDE
+#if TSL_IS_IN_OSS
+#define HWY_TARGET_INCLUDE "tensorflow/core/kernels/sparse_matmul_op.cc"
+#else
+#define HWY_TARGET_INCLUDE \
+  "third_party/tensorflow/core/kernels/sparse_matmul_op.cc"
+#endif
+#include "hwy/foreach_target.h"  // from @highway  // IWYU pragma: keep
+#include "hwy/highway.h"  // from @highway
+
+HWY_BEFORE_NAMESPACE();
+namespace tensorflow {
+namespace {
+namespace HWY_NAMESPACE {
+
+namespace hn = hwy::HWY_NAMESPACE;
+
+// Use up to 512-bit float vectors (16 lanes on AVX-512 / Genoa, 8 on AVX2 /
+// SVE_256, 4 on SSE4 / NEON / SVE2_128 / scalable SVE/RVV, 1 on HWY_SCALAR).
+// On scalable targets (HWY_HAVE_SCALABLE == 1), cap to 4 lanes (128 bits, the
+// minimum SVE/RVV vector width) so kNumOperands matches hn::Lanes(D()) at
+// runtime.
+using D = hn::CappedTag<float, HWY_HAVE_SCALABLE ? 4 : 16>;
+using Packet = hn::Vec<D>;
+constexpr int kNumOperands = hn::MaxLanes(D());
+
 ALWAYS_INLINE void LoadSingleScalar(const bfloat16** data, Packet* l) {
-  auto tmp = ConvertBfloat16ToFloat(*data);
-  *l = Eigen::internal::pset1<Packet>(tmp);
+  const D d;
+  *l = hn::Set(d, ConvertBfloat16ToFloat(*data));
   ++*data;
 }
 
 ALWAYS_INLINE void LoadTwoScalars(const bfloat16** data, Packet* l1,
                                   Packet* l2) {
-  if (kNumOperands >= 2) {
-    auto tmp = ConvertTwoBfloat16ToFloat(*data);
-    *l1 = Eigen::internal::pbroadcast_first<Packet>(tmp);
-    *l2 = Eigen::internal::pbroadcast_second<Packet>(tmp);
-    *data += 2;
-  } else {
-    LoadSingleScalar(data, l1);
-    LoadSingleScalar(data, l2);
-  }
+  LoadSingleScalar(data, l1);
+  LoadSingleScalar(data, l2);
 }
 
 ALWAYS_INLINE void LoadFourScalars(const bfloat16** data, Packet* l1,
                                    Packet* l2, Packet* l3, Packet* l4) {
-  if (kNumOperands >= 4) {
-    auto tmp = ConvertFourBfloat16ToFloat(*data);
-    *l1 = Eigen::internal::pbroadcast_first<Packet>(tmp);
-    *l2 = Eigen::internal::pbroadcast_second<Packet>(tmp);
-    *l3 = Eigen::internal::pbroadcast_third<Packet>(tmp);
-    *l4 = Eigen::internal::pbroadcast_fourth<Packet>(tmp);
-    *data += 4;
-  } else {
-    LoadTwoScalars(data, l1, l2);
-    LoadTwoScalars(data, l3, l4);
-  }
+  LoadTwoScalars(data, l1, l2);
+  LoadTwoScalars(data, l3, l4);
 }
 
 ALWAYS_INLINE void LoadSingleScalar(const float** data, Packet* l) {
-  *l = Eigen::internal::pload1<Packet>(*data);
+  const D d;
+  *l = hn::Set(d, **data);
   ++(*data);
 }
 
@@ -424,19 +426,37 @@ ALWAYS_INLINE void LoadSixScalars(const T** data, Packet* l1, Packet* l2,
   LoadTwoScalars(data, l5, l6);
 }
 
+ALWAYS_INLINE void ExpandBfloat16(const bfloat16* inp, Packet* b_0,
+                                  Packet* b_1) {
+  const D d;
+#if HWY_TARGET != HWY_SCALAR && !HWY_IS_BIG_ENDIAN
+  if (kNumOperands >= 4) {
+    const hn::Repartition<uint16_t, D> d16;
+    const auto zero = hn::Zero(d16);
+    const auto b = hn::Load(d16, reinterpret_cast<const uint16_t*>(inp));
+    *b_0 = hn::BitCast(d, hn::InterleaveLower(d16, zero, b));
+    *b_1 = hn::BitCast(d, hn::InterleaveUpper(d16, zero, b));
+    return;
+  }
+#endif
+  const hn::Rebind<hwy::bfloat16_t, D> d_bf16;
+  const auto* bf_inp = reinterpret_cast<const hwy::bfloat16_t*>(inp);
+  *b_0 = hn::PromoteTo(d, hn::LoadU(d_bf16, bf_inp));
+  *b_1 = hn::PromoteTo(d, hn::LoadU(d_bf16, bf_inp + kNumOperands));
+}
+
 // Vectorized version of ScalarMulAdd.
 ALWAYS_INLINE void MulAdd(const Packet a, const bfloat16** binp, float** out) {
-  auto inp = reinterpret_cast<const float*>(*binp);
-  const auto b = LOAD(inp);
-  EXPAND_BFLOAT_L(b, b_0);
-  EXPAND_BFLOAT_U(b, b_1);
+  const D d;
+  Packet b_0, b_1;
+  ExpandBfloat16(*binp, &b_0, &b_1);
   *binp += 2 * kNumOperands;
-  auto c1 = LOAD(*out);
-  auto c2 = LOAD(*out + kNumOperands);
-  FMA(a, b_0, c1, c1);
-  FMA(a, b_1, c2, c2);
-  STORE(*out, c1);
-  STORE(*out + kNumOperands, c2);
+  Packet c1 = hn::Load(d, *out);
+  Packet c2 = hn::Load(d, *out + kNumOperands);
+  c1 = hn::MulAdd(a, b_0, c1);
+  c2 = hn::MulAdd(a, b_1, c2);
+  hn::Store(c1, d, *out);
+  hn::Store(c2, d, *out + kNumOperands);
   *out += 2 * kNumOperands;
 }
 
@@ -444,31 +464,24 @@ ALWAYS_INLINE void MulAdd(const Packet a, const bfloat16** binp, float** out) {
 ALWAYS_INLINE void MulAdd3Way(const Packet a1, const Packet a2, const Packet a3,
                               const bfloat16** binp1, const bfloat16** binp2,
                               const bfloat16** binp3, float** out) {
-  auto inp1 = reinterpret_cast<const float*>(*binp1);
-  auto inp2 = reinterpret_cast<const float*>(*binp2);
-  auto inp3 = reinterpret_cast<const float*>(*binp3);
-  auto c1 = LOAD(*out);
-  auto c2 = LOAD(*out + kNumOperands);
-  const auto b1 = LOAD(inp1);
-  EXPAND_BFLOAT_L(b1, b1_0);
-  EXPAND_BFLOAT_U(b1, b1_1);
+  const D d;
+  Packet c1 = hn::Load(d, *out);
+  Packet c2 = hn::Load(d, *out + kNumOperands);
+  Packet b1_0, b1_1, b2_0, b2_1, b3_0, b3_1;
+  ExpandBfloat16(*binp1, &b1_0, &b1_1);
   *binp1 += 2 * kNumOperands;
-  const auto b2 = LOAD(inp2);
-  EXPAND_BFLOAT_L(b2, b2_0);
-  EXPAND_BFLOAT_U(b2, b2_1);
+  ExpandBfloat16(*binp2, &b2_0, &b2_1);
   *binp2 += 2 * kNumOperands;
-  const auto b3 = LOAD(inp3);
-  EXPAND_BFLOAT_L(b3, b3_0);
-  EXPAND_BFLOAT_U(b3, b3_1);
+  ExpandBfloat16(*binp3, &b3_0, &b3_1);
   *binp3 += 2 * kNumOperands;
-  FMA(a1, b1_0, c1, c1);
-  FMA(a1, b1_1, c2, c2);
-  FMA(a2, b2_0, c1, c1);
-  FMA(a2, b2_1, c2, c2);
-  FMA(a3, b3_0, c1, c1);
-  FMA(a3, b3_1, c2, c2);
-  STORE(*out, c1);
-  STORE(*out + kNumOperands, c2);
+  c1 = hn::MulAdd(a1, b1_0, c1);
+  c2 = hn::MulAdd(a1, b1_1, c2);
+  c1 = hn::MulAdd(a2, b2_0, c1);
+  c2 = hn::MulAdd(a2, b2_1, c2);
+  c1 = hn::MulAdd(a3, b3_0, c1);
+  c2 = hn::MulAdd(a3, b3_1, c2);
+  hn::Store(c1, d, *out);
+  hn::Store(c2, d, *out + kNumOperands);
   *out += 2 * kNumOperands;
 }
 
@@ -477,50 +490,37 @@ ALWAYS_INLINE void TwoMulAdd3Way(const Packet a1, const Packet a2,
                                  const Packet a3, const bfloat16** binp1,
                                  const bfloat16** binp2, const bfloat16** binp3,
                                  float** out) {
-  auto inp1 = reinterpret_cast<const float*>(*binp1);
-  auto inp2 = reinterpret_cast<const float*>(*binp2);
-  auto inp3 = reinterpret_cast<const float*>(*binp3);
-  auto c1 = LOAD(*out);
-  auto c2 = LOAD(*out + kNumOperands);
-  const auto b1 = LOAD(inp1);
-  const auto b2 = LOAD(inp2);
-  const auto b3 = LOAD(inp3);
+  const D d;
+  Packet c1 = hn::Load(d, *out);
+  Packet c2 = hn::Load(d, *out + kNumOperands);
+  Packet b1_0, b1_1, b2_0, b2_1, b3_0, b3_1;
+  ExpandBfloat16(*binp1, &b1_0, &b1_1);
+  ExpandBfloat16(*binp2, &b2_0, &b2_1);
+  ExpandBfloat16(*binp3, &b3_0, &b3_1);
 
-  EXPAND_BFLOAT_L(b1, b1_0);
-  EXPAND_BFLOAT_U(b1, b1_1);
-  EXPAND_BFLOAT_L(b2, b2_0);
-  EXPAND_BFLOAT_U(b2, b2_1);
-  EXPAND_BFLOAT_L(b3, b3_0);
-  EXPAND_BFLOAT_U(b3, b3_1);
-  auto c3 = LOAD(*out + 2 * kNumOperands);
-  auto c4 = LOAD(*out + 3 * kNumOperands);
-  const auto b4 = LOAD(inp1 + kNumOperands);
-  const auto b5 = LOAD(inp2 + kNumOperands);
-  const auto b6 = LOAD(inp3 + kNumOperands);
+  Packet c3 = hn::Load(d, *out + 2 * kNumOperands);
+  Packet c4 = hn::Load(d, *out + 3 * kNumOperands);
+  Packet b4_0, b4_1, b5_0, b5_1, b6_0, b6_1;
+  ExpandBfloat16(*binp1 + 2 * kNumOperands, &b4_0, &b4_1);
+  ExpandBfloat16(*binp2 + 2 * kNumOperands, &b5_0, &b5_1);
+  ExpandBfloat16(*binp3 + 2 * kNumOperands, &b6_0, &b6_1);
 
-  EXPAND_BFLOAT_L(b4, b4_0);
-  EXPAND_BFLOAT_U(b4, b4_1);
-  EXPAND_BFLOAT_L(b5, b5_0);
-  EXPAND_BFLOAT_U(b5, b5_1);
-  EXPAND_BFLOAT_L(b6, b6_0);
-  EXPAND_BFLOAT_U(b6, b6_1);
-
-  FMA(a1, b1_0, c1, c1);
-  FMA(a1, b1_1, c2, c2);
-  FMA(a1, b4_0, c3, c3);
-  FMA(a1, b4_1, c4, c4);
-  FMA(a2, b2_0, c1, c1);
-  FMA(a2, b2_1, c2, c2);
-  FMA(a2, b5_0, c3, c3);
-  FMA(a2, b5_1, c4, c4);
-  FMA(a3, b3_0, c1, c1);
-  FMA(a3, b3_1, c2, c2);
-  FMA(a3, b6_0, c3, c3);
-  FMA(a3, b6_1, c4, c4);
-  STORE(*out, c1);
-  STORE(*out + kNumOperands, c2);
-  STORE(*out + 2 * kNumOperands, c3);
-  STORE(*out + 3 * kNumOperands, c4);
+  c1 = hn::MulAdd(a1, b1_0, c1);
+  c2 = hn::MulAdd(a1, b1_1, c2);
+  c3 = hn::MulAdd(a1, b4_0, c3);
+  c4 = hn::MulAdd(a1, b4_1, c4);
+  c1 = hn::MulAdd(a2, b2_0, c1);
+  c2 = hn::MulAdd(a2, b2_1, c2);
+  c3 = hn::MulAdd(a2, b5_0, c3);
+  c4 = hn::MulAdd(a2, b5_1, c4);
+  c1 = hn::MulAdd(a3, b3_0, c1);
+  c2 = hn::MulAdd(a3, b3_1, c2);
+  c3 = hn::MulAdd(a3, b6_0, c3);
+  c4 = hn::MulAdd(a3, b6_1, c4);
+  hn::Store(c1, d, *out);
+  hn::Store(c2, d, *out + kNumOperands);
+  hn::Store(c3, d, *out + 2 * kNumOperands);
+  hn::Store(c4, d, *out + 3 * kNumOperands);
   *out += 4 * kNumOperands;
   *binp1 += 4 * kNumOperands;
   *binp2 += 4 * kNumOperands;
@@ -540,11 +540,12 @@ ALWAYS_INLINE void MulAdd3Way128(const Packet a1, const Packet a2,
 
 // Vectorized version of ScalarMulAdd
 ALWAYS_INLINE void MulAdd(const Packet a, const float** inp, float** out) {
-  const auto b = LOAD(*inp);
+  const D d;
+  const Packet b = hn::Load(d, *inp);
   *inp += kNumOperands;
-  auto c = LOAD(*out);
-  FMA(a, b, c, c);
-  STORE(*out, c);
+  Packet c = hn::Load(d, *out);
+  c = hn::MulAdd(a, b, c);
+  hn::Store(c, d, *out);
   *out += kNumOperands;
 }
 
@@ -552,17 +553,18 @@ ALWAYS_INLINE void MulAdd(const Packet a, const float** inp, float** out) {
 ALWAYS_INLINE void MulAdd3Way(const Packet a1, const Packet a2, const Packet a3,
                               const float** inp1, const float** inp2,
                               const float** inp3, float** out) {
-  auto c = LOAD(*out);
-  const auto b1 = LOAD(*inp1);
+  const D d;
+  Packet c = hn::Load(d, *out);
+  const Packet b1 = hn::Load(d, *inp1);
   *inp1 += kNumOperands;
-  const auto b2 = LOAD(*inp2);
+  const Packet b2 = hn::Load(d, *inp2);
   *inp2 += kNumOperands;
-  const auto b3 = LOAD(*inp3);
+  const Packet b3 = hn::Load(d, *inp3);
   *inp3 += kNumOperands;
-  FMA(a1, b1, c, c);
-  FMA(a2, b2, c, c);
-  FMA(a3, b3, c, c);
-  STORE(*out, c);
+  c = hn::MulAdd(a1, b1, c);
+  c = hn::MulAdd(a2, b2, c);
+  c = hn::MulAdd(a3, b3, c);
+  hn::Store(c, d, *out);
   *out += kNumOperands;
 }
 
@@ -571,24 +573,25 @@ ALWAYS_INLINE void TwoMulAdd3Way(const Packet a1, const Packet a2,
                                  const Packet a3, const float** inp1,
                                  const float** inp2, const float** inp3,
                                  float** out) {
-  auto c1 = LOAD(*out);
-  const auto b1 = LOAD(*inp1);
-  const auto b2 = LOAD(*inp2);
-  const auto b3 = LOAD(*inp3);
+  const D d;
+  Packet c1 = hn::Load(d, *out);
+  const Packet b1 = hn::Load(d, *inp1);
+  const Packet b2 = hn::Load(d, *inp2);
+  const Packet b3 = hn::Load(d, *inp3);
 
-  auto c2 = LOAD(*out + kNumOperands);
-  const auto b4 = LOAD(*inp1 + kNumOperands);
-  const auto b5 = LOAD(*inp2 + kNumOperands);
-  const auto b6 = LOAD(*inp3 + kNumOperands);
+  Packet c2 = hn::Load(d, *out + kNumOperands);
+  const Packet b4 = hn::Load(d, *inp1 + kNumOperands);
+  const Packet b5 = hn::Load(d, *inp2 + kNumOperands);
+  const Packet b6 = hn::Load(d, *inp3 + kNumOperands);
 
-  FMA(a1, b1, c1, c1);
-  FMA(a1, b4, c2, c2);
-  FMA(a2, b2, c1, c1);
-  FMA(a2, b5, c2, c2);
-  FMA(a3, b3, c1, c1);
-  FMA(a3, b6, c2, c2);
-  STORE(*out, c1);
-  STORE(*out + kNumOperands, c2);
+  c1 = hn::MulAdd(a1, b1, c1);
+  c2 = hn::MulAdd(a1, b4, c2);
+  c1 = hn::MulAdd(a2, b2, c1);
+  c2 = hn::MulAdd(a2, b5, c2);
+  c1 = hn::MulAdd(a3, b3, c1);
+  c2 = hn::MulAdd(a3, b6, c2);
+  hn::Store(c1, d, *out);
+  hn::Store(c2, d, *out + kNumOperands);
   *out += 2 * kNumOperands;
   *inp1 += 2 * kNumOperands;
   *inp2 += 2 * kNumOperands;
@@ -600,8 +603,51 @@ ALWAYS_INLINE void FourMulAdd3Way(const Packet a1, const Packet a2,
                                   const Packet a3, const float** inp1,
                                   const float** inp2, const float** inp3,
                                   float** out) {
-  TwoMulAdd3Way(a1, a2, a3, inp1, inp2, inp3, out);
-  TwoMulAdd3Way(a1, a2, a3, inp1, inp2, inp3, out);
+  if constexpr (kNumOperands < 16) {
+    TwoMulAdd3Way(a1, a2, a3, inp1, inp2, inp3, out);
+    TwoMulAdd3Way(a1, a2, a3, inp1, inp2, inp3, out);
+    return;
+  }
+  const D d;
+  Packet c1 = hn::Load(d, *out);
+  Packet c2 = hn::Load(d, *out + kNumOperands);
+  Packet c3 = hn::Load(d, *out + 2 * kNumOperands);
+  Packet c4 = hn::Load(d, *out + 3 * kNumOperands);
+
+  const Packet b1_0 = hn::Load(d, *inp1);
+  const Packet b1_1 = hn::Load(d, *inp1 + kNumOperands);
+  const Packet b2_0 = hn::Load(d, *inp2);
+  const Packet b2_1 = hn::Load(d, *inp2 + kNumOperands);
+  const Packet b3_0 = hn::Load(d, *inp3);
+  const Packet b3_1 = hn::Load(d, *inp3 + kNumOperands);
+
+  const Packet b4_0 = hn::Load(d, *inp1 + 2 * kNumOperands);
+  const Packet b4_1 = hn::Load(d, *inp1 + 3 * kNumOperands);
+  const Packet b5_0 = hn::Load(d, *inp2 + 2 * kNumOperands);
+  const Packet b5_1 = hn::Load(d, *inp2 + 3 * kNumOperands);
+  const Packet b6_0 = hn::Load(d, *inp3 + 2 * kNumOperands);
+  const Packet b6_1 = hn::Load(d, *inp3 + 3 * kNumOperands);
+
+  c1 = hn::MulAdd(a1, b1_0, c1);
+  c2 = hn::MulAdd(a1, b1_1, c2);
+  c3 = hn::MulAdd(a1, b4_0, c3);
+  c4 = hn::MulAdd(a1, b4_1, c4);
+  c1 = hn::MulAdd(a2, b2_0, c1);
+  c2 = hn::MulAdd(a2, b2_1, c2);
+  c3 = hn::MulAdd(a2, b5_0, c3);
+  c4 = hn::MulAdd(a2, b5_1, c4);
+  c1 = hn::MulAdd(a3, b3_0, c1);
+  c2 = hn::MulAdd(a3, b3_1, c2);
+  c3 = hn::MulAdd(a3, b6_0, c3);
+  c4 = hn::MulAdd(a3, b6_1, c4);
+  hn::Store(c1, d, *out);
+  hn::Store(c2, d, *out + kNumOperands);
+  hn::Store(c3, d, *out + 2 * kNumOperands);
+  hn::Store(c4, d, *out + 3 * kNumOperands);
+  *out += 4 * kNumOperands;
+  *inp1 += 4 * kNumOperands;
+  *inp2 += 4 * kNumOperands;
+  *inp3 += 4 * kNumOperands;
 }
 
 // Apply MulAdd3Way on 128 operands.
@@ -609,7 +655,10 @@ ALWAYS_INLINE void MulAdd3Way128(const Packet a1, const Packet a2,
                                  const Packet a3, const float** inp1,
                                  const float** inp2, const float** inp3,
                                  float** out) {
-  if (kNumOperands == 8) {
+  if (kNumOperands == 16) {
+    FourMulAdd3Way(a1, a2, a3, inp1, inp2, inp3, out);
+    FourMulAdd3Way(a1, a2, a3, inp1, inp2, inp3, out);
+  } else if (kNumOperands == 8) {
     FourMulAdd3Way(a1, a2, a3, inp1, inp2, inp3, out);
     FourMulAdd3Way(a1, a2, a3, inp1, inp2, inp3, out);
     FourMulAdd3Way(a1, a2, a3, inp1, inp2, inp3, out);
@@ -624,6 +673,7 @@ ALWAYS_INLINE void MulAdd3Way128(const Packet a1, const Packet a2,
     }
   }
 }
+
 // Computes product of "left_slices" with "num_cols" columns of "right", and
 // stores the output in *"output".
 // Note that left_slices is a list of SparseSlices, which are conceptually
@@ -635,12 +685,13 @@ inline void GEPP(
     const std::vector<SparseSlice<TL>*>& left_slices,
     const Eigen::TensorMap<Eigen::Tensor<const TR, 2, Eigen::RowMajor>,
                            Eigen::Aligned>& right,
-    const int num_cols, Matrix* output) {
+    const int num_cols, MatrixMap* output) {
   const int cols = (Cols == -1) ? num_cols : Cols;
   DCHECK_EQ(num_cols, cols);
   const int right_num_cols = right.dimension(1);
   const int output_num_cols = output->dimension(1);
-  static const int kNumOperandsR = kNumOperands * sizeof(float) / sizeof(TR);
+  static constexpr int kNumOperandsR =
+      kNumOperands * sizeof(float) / sizeof(TR);
   const int cols_mod = cols % kNumOperandsR;
   int k_offset = 0;
   // Pre-compute pointers for output matrix.
@@ -691,12 +742,12 @@ inline void GEPP(
             MulAdd3Way(nl1, nl2, nl3, &nr1, &nr2, &nr3, &nout);
           }
 
-          const float sl1 = Eigen::internal::pfirst<Packet>(l1);
-          const float sl2 = Eigen::internal::pfirst<Packet>(l2);
-          const float sl3 = Eigen::internal::pfirst<Packet>(l3);
-          const float nsl1 = Eigen::internal::pfirst<Packet>(nl1);
-          const float nsl2 = Eigen::internal::pfirst<Packet>(nl2);
-          const float nsl3 = Eigen::internal::pfirst<Packet>(nl3);
+          const float sl1 = hn::GetLane(l1);
+          const float sl2 = hn::GetLane(l2);
+          const float sl3 = hn::GetLane(l3);
+          const float nsl1 = hn::GetLane(nl1);
+          const float nsl2 = hn::GetLane(nl2);
+          const float nsl3 = hn::GetLane(nl3);
           for (int k = 0; k < cols_mod; ++k) {
             ScalarMulAdd3Way(sl1, sl2, sl3, &r1, &r2, &r3, &out);
             ScalarMulAdd3Way(nsl1, nsl2, nsl3, &nr1, &nr2, &nr3, &nout);
@@ -718,9 +769,9 @@ inline void GEPP(
           for (int n = 0; n < cols / kNumOperandsR; ++n) {
             MulAdd3Way(l1, l2, l3, &r1, &r2, &r3, &out);
           }
-          const float sl1 = Eigen::internal::pfirst<Packet>(l1);
-          const float sl2 = Eigen::internal::pfirst<Packet>(l2);
-          const float sl3 = Eigen::internal::pfirst<Packet>(l3);
+          const float sl1 = hn::GetLane(l1);
+          const float sl2 = hn::GetLane(l2);
+          const float sl3 = hn::GetLane(l3);
           for (int k = 0; k < cols_mod; ++k) {
             ScalarMulAdd3Way(sl1, sl2, sl3, &r1, &r2, &r3, &out);
           }
@@ -754,10 +805,10 @@ inline void GEPP(
           MulAdd(n3l, &n3r, &n3out);
         }
 
-        const float sl1 = Eigen::internal::pfirst<Packet>(l);
-        const float sl2 = Eigen::internal::pfirst<Packet>(nl);
-        const float sl3 = Eigen::internal::pfirst<Packet>(n2l);
-        const float sl4 = Eigen::internal::pfirst<Packet>(n3l);
+        const float sl1 = hn::GetLane(l);
+        const float sl2 = hn::GetLane(nl);
+        const float sl3 = hn::GetLane(n2l);
+        const float sl4 = hn::GetLane(n3l);
         for (int k = 0; k < cols_mod; ++k) {
           ScalarMulAdd(sl1, &r, &out);
           ScalarMulAdd(sl2, &nr, &nout);
@@ -774,7 +825,7 @@ inline void GEPP(
         for (int n = 0; n < cols / kNumOperandsR; ++n) {
           MulAdd(l, &r, &out);
         }
-        const float sl = Eigen::internal::pfirst<Packet>(l);
+        const float sl = hn::GetLane(l);
         for (int k = 0; k < cols_mod; ++k) {
           ScalarMulAdd(sl, &r, &out);
         }
@@ -786,13 +837,172 @@ inline void GEPP(
   }
 }
 
-#undef LOAD
-#undef EXPAND_BFLOAT_L
-#undef EXPAND_BFLOAT_U
-#undef STORE
-#undef FMA
+template <int NUM_ELEM = -1>
+ALWAYS_INLINE void CopyAndMayBeInterleaveBfloat16(void* bdst, const void* bsrc,
+                                                  int num_elements) {
+#if HWY_TARGET != HWY_SCALAR && !HWY_IS_BIG_ENDIAN
+  if constexpr (kNumOperands >= 8) {
+    static constexpr int kStep =
+        kNumOperands * sizeof(float) / sizeof(bfloat16);
+    const int num = (NUM_ELEM == -1) ? num_elements : NUM_ELEM;
+    DCHECK_EQ(num, num_elements);
+    const hn::Repartition<uint64_t, D> d64;
+    const uint64_t* src = reinterpret_cast<const uint64_t*>(bsrc);
+    uint64_t* dst = reinterpret_cast<uint64_t*>(bdst);
+    if constexpr (kNumOperands == 16) {
+      alignas(64) static constexpr uint64_t kIdx[8] = {0, 4, 1, 5, 2, 6, 3, 7};
+      const auto indices = hn::SetTableIndices(d64, kIdx);
+      for (int index = 0; index + kStep <= num; index += kStep) {
+        auto in = hn::LoadU(d64, src);
+        in = hn::TableLookupLanes(in, indices);
+        hn::Store(in, d64, dst);
+        src += kNumOperands / 2;
+        dst += kNumOperands / 2;
+      }
+    } else if constexpr (kNumOperands == 8) {
+      for (int index = 0; index + kStep <= num; index += kStep) {
+        auto in = hn::LoadU(d64, src);
+        in = hn::Per4LaneBlockShuffle<3, 1, 2, 0>(in);
+        hn::Store(in, d64, dst);
+        src += kNumOperands / 2;
+        dst += kNumOperands / 2;
+      }
+    }
+    if (num % kStep != 0) {
+      memcpy(reinterpret_cast<void*>(dst), reinterpret_cast<const void*>(src),
+             (num % kStep) * sizeof(bfloat16));
+    }
+    return;
+  }
+#endif
+  memcpy(bdst, bsrc, num_elements * sizeof(bfloat16));
+}
 
+template <typename T>
+ALWAYS_INLINE void CopyAndMayBeInterleave(void* dst, const void* src,
+                                          int num_elements) {
+  if (std::is_same<T, float>::value || kNumOperands < 8) {
+    memcpy(dst, src, num_elements * sizeof(T));
+  } else if (std::is_same<T, bfloat16>::value) {
+    if (num_elements == N) {
+      CopyAndMayBeInterleaveBfloat16<N>(dst, src, num_elements);
+    } else {
+      CopyAndMayBeInterleaveBfloat16<-1>(dst, src, num_elements);
+    }
+  } else {
+    LOG(FATAL) << "Unsupported type";
+  }
+}
+
+template <typename TR>
+void ShuffleMatrixWorkImpl(const BasicMatrixMap<const TR>& mat,
+                           int slice_row_start, int slice_num_rows,
+                           int slice_col_start, int slice_num_cols, int s,
+                           int e, BasicMatrixMap<TR>* buffer) {
+  const int row_start = s % slice_num_rows + slice_row_start;
+  const int col_start = s / slice_num_rows * N + slice_col_start;
+  auto* out_start = &(*buffer)(s, 0);
+  const auto* input_start = &mat(row_start, col_start);
+  const auto* input_end = &mat(slice_row_start + slice_num_rows - 1,
+                               slice_col_start + slice_num_cols - 1);
+  const int mat_num_cols = mat.dimension(1);
+  const int row_slice_size = slice_num_rows * mat_num_cols;
+
+  const int aligned_end = slice_num_cols / N * slice_num_rows;
+  const int e1 = std::min(e, aligned_end);
+  while (s < e1) {
+    CopyAndMayBeInterleave<TR>(out_start, input_start, N);
+    out_start += N;
+    input_start += mat_num_cols;
+    if (input_start > input_end) {
+      input_start = input_start - row_slice_size + N;
+    }
+    ++s;
+  }
+  int s1 = std::max(s, aligned_end);
+  const int copy_num_cols = slice_num_cols % N;
+  while (s1 < e) {
+    CopyAndMayBeInterleave<TR>(out_start, input_start, copy_num_cols);
+    out_start += N;
+    input_start += mat_num_cols;
+    ++s1;
+  }
+}
+
+template <typename Pair>
+void ComputeOutputBlockImpl(
+    const std::vector<SparseSlice<typename Pair::TL>*>& left,
+    const BasicMatrixMap<const typename Pair::TR>& right, int num_cols,
+    int output_row_offset, int output_col_offset, bool assign,
+    bool transpose_output, MatrixMap* output) {
+  using TL = typename Pair::TL;
+  using TR = typename Pair::TR;
+  int num_rows = left[0]->num_rows;
+  const int rhs_num_cols = right.dimension(1);
+  DCHECK_LE(num_cols, rhs_num_cols);
+#if HWY_ARCH_X86 && !defined(__ANDROID__)
+  alignas(64) float out_data[M * N];
+#else
+  // Mobile/ARM builds enforce -Werror,-Wframe-larger-than=25344, which is
+  // smaller than sizeof(float) * M * N (32 KiB).
+  float* out_data =
+      static_cast<float*>(Eigen::internal::handmade_aligned_malloc(
+          sizeof(float) * num_rows * rhs_num_cols, 64));
+  std::unique_ptr<float, void (*)(void*)> out_heap(
+      out_data, Eigen::internal::handmade_aligned_free);
+#endif
+  memset(out_data, 0, sizeof(float) * num_rows * rhs_num_cols);
+  MatrixMap out(out_data, num_rows, rhs_num_cols);
+  if (num_cols == N) {
+    GEPP<TL, TR, N>(left, right, num_cols, &out);
+  } else {
+    GEPP<TL, TR, -1>(left, right, num_cols, &out);
+  }
+  if (!assign) {
+    const D d;
+    if (transpose_output) {
+      for (int i = 0; i < num_rows; ++i) {
+        for (int j = 0; j < num_cols; ++j) {
+          (*output)(output_col_offset + j, output_row_offset + i) += out(i, j);
+        }
+      }
+    } else {
+      for (int i = 0; i < num_rows; ++i) {
+        float* dst = &(*output)(output_row_offset + i, output_col_offset);
+        const float* src = &out(i, 0);
+        int j = 0;
+        for (; j + kNumOperands <= num_cols; j += kNumOperands) {
+          hn::StoreU(hn::Add(hn::LoadU(d, dst + j), hn::Load(d, src + j)), d,
+                     dst + j);
+        }
+        for (; j < num_cols; ++j) {
+          dst[j] += src[j];
+        }
+      }
+    }
+  } else {
+    if (transpose_output) {
+      for (int i = 0; i < num_rows; ++i) {
+        for (int j = 0; j < num_cols; ++j) {
+          (*output)(output_col_offset + j, output_row_offset + i) = out(i, j);
+        }
+      }
+    } else {
+      for (int i = 0; i < num_rows; ++i) {
+        memcpy(&(*output)(output_row_offset + i, output_col_offset), &out(i, 0),
+               num_cols * sizeof(float));
+      }
+    }
+  }
+}
+
+}  // namespace HWY_NAMESPACE
 }  // namespace
+}  // namespace tensorflow
+HWY_AFTER_NAMESPACE();
+
+#if HWY_ONCE
+namespace tensorflow {
 
 template <typename TL, typename TR>
 class SparseMatMul {
@@ -839,18 +1049,18 @@ class SparseMatMul {
   static inline std::unique_ptr<BlockingCounter> CreateDenseSlices(
       const ConstMatrixMapR& mat, int row_start, int num_rows, int col_start,
       int num_cols, const DeviceBase::CpuWorkerThreads* thread_pool,
-      MatrixR* buffer, std::vector<ConstMatrixMapR*>* slices);
+      MatrixMapR* buffer, std::vector<ConstMatrixMapR*>* slices);
 
   // Helper function for CreateDenseSlices to move the data around. It returns a
   // BlockingCounter which should be used to wait for the shuffle operations to
   // complete.
   static inline BlockingCounter* ShuffleMatrix(
       const ConstMatrixMapR& mat, int slice_row_start, int slice_num_rows,
-      int slice_col_start, int slice_num_cols, const int N,
-      const DeviceBase::CpuWorkerThreads* thread_pool, MatrixR* buffer);
+      int slice_col_start, int slice_num_cols,
+      const DeviceBase::CpuWorkerThreads* thread_pool, MatrixMapR* buffer);
 
   // Helper function for CreateDenseSlices to create slices.
-  static inline void SliceMatrix(const MatrixR& mat, const int num_rows,
+  static inline void SliceMatrix(const MatrixMapR& mat, const int num_rows,
                                  const int num_slices,
                                  std::vector<ConstMatrixMapR*>* slices);
 
@@ -1021,49 +1231,10 @@ inline void SparseMatMul<TL, TR>::ComputeOutputBlock(
     const typename SparseMatMul<TL, TR>::ConstMatrixMapR& right, int num_cols,
     int output_row_offset, int output_col_offset, bool assign,
     bool transpose_output, MatrixMap* output) {
-  const auto perm = dsizes_10();
-  int num_rows = left[0]->num_rows;
-  const int rhs_num_cols = right.dimension(1);
-  DCHECK_LE(num_cols, rhs_num_cols);
-  Matrix out(num_rows, rhs_num_cols);
-  out.setZero();
-  if (num_cols == N) {
-    GEPP<TL, TR, N>(left, right, num_cols, &out);
-  } else {
-    GEPP<TL, TR, -1>(left, right, num_cols, &out);
-  }
-  if (!assign) {
-    const DSizes begin(output_row_offset, output_col_offset);
-    const DSizes sizes(num_rows, num_cols);
-    if (transpose_output) {
-      if (num_cols == rhs_num_cols) {
-        output->shuffle(perm).slice(begin, sizes) += out;
-      } else {
-        const auto zero = dsizes_00();
-        output->shuffle(perm).slice(begin, sizes) += out.slice(zero, sizes);
-      }
-    } else {
-      if (num_cols == rhs_num_cols) {
-        output->slice(begin, sizes) += out;
-      } else {
-        const auto zero = dsizes_00();
-        output->slice(begin, sizes) += out.slice(zero, sizes);
-      }
-    }
-  } else {
-    std::unique_ptr<Matrix> out_tr;
-    if (transpose_output) {
-      out_tr.reset(new Matrix(rhs_num_cols, num_rows));
-      *out_tr = out.shuffle(perm);
-      std::swap(output_row_offset, output_col_offset);
-      std::swap(num_rows, num_cols);
-    }
-    const Matrix& final_out = transpose_output ? *out_tr : out;
-    for (int i = 0; i < num_rows; ++i) {
-      memcpy(&(*output)(output_row_offset + i, output_col_offset),
-             &final_out(i, 0), num_cols * sizeof(float));
-    }
-  }
+  using Pair = TypePair<TL, TR>;
+  HWY_EXPORT_AND_DYNAMIC_DISPATCH_T(ComputeOutputBlockImpl<Pair>)(
+      left, right, num_cols, output_row_offset, output_col_offset, assign,
+      transpose_output, output);
 }
 
 template <typename TL, typename TR>
@@ -1118,95 +1289,24 @@ SparseMatMul<TL, TR>::CreateSparseSlices(
   }
   return std::unique_ptr<BlockingCounter>(counter);
 }
-#define LOAD(x) Eigen::internal::ploadu<Packet>((x));
-#define INTERLEAVE(x) Eigen::internal::pinterleave4x64<Packet>(x);
-#define STORE(x, y) Eigen::internal::pstoreu<float>(x, y);
-
-template <int NUM_ELEM = -1>
-ALWAYS_INLINE void CopyAndMayBeInterleaveBfloat16(void* bdst, const void* bsrc,
-                                                  int num_elements) {
-  DCHECK_GE(kNumOperands, 8);
-  static const int kStep = kNumOperands * sizeof(float) / sizeof(bfloat16);
-  const int num = (NUM_ELEM == -1) ? num_elements : NUM_ELEM;
-  DCHECK_EQ(num, num_elements);
-  const float* src = reinterpret_cast<const float*>(bsrc);
-  float* dst = reinterpret_cast<float*>(bdst);
-  for (int index = 0; index + kStep <= num; index += kStep) {
-    auto in = LOAD(src);
-    auto tmp = INTERLEAVE(in);
-    STORE(dst, tmp);
-    src += kNumOperands;
-    dst += kNumOperands;
-  }
-  if (num % kStep != 0) {
-    memcpy(reinterpret_cast<void*>(dst), reinterpret_cast<const void*>(src),
-           (num % kStep) * sizeof(bfloat16));
-  }
-}
-
-template <typename T>
-ALWAYS_INLINE void CopyAndMayBeInterleave(void* dst, const void* src,
-                                          int num_elements) {
-  if (std::is_same<T, float>::value || kNumOperands < 8) {
-    memcpy(dst, src, num_elements * sizeof(T));
-  } else if (std::is_same<T, bfloat16>::value) {
-    if (num_elements == N) {
-      CopyAndMayBeInterleaveBfloat16<N>(dst, src, num_elements);
-    } else {
-      CopyAndMayBeInterleaveBfloat16<-1>(dst, src, num_elements);
-    }
-  } else {
-    LOG(FATAL) << "Unsupported type";
-  }
-}
-
-#undef LOAD
-#undef Interleave
-#undef Store
 
 template <typename TL, typename TR>
 inline BlockingCounter* SparseMatMul<TL, TR>::ShuffleMatrix(
     const typename SparseMatMul<TL, TR>::ConstMatrixMapR& mat,
     int slice_row_start, int slice_num_rows, int slice_col_start,
-    int slice_num_cols, const int N,
-    const DeviceBase::CpuWorkerThreads* thread_pool, MatrixR* buffer) {
+    int slice_num_cols, const DeviceBase::CpuWorkerThreads* thread_pool,
+    MatrixMapR* buffer) {
   DCHECK_EQ(N % 2, 0);
-  DCHECK_LE(kNumOperands * sizeof(float) / sizeof(TR), N);
   // Note(nikhilsarda): This heuristic is optimal in benchmarks as of
   // Jan 21, 2020.
   int num_threads = std::min(thread_pool->num_threads, 8);
   BlockingCounter* counter = new BlockingCounter(num_threads);
   DCHECK_EQ(N, buffer->dimension(1));
   auto shuffle_work = [&mat, slice_row_start, slice_num_rows, slice_col_start,
-                       slice_num_cols, N, buffer, counter](int s, int e) {
-    const int row_start = s % slice_num_rows + slice_row_start;
-    const int col_start = s / slice_num_rows * N + slice_col_start;
-    auto* out_start = &(*buffer)(s, 0);
-    const auto* input_start = &mat(row_start, col_start);
-    const auto* input_end = &mat(slice_row_start + slice_num_rows - 1,
-                                 slice_col_start + slice_num_cols - 1);
-    const int mat_num_cols = mat.dimension(1);
-    const int row_slice_size = slice_num_rows * mat_num_cols;
-
-    const int aligned_end = slice_num_cols / N * slice_num_rows;
-    const int e1 = std::min(e, aligned_end);
-    while (s < e1) {
-      CopyAndMayBeInterleave<TR>(out_start, input_start, N);
-      out_start += N;
-      input_start += mat_num_cols;
-      if (input_start > input_end) {
-        input_start = input_start - row_slice_size + N;
-      }
-      ++s;
-    }
-    int s1 = std::max(s, aligned_end);
-    const int copy_num_cols = slice_num_cols % N;
-    while (s1 < e) {
-      CopyAndMayBeInterleave<TR>(out_start, input_start, copy_num_cols);
-      out_start += N;
-      input_start += mat_num_cols;
-      ++s1;
-    }
+                       slice_num_cols, buffer, counter](int s, int e) {
+    HWY_EXPORT_AND_DYNAMIC_DISPATCH_T(ShuffleMatrixWorkImpl<TR>)(
+        mat, slice_row_start, slice_num_rows, slice_col_start, slice_num_cols,
+        s, e, buffer);
     if (counter) counter->DecrementCount();
   };
 
@@ -1225,7 +1325,7 @@ inline BlockingCounter* SparseMatMul<TL, TR>::ShuffleMatrix(
 
 template <typename TL, typename TR>
 inline void SparseMatMul<TL, TR>::SliceMatrix(
-    const MatrixR& mat, const int num_rows, const int num_slices,
+    const MatrixMapR& mat, const int num_rows, const int num_slices,
     std::vector<typename SparseMatMul<TL, TR>::ConstMatrixMapR*>* slices) {
   slices->resize(num_slices);
   DSizes d(num_rows, mat.dimension(1));
@@ -1239,10 +1339,10 @@ template <typename TL, typename TR>
 inline std::unique_ptr<BlockingCounter> SparseMatMul<TL, TR>::CreateDenseSlices(
     const typename SparseMatMul<TL, TR>::ConstMatrixMapR& mat, int row_start,
     int num_rows, int col_start, int num_cols,
-    const DeviceBase::CpuWorkerThreads* thread_pool, MatrixR* buffer,
+    const DeviceBase::CpuWorkerThreads* thread_pool, MatrixMapR* buffer,
     std::vector<typename SparseMatMul<TL, TR>::ConstMatrixMapR*>* slices) {
   std::unique_ptr<BlockingCounter> shuffle_counter(ShuffleMatrix(
-      mat, row_start, num_rows, col_start, num_cols, N, thread_pool, buffer));
+      mat, row_start, num_rows, col_start, num_cols, thread_pool, buffer));
   const int num_slices = (num_cols + N - 1) / N;
   SliceMatrix(*buffer, num_rows, num_slices, slices);
   return shuffle_counter;
@@ -1339,12 +1439,16 @@ inline void SparseMatMul<TL, TR>::Compute(
 
   const int right_dim0 = right.dimension(0);
   const int right_dim1 = right.dimension(1);
-  // Allocate buffer for storing slices of right matrix.
+  // Allocate 64-byte aligned buffer for storing slices of right matrix.
   // Note buffer needs enough space to hold at most a KR * NR matrix since that
   // is the block size per iteration.
   const int buffer_num_rows =
       std::min(KR, right_dim0) * ((std::min(NR, right_dim1) + N - 1) / N);
-  MatrixR buffer(buffer_num_rows, N);
+  std::unique_ptr<TR, void (*)(void*)> buffer_data(
+      static_cast<TR*>(Eigen::internal::handmade_aligned_malloc(
+          sizeof(TR) * buffer_num_rows * N, 64)),
+      Eigen::internal::handmade_aligned_free);
+  MatrixMapR buffer(buffer_data.get(), buffer_num_rows, N);
   std::vector<ConstMatrixMapR*> right_slices;
 
   std::vector<SparseSlice<TL>*> block_left_slices;
@@ -1434,3 +1538,4 @@ REGISTER_SPARSE_MATMUL(bfloat16, float);
 #undef REGISTER_SPARSE_MATMUL
 
 }  // end namespace tensorflow
+#endif  // HWY_ONCE
diff --git a/tensorflow/core/kernels/sparse_matmul_op_test.cc b/tensorflow/core/kernels/sparse_matmul_op_test.cc
index 9441bfc5941..ead70797040 100644
--- a/tensorflow/core/kernels/sparse_matmul_op_test.cc
+++ b/tensorflow/core/kernels/sparse_matmul_op_test.cc
@@ -14,11 +14,22 @@ limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/core/kernels/sparse_matmul_op.h"
+
+#include <cstdint>
+#include <vector>
+
+#include "hwy/targets.h"  // from @highway
 #include "tensorflow/core/common_runtime/kernel_benchmark_testlib.h"
 #include "tensorflow/core/framework/bfloat16.h"
+#include "tensorflow/core/framework/fake_input.h"
+#include "tensorflow/core/framework/node_def_builder.h"
 #include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/graph/node_builder.h"
+#include "tensorflow/core/kernels/ops_testutil.h"
 #include "tensorflow/core/lib/random/simple_philox.h"
 #include "tensorflow/core/lib/strings/stringprintf.h"
 #include "tensorflow/core/platform/test.h"
@@ -252,166 +263,79 @@ BM_SPARSE_MULTI(2048, 512, 256, 85, 85, 6);
 
 }  // end namespace tensorflow
 
-namespace Eigen {
-namespace internal {
+namespace tensorflow {
+namespace {
 
-class SparseMatmulOpTest : public ::testing::Test {
+// Test on all highway targets supported by test machine
+class SparseMatMulHighwayOpTest : public OpsTestBase {
  protected:
-  SparseMatmulOpTest()
-      : PacketSize(Eigen::internal::packet_traits<float>::size) {
-    typedef typename NumTraits<float>::Real RealFloat;
-
-    for (int i = 0; i < kMaxPacketSize; ++i) {
-      data1[i] = internal::random<float>() / RealFloat(PacketSize);
-      data2[i] = internal::random<float>() / RealFloat(PacketSize);
-      data3[i] = internal::random<float>() / RealFloat(PacketSize);
+  template <typename TA, typename TB>
+  void VerifyAcrossAllTargets(int m, int k, int n, bool transpose_a,
+                              bool transpose_b, bool a_is_sparse,
+                              bool b_is_sparse) {
+    TensorShape shape_a =
+        transpose_a ? TensorShape({k, m}) : TensorShape({m, k});
+    TensorShape shape_b =
+        transpose_b ? TensorShape({n, k}) : TensorShape({k, n});
+    std::vector<TA> data_a(shape_a.num_elements());
+    std::vector<TB> data_b(shape_b.num_elements());
+    for (int64_t i = 0; i < shape_a.num_elements(); ++i) {
+      data_a[i] = (i % 5 == 0) ? TA(0) : TA(static_cast<float>((i % 7) - 3));
     }
-    for (int i = kMaxPacketSize; i < kMaxPacketSize * 2; ++i) {
-      data3[i] = internal::random<float>() / RealFloat(PacketSize);
+    for (int64_t i = 0; i < shape_b.num_elements(); ++i) {
+      data_b[i] = (i % 3 == 0) ? TB(0) : TB(static_cast<float>((i % 5) - 2));
     }
 
-    // zero out lower 16-bits of mantissa of data3 values
-    // copy bfloat representation to data3_bfloat16
-    for (int i = 0; i < kMaxPacketSize * 2; ++i) {
-      uint16_t* data3_p = reinterpret_cast<uint16_t*>(&data3[i]);
-      uint16_t* data3_bfloat16_p =
-          reinterpret_cast<uint16_t*>(data3_bfloat16) + i;
-#if __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
-      data3_p[1] = 0;
-      data3_bfloat16_p[0] = data3_p[0];
-#else
-      data3_p[0] = 0;
-      data3_bfloat16_p[0] = data3_p[1];
-#endif
-    }
-  }
-
-  bool areApprox(const float* a, const float* b, int size) {
-    for (int i = 0; i < size; ++i) {
-      if (a[i] != b[i] && !internal::isApprox(a[i], b[i])) {
-        auto ma = Map<const Matrix<float, 1, Dynamic> >(a, size);
-        auto mb = Map<const Matrix<float, 1, Dynamic> >(b, size);
-        std::cout << "[" << ma << "]"
-                  << " != [" << mb << "], differences: [" << (mb - ma) << "]\n";
-        return false;
+    Tensor expected(allocator(), DT_FLOAT, TensorShape({m, n}));
+    auto expected_mat = expected.matrix<float>();
+    expected_mat.setZero();
+    for (int row = 0; row < m; ++row) {
+      for (int depth = 0; depth < k; ++depth) {
+        const float a_val = static_cast<float>(
+            transpose_a ? data_a[depth * m + row] : data_a[row * k + depth]);
+        for (int col = 0; col < n; ++col) {
+          const float b_val = static_cast<float>(
+              transpose_b ? data_b[col * k + depth] : data_b[depth * n + col]);
+          expected_mat(row, col) += a_val * b_val;
+        }
       }
     }
-    return true;
+
+    hwy::SetSupportedTargetsForTest(0);
+    const std::vector<int64_t> targets = hwy::SupportedAndGeneratedTargets();
+    ASSERT_FALSE(targets.empty());
+
+    for (int64_t target : targets) {
+      hwy::SetSupportedTargetsForTest(target);
+      inputs_.clear();
+      ASSERT_TRUE(NodeDefBuilder("sparse_matmul_op", "SparseMatMul")
+                      .Input(FakeInput(DataTypeToEnum<TA>::v()))
+                      .Input(FakeInput(DataTypeToEnum<TB>::v()))
+                      .Attr("transpose_a", transpose_a)
+                      .Attr("transpose_b", transpose_b)
+                      .Attr("a_is_sparse", a_is_sparse)
+                      .Attr("b_is_sparse", b_is_sparse)
+                      .Finalize(node_def())
+                      .ok());
+      ASSERT_TRUE(InitOp().ok());
+      AddInputFromArray<TA>(shape_a, data_a);
+      AddInputFromArray<TB>(shape_b, data_b);
+      ASSERT_TRUE(RunOpKernel().ok());
+      test::ExpectTensorNear<float>(expected, *GetOutput(0), 1e-4);
+    }
+    hwy::SetSupportedTargetsForTest(0);
   }
-
-#ifdef EIGEN_VECTORIZE_AVX512
-  static const int kMaxPacketSize = 16;
-#elif defined EIGEN_VECTORIZE_AVX || defined EIGEN_VECTORIZE_AVX2
-  static const int kMaxPacketSize = 8;
-#else
-  static constexpr int kMaxPacketSize = 4;
-#endif
-  typedef typename Eigen::internal::packet_traits<float>::type Packet;
-  const int PacketSize;
-  // float values
-  EIGEN_ALIGN_MAX float data1[kMaxPacketSize];
-  // output of intrinsics
-  EIGEN_ALIGN_MAX float data2[kMaxPacketSize];
-  // float values with only 7 mantissa bits (bfloat representable)
-  EIGEN_ALIGN_MAX float data3[kMaxPacketSize * 2];
-  // bfloat16 representation of data3
-  EIGEN_ALIGN_MAX float data3_bfloat16[kMaxPacketSize];
-  EIGEN_ALIGN_MAX float ref[kMaxPacketSize];
-
- public:
-  EIGEN_MAKE_ALIGNED_OPERATOR_NEW
 };
 
-TEST_F(SparseMatmulOpTest, BroadcastPacketTest) {
-  for (int i = 0; i < PacketSize; ++i) ref[i] = data1[0];
-  internal::pstoreu(data2, internal::pbroadcast_first<Packet>(
-                               internal::ploadu<Packet>(data1)));
-  ASSERT_TRUE(areApprox(ref, data2, PacketSize));
-  if (PacketSize > 1) {
-    for (int i = 0; i < PacketSize; ++i) ref[i] = data1[1];
-    internal::pstoreu(data2, internal::pbroadcast_second<Packet>(
-                                 internal::ploadu<Packet>(data1)));
-    ASSERT_TRUE(areApprox(ref, data2, PacketSize));
-
-    if (PacketSize > 2) {
-      for (int i = 0; i < PacketSize; ++i) ref[i] = data1[2];
-      internal::pstoreu(data2, internal::pbroadcast_third<Packet>(
-                                   internal::ploadu<Packet>(data1)));
-      ASSERT_TRUE(areApprox(ref, data2, PacketSize));
-
-      if (PacketSize > 3) {
-        for (int i = 0; i < PacketSize; ++i) ref[i] = data1[3];
-        internal::pstoreu(data2, internal::pbroadcast_fourth<Packet>(
-                                     internal::ploadu<Packet>(data1)));
-        ASSERT_TRUE(areApprox(ref, data2, PacketSize));
-      }
-    }
-  }
+TEST_F(SparseMatMulHighwayOpTest, AllSupportedTargetsMatchReference) {
+  VerifyAcrossAllTargets<float, float>(67, 133, 263, false, false, true, false);
+  VerifyAcrossAllTargets<bfloat16, bfloat16>(67, 133, 263, false, false, true,
+                                             false);
+  VerifyAcrossAllTargets<float, bfloat16>(67, 133, 128, true, true, true,
+                                          false);
+  VerifyAcrossAllTargets<bfloat16, float>(67, 133, 263, false, true, false,
+                                          true);
 }
 
-TEST_F(SparseMatmulOpTest, InterleavePacketTest) {
-  if (PacketSize == 8) {  // AVX
-    for (int i = 0; i < PacketSize / 4; ++i) ref[i] = data1[i];
-    for (int i = PacketSize / 4; i < PacketSize / 2; ++i)
-      ref[i] = data1[i + PacketSize / 4];
-    for (int i = PacketSize / 2; i < 3 * PacketSize / 4; ++i)
-      ref[i] = data1[i - PacketSize / 4];
-    for (int i = 3 * PacketSize / 4; i < PacketSize; ++i) ref[i] = data1[i];
-  } else {
-    // No interleaving done for smaller packets
-    for (int i = 0; i < PacketSize; ++i) ref[i] = data1[i];
-  }
-
-  internal::pstoreu(data2, internal::pinterleave4x64<Packet>(
-                               internal::ploadu<Packet>(data1)));
-  ASSERT_TRUE(areApprox(ref, data2, PacketSize));
-}
-
-TEST_F(SparseMatmulOpTest, Bfloat16ExpandTest) {
-  if (PacketSize == 8) {  // AVX
-    for (int i = 0; i < PacketSize / 2; ++i) {
-      ref[i] = data3[i];
-    }
-    for (int i = 0; i < PacketSize / 2; ++i) {
-      ref[i + PacketSize / 2] = data3[i + PacketSize];
-    }
-  } else {
-    for (int i = 0; i < PacketSize; ++i) {
-      ref[i] = data3[i];
-    }
-  }
-  internal::pstoreu(data2, internal::pexpand_bf16_l<Packet>(
-                               internal::ploadu<Packet>(data3_bfloat16)));
-  ASSERT_TRUE(areApprox(ref, data2, PacketSize));
-
-  if (PacketSize == 8) {  // AVX
-    for (int i = 0; i < PacketSize / 2; ++i) {
-      ref[i] = data3[i + PacketSize / 2];
-    }
-    for (int i = 0; i < PacketSize / 2; ++i) {
-      ref[i + PacketSize / 2] = data3[i + 3 * PacketSize / 2];
-    }
-  } else {
-    for (int i = 0; i < PacketSize; ++i) {
-      ref[i] = data3[i + PacketSize];
-    }
-  }
-
-  internal::pstoreu(data2, internal::pexpand_bf16_u<Packet>(
-                               internal::ploadu<Packet>(data3_bfloat16)));
-  ASSERT_TRUE(areApprox(ref, data2, PacketSize));
-}
-
-TEST_F(SparseMatmulOpTest, Bfloat16LoadTest) {
-  if (PacketSize >= 4) {
-    for (int i = 0; i < 4; ++i) ref[i] = data3[i];
-    internal::pstoreu(data2, internal::pload4bf16<Packet>(data3_bfloat16));
-    ASSERT_TRUE(areApprox(ref, data2, 4));
-
-    internal::pstoreu(data2, internal::pload2bf16<Packet>(data3_bfloat16));
-    ASSERT_TRUE(areApprox(ref, data2, 2));
-  }
-}
-
-}  // namespace internal
-}  // namespace Eigen
+}  // namespace
+}  // namespace tensorflow

From f92bcc6c34fd99062e115d78f0d0d1e337910f86 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 9 Oct 2026 13:07:42 -0700
Subject: [PATCH 0387/1000] Move `optimization_parameters.proto` and
 `tpu_embdding_configuration.proto` to `xla/tpu`

Reverts 1e1629fd562d06d7c8d485bba0942b76e5f96685

PiperOrigin-RevId: 996710734
---
 tensorflow/core/protobuf/tpu/BUILD            |  12 +-
 .../tpu/optimization_parameters.proto         | 639 ++++++++++++++++-
 .../tpu/tpu_embedding_configuration.proto     | 158 ++++-
 third_party/xla/xla/tpu/protobuf/BUILD        |  60 --
 .../protobuf/optimization_parameters.proto    | 657 ------------------
 .../tpu_embedding_configuration.proto         | 176 -----
 6 files changed, 797 insertions(+), 905 deletions(-)
 delete mode 100644 third_party/xla/xla/tpu/protobuf/BUILD
 delete mode 100644 third_party/xla/xla/tpu/protobuf/optimization_parameters.proto
 delete mode 100644 third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration.proto

diff --git a/tensorflow/core/protobuf/tpu/BUILD b/tensorflow/core/protobuf/tpu/BUILD
index 08bf50b04d7..35a385fdb85 100644
--- a/tensorflow/core/protobuf/tpu/BUILD
+++ b/tensorflow/core/protobuf/tpu/BUILD
@@ -32,12 +32,9 @@ tf_proto_library(
         "tpu_embedding_configuration.proto",
     ],
     protodeps = [
-        "@xla//xla/tpu/protobuf:tpu_embedding_configuration_proto",
+        ":optimization_parameters_proto",
     ],
     visibility = ["//visibility:public"],
-    exports = [
-        "@xla//xla/tpu/protobuf:tpu_embedding_configuration_proto",
-    ],
 )
 
 tf_proto_library(
@@ -45,13 +42,8 @@ tf_proto_library(
     srcs = [
         "optimization_parameters.proto",
     ],
-    protodeps = [
-        "@xla//xla/tpu/protobuf:optimization_parameters_proto",
-    ],
+    protodeps = ["@xla//xla/service:hlo_proto"],
     visibility = ["//visibility:public"],
-    exports = [
-        "@xla//xla/tpu/protobuf:optimization_parameters_proto",
-    ],
 )
 
 tf_proto_library(
diff --git a/tensorflow/core/protobuf/tpu/optimization_parameters.proto b/tensorflow/core/protobuf/tpu/optimization_parameters.proto
index 869e9222049..6899a183f58 100644
--- a/tensorflow/core/protobuf/tpu/optimization_parameters.proto
+++ b/tensorflow/core/protobuf/tpu/optimization_parameters.proto
@@ -17,4 +17,641 @@ syntax = "proto3";
 
 package tensorflow.tpu;
 
-import public "xla/tpu/protobuf/optimization_parameters.proto";
+import "google/protobuf/wrappers.proto";
+import "xla/service/hlo.proto";
+
+message ClippingLimits {
+  google.protobuf.FloatValue lower = 1;  // -inf if not set
+  google.protobuf.FloatValue upper = 2;  // +inf if not set
+}
+
+// Configuration for simulated quantization; simulated quantization is used to
+// reduce training/serving skew when the serving variables are quantized. The
+// same quantization operations are executed during training to minimize
+// differences with serving.
+//
+// Simulated quantization inserts the following operations on the forward pass
+// after gathering the embedding vector from HBM. The backward pass operations
+// are unchanged.
+//
+// clipped_val = clip(input, clipping_limits)
+// quantum = clipping_limits.range() / (num_buckets - 1)
+// quantized_val = floor((clipped_val - clipping_limits.lower()) / quantum + .5)
+// return quantized_val * quantum + clipping_limits.lower().
+message SimulatedQuantization {
+  // Whether simulated quantization is enabled.
+  bool enabled = 1;
+
+  // Minimum and maximum values of the range used for quantization.
+  ClippingLimits clipping_limits = 2;
+
+  // Number of possible quantized values.
+  int32 num_buckets = 3;
+}
+
+// Dynamic input specification for optimizers in the TPUEmbeddingConfiguration.
+// The actual dynamic inputs are provided as a scalar input list to the
+// SendTPUEmbeddingGradients Op indexed by their tag specified through the
+// following proto.
+message OptimizerDynamicInput {
+  // For tables where dynamic inputs are needed (e.g., learning rates or other
+  // dynamic hyperparameters used in optimizers), a tag must be specified for
+  // the input.
+  //
+  // The tag must be a non-negative  integer. All tags in the range
+  // [0, number_of_unique_tags) must be present in the TPU embedding
+  // configuration, i.e. a tag cannot be skipped if a different tag numerically
+  // greater than it is used in the configuration.
+  //
+  // If multiple tables specify the same tag, they *MUST* have
+  // the same dynamic input, for example, their dynamic learning rate could be
+  // computed by the same TensorFlow sub-graph. The partitioning of the
+  // embedding layer would be more optimal if the number_of_unique_tags is as
+  // *LOW* as possible, i.e., if many tables share the same tag.
+  //
+  // The hyper_parameters input of the SendTPUEmbeddingGradients op is used to
+  // communicate dynamic hyper-parameters to the TPU embedding program.
+  // The hyper_parameters input is a list of scalars where the size of the list
+  // is equal to the number of unique tags. The hyper-parameter associated with
+  // a particular tag is specified by populating its corresponding index in the
+  // list of scalars.
+  int32 tag = 1;
+}
+
+// Source of learning rate to use.
+message LearningRate {
+  oneof learning_rate {
+    float constant = 1;
+    OptimizerDynamicInput dynamic = 2;
+  }
+}
+
+// Each optimizer's parameter proto has a link to its documentation and CPU
+// implementation (if available) for user reference.
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adagrad
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L1634
+message AdagradParameters {
+  // Old initial accumulator parameter.
+  reserved "initial_accumulator";
+  reserved 1;
+}
+
+// This optimizer combines the Adagrad and Momentum update rules.
+// accum(new) = beta2 == 1.0 ?
+//              accum(old) + grad^2 :
+//              beta2 * accum(old) + (1 - beta2) * grad^2
+// accum_with_exponent = (accum(new) + epsilon)^(-1.0 / exponent)
+// mom_accum(new) = momentum * mom_accum(old) + accum_with_exponent
+// update = use_nesterov ?
+//          momentum * mom_accum(new) + accum_with_exponent :
+//          mom_accum(new)
+// var(new) = var(old) - lr * grad * update
+// Algorithm described in https://arxiv.org/abs/2002.11803.
+message AdagradMomentumParameters {
+  // Moving average parameter for the momentum accumulator.
+  float momentum = 1;
+
+  // Whether to use the Nesterov variant of momentum.
+  bool use_nesterov = 2;
+
+  // Exponent for the gradient^2 accumulator.
+  float exponent = 3;
+
+  // Moving average parameter for the gradient^2 accumulator.
+  float beta2 = 4;
+
+  // Offset added to the Adagrad accumulator.
+  float epsilon = 5;
+}
+
+// Algorithm in http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf.
+message BoundedAdagradParameters {
+  // Whether to use the updated or the old value of the accumulator when
+  // computing the effective learning rate. When update_accumulator_first is set
+  // to True, the updated value of the accumulator is used.
+  bool update_accumulator_first = 1;
+
+  // The max_var_update value to use. Set value to 0 (default) to disable using
+  // max_var_update to clip the gradient.
+  float max_var_update = 2;
+
+  // The maximum value of the accumulator. Set max_accumulator to 0 (default)
+  // to disable using max_accumulator to clip the accumulator.
+  float max_accumulator = 3;
+}
+
+// Frequency Aware Adagrad optimizer. This optimizer implements the AdaGrad
+// algorithm and further allows to:
+// * Scale the learning rate based on frequency of the update. Sparsely updated
+//   rows are updated with a higher effective learning rate, and frequently
+//   updated rows are updated with a lower effective learning rate.
+// * Decay the growth of the accumulator values.
+// * Use L1 / L2 regularization for the weight updates.
+//
+// The optimization algorithm is shown below.
+// counter(new) = counter(old) + 1
+// accum(new) = max(accumulator_decay * accum(old) + grad^2,
+//                  initial_accumulator_value)
+// lr_scale = min((step_counter / accum(new)) ^ probability_exponent,
+// max_lr_multiplier) update = grad * lr_scale / sqrt(accum(new)) if
+// (l1_regularization_strength > 0.0):
+//   update = update + l1_regularization_strength * sign(var(old))
+// if (l2_regularization_strength > 0.0):
+//   update = update + l2_regularization_strength * var(old)
+// var(new) = var(old) - lr_scale * grad * update
+
+message FrequencyAwareAdagradParameters {
+  // The L1 regularization parameter for adjusting the update based on the sign
+  // of the variable.
+  float l1_regularization_strength = 1;
+
+  // The L2 regularization parameter for adjusting the update based on the
+  // variable.
+  float l2_regularization_strength = 2;
+
+  // The exponent used for scaling the learning rate based on the sparsity of
+  // updates.
+  float probability_exponent = 4;
+
+  // The maximum value of the learning rate scale.
+  float max_lr_multiplier = 3;
+
+  // The decay for the Adagrad accumulator.
+  float accumulator_decay = 5;
+
+  // The initial and minimum value for the Adagrad accumulator.
+  float initial_accumulator_value = 6;
+
+  // The tag for identifying the step counter used for the frequency aware
+  // Adagrad optimizer.
+  OptimizerDynamicInput step_counter = 7;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L629
+message StochasticGradientDescentParameters {}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Ftrl
+// https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/41159.pdf
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L2646
+//
+// The hyperparameters for FTRL are the same as for the Keras implementation,
+// with some additions. The "beta" parameter matches the behavior described in
+// the second link above; "beta" / (2 * learning rate) should be added to "l2"
+// to get equivalent behavior in the other TensorFlow implementations of this
+// optimizer. When the multiply_linear_by_lr field is set to true, a modified
+// formula is used for FTRL that treats the "linear" accumulator as being
+// pre-multiplied by the learning rate (i.e., the accumulator named "linear"
+// actually stores "linear * learning_rate"). Other than checkpoint
+// compatibility, this is mathematically equivalent for a static learning rate;
+// for a dynamic learning rate, it is nearly the same as long as the learning
+// rate does not change quickly. The benefit of setting multiply_linear_by_lr to
+// true is that the modified formula handles zero and near-zero learning rates
+// without producing NaNs, improving flexibility for learning rate ramp-up.
+message FtrlParameters {
+  float l1 = 1;
+  float l2 = 2;
+  float lr_power = 3;
+  float beta = 7;
+  bool multiply_linear_by_lr = 6;
+
+  // Previously, allow_zero_accumulator parameter changed some internal formulas
+  // to allow zero and near-zero accumulator values at the cost of some
+  // performance. The current implementation ignores this parameter; zero or
+  // near-zero accumulator values are now always supported.
+  bool allow_zero_accumulator = 8 [deprecated = true];
+
+  // Old initial accumulator parameters.
+  reserved "initial_accum", "initial_linear";
+  reserved 4, 5;
+}
+
+// The Adam optimizer does not implement hyper-parameter update due to hardware
+// limitations; use the dynamic learning rate feature instead, setting the
+// learning rate to: user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
+// Here, t is the current timestep.
+//
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adam
+// https://github.com/tensorflow/tensorflow/blob/ab51450c817674c8ff08a7ae4f8ac50cdc4bed8b/tensorflow/python/training/adam.py#L32
+//
+// Note that the code by default implements the lazy version of Adam
+// (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/LazyAdamOptimizer)
+// unless the use_non_lazy_adam parameter is set, in which case it implements
+// the normal version of Adam that updates all parameters in the embedding
+// table, even for entries that are not used in the current minibatch
+// (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/AdamOptimizer). If
+// use_non_lazy_adam is enabled, gradient accumulation is also required to be
+// enabled in order to get correct results; a warning will be printed otherwise
+// (which may change to an error in the future). If use_sum_inside_sqrt is set,
+// the Adam variable update formula will be changed from m / (sqrt(v) + epsilon)
+// to m / sqrt(v + epsilon**2); this option improves the performance of TPU
+// training and is not expected to harm model quality.
+message AdamParameters {
+  float beta1 = 3;
+  float beta2 = 4;
+  float epsilon = 5;
+  bool use_non_lazy_adam = 8;
+  bool use_sum_inside_sqrt = 10;
+
+  // Old initial accumulator parameters.
+  reserved "initial_m", "initial_v";
+  reserved 6, 7;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L3068
+message MomentumParameters {
+  float momentum = 1;
+  bool use_nesterov = 2;
+
+  // Old initial accumulator parameter.
+  reserved "initial_accum";
+  reserved 3;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Lion
+// momenta(new) = beta2 * momenta(old) + (1 - beta2) * grad
+// momenta_t = beta1 * momenta(old) + (1 - beta1) * grad
+// var(new) = var(old) - lr * sign(momenta_t)
+// Algorithm described in https://arxiv.org/abs/2302.06675.
+message LionParameters {
+  float beta1 = 1;
+  float beta2 = 2;
+  bool use_non_lazy_lion = 3;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L4229
+message RmsPropParameters {
+  float rho = 1;
+  float momentum = 2;
+  float epsilon = 3;
+
+  // Old initial accumulator parameters.
+  reserved "initial_ms", "initial_mom";
+  reserved 4, 5;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L4358
+message CenteredRmsPropParameters {
+  float rho = 1;
+  float momentum = 2;
+  float epsilon = 3;
+
+  // Old initial accumulator parameters.
+  reserved "initial_ms", "initial_mom", "initial_mg";
+  reserved 4, 5, 6;
+}
+
+// Variant of algorithm in http://proceedings.mlr.press/v44/shamir15.pdf
+message MdlAdagradLightParameters {
+  float l2 = 1;
+  float lr_power = 2;
+  float min_servable_mdl_benefit = 3;
+  float mdl_mix_in_margin = 4;
+  float mdl_benefit_rampup_coeff = 5;
+  float mdl_min_weight = 6;
+  float benefit_revisit_scale = 7;
+  float max_event_benefit = 8;
+  float max_total_benefit = 9;
+  float mdl_hard_limit = 10;
+  bool hard_limit_min_benefit = 11;
+  bool mdl_regularize = 12;
+
+  // Old initial accumulator parameters.
+  reserved "initial_accumulator", "initial_weight", "initial_benefit";
+  reserved 13, 14, 15;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adadelta
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L933
+message AdadeltaParameters {
+  float rho = 1;
+  float epsilon = 2;
+
+  // Old initial accumulator parameters.
+  reserved "initial_accumulator", "initial_update";
+  reserved 3, 4;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/compat/v1/train/ProximalAdagradOptimizer
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L1961
+message ProximalAdagradParameters {
+  float l1 = 1;
+  float l2 = 2;
+
+  // Old initial accumulator parameter.
+  reserved "initial_accumulator";
+  reserved 3;
+}
+
+// The online Yogi optimizer does not implement hyper-parameter update; use the
+// dynamic learning rate feature instead, setting the learning rate to:
+// user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
+// Here, t is the current timestep.
+//
+// https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf
+// plus some extensions based on FTRL.
+//
+// Note that the code by default implements the lazy version of online Yogi.
+message OnlineYogiParameters {
+  // The L1 regularization parameter (used analogously to the one in FTRL).
+  float l1 = 1;
+
+  // The L2 regularization parameter (used analogously to the one in FTRL).
+  float l2 = 2;
+
+  // \beta_2 from Algorithm 2 in the paper.
+  float beta2 = 3;
+
+  // Reserved ids corresponding to removed tanh activation.
+  reserved 6;  // sign
+  reserved 7;  // tanh
+}
+
+// The online Yogi optimizer does not implement hyper-parameter update; use the
+// dynamic learning rate feature instead, setting the learning rate to:
+// user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
+// Here, t is the current timestep.
+//
+// https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf
+// plus some extensions based on FTRL.
+//
+// Note that the code by default implements the lazy version of proximal Yogi.
+message ProximalYogiParameters {
+  // The L1 regularization parameter.
+  float l1 = 1;
+
+  // The L2 regularization parameter.
+  float l2 = 2;
+
+  // The exponential decay rate for the 1st moment estimates.
+  float beta1 = 3;
+
+  // The exponential decay rate for the 2nd moment estimates.
+  float beta2 = 4;
+
+  // A constant trading off adaptivity and noise.
+  float epsilon = 5;
+
+  // Reserved ids corresponding to removed tanh activation.
+  reserved 8;  // sign
+  reserved 9;  // tanh
+}
+
+// Estimator for the frequency of updates to a lookup table. It maintains an
+// array (tf.Variable) D, where each element records the average number of
+// global steps between two consecutive batches that hit the corresponding
+// bucket. Once an item with bucket id i is sampled, D[i] is updated by:
+//   D[i] <- D[i] * (1 - tau) + delta[i] * tau,
+//
+// where tau is a learning rate between 0 and 1 (exclusive), and
+//   delta[i] = current global step - last step i is sampled.
+//
+// The estimated frequency (sampling rate in a batch) is thus 1 / D[i].
+//
+// Elements in D are initialized with a large value max_delta. delta[i] will
+// also be capped by this value.
+//
+// The exact sequence of operations used in the optimizer is shown below.
+// last_hit_step[i] is a tf.Variable that holds the last global step at which i
+// was sampled.
+//
+//   delta = global_step - last_hit_step[i]
+//   clipped_delta = min(delta, params.max_delta)
+//   is_outlier = (delta >= params.outlier_threshold * D[i])
+//   D[i] <- is_outlier ? clipped_delta
+//                      : D[i] * (1 - params.tau) + clipped_delta * params.tau
+//   last_hit_step[i] <- global_step
+message FrequencyEstimatorParameters {
+  // Learning rate between (0, 1) that is used to update the array D.
+  float tau = 1;
+
+  // Maximum value of delta: difference between the current global step and the
+  // last global step at which the row was sampled.
+  float max_delta = 2;
+
+  // Threshold used to determine whether the current update is an outlier.
+  float outlier_threshold = 3;
+
+  // The weight exponent used to transform the estimated delta into weights.
+  // The transformation function is: (delta / max_delta) ^ (weight_exponent)
+  float weight_exponent = 4;
+}
+
+// A user-defined optimizer.
+// The contained HLO program must take the following arguments in the following
+// order:
+// 1.  gradients
+// 2.  table weights
+// 3.  slot variables
+// 4.  an optional scalar input that is passed in via the dynamic learning
+//     rate mechanism.
+//
+// It must return/end in a tuple op that contains the following values in the
+// following order:
+// 1.  new table values
+// 2.  new slot variable value
+//
+// The program must have shape (1,1) with dtype float32 throughout and only use
+// HLO that operate elementwise (e.g., no reduce, no variables, no control flow
+// and no broadcasting outside of the single scalar input).
+// The HLO program should be written as if it were a dense update. It will be
+// called on each row that needs an update and will applied elementwise.
+message UserDefinedProgramParameters {
+  xla.HloModuleProto program = 1;
+  reserved 2;  // Was padding_values
+}
+
+// Optimizer that just sets the variable to the value of the gradient. To be
+// correct, this requires either gradient accumulation (to sum the values of a
+// computed expression across the samples) or to deduplicate IDs within a single
+// host (to assign the value from an arbitrary sample).
+message AssignParameters {}
+
+// Status of using gradient accumulation (doing two passes over the input
+// gradients: one to accumulate them into a temporary array and another to apply
+// them using the actual optimization algorithm). The extra message is to wrap
+// the enum for scoping.
+message GradientAccumulationStatus {
+  // if UNSPECIFIED (default), gradient accumulation is ENABLED.
+  enum Status {
+    UNSPECIFIED = 0;
+    ENABLED = 1;
+    DISABLED = 2;
+  }
+}
+
+// Whether to optimize the packing of low-dimensional embedding tables in HBM
+// (high bandwidth memory). TPUs access HBM at 32-byte (8-float) granularity.
+// For functional correctness, the TPU software internally pads the embedding
+// dimension to a multiple of 8. This can sometimes lead to significant memory
+// wastage due to padding. For 1-dimensional, 2-dimensional, and 4-dimensional,
+// the TPU software can remove this padding by packing multiple rows into the
+// same 8-float HBM chunk. For example, 8 rows could be packed into the same
+// 8-float chunk for a 1-dimensional embedding table.
+
+// There is one important limitation for this HBM packing though. When only a
+// subset of rows in an 8-float chunk are accessed on a particular step, the
+// adjoining rows in the same chunk are updated with zero gradients on the
+// backward pass even if they are not touched. This is an artifact of the
+// packing implementation. This operation is NOT functionally correct for
+// optimizers where zero gradients change the embeddings/slot-variable values,
+// e.g., momentum-based optimizers. Hence, this HBM packing cannot be enabled
+// for embedding tables with such optimizers. The TPU software automatically
+// recognizes that a zero gradient can modify state and turns off the low
+// dimensional embedding packing in that scenario.
+//
+// However, for optimizers where a zero gradient is a NoOp, such as SGD,
+// Adagrad, and FTRL, this packing optimization can be used. However, there are
+// some important considerations:
+// * Clipping limits: The initial values for such embeddings should fall within
+//   the clipping limits specified in the optimization parameters. Otherwise, a
+//   zero gradient will cause the embeddings to be clipped. This changes state
+//   and hence, is not a NoOp.
+// * FTRL: The embedding vector is computed directly from the values of the
+//   accumulator and linear slot variables. Hence, the initial embedding values
+//   should match that computed from the initial values of the accumulator and
+//   linear slot variables. Note that in nearly all cases, the linear value is
+//   initialized to zero; this corresponds to an embedding value of zero.
+//
+// Performance: The TPU has to perform additional work when low dimensional
+// packing is enabled. In certain situations when the vocabulary size is small,
+// it may not make sense to turn on this packing since the total memory usage
+// due to padding is extremely low. Hence, the TPU software automatically turns
+// off the packing optimization in such scenarios.
+message LowDimensionalPackingStatus {
+  // if UNSPECIFIED (default), the low dimension packing status is DISABLED.
+  // This can change in future.
+  //
+  // if ENABLED, the low dimension packing is enabled only if the following
+  // three additional conditions are true:
+  //  * The optimizer treats the zero gradient as a NoOp.
+  //  * The embedding dimension is 1, 2, or 4.
+  //  * The vocabulary size is large enough to avoid performance issues.
+  //
+  // if DISABLED, the low dimension packing is always disabled.
+  enum Status {
+    UNSPECIFIED = 0;
+    ENABLED = 1;
+    DISABLED = 2;
+  }
+}
+
+// Configuration proto for hot ID optimization. This is an experimental feature
+// that is currently disabled (by default).
+message HotIdReplicationConfiguration {
+  // Whether to enable or disable hot ID optimization.
+  // If set to UNSPECIFIED (default), hot ID optimization is DISABLED.
+  // If set to ENABLED, hot ID replication is turned ON.
+  // If set to MIGRATION_ONLY, hot ID migration is turned ON.
+  enum Status {
+    UNSPECIFIED = 0;
+    ENABLED = 1;
+    DISABLED = 2;
+    MIGRATION_ONLY = 3;
+  }
+  Status status = 1;
+}
+
+message OptimizationParameters {
+  // Learning rate used for updating the embedding layer parameters.
+  LearningRate learning_rate = 13;
+
+  // Limits to which to clip the weight values after the backward pass; not
+  // present means no limits are applied.
+  ClippingLimits clipping_limits = 2;
+
+  // Limits to which to clip the backward pass gradient before using it for
+  // updates; not present means no limits are applied.
+  ClippingLimits gradient_clipping_limits = 7;
+
+  // Amount of weight decay to apply; see weight_decay_optimizers.py for
+  // details. All optimizers except MDL Adagrad Light are supported with this
+  // option. Although there is no check, users who want weight decay will also
+  // want to ensure that gradient accumulation is enabled so that the decay will
+  // happen once per global batch.
+  float weight_decay_factor = 16;
+
+  // If true, the weight decay factor is multiplied by the current learning rate
+  // before use; this is to match the note in DecoupledWeightDecayExtension in
+  // weight_decay_optimizers.py.
+  bool multiply_weight_decay_factor_by_learning_rate = 22;
+
+  // Configuration for simulated quantization which is used to reduce
+  // training/serving skew when the serving variables are quantized. The same
+  // quantization operations are executed during training to minimize
+  // differences with serving.
+  SimulatedQuantization simulated_quantization = 27;
+
+  // Status of using gradient accumulation (doing two passes over the input
+  // gradients: one to accumulate them into a temporary array and another to
+  // apply them using the actual optimization algorithm).
+  GradientAccumulationStatus.Status gradient_accumulation_status = 17;
+
+  // Status of the low-dimensional embedding packing optimization. This controls
+  // whether to optimize the packing of 1-dimensional, 2-dimensional, and
+  // 4-dimensional embedding tables in memory.
+  LowDimensionalPackingStatus.Status low_dimensional_packing_status = 28;
+
+  // Configuration proto for hot ID replication. This is an experimental
+  // feature that is currently disabled (by default).
+  HotIdReplicationConfiguration hot_id_replication_configuration = 18;
+
+  // Optimization algorithm parameters; which field is selected determines which
+  // algorithm to use.
+  oneof parameters {
+    AdagradParameters adagrad = 3;
+    AdagradMomentumParameters adagrad_momentum = 26;
+    BoundedAdagradParameters bounded_adagrad = 19;
+    FrequencyAwareAdagradParameters frequency_aware_adagrad = 30;
+    StochasticGradientDescentParameters stochastic_gradient_descent = 4;
+    FtrlParameters ftrl = 5;
+    AdamParameters adam = 6;
+    MomentumParameters momentum = 8;
+    LionParameters lion = 29;
+    RmsPropParameters rms_prop = 9;
+    CenteredRmsPropParameters centered_rms_prop = 10;
+    MdlAdagradLightParameters mdl_adagrad_light = 11;
+    AdadeltaParameters adadelta = 12;
+    ProximalAdagradParameters proximal_adagrad = 14;
+    OnlineYogiParameters online_yogi = 20;
+    ProximalYogiParameters proximal_yogi = 21;
+    FrequencyEstimatorParameters frequency_estimator = 23;
+    UserDefinedProgramParameters user_defined_program = 24;
+    AssignParameters assign = 25;
+  }
+
+  reserved 1, 15;
+
+  // NEXT_ID: 31
+}
+
+// Specification of an optimization algorithm's state variables (both the main
+// value vector and any extra accumulators, etc.). This proto is only used
+// internally by the TPU software and is not exposed directly to the TF model.
+message StateVariableSpecification {
+  // Parameter name for the state variable.
+  string name = 1;
+
+  // A normal state variable that should be saved and restored in checkpoints
+  // and used as an input or output to non-debug TensorFlow ops.
+  message UserDefined {
+    reserved 1;  // Was padding_initial_value.
+  }
+
+  // A state variable that should be filled with a constant and normally hidden
+  // from users (used for intermediate gradients being accumulated, for
+  // example).
+  message FillWithConstant {
+    double initial_value = 1;
+  }
+
+  // Usage type of this state variable.
+  oneof usage {
+    UserDefined user_defined = 2;
+    FillWithConstant fill_with_constant = 3;
+  }
+}
diff --git a/tensorflow/core/protobuf/tpu/tpu_embedding_configuration.proto b/tensorflow/core/protobuf/tpu/tpu_embedding_configuration.proto
index ffc8d2e26ab..c11ece0c6f4 100644
--- a/tensorflow/core/protobuf/tpu/tpu_embedding_configuration.proto
+++ b/tensorflow/core/protobuf/tpu/tpu_embedding_configuration.proto
@@ -17,4 +17,160 @@ syntax = "proto3";
 
 package tensorflow.tpu;
 
-import public "xla/tpu/protobuf/tpu_embedding_configuration.proto";
+import "tensorflow/core/protobuf/tpu/optimization_parameters.proto";
+
+message TPUEmbeddingConfiguration {
+  // Description of the various embedding tables.
+  message TableDescriptor {
+    // Name of the table.
+    string name = 1;
+
+    // Size of the vocabulary (i.e., number of rows) in the table.
+    int64 vocabulary_size = 2;
+
+    // The embedding dimension (i.e., the width of the embedding table).
+    int32 dimension = 3;
+
+    // Number of features mapped to this table.
+    int32 num_features = 4;
+
+    // Details of the learning algorithm used to update the embedding
+    // parameters.
+    OptimizationParameters optimization_parameters = 5;
+  }
+  repeated TableDescriptor table_descriptor = 1;
+
+  // Mode. Should the embedding layer program be run for inference (just forward
+  // pass), training (both forward and backward pass) or just the backward_pass.
+  enum Mode {
+    UNSPECIFIED = 0;
+    INFERENCE = 1;
+    TRAINING = 2;
+    BACKWARD_PASS_ONLY = 3;
+  }
+  Mode mode = 2;
+
+  // Number of samples in each batch of embedding layer activations sent to
+  // the TensorCore.
+  int32 batch_size_per_tensor_core = 3;
+
+  // Number of TPU hosts used for inference/training.
+  int32 num_hosts = 4;
+
+  // Number of TensorCore used for inference/training.
+  int32 num_tensor_cores = 5;
+
+  // Sharding strategy of the embedding tables among the hosts.
+  // If the sharding_strategy is "mod", each id is assigned to host
+  // "id % num_hosts". For instance, 13 ids are split across 5 hosts as:
+  // [[0, 5, 10], [1, 6, 11], [2, 7, 12], [3, 8], [4, 9]].
+  // If the sharding_strategy is "div", ids are assigned to hosts in a
+  // contiguous manner. In this case, 13 ids are split across 5 hosts as:
+  // [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11, 12]].
+  // In both the strategies, if the id space does not evenly divide the number
+  // of hosts, each of the first "table_descriptor.vocabulary_size % num_hosts"
+  // hosts will be assigned one more id.
+  // This partitioning strategy exactly follows that in the embedding_lookup
+  // TensorFlow function at tensorflow/python/ops/embedding_ops.py.
+  enum ShardingStrategy {
+    DIV_DEFAULT = 0;
+    MOD = 1;
+  }
+  ShardingStrategy sharding_strategy = 6;
+
+  // This parameter determines if the execution of the sparse core will be
+  // pipelined with that of the TensorCore. This parameter only affects results
+  // when mode=TRAINING. If mode=INFERENCE or BACKWARD_PASS_ONLY, this parameter
+  // does not affect execution and hence, is a don't care value.
+  //
+  // false: The execution of the sparse core is not pipelined with that of the
+  // TensorCore. The forward pass of every step on the sparse core is executed
+  // only after the backward pass of the previous step is complete. And the
+  // backward pass on the sparse core is executed only after the embedding
+  // gradients have been computed on the TensorCore on every step. This ensures
+  // that the activations on every step observe the gradient updates from the
+  // previous step on both the sparse core and the TensorCore.
+  //
+  // true: The execution of the sparse core is pipelined with that of the
+  // TensorCore. The forward pass of every step on the sparse core can be
+  // executed after the forward pass of the previous step is complete without
+  // waiting for the backward pass. This improves the utilization of the sparse
+  // core allowing it to process step N+1 while the embedding gradients for step
+  // N are computed on the TensorCore. The backward pass of every step on the
+  // sparse core is executed directly after the forward pass for the next step
+  // is complete. The drawback is that embedding activations for step N+1 do not
+  // observe the embedding gradient updates from step N. This could affect model
+  // quality if step N and N+1 involve the same set of embedding IDs. However,
+  // since the embedding updates are sparse, this is generally not considered a
+  // problem.
+  bool pipeline_execution_with_tensor_core = 7;
+
+  // Directory where embedding lookup statistics are stored. These statistics
+  // summarize information about the inputs to the embedding lookup
+  // operation, in particular, the average number of embedding IDs per example
+  // and how well the embedding IDs are load balanced across the system. The
+  // lookup statistics are used during TPU initialization for embedding table
+  // partitioning. Collection of lookup statistics is done at runtime by
+  // profiling the embedding inputs: only 3% of input samples are profiled to
+  // minimize host CPU overhead. Once a suitable number of samples are
+  // profiled, the lookup statistics are saved to table-specific files in the
+  // profile data directory generally at the end of a TPU training loop. The
+  // filename corresponding to each table is obtained by hashing table specific
+  // parameters (e.g., table name and number of features) and global
+  // configuration parameters (e.g., sharding strategy and TPU worker task
+  // count). The same profile data directory can be shared amongst several
+  // models to reuse embedding lookup statistics.
+  string profile_data_directory = 9;
+
+  // Description of different input features.
+  message FeatureDescriptor {
+    // Name of the input feature.
+    string name = 1;
+
+    // Index of the corresponding table in the TableDescriptor list.
+    int32 table_id = 2;
+
+    // Static shape of the inputs (excluding the reduction axis). Note that
+    // the shape of the actual inputs provided using the infeed op must be
+    // strictly smaller than input_shape. The outputs received at the TensorCore
+    // will have rank = input_shape.size() + 1. The innermost axis corresponds
+    // to the embedding dimension. If the input has shape [m, n, k] (excluding
+    // the reduction axis) and the embedding dimension is d, the output received
+    // at the TensorCore will have shape [m, n, k, d].
+    repeated int32 input_shape = 3;
+  }
+  // If the feature_descriptor field is populated, the model should NOT populate
+  // TableDescriptor.num_features and batch_size_per_tensor_core. These two
+  // fields will be auto-populated by the TPUEmbedding rewrite passes.
+  repeated FeatureDescriptor feature_descriptor = 10;
+
+  // SPMD (Single Program Multiple Data) sharding configuration for
+  // TPUEmbedding. When model parallelism is used on the TensorCore, the number
+  // of cores per replica must be passed to TPUEmbedding so that the right
+  // shapes can be computed in the TF/XLA bridge.
+  message SpmdSharding {
+    // Whether SPMD sharding is enabled.
+    bool enabled = 1;
+
+    // Number of cores per replica.
+    int32 num_cores_per_replica = 2;
+
+    // If true, the tensors are manually partitioned. Otherwise, use the
+    // automatic SPMD partitioning. This should be true when users use
+    // `shard_map`.
+    bool use_manual_partitioning = 3;
+
+    // If true, lowers to Shardy custom-call for manual partitioning.
+    // Default is false.
+    bool use_shardy_partitioner = 4;
+  }
+  SpmdSharding spmd_sharding = 11;
+
+  // Old TPU embedding output layout.
+  reserved "output_layout";
+  reserved 8;
+}
+
+// A placeholder message that is used to define a unique Status payload
+// URL for TPU embedding errors.
+message TPUEmbeddingError {}
diff --git a/third_party/xla/xla/tpu/protobuf/BUILD b/third_party/xla/xla/tpu/protobuf/BUILD
deleted file mode 100644
index f732b3e19e7..00000000000
--- a/third_party/xla/xla/tpu/protobuf/BUILD
+++ /dev/null
@@ -1,60 +0,0 @@
-# Copyright 2026 The OpenXLA 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.
-# ==============================================================================
-
-# copybara:uncomment load("@rules_python//python:proto.bzl", "py_proto_library")
-load(
-    "//xla/tsl/platform:build_config.bzl",
-    "tf_proto_library",
-)
-
-package(
-    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
-    licenses = ["notice"],
-)
-
-tf_proto_library(
-    name = "tpu_embedding_configuration_proto",
-    srcs = [
-        "tpu_embedding_configuration.proto",
-    ],
-    protodeps = [
-        ":optimization_parameters_proto",
-    ],
-    visibility = ["//visibility:public"],
-)
-
-tf_proto_library(
-    name = "optimization_parameters_proto",
-    srcs = [
-        "optimization_parameters.proto",
-    ],
-    protodeps = ["//xla/service:hlo_proto"],
-    visibility = ["//visibility:public"],
-    deps = ["@com_google_protobuf//:wrappers_proto"],
-)
-
-# copybara:uncomment_begin(google-only)
-# py_proto_library(
-#     name = "tpu_embedding_configuration_py_pb2",
-#     visibility = ["//visibility:public"],
-#     deps = [":tpu_embedding_configuration_proto"],
-# )
-#
-# py_proto_library(
-#     name = "optimization_parameters_py_pb2",
-#     visibility = ["//visibility:public"],
-#     deps = [":optimization_parameters_proto"],
-# )
-# copybara:uncomment_end
diff --git a/third_party/xla/xla/tpu/protobuf/optimization_parameters.proto b/third_party/xla/xla/tpu/protobuf/optimization_parameters.proto
deleted file mode 100644
index 424110f3d89..00000000000
--- a/third_party/xla/xla/tpu/protobuf/optimization_parameters.proto
+++ /dev/null
@@ -1,657 +0,0 @@
-// Copyright 2026 The OpenXLA 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.
-// ==============================================================================
-
-syntax = "proto3";
-
-package tensorflow.tpu;
-
-import "google/protobuf/wrappers.proto";
-import "xla/service/hlo.proto";
-
-message ClippingLimits {
-  google.protobuf.FloatValue lower = 1;  // -inf if not set
-  google.protobuf.FloatValue upper = 2;  // +inf if not set
-}
-
-// Configuration for simulated quantization; simulated quantization is used to
-// reduce training/serving skew when the serving variables are quantized. The
-// same quantization operations are executed during training to minimize
-// differences with serving.
-//
-// Simulated quantization inserts the following operations on the forward pass
-// after gathering the embedding vector from HBM. The backward pass operations
-// are unchanged.
-//
-// clipped_val = clip(input, clipping_limits)
-// quantum = clipping_limits.range() / (num_buckets - 1)
-// quantized_val = floor((clipped_val - clipping_limits.lower()) / quantum + .5)
-// return quantized_val * quantum + clipping_limits.lower().
-message SimulatedQuantization {
-  // Whether simulated quantization is enabled.
-  bool enabled = 1;
-
-  // Minimum and maximum values of the range used for quantization.
-  ClippingLimits clipping_limits = 2;
-
-  // Number of possible quantized values.
-  int32 num_buckets = 3;
-}
-
-// Dynamic input specification for optimizers in the TPUEmbeddingConfiguration.
-// The actual dynamic inputs are provided as a scalar input list to the
-// SendTPUEmbeddingGradients Op indexed by their tag specified through the
-// following proto.
-message OptimizerDynamicInput {
-  // For tables where dynamic inputs are needed (e.g., learning rates or other
-  // dynamic hyperparameters used in optimizers), a tag must be specified for
-  // the input.
-  //
-  // The tag must be a non-negative  integer. All tags in the range
-  // [0, number_of_unique_tags) must be present in the TPU embedding
-  // configuration, i.e. a tag cannot be skipped if a different tag numerically
-  // greater than it is used in the configuration.
-  //
-  // If multiple tables specify the same tag, they *MUST* have
-  // the same dynamic input, for example, their dynamic learning rate could be
-  // computed by the same TensorFlow sub-graph. The partitioning of the
-  // embedding layer would be more optimal if the number_of_unique_tags is as
-  // *LOW* as possible, i.e., if many tables share the same tag.
-  //
-  // The hyper_parameters input of the SendTPUEmbeddingGradients op is used to
-  // communicate dynamic hyper-parameters to the TPU embedding program.
-  // The hyper_parameters input is a list of scalars where the size of the list
-  // is equal to the number of unique tags. The hyper-parameter associated with
-  // a particular tag is specified by populating its corresponding index in the
-  // list of scalars.
-  int32 tag = 1;
-}
-
-// Source of learning rate to use.
-message LearningRate {
-  oneof learning_rate {
-    float constant = 1;
-    OptimizerDynamicInput dynamic = 2;
-  }
-}
-
-// Each optimizer's parameter proto has a link to its documentation and CPU
-// implementation (if available) for user reference.
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adagrad
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L1634
-message AdagradParameters {
-  // Old initial accumulator parameter.
-  reserved "initial_accumulator";
-  reserved 1;
-}
-
-// This optimizer combines the Adagrad and Momentum update rules.
-// accum(new) = beta2 == 1.0 ?
-//              accum(old) + grad^2 :
-//              beta2 * accum(old) + (1 - beta2) * grad^2
-// accum_with_exponent = (accum(new) + epsilon)^(-1.0 / exponent)
-// mom_accum(new) = momentum * mom_accum(old) + accum_with_exponent
-// update = use_nesterov ?
-//          momentum * mom_accum(new) + accum_with_exponent :
-//          mom_accum(new)
-// var(new) = var(old) - lr * grad * update
-// Algorithm described in https://arxiv.org/abs/2002.11803.
-message AdagradMomentumParameters {
-  // Moving average parameter for the momentum accumulator.
-  float momentum = 1;
-
-  // Whether to use the Nesterov variant of momentum.
-  bool use_nesterov = 2;
-
-  // Exponent for the gradient^2 accumulator.
-  float exponent = 3;
-
-  // Moving average parameter for the gradient^2 accumulator.
-  float beta2 = 4;
-
-  // Offset added to the Adagrad accumulator.
-  float epsilon = 5;
-}
-
-// Algorithm in http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf.
-message BoundedAdagradParameters {
-  // Whether to use the updated or the old value of the accumulator when
-  // computing the effective learning rate. When update_accumulator_first is set
-  // to True, the updated value of the accumulator is used.
-  bool update_accumulator_first = 1;
-
-  // The max_var_update value to use. Set value to 0 (default) to disable using
-  // max_var_update to clip the gradient.
-  float max_var_update = 2;
-
-  // The maximum value of the accumulator. Set max_accumulator to 0 (default)
-  // to disable using max_accumulator to clip the accumulator.
-  float max_accumulator = 3;
-}
-
-// Frequency Aware Adagrad optimizer. This optimizer implements the AdaGrad
-// algorithm and further allows to:
-// * Scale the learning rate based on frequency of the update. Sparsely updated
-//   rows are updated with a higher effective learning rate, and frequently
-//   updated rows are updated with a lower effective learning rate.
-// * Decay the growth of the accumulator values.
-// * Use L1 / L2 regularization for the weight updates.
-//
-// The optimization algorithm is shown below.
-// counter(new) = counter(old) + 1
-// accum(new) = max(accumulator_decay * accum(old) + grad^2,
-//                  initial_accumulator_value)
-// lr_scale = min((step_counter / accum(new)) ^ probability_exponent,
-// max_lr_multiplier) update = grad * lr_scale / sqrt(accum(new)) if
-// (l1_regularization_strength > 0.0):
-//   update = update + l1_regularization_strength * sign(var(old))
-// if (l2_regularization_strength > 0.0):
-//   update = update + l2_regularization_strength * var(old)
-// var(new) = var(old) - lr_scale * grad * update
-
-message FrequencyAwareAdagradParameters {
-  // The L1 regularization parameter for adjusting the update based on the sign
-  // of the variable.
-  float l1_regularization_strength = 1;
-
-  // The L2 regularization parameter for adjusting the update based on the
-  // variable.
-  float l2_regularization_strength = 2;
-
-  // The exponent used for scaling the learning rate based on the sparsity of
-  // updates.
-  float probability_exponent = 4;
-
-  // The maximum value of the learning rate scale.
-  float max_lr_multiplier = 3;
-
-  // The decay for the Adagrad accumulator.
-  float accumulator_decay = 5;
-
-  // The initial and minimum value for the Adagrad accumulator.
-  float initial_accumulator_value = 6;
-
-  // The tag for identifying the step counter used for the frequency aware
-  // Adagrad optimizer.
-  OptimizerDynamicInput step_counter = 7;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L629
-message StochasticGradientDescentParameters {}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Ftrl
-// https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/41159.pdf
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L2646
-//
-// The hyperparameters for FTRL are the same as for the Keras implementation,
-// with some additions. The "beta" parameter matches the behavior described in
-// the second link above; "beta" / (2 * learning rate) should be added to "l2"
-// to get equivalent behavior in the other TensorFlow implementations of this
-// optimizer. When the multiply_linear_by_lr field is set to true, a modified
-// formula is used for FTRL that treats the "linear" accumulator as being
-// pre-multiplied by the learning rate (i.e., the accumulator named "linear"
-// actually stores "linear * learning_rate"). Other than checkpoint
-// compatibility, this is mathematically equivalent for a static learning rate;
-// for a dynamic learning rate, it is nearly the same as long as the learning
-// rate does not change quickly. The benefit of setting multiply_linear_by_lr to
-// true is that the modified formula handles zero and near-zero learning rates
-// without producing NaNs, improving flexibility for learning rate ramp-up.
-message FtrlParameters {
-  float l1 = 1;
-  float l2 = 2;
-  float lr_power = 3;
-  float beta = 7;
-  bool multiply_linear_by_lr = 6;
-
-  // Previously, allow_zero_accumulator parameter changed some internal formulas
-  // to allow zero and near-zero accumulator values at the cost of some
-  // performance. The current implementation ignores this parameter; zero or
-  // near-zero accumulator values are now always supported.
-  bool allow_zero_accumulator = 8 [deprecated = true];
-
-  // Old initial accumulator parameters.
-  reserved "initial_accum", "initial_linear";
-  reserved 4, 5;
-}
-
-// The Adam optimizer does not implement hyper-parameter update due to hardware
-// limitations; use the dynamic learning rate feature instead, setting the
-// learning rate to: user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
-// Here, t is the current timestep.
-//
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adam
-// https://github.com/tensorflow/tensorflow/blob/ab51450c817674c8ff08a7ae4f8ac50cdc4bed8b/tensorflow/python/training/adam.py#L32
-//
-// Note that the code by default implements the lazy version of Adam
-// (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/LazyAdamOptimizer)
-// unless the use_non_lazy_adam parameter is set, in which case it implements
-// the normal version of Adam that updates all parameters in the embedding
-// table, even for entries that are not used in the current minibatch
-// (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/AdamOptimizer). If
-// use_non_lazy_adam is enabled, gradient accumulation is also required to be
-// enabled in order to get correct results; a warning will be printed otherwise
-// (which may change to an error in the future). If use_sum_inside_sqrt is set,
-// the Adam variable update formula will be changed from m / (sqrt(v) + epsilon)
-// to m / sqrt(v + epsilon**2); this option improves the performance of TPU
-// training and is not expected to harm model quality.
-message AdamParameters {
-  float beta1 = 3;
-  float beta2 = 4;
-  float epsilon = 5;
-  bool use_non_lazy_adam = 8;
-  bool use_sum_inside_sqrt = 10;
-
-  // Old initial accumulator parameters.
-  reserved "initial_m", "initial_v";
-  reserved 6, 7;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L3068
-message MomentumParameters {
-  float momentum = 1;
-  bool use_nesterov = 2;
-
-  // Old initial accumulator parameter.
-  reserved "initial_accum";
-  reserved 3;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Lion
-// momenta(new) = beta2 * momenta(old) + (1 - beta2) * grad
-// momenta_t = beta1 * momenta(old) + (1 - beta1) * grad
-// var(new) = var(old) - lr * sign(momenta_t)
-// Algorithm described in https://arxiv.org/abs/2302.06675.
-message LionParameters {
-  float beta1 = 1;
-  float beta2 = 2;
-  bool use_non_lazy_lion = 3;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L4229
-message RmsPropParameters {
-  float rho = 1;
-  float momentum = 2;
-  float epsilon = 3;
-
-  // Old initial accumulator parameters.
-  reserved "initial_ms", "initial_mom";
-  reserved 4, 5;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L4358
-message CenteredRmsPropParameters {
-  float rho = 1;
-  float momentum = 2;
-  float epsilon = 3;
-
-  // Old initial accumulator parameters.
-  reserved "initial_ms", "initial_mom", "initial_mg";
-  reserved 4, 5, 6;
-}
-
-// Variant of algorithm in http://proceedings.mlr.press/v44/shamir15.pdf
-message MdlAdagradLightParameters {
-  float l2 = 1;
-  float lr_power = 2;
-  float min_servable_mdl_benefit = 3;
-  float mdl_mix_in_margin = 4;
-  float mdl_benefit_rampup_coeff = 5;
-  float mdl_min_weight = 6;
-  float benefit_revisit_scale = 7;
-  float max_event_benefit = 8;
-  float max_total_benefit = 9;
-  float mdl_hard_limit = 10;
-  bool hard_limit_min_benefit = 11;
-  bool mdl_regularize = 12;
-
-  // Old initial accumulator parameters.
-  reserved "initial_accumulator", "initial_weight", "initial_benefit";
-  reserved 13, 14, 15;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adadelta
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L933
-message AdadeltaParameters {
-  float rho = 1;
-  float epsilon = 2;
-
-  // Old initial accumulator parameters.
-  reserved "initial_accumulator", "initial_update";
-  reserved 3, 4;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/compat/v1/train/ProximalAdagradOptimizer
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L1961
-message ProximalAdagradParameters {
-  float l1 = 1;
-  float l2 = 2;
-
-  // Old initial accumulator parameter.
-  reserved "initial_accumulator";
-  reserved 3;
-}
-
-// The online Yogi optimizer does not implement hyper-parameter update; use the
-// dynamic learning rate feature instead, setting the learning rate to:
-// user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
-// Here, t is the current timestep.
-//
-// https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf
-// plus some extensions based on FTRL.
-//
-// Note that the code by default implements the lazy version of online Yogi.
-message OnlineYogiParameters {
-  // The L1 regularization parameter (used analogously to the one in FTRL).
-  float l1 = 1;
-
-  // The L2 regularization parameter (used analogously to the one in FTRL).
-  float l2 = 2;
-
-  // \beta_2 from Algorithm 2 in the paper.
-  float beta2 = 3;
-
-  // Reserved ids corresponding to removed tanh activation.
-  reserved 6;  // sign
-  reserved 7;  // tanh
-}
-
-// The online Yogi optimizer does not implement hyper-parameter update; use the
-// dynamic learning rate feature instead, setting the learning rate to:
-// user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
-// Here, t is the current timestep.
-//
-// https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf
-// plus some extensions based on FTRL.
-//
-// Note that the code by default implements the lazy version of proximal Yogi.
-message ProximalYogiParameters {
-  // The L1 regularization parameter.
-  float l1 = 1;
-
-  // The L2 regularization parameter.
-  float l2 = 2;
-
-  // The exponential decay rate for the 1st moment estimates.
-  float beta1 = 3;
-
-  // The exponential decay rate for the 2nd moment estimates.
-  float beta2 = 4;
-
-  // A constant trading off adaptivity and noise.
-  float epsilon = 5;
-
-  // Reserved ids corresponding to removed tanh activation.
-  reserved 8;  // sign
-  reserved 9;  // tanh
-}
-
-// Estimator for the frequency of updates to a lookup table. It maintains an
-// array (tf.Variable) D, where each element records the average number of
-// global steps between two consecutive batches that hit the corresponding
-// bucket. Once an item with bucket id i is sampled, D[i] is updated by:
-//   D[i] <- D[i] * (1 - tau) + delta[i] * tau,
-//
-// where tau is a learning rate between 0 and 1 (exclusive), and
-//   delta[i] = current global step - last step i is sampled.
-//
-// The estimated frequency (sampling rate in a batch) is thus 1 / D[i].
-//
-// Elements in D are initialized with a large value max_delta. delta[i] will
-// also be capped by this value.
-//
-// The exact sequence of operations used in the optimizer is shown below.
-// last_hit_step[i] is a tf.Variable that holds the last global step at which i
-// was sampled.
-//
-//   delta = global_step - last_hit_step[i]
-//   clipped_delta = min(delta, params.max_delta)
-//   is_outlier = (delta >= params.outlier_threshold * D[i])
-//   D[i] <- is_outlier ? clipped_delta
-//                      : D[i] * (1 - params.tau) + clipped_delta * params.tau
-//   last_hit_step[i] <- global_step
-message FrequencyEstimatorParameters {
-  // Learning rate between (0, 1) that is used to update the array D.
-  float tau = 1;
-
-  // Maximum value of delta: difference between the current global step and the
-  // last global step at which the row was sampled.
-  float max_delta = 2;
-
-  // Threshold used to determine whether the current update is an outlier.
-  float outlier_threshold = 3;
-
-  // The weight exponent used to transform the estimated delta into weights.
-  // The transformation function is: (delta / max_delta) ^ (weight_exponent)
-  float weight_exponent = 4;
-}
-
-// A user-defined optimizer.
-// The contained HLO program must take the following arguments in the following
-// order:
-// 1.  gradients
-// 2.  table weights
-// 3.  slot variables
-// 4.  an optional scalar input that is passed in via the dynamic learning
-//     rate mechanism.
-//
-// It must return/end in a tuple op that contains the following values in the
-// following order:
-// 1.  new table values
-// 2.  new slot variable value
-//
-// The program must have shape (1,1) with dtype float32 throughout and only use
-// HLO that operate elementwise (e.g., no reduce, no variables, no control flow
-// and no broadcasting outside of the single scalar input).
-// The HLO program should be written as if it were a dense update. It will be
-// called on each row that needs an update and will applied elementwise.
-message UserDefinedProgramParameters {
-  xla.HloModuleProto program = 1;
-  reserved 2;  // Was padding_values
-}
-
-// Optimizer that just sets the variable to the value of the gradient. To be
-// correct, this requires either gradient accumulation (to sum the values of a
-// computed expression across the samples) or to deduplicate IDs within a single
-// host (to assign the value from an arbitrary sample).
-message AssignParameters {}
-
-// Status of using gradient accumulation (doing two passes over the input
-// gradients: one to accumulate them into a temporary array and another to apply
-// them using the actual optimization algorithm). The extra message is to wrap
-// the enum for scoping.
-message GradientAccumulationStatus {
-  // if UNSPECIFIED (default), gradient accumulation is ENABLED.
-  enum Status {
-    UNSPECIFIED = 0;
-    ENABLED = 1;
-    DISABLED = 2;
-  }
-}
-
-// Whether to optimize the packing of low-dimensional embedding tables in HBM
-// (high bandwidth memory). TPUs access HBM at 32-byte (8-float) granularity.
-// For functional correctness, the TPU software internally pads the embedding
-// dimension to a multiple of 8. This can sometimes lead to significant memory
-// wastage due to padding. For 1-dimensional, 2-dimensional, and 4-dimensional,
-// the TPU software can remove this padding by packing multiple rows into the
-// same 8-float HBM chunk. For example, 8 rows could be packed into the same
-// 8-float chunk for a 1-dimensional embedding table.
-
-// There is one important limitation for this HBM packing though. When only a
-// subset of rows in an 8-float chunk are accessed on a particular step, the
-// adjoining rows in the same chunk are updated with zero gradients on the
-// backward pass even if they are not touched. This is an artifact of the
-// packing implementation. This operation is NOT functionally correct for
-// optimizers where zero gradients change the embeddings/slot-variable values,
-// e.g., momentum-based optimizers. Hence, this HBM packing cannot be enabled
-// for embedding tables with such optimizers. The TPU software automatically
-// recognizes that a zero gradient can modify state and turns off the low
-// dimensional embedding packing in that scenario.
-//
-// However, for optimizers where a zero gradient is a NoOp, such as SGD,
-// Adagrad, and FTRL, this packing optimization can be used. However, there are
-// some important considerations:
-// * Clipping limits: The initial values for such embeddings should fall within
-//   the clipping limits specified in the optimization parameters. Otherwise, a
-//   zero gradient will cause the embeddings to be clipped. This changes state
-//   and hence, is not a NoOp.
-// * FTRL: The embedding vector is computed directly from the values of the
-//   accumulator and linear slot variables. Hence, the initial embedding values
-//   should match that computed from the initial values of the accumulator and
-//   linear slot variables. Note that in nearly all cases, the linear value is
-//   initialized to zero; this corresponds to an embedding value of zero.
-//
-// Performance: The TPU has to perform additional work when low dimensional
-// packing is enabled. In certain situations when the vocabulary size is small,
-// it may not make sense to turn on this packing since the total memory usage
-// due to padding is extremely low. Hence, the TPU software automatically turns
-// off the packing optimization in such scenarios.
-message LowDimensionalPackingStatus {
-  // if UNSPECIFIED (default), the low dimension packing status is DISABLED.
-  // This can change in future.
-  //
-  // if ENABLED, the low dimension packing is enabled only if the following
-  // three additional conditions are true:
-  //  * The optimizer treats the zero gradient as a NoOp.
-  //  * The embedding dimension is 1, 2, or 4.
-  //  * The vocabulary size is large enough to avoid performance issues.
-  //
-  // if DISABLED, the low dimension packing is always disabled.
-  enum Status {
-    UNSPECIFIED = 0;
-    ENABLED = 1;
-    DISABLED = 2;
-  }
-}
-
-// Configuration proto for hot ID optimization. This is an experimental feature
-// that is currently disabled (by default).
-message HotIdReplicationConfiguration {
-  // Whether to enable or disable hot ID optimization.
-  // If set to UNSPECIFIED (default), hot ID optimization is DISABLED.
-  // If set to ENABLED, hot ID replication is turned ON.
-  // If set to MIGRATION_ONLY, hot ID migration is turned ON.
-  enum Status {
-    UNSPECIFIED = 0;
-    ENABLED = 1;
-    DISABLED = 2;
-    MIGRATION_ONLY = 3;
-  }
-  Status status = 1;
-}
-
-message OptimizationParameters {
-  // Learning rate used for updating the embedding layer parameters.
-  LearningRate learning_rate = 13;
-
-  // Limits to which to clip the weight values after the backward pass; not
-  // present means no limits are applied.
-  ClippingLimits clipping_limits = 2;
-
-  // Limits to which to clip the backward pass gradient before using it for
-  // updates; not present means no limits are applied.
-  ClippingLimits gradient_clipping_limits = 7;
-
-  // Amount of weight decay to apply; see weight_decay_optimizers.py for
-  // details. All optimizers except MDL Adagrad Light are supported with this
-  // option. Although there is no check, users who want weight decay will also
-  // want to ensure that gradient accumulation is enabled so that the decay will
-  // happen once per global batch.
-  float weight_decay_factor = 16;
-
-  // If true, the weight decay factor is multiplied by the current learning rate
-  // before use; this is to match the note in DecoupledWeightDecayExtension in
-  // weight_decay_optimizers.py.
-  bool multiply_weight_decay_factor_by_learning_rate = 22;
-
-  // Configuration for simulated quantization which is used to reduce
-  // training/serving skew when the serving variables are quantized. The same
-  // quantization operations are executed during training to minimize
-  // differences with serving.
-  SimulatedQuantization simulated_quantization = 27;
-
-  // Status of using gradient accumulation (doing two passes over the input
-  // gradients: one to accumulate them into a temporary array and another to
-  // apply them using the actual optimization algorithm).
-  GradientAccumulationStatus.Status gradient_accumulation_status = 17;
-
-  // Status of the low-dimensional embedding packing optimization. This controls
-  // whether to optimize the packing of 1-dimensional, 2-dimensional, and
-  // 4-dimensional embedding tables in memory.
-  LowDimensionalPackingStatus.Status low_dimensional_packing_status = 28;
-
-  // Configuration proto for hot ID replication. This is an experimental
-  // feature that is currently disabled (by default).
-  HotIdReplicationConfiguration hot_id_replication_configuration = 18;
-
-  // Optimization algorithm parameters; which field is selected determines which
-  // algorithm to use.
-  oneof parameters {
-    AdagradParameters adagrad = 3;
-    AdagradMomentumParameters adagrad_momentum = 26;
-    BoundedAdagradParameters bounded_adagrad = 19;
-    FrequencyAwareAdagradParameters frequency_aware_adagrad = 30;
-    StochasticGradientDescentParameters stochastic_gradient_descent = 4;
-    FtrlParameters ftrl = 5;
-    AdamParameters adam = 6;
-    MomentumParameters momentum = 8;
-    LionParameters lion = 29;
-    RmsPropParameters rms_prop = 9;
-    CenteredRmsPropParameters centered_rms_prop = 10;
-    MdlAdagradLightParameters mdl_adagrad_light = 11;
-    AdadeltaParameters adadelta = 12;
-    ProximalAdagradParameters proximal_adagrad = 14;
-    OnlineYogiParameters online_yogi = 20;
-    ProximalYogiParameters proximal_yogi = 21;
-    FrequencyEstimatorParameters frequency_estimator = 23;
-    UserDefinedProgramParameters user_defined_program = 24;
-    AssignParameters assign = 25;
-  }
-
-  reserved 1, 15;
-
-  // NEXT_ID: 31
-}
-
-// Specification of an optimization algorithm's state variables (both the main
-// value vector and any extra accumulators, etc.). This proto is only used
-// internally by the TPU software and is not exposed directly to the TF model.
-message StateVariableSpecification {
-  // Parameter name for the state variable.
-  string name = 1;
-
-  // A normal state variable that should be saved and restored in checkpoints
-  // and used as an input or output to non-debug TensorFlow ops.
-  message UserDefined {
-    reserved 1;  // Was padding_initial_value.
-  }
-
-  // A state variable that should be filled with a constant and normally hidden
-  // from users (used for intermediate gradients being accumulated, for
-  // example).
-  message FillWithConstant {
-    double initial_value = 1;
-  }
-
-  // Usage type of this state variable.
-  oneof usage {
-    UserDefined user_defined = 2;
-    FillWithConstant fill_with_constant = 3;
-  }
-}
diff --git a/third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration.proto b/third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration.proto
deleted file mode 100644
index 14f64bcc391..00000000000
--- a/third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration.proto
+++ /dev/null
@@ -1,176 +0,0 @@
-// 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.
-// ==============================================================================
-
-syntax = "proto3";
-
-package tensorflow.tpu;
-
-import "xla/tpu/protobuf/optimization_parameters.proto";
-
-message TPUEmbeddingConfiguration {
-  // Description of the various embedding tables.
-  message TableDescriptor {
-    // Name of the table.
-    string name = 1;
-
-    // Size of the vocabulary (i.e., number of rows) in the table.
-    int64 vocabulary_size = 2;
-
-    // The embedding dimension (i.e., the width of the embedding table).
-    int32 dimension = 3;
-
-    // Number of features mapped to this table.
-    int32 num_features = 4;
-
-    // Details of the learning algorithm used to update the embedding
-    // parameters.
-    OptimizationParameters optimization_parameters = 5;
-  }
-  repeated TableDescriptor table_descriptor = 1;
-
-  // Mode. Should the embedding layer program be run for inference (just forward
-  // pass), training (both forward and backward pass) or just the backward_pass.
-  enum Mode {
-    UNSPECIFIED = 0;
-    INFERENCE = 1;
-    TRAINING = 2;
-    BACKWARD_PASS_ONLY = 3;
-  }
-  Mode mode = 2;
-
-  // Number of samples in each batch of embedding layer activations sent to
-  // the TensorCore.
-  int32 batch_size_per_tensor_core = 3;
-
-  // Number of TPU hosts used for inference/training.
-  int32 num_hosts = 4;
-
-  // Number of TensorCore used for inference/training.
-  int32 num_tensor_cores = 5;
-
-  // Sharding strategy of the embedding tables among the hosts.
-  // If the sharding_strategy is "mod", each id is assigned to host
-  // "id % num_hosts". For instance, 13 ids are split across 5 hosts as:
-  // [[0, 5, 10], [1, 6, 11], [2, 7, 12], [3, 8], [4, 9]].
-  // If the sharding_strategy is "div", ids are assigned to hosts in a
-  // contiguous manner. In this case, 13 ids are split across 5 hosts as:
-  // [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11, 12]].
-  // In both the strategies, if the id space does not evenly divide the number
-  // of hosts, each of the first "table_descriptor.vocabulary_size % num_hosts"
-  // hosts will be assigned one more id.
-  // This partitioning strategy exactly follows that in the embedding_lookup
-  // TensorFlow function at tensorflow/python/ops/embedding_ops.py.
-  enum ShardingStrategy {
-    DIV_DEFAULT = 0;
-    MOD = 1;
-  }
-  ShardingStrategy sharding_strategy = 6;
-
-  // This parameter determines if the execution of the sparse core will be
-  // pipelined with that of the TensorCore. This parameter only affects results
-  // when mode=TRAINING. If mode=INFERENCE or BACKWARD_PASS_ONLY, this parameter
-  // does not affect execution and hence, is a don't care value.
-  //
-  // false: The execution of the sparse core is not pipelined with that of the
-  // TensorCore. The forward pass of every step on the sparse core is executed
-  // only after the backward pass of the previous step is complete. And the
-  // backward pass on the sparse core is executed only after the embedding
-  // gradients have been computed on the TensorCore on every step. This ensures
-  // that the activations on every step observe the gradient updates from the
-  // previous step on both the sparse core and the TensorCore.
-  //
-  // true: The execution of the sparse core is pipelined with that of the
-  // TensorCore. The forward pass of every step on the sparse core can be
-  // executed after the forward pass of the previous step is complete without
-  // waiting for the backward pass. This improves the utilization of the sparse
-  // core allowing it to process step N+1 while the embedding gradients for step
-  // N are computed on the TensorCore. The backward pass of every step on the
-  // sparse core is executed directly after the forward pass for the next step
-  // is complete. The drawback is that embedding activations for step N+1 do not
-  // observe the embedding gradient updates from step N. This could affect model
-  // quality if step N and N+1 involve the same set of embedding IDs. However,
-  // since the embedding updates are sparse, this is generally not considered a
-  // problem.
-  bool pipeline_execution_with_tensor_core = 7;
-
-  // Directory where embedding lookup statistics are stored. These statistics
-  // summarize information about the inputs to the embedding lookup
-  // operation, in particular, the average number of embedding IDs per example
-  // and how well the embedding IDs are load balanced across the system. The
-  // lookup statistics are used during TPU initialization for embedding table
-  // partitioning. Collection of lookup statistics is done at runtime by
-  // profiling the embedding inputs: only 3% of input samples are profiled to
-  // minimize host CPU overhead. Once a suitable number of samples are
-  // profiled, the lookup statistics are saved to table-specific files in the
-  // profile data directory generally at the end of a TPU training loop. The
-  // filename corresponding to each table is obtained by hashing table specific
-  // parameters (e.g., table name and number of features) and global
-  // configuration parameters (e.g., sharding strategy and TPU worker task
-  // count). The same profile data directory can be shared amongst several
-  // models to reuse embedding lookup statistics.
-  string profile_data_directory = 9;
-
-  // Description of different input features.
-  message FeatureDescriptor {
-    // Name of the input feature.
-    string name = 1;
-
-    // Index of the corresponding table in the TableDescriptor list.
-    int32 table_id = 2;
-
-    // Static shape of the inputs (excluding the reduction axis). Note that
-    // the shape of the actual inputs provided using the infeed op must be
-    // strictly smaller than input_shape. The outputs received at the TensorCore
-    // will have rank = input_shape.size() + 1. The innermost axis corresponds
-    // to the embedding dimension. If the input has shape [m, n, k] (excluding
-    // the reduction axis) and the embedding dimension is d, the output received
-    // at the TensorCore will have shape [m, n, k, d].
-    repeated int32 input_shape = 3;
-  }
-  // If the feature_descriptor field is populated, the model should NOT populate
-  // TableDescriptor.num_features and batch_size_per_tensor_core. These two
-  // fields will be auto-populated by the TPUEmbedding rewrite passes.
-  repeated FeatureDescriptor feature_descriptor = 10;
-
-  // SPMD (Single Program Multiple Data) sharding configuration for
-  // TPUEmbedding. When model parallelism is used on the TensorCore, the number
-  // of cores per replica must be passed to TPUEmbedding so that the right
-  // shapes can be computed in the TF/XLA bridge.
-  message SpmdSharding {
-    // Whether SPMD sharding is enabled.
-    bool enabled = 1;
-
-    // Number of cores per replica.
-    int32 num_cores_per_replica = 2;
-
-    // If true, the tensors are manually partitioned. Otherwise, use the
-    // automatic SPMD partitioning. This should be true when users use
-    // `shard_map`.
-    bool use_manual_partitioning = 3;
-
-    // If true, lowers to Shardy custom-call for manual partitioning.
-    // Default is false.
-    bool use_shardy_partitioner = 4;
-  }
-  SpmdSharding spmd_sharding = 11;
-
-  // Old TPU embedding output layout.
-  reserved "output_layout";
-  reserved 8;
-}
-
-// A placeholder message that is used to define a unique Status payload
-// URL for TPU embedding errors.
-message TPUEmbeddingError {}

From 19ef798295d8bf1573a630b0bf8af3a49a5e9347 Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Fri, 9 Oct 2026 13:13:15 -0700
Subject: [PATCH 0388/1000] [XLA] Fix kLt and kLe trip counts in
 MatchTrivialLoopTripCount

`MatchTrivialLoopTripCount` in `while_loop_analysis.cc` miscomputes trip counts for `kLt` and `kLe` loops with negative bounds or steps greater than one. In the `kLt` branch, the remainder adjustment checks `*trips < *while_cond_bound_val`, comparing the candidate trip count against the loop bound `N`. When the candidate trip count is at least `N`, such as when `N <= 0`, that check fails and undercounts by one. For zero trip `kLt` and `kLe` loops with `0 > N - init > -step`, integer division truncates toward zero and returns `1`.

Return `0` immediately when `*while_cond_bound_val <= *indvar_init_val` for `kLt` or `*while_cond_bound_val < *indvar_init_val` for `kLe`. For positive `kLt` spans, compute the trip count as `(*trips - 1) / trip_count_step + 1`. For `kLe` loops with `*trips >= 0`, compute the trip count as `CheckedAdd(*trips / trip_count_step, 1)`. Existing programs with `kLt` or `kLe` loops that have negative bounds or zero trips now receive exact static trip counts.

Tests:
- `//third_party/tensorflow/compiler/xla/hlo/analysis:while_loop_analysis_test`: `NegativeAndZeroTripCountNonUnitStep` verifies `kLt` and `kLe` trip counts for negative bounds and zero trip loops.
PiperOrigin-RevId: 996714020
---
 .../xla/hlo/analysis/while_loop_analysis.cc   | 32 +++++++------------
 .../hlo/analysis/while_loop_analysis_test.cc  | 18 +++++++++++
 2 files changed, 30 insertions(+), 20 deletions(-)

diff --git a/third_party/xla/xla/hlo/analysis/while_loop_analysis.cc b/third_party/xla/xla/hlo/analysis/while_loop_analysis.cc
index b54722a9118..dd03be9570c 100644
--- a/third_party/xla/xla/hlo/analysis/while_loop_analysis.cc
+++ b/third_party/xla/xla/hlo/analysis/while_loop_analysis.cc
@@ -16,7 +16,6 @@ limitations under the License.
 #include "xla/hlo/analysis/while_loop_analysis.h"
 
 #include <algorithm>
-#include <cmath>
 #include <cstdint>
 #include <map>
 #include <memory>
@@ -1027,26 +1026,16 @@ optional<int64_t> MatchTrivialLoopTripCount(const HloInstruction* while_op,
                 .WithOperand(0, m::Op().Is(while_cond_indvar)))) {
     VLOG(2) << "Pattern-match succeeded: loop condition is i < N: "
             << while_cond_root->ToString();
+    if (*while_cond_bound_val <= *indvar_init_val) {
+      return 0;
+    }
     optional<int64_t> trips =
         CheckedSubtract(*while_cond_bound_val, *indvar_init_val);
-    if (trips) {
-      const int64_t remainder = std::remainder(*trips, trip_count_step);
-      const int64_t div = std::floor(*trips / trip_count_step);
-      if (remainder == 0) {
-        return std::max(int64_t{0}, div);
-      }
-      trips = CheckedAdd(div, 1);
-      if (!trips) {
-        VLOG(2) << "Pattern-match failed: Trip count exceeds INT64_MAX.";
-        return nullopt;
-      }
-      if (*trips < *while_cond_bound_val) {
-        return std::max(int64_t{0}, *trips);
-      }
-      return std::max(int64_t{0}, div);
+    if (!trips) {
+      VLOG(2) << "Pattern-match failed: Trip count exceeds INT64_MAX.";
+      return nullopt;
     }
-    VLOG(2) << "Pattern-match failed: Trip count exceeds INT64_MAX.";
-    return nullopt;
+    return (*trips - 1) / trip_count_step + 1;
   }
 
   // Handle `i = init; i <= N; i+=k`.
@@ -1056,18 +1045,21 @@ optional<int64_t> MatchTrivialLoopTripCount(const HloInstruction* while_op,
                 .WithOperand(0, m::Op().Is(while_cond_indvar)))) {
     VLOG(2) << "Pattern-match succeeded: loop condition is i <= N: "
             << while_cond_root->ToString();
+    if (*while_cond_bound_val < *indvar_init_val) {
+      return 0;
+    }
     optional<int64_t> trips =
         CheckedSubtract(*while_cond_bound_val, *indvar_init_val);
     if (!trips) {
       VLOG(2) << "Pattern-match failed: Trip count exceeds INT64_MAX";
       return nullopt;
     }
-    trips = CheckedAdd(std::floor(*trips / trip_count_step), 1);
+    trips = CheckedAdd(*trips / trip_count_step, 1);
     if (!trips) {
       VLOG(2) << "Pattern-match failed: Trip count exceeds INT64_MAX";
       return nullopt;
     }
-    return std::max<int64_t>(0, *trips);
+    return *trips;
   }
 
   VLOG(2) << "Pattern-match failed: while condition follows unknown pattern: "
diff --git a/third_party/xla/xla/hlo/analysis/while_loop_analysis_test.cc b/third_party/xla/xla/hlo/analysis/while_loop_analysis_test.cc
index c6b0b3a2a83..f99819e5b9b 100644
--- a/third_party/xla/xla/hlo/analysis/while_loop_analysis_test.cc
+++ b/third_party/xla/xla/hlo/analysis/while_loop_analysis_test.cc
@@ -16,9 +16,11 @@ limitations under the License.
 #include "xla/hlo/analysis/while_loop_analysis.h"
 
 #include <cstdint>
+#include <initializer_list>
 #include <memory>
 #include <optional>
 #include <string>
+#include <tuple>
 #include <vector>
 
 #include <gmock/gmock.h>
@@ -423,6 +425,22 @@ TEST_F(WhileLoopAnalysisTest, ExactBoundTrivialTripCount) {
       CalculateTripCount(0, 40, 5, ComparisonDirection::kLe));
 }
 
+TEST_F(WhileLoopAnalysisTest, NegativeAndZeroTripCountNonUnitStep) {
+  for (auto [init, bound, step, dir] : std::initializer_list<
+           std::tuple<int32_t, int32_t, int32_t, ComparisonDirection>>{
+           {-5, 0, 2, ComparisonDirection::kLt},
+           {-10, -3, 2, ComparisonDirection::kLt},
+           {-2, 1, 2, ComparisonDirection::kLt},
+           {5, 2, 4, ComparisonDirection::kLt},
+           {-5, 0, 2, ComparisonDirection::kLe},
+           {1, 0, 2, ComparisonDirection::kLe},
+           {5, 2, 4, ComparisonDirection::kLe}}) {
+    ASSERT_OK_AND_ASSIGN(int64_t tc,
+                         MakeWhileLoopAndGetTripCount(init, bound, step, dir));
+    EXPECT_EQ(tc, CalculateTripCount(init, bound, step, dir));
+  }
+}
+
 TEST_F(WhileLoopAnalysisTest, NoAIVNoConstChain) {
   absl::string_view kHloModule = R"(
     HloModule ModuleWithWhile

From 56c0529fc9c1db67d4877b2160d12fc6f4a14af4 Mon Sep 17 00:00:00 2001
From: Parker Schuh <parkers@google.com>
Date: Fri, 9 Oct 2026 13:16:40 -0700
Subject: [PATCH 0389/1000] Add temporary overload to CommonPjRtClient for
 updating device assignment to remove the last need to subclass the
 CommonPjRtLoadedExecutable.

PiperOrigin-RevId: 996715961
---
 .../xla/xla/pjrt/common_pjrt_client.cc        | 30 ++++++++++---------
 third_party/xla/xla/pjrt/common_pjrt_client.h | 11 ++++---
 2 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.cc b/third_party/xla/xla/pjrt/common_pjrt_client.cc
index bd125f8448a..0ed4110ee4d 100644
--- a/third_party/xla/xla/pjrt/common_pjrt_client.cc
+++ b/third_party/xla/xla/pjrt/common_pjrt_client.cc
@@ -2665,32 +2665,33 @@ std::vector<std::unique_ptr<PjRtBuffer>> CommonPjRtClient::CreateOutputs(
   return res;
 }
 
-absl::StatusOr<CommonPjRtLoadedExecutable::DeviceAndAssignment>
-CommonPjRtLoadedExecutable::LookupDeviceAndAssignment(
+absl::StatusOr<PjRtExecutableLoadState::DeviceAndAssignment>
+CommonPjRtClient::LookupDeviceAndAssignment(
     const ExecuteOptions& options, int replica, int partition,
-    PjRtDevice* device) const {
-  DeviceAndAssignment result;
+    PjRtDevice* device,
+    const std::shared_ptr<DeviceAssignment>& device_assignment,
+    PjRtExecutableLoadState* load_state) const {
+  PjRtExecutableLoadState::DeviceAndAssignment result;
   if (device == nullptr) {
-    if (device_assignment_ == nullptr) {
+    if (device_assignment == nullptr) {
       return InvalidArgument(
           "device_assignment_ must be set if device is not provided.");
     }
-    const int64_t device_id = (*device_assignment_)(replica, partition);
+    const int64_t device_id = (*device_assignment)(replica, partition);
     GlobalDeviceId global_device_id(device_id);
-    ABSL_ASSIGN_OR_RETURN(device, client()->LookupDevice(global_device_id));
-    result.device_assignment = device_assignment_;
+    ABSL_ASSIGN_OR_RETURN(device, LookupDevice(global_device_id));
+    result.device_assignment = device_assignment;
   } else {
-    if (device_assignment_ != nullptr) {
+    if (device_assignment != nullptr) {
       return InvalidArgument(
           "device_assignment_ must not be set if device is provided.");
     }
     CHECK_EQ(replica, 0);
     CHECK_EQ(partition, 0);
-    CHECK(addressable_devices_.empty());
     result.device_assignment = std::make_shared<DeviceAssignment>(1, 1);
     (*result.device_assignment)(0, 0) = device->id();
   }
-  CHECK_EQ(device->process_index(), client()->process_index());
+  CHECK_EQ(device->process_index(), process_index());
   result.device = device;
   result.local_device_id = device->local_device_id();
   result.global_device_id = device->global_device_id();
@@ -2706,9 +2707,10 @@ absl::Status CommonPjRtLoadedExecutable::ExecutePrepare(
     int replica, int partition, const ExecuteOptions& options,
     size_t host_callback_idx, PjRtDevice* device, int attempt) const {
   tsl::profiler::TraceMe traceme("CommonPjRtLoadedExecutable::ExecutePrepare");
-  ABSL_ASSIGN_OR_RETURN(
-      DeviceAndAssignment device_and_assign,
-      LookupDeviceAndAssignment(options, replica, partition, device));
+  ABSL_ASSIGN_OR_RETURN(DeviceAndAssignment device_and_assign,
+                   client()->LookupDeviceAndAssignment(
+                       options, replica, partition, device, device_assignment_,
+                       load_state_.get()));
   device = device_and_assign.device;
   // Fill in device to launch_args so it will be present even if ExecutePrepare
   // fails with OOM.
diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.h b/third_party/xla/xla/pjrt/common_pjrt_client.h
index 27bbc4929d5..2eeffdb8254 100644
--- a/third_party/xla/xla/pjrt/common_pjrt_client.h
+++ b/third_party/xla/xla/pjrt/common_pjrt_client.h
@@ -130,6 +130,13 @@ class CommonPjRtClient : public PjRtClient {
     return false;
   }
 
+  virtual absl::StatusOr<PjRtExecutableLoadState::DeviceAndAssignment>
+  LookupDeviceAndAssignment(
+      const ExecuteOptions& options, int replica, int partition,
+      PjRtDevice* device,
+      const std::shared_ptr<DeviceAssignment>& device_assignment,
+      PjRtExecutableLoadState* load_state) const;
+
   absl::StatusOr<std::unique_ptr<HloCostAnalysis>> GetHloCostAnalysis()
       const override;
 
@@ -817,10 +824,6 @@ class CommonPjRtLoadedExecutable : public PjRtLoadedExecutable {
 
   using DeviceAndAssignment = PjRtExecutableLoadState::DeviceAndAssignment;
 
-  virtual absl::StatusOr<DeviceAndAssignment> LookupDeviceAndAssignment(
-      const ExecuteOptions& options, int replica, int partition,
-      PjRtDevice* device) const;
-
   virtual absl::StatusOr<std::unique_ptr<PjRtRawLoadedExecutable>>
   LoadRawExecutable(const ExecuteOptions& options, size_t host_callback_idx,
                     xla::RunId run_id, DeviceAndAssignment device_and_assign,

From 6c33de97203cc70d752ac726c6d092a5eea15651 Mon Sep 17 00:00:00 2001
From: Bryan Massoth <bmassoth@google.com>
Date: Fri, 9 Oct 2026 13:54:37 -0700
Subject: [PATCH 0390/1000] Fix tsan issues in threadpool_listener.cc Remove
 unnecessary pointer hold to unique_ptr wrapped EventCollector in executor.cc

PiperOrigin-RevId: 996738272
---
 tensorflow/core/common_runtime/BUILD           |  3 ++-
 tensorflow/core/common_runtime/executor.cc     | 18 ++++++++----------
 .../backends/cpu/threadpool_listener.cc        |  6 ------
 3 files changed, 10 insertions(+), 17 deletions(-)

diff --git a/tensorflow/core/common_runtime/BUILD b/tensorflow/core/common_runtime/BUILD
index ed2680bb2aa..eb6715d4603 100644
--- a/tensorflow/core/common_runtime/BUILD
+++ b/tensorflow/core/common_runtime/BUILD
@@ -717,14 +717,15 @@ cc_library(
         "//tensorflow/core/platform:error_logging",
         "//tensorflow/core/profiler/lib:annotated_traceme",
         "//tensorflow/core/profiler/lib:connected_traceme",
+        "//tensorflow/core/profiler/lib:context_types_hdrs",
         "//tensorflow/core/profiler/lib:scoped_annotation",
+        "//tensorflow/core/profiler/lib:traceme",
         "//tensorflow/core/profiler/lib:traceme_encode",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
         "@com_google_absl//absl/types:optional",
-        "@xla//xla/tsl/platform:env",
     ],
     alwayslink = 1,
 )
diff --git a/tensorflow/core/common_runtime/executor.cc b/tensorflow/core/common_runtime/executor.cc
index 65aec1d74bf..615c58f47ab 100644
--- a/tensorflow/core/common_runtime/executor.cc
+++ b/tensorflow/core/common_runtime/executor.cc
@@ -27,6 +27,7 @@ limitations under the License.
 #include "absl/strings/str_join.h"
 #include "absl/time/time.h"
 #include "absl/types/optional.h"
+#include "xla/tsl/platform/macros.h"
 #include "tensorflow/core/activity_watcher/activity.h"
 #include "tensorflow/core/common_runtime/costmodel_manager.h"
 #include "tensorflow/core/common_runtime/entry.h"
@@ -419,7 +420,6 @@ class ExecutorState {
   // Step-local container.
   ScopedStepContainer* step_container_;
   StepStatsCollectorInterface* const stats_collector_;
-  const tsl::tracing::EventCollector* const event_collector_;
   Context context_;
 
   // QUESTION: Make it a checkpoint::TensorSliceReaderCacheWrapper
@@ -473,8 +473,6 @@ ExecutorState<PropagatorStateType>::ExecutorState(
       tensor_store_(args.tensor_store),
       step_container_(args.step_container),
       stats_collector_(args.stats_collector),
-      event_collector_(tsl::tracing::GetEventCollector(
-          tsl::tracing::EventCategory::kCompute)),
       context_(ContextKind::kThread),
       slice_reader_cache_(new checkpoint::TensorSliceReaderCacheWrapper),
       call_frame_(args.call_frame),
@@ -601,16 +599,16 @@ struct ExecutorState<PropagatorStateType>::AsyncState {
 
 // Returns true if `item` might be traced by the given trace and event
 // collectors. Returns false only if `item` definitely will not be traced.
-bool MightTrace(const tsl::tracing::EventCollector* event_collector,
-                bool is_expensive) {
-  // Tracing will only be enabled if either `event_collector` is non null,
-  // or `trace_collector` is non-null and enabled for this particular kernel.
-  // Although `profiler::TraceMe`, `profiler::ScopedAnnotation`, and
+bool MightTrace(bool is_expensive) {
+  // Tracing will only be enabled if either
+  // `tsl::tracing::EventCollector::IsEnabled()` is true, or `trace_collector`
+  // is non-null and enabled for this particular kernel. Although
+  // `profiler::TraceMe`, `profiler::ScopedAnnotation`, and
   // `tsl::tracing::ScopedRegion` check subsets of these properties internally
   // in their constructors, the cost of passing the necessary arguments to them
   // can be significant, so we avoid constructing them in the common case (when
   // we know they will not be used).
-  if (event_collector != nullptr) {
+  if (tsl::tracing::EventCollector::IsEnabled()) {
     return true;
   }
 
@@ -632,7 +630,7 @@ absl::Status ExecutorState<PropagatorStateType>::ProcessSync(
   Device* device = immutable_state_.params().device;
   const bool is_expensive = kernel_stats_->IsExpensive(item);
 
-  if (TF_PREDICT_FALSE(MightTrace(event_collector_, is_expensive))) {
+  if (TF_PREDICT_FALSE(MightTrace(is_expensive))) {
     tsl::tracing::ScopedRegion region(tsl::tracing::EventCategory::kCompute,
                                       op_kernel->name_view());
     profiler::AnnotatedTraceMe activity(
diff --git a/third_party/xla/xla/tsl/profiler/backends/cpu/threadpool_listener.cc b/third_party/xla/xla/tsl/profiler/backends/cpu/threadpool_listener.cc
index 25be46ebe46..519e9157d6c 100644
--- a/third_party/xla/xla/tsl/profiler/backends/cpu/threadpool_listener.cc
+++ b/third_party/xla/xla/tsl/profiler/backends/cpu/threadpool_listener.cc
@@ -46,11 +46,6 @@ void RegisterThreadpoolEventCollector(
   tracing::SetEventCollector(tracing::EventCategory::kRunClosure, collector);
 }
 
-void UnregisterThreadpoolEventCollector() {
-  tracing::SetEventCollector(tracing::EventCategory::kScheduleClosure, nullptr);
-  tracing::SetEventCollector(tracing::EventCategory::kRunClosure, nullptr);
-}
-
 }  // namespace
 
 void ThreadpoolEventCollector::RecordEvent(uint64_t arg) const {
@@ -89,7 +84,6 @@ absl::Status ThreadpoolProfilerInterface::Start() {
 
 absl::Status ThreadpoolProfilerInterface::Stop() {
   threadpool_listener::Deactivate();
-  UnregisterThreadpoolEventCollector();
   return absl::OkStatus();
 }
 

From 466ab4c1b955acdef99ef7713b8dd1a107abc630 Mon Sep 17 00:00:00 2001
From: Quoc Truong <quoct@google.com>
Date: Fri, 9 Oct 2026 14:17:02 -0700
Subject: [PATCH 0391/1000] Use linux-x86-n4-16 for XLA BAP benchmark
 orchestration workflows.

PiperOrigin-RevId: 996752435
---
 third_party/xla/.github/workflows/nightly_benchmarks.yml   | 1 +
 third_party/xla/.github/workflows/postsubmit_benchmark.yml | 1 +
 third_party/xla/.github/workflows/presubmit_benchmark.yml  | 1 +
 third_party/xla/.github/workflows/run_benchmarks.yml       | 1 +
 4 files changed, 4 insertions(+)

diff --git a/third_party/xla/.github/workflows/nightly_benchmarks.yml b/third_party/xla/.github/workflows/nightly_benchmarks.yml
index 9cb42a2f529..32cfddea1c1 100644
--- a/third_party/xla/.github/workflows/nightly_benchmarks.yml
+++ b/third_party/xla/.github/workflows/nightly_benchmarks.yml
@@ -35,3 +35,4 @@ jobs:
       registry_file: "xla/tools/benchmarks/benchmark_registry.pbtxt"
       tag_filter: "scheduled"
       bap_ref: "aabe067ccf0612e9fe5dfb280176df43b12babf9"
+      runner: "linux-x86-n4-16"
diff --git a/third_party/xla/.github/workflows/postsubmit_benchmark.yml b/third_party/xla/.github/workflows/postsubmit_benchmark.yml
index b8956cb1704..aaffe7e599d 100644
--- a/third_party/xla/.github/workflows/postsubmit_benchmark.yml
+++ b/third_party/xla/.github/workflows/postsubmit_benchmark.yml
@@ -64,6 +64,7 @@ jobs:
       registry_file: "xla/tools/benchmarks/benchmark_registry.pbtxt"
       tag_filter: "postsubmit"
       bap_ref: "aabe067ccf0612e9fe5dfb280176df43b12babf9"
+      runner: "linux-x86-n4-16"
       publish_metrics: true
       pub_sub_gcp_project_id: "ml-oss-benchmarking-production"
       pub_sub_gcp_topic_id: "public-results-prod"
diff --git a/third_party/xla/.github/workflows/presubmit_benchmark.yml b/third_party/xla/.github/workflows/presubmit_benchmark.yml
index d8987ad574f..f836cb1b821 100644
--- a/third_party/xla/.github/workflows/presubmit_benchmark.yml
+++ b/third_party/xla/.github/workflows/presubmit_benchmark.yml
@@ -58,6 +58,7 @@ jobs:
       registry_file: "xla/tools/benchmarks/benchmark_registry.pbtxt"
       tag_filter: "presubmit"
       bap_ref: "aabe067ccf0612e9fe5dfb280176df43b12babf9"
+      runner: "linux-x86-n4-16"
       publish_metrics: false
       benchmark_filter: ${{ inputs.benchmark_filter || '' }}
       environment_filter: ${{ inputs.environment_filter || '' }}
diff --git a/third_party/xla/.github/workflows/run_benchmarks.yml b/third_party/xla/.github/workflows/run_benchmarks.yml
index e71c49ff874..e0538921e53 100644
--- a/third_party/xla/.github/workflows/run_benchmarks.yml
+++ b/third_party/xla/.github/workflows/run_benchmarks.yml
@@ -61,6 +61,7 @@ jobs:
     with:
       registry_file: "xla/tools/benchmarks/benchmark_registry.pbtxt"
       bap_ref: "aabe067ccf0612e9fe5dfb280176df43b12babf9"
+      runner: "linux-x86-n4-16"
       ab_mode: ${{ inputs.ab_mode }}
       baseline_ref: ${{ inputs.baseline_ref }}
       experiment_ref: ${{ inputs.experiment_ref }}

From 6905192301c41455c5b30dca7038a44f9500aa02 Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Fri, 9 Oct 2026 14:34:02 -0700
Subject: [PATCH 0392/1000] `AlgebraicSimplifierVisitor::HandleDynamicSlice`
 folds a nested `dynamic-slice` of a `dynamic-slice` into a single
 `dynamic-slice` on the base operand. The pass clamps the inner and outer
 start indices to their respective slice bounds and adds them in the outer
 index type. When either index uses a narrow integer type such as `s8` or
 `u8`, the clamp bound or the combined sum can overflow that type. An overflow
 wraps the start offset and causes the folded slice to read from the wrong
 coordinates.

Check across every slice dimension that the inner clamp bound fits in the inner index element type using `primitive_util::FitsInIntegralType`. Also check that the maximum combined offset, equal to the base dimension minus the outer slice size, fits in the outer index element type. Skip the folding rewrite whenever any dimension fails this check. Existing programs with wide index types are unaffected, while programs that previously suffered integer overflow now preserve the nested operations. No flags gate this bug fix.

Tests:
- `//third_party/tensorflow/compiler/xla/hlo/transforms/simplifiers:algebraic_simplifier_test` pins skipping the fold when narrow start indices would overflow.
PiperOrigin-RevId: 996762580
---
 .../simplifiers/algebraic_simplifier.cc       | 31 ++++++++++++----
 .../simplifiers/algebraic_simplifier_test.cc  | 37 +++++++++++++++++++
 2 files changed, 61 insertions(+), 7 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
index 8bf9b0f3581..6ab2eacd57c 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
@@ -8392,22 +8392,39 @@ absl::Status AlgebraicSimplifierVisitor::HandleDynamicSlice(
   //   ds(x, clamp(0, id, inner_size - outer_size) +
   //         clamp(0, inner_id, operand_size - inner_size))
   if (operand->opcode() == HloOpcode::kDynamicSlice) {
+    const int64_t first_index = Cast<HloDynamicSliceInstruction>(dynamic_slice)
+                                    ->first_index_operand_number();
+    for (int64_t i = first_index; i < dynamic_slice->operand_count(); ++i) {
+      int64_t operand_size =
+          operand->operand(0)->shape().dimensions(i - first_index);
+      if (!primitive_util::FitsInIntegralType(
+              operand_size - operand->dynamic_slice_sizes()[i - first_index],
+              operand->operand(i)->shape().element_type()) ||
+          !primitive_util::FitsInIntegralType(
+              operand_size -
+                  dynamic_slice->dynamic_slice_sizes()[i - first_index],
+              dynamic_slice->operand(i)->shape().element_type())) {
+        return absl::OkStatus();
+      }
+    }
     ABSL_RETURN_IF_ERROR(dynamic_slice->ReplaceOperandWithDifferentShape(
         0, operand->mutable_operand(0)));
-    for (int64_t i = 1; i < dynamic_slice->operand_count(); ++i) {
+    for (int64_t i = first_index; i < dynamic_slice->operand_count(); ++i) {
       HloInstruction* index = dynamic_slice->mutable_operand(i);
       index = index->AddInstruction(HloInstruction::CreateTernary(
           index->shape(), HloOpcode::kClamp, MakeScalarLike(index, 0), index,
-          MakeScalarLike(index,
-                         operand->dynamic_slice_sizes()[i - 1] -
-                             dynamic_slice->dynamic_slice_sizes()[i - 1])));
+          MakeScalarLike(
+              index,
+              operand->dynamic_slice_sizes()[i - first_index] -
+                  dynamic_slice->dynamic_slice_sizes()[i - first_index])));
       HloInstruction* inner_index = operand->mutable_operand(i);
       inner_index = inner_index->AddInstruction(HloInstruction::CreateTernary(
           inner_index->shape(), HloOpcode::kClamp,
           MakeScalarLike(inner_index, 0), inner_index,
-          MakeScalarLike(inner_index,
-                         operand->operand(0)->shape().dimensions(i - 1) -
-                             operand->dynamic_slice_sizes()[i - 1])));
+          MakeScalarLike(
+              inner_index,
+              operand->operand(0)->shape().dimensions(i - first_index) -
+                  operand->dynamic_slice_sizes()[i - first_index])));
       if (inner_index->shape().element_type() !=
           index->shape().element_type()) {
         inner_index = inner_index->AddInstruction(
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
index ec5e5673995..7c99815d8a8 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
@@ -15438,6 +15438,43 @@ TEST_F(AlgebraicSimplifierTest, DynamicSliceOfDynamicSlice) {
                                   m::ConstantScalar(6))))));
 }
 
+TEST_F(AlgebraicSimplifierTest, DynamicSliceOfDynamicSliceNarrowIndexOverflow) {
+  // Each individual clamp bound (70) fits in s8, but their sum (140) overflows
+  // s8.
+  constexpr absl::string_view hlo_outer_overflow = R"(
+    HloModule module
+
+    ENTRY test {
+      operand = f32[150] parameter(0)
+      i = s8[] parameter(1)
+      j = s8[] parameter(2)
+      inner_ds = f32[80] dynamic-slice(operand, i), dynamic_slice_sizes={80}
+      ROOT outer_ds = f32[10] dynamic-slice(inner_ds, j), dynamic_slice_sizes={10}
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto module1,
+                       ParseAndReturnVerifiedModule(hlo_outer_overflow));
+  AlgebraicSimplifier simplifier(default_options_);
+  EXPECT_THAT(simplifier.Run(module1.get()), absl_testing::IsOkAndHolds(false));
+
+  // The combined bound (190) fits in the outer index type s32, but the inner
+  // clamp bound (150) overflows the inner index type s8.
+  constexpr absl::string_view hlo_inner_overflow = R"(
+    HloModule module
+
+    ENTRY test {
+      operand = f32[200] parameter(0)
+      i = s8[] parameter(1)
+      j = s32[] parameter(2)
+      inner_ds = f32[50] dynamic-slice(operand, i), dynamic_slice_sizes={50}
+      ROOT outer_ds = f32[10] dynamic-slice(inner_ds, j), dynamic_slice_sizes={10}
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto module2,
+                       ParseAndReturnVerifiedModule(hlo_inner_overflow));
+  EXPECT_THAT(simplifier.Run(module2.get()), absl_testing::IsOkAndHolds(false));
+}
+
 TEST_F(AlgebraicSimplifierTest, DynamicUpdateSliceOfDynamicUpdateSlice) {
   constexpr absl::string_view hlo_string = R"(
     HloModule module

From d2c836215012a3b0fb015b98b6e594cc58996fd8 Mon Sep 17 00:00:00 2001
From: Ankit Vishwakarma <ankit99aug@gmail.com>
Date: Fri, 9 Oct 2026 14:35:28 -0700
Subject: [PATCH 0393/1000] PR #48156: Fix sm107 cuda ptx 47982
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/48156

Fixes #47982

📝 Summary of Changes
Fixes a crash when compiling for SM107 GPUs (Compute Capability 10.7) on CUDA 12.9.
- Updated `ResolveSupportedComputeCapability` and `GetSmName` in `xla/service/gpu/llvm_gpu_backend/nvptx_backend.h` and `nvptx_backend.cc` to accept an optional `ptx_version` parameter.
- Added `GetMinimumRequiredPtxVersion` to filter out target architectures whose required minimum PTX ISA version exceeds the PTX version supported by the active CUDA version.
- On SM107 GPUs with CUDA 12.9 (PTX 8.8), target resolution now safely falls back to `sm_103f` (the Blackwell family-compatible target). On CUDA 13.4+ (PTX 9.4+), it resolves directly to `sm_107a`.
- Updated `NVPTXCompiler::CompileTargetBinary` and `NVPTXCompiler::LinkModules` in `xla/service/gpu/nvptx_compiler.cc` to pass the active PTX version to target resolution.

🎯 Justification
On SM107 GPUs with CUDA 12.9, XLA previously resolved the device compute capability to `sm_107a`. However, LLVM's NVPTX target requires PTX 9.4+ for target `sm_107a`, whereas CUDA 12.9 only supports up to PTX 8.8. Passing `sm_107a` and `+ptx88` to LLVM caused a `report_fatal_error`, aborting the process. By incorporating PTX version constraints into target resolution, XLA safely falls back to the Blackwell family-compatible target (`sm_103f`) on CUDA 12.9, enabling JAX/XLA programs to compile and execute without process failure.

🚀 Kind of Contribution
🐛 Bug Fix

📊 Benchmark (for Performance Improvements)
N/A

🧪 Unit Tests:
Added `UtilsTest.SM107FallsBackWhenPtxVersionTooLow` in `xla/service/gpu/llvm_gpu_backend/nvptx_backend_test.cc` to verify:
- SM107 (`{10, 7, kAcceleratedFeatures}`) with PTX 8.8 resolves to `sm_103f` (`{10, 3, kFamilyCompatibleFeatures}`).
- SM107 (`{10, 7, kAcceleratedFeatures}`) with PTX 9.4 resolves to `sm_107a` (`{10, 7, kAcceleratedFeatures}`).

🧪 Execution Tests:
Target selection logic and PTX version mapping verified via unit tests in `nvptx_backend_test.cc`.

Copybara import of the project:

--
381e33f31960220eaa8328bf3ad194da2cc5c3d0 by Ankit Vishwakarma <ankit99aug@gmail.com>:

Fix SM107 compilation failure under CUDA 12.9 (Issue #47982)

--
f4e5f7a52c57b804494f1a476c9fba1f40b67450 by ankit vishwakarma <ankit99aug@gmail.com>:

Fix SM107 PTX change build dependencies and formatting

--
47e8b93e30512754bbda0e216d0032f39b823690 by ankit vishwakarma <ankit99aug@gmail.com>:

Address SM107 PTX review comments

--
eae21105b863136de9d921cd75bbf647407c14b9 by ankit vishwakarma <ankit99aug@gmail.com>:

Fix PTX target version selection after rebase

Merging this change closes #48156

PiperOrigin-RevId: 996763427
---
 third_party/xla/xla/service/gpu/BUILD         |  3 +
 .../xla/service/gpu/llvm_gpu_backend/BUILD    |  1 +
 .../gpu/llvm_gpu_backend/nvptx_backend.cc     | 91 ++++++++++++++++++-
 .../gpu/llvm_gpu_backend/nvptx_backend.h      |  7 +-
 .../llvm_gpu_backend/nvptx_backend_test.cc    | 30 ++++++
 .../xla/xla/service/gpu/nvptx_compiler.cc     | 39 +++++++-
 6 files changed, 163 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index b7ed894874e..8fcf2718518 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -2520,6 +2520,7 @@ cc_library(
         "//xla/service:hlo_verifier",
         "//xla/service/gpu/llvm_gpu_backend:nvptx_backend",
         "//xla/service/gpu/llvm_gpu_backend:nvptx_utils",
+        "//xla/service/gpu/llvm_gpu_backend:ptx_version_util",
         "//xla/service/llvm_ir:llvm_util",
         "//xla/stream_executor:device_description",
         "//xla/stream_executor:dnn",
@@ -2532,6 +2533,7 @@ cc_library(
         "//xla/stream_executor/cuda:cuda_compute_capability",
         "//xla/stream_executor/cuda:cuda_diagnostics",
         "//xla/stream_executor/cuda:cuda_platform_id",
+        "//xla/stream_executor/cuda:subprocess_compilation",
         "//xla/tsl/platform:env",
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:statusor",
@@ -2554,6 +2556,7 @@ cc_library(
         "@llvm-project//llvm:Support",
         "@llvm-project//llvm:ir_headers",
         "@llvm-project//mlir:IR",
+        "@local_config_cuda//cuda:cuda_headers",
         "@tsl//tsl/platform:path",
         "@tsl//tsl/profiler/lib:scoped_annotation",
         "@tsl//tsl/profiler/lib:traceme",
diff --git a/third_party/xla/xla/service/gpu/llvm_gpu_backend/BUILD b/third_party/xla/xla/service/gpu/llvm_gpu_backend/BUILD
index cf6534c147e..3d4d9b2fdd6 100644
--- a/third_party/xla/xla/service/gpu/llvm_gpu_backend/BUILD
+++ b/third_party/xla/xla/service/gpu/llvm_gpu_backend/BUILD
@@ -86,6 +86,7 @@ cc_library(
         "//xla/service/gpu:metrics",
         "//xla/service/llvm_ir:llvm_command_line_options",
         "//xla/stream_executor:device_description",
+        "//xla/stream_executor:semantic_version",
         "//xla/stream_executor/cuda:cuda_compute_capability",
         "//xla/tsl/platform:env",
         "@com_google_absl//absl/base",
diff --git a/third_party/xla/xla/service/gpu/llvm_gpu_backend/nvptx_backend.cc b/third_party/xla/xla/service/gpu/llvm_gpu_backend/nvptx_backend.cc
index a6d069eb259..5deb684a808 100644
--- a/third_party/xla/xla/service/gpu/llvm_gpu_backend/nvptx_backend.cc
+++ b/third_party/xla/xla/service/gpu/llvm_gpu_backend/nvptx_backend.cc
@@ -162,7 +162,12 @@ absl::StatusOr<std::unique_ptr<llvm::TargetMachine>> NVPTXGetTargetMachine(
   std::string feature_str =
       absl::StrFormat("+ptx%d", highest_supported_ptx_version);
 
-  return GetTargetMachine(target_triple, nvptx::GetSmName(compute_capability),
+  const se::SemanticVersion ptx_version{
+      static_cast<unsigned>(highest_supported_ptx_version / 10),
+      static_cast<unsigned>(highest_supported_ptx_version % 10), 0};
+
+  return GetTargetMachine(target_triple,
+                          nvptx::GetSmName(compute_capability, ptx_version),
                           debug_options, feature_str);
 }
 
@@ -271,8 +276,80 @@ constexpr se::CudaComputeCapability kSupportedVersions[] = {
     {8, 7},  {8, 6},  {8, 0},  {7, 5},  {7, 2},  {7, 0},  {6, 2}, {6, 1},
     {6, 0},  {5, 3},  {5, 2},  {5, 0},  {3, 7},  {3, 5},  {3, 2}, {3, 0}};
 
+stream_executor::SemanticVersion GetMinimumRequiredPtxVersion(
+    se::CudaComputeCapability cc) {
+  if (cc.major == 12) {
+    if (cc.minor >= 1) {
+      return {8, 8, 0};
+    }
+    return {8, 7, 0};
+  }
+  if (cc.major == 11) {
+    return {9, 0, 0};
+  }
+  if (cc.major == 10) {
+    if (cc.minor >= 7) {
+      return {9, 4, 0};
+    }
+    if (cc.minor >= 3) {
+      return {8, 8, 0};
+    }
+    return {8, 7, 0};
+  }
+  if (cc.major == 9) {
+    return {7, 8, 0};
+  }
+  if (cc.major == 8) {
+    if (cc.minor >= 9) {
+      return {7, 8, 0};
+    }
+    if (cc.minor >= 7) {
+      return {7, 4, 0};
+    }
+    if (cc.minor >= 6) {
+      return {7, 1, 0};
+    }
+    return {7, 0, 0};
+  }
+  if (cc.major == 7) {
+    if (cc.minor >= 5) {
+      return {6, 3, 0};
+    }
+    if (cc.minor >= 2) {
+      return {6, 1, 0};
+    }
+    return {6, 0, 0};
+  }
+  if (cc.major == 6) {
+    return {5, 0, 0};
+  }
+  if (cc.major == 5) {
+    if (cc.minor >= 3) {
+      return {4, 2, 0};
+    }
+    if (cc.minor >= 2) {
+      return {4, 1, 0};
+    }
+    return {4, 0, 0};
+  }
+  if (cc.major == 3) {
+    if (cc.minor >= 7) {
+      return {4, 1, 0};
+    }
+    if (cc.minor >= 5) {
+      return {3, 5, 0};
+    }
+    if (cc.minor >= 2) {
+      return {3, 2, 0};
+    }
+    return {3, 0, 0};
+  }
+  return {3, 0, 0};
+}
+
 se::CudaComputeCapability ResolveSupportedComputeCapability(
-    se::CudaComputeCapability compute_capability) {
+    se::CudaComputeCapability compute_capability,
+    std::optional<stream_executor::SemanticVersion> ptx_version) {
   using CudaComputeCapabilities =
       se::CudaComputeCapability::CudaComputeCapabilities;
 
@@ -287,6 +364,10 @@ se::CudaComputeCapability ResolveSupportedComputeCapability(
 
   for (const auto& v : kSupportedVersions) {
     if (gpu_compute_capability.SupportsAllFeaturesOf(v)) {
+      if (ptx_version.has_value() &&
+          *ptx_version < GetMinimumRequiredPtxVersion(v)) {
+        continue;
+      }
       // Found the most advanced supported capability
       target_compute_capability = v;
       break;
@@ -316,9 +397,11 @@ se::CudaComputeCapability ResolveSupportedComputeCapability(
   return target_compute_capability;
 }
 
-std::string GetSmName(se::CudaComputeCapability compute_capability) {
+std::string GetSmName(
+    se::CudaComputeCapability compute_capability,
+    std::optional<stream_executor::SemanticVersion> ptx_version) {
   se::CudaComputeCapability target_compute_capability =
-      ResolveSupportedComputeCapability(compute_capability);
+      ResolveSupportedComputeCapability(compute_capability, ptx_version);
 
   // If the current CC isn't supported by LLVM and it is newer then
   // the max supported LLVM version, do not warn about it. The end
diff --git a/third_party/xla/xla/service/gpu/llvm_gpu_backend/nvptx_backend.h b/third_party/xla/xla/service/gpu/llvm_gpu_backend/nvptx_backend.h
index d4737daa797..41ac1e5284a 100644
--- a/third_party/xla/xla/service/gpu/llvm_gpu_backend/nvptx_backend.h
+++ b/third_party/xla/xla/service/gpu/llvm_gpu_backend/nvptx_backend.h
@@ -28,6 +28,7 @@ limitations under the License.
 #include "llvm/TargetParser/Triple.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/device_description.h"
+#include "xla/stream_executor/semantic_version.h"
 #include "xla/xla.pb.h"
 
 namespace xla::gpu::nvptx {
@@ -44,13 +45,15 @@ absl::StatusOr<int> GetMaxPtxVersionSupportedByLlvm(
 // advanced supported compute capability that the device can run, potentially
 // with the family ("f") feature extension enabled.
 stream_executor::CudaComputeCapability ResolveSupportedComputeCapability(
-    stream_executor::CudaComputeCapability compute_capability);
+    stream_executor::CudaComputeCapability compute_capability,
+    std::optional<stream_executor::SemanticVersion> ptx_version = std::nullopt);
 
 // Gets the GPU name as it's known to LLVM for a given compute
 // capability.  If we see an unrecognized compute capability, we
 // return the highest one that is known and below the selected device.
 std::string GetSmName(
-    stream_executor::CudaComputeCapability compute_capability);
+    stream_executor::CudaComputeCapability compute_capability,
+    std::optional<stream_executor::SemanticVersion> ptx_version = std::nullopt);
 
 // Compiles the argument module and returns it. libdevice_dir_path is the
 // parent directory of the libdevice bitcode libraries. The contents of the
diff --git a/third_party/xla/xla/service/gpu/llvm_gpu_backend/nvptx_backend_test.cc b/third_party/xla/xla/service/gpu/llvm_gpu_backend/nvptx_backend_test.cc
index d06de9254e8..50754b7af00 100644
--- a/third_party/xla/xla/service/gpu/llvm_gpu_backend/nvptx_backend_test.cc
+++ b/third_party/xla/xla/service/gpu/llvm_gpu_backend/nvptx_backend_test.cc
@@ -116,6 +116,36 @@ TEST(UtilsTest, UnknownCapabilityFallsBackToFamilyCompatible) {
             (se::CudaComputeCapability{9, 0, FeatureExtension::kNone}));
 }
 
+TEST(UtilsTest, SM107FallsBackWhenPtxVersionTooLow) {
+  using FeatureExtension = se::CudaComputeCapability::FeatureExtension;
+  // SM107 requires PTX 9.4. With PTX 8.8 (e.g. CUDA 12.9), it falls back to
+  // sm_103f.
+  EXPECT_EQ(nvptx::ResolveSupportedComputeCapability(
+                se::CudaComputeCapability{
+                    10, 7, FeatureExtension::kAcceleratedFeatures},
+                se::SemanticVersion{8, 8, 0}),
+            (se::CudaComputeCapability{
+                10, 3, FeatureExtension::kFamilyCompatibleFeatures}));
+  // With PTX 9.4, SM107 remains sm_107a.
+  EXPECT_EQ(nvptx::ResolveSupportedComputeCapability(
+                se::CudaComputeCapability{
+                    10, 7, FeatureExtension::kAcceleratedFeatures},
+                se::SemanticVersion{9, 4, 0}),
+            (se::CudaComputeCapability{
+                10, 7, FeatureExtension::kAcceleratedFeatures}));
+
+  EXPECT_EQ(nvptx::GetSmName(
+                se::CudaComputeCapability{
+                    10, 7, FeatureExtension::kAcceleratedFeatures},
+                se::SemanticVersion{8, 8, 0}),
+            "sm_103f");
+  EXPECT_EQ(nvptx::GetSmName(
+                se::CudaComputeCapability{
+                    10, 7, FeatureExtension::kAcceleratedFeatures},
+                se::SemanticVersion{9, 4, 0}),
+            "sm_107a");
+}
+
 using VersionPair = std::pair<se::SemanticVersion, se::SemanticVersion>;
 using PtxVersionFromCudaVersionTest = ::testing::TestWithParam<VersionPair>;
 
diff --git a/third_party/xla/xla/service/gpu/nvptx_compiler.cc b/third_party/xla/xla/service/gpu/nvptx_compiler.cc
index d1263e2cebf..516681fedb4 100644
--- a/third_party/xla/xla/service/gpu/nvptx_compiler.cc
+++ b/third_party/xla/xla/service/gpu/nvptx_compiler.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/service/gpu/nvptx_compiler.h"
 
+#include <algorithm>
 #include <array>
 #include <cstdint>
 #include <fstream>
@@ -32,12 +33,14 @@ limitations under the License.
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
+#include "third_party/gpus/cuda/include/cuda.h"
 #include "llvm/IRReader/IRReader.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_ostream.h"
@@ -83,6 +86,7 @@ limitations under the License.
 #include "xla/service/gpu/ir_emission_utils.h"
 #include "xla/service/gpu/llvm_gpu_backend/nvptx_backend.h"
 #include "xla/service/gpu/llvm_gpu_backend/nvptx_utils.h"
+#include "xla/service/gpu/llvm_gpu_backend/ptx_version_util.h"
 #include "xla/service/gpu/metrics.h"
 #include "xla/service/gpu/nvptx_alias_info.h"
 #include "xla/service/gpu/ptx_compile_options_from_debug_options.h"
@@ -99,6 +103,7 @@ limitations under the License.
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/cuda/cuda_diagnostics.h"
 #include "xla/stream_executor/cuda/cuda_platform_id.h"
+#include "xla/stream_executor/cuda/subprocess_compilation.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/dnn.h"
 #include "xla/stream_executor/semantic_version.h"
@@ -657,8 +662,23 @@ NVPTXCompiler::CompileTargetBinary(
   se::cuda::CompilationOptions compilation_options =
       PtxCompileOptionsFromDebugOptions(module_config.debug_options());
 
+  absl::StatusOr<stream_executor::SemanticVersion> runtime_cuda_version =
+      stream_executor::GetAsmCompilerVersion(
+          module_config.debug_options().xla_gpu_cuda_data_dir());
+  constexpr stream_executor::SemanticVersion kCompileTimeCudaVersion{
+      CUDA_VERSION / 1000, (CUDA_VERSION / 10) % 100, CUDA_VERSION % 10};
+  auto highest_supported_cuda_version = [&] {
+    if (runtime_cuda_version.ok()) {
+      return std::min(runtime_cuda_version.value(), kCompileTimeCudaVersion);
+    }
+    return kCompileTimeCudaVersion;
+  }();
+  auto ptx_version = nvptx::DetermineHighestSupportedPtxVersionFromCudaVersion(
+      highest_supported_cuda_version);
+
   se::CudaComputeCapability cc = nvptx::ResolveSupportedComputeCapability(
-      *device_description.gpu_compute_capability().cuda_compute_capability());
+      *device_description.gpu_compute_capability().cuda_compute_capability(),
+      ptx_version);
 
   // This may print multiple lines per HLO compilation because of the
   // parallelized compilation of LLVM modules.
@@ -719,8 +739,23 @@ absl::StatusOr<std::vector<uint8_t>> NVPTXCompiler::LinkModules(
     return std::vector<uint8_t>{};
   }
 
+  absl::StatusOr<stream_executor::SemanticVersion> runtime_cuda_version =
+      stream_executor::GetAsmCompilerVersion(
+          debug_options.xla_gpu_cuda_data_dir());
+  constexpr stream_executor::SemanticVersion kCompileTimeCudaVersion{
+      CUDA_VERSION / 1000, (CUDA_VERSION / 10) % 100, CUDA_VERSION % 10};
+  auto highest_supported_cuda_version = [&] {
+    if (runtime_cuda_version.ok()) {
+      return std::min(runtime_cuda_version.value(), kCompileTimeCudaVersion);
+    }
+    return kCompileTimeCudaVersion;
+  }();
+  auto ptx_version = nvptx::DetermineHighestSupportedPtxVersionFromCudaVersion(
+      highest_supported_cuda_version);
+
   se::CudaComputeCapability cc = nvptx::ResolveSupportedComputeCapability(
-      *device_description.gpu_compute_capability().cuda_compute_capability());
+      *device_description.gpu_compute_capability().cuda_compute_capability(),
+      ptx_version);
 
   ABSL_ASSIGN_OR_RETURN(const se::cuda::CompilationProvider* compilation_provider,
                    GetCompilationProvider(debug_options, stream_exec));

From 986c4ec584798a6605c8e290fcd97a869208ecc9 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 9 Oct 2026 15:23:22 -0700
Subject: [PATCH 0394/1000] [XLA][BufferAssignment] Split global and
 thread-local buffer assignment.

Sharing `AssignBuffersForComputations` between global and thread-local
computations via an `is_thread_local` flag overly complicated thread-local
buffer assignment in favor of code reuse, even though very little code was
actually shared. Thread-local computations simply assign a dedicated
allocation to each buffer and do not use call-graph post-order positions,
buffer sorting, allocation managers, operand buffer reuse, or delayed
sequential heap simulation, yet ran through that setup and branching anyway.

Split `AssignBuffersForComputations` into `AssignBuffersForGlobalComputations`
and `AssignBuffersForThreadLocalComputations`. Extract
`AssignViewOrConstantBuffer` for the only logic shared by both paths (DUS view
buffer skipping and `allocate_buffers_for_constants` handling), remove the
`is_thread_local` parameter from `AssignSpecialHloBuffer` and
`AssignSingleHloBuffer`, and return early when there are no thread-local
computations.

PiperOrigin-RevId: 996791897
---
 .../xla/xla/service/buffer_assignment.cc      | 114 ++++++++++++------
 .../xla/xla/service/buffer_assignment.h       |  30 +++--
 .../xla/xla/service/buffer_assignment_test.cc |  59 +++++++++
 3 files changed, 156 insertions(+), 47 deletions(-)

diff --git a/third_party/xla/xla/service/buffer_assignment.cc b/third_party/xla/xla/service/buffer_assignment.cc
index 654573c227f..163071d7d73 100644
--- a/third_party/xla/xla/service/buffer_assignment.cc
+++ b/third_party/xla/xla/service/buffer_assignment.cc
@@ -2087,10 +2087,8 @@ absl::StatusOr<bool> BufferAssigner::MaybeAssignBuffer(
   return true;
 }
 
-absl::StatusOr<bool> BufferAssigner::AssignSpecialHloBuffer(
-    const HloBuffer* hlo_buffer, bool is_thread_local,
-    BufferAllocationsManagerForComputationsWithoutOrdering* allocation_manager,
-    BufferAssignment* assignment) {
+absl::StatusOr<bool> BufferAssigner::AssignViewOrConstantBuffer(
+    const HloBuffer* hlo_buffer, BufferAssignment* assignment) {
   // "View" buffers are pointer stand-ins that alias into another allocation, so
   // they get no allocation of their own.
   if (opts_.dus_view_color.has_value()) {
@@ -2101,12 +2099,13 @@ absl::StatusOr<bool> BufferAssigner::AssignSpecialHloBuffer(
     }
   }
 
-  const int64_t buffer_size = assignment->HloBufferSize(*hlo_buffer);
   for (const HloValue* value : hlo_buffer->values()) {
     if (value->instruction()->opcode() == HloOpcode::kConstant) {
       if (opts_.allocate_buffers_for_constants) {
-        ABSL_ASSIGN_OR_RETURN(BufferAllocation * allocation,
-                         assignment->NewAllocation(*hlo_buffer, buffer_size));
+        ABSL_ASSIGN_OR_RETURN(
+            BufferAllocation * allocation,
+            assignment->NewAllocation(*hlo_buffer,
+                                      assignment->HloBufferSize(*hlo_buffer)));
         allocation->set_constant(true);
         VLOG(3) << "New allocation #" << allocation->index() << " for constant "
                 << *hlo_buffer << " value ptr: " << value;
@@ -2114,7 +2113,21 @@ absl::StatusOr<bool> BufferAssigner::AssignSpecialHloBuffer(
       VLOG(3) << "Not allocating buffer for constant";
       return true;
     }
+  }
+  return false;
+}
 
+absl::StatusOr<bool> BufferAssigner::AssignSpecialHloBuffer(
+    const HloBuffer* hlo_buffer,
+    BufferAllocationsManagerForComputationsWithoutOrdering* allocation_manager,
+    BufferAssignment* assignment) {
+  ABSL_ASSIGN_OR_RETURN(bool handled,
+                   AssignViewOrConstantBuffer(hlo_buffer, assignment));
+  if (handled) {
+    return true;
+  }
+
+  for (const HloValue* value : hlo_buffer->values()) {
     const HloInstruction* instruction = value->instruction();
     const bool is_entry_parameter =
         instruction->opcode() == HloOpcode::kParameter &&
@@ -2128,8 +2141,10 @@ absl::StatusOr<bool> BufferAssigner::AssignSpecialHloBuffer(
       // allocation and sets its parameter number. Parameters of non-entry
       // computations do not need special allocations because they live inside
       // callers.
-      ABSL_ASSIGN_OR_RETURN(BufferAllocation * allocation,
-                       assignment->NewAllocation(*hlo_buffer, buffer_size));
+      ABSL_ASSIGN_OR_RETURN(
+          BufferAllocation * allocation,
+          assignment->NewAllocation(*hlo_buffer,
+                                    assignment->HloBufferSize(*hlo_buffer)));
 
       allocation->set_entry_computation_parameter(
           instruction->parameter_number(), value->index(), parameter_has_alias);
@@ -2143,19 +2158,12 @@ absl::StatusOr<bool> BufferAssigner::AssignSpecialHloBuffer(
     }
   }
 
-  if (is_thread_local) {
-    ABSL_ASSIGN_OR_RETURN(BufferAllocation * allocation,
-                     assignment->NewAllocation(*hlo_buffer, buffer_size));
-    allocation->set_is_thread_local(true);
-    VLOG(3) << "New allocation #" << allocation->index()
-            << " for thread-local: " << *hlo_buffer;
-    return true;
-  }
-
   for (const HloValue* value : hlo_buffer->values()) {
     if (value->shape().IsTuple()) {
-      ABSL_ASSIGN_OR_RETURN(BufferAllocation * allocation,
-                       assignment->NewAllocation(*hlo_buffer, buffer_size));
+      ABSL_ASSIGN_OR_RETURN(
+          BufferAllocation * allocation,
+          assignment->NewAllocation(*hlo_buffer,
+                                    assignment->HloBufferSize(*hlo_buffer)));
       allocation->set_is_tuple(true);
       VLOG(3) << "New allocation #" << allocation->index()
               << " for tuple-shaped buffer: " << *hlo_buffer;
@@ -2193,8 +2201,8 @@ bool BufferAssigner::DelayTemporaryBufferAssignment(
       // decide to create a new allocation, to ensure we've exhausted all
       // the buffer re-use cases above.
       //
-      // Entry parameters and thread local buffers were already handled
-      // earlier in this loop iteration.  See
+      // Special buffers (entry parameters, constants, tuples) were already
+      // handled earlier in this loop iteration.  See
       // BufferAllocation::IsPreallocatedTempBuffer for the definition of
       // temp buffers.
       (*buffers_to_assign_sequentially)[computation].insert(hlo_value);
@@ -2206,7 +2214,7 @@ bool BufferAssigner::DelayTemporaryBufferAssignment(
 }
 
 absl::Status BufferAssigner::AssignSingleHloBuffer(
-    const HloBuffer* hlo_buffer, bool is_thread_local,
+    const HloBuffer* hlo_buffer,
     absl::flat_hash_map<const HloComputation*,
                         absl::flat_hash_set<const HloValue*>>*
         buffers_to_assign_sequentially,
@@ -2261,9 +2269,42 @@ absl::Status BufferAssigner::AssignSingleHloBuffer(
   return absl::OkStatus();
 }
 
-absl::Status BufferAssigner::AssignBuffersForComputations(
+absl::Status BufferAssigner::AssignBuffersForThreadLocalComputations(
+    const std::vector<const HloComputation*>& computations,
+    BufferAssignment* assignment) {
+  if (computations.empty()) {
+    return absl::OkStatus();
+  }
+
+  const HloAliasAnalysis& alias_analysis = assignment->alias_analysis();
+  absl::flat_hash_set<const HloComputation*> computations_set(
+      computations.begin(), computations.end());
+
+  for (const HloBuffer& buffer : alias_analysis.buffers()) {
+    TF_RET_CHECK(!buffer.values().empty());
+    const HloComputation* comp = buffer.values()[0]->instruction()->parent();
+    if (!computations_set.contains(comp) || assignment->HasAllocation(buffer)) {
+      continue;
+    }
+
+    ABSL_ASSIGN_OR_RETURN(bool handled,
+                     AssignViewOrConstantBuffer(&buffer, assignment));
+    if (handled) {
+      continue;
+    }
+
+    ABSL_ASSIGN_OR_RETURN(
+        BufferAllocation * allocation,
+        assignment->NewAllocation(buffer, assignment->HloBufferSize(buffer)));
+    allocation->set_is_thread_local(true);
+    VLOG(3) << "New allocation #" << allocation->index()
+            << " for thread-local: " << buffer;
+  }
+  return absl::OkStatus();
+}
+
+absl::Status BufferAssigner::AssignBuffersForGlobalComputations(
     const std::vector<const HloComputation*>& computations,
-    bool is_thread_local,
     absl::flat_hash_map<const HloComputation*,
                         absl::flat_hash_set<const HloValue*>>*
         buffers_to_assign_sequentially,
@@ -2273,6 +2314,7 @@ absl::Status BufferAssigner::AssignBuffersForComputations(
   if (computations.empty()) {
     return absl::OkStatus();
   }
+  TF_RET_CHECK(buffers_to_assign_sequentially != nullptr);
   std::vector<const HloBuffer*> sorted_buffers;
 
   // First assign the preset allocations.
@@ -2323,7 +2365,7 @@ absl::Status BufferAssigner::AssignBuffersForComputations(
     const HloInstructionSequence* instruction_sequence =
         assignment->hlo_ordering().SequentialOrder(*computation);
     const bool has_sequential_order = instruction_sequence != nullptr;
-    if (has_sequential_order && buffers_to_assign_sequentially != nullptr) {
+    if (has_sequential_order) {
       // Every sequential computation must get an entry in the
       // buffers_to_assign_sequentially map, even if we end up with an empty
       // set of buffers. This ensures we can correctly determine whether to
@@ -2406,11 +2448,10 @@ absl::Status BufferAssigner::AssignBuffersForComputations(
     BufferAllocationsManagerForComputationsWithoutOrdering* allocation_manager =
         get_manager(color);
 
-    ABSL_ASSIGN_OR_RETURN(bool special,
-                     AssignSpecialHloBuffer(buffer, is_thread_local,
-                                            allocation_manager, assignment));
+    ABSL_ASSIGN_OR_RETURN(bool special, AssignSpecialHloBuffer(
+                                       buffer, allocation_manager, assignment));
     if (!special) {
-      ABSL_RETURN_IF_ERROR(AssignSingleHloBuffer(buffer, is_thread_local,
+      ABSL_RETURN_IF_ERROR(AssignSingleHloBuffer(buffer,
                                             buffers_to_assign_sequentially,
                                             allocation_manager, assignment));
     }
@@ -3310,9 +3351,9 @@ absl::Status BufferAssigner::RunAssignBuffers(
   // 'buffers_to_assign_sequentially'.
   flat_hash_map<const HloComputation*, flat_hash_set<const HloValue*>>
       buffers_to_assign_sequentially;
-  ABSL_RETURN_IF_ERROR(AssignBuffersForComputations(
-      global_computations, /*is_thread_local=*/false,
-      &buffers_to_assign_sequentially, assignment, non_sequential_algorithm));
+  ABSL_RETURN_IF_ERROR(AssignBuffersForGlobalComputations(
+      global_computations, &buffers_to_assign_sequentially, assignment,
+      non_sequential_algorithm));
   // Assign buffers with sequential ordering, if any. If all global
   // computations are sequential, we can run heap simulation on the whole
   // module, which reduces memory usage.
@@ -3341,11 +3382,8 @@ absl::Status BufferAssigner::RunAssignBuffers(
     thread_local_computations_no_fusion.push_back(computation);
   }
 
-  ABSL_RETURN_IF_ERROR(AssignBuffersForComputations(
-      thread_local_computations_no_fusion, /*is_thread_local=*/true,
-      /*buffers_to_assign_sequentially=*/nullptr, assignment,
-      non_sequential_algorithm));
-
+  ABSL_RETURN_IF_ERROR(AssignBuffersForThreadLocalComputations(
+      thread_local_computations_no_fusion, assignment));
   // Mark all buffers which may be live out of the entry computation as
   // "liveout".
   for (const HloBuffer* buffer :
diff --git a/third_party/xla/xla/service/buffer_assignment.h b/third_party/xla/xla/service/buffer_assignment.h
index 7c161803fa8..31e1cdec889 100644
--- a/third_party/xla/xla/service/buffer_assignment.h
+++ b/third_party/xla/xla/service/buffer_assignment.h
@@ -1034,12 +1034,12 @@ class BufferAssigner {
       const std::vector<const HloComputation*>& thread_local_computations,
       BufferAssignment* assignment);
 
-  // is modified to reflect the new buffer assignments. If is_thread_local is
-  // true, then all assigned buffers have the is_thread_local flag set to
-  // true.
-  absl::Status AssignBuffersForComputations(
+  // Assigns buffers for the given global computations. `assignment` is modified
+  // to reflect the new buffer assignments. Temporary buffers for computations
+  // with sequential ordering are collected in `buffers_to_assign_sequentially`
+  // for later heap simulation.
+  absl::Status AssignBuffersForGlobalComputations(
       const std::vector<const HloComputation*>& computations,
-      bool is_thread_local,
       absl::flat_hash_map<const HloComputation*,
                           absl::flat_hash_set<const HloValue*>>*
           buffers_to_assign_sequentially,
@@ -1048,6 +1048,13 @@ class BufferAssigner {
           AssignmentAlgorithmForComputationsWithoutOrderingProto::Value
               algorithm);
 
+  // Assigns dedicated thread-local buffers for the given thread-local
+  // computations. `assignment` is modified to reflect the new buffer
+  // assignments.
+  absl::Status AssignBuffersForThreadLocalComputations(
+      const std::vector<const HloComputation*>& computations,
+      BufferAssignment* assignment);
+
   // Returns true if buffer's live range interferences with buffer2's.
   bool LiveRangeInterferes(const HloValue* buffer1,
                            const HloLiveRange::LiveRangeBounds& live_range1,
@@ -1061,17 +1068,22 @@ class BufferAssigner {
       absl::flat_hash_set<const HloBuffer*>* assigned_buffers,
       BufferAssignment* assignment);
 
-  // Assigns HloBuffers that require dedicated allocations upfront (constants,
-  // entry parameters, thread-local, tuples).
+  // Handles HloBuffers that represent DUS views or constants. Returns true if
+  // the buffer was handled (either allocated as a constant or skipped).
+  absl::StatusOr<bool> AssignViewOrConstantBuffer(const HloBuffer* hlo_buffer,
+                                                  BufferAssignment* assignment);
+
+  // Assigns HloBuffers that require dedicated allocations upfront (views,
+  // constants, entry parameters, tuples).
   absl::StatusOr<bool> AssignSpecialHloBuffer(
-      const HloBuffer* hlo_buffer, bool is_thread_local,
+      const HloBuffer* hlo_buffer,
       BufferAllocationsManagerForComputationsWithoutOrdering*
           allocation_manager,
       BufferAssignment* assignment);
 
   // Assigns a single hlo buffer to an HLO allocation.
   absl::Status AssignSingleHloBuffer(
-      const HloBuffer* hlo_buffer, bool is_thread_local,
+      const HloBuffer* hlo_buffer,
       absl::flat_hash_map<const HloComputation*,
                           absl::flat_hash_set<const HloValue*>>*
           buffers_to_assign_sequentially,
diff --git a/third_party/xla/xla/service/buffer_assignment_test.cc b/third_party/xla/xla/service/buffer_assignment_test.cc
index 8d45806b474..b1cacb823b5 100644
--- a/third_party/xla/xla/service/buffer_assignment_test.cc
+++ b/third_party/xla/xla/service/buffer_assignment_test.cc
@@ -2257,6 +2257,65 @@ TEST_F(BufferAssignmentTest, EmbeddedComputationBuffers) {
   EXPECT_FALSE(map_alloc.is_thread_local());
 }
 
+TEST_F(BufferAssignmentTest, EmbeddedComputationConstantAndTupleBuffers) {
+  // Verify that constants in thread-local computations respect
+  // allocate_buffers_for_constants (marked as constant rather than
+  // thread-local, or omitted when disabled), and tuple buffers in thread-local
+  // computations are marked as thread-local.
+  auto module = CreateNewVerifiedModule();
+  Shape scalar_shape = ShapeUtil::MakeShape(F32, {});
+
+  auto map_builder = HloComputation::Builder(TestName() + "_map");
+  auto map_param = map_builder.AddInstruction(
+      HloInstruction::CreateParameter(0, scalar_shape, "map_param"));
+  auto map_const = map_builder.AddInstruction(
+      HloInstruction::CreateConstant(LiteralUtil::CreateR0<float>(2.0f)));
+  auto map_tuple = map_builder.AddInstruction(
+      HloInstruction::CreateTuple({map_param, map_const}));
+  auto map_gte0 = map_builder.AddInstruction(
+      HloInstruction::CreateGetTupleElement(scalar_shape, map_tuple, 0));
+  auto map_gte1 = map_builder.AddInstruction(
+      HloInstruction::CreateGetTupleElement(scalar_shape, map_tuple, 1));
+  auto map_root = map_builder.AddInstruction(HloInstruction::CreateBinary(
+      scalar_shape, HloOpcode::kAdd, map_gte0, map_gte1));
+  HloComputation* map_computation =
+      module->AddEmbeddedComputation(map_builder.Build());
+
+  auto builder = HloComputation::Builder(TestName());
+  auto param = builder.AddInstruction(
+      HloInstruction::CreateParameter(0, scalar_shape, "param"));
+  builder.AddInstruction(
+      HloInstruction::CreateMap(scalar_shape, {param}, map_computation));
+  module->AddEntryComputation(builder.Build());
+
+  // 1. With allocate_buffers_for_constants = true:
+  std::unique_ptr<BufferAssignment> assignment_with_consts =
+      RunBufferAssignment(module.get());
+  EXPECT_TRUE(assignment_with_consts->HasTopLevelAllocation(map_const));
+  const BufferAllocation& const_alloc =
+      GetTopLevelAllocation(*assignment_with_consts, map_const);
+  EXPECT_TRUE(const_alloc.is_constant());
+  EXPECT_FALSE(const_alloc.is_thread_local());
+
+  const BufferAllocation& tuple_alloc =
+      GetTopLevelAllocation(*assignment_with_consts, map_tuple);
+  EXPECT_TRUE(tuple_alloc.is_thread_local());
+  EXPECT_FALSE(tuple_alloc.is_tuple());
+
+  const BufferAllocation& root_alloc =
+      GetTopLevelAllocation(*assignment_with_consts, map_root);
+  EXPECT_TRUE(root_alloc.is_thread_local());
+
+  // 2. With allocate_buffers_for_constants = false:
+  std::unique_ptr<BufferAssignment> assignment_no_consts =
+      RunBufferAssignmentNoBuffersForConstants(module.get());
+  EXPECT_FALSE(assignment_no_consts->HasTopLevelAllocation(map_const));
+  EXPECT_TRUE(GetTopLevelAllocation(*assignment_no_consts, map_tuple)
+                  .is_thread_local());
+  EXPECT_TRUE(
+      GetTopLevelAllocation(*assignment_no_consts, map_root).is_thread_local());
+}
+
 TEST_F(BufferAssignmentTest, CustomCallEmbeddedComputationBuffers) {
   // Verify that buffers for embedded computations in a custom call are properly
   // marked as thread-local.

From 6ed5641472576871388e78fbaa3561ea88c919b2 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 9 Oct 2026 16:00:35 -0700
Subject: [PATCH 0395/1000] [PjRt-IFRT] Fix `OutputSpec::shard_shape` in
 `PjRtExecutable::CommonMetadata::Serialize`

`PjRtExecutable::CommonMetadata::Serialize` computed the tiled per-shard shape `xla_shard_shape = xla::hlo_sharding_util::TileShape((*output_hlo_shardings)[i], xla_shape)` for sharded outputs, but then initialized `shard_shape` from the unsharded `xla_shape.dimensions()` instead of `xla_shard_shape.dimensions()`.

Use `xla_shard_shape.dimensions()` and add unit test assertions in `xla/python/pjrt_ifrt/xla_executable_impl_test_lib.cc` checking `output_specs(i).shape()` and `output_specs(i).shard_shape()`.

PiperOrigin-RevId: 996811910
---
 third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc       | 2 +-
 .../xla/xla/python/pjrt_ifrt/xla_executable_impl_test_lib.cc  | 4 ++++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc
index 4c14b6041e6..abc6d30dbbe 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc
@@ -559,7 +559,7 @@ absl::StatusOr<std::string> PjRtExecutable::CommonMetadata::Serialize(
         xla::Shape xla_shape(element_type, output_shapes[i].dims());
         xla::Shape xla_shard_shape = xla::hlo_sharding_util::TileShape(
             (*output_hlo_shardings)[i], xla_shape);
-        shard_shape = Shape(xla_shape.dimensions());
+        shard_shape = Shape(xla_shard_shape.dimensions());
       }
     }
     *output_spec.mutable_shard_shape() = shard_shape->ToProto(serdes_version);
diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_executable_impl_test_lib.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_executable_impl_test_lib.cc
index ad7061c97fe..08923015c8c 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/xla_executable_impl_test_lib.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/xla_executable_impl_test_lib.cc
@@ -1065,6 +1065,10 @@ TEST(ExecutableTest, ExecutableSerialization) {
       EXPECT_EQ(pjrt_layout->ToString(), output_layouts[i]->ToString());
     }
   }
+  EXPECT_THAT(metadata.output_specs(0).shape().dims(), ElementsAre(2, 3));
+  EXPECT_THAT(metadata.output_specs(0).shard_shape().dims(), ElementsAre(2, 3));
+  EXPECT_THAT(metadata.output_specs(1).shape().dims(), ElementsAre(2, 3));
+  EXPECT_THAT(metadata.output_specs(1).shard_shape().dims(), ElementsAre(1, 3));
 
   int kNumParameters = 2;
   auto parameter_shardings = loaded_executable->GetParameterShardings();

From 99215db921c70b2b002d626cdc359a06bb3d25c8 Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Fri, 9 Oct 2026 16:18:07 -0700
Subject: [PATCH 0396/1000] 
 `AlgebraicSimplifierVisitor::HandleDynamicUpdateSlice` folds a nested
 `dynamic-update-slice` on a `dynamic-slice` of the same base operand into a
 single `dynamic-update-slice`. The pass clamps the outer and inner start
 indices to their slice bounds and adds them in the outer index type. When
 either index uses a narrow integer type such as `s8` or `u8`, the clamp bound
 or the combined sum can overflow that type. An overflow wraps the start
 offset and causes the folded update to write to the wrong coordinates.

Check across every slice dimension that the inner clamp bound fits in the inner index element type using `primitive_util::FitsInIntegralType`. Also check that the maximum combined offset, equal to the base dimension minus the inner update size, fits in the outer index element type. Skip the folding rewrite whenever any dimension fails this check. Existing programs with wide index types are unaffected, while programs that previously suffered integer overflow now preserve the nested operations. No flags gate this bug fix.

Tests:
- `//third_party/tensorflow/compiler/xla/hlo/transforms/simplifiers:algebraic_simplifier_test` pins skipping the fold when narrow start indices would overflow.
PiperOrigin-RevId: 996821793
---
 .../simplifiers/algebraic_simplifier.cc       | 33 ++++++++++++++-----
 .../simplifiers/algebraic_simplifier_test.cc  | 20 +++++++++++
 2 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
index 6ab2eacd57c..81e651adcc9 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
@@ -8683,23 +8683,40 @@ absl::Status AlgebraicSimplifierVisitor::HandleDynamicUpdateSlice(
            absl::MakeConstSpan(dynamic_update_slice->operands()).subspan(2),
            absl::MakeConstSpan(dus_update->operand(0)->operands())
                .subspan(1)))) {
+    const int64_t first_index =
+        Cast<HloDynamicUpdateSliceInstruction>(dynamic_update_slice)
+            ->first_index_operand_number();
+    for (int64_t i = first_index; i < dynamic_update_slice->operand_count();
+         ++i) {
+      if (!primitive_util::FitsInIntegralType(
+              dus_update->shape().dimensions(i - first_index) -
+                  dus_update->operand(1)->shape().dimensions(i - first_index),
+              dus_update->operand(i)->shape().element_type()) ||
+          !primitive_util::FitsInIntegralType(
+              dynamic_update_slice->shape().dimensions(i - first_index) -
+                  dus_update->operand(1)->shape().dimensions(i - first_index),
+              dynamic_update_slice->operand(i)->shape().element_type())) {
+        return absl::OkStatus();
+      }
+    }
     ABSL_RETURN_IF_ERROR(dynamic_update_slice->ReplaceOperandWithDifferentShape(
         1, dus_update->mutable_operand(1)));
-    for (int64_t i = 2; i < dynamic_update_slice->operand_count(); ++i) {
+    for (int64_t i = first_index; i < dynamic_update_slice->operand_count();
+         ++i) {
       HloInstruction* index = dynamic_update_slice->mutable_operand(i);
       index = index->AddInstruction(HloInstruction::CreateTernary(
           index->shape(), HloOpcode::kClamp, MakeScalarLike(index, 0), index,
-          MakeScalarLike(index,
-                         dynamic_update_slice->shape().dimensions(i - 2) -
-                             dus_update->shape().dimensions(i - 2))));
+          MakeScalarLike(
+              index, dynamic_update_slice->shape().dimensions(i - first_index) -
+                         dus_update->shape().dimensions(i - first_index))));
       HloInstruction* inner_index = dus_update->mutable_operand(i);
       inner_index = inner_index->AddInstruction(HloInstruction::CreateTernary(
           inner_index->shape(), HloOpcode::kClamp,
           MakeScalarLike(inner_index, 0), inner_index,
-          MakeScalarLike(
-              inner_index,
-              dus_update->shape().dimensions(i - 2) -
-                  dus_update->operand(1)->shape().dimensions(i - 2))));
+          MakeScalarLike(inner_index,
+                         dus_update->shape().dimensions(i - first_index) -
+                             dus_update->operand(1)->shape().dimensions(
+                                 i - first_index))));
       if (inner_index->shape().element_type() !=
           index->shape().element_type()) {
         inner_index = inner_index->AddInstruction(
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
index 7c99815d8a8..e1b15ed9059 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
@@ -15501,6 +15501,26 @@ TEST_F(AlgebraicSimplifierTest, DynamicUpdateSliceOfDynamicUpdateSlice) {
                                   m::ConstantScalar(2))))));
 }
 
+TEST_F(AlgebraicSimplifierTest,
+       DynamicUpdateSliceOfDynamicUpdateSliceNarrowIndexOverflow) {
+  constexpr absl::string_view hlo_string = R"(
+    HloModule module
+
+    ENTRY test {
+      operand = f32[150] parameter(0)
+      update = f32[10] parameter(1)
+      i = s8[] parameter(2)
+      j = s8[] parameter(3)
+      ds = f32[80] dynamic-slice(operand, i), dynamic_slice_sizes={80}
+      inner_dus = f32[80] dynamic-update-slice(ds, update, j)
+      ROOT outer_dus = f32[150] dynamic-update-slice(operand, inner_dus, i)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  AlgebraicSimplifier simplifier(default_options_);
+  EXPECT_THAT(simplifier.Run(module.get()), absl_testing::IsOkAndHolds(false));
+}
+
 TEST_F(AlgebraicSimplifierTest, FusesShuffleRotates) {
   constexpr absl::string_view kModuleStr = R"(
     HloModule shuffle_module

From 5501fa7ff70bced8306c263f298fa2e2ebb35422 Mon Sep 17 00:00:00 2001
From: Bill Varcho <varcho@google.com>
Date: Fri, 9 Oct 2026 16:26:28 -0700
Subject: [PATCH 0397/1000] [XLA:SPMD][HloShardingV3][Bug-Fix] Avoid unguarded
 tile_assignment() calls on NamedSharding in ShardingPropagation and
 SpmdPartitioner.

PiperOrigin-RevId: 996825649
---
 .../xla/xla/service/sharding_propagation.cc   |  16 ++-
 .../xla/service/sharding_propagation_test.cc  | 136 ++++++++++++++++++
 .../xla/xla/service/spmd/spmd_partitioner.cc  |   6 +-
 .../xla/service/spmd/spmd_partitioner_test.cc |  25 ++++
 4 files changed, 172 insertions(+), 11 deletions(-)

diff --git a/third_party/xla/xla/service/sharding_propagation.cc b/third_party/xla/xla/service/sharding_propagation.cc
index fdda41b226a..a4e83853bd6 100644
--- a/third_party/xla/xla/service/sharding_propagation.cc
+++ b/third_party/xla/xla/service/sharding_propagation.cc
@@ -1275,8 +1275,7 @@ bool InferConvolutionShardingFromOperands(HloInstruction* instruction,
   if (IsConvolutionKernelSmall(instruction)) {
     // If the kernel is small compared to the input then we can generate an
     // output what is sharded the same way as the input.
-    const auto& tile_assignment = lhs->sharding().tile_assignment();
-    if (tile_assignment.dim(dnums.input_feature_dimension()) > 1) {
+    if (lhs->sharding().dimension(dnums.input_feature_dimension()) > 1) {
       return false;
     }
     return MaybeImproveInstructionSharding(get_tiled_sharding_based_on_lhs(),
@@ -1696,8 +1695,7 @@ std::optional<HloSharding> ShardingPropagation::GetShardingFromUser(
       }
 
       const int64_t cdim = user.concatenate_dimension();
-      auto& tile_assignment = user.sharding().tile_assignment();
-      if (tile_assignment.dim(cdim) == 1) {
+      if (user.sharding().dimension(cdim) == 1) {
         // If we are concatenating along a non-sharded dimension then the
         // operands should have the same sharding as the result.
         return user.sharding();
@@ -1709,6 +1707,12 @@ std::optional<HloSharding> ShardingPropagation::GetShardingFromUser(
         return user.sharding();
       }
 
+      HloSharding user_sharding =
+          user.sharding().UseNamedShardingLeaf()
+              ? HloSharding::V3ToV2Sharding(user.sharding().named_sharding())
+              : user.sharding();
+      const auto& tile_assignment = user_sharding.tile_assignment();
+
       // If we are concatenating along a sharded dimension then we want the
       // operands to be distributed among the devices their data is used.
       int64_t start_offset = 0;
@@ -2394,12 +2398,10 @@ bool ShardingPropagation::InferShardingFromOperands(
           if (!AggressiveConcatOperandShardingCanPassThrough(concat_operand)) {
             return false;
           }
-          const auto& tile_assignment =
-              concat_operand->sharding().tile_assignment();
           for (int64_t i = 0; i < instruction->shape().dimensions().size();
                ++i) {
             if (absl::c_linear_search(instruction->dimensions(), i) &&
-                tile_assignment.dim(i) > 1) {
+                concat_operand->sharding().dimension(i) > 1) {
               return false;
             }
           }
diff --git a/third_party/xla/xla/service/sharding_propagation_test.cc b/third_party/xla/xla/service/sharding_propagation_test.cc
index 15d14178440..2e655ca0bcb 100644
--- a/third_party/xla/xla/service/sharding_propagation_test.cc
+++ b/third_party/xla/xla/service/sharding_propagation_test.cc
@@ -12366,5 +12366,141 @@ ENTRY %main.6 (Arg_0.1: s32[8,2]) -> s32[8,2] {
   EXPECT_THAT(sharded_custom_call, op::Sharding("{devices=[4,2]<=[8]}"));
 }
 
+TEST_F(ShardingPropagationTest, ConvolutionWithNamedSharding) {
+  const char* const hlo_string = R"(
+HloModule module
+
+ENTRY conv {
+  %lhs = f32[8,16,512] parameter(0), sharding={mesh['a'=2] [{}, {'a'}, {}]}
+  %rhs = f32[8,2,512] parameter(1)
+  %conv = f32[3,512,512] convolution(%lhs, %rhs),
+    window={size=2 stride=5},
+    dim_labels=f0b_i0o->0bf
+  ROOT %tuple = (f32[3,512,512]) tuple(%conv)
+})";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  ASSERT_OK_AND_ASSIGN(
+      bool changed,
+      ShardingPropagation(/*is_spmd=*/true, /*propagate_metadata=*/false)
+          .Run(module.get()));
+  EXPECT_TRUE(changed);
+  auto* conv = FindInstruction(module.get(), "conv");
+  ASSERT_NE(conv, nullptr);
+  EXPECT_THAT(conv, op::Sharding("{mesh['a'=2] [{'a'}, {}, {}]}"));
+}
+
+TEST_F(ShardingPropagationTest,
+       ConvolutionWithNamedShardingShardedInputFeature) {
+  const char* const hlo_string = R"(
+HloModule module
+
+ENTRY conv {
+  %lhs = f32[8,16,512] parameter(0), sharding={mesh['a'=2,'b'=2] [{'a'}, {'b'}, {}]}
+  %rhs = f32[8,2,512] parameter(1)
+  %conv = f32[3,512,512] convolution(%lhs, %rhs),
+    window={size=2 stride=5},
+    dim_labels=f0b_i0o->0bf
+  ROOT %tuple = (f32[3,512,512]) tuple(%conv)
+})";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  ASSERT_OK_AND_ASSIGN(
+      bool changed,
+      ShardingPropagation(/*is_spmd=*/true, /*propagate_metadata=*/false)
+          .Run(module.get()));
+  EXPECT_FALSE(changed);
+  auto* conv = FindInstruction(module.get(), "conv");
+  ASSERT_NE(conv, nullptr);
+  EXPECT_THAT(conv, op::NoSharding());
+}
+
+TEST_F(ShardingPropagationTest, ConcatenateForwardWithNamedSharding) {
+  const char* const hlo_string = R"(
+HloModule module
+
+ENTRY concat {
+  %param.0 = f32[5,7] parameter(0), sharding={mesh['a'=2] [{'a'}, {}]}
+  %param.1 = f32[5,9] parameter(1), sharding={mesh['a'=2] [{'a'}, {}]}
+  %concat = f32[5,16] concatenate(%param.0, %param.1), dimensions={1}
+  ROOT %tuple = (f32[5,16]) tuple(%concat)
+})";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  ASSERT_OK_AND_ASSIGN(
+      bool changed,
+      ShardingPropagation(/*is_spmd=*/true, /*propagate_metadata=*/false)
+          .Run(module.get()));
+  EXPECT_TRUE(changed);
+  auto* concat = FindInstruction(module.get(), "concat");
+  ASSERT_NE(concat, nullptr);
+  EXPECT_THAT(concat, op::Sharding("{mesh['a'=2] [{'a'}, {}]}"));
+}
+
+TEST_F(ShardingPropagationTest,
+       ConcatenateBackwardWithNamedShardingNonShardedConcatDim) {
+  const char* const hlo_string = R"(
+HloModule module
+
+ENTRY concat {
+  %param.0 = f32[4,7] parameter(0)
+  %copy.0 = f32[4,7] copy(%param.0)
+  %param.1 = f32[4,9] parameter(1)
+  %copy.1 = f32[4,9] copy(%param.1)
+  %concat = f32[4,16] concatenate(%copy.0, %copy.1),
+    dimensions={1}, sharding={mesh['a'=2] [{'a'}, {}]}
+  ROOT %tuple = (f32[4,16]) tuple(%concat)
+})";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  ASSERT_OK_AND_ASSIGN(
+      bool changed,
+      ShardingPropagation(/*is_spmd=*/false, /*propagate_metadata=*/false)
+          .Run(module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_THAT(FindInstruction(module.get(), "copy.0"),
+              op::Sharding("{mesh['a'=2] [{'a'}, {}]}"));
+  EXPECT_THAT(FindInstruction(module.get(), "copy.1"),
+              op::Sharding("{mesh['a'=2] [{'a'}, {}]}"));
+}
+
+TEST_F(ShardingPropagationTest,
+       ConcatenateBackwardWithNamedShardingShardedConcatDim) {
+  const char* const hlo_string = R"(
+HloModule module
+
+ENTRY concat {
+  %param.0 = f32[4,7] parameter(0)
+  %copy.0 = f32[4,7] copy(%param.0)
+  %param.1 = f32[4,7] parameter(1)
+  %copy.1 = f32[4,7] copy(%param.1)
+  %concat = f32[8,7] concatenate(%copy.0, %copy.1),
+    dimensions={0}, sharding={mesh['a'=4] [{'a'}, {}]}
+  ROOT %tuple = (f32[8,7]) tuple(%concat)
+})";
+  {
+    ASSERT_OK_AND_ASSIGN(auto spmd_module,
+                         ParseAndReturnVerifiedModule(hlo_string));
+    ASSERT_OK_AND_ASSIGN(
+        bool changed,
+        ShardingPropagation(/*is_spmd=*/true, /*propagate_metadata=*/false)
+            .Run(spmd_module.get()));
+    EXPECT_TRUE(changed);
+    EXPECT_THAT(FindInstruction(spmd_module.get(), "copy.0"),
+                op::Sharding("{mesh['a'=4] [{'a'}, {}]}"));
+    EXPECT_THAT(FindInstruction(spmd_module.get(), "copy.1"),
+                op::Sharding("{mesh['a'=4] [{'a'}, {}]}"));
+  }
+  {
+    ASSERT_OK_AND_ASSIGN(auto non_spmd_module,
+                         ParseAndReturnVerifiedModule(hlo_string));
+    ASSERT_OK_AND_ASSIGN(
+        bool changed,
+        ShardingPropagation(/*is_spmd=*/false, /*propagate_metadata=*/false)
+            .Run(non_spmd_module.get()));
+    EXPECT_TRUE(changed);
+    EXPECT_THAT(FindInstruction(non_spmd_module.get(), "copy.0"),
+                op::Sharding("{devices=[2,1]0,1}"));
+    EXPECT_THAT(FindInstruction(non_spmd_module.get(), "copy.1"),
+                op::Sharding("{devices=[2,1]2,3}"));
+  }
+}
+
 }  // namespace
 }  // namespace xla
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.cc b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
index 3d7a5863bf0..b4e9dc67d53 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
@@ -5327,12 +5327,10 @@ SpmdPartitioningVisitor::ProcessUpdatePieceExtractOperand(
               }
             }
             if (is_communication_free) {
-              for (int64_t i = 0;
-                   i < rep_sharding.tile_assignment().dimensions().size();
-                   ++i) {
+              for (int64_t i = 0; i < rep_sharding.TiledDataRank(); ++i) {
                 if (std::find(post_to_pre.begin(), post_to_pre.end(), i) ==
                     post_to_pre.end()) {
-                  if (rep_sharding.tile_assignment().dimensions()[i] > 1) {
+                  if (rep_sharding.dimension(i) > 1) {
                     is_communication_free = false;
                     break;
                   }
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
index 85d1bfed107..5cb9571e86c 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
@@ -9269,6 +9269,31 @@ ENTRY entry {
   EXPECT_THAT(root, op::Select());
 }
 
+TEST_P(SpmdPartitioningTest,
+       DynamicUpdateSliceOfReshapeOfSliceShardedWithEnzymeOpt) {
+  absl::string_view hlo_string = R"(
+HloModule module
+
+ENTRY entry {
+  %operand = f32[1,4,4] parameter(1), sharding={devices=[1,2,1]<=[2]}
+  %slice = f32[1,2,4] slice(%operand), slice={[0:1], [0:2], [0:4]}, sharding={devices=[1,2,1]<=[2]}
+  %reshape = f32[2,4] reshape(%slice), sharding={devices=[2,1]<=[2]}
+  %input = f32[4,4] parameter(0), sharding={devices=[2,1]<=[2]}
+  %constant.0 = s32[] constant(0)
+  %constant.1 = s32[] constant(0)
+  ROOT %dus = f32[4,4] dynamic-update-slice(%input, %reshape, %constant.0, %constant.1), sharding={devices=[2,1]<=[2]}
+})";
+
+  ASSERT_OK_AND_ASSIGN(auto module,
+                       PartitionComputation(hlo_string, /*num_devices=*/2,
+                                            SpmdPartitionerOptions(),
+                                            /*enable_enzyme_opt=*/true));
+  VLOG(1) << module->ToString();
+
+  const auto root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(root, op::Select());
+}
+
 TEST_P(SpmdPartitioningTest, DynamicUpdateSliceOfConstantInRange) {
   absl::string_view hlo_string = R"(
   HloModule module

From f8f46a6a9a0f4a26b1be5b9259de126f08c4f766 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Fri, 9 Oct 2026 16:58:07 -0700
Subject: [PATCH 0398/1000] PR #3024: Improve precision of CHLO `asin` fallback
 expansion

Imported from GitHub PR https://github.com/openxla/stablehlo/pull/3024

Replaces the half-angle formula in the CHLO `asin` fallback decomposition with `atan2(x, sqrt((1 - x) * (1 + x)))` to prevent underflow near zero under FTZ mode.

This improves the exhaustive f32 ULP error in JAX's tests on CPU to 1.5 ULP (down from 8388608 ULP in the worst f32 case).
Copybara import of the project:

--
252615efe7c8a8ae99a773a903ba09843c080710 by Peter Hawkins <phawkins@google.com>:

Improve precision of CHLO `asin` fallback expansion

Replaces the half-angle formula in the CHLO `asin` fallback decomposition with `atan2(x, sqrt((1 - x) * (1 + x)))` to prevent underflow near zero under FTZ mode.

Merging this change closes #3024

PiperOrigin-RevId: 996840812
---
 .../xla/third_party/stablehlo/temporary.patch | 88 ++++++++++++++++++-
 .../Dialect/chlo/chlo_legalize_to_mhlo.mlir   | 24 ++---
 2 files changed, 95 insertions(+), 17 deletions(-)

diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch
index 680520db88d..c0bd29a0066 100644
--- a/third_party/xla/third_party/stablehlo/temporary.patch
+++ b/third_party/xla/third_party/stablehlo/temporary.patch
@@ -4626,7 +4626,59 @@ diff --ruN a/stablehlo/stablehlo/tests/TestUtils.cpp b/stablehlo/stablehlo/tests
 diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir
 --- stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir
 +++ stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir
-@@ -5424,6 +5424,55 @@
+@@ -8,10 +8,8 @@
+ // CHECK:           %[[ADD_0:.*]] = stablehlo.add %[[CONSTANT_0]], %[[ARG0]] : tensor<bf16>
+ // CHECK:           %[[MULTIPLY_0:.*]] = stablehlo.multiply %[[SUBTRACT_0]], %[[ADD_0]] : tensor<bf16>
+ // CHECK:           %[[SQRT_0:.*]] = stablehlo.sqrt %[[MULTIPLY_0]] : tensor<bf16>
+-// CHECK:           %[[ADD_1:.*]] = stablehlo.add %[[CONSTANT_0]], %[[SQRT_0]] : tensor<bf16>
+-// CHECK:           %[[VAL_0:.*]] = stablehlo.atan2 %[[ARG0]], %[[ADD_1]] : tensor<bf16>
+-// CHECK:           %[[ADD_2:.*]] = stablehlo.add %[[VAL_0]], %[[VAL_0]] : tensor<bf16>
+-// CHECK:           return %[[ADD_2]] : tensor<bf16>
++// CHECK:           %[[VAL_0:.*]] = stablehlo.atan2 %[[ARG0]], %[[SQRT_0]] : tensor<bf16>
++// CHECK:           return %[[VAL_0]] : tensor<bf16>
+ // CHECK:         }
+ func.func @asin_bf16(%arg : tensor<bf16>) -> tensor<bf16> {
+   %result = "chlo.asin"(%arg) : (tensor<bf16>) -> tensor<bf16>
+@@ -27,10 +25,8 @@
+ // CHECK:           %[[ADD_0:.*]] = stablehlo.add %[[CONSTANT_0]], %[[ARG0]] : tensor<f16>
+ // CHECK:           %[[MULTIPLY_0:.*]] = stablehlo.multiply %[[SUBTRACT_0]], %[[ADD_0]] : tensor<f16>
+ // CHECK:           %[[SQRT_0:.*]] = stablehlo.sqrt %[[MULTIPLY_0]] : tensor<f16>
+-// CHECK:           %[[ADD_1:.*]] = stablehlo.add %[[CONSTANT_0]], %[[SQRT_0]] : tensor<f16>
+-// CHECK:           %[[VAL_0:.*]] = stablehlo.atan2 %[[ARG0]], %[[ADD_1]] : tensor<f16>
+-// CHECK:           %[[ADD_2:.*]] = stablehlo.add %[[VAL_0]], %[[VAL_0]] : tensor<f16>
+-// CHECK:           return %[[ADD_2]] : tensor<f16>
++// CHECK:           %[[VAL_0:.*]] = stablehlo.atan2 %[[ARG0]], %[[SQRT_0]] : tensor<f16>
++// CHECK:           return %[[VAL_0]] : tensor<f16>
+ // CHECK:         }
+ func.func @asin_f16(%arg : tensor<f16>) -> tensor<f16> {
+   %result = "chlo.asin"(%arg) : (tensor<f16>) -> tensor<f16>
+@@ -46,10 +42,8 @@
+ // CHECK:           %[[ADD_0:.*]] = stablehlo.add %[[CONSTANT_0]], %[[ARG0]] : tensor<f32>
+ // CHECK:           %[[MULTIPLY_0:.*]] = stablehlo.multiply %[[SUBTRACT_0]], %[[ADD_0]] : tensor<f32>
+ // CHECK:           %[[SQRT_0:.*]] = stablehlo.sqrt %[[MULTIPLY_0]] : tensor<f32>
+-// CHECK:           %[[ADD_1:.*]] = stablehlo.add %[[CONSTANT_0]], %[[SQRT_0]] : tensor<f32>
+-// CHECK:           %[[VAL_0:.*]] = stablehlo.atan2 %[[ARG0]], %[[ADD_1]] : tensor<f32>
+-// CHECK:           %[[ADD_2:.*]] = stablehlo.add %[[VAL_0]], %[[VAL_0]] : tensor<f32>
+-// CHECK:           return %[[ADD_2]] : tensor<f32>
++// CHECK:           %[[VAL_0:.*]] = stablehlo.atan2 %[[ARG0]], %[[SQRT_0]] : tensor<f32>
++// CHECK:           return %[[VAL_0]] : tensor<f32>
+ // CHECK:         }
+ func.func @asin_f32(%arg : tensor<f32>) -> tensor<f32> {
+   %result = "chlo.asin"(%arg) : (tensor<f32>) -> tensor<f32>
+@@ -65,10 +59,8 @@
+ // CHECK:           %[[ADD_0:.*]] = stablehlo.add %[[CONSTANT_0]], %[[ARG0]] : tensor<f64>
+ // CHECK:           %[[MULTIPLY_0:.*]] = stablehlo.multiply %[[SUBTRACT_0]], %[[ADD_0]] : tensor<f64>
+ // CHECK:           %[[SQRT_0:.*]] = stablehlo.sqrt %[[MULTIPLY_0]] : tensor<f64>
+-// CHECK:           %[[ADD_1:.*]] = stablehlo.add %[[CONSTANT_0]], %[[SQRT_0]] : tensor<f64>
+-// CHECK:           %[[VAL_0:.*]] = stablehlo.atan2 %[[ARG0]], %[[ADD_1]] : tensor<f64>
+-// CHECK:           %[[ADD_2:.*]] = stablehlo.add %[[VAL_0]], %[[VAL_0]] : tensor<f64>
+-// CHECK:           return %[[ADD_2]] : tensor<f64>
++// CHECK:           %[[VAL_0:.*]] = stablehlo.atan2 %[[ARG0]], %[[SQRT_0]] : tensor<f64>
++// CHECK:           return %[[VAL_0]] : tensor<f64>
+ // CHECK:         }
+ func.func @asin_f64(%arg : tensor<f64>) -> tensor<f64> {
+   %result = "chlo.asin"(%arg) : (tensor<f64>) -> tensor<f64>
+@@ -5424,6 +5416,55 @@
  
  // -----
  
@@ -12408,6 +12460,40 @@ diff --ruN a/stablehlo/stablehlo/tools/StablehloTranslateMain.cpp b/stablehlo/st
    return failed(
        mlir::mlirTranslateMain(argc, argv, "StableHLO interpreter driver\n"));
  }
+diff --ruN a/stablehlo/stablehlo/transforms/ChloDecompositionPatternsMath.td b/stablehlo/stablehlo/transforms/ChloDecompositionPatternsMath.td
+--- stablehlo/stablehlo/transforms/ChloDecompositionPatternsMath.td
++++ stablehlo/stablehlo/transforms/ChloDecompositionPatternsMath.td
+@@ -404,23 +404,19 @@
+ 
+ // Arcus sine on real input:
+ //
+-//     arcsin(x) = 2 * arctan2(x, 1 + sqrt(1 - x * x))
++//     arcsin(x) = atan2(x, sqrt((1 - x) * (1 + x)))
+ //
+ //     To avoid cancellation errors at abs(x) close to 1, we'll use
+ //
+ //       1 - x * x == (1 - x) * (1 + x)
+ //
+ def : Pat<(CHLO_AsinOp NonComplexElementType:$x),
+-  (StableHLO_AddOp
+-    (StableHLO_Atan2Op:$ta
+-      $x,
+-      (StableHLO_AddOp
+-        (StableHLO_ConstantLike<"1">:$one $x),
+-        (StableHLO_SqrtOp
+-          (StableHLO_MulOp
+-            (StableHLO_SubtractOp $one, $x),
+-            (StableHLO_AddOp $one, $x)), ConstDefaultResultAccuracyAttr))),
+-    $ta)>;
++  (StableHLO_Atan2Op
++    $x,
++    (StableHLO_SqrtOp
++      (StableHLO_MulOp
++        (StableHLO_SubtractOp (StableHLO_ConstantLike<"1">:$one $x), $x),
++        (StableHLO_AddOp $one, $x)), ConstDefaultResultAccuracyAttr))>;
+ 
+ // Arcus cosine on complex input
+ //
 diff --ruN a/stablehlo/stablehlo/transforms/ChloDecompositionUtils.h b/stablehlo/stablehlo/transforms/ChloDecompositionUtils.h
 --- stablehlo/stablehlo/transforms/ChloDecompositionUtils.h
 +++ stablehlo/stablehlo/transforms/ChloDecompositionUtils.h
diff --git a/third_party/xla/xla/mlir_hlo/tests/Dialect/chlo/chlo_legalize_to_mhlo.mlir b/third_party/xla/xla/mlir_hlo/tests/Dialect/chlo/chlo_legalize_to_mhlo.mlir
index 14df651477f..09752e3d03b 100644
--- a/third_party/xla/xla/mlir_hlo/tests/Dialect/chlo/chlo_legalize_to_mhlo.mlir
+++ b/third_party/xla/xla/mlir_hlo/tests/Dialect/chlo/chlo_legalize_to_mhlo.mlir
@@ -22,10 +22,8 @@
 // CHECK-NEXT:     %[[TMP_3:.*]] = mhlo.add %[[TMP_1]], %[[TMP_arg0]] : tensor<bf16>
 // CHECK-NEXT:     %[[TMP_4:.*]] = mhlo.multiply %[[TMP_2]], %[[TMP_3]] : tensor<bf16>
 // CHECK-NEXT:     %[[TMP_5:.*]] = mhlo.sqrt %[[TMP_4]] : tensor<bf16>
-// CHECK-NEXT:     %[[TMP_6:.*]] = mhlo.add %[[TMP_1]], %[[TMP_5]] : tensor<bf16>
-// CHECK-NEXT:     %[[TMP_7:.*]] = mhlo.atan2 %[[TMP_arg0]], %[[TMP_6]] : tensor<bf16>
-// CHECK-NEXT:     %[[TMP_8:.*]] = mhlo.add %[[TMP_7]], %[[TMP_7]] : tensor<bf16>
-// CHECK-NEXT:     return %[[TMP_8]] : tensor<bf16>
+// CHECK-NEXT:     %[[TMP_6:.*]] = mhlo.atan2 %[[TMP_arg0]], %[[TMP_5]] : tensor<bf16>
+// CHECK-NEXT:     return %[[TMP_6]] : tensor<bf16>
 func.func @asin_bf16(%arg : tensor<bf16>) -> tensor<bf16> {
   %result = "chlo.asin"(%arg) : (tensor<bf16>) -> tensor<bf16>
   func.return %result : tensor<bf16>
@@ -41,10 +39,8 @@ func.func @asin_bf16(%arg : tensor<bf16>) -> tensor<bf16> {
 // CHECK-NEXT:    %[[TMP_3:.*]] = mhlo.add %[[TMP_1]], %[[TMP_arg0]] : tensor<f16>
 // CHECK-NEXT:    %[[TMP_4:.*]] = mhlo.multiply %[[TMP_2]], %[[TMP_3]] : tensor<f16>
 // CHECK-NEXT:    %[[TMP_5:.*]] = mhlo.sqrt %[[TMP_4]] : tensor<f16>
-// CHECK-NEXT:    %[[TMP_6:.*]] = mhlo.add %[[TMP_1]], %[[TMP_5]] : tensor<f16>
-// CHECK-NEXT:    %[[TMP_7:.*]] = mhlo.atan2 %[[TMP_arg0]], %[[TMP_6]] : tensor<f16>
-// CHECK-NEXT:    %[[TMP_8:.*]] = mhlo.add %[[TMP_7]], %[[TMP_7]] : tensor<f16>
-// CHECK-NEXT:    return %[[TMP_8]] : tensor<f16>
+// CHECK-NEXT:    %[[TMP_6:.*]] = mhlo.atan2 %[[TMP_arg0]], %[[TMP_5]] : tensor<f16>
+// CHECK-NEXT:    return %[[TMP_6]] : tensor<f16>
 func.func @asin_f16(%arg : tensor<f16>) -> tensor<f16> {
   // CHECK-HIGH-LEVEL: mhlo.asin
   %result = "chlo.asin"(%arg) : (tensor<f16>) -> tensor<f16>
@@ -61,10 +57,8 @@ func.func @asin_f16(%arg : tensor<f16>) -> tensor<f16> {
 // CHECK-NEXT:    %[[TMP_3:.*]] = mhlo.add %[[TMP_1]], %[[TMP_arg0]] : tensor<f32>
 // CHECK-NEXT:    %[[TMP_4:.*]] = mhlo.multiply %[[TMP_2]], %[[TMP_3]] : tensor<f32>
 // CHECK-NEXT:    %[[TMP_5:.*]] = mhlo.sqrt %[[TMP_4]] : tensor<f32>
-// CHECK-NEXT:    %[[TMP_6:.*]] = mhlo.add %[[TMP_1]], %[[TMP_5]] : tensor<f32>
-// CHECK-NEXT:    %[[TMP_7:.*]] = mhlo.atan2 %[[TMP_arg0]], %[[TMP_6]] : tensor<f32>
-// CHECK-NEXT:    %[[TMP_8:.*]] = mhlo.add %[[TMP_7]], %[[TMP_7]] : tensor<f32>
-// CHECK-NEXT:    return %[[TMP_8]] : tensor<f32>
+// CHECK-NEXT:    %[[TMP_6:.*]] = mhlo.atan2 %[[TMP_arg0]], %[[TMP_5]] : tensor<f32>
+// CHECK-NEXT:    return %[[TMP_6]] : tensor<f32>
 func.func @asin_f32(%arg : tensor<f32>) -> tensor<f32> {
   // CHECK-HIGH-LEVEL: mhlo.asin
   %result = "chlo.asin"(%arg) : (tensor<f32>) -> tensor<f32>
@@ -80,10 +74,8 @@ func.func @asin_f32(%arg : tensor<f32>) -> tensor<f32> {
 // CHECK-NEXT:    %[[TMP_3:.*]] = mhlo.add %[[TMP_1]], %[[TMP_arg0]] : tensor<f64>
 // CHECK-NEXT:    %[[TMP_4:.*]] = mhlo.multiply %[[TMP_2]], %[[TMP_3]] : tensor<f64>
 // CHECK-NEXT:    %[[TMP_5:.*]] = mhlo.sqrt %[[TMP_4]] : tensor<f64>
-// CHECK-NEXT:    %[[TMP_6:.*]] = mhlo.add %[[TMP_1]], %[[TMP_5]] : tensor<f64>
-// CHECK-NEXT:    %[[TMP_7:.*]] = mhlo.atan2 %[[TMP_arg0]], %[[TMP_6]] : tensor<f64>
-// CHECK-NEXT:    %[[TMP_8:.*]] = mhlo.add %[[TMP_7]], %[[TMP_7]] : tensor<f64>
-// CHECK-NEXT:    return %[[TMP_8]] : tensor<f64>
+// CHECK-NEXT:    %[[TMP_6:.*]] = mhlo.atan2 %[[TMP_arg0]], %[[TMP_5]] : tensor<f64>
+// CHECK-NEXT:    return %[[TMP_6]] : tensor<f64>
 func.func @asin_f64(%arg : tensor<f64>) -> tensor<f64> {
   // CHECK-HIGH-LEVEL: mhlo.asin
   %result = "chlo.asin"(%arg) : (tensor<f64>) -> tensor<f64>

From 70f3520b8175be93049691f68fd1ce3991cf173b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 9 Oct 2026 17:10:46 -0700
Subject: [PATCH 0399/1000] Ensure LookupTableInsertOp::ComputeAsync invokes
 DoneCallback on all error paths.

Use absl::MakeCleanup in LookupTableInsertOp::ComputeAsync so DoneCallback is always invoked if an early return occurs.

PiperOrigin-RevId: 996846906
---
 .../kernels/data/experimental/lookup_ops.cc   | 22 +++++++++++--------
 1 file changed, 13 insertions(+), 9 deletions(-)

diff --git a/tensorflow/core/kernels/data/experimental/lookup_ops.cc b/tensorflow/core/kernels/data/experimental/lookup_ops.cc
index 470e9ed8e44..c1ea8df7ada 100644
--- a/tensorflow/core/kernels/data/experimental/lookup_ops.cc
+++ b/tensorflow/core/kernels/data/experimental/lookup_ops.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include "tensorflow/core/framework/dataset.h"
 #include "tensorflow/core/framework/function_handle_cache.h"
 #include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
@@ -35,6 +36,7 @@ limitations under the License.
 #include "tensorflow/core/kernels/lookup_util.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
 #include "tensorflow/core/lib/io/inputbuffer.h"
 #include "tensorflow/core/lib/strings/numbers.h"
 #include "tensorflow/core/lib/strings/str_util.h"
@@ -152,41 +154,41 @@ void InitializeTableFromDataset(OpKernelContext* ctx,
   OP_REQUIRES_ASYNC(
       ctx, dataset_types.size() == 2,
       absl::InvalidArgumentError("Dataset should have two output types only"),
-      done);
+      []() {});
   OP_REQUIRES_ASYNC(
       ctx, dataset_types[0] == table->key_dtype(),
       absl::InvalidArgumentError(absl::StrCat(
           "Key dtype expected: ", DataTypeString(table->key_dtype()),
           " but obtained: ", DataTypeString(dataset_types[0]),
           " from the dataset")),
-      done);
+      []() {});
   OP_REQUIRES_ASYNC(
       ctx, dataset_types[1] == table->value_dtype(),
       absl::InvalidArgumentError(absl::StrCat(
           "Value dtype expected: ", DataTypeString(table->value_dtype()),
           " but obtained: ", DataTypeString(dataset_types[1]),
           " from the dataset")),
-      done);
+      []() {});
   // Assert that the dataset output shapes are scalars.
   const auto& dataset_shapes = dataset->output_shapes();
   OP_REQUIRES_ASYNC(
       ctx, dataset_shapes.size() == 2,
       absl::InvalidArgumentError("Dataset should have two output shapes only"),
-      done);
+      []() {});
   OP_REQUIRES_ASYNC(ctx,
                     dataset_shapes[0].IsCompatibleWith(PartialTensorShape({})),
                     absl::InvalidArgumentError(
                         absl::StrCat("Expected scalar for key. Obtained: ",
                                      dataset_shapes[0].DebugString())),
-                    done);
+                    []() {});
   OP_REQUIRES_ASYNC(ctx,
                     dataset_shapes[1].IsCompatibleWith(PartialTensorShape({})),
                     absl::InvalidArgumentError(
                         absl::StrCat("Expected scalar for key. Obtained: ",
                                      dataset_shapes[1].DebugString())),
-                    done);
+                    []() {});
   DatasetIterator iter(dataset);
-  OP_REQUIRES_OK_ASYNC(ctx, iter.Init(ctx), done);
+  OP_REQUIRES_OK_ASYNC(ctx, iter.Init(ctx), []() {});
   absl::Status s =
       table->Initialize(iter, MakeDatasetInitializerSerializer(ctx, dataset));
   if (absl::IsFailedPrecondition(s) && table->is_initialized()) {
@@ -203,13 +205,15 @@ class InitializeTableFromDatasetOp : public AsyncOpKernel {
         background_worker_(ctx->env(), "initialize_table_from_dataset") {}
 
   void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override {
+    auto cleanup = gtl::MakeCleanup([&done]() { done(); });
     lookup::InitializableLookupTable* table;
     OP_REQUIRES_OK_ASYNC(
-        ctx, GetInitializableLookupTable("table_handle", ctx, &table), done);
+        ctx, GetInitializableLookupTable("table_handle", ctx, &table), []() {});
     core::ScopedUnref unref_me(table);
     data::DatasetBase* dataset;
     OP_REQUIRES_OK_ASYNC(
-        ctx, GetDatasetFromVariantTensor(ctx->input(1), &dataset), done);
+        ctx, GetDatasetFromVariantTensor(ctx->input(1), &dataset), []() {});
+    cleanup.release();
     background_worker_.Schedule([ctx, dataset, table, done]() {
       InitializeTableFromDataset(ctx, dataset, table, done);
     });

From 21f39e2eccdfbac631aaf6c5073127cdedb90906 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 9 Oct 2026 17:37:38 -0700
Subject: [PATCH 0400/1000] Reverts 8e6a81187d36bfd6d601e06f3c911cbc94953e95

PiperOrigin-RevId: 996857745
---
 .../xla/xla/hlo/transforms/collectives/BUILD  |  2 -
 .../all_gather_broadcast_reorder_test.cc      | 89 -------------------
 2 files changed, 91 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/collectives/BUILD b/third_party/xla/xla/hlo/transforms/collectives/BUILD
index 0011e156f6d..5692dd75556 100644
--- a/third_party/xla/xla/hlo/transforms/collectives/BUILD
+++ b/third_party/xla/xla/hlo/transforms/collectives/BUILD
@@ -267,10 +267,8 @@ xla_cc_test(
     srcs = ["all_gather_broadcast_reorder_test.cc"],
     deps = [
         ":all_gather_broadcast_reorder",
-        "//xla/hlo/testlib:filecheck",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
         "//xla/hlo/utils:hlo_matchers",
-        "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest_main",
         "@tsl//tsl/platform:statusor",
diff --git a/third_party/xla/xla/hlo/transforms/collectives/all_gather_broadcast_reorder_test.cc b/third_party/xla/xla/hlo/transforms/collectives/all_gather_broadcast_reorder_test.cc
index fc7921c4938..6e4a14b3a1c 100644
--- a/third_party/xla/xla/hlo/transforms/collectives/all_gather_broadcast_reorder_test.cc
+++ b/third_party/xla/xla/hlo/transforms/collectives/all_gather_broadcast_reorder_test.cc
@@ -17,9 +17,7 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
-#include "absl/status/status_matchers.h"
 #include "absl/strings/string_view.h"
-#include "xla/hlo/testlib/filecheck.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/hlo/utils/hlo_matchers.h"
 #include "tsl/platform/statusor.h"
@@ -27,8 +25,6 @@ limitations under the License.
 namespace xla {
 namespace {
 
-using ::absl_testing::IsOkAndHolds;
-
 namespace m = xla::testing::opcode_matchers;
 
 class AllGatherBroadcastReorderTest : public HloHardwareIndependentTestBase {
@@ -122,90 +118,5 @@ ENTRY main {
   RunPass(hlo_string, PassOutput::NoChange);
 }
 
-TEST_F(AllGatherBroadcastReorderTest, DoesNotCrossCallBoundaries) {
-  absl::string_view hlo_string = R"(
-HloModule m
-
-foo {
-  param = f32[5, 4, 8, 128] parameter(0)
-  ROOT ag = f32[5, 4, 8, 256] all-gather(param), dimensions={3}, replica_groups={{0, 1}}
-}
-
-ENTRY main {
-  x = f32[128, 5] parameter(0)
-  bc = f32[5, 4, 8, 128] broadcast(x), dimensions={3, 0}
-  ROOT call = f32[5, 4, 8, 256] call(bc), to_apply=foo
-}
-)";
-  RunPass(hlo_string, PassOutput::NoChange);
-}
-
-TEST_F(AllGatherBroadcastReorderTest, EntryCallsFooTransformationOnFoo) {
-  absl::string_view hlo_string = R"(
-// CHECK: %foo (
-// CHECK: %[[AG:.*]] = f32[256,5]{{.*}} all-gather(%x_foo), {{.*}}dimensions={0}
-// CHECK: ROOT %[[BC:.*]] = f32[5,4,8,256]{{.*}} broadcast(%[[AG]]), dimensions={3,0}
-// CHECK: }
-// CHECK: ENTRY %main (
-// CHECK: ROOT %[[CALL:.*]] = f32[5,4,8,256]{{.*}} call(%x_main), to_apply=%foo
-HloModule m
-
-foo {
-  x_foo = f32[128, 5] parameter(0)
-  bc = f32[5, 4, 8, 128] broadcast(x_foo), dimensions={3, 0}
-  ROOT ag = f32[5, 4, 8, 256] all-gather(bc), dimensions={3}, replica_groups={{0, 1}}
-}
-
-ENTRY main {
-  x_main = f32[128, 5] parameter(0)
-  ROOT call = f32[5, 4, 8, 256] call(x_main), to_apply=foo
-}
-)";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
-  EXPECT_THAT(AllGatherBroadcastReorder().Run(module.get()),
-              IsOkAndHolds(true));
-  EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true));
-}
-
-TEST_F(AllGatherBroadcastReorderTest,
-       EntryCallsFooAndBarFooCallsBarTransformationOnBar) {
-  absl::string_view hlo_string = R"(
-// CHECK: %bar (
-// CHECK: %[[AG:.*]] = f32[256,5]{{.*}} all-gather(%x_bar), {{.*}}dimensions={0}
-// CHECK: ROOT %[[BC:.*]] = f32[5,4,8,256]{{.*}} broadcast(%[[AG]]), dimensions={3,0}
-// CHECK: }
-// CHECK: %foo (
-// CHECK: ROOT %[[CALL_BAR:.*]] = f32[5,4,8,256]{{.*}} call(%x_foo), to_apply=%bar
-// CHECK: }
-// CHECK: ENTRY %main (
-// CHECK: %[[CALL_FOO:.*]] = f32[5,4,8,256]{{.*}} call(%x_main), to_apply=%foo
-// CHECK: %[[CALL_BAR2:.*]] = f32[5,4,8,256]{{.*}} call(%x_main), to_apply=%bar
-// CHECK: ROOT %[[ADD:.*]] = f32[5,4,8,256]{{.*}} add(%[[CALL_FOO]], %[[CALL_BAR2]])
-HloModule m
-
-bar {
-  x_bar = f32[128, 5] parameter(0)
-  bc = f32[5, 4, 8, 128] broadcast(x_bar), dimensions={3, 0}
-  ROOT ag = f32[5, 4, 8, 256] all-gather(bc), dimensions={3}, replica_groups={{0, 1}}
-}
-
-foo {
-  x_foo = f32[128, 5] parameter(0)
-  ROOT call_bar = f32[5, 4, 8, 256] call(x_foo), to_apply=bar
-}
-
-ENTRY main {
-  x_main = f32[128, 5] parameter(0)
-  call_foo = f32[5, 4, 8, 256] call(x_main), to_apply=foo
-  call_bar = f32[5, 4, 8, 256] call(x_main), to_apply=bar
-  ROOT add = f32[5, 4, 8, 256] add(call_foo, call_bar)
-}
-)";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
-  EXPECT_THAT(AllGatherBroadcastReorder().Run(module.get()),
-              IsOkAndHolds(true));
-  EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true));
-}
-
 }  // namespace
 }  // namespace xla

From ba7bcf1ba45194109d4ecebf2fe4e8f91a6fcdfa Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 9 Oct 2026 19:14:19 -0700
Subject: [PATCH 0401/1000] Ensure RpcCheckStatusOp and RpcGetValueOp invoke
 DoneCallback on all error paths.

Use gtl::MakeCleanup in RpcCheckStatusOp::ComputeAsync and RpcGetValueOp::ComputeAsync so DoneCallback is always invoked if LookupResource fails early.

PiperOrigin-RevId: 996890198
---
 .../experimental/rpc/kernels/rpc_ops.cc          | 11 +++++++----
 .../distribute/experimental/rpc/rpc_ops_test.py  | 16 ++++++++++++++++
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/tensorflow/distribute/experimental/rpc/kernels/rpc_ops.cc b/tensorflow/distribute/experimental/rpc/kernels/rpc_ops.cc
index 235c40203cf..fb4863cb641 100644
--- a/tensorflow/distribute/experimental/rpc/kernels/rpc_ops.cc
+++ b/tensorflow/distribute/experimental/rpc/kernels/rpc_ops.cc
@@ -54,6 +54,7 @@ limitations under the License.
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/framework/variant.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
 #include "tensorflow/core/lib/gtl/flatmap.h"
 #include "tensorflow/core/platform/env.h"
 #include "tensorflow/core/platform/mutex.h"
@@ -770,6 +771,7 @@ RpcCheckStatusOp::RpcCheckStatusOp(OpKernelConstruction* ctx)
     : AsyncOpKernel(ctx) {}
 
 void RpcCheckStatusOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
+  auto cleanup = gtl::MakeCleanup([&done]() { done(); });
   core::RefCountPtr<RpcFutureResource> future_resource;
   auto handle = HandleFromInput(ctx, 0);
   {
@@ -779,14 +781,14 @@ void RpcCheckStatusOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
         ctx->SetStatus(absl::NotFoundError(
             "Future resource no longer exists. Please make sure "
             "resource is not already deleted."));
-        done();
-        return;
       } else {
         ctx->SetStatus(status);
       }
+      return;
     }
   }
 
+  cleanup.release();
   future_resource->AddDoneCallback(
       [ctx, done, handle](const absl::Status& status,
                           const CallResponse& response) {
@@ -805,6 +807,7 @@ void RpcCheckStatusOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
 RpcGetValueOp::RpcGetValueOp(OpKernelConstruction* ctx) : AsyncOpKernel(ctx) {}
 
 void RpcGetValueOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
+  auto cleanup = gtl::MakeCleanup([&done]() { done(); });
   core::RefCountPtr<RpcFutureResource> future_resource;
   auto handle = HandleFromInput(ctx, 0);
   {
@@ -814,14 +817,14 @@ void RpcGetValueOp::ComputeAsync(OpKernelContext* ctx, DoneCallback done) {
         ctx->SetStatus(absl::NotFoundError(
             "Future resource no longer exists. Please ensure "
             "resource is not already deleted."));
-        done();
-        return;
       } else {
         ctx->SetStatus(status);
       }
+      return;
     }
   }
 
+  cleanup.release();
   future_resource->AddDoneCallback([ctx, done, handle](
                                        const absl::Status& status,
                                        const CallResponse& response) {
diff --git a/tensorflow/python/distribute/experimental/rpc/rpc_ops_test.py b/tensorflow/python/distribute/experimental/rpc/rpc_ops_test.py
index d239e07fb85..18a4f17b439 100644
--- a/tensorflow/python/distribute/experimental/rpc/rpc_ops_test.py
+++ b/tensorflow/python/distribute/experimental/rpc/rpc_ops_test.py
@@ -93,6 +93,22 @@ class RpcOpsTest(test.TestCase):
 
     rpc_ops.gen_rpc_ops.delete_rpc_future_resource(future_resource, deleter)
 
+    with self.assertRaisesRegex(
+        errors.NotFoundError, "Future resource no longer exists"
+    ):
+      rpc_ops.gen_rpc_ops.rpc_check_status(future_resource)
+
+    with self.assertRaisesRegex(
+        errors.NotFoundError, "Future resource no longer exists"
+    ):
+      rpc_ops.gen_rpc_ops.rpc_get_value(future_resource, Tout=[dtypes.int32])
+
+    with self.assertRaises(errors.InvalidArgumentError):
+      rpc_ops.gen_rpc_ops.rpc_check_status(server_resource)
+
+    with self.assertRaises(errors.InvalidArgumentError):
+      rpc_ops.gen_rpc_ops.rpc_get_value(server_resource, Tout=[dtypes.int32])
+
   def test_exported_rpc_api_static_factory(self):
 
     @eager_def_function.function(input_signature=[

From 383d4dde3228be949f48e4cf2bc960174a3e8cf6 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Sat, 10 Oct 2026 08:23:48 +0530
Subject: [PATCH 0402/1000] Refactor DVariable for improved layout handling

---
 tensorflow/dtensor/python/d_variable.py | 259 ++++++++++++++++++++++++
 1 file changed, 259 insertions(+)

diff --git a/tensorflow/dtensor/python/d_variable.py b/tensorflow/dtensor/python/d_variable.py
index fe9a5ff8c8c..0c8884eb87a 100644
--- a/tensorflow/dtensor/python/d_variable.py
+++ b/tensorflow/dtensor/python/d_variable.py
@@ -13,6 +13,21 @@
 # limitations under the License.
 # ==============================================================================
 """DTensor variable and saveable."""
+# Copyright 2022 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.
+# ==============================================================================
+"""DTensor variable and saveable."""
 
 import functools
 
@@ -140,6 +155,250 @@ class _DVariableSaveable(saveable_object.SaveableObject):
         .save_as_bf16 else tensor)
 
 
+@tf_export('experimental.dtensor.DVariable', v1=[])
+class DVariable(resource_variable_ops.ResourceVariable):
+  """A replacement for tf.Variable which follows initial value placement.
+
+    The class also handles restore/save operations in DTensor. Note that,
+    DVariable may fall back to normal tf.Variable at this moment if
+    `initial_value` is not a DTensor.
+  """
+
+  def __init__(self, initial_value, *args, dtype=None, **kwargs):
+    """Overrides tf.Variable to fix VarHandleOp placements."""
+    # Variables by default use the current device scope for placement. This
+    # wrapper has them follow the initial value's placement instead (which will
+    # be the DTensor device if the initial value has a layout).
+
+    # Pop layout from kwargs since keras make_variable may pass a 'layout'
+    # keyword argument. We need to pop it because we are passing kwargs to
+    # super class constructor.
+    layout = kwargs.pop('layout', None)
+    shape = kwargs.get('shape', None)
+
+    if callable(initial_value):
+      unwrapped = initial_value
+      if issubclass(type(initial_value), functools.partial):
+        unwrapped = initial_value.func  # pyrefly: ignore[missing-attribute]
+
+      # If wrapped is a CheckpointInitialValueCallable, this means that
+      # we are creating a Variable during a checkpoint restore.
+      # Thus the restore will happen now through this callable
+      # and we will create the DVariable with the restored dtensor.
+      if issubclass(type(unwrapped), trackable.CheckpointInitialValueCallable):
+        if shape is None or layout is None:
+          raise ValueError('Expected shape and layout to be not None.')
+
+        # CheckpointInitialValueCallable will call an eager tf.RestoreV2,
+        # which does not have any shape information or layout information
+        # attached. Thus we will do two things to have them correctly specified:
+        #
+        # The default layout scope allows us to correctly specify the output
+        # layout of the tf.RestoreV2 that will be called
+        #
+        # Passing shard_info with the correct shape allows the tf.RestoreV2
+        # ShapeInference to extract the shape.
+        initial_value = api.call_with_layout(
+            initial_value,
+            layout,
+            shard_info=trackable.ShardInfo(
+                shape=shape, offset=[0] * len(shape)))
+      else:
+        initial_value = initial_value()
+
+    # When the initial value came from a Checkpoint restoration, fetch tensor.
+    if isinstance(initial_value, trackable.CheckpointInitialValue):
+      initial_value = initial_value.wrapped_value
+
+    initial_value = ops.convert_to_tensor(initial_value, dtype=dtype)
+    variable_device = initial_value.device
+    self._save_as_bf16 = False
+    # TODO(b/159035705): The following code enables variable creation inside
+    # a tf.function. However, it requires a global dtensor device.
+    # if not variable_device and not tf.executing_eagerly():
+    #   try:
+    #     initial_value.op.get_attr("_layout")
+    #   except ValueError:
+    #     pass
+    #   else:
+    #     # The initial value is a DTensor, but because the DTensor device is
+    #     # only active during eager execution at the moment we need to
+    #     # translate that into a placement for the eager VarHandleOp.
+    #     variable_device = _dtensor_device().name
+    with ops.device(variable_device):
+      # If initial tensor assigned to DVariable is DTensor, record the layout of
+      # the resource so that this can be queried.
+      if context.executing_eagerly():
+        if api.is_dtensor(initial_value):
+          value_layout = api.fetch_layout(initial_value)
+          if layout is not None and layout != value_layout:
+            raise errors_impl.InvalidArgumentError(
+                None,
+                None,
+                'Conflicting layout are provided for initial '
+                f'value layout ({value_layout}) and variable ({layout}).',
+            )
+          layout = value_layout
+        elif layout is not None:
+          initial_value = api.relayout(initial_value, layout)
+        else:
+          raise errors_impl.InvalidArgumentError(
+              None,
+              None,
+              'Neither layout nor DTensor initial value are provided.',
+          )
+        self.layout = layout
+        with api.default_mesh(layout.mesh):
+          super(DVariable, self).__init__(
+              initial_value, *args, dtype=dtype, **kwargs
+          )
+      else:
+        # FIXME(175928457): Record value layout in graph mode.
+        if layout is not None:
+          initial_value = api.relayout(initial_value, layout)
+        super(DVariable, self).__init__(
+            initial_value, *args, dtype=dtype, **kwargs)
+
+  @property
+  def save_as_bf16(self):
+    return self._save_as_bf16
+
+  @save_as_bf16.setter
+  def save_as_bf16(self, save_as_bf16):
+    """Enables saving float32 as bfloat16."""
+    self._save_as_bf16 = save_as_bf16 and self.dtype == dtypes.float32
+
+  def _gather_saveables_for_checkpoint(self):
+    return {
+        trackable.VARIABLE_VALUE_KEY:
+            functools.partial(_DVariableSaveable, self)
+    }
+import functools
+
+from tensorflow.dtensor.python import api
+from tensorflow.dtensor.python import layout as layout_lib
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.trackable import base as trackable
+from tensorflow.python.training.saving import saveable_object
+from tensorflow.python.util.tf_export import tf_export
+
+
+class DSaveSpec(saveable_object.SaveSpec):
+  """DTensor SaveSpec that additionaly captures global_shape and layout."""
+
+  def __init__(self,
+               tensor,
+               slice_spec,
+               name,
+               global_shape,
+               layout,
+               dtype=None,
+               device=None):
+    super().__init__(
+        tensor=tensor,
+        slice_spec=slice_spec,
+        name=name,
+        dtype=dtype,
+        device=device)
+    self.global_shape = global_shape
+    self.layout = layout
+
+
+class _DVariableSaveable(saveable_object.SaveableObject):
+  """Class for defining how to save/restore DTensor variable."""
+
+  def __init__(self, dvariable, name):
+    with ops.device(dvariable.device):
+      original_layout = api.fetch_layout(dvariable)
+    # Record original layout to allow restore.
+    self._original_layout = original_layout
+    self._dvariable = dvariable
+
+    def pack(tensors, layout):
+      with ops.device(dvariable.device):
+        return api.pack(tensors, layout)
+
+    host_layout = layout_lib.Layout(original_layout.sharding_specs,
+                                    original_layout.mesh.host_mesh())
+
+    def get_host_dtensor():
+      # Copy to host mesh if needed.
+      if original_layout.mesh.device_type().upper() != 'CPU':
+        # Prefer pack and unpack in eager mode because it supports sharded
+        # layouts.
+        if context.executing_eagerly():
+          host_dtensor = api.pack(
+              api.unpack(dvariable.read_value()), host_layout)
+        else:
+          host_dtensor = api.copy_to_mesh(dvariable.read_value(), host_layout)
+      else:
+        host_dtensor = dvariable.read_value()
+      return (math_ops.cast(host_dtensor, dtypes.bfloat16)
+              if self.should_cast(host_dtensor) else host_dtensor)
+
+    num_local_devices = original_layout.mesh.num_local_devices()
+    super(_DVariableSaveable, self).__init__(
+        None,
+        [
+            DSaveSpec(
+                tensor=get_host_dtensor,
+                slice_spec=pack([''] * num_local_devices,
+                                layout_lib.Layout.replicated(
+                                    original_layout.mesh.host_mesh(), rank=0)),
+                name=pack([name] * num_local_devices,
+                          layout_lib.Layout.replicated(
+                              original_layout.mesh.host_mesh(), rank=0)),
+                global_shape=dvariable.shape,
+                # Layout is attached as attribute, no need to put it as a
+                # Tensor on DTensorDevice.
+                layout=host_layout.to_string(),
+                dtype=dtypes.bfloat16
+                if self.should_cast(dvariable) else dvariable.dtype,
+                device=dvariable.device)
+        ],
+        name)
+
+  def should_cast(self, v):
+    """Returns True if v has float32 dtype and is intructed to save as bf16.
+
+    Args:
+      v : The variable that determines whether to cast.
+
+    Returns:
+      True if current savable DVariable is instructed to save as bfloat16 and
+        the variable has dtype float32.
+    """
+    return self._dvariable.save_as_bf16 and v.dtype == dtypes.float32
+
+  def restore(self, restored_tensors, restored_shapes):
+    """Restores the same value into all variables."""
+    tensor, = restored_tensors
+
+    @def_function.function
+    def _restore(t):
+      with ops.device(self._dvariable.device):
+        return api.copy_to_mesh(t, self._original_layout)
+
+    # This assign establishes connections from restored tensor and tensors
+    # being restored to -- so that restore in SPMD can backtrack the DVariable
+    # and its layout, given that we're using tf.function style restore.
+    # Note that the restored dvaraible is on CPU no matter what as the restoreV2
+    # op must run on CPU.
+    # TODO(b/159035705): Allow restore for Tensor objects as well?
+    # Restore the dvariable back to original layout.
+    if self._original_layout.mesh.device_type().upper() != 'CPU':
+      tensor = _restore(tensor)
+    return self._dvariable.assign(
+        math_ops.cast(tensor, dtype=self._dvariable.dtype) if self._dvariable
+        .save_as_bf16 else tensor)
+
+
 @tf_export('experimental.dtensor.DVariable', v1=[])
 class DVariable(resource_variable_ops.ResourceVariable):
   """A replacement for tf.Variable which follows initial value placement.

From 7b40d8c4b5a784fbbd39199c055539c85e78e4f6 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 9 Oct 2026 19:43:47 -0700
Subject: [PATCH 0403/1000] [XLA] Allow overriding live range analysis for
 Buffer Assignment

HloLiveRange for async ops is correct, but pessimistic when it comes to
context shapes. An implementation may define the async computation of
call-start.1 to execute sequentially after the async computation of
call-start.0, in which case there will be no live range overlap between the
two context shapes, hence no interference.

This change introduces `LiveRangeInterferenceOptions` to specify custom
interference rules for buffer assignment and optimizes active allocation
reuse in `FastMerge`.

PiperOrigin-RevId: 996899242
---
 .../xla/xla/service/buffer_assignment.cc      | 156 +++++++++++--
 .../xla/xla/service/buffer_assignment.h       |  29 +++
 .../xla/xla/service/buffer_assignment_test.cc | 207 ++++++++++++++++++
 3 files changed, 377 insertions(+), 15 deletions(-)

diff --git a/third_party/xla/xla/service/buffer_assignment.cc b/third_party/xla/xla/service/buffer_assignment.cc
index 163071d7d73..bf53792fca7 100644
--- a/third_party/xla/xla/service/buffer_assignment.cc
+++ b/third_party/xla/xla/service/buffer_assignment.cc
@@ -205,8 +205,13 @@ class FastMergeBufferAllocationsManagerForComputationsWithoutOrdering
     : public BufferAllocationsManagerForComputationsWithoutOrdering {
  public:
   FastMergeBufferAllocationsManagerForComputationsWithoutOrdering(
-      BufferAssignment* assignment, BufferAssigner* assigner)
-      : assignment_(assignment), assigner_(assigner) {}
+      BufferAssignment* assignment, BufferAssigner* assigner,
+      std::optional<BufferAssignment::LiveRangeInterferenceOptions>
+          live_range_interference_options = std::nullopt)
+      : assignment_(assignment),
+        assigner_(assigner),
+        live_range_interference_options_(
+            std::move(live_range_interference_options)) {}
 
   void RegisterAliasedEntryParameterAllocation(
       BufferAllocation::Index index) override {
@@ -238,30 +243,85 @@ class FastMergeBufferAllocationsManagerForComputationsWithoutOrdering
       ABSL_ASSIGN_OR_RETURN(bool success, assigner_->MaybeAssignBuffer(
                                          allocation, *hlo_buffer, assignment_));
       if (success) {
-        active_allocations_.emplace(
-            ActiveAllocation({live_range.max_end, index}));
+        RecordAllocationUse(allocation, live_range.max_end);
         VLOG(3) << "Reusing allocation #" << allocation->index()
                 << " via "
                    "FastMergeBufferAllocationsManagerForComputationsWithout"
                    "Ordering "
                    "for: "
                 << *hlo_buffer;
-        pool.erase(it);
         return true;
       }
       ++it;
     }
+
+    if (live_range_interference_options_.has_value()) {
+      bool has_custom_interference =
+          absl::c_any_of(hlo_buffer->values(), [&](const HloValue* hlo_value) {
+            return live_range_interference_options_->has_custom_interference(
+                assignment_->alias_analysis(), *hlo_value);
+          });
+      if (has_custom_interference) {
+        // Collect active allocations that are of the same color, have
+        // sufficient size, and are not in the free pool (which was already
+        // searched above). When custom interference rules are present, active
+        // allocations may be safely reusable if custom interference returns
+        // false.
+        std::vector<BufferAllocation*> candidates;
+        for (BufferAllocation::Index index : allocation_indices_) {
+          BufferAllocation* allocation =
+              assignment_->GetMutableAllocation(index);
+          if (allocation->color() == buffer_color &&
+              allocation->size() >= required_size &&
+              !pool.contains({allocation->size(), index})) {
+            candidates.push_back(allocation);
+          }
+        }
+        // Sort candidates by size (best-fit) to preserve larger allocations for
+        // future large buffers.
+        absl::c_sort(candidates,
+                     [](const BufferAllocation* a, const BufferAllocation* b) {
+                       if (a->size() != b->size()) {
+                         return a->size() < b->size();
+                       }
+                       return a->index() < b->index();
+                     });
+        for (BufferAllocation* allocation : candidates) {
+          ABSL_ASSIGN_OR_RETURN(bool success,
+                           assigner_->MaybeAssignBuffer(allocation, *hlo_buffer,
+                                                        assignment_));
+          if (success) {
+            RecordAllocationUse(allocation, live_range.max_end);
+            return true;
+          }
+        }
+      }
+    }
+
     return false;
   }
 
   // Registers a new standalone allocation to be tracked.
   void RegisterNewAllocation(const HloBuffer* hlo_buffer,
                              BufferAllocation::Index index) override {
+    allocation_indices_.push_back(index);
     BufferLiveRange live_range = GetBufferLiveRange(hlo_buffer);
+    allocation_max_end_time_[index] = live_range.max_end;
     active_allocations_.emplace(ActiveAllocation({live_range.max_end, index}));
   }
 
  private:
+  void RecordAllocationUse(BufferAllocation* allocation, int64_t max_end_time) {
+    BufferAllocation::Index index = allocation->index();
+    size_t erased =
+        free_pool_[allocation->color()].erase({allocation->size(), index});
+    int64_t& current_max_end = allocation_max_end_time_[index];
+    if (erased > 0 || max_end_time > current_max_end) {
+      current_max_end = std::max(current_max_end, max_end_time);
+      active_allocations_.emplace(ActiveAllocation({current_max_end, index}));
+    }
+  }
+
   BufferLiveRange GetBufferLiveRange(const HloBuffer* hlo_buffer) const {
     return GetHloBufferLiveRange(
         hlo_buffer, assignment_->hlo_live_range().buffer_live_ranges());
@@ -276,6 +336,11 @@ class FastMergeBufferAllocationsManagerForComputationsWithoutOrdering
            active_allocations_.top().max_end_time <= current_start_time) {
       auto top = active_allocations_.top();
       active_allocations_.pop();
+      auto it = allocation_max_end_time_.find(top.index);
+      if (it != allocation_max_end_time_.end() &&
+          top.max_end_time < it->second) {
+        continue;
+      }
       BufferAllocation* alloc = assignment_->GetMutableAllocation(top.index);
       free_pool_[alloc->color()].emplace(alloc->size(), top.index);
     }
@@ -294,10 +359,14 @@ class FastMergeBufferAllocationsManagerForComputationsWithoutOrdering
 
   BufferAssignment* const assignment_;
   BufferAssigner* const assigner_;
+  std::optional<BufferAssignment::LiveRangeInterferenceOptions>
+      live_range_interference_options_;
   int64_t prev_start_time_ = 0;
   std::priority_queue<ActiveAllocation, std::vector<ActiveAllocation>,
                       std::greater<ActiveAllocation>>
       active_allocations_;
+  absl::flat_hash_map<BufferAllocation::Index, int64_t>
+      allocation_max_end_time_;
   absl::flat_hash_map<
       LogicalBuffer::Color,
       absl::btree_set<std::pair<int64_t, BufferAllocation::Index>>>
@@ -310,7 +379,7 @@ BufferAssigner::CreateFastMergeManagerForTest(BufferAssignment* assignment,
                                               BufferAssigner* assigner) {
   return std::make_unique<
       FastMergeBufferAllocationsManagerForComputationsWithoutOrdering>(
-      assignment, assigner);
+      assignment, assigner, assigner->opts_.live_range_interference_options);
 }
 
 namespace {
@@ -1857,6 +1926,20 @@ bool BufferAssigner::LiveRangeInterferes(
     const HloValue* buffer1, const HloLiveRange::LiveRangeBounds& live_range1,
     const HloValue* buffer2, const HloLiveRange::LiveRangeBounds& live_range2,
     BufferAssignment* assignment) {
+  if (opts_.live_range_interference_options.has_value()) {
+    const auto& options = *opts_.live_range_interference_options;
+    if (options.has_custom_interference(assignment->alias_analysis(),
+                                        *buffer1) &&
+        options.has_custom_interference(assignment->alias_analysis(),
+                                        *buffer2)) {
+      std::optional<bool> custom_interference =
+          options.interferes(assignment->alias_analysis(), *buffer1, *buffer2);
+      if (custom_interference.has_value()) {
+        return *custom_interference;
+      }
+    }
+  }
+
   CHECK((assignment->hlo_live_range().total_order_scheduled()));
 
   // An HloValue can hold multiple instruction positions during its lifetime
@@ -2045,14 +2128,34 @@ absl::StatusOr<bool> BufferAssigner::MaybeAssignBuffer(
                   << new_value->ToShortString();
           return false;
         }
-      } else if (assignment->hlo_ordering().MayInterfere(
-                     assigned_buffer, *new_value,
-                     assignment->dataflow_analysis(), alias_info_)) {
-        // Fallback to partial order based interference detection (slower) when
-        // we don't have a total order scheduled module.
-        VLOG(4) << "Can't assign: assignee " << assigned_buffer
-                << " may interfere with " << new_value->ToShortString();
-        return false;
+      } else {
+        std::optional<bool> custom_interference;
+        if (opts_.live_range_interference_options.has_value()) {
+          const auto& options = *opts_.live_range_interference_options;
+          if (options.has_custom_interference(assignment->alias_analysis(),
+                                              *new_value) &&
+              options.has_custom_interference(assignment->alias_analysis(),
+                                              assigned_buffer)) {
+            custom_interference = options.interferes(
+                assignment->alias_analysis(), *new_value, assigned_buffer);
+          }
+        }
+        if (custom_interference.has_value()) {
+          if (*custom_interference) {
+            VLOG(4) << "Can't assign: assignee " << assigned_buffer
+                    << " custom interference with "
+                    << new_value->ToShortString();
+            return false;
+          }
+        } else if (assignment->hlo_ordering().MayInterfere(
+                       assigned_buffer, *new_value,
+                       assignment->dataflow_analysis(), alias_info_)) {
+          // Fallback to partial order based interference detection (slower)
+          // when we don't have a total order scheduled module.
+          VLOG(4) << "Can't assign: assignee " << assigned_buffer
+                  << " may interfere with " << new_value->ToShortString();
+          return false;
+        }
       }
 
       // Copy instruction don't share a buffer with their input operand.
@@ -2185,6 +2288,19 @@ bool BufferAssigner::DelayTemporaryBufferAssignment(
     return false;
   }
 
+  // Buffers with custom interference overrides must bypass heap simulation,
+  // which only checks interval overlap. Being pessimistic here by only
+  // checking if the buffer is subject to overrides is fine - bypassing delay
+  // simply falls back to standard sequential assignment, where exact pairwise
+  // interference is checked before any reuse and downstream buffers can still
+  // reuse this allocation.
+  if (opts_.live_range_interference_options.has_value() &&
+      absl::c_any_of(hlo_buffer->values(), [&](const HloValue* hlo_value) {
+        return opts_.live_range_interference_options->has_custom_interference(
+            assignment->alias_analysis(), *hlo_value);
+      })) {
+    return false;
+  }
   bool all_computations_have_sequential_order = true;
   for (const HloValue* hlo_value : hlo_buffer->values()) {
     HloComputation* computation = hlo_value->instruction()->parent();
@@ -2424,7 +2540,7 @@ absl::Status BufferAssigner::AssignBuffersForGlobalComputations(
                });
 
   FastMergeBufferAllocationsManagerForComputationsWithoutOrdering fast_manager(
-      assignment, this);
+      assignment, this, opts_.live_range_interference_options);
   DefaultBufferAllocationsManagerForComputationsWithoutOrdering default_manager(
       assignment, this);
 
@@ -2607,6 +2723,16 @@ absl::StatusOr<int64_t> BufferAssigner::ReuseCompatibleTempHeaps(
     // `values`), has the color we are currently processing, and has a known
     // live range.
     auto is_moveable = [&](const HloValue* buffer_value) {
+      // Exclude values with custom interference rules. By the
+      // LiveRangeInterferenceOptions invariant, `interferes(a, b)` can only
+      // return non-nullopt when both `a` and `b` have custom interference, so
+      // excluding them here also guarantees standard live-range checks suffice
+      // in `candidate_interferes_with`.
+      if (opts_.live_range_interference_options.has_value() &&
+          opts_.live_range_interference_options->has_custom_interference(
+              assignment->alias_analysis(), *buffer_value)) {
+        return false;
+      }
       return values->contains(buffer_value) &&
              buffer_value->color() == buffer_color &&
              live_ranges.contains(buffer_value);
diff --git a/third_party/xla/xla/service/buffer_assignment.h b/third_party/xla/xla/service/buffer_assignment.h
index 31e1cdec889..ee33ba8564e 100644
--- a/third_party/xla/xla/service/buffer_assignment.h
+++ b/third_party/xla/xla/service/buffer_assignment.h
@@ -486,6 +486,30 @@ class BufferAssignment {
     buffer_assignment::BufferIsolationConfig config;
   };
 
+  struct LiveRangeInterferenceOptions {
+    // Unary predicate: returns true if this HloValue has custom interference
+    // rules and should bypass HeapSimulator interval packing.
+    // Bypassing HeapSimulator ensures that buffers subject to custom rules
+    // are assigned sequentially where pairwise interference overrides can be
+    // evaluated.
+    std::function<bool(const HloAliasAnalysis&, const HloValue&)>
+        has_custom_interference;
+
+    // Symmetric binary predicate: returns true/false to override interference
+    // between lhs and rhs, or std::nullopt to fall back to default analysis.
+    // Must be symmetric: interferes(alias_analysis, a, b) ==
+    // interferes(alias_analysis, b, a).
+    // Returning true forces interference (preventing allocation reuse), while
+    // returning false forces non-interference (enabling allocation reuse).
+    //
+    // Invariant: If interferes(alias_analysis, a, b) returns a non-nullopt
+    // value, has_custom_interference(alias_analysis, a) and
+    // has_custom_interference(alias_analysis, b) must both return true.
+    std::function<std::optional<bool>(const HloAliasAnalysis&, const HloValue&,
+                                      const HloValue&)>
+        interferes;
+  };
+
   // Returns the vector containing all buffer allocations in this assignment.
   const std::vector<BufferAllocation>& Allocations() const {
     return allocations_;
@@ -871,6 +895,8 @@ class BufferAssigner {
 
   using MustNotLiveOut = std::function<bool(
       const HloAliasAnalysis&, const HloInstruction*, const ShapeIndex&)>;
+  using LiveRangeInterferenceOptions =
+      BufferAssignment::LiveRangeInterferenceOptions;
 
   // The order in which to process buffers during buffer assignment.
   enum class BufferOrder {
@@ -895,6 +921,9 @@ class BufferAssigner {
     // live out of a computation.
     std::optional<MustNotLiveOut> must_not_live_out;
 
+    // Optional callbacks to override live-range interference analysis between
+    // two buffer values (e.g. for target-specific concurrency or barriers).
+    std::optional<LiveRangeInterferenceOptions> live_range_interference_options;
     // Description of any buffer offsets that are already set by an earlier
     // pass.
     std::unique_ptr<memory_space_assignment::PresetAssignments>
diff --git a/third_party/xla/xla/service/buffer_assignment_test.cc b/third_party/xla/xla/service/buffer_assignment_test.cc
index b1cacb823b5..ca87aadb1b5 100644
--- a/third_party/xla/xla/service/buffer_assignment_test.cc
+++ b/third_party/xla/xla/service/buffer_assignment_test.cc
@@ -32,6 +32,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/match.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "xla/comparison_util.h"
@@ -5896,5 +5897,211 @@ TEST_F(BufferAssignmentTest,
   EXPECT_EQ(roundtrip->Allocations().size(), buffers->Allocations().size());
 }
 
+TEST_F(BufferAssignmentTest, HasLiveRangeInterferenceOption) {
+  constexpr absl::string_view kHlo = R"hlo(
+HloModule main, is_scheduled=true
+
+ENTRY main {
+  call-start.0 = ((), (), (s32[4096]{0})) call-start(), to_apply={
+    ROOT tuple.0 = tuple()
+  }
+  call-start.1 = ((), (), (s32[4096]{0})) call-start(), to_apply={
+    ROOT tuple.0 = tuple()
+  }
+  call-done.0 = () call-done(call-start.0)
+  call-done.1 = () call-done(call-start.1)
+  buffer.0 = s32[4096]{0} custom-call(), custom_call_target="AllocateBuffer"
+  ROOT copy.0 = s32[4096]{0} copy(buffer.0)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kHlo));
+
+  BufferAssigner::Options opts;
+  opts.allocate_buffers_for_constants = true;
+  opts.must_not_live_out = [](const HloAliasAnalysis& alias_analysis,
+                              const HloInstruction* instruction,
+                              const ShapeIndex& index) -> bool {
+    return absl::StartsWith(instruction->name(), "call-start");
+  };
+  BufferAssignment::LiveRangeInterferenceOptions live_range_opts;
+  live_range_opts.has_custom_interference =
+      [](const HloAliasAnalysis& alias_analysis,
+         const HloValue& value) -> bool {
+    return absl::StartsWith(value.instruction()->name(), "call-start") &&
+           value.index() == ShapeIndex{2, 0};
+  };
+  live_range_opts.interferes = [](const HloAliasAnalysis& alias_analysis,
+                                  const HloValue& lhs,
+                                  const HloValue& rhs) -> std::optional<bool> {
+    if (absl::StartsWith(lhs.instruction()->name(), "call-start") &&
+        absl::StartsWith(rhs.instruction()->name(), "call-start") &&
+        lhs.index() == ShapeIndex{2, 0} && rhs.index() == ShapeIndex{2, 0}) {
+      return false;
+    }
+    return std::nullopt;
+  };
+  opts.live_range_interference_options = live_range_opts;
+
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<BufferAssignment> buffers,
+      BufferAssigner::Run(
+          module.get(),
+          std::make_unique<SequentialHloOrdering>(module->schedule()),
+          &BufferSizeBytes, &alias_info_,
+          [](LogicalBuffer::Color) { return 1; }, std::move(opts)));
+
+  const HloInstruction* call_start0 =
+      module->entry_computation()->GetInstructionWithName("call-start.0");
+  const HloInstruction* call_start1 =
+      module->entry_computation()->GetInstructionWithName("call-start.1");
+  const HloInstruction* buffer0 =
+      module->entry_computation()->GetInstructionWithName("buffer.0");
+
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice slice0,
+                       buffers->GetUniqueSlice(call_start0, ShapeIndex{2, 0}));
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice slice1,
+                       buffers->GetUniqueSlice(call_start1, ShapeIndex{2, 0}));
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice buffer0_slice,
+                       buffers->GetUniqueTopLevelSlice(buffer0));
+
+  EXPECT_EQ(slice0.index(), slice1.index());
+  EXPECT_EQ(slice0.offset(), slice1.offset());
+  EXPECT_EQ(buffer0_slice.index(), slice0.index());
+  EXPECT_EQ(buffer0_slice.offset(), slice0.offset());
+}
+
+TEST_F(BufferAssignmentTest, CustomInterferenceForcesInterference) {
+  constexpr absl::string_view kHlo = R"hlo(
+HloModule main, is_scheduled=true
+
+ENTRY main {
+  buffer.0 = s32[4096]{0} custom-call(), custom_call_target="AllocateBuffer"
+  copy.0 = s32[4096]{0} copy(buffer.0)
+  buffer.1 = s32[4096]{0} custom-call(), custom_call_target="AllocateBuffer"
+  ROOT copy.1 = s32[4096]{0} copy(buffer.1)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kHlo));
+
+  BufferAssigner::Options opts;
+  opts.allocate_buffers_for_constants = true;
+  BufferAssignment::LiveRangeInterferenceOptions live_range_opts;
+  live_range_opts.has_custom_interference =
+      [](const HloAliasAnalysis& alias_analysis,
+         const HloValue& value) -> bool {
+    return absl::StartsWith(value.instruction()->name(), "buffer.");
+  };
+  live_range_opts.interferes = [](const HloAliasAnalysis& alias_analysis,
+                                  const HloValue& lhs,
+                                  const HloValue& rhs) -> std::optional<bool> {
+    if (absl::StartsWith(lhs.instruction()->name(), "buffer.") &&
+        absl::StartsWith(rhs.instruction()->name(), "buffer.")) {
+      return true;
+    }
+    return std::nullopt;
+  };
+  opts.live_range_interference_options = live_range_opts;
+
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<BufferAssignment> buffers,
+      BufferAssigner::Run(
+          module.get(),
+          std::make_unique<SequentialHloOrdering>(module->schedule()),
+          &BufferSizeBytes, &alias_info_,
+          [](LogicalBuffer::Color) { return 1; }, std::move(opts)));
+
+  const HloInstruction* buffer0 =
+      module->entry_computation()->GetInstructionWithName("buffer.0");
+  const HloInstruction* buffer1 =
+      module->entry_computation()->GetInstructionWithName("buffer.1");
+
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice buffer0_slice,
+                       buffers->GetUniqueTopLevelSlice(buffer0));
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice buffer1_slice,
+                       buffers->GetUniqueTopLevelSlice(buffer1));
+
+  EXPECT_NE(buffer0_slice.index(), buffer1_slice.index());
+}
+
+TEST_F(BufferAssignmentTest, HasLiveRangeInterferenceOptionFastMerge) {
+  constexpr absl::string_view kHlo = R"hlo(
+HloModule main, is_scheduled=true
+
+ENTRY main {
+  call-start.0 = ((), (), (s32[4096]{0})) call-start(), to_apply={
+    ROOT tuple.0 = tuple()
+  }
+  call-start.1 = ((), (), (s32[4096]{0})) call-start(), to_apply={
+    ROOT tuple.0 = tuple()
+  }
+  call-done.0 = () call-done(call-start.0)
+  call-done.1 = () call-done(call-start.1)
+  buffer.0 = s32[4096]{0} custom-call(), custom_call_target="AllocateBuffer"
+  ROOT copy.0 = s32[4096]{0} copy(buffer.0)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kHlo));
+
+  BufferAssigner::Options opts;
+  opts.allocate_buffers_for_constants = true;
+  opts.assignment_algorithm_for_computations_without_ordering =
+      buffer_assignment::
+          AssignmentAlgorithmForComputationsWithoutOrderingProto::FAST_MERGE;
+  opts.color_memory_limit = [](LogicalBuffer::Color) { return 1024 * 1024; };
+  opts.must_not_live_out = [](const HloAliasAnalysis& alias_analysis,
+                              const HloInstruction* instruction,
+                              const ShapeIndex& index) -> bool {
+    return absl::StartsWith(instruction->name(), "call-start");
+  };
+  BufferAssignment::LiveRangeInterferenceOptions live_range_opts;
+  live_range_opts.has_custom_interference =
+      [](const HloAliasAnalysis& alias_analysis,
+         const HloValue& value) -> bool {
+    return absl::StartsWith(value.instruction()->name(), "call-start") &&
+           value.index() == ShapeIndex{2, 0};
+  };
+  live_range_opts.interferes = [](const HloAliasAnalysis& alias_analysis,
+                                  const HloValue& lhs,
+                                  const HloValue& rhs) -> std::optional<bool> {
+    if (absl::StartsWith(lhs.instruction()->name(), "call-start") &&
+        absl::StartsWith(rhs.instruction()->name(), "call-start") &&
+        lhs.index() == ShapeIndex{2, 0} && rhs.index() == ShapeIndex{2, 0}) {
+      return false;
+    }
+    return std::nullopt;
+  };
+  opts.live_range_interference_options = live_range_opts;
+
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<BufferAssignment> buffers,
+      BufferAssigner::Run(
+          module.get(),
+          std::make_unique<SequentialHloOrdering>(module->schedule()),
+          &BufferSizeBytes, &alias_info_,
+          [](LogicalBuffer::Color) { return 1; }, std::move(opts)));
+
+  const HloInstruction* call_start0 =
+      module->entry_computation()->GetInstructionWithName("call-start.0");
+  const HloInstruction* call_start1 =
+      module->entry_computation()->GetInstructionWithName("call-start.1");
+  const HloInstruction* buffer0 =
+      module->entry_computation()->GetInstructionWithName("buffer.0");
+
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice slice0,
+                       buffers->GetUniqueSlice(call_start0, ShapeIndex{2, 0}));
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice slice1,
+                       buffers->GetUniqueSlice(call_start1, ShapeIndex{2, 0}));
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice buffer0_slice,
+                       buffers->GetUniqueTopLevelSlice(buffer0));
+
+  EXPECT_EQ(slice0.index(), slice1.index());
+  EXPECT_EQ(slice0.offset(), slice1.offset());
+  EXPECT_EQ(buffer0_slice.index(), slice0.index());
+  EXPECT_EQ(buffer0_slice.offset(), slice0.offset());
+}
+
 }  // namespace
 }  // namespace xla

From 76aef3bf9a408e2c78429fa96a9f762d431f2111 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Sat, 10 Oct 2026 08:35:49 +0530
Subject: [PATCH 0404/1000] Refactor shutdown_accelerator_system for error
 handling

---
 tensorflow/dtensor/python/accelerator_util.py | 45 +++++++++++--------
 1 file changed, 27 insertions(+), 18 deletions(-)

diff --git a/tensorflow/dtensor/python/accelerator_util.py b/tensorflow/dtensor/python/accelerator_util.py
index 6fa6cafdb31..ce641b9b2b1 100644
--- a/tensorflow/dtensor/python/accelerator_util.py
+++ b/tensorflow/dtensor/python/accelerator_util.py
@@ -293,27 +293,36 @@ def initialize_accelerator_system(
 def shutdown_accelerator_system() -> None:
   """Shuts down the accelerator system."""
   global _INITIALIZED_ACCELERATOR_SYSTEM_TYPE
+  # Only `async_wait()` is wrapped so that an error from it never gets
+  # silently replaced by a later `raise` below (e.g. "not initialized"),
+  # which would otherwise mask the real failure.
   try:
     context.async_wait()
-  finally:
-    if not is_initialized():
-      raise ValueError(
-          "Accelerator system is not initialized. Call "
-          "tf.experimental.dtensor.initialize_accelerator_system first."
-      )
+  except Exception:
+    logging.exception(
+        "Error waiting for async operations to finish before shutting "
+        "down the accelerator system."
+    )
+    raise
 
-    device_type = _INITIALIZED_ACCELERATOR_SYSTEM_TYPE
+  if not is_initialized():
+    raise ValueError(
+        "Accelerator system is not initialized. Call "
+        "tf.experimental.dtensor.initialize_accelerator_system first."
+    )
 
-    if not config.is_local_mode():
-      raise ValueError(
-          "Shutting down accelerator system under multi-client mode is "
-          "not supported."
-      )
+  device_type = _INITIALIZED_ACCELERATOR_SYSTEM_TYPE
 
-    if device_type == "TPU":
-      tpu_util.shutdown_tpu_system()
+  if not config.is_local_mode():
+    raise ValueError(
+        "Shutting down accelerator system under multi-client mode is "
+        "not supported."
+    )
 
-    # reset TF context to stop gRPC servers.
-    context._reset_context()  # pylint: disable=protected-access
-    context.context()._clear_caches()  # pylint: disable=protected-access
-    _INITIALIZED_ACCELERATOR_SYSTEM_TYPE = None
+  if device_type == "TPU":
+    tpu_util.shutdown_tpu_system()
+
+  # reset TF context to stop gRPC servers.
+  context._reset_context()  # pylint: disable=protected-access
+  context.context()._clear_caches()  # pylint: disable=protected-access
+  _INITIALIZED_ACCELERATOR_SYSTEM_TYPE = None

From 0491af768d1d8bbacff8c717104e00a4d5d17efd Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Sat, 10 Oct 2026 08:39:16 +0530
Subject: [PATCH 0405/1000] Refactor jobs function to check d_jobs more clearly

---
 tensorflow/dtensor/python/config.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/dtensor/python/config.py b/tensorflow/dtensor/python/config.py
index 0653170debd..fb90258e3ce 100644
--- a/tensorflow/dtensor/python/config.py
+++ b/tensorflow/dtensor/python/config.py
@@ -156,7 +156,7 @@ def _bns_task_id(job: str) -> Union[int, str]:
 def jobs() -> List[str]:
   """Returns a list of job names of all clients in this DTensor cluster."""
   d_jobs = os.environ.get(_DT_JOBS)
-  if d_jobs is None:
+  if not d_jobs:
     return []
   d_jobs_list = d_jobs.split(",")
 

From f9b15d6243a81bcf93584254493589cbd8fc053b Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Sat, 10 Oct 2026 08:42:06 +0530
Subject: [PATCH 0406/1000] Refactor shutdown_accelerator_system function

---
 tensorflow/dtensor/python/accelerator_util.py | 36 ++++++++++---------
 1 file changed, 19 insertions(+), 17 deletions(-)

diff --git a/tensorflow/dtensor/python/accelerator_util.py b/tensorflow/dtensor/python/accelerator_util.py
index ce641b9b2b1..a279039d232 100644
--- a/tensorflow/dtensor/python/accelerator_util.py
+++ b/tensorflow/dtensor/python/accelerator_util.py
@@ -293,18 +293,9 @@ def initialize_accelerator_system(
 def shutdown_accelerator_system() -> None:
   """Shuts down the accelerator system."""
   global _INITIALIZED_ACCELERATOR_SYSTEM_TYPE
-  # Only `async_wait()` is wrapped so that an error from it never gets
-  # silently replaced by a later `raise` below (e.g. "not initialized"),
-  # which would otherwise mask the real failure.
-  try:
-    context.async_wait()
-  except Exception:
-    logging.exception(
-        "Error waiting for async operations to finish before shutting "
-        "down the accelerator system."
-    )
-    raise
 
+  # Validate first, so a validation error can never mask a real failure
+  # from `async_wait()` below.
   if not is_initialized():
     raise ValueError(
         "Accelerator system is not initialized. Call "
@@ -319,10 +310,21 @@ def shutdown_accelerator_system() -> None:
         "not supported."
     )
 
-  if device_type == "TPU":
-    tpu_util.shutdown_tpu_system()
+  try:
+    context.async_wait()
+  except Exception:
+    logging.exception(
+        "Error waiting for async operations to finish before shutting "
+        "down the accelerator system."
+    )
+    raise
+  finally:
+    # Cleanup always runs, even if `async_wait()` raised, so the system
+    # is never left stuck in an "initialized" state.
+    if device_type == "TPU":
+      tpu_util.shutdown_tpu_system()
 
-  # reset TF context to stop gRPC servers.
-  context._reset_context()  # pylint: disable=protected-access
-  context.context()._clear_caches()  # pylint: disable=protected-access
-  _INITIALIZED_ACCELERATOR_SYSTEM_TYPE = None
+    # reset TF context to stop gRPC servers.
+    context._reset_context()  # pylint: disable=protected-access
+    context.context()._clear_caches()  # pylint: disable=protected-access
+    _INITIALIZED_ACCELERATOR_SYSTEM_TYPE = None

From 614df11177ee80332bc53042f1bda917fbf6c1b6 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Sat, 10 Oct 2026 08:42:12 +0530
Subject: [PATCH 0407/1000] Update tensorflow/dtensor/python/config.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 tensorflow/dtensor/python/config.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/dtensor/python/config.py b/tensorflow/dtensor/python/config.py
index fb90258e3ce..cd60ca6db9d 100644
--- a/tensorflow/dtensor/python/config.py
+++ b/tensorflow/dtensor/python/config.py
@@ -156,7 +156,7 @@ def _bns_task_id(job: str) -> Union[int, str]:
 def jobs() -> List[str]:
   """Returns a list of job names of all clients in this DTensor cluster."""
   d_jobs = os.environ.get(_DT_JOBS)
-  if not d_jobs:
+  if not d_jobs or not d_jobs.strip():
     return []
   d_jobs_list = d_jobs.split(",")
 

From 9cd1c084e041efc9d5a548aea502bbe04cc68553 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Sat, 10 Oct 2026 08:48:23 +0530
Subject: [PATCH 0408/1000] Update mesh_util.py

---
 tensorflow/dtensor/python/mesh_util.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/tensorflow/dtensor/python/mesh_util.py b/tensorflow/dtensor/python/mesh_util.py
index 4755a2f3148..179d0d293b0 100644
--- a/tensorflow/dtensor/python/mesh_util.py
+++ b/tensorflow/dtensor/python/mesh_util.py
@@ -176,18 +176,22 @@ def create_distributed_mesh(
     raise ValueError('Accelerators are uninitialized, please run '
                      'dtensor.initialize_accelerator_system() first.')
 
-  if device_type and device_type.upper() == 'TPU':
+  device_specs, device_type = _make_device_specs(local_devices, device_type)
+
+  if device_type.upper() == 'TPU':
     # TODO(b/185940495): Allow multi-mesh and partial on TPU.
     # TPU meshes can only be configured through environment variables that
     # reflect the actual TPU topology. Do not let users specify custom args.
+    # Checked here (after type inference) rather than against the original
+    # `device_type` argument, so a caller who passes `local_devices` full of
+    # TPU specs without also passing `device_type='TPU'` still gets this
+    # error instead of silently having `local_devices` ignored below.
     if local_devices is not None:
       raise ValueError(
           f'Do not specify devices for {device_type.upper()} meshes. '
           f'Using a partial list of devices for {device_type.upper()} '
           f'is not supported.')
 
-  device_specs, device_type = _make_device_specs(local_devices, device_type)
-
   if device_type.upper() in ['CPU', 'GPU']:
     # For CPU and GPU meshes, user-specified args take precedence over env vars.
     # This is particularly useful on single clients when users want to create

From 82da87bd731404a288d126b1aa7c7dfa0776ae70 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Sat, 10 Oct 2026 08:52:00 +0530
Subject: [PATCH 0409/1000] Update tensorflow/dtensor/python/mesh_util.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 tensorflow/dtensor/python/mesh_util.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/dtensor/python/mesh_util.py b/tensorflow/dtensor/python/mesh_util.py
index 179d0d293b0..a209d48fb27 100644
--- a/tensorflow/dtensor/python/mesh_util.py
+++ b/tensorflow/dtensor/python/mesh_util.py
@@ -178,7 +178,7 @@ def create_distributed_mesh(
 
   device_specs, device_type = _make_device_specs(local_devices, device_type)
 
-  if device_type.upper() == 'TPU':
+  if device_type and device_type.upper() == 'TPU':
     # TODO(b/185940495): Allow multi-mesh and partial on TPU.
     # TPU meshes can only be configured through environment variables that
     # reflect the actual TPU topology. Do not let users specify custom args.

From 4b9df52bd8be0413da54b668833457a82867265d Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 9 Oct 2026 20:41:23 -0700
Subject: [PATCH 0410/1000] Automated Code Change

PiperOrigin-RevId: 996918925
---
 .../core/kernels/linalg/tridiagonal_solve_op_gpu.cu.cc | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/tensorflow/core/kernels/linalg/tridiagonal_solve_op_gpu.cu.cc b/tensorflow/core/kernels/linalg/tridiagonal_solve_op_gpu.cu.cc
index 738df197fb4..b2df6eb132c 100644
--- a/tensorflow/core/kernels/linalg/tridiagonal_solve_op_gpu.cu.cc
+++ b/tensorflow/core/kernels/linalg/tridiagonal_solve_op_gpu.cu.cc
@@ -375,18 +375,16 @@ class TridiagonalSolveOpGpu : public OpKernel {
     for (int index : perm) {
       dims.push_back(lhs.dim_size(index));
     }
-    TensorShape lhs_transposed_shape(
-        absl::Span<const int64_t>(dims.data(), ndims));
+    TensorShape lhs_transposed_shape(absl::MakeConstSpan(dims).first(ndims));
 
     std::unique_ptr<GpuSolver> cublas_solver(new GpuSolver(context));
     OP_REQUIRES_OK(context, cublas_solver->allocate_scoped_tensor(
                                 DataTypeToEnum<Scalar>::value,
                                 lhs_transposed_shape, &lhs_transposed));
     auto device = context->eigen_device<Eigen::GpuDevice>();
-    OP_REQUIRES_OK(
-        context,
-        DoTranspose(device, lhs, absl::Span<const int>(perm.data(), ndims),
-                    &lhs_transposed));
+    OP_REQUIRES_OK(context, DoTranspose(device, lhs,
+                                        absl::MakeConstSpan(perm).first(ndims),
+                                        &lhs_transposed));
   }
 
   TridiagonalSolveOpGpuLinalg<Scalar> linalgOp_;

From 6ce4953c0fb38b60504909c9aacc875270c1e32c Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Sat, 10 Oct 2026 09:32:00 +0530
Subject: [PATCH 0411/1000] Update context.py

---
 tensorflow/python/eager/context.py | 3359 ++++++++++++++++++++++++++--
 1 file changed, 3185 insertions(+), 174 deletions(-)

diff --git a/tensorflow/python/eager/context.py b/tensorflow/python/eager/context.py
index 73be78de959..2435b922fac 100644
--- a/tensorflow/python/eager/context.py
+++ b/tensorflow/python/eager/context.py
@@ -1,4 +1,4 @@
-# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2017 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.
@@ -12,222 +12,3233 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""State management for eager execution."""
 
-import weakref
+import collections
+import contextlib
+import copy
+import gc
+import itertools
+import os
+import random
+import threading
 
-from absl.testing import parameterized
+from absl import logging
 import numpy as np
 
-from xla.service import hlo_pb2
-from tensorflow.python.eager import context
-from tensorflow.python.eager import def_function
-from tensorflow.python.framework import constant_op
-from tensorflow.python.framework import errors
-from tensorflow.python.framework import ops
-from tensorflow.python.framework import test_util
-from tensorflow.python.ops import array_ops
-from tensorflow.python.platform import test
+from xla.tsl.protobuf import coordination_config_pb2
+from tensorflow.core.framework import function_pb2
+from tensorflow.core.framework import graph_debug_info_pb2
+from tensorflow.core.protobuf import config_pb2
+from tensorflow.core.protobuf import rewriter_config_pb2
+from tensorflow.python import pywrap_tfe
+from tensorflow.python import tf2
+from tensorflow.python.client import pywrap_tf_session
+from tensorflow.python.eager import cancellation
+from tensorflow.python.eager import execute
+from tensorflow.python.eager import executor
+from tensorflow.python.eager import monitoring
+from tensorflow.python.framework import c_api_util
+from tensorflow.python.framework import device as pydev
+from tensorflow.python.framework import tfrt_utils
+from tensorflow.python.util import compat
+from tensorflow.python.util import function_utils
+from tensorflow.python.util import is_in_graph_mode
+from tensorflow.python.util import tf_contextlib
+from tensorflow.python.util.deprecation import deprecated
+from tensorflow.python.util.tf_export import tf_export
 
 
-class ContextTest(test.TestCase, parameterized.TestCase):
+# TODO(b/307794935): Remove after a solution is found.
+is_oss = True  # updated by copybara
 
-  def testSetGlobalSeed(self):
-    c = context.Context()
-    c._set_global_seed(123)
-    for t in [np.int32, np.int64, np.uint32, np.uint64]:
-      c._set_global_seed(t(123))
-      c._set_global_seed(np.array(123, dtype=t))
-      c._set_global_seed(ops.convert_to_tensor(123, dtype=t))
+GRAPH_MODE = 0
+EAGER_MODE = 1
 
-  def testContextIsDestroyedAfterTensors(self):
-    # Create a new context
-    new_context = context.Context()
-    weak_c = weakref.ref(new_context)
-    new_context.ensure_initialized()
+default_execution_mode = EAGER_MODE if tf2.enabled() else GRAPH_MODE
 
-    # Create a tensor with the new context as default.
-    # Make sure to restore the original context.
-    original_context = context.context()
+# Cache from (old_device_name, partial_new_device_name) -> (new_device_name,
+# new_device_spec).
+# Note that we do not protect this with a lock and instead rely on python's GIL
+# and the idempotent nature of writes to provide thread safety.
+_device_parsing_cache = {}
+_starting_device_spec = pydev.DeviceSpec.from_string("")
+
+_MAXINT32 = 2**31 - 1
+
+DEVICE_PLACEMENT_EXPLICIT = pywrap_tfe.TFE_DEVICE_PLACEMENT_EXPLICIT
+DEVICE_PLACEMENT_WARN = pywrap_tfe.TFE_DEVICE_PLACEMENT_WARN
+DEVICE_PLACEMENT_SILENT = pywrap_tfe.TFE_DEVICE_PLACEMENT_SILENT
+DEVICE_PLACEMENT_SILENT_FOR_INT32 = (
+    pywrap_tfe.TFE_DEVICE_PLACEMENT_SILENT_FOR_INT32
+)
+
+SYNC = 0
+ASYNC = 1
+
+_KEEP_ALIVE_SECS = 600
+
+_python_eager_context_create_counter = monitoring.Counter(
+    "/tensorflow/api/python/eager_context_create_counter",
+    "Counter for number of eager contexts created in Python.",
+)
+
+# Re-exporting through context.
+is_tfrt_enabled = tfrt_utils.enabled
+
+# This flag and the associated environment var are transient and will eventually
+# be removed, once this experiment is enabled by default.
+_JIT_COMPILE_REWRITE_ENABLED = os.getenv("TF_JIT_COMPILE_REWRITE") == "1"
+
+_XLA_SHARDING_FOR_RESOURCE_VARIABLES = (
+    os.getenv("TF_XLA_SHARDING_FOR_RESOURCE_VARIABLES") == "1"
+)
+_OPTIONALS_IN_GRADIENTS = os.getenv("TF_OPTIONALS_IN_GRADIENTS") != "0"
+
+
+def run_eager_op_as_function_enabled():
+  return True
+
+
+# This method should only be called after the context has beein initialized.
+def enable_jit_compile_rewrite():
+  """Run jit_compile functions through rewrite pass.
+
+  This runs jit_compile functions through all of the multidevice function
+  rewrite passes.
+  """
+  global _JIT_COMPILE_REWRITE_ENABLED
+  _JIT_COMPILE_REWRITE_ENABLED = True
+  if context_safe() is not None:
+    context_safe().jit_compile_rewrite = True
+
+
+# This method should only be called after the context has been initialized.
+def disable_jit_compile_rewrite():
+  global _JIT_COMPILE_REWRITE_ENABLED
+  _JIT_COMPILE_REWRITE_ENABLED = False
+  if context_safe() is not None:
+    context_safe().jit_compile_rewrite = False
+
+
+def jit_compile_rewrite_enabled():
+  if context_safe() is not None:
+    return context_safe().jit_compile_rewrite
+  return _JIT_COMPILE_REWRITE_ENABLED
+
+
+def enable_xla_sharding_for_resource_variables():
+  """Enables support for annotating TF2 ResourceVariables with XLA sharding.
+
+  This allows placing XLA sharding annotations on the TF2 ResourceVariable
+  python object and inserts an XlaShardingOp with the annotation whenever a
+  ReadVariableOp is created.
+  """
+  global _XLA_SHARDING_FOR_RESOURCE_VARIABLES
+  _XLA_SHARDING_FOR_RESOURCE_VARIABLES = True
+  if context_safe() is not None:
+    context_safe().xla_sharding_for_resource_variables = True
+
+
+def disable_xla_sharding_for_resource_variables():
+  global _XLA_SHARDING_FOR_RESOURCE_VARIABLES
+  _XLA_SHARDING_FOR_RESOURCE_VARIABLES = False
+  if context_safe() is not None:
+    context_safe().xla_sharding_for_resource_variables = False
+
+
+def xla_sharding_for_resource_variables_enabled():
+  if context_safe() is not None:
+    return context_safe().xla_sharding_for_resource_variables
+  return _XLA_SHARDING_FOR_RESOURCE_VARIABLES
+
+
+def enable_optionals_in_gradients():
+  global _OPTIONALS_IN_GRADIENTS
+  _OPTIONALS_IN_GRADIENTS = True
+  if context_safe() is not None:
+    context_safe().optionals_in_gradients = True
+
+
+def disable_optionals_in_gradients():
+  global _OPTIONALS_IN_GRADIENTS
+  _OPTIONALS_IN_GRADIENTS = False
+  if context_safe() is not None:
+    context_safe().optionals_in_gradients = False
+
+
+def optionals_in_gradients_enabled():
+  if context_safe() is not None:
+    return context_safe().optionals_in_gradients
+  return _OPTIONALS_IN_GRADIENTS
+
+
+@contextlib.contextmanager
+def temporarily_disable_xla_sharding_for_resource_variables():
+  """Temporarily disables XLA sharding for resource variables.
+
+  Should be a no-op if it is already disabled.
+
+  Yields:
+    None.
+  """
+  previously_enabled = xla_sharding_for_resource_variables_enabled()
+
+  try:
+    disable_xla_sharding_for_resource_variables()
+    yield
+  finally:
+    if previously_enabled:
+      enable_xla_sharding_for_resource_variables()
+
+
+@contextlib.contextmanager
+def temporarily_disable_optionals_in_gradients():
+  """Temporarily disables generation of optionals in gradients.
+
+  Should be a no-op if it is already disabled.
+
+  Yields:
+    None.
+  """
+  previously_enabled = optionals_in_gradients_enabled()
+
+  try:
+    disable_optionals_in_gradients()
+    yield
+  finally:
+    if previously_enabled:
+      enable_optionals_in_gradients()
+
+
+# Expose it as internally public APIs for Keras use cases in b/171080602.
+tf_export("__internal__.is_tfrt_enabled", v1=[])(is_tfrt_enabled)
+
+
+class _EagerTensorCache(object):
+  """Simple cache which evicts items based on length in a FIFO manner."""
+
+  __slots__ = ["_data", "_max_items", "_max_tensor_size"]
+
+  def __init__(self, max_items=256, max_tensor_size=10000):
+    self._data = collections.OrderedDict()
+    self._max_items = max_items
+    self._max_tensor_size = max_tensor_size
+
+  def put(self, key, value):
+    if value._num_elements() > self._max_tensor_size:  # pylint: disable=protected-access
+      return
+
+    self._data[key] = value
+
+    if len(self._data) > self._max_items:
+      self._data.popitem(last=False)
+
+  def get(self, key):
+    return self._data.get(key, None)
+
+  def flush(self):
+    self._data.clear()
+
+
+class FunctionCallOptions:
+  """Options applied at call sites of eager functions.
+
+  Eager functions are functions decorated with tf.contrib.eager.defun.
+  """
+
+  __slots__ = ["_config_proto_serialized", "_executor_type"]
+
+  def __init__(self, executor_type=None, config_proto=None):
+    """Constructor.
+
+    Args:
+      executor_type: (optional) name of the executor to be used to execute the
+        eager function. If None or an empty string, the default Tensorflow
+        executor will be used.
+      config_proto: (optional) a `config_pb2.ConfigProto` proto or a serialized
+        string of that proto. The config used by Grappler when optimizing the
+        function graph. Each concrete function is optimized the first time is
+        called. Changing config_proto after the first call has no effect. If
+        config_proto is None, an empty RewriterConfig will be used.
+    """
+    self.config_proto_serialized = config_proto
+    self.executor_type = executor_type
+
+  @property
+  def executor_type(self):
+    return self._executor_type
+
+  @executor_type.setter
+  def executor_type(self, executor_type):
+    self._executor_type = executor_type
+
+  @property
+  def config_proto_serialized(self):
+    return self._config_proto_serialized
+
+  @config_proto_serialized.setter
+  def config_proto_serialized(self, config):
+    if isinstance(config, config_pb2.ConfigProto):
+      self._config_proto_serialized = config.SerializeToString(
+          deterministic=True
+      )
+    elif isinstance(config, str):
+      self._config_proto_serialized = config
+    elif config is None:
+      self._config_proto_serialized = (
+          config_pb2.ConfigProto().SerializeToString()
+      )
+    else:
+      raise ValueError(
+          "the rewriter config must be either a "
+          "config_pb2.ConfigProto, or a serialized string of that "
+          "proto or None. got: {}".format(type(config))
+      )
+
+  def as_attrs(self):
+    if self.config_proto_serialized is None:
+      config = function_utils.get_disabled_rewriter_config()
+    else:
+      config = self.config_proto_serialized
+    executor_type = self.executor_type or ""
+
+    return {"executor_type": executor_type, "config_proto": config}
+
+
+# Map from context_id (an int) to _TensorCaches.
+# Dicts are thread safe in CPython.
+# TODO(iga): Remove this once TensorCaches are moved to C++.
+_tensor_caches_map = {}
+
+
+class _TensorCaches(threading.local):
+  """Thread local tensor caches."""
+
+  __slots__ = ["_ones_rank_cache", "_zeros_cache"]
+
+  def __init__(self):
+    super().__init__()
+    self._ones_rank_cache = None
+    self._zeros_cache = None
+
+  @property
+  def ones_rank_cache(self):
+    if not self._ones_rank_cache:
+      self._ones_rank_cache = _EagerTensorCache()
+    return self._ones_rank_cache
+
+  @property
+  def zeros_cache(self):
+    if not self._zeros_cache:
+      self._zeros_cache = _EagerTensorCache()
+    return self._zeros_cache
+
+
+ContextSwitch = collections.namedtuple(
+    "ContextSwitch",
+    ["is_building_function", "enter_context_fn", "device_stack"],
+)
+
+
+# `_ContextSwitchStack` is a `threading.local` to match the semantics of
+# ``DefaultGraphStack`, which is also a `threading.local`.
+class _ContextSwitchStack(threading.local):
+  """A thread-local stack of context switches."""
+
+  def __init__(self, eager):
+    super().__init__()
+    self.stack = []
+    if eager:
+      # Initialize the stack with a pointer to enter the eager context; this
+      # ensures that the fact that eager execution was enabled is propagated
+      # across threads, since (1) `enable_eager_execution` modifies a
+      # process-level flag (`default_execution_mode`) and (2) `__init__` is
+      # called each time a threading.local object is used in a separate thread.
+      self.push(
+          is_building_function=False,
+          enter_context_fn=eager_mode,
+          device_stack=None,
+      )
+
+  def push(self, is_building_function, enter_context_fn, device_stack):
+    """Push metadata about a context switch onto the stack.
+
+    A context switch can take any one of the two forms: installing a graph as
+    the default graph, or entering the eager context. For each context switch,
+    we record whether or not the entered context is building a function.
+
+    Args:
+      is_building_function: (bool.) Whether the context is building a function.
+      enter_context_fn: (function.) A callable that executes the context switch.
+        For example, `graph.as_default` or `eager_mode`.
+      device_stack: If applicable, the device function stack for this graph.
+        When breaking out of graphs in init_scope, the innermost nonempty device
+        stack is used. Eager contexts put `None` here and the value is never
+        used.
+    """
+
+    self.stack.append(
+        ContextSwitch(is_building_function, enter_context_fn, device_stack)
+    )
+
+  def pop(self):
+    """Pop the stack."""
+
+    self.stack.pop()
+
+
+@tf_export("config.LogicalDevice")
+class LogicalDevice(
+    collections.namedtuple("LogicalDevice", ["name", "device_type"])
+):
+  """Abstraction for a logical device initialized by the runtime.
+
+  A `tf.config.LogicalDevice` corresponds to an initialized logical device on a
+  `tf.config.PhysicalDevice` or a remote device visible to the cluster. Tensors
+  and operations can be placed on a specific logical device by calling
+  `tf.device` with a specified `tf.config.LogicalDevice`.
+
+  Fields:
+    name: The fully qualified name of the device. Can be used for Op or function
+      placement.
+    device_type: String declaring the type of device such as "CPU" or "GPU".
+  """
+
+
+@tf_export(
+    "config.LogicalDeviceConfiguration",
+    "config.experimental.VirtualDeviceConfiguration",
+)
+class LogicalDeviceConfiguration(
+    collections.namedtuple(
+        "LogicalDeviceConfiguration",
+        [
+            "memory_limit",
+            "experimental_priority",
+            "experimental_device_ordinal",
+        ],
+    )
+):
+  """Configuration class for a logical devices.
+
+  The class specifies the parameters to configure a `tf.config.PhysicalDevice`
+  as it is initialized to a `tf.config.LogicalDevice` during runtime
+  initialization. Not all fields are valid for all device types.
+
+  See `tf.config.get_logical_device_configuration` and
+  `tf.config.set_logical_device_configuration` for usage examples.
+
+  Fields:
+    memory_limit: (optional) Maximum memory (in MB) to allocate on the virtual
+      device. Currently only supported for GPUs.
+    experimental_priority: (optional) Priority to assign to a virtual device.
+      Lower values have higher priorities and 0 is the default.
+      Within a physical GPU, the GPU scheduler will prioritize ops on virtual
+      devices with higher priority. Currently only supported for Nvidia GPUs.
+    experimental_device_ordinal: (optional) Ordinal number to order the virtual
+    device.
+      LogicalDevice with lower ordinal number will receive a lower device id.
+      Physical device id and location in the list is used to break ties.
+      Currently only supported for Nvidia GPUs.
+  """
+
+  def __new__(
+      cls,
+      memory_limit=None,
+      experimental_priority=None,
+      experimental_device_ordinal=None,
+  ):
+    return super().__new__(
+        cls, memory_limit, experimental_priority, experimental_device_ordinal
+    )
+
+
+@tf_export("config.PhysicalDevice")
+class PhysicalDevice(
+    collections.namedtuple("PhysicalDevice", ["name", "device_type"])
+):
+  """Abstraction for a locally visible physical device.
+
+  TensorFlow can utilize various devices such as the CPU or multiple GPUs
+  for computation. Before initializing a local device for use, the user can
+  customize certain properties of the device such as it's visibility or memory
+  configuration.
+
+  Once a visible `tf.config.PhysicalDevice` is initialized one or more
+  `tf.config.LogicalDevice` objects are created. Use
+  `tf.config.set_visible_devices` to configure the visibility of a physical
+  device and `tf.config.set_logical_device_configuration` to configure multiple
+  `tf.config.LogicalDevice` objects for a `tf.config.PhysicalDevice`. This is
+  useful when separation between models is needed or to simulate a multi-device
+  environment.
+
+  Fields:
+    name: Unique identifier for device.
+    device_type: String declaring the type of device such as "CPU" or "GPU".
+  """
+
+  pass
+
+
+class _AtomicCounter(object):
+  """A simple atomic counter."""
+
+  __slots__ = ["_value", "_lock"]
+
+  def __init__(self):
+    self._value = 0
+    self._lock = threading.Lock()
+
+  def increment_and_get(self):
+    with self._lock:
+      self._value += 1
+      return self._value
+
+
+_context_id_counter = _AtomicCounter()
+
+
+class _TensorCacheDeleter(object):
+  """Deletes tensor caches for a given context."""
+
+  __slots__ = ["_context_id"]
+
+  def __init__(self, context_id):
+    self._context_id = context_id
+
+  def __del__(self):
+    if _tensor_caches_map is None:
+      return
+    if self._context_id in _tensor_caches_map:
+      del _tensor_caches_map[self._context_id]
+
+
+# TODO(agarwal): rename to EagerContext / EagerRuntime ?
+# TODO(agarwal): consider keeping the corresponding Graph here.
+class Context:
+  """Environment in which eager operations execute."""
+
+  # TODO(agarwal): create and link in some documentation for `execution_mode`.
+  # pylint: disable=redefined-outer-name
+  def __init__(
+      self,
+      config=None,
+      device_policy=None,
+      execution_mode=None,
+      server_def=None,
+  ):
+    """Creates a new Context.
+
+    Args:
+      config: (Optional.) A `ConfigProto` protocol buffer with configuration
+        options for the Context. Note that a lot of these options may be
+        currently unimplemented or irrelevant when eager execution is enabled.
+      device_policy: (Optional.) What policy to use when trying to run an
+        operation on a device with inputs which are not on that device. When set
+        to None, an appropriate value will be picked automatically. The value
+        picked may change between TensorFlow releases.  Defaults to
+        DEVICE_PLACEMENT_SILENT. Valid values: DEVICE_PLACEMENT_EXPLICIT -
+        raises an error if the placement is not correct. DEVICE_PLACEMENT_WARN -
+        copies the tensors which are not on the right device but raises a
+        warning. DEVICE_PLACEMENT_SILENT - silently copies the tensors. This
+        might hide performance problems. DEVICE_PLACEMENT_SILENT_FOR_INT32 -
+        silently copies int32 tensors, raising errors on the other ones.
+      execution_mode: (Optional.) Policy controlling how operations dispatched
+        are actually executed. When set to None, an appropriate value will be
+        picked automatically. The value picked may change between TensorFlow
+        releases. Valid values: - SYNC: executes each operation synchronously.
+        ASYNC - executes each operation asynchronously. These operations may
+        return "non-ready" handles.
+      server_def: (Optional.) A tensorflow::ServerDef proto. Enables execution
+        on remote devices. GrpcServers need to be started by creating an
+        identical server_def to this, and setting the appropriate task_indexes,
+        so that the servers can communicate. It will then be possible to execute
+        operations on remote devices.
+
+    Raises:
+     ValueError: If execution_mode is not valid.
+    """
+    # This _id is used only to index the tensor caches.
+    # TODO(iga): Remove this when tensor caches are moved to C++.
+    self._id = _context_id_counter.increment_and_get()
+    self._tensor_cache_deleter = _TensorCacheDeleter(self._id)
+    _tensor_caches_map[self._id] = _TensorCaches()
+
+    self._config = config
+    self._thread_local_data = pywrap_tfe.EagerContextThreadLocalData(
+        self,
+        is_eager=lambda: default_execution_mode == EAGER_MODE,
+        device_spec=_starting_device_spec,
+    )
+    self._context_switches = _ContextSwitchStack(self.executing_eagerly())
+    self._context_handle = None
+    self._context_devices = None
+    self._seed = None
+    self._initialize_lock = threading.Lock()
+    self._initialized = False
+    if device_policy is None:
+      device_policy = DEVICE_PLACEMENT_SILENT
+    self._device_policy = device_policy
+    self._mirroring_policy = None
+    if execution_mode not in (None, SYNC, ASYNC):
+      raise ValueError(
+          "execution_mode should be None/SYNC/ASYNC. Got %s" % execution_mode
+      )
+    if execution_mode is None:
+      execution_mode = SYNC
+    self._default_is_async = execution_mode == ASYNC
+    self._use_tfrt = is_tfrt_enabled()
+    self._jit_compile_rewrite = jit_compile_rewrite_enabled()
+    self._xla_sharding_for_resource_variables = (
+        xla_sharding_for_resource_variables_enabled()
+    )
+    self._optionals_in_gradients = optionals_in_gradients_enabled()
+    self._server_def = server_def
+    self._collective_ops_server_def = None
+    self._collective_leader = None
+    self._collective_scoped_allocator_enabled_ops = None
+    self._collective_use_nccl_communication = None
+    self._collective_device_filters = None
+    self._coordination_service_config = None
+
+    self._device_lock = threading.Lock()
+    self._physical_devices = None
+    self._physical_device_to_index = None
+    self._pluggable_devices = None
+    self._visible_device_list = []
+    self._memory_growth_map = None
+    self._virtual_device_map = {}
+
+    # Values set after construction
+    self._optimizer_jit = None
+    self._intra_op_parallelism_threads = None
+    self._inter_op_parallelism_threads = None
+    self._soft_device_placement = None
+    self._log_device_placement = None
+    self._operation_timeout_in_ms = None
+    self._enable_mlir_graph_optimization = None
+    self._optimizer_experimental_options = {}
+
+    _python_eager_context_create_counter.get_cell().increase_by(1)
+
+    self._is_global_context = False
+
+    # Number of retries to give the SetServerDef step. This is useful for fault
+    # tolerant initial connection in high-preemption settings like
+    # ParameterServerStrategy training.
+    self._set_server_def_retries = 0
+
+  # pylint: enable=redefined-outer-name
+
+  def _set_global_seed(self, seed):
+    """Set a global eager mode seed for random ops."""
+    self._seed = seed
+    # `random.Random(seed)` needs `seed` to be hashable, while values of type
+    # e.g. `np.int64` or `np.ndarray` are not. We use `int(...)` to convert them
+    # to int.
     try:
-      context._set_context(new_context)
-      # Use a 2D tensor so that it is not cached.
-      tensor1 = constant_op.constant([[3.]])
-      # Produce a tensor as an operation output. This uses a different code path
-      # from tensors created from Python.
-      tensor2 = tensor1 * tensor1
-      context._set_context(original_context)
-    except:
-      context._set_context(original_context)
-      raise
+      hash(seed)
+      self._rng = random.Random(seed)
+    except TypeError:
+      seed = int(np.array(seed))
+      self._rng = random.Random(seed)
+    # Also clear the kernel cache, to reset any existing seeds
+    if self._context_handle is not None:
+      pywrap_tfe.TFE_ContextClearCaches(self._context_handle)
 
-    # Deleting our context reference should not delete the underlying object.
-    del new_context
-    self.assertIsNot(weak_c(), None)
+  def _internal_operation_seed(self):
+    """Returns a fake operation seed.
 
-    # Deleting the first tensor should not delete the context since there is
-    # another tensor.
-    del tensor1
-    self.assertIsNot(weak_c(), None)
+      In eager mode, user shouldn't set or depend on operation seed.
+      Here, we generate a random seed based on global seed to make
+      operation's randomness different and depend on the global seed.
 
-    # Deleting the last tensor should result in deleting its context.
-    del tensor2
-    self.assertIs(weak_c(), None)
+    Returns:
+      A fake operation seed based on global seed.
+    """
+    return self._rng.randint(0, _MAXINT32)
 
-  def testSimpleGraphCollection(self):
+  def _initialize_logical_devices(self):
+    """Helper to initialize devices."""
+    # Store list of devices
+    logical_devices = []
+    context_devices = []
+    device_list = pywrap_tfe.TFE_ContextListDevices(self._context_handle)
+    try:
+      self._num_gpus = 0
+      current_job, current_task = None, None
+      server_def = self._server_def or self._collective_ops_server_def
+      if server_def is not None:
+        current_job, current_task = server_def.job_name, server_def.task_index
+      for i in range(pywrap_tfe.TF_DeviceListCount(device_list)):
+        dev_name = pywrap_tfe.TF_DeviceListName(device_list, i)
+        context_devices.append(pydev.canonical_name(dev_name))
+        spec = pydev.DeviceSpec.from_string(dev_name)
+        # If the job is localhost, we assume that the cluster has not yet been
+        # configured and thus clear the job, replica & task.
+        if spec.job == "localhost":
+          spec = spec.replace(job=None, replica=None, task=None)
+        logical_devices.append(
+            LogicalDevice(name=spec.to_string(), device_type=spec.device_type)
+        )
+        dev_type = pywrap_tfe.TF_DeviceListType(device_list, i)
+        if (
+            dev_type == "GPU"
+            and spec.job == current_job
+            and spec.task == current_task
+        ):
+          self._num_gpus += 1
 
-    @def_function.function
-    def f(x):
-      with ops.device('CPU:0'):
-        return x + constant_op.constant(1.)
+    finally:
+      self._logical_devices = logical_devices
+      self._context_devices = context_devices
+      pywrap_tfe.TF_DeleteDeviceList(device_list)
 
-    with context.collect_graphs() as graphs:
-      with ops.device('CPU:0'):
-        x = constant_op.constant(1.)
-      f(x)
+  def ensure_initialized(self):
+    """Initialize handle and devices if not already done so."""
+    if self._initialized:
+      return
+    with self._initialize_lock:
+      if self._initialized:
+        return
+      assert self._context_devices is None
+      opts = pywrap_tfe.TFE_NewContextOptions()
+      try:
+        config_str = self.config.SerializeToString()
+        pywrap_tfe.TFE_ContextOptionsSetConfig(opts, config_str)
+        if self._device_policy is not None:
+          pywrap_tfe.TFE_ContextOptionsSetDevicePlacementPolicy(
+              opts, self._device_policy
+          )
+        if self._mirroring_policy is not None:
+          pywrap_tfe.TFE_ContextOptionsSetMirroringPolicy(
+              opts, self._mirroring_policy
+          )
+        if self._default_is_async == ASYNC:
+          pywrap_tfe.TFE_ContextOptionsSetAsync(opts, True)
+        if self._use_tfrt is not None:
+          pywrap_tfe.TFE_ContextOptionsSetTfrt(opts, self._use_tfrt)
+        pywrap_tfe.TFE_ContextOptionsSetRunEagerOpAsFunction(opts, True)
+        pywrap_tfe.TFE_ContextOptionsSetJitCompileRewrite(
+            opts, self._jit_compile_rewrite
+        )
+        context_handle = pywrap_tfe.TFE_NewContext(opts)
+      finally:
+        pywrap_tfe.TFE_DeleteContextOptions(opts)
+      assert not (self._server_def and self._collective_ops_server_def), (
+          "Cannot enable remote execution as well as collective ops at the "
+          "moment. If this is important to you, please file an issue."
+      )
+      if self._server_def is not None:
+        server_def_str = self._server_def.SerializeToString()
+        timeout = 0  # Indicates no timeout.
+        pywrap_tfe.TFE_ContextSetServerDefWithTimeoutAndRetries(
+            context_handle,
+            _KEEP_ALIVE_SECS,
+            server_def_str,
+            timeout,
+            self._set_server_def_retries,
+        )
+      elif self._collective_ops_server_def is not None:
+        server_def_str = self._collective_ops_server_def.SerializeToString()
+        pywrap_tfe.TFE_EnableCollectiveOps(context_handle, server_def_str)
 
-    self.assertLen(graphs, 1)
-    graph, = graphs
-    self.assertIn('CPU:0', graph.node[1].device)
+      self._context_handle = context_handle
+      self._initialize_logical_devices()
+      self._initialized = True
 
-  @test_util.disable_tfrt(
-      'b/171600738: tfrt does not support exporting post-optimization graph')
-  def testGraphCollectionAfterDevicePlacement(self):
+      if self._is_global_context:
+        pywrap_tfe.TFE_Py_SetCEagerContext(self._context_handle)
 
-    @def_function.function
-    def f(x):
-      return x + constant_op.constant(1.)
+  def ensure_uninitialized(self):
+    """Uninitialize handle and devices if not already done so."""
+    with self._initialize_lock:
+      if not self._initialized:
+        return
+      self._context_devices = None
+      self._logical_devices = None
+      self._server_def = None
+      self._initialized = False
 
-    with context.collect_graphs() as graphs:
-      with ops.device('CPU:0'):
-        f(constant_op.constant(1.))
+      if self._is_global_context:
+        pywrap_tfe.TFE_Py_SetCEagerContext(None)
 
-    self.assertLen(graphs, 1)
-    graph, = graphs
-    self.assertIn('CPU:0', graph.node[0].device)
+      self._context_handle = None
 
-  def testGetFunctionDef(self):
+  def mark_as_global_context(self):
+    # If the context was already initialized, publish it. Otherwise wait with
+    # publication until it's initialized.
+    if self._initialized:
+      pywrap_tfe.TFE_Py_SetCEagerContext(self._context_handle)
+    self._is_global_context = True
 
-    @def_function.function
-    def f():
-      return constant_op.constant(1.)
+  def _clear_caches(self):
+    self.ones_rank_cache().flush()
+    self.zeros_cache().flush()
+    pywrap_tfe.TFE_ClearScalarCache()
 
-    concrete = f.get_concrete_function()
-    function_def = context.get_function_def(concrete.name)
+  def get_server_def(self):
+    return self._server_def
 
-    self.assertIsNot(function_def, None)
+  def set_server_def(self, server_def, keep_alive_secs=_KEEP_ALIVE_SECS):
+    """Allow setting a server_def on the context.
 
-    found_const_node = False
-    for node_def in function_def.node_def:
-      if node_def.op == 'Const':
-        found_const_node = True
-        break
-    self.assertTrue(found_const_node)
+    When a server def is replaced, it effectively clears a bunch of caches
+    within the context. If you attempt to use a tensor object that was pointing
+    to a tensor on the remote device, it will raise an error.
 
-    with self.assertRaises(errors.NotFoundError):
-      _ = context.get_function_def('this_should_not_be_found')
+    Args:
+      server_def: A tensorflow::ServerDef proto. Enables execution on remote
+        devices.
+      keep_alive_secs: Num. seconds after which the remote end will hang up. As
+        long as the client is still alive, the server state for the context will
+        be kept alive. If the client is killed (or there is some failure), the
+        server will clean up its context keep_alive_secs after the final RPC it
+        receives.
 
-  @test_util.run_gpu_only
-  @test_util.disable_tfrt('b/169293680: TFE_GetTotalMemoryUsage is unsupported')
-  def testGetMemoryInfo(self):
-    array_ops.zeros([10])  # Allocate some memory on the GPU.
-    self.assertGreater(context.context().get_memory_info('GPU:0')['current'], 0)
+    Raises:
+      ValueError: if server_def is None.
+    """
+    if not server_def:
+      raise ValueError("server_def is None.")
 
-  @test_util.disable_tfrt('b/169293680: TFE_GetTotalMemoryUsage is unsupported')
-  def testGetMemoryInfoCPU(self):
-    if test_util.IsMklEnabled():
-      # TODO(gzmkl) work with Google team to address design issue in allocator.h
-      self.skipTest('MklCPUAllocator does not throw exception. So skip test.')
+    self._server_def = server_def
 
-    with self.assertRaisesRegex(ValueError, 'Allocator stats not available'):
-      context.context().get_memory_info('CPU:0')
+    if self._context_handle:
+      server_def_str = server_def.SerializeToString()
+      pywrap_tfe.TFE_ContextSetServerDef(
+          self._context_handle, keep_alive_secs, server_def_str
+      )
+      self._initialize_logical_devices()
 
-  @test_util.disable_tfrt('b/169293680: TFE_GetTotalMemoryUsage is unsupported')
-  def testGetMemoryInfoUnknownDevice(self):
-    with self.assertRaisesRegex(ValueError, 'No matching devices found'):
-      context.context().get_memory_info('unknown_device:0')
+    # Clear all the caches in case there are remote tensors in them.
+    self._clear_caches()
+    # Also clear the device parsing cache since it caches the resolution of
+    # partial device names, which may become different due to the set_server_def
+    # call as we may have defined different devices.
+    _device_parsing_cache.clear()
 
-  @test_util.disable_tfrt('b/169293680: TFE_GetTotalMemoryUsage is unsupported')
-  def testGetMemoryInfoUnparsableDevice(self):
-    with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
-      context.context().get_memory_info('GPU')
-    with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
-      context.context().get_memory_info('GPU:')
-    with self.assertRaisesRegex(ValueError, 'Failed parsing device name'):
-      context.context().get_memory_info('GPU:CPU')
+  def update_server_def(self, server_def, keep_alive_secs=_KEEP_ALIVE_SECS):
+    """Update a server_def on the context.
 
-  def testListFunctionNames(self):
+    Args:
+      server_def: A tensorflow::ServerDef proto. Enables execution on remote
+        devices.
+      keep_alive_secs: Num. seconds after which the remote end will hang up. As
+        long as the client is still alive, the server state for the context will
+        be kept alive. If the client is killed (or there is some failure), the
+        server will clean up its context keep_alive_secs after the final RPC it
+        receives.
 
-    @def_function.function
-    def f():
-      return constant_op.constant(1.)
+    Raises:
+      ValueError: if server_def is None.
+    """
+    if not server_def:
+      raise ValueError("server_def is None.")
 
-    concrete = f.get_concrete_function()
-    self.assertIn(concrete.name.decode(),
-                  context.context().list_function_names())
+    self._server_def = server_def
 
-  def testCoordinationServiceMethodsInitializeNonGlobalContext(self):
-    # Each method must initialize its own (non-global) context before calling
-    # the C API; without a coordination service it then fails cleanly.
-    with self.assertRaisesRegex(errors.FailedPreconditionError,
-                                'Coordination service agent is not enabled'):
-      context.Context().set_config_key_value('key', 'value')
-    with self.assertRaisesRegex(errors.FailedPreconditionError,
-                                'Coordination service is not enabled'):
-      context.Context().get_config_key_value('key', timeout_in_ms=1)
-    with self.assertRaisesRegex(errors.FailedPreconditionError,
-                                'Coordination service is not enabled'):
-      context.Context().delete_config_key_value('key')
-    with self.assertRaisesRegex(errors.FailedPreconditionError,
-                                'Coordination service is not enabled'):
-      context.Context().wait_at_barrier('barrier', 1)
+    if self._context_handle:
+      server_def_str = server_def.SerializeToString()
+      pywrap_tfe.TFE_ContextUpdateServerDef(
+          self._context_handle, keep_alive_secs, server_def_str
+      )
+      self._initialize_logical_devices()
 
-  def testSetLogicalDeviceAfterContextInitialization(self):
-    ctx = context.Context()
-    ctx.set_logical_cpu_devices(4)
-    self.assertIs(len(ctx.list_logical_devices('CPU')), 4)
+    self._clear_caches()
 
-    # Cannot set logical device twice.
-    with self.assertRaisesRegex(RuntimeError, 'Virtual CPUs already set'):
-      ctx.set_logical_cpu_devices(8)
+  def check_alive(self, worker_name):
+    """Checks whether a remote worker is alive or not.
 
-  def testSetLogicalLocalCpuDevice(self):
-    ctx = context.Context()
+    Args:
+      worker_name: a string representing the remote worker. It must be a fully
+        specified name like "/job:worker/replica:0/task:0".
 
-    # Manually add a remote CPU device into logical device list.
-    ctx._logical_devices = []  # pylint: disable=protected-access
-    dev = context.LogicalDevice(name='/job:worker/replica:0/task:1',
-                                device_type='CPU')
-    ctx._logical_devices.append(dev)  # pylint: disable=protected-access
-    self.assertIs(len(ctx.list_logical_devices('CPU')), 1)
+    Returns:
+      a boolean indicating whether the remote worker is alive or not.
 
-    # This would pass the check since the previously added device is not local.
-    ctx.set_logical_cpu_devices(4)
-    # Logical device list would be overwritten after initialization.
-    self.assertIs(len(ctx.list_logical_devices('CPU')), 4)
+    Raises:
+      ValueError: if context is not initialized.
+    """
+    # TODO(yuefengz): support checking multiple workers.
+    if self._context_handle:
+      return pywrap_tfe.TFE_ContextCheckAlive(self._context_handle, worker_name)
+    else:
+      raise ValueError("Context is not initialized.")
 
-  @parameterized.named_parameters([(f'_{stage}', stage) for stage in [
-      'hlo', 'hlo_serialized', 'optimized_hlo', 'optimized_hlo_serialized',
-      'optimized_hlo_proto_serialized', 'optimized_hlo_dot'
-  ]])
-  def testGetCompilerIr(self, stage):
+  def sync_executors(self):
+    """Sync both local executors and the ones on remote workers.
 
-    @def_function.function(jit_compile=True)
-    def test_func(x):
-      return 2 * x
+    In async execution mode, local function calls can return before the
+    corresponding remote op/function execution requests are completed. Calling
+    this method creates a synchronization barrier for remote executors. It only
+    returns when all remote pending nodes are finished, potentially with errors
+    if any remote executors are in error state.
 
-    a = array_ops.ones((1000, 1000))  # 4 * 1000 * 1000 in bytes
-    result = test_func.experimental_get_compiler_ir(a)(stage=stage)
-    self.assertNotEmpty(result)
-    if stage == 'optimized_hlo_proto_serialized':
-      hlo_proto = hlo_pb2.HloProto.FromString(result)
-      allocations = hlo_proto.buffer_assignment.buffer_allocations
-      buffer_size = sum(
-          getattr(allocation, 'size') for allocation in allocations)
-      # The sizes of input and output are both 4 * 1000 * 1000 in bytes.
-      self.assertGreaterEqual(buffer_size, 2 * 4 * 1000 * 1000)
-      self.assertLess(buffer_size, 4 * 4 * 1000 * 1000)
+    Raises:
+      ValueError: if context is not initialized.
+    """
+    if self._context_handle:
+      pywrap_tfe.TFE_ContextSyncExecutors(self._context_handle)
+    else:
+      raise ValueError("Context is not initialized.")
+
+  def clear_executor_errors(self):
+    """Clear errors in both local executors and remote workers.
+
+    After receiving errors from remote workers, additional requests on the fly
+    could further taint the status on the remote workers due to the async nature
+    of remote execution. Calling this method block on waiting for all pending
+    nodes in remote executors to finish and clear their error statuses.
+
+    Raises:
+      ValueError: if context is not initialized.
+    """
+    if self._context_handle:
+      pywrap_tfe.TFE_ContextClearExecutors(self._context_handle)
+    else:
+      raise ValueError("Context is not initialized.")
+
+  def configure_coordination_service(
+      self,
+      service_type,
+      service_leader="",
+      enable_health_check=True,
+      cluster_register_timeout_in_ms=0,
+      heartbeat_timeout_in_ms=0,
+      shutdown_barrier_timeout_in_ms=0,
+      coordinated_jobs=None,
+      allow_new_incarnation_to_reconnect=False,
+  ):
+    """Enable distributed coordination service with specified configs."""
+    if self._context_handle:
+      logging.warning(
+          "Configuring coordination service type may not be "
+          "effective because the context is already initialized."
+      )
+    config = coordination_config_pb2.CoordinationServiceConfig()
+    config.service_type = service_type
+    if service_leader:
+      config.service_leader = pydev.canonical_name(service_leader)
+    config.enable_health_check = enable_health_check
+    config.cluster_register_timeout_in_ms = cluster_register_timeout_in_ms
+    config.heartbeat_timeout_in_ms = heartbeat_timeout_in_ms
+    config.shutdown_barrier_timeout_in_ms = shutdown_barrier_timeout_in_ms
+    config.allow_new_incarnation_to_reconnect = (
+        allow_new_incarnation_to_reconnect
+    )
+    if coordinated_jobs is not None:
+      if isinstance(coordinated_jobs, list):
+        config.coordinated_job_list.extend(coordinated_jobs)
+      else:
+        raise ValueError(
+            "`coordinated_jobs` must be list[CoordinatedJob] or "
+            "None, but got: %s" % (coordinated_jobs,)
+        )
+    self._coordination_service_config = config
+
+  @property
+  def coordination_service(self):
+    return self._coordination_service_config
+
+  def set_config_key_value(self, key, value):
+    self.ensure_initialized()
+    pywrap_tfe.TFE_InsertConfigKeyValue(self._context_handle, key, value)
+
+  # If `timeout_in_ms=0`, this will block until the key-value is set or the
+  # worker shuts down.
+  def get_config_key_value(self, key, timeout_in_ms=0):
+    self.ensure_initialized()
+    with c_api_util.tf_buffer() as buffer_:
+      pywrap_tfe.TFE_GetConfigKeyValue(
+          self._context_handle, key, timeout_in_ms, buffer_
+      )
+      value = pywrap_tf_session.TF_GetBuffer(buffer_).decode("utf-8")
+    return value
+
+  def delete_config_key_value(self, key):
+    self.ensure_initialized()
+    pywrap_tfe.TFE_DeleteConfigKeyValue(self._context_handle, key)
+
+  def report_error_to_cluster(self, error_code, error_message):
+    """Report error to other members in a multi-client cluster.
+
+    Args:
+      error_code: a `tf.errors` error code.
+      error_message: a string. The error message.
+    """
+    if self._context_handle:
+      pywrap_tfe.TFE_ReportErrorToCluster(
+          self._context_handle, error_code, error_message
+      )
+    else:
+      raise ValueError("Context is not initialized.")
+
+  def get_task_states(self, job_configs):
+    """Get task states from the Coordination Service.
+
+    Args:
+      job_configs: A list of tuples of job name and task number.
+
+    Returns:
+      A list of TF_Status.
+    """
+    if self._context_handle:
+      job_names, task_nums = zip(*job_configs)
+      return pywrap_tfe.TFE_GetTaskStates(
+          self._context_handle, job_names, task_nums
+      )
+    else:
+      raise ValueError("Context is not initialized.")
+
+  def wait_at_barrier(self, barrier_id, timeout_in_ms):
+    """Blocks until all coordinated tasks are at the barrier.
+
+    The barrier may fail if it times out or if one of the tasks is unhealthy.
+
+    Args:
+      barrier_id: Unique string identifying the barrier.
+      timeout_in_ms: Duration before the barrier times out and fails.
+    """
+    self.ensure_initialized()
+    pywrap_tfe.TFE_WaitAtBarrier(
+        self._context_handle, barrier_id, timeout_in_ms
+    )
+
+  def clear_kernel_cache(self):
+    """Clear kernel cache and reset all stateful kernels."""
+    if self._context_handle is not None:
+      pywrap_tfe.TFE_ContextClearCaches(self._context_handle)
+
+  def enable_collective_ops(self, server_def):
+    """Enable distributed collective ops with an appropriate server_def.
+
+    Args:
+      server_def: A tensorflow::ServerDef proto. Enables execution on remote
+        devices.
+
+    Raises:
+      ValueError: if server_def is None.
+      RuntimeError: if this method is not called at program startup.
+    """
+    if not server_def:
+      raise ValueError("server_def is None.")
+
+    self._collective_ops_server_def = server_def
+
+    # TODO(b/129298253): Allow creating datasets/tensors before enabling
+    # collective ops.
+    if self._context_handle is not None:
+      logging.warning(
+          "Enabling collective ops after program startup may cause "
+          "error when accessing previously created tensors."
+      )
+      with self._initialize_lock:
+        assert self._initialized
+        server_def_str = self._collective_ops_server_def.SerializeToString()
+        pywrap_tfe.TFE_EnableCollectiveOps(self._context_handle, server_def_str)
+        self._initialize_logical_devices()
+        self._clear_caches()
+
+  def configure_collective_ops(
+      self,
+      collective_leader="",
+      scoped_allocator_enabled_ops=("CollectiveReduce",),
+      use_nccl_communication=False,
+      device_filters=None,
+  ):
+    """Configure collective ops.
+
+      Collective group leader is necessary for collective ops to run, other
+      configurations are mainly for the purpose of performance.
+
+    Args:
+      collective_leader: a device string for collective leader, e.g.
+        "/job:worker/replica:0/task:0"; empty string means local execution of
+        collective ops.
+      scoped_allocator_enabled_ops: a tuple or a list of op names for scoped
+        allocator to run with.
+      use_nccl_communication: whether to use nccl communication for collective
+        ops.
+      device_filters: a tuple or a list of device strings. If set, corresponding
+        task can only see the devices filtered by these device filters.
+
+    Raises:
+      RuntimeError: if this method is not called at program startup.
+    """
+    if self._collective_leader is not None:
+      if (
+          self._collective_leader != collective_leader
+          or self._collective_scoped_allocator_enabled_ops
+          != scoped_allocator_enabled_ops
+          or self._collective_use_nccl_communication != use_nccl_communication
+          or self._collective_device_filters != device_filters
+      ):
+        raise ValueError("Collective ops are already configured.")
+      else:
+        return
+
+    if self._context_handle is not None:
+      raise RuntimeError("Collective ops must be configured at program startup")
+
+    self._collective_leader = collective_leader
+    self._collective_scoped_allocator_enabled_ops = scoped_allocator_enabled_ops
+    self._collective_use_nccl_communication = use_nccl_communication
+    self._collective_device_filters = device_filters
+
+  def abort_collective_ops(self, code, message):
+    """Abort the collective ops.
+
+    This is intended to be used when a peer failure is detected, which allows
+    the user to handle the case instead of hanging. This aborts all on-going
+    collectives. After all subsequent collectives error immediately, and you
+    need to reset_context() to use collectives again.
+
+    Args:
+      code: a `tf.errors` error code.
+      message: a string. The error message.
+    """
+    self.ensure_initialized()
+    pywrap_tfe.TFE_AbortCollectiveOps(self._handle, code, message)
+
+  def check_collective_ops_peer_health(self, task, timeout_in_ms):
+    """Check collective peer health.
+
+    This probes each task to see if they're still alive. Note that restarted
+    tasks are considered a different one, and they're considered not healthy.
+
+    This should only be used in multi client multi worker training.
+
+    Args:
+      task: a task string, must be in the format of /job:xxx/replica:0/task:N.
+      timeout_in_ms: an integer, the timeout. If zero, there's no timeout.
+
+    Raises:
+      tf.errors.UnavailableError: when a peer is down.
+      tf.errors.FailedPreconditionError: when a peer is a different one from the
+        one this task has talked to, e.g. the peer has restarted.
+      tf.errors.InvalidArgumentError: when the task string is invalid.
+    """
+    self.ensure_initialized()
+    pywrap_tfe.TFE_CollectiveOpsCheckPeerHealth(
+        self._handle, task, timeout_in_ms
+    )
+
+  @property
+  def _handle(self):
+    if self._context_handle is None:
+      raise AssertionError("Context must be initialized first.")
+
+    return self._context_handle
+
+  @property
+  def _devices(self):
+    if self._context_devices is None:
+      raise AssertionError("Context must be initialized first.")
+
+    return self._context_devices
+
+  def __str__(self):
+    if self._context_handle is None:
+      return "Eager TensorFlow Context. Devices currently uninitialized."
+    else:
+      devices = self._devices
+      lines = ["Eager TensorFlow Context with %d devices" % (len(devices))]
+      for i, d in enumerate(devices):
+        lines.append("   Device %d: %s" % (i, d))
+      return "\n".join(lines)
+
+  @tf_contextlib.contextmanager
+  def _mode(self, mode):
+    """A context manager to allow setting the mode to EAGER/GRAPH."""
+    ctx = self._thread_local_data
+    old_is_eager = ctx.is_eager
+    ctx.is_eager = mode == EAGER_MODE
+    if mode == EAGER_MODE:
+      # Entering graph mode does not provide us with sufficient information to
+      # record a context switch; graph-based context switches are only logged
+      # when a graph is registered as the default graph.
+      self.context_switches.push(False, eager_mode, None)
+    try:
+      yield
+    finally:
+      ctx.is_eager = old_is_eager
+      if mode == EAGER_MODE:
+        self.context_switches.pop()
+
+  def executing_eagerly(self):
+    """Returns True if current thread has eager executing enabled."""
+    return self._thread_local_data.is_eager
+
+  def ones_rank_cache(self):
+    """Per-device cache for scalars."""
+    return _tensor_caches_map[self._id].ones_rank_cache
+
+  def zeros_cache(self):
+    """Per-device cache for scalars."""
+    return _tensor_caches_map[self._id].zeros_cache
+
+  @property
+  def scope_name(self):
+    """Returns scope name for the current thread."""
+    return self._thread_local_data.scope_name
+
+  @scope_name.setter
+  def scope_name(self, s):
+    """Sets scope name for the current thread."""
+    self._thread_local_data.scope_name = s
+
+  @property
+  def device_name(self):
+    """Returns the device name for the current thread."""
+    return self._thread_local_data.device_name
+
+  @property
+  def device_spec(self):
+    """Returns the device spec for the current thread."""
+    return self._thread_local_data.device_spec
+
+  def _set_device(self, device_name, device_spec):
+    self._thread_local_data.device_name = device_name
+    self._thread_local_data.device_spec = device_spec
+
+  def device(self, name):
+    """Context-manager to force placement of operations and Tensors on a device.
+
+    Args:
+      name: Name of the device or None to get default placement.
+
+    Returns:
+      Context manager that forces device placement.
+
+    Raises:
+      ValueError: If name is not a string or is an invalid device name.
+      RuntimeError: If device scopes are not properly nested.
+    """
+    if isinstance(name, LogicalDevice):
+      name = name.name
+    elif pydev.is_device_spec(name):
+      name = name.to_string()
+    return _EagerDeviceContext(self, name)
+
+  def devices(self):
+    """List of the names of devices available to execute operations."""
+    return self._devices
+
+  def host_address_space(self):
+    self.ensure_initialized()
+    with c_api_util.tf_buffer() as buffer_:
+      pywrap_tfe.TFE_HostAddressSpace(self._context_handle, buffer_)
+      address_space = pywrap_tf_session.TF_GetBuffer(buffer_).decode("utf-8")
+    return address_space
+
+  # TODO(fishx): remove this property.
+  @property
+  def execution_mode(self):
+    """Gets execution mode for current thread."""
+    return ASYNC if self.is_async() else SYNC
+
+  @execution_mode.setter
+  def execution_mode(self, mode):
+    """Sets execution mode for current thread."""
+    if mode not in (None, SYNC, ASYNC):
+      raise ValueError(
+          "Execution mode should be None/SYNC/ASYNC. Got %s" % mode
+      )
+
+    if mode is None:
+      mode = SYNC
+
+    enable_async = mode == ASYNC
+    if self.is_async() != enable_async:
+      # Only set the execution mode if the context has already been initialized
+      if self._context_handle is not None:
+        self.executor.wait()
+        executor_new = executor.new_executor(enable_async)
+        self._thread_local_data.executor = executor_new
+        pywrap_tfe.TFE_ContextSetExecutorForThread(
+            self._context_handle, executor_new.handle()
+        )
+      else:
+        self._default_is_async = enable_async
+
+  def is_async(self):
+    if self._context_handle is not None:
+      return self.executor.is_async()
+    else:
+      return self._default_is_async
+
+  @property
+  def executor(self):
+    self.ensure_initialized()
+    return executor.Executor(
+        pywrap_tfe.TFE_ContextGetExecutorForThread(self._context_handle)
+    )
+
+  @executor.setter
+  def executor(self, e):
+    self.ensure_initialized()
+    pywrap_tfe.TFE_ContextSetExecutorForThread(self._context_handle, e.handle())
+
+  @property
+  def config(self):
+    """Return the ConfigProto with all runtime deltas applied."""
+    # Ensure physical devices have been discovered and config has been imported
+    self._initialize_physical_devices()
+
+    config = config_pb2.ConfigProto()
+    if self._config is not None:
+      config.CopyFrom(self._config)
+
+    if self._optimizer_jit is not None:
+      config.graph_options.optimizer_options.global_jit_level = (
+          config_pb2.OptimizerOptions.ON_1
+          if self._optimizer_jit
+          else config_pb2.OptimizerOptions.OFF
+      )
+    if self._intra_op_parallelism_threads is not None:
+      config.intra_op_parallelism_threads = self._intra_op_parallelism_threads
+    if self._inter_op_parallelism_threads is not None:
+      config.inter_op_parallelism_threads = self._inter_op_parallelism_threads
+
+    if self._soft_device_placement is not None:
+      config.allow_soft_placement = self._soft_device_placement
+    else:
+      config.allow_soft_placement = self.executing_eagerly()
+
+    if self._log_device_placement is not None:
+      config.log_device_placement = self._log_device_placement
+
+    if self._operation_timeout_in_ms is not None:
+      config.operation_timeout_in_ms = self._operation_timeout_in_ms
+
+    is_mlir_bridge_enabled = pywrap_tfe.TF_IsMlirBridgeEnabled()
+    config.experimental.mlir_bridge_rollout = is_mlir_bridge_enabled
+    if (
+        is_mlir_bridge_enabled
+        == config_pb2.ConfigProto.Experimental.MLIR_BRIDGE_ROLLOUT_ENABLED
+    ):
+      config.experimental.enable_mlir_bridge = True
+
+    if self._enable_mlir_graph_optimization is not None:
+      config.experimental.enable_mlir_graph_optimization = (
+          self._enable_mlir_graph_optimization
+      )
+
+    def rewriter_toggle(option):
+      toggle = self._optimizer_experimental_options.get(option, None)
+      if toggle is None:
+        return
+
+      setattr(
+          config.graph_options.rewrite_options,
+          option,
+          (
+              rewriter_config_pb2.RewriterConfig.ON
+              if toggle
+              else rewriter_config_pb2.RewriterConfig.OFF
+          ),
+      )
+
+    def rewriter_bool(option):
+      toggle = self._optimizer_experimental_options.get(option, None)
+      if toggle is None:
+        return
+
+      setattr(config.graph_options.rewrite_options, option, toggle)
+
+    rewriter_toggle("layout_optimizer")
+    rewriter_toggle("constant_folding")
+    rewriter_toggle("shape_optimization")
+    rewriter_toggle("remapping")
+    rewriter_toggle("arithmetic_optimization")
+    rewriter_toggle("dependency_optimization")
+    rewriter_toggle("loop_optimization")
+    rewriter_toggle("function_optimization")
+    rewriter_toggle("debug_stripper")
+    rewriter_bool("disable_model_pruning")
+    rewriter_toggle("scoped_allocator_optimization")
+    rewriter_toggle("pin_to_host_optimization")
+    rewriter_toggle("implementation_selector")
+    rewriter_toggle("auto_mixed_precision")
+    rewriter_toggle("use_plugin_optimizers")
+    rewriter_bool("disable_meta_optimizer")
+    rewriter_toggle("auto_mixed_precision_onednn_bfloat16")
+    rewriter_toggle("auto_mixed_precision_mkl")
+    nodes = self._optimizer_experimental_options.get("min_graph_nodes", None)
+    if nodes is not None:
+      config.graph_options.rewrite_options.min_graph_nodes = nodes
+
+    # Compute device counts
+    config.device_count["CPU"] = 0
+    config.device_count["GPU"] = 0
+    for dev in self._physical_devices:
+      if dev not in self._visible_device_list:
+        continue
+
+      virtual_devices = self._virtual_device_map.get(dev)
+      if virtual_devices is None:
+        config.device_count[dev.device_type] += 1
+      else:
+        config.device_count[dev.device_type] += len(virtual_devices)
+
+    # Configure gpu_options
+    gpu_options = self._compute_gpu_options()
+    config.gpu_options.MergeFrom(gpu_options)
+
+    # Configure pluggable_device_options
+    pluggable_device_options = self._compute_pluggable_device_options()
+    config.pluggable_device_options.MergeFrom(pluggable_device_options)
+
+    # Configure collective ops
+    if self._collective_leader:
+      config.experimental.collective_group_leader = self._collective_leader
+    if self._collective_scoped_allocator_enabled_ops:
+      rewrite_options = config.graph_options.rewrite_options
+      rewrite_options.scoped_allocator_optimization = (
+          rewriter_config_pb2.RewriterConfig.ON
+      )
+      del rewrite_options.scoped_allocator_opts.enable_op[:]
+      for op in self._collective_scoped_allocator_enabled_ops:
+        rewrite_options.scoped_allocator_opts.enable_op.append(op)
+    if self._collective_use_nccl_communication:
+      config.experimental.collective_nccl = True
+    if self._collective_device_filters:
+      del config.device_filters[:]
+      for f in self._collective_device_filters:
+        config.device_filters.append(f)
+
+    # Configure coordination service
+    if self._coordination_service_config:
+      config.experimental.coordination_config.CopyFrom(
+          self._coordination_service_config
+      )
+
+    return config
+
+  def _compute_device_options(self, device_type="GPU"):
+    """Build the GPUOptions proto for GPU or PluggableDevice."""
+    if device_type not in ["GPU", "PluggableDevice"]:
+      raise ValueError(
+          "device types other than GPU and PluggableDevice are not supported."
+      )
+    visible_device_list = []
+    virtual_devices = []
+    # This mapping is needed to handle multiple sub types of PluggableDevices.
+    device_to_indices = collections.defaultdict(int)
+    memory_growths = set()
+    compatible_devices = (
+        self.list_physical_devices("GPU")
+        if device_type == "GPU"
+        else self._pluggable_devices
+    )
+    support_virtual_devices = device_type == "GPU"
+    for dev in compatible_devices:
+      device_index = device_to_indices[dev.device_type]
+      device_to_indices[dev.device_type] += 1
+
+      if dev not in self._visible_device_list:
+        continue
+
+      growth = self._memory_growth_map[dev]
+      memory_growths.add(growth)
+      if device_type == "PluggableDevice":
+        visible_device_list.append(dev.device_type + ":" + str(device_index))
+      else:
+        visible_device_list.append(str(device_index))
+
+      if support_virtual_devices and self._virtual_device_map:
+        vdevs = self._virtual_device_map.get(dev, [])
+        device_ordinals = []
+        device_limits = []
+        priority = []
+        for virt_dev in vdevs:
+          if virt_dev.experimental_device_ordinal is not None:
+            device_ordinals.append(virt_dev.experimental_device_ordinal)
+          device_limits.append(virt_dev.memory_limit)
+          if virt_dev.experimental_priority is not None:
+            priority.append(virt_dev.experimental_priority)
+        # If priority is specified, it must be specified for all virtual
+        # devices.
+        if priority and len(device_limits) != len(priority):
+          raise ValueError("priority must be specified for all virtual devices")
+        # If device_ordinals is specified, it must be specified for all virtual
+        # devices.
+        if device_ordinals and len(device_limits) != len(device_ordinals):
+          raise ValueError(
+              "device_ordinals must be specified for all virtual devices"
+          )
+
+        virtual_devices.append(
+            config_pb2.GPUOptions.Experimental.VirtualDevices(
+                memory_limit_mb=device_limits,
+                priority=priority,
+                device_ordinal=device_ordinals,
+            )
+        )
+
+    # Only compute growth if virtual devices have not been configured and we
+    # have GPUs
+    if not virtual_devices and memory_growths:
+      if len(memory_growths) > 1:
+        raise ValueError("Memory growth cannot differ between GPU devices")
+      allow_growth = memory_growths.pop()
+    else:
+      allow_growth = None
+
+    return config_pb2.GPUOptions(
+        allow_growth=allow_growth,
+        visible_device_list=",".join(visible_device_list),
+        experimental=config_pb2.GPUOptions.Experimental(
+            virtual_devices=virtual_devices
+        ),
+    )
+
+  def _compute_gpu_options(self):
+    """Build the GPUOptions proto for GPU."""
+    return self._compute_device_options(device_type="GPU")
+
+  def _compute_pluggable_device_options(self):
+    """Build the GPUOptions proto for PluggableDevice."""
+    return self._compute_device_options(device_type="PluggableDevice")
+
+  @property
+  def function_call_options(self):
+    """Returns function call options for current thread.
+
+    Note that the returned object is still referenced by the eager context.
+
+    Returns: the FunctionCallOptions for current thread.
+    """
+    if self._thread_local_data.function_call_options is None:
+      config = self.config
+
+      # Default to soft placement for functions unless specified
+      if self._soft_device_placement is None:
+        config.allow_soft_placement = True
+      self._thread_local_data.function_call_options = FunctionCallOptions(
+          config_proto=config
+      )
+
+    return self._thread_local_data.function_call_options
+
+  @function_call_options.setter
+  def function_call_options(self, options):
+    """Returns function call options for current thread."""
+    self._thread_local_data.function_call_options = options
+
+  def num_gpus(self):
+    """The number of GPUs available to execute operations."""
+    self.ensure_initialized()
+    return self._num_gpus
+
+  def add_c_function(self, c_func):
+    """Add a C API TF_Function to the context.
+
+    Once added, the function (identified by its name) can be executed like any
+    other operation.
+
+    Args:
+      c_func: A wrapped TF_Function (returned from TF_GraphToFunction_wrapper).
+    """
+    self.ensure_initialized()
+    pywrap_tfe.TFE_ContextAddFunction(self._handle, c_func)
+
+  def get_c_function(self, name):
+    """Get a C API TF_Function from the context.
+
+    Args:
+      name: Name of the function to get.
+
+    Returns:
+      A ScopedTFFunction wrapping the C API TF_Function.
+    """
+    self.ensure_initialized()
+    return c_api_util.ScopedTFFunction(
+        pywrap_tfe.TFE_ContextGetFunction(self._handle, name), name
+    )
+
+  def add_function_def(self, fdef):
+    """Add a function definition to the context.
+
+    Once added, the function (identified by its name) can be executed like any
+    other operation.
+
+    Args:
+      fdef: A FunctionDef protocol buffer message.
+    """
+    self.ensure_initialized()
+    if is_oss:
+      fdef_string = fdef.SerializeToString()
+      pywrap_tfe.TFE_ContextAddFunctionDef(
+          self._handle, fdef_string, len(fdef_string)
+      )
+    else:
+      pywrap_tfe.TFE_ContextAddFunctionDefNoSerialization(self._handle, fdef)
+
+  def get_function_def(self, name):
+    """Get a function definition from the context.
+
+    Args:
+      name: function signature name.
+
+    Returns:
+      The requested FunctionDef.
+
+    Raises:
+      tf.errors.NotFoundError: if name is not the name of a registered function.
+    """
+    if is_oss:
+      with c_api_util.tf_buffer() as buffer_:
+        pywrap_tfe.TFE_ContextGetFunctionDef(self._handle, name, buffer_)
+        proto_data = pywrap_tf_session.TF_GetBuffer(buffer_)
+      function_def = function_pb2.FunctionDef()
+      function_def.ParseFromString(proto_data)
+    else:
+      function_def = pywrap_tfe.TFE_ContextGetFunctionDefNoSerialization(
+          self._handle, name
+      )
+    return function_def
+
+  def get_graph_debug_info(self, name):
+    """Get GraphDebugInfo associated with a function from the context.
+
+    Args:
+      name: function signature name.
+
+    Returns:
+      The requested GraphDebugInfo.
+
+    Raises:
+      tf.errors.NotFoundError: if name is not the name of a registered function.
+    """
+    with c_api_util.tf_buffer() as buffer_:
+      pywrap_tfe.TFE_ContextGetGraphDebugInfo(self._handle, name, buffer_)
+      proto_data = pywrap_tf_session.TF_GetBuffer(buffer_)
+    graph_debug_info = graph_debug_info_pb2.GraphDebugInfo()
+    graph_debug_info.ParseFromString(proto_data)
+
+    return graph_debug_info
+
+  def is_custom_device(self, device_name):
+    """Calls TFE_IsCustomDevice. See the non-member function."""
+    self.ensure_initialized()
+    return pywrap_tfe.TFE_Py_IsCustomDevice(self._handle, device_name)
+
+  def register_custom_device(
+      self, device_capsule, device_name, device_info_capsule
+  ):
+    """Calls TFE_RegisterCustomDevice. See the non-member function."""
+    self.ensure_initialized()
+    pywrap_tfe.TFE_Py_RegisterCustomDevice(
+        self._handle, device_capsule, device_name, device_info_capsule
+    )
+
+  def pack_eager_tensors(self, tensors):
+    """Pack multiple `EagerTensor`s of the same dtype and shape.
+
+    Args:
+      tensors: a list of EagerTensors to pack.
+
+    Returns:
+      A packed EagerTensor.
+    """
+    self.ensure_initialized()
+    return pywrap_tfe.TFE_Py_PackEagerTensors(self._handle, tensors)
+
+  def list_function_names(self):
+    """Get a list of names of registered functions.
+
+    Returns:
+      A set of names of all registered functions for the context.
+    """
+    self.ensure_initialized()
+    return set(pywrap_tfe.TFE_ContextListFunctionNames(self._handle))
+
+  def remove_function(self, name):
+    """Remove a function from the context.
+
+    Once removed, the function cannot be executed anymore.
+
+    Args:
+      name: function signature name.
+    """
+    self.ensure_initialized()
+    pywrap_tfe.TFE_ContextRemoveFunction(self._handle, name)
+
+  def has_function(self, name):
+    """Check if a function `name` is registered."""
+    self.ensure_initialized()
+    return bool(pywrap_tfe.TFE_ContextHasFunction(self._handle, name))
+
+  @property
+  def function_scope_id(self):
+    """Returns an id that is unique to each scope holding functions."""
+    return id(self._context_handle)
+
+  def call_function(self, name, tensor_inputs, num_outputs):
+    """Calls the function associated with the given name."""
+    attrs = tuple(
+        itertools.chain(*self.function_call_options.as_attrs().items())
+    )
+
+    cancellation_context = cancellation.context()
+    if cancellation_context is None:
+      outputs = execute.execute(
+          name.decode("utf-8"),
+          num_outputs=num_outputs,
+          inputs=tensor_inputs,
+          attrs=attrs,
+          ctx=self,
+      )
+    else:
+      outputs = execute.execute_with_cancellation(
+          name.decode("utf-8"),
+          num_outputs=num_outputs,
+          inputs=tensor_inputs,
+          attrs=attrs,
+          ctx=self,
+          cancellation_manager=cancellation_context,
+      )
+    # Empty list means no function outputs so return None
+    outputs = outputs or None
+
+    return outputs
+
+  def add_op_callback(self, callback):
+    """Add a post-op callback to the context.
+
+    A post-op callback is invoked immediately after an eager operation or
+    function has finished execution or after a op has been added to a graph,
+    providing access to the op's type, name input and output tensors. Multiple
+    op callbacks can be added, in which case the callbacks will be invoked in
+    the order in which they are added.
+
+    Args:
+      callback: a callable of the signature `f(op_type, inputs, attrs, outputs,
+        op_name=None, graph=None)`. See doc strings in `op_callbacks.py` for
+        details on the function signature and its semantics.
+    """
+    if callback not in self._thread_local_data.op_callbacks:
+      self._thread_local_data.op_callbacks.append(callback)
+
+  def remove_op_callback(self, callback):
+    """Remove an already-registered op callback.
+
+    Args:
+      callback: The op callback to be removed.
+
+    Raises:
+      KeyError: If `callback` is not already registered.
+    """
+    if callback not in self._thread_local_data.op_callbacks:
+      raise KeyError(
+          "The specified op callback has not been registered, "
+          "and hence cannot be removed."
+      )
+    del self._thread_local_data.op_callbacks[
+        self._thread_local_data.op_callbacks.index(callback)
+    ]
+
+  @property
+  def op_callbacks(self):
+    return self._thread_local_data.op_callbacks
+
+  @property
+  def invoking_op_callbacks(self):
+    return self._thread_local_data.invoking_op_callbacks
+
+  @invoking_op_callbacks.setter
+  def invoking_op_callbacks(self, value):
+    self._thread_local_data.invoking_op_callbacks = value
+
+  def _initialize_physical_devices(self, reinitialize=False):
+    """Gets local devices visible to the system.
+
+    Args:
+      reinitialize: If True, reinitializes self._physical_devices  so that
+        dynamic registered devices will also be visible to the python front-end.
+    """
+    # We lazy initialize self._physical_devices since we do not want to do this
+    # the constructor since the backend may not be initialized yet.
+    with self._device_lock:
+      if not reinitialize and self._physical_devices is not None:
+        return
+
+      devs = pywrap_tfe.TF_ListPhysicalDevices()
+      self._physical_devices = [
+          PhysicalDevice(name=d.decode(), device_type=d.decode().split(":")[1])
+          for d in devs
+      ]
+      self._physical_device_to_index = {
+          p: i for i, p in enumerate(self._physical_devices)
+      }
+      # We maintain a separate list just so we can check whether the device in
+      # _physical_devices is a PluggableDevice.
+      pluggable_devs = pywrap_tfe.TF_ListPluggablePhysicalDevices()
+      self._pluggable_devices = [
+          PhysicalDevice(name=d.decode(), device_type=d.decode().split(":")[1])
+          for d in pluggable_devs
+      ]
+
+      self._visible_device_list = list(self._physical_devices)
+      self._memory_growth_map = {
+          d: None
+          for d in self._physical_devices
+          if d.device_type == "GPU" or d in self._pluggable_devices
+      }
+
+    # Import device settings that may have been passed into the constructor
+    self._import_config()
+
+  def reinitialize_physical_devices(self):
+    """Gets local devices visible to the system."""
+    # Reinitialize the physical device list after registering
+    # the pluggable device.
+    self._initialize_physical_devices(True)
+
+  def list_physical_devices(self, device_type=None):
+    """List local devices visible to the system.
+
+    This API allows a client to query the devices before they have been
+    initialized by the eager runtime. Additionally a user can filter by device
+    type, to get only CPUs or GPUs.
+
+    Args:
+      device_type: Optional device type to limit results to
+
+    Returns:
+      List of PhysicalDevice objects.
+    """
+    self._initialize_physical_devices()
+
+    if device_type is None:
+      return list(self._physical_devices)
+
+    return [d for d in self._physical_devices if d.device_type == device_type]
+
+  def get_device_details(self, device):  # pylint: disable=redefined-outer-name
+    """Returns details about a physical devices.
+
+    Args:
+      device: A `tf.config.PhysicalDevice` returned by
+        `tf.config.list_physical_devices` or `tf.config.get_visible_devices`.
+
+    Returns:
+      A dict with string keys.
+    """
+    if not isinstance(device, PhysicalDevice):
+      raise ValueError(
+          "device must be a tf.config.PhysicalDevice, but got: %s" % (device,)
+      )
+    if (
+        self._physical_device_to_index is None
+        or device not in self._physical_device_to_index
+    ):
+      raise ValueError(
+          "The PhysicalDevice must be one obtained from "
+          "calling `tf.config.list_physical_devices`, but got: "
+          "%s" % (device,)
+      )
+    index = self._physical_device_to_index[device]
+    details = pywrap_tfe.TF_GetDeviceDetails(index)
+
+    # Change compute_capability from a string to a tuple
+    if "compute_capability" in details:
+      try:
+        major, minor = details["compute_capability"].split(".")
+        details["compute_capability"] = (int(major), int(minor))
+      except ValueError as exc:
+        raise RuntimeError(
+            "Device returned compute capability an in invalid format: %s"
+            % details["compute_capability"]
+        ) from exc
+    return details
+
+  def _import_config(self):
+    """Import config if passed in during construction.
+
+    If Context was created with a ConfigProto such as when calling
+    tf.compat.v1.enable_eager_execution(), then we need to pull out the
+    various pieces we might be replacing and import then into our internal
+    class representation.
+    """
+    if self._config is None:
+      return
+
+    num_cpus = self._config.device_count.get("CPU", 1)
+    if num_cpus != 1:
+      cpus = [d for d in self._physical_devices if d.device_type == "CPU"]
+      if num_cpus == 0:
+        self.set_visible_devices([], "CPU")
+      elif num_cpus > 1:
+        self.set_logical_device_configuration(
+            cpus[0], [LogicalDeviceConfiguration() for _ in range(num_cpus)]
+        )
+
+    # Parse GPU options
+    gpus = [d for d in self._physical_devices if d.device_type == "GPU"]
+
+    # If there are no GPUs detected, simply ignore all the GPU options passed in
+    # rather than doing any validation checks.
+    if not gpus:
+      return
+
+    gpu_count = self._config.device_count.get("GPU", None)
+
+    visible_gpus = []
+    # TODO(gjn): Handle importing existing virtual GPU configuration
+    visible_indices = self._config.gpu_options.visible_device_list
+    if visible_indices:
+      for index in visible_indices.split(","):
+        if int(index) >= len(gpus):
+          raise ValueError("Invalid visible device index: %s" % index)
+        visible_gpus.append(gpus[int(index)])
+    else:
+      visible_gpus = gpus
+
+    if gpu_count is not None:
+      visible_gpus = visible_gpus[:gpu_count]
+
+    self.set_visible_devices(visible_gpus, "GPU")
+
+  def list_logical_devices(self, device_type=None):
+    """Return logical devices."""
+    self.ensure_initialized()
+    if device_type is None:
+      return list(self._logical_devices)
+
+    return [d for d in self._logical_devices if d.device_type == device_type]
+
+  def get_visible_devices(self, device_type=None):
+    """Get the list of visible devices."""
+    self._initialize_physical_devices()
+
+    if device_type is None:
+      return list(self._visible_device_list)
+
+    return [
+        d for d in self._visible_device_list if d.device_type == device_type
+    ]
+
+  def set_visible_devices(self, devices, device_type=None):
+    """Set the list of visible devices."""
+    self._initialize_physical_devices()
+
+    if not isinstance(devices, list):
+      devices = [devices]
+
+    for d in devices:
+      if d not in self._physical_devices:
+        raise ValueError("Unrecognized device: %s" % repr(d))
+      if device_type is not None and d.device_type != device_type:
+        raise ValueError("Unrecognized device: %s" % repr(d))
+
+    visible_device_list = []
+    if device_type is not None:
+      visible_device_list = [
+          d for d in self._visible_device_list if d.device_type != device_type
+      ]
+
+    visible_device_list += devices
+
+    if self._visible_device_list == visible_device_list:
+      return
+
+    if self._context_handle is not None:
+      raise RuntimeError(
+          "Visible devices cannot be modified after being initialized"
+      )
+
+    self._visible_device_list = visible_device_list
+
+  def get_memory_info(self, dev):
+    """Returns a dict of memory info for the device."""
+    self._initialize_physical_devices()
+    self.ensure_initialized()
+    return pywrap_tfe.TFE_GetMemoryInfo(self._context_handle, dev)
+
+  def reset_memory_stats(self, dev):
+    """Resets the tracked memory stats for the device."""
+    self._initialize_physical_devices()
+    self.ensure_initialized()
+    pywrap_tfe.TFE_ResetMemoryStats(self._context_handle, dev)
+
+  def get_memory_growth(self, dev):
+    """Get if memory growth is enabled for a PhysicalDevice."""
+    self._initialize_physical_devices()
+
+    if dev not in self._physical_devices:
+      raise ValueError("Unrecognized device: %s" % repr(dev))
+
+    return self._memory_growth_map[dev]
+
+  def set_memory_growth(self, dev, enable):
+    """Set if memory growth should be enabled for a PhysicalDevice."""
+    self._initialize_physical_devices()
+
+    if dev not in self._physical_devices:
+      raise ValueError("Unrecognized device: %s" % repr(dev))
+
+    if dev in self._virtual_device_map:
+      raise ValueError(
+          "Cannot set memory growth on device when virtual devices configured"
+      )
+
+    if dev.device_type != "GPU" and dev not in self._pluggable_devices:
+      raise ValueError(
+          "Cannot set memory growth on non-GPU and non-Pluggable devices"
+      )
+
+    if self._memory_growth_map.get(dev) == enable:
+      return
+
+    if self._context_handle is not None:
+      raise RuntimeError(
+          "Physical devices cannot be modified after being initialized"
+      )
+
+    self._memory_growth_map[dev] = enable
+
+  def get_logical_device_configuration(self, dev):
+    """Get the virtual device configuration for a PhysicalDevice."""
+    self._initialize_physical_devices()
+
+    if dev not in self._physical_devices:
+      raise ValueError("Unrecognized device: %s" % repr(dev))
+
+    return self._virtual_device_map.get(dev)
+
+  def set_logical_device_configuration(self, dev, virtual_devices):
+    """Set the virtual device configuration for a PhysicalDevice."""
+    self._initialize_physical_devices()
+
+    if dev not in self._physical_devices:
+      raise ValueError("Unrecognized device: %s" % repr(dev))
+
+    if dev.device_type == "CPU":
+      for vdev in virtual_devices:
+        if vdev.memory_limit is not None:
+          raise ValueError(
+              "Setting memory limit on CPU virtual devices is "
+              "currently not supported"
+          )
+        if vdev.experimental_priority is not None:
+          raise ValueError(
+              "Setting experimental_priority on CPU virtual "
+              " devices is currently not supported"
+          )
+        if vdev.experimental_device_ordinal is not None:
+          raise ValueError(
+              "Setting experimental_device_ordinal on CPU virtual "
+              " devices is currently not supported"
+          )
+    elif dev.device_type == "GPU":
+      for vdev in virtual_devices:
+        if vdev.memory_limit is None:
+          raise ValueError(
+              "Setting memory limit is required for GPU virtual devices"
+          )
+    else:
+      raise ValueError(
+          "Virtual devices are not supported for %s" % dev.device_type
+      )
+
+    if self._virtual_device_map.get(dev) == virtual_devices:
+      return
+
+    if self._context_handle is not None:
+      raise RuntimeError(
+          "Virtual devices cannot be modified after being initialized"
+      )
+
+    self._virtual_device_map[dev] = virtual_devices
+
+  def set_logical_cpu_devices(self, num_cpus, prefix=""):
+    """Set virtual CPU devices in context.
+
+    If virtual CPU devices are already configured at context initialization
+    by tf.config.set_logical_device_configuration(), this method should not be
+    called.
+
+    Args:
+      num_cpus: Number of virtual CPUs.
+      prefix: Device name prefix.
+
+    Raises:
+     RuntimeError: If virtual CPUs are already configured at context
+     initialization.
+    """
+    server_def = self._server_def or self._collective_ops_server_def
+    local_prefix = ["/device"]
+    if server_def is not None:
+      local_prefix.append(
+          "/job:%s/replica:0/task:%d"
+          % (server_def.job_name, server_def.task_index)
+      )
+    logical_local_devices = [
+        d
+        for d in self.list_logical_devices("CPU")
+        if d.name.startswith(tuple(local_prefix))
+    ]
+    self.ensure_initialized()
+    # Error out if there are already multiple logical CPU in the context.
+    if len(logical_local_devices) > 1:
+      raise RuntimeError("Virtual CPUs already set, cannot modify again.")
+
+    pywrap_tfe.TFE_SetLogicalCpuDevices(self._context_handle, num_cpus, prefix)
+    self._initialize_logical_devices()
+
+  def get_compiler_ir(
+      self,
+      device_name,
+      platform_name,
+      function_name,
+      flat_args,
+      captured_inputs,
+      stage="hlo",
+  ):
+    """Get the compiler IR bytes.
+
+    Args:
+      device_name: The name of the device with the form as
+        "/job:localhost/replica:0/task:0/device:CPU:0", "/device:TPU:0" etc.
+        When this is used, actual device is needed for getting the compiler IR.
+      platform_name: The name of the platform, e.g. "TPU". When this is used,
+        first we find a device whose name contains the platform, if it is found
+        we get the compiler IR by device. Otherwise the compiler IR is obtained
+        as if using that device. The former logic of falling back to device is
+        necessary, as there are cases of TF variables that need to access
+        devices, but the upper layer may generally choose platform for getting
+        compiler IR in a device-agnostic way.
+      function_name: The name of the function to get the compiler IR.
+      flat_args: The flat argument inputs.
+      captured_inputs: The inputs that are captured.
+      stage: The exported stage for the given function.
+
+    Returns:
+      The compiler IR bytes.
+    """
+    return pywrap_tfe.TF_GetCompilerIr(
+        self._context_handle,
+        function_name,
+        stage,
+        device_name,
+        flat_args,
+        captured_inputs,
+        platform_name,
+    )
+
+  @deprecated(
+      None, "XLA:CPU and XLA:GPU devices are deprecated", warn_once=True
+  )
+  def enable_xla_devices(self):
+    """Enables XLA:CPU and XLA:GPU devices registration."""
+    pywrap_tfe.TF_EnableXlaDevices()
+
+  @property
+  def enable_mlir_bridge(self):
+    return pywrap_tfe.TF_IsMlirBridgeEnabled()
+
+  @property
+  def enable_mlir_graph_optimization(self):
+    return self._enable_mlir_graph_optimization
+
+  @enable_mlir_bridge.setter
+  def enable_mlir_bridge(self, enabled):
+    pywrap_tfe.TF_EnableMlirBridge(enabled)
+    self._thread_local_data.function_call_options = None
+
+  @enable_mlir_graph_optimization.setter
+  def enable_mlir_graph_optimization(self, enabled):
+    self._enable_mlir_graph_optimization = enabled
+    self._thread_local_data.function_call_options = None
+
+  @property
+  def optimizer_jit(self):
+    level = self.config.graph_options.optimizer_options.global_jit_level
+    return (
+        level == config_pb2.OptimizerOptions.ON_1
+        or level == config_pb2.OptimizerOptions.ON_2
+    )
+
+  @optimizer_jit.setter
+  def optimizer_jit(self, enabled):
+    self._optimizer_jit = enabled
+
+    self._thread_local_data.function_call_options = None
+
+  def get_optimizer_experimental_options(self):
+    """Get experimental options for the optimizer.
+
+    Returns:
+      Dictionary of current option values
+    """
+    rewrite_options = self.config.graph_options.rewrite_options
+    options = {}
+
+    def rewriter_toggle(option):
+      attr = getattr(rewrite_options, option)
+      if attr != 0:
+        options[option] = attr == rewriter_config_pb2.RewriterConfig.ON
+
+    def rewriter_bool(option):
+      options[option] = getattr(rewrite_options, option)
+
+    rewriter_toggle("layout_optimizer")
+    rewriter_toggle("constant_folding")
+    rewriter_toggle("shape_optimization")
+    rewriter_toggle("remapping")
+    rewriter_toggle("arithmetic_optimization")
+    rewriter_toggle("dependency_optimization")
+    rewriter_toggle("loop_optimization")
+    rewriter_toggle("function_optimization")
+    rewriter_toggle("debug_stripper")
+    rewriter_bool("disable_model_pruning")
+    rewriter_toggle("scoped_allocator_optimization")
+    rewriter_toggle("pin_to_host_optimization")
+    rewriter_toggle("implementation_selector")
+    rewriter_toggle("auto_mixed_precision")
+    rewriter_toggle("use_plugin_optimizers")
+    rewriter_bool("disable_meta_optimizer")
+    rewriter_toggle("auto_mixed_precision_onednn_bfloat16")
+    rewriter_toggle("auto_mixed_precision_mkl")
+
+    if rewrite_options.min_graph_nodes != 0:
+      options["min_graph_nodes"] = rewrite_options.min_graph_nodes
+
+    return options
+
+  def set_optimizer_experimental_options(self, options):
+    """Set experimental options for the optimizer.
+
+    Args:
+      options: Dictionary of options to modify
+    """
+    self._optimizer_experimental_options.update(options)
+
+    self._thread_local_data.function_call_options = None
+
+  def _validate_thread_count(self, num_threads):
+    """Validate thread count parameter.
+
+    Args:
+      num_threads: The number of threads to validate.
+
+    Raises:
+      TypeError: If num_threads is not an integer (including booleans).
+      ValueError: If num_threads is not in the range [0, 10000].
+    """
+    if isinstance(num_threads, bool) or not isinstance(
+        num_threads, compat.integral_types
+    ):
+      raise TypeError(
+          f"num_threads must be an integer, got {type(num_threads).__name__}."
+      )
+    if not (0 <= num_threads <= 10000):
+      raise ValueError(
+          f"num_threads must be in the range [0, 10000], got {num_threads}. "
+          "Setting an excessively large number of threads can cause system "
+          "instability or crashes."
+      )
+
+  @property
+  def intra_op_parallelism_threads(self):
+    return self.config.intra_op_parallelism_threads
+
+  @intra_op_parallelism_threads.setter
+  def intra_op_parallelism_threads(self, num_threads):
+    self._validate_thread_count(num_threads)
+    if self._intra_op_parallelism_threads == num_threads:
+      return
+
+    if self._context_handle is not None:
+      raise RuntimeError(
+          "Intra op parallelism cannot be modified after initialization."
+      )
+
+    self._intra_op_parallelism_threads = num_threads
+
+  @property
+  def inter_op_parallelism_threads(self):
+    return self.config.inter_op_parallelism_threads
+
+  @inter_op_parallelism_threads.setter
+  def inter_op_parallelism_threads(self, num_threads):
+    # None keeps the default behavior.
+    if num_threads is not None:
+      self._validate_thread_count(num_threads)
+    if self._inter_op_parallelism_threads == num_threads:
+      return
+
+    if self._context_handle is not None:
+      raise RuntimeError(
+          "Inter op parallelism cannot be modified after initialization."
+      )
+
+    self._inter_op_parallelism_threads = num_threads
+
+  @property
+  def soft_device_placement(self):
+    return self.config.allow_soft_placement
+
+  @soft_device_placement.setter
+  def soft_device_placement(self, enable):
+    if self._context_handle is not None:
+      pywrap_tfe.TFE_ContextSetSoftDevicePlacement(self._handle, enable)
+
+    self._soft_device_placement = enable
+    self._thread_local_data.function_call_options = None
+
+  @property
+  def log_device_placement(self):
+    return self.config.log_device_placement
+
+  @log_device_placement.setter
+  def log_device_placement(self, enable):
+    if self._context_handle is not None:
+      pywrap_tfe.TFE_ContextSetLogDevicePlacement(self._handle, enable)
+
+    self._log_device_placement = enable
+    self._thread_local_data.function_call_options = None
+
+  @property
+  def jit_compile_rewrite(self):
+    return self._jit_compile_rewrite
+
+  @jit_compile_rewrite.setter
+  def jit_compile_rewrite(self, enable):
+    if self._context_handle is not None:
+      pywrap_tfe.TFE_ContextSetJitCompileRewrite(self._handle, enable)
+    self._jit_compile_rewrite = enable
+
+  @property
+  def xla_sharding_for_resource_variables(self):
+    return self._xla_sharding_for_resource_variables
+
+  @xla_sharding_for_resource_variables.setter
+  def xla_sharding_for_resource_variables(self, enable):
+    self._xla_sharding_for_resource_variables = enable
+
+  @property
+  def optionals_in_gradients(self):
+    return self._optionals_in_gradients
+
+  @optionals_in_gradients.setter
+  def optionals_in_gradients(self, enable):
+    self._optionals_in_gradients = enable
+
+  @property
+  def device_policy(self):
+    # Only get the policy from the context if it has already been initialized
+    if self._context_handle is not None:
+      return pywrap_tfe.TFE_ContextGetDevicePlacementPolicy(self._handle)
+
+    return self._device_policy
+
+  @device_policy.setter
+  def device_policy(self, policy):
+    if policy is None:
+      policy = DEVICE_PLACEMENT_SILENT
+
+    if self._device_policy != policy:
+      self._device_policy = policy
+
+      # Only set the policy if the context has already been initialized
+      if self._context_handle is not None:
+        pywrap_tfe.TFE_ContextSetThreadLocalDevicePlacementPolicy(
+            self._handle, self._device_policy
+        )
+
+  @property
+  def use_tfrt(self):
+    return self._use_tfrt
+
+  @use_tfrt.setter
+  def use_tfrt(self, tfrt):
+    """Sets whether to use TFRT."""
+    if not isinstance(tfrt, bool):
+      raise ValueError("Expecting a boolean but got %s" % type(tfrt))
+
+    if self._use_tfrt != tfrt:
+      if self._initialized:
+        raise ValueError("use_tfrt should be set before being initialized.")
+      self._use_tfrt = tfrt
+
+  @property
+  def operation_timeout_in_ms(self):
+    return self.config.operation_timeout_in_ms
+
+  @operation_timeout_in_ms.setter
+  def operation_timeout_in_ms(self, timeout_in_ms):
+    if self._operation_timeout_in_ms == timeout_in_ms:
+      return
+
+    if self._context_handle is not None:
+      raise RuntimeError(
+          "Operation timeout cannot be modified after initialization."
+      )
+
+    self._operation_timeout_in_ms = timeout_in_ms
+
+  def enable_run_metadata(self):
+    """Enables tracing of op execution via RunMetadata.
+
+    To retrieve the accumulated metadata call context.export_run_metadata()
+    and to stop tracing call context.disable_run_metadata().
+    """
+    self.ensure_initialized()
+    pywrap_tfe.TFE_ContextEnableRunMetadata(self._handle)
+
+  def disable_run_metadata(self):
+    """Disables tracing of op execution via RunMetadata."""
+    if not self._context_handle:
+      return
+    pywrap_tfe.TFE_ContextDisableRunMetadata(self._context_handle)
+
+  def enable_graph_collection(self):
+    """Enables graph collection of executed functions.
+
+    To retrieve the accumulated graphs call context.export_run_metadata()
+    and to stop collecting graphs call context.disable_graph_collection().
+    """
+    self.ensure_initialized()
+    pywrap_tfe.TFE_ContextEnableGraphCollection(self._handle)
+
+  def disable_graph_collection(self):
+    """Disables graph collection of executed functions."""
+    if not self._context_handle:
+      return
+    pywrap_tfe.TFE_ContextDisableGraphCollection(self._context_handle)
+
+  def export_run_metadata(self):
+    """Returns a RunMetadata proto with accumulated information.
+
+    The returned protocol buffer contains information since the most recent call
+    to either enable_run_metadata or export_run_metadata.
+
+    Returns:
+      A RunMetadata protocol buffer. Or None if not enabled.
+    """
+    if not self._context_handle:
+      return None
+    with c_api_util.tf_buffer() as buffer_:
+      pywrap_tfe.TFE_ContextExportRunMetadata(self._context_handle, buffer_)
+      proto_data = pywrap_tf_session.TF_GetBuffer(buffer_)
+    run_metadata = config_pb2.RunMetadata()
+    run_metadata.ParseFromString(compat.as_bytes(proto_data))
+    return run_metadata
+
+  def set_server_def_retries(self, retries):
+    """Set the number of retries to use when calling SetServerDef.
+
+    In cases where many servers run in high-preemption environments, jobs could
+    be preempted during startup and initial connection via SetServerDef. Retries
+    allow for more robust connection in these environments.
+
+    Args:
+      retries: int specifying the number of connection retries before failing.
+        Retries follow an exponential backoff waiting period with min value 1ms,
+        max value 10s, and exponent 1.3.
+    """
+    self._set_server_def_retries = retries
+
+  @property
+  def context_switches(self):
+    """Returns a stack of context switches."""
+    return self._context_switches
 
 
-if __name__ == '__main__':
-  ops.enable_eager_execution()
-  test.main()
+class _EagerDeviceContext(object):
+  """Context-manager forcing placement of ops and Tensors on a device."""
+
+  __slots__ = ["_device_name", "_ctx", "_stack"]
+
+  def __init__(self, ctx, device_name):
+    self._device_name = device_name
+    self._ctx = ctx
+    self._stack = []
+
+  # TODO(b/189233748): Consolidate the device string parsing logic with
+  # tensorflow/core/util/device_name_utils.cc.
+  def __enter__(self):
+    ctx = self._ctx
+    old_device_name = ctx.device_name
+    old_device_spec = ctx.device_spec
+    new_device_name = self._device_name
+    cache_key = (old_device_name, new_device_name)
+    try:
+      new_device_name, new_device_spec = _device_parsing_cache[cache_key]
+    except TypeError as exc:
+      # Error while trying to compute the cache key.
+      raise ValueError(
+          "Expecting a string device name. Got %s(%s)"
+          % (type(new_device_name), new_device_name)
+      ) from exc
+    except KeyError as exc:
+      # Handle a cache miss.
+      if new_device_name is not None:
+        if not isinstance(new_device_name, str):
+          raise ValueError(
+              "Expecting a string device name. Got %s(%s)"
+              % (type(new_device_name), new_device_name)
+          ) from exc
+        device_spec = pydev.DeviceSpec.from_string(new_device_name)
+        if old_device_name:
+          new_device_spec = copy.copy(old_device_spec)
+        else:
+          ctx.ensure_initialized()
+          new_device_spec = pydev.DeviceSpec.from_string(
+              ctx._context_devices[0]
+          )  # pylint: disable=protected-access
+        new_device_spec = new_device_spec.make_merged_spec(device_spec)
+      else:
+        new_device_spec = pydev.DeviceSpec.from_string("")
+      new_device_name = new_device_spec.to_string()
+      _device_parsing_cache[cache_key] = (new_device_name, new_device_spec)
+
+    ctx._set_device(new_device_name, new_device_spec)  # pylint: disable=protected-access
+    self._stack.append((old_device_name, old_device_spec, new_device_spec))
+
+  def __exit__(self, *ex_info):
+    ctx = self._ctx
+    old_device_name, old_device_spec, new_device_spec = self._stack[-1]
+    if ctx.device_spec is not new_device_spec:
+      raise RuntimeError("Exiting device scope without proper scope nesting")
+    del self._stack[-1]
+    ctx._set_device(old_device_name, old_device_spec)  # pylint: disable=protected-access
+
+
+# Do not change directly.
+_context = None
+_context_lock = threading.Lock()
+
+
+def _set_context_locked(ctx):
+  global _context
+  pywrap_tfe.TFE_Py_SetEagerContext(ctx)
+  ctx.mark_as_global_context()
+  _context = ctx
+
+
+def _set_context(ctx):
+  with _context_lock:
+    _set_context_locked(ctx)
+
+
+def _create_context():
+  with _context_lock:
+    if _context is None:
+      ctx = Context()
+      _set_context_locked(ctx)
+
+
+def _reset_context():
+  """Clears and re-initializes the singleton context.
+
+  Should only be used for testing.
+  """
+  global _context
+  global _device_parsing_cache
+
+  # Garbage collect and clear scalar cache to avoid Tensor from current context
+  # polluting next context.
+  gc.collect()
+  pywrap_tfe.TFE_ClearScalarCache()
+  with _context_lock:
+    if _context is not None:
+      _context._clear_caches()
+      _context = None
+  _create_context()
+  _device_parsing_cache = {}
+
+
+def _reset_jit_compiler_flags():
+  """Clears and re-initializes the TF JIT compiler flags.
+
+  Should only be used for testing.
+  """
+  pywrap_tfe.TF_ResetJitCompilerFlags()
+
+
+def context():
+  """Returns a singleton context object."""
+  if _context is None:
+    _create_context()
+  return _context
+
+
+def context_safe():
+  """Returns current context (or None if one hasn't been initialized)."""
+  return _context
+
+
+def ensure_initialized():
+  """Initialize the context."""
+  context().ensure_initialized()
+
+
+def initialize_logical_devices():
+  """Initialize the virtual devices."""
+  context()._initialize_logical_devices()  # pylint: disable=protected-access
+
+
+def set_global_seed(seed):
+  """Sets the eager mode seed."""
+  context()._set_global_seed(seed)  # pylint: disable=protected-access
+
+
+def global_seed():
+  """Returns the eager mode seed."""
+  return context()._seed  # pylint: disable=protected-access
+
+
+def internal_operation_seed():
+  """Returns the operation seed generated based on global seed."""
+  return context()._internal_operation_seed()  # pylint: disable=protected-access
+
+
+@tf_export("executing_eagerly", v1=[])
+def executing_eagerly():
+  """Checks whether the current thread has eager execution enabled.
+
+  Eager execution is enabled by default and this API returns `True`
+  in most of cases. However, this API might return `False` in the following use
+  cases.
+
+  *  Executing inside `tf.function`, unless under `tf.init_scope` or
+     `tf.config.run_functions_eagerly(True)` is previously called.
+  *  Executing inside a transformation function for `tf.dataset`.
+  *  `tf.compat.v1.disable_eager_execution()` is called.
+
+  General case:
+
+  >>> print(tf.executing_eagerly())
+  True
+
+  Inside `tf.function`:
+
+  >>> @tf.function
+  ... def fn():
+  ...   with tf.init_scope():
+  ...     print(tf.executing_eagerly())
+  ...   print(tf.executing_eagerly())
+  >>> fn()
+  True
+  False
+
+  Inside `tf.function` after `tf.config.run_functions_eagerly(True)` is called:
+
+  >>> tf.config.run_functions_eagerly(True)
+  >>> @tf.function
+  ... def fn():
+  ...   with tf.init_scope():
+  ...     print(tf.executing_eagerly())
+  ...   print(tf.executing_eagerly())
+  >>> fn()
+  True
+  True
+  >>> tf.config.run_functions_eagerly(False)
+
+  Inside a transformation function for `tf.dataset`:
+
+  >>> def data_fn(x):
+  ...   print(tf.executing_eagerly())
+  ...   return x
+  >>> dataset = tf.data.Dataset.range(100)
+  >>> dataset = dataset.map(data_fn)
+  False
+
+  Returns:
+    `True` if the current thread has eager execution enabled.
+  """
+  ctx = context_safe()
+  if ctx is None:
+    return default_execution_mode == EAGER_MODE
+
+  return ctx.executing_eagerly()
+
+
+@tf_export(v1=["executing_eagerly"])
+def executing_eagerly_v1():
+  """Checks whether the current thread has eager execution enabled.
+
+  Eager execution is typically enabled via
+  `tf.compat.v1.enable_eager_execution`, but may also be enabled within the
+  context of a Python function via tf.contrib.eager.py_func.
+
+  When eager execution is enabled, returns `True` in most cases. However,
+  this API might return `False` in the following use cases.
+
+  *  Executing inside `tf.function`, unless under `tf.init_scope` or
+     `tf.config.run_functions_eagerly(True)` is previously called.
+  *  Executing inside a transformation function for `tf.dataset`.
+  *  `tf.compat.v1.disable_eager_execution()` is called.
+
+  >>> tf.compat.v1.enable_eager_execution()
+
+  General case:
+
+  >>> print(tf.executing_eagerly())
+  True
+
+  Inside `tf.function`:
+
+  >>> @tf.function
+  ... def fn():
+  ...   with tf.init_scope():
+  ...     print(tf.executing_eagerly())
+  ...   print(tf.executing_eagerly())
+  >>> fn()
+  True
+  False
+
+  Inside `tf.function`
+  after  `tf.config.run_functions_eagerly(True)` is called:
+
+  >>> tf.config.run_functions_eagerly(True)
+  >>> @tf.function
+  ... def fn():
+  ...   with tf.init_scope():
+  ...     print(tf.executing_eagerly())
+  ...   print(tf.executing_eagerly())
+  >>> fn()
+  True
+  True
+  >>> tf.config.run_functions_eagerly(False)
+
+  Inside a transformation function for `tf.dataset`:
+
+  >>> def data_fn(x):
+  ...   print(tf.executing_eagerly())
+  ...   return x
+  >>> dataset = tf.data.Dataset.range(100)
+  >>> dataset = dataset.map(data_fn)
+  False
+
+  Returns:
+    `True` if the current thread has eager execution enabled.
+  """
+  return executing_eagerly()
+
+
+def in_eager_mode():
+  """Use executing_eagerly() instead. This function will be removed."""
+  return executing_eagerly()
+
+
+def anonymous_name():
+  """Returns the anonymous shared name.
+
+  In eager mode we create anonymous resources to avoid spurious sharing issues.
+  The runtime generates a unique name on our behalf when the reserved
+  anonymous shared name is used as a shared name.
+
+  Returns:
+    The anonymous shared name.
+  """
+
+  # The magic value is defined as
+  # `tensorflow::ResourceHandle::ANONYMOUS_NAME` in C++.
+  return "cd2c89b7-88b7-44c8-ad83-06c2a9158347"
+
+
+def graph_mode():
+  """Context-manager to disable eager execution for the current thread."""
+  return context()._mode(GRAPH_MODE)  # pylint: disable=protected-access
+
+
+# Used by b/167638505 for keras backend API and Lambda layer.
+@tf_export("__internal__.eager_context.eager_mode", v1=[])
+def eager_mode():
+  """Context-manager to enable eager execution for the current thread."""
+  return context()._mode(EAGER_MODE)  # pylint: disable=protected-access
+
+
+def scope_name():
+  """Name of the current scope."""
+  return context().scope_name
+
+
+def device(name):
+  """Context-manager to force placement of operations and Tensors on a device.
+
+  Example:
+  ```python
+  with tf.device('gpu:0'):
+    with tf.device('cpu:0'):
+      shape = tf.constant([], dtype=tf.int32)
+    x = tf.random.truncated_normal(shape, tf.float32)
+  ```
+  will ensure that the `shape` Tensor is on CPU but the `truncated_normal`
+  operation runs on GPU 0.
+
+  Args:
+    name: Name of the device (see context().devices()), or None to perform
+      automatic placement.
+
+  Returns:
+    Context manager for setting the device.
+  """
+  ensure_initialized()
+  return context().device(name)
+
+
+# Expose some properties of Context as internally public APIs (b/160348781).
+@tf_export("__internal__.eager_context.get_config", v1=[])
+def get_config():
+  """Get the ConfigProto of Context.
+
+  Returns:
+    The ConfigProto of Context.
+  """
+  return context().config
+
+
+@tf_export("__internal__.eager_context.get_device_name", v1=[])
+def get_device_name():
+  """Get the device name for the current thread.
+
+  Returns:
+    The device name for the current thread.
+  """
+  return context().device_name
+
+
+@tf_export("__internal__.eager_context.set_soft_device_placement", v1=[])
+def set_soft_device_placement(enabled):
+  """Set if soft device placements should be allowed.
+
+  Args:
+    enabled: Whether to enable soft device placement.
+  """
+  context().soft_device_placement = enabled
+
+
+@tf_export("__internal__.eager_context.get_executor", v1=[])
+def get_executor():
+  """Get the Executor of the current thread.
+
+  Returns:
+    The Executor of the current thread.
+  """
+  return context().executor
+
+
+@tf_export("debugging.get_log_device_placement")
+def get_log_device_placement():
+  """Get if device placements are logged.
+
+  Returns:
+    If device placements are logged.
+  """
+  return context().log_device_placement
+
+
+@tf_export("debugging.set_log_device_placement")
+def set_log_device_placement(enabled):
+  """Turns logging for device placement decisions on or off.
+
+  Operations execute on a particular device, producing and consuming tensors on
+  that device. This may change the performance of the operation or require
+  TensorFlow to copy data to or from an accelerator, so knowing where operations
+  execute is useful for debugging performance issues.
+
+  For more advanced profiling, use the [TensorFlow
+  profiler](https://www.tensorflow.org/guide/profiler).
+
+  Device placement for operations is typically controlled by a `tf.device`
+  scope, but there are exceptions, for example operations on a `tf.Variable`
+  which follow the initial placement of the variable. Turning off soft device
+  placement (with `tf.config.set_soft_device_placement`) provides more explicit
+  control.
+
+  >>> tf.debugging.set_log_device_placement(True)
+  >>> tf.ones([])
+  >>> # [...] op Fill in device /job:localhost/replica:0/task:0/device:GPU:0
+  >>> with tf.device("CPU"):
+  ...  tf.ones([])
+  >>> # [...] op Fill in device /job:localhost/replica:0/task:0/device:CPU:0
+  >>> tf.debugging.set_log_device_placement(False)
+
+  Turning on `tf.debugging.set_log_device_placement` also logs the placement of
+  ops inside `tf.function` when the function is called.
+
+  Args:
+    enabled: Whether to enabled device placement logging.
+  """
+  context().log_device_placement = enabled
+
+
+@tf_contextlib.contextmanager
+def device_policy(policy):
+  """Context manager for setting device placement policy for current thread."""
+  ctx = context()
+  old_policy = ctx.device_policy
+  try:
+    ctx.device_policy = policy
+    yield
+  finally:
+    ctx.device_policy = old_policy
+
+
+def set_execution_mode(mode):
+  """Sets execution mode for the current thread."""
+  context().execution_mode = mode
+
+
+# TODO(fishx): remove this method.
+@tf_contextlib.contextmanager
+def execution_mode(mode):
+  """Context manager for setting execution mode for current thread."""
+  if mode is None:
+    yield
+  else:
+    ctx = context()
+    executor_new = executor.new_executor(mode == ASYNC)
+    executor_old = ctx.executor
+    try:
+      executor_old.wait()
+      ctx.executor = executor_new
+      yield
+    finally:
+      ctx.executor = executor_old
+      executor_new.wait()
+
+
+@tf_contextlib.contextmanager
+def executor_scope(e):
+  """Context manager for changing executor for current thread.
+
+  Args:
+    e: A Executor to execute eager ops under this scope. Setting it to None will
+      switch back to use the default executor for the context.
+
+  Yields:
+    Context manager for setting the executor for current thread.
+  """
+  ctx = context()
+  executor_old = ctx.executor
+  try:
+    ctx.executor = e
+    yield
+  finally:
+    ctx.executor = executor_old
+
+
+@tf_export("experimental.function_executor_type")
+@tf_contextlib.contextmanager
+def function_executor_type(executor_type):
+  """Context manager for setting the executor of eager defined functions.
+
+  Eager defined functions are functions decorated by tf.contrib.eager.defun.
+
+  Args:
+    executor_type: a string for the name of the executor to be used to execute
+      functions defined by tf.contrib.eager.defun.
+
+  Yields:
+    Context manager for setting the executor of eager defined functions.
+  """
+  current_options = context().function_call_options
+  old_options = copy.copy(current_options)
+  try:
+    current_options.executor_type = executor_type
+    yield
+  finally:
+    context().function_call_options = old_options
+
+
+def is_async():
+  """Returns true if current thread is in async mode."""
+  return context().is_async()
+
+
+def num_gpus():
+  """Get the number of available GPU devices.
+
+  Returns:
+    The number of available GPU devices.
+  """
+  return context().num_gpus()
+
+
+def enable_run_metadata():
+  """Enables tracing of op execution via RunMetadata.
+
+  To retrieve the accumulated metadata call context.export_run_metadata()
+  and to stop tracing call context.disable_run_metadata().
+  """
+  context().enable_run_metadata()
+
+
+def disable_run_metadata():
+  """Disables tracing of op execution via RunMetadata."""
+  context().disable_run_metadata()
+
+
+def enable_graph_collection():
+  """Enables graph collection of executed functions.
+
+  To retrieve the accumulated graphs call context.export_run_metadata()
+  and to stop collecting graphs call context.disable_graph_collection().
+  """
+  context().enable_graph_collection()
+
+
+def disable_graph_collection():
+  """Disables graph collection of executed functions."""
+  context().disable_graph_collection()
+
+
+def export_run_metadata():
+  """Returns a RunMetadata proto with accumulated information.
+
+  The returned protocol buffer contains information since the most recent call
+  to either enable_run_metadata or export_run_metadata.
+
+  Returns:
+    A RunMetadata protocol buffer.
+  """
+  return context().export_run_metadata()
+
+
+@contextlib.contextmanager
+def collect_graphs(optimized=True):
+  """Collects a flat list of pre- or post-optimization graphs.
+
+  The collected graphs include device placements, which can be useful for
+  testing.
+
+  Usage:
+
+  ```
+  @def_function.function
+  def f(x):
+    return x + constant_op.constant(1.)
+
+  with context.collect_graphs() as graphs:
+    with ops.device("CPU:0"):
+      f(constant_op.constant(1.))
+
+  graph, = graphs  # `graph` contains a single GraphDef for inspection
+  ```
+
+  Args:
+    optimized: whether to collect optimized graphs or non-optimized graphs
+
+  Yields:
+    A list of GraphDefs, populated when the context manager exits.
+  """
+  ctx = context()
+  ctx.enable_graph_collection()
+  try:
+    graphs = []
+    yield graphs
+    metadata = ctx.export_run_metadata()
+  finally:
+    ctx.disable_graph_collection()
+  for graph in metadata.function_graphs:
+    if optimized:
+      graphs.append(graph.post_optimization_graph)
+    else:
+      graphs.append(graph.pre_optimization_graph)
+
+
+def get_server_def():
+  return context().get_server_def()
+
+
+def set_server_def(server_def):
+  context().set_server_def(server_def)
+
+
+def set_server_def_retries(retries):
+  """Set the number of retries to use when calling SetServerDef.
+
+  In cases where many servers run in high-preemption environments, jobs could
+  be preempted during startup and initial connection via SetServerDef. Retries
+  allow for more robust connection in these environments.
+
+
+  Args:
+    retries: int specifying the number of connection retries before failing.
+      Retries follow an exponential backoff waiting period with min value 1ms,
+      max value 10s, and exponent 1.3.
+  """
+  context().set_server_def_retries(retries)
+
+
+def update_server_def(server_def):
+  context().update_server_def(server_def)
+
+
+def check_alive(worker_name):
+  return context().check_alive(worker_name)
+
+
+@tf_export("experimental.async_scope")
+@tf_contextlib.contextmanager
+def async_scope():
+  """Context manager for grouping async operations.
+
+  Ops/function calls inside the scope can return before finishing the actual
+  execution. When exiting the async scope, a synchronization barrier will be
+  automatically added to ensure the completion of all async op and function
+  execution, potentially raising exceptions if async execution results in
+  an error state.
+
+  Users may write the following code to asynchronously invoke `train_step_fn`
+  and log the `loss` metric for every `num_steps` steps in a training loop.
+  `train_step_fn` internally consumes data using `iterator.get_next()`, and may
+  throw OutOfRangeError when running out of data. In the case:
+
+  ```
+  try:
+    with tf.experimental.async_scope():
+      for _ in range(num_steps):
+        # Step function updates the metric `loss` internally
+        train_step_fn()
+  except tf.errors.OutOfRangeError:
+    tf.experimental.async_clear_error()
+  logging.info('loss = %s', loss.numpy())
+  ```
+
+  Yields:
+    Context manager for grouping async operations.
+  """
+  # TODO(haoyuzhang): replace env var once we have a config method to turn on
+  # and off async streaming RPC
+  remote_async_env_var = "TF_ENABLE_EAGER_CLIENT_STREAMING_ENQUEUE"
+  old_policy = os.environ.get(remote_async_env_var)
+  try:
+    os.environ[remote_async_env_var] = str(True)
+    yield
+    # Note: sync local and remote executors iff the async block does not raise
+    # an exception. Triggering sync after an exception may lead to derived
+    # runtime errors and unexpected exception types.
+    context().sync_executors()
+  finally:
+    if old_policy is None:
+      del os.environ[remote_async_env_var]
+    else:
+      os.environ[remote_async_env_var] = old_policy
+
+
+def async_wait():
+  """Sync all async operations and raise any errors during execution.
+
+  In async execution mode, an op/function call can return before finishing the
+  actual execution. Calling this method creates a synchronization barrier for
+  all async op and function execution. It only returns when all pending nodes
+  are finished, potentially raising exceptions if async execution results in
+  an error state. It is a no-op if the context is not initialized.
+  """
+  disable_async_executor_env_var = "TF_PS_DISABLE_ASYNC_EXECUTOR_GLOBALLY"
+  if os.environ.get(disable_async_executor_env_var) == str(True):
+    return
+  if context()._context_handle is not None:  # pylint: disable=protected-access
+    context().sync_executors()
+
+
+@tf_export("experimental.async_clear_error")
+def async_clear_error():
+  """Clear pending operations and error statuses in async execution.
+
+  In async execution mode, an error in op/function execution can lead to errors
+  in subsequent ops/functions that are scheduled but not yet executed. Calling
+  this method clears all pending operations and reset the async execution state.
+
+  Example:
+
+  ```
+  while True:
+    try:
+      # Step function updates the metric `loss` internally
+      train_step_fn()
+    except tf.errors.OutOfRangeError:
+      tf.experimental.async_clear_error()
+      break
+  logging.info('loss = %s', loss.numpy())
+  ```
+  """
+  context().clear_executor_errors()
+
+
+def add_c_function(c_func):
+  """Add a C API TF_Function to the context."""
+  context().add_c_function(c_func)
+
+
+def get_c_function(name):
+  """Get a C API TF_Function from the context."""
+  return context().get_c_function(name)
+
+
+def remove_function(name):
+  """Remove a function from the context."""
+  context().remove_function(name)
+
+
+def get_function_def(name):
+  return context().get_function_def(name)
+
+
+def is_custom_device(device_name):
+  """Calls TFE_IsCustomDevice.
+
+  Enables using C extensions specifying a custom device from Python. See the
+  experimental eager C API in tensorflow/c/eager/c_api_experimental.h for
+  details.
+
+  Args:
+    device_name: A string indicating the name to check whether it is a
+      registered custom device.
+
+  Returns:
+    A boolean.
+  """
+  return context().is_custom_device(device_name)
+
+
+def register_custom_device(device_capsule, device_name, device_info_capsule):
+  """Calls TFE_RegisterCustomDevice to register a custom device with Python.
+
+  Enables using C extensions specifying a custom device from Python. See the
+  experimental eager C API in tensorflow/c/eager/c_api_experimental.h for
+  details.
+
+  Note that custom devices are not currently supported inside `tf.function`s.
+
+  Args:
+    device_capsule: A PyCapsule with the name set to 'TFE_CustomDevice'
+      containing a pointer to a TFE_CustomDevice struct. The capsule retains
+      ownership of the memory.
+    device_name: A string indicating the name to register the custom device
+      under, e.g. '/job:localhost/replica:0/task:0/device:CUSTOM:0'. It may
+      subsequently be passed to `with tf.device(...):`.
+    device_info_capsule: A PyCapsule with the name set to
+      'TFE_CustomDevice_DeviceInfo' containing a pointer to a device-specific
+      struct with the initial state of the custom device (the void* device_info
+      argument to TFE_RegisterCustomDevice). This method takes ownership of the
+      memory and clears the capsule destructor.
+  """
+  context().register_custom_device(
+      device_capsule, device_name, device_info_capsule
+  )
+
+
+# Not every user creates a Context via context.context()
+# (for example, enable_eager_execution in python/framework/ops.py),
+# but they do all import this file.  Note that IS_IN_GRAPH_MODE and
+# in_graph_mode are both parameterless functions.
+def _tmp_in_graph_mode():
+  if context_safe() is None:
+    # Context not yet initialized. Assume graph mode following the
+    # default implementation in `is_in_graph_mode`.
+    return True
+  return not executing_eagerly()
+
+
+is_in_graph_mode.IS_IN_GRAPH_MODE = _tmp_in_graph_mode

From d55e35303dde45c4486479d29fb0abefc2b57d70 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Sat, 10 Oct 2026 09:32:21 +0530
Subject: [PATCH 0412/1000] Update context.py


From 1460ed5800de272491261cea082a6fb7051c34ca Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Sat, 10 Oct 2026 09:34:11 +0530
Subject: [PATCH 0413/1000] Add tests for coordination service methods

Added tests for coordination service methods to ensure they initialize a non-global context and handle errors correctly when the coordination service is not enabled.
---
 tensorflow/python/eager/context_test.py | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/tensorflow/python/eager/context_test.py b/tensorflow/python/eager/context_test.py
index 9d93e3650b6..73be78de959 100644
--- a/tensorflow/python/eager/context_test.py
+++ b/tensorflow/python/eager/context_test.py
@@ -165,6 +165,22 @@ class ContextTest(test.TestCase, parameterized.TestCase):
     self.assertIn(concrete.name.decode(),
                   context.context().list_function_names())
 
+  def testCoordinationServiceMethodsInitializeNonGlobalContext(self):
+    # Each method must initialize its own (non-global) context before calling
+    # the C API; without a coordination service it then fails cleanly.
+    with self.assertRaisesRegex(errors.FailedPreconditionError,
+                                'Coordination service agent is not enabled'):
+      context.Context().set_config_key_value('key', 'value')
+    with self.assertRaisesRegex(errors.FailedPreconditionError,
+                                'Coordination service is not enabled'):
+      context.Context().get_config_key_value('key', timeout_in_ms=1)
+    with self.assertRaisesRegex(errors.FailedPreconditionError,
+                                'Coordination service is not enabled'):
+      context.Context().delete_config_key_value('key')
+    with self.assertRaisesRegex(errors.FailedPreconditionError,
+                                'Coordination service is not enabled'):
+      context.Context().wait_at_barrier('barrier', 1)
+
   def testSetLogicalDeviceAfterContextInitialization(self):
     ctx = context.Context()
     ctx.set_logical_cpu_devices(4)

From 2aca32c564afbbe17394a992355a65e086bbcaad Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 9 Oct 2026 21:13:37 -0700
Subject: [PATCH 0414/1000] Automated Code Change

PiperOrigin-RevId: 996929282
---
 tensorflow/lite/delegates/delegate_test.cc | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensorflow/lite/delegates/delegate_test.cc b/tensorflow/lite/delegates/delegate_test.cc
index a888a8b9e94..87156fb2d83 100644
--- a/tensorflow/lite/delegates/delegate_test.cc
+++ b/tensorflow/lite/delegates/delegate_test.cc
@@ -1445,9 +1445,9 @@ TEST_F(TestDelegateWithControlEdges, NoControlEdges) {
   delegate_ = std::make_unique<SimpleDelegate>(std::vector<int>({0, 2}));
   interpreter_->ModifyGraphWithDelegate(delegate_->get_tf_lite_delegate());
   ASSERT_EQ(interpreter_->execution_plan().size(), 3);     // [ {0, 2}, 1, 3]
-  EXPECT_EQ(interpreter_->execution_plan().data()[0], 4);  // new super-node
-  EXPECT_EQ(interpreter_->execution_plan().data()[1], 1);  // undelegated
-  EXPECT_EQ(interpreter_->execution_plan().data()[2], 3);  // undelegated
+  EXPECT_EQ(interpreter_->execution_plan()[0], 4);         // new super-node
+  EXPECT_EQ(interpreter_->execution_plan()[1], 1);         // undelegated
+  EXPECT_EQ(interpreter_->execution_plan()[2], 3);         // undelegated
 }
 
 TEST_F(TestDelegateWithControlEdges, OverrideControlEdges) {
@@ -1462,10 +1462,10 @@ TEST_F(TestDelegateWithControlEdges, OverrideControlEdges) {
   // preserved. Nodes 0 and 2 both get rewritten into new delegate nodes
   // 4 and 5.
   ASSERT_EQ(interpreter_->execution_plan().size(), 4);  // [ 0, 1, 2, 3]
-  EXPECT_EQ(interpreter_->execution_plan().data()[0], 4);
-  EXPECT_EQ(interpreter_->execution_plan().data()[1], 1);
-  EXPECT_EQ(interpreter_->execution_plan().data()[2], 5);
-  EXPECT_EQ(interpreter_->execution_plan().data()[3], 3);
+  EXPECT_EQ(interpreter_->execution_plan()[0], 4);
+  EXPECT_EQ(interpreter_->execution_plan()[1], 1);
+  EXPECT_EQ(interpreter_->execution_plan()[2], 5);
+  EXPECT_EQ(interpreter_->execution_plan()[3], 3);
 }
 
 // Test that empty control edge metadata for subgraph 0 don't change anything.

From a96e7ef5ff36a6626f488209a83d5a99a154ae74 Mon Sep 17 00:00:00 2001
From: RadikHoroshev <rodion111@gmail.com>
Date: Sat, 10 Oct 2026 07:29:31 +0300
Subject: [PATCH 0415/1000] Address TensorList review test feedback

---
 .../data_structures/list_ops_test.py            | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/kernel_tests/data_structures/list_ops_test.py b/tensorflow/python/kernel_tests/data_structures/list_ops_test.py
index 23047b5ab4b..87a668a460c 100644
--- a/tensorflow/python/kernel_tests/data_structures/list_ops_test.py
+++ b/tensorflow/python/kernel_tests/data_structures/list_ops_test.py
@@ -1724,10 +1724,10 @@ class ListOpsTest(test_util.TensorFlowTestCase, parameterized.TestCase):
       self.evaluate(l)
 
   @test_util.run_in_graph_and_eager_modes
-  def testReserveAtMaxElementsSucceeds(self):
+  def testReserveAtModerateElementsSucceeds(self):
     l = list_ops.tensor_list_reserve(
-        element_shape=[], num_elements=1 << 24, element_dtype=dtypes.float32)
-    self.assertEqual(self.evaluate(list_ops.tensor_list_length(l)), 1 << 24)
+        element_shape=[], num_elements=100000, element_dtype=dtypes.float32)
+    self.assertEqual(self.evaluate(list_ops.tensor_list_length(l)), 100000)
 
   @test_util.run_in_graph_and_eager_modes
   def testResizeExceedingMaxElementsFails(self):
@@ -1737,6 +1737,17 @@ class ListOpsTest(test_util.TensorFlowTestCase, parameterized.TestCase):
       l = list_ops.tensor_list_resize(l, (1 << 24) + 1)
       self.evaluate(l)
 
+  @test_util.run_in_graph_and_eager_modes
+  def testResizeShrinkingSucceeds(self):
+    l = list_ops.tensor_list_from_tensor(
+        [1., 2., 3., 4., 5.], element_shape=[])
+    l = list_ops.tensor_list_resize(l, 2)
+    self.assertEqual(self.evaluate(list_ops.tensor_list_length(l)), 2)
+    self.assertAllEqual(
+        self.evaluate(
+            list_ops.tensor_list_stack(l, element_dtype=dtypes.float32)),
+        [1., 2.])
+
   @test_util.run_deprecated_v1
   @test_util.enable_control_flow_v2
   def testSkipEagerResizeGrad(self):

From 5a4440729d80da4f721e40257042980f19b8c189 Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Fri, 9 Oct 2026 21:35:23 -0700
Subject: [PATCH 0416/1000] Convert NMS thresholds only once there are boxes to
 suppress

With no boxes, ComputeResult returns before it uses the thresholds, so
convert them after that return instead of building conversions that
nothing uses.

Also test that a score_threshold whose shape is only known when the op
compiles is rejected if it isn't a scalar.
---
 tensorflow/compiler/tests/image_ops_test.py   | 23 +++++++++++++++++++
 .../compiler/tf2xla/kernels/image_ops.cc      | 23 ++++++++++---------
 2 files changed, 35 insertions(+), 11 deletions(-)

diff --git a/tensorflow/compiler/tests/image_ops_test.py b/tensorflow/compiler/tests/image_ops_test.py
index a41796b7165..49fbdce8c34 100644
--- a/tensorflow/compiler/tests/image_ops_test.py
+++ b/tensorflow/compiler/tests/image_ops_test.py
@@ -1082,6 +1082,29 @@ class NonMaxSuppressionTest(xla_test.XLATestCase):
 
     self._testMixedThresholdDtype(nms)
 
+  def testNMSV3NonScalarScoreThreshold(self):
+    # Shape inference can't reject a score_threshold of unknown shape, so the
+    # XLA kernel checks that it is a scalar when it compiles.
+    with self.session() as sess:
+      boxes = array_ops.placeholder(np.float32, shape=[6, 4])
+      scores = array_ops.placeholder(np.float32, shape=[6])
+      score_threshold = array_ops.placeholder(np.float32)
+      with self.device_scope():
+        selected_indices = image_ops.non_max_suppression_v3(
+            boxes=boxes,
+            scores=scores,
+            max_output_size=6,
+            iou_threshold=0.5,
+            score_threshold=score_threshold)
+      with self.assertRaisesOpError("Score Threshold isn't a scalar"):
+        sess.run(
+            selected_indices,
+            feed_dict={
+                boxes: np.zeros([6, 4], np.float32),
+                scores: np.zeros([6], np.float32),
+                score_threshold: [0.4]
+            })
+
   def testNMSV3EmptyInput(self):
     # Regression test for #117245: with no boxes the suppression loop was
     # built from zero-sized dimensions and segfaulted the compiler. The
diff --git a/tensorflow/compiler/tf2xla/kernels/image_ops.cc b/tensorflow/compiler/tf2xla/kernels/image_ops.cc
index 3ef8757e4e5..fa78850785d 100644
--- a/tensorflow/compiler/tf2xla/kernels/image_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/image_ops.cc
@@ -495,17 +495,6 @@ class NonMaxSuppressionOp : public XlaOpKernel {
     OP_REQUIRES(context, output_size <= std::numeric_limits<int32_t>::max(),
                 absl::InvalidArgumentError(absl::StrCat(
                     "Need output_size <= kint32Max, got ", output_size)));
-    // The thresholds' type, T_threshold, can differ from T, the type of boxes
-    // and scores. Compare in T, as the TensorFlow kernels do, since XLA
-    // doesn't mix floating-point types in one operation.
-    xla::XlaOp score_thresh = context->Input("score_threshold");
-    if (context->InputXlaType("score_threshold") != scores_xla_type) {
-      score_thresh = xla::ConvertElementType(score_thresh, scores_xla_type);
-    }
-    xla::XlaOp iou_thresh = context->Input("iou_threshold");
-    if (context->InputXlaType("iou_threshold") != boxes_xla_type) {
-      iou_thresh = xla::ConvertElementType(iou_thresh, boxes_xla_type);
-    }
     xla::XlaBuilder* const builder = context->builder();
 
     if (num_boxes == 0) {
@@ -534,6 +523,18 @@ class NonMaxSuppressionOp : public XlaOpKernel {
       return;
     }
 
+    // The thresholds' type, T_threshold, can differ from T, the type of boxes
+    // and scores. Compare in T, as the TensorFlow kernels do, since XLA
+    // doesn't mix floating-point types in one operation.
+    xla::XlaOp score_thresh = context->Input("score_threshold");
+    if (context->InputXlaType("score_threshold") != scores_xla_type) {
+      score_thresh = xla::ConvertElementType(score_thresh, scores_xla_type);
+    }
+    xla::XlaOp iou_thresh = context->Input("iou_threshold");
+    if (context->InputXlaType("iou_threshold") != boxes_xla_type) {
+      iou_thresh = xla::ConvertElementType(iou_thresh, boxes_xla_type);
+    }
+
     // Choose a more convenient layout.
     const xla::XlaOp boxes = xla::Transpose(boxes_input, {1, 0});
     const xla::XlaOp boxes_sorted = xla::GetTupleElement(

From 0cdb7ef288a49ce95fefd01427843cd9fa2964e2 Mon Sep 17 00:00:00 2001
From: Raashish Aggarwal <94279692+raashish1601@users.noreply.github.com>
Date: Sat, 10 Oct 2026 10:49:08 +0530
Subject: [PATCH 0417/1000] Check the error message for out-of-bounds axes in
 var

---
 tensorflow/python/ops/numpy_ops/np_array_ops_test.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index 1ea68abd192..cd5e0901c87 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1005,9 +1005,11 @@ class ArrayMethodsTest(test.TestCase):
     run_test(
         np.arange(8).reshape((2, 2, 2)).tolist(), axis=(2, 0), keepdims=True)
     self.assertRaises(ValueError, np_array_ops.var, np.ones([2, 2]), out=[])
-    with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
+    with self.assertRaisesRegex((ValueError, errors_impl.InvalidArgumentError),
+                                'out of bounds'):
       np_array_ops.var(np.ones([2, 2]), axis=-3, ddof=1)
-    with self.assertRaises((ValueError, errors_impl.InvalidArgumentError)):
+    with self.assertRaisesRegex((ValueError, errors_impl.InvalidArgumentError),
+                                'out of bounds'):
       np_array_ops.var(np.ones([2, 2]), axis=(-3, 0), ddof=1)
 
   def testProd(self):

From a52fbebb396e76b6997f0feb076922a65027b7da Mon Sep 17 00:00:00 2001
From: Raashish Aggarwal <94279692+raashish1601@users.noreply.github.com>
Date: Sat, 10 Oct 2026 11:42:43 +0530
Subject: [PATCH 0418/1000] Normalize static var axes without extra ops and
 test more out-of-bounds axes

---
 .../python/ops/numpy_ops/np_array_ops.py      | 21 +++++++++++++++----
 .../python/ops/numpy_ops/np_array_ops_test.py | 21 +++++++++++++------
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index b6c0b3ceab7..48cfa7c81f6 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -834,10 +834,23 @@ def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=None):  # pylint: d
         n = array_ops.size(input_tensor)
       else:
         # axis can be an int or a sequence of ints.
-        axis = math_ops.cast(axis, dtypes.int32)
-        axis = array_ops.where_v2(
-            axis < 0, axis + array_ops.rank(input_tensor), axis
-        )
+        if isinstance(axis, (int, np.integer)):
+          axis = [axis]
+        rank = input_tensor.shape.rank
+        if (
+            rank is not None
+            and isinstance(axis, (list, tuple))
+            and builtins.all(isinstance(ax, (int, np.integer)) for ax in axis)
+        ):
+          # Known rank and axes: normalize them without adding ops.
+          axis = np.array(
+              [ax + rank if ax < 0 else ax for ax in axis], dtype=np.int32
+          )
+        else:
+          axis = math_ops.cast(axis, dtypes.int32)
+          axis = array_ops.where_v2(
+              axis < 0, axis + array_ops.rank(input_tensor), axis
+          )
         n = math_ops.reduce_prod(
             array_ops.gather(array_ops.shape(input_tensor), axis)
         )
diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index cd5e0901c87..eff25b4bd5f 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1005,12 +1005,21 @@ class ArrayMethodsTest(test.TestCase):
     run_test(
         np.arange(8).reshape((2, 2, 2)).tolist(), axis=(2, 0), keepdims=True)
     self.assertRaises(ValueError, np_array_ops.var, np.ones([2, 2]), out=[])
-    with self.assertRaisesRegex((ValueError, errors_impl.InvalidArgumentError),
-                                'out of bounds'):
-      np_array_ops.var(np.ones([2, 2]), axis=-3, ddof=1)
-    with self.assertRaisesRegex((ValueError, errors_impl.InvalidArgumentError),
-                                'out of bounds'):
-      np_array_ops.var(np.ones([2, 2]), axis=(-3, 0), ddof=1)
+    out_of_bounds = 'out of bounds|not in|Invalid reduction dimension'
+    for axis in (-3, (-3, 0), 3, (0, 3)):
+      with self.assertRaisesRegex(
+          (ValueError, errors_impl.InvalidArgumentError), out_of_bounds):
+        np_array_ops.var(np.ones([2, 2]), axis=axis, ddof=1)
+
+    # With an unknown rank, the axis is only checked at run time.
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
+    def var_fn(a):
+      return np_array_ops.var(a, axis=-3, ddof=1)
+
+    with self.assertRaisesRegex(
+        (ValueError, errors_impl.InvalidArgumentError), out_of_bounds):
+      var_fn(np.ones([2, 2], dtype=np.float32))
 
   def testProd(self):
 

From 0ed87b8a727685e7dd28a74bcf282019c44ae6b3 Mon Sep 17 00:00:00 2001
From: Raashish Aggarwal <94279692+raashish1601@users.noreply.github.com>
Date: Sat, 10 Oct 2026 12:30:47 +0530
Subject: [PATCH 0419/1000] Normalize range and NumPy array var axes statically
 and test a dynamic tuple axis

---
 tensorflow/python/ops/numpy_ops/np_array_ops.py     |  4 +++-
 .../python/ops/numpy_ops/np_array_ops_test.py       | 13 +++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index 48cfa7c81f6..331e15259a9 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -836,10 +836,12 @@ def var(a, axis=None, dtype=None, out=None, ddof=0, keepdims=None):  # pylint: d
         # axis can be an int or a sequence of ints.
         if isinstance(axis, (int, np.integer)):
           axis = [axis]
+        elif isinstance(axis, np.ndarray) and axis.ndim == 0:
+          axis = [axis.item()]
         rank = input_tensor.shape.rank
         if (
             rank is not None
-            and isinstance(axis, (list, tuple))
+            and isinstance(axis, (list, tuple, range, np.ndarray))
             and builtins.all(isinstance(ax, (int, np.integer)) for ax in axis)
         ):
           # Known rank and axes: normalize them without adding ops.
diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index eff25b4bd5f..d0fdeb3a712 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -995,6 +995,11 @@ class ArrayMethodsTest(test.TestCase):
     run_test([[1.0j, 2.0], [3.0j, 4.0]], axis=0, ddof=1)
     run_test(np.arange(8.0).reshape((2, 2, 2)).tolist(), axis=(0, 2), ddof=1)
     run_test(np.arange(8.0).reshape((2, 2, 2)).tolist(), axis=(-1, 0), ddof=1)
+    # A range of axes is static too.
+    self.match(
+        np_array_ops.var(np.arange(8.0).reshape((2, 2, 2)), axis=range(2),
+                         ddof=1),
+        np.var(np.arange(8.0).reshape((2, 2, 2)), axis=(0, 1), ddof=1))
     run_test(np.arange(8.0).reshape((2, 2, 2)).tolist(), axis=(), ddof=1)
     run_test(5.0, ddof=1)
     run_test(5.0, axis=(), ddof=1)
@@ -1017,9 +1022,17 @@ class ArrayMethodsTest(test.TestCase):
     def var_fn(a):
       return np_array_ops.var(a, axis=-3, ddof=1)
 
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
+    def var_fn_tuple(a):
+      return np_array_ops.var(a, axis=(-3, 0), ddof=1)
+
     with self.assertRaisesRegex(
         (ValueError, errors_impl.InvalidArgumentError), out_of_bounds):
       var_fn(np.ones([2, 2], dtype=np.float32))
+    with self.assertRaisesRegex(
+        (ValueError, errors_impl.InvalidArgumentError), out_of_bounds):
+      var_fn_tuple(np.ones([2, 2], dtype=np.float32))
 
   def testProd(self):
 

From d2b963a62437f94ef4170b997f5fbcc10d643a60 Mon Sep 17 00:00:00 2001
From: Ce Zheng <zce@google.com>
Date: Sat, 10 Oct 2026 00:19:12 -0700
Subject: [PATCH 0420/1000] Support small contiguous chunks and sub-byte
 chunking in TransposePlan.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

1. Allow `ChooseParallelizationStrategy` to partition tensors smaller than 64 MiB when `chunk_contiguity != ChunkContiguity::kNone`. Previously, `kMinBytesPerThread` was fixed at 64 MiB (or 1 MiB for memcpy plans), which clamped `Parallelism()` to 1 for sub-64 MiB transfers even when a caller explicitly requested multiple contiguous input/output chunks for pipelined DMA staging.

2. Support `Transformation::kPackSubbyte` with `ExecuteChunk()` and `TransposeConstStride1`:
   - Support `Transformation::kPackSubbyte` directly in `TransposeConstStride1` (`inner_kernel_is_memcpy() == true`) by packing contiguous stride-1 spans via `PackIntN`, scaling outer output strides by `elements_per_byte`, and setting `loop.inc = 8 / bits_per_element_` so `PartitionLoops` only splits at byte boundaries.
   - Skip 1-iteration trivial loops in `IdentifyContiguousLoops()` so size-1 inner loops preserved by `RemoveTrivialLoops()` do not overwrite a real outer loop's `contiguity` to 2.
   - Scale `output_chunk_offset_bytes_`, `output_chunk_iteration_offsets_`, and `output_chunk_size_bytes_` by `elements_per_byte` (8 / `bits_per_element_`) in `ComputeChunkSizes()` when `transformation_ == Transformation::kPackSubbyte`.
   - For fallback sub-byte packing (`use_fallback_pack_`), allocate and zero-initialize `OutputNumElems()` bytes in `Execute()` so partial output tiles are handled safely.
   - Fall back to `use_fallback_pack_` whenever `b_is_tiled_` and `b_tiling_.back() % elements_per_byte != 0` (including `b_tiling_.back() == 1`), or when the stride-1 dimension in `B` or its tile size is not divisible by `elements_per_byte`.

3. Include `chunk_contiguity` in `TransposePlanCacheKey` equality, hashing, and lookup in `TransposePlanCache`.

### Benchmark Highlights (Before vs. After)

- **Significant improvement (`-50.9%` to `-82.9%` time/CPU reduction, `2.0x`–`5.8x` speedup) on `BM_Transpose_Pack4` with stride-1 inner dimension (`perm = {1, 0, 2}`)**:
  - `256×256×256` (`perm_1_0_2`):
    - 1 thread: `2.922 ms` → `1.434 ms` (**-50.90%**, `2.04x` faster)
    - 4 threads: `1.945 ms` → `0.458 ms` wall (`-76.46%`), `1.510 ms` → `0.347 ms` CPU (**-77.02%**, `4.35x` faster)
    - 8 threads: `1.750 ms` → `0.354 ms` wall (`-79.79%`), `1.384 ms` → `0.236 ms` CPU (**-82.92%**, `5.85x` faster)
  - `512×512×512` (`perm_1_0_2`):
    - 1 thread: `21.296 ms` → `8.813 ms` (**-58.62%**, `2.42x` faster)
    - 4 threads: `15.047 ms` → `5.363 ms` wall (`-64.36%`), `14.139 ms` → `3.955 ms` CPU (**-72.03%**, `3.57x` faster)
    - 8 threads: `13.832 ms` → `4.320 ms` wall (`-68.77%`), `13.077 ms` → `3.487 ms` CPU (**-73.34%**, `3.75x` faster)
  - **Explanation**: Previously, when `inner_kernel_is_memcpy_` was true (`pos_stride1b_in_a == pos_stride1a_in_a`) and `transformation_ == Transformation::kPackSubbyte`, `TransposePlan::Initialize()` fell back to `use_fallback_pack_ = true`, which allocated a full-tensor temporary buffer (`16 MiB` for `256³`, `128 MiB` for `512³`), copied unpacked 1-byte elements into `temp_buf`, and then ran a single-threaded `PackIntN` pass over the entire buffer. With `TransposeConstStride1<Transformation::kPackSubbyte>`, contiguous stride-1 rows are packed directly from input to output via `PackIntN` in parallel across worker threads without any intermediate buffer allocation.
- **No regressions**: All other `BM_Transpose_*` permutations (`inner_kernel_is_memcpy_ == false` or `Transformation::kNone`) and all `BM_TransposeMicroKernel` cases execute the exact same kernel paths as before when `chunk_contiguity == ChunkContiguity::kNone` (overall suite geomean `-4.30%` CPU time / `-3.71%` wall time, microkernel geomean `-0.05%`).

<details>

<summary>Expand for full before-vs-after CPU benchmark comparison (benchy)</summary>

```text
name                                                            time/op (ref)   time/op (exp)  vs base
BM_Transpose_uint8_threads_1_256_256_perm_1_0                    5.148µ ± ∞      5.058µ ± ∞         ~ (p=0.310 n=5)
BM_Transpose_uint8_threads_4_256_256_perm_1_0                    5.106µ ± ∞      5.046µ ± ∞         ~ (p=0.548 n=5)
BM_Transpose_uint8_threads_8_256_256_perm_1_0                    5.122µ ± ∞      5.014µ ± ∞         ~ (p=0.310 n=5)
BM_Transpose_float_threads_1_256_256_perm_1_0                    8.506µ ± ∞      8.815µ ± ∞         ~ (p=0.095 n=5)
BM_Transpose_float_threads_4_256_256_perm_1_0                    8.750µ ± ∞      8.494µ ± ∞         ~ (p=0.222 n=5)
BM_Transpose_float_threads_8_256_256_perm_1_0                    8.469µ ± ∞      8.829µ ± ∞         ~ (p=0.222 n=5)
BM_Transpose_Pack4_threads_1_256_256_perm_1_0                    19.54µ ± ∞      19.62µ ± ∞         ~ (p=0.548 n=5)
BM_Transpose_Pack4_threads_4_256_256_perm_1_0                    19.58µ ± ∞      19.83µ ± ∞         ~ (p=0.222 n=5)
BM_Transpose_Pack4_threads_8_256_256_perm_1_0                    19.54µ ± ∞      19.61µ ± ∞         ~ (p=1.000 n=5)
BM_Transpose_uint8_threads_1_512_512_perm_1_0                    27.66µ ± ∞      27.43µ ± ∞         ~ (p=0.841 n=5)
BM_Transpose_uint8_threads_4_512_512_perm_1_0                    28.60µ ± ∞      27.58µ ± ∞         ~ (p=0.421 n=5)
BM_Transpose_uint8_threads_8_512_512_perm_1_0                    27.94µ ± ∞      27.68µ ± ∞         ~ (p=0.841 n=5)
BM_Transpose_float_threads_1_512_512_perm_1_0                    53.23µ ± ∞      56.57µ ± ∞    +6.27% (p=0.032 n=5)
BM_Transpose_float_threads_4_512_512_perm_1_0                    52.39µ ± ∞      52.21µ ± ∞         ~ (p=0.690 n=5)
BM_Transpose_float_threads_8_512_512_perm_1_0                    50.12µ ± ∞      54.65µ ± ∞         ~ (p=0.222 n=5)
BM_Transpose_Pack4_threads_1_512_512_perm_1_0                    81.32µ ± ∞      79.67µ ± ∞         ~ (p=0.548 n=5)
BM_Transpose_Pack4_threads_4_512_512_perm_1_0                    79.72µ ± ∞      79.77µ ± ∞         ~ (p=0.690 n=5)
BM_Transpose_Pack4_threads_8_512_512_perm_1_0                    80.11µ ± ∞      80.25µ ± ∞         ~ (p=1.000 n=5)
BM_Transpose_uint8_threads_1_1024_1024_perm_1_0                  210.0µ ± ∞      205.3µ ± ∞         ~ (p=0.548 n=5)
BM_Transpose_uint8_threads_4_1024_1024_perm_1_0                  210.0µ ± ∞      207.9µ ± ∞         ~ (p=0.690 n=5)
BM_Transpose_uint8_threads_8_1024_1024_perm_1_0                  209.2µ ± ∞      203.3µ ± ∞         ~ (p=0.222 n=5)
BM_Transpose_float_threads_1_1024_1024_perm_1_0                  534.5µ ± ∞      528.4µ ± ∞         ~ (p=0.841 n=5)
BM_Transpose_float_threads_4_1024_1024_perm_1_0                  528.8µ ± ∞      535.9µ ± ∞         ~ (p=0.690 n=5)
BM_Transpose_float_threads_8_1024_1024_perm_1_0                  530.6µ ± ∞      536.4µ ± ∞         ~ (p=0.690 n=5)
BM_Transpose_Pack4_threads_1_1024_1024_perm_1_0                  385.1µ ± ∞      384.5µ ± ∞         ~ (p=0.690 n=5)
BM_Transpose_Pack4_threads_4_1024_1024_perm_1_0                  383.3µ ± ∞      379.9µ ± ∞         ~ (p=0.421 n=5)
BM_Transpose_Pack4_threads_8_1024_1024_perm_1_0                  381.8µ ± ∞      376.2µ ± ∞         ~ (p=0.421 n=5)
BM_Transpose_uint8_threads_1_256_256_256_perm_0_2_1              2.999m ± ∞      3.103m ± ∞         ~ (p=0.548 n=5)
BM_Transpose_uint8_threads_4_256_256_256_perm_0_2_1              3.039m ± ∞      3.005m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_uint8_threads_8_256_256_256_perm_0_2_1              3.259m ± ∞      3.242m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_float_threads_1_256_256_256_perm_0_2_1             10.191m ± ∞      9.928m ± ∞         ~ (p=0.690 n=5)
BM_Transpose_float_threads_4_256_256_256_perm_0_2_1              9.899m ± ∞     10.595m ± ∞         ~ (p=0.095 n=5)
BM_Transpose_float_threads_8_256_256_256_perm_0_2_1              10.39m ± ∞      10.82m ± ∞         ~ (p=0.222 n=5)
BM_Transpose_Pack4_threads_1_256_256_256_perm_0_2_1              6.219m ± ∞      6.201m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_Pack4_threads_4_256_256_256_perm_0_2_1              6.310m ± ∞      6.043m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_Pack4_threads_8_256_256_256_perm_0_2_1              6.413m ± ∞      6.074m ± ∞         ~ (p=0.222 n=5)
BM_Transpose_uint8_threads_1_256_256_256_perm_1_0_2              1.765m ± ∞      1.813m ± ∞         ~ (p=0.548 n=5)
BM_Transpose_uint8_threads_4_256_256_256_perm_1_0_2              512.4µ ± ∞      588.7µ ± ∞         ~ (p=0.222 n=5)
BM_Transpose_uint8_threads_8_256_256_256_perm_1_0_2              394.5µ ± ∞      446.2µ ± ∞         ~ (p=0.151 n=5)
BM_Transpose_float_threads_1_256_256_256_perm_1_0_2              5.620m ± ∞      5.814m ± ∞         ~ (p=0.151 n=5)
BM_Transpose_float_threads_4_256_256_256_perm_1_0_2              3.911m ± ∞      3.158m ± ∞         ~ (p=0.222 n=5)
BM_Transpose_float_threads_8_256_256_256_perm_1_0_2              3.555m ± ∞      3.843m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_Pack4_threads_1_256_256_256_perm_1_0_2              2.922m ± ∞      1.434m ± ∞   -50.90% (p=0.008 n=5)
BM_Transpose_Pack4_threads_4_256_256_256_perm_1_0_2             1944.6µ ± ∞      457.8µ ± ∞   -76.46% (p=0.008 n=5)
BM_Transpose_Pack4_threads_8_256_256_256_perm_1_0_2             1749.7µ ± ∞      353.6µ ± ∞   -79.79% (p=0.008 n=5)
BM_Transpose_uint8_threads_1_256_256_256_perm_1_2_0              35.02m ± ∞      35.16m ± ∞         ~ (p=1.000 n=5)
BM_Transpose_uint8_threads_4_256_256_256_perm_1_2_0              34.80m ± ∞      34.80m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_uint8_threads_8_256_256_256_perm_1_2_0              34.99m ± ∞      34.83m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_float_threads_1_256_256_256_perm_1_2_0              12.09m ± ∞      12.29m ± ∞         ~ (p=0.095 n=5)
BM_Transpose_float_threads_4_256_256_256_perm_1_2_0              12.16m ± ∞      12.05m ± ∞         ~ (p=1.000 n=5)
BM_Transpose_float_threads_8_256_256_256_perm_1_2_0              12.21m ± ∞      12.11m ± ∞         ~ (p=1.000 n=5)
BM_Transpose_Pack4_threads_1_256_256_256_perm_1_2_0              26.07m ± ∞      27.02m ± ∞    +3.65% (p=0.032 n=5)
BM_Transpose_Pack4_threads_4_256_256_256_perm_1_2_0              26.51m ± ∞      26.66m ± ∞         ~ (p=0.548 n=5)
BM_Transpose_Pack4_threads_8_256_256_256_perm_1_2_0              26.45m ± ∞      28.39m ± ∞         ~ (p=0.151 n=5)
BM_Transpose_uint8_threads_1_256_256_256_perm_2_0_1              16.06m ± ∞      17.35m ± ∞    +7.99% (p=0.008 n=5)
BM_Transpose_uint8_threads_4_256_256_256_perm_2_0_1              15.99m ± ∞      17.21m ± ∞    +7.61% (p=0.008 n=5)
BM_Transpose_uint8_threads_8_256_256_256_perm_2_0_1              16.05m ± ∞      17.14m ± ∞    +6.79% (p=0.008 n=5)
BM_Transpose_float_threads_1_256_256_256_perm_2_0_1              9.971m ± ∞     10.625m ± ∞         ~ (p=0.548 n=5)
BM_Transpose_float_threads_4_256_256_256_perm_2_0_1              9.907m ± ∞     10.562m ± ∞         ~ (p=0.690 n=5)
BM_Transpose_float_threads_8_256_256_256_perm_2_0_1              10.03m ± ∞      10.50m ± ∞         ~ (p=0.421 n=5)
BM_Transpose_Pack4_threads_1_256_256_256_perm_2_0_1              8.427m ± ∞     10.370m ± ∞         ~ (p=0.095 n=5)
BM_Transpose_Pack4_threads_4_256_256_256_perm_2_0_1              8.768m ± ∞      9.186m ± ∞         ~ (p=1.000 n=5)
BM_Transpose_Pack4_threads_8_256_256_256_perm_2_0_1              8.345m ± ∞      9.124m ± ∞         ~ (p=0.690 n=5)
BM_Transpose_uint8_threads_1_256_256_256_perm_2_1_0              40.88m ± ∞      41.09m ± ∞         ~ (p=0.222 n=5)
BM_Transpose_uint8_threads_4_256_256_256_perm_2_1_0              40.83m ± ∞      41.10m ± ∞         ~ (p=0.310 n=5)
BM_Transpose_uint8_threads_8_256_256_256_perm_2_1_0              40.74m ± ∞      41.56m ± ∞         ~ (p=0.421 n=5)
BM_Transpose_float_threads_1_256_256_256_perm_2_1_0              11.66m ± ∞      11.21m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_float_threads_4_256_256_256_perm_2_1_0              11.75m ± ∞      11.42m ± ∞         ~ (p=0.310 n=5)
BM_Transpose_float_threads_8_256_256_256_perm_2_1_0              11.34m ± ∞      11.49m ± ∞         ~ (p=1.000 n=5)
BM_Transpose_Pack4_threads_1_256_256_256_perm_2_1_0              28.27m ± ∞      28.15m ± ∞         ~ (p=1.000 n=5)
BM_Transpose_Pack4_threads_4_256_256_256_perm_2_1_0              28.77m ± ∞      28.27m ± ∞         ~ (p=0.151 n=5)
BM_Transpose_Pack4_threads_8_256_256_256_perm_2_1_0              28.15m ± ∞      28.12m ± ∞         ~ (p=1.000 n=5)
BM_Transpose_uint8_threads_1_512_512_512_perm_0_2_1              31.35m ± ∞      31.42m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_uint8_threads_4_512_512_512_perm_0_2_1              16.68m ± ∞      16.51m ± ∞         ~ (p=0.690 n=5)
BM_Transpose_uint8_threads_8_512_512_512_perm_0_2_1              17.37m ± ∞      21.77m ± ∞         ~ (p=1.000 n=5)
BM_Transpose_float_threads_1_512_512_512_perm_0_2_1              91.53m ± ∞      88.32m ± ∞         ~ (p=0.056 n=5)
BM_Transpose_float_threads_4_512_512_512_perm_0_2_1              29.80m ± ∞      35.35m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_float_threads_8_512_512_512_perm_0_2_1              21.82m ± ∞      19.56m ± ∞         ~ (p=0.421 n=5)
BM_Transpose_Pack4_threads_1_512_512_512_perm_0_2_1              57.96m ± ∞      57.55m ± ∞         ~ (p=0.421 n=5)
BM_Transpose_Pack4_threads_4_512_512_512_perm_0_2_1              30.05m ± ∞      30.28m ± ∞         ~ (p=1.000 n=5)
BM_Transpose_Pack4_threads_8_512_512_512_perm_0_2_1              29.27m ± ∞      29.55m ± ∞         ~ (p=1.000 n=5)
BM_Transpose_uint8_threads_1_512_512_512_perm_1_0_2              14.16m ± ∞      14.19m ± ∞         ~ (p=0.690 n=5)
BM_Transpose_uint8_threads_4_512_512_512_perm_1_0_2              6.642m ± ∞      5.987m ± ∞         ~ (p=0.548 n=5)
BM_Transpose_uint8_threads_8_512_512_512_perm_1_0_2              6.684m ± ∞      6.247m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_float_threads_1_512_512_512_perm_1_0_2              45.36m ± ∞      45.84m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_float_threads_4_512_512_512_perm_1_0_2              25.78m ± ∞      22.61m ± ∞   -12.30% (p=0.016 n=5)
BM_Transpose_float_threads_8_512_512_512_perm_1_0_2              18.06m ± ∞      17.05m ± ∞         ~ (p=0.151 n=5)
BM_Transpose_Pack4_threads_1_512_512_512_perm_1_0_2             21.296m ± ∞      8.813m ± ∞   -58.62% (p=0.008 n=5)
BM_Transpose_Pack4_threads_4_512_512_512_perm_1_0_2             15.047m ± ∞      5.363m ± ∞   -64.36% (p=0.008 n=5)
BM_Transpose_Pack4_threads_8_512_512_512_perm_1_0_2             13.832m ± ∞      4.320m ± ∞   -68.77% (p=0.008 n=5)
BM_Transpose_uint8_threads_1_512_512_512_perm_1_2_0              62.17m ± ∞      55.24m ± ∞   -11.15% (p=0.032 n=5)
BM_Transpose_uint8_threads_4_512_512_512_perm_1_2_0              33.02m ± ∞      31.56m ± ∞         ~ (p=0.548 n=5)
BM_Transpose_uint8_threads_8_512_512_512_perm_1_2_0              33.13m ± ∞      31.46m ± ∞         ~ (p=0.421 n=5)
BM_Transpose_float_threads_1_512_512_512_perm_1_2_0              131.4m ± ∞      132.9m ± ∞         ~ (p=0.690 n=5)
BM_Transpose_float_threads_4_512_512_512_perm_1_2_0              37.96m ± ∞      37.70m ± ∞         ~ (p=1.000 n=5)
BM_Transpose_float_threads_8_512_512_512_perm_1_2_0              21.85m ± ∞      23.82m ± ∞         ~ (p=0.690 n=5)
BM_Transpose_Pack4_threads_1_512_512_512_perm_1_2_0              81.01m ± ∞      78.26m ± ∞         ~ (p=0.690 n=5)
BM_Transpose_Pack4_threads_4_512_512_512_perm_1_2_0              43.46m ± ∞      43.38m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_Pack4_threads_8_512_512_512_perm_1_2_0              39.00m ± ∞      40.35m ± ∞         ~ (p=0.222 n=5)
BM_Transpose_uint8_threads_1_512_512_512_perm_2_0_1              49.23m ± ∞      51.71m ± ∞         ~ (p=0.222 n=5)
BM_Transpose_uint8_threads_4_512_512_512_perm_2_0_1              31.18m ± ∞      30.72m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_uint8_threads_8_512_512_512_perm_2_0_1              33.80m ± ∞      30.92m ± ∞         ~ (p=1.000 n=5)
BM_Transpose_float_threads_1_512_512_512_perm_2_0_1              88.58m ± ∞      87.20m ± ∞         ~ (p=0.222 n=5)
BM_Transpose_float_threads_4_512_512_512_perm_2_0_1              35.23m ± ∞      33.60m ± ∞         ~ (p=0.056 n=5)
BM_Transpose_float_threads_8_512_512_512_perm_2_0_1              22.25m ± ∞      21.13m ± ∞         ~ (p=0.056 n=5)
BM_Transpose_Pack4_threads_1_512_512_512_perm_2_0_1              144.7m ± 3%     141.5m ± 4%        ~ (p=0.065 n=6)
BM_Transpose_Pack4_threads_4_512_512_512_perm_2_0_1              77.06m ± 22%    81.38m ± 17%       ~ (p=0.699 n=6)
BM_Transpose_Pack4_threads_8_512_512_512_perm_2_0_1              81.68m ± 6%     80.55m ± 7%        ~ (p=0.394 n=6)
BM_Transpose_uint8_threads_1_512_512_512_perm_2_1_0              69.85m ± ∞      67.00m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_uint8_threads_4_512_512_512_perm_2_1_0              35.56m ± ∞      35.46m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_uint8_threads_8_512_512_512_perm_2_1_0              34.49m ± ∞      35.18m ± ∞         ~ (p=0.690 n=5)
BM_Transpose_float_threads_1_512_512_512_perm_2_1_0              72.82m ± ∞      72.71m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_float_threads_4_512_512_512_perm_2_1_0              27.38m ± ∞      28.76m ± ∞         ~ (p=0.421 n=5)
BM_Transpose_float_threads_8_512_512_512_perm_2_1_0              21.81m ± ∞      20.06m ± ∞    -8.00% (p=0.032 n=5)
BM_Transpose_Pack4_threads_1_512_512_512_perm_2_1_0              134.5m ± ∞      139.3m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_Pack4_threads_4_512_512_512_perm_2_1_0              68.28m ± ∞      66.73m ± ∞         ~ (p=0.151 n=5)
BM_Transpose_Pack4_threads_8_512_512_512_perm_2_1_0              71.95m ± ∞      69.90m ± ∞         ~ (p=0.056 n=5)
BM_Transpose_uint8_threads_1_64_224_224_3_perm_1_2_3_0           4.088m ± ∞      3.987m ± ∞    -2.46% (p=0.008 n=5)
BM_Transpose_uint8_threads_4_64_224_224_3_perm_1_2_3_0           4.094m ± ∞      3.954m ± ∞         ~ (p=0.222 n=5)
BM_Transpose_uint8_threads_8_64_224_224_3_perm_1_2_3_0           4.079m ± ∞      3.985m ± ∞         ~ (p=0.151 n=5)
BM_Transpose_float_threads_1_64_224_224_3_perm_1_2_3_0           10.71m ± ∞      10.59m ± ∞         ~ (p=1.000 n=5)
BM_Transpose_float_threads_4_64_224_224_3_perm_1_2_3_0           10.57m ± ∞      11.12m ± ∞         ~ (p=0.690 n=5)
BM_Transpose_float_threads_8_64_224_224_3_perm_1_2_3_0           10.31m ± ∞      10.48m ± ∞         ~ (p=1.000 n=5)
BM_Transpose_Pack4_threads_1_64_224_224_3_perm_1_2_3_0           6.811m ± ∞      6.757m ± ∞         ~ (p=0.690 n=5)
BM_Transpose_Pack4_threads_4_64_224_224_3_perm_1_2_3_0           6.751m ± ∞      6.755m ± ∞         ~ (p=0.690 n=5)
BM_Transpose_Pack4_threads_8_64_224_224_3_perm_1_2_3_0           6.767m ± ∞      6.801m ± ∞         ~ (p=0.548 n=5)
BM_Transpose_uint8_threads_1_256_64_64_3_perm_1_3_2_0            4.508m ± ∞      4.487m ± ∞         ~ (p=0.841 n=5)
BM_Transpose_uint8_threads_4_256_64_64_3_perm_1_3_2_0            4.513m ± ∞      4.491m ± ∞         ~ (p=0.690 n=5)
BM_Transpose_uint8_threads_8_256_64_64_3_perm_1_3_2_0            4.502m ± ∞      4.472m ± ∞         ~ (p=0.421 n=5)
BM_Transpose_float_threads_1_256_64_64_3_perm_1_3_2_0            5.613m ± 14%    6.057m ± 20%       ~ (p=0.240 n=6)
BM_Transpose_float_threads_4_256_64_64_3_perm_1_3_2_0            5.379m ± 6%     5.664m ± 13%       ~ (p=0.240 n=6)
BM_Transpose_float_threads_8_256_64_64_3_perm_1_3_2_0            5.517m ± 7%     5.784m ± 16%       ~ (p=0.065 n=6)
BM_Transpose_Pack4_threads_1_256_64_64_3_perm_1_3_2_0            6.839m ± ∞      6.765m ± ∞         ~ (p=0.151 n=5)
BM_Transpose_Pack4_threads_4_256_64_64_3_perm_1_3_2_0            6.831m ± ∞      6.902m ± ∞         ~ (p=0.690 n=5)
BM_Transpose_Pack4_threads_8_256_64_64_3_perm_1_3_2_0            6.839m ± ∞      6.764m ± ∞    -1.10% (p=0.032 n=5)
BM_TransposeMicroKernel<int8_t, 32>                              102.1n ± ∞      101.6n ± ∞         ~ (p=0.690 n=5)
BM_TransposeMicroKernel<int16_t, 16>                             12.91n ± ∞      12.93n ± ∞         ~ (p=0.690 n=5)
BM_TransposeMicroKernel<float, 8>                                16.11n ± ∞      16.08n ± ∞         ~ (p=1.000 n=5)
BM_TransposeMicroKernel<int64_t, 4>                              8.750n ± ∞      8.783n ± ∞         ~ (p=0.500 n=5)
BM_TransposeMicroKernel<absl::uint128, 2>                        5.947n ± ∞      5.961n ± ∞         ~ (p=0.524 n=5)
BM_TransposeMicroKernel<int8_t, 16>                              10.51n ± ∞      10.52n ± ∞         ~ (p=1.000 n=5)
BM_TransposeMicroKernel<int16_t, 8>                              16.44n ± ∞      16.39n ± ∞         ~ (p=0.841 n=5)
BM_TransposeMicroKernel<float, 4>                                20.38n ± ∞      20.38n ± ∞         ~ (p=0.841 n=5)
BM_TransposeMicroKernel<int64_t, 2>                              13.47n ± ∞      13.55n ± ∞         ~ (p=0.056 n=5)
BM_TransposeMicroKernel<int8_t, 16, StrideMode::kWideSymmetric>  14.16n ± ∞      14.25n ± ∞         ~ (p=0.151 n=5)
BM_TransposeMicroKernel<int16_t, 8, StrideMode::kWideSymmetric>  21.86n ± ∞      21.98n ± ∞         ~ (p=0.310 n=5)
BM_TransposeMicroKernel<float, 4, StrideMode::kWideSymmetric>    43.03n ± ∞      43.02n ± ∞         ~ (p=0.841 n=5)
BM_TransposeMicroKernel<int64_t, 2, StrideMode::kWideSymmetric>  12.24n ± ∞      12.20n ± ∞         ~ (p=0.095 n=5)
BM_TransposeMicroKernel<int8_t, 16, StrideMode::kWideScatter>    13.97n ± ∞      14.01n ± ∞         ~ (p=0.151 n=5)
BM_TransposeMicroKernel<int16_t, 8, StrideMode::kWideScatter>    21.87n ± ∞      21.89n ± ∞         ~ (p=0.841 n=5)
BM_TransposeMicroKernel<float, 4, StrideMode::kWideScatter>      43.31n ± ∞      43.01n ± ∞         ~ (p=0.095 n=5)
BM_TransposeMicroKernel<int64_t, 2, StrideMode::kWideScatter>    12.26n ± ∞      12.28n ± ∞         ~ (p=0.841 n=5)
BM_TransposeMicroKernel<int8_t, 8>                               21.78n ± ∞      21.62n ± ∞         ~ (p=0.222 n=5)
BM_TransposeMicroKernel<int16_t, 4>                              33.17n ± ∞      33.31n ± ∞         ~ (p=0.690 n=5)
BM_TransposeMicroKernel<float, 2>                                79.43n ± ∞      79.77n ± ∞         ~ (p=0.548 n=5)
BM_TransposeMicroKernel<int8_t, 4>                               37.20n ± ∞      37.35n ± ∞         ~ (p=1.000 n=5)
BM_TransposeMicroKernel<int16_t, 2>                              102.4n ± ∞      102.1n ± ∞         ~ (p=0.802 n=5)
BM_TransposeMicroKernel<int8_t, 2>                               104.0n ± ∞      104.0n ± ∞         ~ (p=0.548 n=5)
geomean                                                         824.4µ          793.8µ         -3.71%
```

</details>

PiperOrigin-RevId: 996980931
---
 third_party/xla/xla/pjrt/transpose.cc      | 276 +++++++++++-------
 third_party/xla/xla/pjrt/transpose.h       |   4 +
 third_party/xla/xla/pjrt/transpose_test.cc | 312 +++++++++++++++++++--
 3 files changed, 454 insertions(+), 138 deletions(-)

diff --git a/third_party/xla/xla/pjrt/transpose.cc b/third_party/xla/xla/pjrt/transpose.cc
index fb0513e9f28..c3f57b33ee9 100644
--- a/third_party/xla/xla/pjrt/transpose.cc
+++ b/third_party/xla/xla/pjrt/transpose.cc
@@ -420,59 +420,81 @@ void Transpose(const char* __restrict a, int outer_bs_a, char* __restrict b,
   }
 }
 
+template <TransposePlan::Transformation transformation>
 void TransposeConstStride1(const char* __restrict a, char* __restrict b,
-                           TransposePlan::Node const* __restrict node) {
+                           TransposePlan::Node const* __restrict node,
+                           int bits_per_element) {
+  constexpr bool pack =
+      (transformation == TransposePlan::Transformation::kPackSubbyte);
+  const int elements_per_byte = pack ? (8 / bits_per_element) : 1;
+  auto copy_or_pack = [&](const char* __restrict src, char* __restrict dst,
+                          int64_t num_bytes) {
+    if constexpr (pack) {
+      PackIntN(bits_per_element, absl::MakeConstSpan(src, num_bytes),
+               absl::MakeSpan(dst, num_bytes / elements_per_byte));
+    } else {
+      std::memcpy(dst, src, num_bytes);
+    }
+  };
+  const int64_t ldb0 = pack ? (node[0].ldb / elements_per_byte) : node[0].ldb;
   if (node[0].is_inner_dim_in_a) {
-    int64_t num_bytes = node->end;
-    std::memcpy(b, a, num_bytes);
+    copy_or_pack(a, b, node[0].end);
   } else if (node[1].is_inner_dim_in_a) {
     int64_t num_bytes = node[1].end;
     for (int64_t i = 0; i < node[0].end; ++i) {
-      std::memcpy(b, a, num_bytes);
+      copy_or_pack(a, b, num_bytes);
       a += node[0].lda;
-      b += node[0].ldb;
+      b += ldb0;
     }
     if (node[0].trailing_tile_next_node_inc) {
-      TransposeConstStride1(a, b, node + node[0].trailing_tile_next_node_inc);
+      TransposeConstStride1<transformation>(
+          a, b, node + node[0].trailing_tile_next_node_inc, bits_per_element);
     }
   } else if (node[2].is_inner_dim_in_a) {
     int64_t num_bytes = node[2].end;
+    const int64_t ldb1 = pack ? (node[1].ldb / elements_per_byte) : node[1].ldb;
     for (int64_t i = 0; i < node[0].end; ++i) {
       const char* a1 = a;
       char* b1 = b;
       for (int64_t j = 0; j < node[1].end; ++j) {
-        std::memcpy(b1, a1, num_bytes);
+        copy_or_pack(a1, b1, num_bytes);
         a1 += node[1].lda;
-        b1 += node[1].ldb;
+        b1 += ldb1;
       }
       if (node[1].trailing_tile_next_node_inc) {
-        TransposeConstStride1(a1, b1,
-                              &node[1] + node[1].trailing_tile_next_node_inc);
+        TransposeConstStride1<transformation>(
+            a1, b1, &node[1] + node[1].trailing_tile_next_node_inc,
+            bits_per_element);
       }
       a += node[0].lda;
-      b += node[0].ldb;
+      b += ldb0;
     }
     if (node[0].trailing_tile_next_node_inc) {
-      TransposeConstStride1(a, b, node + node[0].trailing_tile_next_node_inc);
+      TransposeConstStride1<transformation>(
+          a, b, node + node[0].trailing_tile_next_node_inc, bits_per_element);
     }
   } else {
+    const int64_t ldb1 = pack ? (node[1].ldb / elements_per_byte) : node[1].ldb;
     for (int64_t i = 0; i < node[0].end; ++i) {
       const char* a1 = a;
       char* b1 = b;
       for (int64_t j = 0; j < node[1].end; ++j) {
-        TransposeConstStride1(a1, b1, node + 2);
+        TransposeConstStride1<transformation>(a1, b1, node + 2,
+                                              bits_per_element);
         a1 += node[1].lda;
-        b1 += node[1].ldb;
+        b1 += ldb1;
       }
       if (node[1].trailing_tile_next_node_inc) {
-        TransposeConstStride1(a1, b1,
-                              &node[1] + node[1].trailing_tile_next_node_inc);
+        TransposeConstStride1<transformation>(
+            a1, b1, &node[1] + node[1].trailing_tile_next_node_inc,
+            bits_per_element);
       }
       a += node[0].lda;
-      b += node[0].ldb;
+      b += ldb0;
     }
     if (node[0].trailing_tile_next_node_inc) {
-      TransposeConstStride1(a, b, node + node[0].trailing_tile_next_node_inc);
+      TransposeConstStride1<transformation>(
+          a, b, node + node[0].trailing_tile_next_node_inc, bits_per_element);
     }
   }
 }
@@ -578,6 +600,16 @@ ABSL_ATTRIBUTE_FUNC_ALIGN(64)
 void TransposePlan::ExecuteChunk(int chunk_id, const void* a, void* b,
                                  bool input_is_global,
                                  bool output_is_global) const {
+  CHECK(!use_fallback_pack_)
+      << "ExecuteChunk does not support unaligned sub-byte packing; use "
+         "Execute().";
+  ExecuteChunkInternal(chunk_id, a, b, input_is_global, output_is_global);
+}
+
+ABSL_ATTRIBUTE_FUNC_ALIGN(64)
+void TransposePlan::ExecuteChunkInternal(int chunk_id, const void* a, void* b,
+                                         bool input_is_global,
+                                         bool output_is_global) const {
   if (num_elems_ == 0) {
     return;
   }
@@ -592,88 +624,80 @@ void TransposePlan::ExecuteChunk(int chunk_id, const void* a, void* b,
     ac += input_chunk_iteration_offsets_[chunk_id];
   }
   char* bc = static_cast<char*>(b);
-  if (transformation_ == Transformation::kPackSubbyte) {
-    int elements_per_byte = 8 / bits_per_element_;
-    if (output_is_global) {
-      int64_t offset = output_chunk_offset_bytes_[chunk_id] +
-                       output_chunk_iteration_offsets_[chunk_id];
-      DCHECK_EQ(offset % elements_per_byte, 0);
-      bc += offset / elements_per_byte;
-    } else {
-      int64_t offset = output_chunk_iteration_offsets_[chunk_id];
-      DCHECK_EQ(offset % elements_per_byte, 0);
-      bc += offset / elements_per_byte;
-    }
+  if (output_is_global) {
+    bc += output_chunk_offset_bytes_[chunk_id] +
+          output_chunk_iteration_offsets_[chunk_id];
   } else {
-    if (output_is_global) {
-      bc += output_chunk_offset_bytes_[chunk_id] +
-            output_chunk_iteration_offsets_[chunk_id];
-    } else {
-      bc += output_chunk_iteration_offsets_[chunk_id];
-    }
+    bc += output_chunk_iteration_offsets_[chunk_id];
   }
 
   if (inner_kernel_is_memcpy()) {
-    CHECK(transformation_ == Transformation::kNone);
-    // Memcpy-based plans all assume element size 1 (i.e., bytes).
-    TransposeConstStride1(ac, bc, nodes.data());
-    return;
-  }
-
-  switch (elem_size_in_bytes_) {
-    case 1:
-      if (transformation_ == Transformation::kNone) {
-        ExecuteTyped<uint8_t, Transformation::kNone>(ac, bc, nodes, 0);
-      } else if (transformation_ == Transformation::kPackSubbyte) {
-        ExecuteTyped<uint8_t, Transformation::kPackSubbyte>(ac, bc, nodes,
-                                                            bits_per_element_);
-      } else {
-        LOG(FATAL) << "Unsupported transformation for elem_size=1: "
-                   << static_cast<int>(transformation_);
-      }
-      break;
-    case 2:
-      ExecuteTyped<uint16_t, Transformation::kNone>(ac, bc, nodes, 0);
-      break;
-    case 4:
-      if (transformation_ == Transformation::kNone) {
-        ExecuteTyped<uint32_t, Transformation::kNone>(ac, bc, nodes, 0);
-      } else {
-        DCHECK(transformation_ == Transformation::kF64ToEf57);
-        ExecuteTyped<uint32_t, Transformation::kF64ToEf57>(ac, bc, nodes, 0);
-      }
-      break;
-    case 8:
-      ExecuteTyped<uint64_t, Transformation::kNone>(ac, bc, nodes, 0);
-      break;
-    case 16:
-      ExecuteTyped<uint128, Transformation::kNone>(ac, bc, nodes, 0);
-      break;
-    default:
-      LOG(FATAL) << "Unimplemented element size " << elem_size_in_bytes_;
+    if (transformation_ == Transformation::kNone) {
+      // Memcpy-based plans all assume element size 1 (i.e., bytes).
+      TransposeConstStride1<Transformation::kNone>(ac, bc, nodes.data(), 0);
+    } else if (transformation_ == Transformation::kPackSubbyte) {
+      TransposeConstStride1<Transformation::kPackSubbyte>(ac, bc, nodes.data(),
+                                                          bits_per_element_);
+    } else {
+      LOG(FATAL) << "Unsupported transformation for memcpy plan: "
+                 << static_cast<int>(transformation_);
+    }
+  } else {
+    switch (elem_size_in_bytes_) {
+      case 1:
+        if (transformation_ == Transformation::kNone) {
+          ExecuteTyped<uint8_t, Transformation::kNone>(ac, bc, nodes, 0);
+        } else if (transformation_ == Transformation::kPackSubbyte) {
+          ExecuteTyped<uint8_t, Transformation::kPackSubbyte>(
+              ac, bc, nodes, bits_per_element_);
+        } else {
+          LOG(FATAL) << "Unsupported transformation for elem_size=1: "
+                     << static_cast<int>(transformation_);
+        }
+        break;
+      case 2:
+        ExecuteTyped<uint16_t, Transformation::kNone>(ac, bc, nodes, 0);
+        break;
+      case 4:
+        if (transformation_ == Transformation::kNone) {
+          ExecuteTyped<uint32_t, Transformation::kNone>(ac, bc, nodes, 0);
+        } else {
+          DCHECK(transformation_ == Transformation::kF64ToEf57);
+          ExecuteTyped<uint32_t, Transformation::kF64ToEf57>(ac, bc, nodes, 0);
+        }
+        break;
+      case 8:
+        ExecuteTyped<uint64_t, Transformation::kNone>(ac, bc, nodes, 0);
+        break;
+      case 16:
+        ExecuteTyped<uint128, Transformation::kNone>(ac, bc, nodes, 0);
+        break;
+      default:
+        LOG(FATAL) << "Unimplemented element size " << elem_size_in_bytes_;
+    }
   }
 }
 
 void TransposePlan::ExecuteInternal(
     const void* a, void* b,
-    std::optional<absl::FunctionRef<void(std::function<void(void)>)>>
-        schedule_work) const {
+    std::optional<absl::FunctionRef<void(std::function<void()>)>> schedule_work)
+    const {
   if (!schedule_work || Parallelism() <= 1) {
     for (int i = 0; i < Parallelism(); ++i) {
-      ExecuteChunk(i, a, b, /*input_is_global=*/true,
-                   /*output_is_global=*/true);
+      ExecuteChunkInternal(i, a, b, /*input_is_global=*/true,
+                           /*output_is_global=*/true);
     }
   } else {
     absl::BlockingCounter counter(Parallelism() - 1);
     for (size_t i = 1; i < nodes_.size(); ++i) {
       (*schedule_work)([&, i]() {
-        ExecuteChunk(i, a, b, /*input_is_global=*/true,
-                     /*output_is_global=*/true);
+        ExecuteChunkInternal(i, a, b, /*input_is_global=*/true,
+                             /*output_is_global=*/true);
         counter.DecrementCount();
       });
     }
-    ExecuteChunk(0, a, b, /*input_is_global=*/true,
-                 /*output_is_global=*/true);
+    ExecuteChunkInternal(0, a, b, /*input_is_global=*/true,
+                         /*output_is_global=*/true);
     counter.Wait();
   }
 }
@@ -689,13 +713,15 @@ void TransposePlan::Execute(
   if (use_fallback_pack_) {
     tsl::profiler::TraceMe traceme("Transpose::ExecuteAndSubbytePack",
                                    /*level=*/2);
-    std::unique_ptr<char[]> temp_buf(new char[num_elems_]);
+    const int64_t output_num_elems = OutputNumElems();
+    auto temp_buf = std::make_unique<char[]>(output_num_elems);
     ExecuteInternal(a, temp_buf.get(), schedule_work);
     int elements_per_byte = 8 / bits_per_element_;
-    PackIntN(
-        bits_per_element_, absl::MakeConstSpan(temp_buf.get(), num_elems_),
-        absl::MakeSpan(static_cast<char*>(b),
-                       CeilOfRatio<int64_t>(num_elems_, elements_per_byte)));
+    PackIntN(bits_per_element_,
+             absl::MakeConstSpan(temp_buf.get(), output_num_elems),
+             absl::MakeSpan(
+                 static_cast<char*>(b),
+                 CeilOfRatio<int64_t>(output_num_elems, elements_per_byte)));
   } else {
     tsl::profiler::TraceMe traceme("Transpose::Execute", /*level=*/2);
     ExecuteInternal(a, b, schedule_work);
@@ -1141,7 +1167,9 @@ absl::StatusOr<std::unique_ptr<TransposePlan>> TransposePlan::Create(
         int elements_per_byte = 8 / bits;
         plan->bits_per_element_ = bits;
         if (output_dims.empty() ||
-            output_dims.back() % elements_per_byte != 0) {
+            output_dims.back() % elements_per_byte != 0 ||
+            (plan->b_is_tiled_ &&
+             plan->b_tiling_.back() % elements_per_byte != 0)) {
           plan->transformation_ = Transformation::kNone;
           plan->use_fallback_pack_ = true;
         }
@@ -1215,7 +1243,18 @@ absl::Status TransposePlan::Initialize() {
   }
   CHECK_GE(pos_stride1b_in_b, 0);
 
-  const int pos_stride1b_in_a = permutation_[pos_stride1b_in_b];
+  int pos_stride1b_in_a = permutation_[pos_stride1b_in_b];
+  if (transformation_ == Transformation::kPackSubbyte) {
+    int elements_per_byte = 8 / bits_per_element_;
+    int64_t b_stride1_tile_size =
+        std::max(a_tiling_[pos_stride1b_in_a], b_tiling_[pos_stride1b_in_b]);
+    if (b_dims_[pos_stride1b_in_b] % elements_per_byte != 0 ||
+        (b_stride1_tile_size > 1 &&
+         b_stride1_tile_size % elements_per_byte != 0)) {
+      transformation_ = Transformation::kNone;
+      use_fallback_pack_ = true;
+    }
+  }
 
   int64_t stride_pos1a =
       inner_stride(pos_stride1a_in_a, lda_, lda_tile_, a_tiling_);
@@ -1228,14 +1267,10 @@ absl::Status TransposePlan::Initialize() {
                            ? InnerKernelKind::kMemcpy
                            : InnerKernelKind::kDefault;
 
-  if (inner_kernel_is_memcpy() && transformation_ != Transformation::kNone) {
-    if (transformation_ == Transformation::kPackSubbyte) {
-      transformation_ = Transformation::kNone;
-      use_fallback_pack_ = true;
-    } else {
-      return InvalidArgument(
-          "Memcpy-based plans cannot be used with transformations.");
-    }
+  if (inner_kernel_is_memcpy() && transformation_ != Transformation::kNone &&
+      transformation_ != Transformation::kPackSubbyte) {
+    return InvalidArgument(
+        "Memcpy-based plans cannot be used with transformations.");
   }
 
   // Calculate sentinel strides.
@@ -1414,9 +1449,13 @@ absl::Status TransposePlan::Initialize() {
   ChooseLoopOrder(loop_order);
 
   for (Loop& loop : loop_order) {
-    if (!inner_kernel_is_memcpy() &&
-        (loop.tile_interior || loop.tile_size == 1)) {
-      if (loop.is_inner_dim_in_a) {
+    if (loop.tile_interior || loop.tile_size == 1) {
+      if (inner_kernel_is_memcpy()) {
+        if (transformation_ == Transformation::kPackSubbyte &&
+            loop.is_inner_dim_in_a) {
+          loop.inc = 8 / bits_per_element_;
+        }
+      } else if (loop.is_inner_dim_in_a) {
         loop.inc = (inner_kernel_kind_ == InnerKernelKind::kDeinterleave)
                        ? loop.dim_size
                        : inner_block_elems_ * outer_block_elems_a_;
@@ -1460,10 +1499,11 @@ absl::Status TransposePlan::Initialize() {
       DCHECK(!inner_kernel_is_memcpy());
       break;
     case Transformation::kPackSubbyte:
-      scratch_size_ = sizeof(uint8_t) * inner_block_elems_ *
-                      inner_block_elems_ * outer_block_elems_a_ *
-                      outer_block_elems_b_;
-      DCHECK(!inner_kernel_is_memcpy());
+      scratch_size_ = inner_kernel_is_memcpy()
+                          ? 0
+                          : sizeof(uint8_t) * inner_block_elems_ *
+                                inner_block_elems_ * outer_block_elems_a_ *
+                                outer_block_elems_b_;
       break;
   }
 
@@ -1599,6 +1639,8 @@ int TransposePlan::ChooseParallelizationStrategy(
   } else if (!inner_kernel_is_memcpy()) {
     acc *= inner_block_elems_ * inner_block_elems_ * outer_block_elems_a_ *
            outer_block_elems_b_;
+  } else if (transformation_ == Transformation::kPackSubbyte) {
+    acc *= 8 / bits_per_element_;
   }
   auto work_it = work_in_bytes.rbegin();
   for (auto it = loop_order.rbegin(); it != loop_order.rend(); ++it) {
@@ -1639,7 +1681,10 @@ int TransposePlan::ChooseParallelizationStrategy(
       continue;
     }
 
-    int kMinBytesPerThread = inner_kernel_is_memcpy() ? (1 << 20) : (1 << 26);
+    int kMinBytesPerThread =
+        chunk_contiguity_ != ChunkContiguity::kNone
+            ? 1
+            : (inner_kernel_is_memcpy() ? (1 << 20) : (1 << 26));
     int64_t min_iterations_per_thread =
         CeilOfRatio<int64_t>(kMinBytesPerThread, work_in_bytes[i]);
     int64_t parallel_work = CeilOfRatio(iterations, min_iterations_per_thread);
@@ -1803,6 +1848,16 @@ void TransposePlan::ComputeChunkSizes() {
         input_max - input_min + elem_size_in_bytes_;
     output_chunk_size_bytes_[chunk_id] =
         output_max - output_min + elem_size_in_bytes_;
+    if (transformation_ == Transformation::kPackSubbyte) {
+      const int elements_per_byte = 8 / bits_per_element_;
+      DCHECK_EQ(output_chunk_offset_bytes_[chunk_id] % elements_per_byte, 0);
+      DCHECK_EQ(output_chunk_iteration_offsets_[chunk_id] % elements_per_byte,
+                0);
+      DCHECK_EQ(output_chunk_size_bytes_[chunk_id] % elements_per_byte, 0);
+      output_chunk_offset_bytes_[chunk_id] /= elements_per_byte;
+      output_chunk_iteration_offsets_[chunk_id] /= elements_per_byte;
+      output_chunk_size_bytes_[chunk_id] /= elements_per_byte;
+    }
   }
 }
 
@@ -1893,15 +1948,16 @@ bool TransposePlanCacheKey::operator==(
          output_tiling == other.output_tiling &&
          transformation == other.transformation &&
          dest_bits_per_element == other.dest_bits_per_element &&
-         num_threads == other.num_threads;
+         num_threads == other.num_threads &&
+         chunk_contiguity == other.chunk_contiguity;
 }
 
 template <typename H>
 H AbslHashValue(H h, const TransposePlanCacheKey& key) {
   return H::combine(std::move(h), key.elem_size_in_bytes, key.num_threads,
-                    key.transformation, key.dest_bits_per_element, key.dims,
-                    key.permutation, key.input_tiling, key.input_striding,
-                    key.output_tiling);
+                    key.chunk_contiguity, key.transformation,
+                    key.dest_bits_per_element, key.dims, key.permutation,
+                    key.input_tiling, key.input_striding, key.output_tiling);
 }
 
 TransposePlanCache::TransposePlanCache(int capacity)
@@ -1937,6 +1993,7 @@ absl::StatusOr<std::shared_ptr<TransposePlan>> TransposePlanCache::GetOrCreate(
   key.transformation = transformation;
   key.dest_bits_per_element = o.dest_bits_per_element;
   key.num_threads = o.num_threads;
+  key.chunk_contiguity = o.chunk_contiguity;
   return cache_.GetOrCreateIfAbsent(
       key,
       [&](const TransposePlanCacheKey& key)
@@ -1973,7 +2030,8 @@ void TransposePlan::IdentifyContiguousLoops(
     for (size_t idx : indices) {
       Loop& loop = loop_order[idx];
       int64_t stride = contiguous_stride(loop);
-      if (stride == 0) {
+      if (stride == 0 ||
+          (!loop.tile_interior && loop.dim_size == loop.tile_size)) {
         continue;
       }
       if (stride >= target_chunk_bytes) {
diff --git a/third_party/xla/xla/pjrt/transpose.h b/third_party/xla/xla/pjrt/transpose.h
index 4ae5f822e03..5e46dbfde70 100644
--- a/third_party/xla/xla/pjrt/transpose.h
+++ b/third_party/xla/xla/pjrt/transpose.h
@@ -329,6 +329,9 @@ class TransposePlan {
   void ExecuteTyped(const char* a, char* b, absl::Span<Node const> nodes,
                     int bits_per_element) const;
 
+  void ExecuteChunkInternal(int chunk_id, const void* a, void* b,
+                            bool input_is_global, bool output_is_global) const;
+
   void ExecuteInternal(
       const void* a, void* b,
       std::optional<absl::FunctionRef<void(std::function<void()>)>>
@@ -448,6 +451,7 @@ struct TransposePlanCacheKey {
   TransposePlan::Transformation transformation;
   std::optional<int> dest_bits_per_element;
   int num_threads;
+  TransposePlan::ChunkContiguity chunk_contiguity;
 
   bool operator==(const TransposePlanCacheKey& other) const;
 };
diff --git a/third_party/xla/xla/pjrt/transpose_test.cc b/third_party/xla/xla/pjrt/transpose_test.cc
index 25d22f8bd12..0efc1cedd41 100644
--- a/third_party/xla/xla/pjrt/transpose_test.cc
+++ b/third_party/xla/xla/pjrt/transpose_test.cc
@@ -23,6 +23,7 @@ limitations under the License.
 #include <functional>
 #include <memory>
 #include <numeric>
+#include <optional>
 #include <ostream>
 #include <string>
 #include <tuple>
@@ -1488,93 +1489,346 @@ TEST(TransposePlanCache, Basics) {
   EXPECT_TRUE(p3.get() != p1.get());
   TF_ASSERT_OK_AND_ASSIGN(auto p1b, cache.GetOrCreate(o));
   EXPECT_TRUE(p1.get() != p1b.get());
+
+  TransposePlanCache contiguity_cache(4);
+  TransposePlan::Options o_none = o;
+  o_none.chunk_contiguity = TransposePlan::ChunkContiguity::kNone;
+  TransposePlan::Options o_input = o;
+  o_input.chunk_contiguity = TransposePlan::ChunkContiguity::kInput;
+  TransposePlan::Options o_output = o;
+  o_output.chunk_contiguity = TransposePlan::ChunkContiguity::kOutput;
+  ASSERT_OK_AND_ASSIGN(std::shared_ptr<TransposePlan> p_none,
+                       contiguity_cache.GetOrCreate(o_none));
+  ASSERT_OK_AND_ASSIGN(std::shared_ptr<TransposePlan> p_input,
+                       contiguity_cache.GetOrCreate(o_input));
+  ASSERT_OK_AND_ASSIGN(std::shared_ptr<TransposePlan> p_output,
+                       contiguity_cache.GetOrCreate(o_output));
+  EXPECT_NE(p_none, p_input);
+  EXPECT_NE(p_none, p_output);
+  EXPECT_NE(p_input, p_output);
+  ASSERT_OK_AND_ASSIGN(std::shared_ptr<TransposePlan> p_output_again,
+                       contiguity_cache.GetOrCreate(o_output));
+  EXPECT_EQ(p_output, p_output_again);
 }
 
-void TestPackIntN(int bits_per_element, absl::Span<const int64_t> dims,
-                  absl::Span<const int64_t> permutation) {
-  int elements_per_byte = 8 / bits_per_element;
+void TestPackIntN(
+    int bits_per_element, absl::Span<const int64_t> dims,
+    absl::Span<const int64_t> permutation,
+    absl::Span<const int64_t> input_tiling = {},
+    absl::Span<const int64_t> output_tiling = {},
+    absl::Span<const int64_t> input_striding = {}, int num_threads = 1,
+    TransposePlan::ChunkContiguity chunk_contiguity =
+        TransposePlan::ChunkContiguity::kNone,
+    std::optional<int> expected_parallelism = std::nullopt,
+    std::optional<bool> expected_inner_kernel_is_memcpy = std::nullopt) {
+  const int elements_per_byte = 8 / bits_per_element;
+  std::vector<int64_t> output_dims = Permute(dims, permutation);
+  std::vector<int64_t> full_input_tiling = PadTiling(dims, input_tiling);
+  std::vector<int64_t> full_input_strides =
+      input_striding.empty()
+          ? ComputeDefaultStrides(dims, full_input_tiling, 1)
+          : std::vector<int64_t>(input_striding.begin(), input_striding.end());
+  std::vector<int64_t> full_output_tiling =
+      PadTiling(output_dims, output_tiling);
+  std::vector<int64_t> full_output_strides =
+      ComputeDefaultStrides(output_dims, full_output_tiling, 1);
 
-  int64_t num_elems = 1;
-  for (int64_t dim : dims) {
-    num_elems *= dim;
-  }
-  std::vector<int8_t> input(num_elems);
+  int64_t min_offset_bytes;
+  int64_t input_size_bytes = SizeOfTiledArray(
+      dims, full_input_tiling, full_input_strides, 1, &min_offset_bytes);
+  std::vector<int8_t> input(input_size_bytes);
   absl::BitGen bitgen;
   int max_val = (1 << bits_per_element) - 1;
-  for (int64_t i = 0; i < num_elems; ++i) {
-    input[i] = absl::Uniform<int>(bitgen, 0, max_val + 1);
+  for (int8_t& val : input) {
+    val = absl::Uniform<int>(bitgen, 0, max_val + 1);
   }
+  const int64_t input_base_offset_bytes = -min_offset_bytes;
 
   TransposePlan::Options options;
   options.elem_size_in_bytes = 1;
   options.dims = dims;
   options.permutation = permutation;
+  if (!input_tiling.empty()) {
+    options.input_tiling = TransposePlan::Tiling{input_tiling};
+  }
+  if (!output_tiling.empty()) {
+    options.output_tiling = TransposePlan::Tiling{output_tiling};
+  }
+  if (!input_striding.empty()) {
+    options.input_striding = TransposePlan::Striding{input_striding};
+  }
   options.dest_bits_per_element = bits_per_element;
+  options.num_threads = num_threads;
+  options.chunk_contiguity = chunk_contiguity;
 
-  TF_ASSERT_OK_AND_ASSIGN(auto plan, TransposePlan::Create(options));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<TransposePlan> plan,
+                       TransposePlan::Create(options));
+  if (expected_parallelism.has_value()) {
+    EXPECT_EQ(plan->Parallelism(), *expected_parallelism);
+  }
+  if (expected_inner_kernel_is_memcpy.has_value()) {
+    EXPECT_EQ(plan->inner_kernel_is_memcpy(), *expected_inner_kernel_is_memcpy);
+  }
 
-  int64_t output_size_bytes =
-      CeilOfRatio<int64_t>(num_elems, elements_per_byte);
-  std::vector<char> output(output_size_bytes, -1);
+  const int64_t output_num_elems = plan->OutputNumElems();
+  const int64_t output_size_bytes =
+      CeilOfRatio<int64_t>(output_num_elems, elements_per_byte);
 
-  plan->Execute(reinterpret_cast<const char*>(input.data()), output.data());
+  std::vector<int8_t> expected_unpacked_transposed(output_num_elems, 0);
+  ReferenceTranspose<int8_t>(
+      dims, permutation, full_input_tiling, full_input_strides,
+      full_output_tiling, full_output_strides, input,
+      absl::MakeSpan(expected_unpacked_transposed), input_base_offset_bytes);
 
-  std::vector<int8_t> expected_unpacked_transposed(num_elems);
-  std::vector<int64_t> output_dims = Permute(dims, permutation);
-  std::vector<int64_t> input_tiling = PadTiling(dims, {});
-  std::vector<int64_t> input_strides =
-      ComputeDefaultStrides(dims, input_tiling, 1);
-  std::vector<int64_t> output_tiling = PadTiling(output_dims, {});
-  std::vector<int64_t> output_strides =
-      ComputeDefaultStrides(output_dims, output_tiling, 1);
-  ReferenceTranspose<int8_t>(dims, permutation, input_tiling, input_strides,
-                             output_tiling, output_strides, input,
-                             absl::MakeSpan(expected_unpacked_transposed), 0);
-
-  std::vector<char> expected_output(output_size_bytes);
+  std::vector<char> expected_output(output_size_bytes, 0);
   PackIntN(bits_per_element,
            absl::MakeConstSpan(reinterpret_cast<const char*>(
                                    expected_unpacked_transposed.data()),
-                               num_elems),
+                               output_num_elems),
            absl::MakeSpan(expected_output));
 
+  const char* input_base =
+      reinterpret_cast<const char*>(input.data()) + input_base_offset_bytes;
+  std::vector<char> output(output_size_bytes, 0);
+  tsl::thread::ThreadPool threadpool(tsl::Env::Default(), "Transpose",
+                                     std::max(1, num_threads));
+  plan->Execute(input_base, output.data(), [&](std::function<void()> fn) {
+    threadpool.Schedule(std::move(fn));
+  });
   EXPECT_EQ(output, expected_output);
+
+  if (chunk_contiguity == TransposePlan::ChunkContiguity::kOutput) {
+    std::vector<char> global_chunked_output(output_size_bytes, 0);
+    for (int i = 0; i < plan->Parallelism(); ++i) {
+      plan->ExecuteChunk(i, input_base, global_chunked_output.data(),
+                         /*input_is_global=*/true,
+                         /*output_is_global=*/true);
+    }
+    EXPECT_EQ(global_chunked_output, expected_output);
+
+    std::vector<char> local_chunked_output(output_size_bytes, 0);
+    int64_t sum_chunk_bytes = 0;
+    for (int i = 0; i < plan->Parallelism(); ++i) {
+      const int64_t in_offset = plan->InputChunkOffsetBytes(i);
+      const int64_t in_bytes = plan->InputChunkSizeBytes(i);
+      ASSERT_GE(input_base_offset_bytes + in_offset, 0);
+      ASSERT_LE(input_base_offset_bytes + in_offset + in_bytes,
+                static_cast<int64_t>(input.size()));
+      std::vector<char> in_slot(in_bytes, 0);
+      std::memcpy(in_slot.data(), input_base + in_offset, in_bytes);
+
+      const int64_t out_offset = plan->OutputChunkOffsetBytes(i);
+      const int64_t out_bytes = plan->OutputChunkSizeBytes(i);
+      EXPECT_EQ(out_offset, sum_chunk_bytes);
+      sum_chunk_bytes += out_bytes;
+      std::vector<char> out_slot(out_bytes, 0);
+      plan->ExecuteChunk(i, in_slot.data(), out_slot.data(),
+                         /*input_is_global=*/false,
+                         /*output_is_global=*/false);
+      std::memcpy(local_chunked_output.data() + out_offset, out_slot.data(),
+                  out_bytes);
+    }
+    EXPECT_EQ(sum_chunk_bytes, output_size_bytes);
+    EXPECT_EQ(local_chunked_output, expected_output);
+  }
 }
 
 TEST(TransposeTest, PackInt4_2D) {
   TestPackIntN(4, {4, 4}, {1, 0});
   TestPackIntN(4, {4, 4}, {0, 1});
   TestPackIntN(4, {16, 32}, {1, 0});
+  // Aligned partial output tiles on both the stride-1 PackIntN path and the
+  // microkernel path.
+  TestPackIntN(4, /*dims=*/{6, 8}, /*permutation=*/{0, 1},
+               /*input_tiling=*/{}, /*output_tiling=*/{4, 4},
+               /*input_striding=*/{}, /*num_threads=*/2,
+               TransposePlan::ChunkContiguity::kOutput,
+               /*expected_parallelism=*/2,
+               /*expected_inner_kernel_is_memcpy=*/true);
+  TestPackIntN(4, /*dims=*/{6, 8}, /*permutation=*/{1, 0},
+               /*input_tiling=*/{}, /*output_tiling=*/{8, 4},
+               /*input_striding=*/{}, /*num_threads=*/2,
+               TransposePlan::ChunkContiguity::kOutput,
+               /*expected_parallelism=*/2,
+               /*expected_inner_kernel_is_memcpy=*/false);
+  // Innermost output dimension splits only at byte boundaries (multiples of
+  // loop.inc = elements_per_byte on the memcpy path, or macro-block size on the
+  // microkernel path).
+  TestPackIntN(4, /*dims=*/{1, 6}, /*permutation=*/{0, 1},
+               /*input_tiling=*/{}, /*output_tiling=*/{},
+               /*input_striding=*/{}, /*num_threads=*/6,
+               TransposePlan::ChunkContiguity::kOutput,
+               /*expected_parallelism=*/3,
+               /*expected_inner_kernel_is_memcpy=*/true);
+  TestPackIntN(4, /*dims=*/{64, 1}, /*permutation=*/{1, 0},
+               /*input_tiling=*/{}, /*output_tiling=*/{},
+               /*input_striding=*/{2, 1}, /*num_threads=*/8,
+               TransposePlan::ChunkContiguity::kOutput,
+               /*expected_parallelism=*/4,
+               /*expected_inner_kernel_is_memcpy=*/false);
 }
 
 TEST(TransposeTest, PackInt4_3D) {
   TestPackIntN(4, {4, 6, 8}, {2, 0, 1});
   TestPackIntN(4, {8, 12, 16}, {1, 2, 0});
+  TestPackIntN(4, {4, 6, 8}, {1, 0, 2}, /*input_tiling=*/{},
+               /*output_tiling=*/{}, /*input_striding=*/{}, /*num_threads=*/1,
+               TransposePlan::ChunkContiguity::kNone,
+               /*expected_parallelism=*/1,
+               /*expected_inner_kernel_is_memcpy=*/true);
+  TestPackIntN(4, {5, 7, 8}, {1, 0, 2}, /*input_tiling=*/{},
+               /*output_tiling=*/{3, 4, 8}, /*input_striding=*/{},
+               /*num_threads=*/2, TransposePlan::ChunkContiguity::kOutput,
+               /*expected_parallelism=*/2,
+               /*expected_inner_kernel_is_memcpy=*/true);
+  // Tiled TPU layout s4[512, 128]{1,0:T(8,128)(8,1)E(4)} from row-major host
+  // buffer: splitting sublane dim 0 by packing factor 8 yields dims = {64, 128,
+  // 8} with input_striding = {8 * 128, 1, 128} (microkernel path).
+  TestPackIntN(4, /*dims=*/{64, 128, 8}, /*permutation=*/{0, 1, 2},
+               /*input_tiling=*/{}, /*output_tiling=*/{1, 128, 8},
+               /*input_striding=*/{8 * 128, 1, 128}, /*num_threads=*/4,
+               TransposePlan::ChunkContiguity::kOutput,
+               /*expected_parallelism=*/4,
+               /*expected_inner_kernel_is_memcpy=*/false);
+  // Tiled TPU transpose layout s4[1024, 512]{0,1:T(8,128)(8,1)E(4)} from
+  // row-major host buffer: splitting minor-most sublane dim 1 by packing factor
+  // 8 yields dims = {1024, 64, 8}, input_striding = {512, 8, 1}, permutation =
+  // {1, 0, 2}, output_tiling = {1, 128, 8} (stride-1 PackIntN path).
+  TestPackIntN(4, /*dims=*/{1024, 64, 8}, /*permutation=*/{1, 0, 2},
+               /*input_tiling=*/{}, /*output_tiling=*/{1, 128, 8},
+               /*input_striding=*/{512, 8, 1}, /*num_threads=*/4,
+               TransposePlan::ChunkContiguity::kOutput,
+               /*expected_parallelism=*/4,
+               /*expected_inner_kernel_is_memcpy=*/true);
 }
 
 TEST(TransposeTest, PackInt2_2D) {
   TestPackIntN(2, {4, 4}, {1, 0});
+  TestPackIntN(2, {4, 4}, {0, 1});
   TestPackIntN(2, {16, 32}, {1, 0});
 }
 
 TEST(TransposeTest, PackInt1_2D) {
   TestPackIntN(1, {8, 8}, {1, 0});
+  TestPackIntN(1, {8, 8}, {0, 1});
   TestPackIntN(1, {16, 32}, {1, 0});
 }
 
 TEST(TransposeTest, PackInt4_Unaligned) {
   TestPackIntN(4, {5, 5}, {1, 0});
   TestPackIntN(4, {15, 31}, {1, 0});
+  // Partial output tiles where OutputNumElems() (18) > num_elems_ (8).
+  TestPackIntN(4, {2, 4}, {1, 0}, /*input_tiling=*/{},
+               /*output_tiling=*/{3, 3});
+  // Output tiling where b_is_tiled_ is true and b_tiling_.back() == 1.
+  TestPackIntN(4, {24, 80}, {0, 1}, /*input_tiling=*/{},
+               /*output_tiling=*/{15, 1});
+  // Input tiling with aligned stride-1 tile size (uses vectorized kPackSubbyte
+  // across 2 chunks) and unaligned stride-1 tile size (routes to fallback
+  // pack).
+  TestPackIntN(4, {12, 7}, {1, 0}, /*input_tiling=*/{2, 4},
+               /*output_tiling=*/{}, /*input_striding=*/{}, /*num_threads=*/4,
+               TransposePlan::ChunkContiguity::kOutput,
+               /*expected_parallelism=*/2);
+  TestPackIntN(4, {12, 7}, {1, 0}, /*input_tiling=*/{3, 4});
+  // Non-unit innermost input stride with aligned tiled output uses the
+  // microkernel path without dummy-loop contiguity clamping.
+  TestPackIntN(4, {2, 8}, {0, 1}, /*input_tiling=*/{}, /*output_tiling=*/{2, 2},
+               /*input_striding=*/{16, 2}, /*num_threads=*/2,
+               TransposePlan::ChunkContiguity::kOutput,
+               /*expected_parallelism=*/2,
+               /*expected_inner_kernel_is_memcpy=*/false);
+  // Negative input strides and unaligned 1D output route to fallback pack.
+  TestPackIntN(4, {4, 4}, {0, 1}, /*input_tiling=*/{}, /*output_tiling=*/{},
+               /*input_striding=*/{-8, -1});
+  TestPackIntN(4, /*dims=*/{1, 5}, /*permutation=*/{0, 1});
 }
 
 TEST(TransposeTest, PackInt2_Unaligned) {
   TestPackIntN(2, {5, 5}, {1, 0});
   TestPackIntN(2, {15, 31}, {1, 0});
+  TestPackIntN(2, {2, 4}, {1, 0}, /*input_tiling=*/{},
+               /*output_tiling=*/{3, 3});
+  TestPackIntN(2, {24, 80}, {0, 1}, /*input_tiling=*/{},
+               /*output_tiling=*/{15, 1});
 }
 
 TEST(TransposeTest, PackInt1_Unaligned) {
   TestPackIntN(1, {5, 5}, {1, 0});
   TestPackIntN(1, {15, 31}, {1, 0});
+  TestPackIntN(1, {2, 4}, {1, 0}, /*input_tiling=*/{},
+               /*output_tiling=*/{3, 3});
+  TestPackIntN(1, {24, 80}, {0, 1}, /*input_tiling=*/{},
+               /*output_tiling=*/{15, 1});
+}
+
+TEST(TransposeTest, Sub64MiBContiguousChunking) {
+  // 16 MiB tensor: 2048 x 2048 int32s (4 bytes each).
+  // Request 8 chunks (2 MiB each). With ChunkContiguity::kOutput or kInput,
+  // ChooseParallelizationStrategy should honor num_threads = 8 without
+  // clamping to 1 chunk due to the 64 MiB default threshold.
+  const std::vector<int64_t> dims = {2048, 2048};
+  const std::vector<int64_t> permutation = {1, 0};
+  const std::vector<int64_t> tiling = {8, 128};
+  const int64_t num_elems = 2048 * 2048;
+
+  std::vector<int32_t> input(num_elems);
+  absl::c_iota(input, 1);
+
+  // 1. Test ChunkContiguity::kOutput (H2D pattern: global input, local output).
+  TransposePlan::Options h2d_options;
+  h2d_options.elem_size_in_bytes = sizeof(int32_t);
+  h2d_options.dims = dims;
+  h2d_options.permutation = permutation;
+  h2d_options.output_tiling = TransposePlan::Tiling{tiling};
+  h2d_options.num_threads = 8;
+  h2d_options.chunk_contiguity = TransposePlan::ChunkContiguity::kOutput;
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<TransposePlan> h2d_plan,
+                       TransposePlan::Create(h2d_options));
+  EXPECT_EQ(h2d_plan->Parallelism(), 8);
+
+  std::vector<int32_t> expected_tiled(num_elems, -1);
+  h2d_plan->Execute(input.data(), expected_tiled.data());
+
+  std::vector<int32_t> actual_tiled(num_elems, -1);
+  for (int i = 0; i < h2d_plan->Parallelism(); ++i) {
+    const int64_t chunk_bytes = h2d_plan->OutputChunkSizeBytes(i);
+    const int64_t chunk_offset = h2d_plan->OutputChunkOffsetBytes(i);
+    std::vector<uint8_t> slot_buf(chunk_bytes, 0);
+    h2d_plan->ExecuteChunk(i, input.data(), slot_buf.data(),
+                           /*input_is_global=*/true,
+                           /*output_is_global=*/false);
+    std::memcpy(reinterpret_cast<uint8_t*>(actual_tiled.data()) + chunk_offset,
+                slot_buf.data(), chunk_bytes);
+  }
+  EXPECT_EQ(actual_tiled, expected_tiled);
+
+  // 2. Test ChunkContiguity::kInput (D2H pattern: local input, global output).
+  TransposePlan::Options d2h_options;
+  d2h_options.elem_size_in_bytes = sizeof(int32_t);
+  d2h_options.dims = dims;
+  d2h_options.permutation = permutation;
+  d2h_options.input_tiling = TransposePlan::Tiling{tiling};
+  d2h_options.num_threads = 8;
+  d2h_options.chunk_contiguity = TransposePlan::ChunkContiguity::kInput;
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<TransposePlan> d2h_plan,
+                       TransposePlan::Create(d2h_options));
+  EXPECT_EQ(d2h_plan->Parallelism(), 8);
+
+  std::vector<int32_t> actual_delinearized(num_elems, -1);
+  for (int i = 0; i < d2h_plan->Parallelism(); ++i) {
+    const int64_t chunk_bytes = d2h_plan->InputChunkSizeBytes(i);
+    const int64_t chunk_offset = d2h_plan->InputChunkOffsetBytes(i);
+    std::vector<uint8_t> slot_buf(chunk_bytes, 0);
+    std::memcpy(
+        slot_buf.data(),
+        reinterpret_cast<const uint8_t*>(actual_tiled.data()) + chunk_offset,
+        chunk_bytes);
+    d2h_plan->ExecuteChunk(i, slot_buf.data(), actual_delinearized.data(),
+                           /*input_is_global=*/false,
+                           /*output_is_global=*/true);
+  }
+  EXPECT_EQ(actual_delinearized, input);
 }
 
 TEST(TransposeTest, InterleaveAndDeinterleaveSelection) {

From b1ccf9249140ca8c8eadc7c4b5d5038eb6ace0cb Mon Sep 17 00:00:00 2001
From: Raashish Aggarwal <94279692+raashish1601@users.noreply.github.com>
Date: Sat, 10 Oct 2026 13:03:51 +0530
Subject: [PATCH 0421/1000] Test an in-bounds tuple axis for var with an
 unknown rank

---
 tensorflow/python/ops/numpy_ops/np_array_ops_test.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index d0fdeb3a712..79ca155f4c0 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1034,6 +1034,15 @@ class ArrayMethodsTest(test.TestCase):
         (ValueError, errors_impl.InvalidArgumentError), out_of_bounds):
       var_fn_tuple(np.ones([2, 2], dtype=np.float32))
 
+    # In-bounds negative axes with an unknown rank are normalized at run time.
+    @def_function.function(
+        input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
+    def var_fn_tuple_in_bounds(a):
+      return np_array_ops.var(a, axis=(-1, 0), ddof=1)
+
+    x = np.arange(24, dtype=np.float32).reshape((2, 3, 4))
+    self.match(var_fn_tuple_in_bounds(x), np.var(x, axis=(-1, 0), ddof=1))
+
   def testProd(self):
 
     def run_test(arr, *args, **kwargs):

From 0f1f1dd08bc18e6134b17b312a55c410ffbbb7c3 Mon Sep 17 00:00:00 2001
From: abhirupone <abhirup.karmakar00@gmail.com>
Date: Sat, 10 Oct 2026 13:48:55 +0530
Subject: [PATCH 0422/1000] Fix RowPartition value_rowids nrows error message
 and restore operator stubs lint suppression

- Restore historical 'Expected nrows >= value_rowids[-1] + 1' error format in RowPartition.from_value_rowids to fix test regressions

- Update unit tests in row_partition_test.py to match restored error message

- Restore too-many-function-args pylint suppression on RaggedTensor overloaded operator stubs
---
 tensorflow/python/ops/ragged/ragged_tensor.py  |  2 ++
 tensorflow/python/ops/ragged/row_partition.py  | 18 +++++++++---------
 .../python/ops/ragged/row_partition_test.py    |  4 ++--
 3 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/tensorflow/python/ops/ragged/ragged_tensor.py b/tensorflow/python/ops/ragged/ragged_tensor.py
index 32ddb3586e6..cdd1901eda5 100644
--- a/tensorflow/python/ops/ragged/ragged_tensor.py
+++ b/tensorflow/python/ops/ragged/ragged_tensor.py
@@ -2200,6 +2200,7 @@ class RaggedTensor(
 
     return stub
 
+  # pylint: disable=too-many-function-args
   __getitem__ = _overloaded_operator("__getitem__")
   __ge__ = _overloaded_operator("__ge__")
   __gt__ = _overloaded_operator("__gt__")
@@ -2231,6 +2232,7 @@ class RaggedTensor(
   __truediv__ = _overloaded_operator("__truediv__")
   __rtruediv__ = _overloaded_operator("__rtruediv__")
   del _overloaded_operator
+  # pylint: enable=too-many-function-args
 
   # =============================================================================
   # Name Scope
diff --git a/tensorflow/python/ops/ragged/row_partition.py b/tensorflow/python/ops/ragged/row_partition.py
index c2e90df37df..89b3a33c8da 100644
--- a/tensorflow/python/ops/ragged/row_partition.py
+++ b/tensorflow/python/ops/ragged/row_partition.py
@@ -263,15 +263,21 @@ class RowPartition(composite_tensor.CompositeTensor):
       else:
         nrows = ops.convert_to_tensor(nrows, value_rowids.dtype, "nrows")
         const_nrows = tensor_util.constant_value(nrows)
-        if const_nrows is not None and const_nrows < 0:
-          raise ValueError("Expected nrows >= 0; got %d" % const_nrows)
+        if const_nrows is not None:
+          if const_nrows < 0:
+            raise ValueError("Expected nrows >= 0; got %d" % const_nrows)
+          const_rowids = tensor_util.constant_value(value_rowids)
+          if const_rowids is not None and const_rowids.size > 0:
+            if not const_nrows >= const_rowids[-1] + 1:
+              raise ValueError(
+                  "Expected nrows >= value_rowids[-1] + 1; got nrows=%d, "
+                  "value_rowids[-1]=%d" % (const_nrows, const_rowids[-1]))
 
       value_rowids.shape.assert_has_rank(1)
       nrows.shape.assert_has_rank(0)
 
       if validate:
         value_rowids_const = tensor_util.constant_value(value_rowids)
-        nrows_const = tensor_util.constant_value(nrows)
         if value_rowids_const is not None:
           if len(value_rowids_const) > 0 and value_rowids_const[0] < 0:
             raise ValueError(
@@ -288,12 +294,6 @@ class RowPartition(composite_tensor.CompositeTensor):
                 f"({value_rowids_const[bad_idx]}) > "
                 f"value_rowids[{bad_idx + 1}] "
                 f"({value_rowids_const[bad_idx + 1]})")
-          if nrows_const is not None and len(value_rowids_const) > 0:
-            if value_rowids_const[-1] >= nrows_const:
-              raise ValueError(
-                  "Arguments to from_value_rowids do not form a valid "
-                  "RowPartition: value_rowids[-1] must be < nrows, "
-                  f"got {value_rowids_const[-1]} >= {nrows_const}")
         msg = ("Arguments to from_value_rowids do not form a valid "
                "RowPartition")
         checks = [
diff --git a/tensorflow/python/ops/ragged/row_partition_test.py b/tensorflow/python/ops/ragged/row_partition_test.py
index a40a0c8bd49..389b8008bc9 100644
--- a/tensorflow/python/ops/ragged/row_partition_test.py
+++ b/tensorflow/python/ops/ragged/row_partition_test.py
@@ -956,7 +956,7 @@ class RowPartitionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
 
   def testFromValueRowidsStaticValidationNrowsExceeded(self):
     with self.assertRaisesRegex(
-        ValueError, r"value_rowids\[-1\] must be < nrows"):
+        ValueError, r"Expected nrows >= value_rowids\[-1\] \+ 1"):
       RowPartition.from_value_rowids([0, 2], nrows=2)
 
   def testFromRowLengthsStaticValidationNegative(self):
@@ -1009,7 +1009,7 @@ class RowPartitionTest(test_util.TensorFlowTestCase, parameterized.TestCase):
       return RowPartition.from_value_rowids([0, 2], nrows=2)
 
     with self.assertRaisesRegex(
-        ValueError, r"value_rowids\[-1\] must be < nrows"):
+        ValueError, r"Expected nrows >= value_rowids\[-1\] \+ 1"):
       fn(constant_op.constant(0))
 
   def testFromRowLengthsStaticValidationUnderJit(self):

From 1e83e7afeb2ad2418ffb3666fd435d39342dcb85 Mon Sep 17 00:00:00 2001
From: wondertx <wondertx@gmail.com>
Date: Sat, 10 Oct 2026 14:11:26 +0800
Subject: [PATCH 0423/1000] Fix NcclBroadcastRecv int32 overflow

---
 tensorflow/core/kernels/nccl_ops.cc  |  4 +-
 tensorflow/core/nccl/nccl_manager.cc |  2 +-
 tensorflow/core/nccl/nccl_rewrite.cc | 14 ++++--
 tensorflow/core/ops/BUILD            |  1 +
 tensorflow/core/ops/nccl_ops.cc      |  3 +-
 tensorflow/core/ops/nccl_ops_test.cc | 68 ++++++++++++++++++++++++++++
 6 files changed, 83 insertions(+), 9 deletions(-)
 create mode 100644 tensorflow/core/ops/nccl_ops_test.cc

diff --git a/tensorflow/core/kernels/nccl_ops.cc b/tensorflow/core/kernels/nccl_ops.cc
index 6a1a3f880ad..3f8001bdbed 100644
--- a/tensorflow/core/kernels/nccl_ops.cc
+++ b/tensorflow/core/kernels/nccl_ops.cc
@@ -28,6 +28,7 @@ limitations under the License.
 #endif
 #endif
 #include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor_util.h"
 #include "tensorflow/core/nccl/nccl_manager.h"
 
 namespace tensorflow {
@@ -238,8 +239,7 @@ class NcclBroadcastRecvKernel : public NcclAsyncOpBase {
   void ComputeAsync(OpKernelContext* c, DoneCallback done) override {
     const Tensor& shape_t = c->input(0);
     TensorShape shape;
-    OP_REQUIRES_OK_ASYNC(
-        c, TensorShapeUtils::MakeShape(shape_t.vec<int32_t>(), &shape), done);
+    OP_REQUIRES_OK_ASYNC(c, tensor::MakeShape(shape_t, &shape), done);
     Tensor* output;
     OP_REQUIRES_OK_ASYNC(c, c->allocate_output(0, shape, &output), done);
 
diff --git a/tensorflow/core/nccl/nccl_manager.cc b/tensorflow/core/nccl/nccl_manager.cc
index e7271cf82b6..d18f05ae0c1 100644
--- a/tensorflow/core/nccl/nccl_manager.cc
+++ b/tensorflow/core/nccl/nccl_manager.cc
@@ -790,7 +790,7 @@ void NcclManager::LoopKernelLaunches(NcclStream* nccl_stream) {
       case kBroadcast: {
         const void* sendbuff = nullptr;
         void* recvbuff = nullptr;
-        int num_elements = -1;
+        int64_t num_elements = -1;
         if (p->input) {
           sendbuff = p->input->tensor_data().data();
           num_elements = p->input->NumElements();
diff --git a/tensorflow/core/nccl/nccl_rewrite.cc b/tensorflow/core/nccl/nccl_rewrite.cc
index 0fdfa45d830..a384f372d13 100644
--- a/tensorflow/core/nccl/nccl_rewrite.cc
+++ b/tensorflow/core/nccl/nccl_rewrite.cc
@@ -110,9 +110,9 @@ absl::Status ReplaceReduce(Graph* graph, Node* node) {
 
 TensorProto TensorFromShape(const TensorShapeProto& shape) {
   TensorProto result;
-  result.set_dtype(DT_INT32);
+  result.set_dtype(DT_INT64);
   for (const auto& dim : shape.dim()) {
-    result.add_int_val(dim.size());
+    result.add_int64_val(dim.size());
   }
   result.mutable_tensor_shape()->add_dim()->set_size(shape.dim_size());
   return result;
@@ -213,7 +213,7 @@ absl::Status ReplaceBroadcast(Graph* graph, Node* node) {
   Node* shape_node = nullptr;
   if (!is_fully_defined) {
     NodeBuilder shape_builder(shape_name, "Shape");
-    shape_builder.Input(in_node).Attr("out_type", DT_INT32).Attr("T", dtype);
+    shape_builder.Input(in_node).Attr("out_type", DT_INT64).Attr("T", dtype);
     TF_RETURN_IF_ERROR(shape_builder.Finalize(graph, &shape_node));
     shape_node->set_assigned_device_name_index(send_dev);
   }
@@ -226,14 +226,18 @@ absl::Status ReplaceBroadcast(Graph* graph, Node* node) {
     int recv_index = recv_index_map[recv_dev];
     if (is_fully_defined) {
       // If the shape is fully defined, define one const node per device.
-      NodeBuilder shape_builder(absl::StrCat(shape_name, recv_index), "Const");
-      shape_builder.Attr("value", tensor_proto).Attr("dtype", DT_INT32);
+      // HostConst keeps the int64 shape in host memory, which is where
+      // _NcclBroadcastRecv reads it from.
+      NodeBuilder shape_builder(absl::StrCat(shape_name, recv_index),
+                                "HostConst");
+      shape_builder.Attr("value", tensor_proto).Attr("dtype", DT_INT64);
       TF_RETURN_IF_ERROR(shape_builder.Finalize(graph, &shape_node));
       shape_node->set_assigned_device_name_index(recv_dev);
     }
     Node* recv_node;
     TF_RETURN_IF_ERROR(
         make_builder("_NcclBroadcastRecv", absl::StrCat("Recv_", recv_index))
+            .Attr("Tshape", DT_INT64)
             .Input(shape_node)
             .Finalize(graph, &recv_node));
     recv_node->set_assigned_device_name_index(recv_dev);
diff --git a/tensorflow/core/ops/BUILD b/tensorflow/core/ops/BUILD
index 8b781e73c21..222d9f8ad4d 100644
--- a/tensorflow/core/ops/BUILD
+++ b/tensorflow/core/ops/BUILD
@@ -557,6 +557,7 @@ tf_cc_test(
         "linalg_ops_test.cc",
         "lookup_ops_test.cc",
         "math_ops_test.cc",
+        "nccl_ops_test.cc",
         "nn_ops_test.cc",
         "parsing_ops_test.cc",
         "random_ops_test.cc",
diff --git a/tensorflow/core/ops/nccl_ops.cc b/tensorflow/core/ops/nccl_ops.cc
index 47d606917e9..7ab542dd793 100644
--- a/tensorflow/core/ops/nccl_ops.cc
+++ b/tensorflow/core/ops/nccl_ops.cc
@@ -121,9 +121,10 @@ shared_name: Identifier that is shared between ops of the same broadcast.
     )doc");
 
 REGISTER_OP("_NcclBroadcastRecv")
-    .Input("shape: int32")
+    .Input("shape: Tshape")
     .Output("output: T")
     .Attr("T: {half, float, float64, int32, int64}")
+    .Attr("Tshape: {int32, int64} = DT_INT32")
     .Attr("num_devices: int")
     .Attr("shared_name: string")
     .SetIsStateful()
diff --git a/tensorflow/core/ops/nccl_ops_test.cc b/tensorflow/core/ops/nccl_ops_test.cc
new file mode 100644
index 00000000000..70266a8a805
--- /dev/null
+++ b/tensorflow/core/ops/nccl_ops_test.cc
@@ -0,0 +1,68 @@
+/* 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.
+==============================================================================*/
+
+#include <cstdint>
+
+#include "tensorflow/core/framework/node_def_builder.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/shape_inference_testutil.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/framework/types.pb.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+
+namespace {
+
+void BuildNcclBroadcastRecv(DataType shape_type, ShapeInferenceTestOp* op) {
+  TF_ASSERT_OK(NodeDefBuilder("test", "_NcclBroadcastRecv")
+                   .Input("shape", 0, shape_type)
+                   .Attr("T", DT_FLOAT)
+                   .Attr("num_devices", 2)
+                   .Attr("shared_name", "s")
+                   .Finalize(&op->node_def));
+}
+
+}  // namespace
+
+TEST(NcclOpsTest, NcclBroadcastRecv_ShapeFn_Int32) {
+  ShapeInferenceTestOp op("_NcclBroadcastRecv");
+  BuildNcclBroadcastRecv(DT_INT32, &op);
+
+  INFER_OK(op, "[3]", "[?,?,?]");
+  INFER_ERROR("Shape must be rank 1", op, "[1,2]");
+
+  Tensor shape_t = test::AsTensor<int32_t>({2, 3});
+  op.input_tensors.resize(1);
+  op.input_tensors[0] = &shape_t;
+  INFER_OK(op, "[2]", "[2,3]");
+}
+
+TEST(NcclOpsTest, NcclBroadcastRecv_ShapeFn_Int64) {
+  ShapeInferenceTestOp op("_NcclBroadcastRecv");
+  BuildNcclBroadcastRecv(DT_INT64, &op);
+
+  INFER_OK(op, "[3]", "[?,?,?]");
+
+  // A dimension that does not fit in int32 must be preserved.
+  Tensor shape_t = test::AsTensor<int64_t>({int64_t{1} << 32, 2});
+  op.input_tensors.resize(1);
+  op.input_tensors[0] = &shape_t;
+  INFER_OK(op, "[2]", "[4294967296,2]");
+}
+
+}  // namespace tensorflow

From edc0a1dc3bf72008351b9a4e97efbb44c9240b8d Mon Sep 17 00:00:00 2001
From: PrabinDevkota <prabindev789@gmail.com>
Date: Sat, 10 Oct 2026 19:21:12 +0530
Subject: [PATCH 0424/1000] Reject rank above 254 inside TensorShape instead of
 each op.

---
 tensorflow/core/kernels/shape_ops_test.cc | 56 -----------------------
 1 file changed, 56 deletions(-)

diff --git a/tensorflow/core/kernels/shape_ops_test.cc b/tensorflow/core/kernels/shape_ops_test.cc
index 2af70a685bf..1cee9c25366 100644
--- a/tensorflow/core/kernels/shape_ops_test.cc
+++ b/tensorflow/core/kernels/shape_ops_test.cc
@@ -13,21 +13,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include <cstdint>
 #include <functional>
 #include <memory>
-#include <vector>
 
-#include <gmock/gmock.h>
-#include <gtest/gtest.h>
-#include "absl/status/status.h"
 #include "tensorflow/core/common_runtime/kernel_benchmark_testlib.h"
 #include "tensorflow/core/framework/allocator.h"
-#include "tensorflow/core/framework/fake_input.h"
-#include "tensorflow/core/framework/node_def_builder.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/tensor.h"
-#include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
 #include "tensorflow/core/graph/algorithm.h"
@@ -36,60 +28,12 @@ limitations under the License.
 #include "tensorflow/core/kernels/ops_testutil.h"
 #include "tensorflow/core/kernels/ops_util.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
-#include "tensorflow/core/platform/status_matchers.h"
 #include "tensorflow/core/platform/test.h"
 #include "tensorflow/core/platform/test_benchmark.h"
 
 namespace tensorflow {
 namespace {
 
-using ::tensorflow::testing::StatusIs;
-using ::testing::HasSubstr;
-
-absl::Status MakeZeroElementShape(int rank, TensorShape* shape) {
-  std::vector<int64_t> dims(rank, 1);
-  if (rank > 0) {
-    dims[0] = 0;
-  }
-  return TensorShape::BuildTensorShape(dims, shape);
-}
-
-class ExpandDimsOpTest : public OpsTestBase {
- protected:
-  void MakeOp() {
-    TF_ASSERT_OK(NodeDefBuilder("expand", "ExpandDims")
-                     .Input(FakeInput(DT_FLOAT))
-                     .Input(FakeInput(DT_INT32))
-                     .Finalize(node_def()));
-    TF_ASSERT_OK(InitOp());
-  }
-};
-
-TEST_F(ExpandDimsOpTest, ExpandingMaxRankSucceeds) {
-  MakeOp();
-  TensorShape input_shape;
-  TF_ASSERT_OK(
-      MakeZeroElementShape(TensorShape::MaxDimensions() - 1, &input_shape));
-  AddInput<float>(input_shape, [](int) { return 0.0f; });
-  AddInputFromList<int32_t>(TensorShape({}), {0});
-
-  TF_ASSERT_OK(RunOpKernel());
-  EXPECT_EQ(GetOutput(0)->dims(), TensorShape::MaxDimensions());
-  EXPECT_EQ(GetOutput(0)->dim_size(0), 1);
-}
-
-TEST_F(ExpandDimsOpTest, ExpandingBeyondMaxRankFails) {
-  MakeOp();
-  TensorShape input_shape;
-  TF_ASSERT_OK(MakeZeroElementShape(TensorShape::MaxDimensions(), &input_shape));
-  AddInput<float>(input_shape, [](int) { return 0.0f; });
-  AddInputFromList<int32_t>(TensorShape({}), {0});
-
-  EXPECT_THAT(RunOpKernel(),
-              StatusIs(absl::StatusCode::kInvalidArgument,
-                       HasSubstr("maximum supported rank")));
-}
-
 static void BM_ExpandDims(::testing::benchmark::State& state) {
   Graph* g = new Graph(OpRegistry::Global());
 

From 800a62599e184ce4dc76421abc89f6f4c37b8dc6 Mon Sep 17 00:00:00 2001
From: PrabinDevkota <prabindev789@gmail.com>
Date: Sat, 10 Oct 2026 19:24:05 +0530
Subject: [PATCH 0425/1000] Reject rank above 254 inside TensorShape instead of
 each op.

---
 tensorflow/compiler/tf2xla/kernels/pack_op.cc |   7 --
 tensorflow/core/framework/tensor_shape.cc     |   2 +-
 .../core/framework/tensor_shape_test.cc       |   2 +-
 tensorflow/core/kernels/BUILD                 |  24 ----
 tensorflow/core/kernels/pack_op.cc            |   8 --
 tensorflow/core/kernels/pack_op_test.cc       | 104 ------------------
 tensorflow/core/kernels/shape_ops.h           |  10 +-
 tensorflow/core/util/strided_slice_op.cc      |  15 +--
 tensorflow/core/util/strided_slice_op_test.cc |  62 -----------
 .../kernel_tests/array_ops/array_ops_test.py  |   2 +-
 10 files changed, 6 insertions(+), 230 deletions(-)
 delete mode 100644 tensorflow/core/kernels/pack_op_test.cc

diff --git a/tensorflow/compiler/tf2xla/kernels/pack_op.cc b/tensorflow/compiler/tf2xla/kernels/pack_op.cc
index 3b1e83e5c4e..6a33ef26499 100644
--- a/tensorflow/compiler/tf2xla/kernels/pack_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/pack_op.cc
@@ -61,13 +61,6 @@ class PackOp : public XlaOpKernel {
                                         -expanded_num_dims, ", ",
                                         expanded_num_dims, ")"));
 
-    OP_REQUIRES(
-        ctx, shapes[0].dims() < TensorShape::MaxDimensions(),
-        errors::InvalidArgument(
-            "Cannot pack tensors of rank ", shapes[0].dims(),
-            ": packing would exceed the maximum supported rank of ",
-            TensorShape::MaxDimensions(), "."));
-
     std::vector<xla::XlaOp> reshaped_inputs(num);
 
     TensorShape child_shape(shapes[0]);
diff --git a/tensorflow/core/framework/tensor_shape.cc b/tensorflow/core/framework/tensor_shape.cc
index 6ae9fd71218..d24160230ea 100644
--- a/tensorflow/core/framework/tensor_shape.cc
+++ b/tensorflow/core/framework/tensor_shape.cc
@@ -565,7 +565,7 @@ absl::Status TensorShapeBase<Shape>::InsertDimWithStatus(int d, int64_t size) {
         "The insertion index must be at most ", dims(), " got ", d));
   }
   if (TF_PREDICT_FALSE(dims() >= MaxDimensions())) {
-    return absl::InternalError(absl::StrCat(
+    return absl::InvalidArgumentError(absl::StrCat(
         "Shape has ", dims(), " dimensions which is the maximum allowed"));
   }
 
diff --git a/tensorflow/core/framework/tensor_shape_test.cc b/tensorflow/core/framework/tensor_shape_test.cc
index 65262fd2ffa..0d46fff212b 100644
--- a/tensorflow/core/framework/tensor_shape_test.cc
+++ b/tensorflow/core/framework/tensor_shape_test.cc
@@ -263,7 +263,7 @@ TEST(TensorShapeTest, InsertDimWithStatusWithTooManyDims) {
   }
   EXPECT_THAT(s.InsertDimWithStatus(1, 1),
               absl_testing::StatusIs(
-                  error::Code::INTERNAL,
+                  error::Code::INVALID_ARGUMENT,
                   ::testing::ContainsRegex(
                       "Shape has.*dimensions which is the maximum allowed")));
 }
diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 744a0672cbe..ee57a0ede97 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -1138,27 +1138,6 @@ tf_kernel_library(
     ],
 )
 
-tf_cc_test(
-    name = "pack_op_test",
-    size = "small",
-    srcs = ["pack_op_test.cc"],
-    deps = [
-        ":ops_testutil",
-        ":ops_util",
-        ":pack_op",
-        "//tensorflow/core:core_cpu",
-        "//tensorflow/core:framework",
-        "//tensorflow/core:lib",
-        "//tensorflow/core:protos_all_cc",
-        "//tensorflow/core:test",
-        "//tensorflow/core:test_main",
-        "//tensorflow/core:testlib",
-        "//tensorflow/core/platform:status_matchers",
-        "@com_google_absl//absl/status",
-        "@com_google_googletest//:gtest",
-    ],
-)
-
 tf_kernel_library(
     name = "pad_op",
     prefix = "pad_op",
@@ -2416,9 +2395,6 @@ tf_cc_test(
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
         "//tensorflow/core:testlib",
-        "//tensorflow/core/platform:status_matchers",
-        "@com_google_absl//absl/status",
-        "@com_google_googletest//:gtest",
     ],
 )
 
diff --git a/tensorflow/core/kernels/pack_op.cc b/tensorflow/core/kernels/pack_op.cc
index 3dd0be89f9c..959aa01f9fb 100644
--- a/tensorflow/core/kernels/pack_op.cc
+++ b/tensorflow/core/kernels/pack_op.cc
@@ -27,7 +27,6 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_types.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/kernels/concat_lib.h"
-#include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/platform/types.h"
 
@@ -68,13 +67,6 @@ class PackOp : public OpKernel {
                     "axis = ", axis_, " not in [", -expanded_num_dims, ", ",
                     expanded_num_dims, ")")));
 
-    OP_REQUIRES(
-        c, first_input.dims() < TensorShape::MaxDimensions(),
-        errors::InvalidArgument(
-            "Cannot pack tensors of rank ", first_input.dims(),
-            ": packing would exceed the maximum supported rank of ",
-            TensorShape::MaxDimensions(), "."));
-
     TensorShape output_shape(first_input.shape());
     OP_REQUIRES_OK(c, output_shape.InsertDimWithStatus(axis, num));
 
diff --git a/tensorflow/core/kernels/pack_op_test.cc b/tensorflow/core/kernels/pack_op_test.cc
deleted file mode 100644
index 288ba4beb85..00000000000
--- a/tensorflow/core/kernels/pack_op_test.cc
+++ /dev/null
@@ -1,104 +0,0 @@
-/* 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.
-==============================================================================*/
-
-#include <cstdint>
-#include <vector>
-
-#include <gmock/gmock.h>
-#include <gtest/gtest.h>
-#include "absl/status/status.h"
-#include "tensorflow/core/framework/fake_input.h"
-#include "tensorflow/core/framework/node_def_builder.h"
-#include "tensorflow/core/framework/tensor.h"
-#include "tensorflow/core/framework/tensor_shape.h"
-#include "tensorflow/core/framework/types.h"
-#include "tensorflow/core/framework/types.pb.h"
-#include "tensorflow/core/kernels/ops_testutil.h"
-#include "tensorflow/core/kernels/ops_util.h"
-#include "tensorflow/core/lib/core/status_test_util.h"
-#include "tensorflow/core/platform/status_matchers.h"
-#include "tensorflow/core/platform/test.h"
-
-namespace tensorflow {
-namespace {
-
-using ::tensorflow::testing::StatusIs;
-using ::testing::HasSubstr;
-
-absl::Status MakeZeroElementShape(int rank, TensorShape* shape) {
-  std::vector<int64_t> dims(rank, 1);
-  if (rank > 0) {
-    dims[0] = 0;
-  }
-  return TensorShape::BuildTensorShape(dims, shape);
-}
-
-class PackOpTest : public OpsTestBase {
- protected:
-  void MakeOp() {
-    TF_ASSERT_OK(NodeDefBuilder("pack", "Pack")
-                     .Input(FakeInput(1, DT_FLOAT))
-                     .Attr("N", 1)
-                     .Attr("T", DT_FLOAT)
-                     .Attr("axis", 0)
-                     .Finalize(node_def()));
-    TF_ASSERT_OK(InitOp());
-  }
-};
-
-TEST_F(PackOpTest, PackingMaxRankSucceeds) {
-  MakeOp();
-  TensorShape input_shape;
-  TF_ASSERT_OK(
-      MakeZeroElementShape(TensorShape::MaxDimensions() - 1, &input_shape));
-  AddInput<float>(input_shape, [](int) { return 0.0f; });
-
-  TF_ASSERT_OK(RunOpKernel());
-  EXPECT_EQ(GetOutput(0)->dims(), TensorShape::MaxDimensions());
-  EXPECT_EQ(GetOutput(0)->dim_size(0), 1);
-}
-
-TEST_F(PackOpTest, PackingMultipleTensorsMaxRankSucceeds) {
-  TF_ASSERT_OK(NodeDefBuilder("pack", "Pack")
-                   .Input(FakeInput(2, DT_FLOAT))
-                   .Attr("N", 2)
-                   .Attr("T", DT_FLOAT)
-                   .Attr("axis", 0)
-                   .Finalize(node_def()));
-  TF_ASSERT_OK(InitOp());
-  TensorShape input_shape;
-  TF_ASSERT_OK(
-      MakeZeroElementShape(TensorShape::MaxDimensions() - 1, &input_shape));
-  AddInput<float>(input_shape, [](int) { return 0.0f; });
-  AddInput<float>(input_shape, [](int) { return 0.0f; });
-
-  TF_ASSERT_OK(RunOpKernel());
-  EXPECT_EQ(GetOutput(0)->dims(), TensorShape::MaxDimensions());
-  EXPECT_EQ(GetOutput(0)->dim_size(0), 2);
-}
-
-TEST_F(PackOpTest, PackingBeyondMaxRankFails) {
-  MakeOp();
-  TensorShape input_shape;
-  TF_ASSERT_OK(MakeZeroElementShape(TensorShape::MaxDimensions(), &input_shape));
-  AddInput<float>(input_shape, [](int) { return 0.0f; });
-
-  EXPECT_THAT(RunOpKernel(),
-              StatusIs(absl::StatusCode::kInvalidArgument,
-                       HasSubstr("maximum supported rank")));
-}
-
-}  // namespace
-}  // namespace tensorflow
diff --git a/tensorflow/core/kernels/shape_ops.h b/tensorflow/core/kernels/shape_ops.h
index 2ab6c2df5cd..28c73a15f4f 100644
--- a/tensorflow/core/kernels/shape_ops.h
+++ b/tensorflow/core/kernels/shape_ops.h
@@ -169,14 +169,8 @@ class ExpandDimsOp : public OpKernel {
       dim = std::min<Tdim>(dim + input_dims + 1, input_dims);
     }
 
-    OP_REQUIRES(
-        ctx, input_dims < TensorShape::MaxDimensions(),
-        errors::InvalidArgument(
-            "Cannot expand a tensor of rank ", input_dims,
-            ": expanding would exceed the maximum supported rank of ",
-            TensorShape::MaxDimensions(), "."));
-
-    // Compute new shape with an additional dimension.
+    // Compute new shape with an additional dimension. BuildTensorShape
+    // rejects a result whose rank exceeds TensorShape::MaxDimensions().
     absl::InlinedVector<int64_t, 8> output_shape_vec(input_dims + 1);
     for (int64_t i = 0; i < dim; ++i) {
       output_shape_vec[i] = input_shape.dim_size(i);
diff --git a/tensorflow/core/util/strided_slice_op.cc b/tensorflow/core/util/strided_slice_op.cc
index fff42cb430b..ca337966301 100644
--- a/tensorflow/core/util/strided_slice_op.cc
+++ b/tensorflow/core/util/strided_slice_op.cc
@@ -21,7 +21,6 @@ limitations under the License.
 #include <utility>
 
 #include "tensorflow/core/framework/bounds_check.h"
-#include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/status.h"
 
 namespace tensorflow {
@@ -410,19 +409,7 @@ absl::Status ValidateStridedSliceOp(
   // new_axis will increase dimension by 1 (with a one-size dimension)
   // slices like foo[3,...] will reduce dimension by 1.
   // This cannot be done earlier, because it depends on Step 3.
-  int output_rank = 0;
-  for (int32_t gather_index : dense_spec.final_shape_gather_indices) {
-    if (gather_index != kShrinkAxis) {
-      ++output_rank;
-    }
-  }
-  if (output_rank > TensorShape::MaxDimensions()) {
-    return errors::InvalidArgument(
-        "StridedSlice would produce a tensor of rank ", output_rank,
-        ", which exceeds the maximum supported rank of ",
-        TensorShape::MaxDimensions(), ".");
-  }
-
+  // AddDimWithStatus rejects a result past TensorShape::MaxDimensions().
   final_shape->Clear();
   if (shape_spec != nullptr) {
     shape_spec->output_to_sparse_mapping.clear();
diff --git a/tensorflow/core/util/strided_slice_op_test.cc b/tensorflow/core/util/strided_slice_op_test.cc
index cd666e7b275..ee96537d009 100644
--- a/tensorflow/core/util/strided_slice_op_test.cc
+++ b/tensorflow/core/util/strided_slice_op_test.cc
@@ -19,11 +19,9 @@ limitations under the License.
 #include <ostream>
 #include <tuple>
 #include <utility>
-#include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
-#include "absl/status/status.h"
 #include "xla/tsl/lib/core/status_test_util.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
@@ -358,30 +356,6 @@ TensorShape AsTensorShape(absl::Span<const int64_t> dim_sizes) {
   return out;
 }
 
-TensorShape ZeroElementRank(int rank) {
-  std::vector<int64_t> dims(rank, 1);
-  if (rank > 0) {
-    dims[0] = 0;
-  }
-  return AsTensorShape(dims);
-}
-
-absl::Status ValidateNewAxis(const TensorShape& input_shape,
-                             TensorShape* processing_shape,
-                             TensorShape* final_shape) {
-  Tensor begin_tensor = test::AsTensor<int32_t>({0});
-  Tensor end_tensor = test::AsTensor<int32_t>({0});
-  Tensor strides_tensor = test::AsTensor<int32_t>({1});
-  bool is_identity, is_simple_slice, slice_dim0;
-  IntVector begin, end, strides;
-  return ValidateStridedSliceOp(
-      &begin_tensor, &end_tensor, strides_tensor, input_shape,
-      /*begin_mask_spec=*/0, /*end_mask_spec=*/0, /*ellipsis_mask=*/0,
-      /*new_axis_mask=*/1, /*shrink_axis_mask=*/0, processing_shape,
-      final_shape, &is_identity, &is_simple_slice, &slice_dim0, &begin, &end,
-      &strides);
-}
-
 TEST(ValidateStridedSliceOpTest, BasicStride) {
   Tensor begin_tensor = test::AsTensor<int32_t>({1, 1});
   Tensor end_tensor = test::AsTensor<int32_t>({7, 7});
@@ -797,41 +771,5 @@ TEST(ValidateStridedSliceOpTest, PartialInputShape) {
       &is_simple_slice, &slice_dim0, &begin, &end, &strides, &shape_spec));
 }
 
-TEST(ValidateStridedSliceOpTest, NewAxisAtMaxRankSucceeds) {
-  TensorShape processing_shape, final_shape;
-  TF_EXPECT_OK(ValidateNewAxis(ZeroElementRank(TensorShape::MaxDimensions() - 1),
-                               &processing_shape, &final_shape));
-  EXPECT_EQ(final_shape.dims(), TensorShape::MaxDimensions());
-  EXPECT_EQ(final_shape.dim_size(0), 1);
-}
-
-TEST(ValidateStridedSliceOpTest, NewAxisBeyondMaxRankFails) {
-  TensorShape processing_shape, final_shape;
-  EXPECT_THAT(
-      ValidateNewAxis(ZeroElementRank(TensorShape::MaxDimensions()),
-                      &processing_shape, &final_shape),
-      absl_testing::StatusIs(tsl::error::Code::INVALID_ARGUMENT,
-                             ::testing::HasSubstr("maximum supported rank")));
-}
-
-TEST(ValidateStridedSliceOpTest, ManyNewAxesBeyondMaxRankFails) {
-  // Mirrors the public StridedSlice repro: rank 225 plus 30 new axes.
-  Tensor begin_tensor = test::AsTensor<int32_t>(std::vector<int32_t>(31, 0));
-  Tensor end_tensor = test::AsTensor<int32_t>(std::vector<int32_t>(31, 0));
-  Tensor strides_tensor = test::AsTensor<int32_t>(std::vector<int32_t>(31, 1));
-  TensorShape processing_shape, final_shape;
-  bool is_identity, is_simple_slice, slice_dim0;
-  IntVector begin, end, strides;
-  EXPECT_THAT(
-      ValidateStridedSliceOp(
-          &begin_tensor, &end_tensor, strides_tensor, ZeroElementRank(225),
-          /*begin_mask_spec=*/0, /*end_mask_spec=*/0,
-          /*ellipsis_mask=*/(1 << 30), /*new_axis_mask=*/(1 << 30) - 1,
-          /*shrink_axis_mask=*/0, &processing_shape, &final_shape,
-          &is_identity, &is_simple_slice, &slice_dim0, &begin, &end, &strides),
-      absl_testing::StatusIs(tsl::error::Code::INVALID_ARGUMENT,
-                             ::testing::HasSubstr("maximum supported rank")));
-}
-
 }  // namespace
 }  // namespace tensorflow
diff --git a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
index f2e01af0bfd..58951a7a22b 100644
--- a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
@@ -798,7 +798,7 @@ class StridedSliceTest(test_util.TensorFlowTestCase):
         [0] + [1] * 224,
         constant_op.constant(0.0, dtype=dtypes.float32))
     with self.assertRaisesRegex(
-        (errors.InvalidArgumentError, ValueError), "maximum supported rank"):
+        (errors.InvalidArgumentError, ValueError), "Too many dimensions"):
       # Public slice syntax, not the generated op wrapper. evaluate() covers
       # graph mode, where shape inference may defer the error until execution.
       sliced = value[tuple([array_ops.newaxis] * 30 + [Ellipsis])]

From 6cbc1e46c80e5ce5d72184b3ccc7ddb2a691c689 Mon Sep 17 00:00:00 2001
From: Blake Hechtman <blakehechtman@google.com>
Date: Sat, 10 Oct 2026 08:38:19 -0700
Subject: [PATCH 0426/1000] [XLA] Remove TuplePointsToAnalysis and
 LogicalBufferAnalysis

Migrate HloRematerialization to HloDataflowAnalysis, and delete TuplePointsToAnalysis and LogicalBufferAnalysis along with their tests.

PiperOrigin-RevId: 997111694
---
 third_party/xla/xla/hlo/analysis/BUILD        |   85 --
 .../hlo/analysis/logical_buffer_analysis.cc   |  344 -----
 .../hlo/analysis/logical_buffer_analysis.h    |  142 --
 .../analysis/logical_buffer_analysis_test.cc  |  833 -----------
 .../hlo/analysis/tuple_points_to_analysis.cc  | 1105 --------------
 .../hlo/analysis/tuple_points_to_analysis.h   |  435 ------
 .../analysis/tuple_points_to_analysis_test.cc | 1328 -----------------
 .../xla/xla/hlo/transforms/simplifiers/BUILD  |    3 +-
 .../simplifiers/hlo_rematerialization.cc      |  183 +--
 .../simplifiers/hlo_rematerialization.h       |   18 +-
 10 files changed, 102 insertions(+), 4374 deletions(-)
 delete mode 100644 third_party/xla/xla/hlo/analysis/logical_buffer_analysis.cc
 delete mode 100644 third_party/xla/xla/hlo/analysis/logical_buffer_analysis.h
 delete mode 100644 third_party/xla/xla/hlo/analysis/logical_buffer_analysis_test.cc
 delete mode 100644 third_party/xla/xla/hlo/analysis/tuple_points_to_analysis.cc
 delete mode 100644 third_party/xla/xla/hlo/analysis/tuple_points_to_analysis.h
 delete mode 100644 third_party/xla/xla/hlo/analysis/tuple_points_to_analysis_test.cc

diff --git a/third_party/xla/xla/hlo/analysis/BUILD b/third_party/xla/xla/hlo/analysis/BUILD
index 246e764b8b6..72797f003cd 100644
--- a/third_party/xla/xla/hlo/analysis/BUILD
+++ b/third_party/xla/xla/hlo/analysis/BUILD
@@ -521,91 +521,6 @@ xla_cc_test(
     ],
 )
 
-cc_library(
-    name = "logical_buffer_analysis",
-    srcs = ["logical_buffer_analysis.cc"],
-    hdrs = ["logical_buffer_analysis.h"],
-    deps = [
-        "//xla:shape_util",
-        "//xla/hlo/ir:hlo",
-        "//xla/service:logical_buffer",
-        "//xla/tsl/lib/gtl:map_util",
-        "@com_google_absl//absl/container:flat_hash_map",
-        "@com_google_absl//absl/container:flat_hash_set",
-        "@com_google_absl//absl/log:check",
-        "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:status_macros",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings",
-    ],
-)
-
-cc_library(
-    name = "tuple_points_to_analysis",
-    srcs = ["tuple_points_to_analysis.cc"],
-    hdrs = ["tuple_points_to_analysis.h"],
-    deps = [
-        ":logical_buffer_analysis",
-        "//xla:shape_tree",
-        "//xla:shape_util",
-        "//xla:util",
-        "//xla:xla_data_proto_cc",
-        "//xla/hlo/ir:hlo",
-        "//xla/service:logical_buffer",
-        "//xla/tsl/lib/gtl:compactptrset",
-        "//xla/tsl/platform:logging",
-        "@com_google_absl//absl/algorithm:container",
-        "@com_google_absl//absl/container:flat_hash_map",
-        "@com_google_absl//absl/container:flat_hash_set",
-        "@com_google_absl//absl/container:inlined_vector",
-        "@com_google_absl//absl/log",
-        "@com_google_absl//absl/log:check",
-        "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:status_macros",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/strings:str_format",
-        "@com_google_absl//absl/types:span",
-    ],
-)
-
-xla_cc_test(
-    name = "tuple_points_to_analysis_test",
-    srcs = ["tuple_points_to_analysis_test.cc"],
-    deps = [
-        ":tuple_points_to_analysis",
-        "//xla:shape_util",
-        "//xla/hlo/ir:hlo",
-        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
-        "//xla/hlo/testlib:test",
-        "//xla/service:logical_buffer",
-        "//xla/tsl/platform:logging",
-        "@com_google_absl//absl/algorithm:container",
-        "@com_google_absl//absl/log:check",
-        "@com_google_absl//absl/strings:string_view",
-        "@com_google_absl//absl/types:span",
-        "@com_google_googletest//:gtest_main",
-    ],
-)
-
-xla_cc_test(
-    name = "logical_buffer_analysis_test",
-    srcs = ["logical_buffer_analysis_test.cc"],
-    deps = [
-        ":logical_buffer_analysis",
-        "//xla:shape_util",
-        "//xla:xla_data_proto_cc",
-        "//xla/hlo/ir:hlo",
-        "//xla/hlo/parser:hlo_parser",
-        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
-        "//xla/hlo/testlib:test",
-        "//xla/service:logical_buffer",
-        "//xla/tsl/platform:status_matchers",
-        "@com_google_absl//absl/strings:string_view",
-        "@com_google_googletest//:gtest_main",
-    ],
-)
-
 cc_library(
     name = "interval",
     srcs = ["interval.cc"],
diff --git a/third_party/xla/xla/hlo/analysis/logical_buffer_analysis.cc b/third_party/xla/xla/hlo/analysis/logical_buffer_analysis.cc
deleted file mode 100644
index 3793eea5860..00000000000
--- a/third_party/xla/xla/hlo/analysis/logical_buffer_analysis.cc
+++ /dev/null
@@ -1,344 +0,0 @@
-/* Copyright 2017 The OpenXLA Authors.
-
-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.
-==============================================================================*/
-
-#include "xla/hlo/analysis/logical_buffer_analysis.h"
-
-#include <memory>
-#include <utility>
-#include <vector>
-
-#include "absl/container/flat_hash_set.h"
-#include "absl/log/check.h"
-#include "absl/status/status.h"
-#include "absl/status/status_macros.h"
-#include "absl/status/statusor.h"
-#include "absl/strings/str_cat.h"
-#include "xla/hlo/ir/hlo_casting_utils.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/ir/hlo_opcode.h"
-#include "xla/service/logical_buffer.h"
-#include "xla/shape.h"
-#include "xla/shape_util.h"
-#include "xla/tsl/lib/gtl/map_util.h"
-
-namespace xla {
-
-namespace {
-
-// Gather fusion instructions from 'instruction' into 'fusion_instructions'.
-void GatherFusionInstructions(
-    HloInstruction* instruction,
-    std::vector<HloInstruction*>* fusion_instructions) {
-  CHECK_EQ(HloOpcode::kFusion, instruction->opcode());
-  for (auto* fused : instruction->fused_instructions()) {
-    if (fused->opcode() == HloOpcode::kFusion) {
-      GatherFusionInstructions(fused, fusion_instructions);
-    }
-  }
-  fusion_instructions->push_back(instruction);
-}
-
-bool IsEmptyOutputSubshapeInAsync(const Shape& shape, const ShapeIndex& index) {
-  if (index.empty()) {
-    return false;
-  }
-  if (index.front() != 1) {
-    return false;
-  }
-
-  CHECK(ShapeUtil::IndexIsValid(shape, index));
-
-  const Shape& subshape = ShapeUtil::GetSubshape(shape, index);
-  return subshape.IsTuple() && subshape.tuple_shapes().empty();
-}
-
-}  // namespace
-
-/* static */ absl::StatusOr<std::unique_ptr<LogicalBufferAnalysis>>
-LogicalBufferAnalysis::Run(const HloModule* module,
-                           bool alias_buffer_across_dataflow) {
-  std::unique_ptr<LogicalBufferAnalysis> analysis(
-      new LogicalBufferAnalysis(module, alias_buffer_across_dataflow));
-  ABSL_RETURN_IF_ERROR(analysis->Analyze());
-  return analysis;
-}
-
-absl::Status LogicalBufferAnalysis::Analyze() {
-  // Empirically we usually have a few more logical buffers than instructions,
-  // so reserve 10% more than the number of instructions to avoid frequent
-  // resizes.
-  logical_buffers_.clear();
-  logical_buffers_.reserve((module_->instruction_count() * 11) / 10);
-
-  // We filter out fusion computations, and get to them through fusion
-  // instructions. This is because it's possible to have orphaned (unreachable)
-  // fusion computations, and we don't want to try to assign buffers to those.
-  std::vector<HloInstruction*> fusion_instructions;
-  for (auto* computation : module_->MakeNonfusionComputations()) {
-    ABSL_RETURN_IF_ERROR(computation->Accept(this));
-    for (auto* instruction : computation->instructions()) {
-      if (instruction->opcode() != HloOpcode::kFusion) {
-        continue;
-      }
-      GatherFusionInstructions(instruction, &fusion_instructions);
-    }
-  }
-  for (auto* instruction : fusion_instructions) {
-    ABSL_RETURN_IF_ERROR(
-        instruction->fused_instructions_computation()->Accept(this));
-  }
-  return absl::OkStatus();
-}
-
-absl::StatusOr<LogicalBuffer*> LogicalBufferAnalysis::GetBuffer(
-    LogicalBuffer::Id id) const {
-  if (id < 0 || id >= logical_buffers_.size()) {
-    return absl::InvalidArgumentError(
-        absl::StrCat("Invalid logical buffer ID: ", id));
-  }
-  return logical_buffers_[id].get();
-}
-
-absl::StatusOr<LogicalBuffer*> LogicalBufferAnalysis::GetBuffer(
-    HloInstruction* instruction, const ShapeIndex& index) const {
-  LogicalBuffer* buffer = tsl::gtl::FindPtrOrNull(
-      output_buffers_, std::make_pair(instruction, index));
-  if (buffer == nullptr) {
-    return absl::InvalidArgumentError(
-        absl::StrCat("No buffer defined by ", instruction->name(), " at index ",
-                     index.ToString()));
-  }
-  return buffer;
-}
-
-void LogicalBufferAnalysis::NewLogicalBuffer(HloInstruction* instruction,
-                                             const ShapeIndex& index) {
-  LogicalBuffer::Id id = logical_buffers_.size();
-  auto buffer = std::make_unique<LogicalBuffer>(instruction, index, id);
-  auto position = std::make_pair(instruction, index);
-  CHECK(output_buffers_.insert({position, buffer.get()}).second);
-  logical_buffers_.push_back(std::move(buffer));
-}
-
-absl::Status LogicalBufferAnalysis::DefaultAction(
-    HloInstruction* hlo_instruction) {
-  // Create a logical buffer for each output of the instruction.
-  ShapeUtil::ForEachSubshape(
-      hlo_instruction->shape(),
-      [this, hlo_instruction](const Shape& shape, const ShapeIndex& index) {
-        NewLogicalBuffer(hlo_instruction, index);
-      });
-
-  return absl::OkStatus();
-}
-
-absl::Status LogicalBufferAnalysis::HandleGetTupleElement(HloInstruction*) {
-  // GetTupleElement does not create buffers.
-  return absl::OkStatus();
-}
-
-absl::Status LogicalBufferAnalysis::HandleAddDependency(
-    HloInstruction* add_dependency) {
-  // AddDependency just forwards the value of its zero-th operand and does not
-  // create buffers.
-  return absl::OkStatus();
-}
-
-absl::Status LogicalBufferAnalysis::HandleCopy(HloInstruction* copy) {
-  // The top-level buffer (index={}) for kCopy is newly created, but all other
-  // buffers (in the case of a tuple shape) come from the operand
-  NewLogicalBuffer(copy, /*index=*/{});
-  return absl::OkStatus();
-}
-
-absl::Status LogicalBufferAnalysis::HandleBitcast(HloInstruction*) {
-  // A kBitcast instruction aliases its operand. That is, the buffer of its
-  // result *is* the buffer of its operand.
-  return absl::OkStatus();
-}
-
-absl::Status LogicalBufferAnalysis::HandleDomain(HloInstruction*) {
-  // A kDomain instruction aliases its operand. That is, the buffer of its
-  // result *is* the buffer of its operand.
-  return absl::OkStatus();
-}
-
-absl::Status LogicalBufferAnalysis::HandleRecvDone(HloInstruction* recv_done) {
-  // RecvDone produces a two-element tuple containing the data value (which
-  // aliases part of its operand) and a token. Only the tuple index table and
-  // the token are defined by the RecvDone.
-  NewLogicalBuffer(recv_done, /*index=*/{});
-  NewLogicalBuffer(recv_done, /*index=*/{1});
-  return absl::OkStatus();
-}
-
-absl::Status LogicalBufferAnalysis::HandleSend(HloInstruction* send) {
-  // Send creates new buffers for the top-level tuple, the context (tuple
-  // element at {1}), and the token (tuple element at {2}). Tuple element at {0}
-  // is an alias of the Send operand, so we don't need to create a new Logical
-  // Buffer for that.
-  NewLogicalBuffer(send, /*index=*/{});
-  NewLogicalBuffer(send, /*index=*/{1});
-  NewLogicalBuffer(send, /*index=*/{2});
-  return absl::OkStatus();
-}
-
-absl::Status LogicalBufferAnalysis::HandleCopyStart(
-    HloInstruction* copy_start) {
-  // CopyStart defines the tuple, target buffer at index {0}, and context at
-  // index {2}.
-  NewLogicalBuffer(copy_start, /*index=*/{});
-  NewLogicalBuffer(copy_start, /*index=*/{0});
-  NewLogicalBuffer(copy_start, /*index=*/{2});
-  return absl::OkStatus();
-}
-
-absl::Status LogicalBufferAnalysis::HandleCopyDone(HloInstruction* copy_done) {
-  // The output of CopyDone aliases with operand {0}. CopyDone doesn't create
-  // any buffers.
-  return absl::OkStatus();
-}
-
-absl::Status LogicalBufferAnalysis::HandleTuple(HloInstruction* tuple) {
-  // A Tuple instruction only creates the top-level buffer.
-  NewLogicalBuffer(tuple, /*index=*/{});
-  return absl::OkStatus();
-}
-
-absl::Status LogicalBufferAnalysis::HandleCustomCall(
-    HloInstruction* custom_call) {
-  auto ccall = Cast<HloCustomCallInstruction>(custom_call);
-  absl::flat_hash_set<ShapeIndex> aliased_outputs;
-  for (const auto& pair : ccall->output_to_operand_aliasing()) {
-    aliased_outputs.insert(pair.first);
-  }
-  ShapeUtil::ForEachSubshape(ccall->shape(), [&](const Shape& shape,
-                                                 const ShapeIndex& index) {
-    if (!aliased_outputs.contains(index) || !alias_buffer_across_dataflow_) {
-      NewLogicalBuffer(custom_call, index);
-    }
-  });
-  return absl::OkStatus();
-}
-
-// WARNING (b/259460539): output_to_operand_aliasing was moved from
-// HloCustomCallInstruction to HloCallableInstruction so that fusions can
-// also be annotated with this aliasing. This feature might not be complete.
-// Particular to this analysis, we added the HandleFusion function below to
-// accommodate the output-operand aliases (similar to HandleCustomCall).
-// TODO (sacer): We might want to consider the pairs discovered by
-// GetFusionInstructionInPlaceInputOutputPairs() here as well.
-absl::Status LogicalBufferAnalysis::HandleFusion(HloInstruction* fusion) {
-  auto cfusion = Cast<HloFusionInstruction>(fusion);
-  absl::flat_hash_set<ShapeIndex> aliased_outputs;
-  for (const auto& pair : cfusion->output_to_operand_aliasing()) {
-    aliased_outputs.insert(pair.first);
-  }
-  ShapeUtil::ForEachSubshape(cfusion->shape(),
-                             [&](const Shape& shape, const ShapeIndex& index) {
-                               if (!aliased_outputs.contains(index)) {
-                                 NewLogicalBuffer(fusion, index);
-                               }
-                             });
-  return absl::OkStatus();
-}
-
-absl::Status LogicalBufferAnalysis::HandleAsyncStart(
-    HloInstruction* async_start) {
-  const Shape& async_start_shape = async_start->shape();
-  absl::flat_hash_set<ShapeIndex> aliased_outputs;
-  for (const auto& pair : Cast<HloAsyncStartInstruction>(async_start)
-                              ->output_to_operand_aliasing()) {
-    aliased_outputs.insert(pair.first);
-  }
-
-  ShapeUtil::ForEachSubshape(
-      async_start_shape, [&](const Shape& shape, const ShapeIndex& index) {
-        bool has_implicit_alias = (index.size() >= 2 && index.front() == 0);
-        bool has_explicit_alias = aliased_outputs.contains(index);
-        if (!has_implicit_alias &&
-            (!has_explicit_alias ||
-             IsEmptyOutputSubshapeInAsync(async_start_shape, index))) {
-          NewLogicalBuffer(async_start, index);
-        }
-      });
-  return absl::OkStatus();
-}
-
-absl::Status LogicalBufferAnalysis::HandleAsyncUpdate(
-    HloInstruction* async_update) {
-  const Shape& async_update_shape = async_update->shape();
-  const HloInstruction* prev_async_op = async_update->operand(0);
-  const Shape& prev_async_op_shape = prev_async_op->shape();
-
-  HloInstruction* async_start = async_update->async_chain_start();
-  absl::flat_hash_set<ShapeIndex> explicitly_aliased_outputs;
-  for (const auto& pair : Cast<HloAsyncStartInstruction>(async_start)
-                              ->output_to_operand_aliasing()) {
-    explicitly_aliased_outputs.insert(pair.first);
-  }
-
-  absl::flat_hash_set<ShapeIndex> aliased_contexts;
-  for (const auto& [output_index, _] :
-       Cast<HloAsyncUpdateInstruction>(async_update)
-           ->output_to_operand_aliasing()) {
-    aliased_contexts.insert(output_index);
-  }
-
-  ShapeUtil::ForEachSubshape(async_update_shape, [&](const Shape& subshape,
-                                                     const ShapeIndex& index) {
-    // Operands, ignore.
-    if (index.size() >= 2 && index.front() == 0) {
-      return;
-    }
-
-    // If explicitly aliased output, and not an empty tuple, ignore (we want to
-    // avoid creating values for buffer not bound yet).
-    bool has_explicit_alias = explicitly_aliased_outputs.contains(index);
-    if (has_explicit_alias &&
-        !IsEmptyOutputSubshapeInAsync(async_update_shape, index)) {
-      return;
-    }
-
-    // Forwarded operand/output from the previous async op.
-    if (!index.empty() && index.front() < 2 &&
-        ShapeUtil::IndexIsValid(prev_async_op_shape, index)) {
-      const Shape& prev_subshape =
-          ShapeUtil::GetSubshape(prev_async_op_shape, index);
-      if (ShapeUtil::Compatible(prev_subshape, subshape)) {
-        return;
-      }
-    }
-
-    // Aliased context.
-    if (aliased_contexts.contains(index)) {
-      return;
-    }
-
-    // Otherwise, new output or new context or new tuple container
-    NewLogicalBuffer(async_update, index);
-  });
-
-  return absl::OkStatus();
-}
-
-absl::Status LogicalBufferAnalysis::HandleAsyncDone(
-    HloInstruction* async_done) {
-  return DefaultAction(async_done);
-}
-
-}  // namespace xla
diff --git a/third_party/xla/xla/hlo/analysis/logical_buffer_analysis.h b/third_party/xla/xla/hlo/analysis/logical_buffer_analysis.h
deleted file mode 100644
index 534282ae573..00000000000
--- a/third_party/xla/xla/hlo/analysis/logical_buffer_analysis.h
+++ /dev/null
@@ -1,142 +0,0 @@
-/* Copyright 2017 The OpenXLA Authors.
-
-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.
-==============================================================================*/
-
-#ifndef XLA_HLO_ANALYSIS_LOGICAL_BUFFER_ANALYSIS_H_
-#define XLA_HLO_ANALYSIS_LOGICAL_BUFFER_ANALYSIS_H_
-
-#include <cstddef>
-#include <memory>
-#include <utility>
-#include <vector>
-
-#include "absl/container/flat_hash_map.h"
-#include "absl/status/status.h"
-#include "absl/status/statusor.h"
-#include "xla/hlo/ir/dfs_hlo_visitor_with_default.h"
-#include "xla/hlo/ir/hlo_instruction.h"
-#include "xla/hlo/ir/hlo_module.h"
-#include "xla/service/logical_buffer.h"
-#include "xla/shape_util.h"
-
-namespace xla {
-// Identifies and instantiates all the logical buffers defined by HLO
-// instructions in an HLO module.
-//
-// This class visits HLO instructions in post-order DFS order to determine
-// which output shape indices of which HLO instructions define/create a new
-// buffer and which ones alias existing buffers. It assigns a unique
-// `LogicalBuffer::Id` to each newly defined buffer.
-//
-// Analysis steps:
-// 1. Visit each HLO instruction in DFS order.
-// 2. Identify the logical buffers defined/created by the instruction.
-// 3. Populate `logical_buffers_` (for retrieval by ID) and `output_buffers_`
-//    (for retrieval by instruction and subshape index).
-class LogicalBufferAnalysis : public DfsHloVisitorWithDefault {
- public:
-  // Runs the logical buffer analysis on the given `HloModule` and returns
-  // the completed analysis result.
-  static absl::StatusOr<std::unique_ptr<LogicalBufferAnalysis>> Run(
-      const HloModule* module, bool alias_buffer_across_dataflow = false);
-
-  // Returns the `LogicalBuffer` with the given ID.
-  //
-  // Returns an error status if the ID is invalid (not less than
-  // `num_logical_buffers()`).
-  absl::StatusOr<LogicalBuffer*> GetBuffer(LogicalBuffer::Id id) const;
-
-  // Returns the `LogicalBuffer` defined by the instruction at
-  // the specified `ShapeIndex` in its output.
-  //
-  // Returns an error status if the instruction does not define a
-  // logical buffer at the given index.
-  absl::StatusOr<LogicalBuffer*> GetBuffer(HloInstruction* instruction,
-                                           const ShapeIndex& index) const;
-
-  // Returns a view list of all `LogicalBuffer`s created in the module.
-  const std::vector<std::unique_ptr<LogicalBuffer>>& logical_buffers() const {
-    return logical_buffers_;
-  }
-
-  // Returns the total number of unique `LogicalBuffer`s created.
-  size_t num_logical_buffers() const { return logical_buffers_.size(); }
-
- private:
-  explicit LogicalBufferAnalysis(const HloModule* module,
-                                 bool alias_buffer_across_dataflow)
-      : module_(module),
-        alias_buffer_across_dataflow_(alias_buffer_across_dataflow) {}
-
-  // Performs the DFS analysis over all non-fusion and fusion computations
-  // of the module.
-  absl::Status Analyze();
-
-  // The module this analysis is performed on.
-  const HloModule* module_;
-
-  // Instantiates a new `LogicalBuffer` associated with `instruction` at
-  // `index` and adds it to internal collections.
-  void NewLogicalBuffer(HloInstruction* instruction, const ShapeIndex& index);
-
-  // Visitor overrides that identify buffer creation points for each
-  // instruction type.
-  absl::Status DefaultAction(HloInstruction* hlo_instruction) override;
-  absl::Status HandleTuple(HloInstruction* tuple) override;
-  absl::Status HandleGetTupleElement(
-      HloInstruction* get_tuple_element) override;
-  absl::Status HandleBitcast(HloInstruction* bitcast) override;
-  absl::Status HandleDomain(HloInstruction* domain) override;
-  absl::Status HandleCopy(HloInstruction* copy) override;
-  absl::Status HandleCopyStart(HloInstruction* copy_start) override;
-  absl::Status HandleCopyDone(HloInstruction* copy_done) override;
-  absl::Status HandleRecvDone(HloInstruction* recv_done) override;
-  absl::Status HandleSend(HloInstruction* send) override;
-  absl::Status HandleAddDependency(HloInstruction* add_dependency) override;
-  absl::Status HandleCustomCall(HloInstruction* custom_call) override;
-  absl::Status HandleFusion(HloInstruction* fusion) override;
-  absl::Status HandleAsyncStart(HloInstruction* async_start) override;
-  absl::Status HandleAsyncUpdate(HloInstruction* async_update) override;
-  absl::Status HandleAsyncDone(HloInstruction* async_done) override;
-
-  // A collection of all logical buffers, indexed by their ID.
-  std::vector<std::unique_ptr<LogicalBuffer>> logical_buffers_;
-
-  // A map mapping an `(HloInstruction, ShapeIndex)` pair to the `LogicalBuffer`
-  // defined at that position.
-  absl::flat_hash_map<std::pair<const HloInstruction*, const ShapeIndex>,
-                      LogicalBuffer*>
-      output_buffers_;
-
-  // Flag indicating whether to respect output-to-operand aliasing annotations
-  // on custom calls (only).
-  //
-  // If true, outputs explicitly annotated to alias an operand do not define
-  // new logical buffers (the operand's buffer is forwarded). If false, new
-  // logical buffers are created for all outputs, ignoring those annotations.
-  //
-  // For example, given:
-  //   %ccall = custom-call(%param), output_to_operand_aliasing={ {}: (0, {}) }
-  //
-  // - If true: `%ccall` does not define a new buffer; it reuses `%param`'s.
-  // - If false: `%ccall` defines a new `LogicalBuffer(%ccall, {})`.
-  //
-  // Note: This does not affect Call, Fusion, AsyncStart, Parameter,
-  // or While instructions, which always define new buffers.
-  const bool alias_buffer_across_dataflow_ = false;
-};
-
-}  // namespace xla
-
-#endif  // XLA_HLO_ANALYSIS_LOGICAL_BUFFER_ANALYSIS_H_
diff --git a/third_party/xla/xla/hlo/analysis/logical_buffer_analysis_test.cc b/third_party/xla/xla/hlo/analysis/logical_buffer_analysis_test.cc
deleted file mode 100644
index e134ee26876..00000000000
--- a/third_party/xla/xla/hlo/analysis/logical_buffer_analysis_test.cc
+++ /dev/null
@@ -1,833 +0,0 @@
-/* Copyright 2026 The OpenXLA Authors.
-
-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.
-==============================================================================*/
-
-#include "xla/hlo/analysis/logical_buffer_analysis.h"
-
-#include <memory>
-#include <string>
-#include <utility>
-#include <vector>
-
-#include "absl/strings/string_view.h"
-#include "xla/hlo/ir/hlo_instruction.h"
-#include "xla/hlo/parser/hlo_parser.h"
-#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
-#include "xla/hlo/testlib/test.h"
-#include "xla/service/logical_buffer.h"
-#include "xla/shape.h"
-#include "xla/shape_util.h"
-#include "xla/tsl/platform/status_matchers.h"
-#include "xla/xla_data.pb.h"
-
-namespace xla {
-namespace {
-
-using ::testing::UnorderedElementsAre;
-using ::tsl::testing::IsOkAndHolds;
-using ::tsl::testing::StatusIs;
-
-class LogicalBufferAnalysisTest : public HloHardwareIndependentTestBase {
- protected:
-  std::unique_ptr<LogicalBufferAnalysis> analysis_;
-
-  // Verifies that a buffer is defined by `instruction` at `index` and matches.
-  void VerifyBufferDefinedAt(HloInstruction* instruction,
-                             const ShapeIndex& index) {
-    ASSERT_OK_AND_ASSIGN(LogicalBuffer * buffer,
-                         analysis_->GetBuffer(instruction, index));
-    EXPECT_EQ(buffer->instruction(), instruction);
-    EXPECT_EQ(buffer->index(), index);
-    EXPECT_THAT(analysis_->GetBuffer(buffer->id()),
-                absl_testing::IsOkAndHolds(buffer));
-  }
-
-  // Verifies that no buffer is defined by `instruction` at `index`.
-  void VerifyNoBufferDefinedAt(const HloInstruction* instruction,
-                               const ShapeIndex& index) {
-    for (const auto& buf : analysis_->logical_buffers()) {
-      if (buf->instruction() == instruction && buf->index() == index) {
-        ADD_FAILURE() << "Instruction " << instruction->name() << " at index "
-                      << index.ToString()
-                      << " should not define a logical buffer.";
-      }
-    }
-  }
-
-  // Returns all defining locations present in the analysis.
-  std::vector<std::pair<const HloInstruction*, ShapeIndex>> GetDefiningSites() {
-    std::vector<std::pair<const HloInstruction*, ShapeIndex>> defining_sites;
-    for (const auto& buf : analysis_->logical_buffers()) {
-      defining_sites.push_back({buf->instruction(), buf->index()});
-    }
-    return defining_sites;
-  }
-};
-
-TEST_F(LogicalBufferAnalysisTest, BasicAndAccessors) {
-  const absl::string_view hlo_str = R"(
-  HloModule module
-
-  ENTRY entry {
-    p0 = f32[2,3] parameter(0)
-    ROOT const = f32[] constant(1.0)
-  }
-  )";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str));
-  ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(module.get()));
-
-  HloInstruction* const_inst = FindInstruction(module.get(), "const");
-  HloInstruction* param = FindInstruction(module.get(), "p0");
-
-  EXPECT_EQ(analysis_->num_logical_buffers(), 2);
-  EXPECT_FALSE(analysis_->logical_buffers().empty());
-
-  VerifyBufferDefinedAt(const_inst, {});
-  VerifyBufferDefinedAt(param, {});
-
-  EXPECT_THAT(GetDefiningSites(),
-              UnorderedElementsAre(std::make_pair(const_inst, ShapeIndex({})),
-                                   std::make_pair(param, ShapeIndex({}))));
-}
-
-TEST_F(LogicalBufferAnalysisTest, DataflowOps) {
-  absl::string_view hlo_str = R"(
-  HloModule module
-
-  ENTRY entry {
-    p0 = f32[2,3] parameter(0)
-    p1 = f32[2,3] parameter(1)
-    tok = token[] after-all()
-    tuple = (f32[2,3], f32[2,3]) tuple(p0, p1)
-    gte = f32[2,3] get-tuple-element(tuple), index=0
-    bitcast = f32[2,3] bitcast(gte)
-    domain_inst = f32[2,3] domain(bitcast), domain={kind="sharding", entry={maximal device=0}, exit={maximal device=1}}
-    copy = f32[2,3] copy(domain_inst)
-    ROOT dep = f32[2,3] add-dependency(copy, tok)
-  }
-  )";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str));
-  ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(module.get()));
-
-  HloInstruction* p0 = FindInstruction(module.get(), "p0");
-  HloInstruction* p1 = FindInstruction(module.get(), "p1");
-  HloInstruction* tok = FindInstruction(module.get(), "tok");
-  HloInstruction* tuple = FindInstruction(module.get(), "tuple");
-  HloInstruction* gte = FindInstruction(module.get(), "gte");
-  HloInstruction* bitcast = FindInstruction(module.get(), "bitcast");
-  HloInstruction* domain_inst = FindInstruction(module.get(), "domain_inst");
-  HloInstruction* copy = FindInstruction(module.get(), "copy");
-  HloInstruction* dep = FindInstruction(module.get(), "dep");
-
-  EXPECT_EQ(analysis_->num_logical_buffers(), 5);
-
-  VerifyBufferDefinedAt(p0, {});
-  VerifyBufferDefinedAt(p1, {});
-  VerifyBufferDefinedAt(tok, {});
-  VerifyBufferDefinedAt(tuple, {});
-  VerifyBufferDefinedAt(copy, {});
-
-  VerifyNoBufferDefinedAt(tuple, {0});
-  VerifyNoBufferDefinedAt(tuple, {1});
-  VerifyNoBufferDefinedAt(gte, {});
-  VerifyNoBufferDefinedAt(bitcast, {});
-  VerifyNoBufferDefinedAt(domain_inst, {});
-  VerifyNoBufferDefinedAt(dep, {});
-
-  EXPECT_THAT(GetDefiningSites(),
-              UnorderedElementsAre(std::make_pair(p0, ShapeIndex({})),
-                                   std::make_pair(p1, ShapeIndex({})),
-                                   std::make_pair(tok, ShapeIndex({})),
-                                   std::make_pair(tuple, ShapeIndex({})),
-                                   std::make_pair(copy, ShapeIndex({}))));
-}
-
-TEST_F(LogicalBufferAnalysisTest, CopyStartDone) {
-  absl::string_view hlo_str = R"(
-  HloModule module
-
-  ENTRY entry {
-    p0 = f32[2,3] parameter(0)
-    copy-start = (f32[2,3], f32[2,3], u32[]) copy-start(p0)
-    ROOT copy-done = f32[2,3] copy-done(copy-start)
-  }
-  )";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str));
-  ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(module.get()));
-
-  HloInstruction* p0 = FindInstruction(module.get(), "p0");
-  HloInstruction* copy_start = FindInstruction(module.get(), "copy-start");
-  HloInstruction* copy_done = FindInstruction(module.get(), "copy-done");
-
-  EXPECT_EQ(analysis_->num_logical_buffers(), 4);
-
-  VerifyBufferDefinedAt(p0, {});
-  VerifyBufferDefinedAt(copy_start, {});
-  VerifyBufferDefinedAt(copy_start, {0});
-  VerifyBufferDefinedAt(copy_start, {2});
-
-  VerifyNoBufferDefinedAt(copy_start, {1});
-  VerifyNoBufferDefinedAt(copy_done, {});
-
-  EXPECT_THAT(
-      GetDefiningSites(),
-      UnorderedElementsAre(std::make_pair(p0, ShapeIndex({})),
-                           std::make_pair(copy_start, ShapeIndex({})),
-                           std::make_pair(copy_start, ShapeIndex({0})),
-                           std::make_pair(copy_start, ShapeIndex({2}))));
-}
-
-TEST_F(LogicalBufferAnalysisTest, SendRecvDone) {
-  absl::string_view hlo_str = R"(
-  HloModule module
-
-  ENTRY entry {
-    p0 = f32[2,3] parameter(0)
-    tok = token[] after-all()
-    send = (f32[2,3], u32[], token[]) send(p0, tok), channel_id=1
-    send-done = token[] send-done(send), channel_id=1
-    recv = (f32[2,3], u32[], token[]) recv(tok), channel_id=2
-    ROOT recv-done = (f32[2,3], token[]) recv-done(recv), channel_id=2
-  }
-  )";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str));
-  ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(module.get()));
-
-  HloInstruction* p0 = FindInstruction(module.get(), "p0");
-  HloInstruction* tok = FindInstruction(module.get(), "tok");
-  HloInstruction* send = FindInstruction(module.get(), "send");
-  HloInstruction* send_done = FindInstruction(module.get(), "send-done");
-  HloInstruction* recv = FindInstruction(module.get(), "recv");
-  HloInstruction* recv_done = FindInstruction(module.get(), "recv-done");
-
-  EXPECT_EQ(analysis_->num_logical_buffers(), 12);
-
-  VerifyBufferDefinedAt(p0, {});
-  VerifyBufferDefinedAt(tok, {});
-  VerifyBufferDefinedAt(send, {});
-  VerifyBufferDefinedAt(send, {1});
-  VerifyBufferDefinedAt(send, {2});
-  VerifyBufferDefinedAt(send_done, {});
-
-  VerifyBufferDefinedAt(recv, {});
-  VerifyBufferDefinedAt(recv, {0});
-  VerifyBufferDefinedAt(recv, {1});
-  VerifyBufferDefinedAt(recv, {2});
-
-  VerifyBufferDefinedAt(recv_done, {});
-  VerifyBufferDefinedAt(recv_done, {1});
-
-  VerifyNoBufferDefinedAt(send, {0});
-  VerifyNoBufferDefinedAt(recv_done, {0});
-
-  EXPECT_THAT(GetDefiningSites(),
-              UnorderedElementsAre(std::make_pair(p0, ShapeIndex({})),
-                                   std::make_pair(tok, ShapeIndex({})),
-                                   std::make_pair(send, ShapeIndex({})),
-                                   std::make_pair(send, ShapeIndex({1})),
-                                   std::make_pair(send, ShapeIndex({2})),
-                                   std::make_pair(send_done, ShapeIndex({})),
-                                   std::make_pair(recv, ShapeIndex({})),
-                                   std::make_pair(recv, ShapeIndex({0})),
-                                   std::make_pair(recv, ShapeIndex({1})),
-                                   std::make_pair(recv, ShapeIndex({2})),
-                                   std::make_pair(recv_done, ShapeIndex({})),
-                                   std::make_pair(recv_done, ShapeIndex({1}))));
-}
-
-TEST_F(LogicalBufferAnalysisTest, CustomCallFusionAsync) {
-  absl::string_view hlo_str = R"(
-  HloModule module
-
-  fused_computation {
-    p0.fused = f32[2,3] parameter(0)
-    ROOT copy.fused = f32[2,3] copy(p0.fused)
-  }
-
-  ENTRY entry {
-    p0 = f32[2,3] parameter(0)
-    ccall = f32[2,3] custom-call(p0), custom_call_target="cc_target",
-            output_to_operand_aliasing={ {}: (0, {}) }
-    fusion = f32[2,3] fusion(p0), kind=kLoop, calls=fused_computation
-    async-start = ((f32[2,3]), f32[2,3], u32[]) custom-call-start(p0),
-                                                custom_call_target="bar"
-    async-update = ((f32[2,3]), f32[2,3], u32[]) async-update(async-start),
-                            output_to_operand_aliasing={ {2}: (0, {2}) }
-    ROOT async-done = f32[2,3] custom-call-done(async-update)
-  }
-  )";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str));
-  ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(module.get()));
-
-  HloInstruction* p0_fused = FindInstruction(module.get(), "p0.fused");
-  HloInstruction* copy_fused = FindInstruction(module.get(), "copy.fused");
-  HloInstruction* p0 = FindInstruction(module.get(), "p0");
-  HloInstruction* ccall = FindInstruction(module.get(), "ccall");
-  HloInstruction* fusion = FindInstruction(module.get(), "fusion");
-  HloInstruction* async_start = FindInstruction(module.get(), "async-start");
-  HloInstruction* async_update = FindInstruction(module.get(), "async-update");
-  HloInstruction* async_done = FindInstruction(module.get(), "async-done");
-  HloComputation* async_computation = async_start->async_wrapped_computation();
-  HloInstruction* async_wrapper_root = async_computation->root_instruction();
-  ASSERT_EQ(async_wrapper_root->operand_count(), 1);
-  HloInstruction* async_wrapper_param = async_wrapper_root->mutable_operand(0);
-
-  EXPECT_EQ(analysis_->num_logical_buffers(), 13);
-
-  VerifyBufferDefinedAt(p0, {});
-  VerifyBufferDefinedAt(ccall, {});
-  VerifyBufferDefinedAt(fusion, {});
-  VerifyBufferDefinedAt(p0_fused, {});
-  VerifyBufferDefinedAt(copy_fused, {});
-
-  VerifyBufferDefinedAt(async_start, {});
-  VerifyBufferDefinedAt(async_start, {0});
-  VerifyBufferDefinedAt(async_start, {1});
-  VerifyBufferDefinedAt(async_start, {2});
-  // {0, 0} is implicitly aliased to input parameter p0.
-  VerifyNoBufferDefinedAt(async_start, {0, 0});
-  VerifyBufferDefinedAt(async_update, {});
-  VerifyNoBufferDefinedAt(async_update, {0});
-  VerifyNoBufferDefinedAt(async_update, {1});
-  VerifyNoBufferDefinedAt(async_update, {2});
-
-  // There are two more buffers defined in the async computation,
-  // which is created by the custom-call-start and custom-call-done
-  // in async wrapped computation.
-  VerifyBufferDefinedAt(async_wrapper_param, {});
-  VerifyBufferDefinedAt(async_wrapper_root, {});
-  VerifyBufferDefinedAt(async_done, {});
-}
-
-TEST_F(LogicalBufferAnalysisTest, NestedAndAliasedFusion) {
-  absl::string_view hlo_str = R"(
-  HloModule module
-
-  fused_computation_inner {
-    p0.inner = f32[2,3] parameter(0)
-    ROOT copy.inner = f32[2,3] copy(p0.inner)
-  }
-
-  fused_computation_outer {
-    p0.outer = f32[2,3] parameter(0)
-    ROOT fusion.inner = f32[2,3] fusion(p0.outer), kind=kLoop,
-                                 calls=fused_computation_inner
-  }
-
-  fused_computation_alias {
-    p0.alias = f32[2,3] parameter(0)
-    ROOT copy.alias = f32[2,3] copy(p0.alias)
-  }
-
-  ENTRY entry {
-    p0 = f32[2,3] parameter(0)
-    fusion.outer = f32[2,3] fusion(p0), kind=kLoop, calls=fused_computation_outer
-    ROOT fusion.alias = f32[2,3] fusion(p0), kind=kLoop,
-                                 calls=fused_computation_alias,
-                                 output_to_operand_aliasing={ {}: (0, {}) }
-  }
-  )";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str));
-  ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(module.get()));
-
-  HloInstruction* p0 = FindInstruction(module.get(), "p0");
-  HloInstruction* fusion_outer = FindInstruction(module.get(), "fusion.outer");
-
-  HloInstruction* fusion_alias = FindInstruction(module.get(), "fusion.alias");
-
-  VerifyBufferDefinedAt(p0, {});
-  VerifyBufferDefinedAt(fusion_outer, {});
-  // No buffer is defined at the output of fusion.alias because it is explicitly
-  // aliased to the input.
-  VerifyNoBufferDefinedAt(fusion_alias, {});
-}
-
-TEST_F(LogicalBufferAnalysisTest, AsyncStartWithOutputToOperandAliasing) {
-  absl::string_view hlo_str = R"(
-  HloModule module
-
-  fused_computation_async {
-    p0.async = f32[2,3] parameter(0)
-    ROOT copy.async = f32[2,3] copy(p0.async)
-  }
-
-  ENTRY entry {
-    p0 = f32[2,3] parameter(0)
-    async-start = ((f32[2,3]), f32[2,3], u32[]) async-start(p0),
-                  calls=fused_computation_async,
-                  output_to_operand_aliasing={ {1}: (0, {}) }
-    ROOT async-done = f32[2,3] async-done(async-start)
-  }
-  )";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str));
-  ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(module.get()));
-
-  HloInstruction* p0 = FindInstruction(module.get(), "p0");
-  HloInstruction* async_start = FindInstruction(module.get(), "async-start");
-
-  VerifyBufferDefinedAt(p0, {});
-
-  // async-start defines output except implicitly/explicitly aliased:
-  // - {} (top-level tuple) -> defined (not aliased)
-  // - {0} (sub-tuple) -> defined
-  // - {0, 0} -> implicitly aliased
-  // - {1} -> explicitly aliased (to p0)
-  // - {2} (state) -> defined
-  VerifyBufferDefinedAt(async_start, {});
-  VerifyBufferDefinedAt(async_start, {0});
-  VerifyBufferDefinedAt(async_start, {2});
-
-  VerifyNoBufferDefinedAt(async_start, {0, 0});
-  VerifyNoBufferDefinedAt(async_start, {1});
-}
-
-TEST_F(LogicalBufferAnalysisTest, CustomCallAliasingWithDataflowFlag) {
-  absl::string_view hlo_str = R"(
-  HloModule module
-
-  ENTRY entry {
-    p0 = f32[2,3] parameter(0)
-    ROOT ccall = f32[2,3] custom-call(p0), custom_call_target="cc_target",
-                 output_to_operand_aliasing={ {}: (0, {}) }
-  }
-  )";
-
-  // Case 1: alias_buffer_across_dataflow = false (default)
-  {
-    ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str));
-    ASSERT_OK_AND_ASSIGN(
-        analysis_,
-        LogicalBufferAnalysis::Run(module.get(),
-                                   /*alias_buffer_across_dataflow=*/false));
-    HloInstruction* ccall = FindInstruction(module.get(), "ccall");
-    VerifyBufferDefinedAt(ccall, {});
-  }
-
-  // Case 2: alias_buffer_across_dataflow = true
-  {
-    ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str));
-    ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(
-                                        module.get(),
-                                        /*alias_buffer_across_dataflow=*/true));
-    HloInstruction* ccall = FindInstruction(module.get(), "ccall");
-    VerifyNoBufferDefinedAt(ccall, {});
-  }
-}
-
-TEST_F(LogicalBufferAnalysisTest, InvalidGetBufferThrows) {
-  const absl::string_view hlo_str = R"(
-  HloModule module
-
-  ENTRY entry {
-    p0 = f32[2,3] parameter(0)
-    ROOT const = f32[] constant(1.0)
-  }
-  )";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str));
-  ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(module.get()));
-
-  HloInstruction* param = FindInstruction(module.get(), "p0");
-
-  EXPECT_THAT(analysis_->GetBuffer(param, {0}),
-              absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
-}
-
-TEST_F(LogicalBufferAnalysisTest, InvalidGetBufferIdBehavior) {
-  const absl::string_view hlo_str = R"(
-  HloModule module
-
-  ENTRY entry {
-    p0 = f32[2,3] parameter(0)
-    ROOT const = f32[] constant(1.0)
-  }
-  )";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str));
-  ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(module.get()));
-
-  // GetBuffer with invalid ID of 100 on a module containing 2 buffers.
-  EXPECT_THAT(analysis_->GetBuffer(100),
-              absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
-}
-
-TEST_F(LogicalBufferAnalysisTest, SimpleAsyncChain) {
-  absl::string_view hlo_str = R"(
-  HloModule module
-
-  async_computation {
-    p0 = f32[4] parameter(0)
-    ROOT ccall = f32[4] custom-call(p0), custom_call_target="bar"
-  }
-
-  ENTRY entry {
-    p0 = f32[4] parameter(0)
-    start = ((f32[4]), f32[4], s32[]) async-start(p0), calls=async_computation,
-                                       async_execution_thread="sparsecore"
-    update = ((f32[4]), f32[4], s32[]) async-update(start)
-    ROOT done = f32[4] async-done(update)
-  }
-  )";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str));
-  ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(module.get()));
-
-  HloInstruction* p0 = FindInstruction(module.get(), "p0");
-  HloInstruction* start = FindInstruction(module.get(), "start");
-  HloInstruction* update = FindInstruction(module.get(), "update");
-
-  VerifyBufferDefinedAt(p0, {});
-
-  // Verify start buffers:
-  // - {} (top-level tuple) -> defined
-  // - {0} (sub-tuple) -> defined
-  // - {0, 0} -> implicitly aliased (to input parameter p0) -> not defined
-  // - {1} (output) -> defined
-  // - {2} (state) -> defined
-  VerifyBufferDefinedAt(start, {});
-  VerifyBufferDefinedAt(start, {0});
-  VerifyBufferDefinedAt(start, {1});
-  VerifyBufferDefinedAt(start, {2});
-  VerifyNoBufferDefinedAt(start, {0, 0});
-
-  // Verify update buffers:
-  // - {} (top-level tuple) -> defined (always defined for new async-update)
-  // - {2} (state) -> defined
-  // - all other subshapes are compatible with start -> not defined
-  VerifyBufferDefinedAt(update, {});
-  VerifyNoBufferDefinedAt(update, {0});
-  VerifyNoBufferDefinedAt(update, {0, 0});
-  VerifyNoBufferDefinedAt(update, {1});
-  VerifyBufferDefinedAt(update, {2});
-}
-
-TEST_F(LogicalBufferAnalysisTest, AsyncUpdateWithChangedShapes) {
-  absl::string_view hlo_str = R"(
-  HloModule module
-
-  async_computation {
-    p0 = f32[4] parameter(0)
-    ROOT ccall = f32[4] custom-call(p0), custom_call_target="bar"
-  }
-
-  ENTRY entry {
-    p0 = f32[4] parameter(0)
-    start = ((f32[4]), f32[4], s32[]) async-start(p0), calls=async_computation,
-                                       async_execution_thread="sparsecore"
-    update = ((f32[4]), f32[4], s32[]) async-update(start)
-    ROOT done = f32[4] async-done(update)
-  }
-  )";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo_str));
-
-  HloInstruction* update = FindInstruction(module.get(), "update");
-
-  // Mutate the shape of update to ((f32[8]), f32[8], s32[])
-  // This makes it different from start's shape.
-  Shape new_shape = ShapeUtil::MakeTupleShape(
-      {ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(F32, {8})}),
-       ShapeUtil::MakeShape(F32, {8}), ShapeUtil::MakeShape(S32, {})});
-  *update->mutable_shape() = new_shape;
-
-  ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(module.get()));
-
-  // {0, 0} is an input operand (starts with 0, size >= 2).
-  // It is NOT compatible with start's {0, 0} (f32[4] vs f32[8]).
-  // It should NOT define a buffer because it's an input.
-  VerifyNoBufferDefinedAt(update, {0, 0});
-
-  // {1} is an output.
-  // It is NOT compatible with start's {1} (f32[4] vs f32[8]).
-  // It SHOULD define a buffer because it's a new output.
-  VerifyBufferDefinedAt(update, {1});
-}
-
-TEST_F(LogicalBufferAnalysisTest,
-       AsyncUpdateWithExplicitAliasAndChangedShapes) {
-  absl::string_view hlo_str = R"(
-  HloModule module
-
-  async_computation {
-    p0 = f32[4] parameter(0)
-    ROOT ccall = f32[4] custom-call(p0), custom_call_target="bar"
-  }
-
-  ENTRY entry {
-    p0 = f32[4] parameter(0)
-    start = ((f32[4]), f32[4], s32[]) async-start(p0), calls=async_computation,
-                                       async_execution_thread="sparsecore",
-                                       output_to_operand_aliasing={ {1}: (0, {}) }
-    update = ((f32[4]), f32[4], s32[]) async-update(start)
-    ROOT done = f32[4] async-done(update)
-  }
-  )";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo_str));
-
-  HloInstruction* update = FindInstruction(module.get(), "update");
-
-  // Mutate the shape of update to ((f32[4]), f32[8], s32[])
-  // This changes the shape of the aliased output {1} from f32[4] to f32[8].
-  Shape new_shape = ShapeUtil::MakeTupleShape(
-      {ShapeUtil::MakeTupleShape({ShapeUtil::MakeShape(F32, {4})}),
-       ShapeUtil::MakeShape(F32, {8}), ShapeUtil::MakeShape(S32, {})});
-  *update->mutable_shape() = new_shape;
-
-  ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(module.get()));
-
-  // {1} is explicitly aliased.
-  // It is NOT compatible with start's {1} (f32[4] vs f32[8]).
-  // It should NOT define a buffer because it's aliased.
-  VerifyNoBufferDefinedAt(update, {1});
-}
-TEST_F(LogicalBufferAnalysisTest, AsyncAliasedEmptyTupleBecomesValued) {
-  absl::string_view hlo_str = R"(
-  HloModule module
-
-  async_computation {
-    p0 = (f32[4]) parameter(0)
-    ROOT ccall = () custom-call(p0), custom_call_target="bar"
-  }
-
-  ENTRY entry {
-    p0 = (f32[4]) parameter(0)
-    start = (((f32[4])), (), s32[]) async-start(p0), calls=async_computation,
-                                       async_execution_thread="sparsecore",
-                                       output_to_operand_aliasing={ {1}: (0, {}) }
-    update = (((f32[4])), (), s32[]) async-update(start)
-    ROOT done = f32[4] async-done(update)
-  }
-  )";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo_str));
-
-  HloInstruction* start = FindInstruction(module.get(), "start");
-  HloInstruction* update = FindInstruction(module.get(), "update");
-
-  // Mutate update shape to (((f32[4])), (f32[4]), s32[])
-  // This simulates late binding where the empty tuple becomes a non-empty
-  // tuple.
-  Shape f32_4 = ShapeUtil::MakeShape(F32, {4});
-  Shape new_shape = ShapeUtil::MakeTupleShape(
-      {ShapeUtil::MakeTupleShape({ShapeUtil::MakeTupleShape({f32_4})}),
-       ShapeUtil::MakeTupleShape({f32_4}), ShapeUtil::MakeShape(S32, {})});
-  *update->mutable_shape() = new_shape;
-
-  ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(module.get()));
-
-  // For start:
-  // {1} is () and is aliased. It SHOULD get a buffer because it is empty.
-  VerifyBufferDefinedAt(start, {1});
-
-  // For update:
-  // {1} is (f32[4]) (non-empty tuple) and is aliased. It SHOULD NOT get a
-  // buffer.
-  VerifyNoBufferDefinedAt(update, {1});
-
-  // {1, 0} is f32[4] and is NOT aliased. It SHOULD get a buffer.
-  VerifyBufferDefinedAt(update, {1, 0});
-}
-
-TEST_F(LogicalBufferAnalysisTest, AsyncStartLateBinding) {
-  absl::string_view hlo_str = R"(
-  HloModule module
-
-  async_computation {
-    p0 = f32[4] parameter(0)
-    ROOT ccall = f32[4] custom-call(p0), custom_call_target="bar"
-  }
-
-  ENTRY entry {
-    p0 = f32[4] parameter(0)
-    async-start = ((), (), s32[]) async-start(), calls=async_computation,
-                            async_execution_thread="sparsecore"
-    async-update = ((f32[4]), f32[4], s32[]) async-update(async-start, p0)
-    ROOT async-done = f32[4] async-done(async-update)
-  }
-  )";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str));
-  ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(module.get()));
-
-  HloInstruction* p0 = FindInstruction(module.get(), "p0");
-  HloInstruction* async_start = FindInstruction(module.get(), "async-start");
-  HloInstruction* async_update = FindInstruction(module.get(), "async-update");
-  HloInstruction* async_done = FindInstruction(module.get(), "async-done");
-
-  VerifyBufferDefinedAt(p0, {});
-
-  VerifyBufferDefinedAt(async_start, {});
-  VerifyBufferDefinedAt(async_start, {0});
-  VerifyBufferDefinedAt(async_start, {1});
-  VerifyBufferDefinedAt(async_start, {2});
-
-  // Verify update buffers:
-  // - {} (top-level tuple) -> defined (always defined for new async-update)
-  // - {2} (context) -> defined by async-update if it has no explicit aliasing
-  // - all other subshapes are compatible with start -> not defined
-  VerifyBufferDefinedAt(async_update, {});
-  VerifyBufferDefinedAt(async_update, {0});
-  VerifyNoBufferDefinedAt(async_update, {0, 0});
-  // {1} is defined in async-update.
-  VerifyBufferDefinedAt(async_update, {1});
-  VerifyBufferDefinedAt(async_update, {2});
-
-  VerifyBufferDefinedAt(async_done, {});
-}
-
-// Verifies that omitting output_to_operand_aliasing on an AsyncStart wrapping
-// DUS causes LogicalBufferAnalysis to mint a brand-new logical buffer for
-// Output {1} (causing memory doubling and copy insertion), whereas annotating
-// it prevents new buffer definition.
-TEST_F(LogicalBufferAnalysisTest, AsyncDUSBufferAllocationWithoutAliasing) {
-  absl::string_view hlo_str = R"(
-  HloModule module
-
-  async_computation {
-    p0 = f32[1024] parameter(0)
-    p1 = f32[256] parameter(1)
-    p2 = s32[] parameter(2)
-    ROOT dus = f32[1024] dynamic-update-slice(p0, p1, p2)
-  }
-
-  ENTRY entry {
-    p0 = f32[1024] parameter(0)
-    p1 = f32[256] parameter(1)
-    p2 = s32[] parameter(2)
-    start_unannotated = ((f32[1024], f32[256], s32[]), f32[1024], s32[]) async-start(p0, p1, p2),
-        calls=async_computation
-    done_unannotated = f32[1024] async-done(start_unannotated)
-    start_annotated = ((f32[1024], f32[256], s32[]), f32[1024], s32[]) async-start(p0, p1, p2),
-        calls=async_computation, output_to_operand_aliasing={ {1}: (0, {}) }
-    done_annotated = f32[1024] async-done(start_annotated)
-    ROOT tuple = (f32[1024], f32[1024]) tuple(done_unannotated, done_annotated)
-  }
-  )";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str));
-  ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(module.get()));
-
-  HloInstruction* unannotated =
-      FindInstruction(module.get(), "start_unannotated");
-  HloInstruction* annotated = FindInstruction(module.get(), "start_annotated");
-
-  // Output {1} defines a new buffer when unannotated:
-  VerifyBufferDefinedAt(unannotated, {1});
-
-  // Output {1} is aliased to operand 0 when annotated, so NO new buffer is
-  // defined:
-  VerifyNoBufferDefinedAt(annotated, {1});
-}
-
-TEST_F(LogicalBufferAnalysisTest, AsyncStartLateBindingWithAliasing) {
-  absl::string_view hlo_str = R"(
-  HloModule module
-
-  async_computation {
-    p0 = f32[4] parameter(0)
-    ROOT ccall = f32[4] custom-call(p0), custom_call_target="bar"
-  }
-
-  ENTRY entry {
-    p0 = f32[4] parameter(0)
-    async-start = ((), (), s32[]) async-start(), calls=async_computation,
-                            async_execution_thread="sparsecore",
-                            output_to_operand_aliasing={ {1}: (0, {}) }
-    async-update = ((f32[4]), f32[4], s32[]) async-update(async-start, p0)
-    ROOT async-done = f32[4] async-done(async-update)
-  }
-  )";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str));
-  ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(module.get()));
-
-  HloInstruction* p0 = FindInstruction(module.get(), "p0");
-  HloInstruction* async_start = FindInstruction(module.get(), "async-start");
-  HloInstruction* async_update = FindInstruction(module.get(), "async-update");
-  HloInstruction* async_done = FindInstruction(module.get(), "async-done");
-
-  VerifyBufferDefinedAt(p0, {});
-
-  VerifyBufferDefinedAt(async_start, {});
-  VerifyBufferDefinedAt(async_start, {0});
-  // {1} (output) is not bound yet, an empty buffer is defined
-  VerifyBufferDefinedAt(async_start, {1});
-  VerifyBufferDefinedAt(async_start, {2});
-
-  // Verify update buffers:
-  // - {}  (top-level tuple) -> defined (always defined for new async-update)
-  // - {2} (context) -> defined by async-update if it has no explicit aliasing
-  // - all other subshapes are compatible with start -> not defined
-  VerifyBufferDefinedAt(async_update, {});
-  VerifyBufferDefinedAt(async_update, {0});
-  VerifyNoBufferDefinedAt(async_update, {0, 0});
-  VerifyNoBufferDefinedAt(async_update, {1});
-  VerifyBufferDefinedAt(async_update, {2});
-
-  VerifyBufferDefinedAt(async_done, {});
-}
-
-TEST_F(LogicalBufferAnalysisTest, AsyncStartLateBindingMultipleUpdates) {
-  absl::string_view hlo_str = R"(
-  HloModule module
-
-  async_computation {
-    p0 = f32[4] parameter(0)
-    ROOT ccall = f32[4] custom-call(p0), custom_call_target="bar"
-  }
-
-  ENTRY entry {
-    p0 = f32[4] parameter(0)
-    async-start = ((), (), s32[]) async-start(), calls=async_computation,
-                            async_execution_thread="sparsecore",
-                            output_to_operand_aliasing={ {1}: (0, {}) }
-    async-update0 = ((), (), s32[]) async-update(async-start),
-                            output_to_operand_aliasing={ {2}: (0, {2}) }
-    async-update1 = ((f32[4]), f32[4], s32[]) async-update(async-update0, p0),
-                            output_to_operand_aliasing={ {2}: (0, {2}) }
-    ROOT async-done = f32[4] async-done(async-update1)
-  }
-  )";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_str));
-  ASSERT_OK_AND_ASSIGN(analysis_, LogicalBufferAnalysis::Run(module.get()));
-
-  HloInstruction* p0 = FindInstruction(module.get(), "p0");
-  HloInstruction* async_start = FindInstruction(module.get(), "async-start");
-  HloInstruction* async_update0 =
-      FindInstruction(module.get(), "async-update0");
-  HloInstruction* async_update1 =
-      FindInstruction(module.get(), "async-update1");
-  HloInstruction* async_done = FindInstruction(module.get(), "async-done");
-
-  VerifyBufferDefinedAt(p0, {});
-
-  VerifyBufferDefinedAt(async_start, {});
-  VerifyBufferDefinedAt(async_start, {0});
-  // {1} (output) is not bound yet, an empty buffer is defined
-  VerifyBufferDefinedAt(async_start, {1});
-  VerifyBufferDefinedAt(async_start, {2});
-
-  // Verify update buffers:
-  // {} (top-level tuple) -> defined (always defined for new async-update)
-  VerifyBufferDefinedAt(async_update0, {});
-  // Not defined due to implicit aliasing from async-start.
-  VerifyNoBufferDefinedAt(async_update0, {0});
-  VerifyNoBufferDefinedAt(async_update0, {1});
-  // Not defined due to explicit aliasing from async-start.
-  VerifyNoBufferDefinedAt(async_update0, {2});
-
-  VerifyBufferDefinedAt(async_update1, {});
-  VerifyBufferDefinedAt(async_update1, {0});
-  // Not defined due to implicit aliasing from async-start.
-  VerifyNoBufferDefinedAt(async_update1, {0, 0});
-  VerifyNoBufferDefinedAt(async_update1, {1});
-  // Not defined due to explicit aliasing from async_update0.
-  VerifyNoBufferDefinedAt(async_update1, {2});
-
-  VerifyBufferDefinedAt(async_done, {});
-}
-}  // namespace
-}  // namespace xla
diff --git a/third_party/xla/xla/hlo/analysis/tuple_points_to_analysis.cc b/third_party/xla/xla/hlo/analysis/tuple_points_to_analysis.cc
deleted file mode 100644
index e7e42bd3fe9..00000000000
--- a/third_party/xla/xla/hlo/analysis/tuple_points_to_analysis.cc
+++ /dev/null
@@ -1,1105 +0,0 @@
-/* Copyright 2017 The OpenXLA Authors.
-
-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.
-==============================================================================*/
-
-#include "xla/hlo/analysis/tuple_points_to_analysis.h"
-
-#include <cstddef>
-#include <cstdint>
-#include <memory>
-#include <ostream>
-#include <string>
-#include <utility>
-#include <vector>
-
-#include "absl/algorithm/container.h"
-#include "absl/container/flat_hash_map.h"
-#include "absl/container/flat_hash_set.h"
-#include "absl/log/check.h"
-#include "absl/status/status.h"
-#include "absl/status/status_macros.h"
-#include "absl/status/statusor.h"
-#include "absl/strings/str_cat.h"
-#include "absl/strings/str_format.h"
-#include "absl/strings/str_join.h"
-#include "absl/types/span.h"
-#include "xla/hlo/analysis/logical_buffer_analysis.h"
-#include "xla/hlo/ir/hlo_casting_utils.h"
-#include "xla/hlo/ir/hlo_instruction.h"
-#include "xla/hlo/ir/hlo_instructions.h"
-#include "xla/hlo/ir/hlo_opcode.h"
-#include "xla/service/logical_buffer.h"
-#include "xla/shape.h"
-#include "xla/shape_util.h"
-#include "xla/tsl/platform/logging.h"
-#include "xla/util.h"
-
-namespace xla {
-
-std::string BufferAlias::ToString() const {
-  return absl::StrCat("BufferAlias(", instruction_->name(), "[",
-                      absl::StrJoin(index_, ","), "])");
-}
-
-std::ostream& operator<<(std::ostream& out, const BufferAlias& buffer_alias) {
-  out << buffer_alias.ToString();
-  return out;
-}
-
-bool PointsToSet::IsAmbiguous() const {
-  bool ambiguous = false;
-  ForEachElement(
-      [&ambiguous](const ShapeIndex& /*index*/, const BufferList& points_to) {
-        ambiguous |= points_to.size() > 1;
-      });
-  return ambiguous;
-}
-
-bool PointsToSet::IsDistinct() const {
-  bool distinct = true;
-  absl::flat_hash_set<const LogicalBuffer*> all_points_to;
-  ForEachElement([&](const ShapeIndex& /*index*/, const BufferList& points_to) {
-    for (auto& buffer : points_to) {
-      if (all_points_to.contains(buffer)) {
-        distinct = false;
-      }
-      all_points_to.insert(buffer);
-    }
-  });
-  return distinct;
-}
-
-size_t PointsToSet::size() const {
-  // Because pointed-to elements may be duplicated we have to create a flattened
-  // set and return the size.
-  return CreateFlattenedSet().size();
-}
-
-PointsToSet::BufferSet PointsToSet::CreateFlattenedSet() const {
-  BufferSet flat_set;
-  ForEachElement(
-      [&flat_set](const ShapeIndex& /*index*/, const BufferList& buffers) {
-        flat_set.insert(buffers.begin(), buffers.end());
-      });
-  return flat_set;
-}
-
-bool PointsToSet::ContainsBuffer(const LogicalBuffer& buffer) const {
-  bool found = false;
-  ForEachElement([&found, &buffer](const ShapeIndex& /*index*/,
-                                   const BufferList& pointed_to_buffers) {
-    if (!found && absl::c_linear_search(pointed_to_buffers, &buffer)) {
-      found = true;
-    }
-  });
-  return found;
-}
-
-bool PointsToSet::ContainsBufferAtIndex(const LogicalBuffer& buffer,
-                                        const ShapeIndex& index) const {
-  const auto& pointed_to_buffers = element(index);
-  return absl::c_linear_search(pointed_to_buffers, &buffer);
-}
-
-void PointsToSet::AddPointedToBuffer(const LogicalBuffer& buffer,
-                                     const ShapeIndex& index) {
-  if (ContainsBufferAtIndex(buffer, index)) {
-    return;
-  }
-  mutable_element(index)->push_back(&buffer);
-}
-
-const PointsToSet::SourceSet& PointsToSet::tuple_sources(
-    const ShapeIndex& index) const {
-  return tree_.element(index).tuple_sources;
-}
-
-void PointsToSet::add_tuple_source(const ShapeIndex& index,
-                                   HloInstruction* tuple) {
-  tree_.mutable_element(index)->tuple_sources.insert(tuple);
-}
-
-namespace {
-// Gather fusion instructions from 'instruction' into 'fusion_instructions'.
-void GatherFusionInstructions(
-    HloInstruction* instruction,
-    std::vector<HloInstruction*>* fusion_instructions) {
-  CHECK_EQ(HloOpcode::kFusion, instruction->opcode());
-  for (auto* fused : instruction->fused_instructions()) {
-    if (fused->opcode() == HloOpcode::kFusion) {
-      GatherFusionInstructions(fused, fusion_instructions);
-    }
-  }
-  fusion_instructions->push_back(instruction);
-}
-
-}  // namespace
-
-/* static */ absl::StatusOr<std::unique_ptr<TuplePointsToAnalysis>>
-TuplePointsToAnalysis::Run(const HloModule* module) {
-  ABSL_ASSIGN_OR_RETURN(auto logical_buffer_analysis,
-                   LogicalBufferAnalysis::Run(module));
-  std::unique_ptr<TuplePointsToAnalysis> analysis(
-      new TuplePointsToAnalysis(module, std::move(logical_buffer_analysis)));
-  ABSL_RETURN_IF_ERROR(analysis->Analyze());
-  return analysis;
-}
-
-absl::Status TuplePointsToAnalysis::Analyze() {
-  per_instruction_.clear();
-  int64_t max_computation_id = -1;
-  for (const HloComputation* computation : module_->computations()) {
-    if (computation->unique_id() > max_computation_id) {
-      max_computation_id = computation->unique_id();
-    }
-  }
-  per_instruction_.resize(max_computation_id + 1);
-
-  logical_buffer_aliases_.clear();
-  logical_buffer_aliases_.resize(
-      logical_buffer_analysis_->num_logical_buffers());
-
-  std::vector<HloInstruction*> fusion_instructions;
-  for (auto* computation : module_->MakeNonfusionComputations()) {
-    ABSL_RETURN_IF_ERROR(computation->Accept(this));
-    ABSL_RETURN_IF_ERROR(
-        PopulateDefinedBuffersAndAliases(computation->instructions()));
-    for (auto* instruction : computation->instructions()) {
-      if (instruction->opcode() == HloOpcode::kFusion) {
-        GatherFusionInstructions(instruction, &fusion_instructions);
-      }
-    }
-  }
-  // Run points-to analysis on fusion instructions in 'computation'.
-  for (auto* instruction : fusion_instructions) {
-    ABSL_RETURN_IF_ERROR(
-        instruction->fused_instructions_computation()->Accept(this));
-    ABSL_RETURN_IF_ERROR(
-        PopulateDefinedBuffersAndAliases(instruction->fused_instructions()));
-  }
-
-  XLA_VLOG_LINES(3, ToString());
-
-  return absl::OkStatus();
-}
-
-absl::Status TuplePointsToAnalysis::PopulateDefinedBuffersAndAliases(
-    const decltype(std::declval<HloComputation>()
-                       .instructions())& instructions) {
-  for (auto* instruction : instructions) {
-    PerInstruction* pi = PerInst(instruction);
-    ABSL_RETURN_IF_ERROR(GatherBuffersDefinedByInstruction(
-        instruction, &pi->instruction_defined_buffers));
-
-    const PointsToSet& points_to_set = GetPointsToSet(instruction);
-    points_to_set.ForEachElement(
-        [this, &instruction](
-            const ShapeIndex& index,
-            const PointsToSet::BufferList& pointed_to_buffers) {
-          for (const LogicalBuffer* buffer : pointed_to_buffers) {
-            logical_buffer_aliases_[buffer->id()].emplace_back(instruction,
-                                                               index);
-          }
-        });
-  }
-  return absl::OkStatus();
-}
-
-absl::Status TuplePointsToAnalysis::DefaultAction(
-    HloInstruction* hlo_instruction) {
-  // Create trivial points-to set for instruction. Each points-to set at index i
-  // contains a single element LogicalBuffer(hlo_instruction, i). This indicates
-  // that this instruction is the source of all buffers in its own output.
-  PointsToSet& points_to_set = CreateEmptyPointsToSet(hlo_instruction);
-  ABSL_RETURN_IF_ERROR(points_to_set.ForEachMutableElementWithStatus(
-      [this, hlo_instruction](
-          const ShapeIndex& index,
-          PointsToSet::BufferList* buffers) -> absl::Status {
-        ABSL_ASSIGN_OR_RETURN(
-            LogicalBuffer * buffer,
-            logical_buffer_analysis_->GetBuffer(hlo_instruction, index));
-        buffers->push_back(buffer);
-        return absl::OkStatus();
-      }));
-
-  if (hlo_instruction->shape().IsTuple()) {
-    // If the hlo instruction is a tuple-shaped, then trivially the instruction
-    // itself is the source of the tuple.
-    points_to_set.add_tuple_source({}, hlo_instruction);
-  }
-
-  return absl::OkStatus();
-}
-
-absl::Status TuplePointsToAnalysis::HandleGetTupleElement(
-    HloInstruction* get_tuple_element) {
-  // GetTupleElement forwards a pointer to a particular element of the tuple
-  // operand.
-  int64_t element_index = get_tuple_element->tuple_index();
-
-  PointsToSet& points_to_set = CreateEmptyPointsToSet(get_tuple_element);
-  const PointsToSet& operand_points_to_set =
-      *PerInst(get_tuple_element->operand(0))->points_to_set;
-
-  // Forward the points-to set (and tuple sources) at index {element_index}
-  // of the operand to the points-to set for this GetTupleElement instruction.
-  points_to_set.ForEachMutableElement(
-      [&](const ShapeIndex& target_index, PointsToSet::BufferList* points_to) {
-        // Construct an index into the operand by prepending element_index to
-        // the index for the GetTupleElement instruction's points-to set.
-        ShapeIndex src_index({element_index});
-        src_index.insert(src_index.end(), target_index.begin(),
-                         target_index.end());
-
-        // Forward the pointed-to buffers.
-        *points_to = operand_points_to_set.element(src_index);
-
-        // Forward the tuple sources.
-        for (HloInstruction* tuple :
-             operand_points_to_set.tuple_sources(src_index)) {
-          points_to_set.add_tuple_source(target_index, tuple);
-        }
-      });
-
-  return absl::OkStatus();
-}
-
-absl::Status TuplePointsToAnalysis::HandleCopy(HloInstruction* copy) {
-  // A kCopy instruction performs a shallow copy of the operand. The top-level
-  // buffer (index={}) is newly created, but all other buffers (in the case of a
-  // tuple shape) come from the operand
-  PointsToSet& points_to_set = CreateCopiedPointsToSet(copy, copy->operand(0));
-  points_to_set.mutable_element(/*index=*/{})->clear();
-  ABSL_ASSIGN_OR_RETURN(LogicalBuffer * buffer,
-                   logical_buffer_analysis_->GetBuffer(copy, /*index=*/{}));
-  points_to_set.AddPointedToBuffer(*buffer, /*index=*/{});
-
-  return absl::OkStatus();
-}
-
-absl::Status TuplePointsToAnalysis::HandleBitcast(HloInstruction* bitcast) {
-  // A kBitcast instruction aliases its operand. That is, the buffer of its
-  // result *is* the buffer of its operand, so just copy the operands points-to
-  // set.
-  CreateCopiedPointsToSet(bitcast, bitcast->operand(0));
-  return absl::OkStatus();
-}
-
-absl::Status TuplePointsToAnalysis::HandleDomain(HloInstruction* domain) {
-  // A kDomain instruction aliases its operand. That is, the buffer of its
-  // result *is* the buffer of its operand, so just copy the operands points-to
-  // set.
-  CreateCopiedPointsToSet(domain, domain->operand(0));
-  return absl::OkStatus();
-}
-
-absl::Status TuplePointsToAnalysis::HandleAddDependency(
-    HloInstruction* add_dependency) {
-  // AddDependency just forwards the value of its zero-th operand.
-  CreateCopiedPointsToSet(add_dependency, add_dependency->operand(0));
-  return absl::OkStatus();
-}
-
-absl::Status TuplePointsToAnalysis::HandleRecvDone(HloInstruction* recv_done) {
-  // RecvDone aliases its input (Recv) tuple element {0} to element {0} of its
-  // output. The other indices ({} and {1}) define their own buffers.
-  PointsToSet& points_to_set = CreateEmptyPointsToSet(recv_done);
-  ABSL_ASSIGN_OR_RETURN(
-      LogicalBuffer * buffer_0,
-      logical_buffer_analysis_->GetBuffer(recv_done, /*index=*/{}));
-  points_to_set.AddPointedToBuffer(*buffer_0, /*index=*/{});
-  ABSL_ASSIGN_OR_RETURN(
-      LogicalBuffer * buffer_1,
-      logical_buffer_analysis_->GetBuffer(recv_done, /*index=*/{1}));
-  points_to_set.AddPointedToBuffer(*buffer_1, /*index=*/{1});
-
-  const PointsToSet& operand_points_to_set =
-      GetPointsToSet(recv_done->operand(0));
-
-  // Recursively copy the points to set of the operand tuple {0} to the output
-  // element {0}.
-  points_to_set.ForEachMutableElement(
-      [&points_to_set, &operand_points_to_set](
-          const ShapeIndex& index, PointsToSet::BufferList* buffers) {
-        if (index.empty() || index[0] != 0) {
-          return;
-        }
-        *buffers = operand_points_to_set.element(index);
-        for (auto& tuple_source : operand_points_to_set.tuple_sources(index)) {
-          points_to_set.add_tuple_source(index, tuple_source);
-        }
-      });
-  return absl::OkStatus();
-}
-
-absl::Status TuplePointsToAnalysis::HandleAsyncStart(
-    HloInstruction* async_start) {
-  PointsToSet& points_to_set = CreateEmptyPointsToSet(async_start);
-
-  // 1. Get the list of explicitly aliased outputs.
-  absl::flat_hash_map<ShapeIndex, std::pair<int64_t, ShapeIndex>>
-      aliased_outputs;
-  for (const auto& pair : Cast<HloAsyncStartInstruction>(async_start)
-                              ->output_to_operand_aliasing()) {
-    aliased_outputs.insert(pair);
-  }
-  ABSL_RETURN_IF_ERROR(points_to_set.ForEachMutableElementWithStatus(
-      [&](const ShapeIndex& target_index,
-          PointsToSet::BufferList* buffers) -> absl::Status {
-        auto it = aliased_outputs.find(target_index);
-        bool has_implicit_alias =
-            (target_index.size() >= 2 && target_index.front() == 0);
-        bool has_explicit_alias = (it != aliased_outputs.end());
-
-        if (has_implicit_alias) {
-          // Implicit alias: forward from operand.
-          // For async-start, logical input i is operand i.
-          int64_t input_idx = target_index[1];
-          CHECK_LT(input_idx, async_start->operand_count());
-          const PointsToSet& operand_points_to_set =
-              GetPointsToSet(async_start->operand(input_idx));
-          ShapeIndex source_index(target_index.begin() + 2, target_index.end());
-          *buffers = operand_points_to_set.element(source_index);
-          for (HloInstruction* tuple :
-               operand_points_to_set.tuple_sources(source_index)) {
-            points_to_set.add_tuple_source(target_index, tuple);
-          }
-        }
-        if (has_explicit_alias &&
-            it->second.first < async_start->operand_count()) {
-          const PointsToSet& input_set =
-              GetPointsToSet(async_start->operand(it->second.first));
-          for (const LogicalBuffer* input_buffer :
-               input_set.element(it->second.second)) {
-            points_to_set.AddPointedToBuffer(*input_buffer, target_index);
-          }
-          for (HloInstruction* tuple :
-               input_set.tuple_sources(it->second.second)) {
-            points_to_set.add_tuple_source(target_index, tuple);
-          }
-        }
-        if (!has_implicit_alias && !has_explicit_alias) {
-          ABSL_ASSIGN_OR_RETURN(
-              LogicalBuffer * buffer,
-              logical_buffer_analysis_->GetBuffer(async_start, target_index));
-          buffers->push_back(buffer);
-        }
-        return absl::OkStatus();
-      }));
-  points_to_set.add_tuple_source({}, async_start);
-  return absl::OkStatus();
-}
-
-absl::Status TuplePointsToAnalysis::HandleAsyncUpdate(
-    HloInstruction* async_update) {
-  PointsToSet& points_to_set = CreateEmptyPointsToSet(async_update);
-
-  const Shape& async_update_shape = async_update->shape();
-  const HloInstruction* prev_async_op = async_update->operand(0);
-  const Shape& prev_async_op_shape = prev_async_op->shape();
-
-  // The top-level shape must be a tuple.
-  CHECK(async_update_shape.IsTuple());
-  CHECK(prev_async_op_shape.IsTuple());
-
-  // 1. Get explicitly aliased outputs from async-start. When crossing loop
-  // boundaries, async_chain_start() may be null; safely return if unreachable.
-  const HloAsyncInstruction* async_update_inst =
-      (async_update != nullptr) ? DynCast<HloAsyncInstruction>(async_update)
-                                : nullptr;
-  if (async_update_inst == nullptr) {
-    return absl::OkStatus();
-  }
-  HloInstruction* async_start = async_update_inst->async_chain_start();
-  const HloAsyncStartInstruction* async_start_inst =
-      (async_start != nullptr) ? DynCast<HloAsyncStartInstruction>(async_start)
-                               : nullptr;
-  if (async_start_inst == nullptr) {
-    return absl::OkStatus();
-  }
-  absl::flat_hash_set<ShapeIndex> explicitly_aliased_outputs;
-  for (const auto& pair : async_start_inst->output_to_operand_aliasing()) {
-    explicitly_aliased_outputs.insert(pair.first);
-  }
-
-  // Get the list of aliased contexts.
-  absl::flat_hash_map<ShapeIndex, ShapeIndex> aliased_contexts;
-  for (const auto& [output_index, operand_info] :
-       Cast<HloAsyncUpdateInstruction>(async_update)
-           ->output_to_operand_aliasing()) {
-    const auto& [_, operand_index] = operand_info;
-    aliased_contexts.emplace(output_index, operand_index);
-  }
-
-  ABSL_RETURN_IF_ERROR(points_to_set.ForEachMutableElementWithStatus(
-      [&](const ShapeIndex& index,
-          PointsToSet::BufferList* buffers) -> absl::Status {
-        const Shape& new_subshape =
-            ShapeUtil::GetSubshape(async_update_shape, index);
-
-        // 1. Handle values which are forwarded from previous async instruction.
-        bool forwarded_operand_or_output =
-            !index.empty() && index.front() < 2 &&
-            ShapeUtil::IndexIsValid(prev_async_op_shape, index) &&
-            ShapeUtil::Compatible(
-                ShapeUtil::GetSubshape(prev_async_op_shape, index),
-                new_subshape);
-        if (forwarded_operand_or_output) {
-          const PointsToSet& operand_points_to_set =
-              GetPointsToSet(prev_async_op);
-          *buffers = operand_points_to_set.element(index);
-          for (HloInstruction* tuple :
-               operand_points_to_set.tuple_sources(index)) {
-            points_to_set.add_tuple_source(index, tuple);
-          }
-          return absl::OkStatus();
-        }
-        if (auto aliased_context_it = aliased_contexts.find(index);
-            aliased_context_it != aliased_contexts.end()) {
-          const auto& [_, prev_async_op_index] = *aliased_context_it;
-          const PointsToSet& operand_points_to_set =
-              GetPointsToSet(prev_async_op);
-          *buffers = operand_points_to_set.element(prev_async_op_index);
-          for (HloInstruction* tuple :
-               operand_points_to_set.tuple_sources(prev_async_op_index)) {
-            points_to_set.add_tuple_source(index, tuple);
-          }
-          return absl::OkStatus();
-        }
-
-        // 2. If not forwarded, it is a new element.
-        bool has_explicit_alias = explicitly_aliased_outputs.contains(index);
-
-        if (index.size() >= 2 && index.front() == 0) {
-          // New input: forward from the corresponding operand of async_update.
-          // operand(0) is the chain. New inputs are bound starting at
-          // operand(1). We need to find which logical operand this index
-          // corresponds to.
-          int64_t logical_operand_idx = index[1];
-          const Shape& prev_inputs_tuple_shape =
-              prev_async_op_shape.tuple_shapes(0);
-          int64_t prev_inputs_count =
-              prev_inputs_tuple_shape.tuple_shapes().size();
-
-          // The operand index in async_update is 1 + (logical_operand_idx -
-          // prev_inputs_count).
-          int64_t operand_idx = 1 + (logical_operand_idx - prev_inputs_count);
-
-          if (operand_idx > 0 && operand_idx < async_update->operand_count()) {
-            const PointsToSet& input_set =
-                GetPointsToSet(async_update->operand(operand_idx));
-            ShapeIndex source_index(index.begin() + 2, index.end());
-            *buffers = input_set.element(source_index);
-            for (HloInstruction* tuple :
-                 input_set.tuple_sources(source_index)) {
-              points_to_set.add_tuple_source(index, tuple);
-            }
-          }
-        } else if (has_explicit_alias) {
-          // Explicit alias: do nothing for now, will be resolved in step 3 if
-          // active.
-        } else {
-          // New output or new tuple container.
-          ABSL_ASSIGN_OR_RETURN(
-              LogicalBuffer * buffer,
-              logical_buffer_analysis_->GetBuffer(async_update, index));
-          buffers->push_back(buffer);
-          if (new_subshape.IsTuple()) {
-            points_to_set.add_tuple_source(index, async_update);
-          }
-        }
-        return absl::OkStatus();
-      }));
-
-  // 3. Apply deferred aliases.
-  ApplyDeferredAliases(async_update, points_to_set);
-
-  return absl::OkStatus();
-}
-
-absl::Status TuplePointsToAnalysis::HandleAsyncDone(
-    HloInstruction* async_done) {
-  PointsToSet& points_to_set = CreateEmptyPointsToSet(async_done);
-  const HloInstruction* async_op = async_done->operand(0);
-  const PointsToSet& operand_points_to_set = GetPointsToSet(async_op);
-  const Shape& async_op_shape = async_op->shape();
-
-  ABSL_RETURN_IF_ERROR(points_to_set.ForEachMutableElementWithStatus(
-      [&](const ShapeIndex& target_index,
-          PointsToSet::BufferList* buffers) -> absl::Status {
-        ShapeIndex src_index({1});
-        src_index.insert(src_index.end(), target_index.begin(),
-                         target_index.end());
-
-        // AsyncDone forwards the result of the async operation, which is
-        // located at index {1} of the async-start/update output tuple. However,
-        // with late binding, the async operation might not bind the output
-        // (e.g., if the output is not yet known or empty), in which case index
-        // {1} of the async operation shape might be invalid or have an
-        // incompatible shape (like empty tuple ()).
-        //
-        // We only forward if the source index is valid and the shapes are
-        // compatible. Otherwise, we treat AsyncDone as defining its own new
-        // buffer.
-        if (ShapeUtil::IndexIsValid(async_op_shape, src_index) &&
-            ShapeUtil::Compatible(
-                ShapeUtil::GetSubshape(async_done->shape(), target_index),
-                ShapeUtil::GetSubshape(async_op_shape, src_index))) {
-          *buffers = operand_points_to_set.element(src_index);
-          for (HloInstruction* tuple :
-               operand_points_to_set.tuple_sources(src_index)) {
-            points_to_set.add_tuple_source(target_index, tuple);
-          }
-        } else {
-          ABSL_ASSIGN_OR_RETURN(
-              LogicalBuffer * buffer,
-              logical_buffer_analysis_->GetBuffer(async_done, target_index));
-          buffers->push_back(buffer);
-          if (ShapeUtil::GetSubshape(async_done->shape(), target_index)
-                  .IsTuple()) {
-            points_to_set.add_tuple_source(target_index, async_done);
-          }
-        }
-        return absl::OkStatus();
-      }));
-
-  // 2. Apply deferred aliases.
-  ApplyDeferredAliases(async_done, points_to_set);
-
-  return absl::OkStatus();
-}
-
-absl::Status TuplePointsToAnalysis::HandleCopyStart(
-    HloInstruction* copy_start) {
-  // CopyStart forwards its aliased operand to {1}.
-  PointsToSet& points_to_set = CreateEmptyPointsToSet(copy_start);
-  const PointsToSet& operand_points_to_set =
-      GetPointsToSet(copy_start->operand(0));
-
-  ABSL_RETURN_IF_ERROR(points_to_set.ForEachMutableElementWithStatus(
-      [&](const ShapeIndex& target_index,
-          PointsToSet::BufferList* buffers) -> absl::Status {
-        if (target_index == ShapeIndex({1})) {
-          *buffers = operand_points_to_set.element(/*index=*/{});
-        } else {
-          ABSL_ASSIGN_OR_RETURN(
-              LogicalBuffer * buffer,
-              logical_buffer_analysis_->GetBuffer(copy_start, target_index));
-          buffers->push_back(buffer);
-        }
-        return absl::OkStatus();
-      }));
-
-  for (HloInstruction* tuple :
-       operand_points_to_set.tuple_sources(/*index=*/{})) {
-    points_to_set.add_tuple_source(/*index=*/{1}, tuple);
-  }
-
-  return absl::OkStatus();
-}
-
-absl::Status TuplePointsToAnalysis::HandleCopyDone(HloInstruction* copy_done) {
-  // CopyDone forwards its aliased operand.
-  PointsToSet& points_to_set = CreateEmptyPointsToSet(copy_done);
-  const PointsToSet& operand_points_to_set =
-      GetPointsToSet(copy_done->operand(0));
-  operand_points_to_set.ForEachElement(
-      [&points_to_set, &operand_points_to_set](
-          const ShapeIndex& src_index,
-          const PointsToSet::BufferList& points_to) {
-        if (src_index == ShapeIndex({0})) {
-          const ShapeIndex target_index = {};
-          *points_to_set.mutable_element(target_index) = points_to;
-
-          for (HloInstruction* tuple :
-               operand_points_to_set.tuple_sources(src_index)) {
-            points_to_set.add_tuple_source(target_index, tuple);
-          }
-        }
-      });
-
-  return absl::OkStatus();
-}
-
-absl::Status TuplePointsToAnalysis::HandleSend(HloInstruction* send) {
-  // Send creates a tuple of {aliased operand, U32 context, token}.
-  PointsToSet& points_to_set = CreateEmptyPointsToSet(send);
-
-  // Creates the points to set for the tuple and its element at {1}.
-  auto top_buffer = points_to_set.mutable_element(ShapeIndex({}));
-  ABSL_ASSIGN_OR_RETURN(LogicalBuffer * buffer_0,
-                   logical_buffer_analysis_->GetBuffer(send, ShapeIndex({})));
-  top_buffer->push_back(buffer_0);
-  points_to_set.add_tuple_source({}, send);
-
-  auto context_buffer = points_to_set.mutable_element(ShapeIndex({1}));
-  ABSL_ASSIGN_OR_RETURN(LogicalBuffer * buffer_1,
-                   logical_buffer_analysis_->GetBuffer(send, ShapeIndex({1})));
-  context_buffer->push_back(buffer_1);
-
-  auto token_buffer = points_to_set.mutable_element(ShapeIndex({2}));
-  ABSL_ASSIGN_OR_RETURN(LogicalBuffer * buffer_2,
-                   logical_buffer_analysis_->GetBuffer(send, ShapeIndex({2})));
-  token_buffer->push_back(buffer_2);
-
-  // Recursively copy the points to set of the operand to output tuple {0}.
-  const PointsToSet& operand_points_to_set = GetPointsToSet(send->operand(0));
-  operand_points_to_set.ForEachElement(
-      [&points_to_set, &operand_points_to_set](
-          const ShapeIndex& src_index,
-          const PointsToSet::BufferList& points_to) {
-        ShapeIndex target_index({0});
-        for (auto element : src_index) {
-          target_index.push_back(element);
-        }
-        *points_to_set.mutable_element(target_index) = points_to;
-
-        for (HloInstruction* tuple :
-             operand_points_to_set.tuple_sources(src_index)) {
-          points_to_set.add_tuple_source(target_index, tuple);
-        }
-      });
-
-  return absl::OkStatus();
-}
-
-absl::Status TuplePointsToAnalysis::ConstructPointsToSetByAggregatingOperands(
-    HloInstruction* instruction) {
-  absl::Span<HloInstruction* const> operands(instruction->operands());
-  PointsToSet& points_to_set = CreateEmptyPointsToSet(instruction);
-  ABSL_ASSIGN_OR_RETURN(LogicalBuffer * buffer, logical_buffer_analysis_->GetBuffer(
-                                               instruction, /*index=*/{}));
-  points_to_set.AddPointedToBuffer(*buffer, /*index=*/{});
-
-  // A tuple/composite contains references to all input operands and
-  // transitively any references in those operands.
-  for (int64_t i = 0; i < operands.size(); ++i) {
-    const PointsToSet& operand_points_to_set =
-        *PerInst(operands[i])->points_to_set;
-
-    // Forward the points-to set and tuple sources of the operand into the
-    // respective subtree of the output instructions points-to set.
-    operand_points_to_set.ForEachElement(
-        [&points_to_set, &operand_points_to_set, i](
-            const ShapeIndex& src_index,
-            const PointsToSet::BufferList& points_to) {
-          ShapeIndex target_index({i});
-          target_index.insert(target_index.end(), src_index.begin(),
-                              src_index.end());
-
-          // Forward the points-to set
-          *points_to_set.mutable_element(target_index) = points_to;
-
-          // Forward the tuple sources.
-          for (HloInstruction* tuple :
-               operand_points_to_set.tuple_sources(src_index)) {
-            points_to_set.add_tuple_source(target_index, tuple);
-          }
-        });
-  }
-
-  // Add the instruction as a tuple source for the top-level tuple/composite.
-  points_to_set.add_tuple_source({}, instruction);
-
-  return absl::OkStatus();
-}
-
-absl::Status TuplePointsToAnalysis::HandleTuple(HloInstruction* tuple) {
-  return ConstructPointsToSetByAggregatingOperands(tuple);
-}
-
-absl::Status TuplePointsToAnalysis::HandleCustomCall(
-    HloInstruction* custom_call) {
-  auto ccall = Cast<HloCustomCallInstruction>(custom_call);
-  if (ccall->custom_call_target() == kCallMarkerBeforeTarget) {
-    // kCallMarkerBeforeTarget is a CustomCall that takes N operands and returns
-    // a tuple containing those operands. It acts as a metadata placeholder for
-    // function arguments before inlining. Since it behaves exactly like a Tuple
-    // instruction, we reuse ConstructPointsToSetByAggregatingOperands.
-    return ConstructPointsToSetByAggregatingOperands(custom_call);
-  }
-  if (ccall->custom_call_target() == kCallMarkerAfterTarget) {
-    // kCallMarkerAfterTarget is a CustomCall that takes exactly 1 operand (the
-    // result of the function call) and simply forwards it as its output. Thus,
-    // its PointsToSet is simply a direct copy of its operand's PointsToSet.
-    if (custom_call->operand_count() != 1) {
-      return absl::InvalidArgumentError(absl::StrCat(
-          "kCallMarkerAfterTarget must have exactly 1 operand, but got ",
-          custom_call->operand_count()));
-    }
-    CreateCopiedPointsToSet(custom_call, custom_call->operand(0));
-    return absl::OkStatus();
-  }
-  PointsToSet& points_to_set = CreateEmptyPointsToSet(custom_call);
-  absl::flat_hash_map<ShapeIndex, std::pair<int64_t, ShapeIndex>>
-      aliased_outputs;
-  for (const auto& pair : ccall->output_to_operand_aliasing()) {
-    aliased_outputs.emplace(pair.first, pair.second);
-  }
-  ABSL_RETURN_IF_ERROR(points_to_set.ForEachMutableElementWithStatus(
-      [&](const ShapeIndex& index,
-          PointsToSet::BufferList* buffers) -> absl::Status {
-        auto it = aliased_outputs.find(index);
-        if (it == aliased_outputs.end() || !alias_buffer_across_dataflow_) {
-          ABSL_ASSIGN_OR_RETURN(
-              LogicalBuffer * buffer,
-              logical_buffer_analysis_->GetBuffer(custom_call, index));
-          points_to_set.AddPointedToBuffer(*buffer, index);
-        } else {
-          const PointsToSet& input_set =
-              *PerInst(ccall->operand(it->second.first))->points_to_set;
-          for (const LogicalBuffer* input_buffer :
-               input_set.element(it->second.second)) {
-            points_to_set.AddPointedToBuffer(*input_buffer, index);
-          }
-
-          for (HloInstruction* tuple :
-               input_set.tuple_sources(it->second.second)) {
-            points_to_set.add_tuple_source(index, tuple);
-          }
-        }
-        return absl::OkStatus();
-      }));
-  points_to_set.add_tuple_source({}, custom_call);
-  return absl::OkStatus();
-}
-
-// WARNING:
-// Adding this, which essentially does the same thing as HandleCustomCall
-// Not sure if it is really needed or it will break anything
-absl::Status TuplePointsToAnalysis::HandleFusion(HloInstruction* fusion) {
-  auto cfusion = Cast<HloFusionInstruction>(fusion);
-  PointsToSet& points_to_set = CreateEmptyPointsToSet(fusion);
-  absl::flat_hash_map<ShapeIndex, std::pair<int64_t, ShapeIndex>>
-      aliased_outputs;
-  for (const auto& pair : cfusion->output_to_operand_aliasing()) {
-    aliased_outputs.emplace(pair.first, pair.second);
-  }
-  ABSL_RETURN_IF_ERROR(points_to_set.ForEachMutableElementWithStatus(
-      [&](const ShapeIndex& index,
-          PointsToSet::BufferList* buffers) -> absl::Status {
-        auto it = aliased_outputs.find(index);
-        if (it == aliased_outputs.end()) {
-          ABSL_ASSIGN_OR_RETURN(LogicalBuffer * buffer,
-                           logical_buffer_analysis_->GetBuffer(fusion, index));
-          points_to_set.AddPointedToBuffer(*buffer, index);
-        } else {
-          const PointsToSet& input_set =
-              *PerInst(cfusion->operand(it->second.first))->points_to_set;
-          for (const LogicalBuffer* input_buffer :
-               input_set.element(it->second.second)) {
-            points_to_set.AddPointedToBuffer(*input_buffer, index);
-          }
-
-          for (HloInstruction* tuple :
-               input_set.tuple_sources(it->second.second)) {
-            points_to_set.add_tuple_source(index, tuple);
-          }
-        }
-        return absl::OkStatus();
-      }));
-  points_to_set.add_tuple_source({}, fusion);
-  return absl::OkStatus();
-}
-
-absl::Status TuplePointsToAnalysis::HandleOptimizationBarrier(
-    HloInstruction* barrier) {
-  // A kOptimizationBarrier instruction is a no-op.
-  CreateCopiedPointsToSet(barrier, barrier->operand(0));
-  return absl::OkStatus();
-}
-
-const PointsToSet& TuplePointsToAnalysis::GetPointsToSet(
-    const HloInstruction* hlo_instruction) const {
-  return *PerInst(hlo_instruction)->points_to_set;
-}
-
-PointsToSet& TuplePointsToAnalysis::CreateEmptyPointsToSet(
-    const HloInstruction* instruction) {
-  PerInstruction* pi = PerInst(instruction);
-  CHECK(pi->points_to_set == nullptr)
-      << "instruction should not have been present in the map.";
-  auto set = std::make_unique<PointsToSet>(&instruction->shape());
-  pi->points_to_set = std::move(set);
-  // Return *set using the iterator returned by emplace.
-  return *pi->points_to_set;
-}
-
-bool TuplePointsToAnalysis::InstructionDefinesBufferAtIndex(
-    const HloInstruction* instruction, const ShapeIndex& index) const {
-  const auto& buffers = GetPointsToSet(instruction).element(index);
-  return (buffers.size() == 1 && buffers[0]->instruction() == instruction);
-}
-
-absl::Status TuplePointsToAnalysis::VerifyBuffer(
-    const LogicalBuffer& buffer) const {
-  if (!InstructionDefinesBufferAtIndex(buffer.instruction(), buffer.index())) {
-    return FailedPrecondition(
-        "LogicalBuffer %s is ill-defined: instruction %s does not define a "
-        "buffer at that index",
-        buffer.ToString(), buffer.instruction()->name());
-  }
-
-  if (buffer.id() < 0 ||
-      buffer.id() >= logical_buffer_analysis_->num_logical_buffers()) {
-    return FailedPrecondition("LogicalBuffer %s is ill-defined: invalid id %d",
-                              buffer.ToString(), buffer.id());
-  }
-  if (GetBuffer(buffer.id()).instruction() != buffer.instruction() ||
-      GetBuffer(buffer.id()).index() != buffer.index()) {
-    return FailedPrecondition(
-        "LogicalBuffer %s is ill-defined: buffer with same id differs: %s",
-        buffer.ToString(), GetBuffer(buffer.id()).ToString());
-  }
-
-  return absl::OkStatus();
-}
-
-const LogicalBuffer& TuplePointsToAnalysis::GetBuffer(
-    LogicalBuffer::Id id) const {
-  CHECK_GE(id, 0);
-  CHECK_LT(id, logical_buffer_analysis_->num_logical_buffers());
-  absl::StatusOr<LogicalBuffer*> buffer =
-      logical_buffer_analysis_->GetBuffer(id);
-  CHECK_OK(buffer);
-  return **buffer;
-}
-
-absl::StatusOr<const LogicalBuffer*> TuplePointsToAnalysis::GetBufferDefinedAt(
-    const HloInstruction* instruction, const ShapeIndex& index) const {
-  const auto& buffers = GetPointsToSet(instruction).element(index);
-  if (buffers.size() != 1 || buffers[0]->instruction() != instruction) {
-    return FailedPrecondition(
-        "instruction %s does not define buffer at index {%s}",
-        instruction->name(), absl::StrJoin(index, ","));
-  }
-  return buffers[0];
-}
-
-const TuplePointsToAnalysis::BufferAliasVector&
-TuplePointsToAnalysis::GetBufferAliases(const LogicalBuffer& buffer) const {
-  return logical_buffer_aliases_[buffer.id()];
-}
-
-const TuplePointsToAnalysis::BufferDefinitionVector&
-TuplePointsToAnalysis::GetBuffersDefinedByInstruction(
-    const HloInstruction* instruction) const {
-  return PerInst(instruction)->instruction_defined_buffers;
-}
-
-absl::Status TuplePointsToAnalysis::GatherBuffersDefinedByInstruction(
-    const HloInstruction* instruction,
-    TuplePointsToAnalysis::BufferDefinitionVector* buffers) {
-  GetPointsToSet(instruction)
-      .ForEachElement([buffers, instruction](
-                          const ShapeIndex& index,
-                          const PointsToSet::BufferList& source_buffers) {
-        // Add buffers which 'instruction' is the source of.
-        CHECK(!source_buffers.empty());
-        if (source_buffers.size() == 1 &&
-            source_buffers[0]->instruction() == instruction) {
-          // If this instruction is the source of this buffer the
-          // indices must match.
-          DCHECK(source_buffers[0]->index() == index);
-          buffers->push_back(source_buffers[0]);
-        } else {
-          // If the points-to set includes more than one buffer then
-          // necessarily this instruction did not produce the
-          // buffer.
-          for (const LogicalBuffer* source_buffer : source_buffers) {
-            DCHECK(source_buffer->instruction() != instruction);
-          }
-        }
-      });
-  return absl::OkStatus();
-}
-
-PointsToSet& TuplePointsToAnalysis::CreateCopiedPointsToSet(
-    const HloInstruction* instruction, const HloInstruction* src) {
-  // PointsToSet doesn't have a copy constructor so copy over element-by-element
-  // from src PointsToSet.
-  PointsToSet& dst_points_to_set = CreateEmptyPointsToSet(instruction);
-  const PointsToSet& src_points_to_set = GetPointsToSet(src);
-  dst_points_to_set.ForEachMutableElement(
-      [&dst_points_to_set, &src_points_to_set](
-          const ShapeIndex& index, PointsToSet::BufferList* buffers) {
-        *buffers = src_points_to_set.element(index);
-        for (auto& tuple_source : src_points_to_set.tuple_sources(index)) {
-          dst_points_to_set.add_tuple_source(index, tuple_source);
-        }
-      });
-  return *PerInst(instruction)->points_to_set;
-}
-
-std::string TuplePointsToAnalysis::ToString() const {
-  std::string output =
-      absl::StrFormat("TuplePointsToSet for module %s:\n", module_->name());
-  for (const auto* computation : module_->MakeNonfusionComputations()) {
-    const char* entry =
-        computation == module_->entry_computation() ? "entry " : "";
-    absl::StrAppend(&output, entry, "computation ", computation->name(), ":\n");
-    for (const HloInstruction* instruction :
-         computation->MakeInstructionPostOrder()) {
-      InstructionToString(instruction, &output);
-      if (instruction->opcode() == HloOpcode::kFusion) {
-        for (auto* fused : instruction->fused_instructions()) {
-          InstructionToString(fused, &output);
-        }
-      }
-    }
-  }
-
-  absl::StrAppend(&output, "LogicalBuffers:\n");
-  for (const auto& b : logical_buffer_analysis_->logical_buffers()) {
-    absl::StrAppend(&output, "  buffer ", b->ToString(), ":\n");
-    for (const BufferAlias& alias : logical_buffer_aliases_[b->id()]) {
-      absl::StrAppend(&output, "    alias ", alias.ToString(), "\n");
-    }
-  }
-  return output;
-}
-
-void TuplePointsToAnalysis::InstructionToString(
-    const HloInstruction* instruction, std::string* output) const {
-  const std::string prefix = instruction->IsFused() ? "    " : "";
-  absl::StrAppend(output, prefix, "  instruction ",
-                  instruction->ToShortString(), ":\n");
-  const PointsToSet& points_to_set = GetPointsToSet(instruction);
-  points_to_set.ForEachElement(
-      [&prefix, &output](const ShapeIndex& index,
-                         const PointsToSet::BufferList& points_to) {
-        absl::StrAppend(
-            output, prefix, "    {", absl::StrJoin(index, ","), "}: ",
-            absl::StrJoin(points_to, ", ",
-                          [](std::string* out, const LogicalBuffer* source) {
-                            out->append(source->ToString());
-                          }),
-            "\n");
-      });
-}
-
-void TuplePointsToAnalysis::ApplyDeferredAliases(
-    HloInstruction* current_instruction, PointsToSet& points_to_set) {
-  // When an async instruction crosses a while loop boundary, its async-start
-  // may be null or outside scope; safely return if unreachable.
-  const HloAsyncInstruction* async_inst =
-      (current_instruction != nullptr)
-          ? DynCast<HloAsyncInstruction>(current_instruction)
-          : nullptr;
-  if (async_inst == nullptr) {
-    return;
-  }
-  HloInstruction* async_start = async_inst->async_chain_start();
-  const HloAsyncStartInstruction* async_start_inst =
-      (async_start != nullptr) ? DynCast<HloAsyncStartInstruction>(async_start)
-                               : nullptr;
-  if (async_start_inst == nullptr) {
-    return;
-  }
-  for (const auto& pair : async_start_inst->output_to_operand_aliasing()) {
-    const ShapeIndex& logical_output_index = pair.first;
-    int64_t operand_nr = pair.second.first;
-    const ShapeIndex& logical_operand_index = pair.second.second;
-
-    // 1. Construct the source index to the output of the current instruction:
-    //     {0, operand_nr, ...}
-    ShapeIndex chain_source_index;
-    chain_source_index.reserve(2 + logical_operand_index.size());
-    chain_source_index.push_back(0);
-    chain_source_index.push_back(operand_nr);
-    chain_source_index.insert(chain_source_index.end(),
-                              logical_operand_index.begin(),
-                              logical_operand_index.end());
-
-    // 2. Construct the target index in the current instruction.
-    ShapeIndex target_index;
-    if (current_instruction->opcode() == HloOpcode::kAsyncDone) {
-      // For AsyncDone, the target is the logical output without the leading 1.
-      // we should skip cases where the logical output index does not start
-      // with 1, as those are not handled in async done.
-      if (logical_output_index.front() != 1) {
-        // aliasing to chain inputs is not handled in async done.
-        continue;
-      }
-
-      target_index = ShapeIndex(logical_output_index.begin() + 1,
-                                logical_output_index.end());
-    } else {
-      // For AsyncStart/AsyncUpdate, the target is the logical output.
-      target_index = logical_output_index;
-    }
-
-    // 3. Determine where to get the source buffers from, and what shape to
-    // check against.
-    const PointsToSet* source_points_to_set = nullptr;
-    const Shape* source_shape = nullptr;
-    if (current_instruction->opcode() == HloOpcode::kAsyncDone) {
-      // For AsyncDone, we read from the incoming chain (operand 0).
-      const HloInstruction* chain_operand = current_instruction->operand(0);
-      source_points_to_set = &GetPointsToSet(chain_operand);
-      source_shape = &chain_operand->shape();
-    } else {
-      // For AsyncStart/AsyncUpdate, we read "internally" from our own
-      // points_to_set.
-      source_points_to_set = &points_to_set;
-      source_shape = &current_instruction->shape();
-    }
-
-    // 4. If both source and target are valid in their respective shapes, apply
-    // the alias.
-    if (ShapeUtil::IndexIsValid(current_instruction->shape(), target_index) &&
-        ShapeUtil::IndexIsValid(*source_shape, chain_source_index)) {
-      const auto& resolved_buffers =
-          source_points_to_set->element(chain_source_index);
-      *points_to_set.mutable_element(target_index) = resolved_buffers;
-      for (HloInstruction* tuple :
-           source_points_to_set->tuple_sources(chain_source_index)) {
-        points_to_set.add_tuple_source(target_index, tuple);
-      }
-
-      // Propagate the resolved alias back to previous instructions in the
-      // chain, just in case the output is bound earlier than the operand
-      HloInstruction* prev_inst = current_instruction;
-      while (prev_inst->opcode() == HloOpcode::kAsyncUpdate ||
-             prev_inst->opcode() == HloOpcode::kAsyncDone) {
-        prev_inst = prev_inst->mutable_operand(0);
-        if (ShapeUtil::IndexIsValid(prev_inst->shape(), logical_output_index)) {
-          PointsToSet& prev_points_to_set = *PerInst(prev_inst)->points_to_set;
-          *prev_points_to_set.mutable_element(logical_output_index) =
-              resolved_buffers;
-          for (HloInstruction* tuple :
-               source_points_to_set->tuple_sources(chain_source_index)) {
-            prev_points_to_set.add_tuple_source(logical_output_index, tuple);
-          }
-        }
-      }
-    }
-  }
-}
-
-std::string PointsToSet::ToString() const {
-  std::string output;
-  ForEachElement([&output](const ShapeIndex& index,
-                           const BufferList& points_to) {
-    absl::StrAppend(&output, "{", absl::StrJoin(index, ","), "}: ",
-                    absl::StrJoin(points_to, ", ",
-                                  [](std::string* out, const LogicalBuffer* b) {
-                                    out->append(b->ToString());
-                                  }),
-                    "\n");
-  });
-  return output;
-}
-
-}  // namespace xla
diff --git a/third_party/xla/xla/hlo/analysis/tuple_points_to_analysis.h b/third_party/xla/xla/hlo/analysis/tuple_points_to_analysis.h
deleted file mode 100644
index abf46313f08..00000000000
--- a/third_party/xla/xla/hlo/analysis/tuple_points_to_analysis.h
+++ /dev/null
@@ -1,435 +0,0 @@
-/* Copyright 2017 The OpenXLA Authors.
-
-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.
-==============================================================================*/
-
-#ifndef XLA_HLO_ANALYSIS_TUPLE_POINTS_TO_ANALYSIS_H_
-#define XLA_HLO_ANALYSIS_TUPLE_POINTS_TO_ANALYSIS_H_
-
-#include <stddef.h>
-
-#include <cstdint>
-#include <iosfwd>
-#include <memory>
-#include <string>
-#include <utility>
-#include <vector>
-
-#include "absl/container/flat_hash_map.h"
-#include "absl/container/inlined_vector.h"
-#include "absl/log/check.h"
-#include "absl/log/log.h"
-#include "absl/status/status.h"
-#include "absl/status/statusor.h"
-#include "xla/hlo/analysis/logical_buffer_analysis.h"
-#include "xla/hlo/ir/dfs_hlo_visitor_with_default.h"
-#include "xla/hlo/ir/hlo_instruction.h"
-#include "xla/hlo/ir/hlo_module.h"
-#include "xla/service/logical_buffer.h"
-#include "xla/shape.h"
-#include "xla/shape_tree.h"
-#include "xla/shape_util.h"
-#include "xla/tsl/lib/gtl/compactptrset.h"
-#include "xla/xla_data.pb.h"
-
-namespace xla {
-
-// A PointsToSet describes the source(s) of the LogicalBuffer(s) contained in
-// the output of a particular HLO instruction.
-//
-// In XLA, HLO instructions do not necessarily allocate new memory; some may
-// forward or alias existing buffers (e.g., GetTupleElement, Bitcast). To track
-// who defines the buffers, the compiler performs points-to analysis.
-//
-// Since an instruction's output can have a nested or structured shape (like a
-// nested tuple), PointsToSet internally uses a ShapeTree. The structure of this
-// tree mirrors the Shape of the instruction.
-//
-// Each node in the tree (corresponding to a subshape index/path in the output
-// Shape) contains:
-// 1) A list of LogicalBuffers (`BufferList`) that can potentially define the
-//    buffer at that subshape index.
-// 2) A set of HLO instructions (`SourceSet`) representing the instructions that
-//    define/originate the tuple shape at that index (i.e. 'tuple sources').
-//
-// Thread Safety: PointsToSet is not thread-safe due to internal mutation of
-// the ShapeTree without locking.
-class PointsToSet {
- public:
-  // Constructs a PointsToSet for the given shape.
-  //
-  // Note: Accepts a `const Shape*` instead of a copy or reference. This is an
-  // optimization in hot execution paths to avoid copying/creating many
-  // temporary `Shape` objects.
-  explicit PointsToSet(const Shape* shape) : tree_(shape) {}
-
-  // Returns true if any of the points-to sets (at any subshape index) contains
-  // more than one LogicalBuffer.
-  //
-  // An ambiguous set indicates that the compiler cannot statically determine
-  // a single unique source buffer for that part of the output.
-  bool IsAmbiguous() const;
-
-  // Returns true if no LogicalBuffer appears at more than one subshape index
-  // of the instruction's shape.
-  //
-  // A distinct points-to set guarantees that different components/paths of the
-  // output (e.g., elements in a tuple) do not share or alias the same logical
-  // buffers.
-  bool IsDistinct() const;
-
-  // Returns the total number of unique LogicalBuffers across all subshape
-  // indices in this PointsToSet. This is equivalent to
-  // `CreateFlattenedSet().size()`.
-  size_t size() const;
-
-  // Returns the union of all LogicalBuffers contained in all subshape indices
-  // of this PointsToSet.
-  using BufferSet = tsl::gtl::CompactPointerSet<const LogicalBuffer*>;
-  BufferSet CreateFlattenedSet() const;
-
-  // Returns true if the points-to set at the given `ShapeIndex` contains the
-  // specified `LogicalBuffer`.
-  bool ContainsBufferAtIndex(const LogicalBuffer& buffer,
-                             const ShapeIndex& index) const;
-
-  // Returns true if the specified `LogicalBuffer` is present in the points-to
-  // set at any `ShapeIndex` in this PointsToSet.
-  bool ContainsBuffer(const LogicalBuffer& buffer) const;
-
-  // Adds the given `LogicalBuffer` to the points-to set at the specified
-  // `ShapeIndex`. This is a no-op if the buffer is already present at that
-  // index.
-  void AddPointedToBuffer(const LogicalBuffer& buffer, const ShapeIndex& index);
-
-  using SourceSet = tsl::gtl::CompactPointerSet<HloInstruction*>;
-
-  // For the subshape at the given index (which must be a tuple shape), returns
-  // the set of HLO instructions that may define/originate the tuple structure
-  // at that index.
-  //
-  // This is used to track the allocation/creation points of tuple shapes,
-  // which helps analyze how tuple buffers are shared.
-  //
-  // Example:
-  // Given:
-  //   %tuple1 = tuple(...)
-  //   %tuple2 = tuple(...)
-  //   %cond = conditional(...) // selects between %tuple1 and %tuple2
-  //   %nested_tuple = tuple(%cond, %tuple1)
-  //
-  // In `nested_tuple`'s PointsToSet:
-  //   tuple_sources({})   -> {%nested_tuple}
-  //   tuple_sources({0})  -> {%cond}
-  //   tuple_sources({1})  -> {%tuple1}
-  //
-  // If the subshape at the index is an array (not a tuple), this returns an
-  // empty set. The instructions in the set are typically Tuple instructions,
-  // constants, parameters, or control flow instructions (e.g., While,
-  // Conditional).
-  //
-  // Note: One cannot assume that `tuple_sources` is a subset of the defining
-  // instructions of the `buffers` at the same index. For example, a shallow
-  // copy of a tuple `%copy = copy(%original_tuple)` allocates a new top-level
-  // buffer defined by `%copy` at index `{}`, which updates `buffers({})` to
-  // contain `LogicalBuffer(%copy, {})`. However, the tuple structure origin is
-  // unchanged, so `tuple_sources({})` remains `{%original_tuple}`.
-  const SourceSet& tuple_sources(const ShapeIndex& index) const;
-
-  // Adds the given HLO instruction as a tuple source for the subshape index.
-  void add_tuple_source(const ShapeIndex& index, HloInstruction* tuple);
-
-  using BufferList = absl::InlinedVector<const LogicalBuffer*, 1>;
-
-  // Returns the `BufferList` (list of `LogicalBuffer`s) for the subshape at the
-  // given index.
-  const BufferList& element(const ShapeIndex& index) const {
-    return tree_.element(index).buffers;
-  }
-  BufferList* mutable_element(const ShapeIndex& index) {
-    return &tree_.mutable_element(index)->buffers;
-  }
-
-  // Invokes `fn` (a callable with signature compatible with
-  // `void(const ShapeIndex&, const BufferList&)`) on each element in the
-  // points-to set.
-  template <typename Fn>
-  void ForEachElement(const Fn& fn) const {
-    tree_.ForEachElement([&fn](const ShapeIndex& index, const Elem& elem) {
-      fn(index, elem.buffers);
-    });
-  }
-  template <typename Fn>
-  void ForEachMutableElement(const Fn& fn) {
-    tree_.ForEachMutableElement([&fn](const ShapeIndex& index, Elem* elem) {
-      fn(index, &elem->buffers);
-    });
-  }
-  template <typename Fn>
-  absl::Status ForEachMutableElementWithStatus(const Fn& fn) {
-    return tree_.ForEachMutableElementWithStatus(
-        [&fn](const ShapeIndex& index, Elem* elem) {
-          return fn(index, &elem->buffers);
-        });
-  }
-  template <typename Fn>
-  absl::Status ForEachElementWithStatus(const Fn& fn) const {
-    return tree_.ForEachElementWithStatus(
-        [&fn](const ShapeIndex& index, const Elem& elem) {
-          return fn(index, elem.buffers);
-        });
-  }
-
-  // Returns a string representation of the PointsToSet for debugging.
-  std::string ToString() const;
-
- private:
-  struct Elem {
-    // The list of LogicalBuffers that may potentially define/provide the value
-    // at this subshape.
-    BufferList buffers;
-    // The set of HLO instructions that define/originate the tuple shape at
-    // this subshape (populated only if this subshape is tuple-shaped). Note
-    // that this set is not necessarily a subset of the instructions defining
-    // the buffers in `buffers` (e.g., in the case of a shallow copy).
-    SourceSet tuple_sources;
-  };
-  ShapeTree<Elem> tree_;
-
-  // PointsToSet contains references (const LogicalBuffer*) to elements within
-  // TuplePointsToAnalysis, so disable copying.
-  PointsToSet(const PointsToSet&) = delete;
-  PointsToSet& operator=(const PointsToSet&) = delete;
-};
-
-// This class describes a particular subshape in a computation (instruction and
-// shape index) and the logical buffer which may be a source of the subshape
-// value.
-class BufferAlias {
- public:
-  BufferAlias(HloInstruction* instruction, const ShapeIndex& index)
-      : instruction_(instruction), index_(index) {}
-
-  // Return the instruction/index of the subshape.
-  HloInstruction* instruction() const { return instruction_; }
-  const ShapeIndex& index() const { return index_; }
-
-  bool operator==(const BufferAlias& other) const {
-    return instruction_ == other.instruction_ && index_ == other.index_;
-  }
-  bool operator!=(const BufferAlias& other) const { return !(*this == other); }
-
-  std::string ToString() const;
-
- private:
-  HloInstruction* instruction_;
-  ShapeIndex index_;
-};
-
-std::ostream& operator<<(std::ostream& out, const BufferAlias& buffer_alias);
-
-// DFS visitor that performs tuple points-to analysis. This analysis determines
-// the potential sources of each buffer in each instruction's output.
-class TuplePointsToAnalysis : public DfsHloVisitorWithDefault {
- public:
-  // Runs points-to analysis on 'module'.
-  static absl::StatusOr<std::unique_ptr<TuplePointsToAnalysis>> Run(
-      const HloModule* module);
-
-  // Return the points-to set of an instruction. This describes the potential
-  // sources of each buffer in the instruction's output.
-  const PointsToSet& GetPointsToSet(
-      const HloInstruction* hlo_instruction) const;
-
-  // Returns the logical buffer with the given ID.
-  const LogicalBuffer& GetBuffer(LogicalBuffer::Id id) const;
-
-  // Returns the buffer defined at the given instruction and index. An error is
-  // returned if no buffer is defined at that point.
-  absl::StatusOr<const LogicalBuffer*> GetBufferDefinedAt(
-      const HloInstruction* instruction, const ShapeIndex& index) const;
-
-  // Return a (possibly empty) vector containing all BufferAliases of the given
-  // logical buffer The buffer alias set is the inverse of the points-to set.
-  // That is, LogicalBuffer B is in the points-to set of instruction I at index
-  // N iff instruction I, index N is a BufferAlias of B.
-  using BufferAliasVector = absl::InlinedVector<BufferAlias, 1>;
-  const BufferAliasVector& GetBufferAliases(const LogicalBuffer& buffer) const;
-
-  // Returns the number of logical buffers in the module
-  LogicalBuffer::Id num_logical_buffers() const {
-    return logical_buffer_analysis_->num_logical_buffers();
-  }
-
-  // Returns a vector of buffers that the instruction produces. Most
-  // instructions produce a single buffer (the top-level buffer), some produce
-  // no buffers (eg bitcast), and some produce more than one buffer (eg,
-  // tuple-shaped parameters).
-  using BufferDefinitionVector = absl::InlinedVector<const LogicalBuffer*, 1>;
-  const BufferDefinitionVector& GetBuffersDefinedByInstruction(
-      const HloInstruction* instruction) const;
-
-  // Returns true if the given instruction defines a buffer at the given index.
-  bool InstructionDefinesBufferAtIndex(const HloInstruction* instruction,
-                                       const ShapeIndex& index) const;
-
-  // Returns an OK status if the given buffer is defined by instruction
-  // 'buffer.instruction()' at index 'buffer.index()' and if the given buffer
-  // matches the TuplePointsToAnalysis' LogicalBuffer with 'buffer.id'. Returns
-  // an FailedPrecondition error status otherwise. An example of a LogicalBuffer
-  // which is not defined is a tuple element in a Tuple instruction. In this
-  // case, the Tuple instruction does not define the LogicalBuffer, rather that
-  // index aliases one of its operands.
-  absl::Status VerifyBuffer(const LogicalBuffer& buffer) const;
-
-  absl::Status DefaultAction(HloInstruction* hlo_instruction) override;
-  absl::Status HandleTuple(HloInstruction* tuple) override;
-  absl::Status HandleGetTupleElement(
-      HloInstruction* get_tuple_element) override;
-  absl::Status HandleAsyncStart(HloInstruction* async_start) override;
-  absl::Status HandleAsyncUpdate(HloInstruction* async_update) override;
-  absl::Status HandleAsyncDone(HloInstruction* async_done) override;
-  absl::Status HandleBitcast(HloInstruction* bitcast) override;
-  absl::Status HandleDomain(HloInstruction* domain) override;
-  absl::Status HandleCopy(HloInstruction* copy) override;
-  absl::Status HandleCopyStart(HloInstruction* copy_start) override;
-  absl::Status HandleCopyDone(HloInstruction* copy_done) override;
-  absl::Status HandleRecvDone(HloInstruction* recv_done) override;
-  absl::Status HandleSend(HloInstruction* send) override;
-  absl::Status HandleAddDependency(HloInstruction* add_dependency) override;
-  absl::Status HandleCustomCall(HloInstruction* custom_call) override;
-  absl::Status HandleFusion(HloInstruction* fusion) override;
-  absl::Status HandleOptimizationBarrier(HloInstruction* barrier) override;
-
-  std::string ToString() const;
-
- private:
-  explicit TuplePointsToAnalysis(
-      const HloModule* module,
-      std::unique_ptr<LogicalBufferAnalysis> logical_buffer_analysis)
-      : module_(module),
-        logical_buffer_analysis_(std::move(logical_buffer_analysis)) {}
-
-  // Perform the analysis. Should be called immediately after constructing the
-  // object and before calling GetPointsToSet.
-  absl::Status Analyze();
-
-  // Populates instruction-defined buffers and aliases for each instruction
-  // in 'instructions'.
-  absl::Status PopulateDefinedBuffersAndAliases(
-      const decltype(std::declval<HloComputation>()
-                         .instructions())& instructions);
-
-  // Creates an empty PointsToSet in the points_to_ map for the given
-  // instruction.
-  PointsToSet& CreateEmptyPointsToSet(const HloInstruction* instruction);
-
-  // Creates a PointsToSet in the points_to_ map for 'instruction' which is a
-  // copy of the existing PointsToSet for 'src'.
-  PointsToSet& CreateCopiedPointsToSet(const HloInstruction* instruction,
-                                       const HloInstruction* src);
-
-  // Adds the buffers defined by the given instruction to the given vector.
-  absl::Status GatherBuffersDefinedByInstruction(
-      const HloInstruction* instruction, BufferDefinitionVector* buffers);
-
-  // Applies deferred aliases from async-start to the current instruction's
-  // points-to set.
-  void ApplyDeferredAliases(HloInstruction* current_instruction,
-                            PointsToSet& points_to_set);
-
-  // Propagates points-to sets for an instruction that aggregates its operands
-  // into an output tuple structure.
-  absl::Status ConstructPointsToSetByAggregatingOperands(
-      HloInstruction* instruction);
-
-  // Print points-to set for 'instruction' to 'output'.
-  void InstructionToString(const HloInstruction* instruction,
-                           std::string* output) const;
-
-  // Information kept per instruction
-  struct PerInstruction {
-    std::unique_ptr<PointsToSet> points_to_set;
-    // Buffers defined by this instruction.
-    // Empirically, ~92% of instructions have 1
-    // instruction_defined_buffer, and 99% have 0 or 1
-    BufferDefinitionVector instruction_defined_buffers;
-  };
-
-  const PerInstruction* PerInst(const HloInstruction* inst) const {
-    const int64_t computation_id = inst->parent()->unique_id();
-    const int32_t local_id = inst->local_id();
-    DCHECK_GE(computation_id, 0);
-    DCHECK_GE(local_id, 0);
-    if (computation_id >= per_instruction_.size() ||
-        local_id >= per_instruction_[computation_id].size() ||
-        per_instruction_[computation_id][local_id] == nullptr) {
-      LOG(FATAL) << "Expected per-instruction information to already exist";
-    }
-    return per_instruction_[computation_id][local_id].get();
-  }
-  PerInstruction* PerInst(const HloInstruction* inst) {
-    const int64_t computation_id = inst->parent()->unique_id();
-    const int32_t local_id = inst->local_id();
-    DCHECK_GE(computation_id, 0);
-    DCHECK_GE(local_id, 0);
-    if (computation_id >= per_instruction_.size()) {
-      per_instruction_.resize(computation_id + 1);
-    }
-    auto& slots = per_instruction_[computation_id];
-    if (local_id >= slots.size()) {
-      slots.resize(local_id + 1);
-    }
-    auto& slot = slots[local_id];
-    if (slot == nullptr) {
-      slot = std::make_unique<PerInstruction>();
-    }
-    return slot.get();
-  }
-
-  // The module this analysis is performed on.
-  const HloModule* module_;
-
-  // The logical buffers for this module.
-  const std::unique_ptr<LogicalBufferAnalysis> logical_buffer_analysis_;
-
-  // Per-instruction information, stored densely: the outer vector is indexed
-  // by the owning computation's unique_id() and the inner vector by the
-  // instruction's computation-local id. This replaces a large hash map keyed
-  // by instruction unique_id (which was a top find() hotspot on big modules)
-  // with two array indexes.
-  //
-  // Both id spaces are dense when the analysis runs: local ids are
-  // recompacted by HloComputation::Cleanup(), which HloPassPipeline runs
-  // after every pass, and analyses are built at pass start (the same
-  // assumption the post-order DFS VisitMap makes). Instructions deleted
-  // while the analysis is alive leave 8-byte holes that are never queried,
-  // matching the previous hash-map semantics where stale entries lingered
-  // until the analysis was rebuilt.
-  std::vector<std::vector<std::unique_ptr<PerInstruction>>> per_instruction_;
-
-  // A map from LogicalBuffer->id() to alias information about that logical
-  // buffer
-  std::vector<BufferAliasVector> logical_buffer_aliases_;
-
-  TuplePointsToAnalysis(const TuplePointsToAnalysis&) = delete;
-  TuplePointsToAnalysis& operator=(const TuplePointsToAnalysis&) = delete;
-  // Whether to alias buffers connected by dataflow relations. This aliasing
-  // relation should not be recognized if copies can be inserted to break up
-  // the dataflow relation.
-  const bool alias_buffer_across_dataflow_ = false;
-};
-
-}  // namespace xla
-
-#endif  // XLA_HLO_ANALYSIS_TUPLE_POINTS_TO_ANALYSIS_H_
diff --git a/third_party/xla/xla/hlo/analysis/tuple_points_to_analysis_test.cc b/third_party/xla/xla/hlo/analysis/tuple_points_to_analysis_test.cc
deleted file mode 100644
index b3391e3242a..00000000000
--- a/third_party/xla/xla/hlo/analysis/tuple_points_to_analysis_test.cc
+++ /dev/null
@@ -1,1328 +0,0 @@
-/* Copyright 2017 The OpenXLA Authors.
-
-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.
-==============================================================================*/
-
-#include "xla/hlo/analysis/tuple_points_to_analysis.h"
-
-#include <cstdint>
-#include <memory>
-#include <string>
-#include <utility>
-#include <vector>
-
-#include "absl/algorithm/container.h"
-#include "absl/log/check.h"
-#include "absl/strings/string_view.h"
-#include "absl/types/span.h"
-#include "xla/hlo/ir/hlo_computation.h"
-#include "xla/hlo/ir/hlo_instruction.h"
-#include "xla/hlo/ir/hlo_opcode.h"
-#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
-#include "xla/hlo/testlib/test.h"
-#include "xla/service/logical_buffer.h"
-#include "xla/shape_util.h"
-#include "xla/tsl/platform/logging.h"
-
-namespace xla {
-namespace {
-
-using ::testing::UnorderedElementsAre;
-using ::testing::UnorderedElementsAreArray;
-
-class TuplePointsToAnalysisTest : public HloHardwareIndependentTestBase {
- protected:
-  // Builds a module with the given entry computation and runs points to
-  // analysis.
-  void BuildModuleAndRunAnalysis(std::unique_ptr<HloComputation> computation) {
-    BuildModule(std::move(computation));
-    RunAnalysis();
-  }
-
-  void BuildModule(std::unique_ptr<HloComputation> computation) {
-    module_ = CreateNewVerifiedModule();
-    module_->AddEntryComputation(std::move(computation));
-  }
-
-  void RunAnalysis() {
-    CHECK_NOTNULL(module_.get());
-    points_to_analysis_ = TuplePointsToAnalysis::Run(module_.get()).value();
-  }
-
-  // Returns the LogicalBuffer defined at the given instruction and
-  // index. CHECKs if no buffer is defined at that point.
-  const LogicalBuffer* GetBuffer(const HloInstruction* instruction,
-                                 const ShapeIndex& index) {
-    const auto& pointed_to =
-        points_to_analysis_->GetPointsToSet(instruction).element(index);
-    CHECK_EQ(1, pointed_to.size());
-    CHECK_EQ(instruction, pointed_to[0]->instruction());
-    CHECK(index == pointed_to[0]->index());
-    return pointed_to[0];
-  }
-
-  // Checks that the given points-to set contains exactly (unordered) the given
-  // LogicalBuffers.
-  void ExpectHasBuffers(const PointsToSet::BufferList& points_to_set,
-                        absl::Span<const LogicalBuffer* const> buffers) {
-    std::vector<const LogicalBuffer*> vec(buffers.begin(), buffers.end());
-    EXPECT_THAT(points_to_set, UnorderedElementsAreArray(vec));
-  }
-
-  // Checks that the given points-to set contains exactly (unordered) the
-  // top-level buffers of the given instructions.
-  void ExpectHasTopLevelBuffers(
-      const PointsToSet::BufferList& points_to_set,
-      absl::Span<HloInstruction* const> instructions) {
-    PointsToSet::BufferList buffers;
-    for (auto instruction : instructions) {
-      buffers.push_back(GetBuffer(instruction, /*index=*/{}));
-    }
-    ExpectHasBuffers(points_to_set, buffers);
-  }
-
-  // Overload which takes a set instead of a vector.
-  void ExpectHasTopLevelBuffers(
-      const PointsToSet::BufferSet& points_to_set,
-      absl::Span<HloInstruction* const> instructions) {
-    ExpectHasTopLevelBuffers(
-        PointsToSet::BufferList(points_to_set.begin(), points_to_set.end()),
-        instructions);
-  }
-
-  // Checks that the buffer defined at the given instruction and index has
-  // aliases which are exactly (unordered) the given instruction/index pairs.
-  void ExpectHasBufferAliases(
-      const HloInstruction* instruction, const ShapeIndex& index,
-      absl::Span<const std::pair<HloInstruction*, ShapeIndex>> expected) {
-    const LogicalBuffer* buffer =
-        points_to_analysis_->GetBufferDefinedAt(instruction, index).value();
-    std::vector<BufferAlias> expected_aliases;
-    expected_aliases.reserve(expected.size());
-    for (auto& pair : expected) {
-      expected_aliases.push_back(BufferAlias(pair.first, pair.second));
-    }
-    EXPECT_THAT(points_to_analysis_->GetBufferAliases(*buffer),
-                UnorderedElementsAreArray(expected_aliases));
-  }
-
-  std::unique_ptr<HloModule> module_;
-  std::unique_ptr<TuplePointsToAnalysis> points_to_analysis_;
-};
-
-TEST_F(TuplePointsToAnalysisTest, SimpleTuple) {
-  std::string hlo_str = R"(
-HloModule SimpleTuple
-
-ENTRY main {
-  constant1 = f32[] constant(1)
-  constant2 = f32[] constant(2)
-  ROOT tuple = (f32[], f32[]) tuple(constant1, constant2)
-}
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(hlo_str, GetModuleConfigForTest()));
-  RunAnalysis();
-
-  auto* constant1 = FindInstruction(module_.get(), "constant1");
-  auto* constant2 = FindInstruction(module_.get(), "constant2");
-  auto* tuple = FindInstruction(module_.get(), "tuple");
-
-  EXPECT_EQ(1, points_to_analysis_->GetPointsToSet(constant1).size());
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(constant1).element({}), {constant1});
-  EXPECT_TRUE(
-      points_to_analysis_->GetPointsToSet(constant1).tuple_sources({}).empty());
-  EXPECT_TRUE(points_to_analysis_->GetPointsToSet(tuple).IsDistinct());
-
-  EXPECT_EQ(1, points_to_analysis_->GetPointsToSet(constant2).size());
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(constant2).element({}), {constant2});
-  EXPECT_TRUE(
-      points_to_analysis_->GetPointsToSet(constant2).tuple_sources({}).empty());
-
-  EXPECT_EQ(3, points_to_analysis_->GetPointsToSet(tuple).size());
-  EXPECT_FALSE(points_to_analysis_->GetPointsToSet(tuple).IsAmbiguous());
-  EXPECT_THAT(points_to_analysis_->GetPointsToSet(tuple).tuple_sources({}),
-              UnorderedElementsAre(tuple));
-
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(tuple).CreateFlattenedSet(),
-      {constant1, constant2, tuple});
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(tuple).element({}), {tuple});
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(tuple).element({0}), {constant1});
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(tuple).element({1}), {constant2});
-
-  const PointsToSet& tuple_points_to_set =
-      points_to_analysis_->GetPointsToSet(tuple);
-  EXPECT_TRUE(tuple_points_to_set.ContainsBufferAtIndex(
-      *GetBuffer(constant1, {}), {0}));
-  EXPECT_TRUE(tuple_points_to_set.ContainsBufferAtIndex(
-      *GetBuffer(constant2, {}), {1}));
-  EXPECT_FALSE(tuple_points_to_set.ContainsBufferAtIndex(
-      *GetBuffer(constant2, {}), {0}));
-  EXPECT_TRUE(tuple_points_to_set.ContainsBuffer(*GetBuffer(constant1, {})));
-  EXPECT_TRUE(tuple_points_to_set.ContainsBuffer(*GetBuffer(constant2, {})));
-}
-
-TEST_F(TuplePointsToAnalysisTest, NestedTuple) {
-  // Create a (nested) tuple containing an inner tuple. The points-to set of the
-  // outer tuple should contain all elements of the points-to set of the inner
-  // tuple.
-  std::string hlo_str = R"(
-HloModule NestedTuple
-
-ENTRY main {
-  constant1 = f32[] constant(1)
-  constant2 = f32[] constant(2)
-  inner_tuple = (f32[], f32[]) tuple(constant1, constant2)
-  constant3 = f32[] constant(3)
-  ROOT tuple = ((f32[], f32[]), f32[]) tuple(inner_tuple, constant3)
-}
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(hlo_str, GetModuleConfigForTest()));
-  RunAnalysis();
-
-  auto* constant1 = FindInstruction(module_.get(), "constant1");
-  auto* constant2 = FindInstruction(module_.get(), "constant2");
-  auto* constant3 = FindInstruction(module_.get(), "constant3");
-  auto* inner_tuple = FindInstruction(module_.get(), "inner_tuple");
-  auto* tuple = FindInstruction(module_.get(), "tuple");
-
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(constant1).element({}), {constant1});
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(constant2).element({}), {constant2});
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(constant3).element({}), {constant3});
-
-  EXPECT_EQ(3, points_to_analysis_->GetPointsToSet(inner_tuple).size());
-  EXPECT_FALSE(points_to_analysis_->GetPointsToSet(inner_tuple).IsAmbiguous());
-  EXPECT_TRUE(points_to_analysis_->GetPointsToSet(inner_tuple).IsDistinct());
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(inner_tuple).CreateFlattenedSet(),
-      {constant1, constant2, inner_tuple});
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(inner_tuple).element({}),
-      {inner_tuple});
-  EXPECT_THAT(
-      points_to_analysis_->GetPointsToSet(inner_tuple).tuple_sources({}),
-      UnorderedElementsAre(inner_tuple));
-
-  EXPECT_EQ(5, points_to_analysis_->GetPointsToSet(tuple).size());
-  EXPECT_FALSE(points_to_analysis_->GetPointsToSet(tuple).IsAmbiguous());
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(tuple).CreateFlattenedSet(),
-      {constant1, constant2, constant3, inner_tuple, tuple});
-
-  EXPECT_THAT(points_to_analysis_->GetPointsToSet(tuple).tuple_sources({}),
-              UnorderedElementsAre(tuple));
-  EXPECT_THAT(points_to_analysis_->GetPointsToSet(tuple).tuple_sources({0}),
-              UnorderedElementsAre(inner_tuple));
-  EXPECT_TRUE(
-      points_to_analysis_->GetPointsToSet(tuple).tuple_sources({1}).empty());
-
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(tuple).element({0}), {inner_tuple});
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(tuple).element({0, 0}), {constant1});
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(tuple).element({0, 1}), {constant2});
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(tuple).element({1}), {constant3});
-}
-
-TEST_F(TuplePointsToAnalysisTest, GetTupleElement) {
-  // Create a nested tuple, then extract the inner tuple with GetTupleElement.
-  // The points-to set of the GetTupleElement should be the same as the inner
-  // tuple.
-  std::string hlo_str = R"(
-HloModule GetTupleElement
-
-ENTRY main {
-  constant1 = f32[] constant(1)
-  constant2 = f32[] constant(2)
-  inner_tuple = (f32[], f32[]) tuple(constant1, constant2)
-  constant3 = f32[] constant(3)
-  tuple = ((f32[], f32[]), f32[]) tuple(inner_tuple, constant3)
-  ROOT get_tuple_element = (f32[], f32[]) get-tuple-element(tuple), index=0
-}
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(hlo_str, GetModuleConfigForTest()));
-  RunAnalysis();
-
-  auto* constant1 = FindInstruction(module_.get(), "constant1");
-  auto* constant2 = FindInstruction(module_.get(), "constant2");
-  auto* inner_tuple = FindInstruction(module_.get(), "inner_tuple");
-  auto* get_tuple_element = FindInstruction(module_.get(), "get_tuple_element");
-
-  auto& points_to_set = points_to_analysis_->GetPointsToSet(get_tuple_element);
-  EXPECT_EQ(3, points_to_set.size());
-  EXPECT_FALSE(points_to_set.IsAmbiguous());
-  EXPECT_TRUE(points_to_set.IsDistinct());
-  ExpectHasTopLevelBuffers(points_to_set.CreateFlattenedSet(),
-                           {constant1, constant2, inner_tuple});
-  ExpectHasTopLevelBuffers(points_to_set.element({}), {inner_tuple});
-
-  EXPECT_THAT(points_to_set.tuple_sources({}),
-              UnorderedElementsAre(inner_tuple));
-}
-
-TEST_F(TuplePointsToAnalysisTest, AddDependency) {
-  std::string hlo_str = R"(
-HloModule AddDependency
-
-ENTRY main {
-  constant = f32[] constant(1)
-  token0 = token[] after-all()
-  ROOT add_dependency = f32[] add-dependency(constant, token0)
-}
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(hlo_str, GetModuleConfigForTest()));
-  RunAnalysis();
-
-  auto* constant = FindInstruction(module_.get(), "constant");
-  auto* add_dependency = FindInstruction(module_.get(), "add_dependency");
-
-  auto& points_to_set = points_to_analysis_->GetPointsToSet(add_dependency);
-  EXPECT_EQ(1, points_to_set.size());
-  EXPECT_FALSE(points_to_set.IsAmbiguous());
-  EXPECT_TRUE(points_to_set.IsDistinct());
-  ExpectHasTopLevelBuffers(points_to_set.CreateFlattenedSet(), {constant});
-}
-
-TEST_F(TuplePointsToAnalysisTest, DuplicatedElement) {
-  // Create a tuple which contains duplicate elements.
-  std::string hlo_str = R"(
-HloModule DuplicatedElement
-
-ENTRY main {
-  constant = f32[] constant(1)
-  ROOT tuple = (f32[], f32[], f32[]) tuple(constant, constant, constant)
-}
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(hlo_str, GetModuleConfigForTest()));
-  RunAnalysis();
-
-  auto* constant = FindInstruction(module_.get(), "constant");
-  auto* tuple = FindInstruction(module_.get(), "tuple");
-
-  EXPECT_EQ(2, points_to_analysis_->GetPointsToSet(tuple).size());
-  EXPECT_FALSE(points_to_analysis_->GetPointsToSet(tuple).IsAmbiguous());
-  EXPECT_FALSE(points_to_analysis_->GetPointsToSet(tuple).IsDistinct());
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(tuple).element({}), {tuple});
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(tuple).CreateFlattenedSet(),
-      {constant, tuple});
-}
-
-TEST_F(TuplePointsToAnalysisTest, TupleCopy) {
-  // Create a copy (HloOpcode::kCopy) of a tuple. The points to sets should be
-  // the same.
-  std::string hlo_str = R"(
-HloModule TupleCopy
-
-ENTRY main {
-  constant1 = f32[] constant(1)
-  constant2 = f32[] constant(2)
-  tuple = (f32[], f32[]) tuple(constant1, constant2)
-  ROOT copy = (f32[], f32[]) copy(tuple)
-}
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(hlo_str, GetModuleConfigForTest()));
-  RunAnalysis();
-
-  auto* constant1 = FindInstruction(module_.get(), "constant1");
-  auto* constant2 = FindInstruction(module_.get(), "constant2");
-  auto* tuple = FindInstruction(module_.get(), "tuple");
-  auto* copy = FindInstruction(module_.get(), "copy");
-
-  EXPECT_FALSE(points_to_analysis_->GetPointsToSet(copy).IsAmbiguous());
-  EXPECT_TRUE(points_to_analysis_->GetPointsToSet(copy).IsDistinct());
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(tuple).CreateFlattenedSet(),
-      {constant1, constant2, tuple});
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(copy).element({}), {copy});
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(copy).CreateFlattenedSet(),
-      {constant1, constant2, copy});
-}
-
-TEST_F(TuplePointsToAnalysisTest, CopyStartAndCopyDone) {
-  // CopyDone forwards its operand to the output tuple at {0}.
-  std::string hlo_str = R"(
-HloModule CopyStartAndCopyDone
-
-ENTRY main {
-  constant = f32[] constant(1)
-  copy_start = (f32[], f32[], u32[]) copy-start(constant)
-  ROOT copy_done = f32[] copy-done(copy_start)
-}
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(hlo_str, GetModuleConfigForTest()));
-  RunAnalysis();
-
-  auto* constant = FindInstruction(module_.get(), "constant");
-  auto* copy_start = FindInstruction(module_.get(), "copy_start");
-  auto* copy_done = FindInstruction(module_.get(), "copy_done");
-
-  EXPECT_FALSE(points_to_analysis_->GetPointsToSet(copy_start).IsAmbiguous());
-  EXPECT_TRUE(points_to_analysis_->GetPointsToSet(copy_start).IsDistinct());
-  EXPECT_FALSE(points_to_analysis_->GetPointsToSet(copy_done).IsAmbiguous());
-  EXPECT_TRUE(points_to_analysis_->GetPointsToSet(copy_done).IsDistinct());
-
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(copy_start).element({}),
-      {copy_start});
-  ExpectHasBufferAliases(copy_start, {0}, {{copy_start, {0}}, {copy_done, {}}});
-  ExpectHasBufferAliases(constant, {}, {{constant, {}}, {copy_start, {1}}});
-}
-
-TEST_F(TuplePointsToAnalysisTest, AsyncOps) {
-  std::string hlo_str = R"(
-  HloModule module
-
-  ENTRY entry {
-    p0 = f32[2,3] parameter(0)
-    async-start = ((f32[2,3]), f32[2,3], u32[]) custom-call-start(p0), custom_call_target="foo"
-    async-update = ((f32[2,3]), f32[2,3], u32[]) custom-call-update(async-start),
-                      output_to_operand_aliasing={{2}: (0, {2})}
-    ROOT async-done = f32[2,3] custom-call-done(async-update)
-  }
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(hlo_str, GetModuleConfigForTest()));
-
-  HloInstruction* param =
-      module_->entry_computation()->parameter_instruction(0);
-  HloInstruction* async_start = FindInstruction(module_.get(), "async-start");
-  HloInstruction* async_update = FindInstruction(module_.get(), "async-update");
-  HloInstruction* async_done = FindInstruction(module_.get(), "async-done");
-
-  RunAnalysis();
-  EXPECT_FALSE(points_to_analysis_->GetPointsToSet(async_start).IsAmbiguous());
-  EXPECT_TRUE(points_to_analysis_->GetPointsToSet(async_start).IsDistinct());
-  EXPECT_FALSE(points_to_analysis_->GetPointsToSet(async_update).IsAmbiguous());
-  EXPECT_TRUE(points_to_analysis_->GetPointsToSet(async_update).IsDistinct());
-  EXPECT_FALSE(points_to_analysis_->GetPointsToSet(async_done).IsAmbiguous());
-  EXPECT_TRUE(points_to_analysis_->GetPointsToSet(async_done).IsDistinct());
-
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(async_start).element({}),
-      {async_start});
-  ExpectHasBufferAliases(
-      param, {}, {{param, {}}, {async_start, {0, 0}}, {async_update, {0, 0}}});
-  ExpectHasBufferAliases(
-      async_start, {1},
-      {{async_start, {1}}, {async_update, {1}}, {async_done, {}}});
-  ExpectHasBufferAliases(async_start, {2},
-                         {{async_start, {2}}, {async_update, {2}}});
-}
-
-TEST_F(TuplePointsToAnalysisTest, AsyncUpdate_LoopCrossing) {
-  const char* const kHlo = R"(
-HloModule test
-
-while_body {
-  state = ((f32[2,3]), f32[2,3], s32[]) parameter(0)
-  ROOT update = ((f32[2,3]), f32[2,3], s32[]) custom-call-update(state)
-}
-
-while_condition {
-  state = ((f32[2,3]), f32[2,3], s32[]) parameter(0)
-  ROOT predicate = pred[] constant(true)
-}
-
-ENTRY main {
-  p0 = f32[2,3] constant(0)
-  start = ((f32[2,3]), f32[2,3], s32[]) custom-call-start(p0),
-    custom_call_target="baz"
-  ROOT while_op = ((f32[2,3]), f32[2,3], s32[]) while(start), condition=while_condition, body=while_body
-}
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(kHlo, GetModuleConfigForTest()));
-  RunAnalysis();
-  HloInstruction* update = FindInstruction(module_.get(), "update");
-  EXPECT_FALSE(points_to_analysis_->GetPointsToSet(update).IsAmbiguous());
-}
-
-TEST_F(TuplePointsToAnalysisTest, SendAndSendDone) {
-  // Send forwards its operand to the output tuple at {0}.
-  std::string hlo_str = R"(
-HloModule SendAndSendDone
-
-ENTRY main {
-  constant = f32[] constant(1)
-  token0 = token[] after-all()
-  send = (f32[], u32[], token[]) send(constant, token0), channel_id=0
-  ROOT send_done = token[] send-done(send), channel_id=0
-}
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(hlo_str, GetModuleConfigForTest()));
-  RunAnalysis();
-
-  auto* constant = FindInstruction(module_.get(), "constant");
-  auto* send = FindInstruction(module_.get(), "send");
-  auto* send_done = FindInstruction(module_.get(), "send_done");
-
-  EXPECT_FALSE(points_to_analysis_->GetPointsToSet(send).IsAmbiguous());
-  EXPECT_TRUE(points_to_analysis_->GetPointsToSet(send).IsDistinct());
-  EXPECT_FALSE(points_to_analysis_->GetPointsToSet(send_done).IsAmbiguous());
-  EXPECT_TRUE(points_to_analysis_->GetPointsToSet(send_done).IsDistinct());
-
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(send).element({}), {send});
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(send).element({0}), {constant});
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(send_done).CreateFlattenedSet(),
-      {send_done});
-  ExpectHasBufferAliases(constant, {}, {{constant, {}}, {send, {0}}});
-}
-
-TEST_F(TuplePointsToAnalysisTest, RecvAndRecvDone) {
-  // RecvDone forwards its operand tuple element at {0} to the output.
-  std::string hlo_str = R"(
-HloModule RecvAndRecvDone
-
-ENTRY main {
-  token0 = token[] after-all()
-  recv = (f32[1,2,3], u32[], token[]) recv(token0), channel_id=0
-  ROOT recv_done = (f32[1,2,3], token[]) recv-done(recv), channel_id=0
-}
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(hlo_str, GetModuleConfigForTest()));
-  RunAnalysis();
-
-  auto* recv = FindInstruction(module_.get(), "recv");
-  auto* recv_done = FindInstruction(module_.get(), "recv_done");
-
-  EXPECT_FALSE(points_to_analysis_->GetPointsToSet(recv).IsAmbiguous());
-  EXPECT_TRUE(points_to_analysis_->GetPointsToSet(recv).IsDistinct());
-  EXPECT_FALSE(points_to_analysis_->GetPointsToSet(recv_done).IsAmbiguous());
-  EXPECT_TRUE(points_to_analysis_->GetPointsToSet(recv_done).IsDistinct());
-
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(recv).element({}), {recv});
-  ExpectHasBufferAliases(recv, {0}, {{recv, {0}}, {recv_done, {0}}});
-}
-
-TEST_F(TuplePointsToAnalysisTest, TupleWithBitcast) {
-  // Bitcast is an alias of its operand. A tuple with a bitcast element should
-  // have the operand of the bitcast in its points-to set.
-  std::string hlo_str = R"(
-HloModule TupleWithBitcast
-
-ENTRY main {
-  constant1 = f32[] constant(1)
-  constant2 = f32[] constant(2)
-  bitcast = f32[] bitcast(constant2)
-  ROOT tuple = (f32[], f32[]) tuple(constant1, bitcast)
-}
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(hlo_str, GetModuleConfigForTest()));
-  RunAnalysis();
-
-  auto* constant1 = FindInstruction(module_.get(), "constant1");
-  auto* constant2 = FindInstruction(module_.get(), "constant2");
-  auto* bitcast = FindInstruction(module_.get(), "bitcast");
-  auto* tuple = FindInstruction(module_.get(), "tuple");
-
-  EXPECT_EQ(1, points_to_analysis_->GetPointsToSet(bitcast).size());
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(bitcast).element({}), {constant2});
-  EXPECT_TRUE(
-      points_to_analysis_->GetPointsToSet(bitcast).tuple_sources({}).empty());
-
-  EXPECT_EQ(3, points_to_analysis_->GetPointsToSet(tuple).size());
-  EXPECT_FALSE(points_to_analysis_->GetPointsToSet(tuple).IsAmbiguous());
-  EXPECT_THAT(points_to_analysis_->GetPointsToSet(tuple).tuple_sources({}),
-              UnorderedElementsAre(tuple));
-
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(tuple).CreateFlattenedSet(),
-      {constant1, constant2, tuple});
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(tuple).element({}), {tuple});
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(tuple).element({0}), {constant1});
-  ExpectHasTopLevelBuffers(
-      points_to_analysis_->GetPointsToSet(tuple).element({1}), {constant2});
-}
-
-TEST_F(TuplePointsToAnalysisTest, PointsToTupleConstantElements) {
-  // Construct a tuple constant and kCopy it. Verify the points-to set of the
-  // copy correctly points into the nested elements of the constant.
-  std::string hlo_str = R"(
-HloModule PointsToTupleConstantElements
-
-ENTRY %PointsToTupleConstantElements () -> (f32[2,1], f32[2]) {
-  %constant = (f32[2,1], f32[2]) constant( ({ {1.0}, {2.0} }, {2.0, 42.0}) )
-  ROOT %copy = (f32[2,1], f32[2]) copy(%constant)
-}
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(hlo_str, GetModuleConfigForTest()));
-  RunAnalysis();
-
-  auto* tuple_constant = FindInstruction(module_.get(), "constant");
-  auto* copy = FindInstruction(module_.get(), "copy");
-
-  auto& points_to_set = points_to_analysis_->GetPointsToSet(copy);
-
-  ExpectHasBuffers(points_to_set.element({}), {GetBuffer(copy, {})});
-  ExpectHasBuffers(points_to_set.element({0}),
-                   {GetBuffer(tuple_constant, {0})});
-  ExpectHasBuffers(points_to_set.element({1}),
-                   {GetBuffer(tuple_constant, {1})});
-}
-
-TEST_F(TuplePointsToAnalysisTest, BufferAliases) {
-  // Create a nested tuple in which individual elements appear multiple
-  // times. Verify buffer alias sets.
-  std::string hlo_str = R"(
-HloModule BufferAliases
-
-ENTRY main {
-  constant1 = f32[] constant(1)
-  constant2 = f32[] constant(2)
-  inner_tuple = (f32[], f32[]) tuple(constant1, constant2)
-  ROOT tuple = ((f32[], f32[]), f32[]) tuple(inner_tuple, constant2)
-}
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(hlo_str, GetModuleConfigForTest()));
-  RunAnalysis();
-
-  auto* constant1 = FindInstruction(module_.get(), "constant1");
-  auto* constant2 = FindInstruction(module_.get(), "constant2");
-  auto* inner_tuple = FindInstruction(module_.get(), "inner_tuple");
-  auto* tuple = FindInstruction(module_.get(), "tuple");
-
-  ExpectHasBufferAliases(
-      constant1, /*index=*/{},
-      {{constant1, {}}, {inner_tuple, {0}}, {tuple, {0, 0}}});
-  ExpectHasBufferAliases(
-      constant2, /*index=*/{},
-      {{constant2, {}}, {inner_tuple, {1}}, {tuple, {0, 1}}, {tuple, {1}}});
-  ExpectHasBufferAliases(inner_tuple, /*index=*/{},
-                         {{inner_tuple, {}}, {tuple, {0}}});
-  ExpectHasBufferAliases(tuple, /*index=*/{}, {{tuple, {}}});
-}
-
-// TODO(b/277597692): The fix for this bug causes this test to fail.
-TEST_F(TuplePointsToAnalysisTest, DISABLED_CustomCall) {
-  std::string hlo_str = R"(
-HloModule DISABLED_CustomCall
-
-ENTRY main {
-  constant = f32[] constant(1)
-  ccall = (f32[], f32[]) custom-call(f32[] constant), custom_call_target="TestOp", output_to_operand_aliasing={{1}: (0, {})}
-  gte0 = f32[] get-tuple-element(ccall), index=0
-  ROOT gte1 = f32[] get-tuple-element(ccall), index=1
-}
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(hlo_str, GetModuleConfigForTest()));
-  RunAnalysis();
-
-  auto* constant = FindInstruction(module_.get(), "constant");
-  auto* ccall = FindInstruction(module_.get(), "ccall");
-  auto* gte0 = FindInstruction(module_.get(), "gte0");
-  auto* gte1 = FindInstruction(module_.get(), "gte1");
-
-  ExpectHasBufferAliases(ccall, /*index=*/{0}, {{gte0, {}}, {ccall, {0}}});
-  ExpectHasBufferAliases(constant, /*index=*/{},
-                         {{constant, {}}, {gte1, {}}, {ccall, {1}}});
-}
-
-class FusionPointsToAnalysisTest : public TuplePointsToAnalysisTest {
- protected:
-  // Takes a HLO string, parses it, runs points-to analysis, then checks for
-  // expected results (see unit test cases for example computation graphs).
-  void Run(const std::string& hlo_str, int64_t expected_num_users) {
-    ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_str));
-    // Get computation root instruction (should be a kFusion).
-    auto* fusion = module_->entry_computation()->root_instruction();
-    auto* tuple_param0 = fusion->operand(0);
-    // Run points-to analysis (should include fused instructions from 'fusion').
-    RunAnalysis();
-
-    // Check points-to set of fusion parameter associated with 'tuple_param0'.
-    auto* fusion_param = GetFusionParameterForOperand(fusion, tuple_param0);
-    ExpectHasBuffers(
-        points_to_analysis_->GetPointsToSet(fusion_param).element({}),
-        {GetBuffer(fusion_param, {})});
-    ExpectHasBuffers(
-        points_to_analysis_->GetPointsToSet(fusion_param).element({0}),
-        {GetBuffer(fusion_param, {0})});
-    ExpectHasBuffers(
-        points_to_analysis_->GetPointsToSet(fusion_param).element({1}),
-        {GetBuffer(fusion_param, {1})});
-
-    // Check that Gte at tuple_index = 0 points-to fusion_param({0})
-    auto fused_gte0 = GetUniqueFusionParameterUserAt(fusion_param, 0);
-    ExpectHasBuffers(
-        points_to_analysis_->GetPointsToSet(fused_gte0).element({}),
-        {GetBuffer(fusion_param, {0})});
-    // Check that Gte at tuple_index = 1 points-to fusion_param({1})
-    auto fused_gte1 = GetUniqueFusionParameterUserAt(fusion_param, 1);
-    ExpectHasBuffers(
-        points_to_analysis_->GetPointsToSet(fused_gte1).element({}),
-        {GetBuffer(fusion_param, {1})});
-
-    // Check buffer aliases of 'fusion_param' at shape index {0}.
-    ExpectHasBufferAliases(fusion_param, /*index=*/{0},
-                           {{fusion_param, {0}}, {fused_gte0, {}}});
-    // Check buffer aliases of 'fusion_param' at shape index {1}.
-    ExpectHasBufferAliases(fusion_param, /*index=*/{1},
-                           {{fusion_param, {1}}, {fused_gte1, {}}});
-
-    // Check number of users of 'fusion_param' aliases at shape index {0}.
-    ExpectNumUsersOfAliases(fusion_param, {0}, expected_num_users);
-  }
-
-  // Returns fusion parameter (from 'fusion.fused_instructions') corresponding
-  // to fusion 'operand'.
-  HloInstruction* GetFusionParameterForOperand(HloInstruction* fusion,
-                                               const HloInstruction* operand) {
-    const auto& fused_instructions = fusion->fused_instructions();
-    auto it =
-        absl::c_find_if(fused_instructions, [&](const HloInstruction* fused) {
-          return fused->opcode() == HloOpcode::kParameter &&
-                 fusion->operand(fused->parameter_number()) == operand;
-        });
-    CHECK(it != fusion->fused_instructions().end());
-    return *it;
-  }
-
-  // Returns all users of 'fusion_paran' at 'tuple_index'.
-  std::vector<HloInstruction*> GetFusionParameterUsersAt(
-      HloInstruction* fusion_param, int64_t tuple_index) {
-    CHECK(fusion_param->shape().IsTuple());
-    std::vector<HloInstruction*> users_at_tuple_index;
-    for (auto user : fusion_param->users()) {
-      CHECK_EQ(HloOpcode::kGetTupleElement, user->opcode());
-      if (user->tuple_index() == tuple_index) {
-        users_at_tuple_index.push_back(user);
-      }
-    }
-    return users_at_tuple_index;
-  }
-
-  // Returns the unique user of 'fusion_param' at 'tuple_index'.
-  HloInstruction* GetUniqueFusionParameterUserAt(HloInstruction* fusion_param,
-                                                 int64_t tuple_index) {
-    std::vector<HloInstruction*> users =
-        GetFusionParameterUsersAt(fusion_param, tuple_index);
-    CHECK_EQ(1, users.size());
-    return users[0];
-  }
-
-  // Checks that the count of all users of all aliases of 'instruction' at
-  // 'index' match 'expected_num_users'.
-  void ExpectNumUsersOfAliases(const HloInstruction* instruction,
-                               const ShapeIndex& index,
-                               const int64_t expected_num_users) {
-    const auto* buffer = GetBuffer(instruction, index);
-    int64_t num_users = 0;
-    for (const auto& alias : points_to_analysis_->GetBufferAliases(*buffer)) {
-      for (auto user : alias.instruction()->users()) {
-        if (user->opcode() == HloOpcode::kGetTupleElement && !index.empty()) {
-          // Gte instructions only access the top-level buffer of their operand.
-          continue;
-        }
-        ++num_users;
-      }
-    }
-    EXPECT_EQ(expected_num_users, num_users);
-  }
-};
-
-// Tests the points-to set of tuple-shaped fusion parameter 0 and all GTE users.
-// Tests the alias set of tuple-shaped fusion parameter 0 at all shape indices.
-// Tests that there is a single user of the aliases of tuple-shaped fusion
-// parameter 0 at shape index {0}.
-//
-//             Param0    Const
-//                 \      /
-//                  Fusion
-//                 /      \
-//        FusionParam0   FusionParam1
-//        /          \       |
-//     Gte(0) Const  Gte(1)  /
-//        \     |      \    /
-//         \    |       Add
-//          \   |        /
-//           \0 |2      /1
-//          DynamicUpdateSlice  // fused root.
-//
-TEST_F(FusionPointsToAnalysisTest, FusionParam0OneUser) {
-  std::string hlo_str = R"(
-HloModule FusionParam0OneUser
-
-%fused_computation (param_1.2: (f32[8], f32[3])) -> f32[8] {
-  %param_1.2 = (f32[8]{0}, f32[3]{0}) parameter(0)
-  %get-tuple-element.1 = f32[8]{0} get-tuple-element((f32[8]{0}, f32[3]{0}) %param_1.2), index=0
-  %get-tuple-element.2 = f32[3]{0} get-tuple-element((f32[8]{0}, f32[3]{0}) %param_1.2), index=1
-  %constant.3 = f32[3]{0} constant({1, 1, 1})
-  %add.1 = f32[3]{0} add(f32[3]{0} %get-tuple-element.2, f32[3]{0} %constant.3)
-  %constant.2 = s32[] constant(0)
-  ROOT %dynamic-update-slice.1 = f32[8]{0} dynamic-update-slice(f32[8]{0} %get-tuple-element.1, f32[3]{0} %add.1, s32[] %constant.2)
-}
-
-ENTRY %FusionParam0OneUser (param0: (f32[8], f32[3])) -> f32[8] {
-  %param0 = (f32[8]{0}, f32[3]{0}) parameter(0)
-  ROOT %fusion = f32[8]{0} fusion((f32[8]{0}, f32[3]{0}) %param0), kind=kLoop, calls=%fused_computation
-}
-)";
-  Run(hlo_str, /*expected_num_users=*/1);
-}
-
-TEST_F(FusionPointsToAnalysisTest,
-       FusionParam0OneUserWithUnreachableInstructionInFusion) {
-  std::string hlo_str = R"(
-HloModule FusionParam0OneUser
-
-%fused_computation (param_1.2: (f32[8], f32[3], f32[7])) -> f32[8] {
-  %param_1.2 = (f32[8]{0}, f32[3]{0}, f32[7]{0}) parameter(0)
-  %get-tuple-element.1 = f32[8]{0} get-tuple-element(%param_1.2), index=0
-  %get-tuple-element.2 = f32[3]{0} get-tuple-element(%param_1.2), index=1
-  %get-tuple-element.3 = f32[7]{0} get-tuple-element(%param_1.2), index=2
-  %placeholder = f32[7]{0} custom-call(%get-tuple-element.3), custom_call_target="IntermediateBufferDummyConsumer", custom_call_has_side_effect=true
-  %constant.3 = f32[3]{0} constant({1, 1, 1})
-  %add.1 = f32[3]{0} add(f32[3]{0} %get-tuple-element.2, f32[3]{0} %constant.3)
-  %constant.2 = s32[] constant(0)
-  ROOT %dynamic-update-slice.1 = f32[8]{0} dynamic-update-slice(f32[8]{0} %get-tuple-element.1, f32[3]{0} %add.1, s32[] %constant.2)
-}
-
-ENTRY %FusionParam0OneUser (param0: (f32[8], f32[3], f32[7])) -> f32[8] {
-  %param0 = (f32[8]{0}, f32[3]{0}, f32[7]{0}) parameter(0)
-  ROOT %fusion = f32[8]{0} fusion(%param0), kind=kLoop, calls=%fused_computation
-}
-)";
-  Run(hlo_str, /*expected_num_users=*/1);
-}
-
-// Tests the points-to set of tuple-shaped fusion parameter 0 and all GTE users.
-// Tests the alias set of tuple-shaped fusion parameter 0 at all shape indices.
-// Tests that there are two users of the aliases of tuple-shaped fusion
-// parameter 0 at shape index {0}.
-//
-//             Param0    Const
-//                 \      /
-//                  Fusion
-//                 /      \
-//        FusionParam0   FusionParam1
-//              |    \       |
-//            Gte(0) Gte(1)  /
-//              |      \    /
-//              |\      Add
-//        Const | \      /
-//           |  | Slice /
-//           |  |   \  /
-//           |  |   Add
-//           |  |    |
-//           |2 |0   |1
-//          DynamicUpdateSlice  // fused root.
-//
-TEST_F(FusionPointsToAnalysisTest, FusionParam0TwoUsers) {
-  std::string hlo_str = R"(
-HloModule FusionParam0TwoUsers
-
-%fused_computation (param_1.2: (f32[8], f32[3])) -> f32[8] {
-  %param_1.2 = (f32[8]{0}, f32[3]{0}) parameter(0)
-  %get-tuple-element.1 = f32[8]{0} get-tuple-element((f32[8]{0}, f32[3]{0}) %param_1.2), index=0
-  %get-tuple-element.2 = f32[3]{0} get-tuple-element((f32[8]{0}, f32[3]{0}) %param_1.2), index=1
-  %constant.3 = f32[3]{0} constant({1, 1, 1})
-  %add.1 = f32[3]{0} add(f32[3]{0} %get-tuple-element.2, f32[3]{0} %constant.3)
-  %slice = f32[3]{0} slice(f32[8]{0} %get-tuple-element.1), slice={[0:3]}
-  %add.2 = f32[3]{0} add(f32[3]{0} %add.1, f32[3]{0} %slice)
-  %constant.2 = s32[] constant(0)
-  ROOT %dynamic-update-slice.1 = f32[8]{0} dynamic-update-slice(f32[8]{0} %get-tuple-element.1, f32[3]{0} %add.2, s32[] %constant.2)
-}
-
-ENTRY %FusionParam0TwoUsers (param0: (f32[8], f32[3])) -> f32[8] {
-  %param0 = (f32[8]{0}, f32[3]{0}) parameter(0)
-  ROOT %fusion = f32[8]{0} fusion((f32[8]{0}, f32[3]{0}) %param0), kind=kLoop, calls=%fused_computation
-}
-)";
-  Run(hlo_str, /*expected_num_users=*/2);
-}
-
-class PointsToAnalysisTestBase : public HloHardwareIndependentTestBase {
- protected:
-  void BuildModule(std::unique_ptr<HloComputation> computation) {
-    module_ = CreateNewVerifiedModule();
-    computation_ = module_->AddEntryComputation(std::move(computation));
-  }
-
-  void RunAnalysis() {
-    CHECK_NOTNULL(module_.get());
-    points_to_analysis_ = TuplePointsToAnalysis::Run(module_.get()).value();
-  }
-
-  void BuildModuleAndRunAnalysis(std::unique_ptr<HloComputation> computation) {
-    BuildModule(std::move(computation));
-    RunAnalysis();
-  }
-
-  std::unique_ptr<HloModule> module_;
-  HloComputation* computation_ = nullptr;
-  std::unique_ptr<TuplePointsToAnalysis> points_to_analysis_;
-};
-
-
-TEST_F(TuplePointsToAnalysisTest, AsyncUpdateChangesContext) {
-  absl::string_view hlo_string = R"hlo(
-HloModule module
-
-async_computation {
-  p0 = f32[4] parameter(0)
-  p1 = f32[4] parameter(1)
-  ROOT add = f32[4] add(p0, p1)
-}
-
-ENTRY entry {
-  x = f32[4] parameter(0)
-  y = f32[4] parameter(1)
-  start = ((), (), (s32[], s32[4])) async-start(), calls=async_computation,
-                                async_execution_thread="sparsecore"
-  update = ((f32[4], f32[4]), f32[4], (s32[8], s32[])) async-update(start, x, y),
-              output_to_operand_aliasing={{2, 1}: (0, {2, 0})}
-  ROOT done = f32[4] async-done(update)
-}
-  )hlo";
-
-  ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string));
-  RunAnalysis();
-
-  HloInstruction* x = module_->entry_computation()->GetInstructionWithName("x");
-  HloInstruction* y = module_->entry_computation()->GetInstructionWithName("y");
-  HloInstruction* start =
-      module_->entry_computation()->GetInstructionWithName("start");
-  HloInstruction* update =
-      module_->entry_computation()->GetInstructionWithName("update");
-  HloInstruction* done =
-      module_->entry_computation()->GetInstructionWithName("done");
-
-  const LogicalBuffer* B_x = GetBuffer(x, {});
-  const LogicalBuffer* B_y = GetBuffer(y, {});
-
-  const LogicalBuffer* B_start_c0 = GetBuffer(start, {2, 0});
-  const LogicalBuffer* B_start_c1 = GetBuffer(start, {2, 1});
-  EXPECT_EQ(start, B_start_c0->instruction());
-  EXPECT_EQ(start, B_start_c1->instruction());
-
-  const PointsToSet& update_set = points_to_analysis_->GetPointsToSet(update);
-  ExpectHasBuffers(update_set.element({0, 0}), {B_x});
-  ExpectHasBuffers(update_set.element({0, 1}), {B_y});
-  ExpectHasBuffers(update_set.element({2, 1}), {B_start_c0});
-
-  const LogicalBuffer* B_update_out = GetBuffer(update, {1});
-  const LogicalBuffer* B_update_c0 = GetBuffer(update, {2, 0});
-  EXPECT_EQ(update, B_update_out->instruction());
-  EXPECT_EQ(update, B_update_c0->instruction());
-
-  const PointsToSet& done_set = points_to_analysis_->GetPointsToSet(done);
-  ExpectHasBuffers(done_set.element({}), {B_update_out});
-}
-
-TEST_F(TuplePointsToAnalysisTest, LateBindingAsyncUpdate) {
-  absl::string_view hlo_string = R"hlo(
-HloModule module
-
-async_computation {
-  p0 = f32[4] parameter(0)
-  p1 = f32[4] parameter(1)
-  ROOT add = f32[4] add(p0, p1)
-}
-
-ENTRY entry {
-  x = f32[4] parameter(0)
-  y = f32[4] parameter(1)
-  start = ((), (), s32[]) async-start(), calls=async_computation,
-                                async_execution_thread="sparsecore"
-  update = ((f32[4], f32[4]), f32[4], s32[]) async-update(start, x, y),
-              output_to_operand_aliasing={{2}: (0, {2})}
-  ROOT done = f32[4] async-done(update)
-}
-  )hlo";
-
-  ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string));
-  RunAnalysis();
-
-  HloInstruction* x = module_->entry_computation()->GetInstructionWithName("x");
-  HloInstruction* y = module_->entry_computation()->GetInstructionWithName("y");
-  HloInstruction* start =
-      module_->entry_computation()->GetInstructionWithName("start");
-  HloInstruction* update =
-      module_->entry_computation()->GetInstructionWithName("update");
-  HloInstruction* done =
-      module_->entry_computation()->GetInstructionWithName("done");
-
-  const LogicalBuffer* B_x = GetBuffer(x, {});
-  const LogicalBuffer* B_y = GetBuffer(y, {});
-  const LogicalBuffer* B_start_sflag = GetBuffer(start, {2});
-
-  const PointsToSet& update_set = points_to_analysis_->GetPointsToSet(update);
-  ExpectHasBuffers(update_set.element({0, 0}), {B_x});
-  ExpectHasBuffers(update_set.element({0, 1}), {B_y});
-  ExpectHasBuffers(update_set.element({2}), {B_start_sflag});
-
-  const LogicalBuffer* B_update_out = GetBuffer(update, {1});
-  EXPECT_EQ(update, B_update_out->instruction());
-
-  const PointsToSet& done_set = points_to_analysis_->GetPointsToSet(done);
-  ExpectHasBuffers(done_set.element({}), {B_update_out});
-}
-
-TEST_F(TuplePointsToAnalysisTest,
-       LateBindingAsync_InputAliased_OutputBoundEarly) {
-  absl::string_view hlo_string = R"hlo(
-HloModule module
-
-async_computation {
-  p0 = f32[4] parameter(0)
-  ROOT custom-call = (f32[4]) custom-call(p0), custom_call_target="foo"
-}
-
-ENTRY entry {
-  param = f32[4] parameter(0)
-  async-start = ((), (f32[4]), s32[]) async-start(), calls=async_computation,
-                                      output_to_operand_aliasing={{1, 0}: (0, {})}
-  async-update = ((f32[4]), (f32[4]), s32[]) async-update(async-start, param)
-  ROOT async-done = (f32[4]) async-done(async-update)
-}
-  )hlo";
-
-  ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string));
-  RunAnalysis();
-
-  HloInstruction* param =
-      module_->entry_computation()->GetInstructionWithName("param");
-  HloInstruction* start =
-      module_->entry_computation()->GetInstructionWithName("async-start");
-  HloInstruction* update =
-      module_->entry_computation()->GetInstructionWithName("async-update");
-  HloInstruction* done =
-      module_->entry_computation()->GetInstructionWithName("async-done");
-
-  const LogicalBuffer* B_param = GetBuffer(param, {});
-
-  // With back-propagation, the resolved alias is propagated back to start.
-  const PointsToSet& start_set = points_to_analysis_->GetPointsToSet(start);
-  ExpectHasBuffers(start_set.element({1, 0}), {B_param});
-
-  // At async-update, the alias is active.
-  const PointsToSet& update_set = points_to_analysis_->GetPointsToSet(update);
-  ExpectHasBuffers(update_set.element({0, 0}), {B_param});
-  ExpectHasBuffers(update_set.element({1, 0}), {B_param});  // Alias resolved!
-
-  // At async-done, the output is forwarded.
-  const PointsToSet& done_set = points_to_analysis_->GetPointsToSet(done);
-  ExpectHasBuffers(done_set.element({0}), {B_param});
-}
-
-TEST_F(TuplePointsToAnalysisTest,
-       LateBindingAsync_InputAliased_OutputBoundLate) {
-  absl::string_view hlo_string = R"hlo(
-HloModule module
-
-async_computation {
-  p0 = f32[4] parameter(0)
-  p1 = f32[4] parameter(1)
-  ROOT tuple = (f32[4], f32[4]) tuple(p0, p1)
-}
-
-ENTRY entry {
-  param0 = f32[4] parameter(0)
-  param1 = f32[4] parameter(1)
-  async-start = ((f32[4]), (), s32[]) async-start(param0),
-                                      calls=async_computation,
-                                      output_to_operand_aliasing={{1, 1}: (0, {})}
-  async-update = ((f32[4], f32[4]), (f32[4], f32[4]), s32[]) async-update(async-start, param1)
-  ROOT async-done = (f32[4], f32[4]) async-done(async-update)
-}
-  )hlo";
-
-  ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string));
-  RunAnalysis();
-
-  HloInstruction* param0 =
-      module_->entry_computation()->GetInstructionWithName("param0");
-  HloInstruction* param1 =
-      module_->entry_computation()->GetInstructionWithName("param1");
-  HloInstruction* start =
-      module_->entry_computation()->GetInstructionWithName("async-start");
-  HloInstruction* update =
-      module_->entry_computation()->GetInstructionWithName("async-update");
-  HloInstruction* done =
-      module_->entry_computation()->GetInstructionWithName("async-done");
-
-  const LogicalBuffer* B_param0 = GetBuffer(param0, {});
-  const LogicalBuffer* B_param1 = GetBuffer(param1, {});
-
-  // At async-start, input is bound but output {1, 1} doesn't exist.
-  const PointsToSet& start_set = points_to_analysis_->GetPointsToSet(start);
-  ExpectHasBuffers(start_set.element({0, 0}), {B_param0});
-  // {1} is empty tuple, so no {1, 1} index.
-
-  // At async-update, both are bound, alias resolved.
-  const PointsToSet& update_set = points_to_analysis_->GetPointsToSet(update);
-  ExpectHasBuffers(update_set.element({0, 0}), {B_param0});
-  ExpectHasBuffers(update_set.element({0, 1}), {B_param1});
-  ExpectHasBuffers(update_set.element({1, 1}), {B_param0});  // Alias resolved!
-
-  // {1, 0} is new output, not aliased, should have a new buffer.
-  const auto& r0_buffers = update_set.element({1, 0});
-  ASSERT_EQ(1, r0_buffers.size());
-  EXPECT_EQ(update, r0_buffers[0]->instruction());
-
-  // At async-done, outputs are forwarded.
-  const PointsToSet& done_set = points_to_analysis_->GetPointsToSet(done);
-  ExpectHasBuffers(done_set.element({0}), r0_buffers);
-  ExpectHasBuffers(done_set.element({1}), {B_param0});
-}
-
-TEST_F(TuplePointsToAnalysisTest, LateBindingAsync_PostponedOutput) {
-  absl::string_view hlo_string = R"hlo(
-HloModule module
-
-async_computation {
-  p0 = f32[4] parameter(0)
-  ROOT custom-call = f32[4] custom-call(p0), custom_call_target="foo"
-}
-
-ENTRY entry {
-  param = f32[4] parameter(0)
-  async-start = ((), (), s32[]) async-start(), calls=async_computation,
-                                output_to_operand_aliasing={{1}: (0, {})}
-  async-update = ((f32[4]), (), s32[]) async-update(async-start, param)
-  ROOT async-done = f32[4] async-done(async-update)
-}
-  )hlo";
-
-  ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string));
-  RunAnalysis();
-
-  HloInstruction* param =
-      module_->entry_computation()->GetInstructionWithName("param");
-  HloInstruction* start =
-      module_->entry_computation()->GetInstructionWithName("async-start");
-  HloInstruction* update =
-      module_->entry_computation()->GetInstructionWithName("async-update");
-  HloInstruction* done =
-      module_->entry_computation()->GetInstructionWithName("async-done");
-
-  const LogicalBuffer* B_param = GetBuffer(param, {});
-
-  // With back-propagation, the resolved alias is propagated back to start and
-  // update.
-  const PointsToSet& start_set = points_to_analysis_->GetPointsToSet(start);
-  ExpectHasBuffers(start_set.element({1}), {B_param});
-  const PointsToSet& update_set = points_to_analysis_->GetPointsToSet(update);
-  ExpectHasBuffers(update_set.element({1}), {B_param});
-
-  // At async-done, output is bound and alias is resolved.
-  const PointsToSet& done_set = points_to_analysis_->GetPointsToSet(done);
-  ExpectHasBuffers(done_set.element({}),
-                   {B_param});  // Alias resolved at async-done!
-}
-
-TEST_F(TuplePointsToAnalysisTest, LateBindingAsync_MultipleUpdates) {
-  absl::string_view hlo_string = R"hlo(
-HloModule module
-
-async_computation {
-  p0 = f32[4] parameter(0)
-  p1 = f32[4] parameter(1)
-  ROOT tuple = (f32[4], f32[4]) tuple(p0, p1)
-}
-
-ENTRY entry {
-  param0 = f32[4] parameter(0)
-  param1 = f32[4] parameter(1)
-  async-start = ((), (), s32[]) async-start(), calls=async_computation
-  async-update1 = ((f32[4]), (f32[4]), s32[]) async-update(async-start, param0)
-  async-update2 = ((f32[4], f32[4]), (f32[4], f32[4]), s32[]) async-update(async-update1, param1)
-  ROOT async-done = (f32[4], f32[4]) async-done(async-update2)
-}
-  )hlo";
-
-  ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string));
-  RunAnalysis();
-
-  HloInstruction* param0 =
-      module_->entry_computation()->GetInstructionWithName("param0");
-  HloInstruction* param1 =
-      module_->entry_computation()->GetInstructionWithName("param1");
-  HloInstruction* update1 =
-      module_->entry_computation()->GetInstructionWithName("async-update1");
-  HloInstruction* update2 =
-      module_->entry_computation()->GetInstructionWithName("async-update2");
-  HloInstruction* done =
-      module_->entry_computation()->GetInstructionWithName("async-done");
-
-  const LogicalBuffer* B_param0 = GetBuffer(param0, {});
-  const LogicalBuffer* B_param1 = GetBuffer(param1, {});
-
-  // At async-update1, param0 is bound.
-  const PointsToSet& update1_set = points_to_analysis_->GetPointsToSet(update1);
-  ExpectHasBuffers(update1_set.element({0, 0}), {B_param0});
-
-  // At async-update2, param0 is forwarded and param1 is bound.
-  const PointsToSet& update2_set = points_to_analysis_->GetPointsToSet(update2);
-  ExpectHasBuffers(update2_set.element({0, 0}), {B_param0});
-  ExpectHasBuffers(update2_set.element({0, 1}),
-                   {B_param1});  // Verified bug fix!
-
-  // Outputs should be forwarded to async-done.
-  const PointsToSet& done_set = points_to_analysis_->GetPointsToSet(done);
-  ExpectHasBuffers(done_set.element({0}), update2_set.element({1, 0}));
-  ExpectHasBuffers(done_set.element({1}), update2_set.element({1, 1}));
-}
-
-TEST_F(TuplePointsToAnalysisTest, CallMarkerBeforePointsToOutputs) {
-  std::string hlo_str = R"(
-HloModule CallMarkerBeforePointsToOutputs
-
-ENTRY main {
-  p0 = f32[2,2] parameter(0)
-  p1 = f32[4,4] parameter(1)
-  ROOT cc = (f32[2,2], f32[4,4]) custom-call(p0, p1),
-    custom_call_target="__xla_internal_call_marker_before"
-}
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(hlo_str, GetModuleConfigForTest()));
-  RunAnalysis();
-
-  auto* p0 = FindInstruction(module_.get(), "p0");
-  auto* p1 = FindInstruction(module_.get(), "p1");
-  auto* cc = FindInstruction(module_.get(), "cc");
-
-  // The custom call output element {0} should point to p0's buffer.
-  ExpectHasTopLevelBuffers(points_to_analysis_->GetPointsToSet(cc).element({0}),
-                           {p0});
-  // The custom call output element {1} should point to p1's buffer.
-  ExpectHasTopLevelBuffers(points_to_analysis_->GetPointsToSet(cc).element({1}),
-                           {p1});
-}
-
-TEST_F(TuplePointsToAnalysisTest, CallMarkerAfterPointsToOutput) {
-  std::string hlo_str = R"(
-HloModule CallMarkerAfterPointsToOutput
-
-ENTRY main {
-  p0 = f32[2,2] parameter(0)
-  ROOT cc = f32[2,2] custom-call(p0),
-    custom_call_target="__xla_internal_call_marker_after"
-}
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(hlo_str, GetModuleConfigForTest()));
-  RunAnalysis();
-
-  auto* p0 = FindInstruction(module_.get(), "p0");
-  auto* cc = FindInstruction(module_.get(), "cc");
-
-  // The custom call output should point to p0's buffer.
-  ExpectHasTopLevelBuffers(points_to_analysis_->GetPointsToSet(cc).element({}),
-                           {p0});
-}
-
-// Test that if async-done has incompatible output shape (e.g., empty tuple at
-// index {1} in async-start), AsyncDone (which returns a non-tuple array) does
-// not forward the points-to set of the empty tuple from async-start, but
-// instead defines its own buffer.
-TEST_F(TuplePointsToAnalysisTest, AsyncDoneNoBindOutput) {
-  std::string hlo_str = R"(
-HloModule AsyncDoneNoBindOutput
-
-%async_comp (param: f32[8,8]) -> f32[8,8] {
-  %param = f32[8,8]{0,1} parameter(0)
-  ROOT %copy = f32[8,8]{0,1} copy(%param)
-}
-
-ENTRY main {
-  %param = f32[8,8]{0,1} parameter(0)
-  %async_start = ((f32[8,8]{0,1}), (), s32[]) async-start(%param), calls=%async_comp
-  ROOT %async_done = f32[8,8]{0,1} async-done(%async_start)
-}
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(hlo_str, GetModuleConfigForTest()));
-  RunAnalysis();
-
-  auto* async_done = FindInstruction(module_.get(), "async_done");
-  const PointsToSet& points_to_set =
-      points_to_analysis_->GetPointsToSet(async_done);
-  EXPECT_FALSE(points_to_set.element({}).empty());
-  EXPECT_EQ(async_done, points_to_set.element({}).at(0)->instruction());
-}
-
-// Test that if async-done binds a tuple output (empty tuple at index {1} in
-// async-start),
-// and AsyncDone returns a tuple, we do not forward and instead define new
-// buffers for all elements of the AsyncDone tuple.
-TEST_F(TuplePointsToAnalysisTest, AsyncDoneNoBindOutputTuple) {
-  std::string hlo_str = R"(
-HloModule AsyncDoneNoBindOutputTuple
-
-%async_comp (param: f32[8,8]) -> (f32[8,8], f32[4,4]) {
-  %param = f32[8,8]{0,1} parameter(0)
-  ROOT %cc = (f32[8,8]{0,1}, f32[4,4]{0,1}) custom-call(%param), custom_call_target="some_target"
-}
-
-ENTRY main {
-  %param = f32[8,8]{0,1} parameter(0)
-  %async_start = ((f32[8,8]{0,1}), (), s32[]) async-start(%param), calls=%async_comp
-  ROOT %async_done = (f32[8,8]{0,1}, f32[4,4]{0,1}) async-done(%async_start)
-}
-)";
-  ASSERT_OK_AND_ASSIGN(
-      module_, ParseAndReturnVerifiedModule(hlo_str, GetModuleConfigForTest()));
-  RunAnalysis();
-
-  auto* async_done = FindInstruction(module_.get(), "async_done");
-  const PointsToSet& points_to_set =
-      points_to_analysis_->GetPointsToSet(async_done);
-
-  // Element {0} and {1} should have buffers defined by async_done itself.
-  EXPECT_FALSE(points_to_set.element({0}).empty());
-  EXPECT_EQ(async_done, points_to_set.element({0}).at(0)->instruction());
-  EXPECT_EQ(ShapeIndex({0}), points_to_set.element({0}).at(0)->index());
-
-  EXPECT_FALSE(points_to_set.element({1}).empty());
-  EXPECT_EQ(async_done, points_to_set.element({1}).at(0)->instruction());
-  EXPECT_EQ(ShapeIndex({1}), points_to_set.element({1}).at(0)->index());
-}
-
-}  // namespace
-}  // namespace xla
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD
index cb1de69e20d..a6a577a4643 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/BUILD
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/BUILD
@@ -1126,13 +1126,12 @@ cc_library(
         "//xla:util",
         "//xla:xla_data_proto_cc",
         "//xla/hlo/analysis:hlo_dataflow_analysis",
-        "//xla/hlo/analysis:tuple_points_to_analysis",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass",
         "//xla/hlo/utils:hlo_query",
         "//xla/service:call_graph",
         "//xla/service:hlo_cost_analysis",
-        "//xla/service:logical_buffer",
+        "//xla/service:hlo_value",
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/algorithm:container",
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc
index ea62adc53f0..5fcc642cf8d 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc
@@ -48,7 +48,6 @@ limitations under the License.
 #include "absl/time/time.h"
 #include "absl/types/span.h"
 #include "xla/hlo/analysis/hlo_dataflow_analysis.h"
-#include "xla/hlo/analysis/tuple_points_to_analysis.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_clone_context.h"
 #include "xla/hlo/ir/hlo_computation.h"
@@ -65,7 +64,7 @@ limitations under the License.
 #include "xla/map_util.h"
 #include "xla/service/call_graph.h"
 #include "xla/service/hlo_cost_analysis.h"
-#include "xla/service/logical_buffer.h"
+#include "xla/service/hlo_value.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/status_macros.h"
@@ -231,7 +230,7 @@ struct RematStrategy {
 // accesses the buffer at 'shape_idx'.
 bool FusedParameterUsesBuffer(const HloInstruction* fusion, int64_t op_idx,
                               const ShapeIndex& shape_idx,
-                              const TuplePointsToAnalysis& points_to_analysis) {
+                              const HloDataflowAnalysis& dataflow_analysis) {
   DCHECK(fusion->IsLoopFusion());
   DCHECK_GE(op_idx, 0);
   DCHECK_LT(op_idx, fusion->operand_count());
@@ -239,21 +238,20 @@ bool FusedParameterUsesBuffer(const HloInstruction* fusion, int64_t op_idx,
     return false;
   }
   const HloInstruction* fused_param = fusion->fused_parameter(op_idx);
-  absl::StatusOr<const LogicalBuffer*> buffer =
-      points_to_analysis.GetBufferDefinedAt(fused_param, shape_idx);
-  if (!buffer.ok()) {
+  if (!dataflow_analysis.ValueIsDefinedAt(fused_param, shape_idx)) {
     return false;
   }
-  for (const BufferAlias& alias :
-       points_to_analysis.GetBufferAliases(**buffer)) {
+  const HloValue& value =
+      dataflow_analysis.GetValueDefinedAt(fused_param, shape_idx);
+  for (const HloPosition& alias : value.positions()) {
     // If an alias of the buffer is the root of the fusion, the buffer is
     // an output of the fusion and therefore assumed to be used.
-    if (alias.instruction() == fusion->fused_expression_root()) {
+    if (alias.instruction == fusion->fused_expression_root()) {
       return true;
     }
-    for (const HloInstruction* alias_user : alias.instruction()->users()) {
+    for (const HloInstruction* alias_user : alias.instruction->users()) {
       if (alias_user->opcode() == HloOpcode::kGetTupleElement &&
-          !alias.index().empty()) {
+          !alias.index.empty()) {
         // GetTupleElement instructions only access the top-level buffer of
         // their operand.
         continue;
@@ -263,9 +261,9 @@ bool FusedParameterUsesBuffer(const HloInstruction* fusion, int64_t op_idx,
       }
       for (int64_t nested_op_idx = 0;
            nested_op_idx < alias_user->operand_count(); ++nested_op_idx) {
-        if (alias_user->operand(nested_op_idx) == alias.instruction() &&
-            FusedParameterUsesBuffer(alias_user, nested_op_idx, alias.index(),
-                                     points_to_analysis)) {
+        if (alias_user->operand(nested_op_idx) == alias.instruction &&
+            FusedParameterUsesBuffer(alias_user, nested_op_idx, alias.index,
+                                     dataflow_analysis)) {
           return true;
         }
       }
@@ -274,30 +272,28 @@ bool FusedParameterUsesBuffer(const HloInstruction* fusion, int64_t op_idx,
   return false;
 }
 
-// Return the items which use the given LogicalBuffer. Sets
+// Return the items which use the given HloValue. Sets
 // has_indirect_users to whether any of the uses is indirect. A use is indirect
-// if the instruction defining logical_buffer is not an operand of the use. This
+// if the instruction defining value is not an operand of the use. This
 // can happen via buffer aliasing (eg, tuples).
 UsesList GetUsers(const HloRematInstructionList& instruction_list,
-                  const LogicalBuffer* logical_buffer,
-                  const TuplePointsToAnalysis& points_to_analysis,
+                  const HloValue* value,
+                  const HloDataflowAnalysis& dataflow_analysis,
                   bool* has_indirect_users) {
   UsesList users;
   // To identify uses iterate through all HloInstruction users of the
-  // BufferAliases of the logical buffer.
+  // positions of the value.
   *has_indirect_users = false;
   const std::optional<int64_t> user_index =
-      logical_buffer->index().size() == 1
-          ? std::make_optional(logical_buffer->index().front())
-          : std::nullopt;
-  for (const BufferAlias& buffer_alias :
-       points_to_analysis.GetBufferAliases(*logical_buffer)) {
+      value->index().size() == 1 ? std::make_optional(value->index().front())
+                                 : std::nullopt;
+  for (const HloPosition& position : value->positions()) {
     const bool is_unsupported_indirect =
-        buffer_alias.instruction() != logical_buffer->instruction() &&
-        !IsSupportedIndirectUser(buffer_alias.instruction());
-    for (const HloInstruction* user : buffer_alias.instruction()->users()) {
+        position.instruction != value->instruction() &&
+        !IsSupportedIndirectUser(position.instruction);
+    for (const HloInstruction* user : position.instruction->users()) {
       if (user->opcode() == HloOpcode::kGetTupleElement &&
-          !buffer_alias.index().empty()) {
+          !position.index.empty()) {
         // GetTupleElement instructions only access the top-level buffer of
         // their operand.
         continue;
@@ -319,16 +315,16 @@ UsesList GetUsers(const HloRematInstructionList& instruction_list,
 
       if (user->operand_count() == 1) {
         if (!user->IsLoopFusion() ||
-            FusedParameterUsesBuffer(user, 0, buffer_alias.index(),
-                                     points_to_analysis)) {
+            FusedParameterUsesBuffer(user, 0, position.index,
+                                     dataflow_analysis)) {
           record_use(0);
         }
       } else {
         for (int64_t op_idx = 0; op_idx < user->operand_count(); ++op_idx) {
-          if (user->operand(op_idx) == buffer_alias.instruction()) {
+          if (user->operand(op_idx) == position.instruction) {
             if (!user->IsLoopFusion() ||
-                FusedParameterUsesBuffer(user, op_idx, buffer_alias.index(),
-                                         points_to_analysis)) {
+                FusedParameterUsesBuffer(user, op_idx, position.index,
+                                         dataflow_analysis)) {
               record_use(op_idx);
             }
           }
@@ -341,16 +337,16 @@ UsesList GetUsers(const HloRematInstructionList& instruction_list,
 
 // Class for tracking memory usage of a computation as the instructions are
 // placed sequentially. Memory usage is the sum of the sizes of live values
-// (LogicalBuffers) at the current point in the instruction sequence.
+// at the current point in the instruction sequence.
 class MemoryUsageTracker {
  public:
   MemoryUsageTracker(const HloRematerialization::Options& options,
                      const HloComputation* computation,
-                     const TuplePointsToAnalysis& points_to_analysis,
+                     const HloDataflowAnalysis& dataflow_analysis,
                      const HloRematInstructionList& instruction_list);
 
   // Starts the placement of the given instruction. This adds the sizes of the
-  // LogicalBuffers defined by the instruction to the current memory
+  // HloValues defined by the instruction to the current memory
   // usage. Placement is broken into two steps (BeginInstruction and
   // EndInstruction) to accurately model memory usage. At BeginInstruction the
   // memory for the output value(s) of the current instruction is allocated. At
@@ -503,10 +499,10 @@ class MemoryUsageTracker {
   std::string ToString() const;
 
  private:
-  // A Buffer represents a single LogicalBuffer in the computation including
-  // various metadata useful for tracking liveness of the value. A LogicalBuffer
+  // A Buffer represents a single HloValue in the computation including
+  // various metadata useful for tracking liveness of the value. An HloValue
   // is not used directly because the HLO graph is transformed and
-  // TuplePointsToAnalysis which owns all LogicalBuffers cannot be updated after
+  // HloDataflowAnalysis which owns all HloValues cannot be updated after
   // HLO graph transformations.
   struct Buffer {
     // The unique id of this Buffer. This value is equal to the buffer's index
@@ -563,17 +559,17 @@ class MemoryUsageTracker {
   // to avoid computing the shape multiple times.
   absl::StatusOr<const Shape*> GetCompactShape(const HloInstruction* hlo);
 
-  // Creates a Buffer representing the given logical buffer. The buffer is added
+  // Creates a Buffer representing the given HloValue. The buffer is added
   // to buffers_ and a reference is returned.
-  Buffer& CreateBufferFromLogicalBuffer(
-      const LogicalBuffer* logical_buffer,
-      const TuplePointsToAnalysis& points_to_analysis, bool live_out) {
+  Buffer& CreateBufferFromValue(const HloValue* value,
+                                const HloDataflowAnalysis& dataflow_analysis,
+                                bool live_out) {
     bool has_indirect_uses = false;
-    UsesList users = GetUsers(instruction_list_, logical_buffer,
-                              points_to_analysis, &has_indirect_uses);
-    return NewBuffer(instruction_list_.GetItem(logical_buffer->instruction()),
-                     logical_buffer->shape(), logical_buffer->index(),
-                     std::move(users), live_out, has_indirect_uses);
+    UsesList users = GetUsers(instruction_list_, value, dataflow_analysis,
+                              &has_indirect_uses);
+    return NewBuffer(instruction_list_.GetItem(value->instruction()),
+                     value->shape(), value->index(), std::move(users), live_out,
+                     has_indirect_uses);
   }
 
   // Creates a new buffer representing a rematerialization of given buffer for
@@ -697,44 +693,43 @@ class MemoryUsageTracker {
 MemoryUsageTracker::MemoryUsageTracker(
     const HloRematerialization::Options& options,
     const HloComputation* computation,
-    const TuplePointsToAnalysis& points_to_analysis,
+    const HloDataflowAnalysis& dataflow_analysis,
     const HloRematInstructionList& instruction_list)
     : options_(options),
       computation_(computation),
       instruction_list_(instruction_list) {
-  PointsToSet::BufferSet live_out_set =
-      points_to_analysis.GetPointsToSet(computation_->root_instruction())
-          .CreateFlattenedSet();
-  absl::flat_hash_map<const LogicalBuffer*, BufferId>
-      logical_buffer_to_buffer_id;
+  HloValueSet live_out_values =
+      dataflow_analysis.GetFlattenedValueSet(computation_->root_instruction());
+  absl::flat_hash_set<const HloValue*> live_out_set(
+      live_out_values.values().begin(), live_out_values.values().end());
+  absl::flat_hash_map<const HloValue*, BufferId> value_to_buffer_id;
   for (auto* item = instruction_list_.first(); item != nullptr;
        item = instruction_list_.next(item)) {
     const HloInstruction* const instruction = item->instruction;
-    for (const LogicalBuffer* logical_buffer :
-         points_to_analysis.GetBuffersDefinedByInstruction(instruction)) {
+    for (const auto& [index, value_set] :
+         dataflow_analysis.GetInstructionValueSet(instruction)) {
+      if (value_set.values().size() != 1 ||
+          value_set.GetUniqueValue().defining_instruction() != instruction) {
+        continue;
+      }
+      const HloValue* value = &value_set.GetUniqueValue();
       Buffer* buffer;
       if (instruction->opcode() == HloOpcode::kWhile) {
         // The while instruction defines no new buffers. Instead it reuses the
         // buffers of its operand. Find the Buffer of its operand at the
         // proper ShapeIndex.
-        const PointsToSet& operand_points_to =
-            points_to_analysis.GetPointsToSet(instruction->operand(0));
-        CHECK_EQ(operand_points_to.element(logical_buffer->index()).size(), 1);
-        const LogicalBuffer* source_logical_buffer =
-            operand_points_to.element(logical_buffer->index())[0];
-        buffer =
-            &buffers_.at(logical_buffer_to_buffer_id.at(source_logical_buffer));
+        const HloValue& source_value = dataflow_analysis.GetUniqueValueAt(
+            instruction->operand(0), value->index());
+        buffer = &buffers_.at(value_to_buffer_id.at(&source_value));
 
         // Mark buffer as has indirect use and live out.
         buffer->has_indirect_uses = true;
-        buffer->live_out =
-            buffer->live_out || ContainsKey(live_out_set, logical_buffer);
+        buffer->live_out = buffer->live_out || ContainsKey(live_out_set, value);
 
         // Add users of while to Buffer users.
         bool unused;
         for (HloRematItemUse& user_item :
-             GetUsers(instruction_list_, logical_buffer, points_to_analysis,
-                      &unused)) {
+             GetUsers(instruction_list_, value, dataflow_analysis, &unused)) {
           auto existing_user_it =
               absl::c_find_if(buffer->users, [&](const HloRematItemUse& use) {
                 return user_item.user == use.user;
@@ -746,9 +741,8 @@ MemoryUsageTracker::MemoryUsageTracker(
           }
         }
       } else {
-        buffer = &CreateBufferFromLogicalBuffer(
-            logical_buffer, points_to_analysis,
-            ContainsKey(live_out_set, logical_buffer));
+        buffer = &CreateBufferFromValue(value, dataflow_analysis,
+                                        ContainsKey(live_out_set, value));
         item->buffers_defined.push_back(buffer->id);
         for (HloRematItemUse& user : buffer->users) {
           if (!absl::c_linear_search(user.user->buffers_used, buffer->id)) {
@@ -757,15 +751,15 @@ MemoryUsageTracker::MemoryUsageTracker(
         }
       }
 
-      logical_buffer_to_buffer_id[logical_buffer] = buffer->id;
+      value_to_buffer_id[value] = buffer->id;
     }
 
     // Trace the output of each instruction. This is so that we can properly
     // track which outputs does GTEs have.
-    for (const LogicalBuffer* logical_buffer :
-         points_to_analysis.GetPointsToSet(instruction).CreateFlattenedSet()) {
-      item->buffers_output.push_back(
-          logical_buffer_to_buffer_id[logical_buffer]);
+    HloValueSet flattened_values =
+        dataflow_analysis.GetFlattenedValueSet(instruction);
+    for (const HloValue* value : flattened_values.values()) {
+      item->buffers_output.push_back(value_to_buffer_id.at(value));
     }
   }
   XLA_VLOG_LINES(10, ToString());
@@ -888,7 +882,7 @@ int64_t MemoryUsageTracker::MemoryReducedIfRematerialized(
     }
 
     // Compute the amount of memory reduced (if any) by rematerializing
-    // 'item->instruction'. The LogicalBuffers defined by 'item->instruction'
+    // 'item->instruction'. The HloValues defined by 'item->instruction'
     // will no longer be live at this program point, so initially set
     // memory_reduced to the size of its defined values.
     for (BufferId buffer_id : item->buffers_defined) {
@@ -2506,7 +2500,7 @@ HloRematerialization::ComputePeakMemoryAndInstruction(
     const HloComputation* computation, const HloInstructionSequence& order,
     const absl::flat_hash_set<absl::string_view>& execution_threads) const {
   HloRematInstructionList instruction_list(order);
-  MemoryUsageTracker tracker(options_, computation, *points_to_analysis_,
+  MemoryUsageTracker tracker(options_, computation, *dataflow_analysis_,
                              instruction_list);
   int64_t peak_memory = tracker.memory_usage();
   const HloInstruction* peak_instruction =
@@ -2562,10 +2556,10 @@ HloRematerialization::UpdateScheduleFromSequence(
   schedule->set_sequence(computation, sequence);
 
   // TODO(b/398843357): This is expensive and we shouldn't need to recompute
-  // the points_to_analysis_ after rematerializing each computation. Recompute
-  // points_to_analysis_ since the older analysis does not include
+  // the dataflow_analysis_ after rematerializing each computation. Recompute
+  // dataflow_analysis_ since the older analysis does not include
   // rematerialized instructions.
-  ABSL_RETURN_IF_ERROR(UpdatePointsToAnalysis(computation->parent()));
+  ABSL_RETURN_IF_ERROR(UpdateDataflowAnalysis(computation->parent()));
   ABSL_ASSIGN_OR_RETURN(
       MemoryUsageAndInstruction peak_memory_result,
       ComputePeakMemoryAndInstruction(
@@ -2575,15 +2569,22 @@ HloRematerialization::UpdateScheduleFromSequence(
   return peak_memory_result;
 }
 
-absl::Status HloRematerialization::UpdatePointsToAnalysis(HloModule* module) {
-  if (points_to_analysis_ == nullptr) {
-    ABSL_ASSIGN_OR_RETURN(points_to_analysis_, TuplePointsToAnalysis::Run(module));
+absl::Status HloRematerialization::UpdateDataflowAnalysis(HloModule* module) {
+  if (dataflow_analysis_ == nullptr) {
+    ABSL_ASSIGN_OR_RETURN(
+        dataflow_analysis_,
+        HloDataflowAnalysis::Run(*module, /*ssa_form=*/false,
+                                 /*bitcast_defines_value=*/false,
+                                 /*execution_threads=*/{},
+                                 /*propagate_through_calls=*/false,
+                                 /*precompute_uses=*/std::nullopt,
+                                 /*propagate_through_control_flow=*/false));
   }
   return absl::OkStatus();
 }
 
-void HloRematerialization::SetPointsToAnalysisStale() {
-  points_to_analysis_ = nullptr;
+void HloRematerialization::SetDataflowAnalysisStale() {
+  dataflow_analysis_ = nullptr;
 }
 
 absl::StatusOr<bool>
@@ -2695,8 +2696,8 @@ HloRematerialization::PeakPrioritySubPass(
   VLOG(3) << "Creating memory tracker for rematerialization on "
           << computation->name() << " instruction list size "
           << state.instruction_list->size();
-  ABSL_RETURN_IF_ERROR(UpdatePointsToAnalysis(computation->parent()));
-  MemoryUsageTracker memory_tracker(options_, computation, *points_to_analysis_,
+  ABSL_RETURN_IF_ERROR(UpdateDataflowAnalysis(computation->parent()));
+  MemoryUsageTracker memory_tracker(options_, computation, *dataflow_analysis_,
                                     *state.instruction_list);
   state.instruction_list->PromoteNodesToSkip([&](HloRematItem* item) {
     return memory_tracker.AllocatedSize(item) >= min_remat_size;
@@ -2858,7 +2859,7 @@ absl::StatusOr<bool> HloRematerialization::RematerializeComputationPeakPriority(
   int64_t cost_estimate_memory_limit_bytes =
       std::max(kMinimumCostEstimateMemoryLimitBytes, memory_limit_bytes);
 
-  ABSL_RETURN_IF_ERROR(UpdatePointsToAnalysis(computation->parent()));
+  ABSL_RETURN_IF_ERROR(UpdateDataflowAnalysis(computation->parent()));
   ABSL_ASSIGN_OR_RETURN(
       HloRematerialization::MemoryUsageAndInstruction peak_memory_result,
       ComputePeakMemoryAndInstruction(
@@ -2990,7 +2991,7 @@ absl::StatusOr<bool> HloRematerialization::RematerializeComputation(
   CHECK(!ContainsKey(rematerialized_computations_, computation));
 
   HloRematInstructionList instruction_list(schedule->sequence(computation));
-  MemoryUsageTracker memory_tracker(options_, computation, *points_to_analysis_,
+  MemoryUsageTracker memory_tracker(options_, computation, *dataflow_analysis_,
                                     instruction_list);
 
   instruction_list.PromoteNodesToSkip([&](HloRematItem* item) {
@@ -3238,10 +3239,10 @@ absl::StatusOr<bool> HloRematerialization::RunImpl(
   rematerialized_computations_.clear();
   instructions_rematerialized_ = 0;
   net_instructions_added_ = 0;
-  SetPointsToAnalysisStale();
+  SetDataflowAnalysisStale();
 
   TF_RET_CHECK(module->has_schedule());
-  ABSL_RETURN_IF_ERROR(UpdatePointsToAnalysis(module));
+  ABSL_RETURN_IF_ERROR(UpdateDataflowAnalysis(module));
   next_channel_id_ = hlo_query::NextChannelId(*module);
 
   // Adjust memory limit to account for the output of the entry
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.h b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.h
index a146d024213..d38b9e263ee 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.h
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.h
@@ -28,7 +28,7 @@
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/time/time.h"
-#include "xla/hlo/analysis/tuple_points_to_analysis.h"
+#include "xla/hlo/analysis/hlo_dataflow_analysis.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
@@ -226,10 +226,10 @@ class HloRematerialization : public HloPassInterface {
 
   int64_t GetBlockSizeLimit() const { return options_.block_size_limit; }
 
-  // Sets points to analysis to stale. Used by Peak Priority algorithm to
-  // indicate that the points to analysis should be updated before the next
+  // Sets dataflow analysis to stale. Used by Peak Priority algorithm to
+  // indicate that the dataflow analysis should be updated before the next
   // rematerialization subpass.
-  void SetPointsToAnalysisStale();
+  void SetDataflowAnalysisStale();
 
   // Holds references to data structures and some constants that are used during
   // rematerialization. This struct is used to avoid long function signatures.
@@ -283,10 +283,10 @@ class HloRematerialization : public HloPassInterface {
     return absl::OkStatus();
   }
 
-  // Only Peak priority requires constant update of points to analysis.
+  // Only Peak priority requires constant update of dataflow analysis.
   void on_block_rematerialized(int remat_count) {
     if (remat_count > 0 && remat_algorithm() == RematAlgorithm::kPeakPriority) {
-      SetPointsToAnalysisStale();
+      SetDataflowAnalysisStale();
     }
   }
 
@@ -300,8 +300,8 @@ class HloRematerialization : public HloPassInterface {
       const HloInstructionSequence& sequence,
       const absl::flat_hash_set<absl::string_view>& execution_threads);
 
-  // Updates points to analysis if it is stale.
-  absl::Status UpdatePointsToAnalysis(HloModule* module);
+  // Updates dataflow analysis if it is stale.
+  absl::Status UpdateDataflowAnalysis(HloModule* module);
 
   // Cleans up dead rematerialized instructions out of the module. Basically
   // runs DCE and updates the schedule.
@@ -402,7 +402,7 @@ class HloRematerialization : public HloPassInterface {
   // These values are updated as rematerialization occurs.
   absl::flat_hash_map<const HloComputation*, int64_t> computation_peak_memory_;
 
-  std::unique_ptr<TuplePointsToAnalysis> points_to_analysis_;
+  std::unique_ptr<HloDataflowAnalysis> dataflow_analysis_;
 
   // Set of computations which have had rematerialization
   // applied. Rematerialization is only applied once per computation.

From a3836e729f80592cb334fbd366960ef4ec3ee102 Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Tue, 15 Sep 2026 18:58:46 -0700
Subject: [PATCH 0427/1000] Validate axis bounds in experimental.numpy flip

NumPy raises AxisError when the axis passed to np.flip is out of bounds
for the input rank. tf.experimental.numpy.flip silently canonicalized
out-of-bounds axes (e.g. leaving negative entries in the reverse axis
list), producing a confusing backend failure. Add static bounds
validation when the rank and axes are statically known, consistent with
the validation added to moveaxis, compress, and cumsum.
---
 tensorflow/python/ops/numpy_ops/np_array_ops.py | 17 +++++++++++++++++
 .../python/ops/numpy_ops/np_array_ops_test.py   |  7 +++++++
 2 files changed, 24 insertions(+)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index 03791ba2e3a..14c13d26b3b 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -1684,6 +1684,23 @@ def flip(m, axis=None):  # pylint: disable=missing-docstring
   if np_utils.isscalar(axis):
     axis = [axis]
 
+  # NumPy raises AxisError for out-of-bounds axes instead of letting the
+  # backend kernel fail with a confusing error.
+  maybe_rank = m.shape.rank
+  if maybe_rank is not None and builtins.all(
+      isinstance(x, (int, np.integer)) for x in axis
+  ):
+    normalized_axes = []
+    for x in axis:
+      normalized = x + maybe_rank if x < 0 else x
+      if normalized < 0 or normalized >= maybe_rank:
+        raise ValueError(
+            f'Argument `axis` (received axis={x}) is out of bounds '
+            f'for input of rank {maybe_rank}.'
+        )
+      normalized_axes.append(normalized)
+    axis = normalized_axes
+
   axis = np_utils._canonicalize_axes(axis, array_ops.rank(m))  # pylint: disable=protected-access
 
   return array_ops.reverse(m, axis)
diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index dffb95e87b9..6c1f291253f 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1552,6 +1552,13 @@ class ArrayMethodsTest(test.TestCase):
     _test(a, axis=[0, 2])
     _test(a, axis=(0, 1, 2))
     _test(a, axis=range(3))
+    # Out-of-bounds axes raise, matching np.flip's AxisError.
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_array_ops.flip(a, axis=3)
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_array_ops.flip(a, axis=-4)
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_array_ops.flip(a, axis=(0, 3))
 
   def testNdim(self):
     self.assertAllEqual(0, np_array_ops.ndim(0.5))

From 4702e9366b6f016f92457c2d8cd968af5f7ef6bc Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Tue, 15 Sep 2026 19:23:30 -0700
Subject: [PATCH 0428/1000] Guard flip static axis validation on sequence types

The static bounds check iterates the axis argument, which would raise a
TypeError if a caller passed a non-sequence such as a symbolic Tensor in
graph mode. Restrict the static path to list, tuple, range, and ndarray
axes so anything else falls through to the existing dynamic
canonicalization.
---
 tensorflow/python/ops/numpy_ops/np_array_ops.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index 14c13d26b3b..ab0b61fc2e1 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -1687,8 +1687,10 @@ def flip(m, axis=None):  # pylint: disable=missing-docstring
   # NumPy raises AxisError for out-of-bounds axes instead of letting the
   # backend kernel fail with a confusing error.
   maybe_rank = m.shape.rank
-  if maybe_rank is not None and builtins.all(
-      isinstance(x, (int, np.integer)) for x in axis
+  if (
+      maybe_rank is not None
+      and isinstance(axis, (list, tuple, range, np.ndarray))
+      and builtins.all(isinstance(x, (int, np.integer)) for x in axis)
   ):
     normalized_axes = []
     for x in axis:

From 740a66772e679e7d71d600e6a8c959c53394e8ee Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Mon, 5 Oct 2026 21:04:03 -0700
Subject: [PATCH 0429/1000] Address review: correct motivation, skip redundant
 canonicalization, extend testFlip

---
 tensorflow/python/ops/numpy_ops/np_array_ops.py |  6 ++++--
 .../python/ops/numpy_ops/np_array_ops_test.py   | 17 +++++++++++++++++
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index ab0b61fc2e1..9b43fab2a82 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -1702,8 +1702,10 @@ def flip(m, axis=None):  # pylint: disable=missing-docstring
         )
       normalized_axes.append(normalized)
     axis = normalized_axes
-
-  axis = np_utils._canonicalize_axes(axis, array_ops.rank(m))  # pylint: disable=protected-access
+  else:
+    axis = np_utils._canonicalize_axes(
+        axis, array_ops.rank(m)
+    )  # pylint: disable=protected-access
 
   return array_ops.reverse(m, axis)
 
diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index 6c1f291253f..4bc65bb3afa 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1559,6 +1559,23 @@ class ArrayMethodsTest(test.TestCase):
       np_array_ops.flip(a, axis=-4)
     with self.assertRaisesRegex(ValueError, 'out of bounds'):
       np_array_ops.flip(a, axis=(0, 3))
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_array_ops.flip(a, axis=[0, -4])
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_array_ops.flip(a, axis=np.array([3]))
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_array_ops.flip(a, axis=range(4))
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_array_ops.fliplr(np.arange(3))
+    # The check runs at trace time when the rank is static.
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      def_function.function(lambda x: np_array_ops.flip(x, axis=-4))(a)
+    # Unknown-rank inputs keep the previous behavior for valid axes.
+    flip_unknown_rank = def_function.function(
+        lambda x: np_array_ops.flip(x, axis=-1),
+        input_signature=[tensor_spec.TensorSpec(None, dtypes.float64)],
+    )
+    self.assertAllEqual(np.flip(a, -1), flip_unknown_rank(a))
 
   def testNdim(self):
     self.assertAllEqual(0, np_array_ops.ndim(0.5))

From 08a5d7634aaf8012d022833227498febc6d1cf9d Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Mon, 5 Oct 2026 22:28:45 -0700
Subject: [PATCH 0430/1000] Add exact negative-boundary and zero-rank axis
 tests for flip

---
 tensorflow/python/ops/numpy_ops/np_array_ops_test.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index 4bc65bb3afa..397eab22a71 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1546,6 +1546,7 @@ class ArrayMethodsTest(test.TestCase):
     _test(a, axis=0)
     _test(a, axis=2)
     _test(a, axis=-1)
+    _test(a, axis=-3)
     # A tuple, list or range of axes, including negative values.
     _test(a, axis=(0, 1))
     _test(a, axis=(-1, -3))
@@ -1553,6 +1554,8 @@ class ArrayMethodsTest(test.TestCase):
     _test(a, axis=(0, 1, 2))
     _test(a, axis=range(3))
     # Out-of-bounds axes raise, matching np.flip's AxisError.
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_array_ops.flip(np.float64(1.0), axis=0)
     with self.assertRaisesRegex(ValueError, 'out of bounds'):
       np_array_ops.flip(a, axis=3)
     with self.assertRaisesRegex(ValueError, 'out of bounds'):

From cee9a47f2d1a27465a92ace6a4ca6cd317c9c0af Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Tue, 6 Oct 2026 15:04:33 -0700
Subject: [PATCH 0431/1000] Delegate flip axis validation to the ReverseV2
 kernel

---
 .../python/ops/numpy_ops/np_array_ops.py      | 27 ++----------
 .../python/ops/numpy_ops/np_array_ops_test.py | 43 ++++++++++---------
 2 files changed, 26 insertions(+), 44 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index 9b43fab2a82..2fbfb7ae345 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -1681,31 +1681,10 @@ def flip(m, axis=None):  # pylint: disable=missing-docstring
   if axis is None:
     return array_ops.reverse(m, math_ops.range(array_ops.rank(m)))
 
-  if np_utils.isscalar(axis):
-    axis = [axis]
-
-  # NumPy raises AxisError for out-of-bounds axes instead of letting the
-  # backend kernel fail with a confusing error.
-  maybe_rank = m.shape.rank
-  if (
-      maybe_rank is not None
-      and isinstance(axis, (list, tuple, range, np.ndarray))
-      and builtins.all(isinstance(x, (int, np.integer)) for x in axis)
+  if np_utils.isscalar(axis) or (
+      isinstance(axis, np.ndarray) and axis.ndim == 0
   ):
-    normalized_axes = []
-    for x in axis:
-      normalized = x + maybe_rank if x < 0 else x
-      if normalized < 0 or normalized >= maybe_rank:
-        raise ValueError(
-            f'Argument `axis` (received axis={x}) is out of bounds '
-            f'for input of rank {maybe_rank}.'
-        )
-      normalized_axes.append(normalized)
-    axis = normalized_axes
-  else:
-    axis = np_utils._canonicalize_axes(
-        axis, array_ops.rank(m)
-    )  # pylint: disable=protected-access
+    axis = [axis]
 
   return array_ops.reverse(m, axis)
 
diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index 397eab22a71..621da6c58cd 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1553,26 +1553,29 @@ class ArrayMethodsTest(test.TestCase):
     _test(a, axis=[0, 2])
     _test(a, axis=(0, 1, 2))
     _test(a, axis=range(3))
-    # Out-of-bounds axes raise, matching np.flip's AxisError.
-    with self.assertRaisesRegex(ValueError, 'out of bounds'):
-      np_array_ops.flip(np.float64(1.0), axis=0)
-    with self.assertRaisesRegex(ValueError, 'out of bounds'):
-      np_array_ops.flip(a, axis=3)
-    with self.assertRaisesRegex(ValueError, 'out of bounds'):
-      np_array_ops.flip(a, axis=-4)
-    with self.assertRaisesRegex(ValueError, 'out of bounds'):
-      np_array_ops.flip(a, axis=(0, 3))
-    with self.assertRaisesRegex(ValueError, 'out of bounds'):
-      np_array_ops.flip(a, axis=[0, -4])
-    with self.assertRaisesRegex(ValueError, 'out of bounds'):
-      np_array_ops.flip(a, axis=np.array([3]))
-    with self.assertRaisesRegex(ValueError, 'out of bounds'):
-      np_array_ops.flip(a, axis=range(4))
-    with self.assertRaisesRegex(ValueError, 'out of bounds'):
-      np_array_ops.fliplr(np.arange(3))
-    # The check runs at trace time when the rank is static.
-    with self.assertRaisesRegex(ValueError, 'out of bounds'):
-      def_function.function(lambda x: np_array_ops.flip(x, axis=-4))(a)
+    # Out-of-bounds axes raise InvalidArgumentError from the C++ kernel.
+    kernel_err = r'out of (valid )?range'
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
+      self.evaluate(np_array_ops.flip(np.float64(1.0), axis=0))
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
+      self.evaluate(np_array_ops.flip(a, axis=3))
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
+      self.evaluate(np_array_ops.flip(a, axis=-4))
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
+      self.evaluate(np_array_ops.flip(a, axis=(0, 3)))
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
+      self.evaluate(np_array_ops.flip(a, axis=[0, -4]))
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
+      self.evaluate(np_array_ops.flip(a, axis=np.array([3])))
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
+      self.evaluate(np_array_ops.flip(a, axis=range(4)))
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
+      self.evaluate(np_array_ops.fliplr(np.arange(3)))
+    # The C++ kernel raises at execution time for traced functions too.
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
+      self.evaluate(
+          def_function.function(lambda x: np_array_ops.flip(x, axis=-4))(a)
+      )
     # Unknown-rank inputs keep the previous behavior for valid axes.
     flip_unknown_rank = def_function.function(
         lambda x: np_array_ops.flip(x, axis=-1),

From 460830adedf80542130e7cb76e544729b9726e1c Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Wed, 7 Oct 2026 07:13:31 -0700
Subject: [PATCH 0432/1000] Add 0-D/1-D ndarray axis and unknown-rank OOB tests
 for flip

---
 .../python/ops/numpy_ops/np_array_ops_test.py       | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index 621da6c58cd..7c3a1dc70f2 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1547,6 +1547,9 @@ class ArrayMethodsTest(test.TestCase):
     _test(a, axis=2)
     _test(a, axis=-1)
     _test(a, axis=-3)
+    # 0-D and 1-D ndarray axes.
+    _test(a, axis=np.array(1))
+    _test(a, axis=np.array([0, 2]))
     # A tuple, list or range of axes, including negative values.
     _test(a, axis=(0, 1))
     _test(a, axis=(-1, -3))
@@ -1565,6 +1568,8 @@ class ArrayMethodsTest(test.TestCase):
       self.evaluate(np_array_ops.flip(a, axis=(0, 3)))
     with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
       self.evaluate(np_array_ops.flip(a, axis=[0, -4]))
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
+      self.evaluate(np_array_ops.flip(a, axis=np.array(3)))
     with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
       self.evaluate(np_array_ops.flip(a, axis=np.array([3])))
     with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
@@ -1583,6 +1588,14 @@ class ArrayMethodsTest(test.TestCase):
     )
     self.assertAllEqual(np.flip(a, -1), flip_unknown_rank(a))
 
+    # Unknown-rank inputs raise InvalidArgumentError for out-of-bounds axes.
+    flip_unknown_rank_invalid = def_function.function(
+        lambda x: np_array_ops.flip(x, axis=-4),
+        input_signature=[tensor_spec.TensorSpec(None, dtypes.float64)],
+    )
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
+      self.evaluate(flip_unknown_rank_invalid(a))
+
   def testNdim(self):
     self.assertAllEqual(0, np_array_ops.ndim(0.5))
     self.assertAllEqual(1, np_array_ops.ndim([1, 2]))

From 73d5c1d9ad2179e4d2d71460ee6f957096acbf39 Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Wed, 7 Oct 2026 15:21:02 -0700
Subject: [PATCH 0433/1000] Add dynamic-rank axis bounds check to rot90 and
 broaden flip test error matching

---
 tensorflow/python/ops/numpy_ops/np_array_ops.py      | 11 ++++++++++-
 tensorflow/python/ops/numpy_ops/np_array_ops_test.py | 11 ++++++++---
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index 2fbfb7ae345..882cb3b293a 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -1754,7 +1754,16 @@ def rot90(m, k=1, axes=(0, 1)):  # pylint: disable=missing-docstring
           )
 
   m_rank = array_ops.rank(m)
-  ax1, ax2 = np_utils._canonicalize_axes(axes, m_rank)  # pylint: disable=protected-access
+  if maybe_rank is None:
+    rank_t = ops.convert_to_tensor(m_rank)
+    axes_t = ops.convert_to_tensor(axes)
+    control_flow_assert.Assert(
+        math_ops.reduce_all(math_ops.logical_and(
+            axes_t >= -rank_t, axes_t < rank_t)),
+        ['Axes', axes_t, 'out of range for array of ndim', rank_t],
+    )
+  ax1, ax2 = np_utils._canonicalize_axes(
+      axes, m_rank)  # pylint: disable=protected-access
 
   k = k % 4
   if k == 0:
diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index 7c3a1dc70f2..530497a185f 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1557,7 +1557,10 @@ class ArrayMethodsTest(test.TestCase):
     _test(a, axis=(0, 1, 2))
     _test(a, axis=range(3))
     # Out-of-bounds axes raise InvalidArgumentError from the C++ kernel.
-    kernel_err = r'out of (valid )?range'
+    kernel_err = (
+        r'(out of (valid )?range'
+        r'|can not have more elements than input tensor has dimensions)'
+    )
     with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
       self.evaluate(np_array_ops.flip(np.float64(1.0), axis=0))
     with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
@@ -1576,8 +1579,10 @@ class ArrayMethodsTest(test.TestCase):
       self.evaluate(np_array_ops.flip(a, axis=range(4)))
     with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
       self.evaluate(np_array_ops.fliplr(np.arange(3)))
-    # The C++ kernel raises at execution time for traced functions too.
-    with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
+    # The C++ kernel raises at execution time (or shape inference at trace time).
+    with self.assertRaisesRegex(
+        (errors_impl.InvalidArgumentError, ValueError), kernel_err
+    ):
       self.evaluate(
           def_function.function(lambda x: np_array_ops.flip(x, axis=-4))(a)
       )

From 1db0b4f505a6cf919e7e930a9b28fa7aa90fa746 Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Wed, 7 Oct 2026 19:14:36 -0700
Subject: [PATCH 0434/1000] Wire rot90 dynamic-rank asserts into the graph and
 add coverage

Connect the rot90 dynamic-rank Assert operations via
ops.control_dependencies so they are not pruned as disconnected nodes
in graph mode, add a duplicate-axis Assert on the canonicalized axes
(mirroring the static path), and add graph-mode tests exercising both
the out-of-bounds and duplicate-axis checks on unknown-rank inputs.
---
 .../python/ops/numpy_ops/np_array_ops.py      | 28 ++++++++++----
 .../python/ops/numpy_ops/np_array_ops_test.py | 38 +++++++++++++++++++
 2 files changed, 58 insertions(+), 8 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index 882cb3b293a..19c56415a71 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -1754,16 +1754,28 @@ def rot90(m, k=1, axes=(0, 1)):  # pylint: disable=missing-docstring
           )
 
   m_rank = array_ops.rank(m)
-  if maybe_rank is None:
-    rank_t = ops.convert_to_tensor(m_rank)
-    axes_t = ops.convert_to_tensor(axes)
-    control_flow_assert.Assert(
-        math_ops.reduce_all(math_ops.logical_and(
-            axes_t >= -rank_t, axes_t < rank_t)),
-        ['Axes', axes_t, 'out of range for array of ndim', rank_t],
-    )
   ax1, ax2 = np_utils._canonicalize_axes(
       axes, m_rank)  # pylint: disable=protected-access
+  if maybe_rank is None:
+    # In graph mode, disconnected side-effectful Assert nodes are pruned
+    # unless their results flow into the returned computation, so the
+    # assertions must be wired in via `control_dependencies`.
+    rank_t = ops.convert_to_tensor(m_rank)
+    axes_t = ops.convert_to_tensor(axes)
+    assert_bounds = control_flow_assert.Assert(
+        math_ops.reduce_all(
+            math_ops.logical_and(axes_t >= -rank_t, axes_t < rank_t)),
+        ['Axes', axes_t, 'out of range for array of ndim', rank_t],
+    )
+    # The canonicalized axes are tensors here, so the duplicate-axis check
+    # the static path performs natively must be re-done dynamically.
+    assert_diff = control_flow_assert.Assert(
+        math_ops.not_equal(ax1, ax2),
+        ['Axes must be different']
+    )
+    with ops.control_dependencies([assert_bounds, assert_diff]):
+      ax1 = array_ops.identity(ax1)
+      ax2 = array_ops.identity(ax2)
 
   k = k % 4
   if k == 0:
diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index 530497a185f..57b71482fee 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1601,6 +1601,44 @@ class ArrayMethodsTest(test.TestCase):
     with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
       self.evaluate(flip_unknown_rank_invalid(a))
 
+  def testRot90OutOfBoundsDynamicRank(self):
+    # The dynamic-rank bounds Assert in `rot90` must actually run in graph
+    # mode; it is wired into the computation via `control_dependencies`.
+    # Note: this is deliberately not run with `jit_compile=True` — tf2xla
+    # lowers `Assert` to a no-op under XLA compilation.
+    x = np.zeros((2, 3), dtype=np.float32)
+
+    @def_function.function(
+        input_signature=[
+            tensor_spec.TensorSpec(dtype=dtypes.float32, shape=None)
+        ]
+    )
+    def f(a):
+      return np_array_ops.rot90(a, axes=(0, -5))
+
+    with self.assertRaisesRegex(
+        errors_impl.InvalidArgumentError, 'out of range for array of ndim'
+    ):
+      f(x)
+
+  def testRot90DuplicateAxesDynamicRank(self):
+    # The dynamic path re-checks for duplicate axes after canonicalization
+    # (e.g. (1, -2) both canonicalize to axis 1 on a rank-3 array).
+    x = np.zeros((2, 3), dtype=np.float32)
+
+    @def_function.function(
+        input_signature=[
+            tensor_spec.TensorSpec(dtype=dtypes.float32, shape=None)
+        ]
+    )
+    def f(a):
+      return np_array_ops.rot90(a, axes=(1, -2))
+
+    with self.assertRaisesRegex(
+        errors_impl.InvalidArgumentError, 'Axes must be different'
+    ):
+      f(x)
+
   def testNdim(self):
     self.assertAllEqual(0, np_array_ops.ndim(0.5))
     self.assertAllEqual(1, np_array_ops.ndim([1, 2]))

From 2fb99dfa01b489f6e7f7d007cb7b9a8b882fd1db Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Fri, 9 Oct 2026 20:02:38 -0700
Subject: [PATCH 0435/1000] Harden rot90 dynamic checks and complete
 dynamic-rank test coverage

Thread the input tensor through the control dependencies so the
asserts still execute on the k % 4 == 0 early-return path, force the
axes tensor to the rank tensor dtype to avoid a TypeError when an
axis exceeds the int32 range, wrap the dynamic-rank exception tests
in self.evaluate so the asserts actually run, and add positive
out-of-bounds cases for flip and rot90 on unknown-rank inputs.
---
 .../python/ops/numpy_ops/np_array_ops.py      | 10 ++++-
 .../python/ops/numpy_ops/np_array_ops_test.py | 42 +++++++++++++++----
 2 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index 19c56415a71..03304dc0cbc 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -1761,7 +1761,11 @@ def rot90(m, k=1, axes=(0, 1)):  # pylint: disable=missing-docstring
     # unless their results flow into the returned computation, so the
     # assertions must be wired in via `control_dependencies`.
     rank_t = ops.convert_to_tensor(m_rank)
-    axes_t = ops.convert_to_tensor(axes)
+    # Force the axes tensor to the rank tensor's dtype: left to inference, a
+    # tuple containing a value above the int32 range would produce an int64
+    # tensor, and the comparison against `rank_t` would raise a TypeError
+    # instead of triggering the assertions below.
+    axes_t = ops.convert_to_tensor(axes, dtype=rank_t.dtype)
     assert_bounds = control_flow_assert.Assert(
         math_ops.reduce_all(
             math_ops.logical_and(axes_t >= -rank_t, axes_t < rank_t)),
@@ -1776,6 +1780,10 @@ def rot90(m, k=1, axes=(0, 1)):  # pylint: disable=missing-docstring
     with ops.control_dependencies([assert_bounds, assert_diff]):
       ax1 = array_ops.identity(ax1)
       ax2 = array_ops.identity(ax2)
+      # Also thread `m` through the control dependencies: otherwise the
+      # `k % 4 == 0` early return below would return `m` without consuming
+      # the asserts, and graph pruning would disconnect them again.
+      m = array_ops.identity(m)
 
   k = k % 4
   if k == 0:
diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index 57b71482fee..9af6fe7f74b 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1499,7 +1499,7 @@ class ArrayMethodsTest(test.TestCase):
       return np_array_ops.swapaxes(a, -10, 0)
 
     with self.assertRaises(errors_impl.InvalidArgumentError):
-      f(x)
+      self.evaluate(f(x))
 
   def testMoveaxis(self):
 
@@ -1593,13 +1593,21 @@ class ArrayMethodsTest(test.TestCase):
     )
     self.assertAllEqual(np.flip(a, -1), flip_unknown_rank(a))
 
-    # Unknown-rank inputs raise InvalidArgumentError for out-of-bounds axes.
-    flip_unknown_rank_invalid = def_function.function(
+    # Unknown-rank inputs raise InvalidArgumentError for out-of-bounds axes,
+    # both negative and positive.
+    flip_unknown_rank_invalid_neg = def_function.function(
         lambda x: np_array_ops.flip(x, axis=-4),
         input_signature=[tensor_spec.TensorSpec(None, dtypes.float64)],
     )
     with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
-      self.evaluate(flip_unknown_rank_invalid(a))
+      self.evaluate(flip_unknown_rank_invalid_neg(a))
+
+    flip_unknown_rank_invalid_pos = def_function.function(
+        lambda x: np_array_ops.flip(x, axis=3),
+        input_signature=[tensor_spec.TensorSpec(None, dtypes.float64)],
+    )
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
+      self.evaluate(flip_unknown_rank_invalid_pos(a))
 
   def testRot90OutOfBoundsDynamicRank(self):
     # The dynamic-rank bounds Assert in `rot90` must actually run in graph
@@ -1613,13 +1621,33 @@ class ArrayMethodsTest(test.TestCase):
             tensor_spec.TensorSpec(dtype=dtypes.float32, shape=None)
         ]
     )
-    def f(a):
+    def f_neg(a):
       return np_array_ops.rot90(a, axes=(0, -5))
 
+    f_neg = def_function.function(
+        f_neg,
+        input_signature=[
+            tensor_spec.TensorSpec(dtype=dtypes.float32, shape=None)
+        ]
+    )
     with self.assertRaisesRegex(
         errors_impl.InvalidArgumentError, 'out of range for array of ndim'
     ):
-      f(x)
+      self.evaluate(f_neg(x))
+
+    def f_pos(a):
+      return np_array_ops.rot90(a, axes=(0, 5))
+
+    f_pos = def_function.function(
+        f_pos,
+        input_signature=[
+            tensor_spec.TensorSpec(dtype=dtypes.float32, shape=None)
+        ]
+    )
+    with self.assertRaisesRegex(
+        errors_impl.InvalidArgumentError, 'out of range for array of ndim'
+    ):
+      self.evaluate(f_pos(x))
 
   def testRot90DuplicateAxesDynamicRank(self):
     # The dynamic path re-checks for duplicate axes after canonicalization
@@ -1637,7 +1665,7 @@ class ArrayMethodsTest(test.TestCase):
     with self.assertRaisesRegex(
         errors_impl.InvalidArgumentError, 'Axes must be different'
     ):
-      f(x)
+      self.evaluate(f(x))
 
   def testNdim(self):
     self.assertAllEqual(0, np_array_ops.ndim(0.5))

From b64721d5a5b6c6f8173ae0caf7fcae74aeb9b190 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 13:32:59 -0700
Subject: [PATCH 0436/1000] Automated Code Change

PiperOrigin-RevId: 997178958
---
 tensorflow/core/kernels/ragged_tensor_variant.h    |  2 +-
 tensorflow/core/kernels/random_shuffle_queue_op.cc |  2 +-
 tensorflow/core/kernels/range_sampler.cc           |  2 +-
 tensorflow/core/kernels/range_sampler.h            | 10 +++++-----
 tensorflow/core/kernels/reduction_gpu_kernels.cu.h |  4 ++--
 tensorflow/core/kernels/redux_functor.h            |  4 ++--
 tensorflow/core/kernels/regex_replace_op.cc        |  2 +-
 tensorflow/core/kernels/sdca_internal.h            |  6 +++---
 8 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/tensorflow/core/kernels/ragged_tensor_variant.h b/tensorflow/core/kernels/ragged_tensor_variant.h
index 16393dbdbe9..be1e83a8738 100644
--- a/tensorflow/core/kernels/ragged_tensor_variant.h
+++ b/tensorflow/core/kernels/ragged_tensor_variant.h
@@ -36,7 +36,7 @@ namespace tensorflow {
 // Class used to store a RaggedTensor as a Variant scalar.
 class RaggedTensorVariant {
  public:
-  RaggedTensorVariant() {}
+  RaggedTensorVariant() = default;
   RaggedTensorVariant(Tensor values, const std::vector<Tensor>& nested_splits)
       : values_(std::move(values)), nested_splits_(nested_splits) {}
 
diff --git a/tensorflow/core/kernels/random_shuffle_queue_op.cc b/tensorflow/core/kernels/random_shuffle_queue_op.cc
index e55f945ea06..509f77846b0 100644
--- a/tensorflow/core/kernels/random_shuffle_queue_op.cc
+++ b/tensorflow/core/kernels/random_shuffle_queue_op.cc
@@ -67,7 +67,7 @@ class RandomShuffleQueue : public TypedQueue<std::vector<Tensor> > {
   }
 
  private:
-  ~RandomShuffleQueue() override {}
+  ~RandomShuffleQueue() override = default;
 
   // Helper for dequeuing a single random element from queues_.
   void DequeueLocked(OpKernelContext* ctx, Tuple* tuple)
diff --git a/tensorflow/core/kernels/range_sampler.cc b/tensorflow/core/kernels/range_sampler.cc
index fe3d5a56a32..3489528df15 100644
--- a/tensorflow/core/kernels/range_sampler.cc
+++ b/tensorflow/core/kernels/range_sampler.cc
@@ -37,7 +37,7 @@ namespace tensorflow {
 using gtl::ArraySlice;
 using gtl::MutableArraySlice;
 
-RangeSampler::~RangeSampler() {}
+RangeSampler::~RangeSampler() = default;
 
 void RangeSampler::SampleBatch(random::SimplePhilox* rnd, bool unique,
                                absl::Span<int64_t> batch) const {
diff --git a/tensorflow/core/kernels/range_sampler.h b/tensorflow/core/kernels/range_sampler.h
index cecb681cd4e..62515ac13a4 100644
--- a/tensorflow/core/kernels/range_sampler.h
+++ b/tensorflow/core/kernels/range_sampler.h
@@ -109,7 +109,7 @@ class AllSampler : public RangeSampler {
  public:
   explicit AllSampler(int64_t range);
 
-  ~AllSampler() override {}
+  ~AllSampler() override = default;
 
   int64_t Sample(random::SimplePhilox* rnd) const override {
     LOG(FATAL) << "Should not be called";
@@ -132,7 +132,7 @@ class UniformSampler : public RangeSampler {
  public:
   explicit UniformSampler(int64_t range);
 
-  ~UniformSampler() override {}
+  ~UniformSampler() override = default;
 
   int64_t Sample(random::SimplePhilox* rnd) const override;
 
@@ -146,7 +146,7 @@ class LogUniformSampler : public RangeSampler {
  public:
   explicit LogUniformSampler(int64_t range);
 
-  ~LogUniformSampler() override {}
+  ~LogUniformSampler() override = default;
 
   int64_t Sample(random::SimplePhilox* rnd) const override;
 
@@ -160,7 +160,7 @@ class LogUniformSampler : public RangeSampler {
 class ThreadUnsafeUnigramSampler : public RangeSampler {
  public:
   explicit ThreadUnsafeUnigramSampler(int64_t range);
-  ~ThreadUnsafeUnigramSampler() override {}
+  ~ThreadUnsafeUnigramSampler() override = default;
 
   int64_t Sample(random::SimplePhilox* rnd) const override;
 
@@ -177,7 +177,7 @@ class ThreadUnsafeUnigramSampler : public RangeSampler {
 class UnigramSampler : public RangeSampler {
  public:
   explicit UnigramSampler(int64_t range);
-  ~UnigramSampler() override {}
+  ~UnigramSampler() override = default;
 
   int64_t Sample(random::SimplePhilox* rnd) const override;
 
diff --git a/tensorflow/core/kernels/reduction_gpu_kernels.cu.h b/tensorflow/core/kernels/reduction_gpu_kernels.cu.h
index a82e6c47061..ab99de962e4 100644
--- a/tensorflow/core/kernels/reduction_gpu_kernels.cu.h
+++ b/tensorflow/core/kernels/reduction_gpu_kernels.cu.h
@@ -241,7 +241,7 @@ __global__ __launch_bounds__(1024) void RowReduceKernel(
 template <typename T1>
 struct storage_type {
   T1 val;
-  __host__ __device__ storage_type() {}
+  __host__ __device__ storage_type() = default;
   __host__ __device__ operator T1() { return val; }
   __host__ __device__ storage_type<T1>& operator=(const T1& in) {
     val = in;
@@ -253,7 +253,7 @@ template <typename T2>
 struct storage_type<std::complex<T2>> {
   T2 real;
   T2 imag;
-  __host__ __device__ storage_type() {}
+  __host__ __device__ storage_type() = default;
   __host__ __device__ operator std::complex<T2>() {
     return std::complex<T2>(real, imag);
   }
diff --git a/tensorflow/core/kernels/redux_functor.h b/tensorflow/core/kernels/redux_functor.h
index 41ab917a746..788d00d166a 100644
--- a/tensorflow/core/kernels/redux_functor.h
+++ b/tensorflow/core/kernels/redux_functor.h
@@ -38,7 +38,7 @@ namespace functor {
 template <typename InputT, typename AccumT, typename OutputT,
           typename BinaryFunctor>
 struct ReduceOuterDimensions {
-  ReduceOuterDimensions() {}
+  ReduceOuterDimensions() = default;
 
   template <int num_dims>
   void operator()(const CPUDevice& device,
@@ -206,7 +206,7 @@ struct ReduceOuterDimensions {
 template <typename InputT, typename AccumT, typename OutputT,
           typename BinaryFunctor, typename Reducer>
 struct ReduceMiddleDimensions {
-  ReduceMiddleDimensions() {}
+  ReduceMiddleDimensions() = default;
 
   template <int num_dims>
   void operator()(const CPUDevice& device,
diff --git a/tensorflow/core/kernels/regex_replace_op.cc b/tensorflow/core/kernels/regex_replace_op.cc
index 2428ecf306e..b85271b6449 100644
--- a/tensorflow/core/kernels/regex_replace_op.cc
+++ b/tensorflow/core/kernels/regex_replace_op.cc
@@ -69,7 +69,7 @@ class RegexReplaceOp : public OpKernel {
     OP_REQUIRES_OK(ctx, ctx->GetAttr("replace_global", &replace_global_));
   }
 
-  ~RegexReplaceOp() override {}
+  ~RegexReplaceOp() override = default;
 
   void Compute(OpKernelContext* ctx) override {
     const Tensor* pattern_tensor;
diff --git a/tensorflow/core/kernels/sdca_internal.h b/tensorflow/core/kernels/sdca_internal.h
index 8f5ac0384ed..ff87389d9a3 100644
--- a/tensorflow/core/kernels/sdca_internal.h
+++ b/tensorflow/core/kernels/sdca_internal.h
@@ -73,7 +73,7 @@ struct ExampleStatistics {
 
 class Regularizations {
  public:
-  Regularizations() {}
+  Regularizations() = default;
 
   // Initialize() must be called immediately after construction.
   absl::Status Initialize(OpKernelConstruction* const context) {
@@ -278,7 +278,7 @@ class FeatureWeightsSparseStorage {
 // for both sparse and dense features.
 class ModelWeights {
  public:
-  ModelWeights() {}
+  ModelWeights() = default;
 
   bool SparseIndexValid(const int col, const int64_t index) const {
     return sparse_weights_[col].IndexValid(index);
@@ -315,7 +315,7 @@ class ModelWeights {
 // Examples contains all the training examples that SDCA uses for a mini-batch.
 class Examples {
  public:
-  Examples() {}
+  Examples() = default;
 
   // Returns the Example at |example_index|.
   const Example& example(const int example_index) const {

From 9b608867a7bdb282e8f9991e3c82790073d38bf3 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 13:40:53 -0700
Subject: [PATCH 0437/1000] Automated Code Change

PiperOrigin-RevId: 997181383
---
 tensorflow/lite/testing/generated_examples_zip_test.cc | 2 +-
 tensorflow/lite/testing/message.h                      | 4 ++--
 tensorflow/lite/testing/message_test.cc                | 2 +-
 tensorflow/lite/testing/test_runner.h                  | 4 ++--
 tensorflow/lite/testing/tf_driver.h                    | 2 +-
 tensorflow/lite/testing/tokenize.h                     | 2 +-
 6 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/tensorflow/lite/testing/generated_examples_zip_test.cc b/tensorflow/lite/testing/generated_examples_zip_test.cc
index ada9e7a0fdb..894211a8e75 100644
--- a/tensorflow/lite/testing/generated_examples_zip_test.cc
+++ b/tensorflow/lite/testing/generated_examples_zip_test.cc
@@ -134,7 +134,7 @@ const std::map<string, string>& GetKnownQuantizeBrokenTests() {
 // sure those temporary directories are removed later.
 class ArchiveEnvironment : public ::testing::Environment {
  public:
-  ~ArchiveEnvironment() override {}
+  ~ArchiveEnvironment() override = default;
 
   // Delete all temporary directories on teardown.
   void TearDown() override {
diff --git a/tensorflow/lite/testing/message.h b/tensorflow/lite/testing/message.h
index e6566ab11ca..b65650fd1a4 100644
--- a/tensorflow/lite/testing/message.h
+++ b/tensorflow/lite/testing/message.h
@@ -38,8 +38,8 @@ class Message {
   // top-level message. Returns true if the parsing succeeded.
   static bool Read(std::istream* input, Message* message);
 
-  Message() {}
-  virtual ~Message() {}
+  Message() = default;
+  virtual ~Message() = default;
 
   // Called when a new field is found. For example, when:
   //   f : "values"
diff --git a/tensorflow/lite/testing/message_test.cc b/tensorflow/lite/testing/message_test.cc
index f7e8e2cc993..378b0dc4b45 100644
--- a/tensorflow/lite/testing/message_test.cc
+++ b/tensorflow/lite/testing/message_test.cc
@@ -27,7 +27,7 @@ namespace {
 // A hierarchical, key-value store.
 class TestMessage : public Message {
  public:
-  TestMessage() {}
+  TestMessage() = default;
   explicit TestMessage(const std::string& text_to_parse) {
     std::stringstream ss(text_to_parse);
     finished_ = Message::Read(&ss, this);
diff --git a/tensorflow/lite/testing/test_runner.h b/tensorflow/lite/testing/test_runner.h
index 60c79ce02b3..de634d7e263 100644
--- a/tensorflow/lite/testing/test_runner.h
+++ b/tensorflow/lite/testing/test_runner.h
@@ -30,8 +30,8 @@ namespace testing {
 // (e.g. TF Lite's interpreter, or the NNAPI).
 class TestRunner {
  public:
-  TestRunner() {}
-  virtual ~TestRunner() {}
+  TestRunner() = default;
+  virtual ~TestRunner() = default;
 
   // Loads the given model, as a path relative to SetModelBaseDir().
   // DEPRECATED: use LoadModel with signature instead.
diff --git a/tensorflow/lite/testing/tf_driver.h b/tensorflow/lite/testing/tf_driver.h
index f25db13e0f0..fb0bbf1a385 100644
--- a/tensorflow/lite/testing/tf_driver.h
+++ b/tensorflow/lite/testing/tf_driver.h
@@ -41,7 +41,7 @@ class TfDriver : public TestRunner {
                     const std::vector<string>& input_layer_type,
                     const std::vector<string>& input_layer_shape,
                     const std::vector<string>& output_layer);
-  ~TfDriver() override {}
+  ~TfDriver() override = default;
 
   void LoadModel(const string& bin_file_path) override;
   void LoadModel(const string& bin_file_path, const string&) override {
diff --git a/tensorflow/lite/testing/tokenize.h b/tensorflow/lite/testing/tokenize.h
index 7bd2783337a..033e5057c45 100644
--- a/tensorflow/lite/testing/tokenize.h
+++ b/tensorflow/lite/testing/tokenize.h
@@ -24,7 +24,7 @@ namespace testing {
 // Process tokens coming from Tokenize().
 class TokenProcessor {
  public:
-  virtual ~TokenProcessor() {}
+  virtual ~TokenProcessor() = default;
   // Process a single token. The token won't be reused, so it is OK to call
   // token.swap().
   virtual void ConsumeToken(std::string* token) = 0;

From 7964703bc92f439fd64c397fb4580e74b36eb0d9 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 13:41:38 -0700
Subject: [PATCH 0438/1000] Automated Code Change

PiperOrigin-RevId: 997181590
---
 tensorflow/core/debug/debug_callback_registry.cc | 2 +-
 tensorflow/core/debug/debug_io_utils.h           | 2 +-
 tensorflow/core/debug/debugger_state_impl.h      | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/debug/debug_callback_registry.cc b/tensorflow/core/debug/debug_callback_registry.cc
index 5ee0d53d507..93f03f90548 100644
--- a/tensorflow/core/debug/debug_callback_registry.cc
+++ b/tensorflow/core/debug/debug_callback_registry.cc
@@ -17,7 +17,7 @@ limitations under the License.
 
 namespace tensorflow {
 
-DebugCallbackRegistry::DebugCallbackRegistry() {}
+DebugCallbackRegistry::DebugCallbackRegistry() = default;
 
 /*static */ DebugCallbackRegistry* DebugCallbackRegistry::instance_ = nullptr;
 
diff --git a/tensorflow/core/debug/debug_io_utils.h b/tensorflow/core/debug/debug_io_utils.h
index 99107971f0f..50af776f7d3 100644
--- a/tensorflow/core/debug/debug_io_utils.h
+++ b/tensorflow/core/debug/debug_io_utils.h
@@ -287,7 +287,7 @@ class DebugGrpcChannel {
   //     debug_service.proto). E.g., "127.0.0.1:12345".
   explicit DebugGrpcChannel(const std::string& server_stream_addr);
 
-  virtual ~DebugGrpcChannel() {}
+  virtual ~DebugGrpcChannel() = default;
 
   // Attempt to establish connection with server.
   //
diff --git a/tensorflow/core/debug/debugger_state_impl.h b/tensorflow/core/debug/debugger_state_impl.h
index 73e74738d59..c60144f7f96 100644
--- a/tensorflow/core/debug/debugger_state_impl.h
+++ b/tensorflow/core/debug/debugger_state_impl.h
@@ -47,7 +47,7 @@ class DebugGraphDecorator : public DebugGraphDecoratorInterface {
  public:
   DebugGraphDecorator(const DebugOptions& debug_options)
       : debug_options_(debug_options) {}
-  ~DebugGraphDecorator() override {}
+  ~DebugGraphDecorator() override = default;
 
   absl::Status DecorateGraph(Graph* graph, Device* device) override;
   absl::Status PublishGraph(const Graph& graph,

From 7d1e6334c3e765dbee1c547f8b5950312fe788d6 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 13:43:47 -0700
Subject: [PATCH 0439/1000] Automated Code Change

PiperOrigin-RevId: 997182211
---
 tensorflow/lite/delegates/nnapi/nnapi_delegate_test.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/lite/delegates/nnapi/nnapi_delegate_test.cc b/tensorflow/lite/delegates/nnapi/nnapi_delegate_test.cc
index 95f56de6f8f..f03de61372a 100644
--- a/tensorflow/lite/delegates/nnapi/nnapi_delegate_test.cc
+++ b/tensorflow/lite/delegates/nnapi/nnapi_delegate_test.cc
@@ -1809,7 +1809,7 @@ TEST(NNAPIDelegate, DivWithNoActivation) {
 
 class BaseConcatenationOpModel : public SingleOpModelWithNNAPI {
  public:
-  BaseConcatenationOpModel() {}
+  BaseConcatenationOpModel() = default;
   BaseConcatenationOpModel(const TensorData& input_template, int axis,
                            int num_inputs) {
     std::vector<std::vector<int>> all_input_shapes;

From 12cb170c73b7cd526280ddca763e634e32309e8b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 13:52:50 -0700
Subject: [PATCH 0440/1000] Automated Code Change

PiperOrigin-RevId: 997185017
---
 tensorflow/core/kernels/batch_kernels.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/batch_kernels.cc b/tensorflow/core/kernels/batch_kernels.cc
index 58aadb27d12..4af394fe2e4 100644
--- a/tensorflow/core/kernels/batch_kernels.cc
+++ b/tensorflow/core/kernels/batch_kernels.cc
@@ -1081,7 +1081,7 @@ REGISTER_KERNEL_BUILDER(Name("Unbatch").Device(DEVICE_CPU), UnbatchKernel);
 // deterministically for the gradient of unbatch.
 class UnbatchGradResource : public ResourceBase {
  public:
-  UnbatchGradResource() {}
+  UnbatchGradResource() = default;
 
   std::string DebugString() const final { return "UnbatchGradResource"; }
 

From b356eee7fd1b2107f618f44a144d0fc5046b43dc Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 14:06:25 -0700
Subject: [PATCH 0441/1000] Automated Code Change

PiperOrigin-RevId: 997189043
---
 tensorflow/core/grappler/graph_analyzer/graph_analyzer.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/grappler/graph_analyzer/graph_analyzer.cc b/tensorflow/core/grappler/graph_analyzer/graph_analyzer.cc
index 2a19adb4120..a5036201455 100644
--- a/tensorflow/core/grappler/graph_analyzer/graph_analyzer.cc
+++ b/tensorflow/core/grappler/graph_analyzer/graph_analyzer.cc
@@ -36,7 +36,7 @@ namespace graph_analyzer {
 GraphAnalyzer::GraphAnalyzer(const GraphDef& graph, int subgraph_size)
     : graph_(graph), subgraph_size_(subgraph_size) {}
 
-GraphAnalyzer::~GraphAnalyzer() {}
+GraphAnalyzer::~GraphAnalyzer() = default;
 
 absl::Status GraphAnalyzer::Run() {
   // The signature computation code would detect this too, but better

From 21845db935547569bdad455aa88264fcd2327c27 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 14:09:22 -0700
Subject: [PATCH 0442/1000] Automated Code Change

PiperOrigin-RevId: 997190234
---
 .../core/kernels/data/experimental/stats_aggregator_ops.cc    | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/kernels/data/experimental/stats_aggregator_ops.cc b/tensorflow/core/kernels/data/experimental/stats_aggregator_ops.cc
index 6999638e839..ea258d3bbba 100644
--- a/tensorflow/core/kernels/data/experimental/stats_aggregator_ops.cc
+++ b/tensorflow/core/kernels/data/experimental/stats_aggregator_ops.cc
@@ -56,7 +56,7 @@ get_counters_map() {
 
 class StatsAggregatorImpl : public StatsAggregator {
  public:
-  StatsAggregatorImpl() {}
+  StatsAggregatorImpl() = default;
 
   void AddToHistogram(const std::string& name, absl::Span<const double> values,
                       const int64_t steps) override {
@@ -139,7 +139,7 @@ class StatsAggregatorHandleOp
 
 class StatsAggregatorImplV2 : public StatsAggregator {
  public:
-  StatsAggregatorImplV2() {}
+  StatsAggregatorImplV2() = default;
 
   ~StatsAggregatorImplV2() override {
     if (summary_writer_interface_) {

From 3e9c7fbacd83348ea6c81e133f5c3ad9dcd8e6c4 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 14:29:49 -0700
Subject: [PATCH 0443/1000] Automated Code Change

PiperOrigin-RevId: 997196729
---
 tensorflow/c/eager/abstract_context.h                        | 2 +-
 tensorflow/c/eager/abstract_operation.h                      | 2 +-
 tensorflow/c/eager/abstract_tensor_handle.h                  | 2 +-
 tensorflow/c/eager/c_api_distributed_test.cc                 | 2 +-
 tensorflow/c/eager/c_api_unified_experimental_graph.cc       | 2 +-
 tensorflow/c/eager/gradients.h                               | 2 +-
 tensorflow/c/eager/graph_function.cc                         | 2 +-
 tensorflow/c/eager/immediate_execution_context.h             | 2 +-
 tensorflow/c/eager/immediate_execution_distributed_manager.h | 2 +-
 tensorflow/c/eager/immediate_execution_tensor_handle.h       | 2 +-
 10 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/tensorflow/c/eager/abstract_context.h b/tensorflow/c/eager/abstract_context.h
index dce451ceadd..442302d9eb2 100644
--- a/tensorflow/c/eager/abstract_context.h
+++ b/tensorflow/c/eager/abstract_context.h
@@ -34,7 +34,7 @@ class AbstractContext {
  protected:
   enum AbstractContextKind { kGraph, kMlir, kEager, kTfrt, kTape, kOpHandler };
   explicit AbstractContext(AbstractContextKind kind) : kind_(kind) {}
-  virtual ~AbstractContext() {}
+  virtual ~AbstractContext() = default;
 
  public:
   AbstractContextKind getKind() const { return kind_; }
diff --git a/tensorflow/c/eager/abstract_operation.h b/tensorflow/c/eager/abstract_operation.h
index c60f7bb158a..ac6bfa2dea2 100644
--- a/tensorflow/c/eager/abstract_operation.h
+++ b/tensorflow/c/eager/abstract_operation.h
@@ -40,7 +40,7 @@ class AbstractOperation {
     kOpHandler
   };
   explicit AbstractOperation(AbstractOperationKind kind) : kind_(kind) {}
-  virtual ~AbstractOperation() {}
+  virtual ~AbstractOperation() = default;
 
  public:
   AbstractOperationKind getKind() const { return kind_; }
diff --git a/tensorflow/c/eager/abstract_tensor_handle.h b/tensorflow/c/eager/abstract_tensor_handle.h
index 4a40b1c9319..5dcfb1866ab 100644
--- a/tensorflow/c/eager/abstract_tensor_handle.h
+++ b/tensorflow/c/eager/abstract_tensor_handle.h
@@ -31,7 +31,7 @@ class AbstractTensorHandle : public core::RefCounted {
  protected:
   enum AbstractTensorHandleKind { kGraph, kMlir, kEager, kTfrt, kCustomDevice };
   explicit AbstractTensorHandle(AbstractTensorHandleKind kind) : kind_(kind) {}
-  ~AbstractTensorHandle() override {}
+  ~AbstractTensorHandle() override = default;
 
  public:
   // Returns tensor dtype.
diff --git a/tensorflow/c/eager/c_api_distributed_test.cc b/tensorflow/c/eager/c_api_distributed_test.cc
index aff1c5a95d3..694867741e8 100644
--- a/tensorflow/c/eager/c_api_distributed_test.cc
+++ b/tensorflow/c/eager/c_api_distributed_test.cc
@@ -310,7 +310,7 @@ string VariableAddFunction() {
 class GraphErrorInjectionPass : public tensorflow::GraphOptimizationPass {
  public:
   static bool enabled_;
-  GraphErrorInjectionPass() {}
+  GraphErrorInjectionPass() = default;
 
   absl::Status Run(
       const tensorflow::GraphOptimizationPassOptions& options) override {
diff --git a/tensorflow/c/eager/c_api_unified_experimental_graph.cc b/tensorflow/c/eager/c_api_unified_experimental_graph.cc
index 4673500b36c..6e3c2296cff 100644
--- a/tensorflow/c/eager/c_api_unified_experimental_graph.cc
+++ b/tensorflow/c/eager/c_api_unified_experimental_graph.cc
@@ -321,7 +321,7 @@ class GraphOperation : public TracingOperation {
   static bool classof(const AbstractOperation* ptr) {
     return ptr->getKind() == kGraph;
   }
-  ~GraphOperation() override {}
+  ~GraphOperation() override = default;
 
  private:
   friend class GraphContext;  // For access to op_.
diff --git a/tensorflow/c/eager/gradients.h b/tensorflow/c/eager/gradients.h
index 932f7ea059b..2f516088d3c 100644
--- a/tensorflow/c/eager/gradients.h
+++ b/tensorflow/c/eager/gradients.h
@@ -58,7 +58,7 @@ class GradientFunction {
       AbstractContext* ctx,
       absl::Span<AbstractTensorHandle* const> grad_outputs,
       absl::Span<AbstractTensorHandle*> grad_inputs) = 0;
-  virtual ~GradientFunction() {}
+  virtual ~GradientFunction() = default;
 };
 
 // Metadata from the forward operation that is made available to the
diff --git a/tensorflow/c/eager/graph_function.cc b/tensorflow/c/eager/graph_function.cc
index ce11530aa15..3c34a009b79 100644
--- a/tensorflow/c/eager/graph_function.cc
+++ b/tensorflow/c/eager/graph_function.cc
@@ -26,7 +26,7 @@ namespace graph {
 GraphFunction::GraphFunction(FunctionDef fdef)
     : AbstractFunction(kGraph),
       func_record_(new FunctionRecord(std::move(fdef), {}, true)) {}
-GraphFunction::~GraphFunction() {}
+GraphFunction::~GraphFunction() = default;
 absl::Status GraphFunction::GetFunctionDef(const FunctionDef **fdef) {
   *fdef = &(func_record_->fdef());
   return absl::OkStatus();
diff --git a/tensorflow/c/eager/immediate_execution_context.h b/tensorflow/c/eager/immediate_execution_context.h
index 158a42d0e1b..f8df47a893b 100644
--- a/tensorflow/c/eager/immediate_execution_context.h
+++ b/tensorflow/c/eager/immediate_execution_context.h
@@ -272,7 +272,7 @@ class ImmediateExecutionContext : public AbstractContext {
  protected:
   explicit ImmediateExecutionContext(AbstractContextKind kind)
       : AbstractContext(kind) {}
-  ~ImmediateExecutionContext() override {}
+  ~ImmediateExecutionContext() override = default;
 };
 
 namespace internal {
diff --git a/tensorflow/c/eager/immediate_execution_distributed_manager.h b/tensorflow/c/eager/immediate_execution_distributed_manager.h
index f4f4f0931dd..85dadc69575 100644
--- a/tensorflow/c/eager/immediate_execution_distributed_manager.h
+++ b/tensorflow/c/eager/immediate_execution_distributed_manager.h
@@ -33,7 +33,7 @@ class WorkerCacheInterface;
 
 class ImmediateExecutionDistributedManager {
  public:
-  virtual ~ImmediateExecutionDistributedManager() {}
+  virtual ~ImmediateExecutionDistributedManager() = default;
 
   // Set up distributed execution environment on local and remote tasks.
   // When `reset_context` is true, initialize new cluster context state based
diff --git a/tensorflow/c/eager/immediate_execution_tensor_handle.h b/tensorflow/c/eager/immediate_execution_tensor_handle.h
index 61fc0fe8c04..23fe259ef10 100644
--- a/tensorflow/c/eager/immediate_execution_tensor_handle.h
+++ b/tensorflow/c/eager/immediate_execution_tensor_handle.h
@@ -83,7 +83,7 @@ class ImmediateExecutionTensorHandle : public AbstractTensorHandle {
  protected:
   explicit ImmediateExecutionTensorHandle(AbstractTensorHandleKind kind)
       : AbstractTensorHandle(kind) {}
-  ~ImmediateExecutionTensorHandle() override {}
+  ~ImmediateExecutionTensorHandle() override = default;
 };
 
 namespace internal {

From 86290092aad63a832c64308123759d491b432500 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 14:44:08 -0700
Subject: [PATCH 0444/1000] Automated Code Change

PiperOrigin-RevId: 997200939
---
 .../experimental/acceleration/mini_benchmark/libjpeg_handle.h | 2 +-
 .../experimental/acceleration/mini_benchmark/mini_benchmark.h | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_handle.h b/tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_handle.h
index 88fab33872b..969eff4adec 100644
--- a/tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_handle.h
+++ b/tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_handle.h
@@ -63,7 +63,7 @@ class LibjpegHandle {
   boolean (*jpeg_finish_decompress_)(j_decompress_ptr);
 
  private:
-  LibjpegHandle() {}
+  LibjpegHandle() = default;
   void* libjpeg_ = nullptr;
 };
 
diff --git a/tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark.h b/tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark.h
index 2cc952e1bc0..1b592d12b33 100644
--- a/tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark.h
+++ b/tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark.h
@@ -66,8 +66,8 @@ class MiniBenchmark {
   // if the overall mini-benchmark-related setup isn't properly initialized.
   virtual int NumRemainingAccelerationTests() = 0;
 
-  MiniBenchmark() {}
-  virtual ~MiniBenchmark() {}
+  MiniBenchmark() = default;
+  virtual ~MiniBenchmark() = default;
 
   MiniBenchmark(MiniBenchmark&) = delete;
   MiniBenchmark& operator=(const MiniBenchmark&) = delete;

From c18681a1809f244f5490cb57186e37443dc61b66 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 14:49:26 -0700
Subject: [PATCH 0445/1000] Automated Code Change

PiperOrigin-RevId: 997202333
---
 tensorflow/core/common_runtime/base_collective_executor.cc    | 4 ++--
 tensorflow/core/common_runtime/base_collective_executor.h     | 2 +-
 .../core/common_runtime/collective_param_resolver_local.h     | 2 +-
 tensorflow/core/common_runtime/cost_measurement.h             | 2 +-
 tensorflow/core/common_runtime/debugger_state_interface.h     | 4 ++--
 tensorflow/core/common_runtime/device_mgr.cc                  | 2 +-
 tensorflow/core/common_runtime/direct_session.cc              | 2 +-
 tensorflow/core/common_runtime/executor.h                     | 4 ++--
 tensorflow/core/common_runtime/executor_factory.h             | 2 +-
 tensorflow/core/common_runtime/function_body.h                | 2 +-
 10 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/tensorflow/core/common_runtime/base_collective_executor.cc b/tensorflow/core/common_runtime/base_collective_executor.cc
index 11f11c87386..eb8f9a98628 100644
--- a/tensorflow/core/common_runtime/base_collective_executor.cc
+++ b/tensorflow/core/common_runtime/base_collective_executor.cc
@@ -111,7 +111,7 @@ class CollectiveAdapterImpl : public CollectiveAdapter {
     Flatten();
   }
 
-  ~CollectiveAdapterImpl() override {}
+  ~CollectiveAdapterImpl() override = default;
 
   const Tensor& Value() const override { return output_; }
 
@@ -227,7 +227,7 @@ CollectiveAdapter* MakeCollectiveAdapter(Tensor* output, int num_chunks,
   }
 }
 
-BaseCollectiveExecutor::~BaseCollectiveExecutor() {}
+BaseCollectiveExecutor::~BaseCollectiveExecutor() = default;
 
 void BaseCollectiveExecutor::StartAbort(const absl::Status& s) {
   if (flags::Global().enable_fatal_error_on_collective_abort.value()) {
diff --git a/tensorflow/core/common_runtime/base_collective_executor.h b/tensorflow/core/common_runtime/base_collective_executor.h
index afde79dc0f9..6b557867d69 100644
--- a/tensorflow/core/common_runtime/base_collective_executor.h
+++ b/tensorflow/core/common_runtime/base_collective_executor.h
@@ -32,7 +32,7 @@ class Device;
 // Tensors for in-place update.
 class CollectiveAdapter {
  public:
-  virtual ~CollectiveAdapter() {}
+  virtual ~CollectiveAdapter() = default;
 
   // Move the backing tensor to 'output' with its original storage and
   // shape. After this call this CollectiveAdapter object should be
diff --git a/tensorflow/core/common_runtime/collective_param_resolver_local.h b/tensorflow/core/common_runtime/collective_param_resolver_local.h
index ed429020bd4..9a21bb4c3f9 100644
--- a/tensorflow/core/common_runtime/collective_param_resolver_local.h
+++ b/tensorflow/core/common_runtime/collective_param_resolver_local.h
@@ -47,7 +47,7 @@ class CollectiveParamResolverLocal : public ParamResolverInterface {
                                NcclCommunicatorInterface* nccl_communicator,
                                const std::string& task_name);
 
-  ~CollectiveParamResolverLocal() override {}
+  ~CollectiveParamResolverLocal() override = default;
 
   void CompleteParamsAsync(const DeviceAttributes& device, CollectiveParams* cp,
                            CancellationManager* cancel_mgr,
diff --git a/tensorflow/core/common_runtime/cost_measurement.h b/tensorflow/core/common_runtime/cost_measurement.h
index 3da322e5dbc..d39330d37d2 100644
--- a/tensorflow/core/common_runtime/cost_measurement.h
+++ b/tensorflow/core/common_runtime/cost_measurement.h
@@ -33,7 +33,7 @@ class CostMeasurement {
 
   explicit CostMeasurement(const Context& context) {}
 
-  virtual ~CostMeasurement() {}
+  virtual ~CostMeasurement() = default;
 
   virtual absl::Duration GetTotalCost() = 0;
 
diff --git a/tensorflow/core/common_runtime/debugger_state_interface.h b/tensorflow/core/common_runtime/debugger_state_interface.h
index 0c4e55873da..2a550c5d5cb 100644
--- a/tensorflow/core/common_runtime/debugger_state_interface.h
+++ b/tensorflow/core/common_runtime/debugger_state_interface.h
@@ -34,7 +34,7 @@ std::string SummarizeDebugTensorWatches(
 // An abstract interface for storing and retrieving debugging information.
 class DebuggerStateInterface {
  public:
-  virtual ~DebuggerStateInterface() {}
+  virtual ~DebuggerStateInterface() = default;
 
   // Publish metadata about the debugged Session::Run() call.
   //
@@ -57,7 +57,7 @@ class DebuggerStateInterface {
 
 class DebugGraphDecoratorInterface {
  public:
-  virtual ~DebugGraphDecoratorInterface() {}
+  virtual ~DebugGraphDecoratorInterface() = default;
 
   // Insert special-purpose debug nodes to graph and dump the graph for
   // record. See the documentation of DebugNodeInserter::InsertNodes() for
diff --git a/tensorflow/core/common_runtime/device_mgr.cc b/tensorflow/core/common_runtime/device_mgr.cc
index 97055b3d722..43b68da0266 100644
--- a/tensorflow/core/common_runtime/device_mgr.cc
+++ b/tensorflow/core/common_runtime/device_mgr.cc
@@ -26,6 +26,6 @@ limitations under the License.
 
 namespace tensorflow {
 
-DeviceMgr::~DeviceMgr() {}
+DeviceMgr::~DeviceMgr() = default;
 
 }  // namespace tensorflow
diff --git a/tensorflow/core/common_runtime/direct_session.cc b/tensorflow/core/common_runtime/direct_session.cc
index ef72a33d1ff..5b4717ae2e6 100644
--- a/tensorflow/core/common_runtime/direct_session.cc
+++ b/tensorflow/core/common_runtime/direct_session.cc
@@ -204,7 +204,7 @@ std::string GetRendezvousKey(const std::string& tensor_name,
 
 class DirectSessionFactory : public SessionFactory {
  public:
-  DirectSessionFactory() {}
+  DirectSessionFactory() = default;
 
   bool AcceptsOptions(const SessionOptions& options) override {
     return options.target.empty() &&
diff --git a/tensorflow/core/common_runtime/executor.h b/tensorflow/core/common_runtime/executor.h
index 489b2a9ed15..e204cfc673a 100644
--- a/tensorflow/core/common_runtime/executor.h
+++ b/tensorflow/core/common_runtime/executor.h
@@ -55,7 +55,7 @@ class StepStatsCollector;
 // Multiple threads can call Executor::Run concurrently.
 class Executor {
  public:
-  virtual ~Executor() {}
+  virtual ~Executor() = default;
 
   // RunAsync() executes the graph computation. "done" is run when the
   // graph computation completes. If any error happens during the
@@ -182,7 +182,7 @@ class ExecutorBarrier {
   ExecutorBarrier(size_t num, Rendezvous* r, StatusCallback done)
       : rendez_(r), done_cb_(done), pending_(num) {}
 
-  ~ExecutorBarrier() {}
+  ~ExecutorBarrier() = default;
 
   // Returns a closure that Executors must call when they are done
   // computing, passing the status of their execution as an argument.
diff --git a/tensorflow/core/common_runtime/executor_factory.h b/tensorflow/core/common_runtime/executor_factory.h
index 3459a4a38b0..823a7a2d605 100644
--- a/tensorflow/core/common_runtime/executor_factory.h
+++ b/tensorflow/core/common_runtime/executor_factory.h
@@ -34,7 +34,7 @@ class ExecutorFactory {
   virtual absl::Status NewExecutor(const LocalExecutorParams& params,
                                    const Graph& graph,
                                    std::unique_ptr<Executor>* out_executor) = 0;
-  virtual ~ExecutorFactory() {}
+  virtual ~ExecutorFactory() = default;
 
   static void Register(const std::string& executor_type,
                        ExecutorFactory* factory);
diff --git a/tensorflow/core/common_runtime/function_body.h b/tensorflow/core/common_runtime/function_body.h
index 2f868eb8607..e19e8024ff1 100644
--- a/tensorflow/core/common_runtime/function_body.h
+++ b/tensorflow/core/common_runtime/function_body.h
@@ -49,7 +49,7 @@ struct FunctionBody {
   absl::InlinedVector<AllocatorAttributes, 4UL> args_alloc_attrs;
   absl::InlinedVector<AllocatorAttributes, 4UL> rets_alloc_attrs;
 
-  FunctionBody() {}
+  FunctionBody() = default;
   FunctionBody(core::RefCountPtr<FunctionRecord>&& record,
                DataTypeSlice arg_types, DataTypeSlice ret_types, Graph* g);
   ~FunctionBody();

From aaecea2527f21e3310989b0572209316ecd52198 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 20:18:25 -0700
Subject: [PATCH 0446/1000] Automated Code Change

PiperOrigin-RevId: 997280025
---
 tensorflow/core/debug/debug_node_key.cc | 2 +-
 tensorflow/core/debug/debug_node_key.h  | 3 +--
 2 files changed, 2 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/debug/debug_node_key.cc b/tensorflow/core/debug/debug_node_key.cc
index 09510b8df1b..ae641cb5b3c 100644
--- a/tensorflow/core/debug/debug_node_key.cc
+++ b/tensorflow/core/debug/debug_node_key.cc
@@ -54,7 +54,7 @@ bool DebugNodeKey::operator!=(const DebugNodeKey& other) const {
   return !((*this) == other);
 }
 
-const std::string DebugNodeKey::DeviceNameToDevicePath(
+std::string DebugNodeKey::DeviceNameToDevicePath(
     const std::string& device_name) {
   return absl::StrCat(kMetadataFilePrefix, kDeviceTag,
                       str_util::StringReplace(
diff --git a/tensorflow/core/debug/debug_node_key.h b/tensorflow/core/debug/debug_node_key.h
index 867e0809314..2058439678f 100644
--- a/tensorflow/core/debug/debug_node_key.h
+++ b/tensorflow/core/debug/debug_node_key.h
@@ -35,8 +35,7 @@ struct DebugNodeKey {
   // Converts a device name string to a device path string.
   // E.g., /job:localhost/replica:0/task:0/cpu:0 will be converted to
   //   ,job_localhost,replica_0,task_0,cpu_0.
-  static const std::string DeviceNameToDevicePath(
-      const std::string& device_name);
+  static std::string DeviceNameToDevicePath(const std::string& device_name);
 
   bool operator==(const DebugNodeKey& other) const;
   bool operator!=(const DebugNodeKey& other) const;

From 79d677136df017f86ec9f06e8541ee31e35fca64 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 20:27:05 -0700
Subject: [PATCH 0447/1000] Automated Code Change

PiperOrigin-RevId: 997282489
---
 tensorflow/core/grappler/verifiers/graph_verifier.h     | 4 ++--
 tensorflow/core/grappler/verifiers/structure_verifier.h | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/grappler/verifiers/graph_verifier.h b/tensorflow/core/grappler/verifiers/graph_verifier.h
index 0e59d4ed3a2..a725de4054f 100644
--- a/tensorflow/core/grappler/verifiers/graph_verifier.h
+++ b/tensorflow/core/grappler/verifiers/graph_verifier.h
@@ -37,8 +37,8 @@ namespace grappler {
 //     specification of the backend.
 class GraphVerifier {
  public:
-  GraphVerifier() {}
-  virtual ~GraphVerifier() {}
+  GraphVerifier() = default;
+  virtual ~GraphVerifier() = default;
 
   // A name for the verifier.
   virtual std::string name() const = 0;
diff --git a/tensorflow/core/grappler/verifiers/structure_verifier.h b/tensorflow/core/grappler/verifiers/structure_verifier.h
index 968f840b41c..19bfb7292ea 100644
--- a/tensorflow/core/grappler/verifiers/structure_verifier.h
+++ b/tensorflow/core/grappler/verifiers/structure_verifier.h
@@ -29,8 +29,8 @@ namespace grappler {
 // Verifies the structure of a graph to ensure it is valid.
 class StructureVerifier : public GraphVerifier {
  public:
-  StructureVerifier() {}
-  ~StructureVerifier() override {}
+  StructureVerifier() = default;
+  ~StructureVerifier() override = default;
 
   std::string name() const override { return "structure_verifier"; };
 

From a89cbe26c8363a49b28234071159bf145994ff0a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 20:31:05 -0700
Subject: [PATCH 0448/1000] Automated Code Change

PiperOrigin-RevId: 997283523
---
 .../core/distributed_runtime/rpc/eager/grpc_eager_client.cc     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_client.cc b/tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_client.cc
index fba4bd2765f..a93506acf7a 100644
--- a/tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_client.cc
+++ b/tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_client.cc
@@ -116,7 +116,7 @@ class GrpcEagerClientThread : public core::RefCounted {
         }));
   }
 
-  ~GrpcEagerClientThread() override {}
+  ~GrpcEagerClientThread() override = default;
 
   ::grpc::CompletionQueue* completion_queue() { return &completion_queue_; }
 

From baa4e45d3c6a9a922de90ae733ba3637d167f794 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 20:37:15 -0700
Subject: [PATCH 0449/1000] Automated Code Change

PiperOrigin-RevId: 997285033
---
 tensorflow/core/kernels/summary_interface.h  |  2 +-
 tensorflow/core/kernels/winograd_transform.h | 18 +++++++++---------
 2 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/tensorflow/core/kernels/summary_interface.h b/tensorflow/core/kernels/summary_interface.h
index 0f743c4dfa2..cd25ae75d7a 100644
--- a/tensorflow/core/kernels/summary_interface.h
+++ b/tensorflow/core/kernels/summary_interface.h
@@ -29,7 +29,7 @@ class GraphDef;
 // Main interface for the summary writer resource.
 class SummaryWriterInterface : public ResourceBase {
  public:
-  virtual ~SummaryWriterInterface() override {}
+  ~SummaryWriterInterface() override {}
 
   // Flushes all unwritten messages in the queue.
   virtual absl::Status Flush() = 0;
diff --git a/tensorflow/core/kernels/winograd_transform.h b/tensorflow/core/kernels/winograd_transform.h
index 4f4067e37e6..cd36d8b909a 100644
--- a/tensorflow/core/kernels/winograd_transform.h
+++ b/tensorflow/core/kernels/winograd_transform.h
@@ -33,18 +33,18 @@ class WinogradTransform : public DeepConv2DTransform<T> {
   WinogradTransform()
       : filter_shape_(3, 3), input_shape_(4, 4), output_shape_(2, 2) {}
 
-  virtual void GetFilterTransformMatrix(const int64_t rows, const int64_t cols,
-                                        T* transform_matrix) const;
+  void GetFilterTransformMatrix(const int64_t rows, const int64_t cols,
+                                T* transform_matrix) const override;
 
-  virtual void GetInputTransformMatrix(const int64_t rows, const int64_t cols,
-                                       T* transform_matrix) const;
+  void GetInputTransformMatrix(const int64_t rows, const int64_t cols,
+                               T* transform_matrix) const override;
 
-  virtual void GetOutputTransformMatrix(const int64_t rows, const int64_t cols,
-                                        T* transform_matrix) const;
+  void GetOutputTransformMatrix(const int64_t rows, const int64_t cols,
+                                T* transform_matrix) const override;
 
-  virtual const Shape& filter_shape() const { return filter_shape_; }
-  virtual const Shape& input_shape() const { return input_shape_; }
-  virtual const Shape& output_shape() const { return output_shape_; }
+  const Shape& filter_shape() const override { return filter_shape_; }
+  const Shape& input_shape() const override { return input_shape_; }
+  const Shape& output_shape() const override { return output_shape_; }
 
  private:
   const Shape filter_shape_;

From 7899e3cf72f2b49ce81269b90a41f9457b4b4484 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 20:37:54 -0700
Subject: [PATCH 0450/1000] Automated Code Change

PiperOrigin-RevId: 997285167
---
 tensorflow/compiler/tf2xla/resource_util.h           | 2 +-
 tensorflow/compiler/tf2xla/xla_compilation_device.cc | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/compiler/tf2xla/resource_util.h b/tensorflow/compiler/tf2xla/resource_util.h
index e4bdb5112be..d1eee1c9056 100644
--- a/tensorflow/compiler/tf2xla/resource_util.h
+++ b/tensorflow/compiler/tf2xla/resource_util.h
@@ -38,7 +38,7 @@ class ResourceUsageAnalysis {
     std::string node_name_;
     std::string op_;
 
-    NodeInfo() {}
+    NodeInfo() = default;
 
     NodeInfo(const std::optional<std::string>& function_name,
              std::string node_name, std::string op)
diff --git a/tensorflow/compiler/tf2xla/xla_compilation_device.cc b/tensorflow/compiler/tf2xla/xla_compilation_device.cc
index 25dea2ccf95..3ccff46dd13 100644
--- a/tensorflow/compiler/tf2xla/xla_compilation_device.cc
+++ b/tensorflow/compiler/tf2xla/xla_compilation_device.cc
@@ -36,8 +36,8 @@ namespace tensorflow {
 // represented by the Tensor.
 class XlaCompilationAllocator : public Allocator {
  public:
-  XlaCompilationAllocator() {}
-  ~XlaCompilationAllocator() override {}
+  XlaCompilationAllocator() = default;
+  ~XlaCompilationAllocator() override = default;
 
   std::string Name() override { return "xla_compilation"; }
 
@@ -78,7 +78,7 @@ XlaCompilationDevice::XlaCompilationDevice(const SessionOptions& options,
                                             type.type()))),
       allocator_(new XlaCompilationAllocator()) {}
 
-XlaCompilationDevice::~XlaCompilationDevice() {}
+XlaCompilationDevice::~XlaCompilationDevice() = default;
 
 Allocator* XlaCompilationDevice::GetAllocator(AllocatorAttributes attr) {
   return allocator_.get();

From 8e4830bdf0eb2ba62ddf3c57987d20566b3b5b64 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 20:41:08 -0700
Subject: [PATCH 0451/1000] Automated Code Change

PiperOrigin-RevId: 997285914
---
 tensorflow/core/kernels/sparse/add_op.cc | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/kernels/sparse/add_op.cc b/tensorflow/core/kernels/sparse/add_op.cc
index 37331c96026..5e0466d3342 100644
--- a/tensorflow/core/kernels/sparse/add_op.cc
+++ b/tensorflow/core/kernels/sparse/add_op.cc
@@ -281,7 +281,7 @@ struct CSRSparseMatrixAdd<GPUDevice, T>
         beta_(beta),
         initialized_(false) {}
 
-  absl::Status Initialize() {
+  absl::Status Initialize() override {
     TF_RETURN_IF_ERROR(cuda_sparse_.Initialize());
     TF_RETURN_IF_ERROR(descrA_.Initialize());
     TF_RETURN_IF_ERROR(descrB_.Initialize());
@@ -292,7 +292,7 @@ struct CSRSparseMatrixAdd<GPUDevice, T>
 
   absl::Status GetWorkspaceSize(const ConstCSRComponent<T>& a,
                                 const ConstCSRComponent<T>& b,
-                                size_t* bufferSize) {
+                                size_t* bufferSize) override {
     DCHECK(initialized_);
 
     const int m = a.row_ptr.size() - 1;
@@ -320,7 +320,7 @@ struct CSRSparseMatrixAdd<GPUDevice, T>
   absl::Status GetOutputStructure(const ConstCSRComponent<T>& a,
                                   const ConstCSRComponent<T>& b,
                                   TTypes<int32_t>::UnalignedVec c_row_ptr,
-                                  int* output_nnz, void* workspace) {
+                                  int* output_nnz, void* workspace) override {
     DCHECK(initialized_);
 
     const int m = a.row_ptr.size() - 1;
@@ -349,7 +349,7 @@ struct CSRSparseMatrixAdd<GPUDevice, T>
 
   absl::Status Compute(const ConstCSRComponent<T>& a,
                        const ConstCSRComponent<T>& b, CSRComponent<T>* c,
-                       void* workspace) {
+                       void* workspace) override {
     DCHECK(initialized_);
 
     const int m = a.row_ptr.size() - 1;

From 3f6eba0318f34dd9950f19b0d06543c30bf1f69a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 20:45:04 -0700
Subject: [PATCH 0452/1000] Automated Code Change

PiperOrigin-RevId: 997286768
---
 tensorflow/core/distributed_runtime/scheduler.h            | 4 ++--
 tensorflow/core/distributed_runtime/server_lib.h           | 6 +++---
 tensorflow/core/distributed_runtime/session_mgr.h          | 2 +-
 tensorflow/core/distributed_runtime/tensor_coding.cc       | 2 +-
 tensorflow/core/distributed_runtime/tensor_coding.h        | 2 +-
 tensorflow/core/distributed_runtime/worker_cache.h         | 2 +-
 tensorflow/core/distributed_runtime/worker_cache_partial.h | 2 +-
 tensorflow/core/distributed_runtime/worker_interface.h     | 2 +-
 8 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/tensorflow/core/distributed_runtime/scheduler.h b/tensorflow/core/distributed_runtime/scheduler.h
index d277bdab74e..18a2d65398d 100644
--- a/tensorflow/core/distributed_runtime/scheduler.h
+++ b/tensorflow/core/distributed_runtime/scheduler.h
@@ -32,7 +32,7 @@ class SlackAnalysis {
  public:
   SlackAnalysis(const Graph* g, const CostModel* cost_model);
 
-  ~SlackAnalysis() {}
+  ~SlackAnalysis() = default;
 
   // Compute the earliest possible start time for each node, based on
   // a given cost model. 'asap_time' is indexed by node id.
@@ -97,7 +97,7 @@ class PriorityScheduler {
   PriorityScheduler(const DeviceSet* devices, const CostModel* cost_model,
                     const Graph* g);
 
-  ~PriorityScheduler() {}
+  ~PriorityScheduler() = default;
 
   // Computes a schedule of the ideal start time for each node.
   // Returns the makespan (the total running time).
diff --git a/tensorflow/core/distributed_runtime/server_lib.h b/tensorflow/core/distributed_runtime/server_lib.h
index 66e336ff3e7..c09f16f4059 100644
--- a/tensorflow/core/distributed_runtime/server_lib.h
+++ b/tensorflow/core/distributed_runtime/server_lib.h
@@ -44,8 +44,8 @@ class MasterEnv;
 // services.
 class ServerInterface {
  public:
-  ServerInterface() {}
-  virtual ~ServerInterface() {}
+  ServerInterface() = default;
+  virtual ~ServerInterface() = default;
 
   // Starts the server running asynchronously. Returns OK on success, otherwise
   // returns an error.
@@ -107,7 +107,7 @@ class ServerFactory {
   // based on the given `server_def`.
   virtual bool AcceptsOptions(const ServerDef& server_def) = 0;
 
-  virtual ~ServerFactory() {}
+  virtual ~ServerFactory() = default;
 
   // For each `ServerFactory` subclass, an instance of that class must
   // be registered by calling this method.
diff --git a/tensorflow/core/distributed_runtime/session_mgr.h b/tensorflow/core/distributed_runtime/session_mgr.h
index 0a2bddddb1a..e6266f01501 100644
--- a/tensorflow/core/distributed_runtime/session_mgr.h
+++ b/tensorflow/core/distributed_runtime/session_mgr.h
@@ -49,7 +49,7 @@ class SessionMgr {
       std::unique_ptr<WorkerCacheInterface> default_worker_cache,
       WorkerCacheFactory worker_cache_factory,
       tsl::CoordinationServiceRpcHandler* coordination_handler);
-  ~SessionMgr() {}
+  ~SessionMgr() = default;
 
   // Allocates state for a new session.
   absl::Status CreateSession(
diff --git a/tensorflow/core/distributed_runtime/tensor_coding.cc b/tensorflow/core/distributed_runtime/tensor_coding.cc
index 3f0caa745f0..95a55d319c5 100644
--- a/tensorflow/core/distributed_runtime/tensor_coding.cc
+++ b/tensorflow/core/distributed_runtime/tensor_coding.cc
@@ -27,7 +27,7 @@ limitations under the License.
 
 namespace tensorflow {
 
-TensorResponse::Source::~Source() {}
+TensorResponse::Source::~Source() = default;
 
 void TensorResponse::Clear() {
   on_host_ = false;
diff --git a/tensorflow/core/distributed_runtime/tensor_coding.h b/tensorflow/core/distributed_runtime/tensor_coding.h
index 703ce7961c3..68019742da3 100644
--- a/tensorflow/core/distributed_runtime/tensor_coding.h
+++ b/tensorflow/core/distributed_runtime/tensor_coding.h
@@ -33,7 +33,7 @@ class TensorProto;
 // into Tensor contents as well as associated metadata.
 class TensorResponse {
  public:
-  TensorResponse() {}
+  TensorResponse() = default;
 
   // Reset to initial state.
   void Clear();
diff --git a/tensorflow/core/distributed_runtime/worker_cache.h b/tensorflow/core/distributed_runtime/worker_cache.h
index 0612a8321d3..3858529b9c6 100644
--- a/tensorflow/core/distributed_runtime/worker_cache.h
+++ b/tensorflow/core/distributed_runtime/worker_cache.h
@@ -33,7 +33,7 @@ class StepStats;
 
 class WorkerCacheInterface {
  public:
-  virtual ~WorkerCacheInterface() {}
+  virtual ~WorkerCacheInterface() = default;
 
   // Updates *workers with strings naming the remote worker tasks to
   // which open channels have been established.
diff --git a/tensorflow/core/distributed_runtime/worker_cache_partial.h b/tensorflow/core/distributed_runtime/worker_cache_partial.h
index 08e272a3bb6..b6352c827f0 100644
--- a/tensorflow/core/distributed_runtime/worker_cache_partial.h
+++ b/tensorflow/core/distributed_runtime/worker_cache_partial.h
@@ -38,7 +38,7 @@ class WorkerCachePartial : public WorkerCacheInterface {
                               DeviceLocality* locality,
                               StatusCallback) override;
 
-  ~WorkerCachePartial() override {}
+  ~WorkerCachePartial() override = default;
 
   // Clear all entries from the DeviceStatus cache.
   void FlushStatusCache();
diff --git a/tensorflow/core/distributed_runtime/worker_interface.h b/tensorflow/core/distributed_runtime/worker_interface.h
index 9476ad93618..e87a16e59ee 100644
--- a/tensorflow/core/distributed_runtime/worker_interface.h
+++ b/tensorflow/core/distributed_runtime/worker_interface.h
@@ -192,7 +192,7 @@ class WorkerInterface {
  protected:
   // Instances of WorkerInterface must be deleted by a call to
   // WorkerCacheInterface::ReleaseWorker().
-  virtual ~WorkerInterface() {}
+  virtual ~WorkerInterface() = default;
   friend class WorkerCacheInterface;
 
   // NOTE: This should only be called by implementations of this

From 6ddfb5ad0e3f0cc790edb9760a79457108787e70 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 20:46:17 -0700
Subject: [PATCH 0453/1000] Automated Code Change

PiperOrigin-RevId: 997287034
---
 tensorflow/core/common_runtime/collective_executor_mgr.h | 2 +-
 tensorflow/core/common_runtime/process_state.h           | 2 +-
 tensorflow/core/common_runtime/scoped_allocator_mgr.h    | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/common_runtime/collective_executor_mgr.h b/tensorflow/core/common_runtime/collective_executor_mgr.h
index bb676d3d6f5..5e3ae40b5b0 100644
--- a/tensorflow/core/common_runtime/collective_executor_mgr.h
+++ b/tensorflow/core/common_runtime/collective_executor_mgr.h
@@ -31,7 +31,7 @@ class CollectiveExecutorMgr : public CollectiveExecutorMgrInterface {
       std::unique_ptr<ParamResolverInterface> param_resolver,
       std::unique_ptr<NcclCommunicatorInterface> nccl_communicator);
 
-  virtual ~CollectiveExecutorMgr();
+  ~CollectiveExecutorMgr() override;
 
   CollectiveExecutor* FindOrCreate(int64_t step_id) override;
 
diff --git a/tensorflow/core/common_runtime/process_state.h b/tensorflow/core/common_runtime/process_state.h
index 09b5f28ff06..e2ccd4e3382 100644
--- a/tensorflow/core/common_runtime/process_state.h
+++ b/tensorflow/core/common_runtime/process_state.h
@@ -79,7 +79,7 @@ class ProcessState : public ProcessStateInterface {
 
  protected:
   ProcessState();
-  virtual ~ProcessState() {}
+  ~ProcessState() override {}
   friend class GPUProcessState;
   friend class PluggableDeviceProcessState;
 
diff --git a/tensorflow/core/common_runtime/scoped_allocator_mgr.h b/tensorflow/core/common_runtime/scoped_allocator_mgr.h
index 070e5191a5e..a69a295e3ba 100644
--- a/tensorflow/core/common_runtime/scoped_allocator_mgr.h
+++ b/tensorflow/core/common_runtime/scoped_allocator_mgr.h
@@ -50,7 +50,7 @@ class ScopedAllocatorContainer : public core::RefCounted {
   friend class ScopedAllocatorMgr;
   ScopedAllocatorContainer(const ScopedAllocatorMgr* mgr, int64_t step_id)
       : mgr_(mgr), step_id_(step_id) {}
-  ~ScopedAllocatorContainer();
+  ~ScopedAllocatorContainer() override;
 
  private:
   const ScopedAllocatorMgr* mgr_;

From 803adb83ba3a65e10e20d0981c39a44ff5bdeda3 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 20:47:23 -0700
Subject: [PATCH 0454/1000] Automated Code Change

PiperOrigin-RevId: 997287249
---
 .../tools/evaluation/tasks/coco_object_detection/run_eval.cc    | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/lite/tools/evaluation/tasks/coco_object_detection/run_eval.cc b/tensorflow/lite/tools/evaluation/tasks/coco_object_detection/run_eval.cc
index 45043c88572..4e0f4993dfd 100644
--- a/tensorflow/lite/tools/evaluation/tasks/coco_object_detection/run_eval.cc
+++ b/tensorflow/lite/tools/evaluation/tasks/coco_object_detection/run_eval.cc
@@ -51,7 +51,7 @@ std::string GetNameFromPath(const std::string& str) {
 class CocoObjectDetection : public TaskExecutor {
  public:
   CocoObjectDetection() : debug_mode_(false), num_interpreter_threads_(1) {}
-  ~CocoObjectDetection() override {}
+  ~CocoObjectDetection() override = default;
 
  protected:
   std::vector<Flag> GetFlags() final;

From 13ccd6a8762b1577eaf9298b624937b6d92cc705 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 20:51:03 -0700
Subject: [PATCH 0455/1000] Automated Code Change

PiperOrigin-RevId: 997288024
---
 .../core/distributed_runtime/rpc_collective_executor_mgr.h    | 4 ++--
 tensorflow/core/distributed_runtime/test_utils.h              | 2 +-
 tensorflow/core/distributed_runtime/worker.h                  | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/distributed_runtime/rpc_collective_executor_mgr.h b/tensorflow/core/distributed_runtime/rpc_collective_executor_mgr.h
index aadbaf33796..b832631619e 100644
--- a/tensorflow/core/distributed_runtime/rpc_collective_executor_mgr.h
+++ b/tensorflow/core/distributed_runtime/rpc_collective_executor_mgr.h
@@ -41,7 +41,7 @@ class RpcCollectiveExecutorMgr : public CollectiveExecutorMgr {
       std::unique_ptr<NcclCommunicatorInterface> nccl_communicator,
       WorkerCacheInterface* worker_cache, const std::string& task_name);
 
-  virtual ~RpcCollectiveExecutorMgr();
+  ~RpcCollectiveExecutorMgr() override;
 
   // This function should only be called at the group_leader, by an RPC.
   // Other needs for StepIds should be satisfied by NextStepId.
@@ -57,7 +57,7 @@ class RpcCollectiveExecutorMgr : public CollectiveExecutorMgr {
   void RetireStepId(int64_t graph_key, int64_t step_id) override;
 
  protected:
-  virtual CollectiveExecutor* Create(int64_t step_id) override;
+  CollectiveExecutor* Create(int64_t step_id) override;
 
   WorkerCacheInterface* const worker_cache_;  // Not owned.
   const std::string task_name_;
diff --git a/tensorflow/core/distributed_runtime/test_utils.h b/tensorflow/core/distributed_runtime/test_utils.h
index 1db7cfb7da7..eab656bdf81 100644
--- a/tensorflow/core/distributed_runtime/test_utils.h
+++ b/tensorflow/core/distributed_runtime/test_utils.h
@@ -122,7 +122,7 @@ class TestWorkerInterface : public WorkerInterface {
 
 class TestWorkerCache : public WorkerCacheInterface {
  public:
-  virtual ~TestWorkerCache() {}
+  ~TestWorkerCache() override {}
 
   void AddWorker(const std::string& target, WorkerInterface* wi) {
     workers_[target] = wi;
diff --git a/tensorflow/core/distributed_runtime/worker.h b/tensorflow/core/distributed_runtime/worker.h
index 4c55e1b9612..02d343f6745 100644
--- a/tensorflow/core/distributed_runtime/worker.h
+++ b/tensorflow/core/distributed_runtime/worker.h
@@ -43,7 +43,7 @@ class WorkerSession;
 class Worker : public WorkerInterface {
  public:
   Worker(WorkerEnv* env);
-  virtual ~Worker() {}
+  ~Worker() override {}
 
   void GetStatusAsync(CallOptions* opts, const GetStatusRequest* request,
                       GetStatusResponse* response, bool fail_fast,

From d544e1123e469280b8fdf915f42ccc626af0e793 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 20:58:12 -0700
Subject: [PATCH 0456/1000] Automated Code Change

PiperOrigin-RevId: 997289626
---
 tensorflow/core/util/gpu_kernel_helper_test.cu.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/util/gpu_kernel_helper_test.cu.cc b/tensorflow/core/util/gpu_kernel_helper_test.cu.cc
index 9af28ae163b..b2b61fb4db8 100644
--- a/tensorflow/core/util/gpu_kernel_helper_test.cu.cc
+++ b/tensorflow/core/util/gpu_kernel_helper_test.cu.cc
@@ -170,7 +170,7 @@ class GpuLaunchConfigTest : public ::testing::Test {
     ASSERT_EQ(cudaSuccess, err) << cudaGetErrorString(err);
 #endif
   }
-  virtual void SetUp() {
+  void SetUp() override {
 #if GOOGLE_CUDA
     cudaError_t err = cudaMallocManaged(&outbuf, sizeof(int) * bufsize);
     outbuf_host = outbuf;
@@ -181,7 +181,7 @@ class GpuLaunchConfigTest : public ::testing::Test {
     ASSERT_EQ(cudaSuccess, err) << cudaGetErrorString(err);
   }
 
-  virtual void TearDown() {
+  void TearDown() override {
     ASSERT_EQ(gpuDeviceSynchronize(), cudaSuccess);
     ASSERT_EQ(gpuFree(outbuf), cudaSuccess);
     outbuf = nullptr;

From 8f4ea1d313bd8ab5bff3e8c1ff49f15d2cc593d1 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 21:10:14 -0700
Subject: [PATCH 0457/1000] Automated Code Change

PiperOrigin-RevId: 997292484
---
 .../distributed_runtime/rpc/eager/grpc_eager_service_impl.h     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_service_impl.h b/tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_service_impl.h
index 62ee6e9f13a..268a40c0c0c 100644
--- a/tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_service_impl.h
+++ b/tensorflow/core/distributed_runtime/rpc/eager/grpc_eager_service_impl.h
@@ -47,7 +47,7 @@ class GrpcEagerServiceImpl : public tsl::AsyncServiceInterface {
                                             RequestMessage, ResponseMessage>;
 
   GrpcEagerServiceImpl(WorkerEnv* env, ::grpc::ServerBuilder* server_builder);
-  virtual ~GrpcEagerServiceImpl() {}
+  ~GrpcEagerServiceImpl() override {}
 
   // Create a master context in eager service.
   absl::Status CreateMasterContext(uint64_t context_id, EagerContext* context);

From b83b7446bd989b2db45e1e5d09ac40e2d79f591f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 21:28:55 -0700
Subject: [PATCH 0458/1000] Automated Code Change

PiperOrigin-RevId: 997297352
---
 .../lite/tools/evaluation/stages/image_preprocessing_stage.h    | 2 +-
 .../lite/tools/evaluation/stages/tflite_inference_stage.h       | 2 +-
 .../lite/tools/evaluation/stages/topk_accuracy_eval_stage.h     | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/lite/tools/evaluation/stages/image_preprocessing_stage.h b/tensorflow/lite/tools/evaluation/stages/image_preprocessing_stage.h
index 39ae99686c7..02172f2874d 100644
--- a/tensorflow/lite/tools/evaluation/stages/image_preprocessing_stage.h
+++ b/tensorflow/lite/tools/evaluation/stages/image_preprocessing_stage.h
@@ -46,7 +46,7 @@ class ImagePreprocessingStage : public EvaluationStage {
 
   EvaluationStageMetrics LatestMetrics() override;
 
-  ~ImagePreprocessingStage() override {}
+  ~ImagePreprocessingStage() override = default;
 
   // Call before Run().
   void SetImagePath(std::string* image_path) { image_path_ = image_path; }
diff --git a/tensorflow/lite/tools/evaluation/stages/tflite_inference_stage.h b/tensorflow/lite/tools/evaluation/stages/tflite_inference_stage.h
index 688b4cbc908..6526a4c2e21 100644
--- a/tensorflow/lite/tools/evaluation/stages/tflite_inference_stage.h
+++ b/tensorflow/lite/tools/evaluation/stages/tflite_inference_stage.h
@@ -51,7 +51,7 @@ class TfliteInferenceStage : public EvaluationStage {
   // EvaluationStageMetrics.num_runs denotes the number of inferences run.
   EvaluationStageMetrics LatestMetrics() override;
 
-  ~TfliteInferenceStage() override {}
+  ~TfliteInferenceStage() override = default;
 
   // Call before Run().
   // This class does not take ownership of raw_input_ptrs.
diff --git a/tensorflow/lite/tools/evaluation/stages/topk_accuracy_eval_stage.h b/tensorflow/lite/tools/evaluation/stages/topk_accuracy_eval_stage.h
index d6da049380d..5fd51f8714d 100644
--- a/tensorflow/lite/tools/evaluation/stages/topk_accuracy_eval_stage.h
+++ b/tensorflow/lite/tools/evaluation/stages/topk_accuracy_eval_stage.h
@@ -42,7 +42,7 @@ class TopkAccuracyEvalStage : public EvaluationStage {
 
   EvaluationStageMetrics LatestMetrics() override;
 
-  ~TopkAccuracyEvalStage() override {}
+  ~TopkAccuracyEvalStage() override = default;
 
   // Call before Init().
   // model_output_shape is not owned, so this class does not free the

From 8400cc834adeea02ff0d17b74a96c42b12052e46 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 21:30:13 -0700
Subject: [PATCH 0459/1000] Automated Code Change

PiperOrigin-RevId: 997297721
---
 tensorflow/core/kernels/conv_ops.h                     | 4 ++--
 tensorflow/core/kernels/dynamic_partition_op_gpu.cu.cc | 2 +-
 tensorflow/core/kernels/l2loss_op.h                    | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/kernels/conv_ops.h b/tensorflow/core/kernels/conv_ops.h
index ebbfa545d05..7a3dde5e558 100644
--- a/tensorflow/core/kernels/conv_ops.h
+++ b/tensorflow/core/kernels/conv_ops.h
@@ -80,12 +80,12 @@ struct Im2ColBufferResource : public ResourceBase {
   Im2ColBufferResource<T, size>() {
     data = static_cast<T*>(port::Malloc(size * sizeof(T)));
   }
-  ~Im2ColBufferResource<T, size>() { port::Free(data); }
+  ~Im2ColBufferResource<T, size>() override { port::Free(data); }
   // This mutex ensures that only a single operation at a time is able to use
   // the buffer memory held by this resource.
   mutex mu;
   T* data;
-  std::string DebugString() const { return "Im2ColBufferResource"; }
+  std::string DebugString() const override { return "Im2ColBufferResource"; }
 };
 
 // Convolution parameters specified by Op attributes.
diff --git a/tensorflow/core/kernels/dynamic_partition_op_gpu.cu.cc b/tensorflow/core/kernels/dynamic_partition_op_gpu.cu.cc
index fb88e09e5e3..37887759784 100644
--- a/tensorflow/core/kernels/dynamic_partition_op_gpu.cu.cc
+++ b/tensorflow/core/kernels/dynamic_partition_op_gpu.cu.cc
@@ -220,7 +220,7 @@ class DynamicPartitionOpGPU : public AsyncOpKernel {
     }
   }
 
-  void ComputeAsync(OpKernelContext* c, DoneCallback done) {
+  void ComputeAsync(OpKernelContext* c, DoneCallback done) override {
     const Tensor& data = c->input(0);
     const Tensor& partitions = c->input(1);
 
diff --git a/tensorflow/core/kernels/l2loss_op.h b/tensorflow/core/kernels/l2loss_op.h
index 2adaacbb776..430b32d95ff 100644
--- a/tensorflow/core/kernels/l2loss_op.h
+++ b/tensorflow/core/kernels/l2loss_op.h
@@ -25,7 +25,7 @@ template <typename Device, typename T>
 struct L2LossOp : public OpKernel {
   explicit L2LossOp(OpKernelConstruction* context) : OpKernel(context) {}
 
-  void Compute(OpKernelContext* context) {}
+  void Compute(OpKernelContext* context) override {}
 };
 
 }  // namespace tensorflow

From 104bfb79db8f6d5433e0c846c30a686d110e377f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 21:30:15 -0700
Subject: [PATCH 0460/1000] Automated Code Change

PiperOrigin-RevId: 997297730
---
 tensorflow/core/nccl/nccl_manager.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/nccl/nccl_manager.cc b/tensorflow/core/nccl/nccl_manager.cc
index e7271cf82b6..e3d63c6f2a6 100644
--- a/tensorflow/core/nccl/nccl_manager.cc
+++ b/tensorflow/core/nccl/nccl_manager.cc
@@ -67,7 +67,7 @@ int NcclManager::instance_count = 0;
 struct NcclManager::NcclStream : public core::RefCounted {
  public:
   NcclStream() = default;
-  ~NcclStream() = default;
+  ~NcclStream() override = default;
 
   se::StreamExecutor* executor = nullptr;
 

From 9fa567ed2df4b579f2e2bcd2ca9a9cea8054bc4c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 21:31:29 -0700
Subject: [PATCH 0461/1000] Automated Code Change

PiperOrigin-RevId: 997298060
---
 tensorflow/core/common_runtime/function_optimization_registry.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/common_runtime/function_optimization_registry.h b/tensorflow/core/common_runtime/function_optimization_registry.h
index ba501d3eba4..ec8ef3f85c2 100644
--- a/tensorflow/core/common_runtime/function_optimization_registry.h
+++ b/tensorflow/core/common_runtime/function_optimization_registry.h
@@ -51,7 +51,7 @@ class FunctionOptimizationPass {
     bool allow_soft_placement = false;
   };
 
-  virtual ~FunctionOptimizationPass() {}
+  virtual ~FunctionOptimizationPass() = default;
   virtual absl::Status Run(const std::string& function_name,
                            const DeviceSet& device_set,
                            const ConfigProto& config_proto,

From 4a8d03a5f3fa52e8af10077b5f196c2fb516d28f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sat, 10 Oct 2026 22:07:01 -0700
Subject: [PATCH 0462/1000] Automated Code Change

PiperOrigin-RevId: 997307474
---
 tensorflow/core/kernels/tile_functor_gpu_complex64.cu.cc | 1 +
 tensorflow/core/kernels/tile_functor_gpu_double.cu.cc    | 1 +
 tensorflow/core/kernels/tile_functor_gpu_float.cu.cc     | 1 +
 tensorflow/core/kernels/tile_functor_gpu_half.cu.cc      | 1 +
 tensorflow/core/kernels/training_ops.cc                  | 1 +
 tensorflow/core/kernels/training_ops.h                   | 1 +
 6 files changed, 6 insertions(+)

diff --git a/tensorflow/core/kernels/tile_functor_gpu_complex64.cu.cc b/tensorflow/core/kernels/tile_functor_gpu_complex64.cu.cc
index eade3aea6f8..75e67a5ba9f 100644
--- a/tensorflow/core/kernels/tile_functor_gpu_complex64.cu.cc
+++ b/tensorflow/core/kernels/tile_functor_gpu_complex64.cu.cc
@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <cstdint>
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 #define EIGEN_USE_GPU
 
diff --git a/tensorflow/core/kernels/tile_functor_gpu_double.cu.cc b/tensorflow/core/kernels/tile_functor_gpu_double.cu.cc
index a0b1a289e5f..565ef18e664 100644
--- a/tensorflow/core/kernels/tile_functor_gpu_double.cu.cc
+++ b/tensorflow/core/kernels/tile_functor_gpu_double.cu.cc
@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <cstdint>
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 #define EIGEN_USE_GPU
 
diff --git a/tensorflow/core/kernels/tile_functor_gpu_float.cu.cc b/tensorflow/core/kernels/tile_functor_gpu_float.cu.cc
index 5fc9142514a..f0e5692ed11 100644
--- a/tensorflow/core/kernels/tile_functor_gpu_float.cu.cc
+++ b/tensorflow/core/kernels/tile_functor_gpu_float.cu.cc
@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <cstdint>
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 #define EIGEN_USE_GPU
 
diff --git a/tensorflow/core/kernels/tile_functor_gpu_half.cu.cc b/tensorflow/core/kernels/tile_functor_gpu_half.cu.cc
index 9fc577d8bde..8967ab863a6 100644
--- a/tensorflow/core/kernels/tile_functor_gpu_half.cu.cc
+++ b/tensorflow/core/kernels/tile_functor_gpu_half.cu.cc
@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <cstdint>
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 #define EIGEN_USE_GPU
 
diff --git a/tensorflow/core/kernels/training_ops.cc b/tensorflow/core/kernels/training_ops.cc
index ef7bd22e57f..003ebb909a1 100644
--- a/tensorflow/core/kernels/training_ops.cc
+++ b/tensorflow/core/kernels/training_ops.cc
@@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <type_traits>
 #define EIGEN_USE_THREADS
 #include "tensorflow/core/kernels/training_ops.h"
 
diff --git a/tensorflow/core/kernels/training_ops.h b/tensorflow/core/kernels/training_ops.h
index 8f986d13a9a..382957cea12 100644
--- a/tensorflow/core/kernels/training_ops.h
+++ b/tensorflow/core/kernels/training_ops.h
@@ -16,6 +16,7 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_KERNELS_TRAINING_OPS_H_
 #define TENSORFLOW_CORE_KERNELS_TRAINING_OPS_H_
 
+#include "absl/status/status.h"
 #include "unsupported/Eigen/CXX11/Tensor"  // from @eigen_archive
 #include "tensorflow/core/framework/tensor_types.h"
 #include "tensorflow/core/lib/core/status.h"

From 94abc44740f5e2280ae0c71215f320859a33082c Mon Sep 17 00:00:00 2001
From: nishad shabbir <nishadshabbir28@gmail.com>
Date: Sun, 11 Oct 2026 13:15:35 +0530
Subject: [PATCH 0463/1000] Put the ops hessians() adds under its name argument

tf.hessians documents name as "Optional name to use for grouping all the
gradient ops together. defaults to 'hessians'", and HessiansV2 forwards
the caller's value into it, but nothing in the body ever reads it. There
has never been a name_scope here, so passing name has no effect at all:
the graph comes out identical either way, and even the documented default
scope is not created, so the ops land in whatever scope happens to be
current.

Wrap the body in ops.name_scope(name, "hessians", ...), the same way
_GradientsHelper() does for tf.gradients. The ops are unchanged, they
just get the prefix the docstring promises instead of leaking into the
enclosing scope.
---
 tensorflow/python/ops/gradients_impl.py | 50 +++++++++++++------------
 tensorflow/python/ops/gradients_test.py | 18 +++++++++
 2 files changed, 45 insertions(+), 23 deletions(-)

diff --git a/tensorflow/python/ops/gradients_impl.py b/tensorflow/python/ops/gradients_impl.py
index 14e316a6433..ad0b40c4847 100644
--- a/tensorflow/python/ops/gradients_impl.py
+++ b/tensorflow/python/ops/gradients_impl.py
@@ -417,7 +417,10 @@ def hessians(ys,
     LookupError: if one of the operations between `xs` and `ys` does not
       have a registered gradient function.
   """
-  xs = gradients_util._AsList(xs)  # pylint: disable=protected-access
+  # pylint: disable=protected-access
+  ys_list = gradients_util._AsList(ys)
+  xs = gradients_util._AsList(xs)
+  # pylint: enable=protected-access
   kwargs = {
       "colocate_gradients_with_ops": colocate_gradients_with_ops,
       "gate_gradients": gate_gradients,
@@ -425,30 +428,31 @@ def hessians(ys,
   }
   # Compute first-order derivatives and iterate for each x in xs.
   hessians = []
-  _gradients = gradients(ys, xs, **kwargs)
-  for gradient, x in zip(_gradients, xs):
-    # change shape to one-dimension without graph branching
-    gradient = array_ops.reshape(gradient, [-1])
+  with ops.name_scope(name, "hessians", list(ys_list) + list(xs)):
+    _gradients = gradients(ys, xs, **kwargs)
+    for gradient, x in zip(_gradients, xs):
+      # change shape to one-dimension without graph branching
+      gradient = array_ops.reshape(gradient, [-1])
 
-    # Declare an iterator and tensor array loop variables for the gradients.
-    n = array_ops.size(x)
-    loop_vars = [
-        array_ops.constant(0, dtypes.int32),
-        tensor_array_ops.TensorArray(x.dtype, n)
-    ]
-    # Iterate over all elements of the gradient and compute second order
-    # derivatives.
-    _, hessian = while_loop.while_loop(
-        lambda j, _: j < n,
-        lambda j, result: (j + 1,
-                           result.write(j, gradients(gradient[j], x)[0])),
-        loop_vars
-    )
+      # Declare an iterator and tensor array loop variables for the gradients.
+      n = array_ops.size(x)
+      loop_vars = [
+          array_ops.constant(0, dtypes.int32),
+          tensor_array_ops.TensorArray(x.dtype, n)
+      ]
+      # Iterate over all elements of the gradient and compute second order
+      # derivatives.
+      _, hessian = while_loop.while_loop(
+          lambda j, _: j < n,
+          lambda j, result: (j + 1,
+                             result.write(j, gradients(gradient[j], x)[0])),
+          loop_vars
+      )
 
-    _shape = array_ops.shape(x)
-    _reshaped_hessian = array_ops.reshape(hessian.stack(),
-                                          array_ops.concat((_shape, _shape), 0))
-    hessians.append(_reshaped_hessian)
+      _shape = array_ops.shape(x)
+      _reshaped_hessian = array_ops.reshape(
+          hessian.stack(), array_ops.concat((_shape, _shape), 0))
+      hessians.append(_reshaped_hessian)
   return hessians
 
 
diff --git a/tensorflow/python/ops/gradients_test.py b/tensorflow/python/ops/gradients_test.py
index 05f3093bef3..e99e4827197 100644
--- a/tensorflow/python/ops/gradients_test.py
+++ b/tensorflow/python/ops/gradients_test.py
@@ -862,6 +862,24 @@ class HessianTest(test_util.TensorFlowTestCase):
     self.assertAllEqual((m, n, m, n), hess_actual.shape)
     self.assertAllClose(hess_value, hess_actual.reshape((m * n, m * n)))
 
+  def testHessianNameScope(self):
+    # `name` is documented as grouping all the gradient ops together, so the
+    # ops hessians() adds must land under it rather than in the enclosing
+    # scope.
+    for name, expected in [(None, "hessians"), ("my_hessians", "my_hessians")]:
+      with ops.Graph().as_default() as g:
+        x = constant_op.constant([1.0, 2.0, 3.0])
+        y = math_ops.reduce_sum(x * x * x)
+        if name is None:
+          gradients.hessians(y, x)
+        else:
+          gradients.hessians(y, x, name=name)
+      added = [op.name for op in g.get_operations()
+               if not op.name.startswith(("Const", "mul", "Sum"))]
+      self.assertNotEmpty(added)
+      for op_name in added:
+        self.assertEqual(expected, op_name.split("/")[0])
+
 
 class IndexedSlicesToTensorTest(test_util.TensorFlowTestCase):
 

From 9a33f0b019ea0280d253256d62f45f88e4481773 Mon Sep 17 00:00:00 2001
From: Aditya Gaurav <mr.x.l.r.8.pride@gmail.com>
Date: Tue, 29 Sep 2026 22:38:13 +0000
Subject: [PATCH 0464/1000] fix(gpu): compute RollKernel launch config before
 allocating buffers

---
 tensorflow/core/kernels/roll_op_gpu.cu.cc | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/kernels/roll_op_gpu.cu.cc b/tensorflow/core/kernels/roll_op_gpu.cu.cc
index 801873d4372..7364aaf7394 100644
--- a/tensorflow/core/kernels/roll_op_gpu.cu.cc
+++ b/tensorflow/core/kernels/roll_op_gpu.cu.cc
@@ -62,6 +62,10 @@ struct Roll<GPUDevice, T> {
     if (!num_elements) return;
     const GPUDevice& d = context->eigen_device<GPUDevice>();
 
+    auto config_or = GetGpuLaunchConfig64(num_elements, d);
+    CHECK_OK(config_or.status());  // Crash OK
+    const GpuLaunchConfig64& cfg = *config_or;
+
     auto dim_bytes = sizeof(int32_t) * dim_size.size();
     auto dim_buf = d.allocate(dim_bytes);
 
@@ -75,10 +79,6 @@ struct Roll<GPUDevice, T> {
     d.memcpyHostToDevice(thres_buf, threshold.data(), thres_bytes);
     d.memcpyHostToDevice(range_buf, dim_range.data(), range_bytes);
 
-    auto config_or = GetGpuLaunchConfig64(num_elements, d);
-    CHECK_OK(config_or.status());  // Crash OK
-    const GpuLaunchConfig64& cfg = *config_or;
-
     TF_CHECK_OK(
         GpuLaunchKernel(RollKernel<T>, cfg.block_count, cfg.thread_per_block, 0,
                         d.stream(), cfg.virtual_thread_count, num_dims, input,

From ad418283d9ccf2d90d2d969481b5f2adfc794c7a Mon Sep 17 00:00:00 2001
From: Aditya Gaurav <mr.x.l.r.8.pride@gmail.com>
Date: Fri, 2 Oct 2026 18:52:25 +0000
Subject: [PATCH 0465/1000] fix(gpu): propagate RollKernel launch errors with
 OP_REQUIRES_OK

---
 tensorflow/core/kernels/roll_op.cc        |  2 +-
 tensorflow/core/kernels/roll_op.h         |  2 +-
 tensorflow/core/kernels/roll_op_gpu.cu.cc | 20 ++++++++++++--------
 3 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/tensorflow/core/kernels/roll_op.cc b/tensorflow/core/kernels/roll_op.cc
index 6b2e8b1d4ee..02ea5c89ba9 100644
--- a/tensorflow/core/kernels/roll_op.cc
+++ b/tensorflow/core/kernels/roll_op.cc
@@ -319,7 +319,7 @@ void DoRollWithMemcpy(const OpKernelContext* context,
 
 template <typename T>
 struct Roll<CPUDevice, T> {
-  void operator()(const OpKernelContext* context, const int64_t num_elements,
+  void operator()(OpKernelContext* context, const int64_t num_elements,
                   const int num_dims, const absl::Span<const int32_t> dim_size,
                   const T* input, T* output,
                   const absl::Span<const int32_t> threshold,
diff --git a/tensorflow/core/kernels/roll_op.h b/tensorflow/core/kernels/roll_op.h
index 9da44aeeba2..b35a87e531d 100644
--- a/tensorflow/core/kernels/roll_op.h
+++ b/tensorflow/core/kernels/roll_op.h
@@ -33,7 +33,7 @@ struct Roll {
   // threshold - the index for each dimension that the roll starts to wrap
   //    back to the front
   // isd - inner shift dimension
-  void operator()(const OpKernelContext* context, const int64_t num_elements,
+  void operator()(OpKernelContext* context, const int64_t num_elements,
                   const int num_dims, const absl::Span<const int32_t> dim_size,
                   const T* input, T* output,
                   const absl::Span<const int32_t> threshold,
diff --git a/tensorflow/core/kernels/roll_op_gpu.cu.cc b/tensorflow/core/kernels/roll_op_gpu.cu.cc
index 7364aaf7394..f6a1de6e49b 100644
--- a/tensorflow/core/kernels/roll_op_gpu.cu.cc
+++ b/tensorflow/core/kernels/roll_op_gpu.cu.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #define EIGEN_USE_GPU
 
+#include "absl/status/statusor.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/tensor_types.h"
 #include "tensorflow/core/kernels/roll_op.h"
@@ -53,7 +54,7 @@ namespace functor {
 
 template <typename T>
 struct Roll<GPUDevice, T> {
-  void operator()(const OpKernelContext* context, const int64_t num_elements,
+  void operator()(OpKernelContext* context, const int64_t num_elements,
                   const int num_dims, const absl::Span<const int32_t> dim_size,
                   const T* input, T* output,
                   const absl::Span<const int32_t> threshold,
@@ -63,7 +64,7 @@ struct Roll<GPUDevice, T> {
     const GPUDevice& d = context->eigen_device<GPUDevice>();
 
     auto config_or = GetGpuLaunchConfig64(num_elements, d);
-    CHECK_OK(config_or.status());  // Crash OK
+    OP_REQUIRES_OK(context, config_or.status());
     const GpuLaunchConfig64& cfg = *config_or;
 
     auto dim_bytes = sizeof(int32_t) * dim_size.size();
@@ -79,16 +80,19 @@ struct Roll<GPUDevice, T> {
     d.memcpyHostToDevice(thres_buf, threshold.data(), thres_bytes);
     d.memcpyHostToDevice(range_buf, dim_range.data(), range_bytes);
 
-    TF_CHECK_OK(
-        GpuLaunchKernel(RollKernel<T>, cfg.block_count, cfg.thread_per_block, 0,
-                        d.stream(), cfg.virtual_thread_count, num_dims, input,
-                        output, reinterpret_cast<const int32_t*>(dim_buf),
-                        reinterpret_cast<const int32_t*>(thres_buf),
-                        reinterpret_cast<const int64_t*>(range_buf)));
+    absl::Status launch_status = GpuLaunchKernel(
+        RollKernel<T>, cfg.block_count, cfg.thread_per_block, 0,
+        d.stream(), cfg.virtual_thread_count, num_dims, input, output,
+        reinterpret_cast<const int32_t*>(dim_buf),
+        reinterpret_cast<const int32_t*>(thres_buf),
+        reinterpret_cast<const int64_t*>(range_buf));
+>>>>>>> c841948c8a2 (fix(gpu): propagate RollKernel launch errors with OP_REQUIRES_OK)
 
     d.deallocate(dim_buf);
     d.deallocate(thres_buf);
     d.deallocate(range_buf);
+
+    OP_REQUIRES_OK(context, launch_status);
   }
 };
 

From 2a0ecb0c74289476efb2089ccccbebeae8d3a70c Mon Sep 17 00:00:00 2001
From: Aditya Gaurav <mr.x.l.r.8.pride@gmail.com>
Date: Sat, 3 Oct 2026 10:44:13 +0000
Subject: [PATCH 0466/1000] Use int64_t for Roll dim_size, threshold and
 shift_mod_sum

Avoid truncating dimensions larger than INT32_MAX, add the absl status deps and include for roll_op, and cover rolling an empty tensor.
---
 tensorflow/core/kernels/BUILD                 |  2 ++
 tensorflow/core/kernels/roll_op.cc            | 29 ++++++++++---------
 tensorflow/core/kernels/roll_op.h             |  4 +--
 tensorflow/core/kernels/roll_op_gpu.cu.cc     | 24 +++++++--------
 .../kernel_tests/array_ops/manip_ops_test.py  |  6 ++++
 5 files changed, 37 insertions(+), 28 deletions(-)

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 12812fefc24..fdeef8c7024 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -3623,6 +3623,8 @@ tf_kernel_library(
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core/framework:bounds_check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@eigen_archive//:eigen3",
     ],
 )
diff --git a/tensorflow/core/kernels/roll_op.cc b/tensorflow/core/kernels/roll_op.cc
index 02ea5c89ba9..aaad1a31a14 100644
--- a/tensorflow/core/kernels/roll_op.cc
+++ b/tensorflow/core/kernels/roll_op.cc
@@ -69,7 +69,7 @@ class RollOp : public OpKernel {
 
     // if there are any duplicate axes, shift_mod_sum will have the
     // total modulo sum of shifts for each dimension
-    absl::InlinedVector<int32_t, 4> shift_mod_sum(num_dims, 0);
+    absl::InlinedVector<int64_t, 4> shift_mod_sum(num_dims, 0);
     for (int i = 0; i < num_shifts; i++) {
       int axis = axis_flat(i);
       if (axis < 0) {
@@ -78,15 +78,16 @@ class RollOp : public OpKernel {
       OP_REQUIRES(context, FastBoundsCheck(axis, num_dims),
                   absl::InvalidArgumentError(
                       absl::StrCat("axis ", axis, " is out of range")));
-      const int ds = std::max<int>(static_cast<int>(input.dim_size(axis)), 1);
-      const int sum = shift_mod_sum[axis] + static_cast<int>(shift_flat(i));
+      const int64_t ds = std::max<int64_t>(input.dim_size(axis), 1);
+      const int64_t sum =
+          shift_mod_sum[axis] + static_cast<int64_t>(shift_flat(i));
       // modulo that works with negatives: ((x % y) + y) % y
       shift_mod_sum[axis] = (sum % ds + ds) % ds;
     }
     // the size of each dimension
-    absl::InlinedVector<int32_t, 4> dim_size(num_dims);
+    absl::InlinedVector<int64_t, 4> dim_size(num_dims);
     // threshold[i] is the index that the roll starts to wrap back to the front
-    absl::InlinedVector<int32_t, 4> threshold(num_dims);
+    absl::InlinedVector<int64_t, 4> threshold(num_dims);
     // dim_range is the number of indices over in the flattened tensor
     // you need to skip in order to make it over from one side of a dimension
     // to the other. Used to make the shifts wrap around after a threshold.
@@ -96,7 +97,7 @@ class RollOp : public OpKernel {
     int64_t isd = 0;
     for (int i = num_dims - 1; i >= 0; i--) {
       if (isd == 0 && shift_mod_sum[i] != 0) isd = i;
-      const int ds = std::max<int>(static_cast<int>(input.dim_size(i)), 1);
+      const int64_t ds = std::max<int64_t>(input.dim_size(i), 1);
       dim_size[i] = ds;
       threshold[i] = (ds - shift_mod_sum[i]) % ds;
       dim_size_prod *= static_cast<int64_t>(input.dim_size(i));
@@ -125,9 +126,9 @@ namespace functor {
 //    back to the front
 template <typename T>
 void DoRoll(const OpKernelContext* context, const int64_t num_elements,
-            const int num_dims, const absl::Span<const int32_t> dim_size,
+            const int num_dims, const absl::Span<const int64_t> dim_size,
             const T* input, T* output,
-            const absl::Span<const int32_t> threshold,
+            const absl::Span<const int64_t> threshold,
             const absl::Span<const int64_t> dim_range) {
   auto work = [input, output, num_dims, &dim_size, &threshold, &dim_range](
                   int64_t start, int64_t end) {
@@ -191,8 +192,8 @@ template <typename T>
 // Use memcpy to copy memory in groups when the data type supports memcpy
 void DoRollWithMemcpy(const OpKernelContext* context,
                       const int64_t num_elements, const int num_dims,
-                      const absl::Span<const int32_t> dim_size, const T* input,
-                      T* output, const absl::Span<const int32_t> threshold,
+                      const absl::Span<const int64_t> dim_size, const T* input,
+                      T* output, const absl::Span<const int64_t> threshold,
                       const absl::Span<const int64_t> dim_range,
                       const int64_t isd) {
   auto work = [input, output, num_dims, &dim_size, &threshold, &dim_range, isd](
@@ -282,9 +283,9 @@ void DoRollWithMemcpy(const OpKernelContext* context,
       // indices = [i, j, k, l, 0, 0]
       //                      ^isd
       for (int j = isd; j >= 0; j--) {
-        int inc = 1;
+        int64_t inc = 1;
         if (j == isd) inc = isd_indx_skip;
-        const int indx = (indices[j] + inc) % dim_size[j];
+        const int64_t indx = (indices[j] + inc) % dim_size[j];
         indices[j] = indx;
         if (indx != 0) {
           if (indx == threshold[j]) {
@@ -320,9 +321,9 @@ void DoRollWithMemcpy(const OpKernelContext* context,
 template <typename T>
 struct Roll<CPUDevice, T> {
   void operator()(OpKernelContext* context, const int64_t num_elements,
-                  const int num_dims, const absl::Span<const int32_t> dim_size,
+                  const int num_dims, const absl::Span<const int64_t> dim_size,
                   const T* input, T* output,
-                  const absl::Span<const int32_t> threshold,
+                  const absl::Span<const int64_t> threshold,
                   const absl::Span<const int64_t> dim_range,
                   const int64_t isd) {
     if (DataTypeCanUseMemcpy(DataTypeToEnum<T>::v())) {
diff --git a/tensorflow/core/kernels/roll_op.h b/tensorflow/core/kernels/roll_op.h
index b35a87e531d..d5a20c85039 100644
--- a/tensorflow/core/kernels/roll_op.h
+++ b/tensorflow/core/kernels/roll_op.h
@@ -34,9 +34,9 @@ struct Roll {
   //    back to the front
   // isd - inner shift dimension
   void operator()(OpKernelContext* context, const int64_t num_elements,
-                  const int num_dims, const absl::Span<const int32_t> dim_size,
+                  const int num_dims, const absl::Span<const int64_t> dim_size,
                   const T* input, T* output,
-                  const absl::Span<const int32_t> threshold,
+                  const absl::Span<const int64_t> threshold,
                   const absl::Span<const int64_t> dim_range, const int64_t isd);
 };
 
diff --git a/tensorflow/core/kernels/roll_op_gpu.cu.cc b/tensorflow/core/kernels/roll_op_gpu.cu.cc
index f6a1de6e49b..036bda0dd4a 100644
--- a/tensorflow/core/kernels/roll_op_gpu.cu.cc
+++ b/tensorflow/core/kernels/roll_op_gpu.cu.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #define EIGEN_USE_GPU
 
+#include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/tensor_types.h"
@@ -33,16 +34,16 @@ namespace {
 template <typename T>
 __global__ void RollKernel(const int64_t nthreads, const int32_t num_dims,
                            const T* __restrict__ input, T* __restrict__ output,
-                           const int32_t* __restrict__ dim_size,
-                           const int32_t* __restrict__ threshold,
+                           const int64_t* __restrict__ dim_size,
+                           const int64_t* __restrict__ threshold,
                            const int64_t* __restrict__ dim_range) {
   CUDA_1D_KERNEL_LOOP(out_idx, nthreads, int64_t) {
     int64_t offset = 0;
     for (int i = 0; i < num_dims; i++) {
       const int64_t stride = dim_range[i] / dim_size[i];
-      const int shift = dim_size[i] - threshold[i];
-      const int indx = (out_idx / stride) % dim_size[i];
-      const int shifted_indx = (indx + shift) % dim_size[i];
+      const int64_t shift = dim_size[i] - threshold[i];
+      const int64_t indx = (out_idx / stride) % dim_size[i];
+      const int64_t shifted_indx = (indx + shift) % dim_size[i];
       offset += (shifted_indx - indx) * stride;
     }
     output[out_idx + offset] = input[out_idx];
@@ -55,9 +56,9 @@ namespace functor {
 template <typename T>
 struct Roll<GPUDevice, T> {
   void operator()(OpKernelContext* context, const int64_t num_elements,
-                  const int num_dims, const absl::Span<const int32_t> dim_size,
+                  const int num_dims, const absl::Span<const int64_t> dim_size,
                   const T* input, T* output,
-                  const absl::Span<const int32_t> threshold,
+                  const absl::Span<const int64_t> threshold,
                   const absl::Span<const int64_t> dim_range,
                   const int64_t isd) {
     if (!num_elements) return;
@@ -67,10 +68,10 @@ struct Roll<GPUDevice, T> {
     OP_REQUIRES_OK(context, config_or.status());
     const GpuLaunchConfig64& cfg = *config_or;
 
-    auto dim_bytes = sizeof(int32_t) * dim_size.size();
+    auto dim_bytes = sizeof(int64_t) * dim_size.size();
     auto dim_buf = d.allocate(dim_bytes);
 
-    auto thres_bytes = sizeof(int32_t) * threshold.size();
+    auto thres_bytes = sizeof(int64_t) * threshold.size();
     auto thres_buf = d.allocate(thres_bytes);
 
     auto range_bytes = sizeof(int64_t) * dim_range.size();
@@ -83,10 +84,9 @@ struct Roll<GPUDevice, T> {
     absl::Status launch_status = GpuLaunchKernel(
         RollKernel<T>, cfg.block_count, cfg.thread_per_block, 0,
         d.stream(), cfg.virtual_thread_count, num_dims, input, output,
-        reinterpret_cast<const int32_t*>(dim_buf),
-        reinterpret_cast<const int32_t*>(thres_buf),
+        reinterpret_cast<const int64_t*>(dim_buf),
+        reinterpret_cast<const int64_t*>(thres_buf),
         reinterpret_cast<const int64_t*>(range_buf));
->>>>>>> c841948c8a2 (fix(gpu): propagate RollKernel launch errors with OP_REQUIRES_OK)
 
     d.deallocate(dim_buf);
     d.deallocate(thres_buf);
diff --git a/tensorflow/python/kernel_tests/array_ops/manip_ops_test.py b/tensorflow/python/kernel_tests/array_ops/manip_ops_test.py
index 382d3ebc295..4c611130b84 100644
--- a/tensorflow/python/kernel_tests/array_ops/manip_ops_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/manip_ops_test.py
@@ -100,6 +100,12 @@ class RollTest(test_util.TensorFlowTestCase):
     self._testAll(np.zeros([0, 1]), 1, 1)
     self._testAll(np.zeros([1, 0]), 1, 1)
 
+  def testRollEmptyInput(self):
+    with self.cached_session():
+      empty_tensor = array_ops.zeros([0, 5], dtype=dtypes.float32)
+      result = manip_ops.roll(empty_tensor, shift=1, axis=0)
+      self.assertEqual(self.evaluate(result).shape, (0, 5))
+
   @test_util.run_v2_only
   def testLargeInput(self):
     with test_util.force_cpu():

From 7ac5687262ca93e1922b607772c1d6d85416060d Mon Sep 17 00:00:00 2001
From: Aditya Gaurav <mr.x.l.r.8.pride@gmail.com>
Date: Sun, 4 Oct 2026 22:48:18 +0000
Subject: [PATCH 0467/1000] Normalize Roll shifts before summing and look up
 GPU strides

Take each shift modulo the dimension size before adding it to the
per-axis sum, so repeated-axis int64 shifts cannot overflow. Read the
GPU kernel stride from dim_range[i + 1] instead of dividing.

Run testRollEmptyInput in graph and eager modes, and add large
repeated-axis shift tests to roll_op_test.cc and manip_ops_test.py.
---
 tensorflow/core/kernels/roll_op.cc             |  3 ++-
 tensorflow/core/kernels/roll_op_gpu.cu.cc      |  2 +-
 tensorflow/core/kernels/roll_op_test.cc        | 17 +++++++++++++++++
 .../kernel_tests/array_ops/manip_ops_test.py   | 18 ++++++++++++++----
 4 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/tensorflow/core/kernels/roll_op.cc b/tensorflow/core/kernels/roll_op.cc
index aaad1a31a14..bd12d5d914d 100644
--- a/tensorflow/core/kernels/roll_op.cc
+++ b/tensorflow/core/kernels/roll_op.cc
@@ -79,8 +79,9 @@ class RollOp : public OpKernel {
                   absl::InvalidArgumentError(
                       absl::StrCat("axis ", axis, " is out of range")));
       const int64_t ds = std::max<int64_t>(input.dim_size(axis), 1);
+      // reduce the shift first so the sum cannot overflow int64_t
       const int64_t sum =
-          shift_mod_sum[axis] + static_cast<int64_t>(shift_flat(i));
+          shift_mod_sum[axis] + static_cast<int64_t>(shift_flat(i)) % ds;
       // modulo that works with negatives: ((x % y) + y) % y
       shift_mod_sum[axis] = (sum % ds + ds) % ds;
     }
diff --git a/tensorflow/core/kernels/roll_op_gpu.cu.cc b/tensorflow/core/kernels/roll_op_gpu.cu.cc
index 036bda0dd4a..3c0ae060452 100644
--- a/tensorflow/core/kernels/roll_op_gpu.cu.cc
+++ b/tensorflow/core/kernels/roll_op_gpu.cu.cc
@@ -40,7 +40,7 @@ __global__ void RollKernel(const int64_t nthreads, const int32_t num_dims,
   CUDA_1D_KERNEL_LOOP(out_idx, nthreads, int64_t) {
     int64_t offset = 0;
     for (int i = 0; i < num_dims; i++) {
-      const int64_t stride = dim_range[i] / dim_size[i];
+      const int64_t stride = i + 1 < num_dims ? dim_range[i + 1] : 1;
       const int64_t shift = dim_size[i] - threshold[i];
       const int64_t indx = (out_idx / stride) % dim_size[i];
       const int64_t shifted_indx = (indx + shift) % dim_size[i];
diff --git a/tensorflow/core/kernels/roll_op_test.cc b/tensorflow/core/kernels/roll_op_test.cc
index 78b06d6acc3..dde6958897a 100644
--- a/tensorflow/core/kernels/roll_op_test.cc
+++ b/tensorflow/core/kernels/roll_op_test.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 
 #include <functional>
+#include <limits>
 #include <memory>
 
 #include "tensorflow/core/common_runtime/device.h"
@@ -366,6 +367,22 @@ TEST_F(RollOpTest, MultiShifts_TwoD32_NoMemcpy) {
   test::ExpectTensorEqual<tstring>(expected, *GetOutput(0));
 }
 
+TEST_F(RollOpTest, LargeShifts_SameAxis64) {
+  MakeOp(DT_FLOAT, DT_INT64);
+
+  // Feed and run
+  AddInputFromArray<float>(TensorShape({5}), {0, 1, 2, 3, 4});
+  AddInputFromArray<int64_t>(TensorShape({2}),
+                             {1, std::numeric_limits<int64_t>::max()});
+  AddInputFromArray<int64_t>(TensorShape({2}), {0, 0});
+  TF_ASSERT_OK(RunOpKernel());
+
+  // Check the output. The total shift is 2**63, which is 3 modulo 5.
+  Tensor expected(allocator(), DT_FLOAT, TensorShape({5}));
+  test::FillValues<float>(&expected, {2, 3, 4, 0, 1});
+  test::ExpectTensorEqual<float>(expected, *GetOutput(0));
+}
+
 TEST_F(RollOpTest, Error_InputMustBeVectorOrHigher) {
   MakeOp(DT_FLOAT, DT_INT32);
 
diff --git a/tensorflow/python/kernel_tests/array_ops/manip_ops_test.py b/tensorflow/python/kernel_tests/array_ops/manip_ops_test.py
index 4c611130b84..d04fc510437 100644
--- a/tensorflow/python/kernel_tests/array_ops/manip_ops_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/manip_ops_test.py
@@ -100,11 +100,21 @@ class RollTest(test_util.TensorFlowTestCase):
     self._testAll(np.zeros([0, 1]), 1, 1)
     self._testAll(np.zeros([1, 0]), 1, 1)
 
+  @test_util.run_in_graph_and_eager_modes
   def testRollEmptyInput(self):
-    with self.cached_session():
-      empty_tensor = array_ops.zeros([0, 5], dtype=dtypes.float32)
-      result = manip_ops.roll(empty_tensor, shift=1, axis=0)
-      self.assertEqual(self.evaluate(result).shape, (0, 5))
+    empty_tensor = array_ops.zeros([0, 5], dtype=dtypes.float32)
+    result = manip_ops.roll(empty_tensor, shift=1, axis=0)
+    self.assertEqual(self.evaluate(result).shape, (0, 5))
+
+  @test_util.run_in_graph_and_eager_modes
+  def testRollLargeShiftsSameAxis(self):
+    np_input = np.arange(5, dtype=np.float32)
+    for shifts in ([1, 2**63 - 1], [2**62] * 3):
+      # Sum the shifts as Python ints, since np.roll would overflow int64.
+      expected = np.roll(np_input, sum(shifts) % 5)
+      result = manip_ops.roll(
+          np_input, np.array(shifts, dtype=np.int64), [0] * len(shifts))
+      self.assertAllEqual(self.evaluate(result), expected)
 
   @test_util.run_v2_only
   def testLargeInput(self):

From 544e74b919158dabdd68d54d67c68907457a5bdc Mon Sep 17 00:00:00 2001
From: Javid Khan <dxbjavid@gmail.com>
Date: Sun, 11 Oct 2026 17:27:59 +0530
Subject: [PATCH 0468/1000] check weights length in sparse and ragged bincount

---
 tensorflow/core/kernels/bincount_op.cc        | 23 ++++++++++++++
 .../kernel_tests/math_ops/bincount_op_test.py | 30 +++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/tensorflow/core/kernels/bincount_op.cc b/tensorflow/core/kernels/bincount_op.cc
index 86dce6aa123..e660801addc 100644
--- a/tensorflow/core/kernels/bincount_op.cc
+++ b/tensorflow/core/kernels/bincount_op.cc
@@ -410,6 +410,18 @@ class SparseBincountOp : public OpKernel {
                             indices, values, dense_shape,
                             sparse_utils::IndexValidation::kUnordered));
 
+    // The 2-D path below reads `weights(i)` with the value index `i`, so a
+    // `weights` shorter than `values` would read past its buffer. The 1-D
+    // path and DenseBincount already enforce this; mirror it here.
+    OP_REQUIRES(
+        ctx, weights_size == 0 || weights_size == values.NumElements(),
+        absl::InvalidArgumentError(absl::StrCat(
+            "`weights` must be the same length as `values` or a length-0 "
+            "`Tensor`, in which case it acts as all weights equal to 1. Got a "
+            "length-",
+            weights_size, " `weights` for ", values.NumElements(),
+            " values.")));
+
     OP_REQUIRES(ctx, dense_shape.NumElements() > 0,
                 absl::InvalidArgumentError(absl::StrCat(
                     "dense_shape must have at least 1 dimension, got ",
@@ -519,6 +531,17 @@ class RaggedBincountOp : public OpKernel {
                     "Splits must end with the number of values, got ",
                     splits(num_rows), " instead of ", num_values)));
 
+    // The loop below reads `weights(idx)` with the value index `idx`, so a
+    // `weights` shorter than `values` would read past its buffer. The 1-D
+    // path and DenseBincount already enforce this; mirror it here.
+    OP_REQUIRES(
+        ctx, weights_size == 0 || weights_size == num_values,
+        absl::InvalidArgumentError(absl::StrCat(
+            "`weights` must be the same length as `values` or a length-0 "
+            "`Tensor`, in which case it acts as all weights equal to 1. Got a "
+            "length-",
+            weights_size, " `weights` for ", num_values, " values.")));
+
     Tensor* out_t;
     OP_REQUIRES_OK(
         ctx, ctx->allocate_output(0, TensorShape({num_rows, size}), &out_t));
diff --git a/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py b/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py
index 36760fd2227..97f6ba9bc6b 100644
--- a/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py
@@ -780,6 +780,21 @@ class SparseBincountOpTest(test_util.TensorFlowTestCase,
           )
       )
 
+  def test_sparse_bincount_invalid_weights_length(self):
+    # Regression test: the 2-D SparseBincount path reads `weights(i)` with the
+    # value index `i`, so a `weights` shorter than the values would be read out
+    # of bounds. It must be rejected like the 1-D path and DenseBincount do.
+    with self.assertRaisesRegex(
+        (ValueError, errors.InvalidArgumentError),
+        "must be the same length as `values`"):
+      self.evaluate(
+          gen_math_ops.sparse_bincount(
+              indices=[[0, 0], [1, 0], [2, 0]],
+              values=[1, 1, 1],
+              dense_shape=[3, 2],
+              size=2,
+              weights=[1.0]))
+
 
 class RaggedBincountOpTest(test_util.TensorFlowTestCase,
                            parameterized.TestCase):
@@ -945,5 +960,20 @@ class RaggedBincountOpTest(test_util.TensorFlowTestCase,
               binary_output=False,
               name=None))
 
+  @test_util.run_in_graph_and_eager_modes
+  def test_ragged_bincount_invalid_weights_length(self):
+    # Regression test: RaggedBincount reads `weights(idx)` with the value
+    # index `idx`, so a `weights` shorter than the values would be read out of
+    # bounds. It must be rejected like the 1-D path and DenseBincount do.
+    with self.assertRaisesRegex(
+        (ValueError, errors.InvalidArgumentError),
+        "must be the same length as `values`"):
+      self.evaluate(
+          gen_math_ops.ragged_bincount(
+              splits=[0, 3],
+              values=[1, 1, 1],
+              size=2,
+              weights=[1.0]))
+
 if __name__ == "__main__":
   googletest.main()

From 88047023a7d1e7b0caec2627f439428d61c701b6 Mon Sep 17 00:00:00 2001
From: Daniel Liu <liujingtai2001@gmail.com>
Date: Sun, 11 Oct 2026 09:52:30 -0400
Subject: [PATCH 0469/1000] Evaluate the half-precision gradients in SwishTest

The test checked only dtype and shape, so in graph mode the backward pass
never ran. Evaluate both gradients, as requested in review.
---
 tensorflow/python/ops/nn_test.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/python/ops/nn_test.py b/tensorflow/python/ops/nn_test.py
index 2a3e2b6ec28..ebd198d2d96 100644
--- a/tensorflow/python/ops/nn_test.py
+++ b/tensorflow/python/ops/nn_test.py
@@ -1180,6 +1180,7 @@ class SwishTest(test_lib.TestCase, parameterized.TestCase):
       self.assertEqual(g_beta.dtype, dtype)
       self.assertEqual(g_features.shape.as_list(), [2, 2])
       self.assertEqual(g_beta.shape.as_list(), [2])
+      self.evaluate([g_features, g_beta])
 
 
 class MomentsTest(test_lib.TestCase):

From 4ba9e883435d6e010c70128664c34b6a839a2166 Mon Sep 17 00:00:00 2001
From: Saurabh Singh <geniusthegreat0963@gmail.com>
Date: Sun, 11 Oct 2026 19:23:07 +0530
Subject: [PATCH 0470/1000] Fix complex finiteness overflow, graph-mode tests,
 and equal_nan coverage

- Check is_finite on real/imag parts instead of abs() to avoid magnitude
  overflow for large finite complex values (e.g. np.complex64(3e38+3e38j))
- Replace bool(tensor) with self.match() in testIsCloseInfinities for
  graph-mode compatibility
- Parameterize testIsCloseEqualNan over float and complex dtypes with
  NaN in real, imag, and both components
---
 .../python/ops/numpy_ops/np_math_ops.py       | 12 ++++---
 .../python/ops/numpy_ops/np_math_ops_test.py  | 34 ++++++++++++++-----
 2 files changed, 33 insertions(+), 13 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py
index 3fc90a529f9..aa6b3b8e89a 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@@ -611,12 +611,14 @@ def isclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False):  # pylint: disable=m
       # both finite and within tolerance, or if they are equal.
       if dtype.is_floating or dtype.is_complex:
         if dtype.is_complex:
-          a_check = math_ops.abs(a)
-          b_check = b_abs
+          both_finite = (
+              math_ops.is_finite(math_ops.real(a))
+              & math_ops.is_finite(math_ops.imag(a))
+              & math_ops.is_finite(math_ops.real(b))
+              & math_ops.is_finite(math_ops.imag(b))
+          )
         else:
-          a_check = a
-          b_check = b
-        both_finite = math_ops.is_finite(a_check) & math_ops.is_finite(b_check)
+          both_finite = math_ops.is_finite(a) & math_ops.is_finite(b)
         result = (both_finite & result) | math_ops.equal(a, b)
 
       if equal_nan:
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 f570ad39929..e45d50ce6b0 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
@@ -375,10 +375,28 @@ class MathTest(test.TestCase, parameterized.TestCase):
     with self.assertRaisesRegex(ValueError, 'out of bounds'):
       np_math_ops.argmin(a, axis=-3)
 
-  @parameterized.parameters([False, True])
-  def testIsCloseEqualNan(self, equal_nan):
-    a = np.asarray([1, 1, np.nan, 1, np.nan], np.float32)
-    b = np.asarray([1, 2, 1, np.nan, np.nan], np.float32)
+  @parameterized.parameters(
+      itertools.product(
+          [False, True],
+          [np.float32, np.float64, np.complex64, np.complex128]
+      )
+  )
+  def testIsCloseEqualNan(self, equal_nan, dtype):
+    a_vals = [1.0, 1.0, np.nan, 1.0, np.nan]
+    b_vals = [1.0, 2.0, 1.0, np.nan, np.nan]
+    if np.issubdtype(dtype, np.complexfloating):
+      a_vals.extend([
+          complex(1.0, np.nan),
+          complex(np.nan, 1.0),
+          complex(np.nan, np.nan),
+      ])
+      b_vals.extend([
+          complex(np.nan, 1.0),
+          complex(1.0, np.nan),
+          complex(np.nan, np.nan),
+      ])
+    a = np.asarray(a_vals, dtype)
+    b = np.asarray(b_vals, dtype)
     self.match(
         np_math_ops.isclose(a, b, equal_nan=equal_nan),
         np.isclose(a, b, equal_nan=equal_nan))
@@ -402,12 +420,12 @@ class MathTest(test.TestCase, parameterized.TestCase):
         np_math_ops.isclose(a, b),
         np.isclose(a, b),
     )
-    self.assertEqual(
-        bool(np_math_ops.allclose([dtype(np.inf)], [dtype(np.inf)])),
+    self.match(
+        np_math_ops.allclose([dtype(np.inf)], [dtype(np.inf)]),
         np.allclose([dtype(np.inf)], [dtype(np.inf)]),
     )
-    self.assertEqual(
-        bool(np_math_ops.allclose([dtype(-np.inf)], [dtype(np.inf)])),
+    self.match(
+        np_math_ops.allclose([dtype(-np.inf)], [dtype(np.inf)]),
         np.allclose([dtype(-np.inf)], [dtype(np.inf)]),
     )
 

From 4ca59e71fe5f1499ae08c260a6e74931fa118949 Mon Sep 17 00:00:00 2001
From: PrabinDevkota <prabindev789@gmail.com>
Date: Sun, 11 Oct 2026 22:10:07 +0530
Subject: [PATCH 0471/1000] Accept both graph and eager rank-limit error
 messages.

---
 tensorflow/python/kernel_tests/array_ops/array_ops_test.py | 3 ++-
 tensorflow/python/kernel_tests/array_ops/stack_op_test.py  | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
index 58951a7a22b..b4f3ac99d10 100644
--- a/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/array_ops_test.py
@@ -446,7 +446,8 @@ class OperatorShapeTest(test_util.TensorFlowTestCase):
   def testExpandDimsExceedsMaxRank(self):
     value = self._zero_high_rank(254)
     with self.assertRaisesRegex(
-        (errors.InvalidArgumentError, ValueError), "maximum supported rank"):
+        (errors.InvalidArgumentError, ValueError),
+        "maximum supported rank|Too many dimensions"):
       self.evaluate(array_ops.expand_dims(value, 0))
 
 
diff --git a/tensorflow/python/kernel_tests/array_ops/stack_op_test.py b/tensorflow/python/kernel_tests/array_ops/stack_op_test.py
index e2ae6ab4685..8b1832bf33c 100644
--- a/tensorflow/python/kernel_tests/array_ops/stack_op_test.py
+++ b/tensorflow/python/kernel_tests/array_ops/stack_op_test.py
@@ -84,7 +84,8 @@ class StackOpTest(test.TestCase):
         [0] + [1] * 253,
         constant_op.constant(0.0, dtype=dtypes.float32))
     with self.assertRaisesRegex(
-        (errors.InvalidArgumentError, ValueError), "maximum supported rank"):
+        (errors.InvalidArgumentError, ValueError),
+        "maximum supported rank|maximum allowed"):
       self.evaluate(array_ops_stack.stack([value], axis=0))
 
   def testSimpleParallelCPU(self):

From 9cbbfad5aaffe2825a03ab5d3fc6ef55cfbbf3ea Mon Sep 17 00:00:00 2001
From: Saurabh Singh <geniusthegreat0963@gmail.com>
Date: Sun, 11 Oct 2026 22:25:04 +0530
Subject: [PATCH 0472/1000] Support scalar and bool expected values in match
 test helper

np.allclose returns a python bool rather than an ndarray, which lacks
dtype and shape attributes and caused AttributeError in match.
Converting expected with np.asarray ensures match handles both array
and scalar/bool results.
---
 tensorflow/python/ops/numpy_ops/np_math_ops_test.py | 1 +
 1 file changed, 1 insertion(+)

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 e45d50ce6b0..f58808ec7ab 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
@@ -254,6 +254,7 @@ class MathTest(test.TestCase, parameterized.TestCase):
 
   def match(self, actual, expected, msg='', check_dtype=True):
     self.assertIsInstance(actual, np_arrays.ndarray)
+    expected = np.asarray(expected)
     if check_dtype:
       self.assertEqual(
           actual.dtype, expected.dtype,

From c1c3206b5b395872a00d21285d109556a07f6793 Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Sun, 11 Oct 2026 10:56:49 -0700
Subject: [PATCH 0473/1000] # Flip PR: restore rot90 to master and drop dynamic
 rot90 tests

final asks from review 5482731933, which retracts every earlier rot90 and
swapaxes change request: dynamic-rank validation of this kind belongs in
the op kernels rather than in extra Python-side Assert nodes, and any
follow-up rot90/swapaxes hardening belongs in a separate change.

- np_array_ops.py: rot90 returns to the exact master form, including the
  identical np_utils._canonicalize_axes(axes, m_rank) statement; the
  maybe_rank is None Assert/control_dependencies block is removed;
- np_array_ops_test.py: testRot90OutOfBoundsDynamicRank and
  testRot90DuplicateAxesDynamicRank removed; all flip tests kept; the
  one-line self.evaluate(f(x)) change in the existing swapaxes test stays.

flip (axis->array_ops.reverse delegation + 0-D/1-D/unknown-rank OOB tests)
is unaffected and remains the substance of this PR.
---
 .../python/ops/numpy_ops/np_array_ops.py      | 31 +---------
 .../python/ops/numpy_ops/np_array_ops_test.py | 58 -------------------
 2 files changed, 1 insertion(+), 88 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index 03304dc0cbc..2fbfb7ae345 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -1754,36 +1754,7 @@ def rot90(m, k=1, axes=(0, 1)):  # pylint: disable=missing-docstring
           )
 
   m_rank = array_ops.rank(m)
-  ax1, ax2 = np_utils._canonicalize_axes(
-      axes, m_rank)  # pylint: disable=protected-access
-  if maybe_rank is None:
-    # In graph mode, disconnected side-effectful Assert nodes are pruned
-    # unless their results flow into the returned computation, so the
-    # assertions must be wired in via `control_dependencies`.
-    rank_t = ops.convert_to_tensor(m_rank)
-    # Force the axes tensor to the rank tensor's dtype: left to inference, a
-    # tuple containing a value above the int32 range would produce an int64
-    # tensor, and the comparison against `rank_t` would raise a TypeError
-    # instead of triggering the assertions below.
-    axes_t = ops.convert_to_tensor(axes, dtype=rank_t.dtype)
-    assert_bounds = control_flow_assert.Assert(
-        math_ops.reduce_all(
-            math_ops.logical_and(axes_t >= -rank_t, axes_t < rank_t)),
-        ['Axes', axes_t, 'out of range for array of ndim', rank_t],
-    )
-    # The canonicalized axes are tensors here, so the duplicate-axis check
-    # the static path performs natively must be re-done dynamically.
-    assert_diff = control_flow_assert.Assert(
-        math_ops.not_equal(ax1, ax2),
-        ['Axes must be different']
-    )
-    with ops.control_dependencies([assert_bounds, assert_diff]):
-      ax1 = array_ops.identity(ax1)
-      ax2 = array_ops.identity(ax2)
-      # Also thread `m` through the control dependencies: otherwise the
-      # `k % 4 == 0` early return below would return `m` without consuming
-      # the asserts, and graph pruning would disconnect them again.
-      m = array_ops.identity(m)
+  ax1, ax2 = np_utils._canonicalize_axes(axes, m_rank)  # pylint: disable=protected-access
 
   k = k % 4
   if k == 0:
diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index 9af6fe7f74b..4e33e5d25d1 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1609,64 +1609,6 @@ class ArrayMethodsTest(test.TestCase):
     with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
       self.evaluate(flip_unknown_rank_invalid_pos(a))
 
-  def testRot90OutOfBoundsDynamicRank(self):
-    # The dynamic-rank bounds Assert in `rot90` must actually run in graph
-    # mode; it is wired into the computation via `control_dependencies`.
-    # Note: this is deliberately not run with `jit_compile=True` — tf2xla
-    # lowers `Assert` to a no-op under XLA compilation.
-    x = np.zeros((2, 3), dtype=np.float32)
-
-    @def_function.function(
-        input_signature=[
-            tensor_spec.TensorSpec(dtype=dtypes.float32, shape=None)
-        ]
-    )
-    def f_neg(a):
-      return np_array_ops.rot90(a, axes=(0, -5))
-
-    f_neg = def_function.function(
-        f_neg,
-        input_signature=[
-            tensor_spec.TensorSpec(dtype=dtypes.float32, shape=None)
-        ]
-    )
-    with self.assertRaisesRegex(
-        errors_impl.InvalidArgumentError, 'out of range for array of ndim'
-    ):
-      self.evaluate(f_neg(x))
-
-    def f_pos(a):
-      return np_array_ops.rot90(a, axes=(0, 5))
-
-    f_pos = def_function.function(
-        f_pos,
-        input_signature=[
-            tensor_spec.TensorSpec(dtype=dtypes.float32, shape=None)
-        ]
-    )
-    with self.assertRaisesRegex(
-        errors_impl.InvalidArgumentError, 'out of range for array of ndim'
-    ):
-      self.evaluate(f_pos(x))
-
-  def testRot90DuplicateAxesDynamicRank(self):
-    # The dynamic path re-checks for duplicate axes after canonicalization
-    # (e.g. (1, -2) both canonicalize to axis 1 on a rank-3 array).
-    x = np.zeros((2, 3), dtype=np.float32)
-
-    @def_function.function(
-        input_signature=[
-            tensor_spec.TensorSpec(dtype=dtypes.float32, shape=None)
-        ]
-    )
-    def f(a):
-      return np_array_ops.rot90(a, axes=(1, -2))
-
-    with self.assertRaisesRegex(
-        errors_impl.InvalidArgumentError, 'Axes must be different'
-    ):
-      self.evaluate(f(x))
-
   def testNdim(self):
     self.assertAllEqual(0, np_array_ops.ndim(0.5))
     self.assertAllEqual(1, np_array_ops.ndim([1, 2]))

From 85234f43c8d08bee2c89f0ab8d544e019ec9fd46 Mon Sep 17 00:00:00 2001
From: Ce Zheng CA <zce@google.com>
Date: Sun, 11 Oct 2026 14:05:40 -0700
Subject: [PATCH 0474/1000] [XLA] Move type-independent setup for Convolution,
 Dot, RaggedDot, ScaledDot, and Pad out of `HloEvaluatorTypedVisitor`
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

`HloEvaluatorTypedVisitor<ReturnT, ElementwiseT>` is instantiated 32 times across all XLA array primitive types in `xla/hlo/evaluator/`. The outer setup, shape validation, sparse operand materialization, operand type conversion, dimension indexing setup (`GetNonContractingDims`, `DotDimensionNumbers` extraction, `ShapeInfo` construction), and `ElementwiseT == float` `MatmulArray2D` convolution fast path in `HandleConvolution`, `HandleDotSlowPath`, `HandleRaggedDot`, `HandleScaledDot`, and `HandlePad` do not depend on `ReturnT` or `ElementwiseT`, yet were duplicated across all 32 visitor instantiations. In addition, `HandleRaggedDot{NonContracting,BatchMode,Contracting}WithLiterals` copied `RaggedDotDimensionNumbers`, `DotDimensionNumbers`, and `RepeatedField<int64_t>` protobuf fields by value across all 32 instantiations.

Move this type-independent setup and validation into non-template `HloEvaluator::DotDimensionInfo` and `HloEvaluator::Evaluate{Convolution,DotSlowPath,RaggedDot,ScaledDot,Pad}` methods in `xla/hlo/evaluator/hlo_evaluator.cc` taking `absl::FunctionRef` callbacks, while keeping all element-wise inner contraction, convolution, and padding loops (`PopulateParallel<ReturnT>`, `PopulateLinearParallel<ReturnT>`, `ForEachIndexNoStatus`, and `IncrementContractingIndexes`) inline in `HloEvaluatorTypedVisitor`.

### Runtime & Memory Performance Impact

**TL;DR:** Moving the outer setup and validation for `Convolution`, `Dot`, `RaggedDot`, `ScaledDot`, and `Pad` to `HloEvaluator` while keeping all inner element-wise `PopulateParallel<ReturnT>` loops inline is CPU-neutral (`geomean +0.09%` across `n=12` runs, with 17 of 19 benchmarks statistically unchanged `~`) and reduces `libhlo_evaluator.a` (`-c opt`) by `-737,710 B` (`-693,566 B` `.text`, `-44,144 B` `.data`).

- **Microbenchmarks (`//xla/hlo/evaluator:hlo_evaluator_test`, `n=12` runs):**

  ```
  name                                     base cpu/op   after cpu/op    vs base
  BM_ReducePrecisely                       210.2m ± 3%   209.1m ± 3%       ~ (p=0.630 n=12)
  BM_UnaryOp/64/process_time               135.8µ ± 2%   135.9µ ± 1%       ~ (p=0.977 n=12)
  BM_UnaryOp/128/process_time              286.0µ ± 2%   285.3µ ± 1%       ~ (p=0.551 n=12)
  BM_UnaryOp/512/process_time              3.232m ± 2%   3.236m ± 2%       ~ (p=0.977 n=12)
  BM_UnaryOp/1024/process_time             12.76m ± 3%   12.63m ± 2%       ~ (p=0.219 n=12)
  BM_BinaryOp/64/process_time              107.2µ ± 1%   107.8µ ± 1%       ~ (p=0.551 n=12)
  BM_BinaryOp/128/process_time             175.2µ ± 1%   176.1µ ± 2%       ~ (p=0.843 n=12)
  BM_BinaryOp/512/process_time             1.632m ± 1%   1.635m ± 1%       ~ (p=0.713 n=12)
  BM_BinaryOp/1024/process_time            6.187m ± 2%   6.255m ± 1%       ~ (p=0.198 n=12)
  BM_EvaluateConvolution/4/process_time    9.220µ ± 3%   9.177µ ± 1%       ~ (p=0.514 n=12)
  BM_EvaluateConvolution/16/process_time   295.8µ ± 1%   294.7µ ± 2%       ~ (p=0.932 n=12)
  BM_EvaluateDotSlowPath/4/process_time    3.495µ ± 3%   3.514µ ± 1%  +0.54% (p=0.039 n=12)
  BM_EvaluateDotSlowPath/32/process_time   1.731m ± 1%   1.741m ± 1%       ~ (p=0.410 n=12)
  BM_EvaluateRaggedDot/4/process_time      4.478µ ± 2%   4.538µ ± 3%       ~ (p=0.630 n=12)
  BM_EvaluateRaggedDot/32/process_time     1.805m ± 1%   1.829m ± 1%  +1.31% (p=0.000 n=12)
  BM_EvaluateScaledDot/4/process_time      7.248µ ± 2%   7.222µ ± 3%       ~ (p=0.977 n=12)
  BM_EvaluateScaledDot/32/process_time     2.958m ± 1%   2.922m ± 1%       ~ (p=0.078 n=12)
  BM_EvaluatePad/4/process_time            84.98µ ± 2%   84.71µ ± 2%       ~ (p=0.443 n=12)
  BM_EvaluatePad/64/process_time           283.9µ ± 2%   283.6µ ± 2%       ~ (p=0.843 n=12)
  geomean                                  363.2µ        363.6µ       +0.09%
  ```

PiperOrigin-RevId: 997547513
---
 .../xla/xla/hlo/evaluator/hlo_evaluator.cc    |  467 ++++++++
 .../xla/xla/hlo/evaluator/hlo_evaluator.h     |   68 ++
 .../xla/hlo/evaluator/hlo_evaluator_test.cc   |  143 +++
 .../evaluator/hlo_evaluator_typed_visitor.h   | 1028 +++++------------
 4 files changed, 991 insertions(+), 715 deletions(-)

diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
index a1ac820152b..3e01a8c8ba5 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
@@ -5494,4 +5494,471 @@ HloEvaluator::ShapeInfo::ShapeInfo(
 
 HloEvaluator::ShapeInfo::~ShapeInfo() = default;
 
+HloEvaluator::DotDimensionInfo::DotDimensionInfo(
+    const Shape& lhs_shape, const Shape& rhs_shape,
+    const DotDimensionNumbers& dnums,
+    std::optional<int64_t> extra_rhs_contracting_dim)
+    : lhs_rank(lhs_shape.dimensions().size()),
+      rhs_rank(rhs_shape.dimensions().size()) {
+  lhs_contracting_dims.assign(dnums.lhs_contracting_dimensions().begin(),
+                              dnums.lhs_contracting_dimensions().end());
+  rhs_contracting_dims.assign(dnums.rhs_contracting_dimensions().begin(),
+                              dnums.rhs_contracting_dimensions().end());
+  if (extra_rhs_contracting_dim.has_value()) {
+    // Group Dimension is also a contracting dimension.
+    rhs_contracting_dims.push_back(*extra_rhs_contracting_dim);
+  }
+  lhs_non_contracting_dims = GetNonContractingDims(
+      lhs_rank, lhs_contracting_dims, dnums.lhs_batch_dimensions());
+  rhs_non_contracting_dims = GetNonContractingDims(
+      rhs_rank, rhs_contracting_dims, dnums.rhs_batch_dimensions());
+  contracting_dim_sizes.reserve(lhs_contracting_dims.size());
+  for (int64_t lhs_dnum : lhs_contracting_dims) {
+    contracting_dim_sizes.push_back(lhs_shape.dimensions(lhs_dnum));
+  }
+  total_contraction_size = Product(contracting_dim_sizes);
+}
+
+HloEvaluator::DotDimensionInfo::~DotDimensionInfo() = default;
+
+absl::Status HloEvaluator::EvaluateConvolution(
+    const HloInstruction* conv, bool use_f32_fast_path,
+    absl::FunctionRef<absl::Status(
+        const Literal& lhs_literal, const Literal& rhs_literal,
+        const Shape& window_shape, const DimensionVector& lhs_dim_multipliers,
+        const DimensionVector& rhs_dim_multipliers, Literal* result)>
+        evaluate_impl) {
+  auto lhs = conv->operand(0);
+  auto rhs = conv->operand(1);
+  const auto& window = conv->window();
+  Shape result_shape = GetShapeWithLayout(conv->shape());
+  Shape lhs_shape;
+  Shape rhs_shape;
+
+  std::optional<Literal> decompressed_lhs;
+  const Literal* lhs_literal_ptr = &GetEvaluatedLiteralFor(lhs);
+  std::optional<Literal> decompressed_rhs;
+  const Literal* rhs_literal_ptr = &GetEvaluatedLiteralFor(rhs);
+
+  if (conv->sparsity_config().has_lhs()) {
+    auto lhs_indices_op = conv->operand(conv->sparsity_config().lhs().idx());
+    const Literal* lhs_indices = &GetEvaluatedLiteralFor(lhs_indices_op);
+    ABSL_ASSIGN_OR_RETURN(decompressed_lhs, xla::MaterializeSparseOperand(
+                                           *lhs_literal_ptr, *lhs_indices,
+                                           conv->sparsity_config().lhs()));
+    lhs_literal_ptr = &decompressed_lhs.value();
+    lhs_shape = lhs_literal_ptr->shape();
+  } else {
+    lhs_shape = GetShapeWithLayout(lhs->shape());
+  }
+
+  if (conv->sparsity_config().has_rhs()) {
+    auto rhs_indices_op = conv->operand(conv->sparsity_config().rhs().idx());
+    const Literal* rhs_indices = &GetEvaluatedLiteralFor(rhs_indices_op);
+    ABSL_ASSIGN_OR_RETURN(decompressed_rhs, xla::MaterializeSparseOperand(
+                                           *rhs_literal_ptr, *rhs_indices,
+                                           conv->sparsity_config().rhs()));
+    rhs_literal_ptr = &decompressed_rhs.value();
+    rhs_shape = rhs_literal_ptr->shape();
+  } else {
+    rhs_shape = GetShapeWithLayout(rhs->shape());
+  }
+
+  CHECK_OK(ShapeUtil::ValidateShape(lhs_shape));
+  CHECK_OK(ShapeUtil::ValidateShape(rhs_shape));
+  CHECK(lhs_shape.IsArray());
+  CHECK(rhs_shape.IsArray());
+
+  const auto& dnums = conv->convolution_dimension_numbers();
+  const int64_t num_spatial_dims = dnums.output_spatial_dimensions_size();
+  CHECK_EQ(num_spatial_dims, dnums.input_spatial_dimensions_size());
+  CHECK_EQ(num_spatial_dims, dnums.kernel_spatial_dimensions_size());
+  CHECK_GE(num_spatial_dims, 0);
+  CHECK_EQ(window.dimensions_size(), num_spatial_dims);
+
+  const auto lhs_rank = lhs_shape.dimensions().size();
+  const auto rhs_rank = rhs_shape.dimensions().size();
+
+  CHECK_EQ(num_spatial_dims + 2, lhs_rank);
+  CHECK_EQ(num_spatial_dims + 2, rhs_rank);
+
+  ABSL_ASSIGN_OR_RETURN(
+      auto inferred_return_shape,
+      ShapeInference::InferConvolveShape(
+          lhs_shape, rhs_shape, conv->feature_group_count(),
+          conv->batch_group_count(), window, dnums, SparsityConfig(),
+          /*preferred_element_type=*/conv->shape().element_type()));
+  CHECK(ShapeUtil::Compatible(result_shape, inferred_return_shape))
+      << "return shape set to: " << ShapeUtil::HumanString(result_shape)
+      << " but is inferred to be: "
+      << ShapeUtil::HumanString(inferred_return_shape);
+
+  std::optional<Literal> converted_lhs;
+  if (!ShapeUtil::SameElementType(lhs_shape, result_shape)) {
+    converted_lhs =
+        lhs_literal_ptr->Convert(result_shape.element_type()).value();
+    lhs_literal_ptr = &converted_lhs.value();
+  }
+  std::optional<Literal> converted_rhs;
+  if (!ShapeUtil::SameElementType(rhs_shape, result_shape)) {
+    converted_rhs =
+        rhs_literal_ptr->Convert(result_shape.element_type()).value();
+    rhs_literal_ptr = &converted_rhs.value();
+  }
+
+  const Literal& lhs_literal = *lhs_literal_ptr;
+  const Literal& rhs_literal = *rhs_literal_ptr;
+  const Shape& eval_lhs_shape = lhs_literal.shape();
+  const Shape& eval_rhs_shape = rhs_literal.shape();
+
+  CHECK_OK(ShapeUtil::ValidateShape(eval_lhs_shape));
+  CHECK_OK(ShapeUtil::ValidateShape(eval_rhs_shape));
+  CHECK(eval_lhs_shape.IsArray());
+  CHECK(eval_rhs_shape.IsArray());
+  CHECK(ShapeUtil::SameElementType(eval_lhs_shape, eval_rhs_shape));
+  CHECK(ShapeUtil::SameElementType(eval_lhs_shape, result_shape));
+
+  const int64_t feature_group_count = conv->feature_group_count();
+  const int64_t batch_group_count = conv->batch_group_count();
+
+  if (use_f32_fast_path) {
+    auto is_row_major_r2 = [](const Shape& s) {
+      return s.dimensions_size() == 2 &&
+             (!s.has_layout() ||
+              LayoutUtil::IsMonotonicWithDim0Major(s.layout()));
+    };
+
+    if (trace_mac_handler_ == nullptr && feature_group_count == 1 &&
+        batch_group_count == 1 && num_spatial_dims == 0 &&
+        dnums.input_batch_dimension() == 0 &&
+        dnums.input_feature_dimension() == 1 &&
+        dnums.kernel_input_feature_dimension() == 0 &&
+        dnums.kernel_output_feature_dimension() == 1 &&
+        dnums.output_batch_dimension() == 0 &&
+        dnums.output_feature_dimension() == 1 &&
+        is_row_major_r2(eval_lhs_shape) && is_row_major_r2(eval_rhs_shape) &&
+        is_row_major_r2(result_shape)) {
+      const int64_t m = eval_lhs_shape.dimensions(0);
+      const int64_t k = eval_lhs_shape.dimensions(1);
+      const int64_t n = eval_rhs_shape.dimensions(1);
+
+      if (m > 0 && k > 0 && n > 0) {
+        Literal lhs_f32 = lhs_literal.Convert(F32).value();
+        Literal rhs_f32 = rhs_literal.Convert(F32).value();
+
+        Array2D<float> lhs_array(m, k);
+        lhs_array.SetValues(lhs_f32.data<float>());
+        Array2D<float> rhs_array(k, n);
+        rhs_array.SetValues(rhs_f32.data<float>());
+
+        std::unique_ptr<Array2D<float>> result_array =
+            HloEvaluator::MatmulArray2D(lhs_array, rhs_array);
+
+        Literal result_f32(ShapeUtil::MakeShape(F32, {m, n}));
+        result_f32.PopulateR2FromArray2D(*result_array);
+
+        SetEvaluatedLiteralFor(
+            conv,
+            std::move(result_f32).Convert(result_shape.element_type()).value());
+        return absl::OkStatus();
+      }
+    }
+  }
+
+  std::vector<int64_t> window_dimension_sizes;
+  window_dimension_sizes.reserve(num_spatial_dims);
+  for (auto i : dnums.kernel_spatial_dimensions()) {
+    window_dimension_sizes.push_back(
+        ShapeUtil::GetDimension(eval_rhs_shape, i));
+  }
+
+  const Shape window_shape = ShapeUtil::MakeShape(eval_rhs_shape.element_type(),
+                                                  window_dimension_sizes);
+
+  DimensionVector lhs_dim_multipliers = MakeDimMultipliers(eval_lhs_shape);
+  DimensionVector rhs_dim_multipliers = MakeDimMultipliers(eval_rhs_shape);
+
+  Literal result(result_shape);
+  ABSL_RETURN_IF_ERROR(evaluate_impl(lhs_literal, rhs_literal, window_shape,
+                                lhs_dim_multipliers, rhs_dim_multipliers,
+                                &result));
+  SetEvaluatedLiteralFor(conv, std::move(result));
+  return absl::OkStatus();
+}
+
+absl::Status HloEvaluator::EvaluateDotSlowPath(
+    const HloInstruction* dot,
+    absl::FunctionRef<
+        absl::Status(const Literal& lhs_literal, const Literal& rhs_literal,
+                     const DotDimensionInfo& dot_info, Literal* result)>
+        evaluate_impl) {
+  auto lhs = dot->operand(0);
+  auto rhs = dot->operand(1);
+  CHECK(dot->shape().IsArray());
+  CHECK(lhs->shape().IsArray());
+  CHECK(rhs->shape().IsArray());
+
+  const Literal* lhs_literal_ptr = &GetEvaluatedLiteralFor(lhs);
+  const Literal* rhs_literal_ptr = &GetEvaluatedLiteralFor(rhs);
+  std::optional<Literal> converted_lhs;
+  if (!ShapeUtil::SameElementType(lhs->shape(), dot->shape())) {
+    converted_lhs =
+        lhs_literal_ptr->Convert(dot->shape().element_type()).value();
+    lhs_literal_ptr = &converted_lhs.value();
+  }
+  std::optional<Literal> converted_rhs;
+  if (!ShapeUtil::SameElementType(rhs->shape(), dot->shape())) {
+    converted_rhs =
+        rhs_literal_ptr->Convert(dot->shape().element_type()).value();
+    rhs_literal_ptr = &converted_rhs.value();
+  }
+
+  const Literal& lhs_literal = *lhs_literal_ptr;
+  const Literal& rhs_literal = *rhs_literal_ptr;
+  const auto& dnums = dot->dot_dimension_numbers();
+
+  CHECK(ShapeUtil::SameElementType(lhs_literal.shape(), rhs_literal.shape()));
+  CHECK(ShapeUtil::SameElementType(lhs_literal.shape(), dot->shape()));
+
+  CHECK_EQ(dnums.lhs_batch_dimensions_size(),
+           dnums.rhs_batch_dimensions_size());
+
+  DotDimensionInfo dot_info(lhs_literal.shape(), rhs_literal.shape(), dnums);
+  Shape dot_shape = GetShapeWithLayout(dot->shape());
+  Literal result(dot_shape);
+  ABSL_RETURN_IF_ERROR(evaluate_impl(lhs_literal, rhs_literal, dot_info, &result));
+  SetEvaluatedLiteralFor(dot, std::move(result));
+  return absl::OkStatus();
+}
+
+absl::Status HloEvaluator::EvaluateRaggedDot(
+    const HloInstruction* dot,
+    absl::FunctionRef<absl::Status(
+        const Literal& lhs_literal, const Literal& rhs_literal,
+        const Literal& gs_literal, const DotDimensionInfo& dot_info,
+        int64_t ragged_dim_as_contracting_dim, Literal* result)>
+        contracting_impl,
+    absl::FunctionRef<
+        absl::Status(const Literal& lhs_literal, const Literal& rhs_literal,
+                     const Literal& gs_literal,
+                     const DotDimensionInfo& dot_info, Literal* result)>
+        batch_impl,
+    absl::FunctionRef<
+        absl::Status(const Literal& lhs_literal, const Literal& rhs_literal,
+                     const Literal& gs_literal,
+                     const DotDimensionInfo& dot_info, Literal* result)>
+        non_contracting_impl) {
+  auto lhs = dot->operand(0);
+  auto rhs = dot->operand(1);
+  auto group_sizes = dot->operand(2);
+
+  CHECK(dot->shape().IsArray());
+  CHECK(lhs->shape().IsArray());
+  CHECK(rhs->shape().IsArray());
+  CHECK(group_sizes->shape().IsArray());
+
+  const Literal* lhs_literal_ptr = &GetEvaluatedLiteralFor(lhs);
+  const Literal* rhs_literal_ptr = &GetEvaluatedLiteralFor(rhs);
+  const Literal* gs_literal_ptr = &GetEvaluatedLiteralFor(group_sizes);
+
+  // LHS and RHS may have a different initial precision than our output.
+  std::optional<Literal> converted_lhs;
+  if (!ShapeUtil::SameElementType(lhs->shape(), dot->shape())) {
+    converted_lhs =
+        lhs_literal_ptr->Convert(dot->shape().element_type()).value();
+    lhs_literal_ptr = &converted_lhs.value();
+  }
+  std::optional<Literal> converted_rhs;
+  if (!ShapeUtil::SameElementType(rhs->shape(), dot->shape())) {
+    converted_rhs =
+        rhs_literal_ptr->Convert(dot->shape().element_type()).value();
+    rhs_literal_ptr = &converted_rhs.value();
+  }
+  // Upcast group sizes to S64, since they could be e.g. S32.
+  std::optional<Literal> converted_gs;
+  if (group_sizes->shape().element_type() != PrimitiveType::S64) {
+    converted_gs = gs_literal_ptr->Convert(PrimitiveType::S64).value();
+    gs_literal_ptr = &converted_gs.value();
+  }
+
+  const Literal& lhs_literal = *lhs_literal_ptr;
+  const Literal& rhs_literal = *rhs_literal_ptr;
+  const Literal& gs_literal = *gs_literal_ptr;
+
+  const auto& ragged_dims = dot->ragged_dot_dimension_numbers();
+  const auto& dot_dims = ragged_dims.dot_dimension_numbers();
+  // Shape inference should have checked that there is exactly one ragged dim.
+  const int64_t lhs_ragged_dim = ragged_dims.lhs_ragged_dimensions(0);
+
+  Shape dot_shape = GetShapeWithLayout(dot->shape());
+  Literal result(dot_shape);
+
+  auto contracting_it =
+      std::find(dot_dims.lhs_contracting_dimensions().begin(),
+                dot_dims.lhs_contracting_dimensions().end(), lhs_ragged_dim);
+  if (contracting_it != dot_dims.lhs_contracting_dimensions().end()) {
+    const int64_t ragged_dim_as_contracting_dim =
+        contracting_it - dot_dims.lhs_contracting_dimensions().begin();
+    DotDimensionInfo dot_info(lhs_literal.shape(), rhs_literal.shape(),
+                              dot_dims);
+    ABSL_RETURN_IF_ERROR(contracting_impl(lhs_literal, rhs_literal, gs_literal,
+                                     dot_info, ragged_dim_as_contracting_dim,
+                                     &result));
+    SetEvaluatedLiteralFor(dot, std::move(result));
+    return absl::OkStatus();
+  }
+
+  if (std::find(dot_dims.lhs_batch_dimensions().begin(),
+                dot_dims.lhs_batch_dimensions().end(),
+                lhs_ragged_dim) != dot_dims.lhs_batch_dimensions().end()) {
+    DotDimensionInfo dot_info(lhs_literal.shape(), rhs_literal.shape(),
+                              dot_dims);
+    ABSL_RETURN_IF_ERROR(
+        batch_impl(lhs_literal, rhs_literal, gs_literal, dot_info, &result));
+    SetEvaluatedLiteralFor(dot, std::move(result));
+    return absl::OkStatus();
+  }
+
+  // Shape inference should have checked that this mode has a group dim.
+  const int64_t rhs_group_dim = ragged_dims.rhs_group_dimensions(0);
+  DotDimensionInfo dot_info(lhs_literal.shape(), rhs_literal.shape(), dot_dims,
+                            rhs_group_dim);
+  ABSL_RETURN_IF_ERROR(non_contracting_impl(lhs_literal, rhs_literal, gs_literal,
+                                       dot_info, &result));
+  SetEvaluatedLiteralFor(dot, std::move(result));
+  return absl::OkStatus();
+}
+
+absl::Status HloEvaluator::EvaluateScaledDot(
+    const HloInstruction* dot,
+    absl::FunctionRef<absl::Status(
+        const Literal& lhs_literal, const Literal& rhs_literal,
+        const Literal& lhs_scale_literal, const Literal& rhs_scale_literal,
+        const ShapeInfo& lhs_info, const ShapeInfo& rhs_info,
+        int64_t total_contraction_size, Literal* result)>
+        evaluate_impl) {
+  auto lhs = dot->operand(0);
+  auto rhs = dot->operand(1);
+  auto lhs_scale = dot->operand(2);
+  auto rhs_scale = dot->operand(3);
+  CHECK(dot->shape().IsArray());
+  CHECK(lhs->shape().IsArray());
+  CHECK(rhs->shape().IsArray());
+  CHECK(lhs_scale->shape().IsArray());
+  CHECK(rhs_scale->shape().IsArray());
+  CHECK(lhs_scale->shape().dimensions().size() == 0 ||
+        lhs_scale->shape().dimensions().size() ==
+            lhs->shape().dimensions().size());
+  CHECK(rhs_scale->shape().dimensions().size() == 0 ||
+        rhs_scale->shape().dimensions().size() ==
+            rhs->shape().dimensions().size());
+  ABSL_ASSIGN_OR_RETURN(
+      const Literal lhs_literal,
+      GetEvaluatedLiteralFor(lhs).Convert(dot->shape().element_type()));
+  ABSL_ASSIGN_OR_RETURN(
+      const Literal rhs_literal,
+      GetEvaluatedLiteralFor(rhs).Convert(dot->shape().element_type()));
+
+  // If the scale is a scalar, we can just use 1.0. Otherwise, we need to
+  // evaluate the scale.
+  auto evaluate_scale =
+      [&](const HloInstruction* operand,
+          const HloInstruction* scale) -> absl::StatusOr<Literal> {
+    if (scale->shape().IsArray() && scale->shape().dimensions().size() > 0) {
+      return GetEvaluatedLiteralFor(scale).Convert(dot->shape().element_type());
+    }
+    std::vector<int64_t> ones(operand->shape().dimensions().size(), 1);
+    Literal scale_f32(ShapeUtil::MakeShape(F32, ones));
+    scale_f32.PopulateWithValue(1.0f);
+    return scale_f32.Convert(dot->shape().element_type());
+  };
+  ABSL_ASSIGN_OR_RETURN(Literal lhs_scale_literal, evaluate_scale(lhs, lhs_scale));
+  ABSL_ASSIGN_OR_RETURN(Literal rhs_scale_literal, evaluate_scale(rhs, rhs_scale));
+
+  const auto& dnums = dot->dot_dimension_numbers();
+  CHECK(ShapeUtil::SameElementType(lhs_literal.shape(), rhs_literal.shape()));
+  CHECK(ShapeUtil::SameElementType(lhs_literal.shape(), dot->shape()));
+
+  CHECK_EQ(dnums.lhs_batch_dimensions_size(),
+           dnums.rhs_batch_dimensions_size());
+
+  ShapeInfo lhs_info(lhs_literal, lhs_scale_literal,
+                     dnums.lhs_contracting_dimensions(),
+                     dnums.lhs_batch_dimensions());
+  ShapeInfo rhs_info(rhs_literal, rhs_scale_literal,
+                     dnums.rhs_contracting_dimensions(),
+                     dnums.rhs_batch_dimensions());
+  const int64_t total_contraction_size =
+      Product(lhs_info.contracting_dim_sizes);
+  Shape dot_shape = GetShapeWithLayout(dot->shape());
+
+  ABSL_ASSIGN_OR_RETURN(Literal result, Literal::Make(dot_shape));
+  ABSL_RETURN_IF_ERROR(evaluate_impl(lhs_literal, rhs_literal, lhs_scale_literal,
+                                rhs_scale_literal, lhs_info, rhs_info,
+                                total_contraction_size, &result));
+  SetEvaluatedLiteralFor(dot, std::move(result));
+  return absl::OkStatus();
+}
+
+absl::Status HloEvaluator::EvaluatePad(
+    const HloInstruction* pad,
+    absl::FunctionRef<absl::Status(
+        const Literal& evaluated_operand, const Literal& evaluated_padding,
+        absl::Span<int64_t> target_index, absl::Span<const int64_t> zero_base,
+        absl::Span<const int64_t> step, Literal* result)>
+        evaluate_impl) {
+  CHECK(pad->operand(0)->shape().IsArray());
+  // Padding value must be scalar.
+  CHECK(ShapeUtil::IsScalar(pad->operand(1)->shape()));
+  CHECK_EQ(pad->operand(0)->shape().dimensions().size(),
+           pad->padding_config().dimensions_size());
+
+  ABSL_ASSIGN_OR_RETURN(auto inferred_return_shape,
+                   ShapeInference::InferPadShape(
+                       /*operand_shape=*/pad->operand(0)->shape(),
+                       /*padding_value_shape=*/pad->operand(1)->shape(),
+                       /*padding_config=*/pad->padding_config()));
+  CHECK(ShapeUtil::CompatibleIgnoringElementType(pad->shape(),
+                                                 inferred_return_shape))
+      << "return shape is set to: " << ShapeUtil::HumanString(pad->shape())
+      << " but is inferred to be: "
+      << ShapeUtil::HumanString(inferred_return_shape);
+
+  PrimitiveType result_type = pad->shape().element_type();
+  PrimitiveType padding_type = pad->operand(1)->shape().element_type();
+  std::optional<Literal> converted_padding;
+  const Literal* evaluated_padding_ptr =
+      &GetEvaluatedLiteralFor(pad->operand(1));
+  if (padding_type != result_type) {
+    ABSL_ASSIGN_OR_RETURN(converted_padding,
+                     evaluated_padding_ptr->Convert(result_type));
+    evaluated_padding_ptr = &converted_padding.value();
+  }
+
+  // Create new HLO of padded shape with padding value.
+  Literal result(GetShapeWithLayout(pad->shape()));
+
+  PrimitiveType operand_type = pad->operand(0)->shape().element_type();
+  std::optional<Literal> converted_operand;
+  const Literal* evaluated_operand_ptr =
+      &GetEvaluatedLiteralFor(pad->operand(0));
+  if (operand_type != result_type) {
+    ABSL_ASSIGN_OR_RETURN(converted_operand,
+                     evaluated_operand_ptr->Convert(result_type));
+    evaluated_operand_ptr = &converted_operand.value();
+  }
+  const Literal& evaluated_operand = *evaluated_operand_ptr;
+
+  std::vector<int64_t> target_index(result.shape().dimensions().size(), 0);
+  std::vector<int64_t> zero_base(evaluated_operand.shape().dimensions().size(),
+                                 0);
+  std::vector<int64_t> step(evaluated_operand.shape().dimensions().size(), 1);
+
+  ABSL_RETURN_IF_ERROR(evaluate_impl(evaluated_operand, *evaluated_padding_ptr,
+                                absl::MakeSpan(target_index), zero_base, step,
+                                &result));
+  SetEvaluatedLiteralFor(pad, std::move(result));
+  return absl::OkStatus();
+}
+
 }  // namespace xla
diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h
index 184864e96bf..58e43dbf27b 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h
@@ -34,6 +34,7 @@ limitations under the License.
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/inlined_vector.h"
 #include "absl/container/node_hash_map.h"
+#include "absl/functional/function_ref.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
@@ -520,6 +521,73 @@ class HloEvaluator : public ConstDfsHloVisitorWithDefault,
     DimensionVector contracting_dim_scale_divisors;
   };
 
+  struct DotDimensionInfo {
+    DotDimensionInfo(
+        const Shape& lhs_shape, const Shape& rhs_shape,
+        const DotDimensionNumbers& dnums,
+        std::optional<int64_t> extra_rhs_contracting_dim = std::nullopt);
+    ~DotDimensionInfo();
+
+    int64_t lhs_rank;
+    int64_t rhs_rank;
+    DimensionVector lhs_non_contracting_dims;
+    DimensionVector rhs_non_contracting_dims;
+    DimensionVector lhs_contracting_dims;
+    DimensionVector rhs_contracting_dims;
+    DimensionVector contracting_dim_sizes;
+    int64_t total_contraction_size;
+  };
+
+  absl::Status EvaluateConvolution(
+      const HloInstruction* conv, bool use_f32_fast_path,
+      absl::FunctionRef<absl::Status(
+          const Literal& lhs_literal, const Literal& rhs_literal,
+          const Shape& window_shape, const DimensionVector& lhs_dim_multipliers,
+          const DimensionVector& rhs_dim_multipliers, Literal* result)>
+          evaluate_impl);
+
+  absl::Status EvaluateDotSlowPath(
+      const HloInstruction* dot,
+      absl::FunctionRef<
+          absl::Status(const Literal& lhs_literal, const Literal& rhs_literal,
+                       const DotDimensionInfo& dot_info, Literal* result)>
+          evaluate_impl);
+
+  absl::Status EvaluateRaggedDot(
+      const HloInstruction* dot,
+      absl::FunctionRef<absl::Status(
+          const Literal& lhs_literal, const Literal& rhs_literal,
+          const Literal& gs_literal, const DotDimensionInfo& dot_info,
+          int64_t ragged_dim_as_contracting_dim, Literal* result)>
+          contracting_impl,
+      absl::FunctionRef<
+          absl::Status(const Literal& lhs_literal, const Literal& rhs_literal,
+                       const Literal& gs_literal,
+                       const DotDimensionInfo& dot_info, Literal* result)>
+          batch_impl,
+      absl::FunctionRef<
+          absl::Status(const Literal& lhs_literal, const Literal& rhs_literal,
+                       const Literal& gs_literal,
+                       const DotDimensionInfo& dot_info, Literal* result)>
+          non_contracting_impl);
+
+  absl::Status EvaluateScaledDot(
+      const HloInstruction* dot,
+      absl::FunctionRef<absl::Status(
+          const Literal& lhs_literal, const Literal& rhs_literal,
+          const Literal& lhs_scale_literal, const Literal& rhs_scale_literal,
+          const ShapeInfo& lhs_info, const ShapeInfo& rhs_info,
+          int64_t total_contraction_size, Literal* result)>
+          evaluate_impl);
+
+  absl::Status EvaluatePad(
+      const HloInstruction* pad,
+      absl::FunctionRef<absl::Status(
+          const Literal& evaluated_operand, const Literal& evaluated_padding,
+          absl::Span<int64_t> target_index, absl::Span<const int64_t> zero_base,
+          absl::Span<const int64_t> step, Literal* result)>
+          evaluate_impl);
+
   // Make HloEvaluatorTypedVisitor a friend because it is logically part of this
   // class.
   //
diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc
index 7249ecb44c7..bacb11ead17 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc
@@ -8998,5 +8998,148 @@ BENCHMARK(BM_BinaryOp)
     ->Arg(512)
     ->Arg(1024);
 
+static void BM_EvaluateConvolution(benchmark::State& state) {
+  int64_t d = state.range(0);
+
+  std::unique_ptr<HloInstruction> lhs = HloInstruction::CreateConstant(
+      LiteralUtil::CreateFull({1, 2, d, d}, 1.0f));
+  std::unique_ptr<HloInstruction> rhs = HloInstruction::CreateConstant(
+      LiteralUtil::CreateFull({2, 2, 3, 3}, 1.0f));
+
+  Window window;
+  WindowDimension dim;
+  dim.set_size(3);
+  dim.set_stride(1);
+  dim.set_padding_low(1);
+  dim.set_padding_high(1);
+  dim.set_window_dilation(1);
+  dim.set_base_dilation(1);
+  *window.add_dimensions() = dim;
+  *window.add_dimensions() = dim;
+
+  ConvolutionDimensionNumbers dnums =
+      XlaBuilder::CreateDefaultConvDimensionNumbers(2);
+  std::unique_ptr<HloInstruction> conv = HloInstruction::CreateConvolve(
+      ShapeUtil::MakeShape(F32, {1, 2, d, d}), {lhs.get(), rhs.get()},
+      /*feature_group_count=*/1, /*batch_group_count=*/1, window, dnums,
+      HloHardwareIndependentTestBase::DefaultPrecisionConfig(2));
+
+  HloEvaluator evaluator;
+  for (auto s : state) {
+    CHECK_OK(evaluator.Evaluate(conv.get()).status());
+  }
+}
+
+BENCHMARK(BM_EvaluateConvolution)->MeasureProcessCPUTime()->Arg(4)->Arg(16);
+
+static void BM_EvaluateDotSlowPath(benchmark::State& state) {
+  int64_t d = state.range(0);
+
+  std::unique_ptr<HloInstruction> lhs = HloInstruction::CreateConstant(
+      LiteralUtil::CreateFull<int32_t>({d, d}, 1));
+  std::unique_ptr<HloInstruction> rhs = HloInstruction::CreateConstant(
+      LiteralUtil::CreateFull<int32_t>({d, d}, 2));
+
+  DotDimensionNumbers dot_dnums;
+  dot_dnums.add_lhs_contracting_dimensions(1);
+  dot_dnums.add_rhs_contracting_dimensions(0);
+  std::unique_ptr<HloInstruction> dot = HloInstruction::CreateDot(
+      ShapeUtil::MakeShape(S32, {d, d}), lhs.get(), rhs.get(), dot_dnums,
+      HloHardwareIndependentTestBase::DefaultPrecisionConfig(2));
+
+  HloEvaluator evaluator;
+  for (auto s : state) {
+    CHECK_OK(evaluator.Evaluate(dot.get()).status());
+  }
+}
+
+BENCHMARK(BM_EvaluateDotSlowPath)->MeasureProcessCPUTime()->Arg(4)->Arg(32);
+
+static void BM_EvaluateRaggedDot(benchmark::State& state) {
+  int64_t d = state.range(0);
+
+  HloComputation::Builder b("BM_EvaluateRaggedDot");
+  HloModuleConfig config;
+  config.set_seed(1);
+  HloModule module("BM_EvaluateRaggedDot", config);
+
+  HloInstruction* lhs = b.AddInstruction(
+      HloInstruction::CreateConstant(LiteralUtil::CreateFull({d, d}, 1.0f)));
+  HloInstruction* rhs = b.AddInstruction(
+      HloInstruction::CreateConstant(LiteralUtil::CreateFull({2, d, d}, 2.0f)));
+  HloInstruction* gs = b.AddInstruction(HloInstruction::CreateConstant(
+      LiteralUtil::CreateR1<int64_t>({d / 2, d / 2})));
+
+  RaggedDotDimensionNumbers ragged_dnums;
+  ragged_dnums.add_lhs_ragged_dimensions(0);
+  ragged_dnums.add_rhs_group_dimensions(0);
+  ragged_dnums.mutable_dot_dimension_numbers()->add_lhs_contracting_dimensions(
+      1);
+  ragged_dnums.mutable_dot_dimension_numbers()->add_rhs_contracting_dimensions(
+      1);
+  b.AddInstruction(HloInstruction::CreateRaggedDot(
+      ShapeUtil::MakeShape(F32, {d, d}), lhs, rhs, gs, ragged_dnums,
+      HloHardwareIndependentTestBase::DefaultPrecisionConfig(2)));
+  module.AddEntryComputation(b.Build());
+
+  HloEvaluator evaluator;
+  for (auto s : state) {
+    evaluator.ResetVisitStates();
+    CHECK_OK(evaluator.Evaluate(*module.entry_computation(), {}).status());
+  }
+}
+
+BENCHMARK(BM_EvaluateRaggedDot)->MeasureProcessCPUTime()->Arg(4)->Arg(32);
+
+static void BM_EvaluateScaledDot(benchmark::State& state) {
+  int64_t d = state.range(0);
+
+  std::unique_ptr<HloInstruction> lhs =
+      HloInstruction::CreateConstant(LiteralUtil::CreateFull({d, d}, 1.0f));
+  std::unique_ptr<HloInstruction> lhs_scale =
+      HloInstruction::CreateConstant(LiteralUtil::CreateFull({d, 2}, 2.0f));
+  std::unique_ptr<HloInstruction> rhs =
+      HloInstruction::CreateConstant(LiteralUtil::CreateFull({d, d}, 1.0f));
+  std::unique_ptr<HloInstruction> rhs_scale =
+      HloInstruction::CreateConstant(LiteralUtil::CreateFull({2, d}, 3.0f));
+
+  DotDimensionNumbers dot_dnums;
+  dot_dnums.add_lhs_contracting_dimensions(1);
+  dot_dnums.add_rhs_contracting_dimensions(0);
+  std::unique_ptr<HloInstruction> scaled_dot = HloInstruction::CreateScaledDot(
+      ShapeUtil::MakeShape(F32, {d, d}), lhs.get(), rhs.get(), lhs_scale.get(),
+      rhs_scale.get(), dot_dnums,
+      HloHardwareIndependentTestBase::DefaultPrecisionConfig(4));
+
+  HloEvaluator evaluator;
+  for (auto s : state) {
+    CHECK_OK(evaluator.Evaluate(scaled_dot.get()).status());
+  }
+}
+
+BENCHMARK(BM_EvaluateScaledDot)->MeasureProcessCPUTime()->Arg(4)->Arg(32);
+
+static void BM_EvaluatePad(benchmark::State& state) {
+  int64_t d = state.range(0);
+
+  std::unique_ptr<HloInstruction> operand =
+      HloInstruction::CreateConstant(LiteralUtil::CreateFull({d, d}, 1.0f));
+  std::unique_ptr<HloInstruction> pad_value =
+      HloInstruction::CreateConstant(LiteralUtil::CreateR0<float>(0.0f));
+
+  PaddingConfig padding_config =
+      CreatePaddingConfig({{{1, 1, 0}}, {{1, 1, 0}}});
+  std::unique_ptr<HloInstruction> pad =
+      HloInstruction::CreatePad(ShapeUtil::MakeShape(F32, {d + 2, d + 2}),
+                                operand.get(), pad_value.get(), padding_config);
+
+  HloEvaluator evaluator;
+  for (auto s : state) {
+    CHECK_OK(evaluator.Evaluate(pad.get()).status());
+  }
+}
+
+BENCHMARK(BM_EvaluatePad)->MeasureProcessCPUTime()->Arg(4)->Arg(64);
+
 }  // namespace
 }  // namespace xla
diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h
index fecf11523a8..98125e7e21a 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h
@@ -904,96 +904,28 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return absl::OkStatus();
   }
 
-  absl::Status HandleConvolutionWithLiterals(const HloInstruction* conv,
-                                             const Literal& lhs_literal,
-                                             const Literal& rhs_literal) {
+  absl::Status HandleConvolutionWithLiterals(
+      const HloInstruction* conv, const Literal& lhs_literal,
+      const Literal& rhs_literal, const Shape& window_shape,
+      const DimensionVector& lhs_dim_multipliers,
+      const DimensionVector& rhs_dim_multipliers, Literal* result) {
     const auto& window = conv->window();
-    Shape result_shape = GetShapeWithLayout(conv->shape());
+    const Shape& result_shape = result->shape();
     const Shape& lhs_shape = lhs_literal.shape();
     const Shape& rhs_shape = rhs_literal.shape();
 
-    CHECK_OK(ShapeUtil::ValidateShape(lhs_shape));
-    CHECK_OK(ShapeUtil::ValidateShape(rhs_shape));
-    CHECK(lhs_shape.IsArray());
-    CHECK(rhs_shape.IsArray());
-    CHECK(ShapeUtil::SameElementType(lhs_shape, rhs_shape));
-    CHECK(ShapeUtil::SameElementType(lhs_shape, result_shape));
-
     const auto& dnums = conv->convolution_dimension_numbers();
-    const int64_t num_spatial_dims = dnums.output_spatial_dimensions_size();
-    CHECK_EQ(num_spatial_dims, dnums.input_spatial_dimensions_size());
-    CHECK_EQ(num_spatial_dims, dnums.kernel_spatial_dimensions_size());
-    CHECK_GE(num_spatial_dims, 0);
-    CHECK_EQ(window.dimensions_size(), num_spatial_dims);
-
-    std::vector<int64_t> window_dimension_sizes;
-    for (auto i : dnums.kernel_spatial_dimensions()) {
-      window_dimension_sizes.push_back(ShapeUtil::GetDimension(rhs_shape, i));
-    }
-
-    const Shape& window_shape =
-        ShapeUtil::MakeShape(rhs_shape.element_type(), window_dimension_sizes);
-
-    DimensionVector lhs_dim_multipliers =
-        HloEvaluator::MakeDimMultipliers(lhs_shape);
-    DimensionVector rhs_dim_multipliers =
-        HloEvaluator::MakeDimMultipliers(rhs_shape);
-
     auto lhs_literal_data = lhs_literal.data<ReturnT>();
     auto rhs_literal_data = rhs_literal.data<ReturnT>();
 
     const int64_t feature_group_count = conv->feature_group_count();
     const int64_t batch_group_count = conv->batch_group_count();
 
-    if constexpr (std::is_same_v<ElementwiseT, float>) {
-      auto is_row_major_r2 = [](const Shape& s) {
-        return s.dimensions_size() == 2 &&
-               (!s.has_layout() ||
-                LayoutUtil::IsMonotonicWithDim0Major(s.layout()));
-      };
-
-      if (parent_->trace_mac_handler_ == nullptr && feature_group_count == 1 &&
-          batch_group_count == 1 && num_spatial_dims == 0 &&
-          dnums.input_batch_dimension() == 0 &&
-          dnums.input_feature_dimension() == 1 &&
-          dnums.kernel_input_feature_dimension() == 0 &&
-          dnums.kernel_output_feature_dimension() == 1 &&
-          dnums.output_batch_dimension() == 0 &&
-          dnums.output_feature_dimension() == 1 && is_row_major_r2(lhs_shape) &&
-          is_row_major_r2(rhs_shape) && is_row_major_r2(result_shape)) {
-        const int64_t m = lhs_shape.dimensions(0);
-        const int64_t k = lhs_shape.dimensions(1);
-        const int64_t n = rhs_shape.dimensions(1);
-
-        if (m > 0 && k > 0 && n > 0) {
-          Literal lhs_f32 = lhs_literal.Convert(F32).value();
-          Literal rhs_f32 = rhs_literal.Convert(F32).value();
-
-          Array2D<float> lhs_array(m, k);
-          lhs_array.SetValues(lhs_f32.data<float>());
-          Array2D<float> rhs_array(k, n);
-          rhs_array.SetValues(rhs_f32.data<float>());
-
-          std::unique_ptr<Array2D<float>> result_array =
-              HloEvaluator::MatmulArray2D(lhs_array, rhs_array);
-
-          Literal result_f32(ShapeUtil::MakeShape(F32, {m, n}));
-          result_f32.PopulateR2FromArray2D(*result_array);
-
-          parent_->SetEvaluatedLiteralFor(
-              conv, std::move(result_f32)
-                        .Convert(result_shape.element_type())
-                        .value());
-          return absl::OkStatus();
-        }
-      }
-    }
-
     auto func = [&window_shape, &dnums, &lhs_shape, &rhs_shape, &window,
                  &lhs_dim_multipliers, &rhs_dim_multipliers, lhs_literal_data,
                  rhs_literal_data, feature_group_count, batch_group_count,
-                 result_shape, this](const absl::Span<const int64_t> out_index,
-                                     int /*thread_id*/) {
+                 &result_shape, this](const absl::Span<const int64_t> out_index,
+                                      int /*thread_id*/) {
       // Dimension number applicable for input (lhs).
       const int64_t input_batch_dim = dnums.input_batch_dimension();
       const int64_t input_z_dim = dnums.input_feature_dimension();
@@ -1138,101 +1070,21 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
       return static_cast<ReturnT>(result_val);
     };
 
-    Literal result(result_shape);
-    ABSL_RETURN_IF_ERROR(result.PopulateParallel<ReturnT>(func));
-
-    parent_->SetEvaluatedLiteralFor(conv, std::move(result));
-    return absl::OkStatus();
+    return result->PopulateParallel<ReturnT>(func);
   }
 
   absl::Status HandleConvolution(const HloInstruction* conv) override {
-    auto lhs = conv->operand(0);
-    auto rhs = conv->operand(1);
-    const auto& window = conv->window();
-    Shape result_shape = GetShapeWithLayout(conv->shape());
-    Shape lhs_shape;
-    Shape rhs_shape;
-
-    std::optional<Literal> decompressed_lhs;
-    const Literal* lhs_literal_ptr = &parent_->GetEvaluatedLiteralFor(lhs);
-    std::optional<Literal> decompressed_rhs;
-    const Literal* rhs_literal_ptr = &parent_->GetEvaluatedLiteralFor(rhs);
-
-    if (conv->sparsity_config().has_lhs()) {
-      auto lhs_indices_op = conv->operand(conv->sparsity_config().lhs().idx());
-      const Literal* lhs_indices =
-          &parent_->GetEvaluatedLiteralFor(lhs_indices_op);
-      ABSL_ASSIGN_OR_RETURN(decompressed_lhs, xla::MaterializeSparseOperand(
-                                             *lhs_literal_ptr, *lhs_indices,
-                                             conv->sparsity_config().lhs()));
-      lhs_literal_ptr = &decompressed_lhs.value();
-      lhs_shape = lhs_literal_ptr->shape();
-    } else {
-      lhs_shape = GetShapeWithLayout(lhs->shape());
-    }
-
-    if (conv->sparsity_config().has_rhs()) {
-      auto rhs_indices_op = conv->operand(conv->sparsity_config().rhs().idx());
-      const Literal* rhs_indices =
-          &parent_->GetEvaluatedLiteralFor(rhs_indices_op);
-      ABSL_ASSIGN_OR_RETURN(decompressed_rhs, xla::MaterializeSparseOperand(
-                                             *rhs_literal_ptr, *rhs_indices,
-                                             conv->sparsity_config().rhs()));
-      rhs_literal_ptr = &decompressed_rhs.value();
-      rhs_shape = rhs_literal_ptr->shape();
-    } else {
-      rhs_shape = GetShapeWithLayout(rhs->shape());
-    }
-
-    CHECK_OK(ShapeUtil::ValidateShape(lhs_shape));
-    CHECK_OK(ShapeUtil::ValidateShape(rhs_shape));
-    CHECK(lhs_shape.IsArray());
-    CHECK(rhs_shape.IsArray());
-
-    const auto& dnums = conv->convolution_dimension_numbers();
-    const int64_t num_spatial_dims = dnums.output_spatial_dimensions_size();
-    CHECK_EQ(num_spatial_dims, dnums.input_spatial_dimensions_size());
-    CHECK_EQ(num_spatial_dims, dnums.kernel_spatial_dimensions_size());
-    CHECK_GE(num_spatial_dims, 0);
-    CHECK_EQ(window.dimensions_size(), num_spatial_dims);
-
-    const auto lhs_rank = lhs_shape.dimensions().size();
-    const auto rhs_rank = rhs_shape.dimensions().size();
-
-    CHECK_EQ(num_spatial_dims + 2, lhs_rank);
-    CHECK_EQ(num_spatial_dims + 2, rhs_rank);
-
-    ABSL_ASSIGN_OR_RETURN(
-        auto inferred_return_shape,
-        ShapeInference::InferConvolveShape(
-            lhs_shape, rhs_shape, conv->feature_group_count(),
-            conv->batch_group_count(), window, dnums, SparsityConfig(),
-            /*preferred_element_type=*/conv->shape().element_type()));
-    CHECK(ShapeUtil::Compatible(result_shape, inferred_return_shape))
-        << "return shape set to: " << ShapeUtil::HumanString(result_shape)
-        << " but is inferred to be: "
-        << ShapeUtil::HumanString(inferred_return_shape);
-
-    const Literal& lhs_literal = *lhs_literal_ptr;
-    const Literal& rhs_literal = *rhs_literal_ptr;
-    const bool lhs_same = ShapeUtil::SameElementType(lhs_shape, result_shape);
-    const bool rhs_same = ShapeUtil::SameElementType(rhs_shape, result_shape);
-    if (rhs_same && lhs_same) {
-      return HandleConvolutionWithLiterals(conv, lhs_literal, rhs_literal);
-    }
-    if (rhs_same) {
-      return HandleConvolutionWithLiterals(
-          conv, lhs_literal.Convert(result_shape.element_type()).value(),
-          rhs_literal);
-    }
-    if (lhs_same) {
-      return HandleConvolutionWithLiterals(
-          conv, lhs_literal,
-          rhs_literal.Convert(result_shape.element_type()).value());
-    }
-    return HandleConvolutionWithLiterals(
-        conv, lhs_literal.Convert(result_shape.element_type()).value(),
-        rhs_literal.Convert(result_shape.element_type()).value());
+    return parent_->EvaluateConvolution(
+        conv, /*use_f32_fast_path=*/std::is_same_v<ElementwiseT, float>,
+        [this, conv](const Literal& lhs_literal, const Literal& rhs_literal,
+                     const Shape& window_shape,
+                     const DimensionVector& lhs_dim_multipliers,
+                     const DimensionVector& rhs_dim_multipliers,
+                     Literal* result) {
+          return HandleConvolutionWithLiterals(
+              conv, lhs_literal, rhs_literal, window_shape, lhs_dim_multipliers,
+              rhs_dim_multipliers, result);
+        });
   }
 
   absl::Status HandleDot(const HloInstruction* dot) override {
@@ -1270,381 +1122,254 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     }
   }
 
-  absl::Status HandleDotSlowPathWithLiterals(const HloInstruction* dot,
-                                             const Literal& lhs_literal,
-                                             const Literal& rhs_literal) {
+  absl::Status HandleDotSlowPathWithLiterals(
+      const HloInstruction* dot, const Literal& lhs_literal,
+      const Literal& rhs_literal,
+      const HloEvaluator::DotDimensionInfo& dot_info, Literal* result) {
     const auto& dnums = dot->dot_dimension_numbers();
+    const Shape& dot_shape = result->shape();
+    return result->PopulateParallel<ReturnT>([&](absl::Span<const int64_t>
+                                                     result_index,
+                                                 int /*thread_id*/) {
+      // Locations in LHS and RHS that we read from.
+      DimensionVector lhs_index(dot_info.lhs_rank);
+      DimensionVector rhs_index(dot_info.rhs_rank);
 
-    const auto lhs_rank = lhs_literal.shape().dimensions().size();
-    const auto rhs_rank = rhs_literal.shape().dimensions().size();
+      // First come the batch dimensions.
+      int64_t idx = 0;
+      for (int64_t i = 0; i < dnums.lhs_batch_dimensions_size(); i++) {
+        lhs_index[dnums.lhs_batch_dimensions(i)] = result_index[idx];
+        rhs_index[dnums.rhs_batch_dimensions(i)] = result_index[idx];
+        idx++;
+      }
 
-    CHECK(ShapeUtil::SameElementType(lhs_literal.shape(), rhs_literal.shape()));
-    CHECK(ShapeUtil::SameElementType(lhs_literal.shape(), dot->shape()));
+      // Next we have non-contracting dimensions, if any.
+      for (int64_t i = 0; i < dot_info.lhs_non_contracting_dims.size(); i++) {
+        lhs_index[dot_info.lhs_non_contracting_dims[i]] = result_index[idx++];
+      }
+      for (int64_t i = 0; i < dot_info.rhs_non_contracting_dims.size(); i++) {
+        rhs_index[dot_info.rhs_non_contracting_dims[i]] = result_index[idx++];
+      }
 
-    CHECK_EQ(dnums.lhs_batch_dimensions_size(),
-             dnums.rhs_batch_dimensions_size());
+      // Accumulate resulting product along the contracting dimensions.
+      ElementwiseT result_val = static_cast<ElementwiseT>(0);
+      for (int64_t k = 0; k < dot_info.total_contraction_size; k++) {
+        const auto lhs =
+            static_cast<ElementwiseT>(lhs_literal.Get<ReturnT>(lhs_index));
+        const auto rhs =
+            static_cast<ElementwiseT>(rhs_literal.Get<ReturnT>(rhs_index));
+        result_val += ToArithmeticSafeType(lhs) * ToArithmeticSafeType(rhs);
 
-    DimensionVector lhs_non_contracting_dims =
-        GetNonContractingDims(lhs_rank, dnums.lhs_contracting_dimensions(),
-                              dnums.lhs_batch_dimensions());
-    DimensionVector rhs_non_contracting_dims =
-        GetNonContractingDims(rhs_rank, dnums.rhs_contracting_dimensions(),
-                              dnums.rhs_batch_dimensions());
+        if (parent_->trace_mac_handler_ != nullptr) {
+          const int64_t result_linear_index =
+              IndexUtil::MultidimensionalIndexToLinearIndex(dot_shape,
+                                                            result_index);
+          const int64_t lhs_linear_index =
+              IndexUtil::MultidimensionalIndexToLinearIndex(lhs_literal.shape(),
+                                                            lhs_index);
+          const int64_t rhs_linear_index =
+              IndexUtil::MultidimensionalIndexToLinearIndex(rhs_literal.shape(),
+                                                            rhs_index);
 
-    DimensionVector contracting_dim_sizes;
-    contracting_dim_sizes.reserve(dnums.lhs_contracting_dimensions_size());
-    DimensionVector lhs_contracting_dims;
-    DimensionVector rhs_contracting_dims;
-    for (int64_t i = 0; i < dnums.lhs_contracting_dimensions_size(); ++i) {
-      const int64_t lhs_dnum = dnums.lhs_contracting_dimensions(i);
-      const int64_t rhs_dnum = dnums.rhs_contracting_dimensions(i);
-      lhs_contracting_dims.push_back(lhs_dnum);
-      rhs_contracting_dims.push_back(rhs_dnum);
-      const int64_t dim_size = lhs_literal.shape().dimensions(lhs_dnum);
-      contracting_dim_sizes.push_back(dim_size);
-    }
-    const int64_t total_contraction_size = Product(contracting_dim_sizes);
-    Shape dot_shape = GetShapeWithLayout(dot->shape());
-    Literal result(dot_shape);
-    ABSL_RETURN_IF_ERROR(result.PopulateParallel<ReturnT>(
-        [&](absl::Span<const int64_t> result_index, int /*thread_id*/) {
-          // Locations in LHS and RHS that we read from.
-          DimensionVector lhs_index(lhs_rank);
-          DimensionVector rhs_index(rhs_rank);
+          parent_->trace_mac_handler_(result_linear_index, lhs_linear_index,
+                                      rhs_linear_index);
+        }
 
-          // First come the batch dimensions.
-          int64_t idx = 0;
-          for (int64_t i = 0; i < dnums.lhs_batch_dimensions_size(); i++) {
-            lhs_index[dnums.lhs_batch_dimensions(i)] = result_index[idx];
-            rhs_index[dnums.rhs_batch_dimensions(i)] = result_index[idx];
-            idx++;
-          }
+        IncrementContractingIndexes(lhs_index, dot_info.lhs_contracting_dims,
+                                    dot_info.contracting_dim_sizes);
+        IncrementContractingIndexes(rhs_index, dot_info.rhs_contracting_dims,
+                                    dot_info.contracting_dim_sizes);
+      }
 
-          // Next we have non-contracting dimensions, if any.
-          for (int64_t i = 0; i < lhs_non_contracting_dims.size(); i++) {
-            lhs_index[lhs_non_contracting_dims[i]] = result_index[idx++];
-          }
-          for (int64_t i = 0; i < rhs_non_contracting_dims.size(); i++) {
-            rhs_index[rhs_non_contracting_dims[i]] = result_index[idx++];
-          }
-
-          // Accumulate resulting product along the contracting dimensions.
-          ElementwiseT result_val = static_cast<ElementwiseT>(0);
-          for (int64_t k = 0; k < total_contraction_size; k++) {
-            const auto lhs =
-                static_cast<ElementwiseT>(lhs_literal.Get<ReturnT>(lhs_index));
-            const auto rhs =
-                static_cast<ElementwiseT>(rhs_literal.Get<ReturnT>(rhs_index));
-            result_val += ToArithmeticSafeType(lhs) * ToArithmeticSafeType(rhs);
-
-            if (parent_->trace_mac_handler_ != nullptr) {
-              const int64_t result_linear_index =
-                  IndexUtil::MultidimensionalIndexToLinearIndex(dot_shape,
-                                                                result_index);
-              const int64_t lhs_linear_index =
-                  IndexUtil::MultidimensionalIndexToLinearIndex(
-                      lhs_literal.shape(), lhs_index);
-              const int64_t rhs_linear_index =
-                  IndexUtil::MultidimensionalIndexToLinearIndex(
-                      rhs_literal.shape(), rhs_index);
-
-              parent_->trace_mac_handler_(result_linear_index, lhs_linear_index,
-                                          rhs_linear_index);
-            }
-
-            IncrementContractingIndexes(lhs_index, lhs_contracting_dims,
-                                        contracting_dim_sizes);
-            IncrementContractingIndexes(rhs_index, rhs_contracting_dims,
-                                        contracting_dim_sizes);
-          }
-
-          return static_cast<ReturnT>(result_val);
-        }));
-
-    parent_->SetEvaluatedLiteralFor(dot, std::move(result));
-    return absl::OkStatus();
+      return static_cast<ReturnT>(result_val);
+    });
   }
 
   absl::Status HandleDotSlowPath(const HloInstruction* dot) {
-    auto lhs = dot->operand(0);
-    auto rhs = dot->operand(1);
-    CHECK(dot->shape().IsArray());
-    CHECK(lhs->shape().IsArray());
-    CHECK(rhs->shape().IsArray());
-    const bool lhs_same =
-        ShapeUtil::SameElementType(lhs->shape(), dot->shape());
-    const bool rhs_same =
-        ShapeUtil::SameElementType(rhs->shape(), dot->shape());
-    const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs);
-    const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs);
-    if (lhs_same && rhs_same) {
-      return HandleDotSlowPathWithLiterals(dot, lhs_literal, rhs_literal);
-    }
-    if (lhs_same) {
-      return HandleDotSlowPathWithLiterals(
-          dot, lhs_literal,
-          rhs_literal.Convert(dot->shape().element_type()).value());
-    }
-    if (rhs_same) {
-      return HandleDotSlowPathWithLiterals(
-          dot, lhs_literal.Convert(dot->shape().element_type()).value(),
-          rhs_literal);
-    }
-    return HandleDotSlowPathWithLiterals(
-        dot, lhs_literal.Convert(dot->shape().element_type()).value(),
-        rhs_literal.Convert(dot->shape().element_type()).value());
+    return parent_->EvaluateDotSlowPath(
+        dot, [this, dot](const Literal& lhs_literal, const Literal& rhs_literal,
+                         const HloEvaluator::DotDimensionInfo& dot_info,
+                         Literal* result) {
+          return HandleDotSlowPathWithLiterals(dot, lhs_literal, rhs_literal,
+                                               dot_info, result);
+        });
   }
 
   absl::Status HandleRaggedDotNonContractingWithLiterals(
       const HloInstruction* dot, const Literal& lhs_literal,
-      const Literal& rhs_literal, const Literal& gs_literal) {
-    auto ragged_dims = dot->ragged_dot_dimension_numbers();
-    auto dot_dims = ragged_dims.dot_dimension_numbers();
+      const Literal& rhs_literal, const Literal& gs_literal,
+      const HloEvaluator::DotDimensionInfo& dot_info, Literal* result) {
+    const auto& ragged_dims = dot->ragged_dot_dimension_numbers();
+    const auto& dot_dims = ragged_dims.dot_dimension_numbers();
     // Shape inference should have checked that there is exactly one ragged dim.
     int64_t lhs_ragged_dim = ragged_dims.lhs_ragged_dimensions(0);
-
-    int64_t lhs_rank = lhs_literal.shape().dimensions().size();
-    int64_t rhs_rank = rhs_literal.shape().dimensions().size();
-    int64_t gs_rank = gs_literal.shape().dimensions().size();
-    int64_t num_groups = gs_literal.shape().dimensions(gs_rank - 1);
-
-    auto lhs_contracting = dot_dims.lhs_contracting_dimensions();
-    auto lhs_non_contracting = GetNonContractingDims(
-        lhs_rank, lhs_contracting, dot_dims.lhs_batch_dimensions());
-
     // Shape inference should have checked that this mode has a group dim.
     int64_t rhs_group_dim = ragged_dims.rhs_group_dimensions(0);
 
-    auto rhs_contracting = dot_dims.rhs_contracting_dimensions();
-    // Group Dimension is also a contracting dimension.
-    rhs_contracting.Add(rhs_group_dim);
-    auto rhs_non_contracting = GetNonContractingDims(
-        rhs_rank, rhs_contracting, dot_dims.rhs_batch_dimensions());
+    int64_t gs_rank = gs_literal.shape().dimensions().size();
+    int64_t num_groups = gs_literal.shape().dimensions(gs_rank - 1);
 
-    DimensionVector contracting_dim_sizes;
-    contracting_dim_sizes.reserve(lhs_contracting.size());
-    for (int64_t i = 0; i < lhs_contracting.size(); ++i) {
-      int64_t dim_size = lhs_literal.shape().dimensions(lhs_contracting[i]);
-      contracting_dim_sizes.push_back(dim_size);
-    }
-    const int64_t total_contracting_size = Product(contracting_dim_sizes);
+    return result->PopulateParallel<ReturnT>([&](absl::Span<const int64_t>
+                                                     result_index,
+                                                 int /*thread_id*/) {
+      // Locations in each operand that we read from to calculate the result
+      // at result_index.
+      DimensionVector lhs_index(dot_info.lhs_rank);
+      DimensionVector rhs_index(dot_info.rhs_rank);
+      DimensionVector group_index(gs_rank);
 
-    Shape dot_shape = GetShapeWithLayout(dot->shape());
-    Literal result(dot_shape);
-    ABSL_RETURN_IF_ERROR(result.PopulateParallel<ReturnT>(
-        [&](absl::Span<const int64_t> result_index, int /*thread_id*/) {
-          // Locations in each operand that we read from to calculate the result
-          // at result_index.
-          DimensionVector lhs_index(lhs_rank);
-          DimensionVector rhs_index(rhs_rank);
-          DimensionVector group_index(gs_rank);
+      // Batch dimensions will always be first in the final product.
+      const int64_t group_dim_index = gs_rank - 1;
+      int64_t idx = 0;
+      int64_t gs_idx = 0;
+      for (int64_t i = 0; i < dot_dims.lhs_batch_dimensions_size(); ++i) {
+        lhs_index[dot_dims.lhs_batch_dimensions(i)] = result_index[idx];
+        rhs_index[dot_dims.rhs_batch_dimensions(i)] = result_index[idx];
+        if (gs_idx < group_dim_index) {
+          group_index[gs_idx++] = result_index[idx];
+        }
+        idx++;
+      }
 
-          // Batch dimensions will always be first in the final product.
-          const int64_t group_dim_index = gs_rank - 1;
-          int64_t idx = 0;
-          int64_t gs_idx = 0;
-          for (int64_t i = 0; i < dot_dims.lhs_batch_dimensions_size(); ++i) {
-            lhs_index[dot_dims.lhs_batch_dimensions(i)] = result_index[idx];
-            rhs_index[dot_dims.rhs_batch_dimensions(i)] = result_index[idx];
-            if (gs_idx < group_dim_index) {
-              group_index[gs_idx++] = result_index[idx];
-            }
-            idx++;
+      // Non-contracting dimensions - lhs, then rhs.
+      for (int64_t i = 0; i < dot_info.lhs_non_contracting_dims.size(); ++i) {
+        // If there is a non-contracting lhs dimension that is not ragged,
+        // then there will also be a dimension for this in group_sizes.
+        if (lhs_ragged_dim != dot_info.lhs_non_contracting_dims[i]) {
+          if (gs_idx < group_dim_index) {
+            group_index[gs_idx++] = result_index[idx];
           }
+        }
+        lhs_index[dot_info.lhs_non_contracting_dims[i]] = result_index[idx++];
+      }
+      for (int64_t i = 0; i < dot_info.rhs_non_contracting_dims.size(); ++i) {
+        rhs_index[dot_info.rhs_non_contracting_dims[i]] = result_index[idx++];
+      }
 
-          // Non-contracting dimensions - lhs, then rhs.
-          for (int64_t i = 0; i < lhs_non_contracting.size(); ++i) {
-            // If there is a non-contracting lhs dimension that is not ragged,
-            // then there will also be a dimension for this in group_sizes.
-            if (lhs_ragged_dim != lhs_non_contracting[i]) {
-              if (gs_idx < group_dim_index) {
-                group_index[gs_idx++] = result_index[idx];
-              }
-            }
-            lhs_index[lhs_non_contracting[i]] = result_index[idx++];
-          }
-          for (int64_t i = 0; i < rhs_non_contracting.size(); ++i) {
-            rhs_index[rhs_non_contracting[i]] = result_index[idx++];
-          }
+      // Calculate which group the current lhs-row belongs to.
+      int64_t lhs_ragged_index = lhs_index[lhs_ragged_dim];
+      int64_t group_row_end = 0;
+      for (int64_t i = 0; i < num_groups; ++i) {
+        group_index[group_dim_index] = i;
+        group_row_end += gs_literal.Get<int64_t>(group_index);
+        if (lhs_ragged_index < group_row_end) {
+          break;
+        }
+      }
+      // If lhs_ragged_index > the sum(groups), then there is no
+      // corresponding rhs value, and there is no result.
+      if (lhs_ragged_index >= group_row_end) {
+        return static_cast<ReturnT>(0);
+      }
+      rhs_index[rhs_group_dim] = group_index[gs_idx];
 
-          // Calculate which group the current lhs-row belongs to.
-          int64_t lhs_ragged_index = lhs_index[lhs_ragged_dim];
-          int64_t group_row_end = 0;
-          for (int64_t i = 0; i < num_groups; ++i) {
-            group_index[group_dim_index] = i;
-            group_row_end += gs_literal.Get<int64_t>(group_index);
-            if (lhs_ragged_index < group_row_end) {
-              break;
-            }
-          }
-          // If lhs_ragged_index > the sum(groups), then there is no
-          // corresponding rhs value, and there is no result.
-          if (lhs_ragged_index >= group_row_end) {
-            return static_cast<ReturnT>(0);
-          }
-          rhs_index[rhs_group_dim] = group_index[gs_idx];
+      // Accumulate resulting product along the contracting dimensions.
+      ElementwiseT result_val = static_cast<ElementwiseT>(0);
+      for (int64_t i = 0; i < dot_info.total_contraction_size; ++i) {
+        const auto lhs =
+            static_cast<ElementwiseT>(lhs_literal.Get<ReturnT>(lhs_index));
+        const auto rhs =
+            static_cast<ElementwiseT>(rhs_literal.Get<ReturnT>(rhs_index));
+        result_val += ToArithmeticSafeType(lhs) * ToArithmeticSafeType(rhs);
 
-          // Accumulate resulting product along the contracting dimensions.
-          ElementwiseT result_val = static_cast<ElementwiseT>(0);
-          for (int64_t i = 0; i < total_contracting_size; ++i) {
-            const auto lhs =
-                static_cast<ElementwiseT>(lhs_literal.Get<ReturnT>(lhs_index));
-            const auto rhs =
-                static_cast<ElementwiseT>(rhs_literal.Get<ReturnT>(rhs_index));
-            result_val += ToArithmeticSafeType(lhs) * ToArithmeticSafeType(rhs);
-
-            IncrementContractingIndexes(lhs_index, lhs_contracting,
-                                        contracting_dim_sizes);
-            IncrementContractingIndexes(rhs_index, rhs_contracting,
-                                        contracting_dim_sizes);
-          }
-          return static_cast<ReturnT>(result_val);
-        }));
-
-    parent_->SetEvaluatedLiteralFor(dot, std::move(result));
-    return absl::OkStatus();
+        IncrementContractingIndexes(lhs_index, dot_info.lhs_contracting_dims,
+                                    dot_info.contracting_dim_sizes);
+        IncrementContractingIndexes(rhs_index, dot_info.rhs_contracting_dims,
+                                    dot_info.contracting_dim_sizes);
+      }
+      return static_cast<ReturnT>(result_val);
+    });
   }
 
-  absl::Status HandleRaggedDotBatchModeWithLiterals(const HloInstruction* dot,
-                                                    const Literal& lhs_literal,
-                                                    const Literal& rhs_literal,
-                                                    const Literal& gs_literal) {
-    auto ragged_dims = dot->ragged_dot_dimension_numbers();
-    auto dot_dims = ragged_dims.dot_dimension_numbers();
-    int64_t lhs_rank = lhs_literal.shape().dimensions().size();
-    int64_t rhs_rank = rhs_literal.shape().dimensions().size();
+  absl::Status HandleRaggedDotBatchModeWithLiterals(
+      const HloInstruction* dot, const Literal& lhs_literal,
+      const Literal& rhs_literal, const Literal& gs_literal,
+      const HloEvaluator::DotDimensionInfo& dot_info, Literal* result) {
+    const auto& dot_dims =
+        dot->ragged_dot_dimension_numbers().dot_dimension_numbers();
     int64_t gs_rank = gs_literal.shape().dimensions().size();
-
-    auto lhs_contracting = dot_dims.lhs_contracting_dimensions();
-    auto lhs_non_contracting = GetNonContractingDims(
-        lhs_rank, lhs_contracting, dot_dims.lhs_batch_dimensions());
-
-    auto rhs_contracting = dot_dims.rhs_contracting_dimensions();
-    auto rhs_non_contracting = GetNonContractingDims(
-        rhs_rank, rhs_contracting, dot_dims.rhs_batch_dimensions());
-
-    DimensionVector contracting_dim_sizes;
-    contracting_dim_sizes.reserve(lhs_contracting.size());
-    for (int64_t i = 0; i < lhs_contracting.size(); ++i) {
-      int64_t dim_size = lhs_literal.shape().dimensions(lhs_contracting[i]);
-      contracting_dim_sizes.push_back(dim_size);
-    }
-    const int64_t total_contracting_size = Product(contracting_dim_sizes);
     const int64_t group_dim_index = gs_rank - 1;
     const int64_t num_groups = gs_literal.shape().dimensions(group_dim_index);
 
-    Shape dot_shape = GetShapeWithLayout(dot->shape());
-    Literal result(dot_shape);
-    ABSL_RETURN_IF_ERROR(result.PopulateParallel<ReturnT>(
-        [&](absl::Span<const int64_t> result_index, int /*thread_id*/) {
-          // Locations in each operand that we read from to calculate the
-          // result at result_index.
-          DimensionVector lhs_index(lhs_rank);
-          DimensionVector rhs_index(rhs_rank);
-          DimensionVector group_index(gs_rank);
+    return result->PopulateParallel<ReturnT>([&](absl::Span<const int64_t>
+                                                     result_index,
+                                                 int /*thread_id*/) {
+      // Locations in each operand that we read from to calculate the
+      // result at result_index.
+      DimensionVector lhs_index(dot_info.lhs_rank);
+      DimensionVector rhs_index(dot_info.rhs_rank);
+      DimensionVector group_index(gs_rank);
 
-          // The batch dimensions will always be first in the final product.
-          int64_t gs_idx = 0;
-          int64_t idx = 0;
-          for (int64_t i = 0; i < dot_dims.lhs_batch_dimensions_size(); ++i) {
-            lhs_index[dot_dims.lhs_batch_dimensions(i)] = result_index[idx];
-            rhs_index[dot_dims.rhs_batch_dimensions(i)] = result_index[idx];
-            // gs only contains the batch dimensions outer to the ragged dim.
-            if (gs_idx < group_dim_index) {
-              group_index[gs_idx++] = result_index[idx];
-            }
-            ++idx;
-          }
+      // The batch dimensions will always be first in the final product.
+      int64_t gs_idx = 0;
+      int64_t idx = 0;
+      for (int64_t i = 0; i < dot_dims.lhs_batch_dimensions_size(); ++i) {
+        lhs_index[dot_dims.lhs_batch_dimensions(i)] = result_index[idx];
+        rhs_index[dot_dims.rhs_batch_dimensions(i)] = result_index[idx];
+        // gs only contains the batch dimensions outer to the ragged dim.
+        if (gs_idx < group_dim_index) {
+          group_index[gs_idx++] = result_index[idx];
+        }
+        ++idx;
+      }
 
-          // Check whether this batch is handled by some group.
-          int64_t batches_handled = 0;
-          for (int i = 0; i < num_groups; ++i) {
-            group_index[group_dim_index] = i;
-            batches_handled += gs_literal.Get<int64_t>(group_index);
-          }
-          if (lhs_index[dot_dims.lhs_batch_dimensions(group_dim_index)] >=
-              batches_handled) {
-            return static_cast<ReturnT>(0);
-          }
+      // Check whether this batch is handled by some group.
+      int64_t batches_handled = 0;
+      for (int i = 0; i < num_groups; ++i) {
+        group_index[group_dim_index] = i;
+        batches_handled += gs_literal.Get<int64_t>(group_index);
+      }
+      if (lhs_index[dot_dims.lhs_batch_dimensions(group_dim_index)] >=
+          batches_handled) {
+        return static_cast<ReturnT>(0);
+      }
 
-          // Non-contracting dimensions - lhs, then rhs.
-          for (int64_t i = 0; i < lhs_non_contracting.size(); ++i) {
-            lhs_index[lhs_non_contracting[i]] = result_index[idx++];
-          }
-          for (int64_t i = 0; i < rhs_non_contracting.size(); ++i) {
-            rhs_index[rhs_non_contracting[i]] = result_index[idx++];
-          }
+      // Non-contracting dimensions - lhs, then rhs.
+      for (int64_t i = 0; i < dot_info.lhs_non_contracting_dims.size(); ++i) {
+        lhs_index[dot_info.lhs_non_contracting_dims[i]] = result_index[idx++];
+      }
+      for (int64_t i = 0; i < dot_info.rhs_non_contracting_dims.size(); ++i) {
+        rhs_index[dot_info.rhs_non_contracting_dims[i]] = result_index[idx++];
+      }
 
-          // Accumulate resulting product along the contracting dimensions.
-          ElementwiseT result_val = static_cast<ElementwiseT>(0);
-          for (int64_t i = 0; i < total_contracting_size; ++i) {
-            const auto lhs =
-                static_cast<ElementwiseT>(lhs_literal.Get<ReturnT>(lhs_index));
-            const auto rhs =
-                static_cast<ElementwiseT>(rhs_literal.Get<ReturnT>(rhs_index));
-            result_val += ToArithmeticSafeType(lhs) * ToArithmeticSafeType(rhs);
+      // Accumulate resulting product along the contracting dimensions.
+      ElementwiseT result_val = static_cast<ElementwiseT>(0);
+      for (int64_t i = 0; i < dot_info.total_contraction_size; ++i) {
+        const auto lhs =
+            static_cast<ElementwiseT>(lhs_literal.Get<ReturnT>(lhs_index));
+        const auto rhs =
+            static_cast<ElementwiseT>(rhs_literal.Get<ReturnT>(rhs_index));
+        result_val += ToArithmeticSafeType(lhs) * ToArithmeticSafeType(rhs);
 
-            IncrementContractingIndexes(lhs_index, lhs_contracting,
-                                        contracting_dim_sizes);
-            IncrementContractingIndexes(rhs_index, rhs_contracting,
-                                        contracting_dim_sizes);
-          }
-          return static_cast<ReturnT>(result_val);
-        }));
-
-    parent_->SetEvaluatedLiteralFor(dot, std::move(result));
-    return absl::OkStatus();
+        IncrementContractingIndexes(lhs_index, dot_info.lhs_contracting_dims,
+                                    dot_info.contracting_dim_sizes);
+        IncrementContractingIndexes(rhs_index, dot_info.rhs_contracting_dims,
+                                    dot_info.contracting_dim_sizes);
+      }
+      return static_cast<ReturnT>(result_val);
+    });
   }
 
   absl::Status HandleRaggedDotContractingWithLiterals(
       const HloInstruction* dot, const Literal& lhs_literal,
-      const Literal& rhs_literal, const Literal& gs_literal) {
-    auto ragged_dims = dot->ragged_dot_dimension_numbers();
-    auto dot_dims = ragged_dims.dot_dimension_numbers();
-    // Shape inference should have checked that there is exactly one ragged dim.
-    int64_t lhs_ragged_dim = ragged_dims.lhs_ragged_dimensions(0);
-
-    int64_t lhs_rank = lhs_literal.shape().dimensions().size();
-    int64_t rhs_rank = rhs_literal.shape().dimensions().size();
+      const Literal& rhs_literal, const Literal& gs_literal,
+      const HloEvaluator::DotDimensionInfo& dot_info,
+      int64_t ragged_dim_as_contracting_dim, Literal* result) {
+    const auto& dot_dims =
+        dot->ragged_dot_dimension_numbers().dot_dimension_numbers();
     int64_t gs_rank = gs_literal.shape().dimensions().size();
-
-    auto lhs_contracting = dot_dims.lhs_contracting_dimensions();
-    auto lhs_non_contracting = GetNonContractingDims(
-        lhs_rank, lhs_contracting, dot_dims.lhs_batch_dimensions());
-
-    auto rhs_contracting = dot_dims.rhs_contracting_dimensions();
-    auto rhs_non_contracting = GetNonContractingDims(
-        rhs_rank, rhs_contracting, dot_dims.rhs_batch_dimensions());
-
-    DimensionVector contracting_dim_sizes;
-    contracting_dim_sizes.reserve(lhs_contracting.size());
-    for (int64_t i = 0; i < lhs_contracting.size(); ++i) {
-      int64_t dim_size = lhs_literal.shape().dimensions(lhs_contracting[i]);
-      contracting_dim_sizes.push_back(dim_size);
-    }
-
-    // We must find a match because we are in contracting mode.
-    std::optional<int64_t> ragged_dim_as_contracting_dim;
-    for (int64_t i = 0; i < lhs_contracting.size(); ++i) {
-      if (lhs_ragged_dim == lhs_contracting[i]) {
-        ragged_dim_as_contracting_dim = i;
-        break;
-      }
-    }
     const int64_t ragged_dim_size =
-        contracting_dim_sizes[ragged_dim_as_contracting_dim.value()];
+        dot_info.contracting_dim_sizes[ragged_dim_as_contracting_dim];
     const int64_t total_contracting_size_excluding_ragged_dim =
-        Product(contracting_dim_sizes) / ragged_dim_size;
+        dot_info.total_contraction_size / ragged_dim_size;
 
-    Shape dot_shape = GetShapeWithLayout(dot->shape());
-    Literal result(dot_shape);
-    ABSL_RETURN_IF_ERROR(result.PopulateParallel<
-                    ReturnT>([&](absl::Span<const int64_t> result_index,
-                                 int /*thread_id*/) {
+    return result->PopulateParallel<ReturnT>([&](absl::Span<const int64_t>
+                                                     result_index,
+                                                 int /*thread_id*/) {
       // Locations in each operand that we read from to calculate the
       // result at result_index.
-      DimensionVector lhs_index(lhs_rank);
-      DimensionVector rhs_index(rhs_rank);
+      DimensionVector lhs_index(dot_info.lhs_rank);
+      DimensionVector rhs_index(dot_info.rhs_rank);
       DimensionVector group_index(gs_rank);
 
       // The group dimension will always be first in the final product. We
@@ -1679,11 +1404,11 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
       group_row_end = std::min(group_row_end, ragged_dim_size);
 
       // Non-contracting dimensions - lhs, then rhs.
-      for (int64_t i = 0; i < lhs_non_contracting.size(); ++i) {
-        lhs_index[lhs_non_contracting[i]] = result_index[idx++];
+      for (int64_t i = 0; i < dot_info.lhs_non_contracting_dims.size(); ++i) {
+        lhs_index[dot_info.lhs_non_contracting_dims[i]] = result_index[idx++];
       }
-      for (int64_t i = 0; i < rhs_non_contracting.size(); ++i) {
-        rhs_index[rhs_non_contracting[i]] = result_index[idx++];
+      for (int64_t i = 0; i < dot_info.rhs_non_contracting_dims.size(); ++i) {
+        rhs_index[dot_info.rhs_non_contracting_dims[i]] = result_index[idx++];
       }
 
       // Accumulate resulting product along the contracting dimensions.
@@ -1691,8 +1416,12 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
       for (int64_t i = 0; i < total_contracting_size_excluding_ragged_dim;
            ++i) {
         for (int64_t j = group_row_start; j < group_row_end; ++j) {
-          lhs_index[lhs_contracting[ragged_dim_as_contracting_dim.value()]] = j;
-          rhs_index[rhs_contracting[ragged_dim_as_contracting_dim.value()]] = j;
+          lhs_index[dot_info
+                        .lhs_contracting_dims[ragged_dim_as_contracting_dim]] =
+              j;
+          rhs_index[dot_info
+                        .rhs_contracting_dims[ragged_dim_as_contracting_dim]] =
+              j;
           const auto lhs =
               static_cast<ElementwiseT>(lhs_literal.Get<ReturnT>(lhs_index));
           const auto rhs =
@@ -1700,44 +1429,15 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
           result_val += ToArithmeticSafeType(lhs) * ToArithmeticSafeType(rhs);
         }
 
-        IncrementContractingIndexes(lhs_index, lhs_contracting,
-                                    contracting_dim_sizes,
+        IncrementContractingIndexes(lhs_index, dot_info.lhs_contracting_dims,
+                                    dot_info.contracting_dim_sizes,
                                     ragged_dim_as_contracting_dim);
-        IncrementContractingIndexes(rhs_index, rhs_contracting,
-                                    contracting_dim_sizes,
+        IncrementContractingIndexes(rhs_index, dot_info.rhs_contracting_dims,
+                                    dot_info.contracting_dim_sizes,
                                     ragged_dim_as_contracting_dim);
       }
       return static_cast<ReturnT>(result_val);
-    }));
-
-    parent_->SetEvaluatedLiteralFor(dot, std::move(result));
-    return absl::OkStatus();
-  }
-
-  absl::Status HandleRaggedDotWithLiterals(const HloInstruction* dot,
-                                           const Literal& lhs_literal,
-                                           const Literal& rhs_literal,
-                                           const Literal& gs_literal) {
-    auto ragged_dims = dot->ragged_dot_dimension_numbers();
-    auto dot_dims = ragged_dims.dot_dimension_numbers();
-    // Shape inference should have checked that there is exactly one ragged dim.
-    int64_t lhs_ragged_dim = ragged_dims.lhs_ragged_dimensions(0);
-
-    if (std::find(dot_dims.lhs_contracting_dimensions().begin(),
-                  dot_dims.lhs_contracting_dimensions().end(),
-                  lhs_ragged_dim) !=
-        dot_dims.lhs_contracting_dimensions().end()) {
-      return HandleRaggedDotContractingWithLiterals(dot, lhs_literal,
-                                                    rhs_literal, gs_literal);
-    }
-    if (std::find(dot_dims.lhs_batch_dimensions().begin(),
-                  dot_dims.lhs_batch_dimensions().end(),
-                  lhs_ragged_dim) != dot_dims.lhs_batch_dimensions().end()) {
-      return HandleRaggedDotBatchModeWithLiterals(dot, lhs_literal, rhs_literal,
-                                                  gs_literal);
-    }
-    return HandleRaggedDotNonContractingWithLiterals(dot, lhs_literal,
-                                                     rhs_literal, gs_literal);
+    });
   }
 
   // This is currently only implemented for the ragged dimension being a
@@ -1745,88 +1445,44 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
   // non-batch, non-contracting dimension). For the batch mode, this will throw
   // an unimplemented error.
   absl::Status HandleRaggedDot(const HloInstruction* dot) override {
-    auto lhs = dot->operand(0);
-    auto rhs = dot->operand(1);
-    auto group_sizes = dot->operand(2);
-
-    CHECK(dot->shape().IsArray());
-    CHECK(lhs->shape().IsArray());
-    CHECK(rhs->shape().IsArray());
-    CHECK(group_sizes->shape().IsArray());
-
-    const Literal& lhs_literal = parent_->GetEvaluatedLiteralFor(lhs);
-    const Literal& rhs_literal = parent_->GetEvaluatedLiteralFor(rhs);
-    const Literal& gs_literal = parent_->GetEvaluatedLiteralFor(group_sizes);
-
-    // LHS and RHS may have a different initial precision than our output.
-    const bool lhs_same =
-        ShapeUtil::SameElementType(lhs->shape(), dot->shape());
-    const bool rhs_same =
-        ShapeUtil::SameElementType(rhs->shape(), dot->shape());
-    // Upcast group sizes to S64, since they could be e.g. S32.
-    const bool gs_same =
-        group_sizes->shape().element_type() == PrimitiveType::S64;
-
-    return HandleRaggedDotWithLiterals(
+    return parent_->EvaluateRaggedDot(
         dot,
-        lhs_same
-            ? lhs_literal
-            : static_cast<const Literal&>(
-                  lhs_literal.Convert(dot->shape().element_type()).value()),
-        rhs_same
-            ? rhs_literal
-            : static_cast<const Literal&>(
-                  rhs_literal.Convert(dot->shape().element_type()).value()),
-        gs_same ? gs_literal
-                : static_cast<const Literal&>(
-                      gs_literal.Convert(PrimitiveType::S64).value()));
+        [this, dot](const Literal& lhs_literal, const Literal& rhs_literal,
+                    const Literal& gs_literal,
+                    const HloEvaluator::DotDimensionInfo& dot_info,
+                    int64_t ragged_dim_as_contracting_dim, Literal* result) {
+          return HandleRaggedDotContractingWithLiterals(
+              dot, lhs_literal, rhs_literal, gs_literal, dot_info,
+              ragged_dim_as_contracting_dim, result);
+        },
+        [this, dot](const Literal& lhs_literal, const Literal& rhs_literal,
+                    const Literal& gs_literal,
+                    const HloEvaluator::DotDimensionInfo& dot_info,
+                    Literal* result) {
+          return HandleRaggedDotBatchModeWithLiterals(
+              dot, lhs_literal, rhs_literal, gs_literal, dot_info, result);
+        },
+        [this, dot](const Literal& lhs_literal, const Literal& rhs_literal,
+                    const Literal& gs_literal,
+                    const HloEvaluator::DotDimensionInfo& dot_info,
+                    Literal* result) {
+          return HandleRaggedDotNonContractingWithLiterals(
+              dot, lhs_literal, rhs_literal, gs_literal, dot_info, result);
+        });
   }
 
   absl::Status HandleScaledDot(const HloInstruction* dot) override {
-    auto lhs = dot->operand(0);
-    auto rhs = dot->operand(1);
-    auto lhs_scale = dot->operand(2);
-    auto rhs_scale = dot->operand(3);
-    CHECK(dot->shape().IsArray());
-    CHECK(lhs->shape().IsArray());
-    CHECK(rhs->shape().IsArray());
-    CHECK(lhs_scale->shape().IsArray());
-    CHECK(rhs_scale->shape().IsArray());
-    CHECK(lhs_scale->shape().dimensions().size() == 0 ||
-          lhs_scale->shape().dimensions().size() ==
-              lhs->shape().dimensions().size());
-    CHECK(rhs_scale->shape().dimensions().size() == 0 ||
-          rhs_scale->shape().dimensions().size() ==
-              rhs->shape().dimensions().size());
-    ABSL_ASSIGN_OR_RETURN(const Literal lhs_literal,
-                     parent_->GetEvaluatedLiteralFor(lhs).Convert(
-                         dot->shape().element_type()));
-    ABSL_ASSIGN_OR_RETURN(const Literal rhs_literal,
-                     parent_->GetEvaluatedLiteralFor(rhs).Convert(
-                         dot->shape().element_type()));
-
-    // If the scale is a scalar, we can just use 1.0. Otherwise, we need to
-    // evaluate the scale.
-    auto evaluate_scale =
-        [&](const HloInstruction* operand,
-            const HloInstruction* scale) -> absl::StatusOr<Literal> {
-      if (scale->shape().IsArray() && scale->shape().dimensions().size() > 0) {
-        ABSL_ASSIGN_OR_RETURN(Literal scale_literal,
-                         parent_->GetEvaluatedLiteralFor(scale).Convert(
-                             dot->shape().element_type()));
-        return scale_literal;
-      }
-      std::vector<int64_t> ones(operand->shape().dimensions().size(), 1);
-      Shape scale_shape =
-          ShapeUtil::MakeShape(dot->shape().element_type(), ones);
-      Literal scale_literal = Literal::CreateFromShape(scale_shape);
-      scale_literal.PopulateWithValue(static_cast<ReturnT>(1.0f));
-      return scale_literal;
-    };
-    ABSL_ASSIGN_OR_RETURN(Literal lhs_scale_literal, evaluate_scale(lhs, lhs_scale));
-    ABSL_ASSIGN_OR_RETURN(Literal rhs_scale_literal, evaluate_scale(rhs, rhs_scale));
-    return HandleScaledDotSlowPathWithLiterals(
-        dot, lhs_literal, rhs_literal, lhs_scale_literal, rhs_scale_literal);
+    return parent_->EvaluateScaledDot(
+        dot, [this, dot](const Literal& lhs_literal, const Literal& rhs_literal,
+                         const Literal& lhs_scale_literal,
+                         const Literal& rhs_scale_literal,
+                         const ShapeInfo& lhs_info, const ShapeInfo& rhs_info,
+                         int64_t total_contraction_size, Literal* result) {
+          return HandleScaledDotSlowPathWithLiterals(
+              dot, lhs_literal, rhs_literal, lhs_scale_literal,
+              rhs_scale_literal, lhs_info, rhs_info, total_contraction_size,
+              result);
+        });
   }
 
  private:
@@ -1835,26 +1491,12 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
   absl::Status HandleScaledDotSlowPathWithLiterals(
       const HloInstruction* dot, const Literal& lhs_literal,
       const Literal& rhs_literal, const Literal& lhs_scale_literal,
-      const Literal& rhs_scale_literal) {
+      const Literal& rhs_scale_literal, const ShapeInfo& lhs_info,
+      const ShapeInfo& rhs_info, int64_t total_contraction_size,
+      Literal* result) {
     const auto& dnums = dot->dot_dimension_numbers();
-    CHECK(ShapeUtil::SameElementType(lhs_literal.shape(), rhs_literal.shape()));
-    CHECK(ShapeUtil::SameElementType(lhs_literal.shape(), dot->shape()));
-
-    CHECK_EQ(dnums.lhs_batch_dimensions_size(),
-             dnums.rhs_batch_dimensions_size());
-
-    ShapeInfo lhs_info(lhs_literal, lhs_scale_literal,
-                       dnums.lhs_contracting_dimensions(),
-                       dnums.lhs_batch_dimensions());
-    ShapeInfo rhs_info(rhs_literal, rhs_scale_literal,
-                       dnums.rhs_contracting_dimensions(),
-                       dnums.rhs_batch_dimensions());
-    const int64_t total_contraction_size =
-        Product(lhs_info.contracting_dim_sizes);
-    Shape dot_shape = GetShapeWithLayout(dot->shape());
-
-    ABSL_ASSIGN_OR_RETURN(Literal result, Literal::Make(dot_shape));
-    ABSL_RETURN_IF_ERROR(result.PopulateParallel<ReturnT>(
+    const Shape& dot_shape = result->shape();
+    return result->PopulateParallel<ReturnT>(
         [&](absl::Span<const int64_t> result_index, int /*thread_id*/) {
           // Locations in LHS and RHS that we read from.
           DimensionVector lhs_index(lhs_info.rank);
@@ -1947,97 +1589,53 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
           }
 
           return static_cast<ReturnT>(result_val);
-        }));
-
-    parent_->SetEvaluatedLiteralFor(dot, std::move(result));
-    return absl::OkStatus();
+        });
   }
 
  public:
   absl::Status HandlePad(const HloInstruction* pad) override {
-    CHECK(pad->operand(0)->shape().IsArray());
-    // Padding value must be scalar.
-    CHECK(ShapeUtil::IsScalar(pad->operand(1)->shape()));
-    CHECK_EQ(pad->operand(0)->shape().dimensions().size(),
-             pad->padding_config().dimensions_size());
+    return parent_->EvaluatePad(
+        pad,
+        [pad](const Literal& evaluated_operand,
+              const Literal& evaluated_padding,
+              absl::Span<int64_t> target_index,
+              absl::Span<const int64_t> zero_base,
+              absl::Span<const int64_t> step, Literal* result) -> absl::Status {
+          const ReturnT scalar = evaluated_padding.Get<ReturnT>({});
+          ABSL_RETURN_IF_ERROR(result->PopulateLinearParallel<ReturnT>(
+              [&scalar](int64_t linear_index, int) { return scalar; }));
 
-    ABSL_ASSIGN_OR_RETURN(auto inferred_return_shape,
-                     ShapeInference::InferPadShape(
-                         /*operand_shape=*/pad->operand(0)->shape(),
-                         /*padding_value_shape=*/pad->operand(1)->shape(),
-                         /*padding_config=*/pad->padding_config()));
-    CHECK(ShapeUtil::CompatibleIgnoringElementType(pad->shape(),
-                                                   inferred_return_shape))
-        << "return shape is set to: " << ShapeUtil::HumanString(pad->shape())
-        << " but is inferred to be: "
-        << ShapeUtil::HumanString(inferred_return_shape);
-    ReturnT scalar;
-    PrimitiveType result_type = pad->shape().element_type();
-    PrimitiveType padding_type = pad->operand(1)->shape().element_type();
-    if (padding_type != result_type) {
-      ABSL_ASSIGN_OR_RETURN(auto literal,
-                       parent_->GetEvaluatedLiteralFor(pad->operand(1))
-                           .Convert(result_type));
-      scalar = literal.Get<ReturnT>({});
-    } else {
-      scalar =
-          parent_->GetEvaluatedLiteralFor(pad->operand(1)).Get<ReturnT>({});
-    }
+          // Loop through each element of the operand, assign them to the
+          // corresponding index of the resulting padded literal.
+          const PaddingConfig& pad_config = pad->padding_config();
 
-    // Create new HLO of padded shape with padding value.
-    Literal result(GetShapeWithLayout(pad->shape()));
-    ABSL_RETURN_IF_ERROR(result.PopulateLinearParallel<ReturnT>(
-        [&scalar](int64_t linear_index, int) { return scalar; }));
+          auto func = [&](absl::Span<const int64_t> input_index) {
+            for (auto i = 0; i < input_index.size(); ++i) {
+              // Interior padding occurs logically before edge padding, so in
+              // the case of negative edge padding elements are removed from the
+              // interior-padded operand.
+              target_index[i] =
+                  pad_config.dimensions(i).edge_padding_low() +
+                  input_index[i] *
+                      (pad_config.dimensions(i).interior_padding() + 1);
 
-    Literal converted_operand;
-    PrimitiveType operand_type = pad->operand(0)->shape().element_type();
-    if (operand_type != result_type) {
-      ABSL_ASSIGN_OR_RETURN(converted_operand,
-                       parent_->GetEvaluatedLiteralFor(pad->operand(0))
-                           .Convert(result_type));
-    }
-    const Literal& evaluated_operand =
-        operand_type != result_type
-            ? converted_operand
-            : parent_->GetEvaluatedLiteralFor(pad->operand(0));
+              // Account for negative low and high padding: skip assignment if
+              // the any target index is out of range.
+              if (!(target_index[i] >= 0 &&
+                    target_index[i] < pad->shape().dimensions(i))) {
+                return true;
+              }
+            }
+            result->Set<ReturnT>(target_index,
+                                 evaluated_operand.Get<ReturnT>(input_index));
+            return true;
+          };
 
-    std::vector<int64_t> target_index(result.shape().dimensions().size(), 0);
-
-    // Loop through each element of the operand, assign them to the
-    // corresponding index of the resulting padded literal.
-    const PaddingConfig& pad_config = pad->padding_config();
-
-    auto func = [&](absl::Span<const int64_t> input_index) {
-      for (auto i = 0; i < input_index.size(); ++i) {
-        // Interior padding occurs logically before edge padding, so in the case
-        // of negative edge padding elements are removed from the
-        // interior-padded operand.
-        target_index[i] =
-            pad_config.dimensions(i).edge_padding_low() +
-            input_index[i] * (pad_config.dimensions(i).interior_padding() + 1);
-
-        // Account for negative low and high padding: skip assignment if the
-        // any target index is out of range.
-        if (!(target_index[i] >= 0 &&
-              target_index[i] < pad->shape().dimensions(i))) {
-          return true;
-        }
-      }
-      result.Set<ReturnT>(target_index,
-                          evaluated_operand.Get<ReturnT>(input_index));
-      return true;
-    };
-
-    std::vector<int64_t> zero_base(
-        evaluated_operand.shape().dimensions().size(), 0);
-    std::vector<int64_t> step(evaluated_operand.shape().dimensions().size(), 1);
-
-    ShapeUtil::ForEachIndexNoStatus(evaluated_operand.shape(), zero_base,
-                                    evaluated_operand.shape().dimensions(),
-                                    step, func);
-
-    parent_->SetEvaluatedLiteralFor(pad, std::move(result));
-    return absl::OkStatus();
+          ShapeUtil::ForEachIndexNoStatus(
+              evaluated_operand.shape(), zero_base,
+              evaluated_operand.shape().dimensions(), step, func);
+          return absl::OkStatus();
+        });
   }
 
   absl::Status HandleClz(const HloInstruction* clz) override {

From d038ceef421a5a010863993e6237ef272f2629b7 Mon Sep 17 00:00:00 2001
From: Ce Zheng CA <zce@google.com>
Date: Sun, 11 Oct 2026 14:08:33 -0700
Subject: [PATCH 0475/1000] [XLA] Move shape and type validation from inline
 Populate templates to PopulateInplaceInternal

Move the dense-array `TF_RET_CHECK(shape().IsArray())` check, the element-type `TF_RET_CHECK` check, and the `element_count() > 32` parallelism threshold check out of the 10 inline `MutableLiteralBase::Populate*` template methods in `xla/literal.h` into the out-of-line `PopulateInplaceInternal` and `PopulateLinearInplaceInternal` implementations in `xla/literal.cc`.

Previously, because `Populate*` and `Populate*Internal` are templated on both `<typename NativeT, typename Generator>`, every caller lambda instantiated duplicate inline `TF_RET_CHECK` error-stream formatting (`MakeErrorStream`, `Shape` printing, and `LowercasePrimitiveTypeName`) and a duplicate `ShapeUtil::ElementsIn(shape())` call prior to entering `Populate*InplaceInternal`. Keeping `Generator&&` templated for the `absl::FunctionRef` populator callback while performing validation once inside `Populate*InplaceInternal` eliminates this per-call-site template bloat without adding any inner-loop indirection.

Runtime & Memory Performance Impact: Neutral to slightly faster CPU performance (`-0.76%` geomean `cpu/op` / `-0.99%` `time/op` across `BM_Populate*` in `xla:literal_test`) and zero dynamic memory change; the generator callable remains inlined into the `Populate*InplaceInternal` callback with zero additional per-element indirection while avoiding a redundant `ShapeUtil::ElementsIn(shape())` call on parallel populate entry.

PiperOrigin-RevId: 997548347
---
 third_party/xla/xla/literal.cc | 37 +++++++++++++----
 third_party/xla/xla/literal.h  | 75 ++++++++--------------------------
 2 files changed, 48 insertions(+), 64 deletions(-)

diff --git a/third_party/xla/xla/literal.cc b/third_party/xla/xla/literal.cc
index 02deeaa916f..e270cc6d0c0 100644
--- a/third_party/xla/xla/literal.cc
+++ b/third_party/xla/xla/literal.cc
@@ -946,18 +946,27 @@ void MutableLiteralBase::PopulateR1(const tsl::core::Bitmap& values) {
   }
 }
 
-void MutableLiteralBase::PopulateInplaceInternal(
+absl::Status MutableLiteralBase::PopulateInplaceInternal(
     absl::FunctionRef<void(void*, absl::Span<const int64_t>, int)> populator,
-    bool parallel) {
+    bool parallel, PrimitiveType element_type) {
   const Shape& this_shape = shape();
+  TF_RET_CHECK(this_shape.IsArray())
+      << __func__ << " is only supported for dense arrays: " << this_shape;
+  if (element_type != PRIMITIVE_TYPE_INVALID) {
+    TF_RET_CHECK(this_shape.element_type() == element_type)
+        << "Failing to populate literal with element type "
+        << primitive_util::LowercasePrimitiveTypeName(this_shape.element_type())
+        << " using data of type "
+        << primitive_util::LowercasePrimitiveTypeName(element_type);
+  }
   const int64_t rank = this_shape.dimensions().size();
-  DCHECK(this_shape.IsArray());
   char* const dest_base = static_cast<char*>(untyped_data());
   if (rank > 0) {
     StrideConfig stride_config(this_shape, this_shape, this_shape.dimensions());
     const int64_t primitive_size =
         ShapeUtil::ByteSizeOfPrimitiveType(shape().element_type());
     const int64_t num_elements = ShapeUtil::ElementsIn(shape());
+    parallel = parallel && num_elements > 32;
     // If we are rank-1 and we are `parallel`, it is better to use a smaller
     // `step` than what `StrideConfig` does: stick the entire dimension in the
     // inner-most loop.
@@ -1009,17 +1018,30 @@ void MutableLiteralBase::PopulateInplaceInternal(
     // For scalars.
     populator(dest_base, {}, /*thread_id=*/-1);
   }
+  return absl::OkStatus();
 }
 
-void MutableLiteralBase::PopulateLinearInplaceInternal(
-    absl::FunctionRef<void(void*, int64_t, int)> populator, bool parallel) {
+absl::Status MutableLiteralBase::PopulateLinearInplaceInternal(
+    absl::FunctionRef<void(void*, int64_t, int)> populator, bool parallel,
+    PrimitiveType element_type) {
   const Shape& this_shape = shape();
+  TF_RET_CHECK(this_shape.IsArray())
+      << __func__ << " is only supported for dense arrays: " << this_shape;
+  if (element_type != PRIMITIVE_TYPE_INVALID) {
+    TF_RET_CHECK(this_shape.element_type() == element_type)
+        << "Failing to populate literal with element type "
+        << primitive_util::LowercasePrimitiveTypeName(this_shape.element_type())
+        << " using data of type "
+        << primitive_util::LowercasePrimitiveTypeName(element_type);
+  }
   const int64_t rank = this_shape.dimensions().size();
-  DCHECK(this_shape.IsArray());
   char* const dest_base = static_cast<char*>(untyped_data());
 
   const int64_t num_elements = ShapeUtil::ElementsIn(shape());
-  if (num_elements == 0) return;
+  if (num_elements == 0) {
+    return absl::OkStatus();
+  }
+  parallel = parallel && num_elements > 32;
 
   if (rank > 0) {
     // Compute initialization function partitioning.
@@ -1072,6 +1094,7 @@ void MutableLiteralBase::PopulateLinearInplaceInternal(
     // For scalars.
     populator(dest_base, 0, /*thread_id=*/-1);
   }
+  return absl::OkStatus();
 }
 
 Literal LiteralBase::Relayout(const Layout& new_layout,
diff --git a/third_party/xla/xla/literal.h b/third_party/xla/xla/literal.h
index 91fcaf6850b..dcce53e8aef 100644
--- a/third_party/xla/xla/literal.h
+++ b/third_party/xla/xla/literal.h
@@ -1542,16 +1542,17 @@ class MutableLiteralBase : public LiteralBase {
   // Implementation details shared between Populate() and PopulateParallel().
   template <typename NativeT, typename Generator>
   absl::Status PopulateInternal(Generator&& generator, bool parallel);
-  void PopulateInplaceInternal(
+  absl::Status PopulateInplaceInternal(
       absl::FunctionRef<void(void*, absl::Span<const int64_t>, int)> populator,
-      bool parallel);
+      bool parallel, PrimitiveType element_type = PRIMITIVE_TYPE_INVALID);
 
   // Implementation details shared between PopulateLinear() and
   // PopulateLinearParallel().
   template <typename NativeT, typename Generator>
   absl::Status PopulateLinearInternal(Generator&& generator, bool parallel);
-  void PopulateLinearInplaceInternal(
-      absl::FunctionRef<void(void*, int64_t, int)> populator, bool parallel);
+  absl::Status PopulateLinearInplaceInternal(
+      absl::FunctionRef<void(void*, int64_t, int)> populator, bool parallel,
+      PrimitiveType element_type = PRIMITIVE_TYPE_INVALID);
 
   friend class LiteralBase;
   friend class MutableBorrowingLiteral;
@@ -2190,28 +2191,16 @@ void MutableLiteralBase::PopulateR4FromArray4D(const Array4D<NativeT>& values) {
 template <typename NativeT, typename Generator>
 absl::Status MutableLiteralBase::PopulateInternal(Generator&& generator,
                                                   bool parallel) {
-  const Shape& this_shape = shape();
-  DCHECK(this_shape.IsArray());
-  TF_RET_CHECK(this_shape.element_type() ==
-               primitive_util::NativeToPrimitiveType<NativeT>())
-      << "Failing to populate literal with element type "
-      << primitive_util::LowercasePrimitiveTypeName(this_shape.element_type())
-      << " using data of type "
-      << primitive_util::LowercasePrimitiveTypeName(
-             primitive_util::NativeToPrimitiveType<NativeT>());
-  PopulateInplaceInternal(
+  return PopulateInplaceInternal(
       [&](void* dest, absl::Span<const int64_t> indices, int thread_id) {
         *static_cast<NativeT*>(dest) = generator(indices, thread_id);
       },
-      parallel);
-  return absl::OkStatus();
+      parallel, primitive_util::NativeToPrimitiveType<NativeT>());
 }
 
 template <typename NativeT, typename Generator,
           MutableLiteralBase::IsGenerator<NativeT, Generator>*>
 absl::Status MutableLiteralBase::Populate(Generator&& generator) {
-  TF_RET_CHECK(shape().IsArray())
-      << __func__ << " is only supported for dense arrays: " << shape();
   return PopulateInternal<NativeT>(
       [&](absl::Span<const int64_t> indexes, int /*thread_id*/) {
         return generator(indexes);
@@ -2221,37 +2210,23 @@ absl::Status MutableLiteralBase::Populate(Generator&& generator) {
 template <typename NativeT, typename Generator,
           MutableLiteralBase::IsParallelGenerator<NativeT, Generator>*>
 absl::Status MutableLiteralBase::PopulateParallel(Generator&& generator) {
-  TF_RET_CHECK(shape().IsArray())
-      << __func__ << " is only supported for dense arrays: " << shape();
-  return PopulateInternal<NativeT>(generator,
-                                   /*parallel=*/element_count() > 32);
+  return PopulateInternal<NativeT>(std::forward<Generator>(generator),
+                                   /*parallel=*/true);
 }
 
 template <typename NativeT, typename Generator>
 absl::Status MutableLiteralBase::PopulateLinearInternal(Generator&& generator,
                                                         bool parallel) {
-  const Shape& this_shape = shape();
-  DCHECK(this_shape.IsArray());
-  TF_RET_CHECK(this_shape.element_type() ==
-               primitive_util::NativeToPrimitiveType<NativeT>())
-      << "Failing to populate literal with element type "
-      << primitive_util::LowercasePrimitiveTypeName(this_shape.element_type())
-      << " using data of type "
-      << primitive_util::LowercasePrimitiveTypeName(
-             primitive_util::NativeToPrimitiveType<NativeT>());
-  PopulateLinearInplaceInternal(
+  return PopulateLinearInplaceInternal(
       [&](void* dest, int64_t linear_index, int thread_id) {
         *static_cast<NativeT*>(dest) = generator(linear_index, thread_id);
       },
-      parallel);
-  return absl::OkStatus();
+      parallel, primitive_util::NativeToPrimitiveType<NativeT>());
 }
 
 template <typename NativeT, typename Generator,
           MutableLiteralBase::IsLinearGenerator<NativeT, Generator>*>
 absl::Status MutableLiteralBase::PopulateLinear(Generator&& generator) {
-  TF_RET_CHECK(shape().IsArray())
-      << __func__ << " is only supported for dense arrays: " << shape();
   return PopulateLinearInternal<NativeT>(
       [&](int64_t linear_index, int /*thread_id*/) {
         return generator(linear_index);
@@ -2261,56 +2236,42 @@ absl::Status MutableLiteralBase::PopulateLinear(Generator&& generator) {
 template <typename NativeT, typename Generator,
           MutableLiteralBase::IsLinearParallelGenerator<NativeT, Generator>*>
 absl::Status MutableLiteralBase::PopulateLinearParallel(Generator&& generator) {
-  TF_RET_CHECK(shape().IsArray())
-      << __func__ << " is only supported for dense arrays: " << shape();
   return PopulateLinearInternal<NativeT>(std::forward<Generator>(generator),
-                                         /*parallel=*/element_count() > 32);
+                                         /*parallel=*/true);
 }
 
 template <typename Populator, MutableLiteralBase::IsPopulator<Populator>*>
 absl::Status MutableLiteralBase::PopulateInplace(Populator&& populator) {
-  TF_RET_CHECK(shape().IsArray())
-      << __func__ << " is only supported for dense arrays: " << shape();
-  PopulateInplaceInternal(
+  return PopulateInplaceInternal(
       [&](void* dest, absl::Span<const int64_t> indexes, int /*thread_id*/) {
         return populator(dest, indexes);
       },
       /*parallel=*/false);
-  return absl::OkStatus();
 }
 
 template <typename Populator,
           MutableLiteralBase::IsParallelPopulator<Populator>*>
 absl::Status MutableLiteralBase::PopulateInplaceParallel(
     Populator&& populator) {
-  TF_RET_CHECK(shape().IsArray())
-      << __func__ << " is only supported for dense arrays: " << shape();
-  PopulateInplaceInternal(std::forward<Populator>(populator),
-                          /*parallel=*/element_count() > 32);
-  return absl::OkStatus();
+  return PopulateInplaceInternal(std::forward<Populator>(populator),
+                                 /*parallel=*/true);
 }
 
 template <typename Populator, MutableLiteralBase::IsLinearPopulator<Populator>*>
 absl::Status MutableLiteralBase::PopulateLinearInplace(Populator&& populator) {
-  TF_RET_CHECK(shape().IsArray())
-      << __func__ << " is only supported for dense arrays: " << shape();
-  PopulateLinearInplaceInternal(
+  return PopulateLinearInplaceInternal(
       [&](void* dest, int64_t linear_index, int /*thread_id*/) {
         return populator(dest, linear_index);
       },
       /*parallel=*/false);
-  return absl::OkStatus();
 }
 
 template <typename Populator,
           MutableLiteralBase::IsLinearParallelPopulator<Populator>*>
 absl::Status MutableLiteralBase::PopulateLinearInplaceParallel(
     Populator&& populator) {
-  TF_RET_CHECK(shape().IsArray())
-      << __func__ << " is only supported for dense arrays: " << shape();
-  PopulateLinearInplaceInternal(std::forward<Populator>(populator),
-                                /*parallel=*/element_count() > 32);
-  return absl::OkStatus();
+  return PopulateLinearInplaceInternal(std::forward<Populator>(populator),
+                                       /*parallel=*/true);
 }
 
 template <typename NativeT>

From 7b92f543cf49a07c1bb229979342bff1171b418c Mon Sep 17 00:00:00 2001
From: Blake Hechtman <blakehechtman@google.com>
Date: Sun, 11 Oct 2026 15:55:08 -0700
Subject: [PATCH 0476/1000] [XLA] Precompute HloValue uses in
 HloRematerialization (NFC)

Enable precompute_uses in HloDataflowAnalysis::Run inside HloRematerialization and simplify FusedParameterUsesBuffer to use HloValue::IsRootOf and HloValue::GetUses.

PiperOrigin-RevId: 997572471
---
 .../simplifiers/hlo_rematerialization.cc      | 43 ++++++-------------
 1 file changed, 14 insertions(+), 29 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc
index 5fcc642cf8d..34fe885e638 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc
@@ -243,31 +243,15 @@ bool FusedParameterUsesBuffer(const HloInstruction* fusion, int64_t op_idx,
   }
   const HloValue& value =
       dataflow_analysis.GetValueDefinedAt(fused_param, shape_idx);
-  for (const HloPosition& alias : value.positions()) {
-    // If an alias of the buffer is the root of the fusion, the buffer is
-    // an output of the fusion and therefore assumed to be used.
-    if (alias.instruction == fusion->fused_expression_root()) {
+  if (value.IsRootOf(fusion->fused_instructions_computation())) {
+    return true;
+  }
+  for (const HloUse& use : value.GetUses()) {
+    if (!use.instruction->IsLoopFusion() ||
+        FusedParameterUsesBuffer(use.instruction, use.operand_number,
+                                 use.operand_index, dataflow_analysis)) {
       return true;
     }
-    for (const HloInstruction* alias_user : alias.instruction->users()) {
-      if (alias_user->opcode() == HloOpcode::kGetTupleElement &&
-          !alias.index.empty()) {
-        // GetTupleElement instructions only access the top-level buffer of
-        // their operand.
-        continue;
-      }
-      if (!alias_user->IsLoopFusion()) {
-        return true;
-      }
-      for (int64_t nested_op_idx = 0;
-           nested_op_idx < alias_user->operand_count(); ++nested_op_idx) {
-        if (alias_user->operand(nested_op_idx) == alias.instruction &&
-            FusedParameterUsesBuffer(alias_user, nested_op_idx, alias.index,
-                                     dataflow_analysis)) {
-          return true;
-        }
-      }
-    }
   }
   return false;
 }
@@ -2573,12 +2557,13 @@ absl::Status HloRematerialization::UpdateDataflowAnalysis(HloModule* module) {
   if (dataflow_analysis_ == nullptr) {
     ABSL_ASSIGN_OR_RETURN(
         dataflow_analysis_,
-        HloDataflowAnalysis::Run(*module, /*ssa_form=*/false,
-                                 /*bitcast_defines_value=*/false,
-                                 /*execution_threads=*/{},
-                                 /*propagate_through_calls=*/false,
-                                 /*precompute_uses=*/std::nullopt,
-                                 /*propagate_through_control_flow=*/false));
+        HloDataflowAnalysis::Run(
+            *module, /*ssa_form=*/false,
+            /*bitcast_defines_value=*/false,
+            /*execution_threads=*/{},
+            /*propagate_through_calls=*/false,
+            /*precompute_uses=*/[](const HloValue&) { return true; },
+            /*propagate_through_control_flow=*/false));
   }
   return absl::OkStatus();
 }

From 557cf192968ef6133b61c413620e5440ec9885fd Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 19:36:53 -0700
Subject: [PATCH 0477/1000] Automated Code Change

PiperOrigin-RevId: 997628228
---
 tensorflow/cc/framework/gradient_checker.cc |  4 +--
 tensorflow/cc/framework/ops.h               | 32 ++++++++++-----------
 2 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/tensorflow/cc/framework/gradient_checker.cc b/tensorflow/cc/framework/gradient_checker.cc
index 49882584e8b..33c335fc405 100644
--- a/tensorflow/cc/framework/gradient_checker.cc
+++ b/tensorflow/cc/framework/gradient_checker.cc
@@ -69,14 +69,14 @@ SET_BASE_UNITS_FOR_TYPE(complex128, ({{1, 0}, {0, 1}}));
 // When T is complex, it assigns the real and complex values to successive rows
 // or columns in the matrix depending on the expand_by_row parameter
 template <typename T, typename JAC_T>
-typename std::enable_if<std::is_floating_point<T>::value>::type SetJacobian(
+std::enable_if_t<std::is_floating_point_v<T>> SetJacobian(
     typename TTypes<JAC_T>::Matrix* jacobian, const int row, const int col,
     const T& value, const bool expand_by_row) {
   (*jacobian)(row, col) = JAC_T{value};
 }
 
 template <typename T, typename JAC_T>
-typename std::enable_if<is_complex<T>::value>::type SetJacobian(
+std::enable_if_t<is_complex<T>::value> SetJacobian(
     typename TTypes<JAC_T>::Matrix* jacobian, const int row, const int col,
     const T& value, const bool expand_by_row) {
   (*jacobian)(row, col) = JAC_T{value.real()};
diff --git a/tensorflow/cc/framework/ops.h b/tensorflow/cc/framework/ops.h
index 784884d8bc1..ced8b3957b7 100644
--- a/tensorflow/cc/framework/ops.h
+++ b/tensorflow/cc/framework/ops.h
@@ -112,9 +112,9 @@ class Input {
   struct Initializer {
     /// Construct from a scalar value of an arithmetic type or a type that can
     /// be converted to a string (eg. a string literal).
-    template <typename T, typename = typename std::enable_if<
-                              std::is_arithmetic<T>::value ||
-                              std::is_convertible<T, std::string>::value>::type>
+    template <typename T, typename = std::enable_if_t<
+                              std::is_arithmetic_v<T> ||
+                              std::is_convertible_v<T, std::string>>>
     Initializer(const T& v) {  // NOLINT(runtime/explicit)
       typedef typename RealType<T>::type RealT;
       Tensor t(DataTypeToEnum<RealT>::v(), TensorShape());
@@ -125,9 +125,9 @@ class Input {
     Initializer(const Tensor& t) : tensor(t) {}  // NOLINT(runtime/explicit)
 
     /// Construct from a scalar value and an explicit shape
-    template <typename T, typename = typename std::enable_if<
-                              std::is_arithmetic<T>::value ||
-                              std::is_convertible<T, std::string>::value>::type>
+    template <typename T, typename = std::enable_if_t<
+                              std::is_arithmetic_v<T> ||
+                              std::is_convertible_v<T, std::string>>>
     Initializer(const T& v, const TensorShape& shape) {
       typedef typename RealType<T>::type RealT;
       Tensor t(DataTypeToEnum<RealT>::v(), shape);
@@ -138,9 +138,9 @@ class Input {
     }
 
     /// Construct from a initializer list of scalars (a one-dimensional tensor).
-    template <typename T, typename = typename std::enable_if<
-                              std::is_arithmetic<T>::value ||
-                              std::is_convertible<T, std::string>::value>::type>
+    template <typename T, typename = std::enable_if_t<
+                              std::is_arithmetic_v<T> ||
+                              std::is_convertible_v<T, std::string>>>
     Initializer(
         const std::initializer_list<T>& v) {  // NOLINT(runtime/explicit)
       typedef typename RealType<T>::type RealT;
@@ -151,9 +151,9 @@ class Input {
     }
 
     /// Construct from a initializer list of scalars and an explicit shape.
-    template <typename T, typename = typename std::enable_if<
-                              std::is_arithmetic<T>::value ||
-                              std::is_convertible<T, std::string>::value>::type>
+    template <typename T, typename = std::enable_if_t<
+                              std::is_arithmetic_v<T> ||
+                              std::is_convertible_v<T, std::string>>>
     Initializer(const std::initializer_list<T>& v, const TensorShape& shape) {
       typedef typename RealType<T>::type RealT;
       Tensor t(DataTypeToEnum<RealT>::v(), shape);
@@ -175,7 +175,7 @@ class Input {
     Initializer(const std::initializer_list<Initializer>& v);
 
     // START_SKIP_DOXYGEN
-    template <typename T, bool = std::is_convertible<T, std::string>::value>
+    template <typename T, bool = std::is_convertible_v<T, std::string>>
     struct RealType {
       typedef tstring type;
     };
@@ -210,9 +210,9 @@ class Input {
   /// * A Tensor object.
   Input(const Output& o) : output_(o) {}  // NOLINT(runtime/explicit)
 
-  template <typename T, typename = typename std::enable_if<
-                            std::is_arithmetic<T>::value ||
-                            std::is_convertible<T, std::string>::value>::type>
+  template <typename T,
+            typename = std::enable_if_t<std::is_arithmetic_v<T> ||
+                                        std::is_convertible_v<T, std::string>>>
   Input(const T& v)  // NOLINT(runtime/explicit)
       : Input(Initializer(v)) {}
 

From b0b91e5c5eec794f544648b6fb88cd0f10f817e4 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 19:37:39 -0700
Subject: [PATCH 0478/1000] Automated Code Change

PiperOrigin-RevId: 997628504
---
 tensorflow/core/kernels/fused_batch_norm_op.cc    | 4 ++--
 tensorflow/core/kernels/fused_batch_norm_op.cu.cc | 2 +-
 tensorflow/core/kernels/gpu_prim_helpers.h        | 3 +--
 3 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/tensorflow/core/kernels/fused_batch_norm_op.cc b/tensorflow/core/kernels/fused_batch_norm_op.cc
index 8a98be1a71e..40fca98b99d 100644
--- a/tensorflow/core/kernels/fused_batch_norm_op.cc
+++ b/tensorflow/core/kernels/fused_batch_norm_op.cc
@@ -478,7 +478,7 @@ struct FusedBatchNormGrad<CPUDevice, T, U> {
 
     // Maybe allocate a temporary workspace of [rest_size, depth] shape.
     Tensor scratch_rest_by_depth;
-    if (std::is_same<T, U>::value) {
+    if (std::is_same_v<T, U>) {
       OP_REQUIRES(context,
                   scratch_rest_by_depth.CopyFrom(transformed_x_backprop_output,
                                                  {rest_size, depth}),
@@ -583,7 +583,7 @@ struct FusedBatchNormFreezeGrad<CPUDevice, T, U> {
 
     // Maybe allocate a temporary workspace of [rest_size, depth] shape.
     Tensor scratch3_tensor;
-    if (std::is_same<T, U>::value) {
+    if (std::is_same_v<T, U>) {
       OP_REQUIRES(
           context,
           scratch3_tensor.CopyFrom(*x_backprop_output, {rest_size, depth}),
diff --git a/tensorflow/core/kernels/fused_batch_norm_op.cu.cc b/tensorflow/core/kernels/fused_batch_norm_op.cu.cc
index 6951285eb49..39a685e6f96 100644
--- a/tensorflow/core/kernels/fused_batch_norm_op.cu.cc
+++ b/tensorflow/core/kernels/fused_batch_norm_op.cu.cc
@@ -147,7 +147,7 @@ struct FusedBatchNormInferenceKernel {
       U var_v = var[channel];
 
       U scaling_factor_v = rsqrt(var_v + epsilon) * scale_v;
-      static_assert(std::is_same<U, float>::value, "U data type must be float");
+      static_assert(std::is_same_v<U, float>, "U data type must be float");
       U shifted_v = fmaf(in_v - mean_v, scaling_factor_v, offset_v);
 
       if (add_side_input) {
diff --git a/tensorflow/core/kernels/gpu_prim_helpers.h b/tensorflow/core/kernels/gpu_prim_helpers.h
index cf24deaa48a..2c457914c18 100644
--- a/tensorflow/core/kernels/gpu_prim_helpers.h
+++ b/tensorflow/core/kernels/gpu_prim_helpers.h
@@ -180,8 +180,7 @@ absl::Status GpuInclusivePrefixSum(OpKernelContext* context, int size,
                                    InputIteratorT input,
                                    OutputIteratorT output) {
   static_assert(
-      !std::is_same<typename std::remove_reference<decltype(*input)>::type,
-                    bool>::value,
+      !std::is_same_v<std::remove_reference_t<decltype(*input)>, bool>,
       "GpuInclusivePrefixSum does not work correct with booleans, please use "
       "TransformInputIterator to explicitly cast to an integer.");
   if (size == 0) return absl::OkStatus();

From b5d39d2a795079bcd8fda8647e1f435d97d08c27 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 19:39:39 -0700
Subject: [PATCH 0479/1000] Automated Code Change

PiperOrigin-RevId: 997629308
---
 tensorflow/core/grappler/utils/graph_view_internal.h |  6 ++----
 tensorflow/core/grappler/utils/graph_view_test.cc    | 11 +++++------
 2 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/tensorflow/core/grappler/utils/graph_view_internal.h b/tensorflow/core/grappler/utils/graph_view_internal.h
index 5ae837ede96..ff965b8e5e6 100644
--- a/tensorflow/core/grappler/utils/graph_view_internal.h
+++ b/tensorflow/core/grappler/utils/graph_view_internal.h
@@ -123,8 +123,7 @@ template <typename FaninViewT, typename FanoutViewT, typename GraphViewT,
           bool IsConst>
 class NodeViewInternal {
  private:
-  using NodeDefT =
-      typename std::conditional<IsConst, const NodeDef, NodeDef>::type;
+  using NodeDefT = std::conditional_t<IsConst, const NodeDef, NodeDef>;
 
  public:
   explicit NodeViewInternal(GraphViewT* graph_view, int node_index)
@@ -275,8 +274,7 @@ template <typename NodeViewT, typename FaninViewT, typename FanoutViewT,
           bool IsConst>
 class GraphViewInternal {
  private:
-  using GraphDefT =
-      typename std::conditional<IsConst, const GraphDef, GraphDef>::type;
+  using GraphDefT = std::conditional_t<IsConst, const GraphDef, GraphDef>;
 
  public:
   explicit GraphViewInternal(GraphDefT* graph) : graph_(graph) {}
diff --git a/tensorflow/core/grappler/utils/graph_view_test.cc b/tensorflow/core/grappler/utils/graph_view_test.cc
index 3857db598ee..5570e845a5a 100644
--- a/tensorflow/core/grappler/utils/graph_view_test.cc
+++ b/tensorflow/core/grappler/utils/graph_view_test.cc
@@ -49,8 +49,7 @@ GraphDef SimpleTestGraph() {
 
 template <typename T>
 const std::string GetGraphViewTypeAsString() {
-  return std::is_same<T, class GraphView>::value ? "GraphView"
-                                                 : "MutableGraphView";
+  return std::is_same_v<T, class GraphView> ? "GraphView" : "MutableGraphView";
 }
 
 using GraphViewTypes = ::testing::Types<GraphView, MutableGraphView>;
@@ -412,7 +411,7 @@ TYPED_TEST(TypedFaninTest, GetControllingFanins) {
   const auto& c_fanins = c_node->GetControllingFanins();
   FanoutViewType d_control_fanin(&graph_view, d_node->node_index(),
                                  Graph::kControlSlot);
-  if (std::is_same<GraphViewType, GraphView>::value) {
+  if (std::is_same_v<GraphViewType, GraphView>) {
     ASSERT_EQ(c_fanins.size(), 2);
     EXPECT_EQ(c_fanins[0], d_control_fanin);
     EXPECT_EQ(c_fanins[1], d_control_fanin);
@@ -545,7 +544,7 @@ TYPED_TEST(TypedFanoutTest, GetControlledFanouts) {
   const auto& d_fanouts = d_node->GetControlledFanouts();
   FaninViewType c_control_fanout(&graph_view, c_node->node_index(),
                                  Graph::kControlSlot);
-  if (std::is_same<GraphViewType, GraphView>::value) {
+  if (std::is_same_v<GraphViewType, GraphView>) {
     ASSERT_EQ(d_fanouts.size(), 2);
     EXPECT_EQ(d_fanouts[0], c_control_fanout);
     EXPECT_EQ(d_fanouts[1], c_control_fanout);
@@ -598,7 +597,7 @@ TYPED_TEST(TypedNodeViewTest, NumControllingFanins) {
 
   EXPECT_EQ(a_node->NumControllingFanins(), 1);
   EXPECT_EQ(b_node->NumControllingFanins(), 1);
-  if (std::is_same<TypeParam, GraphView>::value) {
+  if (std::is_same_v<TypeParam, GraphView>) {
     EXPECT_EQ(c_node->NumControllingFanins(), 2);
   } else {
     EXPECT_EQ(c_node->NumControllingFanins(), 1);
@@ -647,7 +646,7 @@ TYPED_TEST(TypedNodeViewTest, NumControlledFanouts) {
   EXPECT_EQ(a_node->NumControlledFanouts(), 0);
   EXPECT_EQ(b_node->NumControlledFanouts(), 0);
   EXPECT_EQ(c_node->NumControlledFanouts(), 2);
-  if (std::is_same<TypeParam, GraphView>::value) {
+  if (std::is_same_v<TypeParam, GraphView>) {
     EXPECT_EQ(d_node->NumControlledFanouts(), 2);
   } else {
     EXPECT_EQ(d_node->NumControlledFanouts(), 1);

From ab8907ee2d261394eeed87e325d602fc999cf186 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 19:40:15 -0700
Subject: [PATCH 0480/1000] Automated Code Change

PiperOrigin-RevId: 997629504
---
 tensorflow/lite/type_to_tflitetype_test.cc | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/tensorflow/lite/type_to_tflitetype_test.cc b/tensorflow/lite/type_to_tflitetype_test.cc
index 694ca91ee88..f415df02722 100644
--- a/tensorflow/lite/type_to_tflitetype_test.cc
+++ b/tensorflow/lite/type_to_tflitetype_test.cc
@@ -58,13 +58,12 @@ TEST(TypeToTfLiteType, Sanity) {
   EXPECT_EQ(kTfLiteFloat32, typeToTfLiteType<float>());
   EXPECT_EQ(kTfLiteBool, typeToTfLiteType<bool>());
   EXPECT_EQ(kTfLiteString, typeToTfLiteType<std::string>());
-  static_assert(
-      std::is_same<float, TfLiteTypeToType<kTfLiteFloat32>::Type>::value,
-      "TfLiteTypeToType test failure");
-  static_assert(std::is_same<bool, TfLiteTypeToType<kTfLiteBool>::Type>::value,
+  static_assert(std::is_same_v<float, TfLiteTypeToType<kTfLiteFloat32>::Type>,
+                "TfLiteTypeToType test failure");
+  static_assert(std::is_same_v<bool, TfLiteTypeToType<kTfLiteBool>::Type>,
                 "TfLiteTypeToType test failure");
   static_assert(
-      std::is_same<std::string, TfLiteTypeToType<kTfLiteString>::Type>::value,
+      std::is_same_v<std::string, TfLiteTypeToType<kTfLiteString>::Type>,
       "TfLiteTypeToType test failure");
 }
 

From c7353eee119b4f2a193392ea0c3ed0f57e7b3119 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 19:40:30 -0700
Subject: [PATCH 0481/1000] Automated Code Change

PiperOrigin-RevId: 997629609
---
 tensorflow/core/kernels/immutable_constant_op.cc     | 4 ++--
 tensorflow/core/kernels/initializable_lookup_table.h | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/kernels/immutable_constant_op.cc b/tensorflow/core/kernels/immutable_constant_op.cc
index 50dd94ed024..a221f5bd90b 100644
--- a/tensorflow/core/kernels/immutable_constant_op.cc
+++ b/tensorflow/core/kernels/immutable_constant_op.cc
@@ -24,7 +24,7 @@ namespace tensorflow {
 namespace {
 class MemmappedTensorAllocator : public Allocator {
  public:
-  MemmappedTensorAllocator() {}
+  MemmappedTensorAllocator() = default;
 
   absl::Status InitializeFromRegion(const std::string& name, Env* env) {
     const auto status =
@@ -110,7 +110,7 @@ void ImmutableConstantOp::Compute(OpKernelContext* ctx) {
   allocator.release()->set_delete_on_deallocate();
 }
 
-ImmutableConstantOp::~ImmutableConstantOp() {}
+ImmutableConstantOp::~ImmutableConstantOp() = default;
 constexpr char const* ImmutableConstantOp::kDTypeAttr;
 constexpr char const* ImmutableConstantOp::kShapeAttr;
 constexpr char const* ImmutableConstantOp::kMemoryRegionNameAttr;
diff --git a/tensorflow/core/kernels/initializable_lookup_table.h b/tensorflow/core/kernels/initializable_lookup_table.h
index 7d4a1cb7197..5b3bb1e68f3 100644
--- a/tensorflow/core/kernels/initializable_lookup_table.h
+++ b/tensorflow/core/kernels/initializable_lookup_table.h
@@ -110,9 +110,9 @@ class InitializableLookupTable : public LookupInterface {
   // This class is Thread-unsafe.
   class InitTableIterator {
    public:
-    InitTableIterator() {}
+    InitTableIterator() = default;
 
-    virtual ~InitTableIterator() {}
+    virtual ~InitTableIterator() = default;
 
     // Prepares the next batch of key and value tensors.
     virtual void Next() = 0;

From f659596038ea8d44b6474e12b7a9158a12447b41 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 19:40:40 -0700
Subject: [PATCH 0482/1000] Automated Code Change

PiperOrigin-RevId: 997629659
---
 tensorflow/c/experimental/filesystem/plugins/gcs/cleanup.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/cleanup.h b/tensorflow/c/experimental/filesystem/plugins/gcs/cleanup.h
index cc7a7451bb8..7e22553fc45 100644
--- a/tensorflow/c/experimental/filesystem/plugins/gcs/cleanup.h
+++ b/tensorflow/c/experimental/filesystem/plugins/gcs/cleanup.h
@@ -92,14 +92,14 @@ class Cleanup {
   bool is_released() const { return released_; }
 
  private:
-  static_assert(!std::is_reference<F>::value, "F must not be a reference");
+  static_assert(!std::is_reference_v<F>, "F must not be a reference");
 
   bool released_ = false;
   F f_;
 };
 
 template <int&... ExplicitParameterBarrier, typename F,
-          typename DecayF = typename std::decay<F>::type>
+          typename DecayF = std::decay_t<F>>
 Cleanup<DecayF> MakeCleanup(F&& f) {
   return Cleanup<DecayF>(std::forward<F>(f));
 }

From fd4ad047081054f4777545c22fea8152f8bf7d92 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 19:45:35 -0700
Subject: [PATCH 0483/1000] Automated Code Change

PiperOrigin-RevId: 997631224
---
 .../acceleration/mini_benchmark/libjpeg_decoder.h          | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder.h b/tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder.h
index 4121dfe8e43..6b99f81280e 100644
--- a/tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder.h
+++ b/tensorflow/lite/experimental/acceleration/mini_benchmark/libjpeg_decoder.h
@@ -168,10 +168,9 @@ class LibjpegDecoder {
     // Extension of the Run method for non-void JPEG calls when we need to
     // collect the returned value.
     // See Run comments above for details.
-    template <
-        typename Fn, typename... Args,
-        typename ResultType = typename std::result_of_t<Fn>,
-        typename = typename std::enable_if<!std::is_void<ResultType>::value> >
+    template <typename Fn, typename... Args,
+              typename ResultType = typename std::result_of_t<Fn>,
+              typename = typename std::enable_if<!std::is_void_v<ResultType>>>
     TfLiteStatus RunAndSetResult(Fn f, ResultType* result, Args... args) {
       if (setjmp(env_)) return kTfLiteError;
       *result = (handle_->*f)(cinfo_.get(), args...);

From e90814ee05db7d49069c29d252fc5ca69a11fd57 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 20:16:14 -0700
Subject: [PATCH 0484/1000] Automated Code Change

PiperOrigin-RevId: 997642900
---
 tensorflow/core/kernels/batching_util/concat_split_util.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/batching_util/concat_split_util.h b/tensorflow/core/kernels/batching_util/concat_split_util.h
index c83d2b20955..0a1936f9971 100644
--- a/tensorflow/core/kernels/batching_util/concat_split_util.h
+++ b/tensorflow/core/kernels/batching_util/concat_split_util.h
@@ -79,7 +79,7 @@ absl::Status Concat(OpKernelContext* context,
     auto output_flat = output->shaped<T, 2>({1, output->NumElements()});
 #if (defined(GOOGLE_CUDA) && GOOGLE_CUDA) || \
     (defined(TENSORFLOW_USE_ROCM) && TENSORFLOW_USE_ROCM)
-    if (std::is_same<Device, GPUDevice>::value) {
+    if (std::is_same_v<Device, GPUDevice>) {
       ConcatGPU<T>(context, inputs_flat, output, &output_flat);
       return absl::OkStatus();
     }

From 2ac5a9a4d585a5ab7b479a01c6b8602a6cfdc149 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 20:21:40 -0700
Subject: [PATCH 0485/1000] Automated Code Change

PiperOrigin-RevId: 997644748
---
 tensorflow/core/grappler/optimizers/data/graph_utils.h        | 4 ++--
 .../core/grappler/optimizers/data/noop_elimination_test.cc    | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/grappler/optimizers/data/graph_utils.h b/tensorflow/core/grappler/optimizers/data/graph_utils.h
index 78cc5e28182..bbbde7da80c 100644
--- a/tensorflow/core/grappler/optimizers/data/graph_utils.h
+++ b/tensorflow/core/grappler/optimizers/data/graph_utils.h
@@ -63,7 +63,7 @@ template <typename T>
 NodeDef* AddScalarConstNode(T v, MutableGraphView* graph) {
   // is_same is an idiomatic hack for making it compile if not instantiated.
   // Replacing with false will result in a compile-time error.
-  static_assert(!std::is_same<T, T>::value,
+  static_assert(!std::is_same_v<T, T>,
                 "Invalid specialization of this method for type T.");
   return {};
 }
@@ -87,7 +87,7 @@ template <typename T>
 absl::Status GetScalarConstNodeValue(const NodeDef& node, T* value) {
   // is_same is an idiomatic hack for making it compile if not instantiated.
   // Replacing with false will result in a compile-time error.
-  static_assert(!std::is_same<T, T>::value,
+  static_assert(!std::is_same_v<T, T>,
                 "Invalid specialization of this method fo rtype T.");
 }
 
diff --git a/tensorflow/core/grappler/optimizers/data/noop_elimination_test.cc b/tensorflow/core/grappler/optimizers/data/noop_elimination_test.cc
index eadd19ed929..a37cdfa3c18 100644
--- a/tensorflow/core/grappler/optimizers/data/noop_elimination_test.cc
+++ b/tensorflow/core/grappler/optimizers/data/noop_elimination_test.cc
@@ -188,7 +188,7 @@ TEST_P(NoOpMultipleEliminationTest, EliminateMultipleNoOpNode) {
   GrapplerItem item;
   MutableGraphView graph(&item.graph);
 
-  static_assert(std::tuple_size<NodesTypes>::value == 2,
+  static_assert(std::tuple_size_v<NodesTypes> == 2,
                 "Make sure to include everything in the test");
   const std::vector<std::pair<std::string, std::vector<int>>> noop_nodes = {
       std::get<0>(GetParam()), std::get<1>(GetParam())};
@@ -249,7 +249,7 @@ TEST_P(NoOpPlaceholdersTest, NonConstNoOpNode) {
   GrapplerItem item;
   MutableGraphView graph(&item.graph);
 
-  static_assert(std::tuple_size<NodesTypes>::value == 2,
+  static_assert(std::tuple_size_v<NodesTypes> == 2,
                 "Make sure to include everything in the test");
   const std::vector<std::pair<std::string, std::vector<DataType>>> noop_nodes =
       {std::get<0>(GetParam()), std::get<1>(GetParam())};

From c58e1f0bf8a3e0952c8163f131e2096eb3c9aad9 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 20:22:17 -0700
Subject: [PATCH 0486/1000] Automated Code Change

PiperOrigin-RevId: 997644961
---
 tensorflow/core/transforms/constant_folding/pass.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/transforms/constant_folding/pass.cc b/tensorflow/core/transforms/constant_folding/pass.cc
index 37047e14a77..f5a214bc353 100644
--- a/tensorflow/core/transforms/constant_folding/pass.cc
+++ b/tensorflow/core/transforms/constant_folding/pass.cc
@@ -80,7 +80,7 @@ namespace tfg {
 #include "tensorflow/core/transforms/passes.h.inc"
 
 template <typename T>
-static std::enable_if_t<std::is_integral<T>::value, ElementsAttr>
+static std::enable_if_t<std::is_integral_v<T>, ElementsAttr>
 CreateElementsAttrOfTypeValues(Type element_type, ArrayRef<int64_t> shape,
                                ArrayRef<T> values) {
   auto tensor_shape = RankedTensorType::get(shape, element_type);

From 38219c8b97f88db0ebbbb2469e51ff93595d976d Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 20:22:33 -0700
Subject: [PATCH 0487/1000] Automated Code Change

PiperOrigin-RevId: 997645071
---
 tensorflow/tools/proto_text/gen_proto_text_functions.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/tools/proto_text/gen_proto_text_functions.cc b/tensorflow/tools/proto_text/gen_proto_text_functions.cc
index 953a8e8cfa7..e791c297cdf 100644
--- a/tensorflow/tools/proto_text/gen_proto_text_functions.cc
+++ b/tensorflow/tools/proto_text/gen_proto_text_functions.cc
@@ -30,7 +30,7 @@ namespace {
 class CrashOnErrorCollector
     : public tensorflow::protobuf::compiler::MultiFileErrorCollector {
  public:
-  ~CrashOnErrorCollector() override {}
+  ~CrashOnErrorCollector() override = default;
 
   void RecordError(absl::string_view filename, int line, int column,
                    absl::string_view message) override {

From 884a6f1eed5494360fa55dfd3dd044e8eca8f200 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 20:22:33 -0700
Subject: [PATCH 0488/1000] Automated Code Change

PiperOrigin-RevId: 997645072
---
 tensorflow/compiler/jit/deadness_analysis.cc          | 4 ++--
 tensorflow/compiler/jit/encapsulate_util.h            | 2 +-
 tensorflow/compiler/jit/xla_activity_listener.cc      | 2 +-
 tensorflow/compiler/jit/xla_activity_listener_test.cc | 2 +-
 4 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensorflow/compiler/jit/deadness_analysis.cc b/tensorflow/compiler/jit/deadness_analysis.cc
index b7f62da5a92..7b84f00fbb9 100644
--- a/tensorflow/compiler/jit/deadness_analysis.cc
+++ b/tensorflow/compiler/jit/deadness_analysis.cc
@@ -132,7 +132,7 @@ class Predicate {
   virtual absl::Span<Predicate* const> GetOperands() const = 0;
 
   virtual Kind kind() const = 0;
-  virtual ~Predicate() {}
+  virtual ~Predicate() = default;
 
   // Invokes func on p and on all of its operands recursively.  Does not invoke
   // `func` on the same Predicate instance twice.  Aborts the search if `func`
@@ -1566,7 +1566,7 @@ void DeadnessAnalysisImpl::Print() const {
 
 }  // namespace
 
-DeadnessAnalysis::~DeadnessAnalysis() {}
+DeadnessAnalysis::~DeadnessAnalysis() = default;
 
 /*static*/ absl::Status DeadnessAnalysis::Run(
     const Graph& graph, std::unique_ptr<DeadnessAnalysis>* result) {
diff --git a/tensorflow/compiler/jit/encapsulate_util.h b/tensorflow/compiler/jit/encapsulate_util.h
index 81ab31c79dc..8e8a2caa56a 100644
--- a/tensorflow/compiler/jit/encapsulate_util.h
+++ b/tensorflow/compiler/jit/encapsulate_util.h
@@ -94,7 +94,7 @@ struct XlaClusterInfo {
   // TODO(klimek): In c++17 we'll be able to initialize host_compute_core
   // without losing aggregate initialization, which allows us to get rid of
   // the constructor definitions again.
-  XlaClusterInfo() {}
+  XlaClusterInfo() = default;
   XlaClusterInfo(const std::string& cluster_name,
                  const NameAttrList& func_name_attrs, Node* node,
                  const std::map<std::string, int>& host_compute_core)
diff --git a/tensorflow/compiler/jit/xla_activity_listener.cc b/tensorflow/compiler/jit/xla_activity_listener.cc
index 4b0b72b84eb..0debd748ce2 100644
--- a/tensorflow/compiler/jit/xla_activity_listener.cc
+++ b/tensorflow/compiler/jit/xla_activity_listener.cc
@@ -100,6 +100,6 @@ void RegisterXlaActivityListener(
 
 void XlaActivityListener::Flush() {}
 
-XlaActivityListener::~XlaActivityListener() {}
+XlaActivityListener::~XlaActivityListener() = default;
 
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/jit/xla_activity_listener_test.cc b/tensorflow/compiler/jit/xla_activity_listener_test.cc
index c53931eabaa..e17c4d42bef 100644
--- a/tensorflow/compiler/jit/xla_activity_listener_test.cc
+++ b/tensorflow/compiler/jit/xla_activity_listener_test.cc
@@ -48,7 +48,7 @@ class TestListener : public XlaActivityListener {
     return absl::OkStatus();
   }
 
-  ~TestListener() override {}
+  ~TestListener() override = default;
 
   const XlaAutoClusteringActivity& auto_clustering_activity() const {
     return auto_clustering_activity_;

From 75ca790f3ffe1550f49e6e383e8c5390e9688429 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 20:26:41 -0700
Subject: [PATCH 0489/1000] Automated Code Change

PiperOrigin-RevId: 997646449
---
 tensorflow/core/kernels/tensor_array_ops.cc   | 4 ++--
 tensorflow/core/kernels/training_op_helpers.h | 4 ++--
 tensorflow/core/kernels/where_op_gpu.cu.h     | 2 +-
 tensorflow/core/kernels/xent_op.cc            | 2 +-
 4 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensorflow/core/kernels/tensor_array_ops.cc b/tensorflow/core/kernels/tensor_array_ops.cc
index fa69fb20abb..8e3096fd78e 100644
--- a/tensorflow/core/kernels/tensor_array_ops.cc
+++ b/tensorflow/core/kernels/tensor_array_ops.cc
@@ -757,7 +757,7 @@ class TensorArrayPackOrGatherOp : public OpKernel {
     }
 
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-    if (std::is_same<Device, GPUDevice>::value) {
+    if (std::is_same_v<Device, GPUDevice>) {
       ConcatGPU<T>(ctx, input_tensors_flat, output_tensor, &output_flat);
       return;
     }
@@ -984,7 +984,7 @@ class TensorArrayConcatOp : public OpKernel {
       auto output_flat =
           output_tensor->shaped<T, 2>({1, output_shape.num_elements()});
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-      if (std::is_same<Device, GPUDevice>::value) {
+      if (std::is_same_v<Device, GPUDevice>) {
         ConcatGPU<T>(ctx, input_tensors_flat, output_tensor, &output_flat);
         return;
       }
diff --git a/tensorflow/core/kernels/training_op_helpers.h b/tensorflow/core/kernels/training_op_helpers.h
index 1087df8f530..950d42330e3 100644
--- a/tensorflow/core/kernels/training_op_helpers.h
+++ b/tensorflow/core/kernels/training_op_helpers.h
@@ -84,7 +84,7 @@ absl::Status EnsureSparseVariableAccess(OpKernelContext* ctx, Var* var) {
     return absl::OkStatus();
   }
   Tensor tmp;
-  if (std::is_same<T, Variant>::value) {
+  if (std::is_same_v<T, Variant>) {
     tsl::AllocatorAttributes attr;
     attr.set_on_host(true);
     TF_RETURN_IF_ERROR(ctx->allocate_temp(var->tensor()->dtype(),
@@ -246,7 +246,7 @@ absl::Status PrepareToUpdateVariable(OpKernelContext* ctx, Tensor* tensor,
     // Tensor's buffer is in use by some read, so we need to copy before
     // updating.
     Tensor tmp;
-    if (std::is_same<T, Variant>::value) {
+    if (std::is_same_v<T, Variant>) {
       tsl::AllocatorAttributes attr;
       attr.set_on_host(true);
       TF_RETURN_IF_ERROR(
diff --git a/tensorflow/core/kernels/where_op_gpu.cu.h b/tensorflow/core/kernels/where_op_gpu.cu.h
index f9d83d7fe7a..dcfa87cd5b6 100644
--- a/tensorflow/core/kernels/where_op_gpu.cu.h
+++ b/tensorflow/core/kernels/where_op_gpu.cu.h
@@ -286,7 +286,7 @@ struct Where<GPUDevice, NDIM, T, TIndex> {
     typedef std::decay<T> DT;
     CubDeviceSelectFlaggedCounter<
         T, TIndex, decltype(output_iterator) /*OutputIterator*/,
-        std::is_convertible<DT, bool>::value /*IsConvertibleToBool*/>
+        std::is_convertible_v<DT, bool> /*IsConvertibleToBool*/>
         counter;
     auto first_success = counter(/*temp_storage*/ nullptr, temp_storage_bytes,
                                  /*d_flags*/ input.data(),
diff --git a/tensorflow/core/kernels/xent_op.cc b/tensorflow/core/kernels/xent_op.cc
index 498c9d01016..13673ccf4a1 100644
--- a/tensorflow/core/kernels/xent_op.cc
+++ b/tensorflow/core/kernels/xent_op.cc
@@ -66,7 +66,7 @@ class SoftmaxXentWithLogitsOp : public OpKernel {
                                    "2-dimensional, or broadcasted to be "
                                    "2-dimensional"));
 
-    if (std::is_same<Device, GPUDevice>::value) {
+    if (std::is_same_v<Device, GPUDevice>) {
       OP_REQUIRES(context, !OpDeterminismRequired(),
                   absl::UnimplementedError(
                       "The GPU implementation of SoftmaxCrossEntropyWithLogits"

From b2c99c09b6aed3e11d4b92a89e812d44abe04232 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 20:26:58 -0700
Subject: [PATCH 0490/1000] Automated Code Change

PiperOrigin-RevId: 997646533
---
 tensorflow/core/kernels/sparse_cross_op.cc        | 10 +++++-----
 tensorflow/core/kernels/sparse_tensors_map_ops.cc |  2 +-
 tensorflow/core/kernels/spectrogram.h             |  2 +-
 tensorflow/core/kernels/unicode_ops.cc            |  2 +-
 tensorflow/core/kernels/variable_ops.cc           |  2 +-
 tensorflow/core/kernels/while_op_test.cc          |  2 +-
 6 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/tensorflow/core/kernels/sparse_cross_op.cc b/tensorflow/core/kernels/sparse_cross_op.cc
index 3d67021c977..bb6447b542d 100644
--- a/tensorflow/core/kernels/sparse_cross_op.cc
+++ b/tensorflow/core/kernels/sparse_cross_op.cc
@@ -51,7 +51,7 @@ class ColumnInterface {
   virtual InternalType Feature(int64_t batch, int64_t n,
                                bool strong_hash) const = 0;
 
-  virtual ~ColumnInterface() {}
+  virtual ~ColumnInterface() = default;
 };
 
 // A column that is backed by a sparse tensor.
@@ -73,7 +73,7 @@ class SparseTensorColumn : public ColumnInterface<InternalType> {
   InternalType Feature(int64_t batch, int64_t n,
                        bool strong_hash) const override;
 
-  ~SparseTensorColumn() override {}
+  ~SparseTensorColumn() override = default;
 
  private:
   const Tensor& values_;
@@ -103,7 +103,7 @@ class KeyedSparseTensorColumn : public ColumnInterface<InternalType> {
   InternalType Feature(int64_t batch, int64_t n,
                        bool strong_hash) const override;
 
-  ~KeyedSparseTensorColumn() override {}
+  ~KeyedSparseTensorColumn() override = default;
 
  private:
   const Tensor& values_;
@@ -188,7 +188,7 @@ class DenseTensorColumn : public ColumnInterface<InternalType> {
   InternalType Feature(int64_t batch, int64_t n,
                        bool strong_hash) const override;
 
-  ~DenseTensorColumn() override {}
+  ~DenseTensorColumn() override = default;
 
  private:
   const Tensor& tensor_;
@@ -211,7 +211,7 @@ class KeyedDenseTensorColumn : public ColumnInterface<InternalType> {
   InternalType Feature(int64_t batch, int64_t n,
                        bool strong_hash) const override;
 
-  ~KeyedDenseTensorColumn() override {}
+  ~KeyedDenseTensorColumn() override = default;
 
  private:
   const Tensor& tensor_;
diff --git a/tensorflow/core/kernels/sparse_tensors_map_ops.cc b/tensorflow/core/kernels/sparse_tensors_map_ops.cc
index e9e44d3c993..c910a8b9ca2 100644
--- a/tensorflow/core/kernels/sparse_tensors_map_ops.cc
+++ b/tensorflow/core/kernels/sparse_tensors_map_ops.cc
@@ -93,7 +93,7 @@ class SparseTensorsMap : public ResourceBase {
   }
 
  protected:
-  ~SparseTensorsMap() override {}
+  ~SparseTensorsMap() override = default;
 
  private:
   std::string name_;
diff --git a/tensorflow/core/kernels/spectrogram.h b/tensorflow/core/kernels/spectrogram.h
index 4b6b9c8ba68..3185b275fce 100644
--- a/tensorflow/core/kernels/spectrogram.h
+++ b/tensorflow/core/kernels/spectrogram.h
@@ -44,7 +44,7 @@ namespace tensorflow {
 class Spectrogram {
  public:
   Spectrogram() : initialized_(false) {}
-  ~Spectrogram() {}
+  ~Spectrogram() = default;
 
   // Initializes the class with a given window length and step length
   // (both in samples). Internally a Hann window is used as the window
diff --git a/tensorflow/core/kernels/unicode_ops.cc b/tensorflow/core/kernels/unicode_ops.cc
index 79b068422bf..6f8521b7f5b 100644
--- a/tensorflow/core/kernels/unicode_ops.cc
+++ b/tensorflow/core/kernels/unicode_ops.cc
@@ -164,7 +164,7 @@ void IterateUnicodeString(const std::string& str, UConverter* converter,
 // create a specialized fast code path for UTF8.
 class WrappedConverter {
  public:
-  WrappedConverter() {}
+  WrappedConverter() = default;
 
   ~WrappedConverter() {
     if (converter_) {
diff --git a/tensorflow/core/kernels/variable_ops.cc b/tensorflow/core/kernels/variable_ops.cc
index 22cfe373b6f..2c67b7a19de 100644
--- a/tensorflow/core/kernels/variable_ops.cc
+++ b/tensorflow/core/kernels/variable_ops.cc
@@ -62,7 +62,7 @@ class LegacyVar : public ResourceBase {
   mutex mu_;
   Tensor tensor_;
 
-  ~LegacyVar() override {}
+  ~LegacyVar() override = default;
 };
 
 VariableOp::VariableOp(OpKernelConstruction* context) : OpKernel(context) {
diff --git a/tensorflow/core/kernels/while_op_test.cc b/tensorflow/core/kernels/while_op_test.cc
index 8d4443af405..df4ac0c3cb3 100644
--- a/tensorflow/core/kernels/while_op_test.cc
+++ b/tensorflow/core/kernels/while_op_test.cc
@@ -47,7 +47,7 @@ namespace {
 
 class WhileOpTest : public OpsTestBase {
  protected:
-  WhileOpTest() {}
+  WhileOpTest() = default;
   void SetUp() override {
     stream_executor::test_util::PopulateDefaultPlatform(&platform_,
                                                         &platform_fns_);

From 84718586fdc04b8773730fa7a41b61f7d4befa7b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 20:27:40 -0700
Subject: [PATCH 0491/1000] Automated Code Change

PiperOrigin-RevId: 997646743
---
 .../lite/delegates/hexagon/builders/activation_builder.cc       | 2 +-
 .../lite/delegates/hexagon/builders/arg_min_max_builder.cc      | 2 +-
 .../lite/delegates/hexagon/builders/arithmetic_builder.cc       | 2 +-
 tensorflow/lite/delegates/hexagon/builders/cast_builder.cc      | 2 +-
 tensorflow/lite/delegates/hexagon/builders/concat_builder.cc    | 2 +-
 tensorflow/lite/delegates/hexagon/builders/conv_2d_builder.cc   | 2 +-
 tensorflow/lite/delegates/hexagon/builders/hardswish_builder.cc | 2 +-
 .../lite/delegates/hexagon/builders/l2_normalization_builder.cc | 2 +-
 .../lite/delegates/hexagon/builders/mirror_pad_builder.cc       | 2 +-
 tensorflow/lite/delegates/hexagon/builders/op_builder.h         | 2 +-
 10 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/tensorflow/lite/delegates/hexagon/builders/activation_builder.cc b/tensorflow/lite/delegates/hexagon/builders/activation_builder.cc
index ee17b849706..ec84811832b 100644
--- a/tensorflow/lite/delegates/hexagon/builders/activation_builder.cc
+++ b/tensorflow/lite/delegates/hexagon/builders/activation_builder.cc
@@ -63,7 +63,7 @@ TfLiteStatus ActivationOpBuilder::RegisterOutputs(const TfLiteIntArray* outputs,
   return kTfLiteOk;
 }
 
-ActivationOpBuilder::~ActivationOpBuilder() {}
+ActivationOpBuilder::~ActivationOpBuilder() = default;
 
 OpBuilder* CreateActivationBuilder(GraphBuilder* graph_builder, int op_type) {
   return new ActivationOpBuilder(graph_builder, op_type);
diff --git a/tensorflow/lite/delegates/hexagon/builders/arg_min_max_builder.cc b/tensorflow/lite/delegates/hexagon/builders/arg_min_max_builder.cc
index f4b39a02357..b9566718e1c 100644
--- a/tensorflow/lite/delegates/hexagon/builders/arg_min_max_builder.cc
+++ b/tensorflow/lite/delegates/hexagon/builders/arg_min_max_builder.cc
@@ -79,7 +79,7 @@ TfLiteStatus ArgMinMaxOpBuilder::RegisterOutputs(const TfLiteIntArray* outputs,
   return kTfLiteOk;
 }
 
-ArgMinMaxOpBuilder::~ArgMinMaxOpBuilder() {}
+ArgMinMaxOpBuilder::~ArgMinMaxOpBuilder() = default;
 
 OpBuilder* CreateArgMinMaxOpBuilder(GraphBuilder* graph_builder, int op_type) {
   return new ArgMinMaxOpBuilder(graph_builder, op_type);
diff --git a/tensorflow/lite/delegates/hexagon/builders/arithmetic_builder.cc b/tensorflow/lite/delegates/hexagon/builders/arithmetic_builder.cc
index 03af9bebe72..baa20f795af 100644
--- a/tensorflow/lite/delegates/hexagon/builders/arithmetic_builder.cc
+++ b/tensorflow/lite/delegates/hexagon/builders/arithmetic_builder.cc
@@ -121,7 +121,7 @@ TfLiteStatus ArithmeticOpBuilder::RegisterOutputs(const TfLiteIntArray* outputs,
   return kTfLiteOk;
 }
 
-ArithmeticOpBuilder::~ArithmeticOpBuilder() {}
+ArithmeticOpBuilder::~ArithmeticOpBuilder() = default;
 
 OpBuilder* CreateArithmeticBuilder(GraphBuilder* graph_builder, int op_type) {
   return new ArithmeticOpBuilder(graph_builder, op_type);
diff --git a/tensorflow/lite/delegates/hexagon/builders/cast_builder.cc b/tensorflow/lite/delegates/hexagon/builders/cast_builder.cc
index d4b8adb105e..5615fcf3309 100644
--- a/tensorflow/lite/delegates/hexagon/builders/cast_builder.cc
+++ b/tensorflow/lite/delegates/hexagon/builders/cast_builder.cc
@@ -79,7 +79,7 @@ TfLiteStatus CastOpBuilder::RegisterOutputs(const TfLiteIntArray* outputs,
   return kTfLiteOk;
 }
 
-CastOpBuilder::~CastOpBuilder() {}
+CastOpBuilder::~CastOpBuilder() = default;
 
 OpBuilder* CreateCastBuilder(GraphBuilder* graph_builder, int op_type) {
   return new CastOpBuilder(graph_builder, op_type);
diff --git a/tensorflow/lite/delegates/hexagon/builders/concat_builder.cc b/tensorflow/lite/delegates/hexagon/builders/concat_builder.cc
index 47232fab355..371fe987d59 100644
--- a/tensorflow/lite/delegates/hexagon/builders/concat_builder.cc
+++ b/tensorflow/lite/delegates/hexagon/builders/concat_builder.cc
@@ -144,7 +144,7 @@ TfLiteStatus ConcatOpBuilder::RegisterOutputs(const TfLiteIntArray* outputs,
   return kTfLiteOk;
 }
 
-ConcatOpBuilder::~ConcatOpBuilder() {}
+ConcatOpBuilder::~ConcatOpBuilder() = default;
 
 OpBuilder* CreateConcatBuilder(GraphBuilder* graph_builder, int op_type) {
   return new ConcatOpBuilder(graph_builder, op_type);
diff --git a/tensorflow/lite/delegates/hexagon/builders/conv_2d_builder.cc b/tensorflow/lite/delegates/hexagon/builders/conv_2d_builder.cc
index 0c17d2f0baa..2737eafbaf0 100644
--- a/tensorflow/lite/delegates/hexagon/builders/conv_2d_builder.cc
+++ b/tensorflow/lite/delegates/hexagon/builders/conv_2d_builder.cc
@@ -531,7 +531,7 @@ TfLiteStatus Conv2dOpBuilder::RegisterOutputs(const TfLiteIntArray* outputs,
   return kTfLiteOk;
 }
 
-Conv2dOpBuilder::~Conv2dOpBuilder() {}
+Conv2dOpBuilder::~Conv2dOpBuilder() = default;
 
 OpBuilder* CreateConv2DBuilder(GraphBuilder* graph_builder, int op_type) {
   return new Conv2dOpBuilder(graph_builder, op_type);
diff --git a/tensorflow/lite/delegates/hexagon/builders/hardswish_builder.cc b/tensorflow/lite/delegates/hexagon/builders/hardswish_builder.cc
index dc42c8f51ed..ba58741f130 100644
--- a/tensorflow/lite/delegates/hexagon/builders/hardswish_builder.cc
+++ b/tensorflow/lite/delegates/hexagon/builders/hardswish_builder.cc
@@ -57,7 +57,7 @@ TfLiteStatus HardSwishOpBuilder::RegisterOutputs(const TfLiteIntArray* outputs,
   return kTfLiteOk;
 }
 
-HardSwishOpBuilder::~HardSwishOpBuilder() {}
+HardSwishOpBuilder::~HardSwishOpBuilder() = default;
 
 OpBuilder* CreateHardSwishBuilder(GraphBuilder* graph_builder, int op_type) {
   return new HardSwishOpBuilder(graph_builder, op_type);
diff --git a/tensorflow/lite/delegates/hexagon/builders/l2_normalization_builder.cc b/tensorflow/lite/delegates/hexagon/builders/l2_normalization_builder.cc
index e4bb336b6e3..2e4aabaae65 100644
--- a/tensorflow/lite/delegates/hexagon/builders/l2_normalization_builder.cc
+++ b/tensorflow/lite/delegates/hexagon/builders/l2_normalization_builder.cc
@@ -55,7 +55,7 @@ TfLiteStatus L2NormalizationOpBuilder::RegisterOutputs(
   return kTfLiteOk;
 }
 
-L2NormalizationOpBuilder::~L2NormalizationOpBuilder() {}
+L2NormalizationOpBuilder::~L2NormalizationOpBuilder() = default;
 
 OpBuilder* CreateL2NormalizationBuilder(GraphBuilder* graph_builder,
                                         int op_type) {
diff --git a/tensorflow/lite/delegates/hexagon/builders/mirror_pad_builder.cc b/tensorflow/lite/delegates/hexagon/builders/mirror_pad_builder.cc
index bcce11acd02..2397f263038 100644
--- a/tensorflow/lite/delegates/hexagon/builders/mirror_pad_builder.cc
+++ b/tensorflow/lite/delegates/hexagon/builders/mirror_pad_builder.cc
@@ -89,7 +89,7 @@ TfLiteStatus MirrorPadOpBuilder::RegisterOutputs(const TfLiteIntArray* outputs,
   return kTfLiteOk;
 }
 
-MirrorPadOpBuilder::~MirrorPadOpBuilder() {}
+MirrorPadOpBuilder::~MirrorPadOpBuilder() = default;
 
 OpBuilder* CreateMirrorPadBuilder(GraphBuilder* graph_builder, int op_type) {
   return new MirrorPadOpBuilder(graph_builder, op_type);
diff --git a/tensorflow/lite/delegates/hexagon/builders/op_builder.h b/tensorflow/lite/delegates/hexagon/builders/op_builder.h
index 8f835a69c65..b06e1fc2b19 100644
--- a/tensorflow/lite/delegates/hexagon/builders/op_builder.h
+++ b/tensorflow/lite/delegates/hexagon/builders/op_builder.h
@@ -74,7 +74,7 @@ class OpBuilder {
   // node output list.
   using TensorID = std::pair<int, int>;
 
-  virtual ~OpBuilder() {}
+  virtual ~OpBuilder() = default;
 
   // Sets the op type in the hexagon graph.
   void SetOpType(int op_type) { op_node_.op_type = op_type; }

From 68f24f50825bc3507660cda39a44327f5093f0dc Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 20:28:50 -0700
Subject: [PATCH 0492/1000] Automated Code Change

PiperOrigin-RevId: 997647138
---
 .../lite/delegates/hexagon/builders/tests/pack_test.cc      | 6 +++---
 .../delegates/hexagon/builders/tests/transpose_conv_test.cc | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensorflow/lite/delegates/hexagon/builders/tests/pack_test.cc b/tensorflow/lite/delegates/hexagon/builders/tests/pack_test.cc
index e9735e414c9..2a30b1ff77c 100644
--- a/tensorflow/lite/delegates/hexagon/builders/tests/pack_test.cc
+++ b/tensorflow/lite/delegates/hexagon/builders/tests/pack_test.cc
@@ -60,10 +60,10 @@ template <typename InputType>
 struct PackOpTest : public ::testing::Test {
   using TypeToTest = InputType;
   TensorType TENSOR_TYPE =
-      (std::is_same<InputType, int16_t>::value
+      (std::is_same_v<InputType, int16_t>
            ? TensorType_INT16
-           : (std::is_same<InputType, uint8_t>::value ? TensorType_UINT8
-                                                      : TensorType_INT8));
+           : (std::is_same_v<InputType, uint8_t> ? TensorType_UINT8
+                                                 : TensorType_INT8));
 };
 
 using TestTypes = testing::Types<int8_t, uint8_t>;
diff --git a/tensorflow/lite/delegates/hexagon/builders/tests/transpose_conv_test.cc b/tensorflow/lite/delegates/hexagon/builders/tests/transpose_conv_test.cc
index b894691120f..63eb636fce3 100644
--- a/tensorflow/lite/delegates/hexagon/builders/tests/transpose_conv_test.cc
+++ b/tensorflow/lite/delegates/hexagon/builders/tests/transpose_conv_test.cc
@@ -89,9 +89,9 @@ class QuantizedTransposeConvOpModel : public SingleOpModelWithHexagon {
   }
 
   void SetBias(std::initializer_list<float> bias) {
-    if (std::is_same<InputType, uint8_t>::value) {
+    if (std::is_same_v<InputType, uint8_t>) {
       QuantizeAndPopulate<int32_t>(bias_, bias);
-    } else if (std::is_same<InputType, int8_t>::value) {
+    } else if (std::is_same_v<InputType, int8_t>) {
       PerChannelQuantizeBias(bias_, bias);
     }
     // Set allocation type to MmapRo to simulate a 'constant' tensor.

From 5eecaa47e583149e3d033441a211033de35e7d83 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 20:49:26 -0700
Subject: [PATCH 0493/1000] Automated Code Change

PiperOrigin-RevId: 997653768
---
 tensorflow/core/kernels/image/resize_bilinear_op_test.cc     | 2 +-
 tensorflow/core/kernels/image/scale_and_translate_op_test.cc | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/kernels/image/resize_bilinear_op_test.cc b/tensorflow/core/kernels/image/resize_bilinear_op_test.cc
index b664bbb02de..0661be87924 100644
--- a/tensorflow/core/kernels/image/resize_bilinear_op_test.cc
+++ b/tensorflow/core/kernels/image/resize_bilinear_op_test.cc
@@ -169,7 +169,7 @@ class ResizeBilinearOpTestBase
 
 class ResizeBilinearOpTest : public ResizeBilinearOpTestBase {
  public:
-  ResizeBilinearOpTest() {}
+  ResizeBilinearOpTest() = default;
 };
 
 class ResizeBilinearHalfPixelCentersOpTest : public ResizeBilinearOpTestBase {
diff --git a/tensorflow/core/kernels/image/scale_and_translate_op_test.cc b/tensorflow/core/kernels/image/scale_and_translate_op_test.cc
index 0def4456696..40c6e15c992 100644
--- a/tensorflow/core/kernels/image/scale_and_translate_op_test.cc
+++ b/tensorflow/core/kernels/image/scale_and_translate_op_test.cc
@@ -36,7 +36,7 @@ using Eigen::Vector2f;
 
 class DynamicKernel {
  public:
-  virtual ~DynamicKernel() {}
+  virtual ~DynamicKernel() = default;
   virtual float Value(const float x) const = 0;
   virtual float Radius() const = 0;
 };

From d95dbb2188149aa94486b969ed13b0a8eacc5faf Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:01:50 -0700
Subject: [PATCH 0494/1000] Automated Code Change

PiperOrigin-RevId: 997656935
---
 tensorflow/core/kernels/host_constant_op.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/host_constant_op.h b/tensorflow/core/kernels/host_constant_op.h
index 9ba151ba71b..ac3a1b9e6a9 100644
--- a/tensorflow/core/kernels/host_constant_op.h
+++ b/tensorflow/core/kernels/host_constant_op.h
@@ -31,7 +31,7 @@ class _HostConstantOp : public OpKernel {
   void Compute(OpKernelContext* ctx) override;
   bool IsExpensive() override { return false; }
   const Tensor* const_tensor() const override { return &tensor_; };
-  ~_HostConstantOp() override {}
+  ~_HostConstantOp() override = default;
 
  private:
   Tensor tensor_;

From c1a07a258ef8debdab14a1ad08aa61c7eac77fae Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:01:50 -0700
Subject: [PATCH 0495/1000] Automated Code Change

PiperOrigin-RevId: 997656939
---
 tensorflow/core/kernels/image/adjust_contrast_op.cc        | 2 +-
 tensorflow/core/kernels/image/crop_and_resize_op.cc        | 4 ++--
 tensorflow/core/kernels/image/resize_bilinear_op_gpu.cu.cc | 3 +--
 3 files changed, 4 insertions(+), 5 deletions(-)

diff --git a/tensorflow/core/kernels/image/adjust_contrast_op.cc b/tensorflow/core/kernels/image/adjust_contrast_op.cc
index a61bd822829..b9740a473e8 100644
--- a/tensorflow/core/kernels/image/adjust_contrast_op.cc
+++ b/tensorflow/core/kernels/image/adjust_contrast_op.cc
@@ -69,7 +69,7 @@ class AdjustContrastOp : public OpKernel {
         absl::InvalidArgumentError(absl::StrCat(
             "max_value must be scalar: ", max_value.shape().DebugString())));
 
-    if (std::is_same<Device, GPUDevice>::value) {
+    if (std::is_same_v<Device, GPUDevice>) {
       OP_REQUIRES(
           context, !OpDeterminismRequired(),
           absl::UnimplementedError(
diff --git a/tensorflow/core/kernels/image/crop_and_resize_op.cc b/tensorflow/core/kernels/image/crop_and_resize_op.cc
index 4f9063471dd..11d5e40f413 100644
--- a/tensorflow/core/kernels/image/crop_and_resize_op.cc
+++ b/tensorflow/core/kernels/image/crop_and_resize_op.cc
@@ -420,7 +420,7 @@ class CropAndResizeGradImageOp : public AsyncOpKernel {
         absl::InvalidArgumentError("image_size and grads are incompatible"),
         done);
 
-    if (std::is_same<Device, GPUDevice>::value) {
+    if (std::is_same_v<Device, GPUDevice>) {
       OP_REQUIRES_ASYNC(
           context, !OpDeterminismRequired(),
           absl::UnimplementedError(
@@ -642,7 +642,7 @@ class CropAndResizeGradBoxesOp : public AsyncOpKernel {
         absl::InvalidArgumentError("boxes and grads have incompatible shape"),
         done);
 
-    if (std::is_same<Device, GPUDevice>::value) {
+    if (std::is_same_v<Device, GPUDevice>) {
       OP_REQUIRES_ASYNC(
           context, !OpDeterminismRequired(),
           absl::UnimplementedError(
diff --git a/tensorflow/core/kernels/image/resize_bilinear_op_gpu.cu.cc b/tensorflow/core/kernels/image/resize_bilinear_op_gpu.cu.cc
index 59fe629188b..c3c80a2e1c6 100644
--- a/tensorflow/core/kernels/image/resize_bilinear_op_gpu.cu.cc
+++ b/tensorflow/core/kernels/image/resize_bilinear_op_gpu.cu.cc
@@ -430,8 +430,7 @@ struct ResizeBilinear<GPUDevice, T> {
       // access.
       constexpr int max_num_threads_per_pixel = 8;
       constexpr int channels_per_thread = 16 / sizeof(T);
-      if (channels % channels_per_thread == 0 &&
-          std::is_same<float, T>::value) {
+      if (channels % channels_per_thread == 0 && std::is_same_v<float, T>) {
         const int64_t num_threads_per_pixel = std::min<int64_t>(
             max_num_threads_per_pixel, channels / channels_per_thread);
         config = GetGpuLaunchConfig64(static_cast<int64_t>(out_height) *

From abbfc06a04a664f9beb9d1ac8e2bb63f75e9eb72 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:02:43 -0700
Subject: [PATCH 0496/1000] Automated Code Change

PiperOrigin-RevId: 997657202
---
 tensorflow/compiler/tf2xla/kernels/const_op.cc     | 5 ++---
 tensorflow/compiler/tf2xla/kernels/sequence_ops.cc | 2 +-
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/tensorflow/compiler/tf2xla/kernels/const_op.cc b/tensorflow/compiler/tf2xla/kernels/const_op.cc
index 84b24195807..239793fc7ab 100644
--- a/tensorflow/compiler/tf2xla/kernels/const_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/const_op.cc
@@ -36,9 +36,8 @@ DstT CastTo(SrcT src) {
 }
 
 template <typename DstT,
-          typename std::enable_if<std::is_same<DstT, Eigen::half>::value ||
-                                  std::is_same<DstT, bfloat16>::value>::type* =
-              nullptr>
+          std::enable_if_t<std::is_same_v<DstT, Eigen::half> ||
+                           std::is_same_v<DstT, bfloat16>>* = nullptr>
 DstT CastTo(int32_t src) {
   return absl::bit_cast<DstT>(static_cast<uint16_t>(src));
 }
diff --git a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc
index 55336a481ae..cad1be9fc97 100644
--- a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc
@@ -65,7 +65,7 @@ absl::StatusOr<xla::XlaOp> CreateRangeTensor(
     }
   }
   int64_t size;
-  if constexpr (std::is_integral<T>::value) {
+  if constexpr (std::is_integral_v<T>) {
     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);

From d61bc091d6d1ee3c7a747965b78f7ac62147fba5 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:05:11 -0700
Subject: [PATCH 0497/1000] Automated Code Change

PiperOrigin-RevId: 997657928
---
 tensorflow/core/grappler/optimizers/meta_optimizer.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/grappler/optimizers/meta_optimizer.cc b/tensorflow/core/grappler/optimizers/meta_optimizer.cc
index fed3de51070..2b13704aab7 100644
--- a/tensorflow/core/grappler/optimizers/meta_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/meta_optimizer.cc
@@ -1415,7 +1415,7 @@ absl::Status OptimizeGraph(
   tensorflow::grappler::GrapplerItem item;
   item.id = grappler_item_id;
   item.optimization_options() = optimization_options;
-  if (cpu_device && std::is_same<decltype(cpu_device), LocalDevice>::value &&
+  if (cpu_device && std::is_same_v<decltype(cpu_device), LocalDevice> &&
       cpu_device->tensorflow_cpu_worker_threads() != nullptr) {
     // Forward to the optimisation pass number of intra threads that are used to
     // parallelise operations.

From bca47e505dfeb6a6ec30039ea3a5d47afcf0873f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:06:18 -0700
Subject: [PATCH 0498/1000] Automated Code Change

PiperOrigin-RevId: 997658240
---
 tensorflow/core/kernels/scan_ops_gpu.h        | 32 +++++++++----------
 tensorflow/core/kernels/scatter_functor.h     |  2 +-
 tensorflow/core/kernels/scatter_nd_op.cc      | 13 ++++----
 tensorflow/core/kernels/scatter_op.cc         |  2 +-
 .../kernels/segment_reduction_ops_gpu.cu.h    |  6 ++--
 .../core/kernels/segment_reduction_ops_impl.h | 24 +++++---------
 tensorflow/core/kernels/sequence_ops.cc       |  2 +-
 tensorflow/core/kernels/slice_op.cc           |  2 +-
 .../core/kernels/spacetobatch_functor.cc      |  4 +--
 .../core/kernels/spacetobatch_functor.h       |  4 +--
 10 files changed, 41 insertions(+), 50 deletions(-)

diff --git a/tensorflow/core/kernels/scan_ops_gpu.h b/tensorflow/core/kernels/scan_ops_gpu.h
index 15b4e5e1ecf..2f084a89bb0 100644
--- a/tensorflow/core/kernels/scan_ops_gpu.h
+++ b/tensorflow/core/kernels/scan_ops_gpu.h
@@ -115,21 +115,21 @@ struct Prod {
 template <typename T, typename Op>
 struct IsSum {
   constexpr static bool value =
-      (std::is_same<Op, Sum<T>>::value ||
-       std::is_same<Op, Eigen::internal::SumReducer<T>>::value);
+      (std::is_same_v<Op, Sum<T>> ||
+       std::is_same_v<Op, Eigen::internal::SumReducer<T>>);
 };
 
 template <typename T, typename Op>
 struct IsProd {
   constexpr static bool value =
-      (std::is_same<Op, Prod<T>>::value ||
-       std::is_same<Op, Eigen::internal::ProdReducer<T>>::value);
+      (std::is_same_v<Op, Prod<T>> ||
+       std::is_same_v<Op, Eigen::internal::ProdReducer<T>>);
 };
 
 template <typename T, typename Op>
 struct IsLogSumExp {
-  constexpr static bool value = (std::is_same<Op, LogSumExp<T>>::value ||
-                                 std::is_same<Op, LogSumExpReducer<T>>::value);
+  constexpr static bool value = (std::is_same_v<Op, LogSumExp<T>> ||
+                                 std::is_same_v<Op, LogSumExpReducer<T>>);
 };
 
 template <typename T, typename Op>
@@ -139,20 +139,20 @@ struct IdentityValue {
                 "IdentityValue not yet defined for this type.");
 
   template <typename U = T, typename OpCopy = Op>
-  __host__ __device__ U operator()(
-      typename std::enable_if<IsSum<U, OpCopy>::value, U>::type t = U(0)) {
-    return t;
-  }
-
-  template <typename U = T, typename OpCopy = Op>
-  __host__ __device__ U operator()(
-      typename std::enable_if<IsProd<U, OpCopy>::value, U>::type t = U(1)) {
+  __host__ __device__ U
+  operator()(std::enable_if_t<IsSum<U, OpCopy>::value, U> t = U(0)) {
     return t;
   }
 
   template <typename U = T, typename OpCopy = Op>
   __host__ __device__ U
-  operator()(typename std::enable_if<IsLogSumExp<U, OpCopy>::value, U>::type t =
+  operator()(std::enable_if_t<IsProd<U, OpCopy>::value, U> t = U(1)) {
+    return t;
+  }
+
+  template <typename U = T, typename OpCopy = Op>
+  __host__ __device__ U
+  operator()(std::enable_if_t<IsLogSumExp<U, OpCopy>::value, U> t =
                  U(Eigen::NumTraits<U>::lowest())) {
     return t;
   }
@@ -251,7 +251,7 @@ void LaunchScan(const GPUDevice& d, typename TTypes<T, 3>::ConstTensor in,
 
   // There seems to be a bug when the type is not float and block_size 1024.
   // Launch on the smallest power of 2 block size that we can.
-  if (ideal_block_size >= 1024 && std::is_same<T, float>::value) {
+  if (ideal_block_size >= 1024 && std::is_same_v<T, float>) {
     const int block_size = 1024;
     TF_CHECK_OK(
         GpuLaunchKernel(scan_kernel<T, Op, block_size, items_per_thread>,
diff --git a/tensorflow/core/kernels/scatter_functor.h b/tensorflow/core/kernels/scatter_functor.h
index dcfae9b7762..fd8e92243de 100644
--- a/tensorflow/core/kernels/scatter_functor.h
+++ b/tensorflow/core/kernels/scatter_functor.h
@@ -278,7 +278,7 @@ struct ScatterFunctorBase<CPUDevice, T, Index, scatter_op::UpdateOp::ASSIGN> {
     // indices and params sizes were validated in DoCompute().
     const Index N = static_cast<Index>(indices.size());
     const Index limit = static_cast<Index>(params.dimension(0));
-    if (!std::is_same<T, tstring>::value) {
+    if (!std::is_same_v<T, tstring>) {
       for (Index i = 0; i < N; i++) {
         // Grab the index and check its validity.  Do this carefully,
         // to avoid checking the value and grabbing it again from
diff --git a/tensorflow/core/kernels/scatter_nd_op.cc b/tensorflow/core/kernels/scatter_nd_op.cc
index 49d8236dcd0..bc611f9d649 100644
--- a/tensorflow/core/kernels/scatter_nd_op.cc
+++ b/tensorflow/core/kernels/scatter_nd_op.cc
@@ -93,7 +93,7 @@ class ScatterOpBase : public OpKernel {
         BadIndicesPolicyFromString(bad_indices_policy_str);
     OP_REQUIRES_OK(c, bad_indices_policy.status());
     bad_indices_policy_ = *bad_indices_policy;
-    if constexpr (std::is_same<Device, GPUDevice>::value) {
+    if constexpr (std::is_same_v<Device, GPUDevice>) {
       OP_REQUIRES(
           c, bad_indices_policy_ != BadIndicesPolicy::kError,
           absl::InvalidArgumentError(
@@ -260,7 +260,7 @@ class TensorScatterOp : public ScatterOpBase<Device> {
 
     AllocatorAttributes alloc_attr;
     MemoryType memory_type = DEVICE_MEMORY;
-    if (std::is_same<Device, CPUDevice>::value) {
+    if (std::is_same_v<Device, CPUDevice>) {
       alloc_attr.set_on_host(true);
       memory_type = HOST_MEMORY;
     } else {
@@ -971,7 +971,7 @@ absl::Status DoScatterNdImpl(OpKernelContext* c, const Tensor& indices,
 
   if (allocate) {
     AllocatorAttributes alloc_attr;
-    if (std::is_same<Device, CPUDevice>::value) {
+    if (std::is_same_v<Device, CPUDevice>) {
       alloc_attr.set_on_host(true);
     }
     TF_RETURN_IF_ERROR(
@@ -1030,7 +1030,7 @@ absl::Status DoScatterNdImpl(OpKernelContext* c, const Tensor& indices,
     }
   }
   const bool check_bad_indices =
-      ((std::is_same<Device, CPUDevice>::value &&
+      ((std::is_same_v<Device, CPUDevice> &&
         bad_indices_policy == BadIndicesPolicy::kDefault) ||
        bad_indices_policy == BadIndicesPolicy::kError);
   if (check_bad_indices && bad_i >= 0) {
@@ -1145,7 +1145,7 @@ absl::Status DoScatterNd(OpKernelContext* c, const Tensor& indices,
                          Tensor* out, bool allocate,
                          BadIndicesPolicy bad_indices_policy) {
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-  if (std::is_same<Device, GPUDevice>::value &&
+  if (std::is_same_v<Device, GPUDevice> &&
       tensorflow::OpDeterminismRequired() && !DisableScatterOpDeterminism()) {
     return DoScatterNdOnCpu<T, Index, Op>(c, indices, updates, shape, out,
                                           allocate);
@@ -1154,8 +1154,7 @@ absl::Status DoScatterNd(OpKernelContext* c, const Tensor& indices,
 
   // Run on the CPU for integer types, since the GPU implementation uses
   // atomics, which are not supported for all integer types.
-  if constexpr (std::is_same<Device, GPUDevice>::value &&
-                std::is_integral<T>::value) {
+  if constexpr (std::is_same_v<Device, GPUDevice> && std::is_integral_v<T>) {
     return DoScatterNdOnCpu<T, Index, Op>(c, indices, updates, shape, out,
                                           allocate, bad_indices_policy);
   } else {
diff --git a/tensorflow/core/kernels/scatter_op.cc b/tensorflow/core/kernels/scatter_op.cc
index 8cb939b3fc9..825db4f49bf 100644
--- a/tensorflow/core/kernels/scatter_op.cc
+++ b/tensorflow/core/kernels/scatter_op.cc
@@ -80,7 +80,7 @@ class ScatterUpdateOp : public OpKernel {
   //   in the graph?
   explicit ScatterUpdateOp(OpKernelConstruction* c) : OpKernel(c) {
     OP_REQUIRES_OK(c, c->GetAttr("use_locking", &use_exclusive_lock_));
-    if (std::is_same<Device, GPUDevice>::value) {
+    if (std::is_same_v<Device, GPUDevice>) {
       OP_REQUIRES(
           c, !OpDeterminismRequired(),
           absl::UnimplementedError(
diff --git a/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h b/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h
index a906d8b0a26..bbb1e081934 100644
--- a/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h
+++ b/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h
@@ -574,7 +574,7 @@ absl::Status SegmentReduceGPUImplNoInnerDim(
   // when Tvec != Treducevec; the result is only used if Tvec == Treducevec.
   Treducevec* output_raw_ptr = reinterpret_cast<Treducevec*>(output_vec);
   Tensor output_raw;
-  bool need_temp_output = !std::is_same<Tvec, Treducevec>::value;
+  bool need_temp_output = !std::is_same_v<Tvec, Treducevec>;
   if (need_temp_output) {
     // Note: We must allocate and reinterpret as bytes because Treducevec may
     // be a vector type and they are not supported as Tensor dtypes.
@@ -591,7 +591,7 @@ absl::Status SegmentReduceGPUImplNoInnerDim(
   TF_RETURN_IF_ERROR(GpuSegmentedReduce(ctx, nsegments, reduce_op,
                                         Treducevec(initial_value), input_iter,
                                         segment_offsets, output_raw_ptr));
-  bool need_epilogue = !std::is_same<Tvec, Treducevec>::value ||
+  bool need_epilogue = !std::is_same_v<Tvec, Treducevec> ||
                        initial_value != empty_segment_value || is_mean ||
                        is_sqrtn;
   if (need_epilogue) {
@@ -1321,7 +1321,7 @@ struct SparseSegmentGradV2Functor<GPUDevice, T, Tindices, Tsegmentids> {
     using TindicesCompact = int32_t;
     Tensor tmp_indices_internal;
     const TindicesCompact* indices_internal_ptr;
-    if constexpr (std::is_same<Tindices, TindicesCompact>::value) {
+    if constexpr (std::is_same_v<Tindices, TindicesCompact>) {
       indices_internal_ptr = indices_vec.data();
     } else {
       OP_REQUIRES_OK_ASYNC(
diff --git a/tensorflow/core/kernels/segment_reduction_ops_impl.h b/tensorflow/core/kernels/segment_reduction_ops_impl.h
index 6039f589c63..727e04a88bd 100644
--- a/tensorflow/core/kernels/segment_reduction_ops_impl.h
+++ b/tensorflow/core/kernels/segment_reduction_ops_impl.h
@@ -737,15 +737,11 @@ class SparseSegmentReductionOpBase : public OpKernel {
   const DataType dtidx_;
 
   template <typename Tin>
-  using EnableIfBfloat16OrHalf =
-      typename std::enable_if<std::is_same<Tin, bfloat16>::value ||
-                                  std::is_same<Tin, Eigen::half>::value,
-                              int>::type;
+  using EnableIfBfloat16OrHalf = std::enable_if_t<
+      std::is_same_v<Tin, bfloat16> || std::is_same_v<Tin, Eigen::half>, int>;
   template <typename Tin>
-  using EnableIfNotBfloat16OrHalf =
-      typename std::enable_if<!std::is_same<Tin, bfloat16>::value &&
-                                  !std::is_same<Tin, Eigen::half>::value,
-                              int>::type;
+  using EnableIfNotBfloat16OrHalf = std::enable_if_t<
+      !std::is_same_v<Tin, bfloat16> && !std::is_same_v<Tin, Eigen::half>, int>;
 
   template <typename Tin, typename Tindex, EnableIfNotBfloat16OrHalf<Tin> = 0>
   EIGEN_ALWAYS_INLINE auto fetch_val(
@@ -1171,15 +1167,11 @@ struct SparseSegmentGradFunctor<CPUDevice, T, Index, SegmentId> {
 
  private:
   template <typename Tin>
-  using EnableIfBfloat16OrHalf =
-      typename std::enable_if<std::is_same<Tin, bfloat16>::value ||
-                                  std::is_same<Tin, Eigen::half>::value,
-                              int>::type;
+  using EnableIfBfloat16OrHalf = std::enable_if_t<
+      std::is_same_v<Tin, bfloat16> || std::is_same_v<Tin, Eigen::half>, int>;
   template <typename Tin>
-  using EnableIfNotBfloat16OrHalf =
-      typename std::enable_if<!std::is_same<Tin, bfloat16>::value &&
-                                  !std::is_same<Tin, Eigen::half>::value,
-                              int>::type;
+  using EnableIfNotBfloat16OrHalf = std::enable_if_t<
+      !std::is_same_v<Tin, bfloat16> && !std::is_same_v<Tin, Eigen::half>, int>;
 
   template <typename Tin, EnableIfNotBfloat16OrHalf<Tin> = 0>
   void Accumulate(
diff --git a/tensorflow/core/kernels/sequence_ops.cc b/tensorflow/core/kernels/sequence_ops.cc
index 99bc253349d..38e115058c4 100644
--- a/tensorflow/core/kernels/sequence_ops.cc
+++ b/tensorflow/core/kernels/sequence_ops.cc
@@ -95,7 +95,7 @@ class RangeOp : public OpKernel {
               "Requires start >= limit when delta < 0: ", start, "/", limit));
     }
     int64_t size;
-    if constexpr (std::is_integral<T>::value) {
+    if constexpr (std::is_integral_v<T>) {
       uint64_t range;
       if ((limit > 0 && start < 0) || (limit < 0 && start > 0)) {
         range = static_cast<uint64_t>(Eigen::numext::abs(limit)) +
diff --git a/tensorflow/core/kernels/slice_op.cc b/tensorflow/core/kernels/slice_op.cc
index 5cef063c379..acf36775fcc 100644
--- a/tensorflow/core/kernels/slice_op.cc
+++ b/tensorflow/core/kernels/slice_op.cc
@@ -169,7 +169,7 @@ class SliceOp : public OpKernel {
     const int input_dims = input.dims();
 
     if (result->NumElements() > 0) {
-      if (std::is_same<Device, CPUDevice>::value && input_dims == 2 &&
+      if (std::is_same_v<Device, CPUDevice> && input_dims == 2 &&
           DataTypeCanUseMemcpy(DataTypeToEnum<T>::v())) {
         auto input_t = input.tensor<T, 2>();
         auto output_t = result->tensor<T, 2>();
diff --git a/tensorflow/core/kernels/spacetobatch_functor.cc b/tensorflow/core/kernels/spacetobatch_functor.cc
index 6e5319b94bb..04415d23196 100644
--- a/tensorflow/core/kernels/spacetobatch_functor.cc
+++ b/tensorflow/core/kernels/spacetobatch_functor.cc
@@ -88,8 +88,8 @@ struct SpaceToBatchHelper<0, B2S> {
 
 template <typename T, int NUM_BLOCK_DIMS, bool B2S>
 struct SpaceToBatchFunctor<CPUDevice, T, NUM_BLOCK_DIMS, B2S> {
-  using SpaceT = typename std::conditional<B2S, T, const T>::type;
-  using BatchT = typename std::conditional<B2S, const T, T>::type;
+  using SpaceT = std::conditional_t<B2S, T, const T>;
+  using BatchT = std::conditional_t<B2S, const T, T>;
   absl::Status operator()(
       const CPUDevice& d,
       typename TTypes<SpaceT, NUM_BLOCK_DIMS + 2>::Tensor space_tensor,
diff --git a/tensorflow/core/kernels/spacetobatch_functor.h b/tensorflow/core/kernels/spacetobatch_functor.h
index 40037ba9ec0..068af96416e 100644
--- a/tensorflow/core/kernels/spacetobatch_functor.h
+++ b/tensorflow/core/kernels/spacetobatch_functor.h
@@ -84,8 +84,8 @@ namespace functor {
 // true, then this performs the inverse batch-to-space conversion.
 template <typename Device, typename T, int NUM_BLOCK_DIMS, bool B2S = false>
 struct SpaceToBatchFunctor {
-  using InputT = typename std::conditional<B2S, T, const T>::type;
-  using OutputT = typename std::conditional<B2S, const T, T>::type;
+  using InputT = std::conditional_t<B2S, T, const T>;
+  using OutputT = std::conditional_t<B2S, const T, T>;
   // Implements the space to batch conversion.
   //
   // space_tensor: input tensor of space-to-batch operation.  If B2S = false,

From e133c0dc6461e8afa6bb243190e27fc64b4e10f5 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:06:46 -0700
Subject: [PATCH 0499/1000] Automated Code Change

PiperOrigin-RevId: 997658354
---
 tensorflow/core/lib/gtl/cleanup.h       | 4 ++--
 tensorflow/core/lib/gtl/cleanup_test.cc | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/lib/gtl/cleanup.h b/tensorflow/core/lib/gtl/cleanup.h
index 3e54f82860d..e4cea8fe2c0 100644
--- a/tensorflow/core/lib/gtl/cleanup.h
+++ b/tensorflow/core/lib/gtl/cleanup.h
@@ -95,14 +95,14 @@ class Cleanup {
   bool is_released() const { return released_; }
 
  private:
-  static_assert(!std::is_reference<F>::value, "F must not be a reference");
+  static_assert(!std::is_reference_v<F>, "F must not be a reference");
 
   bool released_ = false;
   F f_;
 };
 
 template <int&... ExplicitParameterBarrier, typename F,
-          typename DecayF = typename std::decay<F>::type>
+          typename DecayF = std::decay_t<F>>
 TF_MUST_USE_RESULT Cleanup<DecayF> MakeCleanup(F&& f) {
   return Cleanup<DecayF>(std::forward<F>(f));
 }
diff --git a/tensorflow/core/lib/gtl/cleanup_test.cc b/tensorflow/core/lib/gtl/cleanup_test.cc
index d26df93a99e..220126eb863 100644
--- a/tensorflow/core/lib/gtl/cleanup_test.cc
+++ b/tensorflow/core/lib/gtl/cleanup_test.cc
@@ -28,7 +28,7 @@ using AnyCleanup = gtl::Cleanup<std::function<void()>>;
 
 template <typename T1, typename T2>
 void AssertTypeEq() {
-  static_assert(std::is_same<T1, T2>::value, "unexpected type");
+  static_assert(std::is_same_v<T1, T2>, "unexpected type");
 }
 
 TEST(CleanupTest, BasicLambda) {

From 1ce4bc2bcae4da3bab2ebeff1e8a4567812159c3 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:07:05 -0700
Subject: [PATCH 0500/1000] Automated Code Change

PiperOrigin-RevId: 997658425
---
 tensorflow/core/summary/loader.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/summary/loader.cc b/tensorflow/core/summary/loader.cc
index 08e4ea469b1..598d7ef5d2f 100644
--- a/tensorflow/core/summary/loader.cc
+++ b/tensorflow/core/summary/loader.cc
@@ -36,7 +36,7 @@ namespace {
 
 template <typename T>
 std::string AddCommas(T n) {
-  static_assert(std::is_integral<T>::value, "is_integral");
+  static_assert(std::is_integral_v<T>, "is_integral");
   std::string s = strings::StrCat(n);
   if (s.size() > 3) {
     int extra = s.size() / 3 - (s.size() % 3 == 0 ? 1 : 0);

From 770ef173964e65400dbcb9127d88602597ad16dc Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:08:27 -0700
Subject: [PATCH 0501/1000] Automated Code Change

PiperOrigin-RevId: 997658717
---
 tensorflow/core/common_runtime/session_factory.h           | 2 +-
 tensorflow/core/common_runtime/session_test.cc             | 2 +-
 tensorflow/core/common_runtime/simple_propagator_state.cc  | 2 +-
 tensorflow/core/common_runtime/stats_publisher_interface.h | 2 +-
 tensorflow/core/common_runtime/step_stats_collector.h      | 4 ++--
 tensorflow/core/common_runtime/threadpool_device.cc        | 2 +-
 6 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensorflow/core/common_runtime/session_factory.h b/tensorflow/core/common_runtime/session_factory.h
index 2bde986ab2d..41a0744039b 100644
--- a/tensorflow/core/common_runtime/session_factory.h
+++ b/tensorflow/core/common_runtime/session_factory.h
@@ -65,7 +65,7 @@ class SessionFactory {
     return absl::UnimplementedError("Reset()");
   }
 
-  virtual ~SessionFactory() {}
+  virtual ~SessionFactory() = default;
   static void Register(const std::string& runtime_type,
                        SessionFactory* factory);
   static absl::Status GetFactory(const SessionOptions& options,
diff --git a/tensorflow/core/common_runtime/session_test.cc b/tensorflow/core/common_runtime/session_test.cc
index 0c36356e65f..9dd1ddf0284 100644
--- a/tensorflow/core/common_runtime/session_test.cc
+++ b/tensorflow/core/common_runtime/session_test.cc
@@ -41,7 +41,7 @@ TEST(SessionTest, InvalidTargetReturnsNull) {
 // NewSession().
 class FakeSessionFactory : public SessionFactory {
  public:
-  FakeSessionFactory() {}
+  FakeSessionFactory() = default;
 
   bool AcceptsOptions(const SessionOptions& options) override {
     return absl::StartsWith(options.target, "fake");
diff --git a/tensorflow/core/common_runtime/simple_propagator_state.cc b/tensorflow/core/common_runtime/simple_propagator_state.cc
index 3855c6a3d6c..e951f91a209 100644
--- a/tensorflow/core/common_runtime/simple_propagator_state.cc
+++ b/tensorflow/core/common_runtime/simple_propagator_state.cc
@@ -43,7 +43,7 @@ SimplePropagatorState::SimplePropagatorState(
   immutable_state_.copy_pending_counts(pending_.get());
 }
 
-SimplePropagatorState::~SimplePropagatorState() {}
+SimplePropagatorState::~SimplePropagatorState() = default;
 
 void SimplePropagatorState::ActivateRoots(
     absl::Span<const NodeItem* const> roots, TaggedNodeSeq* ready) {
diff --git a/tensorflow/core/common_runtime/stats_publisher_interface.h b/tensorflow/core/common_runtime/stats_publisher_interface.h
index 2f0e3221be9..f3c98a33d02 100644
--- a/tensorflow/core/common_runtime/stats_publisher_interface.h
+++ b/tensorflow/core/common_runtime/stats_publisher_interface.h
@@ -63,7 +63,7 @@ class StatsPublisherInterface {
   virtual std::unique_ptr<ProfileHandler> GetProfileHandler(
       uint64_t step, int64_t execution_count, const RunOptions& ropts) = 0;
 
-  virtual ~StatsPublisherInterface() {}
+  virtual ~StatsPublisherInterface() = default;
 
   static void RegisterStatsPublisher(StatsPublisherFactory factory_fn);
 
diff --git a/tensorflow/core/common_runtime/step_stats_collector.h b/tensorflow/core/common_runtime/step_stats_collector.h
index 1c3503a8101..2ce73b4abf9 100644
--- a/tensorflow/core/common_runtime/step_stats_collector.h
+++ b/tensorflow/core/common_runtime/step_stats_collector.h
@@ -46,7 +46,7 @@ class Tensor;
 // that interfaces with the `Session` layer.
 class NodeExecStatsInterface {
  public:
-  virtual ~NodeExecStatsInterface() {}
+  virtual ~NodeExecStatsInterface() = default;
 
   // Called when the statistics collection for the node has finished. Once this
   // method is called, the caller should not make assumptions about the validity
@@ -137,7 +137,7 @@ class NodeExecStatsWrapper : public NodeExecStatsInterface {
 // that interfaces with the `Session` layer.
 class StepStatsCollectorInterface {
  public:
-  virtual ~StepStatsCollectorInterface() {}
+  virtual ~StepStatsCollectorInterface() = default;
 
   // Creates an instance of `NodeExecStatsInterface` that should be used for
   // collecting statistics about individual node execution.
diff --git a/tensorflow/core/common_runtime/threadpool_device.cc b/tensorflow/core/common_runtime/threadpool_device.cc
index b4b9e0826db..cfe81857ea3 100644
--- a/tensorflow/core/common_runtime/threadpool_device.cc
+++ b/tensorflow/core/common_runtime/threadpool_device.cc
@@ -101,7 +101,7 @@ ThreadPoolDevice::ThreadPoolDevice(const SessionOptions& options,
 #endif  // defined(ENABLE_ONEDNN_OPENMP) && defined(INTEL_MKL)
 }
 
-ThreadPoolDevice::~ThreadPoolDevice() {}
+ThreadPoolDevice::~ThreadPoolDevice() = default;
 
 Allocator* ThreadPoolDevice::GetAllocator(AllocatorAttributes attr) {
   return allocator_;

From 8a73d629bdddeb872a4327c96ede8dce7e6185fe Mon Sep 17 00:00:00 2001
From: nishad shabbir <nishadshabbir28@gmail.com>
Date: Mon, 12 Oct 2026 12:48:44 +0530
Subject: [PATCH 0502/1000] Update tensorflow/python/ops/gradients_impl.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 tensorflow/python/ops/gradients_impl.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/ops/gradients_impl.py b/tensorflow/python/ops/gradients_impl.py
index ad0b40c4847..50eea59c192 100644
--- a/tensorflow/python/ops/gradients_impl.py
+++ b/tensorflow/python/ops/gradients_impl.py
@@ -428,7 +428,7 @@ def hessians(ys,
   }
   # Compute first-order derivatives and iterate for each x in xs.
   hessians = []
-  with ops.name_scope(name, "hessians", list(ys_list) + list(xs)):
+  with ops.name_scope(name, "hessians", ys_list + xs):
     _gradients = gradients(ys, xs, **kwargs)
     for gradient, x in zip(_gradients, xs):
       # change shape to one-dimension without graph branching

From c064c013c825556e373831d875856ea62571b89d Mon Sep 17 00:00:00 2001
From: nishad shabbir <nishadshabbir28@gmail.com>
Date: Mon, 12 Oct 2026 12:49:02 +0530
Subject: [PATCH 0503/1000] Update tensorflow/python/ops/gradients_test.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
---
 tensorflow/python/ops/gradients_test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/ops/gradients_test.py b/tensorflow/python/ops/gradients_test.py
index e99e4827197..62b2d051d85 100644
--- a/tensorflow/python/ops/gradients_test.py
+++ b/tensorflow/python/ops/gradients_test.py
@@ -870,12 +870,12 @@ class HessianTest(test_util.TensorFlowTestCase):
       with ops.Graph().as_default() as g:
         x = constant_op.constant([1.0, 2.0, 3.0])
         y = math_ops.reduce_sum(x * x * x)
+        existing_ops = set(g.get_operations())
         if name is None:
           gradients.hessians(y, x)
         else:
           gradients.hessians(y, x, name=name)
-      added = [op.name for op in g.get_operations()
-               if not op.name.startswith(("Const", "mul", "Sum"))]
+        added = [op.name for op in g.get_operations() if op not in existing_ops]
       self.assertNotEmpty(added)
       for op_name in added:
         self.assertEqual(expected, op_name.split("/")[0])

From 561a377ad9ee6ee64af846763866ebeb5db66dcd Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:08:57 -0700
Subject: [PATCH 0504/1000] Automated Code Change

PiperOrigin-RevId: 997658818
---
 tensorflow/core/grappler/utils.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/grappler/utils.cc b/tensorflow/core/grappler/utils.cc
index 65518d1124f..c36cb8602fc 100644
--- a/tensorflow/core/grappler/utils.cc
+++ b/tensorflow/core/grappler/utils.cc
@@ -381,7 +381,7 @@ namespace {
 template <typename UniqueContainer>
 void EraseNodesFromGraphImpl(const UniqueContainer& nodes_to_delete,
                              GraphDef* graph) {
-  static_assert(std::is_same<typename UniqueContainer::value_type, int>::value,
+  static_assert(std::is_same_v<typename UniqueContainer::value_type, int>,
                 "Need to pass container of ints");
 
   int last = graph->node_size() - 1;

From e8ae36485b89b4af3eea63870f25c253519e16ee Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:11:23 -0700
Subject: [PATCH 0505/1000] Automated Code Change

PiperOrigin-RevId: 997659376
---
 .../tests/fuse_binary_into_following_affine_test.cc             | 2 +-
 .../tests/fuse_binary_into_preceding_affine_test.cc             | 2 +-
 .../lite/toco/graph_transformations/tests/lstm_utils_test.cc    | 2 +-
 .../tests/remove_successive_transpose_test.cc                   | 2 +-
 .../tests/resolve_constant_concatenation_test.cc                | 2 +-
 .../toco/graph_transformations/tests/unpack_quantize_test.cc    | 2 +-
 6 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensorflow/lite/toco/graph_transformations/tests/fuse_binary_into_following_affine_test.cc b/tensorflow/lite/toco/graph_transformations/tests/fuse_binary_into_following_affine_test.cc
index a66ad270a3d..b09c930cb15 100644
--- a/tensorflow/lite/toco/graph_transformations/tests/fuse_binary_into_following_affine_test.cc
+++ b/tensorflow/lite/toco/graph_transformations/tests/fuse_binary_into_following_affine_test.cc
@@ -39,7 +39,7 @@ std::vector<testing::Matcher<float>> ArrayFloatNear(
 
 class FuseBinaryIntoFollowingAffineTest : public ::testing::Test {
  protected:
-  FuseBinaryIntoFollowingAffineTest() {}
+  FuseBinaryIntoFollowingAffineTest() = default;
 
   void SetUp() override { model_ = std::make_unique<Model>(); }
 
diff --git a/tensorflow/lite/toco/graph_transformations/tests/fuse_binary_into_preceding_affine_test.cc b/tensorflow/lite/toco/graph_transformations/tests/fuse_binary_into_preceding_affine_test.cc
index 35888667d4b..43bf4757d9b 100644
--- a/tensorflow/lite/toco/graph_transformations/tests/fuse_binary_into_preceding_affine_test.cc
+++ b/tensorflow/lite/toco/graph_transformations/tests/fuse_binary_into_preceding_affine_test.cc
@@ -39,7 +39,7 @@ std::vector<testing::Matcher<float>> ArrayFloatNear(
 
 class FuseBinaryIntoPrecedingAffineTest : public ::testing::Test {
  protected:
-  FuseBinaryIntoPrecedingAffineTest() {}
+  FuseBinaryIntoPrecedingAffineTest() = default;
 
   void SetUp() override { model_ = std::make_unique<Model>(); }
 
diff --git a/tensorflow/lite/toco/graph_transformations/tests/lstm_utils_test.cc b/tensorflow/lite/toco/graph_transformations/tests/lstm_utils_test.cc
index a743b26414b..a204040b342 100644
--- a/tensorflow/lite/toco/graph_transformations/tests/lstm_utils_test.cc
+++ b/tensorflow/lite/toco/graph_transformations/tests/lstm_utils_test.cc
@@ -41,7 +41,7 @@ std::vector<testing::Matcher<float>> ArrayFloatNear(
 
 class CopyArrayDataTest : public ::testing::Test {
  public:
-  CopyArrayDataTest() {}
+  CopyArrayDataTest() = default;
 
   void PrepareBuffers(Model* model, std::initializer_list<float> src_data,
                       int src_dim_1, int src_dim_2,
diff --git a/tensorflow/lite/toco/graph_transformations/tests/remove_successive_transpose_test.cc b/tensorflow/lite/toco/graph_transformations/tests/remove_successive_transpose_test.cc
index 561ca830fcb..2463a3f7d8d 100644
--- a/tensorflow/lite/toco/graph_transformations/tests/remove_successive_transpose_test.cc
+++ b/tensorflow/lite/toco/graph_transformations/tests/remove_successive_transpose_test.cc
@@ -26,7 +26,7 @@ using ::testing::Test;
 
 class RemoveSuccessiveTransposeTest : public Test {
  protected:
-  RemoveSuccessiveTransposeTest() {}
+  RemoveSuccessiveTransposeTest() = default;
 
   void SetUp() override { model_ = std::make_unique<toco::Model>(); }
 
diff --git a/tensorflow/lite/toco/graph_transformations/tests/resolve_constant_concatenation_test.cc b/tensorflow/lite/toco/graph_transformations/tests/resolve_constant_concatenation_test.cc
index 159e24743f6..7f5eaab8c1e 100644
--- a/tensorflow/lite/toco/graph_transformations/tests/resolve_constant_concatenation_test.cc
+++ b/tensorflow/lite/toco/graph_transformations/tests/resolve_constant_concatenation_test.cc
@@ -102,7 +102,7 @@ std::vector<testing::Matcher<float>> ArrayFloatNear(
 
 class ResolveConstantConcatenationTest : public ::testing::Test {
  protected:
-  ResolveConstantConcatenationTest() {}
+  ResolveConstantConcatenationTest() = default;
 
   // Prepare a hypothetical TOCO model with one Concatenation operator in it
   // together with 4 arrays as its inputs.
diff --git a/tensorflow/lite/toco/graph_transformations/tests/unpack_quantize_test.cc b/tensorflow/lite/toco/graph_transformations/tests/unpack_quantize_test.cc
index 0b12905cdb1..a516bfd7ec1 100755
--- a/tensorflow/lite/toco/graph_transformations/tests/unpack_quantize_test.cc
+++ b/tensorflow/lite/toco/graph_transformations/tests/unpack_quantize_test.cc
@@ -30,7 +30,7 @@ namespace {
 
 class UnpackQuantizeTest : public ::testing::Test {
  protected:
-  UnpackQuantizeTest() {}
+  UnpackQuantizeTest() = default;
 
   // Prepare a hypothetical TOCO model with one unpack operator in it
   // together with 2 arrays as its outputs.

From 0c19a9e7f02a24aef24e288027bf5c7c80b94873 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:12:59 -0700
Subject: [PATCH 0506/1000] Automated Code Change

PiperOrigin-RevId: 997659900
---
 tensorflow/core/kernels/fuzzing/fuzz_session.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/fuzzing/fuzz_session.h b/tensorflow/core/kernels/fuzzing/fuzz_session.h
index 6da638b3378..333f78ed185 100644
--- a/tensorflow/core/kernels/fuzzing/fuzz_session.h
+++ b/tensorflow/core/kernels/fuzzing/fuzz_session.h
@@ -67,7 +67,7 @@ namespace fuzzing {
 class FuzzSession {
  public:
   FuzzSession() : initialized_(false) {}
-  virtual ~FuzzSession() {}
+  virtual ~FuzzSession() = default;
 
   // Constructs a Graph using the supplied Scope.
   // By convention, the graph should have inputs named "input1", ...

From 99dfb3228affea29af1f1903a9ea364de0d1019f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:13:19 -0700
Subject: [PATCH 0507/1000] Automated Code Change

PiperOrigin-RevId: 997660032
---
 tensorflow/core/common_runtime/graph_view.cc                  | 4 ++--
 .../core/common_runtime/optimize_function_graph_utils.cc      | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/common_runtime/graph_view.cc b/tensorflow/core/common_runtime/graph_view.cc
index 006a3c2b2c1..ac17ff913f9 100644
--- a/tensorflow/core/common_runtime/graph_view.cc
+++ b/tensorflow/core/common_runtime/graph_view.cc
@@ -51,9 +51,9 @@ std::string NodeItem::DebugString() const {
 }
 
 GraphView::~GraphView() {
-  static_assert(std::is_trivially_destructible<AllocatorAttributes>::value,
+  static_assert(std::is_trivially_destructible_v<AllocatorAttributes>,
                 "Update code if AllocatorAttributes gains a destructor");
-  static_assert(std::is_trivially_destructible<EdgeInfo>::value,
+  static_assert(std::is_trivially_destructible_v<EdgeInfo>,
                 "Update code if EdgeInfo gains a destructor");
   for (int i = 0; i < num_nodes_; i++) {
     NodeItem* n = node(i);
diff --git a/tensorflow/core/common_runtime/optimize_function_graph_utils.cc b/tensorflow/core/common_runtime/optimize_function_graph_utils.cc
index c4f0d5b8c55..316ceab326d 100644
--- a/tensorflow/core/common_runtime/optimize_function_graph_utils.cc
+++ b/tensorflow/core/common_runtime/optimize_function_graph_utils.cc
@@ -847,7 +847,7 @@ PreprocessAndPartitionGraph(
 
   // As not all devices might set number of threads for intra op
   // parallelisation we restrict this only to local device which does.
-  if (cpu_device && std::is_same<decltype(cpu_device), LocalDevice>::value &&
+  if (cpu_device && std::is_same_v<decltype(cpu_device), LocalDevice> &&
       cpu_device->tensorflow_cpu_worker_threads() != nullptr) {
     // Forward to the optimisation pass number of intra threads that are used to
     // parallelise operations.

From 316fa90634b29c5e217de9ca0a40bc35ac0a1756 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:13:49 -0700
Subject: [PATCH 0508/1000] Automated Code Change

PiperOrigin-RevId: 997660298
---
 .../c/experimental/stream_executor/stream_executor_test.cc      | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/c/experimental/stream_executor/stream_executor_test.cc b/tensorflow/c/experimental/stream_executor/stream_executor_test.cc
index 5108ffaf056..66ce8a3f40c 100644
--- a/tensorflow/c/experimental/stream_executor/stream_executor_test.cc
+++ b/tensorflow/c/experimental/stream_executor/stream_executor_test.cc
@@ -150,7 +150,7 @@ TEST(StreamExecutor, UnifiedMemoryAllocateNotSet) {
 /*** StreamExecutor behavior tests ***/
 class StreamExecutorTest : public ::testing::Test {
  protected:
-  StreamExecutorTest() {}
+  StreamExecutorTest() = default;
   void SetUp() override {
     test_util::PopulateDefaultPlatform(&platform_, &platform_fns_);
     test_util::PopulateDefaultDeviceFns(&device_fns_);

From 7853e57382e8b348c0a46dfff535e658b512ddb3 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:14:29 -0700
Subject: [PATCH 0509/1000] Automated Code Change

PiperOrigin-RevId: 997660492
---
 tensorflow/lite/delegates/coreml/builders/op_builder.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/lite/delegates/coreml/builders/op_builder.h b/tensorflow/lite/delegates/coreml/builders/op_builder.h
index 802db8b96b9..67e0d8df107 100644
--- a/tensorflow/lite/delegates/coreml/builders/op_builder.h
+++ b/tensorflow/lite/delegates/coreml/builders/op_builder.h
@@ -32,7 +32,7 @@ class OpBuilder;
 // API is experimental and subject to change.
 class TensorID {
  public:
-  TensorID() {}
+  TensorID() = default;
   TensorID(int node, int output_id) : node_(node), output_id_(output_id) {}
 
   std::string ToString() const;
@@ -99,7 +99,7 @@ class OpBuilder {
  public:
   explicit OpBuilder(GraphBuilder* graph_builder)
       : graph_builder_(graph_builder) {}
-  virtual ~OpBuilder() {}
+  virtual ~OpBuilder() = default;
 
   // Returns the Layer this builder responsible for.
   // Ownership is transferred to caller.

From 4cb4004a4e78aa315096dfa401b672eda7e5043a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:15:08 -0700
Subject: [PATCH 0510/1000] Automated Code Change

PiperOrigin-RevId: 997660679
---
 tensorflow/core/kernels/data/random_seed_ops.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/data/random_seed_ops.h b/tensorflow/core/kernels/data/random_seed_ops.h
index 393fe21da63..79ebc40d069 100644
--- a/tensorflow/core/kernels/data/random_seed_ops.h
+++ b/tensorflow/core/kernels/data/random_seed_ops.h
@@ -56,7 +56,7 @@ class RandomSeeds {
 // generated.
 class SeedGenerator {
  public:
-  virtual ~SeedGenerator() {}
+  virtual ~SeedGenerator() = default;
 
   virtual int64_t seed() const = 0;
   virtual int64_t seed2() const = 0;

From 699b6f361b1b3eddb18c9de3c742354711f43772 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:15:47 -0700
Subject: [PATCH 0511/1000] Automated Code Change

PiperOrigin-RevId: 997660921
---
 tensorflow/lite/kernels/ctc/ctc_beam_entry.h  | 2 +-
 tensorflow/lite/kernels/ctc/ctc_beam_scorer.h | 2 +-
 tensorflow/lite/kernels/ctc/ctc_beam_search.h | 2 +-
 tensorflow/lite/kernels/ctc/ctc_decoder.h     | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/lite/kernels/ctc/ctc_beam_entry.h b/tensorflow/lite/kernels/ctc/ctc_beam_entry.h
index eede67c716b..62a43469785 100644
--- a/tensorflow/lite/kernels/ctc/ctc_beam_entry.h
+++ b/tensorflow/lite/kernels/ctc/ctc_beam_entry.h
@@ -134,7 +134,7 @@ class BeamRoot {
 template <class CTCBeamState = EmptyBeamState>
 class BeamComparer {
  public:
-  virtual ~BeamComparer() {}
+  virtual ~BeamComparer() = default;
   virtual bool inline operator()(const BeamEntry<CTCBeamState>* a,
                                  const BeamEntry<CTCBeamState>* b) const {
     return a->newp.total > b->newp.total;
diff --git a/tensorflow/lite/kernels/ctc/ctc_beam_scorer.h b/tensorflow/lite/kernels/ctc/ctc_beam_scorer.h
index 7936cd020ca..ba910f2d918 100644
--- a/tensorflow/lite/kernels/ctc/ctc_beam_scorer.h
+++ b/tensorflow/lite/kernels/ctc/ctc_beam_scorer.h
@@ -39,7 +39,7 @@ namespace ctc {
 template <typename CTCBeamState>
 class BaseBeamScorer {
  public:
-  virtual ~BaseBeamScorer() {}
+  virtual ~BaseBeamScorer() = default;
   // State initialization.
   virtual void InitializeState(CTCBeamState* root) const {}
   // ExpandState is called when expanding a beam to one of its children.
diff --git a/tensorflow/lite/kernels/ctc/ctc_beam_search.h b/tensorflow/lite/kernels/ctc/ctc_beam_search.h
index df442f9794c..93e06705ef9 100644
--- a/tensorflow/lite/kernels/ctc/ctc_beam_search.h
+++ b/tensorflow/lite/kernels/ctc/ctc_beam_search.h
@@ -94,7 +94,7 @@ class CTCBeamSearchDecoder : public CTCDecoder {
     Reset();
   }
 
-  ~CTCBeamSearchDecoder() override {}
+  ~CTCBeamSearchDecoder() override = default;
 
   // Run the hibernating beam search algorithm on the given input.
   bool Decode(const CTCDecoder::SequenceLength& seq_len,
diff --git a/tensorflow/lite/kernels/ctc/ctc_decoder.h b/tensorflow/lite/kernels/ctc/ctc_decoder.h
index 1c5591b08d1..a72a13048b6 100644
--- a/tensorflow/lite/kernels/ctc/ctc_decoder.h
+++ b/tensorflow/lite/kernels/ctc/ctc_decoder.h
@@ -46,7 +46,7 @@ class CTCDecoder {
         batch_size_(batch_size),
         merge_repeated_(merge_repeated) {}
 
-  virtual ~CTCDecoder() {}
+  virtual ~CTCDecoder() = default;
 
   // Dimensionality of the input/output is expected to be:
   //  - seq_len[b] - b = 0 to batch_size_

From b947f7f3fd9c5d0e187c88f65ed316a824b35815 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:16:26 -0700
Subject: [PATCH 0512/1000] Automated Code Change

PiperOrigin-RevId: 997661142
---
 tensorflow/core/util/gpu_device_functions.h |  7 +++----
 tensorflow/core/util/gpu_kernel_helper.h    | 13 +++++--------
 tensorflow/core/util/gpu_launch_config.h    |  2 +-
 3 files changed, 9 insertions(+), 13 deletions(-)

diff --git a/tensorflow/core/util/gpu_device_functions.h b/tensorflow/core/util/gpu_device_functions.h
index bdbc8610c23..de1946bb3a9 100644
--- a/tensorflow/core/util/gpu_device_functions.h
+++ b/tensorflow/core/util/gpu_device_functions.h
@@ -620,7 +620,7 @@ __device__ long long GpuAtomicCasHelper(long long* ptr, F accumulate) {
 
 template <typename From, typename To>
 using ToTypeIfConvertible =
-    typename std::enable_if<std::is_convertible<From, To>::value, To>::type;
+    std::enable_if_t<std::is_convertible_v<From, To>, To>;
 
 template <typename T>
 struct CudaSupportedTypeImpl {
@@ -634,9 +634,8 @@ struct CudaSupportedTypeImpl<long long> {
 
 template <>
 struct CudaSupportedTypeImpl<unsigned long> {
-  using type =
-      typename std::conditional<sizeof(unsigned long) == sizeof(unsigned int),
-                                unsigned int, unsigned long long>::type;
+  using type = std::conditional_t<sizeof(unsigned long) == sizeof(unsigned int),
+                                  unsigned int, unsigned long long>;
 };
 
 template <>
diff --git a/tensorflow/core/util/gpu_kernel_helper.h b/tensorflow/core/util/gpu_kernel_helper.h
index a866bd98975..3a9a1c11c6f 100644
--- a/tensorflow/core/util/gpu_kernel_helper.h
+++ b/tensorflow/core/util/gpu_kernel_helper.h
@@ -250,23 +250,20 @@ class alignas(alignof(T) * N) AlignedVector {
   // Note: This is required for T=Eigen::half because it only supports explicit
   // conversions from other types and its template constructor is too relaxed
   // to be able to use std::is_constructible.
-  template <typename U, typename std::enable_if<std::is_arithmetic<U>::value,
-                                                int>::type = 0>
+  template <typename U, std::enable_if_t<std::is_arithmetic_v<U>, int> = 0>
   __host__ __device__ explicit AlignedVector(U uniform_u) {
     value_type uniform(uniform_u);
     UNROLL_ON_DEVICE for (int i = 0; i < kSize; ++i) { values_[i] = uniform; }
   }
   // Implicit conversion.
-  template <typename U, typename std::enable_if<
-                            std::is_convertible<U, T>::value, int>::type = 0>
+  template <typename U, std::enable_if_t<std::is_convertible_v<U, T>, int> = 0>
   __host__ __device__ AlignedVector(const AlignedVector<U, N>& other) {
     UNROLL_ON_DEVICE for (int i = 0; i < kSize; ++i) { values_[i] = other[i]; }
   }
   // Explicit conversion.
-  template <typename U,
-            typename std::enable_if<!std::is_convertible<U, T>::value &&
-                                        std::is_constructible<T, U>::value,
-                                    int>::type = 0>
+  template <typename U, std::enable_if_t<!std::is_convertible_v<U, T> &&
+                                             std::is_constructible_v<T, U>,
+                                         int> = 0>
   __host__ __device__ explicit AlignedVector(const AlignedVector<U, N>& other) {
     UNROLL_ON_DEVICE for (int i = 0; i < kSize; ++i) {
       values_[i] = T(other[i]);
diff --git a/tensorflow/core/util/gpu_launch_config.h b/tensorflow/core/util/gpu_launch_config.h
index b7485fe315f..f6eb463dccf 100644
--- a/tensorflow/core/util/gpu_launch_config.h
+++ b/tensorflow/core/util/gpu_launch_config.h
@@ -484,7 +484,7 @@ using NoneTrue = std::is_same<BoolPack<Bs..., false>, BoolPack<false, Bs...>>;
 // Returns whether none of the types in Ts is a reference.
 template <typename... Ts>
 constexpr bool NoneIsReference() {
-  return NoneTrue<(std::is_reference<Ts>::value)...>::value;
+  return NoneTrue<(std::is_reference_v<Ts>)...>::value;
 }
 }  // namespace detail
 }  // namespace tensorflow

From b5e0e6bd93f6bb3690a742c240cfdb331cbe0add Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:17:33 -0700
Subject: [PATCH 0513/1000] Automated Code Change

PiperOrigin-RevId: 997661576
---
 tensorflow/core/graph/edgeset.h            | 2 +-
 tensorflow/core/graph/graph.h              | 2 +-
 tensorflow/core/graph/graph_def_builder.cc | 2 +-
 tensorflow/core/graph/graph_test.cc        | 2 +-
 tensorflow/core/graph/subgraph.h           | 2 +-
 tensorflow/core/graph/subgraph_test.cc     | 2 +-
 6 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/tensorflow/core/graph/edgeset.h b/tensorflow/core/graph/edgeset.h
index 8ade339d25e..cdd993e59e0 100644
--- a/tensorflow/core/graph/edgeset.h
+++ b/tensorflow/core/graph/edgeset.h
@@ -98,7 +98,7 @@ class EdgeSet::const_iterator {
   typedef typename EdgeSet::difference_type difference_type;
   typedef std::forward_iterator_tag iterator_category;
 
-  const_iterator() {}
+  const_iterator() = default;
 
   const_iterator& operator++();
   const_iterator operator++(int /*unused*/);
diff --git a/tensorflow/core/graph/graph.h b/tensorflow/core/graph/graph.h
index 10b29e09756..5d6caaf4489 100644
--- a/tensorflow/core/graph/graph.h
+++ b/tensorflow/core/graph/graph.h
@@ -455,7 +455,7 @@ class Edge {
   std::string DebugString() const;
 
  private:
-  Edge() {}
+  Edge() = default;
 
   friend class EdgeSetTest;
   friend class GraphTest;
diff --git a/tensorflow/core/graph/graph_def_builder.cc b/tensorflow/core/graph/graph_def_builder.cc
index a4f08eab66b..711f588fb8c 100644
--- a/tensorflow/core/graph/graph_def_builder.cc
+++ b/tensorflow/core/graph/graph_def_builder.cc
@@ -24,7 +24,7 @@ namespace tensorflow {
 
 GraphDefBuilder::Options::Options(Graph* graph, absl::Status* status)
     : graph_(graph), status_(status) {}
-GraphDefBuilder::Options::~Options() {}
+GraphDefBuilder::Options::~Options() = default;
 
 GraphDefBuilder::Options GraphDefBuilder::Options::WithName(
     absl::string_view name) const {
diff --git a/tensorflow/core/graph/graph_test.cc b/tensorflow/core/graph/graph_test.cc
index fb5ce07959a..1329326e954 100644
--- a/tensorflow/core/graph/graph_test.cc
+++ b/tensorflow/core/graph/graph_test.cc
@@ -57,7 +57,7 @@ REGISTER_OP("TwoInputsOneOutput")
 class GraphTest : public ::testing::Test {
  protected:
   GraphTest() : graph_(OpRegistry::Global()) {}
-  ~GraphTest() override {}
+  ~GraphTest() override = default;
 
   static void VerifyNodes(Node* node, const std::vector<Node*>& expected_in,
                           const std::vector<Node*>& expected_out) {
diff --git a/tensorflow/core/graph/subgraph.h b/tensorflow/core/graph/subgraph.h
index c8843a37d58..a35780ee97e 100644
--- a/tensorflow/core/graph/subgraph.h
+++ b/tensorflow/core/graph/subgraph.h
@@ -53,7 +53,7 @@ class PruneRewrite {
   PruneRewrite(const std::string* endpoint_name,
                const DeviceAttributes* device_info)
       : endpoint_name_(endpoint_name), device_info_(device_info) {}
-  virtual ~PruneRewrite() {}
+  virtual ~PruneRewrite() = default;
 
   // Creates a new node whose output replaces the given `tensor` in graph `g`.
   // The node will be assigned to the device named in `device_info`.
diff --git a/tensorflow/core/graph/subgraph_test.cc b/tensorflow/core/graph/subgraph_test.cc
index a5f4be88e8e..d94a0e7370c 100644
--- a/tensorflow/core/graph/subgraph_test.cc
+++ b/tensorflow/core/graph/subgraph_test.cc
@@ -47,7 +47,7 @@ class SubgraphTest : public ::testing::Test {
     device_info_.set_incarnation(0);
   }
 
-  ~SubgraphTest() override {}
+  ~SubgraphTest() override = default;
 
   void ExpectOK(const std::string& gdef_ascii) {
     CHECK(protobuf::TextFormat::ParseFromString(gdef_ascii, &gdef_));

From 99e11cfd3675835bc5a1829c2746b532d5dada18 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:17:57 -0700
Subject: [PATCH 0514/1000] Automated Code Change

PiperOrigin-RevId: 997661722
---
 tensorflow/c/eager/tape.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/c/eager/tape.h b/tensorflow/c/eager/tape.h
index ad7c455a813..be0ff796bfc 100644
--- a/tensorflow/c/eager/tape.h
+++ b/tensorflow/c/eager/tape.h
@@ -94,7 +94,7 @@ using OpTape =
 template <typename Gradient, typename BackwardFunction, typename TapeTensor>
 class VSpace {
  public:
-  virtual ~VSpace() {}
+  virtual ~VSpace() = default;
 
   // Returns the number of elements in the gradient tensor.
   virtual int64_t NumElements(Gradient* tensor) const = 0;

From b19c244c99bb9576211a0e1e696eba7dfc9ca144 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:18:15 -0700
Subject: [PATCH 0515/1000] Automated Code Change

PiperOrigin-RevId: 997661830
---
 tensorflow/core/distributed_runtime/local_master.h          | 2 +-
 tensorflow/core/distributed_runtime/master_interface.h      | 2 +-
 tensorflow/core/distributed_runtime/master_session.cc       | 2 +-
 tensorflow/core/distributed_runtime/message_wrappers.cc     | 2 +-
 tensorflow/core/distributed_runtime/message_wrappers.h      | 6 +++---
 .../core/distributed_runtime/rendezvous_mgr_interface.h     | 2 +-
 6 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/tensorflow/core/distributed_runtime/local_master.h b/tensorflow/core/distributed_runtime/local_master.h
index b9fe78e8591..39c2feb1592 100644
--- a/tensorflow/core/distributed_runtime/local_master.h
+++ b/tensorflow/core/distributed_runtime/local_master.h
@@ -38,7 +38,7 @@ class Master;
 // for cancellation.
 class LocalMaster : public MasterInterface {
  public:
-  ~LocalMaster() override {}
+  ~LocalMaster() override = default;
 
   absl::Status CreateSession(CallOptions* call_options,
                              const CreateSessionRequest* request,
diff --git a/tensorflow/core/distributed_runtime/master_interface.h b/tensorflow/core/distributed_runtime/master_interface.h
index 89d4a0ffbf6..94b10d885b6 100644
--- a/tensorflow/core/distributed_runtime/master_interface.h
+++ b/tensorflow/core/distributed_runtime/master_interface.h
@@ -32,7 +32,7 @@ namespace tensorflow {
 // roundtrip.
 class MasterInterface {
  public:
-  virtual ~MasterInterface() {}
+  virtual ~MasterInterface() = default;
   virtual absl::Status CreateSession(CallOptions* call_options,
                                      const CreateSessionRequest* request,
                                      CreateSessionResponse* response) = 0;
diff --git a/tensorflow/core/distributed_runtime/master_session.cc b/tensorflow/core/distributed_runtime/master_session.cc
index d3eccb74e4b..6a988a3dd30 100644
--- a/tensorflow/core/distributed_runtime/master_session.cc
+++ b/tensorflow/core/distributed_runtime/master_session.cc
@@ -520,7 +520,7 @@ class RunManyGraphs {
  public:
   explicit RunManyGraphs(int num) : calls_(num), pending_(num) {}
 
-  ~RunManyGraphs() {}
+  ~RunManyGraphs() = default;
 
   // Returns the index-th call.
   struct Call {
diff --git a/tensorflow/core/distributed_runtime/message_wrappers.cc b/tensorflow/core/distributed_runtime/message_wrappers.cc
index 16a437f63a7..d32658cabb6 100644
--- a/tensorflow/core/distributed_runtime/message_wrappers.cc
+++ b/tensorflow/core/distributed_runtime/message_wrappers.cc
@@ -830,7 +830,7 @@ void NonOwnedProtoRunGraphResponse::AddPartitionGraph(
   *graph_def = partition_graph;
 }
 
-MutableRunStepResponseWrapper::~MutableRunStepResponseWrapper() {}
+MutableRunStepResponseWrapper::~MutableRunStepResponseWrapper() = default;
 
 size_t InMemoryRunStepResponse::num_tensors() const { return tensors_.size(); }
 
diff --git a/tensorflow/core/distributed_runtime/message_wrappers.h b/tensorflow/core/distributed_runtime/message_wrappers.h
index b911d23245b..3133f996d49 100644
--- a/tensorflow/core/distributed_runtime/message_wrappers.h
+++ b/tensorflow/core/distributed_runtime/message_wrappers.h
@@ -49,7 +49,7 @@ namespace tensorflow {
 // TensorFlow master.
 class RunStepRequestWrapper {
  public:
-  virtual ~RunStepRequestWrapper() {}
+  virtual ~RunStepRequestWrapper() = default;
 
   // REQUIRED: session_handle must be returned by a CreateSession call
   // to the same master service.
@@ -250,7 +250,7 @@ class ProtoRunStepRequest : public RunStepRequestWrapper {
 // TensorFlow worker.
 class RunGraphRequestWrapper {
  public:
-  virtual ~RunGraphRequestWrapper() {}
+  virtual ~RunGraphRequestWrapper() = default;
 
   // The session handle used to register the graph. If empty, a single global
   // namespace is used.
@@ -475,7 +475,7 @@ class ProtoRunGraphRequest : public RunGraphRequestWrapper {
 // is always used as a mutable pointer.
 class MutableRunGraphResponseWrapper {
  public:
-  virtual ~MutableRunGraphResponseWrapper() {}
+  virtual ~MutableRunGraphResponseWrapper() = default;
 
   // A list of tensors corresponding to those requested by
   // `RunGraphRequest.recv_key`.
diff --git a/tensorflow/core/distributed_runtime/rendezvous_mgr_interface.h b/tensorflow/core/distributed_runtime/rendezvous_mgr_interface.h
index d599b80093c..e5b79366ce1 100644
--- a/tensorflow/core/distributed_runtime/rendezvous_mgr_interface.h
+++ b/tensorflow/core/distributed_runtime/rendezvous_mgr_interface.h
@@ -72,7 +72,7 @@ class RemoteRendezvous : public Rendezvous {
 class RendezvousMgrInterface {
  public:
   RendezvousMgrInterface() = default;
-  virtual ~RendezvousMgrInterface() {}
+  virtual ~RendezvousMgrInterface() = default;
 
   // Returns Rendezvous supporting send and recv among workers in the
   // "step_id".  The caller takes ownership of one reference on the

From 501f127563722e11fbf7dacebf00c8a5fee3b0a5 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:18:24 -0700
Subject: [PATCH 0516/1000] Automated Code Change

PiperOrigin-RevId: 997661885
---
 tensorflow/core/transforms/remapper/pass.cc            | 10 +++++-----
 tensorflow/core/transforms/remapper/remapping_helper.h |  8 ++++----
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/tensorflow/core/transforms/remapper/pass.cc b/tensorflow/core/transforms/remapper/pass.cc
index 189f2f3a666..e1e76e47a2a 100644
--- a/tensorflow/core/transforms/remapper/pass.cc
+++ b/tensorflow/core/transforms/remapper/pass.cc
@@ -395,7 +395,7 @@ class BasePatternActivationRewriter : public BasePatternRewriter {
 
   using BasePattern = typename BasePatternRewriter::Pattern;
   using Pattern = std::conditional_t<
-      std::is_same<BasePatternRewriter, ContractionBiasAddRewriter>::value,
+      std::is_same_v<BasePatternRewriter, ContractionBiasAddRewriter>,
       ContractionBiasAddActivation, void>;
 
   bool matchPattern(Operation *op, BasePattern &base_pattern,
@@ -406,8 +406,8 @@ class BasePatternActivationRewriter : public BasePatternRewriter {
     if (this->helper_.HasControlOperandsOrResultUsers(op)) return false;
 
     // TODO(intel-tf): Add support for more patterns.
-    if constexpr (std::is_same<BasePattern, ContractionBiasAdd>::value &&
-                  std::is_same<Pattern, ContractionBiasAddActivation>::value) {
+    if constexpr (std::is_same_v<BasePattern, ContractionBiasAdd> &&
+                  std::is_same_v<Pattern, ContractionBiasAddActivation>) {
       Operation *bias_add_op = op->getOperand(0).getDefiningOp();
       if (!this->helper_.getDialect()->IsBiasAdd(bias_add_op) ||
           !this->helper_.HaveSameDataType(op, bias_add_op) ||
@@ -431,8 +431,8 @@ class BasePatternActivationRewriter : public BasePatternRewriter {
     BasePattern base_pattern;
     Pattern pattern;
     if (!matchPattern(op, base_pattern, pattern)) return failure();
-    if constexpr (!std::is_same<BasePatternRewriter,
-                                ContractionBiasAddRewriter>::value) {
+    if constexpr (!std::is_same_v<BasePatternRewriter,
+                                  ContractionBiasAddRewriter>) {
       return failure();
     }
     if (!this->helper_.IsDeviceCompatible(pattern)) return failure();
diff --git a/tensorflow/core/transforms/remapper/remapping_helper.h b/tensorflow/core/transforms/remapper/remapping_helper.h
index 1d8db8fc459..117449cb7cf 100644
--- a/tensorflow/core/transforms/remapper/remapping_helper.h
+++ b/tensorflow/core/transforms/remapper/remapping_helper.h
@@ -223,10 +223,10 @@ class OpPropertyHelper : public OpCatHelper {
   template <typename Pattern>
   bool IsDeviceCompatible(const Pattern& pattern) const {
     // Currently, this function is used by contraction based fussion.
-    if constexpr (!std::is_same<Pattern, ContractionBiasAdd>::value &&
-                  !std::is_same<Pattern, ContractionBiasAddActivation>::value &&
-                  !std::is_same<Pattern, ContractionBiasAddAdd>::value &&
-                  !std::is_same<Pattern, ContractionBiasAddActivation>::value) {
+    if constexpr (!std::is_same_v<Pattern, ContractionBiasAdd> &&
+                  !std::is_same_v<Pattern, ContractionBiasAddActivation> &&
+                  !std::is_same_v<Pattern, ContractionBiasAddAdd> &&
+                  !std::is_same_v<Pattern, ContractionBiasAddActivation>) {
       return false;
     }
     return IsGpuCompatible(pattern) || IsCpuCompatible(pattern.contraction);

From 759eb2a5bae80420888e40f94d8c7c542759a4f4 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:19:01 -0700
Subject: [PATCH 0517/1000] Automated Code Change

PiperOrigin-RevId: 997662051
---
 tensorflow/core/common_runtime/gpu/gpu_allocator_retry_test.cc | 2 +-
 tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.h         | 2 +-
 tensorflow/core/common_runtime/gpu/gpu_device.cc               | 2 +-
 tensorflow/core/common_runtime/gpu/gpu_device_factory.cc       | 2 +-
 tensorflow/core/common_runtime/gpu/gpu_process_state.h         | 2 +-
 5 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensorflow/core/common_runtime/gpu/gpu_allocator_retry_test.cc b/tensorflow/core/common_runtime/gpu/gpu_allocator_retry_test.cc
index 5f3f2babf98..1e4d7fa969f 100644
--- a/tensorflow/core/common_runtime/gpu/gpu_allocator_retry_test.cc
+++ b/tensorflow/core/common_runtime/gpu/gpu_allocator_retry_test.cc
@@ -124,7 +124,7 @@ class AlternatingBarrier {
 
 class GPUAllocatorRetryTest : public ::testing::Test {
  protected:
-  GPUAllocatorRetryTest() {}
+  GPUAllocatorRetryTest() = default;
 
   void LaunchConsumerThreads(int num_consumers, int cap_needed) {
     barrier_ = std::make_unique<AlternatingBarrier>(num_consumers);
diff --git a/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.h b/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.h
index e1f4f0d9419..fd97d2bd04d 100644
--- a/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.h
+++ b/tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.h
@@ -52,7 +52,7 @@ class GPUBFCAllocator : public tsl::BFCAllocator {
                   size_t total_memory, const std::string& name,
                   const Options& opts);
 
-  ~GPUBFCAllocator() override {}
+  ~GPUBFCAllocator() override = default;
 
   GPUBFCAllocator(const GPUBFCAllocator&) = delete;
   void operator=(const GPUBFCAllocator&) = delete;
diff --git a/tensorflow/core/common_runtime/gpu/gpu_device.cc b/tensorflow/core/common_runtime/gpu/gpu_device.cc
index b676fbaed6d..5abc7dec326 100644
--- a/tensorflow/core/common_runtime/gpu/gpu_device.cc
+++ b/tensorflow/core/common_runtime/gpu/gpu_device.cc
@@ -168,7 +168,7 @@ class EigenGpuStreamDevice : public ::Eigen::StreamInterface {
         semaphore_(nullptr),
         context_(nullptr),
         device_(this) {}
-  ~EigenGpuStreamDevice() override {}
+  ~EigenGpuStreamDevice() override = default;
 
   void Reinitialize(OpKernelContext* context, gpuStream_t gpu_stream,
                     tsl::PlatformDeviceId platform_device_id,
diff --git a/tensorflow/core/common_runtime/gpu/gpu_device_factory.cc b/tensorflow/core/common_runtime/gpu/gpu_device_factory.cc
index fa7e968cad1..6a49c1f51ed 100644
--- a/tensorflow/core/common_runtime/gpu/gpu_device_factory.cc
+++ b/tensorflow/core/common_runtime/gpu/gpu_device_factory.cc
@@ -103,7 +103,7 @@ class GPUCompatibleCPUDevice : public ThreadPoolDevice {
           options.config.gpu_options().force_gpu_compatible();
     }
   }
-  ~GPUCompatibleCPUDevice() override {}
+  ~GPUCompatibleCPUDevice() override = default;
 
   Allocator* GetAllocator(AllocatorAttributes attr) override {
     GPUProcessState* ps = GPUProcessState::singleton();
diff --git a/tensorflow/core/common_runtime/gpu/gpu_process_state.h b/tensorflow/core/common_runtime/gpu/gpu_process_state.h
index 19f8448ed5e..477a299d6b0 100644
--- a/tensorflow/core/common_runtime/gpu/gpu_process_state.h
+++ b/tensorflow/core/common_runtime/gpu/gpu_process_state.h
@@ -140,7 +140,7 @@ class GPUProcessState {
   // GPUProcessState is a singleton that should not normally be deleted except
   // at process shutdown.
   GPUProcessState();
-  virtual ~GPUProcessState() {}
+  virtual ~GPUProcessState() = default;
   friend class GPUDeviceTest;
 
   // Helper method for unit tests to reset the ProcessState singleton by

From 070006097a9e596cb86a06d857777cc004ce9688 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:19:01 -0700
Subject: [PATCH 0518/1000] Automated Code Change

PiperOrigin-RevId: 997662055
---
 tensorflow/c/experimental/gradients/array_grad.cc | 2 +-
 tensorflow/c/experimental/gradients/nn_grad.cc    | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/c/experimental/gradients/array_grad.cc b/tensorflow/c/experimental/gradients/array_grad.cc
index 6b599495e72..53cdf791de8 100644
--- a/tensorflow/c/experimental/gradients/array_grad.cc
+++ b/tensorflow/c/experimental/gradients/array_grad.cc
@@ -35,7 +35,7 @@ class IdentityNGradientFunction : public GradientFunction {
     }
     return absl::OkStatus();
   }
-  ~IdentityNGradientFunction() override {}
+  ~IdentityNGradientFunction() override = default;
 };
 }  // namespace
 
diff --git a/tensorflow/c/experimental/gradients/nn_grad.cc b/tensorflow/c/experimental/gradients/nn_grad.cc
index a56fa907f17..9ff5c97a379 100644
--- a/tensorflow/c/experimental/gradients/nn_grad.cc
+++ b/tensorflow/c/experimental/gradients/nn_grad.cc
@@ -107,7 +107,7 @@ class SparseSoftmaxCrossEntropyWithLogitsGradientFunction
     grad_inputs[1] = nullptr;
     return absl::OkStatus();
   }
-  ~SparseSoftmaxCrossEntropyWithLogitsGradientFunction() override {}
+  ~SparseSoftmaxCrossEntropyWithLogitsGradientFunction() override = default;
 
  private:
   vector<AbstractTensorHandle*> forward_outputs_;
@@ -146,7 +146,7 @@ class BiasAddGradientFunction : public GradientFunction {
 
     return absl::OkStatus();
   }
-  ~BiasAddGradientFunction() override {}
+  ~BiasAddGradientFunction() override = default;
 
  private:
   AttrBuilder forward_attrs_;

From d054926021a3f6621a644c03b24d40b961e9335a Mon Sep 17 00:00:00 2001
From: destro4evr-rgb <destro4evr@proton.me>
Date: Sat, 10 Oct 2026 04:25:27 -0400
Subject: [PATCH 0519/1000] security: validate batch dimension in
 BarrierInsertMany

TryInsertMany iterates over keys.NumElements() entries and calls
values_matrix.chip<0>(i) for each, but never checked that
values.dim_size(0) == keys.NumElements().  When the values tensor has
fewer rows than keys, chip<0>(i) reads past the end of the values buffer
(heap OOB read; Eigen does not bounds-check in release builds).

Add an OP_REQUIRES_ASYNC that enforces the batch-dimension parity before
entering the per-key loop, matching the shape contract documented for the
op and used in analogous queue kernel implementations.
---
 tensorflow/core/kernels/barrier_ops.cc        | 26 ++++++++++---
 .../data_structures/barrier_ops_test.py       | 38 +++++++++++++++++++
 2 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/tensorflow/core/kernels/barrier_ops.cc b/tensorflow/core/kernels/barrier_ops.cc
index 830d5e3c85f..52027d260fb 100644
--- a/tensorflow/core/kernels/barrier_ops.cc
+++ b/tensorflow/core/kernels/barrier_ops.cc
@@ -98,7 +98,14 @@ class Barrier : public ResourceBase {
                          ": received shape ", element_shape.DebugString())),
         callback);
     if (element_shape.dims() > 0) element_shape.RemoveDim(0);
-    const std::size_t num_inserted = keys.NumElements();
+    const int64_t num_inserted = keys.NumElements();
+    OP_REQUIRES_ASYNC(
+        ctx, values.dims() > 0 && values.dim_size(0) == num_inserted,
+        absl::InvalidArgumentError(absl::StrCat(
+            "keys and values must have the same batch size. Got ", num_inserted,
+            " keys but values has shape ", values.shape().DebugString(),
+            " in barrier ", name_, ".")),
+        callback);
 
     // For each key, update the corresponding incomplete tuple with the
     // the corresponding given value at component_index.
@@ -128,7 +135,7 @@ class Barrier : public ResourceBase {
       // them in a vector
       std::vector<Tuple> ready_tuples;
 
-      for (int i = 0; i < num_inserted; ++i) {
+      for (int64_t i = 0; i < num_inserted; ++i) {
         OP_REQUIRES_OK_ASYNC(
             ctx,
             InsertOneLocked<T>(ctx, keys, values, element_shape,
@@ -158,12 +165,12 @@ class Barrier : public ResourceBase {
       // tuple by slicing entries into new Tensors.  This part is slow
       // but seems the cleanest solution for now.
       insert_tuple.reserve(2 + num_components());  // indices, keys, rest
-      int insertion_size = ready_tuples.size();
-      for (int i = 0; i < 2 + num_components(); ++i) {
+      int64_t insertion_size = ready_tuples.size();
+      for (int64_t i = 0; i < 2 + num_components(); ++i) {
         TensorShape component_shape(ready_tuples[0][i].shape());
         component_shape.InsertDim(0, insertion_size);
         Tensor component(ready_tuples[0][i].dtype(), component_shape);
-        for (int b = 0; b < insertion_size; ++b) {
+        for (int64_t b = 0; b < insertion_size; ++b) {
           OP_REQUIRES_OK_ASYNC(
               ctx,
               batch_util::CopyElementToSlice(std::move(ready_tuples[b][i]),
@@ -307,7 +314,7 @@ class Barrier : public ResourceBase {
   absl::Status InsertOneLocked(OpKernelContext* ctx, const Tensor& keys,
                                const Tensor& values,
                                const TensorShape& element_shape,
-                               int component_index, int i,
+                               int component_index, int64_t i,
                                std::vector<Tuple>* ready_tuples,
                                bool* new_elements)
       TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
@@ -525,6 +532,8 @@ class InsertManyOp : public BarrierOpKernel {
       : BarrierOpKernel(context) {
     OP_REQUIRES_OK(context,
                    context->GetAttr("component_index", &component_index_));
+    OP_REQUIRES(context, component_index_ >= 0,
+                absl::InvalidArgumentError("component_index must be >= 0"));
   }
 
  protected:
@@ -546,6 +555,11 @@ class InsertManyOp : public BarrierOpKernel {
     const Tensor* keys;
     const Tensor* values;
     OP_REQUIRES_OK_ASYNC(ctx, ctx->input("keys", &keys), callback);
+    OP_REQUIRES_ASYNC(
+        ctx, TensorShapeUtils::IsVector(keys->shape()),
+        absl::InvalidArgumentError(absl::StrCat(
+            "keys must be 1-D, got shape ", keys->shape().DebugString())),
+        callback);
     OP_REQUIRES_OK_ASYNC(ctx, ctx->input("values", &values), callback);
     barrier->TryInsertMany<T>(*keys, component_index_, *values, ctx, callback);
   }
diff --git a/tensorflow/python/kernel_tests/data_structures/barrier_ops_test.py b/tensorflow/python/kernel_tests/data_structures/barrier_ops_test.py
index d95a54f6f54..3e09cd604b5 100644
--- a/tensorflow/python/kernel_tests/data_structures/barrier_ops_test.py
+++ b/tensorflow/python/kernel_tests/data_structures/barrier_ops_test.py
@@ -23,6 +23,7 @@ from tensorflow.python.framework import errors_impl
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor
 from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import data_flow_ops
 from tensorflow.python.ops import gen_data_flow_ops
 from tensorflow.python.platform import test
@@ -65,6 +66,43 @@ class BarrierTest(test.TestCase):
       attr { key: 'shared_name' value: { s: 'B' } }
       """, b.barrier_ref.op.node_def)
 
+  @test_util.run_deprecated_v1
+  def testInsertManyShapeMismatch(self):
+    with self.cached_session() as sess:
+      b = data_flow_ops.Barrier(
+          (dtypes.float32,), shapes=((),), shared_name="B")
+      sess.run(b.barrier_ref)
+      keys = array_ops.placeholder_with_default(["k1", "k2"], shape=None)
+      values = array_ops.placeholder_with_default([1.0], shape=None)
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          "keys and values must have the same batch size"):
+        sess.run(b.insert_many(0, keys, values))
+
+  @test_util.run_deprecated_v1
+  def testInsertManyNegativeComponentIndex(self):
+    with self.cached_session() as sess:
+      b = data_flow_ops.Barrier(
+          (dtypes.float32,), shapes=((),), shared_name="B")
+      sess.run(b.barrier_ref)
+      with self.assertRaisesRegex(
+          (ValueError, errors_impl.InvalidArgumentError),
+          "component_index must be >= 0"):
+        sess.run(b.insert_many(-1, ["k1"], [1.0]))
+
+  @test_util.run_deprecated_v1
+  def testInsertManyNonVectorKeys(self):
+    with self.cached_session() as sess:
+      b = data_flow_ops.Barrier(
+          (dtypes.float32,), shapes=((),), shared_name="B")
+      sess.run(b.barrier_ref)
+      keys = array_ops.placeholder_with_default([["k1"]], shape=None)
+      values = array_ops.placeholder_with_default([[1.0]], shape=None)
+      with self.assertRaisesRegex(
+          errors_impl.InvalidArgumentError,
+          "keys must be 1-D"):
+        sess.run(b.insert_many(0, keys, values))
+
   @test_util.run_deprecated_v1
   def testInsertMany(self):
     with self.cached_session():

From 0ec6cf64b111a65e52ff3709071c16a972b0af07 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:20:35 -0700
Subject: [PATCH 0520/1000] Automated Code Change

PiperOrigin-RevId: 997662459
---
 tensorflow/cc/client/client_session.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/cc/client/client_session.cc b/tensorflow/cc/client/client_session.cc
index f776fdd9612..ede2c7eca1f 100644
--- a/tensorflow/cc/client/client_session.cc
+++ b/tensorflow/cc/client/client_session.cc
@@ -61,7 +61,7 @@ ClientSession::ClientSession(const Scope& scope,
 // Define destructor here so we can forward declare `Impl` in client_session.h.
 // If we define a dtor in the header file or use the default dtor,
 // unique_ptr<Impl> needs the complete type.
-ClientSession::~ClientSession() {}
+ClientSession::~ClientSession() = default;
 
 SessionOptions ClientSession::Impl::MakeDefaultSessionOptions(
     const std::string& target) {

From 9134da4b8008c3dcea74bae009236fa9625f9a29 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:21:40 -0700
Subject: [PATCH 0521/1000] Automated Code Change

PiperOrigin-RevId: 997662733
---
 tensorflow/core/grappler/clusters/cluster.cc         | 2 +-
 tensorflow/core/grappler/clusters/virtual_cluster.cc | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/grappler/clusters/cluster.cc b/tensorflow/core/grappler/clusters/cluster.cc
index 18db67281df..71ce09c385c 100644
--- a/tensorflow/core/grappler/clusters/cluster.cc
+++ b/tensorflow/core/grappler/clusters/cluster.cc
@@ -29,7 +29,7 @@ Cluster::Cluster(int timeout_s) : timeout_s_(timeout_s) {
   DisableDetailedStats(false);
 }
 
-Cluster::~Cluster() {}
+Cluster::~Cluster() = default;
 
 void Cluster::AllowSoftPlacement(bool soft_placement_state) {
   options_.config.set_allow_soft_placement(soft_placement_state);
diff --git a/tensorflow/core/grappler/clusters/virtual_cluster.cc b/tensorflow/core/grappler/clusters/virtual_cluster.cc
index 2eba43f4ae2..3f18abb81ec 100644
--- a/tensorflow/core/grappler/clusters/virtual_cluster.cc
+++ b/tensorflow/core/grappler/clusters/virtual_cluster.cc
@@ -67,7 +67,7 @@ VirtualCluster::VirtualCluster(const DeviceSet* device_set)
   }
 }
 
-VirtualCluster::~VirtualCluster() {}
+VirtualCluster::~VirtualCluster() = default;
 
 absl::Status VirtualCluster::Provision() { return absl::OkStatus(); }
 

From 8df3f7381c23f5d2acd1f92b6b70869df35f04ca Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:28:57 -0700
Subject: [PATCH 0522/1000] Automated Code Change

PiperOrigin-RevId: 997664751
---
 tensorflow/cc/saved_model/bundle_v2_test.cc               | 2 +-
 tensorflow/cc/saved_model/reader_test.cc                  | 2 +-
 tensorflow/cc/saved_model/saved_model_bundle_lite_test.cc | 2 +-
 tensorflow/cc/saved_model/saved_model_bundle_test.cc      | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/cc/saved_model/bundle_v2_test.cc b/tensorflow/cc/saved_model/bundle_v2_test.cc
index 138028257b9..170facf393c 100644
--- a/tensorflow/cc/saved_model/bundle_v2_test.cc
+++ b/tensorflow/cc/saved_model/bundle_v2_test.cc
@@ -41,7 +41,7 @@ constexpr char kTestData[] = "cc/saved_model/testdata";
 
 class BundleV2Test : public ::testing::Test {
  protected:
-  BundleV2Test() {}
+  BundleV2Test() = default;
 
   void RestoreVarsAndVerify(SavedModelV2Bundle* bundle,
                             std::vector<std::string> expected_names) {
diff --git a/tensorflow/cc/saved_model/reader_test.cc b/tensorflow/cc/saved_model/reader_test.cc
index c8e153bed36..384351083ee 100644
--- a/tensorflow/cc/saved_model/reader_test.cc
+++ b/tensorflow/cc/saved_model/reader_test.cc
@@ -51,7 +51,7 @@ string NonChunkedSavedModel() {
 
 class ReaderTest : public ::testing::Test {
  protected:
-  ReaderTest() {}
+  ReaderTest() = default;
 
   void CheckMetaGraphDef(const MetaGraphDef& meta_graph_def) {
     const auto& tags = meta_graph_def.meta_info_def().tags();
diff --git a/tensorflow/cc/saved_model/saved_model_bundle_lite_test.cc b/tensorflow/cc/saved_model/saved_model_bundle_lite_test.cc
index 1d7192381c1..78b3bf58e92 100644
--- a/tensorflow/cc/saved_model/saved_model_bundle_lite_test.cc
+++ b/tensorflow/cc/saved_model/saved_model_bundle_lite_test.cc
@@ -42,7 +42,7 @@ constexpr char kTestDataInitOpV2[] =
 
 class LoaderTest : public ::testing::Test {
  protected:
-  LoaderTest() {}
+  LoaderTest() = default;
 
   string MakeSerializedExample(float x) {
     tensorflow::Example example;
diff --git a/tensorflow/cc/saved_model/saved_model_bundle_test.cc b/tensorflow/cc/saved_model/saved_model_bundle_test.cc
index 674d11fb194..97f2b5ce799 100644
--- a/tensorflow/cc/saved_model/saved_model_bundle_test.cc
+++ b/tensorflow/cc/saved_model/saved_model_bundle_test.cc
@@ -57,7 +57,7 @@ constexpr char kV2ModuleSavedModelChecksum[] = "15788619162413586750";
 
 class LoaderTest : public ::testing::Test {
  protected:
-  LoaderTest() {}
+  LoaderTest() = default;
 
   string MakeSerializedExample(float x) {
     tensorflow::Example example;

From 1d650be3a2641de96ad7d111f574a0b02db6f2ef Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:29:19 -0700
Subject: [PATCH 0523/1000] Automated Code Change

PiperOrigin-RevId: 997664841
---
 tensorflow/lite/delegates/telemetry_test.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/lite/delegates/telemetry_test.cc b/tensorflow/lite/delegates/telemetry_test.cc
index 72478f6a74d..44619efad08 100644
--- a/tensorflow/lite/delegates/telemetry_test.cc
+++ b/tensorflow/lite/delegates/telemetry_test.cc
@@ -49,7 +49,7 @@ TEST(TelemetryTest, StatusConversion) {
 // Dummy profiler to test delegate reporting.
 class DelegateProfiler : public Profiler {
  public:
-  DelegateProfiler() {}
+  DelegateProfiler() = default;
   ~DelegateProfiler() override = default;
 
   uint32_t BeginEvent(const char* tag, EventType event_type,

From 93b19f2cc821696a4fec74e618a751a977e07f10 Mon Sep 17 00:00:00 2001
From: bodapatisaikrishna <saikrishnabodapati@gmail.com>
Date: Mon, 12 Oct 2026 17:22:17 +0530
Subject: [PATCH 0524/1000] Safely handle uninitialized resource variables and
 add tests

Unconditionally validate resource variable dtype in training op helpers
and return early when tensors are uninitialized so downstream ops can
gracefully report errors without crashing.

Signed-off-by: bodapatisaikrishna <saikrishnabodapati@gmail.com>
---
 tensorflow/core/kernels/training_op_helpers.h | 12 +++-
 .../python/training/training_ops_test.py      | 71 +++++++++++++++++++
 2 files changed, 80 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/kernels/training_op_helpers.h b/tensorflow/core/kernels/training_op_helpers.h
index 8f7a2fffef9..11981f7aba8 100644
--- a/tensorflow/core/kernels/training_op_helpers.h
+++ b/tensorflow/core/kernels/training_op_helpers.h
@@ -76,6 +76,10 @@ absl::Status EnsureSparseVariableAccess(OpKernelContext* ctx, Var* var) {
         " (resource variable dtype)"));
   }
 
+  if (!var->tensor()->IsInitialized()) {
+    return absl::OkStatus();
+  }
+
   // Once copy-on-read mode is True the refcount is guaranteed to be 1. This can
   // also happen if there are no concurrent reads of the variable and
   // copy-on-read mode is false.
@@ -242,11 +246,14 @@ void MaybeForwardRefInputToRefOutput(OpKernelContext* ctx, int input,
 template <typename Device, typename T>
 absl::Status PrepareToUpdateVariable(OpKernelContext* ctx, Tensor* tensor,
                                      bool copy_on_read_mode) {
-  if (tensor->IsInitialized() && tensor->dtype() != DataTypeToEnum<T>::v()) {
+  if (tensor->dtype() != DataTypeToEnum<T>::v()) {
     return absl::InvalidArgumentError(absl::StrCat(
         "dtype mismatch: expected ", DataTypeString(DataTypeToEnum<T>::v()),
         " but got ", DataTypeString(tensor->dtype())));
   }
+  if (!tensor->IsInitialized()) {
+    return absl::OkStatus();
+  }
   if (copy_on_read_mode || !tensor->RefCountIsOne()) {
     // Tensor's buffer is in use by some read, so we need to copy before
     // updating.
@@ -298,8 +305,7 @@ absl::Status GetInputTensorFromVariable(OpKernelContext* ctx, int input,
     ResourceHandle handle;
     TF_RETURN_IF_ERROR(HandleFromInput(ctx, input, &handle));
     TF_RETURN_IF_ERROR(LookupResource(ctx, handle, &var));
-    if ((var->is_initialized || var->tensor()->IsInitialized()) &&
-        var->tensor()->dtype() != DataTypeToEnum<T>::v()) {
+    if (var->tensor()->dtype() != DataTypeToEnum<T>::v()) {
       return absl::InvalidArgumentError(absl::StrCat(
           "dtype mismatch: expected ", DataTypeString(DataTypeToEnum<T>::v()),
           " but got ", DataTypeString(var->tensor()->dtype()),
diff --git a/tensorflow/python/training/training_ops_test.py b/tensorflow/python/training/training_ops_test.py
index e2aad463ad1..83e56c512dc 100644
--- a/tensorflow/python/training/training_ops_test.py
+++ b/tensorflow/python/training/training_ops_test.py
@@ -748,6 +748,77 @@ class TrainingOpsTest(TensorFlowTestCase):
           )
       )
 
+  @test_util.run_v2_only
+  def testResourceVariableUninitialized(self):
+    from tensorflow.python.eager import context
+    from tensorflow.python.ops import gen_resource_variable_ops
+
+    shared_name = (
+        context.anonymous_name() if context.executing_eagerly() else ""
+    )
+    var_uninit_f32 = gen_resource_variable_ops.var_handle_op(
+        dtype=dtypes.float32, shape=[1, 2], shared_name=shared_name
+    )
+    var_uninit_f16 = gen_resource_variable_ops.var_handle_op(
+        dtype=dtypes.float16, shape=[1, 2], shared_name=shared_name
+    )
+    accum_uninit_f16 = gen_resource_variable_ops.var_handle_op(
+        dtype=dtypes.float16, shape=[1, 2], shared_name=shared_name
+    )
+    lr = constant_op.constant(0.1, dtype=dtypes.float16)
+    grad = constant_op.constant([[0.1, 0.2]], dtype=dtypes.float16)
+    indices = constant_op.constant([0], dtype=dtypes.int32)
+
+    # Mismatched dtype should raise InvalidArgumentError
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError, "dtype mismatch"
+    ):
+      self.evaluate(
+          gen_training_ops.resource_apply_gradient_descent(
+              var=var_uninit_f32,
+              alpha=lr,
+              delta=grad,
+          )
+      )
+
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError, "dtype mismatch"
+    ):
+      self.evaluate(
+          gen_training_ops.resource_sparse_apply_adagrad(
+              var=var_uninit_f32,
+              accum=accum_uninit_f16,
+              lr=lr,
+              grad=grad,
+              indices=indices,
+          )
+      )
+
+    # Matched dtype should raise FailedPreconditionError
+    with self.assertRaisesRegex(
+        errors.FailedPreconditionError, "uninitialized"
+    ):
+      self.evaluate(
+          gen_training_ops.resource_apply_gradient_descent(
+              var=var_uninit_f16,
+              alpha=lr,
+              delta=grad,
+          )
+      )
+
+    with self.assertRaisesRegex(
+        errors.FailedPreconditionError, "uninitialized"
+    ):
+      self.evaluate(
+          gen_training_ops.resource_sparse_apply_adagrad(
+              var=var_uninit_f16,
+              accum=accum_uninit_f16,
+              lr=lr,
+              grad=grad,
+              indices=indices,
+          )
+      )
+
 
 if __name__ == '__main__':
   googletest.main()

From 944b03b256609a9b7709e10e00ff379e7741d5a7 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 17:39:39 +0530
Subject: [PATCH 0525/1000] Update pywrap_tensorflow.py

---
 tensorflow/python/pywrap_tensorflow.py | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/tensorflow/python/pywrap_tensorflow.py b/tensorflow/python/pywrap_tensorflow.py
index 9c6b946c923..54baf9ea1b4 100644
--- a/tensorflow/python/pywrap_tensorflow.py
+++ b/tensorflow/python/pywrap_tensorflow.py
@@ -71,16 +71,15 @@ try:
 
   # pylint: disable=wildcard-import,g-import-not-at-top,line-too-long,undefined-variable
   try:
-    try:
-      from tensorflow.python._pywrap_tensorflow_internal import *
-    # This try catch logic is because there is no bazel equivalent for py_extension.
-    # Externally in opensource we must enable exceptions to load the shared object
-    # by exposing the PyInit symbols with pybind. This error will only be
-    # caught internally or if someone changes the name of the target _pywrap_tensorflow_internal.
+    from tensorflow.python._pywrap_tensorflow_internal import *
+  # This try catch logic is because there is no bazel equivalent for py_extension.
+  # Externally in opensource we must enable exceptions to load the shared object
+  # by exposing the PyInit symbols with pybind. This error will only be
+  # caught internally or if someone changes the name of the target _pywrap_tensorflow_internal.
 
-    # This logic is used in other internal projects using py_extension.
-    except ModuleNotFoundError:
-      pass
+  # This logic is used in other internal projects using py_extension.
+  except ModuleNotFoundError:
+    pass
   finally:
     # Always restore the dlopen flags, even if the import above raised an
     # ImportError that is not a ModuleNotFoundError (e.g. a genuine load

From efd12f18f3dfbb6cb9caf88b51173ee6394a8243 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:30:32 -0700
Subject: [PATCH 0526/1000] Automated Code Change

PiperOrigin-RevId: 997665100
---
 tensorflow/core/kernels/list_kernels.h           | 6 +++---
 tensorflow/core/kernels/matmul_op_fused.cc       | 6 +++---
 tensorflow/core/kernels/matmul_op_impl.h         | 4 ++--
 tensorflow/core/kernels/maxpooling_op.cc         | 4 ++--
 tensorflow/core/kernels/multinomial_op.cc        | 4 ++--
 tensorflow/core/kernels/multinomial_op_gpu.cu.cc | 2 +-
 tensorflow/core/kernels/pack_op.cc               | 4 ++--
 tensorflow/core/kernels/pooling_ops_common.h     | 9 ++++-----
 tensorflow/core/kernels/quantization_utils.h     | 4 ++--
 tensorflow/core/kernels/reduction_ops.h          | 6 +++---
 10 files changed, 24 insertions(+), 25 deletions(-)

diff --git a/tensorflow/core/kernels/list_kernels.h b/tensorflow/core/kernels/list_kernels.h
index 78517dfb6ff..95c6494b0b7 100644
--- a/tensorflow/core/kernels/list_kernels.h
+++ b/tensorflow/core/kernels/list_kernels.h
@@ -257,7 +257,7 @@ class TensorListStack : public OpKernel {
     auto output_flat = output->shaped<T, 2>({1, output->NumElements()});
 
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-    if (std::is_same<Device, Eigen::GpuDevice>::value) {
+    if (std::is_same_v<Device, Eigen::GpuDevice>) {
       ConcatGPU<T>(c, inputs_flat, output, &output_flat);
       return;
     }
@@ -573,7 +573,7 @@ class TensorListConcat : public OpKernel {
     auto output_flat = output->shaped<T, 2>({1, output->NumElements()});
 
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-    if (std::is_same<Device, Eigen::GpuDevice>::value) {
+    if (std::is_same_v<Device, Eigen::GpuDevice>) {
       ConcatGPU<T>(c, inputs_flat, output, &output_flat);
       return;
     }
@@ -764,7 +764,7 @@ class TensorListGather : public OpKernel {
     auto output_flat = output->shaped<T, 2>({1, output->NumElements()});
 
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-    if (std::is_same<Device, Eigen::GpuDevice>::value) {
+    if (std::is_same_v<Device, Eigen::GpuDevice>) {
       ConcatGPU<T>(c, inputs_flat, output, &output_flat);
       return;
     }
diff --git a/tensorflow/core/kernels/matmul_op_fused.cc b/tensorflow/core/kernels/matmul_op_fused.cc
index 85ca1a00cf6..1ff8fded4ea 100644
--- a/tensorflow/core/kernels/matmul_op_fused.cc
+++ b/tensorflow/core/kernels/matmul_op_fused.cc
@@ -635,7 +635,7 @@ class FusedMatMulOp : public OpKernel {
     std::vector<FusedComputationPattern> patterns;
 
     using FCT = FusedComputationType;
-    if (std::is_same<Device, CPUDevice>::value) {
+    if (std::is_same_v<Device, CPUDevice>) {
       patterns = {
           {FCT::kBiasAdd, {"BiasAdd"}},
           {FCT::kBiasAddWithRelu, {"BiasAdd", "Relu"}},
@@ -645,7 +645,7 @@ class FusedMatMulOp : public OpKernel {
           {FCT::kBiasAddWithElu, {"BiasAdd", "Elu"}},
           {FCT::kBiasAddWithLeakyRelu, {"BiasAdd", "LeakyRelu"}},
       };
-    } else if (std::is_same<Device, GPUDevice>::value) {
+    } else if (std::is_same_v<Device, GPUDevice>) {
       patterns = {
           {FCT::kBiasAdd, {"BiasAdd"}},
           {FCT::kBiasAddWithRelu, {"BiasAdd", "Relu"}},
@@ -658,7 +658,7 @@ class FusedMatMulOp : public OpKernel {
     OP_REQUIRES_OK(context, InitializeFusedComputation(
                                 context, "MatMul", patterns,
                                 &fused_computation_, &fused_computation_args_));
-    if (std::is_same<Device, GPUDevice>::value &&
+    if (std::is_same_v<Device, GPUDevice> &&
         (fused_computation_ == FCT::kBiasAddWithGeluExact ||
          fused_computation_ == FCT::kBiasAddWithTanh ||
          fused_computation_ == FCT::kBiasAddWithSigmoid)) {
diff --git a/tensorflow/core/kernels/matmul_op_impl.h b/tensorflow/core/kernels/matmul_op_impl.h
index 4aaf8ef3df2..58740a6c6f6 100644
--- a/tensorflow/core/kernels/matmul_op_impl.h
+++ b/tensorflow/core/kernels/matmul_op_impl.h
@@ -1002,10 +1002,10 @@ class BaseBatchMatMulOp : public OpKernel {
     } else {
       // Cast tensor to desired type to reuse Eigen.
       // TODO(b/178749687): remove this cast if Eigen supports this natively.
-      if constexpr (!std::is_same<Ta, Tout>::value) {
+      if constexpr (!std::is_same_v<Ta, Tout>) {
         in0_reshaped = CastTensor<Ta, Tout>(in0_reshaped);
       }
-      if constexpr (!std::is_same<Tb, Tout>::value) {
+      if constexpr (!std::is_same_v<Tb, Tout>) {
         in1_reshaped = CastTensor<Tb, Tout>(in1_reshaped);
       }
       LaunchBatchMatMul<Device, Tout>::Launch(
diff --git a/tensorflow/core/kernels/maxpooling_op.cc b/tensorflow/core/kernels/maxpooling_op.cc
index e5f94a0ff9b..ca719d9abd0 100644
--- a/tensorflow/core/kernels/maxpooling_op.cc
+++ b/tensorflow/core/kernels/maxpooling_op.cc
@@ -1122,7 +1122,7 @@ class MaxPoolingGradWithArgmaxOp : public OpKernel {
   explicit MaxPoolingGradWithArgmaxOp(OpKernelConstruction* context)
       : OpKernel(context) {
     std::string data_format_str;
-    if (std::is_same<Device, GPUDevice>::value) {
+    if (std::is_same_v<Device, GPUDevice>) {
       OP_REQUIRES(context, !tensorflow::OpDeterminismRequired(),
                   absl::UnimplementedError("Determinism is not yet supported "
                                            "for MaxPoolGradWithArgmax."));
@@ -1339,7 +1339,7 @@ class MaxPoolingNoMaskOp<GPUDevice, T> : public OpKernel {
     }
 
     // Assuming qint8 <--> NCHW_VECT_C (int8x4) here.
-    constexpr bool is_int8x4 = std::is_same<T, qint8>::value;
+    constexpr bool is_int8x4 = std::is_same_v<T, qint8>;
     OP_REQUIRES(context, (is_int8x4 == (data_format_ == FORMAT_NCHW_VECT_C)),
                 absl::InvalidArgumentError(
                     "qint8 should be used with data_format NCHW_VECT_C."));
diff --git a/tensorflow/core/kernels/multinomial_op.cc b/tensorflow/core/kernels/multinomial_op.cc
index 4b76b3b4223..2bc19225a66 100644
--- a/tensorflow/core/kernels/multinomial_op.cc
+++ b/tensorflow/core/kernels/multinomial_op.cc
@@ -188,7 +188,7 @@ class MultinomialOp : public OpKernel {
     // Execute kernel only for nonempty output; otherwise Eigen crashes on GPU.
     if (samples_t->NumElements() > 0) {
       Tensor noises, scores, scratch;  // Scratch space only used for GPU.
-      if (std::is_same<Device, GPUDevice>::value) {
+      if (std::is_same_v<Device, GPUDevice>) {
         OP_REQUIRES_OK(
             ctx,
             ctx->allocate_temp(
@@ -207,7 +207,7 @@ class MultinomialOp : public OpKernel {
 
       int num_samples_ceil_4 = (num_samples + 3) / 4 * 4;
       // CPU generates doubles = 2 samples per number.
-      if (std::is_same<Device, CPUDevice>::value) num_samples_ceil_4 *= 2;
+      if (std::is_same_v<Device, CPUDevice>) num_samples_ceil_4 *= 2;
       auto rng =
           generator->ReserveRandomOutputs(batch_size * num_samples_ceil_4, 256);
       functor::MultinomialFunctor<Device, T, OutputType>()(
diff --git a/tensorflow/core/kernels/multinomial_op_gpu.cu.cc b/tensorflow/core/kernels/multinomial_op_gpu.cu.cc
index 26ad686e7d2..0533e603a41 100644
--- a/tensorflow/core/kernels/multinomial_op_gpu.cu.cc
+++ b/tensorflow/core/kernels/multinomial_op_gpu.cu.cc
@@ -47,7 +47,7 @@ __global__ void MultinomialKernel(int32_t nthreads, const int32_t num_classes,
   GPU_1D_KERNEL_LOOP(index, nthreads) {
     const int maxima_idx = index / num_classes;
     if (ldg(maxima + maxima_idx) == ldg(scores + index)) {
-      using UnsignedOutputType = typename std::make_unsigned<OutputType>::type;
+      using UnsignedOutputType = std::make_unsigned_t<OutputType>;
       GpuAtomicMax(reinterpret_cast<UnsignedOutputType*>(output + maxima_idx),
                    static_cast<UnsignedOutputType>(index % num_classes));
     }
diff --git a/tensorflow/core/kernels/pack_op.cc b/tensorflow/core/kernels/pack_op.cc
index a43723bd438..c5319e213e1 100644
--- a/tensorflow/core/kernels/pack_op.cc
+++ b/tensorflow/core/kernels/pack_op.cc
@@ -108,7 +108,7 @@ class PackOp : public OpKernel {
     // required: Eigen device() assignments must be compiled with
     // EIGEN_USE_GPU defined (i.e. by NVCC/hipcc), otherwise Eigen fires a
     // static assertion even inside discarded if-constexpr branches.
-    if (std::is_same<Device, GPUDevice>::value && first_input.dims() == 0) {
+    if (std::is_same_v<Device, GPUDevice> && first_input.dims() == 0) {
       PackScalarsOnGPU<T>(c, num, output);
       return;
     }
@@ -146,7 +146,7 @@ class PackOp : public OpKernel {
     }
     if (output_size > 0) {
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-      if (std::is_same<Device, GPUDevice>::value) {
+      if (std::is_same_v<Device, GPUDevice>) {
         ConcatGPU<T>(c, inputs_flat, output, &output_flat);
         return;
       }
diff --git a/tensorflow/core/kernels/pooling_ops_common.h b/tensorflow/core/kernels/pooling_ops_common.h
index 835bd654f61..1c1f2884f97 100644
--- a/tensorflow/core/kernels/pooling_ops_common.h
+++ b/tensorflow/core/kernels/pooling_ops_common.h
@@ -164,8 +164,7 @@ class MaxPoolingOp : public OpKernel {
     } else {
       // MaxPoolingOp is only called on the GPU when the eigen_tensor label
       // is used. In this case, explicit padding is not supported
-      if (std::is_same<Device, GPUDevice>::value &&
-          padding_ == Padding::EXPLICIT) {
+      if (std::is_same_v<Device, GPUDevice> && padding_ == Padding::EXPLICIT) {
         context->SetStatus(absl::UnimplementedError(
             "MaxPoolingOp does not support explicit padding."));
         return;
@@ -202,7 +201,7 @@ class MaxPoolingOp : public OpKernel {
     // Spatial MaxPooling implementation.
     //
     // TODO(vrv): Remove this once we no longer need it.
-    if (std::is_same<Device, GPUDevice>::value) {
+    if (std::is_same_v<Device, GPUDevice>) {
       Eigen::PaddingType pt = BrainPadding2EigenPadding(padding);
       functor::SpatialMaxPooling<Device, T>()(
           context->eigen_device<Device>(), output->tensor<T, 4>(),
@@ -469,9 +468,9 @@ class MaxPoolingV2Op : public OpKernel {
     //
     // TODO(vrv): Remove this once we no longer need it.
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
-    if (std::is_same<Device, GPUDevice>::value) {
+    if (std::is_same_v<Device, GPUDevice>) {
       Eigen::PaddingType pt = BrainPadding2EigenPadding(padding);
-      if (std::is_same<T, qint8>::value) {
+      if (std::is_same_v<T, qint8>) {
         LaunchMaxPoolingNoMask_NCHW_VECT_C<GPUDevice>::launch(
             context, params, tensor_in, output);
       } else {
diff --git a/tensorflow/core/kernels/quantization_utils.h b/tensorflow/core/kernels/quantization_utils.h
index f10b5e823d4..6881c21ce40 100644
--- a/tensorflow/core/kernels/quantization_utils.h
+++ b/tensorflow/core/kernels/quantization_utils.h
@@ -71,7 +71,7 @@ inline int64_t FloatToQuantizedUnclamped<float>(float input, float range_min,
 // any over or underflows.
 template <class T>
 T FloatToQuantized(float input, float range_min, float range_max) {
-  if (std::is_same<T, float>::value) {
+  if (std::is_same_v<T, float>) {
     // Specialization for float. This is used in reference implementation
     // for float which is useful to compare performance between float
     // and quantized type.
@@ -89,7 +89,7 @@ T FloatToQuantized(float input, float range_min, float range_max) {
 
 template <class T>
 float QuantizedToFloat(T input, float range_min, float range_max) {
-  if (std::is_same<T, float>::value) {
+  if (std::is_same_v<T, float>) {
     // Specialization for float. This is used in reference implementation
     // for float which is useful to compare performance between float
     // and quantized type.
diff --git a/tensorflow/core/kernels/reduction_ops.h b/tensorflow/core/kernels/reduction_ops.h
index 34f559704ed..0ceefaecf9e 100644
--- a/tensorflow/core/kernels/reduction_ops.h
+++ b/tensorflow/core/kernels/reduction_ops.h
@@ -87,7 +87,7 @@ struct ReduceEigenImpl<Device, OUT_T, IN_T, ReductionAxes,
   void operator()(const Device& d, OUT_T out, IN_T in,
                   const ReductionAxes& reduction_axes,
                   const functor::MeanReducer<Scalar>& reducer) {
-    static_assert(std::is_same<Scalar, typename OUT_T::Scalar>::value, "");
+    static_assert(std::is_same_v<Scalar, typename OUT_T::Scalar>, "");
     Eigen::internal::SumReducer<Scalar> sum_reducer;
     out.device(d) = in.reduce(reduction_axes, sum_reducer) /
                     static_cast<Scalar>(in.size() / out.size());
@@ -132,7 +132,7 @@ struct ReduceEigenImpl<Device, OUT_T, IN_T, ReductionAxes,
   void operator()(const Device& d, OUT_T out, IN_T in,
                   const ReductionAxes& reduction_axes,
                   const functor::EuclideanNormReducer<Scalar>& reducer) {
-    static_assert(std::is_same<Scalar, typename OUT_T::Scalar>::value, "");
+    static_assert(std::is_same_v<Scalar, typename OUT_T::Scalar>, "");
     Eigen::internal::SumReducer<Scalar> sum_reducer;
     out.device(d) =
         (in * in.conjugate()).reduce(reduction_axes, sum_reducer).sqrt();
@@ -146,7 +146,7 @@ struct ReduceEigenImpl<Device, OUT_T, IN_T, ReductionAxes,
   void operator()(const Device& d, OUT_T out, IN_T in,
                   const ReductionAxes& reduction_axes,
                   const functor::EuclideanNormReducer<bfloat16>& reducer) {
-    static_assert(std::is_same<bfloat16, typename OUT_T::Scalar>::value, "");
+    static_assert(std::is_same_v<bfloat16, typename OUT_T::Scalar>, "");
     Eigen::internal::SumReducer<float> sum_reducer;
     auto in_as_float = in.template cast<float>();
     out.device(d) = (in_as_float * in_as_float.conjugate())

From aacabb8b405189eb0b8edbf680ad3804c9996c7d Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 17:47:21 +0530
Subject: [PATCH 0527/1000] Remove exception handling in async wait

Removed exception handling for async wait in accelerator utility.
---
 tensorflow/dtensor/python/accelerator_util.py | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/tensorflow/dtensor/python/accelerator_util.py b/tensorflow/dtensor/python/accelerator_util.py
index a279039d232..24c23121f6a 100644
--- a/tensorflow/dtensor/python/accelerator_util.py
+++ b/tensorflow/dtensor/python/accelerator_util.py
@@ -312,12 +312,6 @@ def shutdown_accelerator_system() -> None:
 
   try:
     context.async_wait()
-  except Exception:
-    logging.exception(
-        "Error waiting for async operations to finish before shutting "
-        "down the accelerator system."
-    )
-    raise
   finally:
     # Cleanup always runs, even if `async_wait()` raised, so the system
     # is never left stuck in an "initialized" state.

From fbf340ce2f73e3b70b10482fdda04d9801a77bb9 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:32:18 -0700
Subject: [PATCH 0528/1000] Automated Code Change

PiperOrigin-RevId: 997665559
---
 tensorflow/lite/tools/evaluation/stages/utils/image_metrics.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/lite/tools/evaluation/stages/utils/image_metrics.h b/tensorflow/lite/tools/evaluation/stages/utils/image_metrics.h
index 2a9aae287fc..973f8ef5204 100644
--- a/tensorflow/lite/tools/evaluation/stages/utils/image_metrics.h
+++ b/tensorflow/lite/tools/evaluation/stages/utils/image_metrics.h
@@ -31,7 +31,7 @@ struct Box2D {
       min = x;
       max = y;
     }
-    Interval() {}
+    Interval() = default;
   };
 
   Interval x;
@@ -69,7 +69,7 @@ struct Detection {
   Box2D box;
   IgnoreType ignore = IgnoreType::kDontIgnore;
 
-  Detection() {}
+  Detection() = default;
   Detection(bool d, int64_t id, float s, Box2D b)
       : difficult(d), imgid(id), score(s), box(b) {}
   Detection(bool d, int64_t id, float s, Box2D b, IgnoreType i)

From 0efe672d5126e0cec481b1916e8537f2812a446b Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 17:50:48 +0530
Subject: [PATCH 0529/1000] Implement async wait error handling test

Add test for shutdown accelerator system with async wait error.
---
 tensorflow/dtensor/python/tests/mesh_util_test.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/tensorflow/dtensor/python/tests/mesh_util_test.py b/tensorflow/dtensor/python/tests/mesh_util_test.py
index b8bc9960915..40d8cbc2b14 100644
--- a/tensorflow/dtensor/python/tests/mesh_util_test.py
+++ b/tensorflow/dtensor/python/tests/mesh_util_test.py
@@ -333,6 +333,18 @@ class InitializedMeshUtilTest(test_util.DTensorBaseTest):
     with self.assertRaisesRegex(ValueError, 'not initialized'):
       accelerator_util.shutdown_accelerator_system()
 
+  def test_shutdown_accelerator_system_async_wait_error(self):
+    accelerator_util.shutdown_accelerator_system()
+    accelerator_util.initialize_accelerator_system(
+        'CPU', num_logical_cpu_devices=1)
+    self.assertTrue(accelerator_util.is_initialized())
+    with mock.patch.object(context, 'async_wait',
+                           side_effect=RuntimeError('async fail')):
+      with self.assertRaisesRegex(RuntimeError, 'async fail'):
+        accelerator_util.shutdown_accelerator_system()
+    # Ensure cleanup still ran in the finally block.
+    self.assertFalse(accelerator_util.is_initialized())
+
   def test_distributed_tpu_mesh_creation(self):
     self.skipForDeviceType(['CPU', 'GPU'], reason='Test is intended for TPUs')
     self.skipForDeviceType(['TPU'],

From 958362fc0cf6767694df2338ff4bbd52df351e55 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 17:52:57 +0530
Subject: [PATCH 0530/1000] Update config.py

---
 tensorflow/dtensor/python/config.py | 285 +++++++++-------------------
 1 file changed, 86 insertions(+), 199 deletions(-)

diff --git a/tensorflow/dtensor/python/config.py b/tensorflow/dtensor/python/config.py
index cd60ca6db9d..217a798518d 100644
--- a/tensorflow/dtensor/python/config.py
+++ b/tensorflow/dtensor/python/config.py
@@ -12,211 +12,98 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""DTensor Configuration API."""
+"""Tests for the open source DTensor Python API."""
 
 import os
-from typing import List, Optional, Union
 
+# pylint: disable=g-direct-tensorflow-import
+from tensorflow.dtensor.python import config
+from tensorflow.dtensor.python.tests import test_util
 from tensorflow.python.eager import context
-from tensorflow.python.framework import config as tf_config
 from tensorflow.python.framework import device as tf_device
-from tensorflow.python.util.tf_export import tf_export
-
-_DT_CLIENT_ID = "DTENSOR_CLIENT_ID"
-# DTENSOR_NUM_CLIENTS is removed, but some DTensor users still use this symbol.
-_DT_NUM_CLIENTS = "DTENSOR_NUM_CLIENTS"
-_DT_JOB_NAME = "DTENSOR_JOB_NAME"
-_DT_JOBS = "DTENSOR_JOBS"
-_DT_HEARTBEAT_ENABLED = "DTENSOR_ENABLE_HEARTBEAT"
+from tensorflow.python.platform import test as tf_test
+# pylint: enable=g-direct-tensorflow-import
 
 
-# All functions in this file can be used before calling
-# `tf.experimental.dtensor.initialize_accelerator_system`.
+class ConfigTest(tf_test.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    test_util.reset_logical_devices('CPU', 2)
+    if test_util.is_gpu_present():
+      test_util.reset_logical_devices('GPU', 2)
+
+  def tearDown(self):
+    os.environ.pop(config._DT_JOBS, [])
+    super().tearDown()
+
+  def test_env_vars(self):
+    self.assertEqual(config.client_id(), 0)
+    self.assertEqual(config.num_clients(), 1)
+    self.assertEqual(config.job_name(), 'localhost')
+    self.assertEqual(config.full_job_name(), 'localhost/replica:0/task:0')
+    self.assertEqual(config.jobs(), [])
+
+  def test_jobs_empty_and_whitespace(self):
+    for empty_val in ('', '   '):
+      os.environ[config._DT_JOBS] = empty_val
+      self.assertEqual(config.jobs(), [])
+
+  def test_list_devices(self):
+    device_type = config.preferred_device_type()
+
+    local_devices = [
+        tf_device.DeviceSpec.from_string(
+            f'/job:localhost/replica:0/task:0/device:{device_type}:0'),
+        tf_device.DeviceSpec.from_string(
+            f'/job:localhost/replica:0/task:0/device:{device_type}:1'),
+    ]
+
+    self.assertEqual(config.local_devices(device_type), local_devices)
+    self.assertEqual(config.num_local_devices(device_type), 2)
+    self.assertEqual(config.num_global_devices(device_type), 2)
+    # The eager context should not be initialized by any of the calls
+    self.assertFalse(context.context()._initialized)  # pylint: disable=protected-access
+
+  def test_sort_jobs_with_bns_names(self):
+    # bns names must be sorted in the bns order.
+    dtensor_jobs = [
+        '/bns/localhost/{task_id}'.format(task_id=i) for i in range(16)
+    ]
+    os.environ[config._DT_JOBS] = ','.join(dtensor_jobs)
+    self.assertListEqual(dtensor_jobs, config.jobs())
+
+    dtensor_jobs = [
+        '/bns/localhost/{task_id}:8888'.format(task_id=i) for i in range(16)
+    ]
+    os.environ[config._DT_JOBS] = ','.join(dtensor_jobs)
+    self.assertListEqual(dtensor_jobs, config.jobs())
+
+    dtensor_jobs = [
+        '/bns/localhost/{task_id}'.format(task_id=100 - i) for i in range(16)
+    ]
+    os.environ[config._DT_JOBS] = ','.join(dtensor_jobs)
+    with self.assertRaisesRegex(ValueError, 'Unexpected DTENSOR_JOBS'):
+      config.jobs()
+
+  def test_jobs_with_ip_port(self):
+    # The ip port format is not a bns address, and not required to sorted.
+    dtensor_jobs = ['localhost:{port}'.format(port=16 - i) for i in range(16)]
+    os.environ[config._DT_JOBS] = ','.join(dtensor_jobs)
+    self.assertListEqual(dtensor_jobs, config.jobs())
+
+  def test_heartbeat_enabled(self):
+    os.environ[config._DT_HEARTBEAT_ENABLED] = 'false'
+    self.assertFalse(config.heartbeat_enabled())
+    os.environ[config._DT_HEARTBEAT_ENABLED] = '0'
+    self.assertFalse(config.heartbeat_enabled())
+    os.environ[config._DT_HEARTBEAT_ENABLED] = 'true'
+    self.assertTrue(config.heartbeat_enabled())
+    os.environ[config._DT_HEARTBEAT_ENABLED] = '1'
+    self.assertTrue(config.heartbeat_enabled())
+    os.environ.pop(config._DT_HEARTBEAT_ENABLED, None)
+    self.assertTrue(config.heartbeat_enabled())
 
 
-# -----------------------------------------------------------------------------
-# Distributed training-related methods.
-#
-# Most users should use DTensor utility methods to create a mesh. The methods
-# here are only for advanced users who want to fully customize their meshes.
-# Note that local_devices and num_local_devices return the actual number of
-# locally attached devices. The others are set through environment variables.
-
-
-@tf_export("experimental.dtensor.local_devices", v1=[])
-def local_devices(
-    device_type: str, for_client_id: Optional[int] = None
-) -> List[tf_device.DeviceSpec]:
-  """Returns a list of device specs configured on this client."""
-  if device_type.upper() not in ["CPU", "GPU", "TPU"]:
-    raise ValueError(f"Device type {device_type} is not CPU, GPU, or TPU.")
-
-  if for_client_id is None:
-    for_client_id = client_id()
-
-  # Return fully qualified device specs, sorted by increasing device index.
-  return [
-      tf_device.DeviceSpec(  # pylint: disable=g-complex-comprehension
-          job=job_name(),
-          replica=0,  # replica is deprecated and mostly hard-coded now.
-          task=for_client_id,
-          device_type=device_type,
-          device_index=i,
-      )
-      for i in range(num_local_devices(device_type))
-  ]
-
-
-@tf_export("experimental.dtensor.num_local_devices", v1=[])
-def num_local_devices(device_type: str) -> int:
-  """Returns the number of devices of device_type configured on this client."""
-
-  # Reads from config because CPU and GPU can use logical devices.
-  if device_type.upper() in ["CPU", "GPU"]:
-    context_config = context.get_config()
-    return context_config.device_count[device_type.upper()]
-
-  return len(tf_config.list_physical_devices(device_type))
-
-
-@tf_export("experimental.dtensor.num_global_devices", v1=[])
-def num_global_devices(device_type: str) -> int:
-  """Returns the number of devices of device_type in this DTensor cluster."""
-  return num_local_devices(device_type) * num_clients()
-
-
-@tf_export("experimental.dtensor.client_id", v1=[])
-def client_id() -> int:
-  """Returns this client's ID."""
-  # If missing, assume running with a single client with client_id of 0.
-  client_id_value = int(os.environ.get(_DT_CLIENT_ID, "0"))
-  if client_id_value < 0:
-    raise ValueError(
-        f"Environment variable {_DT_CLIENT_ID} "
-        f"must be >= 0, got {client_id_value}. "
-    )
-  if client_id_value >= num_clients():
-    raise ValueError(
-        f"Environment variable {_DT_CLIENT_ID} "
-        f"must be < {num_clients()}, got {client_id_value}"
-    )
-  return client_id_value
-
-
-@tf_export("experimental.dtensor.num_clients", v1=[])
-def num_clients() -> int:
-  """Returns the number of clients in this DTensor cluster."""
-  if is_local_mode():
-    return 1
-  return len(jobs())
-
-
-@tf_export("experimental.dtensor.job_name", v1=[])
-def job_name() -> str:
-  """Returns the job name used by all clients in this DTensor cluster."""
-  # If missing, assumes the program runs locally and use localhost as job name
-  # per TensorFlow convention.
-  return os.environ.get(
-      _DT_JOB_NAME, "localhost" if num_clients() == 1 else "worker"
-  )
-
-
-@tf_export("experimental.dtensor.full_job_name", v1=[])
-def full_job_name(task_id: Optional[int] = None) -> str:
-  """Returns the fully qualified TF job name for this or another task."""
-  # If task_id is None, use this client's ID, which is equal to its task ID.
-  if task_id is None:
-    task_id = client_id()
-  # In local runs and unit tests, there should be exactly one client running
-  # on one TF task.
-  if num_clients() == 1 and task_id != 0:
-    raise ValueError(f"Unexpected task ID {task_id} in local runs")
-  return f"{job_name()}/replica:0/task:{task_id}"
-
-
-def _bns_task_id(job: str) -> Union[int, str]:
-  """Tries to extract an integer task ID from a job name.
-
-  For example, for `job` = '/.../tpu_worker/0:port_name', return 0.
-
-  Args:
-    job: A job name to extract task ID from.
-
-  Returns:
-    The task ID on success, or the original job name on failure.
-  """
-  maybe_task_id = job.rsplit("/")[-1].rsplit(":")[0]
-  try:
-    return int(maybe_task_id)
-  except ValueError:
-    return job
-
-
-@tf_export("experimental.dtensor.jobs", v1=[])
-def jobs() -> List[str]:
-  """Returns a list of job names of all clients in this DTensor cluster."""
-  d_jobs = os.environ.get(_DT_JOBS)
-  if not d_jobs or not d_jobs.strip():
-    return []
-  d_jobs_list = d_jobs.split(",")
-
-  # Validate ordering for BNS style job names.
-  # For definition of BNS, refer to https://research.google/pubs/pub43438/.
-  if any([name.startswith("/bns/") for name in d_jobs_list]):
-    if d_jobs_list != sorted(d_jobs_list, key=_bns_task_id):
-      raise ValueError(
-          f"Unexpected DTENSOR_JOBS content {d_jobs}. Sort entries "
-          "in DTENSOR_JOBS because cluster construction relies on "
-          "the order."
-      )
-
-  return d_jobs_list
-
-
-@tf_export("experimental.dtensor.heartbeat_enabled", v1=[])
-def heartbeat_enabled() -> bool:
-  """Returns true if DTensor heartbeat service is enabled."""
-  return os.environ.get(_DT_HEARTBEAT_ENABLED, "true").lower() in ("true", "1")
-
-
-def is_local_mode() -> bool:
-  """Returns true if DTensor shall run in local mode."""
-  return not jobs()
-
-
-def is_tpu_present() -> bool:
-  """Returns true if TPU devices are present."""
-  # Check if TPU is present from initialized context.
-  # TPU_SYSTEM is a device that indicates TPUs are present.
-  tpu_system_devices = tf_config.list_physical_devices("TPU_SYSTEM")
-  return bool(tpu_system_devices)
-
-
-def is_gpu_present() -> bool:
-  """Returns true if GPU devices are present."""
-  return bool(tf_config.list_physical_devices("GPU"))
-
-
-@tf_export("experimental.dtensor.preferred_device_type", v1=[])
-def preferred_device_type() -> str:
-  """Returns the preferred device type for the accelerators.
-
-  The returned device type is determined by checking the first present device
-  type from all supported device types in the order of 'TPU', 'GPU', 'CPU'.
-  """
-  if is_tpu_present():
-    return "TPU"
-  elif is_gpu_present():
-    return "GPU"
-
-  return "CPU"
-
-
-def use_multi_device_mode() -> bool:
-  """Return True if environment indicates multi-device mode is enabled."""
-  return os.environ.get("DTENSOR_ENABLE_MULTI_DEVICE_EXPANSION", "0") != "0"
-
-
-def gpu_use_nccl_communication() -> bool:
-  """Return True if environment indicates NCCL shall be used for GPU."""
-  return os.environ.get("DTENSOR_GPU_USE_NCCL_COMMUNICATION", "0") != "0"
+if __name__ == '__main__':
+  tf_test.main()

From b5155119ea7b1f868668c665ba44837be6066bc0 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 17:57:10 +0530
Subject: [PATCH 0531/1000] Clean up d_variable.py and update type checks

Removed copyright notice and unused imports. Updated type checks for callable initial values.
---
 tensorflow/dtensor/python/d_variable.py | 263 +-----------------------
 1 file changed, 2 insertions(+), 261 deletions(-)

diff --git a/tensorflow/dtensor/python/d_variable.py b/tensorflow/dtensor/python/d_variable.py
index 0c8884eb87a..c501c597fdf 100644
--- a/tensorflow/dtensor/python/d_variable.py
+++ b/tensorflow/dtensor/python/d_variable.py
@@ -13,21 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 """DTensor variable and saveable."""
-# Copyright 2022 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.
-# ==============================================================================
-"""DTensor variable and saveable."""
 
 import functools
 
@@ -178,14 +163,14 @@ class DVariable(resource_variable_ops.ResourceVariable):
 
     if callable(initial_value):
       unwrapped = initial_value
-      if issubclass(type(initial_value), functools.partial):
+      if isinstance(initial_value, functools.partial):
         unwrapped = initial_value.func  # pyrefly: ignore[missing-attribute]
 
       # If wrapped is a CheckpointInitialValueCallable, this means that
       # we are creating a Variable during a checkpoint restore.
       # Thus the restore will happen now through this callable
       # and we will create the DVariable with the restored dtensor.
-      if issubclass(type(unwrapped), trackable.CheckpointInitialValueCallable):
+      if isinstance(unwrapped, trackable.CheckpointInitialValueCallable):
         if shape is None or layout is None:
           raise ValueError('Expected shape and layout to be not None.')
 
@@ -273,247 +258,3 @@ class DVariable(resource_variable_ops.ResourceVariable):
         trackable.VARIABLE_VALUE_KEY:
             functools.partial(_DVariableSaveable, self)
     }
-import functools
-
-from tensorflow.dtensor.python import api
-from tensorflow.dtensor.python import layout as layout_lib
-from tensorflow.python.eager import context
-from tensorflow.python.eager import def_function
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import errors_impl
-from tensorflow.python.framework import ops
-from tensorflow.python.ops import math_ops
-from tensorflow.python.ops import resource_variable_ops
-from tensorflow.python.trackable import base as trackable
-from tensorflow.python.training.saving import saveable_object
-from tensorflow.python.util.tf_export import tf_export
-
-
-class DSaveSpec(saveable_object.SaveSpec):
-  """DTensor SaveSpec that additionaly captures global_shape and layout."""
-
-  def __init__(self,
-               tensor,
-               slice_spec,
-               name,
-               global_shape,
-               layout,
-               dtype=None,
-               device=None):
-    super().__init__(
-        tensor=tensor,
-        slice_spec=slice_spec,
-        name=name,
-        dtype=dtype,
-        device=device)
-    self.global_shape = global_shape
-    self.layout = layout
-
-
-class _DVariableSaveable(saveable_object.SaveableObject):
-  """Class for defining how to save/restore DTensor variable."""
-
-  def __init__(self, dvariable, name):
-    with ops.device(dvariable.device):
-      original_layout = api.fetch_layout(dvariable)
-    # Record original layout to allow restore.
-    self._original_layout = original_layout
-    self._dvariable = dvariable
-
-    def pack(tensors, layout):
-      with ops.device(dvariable.device):
-        return api.pack(tensors, layout)
-
-    host_layout = layout_lib.Layout(original_layout.sharding_specs,
-                                    original_layout.mesh.host_mesh())
-
-    def get_host_dtensor():
-      # Copy to host mesh if needed.
-      if original_layout.mesh.device_type().upper() != 'CPU':
-        # Prefer pack and unpack in eager mode because it supports sharded
-        # layouts.
-        if context.executing_eagerly():
-          host_dtensor = api.pack(
-              api.unpack(dvariable.read_value()), host_layout)
-        else:
-          host_dtensor = api.copy_to_mesh(dvariable.read_value(), host_layout)
-      else:
-        host_dtensor = dvariable.read_value()
-      return (math_ops.cast(host_dtensor, dtypes.bfloat16)
-              if self.should_cast(host_dtensor) else host_dtensor)
-
-    num_local_devices = original_layout.mesh.num_local_devices()
-    super(_DVariableSaveable, self).__init__(
-        None,
-        [
-            DSaveSpec(
-                tensor=get_host_dtensor,
-                slice_spec=pack([''] * num_local_devices,
-                                layout_lib.Layout.replicated(
-                                    original_layout.mesh.host_mesh(), rank=0)),
-                name=pack([name] * num_local_devices,
-                          layout_lib.Layout.replicated(
-                              original_layout.mesh.host_mesh(), rank=0)),
-                global_shape=dvariable.shape,
-                # Layout is attached as attribute, no need to put it as a
-                # Tensor on DTensorDevice.
-                layout=host_layout.to_string(),
-                dtype=dtypes.bfloat16
-                if self.should_cast(dvariable) else dvariable.dtype,
-                device=dvariable.device)
-        ],
-        name)
-
-  def should_cast(self, v):
-    """Returns True if v has float32 dtype and is intructed to save as bf16.
-
-    Args:
-      v : The variable that determines whether to cast.
-
-    Returns:
-      True if current savable DVariable is instructed to save as bfloat16 and
-        the variable has dtype float32.
-    """
-    return self._dvariable.save_as_bf16 and v.dtype == dtypes.float32
-
-  def restore(self, restored_tensors, restored_shapes):
-    """Restores the same value into all variables."""
-    tensor, = restored_tensors
-
-    @def_function.function
-    def _restore(t):
-      with ops.device(self._dvariable.device):
-        return api.copy_to_mesh(t, self._original_layout)
-
-    # This assign establishes connections from restored tensor and tensors
-    # being restored to -- so that restore in SPMD can backtrack the DVariable
-    # and its layout, given that we're using tf.function style restore.
-    # Note that the restored dvaraible is on CPU no matter what as the restoreV2
-    # op must run on CPU.
-    # TODO(b/159035705): Allow restore for Tensor objects as well?
-    # Restore the dvariable back to original layout.
-    if self._original_layout.mesh.device_type().upper() != 'CPU':
-      tensor = _restore(tensor)
-    return self._dvariable.assign(
-        math_ops.cast(tensor, dtype=self._dvariable.dtype) if self._dvariable
-        .save_as_bf16 else tensor)
-
-
-@tf_export('experimental.dtensor.DVariable', v1=[])
-class DVariable(resource_variable_ops.ResourceVariable):
-  """A replacement for tf.Variable which follows initial value placement.
-
-    The class also handles restore/save operations in DTensor. Note that,
-    DVariable may fall back to normal tf.Variable at this moment if
-    `initial_value` is not a DTensor.
-  """
-
-  def __init__(self, initial_value, *args, dtype=None, **kwargs):
-    """Overrides tf.Variable to fix VarHandleOp placements."""
-    # Variables by default use the current device scope for placement. This
-    # wrapper has them follow the initial value's placement instead (which will
-    # be the DTensor device if the initial value has a layout).
-
-    # Pop layout from kwargs since keras make_variable may pass a 'layout'
-    # keyword argument. We need to pop it because we are passing kwargs to
-    # super class constructor.
-    layout = kwargs.pop('layout', None)
-    shape = kwargs.get('shape', None)
-
-    if callable(initial_value):
-      unwrapped = initial_value
-      if issubclass(type(initial_value), functools.partial):
-        unwrapped = initial_value.func  # pyrefly: ignore[missing-attribute]
-
-      # If wrapped is a CheckpointInitialValueCallable, this means that
-      # we are creating a Variable during a checkpoint restore.
-      # Thus the restore will happen now through this callable
-      # and we will create the DVariable with the restored dtensor.
-      if issubclass(type(unwrapped), trackable.CheckpointInitialValueCallable):
-        if not shape or not layout:
-          raise ValueError('Expected shape and layout to be not None.')
-
-        # CheckpointInitialValueCallable will call an eager tf.RestoreV2,
-        # which does not have any shape information or layout information
-        # attached. Thus we will do two things to have them correctly specified:
-        #
-        # The default layout scope allows us to correctly specify the output
-        # layout of the tf.RestoreV2 that will be called
-        #
-        # Passing shard_info with the correct shape allows the tf.RestoreV2
-        # ShapeInference to extract the shape.
-        initial_value = api.call_with_layout(
-            initial_value,
-            layout,
-            shard_info=trackable.ShardInfo(
-                shape=shape, offset=[0] * len(shape)))
-      else:
-        initial_value = initial_value()
-
-    # When the initial value came from a Checkpoint restoration, fetch tensor.
-    if isinstance(initial_value, trackable.CheckpointInitialValue):
-      initial_value = initial_value.wrapped_value
-
-    initial_value = ops.convert_to_tensor(initial_value, dtype=dtype)
-    variable_device = initial_value.device
-    self._save_as_bf16 = False
-    # TODO(b/159035705): The following code enables variable creation inside
-    # a tf.function. However, it requires a global dtensor device.
-    # if not variable_device and not tf.executing_eagerly():
-    #   try:
-    #     initial_value.op.get_attr("_layout")
-    #   except ValueError:
-    #     pass
-    #   else:
-    #     # The initial value is a DTensor, but because the DTensor device is
-    #     # only active during eager execution at the moment we need to
-    #     # translate that into a placement for the eager VarHandleOp.
-    #     variable_device = _dtensor_device().name
-    with ops.device(variable_device):
-      # If initial tensor assigned to DVariable is DTensor, record the layout of
-      # the resource so that this can be queried.
-      if context.executing_eagerly():
-        if api.is_dtensor(initial_value):
-          value_layout = api.fetch_layout(initial_value)
-          if layout is not None and layout != value_layout:
-            raise errors_impl.InvalidArgumentError(
-                None,
-                None,
-                'Conflicting layout are provided for initial '
-                f'value layout ({value_layout}) and variable ({layout}).',
-            )
-          layout = value_layout
-        elif layout is not None:
-          initial_value = api.relayout(initial_value, layout)
-        else:
-          raise errors_impl.InvalidArgumentError(
-              None,
-              None,
-              'Neither layout nor DTensor initial value are provided.',
-          )
-        self.layout = layout
-        with api.default_mesh(layout.mesh):
-          super(DVariable, self).__init__(
-              initial_value, *args, dtype=dtype, **kwargs
-          )
-      else:
-        # FIXME(175928457): Record value layout in graph mode.
-        if layout is not None:
-          initial_value = api.relayout(initial_value, layout)
-        super(DVariable, self).__init__(
-            initial_value, *args, dtype=dtype, **kwargs)
-
-  @property
-  def save_as_bf16(self):
-    return self._save_as_bf16
-
-  @save_as_bf16.setter
-  def save_as_bf16(self, save_as_bf16):
-    """Enables saving float32 as bfloat16."""
-    self._save_as_bf16 = save_as_bf16 and self.dtype == dtypes.float32
-
-  def _gather_saveables_for_checkpoint(self):
-    return {
-        trackable.VARIABLE_VALUE_KEY:
-            functools.partial(_DVariableSaveable, self)
-    }

From acbd31479839c36776be8a407e8394e9197c1419 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 17:58:36 +0530
Subject: [PATCH 0532/1000] Update save_restore_v2_test.py

---
 .../python/tests/save_restore_v2_test.py      | 31 +++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/tensorflow/dtensor/python/tests/save_restore_v2_test.py b/tensorflow/dtensor/python/tests/save_restore_v2_test.py
index e53f5003240..8455cf3b6f0 100644
--- a/tensorflow/dtensor/python/tests/save_restore_v2_test.py
+++ b/tensorflow/dtensor/python/tests/save_restore_v2_test.py
@@ -33,6 +33,7 @@ from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import stateless_random_ops
 from tensorflow.python.platform import test
+from tensorflow.python.trackable import base as trackable
 
 Mesh = layout_lib.Mesh
 Layout = layout_lib.Layout
@@ -122,6 +123,36 @@ class DTensorSaveRestoreV2Test(test_util.DTensorBaseTest):
     self.assertDTensorEqual(tensor_a, layout, dvariable_a.read_value())
     self.assertDTensorEqual(tensor_b, layout, dvariable_b.read_value())
 
+  def test_checkpoint_deferred_restore_scalar_dvariable(self):
+    layout = layout_lib.Layout.replicated(self.mesh, rank=0)
+    v_save = d_variable.DVariable(
+        api.pack(
+            [constant_op.constant(42.0)] * self.mesh.num_local_devices(),
+            layout,
+        )
+    )
+    ckpt = checkpoint.Checkpoint(v=v_save)
+    saved_path = ckpt.save(self.get_temp_dir())
+
+    # Re-create checkpoint on an empty container to test deferred restoration
+    root = trackable.SimpleTrackable()
+    restore_ckpt = checkpoint.Checkpoint(root=root)
+    restore_ckpt.restore(saved_path)
+
+    # Preload simple restoration produces a CheckpointInitialValueCallable
+    init_callable = root._preload_simple_restoration('v')
+    self.assertIsInstance(
+        init_callable, trackable.CheckpointInitialValueCallable
+    )
+
+    # Creating scalar DVariable with shape=() must succeed without ValueError
+    v_restored = d_variable.DVariable(
+        init_callable, shape=(), layout=layout, dtype=dtypes.float32
+    )
+    self.assertDTensorEqual(
+        constant_op.constant(42.0), layout, v_restored.read_value()
+    )
+
   @parameterized.named_parameters(
       ('x_unsharded', [_MESH_DIM_X, UNSHARDED]),
       ('unsharded_x', [UNSHARDED, _MESH_DIM_X]),

From a5479790df5251931db65314fea43c376ed38806 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 18:07:12 +0530
Subject: [PATCH 0533/1000] Update mesh_util.py

---
 tensorflow/dtensor/python/mesh_util.py | 653 ++++++++++++++-----------
 1 file changed, 368 insertions(+), 285 deletions(-)

diff --git a/tensorflow/dtensor/python/mesh_util.py b/tensorflow/dtensor/python/mesh_util.py
index a209d48fb27..931e9ee4a19 100644
--- a/tensorflow/dtensor/python/mesh_util.py
+++ b/tensorflow/dtensor/python/mesh_util.py
@@ -1,4 +1,4 @@
-# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2023 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.
@@ -12,303 +12,386 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Utilities to help with mesh creation."""
+"""Tests for mesh_util."""
+import os
+from unittest import mock
 
-from typing import Dict, List, Optional, Tuple, Union
-
-from absl import logging
-import numpy as np
+from absl.testing import parameterized
 
 from tensorflow.dtensor.python import accelerator_util
-from tensorflow.dtensor.python import api
 from tensorflow.dtensor.python import config
-from tensorflow.dtensor.python import layout
-from tensorflow.dtensor.python import tpu_util
+from tensorflow.dtensor.python import mesh_util
+from tensorflow.dtensor.python.tests import test_util
 from tensorflow.python.eager import context
+from tensorflow.python.framework import config as tf_config
 from tensorflow.python.framework import device as tf_device
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import math_ops
-from tensorflow.python.util.tf_export import tf_export
+from tensorflow.python.platform import test
 
 
-def _print_context(num_global_devices: int, num_clients: int, client_id: int,
-                   device_type: str, mesh: layout.Mesh) -> None:
-  logging.info('This is client %d of %d clients', client_id, num_clients)
-  logging.info('Number of global %s devices: %d', device_type.upper(),
-               num_global_devices)
-  # pylint: disable=protected-access
-  logging.info('Global device IDs: %s', mesh.global_device_ids())
-  logging.info('Local device IDs: %s', mesh.local_device_ids())
-  logging.info('Local devices: %s', mesh.local_devices())
-  # pylint: enable=protected-access
+class MeshUtilTest(test_util.DTensorBaseTest):
+  """Tests for mesh_util that do not require accelerator initialization."""
+
+  def test_mesh_creation(self):
+    self.skipForDeviceType(
+        ['TPU'], reason='Test is intended for CPUs and GPUs.'
+    )
+    mesh = mesh_util.create_mesh()
+    num_devices = len(test_util.list_local_logical_devices(mesh.device_type()))
+    self.assertEqual(mesh.num_local_devices(), num_devices)
+    self.assertEqual(mesh.size, num_devices)
+
+  def test_mesh_dict_creation(self):
+    self.skipForDeviceType(
+        ['TPU'], reason='Test is intended for CPUs and GPUs.'
+    )
+    num_devices = len(test_util.list_local_logical_devices('CPU'))
+    mesh = mesh_util.create_mesh({'x': num_devices, 'y': 1}, device_type='CPU')
+    num_devices = len(test_util.list_local_logical_devices(mesh.device_type()))
+    self.assertEqual(mesh.num_local_devices(), num_devices)
+    self.assertEqual(mesh.dim_names, ['x', 'y'])
+    self.assertEqual(mesh.size, num_devices)
+
+  def test_tpu_mesh_creation(self):
+    self.skipForDeviceType(['CPU', 'GPU'], reason='Test is intended for TPUs.')
+    mesh = mesh_util.create_mesh(mesh_name='1d_mesh', device_type='TPU')
+    num_devices = len(test_util.list_local_logical_devices('TPU'))
+    self.assertEqual(mesh.num_local_devices(), num_devices)
+    self.assertEqual(mesh.size, num_devices)
+
+  @parameterized.named_parameters(('use_xla_spmd', True),
+                                  ('do_not_use_xla_spmd', False))
+  def test_tpu_2d_mesh_creation(self, use_xla_spmd):
+    self.skipForDeviceType(['CPU', 'GPU'], reason='Test is intended for TPUs.')
+    self.skipForDeviceType(['TPU'],
+                           reason='Test requires exactly 2 cores',
+                           unless_device_count_equals_to=2)
+    devices = test_util.list_local_logical_devices('TPU')
+    self.assertLen(devices, 2)
+    mesh = mesh_util.create_mesh([('x', 2), ('y', 1)],
+                                 device_type='TPU',
+                                 use_xla_spmd=use_xla_spmd)
+    self.assertEqual(mesh.num_local_devices(), 2)
+    self.assertEqual(mesh.size, 2)
+    self.assertAllEqual(mesh.dim_names, ['x', 'y'])
+    self.assertEqual(mesh.use_xla_spmd(), use_xla_spmd)
+
+  def test_tpu_2d_mesh_creation_with_devices(self):
+    self.skipForDeviceType(['CPU', 'GPU'], reason='Test is intended for TPUs.')
+    self.skipForDeviceType(['TPU'],
+                           reason='Test requires at least 2 cores',
+                           unless_device_count_equals_to=2)
+    devices = test_util.list_local_logical_devices('TPU')
+    self.assertLen(devices, 2)
+    mesh = mesh_util.create_mesh([('x', 2), ('y', 1)],
+                                 devices=['/device:tpu:0', '/device:tpu:1'])
+    self.assertEqual(mesh.num_local_devices(), 2)
+    self.assertEqual(mesh.size, 2)
+    self.assertAllEqual(mesh.dim_names, ['x', 'y'])
+
+  def test_tpu_2d_mesh_creation_with_device_specs(self):
+    self.skipForDeviceType(['CPU', 'GPU'], reason='Test is intended for TPUs.')
+    self.skipForDeviceType(['TPU'],
+                           reason='Test requires at least 2 cores',
+                           unless_device_count_equals_to=2)
+    devices = test_util.list_local_logical_devices('TPU')
+    self.assertLen(devices, 2)
+    mesh = mesh_util.create_mesh(
+        [('x', 2), ('y', 1)],
+        devices=[
+            tf_device.DeviceSpec.from_string('/tpu:0'),
+            tf_device.DeviceSpec.from_string('/tpu:1'),
+        ],
+    )
+    self.assertEqual(mesh.num_local_devices(), 2)
+    self.assertEqual(mesh.size, 2)
+    self.assertAllEqual(mesh.dim_names, ['x', 'y'])
+
+  def test_single_client_mesh_creation(self):
+    self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
+    num_devices = len(test_util.list_local_logical_devices('CPU'))
+    with mock.patch.object(accelerator_util,
+                           'is_initialized') as is_initialized:
+      is_initialized.return_value = True
+      mesh = mesh_util.create_distributed_mesh(
+          mesh_name='single_client_1d_mesh', mesh_dims=[('x', num_devices)])
+      self.assertEqual(mesh.num_local_devices(), num_devices)
+      self.assertEqual(mesh.size, num_devices)
+
+  def test_single_client_mesh_dict_creation(self):
+    self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
+    num_devices = len(test_util.list_local_logical_devices('CPU'))
+    with mock.patch.object(
+        accelerator_util, 'is_initialized'
+    ) as is_initialized:
+      is_initialized.return_value = True
+      mesh = mesh_util.create_distributed_mesh(
+          mesh_name='single_client_1d_mesh',
+          mesh_dims={'x': num_devices, 'y': 1},
+      )
+      self.assertEqual(mesh.num_local_devices(), num_devices)
+      self.assertEqual(mesh.dim_names, ['x', 'y'])
+      self.assertEqual(mesh.size, num_devices)
+
+  def test_single_client_mesh_with_local_devices(self):
+    self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
+    with mock.patch.object(accelerator_util,
+                           'is_initialized') as is_initialized:
+      is_initialized.return_value = True
+      mesh = mesh_util.create_distributed_mesh(
+          mesh_name='single_client_1d_mesh',
+          mesh_dims=[('x', 1)],
+          local_devices=['CPU:0'])
+      self.assertEqual(mesh.num_local_devices(), 1)
+      self.assertEqual(mesh.size, 1)
+
+  def test_create_distributed_mesh_requires_initialize(self):
+    self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
+
+    with mock.patch.object(accelerator_util,
+                           'is_initialized') as is_initialized:
+      is_initialized.return_value = False
+      with self.assertRaisesRegex(ValueError, 'Accelerators are uninitialized'):
+        _ = mesh_util.create_distributed_mesh(
+            mesh_name='single_client_1d_mesh',
+            mesh_dims=[('x', 1)],
+            local_devices=['CPU:0'])
+
+  def test_create_distributed_mesh_tpu_with_local_devices_raises_error(self):
+    self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
+    with mock.patch.object(accelerator_util, 'is_initialized', return_value=True):
+      with self.assertRaisesRegex(
+          ValueError, 'Do not specify devices for TPU meshes'
+      ):
+        mesh_util.create_distributed_mesh(
+            mesh_name='tpu_mesh',
+            mesh_dims=[('x', 1)],
+            local_devices=['TPU:0'],
+        )
+      with self.assertRaisesRegex(
+          ValueError, 'Do not specify devices for TPU meshes'
+      ):
+        mesh_util.create_distributed_mesh(
+            mesh_name='tpu_mesh',
+            mesh_dims=[('x', 1)],
+            local_devices=['TPU:0'],
+            device_type='TPU',
+        )
+
+  def test_single_client_mesh_creation_wrong_shape(self):
+    self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
+    num_devices = len(test_util.list_local_logical_devices('CPU'))
+    with mock.patch.object(accelerator_util,
+                           'is_initialized') as is_initialized:
+      is_initialized.return_value = True
+      with self.assertRaisesRegex(ValueError,
+                                  'must be equal to total size of the mesh'):
+        mesh_util.create_distributed_mesh(
+            mesh_name='single_client_1d_mesh',
+            mesh_dims=[('x', num_devices * 2)])
+
+  def test_single_client_mesh_creation_using_fewer_devices(self):
+    self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
+    test_util.reset_logical_devices('CPU', 4)
+    with mock.patch.object(accelerator_util,
+                           'is_initialized') as is_initialized:
+      is_initialized.return_value = True
+      mesh = mesh_util.create_distributed_mesh(
+          mesh_name='single_client_1d_mesh',
+          mesh_dims=[('x', 2)],
+          local_devices=['CPU:0', 'CPU:1'])
+      self.assertEqual(mesh.num_local_devices(), 2)
+      self.assertEqual(mesh.size, 2)
+
+      mesh = mesh_util.create_distributed_mesh(
+          mesh_name='single_client_1d_mesh',
+          mesh_dims=[('x', 2)],
+          local_devices=['CPU:0', 'CPU:1'])
+      self.assertEqual(mesh.num_local_devices(), 2)
+      self.assertEqual(mesh.size, 2)
+
+  def test_single_client_mesh_creation_with_xla_spmd_raises_error(self):
+    self.skipForDeviceType(['TPU'],
+                           reason='Test is intended for non TPU devices')
+    test_util.reset_logical_devices('CPU', 4)
+    with mock.patch.object(accelerator_util,
+                           'is_initialized') as is_initialized:
+      is_initialized.return_value = True
+      with self.assertRaisesRegex(
+          ValueError, 'XLA SPMD is not currently not supported for'):
+        mesh_util.create_distributed_mesh(
+            mesh_name='single_client_mesh',
+            mesh_dims=[('x', 2)],
+            local_devices=['CPU:0', 'CPU:1'],
+            use_xla_spmd=True)
+
+  @mock.patch.object(config, 'num_clients')
+  @mock.patch.object(accelerator_util, 'is_initialized')
+  def test_multi_client_mesh_creation(self, num_clients, is_initialized):
+    self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
+    with mock.patch.object(accelerator_util,
+                           'is_initialized') as is_initialized:
+      with mock.patch.object(config, 'num_clients') as num_clients:
+        num_clients.return_value = 2
+        is_initialized.return_value = True
+        test_util.reset_context()
+        cpus = tf_config.list_physical_devices('CPU')
+        tf_config.set_logical_device_configuration(
+            cpus[0], [context.LogicalDeviceConfiguration()] * 4
+        )
+        with mock.patch.object(config, 'client_id', return_value=0):
+          mesh_1 = mesh_util.create_distributed_mesh(
+              mesh_name='multi_client_1d_mesh_1',
+              mesh_dims=[('x', 4)],
+              local_devices=['CPU:0', 'CPU:1'])
+        self.assertEqual(mesh_1.num_local_devices(), 2)
+        self.assertEqual(mesh_1.size, 4)
+        with mock.patch.object(config, 'client_id', return_value=1):
+          mesh_2 = mesh_util.create_distributed_mesh(
+              mesh_name='multi_client_1d_mesh_2',
+              mesh_dims=[('x', 4)],
+              local_devices=['CPU:2', 'CPU:3'])
+        self.assertEqual(mesh_2.num_local_devices(), 2)
+        self.assertEqual(mesh_2.size, 4)
 
 
-def _make_device_specs(
-    devices: Optional[List[Union[tf_device.DeviceSpec, str]]] = None,
-    device_type: Optional[str] = None
-) -> Tuple[List[tf_device.DeviceSpec], str]:
-  """Makes device specs for all local devices or from a provided list."""
+class InitializedMeshUtilTest(test_util.DTensorBaseTest):
+  """Tests for mesh_util that require accelerator initialization."""
 
-  if devices is None:
-    if device_type is None:
-      device_type = 'CPU'
-    devices = config.local_devices(device_type)
-  else:
-    if isinstance(devices[0], str):
-      devices = [tf_device.DeviceSpec.from_string(d) for d in devices]
-    if device_type is None:
-      device_type = devices[0].device_type  # pyrefly: ignore[missing-attribute]
+  def setUp(self):
+    super().setUp()
+    device_type = config.preferred_device_type()
+    accelerator_util.initialize_accelerator_system(device_type)
 
-    if device_type.upper() != devices[0].device_type.upper():  # pyrefly: ignore[missing-attribute]
-      raise ValueError(
-          f'Conflicting devices {str(devices)} and device_type {device_type}'
+  def tearDown(self):
+    super().tearDown()
+    context._reset_context()  # pylint: disable=protected-access
+
+  def test_is_initialized(self):
+    self.assertTrue(accelerator_util.is_initialized())
+
+  def test_initialize_accelerator_system(self):
+    accelerator_util.shutdown_accelerator_system()
+    device_type = accelerator_util.initialize_accelerator_system('CPU')
+    self.assertEqual(device_type, 'CPU')
+
+    # Default uses preferred_device_type.
+    accelerator_util.shutdown_accelerator_system()
+    device_type = accelerator_util.initialize_accelerator_system()
+    self.assertEqual(device_type, config.preferred_device_type())
+
+  @mock.patch.dict(os.environ, {'DTENSOR_GPU_USE_NCCL_COMMUNICATION': '1'})
+  def test_initialize_error_vgpu_with_nccl(self):
+    self.skipForDeviceType(['CPU', 'TPU'], reason='Test is intended for GPUs')
+
+    accelerator_util.shutdown_accelerator_system()
+    num_physical_devices = config.num_local_devices('GPU')
+    test_util.reset_logical_devices('GPU', 2 * num_physical_devices)
+    with self.assertRaisesRegex(ValueError,
+                                'DTENSOR_GPU_USE_NCCL_COMMUNICATION'):
+      accelerator_util.initialize_accelerator_system('GPU')
+
+  @mock.patch.dict(os.environ, {'DTENSOR_GPU_USE_NCCL_COMMUNICATION': '1'})
+  def test_initialize_with_nccl(self):
+    self.skipForDeviceType(['CPU', 'TPU'], reason='Test is intended for GPUs')
+
+    accelerator_util.shutdown_accelerator_system()
+    accelerator_util.initialize_accelerator_system('GPU')
+
+    num_devices = len(test_util.list_local_logical_devices('GPU'))
+    mesh = mesh_util.create_mesh([('dim', num_devices)], device_type='GPU')
+    # The following shall run, but there is no clear way to check if it uses
+    # Collectives backed by NCCL.
+    mesh_util.barrier(mesh)
+
+  def test_initialize_after_tensorflow(self):
+    accelerator_util.shutdown_accelerator_system()
+    context.ensure_initialized()
+
+    with self.assertRaisesRegex(ValueError,
+                                'TensorFlow has already been initialized'):
+      accelerator_util.initialize_accelerator_system('CPU')
+
+  def test_initialize_after_tensorflow_with_reset(self):
+    accelerator_util.shutdown_accelerator_system()
+    test_util.reset_logical_devices('CPU', 32)
+    context.ensure_initialized()
+
+    with self.assertLogs(level='WARNING') as log:
+      accelerator_util.initialize_accelerator_system(
+          'CPU', experimental_reset_context=True
       )
 
-  return devices, device_type
+    self.assertIn('experimental_reset_context', log[0][0].message)
+    # Preserves the original logical device setting.
+    self.assertLen(test_util.list_local_logical_devices('CPU'), 32)
+
+  @parameterized.parameters(
+      dict(
+          device_type='CPU', skip_for=[]
+      ),  # We can create CPU meshes on TPU and GPU platforms!
+      dict(device_type='GPU', skip_for=['CPU', 'TPU']),
+      dict(device_type='TPU', skip_for=['CPU', 'GPU']),
+  )
+  def test_initialize_with_manual_logical_cpu_devices(
+      self, device_type: str, skip_for: list[str]
+  ):
+    self.skipForDeviceType(
+        skip_for,
+        reason=f'Test is not intended for {skip_for}',
+    )
+    accelerator_util.shutdown_accelerator_system()
+    test_util.reset_logical_devices('CPU', 1)
+    accelerator_util.initialize_accelerator_system(
+        device_type, num_logical_cpu_devices=32
+    )
+    self.assertLen(test_util.list_local_logical_devices('CPU'), 32)
+
+  def test_shutdown_accelerator_system(self):
+    self.assertTrue(accelerator_util.is_initialized())
+    accelerator_util.shutdown_accelerator_system()
+    self.assertFalse(accelerator_util.is_initialized())
+
+    with self.assertRaisesRegex(ValueError, 'not initialized'):
+      accelerator_util.shutdown_accelerator_system()
+
+  def test_shutdown_accelerator_system_async_wait_error(self):
+    accelerator_util.shutdown_accelerator_system()
+    accelerator_util.initialize_accelerator_system(
+        'CPU', num_logical_cpu_devices=1)
+    self.assertTrue(accelerator_util.is_initialized())
+    with mock.patch.object(context, 'async_wait',
+                           side_effect=RuntimeError('async fail')):
+      with self.assertRaisesRegex(RuntimeError, 'async fail'):
+        accelerator_util.shutdown_accelerator_system()
+    # Ensure cleanup still ran in the finally block.
+    self.assertFalse(accelerator_util.is_initialized())
+
+  def test_distributed_tpu_mesh_creation(self):
+    self.skipForDeviceType(['CPU', 'GPU'], reason='Test is intended for TPUs')
+    self.skipForDeviceType(['TPU'],
+                           reason='Test requires exactly 8 cores',
+                           unless_device_count_equals_to=8)
+
+    num_devices = len(test_util.list_local_logical_devices('TPU'))
+    mesh = mesh_util.create_distributed_mesh(
+        mesh_name='distributed_1d_mesh',
+        mesh_dims=[('x', num_devices)],
+        device_type='TPU')
+    self.assertEqual(mesh.num_local_devices(), 8)
+    self.assertEqual(mesh.size, 8)
+
+  def test_mesh_barrier(self):
+    device_type = config.preferred_device_type()
+
+    num_devices = len(test_util.list_local_logical_devices(device_type))
+    mesh = mesh_util.create_mesh([('dim', num_devices)],
+                                 device_type=device_type)
+
+    # FIXME(b/235416015): To really test this we'll need a new eager async
+    # API.  The following shall run, but the barrier semantics is not tested.
+    mesh_util.barrier(mesh, 'Name')
+    mesh_util.barrier(mesh)
 
 
-@tf_export('experimental.dtensor.create_mesh', v1=[])
-def create_mesh(
-    mesh_dims: Optional[Union[List[Tuple[str, int]], Dict[str, int]]] = None,
-    mesh_name: str = '',
-    devices: Optional[List[Union[tf_device.DeviceSpec, str]]] = None,
-    device_type: Optional[str] = None,
-    use_xla_spmd: bool = layout.USE_XLA_SPMD,
-) -> layout.Mesh:
-  """Creates a single-client mesh.
-
-  If both `mesh_dims` and `devices` are specified, they must match each otehr.
-  As a special case, when all arguments are missing, this creates a 1D CPU mesh
-  with an empty name, assigning all available devices to that dimension.
-
-  Args:
-    mesh_dims: A dict of dim_name: dim_size, or a list of (dim_name, dim_size)
-      tuples. Defaults to a single batch-parallel dimension called 'x' usin all
-      devices. As a special case, a single-element mesh_dims whose dim_size is
-      -1 also uses all devices.  e.g. `{'x' : 4, 'y' : 1}` or `[('x', 4), ('y',
-      1)]`.
-    mesh_name: Name of the created mesh. Defaults to ''.
-    devices: String representations of devices to use. This is the device part
-      of tf.DeviceSpec, e.g. 'CPU:0'. Defaults to all available logical devices.
-    device_type: If `devices` is missing, the type of devices to use. Defaults
-      to 'CPU'.
-    use_xla_spmd: Boolean when True, will use XLA SPMD instead of DTensor SPMD.
-
-  Returns:
-    A single-client mesh created from specified or default arguments.
-  """
-  device_specs, device_type = _make_device_specs(devices, device_type)
-
-  local_spec = tf_device.DeviceSpec(job=config.job_name(), replica=0, task=0)
-  device_specs = [local_spec.make_merged_spec(d) for d in device_specs]
-
-  if isinstance(mesh_dims, dict):
-    mesh_dims = list(mesh_dims.items())
-  if mesh_dims is None:
-    mesh_dims = [('x', len(device_specs))]
-  elif len(mesh_dims) == 1 and mesh_dims[0][1] == -1:
-    # Replace -1 dim_size in a 1D mesh will the number of all devices.
-    mesh_dims[0] = (mesh_dims[0][0], len(device_specs))
-
-  dim_names = [d[0] for d in mesh_dims]
-  shape = [d[1] for d in mesh_dims]
-
-  if np.prod(shape) != len(device_specs):
-    raise ValueError(f'length of devices ({len(device_specs)}) must be '
-                     f'equal to total size of the mesh of shape {shape}')
-
-  global_device_ids = np.arange(len(device_specs)).reshape(shape)
-  local_device_ids = np.ravel(global_device_ids).tolist()
-  mesh = layout.Mesh(
-      dim_names=dim_names,
-      global_device_ids=global_device_ids,
-      local_device_ids=local_device_ids,
-      local_devices=device_specs,
-      mesh_name=mesh_name,
-      use_xla_spmd=use_xla_spmd)
-  _print_context(
-      num_global_devices=len(device_specs),
-      num_clients=1,
-      client_id=0,
-      device_type=device_type,
-      mesh=mesh)
-  return mesh
-
-
-@tf_export('experimental.dtensor.create_distributed_mesh', v1=[])
-def create_distributed_mesh(
-    mesh_dims: Union[List[Tuple[str, int]], Dict[str, int]],
-    mesh_name: str = '',
-    local_devices: Optional[List[Union[tf_device.DeviceSpec, str]]] = None,
-    device_type: Optional[str] = None,
-    use_xla_spmd: bool = layout.USE_XLA_SPMD,
-) -> layout.Mesh:
-  """Creates a distributed mesh.
-
-  This is similar to `create_mesh`, but with a different set of arguments to
-  create a mesh that spans evenly across a multi-client DTensor cluster.
-
-  For CPU and GPU meshes, users can choose to use fewer local devices than what
-  is available `local_devices`.
-
-  For TPU, only meshes that uses all TPU cores is supported by the DTensor
-  runtime.
-
-  Args:
-    mesh_dims: A dict of dim_name: dim_size, or a list of (dim_name, dim_size)
-      tuples. e.g. `{'x' : 4, 'y' : 1}` or `[('x', 4), ('y', 1)]`.
-    mesh_name: Name of the created mesh. Defaults to ''.
-    local_devices: String representations of devices to use. This is the device
-      part of tf.DeviceSpec, e.g. 'CPU:0'. Defaults to all available local
-      logical devices.
-    device_type: Type of device to build the mesh for. Defaults to 'CPU'.
-      Supported values are 'CPU', 'GPU', 'TPU'.6
-    use_xla_spmd: Boolean when True, will use XLA SPMD instead of DTensor SPMD.
-
-  Returns:
-    A mesh that spans evenly across all DTensor clients in the cluster.
-  """
-  if isinstance(mesh_dims, dict):
-    mesh_dims = list(mesh_dims.items())
-  dim_names, shape = zip(*mesh_dims)
-
-  if not accelerator_util.is_initialized():
-    raise ValueError('Accelerators are uninitialized, please run '
-                     'dtensor.initialize_accelerator_system() first.')
-
-  device_specs, device_type = _make_device_specs(local_devices, device_type)
-
-  if device_type and device_type.upper() == 'TPU':
-    # TODO(b/185940495): Allow multi-mesh and partial on TPU.
-    # TPU meshes can only be configured through environment variables that
-    # reflect the actual TPU topology. Do not let users specify custom args.
-    # Checked here (after type inference) rather than against the original
-    # `device_type` argument, so a caller who passes `local_devices` full of
-    # TPU specs without also passing `device_type='TPU'` still gets this
-    # error instead of silently having `local_devices` ignored below.
-    if local_devices is not None:
-      raise ValueError(
-          f'Do not specify devices for {device_type.upper()} meshes. '
-          f'Using a partial list of devices for {device_type.upper()} '
-          f'is not supported.')
-
-  if device_type.upper() in ['CPU', 'GPU']:
-    # For CPU and GPU meshes, user-specified args take precedence over env vars.
-    # This is particularly useful on single clients when users want to create
-    # meshes that use fewer logical devices than what's available.
-
-    local_spec = tf_device.DeviceSpec(
-        job=config.job_name(), replica=0, task=config.client_id())
-    device_specs = [local_spec.make_merged_spec(d) for d in device_specs]
-
-    # Assumes identical number of local devices per client.
-    num_global_devices = len(device_specs) * config.num_clients()
-
-    if np.prod(shape) != num_global_devices:
-      raise ValueError(
-          f'Global number of devices '
-          f'({len(device_specs)} per client * {config.num_clients()} clients '
-          f'= {num_global_devices}) must be '
-          f'equal to total size of the mesh of shape {shape}')
-
-    global_device_ids = np.arange(num_global_devices).reshape(shape)
-    flattened = np.ravel(global_device_ids).tolist()
-    start_idx = len(device_specs) * config.client_id()
-    local_device_ids = flattened[start_idx:start_idx + len(device_specs)]
-
-    mesh = layout.Mesh(
-        dim_names=dim_names,
-        global_device_ids=global_device_ids,
-        local_device_ids=local_device_ids,
-        local_devices=device_specs,
-        mesh_name=mesh_name,
-        use_xla_spmd=use_xla_spmd)
-    _print_context(num_global_devices, config.num_clients(), config.client_id(),
-                   device_type, mesh)
-    return mesh
-
-  if device_type.upper() == 'TPU':
-    mesh = tpu_util.create_tpu_mesh(
-        mesh_dim_names=dim_names,
-        mesh_shape=shape,
-        mesh_name=mesh_name,
-        use_xla_spmd=use_xla_spmd)
-    _print_context(
-        config.num_global_devices(device_type), config.num_clients(),
-        config.client_id(), device_type, mesh)
-    return mesh
-
-  raise ValueError(f'Device type {device_type} is not CPU, GPU or TPU')
-
-
-_BARRIER_DICT = {}
-
-
-@tf_export('experimental.dtensor.barrier', v1=[])
-def barrier(mesh: layout.Mesh,
-            barrier_name: Optional[str] = None,
-            timeout_in_ms: Optional[int] = None):
-  """Runs a barrier on the mesh.
-
-  Upon returning from the barrier, all operations run before the barrier
-  would have completed across all clients. Currently we allocate a fully
-  sharded tensor with mesh shape and run an all_reduce on it.
-
-  Example:
-
-  A barrier can be used before application exit to ensure completion of pending
-  ops.
-
-  ```python
-
-  x = [1, 2, 3]
-  x = dtensor.relayout(x, dtensor.Layout.batch_sharded(mesh, 'batch', 1))
-  dtensor.barrier(mesh)
-
-  # At this point all devices on all clients in the mesh have completed
-  # operations before the barrier. Therefore it is OK to tear down the clients.
-  sys.exit()
-  ```
-
-  Args:
-    mesh: The mesh to run the barrier on.
-    barrier_name: The name of the barrier. Mainly used for logging purpose.
-    timeout_in_ms: The timeout of the barrier in ms. If omitted, blocks
-      indefinitely till the barrier is reached from all clients.
-  """
-  if barrier_name is None:
-    barrier_name = '(barrier)'
-
-  logging.info('entering barrier before op: %s', barrier_name)
-
-  # Make sure all ops are consumed before running the sync.
-  context.async_wait()
-
-  # Reduction on a fully sharded tensor requires all devices to participate
-  # and serves as a barrier on the mesh.
-  component = array_ops.reshape(1.0, [1] * len(mesh.shape()))
-  ones = api.pack([component] * mesh.num_local_devices(),
-                  layout.Layout(mesh.dim_names, mesh))
-
-  mesh_size = math_ops.reduce_sum(ones)
-  if mesh_size != mesh.size:
-    raise ValueError(
-        'Global barrier produced wrong mesh size : {0} while mesh has actual'
-        'size : {1}'.format(mesh_size, mesh.size))
-
-  # TODO(hthu): This isn't strictly needed but might cause confusing behaviors
-  # from users. Consider dropping this if there is a `big` performance hit.
-  context.async_wait()
-
-  if context.context().coordination_service:
-    if timeout_in_ms is None:
-      timeout_in_ms = 24 * 60 * 60 * 1000  # 24 hours to stand in for infinite.
-
-    num_calls = _BARRIER_DICT.setdefault(barrier_name, 0)
-    _BARRIER_DICT[barrier_name] = num_calls + 1
-
-    barrier_id = f'{barrier_name}:{num_calls}'
-    context.context().wait_at_barrier(barrier_id, timeout_in_ms)
-
-  logging.info('finished running barrier across all clients after '
-               'op: %s', barrier_name)
+if __name__ == '__main__':
+  test.main()

From d835c354e4b83953443f0f2521bb152611e0d9b3 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 18:07:40 +0530
Subject: [PATCH 0534/1000] Update mesh_util.py

---
 tensorflow/dtensor/python/mesh_util.py | 657 +++++++++++--------------
 1 file changed, 289 insertions(+), 368 deletions(-)

diff --git a/tensorflow/dtensor/python/mesh_util.py b/tensorflow/dtensor/python/mesh_util.py
index 931e9ee4a19..6e2e962bc2a 100644
--- a/tensorflow/dtensor/python/mesh_util.py
+++ b/tensorflow/dtensor/python/mesh_util.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2022 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.
@@ -12,386 +12,307 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Tests for mesh_util."""
-import os
-from unittest import mock
+"""Utilities to help with mesh creation."""
 
-from absl.testing import parameterized
+from typing import Dict, List, Optional, Tuple, Union
+
+from absl import logging
+import numpy as np
 
 from tensorflow.dtensor.python import accelerator_util
+from tensorflow.dtensor.python import api
 from tensorflow.dtensor.python import config
-from tensorflow.dtensor.python import mesh_util
-from tensorflow.dtensor.python.tests import test_util
+from tensorflow.dtensor.python import layout
+from tensorflow.dtensor.python import tpu_util
 from tensorflow.python.eager import context
-from tensorflow.python.framework import config as tf_config
 from tensorflow.python.framework import device as tf_device
-from tensorflow.python.platform import test
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.util.tf_export import tf_export
 
 
-class MeshUtilTest(test_util.DTensorBaseTest):
-  """Tests for mesh_util that do not require accelerator initialization."""
-
-  def test_mesh_creation(self):
-    self.skipForDeviceType(
-        ['TPU'], reason='Test is intended for CPUs and GPUs.'
-    )
-    mesh = mesh_util.create_mesh()
-    num_devices = len(test_util.list_local_logical_devices(mesh.device_type()))
-    self.assertEqual(mesh.num_local_devices(), num_devices)
-    self.assertEqual(mesh.size, num_devices)
-
-  def test_mesh_dict_creation(self):
-    self.skipForDeviceType(
-        ['TPU'], reason='Test is intended for CPUs and GPUs.'
-    )
-    num_devices = len(test_util.list_local_logical_devices('CPU'))
-    mesh = mesh_util.create_mesh({'x': num_devices, 'y': 1}, device_type='CPU')
-    num_devices = len(test_util.list_local_logical_devices(mesh.device_type()))
-    self.assertEqual(mesh.num_local_devices(), num_devices)
-    self.assertEqual(mesh.dim_names, ['x', 'y'])
-    self.assertEqual(mesh.size, num_devices)
-
-  def test_tpu_mesh_creation(self):
-    self.skipForDeviceType(['CPU', 'GPU'], reason='Test is intended for TPUs.')
-    mesh = mesh_util.create_mesh(mesh_name='1d_mesh', device_type='TPU')
-    num_devices = len(test_util.list_local_logical_devices('TPU'))
-    self.assertEqual(mesh.num_local_devices(), num_devices)
-    self.assertEqual(mesh.size, num_devices)
-
-  @parameterized.named_parameters(('use_xla_spmd', True),
-                                  ('do_not_use_xla_spmd', False))
-  def test_tpu_2d_mesh_creation(self, use_xla_spmd):
-    self.skipForDeviceType(['CPU', 'GPU'], reason='Test is intended for TPUs.')
-    self.skipForDeviceType(['TPU'],
-                           reason='Test requires exactly 2 cores',
-                           unless_device_count_equals_to=2)
-    devices = test_util.list_local_logical_devices('TPU')
-    self.assertLen(devices, 2)
-    mesh = mesh_util.create_mesh([('x', 2), ('y', 1)],
-                                 device_type='TPU',
-                                 use_xla_spmd=use_xla_spmd)
-    self.assertEqual(mesh.num_local_devices(), 2)
-    self.assertEqual(mesh.size, 2)
-    self.assertAllEqual(mesh.dim_names, ['x', 'y'])
-    self.assertEqual(mesh.use_xla_spmd(), use_xla_spmd)
-
-  def test_tpu_2d_mesh_creation_with_devices(self):
-    self.skipForDeviceType(['CPU', 'GPU'], reason='Test is intended for TPUs.')
-    self.skipForDeviceType(['TPU'],
-                           reason='Test requires at least 2 cores',
-                           unless_device_count_equals_to=2)
-    devices = test_util.list_local_logical_devices('TPU')
-    self.assertLen(devices, 2)
-    mesh = mesh_util.create_mesh([('x', 2), ('y', 1)],
-                                 devices=['/device:tpu:0', '/device:tpu:1'])
-    self.assertEqual(mesh.num_local_devices(), 2)
-    self.assertEqual(mesh.size, 2)
-    self.assertAllEqual(mesh.dim_names, ['x', 'y'])
-
-  def test_tpu_2d_mesh_creation_with_device_specs(self):
-    self.skipForDeviceType(['CPU', 'GPU'], reason='Test is intended for TPUs.')
-    self.skipForDeviceType(['TPU'],
-                           reason='Test requires at least 2 cores',
-                           unless_device_count_equals_to=2)
-    devices = test_util.list_local_logical_devices('TPU')
-    self.assertLen(devices, 2)
-    mesh = mesh_util.create_mesh(
-        [('x', 2), ('y', 1)],
-        devices=[
-            tf_device.DeviceSpec.from_string('/tpu:0'),
-            tf_device.DeviceSpec.from_string('/tpu:1'),
-        ],
-    )
-    self.assertEqual(mesh.num_local_devices(), 2)
-    self.assertEqual(mesh.size, 2)
-    self.assertAllEqual(mesh.dim_names, ['x', 'y'])
-
-  def test_single_client_mesh_creation(self):
-    self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
-    num_devices = len(test_util.list_local_logical_devices('CPU'))
-    with mock.patch.object(accelerator_util,
-                           'is_initialized') as is_initialized:
-      is_initialized.return_value = True
-      mesh = mesh_util.create_distributed_mesh(
-          mesh_name='single_client_1d_mesh', mesh_dims=[('x', num_devices)])
-      self.assertEqual(mesh.num_local_devices(), num_devices)
-      self.assertEqual(mesh.size, num_devices)
-
-  def test_single_client_mesh_dict_creation(self):
-    self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
-    num_devices = len(test_util.list_local_logical_devices('CPU'))
-    with mock.patch.object(
-        accelerator_util, 'is_initialized'
-    ) as is_initialized:
-      is_initialized.return_value = True
-      mesh = mesh_util.create_distributed_mesh(
-          mesh_name='single_client_1d_mesh',
-          mesh_dims={'x': num_devices, 'y': 1},
-      )
-      self.assertEqual(mesh.num_local_devices(), num_devices)
-      self.assertEqual(mesh.dim_names, ['x', 'y'])
-      self.assertEqual(mesh.size, num_devices)
-
-  def test_single_client_mesh_with_local_devices(self):
-    self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
-    with mock.patch.object(accelerator_util,
-                           'is_initialized') as is_initialized:
-      is_initialized.return_value = True
-      mesh = mesh_util.create_distributed_mesh(
-          mesh_name='single_client_1d_mesh',
-          mesh_dims=[('x', 1)],
-          local_devices=['CPU:0'])
-      self.assertEqual(mesh.num_local_devices(), 1)
-      self.assertEqual(mesh.size, 1)
-
-  def test_create_distributed_mesh_requires_initialize(self):
-    self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
-
-    with mock.patch.object(accelerator_util,
-                           'is_initialized') as is_initialized:
-      is_initialized.return_value = False
-      with self.assertRaisesRegex(ValueError, 'Accelerators are uninitialized'):
-        _ = mesh_util.create_distributed_mesh(
-            mesh_name='single_client_1d_mesh',
-            mesh_dims=[('x', 1)],
-            local_devices=['CPU:0'])
-
-  def test_create_distributed_mesh_tpu_with_local_devices_raises_error(self):
-    self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
-    with mock.patch.object(accelerator_util, 'is_initialized', return_value=True):
-      with self.assertRaisesRegex(
-          ValueError, 'Do not specify devices for TPU meshes'
-      ):
-        mesh_util.create_distributed_mesh(
-            mesh_name='tpu_mesh',
-            mesh_dims=[('x', 1)],
-            local_devices=['TPU:0'],
-        )
-      with self.assertRaisesRegex(
-          ValueError, 'Do not specify devices for TPU meshes'
-      ):
-        mesh_util.create_distributed_mesh(
-            mesh_name='tpu_mesh',
-            mesh_dims=[('x', 1)],
-            local_devices=['TPU:0'],
-            device_type='TPU',
-        )
-
-  def test_single_client_mesh_creation_wrong_shape(self):
-    self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
-    num_devices = len(test_util.list_local_logical_devices('CPU'))
-    with mock.patch.object(accelerator_util,
-                           'is_initialized') as is_initialized:
-      is_initialized.return_value = True
-      with self.assertRaisesRegex(ValueError,
-                                  'must be equal to total size of the mesh'):
-        mesh_util.create_distributed_mesh(
-            mesh_name='single_client_1d_mesh',
-            mesh_dims=[('x', num_devices * 2)])
-
-  def test_single_client_mesh_creation_using_fewer_devices(self):
-    self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
-    test_util.reset_logical_devices('CPU', 4)
-    with mock.patch.object(accelerator_util,
-                           'is_initialized') as is_initialized:
-      is_initialized.return_value = True
-      mesh = mesh_util.create_distributed_mesh(
-          mesh_name='single_client_1d_mesh',
-          mesh_dims=[('x', 2)],
-          local_devices=['CPU:0', 'CPU:1'])
-      self.assertEqual(mesh.num_local_devices(), 2)
-      self.assertEqual(mesh.size, 2)
-
-      mesh = mesh_util.create_distributed_mesh(
-          mesh_name='single_client_1d_mesh',
-          mesh_dims=[('x', 2)],
-          local_devices=['CPU:0', 'CPU:1'])
-      self.assertEqual(mesh.num_local_devices(), 2)
-      self.assertEqual(mesh.size, 2)
-
-  def test_single_client_mesh_creation_with_xla_spmd_raises_error(self):
-    self.skipForDeviceType(['TPU'],
-                           reason='Test is intended for non TPU devices')
-    test_util.reset_logical_devices('CPU', 4)
-    with mock.patch.object(accelerator_util,
-                           'is_initialized') as is_initialized:
-      is_initialized.return_value = True
-      with self.assertRaisesRegex(
-          ValueError, 'XLA SPMD is not currently not supported for'):
-        mesh_util.create_distributed_mesh(
-            mesh_name='single_client_mesh',
-            mesh_dims=[('x', 2)],
-            local_devices=['CPU:0', 'CPU:1'],
-            use_xla_spmd=True)
-
-  @mock.patch.object(config, 'num_clients')
-  @mock.patch.object(accelerator_util, 'is_initialized')
-  def test_multi_client_mesh_creation(self, num_clients, is_initialized):
-    self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
-    with mock.patch.object(accelerator_util,
-                           'is_initialized') as is_initialized:
-      with mock.patch.object(config, 'num_clients') as num_clients:
-        num_clients.return_value = 2
-        is_initialized.return_value = True
-        test_util.reset_context()
-        cpus = tf_config.list_physical_devices('CPU')
-        tf_config.set_logical_device_configuration(
-            cpus[0], [context.LogicalDeviceConfiguration()] * 4
-        )
-        with mock.patch.object(config, 'client_id', return_value=0):
-          mesh_1 = mesh_util.create_distributed_mesh(
-              mesh_name='multi_client_1d_mesh_1',
-              mesh_dims=[('x', 4)],
-              local_devices=['CPU:0', 'CPU:1'])
-        self.assertEqual(mesh_1.num_local_devices(), 2)
-        self.assertEqual(mesh_1.size, 4)
-        with mock.patch.object(config, 'client_id', return_value=1):
-          mesh_2 = mesh_util.create_distributed_mesh(
-              mesh_name='multi_client_1d_mesh_2',
-              mesh_dims=[('x', 4)],
-              local_devices=['CPU:2', 'CPU:3'])
-        self.assertEqual(mesh_2.num_local_devices(), 2)
-        self.assertEqual(mesh_2.size, 4)
+def _print_context(num_global_devices: int, num_clients: int, client_id: int,
+                   device_type: str, mesh: layout.Mesh) -> None:
+  logging.info('This is client %d of %d clients', client_id, num_clients)
+  logging.info('Number of global %s devices: %d', device_type.upper(),
+               num_global_devices)
+  # pylint: disable=protected-access
+  logging.info('Global device IDs: %s', mesh.global_device_ids())
+  logging.info('Local device IDs: %s', mesh.local_device_ids())
+  logging.info('Local devices: %s', mesh.local_devices())
+  # pylint: enable=protected-access
 
 
-class InitializedMeshUtilTest(test_util.DTensorBaseTest):
-  """Tests for mesh_util that require accelerator initialization."""
+def _make_device_specs(
+    devices: Optional[List[Union[tf_device.DeviceSpec, str]]] = None,
+    device_type: Optional[str] = None
+) -> Tuple[List[tf_device.DeviceSpec], str]:
+  """Makes device specs for all local devices or from a provided list."""
 
-  def setUp(self):
-    super().setUp()
-    device_type = config.preferred_device_type()
-    accelerator_util.initialize_accelerator_system(device_type)
+  if devices is None:
+    if device_type is None:
+      device_type = 'CPU'
+    devices = config.local_devices(device_type)
+  else:
+    if not devices:
+      raise ValueError(
+          'Received an empty `devices` list; at least one device is '
+          'required to build a mesh.')
+    if isinstance(devices[0], str):
+      devices = [tf_device.DeviceSpec.from_string(d) for d in devices]
+    if device_type is None:
+      device_type = devices[0].device_type  # pyrefly: ignore[missing-attribute]
 
-  def tearDown(self):
-    super().tearDown()
-    context._reset_context()  # pylint: disable=protected-access
-
-  def test_is_initialized(self):
-    self.assertTrue(accelerator_util.is_initialized())
-
-  def test_initialize_accelerator_system(self):
-    accelerator_util.shutdown_accelerator_system()
-    device_type = accelerator_util.initialize_accelerator_system('CPU')
-    self.assertEqual(device_type, 'CPU')
-
-    # Default uses preferred_device_type.
-    accelerator_util.shutdown_accelerator_system()
-    device_type = accelerator_util.initialize_accelerator_system()
-    self.assertEqual(device_type, config.preferred_device_type())
-
-  @mock.patch.dict(os.environ, {'DTENSOR_GPU_USE_NCCL_COMMUNICATION': '1'})
-  def test_initialize_error_vgpu_with_nccl(self):
-    self.skipForDeviceType(['CPU', 'TPU'], reason='Test is intended for GPUs')
-
-    accelerator_util.shutdown_accelerator_system()
-    num_physical_devices = config.num_local_devices('GPU')
-    test_util.reset_logical_devices('GPU', 2 * num_physical_devices)
-    with self.assertRaisesRegex(ValueError,
-                                'DTENSOR_GPU_USE_NCCL_COMMUNICATION'):
-      accelerator_util.initialize_accelerator_system('GPU')
-
-  @mock.patch.dict(os.environ, {'DTENSOR_GPU_USE_NCCL_COMMUNICATION': '1'})
-  def test_initialize_with_nccl(self):
-    self.skipForDeviceType(['CPU', 'TPU'], reason='Test is intended for GPUs')
-
-    accelerator_util.shutdown_accelerator_system()
-    accelerator_util.initialize_accelerator_system('GPU')
-
-    num_devices = len(test_util.list_local_logical_devices('GPU'))
-    mesh = mesh_util.create_mesh([('dim', num_devices)], device_type='GPU')
-    # The following shall run, but there is no clear way to check if it uses
-    # Collectives backed by NCCL.
-    mesh_util.barrier(mesh)
-
-  def test_initialize_after_tensorflow(self):
-    accelerator_util.shutdown_accelerator_system()
-    context.ensure_initialized()
-
-    with self.assertRaisesRegex(ValueError,
-                                'TensorFlow has already been initialized'):
-      accelerator_util.initialize_accelerator_system('CPU')
-
-  def test_initialize_after_tensorflow_with_reset(self):
-    accelerator_util.shutdown_accelerator_system()
-    test_util.reset_logical_devices('CPU', 32)
-    context.ensure_initialized()
-
-    with self.assertLogs(level='WARNING') as log:
-      accelerator_util.initialize_accelerator_system(
-          'CPU', experimental_reset_context=True
+    if device_type.upper() != devices[0].device_type.upper():  # pyrefly: ignore[missing-attribute]
+      raise ValueError(
+          f'Conflicting devices {str(devices)} and device_type {device_type}'
       )
 
-    self.assertIn('experimental_reset_context', log[0][0].message)
-    # Preserves the original logical device setting.
-    self.assertLen(test_util.list_local_logical_devices('CPU'), 32)
-
-  @parameterized.parameters(
-      dict(
-          device_type='CPU', skip_for=[]
-      ),  # We can create CPU meshes on TPU and GPU platforms!
-      dict(device_type='GPU', skip_for=['CPU', 'TPU']),
-      dict(device_type='TPU', skip_for=['CPU', 'GPU']),
-  )
-  def test_initialize_with_manual_logical_cpu_devices(
-      self, device_type: str, skip_for: list[str]
-  ):
-    self.skipForDeviceType(
-        skip_for,
-        reason=f'Test is not intended for {skip_for}',
-    )
-    accelerator_util.shutdown_accelerator_system()
-    test_util.reset_logical_devices('CPU', 1)
-    accelerator_util.initialize_accelerator_system(
-        device_type, num_logical_cpu_devices=32
-    )
-    self.assertLen(test_util.list_local_logical_devices('CPU'), 32)
-
-  def test_shutdown_accelerator_system(self):
-    self.assertTrue(accelerator_util.is_initialized())
-    accelerator_util.shutdown_accelerator_system()
-    self.assertFalse(accelerator_util.is_initialized())
-
-    with self.assertRaisesRegex(ValueError, 'not initialized'):
-      accelerator_util.shutdown_accelerator_system()
-
-  def test_shutdown_accelerator_system_async_wait_error(self):
-    accelerator_util.shutdown_accelerator_system()
-    accelerator_util.initialize_accelerator_system(
-        'CPU', num_logical_cpu_devices=1)
-    self.assertTrue(accelerator_util.is_initialized())
-    with mock.patch.object(context, 'async_wait',
-                           side_effect=RuntimeError('async fail')):
-      with self.assertRaisesRegex(RuntimeError, 'async fail'):
-        accelerator_util.shutdown_accelerator_system()
-    # Ensure cleanup still ran in the finally block.
-    self.assertFalse(accelerator_util.is_initialized())
-
-  def test_distributed_tpu_mesh_creation(self):
-    self.skipForDeviceType(['CPU', 'GPU'], reason='Test is intended for TPUs')
-    self.skipForDeviceType(['TPU'],
-                           reason='Test requires exactly 8 cores',
-                           unless_device_count_equals_to=8)
-
-    num_devices = len(test_util.list_local_logical_devices('TPU'))
-    mesh = mesh_util.create_distributed_mesh(
-        mesh_name='distributed_1d_mesh',
-        mesh_dims=[('x', num_devices)],
-        device_type='TPU')
-    self.assertEqual(mesh.num_local_devices(), 8)
-    self.assertEqual(mesh.size, 8)
-
-  def test_mesh_barrier(self):
-    device_type = config.preferred_device_type()
-
-    num_devices = len(test_util.list_local_logical_devices(device_type))
-    mesh = mesh_util.create_mesh([('dim', num_devices)],
-                                 device_type=device_type)
-
-    # FIXME(b/235416015): To really test this we'll need a new eager async
-    # API.  The following shall run, but the barrier semantics is not tested.
-    mesh_util.barrier(mesh, 'Name')
-    mesh_util.barrier(mesh)
+  return devices, device_type
 
 
-if __name__ == '__main__':
-  test.main()
+@tf_export('experimental.dtensor.create_mesh', v1=[])
+def create_mesh(
+    mesh_dims: Optional[Union[List[Tuple[str, int]], Dict[str, int]]] = None,
+    mesh_name: str = '',
+    devices: Optional[List[Union[tf_device.DeviceSpec, str]]] = None,
+    device_type: Optional[str] = None,
+    use_xla_spmd: bool = layout.USE_XLA_SPMD,
+) -> layout.Mesh:
+  """Creates a single-client mesh.
+
+  If both `mesh_dims` and `devices` are specified, they must match each otehr.
+  As a special case, when all arguments are missing, this creates a 1D CPU mesh
+  with an empty name, assigning all available devices to that dimension.
+
+  Args:
+    mesh_dims: A dict of dim_name: dim_size, or a list of (dim_name, dim_size)
+      tuples. Defaults to a single batch-parallel dimension called 'x' usin all
+      devices. As a special case, a single-element mesh_dims whose dim_size is
+      -1 also uses all devices.  e.g. `{'x' : 4, 'y' : 1}` or `[('x', 4), ('y',
+      1)]`.
+    mesh_name: Name of the created mesh. Defaults to ''.
+    devices: String representations of devices to use. This is the device part
+      of tf.DeviceSpec, e.g. 'CPU:0'. Defaults to all available logical devices.
+    device_type: If `devices` is missing, the type of devices to use. Defaults
+      to 'CPU'.
+    use_xla_spmd: Boolean when True, will use XLA SPMD instead of DTensor SPMD.
+
+  Returns:
+    A single-client mesh created from specified or default arguments.
+  """
+  device_specs, device_type = _make_device_specs(devices, device_type)
+
+  local_spec = tf_device.DeviceSpec(job=config.job_name(), replica=0, task=0)
+  device_specs = [local_spec.make_merged_spec(d) for d in device_specs]
+
+  if isinstance(mesh_dims, dict):
+    mesh_dims = list(mesh_dims.items())
+  if mesh_dims is None:
+    mesh_dims = [('x', len(device_specs))]
+  elif len(mesh_dims) == 1 and mesh_dims[0][1] == -1:
+    # Replace -1 dim_size in a 1D mesh will the number of all devices.
+    mesh_dims[0] = (mesh_dims[0][0], len(device_specs))
+
+  dim_names = [d[0] for d in mesh_dims]
+  shape = [d[1] for d in mesh_dims]
+
+  if np.prod(shape) != len(device_specs):
+    raise ValueError(f'length of devices ({len(device_specs)}) must be '
+                     f'equal to total size of the mesh of shape {shape}')
+
+  global_device_ids = np.arange(len(device_specs)).reshape(shape)
+  local_device_ids = np.ravel(global_device_ids).tolist()
+  mesh = layout.Mesh(
+      dim_names=dim_names,
+      global_device_ids=global_device_ids,
+      local_device_ids=local_device_ids,
+      local_devices=device_specs,
+      mesh_name=mesh_name,
+      use_xla_spmd=use_xla_spmd)
+  _print_context(
+      num_global_devices=len(device_specs),
+      num_clients=1,
+      client_id=0,
+      device_type=device_type,
+      mesh=mesh)
+  return mesh
+
+
+@tf_export('experimental.dtensor.create_distributed_mesh', v1=[])
+def create_distributed_mesh(
+    mesh_dims: Union[List[Tuple[str, int]], Dict[str, int]],
+    mesh_name: str = '',
+    local_devices: Optional[List[Union[tf_device.DeviceSpec, str]]] = None,
+    device_type: Optional[str] = None,
+    use_xla_spmd: bool = layout.USE_XLA_SPMD,
+) -> layout.Mesh:
+  """Creates a distributed mesh.
+
+  This is similar to `create_mesh`, but with a different set of arguments to
+  create a mesh that spans evenly across a multi-client DTensor cluster.
+
+  For CPU and GPU meshes, users can choose to use fewer local devices than what
+  is available `local_devices`.
+
+  For TPU, only meshes that use all TPU cores is supported by the DTensor
+  runtime.
+
+  Args:
+    mesh_dims: A dict of dim_name: dim_size, or a list of (dim_name, dim_size)
+      tuples. e.g. `{'x' : 4, 'y' : 1}` or `[('x', 4), ('y', 1)]`.
+    mesh_name: Name of the created mesh. Defaults to ''.
+    local_devices: String representations of devices to use. This is the device
+      part of tf.DeviceSpec, e.g. 'CPU:0'. Defaults to all available local
+      logical devices.
+    device_type: Type of device to build the mesh for. Defaults to 'CPU'.
+      Supported values are 'CPU', 'GPU', 'TPU'.
+    use_xla_spmd: Boolean when True, will use XLA SPMD instead of DTensor SPMD.
+
+  Returns:
+    A mesh that spans evenly across all DTensor clients in the cluster.
+  """
+  if isinstance(mesh_dims, dict):
+    mesh_dims = list(mesh_dims.items())
+  dim_names, shape = zip(*mesh_dims)
+
+  if not accelerator_util.is_initialized():
+    raise ValueError('Accelerators are uninitialized, please run '
+                     'dtensor.initialize_accelerator_system() first.')
+
+  device_specs, device_type = _make_device_specs(local_devices, device_type)
+
+  if device_type and device_type.upper() == 'TPU':
+    # TODO(b/185940495): Allow multi-mesh and partial on TPU.
+    # TPU meshes can only be configured through environment variables that
+    # reflect the actual TPU topology. Do not let users specify custom args.
+    # Checked here (after type inference) rather than against the original
+    # `device_type` argument, so a caller who passes `local_devices` full of
+    # TPU specs without also passing `device_type='TPU'` still gets this
+    # error instead of silently having `local_devices` ignored below.
+    if local_devices is not None:
+      raise ValueError(
+          f'Do not specify devices for {device_type.upper()} meshes. '
+          f'Using a partial list of devices for {device_type.upper()} '
+          f'is not supported.')
+
+  if device_type.upper() in ['CPU', 'GPU']:
+    # For CPU and GPU meshes, user-specified args take precedence over env vars.
+    # This is particularly useful on single clients when users want to create
+    # meshes that use fewer logical devices than what's available.
+
+    local_spec = tf_device.DeviceSpec(
+        job=config.job_name(), replica=0, task=config.client_id())
+    device_specs = [local_spec.make_merged_spec(d) for d in device_specs]
+
+    # Assumes identical number of local devices per client.
+    num_global_devices = len(device_specs) * config.num_clients()
+
+    if np.prod(shape) != num_global_devices:
+      raise ValueError(
+          f'Global number of devices '
+          f'({len(device_specs)} per client * {config.num_clients()} clients '
+          f'= {num_global_devices}) must be '
+          f'equal to total size of the mesh of shape {shape}')
+
+    global_device_ids = np.arange(num_global_devices).reshape(shape)
+    flattened = np.ravel(global_device_ids).tolist()
+    start_idx = len(device_specs) * config.client_id()
+    local_device_ids = flattened[start_idx:start_idx + len(device_specs)]
+
+    mesh = layout.Mesh(
+        dim_names=dim_names,
+        global_device_ids=global_device_ids,
+        local_device_ids=local_device_ids,
+        local_devices=device_specs,
+        mesh_name=mesh_name,
+        use_xla_spmd=use_xla_spmd)
+    _print_context(num_global_devices, config.num_clients(), config.client_id(),
+                   device_type, mesh)
+    return mesh
+
+  if device_type.upper() == 'TPU':
+    mesh = tpu_util.create_tpu_mesh(
+        mesh_dim_names=dim_names,
+        mesh_shape=shape,
+        mesh_name=mesh_name,
+        use_xla_spmd=use_xla_spmd)
+    _print_context(
+        config.num_global_devices(device_type), config.num_clients(),
+        config.client_id(), device_type, mesh)
+    return mesh
+
+  raise ValueError(f'Device type {device_type} is not CPU, GPU or TPU')
+
+
+_BARRIER_DICT = {}
+
+
+@tf_export('experimental.dtensor.barrier', v1=[])
+def barrier(mesh: layout.Mesh,
+            barrier_name: Optional[str] = None,
+            timeout_in_ms: Optional[int] = None):
+  """Runs a barrier on the mesh.
+
+  Upon returning from the barrier, all operations run before the barrier
+  would have completed across all clients. Currently we allocate a fully
+  sharded tensor with mesh shape and run an all_reduce on it.
+
+  Example:
+
+  A barrier can be used before application exit to ensure completion of pending
+  ops.
+
+  ```python
+
+  x = [1, 2, 3]
+  x = dtensor.relayout(x, dtensor.Layout.batch_sharded(mesh, 'batch', 1))
+  dtensor.barrier(mesh)
+
+  # At this point all devices on all clients in the mesh have completed
+  # operations before the barrier. Therefore it is OK to tear down the clients.
+  sys.exit()
+  ```
+
+  Args:
+    mesh: The mesh to run the barrier on.
+    barrier_name: The name of the barrier. Mainly used for logging purpose.
+    timeout_in_ms: The timeout of the barrier in ms. If omitted, blocks
+      indefinitely till the barrier is reached from all clients.
+  """
+  if barrier_name is None:
+    barrier_name = '(barrier)'
+
+  logging.info('entering barrier before op: %s', barrier_name)
+
+  # Make sure all ops are consumed before running the sync.
+  context.async_wait()
+
+  # Reduction on a fully sharded tensor requires all devices to participate
+  # and serves as a barrier on the mesh.
+  component = array_ops.reshape(1.0, [1] * len(mesh.shape()))
+  ones = api.pack([component] * mesh.num_local_devices(),
+                  layout.Layout(mesh.dim_names, mesh))
+
+  mesh_size = math_ops.reduce_sum(ones)
+  if mesh_size != mesh.size:
+    raise ValueError(
+        'Global barrier produced wrong mesh size : {0} while mesh has actual'
+        'size : {1}'.format(mesh_size, mesh.size))
+
+  # TODO(hthu): This isn't strictly needed but might cause confusing behaviors
+  # from users. Consider dropping this if there is a `big` performance hit.
+  context.async_wait()
+
+  if context.context().coordination_service:
+    if timeout_in_ms is None:
+      timeout_in_ms = 24 * 60 * 60 * 1000  # 24 hours to stand in for infinite.
+
+    num_calls = _BARRIER_DICT.setdefault(barrier_name, 0)
+    _BARRIER_DICT[barrier_name] = num_calls + 1
+
+    barrier_id = f'{barrier_name}:{num_calls}'
+    context.context().wait_at_barrier(barrier_id, timeout_in_ms)
+
+  logging.info('finished running barrier across all clients after '
+               'op: %s', barrier_name)

From 5650689eff462252d2b1808072e4cc9a779845e3 Mon Sep 17 00:00:00 2001
From: wondertx <wondertx@gmail.com>
Date: Mon, 12 Oct 2026 20:38:27 +0800
Subject: [PATCH 0535/1000] fix test coverage

---
 tensorflow/core/ops/nccl_ops_test.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/core/ops/nccl_ops_test.cc b/tensorflow/core/ops/nccl_ops_test.cc
index 70266a8a805..f881a411b82 100644
--- a/tensorflow/core/ops/nccl_ops_test.cc
+++ b/tensorflow/core/ops/nccl_ops_test.cc
@@ -57,6 +57,7 @@ TEST(NcclOpsTest, NcclBroadcastRecv_ShapeFn_Int64) {
   BuildNcclBroadcastRecv(DT_INT64, &op);
 
   INFER_OK(op, "[3]", "[?,?,?]");
+  INFER_ERROR("Shape must be rank 1", op, "[1,2]");
 
   // A dimension that does not fit in int32 must be preserved.
   Tensor shape_t = test::AsTensor<int64_t>({int64_t{1} << 32, 2});

From 517ae8abb54083c92a17776d63befc6a03ffa9f1 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:46:53 -0700
Subject: [PATCH 0536/1000] Automated Code Change

PiperOrigin-RevId: 997669115
---
 tensorflow/core/framework/variant_tensor_data.h | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/tensorflow/core/framework/variant_tensor_data.h b/tensorflow/core/framework/variant_tensor_data.h
index e4fabce2994..01b45e33ab9 100644
--- a/tensorflow/core/framework/variant_tensor_data.h
+++ b/tensorflow/core/framework/variant_tensor_data.h
@@ -48,10 +48,8 @@ class VariantTensorData {
   const std::string& type_name() const { return type_name_; }
   void set_type_name(const std::string& type_name) { type_name_ = type_name; }
 
-  template <
-      typename T,
-      bool = std::is_trivially_copyable<typename std::decay<T>::type>::value &&
-             !std::is_pointer<typename std::decay<T>::type>::value>
+  template <typename T, bool = std::is_trivially_copyable_v<std::decay_t<T>> &&
+                               !std::is_pointer_v<std::decay_t<T>>>
   struct PODResolver {};
 
   // Portions of the object that are not Tensors.
@@ -136,14 +134,14 @@ class VariantTensorData {
   template <typename T>
   void SetMetadata(const T& value, PODResolver<T, false /* is_pod */>) {
     static_assert(
-        std::is_pointer<typename std::decay<T>::type>::value,
+        std::is_pointer_v<std::decay_t<T>>,
         "Only strings and pointers are supported for non-POD SetMetadata");
   }
 
   template <typename T>
   bool GetMetadata(T* value, PODResolver<T, false /* is_pod */>) const {
     static_assert(
-        std::is_pointer<typename std::decay<T>::type>::value,
+        std::is_pointer_v<std::decay_t<T>>,
         "Only strings and pointers are supported for non-POD GetMetadata");
     return false;
   }

From 5b83926ed68f383d1e7155afcea1ea4290d2fe52 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 18:11:56 +0530
Subject: [PATCH 0537/1000] Fix logging format in mesh_util.py


From bdf43f1e1604f48c22714104ee0fcef1e6d4f803 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 18:12:28 +0530
Subject: [PATCH 0538/1000] Update mesh_util_test.py

---
 .../dtensor/python/tests/mesh_util_test.py    | 21 +++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/tensorflow/dtensor/python/tests/mesh_util_test.py b/tensorflow/dtensor/python/tests/mesh_util_test.py
index b8bc9960915..567082357eb 100644
--- a/tensorflow/dtensor/python/tests/mesh_util_test.py
+++ b/tensorflow/dtensor/python/tests/mesh_util_test.py
@@ -156,6 +156,27 @@ class MeshUtilTest(test_util.DTensorBaseTest):
             mesh_dims=[('x', 1)],
             local_devices=['CPU:0'])
 
+  def test_create_distributed_mesh_tpu_with_local_devices_raises_error(self):
+    self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
+    with mock.patch.object(accelerator_util, 'is_initialized', return_value=True):
+      with self.assertRaisesRegex(
+          ValueError, 'Do not specify devices for TPU meshes'
+      ):
+        mesh_util.create_distributed_mesh(
+            mesh_name='tpu_mesh',
+            mesh_dims=[('x', 1)],
+            local_devices=['TPU:0'],
+        )
+      with self.assertRaisesRegex(
+          ValueError, 'Do not specify devices for TPU meshes'
+      ):
+        mesh_util.create_distributed_mesh(
+            mesh_name='tpu_mesh',
+            mesh_dims=[('x', 1)],
+            local_devices=['TPU:0'],
+            device_type='TPU',
+        )
+
   def test_single_client_mesh_creation_wrong_shape(self):
     self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
     num_devices = len(test_util.list_local_logical_devices('CPU'))

From 975d9b6da5742f6664c55ec2c1ea3e37ae0f5a39 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:47:12 -0700
Subject: [PATCH 0539/1000] Automated Code Change

PiperOrigin-RevId: 997669186
---
 tensorflow/core/example/feature_util.h | 29 +++++++++++---------------
 1 file changed, 12 insertions(+), 17 deletions(-)

diff --git a/tensorflow/core/example/feature_util.h b/tensorflow/core/example/feature_util.h
index 2294ed9e4ec..24c1193baea 100644
--- a/tensorflow/core/example/feature_util.h
+++ b/tensorflow/core/example/feature_util.h
@@ -178,24 +178,21 @@ template <typename ValueType, class Enable = void>
 struct FeatureTrait;
 
 template <typename ValueType>
-struct FeatureTrait<ValueType, typename std::enable_if<
-                                   std::is_integral<ValueType>::value>::type> {
+struct FeatureTrait<ValueType,
+                    std::enable_if_t<std::is_integral_v<ValueType>>> {
   using Type = protobuf_int64;
 };
 
 template <typename ValueType>
-struct FeatureTrait<
-    ValueType,
-    typename std::enable_if<std::is_floating_point<ValueType>::value>::type> {
+struct FeatureTrait<ValueType,
+                    std::enable_if_t<std::is_floating_point_v<ValueType>>> {
   using Type = float;
 };
 
 template <typename T>
-struct is_string
-    : public std::integral_constant<
-          bool,
-          std::is_same<char*, typename std::decay<T>::type>::value ||
-              std::is_same<const char*, typename std::decay<T>::type>::value> {
+struct is_string : public std::integral_constant<
+                       bool, std::is_same_v<char*, std::decay_t<T>> ||
+                                 std::is_same_v<const char*, std::decay_t<T>>> {
 };
 
 template <>
@@ -208,15 +205,14 @@ template <>
 struct is_string<tstring> : std::true_type {};
 
 template <typename ValueType>
-struct FeatureTrait<
-    ValueType, typename std::enable_if<is_string<ValueType>::value>::type> {
+struct FeatureTrait<ValueType, std::enable_if_t<is_string<ValueType>::value>> {
   using Type = std::string;
 };
 
 // Port of the C++20 `requires` expressions.
 template <typename... T, typename F>
 constexpr bool Requires(F) {
-  return std::is_invocable<F, T...>::value;
+  return std::is_invocable_v<F, T...>;
 }
 
 struct NoneSuch {};
@@ -281,8 +277,8 @@ struct TypeHasFeatures<Features> : std::true_type {};
 // A family of template functions to return mutable Features proto from a
 // container proto. Supported ProtoTypes: SequenceExample, Example, Features.
 template <typename ProtoType>
-typename std::enable_if<TypeHasFeatures<ProtoType>::value, Features*>::type
-GetFeatures(ProtoType* proto);
+std::enable_if_t<TypeHasFeatures<ProtoType>::value, Features*> GetFeatures(
+    ProtoType* proto);
 
 template <>
 Features* GetFeatures<Features>(Features* proto);
@@ -292,8 +288,7 @@ template <>
 Features* GetFeatures<SequenceExample>(SequenceExample* proto);
 
 template <typename ProtoType>
-typename std::enable_if<TypeHasFeatures<ProtoType>::value,
-                        const Features&>::type
+std::enable_if_t<TypeHasFeatures<ProtoType>::value, const Features&>
 GetFeatures(const ProtoType& proto);
 
 template <>

From e8bc688be68a89b613047aa3113889d938f769ef Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:50:19 -0700
Subject: [PATCH 0540/1000] Automated Code Change

PiperOrigin-RevId: 997670046
---
 tensorflow/cc/framework/scope.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/cc/framework/scope.cc b/tensorflow/cc/framework/scope.cc
index 88b2c24f66b..705af574234 100644
--- a/tensorflow/cc/framework/scope.cc
+++ b/tensorflow/cc/framework/scope.cc
@@ -29,7 +29,7 @@ Scope::Scope(Impl* impl) : impl_(impl) {}
 
 Scope::Scope(const Scope& other) : impl_(new Impl(*other.impl())) {}
 
-Scope::~Scope() {}
+Scope::~Scope() = default;
 
 Scope& Scope::operator=(const Scope& other) {
   // We can't copy Impls because of the const members, use copy ctor instead

From 3efc9d7ca76d3d7d99330b26a405a56b61879021 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 18:30:21 +0530
Subject: [PATCH 0541/1000] Improve context manager for default mesh handling

Refactor context manager to ensure previous default mesh is restored even if an exception occurs.
---
 tensorflow/dtensor/python/dtensor_device.py | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/tensorflow/dtensor/python/dtensor_device.py b/tensorflow/dtensor/python/dtensor_device.py
index 5c6dff23052..5ee21f263c9 100644
--- a/tensorflow/dtensor/python/dtensor_device.py
+++ b/tensorflow/dtensor/python/dtensor_device.py
@@ -360,13 +360,19 @@ class DTensorDevice(object):
         self._device_info,
         mesh.to_string().encode("utf-8"))
     self._current_default_mesh = mesh
-    yield
-    _pywrap_dtensor_device.ExperimentalClearDefaultMesh(self._device_info)
-    if previous_default:
-      _pywrap_dtensor_device.ExperimentalSetDefaultMesh(
-          self._device_info,
-          previous_default.to_string().encode("utf-8"))
-    self._current_default_mesh = previous_default
+    try:
+      yield
+    finally:
+      # Always restore the previous default mesh, even if an exception was
+      # raised inside the `with` block. Otherwise both the native default
+      # mesh and `_current_default_mesh` are left stuck on `mesh` for the
+      # rest of the program.
+      _pywrap_dtensor_device.ExperimentalClearDefaultMesh(self._device_info)
+      if previous_default:
+        _pywrap_dtensor_device.ExperimentalSetDefaultMesh(
+            self._device_info,
+            previous_default.to_string().encode("utf-8"))
+      self._current_default_mesh = previous_default
 
   @contextlib.contextmanager
   def _default_layout(self, layout: layout_lib.Layout):

From 174ffeb936b924d3cdb1c17965c48fd39de10f1f Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 18:33:38 +0530
Subject: [PATCH 0542/1000] Update input_util.py

---
 tensorflow/dtensor/python/input_util.py | 610 ++++++++++++++++++++++++
 1 file changed, 610 insertions(+)

diff --git a/tensorflow/dtensor/python/input_util.py b/tensorflow/dtensor/python/input_util.py
index 280fcdabe7d..488870c2169 100644
--- a/tensorflow/dtensor/python/input_util.py
+++ b/tensorflow/dtensor/python/input_util.py
@@ -284,6 +284,616 @@ def _shard_counts(layout: layout_lib.Layout,
   return shard_counts
 
 
+def _index_matrix(layout: layout_lib.Layout,
+                  elem_spec: tensor_spec.TensorSpec) -> tensor.Tensor:
+  """Computes a utility matrix to derive device-based slice offsets.
+
+  This function builds a matrix of shape `[mesh.rank, layout.rank]` for each
+  dataset element. This matrix can be used to slice the DTensor components
+  returned by the iterator according to the local device that component is to be
+  placed on. This can be done by multiplying the device offsets of shape
+  `[1, mesh.rank]` with this index matrix to get a `[1, layout.rank]` shape
+  tensor containing the slice offsets.
+
+  Note: the index on the batch dim is always 0 since sharding on the batch
+  dimension is handled by either tf.data.Dataset's shard transformation (in the
+  single-client case) or tf.data service's distribute function (in the
+  multi-client case). If there is no sharding on the batch dimension (or any
+  other dimension), the slice index remains 0.
+
+  Args:
+    layout: the layout of the dataset element.
+    elem_spec: the spec of the *batched* dataset element, i.e. its shape must
+      already include the leading batch dimension and have the same rank as
+      `layout`.
+
+  Returns:
+    The index matrix as a tensor.
+  """
+  matrix = []
+  for dim in layout.mesh.dim_names:
+    row = [0]
+    for layout_idx, spec in enumerate(layout.sharding_specs[1:]):
+      if spec == layout_lib.UNSHARDED or spec != dim:
+        row.append(0)
+      else:
+        # `layout_idx` enumerates `layout.sharding_specs[1:]`, i.e. it is
+        # zero-based over the non-batch dims, while `elem_spec.shape` (which
+        # includes the batch dim at index 0) needs the matching tensor axis,
+        # which is one higher.
+        row.append(
+            elem_spec.shape[layout_idx + 1] // layout.mesh.dim_size(dim))
+    matrix.append(row)
+
+  return constant_op.constant(matrix, dtype=dtypes.int32)
+
+
+def _pack_iterator_resource_dtensor(
+    datasets: List[Tuple[int, data_types.DatasetV2]],
+    layouts: Any,
+    mesh: layout_lib.Mesh,
+    num_local_devices_per_replica: int):
+  """Creates a DTensor iterator resource for the per-replica datasets.
+
+  Given a list of replica ID to tf.data.Dataset mappings, this function creates
+  iterators for each device and then packs the underlying iterator resource
+  tensors into a single DTensor. This resource tensor is used by the
+  IteratorGetNext op to retrieve the next element in the dataset.
+
+  Args:
+    datasets: a list of tuples of each unique local replica ID to the dataset
+      object whose elements will be placed on the devices corresponding to that
+      replica.
+    layouts: a structure of DTensor layouts to be applied to the elements
+      returned by the underlying iterators. This can be a single layout or
+      (possibly nested) tuples or dictionaries of layouts, and the structure
+      must match the structure of the iterator elements.
+    mesh: the DTensor mesh to place the iterator batches on.
+    num_local_devices_per_replica: the number of devices in each data-parallel
+      replica.
+
+  Returns:
+    A DTensor of the underlying iterator resource tensors.
+  """
+  host_mesh_devices = mesh.host_mesh().local_devices()
+  device_idx = 0
+
+  iterators = []
+  for _, dataset in datasets:
+    for idx in range(num_local_devices_per_replica):
+      with ops.device_v2(host_mesh_devices[device_idx]):
+        device_dataset = dataset.shard(
+            num_shards=num_local_devices_per_replica, index=idx)
+        iterators.append(iter(device_dataset))
+      device_idx += 1
+
+  if device_idx != len(host_mesh_devices):
+    raise ValueError(
+        'The `datasets` argument does not have the correct number of'
+        f' underlying datasets, found {device_idx} but expected'
+        f' {len(host_mesh_devices)}.')
+
+  host_layouts = nest.map_structure(
+      lambda l: layout_lib.Layout(l.sharding_specs, mesh.host_mesh()), layouts)
+
+  # Pack the iterator resource tensors into a replicated 0-dimensional DTensor
+  # and set the element layouts.
+  iterator_resources = [it._iterator_resource for it in iterators]  # pylint: disable=protected-access
+  d_iterator_resource = api.pack(
+      iterator_resources,
+      layout_lib.Layout.replicated(mesh=mesh.host_mesh(), rank=0))
+  api._dtensor_device().set_iterator_element_layouts(  # pylint: disable=protected-access
+      d_iterator_resource, nest.flatten(host_layouts))
+
+  return d_iterator_resource
+
+
+@tf_export('experimental.dtensor.DTensorDataset', v1=[])
+class DTensorDataset(dataset_ops.UnaryUnchangedStructureDataset):
+  """A dataset of DTensors.
+
+  DTensorDataset encapsulates a `tf.data.Dataset` whose elements are
+  automatically packed and returned as DTensors based on a given mesh and
+  layouts.
+  """
+
+  def __init__(self,
+               dataset: data_types.DatasetV2,
+               *,
+               mesh: layout_lib.Mesh,
+               layouts: Any,
+               global_batch_size: int,
+               dataset_already_batched: bool = False,
+               batch_dim: Optional[str] = None,
+               prefetch: Optional[int] = None,
+               tf_data_service_config: Optional[TFDataServiceConfig] = None):
+    """Creates a DTensorDataset.
+
+    DTensorDataset automatically handles distribution of the dataset elements to
+    each client's devices. It can be used to create an iterator that returns
+    DTensors of the input data on each iteration.
+
+    DTensorDataset works best with unbatched datasets. It takes the mesh and the
+    provided layouts to automatically calculate how to batch the input locally
+    for each replica.
+
+    If the provided dataset is already batched according to the per-replica
+    batch size, then `dataset_already_batched` must be set and DTensorDataset
+    will check that the batch size is consistent with the intended
+    `global_batch_size` using the layout information. Each replica receives a
+    separate slice of the global batch, thus the per-replica batch size can be
+    computed as the global batch size divided by the number of model replicas.
+    For a DTensor mesh, the number of replicas is equal to the size of the
+    mesh's batch dimension.
+
+    Note: `tf.experimental.dtensor.DTensorDataset` instances do *not* implement
+    the full interface of `tf.data.Dataset`. It only supports two usages we will
+    mention below: iteration and `element_spec`. We don't support any other APIs
+    to transform or inspect the dataset.
+
+    TODO(b/223275517): add support for input datasets that are already batched
+    to the global batch size.
+
+    Args:
+      dataset: a `tf.data.Dataset` object.
+      mesh: the DTensor mesh to place the dataset batches on.
+      layouts: a structure of DTensor layouts to be applied to the input dataset
+        values. This can be a single layout or (possibly nested) tuples or
+        dictionaries of layouts, and the structure must match the structure of
+        the dataset. Either all or none of the layouts should be sharded on the
+        batch dimension; having only a subset of layouts batch sharded will not
+        work and raises a ValueError.
+      global_batch_size: the desired global batch size.
+      dataset_already_batched: must be set only if the dataset is already
+        batched to the per-replica batch size. The batched dataset must have
+        `drop_remainder=True` set since DTensor requires static shapes for
+        slicing the input tensors.
+      batch_dim: the mesh dimension on which the input's batch dimension is
+        sharded. Set to None if the input layouts do not shard on the batch
+        dimension.
+      prefetch: number of batches to prefetch using Dataset.prefetch.
+      tf_data_service_config: if operating in multi-client mode, this config
+        specifies the tf.data service configuration to use.
+
+    Raises:
+      ValueError: on any of the following situations,
+        1. if the structures and ranks of layouts and the dataset do not match.
+        2. if the shapes in the dataset's spec are not fully defined.
+        3. if batch_dim is specified and all layouts are not batch-sharded.
+        4. if per_replica_batch_size is specified for an already batched Dataset
+           but it does not match the expected per-replica size based on the
+           provided mesh.
+      TypeError: if type of structures of layouts and the dataset do not match.
+    """
+    super().__init__(dataset, dataset_ops.to_variant(dataset))
+
+    # TODO(b/271162918): fix multi-client use case.
+    if tf_data_service_config is not None:
+      raise NotImplementedError(
+          'Multi-client DTensorDataset is currently not supported.'
+          ' Check b/271162918.')
+
+    self._mesh = mesh
+    self._layouts = layouts
+    self._batch_dim = batch_dim
+    self._prefetch = prefetch
+    self._tf_data_service_config = tf_data_service_config
+
+    nest.assert_same_structure(dataset.element_spec, layouts)
+    flattened_layouts = nest.flatten(layouts)
+    flattened_elem_spec = nest.flatten(dataset.element_spec)
+
+    if batch_dim:
+      self.num_global_replicas = mesh.dim_size(batch_dim)
+      self._local_replica_ids = list(
+          dict.fromkeys(
+              [loc[batch_dim] for loc in mesh.local_device_locations()]))
+
+      for layout in flattened_layouts:
+        if batch_dim != layout.sharding_specs[0]:
+          raise ValueError(
+              ('batch_dim %s was specified but at least one layout did not '
+               'contain it: %s') % (batch_dim, layout))
+    else:
+      # Only one replica since there is no sharding on the batch dimension.
+      self.num_global_replicas = 1
+      self._local_replica_ids = [0]
+
+    # Validate layout and element spec compatibility, and raise ValueError if
+    # invalid.
+    _validate_input(
+        flattened_layouts,
+        flattened_elem_spec,
+        dataset_already_batched=dataset_already_batched)
+
+    expected_batch_size = global_batch_size // self.num_global_replicas
+    if not dataset_already_batched:
+      self._batched_dataset = dataset.batch(
+          expected_batch_size, drop_remainder=True)
+    else:
+      per_replica_batch_size = flattened_elem_spec[0].shape.as_list()[0]
+      if per_replica_batch_size != expected_batch_size:
+        raise ValueError(
+            ('per_replica_batch_size does not matched expected size based on '
+             'the mesh, got %d but expected %d.') %
+            (per_replica_batch_size, expected_batch_size))
+      self._batched_dataset = dataset
+
+    # Construct a global element spec of the dataset.
+    flattened_global_elem_spec = []
+    batch_tensor_shape = tensor_shape.as_shape([global_batch_size])
+    for elem_spec in nest.flatten(self._batched_dataset.element_spec):
+      new_elem_spec = tensor_spec.TensorSpec(
+          shape=operator.concat(batch_tensor_shape, elem_spec.shape[1:]),
+          dtype=elem_spec.dtype,
+          name=elem_spec.name)
+      flattened_global_elem_spec.append(new_elem_spec)
+    self._global_element_spec = nest.pack_sequence_as(
+        dataset.element_spec, flattened_global_elem_spec)
+
+    num_global_devices_per_replica = config.num_global_devices(
+        mesh.device_type()) // self.num_global_replicas
+    self._num_local_replicas = len(self._local_replica_ids)
+    self._num_local_devices_per_replica = mesh.num_local_devices(
+    ) // self._num_local_replicas
+    # The number of clients each replica is split over.
+    self._num_clients_per_replica = (
+        num_global_devices_per_replica // self._num_local_devices_per_replica)
+    # In the case where a replica is split across multiple clients, an offset
+    # needs to be added to the index used by the partitioning logic such that
+    # the local devices on that client can be correctly matched to slices of the
+    # input tensor(s). If replicas are wholly contained within a client, then
+    # this offset is always 0.
+    self._partition_offset = (config.client_id() % self._num_clients_per_replica
+                             ) * self._num_local_devices_per_replica
+
+    # Helper data structures used in partitioning the dataset tensors.
+    self._all_shard_counts = [
+        _shard_counts(layout, batch_dim) for layout in flattened_layouts
+    ]
+    # `_index_matrix` requires the *batched* element spec (shape rank must
+    # match `layout.rank`, including the leading batch dim). `flattened_elem_spec`
+    # is `dataset.element_spec`, which only has that rank when the dataset was
+    # already batched; use `self._batched_dataset.element_spec` instead so this
+    # is correct for both `dataset_already_batched` values.
+    flattened_batched_elem_spec = nest.flatten(self._batched_dataset.element_spec)
+    self._index_matrices = [
+        _index_matrix(layout, elem_spec)
+        for layout, elem_spec in zip(flattened_layouts,
+                                     flattened_batched_elem_spec)
+    ]
+
+  def __iter__(self):
+    datasets: List[Tuple[int, data_types.DatasetV2]] = []
+
+    # Start with the batched the dataset.
+    local_dataset = self._batched_dataset
+
+    if self._batch_dim is not None:
+      if self._num_clients_per_replica > 1:
+        # If a replica is split over multiple clients then each batch needs to
+        # be repeated before distribution as many times as there are clients
+        # corresponding to that replica.
+        local_dataset = self._repeat_batch(local_dataset,
+                                           self._num_clients_per_replica)
+        sharding_policy = data_service_ops.ShardingPolicy.DATA
+      else:
+        # Replicas are unique to each client, so FILE based sharding can be used
+        # which is more performant since each worker does not need to read the
+        # entire dataset.
+        sharding_policy = data_service_ops.ShardingPolicy.FILE
+    else:
+      # No batch dimension sharding specified so disable dataset sharding during
+      # the distribute step.
+      sharding_policy = data_service_ops.ShardingPolicy.OFF
+
+    # Apply distribution here (if specified) so all remaining transformations
+    # are executed locally.
+    if self._tf_data_service_config is not None:
+      local_dataset = local_dataset.apply(
+          data_service_ops.distribute(
+              processing_mode=sharding_policy,
+              service=self._tf_data_service_config.dispatcher_address,
+              job_name=f'{self._tf_data_service_config.job_name}_{config.client_id()}',
+              target_workers='LOCAL'))
+
+    for local_replica_idx, replica_id in enumerate(self._local_replica_ids):
+      # Select the shard for the corresponding replica.
+      dataset = distribute._AutoShardDataset(
+          local_dataset,
+          num_workers=self._num_local_replicas,
+          index=local_replica_idx,
+          num_replicas=self.num_global_replicas)
+
+      # Repeat each batch for each local device in the replica.
+      dataset = self._repeat_batch(dataset, self._num_local_devices_per_replica)
+
+      # Slice each shard further for all non-batch dim shards. If there is no
+      # non-batch dim sharding, this slice is essentially a no-op.
+      dataset = self._partition(dataset)
+
+      # Apply prefetch as the last step. Since each batch is repeated, the
+      # number of elements to prefetch has to be scaled by the same size.
+      if self._prefetch is not None:
+        dataset = dataset.prefetch(
+            self._prefetch * self._num_local_devices_per_replica)
+
+      datasets.append((replica_id, dataset))
+
+    # Convert the datasets into iterators placed on the host.
+    d_iterator_resource = _pack_iterator_resource_dtensor(
+        datasets=datasets,
+        layouts=self._layouts,
+        mesh=self._mesh,
+        num_local_devices_per_replica=self._num_local_devices_per_replica)
+
+    return _DTensorIterator(
+        dtensor_components=(d_iterator_resource,),
+        global_element_spec=self._global_element_spec,
+        layouts=self._layouts)
+
+  def _repeat_batch(self, dataset, repeats):
+    if repeats == 1:
+      # Remove this shortcut if tf.data can optimize this away.
+      return dataset
+
+    def repeat(*x):
+      return dataset_ops.DatasetV2.from_tensors(x).repeat(repeats)
+
+    return dataset.flat_map(repeat)
+
+  def _partition(self, dataset):
+    """Slices each dataset element on any sharded non-batch dimension."""
+    if self._num_local_devices_per_replica == 1 and self._partition_offset == 0:
+      # Remove this shortcut if tf.data can optimize this away.
+      return dataset
+
+    # TODO(b/223275517): decouple from self and make testable.
+    def slice_batch(index, batch):
+      flattened_batch = nest.flatten(batch)
+      flattened_output = []
+
+      norm_index = math_ops.cast(
+          index % self._num_local_devices_per_replica, dtype=dtypes.int32)
+      norm_index += self._partition_offset
+      coords = self._mesh.coords(norm_index)
+      coords = array_ops.reshape(coords, (1, -1))
+
+      for element, shard_counts, idx_matrix in zip(flattened_batch,
+                                                   self._all_shard_counts,
+                                                   self._index_matrices):
+        indexes = math_ops.matmul(coords, idx_matrix)
+        start = array_ops.reshape(indexes, (-1,))
+        size = array_ops.shape_v2(
+            element, out_type=dtypes.int32) // shard_counts
+        flattened_output.append(
+            array_ops.slice(element, begin=start, size=size))
+
+      return nest.pack_sequence_as(batch, flattened_output)
+
+    enumerated_dataset = dataset.enumerate()
+    partitioned_dataset = enumerated_dataset.map(slice_batch)
+    return partitioned_dataset
+
+  @property
+  def element_spec(self):
+    return self._global_element_specfrom tensorflow.python.data.experimental.ops import distribute
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import iterator_ops
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.types import data as data_types
+from tensorflow.python.util import nest
+from tensorflow.python.util.tf_export import tf_export
+
+
+@dataclasses.dataclass
+class TFDataServiceConfig:
+  """Specifies the tf.data service configuration to use.
+
+  Attributes:
+    dispatcher_address: a string specifying the address of the tf.data service
+      dispatcher server.
+    job_name: a non-empty string identifying the shared job that will be created
+      on tf.data service to process this dataset.
+  """
+  dispatcher_address: str
+  job_name: str
+
+
+# TODO(b/223275517): Add support for get_next_as_optional().
+class _DTensorIterator(iterator_ops.OwnedIterator):
+  """An iterator for a tf.data.Dataset distributed using DTensor.
+
+  DTensorIterator encapsulates multiple underlying dataset iterators. It handles
+  retrieving the tensors to be placed on each underlying device and then uses
+  the 'pack' operation to create and return a DTensor. Thus users need only
+  interact with a single DTensorIterator to automatically distribute dataset
+  tensors onto devices.
+  """
+
+  def __init__(
+      self,
+      dtensor_components: Tuple[tensor.Tensor],
+      global_element_spec: tensor_spec.TensorSpec,
+      layouts: Any):
+    """Initializes a distributed iterator for DTensor datasets.
+
+    This iterator encapsulates tf.data iterators for the underlying devices, and
+    treats it as a packed DTensor of iterator resource tensors.
+
+    Args:
+      dtensor_components: a tuple containing the underlying iterator resources
+        packed into a DTensor. This is expected to be a tuple with a single
+        element.
+      global_element_spec: the underlying dataset's element spec from a global
+        view.
+      layouts: a structure of DTensor layouts to be applied to the elements
+        returned by the underlying iterators. This can be a single layout or
+        (possibly nested) tuples or dictionaries of layouts, and the structure
+        must match the structure of the iterator elements.
+    """
+    # dtensor_components is expected to be a single-element tuple.
+    [self._iterator_resource_dtensor] = dtensor_components
+    self._global_element_spec = global_element_spec
+    self._layouts = layouts
+    self._layouts_str = nest.map_structure(
+        lambda layout: layout.to_string(), layouts)
+
+    super().__init__(
+        components=dtensor_components, element_spec=global_element_spec)
+
+  def __next__(self):
+    try:
+      # IteratorGetNext will return a DTensor on the host, so move it to the
+      # device mesh. If the dataset layouts are on the host mesh itself, this
+      # is handled by DTensor as a no-op.
+      host_elem = self._next_internal()
+      context.async_wait()
+      device_elem = nest.map_structure(
+          api.copy_to_mesh, host_elem, self._layouts)
+      context.async_wait()
+      return device_elem
+    except errors.OutOfRangeError as e:
+      # Match TF2 eager executor behavior by raising StopIteration when iterator
+      # is out of range.
+      if context.executing_eagerly():
+        raise StopIteration from e
+      else:
+        raise e
+
+  @property
+  def _type_spec(self):
+    return _DTensorIteratorSpec(self._global_element_spec, self._layouts_str)
+
+
+class _DTensorIteratorSpec(iterator_ops.IteratorSpec):
+  """Type specification for `_DTensorIterator`."""
+
+  __slots__ = ['_global_element_spec', '_layouts_str']
+
+  def __init__(
+      self, global_element_spec: tensor_spec.TensorSpec, layouts_str: Any):
+    super().__init__(global_element_spec)
+    self._global_element_spec = global_element_spec
+    self._layouts_str = layouts_str
+
+  @property
+  def value_type(self):
+    return _DTensorIterator
+
+  def _serialize(self):
+    return (self._global_element_spec, self._layouts_str)
+
+  @property
+  def _component_specs(self):
+    return (tensor_spec.TensorSpec([], dtypes.resource),)
+
+  def _to_components(self, value):
+    return (value._iterator_resource_dtensor,)  # pylint: disable=protected-access
+
+  def _from_components(self, components):
+    layouts = nest.map_structure(
+        layout_lib.Layout.from_string, self._layouts_str)
+    return _DTensorIterator(
+        dtensor_components=components,
+        global_element_spec=self._global_element_spec,
+        layouts=layouts)
+
+  @classmethod
+  def from_value(cls, value):
+    return cls(value._global_element_spec, value._layouts_str)  # pylint: disable=protected-access
+
+
+def _validate_input(flattened_layouts: Sequence[layout_lib.Layout],
+                    flattened_elem_spec: Sequence[tensor_spec.TensorSpec],
+                    dataset_already_batched: bool):
+  """Checks that the dataset's layouts and element specs are compatible.
+
+  Args:
+    flattened_layouts: the flattened list of layouts used to distribute the
+      dataset.
+    flattened_elem_spec: the flattened list of element specs used in the
+      dataset's components.
+    dataset_already_batched: whether the dataset to be validated is already
+      batched.
+
+  Raises:
+    ValueError: if the dataset's inputs are incompatible.
+  """
+  if not flattened_elem_spec:
+    raise ValueError(
+        'Expected input element spec of at least one element, was empty.')
+
+  first_elem_shape = flattened_elem_spec[0].shape
+
+  for layout, elem_spec in zip(flattened_layouts, flattened_elem_spec):
+    if elem_spec.shape.rank is None:
+      raise ValueError(
+          'Dataset element shape must have a valid rank, got spec %s.' %
+          elem_spec)
+
+    # Check that layout's rank matches the element's rank. If dataset is not yet
+    # batched, then the layout's rank must be one greater than the element's
+    # rank.
+    expected_rank = elem_spec.shape.rank
+    if not dataset_already_batched:
+      expected_rank += 1
+    if layout.rank != expected_rank:
+      raise ValueError(
+          ('Expected layout with rank %d for element spec %s, got layout %s. '
+           'Check that the dataset is not batched before passing to '
+           'DTensorDataset.') %
+          (expected_rank, elem_spec, layout.sharding_specs))
+
+    if dataset_already_batched:
+      # Check that the batch dimension size of all dataset elements match.
+      batch_dim_size = first_elem_shape.as_list()[0]
+      if batch_dim_size is None:
+        raise ValueError(
+            ('Size of batch dimension of element spec %s is None. Ensure '
+             'drop_remainder=True when batching the dataset.') % elem_spec)
+
+      if elem_spec.shape.as_list()[0] != batch_dim_size:
+        raise ValueError(
+            ('Size of batch dimension of element spec %s does not match '
+             'expected size %d.') % (elem_spec, batch_dim_size))
+
+
+def _shard_counts(layout: layout_lib.Layout,
+                  batch_dim: Optional[str] = None) -> List[int]:
+  """Computes a list of the number of shards in each dimension of the layout.
+
+  The shard counts are used to slice each dataset element. The batch dimension's
+  count is overridden to 1 since we only consider how many shards to make
+  locally (within each local replica). Sharding across clients is handled by
+  either tf.data.Dataset's shard transformation (in the single-client case) or
+  tf.data service's distribute function (in the multi-client case).
+
+  Args:
+    layout: the layout to compute the shard counts for.
+    batch_dim: the name of the batch dimension of the layout, if present.
+
+  Returns:
+    A list of shard counts, one element per dimension of the layout.
+  """
+  shard_counts = []
+  for spec in layout.sharding_specs:
+    if spec in (batch_dim, layout_lib.UNSHARDED):
+      shard_counts.append(1)
+    else:
+      shard_counts.append(layout.mesh.dim_size(spec))
+  return shard_counts
+
+
 def _index_matrix(layout: layout_lib.Layout,
                   elem_spec: tensor_spec.TensorSpec) -> tensor.Tensor:
   """Computes a utility matrix to derive device-based slice offsets.

From 564a0844fe9b0c0f2c2614bf986038280423dbac Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 18:38:32 +0530
Subject: [PATCH 0543/1000] Fix comparison logic for ring_bounds validation

---
 tensorflow/dtensor/python/tpu_util.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/dtensor/python/tpu_util.py b/tensorflow/dtensor/python/tpu_util.py
index fdfbaec79f1..24bea30bccb 100644
--- a/tensorflow/dtensor/python/tpu_util.py
+++ b/tensorflow/dtensor/python/tpu_util.py
@@ -445,7 +445,7 @@ def _enumerate_core_locations(
 
   # Reorder ring_bounds and validate it's element-wise >= host_bounds.
   ring_bounds = [ring_bounds[i] for i in axes]  # pyrefly: ignore[bad-index]
-  if ring_bounds < host_bounds:
+  if any(r < h for r, h in zip(ring_bounds, host_bounds)):
     raise ValueError(
         "ring_bounds %s should be >= host_bounds %s"
         % (ring_bounds, host_bounds)
@@ -713,7 +713,7 @@ def create_tpu_mesh(
     ring_bounds = topology_shape
   elif len(ring_bounds) != 4:
     raise ValueError("Expected 4 elements in ring_bounds, got %s" % ring_bounds)
-  elif ring_bounds > topology_shape:
+  elif any(r > t for r, t in zip(ring_bounds, topology_shape)):
     raise ValueError(
         "ring_bounds %s should be <= topology sizes %s"
         % (ring_bounds, topology_shape)

From 2fa7723d7c13307ab1480dd9fcd97c6e31cdc38b Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 18:43:11 +0530
Subject: [PATCH 0544/1000] Refactor load context management with stack

---
 tensorflow/python/distribute/load_context.py | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/tensorflow/python/distribute/load_context.py b/tensorflow/python/distribute/load_context.py
index 7680001c2a7..6950c240770 100644
--- a/tensorflow/python/distribute/load_context.py
+++ b/tensorflow/python/distribute/load_context.py
@@ -25,22 +25,25 @@ class LoadContext(threading.local):
 
   def __init__(self):
     super().__init__()
-    self._entered_load_context = []
-    self._load_options = None
+    # Stack of load options, one entry per nested `load_context()` call. The
+    # top of the stack is the currently active load options; popping on exit
+    # restores whatever options (if any) were active before entering, rather
+    # than always resetting to `None`.
+    self._load_options_stack = []
 
   def set_load_options(self, load_options):
-    self._load_options = load_options
-    self._entered_load_context.append(True)
+    self._load_options_stack.append(load_options)
 
   def clear_load_options(self):
-    self._load_options = None
-    self._entered_load_context.pop()
+    self._load_options_stack.pop()
 
   def load_options(self):
-    return self._load_options
+    if not self._load_options_stack:
+      return None
+    return self._load_options_stack[-1]
 
   def in_load_context(self):
-    return self._entered_load_context
+    return bool(self._load_options_stack)
 
 
 _load_context = LoadContext()

From f009a95d85bde5b88e9bcefeac95e1fb9328a253 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 18:45:48 +0530
Subject: [PATCH 0545/1000] Fix typo in experimental_replication_mode check

---
 tensorflow/python/distribute/central_storage_strategy.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/distribute/central_storage_strategy.py b/tensorflow/python/distribute/central_storage_strategy.py
index 72515173904..b85fbf9bf0e 100644
--- a/tensorflow/python/distribute/central_storage_strategy.py
+++ b/tensorflow/python/distribute/central_storage_strategy.py
@@ -98,7 +98,7 @@ class CentralStorageStrategy(distribute_lib.Strategy):
     Returns:
       A "distributed `Dataset`" that the caller can iterate over.
     """
-    if (options and options.experimental_replication_moden ==
+    if (options and options.experimental_replication_mode ==
         distribute_lib.InputReplicationMode.PER_REPLICA):
       raise NotImplementedError(
           'InputReplicationMode.PER_REPLICA '

From 2b08b6af9643c3f2bd71dd8a8c3e9b249897fcb1 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:50:58 -0700
Subject: [PATCH 0546/1000] Automated Code Change

PiperOrigin-RevId: 997670252
---
 .../eager/cluster_function_library_runtime.h                  | 2 +-
 .../distributed_runtime/eager/destroy_tensor_handle_node.h    | 2 +-
 tensorflow/core/distributed_runtime/eager/eager_client.h      | 2 +-
 .../core/distributed_runtime/eager/eager_service_impl_test.cc | 4 ++--
 4 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensorflow/core/distributed_runtime/eager/cluster_function_library_runtime.h b/tensorflow/core/distributed_runtime/eager/cluster_function_library_runtime.h
index 6fb1fc280f0..8e8d34f353f 100644
--- a/tensorflow/core/distributed_runtime/eager/cluster_function_library_runtime.h
+++ b/tensorflow/core/distributed_runtime/eager/cluster_function_library_runtime.h
@@ -44,7 +44,7 @@ class EagerClusterFunctionLibraryRuntime
         ctx_(ctx),
         remote_device_mgr_(remote_device_mgr) {}
 
-  ~EagerClusterFunctionLibraryRuntime() override{};
+  ~EagerClusterFunctionLibraryRuntime() override = default;
 
   // Register a partition (i.e., component function) of a multi-device function
   // on the remote target specified in `options.target`. This should be
diff --git a/tensorflow/core/distributed_runtime/eager/destroy_tensor_handle_node.h b/tensorflow/core/distributed_runtime/eager/destroy_tensor_handle_node.h
index a0991dc601b..dd286a37a0d 100644
--- a/tensorflow/core/distributed_runtime/eager/destroy_tensor_handle_node.h
+++ b/tensorflow/core/distributed_runtime/eager/destroy_tensor_handle_node.h
@@ -39,7 +39,7 @@ class DestroyTensorHandleNode : public tensorflow::AsyncEagerNode {
         eager_client_(std::move(eager_client)),
         ready_(ready) {}
 
-  ~DestroyTensorHandleNode() override {}
+  ~DestroyTensorHandleNode() override = default;
 
   void RunAsync(StatusCallback done) override {
     EnqueueResponse* response = new EnqueueResponse;
diff --git a/tensorflow/core/distributed_runtime/eager/eager_client.h b/tensorflow/core/distributed_runtime/eager/eager_client.h
index a2a3d596bff..b1f2b5c8870 100644
--- a/tensorflow/core/distributed_runtime/eager/eager_client.h
+++ b/tensorflow/core/distributed_runtime/eager/eager_client.h
@@ -86,7 +86,7 @@ class EagerClient : public core::RefCounted {
 // Simple wrapper class that can be used to retrieve EagerClients.
 class EagerClientCache {
  public:
-  virtual ~EagerClientCache() {}
+  virtual ~EagerClientCache() = default;
 
   // If the `target` exists, assign the EagerClient pointer to `client` and
   // increment the refcount of the client. The reference ownership is
diff --git a/tensorflow/core/distributed_runtime/eager/eager_service_impl_test.cc b/tensorflow/core/distributed_runtime/eager/eager_service_impl_test.cc
index 3c2dd7e368d..75415a20ac6 100644
--- a/tensorflow/core/distributed_runtime/eager/eager_service_impl_test.cc
+++ b/tensorflow/core/distributed_runtime/eager/eager_service_impl_test.cc
@@ -77,8 +77,8 @@ class TestEagerServiceImpl : public EagerServiceImpl {
 
 class FakeEagerClient : public EagerClient {
  public:
-  FakeEagerClient() {}
-  ~FakeEagerClient() override {}
+  FakeEagerClient() = default;
+  ~FakeEagerClient() override = default;
 
   void SetServiceImpl(TestEagerServiceImpl* impl) { impl_ = impl; }
 

From 0d1b2fb6049c5b2f158ed35d67103620e403de50 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:51:18 -0700
Subject: [PATCH 0547/1000] Automated Code Change

PiperOrigin-RevId: 997670391
---
 tensorflow/core/util/ctc/ctc_beam_entry.h  | 2 +-
 tensorflow/core/util/ctc/ctc_beam_scorer.h | 2 +-
 tensorflow/core/util/ctc/ctc_beam_search.h | 2 +-
 tensorflow/core/util/ctc/ctc_decoder.h     | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/util/ctc/ctc_beam_entry.h b/tensorflow/core/util/ctc/ctc_beam_entry.h
index c8424c41fc2..543ca0165ea 100644
--- a/tensorflow/core/util/ctc/ctc_beam_entry.h
+++ b/tensorflow/core/util/ctc/ctc_beam_entry.h
@@ -138,7 +138,7 @@ class BeamRoot {
 template <class T, class CTCBeamState = EmptyBeamState>
 class BeamComparer {
  public:
-  virtual ~BeamComparer() {}
+  virtual ~BeamComparer() = default;
   virtual bool operator()(const BeamEntry<T, CTCBeamState>* a,
                           const BeamEntry<T, CTCBeamState>* b) const {
     return a->newp.total > b->newp.total;
diff --git a/tensorflow/core/util/ctc/ctc_beam_scorer.h b/tensorflow/core/util/ctc/ctc_beam_scorer.h
index 1ea370f4117..ec963e72ca9 100644
--- a/tensorflow/core/util/ctc/ctc_beam_scorer.h
+++ b/tensorflow/core/util/ctc/ctc_beam_scorer.h
@@ -37,7 +37,7 @@ namespace ctc {
 template <typename T, typename CTCBeamState>
 class BaseBeamScorer {
  public:
-  virtual ~BaseBeamScorer() {}
+  virtual ~BaseBeamScorer() = default;
   // State initialization.
   virtual void InitializeState(CTCBeamState* root) const {}
   // ExpandState is called when expanding a beam to one of its children.
diff --git a/tensorflow/core/util/ctc/ctc_beam_search.h b/tensorflow/core/util/ctc/ctc_beam_search.h
index 66c98b84643..c93cdf8f7fc 100644
--- a/tensorflow/core/util/ctc/ctc_beam_search.h
+++ b/tensorflow/core/util/ctc/ctc_beam_search.h
@@ -96,7 +96,7 @@ class CTCBeamSearchDecoder : public CTCDecoder<T> {
     Reset();
   }
 
-  ~CTCBeamSearchDecoder() override {}
+  ~CTCBeamSearchDecoder() override = default;
 
   // Run the hibernating beam search algorithm on the given input.
   absl::Status Decode(const typename CTCDecoder<T>::SequenceLength& seq_len,
diff --git a/tensorflow/core/util/ctc/ctc_decoder.h b/tensorflow/core/util/ctc/ctc_decoder.h
index 714b812ede1..523ff99918f 100644
--- a/tensorflow/core/util/ctc/ctc_decoder.h
+++ b/tensorflow/core/util/ctc/ctc_decoder.h
@@ -49,7 +49,7 @@ class CTCDecoder {
         batch_size_(batch_size),
         merge_repeated_(merge_repeated) {}
 
-  virtual ~CTCDecoder() {}
+  virtual ~CTCDecoder() = default;
 
   // Dimensionality of the input/output is expected to be:
   //  - seq_len[b] - b = 0 to batch_size_

From 7cdd38b7bc6e9522f84f518f62397b05dd2bb695 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:51:42 -0700
Subject: [PATCH 0548/1000] Automated Code Change

PiperOrigin-RevId: 997670533
---
 tensorflow/core/transforms/utils/eval_utils.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/transforms/utils/eval_utils.cc b/tensorflow/core/transforms/utils/eval_utils.cc
index ff87fd74fff..d8f9fc7dfa8 100644
--- a/tensorflow/core/transforms/utils/eval_utils.cc
+++ b/tensorflow/core/transforms/utils/eval_utils.cc
@@ -64,7 +64,7 @@ SimpleDevice::SimpleDevice() : DeviceBase(tensorflow::Env::Default()) {
   set_eigen_cpu_device(eigen_device_.get());
 }
 
-SimpleDevice::~SimpleDevice() {}
+SimpleDevice::~SimpleDevice() = default;
 
 tensorflow::Allocator *SimpleDevice::GetAllocator(
     tensorflow::AllocatorAttributes attr) {

From 457d2a1bc4b1b2ab03c297489f9d6f1120eac2b5 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:52:55 -0700
Subject: [PATCH 0549/1000] Automated Code Change

PiperOrigin-RevId: 997671110
---
 tensorflow/core/framework/bounds_check.h      |  4 +-
 tensorflow/core/framework/op_kernel.h         |  2 +-
 tensorflow/core/framework/resource_mgr.h      |  2 +-
 tensorflow/core/framework/tensor.cc           |  4 +-
 tensorflow/core/framework/tensor.h            |  8 +-
 tensorflow/core/framework/tensor_shape.h      |  5 +-
 tensorflow/core/framework/tensor_util.cc      |  6 +-
 tensorflow/core/framework/tensor_util.h       |  9 +-
 tensorflow/core/framework/variant.h           | 59 +++++++------
 .../core/framework/variant_encode_decode.h    | 86 ++++++++-----------
 10 files changed, 84 insertions(+), 101 deletions(-)

diff --git a/tensorflow/core/framework/bounds_check.h b/tensorflow/core/framework/bounds_check.h
index 76e6e6dd058..48c696f8e67 100644
--- a/tensorflow/core/framework/bounds_check.h
+++ b/tensorflow/core/framework/bounds_check.h
@@ -29,7 +29,7 @@ namespace tensorflow {
 template <typename Ta, typename Tb>
 EIGEN_ALWAYS_INLINE EIGEN_DEVICE_FUNC bool FastBoundsCheck(const Ta index,
                                                            const Tb limit) {
-  static_assert(std::is_integral<Ta>::value && std::is_integral<Tb>::value,
+  static_assert(std::is_integral_v<Ta> && std::is_integral_v<Tb>,
                 "FastBoundsCheck can only be used on integer types.");
   typedef typename std::make_unsigned<decltype(index + limit)>::type UIndex;
   return TF_PREDICT_TRUE(static_cast<UIndex>(index) <
@@ -43,7 +43,7 @@ namespace internal {
 // etc).  It does not guarantee any atomicity or barriers.
 template <typename T>
 EIGEN_ALWAYS_INLINE EIGEN_DEVICE_FUNC const T SubtleMustCopy(const T &x) {
-  static_assert(std::is_integral<T>::value,
+  static_assert(std::is_integral_v<T>,
                 "SubtleMustCopy can only be used on integer types.");
   auto *to_x = reinterpret_cast<const volatile T *>(&x);
   return *to_x;
diff --git a/tensorflow/core/framework/op_kernel.h b/tensorflow/core/framework/op_kernel.h
index e90ff055c9c..c38a9f436e4 100644
--- a/tensorflow/core/framework/op_kernel.h
+++ b/tensorflow/core/framework/op_kernel.h
@@ -1675,7 +1675,7 @@ inline TensorValue OpKernelContext::release_output(int index) {
 
 template <typename T>
 T* OpKernelContext::op_device_context() {
-  static_assert(std::is_base_of<DeviceContext, T>::value,
+  static_assert(std::is_base_of_v<DeviceContext, T>,
                 "T is not a subclass of DeviceContext");
   return static_cast<T*>(op_device_context());
 }
diff --git a/tensorflow/core/framework/resource_mgr.h b/tensorflow/core/framework/resource_mgr.h
index e1b67bcf58d..5ba36fa019e 100644
--- a/tensorflow/core/framework/resource_mgr.h
+++ b/tensorflow/core/framework/resource_mgr.h
@@ -649,7 +649,7 @@ class ResourceDeleter {
 
 template <typename T>
 void CheckDeriveFromResourceBase() {
-  static_assert(std::is_base_of<ResourceBase, T>::value,
+  static_assert(std::is_base_of_v<ResourceBase, T>,
                 "T must derive from ResourceBase");
 }
 
diff --git a/tensorflow/core/framework/tensor.cc b/tensorflow/core/framework/tensor.cc
index 3ceb67b735a..00b27ac3d10 100644
--- a/tensorflow/core/framework/tensor.cc
+++ b/tensorflow/core/framework/tensor.cc
@@ -208,7 +208,7 @@ struct Helper {
     }
     port::CopyToArray(in, data);
 
-    if constexpr (std::is_same_v<typename std::remove_cv<T>::type, bool>) {
+    if constexpr (std::is_same_v<std::remove_cv_t<T>, bool>) {
       // Check that contents are valid and not trap representations for bool
       // TODO(tlongeri): do we need this for any other types?
       static constexpr bool true_value = true;
@@ -701,7 +701,7 @@ TensorBuffer* FromProtoField(Allocator* a, const TensorProto& in, int64_t n) {
       std::copy_n(begin, n, data);
     } else {
       std::copy_n(begin, in_n, data);
-      if (std::is_trivially_copyable<T>::value) {
+      if (std::is_trivially_copyable_v<T>) {
         const T last = *(data + in_n - 1);
         std::fill_n(data + in_n, n - in_n, last);
       } else {
diff --git a/tensorflow/core/framework/tensor.h b/tensorflow/core/framework/tensor.h
index 90c3f288041..785daa41d6e 100644
--- a/tensorflow/core/framework/tensor.h
+++ b/tensorflow/core/framework/tensor.h
@@ -275,8 +275,8 @@ class Tensor {
 
   // Explicitly delete constructor that take a pointer (except char*)
   // so that the pointer doesn't get implicitly cast to bool.
-  template <typename T, typename std::enable_if<!std::is_same<T, char>::value,
-                                                T>::type* = nullptr>
+  template <typename T,
+            std::enable_if_t<!std::is_same_v<T, char>, T>* = nullptr>
   explicit Tensor(T* t) = delete;
 
   ~Tensor();
@@ -988,7 +988,7 @@ typename TTypes<T, NDIMS>::UnalignedConstTensor Tensor::unaligned_shaped(
 template <typename T>
 typename TTypes<T>::Scalar Tensor::scalar() {
   static_assert(
-      !std::is_same<T, std::string>::value,
+      !std::is_same_v<T, std::string>,
       "std::string is no longer a scalar type, use tensorflow::tstring");
   CheckIsAlignedAndSingleElement();
   return typename TTypes<T>::Scalar(base<T>());
@@ -997,7 +997,7 @@ typename TTypes<T>::Scalar Tensor::scalar() {
 template <typename T>
 typename TTypes<T>::ConstScalar Tensor::scalar() const {
   static_assert(
-      !std::is_same<T, std::string>::value,
+      !std::is_same_v<T, std::string>,
       "std::string is no longer a scalar type, use tensorflow::tstring");
   CheckIsAlignedAndSingleElement();
   return typename TTypes<T>::ConstScalar(base<T>());
diff --git a/tensorflow/core/framework/tensor_shape.h b/tensorflow/core/framework/tensor_shape.h
index bb0104e8369..9068b3d9268 100644
--- a/tensorflow/core/framework/tensor_shape.h
+++ b/tensorflow/core/framework/tensor_shape.h
@@ -330,9 +330,8 @@ class TensorShapeBase : public TensorShapeRep {
   absl::Status InitDims(absl::Span<const int64_t> dim_sizes);
 
   // True for PartialTensorShape, false for TensorShape
-  static constexpr bool kIsPartial =
-      std::is_same<Shape, PartialTensorShape>::value;
-  static_assert(kIsPartial || std::is_same<Shape, TensorShape>::value,
+  static constexpr bool kIsPartial = std::is_same_v<Shape, PartialTensorShape>;
+  static_assert(kIsPartial || std::is_same_v<Shape, TensorShape>,
                 "Shape is neither TensorShape nor PartialTensorShape");
 
   // Used by AddDim and MakeShapeHelper.  Does no error checking.
diff --git a/tensorflow/core/framework/tensor_util.cc b/tensorflow/core/framework/tensor_util.cc
index 081bbb080c8..28d2418c860 100644
--- a/tensorflow/core/framework/tensor_util.cc
+++ b/tensorflow/core/framework/tensor_util.cc
@@ -279,14 +279,12 @@ inline bool PackedValuesNotEqual(const std::complex<RealType>& a,
 }
 
 // Integer can't be negative zero.
-template <typename T,
-          typename std::enable_if<std::is_integral<T>::value>::type* = nullptr>
+template <typename T, std::enable_if_t<std::is_integral_v<T>>* = nullptr>
 static bool IsNegativeZero(T value) {
   return false;
 }
 
-template <typename T,
-          typename std::enable_if<!std::is_integral<T>::value>::type* = nullptr>
+template <typename T, std::enable_if_t<!std::is_integral_v<T>>* = nullptr>
 static bool IsNegativeZero(T value) {
   return value == T(0) && std::signbit(value);
 }
diff --git a/tensorflow/core/framework/tensor_util.h b/tensorflow/core/framework/tensor_util.h
index 7b9aaba5577..b110e67cda4 100644
--- a/tensorflow/core/framework/tensor_util.h
+++ b/tensorflow/core/framework/tensor_util.h
@@ -265,8 +265,7 @@ class TensorProtoHelper<std::string> : public std::true_type {
 };
 
 template <typename Type, typename IterType>
-typename std::enable_if<internal::TensorProtoHelper<Type>::value,
-                        TensorProto>::type
+std::enable_if_t<internal::TensorProtoHelper<Type>::value, TensorProto>
 CreateTensorProto(IterType values_begin, IterType values_end,
                   const size_t values_size,
                   const absl::Span<const size_t> shape) {
@@ -294,8 +293,7 @@ CreateTensorProto(IterType values_begin, IterType values_end,
 // compiler can't infer it:
 //   auto proto = CreateTensorProtoSpan<float>(my_array, shape);
 template <typename Type>
-typename std::enable_if<internal::TensorProtoHelper<Type>::value,
-                        TensorProto>::type
+std::enable_if_t<internal::TensorProtoHelper<Type>::value, TensorProto>
 CreateTensorProtoSpan(const absl::Span<const Type> values,
                       const absl::Span<const size_t> shape) {
   return internal::CreateTensorProto<Type>(values.begin(), values.end(),
@@ -306,8 +304,7 @@ CreateTensorProtoSpan(const absl::Span<const Type> values,
 // which case Type can automatically be inferred:
 //   auto proto = CreateTensorProto(my_vector, shape);
 template <typename Type>
-typename std::enable_if<internal::TensorProtoHelper<Type>::value,
-                        TensorProto>::type
+std::enable_if_t<internal::TensorProtoHelper<Type>::value, TensorProto>
 CreateTensorProto(const std::vector<Type>& values,
                   const absl::Span<const size_t> shape) {
   // This awkward iterator passing is essentially just to support vector<bool>,
diff --git a/tensorflow/core/framework/variant.h b/tensorflow/core/framework/variant.h
index 07038180420..f63daca565c 100644
--- a/tensorflow/core/framework/variant.h
+++ b/tensorflow/core/framework/variant.h
@@ -166,28 +166,28 @@ class Variant {
   // Make sure that the type is CopyConstructible and not a
   // tensorflow::Variant object itself. We want the copy constructor to be
   // chosen for the tensorflow::Variant case.
-  template <typename T, typename VT = typename std::decay<T>::type,
-            typename std::enable_if<!std::is_same<Variant, VT>::value &&
-                                        std::is_move_constructible<VT>::value,
-                                    void>::type* = nullptr>
+  template <typename T, typename VT = std::decay_t<T>,
+            std::enable_if_t<!std::is_same_v<Variant, VT> &&
+                                 std::is_move_constructible_v<VT>,
+                             void>* = nullptr>
   Variant(T&& value);
 
-  template <typename T, typename VT = typename std::decay<T>::type,
-            typename std::enable_if<!std::is_same<Variant, VT>::value &&
-                                        std::is_copy_constructible<VT>::value,
-                                    void>::type* = nullptr>
+  template <typename T, typename VT = std::decay_t<T>,
+            std::enable_if_t<!std::is_same_v<Variant, VT> &&
+                                 std::is_copy_constructible_v<VT>,
+                             void>* = nullptr>
   Variant(const T& value);
 
-  template <typename T, typename VT = typename std::decay<T>::type,
-            typename std::enable_if<!std::is_same<Variant, VT>::value &&
-                                        std::is_copy_constructible<VT>::value,
-                                    void>::type* = nullptr>
+  template <typename T, typename VT = std::decay_t<T>,
+            std::enable_if_t<!std::is_same_v<Variant, VT> &&
+                                 std::is_copy_constructible_v<VT>,
+                             void>* = nullptr>
   Variant& operator=(const T& value);
 
-  template <typename T, typename VT = typename std::decay<T>::type,
-            typename std::enable_if<!std::is_same<Variant, VT>::value &&
-                                        std::is_move_constructible<VT>::value,
-                                    void>::type* = nullptr>
+  template <typename T, typename VT = std::decay_t<T>,
+            std::enable_if_t<!std::is_same_v<Variant, VT> &&
+                                 std::is_move_constructible_v<VT>,
+                             void>* = nullptr>
   Variant& operator=(T&& value);
 
   Variant& operator=(const Variant& rhs) {
@@ -343,8 +343,7 @@ class Variant {
     ~Value() final = default;
 
     TypeIndex TypeId() const final {
-      const TypeIndex value_type_index =
-          TypeIndex::Make<typename std::decay<T>::type>();
+      const TypeIndex value_type_index = TypeIndex::Make<std::decay_t<T>>();
       return value_type_index;
     }
 
@@ -551,25 +550,25 @@ inline Variant::Variant(Variant&& other) noexcept
 }
 
 template <typename T, typename VT,
-          typename std::enable_if<!std::is_same<Variant, VT>::value &&
-                                      std::is_move_constructible<VT>::value,
-                                  void>::type*>
+          std::enable_if_t<!std::is_same_v<Variant, VT> &&
+                               std::is_move_constructible_v<VT>,
+                           void>*>
 inline Variant::Variant(T&& value) : is_inline_(CanInlineType<VT>()) {
   InsertValue<VT>(std::forward<T>(value));
 }
 
 template <typename T, typename VT,
-          typename std::enable_if<!std::is_same<Variant, VT>::value &&
-                                      std::is_copy_constructible<VT>::value,
-                                  void>::type*>
+          std::enable_if_t<!std::is_same_v<Variant, VT> &&
+                               std::is_copy_constructible_v<VT>,
+                           void>*>
 inline Variant::Variant(const T& value) : is_inline_(CanInlineType<VT>()) {
   InsertValue<VT>(value);
 }
 
 template <typename T, typename VT,
-          typename std::enable_if<!std::is_same<Variant, VT>::value &&
-                                      std::is_move_constructible<VT>::value,
-                                  void>::type*>
+          std::enable_if_t<!std::is_same_v<Variant, VT> &&
+                               std::is_move_constructible_v<VT>,
+                           void>*>
 inline Variant& Variant::operator=(T&& value) {
   ResetMemory();
   is_inline_ = CanInlineType<VT>();
@@ -578,9 +577,9 @@ inline Variant& Variant::operator=(T&& value) {
 }
 
 template <typename T, typename VT,
-          typename std::enable_if<!std::is_same<Variant, VT>::value &&
-                                      std::is_copy_constructible<VT>::value,
-                                  void>::type*>
+          std::enable_if_t<!std::is_same_v<Variant, VT> &&
+                               std::is_copy_constructible_v<VT>,
+                           void>*>
 inline Variant& Variant::operator=(const T& value) {
   ResetMemory();
   is_inline_ = CanInlineType<VT>();
diff --git a/tensorflow/core/framework/variant_encode_decode.h b/tensorflow/core/framework/variant_encode_decode.h
index 794840afa1e..c37ae676f95 100644
--- a/tensorflow/core/framework/variant_encode_decode.h
+++ b/tensorflow/core/framework/variant_encode_decode.h
@@ -40,14 +40,11 @@ namespace tensorflow {
 // * A protocol buffer (TypeResolver<T, false, false, true>)
 // * None of the above (TypeResolver<T, false, false, false>)
 //
-template <
-    typename T,
-    bool = std::is_trivially_copyable<typename std::decay<T>::type>::value &&
-           !std::is_pointer<typename std::decay<T>::type>::value,
-    bool =
-        std::is_same<typename std::decay<T>::type, ::tensorflow::Tensor>::value,
-    bool = std::is_base_of<protobuf::MessageLite,
-                           typename std::decay<T>::type>::value>
+template <typename T,
+          bool = std::is_trivially_copyable_v<std::decay_t<T>> &&
+                 !std::is_pointer_v<std::decay_t<T>>,
+          bool = std::is_same_v<std::decay_t<T>, ::tensorflow::Tensor>,
+          bool = std::is_base_of_v<protobuf::MessageLite, std::decay_t<T>>>
 struct TypeResolver {};
 
 // Specialization for POD type
@@ -84,12 +81,11 @@ bool EncodeVariantImpl(const T& value,
 
 // Specialization for other types
 template <typename T>
-typename std::enable_if<!std::is_pointer<typename std::decay<T>::type>::value,
-                        bool>::type
-EncodeVariantImpl(const T& value,
-                  TypeResolver<T, false /* is_pod */, false /* Tensor */,
-                               false /* protobuf */>,
-                  VariantTensorData* data) {
+std::enable_if_t<!std::is_pointer_v<std::decay_t<T>>, bool> EncodeVariantImpl(
+    const T& value,
+    TypeResolver<T, false /* is_pod */, false /* Tensor */,
+                 false /* protobuf */>,
+    VariantTensorData* data) {
   if constexpr (std::is_void_v<decltype(value.Encode(data))>) {
     value.Encode(data);
     return true;
@@ -100,12 +96,11 @@ EncodeVariantImpl(const T& value,
 
 // Specialization for pointers
 template <typename T>
-typename std::enable_if<std::is_pointer<typename std::decay<T>::type>::value,
-                        bool>::type
-EncodeVariantImpl(const T& value,
-                  TypeResolver<T, false /* is_pod */, false /* Tensor */,
-                               false /* protobuf */>,
-                  VariantTensorData* data) {
+std::enable_if_t<std::is_pointer_v<std::decay_t<T>>, bool> EncodeVariantImpl(
+    const T& value,
+    TypeResolver<T, false /* is_pod */, false /* Tensor */,
+                 false /* protobuf */>,
+    VariantTensorData* data) {
   // Pointers cannot be encoded.
   return false;
 }
@@ -143,23 +138,21 @@ bool DecodeVariantImpl(VariantTensorData data,
 
 // Specialization for other types
 template <typename T>
-typename std::enable_if<!std::is_pointer<typename std::decay<T>::type>::value,
-                        bool>::type
-DecodeVariantImpl(VariantTensorData data,
-                  TypeResolver<T, false /* is_pod */, false /* Tensor */,
-                               false /* protobuf */>,
-                  T* value) {
+std::enable_if_t<!std::is_pointer_v<std::decay_t<T>>, bool> DecodeVariantImpl(
+    VariantTensorData data,
+    TypeResolver<T, false /* is_pod */, false /* Tensor */,
+                 false /* protobuf */>,
+    T* value) {
   return value->Decode(std::move(data));
 }
 
 // Specialization for pointers
 template <typename T>
-typename std::enable_if<std::is_pointer<typename std::decay<T>::type>::value,
-                        bool>::type
-DecodeVariantImpl(VariantTensorData data,
-                  TypeResolver<T, false /* is_pod */, false /* Tensor */,
-                               false /* protobuf */>,
-                  T* value) {
+std::enable_if_t<std::is_pointer_v<std::decay_t<T>>, bool> DecodeVariantImpl(
+    VariantTensorData data,
+    TypeResolver<T, false /* is_pod */, false /* Tensor */,
+                 false /* protobuf */>,
+    T* value) {
   // Pointers cannot be decoded.
   return false;
 }
@@ -169,15 +162,13 @@ struct has_type_name : std::false_type {};
 
 template <typename C>
 struct has_type_name<
-    C, typename std::enable_if<std::is_same<
-           decltype(std::declval<C>().TypeName()), std::string>::value>::type>
+    C, std::enable_if_t<
+           std::is_same_v<decltype(std::declval<C>().TypeName()), std::string>>>
     : std::true_type {};
 
-template <typename T, bool = has_type_name<typename std::decay<T>::type>::value,
-          bool = std::is_same<typename std::decay<T>::type,
-                              ::tensorflow::Tensor>::value,
-          bool = std::is_base_of<protobuf::MessageLite,
-                                 typename std::decay<T>::type>::value>
+template <typename T, bool = has_type_name<std::decay_t<T>>::value,
+          bool = std::is_same_v<std::decay_t<T>, ::tensorflow::Tensor>,
+          bool = std::is_base_of_v<protobuf::MessageLite, std::decay_t<T>>>
 struct TypeNameResolver {};
 
 template <typename T>
@@ -218,22 +209,21 @@ struct has_debug_string : std::false_type {};
 
 template <typename C>
 struct has_debug_string<
-    C,
-    typename std::enable_if<std::is_same<
-        decltype(std::declval<C>().DebugString()), std::string>::value>::type>
+    C, std::enable_if_t<std::is_same_v<
+           decltype(std::declval<C>().DebugString()), std::string>>>
     : std::true_type {};
 
 template <typename C, typename = void>
 struct can_strcat : std::false_type {};
 
 template <typename C>
-struct can_strcat<C, typename std::enable_if<std::is_same<
-                         decltype(strings::StrCat(std::declval<C>())),
-                         std::string>::value>::type> : std::true_type {};
+struct can_strcat<
+    C, std::enable_if_t<std::is_same_v<
+           decltype(strings::StrCat(std::declval<C>())), std::string>>>
+    : std::true_type {};
 
-template <typename T,
-          bool = has_debug_string<typename std::decay<T>::type>::value,
-          bool = can_strcat<typename std::decay<T>::type>::value>
+template <typename T, bool = has_debug_string<std::decay_t<T>>::value,
+          bool = can_strcat<std::decay_t<T>>::value>
 struct DebugStringResolver {};
 
 // TODO(ebrevdo): Expand DebugStringResolver to return TypeString if

From d6e622f1a829393c6c66d27b1e9a96553b3fdef5 Mon Sep 17 00:00:00 2001
From: wuren <928784910@qq.com>
Date: Mon, 12 Oct 2026 22:14:35 +0800
Subject: [PATCH 0550/1000] fix: compute static matrix norm permutations for
 XLA

Compute permutations in Python when tensor rank is known so matrix
2-norm compilation does not require a dynamic ListDiff. Preserve the
existing graph implementation for tensors with unknown rank.

Add rank-2/3/4 XLA coverage against NumPy, including negative,
non-adjacent and reversed axes and both keepdims settings.

Adapt the original fix from tensorflow/tensorflow#123946, with the
additional rank coverage previously proposed in Chaithra3-6/tensorflow#1.

Fixes #122865

Co-authored-by: Chaithra3-6 <chaithrashree365kp@gmail.com>
---
 tensorflow/python/kernel_tests/linalg/BUILD   |  2 +
 .../kernel_tests/linalg/norm_op_test.py       | 35 ++++++++++++++
 tensorflow/python/ops/linalg_ops.py           | 46 +++++++++++--------
 3 files changed, 65 insertions(+), 18 deletions(-)

diff --git a/tensorflow/python/kernel_tests/linalg/BUILD b/tensorflow/python/kernel_tests/linalg/BUILD
index fef7d0c72aa..5c34c886a37 100644
--- a/tensorflow/python/kernel_tests/linalg/BUILD
+++ b/tensorflow/python/kernel_tests/linalg/BUILD
@@ -793,12 +793,14 @@ cuda_py_strict_test(
     # TODO(b/208263392): Re-enable. tf.Squeeze op after tf.Where op doesn't reshape.
     xla_enable_strict_auto_jit = False,
     deps = [
+        "//tensorflow/python/eager:def_function",
         "//tensorflow/python/framework:constant_op",
         "//tensorflow/python/framework:test_lib",
         "//tensorflow/python/ops:array_ops",
         "//tensorflow/python/ops:linalg_ops",
         "//tensorflow/python/platform:client_testlib",
         "//third_party/py/numpy",
+        "@absl_py//absl/testing:parameterized",
     ],
 )
 
diff --git a/tensorflow/python/kernel_tests/linalg/norm_op_test.py b/tensorflow/python/kernel_tests/linalg/norm_op_test.py
index 5917a322556..c4e0847f19b 100644
--- a/tensorflow/python/kernel_tests/linalg/norm_op_test.py
+++ b/tensorflow/python/kernel_tests/linalg/norm_op_test.py
@@ -14,8 +14,10 @@
 # ==============================================================================
 """Tests for tensorflow.ops.tf.norm."""
 
+from absl.testing import parameterized
 import numpy as np
 
+from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
@@ -70,6 +72,39 @@ class NormOpTest(test_lib.TestCase):
     self.evaluate(linalg_ops.norm(matrix, axis=(-2, -1)))
 
 
+class MatrixNormXlaTest(test_lib.TestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      (f"{name}_keepdims_{keepdims}", shape, axis, keepdims)
+      for name, shape, axis in (
+          ("rank2", (3, 4), (0, 1)),
+          ("rank3_nonadjacent", (2, 3, 4), (0, 2)),
+          ("rank3_reversed", (2, 3, 4), (2, 0)),
+          ("rank4", (2, 3, 4, 5), (1, 2)),
+          ("rank4_negative", (2, 3, 4, 5), (-3, -2)),
+      )
+      for keepdims in (False, True)
+  )
+  @test_util.run_v2_only
+  def testMatrixNormOrd2XlaCompile(self, shape, axis, keepdims):
+    """Matrix 2-norm with static rank must compile without dynamic ListDiff."""
+    if not test_lib.is_built_with_xla():
+      self.skipTest("Test only applicable when running with XLA support")
+
+    matrix = np.random.RandomState(0).randn(*shape).astype(np.float32)
+    tensor = constant_op.constant(matrix)
+    expected = np.linalg.norm(matrix, ord=2, axis=axis, keepdims=keepdims)
+
+    def norm_fn(t):
+      return linalg_ops.norm(t, ord=2, axis=axis, keepdims=keepdims)
+
+    eager_result = norm_fn(tensor)
+    xla_result = def_function.function(norm_fn, jit_compile=True)(tensor)
+    self.assertAllEqual(expected.shape, xla_result.shape.as_list())
+    self.assertAllClose(expected, eager_result, rtol=1e-5, atol=1e-5)
+    self.assertAllClose(expected, xla_result, rtol=1e-5, atol=1e-5)
+
+
 def _GetNormOpTest(dtype_, shape_, ord_, axis_, keep_dims_, use_static_shape_):
 
   def _CompareNorm(self, matrix):
diff --git a/tensorflow/python/ops/linalg_ops.py b/tensorflow/python/ops/linalg_ops.py
index 9e545304aa6..ce61957df34 100644
--- a/tensorflow/python/ops/linalg_ops.py
+++ b/tensorflow/python/ops/linalg_ops.py
@@ -811,24 +811,34 @@ def norm(tensor,
 
     if ord in ['fro', 'euclidean', 2, 2.0]:
       if is_matrix_norm and ord in [2, 2.0]:
-        rank = array_ops.rank(tensor)
-        positive_axis = map_fn.map_fn(
-            lambda i: cond.cond(
-                i >= 0,
-                lambda: i,
-                lambda: i + rank),
-            ops.convert_to_tensor(axis))
-        axes = math_ops.range(rank)
-        perm_before = array_ops.concat([
-            gen_array_ops.list_diff(axes, positive_axis, dtypes.int32)[0],
-            positive_axis
-        ],
-                                       axis=0)
-        perm_after = map_fn.map_fn(
-            lambda i: math_ops.cast(
-                array_ops.squeeze(
-                    array_ops.where_v2(math_ops.equal(perm_before, i))),
-                dtype=dtypes.int32), axes)
+        static_rank = tensor.shape.rank
+        if static_rank is not None:
+          # axis is already validated to contain Python ints. Compute static
+          # permutations in Python so XLA need not compile a dynamic ListDiff.
+          positive_axis = tuple(a + static_rank if a < 0 else a for a in axis)
+          axes = list(range(static_rank))
+          remaining_axes = [a for a in axes if a not in positive_axis]
+          perm_before = remaining_axes + list(positive_axis)
+          perm_after = [perm_before.index(a) for a in axes]
+        else:
+          rank = array_ops.rank(tensor)
+          positive_axis = map_fn.map_fn(
+              lambda i: cond.cond(
+                  i >= 0,
+                  lambda: i,
+                  lambda: i + rank),
+              ops.convert_to_tensor(axis))
+          axes = math_ops.range(rank)
+          perm_before = array_ops.concat([
+              gen_array_ops.list_diff(axes, positive_axis, dtypes.int32)[0],
+              positive_axis
+          ],
+                                         axis=0)
+          perm_after = map_fn.map_fn(
+              lambda i: math_ops.cast(
+                  array_ops.squeeze(
+                      array_ops.where_v2(math_ops.equal(perm_before, i))),
+                  dtype=dtypes.int32), axes)
         permed = array_ops.transpose(tensor, perm=perm_before)
         matrix_2_norm = array_ops.expand_dims(
             math_ops.reduce_max(

From 093b3e88ae7bcd7210adb436f09e4060b09d8e0c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 21:56:11 -0700
Subject: [PATCH 0551/1000] Automated Code Change

PiperOrigin-RevId: 997672362
---
 tensorflow/core/grappler/optimizers/arithmetic_optimizer.h    | 2 +-
 tensorflow/core/grappler/optimizers/auto_mixed_precision.cc   | 2 +-
 tensorflow/core/grappler/optimizers/auto_mixed_precision.h    | 2 +-
 .../core/grappler/optimizers/auto_mixed_precision_lists.h     | 4 ++--
 tensorflow/core/grappler/optimizers/auto_parallel.h           | 2 +-
 .../core/grappler/optimizers/common_subgraph_elimination.h    | 4 ++--
 tensorflow/core/grappler/optimizers/constant_folding.h        | 2 +-
 tensorflow/core/grappler/optimizers/custom_graph_optimizer.h  | 2 +-
 .../optimizers/generic_layout_optimizer_transposer_factory.h  | 2 +-
 tensorflow/core/grappler/optimizers/graph_optimizer.h         | 2 +-
 10 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h
index 2a2a9dd95e8..db3f1c30fd8 100644
--- a/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h
+++ b/tensorflow/core/grappler/optimizers/arithmetic_optimizer.h
@@ -41,7 +41,7 @@ class ArithmeticOptimizer : public GraphOptimizer {
       : opt_level_(opt_level),
         options_(ArithmeticOptimizerOptions::Default(opt_level)) {}
 
-  ~ArithmeticOptimizer() override {}
+  ~ArithmeticOptimizer() override = default;
 
   std::string name() const override { return "arithmetic_optimizer"; };
 
diff --git a/tensorflow/core/grappler/optimizers/auto_mixed_precision.cc b/tensorflow/core/grappler/optimizers/auto_mixed_precision.cc
index f2de68dba74..4e5b3e474f5 100644
--- a/tensorflow/core/grappler/optimizers/auto_mixed_precision.cc
+++ b/tensorflow/core/grappler/optimizers/auto_mixed_precision.cc
@@ -286,7 +286,7 @@ std::vector<int> NonControlInputs(const NodeDef& node) {
 // input/output port mappings.
 class NodeTypeAttrMap {
  public:
-  NodeTypeAttrMap() {}
+  NodeTypeAttrMap() = default;
 
   explicit NodeTypeAttrMap(const GraphDef& graph) { TF_CHECK_OK(Init(graph)); }
 
diff --git a/tensorflow/core/grappler/optimizers/auto_mixed_precision.h b/tensorflow/core/grappler/optimizers/auto_mixed_precision.h
index 4343787f136..cc9da384323 100644
--- a/tensorflow/core/grappler/optimizers/auto_mixed_precision.h
+++ b/tensorflow/core/grappler/optimizers/auto_mixed_precision.h
@@ -40,7 +40,7 @@ class AutoMixedPrecision : public GraphOptimizer {
       AutoMixedPrecisionMode mode = AutoMixedPrecisionMode::CUDA)
       : mode_(mode) {}
 
-  ~AutoMixedPrecision() override {}
+  ~AutoMixedPrecision() override = default;
 
   std::string name() const override {
     switch (mode_) {
diff --git a/tensorflow/core/grappler/optimizers/auto_mixed_precision_lists.h b/tensorflow/core/grappler/optimizers/auto_mixed_precision_lists.h
index d2f86c54ca1..376b1516eb8 100644
--- a/tensorflow/core/grappler/optimizers/auto_mixed_precision_lists.h
+++ b/tensorflow/core/grappler/optimizers/auto_mixed_precision_lists.h
@@ -31,7 +31,7 @@ namespace grappler {
 // (referred to as 'f16' for short) and which ops stay as fp32.
 class AutoMixedPrecisionLists {
  public:
-  virtual ~AutoMixedPrecisionLists() {}
+  virtual ~AutoMixedPrecisionLists() = default;
 
   // Returns the set of ops that are considered numerically-safe (for execution
   // in f16), performance-critical, and can run in f16. These ops are always
@@ -390,7 +390,7 @@ using AutoMixedPrecisionListsCuda = AutoMixedPrecisionListsFp16;
 
 class AutoMixedPrecisionListsMkl : public AutoMixedPrecisionLists {
  public:
-  AutoMixedPrecisionListsMkl() {}
+  AutoMixedPrecisionListsMkl() = default;
 
   // Only ops which are supported by MKL in bfloat16 should be added to the
   // allow list, infer list, or clear list.
diff --git a/tensorflow/core/grappler/optimizers/auto_parallel.h b/tensorflow/core/grappler/optimizers/auto_parallel.h
index 9ced1381c76..c47313b2027 100644
--- a/tensorflow/core/grappler/optimizers/auto_parallel.h
+++ b/tensorflow/core/grappler/optimizers/auto_parallel.h
@@ -29,7 +29,7 @@ class AutoParallel : public GraphOptimizer {
   AutoParallel(int num_replicas) : num_replicas_(num_replicas) {
     CHECK(num_replicas_ >= 2);
   }
-  ~AutoParallel() override {}
+  ~AutoParallel() override = default;
 
   std::string name() const override { return "autoparallel"; };
 
diff --git a/tensorflow/core/grappler/optimizers/common_subgraph_elimination.h b/tensorflow/core/grappler/optimizers/common_subgraph_elimination.h
index b4e7e11440c..2fa3c6d706e 100644
--- a/tensorflow/core/grappler/optimizers/common_subgraph_elimination.h
+++ b/tensorflow/core/grappler/optimizers/common_subgraph_elimination.h
@@ -35,12 +35,12 @@ struct GrapplerItem;
 
 class CommonSubgraphElimination : public GraphOptimizer {
  public:
-  CommonSubgraphElimination() {}
+  CommonSubgraphElimination() = default;
 
   explicit CommonSubgraphElimination(RewriterConfig::Toggle opt_level)
       : opt_level_(opt_level) {}
 
-  ~CommonSubgraphElimination() override {}
+  ~CommonSubgraphElimination() override = default;
 
   std::string name() const override { return "common_subgraph_elimination"; };
 
diff --git a/tensorflow/core/grappler/optimizers/constant_folding.h b/tensorflow/core/grappler/optimizers/constant_folding.h
index 5e9113414aa..d283f87ff81 100644
--- a/tensorflow/core/grappler/optimizers/constant_folding.h
+++ b/tensorflow/core/grappler/optimizers/constant_folding.h
@@ -53,7 +53,7 @@ class ConstantFolding : public GraphOptimizer {
                   bool disable_compressed_tensor_optimization = false,
                   bool fold_quantization_emulation = true);
 
-  ~ConstantFolding() override {}
+  ~ConstantFolding() override = default;
 
   std::string name() const override { return "constant_folding"; };
 
diff --git a/tensorflow/core/grappler/optimizers/custom_graph_optimizer.h b/tensorflow/core/grappler/optimizers/custom_graph_optimizer.h
index 9a8f59f14af..2fb8598f227 100644
--- a/tensorflow/core/grappler/optimizers/custom_graph_optimizer.h
+++ b/tensorflow/core/grappler/optimizers/custom_graph_optimizer.h
@@ -27,7 +27,7 @@ namespace grappler {
 // A custom optimizer that can be registered.
 class CustomGraphOptimizer : public GraphOptimizer {
  public:
-  ~CustomGraphOptimizer() override {}
+  ~CustomGraphOptimizer() override = default;
   virtual absl::Status Init(
       const tensorflow::RewriterConfig_CustomGraphOptimizer* config =
           nullptr) = 0;
diff --git a/tensorflow/core/grappler/optimizers/generic_layout_optimizer_transposer_factory.h b/tensorflow/core/grappler/optimizers/generic_layout_optimizer_transposer_factory.h
index 79e3adeddf8..8458d2e0781 100644
--- a/tensorflow/core/grappler/optimizers/generic_layout_optimizer_transposer_factory.h
+++ b/tensorflow/core/grappler/optimizers/generic_layout_optimizer_transposer_factory.h
@@ -26,7 +26,7 @@ namespace grappler {
 
 class TransposerFactory {
  public:
-  explicit TransposerFactory() {}
+  explicit TransposerFactory() = default;
 
   std::shared_ptr<Transposer> GetTransposer(const NodeDef& node);
 
diff --git a/tensorflow/core/grappler/optimizers/graph_optimizer.h b/tensorflow/core/grappler/optimizers/graph_optimizer.h
index 786573a4585..f8cbfa56c48 100644
--- a/tensorflow/core/grappler/optimizers/graph_optimizer.h
+++ b/tensorflow/core/grappler/optimizers/graph_optimizer.h
@@ -36,7 +36,7 @@ struct GrapplerItem;
 class GraphOptimizer {
  public:
   GraphOptimizer() : deadline_usec_(0) {}
-  virtual ~GraphOptimizer() {}
+  virtual ~GraphOptimizer() = default;
 
   virtual std::string name() const = 0;
 

From f073d70daa53752a8fa20aadbadabeaa00e284f5 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 22:03:49 -0700
Subject: [PATCH 0552/1000] Automated Code Change

PiperOrigin-RevId: 997675439
---
 tensorflow/core/util/cuda_sparse.h | 2 +-
 tensorflow/core/util/gpu_solvers.h | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/util/cuda_sparse.h b/tensorflow/core/util/cuda_sparse.h
index 0538b919aeb..0a5b865ad9a 100644
--- a/tensorflow/core/util/cuda_sparse.h
+++ b/tensorflow/core/util/cuda_sparse.h
@@ -310,7 +310,7 @@ class GpuSparse {
  public:
   // This object stores a pointer to context, which must outlive it.
   explicit GpuSparse(OpKernelContext* context);
-  virtual ~GpuSparse() {}
+  virtual ~GpuSparse() = default;
 
   // This initializes the GpuSparse class if it hasn't
   // been initialized yet.  All following public methods require the
diff --git a/tensorflow/core/util/gpu_solvers.h b/tensorflow/core/util/gpu_solvers.h
index 327e39a2161..df6a65014b9 100644
--- a/tensorflow/core/util/gpu_solvers.h
+++ b/tensorflow/core/util/gpu_solvers.h
@@ -619,7 +619,7 @@ class ScratchSpace {
                                        &scratch_tensor_, alloc_attr));
   }
 
-  virtual ~ScratchSpace() {}
+  virtual ~ScratchSpace() = default;
 
   Scalar* mutable_data() {
     return scratch_tensor_.template flat<Scalar>().data();

From d1fb07a80a4b638efd898867bfc7bf3b9c4cbfa8 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 22:07:58 -0700
Subject: [PATCH 0553/1000] Automated Code Change

PiperOrigin-RevId: 997677221
---
 tensorflow/lite/tools/benchmark/benchmark_model.h | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tensorflow/lite/tools/benchmark/benchmark_model.h b/tensorflow/lite/tools/benchmark/benchmark_model.h
index 93072ffcddd..ed8b7a512c3 100644
--- a/tensorflow/lite/tools/benchmark/benchmark_model.h
+++ b/tensorflow/lite/tools/benchmark/benchmark_model.h
@@ -43,7 +43,7 @@ enum RunType {
 
 class BenchmarkResults {
  public:
-  BenchmarkResults() {}
+  BenchmarkResults() = default;
   BenchmarkResults(double model_size_mb, int64_t startup_latency_us,
                    uint64_t input_bytes,
                    tensorflow::StatWithPercentiles<int64_t> warmup_time_us,
@@ -110,7 +110,7 @@ class BenchmarkListener {
   virtual void OnSingleRunEnd() {}
   // Called after the (outer) inference loop begins.
   virtual void OnBenchmarkEnd(const BenchmarkResults& results) {}
-  virtual ~BenchmarkListener() {}
+  virtual ~BenchmarkListener() = default;
 };
 
 // A listener that forwards its method calls to a collection of listeners.
@@ -156,7 +156,7 @@ class BenchmarkListeners : public BenchmarkListener {
     }
   }
 
-  ~BenchmarkListeners() override {}
+  ~BenchmarkListeners() override = default;
 
  private:
   // Use vector so listeners are invoked in the order they are added.
@@ -190,7 +190,7 @@ class BenchmarkModel {
   BenchmarkModel();
   explicit BenchmarkModel(BenchmarkParams params)
       : params_(std::move(params)) {}
-  virtual ~BenchmarkModel() {}
+  virtual ~BenchmarkModel() = default;
   virtual TfLiteStatus Init() = 0;
   virtual TfLiteStatus Run(int argc, char** argv);
   virtual TfLiteStatus Run();

From 5a3b649f3f80afecbb86e860866e6867cbc616d2 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 22:42:08 -0700
Subject: [PATCH 0554/1000] Automated Code Change

PiperOrigin-RevId: 997689678
---
 tensorflow/core/common_runtime/placer.cc                      | 2 +-
 tensorflow/core/common_runtime/pool_allocator.h               | 4 ++--
 .../core/common_runtime/process_function_library_runtime.h    | 2 +-
 .../common_runtime/process_function_library_runtime_test.cc   | 2 +-
 tensorflow/core/common_runtime/profile_handler.h              | 4 ++--
 tensorflow/core/common_runtime/rendezvous_mgr.cc              | 2 +-
 tensorflow/core/common_runtime/request_cost_accessor.h        | 2 +-
 tensorflow/core/common_runtime/ring_alg.h                     | 2 +-
 tensorflow/core/common_runtime/ring_gatherer.h                | 2 +-
 tensorflow/core/common_runtime/session.cc                     | 4 ++--
 10 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/tensorflow/core/common_runtime/placer.cc b/tensorflow/core/common_runtime/placer.cc
index 93fe468c46e..23e97f0c043 100644
--- a/tensorflow/core/common_runtime/placer.cc
+++ b/tensorflow/core/common_runtime/placer.cc
@@ -218,7 +218,7 @@ Placer::Placer(Graph* graph, const std::string& function_name,
                const DeviceSet* devices)
     : Placer(graph, function_name, flib_def, devices, nullptr, true, false) {}
 
-Placer::~Placer() {}
+Placer::~Placer() = default;
 
 absl::Status Placer::Run() {
   GraphOptimizationPassOptions options;
diff --git a/tensorflow/core/common_runtime/pool_allocator.h b/tensorflow/core/common_runtime/pool_allocator.h
index 69c1e7a75b8..295638cecc4 100644
--- a/tensorflow/core/common_runtime/pool_allocator.h
+++ b/tensorflow/core/common_runtime/pool_allocator.h
@@ -35,7 +35,7 @@ namespace tensorflow {
 // Interface of an object that rounds up integers.
 class RoundUpInterface {
  public:
-  virtual ~RoundUpInterface() {}
+  virtual ~RoundUpInterface() = default;
   virtual size_t RoundUp(size_t num_bytes) = 0;
 };
 
@@ -157,7 +157,7 @@ class BasicCPUAllocator : public SubAllocator {
                     const std::vector<Visitor>& free_visitors)
       : SubAllocator(alloc_visitors, free_visitors), numa_node_(numa_node) {}
 
-  ~BasicCPUAllocator() override {}
+  ~BasicCPUAllocator() override = default;
 
   void* Alloc(size_t alignment, size_t num_bytes,
               size_t* bytes_received) override;
diff --git a/tensorflow/core/common_runtime/process_function_library_runtime.h b/tensorflow/core/common_runtime/process_function_library_runtime.h
index 638f9cf5c14..bc3e3c4f3d7 100644
--- a/tensorflow/core/common_runtime/process_function_library_runtime.h
+++ b/tensorflow/core/common_runtime/process_function_library_runtime.h
@@ -46,7 +46,7 @@ namespace tensorflow {
 
 class FunctionArgsInterface {
  public:
-  virtual ~FunctionArgsInterface() {}
+  virtual ~FunctionArgsInterface() = default;
 
   virtual bool HasRemoteOrPackedInputs() const = 0;
 
diff --git a/tensorflow/core/common_runtime/process_function_library_runtime_test.cc b/tensorflow/core/common_runtime/process_function_library_runtime_test.cc
index f1d76218492..5edcd2f734a 100644
--- a/tensorflow/core/common_runtime/process_function_library_runtime_test.cc
+++ b/tensorflow/core/common_runtime/process_function_library_runtime_test.cc
@@ -981,7 +981,7 @@ class TestFunctionPackedArgs : public FunctionArgsInterface {
     packed_args_.emplace(index, std::move(tensor_args));
   }
 
-  ~TestFunctionPackedArgs() override{};
+  ~TestFunctionPackedArgs() override = default;
 
   bool HasRemoteOrPackedInputs() const override { return true; };
 
diff --git a/tensorflow/core/common_runtime/profile_handler.h b/tensorflow/core/common_runtime/profile_handler.h
index 28ae706c3f0..1651b7b0228 100644
--- a/tensorflow/core/common_runtime/profile_handler.h
+++ b/tensorflow/core/common_runtime/profile_handler.h
@@ -26,8 +26,8 @@ namespace tensorflow {
 // A profile handler collects event stats from a running step.
 class ProfileHandler {
  public:
-  ProfileHandler() {}
-  virtual ~ProfileHandler() {}
+  ProfileHandler() = default;
+  virtual ~ProfileHandler() = default;
 
   // Records that a single Op was executed in the current step.
   //
diff --git a/tensorflow/core/common_runtime/rendezvous_mgr.cc b/tensorflow/core/common_runtime/rendezvous_mgr.cc
index 8ff32de5a2f..215369359a1 100644
--- a/tensorflow/core/common_runtime/rendezvous_mgr.cc
+++ b/tensorflow/core/common_runtime/rendezvous_mgr.cc
@@ -209,7 +209,7 @@ PrivateIntraProcessRendezvous::PrivateIntraProcessRendezvous(
     : device_mgr_(device_mgr),
       local_(nullptr, /* num_shards= */ device_mgr->NumDevices()) {}
 
-PrivateIntraProcessRendezvous::~PrivateIntraProcessRendezvous() {}
+PrivateIntraProcessRendezvous::~PrivateIntraProcessRendezvous() = default;
 
 absl::Status PrivateIntraProcessRendezvous::Send(const ParsedKey& key,
                                                  const Rendezvous::Args& args,
diff --git a/tensorflow/core/common_runtime/request_cost_accessor.h b/tensorflow/core/common_runtime/request_cost_accessor.h
index ba64da4b296..738c6b0c54b 100644
--- a/tensorflow/core/common_runtime/request_cost_accessor.h
+++ b/tensorflow/core/common_runtime/request_cost_accessor.h
@@ -28,7 +28,7 @@ namespace tensorflow {
 // request.
 class RequestCostAccessor {
  public:
-  virtual ~RequestCostAccessor() {}
+  virtual ~RequestCostAccessor() = default;
   virtual RequestCost* GetRequestCost() const = 0;
 };
 
diff --git a/tensorflow/core/common_runtime/ring_alg.h b/tensorflow/core/common_runtime/ring_alg.h
index ebb53adf49f..ba6339bb912 100644
--- a/tensorflow/core/common_runtime/ring_alg.h
+++ b/tensorflow/core/common_runtime/ring_alg.h
@@ -33,7 +33,7 @@ class Device;
 class RingAlg : public CollectiveImplementationInterface {
  public:
   explicit RingAlg(CollectiveType type, const std::string& name);
-  ~RingAlg() override {}
+  ~RingAlg() override = default;
 
   // Establishes the requested number of subdivision permutations based on the
   // ring order implicit in the device order.
diff --git a/tensorflow/core/common_runtime/ring_gatherer.h b/tensorflow/core/common_runtime/ring_gatherer.h
index 875ac86a6d9..c48c0dfb957 100644
--- a/tensorflow/core/common_runtime/ring_gatherer.h
+++ b/tensorflow/core/common_runtime/ring_gatherer.h
@@ -32,7 +32,7 @@ class Device;
 class RingGatherer : public RingAlg {
  public:
   RingGatherer() : RingAlg(GATHER_COLLECTIVE, "Gather") {}
-  ~RingGatherer() override {}
+  ~RingGatherer() override = default;
 
   absl::Status InitializeCollectiveParams(
       CollectiveParams* col_params) override;
diff --git a/tensorflow/core/common_runtime/session.cc b/tensorflow/core/common_runtime/session.cc
index c6680f900ee..fab56c07194 100644
--- a/tensorflow/core/common_runtime/session.cc
+++ b/tensorflow/core/common_runtime/session.cc
@@ -32,9 +32,9 @@ auto* session_created = monitoring::Gauge<bool, 0>::New(
 
 void SetSessionCreatedMetric() { session_created->GetCell()->Set(true); }
 
-Session::Session() {}
+Session::Session() = default;
 
-Session::~Session() {}
+Session::~Session() = default;
 
 absl::Status Session::Run(
     const RunOptions& run_options,

From e79e23ac9c45cdd7fc22fb443a8584af66f0abbb Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Sun, 11 Oct 2026 22:49:38 -0700
Subject: [PATCH 0555/1000] Automated Code Change

PiperOrigin-RevId: 997692312
---
 .../lite/tools/evaluation/tasks/inference_diff/run_eval.cc      | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/lite/tools/evaluation/tasks/inference_diff/run_eval.cc b/tensorflow/lite/tools/evaluation/tasks/inference_diff/run_eval.cc
index 897cccd5387..27e4bdf5073 100644
--- a/tensorflow/lite/tools/evaluation/tasks/inference_diff/run_eval.cc
+++ b/tensorflow/lite/tools/evaluation/tasks/inference_diff/run_eval.cc
@@ -40,7 +40,7 @@ constexpr char kDelegateFlag[] = "delegate";
 class InferenceDiff : public TaskExecutor {
  public:
   InferenceDiff() : num_runs_(50), num_interpreter_threads_(1) {}
-  ~InferenceDiff() override {}
+  ~InferenceDiff() override = default;
 
  protected:
   std::vector<Flag> GetFlags() final;

From 89d193110483ffb81e62885c1dd50b463bba3217 Mon Sep 17 00:00:00 2001
From: Bixia Zheng <bixia@google.com>
Date: Mon, 12 Oct 2026 00:18:26 -0700
Subject: [PATCH 0556/1000] Avoid adding missing control-flow shardings when
 Shardy generates device code.

PiperOrigin-RevId: 997726468
---
 third_party/xla/xla/service/spmd/shardy/BUILD            | 2 +-
 .../xla/xla/service/spmd/shardy/shardy_xla_pass.cc       | 9 ++++++++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/shardy/BUILD b/third_party/xla/xla/service/spmd/shardy/BUILD
index bd0a029bdc3..c387d7b6729 100644
--- a/third_party/xla/xla/service/spmd/shardy/BUILD
+++ b/third_party/xla/xla/service/spmd/shardy/BUILD
@@ -58,7 +58,6 @@ cc_library(
         "//xla/service/spmd/shardy/stablehlo_round_trip:stablehlo_import",
         "//xla/tsl/framework/mlir:status_scoped_diagnostic_handler",
         "//xla/tsl/platform:env",
-        "//xla/tsl/platform:errors",
         "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -77,6 +76,7 @@ cc_library(
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:Support",
         "@shardy//shardy/common:file_utils",
+        "@shardy//shardy/dialect/sdy/transforms/common:partitioner_stage",
         "@shardy//shardy/dialect/sdy/transforms/common:propagation_options",
         "@shardy//shardy/dialect/sdy/transforms/export:utils",
         "@shardy//shardy/dialect/sdy/transforms/propagation:passes",
diff --git a/third_party/xla/xla/service/spmd/shardy/shardy_xla_pass.cc b/third_party/xla/xla/service/spmd/shardy/shardy_xla_pass.cc
index c05f855e241..f1ed9e0fc12 100644
--- a/third_party/xla/xla/service/spmd/shardy/shardy_xla_pass.cc
+++ b/third_party/xla/xla/service/spmd/shardy/shardy_xla_pass.cc
@@ -43,6 +43,7 @@ limitations under the License.
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/LLVM.h"
 #include "shardy/common/file_utils.h"
+#include "shardy/dialect/sdy/transforms/common/partitioner_stage.h"
 #include "shardy/dialect/sdy/transforms/common/propagation_options.h"
 #include "shardy/dialect/sdy/transforms/export/utils.h"
 #include "shardy/dialect/sdy/transforms/propagation/passes.h"
@@ -73,7 +74,6 @@ limitations under the License.
 #include "xla/status_macros.h"
 #include "xla/tsl/framework/mlir/status_scoped_diagnostic_handler.h"
 #include "xla/tsl/platform/env.h"
-#include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"
 #include "xla/xla.pb.h"
@@ -426,7 +426,14 @@ absl::Status runShardingPropagation(HloModule* hloModule,
   options.enableNativeNonFlatSupport = enableNativeNonFlatSupport;
   mlir::sdy::addPropagationPipeline(pm, dumpIndex, options);
 
+  bool shardyGeneratesDeviceCode =
+      !options.avoidExportForPartitioning &&
+      !options.enablePerInstructionPartitioning &&
+      options.partitionerStage ==
+          mlir::sdy::PartitionerStage::kConvertGlobalToLocal;
   xla::sdy::StablehloExportPipelineOptions stablehloExportPipelineOptions;
+  stablehloExportPipelineOptions.addMissingShardingToControlFlow =
+      !shardyGeneratesDeviceCode;
   stablehloExportPipelineOptions.enableHloShardingV3 = enableHloShardingV3;
   stablehloExportPipelineOptions.exportAllReduceScatter =
       debugOptions.xla_sdy_export_all_reduce_scatter();

From e890233e44771281c76c2226ae54c90c18bb1b01 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Mon, 12 Oct 2026 00:31:19 -0700
Subject: [PATCH 0557/1000] [XLA:GPU]: Migrate AllToAllThunk to use
 PerDeviceState

PiperOrigin-RevId: 997732140
---
 .../xla/xla/backends/gpu/runtime/BUILD        |  5 +-
 .../backends/gpu/runtime/all_gather_thunk.h   |  1 -
 .../backends/gpu/runtime/all_to_all_thunk.cc  | 94 ++++++++-----------
 .../backends/gpu/runtime/all_to_all_thunk.h   | 45 ++++-----
 .../runtime/all_to_all_thunk_multigpu_test.cc | 41 +++++++-
 .../gpu/runtime/all_to_all_thunk_test.cc      |  4 +-
 .../runtime/thunk_proto_deserialization.cc    |  7 +-
 .../xla/xla/service/gpu/thunk_emitter.cc      |  9 +-
 8 files changed, 110 insertions(+), 96 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index 8f30fdfe46d..01634f1e2ef 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -1807,7 +1807,6 @@ cc_library(
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:str_format",
-        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
         "@tsl//tsl/platform:casts",
     ],
@@ -2229,6 +2228,7 @@ cc_library(
         ":collective_execution",
         ":collective_thunk",
         ":collective_thunk_proto_cc",
+        ":per_device_state",
         ":thunk",
         ":thunk_proto_cc",
         "//xla:shape_util",
@@ -2250,15 +2250,12 @@ cc_library(
         "//xla/stream_executor:stream",
         "//xla/tsl/platform:logging",
         "@com_google_absl//absl/algorithm:container",
-        "@com_google_absl//absl/base:core_headers",
-        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
-        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
     ],
 )
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_gather_thunk.h b/third_party/xla/xla/backends/gpu/runtime/all_gather_thunk.h
index dbaa7027dad..f6f210cdccf 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_gather_thunk.h
+++ b/third_party/xla/xla/backends/gpu/runtime/all_gather_thunk.h
@@ -24,7 +24,6 @@ limitations under the License.
 #include "absl/container/node_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
 #include "xla/backends/gpu/collectives/gpu_clique_key.h"
 #include "xla/backends/gpu/runtime/collective_thunk.h"
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk.cc
index be45e7e369a..d25bc5653a3 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk.cc
@@ -17,7 +17,6 @@ limitations under the License.
 
 #include <cstdint>
 #include <cstdlib>
-#include <iterator>
 #include <memory>
 #include <optional>
 #include <string>
@@ -25,20 +24,19 @@ limitations under the License.
 #include <vector>
 
 #include "absl/algorithm/container.h"
-#include "absl/container/flat_hash_map.h"
 #include "absl/container/inlined_vector.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
 #include "absl/strings/substitute.h"
-#include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
 #include "xla/backends/gpu/collectives/gpu_clique_key.h"
 #include "xla/backends/gpu/collectives/gpu_collectives.h"
 #include "xla/backends/gpu/runtime/collective_execution.h"
 #include "xla/backends/gpu/runtime/collective_thunk.h"
 #include "xla/backends/gpu/runtime/collective_thunk.pb.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk.pb.h"
 #include "xla/backends/gpu/transforms/collectives/collective_ops_utils.h"
@@ -74,26 +72,29 @@ AllToAllConfig GetAllToAllConfig(const HloAllToAllInstruction* instr) {
 struct BufferRendezvousValue {
   uint16_t rank;
   uint64_t buffer;
+  se::Event* event;
 };
 
 }  // namespace
 
 AllToAllThunk::AllToAllThunk(ThunkInfo thunk_info, const AllToAllConfig& config,
                              std::vector<CollectiveThunk::Buffer> buffers,
-                             bool p2p_memcpy_enabled)
+                             bool p2p_memcpy_enabled, int devices_per_host)
     : CollectiveThunk(Thunk::kAllToAll, thunk_info, std::move(buffers)),
       config_(config),
-      p2p_memcpy_enabled_(p2p_memcpy_enabled) {
+      p2p_memcpy_enabled_(p2p_memcpy_enabled),
+      per_device_states_(devices_per_host) {
   CHECK_EQ(config_.config.operand_element_type.size(), this->buffers().size());
 }
 
 AllToAllThunk::AllToAllThunk(ThunkInfo thunk_info,
                              const HloAllToAllInstruction* instr,
                              std::vector<CollectiveThunk::Buffer> buffers,
-                             bool p2p_memcpy_enabled)
+                             bool p2p_memcpy_enabled, int devices_per_host)
     : CollectiveThunk(Thunk::kAllToAll, thunk_info, std::move(buffers)),
       config_(GetAllToAllConfig(instr)),
-      p2p_memcpy_enabled_(p2p_memcpy_enabled) {
+      p2p_memcpy_enabled_(p2p_memcpy_enabled),
+      per_device_states_(devices_per_host) {
   CHECK_EQ(config_.config.operand_element_type.size(), this->buffers().size());
 }
 
@@ -146,24 +147,15 @@ absl::Status AllToAllThunk::Initialize(const InitializeParams& params) {
     ABSL_ASSIGN_OR_RETURN(int32_t num_ranks, comm->NumRanks());
 
     se::StreamExecutor* executor = params.executor;
-    {
-      absl::MutexLock lock(pointer_maps_mutex_);
-      if (!receive_pointer_maps_.count(executor)) {
-        ABSL_ASSIGN_OR_RETURN(
-            std::unique_ptr<se::MemoryAllocation> alloc,
-            executor->HostMemoryAllocate(num_ranks * sizeof(uint64_t)));
-        bool inserted =
-            receive_pointer_maps_.insert({executor, std::move(alloc)}).second;
-        CHECK(inserted);
-      }
+    ABSL_ASSIGN_OR_RETURN(DeviceState * state, per_device_states_.GetOrCreate(
+                                              executor->device_ordinal()));
+    if (!state->receive_pointer_map) {
+      ABSL_ASSIGN_OR_RETURN(
+          state->receive_pointer_map,
+          executor->HostMemoryAllocate(num_ranks * sizeof(uint64_t)));
     }
-    {
-      absl::MutexLock lock(events_mutex_);
-      if (!events_.count(executor)) {
-        ABSL_ASSIGN_OR_RETURN(std::unique_ptr<se::Event> event,
-                         executor->CreateEvent());
-        events_.insert({executor, std::move(event)});
-      }
+    if (!state->event) {
+      ABSL_ASSIGN_OR_RETURN(state->event, executor->CreateEvent());
     }
     std::optional<RankId> rank =
         clique_key.rank(params.collective_params->global_device_id);
@@ -178,6 +170,7 @@ absl::Status AllToAllThunk::Initialize(const InitializeParams& params) {
     for (int peer = 0; peer < num_ranks; ++peer) {
       BufferRendezvousValue buffer_rendezvous_value;
       buffer_rendezvous_value.rank = rank.value().value();
+      buffer_rendezvous_value.event = state->event.get();
       int buffer_idx = (rank.value().value() + peer) % num_ranks;
       if (config_.has_split_dimension) {
         buffer_rendezvous_value.buffer = reinterpret_cast<uint64_t>(
@@ -208,14 +201,16 @@ absl::Status AllToAllThunk::Initialize(const InitializeParams& params) {
               }));
       int peer_buffer_idx =
           (rank.value().value() - peer + num_ranks) % num_ranks;
-      uint64_t* recv_ptr;
-      {
-        absl::MutexLock lock(pointer_maps_mutex_);
-        recv_ptr = reinterpret_cast<uint64_t*>(
-            receive_pointer_maps_[executor]->address().opaque());
-      }
+      uint64_t* recv_ptr = reinterpret_cast<uint64_t*>(
+          state->receive_pointer_map->address().opaque());
       recv_ptr[config_.has_split_dimension ? peer_buffer_idx : peer] =
           (*rendezvous_results)[peer_buffer_idx].buffer;
+      if (peer == 0 && state->events.empty()) {
+        state->events.reserve(num_ranks);
+        for (const BufferRendezvousValue& result : *rendezvous_results) {
+          state->events.push_back(result.event);
+        }
+      }
     }
   }
   return absl::OkStatus();
@@ -233,28 +228,17 @@ absl::Status AllToAllThunk::RunCollective(const ExecuteParams& params,
                          config_.config.replica_groups,
                          config_.config.group_mode) &&
       p2p_memcpy_enabled_) {
-    uint64_t* receive_pointer_map = nullptr;
-    {
-      absl::MutexLock lock(pointer_maps_mutex_);
-      receive_pointer_map = reinterpret_cast<uint64_t*>(
-          receive_pointer_maps_[stream.parent()]->address().opaque());
-    }
+    DeviceState* state =
+        per_device_states_.Find(stream.parent()->device_ordinal());
+    TF_RET_CHECK(state != nullptr);
+    uint64_t* receive_pointer_map = reinterpret_cast<uint64_t*>(
+        state->receive_pointer_map->address().opaque());
     std::optional<RankId> rank =
         clique_key.rank(params.collective_params->global_device_id);
-    se::Event* event = nullptr;
-    {
-      absl::MutexLock lock(events_mutex_);
-      event = events_[stream.parent()].get();
-    }
-    std::vector<se::Event*> events;
-    {
-      absl::MutexLock lock(events_mutex_);
-      absl::c_transform(events_, std::back_inserter(events),
-                        [](const auto& pair) { return pair.second.get(); });
-    }
-    return xla::gpu::RunMemCpyAllToAll(
-        config_.has_split_dimension, device_buffers, stream, comm,
-        receive_pointer_map, clique_key, *rank, event, events);
+    return xla::gpu::RunMemCpyAllToAll(config_.has_split_dimension,
+                                       device_buffers, stream, comm,
+                                       receive_pointer_map, clique_key, *rank,
+                                       state->event.get(), state->events);
   }
   return xla::gpu::RunAllToAll(config_.has_split_dimension, device_buffers,
                                stream, comm,
@@ -263,7 +247,8 @@ absl::Status AllToAllThunk::RunCollective(const ExecuteParams& params,
 
 absl::StatusOr<std::unique_ptr<AllToAllThunk>> AllToAllThunk::FromProto(
     ThunkInfo thunk_info, const AllToAllThunkProto& thunk_proto,
-    absl::Span<const BufferAllocation> buffer_allocations) {
+    absl::Span<const BufferAllocation> buffer_allocations,
+    int devices_per_host) {
   std::vector<CollectiveThunk::Buffer> buffers;
   buffers.reserve(thunk_proto.buffers_size());
   for (const CollectiveBufferProto& proto : thunk_proto.buffers()) {
@@ -279,7 +264,7 @@ absl::StatusOr<std::unique_ptr<AllToAllThunk>> AllToAllThunk::FromProto(
   return std::make_unique<AllToAllThunk>(
       std::move(thunk_info),
       AllToAllConfig{config, thunk_proto.has_split_dimension()}, buffers,
-      thunk_proto.p2p_memcpy_enabled());
+      thunk_proto.p2p_memcpy_enabled(), devices_per_host);
 }
 
 absl::StatusOr<ThunkProto> AllToAllThunk::ToProto() const {
@@ -371,7 +356,8 @@ absl::Status RunAllToAll(bool has_split_dimension,
 absl::Status SyncProgress(absl::string_view name,
                           const GpuCliqueKey& clique_key, RankId rank,
                           int64_t num_ranks, se::Stream& stream,
-                          se::Event* event, std::vector<se::Event*>& events) {
+                          se::Event* event,
+                          absl::Span<se::Event* const> events) {
   // Record event for this device.
   ABSL_RETURN_IF_ERROR(stream.RecordEvent(event));
 
@@ -397,7 +383,7 @@ absl::Status RunMemCpyAllToAll(bool has_split_dimension,
                                uint64_t receive_pointer_map[],
                                const GpuCliqueKey& clique_key, RankId rank,
                                se::Event* event,
-                               std::vector<se::Event*>& events) {
+                               absl::Span<se::Event* const> events) {
   int device_ordinal = stream.parent()->device_ordinal();
   XLA_VLOG_DEVICE(3, device_ordinal) << "Performing mem-copy-all-to-all";
   ABSL_ASSIGN_OR_RETURN(int32_t num_ranks, comm.NumRanks());
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk.h b/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk.h
index 62e2ae39218..bdad0f8d897 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk.h
+++ b/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk.h
@@ -20,15 +20,13 @@ limitations under the License.
 #include <memory>
 #include <vector>
 
-#include "absl/base/thread_annotations.h"
-#include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
-#include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
 #include "xla/backends/gpu/collectives/gpu_clique_key.h"
 #include "xla/backends/gpu/runtime/collective_thunk.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/backends/gpu/runtime/thunk.pb.h"
 #include "xla/core/collectives/communicator.h"
 #include "xla/core/collectives/rank_id.h"
@@ -52,11 +50,12 @@ struct AllToAllConfig {
 class AllToAllThunk : public CollectiveThunk {
  public:
   AllToAllThunk(ThunkInfo thunk_info, const HloAllToAllInstruction* instr,
-                std::vector<Buffer> buffers, bool p2p_memcpy_enabled);
+                std::vector<Buffer> buffers, bool p2p_memcpy_enabled,
+                int devices_per_host);
 
   AllToAllThunk(ThunkInfo thunk_info, const AllToAllConfig& config,
                 std::vector<CollectiveThunk::Buffer> buffers,
-                bool p2p_memcpy_enabled);
+                bool p2p_memcpy_enabled, int devices_per_host);
 
   // Returns whether the given instruction can be lowered to an all-to-all
   // call.
@@ -73,7 +72,8 @@ class AllToAllThunk : public CollectiveThunk {
 
   static absl::StatusOr<std::unique_ptr<AllToAllThunk>> FromProto(
       ThunkInfo thunk_info, const AllToAllThunkProto& thunk_proto,
-      absl::Span<const BufferAllocation> buffer_allocations);
+      absl::Span<const BufferAllocation> buffer_allocations,
+      int devices_per_host);
 
   absl::StatusOr<ThunkProto> ToProto() const override;
 
@@ -91,23 +91,24 @@ class AllToAllThunk : public CollectiveThunk {
   bool CanUseSymmetricBuffer() const override { return true; }
 
  private:
+  struct DeviceState {
+    // A uint64_t array of size num_devices. The array is used in each call to
+    // RunCollective(), but is preallocated as CUDA host memory and written to
+    // in the first call to Initialize(), since addresses won't change across
+    // calls to RunCollective().
+    std::unique_ptr<se::MemoryAllocation> receive_pointer_map;
+    // Event to synchronize streams on different devices at the start/end of the
+    // kernel.
+    std::unique_ptr<se::Event> event;
+    // Events for all ranks in the clique, populated once during Initialize().
+    // Not internally synchronized; relies on host-side thunk initialization and
+    // execution being serialized per device.
+    std::vector<se::Event*> events;
+  };
+
   const AllToAllConfig config_;
   bool p2p_memcpy_enabled_ = false;
-
-  absl::Mutex pointer_maps_mutex_;
-  // Maps from a device to a uint64_t array of size num_devices. The array is
-  // used in each call to RunCollective(), but is preallocated as CUDA host
-  // memory and written to in the first call to Initialize(), since addresses
-  // won't change across calls to RunCollective().
-  absl::flat_hash_map<se::StreamExecutor*,
-                      std::unique_ptr<se::MemoryAllocation>>
-      receive_pointer_maps_ ABSL_GUARDED_BY(pointer_maps_mutex_);
-
-  absl::Mutex events_mutex_;
-  // Events to synchronize steams on different devices at the start/end of the
-  // kernel.
-  absl::flat_hash_map<se::StreamExecutor*, std::unique_ptr<se::Event>> events_
-      ABSL_GUARDED_BY(events_mutex_);
+  PerDeviceState<DeviceState> per_device_states_;
 };
 
 absl::Status RunAllToAll(bool has_split_dimension,
@@ -121,7 +122,7 @@ absl::Status RunMemCpyAllToAll(bool has_split_dimension,
                                uint64_t receive_pointer_map[],
                                const GpuCliqueKey& clique_key, RankId rank,
                                se::Event* event,
-                               std::vector<se::Event*>& events);
+                               absl::Span<se::Event* const> events);
 
 }  // namespace gpu
 }  // namespace xla
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk_multigpu_test.cc b/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk_multigpu_test.cc
index 1b4babf5f99..33a8025e1f9 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk_multigpu_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk_multigpu_test.cc
@@ -70,7 +70,8 @@ static std::vector<BufferAllocation> MakeThunkBufferAllocations() {
   return allocations;
 }
 
-static AllToAllThunk MakeThunk(absl::Span<const BufferAllocation> allocations) {
+static AllToAllThunk MakeThunk(absl::Span<const BufferAllocation> allocations,
+                               bool p2p_memcpy_enabled = false) {
   std::vector<CollectiveThunk::Buffer> buffers;
   buffers.reserve(kNumBuffers);
   for (int i = 0; i < kNumBuffers; ++i) {
@@ -88,8 +89,8 @@ static AllToAllThunk MakeThunk(absl::Span<const BufferAllocation> allocations) {
   }
 
   return AllToAllThunk(Thunk::ThunkInfo(), MakeAllToAllConfig(),
-                       std::move(buffers),
-                       /*p2p_memcpy_enabled=*/false);
+                       std::move(buffers), p2p_memcpy_enabled,
+                       /*devices_per_host=*/kNumDevices);
 }
 
 using DeviceTestSlot = CollectiveThunkMultiGpuTestState;
@@ -254,6 +255,40 @@ TEST(AllToAllThunkMultiGpuTest, ExecuteOnStream) {
       }));
 }
 
+TEST(AllToAllThunkMultiGpuTest, ExecuteOnStreamP2PMemcpy) {
+  if (!HasEnoughGpus(kNumDevices)) {
+    GTEST_SKIP() << "Test requires at least " << kNumDevices << " GPUs";
+  }
+
+  DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices);
+  std::vector<BufferAllocation> buffer_allocations =
+      MakeThunkBufferAllocations();
+  AllToAllThunk thunk =
+      MakeThunk(buffer_allocations, /*p2p_memcpy_enabled=*/true);
+  std::vector<DeviceTestSlot> slots(kNumDevices);
+
+  ASSERT_OK(RunOnDevices(
+      kNumDevices, "alltoall_p2p_execute", [&](int d) -> absl::Status {
+        ABSL_RETURN_IF_ERROR(SetupDeviceSlot(d, slots[d], thunk, device_assignment));
+        ABSL_RETURN_IF_ERROR(
+            RunExecuteOnStreamPhase(slots[d], thunk, d, /*phase=*/1));
+
+        BufferAllocations allocations =
+            MakeBufferAllocations(slots[d], slots[d].create_buffers);
+        Thunk::InitializeParams init_params;
+        init_params.executor = slots[d].executor;
+        init_params.stream = slots[d].stream.get();
+        init_params.command_buffer_trace_stream = slots[d].stream.get();
+        init_params.buffer_allocations = &allocations;
+        init_params.collective_params = &*slots[d].collective_params;
+        init_params.collective_cliques = &slots[d].collective_cliques;
+        init_params.local_device_count = kNumDevices;
+        ABSL_RETURN_IF_ERROR(thunk.Initialize(init_params));
+
+        return RunExecuteOnStreamPhase(slots[d], thunk, d, /*phase=*/2);
+      }));
+}
+
 TEST(AllToAllThunkMultiGpuTest, RecordCommandBufferCreate) {
   if (!HasEnoughGpus(kNumDevices)) {
     GTEST_SKIP() << "Test requires at least " << kNumDevices << " GPUs";
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk_test.cc
index 4238c814879..c5c687086ac 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk_test.cc
@@ -200,7 +200,7 @@ class MemzeroAllToAllThunk : public AllToAllThunk {
   MemzeroAllToAllThunk(Thunk::ThunkInfo thunk_info, AllToAllConfig config,
                        std::vector<CollectiveThunk::Buffer> buffers)
       : AllToAllThunk(std::move(thunk_info), config, std::move(buffers),
-                      /*p2p_memcpy_enabled=*/false) {}
+                      /*p2p_memcpy_enabled=*/false, /*devices_per_host=*/1) {}
 
  protected:
   absl::Status RunCollective(const ExecuteParams& params, const GpuCliqueKey&,
@@ -241,7 +241,7 @@ TEST(CollectiveThunkTest, ProtoRoundTrip) {
   ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<AllToAllThunk> thunk,
       AllToAllThunk::FromProto(thunk_info, proto.all_to_all_thunk(),
-                               buffer_allocations));
+                               buffer_allocations, /*devices_per_host=*/1));
 
   ASSERT_OK_AND_ASSIGN(ThunkProto round_trip_proto, thunk->ToProto());
 
diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
index 7d047a3feac..d71fc6b2709 100644
--- a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
@@ -291,9 +291,10 @@ absl::StatusOr<std::unique_ptr<Thunk>> DeserializeThunkProtoImpl(
                                            thunk_proto.reduce_scatter_thunk(),
                                            buffer_allocations);
     case ThunkProto::kAllToAllThunk:
-      return AllToAllThunk::FromProto(std::move(thunk_info),
-                                      thunk_proto.all_to_all_thunk(),
-                                      buffer_allocations);
+      return AllToAllThunk::FromProto(
+          std::move(thunk_info), thunk_proto.all_to_all_thunk(),
+          buffer_allocations,
+          gpu_topology.has_value() ? gpu_topology->num_devices_per_host() : 0);
     case ThunkProto::kRaggedAllToAllThunk:
       return RaggedAllToAllThunk::FromProto(
           std::move(thunk_info), thunk_proto.ragged_all_to_all_thunk(),
diff --git a/third_party/xla/xla/service/gpu/thunk_emitter.cc b/third_party/xla/xla/service/gpu/thunk_emitter.cc
index 285180bd727..f6304494b57 100644
--- a/third_party/xla/xla/service/gpu/thunk_emitter.cc
+++ b/third_party/xla/xla/service/gpu/thunk_emitter.cc
@@ -2273,12 +2273,6 @@ Future<ThunkSequence> ThunkEmitter::EmitCollective(
         ir_emitter_context_->gpu_topology().num_devices_per_host(),
         ir_emitter_context_->debug_options().xla_gpu_use_memcpy_local_p2p(),
         has_dynamic_root);
-  } else if constexpr (std::is_same_v<CollectiveThunkType,
-                                      RaggedAllToAllThunk>) {
-    thunks = ThunkSequence::Of<CollectiveThunkType>(
-        info, inst, /*buffers=*/std::move(buffers),
-        ir_emitter_context_->debug_options().xla_gpu_use_memcpy_local_p2p(),
-        ir_emitter_context_->gpu_topology().num_devices_per_host());
   } else if constexpr (std::is_constructible_v<
                            CollectiveThunkType, Thunk::ThunkInfo,
                            decltype(inst),
@@ -2288,7 +2282,8 @@ Future<ThunkSequence> ThunkEmitter::EmitCollective(
   } else {
     thunks = ThunkSequence::Of<CollectiveThunkType>(
         info, inst, /*buffers=*/std::move(buffers),
-        ir_emitter_context_->debug_options().xla_gpu_use_memcpy_local_p2p());
+        ir_emitter_context_->debug_options().xla_gpu_use_memcpy_local_p2p(),
+        ir_emitter_context_->gpu_topology().num_devices_per_host());
   }
   return thunks;
 }

From d6d055ff1805b8b7a68a7ebbd3ea13133bb0737a Mon Sep 17 00:00:00 2001
From: Ilya Tikhonovskiy <loislo@google.com>
Date: Mon, 12 Oct 2026 00:48:06 -0700
Subject: [PATCH 0558/1000] [XLA] Optimize HloEvaluator::HandleReduceWindow and
 ReferenceUtil::ReduceWindow4DGeneric.

Profiling `reduce_window_test_gpu` showed that over 90% of total CPU time (248.3s of 275.9s CPU samples) was spent computing reference outputs on the CPU rather than compiling or running on the GPU:
1. `HloEvaluator::HandleReduceWindow` (81.7% of CPU) invoked `HloEvaluator::Evaluate` (a full `PostOrderDFS` traversal plus `Literal` allocation, copy, and destruction) and `ResetVisitStates()` for every window element, and `IterateThroughWindow` iterated through out-of-bounds padded positions and base-dilation holes while calling `IndexUtil::BumpIndices` per element.
2. `ReferenceUtil::ReduceWindow4DGeneric` (10.4% of CPU) checked 8 bounds conditions in the innermost loop across padded regions and recomputed 4D linear indices per element.

Precompute the valid in-bounds arithmetic progression `(first_base, base_step, count)` per dimension in `IterateThroughWindow`, add a direct native-type fast path for scalar binary reducers (`kAdd`, `kMultiply`, `kMaximum`, `kMinimum`, `kAnd`, `kOr`, `kXor`) in `HloEvaluator::HandleReduceWindow`, and clamp the 4D window loop bounds with contiguous row pointers in `ReferenceUtil::ReduceWindow4DGeneric`.

This reduces `hlo_evaluator_slow_reduce_window_test` from 13,546 ms to 77 ms (176x speedup) and reduces unsharded `reduce_window_test_gpu` total CPU time from 275.9s to 39.1s (wall time from 127.5s to 35.8s).

PiperOrigin-RevId: 997738801
---
 .../xla/xla/hlo/evaluator/hlo_evaluator.cc    | 248 +++++++++++++++---
 third_party/xla/xla/reference_util.cc         |  45 ++--
 2 files changed, 233 insertions(+), 60 deletions(-)

diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
index 3e01a8c8ba5..ed3179972b7 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
@@ -27,9 +27,11 @@ limitations under the License.
 #include <iterator>
 #include <limits>
 #include <memory>
+#include <numeric>
 #include <optional>
 #include <random>
 #include <string>
+#include <type_traits>
 #include <utility>
 #include <variant>
 #include <vector>
@@ -4005,54 +4007,81 @@ absl::StatusOr<Literal> ExtractFromIndexPositions(
 void IterateThroughWindow(
     const Shape& window_shape, const Window& window, const Shape& base_shape,
     const absl::Span<const int64_t> window_count_index,
-    const std::function<void(absl::Span<const int64_t>)>& f) {
+    absl::FunctionRef<void(absl::Span<const int64_t>)> f) {
   const size_t rank = base_shape.dimensions().size();
-  DimensionVector window_index(rank);
-  absl::Span<int64_t> window_index_span = absl::MakeSpan(window_index);
-  std::fill(window_index_span.begin(), window_index_span.end(), 0);
+  if (rank == 0) {
+    if (!ShapeUtil::IsZeroElementArray(window_shape) &&
+        !ShapeUtil::IsZeroElementArray(base_shape)) {
+      f({});
+    }
+    return;
+  }
+  DimensionVector first_base(rank);
+  DimensionVector base_step(rank);
+  DimensionVector count(rank);
+  for (size_t i = 0; i < rank; ++i) {
+    const int64_t win_size = window_shape.dimensions(i);
+    const int64_t base_limit = base_shape.dimensions(i);
+    if (win_size <= 0 || base_limit <= 0) {
+      return;
+    }
+    const WindowDimension& dim = window.dimensions(i);
+    const int64_t win_dilation = dim.window_dilation();
+    const int64_t base_dilation = dim.base_dilation();
+    const int64_t start_unscaled =
+        window_count_index[i] * dim.stride() - dim.padding_low();
 
-  do {
-    DimensionVector base_index(rank);
-    absl::Span<int64_t> base_index_span = absl::MakeSpan(base_index);
-    bool out_of_bound = false;
-    for (size_t i = 0; i < rank; ++i) {
-      // Padding is applied to the dilated base. Say that padding is 3 and
-      // dilation is 2 for some dimension. After applying base dilation and
-      // padding, the dimension looks like:
-      // P P P E D D E D D ... E D D E P P P
-      // where E are the elements and D are the holes. So, the elements are
-      // located in indices: padding + k*base_dilation for k = {0, 1, 2, ...}.
-      // We are accessing elements in the transformed base at indices:
-      // window_count_index * stride + window_index * window_dilation.
-      // Solving for k gives us
-      // (win_count_i * stride + win_i * win_dilation - pad) / base_dilation
-      // When this is a natural number, we index an original element.
-      // Otherwise, we index a 0 (pad or hole), and we don't need to apply
-      // the callback f.
-      const xla::WindowDimension& window_dimension = window.dimensions(i);
-      int64_t base_index_i =
-          window_count_index[i] * window_dimension.stride() +
-          window_index_span[i] * window_dimension.window_dilation() -
-          window_dimension.padding_low();
-      const int64_t base_dilation = window_dimension.base_dilation();
-      if (base_dilation != 1) {
-        if (base_index_i % base_dilation != 0) {
-          out_of_bound = true;
-          base_index_span[i] = base_index_i;
-          break;
-        }
-        base_index_i /= base_dilation;
+    int64_t w_min = 0;
+    if (start_unscaled < 0) {
+      w_min = (-start_unscaled + win_dilation - 1) / win_dilation;
+    }
+    if (base_dilation != 1) {
+      const int64_t w_search_limit = std::min(win_size, w_min + base_dilation);
+      while (w_min < w_search_limit &&
+             (start_unscaled + w_min * win_dilation) % base_dilation != 0) {
+        ++w_min;
       }
-      base_index_span[i] = base_index_i;
-      if (base_index_i < 0 || base_index_i >= base_shape.dimensions(i)) {
-        out_of_bound = true;
+      if (w_min >= w_search_limit) {
+        return;
+      }
+    }
+    if (w_min >= win_size) {
+      return;
+    }
+    const int64_t first_base_i =
+        (start_unscaled + w_min * win_dilation) / base_dilation;
+    if (first_base_i >= base_limit) {
+      return;
+    }
+    const int64_t g = std::gcd(win_dilation, base_dilation);
+    const int64_t w_step = base_dilation / g;
+    const int64_t base_step_i = win_dilation / g;
+    const int64_t count_by_win = (win_size - 1 - w_min) / w_step + 1;
+    const int64_t count_by_base =
+        (base_limit - 1 - first_base_i) / base_step_i + 1;
+    first_base[i] = first_base_i;
+    base_step[i] = base_step_i;
+    count[i] = std::min(count_by_win, count_by_base);
+  }
+  DimensionVector base_index = first_base;
+  absl::Span<const int64_t> base_index_span(base_index);
+  DimensionVector idx(rank, 0);
+  while (true) {
+    f(base_index_span);
+    int64_t d = static_cast<int64_t>(rank) - 1;
+    while (d >= 0) {
+      if (++idx[d] < count[d]) {
+        base_index[d] += base_step[d];
         break;
       }
+      idx[d] = 0;
+      base_index[d] = first_base[d];
+      --d;
     }
-    if (!out_of_bound) {
-      f(base_index_span);
+    if (d < 0) {
+      break;
     }
-  } while (IndexUtil::BumpIndices(window_shape, window_index_span));
+  }
 }
 
 template <typename Fp, typename Uint, typename ResultT>
@@ -4862,6 +4891,126 @@ absl::Status HloEvaluator::HandleReduce(const HloInstruction* hlo) {
   return absl::OkStatus();
 }
 
+namespace {
+
+std::optional<HloOpcode> MatchSimpleBinaryReducer(
+    const HloComputation* computation, PrimitiveType elem_type) {
+  if (computation->num_parameters() != 2) {
+    return std::nullopt;
+  }
+  const HloInstruction* op = computation->root_instruction();
+  if (op->operand_count() != 2 ||
+      op->operand(0) != computation->parameter_instruction(0) ||
+      op->operand(1) != computation->parameter_instruction(1) ||
+      !ShapeUtil::IsScalarWithElementType(op->shape(), elem_type)) {
+    return std::nullopt;
+  }
+  const auto op_code = std::make_optional(op->opcode());
+  switch (op->opcode()) {
+    case HloOpcode::kAdd:
+    case HloOpcode::kMultiply:
+      return elem_type != PRED ? op_code : std::nullopt;
+    case HloOpcode::kMaximum:
+    case HloOpcode::kMinimum:
+      return !primitive_util::IsComplexType(elem_type) ? op_code : std::nullopt;
+    case HloOpcode::kAnd:
+    case HloOpcode::kOr:
+    case HloOpcode::kXor:
+      return (primitive_util::IsIntegralType(elem_type) || elem_type == PRED)
+                 ? op_code
+                 : std::nullopt;
+    default:
+      return std::nullopt;
+  }
+}
+
+template <typename ReturnT, typename ElementwiseT>
+ReturnT ApplyBinaryOp(HloOpcode opcode, ReturnT a, ReturnT b) {
+  ElementwiseT lhs = static_cast<ElementwiseT>(a);
+  ElementwiseT rhs = static_cast<ElementwiseT>(b);
+  switch (opcode) {
+    case HloOpcode::kAdd:
+      return static_cast<ReturnT>(
+          ElementwiseT(ToArithmeticSafeType(lhs) + ToArithmeticSafeType(rhs)));
+    case HloOpcode::kMultiply:
+      return static_cast<ReturnT>(
+          ElementwiseT(ToArithmeticSafeType(lhs) * ToArithmeticSafeType(rhs)));
+    case HloOpcode::kMaximum:
+      if constexpr (!is_complex_v<ElementwiseT>) {
+        if constexpr (std::numeric_limits<ElementwiseT>::has_quiet_NaN) {
+          if (std::isnan(lhs)) {
+            return static_cast<ReturnT>(lhs);
+          }
+          if (std::isnan(rhs)) {
+            return static_cast<ReturnT>(rhs);
+          }
+        }
+        return static_cast<ReturnT>(std::max(lhs, rhs));
+      }
+      break;
+    case HloOpcode::kMinimum:
+      if constexpr (!is_complex_v<ElementwiseT>) {
+        if constexpr (std::numeric_limits<ElementwiseT>::has_quiet_NaN) {
+          if (std::isnan(lhs)) {
+            return static_cast<ReturnT>(lhs);
+          }
+          if (std::isnan(rhs)) {
+            return static_cast<ReturnT>(rhs);
+          }
+        }
+        return static_cast<ReturnT>(std::min(lhs, rhs));
+      }
+      break;
+    case HloOpcode::kAnd:
+      if constexpr (std::is_integral_v<ElementwiseT>) {
+        return static_cast<ReturnT>(lhs & rhs);
+      }
+      break;
+    case HloOpcode::kOr:
+      if constexpr (std::is_integral_v<ElementwiseT>) {
+        return static_cast<ReturnT>(lhs | rhs);
+      }
+      break;
+    case HloOpcode::kXor:
+      if constexpr (std::is_integral_v<ElementwiseT>) {
+        return static_cast<ReturnT>(lhs ^ rhs);
+      }
+      break;
+    default:
+      break;
+  }
+  LOG(FATAL) << "Unsupported opcode in ApplyBinaryOp: "
+             << HloOpcodeString(opcode);
+}
+
+template <PrimitiveType kType>
+absl::Status EvaluateSimpleReduceWindowTyped(
+    HloOpcode opcode, const Literal& input, const Literal& init,
+    const Shape& window_shape, const Window& window, Literal& result) {
+  using ReturnT = NativeTypeOf<kType>;
+  using ElementwiseT = std::conditional_t<
+      primitive_util::IsSignedIntegralType(kType), int64_t,
+      std::conditional_t<
+          primitive_util::IsUnsignedIntegralType(kType), uint64_t,
+          std::conditional_t<primitive_util::IsFloatingPointType(kType) &&
+                                 sizeof(ReturnT) < sizeof(float),
+                             float, ReturnT>>>;
+  const ReturnT init_val = init.Get<ReturnT>({});
+  return result.PopulateParallel<ReturnT>(
+      [&](absl::Span<const int64_t> output_index, int) -> ReturnT {
+        ReturnT acc = init_val;
+        IterateThroughWindow(window_shape, window, input.shape(), output_index,
+                             [&](absl::Span<const int64_t> input_index) {
+                               acc = ApplyBinaryOp<ReturnT, ElementwiseT>(
+                                   opcode, acc,
+                                   input.Get<ReturnT>(input_index));
+                             });
+        return acc;
+      });
+}
+
+}  // namespace
+
 absl::Status HloEvaluator::HandleReduceWindow(const HloInstruction* hlo) {
   auto* reduce_window = Cast<HloReduceWindowInstruction>(hlo);
   const Window& window = reduce_window->window();
@@ -4899,6 +5048,23 @@ absl::Status HloEvaluator::HandleReduceWindow(const HloInstruction* hlo) {
   const Shape window_shape = ShapeUtil::MakeShape(
       input_arrays[0]->shape().element_type(), window_dimension_sizes);
 
+  if (!inferred_return_shape.IsTuple() && num_args == 1) {
+    if (std::optional<HloOpcode> fast_op = MatchSimpleBinaryReducer(
+            function, inferred_return_shape.element_type())) {
+      Literal result(inferred_return_shape);
+      ABSL_RETURN_IF_ERROR(primitive_util::ArrayTypeSwitch(
+          [&](auto type) {
+            return EvaluateSimpleReduceWindowTyped<type>(
+                *fast_op, *input_literal_vec[0], *init_literal_vec[0],
+                window_shape, window, result);
+          },
+          inferred_return_shape.element_type()));
+      VLOG(2) << "Final result is:" << result.ToString() << "\n";
+      SetEvaluatedLiteralFor(reduce_window, std::move(result));
+      return absl::OkStatus();
+    }
+  }
+
   const int num_threads = ShapeUtil::GetForEachIndexParallelThreadCount() + 1;
   std::vector<std::unique_ptr<HloEvaluator>> embedded_evaluators;
   embedded_evaluators.reserve(num_threads);
diff --git a/third_party/xla/xla/reference_util.cc b/third_party/xla/xla/reference_util.cc
index 78176c2561b..031c83bb1e6 100644
--- a/third_party/xla/xla/reference_util.cc
+++ b/third_party/xla/xla/reference_util.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/reference_util.h"
 
+#include <algorithm>
 #include <array>
 #include <cmath>
 #include <cstdint>
@@ -220,36 +221,42 @@ ReferenceUtil::ReduceWindow4DGeneric(
   }
   auto result = std::make_unique<Array4D<float>>(
       window_counts[0], window_counts[1], window_counts[2], window_counts[3]);
+  const int64_t n2 = operand.n2();
+  const int64_t n3 = operand.n3();
+  const int64_t n4 = operand.n4();
+  const float* operand_data = operand.data();
+  float* out_ptr = result->data();
   // Do a full 4D reduce window.
   for (int64_t i0 = 0; i0 < window_counts[0]; ++i0) {
+    const int64_t i0_base = i0 * stride[0] - pad_low[0];
+    const int64_t b0_start = std::max<int64_t>(0, i0_base);
+    const int64_t b0_end = std::min<int64_t>(operand.n1(), i0_base + window[0]);
     for (int64_t i1 = 0; i1 < window_counts[1]; ++i1) {
+      const int64_t i1_base = i1 * stride[1] - pad_low[1];
+      const int64_t b1_start = std::max<int64_t>(0, i1_base);
+      const int64_t b1_end = std::min<int64_t>(n2, i1_base + window[1]);
       for (int64_t i2 = 0; i2 < window_counts[2]; ++i2) {
+        const int64_t i2_base = i2 * stride[2] - pad_low[2];
+        const int64_t b2_start = std::max<int64_t>(0, i2_base);
+        const int64_t b2_end = std::min<int64_t>(n3, i2_base + window[2]);
         for (int64_t i3 = 0; i3 < window_counts[3]; ++i3) {
-          int64_t i0_base = i0 * stride[0] - pad_low[0];
-          int64_t i1_base = i1 * stride[1] - pad_low[1];
-          int64_t i2_base = i2 * stride[2] - pad_low[2];
-          int64_t i3_base = i3 * stride[3] - pad_low[3];
+          const int64_t i3_base = i3 * stride[3] - pad_low[3];
+          const int64_t b3_start = std::max<int64_t>(0, i3_base);
+          const int64_t b3_end = std::min<int64_t>(n4, i3_base + window[3]);
 
           float val = init;
-          for (int64_t i0_win = 0; i0_win < window[0]; ++i0_win) {
-            for (int64_t i1_win = 0; i1_win < window[1]; ++i1_win) {
-              for (int64_t i2_win = 0; i2_win < window[2]; ++i2_win) {
-                for (int64_t i3_win = 0; i3_win < window[3]; ++i3_win) {
-                  if (i0_base + i0_win >= 0 && i1_base + i1_win >= 0 &&
-                      i2_base + i2_win >= 0 && i3_base + i3_win >= 0 &&
-                      i0_base + i0_win < operand.n1() &&
-                      i1_base + i1_win < operand.n2() &&
-                      i2_base + i2_win < operand.n3() &&
-                      i3_base + i3_win < operand.n4()) {
-                    val = reduce_func(
-                        val, operand(i0_base + i0_win, i1_base + i1_win,
-                                     i2_base + i2_win, i3_base + i3_win));
-                  }
+          for (int64_t b0 = b0_start; b0 < b0_end; ++b0) {
+            for (int64_t b1 = b1_start; b1 < b1_end; ++b1) {
+              const int64_t b01_offset = (b0 * n2 + b1) * n3;
+              for (int64_t b2 = b2_start; b2 < b2_end; ++b2) {
+                const float* row = operand_data + (b01_offset + b2) * n4;
+                for (int64_t b3 = b3_start; b3 < b3_end; ++b3) {
+                  val = reduce_func(val, row[b3]);
                 }
               }
             }
           }
-          (*result)(i0, i1, i2, i3) = val;
+          *out_ptr++ = val;
         }
       }
     }

From 26b91b0869a7e74a10d8d9075b814377efe6e98a Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 21:20:48 +0530
Subject: [PATCH 0559/1000] Address review: add empty-local_devices test, fix
 subject-verb agreement

- Add test_create_distributed_mesh_empty_local_devices, asserting the
  ValueError('Received an empty `devices` list...') guard added in
  _make_device_specs for an empty local_devices list.
- Fix "meshes that use all TPU cores is supported" -> "...are supported"
  in create_distributed_mesh's docstring (plural subject, plural verb).
---
 tensorflow/dtensor/python/mesh_util.py            |  2 +-
 tensorflow/dtensor/python/tests/mesh_util_test.py | 12 ++++++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/tensorflow/dtensor/python/mesh_util.py b/tensorflow/dtensor/python/mesh_util.py
index 6e2e962bc2a..45771457341 100644
--- a/tensorflow/dtensor/python/mesh_util.py
+++ b/tensorflow/dtensor/python/mesh_util.py
@@ -155,7 +155,7 @@ def create_distributed_mesh(
   For CPU and GPU meshes, users can choose to use fewer local devices than what
   is available `local_devices`.
 
-  For TPU, only meshes that use all TPU cores is supported by the DTensor
+  For TPU, only meshes that use all TPU cores are supported by the DTensor
   runtime.
 
   Args:
diff --git a/tensorflow/dtensor/python/tests/mesh_util_test.py b/tensorflow/dtensor/python/tests/mesh_util_test.py
index 567082357eb..07503c6352f 100644
--- a/tensorflow/dtensor/python/tests/mesh_util_test.py
+++ b/tensorflow/dtensor/python/tests/mesh_util_test.py
@@ -177,6 +177,18 @@ class MeshUtilTest(test_util.DTensorBaseTest):
             device_type='TPU',
         )
 
+  def test_create_distributed_mesh_empty_local_devices(self):
+    self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
+    with mock.patch.object(accelerator_util, 'is_initialized', return_value=True):
+      with self.assertRaisesRegex(
+          ValueError, 'Received an empty `devices` list'
+      ):
+        mesh_util.create_distributed_mesh(
+            mesh_name='empty_mesh',
+            mesh_dims=[('x', 1)],
+            local_devices=[],
+        )
+
   def test_single_client_mesh_creation_wrong_shape(self):
     self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
     num_devices = len(test_util.list_local_logical_devices('CPU'))

From d0f716cb853cd58f7d0d1219a71ebebaa1e55ee5 Mon Sep 17 00:00:00 2001
From: Maxime France-Pillois <mfrancep@amd.com>
Date: Mon, 12 Oct 2026 00:52:41 -0700
Subject: [PATCH 0560/1000] PR #45998: [GPU][Triton][RaggedDot] Emitter:
 kRaggedContracting mode (PR 4/8)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/45998

📝 Summary of Changes
This PR is the fourth in a series of eight aiming to enable the Triton backend for the `ragged-dot` operation using triton group-gemm kernels, based on the experimental tiling strategy.

PR 1/8: https://github.com/openxla/xla/pull/45995
PR 2/8: https://github.com/openxla/xla/pull/45996
PR 3/8: https://github.com/openxla/xla/pull/45997
PR 4/8: https://github.com/openxla/xla/pull/45998  <---- This PR
PR 5/8: https://github.com/openxla/xla/pull/45999
PR 6/8: https://github.com/openxla/xla/pull/46000
PR 7/8: https://github.com/openxla/xla/pull/46001
PR 8/8: https://github.com/openxla/xla/pull/46002

This PR adds the kRaggedContracting implementation in EmitRaggedDot().
In kRaggedContracting the ragged dimension is the K (contracting) axis: Grid = G x K_tiles x N_tiles. Each program uses a prefix-sum over group_sizes to locate its start_m offset, then accumulates a partial dot for the slice of M rows belonging to its group.

🎯 Justification
The default way for handling ragged-dot ops is to expand the ragged-dot matrix into zero-padded matrices then rewriting the ragged-dot into a batched dot between the padded matrices and the second operand, which results in numerous unnecessary calculations.
Although there are other, more efficient backends (CuDNN and Hipblast), these are not capable of handling all possible configurations.
In this context, having another, more generic Triton backend could enable good performance for ‘ragged-dot’ operations on all targets and for all input configurations.
The order in which backends are selected is as follows:
1. Cudnn if compatible Cuda target or Hipblaslt if compatible HIP target
2. Triton backend if compatible input configuration
3. Default rewriting

🚀 Kind of Contribution
Please remove what does not apply: ⚡️ Performance Improvement, ✨ New Feature

📊 Benchmark (for Performance Improvements)

A few performance numbers:

### On H100:

#### Ragged dimension in non-contracting dimension, dtype = f16

| M | K | N | G | Triton (µs) | Triton (TFLOPS) | Baseline (µs) | Baseline (TFLOPS) |
|---:|---:|---:|---:|---:|---:|---:|---:|
| 32,768 | 6,144 | 16,384 | 8 | 20,966 | 314.65 | 65,785 | 100.28 |
| 32,768 | 16,384 | 6,144 | 8 | 22,489 | 293.35 | 65,128 | 101.29 |
| 49,152 | 1,408 | 2,048 | 64 | 1,307 | 216.95 | 35,787 | 7.92 |
| 393,216 | 2,048 | 1,408 | 64 | 7,315 | 310.03 | – | – |
| 3,145,728 | 2,048 | 1,408 | 8 | 53,108 | 341.60 | – | – |

#### Ragged dimension in contracting dimension, dtype = f16

| M | K | N | G | Triton (µs) | Triton (TFLOPS) | Baseline (µs) | Baseline (TFLOPS) |
|---:|---:|---:|---:|---:|---:|---:|---:|
| 32,768 | 6,144 | 16,384 | 8 | 20,041 | 329.18 | 71,325 | 92.49 |
| 32,768 | 16,384 | 6,144 | 8 | 20,898 | 315.68 | 72,130 | 91.46 |
| 49,152 | 1,408 | 2,048 | 64 | 954 | 297.16 | 51,334 | 5.52 |
| 393,216 | 2,048 | 1,408 | 64 | 6,881 | 329.56 | – | – |
| 3,145,728 | 2,048 | 1,408 | 8 | 61,241 | 296.24 | – | – |

### On MI355:

#### Ragged dimension in non-contracting dimension, dtype = bf16

| M | K | N | G | Triton (µs) | Triton (TFLOPS) | Baseline (µs) | Baseline (TFLOPS) |
|---:|---:|---:|---:|---:|---:|---:|---:|
| 32,768 | 6,144 | 16,384 | 8 | 7,306 | 902.92 | 32,238 | 204.63 |
| 32,768 | 16,384 | 6,144 | 8 | 7,265 | 908.09 | 33,316 | 198.01 |
| 49,152 | 1,408 | 2,048 | 64 | 374 | 757.35 | 12,462 | 22.75 |
| 393,216 | 2,048 | 1,408 | 64 | 2,810 | 806.90 | 125,435 | 18.08 |

#### Ragged dimension in contracting dimension, dtype = bf16

| M | K | N | G | Triton (µs) | Triton (TFLOPS) | Baseline (µs) | Baseline (TFLOPS) |
|---:|---:|---:|---:|---:|---:|---:|---:|
| 32,768 | 6,144 | 16,384 | 8 | 8,659 | 761.88 | 31,464 | 209.67 |
| 32,768 | 16,384 | 6,144 | 8 | 8,770 | 752.23 | 35,697 | 184.81 |
| 49,152 | 1,408 | 2,048 | 64 | 489 | 579.42 | 17,322 | 16.36 |
| 393,216 | 2,048 | 1,408 | 64 | 3,397 | 667.59 | 179,456 | 12.64 |

### On MI300:

#### Ragged dimension in non-contracting dimension, dtype = bf16

| M | K | N | G | Triton (µs) | Triton (TFLOPS) | hipBLASLt (µs) | hipBLASLt (TFLOPS) | Baseline (µs) | Baseline (TFLOPS) |
|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|
| 32,768 | 6,144 | 16,384 | 8 | 15,904 | 414.81 | 12,337 | 534.72 | 67,215 | 98.15 |
| 32,768 | 16,384 | 6,144 | 8 | 17,399 | 379.16 | 12,427 | 530.85 | 71,709 | 92.00 |
| 49,152 | 1,408 | 2,048 | 64 | 803 | 352.90 | 675 | 419.74 | 24,735 | 11.46 |
| 393,216 | 2,048 | 1,408 | 64 | 5,584 | 406.09 | 5,115 | 443.35 | 247,734 | 9.15 |

#### Ragged dimension in contracting dimension, dtype = bf16

| M | K | N | G | Triton (µs) | Triton (TFLOPS) | hipBLASLt (µs) | hipBLASLt (TFLOPS) | Baseline (µs) | Baseline (TFLOPS) |
|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|
| 32,768 | 6,144 | 16,384 | 8 | 17,343 | 380.38 | 11,481 | 574.63 | 70,294 | 93.85 |
| 32,768 | 16,384 | 6,144 | 8 | 17,665 | 373.45 | 11,380 | 579.69 | 71,224 | 92.62 |
| 49,152 | 1,408 | 2,048 | 64 | 918 | 308.93 | 656 | 432.11 | 36,757 | 7.71 |
| 393,216 | 2,048 | 1,408 | 64 | 6,028 | 376.18 | 5,169 | 438.73 | – | – |

Copybara import of the project:

--
900cae31530cb15d9b6b15abd1316d91ffa95dc9 by Maxime France-Pillois <mfrancep@amd.com>:

[XTile][RaggedDot] Emitter: kRaggedContracting mode

Add the kRaggedContracting implementation in EmitRaggedDot().

In kRaggedContracting the ragged dimension is the K (contracting) axis:
Grid = G x K_tiles x N_tiles. Each program uses a prefix-sum over
group_sizes to locate its start_m offset, then accumulates a partial
dot for the slice of M rows belonging to its group.

kRaggedBatch still returns Unimplemented for now.

Merging this change closes #45998

PiperOrigin-RevId: 997740617
---
 .../codegen/experimental_fusion_emitter.cc    | 1202 +++++++++++------
 1 file changed, 812 insertions(+), 390 deletions(-)

diff --git a/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc b/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc
index b9d9be8059d..e8b4bb74d0a 100644
--- a/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc
+++ b/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc
@@ -134,6 +134,16 @@ ArrayRef<T> MakeArrayRef(const absl::Span<const T> span) {
   return ArrayRef(span.data(), span.size());
 }
 
+// Reshapes `input` to `new_shape`, preserving its element type, and returns the
+// result as a TensorValue.
+TensorValue ReshapeTo(mlir::ImplicitLocOpBuilder& b, TensorValue input,
+                      ArrayRef<int64_t> new_shape) {
+  auto result_type =
+      mlir::RankedTensorType::get(new_shape, input.getType().getElementType());
+  return mlir::cast<TensorValue>(
+      stablehlo::ReshapeOp::create(b, result_type, input).getResult());
+}
+
 absl::StatusOr<TensorValue> EmitAllGather(
     EmitterContext& emitter_ctx, const HloAllGatherInstruction* all_gather,
     const ge::TiledHloInstruction& tiled_all_gather, ValueRange operands) {
@@ -700,17 +710,63 @@ absl::StatusOr<TensorValue> EmitRegionInstructionWithDependencies(
   return emitter_ctx.TiledHloToTensorValue(*root);
 }
 
-// Emits a kRaggedDot instruction.
-//
-// kRaggedNonContracting (G is kSequential outer loop):
-//   Output (M_total, N).  For each group g in [0, G):
-//     - Load group_size_g from group_sizes[g] (1-element tile at G-loop IV).
-//     - Check if this M tile (pid_M * BLOCK_M) belongs to group g.
-//     - If yes, emit K-loop and accumulate partial dot into output.
-//   Loop-carried: accumulator (M, N) and last_m (prefix sum of group sizes).
-//
-// kRaggedContracting: Grid = G × K × N. Prefix-sum per program to find start_m.
-absl::StatusOr<TensorValue> EmitRaggedDot(
+// Reshapes a rank-1 (or higher, size-1) group_sizes tile to a rank-0 tensor,
+// extracts the scalar element, and index-casts it to `index` type. Shared by
+// both ragged-dot mode emitters, which each need to turn a group_sizes tile
+// into an `index` scalar group_size_g.
+Value ExtractGroupSizeScalar(mlir::ImplicitLocOpBuilder& b,
+                             TensorValue gs_tile) {
+  TensorValue gs_tile_scalar = ReshapeTo(b, gs_tile, /*new_shape=*/{});
+  Value gs_raw = mlir::tensor::ExtractOp::create(b, gs_tile_scalar);
+  return mlir::arith::IndexCastOp::create(b, b.getIndexType(), gs_raw);
+}
+
+// Resolves `hlo` (traversing through a leading kConvert and up any nested
+// fusion parameter chain) to the entry-function buffer argument that backs it.
+// The returned buffer Value carries the underlying (pre-convert) element type,
+// which callers can obtain via `getElementTypeOrSelf`.
+mlir::Value ResolveBufferArgument(EmitterContext& emitter_ctx,
+                                  const HloFusionInstruction& fusion,
+                                  const HloInstruction* hlo) {
+  if (hlo->opcode() == HloOpcode::kConvert) {
+    hlo = hlo->operand(0);
+  }
+  if (hlo->opcode() == HloOpcode::kParameter && !fusion.IsUserOf(hlo)) {
+    hlo = hlo->parent()->FusionInstruction()->operand(hlo->parameter_number());
+  }
+  int64_t arg_idx = fusion.operand_index(hlo);
+  return emitter_ctx.entry_func().getArgument(arg_idx);
+}
+
+// Applies a 1-D boolean row-mask (over `mask_dim`) to `operand`, zeroing out
+// masked-out elements. `mask_1d` must be a rank-1 i1 tensor whose length equals
+// the size of `operand` along `mask_dim`. Shared masking primitive used by the
+// ragged-dot emitters to zero rows/columns outside the current group.
+absl::StatusOr<TensorValue> ApplyRowMask(mlir::ImplicitLocOpBuilder& b,
+                                         TensorValue operand, Value mask_1d,
+                                         int64_t mask_dim, Type elem_type) {
+  auto shape = operand.getType().getShape();
+  Value mask = xtile::BroadcastInDims(b, mlir::cast<TensorValue>(mask_1d),
+                                      shape, {mask_dim});
+  TensorValue zero = CreateConst(b, elem_type, 0.0f, shape);
+  return mlir::cast<TensorValue>(
+      arith::SelectOp::create(b, mask, operand, zero).getResult());
+}
+
+// Common values derived once and shared by both ragged-dot mode emitters.
+struct RaggedDotParams {
+  const HloRaggedDotInstruction* ragged_dot_instr;
+  const RaggedDotDimensionNumbers& ragged_dims;
+  const DotDimensionNumbers& dot_dims;
+  const ge::TilingSpace& tiling_space;
+  int64_t output_rank;
+  int64_t lhs_ragged_dim;
+  SmallVector<int64_t> padded_tile_sizes;
+  Type acc_type;
+  TensorValue accumulator;
+};
+
+absl::StatusOr<RaggedDotParams> BuildRaggedDotContext(
     EmitterContext& emitter_ctx,
     const ge::TiledHloInstruction& tiled_ragged_dot) {
   auto& b = emitter_ctx.b();
@@ -720,21 +776,11 @@ absl::StatusOr<TensorValue> EmitRaggedDot(
       ragged_dot_instr->ragged_dot_dimension_numbers();
   const DotDimensionNumbers& dot_dims = ragged_dims.dot_dimension_numbers();
 
-  const int64_t lhs_ragged_dim = ragged_dims.lhs_ragged_dimensions(0);
-  const bool is_contracting =
-      absl::c_count(dot_dims.lhs_contracting_dimensions(), lhs_ragged_dim) > 0;
-  const bool is_batch =
-      absl::c_count(dot_dims.lhs_batch_dimensions(), lhs_ragged_dim) > 0;
-
   TF_RET_CHECK(tiled_ragged_dot.hlo_regions().size() == 1);
 
   ABSL_ASSIGN_OR_RETURN(SmallVector<int64_t> padded_tile_sizes,
                    tiled_ragged_dot.tile().GetStaticTileSizes());
 
-  const ge::TilingSpace& tiling_space = tiled_ragged_dot.tile().tiling_space();
-  const int64_t output_rank =
-      static_cast<int64_t>(tiled_ragged_dot.hlo()->shape().dimensions().size());
-
   // Accumulator type: derived from the ragged dot's precision config so that
   // user-specified algorithms (e.g. ALG_DOT_BF16_BF16_F32) are respected,
   // mirroring the regular EmitDot path.
@@ -742,382 +788,758 @@ absl::StatusOr<TensorValue> EmitRaggedDot(
                    xtile::GetDotAccumulatorType(b, *ragged_dot_instr));
   TensorValue accumulator = CreateConst(b, acc_type, 0.0f, padded_tile_sizes);
 
-  const ge::TiledHloRegion& region = tiled_ragged_dot.hlo_regions().front();
-  // Emits a tiled operand and its transitive deps within the region.
-  auto emit_operand = [&](const ge::TiledHloInstruction* operand_t) {
-    return EmitRegionInstructionWithDependencies(emitter_ctx, region,
-                                                 operand_t);
+  return RaggedDotParams{
+      ragged_dot_instr,
+      ragged_dims,
+      dot_dims,
+      tiled_ragged_dot.tile().tiling_space(),
+      static_cast<int64_t>(tiled_ragged_dot.hlo()->shape().dimensions().size()),
+      ragged_dims.lhs_ragged_dimensions(0),
+      std::move(padded_tile_sizes),
+      acc_type,
+      accumulator,
   };
+}
+
+// Bundles the loop-invariant values that the kRaggedNonContracting K-loop body
+// needs, computed once by EmitRaggedDotNonContracting before the G/K loops.
+struct NonContractingKLoopParams {
+  const ge::TiledHloInstruction* lhs_tiled;
+  const ge::TiledHloInstruction* rhs_tiled;
+  int64_t lhs_ragged_dim;
+  int64_t num_batch_dims;
+  // i32 comparison values precomputed before the G loop registered its IV.
+  Value tile_m_i32_cmp;
+  Value last_m_i32_cmp;
+  Value last_m_plus_i32_cmp;
+};
+
+// Zeroes LHS rows whose absolute M index falls outside the half-open range
+// [last_m, last_m_plus_gs) that belongs to the current group g. The mask is
+// built over [BLOCK_M] rows (using the pre-computed i32 comparison values from
+// `params`) and broadcast along `lhs_ragged_dim`.
+absl::StatusOr<TensorValue> ApplyGroupRowMask(
+    ImplicitLocOpBuilder& b, TensorValue lhs_tensor,
+    const NonContractingKLoopParams& params) {
+  auto lhs_shape = lhs_tensor.getType().getShape();
+  // lhs_ragged_dim is the M dimension in the LHS tile shape.
+  // For no-batch: lhs_ragged_dim=0 → lhs_shape[0]=BLOCK_M.
+  // For batched:  lhs_ragged_dim=1 → lhs_shape[1]=BLOCK_M.
+  int64_t lhs_m_dim = lhs_shape[params.lhs_ragged_dim];  // == BLOCK_M
+  TensorValue row_iota = Iota(b, static_cast<int32_t>(lhs_m_dim));
+  // abs_m[r] = tile_m_abs_i32 + r. Use the pre-computed comparison values
+  // (avoids re-evaluating the M offset inside the G-loop where sequential
+  // dims are registered).
+  TensorValue tile_m_splat =
+      xtile::Splat(b, params.tile_m_i32_cmp, {lhs_m_dim});
+  Value abs_m = arith::AddIOp::create(b, row_iota, tile_m_splat);
+  TensorValue lm_splat = xtile::Splat(b, params.last_m_i32_cmp, {lhs_m_dim});
+  TensorValue lmp_splat =
+      xtile::Splat(b, params.last_m_plus_i32_cmp, {lhs_m_dim});
+  Value mask_lo =
+      arith::CmpIOp::create(b, arith::CmpIPredicate::sge, abs_m, lm_splat);
+  Value mask_hi =
+      arith::CmpIOp::create(b, arith::CmpIPredicate::slt, abs_m, lmp_splat);
+  Value row_mask = arith::AndIOp::create(b, mask_lo, mask_hi);
+  ABSL_ASSIGN_OR_RETURN(Type lhs_elem_ty,
+                   PrimitiveTypeToMlirType(
+                       b, params.lhs_tiled->hlo()->shape().element_type()));
+  return ApplyRowMask(b, lhs_tensor, row_mask, params.lhs_ragged_dim,
+                      lhs_elem_ty);
+}
+
+// Masks `operand` along `dim` (whose tile size is `operand`'s extent at `dim`)
+// so that elements whose absolute index reaches `dimension_size` are zeroed.
+// This is a shape-driven variant of MaskOperand that does not consult the HLO
+// shape, so it can be applied to a tile that was reshaped/squeezed after being
+// loaded (where the HLO shape no longer matches the tile rank).
+absl::StatusOr<TensorValue> MaskTileAtKBoundary(
+    ImplicitLocOpBuilder& b, TensorValue operand, Value tile_index_i32,
+    int64_t dim, int64_t dimension_size, Type elem_type) {
+  llvm::ArrayRef<int64_t> tile_shape = operand.getType().getShape();
+  const int64_t tile_size = tile_shape[dim];
+  if (dimension_size % tile_size == 0) {
+    return operand;
+  }
+  Value tile_size_value = CreateConst(b, b.getI32Type(), tile_size);
+  Value num_full_tiles = arith::DivSIOp::create(
+      b, CreateConst(b, b.getI32Type(), dimension_size), tile_size_value);
+  auto cond = arith::CmpIOp::create(b, arith::CmpIPredicate::sge,
+                                    tile_index_i32, num_full_tiles);
+  auto if_mask = mlir::scf::IfOp::create(b, operand.getType(), cond,
+                                         /*withElse=*/true);
+  {
+    mlir::OpBuilder::InsertionGuard mg(b);
+    b.setInsertionPointToStart(if_mask.thenBlock());
+    Value tile_offset =
+        arith::MulIOp::create(b, tile_index_i32, tile_size_value);
+    TensorValue range = Iota(b, tile_size);
+    TensorValue bcast_off = xtile::Splat(b, tile_offset, {tile_size});
+    Value indices = arith::AddIOp::create(b, range, bcast_off);
+    Value boundary =
+        CreateConst(b, b.getI32Type(), dimension_size, {tile_size});
+    Value mask =
+        arith::CmpIOp::create(b, arith::CmpIPredicate::slt, indices, boundary);
+    mask = xtile::BroadcastInDims(b, mlir::cast<TensorValue>(mask), tile_shape,
+                                  {dim});
+    TensorValue zero = CreateConst(b, elem_type, 0.0f, tile_shape);
+    Value masked = arith::SelectOp::create(b, mask, operand, zero);
+    mlir::scf::YieldOp::create(b, masked);
+  }
+  {
+    mlir::OpBuilder::InsertionGuard mg(b);
+    b.setInsertionPointToStart(if_mask.elseBlock());
+    mlir::scf::YieldOp::create(b, operand);
+  }
+  b.setInsertionPointAfter(if_mask);
+  return mlir::cast<TensorValue>(if_mask.getResult(0));
+}
+
+// Runs the inner partial dot lhs × rhs → acc. For batched ragged dot
+// (num_batch_dims > 0) the operands carry leading size-1 batch dims that
+// Triton's tt.dot cannot handle, so they are squeezed before the dot and the
+// result is unsqueezed back. The non-batched path is a plain single-tile dot.
+absl::StatusOr<Value> RunSqueezedTileDot(
+    ImplicitLocOpBuilder& b, const HloRaggedDotInstruction& ragged_dot,
+    const DotDimensionNumbers& dot_dims, int64_t rhs_k_dim_after_squeeze,
+    int64_t num_batch_dims, TensorValue lhs, TensorValue rhs, TensorValue acc) {
+  auto acc_shape_ref = acc.getType().getShape();
+  llvm::SmallVector<int64_t> lhs_2d_shape(
+      lhs.getType().getShape().begin() + num_batch_dims,
+      lhs.getType().getShape().end());
+  llvm::SmallVector<int64_t> rhs_2d_shape(
+      rhs.getType().getShape().begin() + num_batch_dims,
+      rhs.getType().getShape().end());
+  llvm::SmallVector<int64_t> acc_2d_shape(
+      acc_shape_ref.begin() + num_batch_dims, acc_shape_ref.end());
+
+  TensorValue dot_lhs = lhs;
+  TensorValue dot_rhs = rhs;
+  TensorValue dot_acc = acc;
+  if (num_batch_dims > 0) {
+    dot_lhs = ReshapeTo(b, lhs, lhs_2d_shape);
+    dot_rhs = ReshapeTo(b, rhs, rhs_2d_shape);
+    dot_acc = ReshapeTo(b, acc, acc_2d_shape);
+  }
+
+  // Build inner dot dimension numbers. When batch dims are squeezed out,
+  // contracting dim indices decrease by num_batch_dims (the batch dims occupied
+  // the leading positions).
+  DotDimensionNumbers inner_dot_dims;
+  inner_dot_dims.add_lhs_contracting_dimensions(
+      dot_dims.lhs_contracting_dimensions(0) - num_batch_dims);
+  inner_dot_dims.add_rhs_contracting_dimensions(rhs_k_dim_after_squeeze -
+                                                num_batch_dims);
+
+  // Partial dot: LHS [M, K] × RHS [K, N] → acc [M, N].
+  ABSL_ASSIGN_OR_RETURN(
+      Value acc_next_2d,
+      xtile::EmitSingleTileDot(b, ragged_dot, inner_dot_dims,
+                               xtile::DotOperands{dot_lhs, dot_rhs, dot_acc}));
+
+  // Unsqueeze batch dims back if they were squeezed.
+  if (num_batch_dims == 0) {
+    return acc_next_2d;
+  }
+  return ReshapeTo(b, mlir::cast<TensorValue>(acc_next_2d), acc_shape_ref);
+}
+
+// Emits the body of a single K-loop iteration for the kRaggedNonContracting
+// variant: emit LHS/RHS tiles, mask LHS rows outside group g, squeeze the RHS
+// group dim, apply the K-boundary masks, run the partial dot, and yield the
+// updated accumulator. See the extracted helpers above for each phase.
+absl::StatusOr<Value> EmitNonContractingKLoopBody(
+    EmitterContext& emitter_ctx,
+    const ge::TiledHloInstruction& tiled_ragged_dot,
+    const RaggedDotParams& rctx, const NonContractingKLoopParams& params,
+    Value k_iv, TensorValue acc_k_in) {
+  auto& b = emitter_ctx.b();
+  const ge::TiledHloInstruction* lhs_tiled = params.lhs_tiled;
+  const ge::TiledHloInstruction* rhs_tiled = params.rhs_tiled;
+  const ge::TiledHloRegion& region = tiled_ragged_dot.hlo_regions().front();
+
+  // Emit LHS and RHS tiles (K-scoped).
+  // EmitRegionInstructionWithDependencies handles the full instruction chain
+  // (parameter, bitcast, transpose, etc.) that layout normalization may insert
+  // inside the fusion for non-standard memory layouts.
+  ABSL_ASSIGN_OR_RETURN(
+      TensorValue lhs_tensor,
+      EmitRegionInstructionWithDependencies(emitter_ctx, region, lhs_tiled));
+  ABSL_ASSIGN_OR_RETURN(
+      TensorValue rhs_tensor,
+      EmitRegionInstructionWithDependencies(emitter_ctx, region, rhs_tiled));
+
+  // Mask LHS rows that are outside group g.
+  ABSL_ASSIGN_OR_RETURN(lhs_tensor, ApplyGroupRowMask(b, lhs_tensor, params));
+
+  // The RHS tile has shape [...batch..., G_tile=1, K_tile, N]. Squeeze out the
+  // G dimension (one reshape + a pure-integer K index shift) so the inner dot
+  // sees [...batch..., K_tile, N].
+  const int64_t rhs_group_dim = rctx.ragged_dims.rhs_group_dimensions(0);
+  const int64_t rhs_k_orig = rctx.dot_dims.rhs_contracting_dimensions(0);
+  llvm::SmallVector<int64_t> rhs_squeezed_shape;
+  for (auto [i, dim] : llvm::enumerate(rhs_tensor.getType().getShape())) {
+    if (static_cast<int64_t>(i) != rhs_group_dim) {
+      rhs_squeezed_shape.push_back(dim);
+    }
+  }
+  rhs_tensor = ReshapeTo(b, rhs_tensor, rhs_squeezed_shape);
+  // After squeezing G, the K contracting dim shifts left by one iff it was to
+  // the right of the (removed) group dim.
+  const int64_t rhs_k_dim_after_squeeze =
+      (rhs_k_orig > rhs_group_dim) ? rhs_k_orig - 1 : rhs_k_orig;
+
+  // Mask LHS and (squeezed) RHS at the K-boundary.
+  Value k_iv_i32 = Cast(b, k_iv, b.getI32Type());
+  ABSL_ASSIGN_OR_RETURN(
+      Type lhs_elem_ty,
+      PrimitiveTypeToMlirType(b, lhs_tiled->hlo()->shape().element_type()));
+  ABSL_ASSIGN_OR_RETURN(lhs_tensor,
+                   MaskOperand(b, *lhs_tiled, lhs_tensor, k_iv_i32,
+                               rctx.dot_dims.lhs_contracting_dimensions(0),
+                               CreateConst(b, lhs_elem_ty, 0.0f)));
+  ABSL_ASSIGN_OR_RETURN(
+      Type rhs_elem_ty,
+      PrimitiveTypeToMlirType(b, rhs_tiled->hlo()->shape().element_type()));
+  const int64_t K = rhs_tiled->hlo()->shape().dimensions(rhs_k_orig);
+  ABSL_ASSIGN_OR_RETURN(
+      rhs_tensor, MaskTileAtKBoundary(b, rhs_tensor, k_iv_i32,
+                                      rhs_k_dim_after_squeeze, K, rhs_elem_ty));
+
+  // Batch-squeeze → partial dot → batch-unsqueeze.
+  return RunSqueezedTileDot(b, *rctx.ragged_dot_instr, rctx.dot_dims,
+                            rhs_k_dim_after_squeeze, params.num_batch_dims,
+                            lhs_tensor, rhs_tensor, acc_k_in);
+}
+
+// Emits the kRaggedNonContracting variant.
+//
+// G is a kSequential outer loop.  Output (M_total, N).  For each group g in
+// [0, G): load group_size_g, check whether this M tile belongs to group g and,
+// if so, run the K-loop accumulating the partial dot.  Loop-carried state is
+// (accumulator, last_m) where last_m is the running prefix sum of group sizes.
+absl::StatusOr<TensorValue> EmitRaggedDotNonContracting(
+    EmitterContext& emitter_ctx,
+    const ge::TiledHloInstruction& tiled_ragged_dot,
+    const RaggedDotParams& rctx) {
+  auto& b = emitter_ctx.b();
+  const auto* ragged_dot_instr = rctx.ragged_dot_instr;
+  const DotDimensionNumbers& dot_dims = rctx.dot_dims;
+  const ge::TilingSpace& tiling_space = rctx.tiling_space;
+  const int64_t output_rank = rctx.output_rank;
+  const SmallVector<int64_t>& padded_tile_sizes = rctx.padded_tile_sizes;
+  const Type acc_type = rctx.acc_type;
+  TensorValue accumulator = rctx.accumulator;
+  // --- kRaggedNonContracting ---
+  // TilingSpace layout:
+  //   dim[output_rank+0]: G  kSequential (tile_size=1, outer group loop)
+  //   dim[output_rank+1]: K  kSequential (tile_size=BLOCK_K, inner K-loop)
+  // Operands in the region: [lhs, rhs, group_sizes]
+
+  const ge::TilingSpace::DimensionInfo& g_dim_info =
+      tiling_space.GetDimensionInfo(*ragged_dot_instr, output_rank);
+  const ge::TilingSpace::DimensionInfo& k_dim_info =
+      tiling_space.GetDimensionInfo(*ragged_dot_instr, output_rank + 1);
+
+  const int64_t G = g_dim_info.dimension_size;
+  CHECK(k_dim_info.tile_size.has_value())
+      << "K tile size must be set before emitting ragged dot.";
+  const int64_t K_tiles =
+      CeilOfRatio(k_dim_info.dimension_size, *k_dim_info.tile_size);
+
+  const ge::TiledHloInstruction* gs_tiled = tiled_ragged_dot.operand(2);
+
+  // M is at output dim `num_batch_dims` (after any batch dims).
+  // `lhs_ragged_dim` is also the M dim index in the LHS tile shape.
+  const int64_t num_batch_dims = dot_dims.lhs_batch_dimensions_size();
+
+  // Compute tile_m_abs BEFORE registering G so M's parallel dim is
+  // evaluated cleanly without any sequential dim interference
+  // (especially important when the M/N schedule swap is active).
+  const int64_t BLOCK_M = padded_tile_sizes[num_batch_dims];
+  const ge::DimTile& m_output_dim_pre =
+      tiled_ragged_dot.tile().dim_tiles()[num_batch_dims];
+  ABSL_ASSIGN_OR_RETURN(
+      SmallVector<Value> m_abs_pre,
+      emitter_ctx.EvaluateTilingParameters({m_output_dim_pre.offset}));
+  Value tile_m_abs = m_abs_pre[0];  // index, computed before G IV registered
+  Value tile_m_i32_cmp = Cast(b, tile_m_abs, b.getI32Type());
+  Value tile_end_i32 = arith::AddIOp::create(
+      b, tile_m_i32_cmp, CreateConst(b, b.getI32Type(), BLOCK_M));
+
+  // Outer G-loop: iter_args = (accumulator, last_m=0).
+  auto g_for_op = mlir::scf::ForOp::create(
+      b, MakeIndex(b, 0), MakeIndex(b, G), MakeIndex(b, 1),
+      mlir::ValueRange{accumulator, MakeIndex(b, 0)});
+
+  {
+    mlir::OpBuilder::InsertionGuard g_guard(b);
+    b.setInsertionPointToStart(g_for_op.getBody());
+    Value g_iv = g_for_op.getInductionVar();
+    Value last_m = g_for_op.getRegionIterArgs()[1];
+    TensorValue acc_in =
+        mlir::cast<TensorValue>(g_for_op.getRegionIterArgs()[0]);
+
+    // Register G sequential dim so EvaluateTilingParameters can resolve
+    // RTVar offsets that depend on the G loop IV.
+    ABSL_ASSIGN_OR_RETURN(ScopedSequentialDimBinding g_iv_binding,
+                     emitter_ctx.MapSymbolIdToSequentialDimValueScoped(
+                         g_dim_info.id, g_iv, Interval{0, G - 1}));
+
+    // Emit group_sizes tile (G-scoped) and extract group_size_g.
+    // Handles any HLO form: parameter, inlined constant, broadcast of scalar
+    // constant, or any other op-defined group_sizes.
+    ABSL_ASSIGN_OR_RETURN(
+        TensorValue gs_tile,
+        EmitRegionInstructionWithDependencies(
+            emitter_ctx, tiled_ragged_dot.hlo_regions().front(), gs_tiled));
+    Value group_size_g = ExtractGroupSizeScalar(b, gs_tile);
+
+    // Determine whether this program's M tile OVERLAPS with group g.
+    // tile_m_abs was computed before the G-loop to avoid interference
+    // from the G sequential dim registration (see comment above).
+    Value last_m_plus_gs = mlir::arith::AddIOp::create(b, last_m, group_size_g);
+    // Signed overlap check in i32 (works for all schedule permutations):
+    //   tile_m_abs_i32 < last_m_plus_gs_i32  AND
+    //   tile_end_i32   > last_m_i32
+    Value last_m_i32_cmp = Cast(b, last_m, b.getI32Type());
+    Value last_m_plus_i32_cmp = Cast(b, last_m_plus_gs, b.getI32Type());
+    Value overlaps = mlir::arith::AndIOp::create(
+        b,
+        mlir::arith::CmpIOp::create(b, mlir::arith::CmpIPredicate::slt,
+                                    tile_m_i32_cmp, last_m_plus_i32_cmp),
+        mlir::arith::CmpIOp::create(b, mlir::arith::CmpIPredicate::sgt,
+                                    tile_end_i32, last_m_i32_cmp));
+
+    // Conditional: run K-loop if this M tile overlaps group g at all.
+    auto if_op = mlir::scf::IfOp::create(b, acc_in.getType(), overlaps,
+                                         /*hasElse=*/true);
+
+    {  // Then block: K-loop accumulation.
+      mlir::OpBuilder::InsertionGuard if_guard(b);
+      b.setInsertionPointToStart(if_op.thenBlock());
+
+      auto k_for_op =
+          mlir::scf::ForOp::create(b, MakeIndex(b, 0), MakeIndex(b, K_tiles),
+                                   MakeIndex(b, 1), mlir::ValueRange{acc_in});
+
+      {  // K-loop body.
+        mlir::OpBuilder::InsertionGuard k_guard(b);
+        b.setInsertionPointToStart(k_for_op.getBody());
+        Value k_iv = k_for_op.getInductionVar();
+        TensorValue acc_k_in =
+            mlir::cast<TensorValue>(k_for_op.getRegionIterArgs()[0]);
+
+        // Register K sequential dim IV.
+        ABSL_ASSIGN_OR_RETURN(ScopedSequentialDimBinding k_iv_binding,
+                         emitter_ctx.MapSymbolIdToSequentialDimValueScoped(
+                             k_dim_info.id, k_iv, Interval{0, K_tiles - 1}));
+
+        NonContractingKLoopParams k_params{tiled_ragged_dot.operand(0),
+                                           tiled_ragged_dot.operand(1),
+                                           rctx.lhs_ragged_dim,
+                                           num_batch_dims,
+                                           tile_m_i32_cmp,
+                                           last_m_i32_cmp,
+                                           last_m_plus_i32_cmp};
+        ABSL_ASSIGN_OR_RETURN(Value acc_next, EmitNonContractingKLoopBody(
+                                             emitter_ctx, tiled_ragged_dot,
+                                             rctx, k_params, k_iv, acc_k_in));
+        mlir::scf::YieldOp::create(b, acc_next);
+      }
+      b.setInsertionPointAfter(k_for_op);
+      mlir::scf::YieldOp::create(b, k_for_op.getResult(0));
+    }
+
+    {  // Else block: accumulator unchanged.
+      mlir::OpBuilder::InsertionGuard else_guard(b);
+      b.setInsertionPointToStart(if_op.elseBlock());
+      mlir::scf::YieldOp::create(b, acc_in);
+    }
+    b.setInsertionPointAfter(if_op);
+
+    // Update last_m for next G iteration.
+    Value last_m_next = mlir::arith::AddIOp::create(b, last_m, group_size_g);
+    mlir::scf::YieldOp::create(
+        b, mlir::ValueRange{if_op.getResult(0), last_m_next});
+  }
+  b.setInsertionPointAfter(g_for_op);
+
+  TensorValue result = mlir::cast<TensorValue>(g_for_op.getResult(0));
+
+  // Cast accumulator to the declared output element type if needed.
+  ABSL_ASSIGN_OR_RETURN(Type out_type,
+                   PrimitiveTypeToMlirType(
+                       b, tiled_ragged_dot.hlo()->shape().element_type()));
+  if (out_type != acc_type) {
+    result = mlir::cast<TensorValue>(Cast(b, result, out_type));
+  }
+  return result;
+}
+
+// Emits the kRaggedContracting variant.
+//
+// Grid = G × K_tiles × N_tiles.  Each program owns exactly one group, so it
+// reconstructs start_m = sum(group_sizes[0..g-1]) via an explicit O(G)
+// prefix-sum before running the M accumulation loop.
+absl::StatusOr<TensorValue> EmitRaggedDotContracting(
+    EmitterContext& emitter_ctx,
+    const ge::TiledHloInstruction& tiled_ragged_dot,
+    const RaggedDotParams& rctx) {
+  auto& b = emitter_ctx.b();
+  const auto* ragged_dot_instr = rctx.ragged_dot_instr;
+  const DotDimensionNumbers& dot_dims = rctx.dot_dims;
+  const ge::TilingSpace& tiling_space = rctx.tiling_space;
+  const int64_t output_rank = rctx.output_rank;
+  const SmallVector<int64_t>& padded_tile_sizes = rctx.padded_tile_sizes;
+  const Type acc_type = rctx.acc_type;
+  // --- kRaggedContracting ---
+  // Grid = G × K_tiles × N_tiles. Prefix-sum per program to find start_m.
+  //
+  // How this differs from kRaggedNonContracting in its use of the RTVars
+  // registered by TilingSpace::ProcessRaggedDot:
+  //   * Both variants need the same conceptual per-group M offset
+  //     start_m[g] = sum(group_sizes[0..g-1]).
+  //   * kRaggedNonContracting registers G as a *sequential* outer loop, so a
+  //     single program sweeps every group in order and the offset is simply
+  //     carried as a loop-carried iter_arg (`last_m`) — no start_m RTVar is
+  //     needed and only the array RTVar group_sizes[g] (operand_id=2) is
+  //     registered.
+  //   * kRaggedContracting registers G as a *parallel* grid dim, so each
+  //     program owns exactly one group and cannot accumulate the offset
+  //     incrementally.  It must instead reconstruct start_m from scratch via
+  //     the explicit O(G) prefix-sum loop below.  That is why
+  //     ProcessRaggedDot registers a second RTVar (sentinel operand_id=-1)
+  //     alongside the group_size RTVar (operand_id=2), both referencing the
+  //     same group_sizes hlo.
+  // The RTVars are looked up here only to assert the tiling space knows about
+  // these runtime-dependent quantities (their bounds are consumed by tile
+  // propagation / tiled-hlo bounds recording); the concrete arithmetic is
+  // emitted directly below.
+
+  // M sequential dim: output_rank + 0.
+  const ge::TilingSpace::DimensionInfo& m_dim_info =
+      tiling_space.GetDimensionInfo(*ragged_dot_instr, output_rank);
+  CHECK(m_dim_info.tile_size.has_value())
+      << "M tile size not set for kRaggedContracting.";
+  const int64_t BLOCK_M = *m_dim_info.tile_size;
+  const int64_t M_total = m_dim_info.dimension_size;
+  auto gs_rtvar_or = tiling_space.GetRTVarInfo(*ragged_dot_instr, 2);
+  auto sm_rtvar_or = tiling_space.GetRTVarInfo(*ragged_dot_instr, -1);
+  CHECK(gs_rtvar_or.has_value()) << "Missing group_size RTVar.";
+  CHECK(sm_rtvar_or.has_value()) << "Missing start_m RTVar.";
+
+  const int64_t num_batch_dims = dot_dims.lhs_batch_dimensions_size();
+
+  // g_abs: the group index for this tile (kParallel dim 0 of output [G,...]).
+  // tile_propagation layout: [G, batch..., K, N]  → G at output dim 0.
+  const ge::DimTile& g_out_dim = tiled_ragged_dot.tile().dim_tiles()[0];
+  ABSL_ASSIGN_OR_RETURN(SmallVector<Value> g_abs_vec,
+                   emitter_ctx.EvaluateTilingParameters({g_out_dim.offset}));
+  Value g_abs = g_abs_vec[0];  // index type, = g * G_tile_size(=1) = g
+
+  // Resolve group_sizes to its fusion argument buffer (needed for the
+  // O(G) prefix sum: we must load gs[0..g-1] at arbitrary offsets).
+  const ge::TiledHloInstruction* gs_tiled = tiled_ragged_dot.operand(2);
+  const HloFusionInstruction& fusion = emitter_ctx.fusion();
+  mlir::Value gs_buf =
+      ResolveBufferArgument(emitter_ctx, fusion, gs_tiled->hlo());
+  ABSL_ASSIGN_OR_RETURN(
+      Type gs_elem_ty,
+      PrimitiveTypeToMlirType(b, gs_tiled->hlo()->shape().element_type()));
+
+  // Helper: load a single group_sizes element from the buffer at index `idx`.
+  // For batched group_sizes [B, G], this is a follow-up; for now we assert
+  // no batch dims to keep the prefix sum simple.
+  // TODO(ragged_dot): support batched kRaggedContracting.
+  CHECK_EQ(num_batch_dims, 0)
+      << "Batched kRaggedContracting not yet supported.";
+
+  // Compute start_m = sum(group_sizes[0..g-1]) via a scalar prefix-sum
+  // loop.  This is O(G) per program; acceptable for small G (~8–128 experts).
+  // A vectorized approach (load all G elements, masked tl.sum) can be added
+  // as a follow-up optimization.
+  auto load_gs_scalar = [&](Value idx) -> Value {
+    SmallVector<int64_t> sizes_1{1};
+    auto gs_1_type = mlir::RankedTensorType::get(sizes_1, gs_elem_ty);
+    TensorValue g1 = mlir::cast<TensorValue>(
+        xtile::ExtractTileOp::create(b, gs_1_type, gs_buf,
+                                     SmallVector<Value>{idx}, sizes_1,
+                                     SmallVector<int64_t>{1})
+            .getResult());
+    TensorValue g0 = ReshapeTo(b, g1, /*new_shape=*/{});
+    Value gv = mlir::tensor::ExtractOp::create(b, g0);
+    return mlir::arith::IndexCastOp::create(b, b.getIndexType(), gv);
+  };
+
+  Value start_m = MakeIndex(b, 0);
+  {
+    auto pfx = mlir::scf::ForOp::create(b, MakeIndex(b, 0), g_abs,
+                                        MakeIndex(b, 1), start_m);
+    mlir::OpBuilder::InsertionGuard guard(b);
+    b.setInsertionPointToStart(pfx.getBody());
+    Value new_sum = arith::AddIOp::create(
+        b, pfx.getRegionIterArgs()[0], load_gs_scalar(pfx.getInductionVar()));
+    mlir::scf::YieldOp::create(b, new_sum);
+    b.setInsertionPointAfter(pfx);
+    start_m = pfx.getResult(0);
+  }
+
+  // Load group_size_g = group_sizes[g_abs] using the tiled group_sizes
+  // operand (which already has the correct kParallel G offset from tile_id).
+  // EmitRegionInstructionWithDependencies resolves any HLO form (param,
+  // constant, broadcast, etc.).
+  ABSL_ASSIGN_OR_RETURN(
+      TensorValue gs_tile,
+      EmitRegionInstructionWithDependencies(
+          emitter_ctx, tiled_ragged_dot.hlo_regions().front(), gs_tiled));
+  Value group_size_g = ExtractGroupSizeScalar(b, gs_tile);
+
+  // M loop count = cdiv(group_size_g, BLOCK_M).
+  Value bm_val = MakeIndex(b, BLOCK_M);
+  Value bm_m1 = MakeIndex(b, BLOCK_M - 1);
+  Value m_loop_count = arith::DivSIOp::create(
+      b, arith::AddIOp::create(b, group_size_g, bm_m1), bm_val);
+
+  // Resolve K and N parallel dim offsets from the tile_id.
+  // Output dim ordering: [G=0, batch..., K, N]
+  // out_k_start = 1 + num_batch_dims, out_n_start = out_k_start + 1
+  const int64_t out_k_start = 1 + num_batch_dims;
+  const ge::DimTile& k_out_dim =
+      tiled_ragged_dot.tile().dim_tiles()[out_k_start];
+  ABSL_ASSIGN_OR_RETURN(SmallVector<Value> k_abs_vec,
+                   emitter_ctx.EvaluateTilingParameters({k_out_dim.offset}));
+  Value k_abs = k_abs_vec[0];
+
+  const int64_t out_n_start = out_k_start + 1;
+  const ge::DimTile& n_out_dim =
+      tiled_ragged_dot.tile().dim_tiles()[out_n_start];
+  ABSL_ASSIGN_OR_RETURN(SmallVector<Value> n_abs_vec,
+                   emitter_ctx.EvaluateTilingParameters({n_out_dim.offset}));
+  Value n_abs = n_abs_vec[0];
+
+  // LHS and RHS buffer arguments for direct ExtractTileOp loads.
+  // For FP8 inputs, the operand of the ragged-dot inside the fusion may be a
+  // convert(parameter) instruction (e.g. f8→f32 widening).
+  // ResolveBufferArgument traverses through any kConvert to reach the
+  // underlying buffer parameter; its element type (before any convert,
+  // e.g. f8e4m3fnuz for FP8 inputs) is carried by the returned buffer Value.
+  // EmitSingleTileDot will cast to the required accumulation type (f32) if
+  // needed.
+  mlir::Value lhs_buf = ResolveBufferArgument(
+      emitter_ctx, fusion, tiled_ragged_dot.operand(0)->hlo());
+  Type lhs_elem_mlir_ty = getElementTypeOrSelf(lhs_buf.getType());
+
+  mlir::Value rhs_buf = ResolveBufferArgument(
+      emitter_ctx, fusion, tiled_ragged_dot.operand(1)->hlo());
+  Type rhs_elem_mlir_ty = getElementTypeOrSelf(rhs_buf.getType());
+
+  // Determine BLOCK_K and BLOCK_N from the output tile padded sizes.
+  const int64_t BLOCK_K = padded_tile_sizes[out_k_start];
+  const int64_t BLOCK_N = padded_tile_sizes[out_n_start];
+
+  // Output tile shape for this program: padded_tile_sizes = [1, K, N] (or
+  // [1, B..., K, N] for batched).  The leading 1 is the G tile size.
+  // Inner accumulator: 2-D [K_tile, N_tile] (drop the G=1 leading dim).
+  // After the M loop we reshape back to padded_tile_sizes.
+  llvm::SmallVector<int64_t> acc_2d_shape(padded_tile_sizes.begin() + 1,
+                                          padded_tile_sizes.end());
+  TensorValue accumulator_2d = CreateConst(b, acc_type, 0.0f, acc_2d_shape);
+
+  // lhs_contracting_dim_in_tile: position of M in the LHS tile
+  // rhs_contracting_dim_in_tile: position of M in the RHS tile
+  const int64_t lhs_m_tile_dim =
+      dot_dims.lhs_contracting_dimensions(0) - num_batch_dims;
+  const int64_t rhs_m_tile_dim =
+      dot_dims.rhs_contracting_dimensions(0) - num_batch_dims;
+
+  // M sequential accumulation loop.
+  auto m_for = mlir::scf::ForOp::create(b, MakeIndex(b, 0), m_loop_count,
+                                        MakeIndex(b, 1), accumulator_2d);
+  {
+    mlir::OpBuilder::InsertionGuard m_guard(b);
+    b.setInsertionPointToStart(m_for.getBody());
+    Value m_iv = m_for.getInductionVar();
+    TensorValue acc_in = mlir::cast<TensorValue>(m_for.getRegionIterArgs()[0]);
+
+    // Register the M sequential dim IV so that the relative tile offset
+    // (m_iv * BLOCK_M) can be evaluated by EvaluateTilingParameters if
+    // needed elsewhere. The actual buffer address uses start_m + m_iv*BLOCK_M
+    // computed below.
+    //
+    // Use the per-group maximum loop count as the interval upper bound.
+    // M_total/BLOCK_M is G times too large: Triton uses this interval for
+    // speculative prefetch analysis and generates loads at m_abs = start_m +
+    // (M_total/BLOCK_M)*BLOCK_M which is far beyond the LHS/RHS buffer →
+    // GPU memory access fault when those speculative addresses hit unmapped
+    // pages. The correct bound is cdiv(cdiv(M_total, G), BLOCK_M) - 1.
+    {
+      const int64_t G_nc_iv =
+          ragged_dot_instr->operand(2)->shape().dimensions(0);
+      const int64_t q_max_nc_iv =
+          (G_nc_iv > 0) ? (M_total + G_nc_iv - 1) / G_nc_iv : M_total;
+      const int64_t m_loop_max_nc_iv = (q_max_nc_iv + BLOCK_M - 1) / BLOCK_M;
+      ABSL_ASSIGN_OR_RETURN(
+          ScopedSequentialDimBinding m_iv_binding,
+          emitter_ctx.MapSymbolIdToSequentialDimValueScoped(
+              m_dim_info.id, m_iv,
+              Interval{0, m_loop_max_nc_iv > 0 ? m_loop_max_nc_iv - 1 : 0}));
+    }
+
+    // Absolute M offset = start_m + m_iv * BLOCK_M.
+    // Clamp to prevent Triton speculative loads beyond the group's slice.
+    // The clamp: m_abs_upper = start_m + (m_loop_count - 1) * BLOCK_M
+    // = start of the LAST VALID tile. This is correct for all group sizes:
+    //   - group_size_g=10, BLOCK_M=32: m_loop_count=1, upper=start_m ✓
+    //   - group_size_g=56, BLOCK_M=32: m_loop_count=2, upper=start_m+32 ✓
+    //   - group_size_g=128, BLOCK_M=64: m_loop_count=2, upper=start_m+64 ✓
+    // The clamp only activates for OOB-speculation (m_iv ≥ m_loop_count).
+    // Those iterations are masked to zero by M-boundary masking below.
+    Value m_abs_raw = arith::AddIOp::create(
+        b, start_m, arith::MulIOp::create(b, m_iv, MakeIndex(b, BLOCK_M)));
+    {
+      // m_loop_count was computed above: cdiv(group_size_g, BLOCK_M).
+      Value last_valid_idx =
+          arith::SubIOp::create(b, m_loop_count, MakeIndex(b, 1));
+      Value m_abs_upper = arith::AddIOp::create(
+          b, start_m,
+          arith::MulIOp::create(b, last_valid_idx, MakeIndex(b, BLOCK_M)));
+      m_abs_raw = arith::MinSIOp::create(b, m_abs_raw, m_abs_upper);
+    }
+    Value m_abs = m_abs_raw;
+
+    // Load LHS tile directly with absolute M offset (no emit_operand).
+    // The tile propagation uses relative M offset (no sm_sym) since
+    // start_m is a synthetic prefix-sum RTVar with no backing tiled HLO.
+    // We add start_m here manually to get the correct absolute buffer offset.
+    //
+    // LHS offset array: for each dim, use m_abs at the contracting dim and
+    // k_abs at the non-contracting dim.
+    // For LHS [M, K] (contracting_dim=0): offsets=[m_abs, k_abs]
+    // For LHS [K, M] (contracting_dim=1): offsets=[k_abs, m_abs]
+    TensorValue lhs_tensor;
+    {
+      SmallVector<Value> lhs_offsets(2);
+      SmallVector<int64_t> lhs_sizes(2);
+      lhs_offsets[lhs_m_tile_dim] = m_abs;
+      lhs_sizes[lhs_m_tile_dim] = BLOCK_M;
+      lhs_offsets[1 - lhs_m_tile_dim] = k_abs;
+      lhs_sizes[1 - lhs_m_tile_dim] = BLOCK_K;
+      auto lhs_tile_type =
+          mlir::RankedTensorType::get(lhs_sizes, lhs_elem_mlir_ty);
+      lhs_tensor = mlir::cast<TensorValue>(
+          xtile::ExtractTileOp::create(b, lhs_tile_type, lhs_buf, lhs_offsets,
+                                       lhs_sizes, SmallVector<int64_t>{1, 1})
+              .getResult());
+    }
+
+    // Load RHS tile directly with absolute M offset.
+    // RHS offset array: for each dim, use m_abs at the contracting dim and
+    // n_abs at the non-contracting dim.
+    TensorValue rhs_tensor;
+    {
+      SmallVector<Value> rhs_offsets(2);
+      SmallVector<int64_t> rhs_sizes(2);
+      rhs_offsets[rhs_m_tile_dim] = m_abs;
+      rhs_sizes[rhs_m_tile_dim] = BLOCK_M;
+      rhs_offsets[1 - rhs_m_tile_dim] = n_abs;
+      rhs_sizes[1 - rhs_m_tile_dim] = BLOCK_N;
+      auto rhs_tile_type =
+          mlir::RankedTensorType::get(rhs_sizes, rhs_elem_mlir_ty);
+      rhs_tensor = mlir::cast<TensorValue>(
+          xtile::ExtractTileOp::create(b, rhs_tile_type, rhs_buf, rhs_offsets,
+                                       rhs_sizes, SmallVector<int64_t>{1, 1})
+              .getResult());
+    }
+
+    // M-boundary masking: zero elements where m_iv*BLOCK_M + idx >=
+    // group_size_g. Applied to both LHS and RHS since they share the same M
+    // contracting dim.
+    {
+      Value m_iv_i32 = Cast(b, m_iv, b.getI32Type());
+      Value gs_g_i32 = Cast(b, group_size_g, b.getI32Type());
+      Value bm_i32 = CreateConst(b, b.getI32Type(), BLOCK_M);
+      Value m_base = arith::MulIOp::create(b, m_iv_i32, bm_i32);
+      TensorValue row_iota = Iota(b, static_cast<int32_t>(BLOCK_M));
+      Value abs_row = arith::AddIOp::create(b, row_iota,
+                                            xtile::Splat(b, m_base, {BLOCK_M}));
+      Value m_mask_1d =
+          arith::CmpIOp::create(b, arith::CmpIPredicate::slt, abs_row,
+                                xtile::Splat(b, gs_g_i32, {BLOCK_M}));
+
+      // Broadcast along the M (contracting) dim of the LHS tile.
+      {
+        auto lhs_shape = lhs_tensor.getType().getShape();
+        Value lhs_mask = xtile::BroadcastInDims(
+            b, mlir::cast<TensorValue>(m_mask_1d), lhs_shape, {lhs_m_tile_dim});
+        // lhs_elem_mlir_ty was computed before the M loop.
+        lhs_tensor = mlir::cast<TensorValue>(
+            arith::SelectOp::create(
+                b, lhs_mask, lhs_tensor,
+                CreateConst(b, lhs_elem_mlir_ty, 0.0f, lhs_shape))
+                .getResult());
+      }
+      // Same for RHS.
+      {
+        auto rhs_shape = rhs_tensor.getType().getShape();
+        Value rhs_mask = xtile::BroadcastInDims(
+            b, mlir::cast<TensorValue>(m_mask_1d), rhs_shape, {rhs_m_tile_dim});
+        rhs_tensor = mlir::cast<TensorValue>(
+            arith::SelectOp::create(
+                b, rhs_mask, rhs_tensor,
+                CreateConst(b, rhs_elem_mlir_ty, 0.0f, rhs_shape))
+                .getResult());
+      }
+    }
+
+    // Inner dot: LHS[..., BLOCK_M, ...] contracted at lhs_m_tile_dim
+    //            with RHS[..., BLOCK_M, ...] → acc[BLOCK_K, BLOCK_N].
+    // The dot accumulates into the 2-D acc_in [K_tile, N_tile].
+    DotDimensionNumbers inner_dot_dims;
+    inner_dot_dims.add_lhs_contracting_dimensions(lhs_m_tile_dim);
+    inner_dot_dims.add_rhs_contracting_dimensions(rhs_m_tile_dim);
+
+    ABSL_ASSIGN_OR_RETURN(Value acc_next,
+                     xtile::EmitSingleTileDot(
+                         b, *ragged_dot_instr, inner_dot_dims,
+                         xtile::DotOperands{lhs_tensor, rhs_tensor, acc_in}));
+    mlir::scf::YieldOp::create(b, acc_next);
+  }
+  b.setInsertionPointAfter(m_for);
+
+  // Reshape 2-D result [K_tile, N_tile] → padded_tile_sizes [1, K, N]
+  // to match what EmitGeneric expects for the InsertTileOp.
+  Value result_2d = m_for.getResult(0);
+  ABSL_ASSIGN_OR_RETURN(Type out_type,
+                   PrimitiveTypeToMlirType(
+                       b, tiled_ragged_dot.hlo()->shape().element_type()));
+  if (out_type != acc_type) {
+    result_2d = Cast(b, result_2d, out_type);
+  }
+  return ReshapeTo(b, mlir::cast<TensorValue>(result_2d), padded_tile_sizes);
+}
+
+// Emits a kRaggedDot instruction by dispatching to the appropriate mode.
+absl::StatusOr<TensorValue> EmitRaggedDot(
+    EmitterContext& emitter_ctx,
+    const ge::TiledHloInstruction& tiled_ragged_dot) {
+  ABSL_ASSIGN_OR_RETURN(RaggedDotParams rctx,
+                   BuildRaggedDotContext(emitter_ctx, tiled_ragged_dot));
+
+  const int64_t lhs_ragged_dim = rctx.lhs_ragged_dim;
+  const bool is_contracting =
+      absl::c_count(rctx.dot_dims.lhs_contracting_dimensions(),
+                    lhs_ragged_dim) > 0;
+  const bool is_batch =
+      absl::c_count(rctx.dot_dims.lhs_batch_dimensions(), lhs_ragged_dim) > 0;
 
   if (is_batch) {
     return absl::UnimplementedError("kRaggedBatch: not yet implemented");
   }
-
   if (!is_contracting) {
-    // --- kRaggedNonContracting ---
-    // TilingSpace layout:
-    //   dim[output_rank+0]: G  kSequential (tile_size=1, outer group loop)
-    //   dim[output_rank+1]: K  kSequential (tile_size=BLOCK_K, inner K-loop)
-    // Operands in the region: [lhs, rhs, group_sizes]
-
-    const ge::TilingSpace::DimensionInfo& g_dim_info =
-        tiling_space.GetDimensionInfo(*ragged_dot_instr, output_rank);
-    const ge::TilingSpace::DimensionInfo& k_dim_info =
-        tiling_space.GetDimensionInfo(*ragged_dot_instr, output_rank + 1);
-
-    const int64_t G = g_dim_info.dimension_size;
-    CHECK(k_dim_info.tile_size.has_value())
-        << "K tile size must be set before emitting ragged dot.";
-    const int64_t K_tiles =
-        CeilOfRatio(k_dim_info.dimension_size, *k_dim_info.tile_size);
-
-    const ge::TiledHloInstruction* lhs_tiled = tiled_ragged_dot.operand(0);
-    const ge::TiledHloInstruction* rhs_tiled = tiled_ragged_dot.operand(1);
-    const ge::TiledHloInstruction* gs_tiled = tiled_ragged_dot.operand(2);
-
-    // M is at output dim `num_batch_dims` (after any batch dims).
-    // `lhs_ragged_dim` is also the M dim index in the LHS tile shape.
-    const int64_t num_batch_dims = dot_dims.lhs_batch_dimensions_size();
-
-    // Compute tile_m_abs BEFORE registering G so M's parallel dim is
-    // evaluated cleanly without any sequential dim interference
-    // (especially important when the M/N schedule swap is active).
-    const int64_t BLOCK_M = padded_tile_sizes[num_batch_dims];
-    const ge::DimTile& m_output_dim_pre =
-        tiled_ragged_dot.tile().dim_tiles()[num_batch_dims];
-    ABSL_ASSIGN_OR_RETURN(
-        SmallVector<Value> m_abs_pre,
-        emitter_ctx.EvaluateTilingParameters({m_output_dim_pre.offset}));
-    Value tile_m_abs = m_abs_pre[0];  // index, computed before G IV registered
-    Value tile_m_i32_cmp = Cast(b, tile_m_abs, b.getI32Type());
-    Value tile_end_i32 = arith::AddIOp::create(
-        b, tile_m_i32_cmp, CreateConst(b, b.getI32Type(), BLOCK_M));
-
-    // Outer G-loop: iter_args = (accumulator, last_m=0).
-    auto g_for_op = mlir::scf::ForOp::create(
-        b, MakeIndex(b, 0), MakeIndex(b, G), MakeIndex(b, 1),
-        mlir::ValueRange{accumulator, MakeIndex(b, 0)});
-
-    {
-      mlir::OpBuilder::InsertionGuard g_guard(b);
-      b.setInsertionPointToStart(g_for_op.getBody());
-      Value g_iv = g_for_op.getInductionVar();
-      Value last_m = g_for_op.getRegionIterArgs()[1];
-      TensorValue acc_in =
-          mlir::cast<TensorValue>(g_for_op.getRegionIterArgs()[0]);
-
-      // Register G sequential dim so EvaluateTilingParameters can resolve
-      // RTVar offsets that depend on the G loop IV.
-      ABSL_ASSIGN_OR_RETURN(ScopedSequentialDimBinding g_iv_binding,
-                       emitter_ctx.MapSymbolIdToSequentialDimValueScoped(
-                           g_dim_info.id, g_iv, Interval{0, G - 1}));
-
-      // Emit group_sizes tile (G-scoped). `emit_operand` handles any HLO form:
-      // parameter, inlined constant, broadcast of scalar constant, or any
-      // other op-defined group_sizes.
-      ABSL_ASSIGN_OR_RETURN(TensorValue gs_tile, emit_operand(gs_tiled));
-
-      // Extract group_size_g from gs_tile.
-      // gs_tile is tensor<1xi{32|64}>; XTile/Triton's tensor.extract requires
-      // a rank-0 tensor — reshape to tensor<elem_type> (rank-0) first.
-      auto gs_elem_type = gs_tile.getType().getElementType();
-      TensorValue gs_tile_scalar = mlir::cast<TensorValue>(
-          stablehlo::ReshapeOp::create(
-              b, mlir::RankedTensorType::get({}, gs_elem_type), gs_tile)
-              .getResult());
-      Value gs_raw = mlir::tensor::ExtractOp::create(b, gs_tile_scalar);
-      Value group_size_g =
-          mlir::arith::IndexCastOp::create(b, b.getIndexType(), gs_raw);
-
-      // Determine whether this program's M tile OVERLAPS with group g.
-      // tile_m_abs was computed before the G-loop to avoid interference
-      // from the G sequential dim registration (see comment above).
-      Value last_m_plus_gs =
-          mlir::arith::AddIOp::create(b, last_m, group_size_g);
-      // Signed overlap check in i32 (works for all schedule permutations):
-      //   tile_m_abs_i32 < last_m_plus_gs_i32  AND
-      //   tile_end_i32   > last_m_i32
-      Value last_m_i32_cmp = Cast(b, last_m, b.getI32Type());
-      Value last_m_plus_i32_cmp = Cast(b, last_m_plus_gs, b.getI32Type());
-      Value overlaps = mlir::arith::AndIOp::create(
-          b,
-          mlir::arith::CmpIOp::create(b, mlir::arith::CmpIPredicate::slt,
-                                      tile_m_i32_cmp, last_m_plus_i32_cmp),
-          mlir::arith::CmpIOp::create(b, mlir::arith::CmpIPredicate::sgt,
-                                      tile_end_i32, last_m_i32_cmp));
-
-      // Conditional: run K-loop if this M tile overlaps group g at all.
-      auto if_op = mlir::scf::IfOp::create(b, acc_in.getType(), overlaps,
-                                           /*hasElse=*/true);
-
-      {  // Then block: K-loop accumulation.
-        mlir::OpBuilder::InsertionGuard if_guard(b);
-        b.setInsertionPointToStart(if_op.thenBlock());
-
-        auto k_for_op =
-            mlir::scf::ForOp::create(b, MakeIndex(b, 0), MakeIndex(b, K_tiles),
-                                     MakeIndex(b, 1), mlir::ValueRange{acc_in});
-
-        {  // K-loop body.
-          mlir::OpBuilder::InsertionGuard k_guard(b);
-          b.setInsertionPointToStart(k_for_op.getBody());
-          Value k_iv = k_for_op.getInductionVar();
-          TensorValue acc_k_in =
-              mlir::cast<TensorValue>(k_for_op.getRegionIterArgs()[0]);
-
-          // Register K sequential dim IV.
-          ABSL_ASSIGN_OR_RETURN(ScopedSequentialDimBinding k_iv_binding,
-                           emitter_ctx.MapSymbolIdToSequentialDimValueScoped(
-                               k_dim_info.id, k_iv, Interval{0, K_tiles - 1}));
-
-          // Emit LHS and RHS tiles (K-scoped).
-          // emit_operand handles the full instruction chain (parameter,
-          // bitcast, transpose, etc.) that layout normalization may insert
-          // inside the fusion for non-standard memory layouts.
-          ABSL_ASSIGN_OR_RETURN(TensorValue lhs_tensor, emit_operand(lhs_tiled));
-          ABSL_ASSIGN_OR_RETURN(TensorValue rhs_tensor, emit_operand(rhs_tiled));
-
-          // Mask LHS rows that are outside group g.
-          // Row r (0..BLOCK_M-1) is valid iff:
-          //   tile_m_abs + r >= last_m  AND  tile_m_abs + r < last_m_plus_gs
-          // Build an i32 mask over [BLOCK_M] rows, then broadcast to [M, K].
-          {
-            auto lhs_shape = lhs_tensor.getType().getShape();
-            // lhs_ragged_dim is the M dimension in the LHS tile shape.
-            // For no-batch: lhs_ragged_dim=0 → lhs_shape[0]=BLOCK_M.
-            // For batched:  lhs_ragged_dim=1 → lhs_shape[1]=BLOCK_M.
-            int64_t lhs_m_tile_size = lhs_shape[lhs_ragged_dim];  // == BLOCK_M
-            int64_t lhs_m_dim = lhs_m_tile_size;
-            // iota over M rows: 0..BLOCK_M-1
-            TensorValue row_iota = Iota(b, static_cast<int32_t>(lhs_m_dim));
-            // abs_m[r] = tile_m_abs_i32 + r
-            // Use pre-computed tile_m_i32_cmp (avoids re-evaluating M offset
-            // inside G-loop where sequential dims are registered).
-            TensorValue tile_m_splat =
-                xtile::Splat(b, tile_m_i32_cmp, {lhs_m_dim});
-            Value abs_m = arith::AddIOp::create(b, row_iota, tile_m_splat);
-            // last_m and last_m_plus_gs in i32 (already computed above).
-            TensorValue lm_splat = xtile::Splat(b, last_m_i32_cmp, {lhs_m_dim});
-            TensorValue lmp_splat =
-                xtile::Splat(b, last_m_plus_i32_cmp, {lhs_m_dim});
-            Value mask_lo = arith::CmpIOp::create(b, arith::CmpIPredicate::sge,
-                                                  abs_m, lm_splat);
-            Value mask_hi = arith::CmpIOp::create(b, arith::CmpIPredicate::slt,
-                                                  abs_m, lmp_splat);
-            Value row_mask = arith::AndIOp::create(b, mask_lo, mask_hi);
-            // Broadcast [BLOCK_M] mask along the M dim of the LHS tile.
-            row_mask =
-                xtile::BroadcastInDims(b, mlir::cast<TensorValue>(row_mask),
-                                       lhs_shape, {lhs_ragged_dim});
-            ABSL_ASSIGN_OR_RETURN(Type lhs_elem_ty,
-                             PrimitiveTypeToMlirType(
-                                 b, lhs_tiled->hlo()->shape().element_type()));
-            TensorValue lhs_zero = CreateConst(b, lhs_elem_ty, 0.0f, lhs_shape);
-            lhs_tensor = mlir::cast<TensorValue>(
-                arith::SelectOp::create(b, row_mask, lhs_tensor, lhs_zero)
-                    .getResult());
-          }
-
-          // The RHS tile has shape [...batch..., G_tile=1, K_tile, N].
-          // Squeeze out the G dimension (at rhs_group_dim) so the inner
-          // dot sees [...batch..., K_tile, N].
-          const int64_t rhs_group_dim_orig =
-              ragged_dims.rhs_group_dimensions(0);
-          auto rhs_full_shape = rhs_tensor.getType().getShape();
-          llvm::SmallVector<int64_t> rhs_squeezed_shape;
-          for (int64_t i = 0; i < static_cast<int64_t>(rhs_full_shape.size());
-               ++i) {
-            if (i != rhs_group_dim_orig) {
-              rhs_squeezed_shape.push_back(rhs_full_shape[i]);
-            }
-          }
-          rhs_tensor = mlir::cast<TensorValue>(
-              stablehlo::ReshapeOp::create(
-                  b,
-                  mlir::RankedTensorType::get(
-                      rhs_squeezed_shape,
-                      rhs_tensor.getType().getElementType()),
-                  rhs_tensor)
-                  .getResult());
-
-          // After squeezing G, the K contracting dim shifts:
-          //   If K_orig > rhs_group_dim → k_idx decreases by 1
-          //   Otherwise stays the same.
-          const int64_t rhs_k_orig = dot_dims.rhs_contracting_dimensions(0);
-          const int64_t rhs_k_dim_after_squeeze =
-              (rhs_k_orig > rhs_group_dim_orig) ? rhs_k_orig - 1 : rhs_k_orig;
-
-          // Mask LHS at K-boundary (same as EmitDot).
-          Value k_iv_i32 = Cast(b, k_iv, b.getI32Type());
-          {
-            ABSL_ASSIGN_OR_RETURN(Type lhs_k_elem_ty,
-                             PrimitiveTypeToMlirType(
-                                 b, lhs_tiled->hlo()->shape().element_type()));
-            Value lhs_k_zero = CreateConst(b, lhs_k_elem_ty, 0.0f);
-            ABSL_ASSIGN_OR_RETURN(lhs_tensor,
-                             MaskOperand(b, *lhs_tiled, lhs_tensor, k_iv_i32,
-                                         dot_dims.lhs_contracting_dimensions(0),
-                                         lhs_k_zero));
-          }
-
-          // Mask squeezed RHS at K-boundary inline (MaskDotOperand reads the
-          // pre-squeeze HLO shape, so we replicate its logic here with the
-          // squeezed shape).
-          {
-            int64_t K = rhs_tiled->hlo()->shape().dimensions(rhs_k_orig);
-            int64_t tile_k = rhs_squeezed_shape[rhs_k_dim_after_squeeze];
-            if (K % tile_k != 0) {
-              Value tile_size_value = CreateConst(b, b.getI32Type(), tile_k);
-              Value num_full_tiles = arith::DivSIOp::create(
-                  b, CreateConst(b, b.getI32Type(), K), tile_size_value);
-              auto cond = arith::CmpIOp::create(b, arith::CmpIPredicate::sge,
-                                                k_iv_i32, num_full_tiles);
-              auto if_mask = mlir::scf::IfOp::create(b, rhs_tensor.getType(),
-                                                     cond, /*withElse=*/true);
-              {
-                mlir::OpBuilder::InsertionGuard mg(b);
-                b.setInsertionPointToStart(if_mask.thenBlock());
-                Value tile_offset =
-                    arith::MulIOp::create(b, k_iv_i32, tile_size_value);
-                TensorValue range = Iota(b, tile_k);
-                TensorValue bcast_off = xtile::Splat(b, tile_offset, {tile_k});
-                Value indices = arith::AddIOp::create(b, range, bcast_off);
-                Value boundary = CreateConst(b, b.getI32Type(), K, {tile_k});
-                Value mask = arith::CmpIOp::create(b, arith::CmpIPredicate::slt,
-                                                   indices, boundary);
-                llvm::ArrayRef<int64_t> tile_shape =
-                    rhs_tensor.getType().getShape();
-                mask = xtile::BroadcastInDims(b, mlir::cast<TensorValue>(mask),
-                                              tile_shape,
-                                              {rhs_k_dim_after_squeeze});
-                ABSL_ASSIGN_OR_RETURN(
-                    Type elem_ty,
-                    PrimitiveTypeToMlirType(
-                        b, rhs_tiled->hlo()->shape().element_type()));
-                TensorValue zero = CreateConst(b, elem_ty, 0.0f, tile_shape);
-                Value masked =
-                    arith::SelectOp::create(b, mask, rhs_tensor, zero);
-                mlir::scf::YieldOp::create(b, masked);
-              }
-              {
-                mlir::OpBuilder::InsertionGuard mg(b);
-                b.setInsertionPointToStart(if_mask.elseBlock());
-                mlir::scf::YieldOp::create(b, rhs_tensor);
-              }
-              b.setInsertionPointAfter(if_mask);
-              rhs_tensor = mlir::cast<TensorValue>(if_mask.getResult(0));
-            }
-          }
-
-          // For batched ragged dot (num_batch_dims > 0), the lhs/rhs/acc have
-          // leading batch dims of size 1 (one batch element per tile). Triton's
-          // tt.dot operates on 2D tensors. Squeeze those size-1 batch dims
-          // before the dot, then unsqueeze the result.
-          //
-          // Non-batched path (num_batch_dims == 0) is unchanged.
-          auto lhs_type = lhs_tensor.getType();
-          auto acc_type_cast =
-              mlir::cast<mlir::RankedTensorType>(acc_k_in.getType());
-          auto acc_shape_ref = acc_type_cast.getShape();
-
-          // Compute 2-D (squeezed) shapes by dropping batch dims.
-          llvm::SmallVector<int64_t> lhs_2d_shape(
-              lhs_type.getShape().begin() + num_batch_dims,
-              lhs_type.getShape().end());
-          llvm::SmallVector<int64_t> rhs_2d_shape(
-              rhs_squeezed_shape.begin() + num_batch_dims,
-              rhs_squeezed_shape.end());
-          llvm::SmallVector<int64_t> acc_2d_shape(
-              acc_shape_ref.begin() + num_batch_dims, acc_shape_ref.end());
-
-          TensorValue dot_lhs = lhs_tensor;
-          TensorValue dot_rhs = rhs_tensor;
-          Value dot_acc = acc_k_in;
-          if (num_batch_dims > 0) {
-            dot_lhs = mlir::cast<TensorValue>(
-                stablehlo::ReshapeOp::create(
-                    b,
-                    mlir::RankedTensorType::get(lhs_2d_shape,
-                                                lhs_type.getElementType()),
-                    lhs_tensor)
-                    .getResult());
-            dot_rhs = mlir::cast<TensorValue>(
-                stablehlo::ReshapeOp::create(
-                    b,
-                    mlir::RankedTensorType::get(
-                        rhs_2d_shape, rhs_tensor.getType().getElementType()),
-                    rhs_tensor)
-                    .getResult());
-            dot_acc = stablehlo::ReshapeOp::create(
-                          b,
-                          mlir::RankedTensorType::get(
-                              acc_2d_shape, acc_type_cast.getElementType()),
-                          acc_k_in)
-                          .getResult();
-          }
-
-          // Build inner dot dimension numbers.
-          // When batch dims are squeezed out, contracting dim indices decrease
-          // by num_batch_dims (the batch dims occupied the leading positions).
-          DotDimensionNumbers inner_dot_dims;
-          inner_dot_dims.add_lhs_contracting_dimensions(
-              dot_dims.lhs_contracting_dimensions(0) - num_batch_dims);
-          inner_dot_dims.add_rhs_contracting_dimensions(
-              rhs_k_dim_after_squeeze - num_batch_dims);
-
-          // Partial dot: LHS [M, K] × RHS [K, N] → acc [M, N].
-          ABSL_ASSIGN_OR_RETURN(
-              Value acc_next_2d,
-              xtile::EmitSingleTileDot(
-                  b, *ragged_dot_instr, inner_dot_dims,
-                  xtile::DotOperands{dot_lhs, dot_rhs,
-                                     mlir::cast<TensorValue>(dot_acc)}));
-
-          // Unsqueeze batch dims back if they were squeezed.
-          Value acc_next = acc_next_2d;
-          if (num_batch_dims > 0) {
-            acc_next = stablehlo::ReshapeOp::create(
-                           b,
-                           mlir::RankedTensorType::get(
-                               llvm::SmallVector<int64_t>(acc_shape_ref.begin(),
-                                                          acc_shape_ref.end()),
-                               acc_type_cast.getElementType()),
-                           acc_next_2d)
-                           .getResult();
-          }
-          mlir::scf::YieldOp::create(b, acc_next);
-        }
-        b.setInsertionPointAfter(k_for_op);
-        mlir::scf::YieldOp::create(b, k_for_op.getResult(0));
-      }
-
-      {  // Else block: accumulator unchanged.
-        mlir::OpBuilder::InsertionGuard else_guard(b);
-        b.setInsertionPointToStart(if_op.elseBlock());
-        mlir::scf::YieldOp::create(b, acc_in);
-      }
-      b.setInsertionPointAfter(if_op);
-
-      // Update last_m for next G iteration.
-      Value last_m_next = mlir::arith::AddIOp::create(b, last_m, group_size_g);
-      mlir::scf::YieldOp::create(
-          b, mlir::ValueRange{if_op.getResult(0), last_m_next});
-    }
-    b.setInsertionPointAfter(g_for_op);
-
-    TensorValue result = mlir::cast<TensorValue>(g_for_op.getResult(0));
-
-    // Cast accumulator to the declared output element type if needed.
-    ABSL_ASSIGN_OR_RETURN(Type out_type,
-                     PrimitiveTypeToMlirType(
-                         b, tiled_ragged_dot.hlo()->shape().element_type()));
-    if (out_type != acc_type) {
-      result = mlir::cast<TensorValue>(Cast(b, result, out_type));
-    }
-    return result;
+    return EmitRaggedDotNonContracting(emitter_ctx, tiled_ragged_dot, rctx);
   }
-
-  return absl::UnimplementedError("kRaggedContracting: not yet implemented");
+  return EmitRaggedDotContracting(emitter_ctx, tiled_ragged_dot, rctx);
 }
 
 absl::StatusOr<TensorValue> EmitIota(

From 8948ff83f35a9b87f7c6b49d0278f77c1cfc3a95 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Mon, 12 Oct 2026 01:18:03 -0700
Subject: [PATCH 0561/1000] [XLA:GPU]: Migrate HostSendRecvAsyncEvents to use
 PerDeviceState

PiperOrigin-RevId: 997750963
---
 .../xla/xla/backends/gpu/runtime/BUILD        |  8 ++-
 .../gpu/runtime/host_send_recv_thunk.cc       | 64 +++++++++++--------
 .../gpu/runtime/host_send_recv_thunk.h        | 27 ++++----
 .../gpu/runtime/host_send_recv_thunk_test.cc  | 35 ++++++++++
 .../runtime/thunk_proto_deserialization.cc    | 12 ++--
 .../xla/xla/service/gpu/thunk_emitter.cc      |  3 +-
 6 files changed, 106 insertions(+), 43 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index 01634f1e2ef..63263c18fb0 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -3496,6 +3496,7 @@ cc_library(
     hdrs = ["host_send_recv_thunk.h"],
     deps = [
         ":host_async_thunk",
+        ":per_device_state",
         ":thunk",
         ":thunk_proto_cc",
         "//xla:shape_util",
@@ -3510,7 +3511,6 @@ cc_library(
         "//xla/stream_executor:stream_executor_h",
         "//xla/tsl/concurrency:async_value",
         "@com_google_absl//absl/base",
-        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
@@ -3518,7 +3518,6 @@ cc_library(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/strings:string_view",
-        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
         "@tsl//tsl/profiler/lib:traceme",
     ],
@@ -3534,7 +3533,12 @@ xla_cc_test(
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",
         "//xla/service:buffer_assignment",
+        "//xla/stream_executor:event",
+        "//xla/stream_executor:mock_stream_executor",
+        "//xla/tsl/concurrency:async_value",
         "//xla/tsl/util/proto:parse_text_proto",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
         "@com_google_googletest//:gtest_main",
     ],
 )
diff --git a/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.cc
index d1c0ab247b1..f14a3a53d67 100644
--- a/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.cc
@@ -29,9 +29,9 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
-#include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
 #include "xla/backends/gpu/runtime/host_async_thunk.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk.pb.h"
 #include "xla/runtime/device_id.h"
@@ -78,10 +78,10 @@ static absl::StatusOr<bool> ShouldSkip(
 absl::Status HostSendRecvAsyncEvents::Emplace(
     se::StreamExecutor* executor, int32_t channel_id,
     tsl::AsyncValueRef<std::unique_ptr<se::Event>> event) {
-  Key key = {executor, channel_id};
-
-  absl::MutexLock lock(mutex_);
-  if (auto it = events_.try_emplace(key, std::move(event)); it.second) {
+  ABSL_ASSIGN_OR_RETURN(DeviceEvents * device_events,
+                   device_events_.GetOrCreate(executor->device_ordinal()));
+  if (auto it = device_events->events.try_emplace(channel_id, std::move(event));
+      it.second) {
     return absl::OkStatus();
   }
 
@@ -92,10 +92,12 @@ absl::Status HostSendRecvAsyncEvents::Emplace(
 absl::StatusOr<AsyncValueRef<std::unique_ptr<se::Event>>>
 HostSendRecvAsyncEvents::Extract(se::StreamExecutor* executor,
                                  int32_t channel_id) {
-  Key key = {executor, channel_id};
-
-  absl::MutexLock lock(mutex_);
-  if (auto event = events_.extract(key)) {
+  DeviceEvents* device_events = device_events_.Find(executor->device_ordinal());
+  if (device_events == nullptr) {
+    return absl::InternalError(absl::StrFormat(
+        "Async send/recv event was not found (channel_id==%d)", channel_id));
+  }
+  if (auto event = device_events->events.extract(channel_id)) {
     return std::move(event.mapped());
   }
 
@@ -145,7 +147,7 @@ absl::StatusOr<ThunkProto> HostSendThunk::ToProto() const {
 absl::StatusOr<std::unique_ptr<HostSendThunk>> HostSendThunk::FromProto(
     ThunkInfo thunk_info, const HostSendThunkProto& proto,
     absl::Span<const BufferAllocation> allocations,
-    HostSendRecvAsyncEventsMap& async_events_map) {
+    HostSendRecvAsyncEventsMap& async_events_map, int devices_per_host) {
   ABSL_ASSIGN_OR_RETURN(Shape shape, Shape::FromProto(proto.shape()));
   ABSL_ASSIGN_OR_RETURN(
       BufferAllocation::Slice buffer,
@@ -157,9 +159,12 @@ absl::StatusOr<std::unique_ptr<HostSendThunk>> HostSendThunk::FromProto(
   absl::flat_hash_map<std::string, std::string> frontend_attrs(
       proto.frontend_attrs().begin(), proto.frontend_attrs().end());
 
-  auto [async_event_it, _] = async_events_map.try_emplace(
-      AsyncEventsUniqueId(proto.async_events_unique_id()),
-      std::make_shared<HostSendRecvAsyncEvents>());
+  auto [async_event_it, inserted] = async_events_map.try_emplace(
+      AsyncEventsUniqueId(proto.async_events_unique_id()), nullptr);
+  if (inserted) {
+    async_event_it->second =
+        std::make_shared<HostSendRecvAsyncEvents>(devices_per_host);
+  }
   return std::make_unique<HostSendThunk>(
       thunk_info, ShapedSlice{buffer, shape}, proto.channel_id(),
       async_event_it->second, std::move(frontend_attrs), device_constraint);
@@ -250,15 +255,18 @@ absl::StatusOr<ThunkProto> HostSendDoneThunk::ToProto() const {
 absl::StatusOr<std::unique_ptr<HostSendDoneThunk>> HostSendDoneThunk::FromProto(
     ThunkInfo thunk_info, const HostSendDoneThunkProto& proto,
     absl::Span<const BufferAllocation> allocations,
-    HostSendRecvAsyncEventsMap& async_events_map) {
+    HostSendRecvAsyncEventsMap& async_events_map, int devices_per_host) {
   std::optional<GlobalDeviceId> device_constraint;
   if (proto.has_device_constraint()) {
     device_constraint = GlobalDeviceId(proto.device_constraint());
   }
 
-  auto [async_event_it, _] = async_events_map.try_emplace(
-      AsyncEventsUniqueId(proto.async_events_unique_id()),
-      std::make_shared<HostSendRecvAsyncEvents>());
+  auto [async_event_it, inserted] = async_events_map.try_emplace(
+      AsyncEventsUniqueId(proto.async_events_unique_id()), nullptr);
+  if (inserted) {
+    async_event_it->second =
+        std::make_shared<HostSendRecvAsyncEvents>(devices_per_host);
+  }
 
   std::optional<ShapedSlice> slice;
   if (proto.has_buffer()) {
@@ -349,7 +357,7 @@ absl::StatusOr<ThunkProto> HostRecvThunk::ToProto() const {
 absl::StatusOr<std::unique_ptr<HostRecvThunk>> HostRecvThunk::FromProto(
     ThunkInfo thunk_info, const HostRecvThunkProto& proto,
     absl::Span<const BufferAllocation> allocations,
-    HostSendRecvAsyncEventsMap& async_events_map) {
+    HostSendRecvAsyncEventsMap& async_events_map, int devices_per_host) {
   ABSL_ASSIGN_OR_RETURN(Shape shape, Shape::FromProto(proto.shape()));
   ABSL_ASSIGN_OR_RETURN(
       BufferAllocation::Slice buffer,
@@ -361,9 +369,12 @@ absl::StatusOr<std::unique_ptr<HostRecvThunk>> HostRecvThunk::FromProto(
   absl::flat_hash_map<std::string, std::string> frontend_attrs(
       proto.frontend_attrs().begin(), proto.frontend_attrs().end());
 
-  auto [async_event_it, _] = async_events_map.try_emplace(
-      AsyncEventsUniqueId(proto.async_events_unique_id()),
-      std::make_shared<HostSendRecvAsyncEvents>());
+  auto [async_event_it, inserted] = async_events_map.try_emplace(
+      AsyncEventsUniqueId(proto.async_events_unique_id()), nullptr);
+  if (inserted) {
+    async_event_it->second =
+        std::make_shared<HostSendRecvAsyncEvents>(devices_per_host);
+  }
   return std::make_unique<HostRecvThunk>(
       thunk_info, std::move(shape), buffer, proto.channel_id(),
       async_event_it->second, std::move(frontend_attrs), device_constraint);
@@ -456,15 +467,18 @@ absl::StatusOr<ThunkProto> HostRecvDoneThunk::ToProto() const {
 absl::StatusOr<std::unique_ptr<HostRecvDoneThunk>> HostRecvDoneThunk::FromProto(
     ThunkInfo thunk_info, const HostRecvDoneThunkProto& proto,
     absl::Span<const BufferAllocation> allocations,
-    HostSendRecvAsyncEventsMap& async_events_map) {
+    HostSendRecvAsyncEventsMap& async_events_map, int devices_per_host) {
   std::optional<GlobalDeviceId> device_constraint;
   if (proto.has_device_constraint()) {
     device_constraint = GlobalDeviceId(proto.device_constraint());
   }
 
-  auto [async_event_it, _] = async_events_map.try_emplace(
-      AsyncEventsUniqueId(proto.async_events_unique_id()),
-      std::make_shared<HostSendRecvAsyncEvents>());
+  auto [async_event_it, inserted] = async_events_map.try_emplace(
+      AsyncEventsUniqueId(proto.async_events_unique_id()), nullptr);
+  if (inserted) {
+    async_event_it->second =
+        std::make_shared<HostSendRecvAsyncEvents>(devices_per_host);
+  }
 
   std::optional<ShapedSlice> slice;
   if (proto.has_buffer()) {
diff --git a/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.h b/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.h
index 9a83f160943..78ec36fb54b 100644
--- a/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.h
+++ b/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk.h
@@ -20,15 +20,13 @@ limitations under the License.
 #include <memory>
 #include <optional>
 #include <string>
-#include <utility>
 
-#include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
 #include "xla/backends/gpu/runtime/host_async_thunk.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk.pb.h"
 #include "xla/runtime/buffer_use.h"
@@ -67,6 +65,9 @@ namespace xla::gpu {
 // runtime, as it has name conflict.
 class HostSendRecvAsyncEvents {
  public:
+  explicit HostSendRecvAsyncEvents(int devices_per_host = 0)
+      : device_events_(devices_per_host) {}
+
   // Emplace a new send/recv completion event.
   absl::Status Emplace(se::StreamExecutor* executor, int32_t channel_id,
                        tsl::AsyncValueRef<std::unique_ptr<se::Event>> event);
@@ -76,11 +77,15 @@ class HostSendRecvAsyncEvents {
       se::StreamExecutor* executor, int32_t channel_id);
 
  private:
-  using Key = std::pair<se::StreamExecutor*, /*channel_id=*/int64_t>;
+  // Not internally synchronized. Relies on the caller serializing host-side
+  // thunk execution per device (e.g., PJRT per-device execute_thread).
+  struct DeviceEvents {
+    absl::flat_hash_map</*channel_id=*/int64_t,
+                        tsl::AsyncValueRef<std::unique_ptr<se::Event>>>
+        events;
+  };
 
-  absl::Mutex mutex_;
-  absl::flat_hash_map<Key, tsl::AsyncValueRef<std::unique_ptr<se::Event>>>
-      events_ ABSL_GUARDED_BY(mutex_);
+  PerDeviceState<DeviceEvents> device_events_;
 };
 
 // A map from a unique id to a shared pointer to HostSendRecvAsyncEvents.
@@ -99,7 +104,7 @@ class HostSendThunk : public HostAsyncThunk {
   static absl::StatusOr<std::unique_ptr<HostSendThunk>> FromProto(
       ThunkInfo thunk_info, const HostSendThunkProto& proto,
       absl::Span<const BufferAllocation> allocations,
-      HostSendRecvAsyncEventsMap& async_events_map);
+      HostSendRecvAsyncEventsMap& async_events_map, int devices_per_host = 0);
 
   HostSendThunk(ThunkInfo thunk_info, ShapedSlice buffer, int64_t channel_id,
                 std::shared_ptr<HostSendRecvAsyncEvents> events,
@@ -139,7 +144,7 @@ class HostSendDoneThunk : public HostAsyncThunk {
   static absl::StatusOr<std::unique_ptr<HostSendDoneThunk>> FromProto(
       ThunkInfo thunk_info, const HostSendDoneThunkProto& proto,
       absl::Span<const BufferAllocation> allocations,
-      HostSendRecvAsyncEventsMap& async_events_map);
+      HostSendRecvAsyncEventsMap& async_events_map, int devices_per_host = 0);
 
   HostSendDoneThunk(ThunkInfo thunk_info, std::optional<ShapedSlice> slice,
                     int64_t channel_id,
@@ -180,7 +185,7 @@ class HostRecvThunk : public HostAsyncThunk {
   static absl::StatusOr<std::unique_ptr<HostRecvThunk>> FromProto(
       ThunkInfo thunk_info, const HostRecvThunkProto& proto,
       absl::Span<const BufferAllocation> allocations,
-      HostSendRecvAsyncEventsMap& async_events_map);
+      HostSendRecvAsyncEventsMap& async_events_map, int devices_per_host = 0);
 
   HostRecvThunk(ThunkInfo thunk_info, Shape shape,
                 BufferAllocation::Slice buffer, int64_t channel_id,
@@ -222,7 +227,7 @@ class HostRecvDoneThunk : public HostAsyncThunk {
   static absl::StatusOr<std::unique_ptr<HostRecvDoneThunk>> FromProto(
       ThunkInfo thunk_info, const HostRecvDoneThunkProto& proto,
       absl::Span<const BufferAllocation> allocations,
-      HostSendRecvAsyncEventsMap& async_events_map);
+      HostSendRecvAsyncEventsMap& async_events_map, int devices_per_host = 0);
 
   HostRecvDoneThunk(ThunkInfo thunk_info, std::optional<ShapedSlice> slice,
                     int64_t channel_id,
diff --git a/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk_test.cc
index 5f0dfbabcdd..07c64b9650a 100644
--- a/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/host_send_recv_thunk_test.cc
@@ -20,19 +20,54 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk.pb.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
+#include "xla/stream_executor/event.h"
+#include "xla/stream_executor/mock_stream_executor.h"
+#include "xla/tsl/concurrency/async_value_ref.h"
 #include "xla/tsl/util/proto/parse_text_proto.h"
 #include "xla/xla_data.pb.h"
 
 namespace xla::gpu {
 namespace {
 
+using ::testing::Return;
 using ::tsl::proto_testing::ParseTextProtoOrDie;
 
+TEST(HostSendRecvAsyncEventsTest, EmplaceAndExtractAcrossDevices) {
+  se::MockStreamExecutor executor0;
+  se::MockStreamExecutor executor1;
+  EXPECT_CALL(executor0, device_ordinal()).WillRepeatedly(Return(0));
+  EXPECT_CALL(executor1, device_ordinal()).WillRepeatedly(Return(1));
+
+  HostSendRecvAsyncEvents events(/*devices_per_host=*/2);
+
+  auto event0 =
+      tsl::MakeConstructedAsyncValueRef<std::unique_ptr<se::Event>>(nullptr);
+  auto event1 =
+      tsl::MakeConstructedAsyncValueRef<std::unique_ptr<se::Event>>(nullptr);
+
+  ASSERT_OK(events.Emplace(&executor0, /*channel_id=*/42, event0));
+  ASSERT_OK(events.Emplace(&executor1, /*channel_id=*/42, event1));
+  EXPECT_THAT(events.Emplace(&executor0, /*channel_id=*/42, event0),
+              absl_testing::StatusIs(absl::StatusCode::kInternal));
+
+  ASSERT_OK_AND_ASSIGN(auto extracted0,
+                       events.Extract(&executor0, /*channel_id=*/42));
+  EXPECT_EQ(extracted0.GetAsyncValue(), event0.GetAsyncValue());
+  EXPECT_THAT(events.Extract(&executor0, /*channel_id=*/42),
+              absl_testing::StatusIs(absl::StatusCode::kInternal));
+
+  ASSERT_OK_AND_ASSIGN(auto extracted1,
+                       events.Extract(&executor1, /*channel_id=*/42));
+  EXPECT_EQ(extracted1.GetAsyncValue(), event1.GetAsyncValue());
+}
+
 TEST(HostSendRecvThunkTest, LinkSendAndSendDone) {
   Thunk::ThunkInfo thunk_info;
   Shape shape = ShapeUtil::MakeShape(F32, {1});
diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
index d71fc6b2709..55117838f60 100644
--- a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
@@ -254,19 +254,23 @@ absl::StatusOr<std::unique_ptr<Thunk>> DeserializeThunkProtoImpl(
     case ThunkProto::kHostSendThunk:
       return HostSendThunk::FromProto(
           std::move(thunk_info), thunk_proto.host_send_thunk(),
-          buffer_allocations, host_send_recv_async_events_map);
+          buffer_allocations, host_send_recv_async_events_map,
+          gpu_topology.has_value() ? gpu_topology->num_devices_per_host() : 0);
     case ThunkProto::kHostSendDoneThunk:
       return HostSendDoneThunk::FromProto(
           std::move(thunk_info), thunk_proto.host_send_done_thunk(),
-          buffer_allocations, host_send_recv_async_events_map);
+          buffer_allocations, host_send_recv_async_events_map,
+          gpu_topology.has_value() ? gpu_topology->num_devices_per_host() : 0);
     case ThunkProto::kHostRecvThunk:
       return HostRecvThunk::FromProto(
           std::move(thunk_info), thunk_proto.host_recv_thunk(),
-          buffer_allocations, host_send_recv_async_events_map);
+          buffer_allocations, host_send_recv_async_events_map,
+          gpu_topology.has_value() ? gpu_topology->num_devices_per_host() : 0);
     case ThunkProto::kHostRecvDoneThunk:
       return HostRecvDoneThunk::FromProto(
           std::move(thunk_info), thunk_proto.host_recv_done_thunk(),
-          buffer_allocations, host_send_recv_async_events_map);
+          buffer_allocations, host_send_recv_async_events_map,
+          gpu_topology.has_value() ? gpu_topology->num_devices_per_host() : 0);
     case ThunkProto::kOutfeedThunk:
       return OutfeedThunk::FromProto(std::move(thunk_info),
                                      thunk_proto.outfeed_thunk(),
diff --git a/third_party/xla/xla/service/gpu/thunk_emitter.cc b/third_party/xla/xla/service/gpu/thunk_emitter.cc
index f6304494b57..c75773a320d 100644
--- a/third_party/xla/xla/service/gpu/thunk_emitter.cc
+++ b/third_party/xla/xla/service/gpu/thunk_emitter.cc
@@ -540,7 +540,8 @@ ThunkEmitter::ThunkEmitter(
     llvm_ir::LLVMCommandLineOptionsReleasableLock* absl_nonnull
         llvm_options_lock)
     : ir_emitter_context_(ir_emitter_context),
-      send_recv_events_(std::make_shared<HostSendRecvAsyncEvents>()),
+      send_recv_events_(std::make_shared<HostSendRecvAsyncEvents>(
+          ir_emitter_context_->gpu_topology().num_devices_per_host())),
       call_graph_(CallGraph::Build(&ir_emitter_context->hlo_module())),
       constants_module_context_(std::make_unique<llvm::LLVMContext>()),
       constants_module_(ir_emitter_context_->CreateLLVMModule(

From d90fda9025fec3f3ab8c9c4dad9cb88428e0de5f Mon Sep 17 00:00:00 2001
From: Bill Varcho <varcho@google.com>
Date: Mon, 12 Oct 2026 01:52:39 -0700
Subject: [PATCH 0562/1000] [SDY][XLA][1/3] bump shardy commit.

PiperOrigin-RevId: 997765575
---
 .../xla/third_party/shardy/temporary.patch    | 32 -------------------
 .../xla/third_party/shardy/workspace.bzl      |  4 +--
 2 files changed, 2 insertions(+), 34 deletions(-)

diff --git a/third_party/xla/third_party/shardy/temporary.patch b/third_party/xla/third_party/shardy/temporary.patch
index 8fcc0f46504..e69de29bb2d 100644
--- a/third_party/xla/third_party/shardy/temporary.patch
+++ b/third_party/xla/third_party/shardy/temporary.patch
@@ -1,32 +0,0 @@
-diff --git a/third_party/llvm/generated.patch b/third_party/llvm/generated.patch
-index 509398da..e9c206ce 100644
---- a/third_party/llvm/generated.patch
-+++ b/third_party/llvm/generated.patch
-@@ -1 +1,12 @@
- Auto generated patch. Do not edit or delete it, even if empty.
-+diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll b/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll
-+--- a/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll
-++++ b/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll
-+@@ -1,5 +1,7 @@
-+ ; Checks if conversion is only inserted for the spilled register
-+ ; instead of 2 conversions for the W register
-++
-++; REQUIRES: asserts
-+ ; RUN: llc -O2 -march=hexagon -mcpu=hexagonv81 -enable-xqf-gen=true \
-+ ; RUN: -hexagon-qfloat-mode=ieee -mattr=+hvxv81,+hvx-length128B \
-+ ; RUN: -enable-postra-xqf-check -debug-only=handle-qfp 2>&1 < %s -o /dev/null | FileCheck %s
-diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
-index 7b8ac881..85cdf966 100644
---- a/third_party/llvm/workspace.bzl
-+++ b/third_party/llvm/workspace.bzl
-@@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
- 
- def repo(name):
-     """Imports LLVM."""
--    LLVM_COMMIT = "c4d9267a31897dd76b2aab5196b75bffb49e6846"
--    LLVM_SHA256 = "a96c8b33f0d8d3be008b8c9f74ef310254c4f7f7edf8c04c76f4d99ed9d34f1a"
-+    LLVM_COMMIT = "b08e005995ed29a2f89aa9a3f599a9d45db65523"
-+    LLVM_SHA256 = "3f4f1765816e22a21ac41dbf71d91d519b0b371b05d7e33ef2a0a862c838420c"
- 
-     tf_http_archive(
-         name = name,
diff --git a/third_party/xla/third_party/shardy/workspace.bzl b/third_party/xla/third_party/shardy/workspace.bzl
index 83345d191e0..659b180a7ba 100644
--- a/third_party/xla/third_party/shardy/workspace.bzl
+++ b/third_party/xla/third_party/shardy/workspace.bzl
@@ -18,8 +18,8 @@
 load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 
 def repo():
-    SHARDY_COMMIT = "ba51e5d4ab46ef7b2f374d47f64c3317b9134441"
-    SHARDY_SHA256 = "d3c79a9a3c219cd0e68f696af200255eee185a5b128e274443d4ecf92e1e5fb5"
+    SHARDY_COMMIT = "b5950641f8b5cbbebdf437e39197e1a170611587"
+    SHARDY_SHA256 = "ff228c5220c99e8de7fa8a480104ffaf1f6801c3b2641b05b6f04bc626779f12"
 
     tf_http_archive(
         name = "shardy",

From 5b17a939cc8a81ee72c2947f867d1711df9ef86a Mon Sep 17 00:00:00 2001
From: Mikhail Goncharov <goncharov@google.com>
Date: Mon, 12 Oct 2026 02:03:53 -0700
Subject: [PATCH 0563/1000] [XLA:GPU] Standardize DebugOptions init for
 autotuner tests

Previously, autotuner tests default-constructed DebugOptions, resulting in empty/zeroed protobuf options.
Now they initialize the fixture's debug_options_ backing field with HloHardwareIndependentTestBase::GetDebugOptionsForTest().

PiperOrigin-RevId: 997771067
---
 third_party/xla/xla/backends/gpu/autotuner/BUILD    |  1 -
 .../gpu/autotuner/block_level_emitter_test.cc       |  4 ++--
 .../xla/xla/backends/gpu/autotuner/cublaslt_test.cc |  4 +++-
 .../xla/xla/backends/gpu/autotuner/cudnn_test.cc    | 10 ++++++----
 .../xla/xla/backends/gpu/autotuner/factory_test.cc  |  2 ++
 .../backends/gpu/autotuner/fission_backend_test.cc  | 13 +++++++++----
 .../xla/backends/gpu/autotuner/hipblaslt_test.cc    |  4 +++-
 .../xla/xla/backends/gpu/autotuner/miopen_test.cc   |  4 +++-
 .../backends/gpu/autotuner/native_emitter_test.cc   |  4 +++-
 .../xla/xla/backends/gpu/autotuner/triton_test.cc   | 11 +++++------
 10 files changed, 36 insertions(+), 21 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/autotuner/BUILD b/third_party/xla/xla/backends/gpu/autotuner/BUILD
index 2e39f0f355b..68a93676691 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/BUILD
+++ b/third_party/xla/xla/backends/gpu/autotuner/BUILD
@@ -121,7 +121,6 @@ xla_test(
     use_legacy_runtime = True,
     deps = [
         ":block_level_emitter",
-        "//xla:debug_options_flags",
         "//xla:xla_proto_cc",
         "//xla/backends/autotuner:codegen_backend",
         "//xla/backends/gpu/codegen/emitters:mlir_kernel_emitter",
diff --git a/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter_test.cc b/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter_test.cc
index 4cbf832dd08..958e19de43e 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter_test.cc
@@ -28,7 +28,6 @@ limitations under the License.
 #include "xla/backends/autotuner/codegen_backend.h"
 #include "xla/backends/gpu/codegen/emitters/mlir_kernel_emitter.h"
 #include "xla/codegen/xtile/xtile_config.pb.h"
-#include "xla/debug_options_flags.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/service/compiler.h"
@@ -80,7 +79,8 @@ class TritonBlockLevelFusionEmitterBackendTest
     : public HloHardwareIndependentTestBase {
  protected:
   TritonBlockLevelFusionEmitterBackendTest()
-      : debug_options_(GetDebugOptionsFromFlags()),
+      : debug_options_(
+            HloHardwareIndependentTestBase::GetDebugOptionsForTest()),
         stream_executor_(GetDefaultStreamExecutor()),
         target_config_(stream_executor_),
         backend_(&debug_options_, &compiler_,
diff --git a/third_party/xla/xla/backends/gpu/autotuner/cublaslt_test.cc b/third_party/xla/xla/backends/gpu/autotuner/cublaslt_test.cc
index afb797c263a..4e6fecbe34f 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/cublaslt_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/cublaslt_test.cc
@@ -117,7 +117,9 @@ class CublasLtBackendTest : public HloHardwareIndependentTestBase {
   CublasLtBackend backend_;
 
   CublasLtBackendTest()
-      : stream_executor_(PlatformUtil::GetDefaultPlatform()
+      : debug_options_(
+            HloHardwareIndependentTestBase::GetDebugOptionsForTest()),
+        stream_executor_(PlatformUtil::GetDefaultPlatform()
                              .value()
                              ->ExecutorForDevice(0)
                              .value()),
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 2aaa51d83b9..c424c28f5f8 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/cudnn_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/cudnn_test.cc
@@ -331,10 +331,11 @@ TEST_F(CudnnBackendTest, GetDefaultConfigSucceedsWithNullStreamExecutor) {
                        GetGpuTargetConfig(GpuModel::H100_SXM));
   ASSERT_OK_AND_ASSIGN(Compiler::GpuTargetConfig target_config,
                        Compiler::GpuTargetConfig::FromProto(proto));
-  debug_options_.set_xla_gpu_cudnn_deviceless_compilation_mode(
+  DebugOptions debug_options = debug_options_;
+  debug_options.set_xla_gpu_cudnn_deviceless_compilation_mode(
       DebugOptions::CUDNN_DEVICELESS_COMPILATION_AUTO);
   CudnnBackend backend_without_stream_executor(
-      /*stream_executor=*/nullptr, &debug_options_, &compiler_, &target_config);
+      /*stream_executor=*/nullptr, &debug_options, &compiler_, &target_config);
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> hlo_module,
                        ParseAndReturnVerifiedModule(kCudnnFusionHlo));
 
@@ -350,10 +351,11 @@ TEST_F(CudnnBackendTest, CudnnDevicelessCompilationDisabledFails) {
                        GetGpuTargetConfig(GpuModel::H100_SXM));
   ASSERT_OK_AND_ASSIGN(Compiler::GpuTargetConfig target_config,
                        Compiler::GpuTargetConfig::FromProto(proto));
-  debug_options_.set_xla_gpu_cudnn_deviceless_compilation_mode(
+  DebugOptions debug_options = debug_options_;
+  debug_options.set_xla_gpu_cudnn_deviceless_compilation_mode(
       DebugOptions::CUDNN_DEVICELESS_COMPILATION_DISABLED);
   CudnnBackend backend_without_stream_executor(
-      /*stream_executor=*/nullptr, &debug_options_, &compiler_, &target_config);
+      /*stream_executor=*/nullptr, &debug_options, &compiler_, &target_config);
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> hlo_module,
                        ParseAndReturnVerifiedModule(kCudnnFusionHlo));
 
diff --git a/third_party/xla/xla/backends/gpu/autotuner/factory_test.cc b/third_party/xla/xla/backends/gpu/autotuner/factory_test.cc
index 2541c00c387..0b2c889100b 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/factory_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/factory_test.cc
@@ -97,6 +97,8 @@ class FactoryTest : public xla::HloHardwareIndependentTestBase,
         compiler_(GetCompilerForPlatform(platform_)),
         stream_executor_(GetStreamExecutor(platform_)),
         target_config_(stream_executor_),
+        debug_options_(
+            HloHardwareIndependentTestBase::GetDebugOptionsForTest()),
         allocator_(stream_executor_) {}
 };
 
diff --git a/third_party/xla/xla/backends/gpu/autotuner/fission_backend_test.cc b/third_party/xla/xla/backends/gpu/autotuner/fission_backend_test.cc
index 555ad67d03c..13ce51f17de 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/fission_backend_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/fission_backend_test.cc
@@ -280,7 +280,9 @@ class FissionTest : public HloHardwareIndependentTestBase,
   std::unique_ptr<FissionBackend> fission_backend_;
 
   FissionTest()
-      : platform_(PlatformUtil::GetDefaultPlatform().value()),
+      : debug_options_(
+            HloHardwareIndependentTestBase::GetDebugOptionsForTest()),
+        platform_(PlatformUtil::GetDefaultPlatform().value()),
         stream_executor_(platform_->ExecutorForDevice(0).value()),
         compiler_(Compiler::GetForPlatform(platform_->id()).value()),
         target_config_(stream_executor_),
@@ -496,7 +498,9 @@ class CublasFissionBackendTest : public HloHardwareIndependentTestBase {
   std::unique_ptr<FissionBackend> fission_backend_;
 
   CublasFissionBackendTest()
-      : platform_(PlatformUtil::GetDefaultPlatform().value()),
+      : debug_options_(
+            HloHardwareIndependentTestBase::GetDebugOptionsForTest()),
+        platform_(PlatformUtil::GetDefaultPlatform().value()),
         stream_executor_(platform_->ExecutorForDevice(0).value()),
         compiler_(Compiler::GetForPlatform(platform_->id()).value()),
         target_config_(stream_executor_),
@@ -687,8 +691,9 @@ TEST_F(CublasFissionBackendTest, CublasFallbackForBf16Bf16F32Algorithm) {
 //     returned.
 TEST_F(CublasFissionBackendTest,
        GetSupportedConfigsRespectsUserGemmRewriteSizeThreshold) {
-  // Confirm the fixture's debug options use the proto default (0), which
-  // differs from DefaultDebugOptionsIgnoringFlags() that sets it to 100.
+  // Set the fixture's debug options threshold to 0, which differs from
+  // DefaultDebugOptionsIgnoringFlags() that sets it to 100.
+  debug_options_.set_xla_gpu_gemm_rewrite_size_threshold(0);
   EXPECT_EQ(debug_options_.xla_gpu_gemm_rewrite_size_threshold(), 0);
 
   // A tiny f32 fusion whose dot's combined size is 50 < 100
diff --git a/third_party/xla/xla/backends/gpu/autotuner/hipblaslt_test.cc b/third_party/xla/xla/backends/gpu/autotuner/hipblaslt_test.cc
index d0abbd8ffb0..9260089f530 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/hipblaslt_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/hipblaslt_test.cc
@@ -107,7 +107,9 @@ class HipblasLtBackendTest : public HloHardwareIndependentTestBase {
   HipblasLtBackend backend_;
 
   HipblasLtBackendTest()
-      : stream_executor_(PlatformUtil::GetDefaultPlatform()
+      : debug_options_(
+            HloHardwareIndependentTestBase::GetDebugOptionsForTest()),
+        stream_executor_(PlatformUtil::GetDefaultPlatform()
                              .value()
                              ->ExecutorForDevice(0)
                              .value()),
diff --git a/third_party/xla/xla/backends/gpu/autotuner/miopen_test.cc b/third_party/xla/xla/backends/gpu/autotuner/miopen_test.cc
index b675f14abcc..8284117472f 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/miopen_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/miopen_test.cc
@@ -81,7 +81,9 @@ class MIOpenBackendTest : public HloHardwareIndependentTestBase {
   MIOpenBackend backend_;
 
   MIOpenBackendTest()
-      : stream_executor_(PlatformUtil::GetDefaultPlatform()
+      : debug_options_(
+            HloHardwareIndependentTestBase::GetDebugOptionsForTest()),
+        stream_executor_(PlatformUtil::GetDefaultPlatform()
                              .value()
                              ->ExecutorForDevice(0)
                              .value()),
diff --git a/third_party/xla/xla/backends/gpu/autotuner/native_emitter_test.cc b/third_party/xla/xla/backends/gpu/autotuner/native_emitter_test.cc
index 25eb45f1bfd..c53c7115c03 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/native_emitter_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/native_emitter_test.cc
@@ -116,7 +116,9 @@ ENTRY %entry_computation (p0: f32[32,16], p1: f32[32,16]) -> (f32[32,16], f32[32
 class NativeEmitterBackendTest : public HloHardwareIndependentTestBase {
  protected:
   NativeEmitterBackendTest()
-      : platform_(PlatformUtil::GetDefaultPlatform().value()),
+      : debug_options_(
+            HloHardwareIndependentTestBase::GetDebugOptionsForTest()),
+        platform_(PlatformUtil::GetDefaultPlatform().value()),
         stream_executor_(platform_->ExecutorForDevice(0).value()),
         target_config_(stream_executor_),
         compiler_(Compiler::GetForPlatform(platform_->id()).value()),
diff --git a/third_party/xla/xla/backends/gpu/autotuner/triton_test.cc b/third_party/xla/xla/backends/gpu/autotuner/triton_test.cc
index 5a03c976d14..8fccdd56f79 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/triton_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/triton_test.cc
@@ -131,7 +131,9 @@ ENTRY %entry (lhs: bf16[4,4], rhs: bf16[4,4], lhs_scale: bf16[1,1], rhs_scale: b
 class TritonBackendTest : public HloHardwareIndependentTestBase {
  protected:
   TritonBackendTest()
-      : platform_(PlatformUtil::GetDefaultPlatform().value()),
+      : debug_options_(
+            HloHardwareIndependentTestBase::GetDebugOptionsForTest()),
+        platform_(PlatformUtil::GetDefaultPlatform().value()),
         stream_executor_(platform_->ExecutorForDevice(0).value()),
         target_config_(stream_executor_),
         alias_info_(stream_executor_->GetDeviceDescription()),
@@ -599,11 +601,8 @@ TEST_F(TritonBackendTest, WarpSpecializationConfigsDoNotHaveNumStagesTwo) {
       true);
   debug_options_.set_xla_gpu_exhaustive_tiling_search(true);
 
-  HloModuleConfig config;
-  config.set_debug_options(debug_options_);
-  ASSERT_OK_AND_ASSIGN(
-      std::unique_ptr<HloModule> module,
-      ParseAndReturnVerifiedModule(kSimpleGemmFusionHlo, config));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kSimpleGemmFusionHlo));
 
   ASSERT_OK_AND_ASSIGN(std::vector<std::unique_ptr<BackendConfig>> configs,
                        backend_.GetSupportedConfigs(

From c38e62a390ebb964a4004a290568b12eae0a4388 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Mon, 12 Oct 2026 02:19:07 -0700
Subject: [PATCH 0564/1000] [XLA:GPU]: Migrate HostExecuteAsyncEvents to use
 PerDeviceState

PiperOrigin-RevId: 997778396
---
 .../xla/xla/backends/gpu/runtime/BUILD        |  5 +-
 .../gpu/runtime/host_execute_thunk.cc         | 59 ++++++++++++-------
 .../backends/gpu/runtime/host_execute_thunk.h | 29 +++++----
 .../gpu/runtime/host_execute_thunk_test.cc    | 42 +++++++++++++
 .../runtime/thunk_proto_deserialization.cc    |  6 +-
 .../xla/xla/service/gpu/thunk_emitter.cc      | 14 +++--
 6 files changed, 112 insertions(+), 43 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index 63263c18fb0..ad6d968a9af 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -4754,6 +4754,7 @@ cc_library(
     deps = [
         ":command",
         ":host_async_thunk",
+        ":per_device_state",
         ":thunk",
         ":thunk_proto_cc",
         "//xla:executable_run_options",
@@ -4782,7 +4783,6 @@ cc_library(
         "//xla/tsl/platform:env",
         "//xla/tsl/util:unique_any",
         "@com_google_absl//absl/base",
-        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/log",
@@ -4792,7 +4792,6 @@ cc_library(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
-        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
         "@tsl//tsl/platform:platform_port",
         "@tsl//tsl/profiler/lib:traceme",
@@ -4826,7 +4825,9 @@ xla_test(
         "//xla/service/cpu:cpu_aot_compilation_result",
         "//xla/service/gpu:buffer_allocations",
         "//xla/stream_executor:device_address",
+        "//xla/stream_executor:event",
         "//xla/stream_executor:memory_space",
+        "//xla/stream_executor:mock_stream_executor",
         "//xla/stream_executor:platform_manager",
         "//xla/stream_executor:stream_executor_address_allocator",
         "//xla/stream_executor:stream_executor_h",
diff --git a/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk.cc
index 7aee81ab416..180ea4ae08f 100644
--- a/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk.cc
@@ -27,6 +27,7 @@ limitations under the License.
 
 #include "absl/base/call_once.h"
 #include "absl/base/casts.h"
+#include "absl/container/flat_hash_map.h"
 #include "absl/container/inlined_vector.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
@@ -35,10 +36,11 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
-#include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
 #include "xla/backends/gpu/host_offloading/gpu_host_offloading_allocator.h"
 #include "xla/backends/gpu/runtime/command.h"
+#include "xla/backends/gpu/runtime/host_async_thunk.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk.pb.h"
 #include "xla/core/host_offloading/host_offloading_allocator.h"
@@ -410,9 +412,9 @@ HostExecuteAsyncEvents::CreateEvent(se::StreamExecutor* executor,
   auto event = tsl::MakeConstructedAsyncValueRef<std::unique_ptr<se::Event>>(
       std::move(host_to_device_stream_event));
 
-  absl::MutexLock lock(events_mu_);
-  auto [it, inserted] =
-      events_.emplace(std::make_pair(executor, run_id), event);
+  ABSL_ASSIGN_OR_RETURN(DeviceEvents * device_events,
+                   device_events_.GetOrCreate(executor->device_ordinal()));
+  auto [it, inserted] = device_events->events.emplace(run_id, event);
 
   if (!inserted) {
     return FailedPrecondition(
@@ -429,15 +431,20 @@ HostExecuteAsyncEvents::ExtractEvent(se::StreamExecutor* executor,
   VLOG(6) << "Extracting event for executor at address " << executor
           << " and event id " << run_id.ToInt();
 
-  absl::MutexLock lock(events_mu_);
-  auto it = events_.find(std::make_pair(executor, run_id));
-  if (it == events_.end()) {
+  DeviceEvents* device_events = device_events_.Find(executor->device_ordinal());
+  if (device_events == nullptr) {
+    return FailedPrecondition(
+        "Event does not exist for executor at address %p and event id %d",
+        executor, run_id.ToInt());
+  }
+  auto it = device_events->events.find(run_id);
+  if (it == device_events->events.end()) {
     return FailedPrecondition(
         "Event does not exist for executor at address %p and event id %d",
         executor, run_id.ToInt());
   }
   auto event = std::move(it->second);
-  events_.erase(it);
+  device_events->events.erase(it);
   return event;
 }
 
@@ -446,10 +453,10 @@ HostExecuteStartThunk::Create(
     Thunk::ThunkInfo thunk_info,
     const HostOffloadingExecutableProto& host_offloading_executable_proto,
     absl::InlinedVector<ShapedSlice, 4> args,
-    absl::InlinedVector<ShapedSlice, 4> results) {
+    absl::InlinedVector<ShapedSlice, 4> results, int devices_per_host) {
   auto thunk = std::make_unique<HostExecuteStartThunk>(
       std::move(thunk_info), host_offloading_executable_proto, std::move(args),
-      std::move(results));
+      std::move(results), /*async_events=*/nullptr, devices_per_host);
   if (host_offloading_executable_proto.has_aot_compilation_result()) {
     ABSL_RETURN_IF_ERROR(thunk->LoadExecutable());
   }
@@ -474,11 +481,12 @@ absl::Status HostExecuteStartThunk::LoadExecutable() {
 HostExecuteStartThunk::HostExecuteStartThunk(
     Thunk::ThunkInfo thunk_info, const HloModule& hlo_module,
     absl::InlinedVector<ShapedSlice, 4> args,
-    absl::InlinedVector<ShapedSlice, 4> results)
+    absl::InlinedVector<ShapedSlice, 4> results, int devices_per_host)
     : Command(Thunk::Kind::kHostExecuteStart, std::move(thunk_info)),
       args_(std::move(args)),
       results_(std::move(results)),
-      async_events_(std::make_shared<HostExecuteAsyncEvents>()) {
+      async_events_(
+          std::make_shared<HostExecuteAsyncEvents>(devices_per_host)) {
   HostOffloadingExecutableProto host_offloading_executable_proto;
   *host_offloading_executable_proto.mutable_hlo_module() = hlo_module.ToProto();
   host_offloading_executable_proto.set_executable_type(
@@ -491,13 +499,14 @@ HostExecuteStartThunk::HostExecuteStartThunk(
     const HostOffloadingExecutableProto& host_offloading_executable_proto,
     absl::InlinedVector<ShapedSlice, 4> args,
     absl::InlinedVector<ShapedSlice, 4> results,
-    std::shared_ptr<HostExecuteAsyncEvents> async_events)
+    std::shared_ptr<HostExecuteAsyncEvents> async_events, int devices_per_host)
     : Command(Thunk::Kind::kHostExecuteStart, std::move(thunk_info)),
       args_(std::move(args)),
       results_(std::move(results)),
       executable_proto_(host_offloading_executable_proto) {
   async_events_ =
-      async_events ? async_events : std::make_shared<HostExecuteAsyncEvents>();
+      async_events ? std::move(async_events)
+                   : std::make_shared<HostExecuteAsyncEvents>(devices_per_host);
 }
 
 std::string HostExecuteStartThunk::ToString(int indent) const { return ""; }
@@ -536,7 +545,7 @@ absl::StatusOr<std::unique_ptr<HostExecuteStartThunk>>
 HostExecuteStartThunk::FromProto(
     ThunkInfo thunk_info, const HostExecuteStartThunkProto& proto,
     absl::Span<const BufferAllocation> buffer_allocations,
-    HostExecuteAsyncEventsMap& async_events_map) {
+    HostExecuteAsyncEventsMap& async_events_map, int devices_per_host) {
   absl::InlinedVector<ShapedSlice, 4> args, results;
 
   for (const ShapedSliceProto& proto : proto.args()) {
@@ -553,9 +562,12 @@ HostExecuteStartThunk::FromProto(
   // If async_events_map already contains an entry for the given unique id,
   // that means that the pairing done thunk is already serialized and we reuse
   // the id to connect them. Otherwise, create a new entry.
-  auto [async_event_it, _] = async_events_map.try_emplace(
-      AsyncEventsUniqueId(proto.async_events_unique_id()),
-      std::make_shared<HostExecuteAsyncEvents>());
+  auto [async_event_it, inserted] = async_events_map.try_emplace(
+      AsyncEventsUniqueId(proto.async_events_unique_id()), nullptr);
+  if (inserted) {
+    async_event_it->second =
+        std::make_shared<HostExecuteAsyncEvents>(devices_per_host);
+  }
   return std::make_unique<HostExecuteStartThunk>(
       thunk_info, proto.executable_proto(), std::move(args), std::move(results),
       async_event_it->second);
@@ -793,13 +805,16 @@ absl::StatusOr<std::unique_ptr<HostExecuteDoneThunk>>
 HostExecuteDoneThunk::FromProto(
     ThunkInfo thunk_info, const HostExecuteDoneThunkProto& proto,
     absl::Span<const BufferAllocation> buffer_allocations,
-    HostExecuteAsyncEventsMap& async_events_map) {
+    HostExecuteAsyncEventsMap& async_events_map, int devices_per_host) {
   // If async_events_map already contains an entry for the given unique id,
   // that means that the pairing start thunk is already serialized and we reuse
   // the id to connect them. Otherwise, create a new entry.
-  auto [async_event_it, _] = async_events_map.try_emplace(
-      AsyncEventsUniqueId(proto.async_events_unique_id()),
-      std::make_shared<HostExecuteAsyncEvents>());
+  auto [async_event_it, inserted] = async_events_map.try_emplace(
+      AsyncEventsUniqueId(proto.async_events_unique_id()), nullptr);
+  if (inserted) {
+    async_event_it->second =
+        std::make_shared<HostExecuteAsyncEvents>(devices_per_host);
+  }
 
   absl::InlinedVector<ShapedSlice, 4> results;
   results.reserve(proto.results().size());
diff --git a/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk.h b/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk.h
index e88f0177590..b6ff631474d 100644
--- a/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk.h
+++ b/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk.h
@@ -19,18 +19,16 @@ limitations under the License.
 #include <memory>
 #include <optional>
 #include <string>
-#include <utility>
 
 #include "absl/base/call_once.h"
-#include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/inlined_vector.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
 #include "xla/backends/gpu/runtime/command.h"
 #include "xla/backends/gpu/runtime/host_async_thunk.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk.pb.h"
 #include "xla/core/host_offloading/host_offloading_allocator.h"
@@ -62,6 +60,9 @@ class HostExecuteAsyncEvents {
   // published results.
   using HostExecuteEvent = tsl::AsyncValueRef<std::unique_ptr<se::Event>>;
 
+  explicit HostExecuteAsyncEvents(int devices_per_host = 0)
+      : device_events_(devices_per_host) {}
+
   // Creates an event for the given executor and run id and returns it to the
   // user if the event was created successfully.
   absl::StatusOr<HostExecuteEvent> CreateEvent(se::StreamExecutor* executor,
@@ -71,9 +72,13 @@ class HostExecuteAsyncEvents {
                                                 RunId run_id);
 
  private:
-  absl::Mutex events_mu_;
-  absl::flat_hash_map<std::pair<se::StreamExecutor*, RunId>, HostExecuteEvent>
-      events_ ABSL_GUARDED_BY(events_mu_);
+  // Not internally synchronized. Relies on the caller serializing host-side
+  // thunk execution per device (e.g., PJRT per-device execute_thread).
+  struct DeviceEvents {
+    absl::flat_hash_map<RunId, HostExecuteEvent> events;
+  };
+
+  PerDeviceState<DeviceEvents> device_events_;
 };
 
 using HostExecuteAsyncEventsMap =
@@ -85,13 +90,14 @@ class HostExecuteStartThunk : public Command {
   HostExecuteStartThunk(Thunk::ThunkInfo thunk_info,
                         const HloModule& hlo_module,
                         absl::InlinedVector<ShapedSlice, 4> args,
-                        absl::InlinedVector<ShapedSlice, 4> results);
+                        absl::InlinedVector<ShapedSlice, 4> results,
+                        int devices_per_host = 0);
 
   static absl::StatusOr<std::unique_ptr<HostExecuteStartThunk>> Create(
       Thunk::ThunkInfo thunk_info,
       const HostOffloadingExecutableProto& host_offloading_executable_proto,
       absl::InlinedVector<ShapedSlice, 4> args,
-      absl::InlinedVector<ShapedSlice, 4> results);
+      absl::InlinedVector<ShapedSlice, 4> results, int devices_per_host = 0);
 
   HostExecuteStartThunk(const HostExecuteStartThunk&) = delete;
   HostExecuteStartThunk& operator=(const HostExecuteStartThunk&) = delete;
@@ -107,7 +113,7 @@ class HostExecuteStartThunk : public Command {
   static absl::StatusOr<std::unique_ptr<HostExecuteStartThunk>> FromProto(
       ThunkInfo thunk_info, const HostExecuteStartThunkProto& proto,
       absl::Span<const BufferAllocation> buffer_allocations,
-      HostExecuteAsyncEventsMap& async_events_map);
+      HostExecuteAsyncEventsMap& async_events_map, int devices_per_host = 0);
 
   absl::Status Initialize(const InitializeParams& params) override;
   absl::Status ExecuteOnStream(const ExecuteParams& params) override;
@@ -153,7 +159,8 @@ class HostExecuteStartThunk : public Command {
       const HostOffloadingExecutableProto& host_offloading_executable_proto,
       absl::InlinedVector<ShapedSlice, 4> args,
       absl::InlinedVector<ShapedSlice, 4> results,
-      std::shared_ptr<HostExecuteAsyncEvents> async_events = nullptr);
+      std::shared_ptr<HostExecuteAsyncEvents> async_events = nullptr,
+      int devices_per_host = 0);
 
  private:
   absl::once_flag executable_init_flag_;
@@ -180,7 +187,7 @@ class HostExecuteDoneThunk : public Command {
   static absl::StatusOr<std::unique_ptr<HostExecuteDoneThunk>> FromProto(
       ThunkInfo thunk_info, const HostExecuteDoneThunkProto& proto,
       absl::Span<const BufferAllocation> buffer_allocations,
-      HostExecuteAsyncEventsMap& async_events_map);
+      HostExecuteAsyncEventsMap& async_events_map, int devices_per_host = 0);
 
   absl::Status Initialize(const InitializeParams& params) override;
   absl::Status ExecuteOnStream(const ExecuteParams& params) override;
diff --git a/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk_test.cc
index 76422d203a0..eb36be79ddc 100644
--- a/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/host_execute_thunk_test.cc
@@ -51,7 +51,9 @@ limitations under the License.
 #include "xla/service/shaped_slice.h"
 #include "xla/shape_util.h"
 #include "xla/stream_executor/device_address.h"
+#include "xla/stream_executor/event.h"
 #include "xla/stream_executor/memory_space.h"
+#include "xla/stream_executor/mock_stream_executor.h"
 #include "xla/stream_executor/platform_manager.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/stream_executor/stream_executor_address_allocator.h"
@@ -708,6 +710,46 @@ TEST(HostExecuteDoneThunkTest, WaitingOnErrorEvent) {
               absl_testing::StatusIs(absl::StatusCode::kInternal));
 }
 
+TEST(HostExecuteAsyncEventsTest, CreateAndExtractEvents) {
+  se::MockStreamExecutor executor0;
+  se::MockStreamExecutor executor1;
+  EXPECT_CALL(executor0, device_ordinal()).WillRepeatedly(::testing::Return(0));
+  EXPECT_CALL(executor1, device_ordinal()).WillRepeatedly(::testing::Return(1));
+  EXPECT_CALL(executor0, CreateEvent()).WillRepeatedly([]() {
+    return std::unique_ptr<se::Event>(nullptr);
+  });
+  EXPECT_CALL(executor1, CreateEvent()).WillRepeatedly([]() {
+    return std::unique_ptr<se::Event>(nullptr);
+  });
+
+  HostExecuteAsyncEvents events(/*devices_per_host=*/2);
+
+  ASSERT_OK_AND_ASSIGN(auto event0_run1,
+                       events.CreateEvent(&executor0, RunId(1)));
+  ASSERT_OK_AND_ASSIGN(auto event1_run1,
+                       events.CreateEvent(&executor1, RunId(1)));
+  ASSERT_OK_AND_ASSIGN(auto event0_run2,
+                       events.CreateEvent(&executor0, RunId(2)));
+  EXPECT_THAT(events.CreateEvent(&executor0, RunId(1)),
+              absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition));
+  EXPECT_THAT(events.ExtractEvent(&executor1, RunId(2)),
+              absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition));
+
+  ASSERT_OK_AND_ASSIGN(auto extracted0_run1,
+                       events.ExtractEvent(&executor0, RunId(1)));
+  EXPECT_EQ(extracted0_run1.GetAsyncValue(), event0_run1.GetAsyncValue());
+  EXPECT_THAT(events.ExtractEvent(&executor0, RunId(1)),
+              absl_testing::StatusIs(absl::StatusCode::kFailedPrecondition));
+
+  ASSERT_OK_AND_ASSIGN(auto extracted1_run1,
+                       events.ExtractEvent(&executor1, RunId(1)));
+  EXPECT_EQ(extracted1_run1.GetAsyncValue(), event1_run1.GetAsyncValue());
+
+  ASSERT_OK_AND_ASSIGN(auto extracted0_run2,
+                       events.ExtractEvent(&executor0, RunId(2)));
+  EXPECT_EQ(extracted0_run2.GetAsyncValue(), event0_run2.GetAsyncValue());
+}
+
 TEST(HostExecuteStartThunkTest, ProtoRoundTrip) {
   static constexpr char const* kHloModule = R"(
     HloModule module
diff --git a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
index 55117838f60..e6b88fdfa4e 100644
--- a/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/thunk_proto_deserialization.cc
@@ -241,11 +241,13 @@ absl::StatusOr<std::unique_ptr<Thunk>> DeserializeThunkProtoImpl(
     case ThunkProto::kHostExecuteStartThunk:
       return HostExecuteStartThunk::FromProto(
           std::move(thunk_info), thunk_proto.host_execute_start_thunk(),
-          buffer_allocations, host_executable_async_events_map);
+          buffer_allocations, host_executable_async_events_map,
+          gpu_topology.has_value() ? gpu_topology->num_devices_per_host() : 0);
     case ThunkProto::kHostExecuteDoneThunk:
       return HostExecuteDoneThunk::FromProto(
           std::move(thunk_info), thunk_proto.host_execute_done_thunk(),
-          buffer_allocations, host_executable_async_events_map);
+          buffer_allocations, host_executable_async_events_map,
+          gpu_topology.has_value() ? gpu_topology->num_devices_per_host() : 0);
     // TODO: Remove this case on Apr 30, 2027
     case ThunkProto::kSelectKThunk:
       return DeserializeSelectKThunkProto(
diff --git a/third_party/xla/xla/service/gpu/thunk_emitter.cc b/third_party/xla/xla/service/gpu/thunk_emitter.cc
index c75773a320d..5fb44510496 100644
--- a/third_party/xla/xla/service/gpu/thunk_emitter.cc
+++ b/third_party/xla/xla/service/gpu/thunk_emitter.cc
@@ -2675,12 +2675,14 @@ absl::StatusOr<ThunkSequence> ThunkEmitter::EmitHostExecuteStart(
   host_offloading_executable_proto.set_executable_type(
       HostOffloadingExecutableProto::EXECUTABLE_TYPE_NANORT);
 
-  ABSL_ASSIGN_OR_RETURN(auto thunk,
-                   HostExecuteStartThunk::Create(
-                       Thunk::ThunkInfo::WithProfileAnnotation(
-                           async_start, ir_emitter_context_->GetNextThunkId()),
-                       std::move(host_offloading_executable_proto),
-                       std::move(operand_slices), std::move(result_slices)));
+  ABSL_ASSIGN_OR_RETURN(
+      auto thunk,
+      HostExecuteStartThunk::Create(
+          Thunk::ThunkInfo::WithProfileAnnotation(
+              async_start, ir_emitter_context_->GetNextThunkId()),
+          std::move(host_offloading_executable_proto),
+          std::move(operand_slices), std::move(result_slices),
+          ir_emitter_context_->gpu_topology().num_devices_per_host()));
 
   auto [it, inserted] = GetInstructionToHostExecuteAsyncEvents().emplace(
       host_execute, thunk->async_events());

From 1565922a1d5c1b957fb3a6adbfa0b9834295ae4f Mon Sep 17 00:00:00 2001
From: Enver Kayaaslan <enver@google.com>
Date: Mon, 12 Oct 2026 02:39:04 -0700
Subject: [PATCH 0565/1000] Add unit tests for non-flat support on
 AllGatherBroadcastReorder.

PiperOrigin-RevId: 997787308
---
 .../xla/xla/hlo/transforms/collectives/BUILD  |  2 +
 .../all_gather_broadcast_reorder_test.cc      | 89 +++++++++++++++++++
 2 files changed, 91 insertions(+)

diff --git a/third_party/xla/xla/hlo/transforms/collectives/BUILD b/third_party/xla/xla/hlo/transforms/collectives/BUILD
index 5692dd75556..0011e156f6d 100644
--- a/third_party/xla/xla/hlo/transforms/collectives/BUILD
+++ b/third_party/xla/xla/hlo/transforms/collectives/BUILD
@@ -267,8 +267,10 @@ xla_cc_test(
     srcs = ["all_gather_broadcast_reorder_test.cc"],
     deps = [
         ":all_gather_broadcast_reorder",
+        "//xla/hlo/testlib:filecheck",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
         "//xla/hlo/utils:hlo_matchers",
+        "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest_main",
         "@tsl//tsl/platform:statusor",
diff --git a/third_party/xla/xla/hlo/transforms/collectives/all_gather_broadcast_reorder_test.cc b/third_party/xla/xla/hlo/transforms/collectives/all_gather_broadcast_reorder_test.cc
index 6e4a14b3a1c..fc7921c4938 100644
--- a/third_party/xla/xla/hlo/transforms/collectives/all_gather_broadcast_reorder_test.cc
+++ b/third_party/xla/xla/hlo/transforms/collectives/all_gather_broadcast_reorder_test.cc
@@ -17,7 +17,9 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/status/status_matchers.h"
 #include "absl/strings/string_view.h"
+#include "xla/hlo/testlib/filecheck.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/hlo/utils/hlo_matchers.h"
 #include "tsl/platform/statusor.h"
@@ -25,6 +27,8 @@ limitations under the License.
 namespace xla {
 namespace {
 
+using ::absl_testing::IsOkAndHolds;
+
 namespace m = xla::testing::opcode_matchers;
 
 class AllGatherBroadcastReorderTest : public HloHardwareIndependentTestBase {
@@ -118,5 +122,90 @@ ENTRY main {
   RunPass(hlo_string, PassOutput::NoChange);
 }
 
+TEST_F(AllGatherBroadcastReorderTest, DoesNotCrossCallBoundaries) {
+  absl::string_view hlo_string = R"(
+HloModule m
+
+foo {
+  param = f32[5, 4, 8, 128] parameter(0)
+  ROOT ag = f32[5, 4, 8, 256] all-gather(param), dimensions={3}, replica_groups={{0, 1}}
+}
+
+ENTRY main {
+  x = f32[128, 5] parameter(0)
+  bc = f32[5, 4, 8, 128] broadcast(x), dimensions={3, 0}
+  ROOT call = f32[5, 4, 8, 256] call(bc), to_apply=foo
+}
+)";
+  RunPass(hlo_string, PassOutput::NoChange);
+}
+
+TEST_F(AllGatherBroadcastReorderTest, EntryCallsFooTransformationOnFoo) {
+  absl::string_view hlo_string = R"(
+// CHECK: %foo (
+// CHECK: %[[AG:.*]] = f32[256,5]{{.*}} all-gather(%x_foo), {{.*}}dimensions={0}
+// CHECK: ROOT %[[BC:.*]] = f32[5,4,8,256]{{.*}} broadcast(%[[AG]]), dimensions={3,0}
+// CHECK: }
+// CHECK: ENTRY %main (
+// CHECK: ROOT %[[CALL:.*]] = f32[5,4,8,256]{{.*}} call(%x_main), to_apply=%foo
+HloModule m
+
+foo {
+  x_foo = f32[128, 5] parameter(0)
+  bc = f32[5, 4, 8, 128] broadcast(x_foo), dimensions={3, 0}
+  ROOT ag = f32[5, 4, 8, 256] all-gather(bc), dimensions={3}, replica_groups={{0, 1}}
+}
+
+ENTRY main {
+  x_main = f32[128, 5] parameter(0)
+  ROOT call = f32[5, 4, 8, 256] call(x_main), to_apply=foo
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_THAT(AllGatherBroadcastReorder().Run(module.get()),
+              IsOkAndHolds(true));
+  EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true));
+}
+
+TEST_F(AllGatherBroadcastReorderTest,
+       EntryCallsFooAndBarFooCallsBarTransformationOnBar) {
+  absl::string_view hlo_string = R"(
+// CHECK: %bar (
+// CHECK: %[[AG:.*]] = f32[256,5]{{.*}} all-gather(%x_bar), {{.*}}dimensions={0}
+// CHECK: ROOT %[[BC:.*]] = f32[5,4,8,256]{{.*}} broadcast(%[[AG]]), dimensions={3,0}
+// CHECK: }
+// CHECK: %foo (
+// CHECK: ROOT %[[CALL_BAR:.*]] = f32[5,4,8,256]{{.*}} call(%x_foo), to_apply=%bar
+// CHECK: }
+// CHECK: ENTRY %main (
+// CHECK: %[[CALL_FOO:.*]] = f32[5,4,8,256]{{.*}} call(%x_main), to_apply=%foo
+// CHECK: %[[CALL_BAR2:.*]] = f32[5,4,8,256]{{.*}} call(%x_main), to_apply=%bar
+// CHECK: ROOT %[[ADD:.*]] = f32[5,4,8,256]{{.*}} add(%[[CALL_FOO]], %[[CALL_BAR2]])
+HloModule m
+
+bar {
+  x_bar = f32[128, 5] parameter(0)
+  bc = f32[5, 4, 8, 128] broadcast(x_bar), dimensions={3, 0}
+  ROOT ag = f32[5, 4, 8, 256] all-gather(bc), dimensions={3}, replica_groups={{0, 1}}
+}
+
+foo {
+  x_foo = f32[128, 5] parameter(0)
+  ROOT call_bar = f32[5, 4, 8, 256] call(x_foo), to_apply=bar
+}
+
+ENTRY main {
+  x_main = f32[128, 5] parameter(0)
+  call_foo = f32[5, 4, 8, 256] call(x_main), to_apply=foo
+  call_bar = f32[5, 4, 8, 256] call(x_main), to_apply=bar
+  ROOT add = f32[5, 4, 8, 256] add(call_foo, call_bar)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_THAT(AllGatherBroadcastReorder().Run(module.get()),
+              IsOkAndHolds(true));
+  EXPECT_THAT(RunFileCheck(module->ToString(), hlo_string), IsOkAndHolds(true));
+}
+
 }  // namespace
 }  // namespace xla

From 16a1ba92ad57e0a6bef393d216f102dc2b4c619e Mon Sep 17 00:00:00 2001
From: "Sevin F. Varoglu" <svaroglu@nvidia.com>
Date: Mon, 12 Oct 2026 03:13:05 -0700
Subject: [PATCH 0566/1000] PR #50394: [XLA:GPU] Treat stream-annotated
 collective calls as opaque in LHS
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50394

📝 Summary of Changes
Extend the LHS "opaque launch" resource-counting rule from _collectives_group="" wrappers to also cover _xla_stream_annotation="collective" wrappers. One-line condition change in GpuAsyncTracker::GetNumResourcesPerInstruction.

🎯 Justification
Without the fix, LHS double-counts a stream-annotated wrapper as two collective-stream users (itself + its nested collective). Under tight overlap limits this makes the schedule infeasible and compilation fails.

🚀 Kind of Contribution
🐛 Bug Fix, 🧪 Tests

📊 Benchmark (for Performance Improvements)
N/A

🧪 Unit Tests:
GpuLatencyHidingSchedulerBaseTest.CollectiveStreamAnnotatedCallUsesOneResource
GpuLatencyHidingSchedulerBaseTest.CollectiveStreamAnnotatedCallOverlapsNativeCollective

🧪 Execution Tests:
N/A
Copybara import of the project:

--
27a0e3d01751924be711a955af984a4315194238 by Sevin F. Varoglu <svaroglu@nvidia.com>:

[XLA:GPU] Treat stream-annotated collective calls as opaque in LHS

Merging this change closes #50394

PiperOrigin-RevId: 997801829
---
 .../gpu/gpu_latency_hiding_scheduler.cc       |  2 +-
 .../gpu/gpu_latency_hiding_scheduler_test.cc  | 79 +++++++++++++++++++
 2 files changed, 80 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/gpu/gpu_latency_hiding_scheduler.cc b/third_party/xla/xla/service/gpu/gpu_latency_hiding_scheduler.cc
index 96ce2cb5fbe..85785c8a602 100644
--- a/third_party/xla/xla/service/gpu/gpu_latency_hiding_scheduler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_latency_hiding_scheduler.cc
@@ -746,7 +746,7 @@ ResourcesVector GpuAsyncTracker::GetResourcesFromInstructionImpl(
 absl::flat_hash_map<int64_t, int64_t>
 GpuAsyncTracker::GetNumResourcesPerInstruction(
     const HloInstruction& instr) const {
-  if (!IsCollectivesGroupAsyncOp(instr)) {
+  if (!IsCollectivesGroupAsyncOp(instr) && !IsCustomCollectiveOp(&instr)) {
     return GpuAsyncTrackerBase::GetNumResourcesPerInstruction(instr);
   }
 
diff --git a/third_party/xla/xla/service/gpu/gpu_latency_hiding_scheduler_test.cc b/third_party/xla/xla/service/gpu/gpu_latency_hiding_scheduler_test.cc
index 54fcecdfdc5..b30f3efec33 100644
--- a/third_party/xla/xla/service/gpu/gpu_latency_hiding_scheduler_test.cc
+++ b/third_party/xla/xla/service/gpu/gpu_latency_hiding_scheduler_test.cc
@@ -1297,6 +1297,55 @@ ENTRY main {
   }
 }
 
+TEST_F(GpuLatencyHidingSchedulerBaseTest,
+       CollectiveStreamAnnotatedCallUsesOneResource) {
+  constexpr absl::string_view kHloModule = R"(
+HloModule test, is_scheduled=true
+
+body {
+  p0 = f32[1] parameter(0)
+  ag-start = (f32[1], f32[1]) all-gather-start(p0), dimensions={0}
+  ROOT ag-done = f32[1] all-gather-done(ag-start)
+}
+
+ENTRY main {
+  p0 = f32[1] parameter(0)
+  call-start = ((f32[1]), f32[1]) call-start(p0), to_apply=body,
+    frontend_attributes={_xla_stream_annotation="collective"}
+  ROOT call-done = f32[1] call-done(call-start),
+    frontend_attributes={_xla_stream_annotation="collective"}
+})";
+
+  HloModuleConfig config = GetModuleConfig(/*fdo_profile=*/"");
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kHloModule, config));
+  HloInstruction* call_start = FindInstruction(module.get(), "call-start");
+  HloInstruction* call_done = FindInstruction(module.get(), "call-done");
+  ASSERT_NE(call_start, nullptr);
+  ASSERT_NE(call_done, nullptr);
+
+  SchedulerConfig sched_config;
+  GpuAsyncTracker async_tracker(sched_config);
+  const int64_t collective_resource =
+      ResourceTypeToIndex(GpuResourceType::kGpuAsyncStreamCollectives);
+
+  EXPECT_EQ(async_tracker.GetNumResourcesPerInstruction(collective_resource,
+                                                        *call_start),
+            0);
+  EXPECT_EQ(async_tracker.GetNumResourcesPerInstruction(collective_resource,
+                                                        *call_done),
+            1);
+
+  auto start_resource_counts =
+      async_tracker.GetNumResourcesPerInstruction(*call_start);
+  EXPECT_TRUE(start_resource_counts.empty());
+
+  auto done_resource_counts =
+      async_tracker.GetNumResourcesPerInstruction(*call_done);
+  ASSERT_EQ(done_resource_counts.size(), 1);
+  EXPECT_EQ(done_resource_counts.at(collective_resource), 1);
+}
+
 TEST_F(GpuLatencyHidingSchedulerBaseTest,
        ExplicitCollectivesGroupSerializesWithNativeCollective) {
   absl::string_view kFdoProfile = R"pb(
@@ -2321,6 +2370,36 @@ ENTRY main {
               GetIndexByName(instruction_sequence, "ar_0"));
 }
 
+TEST_F(GpuLatencyHidingSchedulerBaseTest,
+       CollectiveStreamAnnotatedCallOverlapsNativeCollective) {
+  absl::string_view kHloModule = R"(
+HloModule m
+
+body {
+  p0 = f32[2] parameter(0)
+  ag-start = (f32[2], f32[2]) all-gather-start(p0), dimensions={0}
+  ROOT ag-done = f32[2] all-gather-done(ag-start)
+}
+
+ENTRY main {
+  p0 = f32[2] parameter(0)
+  native-start = (f32[2], f32[2]) all-gather-start(p0), dimensions={0}
+  call-start = ((f32[2]), f32[2]) call-start(p0), to_apply=body,
+    frontend_attributes={_xla_stream_annotation="collective"}
+  call-done = f32[2] call-done(call-start),
+    frontend_attributes={_xla_stream_annotation="collective"}
+  native-done = f32[2] all-gather-done(native-start)
+  ROOT result = (f32[2], f32[2]) tuple(call-done, native-done)
+}
+)";
+
+  auto config = GetModuleConfig(/*fdo_profile=*/"");
+  ASSERT_OK_AND_ASSIGN(auto module,
+                       ParseAndReturnVerifiedModule(kHloModule, config));
+
+  EXPECT_OK(ScheduleModule(module.get(), /*num_parallel_resources=*/1));
+}
+
 TEST_F(GpuLatencyHidingSchedulerBaseTest, LargeIndependentFillsAreForcedEarly) {
   constexpr absl::string_view kHloModule = R"(
 HloModule test

From 14d068a63e67cf73a4b94216c0b1cfafe6564384 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 12 Oct 2026 03:15:16 -0700
Subject: [PATCH 0567/1000] [XLA:TPU] Avoid repeated DCEs in
 `FindAndWrapOffloadedComputations`.

In every iteration in `FindAndWrapOffloadedComputations`,  there was a new `HloDCE().Run(computation.parent())` across the entire `HloModule`, which was very expensive.  In particular, in every iteration this repeatedly rebuilt the full computation post-order and `unmerged_ancestors` via `std::list`

This CL
1. pre-filters offload candidates once,
2. tracks `unmerged_ancestors` incrementally in an `absl::flat_hash_set` via a worklist,
3. eagerly removes dead fused instructions without control successors during wrapping, and
4. runs a single `HloDCE::RunOnComputation(&computation)` and `RemoveUnusedComputations()` at the end.

PiperOrigin-RevId: 997802624
---
 third_party/xla/xla/hlo/transforms/BUILD      |   2 -
 .../offloaded_instruction_wrapper.cc          | 242 ++++++++++--------
 2 files changed, 131 insertions(+), 113 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/BUILD b/third_party/xla/xla/hlo/transforms/BUILD
index 2fdbd339ffb..05cddecc1b1 100644
--- a/third_party/xla/xla/hlo/transforms/BUILD
+++ b/third_party/xla/xla/hlo/transforms/BUILD
@@ -623,7 +623,6 @@ cc_library(
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/transforms/simplifiers:hlo_dce",
-        "//xla/tsl/platform:errors",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base",
         "@com_google_absl//absl/container:flat_hash_map",
@@ -635,7 +634,6 @@ cc_library(
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:string_view",
-        "@tsl//tsl/platform:casts",
     ],
 )
 
diff --git a/third_party/xla/xla/hlo/transforms/offloaded_instruction_wrapper.cc b/third_party/xla/xla/hlo/transforms/offloaded_instruction_wrapper.cc
index 3083c8a6ac1..ed5145ea7c5 100644
--- a/third_party/xla/xla/hlo/transforms/offloaded_instruction_wrapper.cc
+++ b/third_party/xla/xla/hlo/transforms/offloaded_instruction_wrapper.cc
@@ -15,6 +15,8 @@ limitations under the License.
 
 #include "xla/hlo/transforms/offloaded_instruction_wrapper.h"
 
+#include <algorithm>
+#include <cstddef>
 #include <cstdint>
 #include <utility>
 #include <vector>
@@ -35,10 +37,8 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/transforms/simplifiers/hlo_dce.h"
 #include "xla/side_effect_util.h"
-#include "xla/tsl/platform/errors.h"
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
-#include "tsl/platform/casts.h"
 
 namespace xla::offloader_util {
 
@@ -90,138 +90,159 @@ FindAndWrapOffloadedComputations(
     }
   }
 
+  auto collect_offload_candidates = [&]() {
+    std::vector<HloInstruction*> post_order =
+        computation.MakeInstructionPostOrder();
+    std::vector<HloInstruction*> candidates;
+    for (auto it = post_order.rbegin(); it != post_order.rend(); ++it) {
+      if (should_offload(*it)) {
+        candidates.push_back(*it);
+      }
+    }
+    return candidates;
+  };
+
+  std::vector<HloInstruction*> offload_candidates =
+      collect_offload_candidates();
+  size_t candidate_idx = 0;
+
   std::vector<std::pair<HloInstruction*, int64_t>>
       offloaded_instructions_and_calls;
   // On each iteration of the outer loop, try to create one offloaded
   // computation out of a connected set of offloaded instructions.
-  while (true) {
-    HloInstruction* offloaded_instr = nullptr;
-    HloCallInstruction* offloaded_call_instr = nullptr;
-    // Stores all the ancestor instructions of offloaded_call_instr which were
-    // not added to the current offloaded computation.
-    absl::flat_hash_set<HloInstruction*> unmerged_ancestors;
-    std::vector<HloInstruction*> post_order =
-        computation.MakeInstructionPostOrder();
-    for (auto it = post_order.rbegin(); it != post_order.rend(); ++it) {
-      HloInstruction* instr = *it;
-      // The current instruction is not an offload instruction.
-      if (!should_offload(instr)) {
-        if (absl::c_any_of(instr->users(), [&](const HloInstruction* user) {
-              return user == offloaded_call_instr ||
-                     unmerged_ancestors.contains(user);
-            })) {
-          unmerged_ancestors.insert(instr);
+  while (candidate_idx < offload_candidates.size()) {
+    HloInstruction* root_instr = offload_candidates[candidate_idx++];
+    if (root_instr->parent() == nullptr || !should_offload(root_instr)) {
+      continue;
+    }
+
+    VLOG(2) << "Offloading instruction: " << root_instr->ToString();
+    VLOG(2) << root_instr->name()
+            << " is the root of a new offloaded computation";
+
+    HloInstruction* call_instr;
+    if (root_instr->opcode() == HloOpcode::kCall) {
+      call_instr = root_instr;
+    } else {
+      call_instr = computation.CreateCallInstruction({root_instr});
+      call_instr->SetAndSanitizeName(new_call_name_prefix);
+      call_instr->UniquifyName(computation.parent());
+      call_instr->set_frontend_attributes(root_instr->frontend_attributes());
+    }
+    HloCallInstruction* offloaded_call_instr =
+        absl::down_cast<HloCallInstruction*>(call_instr);
+    CHECK_NE(offloaded_call_instr, nullptr);
+    ABSL_RETURN_IF_ERROR(clear_backend_config_device_type(offloaded_call_instr));
+    ABSL_RETURN_IF_ERROR(ClearComputeTypeFrontendAttribute(offloaded_call_instr));
+    ClearSideEffects(root_instr);
+    HloInstruction* offloaded_instr = root_instr;
+
+    const bool can_fuse_any = std::any_of(
+        offload_candidates.begin() + candidate_idx, offload_candidates.end(),
+        [&](HloInstruction* cand) {
+          return cand->parent() != nullptr && should_offload(cand) &&
+                 should_fuse(*offloaded_call_instr, *cand);
+        });
+
+    bool fused = false;
+    bool has_unmerged_offload_instr = true;
+    if (can_fuse_any) {
+      // Stores all the ancestor instructions of offloaded_call_instr which were
+      // not added to the current offloaded computation.
+      absl::flat_hash_set<HloInstruction*> unmerged_ancestors;
+      auto add_unmerged_ancestor = [&](HloInstruction* ancestor) {
+        if (!unmerged_ancestors.insert(ancestor).second) {
+          return;
         }
-        continue;
-      }
-
-      VLOG(2) << "Offloading instruction: " << instr->ToString();
-
-      // The current instruction is the root of a new offloaded computation.
-      if (offloaded_call_instr == nullptr) {
-        VLOG(2) << instr->name()
-                << " is the root of a new offloaded computation";
-
-        HloInstruction* call_instr;
-        if (instr->opcode() == HloOpcode::kCall) {
-          call_instr = instr;
-        } else {
-          call_instr = computation.CreateCallInstruction({instr});
-          call_instr->SetAndSanitizeName(new_call_name_prefix);
-          call_instr->UniquifyName(computation.parent());
-          call_instr->set_frontend_attributes(instr->frontend_attributes());
+        std::vector<HloInstruction*> worklist = {ancestor};
+        while (!worklist.empty()) {
+          HloInstruction* curr = worklist.back();
+          worklist.pop_back();
+          for (HloInstruction* op : curr->operands()) {
+            if (unmerged_ancestors.insert(op).second) {
+              worklist.push_back(op);
+            }
+          }
         }
-        offloaded_call_instr = absl::down_cast<HloCallInstruction*>(call_instr);
-        CHECK_NE(offloaded_call_instr, nullptr);
-        ABSL_RETURN_IF_ERROR(clear_backend_config_device_type(offloaded_call_instr));
-        ABSL_RETURN_IF_ERROR(
-            ClearComputeTypeFrontendAttribute(offloaded_call_instr));
-        ClearSideEffects(instr);
-        offloaded_instr = instr;
-        continue;
-      }
+      };
+      auto add_non_offload_call_operands = [&]() {
+        for (HloInstruction* op : offloaded_call_instr->operands()) {
+          if (!should_offload(op)) {
+            add_unmerged_ancestor(op);
+          }
+        }
+      };
+      add_non_offload_call_operands();
 
-      // If the current instruction is indirectly connected to the current
-      // offloaded computation, it must go in a separate offload computation.
-      if (absl::c_any_of(instr->users(), [&](const HloInstruction* user) {
-            return unmerged_ancestors.contains(user);
-          })) {
-        VLOG(2) << instr->name()
-                << " is indirectly connected to the current offloaded "
-                   "computation, it must go in a separate offload computation";
+      has_unmerged_offload_instr = false;
+      for (size_t i = candidate_idx; i < offload_candidates.size(); ++i) {
+        HloInstruction* instr = offload_candidates[i];
+        if (instr->parent() == nullptr || !should_offload(instr)) {
+          continue;
+        }
 
-        unmerged_ancestors.insert(instr);
-        continue;
-      }
-
-      // The current instruction is directly connected to the current offloaded
-      // computation.
-      if (offloaded_call_instr->IsUserOf(instr)) {
-        VLOG(2)
-            << instr->name()
-            << " is directly connected to the current offloaded computation";
-        if (should_fuse(*offloaded_call_instr, *instr)) {
-          bool instr_escapes_offloaded_computation =
+        if (!unmerged_ancestors.contains(instr) &&
+            should_fuse(*offloaded_call_instr, *instr)) {
+          bool add_output =
+              !offloaded_call_instr->IsUserOf(instr) ||
               absl::c_any_of(instr->users(), [&](const HloInstruction* user) {
                 return user != offloaded_call_instr && !should_offload(user);
               });
-
-          VLOG(3) << instr->name() << " fusing into existing computation";
-          VLOG(6) << "instr_escapes_offloaded_computation: "
-                  << instr_escapes_offloaded_computation;
-
           offloaded_call_instr->AppendInstructionIntoCalledComputation(
-              instr, /*add_output=*/instr_escapes_offloaded_computation);
+              instr, add_output);
           ClearSideEffects(instr);
+          if (instr->IsDead() && !instr->HasSuccessorControlDependencies()) {
+            ABSL_RETURN_IF_ERROR(instr->SafelyDropAllControlDependencies());
+            ABSL_RETURN_IF_ERROR(computation.RemoveInstruction(instr));
+          }
           offloaded_instr = instr;
+          fused = true;
+          add_non_offload_call_operands();
         } else {
-          unmerged_ancestors.insert(instr);
+          add_unmerged_ancestor(instr);
+          has_unmerged_offload_instr = true;
         }
-        continue;
-      }
-
-      if (should_fuse(*offloaded_call_instr, *instr)) {
-        VLOG(2) << instr->ToString()
-                << " current instruction is disconnected from the current "
-                   "offload computation";
-        // If the current instruction is disconnected from the current offload
-        // computation, include it anyway.
-        offloaded_call_instr->AppendInstructionIntoCalledComputation(
-            instr, /*add_output=*/true);
-        ClearSideEffects(instr);
-        offloaded_instr = instr;
-      } else {
-        unmerged_ancestors.insert(instr);
       }
     }
 
-    // If there are no offload instructions left in the computation, stop
-    // looking.
-    if (offloaded_call_instr == nullptr) {
+    offloaded_instructions_and_calls.push_back(
+        {offloaded_instr, offloaded_call_instr->unique_id()});
+    if (!has_unmerged_offload_instr) {
       break;
     }
-    offloaded_instructions_and_calls.push_back(
-        std::pair(offloaded_instr, offloaded_call_instr->unique_id()));
-
-    for (HloInstruction* instr : computation.instructions()) {
-      // If an offloaded instruction is a Sharding custom call or has control
-      // dependencies (such as those around elided copies), remove it
-      // explicitly since it won't be removed by HloDCE.
-      if (instr->IsDead() && (instr->IsCustomCall("Sharding") ||
-                              (instr->HasControlDependencies() &&
-                               !instr->HasSuccessorControlDependencies()))) {
-        ABSL_RETURN_IF_ERROR(instr->SafelyDropAllControlDependencies());
-        ABSL_RETURN_IF_ERROR(computation.RemoveInstruction(instr));
-      }
+    if (fused) {
+      offload_candidates = collect_offload_candidates();
+      candidate_idx = 0;
     }
-
-    // DCE any offloaded instructions that have no remaining un-wrapped uses.
-    ABSL_RETURN_IF_ERROR(HloDCE().Run(computation.parent()).status());
-
-    VLOG(6) << "After offloading computation after DCE:";
-    XLA_VLOG_LINES(6, computation.parent()->ToString());
   }
 
+  std::vector<std::pair<HloInstruction*, HloCallInstruction*>> alive_calls;
+  if (offloaded_instructions_and_calls.empty()) {
+    return alive_calls;
+  }
+
+  for (HloInstruction* instr : computation.instructions()) {
+    // If an offloaded instruction is a Sharding custom call or has control
+    // dependencies (such as those around elided copies), remove it
+    // explicitly since it won't be removed by HloDCE.
+    if (instr->IsDead() && (instr->IsCustomCall("Sharding") ||
+                            (instr->HasControlDependencies() &&
+                             !instr->HasSuccessorControlDependencies()))) {
+      ABSL_RETURN_IF_ERROR(instr->SafelyDropAllControlDependencies());
+      ABSL_RETURN_IF_ERROR(computation.RemoveInstruction(instr));
+    }
+  }
+
+  // DCE any offloaded instructions that have no remaining un-wrapped uses.
+  ABSL_RETURN_IF_ERROR(HloDCE::RunOnComputation(&computation).status());
+  ABSL_RETURN_IF_ERROR(computation.parent()->RemoveUnusedComputations());
+  if (computation.parent()->has_schedule()) {
+    ABSL_RETURN_IF_ERROR(computation.parent()->schedule().Update());
+  }
+
+  VLOG(6) << "After offloading computation after DCE:";
+  XLA_VLOG_LINES(6, computation.parent()->ToString());
+
   // Filter out any offloaded call instructions that were removed by HloDCE.
   // Since HloDCE deletes instructions and subsequent iterations of the while
   // loop may allocate new ones reusing the same pointer address, we must use
@@ -231,7 +252,6 @@ FindAndWrapOffloadedComputations(
     alive_instructions[instr->unique_id()] = instr;
   }
 
-  std::vector<std::pair<HloInstruction*, HloCallInstruction*>> alive_calls;
   alive_calls.reserve(offloaded_instructions_and_calls.size());
   for (const auto& pair : offloaded_instructions_and_calls) {
     if (auto it = alive_instructions.find(pair.second);

From ce44a13880121547dd6c38b57188d09134e48713 Mon Sep 17 00:00:00 2001
From: Mohammed Anany <manany@google.com>
Date: Mon, 12 Oct 2026 03:21:55 -0700
Subject: [PATCH 0568/1000] [XLA:GPU][Bug Fix] Pass `std::shared_ptr<const
 BaseTilingSpace>` by value to `EvaluateCandidate` via `std::move(base)` so
 all worker references to `BaseTilingSpace` are destroyed when
 `EvaluateCandidate` returns, before `promise.Set(...)` completes the
 candidate futures.

Verified by running flaky test for 2000 iterations.

PiperOrigin-RevId: 997805524
---
 .../model/gpu_indexing_performance_model.cc   | 30 ++++++++++---------
 1 file changed, 16 insertions(+), 14 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc
index f065f03dcce..ef9ae0d2493 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc
+++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc
@@ -691,7 +691,12 @@ FilterAndSortCandidates(
   return candidates;
 }
 
-// Evaluates a single tiling candidate: clones `base_tiling_space`, assigns
+struct BaseTilingSpace {
+  PooledOrFallbackMlirContext context;
+  std::unique_ptr<experimental::TilingSpace> tiling_space;
+};
+
+// Evaluates a single tiling candidate: clones `base->tiling_space`, assigns
 // `padded_tile_sizes` to the clone, tiles the fusion and estimates its run
 // time. Returns std::nullopt if the candidate cannot be tiled or if it
 // violates Triton constraints.
@@ -701,7 +706,7 @@ FilterAndSortCandidates(
 // single mlir::MLIRContext.
 absl::StatusOr<std::optional<TiledRunTimeData>> EvaluateCandidate(
     const HloFusionAdaptor& fusion_adaptor,
-    const experimental::TilingSpace& base_tiling_space,
+    std::shared_ptr<const BaseTilingSpace> base,
     absl::Span<const int64_t> padded_tile_sizes,
     MlirContextPool* mlir_context_pool,
     const se::DeviceDescription& device_info,
@@ -709,11 +714,11 @@ absl::StatusOr<std::optional<TiledRunTimeData>> EvaluateCandidate(
     absl::FunctionRef<int64_t(const HloInstruction*)> flops_fn) {
   // Cloning is faster than calling TilingSpace::Create() for each tiling
   // candidate.
-  ABSL_ASSIGN_OR_RETURN(
-      PooledOrFallbackMlirContext context,
-      BorrowMlirContextOr(mlir_context_pool, base_tiling_space.mlir_context()));
+  ABSL_ASSIGN_OR_RETURN(PooledOrFallbackMlirContext context,
+                   BorrowMlirContextOr(mlir_context_pool,
+                                       base->tiling_space->mlir_context()));
   std::unique_ptr<experimental::TilingSpace> tiling_space =
-      base_tiling_space.Clone(context.get());
+      base->tiling_space->Clone(context.get());
   ABSL_RETURN_IF_ERROR(tiling_space->AssignTileSizes(padded_tile_sizes));
   VLOG(3) << "Trying tile sizes " << absl::StrJoin(padded_tile_sizes, ",");
 
@@ -765,10 +770,6 @@ tsl::Future<TopKTiledRunTimeDataOrError> TryFindTopKBestTilingsWithTilingSpace(
   // single-threaded `mlir_context`. Candidates clone from the base TilingSpace
   // asynchronously, so the context stays borrowed until the last candidate
   // releases `base`.
-  struct BaseTilingSpace {
-    PooledOrFallbackMlirContext context;
-    std::unique_ptr<experimental::TilingSpace> tiling_space;
-  };
   ABSL_ASSIGN_OR_RETURN(PooledOrFallbackMlirContext context,
                    BorrowMlirContextOr(mlir_context_pool, mlir_context));
   auto base = std::make_shared<BaseTilingSpace>(
@@ -808,10 +809,11 @@ tsl::Future<TopKTiledRunTimeDataOrError> TryFindTopKBestTilingsWithTilingSpace(
     }
 
     candidate_futures.push_back(tsl::MakeFutureOn(
-        *executor, [base, &fusion_adaptor, device_info = &device_info,
-                    shape_size, flops_fn, mlir_context_pool,
-                    padded_tile_sizes = std::move(padded_tile_sizes)]() {
-          return EvaluateCandidate(fusion_adaptor, *base->tiling_space,
+        *executor,
+        [base, &fusion_adaptor, device_info = &device_info, shape_size,
+         flops_fn, mlir_context_pool,
+         padded_tile_sizes = std::move(padded_tile_sizes)]() mutable {
+          return EvaluateCandidate(fusion_adaptor, std::move(base),
                                    padded_tile_sizes, mlir_context_pool,
                                    *device_info, shape_size, flops_fn);
         }));

From c85a1da72d368659c945a9601e0580b06daac475 Mon Sep 17 00:00:00 2001
From: Mikhail Goncharov <goncharov@google.com>
Date: Mon, 12 Oct 2026 03:32:58 -0700
Subject: [PATCH 0569/1000] [XLA:GPU] Fix unchecked access to absl::StatusOr in
 autotuner tests

PiperOrigin-RevId: 997810586
---
 .../xla/xla/backends/gpu/autotuner/BUILD      | 10 ++++++
 .../backends/gpu/autotuner/cublaslt_test.cc   | 17 ++++++---
 .../xla/backends/gpu/autotuner/cudnn_test.cc  | 20 ++++++++---
 .../gpu/autotuner/fission_backend_test.cc     | 36 +++++++++++++++----
 .../backends/gpu/autotuner/hipblaslt_test.cc  | 20 ++++++++---
 .../xla/backends/gpu/autotuner/miopen_test.cc | 19 +++++++---
 .../gpu/autotuner/native_emitter_test.cc      | 31 ++++++++++++++--
 .../xla/backends/gpu/autotuner/triton_test.cc | 30 ++++++++++++++--
 8 files changed, 155 insertions(+), 28 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/autotuner/BUILD b/third_party/xla/xla/backends/gpu/autotuner/BUILD
index 68a93676691..4fa20a5570a 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/BUILD
+++ b/third_party/xla/xla/backends/gpu/autotuner/BUILD
@@ -215,7 +215,9 @@ xla_test(
         "//xla/service/gpu:nvptx_compiler_impl",
         "//xla/stream_executor:blas",
         "//xla/stream_executor:device_description_proto_cc",
+        "//xla/stream_executor:platform",
         "//xla/stream_executor:stream_executor_h",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",
@@ -298,11 +300,13 @@ xla_test(
         "//xla/service/gpu:ir_emission_utils",
         "//xla/service/gpu:nvptx_compiler_impl",
         "//xla/stream_executor:device_description_proto_cc",
+        "//xla/stream_executor:platform",
         "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor/cuda:cuda_compute_capability",
         "//xla/stream_executor/cuda:cudnn_plugin",
         "//xla/tsl/protobuf:dnn_proto_cc",
         "//xla/tsl/util/proto:proto_matchers",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",
@@ -407,6 +411,7 @@ xla_test(
         "//xla/tsl/platform:env",
         "//xla/tsl/testing:temporary_directory",
         "//xla/tsl/util/proto:proto_matchers",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:status_matchers",
@@ -525,6 +530,7 @@ xla_test(
         "//xla/stream_executor:stream_executor_h",
         "//xla/tsl/platform:statusor",
         "//xla/tsl/util/proto:proto_matchers",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",
@@ -660,7 +666,9 @@ xla_test(
         "//xla/service/gpu:amdgpu_compiler_impl",
         "//xla/stream_executor:blas",
         "//xla/stream_executor:device_description_proto_cc",
+        "//xla/stream_executor:platform",
         "//xla/stream_executor:stream_executor_h",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",
@@ -866,6 +874,7 @@ xla_test(
         "//xla/stream_executor:platform",
         "//xla/stream_executor:stream_executor_h",
         "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",
@@ -1113,6 +1122,7 @@ xla_test(
         "//xla/stream_executor/rocm:rocm_platform_id",
         "//xla/tsl/protobuf:dnn_proto_cc",
         "//xla/tsl/util/proto:proto_matchers",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",
diff --git a/third_party/xla/xla/backends/gpu/autotuner/cublaslt_test.cc b/third_party/xla/xla/backends/gpu/autotuner/cublaslt_test.cc
index 4e6fecbe34f..5fb2cea2d2c 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/cublaslt_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/cublaslt_test.cc
@@ -21,6 +21,7 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
 #include "absl/status/statusor.h"
@@ -38,6 +39,7 @@ limitations under the License.
 #include "xla/shape.h"
 #include "xla/stream_executor/blas.h"
 #include "xla/stream_executor/device_description.pb.h"
+#include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/xla.pb.h"
 
@@ -52,6 +54,16 @@ using ::testing::Optional;
 
 using CublasLtBackendConfig = AutotuneResult::GemmKey;
 
+// Returns the first device executor of the default platform.
+se::StreamExecutor* GetDefaultStreamExecutor() {
+  absl::StatusOr<se::Platform*> platform = PlatformUtil::GetDefaultPlatform();
+  CHECK_OK(platform.status());
+  absl::StatusOr<se::StreamExecutor*> stream_executor =
+      (*platform)->ExecutorForDevice(0);
+  CHECK_OK(stream_executor.status());
+  return *stream_executor;
+}
+
 const char kCublasLtCustomCallHlo[] = R"(
 HloModule module
 
@@ -119,10 +131,7 @@ class CublasLtBackendTest : public HloHardwareIndependentTestBase {
   CublasLtBackendTest()
       : debug_options_(
             HloHardwareIndependentTestBase::GetDebugOptionsForTest()),
-        stream_executor_(PlatformUtil::GetDefaultPlatform()
-                             .value()
-                             ->ExecutorForDevice(0)
-                             .value()),
+        stream_executor_(GetDefaultStreamExecutor()),
         target_config_(stream_executor_),
         backend_(stream_executor_, &debug_options_, &compiler_, &target_config_,
                  &mlir_context_) {}
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 c424c28f5f8..209154e9491 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/cudnn_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/cudnn_test.cc
@@ -20,6 +20,7 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
 #include "absl/status/statusor.h"
@@ -38,6 +39,7 @@ limitations under the License.
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/cuda/cuda_dnn.h"
 #include "xla/stream_executor/device_description.pb.h"
+#include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/protobuf/dnn.pb.h"
 #include "xla/tsl/util/proto/proto_matchers.h"
@@ -45,6 +47,19 @@ limitations under the License.
 
 namespace xla {
 namespace gpu {
+namespace {
+
+// Returns the first device executor of the default platform.
+se::StreamExecutor* GetDefaultStreamExecutor() {
+  absl::StatusOr<se::Platform*> platform = PlatformUtil::GetDefaultPlatform();
+  CHECK_OK(platform.status());
+  absl::StatusOr<se::StreamExecutor*> stream_executor =
+      (*platform)->ExecutorForDevice(0);
+  CHECK_OK(stream_executor.status());
+  return *stream_executor;
+}
+
+}  // namespace
 
 using CudnnBackendConfig = stream_executor::dnn::AlgorithmProto;
 
@@ -183,10 +198,7 @@ absl::string_view kUnsupportedHlo = R"hlo(
 class CudnnBackendTest : public HloHardwareIndependentTestBase {
  protected:
   CudnnBackendTest()
-      : stream_executor_(PlatformUtil::GetDefaultPlatform()
-                             .value()
-                             ->ExecutorForDevice(0)
-                             .value()),
+      : stream_executor_(GetDefaultStreamExecutor()),
         target_config_(stream_executor_),
         debug_options_(
             HloHardwareIndependentTestBase::GetDebugOptionsForTest()) {
diff --git a/third_party/xla/xla/backends/gpu/autotuner/fission_backend_test.cc b/third_party/xla/xla/backends/gpu/autotuner/fission_backend_test.cc
index 13ce51f17de..3f217f38d44 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/fission_backend_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/fission_backend_test.cc
@@ -24,6 +24,7 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
@@ -74,6 +75,29 @@ using ::testing::HasSubstr;
 using ::testing::IsEmpty;
 using ::testing::Not;
 
+// Returns the default platform.
+se::Platform* GetDefaultPlatform() {
+  absl::StatusOr<se::Platform*> platform = PlatformUtil::GetDefaultPlatform();
+  CHECK_OK(platform.status());
+  return *platform;
+}
+
+// Returns the compiler registered for `platform`.
+std::unique_ptr<Compiler> GetCompilerForPlatform(se::Platform* platform) {
+  absl::StatusOr<std::unique_ptr<Compiler>> compiler =
+      Compiler::GetForPlatform(platform->id());
+  CHECK_OK(compiler.status());
+  return *std::move(compiler);
+}
+
+// Returns the first device executor of `platform`.
+se::StreamExecutor* GetStreamExecutor(se::Platform* platform) {
+  absl::StatusOr<se::StreamExecutor*> stream_executor =
+      platform->ExecutorForDevice(0);
+  CHECK_OK(stream_executor.status());
+  return *stream_executor;
+}
+
 const char kTritonFusionHlo[] = R"(
   HloModule module
 
@@ -282,9 +306,9 @@ class FissionTest : public HloHardwareIndependentTestBase,
   FissionTest()
       : debug_options_(
             HloHardwareIndependentTestBase::GetDebugOptionsForTest()),
-        platform_(PlatformUtil::GetDefaultPlatform().value()),
-        stream_executor_(platform_->ExecutorForDevice(0).value()),
-        compiler_(Compiler::GetForPlatform(platform_->id()).value()),
+        platform_(GetDefaultPlatform()),
+        stream_executor_(GetStreamExecutor(platform_)),
+        compiler_(GetCompilerForPlatform(platform_)),
         target_config_(stream_executor_),
         device_description_(stream_executor_->GetDeviceDescription()),
         rewriter_pipeline_(GetParam().pipeline_factory(device_description_)),
@@ -500,9 +524,9 @@ class CublasFissionBackendTest : public HloHardwareIndependentTestBase {
   CublasFissionBackendTest()
       : debug_options_(
             HloHardwareIndependentTestBase::GetDebugOptionsForTest()),
-        platform_(PlatformUtil::GetDefaultPlatform().value()),
-        stream_executor_(platform_->ExecutorForDevice(0).value()),
-        compiler_(Compiler::GetForPlatform(platform_->id()).value()),
+        platform_(GetDefaultPlatform()),
+        stream_executor_(GetStreamExecutor(platform_)),
+        compiler_(GetCompilerForPlatform(platform_)),
         target_config_(stream_executor_),
         device_description_(stream_executor_->GetDeviceDescription()),
         alias_info_(device_description_),
diff --git a/third_party/xla/xla/backends/gpu/autotuner/hipblaslt_test.cc b/third_party/xla/xla/backends/gpu/autotuner/hipblaslt_test.cc
index 9260089f530..ee2a5acbbfb 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/hipblaslt_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/hipblaslt_test.cc
@@ -21,6 +21,7 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
 #include "absl/status/statusor.h"
@@ -35,11 +36,25 @@ limitations under the License.
 #include "xla/service/platform_util.h"
 #include "xla/stream_executor/blas.h"
 #include "xla/stream_executor/device_description.pb.h"
+#include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/xla.pb.h"
 
 namespace xla {
 namespace gpu {
+namespace {
+
+// Returns the first device executor of the default platform.
+se::StreamExecutor* GetDefaultStreamExecutor() {
+  absl::StatusOr<se::Platform*> platform = PlatformUtil::GetDefaultPlatform();
+  CHECK_OK(platform.status());
+  absl::StatusOr<se::StreamExecutor*> stream_executor =
+      (*platform)->ExecutorForDevice(0);
+  CHECK_OK(stream_executor.status());
+  return *stream_executor;
+}
+
+}  // namespace
 
 using HipblasLtBackendConfig = AutotuneResult::GemmKey;
 
@@ -109,10 +124,7 @@ class HipblasLtBackendTest : public HloHardwareIndependentTestBase {
   HipblasLtBackendTest()
       : debug_options_(
             HloHardwareIndependentTestBase::GetDebugOptionsForTest()),
-        stream_executor_(PlatformUtil::GetDefaultPlatform()
-                             .value()
-                             ->ExecutorForDevice(0)
-                             .value()),
+        stream_executor_(GetDefaultStreamExecutor()),
         target_config_(stream_executor_),
         backend_(stream_executor_, &debug_options_, &compiler_,
                  &target_config_) {}
diff --git a/third_party/xla/xla/backends/gpu/autotuner/miopen_test.cc b/third_party/xla/xla/backends/gpu/autotuner/miopen_test.cc
index 8284117472f..ec1f809f82f 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/miopen_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/miopen_test.cc
@@ -20,6 +20,7 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
 #include "absl/status/statusor.h"
@@ -43,6 +44,19 @@ limitations under the License.
 
 namespace xla {
 namespace gpu {
+namespace {
+
+// Returns the first device executor of the default platform.
+se::StreamExecutor* GetDefaultStreamExecutor() {
+  absl::StatusOr<se::Platform*> platform = PlatformUtil::GetDefaultPlatform();
+  CHECK_OK(platform.status());
+  absl::StatusOr<se::StreamExecutor*> stream_executor =
+      (*platform)->ExecutorForDevice(0);
+  CHECK_OK(stream_executor.status());
+  return *stream_executor;
+}
+
+}  // namespace
 
 using MIOpenBackendConfig = stream_executor::dnn::AlgorithmProto;
 
@@ -83,10 +97,7 @@ class MIOpenBackendTest : public HloHardwareIndependentTestBase {
   MIOpenBackendTest()
       : debug_options_(
             HloHardwareIndependentTestBase::GetDebugOptionsForTest()),
-        stream_executor_(PlatformUtil::GetDefaultPlatform()
-                             .value()
-                             ->ExecutorForDevice(0)
-                             .value()),
+        stream_executor_(GetDefaultStreamExecutor()),
         target_config_(stream_executor_),
         allocator_(stream_executor_),
         backend_(
diff --git a/third_party/xla/xla/backends/gpu/autotuner/native_emitter_test.cc b/third_party/xla/xla/backends/gpu/autotuner/native_emitter_test.cc
index c53c7115c03..7d6cfe244c9 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/native_emitter_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/native_emitter_test.cc
@@ -17,10 +17,12 @@ limitations under the License.
 
 #include <cstdint>
 #include <memory>
+#include <utility>
 #include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
 #include "absl/status/statusor.h"
@@ -50,6 +52,29 @@ using ::testing::UnorderedElementsAre;
 using ::tsl::proto_testing::EqualsProto;
 using ::xla::xtile::BlockLevelFusionConfig;
 
+// Returns the default platform.
+se::Platform* GetDefaultPlatform() {
+  absl::StatusOr<se::Platform*> platform = PlatformUtil::GetDefaultPlatform();
+  CHECK_OK(platform.status());
+  return *platform;
+}
+
+// Returns the compiler registered for `platform`.
+std::unique_ptr<Compiler> GetCompilerForPlatform(se::Platform* platform) {
+  absl::StatusOr<std::unique_ptr<Compiler>> compiler =
+      Compiler::GetForPlatform(platform->id());
+  CHECK_OK(compiler.status());
+  return *std::move(compiler);
+}
+
+// Returns the first device executor of `platform`.
+se::StreamExecutor* GetStreamExecutor(se::Platform* platform) {
+  absl::StatusOr<se::StreamExecutor*> stream_executor =
+      platform->ExecutorForDevice(0);
+  CHECK_OK(stream_executor.status());
+  return *stream_executor;
+}
+
 const char kReductionFusionHlo[] = R"(
 HloModule m
 
@@ -118,10 +143,10 @@ class NativeEmitterBackendTest : public HloHardwareIndependentTestBase {
   NativeEmitterBackendTest()
       : debug_options_(
             HloHardwareIndependentTestBase::GetDebugOptionsForTest()),
-        platform_(PlatformUtil::GetDefaultPlatform().value()),
-        stream_executor_(platform_->ExecutorForDevice(0).value()),
+        platform_(GetDefaultPlatform()),
+        stream_executor_(GetStreamExecutor(platform_)),
         target_config_(stream_executor_),
-        compiler_(Compiler::GetForPlatform(platform_->id()).value()),
+        compiler_(GetCompilerForPlatform(platform_)),
         backend_(&debug_options_, compiler_.get(), &target_config_) {
     debug_options_.set_xla_gpu_libnvjitlink_mode(
         xla::DebugOptions::LIB_NV_JIT_LINK_MODE_DISABLED);
diff --git a/third_party/xla/xla/backends/gpu/autotuner/triton_test.cc b/third_party/xla/xla/backends/gpu/autotuner/triton_test.cc
index 8fccdd56f79..7053126dc90 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/triton_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/triton_test.cc
@@ -26,6 +26,7 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/status_matchers.h"
@@ -71,6 +72,29 @@ using ::testing::Optional;
 using ::testing::SizeIs;
 using ::tsl::proto_testing::EqualsProto;
 
+// Returns the default platform.
+se::Platform* GetDefaultPlatform() {
+  absl::StatusOr<se::Platform*> platform = PlatformUtil::GetDefaultPlatform();
+  CHECK_OK(platform.status());
+  return *platform;
+}
+
+// Returns the compiler registered for `platform`.
+std::unique_ptr<Compiler> GetCompilerForPlatform(se::Platform* platform) {
+  absl::StatusOr<std::unique_ptr<Compiler>> compiler =
+      Compiler::GetForPlatform(platform->id());
+  CHECK_OK(compiler.status());
+  return *std::move(compiler);
+}
+
+// Returns the first device executor of `platform`.
+se::StreamExecutor* GetStreamExecutor(se::Platform* platform) {
+  absl::StatusOr<se::StreamExecutor*> stream_executor =
+      platform->ExecutorForDevice(0);
+  CHECK_OK(stream_executor.status());
+  return *stream_executor;
+}
+
 const char kHlo[] = R"(
   HloModule module
 
@@ -133,11 +157,11 @@ class TritonBackendTest : public HloHardwareIndependentTestBase {
   TritonBackendTest()
       : debug_options_(
             HloHardwareIndependentTestBase::GetDebugOptionsForTest()),
-        platform_(PlatformUtil::GetDefaultPlatform().value()),
-        stream_executor_(platform_->ExecutorForDevice(0).value()),
+        platform_(GetDefaultPlatform()),
+        stream_executor_(GetStreamExecutor(platform_)),
         target_config_(stream_executor_),
         alias_info_(stream_executor_->GetDeviceDescription()),
-        compiler_(Compiler::GetForPlatform(platform_->id()).value()),
+        compiler_(GetCompilerForPlatform(platform_)),
         backend_(&debug_options_, compiler_.get(), &target_config_,
                  &alias_info_, &mlir_context_pool_) {}
 

From 126bc38b942995a45c8da03fc8adf1e126086dd5 Mon Sep 17 00:00:00 2001
From: Dmitri Latushko <latushko@google.com>
Date: Mon, 12 Oct 2026 04:20:44 -0700
Subject: [PATCH 0570/1000] Increase CPU shard_count of multi_device_spmd_test
 from 25 to 50.

With `DTENSOR_ENABLE_MULTI_DEVICE_EXPANSION=1`, `multi_device_spmd_test_cpu` compiles and caches multi-device SPMD MLIR modules and execution functions across ~35 parameterized test cases per shard at `shard_count = 25`. Under ASAN and MSAN, individual shards exceed the per-shard memory limit near the final `testUnaryOps*` test cases. Increasing `shard_count["cpu"]` from `25` to `50` halves the per-shard test count and peak memory footprint.

PiperOrigin-RevId: 997830441
---
 tensorflow/dtensor/python/tests/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/dtensor/python/tests/BUILD b/tensorflow/dtensor/python/tests/BUILD
index 0db3a7cf26e..23e1fb709e8 100644
--- a/tensorflow/dtensor/python/tests/BUILD
+++ b/tensorflow/dtensor/python/tests/BUILD
@@ -763,7 +763,7 @@ dtensor_test(
     },
     main = "spmd_test.py",
     shard_count = {
-        "cpu": 25,
+        "cpu": 50,
         "gpu": 10,
         "tpu": 10,
         TPU_V3_DONUT_BACKEND: 32,

From 823d348ecd19d09e00491bcaeeeef8681c5c4a5a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 12 Oct 2026 04:46:11 -0700
Subject: [PATCH 0571/1000] Replace the `GetMemoryLimit(*assignment, color) >
 0` check for fast buffer assignment with an explicit list of memory spaces.
 Previously, this condition only captured HBM and was unclear because it was
 used both as an enablement check and to determine the fallback memory limit.

In this CL we explicitly enable fast buffer assignment for the `kNone` memory space (color 6). `kNone` indicates no mandatory memory assignment (which typically places the value in an `SREG`). Enabling fast BA for this space can save up to several minutes during compilation.

PiperOrigin-RevId: 997839830
---
 third_party/xla/xla/service/buffer_assignment.cc | 8 +++++---
 third_party/xla/xla/service/buffer_assignment.h  | 5 +++++
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/service/buffer_assignment.cc b/third_party/xla/xla/service/buffer_assignment.cc
index bf53792fca7..e3f1da7179a 100644
--- a/third_party/xla/xla/service/buffer_assignment.cc
+++ b/third_party/xla/xla/service/buffer_assignment.cc
@@ -2550,7 +2550,7 @@ absl::Status BufferAssigner::AssignBuffersForGlobalComputations(
         buffer_assignment::
             AssignmentAlgorithmForComputationsWithoutOrderingProto::
                 FAST_MERGE) {
-      if (GetMemoryLimit(*assignment, color) > 0) {
+      if (opts_.supports_fast_buffer_assignment(color)) {
         return &fast_manager;
       }
     }
@@ -2966,8 +2966,10 @@ absl::Status BufferAssigner::AssignBuffersWithSequentialOrdering(
     buffer_assignment::BufferAssignmentAlgorithmProto::Value algo_to_use =
         buffer_assignment_algorithm;
     if (algo_to_use ==
-        buffer_assignment::BufferAssignmentAlgorithmProto::FAST_MERGE) {
-      if (GetMemoryLimit(*assignment, color) == 0) {
+            buffer_assignment::BufferAssignmentAlgorithmProto::FAST_MERGE ||
+        algo_to_use ==
+            buffer_assignment::BufferAssignmentAlgorithmProto::FAST_SPLIT) {
+      if (!opts_.supports_fast_buffer_assignment(color)) {
         algo_to_use =
             buffer_assignment::BufferAssignmentAlgorithmProto::DEFAULT;
       }
diff --git a/third_party/xla/xla/service/buffer_assignment.h b/third_party/xla/xla/service/buffer_assignment.h
index ee33ba8564e..38a89223cd4 100644
--- a/third_party/xla/xla/service/buffer_assignment.h
+++ b/third_party/xla/xla/service/buffer_assignment.h
@@ -974,6 +974,11 @@ class BufferAssigner {
     // If set and returns > 0, the returned limit is used instead of the
     // default module config's device memory size.
     std::function<int64_t(LogicalBuffer::Color)> color_memory_limit;
+
+    // Whether a buffer color supports fast buffer assignment (FAST_MERGE /
+    // FAST_SPLIT); unsupported colors use the DEFAULT algorithm.
+    std::function<bool(LogicalBuffer::Color)> supports_fast_buffer_assignment =
+        [](LogicalBuffer::Color) { return true; };
   };
 
   static std::unique_ptr<BufferAllocationsManagerForComputationsWithoutOrdering>

From 93925a442ccb4ebb95efe2bdc0af3cacdbf6d676 Mon Sep 17 00:00:00 2001
From: Saurabh Singh <geniusthegreat0963@gmail.com>
Date: Mon, 12 Oct 2026 23:54:26 +0530
Subject: [PATCH 0572/1000] Use self.match() in testIsCloseIntegerTolerances
 for graph-mode compatibility

Replace bool(...) call on allclose with self.match(...) to prevent
TypeError in graph execution mode.
---
 tensorflow/python/ops/numpy_ops/np_math_ops_test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

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 f58808ec7ab..9a5bf467bfd 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
@@ -442,8 +442,8 @@ class MathTest(test.TestCase, parameterized.TestCase):
         np_math_ops.isclose(a, b, rtol=9, atol=0),
         np.isclose(a, b, rtol=9, atol=0),
     )
-    self.assertEqual(
-        bool(np_math_ops.allclose(a, b, rtol=9, atol=0)),
+    self.match(
+        np_math_ops.allclose(a, b, rtol=9, atol=0),
         np.allclose(a, b, rtol=9, atol=0),
     )
     c = np.array([100, 200, -300, 0], dtype)

From 6c38a24ccbf10924922beead70881b0ba83b22e9 Mon Sep 17 00:00:00 2001
From: Kingston Mandisodza <kingstonm@google.com>
Date: Mon, 12 Oct 2026 04:48:46 -0700
Subject: [PATCH 0573/1000] Initialise the windowed einsum reduce buffer with
 the reduction's initial value.

When a reduce over a windowed dimension in windowed einsum is moved into the loop and accumulated across iterations, the loop-carried result buffer was previously initialised with zero. This causes incorrect results for reductions whose identity value is not zero (such as max, min, or multiply).

PiperOrigin-RevId: 997840757
---
 .../xla/xla/service/spmd/dot_handler.cc       |  5 +--
 .../xla/service/spmd/spmd_partitioner_test.cc | 42 +++++++++++++++++++
 2 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/dot_handler.cc b/third_party/xla/xla/service/spmd/dot_handler.cc
index c1c5fabe94c..937a7445552 100644
--- a/third_party/xla/xla/service/spmd/dot_handler.cc
+++ b/third_party/xla/xla/service/spmd/dot_handler.cc
@@ -4784,10 +4784,7 @@ absl::Status MoveUsersIntoWindowedDotGeneralLoopOnNonContractingDimensions(
     }
     auto* broadcast =
         computation->AddInstruction(HloInstruction::CreateBroadcast(
-            padded_out_shape,
-            computation->AddInstruction(HloInstruction::CreateConstant(
-                LiteralUtil::Zero(out->shape().element_type()))),
-            {}));
+            padded_out_shape, out->mutable_operand(1), {}));
     new_operands.push_back(broadcast);
   }
 
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
index 5cb9571e86c..eaffda4c64c 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include <algorithm>
 #include <cstdint>
 #include <iostream>
+#include <limits>
 #include <memory>
 #include <optional>
 #include <string>
@@ -8620,6 +8621,47 @@ ENTRY entry {
                 output_tuple, op::GetTupleElement(op::Parameter(0)), next_i));
 }
 
+// Verify that a reduce moved into a windowed einsum loop initialises its
+// accumulator buffer with the reduce's init value.
+TEST_P(SpmdPartitioningTest, EinsumWindowedNonContractingMaxReduceUsesInit) {
+  absl::string_view hlo_string = R"(
+HloModule module
+
+max {
+  a = f32[] parameter(0)
+  b = f32[] parameter(1)
+  ROOT maximum = f32[] maximum(a, b)
+}
+
+ENTRY entry {
+  %x = f32[8,16] parameter(0), sharding={devices=[4,1]<=[4]}
+  %w = f32[16,256] parameter(1), sharding={devices=[1,4]<=[4]}
+  %dot = f32[8,256] dot(%x, %w), lhs_contracting_dims={1},
+    rhs_contracting_dims={0}, sharding={devices=[1,4]<=[4]}
+  %init = f32[] constant(-inf)
+  ROOT %reduce = f32[256] reduce(%dot, %init), dimensions={0}, to_apply=max,
+    sharding={devices=[4]<=[4]}
+})";
+
+  SpmdPartitionerOptions options;
+  options.threshold_for_windowed_einsum_mib = 0;
+  ASSERT_OK_AND_ASSIGN(
+      auto module,
+      PartitionComputation(hlo_string, /*num_devices=*/4, options));
+  VLOG(1) << module->ToString();
+
+  const auto root = module->entry_computation()->root_instruction();
+  const auto neg_inf = op::Constant(
+      LiteralUtil::CreateR0<float>(-std::numeric_limits<float>::infinity()));
+  const auto input_subtuple =
+      op::Tuple(neg_inf, AllOf(op::Broadcast(neg_inf), op::Shape("f32[64]")));
+
+  EXPECT_THAT(root,
+              AllOf(op::GetTupleElement(op::GetTupleElement(op::While(
+                        op::Tuple(_, _, input_subtuple, _, op::Constant())))),
+                    op::Shape("f32[64]")));
+}
+
 TEST_P(SpmdPartitioningTest,
        BidirectionalEinsumRHSWindowedNonContractingReduce2) {
   absl::string_view hlo_string = R"(

From 8f44f6a7ab7bd9457dda6f6d33b622d507f9e3b4 Mon Sep 17 00:00:00 2001
From: Akhil Goel <akhil.goel@intel.com>
Date: Mon, 12 Oct 2026 04:56:45 -0700
Subject: [PATCH 0574/1000] PR #50356: [XLA:GPU][oneAPI] Align XPU Triton hook
 with upstream LLVM

Imported from GitHub PR https://github.com/openxla/xla/pull/50356

This PR aligns the XPU Triton hook with https://github.com/llvm/llvm-project/commit/c4d9267a31897dd76b2aab5196b75bffb49e6846
It replaces the usages of `amdgpu::Chipset` with `ROCDL::TargetInfo` in GPU to ROCDL conversion patterns. This is needed to remain compatible with this breaking LLVM change https://github.com/llvm/llvm-project/commit/155462f440ffd574852b6ddc2b364907c9e235fc, included in the most recent LLVM update.
Copybara import of the project:

--
a904ef553a54fcdb79ead0150424716ea43af9d7 by Akhil Goel <akhil.goel@intel.com>:

Update XPU Triton hook

Merging this change closes #50356

PiperOrigin-RevId: 997843756
---
 third_party/xla/third_party/triton/intel_xpu/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/third_party/triton/intel_xpu/workspace.bzl b/third_party/xla/third_party/triton/intel_xpu/workspace.bzl
index 5f8ef5a1c44..1c3f756406b 100644
--- a/third_party/xla/third_party/triton/intel_xpu/workspace.bzl
+++ b/third_party/xla/third_party/triton/intel_xpu/workspace.bzl
@@ -1,7 +1,7 @@
 """Intel XPU Triton archive metadata."""
 
-XPU_TRITON_COMMIT = "9b507bc72f26026254e572c8d9e59c22497ae522"
-XPU_TRITON_SHA256 = "98f63c0a7fbbf5c6a13a76fe976e3d2f352d731f56bf75c8043fa77e302a9772"
+XPU_TRITON_COMMIT = "6046064c56a27193423475e0a6f412a08b13c6e6"
+XPU_TRITON_SHA256 = "158213c620f555aaf295613f48240c63a6b40dbe7db3805f97f4c6f2361b7f7d"
 
 def use_xpu_triton(repository_ctx):
     return repository_ctx.getenv("ENABLE_INTEL_XPU_TRITON", "").strip() == "1"

From 3a32173ca074b01c8176c59a3baaeb896c970ab8 Mon Sep 17 00:00:00 2001
From: Olli Lupton <olupton@nvidia.com>
Date: Mon, 12 Oct 2026 04:58:00 -0700
Subject: [PATCH 0575/1000] PR #50156: Fill AnnotationStack even when NVTX is
 active.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50156

📝 Summary of Changes

Fill AnnotationStack even when NVTX is active.

🎯 Justification

Previously this was not useful because CUPTI only supported a single subscriber, so attempting to use an external tool such as Nsight Systems (which loads an NVTX implementation and makes `domain` non-null) and the built-in XLA profiler was doomed. Now, with multi-subscriber support, this is possible, but without this change AutoPGLE-under-nsys did not work because the lack of AnnotationStack broke PGLE profile collection.

🚀 Kind of Contribution

🐛 Bug Fix

📊 Benchmark (for Performance Improvements)

n/a

🧪 Unit Tests:

None.

🧪 Execution Tests:

It would be difficult to test this without the ability to execute tests under nsight systems.
Copybara import of the project:

--
f2b8c0de1c891686e38a6937d4c2cabcb54a19a8 by Olli Lupton <olupton@nvidia.com>:

Fill AnnotationStack even when NVTX is active.

Previously this was not useful because CUPTI only supported a single
subscriber, so attempting to use an external tool such as Nsight Systems
(which loads an NVTX implementation and makes `domain` non-null) and the
built-in XLA profiler was doomed. Now, with multi-subscriber support,
this is possible, but without this change AutoPGLE-under-nsys did not
work because the lack of AnnotationStack broke PGLE profile collection.

--
a890e0fcb7f1bd2bf106a173c62eada476808727 by Olli Lupton <olupton@nvidia.com>:

Address code review

Merging this change closes #50156

PiperOrigin-RevId: 997844185
---
 .../third_party/tsl/tsl/profiler/lib/BUILD    |  1 -
 .../tsl/tsl/profiler/lib/scoped_annotation.h  | 43 +++++++++++--------
 2 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD b/third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD
index 238caa7630a..062921a2f27 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD
@@ -412,7 +412,6 @@ cc_library(
         "//tsl/platform",
         "@com_google_absl//absl/strings",
         "@xla//xla/tsl/platform:macros",
-        "@xla//xla/tsl/platform:types",
     ] + if_not_android([
         "@xla//xla/tsl/profiler/backends/cpu:annotation_stack",
     ]),
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/scoped_annotation.h b/third_party/xla/third_party/tsl/tsl/profiler/lib/scoped_annotation.h
index 86157d38adc..ba11bbbd8b5 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/scoped_annotation.h
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/scoped_annotation.h
@@ -37,18 +37,27 @@ namespace tsl::profiler {
 // TraceCollector until PopAnnotation() is called.
 template <typename Generator>
 void PushAnnotation(const Generator& generator) {
-  if (auto domain = DefaultProfilerDomain();
-      TF_PREDICT_FALSE(domain != nullptr)) {
-    RangePush(domain, generator());
-    return;
-  }
-
+  auto domain = DefaultProfilerDomain();
+  const bool push_nvtx = TF_PREDICT_FALSE(domain != nullptr);
 #if !defined(IS_MOBILE_PLATFORM)
-  if (TF_PREDICT_FALSE(AnnotationStack::IsEnabled())) {
-    AnnotationStack::PushAnnotation(
-        static_cast<absl::string_view>(generator()));
-  }
+  const bool push_stack = TF_PREDICT_FALSE(AnnotationStack::IsEnabled());
+#else
+  constexpr bool push_stack = false;
 #endif
+  if (TF_PREDICT_FALSE(push_nvtx || push_stack)) {
+    // Now CUPTI supports multiple subscribers we may need both sets of
+    // annotations at the same time.
+    decltype(auto) annotation = generator();
+    if (push_nvtx) {
+      RangePush(domain, annotation);
+    }
+#if !defined(IS_MOBILE_PLATFORM)
+    if (push_stack) {
+      AnnotationStack::PushAnnotation(
+          static_cast<absl::string_view>(annotation));
+    }
+#endif
+  }
 }
 
 template <typename AnnotationGenerator, typename RangeGenerator>
@@ -57,7 +66,6 @@ void PushAnnotation(const AnnotationGenerator& annotation_generator,
   if (auto domain = DefaultProfilerDomain();
       TF_PREDICT_FALSE(domain != nullptr)) {
     RangePush(domain, range_generator());
-    return;
   }
 
 #if !defined(IS_MOBILE_PLATFORM)
@@ -80,17 +88,16 @@ inline void PopAnnotation() {
   // fail probably due to compiler in that presubmit config.
   std::atomic_thread_fence(std::memory_order_acquire);
 
-  if (auto domain = DefaultProfilerDomain();
-      TF_PREDICT_FALSE(domain != nullptr)) {
-    RangePop(domain);
-    return;
-  }
-
 #if !defined(IS_MOBILE_PLATFORM)
   if (TF_PREDICT_FALSE(AnnotationStack::IsEnabled())) {
     AnnotationStack::PopAnnotation();
   }
 #endif
+
+  if (auto domain = DefaultProfilerDomain();
+      TF_PREDICT_FALSE(domain != nullptr)) {
+    RangePop(domain);
+  }
 }
 
 // Adds an annotation to all activities for the duration of the instance
@@ -120,7 +127,7 @@ class ScopedAnnotation {
                    std::forward<RangeGenerator>(range_generator));
   }
 
-  // Pops the annotation or range pushed by the constructor.
+  // Pops the annotation and/or range pushed by the constructor.
   ~ScopedAnnotation() { PopAnnotation(); }
 
   static bool IsEnabled() {

From bd65780e9e266b8b89d9a9bfb66df3c472545a27 Mon Sep 17 00:00:00 2001
From: Nikita Putikhin <nputikhin@google.com>
Date: Mon, 12 Oct 2026 06:10:28 -0700
Subject: [PATCH 0576/1000] Add flag to keep Triton GEMMs enabled in
 deterministic mode

We disable Triton with `--xla_gpu_exclude_nondeterministic_ops` or `--xla_gpu_deterministic_ops`: its default config is very slow and the configs are not checked for correctness like in normal autotuning.

The new flag keeps Triton dsiabled by default but adds an option to enable it (without the correctness check).

Checked mode is coming later.

PiperOrigin-RevId: 997873087
---
 third_party/xla/xla/debug_options_flags.cc    | 24 +++++++++++++++
 .../xla/xla/debug_options_flags_test.cc       | 25 ++++++++++++++++
 .../gpu/autotuning/config_assigner_pass.cc    |  6 ++--
 .../autotuning/config_assigner_pass_test.cc   | 30 ++++++++++++++-----
 third_party/xla/xla/xla.proto                 | 13 +++++++-
 5 files changed, 88 insertions(+), 10 deletions(-)

diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index fa090afea19..6f063f4cf77 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -629,6 +629,8 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() {
       DebugOptions::DEVICELESS_CUB_WITH_FALLBACK);
   opts.set_xla_gpu_cudnn_deviceless_compilation_mode(
       DebugOptions::CUDNN_DEVICELESS_COMPILATION_DISABLED);
+  opts.set_xla_gpu_experimental_deterministic_ops_triton(
+      DebugOptions::DETERMINISTIC_OPS_TRITON_MODE_OFF);
   // Pre-populate all default autotune backends so that modifier flags (e.g.
   // +cudnn, -triton) can apply to the full default set.
   for (int i = 0; i < autotuner::Backend_descriptor()->value_count(); ++i) {
@@ -3998,6 +4000,28 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
       /*default_value_for_display=*/"",
       "Whitespace-separated extra flags to pass to the ptxas compiler, e.g. "
       "'--maxregcount=32'."));
+  auto setter_for_xla_gpu_experimental_deterministic_ops_triton =
+      [debug_options](absl::string_view input) {
+        DebugOptions::DeterministicOpsTritonMode mode;
+        if (!DebugOptions::DeterministicOpsTritonMode_Parse(
+                absl::StrCat("DETERMINISTIC_OPS_TRITON_MODE_",
+                             absl::AsciiStrToUpper(input)),
+                &mode) ||
+            mode == DebugOptions::DETERMINISTIC_OPS_TRITON_MODE_UNSET) {
+          return false;
+        }
+        debug_options->set_xla_gpu_experimental_deterministic_ops_triton(mode);
+        return true;
+      };
+  flag_list->push_back(tsl::Flag(
+      "xla_gpu_experimental_deterministic_ops_triton",
+      setter_for_xla_gpu_experimental_deterministic_ops_triton,
+      /*default_value_for_display=*/"off",
+      "Controls whether Triton is enabled when deterministic ops are requested "
+      "(--xla_gpu_deterministic_ops or "
+      "--xla_gpu_exclude_nondeterministic_ops). "
+      "Supported modes: 'off' (default) and 'unchecked' (enables Triton and "
+      "skips correctness checks)."));
 }  // NOLINT(readability/fn_size)
 
 // Allocates flag_values and flag_objects; this function must not be called more
diff --git a/third_party/xla/xla/debug_options_flags_test.cc b/third_party/xla/xla/debug_options_flags_test.cc
index e7ff6c2f72a..68064de02b6 100644
--- a/third_party/xla/xla/debug_options_flags_test.cc
+++ b/third_party/xla/xla/debug_options_flags_test.cc
@@ -499,5 +499,30 @@ TEST(DebugOptions, DeduplicateBackendConfigsMinSizeFlagsParsing) {
   EXPECT_EQ(opts.xla_deduplicate_backend_configs_min_size(), 128);
 }
 
+TEST(DebugOptions, DeterministicOpsTriton) {
+  DebugOptions opts = DefaultDebugOptionsIgnoringFlags();
+  EXPECT_EQ(opts.xla_gpu_experimental_deterministic_ops_triton(),
+            DebugOptions::DETERMINISTIC_OPS_TRITON_MODE_OFF);
+
+  std::vector<tsl::Flag> flags;
+  MakeDebugOptionsFlags(&flags, &opts);
+  auto parse = [&](std::string value) {
+    std::vector<std::string> args = {
+        "--xla_gpu_experimental_deterministic_ops_triton=" + value};
+    return tsl::Flags::Parse(args, flags);
+  };
+
+  EXPECT_TRUE(parse("unchecked"));
+  EXPECT_EQ(opts.xla_gpu_experimental_deterministic_ops_triton(),
+            DebugOptions::DETERMINISTIC_OPS_TRITON_MODE_UNCHECKED);
+
+  EXPECT_TRUE(parse("off"));
+  EXPECT_EQ(opts.xla_gpu_experimental_deterministic_ops_triton(),
+            DebugOptions::DETERMINISTIC_OPS_TRITON_MODE_OFF);
+
+  EXPECT_FALSE(parse("unset"));
+  EXPECT_FALSE(parse("invalid"));
+}
+
 }  // namespace
 }  // namespace xla
diff --git a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc
index 33973b180ae..177debd715b 100644
--- a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc
+++ b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc
@@ -420,8 +420,10 @@ ConfigAssignerPass::GetEnabledBackends(
         autotuner::Backend::BLOCK_LEVEL_EMITTER);
   }
 
-  if (debug_options.xla_gpu_exclude_nondeterministic_ops() ||
-      debug_options.xla_gpu_deterministic_ops()) {
+  if ((debug_options.xla_gpu_exclude_nondeterministic_ops() ||
+       debug_options.xla_gpu_deterministic_ops()) &&
+      debug_options.xla_gpu_experimental_deterministic_ops_triton() ==
+          DebugOptions::DETERMINISTIC_OPS_TRITON_MODE_OFF) {
     disabled_autotune_backends.push_back(autotuner::Backend::TRITON);
   }
 
diff --git a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc
index 65670c3a8a2..8a9ee807834 100644
--- a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc
+++ b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc
@@ -702,19 +702,35 @@ TEST_F(AutotunerFlagsTest, GetEnabledBackendsRespectsDeterminism) {
   RegisterSymbolicExprStorage(&mlir_context);
   MlirContextPool mlir_context_pool(CreateMlirContext);
 
-  ASSERT_OK_AND_ASSIGN(std::vector<std::unique_ptr<CodegenBackend>> backends,
-                       ConfigAssignerPass::GetEnabledBackends(
-                           stream_executor_, allocator_.get(), &target_config,
-                           &alias_info, debug_options, &mlir_context,
-                           /*shape_size_fn=*/[](const Shape&) { return 0; },
-                           &compiler_, stream_executor_->GetPlatform()->id(),
-                           /*thread_pool=*/nullptr, &mlir_context_pool));
+  auto get_enabled_backends = [&]() {
+    return ConfigAssignerPass::GetEnabledBackends(
+        stream_executor_, allocator_.get(), &target_config, &alias_info,
+        debug_options, &mlir_context,
+        /*shape_size_fn=*/[](const Shape&) { return 0; }, &compiler_,
+        stream_executor_->GetPlatform()->id(),
+        /*thread_pool=*/nullptr, &mlir_context_pool);
+  };
 
+  ASSERT_OK_AND_ASSIGN(std::vector<std::unique_ptr<CodegenBackend>> backends,
+                       get_enabled_backends());
   for (const auto& backend : backends) {
     EXPECT_NE(backend->backend(), autotuner::Backend::TRITON);
     EXPECT_NE(backend->backend(), autotuner::Backend::NATIVE_EMITTER);
     EXPECT_NE(backend->backend(), autotuner::Backend::BLOCK_LEVEL_EMITTER);
   }
+
+  debug_options.set_xla_gpu_experimental_deterministic_ops_triton(
+      DebugOptions::DETERMINISTIC_OPS_TRITON_MODE_UNCHECKED);
+  ASSERT_OK_AND_ASSIGN(backends, get_enabled_backends());
+  bool has_triton = false;
+  for (const auto& backend : backends) {
+    if (backend->backend() == autotuner::Backend::TRITON) {
+      has_triton = true;
+    }
+    EXPECT_NE(backend->backend(), autotuner::Backend::NATIVE_EMITTER);
+    EXPECT_NE(backend->backend(), autotuner::Backend::BLOCK_LEVEL_EMITTER);
+  }
+  EXPECT_TRUE(has_triton);
 }
 
 TEST_F(ConfigAssignerPassTest, CublasLtSelectFirstConfig) {
diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index eba8f167462..aea36a7d7e3 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -535,6 +535,12 @@ message DebugOptions {
     EARLY_EXIT_POINT_AFTER_CONFIG_ASSIGNMENT = 1;
   }
 
+  enum DeterministicOpsTritonMode {
+    DETERMINISTIC_OPS_TRITON_MODE_UNSET = 0;
+    DETERMINISTIC_OPS_TRITON_MODE_OFF = 1;
+    DETERMINISTIC_OPS_TRITON_MODE_UNCHECKED = 2;
+  }
+
   // Filter which allows enabling NCCL symmetric buffers for a subset of
   // collectives based on their type and size.
   message CollectiveFilter {
@@ -996,6 +1002,11 @@ message DebugOptions {
   optional bool xla_gpu_experimental_collective_start_as_early_as_possible =
       488;
 
+  // Controls whether Triton is enabled when deterministic ops are requested
+  // (`xla_gpu_deterministic_ops` or `xla_gpu_exclude_nondeterministic_ops`).
+  optional DeterministicOpsTritonMode
+      xla_gpu_experimental_deterministic_ops_triton = 556;
+
   // Experimentally disables binary libraries in GPU compiler passes.
   optional bool xla_gpu_experimental_disable_binary_libraries = 329;
 
@@ -1931,7 +1942,7 @@ message DebugOptions {
   // Enables Raft for stable TopK.
   optional bool xla_gpu_experimental_enable_raft_for_stable_topk = 546;
 
-  // Next id: 556
+  // Next id: 557
 
   // Extra options to pass to the compilation backend (e.g. LLVM); specific
   // interpretation of these values is left to the backend.

From f6ab8644e83feba2c0d445ec611e0e8c7dda65d7 Mon Sep 17 00:00:00 2001
From: Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>
Date: Mon, 12 Oct 2026 07:16:59 -0700
Subject: [PATCH 0577/1000] PR #50460: [ROCm] Add .kpack entries for rocm 10.1
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50460

📝 Summary of Changes
Add missing miopen and rocsolver .kpack deps

🎯 Justification
Needed for rocm 10.1 and onwards

🚀 Kind of Contribution
🐛 Bug Fix

📊 Benchmark (for Performance Improvements)
N\A

🧪 Unit Tests:
None

🧪 Execution Tests:
None

Copybara import of the project:

--
7774e649ed8290ecc3d5513288fe23f1af44dd0a by Dragan Mladjenovic <Dragan.Mladjenovic@amd.com>:

[ROCm] Add .kpack entries for rocm 10.1

Merging this change closes #50460

PiperOrigin-RevId: 997901974
---
 third_party/xla/third_party/gpus/rocm/BUILD.tpl | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/third_party/gpus/rocm/BUILD.tpl b/third_party/xla/third_party/gpus/rocm/BUILD.tpl
index 78c923ea0e2..07b56ccf1f5 100644
--- a/third_party/xla/third_party/gpus/rocm/BUILD.tpl
+++ b/third_party/xla/third_party/gpus/rocm/BUILD.tpl
@@ -305,6 +305,9 @@ rocm_lib_import(
     ]) + glob([
         "%{rocm_root}/lib/libMIOpenCKGroupedConv_" + arch + ".so"
         for arch in rocm_gpu_architectures()
+    ]) + glob([
+        "%{rocm_root}/.kpack/miopen_lib_" + arch + ".kpack"
+        for arch in rocm_gpu_architectures()
     ]),
     interface_library = "%{rocm_root}/lib/libMIOpen.so",
     deps = [
@@ -395,8 +398,12 @@ rocm_lib_import(
 cc_library(
     name = "rocsparse_libs",
     data = glob(["%{rocm_root}/lib/librocsparse.so*"]) + glob([
-        "%{rocm_root}/.kpack/blas_lib_" + arch + ".kpack"
+        pattern
         for arch in rocm_gpu_architectures()
+        for pattern in (
+            "%{rocm_root}/.kpack/blas_lib_" + arch + ".kpack",
+            "%{rocm_root}/.kpack/sparse_lib_" + arch + ".kpack",
+        )
     ]),
     deps = [
         ":hip_runtime_libs",
@@ -451,8 +458,12 @@ rocm_lib_import(
         "%{rocm_root}/lib/librocsolver.so*",
         "%{rocm_root}/lib/host-math/lib/*.so*",
     ]) + glob([
-        "%{rocm_root}/.kpack/blas_lib_" + arch + ".kpack"
+        pattern
         for arch in rocm_gpu_architectures()
+        for pattern in (
+            "%{rocm_root}/.kpack/blas_lib_" + arch + ".kpack",
+            "%{rocm_root}/.kpack/solver_lib_" + arch + ".kpack",
+        )
     ]),
     interface_library = "%{rocm_root}/lib/librocsolver.so",
     deps = [

From f09069e81cd7f08c24760f890fc9cc4cdeb113de Mon Sep 17 00:00:00 2001
From: greypilgrim-083 <prashantchaudhary7353@gmail.com>
Date: Tue, 13 Oct 2026 00:38:07 +0530
Subject: [PATCH 0578/1000] Add missing //tensorflow/core:framework dep to
 sliding_window_dataset_op_test

---
 tensorflow/core/kernels/data/experimental/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/core/kernels/data/experimental/BUILD b/tensorflow/core/kernels/data/experimental/BUILD
index 532d99074d1..84eae61f1a4 100644
--- a/tensorflow/core/kernels/data/experimental/BUILD
+++ b/tensorflow/core/kernels/data/experimental/BUILD
@@ -861,6 +861,7 @@ tf_cc_test(
     deps = [
         ":sliding_window_dataset_op",
         "//tensorflow/core:experimental_dataset_ops_op_lib",
+        "//tensorflow/core:framework",
         "//tensorflow/core:test_main",
         "//tensorflow/core/data:dataset_test_base",
         "//tensorflow/core/data:name_utils",

From 7908588528658a2d832ff2b2fecbdfe7e313a65d Mon Sep 17 00:00:00 2001
From: Oleg Shyshkov <shyshkov@google.com>
Date: Mon, 12 Oct 2026 07:17:40 -0700
Subject: [PATCH 0579/1000] [XLA:GPU] Consolidate collective input and output
 memory requirement checks.

Previously, gpu_compiler.cc and gpu_memory_space_assignment.cc had duplicate and overlapping logic for identifying when HLO values or uses require collective memory space.

PiperOrigin-RevId: 997902338
---
 third_party/xla/xla/service/gpu/BUILD         |   2 +
 .../xla/xla/service/gpu/gpu_compiler.cc       | 129 +-----------
 .../gpu/gpu_memory_space_assignment.cc        | 148 +++++++------
 .../service/gpu/gpu_memory_space_assignment.h |   9 +
 .../gpu/gpu_memory_space_assignment_test.cc   | 199 ++++++++++++++++++
 5 files changed, 289 insertions(+), 198 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 8fcf2718518..98221e9f0be 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -141,6 +141,7 @@ xla_cc_test(
     srcs = ["gpu_memory_space_assignment_test.cc"],
     deps = [
         ":gpu_memory_space_assignment",
+        "//xla:shape_util",
         "//xla:xla_data_proto_cc",
         "//xla:xla_proto_cc",
         "//xla/hlo/analysis:alias_info",
@@ -150,6 +151,7 @@ xla_cc_test(
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
         "//xla/hlo/testlib:verified_hlo_module",
         "//xla/service:buffer_assignment",
+        "//xla/service:hlo_buffer",
         "//xla/service:hlo_module_config",
         "//xla/service:hlo_value",
         "//xla/tsl/lib/core:status_test_util",
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 599e4ab68f7..62fdb262f45 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -247,7 +247,6 @@ limitations under the License.
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/buffer_value.h"
 #include "xla/service/call_inliner.h"
-#include "xla/service/collective_ops_utils.h"
 #include "xla/service/collective_pipeliner.h"
 #include "xla/service/collective_pipeliner_utils.h"
 #include "xla/service/collective_utils.h"
@@ -2429,132 +2428,6 @@ absl::StatusOr<std::unique_ptr<HloModule>> GpuCompiler::RunHloPasses(
   return std::move(module);
 }
 
-namespace {
-
-bool UsesCollectiveMemorySpaceFrontendAttr(const HloUse& use) {
-  if (use.instruction->opcode() != HloOpcode::kCustomCall) {
-    return false;
-  }
-  auto attr =
-      use.instruction->get_frontend_attribute(kOperandsMemorySpacesAttr);
-  if (!attr.has_value()) {
-    return false;
-  }
-  auto pairs = ParseIndexMemorySpacePairs(*attr);
-  if (!pairs.ok()) {
-    return false;
-  }
-  for (auto [index, memory_space] : *pairs) {
-    if (index == use.operand_number &&
-        memory_space == MemorySpaceColor::kCollective) {
-      return true;
-    }
-  }
-  return false;
-}
-
-bool DefinesCollectiveMemorySpaceFrontendAttr(const HloValue* value) {
-  const HloInstruction* def = value->defining_instruction();
-  if (def->opcode() != HloOpcode::kCustomCall) {
-    return false;
-  }
-
-  auto attr = def->get_frontend_attribute(kResultsMemorySpacesAttr);
-  if (!attr.has_value()) {
-    return false;
-  }
-
-  auto pairs = ParseIndexMemorySpacePairs(*attr);
-  if (!pairs.ok()) {
-    return false;
-  }
-
-  // Determine the logical result index. If the custom call returns a tuple,
-  // we look at the top-level index (e.g., element 0 or 1 of the tuple).
-  int64_t result_index = 0;
-  if (def->shape().IsTuple()) {
-    if (value->defining_index().empty()) {
-      // The buffer for the tuple pointer array itself is not S1.
-      return false;
-    }
-    result_index = value->defining_index()[0];
-  }
-
-  for (auto [index, memory_space] : *pairs) {
-    if (index == result_index &&
-        memory_space == MemorySpaceColor::kCollective) {
-      return true;
-    }
-  }
-
-  return false;
-}
-
-bool RequiresCollectiveInput(const HloUse& use, const DebugOptions& opts) {
-  HloInstruction* user = use.instruction;
-
-  // Handle standard non-fusion/fusion collectives under NCCL user buffers
-  if (IsCollective(user) &&
-      (opts.xla_gpu_enable_nccl_user_buffers() ||
-       IsNcclSymmetricBuffersEnabledForCollective(user, opts))) {
-    return true;
-  }
-  // Handle one-shot RaggedAllToAll with NCCL enabled.
-  if (IsOneShotRaggedAllToAllWithNcclEnabled(opts)) {
-    // User is RA2A or AsyncStart(RA2A). Used by operand 1 (output).
-    if ((IsRaggedAllToAllOrAsyncStartRaggedAllToAll(user)) &&
-        use.operand_number == 1) {
-      return true;
-    }
-  }
-
-  // Device-initiated and one-sided collectives require S1 memory for all
-  // buffers.
-  if (RequiresCollectiveSymmetricMemorySpace(user)) {
-    return true;
-  }
-
-  // Check custom calls with operands_memory_spaces attribute
-  if (UsesCollectiveMemorySpaceFrontendAttr(use)) {
-    return true;
-  }
-
-  return false;
-}
-
-bool RequiresCollectiveOutput(const HloValue* value, const DebugOptions& opts) {
-  HloInstruction* def = value->defining_instruction();
-
-  // Handle standard non-fusion/fusion collectives under NCCL user buffers
-  if (IsCollective(def) &&
-      (opts.xla_gpu_enable_nccl_user_buffers() ||
-       IsNcclSymmetricBuffersEnabledForCollective(def, opts))) {
-    return true;
-  }
-  // Handle one-shot RaggedAllToAll with NCCL enabled.
-  if (IsOneShotRaggedAllToAllWithNcclEnabled(opts)) {
-    // Defining Instruction is RA2A or AsyncDone(RA2A)
-    if (IsRaggedAllToAllOrAsyncDoneRaggedAllToAll(def)) {
-      return true;
-    }
-  }
-
-  // If the instruction requires symmetric memory, its output value is S1.
-  // If the output value is live-out of the module, we must convert is to S0.
-  if (RequiresCollectiveSymmetricMemorySpace(def)) {
-    return true;
-  }
-
-  // Check custom calls with results_memory_spaces attribute
-  if (DefinesCollectiveMemorySpaceFrontendAttr(value)) {
-    return true;
-  }
-
-  return false;
-}
-
-}  // namespace
-
 void GpuCollectiveBufferAnalysis(
     HloModule* module, const HloAliasAnalysis& alias_analysis,
     std::function<void(HloInstruction*, const ShapeIndex&)> add_index_to_copy,
@@ -2606,7 +2479,7 @@ void GpuCollectiveBufferAnalysis(
         live_out_values.push_back(value);
       }
 
-      if (RequiresCollectiveOutput(value, opts)) {
+      if (RequiresCollectiveOutput(*value, opts)) {
         defined_by_collective = true;
       }
     }
diff --git a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.cc b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.cc
index d598c075274..593b2db41b9 100644
--- a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.cc
+++ b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.cc
@@ -149,39 +149,6 @@ bool IsNcclSymmetricOrUserBuffersEnabledForInstruction(
          IsNcclSymmetricBuffersEnabledForCollective(inst, option);
 }
 
-bool HasCollectiveMemoryInstruction(const HloValue& input_alias,
-                                    const DebugOptions& option) {
-  // Tuple-shaped values are pointer containers and never hold data that needs
-  // to live in collective memory. Only array sub-elements do.
-  if (input_alias.shape().IsTuple()) {
-    return false;
-  }
-  // If any use is a collective instruction, we must color the value to use
-  // collective memory space.
-  for (const HloUse& use : input_alias.GetUses()) {
-    if (IsNcclSymmetricOrUserBuffersEnabledForInstruction(use.instruction,
-                                                          option)) {
-      return true;
-    }
-  }
-  return IsNcclSymmetricOrUserBuffersEnabledForInstruction(
-      input_alias.instruction(), option);
-}
-
-bool HasSymmetricMemoryInstruction(const HloValue& input_alias) {
-  // Tuple-shaped values are pointer containers and never hold data that needs
-  // to live in collective memory. Only array sub-elements do.
-  if (input_alias.shape().IsTuple()) {
-    return false;
-  }
-  for (const HloUse& use : input_alias.GetUses()) {
-    if (RequiresCollectiveSymmetricMemorySpace(use.instruction)) {
-      return true;
-    }
-  }
-  return RequiresCollectiveSymmetricMemorySpace(input_alias.instruction());
-}
-
 // Returns the memory space requested for the given custom call use, or
 // MemorySpaceColor::kDefault if none is specified.
 static absl::StatusOr<MemorySpaceColor> GetCustomCallOperandMemorySpace(
@@ -207,29 +174,6 @@ static absl::StatusOr<MemorySpaceColor> GetCustomCallOperandMemorySpace(
   return MemorySpaceColor::kDefault;
 }
 
-// Returns true if the value is
-// 1. Used by a RaggedAllToAll as Operand(1)
-// 2. RaggedAllToAll result
-bool IsRaggedAllToAllCollectiveOperandOrResult(const HloValue& value) {
-  if (value.shape().IsTuple()) {
-    return false;
-  }
-
-  // Check if the value is DEFINED by an RA2A (the result)
-  if (IsRaggedAllToAllOrAsyncDoneRaggedAllToAll(value.defining_instruction())) {
-    return true;
-  }
-
-  // Check if the value is USED by an RA2A as the destination (Operand 1)
-  for (const HloUse& use : value.GetUses()) {
-    if (IsRaggedAllToAllOrAsyncStartRaggedAllToAll(use.instruction) &&
-        use.operand_number == 1) {
-      return true;
-    }
-  }
-  return false;
-}
-
 // Returns the memory space requested for a custom call result value, or
 // MemorySpaceColor::kDefault if none is specified.
 static absl::StatusOr<MemorySpaceColor> GetCustomCallResultMemorySpace(
@@ -256,13 +200,87 @@ static absl::StatusOr<MemorySpaceColor> GetCustomCallResultMemorySpace(
   return MemorySpaceColor::kDefault;
 }
 
+bool RequiresCollectiveInput(const HloUse& use, const DebugOptions& option) {
+  const HloInstruction* user = use.instruction;
+  if (ShapeUtil::GetSubshape(user->operand(use.operand_number)->shape(),
+                             use.operand_index)
+          .IsTuple()) {
+    return false;
+  }
+
+  // Handle standard collectives under NCCL user/symmetric buffers, or
+  // device-initiated/one-sided collectives requiring symmetric memory.
+  if (IsNcclSymmetricOrUserBuffersEnabledForInstruction(user, option) ||
+      RequiresCollectiveSymmetricMemorySpace(user)) {
+    return true;
+  }
+
+  // Handle one-shot RaggedAllToAll with NCCL enabled: operand 1 (the output
+  // buffer) requires collective memory.
+  if (IsOneShotRaggedAllToAllWithNcclEnabled(option) &&
+      IsRaggedAllToAllOrAsyncStartRaggedAllToAll(user) &&
+      use.operand_number == 1) {
+    return true;
+  }
+
+  // Check custom calls with operands_memory_spaces attribute.
+  absl::StatusOr<MemorySpaceColor> operand_ms =
+      GetCustomCallOperandMemorySpace(use);
+  return operand_ms.ok() && *operand_ms == MemorySpaceColor::kCollective;
+}
+
+bool RequiresCollectiveOutput(const HloValue& value,
+                              const DebugOptions& option) {
+  // Tuple-shaped values are pointer containers and never hold data that needs
+  // to live in collective memory. Only array sub-elements do.
+  if (value.shape().IsTuple()) {
+    return false;
+  }
+
+  const HloInstruction* def = value.defining_instruction();
+
+  // Handle standard collectives under NCCL user/symmetric buffers, or
+  // device-initiated/one-sided collectives requiring symmetric memory.
+  if (IsNcclSymmetricOrUserBuffersEnabledForInstruction(def, option) ||
+      RequiresCollectiveSymmetricMemorySpace(def)) {
+    return true;
+  }
+
+  // Handle one-shot RaggedAllToAll with NCCL enabled.
+  if (IsOneShotRaggedAllToAllWithNcclEnabled(option) &&
+      IsRaggedAllToAllOrAsyncDoneRaggedAllToAll(def)) {
+    return true;
+  }
+
+  // Check custom calls with results_memory_spaces attribute.
+  absl::StatusOr<MemorySpaceColor> result_ms =
+      GetCustomCallResultMemorySpace(value);
+  return result_ms.ok() && *result_ms == MemorySpaceColor::kCollective;
+}
+
 namespace {
+
+bool HasCollectiveMemoryInstruction(const HloValue& input_alias,
+                                    const DebugOptions& option) {
+  // Tuple-shaped values are pointer containers and never hold data that needs
+  // to live in collective memory. Only array sub-elements do.
+  if (input_alias.shape().IsTuple()) {
+    return false;
+  }
+  if (RequiresCollectiveOutput(input_alias, option)) {
+    return true;
+  }
+  for (const HloUse& use : input_alias.GetUses()) {
+    if (RequiresCollectiveInput(use, option)) {
+      return true;
+    }
+  }
+  return false;
+}
+
 // Determines the memory space color for the given HLO buffer
 absl::StatusOr<BufferValue::Color> DetermineBufferColor(
     const HloBuffer& buffer, const DebugOptions& option) {
-  // Is one-shot RaggedAllToAll with NCCL feature is enabled.
-  const bool is_one_shot_ra2a_with_nccl =
-      IsOneShotRaggedAllToAllWithNcclEnabled(option);
   // Collect Color Candidates
   absl::InlinedVector<BufferValue::Color, 4> candidates;
   for (const HloValue* value : buffer.values()) {
@@ -303,17 +321,7 @@ absl::StatusOr<BufferValue::Color> DetermineBufferColor(
     }
 
     // Collective Candidates
-    if (is_one_shot_ra2a_with_nccl &&
-        IsRaggedAllToAllCollectiveOperandOrResult(*value)) {
-      // One-shot RaggedAllToAll with NCCL requires collective memory for
-      // both operand 1 and the result.
-      candidates.push_back(
-          static_cast<BufferValue::Color>(MemorySpaceColor::kCollective));
-    } else if (HasSymmetricMemoryInstruction(*value)) {
-      // Device-initiated and one-sided collectives require symmetric memory.
-      candidates.push_back(
-          static_cast<BufferValue::Color>(MemorySpaceColor::kCollective));
-    } else if (HasCollectiveMemoryInstruction(*value, option)) {
+    if (HasCollectiveMemoryInstruction(*value, option)) {
       candidates.push_back(
           static_cast<BufferValue::Color>(MemorySpaceColor::kCollective));
     }
diff --git a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.h b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.h
index 8433f555458..78308c86510 100644
--- a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.h
+++ b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment.h
@@ -25,6 +25,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/layout.h"
 #include "xla/service/buffer_assignment.h"
+#include "xla/service/hlo_value.h"
 #include "xla/xla.pb.h"
 
 namespace xla::gpu {
@@ -67,6 +68,14 @@ ParseIndexMemorySpacePairs(absl::string_view str);
 // buffers registered with the collective runtime ahead of time.
 bool RequiresCollectiveSymmetricMemorySpace(const HloInstruction* inst);
 
+// Returns true if the given operand use requires collective memory space.
+bool RequiresCollectiveInput(const HloUse& use, const DebugOptions& option);
+
+// Returns true if the given HLO value definition requires collective memory
+// space.
+bool RequiresCollectiveOutput(const HloValue& value,
+                              const DebugOptions& option);
+
 // Creates a buffer colorer that assigns memory space colors to HLO values
 // during buffer assignment. It handles:
 //  - Collective operations (all-reduce, all-gather, etc.) → kCollective
diff --git a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment_test.cc b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment_test.cc
index e2605cbe591..1a6595ea006 100644
--- a/third_party/xla/xla/service/gpu/gpu_memory_space_assignment_test.cc
+++ b/third_party/xla/xla/service/gpu/gpu_memory_space_assignment_test.cc
@@ -29,8 +29,10 @@ limitations under the License.
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/hlo/testlib/verified_hlo_module.h"
 #include "xla/service/buffer_assignment.h"
+#include "xla/service/hlo_buffer.h"
 #include "xla/service/hlo_module_config.h"
 #include "xla/service/hlo_value.h"
+#include "xla/shape_util.h"
 #include "xla/tsl/lib/core/status_test_util.h"
 #include "xla/tsl/platform/statusor.h"
 #include "xla/xla.pb.h"
@@ -40,7 +42,9 @@ namespace xla::gpu {
 namespace {
 
 using ::testing::Eq;
+using ::testing::IsEmpty;
 using ::testing::IsTrue;
+using ::testing::Not;
 using ::testing::NotNull;
 using ::testing::SizeIs;
 
@@ -689,5 +693,200 @@ INSTANTIATE_TEST_SUITE_P(
                         : "mosaic_does_not_use_collective_metadata";
     });
 
+TEST_F(GpuMemorySpaceAssignmentTest,
+       SymmetricMemoryCollectivesModeColorsInputsAndOutputs) {
+  constexpr absl::string_view kHloModule = R"(
+    HloModule m, replica_count=2
+
+    async_comp {
+      param_0 = f32[512]{0} parameter(0)
+      ROOT ag = f32[1024]{0} all-gather(param_0),
+        dimensions={0},
+        replica_groups={{0,1}},
+        backend_config={"collective_backend_config":{
+          "collectives_mode":"COLLECTIVES_SYMMETRIC_MEMORY"
+        }}
+    }
+
+    ENTRY main {
+      p0 = f32[512]{0} parameter(0)
+      ag_start = ((f32[512]{0}), f32[1024]{0}) async-start(p0),
+        calls=async_comp,
+        backend_config={"collective_backend_config":{
+          "collectives_mode":"COLLECTIVES_SYMMETRIC_MEMORY"
+        }}
+      ROOT ag_done = f32[1024]{0} async-done(ag_start),
+        backend_config={"collective_backend_config":{
+          "collectives_mode":"COLLECTIVES_SYMMETRIC_MEMORY"
+        }}
+    }
+  )";
+
+  HloModuleConfig config = GetModuleConfigForTest();
+  BufferAssigner::Colorer colorer = CreateColorer(config.debug_options());
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(kHloModule, config));
+  AliasInfo alias_info;
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloAliasAnalysis> alias_analysis,
+                       HloAliasAnalysis::Run(module.get(), &alias_info));
+  DependencyHloOrdering ordering(module.get());
+  EXPECT_OK(colorer(alias_analysis.get(), ordering));
+
+  auto get_color = [&](absl::string_view name) {
+    const HloInstruction* instr =
+        module->entry_computation()->GetInstructionWithName(name);
+    const HloBuffer& buffer = alias_analysis->GetUniqueBufferAt(instr);
+    EXPECT_EQ(buffer.values().size(), 1);
+    EXPECT_TRUE(buffer.values()[0]->has_color());
+    return buffer.values()[0]->color();
+  };
+
+  EXPECT_EQ(get_color("p0"), static_cast<int>(MemorySpaceColor::kCollective));
+  EXPECT_EQ(get_color("ag_done"),
+            static_cast<int>(MemorySpaceColor::kCollective));
+
+  const HloInstruction* p0 =
+      module->entry_computation()->GetInstructionWithName("p0");
+  const HloInstruction* ag_done =
+      module->entry_computation()->GetInstructionWithName("ag_done");
+  const HloValue& p0_val =
+      *alias_analysis->GetUniqueBufferAt(p0).values().front();
+  const HloValue& ag_done_val =
+      *alias_analysis->GetUniqueBufferAt(ag_done).values().front();
+  ASSERT_THAT(p0_val.GetUses(), Not(IsEmpty()));
+  for (const HloUse& use : p0_val.GetUses()) {
+    EXPECT_TRUE(RequiresCollectiveInput(use, config.debug_options()));
+  }
+  EXPECT_FALSE(RequiresCollectiveOutput(p0_val, config.debug_options()));
+  EXPECT_TRUE(RequiresCollectiveOutput(ag_done_val, config.debug_options()));
+
+  // Tuple-shaped pointer containers on ag_start must remain in kDefault.
+  const HloInstruction* ag_start =
+      module->entry_computation()->GetInstructionWithName("ag_start");
+  const HloBuffer& tuple_buffer =
+      alias_analysis->GetUniqueBufferAt(ag_start, /*index=*/{});
+  const HloValue& tuple_val = *tuple_buffer.values()[0];
+  EXPECT_EQ(tuple_val.color(), static_cast<int>(MemorySpaceColor::kDefault));
+  EXPECT_FALSE(RequiresCollectiveOutput(tuple_val, config.debug_options()));
+  ASSERT_THAT(tuple_val.GetUses(), SizeIs(1));
+  EXPECT_FALSE(
+      RequiresCollectiveInput(tuple_val.GetUses()[0], config.debug_options()));
+}
+
+TEST_F(GpuMemorySpaceAssignmentTest,
+       OneShotRaggedAllToAllWithNcclColorsOnlyOperand1AndOutput) {
+  constexpr absl::string_view kHloModule = R"(
+    HloModule m, replica_count=2
+
+    ENTRY main {
+      input = f32[16]{0} parameter(0)
+      output = f32[16]{0} parameter(1)
+      input_offsets = s64[2]{0} parameter(2)
+      send_sizes = s64[2]{0} parameter(3)
+      output_offsets = s64[2]{0} parameter(4)
+      recv_sizes = s64[2]{0} parameter(5)
+      ROOT ra2a = f32[16]{0} ragged-all-to-all(
+          input, output, input_offsets, send_sizes, output_offsets, recv_sizes),
+        replica_groups={{0,1}}
+    }
+  )";
+
+  HloModuleConfig config = GetModuleConfigForTest();
+  DebugOptions debug_options = config.debug_options();
+  debug_options
+      .set_xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl(true);
+  config.set_debug_options(debug_options);
+  BufferAssigner::Colorer colorer = CreateColorer(config.debug_options());
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(kHloModule, config));
+  AliasInfo alias_info;
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloAliasAnalysis> alias_analysis,
+                       HloAliasAnalysis::Run(module.get(), &alias_info));
+  DependencyHloOrdering ordering(module.get());
+  EXPECT_OK(colorer(alias_analysis.get(), ordering));
+
+  EXPECT_EQ(FindColorByName(*alias_analysis, "input"),
+            static_cast<int>(MemorySpaceColor::kDefault));
+  EXPECT_EQ(FindColorByName(*alias_analysis, "output"),
+            static_cast<int>(MemorySpaceColor::kCollective));
+  EXPECT_EQ(FindColorByName(*alias_analysis, "input_offsets"),
+            static_cast<int>(MemorySpaceColor::kDefault));
+  EXPECT_EQ(FindColorByName(*alias_analysis, "send_sizes"),
+            static_cast<int>(MemorySpaceColor::kDefault));
+  EXPECT_EQ(FindColorByName(*alias_analysis, "output_offsets"),
+            static_cast<int>(MemorySpaceColor::kDefault));
+  EXPECT_EQ(FindColorByName(*alias_analysis, "recv_sizes"),
+            static_cast<int>(MemorySpaceColor::kDefault));
+  EXPECT_EQ(FindColorByName(*alias_analysis, "ra2a"),
+            static_cast<int>(MemorySpaceColor::kCollective));
+}
+
+TEST_F(GpuMemorySpaceAssignmentTest,
+       RequiresCollectiveInputAndOutputDistinguishesColorsAndTuples) {
+  constexpr absl::string_view kHloModule = R"(
+    HloModule m
+
+    ENTRY main {
+      p0 = f32[16]{0} parameter(0)
+      p1 = f32[16]{0} parameter(1)
+      ROOT cc = (f32[16]{0}, f32[16]{0}) custom-call(p0, p1),
+        custom_call_target="my_custom_call",
+        frontend_attributes={
+          operands_memory_spaces="{0:7,1:2}",
+          results_memory_spaces="{0:7,1:2}"
+        }
+    }
+  )";
+
+  HloModuleConfig config = GetModuleConfigForTest();
+  BufferAssigner::Colorer colorer = CreateColorer(config.debug_options());
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(kHloModule, config));
+  AliasInfo alias_info;
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloAliasAnalysis> alias_analysis,
+                       HloAliasAnalysis::Run(module.get(), &alias_info));
+  DependencyHloOrdering ordering(module.get());
+  EXPECT_OK(colorer(alias_analysis.get(), ordering));
+
+  const HloInstruction* p0 =
+      module->entry_computation()->GetInstructionWithName("p0");
+  const HloInstruction* p1 =
+      module->entry_computation()->GetInstructionWithName("p1");
+  const HloInstruction* cc =
+      module->entry_computation()->GetInstructionWithName("cc");
+
+  const HloValue& p0_val =
+      *alias_analysis->GetUniqueBufferAt(p0).values().front();
+  const HloValue& p1_val =
+      *alias_analysis->GetUniqueBufferAt(p1).values().front();
+  ASSERT_THAT(p0_val.GetUses(), SizeIs(1));
+  ASSERT_THAT(p1_val.GetUses(), SizeIs(1));
+
+  EXPECT_TRUE(
+      RequiresCollectiveInput(p0_val.GetUses()[0], config.debug_options()));
+  EXPECT_FALSE(
+      RequiresCollectiveInput(p1_val.GetUses()[0], config.debug_options()));
+  EXPECT_EQ(p0_val.color(), static_cast<int>(MemorySpaceColor::kCollective));
+  EXPECT_EQ(p1_val.color(), static_cast<int>(MemorySpaceColor::kTempBuffer));
+
+  const HloValue& cc_tuple_val =
+      *alias_analysis->GetUniqueBufferAt(cc, /*index=*/{}).values().front();
+  const HloValue& cc_out0_val =
+      *alias_analysis->GetUniqueBufferAt(cc, /*index=*/{0}).values().front();
+  const HloValue& cc_out1_val =
+      *alias_analysis->GetUniqueBufferAt(cc, /*index=*/{1}).values().front();
+
+  EXPECT_FALSE(RequiresCollectiveOutput(cc_tuple_val, config.debug_options()));
+  EXPECT_TRUE(RequiresCollectiveOutput(cc_out0_val, config.debug_options()));
+  EXPECT_FALSE(RequiresCollectiveOutput(cc_out1_val, config.debug_options()));
+  EXPECT_EQ(cc_tuple_val.color(), static_cast<int>(MemorySpaceColor::kDefault));
+  EXPECT_EQ(cc_out0_val.color(),
+            static_cast<int>(MemorySpaceColor::kCollective));
+  EXPECT_EQ(cc_out1_val.color(),
+            static_cast<int>(MemorySpaceColor::kTempBuffer));
+}
+
 }  // namespace
 }  // namespace xla::gpu

From f0a307c3fdd65ab4b1d62a4e181bbf9759270e47 Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Mon, 12 Oct 2026 12:12:13 -0700
Subject: [PATCH 0580/1000] Support more than 5 dimensions in the oneDNN
 softmax kernel

MklSoftmaxOp only maps 1 to 5 dimensions to a oneDNN memory format and
failed with "Input dims must be <= 5 and >=1" otherwise, so on CPU with
oneDNN enabled, tf.nn.softmax of a tensor of rank 6 or more failed.
Softmax normalizes along the last dimension, so for more than 5
dimensions, collapse the leading ones into one and use the 2-D format.
---
 tensorflow/core/kernels/mkl/mkl_softmax_op.cc     | 15 ++++++++++++---
 .../python/kernel_tests/nn_ops/softmax_op_test.py |  8 ++++++++
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/kernels/mkl/mkl_softmax_op.cc b/tensorflow/core/kernels/mkl/mkl_softmax_op.cc
index 4bee6cb0a87..c58684ecdcf 100644
--- a/tensorflow/core/kernels/mkl/mkl_softmax_op.cc
+++ b/tensorflow/core/kernels/mkl/mkl_softmax_op.cc
@@ -227,10 +227,19 @@ class MklSoftmaxOp : public OpKernel {
   void Compute(OpKernelContext* context) override {
     try {
       const Tensor& src_tensor = context->input(0);
-      auto src_shape = src_tensor.shape();
+      TensorShape src_shape = src_tensor.shape();
+      // The formats below go up to 5 dimensions. Softmax normalizes along the
+      // last dimension, so collapse the others into one for more dimensions.
+      if (src_shape.dims() > 5) {
+        int64_t num_rows = 1;
+        for (int i = 0; i < src_shape.dims() - 1; ++i) {
+          num_rows *= src_shape.dim_size(i);
+        }
+        src_shape =
+            TensorShape({num_rows, src_shape.dim_size(src_shape.dims() - 1)});
+      }
       const int input_dims = src_shape.dims();
       memory::format_tag src_fmt;
-      // TODO(intel-tf): Add support for dimensions larger than 5.
       switch (input_dims) {
         case 1:
           src_fmt = memory::format_tag::a;
@@ -249,7 +258,7 @@ class MklSoftmaxOp : public OpKernel {
           break;
         default:
           OP_REQUIRES_OK(context,
-                         absl::AbortedError("Input dims must be <= 5 and >=1"));
+                         absl::AbortedError("Input dims must be >= 1"));
           return;
       }
 
diff --git a/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py b/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py
index 21190d6072c..1eb898cf4f7 100644
--- a/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py
@@ -276,6 +276,14 @@ class SoftmaxTest(test.TestCase):
         use_gpu=False)
     self._testOverflow(use_gpu=False)
 
+  def testHighRank(self):
+    # Regression test for GitHub issue 128814: the oneDNN kernel rejected
+    # inputs of more than 5 dimensions.
+    for rank in (6, 7):
+      features = np.random.rand(*([2] * (rank - 1) + [3])).astype(np.float32)
+      self._testSoftmax(features)
+      self._testSoftmax(features, dim=1)
+
   def testShapeInference(self):
     op = nn_ops.softmax([[[1., 1., 1., 1.], [1., 2., 3., 4.]],
                          [[2., 3., 4., 5.], [6., 7., 8., 9.]],

From 91b5a0ea5eeb59ea584e3e6b8be1bb58917f9114 Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Mon, 12 Oct 2026 12:21:18 -0700
Subject: [PATCH 0581/1000] Fix XLA Range with a dynamic start or a negative
 delta

When start or limit is dynamic, Range sizes its output from bounds of
them, but always took start's lower bound and limit's upper bound,
which only give the longest range when counting up. With a negative
delta the output was too short, often empty, as for
tf.range(n - 1, -1, -1). It also computed the values from start's
bound, so a dynamic start shifted them, as for tf.range(n, 5).

Take the bounds that give the longest range for delta's sign, and add
the input start to delta * iota when start is dynamic. A static start
builds the same HLO as before.
---
 tensorflow/compiler/tests/BUILD               |  1 +
 tensorflow/compiler/tests/ternary_ops_test.py | 29 +++++++++++++++
 .../compiler/tf2xla/kernels/sequence_ops.cc   | 35 ++++++++++++++-----
 3 files changed, 56 insertions(+), 9 deletions(-)

diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD
index deb8ce4f21f..a7f77a8198f 100644
--- a/tensorflow/compiler/tests/BUILD
+++ b/tensorflow/compiler/tests/BUILD
@@ -1872,6 +1872,7 @@ tf_xla_py_strict_test(
     deps = [
         ":xla_test",
         "//tensorflow/python:extra_py_tests_deps",
+        "//tensorflow/python/eager:def_function",
         "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/framework:test_lib",
         "//tensorflow/python/ops:array_ops",
diff --git a/tensorflow/compiler/tests/ternary_ops_test.py b/tensorflow/compiler/tests/ternary_ops_test.py
index 8e3dad56d9f..1d97c18df6f 100644
--- a/tensorflow/compiler/tests/ternary_ops_test.py
+++ b/tensorflow/compiler/tests/ternary_ops_test.py
@@ -19,6 +19,7 @@ import numpy as np
 import scipy.special as sps
 
 from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import def_function
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
@@ -76,6 +77,34 @@ class TernaryOpsTest(xla_test.XLATestCase, parameterized.TestCase):
           expected=np.array([1, 3, 5], dtype=dtype),
       )
 
+  def _rangeOfMaskedSize(self, args):
+    # n is the number of true entries in mask, a size only known at run time,
+    # and args maps it to the range's start, limit and delta.
+    @def_function.function(jit_compile=True)
+    def f(mask):
+      indices = array_ops.reshape(array_ops.where(mask), [-1])
+      n = array_ops.size(array_ops.gather(math_ops.range(6), indices))
+      return math_ops.range(*args(n))
+
+    with self.session() as session:
+      with self.device_scope():
+        mask = array_ops.placeholder(dtypes.bool, [6])
+        output = f(mask)
+      return session.run(output,
+                         {mask: [True, False, True, False, True, False]})
+
+  def testRangeDynamicBounds(self):
+    # Regression test for GitHub issue 128950: with a negative delta, the
+    # output was sized from the wrong bounds of start and limit, and with a
+    # dynamic start, its values began at start's bound. Here n is 3, bound 6.
+    self.assertAllEqual(
+        self._rangeOfMaskedSize(lambda n: (n - 1, -1, -1)), [2, 1, 0])
+    self.assertAllEqual(
+        self._rangeOfMaskedSize(lambda n: (6, n, -1)), [6, 5, 4])
+    self.assertAllEqual(self._rangeOfMaskedSize(lambda n: (n, 5, 1)), [3, 4])
+    self.assertAllEqual(
+        self._rangeOfMaskedSize(lambda n: (n, -n, -2)), [3, 1, -1])
+
   def testSelect(self):
     for dtype in self.numeric_types:
       self._testTernary(
diff --git a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc
index cad1be9fc97..a9998e90e0a 100644
--- a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc
@@ -40,7 +40,8 @@ limitations under the License.
 namespace tensorflow {
 namespace {
 
-// The type-specific part of the implementation of Range.
+// The type-specific part of the implementation of Range: checks the arguments
+// and returns delta * [0, 1, ..., size - 1] for the range's size.
 template <typename T>
 absl::StatusOr<xla::XlaOp> CreateRangeTensor(
     const xla::LiteralSlice& start_literal,
@@ -80,10 +81,9 @@ absl::StatusOr<xla::XlaOp> CreateRangeTensor(
         std::ceil(std::abs((limit_f - start_f) / delta_f)));
   }
 
-  return xla::ConstantR0(builder, start) +
-         xla::ConstantR0(builder, delta) *
-             xla::Iota(builder, xla::primitive_util::NativeToPrimitiveType<T>(),
-                       size);
+  return xla::ConstantR0(builder, delta) *
+         xla::Iota(builder, xla::primitive_util::NativeToPrimitiveType<T>(),
+                   size);
 }
 
 class RangeOp : public XlaOpKernel {
@@ -104,11 +104,24 @@ class RangeOp : public XlaOpKernel {
                 errors::InvalidArgument("delta must be a scalar, not shape ",
                                         delta_in_shape.DebugString()));
     xla::Literal start, limit, delta;
-    OP_REQUIRES_OK(ctx, ctx->ConstantInput(
-                            0, &start, xla::ValueInferenceMode::kLowerBound));
-    OP_REQUIRES_OK(ctx, ctx->ConstantInput(
-                            1, &limit, xla::ValueInferenceMode::kUpperBound));
     OP_REQUIRES_OK(ctx, ctx->ConstantInput(2, &delta));
+    // Size the output from the bounds of start and limit that give the longest
+    // range: the lowest start and highest limit when counting up, and the
+    // highest start and lowest limit when counting down.
+    const xla::PrimitiveType delta_type = delta.shape().element_type();
+    const bool count_down =
+        xla::primitive_util::IsFloatingPointType(delta_type)
+            ? delta.GetAsDouble({}).value_or(0) < 0
+            : xla::primitive_util::IsSignedIntegralType(delta_type) &&
+                  delta.GetIntegralAsS64({}).value_or(0) < 0;
+    OP_REQUIRES_OK(ctx, ctx->ConstantInput(
+                            0, &start,
+                            count_down ? xla::ValueInferenceMode::kUpperBound
+                                       : xla::ValueInferenceMode::kLowerBound));
+    OP_REQUIRES_OK(ctx, ctx->ConstantInput(
+                            1, &limit,
+                            count_down ? xla::ValueInferenceMode::kLowerBound
+                                       : xla::ValueInferenceMode::kUpperBound));
 
     DataType type = input_type(0);
     absl::StatusOr<xla::XlaOp> output;
@@ -165,6 +178,10 @@ class RangeOp : public XlaOpKernel {
     bool limit_is_dynamic = false;
     OP_REQUIRES_OK(ctx,
                    ctx->ResolveInputDynamismIntoPred(1, &limit_is_dynamic));
+    // If start is dynamic, its literal is a bound, so add the input instead.
+    output = (start_is_dynamic ? ctx->Input(0)
+                               : xla::ConstantLiteral(ctx->builder(), start)) +
+             output.value();
 
     if (start_is_dynamic || limit_is_dynamic) {
       xla::XlaOp delta = ctx->Input(2);

From 2f416f5ca239f7e23a2c28726889832541e18172 Mon Sep 17 00:00:00 2001
From: Sevin Varoglu <svaroglu@nvidia.com>
Date: Mon, 12 Oct 2026 07:43:01 -0700
Subject: [PATCH 0582/1000] PR #50089: [XLA:GPU] ragged_all_to_all dev_kernel:
 fan remote GIN puts across contexts
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50089

📝 Summary of Changes
`RaggedAllToAllDeviceKernelImpl` today hard-codes `gin_context = 0` and routes every remote-peer put through it — funneling the entire remote payload through a single fabric queue on multi-rail hardware while other rails sit idle.

This PR routes each CTA to its own GIN context by `blockIdx.x % dev_comm.ginContextCount`, mirroring NCCL's own single-shot symmetric collectives (e.g. `all_gather_gin.cuh`). Also swaps the exit synchronization from per-signal `waitSignal` (which does not carry across contexts) to a `Put`-fenced release barrier that fences all bound contexts on every rank.

🎯 Justification
On multi-rail hardware, one context limits bandwidth to one QP, regardless of CTA count. Using ginContextCount contexts lets traffic use the other rails. This also requires replacing waitSignal with a Put fence because a context’s signal storage cannot see writes from other contexts.

🚀 Kind of Contribution
⚡️ Performance Improvement

📊 Benchmark (for Performance Improvements)
N/A

🧪 Unit Tests:
N/A

🧪 Execution Tests:
N/A
Copybara import of the project:

--
6a8bef0d8130afc31171424674f04dbec03f4340 by Sevin Varoglu <svaroglu@nvidia.com>:

[XLA:GPU] ragged_all_to_all dev_kernel: fan remote GIN puts across contexts

Merging this change closes #50089

PiperOrigin-RevId: 997913414
---
 .../ragged_all_to_all_device_kernel_lib.cu.h  | 19 ++++---------------
 1 file changed, 4 insertions(+), 15 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/gpu/ragged_all_to_all_device_kernel_lib.cu.h b/third_party/xla/xla/stream_executor/gpu/ragged_all_to_all_device_kernel_lib.cu.h
index 0e581fb55fc..2fbc76f4e6f 100644
--- a/third_party/xla/xla/stream_executor/gpu/ragged_all_to_all_device_kernel_lib.cu.h
+++ b/third_party/xla/xla/stream_executor/gpu/ragged_all_to_all_device_kernel_lib.cu.h
@@ -177,9 +177,6 @@ __global__ void __launch_bounds__(512) RaggedAllToAllDeviceKernelImpl(
     const int64_t* __restrict__ output_offsets_ptr,
     int64_t num_updates_per_replica, int64_t num_row_elements,
     int64_t input_buffer_offset_bytes, int64_t output_buffer_offset_bytes) {
-  // NCCL device barrier/GIN APIs emit scope-qualified atomics that require
-  // sm_60+. Lower architectures compile to an empty stub; the kernel is only
-  // launched when the device supports NCCL device comms.
 #if __CUDA_ARCH__ >= 600
   ncclTeam world = ncclTeamWorld(dev_comm);
   ncclTeam lsa = ncclTeamLsa(dev_comm);
@@ -189,12 +186,11 @@ __global__ void __launch_bounds__(512) RaggedAllToAllDeviceKernelImpl(
   const bool has_remote_peers = (lsa_size < num_ranks);
 
   if (has_remote_peers) {
-    const int gin_context = 0;
+    const int n_ctxs = static_cast<int>(dev_comm.ginContextCount);
+    const int my_ctx = static_cast<int>(blockIdx.x) % n_ctxs;
     const unsigned int signal_index = 0;
 
-    ncclGin gin{dev_comm, gin_context};
-    uint64_t signal_value =
-        (blockIdx.x == 0) ? gin.readSignal(signal_index) : 0;
+    ncclGin gin{dev_comm, my_ctx};
 
     ncclBarrierSession<ncclCoopCta> bar{ncclCoopCta(), ncclTeamTagWorld(), gin,
                                         blockIdx.x};
@@ -207,16 +203,9 @@ __global__ void __launch_bounds__(512) RaggedAllToAllDeviceKernelImpl(
         input_buffer_offset_bytes, output_buffer_offset_bytes, start_lsa,
         lsa_size, num_ranks, &gin, world, signal_index);
 
-    const int num_remote_peers =
-        (num_ranks - lsa_size) * num_updates_per_replica;
-    if (blockIdx.x == 0) {
-      gin.waitSignal(ncclCoopCta(), signal_index,
-                     signal_value + num_remote_peers);
-    }
-
     gin.flush(ncclCoopCta());
     bar.sync(ncclCoopCta(), ::cuda::memory_order_release,
-             ncclGinFenceLevel::Relaxed);
+             ncclGinFenceLevel::Put | ncclGinFenceLevel::Get);
   } else {
     ncclLsaBarrierSession<ncclCoopCta> bar{ncclCoopCta(), dev_comm,
                                            ncclTeamTagLsa{}, blockIdx.x};

From 2f1c8803c8191db280d6c97068d7c7aaa89e67f2 Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Mon, 12 Oct 2026 07:44:20 -0700
Subject: [PATCH 0583/1000] Fix `AlgebraicSimplifierVisitor::HandleRemainder`
 on `iota % N` and `(iota + N) % N` to reject negative divisors and wrapping
 narrow iota shapes, and use `GetFirstInteger()` for effective scalar
 divisors.

PiperOrigin-RevId: 997914125
---
 .../simplifiers/algebraic_simplifier.cc       | 31 +++++++------
 .../simplifiers/algebraic_simplifier_test.cc  | 44 +++++++++++++++++++
 2 files changed, 62 insertions(+), 13 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
index 81e651adcc9..40e39ee6c29 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
@@ -6668,14 +6668,17 @@ absl::Status AlgebraicSimplifierVisitor::HandleRemainder(
   if (Match(remainder,
             m::Remainder(m::Iota(&iota),
                          m::Broadcast(m::ConstantEffectiveScalar(&divisor))))) {
-    // The iota counts {0, ..., iota_upper_bound - 1}.  (Actually this is
-    // conservative; the iota may overflow and count up to a smaller value than
-    // this.  But that's OK for our purposes here.)
+    // The iota counts {0, ..., iota_upper_bound - 1}. Require that
+    // iota_upper_bound - 1 fits in the element type, so the iota never wraps
+    // to negative values, and that N is positive, since X % N for negative N
+    // is not the identity.
     int64_t iota_upper_bound = iota->shape().dimensions(
         Cast<HloIotaInstruction>(iota)->iota_dimension());
-    std::optional<int64_t> divisor_val = divisor->literal().GetIntegralAsS64(
-        std::vector<int64_t>(0, divisor->shape().dimensions().size()));
-    if (divisor_val && *divisor_val >= iota_upper_bound) {
+    std::optional<int64_t> divisor_val = divisor->literal().GetFirstInteger();
+    if (divisor_val && *divisor_val > 0 && *divisor_val >= iota_upper_bound &&
+        (iota_upper_bound == 0 ||
+         primitive_util::FitsInIntegralType(iota_upper_bound - 1,
+                                            iota->shape().element_type()))) {
       return ReplaceInstruction(remainder, iota);
     }
   }
@@ -6696,17 +6699,19 @@ absl::Status AlgebraicSimplifierVisitor::HandleRemainder(
                              m::Broadcast(m::ConstantEffectiveScalar(&addend))),
               m::Broadcast(&bcast, m::ConstantEffectiveScalar(&divisor)))) &&
       addend == divisor) {
-    // The iota counts {0, ...iota_upper_bound - 1}, with the same caveat above
-    // that iota_upper_bound is conservative, and the true upper bound may be
-    // smaller.
+    // The iota counts {0, ..., iota_upper_bound - 1}. As above, require a
+    // positive N and an iota that does not wrap. Then the largest value of
+    // X + N is N + iota_upper_bound - 1, which must also fit in the element
+    // type.
     int64_t iota_upper_bound = iota->shape().dimensions(
         Cast<HloIotaInstruction>(iota)->iota_dimension());
-    std::optional<int64_t> divisor_val = divisor->literal().GetIntegralAsS64(
-        std::vector<int64_t>(0, divisor->shape().dimensions().size()));
-    if (divisor_val) {
+    std::optional<int64_t> divisor_val = divisor->literal().GetFirstInteger();
+    if (divisor_val && *divisor_val > 0 && iota_upper_bound > 0 &&
+        primitive_util::FitsInIntegralType(iota_upper_bound - 1,
+                                           iota->shape().element_type())) {
       // Check whether divisor_val + iota_upper_bound - 1 overflows.
       std::optional<int64_t> max_val =
-          OverflowSafeAdd(*divisor_val, iota_upper_bound);
+          OverflowSafeAdd(*divisor_val, iota_upper_bound - 1);
       if (max_val.has_value() && primitive_util::FitsInIntegralType(
                                      *max_val, iota->shape().element_type())) {
         return ReplaceWithNewInstruction(
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
index e1b15ed9059..73d8601dc7f 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
@@ -11541,6 +11541,50 @@ TEST_F(AlgebraicSimplifierTest, RemainderOfNPlusIotaOverflow) {
   ASSERT_FALSE(AlgebraicSimplifier(default_options_).Run(m.get()).value());
 }
 
+TEST_F(AlgebraicSimplifierTest, RemainderOfNegativeNPlusIotaNotSimplified) {
+  constexpr absl::string_view kNegativeNModuleStr = R"(
+    HloModule m
+    test {
+      iota = s32[5] iota(), iota_dimension=0
+      neg_three = s32[] constant(-3)
+      bcast = s32[5] broadcast(neg_three), dimensions={}
+      sum = s32[5] add(iota, bcast)
+      ROOT remainder = s32[5] remainder(sum, bcast)
+    })";
+  ASSERT_OK_AND_ASSIGN(auto m1,
+                       ParseAndReturnVerifiedModule(kNegativeNModuleStr));
+  EXPECT_FALSE(AlgebraicSimplifier(default_options_).Run(m1.get()).value());
+
+  constexpr absl::string_view kNarrowIotaModuleStr = R"(
+    HloModule m
+    test {
+      iota = s8[200] iota(), iota_dimension=0
+      n = s8[] constant(-100)
+      bcast = s8[200] broadcast(n), dimensions={}
+      sum = s8[200] add(iota, bcast)
+      ROOT remainder = s8[200] remainder(sum, bcast)
+    })";
+  ASSERT_OK_AND_ASSIGN(auto m2,
+                       ParseAndReturnVerifiedModule(kNarrowIotaModuleStr));
+  EXPECT_FALSE(AlgebraicSimplifier(default_options_).Run(m2.get()).value());
+
+  constexpr absl::string_view kRank1EffectiveScalarModuleStr = R"(
+    HloModule m
+    test {
+      iota = s32[5,1] iota(), iota_dimension=0
+      five = s32[1] constant({5})
+      five_bcast = s32[5,1] broadcast(five), dimensions={1}
+      ROOT remainder = s32[5,1] remainder(iota, five_bcast)
+    })";
+  ASSERT_OK_AND_ASSIGN(
+      auto m3, ParseAndReturnVerifiedModule(kRank1EffectiveScalarModuleStr));
+  AlgebraicSimplifierOptions options = default_options_;
+  options.set_enable_broadcast_degenerate_dimension(false);
+  ASSERT_TRUE(AlgebraicSimplifier(options).Run(m3.get()).value());
+  EXPECT_THAT(m3->entry_computation()->root_instruction(),
+              GmockMatch(m::Iota()));
+}
+
 TEST_F(AlgebraicSimplifierTest, RepeatedRemainder) {
   constexpr absl::string_view kModuleStr = R"(
     HloModule m

From fb7e962fc32ae0cdde91d33f2cbdb19a6dd3fd38 Mon Sep 17 00:00:00 2001
From: Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>
Date: Mon, 12 Oct 2026 07:45:00 -0700
Subject: [PATCH 0584/1000] PR #50546: [ROCM] removing rocm_smi from MORI deps

Imported from GitHub PR https://github.com/openxla/xla/pull/50546

MORI resolves smi library at runtime therefore this dependency can be dropped. This is also needed for therock10.1 support
Copybara import of the project:

--
890b4d6124b6a582f3b4c22882f9890c7de37329 by Pavel Emeliyanenko <pavel.emeliyanenko@amd.com>:

removing rocm_smi from MORI deps since now it resolves smi lib at runtime dynamically

Merging this change closes #50546

PiperOrigin-RevId: 997914445
---
 third_party/xla/third_party/roc_mori/src_application.BUILD | 2 --
 1 file changed, 2 deletions(-)

diff --git a/third_party/xla/third_party/roc_mori/src_application.BUILD b/third_party/xla/third_party/roc_mori/src_application.BUILD
index 7adc03693f0..331a34cd6e4 100644
--- a/third_party/xla/third_party/roc_mori/src_application.BUILD
+++ b/third_party/xla/third_party/roc_mori/src_application.BUILD
@@ -41,8 +41,6 @@ cc_library(
         "@roc_mori//:mori_shmem_headers",
         # CMake hip::host: libamdhip64.so + HIP host headers.
         "@local_config_rocm//rocm:hip",
-        # CMake find_library(ROCM_SMI_LIB rocm_smi64): librocm_smi64.so.
-        "@local_config_rocm//rocm:rocm_smi",
         "@local_config_rocm//rocm:hsa_runtime",
         "@local_config_rocm//rocm:hsakmt",
         # CMake ibverbs: system libibverbs.so (rdma-core).

From 8f5719d1c2cfbc9f88384ceecdb7f3db52cc7a41 Mon Sep 17 00:00:00 2001
From: Kingston Mandisodza <kingstonm@google.com>
Date: Mon, 12 Oct 2026 08:05:04 -0700
Subject: [PATCH 0585/1000] [XLA:SPMD] Key WindowReshardCacheEntry on pad_value
 and masking flags.

Include `pad_value`, `mask_invalid_region`, and `force_mask_in_compact` in `WindowReshardCacheEntry` (matching `pad_value` by pointer equality or equal constant literal), and forward `mask_invalid_region` and `force_mask_in_compact` in recursive `ReshardAsWindowedInput` calls so windowed operations on the same operand with different fill values or masking requirements do not reuse each other's padded halo exchange.

PiperOrigin-RevId: 997924359
---
 .../xla/xla/service/spmd/spmd_partitioner.cc  | 19 ++--
 .../xla/xla/service/spmd/spmd_partitioner.h   | 12 ++-
 .../xla/service/spmd/spmd_partitioner_test.cc | 93 +++++++++++++++++++
 3 files changed, 114 insertions(+), 10 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.cc b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
index b4e9dc67d53..9dff289d687 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
@@ -879,14 +879,17 @@ PartitionedHlo::ReshardAsWindowedInput(const Window& window,
           ? HloSharding::V3ToV2Sharding(raw_target.named_sharding())
           : raw_target;
   for (auto& entry : cache) {
-    if (std::get<0>(entry) == target &&
-        protobuf_util::HaveSameSerialization(std::get<1>(entry), window)) {
-      return std::get<2>(entry);
+    if (entry.mask_invalid_region == mask_invalid_region &&
+        entry.force_mask_in_compact == force_mask_in_compact &&
+        entry.sharding == target && entry.pad_value->Identical(*pad_value) &&
+        protobuf_util::HaveSameSerialization(entry.window, window)) {
+      return entry.result;
     }
   }
   auto update_cache = [&](WindowedInputShardReturnValue result) {
-    cache.emplace_back(target, window, std::move(result));
-    return std::get<2>(cache.back());
+    cache.push_back({target, window, pad_value, mask_invalid_region,
+                     force_mask_in_compact, std::move(result)});
+    return cache.back().result;
   };
   VLOG(2) << "ReshardAsWindowedInput()\n"
           << "\twindow:" << window_util::ToString(window)
@@ -1262,7 +1265,8 @@ PartitionedHlo::ReshardAsWindowedInput(const Window& window,
   }
   if (target != sharding() &&
       sharding_with_windowed_dims_replicated != sharding()) {
-    return Reshard(target).ReshardAsWindowedInput(window, target, pad_value);
+    return Reshard(target).ReshardAsWindowedInput(
+        window, target, pad_value, mask_invalid_region, force_mask_in_compact);
   }
   if (Product(trimmed_target_sharding_tile_shape) == 1) {
     // The trimmed sharding may have just one shard left. We can simply return
@@ -1370,7 +1374,8 @@ PartitionedHlo::ReshardAsWindowedInput(const Window& window,
         return handle_all_windowed_dimensions_are_replicated();
       }
       return Reshard(sharding_with_windowed_dims_replicated)
-          .ReshardAsWindowedInput(window, target, pad_value);
+          .ReshardAsWindowedInput(window, target, pad_value,
+                                  mask_invalid_region, force_mask_in_compact);
     }
     visiting_hlo = *resharded;
   }
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.h b/third_party/xla/xla/service/spmd/spmd_partitioner.h
index 8ecaec55fa4..265dffbe507 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner.h
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner.h
@@ -449,11 +449,17 @@ class PartitionedHlo {
   };
   // A cache for resharding each partitioned HLO.
   struct ReshardCache {
+    struct WindowReshardCacheEntry {
+      HloSharding sharding;
+      Window window;
+      const HloInstruction* pad_value;
+      bool mask_invalid_region;
+      bool force_mask_in_compact;
+      WindowedInputShardReturnValue result;
+    };
     struct PerHloCache {
       absl::flat_hash_map<HloSharding, PartitionedHlo> reshard_cache;
-      std::vector<
-          std::tuple<HloSharding, Window, WindowedInputShardReturnValue>>
-          window_reshard_cache;
+      std::vector<WindowReshardCacheEntry> window_reshard_cache;
     };
     // Use absl::node_hash_map for pointer stability.
     absl::node_hash_map<HloInstruction*, PerHloCache> per_hlo_cache;
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
index eaffda4c64c..017a60bbcf4 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
@@ -3092,6 +3092,99 @@ ENTRY entry {
                           op::Shape("f32[27,11]")));
 }
 
+TEST_P(SpmdPartitioningTest,
+       PadsWithDifferentFillValuesDoNotShareWindowReshardCache) {
+  absl::string_view hlo_string = R"(
+HloModule module
+
+ENTRY entry {
+  %param0 = f32[8] parameter(0), sharding={devices=[4]<=[4]}
+  %zero = f32[] constant(0)
+  %five = f32[] constant(5)
+  %pad0 = f32[12] pad(%param0, %zero), padding=2_2, sharding={devices=[4]<=[4]}
+  %pad5 = f32[12] pad(%param0, %five), padding=2_2, sharding={devices=[4]<=[4]}
+  ROOT %tuple = (f32[12], f32[12]) tuple(%pad0, %pad5),
+    sharding={{devices=[4]<=[4]}, {devices=[4]<=[4]}}
+})";
+
+  ASSERT_OK_AND_ASSIGN(auto module,
+                       PartitionComputation(hlo_string, /*num_devices=*/4));
+  VLOG(1) << module->ToString();
+
+  const auto* root = module->entry_computation()->root_instruction();
+  auto zero = op::Constant(LiteralUtil::CreateR0<float>(0));
+  auto five = op::Constant(LiteralUtil::CreateR0<float>(5));
+
+  EXPECT_THAT(
+      root,
+      op::Tuple(
+          op::DynamicSlice(op::Pad(op::Concatenate(), zero), _),
+          op::Select(_, op::DynamicSlice(op::Pad(op::Concatenate(), five), _),
+                     op::Broadcast(five))));
+}
+
+TEST_P(SpmdPartitioningTest,
+       DifferentMaskInvalidRegionDoNotShareWindowReshardCache) {
+  absl::string_view hlo_string = R"(
+HloModule module
+
+add {
+  lhs = f32[] parameter(0)
+  rhs = f32[] parameter(1)
+  ROOT add = f32[] add(lhs, rhs)
+}
+
+ENTRY entry {
+  %param0 = f32[6,1] parameter(0), sharding={devices=[4,1]<=[4]}
+  %reshape = f32[3,2] reshape(%param0), sharding={devices=[4,1]<=[4]}
+  %zero = f32[] constant(0)
+  %rw = f32[8,1] reduce-window(%param0, %zero), window={size=1x1 pad=0_2x0_0},
+    to_apply=add, sharding={devices=[4,1]<=[4]}
+  ROOT %tuple = (f32[3,2], f32[8,1]) tuple(%reshape, %rw),
+    sharding={{devices=[4,1]<=[4]}, {devices=[4,1]<=[4]}}
+})";
+
+  ASSERT_OK_AND_ASSIGN(auto module,
+                       PartitionComputation(hlo_string, /*num_devices=*/4));
+  VLOG(1) << module->ToString();
+
+  const auto* root = module->entry_computation()->root_instruction();
+  auto zero = op::Constant(LiteralUtil::CreateR0<float>(0));
+
+  EXPECT_THAT(root, op::Tuple(op::Reshape(op::Parameter(0)),
+                              op::ReduceWindow(op::Select(_, op::Parameter(0),
+                                                          op::Broadcast(zero)),
+                                               zero)));
+}
+
+TEST_P(SpmdPartitioningTest,
+       DifferentForceMaskInCompactDoNotShareWindowReshardCache) {
+  absl::string_view hlo_string = R"(
+HloModule module
+
+ENTRY entry {
+  %param0 = f32[24,1] parameter(0), sharding={devices=[8,1]<=[8]}
+  %reshape = f32[3,8] reshape(%param0), sharding={devices=[8,1]<=[8]}
+  %zero = f32[] constant(0)
+  %pad = f32[64,1] pad(%param0, %zero), padding=0_40x0_0,
+    sharding={devices=[8,1]<=[8]}
+  ROOT %tuple = (f32[3,8], f32[64,1]) tuple(%reshape, %pad),
+    sharding={{devices=[8,1]<=[8]}, {devices=[8,1]<=[8]}}
+})";
+
+  ASSERT_OK_AND_ASSIGN(auto module,
+                       PartitionComputation(hlo_string, /*num_devices=*/8));
+  VLOG(1) << module->ToString();
+
+  const auto* root = module->entry_computation()->root_instruction();
+  auto zero = op::Constant(LiteralUtil::CreateR0<float>(0));
+
+  EXPECT_THAT(root,
+              op::Tuple(op::Reshape(op::DynamicSlice(op::Concatenate(), _, _)),
+                        op::Select(_, op::DynamicSlice(op::Concatenate(), _, _),
+                                   op::Broadcast(zero))));
+}
+
 TEST_P(SpmdPartitioningTest, SliceAlongNonPartitionedDimension) {
   absl::string_view hlo_string = R"(
 HloModule module

From 5080ef597c41f7d39385bd8b3fc3ab8b66c894e0 Mon Sep 17 00:00:00 2001
From: Levon Ter-Grigoryan <patrios@google.com>
Date: Mon, 12 Oct 2026 08:17:48 -0700
Subject: [PATCH 0586/1000] [XLA:GPU] Skip cross-device barrier when a single
 device RaggedAllToAll is used.

Also add test for this.

PiperOrigin-RevId: 997930903
---
 .../xla/xla/backends/gpu/runtime/BUILD        |  14 ++
 .../gpu/runtime/ragged_all_to_all_thunk.cc    |  13 +-
 .../ragged_all_to_all_thunk_multigpu_test.cc  | 138 ++++++++++++++++++
 3 files changed, 161 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index ad6d968a9af..f2d11fbc751 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -2453,21 +2453,35 @@ xla_test(
     },
     backends = ["gpu"],
     deps = [
+        ":collective_clique_requests",
+        ":collective_memory_requests",
+        ":collective_params",
         ":collective_thunk",
         ":collective_thunk_multigpu_test_utils",
         ":ragged_all_to_all_thunk",
         ":thunk",
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",
+        "//xla:xla_proto_cc",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/testlib:verified_hlo_module",
+        "//xla/runtime:device_id",
         "//xla/service:buffer_assignment",
         "//xla/service:device_assignment",
+        "//xla/service:hlo_module_config",
+        "//xla/service:service_executable_run_options",
         "//xla/service:shaped_slice",
         "//xla/service/gpu:buffer_allocations",
+        "//xla/service/gpu:gpu_executable_run_options",
         "//xla/stream_executor:device_address",
         "//xla/stream_executor:stream",
+        "//xla/stream_executor:stream_executor_address_allocator",
+        "//xla/stream_executor:stream_executor_h",
+        "@com_google_absl//absl/base",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest_main",
     ],
diff --git a/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc
index 2c907a0b4a0..4c3570d6262 100644
--- a/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc
@@ -601,7 +601,8 @@ RaggedAllToAllThunk::GetCliqueRequirements(const GpuCliqueKey& clique_key,
     const int core_count = params.executor->GetDeviceDescription().core_count();
     clique_reqs.dev_comm = DeviceKernelLsaDevCommRequirements(core_count);
   }
-  if (use_multi_gpu_barrier_with_nccl_in_one_shot_kernel()) {
+  if (use_multi_gpu_barrier_with_nccl_in_one_shot_kernel() &&
+      clique_key.num_devices() > 1) {
     clique_reqs.barrier_reqs = CollectiveCliqueRequests::BarrierRequirements();
     clique_reqs.barrier_reqs->use_cross_device_barrier = true;
   }
@@ -649,7 +650,8 @@ absl::Status RaggedAllToAllThunk::Initialize(const InitializeParams& params) {
               "Failed to allocate output offsets buffer.");
         }
 
-        if (use_multi_gpu_barrier_with_nccl_in_one_shot_kernel() ||
+        if ((use_multi_gpu_barrier_with_nccl_in_one_shot_kernel() &&
+             state->clique_key.num_devices() > 1) ||
             config_.use_device_kernel) {
           if (config_.fast_interconnect_slice_size_override.has_value()) {
             state->lsa_size =
@@ -715,7 +717,8 @@ absl::Status RaggedAllToAllThunk::Initialize(const InitializeParams& params) {
   //      it's safe to cache the symmetric handler.
   if ((!state->lsa_size.has_value() ||
        state->lsa_size.value() != state->clique_key.num_devices() ||
-       !config_.use_multi_gpu_barrier_with_nccl_in_one_shot_kernel) &&
+       !config_.use_multi_gpu_barrier_with_nccl_in_one_shot_kernel ||
+       state->clique_key.num_devices() < 2) &&
       is_local(params.local_device_count)) {
     // Rendezvous - Exchange output pointers and barrier signal buffers.
     ABSL_ASSIGN_OR_RETURN(
@@ -1025,6 +1028,7 @@ absl::Status RaggedAllToAllThunk::RunCollective(const ExecuteParams& params,
 
   if (is_one_shot_kernel_enabled() && fabric_safe) {
     if (use_multi_gpu_barrier_with_nccl_in_one_shot_kernel() &&
+        clique_key.num_devices() > 1 &&
         IsRaggedAllToAllWithSymmetricMemoryKernelSupported(
             config_.config.operand_element_type[0]) &&
         state->lsa_size.has_value() &&
@@ -1141,7 +1145,8 @@ RendezvousResources(int device_ordinal, RankId rank,
 
 absl::Status RaggedAllToAllThunk::PrepareCollective(
     const PrepareParams& params, const GpuCliqueKey& clique_key) {
-  if (config_.use_multi_gpu_barrier_with_nccl_in_one_shot_kernel ||
+  if ((config_.use_multi_gpu_barrier_with_nccl_in_one_shot_kernel &&
+       clique_key.num_devices() > 1) ||
       use_symmetric_memory()) {
     // Request symmetric memory for the output buffer only.
     const Buffer& output_buffer = buffers()[1];
diff --git a/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk_multigpu_test.cc b/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk_multigpu_test.cc
index b8eba79c2ec..f2032530d65 100644
--- a/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk_multigpu_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk_multigpu_test.cc
@@ -24,22 +24,38 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/base/casts.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
 #include "absl/types/span.h"
+#include "xla/backends/gpu/runtime/collective_clique_requests.h"
+#include "xla/backends/gpu/runtime/collective_memory_requests.h"
+#include "xla/backends/gpu/runtime/collective_params.h"
 #include "xla/backends/gpu/runtime/collective_thunk.h"
 #include "xla/backends/gpu/runtime/collective_thunk_multigpu_test_utils.h"
 #include "xla/backends/gpu/runtime/ragged_all_to_all_thunk.h"
 #include "xla/backends/gpu/runtime/thunk.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/hlo/testlib/verified_hlo_module.h"
+#include "xla/runtime/device_id.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/device_assignment.h"
 #include "xla/service/gpu/buffer_allocations.h"
+#include "xla/service/gpu/gpu_executable_run_options.h"
+#include "xla/service/hlo_module_config.h"
+#include "xla/service/service_executable_run_options.h"
 #include "xla/service/shaped_slice.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/stream_executor/device_address.h"
 #include "xla/stream_executor/stream.h"
+#include "xla/stream_executor/stream_executor.h"
+#include "xla/stream_executor/stream_executor_address_allocator.h"
+#include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
 
 namespace xla::gpu {
@@ -334,5 +350,127 @@ TEST(RaggedAllToAllThunkMultiGpuTest, RecordCommandBufferUpdate) {
   }));
 }
 
+TEST(RaggedAllToAllThunkMultiGpuTest,
+     SingletonRaggedAllToAllDoesNotRequestBarrierOrSymmetricMemory) {
+  if (!HasEnoughGpus(kNumDevices)) {
+    GTEST_SKIP() << "Test requires at least " << kNumDevices << " GPUs";
+  }
+
+  constexpr absl::string_view kModuleStr = R"hlo(
+HloModule module, replica_count=2
+
+ENTRY entry {
+  input = bf16[4,16] parameter(0)
+  output = bf16[4,16] parameter(1)
+  input_offsets = s64[1] parameter(2)
+  send_sizes = s64[1] parameter(3)
+  output_offsets = s64[1] parameter(4)
+  recv_sizes = s64[1] parameter(5)
+  ROOT ra2a = bf16[4,16] ragged-all-to-all(
+      input, output, input_offsets, send_sizes, output_offsets, recv_sizes),
+      replica_groups={{0},{1}}
+}
+)hlo";
+
+  HloModuleConfig config;
+  DebugOptions debug_options = config.debug_options();
+  debug_options.set_xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel(
+      true);
+  debug_options
+      .set_xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl(true);
+  config.set_debug_options(debug_options);
+  config.set_replica_count(kNumDevices);
+
+  VerifiedHloModule module("test_module", config,
+                           /*verifier_layout_sensitive=*/false,
+                           /*allow_mixed_precision_in_hlo_verifier=*/false,
+                           ShapeUtil::ByteSizeOfElements);
+  ASSERT_OK(module.ParseHloStringAndVerifyModule(kModuleStr));
+
+  const HloInstruction* root = module.entry_computation()->root_instruction();
+  ASSERT_EQ(root->opcode(), HloOpcode::kRaggedAllToAll);
+  const auto* ra2a_instr =
+      absl::down_cast<const HloRaggedAllToAllInstruction*>(root);
+
+  std::vector<BufferAllocation> allocations;
+  allocations.reserve(ra2a_instr->operand_count());
+  for (int i = 0; i < ra2a_instr->operand_count(); ++i) {
+    allocations.emplace_back(
+        /*index=*/i, ShapeUtil::ByteSizeOf(ra2a_instr->operand(i)->shape()),
+        /*color=*/0);
+  }
+
+  std::vector<CollectiveThunk::Buffer> buffers;
+  buffers.reserve(ra2a_instr->operand_count());
+  for (int i = 0; i < ra2a_instr->operand_count(); ++i) {
+    const Shape& shape = ra2a_instr->operand(i)->shape();
+    int64_t byte_size = ShapeUtil::ByteSizeOf(shape);
+    ShapedSlice slice{BufferAllocation::Slice(&allocations[i], 0, byte_size),
+                      shape};
+    buffers.push_back(CollectiveThunk::Buffer{
+        /*element_count=*/ShapeUtil::ElementsIn(shape),
+        /*source_buffer=*/slice,
+        /*destination_buffer=*/slice,
+        /*source_memory_space=*/0,
+        /*destination_memory_space=*/0,
+    });
+  }
+
+  RaggedAllToAllThunk thunk(Thunk::ThunkInfo(), ra2a_instr, std::move(buffers),
+                            /*p2p_memcpy_enabled=*/true,
+                            /*devices_per_host=*/kNumDevices);
+
+  DeviceAssignment device_assignment = MakeDeviceAssignment(kNumDevices);
+  GpuExecutableRunOptions::DeviceIdMap id_map;
+  for (int i = 0; i < kNumDevices; ++i) {
+    id_map[LocalDeviceId(i)] = GlobalDeviceId(i);
+  }
+
+  ASSERT_OK(RunOnDevices(
+      kNumDevices, "singleton_ra2a_prepare", [&](int d) -> absl::Status {
+        se::StreamExecutor* executor = GetGpuExecutor(d);
+        ABSL_ASSIGN_OR_RETURN(std::unique_ptr<se::Stream> stream,
+                         executor->CreateStream());
+        se::StreamExecutorAddressAllocator allocator(executor);
+
+        GpuExecutableRunOptions gpu_options;
+        gpu_options.set_gpu_global_device_ids(id_map);
+
+        ServiceExecutableRunOptions run_options;
+        run_options.mutable_run_options()->set_stream(stream.get());
+        run_options.mutable_run_options()->set_device_assignment(
+            &device_assignment);
+        run_options.mutable_run_options()->set_gpu_executable_run_options(
+            &gpu_options);
+        run_options.mutable_run_options()->set_local_device_count(kNumDevices);
+
+        ABSL_ASSIGN_OR_RETURN(
+            CollectiveParams collective_params,
+            CollectiveParams::Create(run_options, /*async_streams=*/{},
+                                     LocalDeviceId(d)));
+
+        std::vector<se::DeviceAddressBase> allocated_buffers(
+            allocations.size());
+        BufferAllocations buffer_allocs(allocated_buffers, d, &allocator);
+        CollectiveCliqueRequests clique_requests;
+        CollectiveMemoryRequests memory_requests(buffer_allocs);
+        Thunk::PrepareParams prepare_params{&collective_params,
+                                            &clique_requests, &memory_requests,
+                                            executor, &buffer_allocs};
+
+        ABSL_RETURN_IF_ERROR(thunk.Prepare(prepare_params));
+
+        std::vector<CollectiveCliqueRequests::CliqueRequest> requested_cliques =
+            clique_requests.OrderedRequestedCliques();
+        EXPECT_EQ(requested_cliques.size(), 1);
+        if (!requested_cliques.empty()) {
+          EXPECT_EQ(requested_cliques[0].key.num_devices(), 1);
+          EXPECT_FALSE(requested_cliques[0].use_cross_device_barrier_requested);
+        }
+        EXPECT_EQ(memory_requests.symmetric_size(), 0);
+        return absl::OkStatus();
+      }));
+}
+
 }  // namespace
 }  // namespace xla::gpu

From bb5b68d067cf3cd38364303587a1944356c72dda Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Mon, 12 Oct 2026 08:21:24 -0700
Subject: [PATCH 0587/1000] [XLA] Preserve conversions and scalar constants in
 cumsum HandleReduceWindow

PiperOrigin-RevId: 997932633
---
 .../simplifiers/algebraic_simplifier.cc       |  44 +++++---
 .../simplifiers/algebraic_simplifier_test.cc  | 105 ++++++++++++++++++
 2 files changed, 132 insertions(+), 17 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
index 40e39ee6c29..5a52d6f43c2 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
@@ -9657,9 +9657,18 @@ absl::Status AlgebraicSimplifierVisitor::HandleReduceWindow(
   const Window& window = reduce_window->window();
 
   // Optimize cumsum-like reduce-window on broadcasted constants.
+  //
+  // A bitcast that changes the element type reinterprets bits, and may change
+  // the element count when the bit widths differ. It does not preserve values,
+  // so we do not look through it.
+  auto is_same_type_bitcast = [](const HloInstruction* instr) {
+    return instr->opcode() == HloOpcode::kBitcast &&
+           instr->shape().element_type() ==
+               instr->operand(0)->shape().element_type();
+  };
   HloInstruction* real_operand = operand;
   while (real_operand->opcode() == HloOpcode::kReshape ||
-         real_operand->opcode() == HloOpcode::kBitcast ||
+         is_same_type_bitcast(real_operand) ||
          real_operand->opcode() == HloOpcode::kConvert ||
          real_operand->opcode() == HloOpcode::kSlice ||
          real_operand->opcode() == HloOpcode::kCopy) {
@@ -9674,7 +9683,7 @@ absl::Status AlgebraicSimplifierVisitor::HandleReduceWindow(
     HloInstruction* val_const = real_operand->mutable_operand(0);
     while (val_const->opcode() == HloOpcode::kConvert ||
            val_const->opcode() == HloOpcode::kReshape ||
-           val_const->opcode() == HloOpcode::kBitcast) {
+           is_same_type_bitcast(val_const)) {
       val_const = val_const->mutable_operand(0);
     }
 
@@ -9682,15 +9691,6 @@ absl::Status AlgebraicSimplifierVisitor::HandleReduceWindow(
       return absl::OkStatus();
     }
 
-    // In case the source operand is not a scalar, we reshape it to a scalar
-    // before we attempt to broadcast it.
-    if (!ShapeUtil::IsScalar(val_const->shape())) {
-      TF_RET_CHECK(ShapeUtil::IsEffectiveScalar(val_const->shape()));
-      val_const = reduce_window->AddInstruction(HloInstruction::CreateReshape(
-          ShapeUtil::MakeScalarShape(val_const->shape().element_type()),
-          val_const));
-    }
-
     int64_t reduction_dim = -1;
     bool valid_pattern = true;
 
@@ -9726,13 +9726,23 @@ absl::Status AlgebraicSimplifierVisitor::HandleReduceWindow(
     }
 
     if (valid_pattern && reduction_dim != -1) {
-      if (val_const->shape().element_type() !=
-          reduce_window->shape().element_type()) {
-        ABSL_ASSIGN_OR_RETURN(Literal dest_literal,
-                         val_const->literal().Convert(
-                             reduce_window->shape().element_type()));
+      TF_RET_CHECK(ShapeUtil::IsEffectiveScalar(val_const->shape()));
+      absl::InlinedVector<PrimitiveType, 4> convert_types;
+      for (const HloInstruction* i = operand; i != val_const;
+           i = i->operand(0)) {
+        if (i->opcode() == HloOpcode::kConvert) {
+          convert_types.push_back(i->shape().element_type());
+        }
+      }
+      if (!ShapeUtil::IsScalar(val_const->shape()) || !convert_types.empty()) {
+        Literal scalar =
+            LiteralUtil::GetFirstScalarLiteral(val_const->literal());
+        for (auto it = convert_types.rbegin(); it != convert_types.rend();
+             ++it) {
+          ABSL_ASSIGN_OR_RETURN(scalar, scalar.Convert(*it));
+        }
         val_const = reduce_window->AddInstruction(
-            HloInstruction::CreateConstant(std::move(dest_literal)));
+            HloInstruction::CreateConstant(std::move(scalar)));
       }
 
       // Create iota for the reduction dimension.
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
index 73d8601dc7f..c9f2b454c59 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
@@ -1995,6 +1995,111 @@ TEST_F(AlgebraicSimplifierTest,
   EXPECT_EQ(root->opcode(), HloOpcode::kMultiply);
 }
 
+TEST_F(AlgebraicSimplifierTest,
+       ReduceWindowCumSumBroadcastOfConstantWithNarrowingAndNonScalarConvert) {
+  constexpr absl::string_view kNarrowingModuleStr = R"(
+    HloModule m
+    add_f32 {
+      p0 = f32[] parameter(0)
+      p1 = f32[] parameter(1)
+      ROOT a = f32[] add(p0, p1)
+    }
+
+    ENTRY test {
+      c = f32[1] constant({1.5})
+      c_scalar = f32[] reshape(c)
+      c_s32 = s32[] convert(c_scalar)
+      c_f32 = f32[] convert(c_s32)
+      b = f32[10] broadcast(c_f32), dimensions={}
+      c_zero = f32[] constant(0.0)
+      ROOT rw = f32[10] reduce-window(b, c_zero), window={size=10 pad=9_0}, to_apply=add_f32
+    }
+  )";
+
+  ASSERT_OK_AND_ASSIGN(auto m1,
+                       ParseAndReturnVerifiedModule(kNarrowingModuleStr));
+  AlgebraicSimplifierOptions options = default_options_;
+  ASSERT_THAT(AlgebraicSimplifier(options).Run(m1.get()),
+              absl_testing::IsOkAndHolds(true));
+  EXPECT_THAT(
+      m1->entry_computation()->root_instruction(),
+      GmockMatch(m::Multiply(m::Broadcast(m::ConstantScalar(1.0f)), m::Op())));
+
+  constexpr absl::string_view kNonScalarConvertModuleStr = R"(
+    HloModule m
+    add_bf16 {
+      p0 = bf16[] parameter(0)
+      p1 = bf16[] parameter(1)
+      ROOT a = bf16[] add(p0, p1)
+    }
+
+    ENTRY test {
+      c = f32[1] constant({1.0})
+      c_scalar = f32[] reshape(c)
+      conv = bf16[] convert(c_scalar)
+      b = bf16[10] broadcast(conv), dimensions={}
+      c_zero = bf16[] constant(0.0)
+      ROOT rw = bf16[10] reduce-window(b, c_zero), window={size=10 pad=9_0}, to_apply=add_bf16
+    }
+  )";
+
+  ASSERT_OK_AND_ASSIGN(
+      auto m2, ParseAndReturnVerifiedModule(kNonScalarConvertModuleStr));
+  ASSERT_THAT(AlgebraicSimplifier(options).Run(m2.get()),
+              absl_testing::IsOkAndHolds(true));
+  EXPECT_THAT(m2->entry_computation()->root_instruction(),
+              GmockMatch(m::Multiply(m::Broadcast(m::Constant()), m::Op())));
+}
+
+// A bitcast that changes the element type reinterprets bits and may change the
+// element count, so the cumsum rewrite must not look through it.
+TEST_F(AlgebraicSimplifierTest,
+       ReduceWindowCumSumBroadcastOfConstantWithElementTypeChangingBitcast) {
+  constexpr absl::string_view kBitcastBelowBroadcastModuleStr = R"(
+    HloModule m
+    add_f32 {
+      p0 = f32[] parameter(0)
+      p1 = f32[] parameter(1)
+      ROOT a = f32[] add(p0, p1)
+    }
+
+    ENTRY test {
+      c = s32[] constant(1065353216)
+      c_f32 = f32[] bitcast(c)
+      b = f32[10] broadcast(c_f32), dimensions={}
+      c_zero = f32[] constant(0.0)
+      ROOT rw = f32[10] reduce-window(b, c_zero), window={size=10 pad=9_0}, to_apply=add_f32
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(
+      auto m1, ParseAndReturnVerifiedModule(kBitcastBelowBroadcastModuleStr));
+  ASSERT_OK(AlgebraicSimplifier(default_options_).Run(m1.get()));
+  EXPECT_THAT(m1->entry_computation()->root_instruction(),
+              GmockMatch(m::ReduceWindow()));
+
+  constexpr absl::string_view kBitcastAboveBroadcastModuleStr = R"(
+    HloModule m
+    add_s16 {
+      p0 = s16[] parameter(0)
+      p1 = s16[] parameter(1)
+      ROOT a = s16[] add(p0, p1)
+    }
+
+    ENTRY test {
+      c = s32[] constant(1)
+      b = s32[20] broadcast(c), dimensions={}
+      b_s16 = s16[40] bitcast(b)
+      c_zero = s16[] constant(0)
+      ROOT rw = s16[40] reduce-window(b_s16, c_zero), window={size=40 pad=39_0}, to_apply=add_s16
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(
+      auto m2, ParseAndReturnVerifiedModule(kBitcastAboveBroadcastModuleStr));
+  ASSERT_OK(AlgebraicSimplifier(default_options_).Run(m2.get()));
+  EXPECT_THAT(m2->entry_computation()->root_instruction(),
+              GmockMatch(m::ReduceWindow()));
+}
+
 // Test that Const + A is canonicalized to A + Const.
 TEST_F(AlgebraicSimplifierTest, AddConstOnLHS) {
   auto m = CreateNewVerifiedModule();

From 67b4ceade9f13ba8d999e8e4da108b50b178d1d6 Mon Sep 17 00:00:00 2001
From: Ilia Sergachev <isergachev@nvidia.com>
Date: Mon, 12 Oct 2026 08:21:43 -0700
Subject: [PATCH 0588/1000] PR #50548: Rubin upstreaming part two
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50548

📝 Summary of Changes

Second part of changes enabling Vera Rubin NVL72 support.

🎯 Justification

New hardware enablement.

🚀 Kind of Contribution

✨ New Feature, 🧪 Tests

📊 Benchmark (for Performance Improvements)

no

🧪 Unit Tests:

yes

🧪 Execution Tests:

no
Copybara import of the project:

--
5a8537f24435aefc31f604876f724701a4cfea96 by Andrei Ivanov <anivanov@nvidia.com>:

[XLA:GPU] Add VR NVL72 collective performance profiles

--
d646c9a111d90b19bbd01cf0ba4eb3a9f0ed1ddc by Andrei Ivanov <anivanov@nvidia.com>:

[XLA:GPU] Add SM107 matmul performance data

--
980c0c01e0b8aed1dedf420169a5b9a093010677 by Andrei Ivanov <anivanov@nvidia.com>:

[XLA:GPU] Add VR NVL72 CUB scratch size data

--
e1effdb1df794f7ebfd4a8147f00e9d1838f888b by Ilia Sergachev <isergachev@nvidia.com>:

[GPU] Add VR NVL72 oversized shared memory per block value

Merging this change closes #50548

PiperOrigin-RevId: 997932787
---
 ...cub_scratch_size_deviceless_lookup_test.cc |    12 +
 .../cub/data/NVIDIA_VR_NVL72.textproto        | 81640 ++++++++++++++++
 .../gpu/target_config/specs/vr_nvl72.txtpb    |     1 +
 .../service/gpu/gpu_device_info_for_tests.cc  |     1 +
 .../model/default_collective_perf_table.txtpb | 41885 ++++++++
 .../gpu/model/default_matmul_perf_table.txtpb |  3005 +
 .../gpu/model/matmul_interpolator_test.cc     |    45 +
 7 files changed, 126589 insertions(+)
 create mode 100644 third_party/xla/xla/backends/gpu/libraries/cub/data/NVIDIA_VR_NVL72.textproto

diff --git a/third_party/xla/xla/backends/gpu/libraries/cub/cub_scratch_size_deviceless_lookup_test.cc b/third_party/xla/xla/backends/gpu/libraries/cub/cub_scratch_size_deviceless_lookup_test.cc
index 88c81e2cc5f..f7c7fee3cd1 100644
--- a/third_party/xla/xla/backends/gpu/libraries/cub/cub_scratch_size_deviceless_lookup_test.cc
+++ b/third_party/xla/xla/backends/gpu/libraries/cub/cub_scratch_size_deviceless_lookup_test.cc
@@ -35,6 +35,7 @@ using ::tsl::proto_testing::ParseTextProtoOrDie;
 
 constexpr stream_executor::SemanticVersion kCubVersion1_15_0{1, 15, 0};
 constexpr stream_executor::SemanticVersion kCubVersion1_16_0{1, 16, 0};
+constexpr stream_executor::SemanticVersion kCubVersion3_2_0{3, 2, 0};
 
 TEST(CubScratchSizeDevicelessLookupTest, LookupExactMatch) {
   ASSERT_OK_AND_ASSIGN(auto lookup,
@@ -311,5 +312,16 @@ TEST(CubScratchSizeDevicelessLookupTest, CanBeLoadedFromBundledData) {
   ASSERT_OK(lookup) << lookup.status().message();
 }
 
+TEST(CubScratchSizeDevicelessLookupTest, BundledDataContainsVrNvl72) {
+  absl::StatusOr<const CubScratchSizeDevicelessLookup&> lookup =
+      CubScratchSizeDevicelessLookup::GetInstance();
+  ASSERT_OK(lookup) << lookup.status().message();
+
+  EXPECT_TRUE(lookup->CanLookup(kCubVersion3_2_0, "NVIDIA VR NVL72",
+                                /*key_type_size=*/1,
+                                /*value_type_size=*/std::nullopt,
+                                /*num_items=*/4860, /*batch_size=*/1));
+}
+
 }  // namespace
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/libraries/cub/data/NVIDIA_VR_NVL72.textproto b/third_party/xla/xla/backends/gpu/libraries/cub/data/NVIDIA_VR_NVL72.textproto
new file mode 100644
index 00000000000..a4a32932fb4
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/libraries/cub/data/NVIDIA_VR_NVL72.textproto
@@ -0,0 +1,81640 @@
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 1
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 1
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 3583
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 4607
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 5631
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 6655
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 7679
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 8703
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 9727
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 10751
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 11775
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 12799
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 13823
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 14847
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 15871
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 16895
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 17919
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 18943
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 19967
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 20991
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 22015
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 23039
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 24063
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 25087
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 26111
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 27135
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 28159
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 29183
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 30207
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 31231
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 32255
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 33279
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 34303
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 35327
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 36351
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 37375
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 38399
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 39423
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 40447
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 41471
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 42495
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 43519
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 44543
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 45567
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 46591
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 47615
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 48639
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 49663
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 50687
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 51711
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 52735
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 53759
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 54783
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 55807
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 56831
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 57855
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 58879
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 59903
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 60927
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 61951
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 62975
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 63999
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 66047
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 67071
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 68095
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 69119
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 71167
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 72191
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 73215
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 75263
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 76287
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 78335
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 79359
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 81407
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 82431
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 84479
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 85503
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 87551
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 88575
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 90623
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 92671
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 94719
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 95743
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 97791
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 99839
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 101887
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 103935
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 105983
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 108031
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 110079
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 112127
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 114175
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 116223
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 118271
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 121343
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 123391
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 125439
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 128511
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 130559
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 133631
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 135679
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 138751
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 141823
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 143871
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 146943
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 150015
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 153087
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 156159
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 159231
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 162303
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 165375
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 168447
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 171519
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 174591
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 178687
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 181759
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 185855
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 188927
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 193023
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 197119
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 200191
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 204287
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 208383
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 212479
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 216575
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 220671
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 225791
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 229887
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 233983
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 239103
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 244223
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 248319
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 253439
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 258559
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 263679
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 268799
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 273919
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 280063
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 285183
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 290303
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 296447
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 302591
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 308735
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 314879
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 321023
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 327167
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 333311
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 340479
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 346623
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 353791
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 360959
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 368127
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 375295
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 382463
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 390655
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 397823
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 406015
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 414207
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 422399
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 430591
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 438783
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 447999
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 457215
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 465407
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 474623
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 484863
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 494079
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 504319
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 513535
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 523775
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 535039
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 545279
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 556543
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 566783
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 578047
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 590335
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 601599
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 613887
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 626175
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 638463
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 650751
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 664063
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 677375
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 690687
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 703999
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 718335
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 732671
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 747007
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 762367
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 777727
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 793087
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 808447
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 824831
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 841215
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 857599
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 875007
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 892415
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 910847
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 928255
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 946687
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 966143
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 985599
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 1005055
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 1024511
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 1044991
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 1066495
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 1086975
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 1109503
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 1131007
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 1153535
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 1177087
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 1200639
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 1224191
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 1248767
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 1273343
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 1298943
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 1324543
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 1351167
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 1378815
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 1405439
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 1434111
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 1462783
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 1491455
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 1521151
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 1551871
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 1582591
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 1614335
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 1647103
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 1679871
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 1712639
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 1747455
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 1782271
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 1818111
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 1853951
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 1890815
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 1928703
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 1967615
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 2006527
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 2046463
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 2087423
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 2129407
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 2171391
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 2215423
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 2259455
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 2304511
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 2350591
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 2397695
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 2444799
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 2493951
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 2544127
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 2594303
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 2646527
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 2699775
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 2753023
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 2808319
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 2864639
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 2921983
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 2980351
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 3039743
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 3100159
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 3162623
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 3225087
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 3289599
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 3356159
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 3422719
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 3491327
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 3560959
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 3631615
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 3704319
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 3779071
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 3853823
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 3931647
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 4009471
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 4090367
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 4171263
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 4255231
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 4340223
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 4427263
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 4515327
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 4605439
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 4697599
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 4791807
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 4887039
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 4985343
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 5084671
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 5186047
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 5289471
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 5395967
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 5503487
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 5613055
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 5725695
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 5840383
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 5957119
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 6075903
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 6197759
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 6321663
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 6447615
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 6576639
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 6707711
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 6841855
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 6979071
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 7118335
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 7260671
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 7406079
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 7553535
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 7705087
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 7858687
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 8016383
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 8176127
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 8339967
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 8506879
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 8676863
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 8849919
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 9027071
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 9207295
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 9391615
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 9579007
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 9770495
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 9966079
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 10165759
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 10368511
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 10576383
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 10787327
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 11003391
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 11223551
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 11447807
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 11676159
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 11909631
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 12148223
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 12390911
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 12638719
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 12891647
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 13149695
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 13411839
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 13680127
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 13953535
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 14233087
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 14517759
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 14807551
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 15103487
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 15405567
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 15713791
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 16028159
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 16348671
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 16675327
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 17009151
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 17349119
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 17696255
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 18049535
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 18411007
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 18778623
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 19154431
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 19537407
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 19928575
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 20326911
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 20733439
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 21148159
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 21571071
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 22002175
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 22442495
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 22891007
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 23348735
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 23815679
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 24291839
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 24777215
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 25272831
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 25778687
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 26293759
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 26820095
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 27355647
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 27903487
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 28260863
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 28262911
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 28264959
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 1
+  is_segmented: true
+  scratch_size_recordings {
+    num_items: 251
+    scratch_space_bytes: 511
+  }
+  scratch_size_recordings {
+    num_items: 511
+    scratch_space_bytes: 767
+  }
+  scratch_size_recordings {
+    num_items: 762
+    scratch_space_bytes: 1023
+  }
+  scratch_size_recordings {
+    num_items: 1017
+    scratch_space_bytes: 1279
+  }
+  scratch_size_recordings {
+    num_items: 1258
+    scratch_space_bytes: 1535
+  }
+  scratch_size_recordings {
+    num_items: 1528
+    scratch_space_bytes: 1791
+  }
+  scratch_size_recordings {
+    num_items: 1787
+    scratch_space_bytes: 2047
+  }
+  scratch_size_recordings {
+    num_items: 2009
+    scratch_space_bytes: 2303
+  }
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 2559
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 2815
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 3071
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 3327
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 3583
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 3839
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 4095
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 4351
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 4607
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 4863
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 5119
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 5375
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 5631
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 5887
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 6143
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 6399
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 6655
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 6911
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 7167
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 7423
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 7679
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 7935
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 8191
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 8447
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 8703
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 8959
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 9215
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 9471
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 9727
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 9983
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 10239
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 10495
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 10751
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 11007
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 11263
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 11519
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 11775
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 12031
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 12287
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 12543
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 12799
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 13055
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 13311
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 13567
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 14079
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 14335
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 14591
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 14847
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 15103
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 15359
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 15615
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 15871
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 16383
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 16639
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 16895
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 17151
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 17663
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 17919
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 18175
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 18687
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 18943
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 19455
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 19711
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 20223
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 20479
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 20991
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 21247
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 21759
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 22271
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 22527
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 23039
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 23551
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 24063
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 24575
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 24831
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 25343
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 25855
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 26367
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 26879
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 27391
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 28159
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 28671
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 29183
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 29695
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 30207
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 30975
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 31487
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 32255
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 32767
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 33535
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 34047
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 34815
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 35327
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 36095
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 36863
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 37631
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 38399
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 39167
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 39935
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 40703
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 41471
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 42239
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 43007
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 44031
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 44799
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 45823
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 46591
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 47615
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 48639
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 49407
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 50431
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 51455
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 52479
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 53503
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 54527
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 55551
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 56831
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 57855
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 59135
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 60159
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 61439
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 62719
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 63999
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 65279
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 66559
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 67839
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 69119
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 70399
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 71935
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 73215
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 74751
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 76287
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 77823
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 79359
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 80895
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 82431
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 84223
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 85759
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 87551
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 89343
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 91135
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 92927
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 94719
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 96511
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 98559
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 100607
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 102399
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 104447
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 106751
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 108799
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 110847
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 113151
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 115455
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 117759
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 120063
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 122367
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 124927
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 127487
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 129791
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 132607
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 135167
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 137727
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 140543
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 143359
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 146175
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 149247
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 152063
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 155135
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 158207
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 161535
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 164607
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 167935
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 171263
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 174591
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 178175
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 181759
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 185343
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 189183
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 192767
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 196607
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 200703
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 204543
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 208639
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 212735
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 217087
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 221439
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 225791
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 230399
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 235007
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 239615
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 244479
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 249343
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 254207
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 259327
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 264447
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 269823
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 275199
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 280831
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 286463
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 292095
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 297983
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 303871
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 310015
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 316159
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 322303
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 328959
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 335359
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 342015
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 348927
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 355839
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 363007
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 370175
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 377599
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 385279
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 392959
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 400895
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 408831
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 417023
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 425215
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 433663
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 442367
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 451327
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 460287
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 469503
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 478975
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 488447
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 498175
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 508159
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 518399
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 528639
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 539391
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 550143
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 561151
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 572159
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 583679
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 595455
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 607231
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 619263
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 631807
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 644351
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 657151
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 670463
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 683775
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 697599
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 711423
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 725759
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 740095
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 754943
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 770047
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 785407
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 801023
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 817151
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 833535
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 850175
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 867071
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 884479
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 902143
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 920063
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 938495
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 957439
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 976383
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 996095
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 1015807
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 1036287
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 1057023
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 1078015
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 1099519
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 1121535
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 1144063
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 1166847
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 1190143
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 1213951
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 1238271
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 1263103
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 1288191
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 1314047
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 1340415
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 1367039
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 1394431
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 1422335
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 1450751
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 1479679
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 1509375
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 1539583
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 1570303
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 1601791
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 1633791
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 1666559
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 1699839
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 1733631
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 1768447
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 1803775
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 1839871
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 1876735
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 1914111
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 1952511
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 1991423
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 2031359
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 2071807
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 2113279
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 2155519
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 2198783
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 2242559
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 2287615
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 2333183
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 2380031
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 2427391
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 2476031
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 2525439
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 2576127
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 2627583
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 2680063
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 2733823
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 2788351
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 2844159
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 2900991
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 2959103
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 3018239
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 3078655
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 3140095
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 3202815
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 3267071
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 3332351
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 3398911
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 3467007
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 3536127
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 3607039
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 3678975
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 3752703
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 3827711
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 3904255
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 3982335
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 4061951
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 4143103
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 4226047
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 4310527
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 4396799
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 4484607
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 4574463
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 4665855
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 4759039
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 4854271
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 4951295
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 5050367
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 5151487
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 5254399
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 5359615
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 5466623
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 5575935
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 5687551
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 5801215
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 5917439
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 6035711
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 6156287
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 6279423
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 6405119
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 6533119
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 6663679
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 6797055
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 6932991
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 7071743
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 7213055
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 7357439
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 7504383
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 7654655
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 7807743
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 7963903
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 8123135
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 8285439
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 8451327
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 8620287
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 8792575
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 8968447
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 9147903
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 9330687
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 9517311
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 9707775
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 9901823
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 10099967
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 10301951
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 10508031
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 10717951
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 10932479
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 11151103
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 11374079
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 11601663
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 11833599
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 12070143
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 12311551
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 12557823
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 12808959
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 13065215
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 13326335
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 13593087
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 13864959
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 14142207
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 14424831
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 14713343
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 15007743
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 15307775
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 15613951
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 15926271
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 16244735
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 16569599
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 16901119
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 17239039
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 17583871
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 17935615
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 18294271
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 18660095
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 19033343
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 19414015
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 19802111
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 20198143
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 20602111
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 21014271
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 21434623
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 21863167
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 22300415
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 22746367
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 23201279
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 23665407
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 24138751
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 24621567
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 25113855
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 25616127
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 26128383
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 26651135
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 27184127
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 27727615
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 28282367
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 28847871
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 29424895
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 30013439
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 30613503
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 31225855
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 31850495
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 32487423
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 33137151
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 33799935
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 34475775
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 35165439
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 35868671
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 36585983
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 37317631
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 38064127
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 38825471
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 39601919
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 40393983
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 41201663
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 42025727
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 42866175
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 43723519
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 44598015
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 45489919
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 46399743
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 47327743
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 48274431
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 49239807
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 50224639
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 51229183
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 52253695
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 53298687
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 54364671
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 55451903
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 56560895
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 57692159
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 58845951
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 60023039
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 61223423
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 62447871
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 63696895
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 64970751
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 66270207
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 67595519
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 68947455
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 70326527
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 71732991
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 73167615
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 74630911
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 76123647
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 77646079
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 79198975
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 80782847
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 82398463
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 84046591
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 85727487
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 87441919
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 89190655
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 90974463
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 92794111
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 94649855
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 96542975
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 98473727
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 100443135
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 102452223
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 104501247
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 106591231
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 108722943
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 110897407
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 113115391
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 115377663
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 117685247
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 120038911
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 122439679
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 124888575
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 127386111
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 129933823
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 132532735
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 135183359
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 137886975
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 140644607
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 143457535
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 146326783
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 149253119
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 152238335
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 155282943
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 158388735
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 161556479
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 164787455
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 168083199
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 171444991
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 174873855
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 178371327
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 181938687
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 185577471
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 189288959
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 193074687
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 196936191
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 200875007
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 204892415
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 208990463
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 213170175
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 217433599
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 221782271
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 226217727
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 230742271
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 235356927
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 240064255
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 244865535
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 249762815
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 254757887
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 259853055
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 265050111
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 270351103
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 275758079
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 281273343
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 286898687
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 292636671
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 298489599
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 304459263
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 310548479
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 316759295
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 323094527
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 329556479
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 336147711
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 342870527
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 349727999
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 356722431
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 363856895
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 371133951
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 378556671
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 386127871
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 393850367
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 401727487
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 409762047
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 417957119
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 426316287
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 434842623
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 443539455
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 452410367
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 461458431
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 470687743
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 480101375
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 489703423
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 499497471
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 509487359
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 519677183
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 530070783
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 540671999
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 551485439
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 562515199
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 573765631
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 585240831
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 596945663
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 608884479
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 621062143
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 633483519
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 646153215
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 659076095
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 672257791
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 685702911
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 699416831
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 713405183
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 727673343
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 742226687
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 757071359
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 772212735
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 787656959
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 803410175
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 819478271
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 835867903
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 852585215
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 869636863
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 887029503
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 904770303
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 922865663
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 941322751
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 960149247
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 979352319
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 998939391
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 2
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 1
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 15615
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 15871
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 16127
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 16383
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 16639
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 16895
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 17151
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 17407
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 17663
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 17919
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 18175
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 18431
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 18687
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 18943
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 19199
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 19455
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 19711
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 19967
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 20223
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 20479
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 20735
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 20991
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 22271
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 22783
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 23039
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 23295
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 23551
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 24063
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 24319
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 24831
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 25087
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 25343
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 25855
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 26111
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 26623
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 26879
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 27391
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 27903
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 28159
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 28671
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 29183
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 29439
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 29951
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 30463
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 30975
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 31487
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 31743
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 32255
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 32767
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 33279
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 34047
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 34559
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 35071
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 35583
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 36095
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 36863
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 37375
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 38911
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 39679
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 40191
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 40959
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 41471
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 42239
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 43007
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 43519
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 44287
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 45055
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 45823
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 46591
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 47359
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 48127
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 48895
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 49663
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 50687
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 51455
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 52223
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 53247
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 55039
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 56063
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 57087
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 57855
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 58879
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 59903
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 60927
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 61951
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 62975
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 64255
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 65279
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 66303
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 67583
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 68607
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 69887
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 72191
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 73471
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 74751
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 76031
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 77311
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 78591
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 79871
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 81407
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 82943
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 84223
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 85759
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 88319
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 89855
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 91391
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 92927
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 94719
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 96255
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 98047
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 99839
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 101631
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 104447
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 106239
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 108287
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 110079
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 112127
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 114175
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 116223
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 118271
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 121343
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 123391
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 125695
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 127999
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 130303
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 132607
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 134911
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 138495
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 141055
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 143615
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 146175
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 148735
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 151551
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 155135
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 157951
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 160767
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 163839
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 166655
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 170751
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 173823
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 177151
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 180223
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 183551
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 187903
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 191487
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 194815
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 198399
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 203007
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 206847
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 210431
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 214271
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 219391
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 223231
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 227327
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 231679
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 236799
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 241151
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 245503
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 250111
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 255743
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 260351
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 265215
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 271103
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 276223
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 281087
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 287487
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 292607
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 297983
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 304639
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 310271
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 315903
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 322815
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 328703
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 335871
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 342015
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 348415
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 355839
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 362239
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 370175
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 376831
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 385023
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 391935
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 400383
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 407807
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 416255
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 423935
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 432895
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 440831
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 450047
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 458239
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 467967
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 476671
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 486399
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 495359
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 505599
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 516095
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 525567
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 536319
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 547327
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 557567
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 568831
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 580607
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 591359
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 603391
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 615679
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 626943
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 639743
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 652543
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 665855
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 679167
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 691711
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 705791
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 719871
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 734207
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 749055
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 763903
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 779263
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 794623
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 810495
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 826623
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 843007
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 859903
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 876799
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 894207
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 911871
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 929791
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 948223
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 966911
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 986111
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 1006591
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 1026303
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 1046527
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 1067007
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 1089023
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 1110271
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 1132031
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 1155327
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 1177855
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 1201919
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 1225215
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 1249279
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 1274623
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 1300479
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 1325567
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 1352447
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 1378815
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 1406463
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 1434879
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 1462783
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 1492223
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 1521919
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 1552639
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 1583615
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 1615103
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 1647359
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 1680383
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 1713663
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 1747711
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 1782527
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 1817855
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 1853951
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 1891583
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 1928959
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 1967103
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 2007039
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 2046719
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 2087935
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 2129151
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 2171903
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 2214655
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 2259199
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 2304511
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 2350591
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 2396671
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 2444543
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 2493439
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 2543103
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 2594815
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 2646527
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 2699007
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 2752511
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 2808319
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 2863871
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 2920447
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 2979327
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 3039231
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 3099135
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 3161087
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 3224575
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 3288831
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 3354623
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 3421439
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 3489535
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 3559167
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 3630847
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 3702783
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 3777279
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 3852287
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 3929343
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 4008191
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 4088319
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 4169983
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 4253183
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 4337919
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 4424191
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 4513279
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 4602879
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 4695295
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 4788223
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 4884223
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 4981759
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 5081343
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 5183487
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 5286911
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 5391871
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 5500159
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 5610239
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 5721599
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 5836031
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 5953535
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 6072063
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 6193151
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 6317311
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 6442751
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 6571775
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 6703103
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 6836735
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 6974207
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 7112959
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 7255551
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 7400703
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 7548415
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 7699199
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 7852543
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 8009983
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 8170239
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 8333567
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 8499967
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 8669695
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 8843263
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 9020159
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 9200639
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 9384191
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 9571327
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 9763071
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 9958143
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 10157055
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 10360831
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 10567167
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 10778623
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 10993919
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 11214335
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 11437823
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 11666431
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 11899391
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 12137983
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 12380671
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 12628223
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 12880127
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 13137919
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 13400575
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 13668095
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 13941759
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 14220287
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 14504447
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 14794751
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 15090431
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 15392767
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 15699711
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 16014335
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 16334079
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 16660735
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 16993535
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 17333503
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 17679871
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 18033919
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 18394367
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 18761727
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 19137023
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 19520255
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 19909631
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 20308479
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 20713727
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 21128703
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 21551359
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 21981951
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 22421759
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 22869503
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 23326975
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 23792895
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 24268543
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 24754431
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 25249023
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 25754111
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 26269695
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 26795007
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 27330047
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 27877119
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 28434431
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 29002495
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 29583103
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 30174463
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 30777855
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 31392767
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 32021247
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 32661503
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 33314815
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 33980415
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 34660351
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 35353087
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 36059903
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 36781055
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 37516543
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 38267135
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 39032575
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 39812607
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 40609023
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 41420543
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 42249215
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 43093503
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 43955711
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 44835071
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 45731071
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 46645759
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 47578623
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 48529919
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 49500159
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 50490879
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 51500799
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 52529919
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 53580287
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 54652415
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 55745535
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 56860415
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 57997055
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 59156991
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 60339967
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 61547007
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 62777599
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 64032511
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 65313791
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 66619391
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 67952383
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 69311231
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 70696703
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 72110335
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 73552383
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 75023615
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 76524287
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 78054399
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 79615999
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 81208063
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 82832127
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 84488447
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 86178047
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 87901183
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 89659903
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 91452415
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 93281791
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 95146495
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 97050111
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 98990591
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 100970495
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 102989567
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 105049599
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 107150079
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 109293311
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 111478527
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 113708031
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 115982335
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 118302207
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 120667647
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 123080703
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 125542911
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 128053503
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 130614271
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 133227007
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 135890687
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 138608639
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 141380351
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 144208383
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 147092479
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 150034687
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 153035007
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 156094719
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 159217151
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 162401023
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 165649663
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 168962303
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 172341247
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 175787519
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 179303679
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 182889727
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 186547199
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 190278143
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 194083327
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 197965055
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 201924607
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 205962495
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 210082303
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 214283263
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 218569471
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 222939903
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 227399423
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 231946751
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 236585471
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 241317119
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 246143999
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 251066111
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 256087807
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 261208831
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 266433279
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 271761407
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 277196799
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 282740479
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 288395007
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 294162687
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 300046847
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 306047231
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 312167679
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 318410751
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 324779263
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 331275263
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 337900287
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 344657663
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 351550975
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 358581759
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 365753599
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 373068031
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 380529919
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 388140031
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 395902975
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 403820543
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 411896831
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 420134911
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 428537855
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 437108223
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 445850879
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 454767359
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 463862527
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 473139455
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 482602495
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 492254463
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 502099455
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 512141567
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 522384127
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 532831999
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 543487999
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 554357247
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 565444607
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 576501247
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 587315199
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 598345727
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 609596415
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 621072383
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 632777983
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 644717567
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 656895743
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 669317631
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 681988095
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 694911999
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 708094207
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 721540095
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 735254783
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 749243903
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 763512831
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 778066943
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 792912383
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 808054783
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 823499775
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 839253759
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 855322879
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 871713279
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 888431615
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 905484287
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 922877951
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 940619519
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 958715903
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 977174271
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 996001791
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 1015205631
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 1034793727
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 1054773759
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 1075153151
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 1095940351
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 1117147135
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 1138774015
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 1160833535
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 1183334143
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 1206284543
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 1229694207
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 1253572095
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 1277927423
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 1302769919
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 1328109311
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 1353955327
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 1380318463
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 1407208703
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 1434636799
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 1462613503
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 1491149567
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 1520256511
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 1549945599
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 1580228351
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 1611116799
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 1642623231
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 1674763519
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 1707542783
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 1740977407
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 1775080703
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 1809866239
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 1845347583
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 1881538303
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 1918452735
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 1956105727
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 1994511615
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 2033685759
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 2
+  is_segmented: true
+  scratch_size_recordings {
+    num_items: 121
+    scratch_space_bytes: 511
+  }
+  scratch_size_recordings {
+    num_items: 251
+    scratch_space_bytes: 767
+  }
+  scratch_size_recordings {
+    num_items: 381
+    scratch_space_bytes: 1023
+  }
+  scratch_size_recordings {
+    num_items: 511
+    scratch_space_bytes: 1279
+  }
+  scratch_size_recordings {
+    num_items: 630
+    scratch_space_bytes: 1535
+  }
+  scratch_size_recordings {
+    num_items: 762
+    scratch_space_bytes: 1791
+  }
+  scratch_size_recordings {
+    num_items: 889
+    scratch_space_bytes: 2047
+  }
+  scratch_size_recordings {
+    num_items: 1017
+    scratch_space_bytes: 2303
+  }
+  scratch_size_recordings {
+    num_items: 1142
+    scratch_space_bytes: 2559
+  }
+  scratch_size_recordings {
+    num_items: 1258
+    scratch_space_bytes: 2815
+  }
+  scratch_size_recordings {
+    num_items: 1387
+    scratch_space_bytes: 3071
+  }
+  scratch_size_recordings {
+    num_items: 1528
+    scratch_space_bytes: 3327
+  }
+  scratch_size_recordings {
+    num_items: 1652
+    scratch_space_bytes: 3583
+  }
+  scratch_size_recordings {
+    num_items: 1787
+    scratch_space_bytes: 3839
+  }
+  scratch_size_recordings {
+    num_items: 1895
+    scratch_space_bytes: 4095
+  }
+  scratch_size_recordings {
+    num_items: 2009
+    scratch_space_bytes: 4351
+  }
+  scratch_size_recordings {
+    num_items: 2172
+    scratch_space_bytes: 4607
+  }
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 4863
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 5119
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 5375
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 5631
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 5887
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 6143
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 6399
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 6655
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 6911
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 7167
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 7423
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 7679
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 7935
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 8191
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 8447
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 8703
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 8959
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 9215
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 9471
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 9727
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 9983
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 10239
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 10495
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 10751
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 11007
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 11263
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 11519
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 11775
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 12031
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 12287
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 12543
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 12799
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 13055
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 13311
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 13567
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 13823
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 14079
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 14335
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 14591
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 14847
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 15103
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 15359
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 15615
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 15871
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 16383
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 16639
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 16895
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 17151
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 17663
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 17919
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 18431
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 18687
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 18943
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 19455
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 19711
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 20223
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 20479
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 20991
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 21503
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 21759
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 22271
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 22783
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 23039
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 23551
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 24063
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 24575
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 25087
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 25343
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 25855
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 26367
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 26879
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 27647
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 28159
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 28671
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 29183
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 29695
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 30463
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 30975
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 31487
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 32255
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 32767
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 33535
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 34047
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 34815
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 35583
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 36095
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 36863
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 37631
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 38399
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 39167
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 39935
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 40703
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 41471
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 42239
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 43263
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 44031
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 44799
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 45823
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 46591
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 47615
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 48639
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 49407
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 50431
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 51455
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 52479
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 53503
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 54527
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 55807
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 56831
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 57855
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 59135
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 60159
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 61439
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 62719
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 63999
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 65279
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 66559
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 67839
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 69119
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 70399
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 71935
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 73471
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 74751
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 76287
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 77823
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 79359
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 80895
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 82431
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 84223
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 85759
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 87551
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 89343
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 91135
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 92927
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 94719
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 96767
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 98559
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 100607
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 102655
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 104703
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 106751
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 108799
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 110847
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 113151
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 115455
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 117759
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 120063
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 122367
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 124927
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 127487
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 130047
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 132607
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 135167
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 137983
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 140543
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 143359
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 146175
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 149247
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 152063
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 155135
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 158207
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 161535
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 164607
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 167935
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 171263
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 174847
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 178175
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 181759
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 185343
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 189183
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 192767
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 196607
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 200703
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 204543
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 208639
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 212991
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 217087
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 221439
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 225791
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 230399
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 235007
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 239615
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 244479
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 249343
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 254463
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 259327
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 264703
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 269823
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 275199
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 280831
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 286463
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 292095
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 297983
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 303871
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 310015
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 316159
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 322559
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 328959
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 335359
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 342271
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 348927
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 356095
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 363007
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 370431
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 377855
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 385279
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 392959
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 400895
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 408831
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 417023
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 425215
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 433919
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 442623
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 451327
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 460287
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 469503
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 478975
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 488447
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 498175
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 508159
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 518399
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 528639
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 539391
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 550143
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 561151
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 572415
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 583679
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 595455
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 607231
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 619519
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 631807
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 644351
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 657407
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 670463
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 683775
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 697599
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 711423
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 725759
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 740095
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 754943
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 770047
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 785407
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 801279
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 817151
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 833535
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 850175
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 867071
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 884479
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 902143
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 920319
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 938751
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 957439
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 976639
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 996095
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 1016063
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 1036287
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 1057023
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 1078271
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 1099775
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 1121791
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 1144063
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 1167103
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 1190399
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 1214207
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 1238271
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 1263103
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 1288447
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 1314047
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 1340415
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 1367295
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 1394687
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 1422335
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 1451007
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 1479935
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 1509375
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 1539583
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 1570559
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 1601791
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 1633791
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 1666559
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 1699839
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 1733887
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 1768447
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 1803775
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 1839871
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 1876735
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 1914367
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 1952511
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 1991679
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 2031359
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 2072063
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 2113535
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 2155775
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 2198783
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 2242815
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 2287615
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 2333439
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 2380031
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 2427647
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 2476287
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 2525695
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 2576127
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 2627839
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 2680319
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 2733823
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 2788607
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 2844415
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 2901247
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 2959103
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 3018495
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 3078655
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 3140351
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 3203071
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 3267071
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 3332607
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 3399167
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 3467007
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 3536383
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 3607295
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 3679231
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 3752959
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 3827967
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 3904511
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 3982591
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 4062207
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 4143359
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 4226303
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 4310783
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 4397055
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 4484863
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 4574719
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 4666111
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 4759551
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 4854527
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 4951807
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 5050623
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 5151743
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 5254655
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 5359871
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 5467135
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 5576447
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 5687807
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 5801727
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 5917695
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 6035967
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 6156799
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 6279935
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 6405375
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 6533631
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 6664191
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 6797567
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 6933503
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 7071999
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 7213567
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 7357695
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 7504895
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 7654911
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 7807999
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 7964159
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 8123647
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 8285951
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 8451583
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 8620799
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 8793087
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 8968959
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 9148415
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 9331199
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 9517823
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 9708287
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 9902335
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 10100479
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 10302463
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 10508543
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 10718719
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 10932991
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 11151615
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 11374847
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 11602175
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 11834367
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 12070911
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 12312319
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 12558591
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 12809727
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 13065983
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 13327103
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 13593855
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 13865727
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 14142975
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 14425855
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 14714367
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 15008511
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 15308799
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 15614975
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 15927295
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 16245759
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 16570623
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 16902143
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 17240063
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 17584895
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 17936639
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 18295295
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 18661119
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 19034367
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 19415039
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 19803391
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 20199423
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 20603391
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 21015551
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 21435647
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 21864447
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 22301695
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 22747647
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 23202815
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 23666687
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 24140031
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 24622847
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 25115391
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 25617663
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 26129919
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 26652415
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 27185663
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 27729407
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 28283903
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 28849407
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 29426431
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 30014975
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 30615295
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 31227647
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 31852287
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 32489215
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 33138943
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 33801727
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 34477823
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 35167231
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 35870719
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 36588031
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 37319935
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 38066175
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 38827519
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 39603967
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 40396031
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 41203967
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 42028031
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 42868735
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 43726079
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 44600575
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 45492479
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 46402303
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 47330559
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 48276991
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 49242623
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 50227455
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 51231999
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 52256511
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 53301759
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 54367743
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 55454975
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 56564223
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 57695487
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 58849279
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 60026367
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 61226751
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 62451455
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 63700479
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 64974335
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 66273791
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 67599359
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 68951295
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 70330367
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 71736831
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 73171711
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 74635007
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 76127743
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 77650431
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 79203327
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 80787455
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 82403071
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 84051199
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 85732095
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 87446783
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 89195775
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 90979583
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 92799231
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 94655231
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 96548351
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 98479359
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 100448767
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 102457855
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 104506879
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 106597119
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 108729087
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 110903551
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 113121535
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 115384063
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 117691647
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 120045567
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 122446591
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 124895487
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 127393279
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 129941247
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 132539903
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 135190783
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 137894655
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 140652543
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 143465471
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 146334719
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 149261567
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 152246783
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 155291647
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 158397439
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 161565439
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 164796671
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 168092671
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 171454463
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 174883583
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 178381055
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 181948671
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 185587711
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 189299455
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 193085439
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 196947199
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 200886015
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 204903935
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 209001983
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 213181951
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 217445631
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 221794559
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 226230271
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 230755071
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 235369983
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 240077311
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 244878847
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 249776639
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 254771967
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 259867391
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 265064959
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 270366207
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 275773439
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 281288959
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 286914559
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 292653055
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 298505983
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 304476159
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 310565631
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 316776959
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 323112447
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 329574655
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 336166143
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 342889471
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 349747199
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 356742143
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 363877119
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 371154687
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 378577663
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 386149119
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 393872127
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 401749759
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 409784575
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 417980415
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 426339839
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 434866687
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 443564031
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 452435199
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 461484031
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 470713599
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 480127999
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 489730559
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 499525119
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 509515519
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 519705855
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 530099967
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 540701951
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 551515903
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 562546431
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 573797119
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 585273087
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 596978687
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 608918271
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 621096447
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 633518335
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 646188799
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 659112703
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 672294911
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 685740799
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 699455487
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 713444607
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 727713535
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 742267647
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 757113087
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 772255487
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 787700479
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 803454463
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 819523583
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 835913983
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 852632319
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 869684991
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 887078655
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 904820223
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 922916607
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 941374975
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 960202495
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 979406335
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 998994431
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 1018974463
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 1039353855
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 1060141055
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 1081343743
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 1102970623
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 1125030143
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 1147530751
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 1170481151
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 1193890815
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 1217768703
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 1242124031
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 1266966527
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 1292305919
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 1318151935
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 1344515071
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 1371405311
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 1398833407
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 1426810111
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 1455346175
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 1484453119
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 1514142207
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 1544424959
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 1575313407
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 1606819839
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 1638956031
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 1671735295
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 1705169919
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 1739273215
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 1774058751
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 1809540095
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 1845730815
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 1882645247
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 1920298239
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 1958704127
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 1997878271
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 4
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 1
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 29695
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 29951
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 30463
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 30975
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 31231
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 31743
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 32255
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 32511
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 33023
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 33535
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 34047
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 34559
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 35071
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 35583
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 36095
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 36607
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 37119
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 37631
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 38143
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 38655
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 40191
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 40959
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 41471
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 41983
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 42751
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 43263
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 44031
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 44543
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 45311
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 46079
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 46847
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 47359
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 48127
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 48895
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 49663
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 50431
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 51199
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 51967
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 52991
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 53759
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 54527
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 55551
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 56319
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 57343
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 58111
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 59135
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 60159
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 60927
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 61951
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 62975
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 63999
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 65279
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 66303
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 67327
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 68351
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 70655
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 71935
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 72959
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 74239
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 75519
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 76799
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 78079
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 79359
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 80639
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 82175
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 83455
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 84991
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 86271
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 87807
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 89343
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 90879
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 92415
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 94207
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 95743
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 97535
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 99071
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 101887
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 103679
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 105471
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 107263
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 109311
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 111103
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 113151
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 115199
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 117247
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 119295
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 121343
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 123647
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 125695
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 127999
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 131327
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 133631
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 136191
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 138495
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 141055
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 143615
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 146175
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 148735
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 151551
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 154111
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 156927
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 161023
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 163839
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 166911
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 169727
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 173055
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 176127
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 179199
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 182527
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 185855
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 189439
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 193791
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 197375
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 200959
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 204799
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 208639
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 212479
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 216319
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 221439
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 225535
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 229631
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 233727
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 238079
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 242687
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 247039
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 252671
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 257535
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 262143
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 267007
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 272127
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 277247
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 283391
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 288511
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 294143
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 299519
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 305151
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 310783
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 317695
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 323583
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 329727
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 335871
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 343295
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 349695
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 356351
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 363007
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 369919
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 377855
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 385023
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 392191
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 399615
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 408063
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 415743
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 423679
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 431615
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 440831
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 449279
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 457727
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 467455
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 476159
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 485375
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 495615
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 504831
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 514559
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 525311
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 535295
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 545279
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 556799
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 567295
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 578047
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 590079
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 601343
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 612863
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 625407
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 637439
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 650495
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 662783
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 676607
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 689407
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 702463
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 717055
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 730623
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 745727
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 759807
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 775423
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 790271
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 806399
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 821759
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 838399
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 854527
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 871679
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 889343
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 906495
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 924927
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 942591
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 961535
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 980991
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 999679
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 1019903
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 1040383
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 1060351
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 1081599
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 1103359
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 1124607
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 1147135
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 1170175
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 1193727
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 1216511
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 1240831
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 1265663
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 1291007
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 1316863
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 1342975
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 1369855
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 1396223
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 1423871
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 1452287
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 1481215
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 1510655
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 1540607
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 1571327
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 1602559
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 1634303
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 1666815
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 1699839
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 1734655
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 1768959
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 1803775
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 1839615
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 1875967
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 1914111
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 1951999
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 1990399
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 2029823
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 2070783
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 2111743
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 2153471
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 2196735
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 2240255
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 2285311
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 2330367
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 2377215
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 2424063
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 2472703
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 2522367
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 2571775
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 2623487
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 2675967
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 2729471
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 2783231
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 2838783
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 2895615
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 2953215
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 3012351
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 3072255
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 3133439
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 3195903
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 3259647
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 3324415
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 3390463
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 3458815
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 3527679
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 3597567
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 3670015
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 3742719
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 3817983
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 3893759
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 3971839
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 4050431
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 4131583
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 4214271
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 4297471
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 4383487
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 4471039
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 4560383
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 4651263
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 4744191
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 4838655
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 4936191
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 5034495
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 5134591
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 5237759
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 5341695
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 5448959
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 5556991
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 5668351
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 5781759
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 5896447
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 6014207
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 6134271
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 6256895
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 6381567
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 6509823
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 6639615
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 6771711
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 6907647
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 7044863
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 7185919
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 7329791
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 7476479
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 7625727
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 7778047
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 7933183
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 8091391
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 8252671
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 8418047
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 8585727
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 8757503
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 8932863
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 9111295
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 9293311
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 9478911
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 9668095
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 9861887
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 10058239
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 10259711
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 10464767
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 10673919
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 10887167
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 11104511
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 11326207
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 11553023
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 11783167
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 12018943
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 12259071
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 12504063
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 12753663
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 13009151
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 13268479
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 13533951
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 13804543
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 14080255
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 14362367
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 14648831
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 14941951
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 15240703
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 15545343
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 15855615
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 16173055
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 16495615
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 16825599
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 17161727
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 17505279
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 17854719
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 18211839
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 18575871
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 18946815
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 19325951
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 19712511
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 20106495
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 20508159
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 20918527
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 21337087
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 21763327
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 22198015
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 22642175
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 23094783
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 23556351
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 24027647
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 24507135
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 24997887
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 25497087
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 26006783
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 26526719
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 27057407
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 27598335
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 28150015
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 28713215
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 29287167
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 29872383
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 30469887
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 31078911
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 31700735
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 32334591
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 32980735
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 33640191
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 34312703
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 34999039
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 35698431
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 36412671
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 37140735
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 37882623
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 38640127
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 39413247
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 40201215
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 41005311
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 41825023
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 42660607
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 43514623
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 44383743
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 45271807
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 46177279
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 47099903
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 48042751
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 49002495
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 49982975
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 50982399
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 52002047
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 53041407
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 54102015
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 55183871
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 56287999
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 57413375
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 58561279
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 59732479
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 60926463
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 62144511
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 63388159
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 64655615
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 65947903
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 67266815
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 68612607
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 69984255
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 71383295
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 72811519
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 74267391
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 75752447
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 77267199
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 78812159
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 80388095
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 81996287
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 83635711
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 85307903
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 87014399
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 88753919
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 90529023
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 92339455
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 94185983
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 96069119
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 97990911
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 99950591
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 101949183
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 103988479
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 106067199
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 108189183
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 110352127
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 112558847
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 114810367
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 117106431
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 119447807
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 121836543
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 124273407
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 126758399
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 129293567
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 131879679
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 134516735
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 137206783
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 139950847
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 142750207
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 145604351
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 148516863
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 151486719
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 154515967
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 157606911
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 160758527
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 163972863
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 167252735
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 170597887
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 174009599
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 177489407
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 181038591
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 184659199
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 188351999
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 192119039
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 195961087
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 199880447
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 203877631
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 207954943
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 212114175
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 216355839
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 220683519
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 225096959
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 229598719
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 234190335
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 238874111
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 243651071
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 248523519
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 253494527
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 258564351
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 263735295
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 269009151
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 274390015
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 279876863
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 285474047
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 291183359
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 297007103
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 302947327
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 309005567
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 315185919
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 321489663
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 327919359
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 334477567
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 341166335
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 347989759
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 354948863
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 362047743
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 369288703
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 376674815
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 384207871
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 391891455
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 399729663
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 407723263
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 415877631
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 424195839
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 432679679
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 441332479
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 450158847
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 459161855
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 468344831
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 477711615
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 487266303
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 497010687
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 506950655
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 517090047
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 527431167
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 537979391
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 548739839
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 559714047
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 570907647
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 582325503
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 593972223
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 605851391
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 617968127
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 630327039
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 642934015
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 655791871
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 668908543
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 682285567
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 695931391
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 709850111
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 724046847
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 738527231
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 753297663
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 768363775
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 783731199
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 799405567
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 815393279
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 831701247
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 848334335
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 865301247
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 882607103
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 900258815
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 918263551
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 936628735
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 955361023
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 974468607
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 993957631
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 1013836543
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 1034112767
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 1054795263
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 1075890687
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 1097409023
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 1119095295
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 1140723199
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 1162783999
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 1185285631
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 1208237567
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 1231648511
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 1255527679
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 1279884287
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 1304728063
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 1330068991
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 1355916543
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 1382280959
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 1409172735
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 1436602367
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 1464580607
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 1493118207
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 1522226687
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 1551917567
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 1582202111
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 1613092095
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 1644600319
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 1676738303
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 1709519359
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 1742955775
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 1777061119
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 1811848447
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 1847331583
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 1883524351
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 1920441087
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 1958096127
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 1996504063
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 2035680255
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 2075640063
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 2116399103
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 2157973247
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 2200387071
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 2243640831
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 2287759615
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 2332760831
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 2378661887
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 2425481215
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 2473236735
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 2521947391
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 2571632383
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 2622311167
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 2674003199
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 2726729471
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 2780509951
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 2835366143
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 2891319551
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 2948391935
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 3006605567
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 3065983743
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 3126549503
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 3188326399
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 3251339007
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 3315619839
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 3381178111
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 3448047615
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 3516254207
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 3585825279
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 3656787711
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 3729169151
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 3802998271
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 3878304255
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 3955116031
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 4033464319
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 4
+  is_segmented: true
+  scratch_size_recordings {
+    num_items: 61
+    scratch_space_bytes: 511
+  }
+  scratch_size_recordings {
+    num_items: 121
+    scratch_space_bytes: 767
+  }
+  scratch_size_recordings {
+    num_items: 191
+    scratch_space_bytes: 1023
+  }
+  scratch_size_recordings {
+    num_items: 251
+    scratch_space_bytes: 1279
+  }
+  scratch_size_recordings {
+    num_items: 311
+    scratch_space_bytes: 1535
+  }
+  scratch_size_recordings {
+    num_items: 381
+    scratch_space_bytes: 1791
+  }
+  scratch_size_recordings {
+    num_items: 441
+    scratch_space_bytes: 2047
+  }
+  scratch_size_recordings {
+    num_items: 511
+    scratch_space_bytes: 2303
+  }
+  scratch_size_recordings {
+    num_items: 573
+    scratch_space_bytes: 2559
+  }
+  scratch_size_recordings {
+    num_items: 630
+    scratch_space_bytes: 2815
+  }
+  scratch_size_recordings {
+    num_items: 693
+    scratch_space_bytes: 3071
+  }
+  scratch_size_recordings {
+    num_items: 762
+    scratch_space_bytes: 3327
+  }
+  scratch_size_recordings {
+    num_items: 823
+    scratch_space_bytes: 3583
+  }
+  scratch_size_recordings {
+    num_items: 889
+    scratch_space_bytes: 3839
+  }
+  scratch_size_recordings {
+    num_items: 960
+    scratch_space_bytes: 4095
+  }
+  scratch_size_recordings {
+    num_items: 1017
+    scratch_space_bytes: 4351
+  }
+  scratch_size_recordings {
+    num_items: 1078
+    scratch_space_bytes: 4607
+  }
+  scratch_size_recordings {
+    num_items: 1142
+    scratch_space_bytes: 4863
+  }
+  scratch_size_recordings {
+    num_items: 1210
+    scratch_space_bytes: 5119
+  }
+  scratch_size_recordings {
+    num_items: 1258
+    scratch_space_bytes: 5375
+  }
+  scratch_size_recordings {
+    num_items: 1334
+    scratch_space_bytes: 5631
+  }
+  scratch_size_recordings {
+    num_items: 1387
+    scratch_space_bytes: 5887
+  }
+  scratch_size_recordings {
+    num_items: 1470
+    scratch_space_bytes: 6143
+  }
+  scratch_size_recordings {
+    num_items: 1528
+    scratch_space_bytes: 6399
+  }
+  scratch_size_recordings {
+    num_items: 1589
+    scratch_space_bytes: 6655
+  }
+  scratch_size_recordings {
+    num_items: 1652
+    scratch_space_bytes: 6911
+  }
+  scratch_size_recordings {
+    num_items: 1718
+    scratch_space_bytes: 7167
+  }
+  scratch_size_recordings {
+    num_items: 1787
+    scratch_space_bytes: 7423
+  }
+  scratch_size_recordings {
+    num_items: 1822
+    scratch_space_bytes: 7679
+  }
+  scratch_size_recordings {
+    num_items: 1895
+    scratch_space_bytes: 7935
+  }
+  scratch_size_recordings {
+    num_items: 1970
+    scratch_space_bytes: 8191
+  }
+  scratch_size_recordings {
+    num_items: 2009
+    scratch_space_bytes: 8447
+  }
+  scratch_size_recordings {
+    num_items: 2089
+    scratch_space_bytes: 8703
+  }
+  scratch_size_recordings {
+    num_items: 2172
+    scratch_space_bytes: 8959
+  }
+  scratch_size_recordings {
+    num_items: 2215
+    scratch_space_bytes: 9215
+  }
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 9471
+  }
+  scratch_size_recordings {
+    num_items: 2350
+    scratch_space_bytes: 9727
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 9983
+  }
+  scratch_size_recordings {
+    num_items: 2492
+    scratch_space_bytes: 10239
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 10495
+  }
+  scratch_size_recordings {
+    num_items: 2591
+    scratch_space_bytes: 10751
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 11007
+  }
+  scratch_size_recordings {
+    num_items: 2747
+    scratch_space_bytes: 11263
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 11519
+  }
+  scratch_size_recordings {
+    num_items: 2857
+    scratch_space_bytes: 11775
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 12031
+  }
+  scratch_size_recordings {
+    num_items: 2972
+    scratch_space_bytes: 12287
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 12543
+  }
+  scratch_size_recordings {
+    num_items: 3091
+    scratch_space_bytes: 12799
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 13055
+  }
+  scratch_size_recordings {
+    num_items: 3215
+    scratch_space_bytes: 13311
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 13567
+  }
+  scratch_size_recordings {
+    num_items: 3344
+    scratch_space_bytes: 13823
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 14079
+  }
+  scratch_size_recordings {
+    num_items: 3478
+    scratch_space_bytes: 14335
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 14591
+  }
+  scratch_size_recordings {
+    num_items: 3617
+    scratch_space_bytes: 14847
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 15103
+  }
+  scratch_size_recordings {
+    num_items: 3762
+    scratch_space_bytes: 15359
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 15615
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 16127
+  }
+  scratch_size_recordings {
+    num_items: 3991
+    scratch_space_bytes: 16383
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 16639
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 16895
+  }
+  scratch_size_recordings {
+    num_items: 4234
+    scratch_space_bytes: 17407
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 17663
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 17919
+  }
+  scratch_size_recordings {
+    num_items: 4492
+    scratch_space_bytes: 18431
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 18687
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 18943
+  }
+  scratch_size_recordings {
+    num_items: 4765
+    scratch_space_bytes: 19455
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 19711
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 20223
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 20479
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 20991
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 21503
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 21759
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 22271
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 22783
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 23039
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 23551
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 24063
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 24575
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 25087
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 25599
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 26111
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 26623
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 27135
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 27647
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 28159
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 28671
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 29183
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 29695
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 30463
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 30975
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 31487
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 32255
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 32767
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 33535
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 34047
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 34815
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 35583
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 36351
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 36863
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 37631
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 38399
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 39167
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 39935
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 40703
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 41471
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 42495
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 43263
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 44031
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 45055
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 45823
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 46847
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 47615
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 48639
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 49663
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 50431
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 51455
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 52479
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 53503
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 54783
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 55807
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 56831
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 57855
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 59135
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 60415
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 61439
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 62719
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 63999
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 65279
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 66559
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 67839
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 69119
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 70655
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 71935
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 73471
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 74751
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 76287
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 77823
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 79359
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 80895
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 82687
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 84223
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 86015
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 87551
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 89343
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 91135
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 92927
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 94719
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 96767
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 98559
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 100607
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 102655
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 104703
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 106751
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 108799
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 111103
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 113151
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 115455
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 117759
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 120063
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 122623
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 124927
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 127487
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 130047
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 132607
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 135167
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 137983
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 140543
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 143359
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 146431
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 149247
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 152319
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 155135
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 158463
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 161535
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 164607
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 167935
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 171263
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 174847
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 178175
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 181759
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 185343
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 189183
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 193023
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 196863
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 200703
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 204799
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 208895
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 212991
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 217087
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 221439
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 226047
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 230399
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 235007
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 239871
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 244479
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 249343
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 254463
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 259583
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 264703
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 269823
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 275455
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 280831
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 286463
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 292095
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 297983
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 303871
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 310015
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 316159
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 322559
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 328959
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 335615
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 342271
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 349183
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 356095
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 363263
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 370431
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 377855
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 385279
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 392959
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 400895
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 408831
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 417023
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 425471
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 433919
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 442623
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 451327
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 460543
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 469759
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 478975
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 488703
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 498431
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 508415
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 518399
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 528895
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 539391
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 550143
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 561151
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 572415
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 583935
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 595455
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 607487
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 619519
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 631807
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 644607
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 657407
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 670463
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 684031
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 697599
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 711679
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 725759
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 740351
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 755199
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 770303
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 785663
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 801279
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 817407
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 833535
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 850175
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 867327
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 884735
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 902399
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 920319
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 938751
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 957439
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 976639
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 996095
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 1016063
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 1036287
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 1057023
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 1078271
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 1099775
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 1121791
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 1144319
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 1167103
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 1190399
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 1214207
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 1238527
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 1263359
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 1288447
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 1314303
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 1340671
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 1367295
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 1394687
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 1422591
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 1451007
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 1479935
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 1509631
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 1539839
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 1570559
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 1602047
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 1634047
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 1666815
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 1700095
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 1733887
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 1768703
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 1804031
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 1840127
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 1876991
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 1914367
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 1952767
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 1991679
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 2031615
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 2072319
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 2113535
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 2156031
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 2199039
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 2243071
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 2287871
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 2333695
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 2380287
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 2427903
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 2476287
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 2525951
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 2576383
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 2627839
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 2680575
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 2734079
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 2788863
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 2844415
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 2901503
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 2959359
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 3018495
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 3078911
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 3140607
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 3203327
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 3267327
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 3332607
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 3399423
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 3467263
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 3536639
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 3607295
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 3679487
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 3753215
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 3828223
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 3904767
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 3982847
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 4062463
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 4143615
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 4226559
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 4311039
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 4397311
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 4485119
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 4574975
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 4666367
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 4759807
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 4855039
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 4952063
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 5051135
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 5151999
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 5255167
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 5360127
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 5467391
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 5576703
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 5688319
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 5801983
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 5917951
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 6036479
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 6157055
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 6280191
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 6405887
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 6533887
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 6664703
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 6797823
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 6933759
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 7072511
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 7214079
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 7358207
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 7505407
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 7655423
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 7808511
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 7964671
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 8123903
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 8286463
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 8452095
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 8621311
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 8793599
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 8969471
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 9148927
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 9331967
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 9518591
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 9708799
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 9903103
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 10100991
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 10302975
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 10509055
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 10719231
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 10933759
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 11152383
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 11375359
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 11602943
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 11834879
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 12071679
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 12313087
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 12559359
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 12810495
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 13066751
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 13328127
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 13594623
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 13866495
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 14143743
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 14426623
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 14715135
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 15009535
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 15309567
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 15615743
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 15928063
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 16246783
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 16571647
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 16902911
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 17241087
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 17585919
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 17937663
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 18296319
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 18662143
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 19035391
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 19416063
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 19804415
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 20200447
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 20604671
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 21016575
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 21436927
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 21865727
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 22302975
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 22749183
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 23204095
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 23668223
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 24141567
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 24624383
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 25116671
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 25619199
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 26131455
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 26653951
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 27187199
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 27730943
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 28285439
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 28851199
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 29428223
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 30016767
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 30617087
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 31229439
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 31854079
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 32491007
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 33140991
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 33803775
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 34479615
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 35169279
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 35872767
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 36590079
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 37321983
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 38068479
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 38829823
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 39606271
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 40398335
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 41206271
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 42030591
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 42871039
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 43728383
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 44603135
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 45495039
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 46405119
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 47333119
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 48279807
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 49245439
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 50230271
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 51234815
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 52259583
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 53304575
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 54370815
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 55458303
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 56567295
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 57698559
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 58852607
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 60029695
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 61230335
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 62454783
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 63704063
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 64977919
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 66277631
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 67603199
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 68955135
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 70334207
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 71740927
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 73175807
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 74639359
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 76132095
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 77654783
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 79207679
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 80791807
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 82407679
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 84055807
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 85736959
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 87451647
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 89200639
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 90984703
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 92804351
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 94660607
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 96553727
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 98484735
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 100454399
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 102463487
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 104512767
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 106603007
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 108734975
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 110909695
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 113127935
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 115390463
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 117698303
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 120052223
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 122453247
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 124902399
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 127400447
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 129948415
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 132547327
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 135198207
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 137902335
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 140660223
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 143473407
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 146342911
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 149269759
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 152255231
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 155300351
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 158406143
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 161574399
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 164805887
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 168101887
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 171463935
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 174893311
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 178391039
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 181958911
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 185597951
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 189309951
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 193096191
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 196958207
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 200897279
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 204915199
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 209013503
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 213193727
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 217457663
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 221806847
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 226242815
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 230767871
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 235383039
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 240090623
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 244892671
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 249790463
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 254786303
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 259881983
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 265079551
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 270381055
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 275788799
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 281304575
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 286930687
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 292669183
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 298522623
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 304493055
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 310582783
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 316794367
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 323130367
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 329592831
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 336184831
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 342908415
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 349766655
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 356761855
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 363897087
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 371175167
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 378598655
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 386170623
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 393893887
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 401771775
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 409807359
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 418003455
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 426363391
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 434890751
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 443588607
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 452460287
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 461509631
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 470739711
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 480154367
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 489757439
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 499552767
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 509543679
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 519734527
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 530129407
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 540731903
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 551546623
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 562577407
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 573828863
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 585305599
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 597011711
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 608951807
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 621130751
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 633553407
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 646224639
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 659149055
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 672332031
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 685778687
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 699494143
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 713484031
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 727753727
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 742308863
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 757155071
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 772297983
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 787743999
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 803499007
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 819568895
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 835960319
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 852679423
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 869733119
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 887127551
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 904870143
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 922967551
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 941426943
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 960255487
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 979460607
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 999049727
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 1019030783
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 1039411455
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 1060199679
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 1081403647
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 1103031551
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 1125092351
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 1147593983
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 1170545919
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 1193956863
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 1217836031
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 1242192639
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 1267036415
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 1292377343
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 1318224895
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 1344589311
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 1371481087
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 1398910719
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 1426888959
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 1455426559
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 1484535039
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 1514225919
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 1544510463
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 1575400447
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 1606908671
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 1639046655
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 1671827711
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 1705264127
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 1739369471
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 1774156799
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 1809639935
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 1845832703
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 1882749439
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 1920404479
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 1958812415
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 1997988607
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 2037948415
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 2078707455
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 2120281599
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 2162687231
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 2205940991
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 2250059775
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 2295060991
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 2340962047
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 2387781375
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 2435536895
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 2484247551
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 2533932543
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 2584611327
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 2636303359
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 2689029631
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 2742810111
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 2797666303
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 2853619711
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 2910692095
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 2968905727
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 3028283903
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 3088849663
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 3150626559
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 3213639167
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 3277911807
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 3343470079
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 3410339583
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 3478546175
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 3548117247
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 3619079679
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 3691461119
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 3765290239
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 3840596223
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 3917407999
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 3995756287
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 8
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 1
+  }
+  scratch_size_recordings {
+    num_items: 2350
+    scratch_space_bytes: 36863
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 37119
+  }
+  scratch_size_recordings {
+    num_items: 2444
+    scratch_space_bytes: 37631
+  }
+  scratch_size_recordings {
+    num_items: 2492
+    scratch_space_bytes: 37887
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 38399
+  }
+  scratch_size_recordings {
+    num_items: 2591
+    scratch_space_bytes: 38655
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 39167
+  }
+  scratch_size_recordings {
+    num_items: 2694
+    scratch_space_bytes: 39679
+  }
+  scratch_size_recordings {
+    num_items: 2747
+    scratch_space_bytes: 39935
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 40447
+  }
+  scratch_size_recordings {
+    num_items: 2857
+    scratch_space_bytes: 40959
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 41471
+  }
+  scratch_size_recordings {
+    num_items: 2972
+    scratch_space_bytes: 41727
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 42239
+  }
+  scratch_size_recordings {
+    num_items: 3091
+    scratch_space_bytes: 42751
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 43263
+  }
+  scratch_size_recordings {
+    num_items: 3215
+    scratch_space_bytes: 43775
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 44287
+  }
+  scratch_size_recordings {
+    num_items: 3344
+    scratch_space_bytes: 44799
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 45311
+  }
+  scratch_size_recordings {
+    num_items: 3478
+    scratch_space_bytes: 45823
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 46335
+  }
+  scratch_size_recordings {
+    num_items: 3617
+    scratch_space_bytes: 47103
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 47615
+  }
+  scratch_size_recordings {
+    num_items: 3762
+    scratch_space_bytes: 48127
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 48639
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 49407
+  }
+  scratch_size_recordings {
+    num_items: 3991
+    scratch_space_bytes: 49919
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 50687
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 51199
+  }
+  scratch_size_recordings {
+    num_items: 4234
+    scratch_space_bytes: 51967
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 52479
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 53247
+  }
+  scratch_size_recordings {
+    num_items: 4492
+    scratch_space_bytes: 54015
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 54783
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 55295
+  }
+  scratch_size_recordings {
+    num_items: 4765
+    scratch_space_bytes: 56063
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 56831
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 57599
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 58367
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 59391
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 60159
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 60927
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 61695
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 62719
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 63487
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 65535
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 66559
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 67327
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 68351
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 69375
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 70399
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 71423
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 72447
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 73471
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 74495
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 75519
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 76799
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 77823
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 79103
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 80383
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 81407
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 82687
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 83967
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 85247
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 86527
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 88063
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 89343
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 90879
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 92159
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 93695
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 95231
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 96767
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 98303
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 99839
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 101375
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 103167
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 104703
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 106495
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 108287
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 109823
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 112895
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 114687
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 116479
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 118527
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 120319
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 122367
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 124415
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 126463
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 128767
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 130815
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 133119
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 135167
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 137471
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 140031
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 142335
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 144895
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 147199
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 149759
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 152319
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 155135
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 157695
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 161535
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 164351
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 167167
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 169983
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 173055
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 176127
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 179199
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 182271
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 185599
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 188927
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 192255
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 195583
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 199167
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 202751
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 207359
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 210943
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 214783
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 218623
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 222463
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 226559
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 230655
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 234751
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 239103
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 243455
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 247807
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 252415
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 258047
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 262655
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 267519
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 272383
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 277247
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 282367
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 287487
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 292863
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 298239
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 304639
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 310271
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 316159
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 322047
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 327935
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 333823
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 340223
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 346367
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 353791
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 360447
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 367103
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 374015
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 380927
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 388095
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 396287
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 403711
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 411391
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 419071
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 426751
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 434943
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 444159
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 452351
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 460799
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 469503
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 478463
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 487423
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 497663
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 507135
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 516607
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 526335
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 537343
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 547583
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 557823
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 568319
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 579327
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 591359
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 602367
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 613887
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 625663
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 638463
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 650495
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 663039
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 675583
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 689663
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 702719
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 716031
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 730879
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 744703
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 759039
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 774399
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 789247
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 804351
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 820735
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 836607
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 852479
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 869887
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 886527
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 903423
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 921855
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 939519
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 957695
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 977151
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 995839
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 1016063
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 1035519
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 1056511
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 1076735
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 1097471
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 1119743
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 1141247
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 1164287
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 1186559
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 1210623
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 1233919
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 1258751
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 1283071
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 1308671
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 1334015
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 1360895
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 1388031
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 1414911
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 1443327
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 1471231
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 1500671
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 1530623
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 1560319
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 1591551
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 1623295
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 1654527
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 1687807
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 1721343
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 1754623
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 1789695
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 1825535
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 1861887
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 1897983
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 1935615
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 1974271
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 2013439
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 2053631
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 2094591
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 2135039
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 2177535
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 2220799
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 2265087
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 2309887
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 2355967
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 2402559
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 2450431
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 2499071
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 2548479
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 2599167
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 2650623
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 2704127
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 2757887
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 2812415
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 2867967
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 2924799
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 2982655
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 3042815
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 3102975
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 3164415
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 3228159
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 3291903
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 3356927
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 3424511
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 3492095
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 3561983
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 3632383
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 3705343
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 3778303
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 3854079
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 3931135
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 4008703
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 4088831
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 4170495
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 4252927
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 4337663
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 4424447
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 4512511
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 4602623
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 4694271
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 4787967
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 4883199
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 4980479
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 5079551
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 5180671
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 5283839
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 5389823
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 5497087
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 5606399
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 5718783
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 5832191
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 5948159
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 6067199
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 6187519
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 6311423
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 6437631
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 6566399
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 6696447
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 6830079
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 6966527
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 7105535
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 7247359
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 7391743
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 7539199
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 7689471
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 7843583
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 7999743
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 8158975
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 8322559
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 8487935
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 8657919
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 8830975
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 9006335
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 9186303
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 9369855
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 9556991
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 9747711
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 9942271
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 10140415
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 10343679
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 10549759
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 10760959
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 10975231
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 11194623
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 11418623
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 11645695
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 11878399
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 12115455
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 12357375
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 12604927
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 12856319
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 13113599
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 13374975
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 13642495
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 13915135
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 14192127
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 14475775
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 14765823
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 15060479
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 15360767
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 15668223
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 15980543
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 16300031
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 16625663
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 16957951
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 17296383
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 17642751
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 17994751
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 18354687
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 18721535
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 19095807
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 19477247
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 19866111
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 20262655
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 20667903
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 21081087
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 21502463
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 21932031
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 22370047
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 22817535
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 23272959
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 23738111
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 24212479
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 24696063
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 25190143
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 25693951
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 26206463
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 26731263
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 27265279
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 27810559
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 28365567
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 28932607
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 29511423
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 30100991
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 30702847
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 31316479
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 31942143
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 32581119
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 33232639
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 33896703
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 34573823
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 35265279
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 35970303
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 36689151
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 37423103
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 38170111
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 38934015
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 39711743
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 40505855
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 41315327
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 42141951
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 42984447
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 43843583
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 44720383
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 45613567
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 46526207
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 47456511
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 48405247
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 49372415
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 50359551
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 51366911
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 52393215
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 53441279
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 54509055
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 55598847
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 56710911
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 57844735
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 59001599
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 60181247
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 61383935
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 62611199
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 63863551
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 65140479
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 66443263
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 67771647
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 69126911
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 70508543
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 71918591
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 73356031
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 74823679
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 76318975
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 77845759
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 79402239
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 80990207
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 82609151
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 84260863
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 85946111
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 87664127
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 89416703
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 91204863
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 93028863
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 94889727
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 96786687
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 98722047
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 100696319
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 102710015
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 104763391
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 106858751
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 108995071
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 111174399
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 113398527
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 115665919
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 117978623
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 120337919
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 122744063
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 125198591
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 127702015
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 130255615
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 132860927
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 135517439
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 138227967
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 140991487
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 143811071
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 146686975
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 149620991
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 152612607
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 155664639
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 158777599
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 161952255
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 165191167
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 168494591
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 171864575
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 175301631
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 178806783
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 182382591
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 186029567
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 189749759
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 193544959
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 197415167
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 201363455
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 205389823
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 209498111
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 213687807
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 217960447
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 222319615
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 226765311
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 231301119
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 235926527
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 240644351
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 245456895
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 250365439
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 255373055
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 260479743
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 265689343
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 271002111
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 276421631
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 281950463
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 287589375
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 293340671
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 299206655
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 305190399
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 311294207
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 317519615
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 323869695
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 330346239
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 336953599
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 343692287
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 350564863
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 357575935
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 364727295
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 372021503
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 379461631
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 387050495
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 394790911
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 402686719
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 410739967
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 418954239
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 427333119
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 435879935
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 444596991
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 453488895
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 462557695
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 471808255
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 481244671
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 490868735
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 500686591
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 510699263
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 520913151
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 531330815
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 541957375
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 552796159
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 563851007
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 575128319
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 586630399
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 598362367
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 610329343
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 622535423
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 634986495
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 647685119
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 660638463
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 673851391
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 687327999
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 701074431
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 715094783
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 729396735
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 743984639
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 758864127
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 774040831
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 789520639
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 805311487
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 821416447
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 837844735
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 854601471
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 871693311
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 889126911
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 906908415
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 925047295
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 943547135
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 962417919
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 981666303
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 1001298431
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 1021324031
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 1041750527
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 1062584831
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 1083837183
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 1105513215
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 1127622911
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 1150175231
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 1173177855
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 1196641023
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 1220573439
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 1244985343
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 1269883903
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 1295281151
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 1321186559
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 1347609855
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 1374562559
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 1402052863
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 1430093311
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 1458695679
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 1487868927
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 1517626111
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 1547977983
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 1578936831
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 1610515967
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 1642726143
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 1675579647
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 1709090559
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 1743272191
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 1778137343
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 1813699583
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 1849973247
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 1886972927
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 1924711679
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 1963205375
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 2002469631
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 2042518527
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 2083368191
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 2125035775
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 2167535871
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 2210561279
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 2253817343
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 2297938687
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 2342942207
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 2388846079
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 2435667711
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 2483426047
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 2532139519
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 2581827071
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 2632508671
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 2684203775
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 2736932607
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 2790716159
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 2845575423
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 2901531903
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 2958607359
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 3016824319
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 3076205823
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 3136774911
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 3198555135
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 3261571327
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 3325847551
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 3391409407
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 3458282495
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 3526492927
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 3596067839
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 3667034111
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 3739419647
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 3813252863
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 3888562943
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 3965379071
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 4043731455
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 4123651071
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 4205168895
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 4288317183
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 4373144831
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 4459652351
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 4547889919
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 4637892351
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 4729694719
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 4823333119
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 4918844415
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 5016265727
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 5115635711
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 5216993023
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 5320377343
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 5425829631
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 5533390847
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 5643103231
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 5755009791
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 5869154559
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 5985582079
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 6104338431
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 6225469695
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 6349023743
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 6475048703
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 6603610623
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 6734727167
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 6868465919
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 7004879359
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 7144021247
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 7285946111
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 7430709247
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 7578367487
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 7728979199
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 7882603007
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 8039299327
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 8
+  is_segmented: true
+  scratch_size_recordings {
+    num_items: 31
+    scratch_space_bytes: 511
+  }
+  scratch_size_recordings {
+    num_items: 61
+    scratch_space_bytes: 767
+  }
+  scratch_size_recordings {
+    num_items: 91
+    scratch_space_bytes: 1023
+  }
+  scratch_size_recordings {
+    num_items: 121
+    scratch_space_bytes: 1279
+  }
+  scratch_size_recordings {
+    num_items: 151
+    scratch_space_bytes: 1535
+  }
+  scratch_size_recordings {
+    num_items: 191
+    scratch_space_bytes: 1791
+  }
+  scratch_size_recordings {
+    num_items: 221
+    scratch_space_bytes: 2047
+  }
+  scratch_size_recordings {
+    num_items: 251
+    scratch_space_bytes: 2303
+  }
+  scratch_size_recordings {
+    num_items: 281
+    scratch_space_bytes: 2559
+  }
+  scratch_size_recordings {
+    num_items: 311
+    scratch_space_bytes: 2815
+  }
+  scratch_size_recordings {
+    num_items: 351
+    scratch_space_bytes: 3071
+  }
+  scratch_size_recordings {
+    num_items: 381
+    scratch_space_bytes: 3327
+  }
+  scratch_size_recordings {
+    num_items: 411
+    scratch_space_bytes: 3583
+  }
+  scratch_size_recordings {
+    num_items: 441
+    scratch_space_bytes: 3839
+  }
+  scratch_size_recordings {
+    num_items: 471
+    scratch_space_bytes: 4095
+  }
+  scratch_size_recordings {
+    num_items: 511
+    scratch_space_bytes: 4351
+  }
+  scratch_size_recordings {
+    num_items: 541
+    scratch_space_bytes: 4607
+  }
+  scratch_size_recordings {
+    num_items: 573
+    scratch_space_bytes: 4863
+  }
+  scratch_size_recordings {
+    num_items: 606
+    scratch_space_bytes: 5119
+  }
+  scratch_size_recordings {
+    num_items: 630
+    scratch_space_bytes: 5375
+  }
+  scratch_size_recordings {
+    num_items: 667
+    scratch_space_bytes: 5631
+  }
+  scratch_size_recordings {
+    num_items: 693
+    scratch_space_bytes: 5887
+  }
+  scratch_size_recordings {
+    num_items: 734
+    scratch_space_bytes: 6143
+  }
+  scratch_size_recordings {
+    num_items: 762
+    scratch_space_bytes: 6399
+  }
+  scratch_size_recordings {
+    num_items: 792
+    scratch_space_bytes: 6655
+  }
+  scratch_size_recordings {
+    num_items: 823
+    scratch_space_bytes: 6911
+  }
+  scratch_size_recordings {
+    num_items: 855
+    scratch_space_bytes: 7167
+  }
+  scratch_size_recordings {
+    num_items: 889
+    scratch_space_bytes: 7423
+  }
+  scratch_size_recordings {
+    num_items: 924
+    scratch_space_bytes: 7679
+  }
+  scratch_size_recordings {
+    num_items: 960
+    scratch_space_bytes: 7935
+  }
+  scratch_size_recordings {
+    num_items: 979
+    scratch_space_bytes: 8191
+  }
+  scratch_size_recordings {
+    num_items: 1017
+    scratch_space_bytes: 8447
+  }
+  scratch_size_recordings {
+    num_items: 1037
+    scratch_space_bytes: 8703
+  }
+  scratch_size_recordings {
+    num_items: 1078
+    scratch_space_bytes: 8959
+  }
+  scratch_size_recordings {
+    num_items: 1120
+    scratch_space_bytes: 9215
+  }
+  scratch_size_recordings {
+    num_items: 1142
+    scratch_space_bytes: 9471
+  }
+  scratch_size_recordings {
+    num_items: 1164
+    scratch_space_bytes: 9727
+  }
+  scratch_size_recordings {
+    num_items: 1210
+    scratch_space_bytes: 9983
+  }
+  scratch_size_recordings {
+    num_items: 1234
+    scratch_space_bytes: 10239
+  }
+  scratch_size_recordings {
+    num_items: 1258
+    scratch_space_bytes: 10495
+  }
+  scratch_size_recordings {
+    num_items: 1308
+    scratch_space_bytes: 10751
+  }
+  scratch_size_recordings {
+    num_items: 1334
+    scratch_space_bytes: 11007
+  }
+  scratch_size_recordings {
+    num_items: 1360
+    scratch_space_bytes: 11263
+  }
+  scratch_size_recordings {
+    num_items: 1387
+    scratch_space_bytes: 11519
+  }
+  scratch_size_recordings {
+    num_items: 1414
+    scratch_space_bytes: 11775
+  }
+  scratch_size_recordings {
+    num_items: 1470
+    scratch_space_bytes: 12031
+  }
+  scratch_size_recordings {
+    num_items: 1499
+    scratch_space_bytes: 12287
+  }
+  scratch_size_recordings {
+    num_items: 1528
+    scratch_space_bytes: 12543
+  }
+  scratch_size_recordings {
+    num_items: 1558
+    scratch_space_bytes: 12799
+  }
+  scratch_size_recordings {
+    num_items: 1589
+    scratch_space_bytes: 13055
+  }
+  scratch_size_recordings {
+    num_items: 1620
+    scratch_space_bytes: 13311
+  }
+  scratch_size_recordings {
+    num_items: 1652
+    scratch_space_bytes: 13567
+  }
+  scratch_size_recordings {
+    num_items: 1685
+    scratch_space_bytes: 13823
+  }
+  scratch_size_recordings {
+    num_items: 1718
+    scratch_space_bytes: 14079
+  }
+  scratch_size_recordings {
+    num_items: 1752
+    scratch_space_bytes: 14335
+  }
+  scratch_size_recordings {
+    num_items: 1787
+    scratch_space_bytes: 14591
+  }
+  scratch_size_recordings {
+    num_items: 1822
+    scratch_space_bytes: 14847
+  }
+  scratch_size_recordings {
+    num_items: 1858
+    scratch_space_bytes: 15359
+  }
+  scratch_size_recordings {
+    num_items: 1895
+    scratch_space_bytes: 15615
+  }
+  scratch_size_recordings {
+    num_items: 1932
+    scratch_space_bytes: 15871
+  }
+  scratch_size_recordings {
+    num_items: 1970
+    scratch_space_bytes: 16127
+  }
+  scratch_size_recordings {
+    num_items: 2009
+    scratch_space_bytes: 16383
+  }
+  scratch_size_recordings {
+    num_items: 2049
+    scratch_space_bytes: 16895
+  }
+  scratch_size_recordings {
+    num_items: 2089
+    scratch_space_bytes: 17151
+  }
+  scratch_size_recordings {
+    num_items: 2130
+    scratch_space_bytes: 17407
+  }
+  scratch_size_recordings {
+    num_items: 2172
+    scratch_space_bytes: 17663
+  }
+  scratch_size_recordings {
+    num_items: 2215
+    scratch_space_bytes: 18175
+  }
+  scratch_size_recordings {
+    num_items: 2259
+    scratch_space_bytes: 18431
+  }
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 18687
+  }
+  scratch_size_recordings {
+    num_items: 2350
+    scratch_space_bytes: 19199
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 19455
+  }
+  scratch_size_recordings {
+    num_items: 2444
+    scratch_space_bytes: 19967
+  }
+  scratch_size_recordings {
+    num_items: 2492
+    scratch_space_bytes: 20223
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 20735
+  }
+  scratch_size_recordings {
+    num_items: 2591
+    scratch_space_bytes: 20991
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 21503
+  }
+  scratch_size_recordings {
+    num_items: 2694
+    scratch_space_bytes: 22015
+  }
+  scratch_size_recordings {
+    num_items: 2747
+    scratch_space_bytes: 22271
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 22783
+  }
+  scratch_size_recordings {
+    num_items: 2857
+    scratch_space_bytes: 23295
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 23807
+  }
+  scratch_size_recordings {
+    num_items: 2972
+    scratch_space_bytes: 24063
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 24575
+  }
+  scratch_size_recordings {
+    num_items: 3091
+    scratch_space_bytes: 25087
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 25599
+  }
+  scratch_size_recordings {
+    num_items: 3215
+    scratch_space_bytes: 26111
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 26623
+  }
+  scratch_size_recordings {
+    num_items: 3344
+    scratch_space_bytes: 27135
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 27647
+  }
+  scratch_size_recordings {
+    num_items: 3478
+    scratch_space_bytes: 28159
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 28671
+  }
+  scratch_size_recordings {
+    num_items: 3617
+    scratch_space_bytes: 29439
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 29951
+  }
+  scratch_size_recordings {
+    num_items: 3762
+    scratch_space_bytes: 30463
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 30975
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 31743
+  }
+  scratch_size_recordings {
+    num_items: 3991
+    scratch_space_bytes: 32255
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 33023
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 33535
+  }
+  scratch_size_recordings {
+    num_items: 4234
+    scratch_space_bytes: 34303
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 34815
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 35583
+  }
+  scratch_size_recordings {
+    num_items: 4492
+    scratch_space_bytes: 36351
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 37119
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 37631
+  }
+  scratch_size_recordings {
+    num_items: 4765
+    scratch_space_bytes: 38399
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 39167
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 39935
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 40703
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 41727
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 42495
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 43263
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 44031
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 45055
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 45823
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 46847
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 47871
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 48639
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 49663
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 50687
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 51711
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 52735
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 53759
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 54783
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 55807
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 56831
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 58111
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 59135
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 60415
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 61695
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 62719
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 63999
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 65279
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 66559
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 67839
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 69375
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 70655
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 72191
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 73471
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 75007
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 76543
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 78079
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 79615
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 81151
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 82687
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 84479
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 86015
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 87807
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 89599
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 91135
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 93183
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 94975
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 96767
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 98815
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 100607
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 102655
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 104703
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 106751
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 109055
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 111103
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 113407
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 115455
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 117759
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 120319
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 122623
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 125183
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 127487
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 130047
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 132607
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 135423
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 137983
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 140799
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 143615
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 146431
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 149247
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 152319
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 155391
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 158463
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 161535
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 164863
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 168191
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 171519
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 174847
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 178431
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 182015
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 185599
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 189183
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 193023
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 196863
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 200703
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 204799
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 208895
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 212991
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 217343
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 221695
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 226047
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 230655
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 235263
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 239871
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 244735
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 249599
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 254463
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 259583
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 264703
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 270079
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 275455
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 280831
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 286463
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 292351
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 298239
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 304127
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 310015
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 316415
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 322559
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 328959
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 335615
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 342271
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 349183
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 356095
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 363263
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 370431
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 377855
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 385535
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 393215
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 400895
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 409087
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 417279
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 425471
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 433919
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 442623
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 451583
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 460543
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 469759
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 479231
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 488703
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 498431
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 508415
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 518655
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 528895
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 539391
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 550399
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 561407
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 572415
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 583935
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 595711
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 607487
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 619519
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 632063
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 644607
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 657663
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 670719
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 684031
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 697855
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 711679
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 726015
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 740351
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 755199
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 770303
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 785663
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 801535
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 817407
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 833791
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 850431
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 867327
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 884735
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 902399
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 920575
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 939007
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 957695
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 976895
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 996351
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 1016319
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 1036543
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 1057279
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 1078527
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 1100031
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 1122047
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 1144319
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 1167359
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 1190655
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 1214463
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 1238783
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 1263359
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 1288703
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 1314559
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 1340671
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 1367551
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 1394943
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 1422847
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 1451263
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 1480191
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 1509887
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 1540095
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 1570815
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 1602047
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 1634303
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 1666815
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 1700095
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 1734143
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 1768959
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 1804287
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 1840383
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 1876991
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 1914623
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 1952767
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 1991935
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 2031871
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 2072319
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 2113791
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 2156031
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 2199295
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 2243071
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 2288127
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 2333695
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 2380543
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 2428159
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 2476543
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 2526207
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 2576639
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 2628095
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 2680831
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 2734335
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 2788863
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 2844671
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 2901503
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 2959615
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 3018751
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 3079167
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 3140863
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 3203583
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 3267583
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 3333119
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 3399679
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 3467519
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 3536895
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 3607807
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 3679743
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 3753471
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 3828479
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 3905023
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 3983103
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 4062719
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 4144127
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 4226815
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 4311551
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 4397567
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 4485631
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 4575231
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 4666879
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 4760063
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 4855295
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 4952319
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 5051391
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 5152511
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 5255423
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 5360639
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 5467903
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 5577215
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 5688575
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 5802495
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 5918463
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 6036735
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 6157567
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 6280703
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 6406399
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 6534399
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 6664959
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 6798335
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 6934271
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 7073023
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 7214335
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 7358719
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 7505919
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 7655935
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 7809023
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 7965183
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 8124671
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 8286975
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 8452863
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 8621823
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 8794111
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 8969983
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 9149439
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 9332479
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 9519103
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 9709567
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 9903615
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 10101759
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 10303743
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 10509823
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 10719999
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 10934271
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 11153151
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 11376127
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 11603711
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 11835647
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 12072447
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 12313855
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 12560127
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 12811263
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 13067519
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 13328895
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 13595391
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 13867263
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 14144767
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 14427647
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 14716159
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 15010303
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 15310591
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 15616767
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 15929087
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 16247551
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 16572671
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 16903935
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 17242111
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 17586943
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 17938687
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 18297343
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 18663423
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 19036671
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 19417343
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 19805695
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 20201727
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 20605695
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 21017855
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 21438207
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 21867007
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 22304255
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 22750463
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 23205375
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 23669503
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 24142847
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 24625663
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 25118207
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 25620479
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 26132991
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 26655743
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 27188735
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 27732479
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 28287231
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 28852735
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 29430015
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 30018559
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 30618879
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 31231231
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 31855871
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 32493055
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 33142783
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 33805567
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 34481663
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 35171327
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 35874815
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 36592383
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 37324031
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 38070527
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 38831871
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 39608575
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 40400639
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 41208831
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 42032895
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 42873599
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 43730943
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 44605695
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 45497855
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 46407679
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 47335935
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 48282623
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 49248255
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 50233087
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 51237887
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 52262399
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 53307647
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 54373887
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 55461375
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 56570623
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 57701887
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 58855935
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 60033023
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 61233663
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 62458367
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 63707647
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 64981759
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 66281471
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 67607039
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 68958975
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 70338303
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 71745023
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 73179903
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 74643455
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 76136447
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 77659135
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 79212287
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 80796415
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 82412287
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 84060671
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 85741823
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 87456511
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 89205759
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 90989823
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 92809727
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 94665727
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 96559103
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 98490367
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 100460031
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 102469375
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 104518655
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 106608895
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 108741119
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 110916095
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 113134335
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 115396863
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 117704959
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 120058879
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 122460159
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 124909311
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 127407615
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 129955583
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 132554751
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 135205887
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 137910015
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 140668159
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 143481599
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 146351103
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 149278207
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 152263679
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 155309055
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 158415103
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 161583359
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 164815103
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 168111359
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 171473663
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 174903039
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 178401023
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 181969151
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 185608447
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 189320703
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 193106943
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 196969215
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 200908543
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 204926719
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 209025279
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 213205759
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 217469695
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 221819135
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 226255615
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 230780671
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 235396095
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 240104191
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 244906239
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 249804287
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 254800383
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 259896319
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 265094399
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 270396159
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 275804159
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 281320191
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 286946559
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 292685567
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 298539263
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 304509951
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 310600191
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 316812031
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 323148287
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 329611263
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 336203519
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 342927615
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 349786111
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 356781823
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 363917311
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 371195647
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 378619647
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 386192127
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 393915903
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 401794047
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 409830143
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 418026751
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 426387199
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 434914815
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 443613183
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 452485375
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 461535231
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 470765823
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 480180991
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 489784831
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 499580415
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 509572095
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 519763455
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 530158591
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 540761855
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 551577087
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 562608639
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 573860863
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 585337855
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 597044735
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 608985599
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 621165311
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 633588479
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 646260223
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 659185407
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 672369151
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 685816575
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 699532799
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 713523455
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 727793919
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 742349823
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 757196799
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 772340735
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 787787519
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 803543295
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 819614207
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 836006399
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 852726527
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 869781247
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 887176703
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 904920319
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 923018751
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 941478911
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 960308479
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 979514623
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 999105023
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 1019087103
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 1039468799
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 1060258303
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 1081463295
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 1103092735
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 1125154559
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 1147657471
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 1170610687
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 1194022911
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 1217903359
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 1242261247
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 1267106559
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 1292448767
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 1318297599
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 1344663551
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 1371556863
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 1398988031
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 1426967807
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 1455507199
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 1484617215
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 1514309631
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 1544595711
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 1575487743
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 1606997503
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 1639137279
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 1671920127
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 1705358335
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 1739465727
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 1774254847
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 1809740031
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 1845934847
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 1882853375
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 1920510463
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 1958920703
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 1998099199
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 2038061055
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 2078822399
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 2120398847
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 2162806783
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 2206062847
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 2250184191
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 2295187711
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 2341091583
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 2387913215
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 2435671551
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 2484385023
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 2534072575
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 2584754175
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 2636449279
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 2689178111
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 2742961663
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 2797820927
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 2853777407
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 2910852863
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 2969069823
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 3028451327
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 3089020415
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 3150800639
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 3213816831
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 3278093055
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 3343654911
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 3410527999
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 3478738431
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 3548313343
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 3619279615
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 3691665151
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 3765498367
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 3840808447
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 3917624575
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 3995976959
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 4075896575
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 4157414399
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 4240562687
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 4325373951
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 4411881471
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 4500119039
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 4590121471
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 4681923839
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 4775562239
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 4871073535
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 4968494847
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 5067864831
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 5169222143
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 5272606463
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 5378058751
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 5485619967
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 5595332351
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 5707238911
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 5821383679
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 5937811199
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 6056567551
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 6177698815
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 6301252863
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 6427277823
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 6555823359
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 6686939903
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 6820678655
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 6957092095
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 7096233983
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 7238158847
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 7382921983
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 7530580223
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 7681191935
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 7834815743
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 7991512063
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 1
+  value_type_size: 2
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 1
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 3583
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 4607
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 5631
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 6655
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 7679
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 8703
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 9727
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 10751
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 11775
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 12799
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 13823
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 14847
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 15871
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 16895
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 17919
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 18943
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 19967
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 20991
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 22015
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 23039
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 24063
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 25087
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 26111
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 27135
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 28159
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 29183
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 30207
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 31231
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 32255
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 33279
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 34303
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 35327
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 36351
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 37375
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 38399
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 39423
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 40447
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 41471
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 42495
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 43519
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 44543
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 45567
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 46591
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 47615
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 48639
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 49663
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 50687
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 51711
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 52735
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 53759
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 54783
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 55807
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 56831
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 57855
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 58879
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 60927
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 61951
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 62975
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 63999
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 65023
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 66047
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 68095
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 69119
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 70143
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 71167
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 73215
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 74239
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 76287
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 77311
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 78335
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 80383
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 81407
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 83455
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 85503
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 86527
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 88575
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 89599
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 91647
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 93695
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 95743
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 96767
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 98815
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 100863
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 102911
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 104959
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 107007
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 109055
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 111103
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 113151
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 115199
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 118271
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 120319
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 122367
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 124415
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 127487
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 129535
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 132607
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 134655
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 137727
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 139775
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 142847
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 145919
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 148991
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 152063
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 154111
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 157183
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 160255
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 163327
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 167423
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 170495
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 173567
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 176639
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 180735
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 183807
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 187903
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 190975
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 195071
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 199167
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 203263
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 207359
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 211455
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 215551
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 219647
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 223743
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 227839
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 232959
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 237055
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 242175
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 246271
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 251391
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 256511
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 261631
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 266751
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 271871
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 276991
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 283135
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 288255
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 294399
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 300543
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 305663
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 311807
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 317951
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 324095
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 331263
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 337407
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 344575
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 350719
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 357887
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 365055
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 372223
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 379391
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 387583
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 394751
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 402943
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 411135
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 419327
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 427519
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 435711
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 443903
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 453119
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 462335
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 471551
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 480767
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 489983
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 500223
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 510463
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 519679
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 530943
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 541183
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 551423
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 562687
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 573951
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 585215
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 597503
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 608767
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 621055
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 633343
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 645631
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 658943
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 672255
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 685567
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 698879
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 713215
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 727551
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 741887
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 756223
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 771583
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 786943
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 802303
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 818687
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 835071
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 851455
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 867839
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 885247
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 903679
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 921087
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 939519
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 957951
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 977407
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 996863
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 1016319
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 1036799
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 1057279
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 1078783
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 1100287
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 1122815
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 1144319
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 1167871
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 1190399
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 1214975
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 1238527
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 1263103
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 1288703
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 1314303
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 1340927
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 1367551
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 1395199
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 1422847
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 1450495
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 1480191
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 1509887
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 1539583
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 1570303
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 1602047
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 1633791
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 1666559
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 1699327
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 1733119
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 1767935
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 1803775
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 1839615
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 1876479
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 1913343
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 1951231
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 1991167
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 2030079
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 2071039
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 2111999
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 2155007
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 2198015
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 2241023
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 2286079
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 2332159
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 2378239
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 2426367
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 2474495
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 2523647
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 2573823
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 2626047
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 2678271
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 2731519
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 2785791
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 2842111
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 2898431
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 2956799
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 3015167
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 3075583
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 3137023
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 3200511
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 3263999
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 3329535
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 3396095
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 3463679
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 3533311
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 3602943
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 3675647
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 3748351
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 3824127
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 3899903
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 3977727
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 4057599
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 4138495
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 4221439
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 4305407
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 4391423
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 4479487
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 4569599
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 4660735
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 4753919
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 4849151
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 4945407
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 5044735
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 5145087
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 5248511
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 5352959
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 5460479
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 5569023
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 5680639
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 5794303
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 5910015
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 6027775
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 6148607
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 6271487
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 6396415
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 6524415
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 6655487
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 6787583
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 6923775
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 7062015
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 7203327
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 7347711
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 7494143
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 7643647
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 7797247
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 7952895
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 8111615
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 8273407
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 8439295
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 8608255
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 8780287
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 8955391
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 9134591
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 9317887
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 9503231
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 9693695
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 9887231
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 10084863
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 10286591
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 10492415
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 10702335
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 10916351
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 11134463
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 11356671
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 11583999
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 11815423
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 12051967
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 12292607
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 12539391
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 12789247
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 13045247
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 13306367
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 13572607
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 13843967
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 14120447
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 14403071
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 14690815
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 14984703
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 15283711
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 15589887
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 15901183
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 16219647
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 16544255
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 16875007
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 17211903
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 17555967
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 17907199
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 18265599
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 18631167
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 19002879
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 19382783
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 19770879
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 20166143
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 20569599
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 20981247
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 21400063
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 21828095
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 22264319
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 22709759
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 23164415
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 23627263
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 24100351
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 24581631
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 25073151
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 25574911
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 26085887
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 26608127
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 27139583
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 27682303
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 28236287
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 28800511
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 29377023
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 29964799
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 30563839
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 31175167
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 31797759
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 32433663
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 33082879
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 33744383
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 34419199
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 35107327
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 35809791
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 36525567
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 37255679
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 38001151
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 38760959
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 39536127
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 40326655
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 41133567
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 41955839
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 42795519
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 43650559
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 44524031
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 45414911
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 46322175
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 47248895
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 48194047
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 48810495
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 48812543
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 48814591
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 1
+  value_type_size: 2
+  is_segmented: true
+  scratch_size_recordings {
+    num_items: 121
+    scratch_space_bytes: 767
+  }
+  scratch_size_recordings {
+    num_items: 251
+    scratch_space_bytes: 1023
+  }
+  scratch_size_recordings {
+    num_items: 381
+    scratch_space_bytes: 1535
+  }
+  scratch_size_recordings {
+    num_items: 511
+    scratch_space_bytes: 1791
+  }
+  scratch_size_recordings {
+    num_items: 630
+    scratch_space_bytes: 2303
+  }
+  scratch_size_recordings {
+    num_items: 762
+    scratch_space_bytes: 2559
+  }
+  scratch_size_recordings {
+    num_items: 889
+    scratch_space_bytes: 3071
+  }
+  scratch_size_recordings {
+    num_items: 1017
+    scratch_space_bytes: 3327
+  }
+  scratch_size_recordings {
+    num_items: 1142
+    scratch_space_bytes: 3839
+  }
+  scratch_size_recordings {
+    num_items: 1258
+    scratch_space_bytes: 4095
+  }
+  scratch_size_recordings {
+    num_items: 1387
+    scratch_space_bytes: 4607
+  }
+  scratch_size_recordings {
+    num_items: 1528
+    scratch_space_bytes: 4863
+  }
+  scratch_size_recordings {
+    num_items: 1652
+    scratch_space_bytes: 5375
+  }
+  scratch_size_recordings {
+    num_items: 1787
+    scratch_space_bytes: 5631
+  }
+  scratch_size_recordings {
+    num_items: 1895
+    scratch_space_bytes: 6143
+  }
+  scratch_size_recordings {
+    num_items: 2009
+    scratch_space_bytes: 6399
+  }
+  scratch_size_recordings {
+    num_items: 2172
+    scratch_space_bytes: 6911
+  }
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 7167
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 7679
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 7935
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 8447
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 8703
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 9215
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 9471
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 9983
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 10239
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 10751
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 11007
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 11519
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 11775
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 12287
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 12543
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 13055
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 13311
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 13823
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 14079
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 14591
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 14847
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 15359
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 15615
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 16127
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 16383
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 16895
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 17151
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 17663
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 17919
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 18431
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 18687
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 19199
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 19455
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 19967
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 20223
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 20735
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 20991
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 21503
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 21759
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 22271
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 22527
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 23039
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 23295
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 23807
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 24575
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 24831
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 25343
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 25599
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 26367
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 26879
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 27647
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 27903
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 28415
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 29183
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 29439
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 30207
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 30719
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 31487
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 32255
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 32511
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 33279
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 34047
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 34559
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 35327
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 36095
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 36863
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 37631
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 37887
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 38655
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 39423
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 40191
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 41471
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 42239
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 43007
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 43775
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 44543
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 45567
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 46335
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 47103
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 48383
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 49151
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 50175
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 50943
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 52223
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 53247
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 54015
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 55295
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 56319
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 57599
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 58623
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 59903
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 60927
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 62207
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 63231
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 64767
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 66047
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 67071
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 68607
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 69887
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 71423
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 72959
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 73983
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 75519
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 77055
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 78591
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 80127
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 81663
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 83711
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 85247
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 86783
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 88575
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 90111
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 92159
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 93951
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 95999
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 97791
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 99839
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 101631
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 103679
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 105471
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 107775
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 110079
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 112127
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 114431
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 116735
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 119039
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 121343
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 123647
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 126207
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 128511
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 131327
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 133887
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 136703
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 139263
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 142079
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 145151
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 147711
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 150783
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 153855
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 156927
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 159999
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 163071
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 166143
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 169727
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 173055
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 176639
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 179967
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 183551
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 187391
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 191231
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 195071
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 198911
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 202751
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 206847
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 210687
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 215039
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 219135
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 223743
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 228095
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 232703
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 237311
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 242175
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 246783
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 251903
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 256767
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 262143
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 267263
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 272639
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 278015
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 283647
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 289023
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 294911
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 301055
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 306687
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 312831
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 319487
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 325631
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 332031
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 338687
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 345599
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 352511
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 359423
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 366591
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 374015
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 381695
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 388863
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 397055
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 404735
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 412671
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 421119
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 429567
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 438015
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 446975
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 455679
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 464895
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 474111
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 483839
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 493311
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 503039
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 513279
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 523263
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 534015
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 544511
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 555519
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 566783
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 577791
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 589311
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 601343
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 613119
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 625407
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 637695
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 650751
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 663807
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 676863
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 690431
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 704255
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 718335
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 732671
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 747263
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 762111
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 777471
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 792831
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 808959
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 825087
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 841727
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 858623
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 875519
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 893183
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 910847
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 929279
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 947711
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 966399
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 986111
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 1005567
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 1025535
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 1046271
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 1067007
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 1088511
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 1110015
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 1132287
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 1155071
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 1178111
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 1201919
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 1225727
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 1250303
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 1275135
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 1300479
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 1326591
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 1353215
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 1380351
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 1407999
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 1436159
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 1464831
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 1494015
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 1523967
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 1554431
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 1585407
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 1617407
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 1649663
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 1682687
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 1715967
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 1750527
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 1785599
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 1821183
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 1857279
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 1894655
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 1932543
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 1970943
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 2010623
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 2050815
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 2092031
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 2133503
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 2176511
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 2219775
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 2264063
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 2309375
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 2355711
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 2402559
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 2450687
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 2499839
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 2549759
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 2600703
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 2652671
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 2705663
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 2759679
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 2814975
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 2871551
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 2928639
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 2987519
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 3046911
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 3108095
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 3170303
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 3233535
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 3298047
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 3364095
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 3431423
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 3500031
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 3569919
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 3641343
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 3714303
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 3788543
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 3864063
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 3941631
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 4020479
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 4100607
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 4182783
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 4266495
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 4351743
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 4438527
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 4527615
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 4617983
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 4710399
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 4804607
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 4900607
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 4998911
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 5098751
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 5200383
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 5304575
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 5410815
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 5518847
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 5629439
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 5741823
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 5856767
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 5973759
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 6093311
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 6214911
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 6339327
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 6466047
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 6595583
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 6727167
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 6862079
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 6999039
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 7139327
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 7281663
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 7427583
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 7575807
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 7727615
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 7881983
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 8039679
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 8200703
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 8364543
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 8531711
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 8702463
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 8876543
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 9053951
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 9235199
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 9419775
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 9607935
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 9800447
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 9996287
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 10196223
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 10400255
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 10607871
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 10820351
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 11036415
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 11257343
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 11482367
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 11711999
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 11946239
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 12185343
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 12428799
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 12677375
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 12931071
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 13189631
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 13453311
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 13722623
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 13996799
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 14276607
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 14562303
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 14853375
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 15150591
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 15453695
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 15762687
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 16078079
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 16399359
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 16727295
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 17062143
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 17403135
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 17751551
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 18106367
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 18468351
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 18837759
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 19214591
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 19598847
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 19990527
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 20390655
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 20798463
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 21214463
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 21638655
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 22071551
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 22512639
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 22963199
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 23422463
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 23890943
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 24368639
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 24855807
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 25353215
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 25860095
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 26377215
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 26904831
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 27442943
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 27991551
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 28551423
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 29122559
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 29704959
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 30299135
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 30905087
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 31523327
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 32153343
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 32796671
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 33452543
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 34121471
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 34804223
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 35500031
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 36209919
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 36934143
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 37672959
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 38426367
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 39194879
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 39978495
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 40778495
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 41594111
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 42425855
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 43273983
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 44139519
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 45022463
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 45922815
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 46841343
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 47778303
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 48733695
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 49708287
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 50702591
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 51716607
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 52750847
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 53806079
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 54882047
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 55979775
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 57099263
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 58241279
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 59405823
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 60593919
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 61805823
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 63042047
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 64303103
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 65588991
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 66900735
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 68238591
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 69603327
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 70995711
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 72415487
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 73863935
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 75341055
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 76847871
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 78384639
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 79952639
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 81551615
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 83182335
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 84846335
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 86543103
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 88273919
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 90039551
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 91839999
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 93677055
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 95550719
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 97461503
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 99410687
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 101399039
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 103426815
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 105495551
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 107605247
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 109757439
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 111952383
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 114191615
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 116475647
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 118804991
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 121181183
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 123604479
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 126076671
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 128598015
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 131170047
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 133793535
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 136469247
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 139198719
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 141982719
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 144822527
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 147718911
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 150673151
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 153686783
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 156760319
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 159895551
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 163093503
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 166355199
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 169682175
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 173075967
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 176537343
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 180068351
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 183669759
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 187343103
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 191089919
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 194911743
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 198809855
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 202786047
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 206841855
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 210978815
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 215198207
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 219502079
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 223892223
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 228370175
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 232937471
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 237596159
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 242348031
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 247194879
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 252139007
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 257181695
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 262325247
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 267571455
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 272922879
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 278381567
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 283949055
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 289628159
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 295420671
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 301328895
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 307355903
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 313502975
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 319772927
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 326168319
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 332691711
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 339345407
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 346132479
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 353054975
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 360115967
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 367318271
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 374664959
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 382157823
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 389800959
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 397597439
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 405549311
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 413660159
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 421933311
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 430371839
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 438979583
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 447758847
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 456714239
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 465848319
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 475165439
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 484668671
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 494361855
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 504249087
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 514334207
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 524620799
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 535113215
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 545815551
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 556731903
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 567866367
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 579223551
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 590808063
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 602624511
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 614676735
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 626970623
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 639509759
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 652300031
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 665346047
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 678652671
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 692226047
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 706070271
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 720191999
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 734595839
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 749287679
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 764273151
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 779558655
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 795149823
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 811052799
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 827273727
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 843819519
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 860695551
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 877909503
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 895468031
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 913377279
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 931644671
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 950277375
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 969283071
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 988668927
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 1008442367
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 1028611071
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 1049183231
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 1070166783
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 1091570175
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 1113401343
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 1135669503
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 1158383103
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 1181550591
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 1205181695
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 1229285375
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 1253870847
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 1278948351
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 1304527359
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 1330617855
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 1357230335
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 1384374783
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 1412062463
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 1440303615
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 1469109503
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 1498491647
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 1528461567
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 1559030783
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 1590211583
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 1622015487
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 1654455807
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 1687545087
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 1721296127
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 1755721727
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 1790836223
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 1826652927
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 1863185919
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 1900449791
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 1938458879
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 1977227775
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 2016772607
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 2057107967
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 2098249983
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 2140215039
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 2183019263
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 2226679551
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 2271213311
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 2316637439
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 2362970111
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 2410229759
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 2458434047
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 2507602943
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 2557754879
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 2608909823
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 2661087999
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 2714310143
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 2768596223
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 2823967743
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 2880447231
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 2938056191
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 2996817407
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 1
+  value_type_size: 4
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 1
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 3583
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 4607
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 5631
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 6655
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 7679
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 8703
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 9727
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 10751
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 11775
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 12799
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 13823
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 14847
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 15871
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 16895
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 17919
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 18943
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 19967
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 20991
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 22015
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 23039
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 24063
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 25087
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 26111
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 27135
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 28159
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 29183
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 30207
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 31231
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 32255
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 33279
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 34303
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 35327
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 36351
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 37375
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 38399
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 39423
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 40447
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 41471
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 42495
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 43519
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 44543
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 45567
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 46591
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 47615
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 48639
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 49663
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 50687
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 51711
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 52735
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 53759
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 54783
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 55807
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 56831
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 57855
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 58879
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 59903
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 61951
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 62975
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 63999
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 65023
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 66047
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 67071
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 69119
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 70143
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 71167
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 73215
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 74239
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 75263
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 77311
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 78335
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 80383
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 81407
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 83455
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 84479
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 86527
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 88575
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 89599
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 91647
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 93695
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 94719
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 96767
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 98815
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 100863
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 102911
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 104959
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 107007
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 109055
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 111103
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 113151
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 115199
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 117247
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 120319
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 122367
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 124415
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 127487
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 129535
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 131583
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 134655
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 136703
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 139775
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 142847
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 145919
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 147967
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 151039
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 154111
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 157183
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 160255
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 163327
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 166399
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 169471
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 173567
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 176639
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 179711
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 183807
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 186879
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 190975
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 195071
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 198143
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 202239
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 206335
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 210431
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 214527
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 218623
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 222719
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 227839
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 231935
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 237055
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 241151
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 246271
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 251391
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 255487
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 260607
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 265727
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 270847
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 276991
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 282111
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 288255
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 293375
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 299519
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 305663
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 310783
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 316927
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 324095
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 330239
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 336383
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 343551
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 349695
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 356863
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 364031
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 371199
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 378367
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 386559
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 393727
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 401919
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 410111
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 418303
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 426495
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 434687
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 442879
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 452095
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 461311
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 470527
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 479743
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 488959
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 499199
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 508415
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 518655
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 528895
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 539135
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 550399
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 561663
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 571903
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 584191
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 595455
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 607743
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 619007
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 631295
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 644607
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 656895
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 670207
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 683519
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 696831
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 711167
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 725503
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 739839
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 754175
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 769535
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 784895
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 800255
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 816639
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 831999
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 849407
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 865791
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 883199
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 900607
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 919039
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 937471
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 955903
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 975359
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 993791
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 1014271
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 1034751
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 1055231
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 1075711
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 1097215
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 1119743
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 1141247
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 1164799
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 1187327
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 1210879
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 1235455
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 1260031
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 1285631
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 1311231
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 1336831
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 1363455
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 1391103
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 1418751
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 1447423
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 1476095
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 1505791
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 1535487
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 1566207
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 1596927
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 1629695
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 1661439
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 1695231
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 1729023
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 1763839
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 1798655
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 1834495
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 1871359
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 1908223
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 1946111
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 1985023
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 2024959
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 2065919
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 2106879
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 2148863
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 2191871
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 2235903
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 2279935
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 2326015
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 2372095
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 2419199
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 2468351
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 2517503
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 2567679
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 2618879
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 2671103
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 2724351
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 2778623
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 2833919
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 2891263
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 2948607
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 3007999
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 3067391
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 3128831
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 3191295
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 3255807
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 3320319
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 3386879
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 3454463
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 3523071
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 3593727
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 3665407
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 3739135
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 3813887
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 3889663
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 3967487
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 4046335
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 4127231
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 4210175
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 4294143
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 4380159
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 4467199
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 4557311
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 4648447
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 4740607
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 4835839
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 4932095
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 5031423
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 5131775
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 5234175
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 5338623
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 5445119
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 5554687
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 5665279
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 5778943
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 5893631
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 6011391
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 6132223
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 6255103
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 6380031
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 6507007
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 6637055
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 6770175
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 6905343
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 7043583
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 7183871
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 7328255
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 7474687
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 7623167
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 7775743
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 7931391
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 8090111
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 8251903
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 8416767
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 8584703
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 8756735
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 8931839
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 9110015
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 9292287
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 9478655
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 9668095
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 9861631
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 10058239
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 10258943
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 10464767
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 10673663
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 10887679
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 11104767
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 11326975
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 11553279
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 11784703
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 12020223
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 12259839
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 12505599
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 12755455
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 13010431
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 13270527
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 13535743
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 13807103
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 14082559
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 14364159
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 14651903
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 14944767
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 15243775
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 15547903
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 15859199
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 16176639
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 16500223
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 16829951
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 17165823
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 17509887
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 17859071
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 18216447
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 18580991
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 18952703
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 19331583
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 19717631
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 20112895
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 20514303
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 20924927
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 21343743
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 21769727
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 22204927
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 22649343
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 23101951
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 23564799
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 24035839
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 24516095
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 25006591
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 25506303
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 26016255
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 26536447
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 27067903
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 27608575
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 28160511
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 28723711
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 29298175
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 29884927
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 30481919
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 31091199
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 31713791
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 32347647
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 32994815
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 33654271
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 34327039
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 35014143
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 35714559
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 36428287
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 37156351
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 37899775
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 38657535
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 39430655
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 40219135
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 41023999
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 41844223
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 42680831
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 43534847
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 44405247
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 45293055
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 46199295
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 47122943
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 47726079
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 47728127
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 47730175
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 1
+  value_type_size: 4
+  is_segmented: true
+  scratch_size_recordings {
+    num_items: 61
+    scratch_space_bytes: 767
+  }
+  scratch_size_recordings {
+    num_items: 121
+    scratch_space_bytes: 1023
+  }
+  scratch_size_recordings {
+    num_items: 191
+    scratch_space_bytes: 1279
+  }
+  scratch_size_recordings {
+    num_items: 251
+    scratch_space_bytes: 1535
+  }
+  scratch_size_recordings {
+    num_items: 311
+    scratch_space_bytes: 2047
+  }
+  scratch_size_recordings {
+    num_items: 381
+    scratch_space_bytes: 2303
+  }
+  scratch_size_recordings {
+    num_items: 441
+    scratch_space_bytes: 2559
+  }
+  scratch_size_recordings {
+    num_items: 511
+    scratch_space_bytes: 2815
+  }
+  scratch_size_recordings {
+    num_items: 573
+    scratch_space_bytes: 3327
+  }
+  scratch_size_recordings {
+    num_items: 630
+    scratch_space_bytes: 3583
+  }
+  scratch_size_recordings {
+    num_items: 693
+    scratch_space_bytes: 3839
+  }
+  scratch_size_recordings {
+    num_items: 762
+    scratch_space_bytes: 4095
+  }
+  scratch_size_recordings {
+    num_items: 823
+    scratch_space_bytes: 4607
+  }
+  scratch_size_recordings {
+    num_items: 889
+    scratch_space_bytes: 4863
+  }
+  scratch_size_recordings {
+    num_items: 960
+    scratch_space_bytes: 5119
+  }
+  scratch_size_recordings {
+    num_items: 1017
+    scratch_space_bytes: 5375
+  }
+  scratch_size_recordings {
+    num_items: 1078
+    scratch_space_bytes: 5887
+  }
+  scratch_size_recordings {
+    num_items: 1142
+    scratch_space_bytes: 6143
+  }
+  scratch_size_recordings {
+    num_items: 1210
+    scratch_space_bytes: 6399
+  }
+  scratch_size_recordings {
+    num_items: 1258
+    scratch_space_bytes: 6655
+  }
+  scratch_size_recordings {
+    num_items: 1334
+    scratch_space_bytes: 7167
+  }
+  scratch_size_recordings {
+    num_items: 1387
+    scratch_space_bytes: 7423
+  }
+  scratch_size_recordings {
+    num_items: 1470
+    scratch_space_bytes: 7679
+  }
+  scratch_size_recordings {
+    num_items: 1528
+    scratch_space_bytes: 7935
+  }
+  scratch_size_recordings {
+    num_items: 1589
+    scratch_space_bytes: 8447
+  }
+  scratch_size_recordings {
+    num_items: 1652
+    scratch_space_bytes: 8703
+  }
+  scratch_size_recordings {
+    num_items: 1718
+    scratch_space_bytes: 8959
+  }
+  scratch_size_recordings {
+    num_items: 1787
+    scratch_space_bytes: 9215
+  }
+  scratch_size_recordings {
+    num_items: 1822
+    scratch_space_bytes: 9727
+  }
+  scratch_size_recordings {
+    num_items: 1895
+    scratch_space_bytes: 9983
+  }
+  scratch_size_recordings {
+    num_items: 1970
+    scratch_space_bytes: 10239
+  }
+  scratch_size_recordings {
+    num_items: 2009
+    scratch_space_bytes: 10495
+  }
+  scratch_size_recordings {
+    num_items: 2089
+    scratch_space_bytes: 11007
+  }
+  scratch_size_recordings {
+    num_items: 2172
+    scratch_space_bytes: 11263
+  }
+  scratch_size_recordings {
+    num_items: 2215
+    scratch_space_bytes: 11519
+  }
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 11775
+  }
+  scratch_size_recordings {
+    num_items: 2350
+    scratch_space_bytes: 12287
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 12543
+  }
+  scratch_size_recordings {
+    num_items: 2492
+    scratch_space_bytes: 12799
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 13055
+  }
+  scratch_size_recordings {
+    num_items: 2591
+    scratch_space_bytes: 13567
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 13823
+  }
+  scratch_size_recordings {
+    num_items: 2747
+    scratch_space_bytes: 14079
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 14335
+  }
+  scratch_size_recordings {
+    num_items: 2857
+    scratch_space_bytes: 14847
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 15103
+  }
+  scratch_size_recordings {
+    num_items: 2972
+    scratch_space_bytes: 15359
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 15615
+  }
+  scratch_size_recordings {
+    num_items: 3091
+    scratch_space_bytes: 16127
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 16383
+  }
+  scratch_size_recordings {
+    num_items: 3215
+    scratch_space_bytes: 16639
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 16895
+  }
+  scratch_size_recordings {
+    num_items: 3344
+    scratch_space_bytes: 17407
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 17663
+  }
+  scratch_size_recordings {
+    num_items: 3478
+    scratch_space_bytes: 17919
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 18175
+  }
+  scratch_size_recordings {
+    num_items: 3617
+    scratch_space_bytes: 18687
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 18943
+  }
+  scratch_size_recordings {
+    num_items: 3762
+    scratch_space_bytes: 19199
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 19455
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 20223
+  }
+  scratch_size_recordings {
+    num_items: 3991
+    scratch_space_bytes: 20479
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 20735
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 21247
+  }
+  scratch_size_recordings {
+    num_items: 4234
+    scratch_space_bytes: 21759
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 22015
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 22527
+  }
+  scratch_size_recordings {
+    num_items: 4492
+    scratch_space_bytes: 23039
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 23295
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 23807
+  }
+  scratch_size_recordings {
+    num_items: 4765
+    scratch_space_bytes: 24319
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 24575
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 25343
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 25599
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 26367
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 26879
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 27135
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 27903
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 28415
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 28927
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 29439
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 30207
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 30719
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 31487
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 31999
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 32767
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 33279
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 34047
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 34559
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 35327
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 35839
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 36607
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 37119
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 38143
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 38655
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 39423
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 40447
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 40959
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 41983
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 42495
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 43519
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 44543
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 45567
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 46079
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 47103
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 48127
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 48895
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 49919
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 50943
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 51967
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 53247
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 54015
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 55039
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 56319
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 57343
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 58623
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 59647
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 60927
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 62207
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 62975
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 64255
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 65535
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 66815
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 68607
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 69887
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 71167
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 72447
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 73983
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 75519
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 76799
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 78335
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 80127
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 81663
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 83199
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 84735
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 86527
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 88319
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 89855
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 91903
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 93439
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 95487
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 97279
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 99327
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 101119
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 103423
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 105215
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 107519
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 109567
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 111615
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 113919
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 116223
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 118527
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 121087
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 123135
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 125695
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 128255
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 130815
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 133375
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 135935
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 139007
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 141567
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 144383
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 147199
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 150015
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 153343
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 156159
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 159487
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 162559
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 165887
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 168959
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 172543
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 175615
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 179199
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 183039
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 186623
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 190463
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 194047
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 198143
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 201983
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 205823
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 209919
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 214015
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 218623
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 222719
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 227327
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 231679
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 236543
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 241407
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 246015
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 250879
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 255999
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 261119
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 266239
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 271359
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 276735
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 282623
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 287999
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 293887
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 299775
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 305663
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 311807
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 318207
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 324607
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 331007
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 337407
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 344319
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 350975
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 358143
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 365055
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 372479
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 379903
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 387583
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 395263
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 403199
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 411135
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 419583
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 427775
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 436479
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 445183
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 454143
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 463103
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 472319
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 481535
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 491263
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 501247
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 510975
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 521215
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 531967
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 542463
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 553215
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 564223
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 575743
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 587263
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 598783
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 610815
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 623103
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 635647
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 647935
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 661247
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 674303
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 687615
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 701439
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 715519
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 729855
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 744447
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 759295
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 774399
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 789759
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 805887
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 821759
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 838143
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 855039
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 871935
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 889599
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 907263
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 925439
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 944127
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 962815
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 982015
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 1001727
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 1021695
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 1041919
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 1062655
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 1084159
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 1105919
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 1127935
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 1150463
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 1173503
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 1196799
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 1220863
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 1245183
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 1270015
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 1295359
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 1321215
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 1347839
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 1374719
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 1402367
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 1430527
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 1458943
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 1488127
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 1517823
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 1548287
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 1579263
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 1610495
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 1643007
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 1675775
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 1709055
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 1743359
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 1778175
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 1813759
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 1849855
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 1886975
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 1924863
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 1963263
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 2002687
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 2042623
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 2083583
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 2125055
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 2167295
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 2210815
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 2255103
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 2300159
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 2346239
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 2393087
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 2440959
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 2489599
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 2539519
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 2590463
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 2641919
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 2695167
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 2748927
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 2803967
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 2859775
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 2917119
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 2975487
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 3034879
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 3095295
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 3157503
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 3220479
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 3284735
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 3350783
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 3417599
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 3486207
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 3555583
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 3627007
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 3699199
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 3773183
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 3848703
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 3925759
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 4004095
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 4084223
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 4165887
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 4249343
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 4334079
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 4420863
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 4509183
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 4599295
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 4691455
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 4785407
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 4880895
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 4978687
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 5078015
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 5179647
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 5283327
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 5388799
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 5496575
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 5606399
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 5718783
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 5832959
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 5949695
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 6068735
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 6190079
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 6313983
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 6439935
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 6568959
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 6700287
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 6834175
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 6970879
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 7110399
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 7252479
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 7397375
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 7545599
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 7696383
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 7850239
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 8007423
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 8167423
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 8331007
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 8497407
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 8667135
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 8840703
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 9017599
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 9197823
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 9381887
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 9569279
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 9760767
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 9955839
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 10155007
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 10358015
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 10565119
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 10776575
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 10992127
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 11211775
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 11436287
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 11664895
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 11898367
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 12135935
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 12378879
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 12626175
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 12878847
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 13136383
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 13399039
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 13667327
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 13940479
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 14219263
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 14503679
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 14793727
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 15089663
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 15391487
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 15699199
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 16013055
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 16333567
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 16660223
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 16993279
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 17333247
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 17679615
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 18033407
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 18393855
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 18761983
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 19137023
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 19519743
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 19910143
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 20308479
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 20714495
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 21128703
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 21551359
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 21982463
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 22422015
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 22870527
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 23327743
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 23794175
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 24270079
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 24755455
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 25250559
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 25755903
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 26270719
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 26796287
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 27332095
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 27878655
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 28436479
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 29005055
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 29585407
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 30177023
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 30780415
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 31395839
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 32024063
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 32664319
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 33317375
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 33983999
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 34663679
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 35356927
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 36063999
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 36785407
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 37520895
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 38271487
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 39036927
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 39817727
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 40613887
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 41426175
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 42254847
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 43099647
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 43961599
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 44840959
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 45737727
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 46652415
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 47585535
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 48537343
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 49507839
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 50498047
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 51507967
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 52538367
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 53588735
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 54660607
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 55753983
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 56868863
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 58006527
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 59166463
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 60349695
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 61556735
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 62787839
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 64043519
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 65324543
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 66630655
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 67963647
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 69323007
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 70709247
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 72123135
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 73565695
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 75037183
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 76537855
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 78068479
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 79630079
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 81222399
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 82846975
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 84504063
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 86193919
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 87917823
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 89676287
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 91469823
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 93299199
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 95165183
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 97068543
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 99009535
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 100989695
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 103009535
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 105069823
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 107171327
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 109314559
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 111500799
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 113730815
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 116005375
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 118325759
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 120692223
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 123106047
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 125567999
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 128079359
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 130640895
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 133253887
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 135918847
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 138637055
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 141410047
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 144238079
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 147122943
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 150065407
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 153066495
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 156127999
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 159250687
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 162435583
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 165684223
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 168997887
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 172377855
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 175825407
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 179341823
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 182928639
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 186587135
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 190319103
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 194125567
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 198007807
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 201968127
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 206007295
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 210127359
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 214329855
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 218616575
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 222988799
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 227448575
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 231997439
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 236637439
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 241370367
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 246197759
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 251121663
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 256144127
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 261266943
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 266492159
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 271822079
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 277258495
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 282803455
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 288459775
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 294228735
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 300113407
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 306115839
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 312238079
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 318482943
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 324852479
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 331349503
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 337976319
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 344735999
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 351630847
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 358663423
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 365836543
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 373153279
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 380616447
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 388228607
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 395993087
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 403912959
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 411991039
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 420231167
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 428635647
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 437208319
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 445952255
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 454871295
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 463969023
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 473248255
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 482713343
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 492367359
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 502214655
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 512259327
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 522504447
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 532954367
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 543613439
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 554485759
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 565575423
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 576887039
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 588424703
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 600193023
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 612196863
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 624441087
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 636929535
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 649668095
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 662661887
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 675915007
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 689433343
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 703221759
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 717286143
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 731632127
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 746264575
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 761189887
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 776413439
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 791941887
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 807780863
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 823936255
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 840414975
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 857223423
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 874367743
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 891855103
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 909692159
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 927886079
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 946443775
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 965372415
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 984679935
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 1004373759
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 1024461055
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 1044950527
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 1065849343
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 1087166463
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 1108909567
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 1131087615
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 1153709567
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 1176783615
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 1200319487
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 1224325887
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 1248812287
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 1273788415
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 1299264255
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 1325249535
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 1351754495
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 1378789375
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 1406365439
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 1434492415
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 1463182335
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 1492446207
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 1522295039
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 1552740863
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 1583795455
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 1615471615
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 1647781119
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 1680736767
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 1714351359
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 1748638463
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 1783611135
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 1819283199
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 1855668735
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 1892782335
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 1930638079
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 1969250559
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 2008635903
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 2048808447
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 2089784575
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 2131580159
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 2174211839
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 2217695999
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 2262050047
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 2307290879
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 2353436927
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 2400505599
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 2448515583
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 2497485823
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 2547435519
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 2598384383
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 2650352127
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 2703358975
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 2757426175
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 2812574719
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 2868826367
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 2926202623
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 2984726783
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 3044421119
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 3105309439
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 3167415807
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 3230764287
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 3295379199
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 3361287167
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 3428512767
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 3497082879
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 3567024639
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 3638365183
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 3711132159
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 3785355007
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 3861062143
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 3938283263
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 4017049087
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 4097389823
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 4179337727
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 4262924543
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 4348182783
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 4435146495
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 4523849727
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 4614326527
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 4706612735
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 4800745215
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 4896760063
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 4994695423
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 1
+  value_type_size: 8
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 1
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 3583
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 4607
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 5631
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 6655
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 7679
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 8703
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 9727
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 10751
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 11775
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 12799
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 13823
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 14847
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 15871
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 16895
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 17919
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 18943
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 19967
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 20991
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 22015
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 23039
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 24063
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 25087
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 26111
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 27135
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 28159
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 29183
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 30207
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 31231
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 32255
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 33279
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 34303
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 35327
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 36351
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 37375
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 38399
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 39423
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 40447
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 41471
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 42495
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 43519
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 44543
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 45567
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 46591
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 47615
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 48639
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 49663
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 50687
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 51711
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 52735
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 53759
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 54783
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 55807
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 56831
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 57855
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 58879
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 59903
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 60927
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 61951
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 63999
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 65023
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 66047
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 67071
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 68095
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 70143
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 71167
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 72191
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 74239
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 75263
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 77311
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 78335
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 79359
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 81407
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 82431
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 84479
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 86527
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 87551
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 89599
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 91647
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 92671
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 94719
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 96767
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 98815
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 100863
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 101887
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 103935
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 105983
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 108031
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 110079
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 113151
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 115199
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 117247
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 119295
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 121343
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 124415
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 126463
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 128511
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 131583
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 133631
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 136703
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 139775
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 141823
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 144895
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 147967
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 151039
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 153087
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 156159
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 159231
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 162303
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 166399
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 169471
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 172543
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 175615
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 179711
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 182783
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 186879
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 189951
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 194047
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 198143
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 201215
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 205311
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 209407
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 213503
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 217599
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 222719
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 226815
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 230911
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 236031
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 240127
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 245247
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 250367
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 254463
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 259583
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 264703
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 269823
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 275967
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 281087
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 286207
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 292351
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 298495
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 303615
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 309759
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 315903
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 322047
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 329215
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 335359
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 341503
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 348671
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 355839
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 363007
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 370175
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 377343
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 384511
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 392703
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 399871
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 408063
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 416255
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 424447
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 432639
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 441855
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 450047
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 459263
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 468479
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 477695
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 486911
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 497151
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 506367
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 516607
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 526847
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 537087
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 548351
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 558591
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 569855
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 581119
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 593407
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 604671
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 616959
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 629247
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 641535
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 654847
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 667135
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 680447
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 694783
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 708095
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 722431
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 736767
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 751103
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 766463
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 781823
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 797183
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 812543
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 828927
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 845311
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 862719
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 880127
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 897535
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 914943
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 933375
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 951807
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 971263
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 990719
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 1010175
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 1030655
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 1051135
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 1071615
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 1093119
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 1114623
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 1137151
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 1159679
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 1183231
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 1206783
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 1230335
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 1254911
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 1280511
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 1306111
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 1331711
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 1358335
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 1384959
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 1412607
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 1441279
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 1469951
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 1499647
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 1529343
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 1560063
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 1590783
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 1622527
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 1655295
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 1688063
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 1721855
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 1756671
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 1791487
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 1827327
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 1863167
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 1901055
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 1938943
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 1977855
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 2016767
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 2056703
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 2098687
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 2140671
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 2182655
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 2226687
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 2270719
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 2316799
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 2362879
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 2409983
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 2458111
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 2507263
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 2557439
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 2608639
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 2659839
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 2713087
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 2767359
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 2822655
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 2878975
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 2937343
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 2995711
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 3055103
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 3116543
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 3179007
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 3242495
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 3307007
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 3372543
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 3440127
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 3508735
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 3579391
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 3651071
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 3723775
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 3798527
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 3874303
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 3951103
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 4030975
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 4110847
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 4192767
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 4276735
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 4362751
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 4449791
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 4538879
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 4628991
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 4722175
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 4816383
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 4912639
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 5010943
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 5111295
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 5212671
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 5317119
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 5423615
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 5532159
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 5642751
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 5755391
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 5870079
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 5987839
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 6107647
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 6229503
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 6354431
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 6481407
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 6610431
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 6742527
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 6877695
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 7014911
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 7155199
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 7298559
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 7443967
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 7592447
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 7745023
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 7899647
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 8057343
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 8218111
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 8382975
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 8550911
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 8721919
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 8895999
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 9073151
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 9255423
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 9439743
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 9629183
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 9821695
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 10017279
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 10217983
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 10422783
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 10630655
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 10843647
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 11059711
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 11280895
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 11507199
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 11736575
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 11971071
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 12210687
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 12455423
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 12704255
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 12958207
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 13217279
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 13481471
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 13750783
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 14026239
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 14306815
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 14592511
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 14884351
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 15182335
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 15485439
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 15794687
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 16111103
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 16433663
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 16761343
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 17097215
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 17439231
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 17787391
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 18143743
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 18506239
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 18875903
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 19253759
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 19638783
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 20030975
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 20431359
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 20839935
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 21256703
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 21682687
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 22115839
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 22558207
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 23008767
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 23469567
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 23938559
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 24417791
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 24905215
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 25403903
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 25911807
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 26429951
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 26958335
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 27497983
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 28047871
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 28607999
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 29180415
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 29764095
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 30359039
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 30966271
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 31585791
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 32217599
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 32861695
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 33519103
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 34188799
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 34872831
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 35570175
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 36281855
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 37006847
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 37747199
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 38501887
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 39271935
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 40057343
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 40858111
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 41675263
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 42508799
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 43358719
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 44226047
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 45110783
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 46012927
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 46932479
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 47871487
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 48828927
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 49805823
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 50801151
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 51816959
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 52853247
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 53911039
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 54989311
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 56089087
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 57210367
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 58354175
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 59521535
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 60711423
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 61925887
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 63164927
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 64427519
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 65715711
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 67030527
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 68370943
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 69737983
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 71132671
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 72556031
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 74007039
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 75486719
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 76996095
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 78536191
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 79540735
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 79542783
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 79544831
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 1
+  value_type_size: 8
+  is_segmented: true
+  scratch_size_recordings {
+    num_items: 31
+    scratch_space_bytes: 767
+  }
+  scratch_size_recordings {
+    num_items: 61
+    scratch_space_bytes: 1023
+  }
+  scratch_size_recordings {
+    num_items: 91
+    scratch_space_bytes: 1279
+  }
+  scratch_size_recordings {
+    num_items: 121
+    scratch_space_bytes: 1535
+  }
+  scratch_size_recordings {
+    num_items: 151
+    scratch_space_bytes: 1791
+  }
+  scratch_size_recordings {
+    num_items: 191
+    scratch_space_bytes: 2047
+  }
+  scratch_size_recordings {
+    num_items: 221
+    scratch_space_bytes: 2303
+  }
+  scratch_size_recordings {
+    num_items: 251
+    scratch_space_bytes: 2559
+  }
+  scratch_size_recordings {
+    num_items: 281
+    scratch_space_bytes: 3071
+  }
+  scratch_size_recordings {
+    num_items: 311
+    scratch_space_bytes: 3327
+  }
+  scratch_size_recordings {
+    num_items: 351
+    scratch_space_bytes: 3583
+  }
+  scratch_size_recordings {
+    num_items: 381
+    scratch_space_bytes: 3839
+  }
+  scratch_size_recordings {
+    num_items: 411
+    scratch_space_bytes: 4095
+  }
+  scratch_size_recordings {
+    num_items: 441
+    scratch_space_bytes: 4351
+  }
+  scratch_size_recordings {
+    num_items: 471
+    scratch_space_bytes: 4607
+  }
+  scratch_size_recordings {
+    num_items: 511
+    scratch_space_bytes: 4863
+  }
+  scratch_size_recordings {
+    num_items: 541
+    scratch_space_bytes: 5375
+  }
+  scratch_size_recordings {
+    num_items: 573
+    scratch_space_bytes: 5631
+  }
+  scratch_size_recordings {
+    num_items: 606
+    scratch_space_bytes: 5887
+  }
+  scratch_size_recordings {
+    num_items: 630
+    scratch_space_bytes: 6143
+  }
+  scratch_size_recordings {
+    num_items: 667
+    scratch_space_bytes: 6399
+  }
+  scratch_size_recordings {
+    num_items: 693
+    scratch_space_bytes: 6655
+  }
+  scratch_size_recordings {
+    num_items: 734
+    scratch_space_bytes: 6911
+  }
+  scratch_size_recordings {
+    num_items: 762
+    scratch_space_bytes: 7167
+  }
+  scratch_size_recordings {
+    num_items: 792
+    scratch_space_bytes: 7679
+  }
+  scratch_size_recordings {
+    num_items: 823
+    scratch_space_bytes: 7935
+  }
+  scratch_size_recordings {
+    num_items: 855
+    scratch_space_bytes: 8191
+  }
+  scratch_size_recordings {
+    num_items: 889
+    scratch_space_bytes: 8447
+  }
+  scratch_size_recordings {
+    num_items: 924
+    scratch_space_bytes: 8703
+  }
+  scratch_size_recordings {
+    num_items: 960
+    scratch_space_bytes: 8959
+  }
+  scratch_size_recordings {
+    num_items: 979
+    scratch_space_bytes: 9215
+  }
+  scratch_size_recordings {
+    num_items: 1017
+    scratch_space_bytes: 9471
+  }
+  scratch_size_recordings {
+    num_items: 1037
+    scratch_space_bytes: 9983
+  }
+  scratch_size_recordings {
+    num_items: 1078
+    scratch_space_bytes: 10239
+  }
+  scratch_size_recordings {
+    num_items: 1120
+    scratch_space_bytes: 10495
+  }
+  scratch_size_recordings {
+    num_items: 1142
+    scratch_space_bytes: 10751
+  }
+  scratch_size_recordings {
+    num_items: 1164
+    scratch_space_bytes: 11007
+  }
+  scratch_size_recordings {
+    num_items: 1210
+    scratch_space_bytes: 11263
+  }
+  scratch_size_recordings {
+    num_items: 1234
+    scratch_space_bytes: 11519
+  }
+  scratch_size_recordings {
+    num_items: 1258
+    scratch_space_bytes: 11775
+  }
+  scratch_size_recordings {
+    num_items: 1308
+    scratch_space_bytes: 12287
+  }
+  scratch_size_recordings {
+    num_items: 1334
+    scratch_space_bytes: 12543
+  }
+  scratch_size_recordings {
+    num_items: 1360
+    scratch_space_bytes: 12799
+  }
+  scratch_size_recordings {
+    num_items: 1387
+    scratch_space_bytes: 13055
+  }
+  scratch_size_recordings {
+    num_items: 1414
+    scratch_space_bytes: 13311
+  }
+  scratch_size_recordings {
+    num_items: 1470
+    scratch_space_bytes: 13567
+  }
+  scratch_size_recordings {
+    num_items: 1499
+    scratch_space_bytes: 13823
+  }
+  scratch_size_recordings {
+    num_items: 1528
+    scratch_space_bytes: 14079
+  }
+  scratch_size_recordings {
+    num_items: 1558
+    scratch_space_bytes: 14591
+  }
+  scratch_size_recordings {
+    num_items: 1589
+    scratch_space_bytes: 14847
+  }
+  scratch_size_recordings {
+    num_items: 1620
+    scratch_space_bytes: 15103
+  }
+  scratch_size_recordings {
+    num_items: 1652
+    scratch_space_bytes: 15359
+  }
+  scratch_size_recordings {
+    num_items: 1685
+    scratch_space_bytes: 15615
+  }
+  scratch_size_recordings {
+    num_items: 1718
+    scratch_space_bytes: 15871
+  }
+  scratch_size_recordings {
+    num_items: 1752
+    scratch_space_bytes: 16127
+  }
+  scratch_size_recordings {
+    num_items: 1787
+    scratch_space_bytes: 16383
+  }
+  scratch_size_recordings {
+    num_items: 1822
+    scratch_space_bytes: 16895
+  }
+  scratch_size_recordings {
+    num_items: 1858
+    scratch_space_bytes: 17407
+  }
+  scratch_size_recordings {
+    num_items: 1895
+    scratch_space_bytes: 17663
+  }
+  scratch_size_recordings {
+    num_items: 1932
+    scratch_space_bytes: 17919
+  }
+  scratch_size_recordings {
+    num_items: 1970
+    scratch_space_bytes: 18175
+  }
+  scratch_size_recordings {
+    num_items: 2009
+    scratch_space_bytes: 18431
+  }
+  scratch_size_recordings {
+    num_items: 2049
+    scratch_space_bytes: 19199
+  }
+  scratch_size_recordings {
+    num_items: 2089
+    scratch_space_bytes: 19455
+  }
+  scratch_size_recordings {
+    num_items: 2130
+    scratch_space_bytes: 19711
+  }
+  scratch_size_recordings {
+    num_items: 2172
+    scratch_space_bytes: 19967
+  }
+  scratch_size_recordings {
+    num_items: 2215
+    scratch_space_bytes: 20479
+  }
+  scratch_size_recordings {
+    num_items: 2259
+    scratch_space_bytes: 20735
+  }
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 20991
+  }
+  scratch_size_recordings {
+    num_items: 2350
+    scratch_space_bytes: 21759
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 22015
+  }
+  scratch_size_recordings {
+    num_items: 2444
+    scratch_space_bytes: 22527
+  }
+  scratch_size_recordings {
+    num_items: 2492
+    scratch_space_bytes: 22783
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 23295
+  }
+  scratch_size_recordings {
+    num_items: 2591
+    scratch_space_bytes: 23807
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 24319
+  }
+  scratch_size_recordings {
+    num_items: 2694
+    scratch_space_bytes: 24831
+  }
+  scratch_size_recordings {
+    num_items: 2747
+    scratch_space_bytes: 25087
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 25599
+  }
+  scratch_size_recordings {
+    num_items: 2857
+    scratch_space_bytes: 26367
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 26879
+  }
+  scratch_size_recordings {
+    num_items: 2972
+    scratch_space_bytes: 27135
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 27647
+  }
+  scratch_size_recordings {
+    num_items: 3091
+    scratch_space_bytes: 28415
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 28927
+  }
+  scratch_size_recordings {
+    num_items: 3215
+    scratch_space_bytes: 29439
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 29951
+  }
+  scratch_size_recordings {
+    num_items: 3344
+    scratch_space_bytes: 30719
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 31231
+  }
+  scratch_size_recordings {
+    num_items: 3478
+    scratch_space_bytes: 31743
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 32255
+  }
+  scratch_size_recordings {
+    num_items: 3617
+    scratch_space_bytes: 33279
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 33791
+  }
+  scratch_size_recordings {
+    num_items: 3762
+    scratch_space_bytes: 34303
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 34815
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 35839
+  }
+  scratch_size_recordings {
+    num_items: 3991
+    scratch_space_bytes: 36351
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 37119
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 37887
+  }
+  scratch_size_recordings {
+    num_items: 4234
+    scratch_space_bytes: 38655
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 39167
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 40191
+  }
+  scratch_size_recordings {
+    num_items: 4492
+    scratch_space_bytes: 40959
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 41727
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 42495
+  }
+  scratch_size_recordings {
+    num_items: 4765
+    scratch_space_bytes: 43263
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 44031
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 45055
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 45823
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 47103
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 47871
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 48639
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 49663
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 50687
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 51711
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 52735
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 54015
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 54783
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 56063
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 57087
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 58367
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 59391
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 60671
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 61695
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 62975
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 63999
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 65535
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 66559
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 68095
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 69375
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 70655
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 72191
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 73471
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 75007
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 76287
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 78079
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 79615
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 81407
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 82687
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 84479
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 86271
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 87807
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 89599
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 91391
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 93183
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 95231
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 96767
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 98815
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 100863
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 102655
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 104959
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 107007
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 109055
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 111359
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 113151
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 115455
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 117759
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 120063
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 122879
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 125183
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 127743
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 130047
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 132607
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 135423
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 137983
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 140799
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 143615
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 146431
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 149247
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 152319
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 155391
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 158463
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 161535
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 164863
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 167935
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 171519
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 174847
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 178431
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 181759
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 185599
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 189183
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 193023
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 196863
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 200703
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 204799
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 208895
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 212991
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 217343
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 221439
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 225791
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 230399
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 235007
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 239615
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 244479
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 249599
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 254463
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 259583
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 264703
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 269823
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 275455
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 280831
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 286463
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 292095
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 297983
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 303871
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 310015
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 315903
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 322303
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 328959
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 335615
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 342271
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 348927
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 356095
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 363007
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 370175
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 377599
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 385023
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 392959
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 400639
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 408831
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 416767
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 425215
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 433919
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 442367
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 451071
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 460287
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 469503
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 478719
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 488191
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 497919
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 508159
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 518143
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 528639
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 539135
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 549887
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 560895
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 572159
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 583679
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 595199
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 606975
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 619263
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 631551
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 644095
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 656895
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 670207
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 683519
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 697087
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 711167
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 725247
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 739839
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 754687
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 769535
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 785151
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 800767
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 816895
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 833023
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 849663
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 866559
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 883967
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 901887
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 919551
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 937983
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 956927
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 975871
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 995327
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 1015295
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 1035775
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 1056511
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 1077503
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 1099007
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 1121023
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 1143551
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 1166079
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 1189631
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 1213439
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 1237503
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 1262335
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 1287423
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 1313279
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 1339647
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 1366271
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 1393663
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 1421311
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 1449983
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 1478911
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 1508351
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 1538559
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 1569279
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 1600767
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 1632767
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 1665279
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 1698815
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 1732607
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 1767167
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 1802495
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 1838591
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 1875199
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 1912575
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 1950975
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 1990143
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 2029823
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 2070527
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 2111743
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 2153983
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 2196991
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 2241023
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 2285823
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 2331391
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 2377983
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 2425599
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 2474239
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 2523647
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 2574335
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 2625535
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 2678271
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 2731775
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 2786303
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 2842111
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 2898687
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 2956799
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 3015935
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 3076095
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 3137535
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 3200255
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 3264255
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 3329535
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 3396095
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 3464191
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 3533567
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 3604223
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 3676159
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 3749887
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 3824639
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 3900927
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 3979007
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 4058879
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 4139775
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 4222719
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 4307199
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 4393215
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 4481023
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 4570623
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 4662271
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 4755199
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 4850687
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 4947455
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 5046527
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 5147135
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 5250303
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 5355263
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 5462271
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 5571327
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 5682943
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 5796607
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 5912319
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 6030847
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 6151423
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 6274559
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 6399743
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 6527999
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 6658303
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 6791423
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 6927359
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 7065855
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 7207167
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 7351295
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 7498239
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 7648255
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 7801087
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 7957247
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 8116223
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 8278527
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 8444159
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 8613119
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 8785151
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 8961023
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 9140223
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 9323007
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 9509631
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 9699583
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 9893375
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 10091263
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 10293247
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 10499071
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 10708991
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 10923263
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 11141631
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 11364607
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 11591679
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 11823615
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 12060159
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 12301055
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 12547327
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 12798207
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 13054207
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 13315071
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 13581567
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 13853183
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 14130175
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 14412799
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 14701055
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 14995199
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 15294975
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 15600639
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 15912959
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 16231167
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 16555775
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 16886783
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 17224447
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 17569023
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 17920255
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 18278655
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 18644223
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 19016959
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 19397375
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 19785471
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 20180991
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 20584703
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 20996351
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 21416447
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 21844479
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 22281471
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 22726911
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 23181567
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 23645183
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 24118015
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 24600575
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 25092351
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 25594367
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 26106111
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 26628351
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 27160831
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 27704063
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 28258047
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 28823039
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 29399807
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 29987839
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 30587391
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 31199231
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 31823103
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 32459519
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 33108735
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 33771007
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 34446335
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 35135231
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 35837951
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 36554751
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 37285631
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 38031359
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 38791935
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 39567871
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 40359167
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 41166591
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 41989631
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 42829311
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 43685887
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 44559615
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 45450751
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 46360063
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 47287039
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 48232959
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 49197311
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 50181375
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 51185151
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 52208639
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 53253119
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 54318079
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 55404287
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 56512255
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 57642751
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 58795263
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 59971071
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 61170687
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 62394111
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 63642111
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 64914687
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 66213119
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 67537151
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 68888063
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 70265855
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 71671295
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 73104639
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 74566655
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 76058111
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 77579007
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 79130623
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 80713215
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 82327551
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 83973887
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 85653503
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 87366655
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 89113855
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 90896127
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 92713983
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 94568447
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 96459519
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 98388735
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 100356607
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 102363647
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 104411135
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 106499071
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 108628991
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 110801663
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 113017599
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 115278079
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 117583615
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 119934975
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 122333951
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 124780799
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 127276287
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 129821439
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 132418047
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 135066367
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 137767679
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 140523007
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 143333631
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 146200063
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 149124095
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 152106751
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 155148799
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 158251775
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 161416959
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 164645119
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 167938047
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 171296767
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 174722815
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 178216959
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 181781247
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 185416959
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 189125375
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 192908031
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 196765951
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 200701183
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 204715263
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 208809471
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 212985855
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 217245439
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 221590527
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 226022143
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 230542591
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 235153407
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 239856639
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 244653567
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 249546495
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 254537727
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 259628287
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 264820735
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 270117375
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 275519487
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 281029887
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 286650623
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 292383487
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 298231295
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 304195839
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 310279679
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 316485375
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 322814975
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 329271295
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 335856639
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 342573823
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 349425407
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 356413695
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 363542015
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 370812671
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 378228991
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 385793535
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 393509375
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 401379583
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 409406975
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 417595135
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 425947135
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 434466303
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 443155455
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 452018431
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 461059071
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 470280191
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 479685631
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 489279231
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 499064831
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 509046015
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 519227135
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 529611519
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 540203775
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 551007999
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 562028031
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 573268735
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 584733951
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 596428543
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 608357119
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 620524287
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 632934911
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 645593599
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 658505215
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 671675391
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 685108991
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 698811135
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 712787199
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 727042815
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 741583615
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 756415487
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 771543807
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 786974463
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 802713855
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 818768127
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 835143679
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 851846399
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 868883455
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 886260991
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 903986175
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 922066175
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 940507391
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 959317503
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 978503935
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 998073855
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 1018035455
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 1038396159
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 1059163903
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 1080347135
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 1101954047
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 1123993343
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 1146472959
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 1169402367
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 1192790783
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 1216646399
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 1240979455
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 1265798911
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 1291114751
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 1316937215
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 1343275775
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 1370141439
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 1397543935
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 1425495039
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 1454004991
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 1483084799
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 1512746495
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 1543001599
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 1573861631
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 1605338879
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 1637445631
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 1670194431
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 1703598335
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 1737670143
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 1772423679
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 1807872255
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 1844029439
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 1880910335
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 1918528255
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 1956899071
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 1996036863
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 2035957503
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 2076676863
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 2118210047
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 2160574463
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 2203785983
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 2247861759
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 2292818687
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 2338675199
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 2385448703
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 2433157631
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 2481820671
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 2531457279
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 2582086143
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 2633727999
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 2686402559
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 2740130559
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 2794933247
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 2850831615
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 2907848447
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 2966005503
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 3025325567
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 3085831935
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 3147548671
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 3210499583
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 3274709503
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 3340203519
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 3407007743
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 3475148031
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 3544650751
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 3615544063
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 3687854847
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 3761611775
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 3836844031
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 3913580799
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 3991852543
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 4071689727
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 4153123327
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 4236185855
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 4320909567
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 4407327743
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 4495474175
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 4585383679
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 4677091327
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 4770633215
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 4866045695
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 4963366655
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 5062633983
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 5163886847
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 5267164415
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 5372507647
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 5479957759
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 5589556735
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 5701348095
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 5815375103
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 5931682303
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 6050316287
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 6171322623
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 6294748927
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 6420643839
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 6549056767
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 6680037631
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 6813638655
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 6949911295
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 7088909567
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 7230687743
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 7375301375
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 7522807551
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 7673263615
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 7826728703
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 7983263231
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 8142928895
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 8305787391
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 8471902719
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 8641340927
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 8814167807
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 8990451199
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 2
+  value_type_size: 2
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 1
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 25599
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 26111
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 26623
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 27135
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 27647
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 28159
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 28671
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 29183
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 29695
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 30207
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 30719
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 31231
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 31743
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 32255
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 32767
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 33279
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 33791
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 34303
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 34815
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 35327
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 35839
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 36351
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 36863
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 37887
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 38399
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 38911
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 39423
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 40447
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 41983
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 43007
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 43519
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 44031
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 45055
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 45567
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 46591
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 47103
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 48127
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 49151
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 49663
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 50687
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 51711
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 52223
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 53247
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 54271
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 55295
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 56319
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 56831
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 57855
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 58879
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 59903
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 61439
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 62463
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 63487
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 64511
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 65535
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 67071
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 68095
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 69119
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 70655
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 71679
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 73215
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 74239
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 75775
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 78335
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 79359
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 80895
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 82431
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 83967
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 85503
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 87039
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 88575
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 90111
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 91647
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 93695
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 95231
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 96767
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 98815
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 100351
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 102399
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 104447
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 105983
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 108031
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 110079
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 113151
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 115199
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 117247
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 119807
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 121855
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 123903
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 126463
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 128511
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 131071
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 133631
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 136191
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 138751
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 141311
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 143871
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 146431
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 150015
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 153087
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 156159
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 158719
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 161791
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 164863
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 167935
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 171007
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 174079
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 177663
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 180735
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 185343
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 188927
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 192511
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 196095
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 199679
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 203775
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 207359
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 211455
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 215551
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 220671
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 224767
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 228863
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 232959
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 237567
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 242175
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 246783
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 251391
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 257023
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 262143
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 267263
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 272383
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 277503
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 282623
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 288255
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 294399
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 300031
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 305663
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 311807
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 317439
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 323583
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 330751
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 337407
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 343551
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 350207
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 356863
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 365055
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 371711
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 378879
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 386047
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 393727
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 401919
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 409599
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 417791
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 425471
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 434687
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 443391
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 451583
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 460287
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 470015
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 479231
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 488447
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 497663
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 508415
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 518143
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 528383
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 538111
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 549887
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 560127
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 570879
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 583167
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 594431
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 605695
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 618495
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 630271
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 642559
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 655871
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 668671
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 682495
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 695295
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 709119
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 723455
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 737791
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 751615
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 767487
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 782335
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 798207
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 813567
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 830463
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 846335
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 863743
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 880127
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 898559
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 915967
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 934399
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 952319
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 971775
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 990719
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 1010687
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 1030143
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 1051135
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 1071615
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 1093119
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 1115647
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 1137151
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 1160191
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 1183743
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 1206271
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 1230847
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 1255423
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 1279999
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 1305599
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 1331711
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 1358847
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 1384959
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 1412607
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 1441279
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 1469951
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 1498623
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 1528319
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 1559039
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 1590271
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 1622015
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 1654783
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 1687551
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 1721343
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 1755647
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 1790463
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 1826303
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 1862655
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 1900031
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 1937919
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 1976319
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 2015743
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 2055679
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 2096639
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 2138111
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 2181631
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 2225151
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 2269183
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 2314239
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 2360831
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 2407935
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 2455551
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 2505215
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 2554367
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 2606079
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 2657791
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 2711039
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 2764799
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 2820607
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 2876415
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 2933759
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 2992127
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 3052031
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 3112959
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 3175423
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 3238399
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 3302911
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 3368959
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 3436543
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 3505151
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 3575295
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 3646463
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 3719167
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 3793407
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 3869183
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 3946495
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 4024831
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 4106239
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 4187647
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 4271103
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 4357119
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 4443647
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 4532735
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 4622847
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 4715519
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 4810239
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 4905471
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 5003775
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 5104127
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 5206015
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 5309951
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 5416447
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 5524479
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 5634559
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 5747199
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 5861887
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 5978623
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 6098431
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 6220287
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 6344703
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 6471167
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 6600703
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 6732799
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 6866943
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 7004159
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 7143935
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 7286783
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 7432703
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 7580671
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 7733247
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 7887359
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 8044543
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 8205823
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 8370175
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 8536575
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 8707583
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 8881663
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 9059327
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 9240063
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 9424895
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 9613823
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 9805823
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 10000895
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 10201599
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 10405375
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 10613759
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 10825727
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 11042303
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 11262975
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 11487743
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 11717631
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 11951615
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 12190719
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 12434431
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 12683263
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 12936703
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 13194751
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 13459455
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 13727743
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 14002687
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 14282751
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 14567935
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 14859263
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 15155711
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 15459327
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 15767551
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 16082943
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 16404479
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 16733695
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 17067519
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 17409023
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 17756671
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 18111487
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 18474495
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 18843647
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 19219455
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 19603967
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 19996159
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 20396543
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 20804095
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 21220351
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 21644799
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 22077439
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 22518271
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 22968831
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 23428607
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 23896575
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 24374271
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 24861695
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 25358847
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 25865727
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 26383359
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 26911231
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 27448831
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 27997695
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 28557823
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 29128703
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 29710847
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 30305279
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 30910975
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 31528959
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 32159743
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 32802815
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 33459199
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 34127359
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 34810879
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 35506175
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 36216319
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 36940287
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 37679103
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 38432255
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 39201279
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 39985151
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 40785407
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 41599999
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 42432511
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 43281407
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 44146175
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 45029375
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 45929471
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 46847999
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 47784959
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 48740351
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 49714687
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 50708991
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 51723775
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 52758015
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 53813247
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 54888959
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 55987199
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 57106431
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 58248191
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 59412991
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 60600831
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 61812735
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 63049215
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 64310783
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 65596927
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 66908671
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 68246015
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 69611519
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 71003647
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 72423423
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 73871359
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 75348991
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 76855807
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 78392319
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 79960063
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 81559039
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 83190271
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 84854271
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 86551551
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 88282111
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 90047999
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 91849215
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 93685247
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 95559167
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 97470975
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 99419135
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 101407743
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 103435775
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 105504255
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 107614719
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 109766655
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 111962111
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 114201087
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 116485119
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 118814207
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 121190399
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 123614207
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 126086655
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 128608767
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 131180543
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 133804031
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 136479743
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 139208703
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 141993471
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 144833535
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 147729407
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 150684159
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 153697791
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 156771327
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 159906815
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 163104767
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 166367231
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 169693695
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 173088255
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 176549375
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 180080127
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 183681535
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 187355647
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 191102463
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 194924031
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 198822911
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 202799615
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 206854143
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 210992127
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 215211007
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 219515391
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 223906303
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 228384255
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 232951295
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 237610495
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 242361855
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 247209983
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 252154367
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 257196543
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 262340607
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 267587071
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 272939007
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 278397439
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 283964927
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 289644543
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 295437311
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 301345791
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 307372543
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 313520127
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 319790079
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 326185983
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 332709375
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 339363327
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 346150911
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 353074175
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 360135167
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 367336959
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 374684159
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 382178303
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 389820927
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 397617151
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 405570047
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 413681151
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 421954559
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 430393343
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 439001087
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 447781375
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 456737279
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 465871871
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 475188735
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 484691967
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 494386175
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 504273919
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 514359295
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 524645887
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 535138815
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 545842175
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 556759039
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 567894015
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 579251711
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 590835711
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 602652671
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 614705663
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 627000319
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 639539711
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 652329983
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 665376767
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 678684159
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 692258303
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 706103295
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 720224255
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 734628863
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 749321727
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 764307967
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 779594239
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 795185663
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 811089407
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 827311615
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 843857407
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 860734463
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 877948927
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 895507455
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 913418239
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 931685887
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 950320127
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 969326079
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 988712959
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 1008486399
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 1028656127
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 1049228799
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 1070213631
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 1091618303
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 1113225727
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 1134853631
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 1156914687
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 1179416063
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 1202367999
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 1225779199
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 1249658367
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 1274014719
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 1298858495
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 1324199423
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 1350047231
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 1376411647
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 1403303423
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 1430732799
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 1458711039
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 1487248895
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 1516357119
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 1546047999
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 1576332799
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 1607222783
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 1638730751
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 1670868991
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 1703649791
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 1737086463
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 1771191807
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 1805979135
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 1841462271
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 1877655039
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 1914571775
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 1952226815
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 1990634495
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 2029810687
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 2069770751
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 2110529535
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 2152103935
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 2194513407
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 2237767167
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 2281886207
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 2326887423
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 2372788223
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 2419607551
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 2467363327
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 2516073983
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 2565758975
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 2616437759
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 2668129791
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 2720856063
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 2774636543
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 2829492735
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 2885446143
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 2942518271
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 3000732159
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 3060110335
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 3120675839
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 3182452735
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 3245465599
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 3309742079
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 3375300607
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 3442169855
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 3510376447
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 3579947519
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 3650910207
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 3723291647
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 3797120511
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 3872426495
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 3949238271
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 4027586559
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 2
+  value_type_size: 2
+  is_segmented: true
+  scratch_size_recordings {
+    num_items: 121
+    scratch_space_bytes: 767
+  }
+  scratch_size_recordings {
+    num_items: 251
+    scratch_space_bytes: 1279
+  }
+  scratch_size_recordings {
+    num_items: 381
+    scratch_space_bytes: 1791
+  }
+  scratch_size_recordings {
+    num_items: 511
+    scratch_space_bytes: 2303
+  }
+  scratch_size_recordings {
+    num_items: 630
+    scratch_space_bytes: 2815
+  }
+  scratch_size_recordings {
+    num_items: 762
+    scratch_space_bytes: 3327
+  }
+  scratch_size_recordings {
+    num_items: 889
+    scratch_space_bytes: 3839
+  }
+  scratch_size_recordings {
+    num_items: 1017
+    scratch_space_bytes: 4351
+  }
+  scratch_size_recordings {
+    num_items: 1142
+    scratch_space_bytes: 4863
+  }
+  scratch_size_recordings {
+    num_items: 1258
+    scratch_space_bytes: 5375
+  }
+  scratch_size_recordings {
+    num_items: 1387
+    scratch_space_bytes: 5887
+  }
+  scratch_size_recordings {
+    num_items: 1528
+    scratch_space_bytes: 6399
+  }
+  scratch_size_recordings {
+    num_items: 1652
+    scratch_space_bytes: 6911
+  }
+  scratch_size_recordings {
+    num_items: 1787
+    scratch_space_bytes: 7423
+  }
+  scratch_size_recordings {
+    num_items: 1895
+    scratch_space_bytes: 7935
+  }
+  scratch_size_recordings {
+    num_items: 2009
+    scratch_space_bytes: 8447
+  }
+  scratch_size_recordings {
+    num_items: 2172
+    scratch_space_bytes: 8959
+  }
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 9471
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 9983
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 10495
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 11007
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 11519
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 12031
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 12543
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 13055
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 13567
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 14079
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 14591
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 15103
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 15615
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 16127
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 16639
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 17151
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 17663
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 18175
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 18687
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 19199
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 19711
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 20223
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 20735
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 21247
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 21759
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 22271
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 22783
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 23295
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 23807
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 24319
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 24831
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 25343
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 25855
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 26367
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 26879
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 27391
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 27903
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 28415
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 28927
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 29439
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 29951
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 30463
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 30975
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 31487
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 32511
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 33023
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 33535
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 34047
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 35071
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 35583
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 36607
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 37119
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 37631
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 38655
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 39167
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 40191
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 40703
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 41727
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 42751
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 43263
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 44287
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 45311
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 45823
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 46847
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 47871
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 48895
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 49919
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 50431
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 51455
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 52479
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 53503
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 55039
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 56063
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 57087
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 58111
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 59135
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 60671
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 61695
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 62719
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 64255
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 65279
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 66815
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 67839
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 69375
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 70911
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 71935
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 73471
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 75007
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 76543
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 78079
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 79615
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 81151
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 82687
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 84223
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 86271
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 87807
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 89343
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 91391
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 92927
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 94975
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 97023
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 98559
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 100607
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 102655
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 104703
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 106751
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 108799
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 111359
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 113407
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 115455
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 118015
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 120063
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 122623
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 125183
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 127743
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 130303
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 132863
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 135423
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 137983
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 140543
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 143615
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 146687
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 149247
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 152319
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 155391
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 158463
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 161535
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 164607
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 168191
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 171263
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 174847
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 178431
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 182015
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 185599
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 189183
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 193279
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 196863
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 200959
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 205055
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 209151
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 213247
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 217343
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 221439
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 226047
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 230655
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 235263
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 239871
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 244479
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 249599
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 254719
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 259839
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 264959
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 270079
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 275711
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 280831
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 286463
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 292095
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 298239
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 303871
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 310015
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 316159
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 322815
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 328959
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 335615
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 342271
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 349439
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 356095
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 363263
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 370431
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 378111
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 385279
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 392959
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 401151
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 408831
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 417023
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 425727
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 433919
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 442623
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 451327
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 460543
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 469759
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 478975
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 488703
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 498431
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 508671
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 518399
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 529151
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 539391
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 550143
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 561407
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 572671
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 583935
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 595711
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 607487
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 619775
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 632063
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 644863
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 657663
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 670463
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 684287
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 697599
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 711935
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 725759
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 740607
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 755455
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 770303
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 785663
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 801535
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 817407
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 833791
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 850175
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 867583
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 884991
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 902399
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 920319
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 938751
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 957695
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 976639
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 996095
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 1016063
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 1036543
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 1057023
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 1078527
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 1100031
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 1122047
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 1144575
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 1167103
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 1190655
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 1214207
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 1238783
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 1263359
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 1288447
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 1314559
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 1340671
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 1367295
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 1394943
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 1422591
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 1451263
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 1479935
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 1509631
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 1539839
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 1570559
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 1602303
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 1634047
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 1666815
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 1700095
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 1733887
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 1768703
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 1804031
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 1840383
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 1877247
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 1914623
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 1953023
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 1991935
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 2031871
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 2072319
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 2113791
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 2156287
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 2199295
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 2243327
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 2287871
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 2333951
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 2380543
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 2428159
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 2476287
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 2525951
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 2576639
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 2627839
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 2680575
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 2734335
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 2789119
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 2844415
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 2901759
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 2959615
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 3018495
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 3078911
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 3140863
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 3203327
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 3267327
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 3332863
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 3399423
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 3467519
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 3536639
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 3607295
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 3679487
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 3753215
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 3828479
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 3904767
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 3983103
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 4062463
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 4143871
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 4226815
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 4311295
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 4397311
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 4485375
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 4574975
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 4666623
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 4759807
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 4855039
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 4952319
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 5051135
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 5151999
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 5255423
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 5360383
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 5467391
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 5576959
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 5688575
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 5802239
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 5917951
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 6036735
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 6157055
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 6280447
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 6405887
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 6533887
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 6664959
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 6798079
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 6933759
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 7072511
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 7214335
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 7358207
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 7505663
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 7655679
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 7808767
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 7964927
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 8124159
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 8286463
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 8452351
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 8621311
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 8793855
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 8969471
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 9149183
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 9331967
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 9518847
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 9708799
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 9903359
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 10100991
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 10303231
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 10509055
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 10719487
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 10934015
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 11152639
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 11375359
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 11603199
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 11835135
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 12071679
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 12313343
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 12559615
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 12810495
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 13067007
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 13328127
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 13594879
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 13866751
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 14143743
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 14426879
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 14715135
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 15009535
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 15309567
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 15615743
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 15928063
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 16247039
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 16571647
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 16902911
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 17241343
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 17585919
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 17937663
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 18296575
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 18662143
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 19035391
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 19416319
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 19804415
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 20200703
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 20604671
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 21016831
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 21437183
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 21865727
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 22302975
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 22749439
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 23204095
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 23668479
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 24141567
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 24624383
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 25116927
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 25619199
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 26131711
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 26653951
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 27187455
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 27731199
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 28285695
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 28851455
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 29428479
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 30016767
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 30617343
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 31229695
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 31854335
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 32491263
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 33140991
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 33804031
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 34479871
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 35169535
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 35873023
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 36590335
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 37321983
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 38068479
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 38829823
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 39606527
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 40398591
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 41206527
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 42030847
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 42871039
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 43728639
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 44603135
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 45495039
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 46405375
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 47333119
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 48279807
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 49245439
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 50230527
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 51235071
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 52259583
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 53304575
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 54371071
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 55458559
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 56567551
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 57698559
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 58852607
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 60029695
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 61230335
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 62455039
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 63704319
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 64978175
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 66277631
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 67603199
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 68955391
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 70334207
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 71741183
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 73175807
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 74639615
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 76132095
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 77654783
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 79207679
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 80791807
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 82407679
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 84055807
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 85737215
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 87451903
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 89200895
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 90984703
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 92804351
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 94660863
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 96553727
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 98484991
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 100454655
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 102463743
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 104512767
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 106603263
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 108735231
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 110909695
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 113128191
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 115390719
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 117698303
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 120052479
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 122453247
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 124902655
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 127400703
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 129948415
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 132547327
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 135198463
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 137902335
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 140660479
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 143473407
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 146343167
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 149269759
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 152255231
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 155300607
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 158406399
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 161574655
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 164805887
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 168102143
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 171463935
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 174893311
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 178391295
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 181958911
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 185598207
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 189310207
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 193096447
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 196958463
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 200897279
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 204915455
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 209013503
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 213193983
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 217457919
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 221806847
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 226242815
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 230767871
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 235383039
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 240090879
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 244892927
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 249790719
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 254786303
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 259882239
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 265079551
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 270381311
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 275789055
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 281304831
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 286930687
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 292669183
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 298522879
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 304493311
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 310583039
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 316794623
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 323130623
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 329593087
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 336185087
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 342908671
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 349766911
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 356761855
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 363897087
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 371175167
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 378598655
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 386170623
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 393894143
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 401771775
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 409807615
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 418003711
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 426363647
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 434891007
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 443588863
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 452460287
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 461509887
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 470739711
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 480154367
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 489757439
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 499553023
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 509543679
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 519734527
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 530129663
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 540732159
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 551546623
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 562577663
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 573828863
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 585305855
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 597011711
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 608952063
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 621131007
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 633553663
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 646224639
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 659149055
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 672332031
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 685778687
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 699494143
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 713484031
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 727753983
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 742309119
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 757155071
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 772297983
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 787743999
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 803499263
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 819568895
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 835960575
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 852679423
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 869733119
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 887127807
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 904870143
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 922967807
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 941426943
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 960255743
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 979460863
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 999049983
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 1019030783
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 1039411455
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 1060199679
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 1081403647
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 1103031551
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 1125092607
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 1147593983
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 1170545919
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 1193957119
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 1217836287
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 1242192639
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 1267036415
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 1292377343
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 1318225151
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 1344589567
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 1371481343
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 1398910719
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 1426888959
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 1455426815
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 1484535039
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 1514225919
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 1544510719
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 1575400703
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 1606908671
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 1639046911
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 1671827711
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 1705264383
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 1739369727
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 1774157055
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 1809640191
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 1845832959
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 1882749695
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 1920404735
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 1958812415
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 1997988607
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 2037948671
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 2078707455
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 2120281855
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 2162687231
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 2205940991
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 2250060031
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 2295061247
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 2340962047
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 2387781375
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 2435537151
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 2484247807
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 2533932799
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 2584611583
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 2636303615
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 2689029887
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 2742810367
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 2797666559
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 2853619967
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 2910692095
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 2968905983
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 3028284159
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 3088849663
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 3150626559
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 3213639423
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 3277911807
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 3343470335
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 3410339583
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 3478546175
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 3548117247
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 3619079935
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 3691461375
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 3765290239
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 3840596223
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 3917407999
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 3995756287
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 2
+  value_type_size: 4
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 1
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 35583
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 36095
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 36863
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 37631
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 37887
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 38655
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 39423
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 39935
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 40703
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 41471
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 42239
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 43007
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 43775
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 44543
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 45311
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 46079
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 46847
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 47615
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 48383
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 49151
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 49919
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 50943
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 51711
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 52479
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 53759
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 54527
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 55551
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 56319
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 57599
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 58623
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 59903
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 60671
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 62719
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 63999
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 65023
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 66303
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 67327
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 68607
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 70143
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 71167
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 72447
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 73983
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 75007
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 76543
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 77823
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 79359
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 80895
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 81919
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 83455
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 84991
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 86527
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 88575
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 90111
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 91647
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 93183
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 94975
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 97023
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 98559
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 100351
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 102399
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 104191
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 106239
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 108031
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 110079
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 112383
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 114175
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 116479
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 119551
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 121855
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 124159
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 126463
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 128767
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 131327
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 133631
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 136447
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 138751
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 141311
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 144127
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 146687
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 149503
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 152575
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 155135
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 158207
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 161279
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 164351
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 167423
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 170495
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 175103
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 178175
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 181503
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 185087
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 188415
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 192255
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 195839
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 199679
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 203519
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 207359
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 211199
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 215295
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 219135
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 223487
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 228095
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 233215
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 237823
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 242175
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 247039
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 251647
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 256255
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 261375
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 266239
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 271615
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 276735
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 282111
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 288511
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 294143
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 300031
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 305663
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 311551
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 317695
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 323839
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 329983
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 336127
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 342527
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 350463
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 357119
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 364031
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 371199
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 378111
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 385535
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 393215
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 401919
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 409599
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 417279
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 425727
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 433663
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 442111
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 450559
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 460543
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 469247
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 478463
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 487679
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 497407
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 506879
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 517887
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 527871
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 538367
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 548607
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 559359
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 571135
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 582399
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 593407
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 604927
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 616959
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 629759
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 642047
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 654847
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 667391
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 680447
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 694527
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 708351
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 722175
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 735999
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 751615
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 766207
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 781311
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 797183
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 813055
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 828671
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 844799
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 862463
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 879359
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 896511
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 914943
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 932863
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 951039
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 970495
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 989695
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 1008895
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 1029375
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 1049855
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 1070079
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 1092351
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 1113343
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 1136383
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 1158655
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 1181183
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 1205247
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 1228799
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 1253887
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 1278207
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 1304063
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 1329919
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 1356031
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 1383423
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 1410303
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 1438975
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 1467135
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 1496831
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 1527039
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 1556991
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 1588479
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 1619455
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 1652479
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 1684735
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 1718783
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 1753599
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 1787647
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 1823743
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 1859327
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 1896959
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 1935103
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 1972735
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 2012671
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 2053119
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 2094079
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 2135295
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 2178047
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 2221823
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 2266111
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 2310655
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 2356991
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 2404095
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 2452479
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 2501375
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 2551551
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 2602495
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 2654207
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 2706431
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 2760447
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 2815743
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 2872063
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 2929151
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 2987775
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 3048191
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 3109119
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 3171071
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 3234047
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 3298815
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 3364351
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 3431423
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 3500543
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 3570431
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 3641343
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 3713791
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 3788287
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 3863807
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 3940607
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 4019711
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 4099839
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 4182271
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 4265471
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 4350719
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 4437503
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 4526335
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 4615935
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 4708607
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 4803327
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 4898303
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 4996351
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 5096447
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 5198591
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 5302527
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 5407487
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 5515519
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 5625855
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 5738495
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 5853183
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 5969919
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 6090239
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 6211583
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 6335487
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 6461951
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 6590719
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 6723071
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 6856959
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 6994687
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 7133951
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 7277055
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 7421951
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 7570687
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 7722239
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 7875583
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 8033535
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 8194047
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 8357887
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 8525055
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 8695551
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 8869119
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 9046015
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 9227007
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 9411839
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 9599743
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 9791231
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 9987327
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 10186751
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 10390527
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 10598399
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 10809599
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 11026175
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 11246335
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 11471359
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 11700479
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 11934207
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 12173567
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 12416511
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 12664319
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 12918015
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 13175807
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 13439487
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 13708287
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 13982719
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 14262271
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 14547455
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 14837759
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 15134463
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 15436287
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 15745535
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 16060671
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 16381183
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 16709119
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 17043199
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 17383679
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 17731327
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 18086399
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 18447615
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 18817023
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 19192575
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 19576319
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 19967999
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 20367103
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 20774143
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 21190143
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 21613055
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 22045695
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 22486527
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 22936319
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 23394559
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 23862015
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 24339711
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 24826111
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 25322495
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 25828607
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 26345215
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 26872575
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 27409407
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 27957759
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 28516607
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 29086719
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 29668095
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 30261759
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 30866175
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 31483647
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 32113919
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 32755711
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 33411071
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 34078207
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 34760959
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 35455487
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 36165119
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 36887295
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 37624831
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 38377727
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 39144703
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 39927551
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 40726527
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 41540863
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 42371839
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 43218943
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 44082943
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 44964863
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 45863935
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 46780927
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 47716863
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 48671231
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 49644031
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 50636287
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 51649023
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 52682239
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 53736191
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 54810879
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 55906303
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 57024511
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 58164735
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 59328255
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 60514559
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 61724159
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 62958591
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 64218623
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 65501695
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 66812415
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 68147967
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 69511423
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 70901247
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 72318719
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 73765375
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 75240447
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 76745471
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 78279935
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 79845631
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 81441791
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 83070975
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 84732927
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 86427135
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 88155647
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 89918463
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 91716863
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 93550591
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 95422207
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 97330943
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 99276543
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 101261823
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 103287039
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 105352959
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 107459327
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 109609471
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 111801599
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 114036991
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 116317951
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 118643967
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 121016319
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 123436543
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 125905407
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 128423679
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 130991615
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 133611775
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 136283391
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 139009535
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 141789183
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 144624895
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 147517695
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 150467583
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 153476351
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 156546303
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 159677439
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 162870271
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 166127615
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 169449983
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 172838911
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 176296191
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 179821311
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 183417855
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 187085823
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 190828287
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 194643967
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 198537471
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 202507519
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 206557695
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 210688511
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 214902271
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 219199743
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 223583999
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 228056063
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 232617215
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 237268991
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 242014719
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 246854143
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 251791359
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 256827647
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 261963775
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 267203071
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 272547071
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 277997055
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 283557119
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 289228543
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 295013375
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 300912895
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 306930943
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 313069567
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 319330559
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 325717759
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 332231423
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 338875903
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 345653759
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 352566271
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 359618047
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 366809599
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 374146815
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 381628671
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 389261567
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 397046527
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 404987135
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 413086975
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 421349119
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 429775359
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 438371071
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 447138559
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 456080639
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 465202175
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 474506751
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 483996159
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 493676287
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 503549183
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 513620479
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 523892223
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 534371071
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 545057279
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 555958783
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 567077631
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 578419711
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 589987583
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 601787647
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 613822463
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 626099199
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 638621695
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 651393023
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 664421631
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 677710079
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 691263231
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 705089279
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 719190271
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 733574399
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 748245247
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 763210751
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 778474495
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 794043391
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 809924607
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 826123263
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 842645247
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 859497727
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 876688127
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 894221823
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 912105471
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 930347519
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 948954879
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 967933951
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 987291903
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 1007037951
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 1027179007
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 1047721727
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 1068675839
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 1090049791
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 1111850751
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 1134087935
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 1156769791
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 1179903999
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 1203502079
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 1227572991
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 1252124159
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 1277166079
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 1302709503
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 1328764159
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 1355338239
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 1382445567
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 1410094591
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 1438295551
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 1467061759
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 1496402943
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 1526330367
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 1556856831
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 1587994111
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 1619753983
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 1651939839
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 1684381695
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 1717473023
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 1751225343
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 1785653247
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 1820769791
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 1856588543
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 1893123327
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 1930388991
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 1968400383
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 2007171839
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 2046718463
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 2087056127
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 2128200447
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 2170167807
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 2212974335
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 2256636927
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 2301173247
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 2346600191
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 2392935167
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 2440197375
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 2488404479
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 2537575935
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 2587730687
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 2638888703
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 2691069695
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 2744294399
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 2798583551
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 2853958655
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 2910441215
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 2968052991
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 3026817279
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 3086757119
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 3147895551
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 3210256895
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 3273869311
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 3338749951
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 3404928255
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 3472430079
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 3541281535
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 3611510527
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 3683143935
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 3756209919
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 3830737407
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 3906755583
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 3984293631
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 4063383039
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 4144053759
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 4226338047
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 4310268159
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 4395876607
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 4483197183
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 4572264447
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 4663112959
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 4755778303
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 4850297343
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 4946710271
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 5045047807
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 5145351935
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 5247661823
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 5352018431
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 5458462207
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 5567034367
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 5677777919
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 5790736895
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 5905954559
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 6023476991
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 2
+  value_type_size: 4
+  is_segmented: true
+  scratch_size_recordings {
+    num_items: 61
+    scratch_space_bytes: 767
+  }
+  scratch_size_recordings {
+    num_items: 121
+    scratch_space_bytes: 1023
+  }
+  scratch_size_recordings {
+    num_items: 191
+    scratch_space_bytes: 1535
+  }
+  scratch_size_recordings {
+    num_items: 251
+    scratch_space_bytes: 1791
+  }
+  scratch_size_recordings {
+    num_items: 311
+    scratch_space_bytes: 2303
+  }
+  scratch_size_recordings {
+    num_items: 381
+    scratch_space_bytes: 2559
+  }
+  scratch_size_recordings {
+    num_items: 441
+    scratch_space_bytes: 3071
+  }
+  scratch_size_recordings {
+    num_items: 511
+    scratch_space_bytes: 3327
+  }
+  scratch_size_recordings {
+    num_items: 573
+    scratch_space_bytes: 3839
+  }
+  scratch_size_recordings {
+    num_items: 630
+    scratch_space_bytes: 4095
+  }
+  scratch_size_recordings {
+    num_items: 693
+    scratch_space_bytes: 4607
+  }
+  scratch_size_recordings {
+    num_items: 762
+    scratch_space_bytes: 4863
+  }
+  scratch_size_recordings {
+    num_items: 823
+    scratch_space_bytes: 5375
+  }
+  scratch_size_recordings {
+    num_items: 889
+    scratch_space_bytes: 5631
+  }
+  scratch_size_recordings {
+    num_items: 960
+    scratch_space_bytes: 6143
+  }
+  scratch_size_recordings {
+    num_items: 1017
+    scratch_space_bytes: 6399
+  }
+  scratch_size_recordings {
+    num_items: 1078
+    scratch_space_bytes: 6911
+  }
+  scratch_size_recordings {
+    num_items: 1142
+    scratch_space_bytes: 7167
+  }
+  scratch_size_recordings {
+    num_items: 1210
+    scratch_space_bytes: 7679
+  }
+  scratch_size_recordings {
+    num_items: 1258
+    scratch_space_bytes: 7935
+  }
+  scratch_size_recordings {
+    num_items: 1334
+    scratch_space_bytes: 8447
+  }
+  scratch_size_recordings {
+    num_items: 1387
+    scratch_space_bytes: 8703
+  }
+  scratch_size_recordings {
+    num_items: 1470
+    scratch_space_bytes: 9215
+  }
+  scratch_size_recordings {
+    num_items: 1528
+    scratch_space_bytes: 9471
+  }
+  scratch_size_recordings {
+    num_items: 1589
+    scratch_space_bytes: 9983
+  }
+  scratch_size_recordings {
+    num_items: 1652
+    scratch_space_bytes: 10239
+  }
+  scratch_size_recordings {
+    num_items: 1718
+    scratch_space_bytes: 10751
+  }
+  scratch_size_recordings {
+    num_items: 1787
+    scratch_space_bytes: 11007
+  }
+  scratch_size_recordings {
+    num_items: 1822
+    scratch_space_bytes: 11519
+  }
+  scratch_size_recordings {
+    num_items: 1895
+    scratch_space_bytes: 11775
+  }
+  scratch_size_recordings {
+    num_items: 1970
+    scratch_space_bytes: 12287
+  }
+  scratch_size_recordings {
+    num_items: 2009
+    scratch_space_bytes: 12543
+  }
+  scratch_size_recordings {
+    num_items: 2089
+    scratch_space_bytes: 13055
+  }
+  scratch_size_recordings {
+    num_items: 2172
+    scratch_space_bytes: 13311
+  }
+  scratch_size_recordings {
+    num_items: 2215
+    scratch_space_bytes: 13823
+  }
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 14079
+  }
+  scratch_size_recordings {
+    num_items: 2350
+    scratch_space_bytes: 14591
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 14847
+  }
+  scratch_size_recordings {
+    num_items: 2492
+    scratch_space_bytes: 15359
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 15615
+  }
+  scratch_size_recordings {
+    num_items: 2591
+    scratch_space_bytes: 16127
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 16383
+  }
+  scratch_size_recordings {
+    num_items: 2747
+    scratch_space_bytes: 16895
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 17151
+  }
+  scratch_size_recordings {
+    num_items: 2857
+    scratch_space_bytes: 17663
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 17919
+  }
+  scratch_size_recordings {
+    num_items: 2972
+    scratch_space_bytes: 18431
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 18687
+  }
+  scratch_size_recordings {
+    num_items: 3091
+    scratch_space_bytes: 19199
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 19455
+  }
+  scratch_size_recordings {
+    num_items: 3215
+    scratch_space_bytes: 19967
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 20223
+  }
+  scratch_size_recordings {
+    num_items: 3344
+    scratch_space_bytes: 20735
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 20991
+  }
+  scratch_size_recordings {
+    num_items: 3478
+    scratch_space_bytes: 21503
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 21759
+  }
+  scratch_size_recordings {
+    num_items: 3617
+    scratch_space_bytes: 22271
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 22527
+  }
+  scratch_size_recordings {
+    num_items: 3762
+    scratch_space_bytes: 23039
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 23295
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 24063
+  }
+  scratch_size_recordings {
+    num_items: 3991
+    scratch_space_bytes: 24575
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 24831
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 25343
+  }
+  scratch_size_recordings {
+    num_items: 4234
+    scratch_space_bytes: 26111
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 26367
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 26879
+  }
+  scratch_size_recordings {
+    num_items: 4492
+    scratch_space_bytes: 27647
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 27903
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 28415
+  }
+  scratch_size_recordings {
+    num_items: 4765
+    scratch_space_bytes: 29183
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 29439
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 30207
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 30719
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 31487
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 32255
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 32511
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 33279
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 34047
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 34559
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 35327
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 36095
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 36863
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 37631
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 38399
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 39167
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 39935
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 40703
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 41471
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 42239
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 43007
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 43775
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 44543
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 45567
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 46335
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 47103
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 48383
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 49151
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 50175
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 50943
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 52223
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 53247
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 54527
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 55295
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 56319
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 57599
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 58623
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 59903
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 60927
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 62207
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 63743
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 64767
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 66047
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 67583
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 68607
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 70143
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 71423
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 72959
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 74495
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 75519
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 77055
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 78591
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 80127
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 82175
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 83711
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 85247
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 86783
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 88575
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 90623
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 92159
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 93951
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 95999
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 97791
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 99839
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 101631
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 103679
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 105983
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 107775
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 110079
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 112127
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 114431
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 116735
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 119039
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 121343
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 123903
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 126207
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 129023
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 131327
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 133887
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 136703
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 139263
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 142079
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 145151
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 147711
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 150783
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 153855
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 156927
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 159999
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 163071
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 166655
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 169727
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 173055
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 176639
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 179967
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 183807
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 187391
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 191231
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 195071
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 198911
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 202751
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 206847
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 210687
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 215039
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 219647
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 223743
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 228351
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 232703
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 237567
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 242175
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 246783
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 251903
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 256767
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 262143
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 267263
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 272639
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 278015
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 283647
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 289535
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 295167
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 301055
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 307199
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 313343
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 319487
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 325631
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 332031
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 338943
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 345599
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 352511
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 359679
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 366591
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 374015
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 381695
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 389375
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 397055
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 404735
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 413183
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 421119
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 429567
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 438015
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 446975
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 455679
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 464895
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 474111
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 483839
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 493311
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 503295
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 513279
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 523775
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 534015
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 544767
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 555519
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 566783
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 577791
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 589311
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 601343
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 613119
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 625407
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 638207
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 650751
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 663807
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 676863
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 690687
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 704511
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 718335
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 732927
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 747519
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 762623
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 777471
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 793343
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 808959
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 825087
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 841727
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 858623
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 875775
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 893183
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 911103
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 929279
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 947711
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 966911
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 986111
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 1005567
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 1026047
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 1046271
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 1067519
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 1088511
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 1110527
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 1132799
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 1155327
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 1178367
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 1201919
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 1225983
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 1250303
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 1275135
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 1300991
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 1327103
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 1353471
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 1380351
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 1407999
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 1436159
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 1464831
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 1494015
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 1523967
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 1554431
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 1585407
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 1617407
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 1649663
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 1682687
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 1716479
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 1750527
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 1785599
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 1821183
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 1857791
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 1894911
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 1932543
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 1971455
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 2010879
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 2050815
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 2092031
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 2133759
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 2176511
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 2219775
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 2264319
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 2309631
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 2355711
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 2403071
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 2450943
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 2500095
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 2550015
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 2600703
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 2652927
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 2705919
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 2760191
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 2815487
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 2871551
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 2929151
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 2987519
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 3047423
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 3108351
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 3170303
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 3234047
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 3298559
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 3364607
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 3431679
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 3500543
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 3570431
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 3641855
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 3714303
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 3788799
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 3864575
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 3941631
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 4020735
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 4101119
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 4183295
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 4266495
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 4352255
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 4439039
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 4527615
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 4618239
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 4710911
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 4804863
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 4900863
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 4998911
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 5099007
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 5200895
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 5304831
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 5410815
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 5519103
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 5629695
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 5742335
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 5857023
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 5974271
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 6093567
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 6215423
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 6339839
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 6466559
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 6595839
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 6727679
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 6862335
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 6999551
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 7139583
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 7282431
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 7428095
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 7576575
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 7727871
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 7882751
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 8040191
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 8200959
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 8365055
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 8532479
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 8702975
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 8876799
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 9054719
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 9235455
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 9420287
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 9608703
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 9800703
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 9997055
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 10196735
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 10400511
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 10608639
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 10821119
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 11037183
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 11258111
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 11483135
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 11712767
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 11947007
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 12185855
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 12429567
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 12678143
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 12931839
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 13190399
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 13454079
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 13723391
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 13997823
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 14277887
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 14563071
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 14854655
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 15151359
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 15454463
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 15763455
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 16078847
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 16400639
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 16728575
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 17062911
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 17404415
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 17752319
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 18107391
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 18469631
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 18839039
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 19215615
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 19600127
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 19992063
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 20391935
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 20799743
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 21215487
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 21639935
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 22072575
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 22514175
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 22964223
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 23423487
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 23891967
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 24370175
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 24857343
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 25354239
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 25861631
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 26378751
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 26906367
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 27444479
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 27993087
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 28552959
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 29124095
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 29706495
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 30300671
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 30906879
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 31524863
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 32155391
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 32798463
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 33454335
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 34123775
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 34806015
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 35502335
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 36212223
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 36936447
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 37675007
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 38428671
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 39197183
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 39980799
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 40780799
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 41596415
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 42428159
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 43276799
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 44142335
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 45025023
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 45925631
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 46844159
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 47781119
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 48736511
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 49711359
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 50705663
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 51719423
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 52753919
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 53809151
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 54885119
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 55982847
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 57102591
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 58244607
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 59409407
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 60597503
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 61809407
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 63045887
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 64306431
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 65592575
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 66904575
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 68242431
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 69607679
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 70999551
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 72419583
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 73868031
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 75345407
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 76852223
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 78389247
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 79956735
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 81556223
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 83187455
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 84850943
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 86547711
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 88278783
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 90044415
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 91845375
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 93682175
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 95556095
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 97466879
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 99416319
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 101404671
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 103432703
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 105501183
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 107611391
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 109763583
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 111959039
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 114198015
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 116482047
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 118811391
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 121187583
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 123611391
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 126083583
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 128605439
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 131177471
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 133800959
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 136476927
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 139206399
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 141990911
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 144830463
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 147727103
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 150681599
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 153695231
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 156769023
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 159904511
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 163102463
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 166364415
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 169691903
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 173085695
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 176547327
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 180078335
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 183679743
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 187353599
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 191100671
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 194922495
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 198820863
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 202797311
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 206853375
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 210990335
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 215209983
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 219514367
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 223904511
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 228382719
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 232950527
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 237609215
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 242361599
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 247208703
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 252152831
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 257195775
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 262339839
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 267586559
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 272938239
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 278396927
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 283964927
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 289644287
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 295437311
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 301345791
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 307372799
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 313520127
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 319790591
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 326186495
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 332710143
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 339364095
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 346151679
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 353074431
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 360135935
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 367339007
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 374685695
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 382179327
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 389822975
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 397619199
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 405571583
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 413683199
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 421956863
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 430395903
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 439003647
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 447783935
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 456739583
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 465874175
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 475191551
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 484695551
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 494389247
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 504277247
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 514362623
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 524649983
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 535142655
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 545845503
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 556762623
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 567897855
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 579255807
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 590840831
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 602657535
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 614711039
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 627005183
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 639545087
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 652336127
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 665382911
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 678690303
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 692264447
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 706109439
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 720231423
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 734636031
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 749329151
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 764315391
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 779601663
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 795194111
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 811097855
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 827319807
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 843866111
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 860743167
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 877958399
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 895517439
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 913427711
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 931696127
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 950330111
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 969336831
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 988723455
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 1008497919
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 1028667903
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 1049241087
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 1070225919
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 1091630591
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 1113463295
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 1135732479
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 1158446847
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 1181615871
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 1205248511
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 1229353215
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 1253940479
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 1279019007
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 1304599551
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 1330691327
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 1357305087
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 1384451327
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 1412140287
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 1440383231
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 1469190911
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 1498574591
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 1528546047
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 1559117055
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 1590299391
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 1622105343
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 1654547199
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 1687638527
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 1721390847
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 1755818751
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 1790935295
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 1826754047
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 1863288831
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 1900554495
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 1938565887
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 1977337343
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 2016883967
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 2057221631
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 2098365951
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 2140333311
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 2183139839
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 2226802431
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 2271338751
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 2316765695
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 2363100671
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 2410362879
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 2458569983
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 2507741439
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 2557896191
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 2609054207
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 2661235199
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 2714459903
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 2768749055
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 2824124159
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 2880606719
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 2938218495
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 2996982783
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 3056922623
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 3118061055
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 3180422399
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 3244030719
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 3308911359
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 3375089663
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 3442591487
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 3511442943
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 3581671935
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 3653305343
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 3726371327
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 3800898815
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 3876916991
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 3954455039
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 4033544447
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 4114215167
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 4196499455
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 4280429567
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 4366038015
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 4453358591
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 4542425855
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 4633274367
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 4725939711
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 4820458751
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 4916867583
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 5015205119
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 5115509247
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 5217819135
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 5322175743
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 5428619519
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 5537191679
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 5647935231
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 5760894207
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 5876111871
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 5993634303
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 2
+  value_type_size: 8
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 1
+  }
+  scratch_size_recordings {
+    num_items: 2350
+    scratch_space_bytes: 29439
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 29695
+  }
+  scratch_size_recordings {
+    num_items: 2444
+    scratch_space_bytes: 30463
+  }
+  scratch_size_recordings {
+    num_items: 2492
+    scratch_space_bytes: 30719
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 31231
+  }
+  scratch_size_recordings {
+    num_items: 2591
+    scratch_space_bytes: 31743
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 32255
+  }
+  scratch_size_recordings {
+    num_items: 2694
+    scratch_space_bytes: 33023
+  }
+  scratch_size_recordings {
+    num_items: 2747
+    scratch_space_bytes: 33279
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 33791
+  }
+  scratch_size_recordings {
+    num_items: 2857
+    scratch_space_bytes: 34559
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 35071
+  }
+  scratch_size_recordings {
+    num_items: 2972
+    scratch_space_bytes: 35583
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 36095
+  }
+  scratch_size_recordings {
+    num_items: 3091
+    scratch_space_bytes: 36863
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 37375
+  }
+  scratch_size_recordings {
+    num_items: 3215
+    scratch_space_bytes: 38143
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 38655
+  }
+  scratch_size_recordings {
+    num_items: 3344
+    scratch_space_bytes: 39423
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 39935
+  }
+  scratch_size_recordings {
+    num_items: 3478
+    scratch_space_bytes: 40703
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 41215
+  }
+  scratch_size_recordings {
+    num_items: 3617
+    scratch_space_bytes: 42239
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 42751
+  }
+  scratch_size_recordings {
+    num_items: 3762
+    scratch_space_bytes: 43519
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 44031
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 45055
+  }
+  scratch_size_recordings {
+    num_items: 3991
+    scratch_space_bytes: 45823
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 46591
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 47359
+  }
+  scratch_size_recordings {
+    num_items: 4234
+    scratch_space_bytes: 48383
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 48895
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 49919
+  }
+  scratch_size_recordings {
+    num_items: 4492
+    scratch_space_bytes: 50943
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 51711
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 52479
+  }
+  scratch_size_recordings {
+    num_items: 4765
+    scratch_space_bytes: 53503
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 54271
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 55295
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 56319
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 57599
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 58623
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 59391
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 60415
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 61695
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 63743
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 65023
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 66303
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 67327
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 68607
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 69887
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 71167
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 72447
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 73727
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 75007
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 76287
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 77567
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 79103
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 80383
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 81919
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 83455
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 84735
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 86527
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 88063
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 89599
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 91135
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 93183
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 94719
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 96767
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 98303
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 100095
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 102143
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 103935
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 105983
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 107775
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 109823
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 112127
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 113919
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 116223
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 118527
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 121343
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 123903
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 126207
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 128511
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 131071
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 133119
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 135679
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 138239
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 140799
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 143871
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 146431
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 149247
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 151807
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 154623
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 157951
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 160767
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 163839
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 166911
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 169983
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 173311
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 176639
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 180991
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 184575
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 187903
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 191487
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 195071
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 198911
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 202751
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 206591
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 210431
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 214527
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 218623
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 222975
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 227071
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 231423
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 237055
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 241407
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 246015
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 250879
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 255487
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 260351
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 265471
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 270591
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 275711
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 281087
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 286719
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 293119
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 298751
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 304639
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 310271
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 316415
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 322559
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 328703
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 335103
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 341503
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 348159
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 355839
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 362495
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 369663
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 377087
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 384255
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 391679
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 399103
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 408063
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 415743
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 423679
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 432127
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 440319
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 449023
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 457727
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 467711
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 476671
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 485887
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 495615
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 505087
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 514815
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 526079
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 536319
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 546559
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 557055
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 567807
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 580095
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 591359
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 602879
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 614655
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 626431
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 639743
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 652287
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 665087
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 677887
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 690943
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 705791
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 719359
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 733183
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 747519
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 763391
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 777983
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 793087
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 809727
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 825599
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 841727
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 858111
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 875775
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 893183
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 910335
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 929279
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 947199
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 965887
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 985599
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 1004799
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 1024767
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 1045503
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 1065983
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 1086975
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 1108991
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 1130751
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 1153791
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 1176575
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 1199615
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 1223935
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 1247999
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 1273343
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 1298431
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 1323519
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 1350655
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 1377023
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 1404927
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 1432575
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 1461503
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 1490175
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 1520383
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 1550079
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 1581567
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 1613311
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 1645055
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 1678335
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 1710847
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 1745663
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 1779711
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 1815807
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 1852159
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 1888511
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 1926655
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 1965311
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 2003711
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 2043903
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 2085119
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 2125823
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 2168319
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 2212095
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 2256639
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 2301695
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 2346751
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 2393599
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 2441727
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 2490367
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 2540287
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 2591231
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 2641919
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 2694655
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 2748671
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 2803711
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 2859519
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 2916863
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 2974719
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 3034367
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 3094783
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 3156479
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 3220479
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 3284479
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 3350015
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 3416831
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 3484671
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 3554047
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 3625727
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 3697919
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 3771391
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 3846399
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 3923967
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 4002047
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 4082687
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 4163583
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 4246527
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 4331775
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 4417535
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 4506367
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 4596991
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 4688127
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 4782335
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 4877055
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 4974847
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 5074431
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 5176063
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 5278719
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 5384191
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 5492223
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 5601791
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 5713919
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 5827839
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 5944575
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 6063103
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 6184191
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 6307327
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 6434303
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 6562815
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 6693375
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 6827007
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 6964223
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 7102975
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 7245055
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 7389695
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 7537663
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 7687423
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 7841535
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 7998719
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 8158719
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 8321791
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 8487167
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 8656895
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 8829951
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 9006335
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 9187071
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 9370623
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 9557759
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 9748479
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 9943807
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 10142207
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 10345471
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 10551551
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 10763007
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 10978303
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 11197695
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 11421695
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 11650047
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 11883007
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 12120319
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 12362495
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 12609279
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 12861951
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 13118463
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 13381375
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 13648127
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 13921023
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 14199807
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 14483711
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 14773247
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 15068159
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 15369471
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 15677183
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 15990271
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 16310271
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 16636671
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 16968703
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 17307903
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 17653759
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 18006783
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 18367743
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 18734335
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 19109375
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 19490815
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 19880703
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 20278271
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 20683519
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 21096959
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 21519359
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 21949183
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 22388479
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 22836223
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 23292927
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 23758591
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 24233471
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 24717311
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 25212159
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 25716223
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 26230527
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 26754815
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 27289599
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 27835903
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 28392703
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 28959487
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 29539583
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 30129919
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 30731775
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 31346687
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 31973631
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 32612607
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 33265663
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 33930239
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 34608895
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 35300863
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 36006399
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 36726783
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 37461503
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 38210559
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 38974207
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 39753983
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 40549119
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 41360127
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 42186495
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 43030271
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 43890943
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 44768255
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 45663999
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 46577407
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 47508223
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 48457727
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 49426943
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 50416127
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 51423743
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 52452351
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 53500927
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 54571263
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 55662335
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 56775167
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 57911039
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 59068671
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 60250623
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 61455359
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 62684927
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 63937535
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 65217023
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 66520319
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 67851263
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 69207807
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 70591743
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 72003839
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 73443583
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 74912767
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 76410623
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 77939199
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 79497727
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 81086975
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 82709247
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 84363519
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 86050303
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 87770879
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 89526783
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 91316991
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 93143295
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 95005439
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 96905471
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 98843647
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 100820991
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 102836991
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 104893695
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 106991103
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 109130239
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 111313663
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 113539071
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 115810047
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 118126079
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 120488959
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 122898687
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 125355775
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 127863039
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 130420735
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 133028351
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 135689215
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 138402815
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 141170687
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 143993599
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 146873855
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 149811199
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 152806911
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 155862783
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 158980607
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 162159871
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 165403135
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 168710399
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 172084735
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 175526399
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 179037439
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 182617855
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 186270207
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 189995775
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 193795071
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 197671423
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 201623551
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 205656575
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 209768959
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 213965055
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 218243839
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 222608383
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 227060735
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 231601663
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 236233727
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 240958207
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 245776895
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 250692863
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 255707135
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 260820991
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 266037247
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 271358207
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 276784383
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 282320383
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 287966463
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 293726207
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 299599871
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 305592575
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 311703551
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 317937919
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 324296703
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 330782207
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 337397503
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 344145663
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 351028479
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 358048511
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 365209343
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 372514303
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 379963391
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 387563007
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 395314175
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 403220223
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 411284991
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 419510015
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 427899903
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 436458495
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 445187327
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 454090751
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 463171839
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 472435711
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 481884671
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 491522303
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 501352191
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 511378943
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 521607423
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 532038399
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 542679551
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 553532671
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 564603135
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 575895039
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 587413247
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 599161343
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 611144447
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 623367423
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 635834111
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 648550911
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 661521663
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 674751999
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 688246783
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 702012671
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 716051967
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 730373119
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 744980479
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 759880447
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 775077887
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 790578943
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 806390527
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 822518015
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 838968319
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 855747071
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 872862463
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 890319103
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 908125695
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 926288383
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 944814079
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 963709695
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 982983679
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 1002643455
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 1022696447
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 1043150847
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 1064012799
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 1085292799
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 1106999295
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 1129138943
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 1151721727
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 1174756607
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 1198250495
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 1222215423
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 1246659327
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 1271592959
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 1297025023
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 1322964991
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 1349425151
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 1376412671
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 1403940863
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 1432019967
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 1460659455
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 1489873407
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 1519670015
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 1550063359
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 1581064703
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 1612685567
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 1644939775
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 1677838591
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 1711395071
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 1745623039
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 1780534527
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 1816145151
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 1852468735
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 1889517823
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 1927308031
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 1965853695
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 2005171199
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 2045274367
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 2086179071
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 2127903231
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 2170460671
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 2213870079
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 2258147071
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 2303310335
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 2349376255
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 2396364031
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 2444290559
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 2493176831
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 2543040767
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 2593900287
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 2645778687
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 2698694399
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 2752322815
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 2806392831
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 2861544703
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 2917798911
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 2975178751
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 3033705983
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 3093403903
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 3154295551
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 3216404991
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 3279757055
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 3344376063
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 3410287103
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 3477516543
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 3546090495
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 3616036095
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 3687380479
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 3760151551
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 3834378495
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 3910089983
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 3987315199
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 4066085375
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 4146430719
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 4228382975
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 4311974399
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 4397237503
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 4484206079
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 4572913919
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 4663395839
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 4755687423
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 4849825023
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 4945844991
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 5043785471
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 5143685119
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 5245582335
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 5349517823
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 5455535871
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 5563670271
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 5673967359
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 5786470399
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 5901223167
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 6018271231
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 6137660415
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 6259437055
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 6383649535
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 6510346239
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 6639576575
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 6771391999
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 6905843455
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 7042983935
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 7182867199
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 7325548031
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 7471082495
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 7619527935
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 7770941951
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 7925384447
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 8082915839
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 8243601663
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 8407497471
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 8574670847
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 8745187583
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 8919115007
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 9096521215
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 9277475071
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 9462047743
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 9650312447
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 9842342143
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 10038212607
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 2
+  value_type_size: 8
+  is_segmented: true
+  scratch_size_recordings {
+    num_items: 31
+    scratch_space_bytes: 767
+  }
+  scratch_size_recordings {
+    num_items: 61
+    scratch_space_bytes: 1023
+  }
+  scratch_size_recordings {
+    num_items: 91
+    scratch_space_bytes: 1279
+  }
+  scratch_size_recordings {
+    num_items: 121
+    scratch_space_bytes: 1535
+  }
+  scratch_size_recordings {
+    num_items: 151
+    scratch_space_bytes: 2047
+  }
+  scratch_size_recordings {
+    num_items: 191
+    scratch_space_bytes: 2303
+  }
+  scratch_size_recordings {
+    num_items: 221
+    scratch_space_bytes: 2559
+  }
+  scratch_size_recordings {
+    num_items: 251
+    scratch_space_bytes: 2815
+  }
+  scratch_size_recordings {
+    num_items: 281
+    scratch_space_bytes: 3327
+  }
+  scratch_size_recordings {
+    num_items: 311
+    scratch_space_bytes: 3583
+  }
+  scratch_size_recordings {
+    num_items: 351
+    scratch_space_bytes: 3839
+  }
+  scratch_size_recordings {
+    num_items: 381
+    scratch_space_bytes: 4095
+  }
+  scratch_size_recordings {
+    num_items: 411
+    scratch_space_bytes: 4607
+  }
+  scratch_size_recordings {
+    num_items: 441
+    scratch_space_bytes: 4863
+  }
+  scratch_size_recordings {
+    num_items: 471
+    scratch_space_bytes: 5119
+  }
+  scratch_size_recordings {
+    num_items: 511
+    scratch_space_bytes: 5375
+  }
+  scratch_size_recordings {
+    num_items: 541
+    scratch_space_bytes: 5887
+  }
+  scratch_size_recordings {
+    num_items: 573
+    scratch_space_bytes: 6143
+  }
+  scratch_size_recordings {
+    num_items: 606
+    scratch_space_bytes: 6399
+  }
+  scratch_size_recordings {
+    num_items: 630
+    scratch_space_bytes: 6655
+  }
+  scratch_size_recordings {
+    num_items: 667
+    scratch_space_bytes: 7167
+  }
+  scratch_size_recordings {
+    num_items: 693
+    scratch_space_bytes: 7423
+  }
+  scratch_size_recordings {
+    num_items: 734
+    scratch_space_bytes: 7679
+  }
+  scratch_size_recordings {
+    num_items: 762
+    scratch_space_bytes: 7935
+  }
+  scratch_size_recordings {
+    num_items: 792
+    scratch_space_bytes: 8447
+  }
+  scratch_size_recordings {
+    num_items: 823
+    scratch_space_bytes: 8703
+  }
+  scratch_size_recordings {
+    num_items: 855
+    scratch_space_bytes: 8959
+  }
+  scratch_size_recordings {
+    num_items: 889
+    scratch_space_bytes: 9215
+  }
+  scratch_size_recordings {
+    num_items: 924
+    scratch_space_bytes: 9727
+  }
+  scratch_size_recordings {
+    num_items: 960
+    scratch_space_bytes: 9983
+  }
+  scratch_size_recordings {
+    num_items: 979
+    scratch_space_bytes: 10239
+  }
+  scratch_size_recordings {
+    num_items: 1017
+    scratch_space_bytes: 10495
+  }
+  scratch_size_recordings {
+    num_items: 1037
+    scratch_space_bytes: 11007
+  }
+  scratch_size_recordings {
+    num_items: 1078
+    scratch_space_bytes: 11263
+  }
+  scratch_size_recordings {
+    num_items: 1120
+    scratch_space_bytes: 11519
+  }
+  scratch_size_recordings {
+    num_items: 1142
+    scratch_space_bytes: 11775
+  }
+  scratch_size_recordings {
+    num_items: 1164
+    scratch_space_bytes: 12287
+  }
+  scratch_size_recordings {
+    num_items: 1210
+    scratch_space_bytes: 12543
+  }
+  scratch_size_recordings {
+    num_items: 1234
+    scratch_space_bytes: 12799
+  }
+  scratch_size_recordings {
+    num_items: 1258
+    scratch_space_bytes: 13055
+  }
+  scratch_size_recordings {
+    num_items: 1308
+    scratch_space_bytes: 13567
+  }
+  scratch_size_recordings {
+    num_items: 1334
+    scratch_space_bytes: 13823
+  }
+  scratch_size_recordings {
+    num_items: 1360
+    scratch_space_bytes: 14079
+  }
+  scratch_size_recordings {
+    num_items: 1387
+    scratch_space_bytes: 14335
+  }
+  scratch_size_recordings {
+    num_items: 1414
+    scratch_space_bytes: 14847
+  }
+  scratch_size_recordings {
+    num_items: 1470
+    scratch_space_bytes: 15103
+  }
+  scratch_size_recordings {
+    num_items: 1499
+    scratch_space_bytes: 15359
+  }
+  scratch_size_recordings {
+    num_items: 1528
+    scratch_space_bytes: 15615
+  }
+  scratch_size_recordings {
+    num_items: 1558
+    scratch_space_bytes: 16127
+  }
+  scratch_size_recordings {
+    num_items: 1589
+    scratch_space_bytes: 16383
+  }
+  scratch_size_recordings {
+    num_items: 1620
+    scratch_space_bytes: 16639
+  }
+  scratch_size_recordings {
+    num_items: 1652
+    scratch_space_bytes: 16895
+  }
+  scratch_size_recordings {
+    num_items: 1685
+    scratch_space_bytes: 17407
+  }
+  scratch_size_recordings {
+    num_items: 1718
+    scratch_space_bytes: 17663
+  }
+  scratch_size_recordings {
+    num_items: 1752
+    scratch_space_bytes: 17919
+  }
+  scratch_size_recordings {
+    num_items: 1787
+    scratch_space_bytes: 18175
+  }
+  scratch_size_recordings {
+    num_items: 1822
+    scratch_space_bytes: 18687
+  }
+  scratch_size_recordings {
+    num_items: 1858
+    scratch_space_bytes: 19199
+  }
+  scratch_size_recordings {
+    num_items: 1895
+    scratch_space_bytes: 19455
+  }
+  scratch_size_recordings {
+    num_items: 1932
+    scratch_space_bytes: 19967
+  }
+  scratch_size_recordings {
+    num_items: 1970
+    scratch_space_bytes: 20223
+  }
+  scratch_size_recordings {
+    num_items: 2009
+    scratch_space_bytes: 20479
+  }
+  scratch_size_recordings {
+    num_items: 2049
+    scratch_space_bytes: 21247
+  }
+  scratch_size_recordings {
+    num_items: 2089
+    scratch_space_bytes: 21503
+  }
+  scratch_size_recordings {
+    num_items: 2130
+    scratch_space_bytes: 21759
+  }
+  scratch_size_recordings {
+    num_items: 2172
+    scratch_space_bytes: 22015
+  }
+  scratch_size_recordings {
+    num_items: 2215
+    scratch_space_bytes: 22783
+  }
+  scratch_size_recordings {
+    num_items: 2259
+    scratch_space_bytes: 23039
+  }
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 23295
+  }
+  scratch_size_recordings {
+    num_items: 2350
+    scratch_space_bytes: 24063
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 24319
+  }
+  scratch_size_recordings {
+    num_items: 2444
+    scratch_space_bytes: 25087
+  }
+  scratch_size_recordings {
+    num_items: 2492
+    scratch_space_bytes: 25343
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 25855
+  }
+  scratch_size_recordings {
+    num_items: 2591
+    scratch_space_bytes: 26367
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 26879
+  }
+  scratch_size_recordings {
+    num_items: 2694
+    scratch_space_bytes: 27647
+  }
+  scratch_size_recordings {
+    num_items: 2747
+    scratch_space_bytes: 27903
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 28415
+  }
+  scratch_size_recordings {
+    num_items: 2857
+    scratch_space_bytes: 29183
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 29695
+  }
+  scratch_size_recordings {
+    num_items: 2972
+    scratch_space_bytes: 30207
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 30719
+  }
+  scratch_size_recordings {
+    num_items: 3091
+    scratch_space_bytes: 31487
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 31999
+  }
+  scratch_size_recordings {
+    num_items: 3215
+    scratch_space_bytes: 32767
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 33279
+  }
+  scratch_size_recordings {
+    num_items: 3344
+    scratch_space_bytes: 34047
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 34559
+  }
+  scratch_size_recordings {
+    num_items: 3478
+    scratch_space_bytes: 35327
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 35839
+  }
+  scratch_size_recordings {
+    num_items: 3617
+    scratch_space_bytes: 36863
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 37375
+  }
+  scratch_size_recordings {
+    num_items: 3762
+    scratch_space_bytes: 38143
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 38655
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 39679
+  }
+  scratch_size_recordings {
+    num_items: 3991
+    scratch_space_bytes: 40447
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 41215
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 41983
+  }
+  scratch_size_recordings {
+    num_items: 4234
+    scratch_space_bytes: 43007
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 43519
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 44543
+  }
+  scratch_size_recordings {
+    num_items: 4492
+    scratch_space_bytes: 45567
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 46335
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 47103
+  }
+  scratch_size_recordings {
+    num_items: 4765
+    scratch_space_bytes: 48127
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 48895
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 49919
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 50943
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 52223
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 53247
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 54015
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 55039
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 56319
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 57343
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 58623
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 59903
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 60927
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 62207
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 63487
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 64767
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 66047
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 67327
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 68607
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 69887
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 71167
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 72703
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 73983
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 75519
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 77055
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 78335
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 80127
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 81663
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 83199
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 84735
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 86783
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 88319
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 90367
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 91903
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 93695
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 95743
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 97535
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 99583
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 101375
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 103423
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 105727
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 107519
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 109823
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 112127
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 113919
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 116479
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 118783
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 121087
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 123647
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 125695
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 128255
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 130815
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 133375
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 136447
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 139007
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 141823
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 144383
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 147199
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 150527
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 153343
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 156415
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 159487
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 162559
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 165887
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 169215
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 172543
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 176127
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 179455
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 183039
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 186623
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 190463
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 194303
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 198143
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 201983
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 206079
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 210175
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 214527
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 218623
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 222975
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 227583
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 231935
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 236543
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 241407
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 246015
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 250879
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 255999
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 261119
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 266239
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 271615
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 277247
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 282623
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 288255
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 294143
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 299775
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 305919
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 312063
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 318207
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 324607
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 331007
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 337663
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 344319
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 350975
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 358143
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 365567
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 372735
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 380159
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 387583
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 395519
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 403199
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 411135
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 419583
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 427775
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 436479
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 445183
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 454143
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 463103
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 472319
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 482047
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 491519
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 501247
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 511487
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 521727
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 531967
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 542463
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 553215
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 564479
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 575743
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 587263
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 599039
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 610815
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 623103
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 635647
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 648447
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 661247
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 674303
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 688127
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 701695
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 715519
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 729855
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 744703
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 759295
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 774399
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 790015
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 805887
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 822015
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 838399
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 855039
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 872447
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 889599
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 907519
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 925439
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 944127
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 962815
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 982015
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 1001983
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 1021695
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 1042175
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 1063167
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 1084159
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 1105919
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 1127935
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 1150719
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 1173759
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 1197055
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 1221119
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 1245439
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 1270527
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 1295615
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 1321727
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 1348095
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 1374975
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 1402623
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 1430527
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 1459199
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 1488383
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 1518079
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 1548543
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 1579263
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 1611007
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 1643263
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 1675775
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 1709567
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 1743615
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 1778687
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 1814015
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 1850367
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 1887487
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 1925119
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 1963519
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 2002687
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 2042879
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 2083583
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 2125055
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 2167807
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 2211327
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 2255359
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 2300415
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 2346239
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 2393343
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 2440959
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 2489855
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 2539775
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 2590463
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 2642175
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 2695167
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 2749183
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 2803967
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 2860287
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 2917119
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 2975743
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 3035135
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 3095807
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 3157759
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 3220735
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 3285247
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 3351039
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 3417855
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 3486207
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 3555839
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 3627007
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 3699455
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 3773439
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 3848959
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 3926015
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 4004607
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 4084479
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 4166399
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 4249599
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 4334335
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 4421119
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 4509695
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 4599807
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 4691967
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 4785663
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 4881407
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 4978943
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 5078527
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 5180159
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 5283583
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 5389567
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 5497087
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 5607167
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 5719039
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 5833727
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 5950207
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 6069247
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 6190335
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 6314239
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 6440703
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 6569215
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 6700799
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 6834943
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 6971647
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 7110655
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 7253247
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 7398143
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 7545855
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 7696895
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 7851007
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 8007935
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 8167935
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 8331263
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 8497919
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 8667903
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 8841215
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 9017855
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 9198335
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 9382399
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 9570047
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 9761279
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 9956607
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 10155775
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 10358783
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 10566143
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 10777343
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 10992639
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 11212543
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 11436799
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 11665663
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 11898879
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 12136959
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 12379647
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 12627199
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 12879615
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 13137407
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 13400063
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 13667839
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 13941503
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 14220287
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 14504703
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 14794495
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 15090687
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 15392255
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 15700223
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 16014079
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 16334335
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 16661247
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 16994303
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 17334015
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 17680895
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 18034687
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 18395135
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 18763007
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 19138303
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 19521023
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 19911423
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 20309503
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 20715775
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 21129983
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 21552639
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 21983743
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 22423295
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 22871807
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 23329279
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 23795967
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 24271615
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 24757247
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 25252095
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 25757183
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 26272255
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 26797823
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 27333887
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 27880447
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 28438015
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 29006847
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 29586943
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 30178559
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 30782207
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 31397887
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 32025599
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 32666367
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 33319679
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 33986047
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 34665727
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 35358975
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 36066047
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 36787455
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 37523199
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 38273535
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 39038975
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 39819775
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 40616447
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 41428479
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 42256895
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 43102207
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 43964159
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 44843519
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 45740543
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 46654975
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 47588095
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 48539903
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 49510655
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 50500863
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 51511039
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 52541183
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 53592063
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 54663679
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 55757055
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 56872447
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 58009599
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 59170047
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 60353279
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 61560319
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 62791423
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 64047359
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 65328127
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 66634495
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 67967487
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 69326847
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 70713343
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 72127487
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 73570047
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 75041279
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 76542207
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 78073087
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 79634687
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 81227263
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 82851839
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 84508927
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 86198783
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 87922943
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 89681407
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 91474943
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 93304319
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 95170559
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 97073919
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 99015423
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 100995583
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 103015423
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 105075967
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 107177215
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 109320703
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 111507199
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 113737215
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 116012287
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 118332159
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 120698879
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 123112959
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 125575167
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 128086783
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 130648319
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 133261055
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 135926527
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 138645247
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 141417983
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 144246015
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 147131135
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 150073599
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 153075199
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 156136703
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 159259647
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 162444543
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 165693439
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 169007359
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 172387583
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 175835135
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 179352063
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 182938879
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 186597887
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 190329599
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 194136319
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 198018815
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 201979135
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 206018815
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 210139135
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 214342143
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 218628863
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 223001343
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 227461375
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 232010495
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 236651007
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 241383679
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 246211583
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 251135743
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 256158463
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 261281535
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 266507263
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 271837183
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 277273855
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 282819583
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 288475903
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 294245119
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 300130303
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 306132735
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 312255487
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 318500607
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 324870399
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 331367935
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 337995263
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 344755199
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 351650303
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 358683135
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 365857023
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 373174015
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 380637439
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 388250367
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 396015103
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 403935487
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 412014079
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 420254463
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 428659455
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 437232639
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 445977343
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 454896639
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 463994623
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 473274623
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 482740223
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 492395007
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 502242559
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 512287743
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 522533375
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 532984063
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 543643647
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 554516479
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 565606655
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 576919039
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 588457215
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 600226303
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 612231167
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 624475647
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 636965119
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 649704447
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 662698239
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 675952383
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 689471487
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 703260927
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 717326079
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 731672319
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 746306047
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 761232127
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 776456703
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 791985663
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 807825407
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 823981823
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 840461567
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 857270783
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 874416127
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 891904255
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 909742335
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 927937279
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 946495999
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 965425919
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 984734463
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 1004429055
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 1024517887
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 1045008127
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 1065908223
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 1087226623
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 1108971007
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 1131150335
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 1153773567
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 1176848639
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 1200385535
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 1224393215
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 1248881407
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 1273858815
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 1299335935
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 1325323007
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 1351829247
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 1378865919
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 1406443263
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 1434571775
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 1463263487
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 1492528639
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 1522379263
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 1552826623
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 1583883263
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 1615560959
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 1647871999
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 1680829439
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 1714446079
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 1748734975
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 1783709695
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 1819384063
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 1855771647
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 1892887039
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 1930744575
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 1969359615
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 2008747007
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 2048921599
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 2089900287
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 2131697919
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 2174332159
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 2217818623
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 2262174975
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 2307418623
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 2353566719
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 2400638207
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 2448651007
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 2497624063
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 2547576319
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 2598527999
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 2650498559
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 2703508479
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 2757578495
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 2812730367
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 2868984575
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 2926364415
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 2984891647
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 3044589567
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 3105481215
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 3167590655
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 3230942719
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 3295561727
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 3361472767
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 3428702207
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 3497276159
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 3567221759
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 3638566143
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 3711337215
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 3785564159
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 3861275647
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 3938500863
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 4017271039
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 4097616383
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 4179568639
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 4263160063
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 4348423167
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 4435391743
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 4524099583
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 4614581503
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 4706873087
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 4801010687
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 4897030655
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 4994971135
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 5094870783
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 5196767999
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 5300703487
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 5406717439
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 5514851839
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 5625148927
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 5737651967
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 5852404735
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 5969452799
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 6088841983
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 6210618623
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 6334831103
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 6461527807
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 6590758143
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 6722573567
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 6857025023
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 6994165503
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 7134048767
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 7276729599
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 7422264063
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 7570709503
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 7722123519
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 7876566015
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 8034097407
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 8194779135
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 8358674943
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 8525848319
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 8696365055
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 8870292479
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 9047698687
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 9228652543
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 9413225215
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 9601489919
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 9793519615
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 9989390079
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 4
+  value_type_size: 2
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 1
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 39679
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 40191
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 40959
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 41727
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 41983
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 42751
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 43519
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 44031
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 44799
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 45567
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 46335
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 47103
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 47871
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 48639
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 50431
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 51199
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 51967
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 52735
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 53503
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 54271
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 55039
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 56063
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 56831
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 57599
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 58879
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 59647
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 60671
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 61439
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 62719
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 63743
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 65023
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 65791
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 66815
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 68095
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 69119
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 70399
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 71423
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 72703
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 74239
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 75263
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 76543
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 78079
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 79103
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 80639
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 81919
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 83455
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 84991
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 86015
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 87551
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 89087
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 91647
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 93695
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 95231
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 96767
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 98303
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 100095
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 102143
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 103679
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 105471
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 107519
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 109311
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 111359
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 113151
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 115199
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 117503
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 119295
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 121599
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 123647
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 125951
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 128255
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 131583
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 133887
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 136447
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 138751
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 141567
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 143871
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 146431
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 149247
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 151807
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 154623
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 157695
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 160255
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 163327
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 166399
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 169471
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 173567
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 176639
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 180223
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 183295
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 186623
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 190207
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 193535
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 197375
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 200959
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 204799
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 208639
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 213503
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 217343
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 221439
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 225279
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 229631
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 234239
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 238335
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 242943
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 247295
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 253183
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 257791
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 262399
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 267519
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 272383
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 277759
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 282879
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 288255
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 294655
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 300287
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 306175
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 311807
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 317695
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 323839
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 329983
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 337151
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 343295
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 349695
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 356607
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 363263
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 370175
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 378367
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 385279
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 392703
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 400383
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 408063
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 416767
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 424447
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 432895
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 440831
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 449279
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 458751
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 467711
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 476415
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 485631
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 495871
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 505599
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 515071
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 525055
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 536063
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 546559
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 556799
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 567551
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 579327
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 590591
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 601599
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 614143
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 626175
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 637951
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 650239
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 664063
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 676607
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 689663
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 703743
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 717567
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 731391
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 746239
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 760831
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 776447
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 791551
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 806399
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 823295
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 838911
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 856063
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 872703
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 889599
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 907775
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 925183
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 944127
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 962303
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 981759
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 1000959
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 1021183
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 1040639
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 1062143
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 1082367
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 1104639
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 1125631
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 1148671
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 1170943
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 1194495
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 1218559
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 1242111
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 1267199
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 1291519
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 1317375
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 1344255
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 1370367
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 1397759
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 1425663
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 1453311
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 1482495
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 1512191
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 1542399
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 1572351
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 1603839
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 1635839
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 1668863
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 1702143
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 1735167
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 1769983
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 1805055
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 1841151
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 1877759
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 1915391
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 1953535
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 1992191
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 2032127
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 2072575
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 2113535
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 2155775
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 2198527
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 2242303
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 2286591
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 2332159
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 2378495
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 2426623
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 2475007
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 2523903
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 2574079
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 2626047
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 2677759
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 2731007
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 2786047
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 2841343
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 2897663
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 2955775
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 3014399
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 3074815
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 3135743
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 3198719
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 3261695
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 3327487
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 3394047
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 3461119
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 3530239
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 3601151
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 3673087
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 3745535
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 3820031
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 3896575
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 3974399
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 4053503
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 4134655
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 4217087
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 4301311
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 4386559
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 4475391
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 4564223
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 4654847
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 4747519
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 4843263
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 4939263
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 5038335
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 5138431
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 5241599
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 5345535
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 5452543
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 5561599
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 5671935
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 5785599
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 5901311
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 6019071
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 6139391
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 6261759
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 6386687
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 6514175
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 6644991
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 6777343
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 6912255
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 7051007
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 7191295
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 7335423
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 7481343
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 7631103
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 7783679
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 7939071
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 8098047
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 8259583
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 8424447
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 8592639
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 8764159
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 8939775
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 9117695
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 9300735
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 9485567
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 9675519
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 9869055
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 10066175
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 10267647
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 10472447
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 10681343
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 10894591
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 11113215
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 11334399
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 11561471
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 11792639
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 12028415
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 12268799
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 12513791
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 12763647
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 13018367
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 13279231
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 13543935
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 13814783
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 14091263
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 14372863
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 14660095
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 14952447
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 15252223
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 15556095
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 15867391
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 16184575
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 16508159
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 16838143
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 17175295
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 17517823
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 17868543
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 18225663
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 18589951
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 18961407
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 19340031
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 19726847
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 20121599
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 20523775
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 20933887
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 21352959
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 21778943
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 22214655
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 22658559
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 23111423
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 23573759
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 24045311
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 24525055
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 25016575
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 25516031
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 26026239
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 26546943
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 27077375
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 27618303
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 28170751
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 28733695
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 29307903
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 29894399
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 30492159
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 31101695
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 31724287
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 32357631
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 33004543
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 33665023
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 34338303
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 35025151
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 35724799
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 36439551
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 37167871
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 37911551
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 38668543
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 39442687
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 40230655
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 41034751
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 41856255
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 42692351
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 43546623
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 44416767
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 45304831
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 46211071
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 47135231
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 48077311
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 49038847
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 50019839
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 51020287
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 52040191
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 53080575
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 54141695
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 55224575
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 56329215
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 57455615
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 58604031
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 59776767
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 60971263
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 62191103
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 63434751
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 64702975
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 65996287
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 67316223
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 68663039
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 70035711
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 71436799
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 72864511
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 74321407
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 75807743
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 77324031
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 78870783
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 80447743
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 82056191
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 83697663
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 85370879
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 87078399
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 88819199
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 90595327
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 92407039
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 94255103
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 96140031
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 98063103
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 100024063
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 102024703
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 104064255
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 106145535
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 108268287
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 110433791
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 112642303
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 114895103
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 117192447
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 119535871
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 121926655
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 124365311
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 126851583
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 129389311
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 131976703
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 134616319
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 137308415
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 140054015
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 142855167
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 145712383
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 148625663
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 151598079
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 154630399
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 157721855
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 160876543
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 164093951
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 167375871
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 170722815
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 174137343
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 177620223
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 181171967
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 184795135
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 188490751
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 192260863
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 196105215
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 200027391
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 204028159
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 208108031
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 212270591
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 216515071
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 220845311
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 225262335
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 229767167
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 234362111
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 239049727
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 243830271
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 248706559
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 253680639
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 258753791
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 263928831
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 269207039
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 274590975
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 280082943
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 285683967
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 291397375
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 297225215
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 303169791
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 309232895
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 315417599
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 321725695
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 328159999
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 334722815
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 341417471
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 348245503
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 355210239
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 362314239
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 369560063
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 376951551
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 384490751
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 392179967
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 400023295
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 408023295
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 416183551
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 424507135
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 432996863
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 441657087
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 450490111
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 459499775
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 468689919
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 478063103
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 487624191
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 497375999
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 507323647
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 517469695
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 527819263
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 538374911
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 549142015
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 560124415
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 571327231
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 582754303
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 594408191
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 606296319
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 618422271
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 630791167
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 643406847
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 656274431
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 669399295
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 682787071
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 696442623
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 710372095
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 724578559
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 739070207
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 753851647
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 768928767
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 784306175
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 799992831
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 815992831
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 832312319
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 848958207
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 865936639
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 883255039
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 900920831
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 918938623
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 937316863
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 956063487
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 975183871
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 994688255
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 1014581759
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 1034872319
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 1055569663
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 1076681471
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 1098214143
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 1120178943
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 1142582015
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 1165433855
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 1188742143
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 1212516351
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 1236767487
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 1261501951
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 1286732287
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 1312466175
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 1338715391
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 1365489151
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 1392799231
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 1420655103
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 1449068031
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 1478049279
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 1507609599
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 1537762303
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 1568517119
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 1599886847
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 1631884287
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 1664228863
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 1696670719
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 1729762047
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 1763514367
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 1797942271
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 1833058815
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 1868877567
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 1905412351
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 1942678015
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 1980689407
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 2019460863
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 2059007487
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 2099345151
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 2140489471
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 2182456831
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 2225263359
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 2268925951
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 2313462271
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 2358889215
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 2405224191
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 2452486399
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 2500693503
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 2549864959
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 2600019711
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 2651177727
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 2703358719
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 2756583423
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 2810872575
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 2866247679
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 2922730239
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 2980342015
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 3039106303
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 3099046143
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 3160184575
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 3222545919
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 3286162431
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 3351043071
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 3417221375
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 3484723199
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 3553574655
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 3623803647
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 3695437055
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 3768503039
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 3843030527
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 3919048703
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 3996586751
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 4075676159
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 4156346879
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 4238631167
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 4322561279
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 4408169727
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 4495490303
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 4584557567
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 4675406079
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 4768071423
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 4862590463
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 4959007487
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 5057345023
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 5157649151
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 5259959039
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 5364315647
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 5470759423
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 5579331583
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 5690075135
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 5803034111
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 5918251775
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 6035774207
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 4
+  value_type_size: 2
+  is_segmented: true
+  scratch_size_recordings {
+    num_items: 61
+    scratch_space_bytes: 767
+  }
+  scratch_size_recordings {
+    num_items: 121
+    scratch_space_bytes: 1023
+  }
+  scratch_size_recordings {
+    num_items: 191
+    scratch_space_bytes: 1535
+  }
+  scratch_size_recordings {
+    num_items: 251
+    scratch_space_bytes: 1791
+  }
+  scratch_size_recordings {
+    num_items: 311
+    scratch_space_bytes: 2303
+  }
+  scratch_size_recordings {
+    num_items: 381
+    scratch_space_bytes: 2559
+  }
+  scratch_size_recordings {
+    num_items: 441
+    scratch_space_bytes: 3071
+  }
+  scratch_size_recordings {
+    num_items: 511
+    scratch_space_bytes: 3327
+  }
+  scratch_size_recordings {
+    num_items: 573
+    scratch_space_bytes: 3839
+  }
+  scratch_size_recordings {
+    num_items: 630
+    scratch_space_bytes: 4095
+  }
+  scratch_size_recordings {
+    num_items: 693
+    scratch_space_bytes: 4607
+  }
+  scratch_size_recordings {
+    num_items: 762
+    scratch_space_bytes: 4863
+  }
+  scratch_size_recordings {
+    num_items: 823
+    scratch_space_bytes: 5375
+  }
+  scratch_size_recordings {
+    num_items: 889
+    scratch_space_bytes: 5631
+  }
+  scratch_size_recordings {
+    num_items: 960
+    scratch_space_bytes: 6143
+  }
+  scratch_size_recordings {
+    num_items: 1017
+    scratch_space_bytes: 6399
+  }
+  scratch_size_recordings {
+    num_items: 1078
+    scratch_space_bytes: 6911
+  }
+  scratch_size_recordings {
+    num_items: 1142
+    scratch_space_bytes: 7167
+  }
+  scratch_size_recordings {
+    num_items: 1210
+    scratch_space_bytes: 7679
+  }
+  scratch_size_recordings {
+    num_items: 1258
+    scratch_space_bytes: 7935
+  }
+  scratch_size_recordings {
+    num_items: 1334
+    scratch_space_bytes: 8447
+  }
+  scratch_size_recordings {
+    num_items: 1387
+    scratch_space_bytes: 8703
+  }
+  scratch_size_recordings {
+    num_items: 1470
+    scratch_space_bytes: 9215
+  }
+  scratch_size_recordings {
+    num_items: 1528
+    scratch_space_bytes: 9471
+  }
+  scratch_size_recordings {
+    num_items: 1589
+    scratch_space_bytes: 9983
+  }
+  scratch_size_recordings {
+    num_items: 1652
+    scratch_space_bytes: 10239
+  }
+  scratch_size_recordings {
+    num_items: 1718
+    scratch_space_bytes: 10751
+  }
+  scratch_size_recordings {
+    num_items: 1787
+    scratch_space_bytes: 11007
+  }
+  scratch_size_recordings {
+    num_items: 1822
+    scratch_space_bytes: 11519
+  }
+  scratch_size_recordings {
+    num_items: 1895
+    scratch_space_bytes: 11775
+  }
+  scratch_size_recordings {
+    num_items: 1970
+    scratch_space_bytes: 12287
+  }
+  scratch_size_recordings {
+    num_items: 2009
+    scratch_space_bytes: 12543
+  }
+  scratch_size_recordings {
+    num_items: 2089
+    scratch_space_bytes: 13055
+  }
+  scratch_size_recordings {
+    num_items: 2172
+    scratch_space_bytes: 13311
+  }
+  scratch_size_recordings {
+    num_items: 2215
+    scratch_space_bytes: 13823
+  }
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 14079
+  }
+  scratch_size_recordings {
+    num_items: 2350
+    scratch_space_bytes: 14591
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 14847
+  }
+  scratch_size_recordings {
+    num_items: 2492
+    scratch_space_bytes: 15359
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 15615
+  }
+  scratch_size_recordings {
+    num_items: 2591
+    scratch_space_bytes: 16127
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 16383
+  }
+  scratch_size_recordings {
+    num_items: 2747
+    scratch_space_bytes: 16895
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 17151
+  }
+  scratch_size_recordings {
+    num_items: 2857
+    scratch_space_bytes: 17663
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 17919
+  }
+  scratch_size_recordings {
+    num_items: 2972
+    scratch_space_bytes: 18431
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 18687
+  }
+  scratch_size_recordings {
+    num_items: 3091
+    scratch_space_bytes: 19199
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 19455
+  }
+  scratch_size_recordings {
+    num_items: 3215
+    scratch_space_bytes: 19967
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 20223
+  }
+  scratch_size_recordings {
+    num_items: 3344
+    scratch_space_bytes: 20735
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 20991
+  }
+  scratch_size_recordings {
+    num_items: 3478
+    scratch_space_bytes: 21503
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 21759
+  }
+  scratch_size_recordings {
+    num_items: 3617
+    scratch_space_bytes: 22271
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 22527
+  }
+  scratch_size_recordings {
+    num_items: 3762
+    scratch_space_bytes: 23039
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 23295
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 24063
+  }
+  scratch_size_recordings {
+    num_items: 3991
+    scratch_space_bytes: 24575
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 24831
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 25343
+  }
+  scratch_size_recordings {
+    num_items: 4234
+    scratch_space_bytes: 26111
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 26367
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 26879
+  }
+  scratch_size_recordings {
+    num_items: 4492
+    scratch_space_bytes: 27647
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 27903
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 28415
+  }
+  scratch_size_recordings {
+    num_items: 4765
+    scratch_space_bytes: 29183
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 29439
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 30207
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 30719
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 31487
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 32255
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 32511
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 33279
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 34047
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 34559
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 35327
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 36095
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 36863
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 37631
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 38399
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 39167
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 39935
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 40703
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 41471
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 42239
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 43007
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 43775
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 44543
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 45567
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 46335
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 47103
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 48383
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 49151
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 50175
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 50943
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 52223
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 53247
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 54527
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 55295
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 56319
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 57599
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 58623
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 59903
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 60927
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 62207
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 63743
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 64767
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 66047
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 67583
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 68607
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 70143
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 71423
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 72959
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 74495
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 75519
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 77055
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 78591
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 80127
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 82175
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 83711
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 85247
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 86783
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 88575
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 90623
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 92159
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 93951
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 95999
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 97791
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 99839
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 101631
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 103679
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 105983
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 107775
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 110079
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 112127
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 114431
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 116735
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 119039
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 121343
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 123903
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 126207
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 129023
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 131327
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 133887
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 136703
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 139263
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 142079
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 145151
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 147711
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 150783
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 153855
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 156927
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 159999
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 163071
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 166655
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 169727
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 173055
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 176639
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 179967
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 183807
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 187391
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 191231
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 195071
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 198911
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 202751
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 206847
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 210687
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 215039
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 219647
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 223743
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 228351
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 232703
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 237567
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 242175
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 246783
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 251903
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 256767
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 262143
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 267263
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 272639
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 278015
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 283647
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 289535
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 295167
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 301055
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 307199
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 313343
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 319487
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 325631
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 332031
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 338943
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 345599
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 352511
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 359679
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 366591
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 374015
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 381695
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 389375
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 397055
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 404735
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 413183
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 421119
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 429567
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 438015
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 446975
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 455679
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 464895
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 474111
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 483839
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 493311
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 503295
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 513279
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 523775
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 534015
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 544767
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 555519
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 566783
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 577791
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 589311
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 601343
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 613119
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 625407
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 638207
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 650751
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 663807
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 676863
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 690687
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 704511
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 718335
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 732927
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 747519
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 762623
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 777471
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 793343
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 808959
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 825087
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 841727
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 858623
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 875775
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 893183
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 911103
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 929279
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 947711
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 966911
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 986111
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 1005567
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 1026047
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 1046271
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 1067519
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 1088511
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 1110527
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 1132799
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 1155327
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 1178367
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 1201919
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 1225983
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 1250303
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 1275135
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 1300991
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 1327103
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 1353471
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 1380351
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 1407999
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 1436159
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 1464831
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 1494015
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 1523967
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 1554431
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 1585407
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 1617407
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 1649663
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 1682687
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 1716479
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 1750527
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 1785599
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 1821183
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 1857791
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 1894911
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 1932543
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 1971455
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 2010879
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 2050815
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 2092031
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 2133759
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 2176511
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 2219775
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 2264319
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 2309631
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 2355711
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 2403071
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 2450943
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 2500095
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 2550015
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 2600703
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 2652927
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 2705919
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 2760191
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 2815487
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 2871551
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 2929151
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 2987519
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 3047423
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 3108351
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 3170303
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 3234047
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 3298559
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 3364607
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 3431679
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 3500543
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 3570431
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 3641855
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 3714303
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 3788799
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 3864575
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 3941631
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 4020735
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 4101119
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 4183295
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 4266495
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 4352255
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 4439039
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 4527615
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 4618239
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 4710911
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 4804863
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 4900863
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 4998911
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 5099007
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 5200895
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 5304831
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 5410815
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 5519103
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 5629695
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 5742335
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 5857023
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 5974271
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 6093567
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 6215423
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 6339839
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 6466559
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 6595839
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 6727679
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 6862335
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 6999551
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 7139583
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 7282431
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 7428095
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 7576575
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 7727871
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 7882751
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 8040191
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 8200959
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 8365055
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 8532479
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 8702975
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 8876799
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 9054719
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 9235455
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 9420287
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 9608703
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 9800703
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 9997055
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 10196735
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 10400511
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 10608639
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 10821119
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 11037183
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 11258111
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 11483135
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 11712767
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 11947007
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 12185855
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 12429567
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 12678143
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 12931839
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 13190399
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 13454079
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 13723391
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 13997823
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 14277887
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 14563071
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 14854655
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 15151359
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 15454463
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 15763455
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 16078847
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 16400639
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 16728575
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 17062911
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 17404415
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 17752319
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 18107391
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 18469631
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 18839039
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 19215615
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 19600127
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 19992063
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 20391935
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 20799743
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 21215487
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 21639935
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 22072575
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 22514175
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 22964223
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 23423487
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 23891967
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 24370175
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 24857343
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 25354239
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 25861631
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 26378751
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 26906367
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 27444479
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 27993087
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 28552959
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 29124095
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 29706495
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 30300671
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 30906879
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 31524863
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 32155391
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 32798463
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 33454335
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 34123775
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 34806015
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 35502335
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 36212223
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 36936447
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 37675007
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 38428671
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 39197183
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 39980799
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 40780799
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 41596415
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 42428159
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 43276799
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 44142335
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 45025023
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 45925631
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 46844159
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 47781119
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 48736511
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 49711359
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 50705663
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 51719423
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 52753919
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 53809151
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 54885119
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 55982847
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 57102591
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 58244607
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 59409407
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 60597503
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 61809407
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 63045887
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 64306431
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 65592575
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 66904575
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 68242431
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 69607679
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 70999551
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 72419583
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 73868031
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 75345407
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 76852223
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 78389247
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 79956735
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 81556223
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 83187455
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 84850943
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 86547711
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 88278783
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 90044415
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 91845375
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 93682175
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 95556095
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 97466879
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 99416319
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 101404671
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 103432703
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 105501183
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 107611391
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 109763583
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 111959039
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 114198015
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 116482047
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 118811391
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 121187583
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 123611391
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 126083583
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 128605439
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 131177471
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 133800959
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 136476927
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 139206399
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 141990911
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 144830463
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 147727103
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 150681599
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 153695231
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 156769023
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 159904511
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 163102463
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 166364415
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 169691903
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 173085695
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 176547327
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 180078335
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 183679743
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 187353599
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 191100671
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 194922495
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 198820863
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 202797311
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 206853375
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 210990335
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 215209983
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 219514367
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 223904511
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 228382719
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 232950527
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 237609215
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 242361599
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 247208703
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 252152831
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 257195775
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 262339839
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 267586559
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 272938239
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 278396927
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 283964927
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 289644287
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 295437311
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 301345791
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 307372799
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 313520127
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 319790591
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 326186495
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 332710143
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 339364095
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 346151679
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 353074431
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 360135935
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 367339007
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 374685695
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 382179327
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 389822975
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 397619199
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 405571583
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 413683199
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 421956863
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 430395903
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 439003647
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 447783935
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 456739583
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 465874175
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 475191551
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 484695551
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 494389247
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 504277247
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 514362623
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 524649983
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 535142655
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 545845503
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 556762623
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 567897855
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 579255807
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 590840831
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 602657535
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 614711039
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 627005183
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 639545087
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 652336127
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 665382911
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 678690303
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 692264447
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 706109439
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 720231423
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 734636031
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 749329151
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 764315391
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 779601663
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 795194111
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 811097855
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 827319807
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 843866111
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 860743167
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 877958399
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 895517439
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 913427711
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 931696127
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 950330111
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 969336831
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 988723455
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 1008497919
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 1028667903
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 1049241087
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 1070225919
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 1091630591
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 1113463295
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 1135732479
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 1158446847
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 1181615871
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 1205248511
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 1229353215
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 1253940479
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 1279019007
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 1304599551
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 1330691327
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 1357305087
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 1384451327
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 1412140287
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 1440383231
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 1469190911
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 1498574591
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 1528546047
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 1559117055
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 1590299391
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 1622105343
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 1654547199
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 1687638527
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 1721390847
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 1755818751
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 1790935295
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 1826754047
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 1863288831
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 1900554495
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 1938565887
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 1977337343
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 2016883967
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 2057221631
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 2098365951
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 2140333311
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 2183139839
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 2226802431
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 2271338751
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 2316765695
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 2363100671
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 2410362879
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 2458569983
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 2507741439
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 2557896191
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 2609054207
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 2661235199
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 2714459903
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 2768749055
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 2824124159
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 2880606719
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 2938218495
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 2996982783
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 3056922623
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 3118061055
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 3180422399
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 3244030719
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 3308911359
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 3375089663
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 3442591487
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 3511442943
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 3581671935
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 3653305343
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 3726371327
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 3800898815
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 3876916991
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 3954455039
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 4033544447
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 4114215167
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 4196499455
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 4280429567
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 4366038015
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 4453358591
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 4542425855
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 4633274367
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 4725939711
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 4820458751
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 4916867583
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 5015205119
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 5115509247
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 5217819135
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 5322175743
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 5428619519
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 5537191679
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 5647935231
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 5760894207
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 5876111871
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 5993634303
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 4
+  value_type_size: 4
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 1
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 49663
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 50175
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 51199
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 52223
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 52735
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 53759
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 54783
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 55295
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 56319
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 57343
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 58367
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 59391
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 60415
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 61439
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 63487
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 64511
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 65535
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 66559
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 67583
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 68607
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 69631
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 71167
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 72191
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 73215
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 74751
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 75775
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 77311
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 78335
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 79871
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 81407
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 82943
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 83967
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 85503
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 87039
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 88575
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 90111
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 91647
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 93183
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 95231
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 96767
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 98303
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 100351
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 101887
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 103935
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 105471
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 107519
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 109567
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 111103
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 113151
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 115199
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 118271
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 120831
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 122879
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 124927
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 126975
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 129535
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 132095
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 134143
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 136703
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 139263
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 141823
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 144383
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 146943
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 149503
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 152575
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 155135
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 158207
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 160767
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 163839
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 166911
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 171007
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 174079
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 177663
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 180735
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 184319
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 187391
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 190975
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 194559
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 198143
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 201727
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 205823
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 209407
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 213503
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 217599
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 221695
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 226815
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 230911
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 235519
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 239615
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 244223
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 248831
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 253439
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 258559
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 263167
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 268287
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 273407
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 279551
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 284671
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 290303
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 295423
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 301055
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 307199
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 312831
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 318975
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 324607
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 332287
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 338431
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 344575
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 351231
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 357887
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 365055
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 371711
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 378879
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 387071
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 394751
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 402431
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 410111
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 417791
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 425983
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 434175
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 443391
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 451583
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 460287
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 469503
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 478207
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 487423
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 498175
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 507391
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 517119
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 527359
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 537599
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 548863
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 559103
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 570367
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 581119
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 592383
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 604671
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 616447
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 628223
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 640511
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 653823
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 666623
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 679423
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 692735
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 707071
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 720895
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 734719
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 749055
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 764415
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 779263
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 794111
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 810495
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 826367
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 842239
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 858623
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 876543
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 893439
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 910847
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 929279
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 947711
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 966143
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 985599
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 1005055
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 1025535
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 1045503
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 1065471
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 1087487
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 1108479
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 1131007
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 1153023
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 1175551
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 1199615
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 1222655
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 1247743
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 1271807
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 1297407
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 1323007
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 1349631
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 1375743
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 1403903
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 1431039
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 1460223
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 1488383
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 1518591
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 1548287
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 1579519
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 1611263
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 1642495
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 1675775
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 1708031
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 1742335
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 1777663
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 1812479
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 1848831
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 1885695
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 1922559
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 1960959
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 2000383
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 2040319
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 2080255
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 2121727
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 2164223
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 2207743
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 2251775
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 2295807
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 2341887
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 2388479
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 2436095
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 2484735
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 2534399
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 2585087
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 2636287
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 2689023
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 2742783
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 2797055
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 2852863
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 2909695
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 2967551
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 3026431
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 3086847
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 3148287
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 3211775
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 3275775
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 3340799
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 3407359
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 3475967
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 3544575
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 3615231
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 3687935
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 3761151
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 3835903
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 3912703
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 3990527
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 4070399
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 4151295
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 4234751
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 4318207
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 4405247
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 4493311
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 4582399
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 4674047
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 4767743
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 4862975
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 4959231
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 5058047
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 5159423
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 5262335
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 5367295
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 5474815
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 5583871
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 5695487
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 5808639
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 5925887
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 6043647
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 6163967
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 6286847
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 6413311
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 6540799
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 6671871
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 6804479
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 6941183
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 7078911
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 7220735
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 7365119
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 7511551
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 7662079
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 7815167
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 7971327
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 8130559
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 8292863
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 8458239
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 8627199
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 8800255
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 8975871
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 9154559
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 9338367
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 9524223
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 9715199
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 9908735
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 10106879
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 10309119
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 10514943
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 10725375
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 10939391
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 11158015
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 11380735
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 11608063
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 11840511
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 12076543
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 12318719
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 12563967
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 12815359
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 13071871
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 13332991
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 13599743
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 13871103
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 14148095
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 14430719
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 14719999
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 15013375
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 15313919
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 15620095
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 15932415
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 16250879
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 16575487
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 16906751
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 17244159
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 17589759
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 17940479
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 18299391
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 18665471
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 19038719
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 19419135
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 19806719
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 20203519
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 20606463
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 21018623
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 21438975
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 21867519
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 22304767
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 22751231
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 23205375
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 23669759
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 24142847
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 24625663
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 25117695
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 25619455
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 26131967
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 26654719
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 27187711
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 27730943
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 28285951
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 28850687
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 29427711
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 30015999
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 30616063
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 31228415
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 31853055
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 32488959
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 33139711
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 33801727
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 34477567
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 35167231
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 35870207
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 36587007
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 37318655
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 38064639
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 38825471
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 39602175
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 40394239
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 41201663
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 42026495
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 42865663
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 43722751
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 44597759
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 45489663
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 46399487
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 47326719
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 48273407
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 49238527
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 50223615
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 51226623
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 52252159
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 53296127
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 54361599
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 55449599
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 56557567
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 57689087
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 58842111
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 60018687
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 61219327
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 62443519
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 63691775
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 64965631
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 66265087
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 67590655
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 68941823
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 70320127
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 71726079
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 73160703
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 74623999
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 76116479
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 77638143
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 79191551
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 80774143
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 82390015
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 84037631
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 85718015
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 87431679
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 89180159
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 90964479
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 92783103
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 94639103
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 96530943
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 98461183
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 100430335
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 102438911
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 104487935
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 106577407
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 108708351
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 110882815
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 113099775
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 115361791
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 117668351
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 120021503
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 122421759
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 124870143
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 127367167
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 129914879
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 132512767
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 135163391
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 137865727
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 140622847
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 143435263
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 146303999
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 149230079
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 152214527
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 155258367
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 158363135
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 161530367
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 164761087
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 168055295
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 171417087
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 174844927
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 178341887
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 181908479
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 185546239
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 189257215
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 193042431
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 196902399
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 200840191
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 204857343
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 208953343
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 213132799
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 217395199
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 221743103
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 226177535
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 230701055
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 235315199
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 240020991
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 244820991
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 249717247
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 254711807
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 259805183
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 265001471
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 270301695
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 275706879
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 281221631
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 286844927
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 292581887
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 298433535
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 304401919
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 310489599
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 316699647
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 323033087
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 329493503
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 336083455
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 342804479
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 349660671
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 356653567
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 363786239
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 371062271
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 378482687
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 386052095
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 393773055
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 401648639
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 409681407
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 417874943
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 426232319
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 434756607
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 443451391
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 452320767
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 461366783
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 470594047
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 480005631
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 489605119
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 499397631
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 509385727
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 519572991
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 529964031
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 540562943
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 551373823
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 562401279
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 573648895
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 585122303
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 596824575
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 608760831
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 620936191
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 633354239
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 646021119
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 658940927
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 672119807
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 685561855
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 699273215
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 713257983
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 727522815
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 742072831
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 756914687
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 772053503
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 787493375
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 803243007
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 819308031
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 835694591
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 852408319
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 869455871
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 886844415
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 904581119
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 922672639
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 941126655
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 959948287
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 979147263
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 998730239
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 1018704895
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 1039077887
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 1059859967
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 1081057279
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 1102678015
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 1124731391
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 1147225087
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 1170169343
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 1193573375
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 1217444351
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 1241792511
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 1266628607
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 1291960319
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 1317800447
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 1344156159
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 1371038207
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 1398458879
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 1426428415
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 1454956031
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 1484055551
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 1513736191
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 1544011263
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 1574891007
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 1606388223
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 1638516735
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 1671286271
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 1704712191
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 1738805759
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 1773581823
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 1809052671
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 1845234175
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 1882138623
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 1919781375
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 1958176767
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 1997339647
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 2037286911
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 2078032383
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 2119592447
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 2161983999
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 2204930559
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 2248186367
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 2292307967
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 2337311231
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 2383215103
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 2430036991
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 2477795327
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 2526508543
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 2576196095
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 2626877951
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 2678573055
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 2731301887
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 2785085439
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 2839944703
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 2895901183
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 2952976383
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 3011193343
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 3070575103
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 3131144191
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 3192924159
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 3255940607
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 3320216575
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 3385778687
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 3452651519
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 3520862207
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 3590436863
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 3661403135
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 3733788671
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 3807622143
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 3882932223
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 3959748095
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 4038100479
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 4118020095
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 4199538175
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 4282686463
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 4367505919
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 4454013439
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 4542251007
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 4632253439
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 4724055551
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 4817694207
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 4913205247
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 5010626559
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 5109996543
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 5211354111
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 5314738175
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 5420190719
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 5527751679
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 5637464063
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 5749370879
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 5863515647
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 5979942911
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 6098699263
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 6219830783
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 6343384575
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 6469409791
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 6597963263
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 6729079807
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 6862818815
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 6999231999
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 7138374143
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 7280299007
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 7425061887
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 7572720127
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 7723332095
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 7876955647
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 8033652223
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 4
+  value_type_size: 4
+  is_segmented: true
+  scratch_size_recordings {
+    num_items: 61
+    scratch_space_bytes: 767
+  }
+  scratch_size_recordings {
+    num_items: 121
+    scratch_space_bytes: 1279
+  }
+  scratch_size_recordings {
+    num_items: 191
+    scratch_space_bytes: 1791
+  }
+  scratch_size_recordings {
+    num_items: 251
+    scratch_space_bytes: 2303
+  }
+  scratch_size_recordings {
+    num_items: 311
+    scratch_space_bytes: 2815
+  }
+  scratch_size_recordings {
+    num_items: 381
+    scratch_space_bytes: 3327
+  }
+  scratch_size_recordings {
+    num_items: 441
+    scratch_space_bytes: 3839
+  }
+  scratch_size_recordings {
+    num_items: 511
+    scratch_space_bytes: 4351
+  }
+  scratch_size_recordings {
+    num_items: 573
+    scratch_space_bytes: 4863
+  }
+  scratch_size_recordings {
+    num_items: 630
+    scratch_space_bytes: 5375
+  }
+  scratch_size_recordings {
+    num_items: 693
+    scratch_space_bytes: 5887
+  }
+  scratch_size_recordings {
+    num_items: 762
+    scratch_space_bytes: 6399
+  }
+  scratch_size_recordings {
+    num_items: 823
+    scratch_space_bytes: 6911
+  }
+  scratch_size_recordings {
+    num_items: 889
+    scratch_space_bytes: 7423
+  }
+  scratch_size_recordings {
+    num_items: 960
+    scratch_space_bytes: 7935
+  }
+  scratch_size_recordings {
+    num_items: 1017
+    scratch_space_bytes: 8447
+  }
+  scratch_size_recordings {
+    num_items: 1078
+    scratch_space_bytes: 8959
+  }
+  scratch_size_recordings {
+    num_items: 1142
+    scratch_space_bytes: 9471
+  }
+  scratch_size_recordings {
+    num_items: 1210
+    scratch_space_bytes: 9983
+  }
+  scratch_size_recordings {
+    num_items: 1258
+    scratch_space_bytes: 10495
+  }
+  scratch_size_recordings {
+    num_items: 1334
+    scratch_space_bytes: 11007
+  }
+  scratch_size_recordings {
+    num_items: 1387
+    scratch_space_bytes: 11519
+  }
+  scratch_size_recordings {
+    num_items: 1470
+    scratch_space_bytes: 12031
+  }
+  scratch_size_recordings {
+    num_items: 1528
+    scratch_space_bytes: 12543
+  }
+  scratch_size_recordings {
+    num_items: 1589
+    scratch_space_bytes: 13055
+  }
+  scratch_size_recordings {
+    num_items: 1652
+    scratch_space_bytes: 13567
+  }
+  scratch_size_recordings {
+    num_items: 1718
+    scratch_space_bytes: 14079
+  }
+  scratch_size_recordings {
+    num_items: 1787
+    scratch_space_bytes: 14591
+  }
+  scratch_size_recordings {
+    num_items: 1822
+    scratch_space_bytes: 15103
+  }
+  scratch_size_recordings {
+    num_items: 1895
+    scratch_space_bytes: 15615
+  }
+  scratch_size_recordings {
+    num_items: 1970
+    scratch_space_bytes: 16127
+  }
+  scratch_size_recordings {
+    num_items: 2009
+    scratch_space_bytes: 16639
+  }
+  scratch_size_recordings {
+    num_items: 2089
+    scratch_space_bytes: 17151
+  }
+  scratch_size_recordings {
+    num_items: 2172
+    scratch_space_bytes: 17663
+  }
+  scratch_size_recordings {
+    num_items: 2215
+    scratch_space_bytes: 18175
+  }
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 18687
+  }
+  scratch_size_recordings {
+    num_items: 2350
+    scratch_space_bytes: 19199
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 19711
+  }
+  scratch_size_recordings {
+    num_items: 2492
+    scratch_space_bytes: 20223
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 20735
+  }
+  scratch_size_recordings {
+    num_items: 2591
+    scratch_space_bytes: 21247
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 21759
+  }
+  scratch_size_recordings {
+    num_items: 2747
+    scratch_space_bytes: 22271
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 22783
+  }
+  scratch_size_recordings {
+    num_items: 2857
+    scratch_space_bytes: 23295
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 23807
+  }
+  scratch_size_recordings {
+    num_items: 2972
+    scratch_space_bytes: 24319
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 24831
+  }
+  scratch_size_recordings {
+    num_items: 3091
+    scratch_space_bytes: 25343
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 25855
+  }
+  scratch_size_recordings {
+    num_items: 3215
+    scratch_space_bytes: 26367
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 26879
+  }
+  scratch_size_recordings {
+    num_items: 3344
+    scratch_space_bytes: 27391
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 27903
+  }
+  scratch_size_recordings {
+    num_items: 3478
+    scratch_space_bytes: 28415
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 28927
+  }
+  scratch_size_recordings {
+    num_items: 3617
+    scratch_space_bytes: 29439
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 29951
+  }
+  scratch_size_recordings {
+    num_items: 3762
+    scratch_space_bytes: 30463
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 30975
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 31999
+  }
+  scratch_size_recordings {
+    num_items: 3991
+    scratch_space_bytes: 32511
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 33023
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 33535
+  }
+  scratch_size_recordings {
+    num_items: 4234
+    scratch_space_bytes: 34559
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 35071
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 35583
+  }
+  scratch_size_recordings {
+    num_items: 4492
+    scratch_space_bytes: 36607
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 37119
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 37631
+  }
+  scratch_size_recordings {
+    num_items: 4765
+    scratch_space_bytes: 38655
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 39167
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 40191
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 40703
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 41727
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 42751
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 43263
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 44287
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 45311
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 45823
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 46847
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 47871
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 48895
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 49919
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 50943
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 51967
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 52991
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 54015
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 55039
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 56063
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 57087
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 58111
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 59135
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 60671
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 61695
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 62719
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 64255
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 65279
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 66815
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 67839
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 69375
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 70911
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 72447
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 73471
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 75007
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 76543
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 78079
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 79615
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 81151
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 82687
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 84735
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 86271
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 87807
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 89855
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 91391
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 93439
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 94975
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 97023
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 99071
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 100607
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 102655
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 104703
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 106751
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 109311
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 111359
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 113407
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 115455
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 118015
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 120575
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 122623
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 125183
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 127743
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 130303
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 132863
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 135423
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 137983
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 141055
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 143615
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 146687
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 149247
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 152319
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 155391
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 158463
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 161535
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 165119
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 168191
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 171775
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 174847
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 178431
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 182015
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 185599
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 189183
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 193279
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 196863
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 200959
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 205055
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 209151
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 213247
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 217343
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 221951
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 226047
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 230655
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 235263
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 239871
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 244991
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 249599
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 254719
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 259839
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 264959
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 270079
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 275711
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 280831
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 286463
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 292607
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 298239
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 304383
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 310015
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 316671
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 322815
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 328959
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 335615
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 342271
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 349439
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 356095
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 363263
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 370431
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 378111
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 385791
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 393471
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 401151
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 409343
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 417535
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 425727
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 433919
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 442623
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 451839
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 460543
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 469759
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 479487
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 488703
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 498431
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 508671
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 518911
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 529151
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 539391
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 550655
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 561407
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 572671
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 583935
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 595711
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 607487
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 619775
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 632063
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 644863
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 657663
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 670975
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 684287
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 698111
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 711935
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 726271
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 740607
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 755455
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 770303
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 785663
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 801535
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 817407
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 833791
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 850687
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 867583
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 884991
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 902399
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 920831
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 939263
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 957695
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 977151
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 996607
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 1016575
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 1036543
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 1057535
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 1078527
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 1100031
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 1122047
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 1144575
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 1167615
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 1190655
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 1214719
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 1238783
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 1263359
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 1288959
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 1314559
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 1340671
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 1367807
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 1394943
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 1423103
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 1451263
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 1480447
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 1510143
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 1540351
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 1571071
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 1602303
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 1634559
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 1666815
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 1700095
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 1734399
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 1769215
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 1804543
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 1840383
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 1877247
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 1914623
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 1953023
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 1991935
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 2031871
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 2072319
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 2113791
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 2156287
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 2199295
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 2243327
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 2288383
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 2333951
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 2380543
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 2428159
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 2476799
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 2526463
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 2576639
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 2628351
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 2681087
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 2734335
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 2789119
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 2844927
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 2901759
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 2959615
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 3019007
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 3079423
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 3140863
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 3203839
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 3267839
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 3333375
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 3399935
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 3467519
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 3537151
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 3607807
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 3679999
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 3753727
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 3828479
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 3905279
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 3983103
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 4062975
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 4144383
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 4226815
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 4311807
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 4397823
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 4485887
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 4575487
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 4667135
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 4760319
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 4855551
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 4952319
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 5051647
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 5152511
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 5255423
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 5360895
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 5467903
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 5577471
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 5688575
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 5802751
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 5918463
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 6036735
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 6157567
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 6280959
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 6406399
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 6534399
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 6664959
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 6798591
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 6934271
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 7073023
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 7214335
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 7358719
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 7506175
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 7656191
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 7809279
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 7965439
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 8124671
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 8286975
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 8452863
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 8621823
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 8794367
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 8969983
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 9149695
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 9332479
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 9519359
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 9709823
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 9903871
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 10102015
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 10303743
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 10510079
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 10719999
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 10934527
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 11153151
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 11376383
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 11603711
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 11835647
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 12072703
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 12313855
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 12560127
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 12811519
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 13067519
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 13329151
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 13595391
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 13867263
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 14144767
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 14427903
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 14716159
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 15010559
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 15310591
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 15616767
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 15929087
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 16247551
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 16572671
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 16903935
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 17242367
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 17586943
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 17938687
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 18297599
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 18663679
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 19036927
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 19417343
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 19805951
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 20201727
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 20605695
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 21017855
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 21438207
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 21867263
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 22304511
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 22750463
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 23205631
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 23669503
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 24143103
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 24625919
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 25118463
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 25620735
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 26133247
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 26655999
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 27188991
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 27732735
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 28287231
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 28852991
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 29430015
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 30018815
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 30618879
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 31231231
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 31855871
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 32493311
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 33143039
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 33805567
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 34481919
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 35171583
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 35875071
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 36592383
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 37324031
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 38070527
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 38831871
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 39608575
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 40400639
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 41209087
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 42032895
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 42873599
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 43731199
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 44605695
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 45498111
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 46407935
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 47336191
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 48282879
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 49248511
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 50233087
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 51238143
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 52262655
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 53307647
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 54374143
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 55461631
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 56570623
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 57702143
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 58856191
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 60033279
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 61233919
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 62458623
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 63707903
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 64981759
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 66281727
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 67607295
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 68958975
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 70338303
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 71745279
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 73179903
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 74643711
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 76136703
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 77659391
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 79212287
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 80796415
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 82412287
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 84060927
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 85741823
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 87456511
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 89206015
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 90989823
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 92809983
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 94665983
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 96559359
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 98490623
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 100460287
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 102469375
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 104518911
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 106608895
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 108741375
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 110916351
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 113134335
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 115396863
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 117704959
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 120059135
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 122460415
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 124909311
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 127407871
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 129955583
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 132555007
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 135206143
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 137910015
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 140668159
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 143481599
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 146351359
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 149278463
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 152263935
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 155309311
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 158415103
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 161583359
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 164815103
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 168111359
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 171473663
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 174903039
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 178401023
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 181969151
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 185608447
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 189320959
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 193107199
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 196969215
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 200908543
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 204926719
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 209025279
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 213205759
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 217469695
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 221819135
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 226255615
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 230780671
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 235396351
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 240104191
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 244906239
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 249804543
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 254800639
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 259896575
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 265094399
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 270396159
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 275804415
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 281320191
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 286946559
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 292685567
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 298539263
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 304510207
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 310600447
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 316812031
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 323148543
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 329611519
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 336203519
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 342927615
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 349786367
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 356781823
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 363917567
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 371195647
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 378619647
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 386192127
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 393916159
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 401794303
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 409830143
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 418026751
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 426387199
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 434915071
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 443613439
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 452485375
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 461535487
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 470765823
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 480180991
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 489785087
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 499580671
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 509572351
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 519763711
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 530158847
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 540761855
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 551577343
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 562608895
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 573861119
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 585338111
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 597044991
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 608985855
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 621165311
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 633588479
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 646260479
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 659185407
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 672369407
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 685816575
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 699533055
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 713523455
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 727793919
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 742350079
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 757197055
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 772340991
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 787787519
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 803543295
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 819614463
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 836006655
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 852726527
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 869781247
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 887176959
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 904920319
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 923019007
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 941479167
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 960308479
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 979514623
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 999105279
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 1019087103
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 1039468799
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 1060258559
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 1081463551
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 1103092991
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 1125154559
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 1147657471
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 1170610943
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 1194023167
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 1217903359
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 1242261247
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 1267106559
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 1292449023
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 1318297855
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 1344663807
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 1371557119
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 1398988031
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 1426967807
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 1455507199
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 1484617471
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 1514309887
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 1544595711
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 1575487743
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 1606997759
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 1639137535
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 1671920383
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 1705358591
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 1739465983
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 1774254847
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 1809740031
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 1845934847
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 1882853631
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 1920510719
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 1958920959
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 1998099199
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 2038061311
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 2078822655
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 2120399103
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 2162807039
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 2206062847
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 2250184447
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 2295187711
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 2341091583
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 2387913471
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 2435671807
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 2484385023
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 2534072575
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 2584754431
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 2636449535
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 2689178367
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 2742961919
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 2797821183
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 2853777663
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 2910852863
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 2969069823
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 3028451583
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 3089020671
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 3150800639
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 3213817087
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 3278093055
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 3343655167
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 3410527999
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 3478738687
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 3548313343
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 3619279615
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 3691665151
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 3765498623
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 3840808703
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 3917624575
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 3995976959
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 4075896575
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 4157414655
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 4240562943
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 4325374207
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 4411881727
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 4500119295
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 4590121727
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 4681923839
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 4775562495
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 4871073535
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 4968494847
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 5067864831
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 5169222399
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 5272606463
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 5378059007
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 5485619967
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 5595332351
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 5707239167
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 5821383935
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 5937811199
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 6056567551
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 6177699071
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 6301252863
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 6427278079
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 6555823359
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 6686939903
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 6820678911
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 6957092095
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 7096234239
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 7238159103
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 7382921983
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 7530580223
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 7681192191
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 7834815743
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 7991512319
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 4
+  value_type_size: 8
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 1
+  }
+  scratch_size_recordings {
+    num_items: 2350
+    scratch_space_bytes: 38143
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 38655
+  }
+  scratch_size_recordings {
+    num_items: 2444
+    scratch_space_bytes: 39423
+  }
+  scratch_size_recordings {
+    num_items: 2492
+    scratch_space_bytes: 39679
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 40447
+  }
+  scratch_size_recordings {
+    num_items: 2591
+    scratch_space_bytes: 40959
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 41727
+  }
+  scratch_size_recordings {
+    num_items: 2694
+    scratch_space_bytes: 42495
+  }
+  scratch_size_recordings {
+    num_items: 2747
+    scratch_space_bytes: 42751
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 43519
+  }
+  scratch_size_recordings {
+    num_items: 2857
+    scratch_space_bytes: 44287
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 45055
+  }
+  scratch_size_recordings {
+    num_items: 2972
+    scratch_space_bytes: 45567
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 46335
+  }
+  scratch_size_recordings {
+    num_items: 3091
+    scratch_space_bytes: 47103
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 47871
+  }
+  scratch_size_recordings {
+    num_items: 3215
+    scratch_space_bytes: 48639
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 49407
+  }
+  scratch_size_recordings {
+    num_items: 3344
+    scratch_space_bytes: 50175
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 50943
+  }
+  scratch_size_recordings {
+    num_items: 3478
+    scratch_space_bytes: 51711
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 52479
+  }
+  scratch_size_recordings {
+    num_items: 3617
+    scratch_space_bytes: 53503
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 54271
+  }
+  scratch_size_recordings {
+    num_items: 3762
+    scratch_space_bytes: 55039
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 55807
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 57087
+  }
+  scratch_size_recordings {
+    num_items: 3991
+    scratch_space_bytes: 57855
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 58879
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 59647
+  }
+  scratch_size_recordings {
+    num_items: 4234
+    scratch_space_bytes: 60927
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 61695
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 62719
+  }
+  scratch_size_recordings {
+    num_items: 4492
+    scratch_space_bytes: 63999
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 65023
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 65791
+  }
+  scratch_size_recordings {
+    num_items: 4765
+    scratch_space_bytes: 67071
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 68095
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 69375
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 70399
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 71935
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 73215
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 74239
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 76543
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 78079
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 79103
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 80639
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 82175
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 83455
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 84991
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 86527
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 88063
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 89599
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 91135
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 92671
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 94207
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 95743
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 97535
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 99071
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 101119
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 102911
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 104447
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 106495
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 108287
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 110335
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 112127
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 114431
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 116479
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 118783
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 120575
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 122879
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 125183
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 127487
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 129791
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 132095
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 134399
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 137215
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 139519
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 143103
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 145919
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 148223
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 151295
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 153855
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 156671
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 159743
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 162303
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 165375
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 168447
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 171519
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 175103
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 178175
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 181503
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 184575
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 188159
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 191999
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 195327
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 199167
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 202751
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 207615
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 211455
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 215551
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 219391
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 223743
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 227839
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 232191
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 236287
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 240895
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 245503
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 250111
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 254719
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 259839
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 264703
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 269823
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 275711
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 281087
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 286463
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 291839
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 297215
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 303103
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 308735
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 314623
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 320767
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 326911
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 333055
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 340479
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 347135
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 353535
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 360447
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 367359
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 374271
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 381695
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 388863
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 396287
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 404991
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 412671
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 420607
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 428799
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 436735
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 445183
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 454143
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 462847
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 472831
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 481535
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 491263
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 500479
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 509951
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 519935
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 529919
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 541439
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 551679
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 562431
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 573183
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 584447
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 595967
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 608511
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 620031
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 632319
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 644607
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 656895
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 670463
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 683519
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 697087
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 710399
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 724223
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 739583
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 753663
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 768255
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 783359
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 799743
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 815103
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 830719
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 847359
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 864767
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 881407
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 898559
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 916223
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 934911
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 953087
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 971775
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 990719
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 1011199
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 1030911
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 1050879
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 1072639
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 1093375
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 1114879
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 1137407
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 1159679
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 1182207
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 1206271
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 1230079
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 1253887
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 1279487
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 1304575
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 1330943
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 1357055
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 1383423
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 1411839
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 1439487
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 1468415
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 1497343
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 1527551
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 1557503
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 1588735
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 1619967
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 1652735
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 1684991
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 1719039
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 1752575
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 1788159
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 1822975
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 1859839
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 1896191
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 1934079
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 1972223
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 2011903
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 2052095
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 2092543
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 2134527
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 2177535
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 2220031
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 2264575
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 2310143
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 2355455
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 2402559
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 2450431
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 2498815
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 2548479
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 2599423
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 2651647
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 2704895
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 2757887
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 2812927
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 2868991
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 2926335
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 2984703
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 3044351
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 3105279
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 3166975
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 3230207
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 3294719
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 3360511
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 3427327
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 3495935
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 3565311
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 3636479
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 3708927
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 3782655
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 3858175
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 3934719
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 4014079
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 4094207
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 4175359
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 4258303
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 4344063
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 4430335
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 4518399
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 4609279
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 4700927
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 4795391
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 4890623
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 4987647
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 5087999
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 5189887
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 5292543
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 5398783
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 5506047
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 5616127
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 5728767
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 5843199
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 5959167
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 6078207
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 6199807
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 6323967
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 6449919
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 6579199
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 6710271
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 6844415
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 6980863
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 7120383
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 7262207
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 7407103
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 7555583
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 7706367
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 7859967
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 8017407
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 8177407
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 8341247
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 8507391
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 8677375
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 8850431
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 9027327
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 9207807
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 9392127
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 9580031
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 9770495
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 9965567
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 10164479
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 10368767
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 10575615
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 10786815
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 11001855
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 11222527
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 11446527
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 11675647
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 11908351
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 12146687
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 12388863
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 12636415
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 12889343
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 13146879
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 13409535
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 13677311
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 13950719
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 14229247
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 14514431
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 14804223
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 15100415
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 15401727
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 15709695
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 16024063
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 16344319
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 16670975
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 17004287
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 17343999
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 17690367
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 18044415
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 18404863
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 18772991
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 19148543
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 19530495
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 19920895
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 20319231
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 20726015
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 21139967
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 21562367
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 21993983
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 22433791
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 22882303
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 23338751
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 23806207
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 24281855
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 24766719
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 25262591
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 25767679
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 26282239
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 26807551
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 27344383
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 27890687
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 28447999
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 29017087
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 29597183
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 30189055
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 30792191
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 31408383
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 32035839
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 32676607
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 33330175
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 33996287
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 34675711
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 35369471
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 36076799
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 36798207
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 37533695
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 38284799
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 39049471
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 39830527
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 40626943
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 41439743
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 42267647
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 43112959
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 43974655
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 44854783
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 45751039
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 46665727
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 47599615
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 48550911
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 49522175
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 50512127
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 51522559
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 52552191
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 53603071
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 54675711
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 55768319
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 56883199
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 58021375
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 59180799
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 60364287
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 61571839
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 62803711
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 64059135
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 65339647
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 66646271
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 67979775
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 69338623
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 70725887
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 72139775
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 73582335
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 75054335
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 76553983
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 78085887
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 79647231
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 81239551
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 82864639
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 84521471
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 86211839
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 87935487
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 89694207
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 91487487
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 93317119
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 95183871
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 97087487
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 99028223
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 101009663
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 103028735
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 105089279
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 107191039
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 109335039
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 111521023
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 113751295
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 116026623
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 118346495
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 120712959
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 123127039
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 125589759
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 128101375
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 130662911
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 133275903
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 135941887
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 138659839
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 141432831
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 144261887
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 147147007
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 150089983
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 153090559
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 156152575
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 159275775
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 162460159
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 165709823
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 169023487
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 172403967
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 175852031
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 179368447
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 182955775
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 186615039
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 190347007
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 194153983
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 198036223
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 201997567
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 206036735
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 210157311
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 214360063
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 218647807
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 223019775
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 227480063
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 232029951
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 236670719
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 241403391
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 246231551
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 251155967
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 256178175
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 261302527
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 266527487
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 271858431
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 277295615
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 282840575
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 288497151
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 294267647
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 300152063
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 306155263
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 312277759
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 318523903
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 324893951
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 331391231
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 338018815
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 344779007
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 351674111
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 358707967
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 365881599
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 373199359
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 380663295
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 388276223
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 396040959
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 403961855
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 412041471
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 420281599
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 428687615
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 437260287
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 446006015
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 454925823
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 464024319
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 473304063
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 482770175
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 492425215
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 502273279
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 512319231
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 522565119
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 533016831
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 543676159
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 554549503
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 565639935
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 576953087
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 588492031
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 600262143
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 612266751
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 624512255
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 637001727
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 649742079
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 662736383
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 675990783
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 689509887
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 703300863
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 717365759
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 731712767
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 746348031
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 761273855
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 776499967
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 792029695
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 807869183
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 824026879
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 840506879
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 857317375
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 874463999
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 891951871
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 909791743
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 927986687
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 946546943
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 965476607
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 984786431
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 1004482303
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 1024571391
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 1045062911
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 1065964031
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 1087282687
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 1109028351
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 1131208447
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 1153832191
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 1176909567
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 1200447231
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 1224455679
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 1248945151
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 1273923583
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 1299401983
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 1325389567
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 1351897599
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 1378934783
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 1406513663
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 1434643199
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 1463335935
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 1492602879
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 1522454527
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 1552903679
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 1583961599
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 1615640575
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 1647952895
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 1680912127
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 1714530047
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 1748819967
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 1783796735
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 1819472127
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 1855862015
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 1892978431
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 1930838271
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 1969454847
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 2008843519
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 2049020159
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 2090000127
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 2131800063
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 2174436351
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 2217924351
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 2262282751
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 2307527935
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 2353678847
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 2400752127
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 2448767487
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 2497742079
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 2547696639
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 2598650623
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 2650623743
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 2703635711
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 2757708287
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 2812861695
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 2869119231
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 2926500863
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 2985030655
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 3044731135
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 3105625343
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 3167738111
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 3231092991
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 3295357183
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 3360241151
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 3426423295
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 3493928447
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 3562784255
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 3633016831
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 3704654335
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 3777724415
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 3852255743
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 3928278271
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 4005820927
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 4084914175
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 4165589503
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 4247878399
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 4331813119
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 4417426175
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 4504751615
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 4593823999
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 4684677631
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 4777347839
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 4871872255
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 4968286463
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 5066629375
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 5166938879
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 5269254655
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 5373616895
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 5480066303
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 5588644607
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 5699394559
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 5812359679
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 5927583743
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 6045112319
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 6164991743
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 6287268607
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 6411991039
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 6539216127
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 6668977407
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 6801333759
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 6936337407
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 7074040831
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 7214498559
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 7357765375
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 7503897343
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 7652952319
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 7804988415
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 7960064767
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 8118243327
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 8279585023
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 8444153599
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 8612013567
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 8783230719
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 8957871871
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 9136006399
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 9317703423
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 9503034367
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 9692071935
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 9884898303
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 10081573119
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 10282181375
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 10486801407
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 10695514367
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 10908401663
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 11125546239
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 11347033599
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 11572951295
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 11803386879
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 12038431487
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 4
+  value_type_size: 8
+  is_segmented: true
+  scratch_size_recordings {
+    num_items: 31
+    scratch_space_bytes: 767
+  }
+  scratch_size_recordings {
+    num_items: 61
+    scratch_space_bytes: 1023
+  }
+  scratch_size_recordings {
+    num_items: 91
+    scratch_space_bytes: 1535
+  }
+  scratch_size_recordings {
+    num_items: 121
+    scratch_space_bytes: 1791
+  }
+  scratch_size_recordings {
+    num_items: 151
+    scratch_space_bytes: 2303
+  }
+  scratch_size_recordings {
+    num_items: 191
+    scratch_space_bytes: 2559
+  }
+  scratch_size_recordings {
+    num_items: 221
+    scratch_space_bytes: 3071
+  }
+  scratch_size_recordings {
+    num_items: 251
+    scratch_space_bytes: 3327
+  }
+  scratch_size_recordings {
+    num_items: 281
+    scratch_space_bytes: 3839
+  }
+  scratch_size_recordings {
+    num_items: 311
+    scratch_space_bytes: 4095
+  }
+  scratch_size_recordings {
+    num_items: 351
+    scratch_space_bytes: 4607
+  }
+  scratch_size_recordings {
+    num_items: 381
+    scratch_space_bytes: 4863
+  }
+  scratch_size_recordings {
+    num_items: 411
+    scratch_space_bytes: 5375
+  }
+  scratch_size_recordings {
+    num_items: 441
+    scratch_space_bytes: 5631
+  }
+  scratch_size_recordings {
+    num_items: 471
+    scratch_space_bytes: 6143
+  }
+  scratch_size_recordings {
+    num_items: 511
+    scratch_space_bytes: 6399
+  }
+  scratch_size_recordings {
+    num_items: 541
+    scratch_space_bytes: 6911
+  }
+  scratch_size_recordings {
+    num_items: 573
+    scratch_space_bytes: 7167
+  }
+  scratch_size_recordings {
+    num_items: 606
+    scratch_space_bytes: 7679
+  }
+  scratch_size_recordings {
+    num_items: 630
+    scratch_space_bytes: 7935
+  }
+  scratch_size_recordings {
+    num_items: 667
+    scratch_space_bytes: 8447
+  }
+  scratch_size_recordings {
+    num_items: 693
+    scratch_space_bytes: 8703
+  }
+  scratch_size_recordings {
+    num_items: 734
+    scratch_space_bytes: 9215
+  }
+  scratch_size_recordings {
+    num_items: 762
+    scratch_space_bytes: 9471
+  }
+  scratch_size_recordings {
+    num_items: 792
+    scratch_space_bytes: 9983
+  }
+  scratch_size_recordings {
+    num_items: 823
+    scratch_space_bytes: 10239
+  }
+  scratch_size_recordings {
+    num_items: 855
+    scratch_space_bytes: 10751
+  }
+  scratch_size_recordings {
+    num_items: 889
+    scratch_space_bytes: 11007
+  }
+  scratch_size_recordings {
+    num_items: 924
+    scratch_space_bytes: 11519
+  }
+  scratch_size_recordings {
+    num_items: 960
+    scratch_space_bytes: 11775
+  }
+  scratch_size_recordings {
+    num_items: 979
+    scratch_space_bytes: 12287
+  }
+  scratch_size_recordings {
+    num_items: 1017
+    scratch_space_bytes: 12543
+  }
+  scratch_size_recordings {
+    num_items: 1037
+    scratch_space_bytes: 13055
+  }
+  scratch_size_recordings {
+    num_items: 1078
+    scratch_space_bytes: 13311
+  }
+  scratch_size_recordings {
+    num_items: 1120
+    scratch_space_bytes: 13823
+  }
+  scratch_size_recordings {
+    num_items: 1142
+    scratch_space_bytes: 14079
+  }
+  scratch_size_recordings {
+    num_items: 1164
+    scratch_space_bytes: 14591
+  }
+  scratch_size_recordings {
+    num_items: 1210
+    scratch_space_bytes: 14847
+  }
+  scratch_size_recordings {
+    num_items: 1234
+    scratch_space_bytes: 15359
+  }
+  scratch_size_recordings {
+    num_items: 1258
+    scratch_space_bytes: 15615
+  }
+  scratch_size_recordings {
+    num_items: 1308
+    scratch_space_bytes: 16127
+  }
+  scratch_size_recordings {
+    num_items: 1334
+    scratch_space_bytes: 16383
+  }
+  scratch_size_recordings {
+    num_items: 1360
+    scratch_space_bytes: 16895
+  }
+  scratch_size_recordings {
+    num_items: 1387
+    scratch_space_bytes: 17151
+  }
+  scratch_size_recordings {
+    num_items: 1414
+    scratch_space_bytes: 17663
+  }
+  scratch_size_recordings {
+    num_items: 1470
+    scratch_space_bytes: 17919
+  }
+  scratch_size_recordings {
+    num_items: 1499
+    scratch_space_bytes: 18431
+  }
+  scratch_size_recordings {
+    num_items: 1528
+    scratch_space_bytes: 18687
+  }
+  scratch_size_recordings {
+    num_items: 1558
+    scratch_space_bytes: 19199
+  }
+  scratch_size_recordings {
+    num_items: 1589
+    scratch_space_bytes: 19455
+  }
+  scratch_size_recordings {
+    num_items: 1620
+    scratch_space_bytes: 19967
+  }
+  scratch_size_recordings {
+    num_items: 1652
+    scratch_space_bytes: 20223
+  }
+  scratch_size_recordings {
+    num_items: 1685
+    scratch_space_bytes: 20735
+  }
+  scratch_size_recordings {
+    num_items: 1718
+    scratch_space_bytes: 20991
+  }
+  scratch_size_recordings {
+    num_items: 1752
+    scratch_space_bytes: 21503
+  }
+  scratch_size_recordings {
+    num_items: 1787
+    scratch_space_bytes: 21759
+  }
+  scratch_size_recordings {
+    num_items: 1822
+    scratch_space_bytes: 22271
+  }
+  scratch_size_recordings {
+    num_items: 1858
+    scratch_space_bytes: 23039
+  }
+  scratch_size_recordings {
+    num_items: 1895
+    scratch_space_bytes: 23295
+  }
+  scratch_size_recordings {
+    num_items: 1932
+    scratch_space_bytes: 23807
+  }
+  scratch_size_recordings {
+    num_items: 1970
+    scratch_space_bytes: 24063
+  }
+  scratch_size_recordings {
+    num_items: 2009
+    scratch_space_bytes: 24575
+  }
+  scratch_size_recordings {
+    num_items: 2049
+    scratch_space_bytes: 25343
+  }
+  scratch_size_recordings {
+    num_items: 2089
+    scratch_space_bytes: 25599
+  }
+  scratch_size_recordings {
+    num_items: 2130
+    scratch_space_bytes: 26111
+  }
+  scratch_size_recordings {
+    num_items: 2172
+    scratch_space_bytes: 26367
+  }
+  scratch_size_recordings {
+    num_items: 2215
+    scratch_space_bytes: 27135
+  }
+  scratch_size_recordings {
+    num_items: 2259
+    scratch_space_bytes: 27647
+  }
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 27903
+  }
+  scratch_size_recordings {
+    num_items: 2350
+    scratch_space_bytes: 28671
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 29183
+  }
+  scratch_size_recordings {
+    num_items: 2444
+    scratch_space_bytes: 29951
+  }
+  scratch_size_recordings {
+    num_items: 2492
+    scratch_space_bytes: 30207
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 30975
+  }
+  scratch_size_recordings {
+    num_items: 2591
+    scratch_space_bytes: 31487
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 32255
+  }
+  scratch_size_recordings {
+    num_items: 2694
+    scratch_space_bytes: 33023
+  }
+  scratch_size_recordings {
+    num_items: 2747
+    scratch_space_bytes: 33279
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 34047
+  }
+  scratch_size_recordings {
+    num_items: 2857
+    scratch_space_bytes: 34815
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 35583
+  }
+  scratch_size_recordings {
+    num_items: 2972
+    scratch_space_bytes: 36095
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 36863
+  }
+  scratch_size_recordings {
+    num_items: 3091
+    scratch_space_bytes: 37631
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 38399
+  }
+  scratch_size_recordings {
+    num_items: 3215
+    scratch_space_bytes: 39167
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 39935
+  }
+  scratch_size_recordings {
+    num_items: 3344
+    scratch_space_bytes: 40703
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 41471
+  }
+  scratch_size_recordings {
+    num_items: 3478
+    scratch_space_bytes: 42239
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 43007
+  }
+  scratch_size_recordings {
+    num_items: 3617
+    scratch_space_bytes: 44031
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 44799
+  }
+  scratch_size_recordings {
+    num_items: 3762
+    scratch_space_bytes: 45567
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 46335
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 47615
+  }
+  scratch_size_recordings {
+    num_items: 3991
+    scratch_space_bytes: 48383
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 49407
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 50175
+  }
+  scratch_size_recordings {
+    num_items: 4234
+    scratch_space_bytes: 51455
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 52223
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 53247
+  }
+  scratch_size_recordings {
+    num_items: 4492
+    scratch_space_bytes: 54527
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 55551
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 56319
+  }
+  scratch_size_recordings {
+    num_items: 4765
+    scratch_space_bytes: 57599
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 58623
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 59903
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 60927
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 62463
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 63743
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 64767
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 66047
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 67583
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 68607
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 70143
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 71679
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 72959
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 74495
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 76031
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 77567
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 79103
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 80639
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 82175
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 83711
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 85247
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 87039
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 88575
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 90623
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 92415
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 93951
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 95999
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 97791
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 99839
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 101631
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 103935
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 105983
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 108287
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 110079
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 112383
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 114687
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 116991
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 119295
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 121599
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 123903
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 126719
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 129023
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 131583
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 134399
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 136703
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 139775
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 142335
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 145151
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 148223
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 150783
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 153855
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 156927
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 159999
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 163583
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 166655
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 169983
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 173055
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 176639
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 180479
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 183807
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 187647
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 191231
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 195071
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 198911
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 203007
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 206847
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 211199
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 215295
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 219647
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 223743
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 228351
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 232959
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 237567
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 242175
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 247295
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 252159
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 257279
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 262143
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 267519
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 272895
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 278271
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 283647
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 289535
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 295167
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 301055
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 307199
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 313343
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 319487
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 325887
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 332543
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 338943
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 345855
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 352767
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 359679
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 367103
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 374271
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 381695
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 389375
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 397055
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 404991
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 413183
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 421119
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 429567
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 438527
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 447231
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 456191
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 464895
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 474623
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 483839
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 493311
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 503295
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 513279
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 523775
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 534015
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 544767
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 555519
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 566783
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 578303
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 589823
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 601343
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 613631
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 625919
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 638207
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 650751
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 663807
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 677375
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 690687
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 704511
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 718847
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 732927
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 747519
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 762623
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 777983
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 793343
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 808959
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 825599
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 841983
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 858623
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 875775
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 893439
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 911103
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 929279
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 947967
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 966911
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 986367
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 1006079
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 1026047
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 1046783
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 1067519
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 1089023
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 1110527
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 1132799
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 1155327
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 1178367
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 1202175
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 1225983
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 1250559
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 1275647
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 1300991
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 1327103
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 1353471
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 1380863
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 1408511
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 1436415
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 1465343
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 1494527
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 1524479
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 1554687
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 1585919
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 1617663
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 1649919
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 1682943
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 1716479
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 1751039
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 1785855
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 1821695
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 1858047
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 1894911
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 1933055
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 1971711
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 2010879
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 2051327
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 2092287
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 2134271
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 2176767
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 2220287
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 2264831
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 2310143
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 2356223
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 2403071
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 2451455
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 2500095
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 2550015
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 2601215
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 2653439
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 2706431
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 2760447
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 2815487
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 2871807
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 2929151
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 2987775
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 3047679
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 3108351
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 3170559
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 3234047
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 3298815
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 3364607
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 3432191
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 3500543
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 3570687
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 3642111
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 3714815
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 3789311
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 3864831
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 3942143
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 4021247
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 4101375
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 4183295
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 4267007
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 4352255
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 4439295
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 4528127
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 4618751
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 4711167
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 4805375
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 4901375
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 4999679
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 5099519
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 5201151
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 5305343
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 5411583
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 5519615
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 5630207
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 5742591
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 5857535
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 5974527
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 6094079
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 6216191
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 6340095
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 6467327
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 6596351
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 6728447
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 6862847
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 7000319
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 7140095
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 7282943
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 7428351
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 7577087
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 7728639
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 7883007
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 8040959
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 8201727
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 8365823
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 8532735
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 8703743
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 8877567
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 9054975
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 9236223
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 9421055
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 9609471
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 9801471
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 9997311
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 10197503
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 10401279
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 10609407
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 10821375
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 11037951
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 11258879
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 11483903
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 11713535
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 11947775
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 12186879
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 12430335
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 12679167
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 12932607
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 13191167
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 13454847
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 13724159
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 13998591
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 14278655
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 14564351
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 14855423
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 15152639
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 15455487
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 15764735
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 16079871
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 16401407
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 16729599
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 17064191
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 17405439
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 17753343
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 18108671
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 18470655
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 18840063
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 19216895
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 19601151
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 19993343
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 20392959
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 20800767
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 21217023
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 21641471
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 22074111
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 22515455
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 22965759
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 23425023
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 23893503
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 24371199
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 24858879
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 25355775
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 25863167
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 26380287
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 26907903
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 27446015
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 27995135
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 28555007
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 29125887
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 29708543
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 30302463
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 30908415
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 31526655
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 32157183
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 32800511
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 33456383
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 34125567
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 34808063
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 35504127
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 36214271
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 36938495
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 37677311
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 38430719
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 39199487
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 39983615
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 40783103
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 41598719
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 42430719
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 43279103
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 44144895
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 45027839
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 45928191
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 46846719
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 47783679
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 48739583
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 49714175
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 50708223
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 51722495
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 52756991
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 53812223
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 54888447
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 55985919
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 57105663
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 58247679
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 59412735
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 60600831
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 61813247
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 63049215
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 64310271
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 65596415
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 66908415
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 68246783
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 69611519
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 71003903
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 72423935
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 73872383
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 75349503
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 76856831
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 78393599
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 79961343
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 81560831
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 83192063
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 84855807
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 86552831
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 88283903
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 90049535
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 91850495
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 93687551
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 95561471
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 97472511
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 99422207
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 101410559
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 103438335
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 105507327
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 107617535
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 109769727
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 111965183
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 114204671
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 116488703
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 118818303
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 121194495
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 123618303
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 126091007
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 128612607
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 131184639
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 133808639
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 136484607
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 139214591
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 141998591
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 144838655
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 147735551
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 150690047
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 153703935
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 156777983
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 159913215
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 163111679
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 166374143
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 169701375
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 173095167
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 176557311
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 180088319
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 183690239
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 187363839
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 191111423
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 194933247
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 198832127
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 202808831
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 206864895
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 211002111
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 215222271
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 219526655
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 223917311
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 228395519
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 232963583
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 237622527
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 242374911
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 247222527
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 252166911
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 257210367
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 262354431
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 267601407
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 272953599
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 278412543
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 283981055
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 289660415
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 295453695
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 301362687
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 307389951
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 313537791
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 319808511
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 326204415
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 332728575
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 339383295
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 346170879
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 353094143
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 360156159
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 367359231
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 374706431
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 382200575
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 389844479
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 397641471
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 405594111
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 413706239
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 421980159
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 430419711
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 439028223
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 447808767
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 456764927
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 465900287
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 475217919
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 484722431
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 494416895
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 504305151
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 514391295
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 524679167
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 535172607
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 545875967
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 556793343
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 567929343
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 579288063
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 590873855
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 602691071
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 614745087
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 627039999
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 639580671
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 652372223
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 665419775
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 678727935
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 692302847
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 706148607
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 720271359
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 734677247
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 749370623
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 764358143
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 779645183
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 795237887
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 811142655
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 827365631
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 843912959
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 860791295
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 878006783
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 895567103
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 913478399
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 931747839
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 950382591
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 969390335
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 988777983
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 1008553727
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 1028724735
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 1049299199
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 1070285055
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 1091690751
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 1113524735
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 1135795199
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 1158511103
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 1181681151
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 1205314815
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 1229421311
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 1254009599
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 1279089663
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 1304671743
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 1330765055
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 1357380351
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 1384528127
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 1412218367
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 1440462591
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 1469271807
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 1498657535
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 1528630527
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 1559203071
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 1590387455
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 1622194943
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 1654639103
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 1687731711
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 1721486079
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 1755916031
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 1791034367
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 1826854911
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 1863391743
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 1900659711
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 1938673151
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 1977446399
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 2016995327
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 2057335295
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 2098481919
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 2140451583
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 2183260671
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 2226925823
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 2271464447
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 2316893439
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 2363231487
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 2410496255
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 2458705919
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 2507880191
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 2558037503
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 2609198591
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 2661382143
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 2714609919
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 2768902143
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 2824280063
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 2880765695
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 2938381055
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 2997148671
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 3057091583
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 3118233599
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 3180598271
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 3244210175
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 3309094143
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 3375276287
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 3442781439
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 3511637247
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 3581869823
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 3653507327
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 3726577407
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 3801108735
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 3877131263
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 3954673919
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 4033767167
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 4114442495
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 4196731391
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 4280666111
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 4366279167
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 4453604607
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 4542676991
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 4633530623
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 4726200831
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 4820725247
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 4917139455
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 5015482367
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 5115791871
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 5218107647
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 5322469887
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 5428919295
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 5537497599
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 5648247551
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 5761212671
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 5876436735
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 5993965311
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 6113844735
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 6236121599
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 6360844031
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 6488060927
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 6617822207
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 6750178559
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 6885182207
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 7022885631
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 7163343359
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 7306610175
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 7452742143
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 7601797119
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 7753833215
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 7908909567
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 8067088127
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 8228429823
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 8392998399
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 8560858367
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 8732075519
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 8906716671
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 9084851199
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 9266548223
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 9451879167
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 9640916735
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 9833734911
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 10030409727
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 10231017983
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 10435638015
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 10644350975
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 10857238271
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 11074382847
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 11295870207
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 11521787903
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 11752223487
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 11987268095
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 8
+  value_type_size: 2
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 1
+  }
+  scratch_size_recordings {
+    num_items: 2350
+    scratch_space_bytes: 41727
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 41983
+  }
+  scratch_size_recordings {
+    num_items: 2444
+    scratch_space_bytes: 42751
+  }
+  scratch_size_recordings {
+    num_items: 2492
+    scratch_space_bytes: 43007
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 43519
+  }
+  scratch_size_recordings {
+    num_items: 2591
+    scratch_space_bytes: 44031
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 44543
+  }
+  scratch_size_recordings {
+    num_items: 2694
+    scratch_space_bytes: 45311
+  }
+  scratch_size_recordings {
+    num_items: 2747
+    scratch_space_bytes: 45567
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 46079
+  }
+  scratch_size_recordings {
+    num_items: 2857
+    scratch_space_bytes: 46847
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 47359
+  }
+  scratch_size_recordings {
+    num_items: 2972
+    scratch_space_bytes: 47871
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 48383
+  }
+  scratch_size_recordings {
+    num_items: 3091
+    scratch_space_bytes: 49151
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 49663
+  }
+  scratch_size_recordings {
+    num_items: 3215
+    scratch_space_bytes: 50431
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 50943
+  }
+  scratch_size_recordings {
+    num_items: 3344
+    scratch_space_bytes: 51711
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 52223
+  }
+  scratch_size_recordings {
+    num_items: 3478
+    scratch_space_bytes: 52991
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 53503
+  }
+  scratch_size_recordings {
+    num_items: 3617
+    scratch_space_bytes: 54527
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 55039
+  }
+  scratch_size_recordings {
+    num_items: 3762
+    scratch_space_bytes: 55807
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 56319
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 57343
+  }
+  scratch_size_recordings {
+    num_items: 3991
+    scratch_space_bytes: 58111
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 58879
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 59647
+  }
+  scratch_size_recordings {
+    num_items: 4234
+    scratch_space_bytes: 60671
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 61183
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 62207
+  }
+  scratch_size_recordings {
+    num_items: 4492
+    scratch_space_bytes: 63231
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 63999
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 64767
+  }
+  scratch_size_recordings {
+    num_items: 4765
+    scratch_space_bytes: 65791
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 66559
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 67583
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 68607
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 69887
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 70911
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 71679
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 72703
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 73983
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 75007
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 77311
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 78591
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 79615
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 80895
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 82175
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 83455
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 84735
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 86015
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 87295
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 88575
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 89855
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 91391
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 92671
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 94207
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 95743
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 97023
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 98815
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 100351
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 101887
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 103423
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 105471
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 107007
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 109055
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 110591
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 112383
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 114431
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 116223
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 118271
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 120063
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 122111
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 124415
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 126207
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 128511
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 130815
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 132607
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 136191
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 138495
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 140799
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 143359
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 145407
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 147967
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 150527
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 153087
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 156159
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 158719
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 161535
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 164095
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 166911
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 170239
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 173055
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 176127
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 179199
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 182271
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 185599
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 188927
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 192255
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 196863
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 200191
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 203775
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 207359
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 211199
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 215039
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 218879
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 222719
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 226815
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 230911
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 235263
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 239359
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 243711
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 248319
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 253695
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 258303
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 263167
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 267775
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 272639
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 277759
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 282879
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 287999
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 293375
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 299007
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 304383
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 310015
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 316927
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 322559
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 328703
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 334847
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 340991
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 347391
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 353791
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 360447
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 367103
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 374783
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 381951
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 389375
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 396543
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 403967
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 411391
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 419327
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 427007
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 435967
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 444415
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 452607
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 461311
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 470015
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 478975
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 488959
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 498175
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 507903
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 517375
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 527103
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 537343
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 548607
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 558847
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 569343
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 580095
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 591359
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 602623
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 615167
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 626943
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 638719
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 651007
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 664575
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 677375
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 690175
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 703231
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 717055
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 731647
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 745471
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 759807
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 774655
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 790271
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 805375
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 820991
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 836863
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 854015
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 870399
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 887039
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 905471
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 922623
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 940543
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 959487
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 978175
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 996863
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 1017087
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 1037055
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 1056767
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 1078271
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 1099263
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 1120255
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 1143039
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 1165055
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 1187839
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 1211903
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 1235199
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 1260287
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 1284607
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 1310719
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 1335807
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 1361919
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 1389311
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 1416191
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 1444863
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 1472767
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 1502463
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 1531647
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 1562367
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 1592831
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 1624575
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 1656319
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 1689599
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 1723135
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 1756927
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 1791999
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 1827071
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 1863423
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 1900799
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 1937919
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 1976575
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 2015999
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 2055167
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 2096383
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 2138111
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 2179583
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 2223359
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 2267903
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 2312959
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 2358015
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 2404863
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 2452991
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 2501631
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 2551551
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 2602495
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 2653183
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 2705919
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 2759935
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 2814975
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 2870783
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 2928127
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 2985983
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 3045631
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 3106047
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 3167743
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 3230719
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 3294719
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 3361279
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 3428095
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 3495935
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 3565311
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 3635967
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 3708159
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 3782655
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 3857663
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 3934207
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 4013311
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 4092927
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 4173823
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 4257791
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 4342015
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 4428799
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 4516607
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 4607231
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 4698367
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 4792575
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 4888319
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 4985087
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 5084671
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 5186303
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 5288959
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 5394431
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 5502463
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 5612031
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 5724159
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 5838079
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 5954815
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 6073343
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 6194431
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 6317567
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 6443519
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 6572031
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 6703615
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 6837247
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 6973439
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 7113215
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 7254271
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 7398911
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 7546879
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 7696639
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 7850751
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 8007935
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 8167935
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 8329983
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 8496383
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 8666111
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 8839167
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 9015551
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 9195263
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 9378815
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 9565951
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 9757695
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 9951999
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 10150399
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 10353663
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 10559743
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 10771199
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 10986495
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 11204863
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 11428863
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 11657215
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 11890175
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 12127487
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 12369663
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 12616447
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 12869119
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 13125631
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 13388543
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 13655295
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 13928191
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 14206975
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 14489855
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 14779391
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 15074303
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 15375615
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 15683327
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 15996415
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 16316415
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 16641791
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 16974847
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 17314047
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 17658879
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 18011903
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 18372863
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 18739455
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 19113471
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 19495935
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 19884799
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 20282367
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 20687615
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 21101055
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 21522431
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 21953279
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 22391551
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 22839295
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 23295999
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 23761663
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 24236543
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 24720383
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 25214207
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 25718271
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 26232575
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 26756863
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 27291647
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 27836927
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 28393727
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 28960511
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 29539583
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 30129919
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 30731775
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 31346687
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 31973631
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 32611583
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 33264639
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 33929215
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 34607871
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 35298815
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 36004351
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 36724735
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 37458431
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 38207487
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 38971135
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 39749887
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 40545023
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 41356031
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 42182399
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 43025151
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 43885823
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 44763135
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 45657855
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 46571263
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 47501055
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 48451583
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 49419775
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 50407935
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 51415551
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 52444159
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 53492735
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 54562047
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 55653119
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 56764927
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 57900799
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 59058431
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 60239359
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 61443071
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 62671615
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 63925247
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 65202687
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 66507007
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 67835903
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 69192447
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 70576383
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 71987455
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 73427199
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 74895359
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 76392191
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 77919743
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 79478271
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 81067519
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 82688767
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 84342015
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 86028799
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 87748351
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 89503231
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 91292415
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 93118719
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 94979839
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 96879871
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 98817023
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 100793343
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 102808319
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 104863999
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 106961407
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 109099519
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 111280895
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 113506303
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 115776255
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 118092287
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 120453119
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 122861823
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 125318911
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 127825151
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 130380799
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 132988415
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 135647231
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 138359807
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 141127679
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 143949567
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 146827775
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 149764095
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 152758783
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 155813631
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 158929407
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 162107647
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 165349887
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 168656127
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 172029439
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 175469055
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 178978047
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 182557439
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 186208767
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 189932287
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 193730559
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 197604863
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 201555967
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 205586943
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 209698303
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 213892351
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 218170111
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 222532607
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 226982911
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 231521791
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 236151807
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 240875263
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 245691903
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 250605823
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 255617023
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 260729855
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 265944063
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 271261951
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 276687103
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 282220031
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 287865087
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 293621759
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 299493375
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 305483007
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 311591935
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 317824255
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 324179967
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 330663423
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 337275647
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 344020735
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 350901503
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 357919487
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 365077247
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 372378111
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 379825151
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 387421695
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 395169791
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 403072767
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 411133439
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 419356415
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 427743231
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 436296703
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 445022463
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 453922815
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 463000831
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 472260607
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 481705471
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 491339007
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 501165823
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 511188479
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 521411839
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 531839743
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 542476799
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 553325823
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 564392191
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 575678975
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 587192063
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 598936063
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 610914047
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 623132927
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 635594495
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 648306175
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 661271807
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 674497023
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 687986687
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 701745407
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 715780607
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 730095615
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 744696831
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 759590655
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 774781951
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 790277887
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 806082303
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 822203647
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 838647807
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 855420415
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 872528639
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 889978111
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 907777535
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 925933055
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 944451583
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 963340031
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 982605823
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 1002258431
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 1022302207
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 1042748415
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 1063603199
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 1084875007
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 1106572287
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 1128702719
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 1151277311
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 1174301951
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 1197787647
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 1221743359
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 1246177023
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 1271100415
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 1296522239
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 1322451967
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 1348901887
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 1375879167
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 1403396095
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 1431463935
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 1460092159
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 1489293823
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 1519079167
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 1549461247
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 1580449279
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 1612057855
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 1644298751
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 1677184255
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 1710728447
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 1744942079
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 1779840255
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 1815437567
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 1851745791
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 1888780543
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 1926555391
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 1965085695
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 2004387839
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 2044475647
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 2085363967
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 2127070719
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 2169611775
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 2213003775
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 2257263359
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 2302408191
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 2348456703
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 2395425023
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 2443333119
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 2492199935
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 2542043391
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 2592883455
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 2644741375
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 2697635583
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 2751262975
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 2805332991
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 2860484863
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 2916739071
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 2974118911
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 3032646143
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 3092344063
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 3153235711
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 3215345151
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 3278697215
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 3343316223
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 3409227263
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 3476456703
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 3545030655
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 3614976255
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 3686320639
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 3759091711
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 3833318655
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 3909030143
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 3986255359
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 4065025535
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 4145370879
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 4227323135
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 4310914559
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 4396177663
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 4483146239
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 4571854079
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 4662335999
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 4754627583
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 4848765183
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 4944785151
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 5042725631
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 5142625279
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 5244522495
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 5348457983
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 5454488319
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 5562622719
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 5672919807
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 5785422847
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 5900175615
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 6017223679
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 6136612863
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 6258389503
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 6382601983
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 6509298687
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 6638529023
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 6770344447
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 6904795903
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 7041936383
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 7181819647
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 7324500479
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 7470034943
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 7618480383
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 7769894399
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 7924336895
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 8081868287
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 8242566399
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 8406462207
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 8573635583
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 8744152319
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 8918079743
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 9095485951
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 9276439807
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 9461012479
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 9649277183
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 9841306879
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 10037177343
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 8
+  value_type_size: 2
+  is_segmented: true
+  scratch_size_recordings {
+    num_items: 31
+    scratch_space_bytes: 767
+  }
+  scratch_size_recordings {
+    num_items: 61
+    scratch_space_bytes: 1023
+  }
+  scratch_size_recordings {
+    num_items: 91
+    scratch_space_bytes: 1279
+  }
+  scratch_size_recordings {
+    num_items: 121
+    scratch_space_bytes: 1535
+  }
+  scratch_size_recordings {
+    num_items: 151
+    scratch_space_bytes: 2047
+  }
+  scratch_size_recordings {
+    num_items: 191
+    scratch_space_bytes: 2303
+  }
+  scratch_size_recordings {
+    num_items: 221
+    scratch_space_bytes: 2559
+  }
+  scratch_size_recordings {
+    num_items: 251
+    scratch_space_bytes: 2815
+  }
+  scratch_size_recordings {
+    num_items: 281
+    scratch_space_bytes: 3327
+  }
+  scratch_size_recordings {
+    num_items: 311
+    scratch_space_bytes: 3583
+  }
+  scratch_size_recordings {
+    num_items: 351
+    scratch_space_bytes: 3839
+  }
+  scratch_size_recordings {
+    num_items: 381
+    scratch_space_bytes: 4095
+  }
+  scratch_size_recordings {
+    num_items: 411
+    scratch_space_bytes: 4607
+  }
+  scratch_size_recordings {
+    num_items: 441
+    scratch_space_bytes: 4863
+  }
+  scratch_size_recordings {
+    num_items: 471
+    scratch_space_bytes: 5119
+  }
+  scratch_size_recordings {
+    num_items: 511
+    scratch_space_bytes: 5375
+  }
+  scratch_size_recordings {
+    num_items: 541
+    scratch_space_bytes: 5887
+  }
+  scratch_size_recordings {
+    num_items: 573
+    scratch_space_bytes: 6143
+  }
+  scratch_size_recordings {
+    num_items: 606
+    scratch_space_bytes: 6399
+  }
+  scratch_size_recordings {
+    num_items: 630
+    scratch_space_bytes: 6655
+  }
+  scratch_size_recordings {
+    num_items: 667
+    scratch_space_bytes: 7167
+  }
+  scratch_size_recordings {
+    num_items: 693
+    scratch_space_bytes: 7423
+  }
+  scratch_size_recordings {
+    num_items: 734
+    scratch_space_bytes: 7679
+  }
+  scratch_size_recordings {
+    num_items: 762
+    scratch_space_bytes: 7935
+  }
+  scratch_size_recordings {
+    num_items: 792
+    scratch_space_bytes: 8447
+  }
+  scratch_size_recordings {
+    num_items: 823
+    scratch_space_bytes: 8703
+  }
+  scratch_size_recordings {
+    num_items: 855
+    scratch_space_bytes: 8959
+  }
+  scratch_size_recordings {
+    num_items: 889
+    scratch_space_bytes: 9215
+  }
+  scratch_size_recordings {
+    num_items: 924
+    scratch_space_bytes: 9727
+  }
+  scratch_size_recordings {
+    num_items: 960
+    scratch_space_bytes: 9983
+  }
+  scratch_size_recordings {
+    num_items: 979
+    scratch_space_bytes: 10239
+  }
+  scratch_size_recordings {
+    num_items: 1017
+    scratch_space_bytes: 10495
+  }
+  scratch_size_recordings {
+    num_items: 1037
+    scratch_space_bytes: 11007
+  }
+  scratch_size_recordings {
+    num_items: 1078
+    scratch_space_bytes: 11263
+  }
+  scratch_size_recordings {
+    num_items: 1120
+    scratch_space_bytes: 11519
+  }
+  scratch_size_recordings {
+    num_items: 1142
+    scratch_space_bytes: 11775
+  }
+  scratch_size_recordings {
+    num_items: 1164
+    scratch_space_bytes: 12287
+  }
+  scratch_size_recordings {
+    num_items: 1210
+    scratch_space_bytes: 12543
+  }
+  scratch_size_recordings {
+    num_items: 1234
+    scratch_space_bytes: 12799
+  }
+  scratch_size_recordings {
+    num_items: 1258
+    scratch_space_bytes: 13055
+  }
+  scratch_size_recordings {
+    num_items: 1308
+    scratch_space_bytes: 13567
+  }
+  scratch_size_recordings {
+    num_items: 1334
+    scratch_space_bytes: 13823
+  }
+  scratch_size_recordings {
+    num_items: 1360
+    scratch_space_bytes: 14079
+  }
+  scratch_size_recordings {
+    num_items: 1387
+    scratch_space_bytes: 14335
+  }
+  scratch_size_recordings {
+    num_items: 1414
+    scratch_space_bytes: 14847
+  }
+  scratch_size_recordings {
+    num_items: 1470
+    scratch_space_bytes: 15103
+  }
+  scratch_size_recordings {
+    num_items: 1499
+    scratch_space_bytes: 15359
+  }
+  scratch_size_recordings {
+    num_items: 1528
+    scratch_space_bytes: 15615
+  }
+  scratch_size_recordings {
+    num_items: 1558
+    scratch_space_bytes: 16127
+  }
+  scratch_size_recordings {
+    num_items: 1589
+    scratch_space_bytes: 16383
+  }
+  scratch_size_recordings {
+    num_items: 1620
+    scratch_space_bytes: 16639
+  }
+  scratch_size_recordings {
+    num_items: 1652
+    scratch_space_bytes: 16895
+  }
+  scratch_size_recordings {
+    num_items: 1685
+    scratch_space_bytes: 17407
+  }
+  scratch_size_recordings {
+    num_items: 1718
+    scratch_space_bytes: 17663
+  }
+  scratch_size_recordings {
+    num_items: 1752
+    scratch_space_bytes: 17919
+  }
+  scratch_size_recordings {
+    num_items: 1787
+    scratch_space_bytes: 18175
+  }
+  scratch_size_recordings {
+    num_items: 1822
+    scratch_space_bytes: 18687
+  }
+  scratch_size_recordings {
+    num_items: 1858
+    scratch_space_bytes: 19199
+  }
+  scratch_size_recordings {
+    num_items: 1895
+    scratch_space_bytes: 19455
+  }
+  scratch_size_recordings {
+    num_items: 1932
+    scratch_space_bytes: 19967
+  }
+  scratch_size_recordings {
+    num_items: 1970
+    scratch_space_bytes: 20223
+  }
+  scratch_size_recordings {
+    num_items: 2009
+    scratch_space_bytes: 20479
+  }
+  scratch_size_recordings {
+    num_items: 2049
+    scratch_space_bytes: 21247
+  }
+  scratch_size_recordings {
+    num_items: 2089
+    scratch_space_bytes: 21503
+  }
+  scratch_size_recordings {
+    num_items: 2130
+    scratch_space_bytes: 21759
+  }
+  scratch_size_recordings {
+    num_items: 2172
+    scratch_space_bytes: 22015
+  }
+  scratch_size_recordings {
+    num_items: 2215
+    scratch_space_bytes: 22783
+  }
+  scratch_size_recordings {
+    num_items: 2259
+    scratch_space_bytes: 23039
+  }
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 23295
+  }
+  scratch_size_recordings {
+    num_items: 2350
+    scratch_space_bytes: 24063
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 24319
+  }
+  scratch_size_recordings {
+    num_items: 2444
+    scratch_space_bytes: 25087
+  }
+  scratch_size_recordings {
+    num_items: 2492
+    scratch_space_bytes: 25343
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 25855
+  }
+  scratch_size_recordings {
+    num_items: 2591
+    scratch_space_bytes: 26367
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 26879
+  }
+  scratch_size_recordings {
+    num_items: 2694
+    scratch_space_bytes: 27647
+  }
+  scratch_size_recordings {
+    num_items: 2747
+    scratch_space_bytes: 27903
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 28415
+  }
+  scratch_size_recordings {
+    num_items: 2857
+    scratch_space_bytes: 29183
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 29695
+  }
+  scratch_size_recordings {
+    num_items: 2972
+    scratch_space_bytes: 30207
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 30719
+  }
+  scratch_size_recordings {
+    num_items: 3091
+    scratch_space_bytes: 31487
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 31999
+  }
+  scratch_size_recordings {
+    num_items: 3215
+    scratch_space_bytes: 32767
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 33279
+  }
+  scratch_size_recordings {
+    num_items: 3344
+    scratch_space_bytes: 34047
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 34559
+  }
+  scratch_size_recordings {
+    num_items: 3478
+    scratch_space_bytes: 35327
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 35839
+  }
+  scratch_size_recordings {
+    num_items: 3617
+    scratch_space_bytes: 36863
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 37375
+  }
+  scratch_size_recordings {
+    num_items: 3762
+    scratch_space_bytes: 38143
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 38655
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 39679
+  }
+  scratch_size_recordings {
+    num_items: 3991
+    scratch_space_bytes: 40447
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 41215
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 41983
+  }
+  scratch_size_recordings {
+    num_items: 4234
+    scratch_space_bytes: 43007
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 43519
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 44543
+  }
+  scratch_size_recordings {
+    num_items: 4492
+    scratch_space_bytes: 45567
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 46335
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 47103
+  }
+  scratch_size_recordings {
+    num_items: 4765
+    scratch_space_bytes: 48127
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 48895
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 49919
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 50943
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 52223
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 53247
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 54015
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 55039
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 56319
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 57343
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 58623
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 59903
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 60927
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 62207
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 63487
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 64767
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 66047
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 67327
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 68607
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 69887
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 71167
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 72703
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 73983
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 75519
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 77055
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 78335
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 80127
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 81663
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 83199
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 84735
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 86783
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 88319
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 90367
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 91903
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 93695
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 95743
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 97535
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 99583
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 101375
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 103423
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 105727
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 107519
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 109823
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 112127
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 113919
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 116479
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 118783
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 121087
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 123647
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 125695
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 128255
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 130815
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 133375
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 136447
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 139007
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 141823
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 144383
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 147199
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 150527
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 153343
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 156415
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 159487
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 162559
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 165887
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 169215
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 172543
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 176127
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 179455
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 183039
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 186623
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 190463
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 194303
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 198143
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 201983
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 206079
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 210175
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 214527
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 218623
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 222975
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 227583
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 231935
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 236543
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 241407
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 246015
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 250879
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 255999
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 261119
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 266239
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 271615
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 277247
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 282623
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 288255
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 294143
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 299775
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 305919
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 312063
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 318207
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 324607
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 331007
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 337663
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 344319
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 350975
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 358143
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 365567
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 372735
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 380159
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 387583
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 395519
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 403199
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 411135
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 419583
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 427775
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 436479
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 445183
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 454143
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 463103
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 472319
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 482047
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 491519
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 501247
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 511487
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 521727
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 531967
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 542463
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 553215
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 564479
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 575743
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 587263
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 599039
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 610815
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 623103
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 635647
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 648447
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 661247
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 674303
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 688127
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 701695
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 715519
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 729855
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 744703
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 759295
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 774399
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 790015
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 805887
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 822015
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 838399
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 855039
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 872447
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 889599
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 907519
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 925439
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 944127
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 962815
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 982015
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 1001983
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 1021695
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 1042175
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 1063167
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 1084159
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 1105919
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 1127935
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 1150719
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 1173759
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 1197055
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 1221119
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 1245439
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 1270527
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 1295615
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 1321727
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 1348095
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 1374975
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 1402623
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 1430527
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 1459199
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 1488383
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 1518079
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 1548543
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 1579263
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 1611007
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 1643263
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 1675775
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 1709567
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 1743615
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 1778687
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 1814015
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 1850367
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 1887487
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 1925119
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 1963519
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 2002687
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 2042879
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 2083583
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 2125055
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 2167807
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 2211327
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 2255359
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 2300415
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 2346239
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 2393343
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 2440959
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 2489855
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 2539775
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 2590463
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 2642175
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 2695167
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 2749183
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 2803967
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 2860287
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 2917119
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 2975743
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 3035135
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 3095807
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 3157759
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 3220735
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 3285247
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 3351039
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 3417855
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 3486207
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 3555839
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 3627007
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 3699455
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 3773439
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 3848959
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 3926015
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 4004607
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 4084479
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 4166399
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 4249599
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 4334335
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 4421119
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 4509695
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 4599807
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 4691967
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 4785663
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 4881407
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 4978943
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 5078527
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 5180159
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 5283583
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 5389567
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 5497087
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 5607167
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 5719039
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 5833727
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 5950207
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 6069247
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 6190335
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 6314239
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 6440703
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 6569215
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 6700799
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 6834943
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 6971647
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 7110655
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 7253247
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 7398143
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 7545855
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 7696895
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 7851007
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 8007935
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 8167935
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 8331263
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 8497919
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 8667903
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 8841215
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 9017855
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 9198335
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 9382399
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 9570047
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 9761279
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 9956607
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 10155775
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 10358783
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 10566143
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 10777343
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 10992639
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 11212543
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 11436799
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 11665663
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 11898879
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 12136959
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 12379647
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 12627199
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 12879615
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 13137407
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 13400063
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 13667839
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 13941503
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 14220287
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 14504703
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 14794495
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 15090687
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 15392255
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 15700223
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 16014079
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 16334335
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 16661247
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 16994303
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 17334015
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 17680895
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 18034687
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 18395135
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 18763007
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 19138303
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 19521023
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 19911423
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 20309503
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 20715775
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 21129983
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 21552639
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 21983743
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 22423295
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 22871807
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 23329279
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 23795967
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 24271615
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 24757247
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 25252095
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 25757183
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 26272255
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 26797823
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 27333887
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 27880447
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 28438015
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 29006847
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 29586943
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 30178559
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 30782207
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 31397887
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 32025599
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 32666367
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 33319679
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 33986047
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 34665727
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 35358975
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 36066047
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 36787455
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 37523199
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 38273535
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 39038975
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 39819775
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 40616447
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 41428479
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 42256895
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 43102207
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 43964159
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 44843519
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 45740543
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 46654975
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 47588095
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 48539903
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 49510655
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 50500863
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 51511039
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 52541183
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 53592063
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 54663679
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 55757055
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 56872447
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 58009599
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 59170047
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 60353279
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 61560319
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 62791423
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 64047359
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 65328127
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 66634495
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 67967487
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 69326847
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 70713343
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 72127487
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 73570047
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 75041279
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 76542207
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 78073087
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 79634687
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 81227263
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 82851839
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 84508927
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 86198783
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 87922943
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 89681407
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 91474943
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 93304319
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 95170559
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 97073919
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 99015423
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 100995583
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 103015423
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 105075967
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 107177215
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 109320703
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 111507199
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 113737215
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 116012287
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 118332159
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 120698879
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 123112959
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 125575167
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 128086783
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 130648319
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 133261055
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 135926527
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 138645247
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 141417983
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 144246015
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 147131135
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 150073599
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 153075199
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 156136703
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 159259647
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 162444543
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 165693439
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 169007359
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 172387583
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 175835135
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 179352063
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 182938879
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 186597887
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 190329599
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 194136319
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 198018815
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 201979135
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 206018815
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 210139135
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 214342143
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 218628863
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 223001343
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 227461375
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 232010495
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 236651007
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 241383679
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 246211583
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 251135743
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 256158463
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 261281535
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 266507263
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 271837183
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 277273855
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 282819583
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 288475903
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 294245119
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 300130303
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 306132735
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 312255487
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 318500607
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 324870399
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 331367935
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 337995263
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 344755199
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 351650303
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 358683135
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 365857023
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 373174015
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 380637439
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 388250367
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 396015103
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 403935487
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 412014079
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 420254463
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 428659455
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 437232639
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 445977343
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 454896639
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 463994623
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 473274623
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 482740223
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 492395007
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 502242559
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 512287743
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 522533375
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 532984063
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 543643647
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 554516479
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 565606655
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 576919039
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 588457215
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 600226303
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 612231167
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 624475647
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 636965119
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 649704447
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 662698239
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 675952383
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 689471487
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 703260927
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 717326079
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 731672319
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 746306047
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 761232127
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 776456703
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 791985663
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 807825407
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 823981823
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 840461567
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 857270783
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 874416127
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 891904255
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 909742335
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 927937279
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 946495999
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 965425919
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 984734463
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 1004429055
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 1024517887
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 1045008127
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 1065908223
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 1087226623
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 1108971007
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 1131150335
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 1153773567
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 1176848639
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 1200385535
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 1224393215
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 1248881407
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 1273858815
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 1299335935
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 1325323007
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 1351829247
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 1378865919
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 1406443263
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 1434571775
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 1463263487
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 1492528639
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 1522379263
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 1552826623
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 1583883263
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 1615560959
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 1647871999
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 1680829439
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 1714446079
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 1748734975
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 1783709695
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 1819384063
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 1855771647
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 1892887039
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 1930744575
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 1969359615
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 2008747007
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 2048921599
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 2089900287
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 2131697919
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 2174332159
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 2217818623
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 2262174975
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 2307418623
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 2353566719
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 2400638207
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 2448651007
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 2497624063
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 2547576319
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 2598527999
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 2650498559
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 2703508479
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 2757578495
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 2812730367
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 2868984575
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 2926364415
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 2984891647
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 3044589567
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 3105481215
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 3167590655
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 3230942719
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 3295561727
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 3361472767
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 3428702207
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 3497276159
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 3567221759
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 3638566143
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 3711337215
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 3785564159
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 3861275647
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 3938500863
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 4017271039
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 4097616383
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 4179568639
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 4263160063
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 4348423167
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 4435391743
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 4524099583
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 4614581503
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 4706873087
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 4801010687
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 4897030655
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 4994971135
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 5094870783
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 5196767999
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 5300703487
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 5406717439
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 5514851839
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 5625148927
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 5737651967
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 5852404735
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 5969452799
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 6088841983
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 6210618623
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 6334831103
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 6461527807
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 6590758143
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 6722573567
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 6857025023
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 6994165503
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 7134048767
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 7276729599
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 7422264063
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 7570709503
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 7722123519
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 7876566015
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 8034097407
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 8194779135
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 8358674943
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 8525848319
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 8696365055
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 8870292479
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 9047698687
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 9228652543
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 9413225215
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 9601489919
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 9793519615
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 9989390079
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 8
+  value_type_size: 4
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 1
+  }
+  scratch_size_recordings {
+    num_items: 2350
+    scratch_space_bytes: 46335
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 46847
+  }
+  scratch_size_recordings {
+    num_items: 2444
+    scratch_space_bytes: 47615
+  }
+  scratch_size_recordings {
+    num_items: 2492
+    scratch_space_bytes: 47871
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 48639
+  }
+  scratch_size_recordings {
+    num_items: 2591
+    scratch_space_bytes: 49151
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 49919
+  }
+  scratch_size_recordings {
+    num_items: 2694
+    scratch_space_bytes: 50687
+  }
+  scratch_size_recordings {
+    num_items: 2747
+    scratch_space_bytes: 50943
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 51711
+  }
+  scratch_size_recordings {
+    num_items: 2857
+    scratch_space_bytes: 52479
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 53247
+  }
+  scratch_size_recordings {
+    num_items: 2972
+    scratch_space_bytes: 53759
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 54527
+  }
+  scratch_size_recordings {
+    num_items: 3091
+    scratch_space_bytes: 55295
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 56063
+  }
+  scratch_size_recordings {
+    num_items: 3215
+    scratch_space_bytes: 56831
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 57599
+  }
+  scratch_size_recordings {
+    num_items: 3344
+    scratch_space_bytes: 58367
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 59135
+  }
+  scratch_size_recordings {
+    num_items: 3478
+    scratch_space_bytes: 59903
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 60671
+  }
+  scratch_size_recordings {
+    num_items: 3617
+    scratch_space_bytes: 61695
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 62463
+  }
+  scratch_size_recordings {
+    num_items: 3762
+    scratch_space_bytes: 63231
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 63999
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 65279
+  }
+  scratch_size_recordings {
+    num_items: 3991
+    scratch_space_bytes: 66047
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 67071
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 67839
+  }
+  scratch_size_recordings {
+    num_items: 4234
+    scratch_space_bytes: 69119
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 69887
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 70911
+  }
+  scratch_size_recordings {
+    num_items: 4492
+    scratch_space_bytes: 72191
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 73215
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 73983
+  }
+  scratch_size_recordings {
+    num_items: 4765
+    scratch_space_bytes: 75263
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 76287
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 77567
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 78591
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 80127
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 81407
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 82431
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 83711
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 85247
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 86271
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 88831
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 90367
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 91647
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 93183
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 94719
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 96255
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 97791
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 99327
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 100863
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 102399
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 103935
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 105727
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 107263
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 109311
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 111103
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 112639
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 114687
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 116479
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 118527
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 120319
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 122623
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 124671
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 126975
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 128767
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 131071
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 133375
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 135679
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 137983
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 140287
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 142591
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 145407
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 147711
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 150271
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 153087
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 155391
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 159487
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 162047
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 164863
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 167935
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 170495
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 173567
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 176639
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 179711
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 183295
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 186367
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 189695
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 192767
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 196351
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 200191
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 203519
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 207359
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 210943
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 214783
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 218623
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 222719
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 226559
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 231935
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 236031
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 240383
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 244479
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 249087
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 253695
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 258303
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 262911
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 268031
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 272895
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 278015
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 282879
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 288255
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 293631
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 300031
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 305407
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 311295
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 316927
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 322815
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 328959
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 335103
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 341247
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 347647
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 354303
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 360703
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 367615
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 375551
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 382463
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 389887
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 397055
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 404479
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 412159
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 419839
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 427775
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 435967
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 444927
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 453375
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 462335
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 471039
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 479999
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 488703
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 498431
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 507647
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 518143
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 528127
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 538111
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 548607
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 558847
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 569599
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 581375
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 592639
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 604159
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 615679
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 627199
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 639487
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 652799
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 665087
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 677631
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 690687
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 704255
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 717567
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 732415
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 746751
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 760831
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 775423
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 791551
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 806911
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 822271
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 837887
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 854527
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 871935
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 888575
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 905727
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 923391
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 942079
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 960255
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 978943
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 997887
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 1018367
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 1038079
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 1058047
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 1079807
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 1100543
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 1122047
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 1144575
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 1166847
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 1189375
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 1213439
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 1237247
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 1261055
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 1286655
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 1311743
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 1337087
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 1364223
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 1390591
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 1417983
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 1446655
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 1474559
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 1504511
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 1533695
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 1564671
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 1594879
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 1626111
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 1658879
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 1691135
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 1725183
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 1758719
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 1794303
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 1829119
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 1865983
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 1902335
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 1940223
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 1978367
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 2018047
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 2058239
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 2098687
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 2140671
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 2182655
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 2226175
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 2270719
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 2315263
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 2361599
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 2408703
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 2455551
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 2504959
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 2554623
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 2604543
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 2656767
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 2710015
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 2764031
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 2818047
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 2874111
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 2931455
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 2989823
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 3049471
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 3110399
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 3171071
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 3234303
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 3298815
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 3364607
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 3431423
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 3500031
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 3569407
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 3640575
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 3713023
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 3786751
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 3862271
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 3938815
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 4018175
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 4098303
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 4179455
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 4262399
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 4347135
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 4433407
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 4522495
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 4612351
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 4703999
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 4798463
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 4893695
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 4990719
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 5091071
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 5191935
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 5295615
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 5400831
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 5509119
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 5618175
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 5730815
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 5845247
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 5961215
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 6080255
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 6201855
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 6324991
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 6450943
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 6580223
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 6711295
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 6845439
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 6981887
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 7121407
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 7263231
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 7408127
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 7555583
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 7706367
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 7859967
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 8017407
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 8177407
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 8340223
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 8507391
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 8676351
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 8849407
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 9026303
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 9205759
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 9390079
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 9577983
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 9769471
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 9963519
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 10162431
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 10365695
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 10572543
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 10783743
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 10998783
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 11218431
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 11442431
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 11671551
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 11904255
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 12141567
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 12384767
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 12631295
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 12884223
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 13141759
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 13403391
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 13671167
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 13944575
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 14223103
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 14507263
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 14797055
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 15092223
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 15394559
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 15701503
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 16015871
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 16335103
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 16661759
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 16995071
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 17333759
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 17680127
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 18033151
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 18393599
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 18761727
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 19136255
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 19519231
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 19908607
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 20306943
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 20712703
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 21125631
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 21548031
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 21979647
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 22418431
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 22865919
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 23323391
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 23788799
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 24264447
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 24749311
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 25244159
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 25748223
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 26263807
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 26788095
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 27323903
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 27870207
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 28427519
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 28995583
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 29574655
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 30165503
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 30768639
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 31383807
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 32011263
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 32651007
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 33303551
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 33969663
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 34648063
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 35340799
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 36047103
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 36767487
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 37502975
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 38253055
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 39016703
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 39797759
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 40593151
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 41404927
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 42231807
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 43076095
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 43937791
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 44815871
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 45712127
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 46625791
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 47557631
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 48508927
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 49479167
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 50468095
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 51476479
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 52506111
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 53555967
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 54626559
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 55719167
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 56831999
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 57969151
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 59127551
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 60310015
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 61515519
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 62746367
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 64000767
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 65280255
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 66585855
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 67916287
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 69275135
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 70660351
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 72073215
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 73513727
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 74983679
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 76483327
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 78012159
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 79572479
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 81162751
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 82785791
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 84441599
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 86129919
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 87852543
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 89609215
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 91400447
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 93228031
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 95092735
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 96994303
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 98934015
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 100912383
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 102930431
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 104987903
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 107087615
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 109228543
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 111413503
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 113640703
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 115913983
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 118231807
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 120596223
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 123007231
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 125466879
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 127976447
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 130534911
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 133144831
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 135807743
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 138523647
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 141294591
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 144119551
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 147001599
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 149941503
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 152940031
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 155997951
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 159118079
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 162299391
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 165544959
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 168856575
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 172232959
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 175676927
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 179190271
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 182773503
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 186428671
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 190156543
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 193959423
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 197838591
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 201794815
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 205830911
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 209946367
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 214145023
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 218427647
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 222796543
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 227251711
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 231796479
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 236432127
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 241159679
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 245982719
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 250902015
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 255920127
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 261038335
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 266258175
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 271582975
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 277014015
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 282553855
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 288205311
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 293968639
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 299847935
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 305843967
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 311961343
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 318200319
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 324563199
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 331054335
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 337674751
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 344428799
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 351316735
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 358342399
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 365508863
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 372818431
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 380275199
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 387879935
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 395637503
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 403549183
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 411619583
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 419852543
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 428249343
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 436813823
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 445549311
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 454459903
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 463549183
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 472819711
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 482275583
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 491920383
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 501759231
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 511793919
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 522028543
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 532468991
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 543118079
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 553980159
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 565059327
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 576360191
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 587886847
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 599644671
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 611636991
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 623869183
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 636346367
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 649073407
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 662054399
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 675295487
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 688800255
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 702575871
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 716627455
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 730959103
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 745579007
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 760489471
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 775699199
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 791212543
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 807036671
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 823176959
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 839639551
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 856432639
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 873560831
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 891031295
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 908851711
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 927028223
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 945569023
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 964479231
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 983768575
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 1003443967
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 1023512575
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 1043982591
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 1064861183
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 1086158335
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 1107881471
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 1130039039
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 1152639231
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 1175691007
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 1199205119
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 1223187967
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 1247651839
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 1272604671
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 1298056447
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 1324017407
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 1350496767
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 1377507327
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 1405056511
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 1433157375
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 1461820415
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 1491055615
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 1520876543
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 1551293951
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 1582319103
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 1613966335
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 1646244863
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 1679169279
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 1712752383
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 1747006463
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 1781946367
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 1817584895
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 1853936895
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 1891014399
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 1928834303
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 1967410943
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 2006758655
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 2046894335
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 2087831295
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 2129587199
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 2172179455
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 2215622399
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 2259934719
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 2305132799
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 2351234559
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 2398259711
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 2446224895
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 2495148287
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 2545050623
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 2595951359
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 2647870207
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 2700826879
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 2754843135
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 2809940223
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 2866138367
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 2923460607
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 2981929983
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 3041567999
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 3102398719
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 3164446975
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 3227735295
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 3291964671
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 3356848639
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 3423030783
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 3490535935
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 3559391743
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 3629624319
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 3701261823
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 3774331903
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 3848863231
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 3924885759
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 4002428415
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 4081521663
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 4162196991
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 4244485887
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 4328420607
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 4414033663
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 4501359103
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 4590431487
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 4681285119
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 4773955327
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 4868479743
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 4964893951
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 5063236863
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 5163546367
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 5265862143
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 5370224383
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 5476673791
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 5585252095
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 5696002047
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 5808967167
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 5924191231
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 6041719807
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 6161599231
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 6283876095
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 6408598527
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 6535831807
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 6665593087
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 6797949439
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 6932953087
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 7070656511
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 7211114239
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 7354381055
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 7500513023
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 7649567999
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 7801604095
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 7956680447
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 8114859007
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 8276200703
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 8440769279
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 8608629247
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 8779846399
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 8954487551
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 9132622079
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 9314319103
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 9499650047
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 9688687615
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 9881522175
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 10078196991
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 10278805247
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 10483425279
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 10692138239
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 10905025535
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 11122170111
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 11343657471
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 11569575167
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 11800010751
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 12035055359
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 8
+  value_type_size: 4
+  is_segmented: true
+  scratch_size_recordings {
+    num_items: 31
+    scratch_space_bytes: 767
+  }
+  scratch_size_recordings {
+    num_items: 61
+    scratch_space_bytes: 1023
+  }
+  scratch_size_recordings {
+    num_items: 91
+    scratch_space_bytes: 1535
+  }
+  scratch_size_recordings {
+    num_items: 121
+    scratch_space_bytes: 1791
+  }
+  scratch_size_recordings {
+    num_items: 151
+    scratch_space_bytes: 2303
+  }
+  scratch_size_recordings {
+    num_items: 191
+    scratch_space_bytes: 2559
+  }
+  scratch_size_recordings {
+    num_items: 221
+    scratch_space_bytes: 3071
+  }
+  scratch_size_recordings {
+    num_items: 251
+    scratch_space_bytes: 3327
+  }
+  scratch_size_recordings {
+    num_items: 281
+    scratch_space_bytes: 3839
+  }
+  scratch_size_recordings {
+    num_items: 311
+    scratch_space_bytes: 4095
+  }
+  scratch_size_recordings {
+    num_items: 351
+    scratch_space_bytes: 4607
+  }
+  scratch_size_recordings {
+    num_items: 381
+    scratch_space_bytes: 4863
+  }
+  scratch_size_recordings {
+    num_items: 411
+    scratch_space_bytes: 5375
+  }
+  scratch_size_recordings {
+    num_items: 441
+    scratch_space_bytes: 5631
+  }
+  scratch_size_recordings {
+    num_items: 471
+    scratch_space_bytes: 6143
+  }
+  scratch_size_recordings {
+    num_items: 511
+    scratch_space_bytes: 6399
+  }
+  scratch_size_recordings {
+    num_items: 541
+    scratch_space_bytes: 6911
+  }
+  scratch_size_recordings {
+    num_items: 573
+    scratch_space_bytes: 7167
+  }
+  scratch_size_recordings {
+    num_items: 606
+    scratch_space_bytes: 7679
+  }
+  scratch_size_recordings {
+    num_items: 630
+    scratch_space_bytes: 7935
+  }
+  scratch_size_recordings {
+    num_items: 667
+    scratch_space_bytes: 8447
+  }
+  scratch_size_recordings {
+    num_items: 693
+    scratch_space_bytes: 8703
+  }
+  scratch_size_recordings {
+    num_items: 734
+    scratch_space_bytes: 9215
+  }
+  scratch_size_recordings {
+    num_items: 762
+    scratch_space_bytes: 9471
+  }
+  scratch_size_recordings {
+    num_items: 792
+    scratch_space_bytes: 9983
+  }
+  scratch_size_recordings {
+    num_items: 823
+    scratch_space_bytes: 10239
+  }
+  scratch_size_recordings {
+    num_items: 855
+    scratch_space_bytes: 10751
+  }
+  scratch_size_recordings {
+    num_items: 889
+    scratch_space_bytes: 11007
+  }
+  scratch_size_recordings {
+    num_items: 924
+    scratch_space_bytes: 11519
+  }
+  scratch_size_recordings {
+    num_items: 960
+    scratch_space_bytes: 11775
+  }
+  scratch_size_recordings {
+    num_items: 979
+    scratch_space_bytes: 12287
+  }
+  scratch_size_recordings {
+    num_items: 1017
+    scratch_space_bytes: 12543
+  }
+  scratch_size_recordings {
+    num_items: 1037
+    scratch_space_bytes: 13055
+  }
+  scratch_size_recordings {
+    num_items: 1078
+    scratch_space_bytes: 13311
+  }
+  scratch_size_recordings {
+    num_items: 1120
+    scratch_space_bytes: 13823
+  }
+  scratch_size_recordings {
+    num_items: 1142
+    scratch_space_bytes: 14079
+  }
+  scratch_size_recordings {
+    num_items: 1164
+    scratch_space_bytes: 14591
+  }
+  scratch_size_recordings {
+    num_items: 1210
+    scratch_space_bytes: 14847
+  }
+  scratch_size_recordings {
+    num_items: 1234
+    scratch_space_bytes: 15359
+  }
+  scratch_size_recordings {
+    num_items: 1258
+    scratch_space_bytes: 15615
+  }
+  scratch_size_recordings {
+    num_items: 1308
+    scratch_space_bytes: 16127
+  }
+  scratch_size_recordings {
+    num_items: 1334
+    scratch_space_bytes: 16383
+  }
+  scratch_size_recordings {
+    num_items: 1360
+    scratch_space_bytes: 16895
+  }
+  scratch_size_recordings {
+    num_items: 1387
+    scratch_space_bytes: 17151
+  }
+  scratch_size_recordings {
+    num_items: 1414
+    scratch_space_bytes: 17663
+  }
+  scratch_size_recordings {
+    num_items: 1470
+    scratch_space_bytes: 17919
+  }
+  scratch_size_recordings {
+    num_items: 1499
+    scratch_space_bytes: 18431
+  }
+  scratch_size_recordings {
+    num_items: 1528
+    scratch_space_bytes: 18687
+  }
+  scratch_size_recordings {
+    num_items: 1558
+    scratch_space_bytes: 19199
+  }
+  scratch_size_recordings {
+    num_items: 1589
+    scratch_space_bytes: 19455
+  }
+  scratch_size_recordings {
+    num_items: 1620
+    scratch_space_bytes: 19967
+  }
+  scratch_size_recordings {
+    num_items: 1652
+    scratch_space_bytes: 20223
+  }
+  scratch_size_recordings {
+    num_items: 1685
+    scratch_space_bytes: 20735
+  }
+  scratch_size_recordings {
+    num_items: 1718
+    scratch_space_bytes: 20991
+  }
+  scratch_size_recordings {
+    num_items: 1752
+    scratch_space_bytes: 21503
+  }
+  scratch_size_recordings {
+    num_items: 1787
+    scratch_space_bytes: 21759
+  }
+  scratch_size_recordings {
+    num_items: 1822
+    scratch_space_bytes: 22271
+  }
+  scratch_size_recordings {
+    num_items: 1858
+    scratch_space_bytes: 23039
+  }
+  scratch_size_recordings {
+    num_items: 1895
+    scratch_space_bytes: 23295
+  }
+  scratch_size_recordings {
+    num_items: 1932
+    scratch_space_bytes: 23807
+  }
+  scratch_size_recordings {
+    num_items: 1970
+    scratch_space_bytes: 24063
+  }
+  scratch_size_recordings {
+    num_items: 2009
+    scratch_space_bytes: 24575
+  }
+  scratch_size_recordings {
+    num_items: 2049
+    scratch_space_bytes: 25343
+  }
+  scratch_size_recordings {
+    num_items: 2089
+    scratch_space_bytes: 25599
+  }
+  scratch_size_recordings {
+    num_items: 2130
+    scratch_space_bytes: 26111
+  }
+  scratch_size_recordings {
+    num_items: 2172
+    scratch_space_bytes: 26367
+  }
+  scratch_size_recordings {
+    num_items: 2215
+    scratch_space_bytes: 27135
+  }
+  scratch_size_recordings {
+    num_items: 2259
+    scratch_space_bytes: 27647
+  }
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 27903
+  }
+  scratch_size_recordings {
+    num_items: 2350
+    scratch_space_bytes: 28671
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 29183
+  }
+  scratch_size_recordings {
+    num_items: 2444
+    scratch_space_bytes: 29951
+  }
+  scratch_size_recordings {
+    num_items: 2492
+    scratch_space_bytes: 30207
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 30975
+  }
+  scratch_size_recordings {
+    num_items: 2591
+    scratch_space_bytes: 31487
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 32255
+  }
+  scratch_size_recordings {
+    num_items: 2694
+    scratch_space_bytes: 33023
+  }
+  scratch_size_recordings {
+    num_items: 2747
+    scratch_space_bytes: 33279
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 34047
+  }
+  scratch_size_recordings {
+    num_items: 2857
+    scratch_space_bytes: 34815
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 35583
+  }
+  scratch_size_recordings {
+    num_items: 2972
+    scratch_space_bytes: 36095
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 36863
+  }
+  scratch_size_recordings {
+    num_items: 3091
+    scratch_space_bytes: 37631
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 38399
+  }
+  scratch_size_recordings {
+    num_items: 3215
+    scratch_space_bytes: 39167
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 39935
+  }
+  scratch_size_recordings {
+    num_items: 3344
+    scratch_space_bytes: 40703
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 41471
+  }
+  scratch_size_recordings {
+    num_items: 3478
+    scratch_space_bytes: 42239
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 43007
+  }
+  scratch_size_recordings {
+    num_items: 3617
+    scratch_space_bytes: 44031
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 44799
+  }
+  scratch_size_recordings {
+    num_items: 3762
+    scratch_space_bytes: 45567
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 46335
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 47615
+  }
+  scratch_size_recordings {
+    num_items: 3991
+    scratch_space_bytes: 48383
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 49407
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 50175
+  }
+  scratch_size_recordings {
+    num_items: 4234
+    scratch_space_bytes: 51455
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 52223
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 53247
+  }
+  scratch_size_recordings {
+    num_items: 4492
+    scratch_space_bytes: 54527
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 55551
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 56319
+  }
+  scratch_size_recordings {
+    num_items: 4765
+    scratch_space_bytes: 57599
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 58623
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 59903
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 60927
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 62463
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 63743
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 64767
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 66047
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 67583
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 68607
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 70143
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 71679
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 72959
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 74495
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 76031
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 77567
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 79103
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 80639
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 82175
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 83711
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 85247
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 87039
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 88575
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 90623
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 92415
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 93951
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 95999
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 97791
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 99839
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 101631
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 103935
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 105983
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 108287
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 110079
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 112383
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 114687
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 116991
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 119295
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 121599
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 123903
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 126719
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 129023
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 131583
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 134399
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 136703
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 139775
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 142335
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 145151
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 148223
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 150783
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 153855
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 156927
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 159999
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 163583
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 166655
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 169983
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 173055
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 176639
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 180479
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 183807
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 187647
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 191231
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 195071
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 198911
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 203007
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 206847
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 211199
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 215295
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 219647
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 223743
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 228351
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 232959
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 237567
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 242175
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 247295
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 252159
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 257279
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 262143
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 267519
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 272895
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 278271
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 283647
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 289535
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 295167
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 301055
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 307199
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 313343
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 319487
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 325887
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 332543
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 338943
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 345855
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 352767
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 359679
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 367103
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 374271
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 381695
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 389375
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 397055
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 404991
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 413183
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 421119
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 429567
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 438527
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 447231
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 456191
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 464895
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 474623
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 483839
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 493311
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 503295
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 513279
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 523775
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 534015
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 544767
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 555519
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 566783
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 578303
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 589823
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 601343
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 613631
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 625919
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 638207
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 650751
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 663807
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 677375
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 690687
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 704511
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 718847
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 732927
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 747519
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 762623
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 777983
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 793343
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 808959
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 825599
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 841983
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 858623
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 875775
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 893439
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 911103
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 929279
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 947967
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 966911
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 986367
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 1006079
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 1026047
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 1046783
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 1067519
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 1089023
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 1110527
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 1132799
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 1155327
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 1178367
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 1202175
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 1225983
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 1250559
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 1275647
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 1300991
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 1327103
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 1353471
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 1380863
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 1408511
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 1436415
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 1465343
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 1494527
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 1524479
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 1554687
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 1585919
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 1617663
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 1649919
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 1682943
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 1716479
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 1751039
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 1785855
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 1821695
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 1858047
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 1894911
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 1933055
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 1971711
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 2010879
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 2051327
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 2092287
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 2134271
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 2176767
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 2220287
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 2264831
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 2310143
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 2356223
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 2403071
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 2451455
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 2500095
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 2550015
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 2601215
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 2653439
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 2706431
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 2760447
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 2815487
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 2871807
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 2929151
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 2987775
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 3047679
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 3108351
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 3170559
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 3234047
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 3298815
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 3364607
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 3432191
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 3500543
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 3570687
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 3642111
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 3714815
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 3789311
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 3864831
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 3942143
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 4021247
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 4101375
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 4183295
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 4267007
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 4352255
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 4439295
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 4528127
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 4618751
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 4711167
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 4805375
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 4901375
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 4999679
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 5099519
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 5201151
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 5305343
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 5411583
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 5519615
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 5630207
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 5742591
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 5857535
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 5974527
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 6094079
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 6216191
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 6340095
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 6467327
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 6596351
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 6728447
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 6862847
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 7000319
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 7140095
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 7282943
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 7428351
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 7577087
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 7728639
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 7883007
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 8040959
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 8201727
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 8365823
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 8532735
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 8703743
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 8877567
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 9054975
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 9236223
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 9421055
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 9609471
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 9801471
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 9997311
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 10197503
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 10401279
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 10609407
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 10821375
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 11037951
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 11258879
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 11483903
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 11713535
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 11947775
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 12186879
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 12430335
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 12679167
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 12932607
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 13191167
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 13454847
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 13724159
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 13998591
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 14278655
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 14564351
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 14855423
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 15152639
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 15455487
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 15764735
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 16079871
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 16401407
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 16729599
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 17064191
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 17405439
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 17753343
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 18108671
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 18470655
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 18840063
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 19216895
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 19601151
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 19993343
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 20392959
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 20800767
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 21217023
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 21641471
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 22074111
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 22515455
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 22965759
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 23425023
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 23893503
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 24371199
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 24858879
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 25355775
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 25863167
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 26380287
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 26907903
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 27446015
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 27995135
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 28555007
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 29125887
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 29708543
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 30302463
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 30908415
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 31526655
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 32157183
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 32800511
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 33456383
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 34125567
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 34808063
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 35504127
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 36214271
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 36938495
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 37677311
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 38430719
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 39199487
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 39983615
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 40783103
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 41598719
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 42430719
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 43279103
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 44144895
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 45027839
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 45928191
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 46846719
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 47783679
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 48739583
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 49714175
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 50708223
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 51722495
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 52756991
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 53812223
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 54888447
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 55985919
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 57105663
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 58247679
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 59412735
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 60600831
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 61813247
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 63049215
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 64310271
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 65596415
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 66908415
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 68246783
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 69611519
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 71003903
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 72423935
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 73872383
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 75349503
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 76856831
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 78393599
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 79961343
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 81560831
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 83192063
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 84855807
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 86552831
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 88283903
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 90049535
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 91850495
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 93687551
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 95561471
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 97472511
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 99422207
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 101410559
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 103438335
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 105507327
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 107617535
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 109769727
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 111965183
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 114204671
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 116488703
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 118818303
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 121194495
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 123618303
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 126091007
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 128612607
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 131184639
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 133808639
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 136484607
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 139214591
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 141998591
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 144838655
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 147735551
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 150690047
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 153703935
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 156777983
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 159913215
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 163111679
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 166374143
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 169701375
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 173095167
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 176557311
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 180088319
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 183690239
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 187363839
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 191111423
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 194933247
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 198832127
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 202808831
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 206864895
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 211002111
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 215222271
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 219526655
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 223917311
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 228395519
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 232963583
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 237622527
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 242374911
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 247222527
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 252166911
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 257210367
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 262354431
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 267601407
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 272953599
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 278412543
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 283981055
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 289660415
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 295453695
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 301362687
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 307389951
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 313537791
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 319808511
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 326204415
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 332728575
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 339383295
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 346170879
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 353094143
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 360156159
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 367359231
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 374706431
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 382200575
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 389844479
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 397641471
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 405594111
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 413706239
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 421980159
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 430419711
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 439028223
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 447808767
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 456764927
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 465900287
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 475217919
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 484722431
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 494416895
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 504305151
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 514391295
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 524679167
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 535172607
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 545875967
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 556793343
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 567929343
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 579288063
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 590873855
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 602691071
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 614745087
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 627039999
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 639580671
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 652372223
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 665419775
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 678727935
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 692302847
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 706148607
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 720271359
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 734677247
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 749370623
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 764358143
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 779645183
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 795237887
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 811142655
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 827365631
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 843912959
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 860791295
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 878006783
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 895567103
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 913478399
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 931747839
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 950382591
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 969390335
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 988777983
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 1008553727
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 1028724735
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 1049299199
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 1070285055
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 1091690751
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 1113524735
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 1135795199
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 1158511103
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 1181681151
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 1205314815
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 1229421311
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 1254009599
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 1279089663
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 1304671743
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 1330765055
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 1357380351
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 1384528127
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 1412218367
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 1440462591
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 1469271807
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 1498657535
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 1528630527
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 1559203071
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 1590387455
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 1622194943
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 1654639103
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 1687731711
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 1721486079
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 1755916031
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 1791034367
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 1826854911
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 1863391743
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 1900659711
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 1938673151
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 1977446399
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 2016995327
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 2057335295
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 2098481919
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 2140451583
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 2183260671
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 2226925823
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 2271464447
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 2316893439
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 2363231487
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 2410496255
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 2458705919
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 2507880191
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 2558037503
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 2609198591
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 2661382143
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 2714609919
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 2768902143
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 2824280063
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 2880765695
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 2938381055
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 2997148671
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 3057091583
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 3118233599
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 3180598271
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 3244210175
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 3309094143
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 3375276287
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 3442781439
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 3511637247
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 3581869823
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 3653507327
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 3726577407
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 3801108735
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 3877131263
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 3954673919
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 4033767167
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 4114442495
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 4196731391
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 4280666111
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 4366279167
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 4453604607
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 4542676991
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 4633530623
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 4726200831
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 4820725247
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 4917139455
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 5015482367
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 5115791871
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 5218107647
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 5322469887
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 5428919295
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 5537497599
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 5648247551
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 5761212671
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 5876436735
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 5993965311
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 6113844735
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 6236121599
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 6360844031
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 6488060927
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 6617822207
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 6750178559
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 6885182207
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 7022885631
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 7163343359
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 7306610175
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 7452742143
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 7601797119
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 7753833215
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 7908909567
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 8067088127
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 8228429823
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 8392998399
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 8560858367
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 8732075519
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 8906716671
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 9084851199
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 9266548223
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 9451879167
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 9640916735
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 9833734911
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 10030409727
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 10231017983
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 10435638015
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 10644350975
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 10857238271
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 11074382847
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 11295870207
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 11521787903
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 11752223487
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 11987268095
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 8
+  value_type_size: 8
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 1
+  }
+  scratch_size_recordings {
+    num_items: 2350
+    scratch_space_bytes: 55807
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 56319
+  }
+  scratch_size_recordings {
+    num_items: 2444
+    scratch_space_bytes: 57343
+  }
+  scratch_size_recordings {
+    num_items: 2492
+    scratch_space_bytes: 57855
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 58879
+  }
+  scratch_size_recordings {
+    num_items: 2591
+    scratch_space_bytes: 59391
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 60415
+  }
+  scratch_size_recordings {
+    num_items: 2694
+    scratch_space_bytes: 61439
+  }
+  scratch_size_recordings {
+    num_items: 2747
+    scratch_space_bytes: 61951
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 62975
+  }
+  scratch_size_recordings {
+    num_items: 2857
+    scratch_space_bytes: 63999
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 65023
+  }
+  scratch_size_recordings {
+    num_items: 2972
+    scratch_space_bytes: 65535
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 66559
+  }
+  scratch_size_recordings {
+    num_items: 3091
+    scratch_space_bytes: 67583
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 68607
+  }
+  scratch_size_recordings {
+    num_items: 3215
+    scratch_space_bytes: 69631
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 70655
+  }
+  scratch_size_recordings {
+    num_items: 3344
+    scratch_space_bytes: 71679
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 72703
+  }
+  scratch_size_recordings {
+    num_items: 3478
+    scratch_space_bytes: 73727
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 74751
+  }
+  scratch_size_recordings {
+    num_items: 3617
+    scratch_space_bytes: 76287
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 77311
+  }
+  scratch_size_recordings {
+    num_items: 3762
+    scratch_space_bytes: 78335
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 79359
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 80895
+  }
+  scratch_size_recordings {
+    num_items: 3991
+    scratch_space_bytes: 81919
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 83455
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 84479
+  }
+  scratch_size_recordings {
+    num_items: 4234
+    scratch_space_bytes: 86015
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 87039
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 88575
+  }
+  scratch_size_recordings {
+    num_items: 4492
+    scratch_space_bytes: 90111
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 91647
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 93695
+  }
+  scratch_size_recordings {
+    num_items: 4765
+    scratch_space_bytes: 95231
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 96767
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 98303
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 99839
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 101887
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 103423
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 104959
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 106495
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 108543
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 110079
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 112127
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 114175
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 115711
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 117759
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 119807
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 121855
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 123903
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 125951
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 127999
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 130047
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 132095
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 134655
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 136703
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 139263
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 141823
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 143871
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 146431
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 148991
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 151551
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 154111
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 157183
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 159743
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 162815
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 165375
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 169471
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 172543
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 175615
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 178687
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 181759
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 184831
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 188415
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 191487
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 195071
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 198655
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 201727
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 205823
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 209407
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 212991
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 217087
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 220671
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 224767
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 228863
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 232959
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 237567
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 242687
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 247295
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 251391
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 255999
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 261119
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 265727
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 270847
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 275455
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 280575
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 285695
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 291327
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 296447
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 302079
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 307711
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 313343
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 319999
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 326143
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 332287
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 338431
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 344575
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 351231
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 357887
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 364543
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 371199
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 378367
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 385535
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 393727
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 400895
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 408575
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 416255
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 423935
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 432127
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 440319
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 448511
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 457215
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 466943
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 475647
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 484863
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 494079
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 503295
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 513023
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 522751
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 532479
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 543743
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 553983
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 564735
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 575487
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 586239
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 597503
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 609279
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 622079
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 633855
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 645631
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 658431
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 670719
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 683519
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 697855
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 711167
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 724991
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 738815
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 753151
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 768511
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 783359
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 798719
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 814079
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 829439
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 846847
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 863231
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 879615
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 896511
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 913919
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 932863
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 950783
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 969215
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 988159
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 1008127
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 1027583
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 1047551
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 1069055
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 1089535
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 1110527
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 1132543
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 1155583
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 1177599
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 1200639
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 1225215
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 1248767
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 1272831
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 1298943
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 1324031
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 1350143
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 1377279
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 1403903
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 1431551
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 1460223
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 1488895
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 1518591
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 1548287
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 1578495
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 1610239
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 1641983
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 1674751
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 1707519
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 1741823
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 1775615
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 1810431
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 1846783
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 1883135
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 1921023
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 1958399
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 1997823
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 2037759
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 2077695
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 2119167
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 2160639
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 2204159
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 2247167
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 2292223
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 2337791
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 2383871
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 2431487
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 2479103
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 2528767
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 2578943
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 2629631
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 2682367
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 2735615
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 2790399
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 2845183
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 2902015
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 2959871
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 3018751
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 3078143
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 3139583
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 3202047
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 3265535
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 3331071
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 3397119
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 3464703
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 3533823
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 3603455
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 3675135
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 3748351
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 3822591
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 3898879
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 3976191
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 4056575
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 4137471
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 4219391
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 4303359
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 4388863
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 4476415
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 4564991
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 4657151
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 4749311
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 4844031
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 4940287
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 5039103
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 5139455
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 5241343
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 5346303
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 5452799
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 5561855
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 5671935
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 5785599
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 5900287
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 6018559
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 6137855
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 6260735
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 6386175
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 6512639
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 6642687
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 6775807
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 6910975
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 7048703
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 7188479
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 7332351
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 7478271
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 7627775
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 7779839
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 7934975
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 8094207
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 8255487
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 8420351
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 8587775
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 8759295
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 8934399
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 9112575
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 9294847
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 9480191
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 9669631
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 9862143
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 10059263
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 10260479
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 10464767
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 10673663
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 10887167
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 11104767
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 11326463
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 11552255
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 11783167
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 12018175
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 12257791
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 12503551
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 12752895
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 13007359
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 13267455
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 13531647
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 13802495
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 14078463
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 14359039
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 14645759
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 14938623
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 15237119
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 15541247
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 15851519
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 16168959
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 16492031
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 16820735
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 17157631
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 17499647
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 17849343
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 18206207
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 18570239
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 18941439
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 19319807
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 19705855
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 20099071
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 20501503
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 20910591
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 21328895
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 21754367
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 22189055
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 22632959
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 23085055
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 23546367
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 24017407
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 24497151
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 24987135
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 25485823
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 25995263
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 26514943
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 27044863
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 27585023
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 28136959
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 28699135
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 29273087
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 29858303
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 30454783
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 31063551
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 31684095
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 32317951
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 32963071
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 33622527
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 34294271
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 34979839
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 35679743
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 36392447
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 37119999
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 37862399
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 38619135
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 39390719
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 40178687
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 40980991
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 41800191
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 42636799
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 43488767
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 44358655
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 45245439
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 46150143
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 47072255
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 48013823
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 48972799
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 49952767
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 50951167
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 51969023
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 53008383
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 54069247
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 55149567
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 56252415
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 57377279
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 58523647
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 59694591
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 60888063
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 62105087
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 63347199
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 64613887
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 65905663
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 67223551
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 68567551
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 69938175
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 71336959
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 72763391
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 74217983
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 75701759
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 77215231
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 78759423
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 80334335
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 81939967
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 83578879
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 85250559
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 86955519
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 88693759
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 90466815
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 92276735
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 94120959
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 96004095
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 97923071
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 99880959
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 101878271
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 103915519
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 105993215
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 108113407
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 110274559
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 112480255
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 114729471
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 117023743
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 119363583
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 121750527
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 124185599
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 126668799
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 129202175
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 131785215
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 134420479
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 137108479
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 139850239
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 142647807
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 145500159
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 148409855
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 151376895
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 154404863
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 157492223
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 160641535
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 163853823
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 167130623
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 170473471
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 173881855
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 177359359
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 180905983
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 184524287
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 188214271
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 191977471
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 195817471
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 199733247
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 203727871
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 207801855
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 211957759
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 216195583
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 220519423
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 224929791
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 229427711
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 234016255
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 238695935
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 243469311
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 248338431
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 253305343
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 258371583
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 263538175
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 268808191
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 274184191
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 279668223
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 285260287
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 290966015
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 296784895
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 302719999
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 308773887
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 314949631
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 321247743
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 327672319
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 334224895
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 340909567
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 347727359
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 354681343
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 361774591
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 369010175
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 376389119
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 383917055
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 391595007
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 399427071
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 407414271
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 415562239
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 423874047
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 432350719
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 440997375
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 449817087
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 458812927
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 467988991
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 477347839
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 486895103
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 496632319
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 506564607
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 516695039
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 527029247
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 537569791
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 548319743
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 559286271
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 570471935
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 581880319
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 593518079
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 605388287
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 617495551
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 629844991
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 642441727
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 655289343
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 668395519
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 681763327
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 695397887
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 709305343
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 723491327
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 737959935
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 752719359
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 767773183
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 783129087
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 798790655
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 814766079
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 831061503
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 847682047
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 864635391
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 881927167
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 899565567
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 917556735
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 935907327
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 954625535
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 973716991
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 993191935
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 1013054463
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 1033316351
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 1053981695
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 1075060735
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 1096562175
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 1118492671
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 1140861951
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 1163678719
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 1186952191
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 1210690559
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 1234904063
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 1259602431
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 1284793855
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 1310488575
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 1336698367
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 1363432447
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 1390700031
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 1418514431
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 1446883839
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 1475821055
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 1505337343
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 1535443967
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 1566152703
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 1597475327
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 1629424127
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 1662012415
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 1695252479
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 1729156607
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 1763740159
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 1799014399
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 1834994175
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 1871693823
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 1909127167
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 1947309567
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 1986255359
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 2025979903
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 2066499071
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 2107828735
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 2149985279
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 2192984063
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 2236844031
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 2281580031
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 2327211519
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 2373754879
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 2421229567
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 2469654015
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 2519046655
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 2569427455
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 2620815359
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 2673231871
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 2726695935
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 2781229055
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 2836853247
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 2893590015
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 2951462399
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 3010490367
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 3070700543
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 3132113407
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 3194755583
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 3258650623
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 3323823103
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 3390299135
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 3458104319
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 3527267327
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 3597811199
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 3669767679
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 3743162879
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 3818024959
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 3894385663
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 3972273151
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 4051717631
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 4132751359
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 4215406591
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 4299714559
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 4385298431
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 4471810559
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 4560053247
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 4650060287
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 4741868031
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 4835511295
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 4931027967
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 5028454911
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 5127830015
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 5229193215
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 5332583423
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 5438041087
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 5545608191
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 5655326719
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 5767239679
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 5881390591
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 5997824511
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 6116587519
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 6237725695
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 6361286143
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 6487318527
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 6615870975
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 6746994687
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 6880740863
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 7017161727
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 7156311551
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 7298244095
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 7443015167
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 7590681599
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 7741301759
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 7894934015
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 8051638783
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 8211478015
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 8374513663
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 8540810239
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 8710449151
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 8883464191
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 9059939327
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 9239944191
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 9423548927
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 9610825727
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 9801848319
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 9996690943
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 10195430911
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 10398145535
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 10604914175
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 10815818751
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 11030941183
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 11250365951
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 11474179071
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 11702468607
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 11935323647
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 12172836351
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 12415098879
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 12662206975
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 12914256895
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 13171364351
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 13433597439
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 13701074943
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 13973901823
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 14252185599
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 14536035327
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 14825561599
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 15120878079
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 15422101503
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 15729349119
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 16042741759
+  }
+}
+entries {
+  cub_version: "3.2.0"
+  device_name: "NVIDIA VR NVL72"
+  key_type_size: 8
+  value_type_size: 8
+  is_segmented: true
+  scratch_size_recordings {
+    num_items: 31
+    scratch_space_bytes: 767
+  }
+  scratch_size_recordings {
+    num_items: 61
+    scratch_space_bytes: 1279
+  }
+  scratch_size_recordings {
+    num_items: 91
+    scratch_space_bytes: 1791
+  }
+  scratch_size_recordings {
+    num_items: 121
+    scratch_space_bytes: 2303
+  }
+  scratch_size_recordings {
+    num_items: 151
+    scratch_space_bytes: 2815
+  }
+  scratch_size_recordings {
+    num_items: 191
+    scratch_space_bytes: 3327
+  }
+  scratch_size_recordings {
+    num_items: 221
+    scratch_space_bytes: 3839
+  }
+  scratch_size_recordings {
+    num_items: 251
+    scratch_space_bytes: 4351
+  }
+  scratch_size_recordings {
+    num_items: 281
+    scratch_space_bytes: 4863
+  }
+  scratch_size_recordings {
+    num_items: 311
+    scratch_space_bytes: 5375
+  }
+  scratch_size_recordings {
+    num_items: 351
+    scratch_space_bytes: 5887
+  }
+  scratch_size_recordings {
+    num_items: 381
+    scratch_space_bytes: 6399
+  }
+  scratch_size_recordings {
+    num_items: 411
+    scratch_space_bytes: 6911
+  }
+  scratch_size_recordings {
+    num_items: 441
+    scratch_space_bytes: 7423
+  }
+  scratch_size_recordings {
+    num_items: 471
+    scratch_space_bytes: 7935
+  }
+  scratch_size_recordings {
+    num_items: 511
+    scratch_space_bytes: 8447
+  }
+  scratch_size_recordings {
+    num_items: 541
+    scratch_space_bytes: 8959
+  }
+  scratch_size_recordings {
+    num_items: 573
+    scratch_space_bytes: 9471
+  }
+  scratch_size_recordings {
+    num_items: 606
+    scratch_space_bytes: 9983
+  }
+  scratch_size_recordings {
+    num_items: 630
+    scratch_space_bytes: 10495
+  }
+  scratch_size_recordings {
+    num_items: 667
+    scratch_space_bytes: 11007
+  }
+  scratch_size_recordings {
+    num_items: 693
+    scratch_space_bytes: 11519
+  }
+  scratch_size_recordings {
+    num_items: 734
+    scratch_space_bytes: 12031
+  }
+  scratch_size_recordings {
+    num_items: 762
+    scratch_space_bytes: 12543
+  }
+  scratch_size_recordings {
+    num_items: 792
+    scratch_space_bytes: 13055
+  }
+  scratch_size_recordings {
+    num_items: 823
+    scratch_space_bytes: 13567
+  }
+  scratch_size_recordings {
+    num_items: 855
+    scratch_space_bytes: 14079
+  }
+  scratch_size_recordings {
+    num_items: 889
+    scratch_space_bytes: 14591
+  }
+  scratch_size_recordings {
+    num_items: 924
+    scratch_space_bytes: 15103
+  }
+  scratch_size_recordings {
+    num_items: 960
+    scratch_space_bytes: 15615
+  }
+  scratch_size_recordings {
+    num_items: 979
+    scratch_space_bytes: 16127
+  }
+  scratch_size_recordings {
+    num_items: 1017
+    scratch_space_bytes: 16639
+  }
+  scratch_size_recordings {
+    num_items: 1037
+    scratch_space_bytes: 17151
+  }
+  scratch_size_recordings {
+    num_items: 1078
+    scratch_space_bytes: 17663
+  }
+  scratch_size_recordings {
+    num_items: 1120
+    scratch_space_bytes: 18175
+  }
+  scratch_size_recordings {
+    num_items: 1142
+    scratch_space_bytes: 18687
+  }
+  scratch_size_recordings {
+    num_items: 1164
+    scratch_space_bytes: 19199
+  }
+  scratch_size_recordings {
+    num_items: 1210
+    scratch_space_bytes: 19711
+  }
+  scratch_size_recordings {
+    num_items: 1234
+    scratch_space_bytes: 20223
+  }
+  scratch_size_recordings {
+    num_items: 1258
+    scratch_space_bytes: 20735
+  }
+  scratch_size_recordings {
+    num_items: 1308
+    scratch_space_bytes: 21247
+  }
+  scratch_size_recordings {
+    num_items: 1334
+    scratch_space_bytes: 21759
+  }
+  scratch_size_recordings {
+    num_items: 1360
+    scratch_space_bytes: 22271
+  }
+  scratch_size_recordings {
+    num_items: 1387
+    scratch_space_bytes: 22783
+  }
+  scratch_size_recordings {
+    num_items: 1414
+    scratch_space_bytes: 23295
+  }
+  scratch_size_recordings {
+    num_items: 1470
+    scratch_space_bytes: 23807
+  }
+  scratch_size_recordings {
+    num_items: 1499
+    scratch_space_bytes: 24319
+  }
+  scratch_size_recordings {
+    num_items: 1528
+    scratch_space_bytes: 24831
+  }
+  scratch_size_recordings {
+    num_items: 1558
+    scratch_space_bytes: 25343
+  }
+  scratch_size_recordings {
+    num_items: 1589
+    scratch_space_bytes: 25855
+  }
+  scratch_size_recordings {
+    num_items: 1620
+    scratch_space_bytes: 26367
+  }
+  scratch_size_recordings {
+    num_items: 1652
+    scratch_space_bytes: 26879
+  }
+  scratch_size_recordings {
+    num_items: 1685
+    scratch_space_bytes: 27391
+  }
+  scratch_size_recordings {
+    num_items: 1718
+    scratch_space_bytes: 27903
+  }
+  scratch_size_recordings {
+    num_items: 1752
+    scratch_space_bytes: 28415
+  }
+  scratch_size_recordings {
+    num_items: 1787
+    scratch_space_bytes: 28927
+  }
+  scratch_size_recordings {
+    num_items: 1822
+    scratch_space_bytes: 29439
+  }
+  scratch_size_recordings {
+    num_items: 1858
+    scratch_space_bytes: 30463
+  }
+  scratch_size_recordings {
+    num_items: 1895
+    scratch_space_bytes: 30975
+  }
+  scratch_size_recordings {
+    num_items: 1932
+    scratch_space_bytes: 31487
+  }
+  scratch_size_recordings {
+    num_items: 1970
+    scratch_space_bytes: 31999
+  }
+  scratch_size_recordings {
+    num_items: 2009
+    scratch_space_bytes: 32511
+  }
+  scratch_size_recordings {
+    num_items: 2049
+    scratch_space_bytes: 33535
+  }
+  scratch_size_recordings {
+    num_items: 2089
+    scratch_space_bytes: 34047
+  }
+  scratch_size_recordings {
+    num_items: 2130
+    scratch_space_bytes: 34559
+  }
+  scratch_size_recordings {
+    num_items: 2172
+    scratch_space_bytes: 35071
+  }
+  scratch_size_recordings {
+    num_items: 2215
+    scratch_space_bytes: 36095
+  }
+  scratch_size_recordings {
+    num_items: 2259
+    scratch_space_bytes: 36607
+  }
+  scratch_size_recordings {
+    num_items: 2304
+    scratch_space_bytes: 37119
+  }
+  scratch_size_recordings {
+    num_items: 2350
+    scratch_space_bytes: 38143
+  }
+  scratch_size_recordings {
+    num_items: 2397
+    scratch_space_bytes: 38655
+  }
+  scratch_size_recordings {
+    num_items: 2444
+    scratch_space_bytes: 39679
+  }
+  scratch_size_recordings {
+    num_items: 2492
+    scratch_space_bytes: 40191
+  }
+  scratch_size_recordings {
+    num_items: 2541
+    scratch_space_bytes: 41215
+  }
+  scratch_size_recordings {
+    num_items: 2591
+    scratch_space_bytes: 41727
+  }
+  scratch_size_recordings {
+    num_items: 2642
+    scratch_space_bytes: 42751
+  }
+  scratch_size_recordings {
+    num_items: 2694
+    scratch_space_bytes: 43775
+  }
+  scratch_size_recordings {
+    num_items: 2747
+    scratch_space_bytes: 44287
+  }
+  scratch_size_recordings {
+    num_items: 2801
+    scratch_space_bytes: 45311
+  }
+  scratch_size_recordings {
+    num_items: 2857
+    scratch_space_bytes: 46335
+  }
+  scratch_size_recordings {
+    num_items: 2914
+    scratch_space_bytes: 47359
+  }
+  scratch_size_recordings {
+    num_items: 2972
+    scratch_space_bytes: 47871
+  }
+  scratch_size_recordings {
+    num_items: 3031
+    scratch_space_bytes: 48895
+  }
+  scratch_size_recordings {
+    num_items: 3091
+    scratch_space_bytes: 49919
+  }
+  scratch_size_recordings {
+    num_items: 3152
+    scratch_space_bytes: 50943
+  }
+  scratch_size_recordings {
+    num_items: 3215
+    scratch_space_bytes: 51967
+  }
+  scratch_size_recordings {
+    num_items: 3279
+    scratch_space_bytes: 52991
+  }
+  scratch_size_recordings {
+    num_items: 3344
+    scratch_space_bytes: 54015
+  }
+  scratch_size_recordings {
+    num_items: 3410
+    scratch_space_bytes: 55039
+  }
+  scratch_size_recordings {
+    num_items: 3478
+    scratch_space_bytes: 56063
+  }
+  scratch_size_recordings {
+    num_items: 3547
+    scratch_space_bytes: 57087
+  }
+  scratch_size_recordings {
+    num_items: 3617
+    scratch_space_bytes: 58623
+  }
+  scratch_size_recordings {
+    num_items: 3689
+    scratch_space_bytes: 59647
+  }
+  scratch_size_recordings {
+    num_items: 3762
+    scratch_space_bytes: 60671
+  }
+  scratch_size_recordings {
+    num_items: 3837
+    scratch_space_bytes: 61695
+  }
+  scratch_size_recordings {
+    num_items: 3913
+    scratch_space_bytes: 63231
+  }
+  scratch_size_recordings {
+    num_items: 3991
+    scratch_space_bytes: 64255
+  }
+  scratch_size_recordings {
+    num_items: 4070
+    scratch_space_bytes: 65791
+  }
+  scratch_size_recordings {
+    num_items: 4151
+    scratch_space_bytes: 66815
+  }
+  scratch_size_recordings {
+    num_items: 4234
+    scratch_space_bytes: 68351
+  }
+  scratch_size_recordings {
+    num_items: 4318
+    scratch_space_bytes: 69375
+  }
+  scratch_size_recordings {
+    num_items: 4404
+    scratch_space_bytes: 70911
+  }
+  scratch_size_recordings {
+    num_items: 4492
+    scratch_space_bytes: 72447
+  }
+  scratch_size_recordings {
+    num_items: 4581
+    scratch_space_bytes: 73983
+  }
+  scratch_size_recordings {
+    num_items: 4672
+    scratch_space_bytes: 75007
+  }
+  scratch_size_recordings {
+    num_items: 4765
+    scratch_space_bytes: 76543
+  }
+  scratch_size_recordings {
+    num_items: 4860
+    scratch_space_bytes: 78079
+  }
+  scratch_size_recordings {
+    num_items: 4957
+    scratch_space_bytes: 79615
+  }
+  scratch_size_recordings {
+    num_items: 5056
+    scratch_space_bytes: 81151
+  }
+  scratch_size_recordings {
+    num_items: 5157
+    scratch_space_bytes: 83199
+  }
+  scratch_size_recordings {
+    num_items: 5260
+    scratch_space_bytes: 84735
+  }
+  scratch_size_recordings {
+    num_items: 5365
+    scratch_space_bytes: 86271
+  }
+  scratch_size_recordings {
+    num_items: 5472
+    scratch_space_bytes: 87807
+  }
+  scratch_size_recordings {
+    num_items: 5581
+    scratch_space_bytes: 89855
+  }
+  scratch_size_recordings {
+    num_items: 5692
+    scratch_space_bytes: 91391
+  }
+  scratch_size_recordings {
+    num_items: 5805
+    scratch_space_bytes: 93439
+  }
+  scratch_size_recordings {
+    num_items: 5921
+    scratch_space_bytes: 95487
+  }
+  scratch_size_recordings {
+    num_items: 6039
+    scratch_space_bytes: 97023
+  }
+  scratch_size_recordings {
+    num_items: 6159
+    scratch_space_bytes: 99071
+  }
+  scratch_size_recordings {
+    num_items: 6282
+    scratch_space_bytes: 101119
+  }
+  scratch_size_recordings {
+    num_items: 6407
+    scratch_space_bytes: 103167
+  }
+  scratch_size_recordings {
+    num_items: 6535
+    scratch_space_bytes: 105215
+  }
+  scratch_size_recordings {
+    num_items: 6665
+    scratch_space_bytes: 107263
+  }
+  scratch_size_recordings {
+    num_items: 6798
+    scratch_space_bytes: 109311
+  }
+  scratch_size_recordings {
+    num_items: 6933
+    scratch_space_bytes: 111359
+  }
+  scratch_size_recordings {
+    num_items: 7071
+    scratch_space_bytes: 113407
+  }
+  scratch_size_recordings {
+    num_items: 7212
+    scratch_space_bytes: 115967
+  }
+  scratch_size_recordings {
+    num_items: 7356
+    scratch_space_bytes: 118015
+  }
+  scratch_size_recordings {
+    num_items: 7503
+    scratch_space_bytes: 120575
+  }
+  scratch_size_recordings {
+    num_items: 7653
+    scratch_space_bytes: 123135
+  }
+  scratch_size_recordings {
+    num_items: 7806
+    scratch_space_bytes: 125183
+  }
+  scratch_size_recordings {
+    num_items: 7962
+    scratch_space_bytes: 127743
+  }
+  scratch_size_recordings {
+    num_items: 8121
+    scratch_space_bytes: 130303
+  }
+  scratch_size_recordings {
+    num_items: 8283
+    scratch_space_bytes: 132863
+  }
+  scratch_size_recordings {
+    num_items: 8448
+    scratch_space_bytes: 135423
+  }
+  scratch_size_recordings {
+    num_items: 8616
+    scratch_space_bytes: 138495
+  }
+  scratch_size_recordings {
+    num_items: 8788
+    scratch_space_bytes: 141055
+  }
+  scratch_size_recordings {
+    num_items: 8963
+    scratch_space_bytes: 144127
+  }
+  scratch_size_recordings {
+    num_items: 9142
+    scratch_space_bytes: 146687
+  }
+  scratch_size_recordings {
+    num_items: 9324
+    scratch_space_bytes: 149759
+  }
+  scratch_size_recordings {
+    num_items: 9510
+    scratch_space_bytes: 152831
+  }
+  scratch_size_recordings {
+    num_items: 9700
+    scratch_space_bytes: 155903
+  }
+  scratch_size_recordings {
+    num_items: 9894
+    scratch_space_bytes: 158975
+  }
+  scratch_size_recordings {
+    num_items: 10091
+    scratch_space_bytes: 162047
+  }
+  scratch_size_recordings {
+    num_items: 10292
+    scratch_space_bytes: 165119
+  }
+  scratch_size_recordings {
+    num_items: 10497
+    scratch_space_bytes: 168703
+  }
+  scratch_size_recordings {
+    num_items: 10706
+    scratch_space_bytes: 171775
+  }
+  scratch_size_recordings {
+    num_items: 10920
+    scratch_space_bytes: 175359
+  }
+  scratch_size_recordings {
+    num_items: 11138
+    scratch_space_bytes: 178943
+  }
+  scratch_size_recordings {
+    num_items: 11360
+    scratch_space_bytes: 182015
+  }
+  scratch_size_recordings {
+    num_items: 11587
+    scratch_space_bytes: 186111
+  }
+  scratch_size_recordings {
+    num_items: 11818
+    scratch_space_bytes: 189695
+  }
+  scratch_size_recordings {
+    num_items: 12054
+    scratch_space_bytes: 193279
+  }
+  scratch_size_recordings {
+    num_items: 12295
+    scratch_space_bytes: 197375
+  }
+  scratch_size_recordings {
+    num_items: 12540
+    scratch_space_bytes: 200959
+  }
+  scratch_size_recordings {
+    num_items: 12790
+    scratch_space_bytes: 205055
+  }
+  scratch_size_recordings {
+    num_items: 13045
+    scratch_space_bytes: 209151
+  }
+  scratch_size_recordings {
+    num_items: 13305
+    scratch_space_bytes: 213247
+  }
+  scratch_size_recordings {
+    num_items: 13571
+    scratch_space_bytes: 217855
+  }
+  scratch_size_recordings {
+    num_items: 13842
+    scratch_space_bytes: 221951
+  }
+  scratch_size_recordings {
+    num_items: 14118
+    scratch_space_bytes: 226559
+  }
+  scratch_size_recordings {
+    num_items: 14400
+    scratch_space_bytes: 230655
+  }
+  scratch_size_recordings {
+    num_items: 14688
+    scratch_space_bytes: 235263
+  }
+  scratch_size_recordings {
+    num_items: 14981
+    scratch_space_bytes: 240383
+  }
+  scratch_size_recordings {
+    num_items: 15280
+    scratch_space_bytes: 244991
+  }
+  scratch_size_recordings {
+    num_items: 15585
+    scratch_space_bytes: 250111
+  }
+  scratch_size_recordings {
+    num_items: 15896
+    scratch_space_bytes: 254719
+  }
+  scratch_size_recordings {
+    num_items: 16213
+    scratch_space_bytes: 259839
+  }
+  scratch_size_recordings {
+    num_items: 16537
+    scratch_space_bytes: 264959
+  }
+  scratch_size_recordings {
+    num_items: 16867
+    scratch_space_bytes: 270591
+  }
+  scratch_size_recordings {
+    num_items: 17204
+    scratch_space_bytes: 275711
+  }
+  scratch_size_recordings {
+    num_items: 17548
+    scratch_space_bytes: 281343
+  }
+  scratch_size_recordings {
+    num_items: 17898
+    scratch_space_bytes: 286975
+  }
+  scratch_size_recordings {
+    num_items: 18255
+    scratch_space_bytes: 292607
+  }
+  scratch_size_recordings {
+    num_items: 18620
+    scratch_space_bytes: 298239
+  }
+  scratch_size_recordings {
+    num_items: 18992
+    scratch_space_bytes: 304383
+  }
+  scratch_size_recordings {
+    num_items: 19371
+    scratch_space_bytes: 310527
+  }
+  scratch_size_recordings {
+    num_items: 19758
+    scratch_space_bytes: 316671
+  }
+  scratch_size_recordings {
+    num_items: 20153
+    scratch_space_bytes: 322815
+  }
+  scratch_size_recordings {
+    num_items: 20556
+    scratch_space_bytes: 329471
+  }
+  scratch_size_recordings {
+    num_items: 20967
+    scratch_space_bytes: 336127
+  }
+  scratch_size_recordings {
+    num_items: 21386
+    scratch_space_bytes: 342783
+  }
+  scratch_size_recordings {
+    num_items: 21813
+    scratch_space_bytes: 349439
+  }
+  scratch_size_recordings {
+    num_items: 22249
+    scratch_space_bytes: 356607
+  }
+  scratch_size_recordings {
+    num_items: 22693
+    scratch_space_bytes: 363775
+  }
+  scratch_size_recordings {
+    num_items: 23146
+    scratch_space_bytes: 370943
+  }
+  scratch_size_recordings {
+    num_items: 23608
+    scratch_space_bytes: 378111
+  }
+  scratch_size_recordings {
+    num_items: 24080
+    scratch_space_bytes: 385791
+  }
+  scratch_size_recordings {
+    num_items: 24561
+    scratch_space_bytes: 393471
+  }
+  scratch_size_recordings {
+    num_items: 25052
+    scratch_space_bytes: 401151
+  }
+  scratch_size_recordings {
+    num_items: 25553
+    scratch_space_bytes: 409343
+  }
+  scratch_size_recordings {
+    num_items: 26064
+    scratch_space_bytes: 417535
+  }
+  scratch_size_recordings {
+    num_items: 26585
+    scratch_space_bytes: 425727
+  }
+  scratch_size_recordings {
+    num_items: 27116
+    scratch_space_bytes: 434431
+  }
+  scratch_size_recordings {
+    num_items: 27658
+    scratch_space_bytes: 443135
+  }
+  scratch_size_recordings {
+    num_items: 28211
+    scratch_space_bytes: 451839
+  }
+  scratch_size_recordings {
+    num_items: 28775
+    scratch_space_bytes: 461055
+  }
+  scratch_size_recordings {
+    num_items: 29350
+    scratch_space_bytes: 470271
+  }
+  scratch_size_recordings {
+    num_items: 29937
+    scratch_space_bytes: 479487
+  }
+  scratch_size_recordings {
+    num_items: 30535
+    scratch_space_bytes: 489215
+  }
+  scratch_size_recordings {
+    num_items: 31145
+    scratch_space_bytes: 498943
+  }
+  scratch_size_recordings {
+    num_items: 31767
+    scratch_space_bytes: 508671
+  }
+  scratch_size_recordings {
+    num_items: 32402
+    scratch_space_bytes: 518911
+  }
+  scratch_size_recordings {
+    num_items: 33050
+    scratch_space_bytes: 529151
+  }
+  scratch_size_recordings {
+    num_items: 33711
+    scratch_space_bytes: 539903
+  }
+  scratch_size_recordings {
+    num_items: 34385
+    scratch_space_bytes: 550655
+  }
+  scratch_size_recordings {
+    num_items: 35072
+    scratch_space_bytes: 561407
+  }
+  scratch_size_recordings {
+    num_items: 35773
+    scratch_space_bytes: 572671
+  }
+  scratch_size_recordings {
+    num_items: 36488
+    scratch_space_bytes: 584447
+  }
+  scratch_size_recordings {
+    num_items: 37217
+    scratch_space_bytes: 596223
+  }
+  scratch_size_recordings {
+    num_items: 37961
+    scratch_space_bytes: 607999
+  }
+  scratch_size_recordings {
+    num_items: 38720
+    scratch_space_bytes: 619775
+  }
+  scratch_size_recordings {
+    num_items: 39494
+    scratch_space_bytes: 632575
+  }
+  scratch_size_recordings {
+    num_items: 40283
+    scratch_space_bytes: 644863
+  }
+  scratch_size_recordings {
+    num_items: 41088
+    scratch_space_bytes: 657663
+  }
+  scratch_size_recordings {
+    num_items: 41909
+    scratch_space_bytes: 670975
+  }
+  scratch_size_recordings {
+    num_items: 42747
+    scratch_space_bytes: 684287
+  }
+  scratch_size_recordings {
+    num_items: 43601
+    scratch_space_bytes: 698111
+  }
+  scratch_size_recordings {
+    num_items: 44473
+    scratch_space_bytes: 711935
+  }
+  scratch_size_recordings {
+    num_items: 45362
+    scratch_space_bytes: 726271
+  }
+  scratch_size_recordings {
+    num_items: 46269
+    scratch_space_bytes: 740607
+  }
+  scratch_size_recordings {
+    num_items: 47194
+    scratch_space_bytes: 755455
+  }
+  scratch_size_recordings {
+    num_items: 48137
+    scratch_space_bytes: 770815
+  }
+  scratch_size_recordings {
+    num_items: 49099
+    scratch_space_bytes: 786175
+  }
+  scratch_size_recordings {
+    num_items: 50080
+    scratch_space_bytes: 801535
+  }
+  scratch_size_recordings {
+    num_items: 51081
+    scratch_space_bytes: 817919
+  }
+  scratch_size_recordings {
+    num_items: 52102
+    scratch_space_bytes: 834303
+  }
+  scratch_size_recordings {
+    num_items: 53144
+    scratch_space_bytes: 850687
+  }
+  scratch_size_recordings {
+    num_items: 54206
+    scratch_space_bytes: 867583
+  }
+  scratch_size_recordings {
+    num_items: 55290
+    scratch_space_bytes: 884991
+  }
+  scratch_size_recordings {
+    num_items: 56395
+    scratch_space_bytes: 902911
+  }
+  scratch_size_recordings {
+    num_items: 57522
+    scratch_space_bytes: 920831
+  }
+  scratch_size_recordings {
+    num_items: 58672
+    scratch_space_bytes: 939263
+  }
+  scratch_size_recordings {
+    num_items: 59845
+    scratch_space_bytes: 958207
+  }
+  scratch_size_recordings {
+    num_items: 61041
+    scratch_space_bytes: 977151
+  }
+  scratch_size_recordings {
+    num_items: 62261
+    scratch_space_bytes: 996607
+  }
+  scratch_size_recordings {
+    num_items: 63506
+    scratch_space_bytes: 1016575
+  }
+  scratch_size_recordings {
+    num_items: 64776
+    scratch_space_bytes: 1037055
+  }
+  scratch_size_recordings {
+    num_items: 66071
+    scratch_space_bytes: 1057535
+  }
+  scratch_size_recordings {
+    num_items: 67392
+    scratch_space_bytes: 1078527
+  }
+  scratch_size_recordings {
+    num_items: 68739
+    scratch_space_bytes: 1100543
+  }
+  scratch_size_recordings {
+    num_items: 70113
+    scratch_space_bytes: 1122559
+  }
+  scratch_size_recordings {
+    num_items: 71515
+    scratch_space_bytes: 1144575
+  }
+  scratch_size_recordings {
+    num_items: 72945
+    scratch_space_bytes: 1167615
+  }
+  scratch_size_recordings {
+    num_items: 74403
+    scratch_space_bytes: 1191167
+  }
+  scratch_size_recordings {
+    num_items: 75891
+    scratch_space_bytes: 1214719
+  }
+  scratch_size_recordings {
+    num_items: 77408
+    scratch_space_bytes: 1238783
+  }
+  scratch_size_recordings {
+    num_items: 78956
+    scratch_space_bytes: 1263871
+  }
+  scratch_size_recordings {
+    num_items: 80535
+    scratch_space_bytes: 1288959
+  }
+  scratch_size_recordings {
+    num_items: 82145
+    scratch_space_bytes: 1315071
+  }
+  scratch_size_recordings {
+    num_items: 83787
+    scratch_space_bytes: 1341183
+  }
+  scratch_size_recordings {
+    num_items: 85462
+    scratch_space_bytes: 1367807
+  }
+  scratch_size_recordings {
+    num_items: 87171
+    scratch_space_bytes: 1395455
+  }
+  scratch_size_recordings {
+    num_items: 88914
+    scratch_space_bytes: 1423103
+  }
+  scratch_size_recordings {
+    num_items: 90692
+    scratch_space_bytes: 1451775
+  }
+  scratch_size_recordings {
+    num_items: 92505
+    scratch_space_bytes: 1480447
+  }
+  scratch_size_recordings {
+    num_items: 94355
+    scratch_space_bytes: 1510143
+  }
+  scratch_size_recordings {
+    num_items: 96242
+    scratch_space_bytes: 1540351
+  }
+  scratch_size_recordings {
+    num_items: 98166
+    scratch_space_bytes: 1571071
+  }
+  scratch_size_recordings {
+    num_items: 100129
+    scratch_space_bytes: 1602815
+  }
+  scratch_size_recordings {
+    num_items: 102131
+    scratch_space_bytes: 1634559
+  }
+  scratch_size_recordings {
+    num_items: 104173
+    scratch_space_bytes: 1667327
+  }
+  scratch_size_recordings {
+    num_items: 106256
+    scratch_space_bytes: 1700607
+  }
+  scratch_size_recordings {
+    num_items: 108381
+    scratch_space_bytes: 1734399
+  }
+  scratch_size_recordings {
+    num_items: 110548
+    scratch_space_bytes: 1769215
+  }
+  scratch_size_recordings {
+    num_items: 112758
+    scratch_space_bytes: 1804543
+  }
+  scratch_size_recordings {
+    num_items: 115013
+    scratch_space_bytes: 1840895
+  }
+  scratch_size_recordings {
+    num_items: 117313
+    scratch_space_bytes: 1877759
+  }
+  scratch_size_recordings {
+    num_items: 119659
+    scratch_space_bytes: 1915135
+  }
+  scratch_size_recordings {
+    num_items: 122052
+    scratch_space_bytes: 1953535
+  }
+  scratch_size_recordings {
+    num_items: 124493
+    scratch_space_bytes: 1992447
+  }
+  scratch_size_recordings {
+    num_items: 126982
+    scratch_space_bytes: 2032383
+  }
+  scratch_size_recordings {
+    num_items: 129521
+    scratch_space_bytes: 2072831
+  }
+  scratch_size_recordings {
+    num_items: 132111
+    scratch_space_bytes: 2114303
+  }
+  scratch_size_recordings {
+    num_items: 134753
+    scratch_space_bytes: 2156799
+  }
+  scratch_size_recordings {
+    num_items: 137448
+    scratch_space_bytes: 2199807
+  }
+  scratch_size_recordings {
+    num_items: 140196
+    scratch_space_bytes: 2243839
+  }
+  scratch_size_recordings {
+    num_items: 142999
+    scratch_space_bytes: 2288383
+  }
+  scratch_size_recordings {
+    num_items: 145858
+    scratch_space_bytes: 2334463
+  }
+  scratch_size_recordings {
+    num_items: 148775
+    scratch_space_bytes: 2381055
+  }
+  scratch_size_recordings {
+    num_items: 151750
+    scratch_space_bytes: 2428671
+  }
+  scratch_size_recordings {
+    num_items: 154785
+    scratch_space_bytes: 2477311
+  }
+  scratch_size_recordings {
+    num_items: 157880
+    scratch_space_bytes: 2526463
+  }
+  scratch_size_recordings {
+    num_items: 161037
+    scratch_space_bytes: 2577151
+  }
+  scratch_size_recordings {
+    num_items: 164257
+    scratch_space_bytes: 2628863
+  }
+  scratch_size_recordings {
+    num_items: 167542
+    scratch_space_bytes: 2681087
+  }
+  scratch_size_recordings {
+    num_items: 170892
+    scratch_space_bytes: 2734847
+  }
+  scratch_size_recordings {
+    num_items: 174309
+    scratch_space_bytes: 2789631
+  }
+  scratch_size_recordings {
+    num_items: 177795
+    scratch_space_bytes: 2845439
+  }
+  scratch_size_recordings {
+    num_items: 181350
+    scratch_space_bytes: 2902271
+  }
+  scratch_size_recordings {
+    num_items: 184977
+    scratch_space_bytes: 2960127
+  }
+  scratch_size_recordings {
+    num_items: 188676
+    scratch_space_bytes: 3019519
+  }
+  scratch_size_recordings {
+    num_items: 192449
+    scratch_space_bytes: 3079935
+  }
+  scratch_size_recordings {
+    num_items: 196297
+    scratch_space_bytes: 3141375
+  }
+  scratch_size_recordings {
+    num_items: 200222
+    scratch_space_bytes: 3203839
+  }
+  scratch_size_recordings {
+    num_items: 204226
+    scratch_space_bytes: 3268351
+  }
+  scratch_size_recordings {
+    num_items: 208310
+    scratch_space_bytes: 3333375
+  }
+  scratch_size_recordings {
+    num_items: 212476
+    scratch_space_bytes: 3399935
+  }
+  scratch_size_recordings {
+    num_items: 216725
+    scratch_space_bytes: 3468031
+  }
+  scratch_size_recordings {
+    num_items: 221059
+    scratch_space_bytes: 3537663
+  }
+  scratch_size_recordings {
+    num_items: 225480
+    scratch_space_bytes: 3608319
+  }
+  scratch_size_recordings {
+    num_items: 229989
+    scratch_space_bytes: 3680511
+  }
+  scratch_size_recordings {
+    num_items: 234588
+    scratch_space_bytes: 3753727
+  }
+  scratch_size_recordings {
+    num_items: 239279
+    scratch_space_bytes: 3828991
+  }
+  scratch_size_recordings {
+    num_items: 244064
+    scratch_space_bytes: 3905279
+  }
+  scratch_size_recordings {
+    num_items: 248945
+    scratch_space_bytes: 3983615
+  }
+  scratch_size_recordings {
+    num_items: 253923
+    scratch_space_bytes: 4063487
+  }
+  scratch_size_recordings {
+    num_items: 259001
+    scratch_space_bytes: 4144383
+  }
+  scratch_size_recordings {
+    num_items: 264181
+    scratch_space_bytes: 4227327
+  }
+  scratch_size_recordings {
+    num_items: 269464
+    scratch_space_bytes: 4311807
+  }
+  scratch_size_recordings {
+    num_items: 274853
+    scratch_space_bytes: 4398335
+  }
+  scratch_size_recordings {
+    num_items: 280350
+    scratch_space_bytes: 4485887
+  }
+  scratch_size_recordings {
+    num_items: 285957
+    scratch_space_bytes: 4575999
+  }
+  scratch_size_recordings {
+    num_items: 291676
+    scratch_space_bytes: 4667135
+  }
+  scratch_size_recordings {
+    num_items: 297509
+    scratch_space_bytes: 4760831
+  }
+  scratch_size_recordings {
+    num_items: 303459
+    scratch_space_bytes: 4856063
+  }
+  scratch_size_recordings {
+    num_items: 309528
+    scratch_space_bytes: 4952831
+  }
+  scratch_size_recordings {
+    num_items: 315718
+    scratch_space_bytes: 5052159
+  }
+  scratch_size_recordings {
+    num_items: 322032
+    scratch_space_bytes: 5153023
+  }
+  scratch_size_recordings {
+    num_items: 328472
+    scratch_space_bytes: 5255935
+  }
+  scratch_size_recordings {
+    num_items: 335041
+    scratch_space_bytes: 5361407
+  }
+  scratch_size_recordings {
+    num_items: 341741
+    scratch_space_bytes: 5468415
+  }
+  scratch_size_recordings {
+    num_items: 348575
+    scratch_space_bytes: 5577471
+  }
+  scratch_size_recordings {
+    num_items: 355546
+    scratch_space_bytes: 5689087
+  }
+  scratch_size_recordings {
+    num_items: 362656
+    scratch_space_bytes: 5802751
+  }
+  scratch_size_recordings {
+    num_items: 369909
+    scratch_space_bytes: 5918975
+  }
+  scratch_size_recordings {
+    num_items: 377307
+    scratch_space_bytes: 6037247
+  }
+  scratch_size_recordings {
+    num_items: 384853
+    scratch_space_bytes: 6158079
+  }
+  scratch_size_recordings {
+    num_items: 392550
+    scratch_space_bytes: 6281471
+  }
+  scratch_size_recordings {
+    num_items: 400401
+    scratch_space_bytes: 6406911
+  }
+  scratch_size_recordings {
+    num_items: 408409
+    scratch_space_bytes: 6534911
+  }
+  scratch_size_recordings {
+    num_items: 416577
+    scratch_space_bytes: 6665983
+  }
+  scratch_size_recordings {
+    num_items: 424908
+    scratch_space_bytes: 6799103
+  }
+  scratch_size_recordings {
+    num_items: 433406
+    scratch_space_bytes: 6934783
+  }
+  scratch_size_recordings {
+    num_items: 442074
+    scratch_space_bytes: 7073535
+  }
+  scratch_size_recordings {
+    num_items: 450915
+    scratch_space_bytes: 7215359
+  }
+  scratch_size_recordings {
+    num_items: 459933
+    scratch_space_bytes: 7359231
+  }
+  scratch_size_recordings {
+    num_items: 469131
+    scratch_space_bytes: 7506687
+  }
+  scratch_size_recordings {
+    num_items: 478513
+    scratch_space_bytes: 7656703
+  }
+  scratch_size_recordings {
+    num_items: 488083
+    scratch_space_bytes: 7809791
+  }
+  scratch_size_recordings {
+    num_items: 497844
+    scratch_space_bytes: 7965951
+  }
+  scratch_size_recordings {
+    num_items: 507800
+    scratch_space_bytes: 8125183
+  }
+  scratch_size_recordings {
+    num_items: 517956
+    scratch_space_bytes: 8287999
+  }
+  scratch_size_recordings {
+    num_items: 528315
+    scratch_space_bytes: 8453375
+  }
+  scratch_size_recordings {
+    num_items: 538881
+    scratch_space_bytes: 8622847
+  }
+  scratch_size_recordings {
+    num_items: 549658
+    scratch_space_bytes: 8794879
+  }
+  scratch_size_recordings {
+    num_items: 560651
+    scratch_space_bytes: 8971007
+  }
+  scratch_size_recordings {
+    num_items: 571864
+    scratch_space_bytes: 9150207
+  }
+  scratch_size_recordings {
+    num_items: 583301
+    scratch_space_bytes: 9333503
+  }
+  scratch_size_recordings {
+    num_items: 594967
+    scratch_space_bytes: 9519871
+  }
+  scratch_size_recordings {
+    num_items: 606866
+    scratch_space_bytes: 9710335
+  }
+  scratch_size_recordings {
+    num_items: 619003
+    scratch_space_bytes: 9904383
+  }
+  scratch_size_recordings {
+    num_items: 631383
+    scratch_space_bytes: 10102527
+  }
+  scratch_size_recordings {
+    num_items: 644010
+    scratch_space_bytes: 10304767
+  }
+  scratch_size_recordings {
+    num_items: 656890
+    scratch_space_bytes: 10510591
+  }
+  scratch_size_recordings {
+    num_items: 670027
+    scratch_space_bytes: 10721023
+  }
+  scratch_size_recordings {
+    num_items: 683427
+    scratch_space_bytes: 10935551
+  }
+  scratch_size_recordings {
+    num_items: 697095
+    scratch_space_bytes: 11154175
+  }
+  scratch_size_recordings {
+    num_items: 711036
+    scratch_space_bytes: 11376895
+  }
+  scratch_size_recordings {
+    num_items: 725256
+    scratch_space_bytes: 11604735
+  }
+  scratch_size_recordings {
+    num_items: 739761
+    scratch_space_bytes: 11836671
+  }
+  scratch_size_recordings {
+    num_items: 754556
+    scratch_space_bytes: 12073215
+  }
+  scratch_size_recordings {
+    num_items: 769647
+    scratch_space_bytes: 12314879
+  }
+  scratch_size_recordings {
+    num_items: 785039
+    scratch_space_bytes: 12561151
+  }
+  scratch_size_recordings {
+    num_items: 800739
+    scratch_space_bytes: 12812543
+  }
+  scratch_size_recordings {
+    num_items: 816753
+    scratch_space_bytes: 13068543
+  }
+  scratch_size_recordings {
+    num_items: 833088
+    scratch_space_bytes: 13329663
+  }
+  scratch_size_recordings {
+    num_items: 849749
+    scratch_space_bytes: 13596415
+  }
+  scratch_size_recordings {
+    num_items: 866743
+    scratch_space_bytes: 13868287
+  }
+  scratch_size_recordings {
+    num_items: 884077
+    scratch_space_bytes: 14145791
+  }
+  scratch_size_recordings {
+    num_items: 901758
+    scratch_space_bytes: 14428415
+  }
+  scratch_size_recordings {
+    num_items: 919793
+    scratch_space_bytes: 14717183
+  }
+  scratch_size_recordings {
+    num_items: 938188
+    scratch_space_bytes: 15011583
+  }
+  scratch_size_recordings {
+    num_items: 956951
+    scratch_space_bytes: 15311615
+  }
+  scratch_size_recordings {
+    num_items: 976090
+    scratch_space_bytes: 15617791
+  }
+  scratch_size_recordings {
+    num_items: 995611
+    scratch_space_bytes: 15930111
+  }
+  scratch_size_recordings {
+    num_items: 1015523
+    scratch_space_bytes: 16249087
+  }
+  scratch_size_recordings {
+    num_items: 1035833
+    scratch_space_bytes: 16573695
+  }
+  scratch_size_recordings {
+    num_items: 1056549
+    scratch_space_bytes: 16905471
+  }
+  scratch_size_recordings {
+    num_items: 1077679
+    scratch_space_bytes: 17243391
+  }
+  scratch_size_recordings {
+    num_items: 1099232
+    scratch_space_bytes: 17587967
+  }
+  scratch_size_recordings {
+    num_items: 1121216
+    scratch_space_bytes: 17939711
+  }
+  scratch_size_recordings {
+    num_items: 1143640
+    scratch_space_bytes: 18298623
+  }
+  scratch_size_recordings {
+    num_items: 1166512
+    scratch_space_bytes: 18664703
+  }
+  scratch_size_recordings {
+    num_items: 1189842
+    scratch_space_bytes: 19037951
+  }
+  scratch_size_recordings {
+    num_items: 1213638
+    scratch_space_bytes: 19418879
+  }
+  scratch_size_recordings {
+    num_items: 1237910
+    scratch_space_bytes: 19806975
+  }
+  scratch_size_recordings {
+    num_items: 1262668
+    scratch_space_bytes: 20203263
+  }
+  scratch_size_recordings {
+    num_items: 1287921
+    scratch_space_bytes: 20607231
+  }
+  scratch_size_recordings {
+    num_items: 1313679
+    scratch_space_bytes: 21019391
+  }
+  scratch_size_recordings {
+    num_items: 1339952
+    scratch_space_bytes: 21439743
+  }
+  scratch_size_recordings {
+    num_items: 1366751
+    scratch_space_bytes: 21868287
+  }
+  scratch_size_recordings {
+    num_items: 1394086
+    scratch_space_bytes: 22306047
+  }
+  scratch_size_recordings {
+    num_items: 1421967
+    scratch_space_bytes: 22751999
+  }
+  scratch_size_recordings {
+    num_items: 1450406
+    scratch_space_bytes: 23207167
+  }
+  scratch_size_recordings {
+    num_items: 1479414
+    scratch_space_bytes: 23671039
+  }
+  scratch_size_recordings {
+    num_items: 1509002
+    scratch_space_bytes: 24144639
+  }
+  scratch_size_recordings {
+    num_items: 1539182
+    scratch_space_bytes: 24627455
+  }
+  scratch_size_recordings {
+    num_items: 1569965
+    scratch_space_bytes: 25119999
+  }
+  scratch_size_recordings {
+    num_items: 1601364
+    scratch_space_bytes: 25622271
+  }
+  scratch_size_recordings {
+    num_items: 1633391
+    scratch_space_bytes: 26134783
+  }
+  scratch_size_recordings {
+    num_items: 1666058
+    scratch_space_bytes: 26657535
+  }
+  scratch_size_recordings {
+    num_items: 1699379
+    scratch_space_bytes: 27190527
+  }
+  scratch_size_recordings {
+    num_items: 1733366
+    scratch_space_bytes: 27734271
+  }
+  scratch_size_recordings {
+    num_items: 1768033
+    scratch_space_bytes: 28289279
+  }
+  scratch_size_recordings {
+    num_items: 1803393
+    scratch_space_bytes: 28855039
+  }
+  scratch_size_recordings {
+    num_items: 1839460
+    scratch_space_bytes: 29432063
+  }
+  scratch_size_recordings {
+    num_items: 1876249
+    scratch_space_bytes: 30020351
+  }
+  scratch_size_recordings {
+    num_items: 1913773
+    scratch_space_bytes: 30620927
+  }
+  scratch_size_recordings {
+    num_items: 1952048
+    scratch_space_bytes: 31233279
+  }
+  scratch_size_recordings {
+    num_items: 1991088
+    scratch_space_bytes: 31857919
+  }
+  scratch_size_recordings {
+    num_items: 2030909
+    scratch_space_bytes: 32494847
+  }
+  scratch_size_recordings {
+    num_items: 2071527
+    scratch_space_bytes: 33145087
+  }
+  scratch_size_recordings {
+    num_items: 2112957
+    scratch_space_bytes: 33807615
+  }
+  scratch_size_recordings {
+    num_items: 2155216
+    scratch_space_bytes: 34483967
+  }
+  scratch_size_recordings {
+    num_items: 2198320
+    scratch_space_bytes: 35173631
+  }
+  scratch_size_recordings {
+    num_items: 2242286
+    scratch_space_bytes: 35877119
+  }
+  scratch_size_recordings {
+    num_items: 2287131
+    scratch_space_bytes: 36594431
+  }
+  scratch_size_recordings {
+    num_items: 2332873
+    scratch_space_bytes: 37326591
+  }
+  scratch_size_recordings {
+    num_items: 2379530
+    scratch_space_bytes: 38073087
+  }
+  scratch_size_recordings {
+    num_items: 2427120
+    scratch_space_bytes: 38834431
+  }
+  scratch_size_recordings {
+    num_items: 2475662
+    scratch_space_bytes: 39611135
+  }
+  scratch_size_recordings {
+    num_items: 2525175
+    scratch_space_bytes: 40403199
+  }
+  scratch_size_recordings {
+    num_items: 2575678
+    scratch_space_bytes: 41211135
+  }
+  scratch_size_recordings {
+    num_items: 2627191
+    scratch_space_bytes: 42035455
+  }
+  scratch_size_recordings {
+    num_items: 2679734
+    scratch_space_bytes: 42876159
+  }
+  scratch_size_recordings {
+    num_items: 2733328
+    scratch_space_bytes: 43733759
+  }
+  scratch_size_recordings {
+    num_items: 2787994
+    scratch_space_bytes: 44608255
+  }
+  scratch_size_recordings {
+    num_items: 2843753
+    scratch_space_bytes: 45500671
+  }
+  scratch_size_recordings {
+    num_items: 2900628
+    scratch_space_bytes: 46410495
+  }
+  scratch_size_recordings {
+    num_items: 2958640
+    scratch_space_bytes: 47338751
+  }
+  scratch_size_recordings {
+    num_items: 3017812
+    scratch_space_bytes: 48285439
+  }
+  scratch_size_recordings {
+    num_items: 3078168
+    scratch_space_bytes: 49251071
+  }
+  scratch_size_recordings {
+    num_items: 3139731
+    scratch_space_bytes: 50236159
+  }
+  scratch_size_recordings {
+    num_items: 3202525
+    scratch_space_bytes: 51240703
+  }
+  scratch_size_recordings {
+    num_items: 3266575
+    scratch_space_bytes: 52265727
+  }
+  scratch_size_recordings {
+    num_items: 3331906
+    scratch_space_bytes: 53311231
+  }
+  scratch_size_recordings {
+    num_items: 3398544
+    scratch_space_bytes: 54377215
+  }
+  scratch_size_recordings {
+    num_items: 3466514
+    scratch_space_bytes: 55464703
+  }
+  scratch_size_recordings {
+    num_items: 3535844
+    scratch_space_bytes: 56574207
+  }
+  scratch_size_recordings {
+    num_items: 3606560
+    scratch_space_bytes: 57705215
+  }
+  scratch_size_recordings {
+    num_items: 3678691
+    scratch_space_bytes: 58859775
+  }
+  scratch_size_recordings {
+    num_items: 3752264
+    scratch_space_bytes: 60036863
+  }
+  scratch_size_recordings {
+    num_items: 3827309
+    scratch_space_bytes: 61237503
+  }
+  scratch_size_recordings {
+    num_items: 3903855
+    scratch_space_bytes: 62462207
+  }
+  scratch_size_recordings {
+    num_items: 3981932
+    scratch_space_bytes: 63711487
+  }
+  scratch_size_recordings {
+    num_items: 4061570
+    scratch_space_bytes: 64985855
+  }
+  scratch_size_recordings {
+    num_items: 4142801
+    scratch_space_bytes: 66285311
+  }
+  scratch_size_recordings {
+    num_items: 4225657
+    scratch_space_bytes: 67610879
+  }
+  scratch_size_recordings {
+    num_items: 4310170
+    scratch_space_bytes: 68963071
+  }
+  scratch_size_recordings {
+    num_items: 4396373
+    scratch_space_bytes: 70342399
+  }
+  scratch_size_recordings {
+    num_items: 4484300
+    scratch_space_bytes: 71749375
+  }
+  scratch_size_recordings {
+    num_items: 4573986
+    scratch_space_bytes: 73184511
+  }
+  scratch_size_recordings {
+    num_items: 4665465
+    scratch_space_bytes: 74647807
+  }
+  scratch_size_recordings {
+    num_items: 4758774
+    scratch_space_bytes: 76140799
+  }
+  scratch_size_recordings {
+    num_items: 4853949
+    scratch_space_bytes: 77663487
+  }
+  scratch_size_recordings {
+    num_items: 4951027
+    scratch_space_bytes: 79216895
+  }
+  scratch_size_recordings {
+    num_items: 5050047
+    scratch_space_bytes: 80801023
+  }
+  scratch_size_recordings {
+    num_items: 5151047
+    scratch_space_bytes: 82417407
+  }
+  scratch_size_recordings {
+    num_items: 5254067
+    scratch_space_bytes: 84065535
+  }
+  scratch_size_recordings {
+    num_items: 5359148
+    scratch_space_bytes: 85746943
+  }
+  scratch_size_recordings {
+    num_items: 5466330
+    scratch_space_bytes: 87461631
+  }
+  scratch_size_recordings {
+    num_items: 5575656
+    scratch_space_bytes: 89211135
+  }
+  scratch_size_recordings {
+    num_items: 5687169
+    scratch_space_bytes: 90995455
+  }
+  scratch_size_recordings {
+    num_items: 5800912
+    scratch_space_bytes: 92815103
+  }
+  scratch_size_recordings {
+    num_items: 5916930
+    scratch_space_bytes: 94671615
+  }
+  scratch_size_recordings {
+    num_items: 6035268
+    scratch_space_bytes: 96564991
+  }
+  scratch_size_recordings {
+    num_items: 6155973
+    scratch_space_bytes: 98496255
+  }
+  scratch_size_recordings {
+    num_items: 6279092
+    scratch_space_bytes: 100465919
+  }
+  scratch_size_recordings {
+    num_items: 6404673
+    scratch_space_bytes: 102475519
+  }
+  scratch_size_recordings {
+    num_items: 6532766
+    scratch_space_bytes: 104524543
+  }
+  scratch_size_recordings {
+    num_items: 6663421
+    scratch_space_bytes: 106615039
+  }
+  scratch_size_recordings {
+    num_items: 6796689
+    scratch_space_bytes: 108747519
+  }
+  scratch_size_recordings {
+    num_items: 6932622
+    scratch_space_bytes: 110922495
+  }
+  scratch_size_recordings {
+    num_items: 7071274
+    scratch_space_bytes: 113140991
+  }
+  scratch_size_recordings {
+    num_items: 7212699
+    scratch_space_bytes: 115403519
+  }
+  scratch_size_recordings {
+    num_items: 7356952
+    scratch_space_bytes: 117711615
+  }
+  scratch_size_recordings {
+    num_items: 7504091
+    scratch_space_bytes: 120065791
+  }
+  scratch_size_recordings {
+    num_items: 7654172
+    scratch_space_bytes: 122467071
+  }
+  scratch_size_recordings {
+    num_items: 7807255
+    scratch_space_bytes: 124916479
+  }
+  scratch_size_recordings {
+    num_items: 7963400
+    scratch_space_bytes: 127415039
+  }
+  scratch_size_recordings {
+    num_items: 8122668
+    scratch_space_bytes: 129963263
+  }
+  scratch_size_recordings {
+    num_items: 8285121
+    scratch_space_bytes: 132562687
+  }
+  scratch_size_recordings {
+    num_items: 8450823
+    scratch_space_bytes: 135213823
+  }
+  scratch_size_recordings {
+    num_items: 8619839
+    scratch_space_bytes: 137917695
+  }
+  scratch_size_recordings {
+    num_items: 8792235
+    scratch_space_bytes: 140676351
+  }
+  scratch_size_recordings {
+    num_items: 8968079
+    scratch_space_bytes: 143489791
+  }
+  scratch_size_recordings {
+    num_items: 9147440
+    scratch_space_bytes: 146359551
+  }
+  scratch_size_recordings {
+    num_items: 9330388
+    scratch_space_bytes: 149286655
+  }
+  scratch_size_recordings {
+    num_items: 9516995
+    scratch_space_bytes: 152272639
+  }
+  scratch_size_recordings {
+    num_items: 9707334
+    scratch_space_bytes: 155318015
+  }
+  scratch_size_recordings {
+    num_items: 9901480
+    scratch_space_bytes: 158424319
+  }
+  scratch_size_recordings {
+    num_items: 10099509
+    scratch_space_bytes: 161592575
+  }
+  scratch_size_recordings {
+    num_items: 10301499
+    scratch_space_bytes: 164824319
+  }
+  scratch_size_recordings {
+    num_items: 10507528
+    scratch_space_bytes: 168121087
+  }
+  scratch_size_recordings {
+    num_items: 10717678
+    scratch_space_bytes: 171483391
+  }
+  scratch_size_recordings {
+    num_items: 10932031
+    scratch_space_bytes: 174912767
+  }
+  scratch_size_recordings {
+    num_items: 11150671
+    scratch_space_bytes: 178411263
+  }
+  scratch_size_recordings {
+    num_items: 11373684
+    scratch_space_bytes: 181979391
+  }
+  scratch_size_recordings {
+    num_items: 11601157
+    scratch_space_bytes: 185619199
+  }
+  scratch_size_recordings {
+    num_items: 11833180
+    scratch_space_bytes: 189331199
+  }
+  scratch_size_recordings {
+    num_items: 12069843
+    scratch_space_bytes: 193117951
+  }
+  scratch_size_recordings {
+    num_items: 12311239
+    scratch_space_bytes: 196980479
+  }
+  scratch_size_recordings {
+    num_items: 12557463
+    scratch_space_bytes: 200919807
+  }
+  scratch_size_recordings {
+    num_items: 12808612
+    scratch_space_bytes: 204938495
+  }
+  scratch_size_recordings {
+    num_items: 13064784
+    scratch_space_bytes: 209037055
+  }
+  scratch_size_recordings {
+    num_items: 13326079
+    scratch_space_bytes: 213217535
+  }
+  scratch_size_recordings {
+    num_items: 13592600
+    scratch_space_bytes: 217481983
+  }
+  scratch_size_recordings {
+    num_items: 13864452
+    scratch_space_bytes: 221831935
+  }
+  scratch_size_recordings {
+    num_items: 14141741
+    scratch_space_bytes: 226268415
+  }
+  scratch_size_recordings {
+    num_items: 14424575
+    scratch_space_bytes: 230793471
+  }
+  scratch_size_recordings {
+    num_items: 14713066
+    scratch_space_bytes: 235409663
+  }
+  scratch_size_recordings {
+    num_items: 15007327
+    scratch_space_bytes: 240117503
+  }
+  scratch_size_recordings {
+    num_items: 15307473
+    scratch_space_bytes: 244920063
+  }
+  scratch_size_recordings {
+    num_items: 15613622
+    scratch_space_bytes: 249818367
+  }
+  scratch_size_recordings {
+    num_items: 15925894
+    scratch_space_bytes: 254814975
+  }
+  scratch_size_recordings {
+    num_items: 16244411
+    scratch_space_bytes: 259910911
+  }
+  scratch_size_recordings {
+    num_items: 16569299
+    scratch_space_bytes: 265109247
+  }
+  scratch_size_recordings {
+    num_items: 16900684
+    scratch_space_bytes: 270411519
+  }
+  scratch_size_recordings {
+    num_items: 17238697
+    scratch_space_bytes: 275819775
+  }
+  scratch_size_recordings {
+    num_items: 17583470
+    scratch_space_bytes: 281336063
+  }
+  scratch_size_recordings {
+    num_items: 17935139
+    scratch_space_bytes: 286962943
+  }
+  scratch_size_recordings {
+    num_items: 18293841
+    scratch_space_bytes: 292701951
+  }
+  scratch_size_recordings {
+    num_items: 18659717
+    scratch_space_bytes: 298556159
+  }
+  scratch_size_recordings {
+    num_items: 19032911
+    scratch_space_bytes: 304527103
+  }
+  scratch_size_recordings {
+    num_items: 19413569
+    scratch_space_bytes: 310617855
+  }
+  scratch_size_recordings {
+    num_items: 19801840
+    scratch_space_bytes: 316829951
+  }
+  scratch_size_recordings {
+    num_items: 20197876
+    scratch_space_bytes: 323166463
+  }
+  scratch_size_recordings {
+    num_items: 20601833
+    scratch_space_bytes: 329629951
+  }
+  scratch_size_recordings {
+    num_items: 21013869
+    scratch_space_bytes: 336222463
+  }
+  scratch_size_recordings {
+    num_items: 21434146
+    scratch_space_bytes: 342947071
+  }
+  scratch_size_recordings {
+    num_items: 21862828
+    scratch_space_bytes: 349805823
+  }
+  scratch_size_recordings {
+    num_items: 22300084
+    scratch_space_bytes: 356801791
+  }
+  scratch_size_recordings {
+    num_items: 22746085
+    scratch_space_bytes: 363938047
+  }
+  scratch_size_recordings {
+    num_items: 23201006
+    scratch_space_bytes: 371216639
+  }
+  scratch_size_recordings {
+    num_items: 23665026
+    scratch_space_bytes: 378641151
+  }
+  scratch_size_recordings {
+    num_items: 24138326
+    scratch_space_bytes: 386213631
+  }
+  scratch_size_recordings {
+    num_items: 24621092
+    scratch_space_bytes: 393938175
+  }
+  scratch_size_recordings {
+    num_items: 25113513
+    scratch_space_bytes: 401816831
+  }
+  scratch_size_recordings {
+    num_items: 25615783
+    scratch_space_bytes: 409853183
+  }
+  scratch_size_recordings {
+    num_items: 26128098
+    scratch_space_bytes: 418050303
+  }
+  scratch_size_recordings {
+    num_items: 26650659
+    scratch_space_bytes: 426411263
+  }
+  scratch_size_recordings {
+    num_items: 27183672
+    scratch_space_bytes: 434939135
+  }
+  scratch_size_recordings {
+    num_items: 27727345
+    scratch_space_bytes: 443638015
+  }
+  scratch_size_recordings {
+    num_items: 28281891
+    scratch_space_bytes: 452510975
+  }
+  scratch_size_recordings {
+    num_items: 28847528
+    scratch_space_bytes: 461561087
+  }
+  scratch_size_recordings {
+    num_items: 29424478
+    scratch_space_bytes: 470791935
+  }
+  scratch_size_recordings {
+    num_items: 30012967
+    scratch_space_bytes: 480208127
+  }
+  scratch_size_recordings {
+    num_items: 30613226
+    scratch_space_bytes: 489812223
+  }
+  scratch_size_recordings {
+    num_items: 31225490
+    scratch_space_bytes: 499608319
+  }
+  scratch_size_recordings {
+    num_items: 31849999
+    scratch_space_bytes: 509600511
+  }
+  scratch_size_recordings {
+    num_items: 32486998
+    scratch_space_bytes: 519792383
+  }
+  scratch_size_recordings {
+    num_items: 33136737
+    scratch_space_bytes: 530188543
+  }
+  scratch_size_recordings {
+    num_items: 33799471
+    scratch_space_bytes: 540792063
+  }
+  scratch_size_recordings {
+    num_items: 34475460
+    scratch_space_bytes: 551608063
+  }
+  scratch_size_recordings {
+    num_items: 35164969
+    scratch_space_bytes: 562640127
+  }
+  scratch_size_recordings {
+    num_items: 35868268
+    scratch_space_bytes: 573892863
+  }
+  scratch_size_recordings {
+    num_items: 36585633
+    scratch_space_bytes: 585370879
+  }
+  scratch_size_recordings {
+    num_items: 37317345
+    scratch_space_bytes: 597078271
+  }
+  scratch_size_recordings {
+    num_items: 38063691
+    scratch_space_bytes: 609019647
+  }
+  scratch_size_recordings {
+    num_items: 38824964
+    scratch_space_bytes: 621200127
+  }
+  scratch_size_recordings {
+    num_items: 39601463
+    scratch_space_bytes: 633623807
+  }
+  scratch_size_recordings {
+    num_items: 40393492
+    scratch_space_bytes: 646296319
+  }
+  scratch_size_recordings {
+    num_items: 41201361
+    scratch_space_bytes: 659222271
+  }
+  scratch_size_recordings {
+    num_items: 42025388
+    scratch_space_bytes: 672406783
+  }
+  scratch_size_recordings {
+    num_items: 42865895
+    scratch_space_bytes: 685854975
+  }
+  scratch_size_recordings {
+    num_items: 43723212
+    scratch_space_bytes: 699571967
+  }
+  scratch_size_recordings {
+    num_items: 44597676
+    scratch_space_bytes: 713563391
+  }
+  scratch_size_recordings {
+    num_items: 45489629
+    scratch_space_bytes: 727834367
+  }
+  scratch_size_recordings {
+    num_items: 46399421
+    scratch_space_bytes: 742391039
+  }
+  scratch_size_recordings {
+    num_items: 47327409
+    scratch_space_bytes: 757239039
+  }
+  scratch_size_recordings {
+    num_items: 48273957
+    scratch_space_bytes: 772383999
+  }
+  scratch_size_recordings {
+    num_items: 49239436
+    scratch_space_bytes: 787831551
+  }
+  scratch_size_recordings {
+    num_items: 50224224
+    scratch_space_bytes: 803587839
+  }
+  scratch_size_recordings {
+    num_items: 51228708
+    scratch_space_bytes: 819660031
+  }
+  scratch_size_recordings {
+    num_items: 52253282
+    scratch_space_bytes: 836053247
+  }
+  scratch_size_recordings {
+    num_items: 53298347
+    scratch_space_bytes: 852774143
+  }
+  scratch_size_recordings {
+    num_items: 54364313
+    scratch_space_bytes: 869829375
+  }
+  scratch_size_recordings {
+    num_items: 55451599
+    scratch_space_bytes: 887226111
+  }
+  scratch_size_recordings {
+    num_items: 56560630
+    scratch_space_bytes: 904970495
+  }
+  scratch_size_recordings {
+    num_items: 57691842
+    scratch_space_bytes: 923070207
+  }
+  scratch_size_recordings {
+    num_items: 58845678
+    scratch_space_bytes: 941531391
+  }
+  scratch_size_recordings {
+    num_items: 60022591
+    scratch_space_bytes: 960361727
+  }
+  scratch_size_recordings {
+    num_items: 61223042
+    scratch_space_bytes: 979569407
+  }
+  scratch_size_recordings {
+    num_items: 62447502
+    scratch_space_bytes: 999160575
+  }
+  scratch_size_recordings {
+    num_items: 63696452
+    scratch_space_bytes: 1019143935
+  }
+  scratch_size_recordings {
+    num_items: 64970381
+    scratch_space_bytes: 1039526655
+  }
+  scratch_size_recordings {
+    num_items: 66269788
+    scratch_space_bytes: 1060316927
+  }
+  scratch_size_recordings {
+    num_items: 67595183
+    scratch_space_bytes: 1081523455
+  }
+  scratch_size_recordings {
+    num_items: 68947086
+    scratch_space_bytes: 1103153919
+  }
+  scratch_size_recordings {
+    num_items: 70326027
+    scratch_space_bytes: 1125217023
+  }
+  scratch_size_recordings {
+    num_items: 71732547
+    scratch_space_bytes: 1147721471
+  }
+  scratch_size_recordings {
+    num_items: 73167197
+    scratch_space_bytes: 1170675455
+  }
+  scratch_size_recordings {
+    num_items: 74630540
+    scratch_space_bytes: 1194089215
+  }
+  scratch_size_recordings {
+    num_items: 76123150
+    scratch_space_bytes: 1217970943
+  }
+  scratch_size_recordings {
+    num_items: 77645613
+    scratch_space_bytes: 1242330367
+  }
+  scratch_size_recordings {
+    num_items: 79198525
+    scratch_space_bytes: 1267176703
+  }
+  scratch_size_recordings {
+    num_items: 80782495
+    scratch_space_bytes: 1292520191
+  }
+  scratch_size_recordings {
+    num_items: 82398144
+    scratch_space_bytes: 1318370559
+  }
+  scratch_size_recordings {
+    num_items: 84046106
+    scratch_space_bytes: 1344738047
+  }
+  scratch_size_recordings {
+    num_items: 85727028
+    scratch_space_bytes: 1371632895
+  }
+  scratch_size_recordings {
+    num_items: 87441568
+    scratch_space_bytes: 1399065343
+  }
+  scratch_size_recordings {
+    num_items: 89190399
+    scratch_space_bytes: 1427046655
+  }
+  scratch_size_recordings {
+    num_items: 90974206
+    scratch_space_bytes: 1455587583
+  }
+  scratch_size_recordings {
+    num_items: 92793690
+    scratch_space_bytes: 1484699391
+  }
+  scratch_size_recordings {
+    num_items: 94649563
+    scratch_space_bytes: 1514393343
+  }
+  scratch_size_recordings {
+    num_items: 96542554
+    scratch_space_bytes: 1544681215
+  }
+  scratch_size_recordings {
+    num_items: 98473405
+    scratch_space_bytes: 1575574783
+  }
+  scratch_size_recordings {
+    num_items: 100442873
+    scratch_space_bytes: 1607086335
+  }
+  scratch_size_recordings {
+    num_items: 102451730
+    scratch_space_bytes: 1639228159
+  }
+  scratch_size_recordings {
+    num_items: 104500764
+    scratch_space_bytes: 1672012543
+  }
+  scratch_size_recordings {
+    num_items: 106590779
+    scratch_space_bytes: 1705452799
+  }
+  scratch_size_recordings {
+    num_items: 108722594
+    scratch_space_bytes: 1739562239
+  }
+  scratch_size_recordings {
+    num_items: 110897045
+    scratch_space_bytes: 1774353151
+  }
+  scratch_size_recordings {
+    num_items: 113114985
+    scratch_space_bytes: 1809840383
+  }
+  scratch_size_recordings {
+    num_items: 115377284
+    scratch_space_bytes: 1846037247
+  }
+  scratch_size_recordings {
+    num_items: 117684829
+    scratch_space_bytes: 1882957567
+  }
+  scratch_size_recordings {
+    num_items: 120038525
+    scratch_space_bytes: 1920616703
+  }
+  scratch_size_recordings {
+    num_items: 122439295
+    scratch_space_bytes: 1959028991
+  }
+  scratch_size_recordings {
+    num_items: 124888080
+    scratch_space_bytes: 1998209791
+  }
+  scratch_size_recordings {
+    num_items: 127385841
+    scratch_space_bytes: 2038173951
+  }
+  scratch_size_recordings {
+    num_items: 129933557
+    scratch_space_bytes: 2078937343
+  }
+  scratch_size_recordings {
+    num_items: 132532228
+    scratch_space_bytes: 2120516351
+  }
+  scratch_size_recordings {
+    num_items: 135182872
+    scratch_space_bytes: 2162926335
+  }
+  scratch_size_recordings {
+    num_items: 137886529
+    scratch_space_bytes: 2206185215
+  }
+  scratch_size_recordings {
+    num_items: 140644259
+    scratch_space_bytes: 2250308863
+  }
+  scratch_size_recordings {
+    num_items: 143457144
+    scratch_space_bytes: 2295314687
+  }
+  scratch_size_recordings {
+    num_items: 146326286
+    scratch_space_bytes: 2341221119
+  }
+  scratch_size_recordings {
+    num_items: 149252811
+    scratch_space_bytes: 2388045567
+  }
+  scratch_size_recordings {
+    num_items: 152237867
+    scratch_space_bytes: 2435806463
+  }
+  scratch_size_recordings {
+    num_items: 155282624
+    scratch_space_bytes: 2484522239
+  }
+  scratch_size_recordings {
+    num_items: 158388276
+    scratch_space_bytes: 2534212863
+  }
+  scratch_size_recordings {
+    num_items: 161556041
+    scratch_space_bytes: 2584897279
+  }
+  scratch_size_recordings {
+    num_items: 164787161
+    scratch_space_bytes: 2636594943
+  }
+  scratch_size_recordings {
+    num_items: 168082904
+    scratch_space_bytes: 2689326847
+  }
+  scratch_size_recordings {
+    num_items: 171444562
+    scratch_space_bytes: 2743113471
+  }
+  scratch_size_recordings {
+    num_items: 174873453
+    scratch_space_bytes: 2797975807
+  }
+  scratch_size_recordings {
+    num_items: 178370922
+    scratch_space_bytes: 2853935359
+  }
+  scratch_size_recordings {
+    num_items: 181938340
+    scratch_space_bytes: 2911014143
+  }
+  scratch_size_recordings {
+    num_items: 185577106
+    scratch_space_bytes: 2969234175
+  }
+  scratch_size_recordings {
+    num_items: 189288648
+    scratch_space_bytes: 3028619007
+  }
+  scratch_size_recordings {
+    num_items: 193074420
+    scratch_space_bytes: 3089191167
+  }
+  scratch_size_recordings {
+    num_items: 196935908
+    scratch_space_bytes: 3150975231
+  }
+  scratch_size_recordings {
+    num_items: 200874626
+    scratch_space_bytes: 3213994751
+  }
+  scratch_size_recordings {
+    num_items: 204892118
+    scratch_space_bytes: 3278274303
+  }
+  scratch_size_recordings {
+    num_items: 208989960
+    scratch_space_bytes: 3343839999
+  }
+  scratch_size_recordings {
+    num_items: 213169759
+    scratch_space_bytes: 3410716415
+  }
+  scratch_size_recordings {
+    num_items: 217433154
+    scratch_space_bytes: 3478931199
+  }
+  scratch_size_recordings {
+    num_items: 221781817
+    scratch_space_bytes: 3548509439
+  }
+  scratch_size_recordings {
+    num_items: 226217453
+    scratch_space_bytes: 3619479807
+  }
+  scratch_size_recordings {
+    num_items: 230741802
+    scratch_space_bytes: 3691869439
+  }
+  scratch_size_recordings {
+    num_items: 235356638
+    scratch_space_bytes: 3765706495
+  }
+  scratch_size_recordings {
+    num_items: 240063770
+    scratch_space_bytes: 3841020671
+  }
+  scratch_size_recordings {
+    num_items: 244865045
+    scratch_space_bytes: 3917841151
+  }
+  scratch_size_recordings {
+    num_items: 249762345
+    scratch_space_bytes: 3996198143
+  }
+  scratch_size_recordings {
+    num_items: 254757591
+    scratch_space_bytes: 4076121855
+  }
+  scratch_size_recordings {
+    num_items: 259852742
+    scratch_space_bytes: 4157644543
+  }
+  scratch_size_recordings {
+    num_items: 265049796
+    scratch_space_bytes: 4240797439
+  }
+  scratch_size_recordings {
+    num_items: 270350791
+    scratch_space_bytes: 4325613311
+  }
+  scratch_size_recordings {
+    num_items: 275757806
+    scratch_space_bytes: 4412125439
+  }
+  scratch_size_recordings {
+    num_items: 281272962
+    scratch_space_bytes: 4500368127
+  }
+  scratch_size_recordings {
+    num_items: 286898421
+    scratch_space_bytes: 4590375167
+  }
+  scratch_size_recordings {
+    num_items: 292636389
+    scratch_space_bytes: 4682182911
+  }
+  scratch_size_recordings {
+    num_items: 298489116
+    scratch_space_bytes: 4775826175
+  }
+  scratch_size_recordings {
+    num_items: 304458898
+    scratch_space_bytes: 4871342847
+  }
+  scratch_size_recordings {
+    num_items: 310548075
+    scratch_space_bytes: 4968769791
+  }
+  scratch_size_recordings {
+    num_items: 316759036
+    scratch_space_bytes: 5068144895
+  }
+  scratch_size_recordings {
+    num_items: 323094216
+    scratch_space_bytes: 5169508095
+  }
+  scratch_size_recordings {
+    num_items: 329556100
+    scratch_space_bytes: 5272898303
+  }
+  scratch_size_recordings {
+    num_items: 336147222
+    scratch_space_bytes: 5378355967
+  }
+  scratch_size_recordings {
+    num_items: 342870166
+    scratch_space_bytes: 5485923071
+  }
+  scratch_size_recordings {
+    num_items: 349727569
+    scratch_space_bytes: 5595641599
+  }
+  scratch_size_recordings {
+    num_items: 356722120
+    scratch_space_bytes: 5707554559
+  }
+  scratch_size_recordings {
+    num_items: 363856562
+    scratch_space_bytes: 5821705471
+  }
+  scratch_size_recordings {
+    num_items: 371133693
+    scratch_space_bytes: 5938139391
+  }
+  scratch_size_recordings {
+    num_items: 378556366
+    scratch_space_bytes: 6056902399
+  }
+  scratch_size_recordings {
+    num_items: 386127493
+    scratch_space_bytes: 6178040575
+  }
+  scratch_size_recordings {
+    num_items: 393850042
+    scratch_space_bytes: 6301601023
+  }
+  scratch_size_recordings {
+    num_items: 401727042
+    scratch_space_bytes: 6427633407
+  }
+  scratch_size_recordings {
+    num_items: 409761582
+    scratch_space_bytes: 6556185855
+  }
+  scratch_size_recordings {
+    num_items: 417956813
+    scratch_space_bytes: 6687309567
+  }
+  scratch_size_recordings {
+    num_items: 426315949
+    scratch_space_bytes: 6821055743
+  }
+  scratch_size_recordings {
+    num_items: 434842267
+    scratch_space_bytes: 6957476607
+  }
+  scratch_size_recordings {
+    num_items: 443539112
+    scratch_space_bytes: 7096626431
+  }
+  scratch_size_recordings {
+    num_items: 452409894
+    scratch_space_bytes: 7238558975
+  }
+  scratch_size_recordings {
+    num_items: 461458091
+    scratch_space_bytes: 7383330047
+  }
+  scratch_size_recordings {
+    num_items: 470687252
+    scratch_space_bytes: 7530996479
+  }
+  scratch_size_recordings {
+    num_items: 480100997
+    scratch_space_bytes: 7681616639
+  }
+  scratch_size_recordings {
+    num_items: 489703016
+    scratch_space_bytes: 7835248895
+  }
+  scratch_size_recordings {
+    num_items: 499497076
+    scratch_space_bytes: 7991953663
+  }
+  scratch_size_recordings {
+    num_items: 509487017
+    scratch_space_bytes: 8151792895
+  }
+  scratch_size_recordings {
+    num_items: 519676757
+    scratch_space_bytes: 8314828543
+  }
+  scratch_size_recordings {
+    num_items: 530070292
+    scratch_space_bytes: 8481125119
+  }
+  scratch_size_recordings {
+    num_items: 540671697
+    scratch_space_bytes: 8650747647
+  }
+  scratch_size_recordings {
+    num_items: 551485130
+    scratch_space_bytes: 8823762687
+  }
+  scratch_size_recordings {
+    num_items: 562514832
+    scratch_space_bytes: 9000237823
+  }
+  scratch_size_recordings {
+    num_items: 573765128
+    scratch_space_bytes: 9180242687
+  }
+  scratch_size_recordings {
+    num_items: 585240430
+    scratch_space_bytes: 9363847423
+  }
+  scratch_size_recordings {
+    num_items: 596945238
+    scratch_space_bytes: 9551124223
+  }
+  scratch_size_recordings {
+    num_items: 608884142
+    scratch_space_bytes: 9742146815
+  }
+  scratch_size_recordings {
+    num_items: 621061824
+    scratch_space_bytes: 9936989439
+  }
+  scratch_size_recordings {
+    num_items: 633483060
+    scratch_space_bytes: 10135729407
+  }
+  scratch_size_recordings {
+    num_items: 646152721
+    scratch_space_bytes: 10338444031
+  }
+  scratch_size_recordings {
+    num_items: 659075775
+    scratch_space_bytes: 10545212671
+  }
+  scratch_size_recordings {
+    num_items: 672257290
+    scratch_space_bytes: 10756117247
+  }
+  scratch_size_recordings {
+    num_items: 685702435
+    scratch_space_bytes: 10971239679
+  }
+  scratch_size_recordings {
+    num_items: 699416483
+    scratch_space_bytes: 11190664447
+  }
+  scratch_size_recordings {
+    num_items: 713404812
+    scratch_space_bytes: 11414477567
+  }
+  scratch_size_recordings {
+    num_items: 727672908
+    scratch_space_bytes: 11642767103
+  }
+  scratch_size_recordings {
+    num_items: 742226366
+    scratch_space_bytes: 11875622143
+  }
+  scratch_size_recordings {
+    num_items: 757070893
+    scratch_space_bytes: 12113134847
+  }
+  scratch_size_recordings {
+    num_items: 772212310
+    scratch_space_bytes: 12355397375
+  }
+  scratch_size_recordings {
+    num_items: 787656556
+    scratch_space_bytes: 12602505471
+  }
+  scratch_size_recordings {
+    num_items: 803409687
+    scratch_space_bytes: 12854555391
+  }
+  scratch_size_recordings {
+    num_items: 819477880
+    scratch_space_bytes: 13111646463
+  }
+  scratch_size_recordings {
+    num_items: 835867437
+    scratch_space_bytes: 13373879551
+  }
+  scratch_size_recordings {
+    num_items: 852584785
+    scratch_space_bytes: 13641357055
+  }
+  scratch_size_recordings {
+    num_items: 869636480
+    scratch_space_bytes: 13914183935
+  }
+  scratch_size_recordings {
+    num_items: 887029209
+    scratch_space_bytes: 14192467711
+  }
+  scratch_size_recordings {
+    num_items: 904769793
+    scratch_space_bytes: 14476317439
+  }
+  scratch_size_recordings {
+    num_items: 922865188
+    scratch_space_bytes: 14765843711
+  }
+  scratch_size_recordings {
+    num_items: 941322491
+    scratch_space_bytes: 15061160191
+  }
+  scratch_size_recordings {
+    num_items: 960148940
+    scratch_space_bytes: 15362383615
+  }
+  scratch_size_recordings {
+    num_items: 979351918
+    scratch_space_bytes: 15669631231
+  }
+  scratch_size_recordings {
+    num_items: 998938956
+    scratch_space_bytes: 15983023871
+  }
+}
diff --git a/third_party/xla/xla/backends/gpu/target_config/specs/vr_nvl72.txtpb b/third_party/xla/xla/backends/gpu/target_config/specs/vr_nvl72.txtpb
index 5c831faa7c5..c810ef8b415 100644
--- a/third_party/xla/xla/backends/gpu/target_config/specs/vr_nvl72.txtpb
+++ b/third_party/xla/xla/backends/gpu/target_config/specs/vr_nvl72.txtpb
@@ -248,6 +248,7 @@ gpu_device_info {
   reserved_shared_memory_per_block: 1024
   max_blocks_per_multiprocessor: 16
   collective_memory_granularity: 2097152
+  oversized_shared_memory_per_block: 334848
 }
 platform_name: "CUDA"
 dnn_version_info {
diff --git a/third_party/xla/xla/service/gpu/gpu_device_info_for_tests.cc b/third_party/xla/xla/service/gpu/gpu_device_info_for_tests.cc
index 8f30470ca7b..9666ee9b39d 100644
--- a/third_party/xla/xla/service/gpu/gpu_device_info_for_tests.cc
+++ b/third_party/xla/xla/service/gpu/gpu_device_info_for_tests.cc
@@ -156,6 +156,7 @@ stream_executor::DeviceDescription TestGpuDeviceInfo::VRNVL72DeviceInfo(
   b.set_threads_per_warp(32);
   b.set_shared_memory_per_block(48 * 1024);
   b.set_shared_memory_per_block_optin(227 * 1024);
+  b.set_oversized_shared_memory_per_block(327 * 1024);
   b.set_shared_memory_per_core(228 * 1024);
   b.set_threads_per_core_limit(1024);
   b.set_core_count(212);
diff --git a/third_party/xla/xla/service/gpu/model/default_collective_perf_table.txtpb b/third_party/xla/xla/service/gpu/model/default_collective_perf_table.txtpb
index 06474f0dc13..de492e32078 100644
--- a/third_party/xla/xla/service/gpu/model/default_collective_perf_table.txtpb
+++ b/third_party/xla/xla/service/gpu/model/default_collective_perf_table.txtpb
@@ -110015,3 +110015,41888 @@ entries {
     }
   }
 }
+entries {
+  key: "sm_107"
+  value {
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7d245427c4bb73f837b88bd0b52224c5"
+      network_throughput_bytes_per_sec: 813273192190
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "92e5852d9f4fa1b5b1585a650f791b7a"
+      network_throughput_bytes_per_sec: 1015229544
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "52844b1253f994022bb07ea0239835f8"
+      network_throughput_bytes_per_sec: 338892578677
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fdb6e897cb8d84a74772188b0b911ee5"
+      network_throughput_bytes_per_sec: 2398081
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "52d89d3a8dc7e60b60026bf7d3879e6b"
+      network_throughput_bytes_per_sec: 22029936038
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d2dfa0b087803a9675b5fc36f0424192"
+      network_throughput_bytes_per_sec: 7877461
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4831838208
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "66001c550718ade61b0ff477bb2a6f8c"
+      network_throughput_bytes_per_sec: 838343857674
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "c85cd0792c4a505c7e4fcaf2667e67f1"
+      network_throughput_bytes_per_sec: 95907090
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2af4570b4a5f461e6746d506091197ea"
+      network_throughput_bytes_per_sec: 8856216216
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "286019cc62525879d568af1a23ed7a0a"
+      network_throughput_bytes_per_sec: 4813160987
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "a989b32b0ff39b4efcfb31066398964f"
+      network_throughput_bytes_per_sec: 11407705399
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0a793c4feebd363812c2ba136fa90288"
+      network_throughput_bytes_per_sec: 301277066460
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "bb6520f235fd6129f12d2fa03f56fecc"
+      network_throughput_bytes_per_sec: 298542274
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "db6c35e32b462a2e1bd5f2dfc8b0fa8d"
+      network_throughput_bytes_per_sec: 34841519
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0f9a6832b99ca7b92783645e040adae9"
+      network_throughput_bytes_per_sec: 3169415
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "077ece68ae856093c873f0c30599f3a1"
+      network_throughput_bytes_per_sec: 1815804056
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "12be67e92fa9aa0cf1c190e32d280536"
+      network_throughput_bytes_per_sec: 55421564482
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "11ced3e1ca1464fb2be85a808b9754eb"
+      network_throughput_bytes_per_sec: 217455935
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ef8731ca93e8c875daba7bca702659df"
+      network_throughput_bytes_per_sec: 336861824022
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3a97aee00a04651d29d5457968ba6682"
+      network_throughput_bytes_per_sec: 2341166010
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ee08073618464f3b8586d329cdb6e226"
+      network_throughput_bytes_per_sec: 846809192183
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a3486caeb11575758c3a094d842ab70e"
+      network_throughput_bytes_per_sec: 8590503932
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b8ebd4ad42977a7f559a48fbe20d382e"
+      network_throughput_bytes_per_sec: 44544435004
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7f8794aa8c1a0102b821efc5646f6ac1"
+      network_throughput_bytes_per_sec: 988452032286
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8a9ebcc6c1e72e9240b5f09ea3d69ab5"
+      network_throughput_bytes_per_sec: 3774342172
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2c833d8dde815c32db190dc164bb1b3e"
+      network_throughput_bytes_per_sec: 2606058463
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "49e024453b7259c5d0613b9e91202956"
+      network_throughput_bytes_per_sec: 41682570249
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6c7654fbf0ef9b84db9d18fcba5a6f43"
+      network_throughput_bytes_per_sec: 545454545
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a5958019093dce310babb62c8e0ab33d"
+      network_throughput_bytes_per_sec: 57442929489
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d1f49366953921c20c991470340a1280"
+      network_throughput_bytes_per_sec: 202271604938
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "cb8ec9a72ca6036b8eaecd91c59bf5ee"
+      network_throughput_bytes_per_sec: 1329762194
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a1a1a6301dcd0ce663ee06a82fb59c2d"
+      network_throughput_bytes_per_sec: 832643726839
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4294967296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f8ef6c3e39c12c2305cefc10d6d5ce4e"
+      network_throughput_bytes_per_sec: 452983971153
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "2a0ef6ccfa56fbbc9a482e6d30d0d031"
+      network_throughput_bytes_per_sec: 3293928427
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "bc6fdc399317ea94c738f6f5b6be2e36"
+      network_throughput_bytes_per_sec: 10240819
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b7eb8be70eed4356d1c3929ecd578c9a"
+      network_throughput_bytes_per_sec: 211832850
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "378d0004c89b2e5f366ff7f6c202f74a"
+      network_throughput_bytes_per_sec: 16973291462
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "0e4090a9e471c44cebf8acc7ef75de18"
+      network_throughput_bytes_per_sec: 62806673
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7b70d21e5a86ab7f6868f87401207f35"
+      network_throughput_bytes_per_sec: 672897196
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b79463ba29f0adc5054211dc799b6948"
+      network_throughput_bytes_per_sec: 1349291753
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f1e40671963bc161628d97e0b7b749e8"
+      network_throughput_bytes_per_sec: 29729125910
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e7d6b1dc87cdbd01297fbb3d07d86e62"
+      network_throughput_bytes_per_sec: 1576327717
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "7e2f9095c85c3a562f564f8f1812d396"
+      network_throughput_bytes_per_sec: 43120665371
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2fcea516f510263a91388a571cbd9c04"
+      network_throughput_bytes_per_sec: 853167350000
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "268f2bd44d20599217136d8d4fd3a59e"
+      network_throughput_bytes_per_sec: 29105551443
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "051078de7c361c47aba866722ca94e28"
+      network_throughput_bytes_per_sec: 497277750726
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "43df6488ac548def3a50f5361bd71a36"
+      network_throughput_bytes_per_sec: 838703624974
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ee68e167f8be71f8d52644c2d7be1a9e"
+      network_throughput_bytes_per_sec: 344869593816
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6a6f130e7a7f23149e7401c7a53c7b8e"
+      network_throughput_bytes_per_sec: 436673773
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "236923961c824e97a2db9d72a37b11a8"
+      network_throughput_bytes_per_sec: 6938087
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f60d25394347c894986103c1a85b8669"
+      network_throughput_bytes_per_sec: 31583614457
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "93f72f2c6d8cc0e071ebd967dc2d3c68"
+      network_throughput_bytes_per_sec: 280172675
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4831838208
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "465da4f0625fa3adff66bf639a3aa176"
+      network_throughput_bytes_per_sec: 471579015748
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 8589934592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "4a8f9666bcdd70a941bad496bc374dbc"
+      network_throughput_bytes_per_sec: 672801129123
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "db03075746ca5ffeba5e030d2fa104a8"
+      network_throughput_bytes_per_sec: 163135255
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7ad74bbbaacfb30618c46aeeb2344210"
+      network_throughput_bytes_per_sec: 66976744
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "bc417a3b6426f73e98100cacb34ff908"
+      network_throughput_bytes_per_sec: 316103266649
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 4831838208
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "29fb26490770be592242dbd9da9e86e1"
+      network_throughput_bytes_per_sec: 499233684523
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "37d25884db4744d18aecbf7785c41e7d"
+      network_throughput_bytes_per_sec: 803318422676
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "32e81cb74181f9f5a025f16621345bfc"
+      network_throughput_bytes_per_sec: 124319876696
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8b85fb82f371d8a3eb331753f59a9bbd"
+      network_throughput_bytes_per_sec: 209175990779
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a4ad197df4d99768f8bc32e1751e4597"
+      network_throughput_bytes_per_sec: 927760912308
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2415919104
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "08b220405980b6b4a64d5353aa4e8477"
+      network_throughput_bytes_per_sec: 1019972500350
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "aea282c3aa662c299080ad2922c88047"
+      network_throughput_bytes_per_sec: 16537467
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b3ad737b92ffb5acd643541409459841"
+      network_throughput_bytes_per_sec: 48901623
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "64b4d0372168e80016afc7ff0efcc1f0"
+      network_throughput_bytes_per_sec: 4481536637
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "280b7268f4d08dd2291b78e159178d93"
+      network_throughput_bytes_per_sec: 457671387002
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "ba45316f3f78c4f286eb58ab196a47d4"
+      network_throughput_bytes_per_sec: 571162963
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "88f147819b2e34a5b9e745ac655f4eaa"
+      network_throughput_bytes_per_sec: 77732793
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "700b351df50bc17e7b2c683bc111b046"
+      network_throughput_bytes_per_sec: 204489360780
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "817daa842b3456bcc26eccd2f4f2cef7"
+      network_throughput_bytes_per_sec: 94126298
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ec33b380012bd504fec002091d1c7177"
+      network_throughput_bytes_per_sec: 112176834
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "16f028ede633bfcce74582b97f388b9c"
+      network_throughput_bytes_per_sec: 61826415094
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "180228a6f07939de0556cab553b95fa1"
+      network_throughput_bytes_per_sec: 31049905243
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1178677d1605f0a58a4d9b9da6d70a1c"
+      network_throughput_bytes_per_sec: 170514025530
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d7dbe32cdcb60969245e8ca733c61fec"
+      network_throughput_bytes_per_sec: 47462096
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "135285d887d0e6b957a1df83a0da95ba"
+      network_throughput_bytes_per_sec: 382955155565
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b1bd9127d36b6258461e0cf0e4a4725c"
+      network_throughput_bytes_per_sec: 12349384
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "8de894443e7b34abdc9c2962e008e0ad"
+      network_throughput_bytes_per_sec: 20912129552
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "72436908d27256d06018a0a63078f8c4"
+      network_throughput_bytes_per_sec: 409280249804
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 128
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2d8269dc7c35329ee43ab94e6959a309"
+      network_throughput_bytes_per_sec: 71910112
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "559eb9ae6b74d21a007ed185074ac61a"
+      network_throughput_bytes_per_sec: 677149155
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "8e1b10ed525f4f91f3f4dbf4ee962e6d"
+      network_throughput_bytes_per_sec: 761146668044
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fddd4ee9cc2e87614656e75b0ec4528a"
+      network_throughput_bytes_per_sec: 440969300499
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "babe4962eb62ecdb1d803fa1dfb33642"
+      network_throughput_bytes_per_sec: 281100730514
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b7f1255682f19ae5e9e013c28176bf4f"
+      network_throughput_bytes_per_sec: 1129978598910
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7862be5ce82b923ecd1a79023b1d5d1a"
+      network_throughput_bytes_per_sec: 159402241
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4831838208
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0059501bcc17c7e8c598360ab52db735"
+      network_throughput_bytes_per_sec: 861021988737
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "797612dcc4d69679f826671c1816d25c"
+      network_throughput_bytes_per_sec: 1157637068626
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "78154c1887c361326d95e7895b6bd2b9"
+      network_throughput_bytes_per_sec: 1268928482483
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fb5ce99db69a4d23d3105a0354a341b9"
+      network_throughput_bytes_per_sec: 739040996123
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e196217412421f80426e2ef7c356a925"
+      network_throughput_bytes_per_sec: 332241836148
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8b27468036b79d0ad7307240aa5b84d2"
+      network_throughput_bytes_per_sec: 402185576236
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "83b17a215a4489bce4d36f4934a2cc2c"
+      network_throughput_bytes_per_sec: 1140841561267
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4294967296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "908577d008491898aaff39e6e46601c9"
+      network_throughput_bytes_per_sec: 1046976970170
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "bbda0a1c70daca8761087c5c5bf00783"
+      network_throughput_bytes_per_sec: 37572595671
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4831838208
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "48ec4761ce1f3338d9300588ae92b759"
+      network_throughput_bytes_per_sec: 387554450175
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "bbb59e05d4fa6fce1ab84be4ae0d0af2"
+      network_throughput_bytes_per_sec: 253744493
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "148c37f6c1e5e4d9f73218d58dac6e5f"
+      network_throughput_bytes_per_sec: 321544922
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "20ae1799763095780e29375223a21208"
+      network_throughput_bytes_per_sec: 208197670314
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "42fcf9ee06ca093d70b9a723e326f0dd"
+      network_throughput_bytes_per_sec: 88402877697
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "65f7eeb9cf6be0bdf4e5d8c39f9105d5"
+      network_throughput_bytes_per_sec: 24556616
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "ed8da9b9db7c5ec6393d977881bceb88"
+      network_throughput_bytes_per_sec: 30852166180
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9416a3f96f3a8fcad45e94afad8b4389"
+      network_throughput_bytes_per_sec: 103328340559
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "08dfcc4679c724891e2ee0b98897d35d"
+      network_throughput_bytes_per_sec: 153354632
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "22feda29f726ee55662cd322bbc79533"
+      network_throughput_bytes_per_sec: 17545930509
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "55807deaa73050d6b7c52bab61aa0b5b"
+      network_throughput_bytes_per_sec: 28972590627
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5b480c944711f77bd91e28cb9ae5cc2e"
+      network_throughput_bytes_per_sec: 11251263
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "95f198f71fd61aa7060d042a62c9a0b3"
+      network_throughput_bytes_per_sec: 7623484
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3be6dbcf4f679d96555d3e59b2030df8"
+      network_throughput_bytes_per_sec: 222657099
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c2400d02eec3e8fa9da21b3bda25c181"
+      network_throughput_bytes_per_sec: 21145950596
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fe149cecc0e0dc1b1866bbbaed14b886"
+      network_throughput_bytes_per_sec: 969132382460
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b64081dec7375c17ded6828c67fb0d50"
+      network_throughput_bytes_per_sec: 2162904290
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d0b7d4cbe6b259a0d8fb5613235c90f5"
+      network_throughput_bytes_per_sec: 793615
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5a3d245166a87846af51e32b7ff2f1ac"
+      network_throughput_bytes_per_sec: 133456904
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "29717226393409071cc6553a498f03d0"
+      network_throughput_bytes_per_sec: 419918044409
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "bb1c7ae518e42a98548b6cccc61f7dfa"
+      network_throughput_bytes_per_sec: 96086503657
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "84ca6f5306ee1f5b9e94f7508406715e"
+      network_throughput_bytes_per_sec: 828622706119
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "78c6190f8a1908d90498ec3135360634"
+      network_throughput_bytes_per_sec: 2645599
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "9ab4044f63157a3746302ba829ee16f3"
+      network_throughput_bytes_per_sec: 40915248946
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "580471ca0fbfb858a1f412e668dc6c9d"
+      network_throughput_bytes_per_sec: 128679139
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a4d0875e07def9e1d305ae2d667c4983"
+      network_throughput_bytes_per_sec: 1969170
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6cde973287343bfb99323311ff1015fd"
+      network_throughput_bytes_per_sec: 1031560677498
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d319d162588e409c7c7867ba1280e79f"
+      network_throughput_bytes_per_sec: 5885057471
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "be68ce14181efd15daeb8470c4fd5c38"
+      network_throughput_bytes_per_sec: 419724206944
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "52341478cac61811c8a7538ecfdf667e"
+      network_throughput_bytes_per_sec: 35667843
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f7cf0106b5cf6d35d6207da0598c3e42"
+      network_throughput_bytes_per_sec: 7932081
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "403c9918c27c798797da8da8e96e5b2b"
+      network_throughput_bytes_per_sec: 1043821971951
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f07b92af61671fa5f6f7401cb5f23aa1"
+      network_throughput_bytes_per_sec: 1242148605485
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d9082c02cc12c1d45f7144b359e77e44"
+      network_throughput_bytes_per_sec: 4702640642
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2d8a5837418d88d554ab0c3019aa0ee4"
+      network_throughput_bytes_per_sec: 1514023009
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f606d53e23b449af40022bc7252fcf19"
+      network_throughput_bytes_per_sec: 42145337620
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "83cf470bf9763758dccc2885ffa027a3"
+      network_throughput_bytes_per_sec: 289461973775
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "23bbba6762b18aefac37c26a6fc2daf1"
+      network_throughput_bytes_per_sec: 392562775
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "15b3717611172866367d45b49353ba87"
+      network_throughput_bytes_per_sec: 144431955922
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "661cf53a46c7ae5915018ac48c72e80f"
+      network_throughput_bytes_per_sec: 1002283051556
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c6a296ba7cf374ab945eea5036328e60"
+      network_throughput_bytes_per_sec: 1159682899
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 128
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "42cbe0d8c6314cb70629f0ca473683d7"
+      network_throughput_bytes_per_sec: 35705568
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "cff7fc7dedf54ceab87334c2b884833e"
+      network_throughput_bytes_per_sec: 87823704586
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "965fab391188bf1ffdf1e374b15a2ed7"
+      network_throughput_bytes_per_sec: 72961900049
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f0bddf3fed669618bb3f032c154daa7c"
+      network_throughput_bytes_per_sec: 1240352
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "62c5030d9e064721548499d2f5888c4b"
+      network_throughput_bytes_per_sec: 939671937718
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "55efb98d498b339d0d3ea3cadcb0af96"
+      network_throughput_bytes_per_sec: 2693555458
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0595689f5413666048eaf87e25d202e0"
+      network_throughput_bytes_per_sec: 42331359672
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8957e28dfcc800e9d866c33ff5383f98"
+      network_throughput_bytes_per_sec: 838551147518
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5f507b241a914b371f32d50a3c950c93"
+      network_throughput_bytes_per_sec: 16209745238
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a99b30ac40f654ab0ebb6e2ab70b98b7"
+      network_throughput_bytes_per_sec: 640177827
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c0b6d9aa70f8b4e907cde4e44072c6d5"
+      network_throughput_bytes_per_sec: 73361194029
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a57bab2fb76d1bf2e1ea71653fb89ee1"
+      network_throughput_bytes_per_sec: 325904077390
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1df1512c171996c15b3e669e44380043"
+      network_throughput_bytes_per_sec: 635586206896
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d93e907dbcee9da8457a78f694c02550"
+      network_throughput_bytes_per_sec: 15372410018
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f8603219a0f8d807cdd492dd42f64165"
+      network_throughput_bytes_per_sec: 722710163
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 4294967296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "10ff17ad1d8c7811815c9a6a3c192b3a"
+      network_throughput_bytes_per_sec: 798591956565
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "111c94cc5fdff6a32719885629d509ad"
+      network_throughput_bytes_per_sec: 677809035
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4c224b166b02a84de723473b8a03ec93"
+      network_throughput_bytes_per_sec: 335128760337
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c5fd652eeb298863873884edcc1e7fe7"
+      network_throughput_bytes_per_sec: 315005933158
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e6d40dbc6d3a52b25ee38b0fdad8faec"
+      network_throughput_bytes_per_sec: 7305565
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9c2f8241aedf9b76756dfb6239125ffe"
+      network_throughput_bytes_per_sec: 956021197789
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "48bb10b6fba1023364baaee71bb124d0"
+      network_throughput_bytes_per_sec: 456139492731
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "1134e6e324219a06169de90552fb3f62"
+      network_throughput_bytes_per_sec: 135513469
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ccc6a6e1ec26b6c3fe7ec4f3acb45551"
+      network_throughput_bytes_per_sec: 12714635625
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1a0c183f29cf18bd558022ea04ea2338"
+      network_throughput_bytes_per_sec: 854487235315
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0154bef373a9f3963b544013bda26128"
+      network_throughput_bytes_per_sec: 39633216162
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "7295cced7959392baf1ebdb1c357affc"
+      network_throughput_bytes_per_sec: 2084832032
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b9f767f0729d189594d088894c861a35"
+      network_throughput_bytes_per_sec: 359840768702
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4294967296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f1836786a256c916b7b805d07187392b"
+      network_throughput_bytes_per_sec: 465636867315
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "48375910d047f621d3832cf0d813978a"
+      network_throughput_bytes_per_sec: 615609877811
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "540d1bc88f70cedad71dad9d41d9c9b7"
+      network_throughput_bytes_per_sec: 321797145926
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c98f7fd06640efd41b87cb498ccfcdfa"
+      network_throughput_bytes_per_sec: 93771701112
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8b97a971d077ade12314e836a9f835d1"
+      network_throughput_bytes_per_sec: 46759241917
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5c425d241a0f4c9ac4d2b3084f4b589c"
+      network_throughput_bytes_per_sec: 5486968
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "5914346efb155ac058df7f18dd3c3aac"
+      network_throughput_bytes_per_sec: 549098465
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8152118a44187fbe20d5282a6e10f006"
+      network_throughput_bytes_per_sec: 196993779484
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "197dba26eac50e613dbff06a8fd2728b"
+      network_throughput_bytes_per_sec: 149304993
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b28a398539c5203dfee099dcbfd44a28"
+      network_throughput_bytes_per_sec: 1161161068798
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8a3dc98a8ad45042257d0715b090af60"
+      network_throughput_bytes_per_sec: 121365360
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "7679a89f295e6b5f52b5e137ab70c40f"
+      network_throughput_bytes_per_sec: 15317294067
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8853c6e32e8c990e1376f18567e842e3"
+      network_throughput_bytes_per_sec: 1018554624646
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "afca04dad0765f7be5a4bac7531c4fed"
+      network_throughput_bytes_per_sec: 2977174
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 2415919104
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "0d90b47de7b72c28274c7fc0f0d3e835"
+      network_throughput_bytes_per_sec: 670463898015
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fc66d3693ecba43c195accf4015d10b7"
+      network_throughput_bytes_per_sec: 2206152
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7222f249ca7b5e2c757c61ced30c8bde"
+      network_throughput_bytes_per_sec: 19539654144
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7ef70b1910bc84b1ab91d637f3075f4d"
+      network_throughput_bytes_per_sec: 249954232
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8b36f06a1872f9d3a229060af2c8d077"
+      network_throughput_bytes_per_sec: 30214845550
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9e8adcfc032993d0ad39074ec26a676d"
+      network_throughput_bytes_per_sec: 29568788
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "84ada5edf7ed6f5bf4e3463e163a1887"
+      network_throughput_bytes_per_sec: 51449096654
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2981c9f2a79e942090eded5ee8497aea"
+      network_throughput_bytes_per_sec: 9448673587
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5ad6446309530c9f5641f23b3a9877f5"
+      network_throughput_bytes_per_sec: 1097787767253
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "075ebe13e11add435abb105d36f6380b"
+      network_throughput_bytes_per_sec: 7061592
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1d3cbc3b67d186e5f015634206bf0ac0"
+      network_throughput_bytes_per_sec: 849404027491
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2415919104
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "cbb4576b2b0afd6093e7ed48c730c3c1"
+      network_throughput_bytes_per_sec: 921681657625
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4eefa4b7b76acdb01f2760d86449591b"
+      network_throughput_bytes_per_sec: 443185145435
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "8e464dbe5d02615512523c3478bb6f1a"
+      network_throughput_bytes_per_sec: 225228197
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2febeffbc8ca442875de625718b798c3"
+      network_throughput_bytes_per_sec: 87009063
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b8b690de350e69fd702836a6c6ee8dc1"
+      network_throughput_bytes_per_sec: 2421439831
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8fbc3894c848f63c50dafa8300d8bf0c"
+      network_throughput_bytes_per_sec: 5741455
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d134b6c7f9cf1c2eef19164ffd136a82"
+      network_throughput_bytes_per_sec: 181205529953
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "bc88a8c5268b89aa141351d264ee2644"
+      network_throughput_bytes_per_sec: 131898328
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d0e1f256495e0cdd4b1d669bb44426de"
+      network_throughput_bytes_per_sec: 3540190146
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8d5ea2e6e0ccd2167bcb20cd2a114c09"
+      network_throughput_bytes_per_sec: 138018954018
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e68b2e1a0fd6cd7ee5817a117f5442cf"
+      network_throughput_bytes_per_sec: 202780119899
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f6a7c0bcc37f7a310c2e1f1f070c7910"
+      network_throughput_bytes_per_sec: 128445993031
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8469777a96a050371c7bb8994b2e225a"
+      network_throughput_bytes_per_sec: 309348674263
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "17ce5344f96734d4e2929184201806d7"
+      network_throughput_bytes_per_sec: 266142373311
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "61c3c8ab2cf1b1ce3071c20dcf54897e"
+      network_throughput_bytes_per_sec: 1786825650
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "93c722a3b287bafde262103ec8cbeb13"
+      network_throughput_bytes_per_sec: 653376686977
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "ede64fe78c8e1412036657e3890747cb"
+      network_throughput_bytes_per_sec: 293590912570
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a9e76cea1d55825a8c38fd557f6533fc"
+      network_throughput_bytes_per_sec: 13557302441
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "d2fd7f54a3ecd90c3f3437449bdbf3cc"
+      network_throughput_bytes_per_sec: 6709865307
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5603f7261f334f3d2d9fcd6208d10d1c"
+      network_throughput_bytes_per_sec: 115107913
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "bc417c35738146c0a69b762725a7816b"
+      network_throughput_bytes_per_sec: 495903395769
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7285efb8962db83fd694e75c5a9704b5"
+      network_throughput_bytes_per_sec: 160520437
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ffd21359f616ef08c31a8331c390df6c"
+      network_throughput_bytes_per_sec: 12787789437
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3bc8c9f323b6d850bde618a845e4d1c8"
+      network_throughput_bytes_per_sec: 63327455006
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "8d3306d9982bb2e68c05322a13bbdeb8"
+      network_throughput_bytes_per_sec: 668722199317
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3632466150fa9c07ebfb64d4f40862d4"
+      network_throughput_bytes_per_sec: 293517790495
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "ba0314d074e3ba684fca125fa7baa5df"
+      network_throughput_bytes_per_sec: 81732341975
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "65d4137d6313d9ac26a649ef79883ad8"
+      network_throughput_bytes_per_sec: 20422635
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8f199b4f4dcafd781f26b7a2b0d2b520"
+      network_throughput_bytes_per_sec: 57793479758
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 64
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a10c2b456acdc09fdeb33fb14a51a0e7"
+      network_throughput_bytes_per_sec: 39776258
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7ab82c071f5ee995b79b32ecfa6d1ad8"
+      network_throughput_bytes_per_sec: 105834323
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "29348ea57254fd8b9bdfecfdf9d459fa"
+      network_throughput_bytes_per_sec: 491603614221
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a16977c8b747af700a7834e7737b30e5"
+      network_throughput_bytes_per_sec: 1430842
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0f0799eb6a308b5fedd5aa1d5a1b2e0c"
+      network_throughput_bytes_per_sec: 517907665290
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "26d8f0d30d8f47f9cda137dee9039d66"
+      network_throughput_bytes_per_sec: 178493957
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2b7fd0260c6b74d71e2c90ab4c15d857"
+      network_throughput_bytes_per_sec: 139443044
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "6e3c09380e1d980a449ff54881db1604"
+      network_throughput_bytes_per_sec: 250519877
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e54be67b5baa099350fc2cf108c2d5b8"
+      network_throughput_bytes_per_sec: 19899595141
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a9f53663ea05102aa9ccd739df92ef37"
+      network_throughput_bytes_per_sec: 1054107958
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0c9c7f770714605590801437c78b3dc5"
+      network_throughput_bytes_per_sec: 686142608837
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1e4effbdadfe2e3d8b54378ad9cd6516"
+      network_throughput_bytes_per_sec: 3095709
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d261d7cb839eae9a058360eb5181c3d8"
+      network_throughput_bytes_per_sec: 1066389012514
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1a71dfc235127520e47791c98a1e3838"
+      network_throughput_bytes_per_sec: 14900471778
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "89a06f01d70f298b1c6b13e1587bf887"
+      network_throughput_bytes_per_sec: 1494768
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 64
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "00059129db60bb6bbc1695af12f2dc9f"
+      network_throughput_bytes_per_sec: 43631405
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "36f62e12db9987fabde870a321d31277"
+      network_throughput_bytes_per_sec: 24525583639
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6f419362035e3ee16b9807102745a6c9"
+      network_throughput_bytes_per_sec: 59687572
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "760c613680e5b24d8ad633f6882095d5"
+      network_throughput_bytes_per_sec: 16120004
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "401b3d5221f15a1cc9549989032e4927"
+      network_throughput_bytes_per_sec: 40796090
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d70f3edf5f59b6f67b5d944133062dcb"
+      network_throughput_bytes_per_sec: 218234530483
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8bafa3a2478a8370a5c072625daa48db"
+      network_throughput_bytes_per_sec: 6881463505
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e043e91cc488e4f36922b87ac05994f3"
+      network_throughput_bytes_per_sec: 389444272785
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 4831838208
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "ce782d8fc6784b1e327b373c0aa6cb07"
+      network_throughput_bytes_per_sec: 671405412803
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4157f535c38555ea72900e96f4891dd4"
+      network_throughput_bytes_per_sec: 66646779661
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "bac9232052ba5fbea726ddea4e869d9c"
+      network_throughput_bytes_per_sec: 24791896095
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "76375ab3fce415b4c53e20cf2f911619"
+      network_throughput_bytes_per_sec: 689523
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "290c97dd59b3a0c25e94528270446766"
+      network_throughput_bytes_per_sec: 116375391
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0daf2a55465221e4c7e33525c21bf8b6"
+      network_throughput_bytes_per_sec: 11037020
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "619cde8ee647556999119b2833818c90"
+      network_throughput_bytes_per_sec: 37647058
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7b57a89f71492d05d768bf7a6dbacddf"
+      network_throughput_bytes_per_sec: 955095980872
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0516710ddb1402b66fbe4eb84bf35670"
+      network_throughput_bytes_per_sec: 168581350482
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4b77979e478aab4e4ce4c82d1aa29a35"
+      network_throughput_bytes_per_sec: 889661164
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "72644f2300af00ecbef802ddc45d18c4"
+      network_throughput_bytes_per_sec: 5051138772
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 8589934592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9e0348f6ad6f10f944bcc151f61e3610"
+      network_throughput_bytes_per_sec: 469583050645
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2533cf77a0d91824fbb4bea8cfd4b0a2"
+      network_throughput_bytes_per_sec: 4801172161
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b88a288b2a30a183fa97fb0d1979b6e7"
+      network_throughput_bytes_per_sec: 516783867589
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fd1292016d25aef5d8ce8bc8daf0bea0"
+      network_throughput_bytes_per_sec: 1526588
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "41f2d7aa0ba4ed22e15782862c4f648a"
+      network_throughput_bytes_per_sec: 33165991902
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "218583cb4f92d5dc7683ed5862d0baf2"
+      network_throughput_bytes_per_sec: 100777482
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c682a06c8e1d73fc431462d84fa7aaff"
+      network_throughput_bytes_per_sec: 513864403318
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "07d9e30199b52c4b0fc8be656f352d28"
+      network_throughput_bytes_per_sec: 858394033108
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "11fb4ba9898d02d7248aacd644f182e2"
+      network_throughput_bytes_per_sec: 578858111
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2415919104
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a85bcc0747986fc3102ab5d1fe9cb91f"
+      network_throughput_bytes_per_sec: 989375064026
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2415919104
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0404e2aea08051499cc8a6f7d08ac7df"
+      network_throughput_bytes_per_sec: 858861614399
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9b0b80f16a110f0da7387ac2a1d93e99"
+      network_throughput_bytes_per_sec: 31636023532
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 36
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a0cca3fbe6016a2471561b6a6675fe99"
+      network_throughput_bytes_per_sec: 30482641
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fe27ff2e266f0881dd6e0df6eabfe113"
+      network_throughput_bytes_per_sec: 461765652200
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2415919104
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "46eda0c5e8152522605961b1bd168f26"
+      network_throughput_bytes_per_sec: 469426639406
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f0026bc85b09ec5b62c5c40c6cb88768"
+      network_throughput_bytes_per_sec: 62581486
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 128
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "09d0cd26f161b09b4cbd15336b2ecd4c"
+      network_throughput_bytes_per_sec: 76230809
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4bcaa8511b034f29b5b5c146752e5ac6"
+      network_throughput_bytes_per_sec: 3348331
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f8eb7552fe9301ceb258e5d7d8539459"
+      network_throughput_bytes_per_sec: 3688146
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7ea5f6694a8b6a85a8f61b5e7d390c26"
+      network_throughput_bytes_per_sec: 141374679789
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6e999a7eef01755eb4334a506d2eee73"
+      network_throughput_bytes_per_sec: 1073430529146
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c48f8bb6120063fce2eca46aae25b967"
+      network_throughput_bytes_per_sec: 37252619
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "73ce2a0224010f131dd5365d9e92883c"
+      network_throughput_bytes_per_sec: 160305146285
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 8589934592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "050987748fb720b789c8c7bf13d9b18c"
+      network_throughput_bytes_per_sec: 932876559753
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8b9d57c55cdbbe0651e648677096a39d"
+      network_throughput_bytes_per_sec: 7125185
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "485e9c29c116f5415fdb02e401a19bb4"
+      network_throughput_bytes_per_sec: 21762815989
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8e4e6c5eb0287fae14196ebbd8a9f4e3"
+      network_throughput_bytes_per_sec: 651874635129
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "04d6b44af802bcedd9bbcd0372f32b96"
+      network_throughput_bytes_per_sec: 2582959641
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "9057ad6a4065d53a8507f95eb083a6d6"
+      network_throughput_bytes_per_sec: 14859273441
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "21108b14c350620915dca0e40fcbb1e9"
+      network_throughput_bytes_per_sec: 4785046728
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "87fc1aab12510238af519cca68f5611b"
+      network_throughput_bytes_per_sec: 484119527021
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "dc44edaaf9720ea4084f2f7e46401542"
+      network_throughput_bytes_per_sec: 870368126167
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "283f43b217051a84a9cb52a2e8436692"
+      network_throughput_bytes_per_sec: 920951
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9064885683b2c7495c29759e58536124"
+      network_throughput_bytes_per_sec: 278506241699
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "59f31274f7a1150a4eaf6ca44cf316cd"
+      network_throughput_bytes_per_sec: 301886792
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b199145371e4632a3a1cf9d663ebe09b"
+      network_throughput_bytes_per_sec: 583475783
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "83c2226e60e550154fda1db8995ebfaf"
+      network_throughput_bytes_per_sec: 51582841401
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "18336c841be8c66cb506b95ed59872d3"
+      network_throughput_bytes_per_sec: 983698265494
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8597a3f9bc7d9d5614638c0d1d36cb4a"
+      network_throughput_bytes_per_sec: 4014592975
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "36f33b4f2dd51e700353ee25473cda7e"
+      network_throughput_bytes_per_sec: 1550340651
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9c05898fa9cbd5fa38126d57358562bb"
+      network_throughput_bytes_per_sec: 738461538
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f5c114e847fa3cff62dd7cfc0d7bb0cb"
+      network_throughput_bytes_per_sec: 288970053279
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "306bbacbff99a909733db77decad0b0a"
+      network_throughput_bytes_per_sec: 36308033240
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0170de3aea358870c84cf65ef29e31fd"
+      network_throughput_bytes_per_sec: 1020673075686
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4294967296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d2513321e1e1a50907ec82f9fff7a1a5"
+      network_throughput_bytes_per_sec: 970847575261
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d4483d085fc5fac55c5e9359481145ba"
+      network_throughput_bytes_per_sec: 730810471751
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "45faed656eab706c8d4525f4f1518e49"
+      network_throughput_bytes_per_sec: 81153549807
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "ae98fad084b0793edf70496c5a00ae2a"
+      network_throughput_bytes_per_sec: 287015397774
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2fb52b70e9e611b7f558d0aee60c78c5"
+      network_throughput_bytes_per_sec: 8397480
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4d2d7bd3d5dcac3e5f6b9e71d1c59bfe"
+      network_throughput_bytes_per_sec: 88048151
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4c56227d18b09691ba1d9ddd91019d7f"
+      network_throughput_bytes_per_sec: 116334259
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "19dddfdd24653164439fdf6314a797f3"
+      network_throughput_bytes_per_sec: 8856216216
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f11aa5914aaf74d2c2c03c4f3d3d454c"
+      network_throughput_bytes_per_sec: 1109769377056
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a370e8984a5a6f1dc9522c8a2355a6aa"
+      network_throughput_bytes_per_sec: 868335982517
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "46b224241fcbc6ada241d7df160d29f8"
+      network_throughput_bytes_per_sec: 21646506165
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "57eda245ae01f780c0c9ab769ac945ec"
+      network_throughput_bytes_per_sec: 1081205452816
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "03bf131e87914ccf9d958e21dce9b83d"
+      network_throughput_bytes_per_sec: 240884497
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "cd29cf380f9f7bfddeacbafed09918a2"
+      network_throughput_bytes_per_sec: 51720799719
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4831838208
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a5f0dc9e9236b6509ca27ce089a7c5a4"
+      network_throughput_bytes_per_sec: 991149602908
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "4ae4c9eda8a5bd76972ad1a465100279"
+      network_throughput_bytes_per_sec: 18974458267
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9a6150c85feb8bda23699f4df7048005"
+      network_throughput_bytes_per_sec: 11800376
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "f497b8dc9f3255d6ca2efb4406d7cae2"
+      network_throughput_bytes_per_sec: 664875562555
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7480bfdac377ee2722d36da3ed5d78a4"
+      network_throughput_bytes_per_sec: 2070451457
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "59bdfc9e9884a5d7714a300f9d8e21db"
+      network_throughput_bytes_per_sec: 1349628
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "4087619b09ae40ff5d22f39cd1bda217"
+      network_throughput_bytes_per_sec: 1475703670
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d44ff7a830e19df7f07cb2b077948310"
+      network_throughput_bytes_per_sec: 375148841851
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1f73d322d2b44142b0df2522290da123"
+      network_throughput_bytes_per_sec: 380381588410
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ac30cb64183ff23f8326dc96e6c5fb64"
+      network_throughput_bytes_per_sec: 302750396997
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "22bc3fdacc931b250651badbc72261f2"
+      network_throughput_bytes_per_sec: 491150677713
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "24485293843c6ea12d120047b050a83b"
+      network_throughput_bytes_per_sec: 671983585639
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "19e7f87e2f693cd752ad9cd3d5fa9539"
+      network_throughput_bytes_per_sec: 26032417133
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a78c87dc9d5a64b8c57c5600a134497a"
+      network_throughput_bytes_per_sec: 38764385
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9c5d7f5cfd10cf419b3e0e19b46e1151"
+      network_throughput_bytes_per_sec: 97278297
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f6064ed8d97e0480fa38fb09878b495e"
+      network_throughput_bytes_per_sec: 144909078
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "b566d4157ae0cb3e8490f19958354bb5"
+      network_throughput_bytes_per_sec: 24729010754
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "393e18fcb3ec1f81623244b56022cb43"
+      network_throughput_bytes_per_sec: 796441679
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1b00154d0e42194917b9893ebcd98560"
+      network_throughput_bytes_per_sec: 35607483
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7abf86582143d654b23cbaad6ed3371b"
+      network_throughput_bytes_per_sec: 293241695
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "aca580f880883a5d39b47a21ce239e50"
+      network_throughput_bytes_per_sec: 570343214577
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5f26d145badf25c25041c28d36acaece"
+      network_throughput_bytes_per_sec: 30871139956
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 8589934592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "31c7cc1ca3cefe3d57ec05da3a4ac12e"
+      network_throughput_bytes_per_sec: 950205200477
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "94db985d92e1bb442ead90e9cae1bd0f"
+      network_throughput_bytes_per_sec: 75122460676
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "aacd486efd475d2631e062978515ddc3"
+      network_throughput_bytes_per_sec: 18073910645
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 4831838208
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "673054edfe3048c95d68c442a88b1952"
+      network_throughput_bytes_per_sec: 801942852604
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7d707c83eb2406826961c0a367c4a613"
+      network_throughput_bytes_per_sec: 85333333
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a645aab2ba6304eaf991fad9bf685073"
+      network_throughput_bytes_per_sec: 667470947464
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "956c2513603ee471adb58c129962bac9"
+      network_throughput_bytes_per_sec: 278943067
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "164bab0b8e4ca9ead6f911eb64267fff"
+      network_throughput_bytes_per_sec: 166558136251
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "cd3404dc4c026ffa7e097817cf7893ae"
+      network_throughput_bytes_per_sec: 181233369181
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c31000e6f19ed731dfbb4f23592520f2"
+      network_throughput_bytes_per_sec: 2701055099
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b28fe3d81a36c074975ec09979307f6a"
+      network_throughput_bytes_per_sec: 13185728337
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8c650186919a04355d46b4e68fe2cc17"
+      network_throughput_bytes_per_sec: 117844009889
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7c49343f4ba56d020e87a657bcb0a206"
+      network_throughput_bytes_per_sec: 2125461254
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c4038d2af60f3e6f660d89a15832a216"
+      network_throughput_bytes_per_sec: 5157246782
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a81acbb8795d9ba131eba921846ec448"
+      network_throughput_bytes_per_sec: 348286156
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2415919104
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9a6fcc63e7ad468f3a089b6ba702c469"
+      network_throughput_bytes_per_sec: 384675747611
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e79b2ea1668746e126645a5673d7a696"
+      network_throughput_bytes_per_sec: 19098143
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7f164896fc11c1adb8d1f9a3731ba223"
+      network_throughput_bytes_per_sec: 863679493
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "235bbfcf390c9045e5183935f016e5d4"
+      network_throughput_bytes_per_sec: 305551395060
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "779b1892b9dba14d2509c82c326ab739"
+      network_throughput_bytes_per_sec: 8475456949
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "e269bf8e3eec308689d257a82088fb8c"
+      network_throughput_bytes_per_sec: 49073788649
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "581cb9885a12fd414b3485e1843f54dd"
+      network_throughput_bytes_per_sec: 2614094
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2f8e59e305de61fd7b9632bdaa786f07"
+      network_throughput_bytes_per_sec: 1401300034
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e255335ed4e7a8364ea63c6af2fd5efd"
+      network_throughput_bytes_per_sec: 2086070791
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c2c8ac97d42ed306917395397541fb48"
+      network_throughput_bytes_per_sec: 109203915850
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a3d313e74e46d0c3ad18f820b05418f3"
+      network_throughput_bytes_per_sec: 130003085739
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "78fed61394c990fef0156f94394cb533"
+      network_throughput_bytes_per_sec: 771129891221
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c5af05b453d29cbcdf83e50193f05e8a"
+      network_throughput_bytes_per_sec: 727073282773
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ac560c42d1ecda86827bab1aecc0e12e"
+      network_throughput_bytes_per_sec: 8294930
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e09383ce428ef3dcc904ba6cb3b7e548"
+      network_throughput_bytes_per_sec: 76762518301
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4294967296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "09614fd44a516e28ca0da499c8aee8cc"
+      network_throughput_bytes_per_sec: 551205298618
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "53884a0163a566b71379a8c94efa4d3c"
+      network_throughput_bytes_per_sec: 1226071989
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8db0c38d5387676469929b53f70ce5b9"
+      network_throughput_bytes_per_sec: 906948640659
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b0eb7fa5be704c1855f53d1a70da6a11"
+      network_throughput_bytes_per_sec: 341017576
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "12feba37cc23c927ce242a454b962b5c"
+      network_throughput_bytes_per_sec: 17084462982
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2de1949ea86295a39740c48bf68fe23d"
+      network_throughput_bytes_per_sec: 5327696935
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "da26fd283c530038b965cfc6de3c9abe"
+      network_throughput_bytes_per_sec: 525085884410
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2fb216ef79ece73a37a82efdf230eaef"
+      network_throughput_bytes_per_sec: 226044057611
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "39e361717257cb21d6b6d471fc657737"
+      network_throughput_bytes_per_sec: 4641359773
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 8589934592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a36f8d5bdd47e56a2d55dc9f8bb03a59"
+      network_throughput_bytes_per_sec: 460567833097
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 32
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9d203c05b8fd3ffe89dee820abef6558"
+      network_throughput_bytes_per_sec: 16253968
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4294967296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e8d7afccc648a0212a7283d5e9e63f6b"
+      network_throughput_bytes_per_sec: 971132948491
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1504dbbdba3b2f75fba3685976f9d4ba"
+      network_throughput_bytes_per_sec: 25669307
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e7d8d83afad2f1b4d532fbd4900ce0d9"
+      network_throughput_bytes_per_sec: 907203460706
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3292227bf039440b853b67423f14d2e8"
+      network_throughput_bytes_per_sec: 5566669
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "29ec550817a2b571d5f7b0b7edfaa7cd"
+      network_throughput_bytes_per_sec: 283564970478
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c80160ac56ed538e56cbd73bf8552ae7"
+      network_throughput_bytes_per_sec: 991047665
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "ba75c1b497c6836df0a57bb98909720d"
+      network_throughput_bytes_per_sec: 763276725907
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6291512936b0f96a6b6f122a8cdaaf51"
+      network_throughput_bytes_per_sec: 569152385
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e0da1e915e1a18820008cca1bfb567a9"
+      network_throughput_bytes_per_sec: 776183946333
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "414aa6de3c21c1950e6724b5fda8fc3a"
+      network_throughput_bytes_per_sec: 661498708
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a7db3e7e26a257a8330c2eb69b548dfa"
+      network_throughput_bytes_per_sec: 34212529002
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "41b263d9cd44ed42684781560bd92aa3"
+      network_throughput_bytes_per_sec: 9240246
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a2e2d6d4902c28916ff21305ac1aecea"
+      network_throughput_bytes_per_sec: 2852367688
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "abed44518bc7f769215d5b14d6ee715a"
+      network_throughput_bytes_per_sec: 661073378318
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0577c1a8958ec9badee3dd2d35e37fb7"
+      network_throughput_bytes_per_sec: 208883675
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 64
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "82cea8d963cc3a035667d3ee8e173db8"
+      network_throughput_bytes_per_sec: 36875800
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "70d87dcdc841db3df5d4331a1311f050"
+      network_throughput_bytes_per_sec: 882188240
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f32d6a77d8f5879d766fb9972b2839f6"
+      network_throughput_bytes_per_sec: 1340996726
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b60bae80b358465835f310fd0a0838ae"
+      network_throughput_bytes_per_sec: 464939413
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "21eb91a153d1421db01865ff847be623"
+      network_throughput_bytes_per_sec: 206390027337
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "92b66efa29ae065b7f11272bb9e8a1ca"
+      network_throughput_bytes_per_sec: 178815825375
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1f62a923de5e56bf7afc1292d82ae7a3"
+      network_throughput_bytes_per_sec: 8057296
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "68ce764b58f8554c0e7644d2c3368cff"
+      network_throughput_bytes_per_sec: 819395602248
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8cb8ea955b44f7201d6b0f81708bf905"
+      network_throughput_bytes_per_sec: 3784096
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 64
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ea284e2e077799d0aa9bcfa47a84fe4b"
+      network_throughput_bytes_per_sec: 30761836
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "ebc9c62bb14516f165cbe17c15b141ce"
+      network_throughput_bytes_per_sec: 1094666825
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f4452ba73ddd7b552eda8a25e6aab30e"
+      network_throughput_bytes_per_sec: 1066666666
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0ee85a76cc2d48770f64a506f88939f2"
+      network_throughput_bytes_per_sec: 423265694226
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "57dd881ff8db24b4dba5f2f8496a5d41"
+      network_throughput_bytes_per_sec: 709687355
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d760f15feb519d7c73fa0b649562cbda"
+      network_throughput_bytes_per_sec: 239537635636
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "37acccf5b6997f88c1ec9e9744245702"
+      network_throughput_bytes_per_sec: 982568367
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9e3eabcd68bb4fa2f7cc4e1c87945ef3"
+      network_throughput_bytes_per_sec: 29341558053
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "071aa0eee20cd9655cf796e0868c90aa"
+      network_throughput_bytes_per_sec: 112701633705
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "092128c2a1323ffc7bc73e11884275c8"
+      network_throughput_bytes_per_sec: 4245048364
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "eacee2998c6f69b93ea0d14ef839ef06"
+      network_throughput_bytes_per_sec: 515465070
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "70c3e18ac591078b59e83f680c55f10b"
+      network_throughput_bytes_per_sec: 98159509
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d7ddabb9fa37d6f0facf66a724cc2e39"
+      network_throughput_bytes_per_sec: 3428571428
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "062f1a7d5697b2405e8f8abce9c4ec09"
+      network_throughput_bytes_per_sec: 10177307
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3033cedcb703a9bf37d50ae42eb6548b"
+      network_throughput_bytes_per_sec: 388865566475
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e0a5906945cdf437e334315c7d20262d"
+      network_throughput_bytes_per_sec: 64419397116
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "dcfaa3b3ac539ae87d625d64076396ae"
+      network_throughput_bytes_per_sec: 359411007985
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f9f96fc8f15dfda377fb720a8bdd33fa"
+      network_throughput_bytes_per_sec: 1101205334438
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "275528cb2db3317c00a406527bc47911"
+      network_throughput_bytes_per_sec: 2641515517
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 32
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "41a83d57ea4967b6b9de0542e975f21b"
+      network_throughput_bytes_per_sec: 10281308
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "02df948e19eacaf86782d4edad5cf291"
+      network_throughput_bytes_per_sec: 267248871012
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2233ea511f381a3fd5a1174beab1a43e"
+      network_throughput_bytes_per_sec: 102432510318
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "4b62b2e467162db51188d86c51eb93da"
+      network_throughput_bytes_per_sec: 4708891683
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f6c2865626ca89f6f3d3609d48a2575a"
+      network_throughput_bytes_per_sec: 890383483727
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2111c0e7f9cb51af4ec2cf582204f95b"
+      network_throughput_bytes_per_sec: 685542931861
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "206f13169754f84255f0f0559a0efc54"
+      network_throughput_bytes_per_sec: 365888518
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "395574b6d4f2d45eb7665bfdeaa341a0"
+      network_throughput_bytes_per_sec: 932604735
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8a7de0fef3ddd91af521e3058947fd19"
+      network_throughput_bytes_per_sec: 552075816075
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5d37af5fd23fbb3dc6e469668a149a1c"
+      network_throughput_bytes_per_sec: 38824644549
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5afeed182ba30b9644e047dd46155bea"
+      network_throughput_bytes_per_sec: 333870515814
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "906e82cf350b67ba0a851a932bec3b30"
+      network_throughput_bytes_per_sec: 293062045835
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6b11a7841d2be9d3eeea688aeb542800"
+      network_throughput_bytes_per_sec: 167514904
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1b7cac971989ac5bb9dc04cd0f3c57f6"
+      network_throughput_bytes_per_sec: 115263394
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "bcabc746a33a78e5c3abfa96b2299823"
+      network_throughput_bytes_per_sec: 385941769
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c3f317c91b962d11ad7e45a345d17c35"
+      network_throughput_bytes_per_sec: 194938836214
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "62e587c11b30e34e7e802e45f2495dd9"
+      network_throughput_bytes_per_sec: 33151079136
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b7aea2a383ae2c961ac99b8c90e3d198"
+      network_throughput_bytes_per_sec: 540363823756
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "05c0ec2144e63e5a203ad802cf445713"
+      network_throughput_bytes_per_sec: 479582610994
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c4136662830eeb95c232f7cf7c44f758"
+      network_throughput_bytes_per_sec: 24686595
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "dc934830a3534472858702e5dec30807"
+      network_throughput_bytes_per_sec: 438731790
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4294967296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "40648d53f94a9e1e51171541e8da5426"
+      network_throughput_bytes_per_sec: 385021793766
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3b4e1b8016271ce801cdcc9b20059fc7"
+      network_throughput_bytes_per_sec: 717474
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "619faf5392f9f10f8105dcac7ebd4d5e"
+      network_throughput_bytes_per_sec: 887537289570
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "813ffe362e12bfa92db9697d1c6d0779"
+      network_throughput_bytes_per_sec: 1274951983666
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "c72131c9668a455ac72e520196c0da14"
+      network_throughput_bytes_per_sec: 13529314616
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "eab320ee0d6266fd6e1841d5a4dfc5c8"
+      network_throughput_bytes_per_sec: 3676095
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b038ee5a2ec0f3a8b17b771f2d3fb62a"
+      network_throughput_bytes_per_sec: 483716295698
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b352062d66a61c94c879a0dfd41d10a2"
+      network_throughput_bytes_per_sec: 890920239317
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 2415919104
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "ebcdb0d5b70ac0413b0d66ac8d8e5cc1"
+      network_throughput_bytes_per_sec: 497986006660
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a08b7e1cb1d4b1127abf284301d0b47d"
+      network_throughput_bytes_per_sec: 68266666
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ddc94467b0bf0e77ce2d40b9d55aa2c6"
+      network_throughput_bytes_per_sec: 329223233908
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d67cf152c02aff26940793a85475a48b"
+      network_throughput_bytes_per_sec: 41308089
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "28f2b15bdb9efca5dede1828987f4411"
+      network_throughput_bytes_per_sec: 22723994452
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d047cf8c51b561295fd93f481f278848"
+      network_throughput_bytes_per_sec: 6860971524
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f8d03aa8ead5fe5fc8fb9ac70b73f120"
+      network_throughput_bytes_per_sec: 420929962993
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "383b127e6df254b203c38b21e611498a"
+      network_throughput_bytes_per_sec: 387978293043
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ad9737cc54aecc460cac503af860efd0"
+      network_throughput_bytes_per_sec: 476842201000
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "88bc13317ad9c8b1c76e53894135f1f1"
+      network_throughput_bytes_per_sec: 2114610221
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "661a3ae9f5fec596c750c51a44489d1f"
+      network_throughput_bytes_per_sec: 112152496850
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ae70c0b5c2956aa7d1027b538a66b9b8"
+      network_throughput_bytes_per_sec: 3289372713
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "81bba3d5e10cda7e18bf20722b1fe356"
+      network_throughput_bytes_per_sec: 1185107131189
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e69f85918573ef82889899bffa674887"
+      network_throughput_bytes_per_sec: 929325014
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8451affa7fb4b9d8aaff412440698b0b"
+      network_throughput_bytes_per_sec: 1000611
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "15aaf10a025b4a25203704e974a7f965"
+      network_throughput_bytes_per_sec: 44390497
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "798cfeca8ce6bf3a9d5e2a350cd07057"
+      network_throughput_bytes_per_sec: 765492588161
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "08b89041e1b8c318aafdd2e6b71190e4"
+      network_throughput_bytes_per_sec: 1684730077
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "de1851dc4a729b8f647a0e7c4d8d71b9"
+      network_throughput_bytes_per_sec: 375534580183
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b103f7436cd3af8a3bc0e3ba1b880332"
+      network_throughput_bytes_per_sec: 38787878
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "e547c43faecb547d8b5f8508594e44ae"
+      network_throughput_bytes_per_sec: 712540590
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3e6fd684ef2e49e505b37b8920d76016"
+      network_throughput_bytes_per_sec: 1108952291421
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fec72f01f91228220fc2de99f6805cdd"
+      network_throughput_bytes_per_sec: 93143596
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "7766103ac69e90986b9a3e03dfe89efe"
+      network_throughput_bytes_per_sec: 1058739903
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e2b1eba18ccba81d81e4667b86b5b73a"
+      network_throughput_bytes_per_sec: 4498627127
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f749bdd44b6420c7413442144358f696"
+      network_throughput_bytes_per_sec: 803737472453
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "cd2470a21534fd3315db48b468340d61"
+      network_throughput_bytes_per_sec: 650939398
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7c9e14facae01784bc6bb5deeb02a728"
+      network_throughput_bytes_per_sec: 377859778
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8108a48f1580d191d8574f82adc623ca"
+      network_throughput_bytes_per_sec: 203346075
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "99195c845f3c07ed076197e54f329a0e"
+      network_throughput_bytes_per_sec: 287825080363
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fce40c6ddc979445780af6e48bde73ec"
+      network_throughput_bytes_per_sec: 172791562911
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "da44656a2c92dea3b84cf0e6aa104017"
+      network_throughput_bytes_per_sec: 846057577266
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b2beeb1caf669b139c8f79484a5e401d"
+      network_throughput_bytes_per_sec: 155165800723
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c3b2097801d91f16d519ce8a9dc75841"
+      network_throughput_bytes_per_sec: 2542520173
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "cb5a48234a76bc7514d5242540b06ade"
+      network_throughput_bytes_per_sec: 156203389
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 32
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b6e8b6e73d229e29e7a227b173df3fc2"
+      network_throughput_bytes_per_sec: 15183466
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ea69877cb0438b96de4191b8d9086025"
+      network_throughput_bytes_per_sec: 147626065
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4d35610d6f63a752dc196b29a4c4babb"
+      network_throughput_bytes_per_sec: 2959537572
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 64
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9d8040edaab0206e73bc4e3a08c9439b"
+      network_throughput_bytes_per_sec: 10269576
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5e4257cec4e8ff5d588b8c1261ca022b"
+      network_throughput_bytes_per_sec: 1222002610
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4294967296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a8e2d82be366acc51db98a1d47efcf01"
+      network_throughput_bytes_per_sec: 856121684621
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "6d3d3ad018b4f52ea94bfed1c49d55df"
+      network_throughput_bytes_per_sec: 46036955
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "479fdeeb18a4a5bfb18e76d14be75176"
+      network_throughput_bytes_per_sec: 88727026569
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "718e70921ca037fe429512a0dc3fd4d9"
+      network_throughput_bytes_per_sec: 63310617
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "878b52eb7f7b73d5f0f223046e3be5b1"
+      network_throughput_bytes_per_sec: 486749851
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "955a032d547eb5c3dcff422be85c4943"
+      network_throughput_bytes_per_sec: 2542507
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3754579f36a495e0eabaa9e3b90dc678"
+      network_throughput_bytes_per_sec: 1514543960
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e3abe3235565038a0e84365ec04f210d"
+      network_throughput_bytes_per_sec: 49824632539
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 32
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "71476c3d711c83704f0bc49620b2e1f4"
+      network_throughput_bytes_per_sec: 9500482
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "96886ca24885d3cb68eb1e09e1e4bc15"
+      network_throughput_bytes_per_sec: 300050362457
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "cc9d75ee5f0b08cdc9c2dc7e203fe53a"
+      network_throughput_bytes_per_sec: 891840400
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "d9be66600f73749162dc078ae1612360"
+      network_throughput_bytes_per_sec: 494624991477
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e94949e67d96d86b05e8941b6e376699"
+      network_throughput_bytes_per_sec: 614021483338
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e65e98c15085d39d390675e74b152690"
+      network_throughput_bytes_per_sec: 528020813808
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "51a2e3773ba92c0daf242639130c6297"
+      network_throughput_bytes_per_sec: 190028271112
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f5af72ab9f895a79586e0dbf8609cab0"
+      network_throughput_bytes_per_sec: 12105163
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4476b7a8f06720ce0437eed9396a603e"
+      network_throughput_bytes_per_sec: 451051646
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3e9dfb90409988159a9ba9aed894da4b"
+      network_throughput_bytes_per_sec: 1552233778
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f8baee334d0240a83b5630990234b8c6"
+      network_throughput_bytes_per_sec: 370616136823
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "683b818f6df80cf6d93a50e68ffda811"
+      network_throughput_bytes_per_sec: 1008717643129
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4294967296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "024f84edeacccf466b137ca868662d60"
+      network_throughput_bytes_per_sec: 1020937895126
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "012d4c97b32666106cc4fbaa089b017f"
+      network_throughput_bytes_per_sec: 110702702702
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c1b369449e3ace0334167f6f83933325"
+      network_throughput_bytes_per_sec: 58832744206
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c6272f96c3a7ba815e81d07a43ad8e69"
+      network_throughput_bytes_per_sec: 3391986
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3e8157fd0d0fb8eeb6a09f5f1d7d47d0"
+      network_throughput_bytes_per_sec: 42435289356
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fb016dfdc6dd7c85247bbf3fe51d5f0c"
+      network_throughput_bytes_per_sec: 309294179339
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a895292b189eb8ffaecb80f68fc27441"
+      network_throughput_bytes_per_sec: 789623394552
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "cfbb1702926835c6f6ecb227e8332338"
+      network_throughput_bytes_per_sec: 227709981269
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "95551cf6a6af2fecf692ad351df2a5e6"
+      network_throughput_bytes_per_sec: 349806422
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "073bde19d69065941da575e56dffe0cc"
+      network_throughput_bytes_per_sec: 402913230914
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ed2bfec2bd6490c0e37bd318b4ef3058"
+      network_throughput_bytes_per_sec: 88656302
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "dc670ba1a51f2d839391e5c3b2e656df"
+      network_throughput_bytes_per_sec: 67753029693
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8845a6732bf4728c1cbddb3892dd8201"
+      network_throughput_bytes_per_sec: 449877248670
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fe21f6b68d9ef4f3c6325fc16ace49e7"
+      network_throughput_bytes_per_sec: 288041754099
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4294967296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "cc2a7f0ee3f139fd2ab6e5fe8e4d6958"
+      network_throughput_bytes_per_sec: 937179792002
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 64
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c835e3d0df8c671fa8a2e0b4401b1de4"
+      network_throughput_bytes_per_sec: 30929080
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e4bb5b62aa36a2e5f5b84e3ed9dfaf69"
+      network_throughput_bytes_per_sec: 13208908596
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f9ad4fcf4482d77380e6c5b77db0bee8"
+      network_throughput_bytes_per_sec: 41587133
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2415919104
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5d10c67dc531a0f61f66272dd657206f"
+      network_throughput_bytes_per_sec: 975268428916
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "151cb6b98574f2cfc7325ef74e90103e"
+      network_throughput_bytes_per_sec: 100301675027
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c66a2f9b9e557b953c1f0bb1453b154b"
+      network_throughput_bytes_per_sec: 336252189
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c9f3fdb466f05b6c1a2c4d68f53a5837"
+      network_throughput_bytes_per_sec: 634304610834
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f79deef119504dd3bdcffd26c9ed603b"
+      network_throughput_bytes_per_sec: 34543246
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 32
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e896204cdd4bfb4372ff4533be872d6d"
+      network_throughput_bytes_per_sec: 24825446
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5b75c758c3832e1e678f4349a9bb2600"
+      network_throughput_bytes_per_sec: 14801196996
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "55d4bbb946712768f4ae060ee5f6dd7d"
+      network_throughput_bytes_per_sec: 1014151413762
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ab7c71f741c59bfa429ff5f6197dc0f5"
+      network_throughput_bytes_per_sec: 832212348019
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "38ce55c6dfe2e33b78c9a069142d2b68"
+      network_throughput_bytes_per_sec: 952888838403
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "e1ef170770478f2271f59801b803bace"
+      network_throughput_bytes_per_sec: 289759710153
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "8387735f93bff09f3dbe2a8273d627fb"
+      network_throughput_bytes_per_sec: 650384069633
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "bde5442a6651e112f1721f721d3b0db8"
+      network_throughput_bytes_per_sec: 19956151035
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "dbb91dcc479983d4b776d1e5a73c35fb"
+      network_throughput_bytes_per_sec: 5022300559
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0216d8ac7e6bd2428e62288a46d31d05"
+      network_throughput_bytes_per_sec: 12270361355
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1bf80c79f1888820ef4715b6a3428882"
+      network_throughput_bytes_per_sec: 533940645736
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "17fcf99dec1c1d55164f7b325bbdf5fa"
+      network_throughput_bytes_per_sec: 275408660
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7fe46d69840288c24e9cb6135f8effcd"
+      network_throughput_bytes_per_sec: 4138168
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c483355ceb90cb07739aa83360adec0a"
+      network_throughput_bytes_per_sec: 790698606592
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "27ae8b5c9ef13a3796dde8a8d3a9276e"
+      network_throughput_bytes_per_sec: 2923495561
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5f84e76780820af349a900f392e7c51a"
+      network_throughput_bytes_per_sec: 3341551849
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "74143f9faa56438941c6eb3761280e3c"
+      network_throughput_bytes_per_sec: 217868316557
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "58c7cdd0ee1c0276fc26b47204ce6dcc"
+      network_throughput_bytes_per_sec: 81736909
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "182bcdb791bdbb0bda91b22d0283a2fb"
+      network_throughput_bytes_per_sec: 129195054
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 8589934592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d944457eae953e50887f514de10fd8b6"
+      network_throughput_bytes_per_sec: 939727748952
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c39fd3990f04dec6477cf25db3447f46"
+      network_throughput_bytes_per_sec: 294387622048
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "94868af1ae6acef22a9a583d7a4b76cd"
+      network_throughput_bytes_per_sec: 222763576493
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2415919104
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ebc4e33a43272350bda9138d47d9f90b"
+      network_throughput_bytes_per_sec: 846240724316
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "13612c140b4b9111b0850f2c4216fdaa"
+      network_throughput_bytes_per_sec: 77715791554
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "51a27c5499d67792b30b55817b61d818"
+      network_throughput_bytes_per_sec: 1024512256
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "744e482522036557be7edb29975f4fd2"
+      network_throughput_bytes_per_sec: 1846153846
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9940e26344bd908558b525a721b3ec88"
+      network_throughput_bytes_per_sec: 3096871751
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ac573e05294bb02f86f2d38b22070361"
+      network_throughput_bytes_per_sec: 245763200041
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0452a0ef46de3f78e120f7f273dde38e"
+      network_throughput_bytes_per_sec: 4673427991
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2da855742dc84606de3a1905d5a9acea"
+      network_throughput_bytes_per_sec: 7842987075
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7e69ccb2e0b490f0a5eea49ff2a1c539"
+      network_throughput_bytes_per_sec: 224124513
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "418332822e4b913a95a0a4c642f85721"
+      network_throughput_bytes_per_sec: 96371935
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "22fbbc2ad25c5232d83a80a7cc6c8204"
+      network_throughput_bytes_per_sec: 744201
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e3edea83af4f3223b2c66b5cd217a10c"
+      network_throughput_bytes_per_sec: 105783203026
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c96e18392029d72c4ab13c9473a29198"
+      network_throughput_bytes_per_sec: 183315431
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "aa31631046e7a7f1be5e0ef5022cf327"
+      network_throughput_bytes_per_sec: 33180918091
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5bc8bb6a32d73cbbeaaf91367bd7bd5a"
+      network_throughput_bytes_per_sec: 130428011692
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "35e9a0db7a08b91bcfbc3fd0cfc1de0f"
+      network_throughput_bytes_per_sec: 477043325
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c6a2c0de4fd14d25039c0101c144a9a6"
+      network_throughput_bytes_per_sec: 449802032815
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "deb2c45b98a39cb988b4faffd1481d8a"
+      network_throughput_bytes_per_sec: 428096985642
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1b1260805cbaa43a4d38fb78e2964334"
+      network_throughput_bytes_per_sec: 991392164333
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "27b5cff22263c4cb3ecd1645952cf998"
+      network_throughput_bytes_per_sec: 460277142104
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2eb4bbc3f9027f2ca45daba3a766681b"
+      network_throughput_bytes_per_sec: 50739957
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 8589934592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "58812a7917e91cd23cd549102a7ca417"
+      network_throughput_bytes_per_sec: 471715373452
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b88e1b600ffd839b2d1d37e3730958b1"
+      network_throughput_bytes_per_sec: 965129380604
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9c95a82232c330bba9268b1c42c69e81"
+      network_throughput_bytes_per_sec: 1647404
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "efaf5ad30cf7cf8ef3d5489b13741ec5"
+      network_throughput_bytes_per_sec: 188254219030
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "bb5e3417de8aee55d83b37a96cdacc30"
+      network_throughput_bytes_per_sec: 46966731
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "b1a6f56757f39a5f3001d27be83015f1"
+      network_throughput_bytes_per_sec: 273959452643
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fad69743c9ef3509515136dcb5a899cc"
+      network_throughput_bytes_per_sec: 395521146683
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a530af17e81596a765ef6e7847dc3cfe"
+      network_throughput_bytes_per_sec: 5352562
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a54b36f2b29b619b2f56de18e263afc3"
+      network_throughput_bytes_per_sec: 1670728598
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 128
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5fcb817735bee0e1d6187e82dec9a4bb"
+      network_throughput_bytes_per_sec: 37566034
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0730fe5c7f6720f4f86d76525e508653"
+      network_throughput_bytes_per_sec: 688606796913
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "768b4ae232da120c8a036634ab3b0038"
+      network_throughput_bytes_per_sec: 378596648179
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e5e6ae741c0209df0363bc99619a38bf"
+      network_throughput_bytes_per_sec: 1028253707966
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a7b20ce7182eb1eb3a104df458a6da97"
+      network_throughput_bytes_per_sec: 656620711
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1e54c32a17f9d0ebe6f6022ada81d1bb"
+      network_throughput_bytes_per_sec: 78506794444
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "5f28526c6ee07defa9561b094ff5be4a"
+      network_throughput_bytes_per_sec: 796619210613
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3ed84c30904cea0f5be27da0576de5da"
+      network_throughput_bytes_per_sec: 21524263075
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8b9e042e84e780aa1afaf2175ff8917d"
+      network_throughput_bytes_per_sec: 1454134
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ffa82afc7c222c6f4fcf12a28be7c194"
+      network_throughput_bytes_per_sec: 943176073757
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8d3f066a90e6b76bff88cc9dd3ed0fc0"
+      network_throughput_bytes_per_sec: 544183100318
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "644e64623ab4514b27e3b7862786d68d"
+      network_throughput_bytes_per_sec: 438105
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "c6cf55c910d729541ea43cd567e569a0"
+      network_throughput_bytes_per_sec: 11497464
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f27edbc187a5d6486c8308e17591ac91"
+      network_throughput_bytes_per_sec: 777327935
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d4c2c607268d9ee0f57c426f8fa47cc2"
+      network_throughput_bytes_per_sec: 24568322399
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "113e0a3e653222d5f2c02d0a04259946"
+      network_throughput_bytes_per_sec: 77353074
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4831838208
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8153c51c9b5efa2cb764325e3418e035"
+      network_throughput_bytes_per_sec: 923752402933
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "78878b5d599605db9f5d5a06f4ba06ca"
+      network_throughput_bytes_per_sec: 794046931396
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7e4938cd3cbd8c9723dada450c5071a1"
+      network_throughput_bytes_per_sec: 3921141
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "186cb576c095eb388e98af06f3b17286"
+      network_throughput_bytes_per_sec: 21249907
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e9579a68b86bdc4bd33ed862a05d2055"
+      network_throughput_bytes_per_sec: 206101552774
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d58b04d8c8b1603f99df40b49e5ef0ab"
+      network_throughput_bytes_per_sec: 1065868653
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5c5dd40203384177dc77e6a8d3417caa"
+      network_throughput_bytes_per_sec: 1127408404
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "38976ecb52eff7e07859e37df2ea8f82"
+      network_throughput_bytes_per_sec: 68400260926
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ed981fa26aae9fb2e27a1d9bb370df1d"
+      network_throughput_bytes_per_sec: 7410594029
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "19d6c3c48a533c9d1064f954f981df85"
+      network_throughput_bytes_per_sec: 72194021
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6e2ce8b5a115b2b170bcffacd31fe57a"
+      network_throughput_bytes_per_sec: 10334735071
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2e0dd0e187fc6f8a499a4e16c5c2c446"
+      network_throughput_bytes_per_sec: 384500651890
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "912177fea91246436b74ddfe0c1b4ac2"
+      network_throughput_bytes_per_sec: 996024058744
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2e87f5c55b2478373381a4e2e5e1843e"
+      network_throughput_bytes_per_sec: 2799726589
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 8589934592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c46f7e557f6b90f0f6c65144923665f6"
+      network_throughput_bytes_per_sec: 856794923426
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "6a83af6e7e2a17bc0798f49560b764d7"
+      network_throughput_bytes_per_sec: 156080555768
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "54340d846737c71c79ed4256a9fbad78"
+      network_throughput_bytes_per_sec: 206079049657
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5e804ac502e093054591afe841a96249"
+      network_throughput_bytes_per_sec: 403812934748
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7d0e2114265c2b2606ebfa92e573f42f"
+      network_throughput_bytes_per_sec: 1060004169957
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "187050fa9ebb7484c41b2fc215051b09"
+      network_throughput_bytes_per_sec: 338456453
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "296f9b9864896db0fd0be6f7d09b22c6"
+      network_throughput_bytes_per_sec: 809606166
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "ddab18a364a1b4236778b5c1f9d237ee"
+      network_throughput_bytes_per_sec: 70622854
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9edc050519818561c5028e752737dab1"
+      network_throughput_bytes_per_sec: 118988097639
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "9b06ace0490dd4d934251316436f41ed"
+      network_throughput_bytes_per_sec: 494651195809
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "8af9e048b0024b4c717f80a07280f3c8"
+      network_throughput_bytes_per_sec: 4304530593
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3e813da898933cadca0c6229e9dc617d"
+      network_throughput_bytes_per_sec: 7863615
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "964d7744a6aa60f36a99ec042bd198a2"
+      network_throughput_bytes_per_sec: 6645642143
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0610284ad448aeb738b484586df1a0b1"
+      network_throughput_bytes_per_sec: 764128985243
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "be8e50419b637a1cbd10689f53440928"
+      network_throughput_bytes_per_sec: 1004591963993
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "28f108e0cf0d74ad9b173d1c4f2e8731"
+      network_throughput_bytes_per_sec: 28956588976
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "dd31fa2a021c6955d3a5ae20372af76c"
+      network_throughput_bytes_per_sec: 72357263
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c1c2392f5c6bf691916c8471ac7bb408"
+      network_throughput_bytes_per_sec: 884371380329
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8523a798691e70880092b7dc5dc81f03"
+      network_throughput_bytes_per_sec: 24763792090
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9acfe70ef1f666c073825d105b6ada1a"
+      network_throughput_bytes_per_sec: 419810894620
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0123f78b1d0d6d1e4786c3a947b49f0f"
+      network_throughput_bytes_per_sec: 1040187582971
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1b6e5cdebf4cb96aef8697238fb6e7f3"
+      network_throughput_bytes_per_sec: 739107920712
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4d7bc6fd03d1833b4042c0e6b9d4ed90"
+      network_throughput_bytes_per_sec: 40119802
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b711d3e60296c168f7e3fa2f0b85abc2"
+      network_throughput_bytes_per_sec: 231162837
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b27c805cd6619c1801b55825d316f6c4"
+      network_throughput_bytes_per_sec: 376639450437
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "44303face361fc0189b66cc04e074d10"
+      network_throughput_bytes_per_sec: 365839044813
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4f35798f0136490d77605bad58190d77"
+      network_throughput_bytes_per_sec: 374512353
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ca5662195dafecaf9e272d6ad2d6ff80"
+      network_throughput_bytes_per_sec: 11773815405
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8beed9415f895e36d993111d52a6b960"
+      network_throughput_bytes_per_sec: 536877652797
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 2415919104
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "639e4d1d9f4f81f26ed6930cdba7cc9a"
+      network_throughput_bytes_per_sec: 799681210445
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "98dba134cb353258c8fe4711922d8f92"
+      network_throughput_bytes_per_sec: 271402067
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a7a1be29024f0e1494edab4b218c6393"
+      network_throughput_bytes_per_sec: 6338717477
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "1147e769ac8a215786b10b91a66d5f42"
+      network_throughput_bytes_per_sec: 632926677600
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4831838208
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "202c139b0b96fc7231ee505180ccee10"
+      network_throughput_bytes_per_sec: 870376278858
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2879ed2ada002e57328a5c87aef1126c"
+      network_throughput_bytes_per_sec: 7800253914
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3121c80c49a241613661eddf4bde073f"
+      network_throughput_bytes_per_sec: 202792355
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5ebad05df925d59d111b6bd34c68d519"
+      network_throughput_bytes_per_sec: 176305338377
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "bc112081346d68bbe7e7d73f99121ae9"
+      network_throughput_bytes_per_sec: 1195328261452
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5552290aa517a98e96ac8c7323a35766"
+      network_throughput_bytes_per_sec: 10377830562
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a432778d5cec2afe05e8cc98d2d67f5c"
+      network_throughput_bytes_per_sec: 512256128
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e6d868ebd3d51232037150fd5759eda6"
+      network_throughput_bytes_per_sec: 142267675701
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f5435413c029d037d8a19c23940706fe"
+      network_throughput_bytes_per_sec: 5853286757
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a42253af0fc00864313045e7d4fbab97"
+      network_throughput_bytes_per_sec: 60756489493
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "07d556652f3f2ad654b56cae0cb6703f"
+      network_throughput_bytes_per_sec: 4662492885
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "16b546b293eb5bdb11df1a4d6c771585"
+      network_throughput_bytes_per_sec: 874205410412
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "cedccf7581acce11c7d73cecc263089e"
+      network_throughput_bytes_per_sec: 8028420923
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a750a342c56fd1677f8a221d2e199edf"
+      network_throughput_bytes_per_sec: 36378573411
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "57b683a75eed5606d04c6a74cca23e6d"
+      network_throughput_bytes_per_sec: 41654613
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "bc347ac4aaf47aeee32bb13dcbbb7deb"
+      network_throughput_bytes_per_sec: 339409820736
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a6ed685e9583556b1f14bdc6faa6b199"
+      network_throughput_bytes_per_sec: 210623746
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "e1269a57114480043180144e81cf2d09"
+      network_throughput_bytes_per_sec: 41342200026
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "eb6c9e7bc87afd148a851fd2fdc9b9d1"
+      network_throughput_bytes_per_sec: 26683961
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4831838208
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5512ded6fba3c9f53245eed7679e099a"
+      network_throughput_bytes_per_sec: 443536136634
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d829422fa5dfaa33261c28ac102bb383"
+      network_throughput_bytes_per_sec: 115867596503
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4831838208
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7dd5ddf4466abcfaf631bbfb7ef0adf7"
+      network_throughput_bytes_per_sec: 554149270872
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5fff5806742a690449230f376b989c8c"
+      network_throughput_bytes_per_sec: 19559238
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0e96645123d7383e80f64dcfbb9fe9c0"
+      network_throughput_bytes_per_sec: 2067178825
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7b26c03b9cdd1cf06de66d0688d73fa1"
+      network_throughput_bytes_per_sec: 424152632643
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9b4bb436c4c1e606f1a29727a81a28cd"
+      network_throughput_bytes_per_sec: 786677836824
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fef7972b13c3169464b826898b352f1a"
+      network_throughput_bytes_per_sec: 11327629418
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6d1637f6f3b1660142e4351c6086deeb"
+      network_throughput_bytes_per_sec: 1182154884112
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "638032a952739f451c552c3c4697226a"
+      network_throughput_bytes_per_sec: 83561859983
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8f8cb9445b99056a6047be17515f22db"
+      network_throughput_bytes_per_sec: 42883036152
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9540f26284aa44bb0791e6cd49a1258e"
+      network_throughput_bytes_per_sec: 127015444249
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "de33fb8fd183a646b7eb27fe9e1dc583"
+      network_throughput_bytes_per_sec: 404377449444
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6c3dcfeb3f17d9cb610d821d0555565a"
+      network_throughput_bytes_per_sec: 2676735405
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6228a5bfb8af0370fc33b1fd40f8bddd"
+      network_throughput_bytes_per_sec: 59196432099
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "135daf9961de5436ee7f17d214db0a0c"
+      network_throughput_bytes_per_sec: 184534059
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e1ddbe870e4d4f1f18c61f244a753cb5"
+      network_throughput_bytes_per_sec: 180028500300
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "08a05c63a661dcb9051b6ed5df9a46d9"
+      network_throughput_bytes_per_sec: 1118418228455
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "19302cea84be4e0e7a4f3b7a99e29581"
+      network_throughput_bytes_per_sec: 3104206138
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "278036a81f9915afd1953815b2cfd216"
+      network_throughput_bytes_per_sec: 175229946524
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "144733171d135cfac307057c493a5334"
+      network_throughput_bytes_per_sec: 226648959
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2757cffbbcb51771cd873d1442e1640f"
+      network_throughput_bytes_per_sec: 981066370387
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1d3fb77d18352c32842f8f21a1ebc5fc"
+      network_throughput_bytes_per_sec: 1008488578985
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fc0d2bc70e0c6cfcbaf86b088a7eeaa3"
+      network_throughput_bytes_per_sec: 15731156985
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "55b42bedaf576f8698c72cf31a772c4a"
+      network_throughput_bytes_per_sec: 178397212
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e2568647e99d89ca3fae70adc8e36f7d"
+      network_throughput_bytes_per_sec: 24792547406
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "52d766cbeaba563dc479f81482c41784"
+      network_throughput_bytes_per_sec: 194108848574
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "56d33580ca50ca57d1d46c4551edd6d8"
+      network_throughput_bytes_per_sec: 17426697394
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ddbe05d3c08d33f2550c6c756f970206"
+      network_throughput_bytes_per_sec: 38271789
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2774cb74b12b0c636cc8de538c4565a8"
+      network_throughput_bytes_per_sec: 354355061580
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "0a8869903964281a33c703ed2efebb8d"
+      network_throughput_bytes_per_sec: 28374401298
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 8589934592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "8ff3b20102ddb82d99ab84f666232ff4"
+      network_throughput_bytes_per_sec: 799187705757
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "aacb45f958fd523d7d1a8b48c2bc674b"
+      network_throughput_bytes_per_sec: 1031920163609
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3605433eb2e5b3d3e7c09cf07fea49df"
+      network_throughput_bytes_per_sec: 242459332909
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "dda02568686126fa839c414ea78960cf"
+      network_throughput_bytes_per_sec: 33882352941
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e04e9385a1a532c6b53091b9d495029f"
+      network_throughput_bytes_per_sec: 146643731207
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e4bef6e93acd21e43df7d56b41a4d00c"
+      network_throughput_bytes_per_sec: 46489736200
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "a1ca18f43e4514d609516bd77ab71d59"
+      network_throughput_bytes_per_sec: 1749119248
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a8681fe97913d33400922723b1780be7"
+      network_throughput_bytes_per_sec: 144134158075
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 32
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "295c79dd5e32aabf4f8f0c2fe91aed68"
+      network_throughput_bytes_per_sec: 16982634
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8681c5d7a4bffaca051bd45077470731"
+      network_throughput_bytes_per_sec: 33190253784
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "344d1816e4512514f41b8a2ff2156713"
+      network_throughput_bytes_per_sec: 894201966125
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4294967296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b462b17ea2f4f8a5d9ae39b0e9c1d98e"
+      network_throughput_bytes_per_sec: 819200039062
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "674abef5b24e56150f4834c94bead97e"
+      network_throughput_bytes_per_sec: 92695898161
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0892f6b842da6e804dce5749aea0fb72"
+      network_throughput_bytes_per_sec: 1052753325
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "3b48323d232b92955c4c7bfb54ce952f"
+      network_throughput_bytes_per_sec: 592135697
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2b082881ba32be0a4cfc6afb9fb5f615"
+      network_throughput_bytes_per_sec: 80239604
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f9b2094c554802a33b27b8237c10fca9"
+      network_throughput_bytes_per_sec: 4904654991
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2fd3f846412f7f67ab09575e5a5926ff"
+      network_throughput_bytes_per_sec: 964577385971
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "55c8e77c4e1c969d602537c47acacd9c"
+      network_throughput_bytes_per_sec: 24396054
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "866d65f3e118830f08479fa04570b0b3"
+      network_throughput_bytes_per_sec: 9782952679
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "de63c7e44fd807a5a9dfabcb323a1b2d"
+      network_throughput_bytes_per_sec: 3432721854
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "488c2fe0b87f63ecdf1a72d392b02db4"
+      network_throughput_bytes_per_sec: 60606060
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a778901cd1f7090c0f4f83d05c6020d1"
+      network_throughput_bytes_per_sec: 747622784443
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "396e317ef47e5b66a7fb78f5a726f59f"
+      network_throughput_bytes_per_sec: 36325922276
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "842c759ba30c44ee9e2a7178d48c3e59"
+      network_throughput_bytes_per_sec: 173051380789
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "f132b3b6ed09f23fce8f2ad75b9c9fca"
+      network_throughput_bytes_per_sec: 438883292374
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fb3094fa513da052004b2ef1939a4570"
+      network_throughput_bytes_per_sec: 1226021019
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a803598230cc537323760e5dd054b2a7"
+      network_throughput_bytes_per_sec: 467722289
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fd4393c7df2b15343c275e71d4068b7e"
+      network_throughput_bytes_per_sec: 15929134708
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "46196ad5c87751165f16f8fee095d94e"
+      network_throughput_bytes_per_sec: 71348958175
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9fc18637504b8d4017d7055feca5a0e4"
+      network_throughput_bytes_per_sec: 8334997297
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7d3166872dd7c04df2eed3459af12ddf"
+      network_throughput_bytes_per_sec: 9902357128
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a78fd1e2e7ec25ec64bbb20938c3d6d4"
+      network_throughput_bytes_per_sec: 377140390840
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f1a82022bff64e8837e8a3e88afb846c"
+      network_throughput_bytes_per_sec: 19031491997
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9d0b32576aceaae25622e12dc312ca3e"
+      network_throughput_bytes_per_sec: 73059360
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0f3d765ea05d650f7701345150b0a3e8"
+      network_throughput_bytes_per_sec: 3655935
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ad966b90da174d324d688d8005dd5b10"
+      network_throughput_bytes_per_sec: 222911564625
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d09a8a0adbf7d70921c72b67354c1471"
+      network_throughput_bytes_per_sec: 743980089243
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "38f07465ee9f86899b2e307ef06bd10c"
+      network_throughput_bytes_per_sec: 37458541
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "21c1e3a914f652d5a816e3baef82edf3"
+      network_throughput_bytes_per_sec: 5532742
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7821c9efabf97e92775d17d59cbb488e"
+      network_throughput_bytes_per_sec: 96062540716
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e98a19174cf16fc36a57b6a0099db702"
+      network_throughput_bytes_per_sec: 69719148936
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ae30fe89309faf66b59d0e23791465ae"
+      network_throughput_bytes_per_sec: 2654280
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c4d11a42e6dc028cc0a9af905ebfb85e"
+      network_throughput_bytes_per_sec: 1037487335
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6495bee700fd5a8ea9f932e186751f39"
+      network_throughput_bytes_per_sec: 940696412773
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "db4c47ecb0509277a2ff679de757eb28"
+      network_throughput_bytes_per_sec: 44281081
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c5e98de9dde91f0dcd4182eeb155868c"
+      network_throughput_bytes_per_sec: 448146456420
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ae08caa85fae89c9c6a82c550cacb7a1"
+      network_throughput_bytes_per_sec: 524029624282
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7ce4526cb2ed59a09f17fa3e044f3523"
+      network_throughput_bytes_per_sec: 15814671814
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 18
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "22ae0ca92cbaef18d868e2276ec2ccda"
+      network_throughput_bytes_per_sec: 10055865
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "eb4c5f75e104336586cff20ab8b21203"
+      network_throughput_bytes_per_sec: 30145354185
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b00824efb464342f9dedaab473c62e14"
+      network_throughput_bytes_per_sec: 20721753794
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 8589934592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "52e66b58e47dce04393606e52c990d4a"
+      network_throughput_bytes_per_sec: 930751440095
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0bcacfa92dba52f7c5ec9af89cebc11f"
+      network_throughput_bytes_per_sec: 34124446758
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "09e42b8b606ec124354fd670abd7be2b"
+      network_throughput_bytes_per_sec: 751421757
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "37c2ae3d4f5ee2fb42f355b130b5f652"
+      network_throughput_bytes_per_sec: 1012383297127
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e5851c265006fc6123f1ccdeadd8a53d"
+      network_throughput_bytes_per_sec: 26373626
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6014748c7b090d01cf3b36d48ec2f75b"
+      network_throughput_bytes_per_sec: 43828645736
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "85c9029e1f9007c95c78844adf41a55f"
+      network_throughput_bytes_per_sec: 148702545557
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "8a71136a005d1079bf54125b237e273f"
+      network_throughput_bytes_per_sec: 6801162308
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d2365863cd0a0aca34bb0806e4a72a1c"
+      network_throughput_bytes_per_sec: 30440957886
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "990edd60bef34c6dc9b7eb732e9ffc14"
+      network_throughput_bytes_per_sec: 23291548
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "aec0f34e1ea9bdc2c3073ee2e37c1566"
+      network_throughput_bytes_per_sec: 12705345
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2b21229a2f1f4b5db63680fbec723b10"
+      network_throughput_bytes_per_sec: 184666249217
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4831838208
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ebf84be3c9bc5cfb8b52fc79e844e303"
+      network_throughput_bytes_per_sec: 907125157190
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "49f39c72a7ea3d783c8d05e421062934"
+      network_throughput_bytes_per_sec: 44625714848
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b88d21b48299fe619f4601a335de13b8"
+      network_throughput_bytes_per_sec: 258218275283
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "561e127a3f1e9a26ca7e11ae7faaa028"
+      network_throughput_bytes_per_sec: 9867237687
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "94c0e571ce8c54c6b7341bf5c8d8a84c"
+      network_throughput_bytes_per_sec: 130442167
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c8f58bbf48f0e0ccca0da43753fac89a"
+      network_throughput_bytes_per_sec: 30012505
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fa66a1984b680d1026e0b4f961b99e16"
+      network_throughput_bytes_per_sec: 384406483371
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "72535d6fee3ec01ee2a455a5d86d22a0"
+      network_throughput_bytes_per_sec: 84502005730
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5d32f935f98a3eae04cbed7a872a0143"
+      network_throughput_bytes_per_sec: 5655958
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "eb4a36e37bf7a362d251b8e7d0c12fae"
+      network_throughput_bytes_per_sec: 84036505344
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "8330d315febace3ee9783e0bf8ad5c92"
+      network_throughput_bytes_per_sec: 165674354506
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4310fea50c1c9a4f46aabd36c0cafa27"
+      network_throughput_bytes_per_sec: 137288599391
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f508d57e47eaca029bf97c6876b28f21"
+      network_throughput_bytes_per_sec: 1074219774
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "244ab456e7ce8411d7828b6f1bef8408"
+      network_throughput_bytes_per_sec: 6495143706
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a266703fec1b80887e365fe93c6fb033"
+      network_throughput_bytes_per_sec: 643935159
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a9ec49fee8d62910516b29c9173646be"
+      network_throughput_bytes_per_sec: 236356175
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "99ad4b0d46e619493ff4e75a3207a4b6"
+      network_throughput_bytes_per_sec: 4898544947
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "de0a7c5b4b27923d55dfd1ffb524dc9c"
+      network_throughput_bytes_per_sec: 461219399926
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "15042d4eea87e387d5aaf5acf461b431"
+      network_throughput_bytes_per_sec: 7536338546
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ff0039cca50563aa0bab3f8939f0a6a8"
+      network_throughput_bytes_per_sec: 751338242904
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2ec71af106fe82a4dc178edd93fa6bf2"
+      network_throughput_bytes_per_sec: 63781995133
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3d2f051f6bfadbaee1326ce0749781af"
+      network_throughput_bytes_per_sec: 28154226184
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b7e8b429f6ddc2482e878fc9b68ddf58"
+      network_throughput_bytes_per_sec: 553719742420
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "03e4cc5292fe3b4189be8ef411e71a28"
+      network_throughput_bytes_per_sec: 824800528490
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f5dd65cae4d8ea0b52e00a52438faec5"
+      network_throughput_bytes_per_sec: 12328646004
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "c8a5ae3ec7e98bc13ae6dc07d08213f6"
+      network_throughput_bytes_per_sec: 730140805841
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3e5770e52bd7e3f72c1c3fb39bf7498d"
+      network_throughput_bytes_per_sec: 346067859074
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "cebf72cb3c138fd2c067c9242fd68fd1"
+      network_throughput_bytes_per_sec: 639278158817
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 36
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ee4fa70393a994f17131bc1522c34b26"
+      network_throughput_bytes_per_sec: 23225806
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c39e349125978aa37ac5feef69007114"
+      network_throughput_bytes_per_sec: 1416103257
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 8589934592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e3b432bd6a074c6174fe512398414bde"
+      network_throughput_bytes_per_sec: 980882196863
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ecd47a7d6d5dd351a439bf3b40f0cf44"
+      network_throughput_bytes_per_sec: 3301894397
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b02b4733f86cab9b9f51ffd871fe1f6c"
+      network_throughput_bytes_per_sec: 195889737628
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "32304e6257208968c4af6fd14577f208"
+      network_throughput_bytes_per_sec: 909520431765
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ec9333b8b0f849d68a35fe13f12b6d13"
+      network_throughput_bytes_per_sec: 610295761
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6b33913249e785974e124d7ba2588033"
+      network_throughput_bytes_per_sec: 136029520295
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "5bec40f94a38e863464b468567f8ddef"
+      network_throughput_bytes_per_sec: 13659380454
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9dbbce82c36d6dd34dce63525e82f2ee"
+      network_throughput_bytes_per_sec: 420476920335
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "04d56708090d972bd7a3be6e5225038f"
+      network_throughput_bytes_per_sec: 423327437740
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "bedae4b4f79dbefd507eafaf6a933928"
+      network_throughput_bytes_per_sec: 43869419
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6c9281bce77dbfe2fb4b43a93aece927"
+      network_throughput_bytes_per_sec: 13091490211
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "795917a487235d7b8503ee24adc13bc4"
+      network_throughput_bytes_per_sec: 8837406870
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b42e9e80502972833d41e08e33708cf8"
+      network_throughput_bytes_per_sec: 1046173198641
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "67ea68efa9f0fb7b7e5c8c862447c38b"
+      network_throughput_bytes_per_sec: 2665124349
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "93a8fb178e1cca83c15b58e29c3cf090"
+      network_throughput_bytes_per_sec: 279777771189
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 8589934592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "42f0db01fb2a5e6ddcbf59749607c83e"
+      network_throughput_bytes_per_sec: 986716001157
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "a1b5e08afc59031dc54a246cad3370ce"
+      network_throughput_bytes_per_sec: 16841120
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b819a7de251e1c62bf88092d00ebfd90"
+      network_throughput_bytes_per_sec: 877673471285
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1402b639efee77b7ed2eae8e6caea1d0"
+      network_throughput_bytes_per_sec: 111455782312
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "25a0c663da11b4c0a52816654be4ad67"
+      network_throughput_bytes_per_sec: 270640999720
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a12521313f1aad79b4fa76dd0c9a5c34"
+      network_throughput_bytes_per_sec: 43523022432
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7480152b462cee9e3645191ff47c37c7"
+      network_throughput_bytes_per_sec: 1077584850772
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 64
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "297846412410fd5c247cd79d1c6bfdb9"
+      network_throughput_bytes_per_sec: 34353193
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "314f7b7b8c2a217a884988ee5f63b1b9"
+      network_throughput_bytes_per_sec: 5320128262
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "34b46885d0c52b38578b3aff13b40759"
+      network_throughput_bytes_per_sec: 663674975724
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "7f084ec07892dfc81b9bd01ce65fc823"
+      network_throughput_bytes_per_sec: 12154866230
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f9e1421da3c88895dc711a9245f49349"
+      network_throughput_bytes_per_sec: 73548724982
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "608c83a561efe5141f7e87e3c0457ff1"
+      network_throughput_bytes_per_sec: 328959286112
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e89041b3e93d292f4be3e5aaca1fdacd"
+      network_throughput_bytes_per_sec: 1011738445128
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7591798af4d69a22a0f5ba1c036ec758"
+      network_throughput_bytes_per_sec: 17873048339
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "169ae915f83bb68b439ab3184180e766"
+      network_throughput_bytes_per_sec: 4268231
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5cb01bf33f56194257eb0172939bfe03"
+      network_throughput_bytes_per_sec: 256799363252
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "cc566177a28456c23e186ddf526f7701"
+      network_throughput_bytes_per_sec: 10364185
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a4f74defa648b708509b9d9aac58dc80"
+      network_throughput_bytes_per_sec: 67302004
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "06b6562a2a50549328825bd2bba8b1e5"
+      network_throughput_bytes_per_sec: 531912073047
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "cdbf4e0987fd4845ea755fa494d2c9b8"
+      network_throughput_bytes_per_sec: 627472340425
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "97f0b4fdd4ee47fbcaab99ca26bff42a"
+      network_throughput_bytes_per_sec: 472005964354
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d0450d74652f1bb0e1668dd0fdbc3c6b"
+      network_throughput_bytes_per_sec: 215860015096
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "178762a3338e7551298bec7a0d57cb6b"
+      network_throughput_bytes_per_sec: 425779625
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b032e0c4ffa5762bf2d9d39f64735f4d"
+      network_throughput_bytes_per_sec: 84480825008
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f9080d4e9dd3968f3ff5b3fb191ec37c"
+      network_throughput_bytes_per_sec: 1413266626
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c096394a000aa1bcc45b19442dac26f3"
+      network_throughput_bytes_per_sec: 6542151235
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3747a8e557a2c2f28a6cd2e83210a3e4"
+      network_throughput_bytes_per_sec: 963911787148
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "be7060f2b6431ec7dae82c0852433943"
+      network_throughput_bytes_per_sec: 4978426010
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6919d2ab1a19acc0820299025f7c8ef1"
+      network_throughput_bytes_per_sec: 1268307787
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d6dd1c5eba55987d325ee6278cc2d534"
+      network_throughput_bytes_per_sec: 715484450130
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "364732bff4314d9fc1fbef7acf1f4228"
+      network_throughput_bytes_per_sec: 1487411233
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "f4a9a9fb7485cc34f378de219fc0562c"
+      network_throughput_bytes_per_sec: 778174033369
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2ed285ac2da7222c03942e03802a6133"
+      network_throughput_bytes_per_sec: 5277287917
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "9eec7e1a4386fee97c05dadb49248d21"
+      network_throughput_bytes_per_sec: 7631508125
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "215574b1f5f7e6fe60eb3cf98400e2f3"
+      network_throughput_bytes_per_sec: 18179881
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "5c2d49eeaca07ce847d8ddff631e0c30"
+      network_throughput_bytes_per_sec: 35631915182
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "57207937c2cd7cc35287c6d5e6326a8d"
+      network_throughput_bytes_per_sec: 787053610945
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "49e718729a6c51e3c1596faa48cf4c7a"
+      network_throughput_bytes_per_sec: 22910568710
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ad2b4fda35aa4e08372798ba9bed0ed6"
+      network_throughput_bytes_per_sec: 48813559
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "696e377d02f5910b959ad225b9ee7295"
+      network_throughput_bytes_per_sec: 7133946552
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "bb3d46f427b415e838e83ae4a4f55c2f"
+      network_throughput_bytes_per_sec: 2368999421
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e209f09945e6c2ad66bdc09df07fd9b3"
+      network_throughput_bytes_per_sec: 419156774
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8b4f40721708026be75d1dedf7092c6c"
+      network_throughput_bytes_per_sec: 7968871595
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "764b0f42e20ee15104851abfa286a07a"
+      network_throughput_bytes_per_sec: 872268690859
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d089115600aa3c1e465ed31821eacc6c"
+      network_throughput_bytes_per_sec: 119956070774
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "cf1ce1c5cbf3cfb48c694f6fdccdc3c1"
+      network_throughput_bytes_per_sec: 49480285
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "715876e9e5ba70e71ad896016bbfd123"
+      network_throughput_bytes_per_sec: 40539952252
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "19ae6d2a375133c917a4e101e520beb0"
+      network_throughput_bytes_per_sec: 591695004320
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "a377e6b7052423281f2f99772aee54ce"
+      network_throughput_bytes_per_sec: 189024144574
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b1a672f35f3d1d53684dc6ceadf3c01f"
+      network_throughput_bytes_per_sec: 187561055
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f522e05ee29d774885d0a49ff707b9eb"
+      network_throughput_bytes_per_sec: 874754909798
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ef47ee53bbace705945409ea69d4f6be"
+      network_throughput_bytes_per_sec: 870327992068
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c4472383e01f1f1953351d001842a7c0"
+      network_throughput_bytes_per_sec: 68027507460
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a30198bdbd0cddf17ef2822e52b98dd8"
+      network_throughput_bytes_per_sec: 101566834560
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "59a7a88fc82a7188a64707eb6ab21462"
+      network_throughput_bytes_per_sec: 262694015595
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "79cd8ecaf8523417daad5633a9914297"
+      network_throughput_bytes_per_sec: 409525677550
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "79495fc7b08d0e5be602b818416f1f30"
+      network_throughput_bytes_per_sec: 36394286243
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "296b63bb758383f86e3c69f6f3ea8eb2"
+      network_throughput_bytes_per_sec: 335308858668
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1fdc074dfe0d5b69a461d08a6c7c2c16"
+      network_throughput_bytes_per_sec: 533965140451
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2db7f2b194739e96e4ac26097d5799ca"
+      network_throughput_bytes_per_sec: 280189711
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "d5c56c4f37c708a3f2c2701b273a3230"
+      network_throughput_bytes_per_sec: 156223429746
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e86f3e234738bed9164db2d44d7702ac"
+      network_throughput_bytes_per_sec: 10886822
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d7708d98099b8408a524704336e251e7"
+      network_throughput_bytes_per_sec: 9429020
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "46b2620ff39acd6110602f421c933701"
+      network_throughput_bytes_per_sec: 432432432
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c57bd8fcff8d6563f532ff5d7999a7df"
+      network_throughput_bytes_per_sec: 164102564
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "825d3a9bdf3343f5d26d0dbaa8dcd049"
+      network_throughput_bytes_per_sec: 15304997664
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 32
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7baae1e1b20a27980cade84146fda7be"
+      network_throughput_bytes_per_sec: 6270821
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3d553dce12f10066f334579a0d6dba00"
+      network_throughput_bytes_per_sec: 4585403
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f16df5644e9c08d4589c33dec6e18979"
+      network_throughput_bytes_per_sec: 1134160840283
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d5a1e1f417c92597e7ffe78052369891"
+      network_throughput_bytes_per_sec: 1207547169
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "20c87f0d3f9848fd31c1bb2d2ecaf7d9"
+      network_throughput_bytes_per_sec: 7637560819
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3c5499b0e1fa7abdab52b6360ff47522"
+      network_throughput_bytes_per_sec: 1138645205625
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "12f0549675b4b5af9a4c826efc3968c3"
+      network_throughput_bytes_per_sec: 489509296644
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c749ae69c44882518d174b6b7dae17ea"
+      network_throughput_bytes_per_sec: 5043076
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7cb07205f6c0348b1deab1275fb68cba"
+      network_throughput_bytes_per_sec: 97765896
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7196b07fb03b514c19871618eb4080bb"
+      network_throughput_bytes_per_sec: 1039190808015
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "44b2c631db850a475163e9aad83021d9"
+      network_throughput_bytes_per_sec: 846495127935
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a1e82b0b3aaaa1efac4641dfc535b772"
+      network_throughput_bytes_per_sec: 2200967221
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a73fbee368556b435412dd565f4fc45b"
+      network_throughput_bytes_per_sec: 69881710
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "c37d2241f3f9dd9d203ed0a5ccde92ce"
+      network_throughput_bytes_per_sec: 6162745016
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f2dc462a4169f47b53435a74c095a2b8"
+      network_throughput_bytes_per_sec: 788173905703
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7ef67b53b672b0ab198e147337aac50d"
+      network_throughput_bytes_per_sec: 50724458204
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fd69e2f832900b720db65be6490457f2"
+      network_throughput_bytes_per_sec: 284898002156
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "d0ebf7d4a9e695450588d892d6d4a2ae"
+      network_throughput_bytes_per_sec: 80898876
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "131888acd40471dd2ce245b2fd30d899"
+      network_throughput_bytes_per_sec: 1234827724695
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b57a1e35eb36c73cf2af785d7c3bb1b9"
+      network_throughput_bytes_per_sec: 462899251566
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e357a5ea9685f2d6694363dd6dcc2cb5"
+      network_throughput_bytes_per_sec: 891730511197
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2415919104
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "94aca6687422774038513b7d0883cd26"
+      network_throughput_bytes_per_sec: 839512661176
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f7a37cb8d0f63d78ed28a2ec489901f2"
+      network_throughput_bytes_per_sec: 5708848
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5d17f172e0e139e0cbe69461f4a4a39e"
+      network_throughput_bytes_per_sec: 879954683730
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "a3cf67340adcf461b001c266c4df9b73"
+      network_throughput_bytes_per_sec: 36277885413
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1a8cd3542fd9e21d763e10c2c7aee170"
+      network_throughput_bytes_per_sec: 10127472527
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7f534b0f9ed3f472099acd491da59685"
+      network_throughput_bytes_per_sec: 546022738449
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7b0ba422a140af9cae28476aa62d619d"
+      network_throughput_bytes_per_sec: 41114178168
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2fb7abe8e2cc03da5b4c5ef22be8258e"
+      network_throughput_bytes_per_sec: 17366136
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "547a604517d21a8e150e3b89b4cc4a2d"
+      network_throughput_bytes_per_sec: 832478024
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "266665901b5f0e28576724fbebe4b3b5"
+      network_throughput_bytes_per_sec: 2004131782
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "cba957744017f18bfbabf213168afba4"
+      network_throughput_bytes_per_sec: 57124427979
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a6ec1ce394c9e3024e240df72ec7e598"
+      network_throughput_bytes_per_sec: 10855123674
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7b6c39845e48a5ff61d4c89a25b1903f"
+      network_throughput_bytes_per_sec: 25047732291
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "50f8f91e357bfebcb2e9256dc7811a41"
+      network_throughput_bytes_per_sec: 396103610795
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0381a1319ef5b2b0af2ad536c3ead13b"
+      network_throughput_bytes_per_sec: 34615644
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9a89d1083b9e12ffac503d187ffe12eb"
+      network_throughput_bytes_per_sec: 222617097565
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 18
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "320c1c3599ef3199379b4d77a575ea25"
+      network_throughput_bytes_per_sec: 17664376
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "66cdc8029a623ae68dacfea29d4f4c97"
+      network_throughput_bytes_per_sec: 124751269035
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "af6d14323a665322c56a35d29602f374"
+      network_throughput_bytes_per_sec: 2820612877
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "806c5ef18a1be16bbc4781bbdba3f840"
+      network_throughput_bytes_per_sec: 1039652546
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5f1582204a89020bc7cbea4d7754e68d"
+      network_throughput_bytes_per_sec: 654354611404
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8a8ec4d7278c768893c0bf7db558bf38"
+      network_throughput_bytes_per_sec: 1605574912
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2f4f94ba7e7bbe8250c600fe4337ac05"
+      network_throughput_bytes_per_sec: 4558708959
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "956288633f94537f620b1f4d8ae51440"
+      network_throughput_bytes_per_sec: 20901371
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4831838208
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "235b13a6a5b5987641e289142da4ad48"
+      network_throughput_bytes_per_sec: 1021862976106
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 8589934592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "72147db6178770ba1d0363440cfa1b89"
+      network_throughput_bytes_per_sec: 386576738001
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ab23bc00f80963c5f188f9de47adfa1e"
+      network_throughput_bytes_per_sec: 1046414137459
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9812a402bf3a6fac277810bcbb5a2ab4"
+      network_throughput_bytes_per_sec: 4302521008
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "03287064f8f38dc86147b5f6bda64f27"
+      network_throughput_bytes_per_sec: 37490203556
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2415919104
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "429db1e7b51367c50b9ae04cf1f42839"
+      network_throughput_bytes_per_sec: 434851334316
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6e22619f5b4259b1a6ca117bc6bcdf8d"
+      network_throughput_bytes_per_sec: 773223767441
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "822923fbb88bd07b8a9ce5d7a64e4c9f"
+      network_throughput_bytes_per_sec: 208393632
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7760941710dc9ff23b15ccd0c67e774a"
+      network_throughput_bytes_per_sec: 603763025073
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b7121b60d95daef4c12b6f1e38253568"
+      network_throughput_bytes_per_sec: 39702746365
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f480352e79a4c23d02492cb303a4e40a"
+      network_throughput_bytes_per_sec: 436018480872
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "db6df62d62a99c9e858cef6fd8687ded"
+      network_throughput_bytes_per_sec: 862714866046
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5265f7f4d78d3cd60c243550ce4eb67f"
+      network_throughput_bytes_per_sec: 70339332
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ef59ed83b98ec0b254373aeb0d4f4336"
+      network_throughput_bytes_per_sec: 27941825761
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9d0033dda8d414acba912bfeff83219c"
+      network_throughput_bytes_per_sec: 26014378335
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1fd18cc9c5657fb74ee88a540b00fb77"
+      network_throughput_bytes_per_sec: 6494500049
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a439ddfa229c6b74c60a303f38c1b32e"
+      network_throughput_bytes_per_sec: 15988723231
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "140b06dc23bfaa27296f5c3c459077bf"
+      network_throughput_bytes_per_sec: 405624876
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0231e3e156d2abdf69bd17f611f9b632"
+      network_throughput_bytes_per_sec: 157928643148
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d4461b902ea1d547aed5524b81b34b24"
+      network_throughput_bytes_per_sec: 28752737
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a10436927db1fcdef22a500c69b0166b"
+      network_throughput_bytes_per_sec: 92012779
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1ee5cb7975b0d3f7fb736115b5f3020f"
+      network_throughput_bytes_per_sec: 4712118
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "21aa8c6f528768bdb35bceae9b9eb223"
+      network_throughput_bytes_per_sec: 152360798
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0694616e03a29a8a4d03dcc411edf25f"
+      network_throughput_bytes_per_sec: 1001583726041
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1c2111d9e8564fee032ea000266fbc5d"
+      network_throughput_bytes_per_sec: 2671304347
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d17c84c474b5d7b7ad45f3678d96c036"
+      network_throughput_bytes_per_sec: 41163438
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "686753e33c07286697750a942ac5ac4d"
+      network_throughput_bytes_per_sec: 11719465
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "90f2b5650dbd29b6c6efc85105c4dd60"
+      network_throughput_bytes_per_sec: 996974766945
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4ed2a1dc3043904e2b39c70ff38af809"
+      network_throughput_bytes_per_sec: 283381898985
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "250b724f593df67500bf15751025148f"
+      network_throughput_bytes_per_sec: 164870440251
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a043abb4110175c528e5dffbff6a7ac4"
+      network_throughput_bytes_per_sec: 277547908946
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "eb52184cea1b0a07af7467d56f5e8438"
+      network_throughput_bytes_per_sec: 10303054
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b9685b404f832aaa635168228f3f9451"
+      network_throughput_bytes_per_sec: 36154468554
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "151fa3463434ec3e4b35d6fcb5be6813"
+      network_throughput_bytes_per_sec: 4376587914
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a636bbf3213207cd8e528f886910049d"
+      network_throughput_bytes_per_sec: 648648291978
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3f0c5adeccac951473c96bc42ab11d59"
+      network_throughput_bytes_per_sec: 2138484630
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2415919104
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "720cc1ac9af759f36d12545b1cdf09ce"
+      network_throughput_bytes_per_sec: 455575254250
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "47664fbfa5070a32a20146ce310dbe30"
+      network_throughput_bytes_per_sec: 1571846
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ae9faa29870faf19b14dee1a7ca64106"
+      network_throughput_bytes_per_sec: 3282387
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c0daeeb32a7370fa17378540179b871c"
+      network_throughput_bytes_per_sec: 1387533875
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 128
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "881d726d05954327d427bf442981065b"
+      network_throughput_bytes_per_sec: 31128404
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5ccfb850b008ba099d6bbb1133a462f7"
+      network_throughput_bytes_per_sec: 18288265662
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "771626896fbf4c85b782821e1a5284f5"
+      network_throughput_bytes_per_sec: 3330081300
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2415919104
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4bfe3d084499efb27520a85553be4694"
+      network_throughput_bytes_per_sec: 828721692233
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1ed676b9b1ebf04e562bb13c48da712e"
+      network_throughput_bytes_per_sec: 166582315
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b320503b9e25630c372d47463173bad5"
+      network_throughput_bytes_per_sec: 309760688
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "63b99c24f4e4021cb3790072ed9d77e7"
+      network_throughput_bytes_per_sec: 2082755990
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "f472090a6fa1b230b91192a7bb664fb5"
+      network_throughput_bytes_per_sec: 10049478634
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f370f2dd44c1b30a12ad54313e8695d9"
+      network_throughput_bytes_per_sec: 20231843791
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e08743be70e61dc63555eebc033a8f3d"
+      network_throughput_bytes_per_sec: 1068978456000
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e08c90a23a7f37bbd12e5dfd668928de"
+      network_throughput_bytes_per_sec: 310215245227
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "62b079e485d58d6f46de5b9181ee0b34"
+      network_throughput_bytes_per_sec: 357414937130
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4831838208
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "bbf07a8fda40a85ced209ab96a345f5e"
+      network_throughput_bytes_per_sec: 879326337037
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "200ef10be6d4b6525c06650c48d7857d"
+      network_throughput_bytes_per_sec: 484215781473
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c73482c9bd671e0dea96de33765ed1af"
+      network_throughput_bytes_per_sec: 2778420
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e93e120fc8e0dab1150081550583ee27"
+      network_throughput_bytes_per_sec: 482961914
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4294967296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "627ce80f3a20517f12256300a5899029"
+      network_throughput_bytes_per_sec: 462286329093
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "e2a93459c02f42d02511fbc8d65f9f80"
+      network_throughput_bytes_per_sec: 629119304779
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4294967296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b6768f5b96a2cf0aa4a49cb5ae380395"
+      network_throughput_bytes_per_sec: 925514289742
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "85fd86e3eddfb22e4b4149db72ce51a0"
+      network_throughput_bytes_per_sec: 83550913
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c88cfdbe9b906502ee28d29106a184a0"
+      network_throughput_bytes_per_sec: 172043010
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "849b840bb8e4aa4ebb59f0614ba6df2e"
+      network_throughput_bytes_per_sec: 227050997
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7c6170bc99b6adbb0ef48247708cfd7d"
+      network_throughput_bytes_per_sec: 663255591464
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 128
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a163bb2094275de50e3f3f664f2bc469"
+      network_throughput_bytes_per_sec: 69095816
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1bb7a06887ec86464873822f3d36fbdf"
+      network_throughput_bytes_per_sec: 68530636
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6f141823228aa4112e884fcb843d91e4"
+      network_throughput_bytes_per_sec: 204982206
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "552ca5e15fdc0eecfbd5c79355c93798"
+      network_throughput_bytes_per_sec: 762913907
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "691d5ae577d7a53b68f477eb318bbaf3"
+      network_throughput_bytes_per_sec: 2933163590
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 8589934592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c8b57c15bacced3b9c5884db97550fda"
+      network_throughput_bytes_per_sec: 554606678987
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1a112f7d38ef757607b73f5b62635b10"
+      network_throughput_bytes_per_sec: 4794725
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "37ce38c049a7082bb64385ce6a68bc48"
+      network_throughput_bytes_per_sec: 1042937871754
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ffdcd092800fda6990c22e7fdf4d252b"
+      network_throughput_bytes_per_sec: 45836493627
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "92769d3711c1072cce77cf468d4b5e8c"
+      network_throughput_bytes_per_sec: 29265716797
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c57531ad09a8edd7407a0b542dca3d1f"
+      network_throughput_bytes_per_sec: 2605413810
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "18ddf13f5ae783b382da2c77d67acc92"
+      network_throughput_bytes_per_sec: 1136637803129
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b8e2be7c0c8fa7833677d7f42971ed37"
+      network_throughput_bytes_per_sec: 28103044
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "beafaa18dc8c97d6d57d2fed08c72e96"
+      network_throughput_bytes_per_sec: 120692449355
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 36864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d695f9cc4236d3fa15ecad0f08a07bf4"
+      network_throughput_bytes_per_sec: 6252374491
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "19cb7f3ee228d3b71b2afb53eb0c4b40"
+      network_throughput_bytes_per_sec: 49884681255
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c3d8fb57732a0e5ab4363c6d7b79dd7b"
+      network_throughput_bytes_per_sec: 306075382201
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "996616258c65d33d3251e2b99965ffdc"
+      network_throughput_bytes_per_sec: 447171821901
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3bf70062b89e5b703e2fbc2aba5c51e3"
+      network_throughput_bytes_per_sec: 1105691897057
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "86a0c1ed1fd8f110b160192897f9b809"
+      network_throughput_bytes_per_sec: 17071853427
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f740248bec6f0001ead1692647854f6a"
+      network_throughput_bytes_per_sec: 2591792
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e5657e766ded6b82fad7f28ca75b1ebb"
+      network_throughput_bytes_per_sec: 13477940
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "9288b9bfbbaca114c2845c23eb69dc2e"
+      network_throughput_bytes_per_sec: 2504432895
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "214bfa359a8fc4ece9f9131074291c98"
+      network_throughput_bytes_per_sec: 2590219224
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ffb7bba447a0ecada447b55e96bc8b35"
+      network_throughput_bytes_per_sec: 16584070
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2766f5983e244e584941a79dc52de1eb"
+      network_throughput_bytes_per_sec: 1580246913
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9c010a1d91a7b9efb2b932d2a35a9944"
+      network_throughput_bytes_per_sec: 184819952410
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "a9191665e1b07d43ebae2b55f6e498d6"
+      network_throughput_bytes_per_sec: 1149163003
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5577bc1bba24e5477bdeb5fdc39118b0"
+      network_throughput_bytes_per_sec: 59128002706
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "f51970c08b1166fa9c16dce41b547514"
+      network_throughput_bytes_per_sec: 356421858
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "cbfd6c707e871327ea792c814c011d57"
+      network_throughput_bytes_per_sec: 175876551492
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "babaf1f4ba7f2fe6d90fcc20e26bb57c"
+      network_throughput_bytes_per_sec: 4081262109
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "617f5f3dffe87f8902ae091e57d1d44e"
+      network_throughput_bytes_per_sec: 85998587525
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "bffd6505c6a4997d4b57267b2fb4867c"
+      network_throughput_bytes_per_sec: 276151497028
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "f064593627c4c175ddd193b71525ebd8"
+      network_throughput_bytes_per_sec: 596976052946
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a4b9982926a69acdaf34a0d7adb59776"
+      network_throughput_bytes_per_sec: 72916800593
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "24b6f45cc3f050ad127a0be69069299a"
+      network_throughput_bytes_per_sec: 117329752713
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 9
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "32c15387fe61dbf5c7c655cc01d7a4c4"
+      network_throughput_bytes_per_sec: 7992895
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c4147e53a9d65387a790beda10d4ca96"
+      network_throughput_bytes_per_sec: 354222055401
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "464701bf7b28bc5fae7e7f4c428d007d"
+      network_throughput_bytes_per_sec: 309639215171
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "939d449f6d7b76b9aaece49ab555c99f"
+      network_throughput_bytes_per_sec: 362720403
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9f8b36645d75a64eb0b44dfa5c004197"
+      network_throughput_bytes_per_sec: 810684992698
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f05054b2857e423b36d19c39b5ad00f5"
+      network_throughput_bytes_per_sec: 290781987089
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d09e3dd89103bfac98cd9faa00957541"
+      network_throughput_bytes_per_sec: 60676287
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e4986cb30e395fdd3a6261c013f47154"
+      network_throughput_bytes_per_sec: 62864268585
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "8cf8122d615325e0d9ee400d08d2b1f3"
+      network_throughput_bytes_per_sec: 669194540618
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "06e306609646a1c51e15861f2493ece3"
+      network_throughput_bytes_per_sec: 382186401523
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 8589934592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "57d1b58884433b278cb49bf84afa5ce2"
+      network_throughput_bytes_per_sec: 498806178571
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 8589934592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0a2102c3eba0e256a27b56adffbbef04"
+      network_throughput_bytes_per_sec: 887341193008
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "be851c3d5ca51bf1d815cfd391a4ef8b"
+      network_throughput_bytes_per_sec: 26944847
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e962cf13cb39cad4116821075b138824"
+      network_throughput_bytes_per_sec: 307276819
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8ac2223697feb59e42b407b363b429e8"
+      network_throughput_bytes_per_sec: 13709063
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "30805fcfc6af817069cb0d66a035b244"
+      network_throughput_bytes_per_sec: 2907542147
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "36b484cddbeaa6e1877ec4c30d18e32a"
+      network_throughput_bytes_per_sec: 979005426854
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9fc35754c8ad852660aa0145feab10d3"
+      network_throughput_bytes_per_sec: 134557890346
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a549a338fdff96ed57198d66936900d7"
+      network_throughput_bytes_per_sec: 1325661680
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "419609f8aef65b258e3f2c0ab8e765d5"
+      network_throughput_bytes_per_sec: 38030465689
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "812d6f6c371ec4b170c3ef75be9f8d9b"
+      network_throughput_bytes_per_sec: 725336561429
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4d6533b1bc5eb3a7b507798dd60ac94a"
+      network_throughput_bytes_per_sec: 17895145631
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4294967296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e516a808200c364d595089621cff3912"
+      network_throughput_bytes_per_sec: 902815443726
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "99498af0ce90ed405692d77995958dbf"
+      network_throughput_bytes_per_sec: 918890411802
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1003c55f712d46432a51abc61abc5655"
+      network_throughput_bytes_per_sec: 194906627563
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9b1abc5b0d9afb03e9d481b1d0ba1943"
+      network_throughput_bytes_per_sec: 6744237102
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "203189d22601bf7472c417edc2e561ba"
+      network_throughput_bytes_per_sec: 174969623
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "98bb874824eb6c1abd4a45b6be90bf2e"
+      network_throughput_bytes_per_sec: 204445060658
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f420db42af18d62392d3d1a09c4f576f"
+      network_throughput_bytes_per_sec: 7721472754
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ce88723318b0252e3e1f4531b187d47a"
+      network_throughput_bytes_per_sec: 235348195
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "79f004724e523e4e428daaef16516fbd"
+      network_throughput_bytes_per_sec: 4431035518
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d9688d25a1898b26008bf0b560b34a1c"
+      network_throughput_bytes_per_sec: 147613993101
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "99e6b58c84203380b7235e64939d052a"
+      network_throughput_bytes_per_sec: 97396566462
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "dbd75500ac48dad7c07c116518e73014"
+      network_throughput_bytes_per_sec: 2490810810
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0cbb75a389e84796d6e4b74e0cb83bbd"
+      network_throughput_bytes_per_sec: 22153952
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 128
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "141affadaa3a6929813dff3e94420cee"
+      network_throughput_bytes_per_sec: 69805489
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "16175a1bd54ff30d3f195f71cde23293"
+      network_throughput_bytes_per_sec: 109469933184
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "a048594e627771297a5b0c7702a922c3"
+      network_throughput_bytes_per_sec: 159153594949
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e42fe7d123014c559fdbd0f343230b46"
+      network_throughput_bytes_per_sec: 212254291967
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c4293d1caf3429c6791b5baac69ac10e"
+      network_throughput_bytes_per_sec: 473635522
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b4295efd8dfb2de304332f97e11ce21d"
+      network_throughput_bytes_per_sec: 2487175
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "e5430dbd2ee86047758efdc48e0edccb"
+      network_throughput_bytes_per_sec: 23259047302
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "40b518cf17bc9388b76f37107af39165"
+      network_throughput_bytes_per_sec: 17208413
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "17e69bc5316190bad901345676dbf885"
+      network_throughput_bytes_per_sec: 572132587641
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "a028a47d2e73afae5b16b9c497c169fe"
+      network_throughput_bytes_per_sec: 181753638
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f22607dee96c4bc0661cb3eb0f71db6b"
+      network_throughput_bytes_per_sec: 13624751
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6752b9beb9f902a20e15f9e1c755d70b"
+      network_throughput_bytes_per_sec: 254776706892
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1d04910387e75301b9d8334a392d4846"
+      network_throughput_bytes_per_sec: 5003822
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 73728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "bb425e81ada134a5bc74dfe69dc717da"
+      network_throughput_bytes_per_sec: 9070866141
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "404f1f94507be53e5974fe2861073566"
+      network_throughput_bytes_per_sec: 949742043
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 8589934592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "467cc9f389d8473e799a8f176396117b"
+      network_throughput_bytes_per_sec: 1046987147159
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c57190dd6fb280d62642e2dc178f908e"
+      network_throughput_bytes_per_sec: 888985887267
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "98ff223effd98c67382004cacbdfbf2f"
+      network_throughput_bytes_per_sec: 9584425
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "751449070a9f655e8dc8a9fef1e16a86"
+      network_throughput_bytes_per_sec: 444396224
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c8a20f3f8dbf7d727fdac860299f19a6"
+      network_throughput_bytes_per_sec: 3329912097
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "30e78409b9f65576193499c1bd934b95"
+      network_throughput_bytes_per_sec: 47922126
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c0d799fd2e5393b458955091b043c278"
+      network_throughput_bytes_per_sec: 97505672307
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "445e7b0bc28c97347b5bf232be9646ee"
+      network_throughput_bytes_per_sec: 9261729790
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "18ab41dcd38c8cb0f267de2bcffcb995"
+      network_throughput_bytes_per_sec: 762891281569
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0420fb60faafd423193ad3a9c235a7ea"
+      network_throughput_bytes_per_sec: 1042228538691
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 9437184
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1806ee0359e2eef206c322d564a5317b"
+      network_throughput_bytes_per_sec: 61758441966
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "76893ffbc1f05a669a18e9d6518c4fc7"
+      network_throughput_bytes_per_sec: 1194382114715
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "7605d652ffcd91cb700053c35ea8f0a1"
+      network_throughput_bytes_per_sec: 20428927680
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "65819a9772b003e926ba878d279d7350"
+      network_throughput_bytes_per_sec: 875260430
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b03e1600dca6481cef26833b4f932b94"
+      network_throughput_bytes_per_sec: 447416288868
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0b9f5f196b1a66e412bf040433b6a781"
+      network_throughput_bytes_per_sec: 16652572
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7f4e2a97fa93c86fcd45c52e5f947d1d"
+      network_throughput_bytes_per_sec: 385245850256
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "82a2702a7628c3fa247a320c163a49c8"
+      network_throughput_bytes_per_sec: 63695896328
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3de997b1e7089afc73b8638024204ee3"
+      network_throughput_bytes_per_sec: 372464968167
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "00c3fdc488e614d48a82944d6e1d0301"
+      network_throughput_bytes_per_sec: 78688524
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e35becba03ed458c09783177a80683d0"
+      network_throughput_bytes_per_sec: 977933860960
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 64 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 66 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 68 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 70 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "ecdf1107c909dfd9d6e970662be2eca9"
+      network_throughput_bytes_per_sec: 177647380495
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b99f91255c235ea04acd7c64cad7906f"
+      network_throughput_bytes_per_sec: 66718790
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2aae498a37d3110f7b8d6fbe52214b52"
+      network_throughput_bytes_per_sec: 10835978835
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2b0a3dbd2a37fee4531ee7c5e8cfa20e"
+      network_throughput_bytes_per_sec: 197074672
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0683d460ad1e181b66aaa43a2262f3cd"
+      network_throughput_bytes_per_sec: 104356687898
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "731095728a7e74f932f0f1fcc4adda8d"
+      network_throughput_bytes_per_sec: 2409588940
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "67b4d6e561280baba89db971ba00d145"
+      network_throughput_bytes_per_sec: 1146686
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "faafbab31952de3b7de9f4ae92cd6b51"
+      network_throughput_bytes_per_sec: 775170325
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "295e1a8c29ca3434b85cfd230294f7a0"
+      network_throughput_bytes_per_sec: 482755905964
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "bf07b7952e355e3c2e02d19e6671a616"
+      network_throughput_bytes_per_sec: 1615380823
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 64
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3cf4d12163c36b4774e5193956d35fd2"
+      network_throughput_bytes_per_sec: 27234042
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 603979776
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "c4f6337fa4392ce6cd580c6c4581a480"
+      network_throughput_bytes_per_sec: 790416501936
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "632949243fd8809732d26e1ebd2d5537"
+      network_throughput_bytes_per_sec: 296527262765
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "53a83416428aab71ad0c3fec32f47538"
+      network_throughput_bytes_per_sec: 2587492103
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d3c46279e15dcd7c2ce5cc599af9bf33"
+      network_throughput_bytes_per_sec: 40416798
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16777216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fdac31ad41ce97e8a4f37d2bc9d46eff"
+      network_throughput_bytes_per_sec: 851710989554
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1df569705fda0bf64cec6697bc12714c"
+      network_throughput_bytes_per_sec: 11928429
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5305db11839fb1f55eb875314deed0f8"
+      network_throughput_bytes_per_sec: 25427095
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "87aa78ee3aca69a0109fb90361a09bef"
+      network_throughput_bytes_per_sec: 5029689
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9e481033ce7178d86b0bcef09b5e16e5"
+      network_throughput_bytes_per_sec: 7047768
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d83b97e06cbaf1ad485ce2b14a10df04"
+      network_throughput_bytes_per_sec: 72701097004
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f52c852c85aab8be5170ab6c86c9b916"
+      network_throughput_bytes_per_sec: 623822316234
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 4294967296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "033a0dc1cc94025cedc2cd57464b6325"
+      network_throughput_bytes_per_sec: 673070042521
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4096
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "acbfac13e9c6c6ea09fe53ede6783701"
+      network_throughput_bytes_per_sec: 763040238
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4e37059564a1997debfb873d9199f951"
+      network_throughput_bytes_per_sec: 10240000000
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ed443fff632ec7dd16f72bcd74ae58a5"
+      network_throughput_bytes_per_sec: 1005185492890
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ba550dd31dba0b762dd3fd5612e92494"
+      network_throughput_bytes_per_sec: 1287755186430
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f3099d3b62c590d7f97cb9c3bded7e74"
+      network_throughput_bytes_per_sec: 90395480
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6b88b367042005770a238a407686639c"
+      network_throughput_bytes_per_sec: 5787926
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9efe55d1f20110dec2e1ddbd25f30e7c"
+      network_throughput_bytes_per_sec: 122777685262
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4b47796614239f45f81f729bf1a450a2"
+      network_throughput_bytes_per_sec: 6502877555
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e9dbe2868e8177d55989bdd71daf868a"
+      network_throughput_bytes_per_sec: 33149886
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5a296c385199d05fbb54eaa5bff3854b"
+      network_throughput_bytes_per_sec: 630196936
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "aa300e0510a1017de1e302215f998825"
+      network_throughput_bytes_per_sec: 308364396810
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "17b143d6bc89e9340fb81e4da846edc6"
+      network_throughput_bytes_per_sec: 267312032630
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4831838208
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "85cf791f151a51abbc2f58609a4151cf"
+      network_throughput_bytes_per_sec: 983630003897
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9211f8f1f532d3d387bed1933873cab8"
+      network_throughput_bytes_per_sec: 678707539
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "33ebbaf4d4843276a44914043fb390a2"
+      network_throughput_bytes_per_sec: 753418440034
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0ac328277b4fe932363f73fd9079ecae"
+      network_throughput_bytes_per_sec: 974498970642
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b7104c1e829ea33bb66b519d636db1fb"
+      network_throughput_bytes_per_sec: 1081158605822
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4718592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "fe1bf82b165549da1e2fd013ee7d3b4e"
+      network_throughput_bytes_per_sec: 34644581497
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 75497472
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9d9f5bca0bb51e5e0d9a937a0934d777"
+      network_throughput_bytes_per_sec: 1106644073759
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "cf069120ed6b5049177dfa8f535d5b6a"
+      network_throughput_bytes_per_sec: 84760169245
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a944bb96b62b8e9f479c53fac1b29434"
+      network_throughput_bytes_per_sec: 104196816
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 32 }
+        source_target_pairs { source: 1 target: 33 }
+        source_target_pairs { source: 2 target: 34 }
+        source_target_pairs { source: 3 target: 35 }
+        source_target_pairs { source: 4 target: 36 }
+        source_target_pairs { source: 5 target: 37 }
+        source_target_pairs { source: 6 target: 38 }
+        source_target_pairs { source: 7 target: 39 }
+        source_target_pairs { source: 8 target: 40 }
+        source_target_pairs { source: 9 target: 41 }
+        source_target_pairs { source: 10 target: 42 }
+        source_target_pairs { source: 11 target: 43 }
+        source_target_pairs { source: 12 target: 44 }
+        source_target_pairs { source: 13 target: 45 }
+        source_target_pairs { source: 14 target: 46 }
+        source_target_pairs { source: 15 target: 47 }
+        source_target_pairs { source: 16 target: 48 }
+        source_target_pairs { source: 17 target: 49 }
+        source_target_pairs { source: 18 target: 50 }
+        source_target_pairs { source: 19 target: 51 }
+        source_target_pairs { source: 20 target: 52 }
+        source_target_pairs { source: 21 target: 53 }
+        source_target_pairs { source: 22 target: 54 }
+        source_target_pairs { source: 23 target: 55 }
+        source_target_pairs { source: 24 target: 56 }
+        source_target_pairs { source: 25 target: 57 }
+        source_target_pairs { source: 26 target: 58 }
+        source_target_pairs { source: 27 target: 59 }
+        source_target_pairs { source: 28 target: 60 }
+        source_target_pairs { source: 29 target: 61 }
+        source_target_pairs { source: 30 target: 62 }
+        source_target_pairs { source: 31 target: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "e6f126b6e8d61ee7816cd03251ad6058"
+      network_throughput_bytes_per_sec: 749669984905
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6bc65a094c4d9f93f5369ab6fc06b7cf"
+      network_throughput_bytes_per_sec: 167718490083
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 37748736
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1376938d5a4a0304bd1e3810c1d0d318"
+      network_throughput_bytes_per_sec: 284089948523
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a862a465752406e2cf9c7ebfa5eeb9a9"
+      network_throughput_bytes_per_sec: 387096774
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 8589934592
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b0002329e6226b53a5b0e1fb342b6958"
+      network_throughput_bytes_per_sec: 1022767628729
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4831838208
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6101050eebab12c8d699525eba07559b"
+      network_throughput_bytes_per_sec: 469005411991
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 8
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7f3e4c0fc4394122f4c70b76d4e9a164"
+      network_throughput_bytes_per_sec: 3723311
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 8388608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "09cacef6e95c98b628462c17b6326782"
+      network_throughput_bytes_per_sec: 84290675241
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 target: 2 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 4 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 6 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 8 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 10 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 12 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 14 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 16 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 18 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 20 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 22 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 24 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 26 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 28 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 30 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 32 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 34 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 36 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 38 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 40 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 42 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 44 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 46 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 48 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 50 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 52 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 54 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 56 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 58 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 60 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 62 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "215e2abd376353c5ddaebe6ad255504f"
+      network_throughput_bytes_per_sec: 12338510778
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1073741824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "513ead36822d49000966dd0365388a54"
+      network_throughput_bytes_per_sec: 716360855142
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 18874368
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "164346133a06bbfdc16cfb5b3c6517d7"
+      network_throughput_bytes_per_sec: 158113862547
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6f098675bd40c914c5c91ef2de01be56"
+      network_throughput_bytes_per_sec: 113728416485
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1024
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3f1d6ad7d12e3ff92151997fdfd0e0de"
+      network_throughput_bytes_per_sec: 454449074
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4831838208
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e4f7ddbd2c5d7fd6627277a5cbcb77a8"
+      network_throughput_bytes_per_sec: 1048976707736
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7165ce5e441106cbf102762355dab929"
+      network_throughput_bytes_per_sec: 21298667533
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0ad8ed44fe9e5a0b0f3c9896b68b7ac7"
+      network_throughput_bytes_per_sec: 41600253907
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "865f5ed95f3467ac18bc089e04cb9694"
+      network_throughput_bytes_per_sec: 663340819231
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0caa0898e5a1a1eb72969f83f30a04f2"
+      network_throughput_bytes_per_sec: 40355916
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1179648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "57747715df2ab61211387d6587fa7879"
+      network_throughput_bytes_per_sec: 21233494132
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4bd23fb067033c02382139c179f6db55"
+      network_throughput_bytes_per_sec: 459882844729
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "5e781037bbd9e6693be1abe36273ea7f"
+      network_throughput_bytes_per_sec: 6454540
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3922a82eabc6e6fb032d12e678d48c29"
+      network_throughput_bytes_per_sec: 354488167680
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a98bd29647c367825f298958182c4ca4"
+      network_throughput_bytes_per_sec: 1201173020
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2147483648
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4727a9ff57acd3df5626f18acd42e385"
+      network_throughput_bytes_per_sec: 894310868572
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "c53f8ac1261e1530974829131da17f1d"
+      network_throughput_bytes_per_sec: 33684210
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7ae3978bd9659c546a6a81dbe9fde3c7"
+      network_throughput_bytes_per_sec: 23842832888
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "27b2d0be5937d561a3a4ea0d57361fac"
+      network_throughput_bytes_per_sec: 37337136
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "102f48d7d7a892157e1e86fc2e501983"
+      network_throughput_bytes_per_sec: 272560938363
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c3e97fa05579fe168b5a25e564a573d6"
+      network_throughput_bytes_per_sec: 505299781642
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b51984717262696e7f8190e85b99a8b7"
+      network_throughput_bytes_per_sec: 52325581
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2eea6728e254d5a128416fc922761a82"
+      network_throughput_bytes_per_sec: 5154835411
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 589824
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c79b2ca1c506cad5778e452edeae6760"
+      network_throughput_bytes_per_sec: 3792007406
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 301989888
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "e27c240c5b40e40ce2372442689c1b19"
+      network_throughput_bytes_per_sec: 390450310302
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 2415919104
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "f268172d901a12911ba69d6857543645"
+      network_throughput_bytes_per_sec: 1047813788233
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "c1328becdecc5a5d6f19dd21c813e175"
+      network_throughput_bytes_per_sec: 43662581620
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "7e22ad68d0974535e42e3b66d346dae9"
+      network_throughput_bytes_per_sec: 253279227053
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 72
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8c6958a54dd555b75f8ee4c73dcb46be"
+      network_throughput_bytes_per_sec: 847497
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "36e4e9dc500587aac6f9d2bf11d0d82e"
+      network_throughput_bytes_per_sec: 597990305104
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 150994944
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "34524c5e07a660c9123821efb36b2cc8"
+      network_throughput_bytes_per_sec: 467063407527
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2048
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "32b1f8166b5665c19a5d2600b262a961"
+      network_throughput_bytes_per_sec: 360563380
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "efd42b664d4648d5963074f1a919e7c0"
+      network_throughput_bytes_per_sec: 430815258
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 128
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "315c4a86dadabb5b27fb5b1737458a98"
+      network_throughput_bytes_per_sec: 30758140
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "72528b7fe47cecc1784cabea1c98ffb1"
+      network_throughput_bytes_per_sec: 312724804
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 67108864
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "32fad1d3215cdacfd691ccea955a7154"
+      network_throughput_bytes_per_sec: 485297387984
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 536870912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "617c0b3b66cf2f8d2031519a87c29380"
+      network_throughput_bytes_per_sec: 981073541558
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4294967296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0fcc8e6e97bf1dcbebfa4ef5bf84e1ae"
+      network_throughput_bytes_per_sec: 912312708766
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 524288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3a16fa78d3b76fd63adfe679304741c1"
+      network_throughput_bytes_per_sec: 55258010118
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1048576
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "391c0c382d8a017fee373038a5e06c81"
+      network_throughput_bytes_per_sec: 217130196200
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "70146ac12eab16a26f89815b5a041177"
+      network_throughput_bytes_per_sec: 21636240
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 2415919104
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3e7979dc1ff371adeccc6d909f821f0b"
+      network_throughput_bytes_per_sec: 547171847609
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "1d6e728dd3bbbdf87ac030c1134e8199"
+      network_throughput_bytes_per_sec: 113923346
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "75c43384471f94fca4ad4920031c7816"
+      network_throughput_bytes_per_sec: 516901955649
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 9216
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b36003690ce2c88a5f8a1976fc2b2e31"
+      network_throughput_bytes_per_sec: 1256270447
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        source_target_pairs { source: 64 target: 65 }
+        source_target_pairs { source: 65 target: 64 }
+        source_target_pairs { source: 66 target: 67 }
+        source_target_pairs { source: 67 target: 66 }
+        source_target_pairs { source: 68 target: 69 }
+        source_target_pairs { source: 69 target: 68 }
+        source_target_pairs { source: 70 target: 71 }
+        source_target_pairs { source: 71 target: 70 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "c20f4d6924358e23021ab526b3ba75bf"
+      network_throughput_bytes_per_sec: 22580452509
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 64
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "bfb252cc94227b2c6d3dfc0d731ec5a2"
+      network_throughput_bytes_per_sec: 24601177767
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 512
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b92923fffb061f879f29b88c4f4eeca0"
+      network_throughput_bytes_per_sec: 192449047
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9ec536aec9020d0190ab4e1fdf867ec3"
+      network_throughput_bytes_per_sec: 106834029546
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 2359296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "9f87629c20a25d52da4f5b28ecd7617d"
+      network_throughput_bytes_per_sec: 43937203195
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 147456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 18
+          num_devices_per_group: 4
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c4e1e1837543409d39fdc7c922b00fe7"
+      network_throughput_bytes_per_sec: 20423268698
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 131072
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "6a12afac7d536eefc0c8d0717b6bb6e8"
+      network_throughput_bytes_per_sec: 47163281624
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "87c3f6f7a4a652daf434fa7b7c51a052"
+      network_throughput_bytes_per_sec: 1111621436055
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "aae75efb062ebbb6203ed57d1433d0b0"
+      network_throughput_bytes_per_sec: 288203940700
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "b8c58d57a75f5ca782b0a607ec08ea7c"
+      network_throughput_bytes_per_sec: 1105831533
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "81ce30857f8e403e7ba9bf71e41e3f86"
+      network_throughput_bytes_per_sec: 226198365
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 4608
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "ae4f6c62e03f26617fd5a86623a92a19"
+      network_throughput_bytes_per_sec: 369349150
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 16384
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "95e3953025dda0afcedb422e79019817"
+      network_throughput_bytes_per_sec: 9663226186
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 36 }
+        source_target_pairs { source: 1 target: 37 }
+        source_target_pairs { source: 2 target: 38 }
+        source_target_pairs { source: 3 target: 39 }
+        source_target_pairs { source: 4 target: 40 }
+        source_target_pairs { source: 5 target: 41 }
+        source_target_pairs { source: 6 target: 42 }
+        source_target_pairs { source: 7 target: 43 }
+        source_target_pairs { source: 8 target: 44 }
+        source_target_pairs { source: 9 target: 45 }
+        source_target_pairs { source: 10 target: 46 }
+        source_target_pairs { source: 11 target: 47 }
+        source_target_pairs { source: 12 target: 48 }
+        source_target_pairs { source: 13 target: 49 }
+        source_target_pairs { source: 14 target: 50 }
+        source_target_pairs { source: 15 target: 51 }
+        source_target_pairs { source: 16 target: 52 }
+        source_target_pairs { source: 17 target: 53 }
+        source_target_pairs { source: 18 target: 54 }
+        source_target_pairs { source: 19 target: 55 }
+        source_target_pairs { source: 20 target: 56 }
+        source_target_pairs { source: 21 target: 57 }
+        source_target_pairs { source: 22 target: 58 }
+        source_target_pairs { source: 23 target: 59 }
+        source_target_pairs { source: 24 target: 60 }
+        source_target_pairs { source: 25 target: 61 }
+        source_target_pairs { source: 26 target: 62 }
+        source_target_pairs { source: 27 target: 63 }
+        source_target_pairs { source: 28 target: 64 }
+        source_target_pairs { source: 29 target: 65 }
+        source_target_pairs { source: 30 target: 66 }
+        source_target_pairs { source: 31 target: 67 }
+        source_target_pairs { source: 32 target: 68 }
+        source_target_pairs { source: 33 target: 69 }
+        source_target_pairs { source: 34 target: 70 }
+        source_target_pairs { source: 35 target: 71 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "413e5c2a54a7ed594164ab2eba359f64"
+      network_throughput_bytes_per_sec: 796609516373
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 32
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "8d5513850b4dc8bc5f900e268071b39a"
+      network_throughput_bytes_per_sec: 4246848
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "4f86b0fc1e82ab3d790807e4232ed4c5"
+      network_throughput_bytes_per_sec: 195011344615
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 1152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0899b9eb66e1a135c6bfb1fbdd0adc3d"
+      network_throughput_bytes_per_sec: 626427406
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "2c36dc65751e760c5d17410abc2b1003"
+      network_throughput_bytes_per_sec: 42666666
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a9bd580013086eb148cb0045adbbda83"
+      network_throughput_bytes_per_sec: 243798186468
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 262144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "18c6e79b40884b4e305b268324e3dfde"
+      network_throughput_bytes_per_sec: 6096939250
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 134217728
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "02e48d6d2b02e52328859c242de7d9e8"
+      network_throughput_bytes_per_sec: 779864082182
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 65536
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3255c96fc31939f1074e5da294b6d707"
+      network_throughput_bytes_per_sec: 5047443006
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 144
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "0073cb1ddd9f1e9a8a156eaea75810b4"
+      network_throughput_bytes_per_sec: 15796401
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 32768
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9ec02154abd4032206757952438221bd"
+      network_throughput_bytes_per_sec: 16984105044
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 288
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d22d113edcb6f644cd88fe769f088756"
+      network_throughput_bytes_per_sec: 25477707
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-gather"
+        shape {
+          element_type: F32
+          dimensions: 268435456
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "a5f85cc53fa4972141a6403bdf6584ae"
+      network_throughput_bytes_per_sec: 862803599897
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 18432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 18
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "3084ed666ed72f64a2c2ad03602e46a6"
+      network_throughput_bytes_per_sec: 369764082
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 33554432
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 8
+          num_devices_per_group: 8
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c6c071537014f06ce64a597725863f80"
+      network_throughput_bytes_per_sec: 968577548133
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 2415919104
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 1
+          num_devices_per_group: 72
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "36761dc51bbd4c01b9bfe63ec96b265e"
+      network_throughput_bytes_per_sec: 790319797127
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-reduce"
+        shape {
+          element_type: F32
+          dimensions: 8192
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d86bd31be86868e17d229a22553c84c1"
+      network_throughput_bytes_per_sec: 1190697674
+    }
+    entries {
+      instruction {
+        name: "collective-permute"
+        opcode: "collective-permute"
+        shape {
+          element_type: F32
+          dimensions: 4294967296
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        channel_id: 1
+        id: 4294967299
+        operand_ids: 4294967298
+        source_target_pairs { target: 1 }
+        source_target_pairs { source: 1 }
+        source_target_pairs { source: 2 target: 3 }
+        source_target_pairs { source: 3 target: 2 }
+        source_target_pairs { source: 4 target: 5 }
+        source_target_pairs { source: 5 target: 4 }
+        source_target_pairs { source: 6 target: 7 }
+        source_target_pairs { source: 7 target: 6 }
+        source_target_pairs { source: 8 target: 9 }
+        source_target_pairs { source: 9 target: 8 }
+        source_target_pairs { source: 10 target: 11 }
+        source_target_pairs { source: 11 target: 10 }
+        source_target_pairs { source: 12 target: 13 }
+        source_target_pairs { source: 13 target: 12 }
+        source_target_pairs { source: 14 target: 15 }
+        source_target_pairs { source: 15 target: 14 }
+        source_target_pairs { source: 16 target: 17 }
+        source_target_pairs { source: 17 target: 16 }
+        source_target_pairs { source: 18 target: 19 }
+        source_target_pairs { source: 19 target: 18 }
+        source_target_pairs { source: 20 target: 21 }
+        source_target_pairs { source: 21 target: 20 }
+        source_target_pairs { source: 22 target: 23 }
+        source_target_pairs { source: 23 target: 22 }
+        source_target_pairs { source: 24 target: 25 }
+        source_target_pairs { source: 25 target: 24 }
+        source_target_pairs { source: 26 target: 27 }
+        source_target_pairs { source: 27 target: 26 }
+        source_target_pairs { source: 28 target: 29 }
+        source_target_pairs { source: 29 target: 28 }
+        source_target_pairs { source: 30 target: 31 }
+        source_target_pairs { source: 31 target: 30 }
+        source_target_pairs { source: 32 target: 33 }
+        source_target_pairs { source: 33 target: 32 }
+        source_target_pairs { source: 34 target: 35 }
+        source_target_pairs { source: 35 target: 34 }
+        source_target_pairs { source: 36 target: 37 }
+        source_target_pairs { source: 37 target: 36 }
+        source_target_pairs { source: 38 target: 39 }
+        source_target_pairs { source: 39 target: 38 }
+        source_target_pairs { source: 40 target: 41 }
+        source_target_pairs { source: 41 target: 40 }
+        source_target_pairs { source: 42 target: 43 }
+        source_target_pairs { source: 43 target: 42 }
+        source_target_pairs { source: 44 target: 45 }
+        source_target_pairs { source: 45 target: 44 }
+        source_target_pairs { source: 46 target: 47 }
+        source_target_pairs { source: 47 target: 46 }
+        source_target_pairs { source: 48 target: 49 }
+        source_target_pairs { source: 49 target: 48 }
+        source_target_pairs { source: 50 target: 51 }
+        source_target_pairs { source: 51 target: 50 }
+        source_target_pairs { source: 52 target: 53 }
+        source_target_pairs { source: 53 target: 52 }
+        source_target_pairs { source: 54 target: 55 }
+        source_target_pairs { source: 55 target: 54 }
+        source_target_pairs { source: 56 target: 57 }
+        source_target_pairs { source: 57 target: 56 }
+        source_target_pairs { source: 58 target: 59 }
+        source_target_pairs { source: 59 target: 58 }
+        source_target_pairs { source: 60 target: 61 }
+        source_target_pairs { source: 61 target: 60 }
+        source_target_pairs { source: 62 target: 63 }
+        source_target_pairs { source: 63 target: 62 }
+        frontend_attributes {}
+        statistics_viz {}
+      }
+      fingerprint: "94ef7649bbef959076271d177e3bdb39"
+      network_throughput_bytes_per_sec: 498347093213
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 2097152
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 4
+          num_devices_per_group: 16
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "d6e99ab11f8c0c4bd10aa136c0b8f7ca"
+      network_throughput_bytes_per_sec: 248220387631
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 294912
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "acfd0dc289a5ea30e9a32af73a1f3fdc"
+      network_throughput_bytes_per_sec: 7482037751
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 256
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 2
+          num_devices_per_group: 36
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "9ac256563978ad06b4d5e4f9d6c7fd50"
+      network_throughput_bytes_per_sec: 335811106
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "reduce-scatter"
+        shape {
+          element_type: F32
+          dimensions: 4194304
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 4294967297
+        operand_ids: 4294967296
+        called_computation_ids: 0
+        frontend_attributes {}
+        use_global_device_ids: true
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 16
+          num_devices_per_group: 4
+          iota_reshape_dims: 64
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "c50508e54147a8633f2835b14b8fed47"
+      network_throughput_bytes_per_sec: 575745229924
+    }
+    entries {
+      instruction {
+        name: "_"
+        opcode: "all-to-all"
+        shape {
+          element_type: F32
+          dimensions: 1207959552
+          layout { minor_to_major: 0 tail_padding_alignment_in_elements: 1 }
+          is_dynamic_dimension: false
+        }
+        metadata {}
+        dimensions: 0
+        channel_id: 1
+        id: 1
+        operand_ids: 0
+        frontend_attributes {}
+        statistics_viz {}
+        iota_collective_device_list {
+          num_replica_groups: 9
+          num_devices_per_group: 8
+          iota_reshape_dims: 72
+          iota_transpose_perm: 0
+        }
+      }
+      fingerprint: "96322e653a0ddf56266cc0e50d59ed10"
+      network_throughput_bytes_per_sec: 964004673344
+    }
+  }
+}
diff --git a/third_party/xla/xla/service/gpu/model/default_matmul_perf_table.txtpb b/third_party/xla/xla/service/gpu/model/default_matmul_perf_table.txtpb
index d32187a924f..a30b0b33b0b 100644
--- a/third_party/xla/xla/service/gpu/model/default_matmul_perf_table.txtpb
+++ b/third_party/xla/xla/service/gpu/model/default_matmul_perf_table.txtpb
@@ -4145,6 +4145,3011 @@ entries {
     }
   }
 }
+entries {
+  key: "sm_107"
+  value {
+    entries {
+      b: 1
+      m: 256
+      n: 256
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 4809981651376 }
+      flops { key: "f32xf32->f32" value: 4702134529147 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 256
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 9300008869179 }
+      flops { key: "f32xf32->f32" value: 9300008869179 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 256
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 17924376068376 }
+      flops { key: "f32xf32->f32" value: 16578276679841 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 256
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 34735436853002 }
+      flops { key: "f32xf32->f32" value: 31714964083175 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 256
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 55007265573770 }
+      flops { key: "f32xf32->f32" value: 63670648956356 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 512
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 9383230425055 }
+      flops { key: "f32xf32->f32" value: 8035065134099 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 512
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 18893261261261 }
+      flops { key: "f32xf32->f32" value: 15477136531365 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 512
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 32961131630648 }
+      flops { key: "f32xf32->f32" value: 25458597875569 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 512
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 61794534069981 }
+      flops { key: "f32xf32->f32" value: 61567765137614 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 512
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 98400093841642 }
+      flops { key: "f32xf32->f32" value: 108590394822006 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 1024
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 33961975708502 }
+      flops { key: "f32xf32->f32" value: 28244471380471 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 1024
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 34663669421487 }
+      flops { key: "f32xf32->f32" value: 29228599303135 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 1024
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 63310249056603 }
+      flops { key: "f32xf32->f32" value: 66841498007968 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 1024
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 93858551048951 }
+      flops { key: "f32xf32->f32" value: 117528658493870 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 1024
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 192841563218390 }
+      flops { key: "f32xf32->f32" value: 206171625192012 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 2048
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 56679783783783 }
+      flops { key: "f32xf32->f32" value: 47393265536723 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 2048
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 107892064308681 }
+      flops { key: "f32xf32->f32" value: 119410790035587 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 2048
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 108065803542673 }
+      flops { key: "f32xf32->f32" value: 116306523396880 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 2048
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 182857940054495 }
+      flops { key: "f32xf32->f32" value: 210043392801251 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 2048
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 284359593220339 }
+      flops { key: "f32xf32->f32" value: 355543650331125 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 4096
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 114912438356164 }
+      flops { key: "f32xf32->f32" value: 115704937931034 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 4096
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 200924742514970 }
+      flops { key: "f32xf32->f32" value: 207767380804953 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 4096
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 311410041763341 }
+      flops { key: "f32xf32->f32" value: 324982392251816 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 4096
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 338079919395466 }
+      flops { key: "f32xf32->f32" value: 354604301188903 }
+    }
+    entries {
+      b: 1
+      m: 256
+      n: 4096
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 528936859113300 }
+      flops { key: "f32xf32->f32" value: 501748515887850 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 256
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 9198035087719 }
+      flops { key: "f32xf32->f32" value: 7810621973929 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 256
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 19784452830188 }
+      flops { key: "f32xf32->f32" value: 17810208067940 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 256
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 32961131630648 }
+      flops { key: "f32xf32->f32" value: 24672376470588 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 256
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 59283448763250 }
+      flops { key: "f32xf32->f32" value: 49344752941176 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 256
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 106184911392405 }
+      flops { key: "f32xf32->f32" value: 106353191759112 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 512
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 33554432000000 }
+      flops { key: "f32xf32->f32" value: 28728109589041 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 512
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 34030864097363 }
+      flops { key: "f32xf32->f32" value: 30783882568807 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 512
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 66182311637080 }
+      flops { key: "f32xf32->f32" value: 62836014981273 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 512
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 99420539259259 }
+      flops { key: "f32xf32->f32" value: 111291648424543 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 512
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 187717102097902 }
+      flops { key: "f32xf32->f32" value: 205855411042944 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 1024
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 64035175572519 }
+      flops { key: "f32xf32->f32" value: 46474282548476 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 1024
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 112977885521885 }
+      flops { key: "f32xf32->f32" value: 114520245733788 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 1024
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 100764060060060 }
+      flops { key: "f32xf32->f32" value: 117118436300174 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 1024
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 200624406576980 }
+      flops { key: "f32xf32->f32" value: 215784128617363 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 1024
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 331401797530864 }
+      flops { key: "f32xf32->f32" value: 342392163265306 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 2048
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 115307326460481 }
+      flops { key: "f32xf32->f32" value: 121574028985507 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 2048
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 172516359897172 }
+      flops { key: "f32xf32->f32" value: 210372614420062 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 2048
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 307134389016018 }
+      flops { key: "f32xf32->f32" value: 328965019607843 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 2048
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 286178524520255 }
+      flops { key: "f32xf32->f32" value: 345921979381443 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 2048
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 501748515887850 }
+      flops { key: "f32xf32->f32" value: 530504853754940 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 4096
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 196800187683284 }
+      flops { key: "f32xf32->f32" value: 201528120120120 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 4096
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 342392163265306 }
+      flops { key: "f32xf32->f32" value: 345032719794344 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 4096
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 529458493096646 }
+      flops { key: "f32xf32->f32" value: 531555358415841 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 4096
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 727467360433604 }
+      flops { key: "f32xf32->f32" value: 726482966170500 }
+    }
+    entries {
+      b: 1
+      m: 512
+      n: 4096
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 941878792982456 }
+      flops { key: "f32xf32->f32" value: 742560044260027 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 256
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 33026015748031 }
+      flops { key: "f32xf32->f32" value: 29127111111111 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 256
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 34030864097363 }
+      flops { key: "f32xf32->f32" value: 30840470588235 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 256
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 60241350089766 }
+      flops { key: "f32xf32->f32" value: 52022375193798 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 256
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 102770082695252 }
+      flops { key: "f32xf32->f32" value: 114912438356164 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 256
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 195367871906841 }
+      flops { key: "f32xf32->f32" value: 198546934911242 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 512
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 63072240601503 }
+      flops { key: "f32xf32->f32" value: 47527524079320 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 512
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 113335636900992 }
+      flops { key: "f32xf32->f32" value: 117734849122807 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 512
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 112977885521885 }
+      flops { key: "f32xf32->f32" value: 120266781362007 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 512
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 198841078518518 }
+      flops { key: "f32xf32->f32" value: 212369822784810 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 512
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 428126724082934 }
+      flops { key: "f32xf32->f32" value: 340222377693282 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 1024
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 114912438356164 }
+      flops { key: "f32xf32->f32" value: 114912438356164 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 1024
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 213044012698412 }
+      flops { key: "f32xf32->f32" value: 215784128617363 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 1024
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 281378884696016 }
+      flops { key: "f32xf32->f32" value: 345921979381443 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 1024
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 412977624615384 }
+      flops { key: "f32xf32->f32" value: 359833050938337 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 1024
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 599186285714285 }
+      flops { key: "f32xf32->f32" value: 539568755778894 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 2048
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 198546934911242 }
+      flops { key: "f32xf32->f32" value: 209061881619937 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 2048
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 302292180180180 }
+      flops { key: "f32xf32->f32" value: 351355308900523 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 2048
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 483667488288288 }
+      flops { key: "f32xf32->f32" value: 536870912000000 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 2048
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 730436614965986 }
+      flops { key: "f32xf32->f32" value: 732429620736698 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 2048
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 941878792982456 }
+      flops { key: "f32xf32->f32" value: 799509921072226 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 4096
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 351355308900523 }
+      flops { key: "f32xf32->f32" value: 323416212048192 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 4096
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 484540534296028 }
+      flops { key: "f32xf32->f32" value: 487178686025408 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 4096
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 719666101876675 }
+      flops { key: "f32xf32->f32" value: 748773935843793 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 4096
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1025541379178605 }
+      flops { key: "f32xf32->f32" value: 1012963984905660 }
+    }
+    entries {
+      b: 1
+      m: 1024
+      n: 4096
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 1807646168350168 }
+      flops { key: "f32xf32->f32" value: 1248537004651162 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 256
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 62836014981273 }
+      flops { key: "f32xf32->f32" value: 46733192200557 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 256
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 106861248407643 }
+      flops { key: "f32xf32->f32" value: 113359567567567 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 256
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 105517081761006 }
+      flops { key: "f32xf32->f32" value: 120699395683453 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 256
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 196782154128104 }
+      flops { key: "f32xf32->f32" value: 205855411042944 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 256
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 328562369645042 }
+      flops { key: "f32xf32->f32" value: 343707370038412 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 512
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 116105301038062 }
+      flops { key: "f32xf32->f32" value: 119410790035587 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 512
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 199136094955489 }
+      flops { key: "f32xf32->f32" value: 186932768802228 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 512
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 302292180180180 }
+      flops { key: "f32xf32->f32" value: 342392163265306 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 512
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 331811441285537 }
+      flops { key: "f32xf32->f32" value: 357437358189081 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 512
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 608697179138322 }
+      flops { key: "f32xf32->f32" value: 551202168377823 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 1024
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 206488812307692 }
+      flops { key: "f32xf32->f32" value: 211699886435331 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 1024
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 300936609865470 }
+      flops { key: "f32xf32->f32" value: 335544320000000 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 1024
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 541200516129032 }
+      flops { key: "f32xf32->f32" value: 533668898608349 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 1024
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1005376239700374 }
+      flops { key: "f32xf32->f32" value: 757222724964739 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 1024
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 958698057142857 }
+      flops { key: "f32xf32->f32" value: 795364314074074 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 2048
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 352277501312335 }
+      flops { key: "f32xf32->f32" value: 317299593380614 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 2048
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 546642140257095 }
+      flops { key: "f32xf32->f32" value: 503631249530956 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 2048
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 833650484472049 }
+      flops { key: "f32xf32->f32" value: 757222724964739 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 2048
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 981482471663619 }
+      flops { key: "f32xf32->f32" value: 1008208285446009 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 2048
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 1810694475548060 }
+      flops { key: "f32xf32->f32" value: 1304668072904009 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 4096
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 442962798679868 }
+      flops { key: "f32xf32->f32" value: 440780715927750 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 4096
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 665267548946716 }
+      flops { key: "f32xf32->f32" value: 691843958762886 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 4096
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 1043480878522837 }
+      flops { key: "f32xf32->f32" value: 954437176888888 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 4096
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1872261244986922 }
+      flops { key: "f32xf32->f32" value: 1209168720720720 }
+    }
+    entries {
+      b: 1
+      m: 2048
+      n: 4096
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 2408843127313516 }
+      flops { key: "f32xf32->f32" value: 1488691248803102 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 256
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 113743837288135 }
+      flops { key: "f32xf32->f32" value: 121135133574007 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 256
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 204600195121951 }
+      flops { key: "f32xf32->f32" value: 206488812307692 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 256
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 300263373601789 }
+      flops { key: "f32xf32->f32" value: 340654131979695 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 256
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 386203335671252 }
+      flops { key: "f32xf32->f32" value: 345032719794344 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 256
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 599186285714285 }
+      flops { key: "f32xf32->f32" value: 521233895145631 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 512
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 201528120120120 }
+      flops { key: "f32xf32->f32" value: 204600195121951 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 512
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 362750616216216 }
+      flops { key: "f32xf32->f32" value: 347714321243523 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 512
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 520223751937984 }
+      flops { key: "f32xf32->f32" value: 525314003913894 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 512
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 990536738007380 }
+      flops { key: "f32xf32->f32" value: 750868408391608 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 512
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 939406670166229 }
+      flops { key: "f32xf32->f32" value: 790097000735835 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 1024
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 340654131979695 }
+      flops { key: "f32xf32->f32" value: 315806418823529 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 1024
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 517216678227360 }
+      flops { key: "f32xf32->f32" value: 515231201535508 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 1024
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 730436614965986 }
+      flops { key: "f32xf32->f32" value: 751920044817927 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 1024
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1205097445566779 }
+      flops { key: "f32xf32->f32" value: 1004435756782039 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 1024
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 1770390476504534 }
+      flops { key: "f32xf32->f32" value: 1254371289719626 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 2048
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 467657588850174 }
+      flops { key: "f32xf32->f32" value: 458080982935153 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 2048
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 726482966170500 }
+      flops { key: "f32xf32->f32" value: 665267548946716 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 2048
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 1285822107926053 }
+      flops { key: "f32xf32->f32" value: 970833475587703 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 2048
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1856079211754537 }
+      flops { key: "f32xf32->f32" value: 1252177054227405 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 2048
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 2400764279485746 }
+      flops { key: "f32xf32->f32" value: 1481534079337702 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 4096
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 653127630170316 }
+      flops { key: "f32xf32->f32" value: 578524689655172 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 4096
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 1229944815578465 }
+      flops { key: "f32xf32->f32" value: 891765854347261 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 4096
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 1887068231985940 }
+      flops { key: "f32xf32->f32" value: 1205097445566779 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 4096
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 2430654949632145 }
+      flops { key: "f32xf32->f32" value: 1437405386880856 }
+    }
+    entries {
+      b: 1
+      m: 4096
+      n: 4096
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 2906915259560067 }
+      flops { key: "f32xf32->f32" value: 1634912906346279 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 256
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 9489375565610 }
+      flops { key: "f32xf32->f32" value: 8507716024340 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 256
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 20213513253012 }
+      flops { key: "f32xf32->f32" value: 17772474576271 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 256
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 28926234482758 }
+      flops { key: "f32xf32->f32" value: 30897267034990 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 256
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 55001630160844 }
+      flops { key: "f32xf32->f32" value: 42853680715197 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 256
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 101989155015197 }
+      flops { key: "f32xf32->f32" value: 75149903695408 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 512
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 32513984496124 }
+      flops { key: "f32xf32->f32" value: 26715312101910 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 512
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 34521020576131 }
+      flops { key: "f32xf32->f32" value: 31775030303030 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 512
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 61008058181818 }
+      flops { key: "f32xf32->f32" value: 49417425625920 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 512
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 107546256410256 }
+      flops { key: "f32xf32->f32" value: 74565404444444 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 512
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 190379756028368 }
+      flops { key: "f32xf32->f32" value: 111384006639004 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 1024
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 58457198606271 }
+      flops { key: "f32xf32->f32" value: 66576253968253 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 1024
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 99864380952380 }
+      flops { key: "f32xf32->f32" value: 118566897526501 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 1024
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 104368373250388 }
+      flops { key: "f32xf32->f32" value: 74400070953436 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 1024
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 195367871906841 }
+      flops { key: "f32xf32->f32" value: 114227853617021 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 1024
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 333046471464019 }
+      flops { key: "f32xf32->f32" value: 187717102097902 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 2048
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 108240103225806 }
+      flops { key: "f32xf32->f32" value: 120266781362007 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 2048
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 191193344729344 }
+      flops { key: "f32xf32->f32" value: 196800187683284 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 2048
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 283720920597172 }
+      flops { key: "f32xf32->f32" value: 297600283813747 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 2048
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 330585536945812 }
+      flops { key: "f32xf32->f32" value: 194800766328011 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 2048
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 489845722627737 }
+      flops { key: "f32xf32->f32" value: 237232030489657 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 4096
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 182361043478260 }
+      flops { key: "f32xf32->f32" value: 197379011764705 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 4096
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 336385283208020 }
+      flops { key: "f32xf32->f32" value: 303660018099547 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 4096
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 509365191650853 }
+      flops { key: "f32xf32->f32" value: 435771844155844 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 4096
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 697234950649350 }
+      flops { key: "f32xf32->f32" value: 621378370370370 }
+    }
+    entries {
+      b: 2
+      m: 256
+      n: 4096
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 689622237636480 }
+      flops { key: "f32xf32->f32" value: 325468773022639 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 256
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 32513984496124 }
+      flops { key: "f32xf32->f32" value: 27685174917491 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 256
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 32896501960784 }
+      flops { key: "f32xf32->f32" value: 31651392052823 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 256
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 59811821746880 }
+      flops { key: "f32xf32->f32" value: 50533783132530 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 256
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 108414966074313 }
+      flops { key: "f32xf32->f32" value: 77762298957126 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 256
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 191193344729344 }
+      flops { key: "f32xf32->f32" value: 105517081761006 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 512
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 59918628571428 }
+      flops { key: "f32xf32->f32" value: 47259763380281 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 512
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 104857600000000 }
+      flops { key: "f32xf32->f32" value: 118149408450704 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 512
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 107546256410256 }
+      flops { key: "f32xf32->f32" value: 77852510440835 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 512
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 197961250737463 }
+      flops { key: "f32xf32->f32" value: 107117101356743 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 512
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 338933656565656 }
+      flops { key: "f32xf32->f32" value: 196512046852123 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 1024
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 112598765100671 }
+      flops { key: "f32xf32->f32" value: 119837257142857 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 1024
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 194518446376811 }
+      flops { key: "f32xf32->f32" value: 204600195121951 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 1024
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 316551245283018 }
+      flops { key: "f32xf32->f32" value: 297600283813747 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 1024
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 283458770855332 }
+      flops { key: "f32xf32->f32" value: 219130984489795 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 1024
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 492542121100917 }
+      flops { key: "f32xf32->f32" value: 267899656686626 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 2048
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 182361043478260 }
+      flops { key: "f32xf32->f32" value: 197379011764705 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 2048
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 329773287469287 }
+      flops { key: "f32xf32->f32" value: 292413350762527 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 2048
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 409825123664122 }
+      flops { key: "f32xf32->f32" value: 429496729600000 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 2048
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 927238189982728 }
+      flops { key: "f32xf32->f32" value: 603904287964004 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 2048
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 763142732054015 }
+      flops { key: "f32xf32->f32" value: 410765808722264 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 4096
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 317252717979022 }
+      flops { key: "f32xf32->f32" value: 270600258064516 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 4096
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 492542121100917 }
+      flops { key: "f32xf32->f32" value: 417473493001555 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 4096
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 780335627906976 }
+      flops { key: "f32xf32->f32" value: 598518296544035 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 4096
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 988712545119705 }
+      flops { key: "f32xf32->f32" value: 798321058736059 }
+    }
+    entries {
+      b: 2
+      m: 512
+      n: 4096
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 1654376155231354 }
+      flops { key: "f32xf32->f32" value: 1008208285446009 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 256
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 61908546125461 }
+      flops { key: "f32xf32->f32" value: 66313106719367 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 256
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 92436451790633 }
+      flops { key: "f32xf32->f32" value: 108942961038961 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 256
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 99273467455621 }
+      flops { key: "f32xf32->f32" value: 77762298957126 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 256
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 170543491740787 }
+      flops { key: "f32xf32->f32" value: 107892064308681 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 256
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 314695728018757 }
+      flops { key: "f32xf32->f32" value: 195083906976744 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 512
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 116914397212543 }
+      flops { key: "f32xf32->f32" value: 121135133574007 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 512
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 186413511111111 }
+      flops { key: "f32xf32->f32" value: 202745812688821 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 512
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 302974555304740 }
+      flops { key: "f32xf32->f32" value: 320328706443914 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 512
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 333460193788819 }
+      flops { key: "f32xf32->f32" value: 229628277159965 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 512
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 496642841813136 }
+      flops { key: "f32xf32->f32" value: 286026058604155 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 1024
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 200324967164179 }
+      flops { key: "f32xf32->f32" value: 197379011764705 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 1024
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 289262344827586 }
+      flops { key: "f32xf32->f32" value: 309971658198614 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 1024
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 442962798679868 }
+      flops { key: "f32xf32->f32" value: 477643160142348 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 1024
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 970833475587703 }
+      flops { key: "f32xf32->f32" value: 664444198019802 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 1024
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 777510372194062 }
+      flops { key: "f32xf32->f32" value: 443511699297810 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 2048
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 321095043062200 }
+      flops { key: "f32xf32->f32" value: 304348589569161 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 2048
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 507439425330812 }
+      flops { key: "f32xf32->f32" value: 451152026890756 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 2048
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 686535693094629 }
+      flops { key: "f32xf32->f32" value: 655520039072039 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 2048
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1003497031775700 }
+      flops { key: "f32xf32->f32" value: 899281259631490 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 2048
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 1771759829216727 }
+      flops { key: "f32xf32->f32" value: 1153940702847931 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 4096
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 439337898527004 }
+      flops { key: "f32xf32->f32" value: 387912508670520 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 4096
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 965595165467625 }
+      flops { key: "f32xf32->f32" value: 610774643913538 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 4096
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 1263225675294117 }
+      flops { key: "f32xf32->f32" value: 853530861685214 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 4096
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1713873621707901 }
+      flops { key: "f32xf32->f32" value: 1110958948784273 }
+    }
+    entries {
+      b: 2
+      m: 1024
+      n: 4096
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 2176871412062848 }
+      flops { key: "f32xf32->f32" value: 1326400369356675 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 256
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 114912438356164 }
+      flops { key: "f32xf32->f32" value: 121574028985507 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 256
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 196800187683284 }
+      flops { key: "f32xf32->f32" value: 208412621118012 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 256
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 314327231850117 }
+      flops { key: "f32xf32->f32" value: 300263373601789 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 256
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 273077778229908 }
+      flops { key: "f32xf32->f32" value: 205225883792048 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 256
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 501280029878618 }
+      flops { key: "f32xf32->f32" value: 254079939422621 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 512
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 188508044943820 }
+      flops { key: "f32xf32->f32" value: 201528120120120 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 512
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 340654131979695 }
+      flops { key: "f32xf32->f32" value: 315065089201877 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 512
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 516222030769230 }
+      flops { key: "f32xf32->f32" value: 463619094991364 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 512
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 960413080500894 }
+      flops { key: "f32xf32->f32" value: 648394821256038 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 512
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 745136588480222 }
+      flops { key: "f32xf32->f32" value: 416815129291311 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 1024
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 309257437788018 }
+      flops { key: "f32xf32->f32" value: 305735143507972 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 1024
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 472597633802816 }
+      flops { key: "f32xf32->f32" value: 454975349152542 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 1024
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 756156214084507 }
+      flops { key: "f32xf32->f32" value: 636857546856465 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 1024
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1355734626262626 }
+      flops { key: "f32xf32->f32" value: 909950698305084 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 1024
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 1751617983686786 }
+      flops { key: "f32xf32->f32" value: 1100673939456065 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 2048
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 459649753424657 }
+      flops { key: "f32xf32->f32" value: 422068327044025 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 2048
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 1011056331450094 }
+      flops { key: "f32xf32->f32" value: 584190328618063 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 2048
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 1239886632794457 }
+      flops { key: "f32xf32->f32" value: 844800805664830 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 2048
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1738853156275303 }
+      flops { key: "f32xf32->f32" value: 1105241198147195 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 2048
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 2200220175327378 }
+      flops { key: "f32xf32->f32" value: 1423588762346702 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 4096
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 849479291139240 }
+      flops { key: "f32xf32->f32" value: 514737211888782 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 4096
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 1279787632896305 }
+      flops { key: "f32xf32->f32" value: 843432136285531 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 4096
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 1844917223367697 }
+      flops { key: "f32xf32->f32" value: 1121401382767624 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 4096
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 2306641941997852 }
+      flops { key: "f32xf32->f32" value: 1372632565036753 }
+    }
+    entries {
+      b: 2
+      m: 2048
+      n: 4096
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 2656903351543621 }
+      flops { key: "f32xf32->f32" value: 1560671255813953 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 256
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 192289008595988 }
+      flops { key: "f32xf32->f32" value: 196800187683284 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 256
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 328160704156479 }
+      flops { key: "f32xf32->f32" value: 329773287469287 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 256
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 494356272559852 }
+      flops { key: "f32xf32->f32" value: 474267590106007 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 256
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 715827882666666 }
+      flops { key: "f32xf32->f32" value: 638372071343638 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 256
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 671928550688360 }
+      flops { key: "f32xf32->f32" value: 340431371921133 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 512
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 309971658198614 }
+      flops { key: "f32xf32->f32" value: 302292180180180 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 512
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 533668898608349 }
+      flops { key: "f32xf32->f32" value: 463619094991364 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 512
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 832358003100775 }
+      flops { key: "f32xf32->f32" value: 655520039072039 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 512
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1191722335183129 }
+      flops { key: "f32xf32->f32" value: 909180206604572 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 512
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 1804608107563025 }
+      flops { key: "f32xf32->f32" value: 1093980462557310 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 1024
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 469292755244755 }
+      flops { key: "f32xf32->f32" value: 410451767584097 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 1024
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 967334976576576 }
+      flops { key: "f32xf32->f32" value: 625723673659673 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 1024
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 1221549287827076 }
+      flops { key: "f32xf32->f32" value: 824686500768049 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 1024
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1705705836378077 }
+      flops { key: "f32xf32->f32" value: 1110958948784273 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 1024
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 2208209406683804 }
+      flops { key: "f32xf32->f32" value: 1416545941952506 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 2048
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 870131137763371 }
+      flops { key: "f32xf32->f32" value: 508882381042654 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 2048
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 1285918352095808 }
+      flops { key: "f32xf32->f32" value: 837552124804992 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 2048
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 1813753081081081 }
+      flops { key: "f32xf32->f32" value: 1111498022450101 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 2048
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 2280841605629128 }
+      flops { key: "f32xf32->f32" value: 1377005845827071 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 2048
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 2634090757843494 }
+      flops { key: "f32xf32->f32" value: 1590728628148148 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 4096
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 1138644564156946 }
+      flops { key: "f32xf32->f32" value: 682173966963151 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 4096
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 1753047875918367 }
+      flops { key: "f32xf32->f32" value: 1071598626746507 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 4096
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 2448670066134549 }
+      flops { key: "f32xf32->f32" value: 1439814715387194 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 4096
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 2946804319725557 }
+      flops { key: "f32xf32->f32" value: 1638362500858287 }
+    }
+    entries {
+      b: 2
+      m: 4096
+      n: 4096
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 3334601937888199 }
+      flops { key: "f32xf32->f32" value: 1775502089202090 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 256
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 34663669421487 }
+      flops { key: "f32xf32->f32" value: 27060025806451 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 256
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 31011489833641 }
+      flops { key: "f32xf32->f32" value: 31418007490636 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 256
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 61794534069981 }
+      flops { key: "f32xf32->f32" value: 49932190476190 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 256
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 105351434850863 }
+      flops { key: "f32xf32->f32" value: 75658245772266 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 256
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 195083906976744 }
+      flops { key: "f32xf32->f32" value: 112598765100671 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 512
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 58661594405594 }
+      flops { key: "f32xf32->f32" value: 46474282548476 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 512
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 97826332361516 }
+      flops { key: "f32xf32->f32" value: 110376421052631 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 512
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 112034831385642 }
+      flops { key: "f32xf32->f32" value: 76783597254004 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 512
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 193119033093525 }
+      flops { key: "f32xf32->f32" value: 113743837288135 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 512
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 334290729763387 }
+      flops { key: "f32xf32->f32" value: 149048004441976 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 1024
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 91929950684931 }
+      flops { key: "f32xf32->f32" value: 118987347517730 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 1024
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 176602273684210 }
+      flops { key: "f32xf32->f32" value: 197379011764705 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 1024
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 296941876106194 }
+      flops { key: "f32xf32->f32" value: 292413350762527 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 1024
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 328562369645042 }
+      flops { key: "f32xf32->f32" value: 188111742116327 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 1024
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 493901482980680 }
+      flops { key: "f32xf32->f32" value: 232009901469317 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 2048
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 195083906976744 }
+      flops { key: "f32xf32->f32" value: 192289008595988 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 2048
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 319566019047619 }
+      flops { key: "f32xf32->f32" value: 292413350762527 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 2048
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 404270265060240 }
+      flops { key: "f32xf32->f32" value: 416179001550387 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 2048
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 697234950649350 }
+      flops { key: "f32xf32->f32" value: 605223320791939 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 2048
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 674884867379006 }
+      flops { key: "f32xf32->f32" value: 320999050523168 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 4096
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 312134251162790 }
+      flops { key: "f32xf32->f32" value: 268973402805611 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 4096
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 464421204152249 }
+      flops { key: "f32xf32->f32" value: 423399772870662 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 4096
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 677012499369483 }
+      flops { key: "f32xf32->f32" value: 539534865397902 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 4096
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1103537331963001 }
+      flops { key: "f32xf32->f32" value: 777510372194062 }
+    }
+    entries {
+      b: 4
+      m: 256
+      n: 4096
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 1494421466945024 }
+      flops { key: "f32xf32->f32" value: 990051530557556 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 256
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 58254222222222 }
+      flops { key: "f32xf32->f32" value: 46603377777777 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 256
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 95325090909090 }
+      flops { key: "f32xf32->f32" value: 112977885521885 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 256
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 103883690402476 }
+      flops { key: "f32xf32->f32" value: 76346830489192 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 256
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 206488812307692 }
+      flops { key: "f32xf32->f32" value: 102139531414982 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 256
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 330585536945812 }
+      flops { key: "f32xf32->f32" value: 172294901155327 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 512
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 115704937931034 }
+      flops { key: "f32xf32->f32" value: 115704937931034 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 512
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 188474955941723 }
+      flops { key: "f32xf32->f32" value: 198546934911242 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 512
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 283159763713080 }
+      flops { key: "f32xf32->f32" value: 288640275268817 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 512
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 332633774473358 }
+      flops { key: "f32xf32->f32" value: 203978310030395 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 512
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 492994409550045 }
+      flops { key: "f32xf32->f32" value: 249012482374768 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 1024
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 179916525469168 }
+      flops { key: "f32xf32->f32" value: 193397302593659 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 1024
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 328160704156479 }
+      flops { key: "f32xf32->f32" value: 289262344827586 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 1024
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 494299378064219 }
+      flops { key: "f32xf32->f32" value: 422733001574803 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 1024
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 546155556459816 }
+      flops { key: "f32xf32->f32" value: 582921728555917 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 1024
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 741534408839779 }
+      flops { key: "f32xf32->f32" value: 360437000335683 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 2048
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 318806954869358 }
+      flops { key: "f32xf32->f32" value: 275601084188911 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 2048
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 445166593698175 }
+      flops { key: "f32xf32->f32" value: 417473493001555 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 2048
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 804903916041979 }
+      flops { key: "f32xf32->f32" value: 595200567627494 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 2048
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1135033640591966 }
+      flops { key: "f32xf32->f32" value: 825320387394312 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 2048
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 1530636955096222 }
+      flops { key: "f32xf32->f32" value: 1045512973709834 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 4096
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 430875531300160 }
+      flops { key: "f32xf32->f32" value: 371794260387811 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 4096
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 922458611683848 }
+      flops { key: "f32xf32->f32" value: 555731033965193 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 4096
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 1177348491228070 }
+      flops { key: "f32xf32->f32" value: 743073926643598 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 4096
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1593088759643917 }
+      flops { key: "f32xf32->f32" value: 1013920513692162 }
+    }
+    entries {
+      b: 4
+      m: 512
+      n: 4096
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 2124118346191889 }
+      flops { key: "f32xf32->f32" value: 1221201960762013 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 256
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 98112374269005 }
+      flops { key: "f32xf32->f32" value: 114912438356164 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 256
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 189039053521126 }
+      flops { key: "f32xf32->f32" value: 211034163522012 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 256
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 314327231850117 }
+      flops { key: "f32xf32->f32" value: 327360312195121 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 256
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 318806954869358 }
+      flops { key: "f32xf32->f32" value: 194800766328011 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 256
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 488953471766848 }
+      flops { key: "f32xf32->f32" value: 235057316987740 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 512
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 187454927374301 }
+      flops { key: "f32xf32->f32" value: 195652664723032 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 512
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 340654131979695 }
+      flops { key: "f32xf32->f32" value: 318051488151658 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 512
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 527378106090373 }
+      flops { key: "f32xf32->f32" value: 479295535766097 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 512
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 680444755386565 }
+      flops { key: "f32xf32->f32" value: 646832424096385 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 512
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 756689093727977 }
+      flops { key: "f32xf32->f32" value: 389177899238854 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 1024
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 318806954869358 }
+      flops { key: "f32xf32->f32" value: 291144746203904 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 1024
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 519217516441005 }
+      flops { key: "f32xf32->f32" value: 455747803056027 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 1024
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 832358003100775 }
+      flops { key: "f32xf32->f32" value: 612866337899543 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 1024
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 985084242201834 }
+      flops { key: "f32xf32->f32" value: 863093151670434 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 1024
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 1719362408326661 }
+      flops { key: "f32xf32->f32" value: 1084587701010101 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 2048
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 431568257234726 }
+      flops { key: "f32xf32->f32" value: 403662339849624 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 2048
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 880116249180327 }
+      flops { key: "f32xf32->f32" value: 622098391657010 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 2048
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 1254371289719626 }
+      flops { key: "f32xf32->f32" value: 858993459200000 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 2048
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1636723592054494 }
+      flops { key: "f32xf32->f32" value: 1106380035033487 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 2048
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 2158275023115578 }
+      flops { key: "f32xf32->f32" value: 1319498401228878 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 4096
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 863136514469453 }
+      flops { key: "f32xf32->f32" value: 501748515887850 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 4096
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 1335499781094527 }
+      flops { key: "f32xf32->f32" value: 856936810853950 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 4096
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 1837026217279726 }
+      flops { key: "f32xf32->f32" value: 1127880067226890 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 4096
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 2294320136752137 }
+      flops { key: "f32xf32->f32" value: 1358736885795634 }
+    }
+    entries {
+      b: 4
+      m: 1024
+      n: 4096
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 2639757101162009 }
+      flops { key: "f32xf32->f32" value: 1548573029024698 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 256
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 193397302593659 }
+      flops { key: "f32xf32->f32" value: 198546934911242 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 256
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 333874945273631 }
+      flops { key: "f32xf32->f32" value: 321865055155875 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 256
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 496183837338262 }
+      flops { key: "f32xf32->f32" value: 458864027350427 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 256
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 715827882666666 }
+      flops { key: "f32xf32->f32" value: 664444198019802 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 256
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 694052001131176 }
+      flops { key: "f32xf32->f32" value: 332119339313331 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 512
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 318051488151658 }
+      flops { key: "f32xf32->f32" value: 302292180180180 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 512
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 525314003913894 }
+      flops { key: "f32xf32->f32" value: 442233041186161 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 512
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 693631669250646 }
+      flops { key: "f32xf32->f32" value: 649964784503632 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 512
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1173488332240437 }
+      flops { key: "f32xf32->f32" value: 923251783319002 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 512
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 1720740102564102 }
+      flops { key: "f32xf32->f32" value: 1154561101075268 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 1024
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 465226093587521 }
+      flops { key: "f32xf32->f32" value: 385129779053084 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 1024
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 983280058608058 }
+      flops { key: "f32xf32->f32" value: 614268778032036 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 1024
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 1272205952606635 }
+      flops { key: "f32xf32->f32" value: 815293716021260 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 1024
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1686947091908876 }
+      flops { key: "f32xf32->f32" value: 1168380657236126 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 1024
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 2179080312531709 }
+      flops { key: "f32xf32->f32" value: 1408618975832735 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 2048
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 860370051282051 }
+      flops { key: "f32xf32->f32" value: 509818659386313 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 2048
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 1299929569007264 }
+      flops { key: "f32xf32->f32" value: 861060003207698 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 2048
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 1798562519262981 }
+      flops { key: "f32xf32->f32" value: 1090646850177755 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 2048
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 2278497239257294 }
+      flops { key: "f32xf32->f32" value: 1352319677581864 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 2048
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 2640619302797418 }
+      flops { key: "f32xf32->f32" value: 1529819161531611 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 4096
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 1150848685959271 }
+      flops { key: "f32xf32->f32" value: 677867313131313 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 4096
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 1727661824617860 }
+      flops { key: "f32xf32->f32" value: 1072133623564653 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 4096
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 2355988642896324 }
+      flops { key: "f32xf32->f32" value: 1393536729381704 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 4096
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 2687714202753442 }
+      flops { key: "f32xf32->f32" value: 1639280942164334 }
+    }
+    entries {
+      b: 4
+      m: 2048
+      n: 4096
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 3068381708162172 }
+      flops { key: "f32xf32->f32" value: 1775502089202090 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 256
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 307134389016018 }
+      flops { key: "f32xf32->f32" value: 296286375275938 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 256
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 470939396491228 }
+      flops { key: "f32xf32->f32" value: 462023160068846 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 256
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 798915047619047 }
+      flops { key: "f32xf32->f32" value: 625723673659673 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 256
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1366083745547073 }
+      flops { key: "f32xf32->f32" value: 908411018612521 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 256
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 1768932164744645 }
+      flops { key: "f32xf32->f32" value: 1088987651115618 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 512
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 440780715927750 }
+      flops { key: "f32xf32->f32" value: 412343250384024 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 512
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 956989147950089 }
+      flops { key: "f32xf32->f32" value: 615633526266752 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 512
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 1298357707376058 }
+      flops { key: "f32xf32->f32" value: 852852918189038 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 512
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 1701651068145800 }
+      flops { key: "f32xf32->f32" value: 1139248619628647 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 512
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 2183511589222165 }
+      flops { key: "f32xf32->f32" value: 1388132041935158 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 1024
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 858993459200000 }
+      flops { key: "f32xf32->f32" value: 498024964749536 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 1024
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 1261741273795534 }
+      flops { key: "f32xf32->f32" value: 833650484472049 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 1024
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 1807646168350168 }
+      flops { key: "f32xf32->f32" value: 1113262648004147 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 1024
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 2266473507124010 }
+      flops { key: "f32xf32->f32" value: 1360027642811906 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 1024
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 2633283265418734 }
+      flops { key: "f32xf32->f32" value: 1576713397944199 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 2048
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 1163238484934660 }
+      flops { key: "f32xf32->f32" value: 679583432911392 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 2048
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 1750190422167889 }
+      flops { key: "f32xf32->f32" value: 1072669154845154 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 2048
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 2350830484948002 }
+      flops { key: "f32xf32->f32" value: 1435453736678294 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 2048
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 2647945311960542 }
+      flops { key: "f32xf32->f32" value: 1644616466297311 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 2048
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 2939713458574408 }
+      flops { key: "f32xf32->f32" value: 1769467359387170 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 4096
+      k: 256
+      flops { key: "bf16xbf16->bf16" value: 1444171921990585 }
+      flops { key: "f32xf32->f32" value: 836899317225253 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 4096
+      k: 512
+      flops { key: "bf16xbf16->bf16" value: 2156108080321285 }
+      flops { key: "f32xf32->f32" value: 1286688824445776 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 4096
+      k: 1024
+      flops { key: "bf16xbf16->bf16" value: 2981580906629642 }
+      flops { key: "f32xf32->f32" value: 1699294676953511 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 4096
+      k: 2048
+      flops { key: "bf16xbf16->bf16" value: 3329391685468926 }
+      flops { key: "f32xf32->f32" value: 1818155947640308 }
+    }
+    entries {
+      b: 4
+      m: 4096
+      n: 4096
+      k: 4096
+      flops { key: "bf16xbf16->bf16" value: 3577648726364015 }
+      flops { key: "f32xf32->f32" value: 1880452103751287 }
+    }
+  }
+}
 entries {
   key: "sm_90"
   value {
diff --git a/third_party/xla/xla/service/gpu/model/matmul_interpolator_test.cc b/third_party/xla/xla/service/gpu/model/matmul_interpolator_test.cc
index d2d7d0f4c04..ae1c71aa388 100644
--- a/third_party/xla/xla/service/gpu/model/matmul_interpolator_test.cc
+++ b/third_party/xla/xla/service/gpu/model/matmul_interpolator_test.cc
@@ -337,6 +337,10 @@ class MatmulInterpolatorDefaultTableTest
     return GetMatmulInterpolator(TestGpuDeviceInfo::AMDMI300DeviceInfo());
   }
 
+  std::unique_ptr<MatmulInterpolator> GetMatmulInterpolatorSm107() {
+    return GetMatmulInterpolator(TestGpuDeviceInfo::VRNVL72DeviceInfo());
+  }
+
   std::unique_ptr<MatmulInterpolator> GetMatmulInterpolatorGfx950() {
     return GetMatmulInterpolator(TestGpuDeviceInfo::AMDMI350DeviceInfo());
   }
@@ -468,6 +472,47 @@ INSTANTIATE_TEST_SUITE_P(
     [](const TestParamInfo<MatmulInterpolatorDefaultTableTest::ParamType>&
            info) { return info.param.test_name; });
 
+using Sm107DefaultTableTest = MatmulInterpolatorDefaultTableTest;
+
+TEST_P(Sm107DefaultTableTest, EstimatesExactRuntime) {
+  const auto& [_, spec, expected_duration] = GetParam();
+  ASSERT_OK_AND_ASSIGN(DotContext context,
+                       Dot(spec.b, spec.m, spec.n, spec.k, spec.lhs_type,
+                           spec.rhs_type, spec.result_type));
+  std::optional<absl::Duration> runtime =
+      GetMatmulInterpolatorSm107()->EstimatedRuntime(*context.dot);
+  ASSERT_TRUE(runtime.has_value());
+  EXPECT_NEAR(absl::ToDoubleNanoseconds(*runtime),
+              absl::ToDoubleNanoseconds(expected_duration), 1.0);
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    MatmulInterpolatorDefaultTableTestInstantiationSm107, Sm107DefaultTableTest,
+    ValuesIn<ParametrizedTestCase>({
+        {
+            /*test_name=*/"bf16_bf16_bf16",
+            /*spec=*/
+            {/*b=*/4, /*m=*/4096, /*n=*/4096, /*k=*/4096,
+             /*lhs_type=*/"bf16",
+             /*rhs_type=*/"bf16",
+             /*result_type=*/"bf16",
+             /*clock_cycles=*/0},
+            /*expected_duration=*/absl::Nanoseconds(153664),
+        },
+        {
+            /*test_name=*/"f32_f32_f32",
+            /*spec=*/
+            {/*b=*/4, /*m=*/4096, /*n=*/4096, /*k=*/4096,
+             /*lhs_type=*/"f32",
+             /*rhs_type=*/"f32",
+             /*result_type=*/"f32",
+             /*clock_cycles=*/0},
+            /*expected_duration=*/absl::Nanoseconds(292353),
+        },
+    }),
+    [](const TestParamInfo<MatmulInterpolatorDefaultTableTest::ParamType>&
+           info) { return info.param.test_name; });
+
 using H100S8Test = MatmulInterpolatorDefaultTableTest;
 
 TEST_P(H100S8Test, EstimatesRuntimeForS8) {

From 09ed628867c4681d299462f69cbad39a097131bc Mon Sep 17 00:00:00 2001
From: Emilio Cota <ecg@google.com>
Date: Mon, 12 Oct 2026 08:36:51 -0700
Subject: [PATCH 0589/1000] [xla:xtile] add bufferization tests for in-bounds
 and unit-size tiles

In preparation for upcoming changes; with these tests we'll clearly
see what these upcoming changes do.

PiperOrigin-RevId: 997940464
---
 .../xtile/ir/transforms/tests/bufferize.mlir  | 164 +++++++++++++++++-
 1 file changed, 160 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/codegen/xtile/ir/transforms/tests/bufferize.mlir b/third_party/xla/xla/codegen/xtile/ir/transforms/tests/bufferize.mlir
index 5093a70b13d..1aa92678f4f 100644
--- a/third_party/xla/xla/codegen/xtile/ir/transforms/tests/bufferize.mlir
+++ b/third_party/xla/xla/codegen/xtile/ir/transforms/tests/bufferize.mlir
@@ -127,16 +127,32 @@ func.func @extract_identity(%source: memref<16xf32>, %tile_id: index) -> tensor<
 
 // CHECK: @extract_5d(%[[SOURCE:.*]]: memref<1x2x1x32768x256xf32>, %[[D1:.*]]: index, %[[D3:.*]]: index, %[[D4:.*]]: index)
 func.func @extract_5d(%source: memref<1x2x1x32768x256xf32>, %d1: index, %d3: index, %d4: index) -> tensor<1x1x1x1x16xf32> {
-  // CHECK: %[[C_0:.*]] = arith.constant 0 : index
-  // CHECK: %[[BUFFER:.*]] = scf.if {{.*}} -> (memref<1x1x1x1x16xf32>) {
+  // CHECK-DAG: %[[C_0:.*]] = arith.constant 0 : index
+  // CHECK-DAG: %[[C_1:.*]] = arith.constant 1 : index
+  // CHECK-DAG: %[[C_2:.*]] = arith.constant 2 : index
+  // CHECK-DAG: %[[C_16:.*]] = arith.constant 16 : index
+  // CHECK-DAG: %[[C_256:.*]] = arith.constant 256 : index
+  // CHECK-DAG: %[[C_32768:.*]] = arith.constant 32768 : index
+  // CHECK: %[[DIFF_1:.*]] = arith.subi %[[C_2]], %[[D1]] : index
+  // CHECK: %[[SIZE_1:.*]] = arith.minsi %[[DIFF_1]], %[[C_1]] : index
+  // CHECK: %[[DIFF_3:.*]] = arith.subi %[[C_32768]], %[[D3]] : index
+  // CHECK: %[[SIZE_3:.*]] = arith.minsi %[[DIFF_3]], %[[C_1]] : index
+  // CHECK: %[[DIFF_4:.*]] = arith.subi %[[C_256]], %[[D4]] : index
+  // CHECK: %[[SIZE_4:.*]] = arith.minsi %[[DIFF_4]], %[[C_16]] : index
+  // CHECK: %[[IS_FULL_1:.*]] = arith.cmpi eq, %[[SIZE_1]], %[[C_1]] : index
+  // CHECK: %[[IS_FULL_3:.*]] = arith.cmpi eq, %[[SIZE_3]], %[[C_1]] : index
+  // CHECK: %[[IS_FULL_13:.*]] = arith.andi %[[IS_FULL_1]], %[[IS_FULL_3]] : i1
+  // CHECK: %[[IS_FULL_4:.*]] = arith.cmpi eq, %[[SIZE_4]], %[[C_16]] : index
+  // CHECK: %[[IS_FULL_TILE:.*]] = arith.andi %[[IS_FULL_13]], %[[IS_FULL_4]] : i1
+  // CHECK: %[[BUFFER:.*]] = scf.if %[[IS_FULL_TILE]] -> (memref<1x1x1x1x16xf32>) {
   // CHECK:   %[[SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[C_0]], %[[D1]], %[[C_0]], %[[D3]], %[[D4]]] [1, 1, 1, 1, 16] [0, 1, 0, 1, 1] : memref<1x2x1x32768x256xf32> to memref<1x1x1x1x16xf32, strided<[0, 8388608, 0, 256, 1], offset: ?>>
   // CHECK:   %[[ALLOC_IF:.*]] = memref.alloc() : memref<1x1x1x1x16xf32>
   // CHECK:   memref.copy %[[SUBVIEW]], %[[ALLOC_IF]]
   // CHECK:   scf.yield %[[ALLOC_IF]] : memref<1x1x1x1x16xf32>
   // CHECK: } else {
-  // CHECK:   %[[INPUT_SUBVIEW:.*]] = memref.subview %{{.*}} : memref<1x2x1x32768x256xf32> to memref<?x?x?x?x?xf32, strided<[0, 8388608, 0, 256, 1], offset: ?>>
+  // CHECK:   %[[INPUT_SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[C_0]], %[[D1]], %[[C_0]], %[[D3]], %[[D4]]] [%[[C_1]], %[[SIZE_1]], %[[C_1]], %[[SIZE_3]], %[[SIZE_4]]] [0, 1, 0, 1, 1] : memref<1x2x1x32768x256xf32> to memref<?x?x?x?x?xf32, strided<[0, 8388608, 0, 256, 1], offset: ?>>
   // CHECK:   %[[ALLOC_ELSE:.*]] = memref.alloc() : memref<1x1x1x1x16xf32>
-  // CHECK:   %[[ALLOC_SUBVIEW:.*]] = memref.subview %[[ALLOC_ELSE]][0, 0, 0, 0, 0] [1, %{{.*}}, 1, %{{.*}}, %{{.*}}] [1, 1, 1, 1, 1] : memref<1x1x1x1x16xf32> to memref<1x?x1x?x?xf32, strided<[16, 16, 16, 16, 1]>>
+  // CHECK:   %[[ALLOC_SUBVIEW:.*]] = memref.subview %[[ALLOC_ELSE]][0, 0, 0, 0, 0] [1, %[[SIZE_1]], 1, %[[SIZE_3]], %[[SIZE_4]]] [1, 1, 1, 1, 1] : memref<1x1x1x1x16xf32> to memref<1x?x1x?x?xf32, strided<[16, 16, 16, 16, 1]>>
   // CHECK:   memref.copy %[[INPUT_SUBVIEW]], %[[ALLOC_SUBVIEW]]
   // CHECK:   scf.yield %[[ALLOC_ELSE]] : memref<1x1x1x1x16xf32>
   // CHECK: }
@@ -161,3 +177,143 @@ func.func @extract_static(%source: memref<16xf32>) -> tensor<8xf32> {
   %tile = xtile.extract %source[%c0][8][1] : memref<16xf32> -> tensor<8xf32>
   return %tile : tensor<8xf32>
 }
+
+// -----
+
+#map = #xla.indexing_map<"(d0) -> (d0 * 8), domain: d0 in [0, 1]">
+// CHECK: #[[MAP:.*]] = #xla.indexing_map<"(d0) -> (d0 * 8), domain: d0 in [0, 1]">
+// CHECK: @insert_in_bounds(
+// CHECK-SAME: %[[SOURCE:.*]]: tensor<8xf32>,
+// CHECK-SAME: %[[DESTINATION:.*]]: memref<16xf32>,
+// CHECK-SAME: %[[TILE_ID:.*]]: index)
+func.func @insert_in_bounds(%source: tensor<8xf32>, %destination: memref<16xf32>, %tile_id: index) {
+  // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
+  // CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index
+  // CHECK: %[[SOURCE_BUFFER:.*]] = bufferization.to_buffer %[[SOURCE]] : tensor<8xf32> to memref<8xf32, strided<[?], offset: ?>>
+  // CHECK: %[[OFFSET:.*]] = xla.apply_indexing #[[MAP]](%[[TILE_ID]])
+  // CHECK: %[[DIFF:.*]] = arith.subi %[[C16]], %[[OFFSET]] : index
+  // CHECK: %[[SIZE:.*]] = arith.minsi %[[DIFF]], %[[C8]] : index
+  // CHECK: %[[IS_FULL_TILE:.*]] = arith.cmpi eq, %[[SIZE]], %[[C8]] : index
+  // CHECK: scf.if %[[IS_FULL_TILE]] {
+  // CHECK:   %[[DESTINATION_SUBVIEW:.*]] = memref.subview %[[DESTINATION]][%[[OFFSET]]] [8] [1] : memref<16xf32> to memref<8xf32, strided<[1], offset: ?>>
+  // CHECK:   memref.copy %[[SOURCE_BUFFER]], %[[DESTINATION_SUBVIEW]] : memref<8xf32, strided<[?], offset: ?>> to memref<8xf32, strided<[1], offset: ?>>
+  // CHECK: } else {
+  // CHECK:   %[[SOURCE_SUBVIEW:.*]] = memref.subview %[[SOURCE_BUFFER]][0] [%[[SIZE]]] [1] : memref<8xf32, strided<[?], offset: ?>> to memref<?xf32, strided<[?], offset: ?>>
+  // CHECK:   %[[DESTINATION_SUBVIEW:.*]] = memref.subview %[[DESTINATION]][%[[OFFSET]]] [%[[SIZE]]] [1] : memref<16xf32> to memref<?xf32, strided<[1], offset: ?>>
+  // CHECK:   memref.copy %[[SOURCE_SUBVIEW]], %[[DESTINATION_SUBVIEW]] : memref<?xf32, strided<[?], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
+  // CHECK: }
+  %offset = xla.apply_indexing #map(%tile_id)
+  xtile.insert %source into %destination[%offset][8][1] : tensor<8xf32> -> memref<16xf32>
+  return
+}
+
+// -----
+
+#map = #xla.indexing_map<"(d0) -> (d0 * 2), domain: d0 in [0, 1]">
+// CHECK: #[[MAP:.*]] = #xla.indexing_map<"(d0) -> (d0 * 2), domain: d0 in [0, 1]">
+// CHECK: @extract_strided_in_bounds(%[[SOURCE:.*]]: memref<17xf32>, %[[TILE_ID:.*]]: index)
+func.func @extract_strided_in_bounds(%source: memref<17xf32>, %tile_id: index) -> tensor<8xf32> {
+  // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+  // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+  // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
+  // CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index
+  // CHECK: %[[OFFSET:.*]] = xla.apply_indexing #[[MAP]](%[[TILE_ID]])
+  // CHECK: %[[SHIFT:.*]] = arith.subi %[[C16]], %[[OFFSET]] : index
+  // CHECK: %[[STRIDED_SHIFT:.*]] = arith.divsi %[[SHIFT]], %[[C2]] : index
+  // CHECK: %[[ELEMENTS_TO_END:.*]] = arith.addi %[[STRIDED_SHIFT]], %[[C1]] : index
+  // CHECK: %[[SIZE:.*]] = arith.minsi %[[ELEMENTS_TO_END]], %[[C8]] : index
+  // CHECK: %[[IS_FULL_TILE:.*]] = arith.cmpi eq, %[[SIZE]], %[[C8]] : index
+  // CHECK: %[[BUFFER:.*]] = scf.if %[[IS_FULL_TILE]] -> (memref<8xf32>) {
+  // CHECK:   %[[SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[OFFSET]]] [8] [2] : memref<17xf32> to memref<8xf32, strided<[2], offset: ?>>
+  // CHECK:   %[[ALLOC_IF:.*]] = memref.alloc() : memref<8xf32>
+  // CHECK:   memref.copy %[[SUBVIEW]], %[[ALLOC_IF]] : memref<8xf32, strided<[2], offset: ?>> to memref<8xf32>
+  // CHECK:   scf.yield %[[ALLOC_IF]] : memref<8xf32>
+  // CHECK: } else {
+  // CHECK:   %[[INPUT_SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[OFFSET]]] [%[[SIZE]]] [2] : memref<17xf32> to memref<?xf32, strided<[2], offset: ?>>
+  // CHECK:   %[[ALLOC_ELSE:.*]] = memref.alloc() : memref<8xf32>
+  // CHECK:   %[[ALLOC_SUBVIEW:.*]] = memref.subview %[[ALLOC_ELSE]][0] [%[[SIZE]]] [1] : memref<8xf32> to memref<?xf32, strided<[1]>>
+  // CHECK:   memref.copy %[[INPUT_SUBVIEW]], %[[ALLOC_SUBVIEW]] : memref<?xf32, strided<[2], offset: ?>> to memref<?xf32, strided<[1]>>
+  // CHECK:   scf.yield %[[ALLOC_ELSE]] : memref<8xf32>
+  // CHECK: }
+  // CHECK: %[[TILE:.*]] = bufferization.to_tensor %[[BUFFER]] : memref<8xf32> to tensor<8xf32>
+  // CHECK: return %[[TILE]] : tensor<8xf32>
+  %offset = xla.apply_indexing #map(%tile_id)
+  %tile = xtile.extract %source[%offset][8][2] : memref<17xf32> -> tensor<8xf32>
+  return %tile : tensor<8xf32>
+}
+
+// -----
+
+#map = #xla.indexing_map<"(d0) -> (d0 * 4), domain: d0 in [0, 1]">
+// CHECK: #[[MAP:.*]] = #xla.indexing_map<"(d0) -> (d0 * 4), domain: d0 in [0, 1]">
+// CHECK: @extract_rank_reduced(%[[SOURCE:.*]]: memref<2x8x16xf32>, %[[D0:.*]]: index, %[[D1:.*]]: index, %[[D2:.*]]: index)
+func.func @extract_rank_reduced(%source: memref<2x8x16xf32>, %d0: index, %d1: index, %d2: index) -> tensor<4x8xf32> {
+  // CHECK-DAG: %[[C_4:.*]] = arith.constant 4 : index
+  // CHECK-DAG: %[[C_8:.*]] = arith.constant 8 : index
+  // CHECK-DAG: %[[C_16:.*]] = arith.constant 16 : index
+  // CHECK: %[[OFFSET_1:.*]] = xla.apply_indexing #[[MAP]](%[[D1]])
+  // CHECK: %[[DIFF_1:.*]] = arith.subi %[[C_8]], %[[OFFSET_1]] : index
+  // CHECK: %[[SIZE_1:.*]] = arith.minsi %[[DIFF_1]], %[[C_4]] : index
+  // CHECK: %[[DIFF_2:.*]] = arith.subi %[[C_16]], %[[D2]] : index
+  // CHECK: %[[SIZE_2:.*]] = arith.minsi %[[DIFF_2]], %[[C_8]] : index
+  // CHECK: %[[IS_FULL_1:.*]] = arith.cmpi eq, %[[SIZE_1]], %[[C_4]] : index
+  // CHECK: %[[IS_FULL_2:.*]] = arith.cmpi eq, %[[SIZE_2]], %[[C_8]] : index
+  // CHECK: %[[IS_FULL_TILE:.*]] = arith.andi %[[IS_FULL_1]], %[[IS_FULL_2]] : i1
+  // CHECK: %[[BUFFER:.*]] = scf.if %[[IS_FULL_TILE]] -> (memref<4x8xf32>) {
+  // CHECK:   %[[SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[D0]], %[[OFFSET_1]], %[[D2]]] [1, 4, 8] [1, 1, 1] : memref<2x8x16xf32> to memref<4x8xf32, strided<[16, 1], offset: ?>>
+  // CHECK:   %[[ALLOC_IF:.*]] = memref.alloc() : memref<4x8xf32>
+  // CHECK:   memref.copy %[[SUBVIEW]], %[[ALLOC_IF]]
+  // CHECK:   scf.yield %[[ALLOC_IF]] : memref<4x8xf32>
+  // CHECK: } else {
+  // CHECK:   %[[INPUT_SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[D0]], %[[OFFSET_1]], %[[D2]]] [1, %[[SIZE_1]], %[[SIZE_2]]] [1, 1, 1] : memref<2x8x16xf32> to memref<?x?xf32, strided<[16, 1], offset: ?>>
+  // CHECK:   %[[ALLOC_ELSE:.*]] = memref.alloc() : memref<4x8xf32>
+  // CHECK:   %[[ALLOC_SUBVIEW:.*]] = memref.subview %[[ALLOC_ELSE]][0, 0] [%[[SIZE_1]], %[[SIZE_2]]] [1, 1] : memref<4x8xf32> to memref<?x?xf32, strided<[8, 1]>>
+  // CHECK:   memref.copy %[[INPUT_SUBVIEW]], %[[ALLOC_SUBVIEW]]
+  // CHECK:   scf.yield %[[ALLOC_ELSE]] : memref<4x8xf32>
+  // CHECK: }
+  // CHECK: bufferization.to_tensor %[[BUFFER]] : memref<4x8xf32> to tensor<4x8xf32>
+  %offset_1 = xla.apply_indexing #map(%d1)
+  %tile = xtile.extract %source[%d0, %offset_1, %d2][1, 4, 8][1, 1, 1] : memref<2x8x16xf32> -> tensor<4x8xf32>
+  return %tile : tensor<4x8xf32>
+}
+
+// -----
+
+#map = #xla.indexing_map<"(d0) -> (d0 * 4), domain: d0 in [0, 1]">
+// CHECK: #[[MAP:.*]] = #xla.indexing_map<"(d0) -> (d0 * 4), domain: d0 in [0, 1]">
+// CHECK: @insert_rank_reduced(
+// CHECK-SAME: %[[SOURCE:.*]]: tensor<4x1x8xf32>,
+// CHECK-SAME: %[[DESTINATION:.*]]: memref<2x8x2x16xf32>,
+// CHECK-SAME: %[[D0:.*]]: index, %[[D1:.*]]: index, %[[D2:.*]]: index, %[[D3:.*]]: index)
+func.func @insert_rank_reduced(%source: tensor<4x1x8xf32>, %destination: memref<2x8x2x16xf32>, %d0: index, %d1: index, %d2: index, %d3: index) {
+  // CHECK-DAG: %[[C_1:.*]] = arith.constant 1 : index
+  // CHECK-DAG: %[[C_2:.*]] = arith.constant 2 : index
+  // CHECK-DAG: %[[C_4:.*]] = arith.constant 4 : index
+  // CHECK-DAG: %[[C_8:.*]] = arith.constant 8 : index
+  // CHECK-DAG: %[[C_16:.*]] = arith.constant 16 : index
+  // CHECK: %[[SOURCE_BUFFER:.*]] = bufferization.to_buffer %[[SOURCE]] : tensor<4x1x8xf32> to memref<4x1x8xf32, strided<[?, ?, ?], offset: ?>>
+  // CHECK: %[[OFFSET_1:.*]] = xla.apply_indexing #[[MAP]](%[[D1]])
+  // CHECK: %[[DIFF_1:.*]] = arith.subi %[[C_8]], %[[OFFSET_1]] : index
+  // CHECK: %[[SIZE_1:.*]] = arith.minsi %[[DIFF_1]], %[[C_4]] : index
+  // CHECK: %[[DIFF_2:.*]] = arith.subi %[[C_2]], %[[D2]] : index
+  // CHECK: %[[SIZE_2:.*]] = arith.minsi %[[DIFF_2]], %[[C_1]] : index
+  // CHECK: %[[DIFF_3:.*]] = arith.subi %[[C_16]], %[[D3]] : index
+  // CHECK: %[[SIZE_3:.*]] = arith.minsi %[[DIFF_3]], %[[C_8]] : index
+  // CHECK: %[[IS_FULL_1:.*]] = arith.cmpi eq, %[[SIZE_1]], %[[C_4]] : index
+  // CHECK: %[[IS_FULL_2:.*]] = arith.cmpi eq, %[[SIZE_2]], %[[C_1]] : index
+  // CHECK: %[[IS_FULL_12:.*]] = arith.andi %[[IS_FULL_1]], %[[IS_FULL_2]] : i1
+  // CHECK: %[[IS_FULL_3:.*]] = arith.cmpi eq, %[[SIZE_3]], %[[C_8]] : index
+  // CHECK: %[[IS_FULL_TILE:.*]] = arith.andi %[[IS_FULL_12]], %[[IS_FULL_3]] : i1
+  // CHECK: scf.if %[[IS_FULL_TILE]] {
+  // CHECK:   %[[DESTINATION_SUBVIEW:.*]] = memref.subview %[[DESTINATION]][%[[D0]], %[[OFFSET_1]], %[[D2]], %[[D3]]] [1, 4, 1, 8] [1, 1, 1, 1] : memref<2x8x2x16xf32> to memref<4x1x8xf32, strided<[32, 16, 1], offset: ?>>
+  // CHECK:   memref.copy %[[SOURCE_BUFFER]], %[[DESTINATION_SUBVIEW]]
+  // CHECK: } else {
+  // CHECK:   %[[SOURCE_SUBVIEW:.*]] = memref.subview %[[SOURCE_BUFFER]][0, 0, 0] [%[[SIZE_1]], %[[SIZE_2]], %[[SIZE_3]]] [1, 1, 1] : memref<4x1x8xf32, strided<[?, ?, ?], offset: ?>> to memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>
+  // CHECK:   %[[DESTINATION_SUBVIEW:.*]] = memref.subview %[[DESTINATION]][%[[D0]], %[[OFFSET_1]], %[[D2]], %[[D3]]] [1, %[[SIZE_1]], %[[SIZE_2]], %[[SIZE_3]]] [1, 1, 1, 1] : memref<2x8x2x16xf32> to memref<?x?x?xf32, strided<[32, 16, 1], offset: ?>>
+  // CHECK:   memref.copy %[[SOURCE_SUBVIEW]], %[[DESTINATION_SUBVIEW]]
+  // CHECK: }
+  %offset_1 = xla.apply_indexing #map(%d1)
+  xtile.insert %source into %destination[%d0, %offset_1, %d2, %d3][1, 4, 1, 8][1, 1, 1, 1] : tensor<4x1x8xf32> -> memref<2x8x2x16xf32>
+  return
+}
+

From 0c984d113201e70d9a897a357f602635a78635fe Mon Sep 17 00:00:00 2001
From: Enver Kayaaslan <enver@google.com>
Date: Mon, 12 Oct 2026 08:36:53 -0700
Subject: [PATCH 0590/1000] Optionally skip cloning computations not-containing
 instructions that the predicate is true.

It hence can be used to selectively flatten when a computation transitively contains some patterns. Note that it will clone a computation for all its callers when it does not skip cloning.

The new predicate is skip_cloning_if_does_not_contain_handler. It is a std::optional and by default it is always nullopt meaning that the default behavior is a no-op. The skipping logic will ignore skip_cloning_if_does_not_contain_handler predicate altogether.

Keeps the single version of a computation on the entire module that skip_cloning_handler predicate is true.

Keeps the single version of a computation on the entire module that
transitively does not contain an instruction that skip_cloning_if_contains_handler
predicate is true.

Note that the two predicates for skipping cloning combines as an OR. That is,
a computation skips being cloned if it holds either of the two predicates.

PiperOrigin-RevId: 997940482
---
 .../simplifiers/flatten_call_graph.cc         |  49 +++-
 .../simplifiers/flatten_call_graph.h          |  27 +-
 .../simplifiers/flatten_call_graph_test.cc    | 258 ++++++++++++++++++
 3 files changed, 320 insertions(+), 14 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/flatten_call_graph.cc b/third_party/xla/xla/hlo/transforms/simplifiers/flatten_call_graph.cc
index 35aec818a69..0ec27026ffe 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/flatten_call_graph.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/flatten_call_graph.cc
@@ -148,10 +148,22 @@ absl::StatusOr<bool> FlattenCallGraph::RunImpl(
   std::vector<HloComputation*> computations =
       module->MakeComputationPostOrder(execution_threads);
 
+  absl::flat_hash_set<const HloComputation*> cloner_candidates;
+  auto is_cloner_candidate = [&](const HloComputation* computation) {
+    return cloner_candidates.contains(computation) ||
+           absl::c_any_of(computation->instructions(),
+                          [&](const HloInstruction* instruction) {
+                            return (*skip_cloning_if_does_not_contain_handler_)(
+                                *instruction);
+                          });
+  };
+  auto should_skip_cloning = [&](const HloComputation& computation) {
+    return (skip_cloning_if_does_not_contain_handler_.has_value() &&
+            !cloner_candidates.contains(&computation)) ||
+           skip_cloning_handler_(computation);
+  };
+
   for (auto* computation : computations) {
-    if (skip_cloning_handler_(*computation)) {
-      continue;
-    }
     absl::InlinedVector<HloInstruction*, 1> callers;
     for (HloInstruction* caller : computation->caller_instructions()) {
       if ((execution_threads.empty() ||
@@ -161,6 +173,18 @@ absl::StatusOr<bool> FlattenCallGraph::RunImpl(
       }
     }
 
+    if (skip_cloning_if_does_not_contain_handler_.has_value() &&
+        is_cloner_candidate(computation)) {
+      cloner_candidates.insert(computation);
+      for (HloInstruction* caller : callers) {
+        cloner_candidates.insert(caller->parent());
+      }
+    }
+
+    if (should_skip_cloning(*computation)) {
+      continue;
+    }
+
     if (callers.empty()) {
       continue;
     }
@@ -192,22 +216,25 @@ absl::StatusOr<bool> FlattenCallGraph::RunImpl(
       }
 
       auto clone_callee = [&](HloComputation* callee) {
+        HloComputation* clone_ptr = nullptr;
         if (!module->has_schedule() ||
             !module->schedule().is_computation_scheduled(callee)) {
-          return module->AddEmbeddedComputation(callee->Clone());
+          clone_ptr = module->AddEmbeddedComputation(callee->Clone());
+        } else {
+          auto [clone, clone_sequence] = callee->CloneWithSchedule();
+          clone_ptr = module->AddEmbeddedComputation(std::move(clone));
+          module->schedule().set_sequence(clone_ptr, clone_sequence);
+        }
+        if (skip_cloning_if_does_not_contain_handler_.has_value()) {
+          cloner_candidates.insert(clone_ptr);
         }
-
-        auto [clone, clone_sequence] = callee->CloneWithSchedule();
-        HloComputation* clone_ptr =
-            module->AddEmbeddedComputation(std::move(clone));
-        module->schedule().set_sequence(clone_ptr, clone_sequence);
         return clone_ptr;
       };
 
       changed = true;
       std::vector<HloComputation*> worklist;
       caller->ReplaceCalledComputations([&](HloComputation* callee) {
-        if (callee == computation && !skip_cloning_handler_(*callee)) {
+        if (callee == computation && !should_skip_cloning(*callee)) {
           HloComputation* clone = clone_callee(callee);
           worklist.push_back(clone);
           return clone;
@@ -224,7 +251,7 @@ absl::StatusOr<bool> FlattenCallGraph::RunImpl(
             continue;
           }
           instruction->ReplaceCalledComputations([&](HloComputation* callee) {
-            if (skip_cloning_handler_(*callee)) {
+            if (should_skip_cloning(*callee)) {
               return callee;
             }
             HloComputation* clone = clone_callee(callee);
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/flatten_call_graph.h b/third_party/xla/xla/hlo/transforms/simplifiers/flatten_call_graph.h
index 7b8864d7f16..5ffbe4da6e4 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/flatten_call_graph.h
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/flatten_call_graph.h
@@ -18,6 +18,7 @@ limitations under the License.
 #ifndef XLA_HLO_TRANSFORMS_SIMPLIFIERS_FLATTEN_CALL_GRAPH_H_
 #define XLA_HLO_TRANSFORMS_SIMPLIFIERS_FLATTEN_CALL_GRAPH_H_
 
+#include <optional>
 #include <utility>
 
 #include "absl/container/flat_hash_set.h"
@@ -32,13 +33,31 @@ namespace xla {
 // Flattening associates each call site with a unique computation (for
 // sequential calling contexts) This simplifies buffer assignment and
 // points-to analysis (see b/36865746 for details).
+//
+// Keeps the single version of a computation on the entire module that
+// `skip_cloning_handler_` predicate is true.
+//
+// If `skip_cloning_if_does_not_contain_handler_` is provided, keeps the single
+// version of a computation on the entire module that does not transitively
+// contain an instruction that `skip_cloning_if_does_not_contain_handler_`
+// predicate is true.
+//
+// Note that the two predicates for skipping cloning combines as an OR. That is,
+// a computation skips being cloned if it holds either of the two predicates.
+//
+// TODO(b/260601110): Combine the two skip handler.
 class FlattenCallGraph : public HloModulePass {
  public:
   FlattenCallGraph()
-      : skip_cloning_handler_([](const HloComputation&) { return false; }) {}
+      : skip_cloning_handler_([](const HloComputation&) { return false; }),
+        skip_cloning_if_does_not_contain_handler_(std::nullopt) {}
   explicit FlattenCallGraph(
-      absl::AnyInvocable<bool(const HloComputation&)> skip_cloning_handler)
-      : skip_cloning_handler_(std::move(skip_cloning_handler)) {}
+      absl::AnyInvocable<bool(const HloComputation&)> skip_cloning_handler,
+      std::optional<absl::AnyInvocable<bool(const HloInstruction&)>>
+          skip_cloning_if_does_not_contain_handler = std::nullopt)
+      : skip_cloning_handler_(std::move(skip_cloning_handler)),
+        skip_cloning_if_does_not_contain_handler_(
+            std::move(skip_cloning_if_does_not_contain_handler)) {}
 
   // A skip cloning handler that skips cloning for computations that are only
   // called by kCall instructions.
@@ -59,6 +78,8 @@ class FlattenCallGraph : public HloModulePass {
 
  private:
   absl::AnyInvocable<bool(const HloComputation&)> skip_cloning_handler_;
+  std::optional<absl::AnyInvocable<bool(const HloInstruction&)>>
+      skip_cloning_if_does_not_contain_handler_;
 };
 
 }  // namespace xla
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/flatten_call_graph_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/flatten_call_graph_test.cc
index ce5c202f18a..56cc7bb4e87 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/flatten_call_graph_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/flatten_call_graph_test.cc
@@ -1809,5 +1809,263 @@ ENTRY main {
   EXPECT_NE(inner_call0->to_apply(), inner_call1->to_apply());
 }
 
+TEST_F(FlattenCallGraphTest, SkipCloningIfDoesNotContain_DirectMatch) {
+  absl::string_view hlo_string = R"hlo(
+HloModule DirectMatch
+
+// CHECK-NOT: %skip_comp.clone
+%skip_comp (param: f32[]) -> f32[] {
+  %param = f32[] parameter(0)
+  ROOT %neg = f32[] negate(%param)
+}
+
+// CHECK: %clone_comp (
+// CHECK: %clone_comp.clone (
+%clone_comp (param: f32[]) -> f32[] {
+  %param = f32[] parameter(0)
+  ROOT %custom = f32[] custom-call(%param), custom_call_target="must_clone"
+}
+
+// CHECK: ENTRY %main
+// CHECK: %call_skip0 = f32[] call({{.*}}), to_apply=%skip_comp
+// CHECK: %call_skip1 = f32[] call({{.*}}), to_apply=%skip_comp
+// CHECK: %call_clone0 = f32[] call({{.*}}), to_apply=%clone_comp
+// CHECK: %call_clone1 = f32[] call({{.*}}), to_apply=%clone_comp.clone
+ENTRY %main (param: f32[]) -> (f32[], f32[], f32[], f32[]) {
+  %param = f32[] parameter(0)
+  %call_skip0 = f32[] call(%param), to_apply=%skip_comp
+  %call_skip1 = f32[] call(%param), to_apply=%skip_comp
+  %call_clone0 = f32[] call(%param), to_apply=%clone_comp
+  %call_clone1 = f32[] call(%param), to_apply=%clone_comp
+  ROOT %tuple = (f32[], f32[], f32[], f32[]) tuple(%call_skip0, %call_skip1, %call_clone0, %call_clone1)
+}
+)hlo";
+
+  FlattenCallGraph pass([](const HloComputation&) { return false; },
+                        [](const HloInstruction& inst) {
+                          return inst.IsCustomCall("must_clone");
+                        });
+  RunAndFilecheckHloRewrite(hlo_string, std::move(pass));
+}
+
+TEST_F(FlattenCallGraphTest, SkipCloningIfDoesNotContain_TransitiveMatch) {
+  absl::string_view hlo_string = R"hlo(
+HloModule TransitiveMatch
+
+// CHECK-NOT: %leaf_skip.clone (
+%leaf_skip (param: f32[]) -> f32[] {
+  %param = f32[] parameter(0)
+  ROOT %neg = f32[] negate(%param)
+}
+
+// CHECK: %leaf_clone.clone (
+%leaf_clone (param: f32[]) -> f32[] {
+  %param = f32[] parameter(0)
+  ROOT %custom = f32[] custom-call(%param), custom_call_target="must_clone"
+}
+
+// CHECK: %mid_clone (
+// CHECK-NEXT: %param.2 = f32[] parameter(0)
+// CHECK-NEXT: %s0_ = f32[] call(%param.2), to_apply=%leaf_skip
+// CHECK-NEXT: %s1_ = f32[] call(%s0_), to_apply=%leaf_skip
+// CHECK-NEXT: %c0 = f32[] call(%s1_), to_apply=%leaf_clone
+// CHECK-NEXT: ROOT %c1 = f32[] call(%c0), to_apply=%leaf_clone.clone
+//
+
+// CHECK: %leaf_clone.clone.1 (
+// CHECK: %leaf_clone.clone.clone (
+// 
+// CHECK: %mid_clone.clone (
+// CHECK-NEXT: %param.5 = f32[] parameter(0)
+// CHECK-NEXT: %s0_.1 = f32[] call(%param.5), to_apply=%leaf_skip
+// CHECK-NEXT: %s1_.1 = f32[] call(%s0_.1), to_apply=%leaf_skip
+// CHECK-NEXT: %c0.1 = f32[] call(%s1_.1), to_apply=%leaf_clone.clone.1
+// CHECK-NEXT: ROOT %c1.1 = f32[] call(%c0.1), to_apply=%leaf_clone.clone.clone
+%mid_clone (param: f32[]) -> f32[] {
+  %param = f32[] parameter(0)
+  %s0_ = f32[] call(%param), to_apply=%leaf_skip
+  %s1_ = f32[] call(%s0_), to_apply=%leaf_skip
+  %c0 = f32[] call(%s1_), to_apply=%leaf_clone
+  ROOT %c1 = f32[] call(%c0), to_apply=%leaf_clone
+}
+
+// CHECK: ENTRY %main (
+// CHECK-NEXT: %param.3 = f32[] parameter(0)
+// CHECK-NEXT: %call0 = f32[] call(%param.3), to_apply=%mid_clone
+// CHECK-NEXT: %call1 = f32[] call(%param.3), to_apply=%mid_clone.clone
+// CHECK-NEXT: ROOT %tuple = (f32[], f32[]) tuple(%call0, %call1)
+ENTRY %main (param: f32[]) -> (f32[], f32[]) {
+  %param = f32[] parameter(0)
+  %call0 = f32[] call(%param), to_apply=%mid_clone
+  %call1 = f32[] call(%param), to_apply=%mid_clone
+  ROOT %tuple = (f32[], f32[]) tuple(%call0, %call1)
+}
+)hlo";
+
+  FlattenCallGraph pass([](const HloComputation&) { return false; },
+                        [](const HloInstruction& inst) {
+                          return inst.IsCustomCall("must_clone");
+                        });
+  RunAndFilecheckHloRewrite(hlo_string, std::move(pass));
+}
+
+TEST_F(FlattenCallGraphTest,
+       SkipCloningIfDoesNotContain_TransitiveSingleCallToBar) {
+  absl::string_view hlo_string = R"hlo(
+HloModule TransitiveSingleCallToBar
+
+// CHECK-NOT: %bar.clone (
+%bar (param: f32[]) -> f32[] {
+  %param = f32[] parameter(0)
+  ROOT %add = f32[] add(%param, %param)
+}
+
+// CHECK: %foo (
+// CHECK: ROOT %call_bar = f32[] call(%custom), to_apply=%bar
+//
+// CHECK: %foo.clone (
+// CHECK: ROOT %call_bar.1 = f32[] call(%custom.1), to_apply=%bar
+%foo (param: f32[]) -> f32[] {
+  %param = f32[] parameter(0)
+  %custom = f32[] custom-call(%param), custom_call_target="must_clone"
+  ROOT %call_bar = f32[] call(%custom), to_apply=%bar
+}
+
+// CHECK: ENTRY %main (
+// CHECK: %call0 = f32[] call(%param.2), to_apply=%foo
+// CHECK: %call1 = f32[] call(%param.2), to_apply=%foo.clone
+ENTRY %main (param: f32[]) -> (f32[], f32[]) {
+  %param = f32[] parameter(0)
+  %call0 = f32[] call(%param), to_apply=%foo
+  %call1 = f32[] call(%param), to_apply=%foo
+  ROOT %tuple = (f32[], f32[]) tuple(%call0, %call1)
+}
+)hlo";
+
+  FlattenCallGraph pass([](const HloComputation&) { return false; },
+                        [](const HloInstruction& inst) {
+                          return inst.IsCustomCall("must_clone");
+                        });
+  RunAndFilecheckHloRewrite(hlo_string, std::move(pass));
+}
+
+TEST_F(FlattenCallGraphTest, SkipCloningIfDoesNotContain_CombinedWithHandler) {
+  absl::string_view hlo_string = R"hlo(
+HloModule CombinedWithHandler
+
+// CHECK-NOT: %skip_by_inst.clone
+%skip_by_inst (param: f32[]) -> f32[] {
+  %param = f32[] parameter(0)
+  ROOT %neg = f32[] negate(%param)
+}
+
+// CHECK-NOT: %skip_by_comp.clone
+%skip_by_comp (param: f32[]) -> f32[] {
+  %param = f32[] parameter(0)
+  ROOT %custom = f32[] custom-call(%param), custom_call_target="must_clone"
+}
+
+// CHECK: %clone_comp (
+// CHECK: ROOT %call_comp = f32[] call({{.*}}), to_apply=%skip_by_comp
+
+// CHECK: %clone_comp.clone (
+// CHECK: ROOT %call_comp{{.*}} = f32[] call({{.*}}), to_apply=%skip_by_comp
+%clone_comp (param: f32[]) -> f32[] {
+  %param = f32[] parameter(0)
+  ROOT %call_comp = f32[] call(%param), to_apply=%skip_by_comp
+}
+
+// CHECK: ENTRY %main
+// CHECK: %call_clone0 = f32[] call({{.*}}), to_apply=%clone_comp
+// CHECK: %call_clone1 = f32[] call({{.*}}), to_apply=%clone_comp.clone
+// CHECK: %call_inst0 = f32[] call({{.*}}), to_apply=%skip_by_inst
+// CHECK: %call_inst1 = f32[] call({{.*}}), to_apply=%skip_by_inst
+ENTRY %main (param: f32[]) -> (f32[], f32[], f32[], f32[]) {
+  %param = f32[] parameter(0)
+  %call_clone0 = f32[] call(%param), to_apply=%clone_comp
+  %call_clone1 = f32[] call(%param), to_apply=%clone_comp
+  %call_inst0 = f32[] call(%param), to_apply=%skip_by_inst
+  %call_inst1 = f32[] call(%param), to_apply=%skip_by_inst
+  ROOT %tuple = (f32[], f32[], f32[], f32[]) tuple(%call_clone0, %call_clone1, %call_inst0, %call_inst1)
+}
+)hlo";
+
+  FlattenCallGraph combined_pass(
+      [](const HloComputation& comp) { return comp.name() == "skip_by_comp"; },
+      [](const HloInstruction& inst) {
+        return inst.IsCustomCall("must_clone");
+      });
+  RunAndFilecheckHloRewrite(hlo_string, std::move(combined_pass));
+}
+
+TEST_F(FlattenCallGraphTest,
+       SkipCloningIfDoesNotContain_EmptyComputation_AllFalse) {
+  absl::string_view hlo_string = R"hlo(
+HloModule EmptyComputation
+
+%empty_comp (param: f32[]) -> f32[] {
+  ROOT %param = f32[] parameter(0)
+}
+
+ENTRY %main (param: f32[]) -> (f32[], f32[]) {
+  %param = f32[] parameter(0)
+  %call0 = f32[] call(%param), to_apply=%empty_comp
+  %call1 = f32[] call(%param), to_apply=%empty_comp
+  ROOT %tuple = (f32[], f32[]) tuple(%call0, %call1)
+}
+)hlo";
+
+  FlattenCallGraph pass([](const HloComputation&) { return false; },
+                        [](const HloInstruction& inst) { return false; });
+  RunAndFilecheckHloRewrite(hlo_string, std::move(pass), std::nullopt);
+}
+
+TEST_F(FlattenCallGraphTest,
+       SkipCloningIfDoesNotContain_EmptyComputation_AllTrue) {
+  absl::string_view hlo_string = R"hlo(
+HloModule EmptyComputation
+
+// CHECK: %empty_comp.clone
+%empty_comp (param: f32[]) -> f32[] {
+  ROOT %param = f32[] parameter(0)
+}
+
+ENTRY %main (param: f32[]) -> (f32[], f32[]) {
+  %param = f32[] parameter(0)
+  %call0 = f32[] call(%param), to_apply=%empty_comp
+  %call1 = f32[] call(%param), to_apply=%empty_comp
+  ROOT %tuple = (f32[], f32[]) tuple(%call0, %call1)
+}
+)hlo";
+
+  FlattenCallGraph pass([](const HloComputation&) { return false; },
+                        [](const HloInstruction& inst) { return true; });
+  RunAndFilecheckHloRewrite(hlo_string, std::move(pass));
+}
+
+TEST_F(FlattenCallGraphTest,
+       SkipCloningIfDoesNotContain_EmptyComputation_NullOpt) {
+  absl::string_view hlo_string = R"hlo(
+HloModule EmptyComputation
+
+// CHECK: %empty_comp.clone
+%empty_comp (param: f32[]) -> f32[] {
+  ROOT %param = f32[] parameter(0)
+}
+
+ENTRY %main (param: f32[]) -> (f32[], f32[]) {
+  %param = f32[] parameter(0)
+  %call0 = f32[] call(%param), to_apply=%empty_comp
+  %call1 = f32[] call(%param), to_apply=%empty_comp
+  ROOT %tuple = (f32[], f32[]) tuple(%call0, %call1)
+}
+)hlo";
+
+  FlattenCallGraph pass(
+      [](const HloComputation&) { return false; },
+      /*skip_cloning_if_does_not_contain_handler_=*/std::nullopt);
+  RunAndFilecheckHloRewrite(hlo_string, std::move(pass));
+}
+
 }  // namespace
 }  // namespace xla

From e0205e3191d84ceda5485a6a27a5ddc9ed2a974d Mon Sep 17 00:00:00 2001
From: Mikhail Goncharov <goncharov@google.com>
Date: Mon, 12 Oct 2026 08:41:18 -0700
Subject: [PATCH 0591/1000] [XLA:GPU] fail in symbolic tile analysis

This path is being depricated and check allows us and users to see "last warning" before we drop the code completely.

updated concat_prune_inactive_small_tile that was added recently.

PiperOrigin-RevId: 997942964
---
 .../concat_prune_inactive_small_tile.hlo      |  9 +----
 .../codegen/tiling/symbolic_tile_analysis.cc  |  9 +----
 .../tiling/symbolic_tile_analysis_test.cc     | 39 +++++++++++++++++++
 .../tiling/tiled_hlo_computation_test.cc      |  4 ++
 .../codegen/tiling/tiled_hlo_schedule_test.cc |  3 ++
 .../tiling/tiling_specification_test.cc       |  4 ++
 .../codegen/xtile/codegen/fusion_emitter.cc   |  3 +-
 .../codegen/tiled_emitter_constraints_test.cc |  4 ++
 .../model/triton_emitter_constraints_test.cc  |  4 ++
 9 files changed, 64 insertions(+), 15 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive_small_tile.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive_small_tile.hlo
index dba7de98b75..a85a1167b3a 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive_small_tile.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive_small_tile.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s 2>&1 | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s 2>&1 | FileCheck %s --check-prefixes=CHECK-TILING
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN: fusion_to_triton %s 2>&1 | FileCheck %s
+// RUN: triton_test_correctness %s
 
 // This is a fork of concat_prune_inactive.hlo that uses a smaller concat-dim
 // tile size (16 instead of 64). With the sign-agnostic boundary-alignment
@@ -17,8 +14,6 @@
 
 // CHECK:      xtile.entry_func @triton_fn
 
-// CHECK-TILING:      xtile.entry_func @triton_fn
-
 HloModule min_repro, is_scheduled=true, entry_computation_layout={(f32[64,10]{1,0}, f32[1,16,10]{2,1,0}, f32[1,16,10]{2,1,0})->f32[64,10]{1,0}}
 
 %fused_computation (param_0: f32[64,10], param_1: f32[1,16,10], param_2: f32[1,16,10]) -> f32[64,10] {
diff --git a/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis.cc b/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis.cc
index 38e915a268b..ddd77bc6e4d 100644
--- a/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis.cc
+++ b/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis.cc
@@ -1398,13 +1398,8 @@ SymbolicTileAnalysis::AnalyzeFromInstruction(
 /*static*/ SymbolicTileAnalysisOrError SymbolicTileAnalysis::AnalyzeFusion(
     const HloFusionAdaptor& fusion, MLIRContext* mlir_context,
     EmitterSpecificConstraintsBuilder emitter_specific_constraints_builder) {
-  CHECK(!fusion.GetRoots()
-             .front()
-             .instruction()
-             .GetModule()
-             ->config()
-             .debug_options()
-             .xla_gpu_experimental_enable_tiling_propagation());
+  LOG(FATAL) << "Calling into symbolic tile analysis is disallowed. "
+             << "Please use experimental tiling propagation instead.";
   RegisterSymbolicExprStorage(mlir_context);
   auto real_root_index_or = GetRealRootIndex(fusion.GetRoots());
   if (!real_root_index_or.ok()) {
diff --git a/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis_test.cc b/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis_test.cc
index d7dc6c536ac..ae4e5bbc456 100644
--- a/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis_test.cc
+++ b/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis_test.cc
@@ -177,10 +177,49 @@ class FakeEmitterSpecificConstraints : public EmitterSpecificConstraints {
   int64_t dim0_tile_size_;
 };
 
+using SymbolicTileAnalysisDeathTest = HloHardwareIndependentTestBase;
+
+TEST_F(SymbolicTileAnalysisDeathTest, CallingSymbolicTileAnalysisFailsCheck) {
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
+HloModule m
+
+fused_computation {
+  p0 = f32[10,10] parameter(0)
+  ROOT r = f32[10,10] abs(p0)
+}
+
+ENTRY main {
+  p0 = f32[10,10] parameter(0)
+  ROOT fusion = f32[10,10] fusion(p0), kind=kCustom, calls=fused_computation
+}
+)"));
+  mlir::MLIRContext mlir_context;
+  EXPECT_DEATH(SymbolicTileAnalysis::AnalyzeComputation(
+                   *module->entry_computation()
+                        ->root_instruction()
+                        ->fused_instructions_computation(),
+                   &mlir_context),
+               "Calling into symbolic tile analysis is disallowed");
+
+  module->mutable_config()
+      .mutable_debug_options()
+      .set_xla_gpu_experimental_enable_tiling_propagation(false);
+  EXPECT_DEATH(SymbolicTileAnalysis::AnalyzeComputation(
+                   *module->entry_computation()
+                        ->root_instruction()
+                        ->fused_instructions_computation(),
+                   &mlir_context),
+               "Calling into symbolic tile analysis is disallowed");
+}
+
 class SymbolicTileAnalysisTest : public HloHardwareIndependentTestBase {
  public:
   SymbolicTileAnalysisTest() = default;
 
+  void SetUp() override {
+    GTEST_SKIP() << "Legacy symbolic tile analysis is disallowed.";
+  }
+
   std::optional<SymbolicTileAnalysis> TryAnalyzeModule(
       HloModule* module,
       EmitterSpecificConstraintsBuilder emitter_specific_constraints_builder =
diff --git a/third_party/xla/xla/codegen/tiling/tiled_hlo_computation_test.cc b/third_party/xla/xla/codegen/tiling/tiled_hlo_computation_test.cc
index 853b69b13c5..d2e3c2e9736 100644
--- a/third_party/xla/xla/codegen/tiling/tiled_hlo_computation_test.cc
+++ b/third_party/xla/xla/codegen/tiling/tiled_hlo_computation_test.cc
@@ -45,6 +45,10 @@ class TiledHloComputationTest : public HloHardwareIndependentTestBase {
  public:
   TiledHloComputationTest() { RegisterSymbolicExprStorage(&mlir_context_); }
 
+  void SetUp() override {
+    GTEST_SKIP() << "Legacy symbolic tile analysis is disallowed.";
+  }
+
  protected:
   DebugOptions GetDebugOptionsForTest() const override {
     DebugOptions debug_options =
diff --git a/third_party/xla/xla/codegen/tiling/tiled_hlo_schedule_test.cc b/third_party/xla/xla/codegen/tiling/tiled_hlo_schedule_test.cc
index a58d9ef83f4..22a324ab1ce 100644
--- a/third_party/xla/xla/codegen/tiling/tiled_hlo_schedule_test.cc
+++ b/third_party/xla/xla/codegen/tiling/tiled_hlo_schedule_test.cc
@@ -52,6 +52,9 @@ using ::testing::HasSubstr;
 class TiledHloScheduleTest : public HloHardwareIndependentTestBase {
  protected:
   TiledHloScheduleTest() { RegisterSymbolicExprStorage(&mlir_context_); }
+  void SetUp() override {
+    GTEST_SKIP() << "Legacy symbolic tile analysis is disallowed.";
+  }
   DebugOptions GetDebugOptionsForTest() const override {
     DebugOptions debug_options =
         HloHardwareIndependentTestBase::GetDebugOptionsForTest();
diff --git a/third_party/xla/xla/codegen/tiling/tiling_specification_test.cc b/third_party/xla/xla/codegen/tiling/tiling_specification_test.cc
index 6a302a41de5..1d72d1c88ce 100644
--- a/third_party/xla/xla/codegen/tiling/tiling_specification_test.cc
+++ b/third_party/xla/xla/codegen/tiling/tiling_specification_test.cc
@@ -51,6 +51,10 @@ class TilingSpecificationTest : public HloHardwareIndependentTestBase {
  public:
   TilingSpecificationTest() { RegisterSymbolicExprStorage(&mlir_context_); }
 
+  void SetUp() override {
+    GTEST_SKIP() << "Legacy symbolic tile analysis is disallowed.";
+  }
+
  protected:
   DebugOptions GetDebugOptionsForTest() const override {
     DebugOptions debug_options =
diff --git a/third_party/xla/xla/codegen/xtile/codegen/fusion_emitter.cc b/third_party/xla/xla/codegen/xtile/codegen/fusion_emitter.cc
index 83b44deb93b..b7c339d07c8 100644
--- a/third_party/xla/xla/codegen/xtile/codegen/fusion_emitter.cc
+++ b/third_party/xla/xla/codegen/xtile/codegen/fusion_emitter.cc
@@ -1256,7 +1256,8 @@ absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> EmitXTileModule(
     return absl::InvalidArgumentError("Multi-output fusion is disabled.");
   }
 
-  CHECK(!debug_options.xla_gpu_experimental_enable_tiling_propagation());
+  LOG(FATAL) << "Calling into legacy fusion emitter is disallowed. "
+             << "Please use experimental tiling propagation instead.";
 
   const HloComputation* hlo_computation =
       fusion.fused_instructions_computation();
diff --git a/third_party/xla/xla/codegen/xtile/codegen/tiled_emitter_constraints_test.cc b/third_party/xla/xla/codegen/xtile/codegen/tiled_emitter_constraints_test.cc
index 863b83b4af1..ef96001e82e 100644
--- a/third_party/xla/xla/codegen/xtile/codegen/tiled_emitter_constraints_test.cc
+++ b/third_party/xla/xla/codegen/xtile/codegen/tiled_emitter_constraints_test.cc
@@ -53,6 +53,10 @@ class TiledEmitterConstraintsTest : public HloHardwareIndependentTestBase {
 
  public:
   TiledEmitterConstraintsTest() { RegisterSymbolicExprStorage(&mlir_context_); }
+
+  void SetUp() override {
+    GTEST_SKIP() << "Legacy symbolic tile analysis is disallowed.";
+  }
   std::optional<SymbolicTileAnalysis> TryAnalyzeModule(
       HloModule* module, bool with_tiled_emitter_specific_constraints = true) {
     EmitterSpecificConstraintsBuilder constraints_builder = nullptr;
diff --git a/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc b/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc
index fe4d40c7fed..95b1abb12a7 100644
--- a/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc
+++ b/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc
@@ -58,6 +58,10 @@ using ::testing::HasSubstr;
 class TritonEmitterConstraintsTest : public HloHardwareIndependentTestBase {
  public:
   TritonEmitterConstraintsTest() = default;
+
+  void SetUp() override {
+    GTEST_SKIP() << "Legacy symbolic tile analysis is disallowed.";
+  }
   std::optional<SymbolicTileAnalysis> TryAnalyzeModule(
       HloModule* module, bool with_triton_emitter_specific_constraints = true) {
     EmitterSpecificConstraintsBuilder constraints_builder = nullptr;

From 142c0acbe062afc8a96ceb08012134156c8e4026 Mon Sep 17 00:00:00 2001
From: Dillon Sharlet <dsharlet@google.com>
Date: Mon, 12 Oct 2026 09:12:43 -0700
Subject: [PATCH 0592/1000] Use YNN_NODE_FLAG_KEEP_DIMS with {-1, -2} in TFLite
 YNNPACK delegate DefineMatMul.

PiperOrigin-RevId: 997960120
---
 tensorflow/lite/delegates/ynnpack/dot.cc | 17 +++++++----------
 1 file changed, 7 insertions(+), 10 deletions(-)

diff --git a/tensorflow/lite/delegates/ynnpack/dot.cc b/tensorflow/lite/delegates/ynnpack/dot.cc
index 881b858302a..2b33160c5ff 100644
--- a/tensorflow/lite/delegates/ynnpack/dot.cc
+++ b/tensorflow/lite/delegates/ynnpack/dot.cc
@@ -195,27 +195,24 @@ TfLiteStatus DefineMatMul(TfLiteContext* context, ynn_subgraph_t subgraph,
         context, subgraph, input_a_tensor, &a_scale_id, &a_zp_id));
   }
 
-  auto transpose = [&](int rank, uint32_t& val_id) -> TfLiteStatus {
+  auto transpose = [&](uint32_t& val_id) -> TfLiteStatus {
     uint32_t transposed_id = YNN_INVALID_VALUE_ID;
-    int32_t perm[YNN_MAX_TENSOR_RANK];
-    std::iota(perm, perm + rank, 0);
-    std::swap(perm[rank - 1], perm[rank - 2]);
-
+    const int32_t perm[2] = {-1, -2};
     TF_LITE_ENSURE_YNN_STATUS(ynn_define_static_transpose(
-        subgraph, rank, perm, val_id, &transposed_id, 0));
+        subgraph, 2, perm, val_id, &transposed_id, YNN_NODE_FLAG_KEEP_DIMS));
     val_id = transposed_id;
     return kTfLiteOk;
   };
 
   if (adj_x) {
-    TF_LITE_ENSURE_STATUS(transpose(rank_a, current_a_id));
+    TF_LITE_ENSURE_STATUS(transpose(current_a_id));
     if (is_dynamically_quantized) {
-      TF_LITE_ENSURE_STATUS(transpose(rank_a, a_zp_id));
-      TF_LITE_ENSURE_STATUS(transpose(rank_a, a_scale_id));
+      TF_LITE_ENSURE_STATUS(transpose(a_zp_id));
+      TF_LITE_ENSURE_STATUS(transpose(a_scale_id));
     }
   }
   if (adj_y) {
-    TF_LITE_ENSURE_STATUS(transpose(rank_b, current_b_id));
+    TF_LITE_ENSURE_STATUS(transpose(current_b_id));
   }
 
   if (mutual_broadcast) {

From eb1a3f85870ae9e8618ca98f426700756a212c63 Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Mon, 12 Oct 2026 14:23:12 -0700
Subject: [PATCH 0593/1000] Sample half-precision logits in float32 in XLA
 Multinomial

The XLA lowering of Multinomial and StatelessMultinomial takes the
argmax of the logits minus Gumbel noise, computed in the logits' type.
For bfloat16 and float16 logits the spacing between values grows with
their size (32 at 4096 for bfloat16), so the noise was rounded away
and equal logits tied, always picking the first class; the bias is
already visible at 512.

Convert half-precision logits to float32 and draw the noise in float32,
as the non-XLA kernels sample in at least float32. The stateless op
already drew bfloat16 uniforms in float32 but converted the noise back,
so drop that special case. Other types build the same HLO as before.
---
 .../compiler/tests/categorical_op_test.py     | 24 +++++++++++++++++++
 .../compiler/tf2xla/kernels/categorical_op.cc | 15 ++++++++----
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/tensorflow/compiler/tests/categorical_op_test.py b/tensorflow/compiler/tests/categorical_op_test.py
index 365ecea97e8..5a3fe22b62b 100644
--- a/tensorflow/compiler/tests/categorical_op_test.py
+++ b/tensorflow/compiler/tests/categorical_op_test.py
@@ -138,6 +138,30 @@ class CategoricalTest(xla_test.XLATestCase):
       chi2 = self._chi2(probs, freqs)
       self.assertLess(chi2, 1e-3)
 
+  def testHalfPrecisionLargeLogits(self):
+    # Regression test for GitHub issue 128958: half-precision logits were
+    # sampled in their own type, where the spacing between values (32 at 4096
+    # for bfloat16) rounds the Gumbel noise away, so equal logits always gave
+    # the first class.
+    num_samples = 40000
+    probs = [[.5, .5]]
+    for dtype in self.float_types:
+      if dtypes.as_dtype(dtype) not in (dtypes.float16, dtypes.bfloat16):
+        continue
+      logits = np.array([[4096., 4096.]], dtype=dtype)
+      with self.subTest(dtype=dtypes.as_dtype(dtype).name):
+        freqs = self._do_sampling(logits, num_samples)
+        self.assertLess(self._chi2(probs, freqs), 1e-3)
+        if dtypes.as_dtype(dtype) == dtypes.bfloat16:
+          with self.session() as sess:
+            with self.test_scope():
+              seed_t = array_ops.placeholder(dtypes.int32, shape=[2])
+              x = stateless_random_ops.stateless_multinomial(
+                  logits, num_samples, seed_t, output_dtype=dtypes.int32)
+            y = sess.run(x, {seed_t: [0x12345678, 0xabcdef12]})
+          freqs = [[np.mean(y == k) for k in range(2)]]
+          self.assertLess(self._chi2(probs, freqs), 1e-3)
+
   def testStatelessMultinomialIsInRange(self):
     for dtype in self.float_types.intersection(
         [dtypes.float32, dtypes.bfloat16]):
diff --git a/tensorflow/compiler/tf2xla/kernels/categorical_op.cc b/tensorflow/compiler/tf2xla/kernels/categorical_op.cc
index 062422d7d65..a1ac2ea7d6a 100644
--- a/tensorflow/compiler/tf2xla/kernels/categorical_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/categorical_op.cc
@@ -101,6 +101,16 @@ class CategoricalOp : public XlaOpKernel {
     }
     xla::PrimitiveType type;
     OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(input_type(0), &type));
+    // Sample in float32 for half-precision logits (the non-XLA kernels sample
+    // in at least float32): the spacing between half-precision values grows
+    // with the logits and can round the Gumbel noise away, which makes equal
+    // logits always pick the first class.
+    xla::XlaOp sampling_logits = logits;
+    if (type == xla::F16 || type == xla::BF16) {
+      type = xla::F32;
+      uniform_shape.set_element_type(type);
+      sampling_logits = xla::ConvertElementType(logits, type);
+    }
     xla::XlaOp log_uniforms = GetLogUniforms(uniform_shape, type, ctx);
 
     if (num_samples_is_dynamic) {
@@ -113,7 +123,7 @@ class CategoricalOp : public XlaOpKernel {
     // https://hips.seas.harvard.edu/blog/2013/04/06/the-gumbel-max-trick-for-discrete-distributions/
     // TODO(b/68769470): Switch to using a cumulative sum approach.
     auto softmax_entries =
-        xla::Sub(logits, log_uniforms,
+        xla::Sub(sampling_logits, log_uniforms,
                  /*broadcast_dimensions=*/{0, class_dimension});
 
     xla::PrimitiveType xla_output_type;
@@ -166,9 +176,6 @@ class StatelessCategoricalOp : public CategoricalOp {
     xla::XlaOp seed = ctx->Input(2);
 
     xla::XlaBuilder* builder = ctx->builder();
-    if (uniform_shape.element_type() == xla::BF16) {
-      uniform_shape.set_element_type(xla::F32);
-    }
     // We want a number in (0, 1) rather than [0, 1) or (0, 1]:
     // * log(-log(0)) is ∞.
     // * log(-log(1)) is -∞.

From e9e53f1e34c88ea0921173c776e2edc1557fbc02 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 12 Oct 2026 09:40:46 -0700
Subject: [PATCH 0594/1000] Log a warning and clarify in documentation that
 `tf.debugging.check_numerics` and `tf.debugging.assert_all_finite` are no-ops
 under XLA compilation (`tf.function(jit_compile=True)`).

Fixes #118075

PiperOrigin-RevId: 997974622
---
 tensorflow/compiler/tf2xla/BUILD              |  1 +
 tensorflow/compiler/tf2xla/kernels/BUILD      |  2 +-
 .../tf2xla/kernels/check_numerics_op.cc       | 30 +++++++++++++++++--
 .../compiler/tf2xla/xla_compiler_test.cc      | 22 ++++++++++++++
 .../python_api/api_def_CheckNumerics.pbtxt    | 21 +++++++++++++
 tensorflow/python/ops/numerics.py             |  3 ++
 6 files changed, 76 insertions(+), 3 deletions(-)

diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD
index 4840af23a2d..e3e2a078863 100644
--- a/tensorflow/compiler/tf2xla/BUILD
+++ b/tensorflow/compiler/tf2xla/BUILD
@@ -1235,6 +1235,7 @@ tf_cc_test(
         "//tensorflow/cc:resource_variable_ops",
         "//tensorflow/cc:scope",
         "//tensorflow/compiler/jit:common",
+        "//tensorflow/compiler/tf2xla/kernels:xla_dummy_ops",
         "//tensorflow/compiler/tf2xla/kernels:xla_ops",
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:framework",
diff --git a/tensorflow/compiler/tf2xla/kernels/BUILD b/tensorflow/compiler/tf2xla/kernels/BUILD
index 9be92d3fed4..b1b6dd231a6 100644
--- a/tensorflow/compiler/tf2xla/kernels/BUILD
+++ b/tensorflow/compiler/tf2xla/kernels/BUILD
@@ -549,10 +549,10 @@ tf_kernel_library(
         "check_numerics_op.cc",
     ],
     deps = [
-        "//tensorflow/compiler/tf2xla:mlir_xla_op_kernel",
         "//tensorflow/compiler/tf2xla:xla_compiler",
         "//tensorflow/compiler/tf2xla:xla_op_registry",
         "//tensorflow/core:array_ops_op_lib",
+        "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:logging_ops_op_lib",
         "@com_google_absl//absl/log",
diff --git a/tensorflow/compiler/tf2xla/kernels/check_numerics_op.cc b/tensorflow/compiler/tf2xla/kernels/check_numerics_op.cc
index 66f1640eb88..73f4823d081 100644
--- a/tensorflow/compiler/tf2xla/kernels/check_numerics_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/check_numerics_op.cc
@@ -13,12 +13,38 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "tensorflow/compiler/tf2xla/mlir_xla_op_kernel.h"
+#include "absl/log/log.h"
+#include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
+#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/platform/mutex.h"
 
 namespace tensorflow {
 namespace {
 
-REGISTER_XLA_OP(Name("CheckNumerics"), MlirXlaOpKernel);
+class CheckNumericsOp : public XlaOpKernel {
+ public:
+  explicit CheckNumericsOp(OpKernelConstruction* ctx) : XlaOpKernel(ctx) {}
+  ~CheckNumericsOp() override = default;
+
+  void Compile(XlaOpKernelContext* ctx) override {
+    static mutex mu(tensorflow::LINKER_INITIALIZED);
+    static int log_counter = 0;
+
+    mutex_lock l(mu);
+    if (log_counter < 20) {
+      ++log_counter;
+      LOG(WARNING) << "Ignoring CheckNumerics operator " << name();
+    }
+    ctx->SetOutput(0, ctx->Input(0));
+  }
+
+ private:
+  CheckNumericsOp(const CheckNumericsOp&) = delete;
+  void operator=(const CheckNumericsOp&) = delete;
+};
+
+REGISTER_XLA_OP(Name("CheckNumerics"), CheckNumericsOp);
 
 }  // anonymous namespace
 }  // namespace tensorflow
diff --git a/tensorflow/compiler/tf2xla/xla_compiler_test.cc b/tensorflow/compiler/tf2xla/xla_compiler_test.cc
index 1539a034e77..8d376b0ebd5 100644
--- a/tensorflow/compiler/tf2xla/xla_compiler_test.cc
+++ b/tensorflow/compiler/tf2xla/xla_compiler_test.cc
@@ -2430,5 +2430,27 @@ TEST_F(XlaCompilerTest, CompileFunctionXlaDeterministicFrontendAttribute) {
   }
 }
 
+TEST_F(XlaCompilerTest, CheckNumerics) {
+  Scope scope = Scope::NewRootScope().ExitOnError();
+  auto a = ops::_Arg(scope.WithOpName("A"), DT_FLOAT, 0);
+  auto check = ops::CheckNumerics(scope.WithOpName("Check"), a, "check");
+  auto ret = ops::_Retval(scope.WithOpName("Ret"), check, 0);
+  auto graph = std::make_unique<Graph>(OpRegistry::Global());
+  TF_ASSERT_OK(scope.ToGraph(graph.get()));
+
+  std::vector<XlaCompiler::Argument> args(1);
+  args[0].kind = XlaCompiler::Argument::kParameter;
+  args[0].type = DT_FLOAT;
+  args[0].shape = TensorShape({2});
+
+  XlaCompiler compiler(DefaultOptions());
+  XlaCompiler::CompilationResult result;
+  TF_ASSERT_OK(compiler.CompileGraph(XlaCompiler::CompileOptions(), "test",
+                                     std::move(graph), args, &result));
+  ASSERT_EQ(result.outputs.size(), 1);
+  EXPECT_EQ(result.outputs[0].type, DT_FLOAT);
+  EXPECT_EQ(result.outputs[0].shape, TensorShape({2}));
+}
+
 }  // namespace
 }  // namespace tensorflow
diff --git a/tensorflow/core/api_def/python_api/api_def_CheckNumerics.pbtxt b/tensorflow/core/api_def/python_api/api_def_CheckNumerics.pbtxt
index cf7a56ec782..c1cf9765651 100644
--- a/tensorflow/core/api_def/python_api/api_def_CheckNumerics.pbtxt
+++ b/tensorflow/core/api_def/python_api/api_def_CheckNumerics.pbtxt
@@ -9,4 +9,25 @@ op {
     name: "check_numerics"
     deprecation_version: 2
   }
+  description: <<END
+When run, reports an `InvalidArgument` error if `tensor` has any values
+that are not a number (NaN) or infinity (Inf). Otherwise, returns the input
+tensor.
+
+Note: Under XLA compilation (e.g. `tf.function(jit_compile=True)`), this op is
+a no-op and does not check for `NaN` or `Inf` values at runtime.
+
+Example usage:
+
+``` python
+a = tf.Variable(1.0)
+tf.debugging.check_numerics(a, message='')
+
+b = tf.Variable(np.nan)
+try:
+  tf.debugging.check_numerics(b, message='Checking b')
+except Exception as e:
+  assert "Checking b : Tensor had NaN values" in e.message
+```
+END
 }
diff --git a/tensorflow/python/ops/numerics.py b/tensorflow/python/ops/numerics.py
index f64f177e710..f7884423858 100644
--- a/tensorflow/python/ops/numerics.py
+++ b/tensorflow/python/ops/numerics.py
@@ -62,6 +62,9 @@ def verify_tensor_all_finite_v2(x, message, name=None):
      ...
   InvalidArgumentError: ...
 
+  Note: Under XLA compilation (e.g. `tf.function(jit_compile=True)`), this
+  check is a no-op and does not check for `NaN` or `Inf` values at runtime.
+
   Args:
     x: Tensor to check.
     message: Message to log on failure.

From fbf889196e9d599444c5b3f08f002af6b578a0e1 Mon Sep 17 00:00:00 2001
From: Emilio Cota <ecg@google.com>
Date: Mon, 12 Oct 2026 09:49:40 -0700
Subject: [PATCH 0595/1000] [xla:xtile] fold full-tile bounds checks statically
 during bufferization

When bufferizing xtile.extract and xtile.insert, we emit clamping arithmetic,
comparisons and an scf.if with both full tile and partial tile branches
unconditionally, even when tile dimensions are statically guaranteed
to fit within the buffer.

Fix that by resolving statically in-bounds dimensions at bufferization time:

1. When computing clamped tile sizes, return a static attribute instead of
   emitting clamping arithmetic for dimensions that have unit size or whose
   offset range proves the tile fits within the buffer.
2. Preserve those static dimensions in the clamped subview types, and only
   emit full-size comparisons for the remaining dynamic dimensions.
3. When all dimensions are statically in-bounds, emit the full-tile path
   directly without an scf.if or partial-tile fallback.

This reduces the amount of code we generate, improving compilation times.

Note: this is a prerequisite for the follow-up change, which initializes
the padding of partial tiles. This change makes the follow-up change
very cheap in terms of compile-time cost.
PiperOrigin-RevId: 997979707
---
 .../tests/elementwise/convert_s8_s32.hlo      |   1 -
 .../tests/elementwise/reduce_precision.hlo    |   1 -
 .../transpose_convert_bf16_f32.hlo            |   1 -
 .../transpose_convert_f32_bf16.hlo            |   1 -
 .../codegen/tiled/tests/reshape/reshape.hlo   |   1 -
 .../tiled/tests/reshape/reshape_c64.hlo       |   1 -
 .../tiled/tests/transpose/transpose_c64.hlo   |   1 -
 third_party/xla/xla/codegen/xtile/ir/BUILD    |   2 +
 .../xtile/ir/transforms/tests/bufferize.mlir  | 103 +++++---------
 .../codegen/xtile/ir/xtile_bufferization.cc   | 128 ++++++++++++------
 10 files changed, 119 insertions(+), 121 deletions(-)

diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_s8_s32.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_s8_s32.hlo
index 48638c87018..b829d715b18 100644
--- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_s8_s32.hlo
+++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/convert_s8_s32.hlo
@@ -41,7 +41,6 @@ fusion {
 
 // CHECK-VECTOR-LABEL: func.func @wrapped_fusion_impl
 // CHECK-VECTOR:         %[[ALLOCA_0:.*]] = memref.alloca() : memref<1x4xi8>
-// CHECK-VECTOR:         %[[ALLOCA_1:.*]] = memref.alloca() : memref<1x4xi8>
 // CHECK-VECTOR:         %[[ALLOCA_2:.*]] = memref.alloca() alignment = 64 : memref<1x4xi32>
 // CHECK-VECTOR:         memref.copy %{{.*}}, %[[ALLOCA_0]]
 
diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/reduce_precision.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/reduce_precision.hlo
index 1d4acd983d4..38947182494 100644
--- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/reduce_precision.hlo
+++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/reduce_precision.hlo
@@ -44,7 +44,6 @@ fusion {
 
 // CHECK-VECTOR-LABEL: func.func @wrapped_fusion_impl
 // CHECK-VECTOR:         %[[ALLOCA_0:.*]] = memref.alloca() : memref<1x4xf32>
-// CHECK-VECTOR:         %[[ALLOCA_1:.*]] = memref.alloca() : memref<1x4xf32>
 // CHECK-VECTOR:         %[[ALLOCA_2:.*]] = memref.alloca() alignment = 64 : memref<1x4xf32>
 // CHECK-VECTOR:         memref.copy %{{.*}}, %[[ALLOCA_0]]
 // CHECK-VECTOR:         memref.copy %{{.*}}, %{{.*}}
diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_bf16_f32.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_bf16_f32.hlo
index bf668ec4edf..4d524ce55f5 100644
--- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_bf16_f32.hlo
+++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_bf16_f32.hlo
@@ -47,7 +47,6 @@ fusion {
 
 // CHECK-VECTOR-LABEL: func.func @wrapped_fusion_impl
 // CHECK-VECTOR:         %[[ALLOCA_0:.*]] = memref.alloca() : memref<1x8x1x8xbf16>
-// CHECK-VECTOR:         %[[ALLOCA_1:.*]] = memref.alloca() : memref<1x8x1x8xbf16>
 // CHECK-VECTOR:         %[[ALLOCA_2:.*]] = memref.alloca() alignment = 64 : memref<1x1x8x8xf32>
 // CHECK-VECTOR:         scf.for
 // CHECK-VECTOR:           memref.copy
diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_f32_bf16.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_f32_bf16.hlo
index ce27287f921..264fbbbbfab 100644
--- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_f32_bf16.hlo
+++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/transpose_convert_f32_bf16.hlo
@@ -47,7 +47,6 @@ fusion {
 
 // CHECK-VECTOR-LABEL: func.func @wrapped_fusion_impl
 // CHECK-VECTOR:         %[[ALLOCA_0:.*]] = memref.alloca() : memref<1x1x8x8xf32>
-// CHECK-VECTOR:         %[[ALLOCA_1:.*]] = memref.alloca() : memref<1x1x8x8xf32>
 // CHECK-VECTOR:         %[[ALLOCA_2:.*]] = memref.alloca() alignment = 64 : memref<1x8x1x8xbf16>
 // CHECK-VECTOR:         scf.for
 // CHECK-VECTOR:           memref.copy
diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/reshape/reshape.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/reshape/reshape.hlo
index dccc1439a19..ecceaeeb2ee 100644
--- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/reshape/reshape.hlo
+++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/reshape/reshape.hlo
@@ -43,7 +43,6 @@ fusion {
 
 // CHECK-VECTOR-LABEL: func.func @wrapped_fusion_impl
 // CHECK-VECTOR:         %[[ALLOCA_0:.*]] = memref.alloca() : memref<1x4xf32>
-// CHECK-VECTOR:         %[[ALLOCA_1:.*]] = memref.alloca() : memref<1x4xf32>
 // CHECK-VECTOR:         memref.copy %{{.*}}, %[[ALLOCA_0]]
 
 // NEW-VECTOR: func.func @wrapped_fusion_impl
diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/reshape/reshape_c64.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/reshape/reshape_c64.hlo
index 75a0fcddbab..f0939135929 100644
--- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/reshape/reshape_c64.hlo
+++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/reshape/reshape_c64.hlo
@@ -41,7 +41,6 @@ fusion {
 
 // CHECK-VECTOR-LABEL: func.func @wrapped_fusion_impl
 // CHECK-VECTOR:         %[[ALLOCA_0:.*]] = memref.alloca() : memref<1x4xcomplex<f32>>
-// CHECK-VECTOR:         %[[ALLOCA_1:.*]] = memref.alloca() : memref<1x4xcomplex<f32>>
 // CHECK-VECTOR:         memref.copy %{{.*}}, %[[ALLOCA_0]]
 
 // NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl
diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_c64.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_c64.hlo
index 30c1ee05acd..b9892d4aa54 100644
--- a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_c64.hlo
+++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/transpose_c64.hlo
@@ -42,7 +42,6 @@ fusion {
 
 // CHECK-VECTOR-LABEL: func.func @wrapped_fusion_impl
 // CHECK-VECTOR:         %[[ALLOCA_0:.*]] = memref.alloca() : memref<16x1x1x16xcomplex<f32>>
-// CHECK-VECTOR:         %[[ALLOCA_1:.*]] = memref.alloca() : memref<16x1x1x16xcomplex<f32>>
 // CHECK-VECTOR:         %[[ALLOCA_2:.*]] = memref.alloca() alignment = 64 : memref<16x1x1x16xcomplex<f32>>
 // CHECK-VECTOR:         memref.copy %{{.*}}, %[[ALLOCA_0]]
 // CHECK-VECTOR:         memref.copy %[[ALLOCA_2]], %{{.*}}
diff --git a/third_party/xla/xla/codegen/xtile/ir/BUILD b/third_party/xla/xla/codegen/xtile/ir/BUILD
index 759a7fe2372..bd43d8f3e8e 100644
--- a/third_party/xla/xla/codegen/xtile/ir/BUILD
+++ b/third_party/xla/xla/codegen/xtile/ir/BUILD
@@ -106,7 +106,9 @@ cc_library(
         ":xtile_dialect_inc_gen",
         ":xtile_ops_inc_gen",
         "//xla/codegen/emitters:implicit_arith_op_builder",
+        "//xla/codegen/emitters/ir:xla",
         "//xla/hlo/analysis:indexing_analysis",
+        "//xla/hlo/analysis:interval",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/strings:string_view",
         "@llvm-project//llvm:Support",
diff --git a/third_party/xla/xla/codegen/xtile/ir/transforms/tests/bufferize.mlir b/third_party/xla/xla/codegen/xtile/ir/transforms/tests/bufferize.mlir
index 1aa92678f4f..a9c4ead3dc3 100644
--- a/third_party/xla/xla/codegen/xtile/ir/transforms/tests/bufferize.mlir
+++ b/third_party/xla/xla/codegen/xtile/ir/transforms/tests/bufferize.mlir
@@ -128,31 +128,20 @@ func.func @extract_identity(%source: memref<16xf32>, %tile_id: index) -> tensor<
 // CHECK: @extract_5d(%[[SOURCE:.*]]: memref<1x2x1x32768x256xf32>, %[[D1:.*]]: index, %[[D3:.*]]: index, %[[D4:.*]]: index)
 func.func @extract_5d(%source: memref<1x2x1x32768x256xf32>, %d1: index, %d3: index, %d4: index) -> tensor<1x1x1x1x16xf32> {
   // CHECK-DAG: %[[C_0:.*]] = arith.constant 0 : index
-  // CHECK-DAG: %[[C_1:.*]] = arith.constant 1 : index
-  // CHECK-DAG: %[[C_2:.*]] = arith.constant 2 : index
   // CHECK-DAG: %[[C_16:.*]] = arith.constant 16 : index
   // CHECK-DAG: %[[C_256:.*]] = arith.constant 256 : index
-  // CHECK-DAG: %[[C_32768:.*]] = arith.constant 32768 : index
-  // CHECK: %[[DIFF_1:.*]] = arith.subi %[[C_2]], %[[D1]] : index
-  // CHECK: %[[SIZE_1:.*]] = arith.minsi %[[DIFF_1]], %[[C_1]] : index
-  // CHECK: %[[DIFF_3:.*]] = arith.subi %[[C_32768]], %[[D3]] : index
-  // CHECK: %[[SIZE_3:.*]] = arith.minsi %[[DIFF_3]], %[[C_1]] : index
   // CHECK: %[[DIFF_4:.*]] = arith.subi %[[C_256]], %[[D4]] : index
   // CHECK: %[[SIZE_4:.*]] = arith.minsi %[[DIFF_4]], %[[C_16]] : index
-  // CHECK: %[[IS_FULL_1:.*]] = arith.cmpi eq, %[[SIZE_1]], %[[C_1]] : index
-  // CHECK: %[[IS_FULL_3:.*]] = arith.cmpi eq, %[[SIZE_3]], %[[C_1]] : index
-  // CHECK: %[[IS_FULL_13:.*]] = arith.andi %[[IS_FULL_1]], %[[IS_FULL_3]] : i1
-  // CHECK: %[[IS_FULL_4:.*]] = arith.cmpi eq, %[[SIZE_4]], %[[C_16]] : index
-  // CHECK: %[[IS_FULL_TILE:.*]] = arith.andi %[[IS_FULL_13]], %[[IS_FULL_4]] : i1
+  // CHECK: %[[IS_FULL_TILE:.*]] = arith.cmpi eq, %[[SIZE_4]], %[[C_16]] : index
   // CHECK: %[[BUFFER:.*]] = scf.if %[[IS_FULL_TILE]] -> (memref<1x1x1x1x16xf32>) {
   // CHECK:   %[[SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[C_0]], %[[D1]], %[[C_0]], %[[D3]], %[[D4]]] [1, 1, 1, 1, 16] [0, 1, 0, 1, 1] : memref<1x2x1x32768x256xf32> to memref<1x1x1x1x16xf32, strided<[0, 8388608, 0, 256, 1], offset: ?>>
   // CHECK:   %[[ALLOC_IF:.*]] = memref.alloc() : memref<1x1x1x1x16xf32>
   // CHECK:   memref.copy %[[SUBVIEW]], %[[ALLOC_IF]]
   // CHECK:   scf.yield %[[ALLOC_IF]] : memref<1x1x1x1x16xf32>
   // CHECK: } else {
-  // CHECK:   %[[INPUT_SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[C_0]], %[[D1]], %[[C_0]], %[[D3]], %[[D4]]] [%[[C_1]], %[[SIZE_1]], %[[C_1]], %[[SIZE_3]], %[[SIZE_4]]] [0, 1, 0, 1, 1] : memref<1x2x1x32768x256xf32> to memref<?x?x?x?x?xf32, strided<[0, 8388608, 0, 256, 1], offset: ?>>
+  // CHECK:   %[[INPUT_SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[C_0]], %[[D1]], %[[C_0]], %[[D3]], %[[D4]]] [1, 1, 1, 1, %[[SIZE_4]]] [0, 1, 0, 1, 1] : memref<1x2x1x32768x256xf32> to memref<1x1x1x1x?xf32, strided<[0, 8388608, 0, 256, 1], offset: ?>>
   // CHECK:   %[[ALLOC_ELSE:.*]] = memref.alloc() : memref<1x1x1x1x16xf32>
-  // CHECK:   %[[ALLOC_SUBVIEW:.*]] = memref.subview %[[ALLOC_ELSE]][0, 0, 0, 0, 0] [1, %[[SIZE_1]], 1, %[[SIZE_3]], %[[SIZE_4]]] [1, 1, 1, 1, 1] : memref<1x1x1x1x16xf32> to memref<1x?x1x?x?xf32, strided<[16, 16, 16, 16, 1]>>
+  // CHECK:   %[[ALLOC_SUBVIEW:.*]] = memref.subview %[[ALLOC_ELSE]][0, 0, 0, 0, 0] [1, 1, 1, 1, %[[SIZE_4]]] [1, 1, 1, 1, 1] : memref<1x1x1x1x16xf32> to memref<1x1x1x1x?xf32, strided<[16, 16, 16, 16, 1]>>
   // CHECK:   memref.copy %[[INPUT_SUBVIEW]], %[[ALLOC_SUBVIEW]]
   // CHECK:   scf.yield %[[ALLOC_ELSE]] : memref<1x1x1x1x16xf32>
   // CHECK: }
@@ -171,7 +160,7 @@ func.func @extract_static(%source: memref<16xf32>) -> tensor<8xf32> {
   // CHECK: %[[SUBVIEW:.*]] = memref.subview %[[SOURCE]][0] [8] [1] : memref<16xf32> to memref<8xf32, strided<[1]>>
   // CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<8xf32>
   // CHECK: memref.copy %[[SUBVIEW]], %[[ALLOC]] : memref<8xf32, strided<[1]>> to memref<8xf32>
-  // CHECK: %[[TILE:.*]] = bufferization.to_tensor %[[ALLOC]] : memref<8xf32> to tensor<8xf32>
+  // CHECK: %[[TILE:.*]] = bufferization.to_tensor %[[ALLOC]] restrict writable : memref<8xf32> to tensor<8xf32>
   // CHECK: return %[[TILE]]
   %c0 = arith.constant 0 : index
   %tile = xtile.extract %source[%c0][8][1] : memref<16xf32> -> tensor<8xf32>
@@ -187,21 +176,11 @@ func.func @extract_static(%source: memref<16xf32>) -> tensor<8xf32> {
 // CHECK-SAME: %[[DESTINATION:.*]]: memref<16xf32>,
 // CHECK-SAME: %[[TILE_ID:.*]]: index)
 func.func @insert_in_bounds(%source: tensor<8xf32>, %destination: memref<16xf32>, %tile_id: index) {
-  // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
-  // CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index
+  // CHECK-NOT: scf.if
   // CHECK: %[[SOURCE_BUFFER:.*]] = bufferization.to_buffer %[[SOURCE]] : tensor<8xf32> to memref<8xf32, strided<[?], offset: ?>>
   // CHECK: %[[OFFSET:.*]] = xla.apply_indexing #[[MAP]](%[[TILE_ID]])
-  // CHECK: %[[DIFF:.*]] = arith.subi %[[C16]], %[[OFFSET]] : index
-  // CHECK: %[[SIZE:.*]] = arith.minsi %[[DIFF]], %[[C8]] : index
-  // CHECK: %[[IS_FULL_TILE:.*]] = arith.cmpi eq, %[[SIZE]], %[[C8]] : index
-  // CHECK: scf.if %[[IS_FULL_TILE]] {
-  // CHECK:   %[[DESTINATION_SUBVIEW:.*]] = memref.subview %[[DESTINATION]][%[[OFFSET]]] [8] [1] : memref<16xf32> to memref<8xf32, strided<[1], offset: ?>>
-  // CHECK:   memref.copy %[[SOURCE_BUFFER]], %[[DESTINATION_SUBVIEW]] : memref<8xf32, strided<[?], offset: ?>> to memref<8xf32, strided<[1], offset: ?>>
-  // CHECK: } else {
-  // CHECK:   %[[SOURCE_SUBVIEW:.*]] = memref.subview %[[SOURCE_BUFFER]][0] [%[[SIZE]]] [1] : memref<8xf32, strided<[?], offset: ?>> to memref<?xf32, strided<[?], offset: ?>>
-  // CHECK:   %[[DESTINATION_SUBVIEW:.*]] = memref.subview %[[DESTINATION]][%[[OFFSET]]] [%[[SIZE]]] [1] : memref<16xf32> to memref<?xf32, strided<[1], offset: ?>>
-  // CHECK:   memref.copy %[[SOURCE_SUBVIEW]], %[[DESTINATION_SUBVIEW]] : memref<?xf32, strided<[?], offset: ?>> to memref<?xf32, strided<[1], offset: ?>>
-  // CHECK: }
+  // CHECK: %[[DESTINATION_SUBVIEW:.*]] = memref.subview %[[DESTINATION]][%[[OFFSET]]] [8] [1] : memref<16xf32> to memref<8xf32, strided<[1], offset: ?>>
+  // CHECK: memref.copy %[[SOURCE_BUFFER]], %[[DESTINATION_SUBVIEW]] : memref<8xf32, strided<[?], offset: ?>> to memref<8xf32, strided<[1], offset: ?>>
   %offset = xla.apply_indexing #map(%tile_id)
   xtile.insert %source into %destination[%offset][8][1] : tensor<8xf32> -> memref<16xf32>
   return
@@ -213,29 +192,12 @@ func.func @insert_in_bounds(%source: tensor<8xf32>, %destination: memref<16xf32>
 // CHECK: #[[MAP:.*]] = #xla.indexing_map<"(d0) -> (d0 * 2), domain: d0 in [0, 1]">
 // CHECK: @extract_strided_in_bounds(%[[SOURCE:.*]]: memref<17xf32>, %[[TILE_ID:.*]]: index)
 func.func @extract_strided_in_bounds(%source: memref<17xf32>, %tile_id: index) -> tensor<8xf32> {
-  // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
-  // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
-  // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
-  // CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index
+  // CHECK-NOT: scf.if
   // CHECK: %[[OFFSET:.*]] = xla.apply_indexing #[[MAP]](%[[TILE_ID]])
-  // CHECK: %[[SHIFT:.*]] = arith.subi %[[C16]], %[[OFFSET]] : index
-  // CHECK: %[[STRIDED_SHIFT:.*]] = arith.divsi %[[SHIFT]], %[[C2]] : index
-  // CHECK: %[[ELEMENTS_TO_END:.*]] = arith.addi %[[STRIDED_SHIFT]], %[[C1]] : index
-  // CHECK: %[[SIZE:.*]] = arith.minsi %[[ELEMENTS_TO_END]], %[[C8]] : index
-  // CHECK: %[[IS_FULL_TILE:.*]] = arith.cmpi eq, %[[SIZE]], %[[C8]] : index
-  // CHECK: %[[BUFFER:.*]] = scf.if %[[IS_FULL_TILE]] -> (memref<8xf32>) {
-  // CHECK:   %[[SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[OFFSET]]] [8] [2] : memref<17xf32> to memref<8xf32, strided<[2], offset: ?>>
-  // CHECK:   %[[ALLOC_IF:.*]] = memref.alloc() : memref<8xf32>
-  // CHECK:   memref.copy %[[SUBVIEW]], %[[ALLOC_IF]] : memref<8xf32, strided<[2], offset: ?>> to memref<8xf32>
-  // CHECK:   scf.yield %[[ALLOC_IF]] : memref<8xf32>
-  // CHECK: } else {
-  // CHECK:   %[[INPUT_SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[OFFSET]]] [%[[SIZE]]] [2] : memref<17xf32> to memref<?xf32, strided<[2], offset: ?>>
-  // CHECK:   %[[ALLOC_ELSE:.*]] = memref.alloc() : memref<8xf32>
-  // CHECK:   %[[ALLOC_SUBVIEW:.*]] = memref.subview %[[ALLOC_ELSE]][0] [%[[SIZE]]] [1] : memref<8xf32> to memref<?xf32, strided<[1]>>
-  // CHECK:   memref.copy %[[INPUT_SUBVIEW]], %[[ALLOC_SUBVIEW]] : memref<?xf32, strided<[2], offset: ?>> to memref<?xf32, strided<[1]>>
-  // CHECK:   scf.yield %[[ALLOC_ELSE]] : memref<8xf32>
-  // CHECK: }
-  // CHECK: %[[TILE:.*]] = bufferization.to_tensor %[[BUFFER]] : memref<8xf32> to tensor<8xf32>
+  // CHECK: %[[SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[OFFSET]]] [8] [2] : memref<17xf32> to memref<8xf32, strided<[2], offset: ?>>
+  // CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<8xf32>
+  // CHECK: memref.copy %[[SUBVIEW]], %[[ALLOC]] : memref<8xf32, strided<[2], offset: ?>> to memref<8xf32>
+  // CHECK: %[[TILE:.*]] = bufferization.to_tensor %[[ALLOC]] restrict writable : memref<8xf32> to tensor<8xf32>
   // CHECK: return %[[TILE]] : tensor<8xf32>
   %offset = xla.apply_indexing #map(%tile_id)
   %tile = xtile.extract %source[%offset][8][2] : memref<17xf32> -> tensor<8xf32>
@@ -244,30 +206,38 @@ func.func @extract_strided_in_bounds(%source: memref<17xf32>, %tile_id: index) -
 
 // -----
 
+#map = #xla.indexing_map<"(d0) -> (d0 * 2), domain: d0 in [0, 1]">
+// CHECK: #[[MAP:.*]] = #xla.indexing_map<"(d0) -> (d0 * 2), domain: d0 in [0, 1]">
+// CHECK: @extract_strided_out_of_bounds(%[[SOURCE:.*]]: memref<16xf32>, %[[TILE_ID:.*]]: index)
+func.func @extract_strided_out_of_bounds(%source: memref<16xf32>, %tile_id: index) -> tensor<8xf32> {
+  // CHECK: %[[OFFSET:.*]] = xla.apply_indexing #[[MAP]](%[[TILE_ID]])
+  // CHECK: scf.if
+  %offset = xla.apply_indexing #map(%tile_id)
+  %tile = xtile.extract %source[%offset][8][2] : memref<16xf32> -> tensor<8xf32>
+  return %tile : tensor<8xf32>
+}
+
+// -----
+
 #map = #xla.indexing_map<"(d0) -> (d0 * 4), domain: d0 in [0, 1]">
 // CHECK: #[[MAP:.*]] = #xla.indexing_map<"(d0) -> (d0 * 4), domain: d0 in [0, 1]">
 // CHECK: @extract_rank_reduced(%[[SOURCE:.*]]: memref<2x8x16xf32>, %[[D0:.*]]: index, %[[D1:.*]]: index, %[[D2:.*]]: index)
 func.func @extract_rank_reduced(%source: memref<2x8x16xf32>, %d0: index, %d1: index, %d2: index) -> tensor<4x8xf32> {
-  // CHECK-DAG: %[[C_4:.*]] = arith.constant 4 : index
   // CHECK-DAG: %[[C_8:.*]] = arith.constant 8 : index
   // CHECK-DAG: %[[C_16:.*]] = arith.constant 16 : index
   // CHECK: %[[OFFSET_1:.*]] = xla.apply_indexing #[[MAP]](%[[D1]])
-  // CHECK: %[[DIFF_1:.*]] = arith.subi %[[C_8]], %[[OFFSET_1]] : index
-  // CHECK: %[[SIZE_1:.*]] = arith.minsi %[[DIFF_1]], %[[C_4]] : index
   // CHECK: %[[DIFF_2:.*]] = arith.subi %[[C_16]], %[[D2]] : index
   // CHECK: %[[SIZE_2:.*]] = arith.minsi %[[DIFF_2]], %[[C_8]] : index
-  // CHECK: %[[IS_FULL_1:.*]] = arith.cmpi eq, %[[SIZE_1]], %[[C_4]] : index
-  // CHECK: %[[IS_FULL_2:.*]] = arith.cmpi eq, %[[SIZE_2]], %[[C_8]] : index
-  // CHECK: %[[IS_FULL_TILE:.*]] = arith.andi %[[IS_FULL_1]], %[[IS_FULL_2]] : i1
+  // CHECK: %[[IS_FULL_TILE:.*]] = arith.cmpi eq, %[[SIZE_2]], %[[C_8]] : index
   // CHECK: %[[BUFFER:.*]] = scf.if %[[IS_FULL_TILE]] -> (memref<4x8xf32>) {
   // CHECK:   %[[SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[D0]], %[[OFFSET_1]], %[[D2]]] [1, 4, 8] [1, 1, 1] : memref<2x8x16xf32> to memref<4x8xf32, strided<[16, 1], offset: ?>>
   // CHECK:   %[[ALLOC_IF:.*]] = memref.alloc() : memref<4x8xf32>
   // CHECK:   memref.copy %[[SUBVIEW]], %[[ALLOC_IF]]
   // CHECK:   scf.yield %[[ALLOC_IF]] : memref<4x8xf32>
   // CHECK: } else {
-  // CHECK:   %[[INPUT_SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[D0]], %[[OFFSET_1]], %[[D2]]] [1, %[[SIZE_1]], %[[SIZE_2]]] [1, 1, 1] : memref<2x8x16xf32> to memref<?x?xf32, strided<[16, 1], offset: ?>>
+  // CHECK:   %[[INPUT_SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[D0]], %[[OFFSET_1]], %[[D2]]] [1, 4, %[[SIZE_2]]] [1, 1, 1] : memref<2x8x16xf32> to memref<4x?xf32, strided<[16, 1], offset: ?>>
   // CHECK:   %[[ALLOC_ELSE:.*]] = memref.alloc() : memref<4x8xf32>
-  // CHECK:   %[[ALLOC_SUBVIEW:.*]] = memref.subview %[[ALLOC_ELSE]][0, 0] [%[[SIZE_1]], %[[SIZE_2]]] [1, 1] : memref<4x8xf32> to memref<?x?xf32, strided<[8, 1]>>
+  // CHECK:   %[[ALLOC_SUBVIEW:.*]] = memref.subview %[[ALLOC_ELSE]][0, 0] [4, %[[SIZE_2]]] [1, 1] : memref<4x8xf32> to memref<4x?xf32, strided<[8, 1]>>
   // CHECK:   memref.copy %[[INPUT_SUBVIEW]], %[[ALLOC_SUBVIEW]]
   // CHECK:   scf.yield %[[ALLOC_ELSE]] : memref<4x8xf32>
   // CHECK: }
@@ -286,30 +256,19 @@ func.func @extract_rank_reduced(%source: memref<2x8x16xf32>, %d0: index, %d1: in
 // CHECK-SAME: %[[DESTINATION:.*]]: memref<2x8x2x16xf32>,
 // CHECK-SAME: %[[D0:.*]]: index, %[[D1:.*]]: index, %[[D2:.*]]: index, %[[D3:.*]]: index)
 func.func @insert_rank_reduced(%source: tensor<4x1x8xf32>, %destination: memref<2x8x2x16xf32>, %d0: index, %d1: index, %d2: index, %d3: index) {
-  // CHECK-DAG: %[[C_1:.*]] = arith.constant 1 : index
-  // CHECK-DAG: %[[C_2:.*]] = arith.constant 2 : index
-  // CHECK-DAG: %[[C_4:.*]] = arith.constant 4 : index
   // CHECK-DAG: %[[C_8:.*]] = arith.constant 8 : index
   // CHECK-DAG: %[[C_16:.*]] = arith.constant 16 : index
   // CHECK: %[[SOURCE_BUFFER:.*]] = bufferization.to_buffer %[[SOURCE]] : tensor<4x1x8xf32> to memref<4x1x8xf32, strided<[?, ?, ?], offset: ?>>
   // CHECK: %[[OFFSET_1:.*]] = xla.apply_indexing #[[MAP]](%[[D1]])
-  // CHECK: %[[DIFF_1:.*]] = arith.subi %[[C_8]], %[[OFFSET_1]] : index
-  // CHECK: %[[SIZE_1:.*]] = arith.minsi %[[DIFF_1]], %[[C_4]] : index
-  // CHECK: %[[DIFF_2:.*]] = arith.subi %[[C_2]], %[[D2]] : index
-  // CHECK: %[[SIZE_2:.*]] = arith.minsi %[[DIFF_2]], %[[C_1]] : index
   // CHECK: %[[DIFF_3:.*]] = arith.subi %[[C_16]], %[[D3]] : index
   // CHECK: %[[SIZE_3:.*]] = arith.minsi %[[DIFF_3]], %[[C_8]] : index
-  // CHECK: %[[IS_FULL_1:.*]] = arith.cmpi eq, %[[SIZE_1]], %[[C_4]] : index
-  // CHECK: %[[IS_FULL_2:.*]] = arith.cmpi eq, %[[SIZE_2]], %[[C_1]] : index
-  // CHECK: %[[IS_FULL_12:.*]] = arith.andi %[[IS_FULL_1]], %[[IS_FULL_2]] : i1
-  // CHECK: %[[IS_FULL_3:.*]] = arith.cmpi eq, %[[SIZE_3]], %[[C_8]] : index
-  // CHECK: %[[IS_FULL_TILE:.*]] = arith.andi %[[IS_FULL_12]], %[[IS_FULL_3]] : i1
+  // CHECK: %[[IS_FULL_TILE:.*]] = arith.cmpi eq, %[[SIZE_3]], %[[C_8]] : index
   // CHECK: scf.if %[[IS_FULL_TILE]] {
   // CHECK:   %[[DESTINATION_SUBVIEW:.*]] = memref.subview %[[DESTINATION]][%[[D0]], %[[OFFSET_1]], %[[D2]], %[[D3]]] [1, 4, 1, 8] [1, 1, 1, 1] : memref<2x8x2x16xf32> to memref<4x1x8xf32, strided<[32, 16, 1], offset: ?>>
   // CHECK:   memref.copy %[[SOURCE_BUFFER]], %[[DESTINATION_SUBVIEW]]
   // CHECK: } else {
-  // CHECK:   %[[SOURCE_SUBVIEW:.*]] = memref.subview %[[SOURCE_BUFFER]][0, 0, 0] [%[[SIZE_1]], %[[SIZE_2]], %[[SIZE_3]]] [1, 1, 1] : memref<4x1x8xf32, strided<[?, ?, ?], offset: ?>> to memref<?x?x?xf32, strided<[?, ?, ?], offset: ?>>
-  // CHECK:   %[[DESTINATION_SUBVIEW:.*]] = memref.subview %[[DESTINATION]][%[[D0]], %[[OFFSET_1]], %[[D2]], %[[D3]]] [1, %[[SIZE_1]], %[[SIZE_2]], %[[SIZE_3]]] [1, 1, 1, 1] : memref<2x8x2x16xf32> to memref<?x?x?xf32, strided<[32, 16, 1], offset: ?>>
+  // CHECK:   %[[SOURCE_SUBVIEW:.*]] = memref.subview %[[SOURCE_BUFFER]][0, 0, 0] [4, 1, %[[SIZE_3]]] [1, 1, 1] : memref<4x1x8xf32, strided<[?, ?, ?], offset: ?>> to memref<4x1x?xf32, strided<[?, ?, ?], offset: ?>>
+  // CHECK:   %[[DESTINATION_SUBVIEW:.*]] = memref.subview %[[DESTINATION]][%[[D0]], %[[OFFSET_1]], %[[D2]], %[[D3]]] [1, 4, 1, %[[SIZE_3]]] [1, 1, 1, 1] : memref<2x8x2x16xf32> to memref<4x1x?xf32, strided<[32, 16, 1], offset: ?>>
   // CHECK:   memref.copy %[[SOURCE_SUBVIEW]], %[[DESTINATION_SUBVIEW]]
   // CHECK: }
   %offset_1 = xla.apply_indexing #map(%d1)
diff --git a/third_party/xla/xla/codegen/xtile/ir/xtile_bufferization.cc b/third_party/xla/xla/codegen/xtile/ir/xtile_bufferization.cc
index b9412c18bb9..72ca4fdd12c 100644
--- a/third_party/xla/xla/codegen/xtile/ir/xtile_bufferization.cc
+++ b/third_party/xla/xla/codegen/xtile/ir/xtile_bufferization.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include <algorithm>
 #include <cassert>
 #include <cstdint>
+#include <optional>
 
 #include "absl/log/check.h"
 #include "llvm/ADT/DenseSet.h"
@@ -42,7 +43,9 @@ limitations under the License.
 #include "mlir/IR/ValueRange.h"
 #include "mlir/Support/LLVM.h"
 #include "xla/codegen/emitters/implicit_arith_op_builder.h"
+#include "xla/codegen/emitters/ir/xla_ops.h"
 #include "xla/codegen/xtile/ir/xtile_ops.h"
+#include "xla/hlo/analysis/interval.h"
 
 namespace xla::xtile {
 
@@ -81,13 +84,27 @@ static llvm::SmallVector<mlir::OpFoldResult> GetClampedTileSize(
       }
       continue;
     }
+    // Tiling guarantees 0 <= offset < buffer_size, so a unit-size tile
+    // is always in-bounds.
+    if (full_tile_size == 1) {
+      tile_size.emplace_back(builder.getIndexAttr(1));
+      continue;
+    }
+    // The stride can be 0 for single element tiles.
+    // TODO(willfroom): Fix tile analysis so this never happens.
+    int64_t clamped_stride = std::max<int64_t>(stride, 1);
+    int64_t max_tile_offset = (full_tile_size - 1) * clamped_stride;
+    std::optional<Interval> offset_range = GetRange(offset);
+    if (offset_range.has_value() && offset_range->IsFeasible() &&
+        offset_range->lower >= 0 &&
+        offset_range->upper < buffer_size - max_tile_offset) {
+      tile_size.emplace_back(builder.getIndexAttr(full_tile_size));
+      continue;
+    }
     emitters::ImplicitArithOpBuilder arith_builder(
         mlir::arith::ConstantIndexOp::create(builder, (buffer_size - 1)),
         &builder);
     auto numerator = arith_builder - offset;
-    // The stride can be 0 for single element tiles.
-    // TODO(willfroom): Fix tile analysis so this never happens.
-    auto clamped_stride = std::max<int64_t>(stride, 1);
     auto bound = numerator / clamped_stride + 1;
     tile_size.emplace_back(bound.min(full_tile_size));
   }
@@ -105,8 +122,16 @@ static mlir::TypedValue<mlir::MemRefType> GetClampedSubView(
   auto strides = GetStaticFoldResult(builder, op.getStrides());
 
   mlir::RankedTensorType tile_type = op.getTile().getType();
-  llvm::SmallVector<int64_t> output_shape(tile_type.getRank(),
-                                          mlir::ShapedType::kDynamic);
+  llvm::SmallDenseSet<unsigned> reduced_dims = op.getReducedDimensions();
+  llvm::SmallVector<int64_t> output_shape;
+  output_shape.reserve(tile_type.getRank());
+  for (auto [idx, size] : llvm::enumerate(tile_size)) {
+    if (!reduced_dims.contains(idx)) {
+      auto attr = size.dyn_cast<mlir::Attribute>();
+      output_shape.push_back(attr ? mlir::cast<mlir::IntegerAttr>(attr).getInt()
+                                  : mlir::ShapedType::kDynamic);
+    }
+  }
   mlir::MemRefType subview_type =
       mlir::memref::SubViewOp::inferRankReducedResultType(
           output_shape, op.getBuffer().getType(), offsets, tile_size, strides);
@@ -168,20 +193,23 @@ static mlir::TypedValue<mlir::RankedTensorType> GetTensorSlice(
                                               tile_size, strides);
 }
 
-static mlir::Value TileIsFullSize(mlir::ImplicitLocOpBuilder& builder,
-                                  TiledBufferInterface op) {
+// Returns the dynamic condition checking whether the tile is full size, or
+// std::nullopt if the tile is statically known to be full size.
+static std::optional<mlir::Value> TileIsFullSize(
+    mlir::ImplicitLocOpBuilder& builder, TiledBufferInterface op) {
   llvm::SmallVector<mlir::OpFoldResult> clamped_tile_size =
       GetClampedTileSize(builder, op, false);
-  mlir::Value is_full_size =
-      mlir::arith::ConstantIntOp::create(builder, builder.getI1Type(), true);
+  std::optional<mlir::Value> is_full_size;
   for (auto [dim_idx, tile_dim_size] : llvm::enumerate(clamped_tile_size)) {
     if (auto value = tile_dim_size.dyn_cast<mlir::Value>()) {
       mlir::Value is_full_size_dim = mlir::arith::CmpIOp::create(
           builder, mlir::arith::CmpIPredicate::eq, value,
           mlir::arith::ConstantIndexOp::create(
               builder, op.getTile().getType().getDimSize(dim_idx)));
-      is_full_size =
-          mlir::arith::AndIOp::create(builder, is_full_size, is_full_size_dim);
+      is_full_size = is_full_size.has_value()
+                         ? mlir::arith::AndIOp::create(builder, *is_full_size,
+                                                       is_full_size_dim)
+                         : is_full_size_dim;
     }
   }
   return is_full_size;
@@ -249,33 +277,39 @@ llvm::LogicalResult ExtractTileOp::bufferize(
     mlir::bufferization::BufferizationState& state) {
   mlir::ImplicitLocOpBuilder builder(getLoc(), rewriter);
 
-  mlir::Value is_full_size = TileIsFullSize(builder, *this);
+  auto emit_full_tile = [&](mlir::ImplicitLocOpBuilder& b) -> mlir::Value {
+    auto buffer = GetFullTileSubView(b, *this);
+    if (buffer.getType().getLayout().isIdentity()) {
+      return mlir::bufferization::ToTensorOp::create(b, getType(), buffer);
+    }
+    // If the buffer doesn't have an identity layout, we can get a
+    // miscompile during bufferization as some ops don't support
+    // non-identity layouts. So we allocate a new buffer with the same
+    // shape but default layout.
+    // TODO(willfroom): Look into how we can remove this constraint.
+    mlir::MemRefType default_buffer_type =
+        mlir::MemRefType::Builder(buffer.getType()).setLayout(nullptr);
+    auto default_buffer = mlir::memref::AllocOp::create(b, default_buffer_type);
+    mlir::memref::CopyOp::create(b, buffer, default_buffer);
+    auto to_tensor_op =
+        mlir::bufferization::ToTensorOp::create(b, getType(), default_buffer);
+    to_tensor_op.setWritable(true);
+    to_tensor_op.setRestrict(true);
+    return to_tensor_op;
+  };
+
+  std::optional<mlir::Value> is_full_size = TileIsFullSize(builder, *this);
+  if (!is_full_size.has_value()) {
+    rewriter.replaceOp(getOperation(), emit_full_tile(builder));
+    return mlir::success();
+  }
+
   auto if_op = mlir::scf::IfOp::create(
-      builder, is_full_size,
+      builder, *is_full_size,
       [&](mlir::OpBuilder& builder, mlir::Location loc) {
         mlir::ImplicitLocOpBuilder then_builder(loc, builder);
-        auto buffer = GetFullTileSubView(then_builder, *this);
-        if (buffer.getType().getLayout().isIdentity()) {
-          auto to_tensor_op = mlir::bufferization::ToTensorOp::create(
-              then_builder, getType(), buffer);
-          mlir::scf::YieldOp::create(then_builder, {to_tensor_op});
-        } else {
-          // If the buffer doesn't have an identity layout, we can get a
-          // miscompile during bufferization as some ops don't support
-          // non-identity layouts. So we allocate a new buffer with the same
-          // shape but default layout.
-          // TODO(willfroom): Look into how we can remove this constraint.
-          mlir::MemRefType default_buffer_type =
-              mlir::MemRefType::Builder(buffer.getType()).setLayout(nullptr);
-          auto default_buffer =
-              mlir::memref::AllocOp::create(then_builder, default_buffer_type);
-          mlir::memref::CopyOp::create(then_builder, buffer, default_buffer);
-          auto to_tensor_op = mlir::bufferization::ToTensorOp::create(
-              then_builder, getType(), default_buffer);
-          to_tensor_op.setWritable(true);
-          to_tensor_op.setRestrict(true);
-          mlir::scf::YieldOp::create(then_builder, {to_tensor_op});
-        }
+        mlir::scf::YieldOp::create(then_builder,
+                                   {emit_full_tile(then_builder)});
       },
       [&](mlir::OpBuilder& builder, mlir::Location loc) {
         mlir::ImplicitLocOpBuilder else_builder(loc, builder);
@@ -335,16 +369,26 @@ llvm::LogicalResult InsertTileOp::bufferize(
     mlir::bufferization::BufferizationState& state) {
   mlir::ImplicitLocOpBuilder builder(getLoc(), rewriter);
 
-  mlir::Value is_full_size = TileIsFullSize(builder, *this);
+  auto emit_full_tile = [&](mlir::ImplicitLocOpBuilder& b) {
+    auto target_buffer_subview = GetFullTileSubView(b, *this);
+    auto materialize_op =
+        mlir::bufferization::MaterializeInDestinationOp::create(
+            b, getSource(), target_buffer_subview);
+    materialize_op.setWritable(true);
+  };
+
+  std::optional<mlir::Value> is_full_size = TileIsFullSize(builder, *this);
+  if (!is_full_size.has_value()) {
+    emit_full_tile(builder);
+    rewriter.eraseOp(getOperation());
+    return mlir::success();
+  }
+
   mlir::scf::IfOp::create(
-      builder, is_full_size,
+      builder, *is_full_size,
       [&](mlir::OpBuilder& builder, mlir::Location loc) {
         mlir::ImplicitLocOpBuilder then_builder(loc, builder);
-        auto target_buffer_subview = GetFullTileSubView(then_builder, *this);
-        auto materialize_op =
-            mlir::bufferization::MaterializeInDestinationOp::create(
-                then_builder, getSource(), target_buffer_subview);
-        materialize_op.setWritable(true);
+        emit_full_tile(then_builder);
         mlir::scf::YieldOp::create(then_builder);
       },
       [&](mlir::OpBuilder& builder, mlir::Location loc) {

From a53a446fc94ce337260a1ba44d8d8f4b24eca088 Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Mon, 12 Oct 2026 15:05:35 -0700
Subject: [PATCH 0596/1000] Seed the high-rank softmax test

Set NumPy's seed before drawing the test's random features, so the
inputs are the same on every run.
---
 tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py b/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py
index 1eb898cf4f7..4f6223ea635 100644
--- a/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py
@@ -279,6 +279,7 @@ class SoftmaxTest(test.TestCase):
   def testHighRank(self):
     # Regression test for GitHub issue 128814: the oneDNN kernel rejected
     # inputs of more than 5 dimensions.
+    np.random.seed(0)
     for rank in (6, 7):
       features = np.random.rand(*([2] * (rank - 1) + [3])).astype(np.float32)
       self._testSoftmax(features)

From 2e27301941712c7ccb99de7575220a0334c9742b Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Mon, 12 Oct 2026 15:11:43 -0700
Subject: [PATCH 0597/1000] Convert NMS thresholds unconditionally

ConvertElementType to the same type is a no-op, so drop the type
checks around the threshold conversions.
---
 tensorflow/compiler/tf2xla/kernels/image_ops.cc | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/tensorflow/compiler/tf2xla/kernels/image_ops.cc b/tensorflow/compiler/tf2xla/kernels/image_ops.cc
index fa78850785d..21aac70d385 100644
--- a/tensorflow/compiler/tf2xla/kernels/image_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/image_ops.cc
@@ -526,14 +526,10 @@ class NonMaxSuppressionOp : public XlaOpKernel {
     // The thresholds' type, T_threshold, can differ from T, the type of boxes
     // and scores. Compare in T, as the TensorFlow kernels do, since XLA
     // doesn't mix floating-point types in one operation.
-    xla::XlaOp score_thresh = context->Input("score_threshold");
-    if (context->InputXlaType("score_threshold") != scores_xla_type) {
-      score_thresh = xla::ConvertElementType(score_thresh, scores_xla_type);
-    }
-    xla::XlaOp iou_thresh = context->Input("iou_threshold");
-    if (context->InputXlaType("iou_threshold") != boxes_xla_type) {
-      iou_thresh = xla::ConvertElementType(iou_thresh, boxes_xla_type);
-    }
+    const xla::XlaOp score_thresh = xla::ConvertElementType(
+        context->Input("score_threshold"), scores_xla_type);
+    const xla::XlaOp iou_thresh = xla::ConvertElementType(
+        context->Input("iou_threshold"), boxes_xla_type);
 
     // Choose a more convenient layout.
     const xla::XlaOp boxes = xla::Transpose(boxes_input, {1, 0});

From 0f1fa24e67c445412fb176194b2aca6fc3584ab8 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 12 Oct 2026 09:52:32 -0700
Subject: [PATCH 0598/1000] Honor `ConverterFlags.enable_x64` in the
 StableHLO-to-TFLite pipeline.

The TF and SavedModel pipelines (`tf_tfl_passes.cc`) run `DowncastX64Pass`
when `enable_x64` is false, but the StableHLO bytecode pipeline
(`stablehlo_tfl_pipeline.cc`, used by `convert_mlir_bytecode`) never did,
so the flag was silently ignored there. Add the pass at the equivalent
point: after `CleanupOptimizationBarrierPass`, before VHLO legalization.

Also make `DowncastX64Pass` saturate 64-bit integer constants instead of
truncating. The pass narrowed constants through `tfl.cast`, whose folder
wraps, so INT64_MAX became -1 and INT64_MIN became 0. PyTorch exports use
those as sentinels (e.g. the fill in
`torch.where(mask, x, iinfo(long).max).min()` or
`masked_fill(~m, iinfo(long).min).argmax()`), and the wrapped value
silently changes the result. Saturating keeps the sentinel meaning; the
pass warns on each constant it clamps.

Integer constants are now narrowed in place (value attribute and result
type) rather than through `tfl.cast`, for both `DenseIntElementsAttr` and
the `DenseResourceElementsAttr` payloads that `lightweight_conversion`
produces for large tensors (a new blob is allocated; the source blob is not
modified). Signed and signless constants clamp to [INT32_MIN, INT32_MAX];
unsigned constants clamp to [0, UINT32_MAX] with unsigned comparisons and
narrow to ui32. Previously ui64 was mapped to signless i32 and routed
through `tfl.cast`, which does not accept a ui64 input. f64 constants still
go through `tfl.cast` as before.

The pass also now narrows block arguments of nested regions. `tfl.while` is
region-based until `WhileOutlinePass`, which runs after the downcast (or
not at all in the StableHLO pipeline), and `WhileOp::verify` only checks
shape compatibility, so an i64 loop counter previously exported with i32
operands feeding i64 block arguments.

PiperOrigin-RevId: 997981709
---
 tensorflow/compiler/mlir/lite/BUILD           |   1 +
 tensorflow/compiler/mlir/lite/python/BUILD    |   1 +
 .../lite/python/stablehlo_tfl_pipeline.cc     |  12 +-
 .../mlir/lite/tests/downcast-x64.mlir         | 134 +++++++++++++-
 .../mlir/lite/transforms/downcast_x64_pass.cc | 174 +++++++++++++++++-
 5 files changed, 315 insertions(+), 7 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/BUILD b/tensorflow/compiler/mlir/lite/BUILD
index 8bf52eb925e..7fa007f74e2 100644
--- a/tensorflow/compiler/mlir/lite/BUILD
+++ b/tensorflow/compiler/mlir/lite/BUILD
@@ -1100,6 +1100,7 @@ cc_library(
         ":pass",
         ":pass_options",
         ":tensorflow_lite_ops",
+        "@com_google_absl//absl/strings",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
diff --git a/tensorflow/compiler/mlir/lite/python/BUILD b/tensorflow/compiler/mlir/lite/python/BUILD
index fbc6c872205..b6547760848 100644
--- a/tensorflow/compiler/mlir/lite/python/BUILD
+++ b/tensorflow/compiler/mlir/lite/python/BUILD
@@ -117,6 +117,7 @@ cc_library(
         ":conversion_failure_reporter",
         "//tensorflow/compiler/mlir/lite:common",
         "//tensorflow/compiler/mlir/lite:converter_flags_proto_cc",
+        "//tensorflow/compiler/mlir/lite:downcast_x64",
         "//tensorflow/compiler/mlir/lite:flatbuffer_export",
         "//tensorflow/compiler/mlir/lite:optimize_broadcast_like_pass_options",
         "//tensorflow/compiler/mlir/lite:tensorflow_lite",
diff --git a/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc b/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc
index adfa5fc8270..fcce229bcda 100644
--- a/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc
+++ b/tensorflow/compiler/mlir/lite/python/stablehlo_tfl_pipeline.cc
@@ -59,6 +59,7 @@ limitations under the License.
 #include "tensorflow/compiler/mlir/lite/quantization/ir/QuantOps.h"
 #include "tensorflow/compiler/mlir/lite/stablehlo/transforms/stablehlo_passes.h"
 #include "tensorflow/compiler/mlir/lite/transforms/cast_bf16_ops_to_f32_pass.h"
+#include "tensorflow/compiler/mlir/lite/transforms/downcast_x64_pass.h"
 #include "tensorflow/compiler/mlir/lite/transforms/large_constant_fold_pass.h"
 #include "tensorflow/compiler/mlir/lite/transforms/optimize_broadcast_like_pass.h"
 #include "tensorflow/compiler/mlir/lite/transforms/optimize_broadcast_like_pass_options.h"
@@ -140,7 +141,8 @@ std::unique_ptr<mlir::Pass> CreatePruneDeadResourcesPass() {
 }  // namespace
 
 void AddPipelinePasses(mlir::OpPassManager& pass_manager,
-                       const mlir::TFL::PassConfig& pass_config) {
+                       const mlir::TFL::PassConfig& pass_config,
+                       const tflite::ConverterFlags& converter_flags) {
   // =========================================================================
   // 1. Skip-to-TFLite & Pre-Lowering Passes
   // =========================================================================
@@ -311,6 +313,12 @@ void AddPipelinePasses(mlir::OpPassManager& pass_manager,
       mlir::createCanonicalizerPass());
   pass_manager.addPass(mlir::createSymbolDCEPass());
   pass_manager.addPass(mlir::TFL::CreateCleanupOptimizationBarrierPass());
+  if (!converter_flags.enable_x64()) {
+    pass_manager.addNestedPass<mlir::func::FuncOp>(
+        mlir::TFL::CreateDowncastX64Pass());
+    pass_manager.addNestedPass<mlir::func::FuncOp>(
+        mlir::createCanonicalizerPass());
+  }
   pass_manager.addPass(mlir::odml::createLegalizeStablehloToVhloPass());
   pass_manager.addPass(mlir::createReconcileUnrealizedCastsPass());
   pass_manager.addPass(CreatePruneDeadResourcesPass());
@@ -430,7 +438,7 @@ absl::Status ConvertStableHloToTFLite(
       tensorflow::InitPassManager(pm, converter_flags.debug_options());
     }
 
-    AddPipelinePasses(pm, pass_config);
+    AddPipelinePasses(pm, pass_config, converter_flags);
 
     mlir::StatusScopedDiagnosticHandler status_handler(context,
                                                        /*propagate=*/false);
diff --git a/tensorflow/compiler/mlir/lite/tests/downcast-x64.mlir b/tensorflow/compiler/mlir/lite/tests/downcast-x64.mlir
index 6476b23a326..9d8b215fd20 100644
--- a/tensorflow/compiler/mlir/lite/tests/downcast-x64.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/downcast-x64.mlir
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 // ==============================================================================
-// RUN: litert-opt %s --tfl-downcast-x64 --canonicalize | FileCheck %s
+// RUN: litert-opt %s --tfl-downcast-x64 --canonicalize -verify-diagnostics | FileCheck %s
 
 // CHECK-LABEL: testFuncSignature
 // CHECK: (%arg0: tensor<4xi32>, %arg1: tensor<2x2xf32>) -> tensor<2x2xf32>
@@ -37,10 +37,140 @@ func.func @testI64ConstantDowncast() -> tensor<i64> {
   func.return %0 : tensor<i64>
 }
 
+// INT64_MAX / INT64_MIN sentinels saturate to INT32_MAX / INT32_MIN instead of
+// wrapping to -1 / 0.
+// CHECK-LABEL: testI64ConstantSaturates
+func.func @testI64ConstantSaturates() -> tensor<4xi64> {
+  // CHECK: %[[CST:.*]] = arith.constant dense<[2147483647, -2147483648, 2147483647, -2147483648]> : tensor<4xi32>
+  // CHECK: return %[[CST]]
+  // expected-warning @+1 {{saturating them to [INT32_MIN, INT32_MAX]}}
+  %0 = arith.constant dense<[9223372036854775807, -9223372036854775808, 2147483648, -2147483649]> : tensor<4xi64>
+  func.return %0 : tensor<4xi64>
+}
+
+// CHECK-LABEL: testI64SplatConstantSaturates
+func.func @testI64SplatConstantSaturates() -> tensor<2x3xi64> {
+  // CHECK: %[[CST:.*]] = arith.constant dense<2147483647> : tensor<2x3xi32>
+  // CHECK: return %[[CST]]
+  // expected-warning @+1 {{saturating them to [INT32_MIN, INT32_MAX]}}
+  %0 = arith.constant dense<9223372036854775807> : tensor<2x3xi64>
+  func.return %0 : tensor<2x3xi64>
+}
+
+// CHECK-LABEL: testTflConstSaturates
+func.func @testTflConstSaturates() -> tensor<2xi64> {
+  // CHECK: %[[CST:.*]] = arith.constant dense<[-2147483648, 7]> : tensor<2xi32>
+  // CHECK: return %[[CST]]
+  // expected-warning @+1 {{saturating them to [INT32_MIN, INT32_MAX]}}
+  %0 = "tfl.pseudo_const"() {value = dense<[-9223372036854775808, 7]> : tensor<2xi64>} : () -> tensor<2xi64>
+  func.return %0 : tensor<2xi64>
+}
+
+// Unsigned constants clamp to [0, UINT32_MAX] with unsigned comparisons: a
+// value with the top bit set saturates to UINT32_MAX rather than being read as
+// negative and clamped to the bottom of the range. (arith.constant requires
+// signless types, so these use tfl.pseudo_const.)
+// CHECK-LABEL: testUI64ConstantSaturates
+func.func @testUI64ConstantSaturates() -> tensor<3xui64> {
+  // CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense<[4294967295, 4294967295, 7]> : tensor<3xui32>}> : () -> tensor<3xui32>
+  // CHECK: return %[[CST]]
+  // expected-warning @+1 {{saturating them to [0, UINT32_MAX]}}
+  %0 = "tfl.pseudo_const"() {value = dense<[18446744073709551615, 4294967296, 7]> : tensor<3xui64>} : () -> tensor<3xui64>
+  func.return %0 : tensor<3xui64>
+}
+
+// Unsigned values that fit in 32 bits but exceed INT32_MAX are kept.
+// CHECK-LABEL: testUI64ConstantInRangeNoWarning
+func.func @testUI64ConstantInRangeNoWarning() -> tensor<2xui64> {
+  // CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense<[4294967295, 2147483648]> : tensor<2xui32>}> : () -> tensor<2xui32>
+  // CHECK: return %[[CST]]
+  %0 = "tfl.pseudo_const"() {value = dense<[4294967295, 2147483648]> : tensor<2xui64>} : () -> tensor<2xui64>
+  func.return %0 : tensor<2xui64>
+}
+
+// ui64 function arguments narrow to ui32, preserving signedness.
+// CHECK-LABEL: testUI64FuncSignature
+// CHECK: (%arg0: tensor<4xui32>) -> tensor<4xui32>
+func.func @testUI64FuncSignature(%arg0: tensor<4xui64>) -> tensor<4xui64> {
+  // CHECK: return %arg0 : tensor<4xui32>
+  func.return %arg0 : tensor<4xui64>
+}
+
+// In-range values pass through unchanged and emit no warning.
+// CHECK-LABEL: testI64ConstantInRangeNoWarning
+func.func @testI64ConstantInRangeNoWarning() -> tensor<2xi64> {
+  // CHECK: %[[CST:.*]] = arith.constant dense<[2147483647, -2147483648]> : tensor<2xi32>
+  // CHECK: return %[[CST]]
+  %0 = arith.constant dense<[2147483647, -2147483648]> : tensor<2xi64>
+  func.return %0 : tensor<2xi64>
+}
+
 // CHECK-LABEL: testGeneralOpForceConvert
 func.func @testGeneralOpForceConvert(%arg0: tensor<2xf64>) -> tensor<2xf64> {
   // CHECK: %[[ADD:.*]] = tfl.add %arg0, %arg0 {fused_activation_function = "NONE"} : tensor<2xf32>
   // CHECK: return %[[ADD]]
   %0 = "tfl.add"(%arg0, %arg0) {fused_activation_function = "NONE"} : (tensor<2xf64>, tensor<2xf64>) -> tensor<2xf64>
   func.return %0 : tensor<2xf64>
-}
\ No newline at end of file
+}
+
+// Block arguments of nested regions (here a region-based tfl.while, which is
+// only outlined into functions after this pass runs) are narrowed along with
+// the operands feeding them.
+// CHECK-LABEL: testWhileRegionBlockArgs
+// CHECK-SAME: (%arg0: tensor<i32>) -> tensor<i32>
+func.func @testWhileRegionBlockArgs(%arg0: tensor<i64>) -> tensor<i64> {
+  // CHECK: "tfl.while"(%arg0) ({
+  // CHECK: ^bb0(%[[CARG:.*]]: tensor<i32>):
+  // CHECK: tfl.less(%[[CARG]], %{{.*}}) : (tensor<i32>, tensor<i32>) -> tensor<i1>
+  // CHECK: ^bb0(%[[BARG:.*]]: tensor<i32>):
+  // CHECK: tfl.add %[[BARG]], %{{.*}} {fused_activation_function = "NONE"} : tensor<i32>
+  // CHECK: }) : (tensor<i32>) -> tensor<i32>
+  %0 = "tfl.while"(%arg0) ({
+    ^bb0(%carg: tensor<i64>):
+      %limit = arith.constant dense<10> : tensor<i64>
+      %cond = "tfl.less"(%carg, %limit) : (tensor<i64>, tensor<i64>) -> tensor<i1>
+      "tfl.yield"(%cond) : (tensor<i1>) -> ()
+  }, {
+    ^bb0(%barg: tensor<i64>):
+      %one = arith.constant dense<1> : tensor<i64>
+      %next = "tfl.add"(%barg, %one) {fused_activation_function = "NONE"} : (tensor<i64>, tensor<i64>) -> tensor<i64>
+      "tfl.yield"(%next) : (tensor<i64>) -> ()
+  }) : (tensor<i64>) -> tensor<i64>
+  func.return %0 : tensor<i64>
+}
+
+// Resource-backed constants (what lightweight_conversion produces for large
+// tensors) are narrowed into a new blob with the same saturation rules.
+// res_i64 = [INT64_MAX, INT64_MIN, 5].
+// CHECK-LABEL: testI64ResourceConstantSaturates
+func.func @testI64ResourceConstantSaturates() -> tensor<3xi64> {
+  // CHECK: %[[CST:.*]] = arith.constant dense_resource<res_i64_x32> : tensor<3xi32>
+  // CHECK: return %[[CST]]
+  // expected-warning @+1 {{saturating them to [INT32_MIN, INT32_MAX]}}
+  %0 = arith.constant dense_resource<res_i64> : tensor<3xi64>
+  func.return %0 : tensor<3xi64>
+}
+
+// res_ui64 = [UINT64_MAX, 7].
+// CHECK-LABEL: testUI64ResourceConstantSaturates
+func.func @testUI64ResourceConstantSaturates() -> tensor<2xui64> {
+  // CHECK: %[[CST:.*]] = "tfl.pseudo_const"() <{value = dense_resource<res_ui64_x32> : tensor<2xui32>}> : () -> tensor<2xui32>
+  // CHECK: return %[[CST]]
+  // expected-warning @+1 {{saturating them to [0, UINT32_MAX]}}
+  %0 = "tfl.pseudo_const"() {value = dense_resource<res_ui64> : tensor<2xui64>} : () -> tensor<2xui64>
+  func.return %0 : tensor<2xui64>
+}
+
+// The narrowed blobs: 64-byte alignment prefix, then
+// [INT32_MAX, INT32_MIN, 5] and [UINT32_MAX, 7] as little-endian 32-bit.
+// CHECK-DAG: res_i64_x32: "0x40000000FFFFFF7F0000008005000000"
+// CHECK-DAG: res_ui64_x32: "0x40000000FFFFFFFF07000000"
+
+{-#
+  dialect_resources: {
+    builtin: {
+      res_i64: "0x08000000FFFFFFFFFFFFFF7F00000000000000800500000000000000",
+      res_ui64: "0x08000000FFFFFFFFFFFFFFFF0700000000000000"
+    }
+  }
+#-}
\ No newline at end of file
diff --git a/tensorflow/compiler/mlir/lite/transforms/downcast_x64_pass.cc b/tensorflow/compiler/mlir/lite/transforms/downcast_x64_pass.cc
index aca4aafdb9d..3bfed9b1734 100644
--- a/tensorflow/compiler/mlir/lite/transforms/downcast_x64_pass.cc
+++ b/tensorflow/compiler/mlir/lite/transforms/downcast_x64_pass.cc
@@ -15,13 +15,27 @@ limitations under the License.
 
 #include "tensorflow/compiler/mlir/lite/transforms/downcast_x64_pass.h"
 
+#include <cstdint>
+#include <cstring>
+#include <limits>
 #include <memory>
+#include <string>
+#include <utility>
 
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/STLExtras.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/AsmState.h"  // from @llvm-project
+#include "mlir/IR/Block.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
 #include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
 #include "mlir/IR/MLIRContext.h"  // from @llvm-project
 #include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
 #include "mlir/IR/Types.h"  // from @llvm-project
 #include "mlir/IR/Value.h"  // from @llvm-project
 #include "mlir/Pass/Pass.h"  // from @llvm-project
@@ -52,12 +66,152 @@ mlir::Type ConvertX64TypeToX32(mlir::Type type) {
   if (type.isF64()) {
     return mlir::Float32Type::get(type.getContext());
   }
-  if (type.isInteger(64)) {
-    return mlir::IntegerType::get(type.getContext(), 32);
+  // Preserve signedness so ui64 narrows to ui32 rather than signless i32.
+  if (auto int_type = mlir::dyn_cast<mlir::IntegerType>(type);
+      int_type && int_type.getWidth() == 64) {
+    return mlir::IntegerType::get(type.getContext(), 32,
+                                  int_type.getSignedness());
   }
   return type;
 }
 
+// Emits the saturation warning on `op`.
+void EmitSaturationWarning(mlir::Operation* op, bool is_unsigned) {
+  op->emitWarning()
+      << "enable_x64=false: 64-bit integer constant has values outside the "
+         "32-bit range; saturating them to "
+      << (is_unsigned ? "[0, UINT32_MAX]" : "[INT32_MIN, INT32_MAX]")
+      << ". Pass enable_x64=true to keep 64-bit types.";
+}
+
+// Narrows a `DenseIntElementsAttr` of 64-bit integers to 32 bits, clamping
+// out-of-range values. Sets `*clamped` if any value was clamped.
+mlir::DenseIntElementsAttr NarrowDenseI64(mlir::DenseIntElementsAttr value,
+                                          mlir::Type new_element_type,
+                                          bool* clamped) {
+  const bool is_unsigned = value.getElementType().isUnsignedInteger();
+  const APInt min_value =
+      is_unsigned ? APInt::getZero(64) : APInt::getSignedMinValue(32).sext(64);
+  const APInt max_value = is_unsigned ? APInt::getMaxValue(32).zext(64)
+                                      : APInt::getSignedMaxValue(32).sext(64);
+  auto less_than = [is_unsigned](const APInt& a, const APInt& b) {
+    return is_unsigned ? a.ult(b) : a.slt(b);
+  };
+  return mlir::cast<mlir::DenseIntElementsAttr>(
+      value.mapValues(new_element_type, [&](const APInt& v) -> APInt {
+        if (less_than(v, min_value)) {
+          *clamped = true;
+          return min_value.trunc(32);
+        }
+        if (less_than(max_value, v)) {
+          *clamped = true;
+          return max_value.trunc(32);
+        }
+        return v.trunc(32);
+      }));
+}
+
+// Narrows a `DenseResourceElementsAttr` of 64-bit integers to 32 bits by
+// allocating a new blob, clamping out-of-range values. Returns a null
+// attribute if the resource's blob is unavailable. Sets `*clamped` if any
+// value was clamped.
+mlir::DenseResourceElementsAttr NarrowResourceI64(
+    mlir::DenseResourceElementsAttr value, mlir::ShapedType new_type,
+    bool* clamped) {
+  mlir::AsmResourceBlob* blob = value.getRawHandle().getBlob();
+  if (!blob) {
+    if (auto* resource = value.getRawHandle().getResource()) {
+      blob = resource->getBlob();
+    }
+  }
+  if (!blob) return nullptr;
+
+  const bool is_unsigned = value.getElementType().isUnsignedInteger();
+  llvm::ArrayRef<char> src_bytes = blob->getData();
+  const int64_t num_elements = new_type.getNumElements();
+  if (src_bytes.size() != num_elements * sizeof(int64_t)) return nullptr;
+  mlir::AsmResourceBlob new_blob = mlir::HeapAsmResourceBlob::allocate(
+      num_elements * sizeof(int32_t), /*align=*/64, /*dataIsMutable=*/true);
+  llvm::MutableArrayRef<char> dst_bytes = new_blob.getMutableData();
+
+  // Copy element-wise through memcpy: the source blob's alignment is not
+  // guaranteed to be suitable for a direct int64_t view.
+  for (int64_t i = 0; i < num_elements; ++i) {
+    const char* src = src_bytes.data() + i * sizeof(int64_t);
+    char* dst = dst_bytes.data() + i * sizeof(int32_t);
+    if (is_unsigned) {
+      uint64_t v;
+      std::memcpy(&v, src, sizeof(v));
+      if (v > std::numeric_limits<uint32_t>::max()) {
+        *clamped = true;
+        v = std::numeric_limits<uint32_t>::max();
+      }
+      const uint32_t narrowed = static_cast<uint32_t>(v);
+      std::memcpy(dst, &narrowed, sizeof(narrowed));
+    } else {
+      int64_t v;
+      std::memcpy(&v, src, sizeof(v));
+      if (v < std::numeric_limits<int32_t>::min()) {
+        *clamped = true;
+        v = std::numeric_limits<int32_t>::min();
+      } else if (v > std::numeric_limits<int32_t>::max()) {
+        *clamped = true;
+        v = std::numeric_limits<int32_t>::max();
+      }
+      const int32_t narrowed = static_cast<int32_t>(v);
+      std::memcpy(dst, &narrowed, sizeof(narrowed));
+    }
+  }
+
+  std::string new_key =
+      absl::StrCat(value.getRawHandle().getKey().str(), "_x32");
+  return mlir::DenseResourceElementsAttr::get(new_type, new_key,
+                                              std::move(new_blob));
+}
+
+// If `op` is a 64-bit integer constant, narrows its `value` attribute and
+// result type to 32 bits in place and returns true. Values outside the 32-bit
+// range are clamped to that range first, with a warning, so the constant
+// saturates instead of wrapping. Handles both `DenseIntElementsAttr` and
+// `DenseResourceElementsAttr` (the latter is what `lightweight_conversion`
+// produces for large constants). Returns false for any other constant (e.g.
+// f64), leaving it for the cast path.
+//
+// Signed and signless constants clamp to [INT32_MIN, INT32_MAX]; unsigned
+// constants clamp to [0, UINT32_MAX] using unsigned comparisons, so a value
+// with the top bit set is treated as large rather than negative. Narrowing
+// directly (rather than through `tfl.cast`) also keeps ui64 working, since
+// `tfl.cast` does not accept a ui64 input.
+//
+// PyTorch exports routinely use INT64_MAX / INT64_MIN as sentinels (e.g. the
+// fill value of `torch.where(mask, x, iinfo(long).max).min()` or
+// `masked_fill(~m, iinfo(long).min).argmax()`). Truncating those to -1 / 0
+// would silently change the result; saturating keeps their meaning.
+bool NarrowI64Constant(mlir::Operation* op) {
+  if (op->getNumResults() != 1) return false;
+  auto value = op->getAttrOfType<mlir::ElementsAttr>("value");
+  if (!value || !value.getElementType().isInteger(64)) return false;
+  if (value.getType() != op->getResult(0).getType()) return false;
+
+  const bool is_unsigned = value.getElementType().isUnsignedInteger();
+  auto new_type =
+      mlir::cast<mlir::ShapedType>(ConvertX64TypeToX32(value.getType()));
+  bool clamped = false;
+  mlir::ElementsAttr new_value;
+  if (auto dense = mlir::dyn_cast<mlir::DenseIntElementsAttr>(value)) {
+    new_value = NarrowDenseI64(dense, new_type.getElementType(), &clamped);
+  } else if (auto resource =
+                 mlir::dyn_cast<mlir::DenseResourceElementsAttr>(value)) {
+    new_value = NarrowResourceI64(resource, new_type, &clamped);
+  }
+  if (!new_value) return false;
+
+  if (clamped) EmitSaturationWarning(op, is_unsigned);
+  op->setAttr("value", new_value);
+  op->getResult(0).setType(new_type);
+  return true;
+}
+
 }  // namespace
 
 void DowncastX64Pass::runOnOperation() {
@@ -97,8 +251,11 @@ void DowncastX64Pass::runOnOperation() {
       return;
     }
 
-    // Handle constant-like operations by inserting CastOps.
+    // Handle constant-like operations. 64-bit integer constants (dense or
+    // resource-backed) are narrowed in place; anything else gets a CastOp for
+    // the constant folder to resolve.
     if (op->hasTrait<mlir::OpTrait::ConstantLike>()) {
+      if (NarrowI64Constant(op)) return;
       for (mlir::OpResult result : op->getResults()) {
         mlir::Type old_type = result.getType();
         mlir::Type new_type = ConvertX64TypeToX32(old_type);
@@ -120,6 +277,17 @@ void DowncastX64Pass::runOnOperation() {
     for (mlir::Value result : op->getResults()) {
       result.setType(ConvertX64TypeToX32(result.getType()));
     }
+
+    // Update block arguments of any nested regions (e.g. the cond/body of a
+    // region-based `tfl.while`, or a residual `stablehlo.while`/`reduce`), so
+    // they stay consistent with the narrowed operands feeding them.
+    for (mlir::Region& region : op->getRegions()) {
+      for (mlir::Block& block : region) {
+        for (mlir::BlockArgument arg : block.getArguments()) {
+          arg.setType(ConvertX64TypeToX32(arg.getType()));
+        }
+      }
+    }
   });
 }
 

From c5475c804621229ecd2c8026259d636961a18bca Mon Sep 17 00:00:00 2001
From: Dmitri Latushko <latushko@google.com>
Date: Mon, 12 Oct 2026 09:53:34 -0700
Subject: [PATCH 0599/1000] Increase shard_count from 8 to 16 for shard_test to
 prevent TSAN timeouts

PiperOrigin-RevId: 997982551
---
 tensorflow/python/data/kernel_tests/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/data/kernel_tests/BUILD b/tensorflow/python/data/kernel_tests/BUILD
index ffc0aea0514..edb0ee1799e 100644
--- a/tensorflow/python/data/kernel_tests/BUILD
+++ b/tensorflow/python/data/kernel_tests/BUILD
@@ -1149,7 +1149,7 @@ tf_py_strict_test(
     name = "shard_test",
     size = "medium",
     srcs = ["shard_test.py"],
-    shard_count = 8,
+    shard_count = 16,
     deps = [
         ":checkpoint_test_base",
         ":test_base",

From c3f2f97fe8657ec91744e4305eb2392d46b34d72 Mon Sep 17 00:00:00 2001
From: Gerardo Carranza <gcarranza@google.com>
Date: Mon, 12 Oct 2026 10:17:38 -0700
Subject: [PATCH 0600/1000] Fix f16 type mismatches in YNNPACK SDPA lowering.

PiperOrigin-RevId: 997998563
---
 tensorflow/lite/delegates/ynnpack/attention.cc | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/tensorflow/lite/delegates/ynnpack/attention.cc b/tensorflow/lite/delegates/ynnpack/attention.cc
index 181392b49b1..28e4d605ecd 100644
--- a/tensorflow/lite/delegates/ynnpack/attention.cc
+++ b/tensorflow/lite/delegates/ynnpack/attention.cc
@@ -189,6 +189,7 @@ TfLiteStatus DefineSdpaNode(TfLiteContext* context, ynn_subgraph_t subgraph,
 
   const TfLiteTensor& q_tensor = context->tensors[sdpa_inputs.q_index];
   const TfLiteTensor& k_tensor = context->tensors[sdpa_inputs.k_index];
+  const TfLiteTensor& v_tensor = context->tensors[sdpa_inputs.v_index];
   const TfLiteTensor& output_tensor = context->tensors[node.outputs[0]];
 
   uint32_t q_val_id = GetOrCreateValueId(context, subgraph, tensor_to_value_id,
@@ -333,9 +334,9 @@ TfLiteStatus DefineSdpaNode(TfLiteContext* context, ynn_subgraph_t subgraph,
   }
 
   uint32_t scale_const_id = YNN_INVALID_VALUE_ID;
-  TF_LITE_ENSURE_YNN_STATUS(
-      ynn_define_tensor(subgraph, ynn_type_fp32, 0, nullptr, &scale_val,
-                        YNN_VALUE_FLAG_COPY_DATA_FP32, &scale_const_id));
+  TF_LITE_ENSURE_YNN_STATUS(ynn_define_tensor(
+      subgraph, GetYnnType(q_tensor.type), 0, nullptr, &scale_val,
+      YNN_VALUE_FLAG_COPY_DATA_FP32, &scale_const_id));
 
   const int q_head_dim = is_seq_major ? 2 : 1;
   const int k_head_dim = is_seq_major ? 2 : 1;
@@ -492,6 +493,8 @@ TfLiteStatus DefineSdpaNode(TfLiteContext* context, ynn_subgraph_t subgraph,
   }
 
   uint32_t probs_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_YNN_STATUS(ynn_define_tensor(
+      subgraph, GetYnnType(v_tensor.type), 0, nullptr, nullptr, 0, &probs_id));
   TF_LITE_ENSURE_YNN_STATUS(
       ynn::define_softmax(subgraph, masked_logits_id, 1.0f, probs_id));
 
@@ -530,6 +533,9 @@ TfLiteStatus DefineSdpaNode(TfLiteContext* context, ynn_subgraph_t subgraph,
                                     &probs_t_id, YNN_NODE_FLAG_KEEP_DIMS));
 
     uint32_t post_bmm_t_id = YNN_INVALID_VALUE_ID;
+    TF_LITE_ENSURE_YNN_STATUS(
+        ynn_define_tensor(subgraph, GetYnnType(output_tensor.type), 0, nullptr,
+                          nullptr, 0, &post_bmm_t_id));
     TF_LITE_ENSURE_YNN_STATUS(
         ynn_define_dot(subgraph, /*num_k_dims=*/1, current_v_val_id, probs_t_id,
                        YNN_INVALID_VALUE_ID, &post_bmm_t_id, 0));
@@ -545,6 +551,11 @@ TfLiteStatus DefineSdpaNode(TfLiteContext* context, ynn_subgraph_t subgraph,
         subgraph, 2, is_seq_major ? seq_major_v_perm : swap_last_two_perm,
         current_v_val_id, &v_trans_id, YNN_NODE_FLAG_KEEP_DIMS));
 
+    if (*post_bmm_ptr == YNN_INVALID_VALUE_ID) {
+      TF_LITE_ENSURE_YNN_STATUS(
+          ynn_define_tensor(subgraph, GetYnnType(output_tensor.type), 0,
+                            nullptr, nullptr, 0, post_bmm_ptr));
+    }
     TF_LITE_ENSURE_YNN_STATUS(
         ynn_define_dot(subgraph, /*num_k_dims=*/1, probs_id, v_trans_id,
                        YNN_INVALID_VALUE_ID, post_bmm_ptr, 0));

From a093f8d5d9f7411da5e0451b08f7769c22ae9d2f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 12 Oct 2026 10:48:36 -0700
Subject: [PATCH 0601/1000] Ensure AsyncOpKernel DoneCallback is invoked on all
 error paths in core/kernels.

Use absl::MakeCleanup (or OP_REQUIRES_OK_ASYNC) across AsyncOpKernel implementations in tensorflow/core/kernels so DoneCallback is always invoked even if a helper or error check returns early.

PiperOrigin-RevId: 998017953
---
 tensorflow/core/kernels/barrier_ops.cc        |  10 +-
 tensorflow/core/kernels/batch_kernels.cc      |  87 ++++++++-----
 tensorflow/core/kernels/check_numerics_op.cc  |  23 ++--
 .../kernels/conditional_accumulator_base.cc   |   2 +-
 .../kernels/conditional_accumulator_base_op.h |  13 +-
 tensorflow/core/kernels/debug_ops.h           |  96 ++++++++------
 .../kernels/dynamic_partition_op_gpu.cu.cc    | 119 +++++++++---------
 tensorflow/core/kernels/functional_ops.cc     |  10 +-
 tensorflow/core/kernels/queue_op.cc           |  13 +-
 .../kernels/segment_reduction_ops_gpu.cu.h    |  46 +++----
 .../core/kernels/segment_reduction_ops_impl.h |  57 +++++----
 .../core/kernels/sparse_to_dense_op_gpu.cu.cc |  89 ++++++-------
 12 files changed, 317 insertions(+), 248 deletions(-)

diff --git a/tensorflow/core/kernels/barrier_ops.cc b/tensorflow/core/kernels/barrier_ops.cc
index 830d5e3c85f..2acebeb6869 100644
--- a/tensorflow/core/kernels/barrier_ops.cc
+++ b/tensorflow/core/kernels/barrier_ops.cc
@@ -30,6 +30,7 @@ limitations under the License.
 #include "tensorflow/core/kernels/queue_base.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/notification.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
 #include "tensorflow/core/lib/gtl/map_util.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/macros.h"
@@ -254,12 +255,12 @@ class Barrier : public ResourceBase {
   void Close(OpKernelContext* ctx, bool cancel_pending_enqueues,
              const DoneCallback& callback) {
     mutex_lock lock(mu_);
+    auto cleanup = gtl::MakeCleanup([&callback]() { callback(); });
     // We're allowed to close twice if the first close wasn't a
     // cancel but the second one is.
     if (closed_ && (cancel_pending_enqueues_ || !cancel_pending_enqueues)) {
       ctx->SetStatus(absl::CancelledError(
           absl::StrCat("Barrier '", name_, "' is already closed.")));
-      callback();
       return;
     }
     cancel_pending_enqueues_ = cancel_pending_enqueues;
@@ -267,10 +268,9 @@ class Barrier : public ResourceBase {
     if (cancel_pending_enqueues_ || incomplete_.empty()) {
       incomplete_.clear();
       // CloseQueueLocked runs the callback
+      cleanup.release();
       CloseQueueLocked(ctx, cancel_pending_enqueues_, callback);
-      return;
     }
-    callback();
   }
 
   int32_t ready_size() { return ready_queue_->size(); }
@@ -403,18 +403,18 @@ class Barrier : public ResourceBase {
   void CloseQueueLocked(OpKernelContext* ctx, bool cancel_pending_enqueues,
                         const DoneCallback& callback)
       TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+    auto cleanup = gtl::MakeCleanup([&callback]() { callback(); });
     // CloseQueueLocked may only be called with mu_ held.
     if (!cancel_pending_enqueues && queue_closed_) {
-      callback();
       return;
     }
     if (cancel_pending_enqueues && queue_cancelled_) {
-      callback();
       return;
     }
     queue_closed_ = true;
     if (cancel_pending_enqueues) queue_cancelled_ = true;
     if (!ready_queue_->is_closed()) {
+      cleanup.release();
       ready_queue_->Close(ctx, cancel_pending_enqueues, callback);
     }
   }
diff --git a/tensorflow/core/kernels/batch_kernels.cc b/tensorflow/core/kernels/batch_kernels.cc
index 4af394fe2e4..e76a964f9c9 100644
--- a/tensorflow/core/kernels/batch_kernels.cc
+++ b/tensorflow/core/kernels/batch_kernels.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include <memory>
 #include <optional>
 #include <string>
+#include <unordered_set>
 #include <utility>
 
 #include "absl/base/attributes.h"
@@ -50,11 +51,13 @@ limitations under the License.
 #include "tensorflow/core/kernels/batching_util/periodic_function.h"
 #include "tensorflow/core/kernels/batching_util/warmup.h"
 #include "tensorflow/core/kernels/ops_util.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
 #include "tensorflow/core/lib/monitoring/gauge.h"
 #include "tensorflow/core/lib/random/random.h"
 #include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/numbers.h"
+#include "tensorflow/core/platform/refcount.h"
 #include "tensorflow/core/platform/status.h"
 #include "tensorflow/core/platform/threadpool.h"
 
@@ -931,15 +934,26 @@ class UnbatchResource : public ResourceBase {
       const uint64_t deadline_micros =
           Env::Default()->NowMicros() + timeout_micros_;
 
-      // Add ourselves to the waitlist for tensors.
-      if (!waiting_callbacks_
-               .emplace(batch_key,
-                        WaitingCallback{deadline_micros, context, done})
-               .second) {
+      if (waiting_callbacks_.find(batch_key) != waiting_callbacks_.end()) {
         return absl::AlreadyExistsError(
             "Multiple session runs with the same batch key.");
       }
 
+      if (nonempty_input) {
+        std::unordered_set<int64_t> seen_keys;
+        for (const int64_t key : batch_keys) {
+          if (!seen_keys.insert(key).second ||
+              waiting_tensors_.find(key) != waiting_tensors_.end()) {
+            return absl::AlreadyExistsError(
+                "Multiple tensors returned for same batch key.");
+          }
+        }
+      }
+
+      // Add ourselves to the waitlist for tensors.
+      waiting_callbacks_.emplace(
+          batch_key, WaitingCallback{deadline_micros, context, done});
+
       // If we have a non-empty tensor, finish the waitlisted runs,
       // and store any remaining pieces.
       if (nonempty_input) {
@@ -953,13 +967,8 @@ class UnbatchResource : public ResourceBase {
             // Note: the deadline here is in case we are arriving late and the
             // kernel that should rendezvous with this tensor has already waited
             // and timed out.
-            if (!waiting_tensors_
-                     .emplace(batch_keys[i],
-                              WaitingTensor{deadline_micros, split_inputs[i]})
-                     .second) {
-              return absl::AlreadyExistsError(
-                  "Multiple tensors returned for same batch key.");
-            }
+            waiting_tensors_.emplace(
+                batch_keys[i], WaitingTensor{deadline_micros, split_inputs[i]});
           }
         }
       }
@@ -1054,6 +1063,7 @@ class UnbatchKernel : public AsyncOpKernel {
   }
 
   void ComputeAsync(OpKernelContext* c, DoneCallback done) final {
+    auto cleanup = gtl::MakeCleanup([&done]() { done(); });
     UnbatchResource* ubr;
     std::function<absl::Status(UnbatchResource**)> creator =
         [this](UnbatchResource** r) {
@@ -1063,11 +1073,10 @@ class UnbatchKernel : public AsyncOpKernel {
     OP_REQUIRES_OK_ASYNC(c,
                          c->resource_manager()->LookupOrCreate(
                              container_, shared_name_, &ubr, creator),
-                         done);
-    auto status = ubr->Compute(c, done);
-    ubr->Unref();
-    OP_REQUIRES_OK_ASYNC(c, status, done);
-    // Assume ubr calls done, so nothing to do here.
+                         []() {});
+    core::ScopedUnref unref_ubr(ubr);
+    OP_REQUIRES_OK_ASYNC(c, ubr->Compute(c, done), []() {});
+    cleanup.release();
   }
 
  private:
@@ -1173,26 +1182,36 @@ class UnbatchGradResource : public ResourceBase {
       if (missing_tensors.empty()) {
         return OutputBatch(context, done);
       }
-      if (!available_batches_
-               .emplace(batch_key, Batch{missing_tensors, context, done})
-               .second) {
+      if (available_batches_.find(batch_key) != available_batches_.end()) {
         return absl::InvalidArgumentError(
             "Batch key with valid batch used twice.");
       }
       for (const int64_t i : missing_tensors) {
-        if (!desired_tensor_to_batch_map_.emplace(i, batch_key).second) {
+        if (desired_tensor_to_batch_map_.find(i) !=
+            desired_tensor_to_batch_map_.end()) {
           return absl::InvalidArgumentError(
               "Missing tensor wanted by more than one batch.");
         }
       }
+      auto desire_it = desired_tensor_to_batch_map_.find(batch_key);
+      if (desire_it != desired_tensor_to_batch_map_.end() &&
+          available_batches_.find(desire_it->second) ==
+              available_batches_.end()) {
+        desired_tensor_to_batch_map_.erase(desire_it);
+        return absl::InvalidArgumentError("Batch no longer exists.");
+      }
+      available_batches_.emplace(batch_key,
+                                 Batch{missing_tensors, context, done});
+      for (const int64_t i : missing_tensors) {
+        desired_tensor_to_batch_map_.emplace(i, batch_key);
+      }
     } else {
       // If we don't have a valid input tensor we can output an empty tensor and
-      // call our done closure.
+      // call our done closure once all bookkeeping succeeds.
       TensorShape output_shape(grad_t.shape());
       output_shape.set_dim(0, 0);
       Tensor* output = nullptr;
       TF_RETURN_IF_ERROR(context->allocate_output(0, output_shape, &output));
-      done();
     }
 
     // Search to see whether our tensor is desired by any existing batch.
@@ -1208,11 +1227,19 @@ class UnbatchGradResource : public ResourceBase {
       // If all tensors are available we should concatenate them and dispatch
       // the batch.
       if (batch_it->second.missing_tensors.empty()) {
-        TF_RETURN_IF_ERROR(
-            OutputBatch(batch_it->second.context, batch_it->second.done));
+        OpKernelContext* batch_context = batch_it->second.context;
+        AsyncOpKernel::DoneCallback batch_done = batch_it->second.done;
         available_batches_.erase(batch_it);
+        absl::Status batch_status = OutputBatch(batch_context, batch_done);
+        if (!batch_status.ok()) {
+          batch_context->CtxFailureWithWarning(batch_status);
+          batch_done();
+        }
       }
     }
+    if (data_t.NumElements() == 0) {
+      done();
+    }
     return absl::OkStatus();
   }
 
@@ -1259,6 +1286,7 @@ class UnbatchGradKernel : public AsyncOpKernel {
   }
 
   void ComputeAsync(OpKernelContext* c, DoneCallback done) final {
+    auto cleanup = gtl::MakeCleanup([&done]() { done(); });
     UnbatchGradResource* ubr;
     std::function<absl::Status(UnbatchGradResource**)> creator =
         [](UnbatchGradResource** r) {
@@ -1268,11 +1296,10 @@ class UnbatchGradKernel : public AsyncOpKernel {
     OP_REQUIRES_OK_ASYNC(c,
                          c->resource_manager()->LookupOrCreate(
                              container_, shared_name_, &ubr, creator),
-                         done);
-    absl::Status status = ubr->Compute(c, done);
-    ubr->Unref();
-    OP_REQUIRES_OK_ASYNC(c, status, done);
-    // Assume ubr calls done, so nothing to do here.
+                         []() {});
+    core::ScopedUnref unref_ubr(ubr);
+    OP_REQUIRES_OK_ASYNC(c, ubr->Compute(c, done), []() {});
+    cleanup.release();
   }
 
  private:
diff --git a/tensorflow/core/kernels/check_numerics_op.cc b/tensorflow/core/kernels/check_numerics_op.cc
index 5535fe04858..7bc207d6ffa 100644
--- a/tensorflow/core/kernels/check_numerics_op.cc
+++ b/tensorflow/core/kernels/check_numerics_op.cc
@@ -31,6 +31,7 @@ limitations under the License.
 
 #if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 #include "tensorflow/core/common_runtime/gpu/gpu_event_mgr.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
 #endif  // GOOGLE_CUDA || TENSORFLOW_USE_ROCM
 
 #if TENSORFLOW_USE_ROCM
@@ -202,10 +203,10 @@ class CheckNumericsOp<GPUDevice, T> : public AsyncOpKernel {
   }
 
   void ComputeAsync(OpKernelContext* context, DoneCallback done) override {
+    auto cleanup = gtl::MakeCleanup([&done]() { done(); });
     // pass along the input to the output
     context->set_output(0, context->input(0));
     if (context->input(0).NumElements() == 0) {
-      done();
       return;
     }
     auto input = context->input(0).flat<T>();
@@ -213,21 +214,24 @@ class CheckNumericsOp<GPUDevice, T> : public AsyncOpKernel {
     // Allocate and initialize the elements to hold the check results
     Tensor abnormal_detected;
     const int abnormal_detected_size = getAnomalyIndicatorSize();
-    OP_REQUIRES_OK(context, context->allocate_temp(
-                                DT_INT32, TensorShape({abnormal_detected_size}),
-                                &abnormal_detected));
+    OP_REQUIRES_OK_ASYNC(
+        context,
+        context->allocate_temp(DT_INT32, TensorShape({abnormal_detected_size}),
+                               &abnormal_detected),
+        []() {});
 
     auto* stream = context->op_device_context()->stream();
     OP_REQUIRES_ASYNC(context, stream != nullptr,
-                      absl::InternalError("No GPU stream available."), done);
+                      absl::InternalError("No GPU stream available."), []() {});
 
     stream_executor::DeviceAddressBase abnormal_detected_ptr(
         abnormal_detected.flat<int>().data(),
         abnormal_detected.flat<int>().size());
-    OP_REQUIRES_OK(
+    OP_REQUIRES_OK_ASYNC(
         context,
         stream->Memset32(&abnormal_detected_ptr, 0,
-                         abnormal_detected.flat<int>().size() * sizeof(int)));
+                         abnormal_detected.flat<int>().size() * sizeof(int)),
+        []() {});
 
     // Call the GPU kernels for the numerical checks
     const Device& d = context->eigen_device<Device>();
@@ -243,7 +247,7 @@ class CheckNumericsOp<GPUDevice, T> : public AsyncOpKernel {
         context,
         context->allocate_temp(DT_INT32, TensorShape({abnormal_detected_size}),
                                &abnormal_detected_host, attr),
-        done);
+        []() {});
     OP_REQUIRES_ASYNC(
         context,
         stream
@@ -251,7 +255,7 @@ class CheckNumericsOp<GPUDevice, T> : public AsyncOpKernel {
                      abnormal_detected_ptr,
                      abnormal_detected_size * sizeof(int))
             .ok(),
-        absl::InternalError("GPU memcpy from device to host failed"), done);
+        absl::InternalError("GPU memcpy from device to host failed"), []() {});
 
     // We have observed crashes on some network stacks when not holding
     // this tensor reference.
@@ -270,6 +274,7 @@ class CheckNumericsOp<GPUDevice, T> : public AsyncOpKernel {
          // Context.
       done();
     };
+    cleanup.release();
     context->device()
         ->tensorflow_accelerator_device_info()
         ->event_mgr->ThenExecute(stream, std::move(check_cb));
diff --git a/tensorflow/core/kernels/conditional_accumulator_base.cc b/tensorflow/core/kernels/conditional_accumulator_base.cc
index 12d74b0b4be..ecc893db9f6 100644
--- a/tensorflow/core/kernels/conditional_accumulator_base.cc
+++ b/tensorflow/core/kernels/conditional_accumulator_base.cc
@@ -85,7 +85,7 @@ void ConditionalAccumulatorBase::TryTakeGrad(int num_required,
               if (counter_ >= attempt->elements_requested) {
                 bool successful_take_grad = TakeGradLockedHelper(
                     attempt->context, attempt->done_callback);
-                if (successful_take_grad) {
+                if (successful_take_grad || !attempt->context->status().ok()) {
                   return kComplete;
                 } else {
                   // Try again
diff --git a/tensorflow/core/kernels/conditional_accumulator_base_op.h b/tensorflow/core/kernels/conditional_accumulator_base_op.h
index e4521e93464..d6fc78e74e7 100644
--- a/tensorflow/core/kernels/conditional_accumulator_base_op.h
+++ b/tensorflow/core/kernels/conditional_accumulator_base_op.h
@@ -18,15 +18,15 @@ limitations under the License.
 
 #define EIGEN_USE_THREADS
 
-#include "tensorflow/core/kernels/conditional_accumulator_base.h"
-
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/resource_mgr.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/kernels/conditional_accumulator_base.h"
 #include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/mutex.h"
 #include "tensorflow/core/platform/thread_annotations.h"
@@ -237,21 +237,24 @@ class ConditionalAccumulatorBaseTakeGradientOp
   void ComputeAsync(OpKernelContext* ctx,
                     ConditionalAccumulatorBase* accumulator,
                     DoneCallback callback) override {
+    auto cleanup = gtl::MakeCleanup([&callback]() { callback(); });
     // Check signature
-    CheckSignature(ctx, accumulator, callback);
+    CheckSignature(ctx, accumulator, []() {});
+    if (!ctx->status().ok()) return;
 
     // Get input num_required
     const Tensor* num_required_tensor;
     OP_REQUIRES_OK_ASYNC(ctx, ctx->input("num_required", &num_required_tensor),
-                         callback);
+                         []() {});
     if (!TensorShapeUtils::IsScalar(num_required_tensor->shape())) {
       ctx->CtxFailureWithWarning(absl::InvalidArgumentError(absl::StrCat(
           "Argument num_required must be scalar, but had bad shape ",
           num_required_tensor->shape().DebugString())));
-      callback();
+      return;
     }
 
     // Actually try to take gradient now
+    cleanup.release();
     accumulator->TryTakeGrad(num_required_tensor->scalar<int32_t>()(), ctx,
                              callback);
   }
diff --git a/tensorflow/core/kernels/debug_ops.h b/tensorflow/core/kernels/debug_ops.h
index 99e30fe8bf1..ce4ecd05270 100644
--- a/tensorflow/core/kernels/debug_ops.h
+++ b/tensorflow/core/kernels/debug_ops.h
@@ -39,6 +39,7 @@ limitations under the License.
 #include "tensorflow/core/framework/device_base.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/framework/tensor_util.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
 #include "tensorflow/core/lib/strings/stringprintf.h"
 #include "tensorflow/core/util/debug_events_writer.h"
 
@@ -747,6 +748,7 @@ class DebugNumericSummaryV2Op<GPUDevice, Tin, Tout> : public AsyncOpKernel {
   }
 
   void ComputeAsync(OpKernelContext* context, DoneCallback done) override {
+    auto cleanup = gtl::MakeCleanup([&done]() { done(); });
     Tensor* output_tensor;
     Tout tensor_id = static_cast<Tout>(tensor_id_);
     const Tensor& tensor = context->input(0);
@@ -764,28 +766,30 @@ class DebugNumericSummaryV2Op<GPUDevice, Tin, Tout> : public AsyncOpKernel {
                                   "(2^",
                                   std::numeric_limits<Tout>::digits,
                                   "). Given tensor_id:", tensor_id_),
-          done);
+          []() {});
     }
 
     if (tensor_debug_mode_ == 2) {  // CURT_HEALTH.
       TensorShape shape({2});
-      OP_REQUIRES_OK(context,
-                     context->allocate_output(0, shape, &output_tensor));
+      OP_REQUIRES_OK_ASYNC(
+          context, context->allocate_output(0, shape, &output_tensor), []() {});
 
       auto* stream = context->op_device_context()->stream();
       OP_REQUIRES_ASYNC(context, stream != nullptr,
-                        absl::InternalError("No GPU stream available."), done);
+                        absl::InternalError("No GPU stream available."),
+                        []() {});
 
       stream_executor::DeviceAddressBase output_tensor_ptr(
           output_tensor->flat<Tout>().data(),
           output_tensor->flat<Tout>().size());
-      OP_REQUIRES_OK(context,
-                     stream->MemZero(&output_tensor_ptr, 2 * sizeof(Tout)));
+      OP_REQUIRES_OK_ASYNC(
+          context, stream->MemZero(&output_tensor_ptr, 2 * sizeof(Tout)),
+          []() {});
       // Copy tensor_id to slot zero
-      OP_REQUIRES_OK(context, stream->Memcpy(&output_tensor_ptr, &tensor_id,
-                                             sizeof(Tout)));
+      OP_REQUIRES_OK_ASYNC(
+          context, stream->Memcpy(&output_tensor_ptr, &tensor_id, sizeof(Tout)),
+          []() {});
       if (num_elem == 0) {
-        done();
         return;
       }
 
@@ -794,34 +798,38 @@ class DebugNumericSummaryV2Op<GPUDevice, Tin, Tout> : public AsyncOpKernel {
       CurtHealthLaunch<Tin, Tout>().Run(d, input.data(), input.size(),
                                         output_tensor->flat<Tout>().data() + 1);
 
+      cleanup.release();
       context->device()
           ->tensorflow_accelerator_device_info()
           ->event_mgr->ThenExecute(stream, std::move(check_cb));
     } else if (tensor_debug_mode_ == 3) {  // CONCISE_HEALTH.
       TensorShape shape({5});
-      OP_REQUIRES_OK(context,
-                     context->allocate_output(0, shape, &output_tensor));
+      OP_REQUIRES_OK_ASYNC(
+          context, context->allocate_output(0, shape, &output_tensor), []() {});
       OP_REQUIRES_ASYNC(context, !tensorflow::OpDeterminismRequired(),
                         absl::UnimplementedError(
                             "Determinism is not yet supported for "
                             "DebugNumericSummaryV2 when tensor_debug_mode is "
                             "CONCISE_HEALTH."),
-                        done);
+                        []() {});
 
       auto* stream = context->op_device_context()->stream();
       OP_REQUIRES_ASYNC(context, stream != nullptr,
-                        absl::InternalError("No GPU stream available."), done);
+                        absl::InternalError("No GPU stream available."),
+                        []() {});
 
       stream_executor::DeviceAddressBase output_tensor_ptr(
           output_tensor->flat<Tout>().data(),
           output_tensor->flat<Tout>().size());
-      OP_REQUIRES_OK(context,
-                     stream->Memset32(&output_tensor_ptr, 0, 5 * sizeof(Tout)));
+      OP_REQUIRES_OK_ASYNC(
+          context, stream->Memset32(&output_tensor_ptr, 0, 5 * sizeof(Tout)),
+          []() {});
       const Tout static_output[] = {tensor_id, num_elem};
-      OP_REQUIRES_OK(context, stream->Memcpy(&output_tensor_ptr, &static_output,
-                                             2 * sizeof(Tout)));
+      OP_REQUIRES_OK_ASYNC(
+          context,
+          stream->Memcpy(&output_tensor_ptr, &static_output, 2 * sizeof(Tout)),
+          []() {});
       if (num_elem == 0) {
-        done();
         return;
       }
 
@@ -830,39 +838,43 @@ class DebugNumericSummaryV2Op<GPUDevice, Tin, Tout> : public AsyncOpKernel {
           d, input.data(), input.size(),
           output_tensor->flat<Tout>().data() + 2);
 
+      cleanup.release();
       context->device()
           ->tensorflow_accelerator_device_info()
           ->event_mgr->ThenExecute(stream, std::move(check_cb));
     } else if (tensor_debug_mode_ == 4) {  // FULL HEALTH
       TensorShape shape({11});
-      OP_REQUIRES_OK(context,
-                     context->allocate_output(0, shape, &output_tensor));
+      OP_REQUIRES_OK_ASYNC(
+          context, context->allocate_output(0, shape, &output_tensor), []() {});
 
       auto* stream = context->op_device_context()->stream();
       OP_REQUIRES_ASYNC(context, stream != nullptr,
-                        absl::InternalError("No GPU stream available."), done);
+                        absl::InternalError("No GPU stream available."),
+                        []() {});
       OP_REQUIRES_ASYNC(context, !tensorflow::OpDeterminismRequired(),
                         absl::UnimplementedError(
                             "Determinism is not yet supported for "
                             "DebugNumericSummaryV2 when tensor_debug_mode is "
                             "FULL_HEALTH."),
-                        done);
+                        []() {});
 
       stream_executor::DeviceAddressBase output_tensor_ptr(
           output_tensor->flat<Tout>().data(),
           output_tensor->flat<Tout>().size());
-      OP_REQUIRES_OK(
-          context, stream->Memset32(&output_tensor_ptr, 0, 11 * sizeof(Tout)));
+      OP_REQUIRES_OK_ASYNC(
+          context, stream->Memset32(&output_tensor_ptr, 0, 11 * sizeof(Tout)),
+          []() {});
 
       int num_dims = tensor.dims();
       const Tout static_output[] = {tensor_id,
                                     -1.0,  // TODO(144919262): Device ID
                                     static_cast<Tout>(tensor.dtype()),
                                     static_cast<Tout>(num_dims), num_elem};
-      OP_REQUIRES_OK(context, stream->Memcpy(&output_tensor_ptr, &static_output,
-                                             5 * sizeof(Tout)));
+      OP_REQUIRES_OK_ASYNC(
+          context,
+          stream->Memcpy(&output_tensor_ptr, &static_output, 5 * sizeof(Tout)),
+          []() {});
       if (num_elem == 0) {
-        done();
         return;
       }
 
@@ -871,17 +883,19 @@ class DebugNumericSummaryV2Op<GPUDevice, Tin, Tout> : public AsyncOpKernel {
       FullHealthLaunch<Tin, Tout>().Run(d, input.data(), input.size(),
                                         output_tensor->flat<Tout>().data() + 5);
 
+      cleanup.release();
       context->device()
           ->tensorflow_accelerator_device_info()
           ->event_mgr->ThenExecute(stream, std::move(check_cb));
     } else if (tensor_debug_mode_ == 5) {  // SHAPE
       TensorShape shape({10});
-      OP_REQUIRES_OK(context,
-                     context->allocate_output(0, shape, &output_tensor));
+      OP_REQUIRES_OK_ASYNC(
+          context, context->allocate_output(0, shape, &output_tensor), []() {});
 
       auto* stream = context->op_device_context()->stream();
       OP_REQUIRES_ASYNC(context, stream != nullptr,
-                        absl::InternalError("No GPU stream available."), done);
+                        absl::InternalError("No GPU stream available."),
+                        []() {});
 
       stream_executor::DeviceAddressBase output_tensor_ptr(
           output_tensor->flat<Tout>().data(),
@@ -904,29 +918,33 @@ class DebugNumericSummaryV2Op<GPUDevice, Tin, Tout> : public AsyncOpKernel {
         static_output[dim_idx++] = static_cast<Tout>(tensor.dim_size(i));
       }
       // Write to device stream
-      OP_REQUIRES_OK(context, stream->Memcpy(&output_tensor_ptr, &static_output,
-                                             sizeof(Tout) * 10));
+      OP_REQUIRES_OK_ASYNC(
+          context,
+          stream->Memcpy(&output_tensor_ptr, &static_output, sizeof(Tout) * 10),
+          []() {});
+      cleanup.release();
       context->device()
           ->tensorflow_accelerator_device_info()
           ->event_mgr->ThenExecute(stream, std::move(check_cb));
     } else if (tensor_debug_mode_ == 8) {  // REDUCE_INF_NAN_THREE_SLOTS.
       TensorShape shape({3});
-      OP_REQUIRES_OK(context,
-                     context->allocate_output(0, shape, &output_tensor));
+      OP_REQUIRES_OK_ASYNC(
+          context, context->allocate_output(0, shape, &output_tensor), []() {});
 
       auto* stream = context->op_device_context()->stream();
       OP_REQUIRES_ASYNC(context, stream != nullptr,
-                        absl::InternalError("No GPU stream available."), done);
+                        absl::InternalError("No GPU stream available."),
+                        []() {});
 
       stream_executor::DeviceAddressBase output_tensor_ptr(
           output_tensor->flat<Tout>().data(),
           output_tensor->flat<Tout>().size());
-      OP_REQUIRES_OK(
+      OP_REQUIRES_OK_ASYNC(
           context,
           stream->Memset32(&output_tensor_ptr, 0,
-                           output_tensor->flat<Tout>().size() * sizeof(Tout)));
+                           output_tensor->flat<Tout>().size() * sizeof(Tout)),
+          []() {});
       if (num_elem == 0) {
-        done();
         return;
       }
 
@@ -935,6 +953,7 @@ class DebugNumericSummaryV2Op<GPUDevice, Tin, Tout> : public AsyncOpKernel {
       ReduceInfNanThreeSlotsLaunch<Tin, Tout>().Run(
           d, input.data(), input.size(), output_tensor->flat<Tout>().data());
 
+      cleanup.release();
       context->device()
           ->tensorflow_accelerator_device_info()
           ->event_mgr->ThenExecute(stream, std::move(check_cb));
@@ -942,7 +961,6 @@ class DebugNumericSummaryV2Op<GPUDevice, Tin, Tout> : public AsyncOpKernel {
       // TODO(cais): Implement other tensor debug modes in debug_event.proto.
       context->SetStatus(absl::UnimplementedError(absl::StrCat(
           "Unimplemented tensor debug mode: ", tensor_debug_mode_)));
-      done();
     }
   }
 
diff --git a/tensorflow/core/kernels/dynamic_partition_op_gpu.cu.cc b/tensorflow/core/kernels/dynamic_partition_op_gpu.cu.cc
index 37887759784..4ae920d9fd9 100644
--- a/tensorflow/core/kernels/dynamic_partition_op_gpu.cu.cc
+++ b/tensorflow/core/kernels/dynamic_partition_op_gpu.cu.cc
@@ -45,6 +45,7 @@ limitations under the License.
 #include "tensorflow/core/kernels/fill_functor.h"
 #include "tensorflow/core/kernels/gather_functor_gpu.cu.h"
 #include "tensorflow/core/kernels/gpu_prim.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
 #include "tensorflow/core/util/gpu_kernel_helper.h"
 #include "tensorflow/core/util/transform_output_iterator.h"
 
@@ -190,25 +191,25 @@ class DynamicPartitionOpGPU : public AsyncOpKernel {
   }
 
   void AllocateTempSpace(OpKernelContext* c, int32_t N, Tensor* indices_in,
-                         Tensor* partitions_out, Tensor* indices_out,
-                         DoneCallback done) {
+                         Tensor* partitions_out, Tensor* indices_out) {
     int32_t M = std::max(N, num_partitions_);
     // indices_in will be made slightly larger to accommodate
     // later computations.
     OP_REQUIRES_OK_ASYNC(
-        c, c->allocate_temp(DT_INT32, TensorShape({M}), indices_in), done);
+        c, c->allocate_temp(DT_INT32, TensorShape({M}), indices_in), []() {});
     OP_REQUIRES_OK_ASYNC(
-        c, c->allocate_temp(DT_INT32, TensorShape({N}), partitions_out), done);
+        c, c->allocate_temp(DT_INT32, TensorShape({N}), partitions_out),
+        []() {});
     OP_REQUIRES_OK_ASYNC(
-        c, c->allocate_temp(DT_INT32, TensorShape({N}), indices_out), done);
+        c, c->allocate_temp(DT_INT32, TensorShape({N}), indices_out), []() {});
   }
 
   void AllocateOutputs(OpKernelContext* c, const Tensor* data,
                        const Tensor* partitions, const Tensor* partition_count,
-                       OpOutputList* Tout, DoneCallback done) {
+                       OpOutputList* Tout) {
     auto e_part_count = partition_count->flat<int32_t>();
     // Allocate output tensors of the right size
-    OP_REQUIRES_OK_ASYNC(c, c->output_list("outputs", Tout), done);
+    OP_REQUIRES_OK_ASYNC(c, c->output_list("outputs", Tout), []() {});
     for (int p = 0; p < num_partitions_; p++) {
       TensorShape shape;
       shape.AddDim(e_part_count(p));
@@ -216,11 +217,12 @@ class DynamicPartitionOpGPU : public AsyncOpKernel {
         shape.AddDim(data->dim_size(i));
       }
       Tensor* out;
-      OP_REQUIRES_OK_ASYNC(c, Tout->allocate(p, shape, &out), done);
+      OP_REQUIRES_OK_ASYNC(c, Tout->allocate(p, shape, &out), []() {});
     }
   }
 
   void ComputeAsync(OpKernelContext* c, DoneCallback done) override {
+    auto cleanup = gtl::MakeCleanup([&done]() { done(); });
     const Tensor& data = c->input(0);
     const Tensor& partitions = c->input(1);
 
@@ -230,7 +232,7 @@ class DynamicPartitionOpGPU : public AsyncOpKernel {
             "data.shape must start with partitions.shape, ",
             "got data.shape = ", data.shape().DebugString(),
             ", partitions.shape = ", partitions.shape().DebugString())),
-        done);
+        []() {});
 
     Tensor partition_count;
 
@@ -243,13 +245,11 @@ class DynamicPartitionOpGPU : public AsyncOpKernel {
           c,
           c->allocate_temp(DT_INT32, TensorShape({num_partitions_}),
                            &partition_count, alloc_attr),
-          done);
+          []() {});
       auto e_part_count = partition_count.flat<int32_t>();
       for (int i = 0; i < num_partitions_; i++) e_part_count(i) = 0;
       OpOutputList outputs;
-      this->AllocateOutputs(c, &data, &partitions, &partition_count, &outputs,
-                            done);
-      if (c->status().ok()) done();
+      this->AllocateOutputs(c, &data, &partitions, &partition_count, &outputs);
       return;
     }
 
@@ -258,20 +258,19 @@ class DynamicPartitionOpGPU : public AsyncOpKernel {
         c,
         c->allocate_temp(DT_INT32, TensorShape({num_partitions_}),
                          &partition_count),
-        done);
+        []() {});
     Tensor indices_out;
     // Count how many times each partition index occurs.
     // Also sort the info in partitions and output it in indices_out,
     // in preparation for the next step.
-    this->CountAndSortParts(c, &partitions, &partition_count, &indices_out,
-                            done);
+    this->CountAndSortParts(c, &partitions, &partition_count, &indices_out);
     if (!c->status().ok()) return;
 
     // In order to allocate the output tensor we have to move partition_count
     // to CPU.
     auto* stream = c->op_device_context()->stream();
     OP_REQUIRES_ASYNC(c, stream,
-                      absl::InternalError("No GPU stream available."), done);
+                      absl::InternalError("No GPU stream available."), []() {});
     Tensor cpu_tensor;
     AllocatorAttributes alloc_attr;
     alloc_attr.set_on_host(true);
@@ -280,7 +279,7 @@ class DynamicPartitionOpGPU : public AsyncOpKernel {
         c,
         c->allocate_temp(partition_count.dtype(), partition_count.shape(),
                          &cpu_tensor, alloc_attr),
-        done);
+        []() {});
     stream_executor::DeviceAddressBase wrapped(
         partition_count.flat<int32_t>().data(),
         num_partitions_ * sizeof(int32_t));
@@ -288,7 +287,7 @@ class DynamicPartitionOpGPU : public AsyncOpKernel {
         c,
         stream->Memcpy(cpu_tensor.flat<int32_t>().data(), wrapped,
                        num_partitions_ * sizeof(int32_t)),
-        done);
+        []() {});
 
     // Keep a reference to partition_count so that the buffer
     // is not deallocated at the end of the function, before
@@ -296,29 +295,27 @@ class DynamicPartitionOpGPU : public AsyncOpKernel {
     TensorReference partition_ref(partition_count);
     auto wrapped_callback = [this, c, &data, &partitions, indices_out,
                              partition_ref, cpu_tensor, done]() {
-      {
-        auto stream = c->op_device_context()->stream();
-        std::unique_ptr<se::ActivateContext> scoped_activation =
-            stream->parent()->Activate();
-
-        OpOutputList outputs;
-        this->AllocateOutputs(c, &data, &partitions, &cpu_tensor, &outputs,
-                              done);
-        if (!c->status().ok()) {
-          partition_ref.Unref();
-          return;
-        }
-        int32_t N = partitions.NumElements();
-        int64_t slice_size = data.NumElements() / N;
-        this->GatherSlices(c, &data, &indices_out, N, slice_size, outputs);
+      // Declared before scoped_activation so ActivateContext is released before
+      // done() runs (preventing deadlock when done inlines another Op kernel).
+      auto cb_cleanup = gtl::MakeCleanup([&partition_ref, &done]() {
         partition_ref.Unref();
-      }  // Release ActivateContext to prevent deadlock when done
-         // inlines another Op kernel, which may assume the original cuda
-         // Context.
+        done();
+      });
+      auto stream = c->op_device_context()->stream();
+      std::unique_ptr<se::ActivateContext> scoped_activation =
+          stream->parent()->Activate();
 
-      done();
+      OpOutputList outputs;
+      this->AllocateOutputs(c, &data, &partitions, &cpu_tensor, &outputs);
+      if (!c->status().ok()) {
+        return;
+      }
+      int32_t N = partitions.NumElements();
+      int64_t slice_size = data.NumElements() / N;
+      this->GatherSlices(c, &data, &indices_out, N, slice_size, outputs);
     };
 
+    cleanup.release();
     c->device()->tensorflow_accelerator_device_info()->event_mgr->ThenExecute(
         stream, wrapped_callback);
   }
@@ -326,7 +323,7 @@ class DynamicPartitionOpGPU : public AsyncOpKernel {
  protected:
   void RadixSort(OpKernelContext* c, const Tensor* partitions,
                  Tensor* indices_in, Tensor* partitions_out,
-                 Tensor* indices_out, DoneCallback done) {
+                 Tensor* indices_out) {
     int32_t N = partitions->NumElements();
     const GPUDevice& device = c->eigen_device<GPUDevice>();
     const auto& cu_stream = GetGpuStream(c);
@@ -345,11 +342,12 @@ class DynamicPartitionOpGPU : public AsyncOpKernel {
         NULL, temp_storage_bytes, partitions_ptr, partitions_out_ptr,
         indices_in_ptr, indices_out_ptr, N, 0, sizeof(int32_t) * 8, cu_stream);
 
-    OP_REQUIRES(
+    OP_REQUIRES_ASYNC(
         c, gpuResult == gpuSuccess,
         absl::InternalError(absl::StrCat(
             "Failed to launch gpuprim::DeviceRadixSort::SortPairs to calculate",
-            "temp_storage_bytes, status: ", GpuGetErrorString(gpuResult))));
+            "temp_storage_bytes, status: ", GpuGetErrorString(gpuResult))),
+        []() {});
 
     // Allocate temporary storage.
     OP_REQUIRES_OK_ASYNC(
@@ -357,24 +355,24 @@ class DynamicPartitionOpGPU : public AsyncOpKernel {
         c->allocate_temp(
             DT_INT8, TensorShape({static_cast<int64_t>(temp_storage_bytes)}),
             &cub_temp_storage),
-        done);
+        []() {});
     // Radix-sort the partition information.
     gpuResult = gpuprim::DeviceRadixSort::SortPairs(
         cub_temp_storage.flat<int8_t>().data(), temp_storage_bytes,
         partitions_ptr, partitions_out_ptr, indices_in_ptr, indices_out_ptr, N,
         0, sizeof(int32_t) * 8, cu_stream);
 
-    OP_REQUIRES(
+    OP_REQUIRES_ASYNC(
         c, gpuResult == gpuSuccess,
         absl::InternalError(absl::StrCat(
             "Failed to launch gpuprim::DeviceRadixSort::SortPairs"
             "temp_storage_bytes: ",
-            temp_storage_bytes, "status: ", GpuGetErrorString(gpuResult))));
+            temp_storage_bytes, "status: ", GpuGetErrorString(gpuResult))),
+        []() {});
   }  // At this point cub_temp_storage will be marked for deallocation.
 
   void CountAndSortParts(OpKernelContext* c, const Tensor* partitions,
-                         Tensor* partition_count, Tensor* indices_out,
-                         DoneCallback done) {
+                         Tensor* partition_count, Tensor* indices_out) {
     const GPUDevice& device = c->eigen_device<GPUDevice>();
     const auto& cu_stream = GetGpuStream(c);
     int32_t N = partitions->NumElements();
@@ -383,11 +381,9 @@ class DynamicPartitionOpGPU : public AsyncOpKernel {
     Tensor aggregates_out;
 
     // Allocate memory for Radix-Sort.
-    this->AllocateTempSpace(c, N, &indices_in, &partitions_out, indices_out,
-                            done);
+    this->AllocateTempSpace(c, N, &indices_in, &partitions_out, indices_out);
     if (!c->status().ok()) return;
-    this->RadixSort(c, partitions, &indices_in, &partitions_out, indices_out,
-                    done);
+    this->RadixSort(c, partitions, &indices_in, &partitions_out, indices_out);
     if (!c->status().ok()) return;
     // We will now apply a reduce operation to count how many times
     // each index appears in partitions.
@@ -400,7 +396,7 @@ class DynamicPartitionOpGPU : public AsyncOpKernel {
         c,
         c->allocate_temp(DT_INT32, TensorShape({num_partitions_}),
                          &aggregates_out),
-        done);
+        []() {});
     // Obtain the pointers to inner buffers.
     int32_t* keys_in_ptr = partitions_out.flat<int32_t>().data();
     // Here we reuse the indices_in tensor for the unique keys output.
@@ -429,7 +425,7 @@ class DynamicPartitionOpGPU : public AsyncOpKernel {
     // Allocate space on GPU for the number of runs. This is required by CUB.
     Tensor num_runs;
     OP_REQUIRES_OK_ASYNC(
-        c, c->allocate_temp(DT_INT32, TensorShape({1}), &num_runs), done);
+        c, c->allocate_temp(DT_INT32, TensorShape({1}), &num_runs), []() {});
     int32_t* num_runs_ptr = num_runs.flat<int32_t>().data();
 
     // Determine temporary device storage requirements
@@ -439,11 +435,12 @@ class DynamicPartitionOpGPU : public AsyncOpKernel {
         NULL, temp_storage_bytes, keys_in_ptr, unique_out_it, values_in,
         aggregates_out_it, num_runs_ptr, reduction_op, N, cu_stream);
 
-    OP_REQUIRES(
+    OP_REQUIRES_ASYNC(
         c, gpuResult == gpuSuccess,
         absl::InternalError(absl::StrCat(
             "Failed to launch gpuprim::DeviceReduce::ReduceByKey ",
-            "temp_storage_bytes, status: ", GpuGetErrorString(gpuResult))));
+            "temp_storage_bytes, status: ", GpuGetErrorString(gpuResult))),
+        []() {});
 
     // Allocate temporary storage.
     OP_REQUIRES_OK_ASYNC(
@@ -451,7 +448,7 @@ class DynamicPartitionOpGPU : public AsyncOpKernel {
         c->allocate_temp(
             DT_INT8, TensorShape({static_cast<int64_t>(temp_storage_bytes)}),
             &cub_temp_storage),
-        done);
+        []() {});
     // Run reduce-by-key. The effect is that we count how many times
     // each index appears in partitions. The distinct indices are stored
     // in unique_out, while the count is stored in aggregates_out.
@@ -461,11 +458,13 @@ class DynamicPartitionOpGPU : public AsyncOpKernel {
         unique_out_it, values_in, aggregates_out_it, num_runs_ptr, reduction_op,
         N, cu_stream);
 
-    OP_REQUIRES(c, gpuResult == gpuSuccess,
-                absl::InternalError(absl::StrCat(
-                    "Failed to launch gpuprim::DeviceReduce::ReduceByKey ",
-                    "temp_storage_bytes: ", temp_storage_bytes,
-                    ", status: ", GpuGetErrorString(gpuResult))));
+    OP_REQUIRES_ASYNC(
+        c, gpuResult == gpuSuccess,
+        absl::InternalError(
+            absl::StrCat("Failed to launch gpuprim::DeviceReduce::ReduceByKey ",
+                         "temp_storage_bytes: ", temp_storage_bytes,
+                         ", status: ", GpuGetErrorString(gpuResult))),
+        []() {});
 
     // We are not done yet. unique_out only contains the indices that appeared
     // at least once in partitions. We move each value from aggregates_out
diff --git a/tensorflow/core/kernels/functional_ops.cc b/tensorflow/core/kernels/functional_ops.cc
index 2c1ab9c01c8..1a6f9515328 100644
--- a/tensorflow/core/kernels/functional_ops.cc
+++ b/tensorflow/core/kernels/functional_ops.cc
@@ -28,6 +28,7 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/lib/core/errors.h"
 #include "tensorflow/core/lib/core/threadpool.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
 #include "tensorflow/core/platform/casts.h"
 #include "tensorflow/core/platform/errors.h"
 #include "tensorflow/core/platform/macros.h"
@@ -162,7 +163,7 @@ class IfOp : public AsyncOpKernel {
     OP_REQUIRES_OK_ASYNC(ctx, GetHandles(ctx, &then_handle, &else_handle),
                          done);
     bool cond;
-    OP_REQUIRES_OK(ctx, ToBool({ctx->input(0)}, &cond));
+    OP_REQUIRES_OK_ASYNC(ctx, ToBool({ctx->input(0)}, &cond), done);
     (new State(this, ctx, cond, then_handle, else_handle, done))->Start();
   }
 
@@ -462,16 +463,17 @@ class WhileOp : public AsyncOpKernel {
   }
 
   void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override {
+    auto cleanup = gtl::MakeCleanup([&done]() { done(); });
     if (ctx->run_all_kernels_inline()) {
       // Use the non-callback-based implementation when kernels (and function
       // callbacks) execute inline to avoid stack overflow.
-      OP_REQUIRES_OK_ASYNC(ctx, DoComputeSync(ctx), done);
-      done();
+      OP_REQUIRES_OK_ASYNC(ctx, DoComputeSync(ctx), []() {});
     } else {
       FHandle cond_handle;
       FHandle body_handle;
       OP_REQUIRES_OK_ASYNC(ctx, GetHandles(ctx, &cond_handle, &body_handle),
-                           done);
+                           []() {});
+      cleanup.release();
       (new State(this, ctx, cond_handle, body_handle, done))->Start();
     }
   }
diff --git a/tensorflow/core/kernels/queue_op.cc b/tensorflow/core/kernels/queue_op.cc
index 96e4755993c..af325b004be 100644
--- a/tensorflow/core/kernels/queue_op.cc
+++ b/tensorflow/core/kernels/queue_op.cc
@@ -25,6 +25,7 @@ limitations under the License.
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
 #include "tensorflow/core/platform/macros.h"
 #include "tensorflow/core/platform/refcount.h"
 #include "tensorflow/core/platform/types.h"
@@ -364,10 +365,11 @@ QueueSizeOp::QueueSizeOp(OpKernelConstruction* context)
 
 void QueueSizeOp::ComputeAsync(OpKernelContext* ctx, QueueInterface* queue,
                                DoneCallback callback) {
+  auto cleanup = gtl::MakeCleanup([&callback]() { callback(); });
   Tensor* Tqueue_size = nullptr;
-  OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({}), &Tqueue_size));
+  OP_REQUIRES_OK_ASYNC(
+      ctx, ctx->allocate_output(0, TensorShape({}), &Tqueue_size), []() {});
   Tqueue_size->flat<int32_t>().setConstant(queue->size());
-  callback();
 }
 
 QueueIsClosedOp::QueueIsClosedOp(OpKernelConstruction* context)
@@ -375,11 +377,12 @@ QueueIsClosedOp::QueueIsClosedOp(OpKernelConstruction* context)
 
 void QueueIsClosedOp::ComputeAsync(OpKernelContext* ctx, QueueInterface* queue,
                                    DoneCallback callback) {
+  auto cleanup = gtl::MakeCleanup([&callback]() { callback(); });
   Tensor* Tqueue_is_closed = nullptr;
-  OP_REQUIRES_OK(ctx,
-                 ctx->allocate_output(0, TensorShape({}), &Tqueue_is_closed));
+  OP_REQUIRES_OK_ASYNC(
+      ctx, ctx->allocate_output(0, TensorShape({}), &Tqueue_is_closed),
+      []() {});
   Tqueue_is_closed->flat<bool>().setConstant(queue->is_closed());
-  callback();
 }
 
 }  // namespace tensorflow
diff --git a/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h b/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h
index bbb1e081934..56641973deb 100644
--- a/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h
+++ b/tensorflow/core/kernels/segment_reduction_ops_gpu.cu.h
@@ -28,6 +28,7 @@ limitations under the License.
 #include "tensorflow/core/kernels/gpu_prim_helpers.h"
 #include "tensorflow/core/kernels/segment_reduction_ops.h"
 #include "tensorflow/core/lib/core/bits.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
 #include "tensorflow/core/util/determinism.h"
 #include "tensorflow/core/util/env_var.h"
 #include "tensorflow/core/util/gpu_device_functions.h"
@@ -1394,6 +1395,7 @@ struct SparseSegmentGradV2Functor<GPUDevice, T, Tindices, Tsegmentids> {
             const Tsegmentids_internal* segment_ids,
             const TensorShape& dense_output_shape,
             typename AsyncOpKernel::DoneCallback done) {
+    auto cleanup = gtl::MakeCleanup([&done]() { done(); });
     const int64_t dense_output_dim0 = dense_output_shape.dim_size(0);
 
     // Allocate and compute segment weights (for Mean/SqrtN operations only).
@@ -1401,7 +1403,8 @@ struct SparseSegmentGradV2Functor<GPUDevice, T, Tindices, Tsegmentids> {
     Tweights* weights_ptr = nullptr;
     if (operation != SparseSegmentReductionOperation::kSum) {
       ComputeSegmentWeights(context, operation, nsegments, nouter, segment_ids,
-                            &tmp_weights, done);
+                            &tmp_weights);
+      if (!context->status().ok()) return;
       weights_ptr = tmp_weights.flat<Tweights>().data();
     }
 
@@ -1415,14 +1418,14 @@ struct SparseSegmentGradV2Functor<GPUDevice, T, Tindices, Tsegmentids> {
           context,
           context->allocate_temp(DataTypeToEnum<TindicesCompact>::value,
                                  TensorShape({nouter}), &tmp_sorted_indices),
-          done);
+          []() {});
       TindicesCompact* tmp_sorted_indices_ptr =
           tmp_sorted_indices.flat<TindicesCompact>().data();
       OP_REQUIRES_OK_ASYNC(
           context,
           context->allocate_temp(DataTypeToEnum<Tsegmentids_internal>::value,
                                  TensorShape({nouter}), &tmp_permuted_segment),
-          done);
+          []() {});
       Tsegmentids_internal* tmp_permuted_segment_ptr =
           tmp_permuted_segment.flat<Tsegmentids_internal>().data();
       OP_REQUIRES_OK_ASYNC(
@@ -1433,7 +1436,7 @@ struct SparseSegmentGradV2Functor<GPUDevice, T, Tindices, Tsegmentids> {
                        /*indices_in=*/segment_ids,
                        /*indices_out=*/tmp_permuted_segment_ptr,
                        /*num_bits=*/Log2Ceiling64(dense_output_dim0)),
-          done);
+          []() {});
       sorted_indices_ptr = tmp_sorted_indices_ptr;
       permuted_segment_ptr = tmp_permuted_segment_ptr;
       // The original tensors are no longer needed.
@@ -1455,18 +1458,18 @@ struct SparseSegmentGradV2Functor<GPUDevice, T, Tindices, Tsegmentids> {
                          context->allocate_temp(DataTypeToEnum<Toffsets>::value,
                                                 TensorShape({nouter}),
                                                 &tmp_sorted_indices_unique_ids),
-                         done);
+                         []() {});
     Toffsets* sorted_indices_unique_ids_ptr =
         tmp_sorted_indices_unique_ids.flat<Toffsets>().data();
     OP_REQUIRES_OK_ASYNC(
         context,
         GpuInclusivePrefixSum(context, nouter, sorted_indices_edge_indicator,
                               sorted_indices_unique_ids_ptr),
-        done);
+        []() {});
 
     se::Stream* stream = context->op_device_context()->stream();
     OP_REQUIRES_ASYNC(context, stream,
-                      absl::InternalError("No GPU stream available."), done);
+                      absl::InternalError("No GPU stream available."), []() {});
 
     // Copy the last element of sorted_indices_unique_ids back to the host to
     // obtain num_unique.
@@ -1480,7 +1483,7 @@ struct SparseSegmentGradV2Functor<GPUDevice, T, Tindices, Tsegmentids> {
                     (nouter - 1),
                 sizeof(*last_idx_host.data())),
             sizeof(*last_idx_host.data())),
-        done);
+        []() {});
 
     auto async_finish_computation =
         [this, context, dense_output_shape, nouter, ninner, input,
@@ -1489,6 +1492,7 @@ struct SparseSegmentGradV2Functor<GPUDevice, T, Tindices, Tsegmentids> {
          segment_ids_tensor, tmp_permuted_segment, permuted_segment_ptr,
          sorted_indices_edge_indicator, tmp_weights, weights_ptr, last_idx_host,
          done]() -> void {
+      auto cb_cleanup = gtl::MakeCleanup([&done]() { done(); });
       const GPUDevice& device = context->eigen_gpu_device();
       Toffsets num_unique = (*last_idx_host.data()) + 1;
 
@@ -1496,12 +1500,12 @@ struct SparseSegmentGradV2Functor<GPUDevice, T, Tindices, Tsegmentids> {
           context->op_device_context()->stream()->parent()->Activate();
 
       TensorShape output_shape = dense_output_shape;
-      OP_REQUIRES_OK_ASYNC(context,
-                           output_shape.SetDimWithStatus(0, num_unique), done);
+      OP_REQUIRES_OK_ASYNC(
+          context, output_shape.SetDimWithStatus(0, num_unique), []() {});
       Tensor* output = nullptr;
       T* output_ptr;
       OP_REQUIRES_OK_ASYNC(
-          context, context->allocate_output(0, output_shape, &output), done);
+          context, context->allocate_output(0, output_shape, &output), []() {});
       output_ptr = output->flat<T>().data();
 
       // Compute the gradient using a weighted SegmentReduceGPU with the segment
@@ -1522,7 +1526,7 @@ struct SparseSegmentGradV2Functor<GPUDevice, T, Tindices, Tsegmentids> {
                                /*indices=*/permuted_segment_ptr,
                                /*weights=*/weights_ptr,
                                /*output=*/output_ptr),
-                           done);
+                           []() {});
 
       Tensor* sorted_unique_indices = nullptr;
       Tindices* sorted_unique_indices_ptr;
@@ -1530,7 +1534,7 @@ struct SparseSegmentGradV2Functor<GPUDevice, T, Tindices, Tsegmentids> {
           context,
           context->allocate_output(1, TensorShape({num_unique}),
                                    &sorted_unique_indices),
-          done);
+          []() {});
       sorted_unique_indices_ptr =
           sorted_unique_indices->flat<Tindices>().data();
 
@@ -1539,11 +1543,10 @@ struct SparseSegmentGradV2Functor<GPUDevice, T, Tindices, Tsegmentids> {
           LaunchScatterUniqueIndicesKernel(
               device, nouter, sorted_indices_edge_indicator, sorted_indices_ptr,
               sorted_indices_unique_ids_ptr, sorted_unique_indices_ptr),
-          done);
-
-      done();
+          []() {});
     };
 
+    cleanup.release();
     context->device()
         ->tensorflow_accelerator_device_info()
         ->event_mgr->ThenExecute(stream, async_finish_computation);
@@ -1554,14 +1557,13 @@ struct SparseSegmentGradV2Functor<GPUDevice, T, Tindices, Tsegmentids> {
                              SparseSegmentReductionOperation operation,
                              Tsegmentids_internal nsegments, Toffsets nouter,
                              const Tsegmentids_internal* segment_ids,
-                             Tensor* tmp_weights,
-                             typename AsyncOpKernel::DoneCallback done) {
+                             Tensor* tmp_weights) {
     const GPUDevice& device = context->eigen_gpu_device();
     OP_REQUIRES_OK_ASYNC(
         context,
         context->allocate_temp(DataTypeToEnum<Tweights>::value,
                                TensorShape({nsegments}), tmp_weights),
-        done);
+        []() {});
     Tweights* weights_ptr = tmp_weights->flat<Tweights>().data();
     // Allocate and compute segment_offsets.
     Tensor tmp_segment_offsets;
@@ -1569,19 +1571,19 @@ struct SparseSegmentGradV2Functor<GPUDevice, T, Tindices, Tsegmentids> {
                          context->allocate_temp(DataTypeToEnum<Toffsets>::value,
                                                 TensorShape({nsegments + 1}),
                                                 &tmp_segment_offsets),
-                         done);
+                         []() {});
     Toffsets* segment_offsets_ptr = tmp_segment_offsets.flat<Toffsets>().data();
     OP_REQUIRES_OK_ASYNC(
         context,
         LaunchSegmentOffsetsKernel(device, nouter, nsegments, segment_ids,
                                    segment_offsets_ptr),
-        done);
+        []() {});
     // Compute the weights based on the segment sizes using segment_offsets.
     OP_REQUIRES_OK_ASYNC(
         context,
         LaunchSegmentWeightsKernel(device, nsegments, operation,
                                    segment_offsets_ptr, weights_ptr),
-        done);
+        []() {});
   }
 };
 
diff --git a/tensorflow/core/kernels/segment_reduction_ops_impl.h b/tensorflow/core/kernels/segment_reduction_ops_impl.h
index 727e04a88bd..ff249d4af5d 100644
--- a/tensorflow/core/kernels/segment_reduction_ops_impl.h
+++ b/tensorflow/core/kernels/segment_reduction_ops_impl.h
@@ -51,6 +51,7 @@ limitations under the License.
 #include "tensorflow/core/kernels/fill_functor.h"
 #include "tensorflow/core/kernels/segment_reduction_ops.h"
 #include "tensorflow/core/lib/core/status.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
 #include "tensorflow/core/platform/bfloat16.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/util/determinism.h"
@@ -240,16 +241,17 @@ class SegmentReductionGPUOp : public AsyncOpKernel {
       : AsyncOpKernel(context) {}
 
   void ComputeAsync(OpKernelContext* context, DoneCallback done) override {
+    auto cleanup = gtl::MakeCleanup([&done]() { done(); });
     const Tensor& input = context->input(0);
     const Tensor& segment_ids = context->input(1);
 
     OP_REQUIRES_ASYNC(
         context, TensorShapeUtils::IsVector(segment_ids.shape()),
-        absl::InvalidArgumentError("segment_ids should be a vector."), done);
+        absl::InvalidArgumentError("segment_ids should be a vector."), []() {});
 
     OP_REQUIRES_ASYNC(
         context, input.dims() >= 1,
-        absl::InvalidArgumentError("Shape must be at least rank 1"), done);
+        absl::InvalidArgumentError("Shape must be at least rank 1"), []() {});
 
     const int64_t num_indices = segment_ids.NumElements();
     OP_REQUIRES_ASYNC(
@@ -257,7 +259,7 @@ class SegmentReductionGPUOp : public AsyncOpKernel {
         absl::InvalidArgumentError(
             "segment_ids should be the same size as dimension 0 of"
             " input."),
-        done);
+        []() {});
 
     if (num_indices == 0) {
       TensorShape output_shape = input.shape();
@@ -265,8 +267,7 @@ class SegmentReductionGPUOp : public AsyncOpKernel {
 
       Tensor* output = nullptr;
       OP_REQUIRES_OK_ASYNC(
-          context, context->allocate_output(0, output_shape, &output), done);
-      done();
+          context, context->allocate_output(0, output_shape, &output), []() {});
       return;
     }
 
@@ -279,11 +280,12 @@ class SegmentReductionGPUOp : public AsyncOpKernel {
     OP_REQUIRES_OK_ASYNC(context,
                          stream->Memcpy(output_rows_host.mutable_data(),
                                         output_rows_device, sizeof(Index)),
-                         done);
+                         []() {});
 
     SegmentReductionFunctor functor_;
     auto create_and_check_output = [context, output_rows_host, &input,
                                     &segment_ids, &functor_, done]() {
+      auto cb_cleanup = gtl::MakeCleanup([&done]() { done(); });
       // Ensure that within the callback, the proper GPU settings are
       // configured.
       auto stream = context->op_device_context()->stream();
@@ -294,17 +296,17 @@ class SegmentReductionGPUOp : public AsyncOpKernel {
       output_rows++;
       OP_REQUIRES_ASYNC(context, output_rows > 0,
                         absl::InvalidArgumentError("segment ids must be >= 0"),
-                        done);
+                        []() {});
 
       TensorShape output_shape = input.shape();
       // Since we're changing the first dimension of the shape, we need to make
       // sure the new shape won't overflow.
-      OP_REQUIRES_OK_ASYNC(context,
-                           output_shape.SetDimWithStatus(0, output_rows), done);
+      OP_REQUIRES_OK_ASYNC(
+          context, output_shape.SetDimWithStatus(0, output_rows), []() {});
 
       Tensor* output = nullptr;
       OP_REQUIRES_OK_ASYNC(
-          context, context->allocate_output(0, output_shape, &output), done);
+          context, context->allocate_output(0, output_shape, &output), []() {});
 
       bool use_deterministic_kernels =
           UseDeterministicSegmentReductions() ||
@@ -324,7 +326,7 @@ class SegmentReductionGPUOp : public AsyncOpKernel {
           absl::UnimplementedError(
               "Deterministic GPU implementation of sorted segment reduction op"
               " not available."),
-          done);
+          []() {});
 
       auto output_flat = output->flat_outer_dims<T>();
       auto data_ptr = input.template flat<T>().data();
@@ -332,10 +334,9 @@ class SegmentReductionGPUOp : public AsyncOpKernel {
       functor_(context, context->eigen_device<GPUDevice>(), output_rows,
                segment_ids.shape(), IsMean, segment_flat, input.NumElements(),
                data_ptr, output_flat);
-
-      done();
     };
 
+    cleanup.release();
     context->device()
         ->tensorflow_accelerator_device_info()
         ->event_mgr->ThenExecute(stream, create_and_check_output);
@@ -939,6 +940,7 @@ class SparseSegmentReductionOpBase<GPUDevice, T, Index, SegmentId>
         default_value_(default_value) {}
 
   void ComputeAsync(OpKernelContext* context, DoneCallback done) override {
+    auto cleanup = gtl::MakeCleanup([&done]() { done(); });
     const Tensor& input = context->input(0);
     const Tensor& indices = context->input(1);
     const Tensor& segment_ids = context->input(2);
@@ -947,12 +949,13 @@ class SparseSegmentReductionOpBase<GPUDevice, T, Index, SegmentId>
         context,
         internal::ValidateSparseSegmentReduction(
             context, input, indices, segment_ids, has_num_segments_),
-        done);
+        []() {});
 
     ScratchSpace<SegmentId> last_segment_id_host(context, 1, /*on_host=*/true);
 
     auto create_and_check_output = [this, context, input, indices, segment_ids,
                                     last_segment_id_host, done]() {
+      auto async_cleanup = gtl::MakeCleanup([&done]() { done(); });
       // Ensure that within the callback, the proper GPU settings are
       // configured.
       auto stream = context->op_device_context()->stream();
@@ -963,14 +966,15 @@ class SparseSegmentReductionOpBase<GPUDevice, T, Index, SegmentId>
       SegmentId output_rows = last_segment_id + 1;
       OP_REQUIRES_ASYNC(context, output_rows > 0,
                         absl::InvalidArgumentError("segment ids must be >= 0"),
-                        done);
+                        []() {});
 
       TensorShape output_shape = input.shape();
-      output_shape.set_dim(0, output_rows);
+      OP_REQUIRES_OK_ASYNC(
+          context, output_shape.SetDimWithStatus(0, output_rows), []() {});
 
       Tensor* output = nullptr;
       OP_REQUIRES_OK_ASYNC(
-          context, context->allocate_output(0, output_shape, &output), done);
+          context, context->allocate_output(0, output_shape, &output), []() {});
 
       auto input_flat = input.flat_outer_dims<T>();
       const auto indices_vec = indices.vec<Index>();
@@ -982,8 +986,7 @@ class SparseSegmentReductionOpBase<GPUDevice, T, Index, SegmentId>
           context,
           functor(context, is_mean_, is_sqrtn_, default_value_, input_flat,
                   indices_vec, segment_ids_vec, output_flat),
-          done);
-      done();
+          []() {});
     };
 
     if (has_num_segments_) {
@@ -994,17 +997,19 @@ class SparseSegmentReductionOpBase<GPUDevice, T, Index, SegmentId>
                                        ? num_segments_t.scalar<int32_t>()()
                                        : num_segments_t.scalar<int64_t>()());
       *last_segment_id_host.mutable_data() = num_segments - 1;
+      cleanup.release();
       create_and_check_output();
     } else {
       const int64_t num_indices = indices.NumElements();
       if (num_indices == 0) {
         TensorShape output_shape = input.shape();
-        output_shape.set_dim(0, 0);
+        OP_REQUIRES_OK_ASYNC(context, output_shape.SetDimWithStatus(0, 0),
+                             []() {});
 
         Tensor* output = nullptr;
-        OP_REQUIRES_OK_ASYNC(
-            context, context->allocate_output(0, output_shape, &output), done);
-        done();
+        OP_REQUIRES_OK_ASYNC(context,
+                             context->allocate_output(0, output_shape, &output),
+                             []() {});
         return;
       }
 
@@ -1018,7 +1023,8 @@ class SparseSegmentReductionOpBase<GPUDevice, T, Index, SegmentId>
           context,
           stream->Memcpy(last_segment_id_host.mutable_data(),
                          last_segment_id_device, sizeof(SegmentId)),
-          done);
+          []() {});
+      cleanup.release();
       context->device()
           ->tensorflow_accelerator_device_info()
           ->event_mgr->ThenExecute(stream, create_and_check_output);
@@ -1447,6 +1453,9 @@ class SparseSegmentGradV2OpCommon {
       Tensor* sorted_unique_indices = nullptr;
       TF_RETURN_IF_ERROR(context->allocate_output(1, TensorShape({0}),
                                                   &sorted_unique_indices));
+      if (done) {
+        done();
+      }
       return absl::OkStatus();
     }
 
diff --git a/tensorflow/core/kernels/sparse_to_dense_op_gpu.cu.cc b/tensorflow/core/kernels/sparse_to_dense_op_gpu.cu.cc
index 8c9e4d2e7ff..79fe4bd4328 100644
--- a/tensorflow/core/kernels/sparse_to_dense_op_gpu.cu.cc
+++ b/tensorflow/core/kernels/sparse_to_dense_op_gpu.cu.cc
@@ -25,6 +25,7 @@ limitations under the License.
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/types.h"
+#include "tensorflow/core/lib/gtl/cleanup.h"
 #include "tensorflow/core/platform/stream_executor.h"
 #include "tensorflow/core/util/gpu_kernel_helper.h"
 
@@ -155,6 +156,7 @@ void LaunchSparseToDense<T, Index>::operator()(
     OpKernelContext* c, AsyncOpKernel::DoneCallback done, AsyncOpKernel* op,
     bool validate_indices, const Tensor& indices, const Tensor& values,
     const Tensor& shape, const T default_value, Tensor* dense) {
+  auto cleanup = gtl::MakeCleanup([&done]() { done(); });
   auto* stream = c->op_device_context()->stream();
   const Eigen::GpuDevice& d = c->eigen_gpu_device();
 
@@ -179,23 +181,27 @@ void LaunchSparseToDense<T, Index>::operator()(
         c,
         c->allocate_temp(DT_INT32, TensorShape({valid_status_size}),
                          &valid_status_tensor),
-        done);
+        []() {});
 
     auto status_ptr = valid_status_tensor.template flat<int>().data();
     stream_executor::DeviceAddressBase valid_status_ptr(status_ptr,
                                                         valid_status_bytes);
 
     GpuLaunchConfig config = GetGpuLaunchConfig(num_elems, d);
-    OP_REQUIRES_OK(
-        c, stream->Memset32(&valid_status_ptr, INT_MAX, valid_status_bytes));
+    OP_REQUIRES_OK_ASYNC(
+        c, stream->Memset32(&valid_status_ptr, INT_MAX, valid_status_bytes),
+        []() {});
     OP_REQUIRES_OK_ASYNC(
         c,
         GpuLaunchKernel(CheckIndicesValid<Index>, config.block_count,
                         config.thread_per_block, 0, d.stream(), indices_ptr,
                         num_elems, shape_ptr, num_dims, status_ptr),
-        done);
-    OP_REQUIRES_OK(c, stream->Memcpy(reinterpret_cast<int*>(valid_status.get()),
-                                     valid_status_ptr, valid_status_bytes));
+        []() {});
+    OP_REQUIRES_OK_ASYNC(
+        c,
+        stream->Memcpy(reinterpret_cast<int*>(valid_status.get()),
+                       valid_status_ptr, valid_status_bytes),
+        []() {});
 
     // We capture 'shape' instead of 'shape_ptr' since this lambda outlives
     // the 'shape' tensor.
@@ -205,48 +211,44 @@ void LaunchSparseToDense<T, Index>::operator()(
                                      dense_size, default_value, indices_ptr,
                                      values_ptr, num_elems, num_values, shape,
                                      num_dims, dense_ptr, done]() {
-      {
-        // Ensure that within the callback, the proper GPU settings are
-        // configured.
-        auto stream = c->op_device_context()->stream();
-        std::unique_ptr<se::ActivateContext> scoped_activation =
-            stream->parent()->Activate();
+      // Declared before scoped_activation so ActivateContext is released before
+      // done() runs (preventing deadlock when done inlines another Op kernel).
+      auto cb_cleanup = gtl::MakeCleanup([&done]() { done(); });
+      auto stream = c->op_device_context()->stream();
+      std::unique_ptr<se::ActivateContext> scoped_activation =
+          stream->parent()->Activate();
 
-        OP_REQUIRES_ASYNC(
-            c, valid_status->valid == INT_MAX,
-            absl::InvalidArgumentError(absl::StrCat(
-                "indices[", valid_status->valid, "] is out of bounds.")),
-            done);
+      OP_REQUIRES_ASYNC(
+          c, valid_status->valid == INT_MAX,
+          absl::InvalidArgumentError(absl::StrCat(
+              "indices[", valid_status->valid, "] is out of bounds.")),
+          []() {});
 
-        OP_REQUIRES_ASYNC(c, valid_status->increasing == INT_MAX,
-                          absl::InvalidArgumentError(absl::StrCat(
-                              "indices[", valid_status->increasing,
-                              "] is out of "
-                              "order. Many sparse ops require sorted indices.\n"
-                              "  Use `tf.sparse.reorder` to create a correctly "
-                              "ordered copy.\n\n")),
-                          done);
+      OP_REQUIRES_ASYNC(c, valid_status->increasing == INT_MAX,
+                        absl::InvalidArgumentError(absl::StrCat(
+                            "indices[", valid_status->increasing,
+                            "] is out of "
+                            "order. Many sparse ops require sorted indices.\n"
+                            "  Use `tf.sparse.reorder` to create a correctly "
+                            "ordered copy.\n\n")),
+                        []() {});
 
-        OP_REQUIRES_ASYNC(c, valid_status->different == INT_MAX,
-                          absl::InvalidArgumentError(
-                              absl::StrCat("indices[", valid_status->different,
-                                           "] is "
-                                           "repeated.")),
-                          done);
+      OP_REQUIRES_ASYNC(c, valid_status->different == INT_MAX,
+                        absl::InvalidArgumentError(
+                            absl::StrCat("indices[", valid_status->different,
+                                         "] is "
+                                         "repeated.")),
+                        []() {});
 
-        OP_REQUIRES_OK_ASYNC(
-            c,
-            LaunchComputeKernels(c, dense_size, default_value, indices_ptr,
-                                 values_ptr, num_elems, num_values,
-                                 shape.flat<Index>().data(), num_dims,
-                                 dense_ptr),
-            done);
-      }  // Release ActivateContext to prevent deadlock when done
-      // inlines another Op kernel, which may assume the original cuda Context.
-
-      done();
+      OP_REQUIRES_OK_ASYNC(
+          c,
+          LaunchComputeKernels(c, dense_size, default_value, indices_ptr,
+                               values_ptr, num_elems, num_values,
+                               shape.flat<Index>().data(), num_dims, dense_ptr),
+          []() {});
     };
 
+    cleanup.release();
     c->device()->tensorflow_accelerator_device_info()->event_mgr->ThenExecute(
         stream, check_status_and_compute);
   } else {
@@ -255,8 +257,7 @@ void LaunchSparseToDense<T, Index>::operator()(
         LaunchComputeKernels(c, dense_size, default_value, indices_ptr,
                              values_ptr, num_elems, num_values, shape_ptr,
                              num_dims, dense_ptr),
-        done);
-    done();
+        []() {});
   }
 }
 

From 74c87ebc304f4fb6ec06979545449d0624a129cf Mon Sep 17 00:00:00 2001
From: Alex Pivovarov <upwind@google.com>
Date: Mon, 12 Oct 2026 10:59:19 -0700
Subject: [PATCH 0602/1000] [XLA:GPU] Constrain __gpu$TopK operands to default
 layout in GpuLayoutAssignment

Previously, `GpuLayoutAssignment::AddBackendConstraints` only constrained the output buffers (`{0}` and `{1}`) of `__gpu$TopK` (`kTopKCustomCallTarget`) to `default_layout` via `SetBufferLayout`, and relied on `GpuLayoutAssignment::InstructionCanChangeLayoutInstance` returning `false` to propagate the output layout backward to the operand.
Because `InstructionCanChangeLayoutInstance == false` marks an instruction as layout-agnostic in both directions (`OperandLayoutAlwaysPropagateToSiblings == true`), an upstream instruction with a mandatory non-default layout (such as an `all-gather` along dimension 1 with `{0, 1}` layout) processed before `__gpu$TopK` propagates `{0, 1}` forward into `__gpu$TopK`'s unconstrained operand and overwrites `__gpu$TopK`'s output buffer layouts with `{0, 1}`. Both the built-in GPU TopK kernel and the RAFT `select_k` FFI handler require contiguous default (`{1, 0}`) layout on both inputs and outputs, so receiving `{0, 1}` corrupts the TopK results.
### Changes:
- Explicitly constrain all operands of `__gpu$TopK` to `default_layout` via `SetOperandLayout` in `GpuLayoutAssignment::AddBackendConstraints`, matching `kSort` and `IsCubDeviceRadixSortNoScratchSize`.
- Remove the `GpuLayoutAssignment::InstructionCanChangeLayoutInstance` override for `kTopKCustomCallTarget` so `LayoutAssignment` does not overwrite `__gpu$TopK` output layouts and inserts a layout-changing copy/transpose when needed.
- Add `TopKAfterAllGatherHasDefaultLayout` regression test in `xla/backends/gpu/transforms/layout_assignment_test.cc`.

PiperOrigin-RevId: 998025140
---
 .../xla/xla/backends/gpu/transforms/BUILD     |  5 +--
 .../gpu/transforms/layout_assignment.cc       | 26 +++--------
 .../gpu/transforms/layout_assignment.h        |  6 ---
 .../gpu/transforms/layout_assignment_test.cc  | 43 +++++++++++++++++++
 4 files changed, 51 insertions(+), 29 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD
index 3ae8c44ac8d..648e57ec637 100644
--- a/third_party/xla/xla/backends/gpu/transforms/BUILD
+++ b/third_party/xla/xla/backends/gpu/transforms/BUILD
@@ -2322,13 +2322,9 @@ cc_library(
         "//xla/service/gpu:cublas_cudnn",
         "//xla/service/gpu:ir_emission_utils",
         "//xla/service/gpu:matmul_utils",
-        "//xla/service/gpu:reduction_utils",
         "//xla/service/gpu:stream_executor_util",
         "//xla/stream_executor:device_description",
-        "//xla/stream_executor:dnn",
         "//xla/stream_executor/cuda:cuda_compute_capability",
-        "//xla/tsl/platform:errors",
-        "//xla/tsl/platform:statusor",
         "//xla/tsl/protobuf:dnn_proto_cc",
         "//xla/tsl/util:env_var",
         "@com_google_absl//absl/algorithm:container",
@@ -2359,6 +2355,7 @@ xla_cc_test(
         "//xla/hlo/testlib:pattern_matcher_gmock",
         "//xla/service:computation_layout",
         "//xla/service:hlo_verifier",
+        "//xla/service:layout_assignment",
         "//xla/service:pattern_matcher",
         "//xla/stream_executor:device_description",
         "//xla/stream_executor/cuda:cuda_compute_capability",
diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment.cc b/third_party/xla/xla/backends/gpu/transforms/layout_assignment.cc
index ff0811a8a21..7736ebe1892 100644
--- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment.cc
@@ -48,7 +48,6 @@ limitations under the License.
 #include "xla/service/gpu/cublas_cudnn.h"
 #include "xla/service/gpu/ir_emission_utils.h"
 #include "xla/service/gpu/matmul_utils.h"
-#include "xla/service/gpu/reduction_utils.h"
 #include "xla/service/gpu/stream_executor_util.h"
 #include "xla/service/hlo_value.h"
 #include "xla/service/layout_assignment.h"
@@ -59,9 +58,6 @@ limitations under the License.
 #include "xla/shape_util.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/tsl/platform/errors.h"
-#include "xla/tsl/platform/statusor.h"
 #include "xla/tsl/protobuf/dnn.pb.h"
 #include "xla/tsl/util/env_var.h"
 #include "xla/util.h"
@@ -654,9 +650,15 @@ absl::Status GpuLayoutAssignment::AddBackendConstraints(
         ABSL_RETURN_IF_ERROR(SetBufferLayout(keys_layout, output_buffer));
       }
     } else if (IsCustomCallToTopK(*instruction)) {
-      // The output of the TopK custom call needs to have default layout.
+      // The operands and outputs of the TopK custom call need to have default
+      // layout.
       Layout default_layout = LayoutUtil::GetDefaultLayoutForRank(
           instruction->operand(0)->shape().dimensions().size());
+      for (int64_t i = 0; i < instruction->operand_count(); ++i) {
+        Shape shape = instruction->operand(i)->shape();
+        *shape.mutable_layout() = default_layout;
+        ABSL_RETURN_IF_ERROR(SetOperandLayout(shape, instruction, i));
+      }
       const HloValue& values_buffer =
           dataflow_analysis_->GetValueDefinedAt(instruction, {0});
       ABSL_RETURN_IF_ERROR(SetBufferLayout(default_layout, values_buffer));
@@ -943,20 +945,6 @@ bool GpuLayoutAssignment::PropagateReductionLayoutToOperand(
   return reduction_size >= 32;
 }
 
-bool GpuLayoutAssignment::InstructionCanChangeLayoutInstance(
-    const HloInstruction* instruction) {
-  // The TopK custom call cannot handle the case if the operand has a different
-  // layout.
-  const HloCustomCallInstruction* custom_call =
-      DynCast<HloCustomCallInstruction>(instruction);
-  if (custom_call != nullptr &&
-      custom_call->custom_call_target() == kTopKCustomCallTarget) {
-    return false;
-  }
-
-  return LayoutAssignment::InstructionCanChangeLayoutInstance(instruction);
-}
-
 // GPU specific layout alignment for reshapes.
 // ShapeUtil::AlignLayouts puts 1-sized dimensions at the end of minor_to_major
 // (most major), which is good for TPU to avoid padding. But on GPU it can
diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment.h b/third_party/xla/xla/backends/gpu/transforms/layout_assignment.h
index 4d0b9ee40f6..fd01d7deb58 100644
--- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment.h
+++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment.h
@@ -19,7 +19,6 @@ limitations under the License.
 #include <cstdint>
 #include <initializer_list>
 #include <memory>
-#include <optional>
 
 #include "absl/status/status.h"
 #include "absl/types/span.h"
@@ -28,9 +27,7 @@ limitations under the License.
 #include "xla/layout.h"
 #include "xla/service/computation_layout.h"
 #include "xla/service/layout_assignment.h"
-#include "xla/shape.h"
 #include "xla/stream_executor/device_description.h"
-#include "xla/stream_executor/dnn.h"
 
 namespace xla {
 namespace gpu {
@@ -81,9 +78,6 @@ class GpuLayoutAssignment : public LayoutAssignment {
 
   bool PropagateReductionLayoutToOperand(const HloInstruction* user) override;
 
-  bool InstructionCanChangeLayoutInstance(
-      const HloInstruction* instruction) override;
-
   std::unique_ptr<Layout> ChooseOperandLayoutFromOutputLayout(
       const Layout& output_layout, const HloInstruction* instruction,
       int64_t operand_no) override;
diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_test.cc b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_test.cc
index f63ad72b827..cf3849fd530 100644
--- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment_test.cc
@@ -37,6 +37,7 @@ limitations under the License.
 #include "xla/layout_util.h"
 #include "xla/service/computation_layout.h"
 #include "xla/service/hlo_verifier.h"
+#include "xla/service/layout_assignment.h"
 #include "xla/service/pattern_matcher.h"
 #include "xla/shape.h"
 #include "xla/shape_layout.h"
@@ -458,6 +459,48 @@ TEST_F(LayoutAssignmentTest, TopKLayout) {
                       .WithShape(F32, {6, 2048}, {1, 0}))));
 }
 
+TEST_F(LayoutAssignmentTest, TopKAfterAllGatherHasDefaultLayout) {
+  const char* hlo_text = R"(
+  HloModule topk
+
+  compare-greater-than {
+    p.1.lhs.3 = s32[] parameter(2)
+    p.1.rhs.4 = s32[] parameter(3)
+    p.0.lhs.1 = f32[] parameter(0)
+    p.0.rhs.2 = f32[] parameter(1)
+    ROOT compare.5 = pred[] compare(p.0.lhs.1, p.0.rhs.2), direction=GT
+  }
+
+  ENTRY main {
+    p0 = f32[32,64]{1,0} parameter(0)
+    ag = f32[32,512]{1,0} all-gather(p0), dimensions={1}, replica_groups={{0,1,2,3,4,5,6,7}}
+    neg = f32[32,512]{1,0} negate(ag)
+    topk = (f32[32,64]{1,0}, s32[32,64]{1,0}) custom-call(neg), custom_call_target="__gpu$TopK", api_version=API_VERSION_TYPED_FFI, called_computations={compare-greater-than}
+    gte0 = f32[32,64]{1,0} get-tuple-element(topk), index=0
+    gte1 = s32[32,64]{1,0} get-tuple-element(topk), index=1
+    ROOT out = (f32[32,64]{1,0}, s32[32,64]{1,0}) tuple(gte0, gte1)
+  })";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_text));
+
+  ComputationLayout computation_layout(
+      module->entry_computation()->ComputeProgramShape(),
+      /*ignore_layouts=*/false);
+  GpuLayoutAssignment layout_assignment(&computation_layout, default_gpu_cc_,
+                                        default_device_description_);
+
+  EXPECT_THAT(layout_assignment.Run(module.get()),
+              absl_testing::IsOkAndHolds(true));
+
+  const HloInstruction* topk = FindInstruction(module.get(), "topk");
+  ASSERT_NE(topk, nullptr);
+  EXPECT_EQ(topk->operand(0)->shape().layout(), LayoutUtil::MakeLayout({1, 0}));
+  EXPECT_EQ(ShapeUtil::GetSubshape(topk->shape(), {0}).layout(),
+            LayoutUtil::MakeLayout({1, 0}));
+  EXPECT_EQ(ShapeUtil::GetSubshape(topk->shape(), {1}).layout(),
+            LayoutUtil::MakeLayout({1, 0}));
+}
+
 TEST_F(LayoutAssignmentTest,
        BitcastConvertFromNarrowerTypeGetsOptimalInputLayout) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,

From c609620cfec9ae686f5af4db543d8f3951c5d329 Mon Sep 17 00:00:00 2001
From: Dmitri Latushko <latushko@google.com>
Date: Mon, 12 Oct 2026 11:05:40 -0700
Subject: [PATCH 0603/1000] Shard framework_run_handler_test across 4 shards to
 prevent timeouts

PiperOrigin-RevId: 998029596
---
 tensorflow/core/framework/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/core/framework/BUILD b/tensorflow/core/framework/BUILD
index 35785f02ee3..81e3b1ba576 100644
--- a/tensorflow/core/framework/BUILD
+++ b/tensorflow/core/framework/BUILD
@@ -1284,6 +1284,7 @@ tf_cc_test(
     name = "framework_run_handler_test",
     size = "medium",
     srcs = ["run_handler_test.cc"],
+    shard_count = 4,
     deps = [
         "//tensorflow/core:core_cpu",
         "//tensorflow/core:direct_session_internal",

From 74c375e3177e8ab603e682a4df210de2271e390a Mon Sep 17 00:00:00 2001
From: David Dunleavy <ddunleavy@google.com>
Date: Mon, 12 Oct 2026 11:07:18 -0700
Subject: [PATCH 0604/1000] Move `optimization_parameters.proto` and
 `tpu_embdding_configuration.proto` to `xla/tpu`

PiperOrigin-RevId: 998030821
---
 tensorflow/core/protobuf/tpu/BUILD            |  12 +-
 .../tpu/optimization_parameters.proto         | 641 +----------------
 .../tpu/tpu_embedding_configuration.proto     | 160 +----
 third_party/xla/xla/tpu/protobuf/BUILD        |  60 ++
 .../protobuf/optimization_parameters.proto    | 657 ++++++++++++++++++
 .../tpu_embedding_configuration.proto         | 176 +++++
 6 files changed, 907 insertions(+), 799 deletions(-)
 create mode 100644 third_party/xla/xla/tpu/protobuf/BUILD
 create mode 100644 third_party/xla/xla/tpu/protobuf/optimization_parameters.proto
 create mode 100644 third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration.proto

diff --git a/tensorflow/core/protobuf/tpu/BUILD b/tensorflow/core/protobuf/tpu/BUILD
index 35a385fdb85..08bf50b04d7 100644
--- a/tensorflow/core/protobuf/tpu/BUILD
+++ b/tensorflow/core/protobuf/tpu/BUILD
@@ -32,9 +32,12 @@ tf_proto_library(
         "tpu_embedding_configuration.proto",
     ],
     protodeps = [
-        ":optimization_parameters_proto",
+        "@xla//xla/tpu/protobuf:tpu_embedding_configuration_proto",
     ],
     visibility = ["//visibility:public"],
+    exports = [
+        "@xla//xla/tpu/protobuf:tpu_embedding_configuration_proto",
+    ],
 )
 
 tf_proto_library(
@@ -42,8 +45,13 @@ tf_proto_library(
     srcs = [
         "optimization_parameters.proto",
     ],
-    protodeps = ["@xla//xla/service:hlo_proto"],
+    protodeps = [
+        "@xla//xla/tpu/protobuf:optimization_parameters_proto",
+    ],
     visibility = ["//visibility:public"],
+    exports = [
+        "@xla//xla/tpu/protobuf:optimization_parameters_proto",
+    ],
 )
 
 tf_proto_library(
diff --git a/tensorflow/core/protobuf/tpu/optimization_parameters.proto b/tensorflow/core/protobuf/tpu/optimization_parameters.proto
index 6899a183f58..1688ab79739 100644
--- a/tensorflow/core/protobuf/tpu/optimization_parameters.proto
+++ b/tensorflow/core/protobuf/tpu/optimization_parameters.proto
@@ -15,643 +15,6 @@
 
 syntax = "proto3";
 
-package tensorflow.tpu;
+package tensorflow.tpu.dummy;
 
-import "google/protobuf/wrappers.proto";
-import "xla/service/hlo.proto";
-
-message ClippingLimits {
-  google.protobuf.FloatValue lower = 1;  // -inf if not set
-  google.protobuf.FloatValue upper = 2;  // +inf if not set
-}
-
-// Configuration for simulated quantization; simulated quantization is used to
-// reduce training/serving skew when the serving variables are quantized. The
-// same quantization operations are executed during training to minimize
-// differences with serving.
-//
-// Simulated quantization inserts the following operations on the forward pass
-// after gathering the embedding vector from HBM. The backward pass operations
-// are unchanged.
-//
-// clipped_val = clip(input, clipping_limits)
-// quantum = clipping_limits.range() / (num_buckets - 1)
-// quantized_val = floor((clipped_val - clipping_limits.lower()) / quantum + .5)
-// return quantized_val * quantum + clipping_limits.lower().
-message SimulatedQuantization {
-  // Whether simulated quantization is enabled.
-  bool enabled = 1;
-
-  // Minimum and maximum values of the range used for quantization.
-  ClippingLimits clipping_limits = 2;
-
-  // Number of possible quantized values.
-  int32 num_buckets = 3;
-}
-
-// Dynamic input specification for optimizers in the TPUEmbeddingConfiguration.
-// The actual dynamic inputs are provided as a scalar input list to the
-// SendTPUEmbeddingGradients Op indexed by their tag specified through the
-// following proto.
-message OptimizerDynamicInput {
-  // For tables where dynamic inputs are needed (e.g., learning rates or other
-  // dynamic hyperparameters used in optimizers), a tag must be specified for
-  // the input.
-  //
-  // The tag must be a non-negative  integer. All tags in the range
-  // [0, number_of_unique_tags) must be present in the TPU embedding
-  // configuration, i.e. a tag cannot be skipped if a different tag numerically
-  // greater than it is used in the configuration.
-  //
-  // If multiple tables specify the same tag, they *MUST* have
-  // the same dynamic input, for example, their dynamic learning rate could be
-  // computed by the same TensorFlow sub-graph. The partitioning of the
-  // embedding layer would be more optimal if the number_of_unique_tags is as
-  // *LOW* as possible, i.e., if many tables share the same tag.
-  //
-  // The hyper_parameters input of the SendTPUEmbeddingGradients op is used to
-  // communicate dynamic hyper-parameters to the TPU embedding program.
-  // The hyper_parameters input is a list of scalars where the size of the list
-  // is equal to the number of unique tags. The hyper-parameter associated with
-  // a particular tag is specified by populating its corresponding index in the
-  // list of scalars.
-  int32 tag = 1;
-}
-
-// Source of learning rate to use.
-message LearningRate {
-  oneof learning_rate {
-    float constant = 1;
-    OptimizerDynamicInput dynamic = 2;
-  }
-}
-
-// Each optimizer's parameter proto has a link to its documentation and CPU
-// implementation (if available) for user reference.
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adagrad
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L1634
-message AdagradParameters {
-  // Old initial accumulator parameter.
-  reserved "initial_accumulator";
-  reserved 1;
-}
-
-// This optimizer combines the Adagrad and Momentum update rules.
-// accum(new) = beta2 == 1.0 ?
-//              accum(old) + grad^2 :
-//              beta2 * accum(old) + (1 - beta2) * grad^2
-// accum_with_exponent = (accum(new) + epsilon)^(-1.0 / exponent)
-// mom_accum(new) = momentum * mom_accum(old) + accum_with_exponent
-// update = use_nesterov ?
-//          momentum * mom_accum(new) + accum_with_exponent :
-//          mom_accum(new)
-// var(new) = var(old) - lr * grad * update
-// Algorithm described in https://arxiv.org/abs/2002.11803.
-message AdagradMomentumParameters {
-  // Moving average parameter for the momentum accumulator.
-  float momentum = 1;
-
-  // Whether to use the Nesterov variant of momentum.
-  bool use_nesterov = 2;
-
-  // Exponent for the gradient^2 accumulator.
-  float exponent = 3;
-
-  // Moving average parameter for the gradient^2 accumulator.
-  float beta2 = 4;
-
-  // Offset added to the Adagrad accumulator.
-  float epsilon = 5;
-}
-
-// Algorithm in http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf.
-message BoundedAdagradParameters {
-  // Whether to use the updated or the old value of the accumulator when
-  // computing the effective learning rate. When update_accumulator_first is set
-  // to True, the updated value of the accumulator is used.
-  bool update_accumulator_first = 1;
-
-  // The max_var_update value to use. Set value to 0 (default) to disable using
-  // max_var_update to clip the gradient.
-  float max_var_update = 2;
-
-  // The maximum value of the accumulator. Set max_accumulator to 0 (default)
-  // to disable using max_accumulator to clip the accumulator.
-  float max_accumulator = 3;
-}
-
-// Frequency Aware Adagrad optimizer. This optimizer implements the AdaGrad
-// algorithm and further allows to:
-// * Scale the learning rate based on frequency of the update. Sparsely updated
-//   rows are updated with a higher effective learning rate, and frequently
-//   updated rows are updated with a lower effective learning rate.
-// * Decay the growth of the accumulator values.
-// * Use L1 / L2 regularization for the weight updates.
-//
-// The optimization algorithm is shown below.
-// counter(new) = counter(old) + 1
-// accum(new) = max(accumulator_decay * accum(old) + grad^2,
-//                  initial_accumulator_value)
-// lr_scale = min((step_counter / accum(new)) ^ probability_exponent,
-// max_lr_multiplier) update = grad * lr_scale / sqrt(accum(new)) if
-// (l1_regularization_strength > 0.0):
-//   update = update + l1_regularization_strength * sign(var(old))
-// if (l2_regularization_strength > 0.0):
-//   update = update + l2_regularization_strength * var(old)
-// var(new) = var(old) - lr_scale * grad * update
-
-message FrequencyAwareAdagradParameters {
-  // The L1 regularization parameter for adjusting the update based on the sign
-  // of the variable.
-  float l1_regularization_strength = 1;
-
-  // The L2 regularization parameter for adjusting the update based on the
-  // variable.
-  float l2_regularization_strength = 2;
-
-  // The exponent used for scaling the learning rate based on the sparsity of
-  // updates.
-  float probability_exponent = 4;
-
-  // The maximum value of the learning rate scale.
-  float max_lr_multiplier = 3;
-
-  // The decay for the Adagrad accumulator.
-  float accumulator_decay = 5;
-
-  // The initial and minimum value for the Adagrad accumulator.
-  float initial_accumulator_value = 6;
-
-  // The tag for identifying the step counter used for the frequency aware
-  // Adagrad optimizer.
-  OptimizerDynamicInput step_counter = 7;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L629
-message StochasticGradientDescentParameters {}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Ftrl
-// https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/41159.pdf
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L2646
-//
-// The hyperparameters for FTRL are the same as for the Keras implementation,
-// with some additions. The "beta" parameter matches the behavior described in
-// the second link above; "beta" / (2 * learning rate) should be added to "l2"
-// to get equivalent behavior in the other TensorFlow implementations of this
-// optimizer. When the multiply_linear_by_lr field is set to true, a modified
-// formula is used for FTRL that treats the "linear" accumulator as being
-// pre-multiplied by the learning rate (i.e., the accumulator named "linear"
-// actually stores "linear * learning_rate"). Other than checkpoint
-// compatibility, this is mathematically equivalent for a static learning rate;
-// for a dynamic learning rate, it is nearly the same as long as the learning
-// rate does not change quickly. The benefit of setting multiply_linear_by_lr to
-// true is that the modified formula handles zero and near-zero learning rates
-// without producing NaNs, improving flexibility for learning rate ramp-up.
-message FtrlParameters {
-  float l1 = 1;
-  float l2 = 2;
-  float lr_power = 3;
-  float beta = 7;
-  bool multiply_linear_by_lr = 6;
-
-  // Previously, allow_zero_accumulator parameter changed some internal formulas
-  // to allow zero and near-zero accumulator values at the cost of some
-  // performance. The current implementation ignores this parameter; zero or
-  // near-zero accumulator values are now always supported.
-  bool allow_zero_accumulator = 8 [deprecated = true];
-
-  // Old initial accumulator parameters.
-  reserved "initial_accum", "initial_linear";
-  reserved 4, 5;
-}
-
-// The Adam optimizer does not implement hyper-parameter update due to hardware
-// limitations; use the dynamic learning rate feature instead, setting the
-// learning rate to: user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
-// Here, t is the current timestep.
-//
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adam
-// https://github.com/tensorflow/tensorflow/blob/ab51450c817674c8ff08a7ae4f8ac50cdc4bed8b/tensorflow/python/training/adam.py#L32
-//
-// Note that the code by default implements the lazy version of Adam
-// (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/LazyAdamOptimizer)
-// unless the use_non_lazy_adam parameter is set, in which case it implements
-// the normal version of Adam that updates all parameters in the embedding
-// table, even for entries that are not used in the current minibatch
-// (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/AdamOptimizer). If
-// use_non_lazy_adam is enabled, gradient accumulation is also required to be
-// enabled in order to get correct results; a warning will be printed otherwise
-// (which may change to an error in the future). If use_sum_inside_sqrt is set,
-// the Adam variable update formula will be changed from m / (sqrt(v) + epsilon)
-// to m / sqrt(v + epsilon**2); this option improves the performance of TPU
-// training and is not expected to harm model quality.
-message AdamParameters {
-  float beta1 = 3;
-  float beta2 = 4;
-  float epsilon = 5;
-  bool use_non_lazy_adam = 8;
-  bool use_sum_inside_sqrt = 10;
-
-  // Old initial accumulator parameters.
-  reserved "initial_m", "initial_v";
-  reserved 6, 7;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L3068
-message MomentumParameters {
-  float momentum = 1;
-  bool use_nesterov = 2;
-
-  // Old initial accumulator parameter.
-  reserved "initial_accum";
-  reserved 3;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Lion
-// momenta(new) = beta2 * momenta(old) + (1 - beta2) * grad
-// momenta_t = beta1 * momenta(old) + (1 - beta1) * grad
-// var(new) = var(old) - lr * sign(momenta_t)
-// Algorithm described in https://arxiv.org/abs/2302.06675.
-message LionParameters {
-  float beta1 = 1;
-  float beta2 = 2;
-  bool use_non_lazy_lion = 3;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L4229
-message RmsPropParameters {
-  float rho = 1;
-  float momentum = 2;
-  float epsilon = 3;
-
-  // Old initial accumulator parameters.
-  reserved "initial_ms", "initial_mom";
-  reserved 4, 5;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L4358
-message CenteredRmsPropParameters {
-  float rho = 1;
-  float momentum = 2;
-  float epsilon = 3;
-
-  // Old initial accumulator parameters.
-  reserved "initial_ms", "initial_mom", "initial_mg";
-  reserved 4, 5, 6;
-}
-
-// Variant of algorithm in http://proceedings.mlr.press/v44/shamir15.pdf
-message MdlAdagradLightParameters {
-  float l2 = 1;
-  float lr_power = 2;
-  float min_servable_mdl_benefit = 3;
-  float mdl_mix_in_margin = 4;
-  float mdl_benefit_rampup_coeff = 5;
-  float mdl_min_weight = 6;
-  float benefit_revisit_scale = 7;
-  float max_event_benefit = 8;
-  float max_total_benefit = 9;
-  float mdl_hard_limit = 10;
-  bool hard_limit_min_benefit = 11;
-  bool mdl_regularize = 12;
-
-  // Old initial accumulator parameters.
-  reserved "initial_accumulator", "initial_weight", "initial_benefit";
-  reserved 13, 14, 15;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adadelta
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L933
-message AdadeltaParameters {
-  float rho = 1;
-  float epsilon = 2;
-
-  // Old initial accumulator parameters.
-  reserved "initial_accumulator", "initial_update";
-  reserved 3, 4;
-}
-
-// https://www.tensorflow.org/api_docs/python/tf/compat/v1/train/ProximalAdagradOptimizer
-// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L1961
-message ProximalAdagradParameters {
-  float l1 = 1;
-  float l2 = 2;
-
-  // Old initial accumulator parameter.
-  reserved "initial_accumulator";
-  reserved 3;
-}
-
-// The online Yogi optimizer does not implement hyper-parameter update; use the
-// dynamic learning rate feature instead, setting the learning rate to:
-// user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
-// Here, t is the current timestep.
-//
-// https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf
-// plus some extensions based on FTRL.
-//
-// Note that the code by default implements the lazy version of online Yogi.
-message OnlineYogiParameters {
-  // The L1 regularization parameter (used analogously to the one in FTRL).
-  float l1 = 1;
-
-  // The L2 regularization parameter (used analogously to the one in FTRL).
-  float l2 = 2;
-
-  // \beta_2 from Algorithm 2 in the paper.
-  float beta2 = 3;
-
-  // Reserved ids corresponding to removed tanh activation.
-  reserved 6;  // sign
-  reserved 7;  // tanh
-}
-
-// The online Yogi optimizer does not implement hyper-parameter update; use the
-// dynamic learning rate feature instead, setting the learning rate to:
-// user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
-// Here, t is the current timestep.
-//
-// https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf
-// plus some extensions based on FTRL.
-//
-// Note that the code by default implements the lazy version of proximal Yogi.
-message ProximalYogiParameters {
-  // The L1 regularization parameter.
-  float l1 = 1;
-
-  // The L2 regularization parameter.
-  float l2 = 2;
-
-  // The exponential decay rate for the 1st moment estimates.
-  float beta1 = 3;
-
-  // The exponential decay rate for the 2nd moment estimates.
-  float beta2 = 4;
-
-  // A constant trading off adaptivity and noise.
-  float epsilon = 5;
-
-  // Reserved ids corresponding to removed tanh activation.
-  reserved 8;  // sign
-  reserved 9;  // tanh
-}
-
-// Estimator for the frequency of updates to a lookup table. It maintains an
-// array (tf.Variable) D, where each element records the average number of
-// global steps between two consecutive batches that hit the corresponding
-// bucket. Once an item with bucket id i is sampled, D[i] is updated by:
-//   D[i] <- D[i] * (1 - tau) + delta[i] * tau,
-//
-// where tau is a learning rate between 0 and 1 (exclusive), and
-//   delta[i] = current global step - last step i is sampled.
-//
-// The estimated frequency (sampling rate in a batch) is thus 1 / D[i].
-//
-// Elements in D are initialized with a large value max_delta. delta[i] will
-// also be capped by this value.
-//
-// The exact sequence of operations used in the optimizer is shown below.
-// last_hit_step[i] is a tf.Variable that holds the last global step at which i
-// was sampled.
-//
-//   delta = global_step - last_hit_step[i]
-//   clipped_delta = min(delta, params.max_delta)
-//   is_outlier = (delta >= params.outlier_threshold * D[i])
-//   D[i] <- is_outlier ? clipped_delta
-//                      : D[i] * (1 - params.tau) + clipped_delta * params.tau
-//   last_hit_step[i] <- global_step
-message FrequencyEstimatorParameters {
-  // Learning rate between (0, 1) that is used to update the array D.
-  float tau = 1;
-
-  // Maximum value of delta: difference between the current global step and the
-  // last global step at which the row was sampled.
-  float max_delta = 2;
-
-  // Threshold used to determine whether the current update is an outlier.
-  float outlier_threshold = 3;
-
-  // The weight exponent used to transform the estimated delta into weights.
-  // The transformation function is: (delta / max_delta) ^ (weight_exponent)
-  float weight_exponent = 4;
-}
-
-// A user-defined optimizer.
-// The contained HLO program must take the following arguments in the following
-// order:
-// 1.  gradients
-// 2.  table weights
-// 3.  slot variables
-// 4.  an optional scalar input that is passed in via the dynamic learning
-//     rate mechanism.
-//
-// It must return/end in a tuple op that contains the following values in the
-// following order:
-// 1.  new table values
-// 2.  new slot variable value
-//
-// The program must have shape (1,1) with dtype float32 throughout and only use
-// HLO that operate elementwise (e.g., no reduce, no variables, no control flow
-// and no broadcasting outside of the single scalar input).
-// The HLO program should be written as if it were a dense update. It will be
-// called on each row that needs an update and will applied elementwise.
-message UserDefinedProgramParameters {
-  xla.HloModuleProto program = 1;
-  reserved 2;  // Was padding_values
-}
-
-// Optimizer that just sets the variable to the value of the gradient. To be
-// correct, this requires either gradient accumulation (to sum the values of a
-// computed expression across the samples) or to deduplicate IDs within a single
-// host (to assign the value from an arbitrary sample).
-message AssignParameters {}
-
-// Status of using gradient accumulation (doing two passes over the input
-// gradients: one to accumulate them into a temporary array and another to apply
-// them using the actual optimization algorithm). The extra message is to wrap
-// the enum for scoping.
-message GradientAccumulationStatus {
-  // if UNSPECIFIED (default), gradient accumulation is ENABLED.
-  enum Status {
-    UNSPECIFIED = 0;
-    ENABLED = 1;
-    DISABLED = 2;
-  }
-}
-
-// Whether to optimize the packing of low-dimensional embedding tables in HBM
-// (high bandwidth memory). TPUs access HBM at 32-byte (8-float) granularity.
-// For functional correctness, the TPU software internally pads the embedding
-// dimension to a multiple of 8. This can sometimes lead to significant memory
-// wastage due to padding. For 1-dimensional, 2-dimensional, and 4-dimensional,
-// the TPU software can remove this padding by packing multiple rows into the
-// same 8-float HBM chunk. For example, 8 rows could be packed into the same
-// 8-float chunk for a 1-dimensional embedding table.
-
-// There is one important limitation for this HBM packing though. When only a
-// subset of rows in an 8-float chunk are accessed on a particular step, the
-// adjoining rows in the same chunk are updated with zero gradients on the
-// backward pass even if they are not touched. This is an artifact of the
-// packing implementation. This operation is NOT functionally correct for
-// optimizers where zero gradients change the embeddings/slot-variable values,
-// e.g., momentum-based optimizers. Hence, this HBM packing cannot be enabled
-// for embedding tables with such optimizers. The TPU software automatically
-// recognizes that a zero gradient can modify state and turns off the low
-// dimensional embedding packing in that scenario.
-//
-// However, for optimizers where a zero gradient is a NoOp, such as SGD,
-// Adagrad, and FTRL, this packing optimization can be used. However, there are
-// some important considerations:
-// * Clipping limits: The initial values for such embeddings should fall within
-//   the clipping limits specified in the optimization parameters. Otherwise, a
-//   zero gradient will cause the embeddings to be clipped. This changes state
-//   and hence, is not a NoOp.
-// * FTRL: The embedding vector is computed directly from the values of the
-//   accumulator and linear slot variables. Hence, the initial embedding values
-//   should match that computed from the initial values of the accumulator and
-//   linear slot variables. Note that in nearly all cases, the linear value is
-//   initialized to zero; this corresponds to an embedding value of zero.
-//
-// Performance: The TPU has to perform additional work when low dimensional
-// packing is enabled. In certain situations when the vocabulary size is small,
-// it may not make sense to turn on this packing since the total memory usage
-// due to padding is extremely low. Hence, the TPU software automatically turns
-// off the packing optimization in such scenarios.
-message LowDimensionalPackingStatus {
-  // if UNSPECIFIED (default), the low dimension packing status is DISABLED.
-  // This can change in future.
-  //
-  // if ENABLED, the low dimension packing is enabled only if the following
-  // three additional conditions are true:
-  //  * The optimizer treats the zero gradient as a NoOp.
-  //  * The embedding dimension is 1, 2, or 4.
-  //  * The vocabulary size is large enough to avoid performance issues.
-  //
-  // if DISABLED, the low dimension packing is always disabled.
-  enum Status {
-    UNSPECIFIED = 0;
-    ENABLED = 1;
-    DISABLED = 2;
-  }
-}
-
-// Configuration proto for hot ID optimization. This is an experimental feature
-// that is currently disabled (by default).
-message HotIdReplicationConfiguration {
-  // Whether to enable or disable hot ID optimization.
-  // If set to UNSPECIFIED (default), hot ID optimization is DISABLED.
-  // If set to ENABLED, hot ID replication is turned ON.
-  // If set to MIGRATION_ONLY, hot ID migration is turned ON.
-  enum Status {
-    UNSPECIFIED = 0;
-    ENABLED = 1;
-    DISABLED = 2;
-    MIGRATION_ONLY = 3;
-  }
-  Status status = 1;
-}
-
-message OptimizationParameters {
-  // Learning rate used for updating the embedding layer parameters.
-  LearningRate learning_rate = 13;
-
-  // Limits to which to clip the weight values after the backward pass; not
-  // present means no limits are applied.
-  ClippingLimits clipping_limits = 2;
-
-  // Limits to which to clip the backward pass gradient before using it for
-  // updates; not present means no limits are applied.
-  ClippingLimits gradient_clipping_limits = 7;
-
-  // Amount of weight decay to apply; see weight_decay_optimizers.py for
-  // details. All optimizers except MDL Adagrad Light are supported with this
-  // option. Although there is no check, users who want weight decay will also
-  // want to ensure that gradient accumulation is enabled so that the decay will
-  // happen once per global batch.
-  float weight_decay_factor = 16;
-
-  // If true, the weight decay factor is multiplied by the current learning rate
-  // before use; this is to match the note in DecoupledWeightDecayExtension in
-  // weight_decay_optimizers.py.
-  bool multiply_weight_decay_factor_by_learning_rate = 22;
-
-  // Configuration for simulated quantization which is used to reduce
-  // training/serving skew when the serving variables are quantized. The same
-  // quantization operations are executed during training to minimize
-  // differences with serving.
-  SimulatedQuantization simulated_quantization = 27;
-
-  // Status of using gradient accumulation (doing two passes over the input
-  // gradients: one to accumulate them into a temporary array and another to
-  // apply them using the actual optimization algorithm).
-  GradientAccumulationStatus.Status gradient_accumulation_status = 17;
-
-  // Status of the low-dimensional embedding packing optimization. This controls
-  // whether to optimize the packing of 1-dimensional, 2-dimensional, and
-  // 4-dimensional embedding tables in memory.
-  LowDimensionalPackingStatus.Status low_dimensional_packing_status = 28;
-
-  // Configuration proto for hot ID replication. This is an experimental
-  // feature that is currently disabled (by default).
-  HotIdReplicationConfiguration hot_id_replication_configuration = 18;
-
-  // Optimization algorithm parameters; which field is selected determines which
-  // algorithm to use.
-  oneof parameters {
-    AdagradParameters adagrad = 3;
-    AdagradMomentumParameters adagrad_momentum = 26;
-    BoundedAdagradParameters bounded_adagrad = 19;
-    FrequencyAwareAdagradParameters frequency_aware_adagrad = 30;
-    StochasticGradientDescentParameters stochastic_gradient_descent = 4;
-    FtrlParameters ftrl = 5;
-    AdamParameters adam = 6;
-    MomentumParameters momentum = 8;
-    LionParameters lion = 29;
-    RmsPropParameters rms_prop = 9;
-    CenteredRmsPropParameters centered_rms_prop = 10;
-    MdlAdagradLightParameters mdl_adagrad_light = 11;
-    AdadeltaParameters adadelta = 12;
-    ProximalAdagradParameters proximal_adagrad = 14;
-    OnlineYogiParameters online_yogi = 20;
-    ProximalYogiParameters proximal_yogi = 21;
-    FrequencyEstimatorParameters frequency_estimator = 23;
-    UserDefinedProgramParameters user_defined_program = 24;
-    AssignParameters assign = 25;
-  }
-
-  reserved 1, 15;
-
-  // NEXT_ID: 31
-}
-
-// Specification of an optimization algorithm's state variables (both the main
-// value vector and any extra accumulators, etc.). This proto is only used
-// internally by the TPU software and is not exposed directly to the TF model.
-message StateVariableSpecification {
-  // Parameter name for the state variable.
-  string name = 1;
-
-  // A normal state variable that should be saved and restored in checkpoints
-  // and used as an input or output to non-debug TensorFlow ops.
-  message UserDefined {
-    reserved 1;  // Was padding_initial_value.
-  }
-
-  // A state variable that should be filled with a constant and normally hidden
-  // from users (used for intermediate gradients being accumulated, for
-  // example).
-  message FillWithConstant {
-    double initial_value = 1;
-  }
-
-  // Usage type of this state variable.
-  oneof usage {
-    UserDefined user_defined = 2;
-    FillWithConstant fill_with_constant = 3;
-  }
-}
+import public "xla/tpu/protobuf/optimization_parameters.proto";
diff --git a/tensorflow/core/protobuf/tpu/tpu_embedding_configuration.proto b/tensorflow/core/protobuf/tpu/tpu_embedding_configuration.proto
index c11ece0c6f4..5130a704ba5 100644
--- a/tensorflow/core/protobuf/tpu/tpu_embedding_configuration.proto
+++ b/tensorflow/core/protobuf/tpu/tpu_embedding_configuration.proto
@@ -15,162 +15,6 @@
 
 syntax = "proto3";
 
-package tensorflow.tpu;
+package tensorflow.tpu.dummy;
 
-import "tensorflow/core/protobuf/tpu/optimization_parameters.proto";
-
-message TPUEmbeddingConfiguration {
-  // Description of the various embedding tables.
-  message TableDescriptor {
-    // Name of the table.
-    string name = 1;
-
-    // Size of the vocabulary (i.e., number of rows) in the table.
-    int64 vocabulary_size = 2;
-
-    // The embedding dimension (i.e., the width of the embedding table).
-    int32 dimension = 3;
-
-    // Number of features mapped to this table.
-    int32 num_features = 4;
-
-    // Details of the learning algorithm used to update the embedding
-    // parameters.
-    OptimizationParameters optimization_parameters = 5;
-  }
-  repeated TableDescriptor table_descriptor = 1;
-
-  // Mode. Should the embedding layer program be run for inference (just forward
-  // pass), training (both forward and backward pass) or just the backward_pass.
-  enum Mode {
-    UNSPECIFIED = 0;
-    INFERENCE = 1;
-    TRAINING = 2;
-    BACKWARD_PASS_ONLY = 3;
-  }
-  Mode mode = 2;
-
-  // Number of samples in each batch of embedding layer activations sent to
-  // the TensorCore.
-  int32 batch_size_per_tensor_core = 3;
-
-  // Number of TPU hosts used for inference/training.
-  int32 num_hosts = 4;
-
-  // Number of TensorCore used for inference/training.
-  int32 num_tensor_cores = 5;
-
-  // Sharding strategy of the embedding tables among the hosts.
-  // If the sharding_strategy is "mod", each id is assigned to host
-  // "id % num_hosts". For instance, 13 ids are split across 5 hosts as:
-  // [[0, 5, 10], [1, 6, 11], [2, 7, 12], [3, 8], [4, 9]].
-  // If the sharding_strategy is "div", ids are assigned to hosts in a
-  // contiguous manner. In this case, 13 ids are split across 5 hosts as:
-  // [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11, 12]].
-  // In both the strategies, if the id space does not evenly divide the number
-  // of hosts, each of the first "table_descriptor.vocabulary_size % num_hosts"
-  // hosts will be assigned one more id.
-  // This partitioning strategy exactly follows that in the embedding_lookup
-  // TensorFlow function at tensorflow/python/ops/embedding_ops.py.
-  enum ShardingStrategy {
-    DIV_DEFAULT = 0;
-    MOD = 1;
-  }
-  ShardingStrategy sharding_strategy = 6;
-
-  // This parameter determines if the execution of the sparse core will be
-  // pipelined with that of the TensorCore. This parameter only affects results
-  // when mode=TRAINING. If mode=INFERENCE or BACKWARD_PASS_ONLY, this parameter
-  // does not affect execution and hence, is a don't care value.
-  //
-  // false: The execution of the sparse core is not pipelined with that of the
-  // TensorCore. The forward pass of every step on the sparse core is executed
-  // only after the backward pass of the previous step is complete. And the
-  // backward pass on the sparse core is executed only after the embedding
-  // gradients have been computed on the TensorCore on every step. This ensures
-  // that the activations on every step observe the gradient updates from the
-  // previous step on both the sparse core and the TensorCore.
-  //
-  // true: The execution of the sparse core is pipelined with that of the
-  // TensorCore. The forward pass of every step on the sparse core can be
-  // executed after the forward pass of the previous step is complete without
-  // waiting for the backward pass. This improves the utilization of the sparse
-  // core allowing it to process step N+1 while the embedding gradients for step
-  // N are computed on the TensorCore. The backward pass of every step on the
-  // sparse core is executed directly after the forward pass for the next step
-  // is complete. The drawback is that embedding activations for step N+1 do not
-  // observe the embedding gradient updates from step N. This could affect model
-  // quality if step N and N+1 involve the same set of embedding IDs. However,
-  // since the embedding updates are sparse, this is generally not considered a
-  // problem.
-  bool pipeline_execution_with_tensor_core = 7;
-
-  // Directory where embedding lookup statistics are stored. These statistics
-  // summarize information about the inputs to the embedding lookup
-  // operation, in particular, the average number of embedding IDs per example
-  // and how well the embedding IDs are load balanced across the system. The
-  // lookup statistics are used during TPU initialization for embedding table
-  // partitioning. Collection of lookup statistics is done at runtime by
-  // profiling the embedding inputs: only 3% of input samples are profiled to
-  // minimize host CPU overhead. Once a suitable number of samples are
-  // profiled, the lookup statistics are saved to table-specific files in the
-  // profile data directory generally at the end of a TPU training loop. The
-  // filename corresponding to each table is obtained by hashing table specific
-  // parameters (e.g., table name and number of features) and global
-  // configuration parameters (e.g., sharding strategy and TPU worker task
-  // count). The same profile data directory can be shared amongst several
-  // models to reuse embedding lookup statistics.
-  string profile_data_directory = 9;
-
-  // Description of different input features.
-  message FeatureDescriptor {
-    // Name of the input feature.
-    string name = 1;
-
-    // Index of the corresponding table in the TableDescriptor list.
-    int32 table_id = 2;
-
-    // Static shape of the inputs (excluding the reduction axis). Note that
-    // the shape of the actual inputs provided using the infeed op must be
-    // strictly smaller than input_shape. The outputs received at the TensorCore
-    // will have rank = input_shape.size() + 1. The innermost axis corresponds
-    // to the embedding dimension. If the input has shape [m, n, k] (excluding
-    // the reduction axis) and the embedding dimension is d, the output received
-    // at the TensorCore will have shape [m, n, k, d].
-    repeated int32 input_shape = 3;
-  }
-  // If the feature_descriptor field is populated, the model should NOT populate
-  // TableDescriptor.num_features and batch_size_per_tensor_core. These two
-  // fields will be auto-populated by the TPUEmbedding rewrite passes.
-  repeated FeatureDescriptor feature_descriptor = 10;
-
-  // SPMD (Single Program Multiple Data) sharding configuration for
-  // TPUEmbedding. When model parallelism is used on the TensorCore, the number
-  // of cores per replica must be passed to TPUEmbedding so that the right
-  // shapes can be computed in the TF/XLA bridge.
-  message SpmdSharding {
-    // Whether SPMD sharding is enabled.
-    bool enabled = 1;
-
-    // Number of cores per replica.
-    int32 num_cores_per_replica = 2;
-
-    // If true, the tensors are manually partitioned. Otherwise, use the
-    // automatic SPMD partitioning. This should be true when users use
-    // `shard_map`.
-    bool use_manual_partitioning = 3;
-
-    // If true, lowers to Shardy custom-call for manual partitioning.
-    // Default is false.
-    bool use_shardy_partitioner = 4;
-  }
-  SpmdSharding spmd_sharding = 11;
-
-  // Old TPU embedding output layout.
-  reserved "output_layout";
-  reserved 8;
-}
-
-// A placeholder message that is used to define a unique Status payload
-// URL for TPU embedding errors.
-message TPUEmbeddingError {}
+import public "xla/tpu/protobuf/tpu_embedding_configuration.proto";
diff --git a/third_party/xla/xla/tpu/protobuf/BUILD b/third_party/xla/xla/tpu/protobuf/BUILD
new file mode 100644
index 00000000000..f732b3e19e7
--- /dev/null
+++ b/third_party/xla/xla/tpu/protobuf/BUILD
@@ -0,0 +1,60 @@
+# Copyright 2026 The OpenXLA 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.
+# ==============================================================================
+
+# copybara:uncomment load("@rules_python//python:proto.bzl", "py_proto_library")
+load(
+    "//xla/tsl/platform:build_config.bzl",
+    "tf_proto_library",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+tf_proto_library(
+    name = "tpu_embedding_configuration_proto",
+    srcs = [
+        "tpu_embedding_configuration.proto",
+    ],
+    protodeps = [
+        ":optimization_parameters_proto",
+    ],
+    visibility = ["//visibility:public"],
+)
+
+tf_proto_library(
+    name = "optimization_parameters_proto",
+    srcs = [
+        "optimization_parameters.proto",
+    ],
+    protodeps = ["//xla/service:hlo_proto"],
+    visibility = ["//visibility:public"],
+    deps = ["@com_google_protobuf//:wrappers_proto"],
+)
+
+# copybara:uncomment_begin(google-only)
+# py_proto_library(
+#     name = "tpu_embedding_configuration_py_pb2",
+#     visibility = ["//visibility:public"],
+#     deps = [":tpu_embedding_configuration_proto"],
+# )
+#
+# py_proto_library(
+#     name = "optimization_parameters_py_pb2",
+#     visibility = ["//visibility:public"],
+#     deps = [":optimization_parameters_proto"],
+# )
+# copybara:uncomment_end
diff --git a/third_party/xla/xla/tpu/protobuf/optimization_parameters.proto b/third_party/xla/xla/tpu/protobuf/optimization_parameters.proto
new file mode 100644
index 00000000000..424110f3d89
--- /dev/null
+++ b/third_party/xla/xla/tpu/protobuf/optimization_parameters.proto
@@ -0,0 +1,657 @@
+// Copyright 2026 The OpenXLA 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.
+// ==============================================================================
+
+syntax = "proto3";
+
+package tensorflow.tpu;
+
+import "google/protobuf/wrappers.proto";
+import "xla/service/hlo.proto";
+
+message ClippingLimits {
+  google.protobuf.FloatValue lower = 1;  // -inf if not set
+  google.protobuf.FloatValue upper = 2;  // +inf if not set
+}
+
+// Configuration for simulated quantization; simulated quantization is used to
+// reduce training/serving skew when the serving variables are quantized. The
+// same quantization operations are executed during training to minimize
+// differences with serving.
+//
+// Simulated quantization inserts the following operations on the forward pass
+// after gathering the embedding vector from HBM. The backward pass operations
+// are unchanged.
+//
+// clipped_val = clip(input, clipping_limits)
+// quantum = clipping_limits.range() / (num_buckets - 1)
+// quantized_val = floor((clipped_val - clipping_limits.lower()) / quantum + .5)
+// return quantized_val * quantum + clipping_limits.lower().
+message SimulatedQuantization {
+  // Whether simulated quantization is enabled.
+  bool enabled = 1;
+
+  // Minimum and maximum values of the range used for quantization.
+  ClippingLimits clipping_limits = 2;
+
+  // Number of possible quantized values.
+  int32 num_buckets = 3;
+}
+
+// Dynamic input specification for optimizers in the TPUEmbeddingConfiguration.
+// The actual dynamic inputs are provided as a scalar input list to the
+// SendTPUEmbeddingGradients Op indexed by their tag specified through the
+// following proto.
+message OptimizerDynamicInput {
+  // For tables where dynamic inputs are needed (e.g., learning rates or other
+  // dynamic hyperparameters used in optimizers), a tag must be specified for
+  // the input.
+  //
+  // The tag must be a non-negative  integer. All tags in the range
+  // [0, number_of_unique_tags) must be present in the TPU embedding
+  // configuration, i.e. a tag cannot be skipped if a different tag numerically
+  // greater than it is used in the configuration.
+  //
+  // If multiple tables specify the same tag, they *MUST* have
+  // the same dynamic input, for example, their dynamic learning rate could be
+  // computed by the same TensorFlow sub-graph. The partitioning of the
+  // embedding layer would be more optimal if the number_of_unique_tags is as
+  // *LOW* as possible, i.e., if many tables share the same tag.
+  //
+  // The hyper_parameters input of the SendTPUEmbeddingGradients op is used to
+  // communicate dynamic hyper-parameters to the TPU embedding program.
+  // The hyper_parameters input is a list of scalars where the size of the list
+  // is equal to the number of unique tags. The hyper-parameter associated with
+  // a particular tag is specified by populating its corresponding index in the
+  // list of scalars.
+  int32 tag = 1;
+}
+
+// Source of learning rate to use.
+message LearningRate {
+  oneof learning_rate {
+    float constant = 1;
+    OptimizerDynamicInput dynamic = 2;
+  }
+}
+
+// Each optimizer's parameter proto has a link to its documentation and CPU
+// implementation (if available) for user reference.
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adagrad
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L1634
+message AdagradParameters {
+  // Old initial accumulator parameter.
+  reserved "initial_accumulator";
+  reserved 1;
+}
+
+// This optimizer combines the Adagrad and Momentum update rules.
+// accum(new) = beta2 == 1.0 ?
+//              accum(old) + grad^2 :
+//              beta2 * accum(old) + (1 - beta2) * grad^2
+// accum_with_exponent = (accum(new) + epsilon)^(-1.0 / exponent)
+// mom_accum(new) = momentum * mom_accum(old) + accum_with_exponent
+// update = use_nesterov ?
+//          momentum * mom_accum(new) + accum_with_exponent :
+//          mom_accum(new)
+// var(new) = var(old) - lr * grad * update
+// Algorithm described in https://arxiv.org/abs/2002.11803.
+message AdagradMomentumParameters {
+  // Moving average parameter for the momentum accumulator.
+  float momentum = 1;
+
+  // Whether to use the Nesterov variant of momentum.
+  bool use_nesterov = 2;
+
+  // Exponent for the gradient^2 accumulator.
+  float exponent = 3;
+
+  // Moving average parameter for the gradient^2 accumulator.
+  float beta2 = 4;
+
+  // Offset added to the Adagrad accumulator.
+  float epsilon = 5;
+}
+
+// Algorithm in http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf.
+message BoundedAdagradParameters {
+  // Whether to use the updated or the old value of the accumulator when
+  // computing the effective learning rate. When update_accumulator_first is set
+  // to True, the updated value of the accumulator is used.
+  bool update_accumulator_first = 1;
+
+  // The max_var_update value to use. Set value to 0 (default) to disable using
+  // max_var_update to clip the gradient.
+  float max_var_update = 2;
+
+  // The maximum value of the accumulator. Set max_accumulator to 0 (default)
+  // to disable using max_accumulator to clip the accumulator.
+  float max_accumulator = 3;
+}
+
+// Frequency Aware Adagrad optimizer. This optimizer implements the AdaGrad
+// algorithm and further allows to:
+// * Scale the learning rate based on frequency of the update. Sparsely updated
+//   rows are updated with a higher effective learning rate, and frequently
+//   updated rows are updated with a lower effective learning rate.
+// * Decay the growth of the accumulator values.
+// * Use L1 / L2 regularization for the weight updates.
+//
+// The optimization algorithm is shown below.
+// counter(new) = counter(old) + 1
+// accum(new) = max(accumulator_decay * accum(old) + grad^2,
+//                  initial_accumulator_value)
+// lr_scale = min((step_counter / accum(new)) ^ probability_exponent,
+// max_lr_multiplier) update = grad * lr_scale / sqrt(accum(new)) if
+// (l1_regularization_strength > 0.0):
+//   update = update + l1_regularization_strength * sign(var(old))
+// if (l2_regularization_strength > 0.0):
+//   update = update + l2_regularization_strength * var(old)
+// var(new) = var(old) - lr_scale * grad * update
+
+message FrequencyAwareAdagradParameters {
+  // The L1 regularization parameter for adjusting the update based on the sign
+  // of the variable.
+  float l1_regularization_strength = 1;
+
+  // The L2 regularization parameter for adjusting the update based on the
+  // variable.
+  float l2_regularization_strength = 2;
+
+  // The exponent used for scaling the learning rate based on the sparsity of
+  // updates.
+  float probability_exponent = 4;
+
+  // The maximum value of the learning rate scale.
+  float max_lr_multiplier = 3;
+
+  // The decay for the Adagrad accumulator.
+  float accumulator_decay = 5;
+
+  // The initial and minimum value for the Adagrad accumulator.
+  float initial_accumulator_value = 6;
+
+  // The tag for identifying the step counter used for the frequency aware
+  // Adagrad optimizer.
+  OptimizerDynamicInput step_counter = 7;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L629
+message StochasticGradientDescentParameters {}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Ftrl
+// https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/41159.pdf
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L2646
+//
+// The hyperparameters for FTRL are the same as for the Keras implementation,
+// with some additions. The "beta" parameter matches the behavior described in
+// the second link above; "beta" / (2 * learning rate) should be added to "l2"
+// to get equivalent behavior in the other TensorFlow implementations of this
+// optimizer. When the multiply_linear_by_lr field is set to true, a modified
+// formula is used for FTRL that treats the "linear" accumulator as being
+// pre-multiplied by the learning rate (i.e., the accumulator named "linear"
+// actually stores "linear * learning_rate"). Other than checkpoint
+// compatibility, this is mathematically equivalent for a static learning rate;
+// for a dynamic learning rate, it is nearly the same as long as the learning
+// rate does not change quickly. The benefit of setting multiply_linear_by_lr to
+// true is that the modified formula handles zero and near-zero learning rates
+// without producing NaNs, improving flexibility for learning rate ramp-up.
+message FtrlParameters {
+  float l1 = 1;
+  float l2 = 2;
+  float lr_power = 3;
+  float beta = 7;
+  bool multiply_linear_by_lr = 6;
+
+  // Previously, allow_zero_accumulator parameter changed some internal formulas
+  // to allow zero and near-zero accumulator values at the cost of some
+  // performance. The current implementation ignores this parameter; zero or
+  // near-zero accumulator values are now always supported.
+  bool allow_zero_accumulator = 8 [deprecated = true];
+
+  // Old initial accumulator parameters.
+  reserved "initial_accum", "initial_linear";
+  reserved 4, 5;
+}
+
+// The Adam optimizer does not implement hyper-parameter update due to hardware
+// limitations; use the dynamic learning rate feature instead, setting the
+// learning rate to: user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
+// Here, t is the current timestep.
+//
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adam
+// https://github.com/tensorflow/tensorflow/blob/ab51450c817674c8ff08a7ae4f8ac50cdc4bed8b/tensorflow/python/training/adam.py#L32
+//
+// Note that the code by default implements the lazy version of Adam
+// (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/LazyAdamOptimizer)
+// unless the use_non_lazy_adam parameter is set, in which case it implements
+// the normal version of Adam that updates all parameters in the embedding
+// table, even for entries that are not used in the current minibatch
+// (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/AdamOptimizer). If
+// use_non_lazy_adam is enabled, gradient accumulation is also required to be
+// enabled in order to get correct results; a warning will be printed otherwise
+// (which may change to an error in the future). If use_sum_inside_sqrt is set,
+// the Adam variable update formula will be changed from m / (sqrt(v) + epsilon)
+// to m / sqrt(v + epsilon**2); this option improves the performance of TPU
+// training and is not expected to harm model quality.
+message AdamParameters {
+  float beta1 = 3;
+  float beta2 = 4;
+  float epsilon = 5;
+  bool use_non_lazy_adam = 8;
+  bool use_sum_inside_sqrt = 10;
+
+  // Old initial accumulator parameters.
+  reserved "initial_m", "initial_v";
+  reserved 6, 7;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L3068
+message MomentumParameters {
+  float momentum = 1;
+  bool use_nesterov = 2;
+
+  // Old initial accumulator parameter.
+  reserved "initial_accum";
+  reserved 3;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Lion
+// momenta(new) = beta2 * momenta(old) + (1 - beta2) * grad
+// momenta_t = beta1 * momenta(old) + (1 - beta1) * grad
+// var(new) = var(old) - lr * sign(momenta_t)
+// Algorithm described in https://arxiv.org/abs/2302.06675.
+message LionParameters {
+  float beta1 = 1;
+  float beta2 = 2;
+  bool use_non_lazy_lion = 3;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L4229
+message RmsPropParameters {
+  float rho = 1;
+  float momentum = 2;
+  float epsilon = 3;
+
+  // Old initial accumulator parameters.
+  reserved "initial_ms", "initial_mom";
+  reserved 4, 5;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L4358
+message CenteredRmsPropParameters {
+  float rho = 1;
+  float momentum = 2;
+  float epsilon = 3;
+
+  // Old initial accumulator parameters.
+  reserved "initial_ms", "initial_mom", "initial_mg";
+  reserved 4, 5, 6;
+}
+
+// Variant of algorithm in http://proceedings.mlr.press/v44/shamir15.pdf
+message MdlAdagradLightParameters {
+  float l2 = 1;
+  float lr_power = 2;
+  float min_servable_mdl_benefit = 3;
+  float mdl_mix_in_margin = 4;
+  float mdl_benefit_rampup_coeff = 5;
+  float mdl_min_weight = 6;
+  float benefit_revisit_scale = 7;
+  float max_event_benefit = 8;
+  float max_total_benefit = 9;
+  float mdl_hard_limit = 10;
+  bool hard_limit_min_benefit = 11;
+  bool mdl_regularize = 12;
+
+  // Old initial accumulator parameters.
+  reserved "initial_accumulator", "initial_weight", "initial_benefit";
+  reserved 13, 14, 15;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adadelta
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L933
+message AdadeltaParameters {
+  float rho = 1;
+  float epsilon = 2;
+
+  // Old initial accumulator parameters.
+  reserved "initial_accumulator", "initial_update";
+  reserved 3, 4;
+}
+
+// https://www.tensorflow.org/api_docs/python/tf/compat/v1/train/ProximalAdagradOptimizer
+// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L1961
+message ProximalAdagradParameters {
+  float l1 = 1;
+  float l2 = 2;
+
+  // Old initial accumulator parameter.
+  reserved "initial_accumulator";
+  reserved 3;
+}
+
+// The online Yogi optimizer does not implement hyper-parameter update; use the
+// dynamic learning rate feature instead, setting the learning rate to:
+// user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
+// Here, t is the current timestep.
+//
+// https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf
+// plus some extensions based on FTRL.
+//
+// Note that the code by default implements the lazy version of online Yogi.
+message OnlineYogiParameters {
+  // The L1 regularization parameter (used analogously to the one in FTRL).
+  float l1 = 1;
+
+  // The L2 regularization parameter (used analogously to the one in FTRL).
+  float l2 = 2;
+
+  // \beta_2 from Algorithm 2 in the paper.
+  float beta2 = 3;
+
+  // Reserved ids corresponding to removed tanh activation.
+  reserved 6;  // sign
+  reserved 7;  // tanh
+}
+
+// The online Yogi optimizer does not implement hyper-parameter update; use the
+// dynamic learning rate feature instead, setting the learning rate to:
+// user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
+// Here, t is the current timestep.
+//
+// https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf
+// plus some extensions based on FTRL.
+//
+// Note that the code by default implements the lazy version of proximal Yogi.
+message ProximalYogiParameters {
+  // The L1 regularization parameter.
+  float l1 = 1;
+
+  // The L2 regularization parameter.
+  float l2 = 2;
+
+  // The exponential decay rate for the 1st moment estimates.
+  float beta1 = 3;
+
+  // The exponential decay rate for the 2nd moment estimates.
+  float beta2 = 4;
+
+  // A constant trading off adaptivity and noise.
+  float epsilon = 5;
+
+  // Reserved ids corresponding to removed tanh activation.
+  reserved 8;  // sign
+  reserved 9;  // tanh
+}
+
+// Estimator for the frequency of updates to a lookup table. It maintains an
+// array (tf.Variable) D, where each element records the average number of
+// global steps between two consecutive batches that hit the corresponding
+// bucket. Once an item with bucket id i is sampled, D[i] is updated by:
+//   D[i] <- D[i] * (1 - tau) + delta[i] * tau,
+//
+// where tau is a learning rate between 0 and 1 (exclusive), and
+//   delta[i] = current global step - last step i is sampled.
+//
+// The estimated frequency (sampling rate in a batch) is thus 1 / D[i].
+//
+// Elements in D are initialized with a large value max_delta. delta[i] will
+// also be capped by this value.
+//
+// The exact sequence of operations used in the optimizer is shown below.
+// last_hit_step[i] is a tf.Variable that holds the last global step at which i
+// was sampled.
+//
+//   delta = global_step - last_hit_step[i]
+//   clipped_delta = min(delta, params.max_delta)
+//   is_outlier = (delta >= params.outlier_threshold * D[i])
+//   D[i] <- is_outlier ? clipped_delta
+//                      : D[i] * (1 - params.tau) + clipped_delta * params.tau
+//   last_hit_step[i] <- global_step
+message FrequencyEstimatorParameters {
+  // Learning rate between (0, 1) that is used to update the array D.
+  float tau = 1;
+
+  // Maximum value of delta: difference between the current global step and the
+  // last global step at which the row was sampled.
+  float max_delta = 2;
+
+  // Threshold used to determine whether the current update is an outlier.
+  float outlier_threshold = 3;
+
+  // The weight exponent used to transform the estimated delta into weights.
+  // The transformation function is: (delta / max_delta) ^ (weight_exponent)
+  float weight_exponent = 4;
+}
+
+// A user-defined optimizer.
+// The contained HLO program must take the following arguments in the following
+// order:
+// 1.  gradients
+// 2.  table weights
+// 3.  slot variables
+// 4.  an optional scalar input that is passed in via the dynamic learning
+//     rate mechanism.
+//
+// It must return/end in a tuple op that contains the following values in the
+// following order:
+// 1.  new table values
+// 2.  new slot variable value
+//
+// The program must have shape (1,1) with dtype float32 throughout and only use
+// HLO that operate elementwise (e.g., no reduce, no variables, no control flow
+// and no broadcasting outside of the single scalar input).
+// The HLO program should be written as if it were a dense update. It will be
+// called on each row that needs an update and will applied elementwise.
+message UserDefinedProgramParameters {
+  xla.HloModuleProto program = 1;
+  reserved 2;  // Was padding_values
+}
+
+// Optimizer that just sets the variable to the value of the gradient. To be
+// correct, this requires either gradient accumulation (to sum the values of a
+// computed expression across the samples) or to deduplicate IDs within a single
+// host (to assign the value from an arbitrary sample).
+message AssignParameters {}
+
+// Status of using gradient accumulation (doing two passes over the input
+// gradients: one to accumulate them into a temporary array and another to apply
+// them using the actual optimization algorithm). The extra message is to wrap
+// the enum for scoping.
+message GradientAccumulationStatus {
+  // if UNSPECIFIED (default), gradient accumulation is ENABLED.
+  enum Status {
+    UNSPECIFIED = 0;
+    ENABLED = 1;
+    DISABLED = 2;
+  }
+}
+
+// Whether to optimize the packing of low-dimensional embedding tables in HBM
+// (high bandwidth memory). TPUs access HBM at 32-byte (8-float) granularity.
+// For functional correctness, the TPU software internally pads the embedding
+// dimension to a multiple of 8. This can sometimes lead to significant memory
+// wastage due to padding. For 1-dimensional, 2-dimensional, and 4-dimensional,
+// the TPU software can remove this padding by packing multiple rows into the
+// same 8-float HBM chunk. For example, 8 rows could be packed into the same
+// 8-float chunk for a 1-dimensional embedding table.
+
+// There is one important limitation for this HBM packing though. When only a
+// subset of rows in an 8-float chunk are accessed on a particular step, the
+// adjoining rows in the same chunk are updated with zero gradients on the
+// backward pass even if they are not touched. This is an artifact of the
+// packing implementation. This operation is NOT functionally correct for
+// optimizers where zero gradients change the embeddings/slot-variable values,
+// e.g., momentum-based optimizers. Hence, this HBM packing cannot be enabled
+// for embedding tables with such optimizers. The TPU software automatically
+// recognizes that a zero gradient can modify state and turns off the low
+// dimensional embedding packing in that scenario.
+//
+// However, for optimizers where a zero gradient is a NoOp, such as SGD,
+// Adagrad, and FTRL, this packing optimization can be used. However, there are
+// some important considerations:
+// * Clipping limits: The initial values for such embeddings should fall within
+//   the clipping limits specified in the optimization parameters. Otherwise, a
+//   zero gradient will cause the embeddings to be clipped. This changes state
+//   and hence, is not a NoOp.
+// * FTRL: The embedding vector is computed directly from the values of the
+//   accumulator and linear slot variables. Hence, the initial embedding values
+//   should match that computed from the initial values of the accumulator and
+//   linear slot variables. Note that in nearly all cases, the linear value is
+//   initialized to zero; this corresponds to an embedding value of zero.
+//
+// Performance: The TPU has to perform additional work when low dimensional
+// packing is enabled. In certain situations when the vocabulary size is small,
+// it may not make sense to turn on this packing since the total memory usage
+// due to padding is extremely low. Hence, the TPU software automatically turns
+// off the packing optimization in such scenarios.
+message LowDimensionalPackingStatus {
+  // if UNSPECIFIED (default), the low dimension packing status is DISABLED.
+  // This can change in future.
+  //
+  // if ENABLED, the low dimension packing is enabled only if the following
+  // three additional conditions are true:
+  //  * The optimizer treats the zero gradient as a NoOp.
+  //  * The embedding dimension is 1, 2, or 4.
+  //  * The vocabulary size is large enough to avoid performance issues.
+  //
+  // if DISABLED, the low dimension packing is always disabled.
+  enum Status {
+    UNSPECIFIED = 0;
+    ENABLED = 1;
+    DISABLED = 2;
+  }
+}
+
+// Configuration proto for hot ID optimization. This is an experimental feature
+// that is currently disabled (by default).
+message HotIdReplicationConfiguration {
+  // Whether to enable or disable hot ID optimization.
+  // If set to UNSPECIFIED (default), hot ID optimization is DISABLED.
+  // If set to ENABLED, hot ID replication is turned ON.
+  // If set to MIGRATION_ONLY, hot ID migration is turned ON.
+  enum Status {
+    UNSPECIFIED = 0;
+    ENABLED = 1;
+    DISABLED = 2;
+    MIGRATION_ONLY = 3;
+  }
+  Status status = 1;
+}
+
+message OptimizationParameters {
+  // Learning rate used for updating the embedding layer parameters.
+  LearningRate learning_rate = 13;
+
+  // Limits to which to clip the weight values after the backward pass; not
+  // present means no limits are applied.
+  ClippingLimits clipping_limits = 2;
+
+  // Limits to which to clip the backward pass gradient before using it for
+  // updates; not present means no limits are applied.
+  ClippingLimits gradient_clipping_limits = 7;
+
+  // Amount of weight decay to apply; see weight_decay_optimizers.py for
+  // details. All optimizers except MDL Adagrad Light are supported with this
+  // option. Although there is no check, users who want weight decay will also
+  // want to ensure that gradient accumulation is enabled so that the decay will
+  // happen once per global batch.
+  float weight_decay_factor = 16;
+
+  // If true, the weight decay factor is multiplied by the current learning rate
+  // before use; this is to match the note in DecoupledWeightDecayExtension in
+  // weight_decay_optimizers.py.
+  bool multiply_weight_decay_factor_by_learning_rate = 22;
+
+  // Configuration for simulated quantization which is used to reduce
+  // training/serving skew when the serving variables are quantized. The same
+  // quantization operations are executed during training to minimize
+  // differences with serving.
+  SimulatedQuantization simulated_quantization = 27;
+
+  // Status of using gradient accumulation (doing two passes over the input
+  // gradients: one to accumulate them into a temporary array and another to
+  // apply them using the actual optimization algorithm).
+  GradientAccumulationStatus.Status gradient_accumulation_status = 17;
+
+  // Status of the low-dimensional embedding packing optimization. This controls
+  // whether to optimize the packing of 1-dimensional, 2-dimensional, and
+  // 4-dimensional embedding tables in memory.
+  LowDimensionalPackingStatus.Status low_dimensional_packing_status = 28;
+
+  // Configuration proto for hot ID replication. This is an experimental
+  // feature that is currently disabled (by default).
+  HotIdReplicationConfiguration hot_id_replication_configuration = 18;
+
+  // Optimization algorithm parameters; which field is selected determines which
+  // algorithm to use.
+  oneof parameters {
+    AdagradParameters adagrad = 3;
+    AdagradMomentumParameters adagrad_momentum = 26;
+    BoundedAdagradParameters bounded_adagrad = 19;
+    FrequencyAwareAdagradParameters frequency_aware_adagrad = 30;
+    StochasticGradientDescentParameters stochastic_gradient_descent = 4;
+    FtrlParameters ftrl = 5;
+    AdamParameters adam = 6;
+    MomentumParameters momentum = 8;
+    LionParameters lion = 29;
+    RmsPropParameters rms_prop = 9;
+    CenteredRmsPropParameters centered_rms_prop = 10;
+    MdlAdagradLightParameters mdl_adagrad_light = 11;
+    AdadeltaParameters adadelta = 12;
+    ProximalAdagradParameters proximal_adagrad = 14;
+    OnlineYogiParameters online_yogi = 20;
+    ProximalYogiParameters proximal_yogi = 21;
+    FrequencyEstimatorParameters frequency_estimator = 23;
+    UserDefinedProgramParameters user_defined_program = 24;
+    AssignParameters assign = 25;
+  }
+
+  reserved 1, 15;
+
+  // NEXT_ID: 31
+}
+
+// Specification of an optimization algorithm's state variables (both the main
+// value vector and any extra accumulators, etc.). This proto is only used
+// internally by the TPU software and is not exposed directly to the TF model.
+message StateVariableSpecification {
+  // Parameter name for the state variable.
+  string name = 1;
+
+  // A normal state variable that should be saved and restored in checkpoints
+  // and used as an input or output to non-debug TensorFlow ops.
+  message UserDefined {
+    reserved 1;  // Was padding_initial_value.
+  }
+
+  // A state variable that should be filled with a constant and normally hidden
+  // from users (used for intermediate gradients being accumulated, for
+  // example).
+  message FillWithConstant {
+    double initial_value = 1;
+  }
+
+  // Usage type of this state variable.
+  oneof usage {
+    UserDefined user_defined = 2;
+    FillWithConstant fill_with_constant = 3;
+  }
+}
diff --git a/third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration.proto b/third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration.proto
new file mode 100644
index 00000000000..14f64bcc391
--- /dev/null
+++ b/third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration.proto
@@ -0,0 +1,176 @@
+// 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.
+// ==============================================================================
+
+syntax = "proto3";
+
+package tensorflow.tpu;
+
+import "xla/tpu/protobuf/optimization_parameters.proto";
+
+message TPUEmbeddingConfiguration {
+  // Description of the various embedding tables.
+  message TableDescriptor {
+    // Name of the table.
+    string name = 1;
+
+    // Size of the vocabulary (i.e., number of rows) in the table.
+    int64 vocabulary_size = 2;
+
+    // The embedding dimension (i.e., the width of the embedding table).
+    int32 dimension = 3;
+
+    // Number of features mapped to this table.
+    int32 num_features = 4;
+
+    // Details of the learning algorithm used to update the embedding
+    // parameters.
+    OptimizationParameters optimization_parameters = 5;
+  }
+  repeated TableDescriptor table_descriptor = 1;
+
+  // Mode. Should the embedding layer program be run for inference (just forward
+  // pass), training (both forward and backward pass) or just the backward_pass.
+  enum Mode {
+    UNSPECIFIED = 0;
+    INFERENCE = 1;
+    TRAINING = 2;
+    BACKWARD_PASS_ONLY = 3;
+  }
+  Mode mode = 2;
+
+  // Number of samples in each batch of embedding layer activations sent to
+  // the TensorCore.
+  int32 batch_size_per_tensor_core = 3;
+
+  // Number of TPU hosts used for inference/training.
+  int32 num_hosts = 4;
+
+  // Number of TensorCore used for inference/training.
+  int32 num_tensor_cores = 5;
+
+  // Sharding strategy of the embedding tables among the hosts.
+  // If the sharding_strategy is "mod", each id is assigned to host
+  // "id % num_hosts". For instance, 13 ids are split across 5 hosts as:
+  // [[0, 5, 10], [1, 6, 11], [2, 7, 12], [3, 8], [4, 9]].
+  // If the sharding_strategy is "div", ids are assigned to hosts in a
+  // contiguous manner. In this case, 13 ids are split across 5 hosts as:
+  // [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11, 12]].
+  // In both the strategies, if the id space does not evenly divide the number
+  // of hosts, each of the first "table_descriptor.vocabulary_size % num_hosts"
+  // hosts will be assigned one more id.
+  // This partitioning strategy exactly follows that in the embedding_lookup
+  // TensorFlow function at tensorflow/python/ops/embedding_ops.py.
+  enum ShardingStrategy {
+    DIV_DEFAULT = 0;
+    MOD = 1;
+  }
+  ShardingStrategy sharding_strategy = 6;
+
+  // This parameter determines if the execution of the sparse core will be
+  // pipelined with that of the TensorCore. This parameter only affects results
+  // when mode=TRAINING. If mode=INFERENCE or BACKWARD_PASS_ONLY, this parameter
+  // does not affect execution and hence, is a don't care value.
+  //
+  // false: The execution of the sparse core is not pipelined with that of the
+  // TensorCore. The forward pass of every step on the sparse core is executed
+  // only after the backward pass of the previous step is complete. And the
+  // backward pass on the sparse core is executed only after the embedding
+  // gradients have been computed on the TensorCore on every step. This ensures
+  // that the activations on every step observe the gradient updates from the
+  // previous step on both the sparse core and the TensorCore.
+  //
+  // true: The execution of the sparse core is pipelined with that of the
+  // TensorCore. The forward pass of every step on the sparse core can be
+  // executed after the forward pass of the previous step is complete without
+  // waiting for the backward pass. This improves the utilization of the sparse
+  // core allowing it to process step N+1 while the embedding gradients for step
+  // N are computed on the TensorCore. The backward pass of every step on the
+  // sparse core is executed directly after the forward pass for the next step
+  // is complete. The drawback is that embedding activations for step N+1 do not
+  // observe the embedding gradient updates from step N. This could affect model
+  // quality if step N and N+1 involve the same set of embedding IDs. However,
+  // since the embedding updates are sparse, this is generally not considered a
+  // problem.
+  bool pipeline_execution_with_tensor_core = 7;
+
+  // Directory where embedding lookup statistics are stored. These statistics
+  // summarize information about the inputs to the embedding lookup
+  // operation, in particular, the average number of embedding IDs per example
+  // and how well the embedding IDs are load balanced across the system. The
+  // lookup statistics are used during TPU initialization for embedding table
+  // partitioning. Collection of lookup statistics is done at runtime by
+  // profiling the embedding inputs: only 3% of input samples are profiled to
+  // minimize host CPU overhead. Once a suitable number of samples are
+  // profiled, the lookup statistics are saved to table-specific files in the
+  // profile data directory generally at the end of a TPU training loop. The
+  // filename corresponding to each table is obtained by hashing table specific
+  // parameters (e.g., table name and number of features) and global
+  // configuration parameters (e.g., sharding strategy and TPU worker task
+  // count). The same profile data directory can be shared amongst several
+  // models to reuse embedding lookup statistics.
+  string profile_data_directory = 9;
+
+  // Description of different input features.
+  message FeatureDescriptor {
+    // Name of the input feature.
+    string name = 1;
+
+    // Index of the corresponding table in the TableDescriptor list.
+    int32 table_id = 2;
+
+    // Static shape of the inputs (excluding the reduction axis). Note that
+    // the shape of the actual inputs provided using the infeed op must be
+    // strictly smaller than input_shape. The outputs received at the TensorCore
+    // will have rank = input_shape.size() + 1. The innermost axis corresponds
+    // to the embedding dimension. If the input has shape [m, n, k] (excluding
+    // the reduction axis) and the embedding dimension is d, the output received
+    // at the TensorCore will have shape [m, n, k, d].
+    repeated int32 input_shape = 3;
+  }
+  // If the feature_descriptor field is populated, the model should NOT populate
+  // TableDescriptor.num_features and batch_size_per_tensor_core. These two
+  // fields will be auto-populated by the TPUEmbedding rewrite passes.
+  repeated FeatureDescriptor feature_descriptor = 10;
+
+  // SPMD (Single Program Multiple Data) sharding configuration for
+  // TPUEmbedding. When model parallelism is used on the TensorCore, the number
+  // of cores per replica must be passed to TPUEmbedding so that the right
+  // shapes can be computed in the TF/XLA bridge.
+  message SpmdSharding {
+    // Whether SPMD sharding is enabled.
+    bool enabled = 1;
+
+    // Number of cores per replica.
+    int32 num_cores_per_replica = 2;
+
+    // If true, the tensors are manually partitioned. Otherwise, use the
+    // automatic SPMD partitioning. This should be true when users use
+    // `shard_map`.
+    bool use_manual_partitioning = 3;
+
+    // If true, lowers to Shardy custom-call for manual partitioning.
+    // Default is false.
+    bool use_shardy_partitioner = 4;
+  }
+  SpmdSharding spmd_sharding = 11;
+
+  // Old TPU embedding output layout.
+  reserved "output_layout";
+  reserved 8;
+}
+
+// A placeholder message that is used to define a unique Status payload
+// URL for TPU embedding errors.
+message TPUEmbeddingError {}

From 8b14a432d7e839c2bc3bb80a8558aa39cb7966e0 Mon Sep 17 00:00:00 2001
From: Jay Gan <jihonggan@google.com>
Date: Mon, 12 Oct 2026 11:07:42 -0700
Subject: [PATCH 0605/1000] Update uv to ~= 0.12.23 in ml-build container
 requirements

`ml-build` currently pins `uv ~= 0.5.30` (from Feb 2025), which caps `uv` at
`0.5.31`. Newer `uv` releases retry transient HTTP/stream errors such as
`BrokenPipe` (added in 0.6.6), mid-body `UnexpectedEof`, and `h2` stream resets,
and support `UV_HTTP_RETRIES` (added in 0.7.21).

PiperOrigin-RevId: 998031103
---
 ci/official/containers/ml_build/builder.requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ci/official/containers/ml_build/builder.requirements.txt b/ci/official/containers/ml_build/builder.requirements.txt
index 3bd28656e7d..055f6609af1 100644
--- a/ci/official/containers/ml_build/builder.requirements.txt
+++ b/ci/official/containers/ml_build/builder.requirements.txt
@@ -11,7 +11,7 @@ pyyaml
 # For JAX
 build ~= 1.2.2
 # uv is faster than pip for installing Python packages.
-uv ~= 0.5.30
+uv ~= 0.12.23
 
 # For running wheel verification script
 immutabledict ~= 4.2.1

From a839108de0e2625cea683f044b2d10984ac7a138 Mon Sep 17 00:00:00 2001
From: Pablo Zimmermann <karupayun@google.com>
Date: Mon, 12 Oct 2026 11:11:54 -0700
Subject: [PATCH 0606/1000] [XLA:GPU] Propagate dot block resource usage in
 GpuIndexingPerformanceModel.

As part of migrating `gpu_dot_fusion_cost_model` into
`GpuIndexingPerformanceModel`, tiled dot fusions need per-block resource
footprints (`shared_memory_per_block_bytes` and `registers_per_thread`) in
`EstimateRunTimeData` so downstream passes and wave-quantized cost modeling can
account for occupancy and launch limits.

Whereas `EstimateRunTimeForDotOpWithBlockParameters` in
`gpu_dot_fusion_cost_model` only evaluates a single `dot` instruction,
propagating and merging `DotEstimates` in `GpuIndexingPerformanceModel` tracks
the maximum per-block shared memory and register footprints across all dots in a
tiled fusion while preserving existing runtime estimates.

PiperOrigin-RevId: 998034026
---
 third_party/xla/xla/service/gpu/model/BUILD   |  1 +
 .../model/gpu_indexing_performance_model.cc   | 95 ++++++++++++-------
 .../gpu_indexing_performance_model_test.cc    | 77 +++++++++++++++
 3 files changed, 140 insertions(+), 33 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/BUILD b/third_party/xla/xla/service/gpu/model/BUILD
index cce38ba2862..4b5ba59841e 100644
--- a/third_party/xla/xla/service/gpu/model/BUILD
+++ b/third_party/xla/xla/service/gpu/model/BUILD
@@ -637,6 +637,7 @@ xla_cc_test(
         "//xla/hlo/analysis:symbolic_map",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+        "//xla/hlo/testlib:verified_hlo_module",
         "//xla/hlo/utils:hlo_traversal",
         "//xla/service:hlo_cost_analysis",
         "//xla/service/gpu:gpu_device_info_for_tests",
diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc
index ef9ae0d2493..92f82ea499b 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc
+++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc
@@ -378,20 +378,40 @@ int64_t GetNumWarps(int64_t largest_live_tile_size) {
   return 8;
 }
 
+// Holds dot launch resource footprints and compute/runtime estimates.
+struct DotEstimates {
+  absl::Duration compute_time = absl::ZeroDuration();
+  absl::Duration exec_time = absl::ZeroDuration();
+  int64_t shared_memory_per_block_bytes = 0;
+  int registers_per_thread = 0;
+};
+
+// Combines two dot estimates within the same fusion by summing sequential
+// compute work and taking the maximum of per-block resource footprints.
+DotEstimates MergeDotEstimates(DotEstimates dst, const DotEstimates& src) {
+  dst.compute_time += src.compute_time;
+  // TODO: b/495346904 - integrate the dot stats more completely.
+  dst.exec_time = std::max(dst.exec_time, src.exec_time);
+  dst.shared_memory_per_block_bytes = std::max(
+      dst.shared_memory_per_block_bytes, src.shared_memory_per_block_bytes);
+  dst.registers_per_thread =
+      std::max(dst.registers_per_thread, src.registers_per_thread);
+  return dst;
+}
+
+// Computes block launch resource footprints and dot runtime estimates.
 template <typename TiledHloInstructionType>
-absl::StatusOr<EstimateRunTimeData> GetDotEstimates(
+absl::StatusOr<DotEstimates> BuildDotEstimates(
     const TiledHloInstructionType* tiled_hlo,
     const se::DeviceDescription& device_info,
     const BlockLevelParameters& block_params) {
-  const auto* dot_instr = Cast<const HloDotInstruction>(tiled_hlo->hlo());
+  const HloDotInstruction* dot_instr =
+      Cast<HloDotInstruction>(tiled_hlo->hlo());
   ABSL_RETURN_IF_ERROR(gpu_dot_fusion_cost_model::IsSupported(dot_instr));
 
-  int64_t block_k = 0;
-  if (!tiled_hlo->operands().empty()) {
-    int64_t lhs_contracting_dim =
-        dot_instr->dot_dimension_numbers().lhs_contracting_dimensions(0);
-    block_k = tiled_hlo->operand(0)->tile_size(lhs_contracting_dim);
-  }
+  int64_t lhs_contracting_dim =
+      dot_instr->dot_dimension_numbers().lhs_contracting_dimensions(0);
+  int64_t block_k = tiled_hlo->operand(0)->tile_size(lhs_contracting_dim);
 
   // Use the provided block parameters which contain num_stages and num_warps.
   // The output_tile_sizes from the block parameters represent the fusion roots,
@@ -403,9 +423,18 @@ absl::StatusOr<EstimateRunTimeData> GetDotEstimates(
   const llvm::SmallVector<int64_t>& tile_sizes = tiled_hlo->tile_sizes();
   dot_block_params.output_tile_sizes.emplace_back(tile_sizes.begin(),
                                                   tile_sizes.end());
+  ABSL_ASSIGN_OR_RETURN(
+      EstimateRunTimeData dot_perf_stats,
+      gpu_dot_fusion_cost_model::EstimateRunTimeForDotOpWithBlockParameters(
+          dot_instr, dot_block_params, device_info, block_k));
 
-  return gpu_dot_fusion_cost_model::EstimateRunTimeForDotOpWithBlockParameters(
-      dot_instr, dot_block_params, device_info, block_k);
+  return DotEstimates{
+      /*compute_time=*/dot_perf_stats.compute_time,
+      /*exec_time=*/dot_perf_stats.exec_time,
+      /*shared_memory_per_block_bytes=*/
+      dot_perf_stats.shared_memory_per_block_bytes,
+      /*registers_per_thread=*/dot_perf_stats.registers_per_thread,
+  };
 }
 
 int64_t GetShapeSizeRecursive(
@@ -438,8 +467,7 @@ absl::StatusOr<EstimateRunTimeData> EstimateRunTimeForTiledHloComputationImpl(
   int64_t bytes_read = 0;
   int64_t num_blocks = tiled_hlo_computation.num_output_tiles();
 
-  absl::Duration dot_compute_time = absl::ZeroDuration();
-  absl::Duration dot_exec_time = absl::ZeroDuration();
+  DotEstimates dot_estimates;
 
   // Check if the computation is too large to fit in registers and would result
   // in spilling.
@@ -467,24 +495,20 @@ absl::StatusOr<EstimateRunTimeData> EstimateRunTimeForTiledHloComputationImpl(
           int64_t num_elements = num_blocks_cur_hlo * padded_tile_size;
 
           if (hlo->opcode() == HloOpcode::kDot) {
-            absl::StatusOr<EstimateRunTimeData> dot_perf_stats =
-                GetDotEstimates(tiled_hlo, device_info, block_level_parameters);
-            if (dot_perf_stats.ok()) {
-              // We're only using compute time for now - memory and L2 access
-              // data needs to be adjusted more carefully so that the model
-              // doesn't overlap their counting.
-              // TODO: b/495346904 - integrate the dot stats more completely.
-              dot_compute_time += dot_perf_stats->compute_time;
-              dot_exec_time =
-                  std::max(dot_exec_time, dot_perf_stats->exec_time);
+            if (absl::StatusOr<DotEstimates> cur_dot_estimates =
+                    BuildDotEstimates(tiled_hlo, device_info,
+                                      block_level_parameters);
+                cur_dot_estimates.ok()) {
+              dot_estimates =
+                  MergeDotEstimates(dot_estimates, *cur_dot_estimates);
 
               // The dot cost model operates on the tile- and wave- quantized
               // FLOPS which is more accurate for performance estimates but
               // the flops reported by the indexing model are naive algorithm
               // flops.
               // Since the compute time for these flops is already accounted
-              // for in dot_compute_time, we're accumulating it separately
-              // from `flops`.
+              // for in dot_estimates.compute_time, we're accumulating it
+              // separately from `flops`.
               dot_flops += flops_per_element_fn(hlo) * num_elements;
               return;
             }
@@ -551,7 +575,7 @@ absl::StatusOr<EstimateRunTimeData> EstimateRunTimeForTiledHloComputationImpl(
       GpuPerformanceModelBase::ComputeTime(
           device_info, flops, num_blocks,
           block_level_parameters.num_warps * WarpSize(device_info)) +
-      dot_compute_time;
+      dot_estimates.compute_time;
 
   absl::Duration memory_access_time = read_time + write_time;
   absl::Duration exec_time =
@@ -563,15 +587,20 @@ absl::StatusOr<EstimateRunTimeData> EstimateRunTimeForTiledHloComputationImpl(
   // accounting for any work in the epilogue or prologue properly. We are
   // planning to do it in the future but the cost model will need to be
   // significantly more complex for that.
-  exec_time = std::max(exec_time, dot_exec_time);
+  exec_time = std::max(exec_time, dot_estimates.exec_time);
 
-  return EstimateRunTimeData{/*flops=*/flops + dot_flops,
-                             /*bytes_read=*/bytes_read,
-                             /*bytes_written=*/bytes_written,
-                             /*read_time=*/read_time,
-                             /*write_time=*/write_time,
-                             /*compute_time=*/compute_time,
-                             /*exec_time=*/exec_time};
+  return EstimateRunTimeData{
+      /*flops=*/flops + dot_flops,
+      /*bytes_read=*/bytes_read,
+      /*bytes_written=*/bytes_written,
+      /*read_time=*/read_time,
+      /*write_time=*/write_time,
+      /*compute_time=*/compute_time,
+      /*exec_time=*/exec_time,
+      /*l2_bytes_read=*/0,
+      /*shared_memory_per_block_bytes=*/
+      dot_estimates.shared_memory_per_block_bytes,
+      /*registers_per_thread=*/dot_estimates.registers_per_thread};
 }
 
 template <typename TiledHloComputationType>
diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc
index d688a5f05eb..67195c26f6d 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc
+++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc
@@ -39,6 +39,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
+#include "xla/hlo/testlib/verified_hlo_module.h"
 #include "xla/hlo/utils/hlo_traversal.h"
 #include "xla/service/gpu/gpu_device_info_for_tests.h"
 #include "xla/service/gpu/ir_emission_utils.h"
@@ -61,6 +62,8 @@ namespace xla {
 namespace gpu {
 namespace {
 
+namespace dot_model = ::xla::gpu::gpu_dot_fusion_cost_model::detail;
+
 using ::testing::ElementsAre;
 using ::xla::xtile::BlockLevelParameters;
 
@@ -723,6 +726,80 @@ ENTRY main {
       4 * (64 * 32) * (num_blocks * CeilOfRatio(257, 64))        // ==> %p1
           + 4 * (64 * 32) * (num_blocks * CeilOfRatio(257, 64))  // ==> %p0
   );
+  // (32 * 64 + 32 * 64) * 4 bytes (f32) * 1 stage = 16384 bytes.
+  EXPECT_EQ(result.shared_memory_per_block_bytes, 2 * (32 * 64 * 4));
+  // 1 accumulator register per thread (32 * 32 / 1024) plus base kernel-state
+  // overhead, staying below one 32-register allocation block.
+  EXPECT_GT(result.registers_per_thread, 1);
+  EXPECT_LT(result.registers_per_thread, 32);
+}
+
+TEST_F(GpuIndexingPerformanceModelTest,
+       EstimateRunTimeForTiledFusion_ChainedDotsMergesDotEstimates) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(R"(
+HloModule m
+
+chained_dot_fusion {
+  p0 = f32[128, 64] parameter(0)
+  p1 = f32[64, 128] parameter(1)
+  p2 = f32[128, 32] parameter(2)
+  dot0 = f32[128, 128] dot(p0, p1),
+    lhs_contracting_dims={1}, rhs_contracting_dims={0},
+    backend_config={sizes:[16]}
+  ROOT dot1 = f32[128, 32] dot(dot0, p2),
+    lhs_contracting_dims={1}, rhs_contracting_dims={0},
+    backend_config={sizes:[64]}
+}
+
+ENTRY main {
+  p0 = f32[128, 64] parameter(0)
+  p1 = f32[64, 128] parameter(1)
+  p2 = f32[128, 32] parameter(2)
+  ROOT fusion = f32[128, 32] fusion(p0, p1, p2),
+    kind=kCustom, calls=chained_dot_fusion
+})"));
+
+  std::unique_ptr<HloFusionAdaptor> fusion_adaptor =
+      HloFusionAdaptor::ForInstruction(
+          module->entry_computation()->root_instruction());
+
+  ASSERT_OK_AND_ASSIGN(
+      EstimateRunTimeData result,
+      indexing_cost_model_.EstimateRunTimeForTiledFusion(
+          *fusion_adaptor,
+          BlockLevelParameters{/*output_tile_sizes=*/{{32, 16}},
+                               /*num_warps=*/2}));
+
+  // Sequential dots reuse shared memory and registers, so the fusion takes the
+  // max of each across dot0 (tile [32, 64], block_k=16) and dot1 (tile [32,
+  // 16], block_k=64):
+  // - dot1 has the larger block_k and dominates shared memory:
+  //   (32 * 64 + 16 * 64) * 4 = 12288 bytes (vs. 6144 for dot0).
+  // - dot0 has the larger output tile and dominates registers: 32 accumulator
+  //   registers per thread (32 * 64 / 64) plus base kernel-state overhead
+  //   (vs. 8 accumulator registers for dot1), staying below 64.
+  EXPECT_EQ(result.shared_memory_per_block_bytes, 12288);
+  EXPECT_GT(result.registers_per_thread, 32);
+  EXPECT_LT(result.registers_per_thread, 64);
+
+  const HloDotInstruction* dot1 = Cast<HloDotInstruction>(
+      module->entry_computation()->root_instruction()->fused_expression_root());
+  const HloDotInstruction* dot0 = Cast<HloDotInstruction>(dot1->operand(0));
+  ASSERT_OK_AND_ASSIGN(
+      dot_model::ComputeAndFlops dot0_compute,
+      dot_model::CalculateComputeTimeWithTileAndWaveQuantization(
+          dot_model::DotProblemInfo(*dot0),
+          dot_model::DotTileSize{/*m=*/32, /*n=*/64, /*k=*/16, /*b=*/1},
+          device_info_));
+  ASSERT_OK_AND_ASSIGN(
+      dot_model::ComputeAndFlops dot1_compute,
+      dot_model::CalculateComputeTimeWithTileAndWaveQuantization(
+          dot_model::DotProblemInfo(*dot1),
+          dot_model::DotTileSize{/*m=*/32, /*n=*/16, /*k=*/64, /*b=*/1},
+          device_info_));
+  EXPECT_EQ(result.compute_time,
+            dot0_compute.compute_time + dot1_compute.compute_time);
 }
 
 TEST_F(GpuIndexingPerformanceModelTest,

From e645b8c81565efb9bf9038cf73cced3350bd1ba7 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Tue, 13 Oct 2026 00:05:07 +0000
Subject: [PATCH 0607/1000] Bump werkzeug in
 /ci/official/requirements_updater/numpy1_requirements

Bumps [werkzeug](https://github.com/pallets/werkzeug) from 3.1.6 to 3.1.9.
- [Release notes](https://github.com/pallets/werkzeug/releases)
- [Changelog](https://github.com/pallets/werkzeug/blob/main/CHANGES.rst)
- [Commits](https://github.com/pallets/werkzeug/compare/3.1.6...3.1.9)

---
updated-dependencies:
- dependency-name: werkzeug
  dependency-version: 3.1.9
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <support@github.com>
---
 .../numpy1_requirements/requirements_lock_3_10.txt          | 6 +++---
 .../numpy1_requirements/requirements_lock_3_11.txt          | 6 +++---
 .../numpy1_requirements/requirements_lock_3_12.txt          | 6 +++---
 3 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt
index 00968ad4260..7ebf036e51c 100644
--- a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt
+++ b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_10.txt
@@ -768,9 +768,9 @@ urllib3==2.8.0 \
     --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
     --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
-werkzeug==3.1.6 \
-    --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
-    --hash=sha256:7ddf3357bb9564e407607f988f683d72038551200c704012bb9a4c523d42f131
+werkzeug==3.1.9 \
+    --hash=sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060 \
+    --hash=sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab
     # via tb-nightly
 wheel==0.41.3 \
     --hash=sha256:488609bc63a29322326e05560731bf7bfea8e48ad646e1f5e40d366607de0942 \
diff --git a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt
index 5f834b529ae..2dac0c07385 100644
--- a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt
+++ b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_11.txt
@@ -767,9 +767,9 @@ urllib3==2.8.0 \
     --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
     --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
-werkzeug==3.1.6 \
-    --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
-    --hash=sha256:7ddf3357bb9564e407607f988f683d72038551200c704012bb9a4c523d42f131
+werkzeug==3.1.9 \
+    --hash=sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060 \
+    --hash=sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab
     # via tb-nightly
 wheel==0.41.3 \
     --hash=sha256:488609bc63a29322326e05560731bf7bfea8e48ad646e1f5e40d366607de0942 \
diff --git a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt
index 6b49c8c337c..c7210b8bafb 100644
--- a/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt
+++ b/ci/official/requirements_updater/numpy1_requirements/requirements_lock_3_12.txt
@@ -767,9 +767,9 @@ urllib3==2.8.0 \
     --hash=sha256:0cf3cae568d36aa9576b28dfb35f11328f1cb974ca7647d9475ebb86c75ac6e3 \
     --hash=sha256:63bf2ead4c879426ebf22ef2a781eeb4aa3b4ae798a0435506f8687fd5bb9b63
     # via requests
-werkzeug==3.1.6 \
-    --hash=sha256:210c6bede5a420a913956b4791a7f4d6843a43b6fcee4dfa08a65e93007d0d25 \
-    --hash=sha256:7ddf3357bb9564e407607f988f683d72038551200c704012bb9a4c523d42f131
+werkzeug==3.1.9 \
+    --hash=sha256:55ca7c70a75689be937aa27f8ff4b018f06ff4838fc73045560bf0f5a1291060 \
+    --hash=sha256:6392e50c78460ba618e5b21f08a71f59c99ce99cdc6cf6e3dd7e6ccca8754fab
     # via tb-nightly
 wheel==0.41.3 \
     --hash=sha256:488609bc63a29322326e05560731bf7bfea8e48ad646e1f5e40d366607de0942 \

From c84e0878de81d2aeafd9dca209e9091f5ec26dd3 Mon Sep 17 00:00:00 2001
From: Sean Talts <talts@google.com>
Date: Mon, 12 Oct 2026 11:25:24 -0700
Subject: [PATCH 0608/1000] [XLA:CPU] Add HLO benchmarks for JAX issue #40150
 and XLA issue #50040

Both are lax.scan gradients whose while loops consist of many small ops and
run much slower when SmallWhileLoopHoistingPass does not outline them:
- https://github.com/jax-ml/jax/issues/40150: 60 scalar carry leaves.
- https://github.com/openxla/xla/issues/50040: f64[3,3] carry whose
  per-iteration byte estimate sits just above the hoisting threshold.

PiperOrigin-RevId: 998042953
---
 .../backends/cpu/benchmarks/hlo/hlo_files.bzl |    2 +
 .../benchmarks/hlo/jax.issue.40150.grad.hlo   | 2795 +++++++++++++++++
 .../hlo/xla.issue.50040.value_and_grad.hlo    |  193 ++
 3 files changed, 2990 insertions(+)
 create mode 100644 third_party/xla/xla/backends/cpu/benchmarks/hlo/jax.issue.40150.grad.hlo
 create mode 100644 third_party/xla/xla/backends/cpu/benchmarks/hlo/xla.issue.50040.value_and_grad.hlo

diff --git a/third_party/xla/xla/backends/cpu/benchmarks/hlo/hlo_files.bzl b/third_party/xla/xla/backends/cpu/benchmarks/hlo/hlo_files.bzl
index c1d7076f67c..838c7a88a6a 100644
--- a/third_party/xla/xla/backends/cpu/benchmarks/hlo/hlo_files.bzl
+++ b/third_party/xla/xla/backends/cpu/benchmarks/hlo/hlo_files.bzl
@@ -44,6 +44,7 @@ REGULAR_HLO_FILES = [
     "jax.issue.33666.linx.slow_0100.module_0013.hlo",
     "jax.issue.33666.linx.slow_0428.module_0003.hlo",
     "jax.issue.33666.linx.slow_0428.module_0006.hlo",
+    "jax.issue.40150.grad.hlo",
     "layer_norm_1x4096x768_bf16.hlo",
     "mean_axis_1x4096x1024_bf16.hlo",
     "mha_block_1x12x128x64_bf16.hlo",
@@ -60,6 +61,7 @@ REGULAR_HLO_FILES = [
     "sum_axis_1x4096x1024_bf16.hlo",
     "topk_logits_k10_1x50000_bf16.hlo",
     "vit_keras_jax_bf16.hlo",
+    "xla.issue.50040.value_and_grad.hlo",
     "xnn.parallel_dots.optimized.hlo",
     "xnn.sequential_dots.optimized.hlo",
     # go/keep-sorted end
diff --git a/third_party/xla/xla/backends/cpu/benchmarks/hlo/jax.issue.40150.grad.hlo b/third_party/xla/xla/backends/cpu/benchmarks/hlo/jax.issue.40150.grad.hlo
new file mode 100644
index 00000000000..3ee5fb8a24f
--- /dev/null
+++ b/third_party/xla/xla/backends/cpu/benchmarks/hlo/jax.issue.40150.grad.hlo
@@ -0,0 +1,2795 @@
+// Reproducer from https://github.com/jax-ml/jax/issues/40150: jax.grad through
+// lax.scan whose carry is 60 scalar leaves (614 steps).
+HloModule jit_model, entry_computation_layout={(f32[3]{0})->f32[3]{0}}, allow_spmd_sharding_propagation_to_parameters={true}, allow_spmd_sharding_propagation_to_output={true}
+
+%dynamic_index_in_dim.1 (Arg_0.1: f32[614], Arg_1.1: s32[]) -> f32[] {
+  %Arg_0.1 = f32[614]{0} parameter(0)
+  %Arg_1.1 = s32[] parameter(1)
+  %dynamic_slice.1 = f32[1]{0} dynamic-slice(%Arg_0.1, %Arg_1.1), dynamic_slice_sizes={1}
+  ROOT %squeeze.61 = f32[] reshape(%dynamic_slice.1)
+}
+
+%dynamic_update_index_in_dim.2 (Arg_0.3: pred[614], Arg_1.3: pred[], Arg_2.1: s32[]) -> pred[614] {
+  %Arg_0.3 = pred[614]{0} parameter(0)
+  %Arg_1.3 = pred[] parameter(1)
+  %broadcast_in_dim.1 = pred[1]{0} reshape(%Arg_1.3)
+  %Arg_2.1 = s32[] parameter(2)
+  ROOT %dynamic_update_slice.1 = pred[614]{0} dynamic-update-slice(%Arg_0.3, %broadcast_in_dim.1, %Arg_2.1)
+}
+
+%dynamic_update_index_in_dim_0.3 (Arg_0.5: f32[614], Arg_1.5: f32[], Arg_2.3: s32[]) -> f32[614] {
+  %Arg_0.5 = f32[614]{0} parameter(0)
+  %Arg_1.5 = f32[] parameter(1)
+  %broadcast_in_dim.3 = f32[1]{0} reshape(%Arg_1.5)
+  %Arg_2.3 = s32[] parameter(2)
+  ROOT %dynamic_update_slice.3 = f32[614]{0} dynamic-update-slice(%Arg_0.5, %broadcast_in_dim.3, %Arg_2.3)
+}
+
+%region_0.4 (arg_tuple.1: (s32[], f32[], f32[], f32[], f32[], /*index=5*/f32[], f32[], f32[], f32[], f32[], /*index=10*/f32[], f32[], f32[], f32[], f32[], /*index=15*/f32[], f32[], f32[], f32[], f32[], /*index=20*/f32[], f32[], f32[], f32[], f32[], /*index=25*/f32[], f32[], f32[], f32[], f32[], /*index=30*/f32[], f32[], f32[], f32[], f32[], /*index=35*/f32[], f32[], f32[], f32[], f32[], /*index=40*/f32[], f32[], f32[], f32[], f32[], /*index=45*/f32[], f32[], f32[], f32[], f32[], /*index=50*/f32[], f32[], f32[], f32[], f32[], /*index=55*/f32[], f32[], f32[], f32[], f32[], /*index=60*/f32[], pred[614], f32[614], f32[614], f32[614], /*index=65*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=70*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=75*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=80*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=85*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=90*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=95*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=100*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=105*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=110*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=115*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=120*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=125*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=130*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=135*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=140*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=145*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=150*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=155*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=160*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=165*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=170*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=175*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=180*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=185*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=190*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=195*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=200*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=205*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=210*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=215*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=220*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=225*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=230*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=235*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=240*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=245*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=250*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=255*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=260*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=265*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=270*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=275*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=280*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=285*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=290*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=295*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=300*/f32[614], f32[], f32[], f32[], f32[], /*index=305*/f32[], f32[], f32[], f32[], f32[], /*index=310*/f32[], f32[], f32[], f32[], f32[], /*index=315*/f32[], f32[], f32[], f32[], f32[], /*index=320*/f32[], f32[], f32[], f32[], f32[], /*index=325*/f32[], f32[], f32[], f32[], f32[], /*index=330*/f32[], f32[], f32[], f32[], f32[], /*index=335*/f32[], f32[], f32[], f32[], f32[], /*index=340*/f32[], f32[], f32[], f32[], f32[], /*index=345*/f32[], f32[], f32[], f32[], f32[], /*index=350*/f32[], f32[], f32[], f32[], f32[], /*index=355*/f32[], f32[], f32[], f32[], f32[], /*index=360*/f32[])) -> (s32[], f32[], f32[], f32[], f32[], /*index=5*/f32[], f32[], f32[], f32[], f32[], /*index=10*/f32[], f32[], f32[], f32[], f32[], /*index=15*/f32[], f32[], f32[], f32[], f32[], /*index=20*/f32[], f32[], f32[], f32[], f32[], /*index=25*/f32[], f32[], f32[], f32[], f32[], /*index=30*/f32[], f32[], f32[], f32[], f32[], /*index=35*/f32[], f32[], f32[], f32[], f32[], /*index=40*/f32[], f32[], f32[], f32[], f32[], /*index=45*/f32[], f32[], f32[], f32[], f32[], /*index=50*/f32[], f32[], f32[], f32[], f32[], /*index=55*/f32[], f32[], f32[], f32[], f32[], /*index=60*/f32[], pred[614], f32[614], f32[614], f32[614], /*index=65*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=70*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=75*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=80*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=85*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=90*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=95*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=100*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=105*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=110*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=115*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=120*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=125*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=130*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=135*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=140*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=145*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=150*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=155*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=160*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=165*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=170*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=175*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=180*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=185*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=190*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=195*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=200*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=205*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=210*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=215*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=220*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=225*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=230*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=235*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=240*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=245*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=250*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=255*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=260*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=265*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=270*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=275*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=280*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=285*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=290*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=295*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=300*/f32[614], f32[], f32[], f32[], f32[], /*index=305*/f32[], f32[], f32[], f32[], f32[], /*index=310*/f32[], f32[], f32[], f32[], f32[], /*index=315*/f32[], f32[], f32[], f32[], f32[], /*index=320*/f32[], f32[], f32[], f32[], f32[], /*index=325*/f32[], f32[], f32[], f32[], f32[], /*index=330*/f32[], f32[], f32[], f32[], f32[], /*index=335*/f32[], f32[], f32[], f32[], f32[], /*index=340*/f32[], f32[], f32[], f32[], f32[], /*index=345*/f32[], f32[], f32[], f32[], f32[], /*index=350*/f32[], f32[], f32[], f32[], f32[], /*index=355*/f32[], f32[], f32[], f32[], f32[], /*index=360*/f32[]) {
+  %arg_tuple.1 = (s32[], f32[], f32[], f32[], f32[], /*index=5*/f32[], f32[], f32[], f32[], f32[], /*index=10*/f32[], f32[], f32[], f32[], f32[], /*index=15*/f32[], f32[], f32[], f32[], f32[], /*index=20*/f32[], f32[], f32[], f32[], f32[], /*index=25*/f32[], f32[], f32[], f32[], f32[], /*index=30*/f32[], f32[], f32[], f32[], f32[], /*index=35*/f32[], f32[], f32[], f32[], f32[], /*index=40*/f32[], f32[], f32[], f32[], f32[], /*index=45*/f32[], f32[], f32[], f32[], f32[], /*index=50*/f32[], f32[], f32[], f32[], f32[], /*index=55*/f32[], f32[], f32[], f32[], f32[], /*index=60*/f32[], pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=65*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=70*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=75*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=80*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=85*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=90*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=95*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=100*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=105*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=110*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=115*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=120*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=125*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=130*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=135*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=140*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=145*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=150*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=155*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=160*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=165*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=170*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=175*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=180*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=185*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=190*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=195*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=200*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=205*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=210*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=215*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=220*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=225*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=230*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=235*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=240*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=245*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=250*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=255*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=260*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=265*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=270*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=275*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=280*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=285*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=290*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=295*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=300*/f32[614]{0}, f32[], f32[], f32[], f32[], /*index=305*/f32[], f32[], f32[], f32[], f32[], /*index=310*/f32[], f32[], f32[], f32[], f32[], /*index=315*/f32[], f32[], f32[], f32[], f32[], /*index=320*/f32[], f32[], f32[], f32[], f32[], /*index=325*/f32[], f32[], f32[], f32[], f32[], /*index=330*/f32[], f32[], f32[], f32[], f32[], /*index=335*/f32[], f32[], f32[], f32[], f32[], /*index=340*/f32[], f32[], f32[], f32[], f32[], /*index=345*/f32[], f32[], f32[], f32[], f32[], /*index=350*/f32[], f32[], f32[], f32[], f32[], /*index=355*/f32[], f32[], f32[], f32[], f32[], /*index=360*/f32[]) parameter(0)
+  %constant.13 = s32[] constant(1)
+  %add.481 = s32[] add(%arg_tuple.1#0, %constant.13)
+  %constant.17 = f32[] constant(0.5)
+  %mul.360 = f32[] multiply(%arg_tuple.1#1, %constant.17)
+  %constant.16 = f32[] constant(0.1)
+  %mul.361 = f32[] multiply(%arg_tuple.1#2, %constant.16)
+  %add.241 = f32[] add(%mul.360, %mul.361)
+  %add.242 = f32[] add(%add.241, %arg_tuple.1#301)
+  %constant.18 = f32[614]{0} constant({0, 0.00999983307, 0.0199986659, 0.0299955, 0.0399893336, 0.0499791652, 0.0599640049, 0.0699428469, 0.0799146891, 0.0898785442, 0.099833414, 0.1097783, 0.119712204, 0.129634142, 0.139543116, 0.149438128, 0.159318209, 0.169182345, 0.179029569, 0.188858896, 0.198669314, 0.208459899, 0.218229622, 0.227977514, 0.237702623, 0.247403964, 0.257080555, 0.266731411, 0.276355654, 0.28595221, 0.295520186, 0.305058628, 0.314566553, 0.324043, 0.333487093, 0.342897803, 0.352274209, 0.361615449, 0.370920479, 0.380188406, 0.389418334, 0.39860931, 0.407760441, 0.416870773, 0.425939471, 0.434965521, 0.44394809, 0.452886283, 0.461779177, 0.470625877, 0.47942555, 0.48817724, 0.496880114, 0.505533338, 0.514135957, 0.522687256, 0.531186223, 0.539632, 0.548023939, 0.556361, 0.564642429, 0.572867453, 0.581035137, 0.589144766, 0.597195446, 0.605186403, 0.613116801, 0.620985925, 0.628793, 0.636537194, 0.64421767, 0.651833773, 0.659384668, 0.666869581, 0.674287915, 0.681638777, 0.688921452, 0.696135223, 0.703279376, 0.710353255, 0.717356086, 0.724287152, 0.731145799, 0.737931371, 0.744643092, 0.751280367, 0.757842541, 0.764328957, 0.7707389, 0.777071714, 0.783326924, 0.789503694, 0.795601606, 0.801619947, 0.807558119, 0.813415527, 0.819191575, 0.824885666, 0.830497324, 0.836025953, 0.841470957, 0.846831858, 0.852108, 0.857299, 0.862404227, 0.867423177, 0.872355461, 0.877200484, 0.881957769, 0.886626959, 0.891207397, 0.895698667, 0.90010047, 0.90441221, 0.908633471, 0.912763953, 0.916803122, 0.920750558, 0.924605966, 0.928368926, 0.932039082, 0.935615957, 0.939099371, 0.94248879, 0.945784032, 0.948984623, 0.952090323, 0.955100834, 0.958015859, 0.96083504, 0.963558197, 0.966184914, 0.968715072, 0.971148372, 0.973484516, 0.975723386, 0.977864623, 0.979908049, 0.981853545, 0.983700812, 0.985449731, 0.987100124, 0.988651752, 0.990104556, 0.991458356, 0.992713, 0.993868351, 0.994924366, 0.995880842, 0.996737778, 0.997495, 0.998152494, 0.998710155, 0.999167919, 0.999525845, 0.999783754, 0.999941707, 0.999999702, 0.999957621, 0.999815643, 0.999573588, 0.999231637, 0.998789728, 0.998248, 0.997606397, 0.996865034, 0.996024, 0.995083332, 0.994043231, 0.99290365, 0.991664827, 0.990326822, 0.988889754, 0.987353861, 0.985719204, 0.98398596, 0.98215431, 0.98022449, 0.978196621, 0.97607094, 0.973847628, 0.971527, 0.969109118, 0.966594398, 0.963983, 0.96127522, 0.958471298, 0.955571532, 0.95257622, 0.9494856, 0.946300089, 0.943019927, 0.939645469, 0.936177075, 0.932615042, 0.928959727, 0.925211549, 0.921370864, 0.917438, 0.913413346, 0.909297407, 0.90509057, 0.900793195, 0.896405756, 0.891928673, 0.887362361, 0.882707357, 0.877964079, 0.873133, 0.868214607, 0.863209426, 0.858117878, 0.852940559, 0.847677886, 0.842330515, 0.836898863, 0.831383526, 0.8257851, 0.820103884, 0.81434083, 0.808496356, 0.802571058, 0.796565473, 0.790480196, 0.784315944, 0.778073192, 0.771752656, 0.765355, 0.758880734, 0.752330601, 0.745705247, 0.739005327, 0.732231498, 0.725384414, 0.718464851, 0.711473405, 0.704410851, 0.697277844, 0.690075099, 0.682803333, 0.675463259, 0.668055713, 0.66058135, 0.653040886, 0.645435, 0.637764692, 0.630030572, 0.62223351, 0.61437422, 0.606453538, 0.598472118, 0.590430915, 0.582330644, 0.574172199, 0.565956235, 0.557683766, 0.549355507, 0.540972292, 0.532534957, 0.524044394, 0.51550144, 0.506906927, 0.49826175, 0.489566714, 0.480822742, 0.472030669, 0.46319139, 0.454305828, 0.445374787, 0.436399043, 0.427379847, 0.418317914, 0.409214139, 0.400069445, 0.390884757, 0.381661, 0.372399062, 0.363099873, 0.353764385, 0.344393492, 0.334988207, 0.325549394, 0.316078037, 0.30657506, 0.297041446, 0.287478089, 0.277886033, 0.268266171, 0.258619457, 0.24894692, 0.239249468, 0.229528099, 0.219783768, 0.210017473, 0.200230166, 0.190422595, 0.180596232, 0.170751795, 0.160890296, 0.151012704, 0.14112, 0.131213203, 0.121293277, 0.11136122, 0.101418026, 0.0914646909, 0.0815022066, 0.0715315789, 0.0615537949, 0.051569853, 0.041580759, 0.0315875039, 0.0215910897, 0.0115925176, 0.0015927864, -0.00840710476, -0.0184061546, -0.0284033641, -0.0383977331, -0.0483882613, -0.0583739541, -0.0683540478, -0.0783270597, -0.0882922485, -0.0982486, -0.108195134, -0.118130848, -0.128054738, -0.137965843, -0.147863135, -0.157745644, -0.167612389, -0.177462354, -0.187294587, -0.19710809, -0.206901878, -0.216674984, -0.226426408, -0.236155197, -0.245860368, -0.255540967, -0.265196, -0.2748245, -0.284425557, -0.293998152, -0.303541332, -0.313054383, -0.322535932, -0.331985205, -0.341401279, -0.350783229, -0.360130101, -0.369440943, -0.378714859, -0.387950897, -0.397148132, -0.406305641, -0.415422559, -0.424497902, -0.433530807, -0.44252035, -0.451465666, -0.460365802, -0.469219923, -0.478027135, -0.486786515, -0.495497227, -0.504158437, -0.512769163, -0.521328628, -0.529836, -0.538290322, -0.54669106, -0.555036962, -0.563327312, -0.571561337, -0.5797382, -0.587857068, -0.595917225, -0.603917718, -0.611857831, -0.619736791, -0.627553761, -0.635307968, -0.642998695, -0.65062505, -0.658186376, -0.665681899, -0.673110843, -0.680472493, -0.687766075, -0.694990873, -0.702146173, -0.709231257, -0.716245472, -0.723188, -0.730058253, -0.736855626, -0.743579149, -0.750228345, -0.756802499, -0.763300836, -0.769723117, -0.776068151, -0.782335877, -0.788525105, -0.794635713, -0.800666571, -0.806617677, -0.812487841, -0.818277061, -0.823984444, -0.829609096, -0.835151136, -0.840609312, -0.845983744, -0.851273298, -0.856478035, -0.861596823, -0.86662972, -0.871575654, -0.876434743, -0.881205916, -0.885889232, -0.890483677, -0.894989371, -0.899405301, -0.903731525, -0.90796715, -0.912112176, -0.916165829, -0.920128047, -0.923998058, -0.92777586, -0.931460679, -0.935052514, -0.938550889, -0.941955268, -0.945265532, -0.948481202, -0.951602101, -0.954627752, -0.957558036, -0.960392416, -0.963130951, -0.965773, -0.968318641, -0.97076726, -0.973118961, -0.975373268, -0.977530122, -0.979589105, -0.981550276, -0.98341316, -0.985177755, -0.986843824, -0.988411248, -0.989879787, -0.991249382, -0.992519796, -0.993690968, -0.994762838, -0.995735168, -0.996607959, -0.997381032, -0.998054445, -0.998628, -0.999101698, -0.999475539, -0.999749362, -0.999923229, -0.999997139, -0.999971032, -0.999844909, -0.999618828, -0.999292791, -0.998866796, -0.998340964, -0.997715235, -0.996989787, -0.99616462, -0.995239854, -0.994215548, -0.993091822, -0.991868794, -0.990546525, -0.989125311, -0.987605095, -0.985986114, -0.984268606, -0.982452571, -0.980538428, -0.978526115, -0.976416051, -0.974208236, -0.971903086, -0.969500661, -0.967001438, -0.964405358, -0.961712956, -0.958924294, -0.956039846, -0.953059614, -0.949984312, -0.946813762, -0.943548739, -0.940189242, -0.936735809, -0.933188617, -0.929548204, -0.925814748, -0.921988785, -0.918070495, -0.914060414, -0.909959078, -0.905766606, -0.901483715, -0.897110462, -0.892647743, -0.888095558, -0.88345474, -0.87872535, -0.873908341, -0.869003713, -0.86401242, -0.858934522, -0.853770912, -0.848521709, -0.843187869, -0.837769508, -0.832267582, -0.82668221, -0.821014404, -0.815264285, -0.809432864, -0.803520203, -0.797527492, -0.791454792, -0.785302877, -0.779072762, -0.772764444, -0.766379118, -0.759916902, -0.753378928, -0.746765375, -0.740077436, -0.73331517, -0.726479888, -0.71957165, -0.712591767, -0.705540299, -0.698418617, -0.69122678, -0.68396616, -0.676636755, -0.669240057, -0.661776125, -0.65424633, -0.646650732, -0.638990879, -0.631266713, -0.623479843, -0.615630209, -0.607719362, -0.599747419, -0.591715515, -0.58362478, -0.575475335, -0.56726867, -0.559004962, -0.550685704, -0.542311, -0.533882439, -0.525400102, -0.516865611, -0.508279085, -0.499642104, -0.490954727, -0.482218713, -0.473434031, -0.464602441, -0.455723941, -0.446800321, -0.437831581, -0.428819507, -0.419764102, -0.410667151, -0.401528716, -0.392350554, -0.383132726, -0.373876572, -0.364583492, -0.355253488, -0.345888406, -0.336488307, -0.327055, -0.317588538, -0.308090746, -0.298561722, -0.289003283, -0.279415488, -0.269800216, -0.260157496, -0.250489235, -0.240795463, -0.231078073, -0.221337095, -0.211574465, -0.201790199, -0.191986233, -0.182162598, -0.172321215, -0.16246213, -0.152587265})
+  %jit_dynamic_index_in_dim_.1 = f32[] call(%constant.18, %arg_tuple.1#0), to_apply=%dynamic_index_in_dim.1
+  %add.243 = f32[] add(%add.242, %jit_dynamic_index_in_dim_.1)
+  %abs.60 = f32[] abs(%add.243)
+  %neg.60 = f32[] negate(%abs.60)
+  %mul.362 = f32[] multiply(%neg.60, %constant.16)
+  %exp.60 = f32[] exponential(%mul.362)
+  %mul.363 = f32[] multiply(%add.243, %exp.60)
+  %tanh.60 = f32[] tanh(%mul.363)
+  %mul.365 = f32[] multiply(%arg_tuple.1#2, %constant.17)
+  %mul.366 = f32[] multiply(%arg_tuple.1#3, %constant.16)
+  %add.245 = f32[] add(%mul.365, %mul.366)
+  %add.246 = f32[] add(%add.245, %arg_tuple.1#302)
+  %add.247 = f32[] add(%add.246, %jit_dynamic_index_in_dim_.1)
+  %abs.61 = f32[] abs(%add.247)
+  %neg.61 = f32[] negate(%abs.61)
+  %mul.367 = f32[] multiply(%neg.61, %constant.16)
+  %exp.61 = f32[] exponential(%mul.367)
+  %mul.368 = f32[] multiply(%add.247, %exp.61)
+  %tanh.61 = f32[] tanh(%mul.368)
+  %mul.370 = f32[] multiply(%arg_tuple.1#3, %constant.17)
+  %mul.371 = f32[] multiply(%arg_tuple.1#4, %constant.16)
+  %add.249 = f32[] add(%mul.370, %mul.371)
+  %add.250 = f32[] add(%add.249, %arg_tuple.1#303)
+  %add.251 = f32[] add(%add.250, %jit_dynamic_index_in_dim_.1)
+  %abs.62 = f32[] abs(%add.251)
+  %neg.62 = f32[] negate(%abs.62)
+  %mul.372 = f32[] multiply(%neg.62, %constant.16)
+  %exp.62 = f32[] exponential(%mul.372)
+  %mul.373 = f32[] multiply(%add.251, %exp.62)
+  %tanh.62 = f32[] tanh(%mul.373)
+  %mul.375 = f32[] multiply(%arg_tuple.1#4, %constant.17)
+  %mul.376 = f32[] multiply(%arg_tuple.1#5, %constant.16)
+  %add.253 = f32[] add(%mul.375, %mul.376)
+  %add.254 = f32[] add(%add.253, %arg_tuple.1#304)
+  %add.255 = f32[] add(%add.254, %jit_dynamic_index_in_dim_.1)
+  %abs.63 = f32[] abs(%add.255)
+  %neg.63 = f32[] negate(%abs.63)
+  %mul.377 = f32[] multiply(%neg.63, %constant.16)
+  %exp.63 = f32[] exponential(%mul.377)
+  %mul.378 = f32[] multiply(%add.255, %exp.63)
+  %tanh.63 = f32[] tanh(%mul.378)
+  %mul.380 = f32[] multiply(%arg_tuple.1#5, %constant.17)
+  %mul.381 = f32[] multiply(%arg_tuple.1#6, %constant.16)
+  %add.257 = f32[] add(%mul.380, %mul.381)
+  %add.258 = f32[] add(%add.257, %arg_tuple.1#305)
+  %add.259 = f32[] add(%add.258, %jit_dynamic_index_in_dim_.1)
+  %abs.64 = f32[] abs(%add.259)
+  %neg.64 = f32[] negate(%abs.64)
+  %mul.382 = f32[] multiply(%neg.64, %constant.16)
+  %exp.64 = f32[] exponential(%mul.382)
+  %mul.383 = f32[] multiply(%add.259, %exp.64)
+  %tanh.64 = f32[] tanh(%mul.383)
+  %mul.385 = f32[] multiply(%arg_tuple.1#6, %constant.17)
+  %mul.386 = f32[] multiply(%arg_tuple.1#7, %constant.16)
+  %add.261 = f32[] add(%mul.385, %mul.386)
+  %add.262 = f32[] add(%add.261, %arg_tuple.1#306)
+  %add.263 = f32[] add(%add.262, %jit_dynamic_index_in_dim_.1)
+  %abs.65 = f32[] abs(%add.263)
+  %neg.65 = f32[] negate(%abs.65)
+  %mul.387 = f32[] multiply(%neg.65, %constant.16)
+  %exp.65 = f32[] exponential(%mul.387)
+  %mul.388 = f32[] multiply(%add.263, %exp.65)
+  %tanh.65 = f32[] tanh(%mul.388)
+  %mul.390 = f32[] multiply(%arg_tuple.1#7, %constant.17)
+  %mul.391 = f32[] multiply(%arg_tuple.1#8, %constant.16)
+  %add.265 = f32[] add(%mul.390, %mul.391)
+  %add.266 = f32[] add(%add.265, %arg_tuple.1#307)
+  %add.267 = f32[] add(%add.266, %jit_dynamic_index_in_dim_.1)
+  %abs.66 = f32[] abs(%add.267)
+  %neg.66 = f32[] negate(%abs.66)
+  %mul.392 = f32[] multiply(%neg.66, %constant.16)
+  %exp.66 = f32[] exponential(%mul.392)
+  %mul.393 = f32[] multiply(%add.267, %exp.66)
+  %tanh.66 = f32[] tanh(%mul.393)
+  %mul.395 = f32[] multiply(%arg_tuple.1#8, %constant.17)
+  %mul.396 = f32[] multiply(%arg_tuple.1#9, %constant.16)
+  %add.269 = f32[] add(%mul.395, %mul.396)
+  %add.270 = f32[] add(%add.269, %arg_tuple.1#308)
+  %add.271 = f32[] add(%add.270, %jit_dynamic_index_in_dim_.1)
+  %abs.67 = f32[] abs(%add.271)
+  %neg.67 = f32[] negate(%abs.67)
+  %mul.397 = f32[] multiply(%neg.67, %constant.16)
+  %exp.67 = f32[] exponential(%mul.397)
+  %mul.398 = f32[] multiply(%add.271, %exp.67)
+  %tanh.67 = f32[] tanh(%mul.398)
+  %mul.400 = f32[] multiply(%arg_tuple.1#9, %constant.17)
+  %mul.401 = f32[] multiply(%arg_tuple.1#10, %constant.16)
+  %add.273 = f32[] add(%mul.400, %mul.401)
+  %add.274 = f32[] add(%add.273, %arg_tuple.1#309)
+  %add.275 = f32[] add(%add.274, %jit_dynamic_index_in_dim_.1)
+  %abs.68 = f32[] abs(%add.275)
+  %neg.68 = f32[] negate(%abs.68)
+  %mul.402 = f32[] multiply(%neg.68, %constant.16)
+  %exp.68 = f32[] exponential(%mul.402)
+  %mul.403 = f32[] multiply(%add.275, %exp.68)
+  %tanh.68 = f32[] tanh(%mul.403)
+  %mul.405 = f32[] multiply(%arg_tuple.1#10, %constant.17)
+  %mul.406 = f32[] multiply(%arg_tuple.1#11, %constant.16)
+  %add.277 = f32[] add(%mul.405, %mul.406)
+  %add.278 = f32[] add(%add.277, %arg_tuple.1#310)
+  %add.279 = f32[] add(%add.278, %jit_dynamic_index_in_dim_.1)
+  %abs.69 = f32[] abs(%add.279)
+  %neg.69 = f32[] negate(%abs.69)
+  %mul.407 = f32[] multiply(%neg.69, %constant.16)
+  %exp.69 = f32[] exponential(%mul.407)
+  %mul.408 = f32[] multiply(%add.279, %exp.69)
+  %tanh.69 = f32[] tanh(%mul.408)
+  %mul.410 = f32[] multiply(%arg_tuple.1#11, %constant.17)
+  %mul.411 = f32[] multiply(%arg_tuple.1#12, %constant.16)
+  %add.281 = f32[] add(%mul.410, %mul.411)
+  %add.282 = f32[] add(%add.281, %arg_tuple.1#311)
+  %add.283 = f32[] add(%add.282, %jit_dynamic_index_in_dim_.1)
+  %abs.70 = f32[] abs(%add.283)
+  %neg.70 = f32[] negate(%abs.70)
+  %mul.412 = f32[] multiply(%neg.70, %constant.16)
+  %exp.70 = f32[] exponential(%mul.412)
+  %mul.413 = f32[] multiply(%add.283, %exp.70)
+  %tanh.70 = f32[] tanh(%mul.413)
+  %mul.415 = f32[] multiply(%arg_tuple.1#12, %constant.17)
+  %mul.416 = f32[] multiply(%arg_tuple.1#13, %constant.16)
+  %add.285 = f32[] add(%mul.415, %mul.416)
+  %add.286 = f32[] add(%add.285, %arg_tuple.1#312)
+  %add.287 = f32[] add(%add.286, %jit_dynamic_index_in_dim_.1)
+  %abs.71 = f32[] abs(%add.287)
+  %neg.71 = f32[] negate(%abs.71)
+  %mul.417 = f32[] multiply(%neg.71, %constant.16)
+  %exp.71 = f32[] exponential(%mul.417)
+  %mul.418 = f32[] multiply(%add.287, %exp.71)
+  %tanh.71 = f32[] tanh(%mul.418)
+  %mul.420 = f32[] multiply(%arg_tuple.1#13, %constant.17)
+  %mul.421 = f32[] multiply(%arg_tuple.1#14, %constant.16)
+  %add.289 = f32[] add(%mul.420, %mul.421)
+  %add.290 = f32[] add(%add.289, %arg_tuple.1#313)
+  %add.291 = f32[] add(%add.290, %jit_dynamic_index_in_dim_.1)
+  %abs.72 = f32[] abs(%add.291)
+  %neg.72 = f32[] negate(%abs.72)
+  %mul.422 = f32[] multiply(%neg.72, %constant.16)
+  %exp.72 = f32[] exponential(%mul.422)
+  %mul.423 = f32[] multiply(%add.291, %exp.72)
+  %tanh.72 = f32[] tanh(%mul.423)
+  %mul.425 = f32[] multiply(%arg_tuple.1#14, %constant.17)
+  %mul.426 = f32[] multiply(%arg_tuple.1#15, %constant.16)
+  %add.293 = f32[] add(%mul.425, %mul.426)
+  %add.294 = f32[] add(%add.293, %arg_tuple.1#314)
+  %add.295 = f32[] add(%add.294, %jit_dynamic_index_in_dim_.1)
+  %abs.73 = f32[] abs(%add.295)
+  %neg.73 = f32[] negate(%abs.73)
+  %mul.427 = f32[] multiply(%neg.73, %constant.16)
+  %exp.73 = f32[] exponential(%mul.427)
+  %mul.428 = f32[] multiply(%add.295, %exp.73)
+  %tanh.73 = f32[] tanh(%mul.428)
+  %mul.430 = f32[] multiply(%arg_tuple.1#15, %constant.17)
+  %mul.431 = f32[] multiply(%arg_tuple.1#16, %constant.16)
+  %add.297 = f32[] add(%mul.430, %mul.431)
+  %add.298 = f32[] add(%add.297, %arg_tuple.1#315)
+  %add.299 = f32[] add(%add.298, %jit_dynamic_index_in_dim_.1)
+  %abs.74 = f32[] abs(%add.299)
+  %neg.74 = f32[] negate(%abs.74)
+  %mul.432 = f32[] multiply(%neg.74, %constant.16)
+  %exp.74 = f32[] exponential(%mul.432)
+  %mul.433 = f32[] multiply(%add.299, %exp.74)
+  %tanh.74 = f32[] tanh(%mul.433)
+  %mul.435 = f32[] multiply(%arg_tuple.1#16, %constant.17)
+  %mul.436 = f32[] multiply(%arg_tuple.1#17, %constant.16)
+  %add.301 = f32[] add(%mul.435, %mul.436)
+  %add.302 = f32[] add(%add.301, %arg_tuple.1#316)
+  %add.303 = f32[] add(%add.302, %jit_dynamic_index_in_dim_.1)
+  %abs.75 = f32[] abs(%add.303)
+  %neg.75 = f32[] negate(%abs.75)
+  %mul.437 = f32[] multiply(%neg.75, %constant.16)
+  %exp.75 = f32[] exponential(%mul.437)
+  %mul.438 = f32[] multiply(%add.303, %exp.75)
+  %tanh.75 = f32[] tanh(%mul.438)
+  %mul.440 = f32[] multiply(%arg_tuple.1#17, %constant.17)
+  %mul.441 = f32[] multiply(%arg_tuple.1#18, %constant.16)
+  %add.305 = f32[] add(%mul.440, %mul.441)
+  %add.306 = f32[] add(%add.305, %arg_tuple.1#317)
+  %add.307 = f32[] add(%add.306, %jit_dynamic_index_in_dim_.1)
+  %abs.76 = f32[] abs(%add.307)
+  %neg.76 = f32[] negate(%abs.76)
+  %mul.442 = f32[] multiply(%neg.76, %constant.16)
+  %exp.76 = f32[] exponential(%mul.442)
+  %mul.443 = f32[] multiply(%add.307, %exp.76)
+  %tanh.76 = f32[] tanh(%mul.443)
+  %mul.445 = f32[] multiply(%arg_tuple.1#18, %constant.17)
+  %mul.446 = f32[] multiply(%arg_tuple.1#19, %constant.16)
+  %add.309 = f32[] add(%mul.445, %mul.446)
+  %add.310 = f32[] add(%add.309, %arg_tuple.1#318)
+  %add.311 = f32[] add(%add.310, %jit_dynamic_index_in_dim_.1)
+  %abs.77 = f32[] abs(%add.311)
+  %neg.77 = f32[] negate(%abs.77)
+  %mul.447 = f32[] multiply(%neg.77, %constant.16)
+  %exp.77 = f32[] exponential(%mul.447)
+  %mul.448 = f32[] multiply(%add.311, %exp.77)
+  %tanh.77 = f32[] tanh(%mul.448)
+  %mul.450 = f32[] multiply(%arg_tuple.1#19, %constant.17)
+  %mul.451 = f32[] multiply(%arg_tuple.1#20, %constant.16)
+  %add.313 = f32[] add(%mul.450, %mul.451)
+  %add.314 = f32[] add(%add.313, %arg_tuple.1#319)
+  %add.315 = f32[] add(%add.314, %jit_dynamic_index_in_dim_.1)
+  %abs.78 = f32[] abs(%add.315)
+  %neg.78 = f32[] negate(%abs.78)
+  %mul.452 = f32[] multiply(%neg.78, %constant.16)
+  %exp.78 = f32[] exponential(%mul.452)
+  %mul.453 = f32[] multiply(%add.315, %exp.78)
+  %tanh.78 = f32[] tanh(%mul.453)
+  %mul.455 = f32[] multiply(%arg_tuple.1#20, %constant.17)
+  %mul.456 = f32[] multiply(%arg_tuple.1#21, %constant.16)
+  %add.317 = f32[] add(%mul.455, %mul.456)
+  %add.318 = f32[] add(%add.317, %arg_tuple.1#320)
+  %add.319 = f32[] add(%add.318, %jit_dynamic_index_in_dim_.1)
+  %abs.79 = f32[] abs(%add.319)
+  %neg.79 = f32[] negate(%abs.79)
+  %mul.457 = f32[] multiply(%neg.79, %constant.16)
+  %exp.79 = f32[] exponential(%mul.457)
+  %mul.458 = f32[] multiply(%add.319, %exp.79)
+  %tanh.79 = f32[] tanh(%mul.458)
+  %mul.460 = f32[] multiply(%arg_tuple.1#21, %constant.17)
+  %mul.461 = f32[] multiply(%arg_tuple.1#22, %constant.16)
+  %add.321 = f32[] add(%mul.460, %mul.461)
+  %add.322 = f32[] add(%add.321, %arg_tuple.1#321)
+  %add.323 = f32[] add(%add.322, %jit_dynamic_index_in_dim_.1)
+  %abs.80 = f32[] abs(%add.323)
+  %neg.80 = f32[] negate(%abs.80)
+  %mul.462 = f32[] multiply(%neg.80, %constant.16)
+  %exp.80 = f32[] exponential(%mul.462)
+  %mul.463 = f32[] multiply(%add.323, %exp.80)
+  %tanh.80 = f32[] tanh(%mul.463)
+  %mul.465 = f32[] multiply(%arg_tuple.1#22, %constant.17)
+  %mul.466 = f32[] multiply(%arg_tuple.1#23, %constant.16)
+  %add.325 = f32[] add(%mul.465, %mul.466)
+  %add.326 = f32[] add(%add.325, %arg_tuple.1#322)
+  %add.327 = f32[] add(%add.326, %jit_dynamic_index_in_dim_.1)
+  %abs.81 = f32[] abs(%add.327)
+  %neg.81 = f32[] negate(%abs.81)
+  %mul.467 = f32[] multiply(%neg.81, %constant.16)
+  %exp.81 = f32[] exponential(%mul.467)
+  %mul.468 = f32[] multiply(%add.327, %exp.81)
+  %tanh.81 = f32[] tanh(%mul.468)
+  %mul.470 = f32[] multiply(%arg_tuple.1#23, %constant.17)
+  %mul.471 = f32[] multiply(%arg_tuple.1#24, %constant.16)
+  %add.329 = f32[] add(%mul.470, %mul.471)
+  %add.330 = f32[] add(%add.329, %arg_tuple.1#323)
+  %add.331 = f32[] add(%add.330, %jit_dynamic_index_in_dim_.1)
+  %abs.82 = f32[] abs(%add.331)
+  %neg.82 = f32[] negate(%abs.82)
+  %mul.472 = f32[] multiply(%neg.82, %constant.16)
+  %exp.82 = f32[] exponential(%mul.472)
+  %mul.473 = f32[] multiply(%add.331, %exp.82)
+  %tanh.82 = f32[] tanh(%mul.473)
+  %mul.475 = f32[] multiply(%arg_tuple.1#24, %constant.17)
+  %mul.476 = f32[] multiply(%arg_tuple.1#25, %constant.16)
+  %add.333 = f32[] add(%mul.475, %mul.476)
+  %add.334 = f32[] add(%add.333, %arg_tuple.1#324)
+  %add.335 = f32[] add(%add.334, %jit_dynamic_index_in_dim_.1)
+  %abs.83 = f32[] abs(%add.335)
+  %neg.83 = f32[] negate(%abs.83)
+  %mul.477 = f32[] multiply(%neg.83, %constant.16)
+  %exp.83 = f32[] exponential(%mul.477)
+  %mul.478 = f32[] multiply(%add.335, %exp.83)
+  %tanh.83 = f32[] tanh(%mul.478)
+  %mul.480 = f32[] multiply(%arg_tuple.1#25, %constant.17)
+  %mul.481 = f32[] multiply(%arg_tuple.1#26, %constant.16)
+  %add.337 = f32[] add(%mul.480, %mul.481)
+  %add.338 = f32[] add(%add.337, %arg_tuple.1#325)
+  %add.339 = f32[] add(%add.338, %jit_dynamic_index_in_dim_.1)
+  %abs.84 = f32[] abs(%add.339)
+  %neg.84 = f32[] negate(%abs.84)
+  %mul.482 = f32[] multiply(%neg.84, %constant.16)
+  %exp.84 = f32[] exponential(%mul.482)
+  %mul.483 = f32[] multiply(%add.339, %exp.84)
+  %tanh.84 = f32[] tanh(%mul.483)
+  %mul.485 = f32[] multiply(%arg_tuple.1#26, %constant.17)
+  %mul.486 = f32[] multiply(%arg_tuple.1#27, %constant.16)
+  %add.341 = f32[] add(%mul.485, %mul.486)
+  %add.342 = f32[] add(%add.341, %arg_tuple.1#326)
+  %add.343 = f32[] add(%add.342, %jit_dynamic_index_in_dim_.1)
+  %abs.85 = f32[] abs(%add.343)
+  %neg.85 = f32[] negate(%abs.85)
+  %mul.487 = f32[] multiply(%neg.85, %constant.16)
+  %exp.85 = f32[] exponential(%mul.487)
+  %mul.488 = f32[] multiply(%add.343, %exp.85)
+  %tanh.85 = f32[] tanh(%mul.488)
+  %mul.490 = f32[] multiply(%arg_tuple.1#27, %constant.17)
+  %mul.491 = f32[] multiply(%arg_tuple.1#28, %constant.16)
+  %add.345 = f32[] add(%mul.490, %mul.491)
+  %add.346 = f32[] add(%add.345, %arg_tuple.1#327)
+  %add.347 = f32[] add(%add.346, %jit_dynamic_index_in_dim_.1)
+  %abs.86 = f32[] abs(%add.347)
+  %neg.86 = f32[] negate(%abs.86)
+  %mul.492 = f32[] multiply(%neg.86, %constant.16)
+  %exp.86 = f32[] exponential(%mul.492)
+  %mul.493 = f32[] multiply(%add.347, %exp.86)
+  %tanh.86 = f32[] tanh(%mul.493)
+  %mul.495 = f32[] multiply(%arg_tuple.1#28, %constant.17)
+  %mul.496 = f32[] multiply(%arg_tuple.1#29, %constant.16)
+  %add.349 = f32[] add(%mul.495, %mul.496)
+  %add.350 = f32[] add(%add.349, %arg_tuple.1#328)
+  %add.351 = f32[] add(%add.350, %jit_dynamic_index_in_dim_.1)
+  %abs.87 = f32[] abs(%add.351)
+  %neg.87 = f32[] negate(%abs.87)
+  %mul.497 = f32[] multiply(%neg.87, %constant.16)
+  %exp.87 = f32[] exponential(%mul.497)
+  %mul.498 = f32[] multiply(%add.351, %exp.87)
+  %tanh.87 = f32[] tanh(%mul.498)
+  %mul.500 = f32[] multiply(%arg_tuple.1#29, %constant.17)
+  %mul.501 = f32[] multiply(%arg_tuple.1#30, %constant.16)
+  %add.353 = f32[] add(%mul.500, %mul.501)
+  %add.354 = f32[] add(%add.353, %arg_tuple.1#329)
+  %add.355 = f32[] add(%add.354, %jit_dynamic_index_in_dim_.1)
+  %abs.88 = f32[] abs(%add.355)
+  %neg.88 = f32[] negate(%abs.88)
+  %mul.502 = f32[] multiply(%neg.88, %constant.16)
+  %exp.88 = f32[] exponential(%mul.502)
+  %mul.503 = f32[] multiply(%add.355, %exp.88)
+  %tanh.88 = f32[] tanh(%mul.503)
+  %mul.505 = f32[] multiply(%arg_tuple.1#30, %constant.17)
+  %mul.506 = f32[] multiply(%arg_tuple.1#31, %constant.16)
+  %add.357 = f32[] add(%mul.505, %mul.506)
+  %add.358 = f32[] add(%add.357, %arg_tuple.1#330)
+  %add.359 = f32[] add(%add.358, %jit_dynamic_index_in_dim_.1)
+  %abs.89 = f32[] abs(%add.359)
+  %neg.89 = f32[] negate(%abs.89)
+  %mul.507 = f32[] multiply(%neg.89, %constant.16)
+  %exp.89 = f32[] exponential(%mul.507)
+  %mul.508 = f32[] multiply(%add.359, %exp.89)
+  %tanh.89 = f32[] tanh(%mul.508)
+  %mul.510 = f32[] multiply(%arg_tuple.1#31, %constant.17)
+  %mul.511 = f32[] multiply(%arg_tuple.1#32, %constant.16)
+  %add.361 = f32[] add(%mul.510, %mul.511)
+  %add.362 = f32[] add(%add.361, %arg_tuple.1#331)
+  %add.363 = f32[] add(%add.362, %jit_dynamic_index_in_dim_.1)
+  %abs.90 = f32[] abs(%add.363)
+  %neg.90 = f32[] negate(%abs.90)
+  %mul.512 = f32[] multiply(%neg.90, %constant.16)
+  %exp.90 = f32[] exponential(%mul.512)
+  %mul.513 = f32[] multiply(%add.363, %exp.90)
+  %tanh.90 = f32[] tanh(%mul.513)
+  %mul.515 = f32[] multiply(%arg_tuple.1#32, %constant.17)
+  %mul.516 = f32[] multiply(%arg_tuple.1#33, %constant.16)
+  %add.365 = f32[] add(%mul.515, %mul.516)
+  %add.366 = f32[] add(%add.365, %arg_tuple.1#332)
+  %add.367 = f32[] add(%add.366, %jit_dynamic_index_in_dim_.1)
+  %abs.91 = f32[] abs(%add.367)
+  %neg.91 = f32[] negate(%abs.91)
+  %mul.517 = f32[] multiply(%neg.91, %constant.16)
+  %exp.91 = f32[] exponential(%mul.517)
+  %mul.518 = f32[] multiply(%add.367, %exp.91)
+  %tanh.91 = f32[] tanh(%mul.518)
+  %mul.520 = f32[] multiply(%arg_tuple.1#33, %constant.17)
+  %mul.521 = f32[] multiply(%arg_tuple.1#34, %constant.16)
+  %add.369 = f32[] add(%mul.520, %mul.521)
+  %add.370 = f32[] add(%add.369, %arg_tuple.1#333)
+  %add.371 = f32[] add(%add.370, %jit_dynamic_index_in_dim_.1)
+  %abs.92 = f32[] abs(%add.371)
+  %neg.92 = f32[] negate(%abs.92)
+  %mul.522 = f32[] multiply(%neg.92, %constant.16)
+  %exp.92 = f32[] exponential(%mul.522)
+  %mul.523 = f32[] multiply(%add.371, %exp.92)
+  %tanh.92 = f32[] tanh(%mul.523)
+  %mul.525 = f32[] multiply(%arg_tuple.1#34, %constant.17)
+  %mul.526 = f32[] multiply(%arg_tuple.1#35, %constant.16)
+  %add.373 = f32[] add(%mul.525, %mul.526)
+  %add.374 = f32[] add(%add.373, %arg_tuple.1#334)
+  %add.375 = f32[] add(%add.374, %jit_dynamic_index_in_dim_.1)
+  %abs.93 = f32[] abs(%add.375)
+  %neg.93 = f32[] negate(%abs.93)
+  %mul.527 = f32[] multiply(%neg.93, %constant.16)
+  %exp.93 = f32[] exponential(%mul.527)
+  %mul.528 = f32[] multiply(%add.375, %exp.93)
+  %tanh.93 = f32[] tanh(%mul.528)
+  %mul.530 = f32[] multiply(%arg_tuple.1#35, %constant.17)
+  %mul.531 = f32[] multiply(%arg_tuple.1#36, %constant.16)
+  %add.377 = f32[] add(%mul.530, %mul.531)
+  %add.378 = f32[] add(%add.377, %arg_tuple.1#335)
+  %add.379 = f32[] add(%add.378, %jit_dynamic_index_in_dim_.1)
+  %abs.94 = f32[] abs(%add.379)
+  %neg.94 = f32[] negate(%abs.94)
+  %mul.532 = f32[] multiply(%neg.94, %constant.16)
+  %exp.94 = f32[] exponential(%mul.532)
+  %mul.533 = f32[] multiply(%add.379, %exp.94)
+  %tanh.94 = f32[] tanh(%mul.533)
+  %mul.535 = f32[] multiply(%arg_tuple.1#36, %constant.17)
+  %mul.536 = f32[] multiply(%arg_tuple.1#37, %constant.16)
+  %add.381 = f32[] add(%mul.535, %mul.536)
+  %add.382 = f32[] add(%add.381, %arg_tuple.1#336)
+  %add.383 = f32[] add(%add.382, %jit_dynamic_index_in_dim_.1)
+  %abs.95 = f32[] abs(%add.383)
+  %neg.95 = f32[] negate(%abs.95)
+  %mul.537 = f32[] multiply(%neg.95, %constant.16)
+  %exp.95 = f32[] exponential(%mul.537)
+  %mul.538 = f32[] multiply(%add.383, %exp.95)
+  %tanh.95 = f32[] tanh(%mul.538)
+  %mul.540 = f32[] multiply(%arg_tuple.1#37, %constant.17)
+  %mul.541 = f32[] multiply(%arg_tuple.1#38, %constant.16)
+  %add.385 = f32[] add(%mul.540, %mul.541)
+  %add.386 = f32[] add(%add.385, %arg_tuple.1#337)
+  %add.387 = f32[] add(%add.386, %jit_dynamic_index_in_dim_.1)
+  %abs.96 = f32[] abs(%add.387)
+  %neg.96 = f32[] negate(%abs.96)
+  %mul.542 = f32[] multiply(%neg.96, %constant.16)
+  %exp.96 = f32[] exponential(%mul.542)
+  %mul.543 = f32[] multiply(%add.387, %exp.96)
+  %tanh.96 = f32[] tanh(%mul.543)
+  %mul.545 = f32[] multiply(%arg_tuple.1#38, %constant.17)
+  %mul.546 = f32[] multiply(%arg_tuple.1#39, %constant.16)
+  %add.389 = f32[] add(%mul.545, %mul.546)
+  %add.390 = f32[] add(%add.389, %arg_tuple.1#338)
+  %add.391 = f32[] add(%add.390, %jit_dynamic_index_in_dim_.1)
+  %abs.97 = f32[] abs(%add.391)
+  %neg.97 = f32[] negate(%abs.97)
+  %mul.547 = f32[] multiply(%neg.97, %constant.16)
+  %exp.97 = f32[] exponential(%mul.547)
+  %mul.548 = f32[] multiply(%add.391, %exp.97)
+  %tanh.97 = f32[] tanh(%mul.548)
+  %mul.550 = f32[] multiply(%arg_tuple.1#39, %constant.17)
+  %mul.551 = f32[] multiply(%arg_tuple.1#40, %constant.16)
+  %add.393 = f32[] add(%mul.550, %mul.551)
+  %add.394 = f32[] add(%add.393, %arg_tuple.1#339)
+  %add.395 = f32[] add(%add.394, %jit_dynamic_index_in_dim_.1)
+  %abs.98 = f32[] abs(%add.395)
+  %neg.98 = f32[] negate(%abs.98)
+  %mul.552 = f32[] multiply(%neg.98, %constant.16)
+  %exp.98 = f32[] exponential(%mul.552)
+  %mul.553 = f32[] multiply(%add.395, %exp.98)
+  %tanh.98 = f32[] tanh(%mul.553)
+  %mul.555 = f32[] multiply(%arg_tuple.1#40, %constant.17)
+  %mul.556 = f32[] multiply(%arg_tuple.1#41, %constant.16)
+  %add.397 = f32[] add(%mul.555, %mul.556)
+  %add.398 = f32[] add(%add.397, %arg_tuple.1#340)
+  %add.399 = f32[] add(%add.398, %jit_dynamic_index_in_dim_.1)
+  %abs.99 = f32[] abs(%add.399)
+  %neg.99 = f32[] negate(%abs.99)
+  %mul.557 = f32[] multiply(%neg.99, %constant.16)
+  %exp.99 = f32[] exponential(%mul.557)
+  %mul.558 = f32[] multiply(%add.399, %exp.99)
+  %tanh.99 = f32[] tanh(%mul.558)
+  %mul.560 = f32[] multiply(%arg_tuple.1#41, %constant.17)
+  %mul.561 = f32[] multiply(%arg_tuple.1#42, %constant.16)
+  %add.401 = f32[] add(%mul.560, %mul.561)
+  %add.402 = f32[] add(%add.401, %arg_tuple.1#341)
+  %add.403 = f32[] add(%add.402, %jit_dynamic_index_in_dim_.1)
+  %abs.100 = f32[] abs(%add.403)
+  %neg.100 = f32[] negate(%abs.100)
+  %mul.562 = f32[] multiply(%neg.100, %constant.16)
+  %exp.100 = f32[] exponential(%mul.562)
+  %mul.563 = f32[] multiply(%add.403, %exp.100)
+  %tanh.100 = f32[] tanh(%mul.563)
+  %mul.565 = f32[] multiply(%arg_tuple.1#42, %constant.17)
+  %mul.566 = f32[] multiply(%arg_tuple.1#43, %constant.16)
+  %add.405 = f32[] add(%mul.565, %mul.566)
+  %add.406 = f32[] add(%add.405, %arg_tuple.1#342)
+  %add.407 = f32[] add(%add.406, %jit_dynamic_index_in_dim_.1)
+  %abs.101 = f32[] abs(%add.407)
+  %neg.101 = f32[] negate(%abs.101)
+  %mul.567 = f32[] multiply(%neg.101, %constant.16)
+  %exp.101 = f32[] exponential(%mul.567)
+  %mul.568 = f32[] multiply(%add.407, %exp.101)
+  %tanh.101 = f32[] tanh(%mul.568)
+  %mul.570 = f32[] multiply(%arg_tuple.1#43, %constant.17)
+  %mul.571 = f32[] multiply(%arg_tuple.1#44, %constant.16)
+  %add.409 = f32[] add(%mul.570, %mul.571)
+  %add.410 = f32[] add(%add.409, %arg_tuple.1#343)
+  %add.411 = f32[] add(%add.410, %jit_dynamic_index_in_dim_.1)
+  %abs.102 = f32[] abs(%add.411)
+  %neg.102 = f32[] negate(%abs.102)
+  %mul.572 = f32[] multiply(%neg.102, %constant.16)
+  %exp.102 = f32[] exponential(%mul.572)
+  %mul.573 = f32[] multiply(%add.411, %exp.102)
+  %tanh.102 = f32[] tanh(%mul.573)
+  %mul.575 = f32[] multiply(%arg_tuple.1#44, %constant.17)
+  %mul.576 = f32[] multiply(%arg_tuple.1#45, %constant.16)
+  %add.413 = f32[] add(%mul.575, %mul.576)
+  %add.414 = f32[] add(%add.413, %arg_tuple.1#344)
+  %add.415 = f32[] add(%add.414, %jit_dynamic_index_in_dim_.1)
+  %abs.103 = f32[] abs(%add.415)
+  %neg.103 = f32[] negate(%abs.103)
+  %mul.577 = f32[] multiply(%neg.103, %constant.16)
+  %exp.103 = f32[] exponential(%mul.577)
+  %mul.578 = f32[] multiply(%add.415, %exp.103)
+  %tanh.103 = f32[] tanh(%mul.578)
+  %mul.580 = f32[] multiply(%arg_tuple.1#45, %constant.17)
+  %mul.581 = f32[] multiply(%arg_tuple.1#46, %constant.16)
+  %add.417 = f32[] add(%mul.580, %mul.581)
+  %add.418 = f32[] add(%add.417, %arg_tuple.1#345)
+  %add.419 = f32[] add(%add.418, %jit_dynamic_index_in_dim_.1)
+  %abs.104 = f32[] abs(%add.419)
+  %neg.104 = f32[] negate(%abs.104)
+  %mul.582 = f32[] multiply(%neg.104, %constant.16)
+  %exp.104 = f32[] exponential(%mul.582)
+  %mul.583 = f32[] multiply(%add.419, %exp.104)
+  %tanh.104 = f32[] tanh(%mul.583)
+  %mul.585 = f32[] multiply(%arg_tuple.1#46, %constant.17)
+  %mul.586 = f32[] multiply(%arg_tuple.1#47, %constant.16)
+  %add.421 = f32[] add(%mul.585, %mul.586)
+  %add.422 = f32[] add(%add.421, %arg_tuple.1#346)
+  %add.423 = f32[] add(%add.422, %jit_dynamic_index_in_dim_.1)
+  %abs.105 = f32[] abs(%add.423)
+  %neg.105 = f32[] negate(%abs.105)
+  %mul.587 = f32[] multiply(%neg.105, %constant.16)
+  %exp.105 = f32[] exponential(%mul.587)
+  %mul.588 = f32[] multiply(%add.423, %exp.105)
+  %tanh.105 = f32[] tanh(%mul.588)
+  %mul.590 = f32[] multiply(%arg_tuple.1#47, %constant.17)
+  %mul.591 = f32[] multiply(%arg_tuple.1#48, %constant.16)
+  %add.425 = f32[] add(%mul.590, %mul.591)
+  %add.426 = f32[] add(%add.425, %arg_tuple.1#347)
+  %add.427 = f32[] add(%add.426, %jit_dynamic_index_in_dim_.1)
+  %abs.106 = f32[] abs(%add.427)
+  %neg.106 = f32[] negate(%abs.106)
+  %mul.592 = f32[] multiply(%neg.106, %constant.16)
+  %exp.106 = f32[] exponential(%mul.592)
+  %mul.593 = f32[] multiply(%add.427, %exp.106)
+  %tanh.106 = f32[] tanh(%mul.593)
+  %mul.595 = f32[] multiply(%arg_tuple.1#48, %constant.17)
+  %mul.596 = f32[] multiply(%arg_tuple.1#49, %constant.16)
+  %add.429 = f32[] add(%mul.595, %mul.596)
+  %add.430 = f32[] add(%add.429, %arg_tuple.1#348)
+  %add.431 = f32[] add(%add.430, %jit_dynamic_index_in_dim_.1)
+  %abs.107 = f32[] abs(%add.431)
+  %neg.107 = f32[] negate(%abs.107)
+  %mul.597 = f32[] multiply(%neg.107, %constant.16)
+  %exp.107 = f32[] exponential(%mul.597)
+  %mul.598 = f32[] multiply(%add.431, %exp.107)
+  %tanh.107 = f32[] tanh(%mul.598)
+  %mul.600 = f32[] multiply(%arg_tuple.1#49, %constant.17)
+  %mul.601 = f32[] multiply(%arg_tuple.1#50, %constant.16)
+  %add.433 = f32[] add(%mul.600, %mul.601)
+  %add.434 = f32[] add(%add.433, %arg_tuple.1#349)
+  %add.435 = f32[] add(%add.434, %jit_dynamic_index_in_dim_.1)
+  %abs.108 = f32[] abs(%add.435)
+  %neg.108 = f32[] negate(%abs.108)
+  %mul.602 = f32[] multiply(%neg.108, %constant.16)
+  %exp.108 = f32[] exponential(%mul.602)
+  %mul.603 = f32[] multiply(%add.435, %exp.108)
+  %tanh.108 = f32[] tanh(%mul.603)
+  %mul.605 = f32[] multiply(%arg_tuple.1#50, %constant.17)
+  %mul.606 = f32[] multiply(%arg_tuple.1#51, %constant.16)
+  %add.437 = f32[] add(%mul.605, %mul.606)
+  %add.438 = f32[] add(%add.437, %arg_tuple.1#350)
+  %add.439 = f32[] add(%add.438, %jit_dynamic_index_in_dim_.1)
+  %abs.109 = f32[] abs(%add.439)
+  %neg.109 = f32[] negate(%abs.109)
+  %mul.607 = f32[] multiply(%neg.109, %constant.16)
+  %exp.109 = f32[] exponential(%mul.607)
+  %mul.608 = f32[] multiply(%add.439, %exp.109)
+  %tanh.109 = f32[] tanh(%mul.608)
+  %mul.610 = f32[] multiply(%arg_tuple.1#51, %constant.17)
+  %mul.611 = f32[] multiply(%arg_tuple.1#52, %constant.16)
+  %add.441 = f32[] add(%mul.610, %mul.611)
+  %add.442 = f32[] add(%add.441, %arg_tuple.1#351)
+  %add.443 = f32[] add(%add.442, %jit_dynamic_index_in_dim_.1)
+  %abs.110 = f32[] abs(%add.443)
+  %neg.110 = f32[] negate(%abs.110)
+  %mul.612 = f32[] multiply(%neg.110, %constant.16)
+  %exp.110 = f32[] exponential(%mul.612)
+  %mul.613 = f32[] multiply(%add.443, %exp.110)
+  %tanh.110 = f32[] tanh(%mul.613)
+  %mul.615 = f32[] multiply(%arg_tuple.1#52, %constant.17)
+  %mul.616 = f32[] multiply(%arg_tuple.1#53, %constant.16)
+  %add.445 = f32[] add(%mul.615, %mul.616)
+  %add.446 = f32[] add(%add.445, %arg_tuple.1#352)
+  %add.447 = f32[] add(%add.446, %jit_dynamic_index_in_dim_.1)
+  %abs.111 = f32[] abs(%add.447)
+  %neg.111 = f32[] negate(%abs.111)
+  %mul.617 = f32[] multiply(%neg.111, %constant.16)
+  %exp.111 = f32[] exponential(%mul.617)
+  %mul.618 = f32[] multiply(%add.447, %exp.111)
+  %tanh.111 = f32[] tanh(%mul.618)
+  %mul.620 = f32[] multiply(%arg_tuple.1#53, %constant.17)
+  %mul.621 = f32[] multiply(%arg_tuple.1#54, %constant.16)
+  %add.449 = f32[] add(%mul.620, %mul.621)
+  %add.450 = f32[] add(%add.449, %arg_tuple.1#353)
+  %add.451 = f32[] add(%add.450, %jit_dynamic_index_in_dim_.1)
+  %abs.112 = f32[] abs(%add.451)
+  %neg.112 = f32[] negate(%abs.112)
+  %mul.622 = f32[] multiply(%neg.112, %constant.16)
+  %exp.112 = f32[] exponential(%mul.622)
+  %mul.623 = f32[] multiply(%add.451, %exp.112)
+  %tanh.112 = f32[] tanh(%mul.623)
+  %mul.625 = f32[] multiply(%arg_tuple.1#54, %constant.17)
+  %mul.626 = f32[] multiply(%arg_tuple.1#55, %constant.16)
+  %add.453 = f32[] add(%mul.625, %mul.626)
+  %add.454 = f32[] add(%add.453, %arg_tuple.1#354)
+  %add.455 = f32[] add(%add.454, %jit_dynamic_index_in_dim_.1)
+  %abs.113 = f32[] abs(%add.455)
+  %neg.113 = f32[] negate(%abs.113)
+  %mul.627 = f32[] multiply(%neg.113, %constant.16)
+  %exp.113 = f32[] exponential(%mul.627)
+  %mul.628 = f32[] multiply(%add.455, %exp.113)
+  %tanh.113 = f32[] tanh(%mul.628)
+  %mul.630 = f32[] multiply(%arg_tuple.1#55, %constant.17)
+  %mul.631 = f32[] multiply(%arg_tuple.1#56, %constant.16)
+  %add.457 = f32[] add(%mul.630, %mul.631)
+  %add.458 = f32[] add(%add.457, %arg_tuple.1#355)
+  %add.459 = f32[] add(%add.458, %jit_dynamic_index_in_dim_.1)
+  %abs.114 = f32[] abs(%add.459)
+  %neg.114 = f32[] negate(%abs.114)
+  %mul.632 = f32[] multiply(%neg.114, %constant.16)
+  %exp.114 = f32[] exponential(%mul.632)
+  %mul.633 = f32[] multiply(%add.459, %exp.114)
+  %tanh.114 = f32[] tanh(%mul.633)
+  %mul.635 = f32[] multiply(%arg_tuple.1#56, %constant.17)
+  %mul.636 = f32[] multiply(%arg_tuple.1#57, %constant.16)
+  %add.461 = f32[] add(%mul.635, %mul.636)
+  %add.462 = f32[] add(%add.461, %arg_tuple.1#356)
+  %add.463 = f32[] add(%add.462, %jit_dynamic_index_in_dim_.1)
+  %abs.115 = f32[] abs(%add.463)
+  %neg.115 = f32[] negate(%abs.115)
+  %mul.637 = f32[] multiply(%neg.115, %constant.16)
+  %exp.115 = f32[] exponential(%mul.637)
+  %mul.638 = f32[] multiply(%add.463, %exp.115)
+  %tanh.115 = f32[] tanh(%mul.638)
+  %mul.640 = f32[] multiply(%arg_tuple.1#57, %constant.17)
+  %mul.641 = f32[] multiply(%arg_tuple.1#58, %constant.16)
+  %add.465 = f32[] add(%mul.640, %mul.641)
+  %add.466 = f32[] add(%add.465, %arg_tuple.1#357)
+  %add.467 = f32[] add(%add.466, %jit_dynamic_index_in_dim_.1)
+  %abs.116 = f32[] abs(%add.467)
+  %neg.116 = f32[] negate(%abs.116)
+  %mul.642 = f32[] multiply(%neg.116, %constant.16)
+  %exp.116 = f32[] exponential(%mul.642)
+  %mul.643 = f32[] multiply(%add.467, %exp.116)
+  %tanh.116 = f32[] tanh(%mul.643)
+  %mul.645 = f32[] multiply(%arg_tuple.1#58, %constant.17)
+  %mul.646 = f32[] multiply(%arg_tuple.1#59, %constant.16)
+  %add.469 = f32[] add(%mul.645, %mul.646)
+  %add.470 = f32[] add(%add.469, %arg_tuple.1#358)
+  %add.471 = f32[] add(%add.470, %jit_dynamic_index_in_dim_.1)
+  %abs.117 = f32[] abs(%add.471)
+  %neg.117 = f32[] negate(%abs.117)
+  %mul.647 = f32[] multiply(%neg.117, %constant.16)
+  %exp.117 = f32[] exponential(%mul.647)
+  %mul.648 = f32[] multiply(%add.471, %exp.117)
+  %tanh.117 = f32[] tanh(%mul.648)
+  %mul.650 = f32[] multiply(%arg_tuple.1#59, %constant.17)
+  %mul.651 = f32[] multiply(%arg_tuple.1#60, %constant.16)
+  %add.473 = f32[] add(%mul.650, %mul.651)
+  %add.474 = f32[] add(%add.473, %arg_tuple.1#359)
+  %add.475 = f32[] add(%add.474, %jit_dynamic_index_in_dim_.1)
+  %abs.118 = f32[] abs(%add.475)
+  %neg.118 = f32[] negate(%abs.118)
+  %mul.652 = f32[] multiply(%neg.118, %constant.16)
+  %exp.118 = f32[] exponential(%mul.652)
+  %mul.653 = f32[] multiply(%add.475, %exp.118)
+  %tanh.118 = f32[] tanh(%mul.653)
+  %mul.655 = f32[] multiply(%arg_tuple.1#60, %constant.17)
+  %mul.656 = f32[] multiply(%arg_tuple.1#1, %constant.16)
+  %add.477 = f32[] add(%mul.655, %mul.656)
+  %add.478 = f32[] add(%add.477, %arg_tuple.1#360)
+  %add.479 = f32[] add(%add.478, %jit_dynamic_index_in_dim_.1)
+  %abs.119 = f32[] abs(%add.479)
+  %neg.119 = f32[] negate(%abs.119)
+  %mul.657 = f32[] multiply(%neg.119, %constant.16)
+  %exp.119 = f32[] exponential(%mul.657)
+  %mul.658 = f32[] multiply(%add.479, %exp.119)
+  %tanh.119 = f32[] tanh(%mul.658)
+  %constant.15 = f32[] constant(0)
+  %ge.60 = pred[] compare(%add.243, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.240 = pred[614]{0} call(%arg_tuple.1#61, %ge.60, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.241 = f32[614]{0} call(%arg_tuple.1#62, %exp.60, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.242 = f32[614]{0} call(%arg_tuple.1#63, %add.243, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %constant.14 = f32[] constant(1)
+  %add.244 = f32[] add(%tanh.60, %constant.14)
+  %sub.60 = f32[] subtract(%constant.14, %tanh.60)
+  %mul.364 = f32[] multiply(%add.244, %sub.60)
+  %jit_dynamic_update_index_in_dim_.243 = f32[614]{0} call(%arg_tuple.1#64, %mul.364, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.61 = pred[] compare(%add.247, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.244 = pred[614]{0} call(%arg_tuple.1#65, %ge.61, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.245 = f32[614]{0} call(%arg_tuple.1#66, %exp.61, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.246 = f32[614]{0} call(%arg_tuple.1#67, %add.247, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.248 = f32[] add(%tanh.61, %constant.14)
+  %sub.61 = f32[] subtract(%constant.14, %tanh.61)
+  %mul.369 = f32[] multiply(%add.248, %sub.61)
+  %jit_dynamic_update_index_in_dim_.247 = f32[614]{0} call(%arg_tuple.1#68, %mul.369, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.62 = pred[] compare(%add.251, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.248 = pred[614]{0} call(%arg_tuple.1#69, %ge.62, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.249 = f32[614]{0} call(%arg_tuple.1#70, %exp.62, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.250 = f32[614]{0} call(%arg_tuple.1#71, %add.251, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.252 = f32[] add(%tanh.62, %constant.14)
+  %sub.62 = f32[] subtract(%constant.14, %tanh.62)
+  %mul.374 = f32[] multiply(%add.252, %sub.62)
+  %jit_dynamic_update_index_in_dim_.251 = f32[614]{0} call(%arg_tuple.1#72, %mul.374, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.63 = pred[] compare(%add.255, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.252 = pred[614]{0} call(%arg_tuple.1#73, %ge.63, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.253 = f32[614]{0} call(%arg_tuple.1#74, %exp.63, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.254 = f32[614]{0} call(%arg_tuple.1#75, %add.255, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.256 = f32[] add(%tanh.63, %constant.14)
+  %sub.63 = f32[] subtract(%constant.14, %tanh.63)
+  %mul.379 = f32[] multiply(%add.256, %sub.63)
+  %jit_dynamic_update_index_in_dim_.255 = f32[614]{0} call(%arg_tuple.1#76, %mul.379, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.64 = pred[] compare(%add.259, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.256 = pred[614]{0} call(%arg_tuple.1#77, %ge.64, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.257 = f32[614]{0} call(%arg_tuple.1#78, %exp.64, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.258 = f32[614]{0} call(%arg_tuple.1#79, %add.259, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.260 = f32[] add(%tanh.64, %constant.14)
+  %sub.64 = f32[] subtract(%constant.14, %tanh.64)
+  %mul.384 = f32[] multiply(%add.260, %sub.64)
+  %jit_dynamic_update_index_in_dim_.259 = f32[614]{0} call(%arg_tuple.1#80, %mul.384, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.65 = pred[] compare(%add.263, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.260 = pred[614]{0} call(%arg_tuple.1#81, %ge.65, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.261 = f32[614]{0} call(%arg_tuple.1#82, %exp.65, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.262 = f32[614]{0} call(%arg_tuple.1#83, %add.263, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.264 = f32[] add(%tanh.65, %constant.14)
+  %sub.65 = f32[] subtract(%constant.14, %tanh.65)
+  %mul.389 = f32[] multiply(%add.264, %sub.65)
+  %jit_dynamic_update_index_in_dim_.263 = f32[614]{0} call(%arg_tuple.1#84, %mul.389, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.66 = pred[] compare(%add.267, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.264 = pred[614]{0} call(%arg_tuple.1#85, %ge.66, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.265 = f32[614]{0} call(%arg_tuple.1#86, %exp.66, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.266 = f32[614]{0} call(%arg_tuple.1#87, %add.267, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.268 = f32[] add(%tanh.66, %constant.14)
+  %sub.66 = f32[] subtract(%constant.14, %tanh.66)
+  %mul.394 = f32[] multiply(%add.268, %sub.66)
+  %jit_dynamic_update_index_in_dim_.267 = f32[614]{0} call(%arg_tuple.1#88, %mul.394, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.67 = pred[] compare(%add.271, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.268 = pred[614]{0} call(%arg_tuple.1#89, %ge.67, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.269 = f32[614]{0} call(%arg_tuple.1#90, %exp.67, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.270 = f32[614]{0} call(%arg_tuple.1#91, %add.271, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.272 = f32[] add(%tanh.67, %constant.14)
+  %sub.67 = f32[] subtract(%constant.14, %tanh.67)
+  %mul.399 = f32[] multiply(%add.272, %sub.67)
+  %jit_dynamic_update_index_in_dim_.271 = f32[614]{0} call(%arg_tuple.1#92, %mul.399, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.68 = pred[] compare(%add.275, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.272 = pred[614]{0} call(%arg_tuple.1#93, %ge.68, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.273 = f32[614]{0} call(%arg_tuple.1#94, %exp.68, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.274 = f32[614]{0} call(%arg_tuple.1#95, %add.275, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.276 = f32[] add(%tanh.68, %constant.14)
+  %sub.68 = f32[] subtract(%constant.14, %tanh.68)
+  %mul.404 = f32[] multiply(%add.276, %sub.68)
+  %jit_dynamic_update_index_in_dim_.275 = f32[614]{0} call(%arg_tuple.1#96, %mul.404, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.69 = pred[] compare(%add.279, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.276 = pred[614]{0} call(%arg_tuple.1#97, %ge.69, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.277 = f32[614]{0} call(%arg_tuple.1#98, %exp.69, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.278 = f32[614]{0} call(%arg_tuple.1#99, %add.279, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.280 = f32[] add(%tanh.69, %constant.14)
+  %sub.69 = f32[] subtract(%constant.14, %tanh.69)
+  %mul.409 = f32[] multiply(%add.280, %sub.69)
+  %jit_dynamic_update_index_in_dim_.279 = f32[614]{0} call(%arg_tuple.1#100, %mul.409, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.70 = pred[] compare(%add.283, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.280 = pred[614]{0} call(%arg_tuple.1#101, %ge.70, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.281 = f32[614]{0} call(%arg_tuple.1#102, %exp.70, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.282 = f32[614]{0} call(%arg_tuple.1#103, %add.283, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.284 = f32[] add(%tanh.70, %constant.14)
+  %sub.70 = f32[] subtract(%constant.14, %tanh.70)
+  %mul.414 = f32[] multiply(%add.284, %sub.70)
+  %jit_dynamic_update_index_in_dim_.283 = f32[614]{0} call(%arg_tuple.1#104, %mul.414, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.71 = pred[] compare(%add.287, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.284 = pred[614]{0} call(%arg_tuple.1#105, %ge.71, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.285 = f32[614]{0} call(%arg_tuple.1#106, %exp.71, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.286 = f32[614]{0} call(%arg_tuple.1#107, %add.287, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.288 = f32[] add(%tanh.71, %constant.14)
+  %sub.71 = f32[] subtract(%constant.14, %tanh.71)
+  %mul.419 = f32[] multiply(%add.288, %sub.71)
+  %jit_dynamic_update_index_in_dim_.287 = f32[614]{0} call(%arg_tuple.1#108, %mul.419, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.72 = pred[] compare(%add.291, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.288 = pred[614]{0} call(%arg_tuple.1#109, %ge.72, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.289 = f32[614]{0} call(%arg_tuple.1#110, %exp.72, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.290 = f32[614]{0} call(%arg_tuple.1#111, %add.291, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.292 = f32[] add(%tanh.72, %constant.14)
+  %sub.72 = f32[] subtract(%constant.14, %tanh.72)
+  %mul.424 = f32[] multiply(%add.292, %sub.72)
+  %jit_dynamic_update_index_in_dim_.291 = f32[614]{0} call(%arg_tuple.1#112, %mul.424, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.73 = pred[] compare(%add.295, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.292 = pred[614]{0} call(%arg_tuple.1#113, %ge.73, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.293 = f32[614]{0} call(%arg_tuple.1#114, %exp.73, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.294 = f32[614]{0} call(%arg_tuple.1#115, %add.295, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.296 = f32[] add(%tanh.73, %constant.14)
+  %sub.73 = f32[] subtract(%constant.14, %tanh.73)
+  %mul.429 = f32[] multiply(%add.296, %sub.73)
+  %jit_dynamic_update_index_in_dim_.295 = f32[614]{0} call(%arg_tuple.1#116, %mul.429, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.74 = pred[] compare(%add.299, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.296 = pred[614]{0} call(%arg_tuple.1#117, %ge.74, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.297 = f32[614]{0} call(%arg_tuple.1#118, %exp.74, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.298 = f32[614]{0} call(%arg_tuple.1#119, %add.299, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.300 = f32[] add(%tanh.74, %constant.14)
+  %sub.74 = f32[] subtract(%constant.14, %tanh.74)
+  %mul.434 = f32[] multiply(%add.300, %sub.74)
+  %jit_dynamic_update_index_in_dim_.299 = f32[614]{0} call(%arg_tuple.1#120, %mul.434, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.75 = pred[] compare(%add.303, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.300 = pred[614]{0} call(%arg_tuple.1#121, %ge.75, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.301 = f32[614]{0} call(%arg_tuple.1#122, %exp.75, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.302 = f32[614]{0} call(%arg_tuple.1#123, %add.303, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.304 = f32[] add(%tanh.75, %constant.14)
+  %sub.75 = f32[] subtract(%constant.14, %tanh.75)
+  %mul.439 = f32[] multiply(%add.304, %sub.75)
+  %jit_dynamic_update_index_in_dim_.303 = f32[614]{0} call(%arg_tuple.1#124, %mul.439, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.76 = pred[] compare(%add.307, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.304 = pred[614]{0} call(%arg_tuple.1#125, %ge.76, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.305 = f32[614]{0} call(%arg_tuple.1#126, %exp.76, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.306 = f32[614]{0} call(%arg_tuple.1#127, %add.307, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.308 = f32[] add(%tanh.76, %constant.14)
+  %sub.76 = f32[] subtract(%constant.14, %tanh.76)
+  %mul.444 = f32[] multiply(%add.308, %sub.76)
+  %jit_dynamic_update_index_in_dim_.307 = f32[614]{0} call(%arg_tuple.1#128, %mul.444, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.77 = pred[] compare(%add.311, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.308 = pred[614]{0} call(%arg_tuple.1#129, %ge.77, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.309 = f32[614]{0} call(%arg_tuple.1#130, %exp.77, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.310 = f32[614]{0} call(%arg_tuple.1#131, %add.311, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.312 = f32[] add(%tanh.77, %constant.14)
+  %sub.77 = f32[] subtract(%constant.14, %tanh.77)
+  %mul.449 = f32[] multiply(%add.312, %sub.77)
+  %jit_dynamic_update_index_in_dim_.311 = f32[614]{0} call(%arg_tuple.1#132, %mul.449, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.78 = pred[] compare(%add.315, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.312 = pred[614]{0} call(%arg_tuple.1#133, %ge.78, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.313 = f32[614]{0} call(%arg_tuple.1#134, %exp.78, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.314 = f32[614]{0} call(%arg_tuple.1#135, %add.315, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.316 = f32[] add(%tanh.78, %constant.14)
+  %sub.78 = f32[] subtract(%constant.14, %tanh.78)
+  %mul.454 = f32[] multiply(%add.316, %sub.78)
+  %jit_dynamic_update_index_in_dim_.315 = f32[614]{0} call(%arg_tuple.1#136, %mul.454, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.79 = pred[] compare(%add.319, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.316 = pred[614]{0} call(%arg_tuple.1#137, %ge.79, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.317 = f32[614]{0} call(%arg_tuple.1#138, %exp.79, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.318 = f32[614]{0} call(%arg_tuple.1#139, %add.319, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.320 = f32[] add(%tanh.79, %constant.14)
+  %sub.79 = f32[] subtract(%constant.14, %tanh.79)
+  %mul.459 = f32[] multiply(%add.320, %sub.79)
+  %jit_dynamic_update_index_in_dim_.319 = f32[614]{0} call(%arg_tuple.1#140, %mul.459, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.80 = pred[] compare(%add.323, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.320 = pred[614]{0} call(%arg_tuple.1#141, %ge.80, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.321 = f32[614]{0} call(%arg_tuple.1#142, %exp.80, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.322 = f32[614]{0} call(%arg_tuple.1#143, %add.323, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.324 = f32[] add(%tanh.80, %constant.14)
+  %sub.80 = f32[] subtract(%constant.14, %tanh.80)
+  %mul.464 = f32[] multiply(%add.324, %sub.80)
+  %jit_dynamic_update_index_in_dim_.323 = f32[614]{0} call(%arg_tuple.1#144, %mul.464, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.81 = pred[] compare(%add.327, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.324 = pred[614]{0} call(%arg_tuple.1#145, %ge.81, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.325 = f32[614]{0} call(%arg_tuple.1#146, %exp.81, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.326 = f32[614]{0} call(%arg_tuple.1#147, %add.327, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.328 = f32[] add(%tanh.81, %constant.14)
+  %sub.81 = f32[] subtract(%constant.14, %tanh.81)
+  %mul.469 = f32[] multiply(%add.328, %sub.81)
+  %jit_dynamic_update_index_in_dim_.327 = f32[614]{0} call(%arg_tuple.1#148, %mul.469, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.82 = pred[] compare(%add.331, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.328 = pred[614]{0} call(%arg_tuple.1#149, %ge.82, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.329 = f32[614]{0} call(%arg_tuple.1#150, %exp.82, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.330 = f32[614]{0} call(%arg_tuple.1#151, %add.331, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.332 = f32[] add(%tanh.82, %constant.14)
+  %sub.82 = f32[] subtract(%constant.14, %tanh.82)
+  %mul.474 = f32[] multiply(%add.332, %sub.82)
+  %jit_dynamic_update_index_in_dim_.331 = f32[614]{0} call(%arg_tuple.1#152, %mul.474, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.83 = pred[] compare(%add.335, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.332 = pred[614]{0} call(%arg_tuple.1#153, %ge.83, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.333 = f32[614]{0} call(%arg_tuple.1#154, %exp.83, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.334 = f32[614]{0} call(%arg_tuple.1#155, %add.335, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.336 = f32[] add(%tanh.83, %constant.14)
+  %sub.83 = f32[] subtract(%constant.14, %tanh.83)
+  %mul.479 = f32[] multiply(%add.336, %sub.83)
+  %jit_dynamic_update_index_in_dim_.335 = f32[614]{0} call(%arg_tuple.1#156, %mul.479, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.84 = pred[] compare(%add.339, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.336 = pred[614]{0} call(%arg_tuple.1#157, %ge.84, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.337 = f32[614]{0} call(%arg_tuple.1#158, %exp.84, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.338 = f32[614]{0} call(%arg_tuple.1#159, %add.339, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.340 = f32[] add(%tanh.84, %constant.14)
+  %sub.84 = f32[] subtract(%constant.14, %tanh.84)
+  %mul.484 = f32[] multiply(%add.340, %sub.84)
+  %jit_dynamic_update_index_in_dim_.339 = f32[614]{0} call(%arg_tuple.1#160, %mul.484, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.85 = pred[] compare(%add.343, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.340 = pred[614]{0} call(%arg_tuple.1#161, %ge.85, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.341 = f32[614]{0} call(%arg_tuple.1#162, %exp.85, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.342 = f32[614]{0} call(%arg_tuple.1#163, %add.343, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.344 = f32[] add(%tanh.85, %constant.14)
+  %sub.85 = f32[] subtract(%constant.14, %tanh.85)
+  %mul.489 = f32[] multiply(%add.344, %sub.85)
+  %jit_dynamic_update_index_in_dim_.343 = f32[614]{0} call(%arg_tuple.1#164, %mul.489, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.86 = pred[] compare(%add.347, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.344 = pred[614]{0} call(%arg_tuple.1#165, %ge.86, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.345 = f32[614]{0} call(%arg_tuple.1#166, %exp.86, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.346 = f32[614]{0} call(%arg_tuple.1#167, %add.347, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.348 = f32[] add(%tanh.86, %constant.14)
+  %sub.86 = f32[] subtract(%constant.14, %tanh.86)
+  %mul.494 = f32[] multiply(%add.348, %sub.86)
+  %jit_dynamic_update_index_in_dim_.347 = f32[614]{0} call(%arg_tuple.1#168, %mul.494, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.87 = pred[] compare(%add.351, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.348 = pred[614]{0} call(%arg_tuple.1#169, %ge.87, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.349 = f32[614]{0} call(%arg_tuple.1#170, %exp.87, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.350 = f32[614]{0} call(%arg_tuple.1#171, %add.351, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.352 = f32[] add(%tanh.87, %constant.14)
+  %sub.87 = f32[] subtract(%constant.14, %tanh.87)
+  %mul.499 = f32[] multiply(%add.352, %sub.87)
+  %jit_dynamic_update_index_in_dim_.351 = f32[614]{0} call(%arg_tuple.1#172, %mul.499, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.88 = pred[] compare(%add.355, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.352 = pred[614]{0} call(%arg_tuple.1#173, %ge.88, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.353 = f32[614]{0} call(%arg_tuple.1#174, %exp.88, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.354 = f32[614]{0} call(%arg_tuple.1#175, %add.355, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.356 = f32[] add(%tanh.88, %constant.14)
+  %sub.88 = f32[] subtract(%constant.14, %tanh.88)
+  %mul.504 = f32[] multiply(%add.356, %sub.88)
+  %jit_dynamic_update_index_in_dim_.355 = f32[614]{0} call(%arg_tuple.1#176, %mul.504, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.89 = pred[] compare(%add.359, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.356 = pred[614]{0} call(%arg_tuple.1#177, %ge.89, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.357 = f32[614]{0} call(%arg_tuple.1#178, %exp.89, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.358 = f32[614]{0} call(%arg_tuple.1#179, %add.359, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.360 = f32[] add(%tanh.89, %constant.14)
+  %sub.89 = f32[] subtract(%constant.14, %tanh.89)
+  %mul.509 = f32[] multiply(%add.360, %sub.89)
+  %jit_dynamic_update_index_in_dim_.359 = f32[614]{0} call(%arg_tuple.1#180, %mul.509, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.90 = pred[] compare(%add.363, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.360 = pred[614]{0} call(%arg_tuple.1#181, %ge.90, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.361 = f32[614]{0} call(%arg_tuple.1#182, %exp.90, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.362 = f32[614]{0} call(%arg_tuple.1#183, %add.363, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.364 = f32[] add(%tanh.90, %constant.14)
+  %sub.90 = f32[] subtract(%constant.14, %tanh.90)
+  %mul.514 = f32[] multiply(%add.364, %sub.90)
+  %jit_dynamic_update_index_in_dim_.363 = f32[614]{0} call(%arg_tuple.1#184, %mul.514, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.91 = pred[] compare(%add.367, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.364 = pred[614]{0} call(%arg_tuple.1#185, %ge.91, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.365 = f32[614]{0} call(%arg_tuple.1#186, %exp.91, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.366 = f32[614]{0} call(%arg_tuple.1#187, %add.367, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.368 = f32[] add(%tanh.91, %constant.14)
+  %sub.91 = f32[] subtract(%constant.14, %tanh.91)
+  %mul.519 = f32[] multiply(%add.368, %sub.91)
+  %jit_dynamic_update_index_in_dim_.367 = f32[614]{0} call(%arg_tuple.1#188, %mul.519, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.92 = pred[] compare(%add.371, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.368 = pred[614]{0} call(%arg_tuple.1#189, %ge.92, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.369 = f32[614]{0} call(%arg_tuple.1#190, %exp.92, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.370 = f32[614]{0} call(%arg_tuple.1#191, %add.371, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.372 = f32[] add(%tanh.92, %constant.14)
+  %sub.92 = f32[] subtract(%constant.14, %tanh.92)
+  %mul.524 = f32[] multiply(%add.372, %sub.92)
+  %jit_dynamic_update_index_in_dim_.371 = f32[614]{0} call(%arg_tuple.1#192, %mul.524, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.93 = pred[] compare(%add.375, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.372 = pred[614]{0} call(%arg_tuple.1#193, %ge.93, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.373 = f32[614]{0} call(%arg_tuple.1#194, %exp.93, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.374 = f32[614]{0} call(%arg_tuple.1#195, %add.375, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.376 = f32[] add(%tanh.93, %constant.14)
+  %sub.93 = f32[] subtract(%constant.14, %tanh.93)
+  %mul.529 = f32[] multiply(%add.376, %sub.93)
+  %jit_dynamic_update_index_in_dim_.375 = f32[614]{0} call(%arg_tuple.1#196, %mul.529, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.94 = pred[] compare(%add.379, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.376 = pred[614]{0} call(%arg_tuple.1#197, %ge.94, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.377 = f32[614]{0} call(%arg_tuple.1#198, %exp.94, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.378 = f32[614]{0} call(%arg_tuple.1#199, %add.379, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.380 = f32[] add(%tanh.94, %constant.14)
+  %sub.94 = f32[] subtract(%constant.14, %tanh.94)
+  %mul.534 = f32[] multiply(%add.380, %sub.94)
+  %jit_dynamic_update_index_in_dim_.379 = f32[614]{0} call(%arg_tuple.1#200, %mul.534, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.95 = pred[] compare(%add.383, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.380 = pred[614]{0} call(%arg_tuple.1#201, %ge.95, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.381 = f32[614]{0} call(%arg_tuple.1#202, %exp.95, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.382 = f32[614]{0} call(%arg_tuple.1#203, %add.383, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.384 = f32[] add(%tanh.95, %constant.14)
+  %sub.95 = f32[] subtract(%constant.14, %tanh.95)
+  %mul.539 = f32[] multiply(%add.384, %sub.95)
+  %jit_dynamic_update_index_in_dim_.383 = f32[614]{0} call(%arg_tuple.1#204, %mul.539, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.96 = pred[] compare(%add.387, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.384 = pred[614]{0} call(%arg_tuple.1#205, %ge.96, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.385 = f32[614]{0} call(%arg_tuple.1#206, %exp.96, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.386 = f32[614]{0} call(%arg_tuple.1#207, %add.387, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.388 = f32[] add(%tanh.96, %constant.14)
+  %sub.96 = f32[] subtract(%constant.14, %tanh.96)
+  %mul.544 = f32[] multiply(%add.388, %sub.96)
+  %jit_dynamic_update_index_in_dim_.387 = f32[614]{0} call(%arg_tuple.1#208, %mul.544, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.97 = pred[] compare(%add.391, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.388 = pred[614]{0} call(%arg_tuple.1#209, %ge.97, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.389 = f32[614]{0} call(%arg_tuple.1#210, %exp.97, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.390 = f32[614]{0} call(%arg_tuple.1#211, %add.391, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.392 = f32[] add(%tanh.97, %constant.14)
+  %sub.97 = f32[] subtract(%constant.14, %tanh.97)
+  %mul.549 = f32[] multiply(%add.392, %sub.97)
+  %jit_dynamic_update_index_in_dim_.391 = f32[614]{0} call(%arg_tuple.1#212, %mul.549, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.98 = pred[] compare(%add.395, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.392 = pred[614]{0} call(%arg_tuple.1#213, %ge.98, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.393 = f32[614]{0} call(%arg_tuple.1#214, %exp.98, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.394 = f32[614]{0} call(%arg_tuple.1#215, %add.395, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.396 = f32[] add(%tanh.98, %constant.14)
+  %sub.98 = f32[] subtract(%constant.14, %tanh.98)
+  %mul.554 = f32[] multiply(%add.396, %sub.98)
+  %jit_dynamic_update_index_in_dim_.395 = f32[614]{0} call(%arg_tuple.1#216, %mul.554, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.99 = pred[] compare(%add.399, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.396 = pred[614]{0} call(%arg_tuple.1#217, %ge.99, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.397 = f32[614]{0} call(%arg_tuple.1#218, %exp.99, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.398 = f32[614]{0} call(%arg_tuple.1#219, %add.399, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.400 = f32[] add(%tanh.99, %constant.14)
+  %sub.99 = f32[] subtract(%constant.14, %tanh.99)
+  %mul.559 = f32[] multiply(%add.400, %sub.99)
+  %jit_dynamic_update_index_in_dim_.399 = f32[614]{0} call(%arg_tuple.1#220, %mul.559, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.100 = pred[] compare(%add.403, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.400 = pred[614]{0} call(%arg_tuple.1#221, %ge.100, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.401 = f32[614]{0} call(%arg_tuple.1#222, %exp.100, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.402 = f32[614]{0} call(%arg_tuple.1#223, %add.403, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.404 = f32[] add(%tanh.100, %constant.14)
+  %sub.100 = f32[] subtract(%constant.14, %tanh.100)
+  %mul.564 = f32[] multiply(%add.404, %sub.100)
+  %jit_dynamic_update_index_in_dim_.403 = f32[614]{0} call(%arg_tuple.1#224, %mul.564, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.101 = pred[] compare(%add.407, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.404 = pred[614]{0} call(%arg_tuple.1#225, %ge.101, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.405 = f32[614]{0} call(%arg_tuple.1#226, %exp.101, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.406 = f32[614]{0} call(%arg_tuple.1#227, %add.407, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.408 = f32[] add(%tanh.101, %constant.14)
+  %sub.101 = f32[] subtract(%constant.14, %tanh.101)
+  %mul.569 = f32[] multiply(%add.408, %sub.101)
+  %jit_dynamic_update_index_in_dim_.407 = f32[614]{0} call(%arg_tuple.1#228, %mul.569, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.102 = pred[] compare(%add.411, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.408 = pred[614]{0} call(%arg_tuple.1#229, %ge.102, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.409 = f32[614]{0} call(%arg_tuple.1#230, %exp.102, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.410 = f32[614]{0} call(%arg_tuple.1#231, %add.411, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.412 = f32[] add(%tanh.102, %constant.14)
+  %sub.102 = f32[] subtract(%constant.14, %tanh.102)
+  %mul.574 = f32[] multiply(%add.412, %sub.102)
+  %jit_dynamic_update_index_in_dim_.411 = f32[614]{0} call(%arg_tuple.1#232, %mul.574, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.103 = pred[] compare(%add.415, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.412 = pred[614]{0} call(%arg_tuple.1#233, %ge.103, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.413 = f32[614]{0} call(%arg_tuple.1#234, %exp.103, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.414 = f32[614]{0} call(%arg_tuple.1#235, %add.415, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.416 = f32[] add(%tanh.103, %constant.14)
+  %sub.103 = f32[] subtract(%constant.14, %tanh.103)
+  %mul.579 = f32[] multiply(%add.416, %sub.103)
+  %jit_dynamic_update_index_in_dim_.415 = f32[614]{0} call(%arg_tuple.1#236, %mul.579, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.104 = pred[] compare(%add.419, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.416 = pred[614]{0} call(%arg_tuple.1#237, %ge.104, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.417 = f32[614]{0} call(%arg_tuple.1#238, %exp.104, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.418 = f32[614]{0} call(%arg_tuple.1#239, %add.419, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.420 = f32[] add(%tanh.104, %constant.14)
+  %sub.104 = f32[] subtract(%constant.14, %tanh.104)
+  %mul.584 = f32[] multiply(%add.420, %sub.104)
+  %jit_dynamic_update_index_in_dim_.419 = f32[614]{0} call(%arg_tuple.1#240, %mul.584, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.105 = pred[] compare(%add.423, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.420 = pred[614]{0} call(%arg_tuple.1#241, %ge.105, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.421 = f32[614]{0} call(%arg_tuple.1#242, %exp.105, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.422 = f32[614]{0} call(%arg_tuple.1#243, %add.423, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.424 = f32[] add(%tanh.105, %constant.14)
+  %sub.105 = f32[] subtract(%constant.14, %tanh.105)
+  %mul.589 = f32[] multiply(%add.424, %sub.105)
+  %jit_dynamic_update_index_in_dim_.423 = f32[614]{0} call(%arg_tuple.1#244, %mul.589, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.106 = pred[] compare(%add.427, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.424 = pred[614]{0} call(%arg_tuple.1#245, %ge.106, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.425 = f32[614]{0} call(%arg_tuple.1#246, %exp.106, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.426 = f32[614]{0} call(%arg_tuple.1#247, %add.427, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.428 = f32[] add(%tanh.106, %constant.14)
+  %sub.106 = f32[] subtract(%constant.14, %tanh.106)
+  %mul.594 = f32[] multiply(%add.428, %sub.106)
+  %jit_dynamic_update_index_in_dim_.427 = f32[614]{0} call(%arg_tuple.1#248, %mul.594, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.107 = pred[] compare(%add.431, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.428 = pred[614]{0} call(%arg_tuple.1#249, %ge.107, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.429 = f32[614]{0} call(%arg_tuple.1#250, %exp.107, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.430 = f32[614]{0} call(%arg_tuple.1#251, %add.431, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.432 = f32[] add(%tanh.107, %constant.14)
+  %sub.107 = f32[] subtract(%constant.14, %tanh.107)
+  %mul.599 = f32[] multiply(%add.432, %sub.107)
+  %jit_dynamic_update_index_in_dim_.431 = f32[614]{0} call(%arg_tuple.1#252, %mul.599, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.108 = pred[] compare(%add.435, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.432 = pred[614]{0} call(%arg_tuple.1#253, %ge.108, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.433 = f32[614]{0} call(%arg_tuple.1#254, %exp.108, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.434 = f32[614]{0} call(%arg_tuple.1#255, %add.435, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.436 = f32[] add(%tanh.108, %constant.14)
+  %sub.108 = f32[] subtract(%constant.14, %tanh.108)
+  %mul.604 = f32[] multiply(%add.436, %sub.108)
+  %jit_dynamic_update_index_in_dim_.435 = f32[614]{0} call(%arg_tuple.1#256, %mul.604, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.109 = pred[] compare(%add.439, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.436 = pred[614]{0} call(%arg_tuple.1#257, %ge.109, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.437 = f32[614]{0} call(%arg_tuple.1#258, %exp.109, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.438 = f32[614]{0} call(%arg_tuple.1#259, %add.439, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.440 = f32[] add(%tanh.109, %constant.14)
+  %sub.109 = f32[] subtract(%constant.14, %tanh.109)
+  %mul.609 = f32[] multiply(%add.440, %sub.109)
+  %jit_dynamic_update_index_in_dim_.439 = f32[614]{0} call(%arg_tuple.1#260, %mul.609, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.110 = pred[] compare(%add.443, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.440 = pred[614]{0} call(%arg_tuple.1#261, %ge.110, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.441 = f32[614]{0} call(%arg_tuple.1#262, %exp.110, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.442 = f32[614]{0} call(%arg_tuple.1#263, %add.443, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.444 = f32[] add(%tanh.110, %constant.14)
+  %sub.110 = f32[] subtract(%constant.14, %tanh.110)
+  %mul.614 = f32[] multiply(%add.444, %sub.110)
+  %jit_dynamic_update_index_in_dim_.443 = f32[614]{0} call(%arg_tuple.1#264, %mul.614, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.111 = pred[] compare(%add.447, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.444 = pred[614]{0} call(%arg_tuple.1#265, %ge.111, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.445 = f32[614]{0} call(%arg_tuple.1#266, %exp.111, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.446 = f32[614]{0} call(%arg_tuple.1#267, %add.447, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.448 = f32[] add(%tanh.111, %constant.14)
+  %sub.111 = f32[] subtract(%constant.14, %tanh.111)
+  %mul.619 = f32[] multiply(%add.448, %sub.111)
+  %jit_dynamic_update_index_in_dim_.447 = f32[614]{0} call(%arg_tuple.1#268, %mul.619, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.112 = pred[] compare(%add.451, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.448 = pred[614]{0} call(%arg_tuple.1#269, %ge.112, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.449 = f32[614]{0} call(%arg_tuple.1#270, %exp.112, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.450 = f32[614]{0} call(%arg_tuple.1#271, %add.451, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.452 = f32[] add(%tanh.112, %constant.14)
+  %sub.112 = f32[] subtract(%constant.14, %tanh.112)
+  %mul.624 = f32[] multiply(%add.452, %sub.112)
+  %jit_dynamic_update_index_in_dim_.451 = f32[614]{0} call(%arg_tuple.1#272, %mul.624, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.113 = pred[] compare(%add.455, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.452 = pred[614]{0} call(%arg_tuple.1#273, %ge.113, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.453 = f32[614]{0} call(%arg_tuple.1#274, %exp.113, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.454 = f32[614]{0} call(%arg_tuple.1#275, %add.455, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.456 = f32[] add(%tanh.113, %constant.14)
+  %sub.113 = f32[] subtract(%constant.14, %tanh.113)
+  %mul.629 = f32[] multiply(%add.456, %sub.113)
+  %jit_dynamic_update_index_in_dim_.455 = f32[614]{0} call(%arg_tuple.1#276, %mul.629, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.114 = pred[] compare(%add.459, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.456 = pred[614]{0} call(%arg_tuple.1#277, %ge.114, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.457 = f32[614]{0} call(%arg_tuple.1#278, %exp.114, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.458 = f32[614]{0} call(%arg_tuple.1#279, %add.459, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.460 = f32[] add(%tanh.114, %constant.14)
+  %sub.114 = f32[] subtract(%constant.14, %tanh.114)
+  %mul.634 = f32[] multiply(%add.460, %sub.114)
+  %jit_dynamic_update_index_in_dim_.459 = f32[614]{0} call(%arg_tuple.1#280, %mul.634, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.115 = pred[] compare(%add.463, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.460 = pred[614]{0} call(%arg_tuple.1#281, %ge.115, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.461 = f32[614]{0} call(%arg_tuple.1#282, %exp.115, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.462 = f32[614]{0} call(%arg_tuple.1#283, %add.463, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.464 = f32[] add(%tanh.115, %constant.14)
+  %sub.115 = f32[] subtract(%constant.14, %tanh.115)
+  %mul.639 = f32[] multiply(%add.464, %sub.115)
+  %jit_dynamic_update_index_in_dim_.463 = f32[614]{0} call(%arg_tuple.1#284, %mul.639, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.116 = pred[] compare(%add.467, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.464 = pred[614]{0} call(%arg_tuple.1#285, %ge.116, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.465 = f32[614]{0} call(%arg_tuple.1#286, %exp.116, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.466 = f32[614]{0} call(%arg_tuple.1#287, %add.467, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.468 = f32[] add(%tanh.116, %constant.14)
+  %sub.116 = f32[] subtract(%constant.14, %tanh.116)
+  %mul.644 = f32[] multiply(%add.468, %sub.116)
+  %jit_dynamic_update_index_in_dim_.467 = f32[614]{0} call(%arg_tuple.1#288, %mul.644, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.117 = pred[] compare(%add.471, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.468 = pred[614]{0} call(%arg_tuple.1#289, %ge.117, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.469 = f32[614]{0} call(%arg_tuple.1#290, %exp.117, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.470 = f32[614]{0} call(%arg_tuple.1#291, %add.471, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.472 = f32[] add(%tanh.117, %constant.14)
+  %sub.117 = f32[] subtract(%constant.14, %tanh.117)
+  %mul.649 = f32[] multiply(%add.472, %sub.117)
+  %jit_dynamic_update_index_in_dim_.471 = f32[614]{0} call(%arg_tuple.1#292, %mul.649, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.118 = pred[] compare(%add.475, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.472 = pred[614]{0} call(%arg_tuple.1#293, %ge.118, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.473 = f32[614]{0} call(%arg_tuple.1#294, %exp.118, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.474 = f32[614]{0} call(%arg_tuple.1#295, %add.475, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.476 = f32[] add(%tanh.118, %constant.14)
+  %sub.118 = f32[] subtract(%constant.14, %tanh.118)
+  %mul.654 = f32[] multiply(%add.476, %sub.118)
+  %jit_dynamic_update_index_in_dim_.475 = f32[614]{0} call(%arg_tuple.1#296, %mul.654, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %ge.119 = pred[] compare(%add.479, %constant.15), direction=GE
+  %jit_dynamic_update_index_in_dim_.476 = pred[614]{0} call(%arg_tuple.1#297, %ge.119, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.477 = f32[614]{0} call(%arg_tuple.1#298, %exp.119, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.478 = f32[614]{0} call(%arg_tuple.1#299, %add.479, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %add.480 = f32[] add(%tanh.119, %constant.14)
+  %sub.119 = f32[] subtract(%constant.14, %tanh.119)
+  %mul.659 = f32[] multiply(%add.480, %sub.119)
+  %jit_dynamic_update_index_in_dim_.479 = f32[614]{0} call(%arg_tuple.1#300, %mul.659, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  ROOT %tuple.1 = (s32[], f32[], f32[], f32[], f32[], /*index=5*/f32[], f32[], f32[], f32[], f32[], /*index=10*/f32[], f32[], f32[], f32[], f32[], /*index=15*/f32[], f32[], f32[], f32[], f32[], /*index=20*/f32[], f32[], f32[], f32[], f32[], /*index=25*/f32[], f32[], f32[], f32[], f32[], /*index=30*/f32[], f32[], f32[], f32[], f32[], /*index=35*/f32[], f32[], f32[], f32[], f32[], /*index=40*/f32[], f32[], f32[], f32[], f32[], /*index=45*/f32[], f32[], f32[], f32[], f32[], /*index=50*/f32[], f32[], f32[], f32[], f32[], /*index=55*/f32[], f32[], f32[], f32[], f32[], /*index=60*/f32[], pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=65*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=70*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=75*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=80*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=85*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=90*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=95*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=100*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=105*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=110*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=115*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=120*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=125*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=130*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=135*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=140*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=145*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=150*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=155*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=160*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=165*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=170*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=175*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=180*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=185*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=190*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=195*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=200*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=205*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=210*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=215*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=220*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=225*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=230*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=235*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=240*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=245*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=250*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=255*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=260*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=265*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=270*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=275*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=280*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=285*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=290*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=295*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=300*/f32[614]{0}, f32[], f32[], f32[], f32[], /*index=305*/f32[], f32[], f32[], f32[], f32[], /*index=310*/f32[], f32[], f32[], f32[], f32[], /*index=315*/f32[], f32[], f32[], f32[], f32[], /*index=320*/f32[], f32[], f32[], f32[], f32[], /*index=325*/f32[], f32[], f32[], f32[], f32[], /*index=330*/f32[], f32[], f32[], f32[], f32[], /*index=335*/f32[], f32[], f32[], f32[], f32[], /*index=340*/f32[], f32[], f32[], f32[], f32[], /*index=345*/f32[], f32[], f32[], f32[], f32[], /*index=350*/f32[], f32[], f32[], f32[], f32[], /*index=355*/f32[], f32[], f32[], f32[], f32[], /*index=360*/f32[]) tuple(%add.481, %tanh.60, %tanh.61, %tanh.62, %tanh.63, /*index=5*/%tanh.64, %tanh.65, %tanh.66, %tanh.67, %tanh.68, /*index=10*/%tanh.69, %tanh.70, %tanh.71, %tanh.72, %tanh.73, /*index=15*/%tanh.74, %tanh.75, %tanh.76, %tanh.77, %tanh.78, /*index=20*/%tanh.79, %tanh.80, %tanh.81, %tanh.82, %tanh.83, /*index=25*/%tanh.84, %tanh.85, %tanh.86, %tanh.87, %tanh.88, /*index=30*/%tanh.89, %tanh.90, %tanh.91, %tanh.92, %tanh.93, /*index=35*/%tanh.94, %tanh.95, %tanh.96, %tanh.97, %tanh.98, /*index=40*/%tanh.99, %tanh.100, %tanh.101, %tanh.102, %tanh.103, /*index=45*/%tanh.104, %tanh.105, %tanh.106, %tanh.107, %tanh.108, /*index=50*/%tanh.109, %tanh.110, %tanh.111, %tanh.112, %tanh.113, /*index=55*/%tanh.114, %tanh.115, %tanh.116, %tanh.117, %tanh.118, /*index=60*/%tanh.119, %jit_dynamic_update_index_in_dim_.240, %jit_dynamic_update_index_in_dim_.241, %jit_dynamic_update_index_in_dim_.242, %jit_dynamic_update_index_in_dim_.243, /*index=65*/%jit_dynamic_update_index_in_dim_.244, %jit_dynamic_update_index_in_dim_.245, %jit_dynamic_update_index_in_dim_.246, %jit_dynamic_update_index_in_dim_.247, %jit_dynamic_update_index_in_dim_.248, /*index=70*/%jit_dynamic_update_index_in_dim_.249, %jit_dynamic_update_index_in_dim_.250, %jit_dynamic_update_index_in_dim_.251, %jit_dynamic_update_index_in_dim_.252, %jit_dynamic_update_index_in_dim_.253, /*index=75*/%jit_dynamic_update_index_in_dim_.254, %jit_dynamic_update_index_in_dim_.255, %jit_dynamic_update_index_in_dim_.256, %jit_dynamic_update_index_in_dim_.257, %jit_dynamic_update_index_in_dim_.258, /*index=80*/%jit_dynamic_update_index_in_dim_.259, %jit_dynamic_update_index_in_dim_.260, %jit_dynamic_update_index_in_dim_.261, %jit_dynamic_update_index_in_dim_.262, %jit_dynamic_update_index_in_dim_.263, /*index=85*/%jit_dynamic_update_index_in_dim_.264, %jit_dynamic_update_index_in_dim_.265, %jit_dynamic_update_index_in_dim_.266, %jit_dynamic_update_index_in_dim_.267, %jit_dynamic_update_index_in_dim_.268, /*index=90*/%jit_dynamic_update_index_in_dim_.269, %jit_dynamic_update_index_in_dim_.270, %jit_dynamic_update_index_in_dim_.271, %jit_dynamic_update_index_in_dim_.272, %jit_dynamic_update_index_in_dim_.273, /*index=95*/%jit_dynamic_update_index_in_dim_.274, %jit_dynamic_update_index_in_dim_.275, %jit_dynamic_update_index_in_dim_.276, %jit_dynamic_update_index_in_dim_.277, %jit_dynamic_update_index_in_dim_.278, /*index=100*/%jit_dynamic_update_index_in_dim_.279, %jit_dynamic_update_index_in_dim_.280, %jit_dynamic_update_index_in_dim_.281, %jit_dynamic_update_index_in_dim_.282, %jit_dynamic_update_index_in_dim_.283, /*index=105*/%jit_dynamic_update_index_in_dim_.284, %jit_dynamic_update_index_in_dim_.285, %jit_dynamic_update_index_in_dim_.286, %jit_dynamic_update_index_in_dim_.287, %jit_dynamic_update_index_in_dim_.288, /*index=110*/%jit_dynamic_update_index_in_dim_.289, %jit_dynamic_update_index_in_dim_.290, %jit_dynamic_update_index_in_dim_.291, %jit_dynamic_update_index_in_dim_.292, %jit_dynamic_update_index_in_dim_.293, /*index=115*/%jit_dynamic_update_index_in_dim_.294, %jit_dynamic_update_index_in_dim_.295, %jit_dynamic_update_index_in_dim_.296, %jit_dynamic_update_index_in_dim_.297, %jit_dynamic_update_index_in_dim_.298, /*index=120*/%jit_dynamic_update_index_in_dim_.299, %jit_dynamic_update_index_in_dim_.300, %jit_dynamic_update_index_in_dim_.301, %jit_dynamic_update_index_in_dim_.302, %jit_dynamic_update_index_in_dim_.303, /*index=125*/%jit_dynamic_update_index_in_dim_.304, %jit_dynamic_update_index_in_dim_.305, %jit_dynamic_update_index_in_dim_.306, %jit_dynamic_update_index_in_dim_.307, %jit_dynamic_update_index_in_dim_.308, /*index=130*/%jit_dynamic_update_index_in_dim_.309, %jit_dynamic_update_index_in_dim_.310, %jit_dynamic_update_index_in_dim_.311, %jit_dynamic_update_index_in_dim_.312, %jit_dynamic_update_index_in_dim_.313, /*index=135*/%jit_dynamic_update_index_in_dim_.314, %jit_dynamic_update_index_in_dim_.315, %jit_dynamic_update_index_in_dim_.316, %jit_dynamic_update_index_in_dim_.317, %jit_dynamic_update_index_in_dim_.318, /*index=140*/%jit_dynamic_update_index_in_dim_.319, %jit_dynamic_update_index_in_dim_.320, %jit_dynamic_update_index_in_dim_.321, %jit_dynamic_update_index_in_dim_.322, %jit_dynamic_update_index_in_dim_.323, /*index=145*/%jit_dynamic_update_index_in_dim_.324, %jit_dynamic_update_index_in_dim_.325, %jit_dynamic_update_index_in_dim_.326, %jit_dynamic_update_index_in_dim_.327, %jit_dynamic_update_index_in_dim_.328, /*index=150*/%jit_dynamic_update_index_in_dim_.329, %jit_dynamic_update_index_in_dim_.330, %jit_dynamic_update_index_in_dim_.331, %jit_dynamic_update_index_in_dim_.332, %jit_dynamic_update_index_in_dim_.333, /*index=155*/%jit_dynamic_update_index_in_dim_.334, %jit_dynamic_update_index_in_dim_.335, %jit_dynamic_update_index_in_dim_.336, %jit_dynamic_update_index_in_dim_.337, %jit_dynamic_update_index_in_dim_.338, /*index=160*/%jit_dynamic_update_index_in_dim_.339, %jit_dynamic_update_index_in_dim_.340, %jit_dynamic_update_index_in_dim_.341, %jit_dynamic_update_index_in_dim_.342, %jit_dynamic_update_index_in_dim_.343, /*index=165*/%jit_dynamic_update_index_in_dim_.344, %jit_dynamic_update_index_in_dim_.345, %jit_dynamic_update_index_in_dim_.346, %jit_dynamic_update_index_in_dim_.347, %jit_dynamic_update_index_in_dim_.348, /*index=170*/%jit_dynamic_update_index_in_dim_.349, %jit_dynamic_update_index_in_dim_.350, %jit_dynamic_update_index_in_dim_.351, %jit_dynamic_update_index_in_dim_.352, %jit_dynamic_update_index_in_dim_.353, /*index=175*/%jit_dynamic_update_index_in_dim_.354, %jit_dynamic_update_index_in_dim_.355, %jit_dynamic_update_index_in_dim_.356, %jit_dynamic_update_index_in_dim_.357, %jit_dynamic_update_index_in_dim_.358, /*index=180*/%jit_dynamic_update_index_in_dim_.359, %jit_dynamic_update_index_in_dim_.360, %jit_dynamic_update_index_in_dim_.361, %jit_dynamic_update_index_in_dim_.362, %jit_dynamic_update_index_in_dim_.363, /*index=185*/%jit_dynamic_update_index_in_dim_.364, %jit_dynamic_update_index_in_dim_.365, %jit_dynamic_update_index_in_dim_.366, %jit_dynamic_update_index_in_dim_.367, %jit_dynamic_update_index_in_dim_.368, /*index=190*/%jit_dynamic_update_index_in_dim_.369, %jit_dynamic_update_index_in_dim_.370, %jit_dynamic_update_index_in_dim_.371, %jit_dynamic_update_index_in_dim_.372, %jit_dynamic_update_index_in_dim_.373, /*index=195*/%jit_dynamic_update_index_in_dim_.374, %jit_dynamic_update_index_in_dim_.375, %jit_dynamic_update_index_in_dim_.376, %jit_dynamic_update_index_in_dim_.377, %jit_dynamic_update_index_in_dim_.378, /*index=200*/%jit_dynamic_update_index_in_dim_.379, %jit_dynamic_update_index_in_dim_.380, %jit_dynamic_update_index_in_dim_.381, %jit_dynamic_update_index_in_dim_.382, %jit_dynamic_update_index_in_dim_.383, /*index=205*/%jit_dynamic_update_index_in_dim_.384, %jit_dynamic_update_index_in_dim_.385, %jit_dynamic_update_index_in_dim_.386, %jit_dynamic_update_index_in_dim_.387, %jit_dynamic_update_index_in_dim_.388, /*index=210*/%jit_dynamic_update_index_in_dim_.389, %jit_dynamic_update_index_in_dim_.390, %jit_dynamic_update_index_in_dim_.391, %jit_dynamic_update_index_in_dim_.392, %jit_dynamic_update_index_in_dim_.393, /*index=215*/%jit_dynamic_update_index_in_dim_.394, %jit_dynamic_update_index_in_dim_.395, %jit_dynamic_update_index_in_dim_.396, %jit_dynamic_update_index_in_dim_.397, %jit_dynamic_update_index_in_dim_.398, /*index=220*/%jit_dynamic_update_index_in_dim_.399, %jit_dynamic_update_index_in_dim_.400, %jit_dynamic_update_index_in_dim_.401, %jit_dynamic_update_index_in_dim_.402, %jit_dynamic_update_index_in_dim_.403, /*index=225*/%jit_dynamic_update_index_in_dim_.404, %jit_dynamic_update_index_in_dim_.405, %jit_dynamic_update_index_in_dim_.406, %jit_dynamic_update_index_in_dim_.407, %jit_dynamic_update_index_in_dim_.408, /*index=230*/%jit_dynamic_update_index_in_dim_.409, %jit_dynamic_update_index_in_dim_.410, %jit_dynamic_update_index_in_dim_.411, %jit_dynamic_update_index_in_dim_.412, %jit_dynamic_update_index_in_dim_.413, /*index=235*/%jit_dynamic_update_index_in_dim_.414, %jit_dynamic_update_index_in_dim_.415, %jit_dynamic_update_index_in_dim_.416, %jit_dynamic_update_index_in_dim_.417, %jit_dynamic_update_index_in_dim_.418, /*index=240*/%jit_dynamic_update_index_in_dim_.419, %jit_dynamic_update_index_in_dim_.420, %jit_dynamic_update_index_in_dim_.421, %jit_dynamic_update_index_in_dim_.422, %jit_dynamic_update_index_in_dim_.423, /*index=245*/%jit_dynamic_update_index_in_dim_.424, %jit_dynamic_update_index_in_dim_.425, %jit_dynamic_update_index_in_dim_.426, %jit_dynamic_update_index_in_dim_.427, %jit_dynamic_update_index_in_dim_.428, /*index=250*/%jit_dynamic_update_index_in_dim_.429, %jit_dynamic_update_index_in_dim_.430, %jit_dynamic_update_index_in_dim_.431, %jit_dynamic_update_index_in_dim_.432, %jit_dynamic_update_index_in_dim_.433, /*index=255*/%jit_dynamic_update_index_in_dim_.434, %jit_dynamic_update_index_in_dim_.435, %jit_dynamic_update_index_in_dim_.436, %jit_dynamic_update_index_in_dim_.437, %jit_dynamic_update_index_in_dim_.438, /*index=260*/%jit_dynamic_update_index_in_dim_.439, %jit_dynamic_update_index_in_dim_.440, %jit_dynamic_update_index_in_dim_.441, %jit_dynamic_update_index_in_dim_.442, %jit_dynamic_update_index_in_dim_.443, /*index=265*/%jit_dynamic_update_index_in_dim_.444, %jit_dynamic_update_index_in_dim_.445, %jit_dynamic_update_index_in_dim_.446, %jit_dynamic_update_index_in_dim_.447, %jit_dynamic_update_index_in_dim_.448, /*index=270*/%jit_dynamic_update_index_in_dim_.449, %jit_dynamic_update_index_in_dim_.450, %jit_dynamic_update_index_in_dim_.451, %jit_dynamic_update_index_in_dim_.452, %jit_dynamic_update_index_in_dim_.453, /*index=275*/%jit_dynamic_update_index_in_dim_.454, %jit_dynamic_update_index_in_dim_.455, %jit_dynamic_update_index_in_dim_.456, %jit_dynamic_update_index_in_dim_.457, %jit_dynamic_update_index_in_dim_.458, /*index=280*/%jit_dynamic_update_index_in_dim_.459, %jit_dynamic_update_index_in_dim_.460, %jit_dynamic_update_index_in_dim_.461, %jit_dynamic_update_index_in_dim_.462, %jit_dynamic_update_index_in_dim_.463, /*index=285*/%jit_dynamic_update_index_in_dim_.464, %jit_dynamic_update_index_in_dim_.465, %jit_dynamic_update_index_in_dim_.466, %jit_dynamic_update_index_in_dim_.467, %jit_dynamic_update_index_in_dim_.468, /*index=290*/%jit_dynamic_update_index_in_dim_.469, %jit_dynamic_update_index_in_dim_.470, %jit_dynamic_update_index_in_dim_.471, %jit_dynamic_update_index_in_dim_.472, %jit_dynamic_update_index_in_dim_.473, /*index=295*/%jit_dynamic_update_index_in_dim_.474, %jit_dynamic_update_index_in_dim_.475, %jit_dynamic_update_index_in_dim_.476, %jit_dynamic_update_index_in_dim_.477, %jit_dynamic_update_index_in_dim_.478, /*index=300*/%jit_dynamic_update_index_in_dim_.479, %arg_tuple.1#301, %arg_tuple.1#302, %arg_tuple.1#303, %arg_tuple.1#304, /*index=305*/%arg_tuple.1#305, %arg_tuple.1#306, %arg_tuple.1#307, %arg_tuple.1#308, %arg_tuple.1#309, /*index=310*/%arg_tuple.1#310, %arg_tuple.1#311, %arg_tuple.1#312, %arg_tuple.1#313, %arg_tuple.1#314, /*index=315*/%arg_tuple.1#315, %arg_tuple.1#316, %arg_tuple.1#317, %arg_tuple.1#318, %arg_tuple.1#319, /*index=320*/%arg_tuple.1#320, %arg_tuple.1#321, %arg_tuple.1#322, %arg_tuple.1#323, %arg_tuple.1#324, /*index=325*/%arg_tuple.1#325, %arg_tuple.1#326, %arg_tuple.1#327, %arg_tuple.1#328, %arg_tuple.1#329, /*index=330*/%arg_tuple.1#330, %arg_tuple.1#331, %arg_tuple.1#332, %arg_tuple.1#333, %arg_tuple.1#334, /*index=335*/%arg_tuple.1#335, %arg_tuple.1#336, %arg_tuple.1#337, %arg_tuple.1#338, %arg_tuple.1#339, /*index=340*/%arg_tuple.1#340, %arg_tuple.1#341, %arg_tuple.1#342, %arg_tuple.1#343, %arg_tuple.1#344, /*index=345*/%arg_tuple.1#345, %arg_tuple.1#346, %arg_tuple.1#347, %arg_tuple.1#348, %arg_tuple.1#349, /*index=350*/%arg_tuple.1#350, %arg_tuple.1#351, %arg_tuple.1#352, %arg_tuple.1#353, %arg_tuple.1#354, /*index=355*/%arg_tuple.1#355, %arg_tuple.1#356, %arg_tuple.1#357, %arg_tuple.1#358, %arg_tuple.1#359, /*index=360*/%arg_tuple.1#360)
+}
+
+%region_1.5 (arg_tuple.3: (s32[], f32[], f32[], f32[], f32[], /*index=5*/f32[], f32[], f32[], f32[], f32[], /*index=10*/f32[], f32[], f32[], f32[], f32[], /*index=15*/f32[], f32[], f32[], f32[], f32[], /*index=20*/f32[], f32[], f32[], f32[], f32[], /*index=25*/f32[], f32[], f32[], f32[], f32[], /*index=30*/f32[], f32[], f32[], f32[], f32[], /*index=35*/f32[], f32[], f32[], f32[], f32[], /*index=40*/f32[], f32[], f32[], f32[], f32[], /*index=45*/f32[], f32[], f32[], f32[], f32[], /*index=50*/f32[], f32[], f32[], f32[], f32[], /*index=55*/f32[], f32[], f32[], f32[], f32[], /*index=60*/f32[], pred[614], f32[614], f32[614], f32[614], /*index=65*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=70*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=75*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=80*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=85*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=90*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=95*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=100*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=105*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=110*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=115*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=120*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=125*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=130*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=135*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=140*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=145*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=150*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=155*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=160*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=165*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=170*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=175*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=180*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=185*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=190*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=195*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=200*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=205*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=210*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=215*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=220*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=225*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=230*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=235*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=240*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=245*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=250*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=255*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=260*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=265*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=270*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=275*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=280*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=285*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=290*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=295*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=300*/f32[614], f32[], f32[], f32[], f32[], /*index=305*/f32[], f32[], f32[], f32[], f32[], /*index=310*/f32[], f32[], f32[], f32[], f32[], /*index=315*/f32[], f32[], f32[], f32[], f32[], /*index=320*/f32[], f32[], f32[], f32[], f32[], /*index=325*/f32[], f32[], f32[], f32[], f32[], /*index=330*/f32[], f32[], f32[], f32[], f32[], /*index=335*/f32[], f32[], f32[], f32[], f32[], /*index=340*/f32[], f32[], f32[], f32[], f32[], /*index=345*/f32[], f32[], f32[], f32[], f32[], /*index=350*/f32[], f32[], f32[], f32[], f32[], /*index=355*/f32[], f32[], f32[], f32[], f32[], /*index=360*/f32[])) -> pred[] {
+  %arg_tuple.3 = (s32[], f32[], f32[], f32[], f32[], /*index=5*/f32[], f32[], f32[], f32[], f32[], /*index=10*/f32[], f32[], f32[], f32[], f32[], /*index=15*/f32[], f32[], f32[], f32[], f32[], /*index=20*/f32[], f32[], f32[], f32[], f32[], /*index=25*/f32[], f32[], f32[], f32[], f32[], /*index=30*/f32[], f32[], f32[], f32[], f32[], /*index=35*/f32[], f32[], f32[], f32[], f32[], /*index=40*/f32[], f32[], f32[], f32[], f32[], /*index=45*/f32[], f32[], f32[], f32[], f32[], /*index=50*/f32[], f32[], f32[], f32[], f32[], /*index=55*/f32[], f32[], f32[], f32[], f32[], /*index=60*/f32[], pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=65*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=70*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=75*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=80*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=85*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=90*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=95*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=100*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=105*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=110*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=115*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=120*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=125*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=130*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=135*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=140*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=145*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=150*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=155*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=160*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=165*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=170*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=175*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=180*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=185*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=190*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=195*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=200*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=205*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=210*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=215*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=220*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=225*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=230*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=235*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=240*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=245*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=250*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=255*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=260*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=265*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=270*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=275*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=280*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=285*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=290*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=295*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=300*/f32[614]{0}, f32[], f32[], f32[], f32[], /*index=305*/f32[], f32[], f32[], f32[], f32[], /*index=310*/f32[], f32[], f32[], f32[], f32[], /*index=315*/f32[], f32[], f32[], f32[], f32[], /*index=320*/f32[], f32[], f32[], f32[], f32[], /*index=325*/f32[], f32[], f32[], f32[], f32[], /*index=330*/f32[], f32[], f32[], f32[], f32[], /*index=335*/f32[], f32[], f32[], f32[], f32[], /*index=340*/f32[], f32[], f32[], f32[], f32[], /*index=345*/f32[], f32[], f32[], f32[], f32[], /*index=350*/f32[], f32[], f32[], f32[], f32[], /*index=355*/f32[], f32[], f32[], f32[], f32[], /*index=360*/f32[]) parameter(0)
+  %constant.20 = s32[] constant(614)
+  ROOT %lt.1 = pred[] compare(%arg_tuple.3#0, %constant.20), direction=LT
+}
+
+%dynamic_index_in_dim_1.6 (Arg_0.7: pred[614], Arg_1.7: s32[]) -> pred[] {
+  %Arg_0.7 = pred[614]{0} parameter(0)
+  %Arg_1.7 = s32[] parameter(1)
+  %dynamic_slice.3 = pred[1]{0} dynamic-slice(%Arg_0.7, %Arg_1.7), dynamic_slice_sizes={1}
+  ROOT %squeeze.63 = pred[] reshape(%dynamic_slice.3)
+}
+
+%region_2.7 (arg_tuple.5: (s32[], f32[3], f32[], f32[], f32[], /*index=5*/f32[], f32[], f32[], f32[], f32[], /*index=10*/f32[], f32[], f32[], f32[], f32[], /*index=15*/f32[], f32[], f32[], f32[], f32[], /*index=20*/f32[], f32[], f32[], f32[], f32[], /*index=25*/f32[], f32[], f32[], f32[], f32[], /*index=30*/f32[], f32[], f32[], f32[], f32[], /*index=35*/f32[], f32[], f32[], f32[], f32[], /*index=40*/f32[], f32[], f32[], f32[], f32[], /*index=45*/f32[], f32[], f32[], f32[], f32[], /*index=50*/f32[], f32[], f32[], f32[], f32[], /*index=55*/f32[], f32[], f32[], f32[], f32[], /*index=60*/f32[], f32[], pred[614], f32[614], f32[614], /*index=65*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=70*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=75*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=80*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=85*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=90*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=95*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=100*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=105*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=110*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=115*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=120*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=125*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=130*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=135*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=140*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=145*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=150*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=155*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=160*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=165*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=170*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=175*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=180*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=185*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=190*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=195*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=200*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=205*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=210*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=215*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=220*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=225*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=230*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=235*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=240*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=245*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=250*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=255*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=260*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=265*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=270*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=275*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=280*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=285*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=290*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=295*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=300*/f32[614], f32[614])) -> (s32[], f32[3], f32[], f32[], f32[], /*index=5*/f32[], f32[], f32[], f32[], f32[], /*index=10*/f32[], f32[], f32[], f32[], f32[], /*index=15*/f32[], f32[], f32[], f32[], f32[], /*index=20*/f32[], f32[], f32[], f32[], f32[], /*index=25*/f32[], f32[], f32[], f32[], f32[], /*index=30*/f32[], f32[], f32[], f32[], f32[], /*index=35*/f32[], f32[], f32[], f32[], f32[], /*index=40*/f32[], f32[], f32[], f32[], f32[], /*index=45*/f32[], f32[], f32[], f32[], f32[], /*index=50*/f32[], f32[], f32[], f32[], f32[], /*index=55*/f32[], f32[], f32[], f32[], f32[], /*index=60*/f32[], f32[], pred[614], f32[614], f32[614], /*index=65*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=70*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=75*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=80*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=85*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=90*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=95*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=100*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=105*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=110*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=115*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=120*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=125*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=130*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=135*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=140*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=145*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=150*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=155*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=160*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=165*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=170*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=175*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=180*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=185*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=190*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=195*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=200*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=205*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=210*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=215*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=220*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=225*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=230*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=235*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=240*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=245*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=250*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=255*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=260*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=265*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=270*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=275*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=280*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=285*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=290*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=295*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=300*/f32[614], f32[614]) {
+  %arg_tuple.5 = (s32[], f32[3]{0}, f32[], f32[], f32[], /*index=5*/f32[], f32[], f32[], f32[], f32[], /*index=10*/f32[], f32[], f32[], f32[], f32[], /*index=15*/f32[], f32[], f32[], f32[], f32[], /*index=20*/f32[], f32[], f32[], f32[], f32[], /*index=25*/f32[], f32[], f32[], f32[], f32[], /*index=30*/f32[], f32[], f32[], f32[], f32[], /*index=35*/f32[], f32[], f32[], f32[], f32[], /*index=40*/f32[], f32[], f32[], f32[], f32[], /*index=45*/f32[], f32[], f32[], f32[], f32[], /*index=50*/f32[], f32[], f32[], f32[], f32[], /*index=55*/f32[], f32[], f32[], f32[], f32[], /*index=60*/f32[], f32[], pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=65*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=70*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=75*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=80*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=85*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=90*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=95*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=100*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=105*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=110*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=115*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=120*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=125*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=130*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=135*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=140*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=145*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=150*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=155*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=160*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=165*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=170*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=175*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=180*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=185*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=190*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=195*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=200*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=205*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=210*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=215*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=220*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=225*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=230*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=235*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=240*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=245*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=250*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=255*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=260*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=265*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=270*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=275*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=280*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=285*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=290*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=295*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=300*/f32[614]{0}, f32[614]{0}) parameter(0)
+  %constant.33 = s32[] constant(1)
+  %add.483 = s32[] add(%arg_tuple.5#0, %constant.33)
+  %constant.32 = f32[] constant(1)
+  %broadcast_in_dim.65 = f32[614]{0} broadcast(%constant.32), dimensions={}
+  %constant.34 = s32[] constant(614)
+  %sub.122 = s32[] subtract(%constant.34, %arg_tuple.5#0)
+  %sub.123 = s32[] subtract(%sub.122, %constant.33)
+  %jit_dynamic_index_in_dim_.243 = f32[] call(%broadcast_in_dim.65, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %add_any.300 = f32[] add(%arg_tuple.5#61, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.483 = f32[] call(%arg_tuple.5#301, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1140 = f32[] multiply(%add_any.300, %jit_dynamic_index_in_dim_.483)
+  %jit_dynamic_index_in_dim_.481 = f32[] call(%arg_tuple.5#299, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1142 = f32[] multiply(%mul.1140, %jit_dynamic_index_in_dim_.481)
+  %jit_dynamic_index_in_dim_.480 = pred[] call(%arg_tuple.5#298, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.482 = f32[] call(%arg_tuple.5#300, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1141 = f32[] multiply(%jit_dynamic_index_in_dim_.482, %mul.1140)
+  %mul.1143 = f32[] multiply(%mul.1141, %jit_dynamic_index_in_dim_.481)
+  %constant.31 = f32[] constant(0.1)
+  %mul.1144 = f32[] multiply(%mul.1143, %constant.31)
+  %neg.240 = f32[] negate(%mul.1144)
+  %constant.30 = f32[] constant(0)
+  %select_n.121 = f32[] select(%jit_dynamic_index_in_dim_.480, %neg.240, %constant.30)
+  %add_any.360 = f32[] add(%mul.1142, %select_n.121)
+  %select_n.120 = f32[] select(%jit_dynamic_index_in_dim_.480, %constant.30, %neg.240)
+  %neg.241 = f32[] negate(%select_n.120)
+  %add_any.361 = f32[] add(%add_any.360, %neg.241)
+  %constant.29 = f32[] constant(0.01)
+  %mul.1145 = f32[] multiply(%add_any.361, %constant.29)
+  %broadcast_in_dim.66 = f32[1]{0} reshape(%mul.1145)
+  %pad.60 = f32[3]{0} pad(%broadcast_in_dim.66, %constant.30), padding=2_0
+  %add_any.362 = f32[3]{0} add(%arg_tuple.5#1, %pad.60)
+  %add_any.301 = f32[] add(%arg_tuple.5#60, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.479 = f32[] call(%arg_tuple.5#297, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1148 = f32[] multiply(%add_any.301, %jit_dynamic_index_in_dim_.479)
+  %jit_dynamic_index_in_dim_.477 = f32[] call(%arg_tuple.5#295, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1150 = f32[] multiply(%mul.1148, %jit_dynamic_index_in_dim_.477)
+  %jit_dynamic_index_in_dim_.476 = pred[] call(%arg_tuple.5#294, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.478 = f32[] call(%arg_tuple.5#296, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1149 = f32[] multiply(%jit_dynamic_index_in_dim_.478, %mul.1148)
+  %mul.1151 = f32[] multiply(%mul.1149, %jit_dynamic_index_in_dim_.477)
+  %mul.1152 = f32[] multiply(%mul.1151, %constant.31)
+  %neg.242 = f32[] negate(%mul.1152)
+  %select_n.123 = f32[] select(%jit_dynamic_index_in_dim_.476, %neg.242, %constant.30)
+  %add_any.363 = f32[] add(%mul.1150, %select_n.123)
+  %select_n.122 = f32[] select(%jit_dynamic_index_in_dim_.476, %constant.30, %neg.242)
+  %neg.243 = f32[] negate(%select_n.122)
+  %add_any.364 = f32[] add(%add_any.363, %neg.243)
+  %mul.1153 = f32[] multiply(%add_any.364, %constant.29)
+  %broadcast_in_dim.67 = f32[1]{0} reshape(%mul.1153)
+  %pad.61 = f32[3]{0} pad(%broadcast_in_dim.67, %constant.30), padding=1_1
+  %add_any.365 = f32[3]{0} add(%add_any.362, %pad.61)
+  %add_any.302 = f32[] add(%arg_tuple.5#59, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.475 = f32[] call(%arg_tuple.5#293, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1156 = f32[] multiply(%add_any.302, %jit_dynamic_index_in_dim_.475)
+  %jit_dynamic_index_in_dim_.473 = f32[] call(%arg_tuple.5#291, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1158 = f32[] multiply(%mul.1156, %jit_dynamic_index_in_dim_.473)
+  %jit_dynamic_index_in_dim_.472 = pred[] call(%arg_tuple.5#290, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.474 = f32[] call(%arg_tuple.5#292, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1157 = f32[] multiply(%jit_dynamic_index_in_dim_.474, %mul.1156)
+  %mul.1159 = f32[] multiply(%mul.1157, %jit_dynamic_index_in_dim_.473)
+  %mul.1160 = f32[] multiply(%mul.1159, %constant.31)
+  %neg.244 = f32[] negate(%mul.1160)
+  %select_n.125 = f32[] select(%jit_dynamic_index_in_dim_.472, %neg.244, %constant.30)
+  %add_any.367 = f32[] add(%mul.1158, %select_n.125)
+  %select_n.124 = f32[] select(%jit_dynamic_index_in_dim_.472, %constant.30, %neg.244)
+  %neg.245 = f32[] negate(%select_n.124)
+  %add_any.368 = f32[] add(%add_any.367, %neg.245)
+  %mul.1161 = f32[] multiply(%add_any.368, %constant.29)
+  %broadcast_in_dim.68 = f32[1]{0} reshape(%mul.1161)
+  %pad.62 = f32[3]{0} pad(%broadcast_in_dim.68, %constant.30), padding=0_2
+  %add_any.369 = f32[3]{0} add(%add_any.365, %pad.62)
+  %add_any.303 = f32[] add(%arg_tuple.5#58, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.471 = f32[] call(%arg_tuple.5#289, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1164 = f32[] multiply(%add_any.303, %jit_dynamic_index_in_dim_.471)
+  %jit_dynamic_index_in_dim_.469 = f32[] call(%arg_tuple.5#287, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1166 = f32[] multiply(%mul.1164, %jit_dynamic_index_in_dim_.469)
+  %jit_dynamic_index_in_dim_.468 = pred[] call(%arg_tuple.5#286, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.470 = f32[] call(%arg_tuple.5#288, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1165 = f32[] multiply(%jit_dynamic_index_in_dim_.470, %mul.1164)
+  %mul.1167 = f32[] multiply(%mul.1165, %jit_dynamic_index_in_dim_.469)
+  %mul.1168 = f32[] multiply(%mul.1167, %constant.31)
+  %neg.246 = f32[] negate(%mul.1168)
+  %select_n.127 = f32[] select(%jit_dynamic_index_in_dim_.468, %neg.246, %constant.30)
+  %add_any.371 = f32[] add(%mul.1166, %select_n.127)
+  %select_n.126 = f32[] select(%jit_dynamic_index_in_dim_.468, %constant.30, %neg.246)
+  %neg.247 = f32[] negate(%select_n.126)
+  %add_any.372 = f32[] add(%add_any.371, %neg.247)
+  %mul.1169 = f32[] multiply(%add_any.372, %constant.29)
+  %broadcast_in_dim.69 = f32[1]{0} reshape(%mul.1169)
+  %pad.63 = f32[3]{0} pad(%broadcast_in_dim.69, %constant.30), padding=2_0
+  %add_any.373 = f32[3]{0} add(%add_any.369, %pad.63)
+  %add_any.304 = f32[] add(%arg_tuple.5#57, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.467 = f32[] call(%arg_tuple.5#285, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1172 = f32[] multiply(%add_any.304, %jit_dynamic_index_in_dim_.467)
+  %jit_dynamic_index_in_dim_.465 = f32[] call(%arg_tuple.5#283, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1174 = f32[] multiply(%mul.1172, %jit_dynamic_index_in_dim_.465)
+  %jit_dynamic_index_in_dim_.464 = pred[] call(%arg_tuple.5#282, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.466 = f32[] call(%arg_tuple.5#284, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1173 = f32[] multiply(%jit_dynamic_index_in_dim_.466, %mul.1172)
+  %mul.1175 = f32[] multiply(%mul.1173, %jit_dynamic_index_in_dim_.465)
+  %mul.1176 = f32[] multiply(%mul.1175, %constant.31)
+  %neg.248 = f32[] negate(%mul.1176)
+  %select_n.129 = f32[] select(%jit_dynamic_index_in_dim_.464, %neg.248, %constant.30)
+  %add_any.375 = f32[] add(%mul.1174, %select_n.129)
+  %select_n.128 = f32[] select(%jit_dynamic_index_in_dim_.464, %constant.30, %neg.248)
+  %neg.249 = f32[] negate(%select_n.128)
+  %add_any.376 = f32[] add(%add_any.375, %neg.249)
+  %mul.1177 = f32[] multiply(%add_any.376, %constant.29)
+  %broadcast_in_dim.70 = f32[1]{0} reshape(%mul.1177)
+  %pad.64 = f32[3]{0} pad(%broadcast_in_dim.70, %constant.30), padding=1_1
+  %add_any.377 = f32[3]{0} add(%add_any.373, %pad.64)
+  %add_any.305 = f32[] add(%arg_tuple.5#56, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.463 = f32[] call(%arg_tuple.5#281, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1180 = f32[] multiply(%add_any.305, %jit_dynamic_index_in_dim_.463)
+  %jit_dynamic_index_in_dim_.461 = f32[] call(%arg_tuple.5#279, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1182 = f32[] multiply(%mul.1180, %jit_dynamic_index_in_dim_.461)
+  %jit_dynamic_index_in_dim_.460 = pred[] call(%arg_tuple.5#278, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.462 = f32[] call(%arg_tuple.5#280, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1181 = f32[] multiply(%jit_dynamic_index_in_dim_.462, %mul.1180)
+  %mul.1183 = f32[] multiply(%mul.1181, %jit_dynamic_index_in_dim_.461)
+  %mul.1184 = f32[] multiply(%mul.1183, %constant.31)
+  %neg.250 = f32[] negate(%mul.1184)
+  %select_n.131 = f32[] select(%jit_dynamic_index_in_dim_.460, %neg.250, %constant.30)
+  %add_any.379 = f32[] add(%mul.1182, %select_n.131)
+  %select_n.130 = f32[] select(%jit_dynamic_index_in_dim_.460, %constant.30, %neg.250)
+  %neg.251 = f32[] negate(%select_n.130)
+  %add_any.380 = f32[] add(%add_any.379, %neg.251)
+  %mul.1185 = f32[] multiply(%add_any.380, %constant.29)
+  %broadcast_in_dim.71 = f32[1]{0} reshape(%mul.1185)
+  %pad.65 = f32[3]{0} pad(%broadcast_in_dim.71, %constant.30), padding=0_2
+  %add_any.381 = f32[3]{0} add(%add_any.377, %pad.65)
+  %add_any.306 = f32[] add(%arg_tuple.5#55, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.459 = f32[] call(%arg_tuple.5#277, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1188 = f32[] multiply(%add_any.306, %jit_dynamic_index_in_dim_.459)
+  %jit_dynamic_index_in_dim_.457 = f32[] call(%arg_tuple.5#275, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1190 = f32[] multiply(%mul.1188, %jit_dynamic_index_in_dim_.457)
+  %jit_dynamic_index_in_dim_.456 = pred[] call(%arg_tuple.5#274, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.458 = f32[] call(%arg_tuple.5#276, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1189 = f32[] multiply(%jit_dynamic_index_in_dim_.458, %mul.1188)
+  %mul.1191 = f32[] multiply(%mul.1189, %jit_dynamic_index_in_dim_.457)
+  %mul.1192 = f32[] multiply(%mul.1191, %constant.31)
+  %neg.252 = f32[] negate(%mul.1192)
+  %select_n.133 = f32[] select(%jit_dynamic_index_in_dim_.456, %neg.252, %constant.30)
+  %add_any.383 = f32[] add(%mul.1190, %select_n.133)
+  %select_n.132 = f32[] select(%jit_dynamic_index_in_dim_.456, %constant.30, %neg.252)
+  %neg.253 = f32[] negate(%select_n.132)
+  %add_any.384 = f32[] add(%add_any.383, %neg.253)
+  %mul.1193 = f32[] multiply(%add_any.384, %constant.29)
+  %broadcast_in_dim.72 = f32[1]{0} reshape(%mul.1193)
+  %pad.66 = f32[3]{0} pad(%broadcast_in_dim.72, %constant.30), padding=2_0
+  %add_any.385 = f32[3]{0} add(%add_any.381, %pad.66)
+  %add_any.307 = f32[] add(%arg_tuple.5#54, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.455 = f32[] call(%arg_tuple.5#273, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1196 = f32[] multiply(%add_any.307, %jit_dynamic_index_in_dim_.455)
+  %jit_dynamic_index_in_dim_.453 = f32[] call(%arg_tuple.5#271, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1198 = f32[] multiply(%mul.1196, %jit_dynamic_index_in_dim_.453)
+  %jit_dynamic_index_in_dim_.452 = pred[] call(%arg_tuple.5#270, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.454 = f32[] call(%arg_tuple.5#272, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1197 = f32[] multiply(%jit_dynamic_index_in_dim_.454, %mul.1196)
+  %mul.1199 = f32[] multiply(%mul.1197, %jit_dynamic_index_in_dim_.453)
+  %mul.1200 = f32[] multiply(%mul.1199, %constant.31)
+  %neg.254 = f32[] negate(%mul.1200)
+  %select_n.135 = f32[] select(%jit_dynamic_index_in_dim_.452, %neg.254, %constant.30)
+  %add_any.387 = f32[] add(%mul.1198, %select_n.135)
+  %select_n.134 = f32[] select(%jit_dynamic_index_in_dim_.452, %constant.30, %neg.254)
+  %neg.255 = f32[] negate(%select_n.134)
+  %add_any.388 = f32[] add(%add_any.387, %neg.255)
+  %mul.1201 = f32[] multiply(%add_any.388, %constant.29)
+  %broadcast_in_dim.73 = f32[1]{0} reshape(%mul.1201)
+  %pad.67 = f32[3]{0} pad(%broadcast_in_dim.73, %constant.30), padding=1_1
+  %add_any.389 = f32[3]{0} add(%add_any.385, %pad.67)
+  %add_any.308 = f32[] add(%arg_tuple.5#53, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.451 = f32[] call(%arg_tuple.5#269, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1204 = f32[] multiply(%add_any.308, %jit_dynamic_index_in_dim_.451)
+  %jit_dynamic_index_in_dim_.449 = f32[] call(%arg_tuple.5#267, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1206 = f32[] multiply(%mul.1204, %jit_dynamic_index_in_dim_.449)
+  %jit_dynamic_index_in_dim_.448 = pred[] call(%arg_tuple.5#266, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.450 = f32[] call(%arg_tuple.5#268, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1205 = f32[] multiply(%jit_dynamic_index_in_dim_.450, %mul.1204)
+  %mul.1207 = f32[] multiply(%mul.1205, %jit_dynamic_index_in_dim_.449)
+  %mul.1208 = f32[] multiply(%mul.1207, %constant.31)
+  %neg.256 = f32[] negate(%mul.1208)
+  %select_n.137 = f32[] select(%jit_dynamic_index_in_dim_.448, %neg.256, %constant.30)
+  %add_any.391 = f32[] add(%mul.1206, %select_n.137)
+  %select_n.136 = f32[] select(%jit_dynamic_index_in_dim_.448, %constant.30, %neg.256)
+  %neg.257 = f32[] negate(%select_n.136)
+  %add_any.392 = f32[] add(%add_any.391, %neg.257)
+  %mul.1209 = f32[] multiply(%add_any.392, %constant.29)
+  %broadcast_in_dim.74 = f32[1]{0} reshape(%mul.1209)
+  %pad.68 = f32[3]{0} pad(%broadcast_in_dim.74, %constant.30), padding=0_2
+  %add_any.393 = f32[3]{0} add(%add_any.389, %pad.68)
+  %add_any.309 = f32[] add(%arg_tuple.5#52, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.447 = f32[] call(%arg_tuple.5#265, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1212 = f32[] multiply(%add_any.309, %jit_dynamic_index_in_dim_.447)
+  %jit_dynamic_index_in_dim_.445 = f32[] call(%arg_tuple.5#263, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1214 = f32[] multiply(%mul.1212, %jit_dynamic_index_in_dim_.445)
+  %jit_dynamic_index_in_dim_.444 = pred[] call(%arg_tuple.5#262, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.446 = f32[] call(%arg_tuple.5#264, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1213 = f32[] multiply(%jit_dynamic_index_in_dim_.446, %mul.1212)
+  %mul.1215 = f32[] multiply(%mul.1213, %jit_dynamic_index_in_dim_.445)
+  %mul.1216 = f32[] multiply(%mul.1215, %constant.31)
+  %neg.258 = f32[] negate(%mul.1216)
+  %select_n.139 = f32[] select(%jit_dynamic_index_in_dim_.444, %neg.258, %constant.30)
+  %add_any.395 = f32[] add(%mul.1214, %select_n.139)
+  %select_n.138 = f32[] select(%jit_dynamic_index_in_dim_.444, %constant.30, %neg.258)
+  %neg.259 = f32[] negate(%select_n.138)
+  %add_any.396 = f32[] add(%add_any.395, %neg.259)
+  %mul.1217 = f32[] multiply(%add_any.396, %constant.29)
+  %broadcast_in_dim.75 = f32[1]{0} reshape(%mul.1217)
+  %pad.69 = f32[3]{0} pad(%broadcast_in_dim.75, %constant.30), padding=2_0
+  %add_any.397 = f32[3]{0} add(%add_any.393, %pad.69)
+  %add_any.310 = f32[] add(%arg_tuple.5#51, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.443 = f32[] call(%arg_tuple.5#261, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1220 = f32[] multiply(%add_any.310, %jit_dynamic_index_in_dim_.443)
+  %jit_dynamic_index_in_dim_.441 = f32[] call(%arg_tuple.5#259, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1222 = f32[] multiply(%mul.1220, %jit_dynamic_index_in_dim_.441)
+  %jit_dynamic_index_in_dim_.440 = pred[] call(%arg_tuple.5#258, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.442 = f32[] call(%arg_tuple.5#260, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1221 = f32[] multiply(%jit_dynamic_index_in_dim_.442, %mul.1220)
+  %mul.1223 = f32[] multiply(%mul.1221, %jit_dynamic_index_in_dim_.441)
+  %mul.1224 = f32[] multiply(%mul.1223, %constant.31)
+  %neg.260 = f32[] negate(%mul.1224)
+  %select_n.141 = f32[] select(%jit_dynamic_index_in_dim_.440, %neg.260, %constant.30)
+  %add_any.399 = f32[] add(%mul.1222, %select_n.141)
+  %select_n.140 = f32[] select(%jit_dynamic_index_in_dim_.440, %constant.30, %neg.260)
+  %neg.261 = f32[] negate(%select_n.140)
+  %add_any.400 = f32[] add(%add_any.399, %neg.261)
+  %mul.1225 = f32[] multiply(%add_any.400, %constant.29)
+  %broadcast_in_dim.76 = f32[1]{0} reshape(%mul.1225)
+  %pad.70 = f32[3]{0} pad(%broadcast_in_dim.76, %constant.30), padding=1_1
+  %add_any.401 = f32[3]{0} add(%add_any.397, %pad.70)
+  %add_any.311 = f32[] add(%arg_tuple.5#50, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.439 = f32[] call(%arg_tuple.5#257, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1228 = f32[] multiply(%add_any.311, %jit_dynamic_index_in_dim_.439)
+  %jit_dynamic_index_in_dim_.437 = f32[] call(%arg_tuple.5#255, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1230 = f32[] multiply(%mul.1228, %jit_dynamic_index_in_dim_.437)
+  %jit_dynamic_index_in_dim_.436 = pred[] call(%arg_tuple.5#254, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.438 = f32[] call(%arg_tuple.5#256, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1229 = f32[] multiply(%jit_dynamic_index_in_dim_.438, %mul.1228)
+  %mul.1231 = f32[] multiply(%mul.1229, %jit_dynamic_index_in_dim_.437)
+  %mul.1232 = f32[] multiply(%mul.1231, %constant.31)
+  %neg.262 = f32[] negate(%mul.1232)
+  %select_n.143 = f32[] select(%jit_dynamic_index_in_dim_.436, %neg.262, %constant.30)
+  %add_any.403 = f32[] add(%mul.1230, %select_n.143)
+  %select_n.142 = f32[] select(%jit_dynamic_index_in_dim_.436, %constant.30, %neg.262)
+  %neg.263 = f32[] negate(%select_n.142)
+  %add_any.404 = f32[] add(%add_any.403, %neg.263)
+  %mul.1233 = f32[] multiply(%add_any.404, %constant.29)
+  %broadcast_in_dim.77 = f32[1]{0} reshape(%mul.1233)
+  %pad.71 = f32[3]{0} pad(%broadcast_in_dim.77, %constant.30), padding=0_2
+  %add_any.405 = f32[3]{0} add(%add_any.401, %pad.71)
+  %add_any.312 = f32[] add(%arg_tuple.5#49, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.435 = f32[] call(%arg_tuple.5#253, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1236 = f32[] multiply(%add_any.312, %jit_dynamic_index_in_dim_.435)
+  %jit_dynamic_index_in_dim_.433 = f32[] call(%arg_tuple.5#251, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1238 = f32[] multiply(%mul.1236, %jit_dynamic_index_in_dim_.433)
+  %jit_dynamic_index_in_dim_.432 = pred[] call(%arg_tuple.5#250, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.434 = f32[] call(%arg_tuple.5#252, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1237 = f32[] multiply(%jit_dynamic_index_in_dim_.434, %mul.1236)
+  %mul.1239 = f32[] multiply(%mul.1237, %jit_dynamic_index_in_dim_.433)
+  %mul.1240 = f32[] multiply(%mul.1239, %constant.31)
+  %neg.264 = f32[] negate(%mul.1240)
+  %select_n.145 = f32[] select(%jit_dynamic_index_in_dim_.432, %neg.264, %constant.30)
+  %add_any.407 = f32[] add(%mul.1238, %select_n.145)
+  %select_n.144 = f32[] select(%jit_dynamic_index_in_dim_.432, %constant.30, %neg.264)
+  %neg.265 = f32[] negate(%select_n.144)
+  %add_any.408 = f32[] add(%add_any.407, %neg.265)
+  %mul.1241 = f32[] multiply(%add_any.408, %constant.29)
+  %broadcast_in_dim.78 = f32[1]{0} reshape(%mul.1241)
+  %pad.72 = f32[3]{0} pad(%broadcast_in_dim.78, %constant.30), padding=2_0
+  %add_any.409 = f32[3]{0} add(%add_any.405, %pad.72)
+  %add_any.313 = f32[] add(%arg_tuple.5#48, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.431 = f32[] call(%arg_tuple.5#249, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1244 = f32[] multiply(%add_any.313, %jit_dynamic_index_in_dim_.431)
+  %jit_dynamic_index_in_dim_.429 = f32[] call(%arg_tuple.5#247, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1246 = f32[] multiply(%mul.1244, %jit_dynamic_index_in_dim_.429)
+  %jit_dynamic_index_in_dim_.428 = pred[] call(%arg_tuple.5#246, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.430 = f32[] call(%arg_tuple.5#248, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1245 = f32[] multiply(%jit_dynamic_index_in_dim_.430, %mul.1244)
+  %mul.1247 = f32[] multiply(%mul.1245, %jit_dynamic_index_in_dim_.429)
+  %mul.1248 = f32[] multiply(%mul.1247, %constant.31)
+  %neg.266 = f32[] negate(%mul.1248)
+  %select_n.147 = f32[] select(%jit_dynamic_index_in_dim_.428, %neg.266, %constant.30)
+  %add_any.411 = f32[] add(%mul.1246, %select_n.147)
+  %select_n.146 = f32[] select(%jit_dynamic_index_in_dim_.428, %constant.30, %neg.266)
+  %neg.267 = f32[] negate(%select_n.146)
+  %add_any.412 = f32[] add(%add_any.411, %neg.267)
+  %mul.1249 = f32[] multiply(%add_any.412, %constant.29)
+  %broadcast_in_dim.79 = f32[1]{0} reshape(%mul.1249)
+  %pad.73 = f32[3]{0} pad(%broadcast_in_dim.79, %constant.30), padding=1_1
+  %add_any.413 = f32[3]{0} add(%add_any.409, %pad.73)
+  %add_any.314 = f32[] add(%arg_tuple.5#47, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.427 = f32[] call(%arg_tuple.5#245, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1252 = f32[] multiply(%add_any.314, %jit_dynamic_index_in_dim_.427)
+  %jit_dynamic_index_in_dim_.425 = f32[] call(%arg_tuple.5#243, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1254 = f32[] multiply(%mul.1252, %jit_dynamic_index_in_dim_.425)
+  %jit_dynamic_index_in_dim_.424 = pred[] call(%arg_tuple.5#242, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.426 = f32[] call(%arg_tuple.5#244, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1253 = f32[] multiply(%jit_dynamic_index_in_dim_.426, %mul.1252)
+  %mul.1255 = f32[] multiply(%mul.1253, %jit_dynamic_index_in_dim_.425)
+  %mul.1256 = f32[] multiply(%mul.1255, %constant.31)
+  %neg.268 = f32[] negate(%mul.1256)
+  %select_n.149 = f32[] select(%jit_dynamic_index_in_dim_.424, %neg.268, %constant.30)
+  %add_any.415 = f32[] add(%mul.1254, %select_n.149)
+  %select_n.148 = f32[] select(%jit_dynamic_index_in_dim_.424, %constant.30, %neg.268)
+  %neg.269 = f32[] negate(%select_n.148)
+  %add_any.416 = f32[] add(%add_any.415, %neg.269)
+  %mul.1257 = f32[] multiply(%add_any.416, %constant.29)
+  %broadcast_in_dim.80 = f32[1]{0} reshape(%mul.1257)
+  %pad.74 = f32[3]{0} pad(%broadcast_in_dim.80, %constant.30), padding=0_2
+  %add_any.417 = f32[3]{0} add(%add_any.413, %pad.74)
+  %add_any.315 = f32[] add(%arg_tuple.5#46, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.423 = f32[] call(%arg_tuple.5#241, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1260 = f32[] multiply(%add_any.315, %jit_dynamic_index_in_dim_.423)
+  %jit_dynamic_index_in_dim_.421 = f32[] call(%arg_tuple.5#239, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1262 = f32[] multiply(%mul.1260, %jit_dynamic_index_in_dim_.421)
+  %jit_dynamic_index_in_dim_.420 = pred[] call(%arg_tuple.5#238, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.422 = f32[] call(%arg_tuple.5#240, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1261 = f32[] multiply(%jit_dynamic_index_in_dim_.422, %mul.1260)
+  %mul.1263 = f32[] multiply(%mul.1261, %jit_dynamic_index_in_dim_.421)
+  %mul.1264 = f32[] multiply(%mul.1263, %constant.31)
+  %neg.270 = f32[] negate(%mul.1264)
+  %select_n.151 = f32[] select(%jit_dynamic_index_in_dim_.420, %neg.270, %constant.30)
+  %add_any.419 = f32[] add(%mul.1262, %select_n.151)
+  %select_n.150 = f32[] select(%jit_dynamic_index_in_dim_.420, %constant.30, %neg.270)
+  %neg.271 = f32[] negate(%select_n.150)
+  %add_any.420 = f32[] add(%add_any.419, %neg.271)
+  %mul.1265 = f32[] multiply(%add_any.420, %constant.29)
+  %broadcast_in_dim.81 = f32[1]{0} reshape(%mul.1265)
+  %pad.75 = f32[3]{0} pad(%broadcast_in_dim.81, %constant.30), padding=2_0
+  %add_any.421 = f32[3]{0} add(%add_any.417, %pad.75)
+  %add_any.316 = f32[] add(%arg_tuple.5#45, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.419 = f32[] call(%arg_tuple.5#237, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1268 = f32[] multiply(%add_any.316, %jit_dynamic_index_in_dim_.419)
+  %jit_dynamic_index_in_dim_.417 = f32[] call(%arg_tuple.5#235, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1270 = f32[] multiply(%mul.1268, %jit_dynamic_index_in_dim_.417)
+  %jit_dynamic_index_in_dim_.416 = pred[] call(%arg_tuple.5#234, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.418 = f32[] call(%arg_tuple.5#236, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1269 = f32[] multiply(%jit_dynamic_index_in_dim_.418, %mul.1268)
+  %mul.1271 = f32[] multiply(%mul.1269, %jit_dynamic_index_in_dim_.417)
+  %mul.1272 = f32[] multiply(%mul.1271, %constant.31)
+  %neg.272 = f32[] negate(%mul.1272)
+  %select_n.153 = f32[] select(%jit_dynamic_index_in_dim_.416, %neg.272, %constant.30)
+  %add_any.423 = f32[] add(%mul.1270, %select_n.153)
+  %select_n.152 = f32[] select(%jit_dynamic_index_in_dim_.416, %constant.30, %neg.272)
+  %neg.273 = f32[] negate(%select_n.152)
+  %add_any.424 = f32[] add(%add_any.423, %neg.273)
+  %mul.1273 = f32[] multiply(%add_any.424, %constant.29)
+  %broadcast_in_dim.82 = f32[1]{0} reshape(%mul.1273)
+  %pad.76 = f32[3]{0} pad(%broadcast_in_dim.82, %constant.30), padding=1_1
+  %add_any.425 = f32[3]{0} add(%add_any.421, %pad.76)
+  %add_any.317 = f32[] add(%arg_tuple.5#44, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.415 = f32[] call(%arg_tuple.5#233, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1276 = f32[] multiply(%add_any.317, %jit_dynamic_index_in_dim_.415)
+  %jit_dynamic_index_in_dim_.413 = f32[] call(%arg_tuple.5#231, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1278 = f32[] multiply(%mul.1276, %jit_dynamic_index_in_dim_.413)
+  %jit_dynamic_index_in_dim_.412 = pred[] call(%arg_tuple.5#230, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.414 = f32[] call(%arg_tuple.5#232, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1277 = f32[] multiply(%jit_dynamic_index_in_dim_.414, %mul.1276)
+  %mul.1279 = f32[] multiply(%mul.1277, %jit_dynamic_index_in_dim_.413)
+  %mul.1280 = f32[] multiply(%mul.1279, %constant.31)
+  %neg.274 = f32[] negate(%mul.1280)
+  %select_n.155 = f32[] select(%jit_dynamic_index_in_dim_.412, %neg.274, %constant.30)
+  %add_any.427 = f32[] add(%mul.1278, %select_n.155)
+  %select_n.154 = f32[] select(%jit_dynamic_index_in_dim_.412, %constant.30, %neg.274)
+  %neg.275 = f32[] negate(%select_n.154)
+  %add_any.428 = f32[] add(%add_any.427, %neg.275)
+  %mul.1281 = f32[] multiply(%add_any.428, %constant.29)
+  %broadcast_in_dim.83 = f32[1]{0} reshape(%mul.1281)
+  %pad.77 = f32[3]{0} pad(%broadcast_in_dim.83, %constant.30), padding=0_2
+  %add_any.429 = f32[3]{0} add(%add_any.425, %pad.77)
+  %add_any.318 = f32[] add(%arg_tuple.5#43, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.411 = f32[] call(%arg_tuple.5#229, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1284 = f32[] multiply(%add_any.318, %jit_dynamic_index_in_dim_.411)
+  %jit_dynamic_index_in_dim_.409 = f32[] call(%arg_tuple.5#227, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1286 = f32[] multiply(%mul.1284, %jit_dynamic_index_in_dim_.409)
+  %jit_dynamic_index_in_dim_.408 = pred[] call(%arg_tuple.5#226, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.410 = f32[] call(%arg_tuple.5#228, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1285 = f32[] multiply(%jit_dynamic_index_in_dim_.410, %mul.1284)
+  %mul.1287 = f32[] multiply(%mul.1285, %jit_dynamic_index_in_dim_.409)
+  %mul.1288 = f32[] multiply(%mul.1287, %constant.31)
+  %neg.276 = f32[] negate(%mul.1288)
+  %select_n.157 = f32[] select(%jit_dynamic_index_in_dim_.408, %neg.276, %constant.30)
+  %add_any.431 = f32[] add(%mul.1286, %select_n.157)
+  %select_n.156 = f32[] select(%jit_dynamic_index_in_dim_.408, %constant.30, %neg.276)
+  %neg.277 = f32[] negate(%select_n.156)
+  %add_any.432 = f32[] add(%add_any.431, %neg.277)
+  %mul.1289 = f32[] multiply(%add_any.432, %constant.29)
+  %broadcast_in_dim.84 = f32[1]{0} reshape(%mul.1289)
+  %pad.78 = f32[3]{0} pad(%broadcast_in_dim.84, %constant.30), padding=2_0
+  %add_any.433 = f32[3]{0} add(%add_any.429, %pad.78)
+  %add_any.319 = f32[] add(%arg_tuple.5#42, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.407 = f32[] call(%arg_tuple.5#225, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1292 = f32[] multiply(%add_any.319, %jit_dynamic_index_in_dim_.407)
+  %jit_dynamic_index_in_dim_.405 = f32[] call(%arg_tuple.5#223, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1294 = f32[] multiply(%mul.1292, %jit_dynamic_index_in_dim_.405)
+  %jit_dynamic_index_in_dim_.404 = pred[] call(%arg_tuple.5#222, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.406 = f32[] call(%arg_tuple.5#224, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1293 = f32[] multiply(%jit_dynamic_index_in_dim_.406, %mul.1292)
+  %mul.1295 = f32[] multiply(%mul.1293, %jit_dynamic_index_in_dim_.405)
+  %mul.1296 = f32[] multiply(%mul.1295, %constant.31)
+  %neg.278 = f32[] negate(%mul.1296)
+  %select_n.159 = f32[] select(%jit_dynamic_index_in_dim_.404, %neg.278, %constant.30)
+  %add_any.435 = f32[] add(%mul.1294, %select_n.159)
+  %select_n.158 = f32[] select(%jit_dynamic_index_in_dim_.404, %constant.30, %neg.278)
+  %neg.279 = f32[] negate(%select_n.158)
+  %add_any.436 = f32[] add(%add_any.435, %neg.279)
+  %mul.1297 = f32[] multiply(%add_any.436, %constant.29)
+  %broadcast_in_dim.85 = f32[1]{0} reshape(%mul.1297)
+  %pad.79 = f32[3]{0} pad(%broadcast_in_dim.85, %constant.30), padding=1_1
+  %add_any.437 = f32[3]{0} add(%add_any.433, %pad.79)
+  %add_any.320 = f32[] add(%arg_tuple.5#41, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.403 = f32[] call(%arg_tuple.5#221, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1300 = f32[] multiply(%add_any.320, %jit_dynamic_index_in_dim_.403)
+  %jit_dynamic_index_in_dim_.401 = f32[] call(%arg_tuple.5#219, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1302 = f32[] multiply(%mul.1300, %jit_dynamic_index_in_dim_.401)
+  %jit_dynamic_index_in_dim_.400 = pred[] call(%arg_tuple.5#218, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.402 = f32[] call(%arg_tuple.5#220, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1301 = f32[] multiply(%jit_dynamic_index_in_dim_.402, %mul.1300)
+  %mul.1303 = f32[] multiply(%mul.1301, %jit_dynamic_index_in_dim_.401)
+  %mul.1304 = f32[] multiply(%mul.1303, %constant.31)
+  %neg.280 = f32[] negate(%mul.1304)
+  %select_n.161 = f32[] select(%jit_dynamic_index_in_dim_.400, %neg.280, %constant.30)
+  %add_any.439 = f32[] add(%mul.1302, %select_n.161)
+  %select_n.160 = f32[] select(%jit_dynamic_index_in_dim_.400, %constant.30, %neg.280)
+  %neg.281 = f32[] negate(%select_n.160)
+  %add_any.440 = f32[] add(%add_any.439, %neg.281)
+  %mul.1305 = f32[] multiply(%add_any.440, %constant.29)
+  %broadcast_in_dim.86 = f32[1]{0} reshape(%mul.1305)
+  %pad.80 = f32[3]{0} pad(%broadcast_in_dim.86, %constant.30), padding=0_2
+  %add_any.441 = f32[3]{0} add(%add_any.437, %pad.80)
+  %add_any.321 = f32[] add(%arg_tuple.5#40, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.399 = f32[] call(%arg_tuple.5#217, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1308 = f32[] multiply(%add_any.321, %jit_dynamic_index_in_dim_.399)
+  %jit_dynamic_index_in_dim_.397 = f32[] call(%arg_tuple.5#215, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1310 = f32[] multiply(%mul.1308, %jit_dynamic_index_in_dim_.397)
+  %jit_dynamic_index_in_dim_.396 = pred[] call(%arg_tuple.5#214, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.398 = f32[] call(%arg_tuple.5#216, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1309 = f32[] multiply(%jit_dynamic_index_in_dim_.398, %mul.1308)
+  %mul.1311 = f32[] multiply(%mul.1309, %jit_dynamic_index_in_dim_.397)
+  %mul.1312 = f32[] multiply(%mul.1311, %constant.31)
+  %neg.282 = f32[] negate(%mul.1312)
+  %select_n.163 = f32[] select(%jit_dynamic_index_in_dim_.396, %neg.282, %constant.30)
+  %add_any.443 = f32[] add(%mul.1310, %select_n.163)
+  %select_n.162 = f32[] select(%jit_dynamic_index_in_dim_.396, %constant.30, %neg.282)
+  %neg.283 = f32[] negate(%select_n.162)
+  %add_any.444 = f32[] add(%add_any.443, %neg.283)
+  %mul.1313 = f32[] multiply(%add_any.444, %constant.29)
+  %broadcast_in_dim.87 = f32[1]{0} reshape(%mul.1313)
+  %pad.81 = f32[3]{0} pad(%broadcast_in_dim.87, %constant.30), padding=2_0
+  %add_any.445 = f32[3]{0} add(%add_any.441, %pad.81)
+  %add_any.322 = f32[] add(%arg_tuple.5#39, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.395 = f32[] call(%arg_tuple.5#213, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1316 = f32[] multiply(%add_any.322, %jit_dynamic_index_in_dim_.395)
+  %jit_dynamic_index_in_dim_.393 = f32[] call(%arg_tuple.5#211, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1318 = f32[] multiply(%mul.1316, %jit_dynamic_index_in_dim_.393)
+  %jit_dynamic_index_in_dim_.392 = pred[] call(%arg_tuple.5#210, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.394 = f32[] call(%arg_tuple.5#212, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1317 = f32[] multiply(%jit_dynamic_index_in_dim_.394, %mul.1316)
+  %mul.1319 = f32[] multiply(%mul.1317, %jit_dynamic_index_in_dim_.393)
+  %mul.1320 = f32[] multiply(%mul.1319, %constant.31)
+  %neg.284 = f32[] negate(%mul.1320)
+  %select_n.165 = f32[] select(%jit_dynamic_index_in_dim_.392, %neg.284, %constant.30)
+  %add_any.447 = f32[] add(%mul.1318, %select_n.165)
+  %select_n.164 = f32[] select(%jit_dynamic_index_in_dim_.392, %constant.30, %neg.284)
+  %neg.285 = f32[] negate(%select_n.164)
+  %add_any.448 = f32[] add(%add_any.447, %neg.285)
+  %mul.1321 = f32[] multiply(%add_any.448, %constant.29)
+  %broadcast_in_dim.88 = f32[1]{0} reshape(%mul.1321)
+  %pad.82 = f32[3]{0} pad(%broadcast_in_dim.88, %constant.30), padding=1_1
+  %add_any.449 = f32[3]{0} add(%add_any.445, %pad.82)
+  %add_any.323 = f32[] add(%arg_tuple.5#38, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.391 = f32[] call(%arg_tuple.5#209, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1324 = f32[] multiply(%add_any.323, %jit_dynamic_index_in_dim_.391)
+  %jit_dynamic_index_in_dim_.389 = f32[] call(%arg_tuple.5#207, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1326 = f32[] multiply(%mul.1324, %jit_dynamic_index_in_dim_.389)
+  %jit_dynamic_index_in_dim_.388 = pred[] call(%arg_tuple.5#206, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.390 = f32[] call(%arg_tuple.5#208, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1325 = f32[] multiply(%jit_dynamic_index_in_dim_.390, %mul.1324)
+  %mul.1327 = f32[] multiply(%mul.1325, %jit_dynamic_index_in_dim_.389)
+  %mul.1328 = f32[] multiply(%mul.1327, %constant.31)
+  %neg.286 = f32[] negate(%mul.1328)
+  %select_n.167 = f32[] select(%jit_dynamic_index_in_dim_.388, %neg.286, %constant.30)
+  %add_any.451 = f32[] add(%mul.1326, %select_n.167)
+  %select_n.166 = f32[] select(%jit_dynamic_index_in_dim_.388, %constant.30, %neg.286)
+  %neg.287 = f32[] negate(%select_n.166)
+  %add_any.452 = f32[] add(%add_any.451, %neg.287)
+  %mul.1329 = f32[] multiply(%add_any.452, %constant.29)
+  %broadcast_in_dim.89 = f32[1]{0} reshape(%mul.1329)
+  %pad.83 = f32[3]{0} pad(%broadcast_in_dim.89, %constant.30), padding=0_2
+  %add_any.453 = f32[3]{0} add(%add_any.449, %pad.83)
+  %add_any.324 = f32[] add(%arg_tuple.5#37, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.387 = f32[] call(%arg_tuple.5#205, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1332 = f32[] multiply(%add_any.324, %jit_dynamic_index_in_dim_.387)
+  %jit_dynamic_index_in_dim_.385 = f32[] call(%arg_tuple.5#203, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1334 = f32[] multiply(%mul.1332, %jit_dynamic_index_in_dim_.385)
+  %jit_dynamic_index_in_dim_.384 = pred[] call(%arg_tuple.5#202, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.386 = f32[] call(%arg_tuple.5#204, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1333 = f32[] multiply(%jit_dynamic_index_in_dim_.386, %mul.1332)
+  %mul.1335 = f32[] multiply(%mul.1333, %jit_dynamic_index_in_dim_.385)
+  %mul.1336 = f32[] multiply(%mul.1335, %constant.31)
+  %neg.288 = f32[] negate(%mul.1336)
+  %select_n.169 = f32[] select(%jit_dynamic_index_in_dim_.384, %neg.288, %constant.30)
+  %add_any.455 = f32[] add(%mul.1334, %select_n.169)
+  %select_n.168 = f32[] select(%jit_dynamic_index_in_dim_.384, %constant.30, %neg.288)
+  %neg.289 = f32[] negate(%select_n.168)
+  %add_any.456 = f32[] add(%add_any.455, %neg.289)
+  %mul.1337 = f32[] multiply(%add_any.456, %constant.29)
+  %broadcast_in_dim.90 = f32[1]{0} reshape(%mul.1337)
+  %pad.84 = f32[3]{0} pad(%broadcast_in_dim.90, %constant.30), padding=2_0
+  %add_any.457 = f32[3]{0} add(%add_any.453, %pad.84)
+  %add_any.325 = f32[] add(%arg_tuple.5#36, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.383 = f32[] call(%arg_tuple.5#201, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1340 = f32[] multiply(%add_any.325, %jit_dynamic_index_in_dim_.383)
+  %jit_dynamic_index_in_dim_.381 = f32[] call(%arg_tuple.5#199, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1342 = f32[] multiply(%mul.1340, %jit_dynamic_index_in_dim_.381)
+  %jit_dynamic_index_in_dim_.380 = pred[] call(%arg_tuple.5#198, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.382 = f32[] call(%arg_tuple.5#200, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1341 = f32[] multiply(%jit_dynamic_index_in_dim_.382, %mul.1340)
+  %mul.1343 = f32[] multiply(%mul.1341, %jit_dynamic_index_in_dim_.381)
+  %mul.1344 = f32[] multiply(%mul.1343, %constant.31)
+  %neg.290 = f32[] negate(%mul.1344)
+  %select_n.171 = f32[] select(%jit_dynamic_index_in_dim_.380, %neg.290, %constant.30)
+  %add_any.459 = f32[] add(%mul.1342, %select_n.171)
+  %select_n.170 = f32[] select(%jit_dynamic_index_in_dim_.380, %constant.30, %neg.290)
+  %neg.291 = f32[] negate(%select_n.170)
+  %add_any.460 = f32[] add(%add_any.459, %neg.291)
+  %mul.1345 = f32[] multiply(%add_any.460, %constant.29)
+  %broadcast_in_dim.91 = f32[1]{0} reshape(%mul.1345)
+  %pad.85 = f32[3]{0} pad(%broadcast_in_dim.91, %constant.30), padding=1_1
+  %add_any.461 = f32[3]{0} add(%add_any.457, %pad.85)
+  %add_any.326 = f32[] add(%arg_tuple.5#35, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.379 = f32[] call(%arg_tuple.5#197, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1348 = f32[] multiply(%add_any.326, %jit_dynamic_index_in_dim_.379)
+  %jit_dynamic_index_in_dim_.377 = f32[] call(%arg_tuple.5#195, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1350 = f32[] multiply(%mul.1348, %jit_dynamic_index_in_dim_.377)
+  %jit_dynamic_index_in_dim_.376 = pred[] call(%arg_tuple.5#194, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.378 = f32[] call(%arg_tuple.5#196, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1349 = f32[] multiply(%jit_dynamic_index_in_dim_.378, %mul.1348)
+  %mul.1351 = f32[] multiply(%mul.1349, %jit_dynamic_index_in_dim_.377)
+  %mul.1352 = f32[] multiply(%mul.1351, %constant.31)
+  %neg.292 = f32[] negate(%mul.1352)
+  %select_n.173 = f32[] select(%jit_dynamic_index_in_dim_.376, %neg.292, %constant.30)
+  %add_any.463 = f32[] add(%mul.1350, %select_n.173)
+  %select_n.172 = f32[] select(%jit_dynamic_index_in_dim_.376, %constant.30, %neg.292)
+  %neg.293 = f32[] negate(%select_n.172)
+  %add_any.464 = f32[] add(%add_any.463, %neg.293)
+  %mul.1353 = f32[] multiply(%add_any.464, %constant.29)
+  %broadcast_in_dim.92 = f32[1]{0} reshape(%mul.1353)
+  %pad.86 = f32[3]{0} pad(%broadcast_in_dim.92, %constant.30), padding=0_2
+  %add_any.465 = f32[3]{0} add(%add_any.461, %pad.86)
+  %add_any.327 = f32[] add(%arg_tuple.5#34, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.375 = f32[] call(%arg_tuple.5#193, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1356 = f32[] multiply(%add_any.327, %jit_dynamic_index_in_dim_.375)
+  %jit_dynamic_index_in_dim_.373 = f32[] call(%arg_tuple.5#191, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1358 = f32[] multiply(%mul.1356, %jit_dynamic_index_in_dim_.373)
+  %jit_dynamic_index_in_dim_.372 = pred[] call(%arg_tuple.5#190, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.374 = f32[] call(%arg_tuple.5#192, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1357 = f32[] multiply(%jit_dynamic_index_in_dim_.374, %mul.1356)
+  %mul.1359 = f32[] multiply(%mul.1357, %jit_dynamic_index_in_dim_.373)
+  %mul.1360 = f32[] multiply(%mul.1359, %constant.31)
+  %neg.294 = f32[] negate(%mul.1360)
+  %select_n.175 = f32[] select(%jit_dynamic_index_in_dim_.372, %neg.294, %constant.30)
+  %add_any.467 = f32[] add(%mul.1358, %select_n.175)
+  %select_n.174 = f32[] select(%jit_dynamic_index_in_dim_.372, %constant.30, %neg.294)
+  %neg.295 = f32[] negate(%select_n.174)
+  %add_any.468 = f32[] add(%add_any.467, %neg.295)
+  %mul.1361 = f32[] multiply(%add_any.468, %constant.29)
+  %broadcast_in_dim.93 = f32[1]{0} reshape(%mul.1361)
+  %pad.87 = f32[3]{0} pad(%broadcast_in_dim.93, %constant.30), padding=2_0
+  %add_any.469 = f32[3]{0} add(%add_any.465, %pad.87)
+  %add_any.328 = f32[] add(%arg_tuple.5#33, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.371 = f32[] call(%arg_tuple.5#189, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1364 = f32[] multiply(%add_any.328, %jit_dynamic_index_in_dim_.371)
+  %jit_dynamic_index_in_dim_.369 = f32[] call(%arg_tuple.5#187, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1366 = f32[] multiply(%mul.1364, %jit_dynamic_index_in_dim_.369)
+  %jit_dynamic_index_in_dim_.368 = pred[] call(%arg_tuple.5#186, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.370 = f32[] call(%arg_tuple.5#188, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1365 = f32[] multiply(%jit_dynamic_index_in_dim_.370, %mul.1364)
+  %mul.1367 = f32[] multiply(%mul.1365, %jit_dynamic_index_in_dim_.369)
+  %mul.1368 = f32[] multiply(%mul.1367, %constant.31)
+  %neg.296 = f32[] negate(%mul.1368)
+  %select_n.177 = f32[] select(%jit_dynamic_index_in_dim_.368, %neg.296, %constant.30)
+  %add_any.471 = f32[] add(%mul.1366, %select_n.177)
+  %select_n.176 = f32[] select(%jit_dynamic_index_in_dim_.368, %constant.30, %neg.296)
+  %neg.297 = f32[] negate(%select_n.176)
+  %add_any.472 = f32[] add(%add_any.471, %neg.297)
+  %mul.1369 = f32[] multiply(%add_any.472, %constant.29)
+  %broadcast_in_dim.94 = f32[1]{0} reshape(%mul.1369)
+  %pad.88 = f32[3]{0} pad(%broadcast_in_dim.94, %constant.30), padding=1_1
+  %add_any.473 = f32[3]{0} add(%add_any.469, %pad.88)
+  %add_any.329 = f32[] add(%arg_tuple.5#32, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.367 = f32[] call(%arg_tuple.5#185, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1372 = f32[] multiply(%add_any.329, %jit_dynamic_index_in_dim_.367)
+  %jit_dynamic_index_in_dim_.365 = f32[] call(%arg_tuple.5#183, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1374 = f32[] multiply(%mul.1372, %jit_dynamic_index_in_dim_.365)
+  %jit_dynamic_index_in_dim_.364 = pred[] call(%arg_tuple.5#182, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.366 = f32[] call(%arg_tuple.5#184, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1373 = f32[] multiply(%jit_dynamic_index_in_dim_.366, %mul.1372)
+  %mul.1375 = f32[] multiply(%mul.1373, %jit_dynamic_index_in_dim_.365)
+  %mul.1376 = f32[] multiply(%mul.1375, %constant.31)
+  %neg.298 = f32[] negate(%mul.1376)
+  %select_n.179 = f32[] select(%jit_dynamic_index_in_dim_.364, %neg.298, %constant.30)
+  %add_any.475 = f32[] add(%mul.1374, %select_n.179)
+  %select_n.178 = f32[] select(%jit_dynamic_index_in_dim_.364, %constant.30, %neg.298)
+  %neg.299 = f32[] negate(%select_n.178)
+  %add_any.476 = f32[] add(%add_any.475, %neg.299)
+  %mul.1377 = f32[] multiply(%add_any.476, %constant.29)
+  %broadcast_in_dim.95 = f32[1]{0} reshape(%mul.1377)
+  %pad.89 = f32[3]{0} pad(%broadcast_in_dim.95, %constant.30), padding=0_2
+  %add_any.477 = f32[3]{0} add(%add_any.473, %pad.89)
+  %add_any.330 = f32[] add(%arg_tuple.5#31, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.363 = f32[] call(%arg_tuple.5#181, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1380 = f32[] multiply(%add_any.330, %jit_dynamic_index_in_dim_.363)
+  %jit_dynamic_index_in_dim_.361 = f32[] call(%arg_tuple.5#179, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1382 = f32[] multiply(%mul.1380, %jit_dynamic_index_in_dim_.361)
+  %jit_dynamic_index_in_dim_.360 = pred[] call(%arg_tuple.5#178, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.362 = f32[] call(%arg_tuple.5#180, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1381 = f32[] multiply(%jit_dynamic_index_in_dim_.362, %mul.1380)
+  %mul.1383 = f32[] multiply(%mul.1381, %jit_dynamic_index_in_dim_.361)
+  %mul.1384 = f32[] multiply(%mul.1383, %constant.31)
+  %neg.300 = f32[] negate(%mul.1384)
+  %select_n.181 = f32[] select(%jit_dynamic_index_in_dim_.360, %neg.300, %constant.30)
+  %add_any.479 = f32[] add(%mul.1382, %select_n.181)
+  %select_n.180 = f32[] select(%jit_dynamic_index_in_dim_.360, %constant.30, %neg.300)
+  %neg.301 = f32[] negate(%select_n.180)
+  %add_any.480 = f32[] add(%add_any.479, %neg.301)
+  %mul.1385 = f32[] multiply(%add_any.480, %constant.29)
+  %broadcast_in_dim.96 = f32[1]{0} reshape(%mul.1385)
+  %pad.90 = f32[3]{0} pad(%broadcast_in_dim.96, %constant.30), padding=2_0
+  %add_any.481 = f32[3]{0} add(%add_any.477, %pad.90)
+  %add_any.331 = f32[] add(%arg_tuple.5#30, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.359 = f32[] call(%arg_tuple.5#177, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1388 = f32[] multiply(%add_any.331, %jit_dynamic_index_in_dim_.359)
+  %jit_dynamic_index_in_dim_.357 = f32[] call(%arg_tuple.5#175, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1390 = f32[] multiply(%mul.1388, %jit_dynamic_index_in_dim_.357)
+  %jit_dynamic_index_in_dim_.356 = pred[] call(%arg_tuple.5#174, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.358 = f32[] call(%arg_tuple.5#176, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1389 = f32[] multiply(%jit_dynamic_index_in_dim_.358, %mul.1388)
+  %mul.1391 = f32[] multiply(%mul.1389, %jit_dynamic_index_in_dim_.357)
+  %mul.1392 = f32[] multiply(%mul.1391, %constant.31)
+  %neg.302 = f32[] negate(%mul.1392)
+  %select_n.183 = f32[] select(%jit_dynamic_index_in_dim_.356, %neg.302, %constant.30)
+  %add_any.483 = f32[] add(%mul.1390, %select_n.183)
+  %select_n.182 = f32[] select(%jit_dynamic_index_in_dim_.356, %constant.30, %neg.302)
+  %neg.303 = f32[] negate(%select_n.182)
+  %add_any.484 = f32[] add(%add_any.483, %neg.303)
+  %mul.1393 = f32[] multiply(%add_any.484, %constant.29)
+  %broadcast_in_dim.97 = f32[1]{0} reshape(%mul.1393)
+  %pad.91 = f32[3]{0} pad(%broadcast_in_dim.97, %constant.30), padding=1_1
+  %add_any.485 = f32[3]{0} add(%add_any.481, %pad.91)
+  %add_any.332 = f32[] add(%arg_tuple.5#29, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.355 = f32[] call(%arg_tuple.5#173, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1396 = f32[] multiply(%add_any.332, %jit_dynamic_index_in_dim_.355)
+  %jit_dynamic_index_in_dim_.353 = f32[] call(%arg_tuple.5#171, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1398 = f32[] multiply(%mul.1396, %jit_dynamic_index_in_dim_.353)
+  %jit_dynamic_index_in_dim_.352 = pred[] call(%arg_tuple.5#170, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.354 = f32[] call(%arg_tuple.5#172, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1397 = f32[] multiply(%jit_dynamic_index_in_dim_.354, %mul.1396)
+  %mul.1399 = f32[] multiply(%mul.1397, %jit_dynamic_index_in_dim_.353)
+  %mul.1400 = f32[] multiply(%mul.1399, %constant.31)
+  %neg.304 = f32[] negate(%mul.1400)
+  %select_n.185 = f32[] select(%jit_dynamic_index_in_dim_.352, %neg.304, %constant.30)
+  %add_any.487 = f32[] add(%mul.1398, %select_n.185)
+  %select_n.184 = f32[] select(%jit_dynamic_index_in_dim_.352, %constant.30, %neg.304)
+  %neg.305 = f32[] negate(%select_n.184)
+  %add_any.488 = f32[] add(%add_any.487, %neg.305)
+  %mul.1401 = f32[] multiply(%add_any.488, %constant.29)
+  %broadcast_in_dim.98 = f32[1]{0} reshape(%mul.1401)
+  %pad.92 = f32[3]{0} pad(%broadcast_in_dim.98, %constant.30), padding=0_2
+  %add_any.489 = f32[3]{0} add(%add_any.485, %pad.92)
+  %add_any.333 = f32[] add(%arg_tuple.5#28, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.351 = f32[] call(%arg_tuple.5#169, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1404 = f32[] multiply(%add_any.333, %jit_dynamic_index_in_dim_.351)
+  %jit_dynamic_index_in_dim_.349 = f32[] call(%arg_tuple.5#167, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1406 = f32[] multiply(%mul.1404, %jit_dynamic_index_in_dim_.349)
+  %jit_dynamic_index_in_dim_.348 = pred[] call(%arg_tuple.5#166, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.350 = f32[] call(%arg_tuple.5#168, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1405 = f32[] multiply(%jit_dynamic_index_in_dim_.350, %mul.1404)
+  %mul.1407 = f32[] multiply(%mul.1405, %jit_dynamic_index_in_dim_.349)
+  %mul.1408 = f32[] multiply(%mul.1407, %constant.31)
+  %neg.306 = f32[] negate(%mul.1408)
+  %select_n.187 = f32[] select(%jit_dynamic_index_in_dim_.348, %neg.306, %constant.30)
+  %add_any.491 = f32[] add(%mul.1406, %select_n.187)
+  %select_n.186 = f32[] select(%jit_dynamic_index_in_dim_.348, %constant.30, %neg.306)
+  %neg.307 = f32[] negate(%select_n.186)
+  %add_any.492 = f32[] add(%add_any.491, %neg.307)
+  %mul.1409 = f32[] multiply(%add_any.492, %constant.29)
+  %broadcast_in_dim.99 = f32[1]{0} reshape(%mul.1409)
+  %pad.93 = f32[3]{0} pad(%broadcast_in_dim.99, %constant.30), padding=2_0
+  %add_any.493 = f32[3]{0} add(%add_any.489, %pad.93)
+  %add_any.334 = f32[] add(%arg_tuple.5#27, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.347 = f32[] call(%arg_tuple.5#165, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1412 = f32[] multiply(%add_any.334, %jit_dynamic_index_in_dim_.347)
+  %jit_dynamic_index_in_dim_.345 = f32[] call(%arg_tuple.5#163, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1414 = f32[] multiply(%mul.1412, %jit_dynamic_index_in_dim_.345)
+  %jit_dynamic_index_in_dim_.344 = pred[] call(%arg_tuple.5#162, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.346 = f32[] call(%arg_tuple.5#164, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1413 = f32[] multiply(%jit_dynamic_index_in_dim_.346, %mul.1412)
+  %mul.1415 = f32[] multiply(%mul.1413, %jit_dynamic_index_in_dim_.345)
+  %mul.1416 = f32[] multiply(%mul.1415, %constant.31)
+  %neg.308 = f32[] negate(%mul.1416)
+  %select_n.189 = f32[] select(%jit_dynamic_index_in_dim_.344, %neg.308, %constant.30)
+  %add_any.495 = f32[] add(%mul.1414, %select_n.189)
+  %select_n.188 = f32[] select(%jit_dynamic_index_in_dim_.344, %constant.30, %neg.308)
+  %neg.309 = f32[] negate(%select_n.188)
+  %add_any.496 = f32[] add(%add_any.495, %neg.309)
+  %mul.1417 = f32[] multiply(%add_any.496, %constant.29)
+  %broadcast_in_dim.100 = f32[1]{0} reshape(%mul.1417)
+  %pad.94 = f32[3]{0} pad(%broadcast_in_dim.100, %constant.30), padding=1_1
+  %add_any.497 = f32[3]{0} add(%add_any.493, %pad.94)
+  %add_any.335 = f32[] add(%arg_tuple.5#26, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.343 = f32[] call(%arg_tuple.5#161, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1420 = f32[] multiply(%add_any.335, %jit_dynamic_index_in_dim_.343)
+  %jit_dynamic_index_in_dim_.341 = f32[] call(%arg_tuple.5#159, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1422 = f32[] multiply(%mul.1420, %jit_dynamic_index_in_dim_.341)
+  %jit_dynamic_index_in_dim_.340 = pred[] call(%arg_tuple.5#158, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.342 = f32[] call(%arg_tuple.5#160, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1421 = f32[] multiply(%jit_dynamic_index_in_dim_.342, %mul.1420)
+  %mul.1423 = f32[] multiply(%mul.1421, %jit_dynamic_index_in_dim_.341)
+  %mul.1424 = f32[] multiply(%mul.1423, %constant.31)
+  %neg.310 = f32[] negate(%mul.1424)
+  %select_n.191 = f32[] select(%jit_dynamic_index_in_dim_.340, %neg.310, %constant.30)
+  %add_any.499 = f32[] add(%mul.1422, %select_n.191)
+  %select_n.190 = f32[] select(%jit_dynamic_index_in_dim_.340, %constant.30, %neg.310)
+  %neg.311 = f32[] negate(%select_n.190)
+  %add_any.500 = f32[] add(%add_any.499, %neg.311)
+  %mul.1425 = f32[] multiply(%add_any.500, %constant.29)
+  %broadcast_in_dim.101 = f32[1]{0} reshape(%mul.1425)
+  %pad.95 = f32[3]{0} pad(%broadcast_in_dim.101, %constant.30), padding=0_2
+  %add_any.501 = f32[3]{0} add(%add_any.497, %pad.95)
+  %add_any.336 = f32[] add(%arg_tuple.5#25, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.339 = f32[] call(%arg_tuple.5#157, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1428 = f32[] multiply(%add_any.336, %jit_dynamic_index_in_dim_.339)
+  %jit_dynamic_index_in_dim_.337 = f32[] call(%arg_tuple.5#155, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1430 = f32[] multiply(%mul.1428, %jit_dynamic_index_in_dim_.337)
+  %jit_dynamic_index_in_dim_.336 = pred[] call(%arg_tuple.5#154, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.338 = f32[] call(%arg_tuple.5#156, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1429 = f32[] multiply(%jit_dynamic_index_in_dim_.338, %mul.1428)
+  %mul.1431 = f32[] multiply(%mul.1429, %jit_dynamic_index_in_dim_.337)
+  %mul.1432 = f32[] multiply(%mul.1431, %constant.31)
+  %neg.312 = f32[] negate(%mul.1432)
+  %select_n.193 = f32[] select(%jit_dynamic_index_in_dim_.336, %neg.312, %constant.30)
+  %add_any.503 = f32[] add(%mul.1430, %select_n.193)
+  %select_n.192 = f32[] select(%jit_dynamic_index_in_dim_.336, %constant.30, %neg.312)
+  %neg.313 = f32[] negate(%select_n.192)
+  %add_any.504 = f32[] add(%add_any.503, %neg.313)
+  %mul.1433 = f32[] multiply(%add_any.504, %constant.29)
+  %broadcast_in_dim.102 = f32[1]{0} reshape(%mul.1433)
+  %pad.96 = f32[3]{0} pad(%broadcast_in_dim.102, %constant.30), padding=2_0
+  %add_any.505 = f32[3]{0} add(%add_any.501, %pad.96)
+  %add_any.337 = f32[] add(%arg_tuple.5#24, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.335 = f32[] call(%arg_tuple.5#153, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1436 = f32[] multiply(%add_any.337, %jit_dynamic_index_in_dim_.335)
+  %jit_dynamic_index_in_dim_.333 = f32[] call(%arg_tuple.5#151, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1438 = f32[] multiply(%mul.1436, %jit_dynamic_index_in_dim_.333)
+  %jit_dynamic_index_in_dim_.332 = pred[] call(%arg_tuple.5#150, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.334 = f32[] call(%arg_tuple.5#152, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1437 = f32[] multiply(%jit_dynamic_index_in_dim_.334, %mul.1436)
+  %mul.1439 = f32[] multiply(%mul.1437, %jit_dynamic_index_in_dim_.333)
+  %mul.1440 = f32[] multiply(%mul.1439, %constant.31)
+  %neg.314 = f32[] negate(%mul.1440)
+  %select_n.195 = f32[] select(%jit_dynamic_index_in_dim_.332, %neg.314, %constant.30)
+  %add_any.507 = f32[] add(%mul.1438, %select_n.195)
+  %select_n.194 = f32[] select(%jit_dynamic_index_in_dim_.332, %constant.30, %neg.314)
+  %neg.315 = f32[] negate(%select_n.194)
+  %add_any.508 = f32[] add(%add_any.507, %neg.315)
+  %mul.1441 = f32[] multiply(%add_any.508, %constant.29)
+  %broadcast_in_dim.103 = f32[1]{0} reshape(%mul.1441)
+  %pad.97 = f32[3]{0} pad(%broadcast_in_dim.103, %constant.30), padding=1_1
+  %add_any.509 = f32[3]{0} add(%add_any.505, %pad.97)
+  %add_any.338 = f32[] add(%arg_tuple.5#23, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.331 = f32[] call(%arg_tuple.5#149, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1444 = f32[] multiply(%add_any.338, %jit_dynamic_index_in_dim_.331)
+  %jit_dynamic_index_in_dim_.329 = f32[] call(%arg_tuple.5#147, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1446 = f32[] multiply(%mul.1444, %jit_dynamic_index_in_dim_.329)
+  %jit_dynamic_index_in_dim_.328 = pred[] call(%arg_tuple.5#146, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.330 = f32[] call(%arg_tuple.5#148, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1445 = f32[] multiply(%jit_dynamic_index_in_dim_.330, %mul.1444)
+  %mul.1447 = f32[] multiply(%mul.1445, %jit_dynamic_index_in_dim_.329)
+  %mul.1448 = f32[] multiply(%mul.1447, %constant.31)
+  %neg.316 = f32[] negate(%mul.1448)
+  %select_n.197 = f32[] select(%jit_dynamic_index_in_dim_.328, %neg.316, %constant.30)
+  %add_any.511 = f32[] add(%mul.1446, %select_n.197)
+  %select_n.196 = f32[] select(%jit_dynamic_index_in_dim_.328, %constant.30, %neg.316)
+  %neg.317 = f32[] negate(%select_n.196)
+  %add_any.512 = f32[] add(%add_any.511, %neg.317)
+  %mul.1449 = f32[] multiply(%add_any.512, %constant.29)
+  %broadcast_in_dim.104 = f32[1]{0} reshape(%mul.1449)
+  %pad.98 = f32[3]{0} pad(%broadcast_in_dim.104, %constant.30), padding=0_2
+  %add_any.513 = f32[3]{0} add(%add_any.509, %pad.98)
+  %add_any.339 = f32[] add(%arg_tuple.5#22, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.327 = f32[] call(%arg_tuple.5#145, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1452 = f32[] multiply(%add_any.339, %jit_dynamic_index_in_dim_.327)
+  %jit_dynamic_index_in_dim_.325 = f32[] call(%arg_tuple.5#143, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1454 = f32[] multiply(%mul.1452, %jit_dynamic_index_in_dim_.325)
+  %jit_dynamic_index_in_dim_.324 = pred[] call(%arg_tuple.5#142, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.326 = f32[] call(%arg_tuple.5#144, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1453 = f32[] multiply(%jit_dynamic_index_in_dim_.326, %mul.1452)
+  %mul.1455 = f32[] multiply(%mul.1453, %jit_dynamic_index_in_dim_.325)
+  %mul.1456 = f32[] multiply(%mul.1455, %constant.31)
+  %neg.318 = f32[] negate(%mul.1456)
+  %select_n.199 = f32[] select(%jit_dynamic_index_in_dim_.324, %neg.318, %constant.30)
+  %add_any.515 = f32[] add(%mul.1454, %select_n.199)
+  %select_n.198 = f32[] select(%jit_dynamic_index_in_dim_.324, %constant.30, %neg.318)
+  %neg.319 = f32[] negate(%select_n.198)
+  %add_any.516 = f32[] add(%add_any.515, %neg.319)
+  %mul.1457 = f32[] multiply(%add_any.516, %constant.29)
+  %broadcast_in_dim.105 = f32[1]{0} reshape(%mul.1457)
+  %pad.99 = f32[3]{0} pad(%broadcast_in_dim.105, %constant.30), padding=2_0
+  %add_any.517 = f32[3]{0} add(%add_any.513, %pad.99)
+  %add_any.340 = f32[] add(%arg_tuple.5#21, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.323 = f32[] call(%arg_tuple.5#141, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1460 = f32[] multiply(%add_any.340, %jit_dynamic_index_in_dim_.323)
+  %jit_dynamic_index_in_dim_.321 = f32[] call(%arg_tuple.5#139, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1462 = f32[] multiply(%mul.1460, %jit_dynamic_index_in_dim_.321)
+  %jit_dynamic_index_in_dim_.320 = pred[] call(%arg_tuple.5#138, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.322 = f32[] call(%arg_tuple.5#140, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1461 = f32[] multiply(%jit_dynamic_index_in_dim_.322, %mul.1460)
+  %mul.1463 = f32[] multiply(%mul.1461, %jit_dynamic_index_in_dim_.321)
+  %mul.1464 = f32[] multiply(%mul.1463, %constant.31)
+  %neg.320 = f32[] negate(%mul.1464)
+  %select_n.201 = f32[] select(%jit_dynamic_index_in_dim_.320, %neg.320, %constant.30)
+  %add_any.519 = f32[] add(%mul.1462, %select_n.201)
+  %select_n.200 = f32[] select(%jit_dynamic_index_in_dim_.320, %constant.30, %neg.320)
+  %neg.321 = f32[] negate(%select_n.200)
+  %add_any.520 = f32[] add(%add_any.519, %neg.321)
+  %mul.1465 = f32[] multiply(%add_any.520, %constant.29)
+  %broadcast_in_dim.106 = f32[1]{0} reshape(%mul.1465)
+  %pad.100 = f32[3]{0} pad(%broadcast_in_dim.106, %constant.30), padding=1_1
+  %add_any.521 = f32[3]{0} add(%add_any.517, %pad.100)
+  %add_any.341 = f32[] add(%arg_tuple.5#20, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.319 = f32[] call(%arg_tuple.5#137, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1468 = f32[] multiply(%add_any.341, %jit_dynamic_index_in_dim_.319)
+  %jit_dynamic_index_in_dim_.317 = f32[] call(%arg_tuple.5#135, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1470 = f32[] multiply(%mul.1468, %jit_dynamic_index_in_dim_.317)
+  %jit_dynamic_index_in_dim_.316 = pred[] call(%arg_tuple.5#134, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.318 = f32[] call(%arg_tuple.5#136, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1469 = f32[] multiply(%jit_dynamic_index_in_dim_.318, %mul.1468)
+  %mul.1471 = f32[] multiply(%mul.1469, %jit_dynamic_index_in_dim_.317)
+  %mul.1472 = f32[] multiply(%mul.1471, %constant.31)
+  %neg.322 = f32[] negate(%mul.1472)
+  %select_n.203 = f32[] select(%jit_dynamic_index_in_dim_.316, %neg.322, %constant.30)
+  %add_any.523 = f32[] add(%mul.1470, %select_n.203)
+  %select_n.202 = f32[] select(%jit_dynamic_index_in_dim_.316, %constant.30, %neg.322)
+  %neg.323 = f32[] negate(%select_n.202)
+  %add_any.524 = f32[] add(%add_any.523, %neg.323)
+  %mul.1473 = f32[] multiply(%add_any.524, %constant.29)
+  %broadcast_in_dim.107 = f32[1]{0} reshape(%mul.1473)
+  %pad.101 = f32[3]{0} pad(%broadcast_in_dim.107, %constant.30), padding=0_2
+  %add_any.525 = f32[3]{0} add(%add_any.521, %pad.101)
+  %add_any.342 = f32[] add(%arg_tuple.5#19, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.315 = f32[] call(%arg_tuple.5#133, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1476 = f32[] multiply(%add_any.342, %jit_dynamic_index_in_dim_.315)
+  %jit_dynamic_index_in_dim_.313 = f32[] call(%arg_tuple.5#131, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1478 = f32[] multiply(%mul.1476, %jit_dynamic_index_in_dim_.313)
+  %jit_dynamic_index_in_dim_.312 = pred[] call(%arg_tuple.5#130, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.314 = f32[] call(%arg_tuple.5#132, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1477 = f32[] multiply(%jit_dynamic_index_in_dim_.314, %mul.1476)
+  %mul.1479 = f32[] multiply(%mul.1477, %jit_dynamic_index_in_dim_.313)
+  %mul.1480 = f32[] multiply(%mul.1479, %constant.31)
+  %neg.324 = f32[] negate(%mul.1480)
+  %select_n.205 = f32[] select(%jit_dynamic_index_in_dim_.312, %neg.324, %constant.30)
+  %add_any.527 = f32[] add(%mul.1478, %select_n.205)
+  %select_n.204 = f32[] select(%jit_dynamic_index_in_dim_.312, %constant.30, %neg.324)
+  %neg.325 = f32[] negate(%select_n.204)
+  %add_any.528 = f32[] add(%add_any.527, %neg.325)
+  %mul.1481 = f32[] multiply(%add_any.528, %constant.29)
+  %broadcast_in_dim.108 = f32[1]{0} reshape(%mul.1481)
+  %pad.102 = f32[3]{0} pad(%broadcast_in_dim.108, %constant.30), padding=2_0
+  %add_any.529 = f32[3]{0} add(%add_any.525, %pad.102)
+  %add_any.343 = f32[] add(%arg_tuple.5#18, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.311 = f32[] call(%arg_tuple.5#129, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1484 = f32[] multiply(%add_any.343, %jit_dynamic_index_in_dim_.311)
+  %jit_dynamic_index_in_dim_.309 = f32[] call(%arg_tuple.5#127, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1486 = f32[] multiply(%mul.1484, %jit_dynamic_index_in_dim_.309)
+  %jit_dynamic_index_in_dim_.308 = pred[] call(%arg_tuple.5#126, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.310 = f32[] call(%arg_tuple.5#128, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1485 = f32[] multiply(%jit_dynamic_index_in_dim_.310, %mul.1484)
+  %mul.1487 = f32[] multiply(%mul.1485, %jit_dynamic_index_in_dim_.309)
+  %mul.1488 = f32[] multiply(%mul.1487, %constant.31)
+  %neg.326 = f32[] negate(%mul.1488)
+  %select_n.207 = f32[] select(%jit_dynamic_index_in_dim_.308, %neg.326, %constant.30)
+  %add_any.531 = f32[] add(%mul.1486, %select_n.207)
+  %select_n.206 = f32[] select(%jit_dynamic_index_in_dim_.308, %constant.30, %neg.326)
+  %neg.327 = f32[] negate(%select_n.206)
+  %add_any.532 = f32[] add(%add_any.531, %neg.327)
+  %mul.1489 = f32[] multiply(%add_any.532, %constant.29)
+  %broadcast_in_dim.109 = f32[1]{0} reshape(%mul.1489)
+  %pad.103 = f32[3]{0} pad(%broadcast_in_dim.109, %constant.30), padding=1_1
+  %add_any.533 = f32[3]{0} add(%add_any.529, %pad.103)
+  %add_any.344 = f32[] add(%arg_tuple.5#17, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.307 = f32[] call(%arg_tuple.5#125, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1492 = f32[] multiply(%add_any.344, %jit_dynamic_index_in_dim_.307)
+  %jit_dynamic_index_in_dim_.305 = f32[] call(%arg_tuple.5#123, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1494 = f32[] multiply(%mul.1492, %jit_dynamic_index_in_dim_.305)
+  %jit_dynamic_index_in_dim_.304 = pred[] call(%arg_tuple.5#122, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.306 = f32[] call(%arg_tuple.5#124, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1493 = f32[] multiply(%jit_dynamic_index_in_dim_.306, %mul.1492)
+  %mul.1495 = f32[] multiply(%mul.1493, %jit_dynamic_index_in_dim_.305)
+  %mul.1496 = f32[] multiply(%mul.1495, %constant.31)
+  %neg.328 = f32[] negate(%mul.1496)
+  %select_n.209 = f32[] select(%jit_dynamic_index_in_dim_.304, %neg.328, %constant.30)
+  %add_any.535 = f32[] add(%mul.1494, %select_n.209)
+  %select_n.208 = f32[] select(%jit_dynamic_index_in_dim_.304, %constant.30, %neg.328)
+  %neg.329 = f32[] negate(%select_n.208)
+  %add_any.536 = f32[] add(%add_any.535, %neg.329)
+  %mul.1497 = f32[] multiply(%add_any.536, %constant.29)
+  %broadcast_in_dim.110 = f32[1]{0} reshape(%mul.1497)
+  %pad.104 = f32[3]{0} pad(%broadcast_in_dim.110, %constant.30), padding=0_2
+  %add_any.537 = f32[3]{0} add(%add_any.533, %pad.104)
+  %add_any.345 = f32[] add(%arg_tuple.5#16, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.303 = f32[] call(%arg_tuple.5#121, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1500 = f32[] multiply(%add_any.345, %jit_dynamic_index_in_dim_.303)
+  %jit_dynamic_index_in_dim_.301 = f32[] call(%arg_tuple.5#119, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1502 = f32[] multiply(%mul.1500, %jit_dynamic_index_in_dim_.301)
+  %jit_dynamic_index_in_dim_.300 = pred[] call(%arg_tuple.5#118, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.302 = f32[] call(%arg_tuple.5#120, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1501 = f32[] multiply(%jit_dynamic_index_in_dim_.302, %mul.1500)
+  %mul.1503 = f32[] multiply(%mul.1501, %jit_dynamic_index_in_dim_.301)
+  %mul.1504 = f32[] multiply(%mul.1503, %constant.31)
+  %neg.330 = f32[] negate(%mul.1504)
+  %select_n.211 = f32[] select(%jit_dynamic_index_in_dim_.300, %neg.330, %constant.30)
+  %add_any.539 = f32[] add(%mul.1502, %select_n.211)
+  %select_n.210 = f32[] select(%jit_dynamic_index_in_dim_.300, %constant.30, %neg.330)
+  %neg.331 = f32[] negate(%select_n.210)
+  %add_any.540 = f32[] add(%add_any.539, %neg.331)
+  %mul.1505 = f32[] multiply(%add_any.540, %constant.29)
+  %broadcast_in_dim.111 = f32[1]{0} reshape(%mul.1505)
+  %pad.105 = f32[3]{0} pad(%broadcast_in_dim.111, %constant.30), padding=2_0
+  %add_any.541 = f32[3]{0} add(%add_any.537, %pad.105)
+  %add_any.346 = f32[] add(%arg_tuple.5#15, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.299 = f32[] call(%arg_tuple.5#117, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1508 = f32[] multiply(%add_any.346, %jit_dynamic_index_in_dim_.299)
+  %jit_dynamic_index_in_dim_.297 = f32[] call(%arg_tuple.5#115, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1510 = f32[] multiply(%mul.1508, %jit_dynamic_index_in_dim_.297)
+  %jit_dynamic_index_in_dim_.296 = pred[] call(%arg_tuple.5#114, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.298 = f32[] call(%arg_tuple.5#116, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1509 = f32[] multiply(%jit_dynamic_index_in_dim_.298, %mul.1508)
+  %mul.1511 = f32[] multiply(%mul.1509, %jit_dynamic_index_in_dim_.297)
+  %mul.1512 = f32[] multiply(%mul.1511, %constant.31)
+  %neg.332 = f32[] negate(%mul.1512)
+  %select_n.213 = f32[] select(%jit_dynamic_index_in_dim_.296, %neg.332, %constant.30)
+  %add_any.543 = f32[] add(%mul.1510, %select_n.213)
+  %select_n.212 = f32[] select(%jit_dynamic_index_in_dim_.296, %constant.30, %neg.332)
+  %neg.333 = f32[] negate(%select_n.212)
+  %add_any.544 = f32[] add(%add_any.543, %neg.333)
+  %mul.1513 = f32[] multiply(%add_any.544, %constant.29)
+  %broadcast_in_dim.112 = f32[1]{0} reshape(%mul.1513)
+  %pad.106 = f32[3]{0} pad(%broadcast_in_dim.112, %constant.30), padding=1_1
+  %add_any.545 = f32[3]{0} add(%add_any.541, %pad.106)
+  %add_any.347 = f32[] add(%arg_tuple.5#14, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.295 = f32[] call(%arg_tuple.5#113, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1516 = f32[] multiply(%add_any.347, %jit_dynamic_index_in_dim_.295)
+  %jit_dynamic_index_in_dim_.293 = f32[] call(%arg_tuple.5#111, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1518 = f32[] multiply(%mul.1516, %jit_dynamic_index_in_dim_.293)
+  %jit_dynamic_index_in_dim_.292 = pred[] call(%arg_tuple.5#110, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.294 = f32[] call(%arg_tuple.5#112, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1517 = f32[] multiply(%jit_dynamic_index_in_dim_.294, %mul.1516)
+  %mul.1519 = f32[] multiply(%mul.1517, %jit_dynamic_index_in_dim_.293)
+  %mul.1520 = f32[] multiply(%mul.1519, %constant.31)
+  %neg.334 = f32[] negate(%mul.1520)
+  %select_n.215 = f32[] select(%jit_dynamic_index_in_dim_.292, %neg.334, %constant.30)
+  %add_any.547 = f32[] add(%mul.1518, %select_n.215)
+  %select_n.214 = f32[] select(%jit_dynamic_index_in_dim_.292, %constant.30, %neg.334)
+  %neg.335 = f32[] negate(%select_n.214)
+  %add_any.548 = f32[] add(%add_any.547, %neg.335)
+  %mul.1521 = f32[] multiply(%add_any.548, %constant.29)
+  %broadcast_in_dim.113 = f32[1]{0} reshape(%mul.1521)
+  %pad.107 = f32[3]{0} pad(%broadcast_in_dim.113, %constant.30), padding=0_2
+  %add_any.549 = f32[3]{0} add(%add_any.545, %pad.107)
+  %add_any.348 = f32[] add(%arg_tuple.5#13, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.291 = f32[] call(%arg_tuple.5#109, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1524 = f32[] multiply(%add_any.348, %jit_dynamic_index_in_dim_.291)
+  %jit_dynamic_index_in_dim_.289 = f32[] call(%arg_tuple.5#107, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1526 = f32[] multiply(%mul.1524, %jit_dynamic_index_in_dim_.289)
+  %jit_dynamic_index_in_dim_.288 = pred[] call(%arg_tuple.5#106, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.290 = f32[] call(%arg_tuple.5#108, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1525 = f32[] multiply(%jit_dynamic_index_in_dim_.290, %mul.1524)
+  %mul.1527 = f32[] multiply(%mul.1525, %jit_dynamic_index_in_dim_.289)
+  %mul.1528 = f32[] multiply(%mul.1527, %constant.31)
+  %neg.336 = f32[] negate(%mul.1528)
+  %select_n.217 = f32[] select(%jit_dynamic_index_in_dim_.288, %neg.336, %constant.30)
+  %add_any.551 = f32[] add(%mul.1526, %select_n.217)
+  %select_n.216 = f32[] select(%jit_dynamic_index_in_dim_.288, %constant.30, %neg.336)
+  %neg.337 = f32[] negate(%select_n.216)
+  %add_any.552 = f32[] add(%add_any.551, %neg.337)
+  %mul.1529 = f32[] multiply(%add_any.552, %constant.29)
+  %broadcast_in_dim.114 = f32[1]{0} reshape(%mul.1529)
+  %pad.108 = f32[3]{0} pad(%broadcast_in_dim.114, %constant.30), padding=2_0
+  %add_any.553 = f32[3]{0} add(%add_any.549, %pad.108)
+  %add_any.349 = f32[] add(%arg_tuple.5#12, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.287 = f32[] call(%arg_tuple.5#105, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1532 = f32[] multiply(%add_any.349, %jit_dynamic_index_in_dim_.287)
+  %jit_dynamic_index_in_dim_.285 = f32[] call(%arg_tuple.5#103, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1534 = f32[] multiply(%mul.1532, %jit_dynamic_index_in_dim_.285)
+  %jit_dynamic_index_in_dim_.284 = pred[] call(%arg_tuple.5#102, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.286 = f32[] call(%arg_tuple.5#104, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1533 = f32[] multiply(%jit_dynamic_index_in_dim_.286, %mul.1532)
+  %mul.1535 = f32[] multiply(%mul.1533, %jit_dynamic_index_in_dim_.285)
+  %mul.1536 = f32[] multiply(%mul.1535, %constant.31)
+  %neg.338 = f32[] negate(%mul.1536)
+  %select_n.219 = f32[] select(%jit_dynamic_index_in_dim_.284, %neg.338, %constant.30)
+  %add_any.555 = f32[] add(%mul.1534, %select_n.219)
+  %select_n.218 = f32[] select(%jit_dynamic_index_in_dim_.284, %constant.30, %neg.338)
+  %neg.339 = f32[] negate(%select_n.218)
+  %add_any.556 = f32[] add(%add_any.555, %neg.339)
+  %mul.1537 = f32[] multiply(%add_any.556, %constant.29)
+  %broadcast_in_dim.115 = f32[1]{0} reshape(%mul.1537)
+  %pad.109 = f32[3]{0} pad(%broadcast_in_dim.115, %constant.30), padding=1_1
+  %add_any.557 = f32[3]{0} add(%add_any.553, %pad.109)
+  %add_any.350 = f32[] add(%arg_tuple.5#11, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.283 = f32[] call(%arg_tuple.5#101, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1540 = f32[] multiply(%add_any.350, %jit_dynamic_index_in_dim_.283)
+  %jit_dynamic_index_in_dim_.281 = f32[] call(%arg_tuple.5#99, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1542 = f32[] multiply(%mul.1540, %jit_dynamic_index_in_dim_.281)
+  %jit_dynamic_index_in_dim_.280 = pred[] call(%arg_tuple.5#98, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.282 = f32[] call(%arg_tuple.5#100, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1541 = f32[] multiply(%jit_dynamic_index_in_dim_.282, %mul.1540)
+  %mul.1543 = f32[] multiply(%mul.1541, %jit_dynamic_index_in_dim_.281)
+  %mul.1544 = f32[] multiply(%mul.1543, %constant.31)
+  %neg.340 = f32[] negate(%mul.1544)
+  %select_n.221 = f32[] select(%jit_dynamic_index_in_dim_.280, %neg.340, %constant.30)
+  %add_any.559 = f32[] add(%mul.1542, %select_n.221)
+  %select_n.220 = f32[] select(%jit_dynamic_index_in_dim_.280, %constant.30, %neg.340)
+  %neg.341 = f32[] negate(%select_n.220)
+  %add_any.560 = f32[] add(%add_any.559, %neg.341)
+  %mul.1545 = f32[] multiply(%add_any.560, %constant.29)
+  %broadcast_in_dim.116 = f32[1]{0} reshape(%mul.1545)
+  %pad.110 = f32[3]{0} pad(%broadcast_in_dim.116, %constant.30), padding=0_2
+  %add_any.561 = f32[3]{0} add(%add_any.557, %pad.110)
+  %add_any.351 = f32[] add(%arg_tuple.5#10, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.279 = f32[] call(%arg_tuple.5#97, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1548 = f32[] multiply(%add_any.351, %jit_dynamic_index_in_dim_.279)
+  %jit_dynamic_index_in_dim_.277 = f32[] call(%arg_tuple.5#95, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1550 = f32[] multiply(%mul.1548, %jit_dynamic_index_in_dim_.277)
+  %jit_dynamic_index_in_dim_.276 = pred[] call(%arg_tuple.5#94, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.278 = f32[] call(%arg_tuple.5#96, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1549 = f32[] multiply(%jit_dynamic_index_in_dim_.278, %mul.1548)
+  %mul.1551 = f32[] multiply(%mul.1549, %jit_dynamic_index_in_dim_.277)
+  %mul.1552 = f32[] multiply(%mul.1551, %constant.31)
+  %neg.342 = f32[] negate(%mul.1552)
+  %select_n.223 = f32[] select(%jit_dynamic_index_in_dim_.276, %neg.342, %constant.30)
+  %add_any.563 = f32[] add(%mul.1550, %select_n.223)
+  %select_n.222 = f32[] select(%jit_dynamic_index_in_dim_.276, %constant.30, %neg.342)
+  %neg.343 = f32[] negate(%select_n.222)
+  %add_any.564 = f32[] add(%add_any.563, %neg.343)
+  %mul.1553 = f32[] multiply(%add_any.564, %constant.29)
+  %broadcast_in_dim.117 = f32[1]{0} reshape(%mul.1553)
+  %pad.111 = f32[3]{0} pad(%broadcast_in_dim.117, %constant.30), padding=2_0
+  %add_any.565 = f32[3]{0} add(%add_any.561, %pad.111)
+  %add_any.352 = f32[] add(%arg_tuple.5#9, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.275 = f32[] call(%arg_tuple.5#93, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1556 = f32[] multiply(%add_any.352, %jit_dynamic_index_in_dim_.275)
+  %jit_dynamic_index_in_dim_.273 = f32[] call(%arg_tuple.5#91, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1558 = f32[] multiply(%mul.1556, %jit_dynamic_index_in_dim_.273)
+  %jit_dynamic_index_in_dim_.272 = pred[] call(%arg_tuple.5#90, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.274 = f32[] call(%arg_tuple.5#92, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1557 = f32[] multiply(%jit_dynamic_index_in_dim_.274, %mul.1556)
+  %mul.1559 = f32[] multiply(%mul.1557, %jit_dynamic_index_in_dim_.273)
+  %mul.1560 = f32[] multiply(%mul.1559, %constant.31)
+  %neg.344 = f32[] negate(%mul.1560)
+  %select_n.225 = f32[] select(%jit_dynamic_index_in_dim_.272, %neg.344, %constant.30)
+  %add_any.567 = f32[] add(%mul.1558, %select_n.225)
+  %select_n.224 = f32[] select(%jit_dynamic_index_in_dim_.272, %constant.30, %neg.344)
+  %neg.345 = f32[] negate(%select_n.224)
+  %add_any.568 = f32[] add(%add_any.567, %neg.345)
+  %mul.1561 = f32[] multiply(%add_any.568, %constant.29)
+  %broadcast_in_dim.118 = f32[1]{0} reshape(%mul.1561)
+  %pad.112 = f32[3]{0} pad(%broadcast_in_dim.118, %constant.30), padding=1_1
+  %add_any.569 = f32[3]{0} add(%add_any.565, %pad.112)
+  %add_any.353 = f32[] add(%arg_tuple.5#8, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.271 = f32[] call(%arg_tuple.5#89, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1564 = f32[] multiply(%add_any.353, %jit_dynamic_index_in_dim_.271)
+  %jit_dynamic_index_in_dim_.269 = f32[] call(%arg_tuple.5#87, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1566 = f32[] multiply(%mul.1564, %jit_dynamic_index_in_dim_.269)
+  %jit_dynamic_index_in_dim_.268 = pred[] call(%arg_tuple.5#86, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.270 = f32[] call(%arg_tuple.5#88, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1565 = f32[] multiply(%jit_dynamic_index_in_dim_.270, %mul.1564)
+  %mul.1567 = f32[] multiply(%mul.1565, %jit_dynamic_index_in_dim_.269)
+  %mul.1568 = f32[] multiply(%mul.1567, %constant.31)
+  %neg.346 = f32[] negate(%mul.1568)
+  %select_n.227 = f32[] select(%jit_dynamic_index_in_dim_.268, %neg.346, %constant.30)
+  %add_any.571 = f32[] add(%mul.1566, %select_n.227)
+  %select_n.226 = f32[] select(%jit_dynamic_index_in_dim_.268, %constant.30, %neg.346)
+  %neg.347 = f32[] negate(%select_n.226)
+  %add_any.572 = f32[] add(%add_any.571, %neg.347)
+  %mul.1569 = f32[] multiply(%add_any.572, %constant.29)
+  %broadcast_in_dim.119 = f32[1]{0} reshape(%mul.1569)
+  %pad.113 = f32[3]{0} pad(%broadcast_in_dim.119, %constant.30), padding=0_2
+  %add_any.573 = f32[3]{0} add(%add_any.569, %pad.113)
+  %add_any.354 = f32[] add(%arg_tuple.5#7, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.267 = f32[] call(%arg_tuple.5#85, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1572 = f32[] multiply(%add_any.354, %jit_dynamic_index_in_dim_.267)
+  %jit_dynamic_index_in_dim_.265 = f32[] call(%arg_tuple.5#83, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1574 = f32[] multiply(%mul.1572, %jit_dynamic_index_in_dim_.265)
+  %jit_dynamic_index_in_dim_.264 = pred[] call(%arg_tuple.5#82, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.266 = f32[] call(%arg_tuple.5#84, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1573 = f32[] multiply(%jit_dynamic_index_in_dim_.266, %mul.1572)
+  %mul.1575 = f32[] multiply(%mul.1573, %jit_dynamic_index_in_dim_.265)
+  %mul.1576 = f32[] multiply(%mul.1575, %constant.31)
+  %neg.348 = f32[] negate(%mul.1576)
+  %select_n.229 = f32[] select(%jit_dynamic_index_in_dim_.264, %neg.348, %constant.30)
+  %add_any.575 = f32[] add(%mul.1574, %select_n.229)
+  %select_n.228 = f32[] select(%jit_dynamic_index_in_dim_.264, %constant.30, %neg.348)
+  %neg.349 = f32[] negate(%select_n.228)
+  %add_any.576 = f32[] add(%add_any.575, %neg.349)
+  %mul.1577 = f32[] multiply(%add_any.576, %constant.29)
+  %broadcast_in_dim.120 = f32[1]{0} reshape(%mul.1577)
+  %pad.114 = f32[3]{0} pad(%broadcast_in_dim.120, %constant.30), padding=2_0
+  %add_any.577 = f32[3]{0} add(%add_any.573, %pad.114)
+  %add_any.355 = f32[] add(%arg_tuple.5#6, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.263 = f32[] call(%arg_tuple.5#81, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1580 = f32[] multiply(%add_any.355, %jit_dynamic_index_in_dim_.263)
+  %jit_dynamic_index_in_dim_.261 = f32[] call(%arg_tuple.5#79, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1582 = f32[] multiply(%mul.1580, %jit_dynamic_index_in_dim_.261)
+  %jit_dynamic_index_in_dim_.260 = pred[] call(%arg_tuple.5#78, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.262 = f32[] call(%arg_tuple.5#80, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1581 = f32[] multiply(%jit_dynamic_index_in_dim_.262, %mul.1580)
+  %mul.1583 = f32[] multiply(%mul.1581, %jit_dynamic_index_in_dim_.261)
+  %mul.1584 = f32[] multiply(%mul.1583, %constant.31)
+  %neg.350 = f32[] negate(%mul.1584)
+  %select_n.231 = f32[] select(%jit_dynamic_index_in_dim_.260, %neg.350, %constant.30)
+  %add_any.579 = f32[] add(%mul.1582, %select_n.231)
+  %select_n.230 = f32[] select(%jit_dynamic_index_in_dim_.260, %constant.30, %neg.350)
+  %neg.351 = f32[] negate(%select_n.230)
+  %add_any.580 = f32[] add(%add_any.579, %neg.351)
+  %mul.1585 = f32[] multiply(%add_any.580, %constant.29)
+  %broadcast_in_dim.121 = f32[1]{0} reshape(%mul.1585)
+  %pad.115 = f32[3]{0} pad(%broadcast_in_dim.121, %constant.30), padding=1_1
+  %add_any.581 = f32[3]{0} add(%add_any.577, %pad.115)
+  %add_any.356 = f32[] add(%arg_tuple.5#5, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.259 = f32[] call(%arg_tuple.5#77, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1588 = f32[] multiply(%add_any.356, %jit_dynamic_index_in_dim_.259)
+  %jit_dynamic_index_in_dim_.257 = f32[] call(%arg_tuple.5#75, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1590 = f32[] multiply(%mul.1588, %jit_dynamic_index_in_dim_.257)
+  %jit_dynamic_index_in_dim_.256 = pred[] call(%arg_tuple.5#74, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.258 = f32[] call(%arg_tuple.5#76, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1589 = f32[] multiply(%jit_dynamic_index_in_dim_.258, %mul.1588)
+  %mul.1591 = f32[] multiply(%mul.1589, %jit_dynamic_index_in_dim_.257)
+  %mul.1592 = f32[] multiply(%mul.1591, %constant.31)
+  %neg.352 = f32[] negate(%mul.1592)
+  %select_n.233 = f32[] select(%jit_dynamic_index_in_dim_.256, %neg.352, %constant.30)
+  %add_any.583 = f32[] add(%mul.1590, %select_n.233)
+  %select_n.232 = f32[] select(%jit_dynamic_index_in_dim_.256, %constant.30, %neg.352)
+  %neg.353 = f32[] negate(%select_n.232)
+  %add_any.584 = f32[] add(%add_any.583, %neg.353)
+  %mul.1593 = f32[] multiply(%add_any.584, %constant.29)
+  %broadcast_in_dim.122 = f32[1]{0} reshape(%mul.1593)
+  %pad.116 = f32[3]{0} pad(%broadcast_in_dim.122, %constant.30), padding=0_2
+  %add_any.585 = f32[3]{0} add(%add_any.581, %pad.116)
+  %add_any.357 = f32[] add(%arg_tuple.5#4, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.255 = f32[] call(%arg_tuple.5#73, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1596 = f32[] multiply(%add_any.357, %jit_dynamic_index_in_dim_.255)
+  %jit_dynamic_index_in_dim_.253 = f32[] call(%arg_tuple.5#71, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1598 = f32[] multiply(%mul.1596, %jit_dynamic_index_in_dim_.253)
+  %jit_dynamic_index_in_dim_.252 = pred[] call(%arg_tuple.5#70, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.254 = f32[] call(%arg_tuple.5#72, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1597 = f32[] multiply(%jit_dynamic_index_in_dim_.254, %mul.1596)
+  %mul.1599 = f32[] multiply(%mul.1597, %jit_dynamic_index_in_dim_.253)
+  %mul.1600 = f32[] multiply(%mul.1599, %constant.31)
+  %neg.354 = f32[] negate(%mul.1600)
+  %select_n.235 = f32[] select(%jit_dynamic_index_in_dim_.252, %neg.354, %constant.30)
+  %add_any.587 = f32[] add(%mul.1598, %select_n.235)
+  %select_n.234 = f32[] select(%jit_dynamic_index_in_dim_.252, %constant.30, %neg.354)
+  %neg.355 = f32[] negate(%select_n.234)
+  %add_any.588 = f32[] add(%add_any.587, %neg.355)
+  %mul.1601 = f32[] multiply(%add_any.588, %constant.29)
+  %broadcast_in_dim.123 = f32[1]{0} reshape(%mul.1601)
+  %pad.117 = f32[3]{0} pad(%broadcast_in_dim.123, %constant.30), padding=2_0
+  %add_any.589 = f32[3]{0} add(%add_any.585, %pad.117)
+  %add_any.359 = f32[] add(%arg_tuple.5#3, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.251 = f32[] call(%arg_tuple.5#69, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1604 = f32[] multiply(%add_any.359, %jit_dynamic_index_in_dim_.251)
+  %jit_dynamic_index_in_dim_.249 = f32[] call(%arg_tuple.5#67, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1606 = f32[] multiply(%mul.1604, %jit_dynamic_index_in_dim_.249)
+  %jit_dynamic_index_in_dim_.248 = pred[] call(%arg_tuple.5#66, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.250 = f32[] call(%arg_tuple.5#68, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1605 = f32[] multiply(%jit_dynamic_index_in_dim_.250, %mul.1604)
+  %mul.1607 = f32[] multiply(%mul.1605, %jit_dynamic_index_in_dim_.249)
+  %mul.1608 = f32[] multiply(%mul.1607, %constant.31)
+  %neg.356 = f32[] negate(%mul.1608)
+  %select_n.237 = f32[] select(%jit_dynamic_index_in_dim_.248, %neg.356, %constant.30)
+  %add_any.591 = f32[] add(%mul.1606, %select_n.237)
+  %select_n.236 = f32[] select(%jit_dynamic_index_in_dim_.248, %constant.30, %neg.356)
+  %neg.357 = f32[] negate(%select_n.236)
+  %add_any.592 = f32[] add(%add_any.591, %neg.357)
+  %mul.1609 = f32[] multiply(%add_any.592, %constant.29)
+  %broadcast_in_dim.124 = f32[1]{0} reshape(%mul.1609)
+  %pad.118 = f32[3]{0} pad(%broadcast_in_dim.124, %constant.30), padding=1_1
+  %add_any.593 = f32[3]{0} add(%add_any.589, %pad.118)
+  %add_any.358 = f32[] add(%arg_tuple.5#2, %jit_dynamic_index_in_dim_.243)
+  %jit_dynamic_index_in_dim_.247 = f32[] call(%arg_tuple.5#65, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1612 = f32[] multiply(%add_any.358, %jit_dynamic_index_in_dim_.247)
+  %jit_dynamic_index_in_dim_.245 = f32[] call(%arg_tuple.5#63, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1614 = f32[] multiply(%mul.1612, %jit_dynamic_index_in_dim_.245)
+  %jit_dynamic_index_in_dim_.244 = pred[] call(%arg_tuple.5#62, %sub.123), to_apply=%dynamic_index_in_dim_1.6
+  %jit_dynamic_index_in_dim_.246 = f32[] call(%arg_tuple.5#64, %sub.123), to_apply=%dynamic_index_in_dim.1
+  %mul.1613 = f32[] multiply(%jit_dynamic_index_in_dim_.246, %mul.1612)
+  %mul.1615 = f32[] multiply(%mul.1613, %jit_dynamic_index_in_dim_.245)
+  %mul.1616 = f32[] multiply(%mul.1615, %constant.31)
+  %neg.358 = f32[] negate(%mul.1616)
+  %select_n.239 = f32[] select(%jit_dynamic_index_in_dim_.244, %neg.358, %constant.30)
+  %add_any.595 = f32[] add(%mul.1614, %select_n.239)
+  %select_n.238 = f32[] select(%jit_dynamic_index_in_dim_.244, %constant.30, %neg.358)
+  %neg.359 = f32[] negate(%select_n.238)
+  %add_any.596 = f32[] add(%add_any.595, %neg.359)
+  %mul.1617 = f32[] multiply(%add_any.596, %constant.29)
+  %broadcast_in_dim.125 = f32[1]{0} reshape(%mul.1617)
+  %pad.119 = f32[3]{0} pad(%broadcast_in_dim.125, %constant.30), padding=0_2
+  %add_any.597 = f32[3]{0} add(%add_any.593, %pad.119)
+  %mul.1146 = f32[] multiply(%add_any.361, %constant.31)
+  %constant.28 = f32[] constant(0.5)
+  %mul.1619 = f32[] multiply(%add_any.596, %constant.28)
+  %add_any.599 = f32[] add(%mul.1146, %mul.1619)
+  %mul.1611 = f32[] multiply(%add_any.592, %constant.28)
+  %mul.1618 = f32[] multiply(%add_any.596, %constant.31)
+  %add_any.598 = f32[] add(%mul.1611, %mul.1618)
+  %mul.1603 = f32[] multiply(%add_any.588, %constant.28)
+  %mul.1610 = f32[] multiply(%add_any.592, %constant.31)
+  %add_any.594 = f32[] add(%mul.1603, %mul.1610)
+  %mul.1595 = f32[] multiply(%add_any.584, %constant.28)
+  %mul.1602 = f32[] multiply(%add_any.588, %constant.31)
+  %add_any.590 = f32[] add(%mul.1595, %mul.1602)
+  %mul.1587 = f32[] multiply(%add_any.580, %constant.28)
+  %mul.1594 = f32[] multiply(%add_any.584, %constant.31)
+  %add_any.586 = f32[] add(%mul.1587, %mul.1594)
+  %mul.1579 = f32[] multiply(%add_any.576, %constant.28)
+  %mul.1586 = f32[] multiply(%add_any.580, %constant.31)
+  %add_any.582 = f32[] add(%mul.1579, %mul.1586)
+  %mul.1571 = f32[] multiply(%add_any.572, %constant.28)
+  %mul.1578 = f32[] multiply(%add_any.576, %constant.31)
+  %add_any.578 = f32[] add(%mul.1571, %mul.1578)
+  %mul.1563 = f32[] multiply(%add_any.568, %constant.28)
+  %mul.1570 = f32[] multiply(%add_any.572, %constant.31)
+  %add_any.574 = f32[] add(%mul.1563, %mul.1570)
+  %mul.1555 = f32[] multiply(%add_any.564, %constant.28)
+  %mul.1562 = f32[] multiply(%add_any.568, %constant.31)
+  %add_any.570 = f32[] add(%mul.1555, %mul.1562)
+  %mul.1547 = f32[] multiply(%add_any.560, %constant.28)
+  %mul.1554 = f32[] multiply(%add_any.564, %constant.31)
+  %add_any.566 = f32[] add(%mul.1547, %mul.1554)
+  %mul.1539 = f32[] multiply(%add_any.556, %constant.28)
+  %mul.1546 = f32[] multiply(%add_any.560, %constant.31)
+  %add_any.562 = f32[] add(%mul.1539, %mul.1546)
+  %mul.1531 = f32[] multiply(%add_any.552, %constant.28)
+  %mul.1538 = f32[] multiply(%add_any.556, %constant.31)
+  %add_any.558 = f32[] add(%mul.1531, %mul.1538)
+  %mul.1523 = f32[] multiply(%add_any.548, %constant.28)
+  %mul.1530 = f32[] multiply(%add_any.552, %constant.31)
+  %add_any.554 = f32[] add(%mul.1523, %mul.1530)
+  %mul.1515 = f32[] multiply(%add_any.544, %constant.28)
+  %mul.1522 = f32[] multiply(%add_any.548, %constant.31)
+  %add_any.550 = f32[] add(%mul.1515, %mul.1522)
+  %mul.1507 = f32[] multiply(%add_any.540, %constant.28)
+  %mul.1514 = f32[] multiply(%add_any.544, %constant.31)
+  %add_any.546 = f32[] add(%mul.1507, %mul.1514)
+  %mul.1499 = f32[] multiply(%add_any.536, %constant.28)
+  %mul.1506 = f32[] multiply(%add_any.540, %constant.31)
+  %add_any.542 = f32[] add(%mul.1499, %mul.1506)
+  %mul.1491 = f32[] multiply(%add_any.532, %constant.28)
+  %mul.1498 = f32[] multiply(%add_any.536, %constant.31)
+  %add_any.538 = f32[] add(%mul.1491, %mul.1498)
+  %mul.1483 = f32[] multiply(%add_any.528, %constant.28)
+  %mul.1490 = f32[] multiply(%add_any.532, %constant.31)
+  %add_any.534 = f32[] add(%mul.1483, %mul.1490)
+  %mul.1475 = f32[] multiply(%add_any.524, %constant.28)
+  %mul.1482 = f32[] multiply(%add_any.528, %constant.31)
+  %add_any.530 = f32[] add(%mul.1475, %mul.1482)
+  %mul.1467 = f32[] multiply(%add_any.520, %constant.28)
+  %mul.1474 = f32[] multiply(%add_any.524, %constant.31)
+  %add_any.526 = f32[] add(%mul.1467, %mul.1474)
+  %mul.1459 = f32[] multiply(%add_any.516, %constant.28)
+  %mul.1466 = f32[] multiply(%add_any.520, %constant.31)
+  %add_any.522 = f32[] add(%mul.1459, %mul.1466)
+  %mul.1451 = f32[] multiply(%add_any.512, %constant.28)
+  %mul.1458 = f32[] multiply(%add_any.516, %constant.31)
+  %add_any.518 = f32[] add(%mul.1451, %mul.1458)
+  %mul.1443 = f32[] multiply(%add_any.508, %constant.28)
+  %mul.1450 = f32[] multiply(%add_any.512, %constant.31)
+  %add_any.514 = f32[] add(%mul.1443, %mul.1450)
+  %mul.1435 = f32[] multiply(%add_any.504, %constant.28)
+  %mul.1442 = f32[] multiply(%add_any.508, %constant.31)
+  %add_any.510 = f32[] add(%mul.1435, %mul.1442)
+  %mul.1427 = f32[] multiply(%add_any.500, %constant.28)
+  %mul.1434 = f32[] multiply(%add_any.504, %constant.31)
+  %add_any.506 = f32[] add(%mul.1427, %mul.1434)
+  %mul.1419 = f32[] multiply(%add_any.496, %constant.28)
+  %mul.1426 = f32[] multiply(%add_any.500, %constant.31)
+  %add_any.502 = f32[] add(%mul.1419, %mul.1426)
+  %mul.1411 = f32[] multiply(%add_any.492, %constant.28)
+  %mul.1418 = f32[] multiply(%add_any.496, %constant.31)
+  %add_any.498 = f32[] add(%mul.1411, %mul.1418)
+  %mul.1403 = f32[] multiply(%add_any.488, %constant.28)
+  %mul.1410 = f32[] multiply(%add_any.492, %constant.31)
+  %add_any.494 = f32[] add(%mul.1403, %mul.1410)
+  %mul.1395 = f32[] multiply(%add_any.484, %constant.28)
+  %mul.1402 = f32[] multiply(%add_any.488, %constant.31)
+  %add_any.490 = f32[] add(%mul.1395, %mul.1402)
+  %mul.1387 = f32[] multiply(%add_any.480, %constant.28)
+  %mul.1394 = f32[] multiply(%add_any.484, %constant.31)
+  %add_any.486 = f32[] add(%mul.1387, %mul.1394)
+  %mul.1379 = f32[] multiply(%add_any.476, %constant.28)
+  %mul.1386 = f32[] multiply(%add_any.480, %constant.31)
+  %add_any.482 = f32[] add(%mul.1379, %mul.1386)
+  %mul.1371 = f32[] multiply(%add_any.472, %constant.28)
+  %mul.1378 = f32[] multiply(%add_any.476, %constant.31)
+  %add_any.478 = f32[] add(%mul.1371, %mul.1378)
+  %mul.1363 = f32[] multiply(%add_any.468, %constant.28)
+  %mul.1370 = f32[] multiply(%add_any.472, %constant.31)
+  %add_any.474 = f32[] add(%mul.1363, %mul.1370)
+  %mul.1355 = f32[] multiply(%add_any.464, %constant.28)
+  %mul.1362 = f32[] multiply(%add_any.468, %constant.31)
+  %add_any.470 = f32[] add(%mul.1355, %mul.1362)
+  %mul.1347 = f32[] multiply(%add_any.460, %constant.28)
+  %mul.1354 = f32[] multiply(%add_any.464, %constant.31)
+  %add_any.466 = f32[] add(%mul.1347, %mul.1354)
+  %mul.1339 = f32[] multiply(%add_any.456, %constant.28)
+  %mul.1346 = f32[] multiply(%add_any.460, %constant.31)
+  %add_any.462 = f32[] add(%mul.1339, %mul.1346)
+  %mul.1331 = f32[] multiply(%add_any.452, %constant.28)
+  %mul.1338 = f32[] multiply(%add_any.456, %constant.31)
+  %add_any.458 = f32[] add(%mul.1331, %mul.1338)
+  %mul.1323 = f32[] multiply(%add_any.448, %constant.28)
+  %mul.1330 = f32[] multiply(%add_any.452, %constant.31)
+  %add_any.454 = f32[] add(%mul.1323, %mul.1330)
+  %mul.1315 = f32[] multiply(%add_any.444, %constant.28)
+  %mul.1322 = f32[] multiply(%add_any.448, %constant.31)
+  %add_any.450 = f32[] add(%mul.1315, %mul.1322)
+  %mul.1307 = f32[] multiply(%add_any.440, %constant.28)
+  %mul.1314 = f32[] multiply(%add_any.444, %constant.31)
+  %add_any.446 = f32[] add(%mul.1307, %mul.1314)
+  %mul.1299 = f32[] multiply(%add_any.436, %constant.28)
+  %mul.1306 = f32[] multiply(%add_any.440, %constant.31)
+  %add_any.442 = f32[] add(%mul.1299, %mul.1306)
+  %mul.1291 = f32[] multiply(%add_any.432, %constant.28)
+  %mul.1298 = f32[] multiply(%add_any.436, %constant.31)
+  %add_any.438 = f32[] add(%mul.1291, %mul.1298)
+  %mul.1283 = f32[] multiply(%add_any.428, %constant.28)
+  %mul.1290 = f32[] multiply(%add_any.432, %constant.31)
+  %add_any.434 = f32[] add(%mul.1283, %mul.1290)
+  %mul.1275 = f32[] multiply(%add_any.424, %constant.28)
+  %mul.1282 = f32[] multiply(%add_any.428, %constant.31)
+  %add_any.430 = f32[] add(%mul.1275, %mul.1282)
+  %mul.1267 = f32[] multiply(%add_any.420, %constant.28)
+  %mul.1274 = f32[] multiply(%add_any.424, %constant.31)
+  %add_any.426 = f32[] add(%mul.1267, %mul.1274)
+  %mul.1259 = f32[] multiply(%add_any.416, %constant.28)
+  %mul.1266 = f32[] multiply(%add_any.420, %constant.31)
+  %add_any.422 = f32[] add(%mul.1259, %mul.1266)
+  %mul.1251 = f32[] multiply(%add_any.412, %constant.28)
+  %mul.1258 = f32[] multiply(%add_any.416, %constant.31)
+  %add_any.418 = f32[] add(%mul.1251, %mul.1258)
+  %mul.1243 = f32[] multiply(%add_any.408, %constant.28)
+  %mul.1250 = f32[] multiply(%add_any.412, %constant.31)
+  %add_any.414 = f32[] add(%mul.1243, %mul.1250)
+  %mul.1235 = f32[] multiply(%add_any.404, %constant.28)
+  %mul.1242 = f32[] multiply(%add_any.408, %constant.31)
+  %add_any.410 = f32[] add(%mul.1235, %mul.1242)
+  %mul.1227 = f32[] multiply(%add_any.400, %constant.28)
+  %mul.1234 = f32[] multiply(%add_any.404, %constant.31)
+  %add_any.406 = f32[] add(%mul.1227, %mul.1234)
+  %mul.1219 = f32[] multiply(%add_any.396, %constant.28)
+  %mul.1226 = f32[] multiply(%add_any.400, %constant.31)
+  %add_any.402 = f32[] add(%mul.1219, %mul.1226)
+  %mul.1211 = f32[] multiply(%add_any.392, %constant.28)
+  %mul.1218 = f32[] multiply(%add_any.396, %constant.31)
+  %add_any.398 = f32[] add(%mul.1211, %mul.1218)
+  %mul.1203 = f32[] multiply(%add_any.388, %constant.28)
+  %mul.1210 = f32[] multiply(%add_any.392, %constant.31)
+  %add_any.394 = f32[] add(%mul.1203, %mul.1210)
+  %mul.1195 = f32[] multiply(%add_any.384, %constant.28)
+  %mul.1202 = f32[] multiply(%add_any.388, %constant.31)
+  %add_any.390 = f32[] add(%mul.1195, %mul.1202)
+  %mul.1187 = f32[] multiply(%add_any.380, %constant.28)
+  %mul.1194 = f32[] multiply(%add_any.384, %constant.31)
+  %add_any.386 = f32[] add(%mul.1187, %mul.1194)
+  %mul.1179 = f32[] multiply(%add_any.376, %constant.28)
+  %mul.1186 = f32[] multiply(%add_any.380, %constant.31)
+  %add_any.382 = f32[] add(%mul.1179, %mul.1186)
+  %mul.1171 = f32[] multiply(%add_any.372, %constant.28)
+  %mul.1178 = f32[] multiply(%add_any.376, %constant.31)
+  %add_any.378 = f32[] add(%mul.1171, %mul.1178)
+  %mul.1163 = f32[] multiply(%add_any.368, %constant.28)
+  %mul.1170 = f32[] multiply(%add_any.372, %constant.31)
+  %add_any.374 = f32[] add(%mul.1163, %mul.1170)
+  %mul.1155 = f32[] multiply(%add_any.364, %constant.28)
+  %mul.1162 = f32[] multiply(%add_any.368, %constant.31)
+  %add_any.370 = f32[] add(%mul.1155, %mul.1162)
+  %mul.1147 = f32[] multiply(%add_any.361, %constant.28)
+  %mul.1154 = f32[] multiply(%add_any.364, %constant.31)
+  %add_any.366 = f32[] add(%mul.1147, %mul.1154)
+  ROOT %tuple.3 = (s32[], f32[3]{0}, f32[], f32[], f32[], /*index=5*/f32[], f32[], f32[], f32[], f32[], /*index=10*/f32[], f32[], f32[], f32[], f32[], /*index=15*/f32[], f32[], f32[], f32[], f32[], /*index=20*/f32[], f32[], f32[], f32[], f32[], /*index=25*/f32[], f32[], f32[], f32[], f32[], /*index=30*/f32[], f32[], f32[], f32[], f32[], /*index=35*/f32[], f32[], f32[], f32[], f32[], /*index=40*/f32[], f32[], f32[], f32[], f32[], /*index=45*/f32[], f32[], f32[], f32[], f32[], /*index=50*/f32[], f32[], f32[], f32[], f32[], /*index=55*/f32[], f32[], f32[], f32[], f32[], /*index=60*/f32[], f32[], pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=65*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=70*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=75*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=80*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=85*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=90*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=95*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=100*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=105*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=110*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=115*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=120*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=125*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=130*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=135*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=140*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=145*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=150*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=155*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=160*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=165*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=170*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=175*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=180*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=185*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=190*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=195*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=200*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=205*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=210*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=215*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=220*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=225*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=230*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=235*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=240*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=245*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=250*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=255*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=260*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=265*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=270*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=275*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=280*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=285*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=290*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=295*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=300*/f32[614]{0}, f32[614]{0}) tuple(%add.483, %add_any.597, %add_any.599, %add_any.598, %add_any.594, /*index=5*/%add_any.590, %add_any.586, %add_any.582, %add_any.578, %add_any.574, /*index=10*/%add_any.570, %add_any.566, %add_any.562, %add_any.558, %add_any.554, /*index=15*/%add_any.550, %add_any.546, %add_any.542, %add_any.538, %add_any.534, /*index=20*/%add_any.530, %add_any.526, %add_any.522, %add_any.518, %add_any.514, /*index=25*/%add_any.510, %add_any.506, %add_any.502, %add_any.498, %add_any.494, /*index=30*/%add_any.490, %add_any.486, %add_any.482, %add_any.478, %add_any.474, /*index=35*/%add_any.470, %add_any.466, %add_any.462, %add_any.458, %add_any.454, /*index=40*/%add_any.450, %add_any.446, %add_any.442, %add_any.438, %add_any.434, /*index=45*/%add_any.430, %add_any.426, %add_any.422, %add_any.418, %add_any.414, /*index=50*/%add_any.410, %add_any.406, %add_any.402, %add_any.398, %add_any.394, /*index=55*/%add_any.390, %add_any.386, %add_any.382, %add_any.378, %add_any.374, /*index=60*/%add_any.370, %add_any.366, %arg_tuple.5#62, %arg_tuple.5#63, %arg_tuple.5#64, /*index=65*/%arg_tuple.5#65, %arg_tuple.5#66, %arg_tuple.5#67, %arg_tuple.5#68, %arg_tuple.5#69, /*index=70*/%arg_tuple.5#70, %arg_tuple.5#71, %arg_tuple.5#72, %arg_tuple.5#73, %arg_tuple.5#74, /*index=75*/%arg_tuple.5#75, %arg_tuple.5#76, %arg_tuple.5#77, %arg_tuple.5#78, %arg_tuple.5#79, /*index=80*/%arg_tuple.5#80, %arg_tuple.5#81, %arg_tuple.5#82, %arg_tuple.5#83, %arg_tuple.5#84, /*index=85*/%arg_tuple.5#85, %arg_tuple.5#86, %arg_tuple.5#87, %arg_tuple.5#88, %arg_tuple.5#89, /*index=90*/%arg_tuple.5#90, %arg_tuple.5#91, %arg_tuple.5#92, %arg_tuple.5#93, %arg_tuple.5#94, /*index=95*/%arg_tuple.5#95, %arg_tuple.5#96, %arg_tuple.5#97, %arg_tuple.5#98, %arg_tuple.5#99, /*index=100*/%arg_tuple.5#100, %arg_tuple.5#101, %arg_tuple.5#102, %arg_tuple.5#103, %arg_tuple.5#104, /*index=105*/%arg_tuple.5#105, %arg_tuple.5#106, %arg_tuple.5#107, %arg_tuple.5#108, %arg_tuple.5#109, /*index=110*/%arg_tuple.5#110, %arg_tuple.5#111, %arg_tuple.5#112, %arg_tuple.5#113, %arg_tuple.5#114, /*index=115*/%arg_tuple.5#115, %arg_tuple.5#116, %arg_tuple.5#117, %arg_tuple.5#118, %arg_tuple.5#119, /*index=120*/%arg_tuple.5#120, %arg_tuple.5#121, %arg_tuple.5#122, %arg_tuple.5#123, %arg_tuple.5#124, /*index=125*/%arg_tuple.5#125, %arg_tuple.5#126, %arg_tuple.5#127, %arg_tuple.5#128, %arg_tuple.5#129, /*index=130*/%arg_tuple.5#130, %arg_tuple.5#131, %arg_tuple.5#132, %arg_tuple.5#133, %arg_tuple.5#134, /*index=135*/%arg_tuple.5#135, %arg_tuple.5#136, %arg_tuple.5#137, %arg_tuple.5#138, %arg_tuple.5#139, /*index=140*/%arg_tuple.5#140, %arg_tuple.5#141, %arg_tuple.5#142, %arg_tuple.5#143, %arg_tuple.5#144, /*index=145*/%arg_tuple.5#145, %arg_tuple.5#146, %arg_tuple.5#147, %arg_tuple.5#148, %arg_tuple.5#149, /*index=150*/%arg_tuple.5#150, %arg_tuple.5#151, %arg_tuple.5#152, %arg_tuple.5#153, %arg_tuple.5#154, /*index=155*/%arg_tuple.5#155, %arg_tuple.5#156, %arg_tuple.5#157, %arg_tuple.5#158, %arg_tuple.5#159, /*index=160*/%arg_tuple.5#160, %arg_tuple.5#161, %arg_tuple.5#162, %arg_tuple.5#163, %arg_tuple.5#164, /*index=165*/%arg_tuple.5#165, %arg_tuple.5#166, %arg_tuple.5#167, %arg_tuple.5#168, %arg_tuple.5#169, /*index=170*/%arg_tuple.5#170, %arg_tuple.5#171, %arg_tuple.5#172, %arg_tuple.5#173, %arg_tuple.5#174, /*index=175*/%arg_tuple.5#175, %arg_tuple.5#176, %arg_tuple.5#177, %arg_tuple.5#178, %arg_tuple.5#179, /*index=180*/%arg_tuple.5#180, %arg_tuple.5#181, %arg_tuple.5#182, %arg_tuple.5#183, %arg_tuple.5#184, /*index=185*/%arg_tuple.5#185, %arg_tuple.5#186, %arg_tuple.5#187, %arg_tuple.5#188, %arg_tuple.5#189, /*index=190*/%arg_tuple.5#190, %arg_tuple.5#191, %arg_tuple.5#192, %arg_tuple.5#193, %arg_tuple.5#194, /*index=195*/%arg_tuple.5#195, %arg_tuple.5#196, %arg_tuple.5#197, %arg_tuple.5#198, %arg_tuple.5#199, /*index=200*/%arg_tuple.5#200, %arg_tuple.5#201, %arg_tuple.5#202, %arg_tuple.5#203, %arg_tuple.5#204, /*index=205*/%arg_tuple.5#205, %arg_tuple.5#206, %arg_tuple.5#207, %arg_tuple.5#208, %arg_tuple.5#209, /*index=210*/%arg_tuple.5#210, %arg_tuple.5#211, %arg_tuple.5#212, %arg_tuple.5#213, %arg_tuple.5#214, /*index=215*/%arg_tuple.5#215, %arg_tuple.5#216, %arg_tuple.5#217, %arg_tuple.5#218, %arg_tuple.5#219, /*index=220*/%arg_tuple.5#220, %arg_tuple.5#221, %arg_tuple.5#222, %arg_tuple.5#223, %arg_tuple.5#224, /*index=225*/%arg_tuple.5#225, %arg_tuple.5#226, %arg_tuple.5#227, %arg_tuple.5#228, %arg_tuple.5#229, /*index=230*/%arg_tuple.5#230, %arg_tuple.5#231, %arg_tuple.5#232, %arg_tuple.5#233, %arg_tuple.5#234, /*index=235*/%arg_tuple.5#235, %arg_tuple.5#236, %arg_tuple.5#237, %arg_tuple.5#238, %arg_tuple.5#239, /*index=240*/%arg_tuple.5#240, %arg_tuple.5#241, %arg_tuple.5#242, %arg_tuple.5#243, %arg_tuple.5#244, /*index=245*/%arg_tuple.5#245, %arg_tuple.5#246, %arg_tuple.5#247, %arg_tuple.5#248, %arg_tuple.5#249, /*index=250*/%arg_tuple.5#250, %arg_tuple.5#251, %arg_tuple.5#252, %arg_tuple.5#253, %arg_tuple.5#254, /*index=255*/%arg_tuple.5#255, %arg_tuple.5#256, %arg_tuple.5#257, %arg_tuple.5#258, %arg_tuple.5#259, /*index=260*/%arg_tuple.5#260, %arg_tuple.5#261, %arg_tuple.5#262, %arg_tuple.5#263, %arg_tuple.5#264, /*index=265*/%arg_tuple.5#265, %arg_tuple.5#266, %arg_tuple.5#267, %arg_tuple.5#268, %arg_tuple.5#269, /*index=270*/%arg_tuple.5#270, %arg_tuple.5#271, %arg_tuple.5#272, %arg_tuple.5#273, %arg_tuple.5#274, /*index=275*/%arg_tuple.5#275, %arg_tuple.5#276, %arg_tuple.5#277, %arg_tuple.5#278, %arg_tuple.5#279, /*index=280*/%arg_tuple.5#280, %arg_tuple.5#281, %arg_tuple.5#282, %arg_tuple.5#283, %arg_tuple.5#284, /*index=285*/%arg_tuple.5#285, %arg_tuple.5#286, %arg_tuple.5#287, %arg_tuple.5#288, %arg_tuple.5#289, /*index=290*/%arg_tuple.5#290, %arg_tuple.5#291, %arg_tuple.5#292, %arg_tuple.5#293, %arg_tuple.5#294, /*index=295*/%arg_tuple.5#295, %arg_tuple.5#296, %arg_tuple.5#297, %arg_tuple.5#298, %arg_tuple.5#299, /*index=300*/%arg_tuple.5#300, %arg_tuple.5#301)
+}
+
+%region_3.8 (arg_tuple.7: (s32[], f32[3], f32[], f32[], f32[], /*index=5*/f32[], f32[], f32[], f32[], f32[], /*index=10*/f32[], f32[], f32[], f32[], f32[], /*index=15*/f32[], f32[], f32[], f32[], f32[], /*index=20*/f32[], f32[], f32[], f32[], f32[], /*index=25*/f32[], f32[], f32[], f32[], f32[], /*index=30*/f32[], f32[], f32[], f32[], f32[], /*index=35*/f32[], f32[], f32[], f32[], f32[], /*index=40*/f32[], f32[], f32[], f32[], f32[], /*index=45*/f32[], f32[], f32[], f32[], f32[], /*index=50*/f32[], f32[], f32[], f32[], f32[], /*index=55*/f32[], f32[], f32[], f32[], f32[], /*index=60*/f32[], f32[], pred[614], f32[614], f32[614], /*index=65*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=70*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=75*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=80*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=85*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=90*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=95*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=100*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=105*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=110*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=115*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=120*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=125*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=130*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=135*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=140*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=145*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=150*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=155*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=160*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=165*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=170*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=175*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=180*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=185*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=190*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=195*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=200*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=205*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=210*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=215*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=220*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=225*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=230*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=235*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=240*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=245*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=250*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=255*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=260*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=265*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=270*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=275*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=280*/f32[614], f32[614], pred[614], f32[614], f32[614], /*index=285*/f32[614], pred[614], f32[614], f32[614], f32[614], /*index=290*/pred[614], f32[614], f32[614], f32[614], pred[614], /*index=295*/f32[614], f32[614], f32[614], pred[614], f32[614], /*index=300*/f32[614], f32[614])) -> pred[] {
+  %arg_tuple.7 = (s32[], f32[3]{0}, f32[], f32[], f32[], /*index=5*/f32[], f32[], f32[], f32[], f32[], /*index=10*/f32[], f32[], f32[], f32[], f32[], /*index=15*/f32[], f32[], f32[], f32[], f32[], /*index=20*/f32[], f32[], f32[], f32[], f32[], /*index=25*/f32[], f32[], f32[], f32[], f32[], /*index=30*/f32[], f32[], f32[], f32[], f32[], /*index=35*/f32[], f32[], f32[], f32[], f32[], /*index=40*/f32[], f32[], f32[], f32[], f32[], /*index=45*/f32[], f32[], f32[], f32[], f32[], /*index=50*/f32[], f32[], f32[], f32[], f32[], /*index=55*/f32[], f32[], f32[], f32[], f32[], /*index=60*/f32[], f32[], pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=65*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=70*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=75*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=80*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=85*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=90*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=95*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=100*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=105*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=110*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=115*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=120*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=125*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=130*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=135*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=140*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=145*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=150*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=155*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=160*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=165*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=170*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=175*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=180*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=185*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=190*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=195*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=200*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=205*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=210*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=215*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=220*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=225*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=230*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=235*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=240*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=245*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=250*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=255*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=260*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=265*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=270*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=275*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=280*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=285*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=290*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=295*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=300*/f32[614]{0}, f32[614]{0}) parameter(0)
+  %constant.36 = s32[] constant(614)
+  ROOT %lt.3 = pred[] compare(%arg_tuple.7#0, %constant.36), direction=LT
+}
+
+ENTRY %main.9 (p.1: f32[3]) -> f32[3] {
+  %constant.40 = s32[] constant(0)
+  %constant.42 = f32[] constant(0.1)
+  %constant.39 = pred[] constant(false)
+  %broadcast.5 = pred[614]{0} broadcast(%constant.39), dimensions={}
+  %constant.38 = f32[] constant(0)
+  %broadcast.4 = f32[614]{0} broadcast(%constant.38), dimensions={}
+  %p.1 = f32[3]{0} parameter(0)
+  %slice.60 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.64 = f32[] reshape(%slice.60)
+  %constant.43 = f32[] constant(0.01)
+  %mul.1620 = f32[] multiply(%squeeze.64, %constant.43)
+  %slice.61 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.65 = f32[] reshape(%slice.61)
+  %mul.1621 = f32[] multiply(%squeeze.65, %constant.43)
+  %slice.62 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.66 = f32[] reshape(%slice.62)
+  %mul.1622 = f32[] multiply(%squeeze.66, %constant.43)
+  %slice.63 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.67 = f32[] reshape(%slice.63)
+  %mul.1623 = f32[] multiply(%squeeze.67, %constant.43)
+  %slice.64 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.68 = f32[] reshape(%slice.64)
+  %mul.1624 = f32[] multiply(%squeeze.68, %constant.43)
+  %slice.65 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.69 = f32[] reshape(%slice.65)
+  %mul.1625 = f32[] multiply(%squeeze.69, %constant.43)
+  %slice.66 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.70 = f32[] reshape(%slice.66)
+  %mul.1626 = f32[] multiply(%squeeze.70, %constant.43)
+  %slice.67 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.71 = f32[] reshape(%slice.67)
+  %mul.1627 = f32[] multiply(%squeeze.71, %constant.43)
+  %slice.68 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.72 = f32[] reshape(%slice.68)
+  %mul.1628 = f32[] multiply(%squeeze.72, %constant.43)
+  %slice.69 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.73 = f32[] reshape(%slice.69)
+  %mul.1629 = f32[] multiply(%squeeze.73, %constant.43)
+  %slice.70 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.74 = f32[] reshape(%slice.70)
+  %mul.1630 = f32[] multiply(%squeeze.74, %constant.43)
+  %slice.71 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.75 = f32[] reshape(%slice.71)
+  %mul.1631 = f32[] multiply(%squeeze.75, %constant.43)
+  %slice.72 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.76 = f32[] reshape(%slice.72)
+  %mul.1632 = f32[] multiply(%squeeze.76, %constant.43)
+  %slice.73 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.77 = f32[] reshape(%slice.73)
+  %mul.1633 = f32[] multiply(%squeeze.77, %constant.43)
+  %slice.74 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.78 = f32[] reshape(%slice.74)
+  %mul.1634 = f32[] multiply(%squeeze.78, %constant.43)
+  %slice.75 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.79 = f32[] reshape(%slice.75)
+  %mul.1635 = f32[] multiply(%squeeze.79, %constant.43)
+  %slice.76 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.80 = f32[] reshape(%slice.76)
+  %mul.1636 = f32[] multiply(%squeeze.80, %constant.43)
+  %slice.77 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.81 = f32[] reshape(%slice.77)
+  %mul.1637 = f32[] multiply(%squeeze.81, %constant.43)
+  %slice.78 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.82 = f32[] reshape(%slice.78)
+  %mul.1638 = f32[] multiply(%squeeze.82, %constant.43)
+  %slice.79 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.83 = f32[] reshape(%slice.79)
+  %mul.1639 = f32[] multiply(%squeeze.83, %constant.43)
+  %slice.80 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.84 = f32[] reshape(%slice.80)
+  %mul.1640 = f32[] multiply(%squeeze.84, %constant.43)
+  %slice.81 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.85 = f32[] reshape(%slice.81)
+  %mul.1641 = f32[] multiply(%squeeze.85, %constant.43)
+  %slice.82 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.86 = f32[] reshape(%slice.82)
+  %mul.1642 = f32[] multiply(%squeeze.86, %constant.43)
+  %slice.83 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.87 = f32[] reshape(%slice.83)
+  %mul.1643 = f32[] multiply(%squeeze.87, %constant.43)
+  %slice.84 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.88 = f32[] reshape(%slice.84)
+  %mul.1644 = f32[] multiply(%squeeze.88, %constant.43)
+  %slice.85 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.89 = f32[] reshape(%slice.85)
+  %mul.1645 = f32[] multiply(%squeeze.89, %constant.43)
+  %slice.86 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.90 = f32[] reshape(%slice.86)
+  %mul.1646 = f32[] multiply(%squeeze.90, %constant.43)
+  %slice.87 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.91 = f32[] reshape(%slice.87)
+  %mul.1647 = f32[] multiply(%squeeze.91, %constant.43)
+  %slice.88 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.92 = f32[] reshape(%slice.88)
+  %mul.1648 = f32[] multiply(%squeeze.92, %constant.43)
+  %slice.89 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.93 = f32[] reshape(%slice.89)
+  %mul.1649 = f32[] multiply(%squeeze.93, %constant.43)
+  %slice.90 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.94 = f32[] reshape(%slice.90)
+  %mul.1650 = f32[] multiply(%squeeze.94, %constant.43)
+  %slice.91 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.95 = f32[] reshape(%slice.91)
+  %mul.1651 = f32[] multiply(%squeeze.95, %constant.43)
+  %slice.92 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.96 = f32[] reshape(%slice.92)
+  %mul.1652 = f32[] multiply(%squeeze.96, %constant.43)
+  %slice.93 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.97 = f32[] reshape(%slice.93)
+  %mul.1653 = f32[] multiply(%squeeze.97, %constant.43)
+  %slice.94 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.98 = f32[] reshape(%slice.94)
+  %mul.1654 = f32[] multiply(%squeeze.98, %constant.43)
+  %slice.95 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.99 = f32[] reshape(%slice.95)
+  %mul.1655 = f32[] multiply(%squeeze.99, %constant.43)
+  %slice.96 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.100 = f32[] reshape(%slice.96)
+  %mul.1656 = f32[] multiply(%squeeze.100, %constant.43)
+  %slice.97 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.101 = f32[] reshape(%slice.97)
+  %mul.1657 = f32[] multiply(%squeeze.101, %constant.43)
+  %slice.98 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.102 = f32[] reshape(%slice.98)
+  %mul.1658 = f32[] multiply(%squeeze.102, %constant.43)
+  %slice.99 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.103 = f32[] reshape(%slice.99)
+  %mul.1659 = f32[] multiply(%squeeze.103, %constant.43)
+  %slice.100 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.104 = f32[] reshape(%slice.100)
+  %mul.1660 = f32[] multiply(%squeeze.104, %constant.43)
+  %slice.101 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.105 = f32[] reshape(%slice.101)
+  %mul.1661 = f32[] multiply(%squeeze.105, %constant.43)
+  %slice.102 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.106 = f32[] reshape(%slice.102)
+  %mul.1662 = f32[] multiply(%squeeze.106, %constant.43)
+  %slice.103 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.107 = f32[] reshape(%slice.103)
+  %mul.1663 = f32[] multiply(%squeeze.107, %constant.43)
+  %slice.104 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.108 = f32[] reshape(%slice.104)
+  %mul.1664 = f32[] multiply(%squeeze.108, %constant.43)
+  %slice.105 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.109 = f32[] reshape(%slice.105)
+  %mul.1665 = f32[] multiply(%squeeze.109, %constant.43)
+  %slice.106 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.110 = f32[] reshape(%slice.106)
+  %mul.1666 = f32[] multiply(%squeeze.110, %constant.43)
+  %slice.107 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.111 = f32[] reshape(%slice.107)
+  %mul.1667 = f32[] multiply(%squeeze.111, %constant.43)
+  %slice.108 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.112 = f32[] reshape(%slice.108)
+  %mul.1668 = f32[] multiply(%squeeze.112, %constant.43)
+  %slice.109 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.113 = f32[] reshape(%slice.109)
+  %mul.1669 = f32[] multiply(%squeeze.113, %constant.43)
+  %slice.110 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.114 = f32[] reshape(%slice.110)
+  %mul.1670 = f32[] multiply(%squeeze.114, %constant.43)
+  %slice.111 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.115 = f32[] reshape(%slice.111)
+  %mul.1671 = f32[] multiply(%squeeze.115, %constant.43)
+  %slice.112 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.116 = f32[] reshape(%slice.112)
+  %mul.1672 = f32[] multiply(%squeeze.116, %constant.43)
+  %slice.113 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.117 = f32[] reshape(%slice.113)
+  %mul.1673 = f32[] multiply(%squeeze.117, %constant.43)
+  %slice.114 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.118 = f32[] reshape(%slice.114)
+  %mul.1674 = f32[] multiply(%squeeze.118, %constant.43)
+  %slice.115 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.119 = f32[] reshape(%slice.115)
+  %mul.1675 = f32[] multiply(%squeeze.119, %constant.43)
+  %slice.116 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.120 = f32[] reshape(%slice.116)
+  %mul.1676 = f32[] multiply(%squeeze.120, %constant.43)
+  %slice.117 = f32[1]{0} slice(%p.1), slice={[0:1]}
+  %squeeze.121 = f32[] reshape(%slice.117)
+  %mul.1677 = f32[] multiply(%squeeze.121, %constant.43)
+  %slice.118 = f32[1]{0} slice(%p.1), slice={[1:2]}
+  %squeeze.122 = f32[] reshape(%slice.118)
+  %mul.1678 = f32[] multiply(%squeeze.122, %constant.43)
+  %slice.119 = f32[1]{0} slice(%p.1), slice={[2:3]}
+  %squeeze.123 = f32[] reshape(%slice.119)
+  %mul.1679 = f32[] multiply(%squeeze.123, %constant.43)
+  %while.367 = (s32[], f32[], f32[], f32[], f32[], /*index=5*/f32[], f32[], f32[], f32[], f32[], /*index=10*/f32[], f32[], f32[], f32[], f32[], /*index=15*/f32[], f32[], f32[], f32[], f32[], /*index=20*/f32[], f32[], f32[], f32[], f32[], /*index=25*/f32[], f32[], f32[], f32[], f32[], /*index=30*/f32[], f32[], f32[], f32[], f32[], /*index=35*/f32[], f32[], f32[], f32[], f32[], /*index=40*/f32[], f32[], f32[], f32[], f32[], /*index=45*/f32[], f32[], f32[], f32[], f32[], /*index=50*/f32[], f32[], f32[], f32[], f32[], /*index=55*/f32[], f32[], f32[], f32[], f32[], /*index=60*/f32[], pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=65*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=70*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=75*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=80*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=85*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=90*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=95*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=100*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=105*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=110*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=115*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=120*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=125*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=130*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=135*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=140*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=145*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=150*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=155*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=160*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=165*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=170*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=175*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=180*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=185*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=190*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=195*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=200*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=205*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=210*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=215*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=220*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=225*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=230*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=235*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=240*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=245*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=250*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=255*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=260*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=265*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=270*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=275*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=280*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=285*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=290*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=295*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=300*/f32[614]{0}, f32[], f32[], f32[], f32[], /*index=305*/f32[], f32[], f32[], f32[], f32[], /*index=310*/f32[], f32[], f32[], f32[], f32[], /*index=315*/f32[], f32[], f32[], f32[], f32[], /*index=320*/f32[], f32[], f32[], f32[], f32[], /*index=325*/f32[], f32[], f32[], f32[], f32[], /*index=330*/f32[], f32[], f32[], f32[], f32[], /*index=335*/f32[], f32[], f32[], f32[], f32[], /*index=340*/f32[], f32[], f32[], f32[], f32[], /*index=345*/f32[], f32[], f32[], f32[], f32[], /*index=350*/f32[], f32[], f32[], f32[], f32[], /*index=355*/f32[], f32[], f32[], f32[], f32[], /*index=360*/f32[]) tuple(%constant.40, %constant.42, %constant.42, %constant.42, %constant.42, /*index=5*/%constant.42, %constant.42, %constant.42, %constant.42, %constant.42, /*index=10*/%constant.42, %constant.42, %constant.42, %constant.42, %constant.42, /*index=15*/%constant.42, %constant.42, %constant.42, %constant.42, %constant.42, /*index=20*/%constant.42, %constant.42, %constant.42, %constant.42, %constant.42, /*index=25*/%constant.42, %constant.42, %constant.42, %constant.42, %constant.42, /*index=30*/%constant.42, %constant.42, %constant.42, %constant.42, %constant.42, /*index=35*/%constant.42, %constant.42, %constant.42, %constant.42, %constant.42, /*index=40*/%constant.42, %constant.42, %constant.42, %constant.42, %constant.42, /*index=45*/%constant.42, %constant.42, %constant.42, %constant.42, %constant.42, /*index=50*/%constant.42, %constant.42, %constant.42, %constant.42, %constant.42, /*index=55*/%constant.42, %constant.42, %constant.42, %constant.42, %constant.42, /*index=60*/%constant.42, %broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, /*index=65*/%broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, /*index=70*/%broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, /*index=75*/%broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, /*index=80*/%broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, /*index=85*/%broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, /*index=90*/%broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, /*index=95*/%broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, /*index=100*/%broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, /*index=105*/%broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, /*index=110*/%broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, /*index=115*/%broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, /*index=120*/%broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, /*index=125*/%broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, /*index=130*/%broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, /*index=135*/%broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, /*index=140*/%broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, /*index=145*/%broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, /*index=150*/%broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, /*index=155*/%broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, /*index=160*/%broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, /*index=165*/%broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, /*index=170*/%broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, /*index=175*/%broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, /*index=180*/%broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, /*index=185*/%broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, /*index=190*/%broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, /*index=195*/%broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, /*index=200*/%broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, /*index=205*/%broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, /*index=210*/%broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, /*index=215*/%broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, /*index=220*/%broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, /*index=225*/%broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, /*index=230*/%broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, /*index=235*/%broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, /*index=240*/%broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, /*index=245*/%broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, /*index=250*/%broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, /*index=255*/%broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, /*index=260*/%broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, /*index=265*/%broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, /*index=270*/%broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, /*index=275*/%broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, /*index=280*/%broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, /*index=285*/%broadcast.5, %broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, /*index=290*/%broadcast.4, %broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, /*index=295*/%broadcast.4, %broadcast.4, %broadcast.5, %broadcast.4, %broadcast.4, /*index=300*/%broadcast.4, %mul.1620, %mul.1621, %mul.1622, %mul.1623, /*index=305*/%mul.1624, %mul.1625, %mul.1626, %mul.1627, %mul.1628, /*index=310*/%mul.1629, %mul.1630, %mul.1631, %mul.1632, %mul.1633, /*index=315*/%mul.1634, %mul.1635, %mul.1636, %mul.1637, %mul.1638, /*index=320*/%mul.1639, %mul.1640, %mul.1641, %mul.1642, %mul.1643, /*index=325*/%mul.1644, %mul.1645, %mul.1646, %mul.1647, %mul.1648, /*index=330*/%mul.1649, %mul.1650, %mul.1651, %mul.1652, %mul.1653, /*index=335*/%mul.1654, %mul.1655, %mul.1656, %mul.1657, %mul.1658, /*index=340*/%mul.1659, %mul.1660, %mul.1661, %mul.1662, %mul.1663, /*index=345*/%mul.1664, %mul.1665, %mul.1666, %mul.1667, %mul.1668, /*index=350*/%mul.1669, %mul.1670, %mul.1671, %mul.1672, %mul.1673, /*index=355*/%mul.1674, %mul.1675, %mul.1676, %mul.1677, %mul.1678, /*index=360*/%mul.1679)
+  %while.368 = (s32[], f32[], f32[], f32[], f32[], /*index=5*/f32[], f32[], f32[], f32[], f32[], /*index=10*/f32[], f32[], f32[], f32[], f32[], /*index=15*/f32[], f32[], f32[], f32[], f32[], /*index=20*/f32[], f32[], f32[], f32[], f32[], /*index=25*/f32[], f32[], f32[], f32[], f32[], /*index=30*/f32[], f32[], f32[], f32[], f32[], /*index=35*/f32[], f32[], f32[], f32[], f32[], /*index=40*/f32[], f32[], f32[], f32[], f32[], /*index=45*/f32[], f32[], f32[], f32[], f32[], /*index=50*/f32[], f32[], f32[], f32[], f32[], /*index=55*/f32[], f32[], f32[], f32[], f32[], /*index=60*/f32[], pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=65*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=70*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=75*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=80*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=85*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=90*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=95*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=100*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=105*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=110*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=115*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=120*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=125*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=130*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=135*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=140*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=145*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=150*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=155*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=160*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=165*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=170*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=175*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=180*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=185*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=190*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=195*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=200*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=205*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=210*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=215*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=220*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=225*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=230*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=235*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=240*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=245*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=250*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=255*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=260*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=265*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=270*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=275*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=280*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=285*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=290*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=295*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=300*/f32[614]{0}, f32[], f32[], f32[], f32[], /*index=305*/f32[], f32[], f32[], f32[], f32[], /*index=310*/f32[], f32[], f32[], f32[], f32[], /*index=315*/f32[], f32[], f32[], f32[], f32[], /*index=320*/f32[], f32[], f32[], f32[], f32[], /*index=325*/f32[], f32[], f32[], f32[], f32[], /*index=330*/f32[], f32[], f32[], f32[], f32[], /*index=335*/f32[], f32[], f32[], f32[], f32[], /*index=340*/f32[], f32[], f32[], f32[], f32[], /*index=345*/f32[], f32[], f32[], f32[], f32[], /*index=350*/f32[], f32[], f32[], f32[], f32[], /*index=355*/f32[], f32[], f32[], f32[], f32[], /*index=360*/f32[]) while(%while.367), condition=%region_1.5, body=%region_0.4
+  %constant.37 = f32[] constant(0)
+  %broadcast.3 = f32[3]{0} broadcast(%constant.37), dimensions={}
+  %constant.41 = f32[] constant(0)
+  %while.670 = (s32[], f32[3]{0}, f32[], f32[], f32[], /*index=5*/f32[], f32[], f32[], f32[], f32[], /*index=10*/f32[], f32[], f32[], f32[], f32[], /*index=15*/f32[], f32[], f32[], f32[], f32[], /*index=20*/f32[], f32[], f32[], f32[], f32[], /*index=25*/f32[], f32[], f32[], f32[], f32[], /*index=30*/f32[], f32[], f32[], f32[], f32[], /*index=35*/f32[], f32[], f32[], f32[], f32[], /*index=40*/f32[], f32[], f32[], f32[], f32[], /*index=45*/f32[], f32[], f32[], f32[], f32[], /*index=50*/f32[], f32[], f32[], f32[], f32[], /*index=55*/f32[], f32[], f32[], f32[], f32[], /*index=60*/f32[], f32[], pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=65*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=70*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=75*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=80*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=85*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=90*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=95*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=100*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=105*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=110*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=115*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=120*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=125*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=130*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=135*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=140*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=145*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=150*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=155*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=160*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=165*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=170*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=175*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=180*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=185*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=190*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=195*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=200*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=205*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=210*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=215*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=220*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=225*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=230*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=235*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=240*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=245*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=250*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=255*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=260*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=265*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=270*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=275*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=280*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=285*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=290*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=295*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=300*/f32[614]{0}, f32[614]{0}) tuple(%constant.40, %broadcast.3, %constant.41, %constant.41, %constant.41, /*index=5*/%constant.41, %constant.41, %constant.41, %constant.41, %constant.41, /*index=10*/%constant.41, %constant.41, %constant.41, %constant.41, %constant.41, /*index=15*/%constant.41, %constant.41, %constant.41, %constant.41, %constant.41, /*index=20*/%constant.41, %constant.41, %constant.41, %constant.41, %constant.41, /*index=25*/%constant.41, %constant.41, %constant.41, %constant.41, %constant.41, /*index=30*/%constant.41, %constant.41, %constant.41, %constant.41, %constant.41, /*index=35*/%constant.41, %constant.41, %constant.41, %constant.41, %constant.41, /*index=40*/%constant.41, %constant.41, %constant.41, %constant.41, %constant.41, /*index=45*/%constant.41, %constant.41, %constant.41, %constant.41, %constant.41, /*index=50*/%constant.41, %constant.41, %constant.41, %constant.41, %constant.41, /*index=55*/%constant.41, %constant.41, %constant.41, %constant.41, %constant.41, /*index=60*/%constant.41, %constant.41, %while.368#61, %while.368#62, %while.368#63, /*index=65*/%while.368#64, %while.368#65, %while.368#66, %while.368#67, %while.368#68, /*index=70*/%while.368#69, %while.368#70, %while.368#71, %while.368#72, %while.368#73, /*index=75*/%while.368#74, %while.368#75, %while.368#76, %while.368#77, %while.368#78, /*index=80*/%while.368#79, %while.368#80, %while.368#81, %while.368#82, %while.368#83, /*index=85*/%while.368#84, %while.368#85, %while.368#86, %while.368#87, %while.368#88, /*index=90*/%while.368#89, %while.368#90, %while.368#91, %while.368#92, %while.368#93, /*index=95*/%while.368#94, %while.368#95, %while.368#96, %while.368#97, %while.368#98, /*index=100*/%while.368#99, %while.368#100, %while.368#101, %while.368#102, %while.368#103, /*index=105*/%while.368#104, %while.368#105, %while.368#106, %while.368#107, %while.368#108, /*index=110*/%while.368#109, %while.368#110, %while.368#111, %while.368#112, %while.368#113, /*index=115*/%while.368#114, %while.368#115, %while.368#116, %while.368#117, %while.368#118, /*index=120*/%while.368#119, %while.368#120, %while.368#121, %while.368#122, %while.368#123, /*index=125*/%while.368#124, %while.368#125, %while.368#126, %while.368#127, %while.368#128, /*index=130*/%while.368#129, %while.368#130, %while.368#131, %while.368#132, %while.368#133, /*index=135*/%while.368#134, %while.368#135, %while.368#136, %while.368#137, %while.368#138, /*index=140*/%while.368#139, %while.368#140, %while.368#141, %while.368#142, %while.368#143, /*index=145*/%while.368#144, %while.368#145, %while.368#146, %while.368#147, %while.368#148, /*index=150*/%while.368#149, %while.368#150, %while.368#151, %while.368#152, %while.368#153, /*index=155*/%while.368#154, %while.368#155, %while.368#156, %while.368#157, %while.368#158, /*index=160*/%while.368#159, %while.368#160, %while.368#161, %while.368#162, %while.368#163, /*index=165*/%while.368#164, %while.368#165, %while.368#166, %while.368#167, %while.368#168, /*index=170*/%while.368#169, %while.368#170, %while.368#171, %while.368#172, %while.368#173, /*index=175*/%while.368#174, %while.368#175, %while.368#176, %while.368#177, %while.368#178, /*index=180*/%while.368#179, %while.368#180, %while.368#181, %while.368#182, %while.368#183, /*index=185*/%while.368#184, %while.368#185, %while.368#186, %while.368#187, %while.368#188, /*index=190*/%while.368#189, %while.368#190, %while.368#191, %while.368#192, %while.368#193, /*index=195*/%while.368#194, %while.368#195, %while.368#196, %while.368#197, %while.368#198, /*index=200*/%while.368#199, %while.368#200, %while.368#201, %while.368#202, %while.368#203, /*index=205*/%while.368#204, %while.368#205, %while.368#206, %while.368#207, %while.368#208, /*index=210*/%while.368#209, %while.368#210, %while.368#211, %while.368#212, %while.368#213, /*index=215*/%while.368#214, %while.368#215, %while.368#216, %while.368#217, %while.368#218, /*index=220*/%while.368#219, %while.368#220, %while.368#221, %while.368#222, %while.368#223, /*index=225*/%while.368#224, %while.368#225, %while.368#226, %while.368#227, %while.368#228, /*index=230*/%while.368#229, %while.368#230, %while.368#231, %while.368#232, %while.368#233, /*index=235*/%while.368#234, %while.368#235, %while.368#236, %while.368#237, %while.368#238, /*index=240*/%while.368#239, %while.368#240, %while.368#241, %while.368#242, %while.368#243, /*index=245*/%while.368#244, %while.368#245, %while.368#246, %while.368#247, %while.368#248, /*index=250*/%while.368#249, %while.368#250, %while.368#251, %while.368#252, %while.368#253, /*index=255*/%while.368#254, %while.368#255, %while.368#256, %while.368#257, %while.368#258, /*index=260*/%while.368#259, %while.368#260, %while.368#261, %while.368#262, %while.368#263, /*index=265*/%while.368#264, %while.368#265, %while.368#266, %while.368#267, %while.368#268, /*index=270*/%while.368#269, %while.368#270, %while.368#271, %while.368#272, %while.368#273, /*index=275*/%while.368#274, %while.368#275, %while.368#276, %while.368#277, %while.368#278, /*index=280*/%while.368#279, %while.368#280, %while.368#281, %while.368#282, %while.368#283, /*index=285*/%while.368#284, %while.368#285, %while.368#286, %while.368#287, %while.368#288, /*index=290*/%while.368#289, %while.368#290, %while.368#291, %while.368#292, %while.368#293, /*index=295*/%while.368#294, %while.368#295, %while.368#296, %while.368#297, %while.368#298, /*index=300*/%while.368#299, %while.368#300)
+  %while.671 = (s32[], f32[3]{0}, f32[], f32[], f32[], /*index=5*/f32[], f32[], f32[], f32[], f32[], /*index=10*/f32[], f32[], f32[], f32[], f32[], /*index=15*/f32[], f32[], f32[], f32[], f32[], /*index=20*/f32[], f32[], f32[], f32[], f32[], /*index=25*/f32[], f32[], f32[], f32[], f32[], /*index=30*/f32[], f32[], f32[], f32[], f32[], /*index=35*/f32[], f32[], f32[], f32[], f32[], /*index=40*/f32[], f32[], f32[], f32[], f32[], /*index=45*/f32[], f32[], f32[], f32[], f32[], /*index=50*/f32[], f32[], f32[], f32[], f32[], /*index=55*/f32[], f32[], f32[], f32[], f32[], /*index=60*/f32[], f32[], pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=65*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=70*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=75*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=80*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=85*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=90*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=95*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=100*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=105*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=110*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=115*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=120*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=125*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=130*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=135*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=140*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=145*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=150*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=155*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=160*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=165*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=170*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=175*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=180*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=185*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=190*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=195*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=200*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=205*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=210*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=215*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=220*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=225*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=230*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=235*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=240*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=245*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=250*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=255*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=260*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=265*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=270*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=275*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=280*/f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, /*index=285*/f32[614]{0}, pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, /*index=290*/pred[614]{0}, f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, /*index=295*/f32[614]{0}, f32[614]{0}, f32[614]{0}, pred[614]{0}, f32[614]{0}, /*index=300*/f32[614]{0}, f32[614]{0}) while(%while.670), condition=%region_3.8, body=%region_2.7
+  ROOT %while.673 = f32[3]{0} get-tuple-element(%while.671), index=1
+}
+
diff --git a/third_party/xla/xla/backends/cpu/benchmarks/hlo/xla.issue.50040.value_and_grad.hlo b/third_party/xla/xla/backends/cpu/benchmarks/hlo/xla.issue.50040.value_and_grad.hlo
new file mode 100644
index 00000000000..377ba42f63f
--- /dev/null
+++ b/third_party/xla/xla/backends/cpu/benchmarks/hlo/xla.issue.50040.value_and_grad.hlo
@@ -0,0 +1,193 @@
+// Reproducer from https://github.com/openxla/xla/issues/50040:
+// jax.value_and_grad through lax.scan with an f64[3,3] carry (2794 steps).
+HloModule jit_objective, entry_computation_layout={(f64[3]{0})->(f64[], f64[3]{0})}, allow_spmd_sharding_propagation_to_parameters={true}, allow_spmd_sharding_propagation_to_output={true,true}
+
+%region_1.1 (reduce_sum.3: f64[], reduce_sum.4: f64[]) -> f64[] {
+  %reduce_sum.3 = f64[] parameter(0)
+  %reduce_sum.4 = f64[] parameter(1)
+  ROOT %reduce_sum.5 = f64[] add(%reduce_sum.3, %reduce_sum.4)
+}
+
+%dynamic_update_index_in_dim.2 (Arg_0.1: f64[2794,3,3], Arg_1.1: f64[3,3], Arg_2.1: s64[]) -> f64[2794,3,3] {
+  %Arg_0.1 = f64[2794,3,3]{2,1,0} parameter(0)
+  %Arg_1.1 = f64[3,3]{1,0} parameter(1)
+  %broadcast_in_dim.3 = f64[1,3,3]{2,1,0} reshape(%Arg_1.1)
+  %Arg_2.1 = s64[] parameter(2)
+  %constant.12 = s64[] constant(0)
+  ROOT %dynamic_update_slice.1 = f64[2794,3,3]{2,1,0} dynamic-update-slice(%Arg_0.1, %broadcast_in_dim.3, %Arg_2.1, %constant.12, %constant.12)
+}
+
+%dynamic_update_index_in_dim_0.3 (Arg_0.3: f64[2794,3,1], Arg_1.3: f64[3,1], Arg_2.3: s64[]) -> f64[2794,3,1] {
+  %Arg_0.3 = f64[2794,3,1]{2,1,0} parameter(0)
+  %Arg_1.3 = f64[3,1]{1,0} parameter(1)
+  %broadcast_in_dim.5 = f64[1,3,1]{2,1,0} reshape(%Arg_1.3)
+  %Arg_2.3 = s64[] parameter(2)
+  %constant.14 = s64[] constant(0)
+  ROOT %dynamic_update_slice.3 = f64[2794,3,1]{2,1,0} dynamic-update-slice(%Arg_0.3, %broadcast_in_dim.5, %Arg_2.3, %constant.14, %constant.14)
+}
+
+%region_0.4 (arg_tuple.1: (s64[], f64[3,3], f64[2794,3,3], f64[2794,3,3], f64[2794,3,1], /*index=5*/f64[2794,3,3], f64[2794,3,1], f64[1,3])) -> (s64[], f64[3,3], f64[2794,3,3], f64[2794,3,3], f64[2794,3,1], /*index=5*/f64[2794,3,3], f64[2794,3,1], f64[1,3]) {
+  %arg_tuple.1 = (s64[], f64[3,3]{1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,1]{2,1,0}, /*index=5*/f64[2794,3,3]{2,1,0}, f64[2794,3,1]{2,1,0}, f64[1,3]{1,0}) parameter(0)
+  %constant.15 = s64[] constant(1)
+  %add.13 = s64[] add(%arg_tuple.1#0, %constant.15)
+  %constant.17 = f64[] constant(0.99)
+  %broadcast.8 = f64[3,3]{1,0} broadcast(%constant.17), dimensions={}
+  %mul.5 = f64[3,3]{1,0} multiply(%arg_tuple.1#1, %broadcast.8)
+  %add.8 = f64[1,3]{1,0} broadcast(%arg_tuple.1#7), dimensions={0,1}
+  %add.9 = f64[3]{0} reshape(%add.8)
+  %add.10 = f64[3,3]{1,0} broadcast(%add.9), dimensions={1}
+  %add.11 = f64[3,3]{1,0} add(%arg_tuple.1#1, %add.10)
+  %sin.1 = f64[3,3]{1,0} sine(%add.11)
+  %constant.16 = f64[] constant(0.01)
+  %broadcast.7 = f64[3,3]{1,0} broadcast(%constant.16), dimensions={}
+  %mul.6 = f64[3,3]{1,0} multiply(%sin.1, %broadcast.7)
+  %mul.4 = f64[3,3]{1,0} multiply(%arg_tuple.1#1, %arg_tuple.1#1)
+  %constant.19 = f64[] constant(0)
+  %reduce_sum.7 = f64[3]{0} reduce(%mul.4, %constant.19), dimensions={1}, to_apply=%region_1.1
+  %broadcast_in_dim.6 = f64[3,1]{1,0} reshape(%reduce_sum.7)
+  %constant.18 = f64[] constant(1)
+  %broadcast.9 = f64[3,1]{1,0} broadcast(%constant.18), dimensions={}
+  %add.7 = f64[3,1]{1,0} add(%broadcast_in_dim.6, %broadcast.9)
+  %div.5 = f64[3,1]{1,0} broadcast(%add.7), dimensions={0,1}
+  %div.6 = f64[3]{0} reshape(%div.5)
+  %div.7 = f64[3,3]{1,0} broadcast(%div.6), dimensions={0}
+  %div.8 = f64[3,3]{1,0} divide(%mul.6, %div.7)
+  %add.12 = f64[3,3]{1,0} add(%mul.5, %div.8)
+  %jit_dynamic_update_index_in_dim_.5 = f64[2794,3,3]{2,1,0} call(%arg_tuple.1#2, %arg_tuple.1#1, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %cos.1 = f64[3,3]{1,0} cosine(%add.11)
+  %jit_dynamic_update_index_in_dim_.6 = f64[2794,3,3]{2,1,0} call(%arg_tuple.1#3, %cos.1, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %jit_dynamic_update_index_in_dim_.7 = f64[2794,3,1]{2,1,0} call(%arg_tuple.1#4, %add.7, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  %jit_dynamic_update_index_in_dim_.8 = f64[2794,3,3]{2,1,0} call(%arg_tuple.1#5, %mul.6, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim.2
+  %mul.7 = f64[3,1]{1,0} multiply(%add.7, %add.7)
+  %div.9 = f64[3,1]{1,0} divide(%broadcast.9, %mul.7)
+  %jit_dynamic_update_index_in_dim_.9 = f64[2794,3,1]{2,1,0} call(%arg_tuple.1#6, %div.9, %arg_tuple.1#0), to_apply=%dynamic_update_index_in_dim_0.3
+  ROOT %tuple.1 = (s64[], f64[3,3]{1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,1]{2,1,0}, /*index=5*/f64[2794,3,3]{2,1,0}, f64[2794,3,1]{2,1,0}, f64[1,3]{1,0}) tuple(%add.13, %add.12, %jit_dynamic_update_index_in_dim_.5, %jit_dynamic_update_index_in_dim_.6, %jit_dynamic_update_index_in_dim_.7, /*index=5*/%jit_dynamic_update_index_in_dim_.8, %jit_dynamic_update_index_in_dim_.9, %arg_tuple.1#7)
+}
+
+%region_2.5 (arg_tuple.3: (s64[], f64[3,3], f64[2794,3,3], f64[2794,3,3], f64[2794,3,1], /*index=5*/f64[2794,3,3], f64[2794,3,1], f64[1,3])) -> pred[] {
+  %arg_tuple.3 = (s64[], f64[3,3]{1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,1]{2,1,0}, /*index=5*/f64[2794,3,3]{2,1,0}, f64[2794,3,1]{2,1,0}, f64[1,3]{1,0}) parameter(0)
+  %constant.21 = s64[] constant(2794)
+  ROOT %lt.1 = pred[] compare(%arg_tuple.3#0, %constant.21), direction=LT
+}
+
+%region_3.6 (reduce_sum.11: f64[], reduce_sum.12: f64[]) -> f64[] {
+  %reduce_sum.11 = f64[] parameter(0)
+  %reduce_sum.12 = f64[] parameter(1)
+  ROOT %reduce_sum.13 = f64[] add(%reduce_sum.11, %reduce_sum.12)
+}
+
+%dynamic_index_in_dim.7 (Arg_0.5: f64[2794,3,3], Arg_1.5: s64[]) -> f64[3,3] {
+  %Arg_0.5 = f64[2794,3,3]{2,1,0} parameter(0)
+  %Arg_1.5 = s64[] parameter(1)
+  %constant.28 = s64[] constant(0)
+  %dynamic_slice.1 = f64[1,3,3]{2,1,0} dynamic-slice(%Arg_0.5, %Arg_1.5, %constant.28, %constant.28), dynamic_slice_sizes={1,3,3}
+  ROOT %squeeze.1 = f64[3,3]{1,0} reshape(%dynamic_slice.1)
+}
+
+%dynamic_index_in_dim_1.8 (Arg_0.7: f64[2794,3,1], Arg_1.7: s64[]) -> f64[3,1] {
+  %Arg_0.7 = f64[2794,3,1]{2,1,0} parameter(0)
+  %Arg_1.7 = s64[] parameter(1)
+  %constant.30 = s64[] constant(0)
+  %dynamic_slice.3 = f64[1,3,1]{2,1,0} dynamic-slice(%Arg_0.7, %Arg_1.7, %constant.30, %constant.30), dynamic_slice_sizes={1,3,1}
+  ROOT %squeeze.3 = f64[3,1]{1,0} reshape(%dynamic_slice.3)
+}
+
+%region_5.9 (reduce_sum.18: f64[], reduce_sum.19: f64[]) -> f64[] {
+  %reduce_sum.18 = f64[] parameter(0)
+  %reduce_sum.19 = f64[] parameter(1)
+  ROOT %reduce_sum.20 = f64[] add(%reduce_sum.18, %reduce_sum.19)
+}
+
+%region_6.10 (reduce_sum.25: f64[], reduce_sum.26: f64[]) -> f64[] {
+  %reduce_sum.25 = f64[] parameter(0)
+  %reduce_sum.26 = f64[] parameter(1)
+  ROOT %reduce_sum.27 = f64[] add(%reduce_sum.25, %reduce_sum.26)
+}
+
+%region_7.11 (reduce_sum.32: f64[], reduce_sum.33: f64[]) -> f64[] {
+  %reduce_sum.32 = f64[] parameter(0)
+  %reduce_sum.33 = f64[] parameter(1)
+  ROOT %reduce_sum.34 = f64[] add(%reduce_sum.32, %reduce_sum.33)
+}
+
+%region_8.12 (reduce_sum.39: f64[], reduce_sum.40: f64[]) -> f64[] {
+  %reduce_sum.39 = f64[] parameter(0)
+  %reduce_sum.40 = f64[] parameter(1)
+  ROOT %reduce_sum.41 = f64[] add(%reduce_sum.39, %reduce_sum.40)
+}
+
+%region_4.13 (arg_tuple.5: (s64[], f64[3], f64[3,3], f64[2794,3,3], f64[2794,3,3], /*index=5*/f64[2794,3,1], f64[2794,3,3], f64[2794,3,1])) -> (s64[], f64[3], f64[3,3], f64[2794,3,3], f64[2794,3,3], /*index=5*/f64[2794,3,1], f64[2794,3,3], f64[2794,3,1]) {
+  %arg_tuple.5 = (s64[], f64[3]{0}, f64[3,3]{1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,3]{2,1,0}, /*index=5*/f64[2794,3,1]{2,1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,1]{2,1,0}) parameter(0)
+  %constant.34 = s64[] constant(1)
+  %add.15 = s64[] add(%arg_tuple.5#0, %constant.34)
+  %constant.35 = s64[] constant(2794)
+  %sub.2 = s64[] subtract(%constant.35, %arg_tuple.5#0)
+  %sub.3 = s64[] subtract(%sub.2, %constant.34)
+  %jit_dynamic_index_in_dim_.7 = f64[3,1]{1,0} call(%arg_tuple.5#5, %sub.3), to_apply=%dynamic_index_in_dim_1.8
+  %div.14 = f64[3,1]{1,0} broadcast(%jit_dynamic_index_in_dim_.7), dimensions={0,1}
+  %div.15 = f64[3]{0} reshape(%div.14)
+  %div.16 = f64[3,3]{1,0} broadcast(%div.15), dimensions={0}
+  %div.17 = f64[3,3]{1,0} divide(%arg_tuple.5#2, %div.16)
+  %constant.32 = f64[] constant(0.01)
+  %broadcast.13 = f64[3,3]{1,0} broadcast(%constant.32), dimensions={}
+  %mul.24 = f64[3,3]{1,0} multiply(%div.17, %broadcast.13)
+  %jit_dynamic_index_in_dim_.6 = f64[3,3]{1,0} call(%arg_tuple.5#4, %sub.3), to_apply=%dynamic_index_in_dim.7
+  %mul.25 = f64[3,3]{1,0} multiply(%mul.24, %jit_dynamic_index_in_dim_.6)
+  %constant.33 = f64[] constant(0)
+  %reduce_sum.44 = f64[3]{0} reduce(%mul.25, %constant.33), dimensions={0}, to_apply=%region_6.10
+  %reshape.3 = f64[1,3]{1,0} reshape(%reduce_sum.44)
+  %reduce_sum.45 = f64[3]{0} reduce(%reshape.3, %constant.33), dimensions={0}, to_apply=%region_7.11
+  %add_any.5 = f64[3]{0} add(%arg_tuple.5#1, %reduce_sum.45)
+  %constant.31 = f64[] constant(0.99)
+  %broadcast.12 = f64[3,3]{1,0} broadcast(%constant.31), dimensions={}
+  %mul.26 = f64[3,3]{1,0} multiply(%arg_tuple.5#2, %broadcast.12)
+  %add_any.6 = f64[3,3]{1,0} add(%mul.25, %mul.26)
+  %jit_dynamic_index_in_dim_.5 = f64[3,3]{1,0} call(%arg_tuple.5#3, %sub.3), to_apply=%dynamic_index_in_dim.7
+  %jit_dynamic_index_in_dim_.9 = f64[3,1]{1,0} call(%arg_tuple.5#7, %sub.3), to_apply=%dynamic_index_in_dim_1.8
+  %mul.19 = f64[3,1]{1,0} broadcast(%jit_dynamic_index_in_dim_.9), dimensions={0,1}
+  %mul.20 = f64[3]{0} reshape(%mul.19)
+  %mul.21 = f64[3,3]{1,0} broadcast(%mul.20), dimensions={0}
+  %mul.22 = f64[3,3]{1,0} multiply(%arg_tuple.5#2, %mul.21)
+  %jit_dynamic_index_in_dim_.8 = f64[3,3]{1,0} call(%arg_tuple.5#6, %sub.3), to_apply=%dynamic_index_in_dim.7
+  %mul.23 = f64[3,3]{1,0} multiply(%mul.22, %jit_dynamic_index_in_dim_.8)
+  %reduce_sum.43 = f64[3]{0} reduce(%mul.23, %constant.33), dimensions={1}, to_apply=%region_5.9
+  %reshape.2 = f64[3,1]{1,0} reshape(%reduce_sum.43)
+  %neg.1 = f64[3,1]{1,0} negate(%reshape.2)
+  %reduce_sum.46 = f64[3]{0} reduce(%neg.1, %constant.33), dimensions={1}, to_apply=%region_8.12
+  %broadcast_in_dim.8 = f64[3,3]{1,0} broadcast(%reduce_sum.46), dimensions={0}
+  %mul.27 = f64[3,3]{1,0} multiply(%jit_dynamic_index_in_dim_.5, %broadcast_in_dim.8)
+  %add_any.7 = f64[3,3]{1,0} add(%add_any.6, %mul.27)
+  %mul.28 = f64[3,3]{1,0} multiply(%broadcast_in_dim.8, %jit_dynamic_index_in_dim_.5)
+  %add_any.8 = f64[3,3]{1,0} add(%add_any.7, %mul.28)
+  ROOT %tuple.3 = (s64[], f64[3]{0}, f64[3,3]{1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,3]{2,1,0}, /*index=5*/f64[2794,3,1]{2,1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,1]{2,1,0}) tuple(%add.15, %add_any.5, %add_any.8, %arg_tuple.5#3, %arg_tuple.5#4, /*index=5*/%arg_tuple.5#5, %arg_tuple.5#6, %arg_tuple.5#7)
+}
+
+%region_9.14 (arg_tuple.7: (s64[], f64[3], f64[3,3], f64[2794,3,3], f64[2794,3,3], /*index=5*/f64[2794,3,1], f64[2794,3,3], f64[2794,3,1])) -> pred[] {
+  %arg_tuple.7 = (s64[], f64[3]{0}, f64[3,3]{1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,3]{2,1,0}, /*index=5*/f64[2794,3,1]{2,1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,1]{2,1,0}) parameter(0)
+  %constant.37 = s64[] constant(2794)
+  ROOT %lt.3 = pred[] compare(%arg_tuple.7#0, %constant.37), direction=LT
+}
+
+ENTRY %main.15 (p.1: f64[3]) -> (f64[], f64[3]) {
+  %constant.41 = s64[] constant(0)
+  %constant.43 = f64[] constant(0.1)
+  %broadcast.17 = f64[3,3]{1,0} broadcast(%constant.43), dimensions={}
+  %constant.40 = f64[] constant(0)
+  %broadcast.16 = f64[2794,3,3]{2,1,0} broadcast(%constant.40), dimensions={}
+  %constant.39 = f64[] constant(0)
+  %broadcast.15 = f64[2794,3,1]{2,1,0} broadcast(%constant.39), dimensions={}
+  %p.1 = f64[3]{0} parameter(0)
+  %broadcast_in_dim.9 = f64[1,3]{1,0} reshape(%p.1)
+  %while.14 = (s64[], f64[3,3]{1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,1]{2,1,0}, /*index=5*/f64[2794,3,3]{2,1,0}, f64[2794,3,1]{2,1,0}, f64[1,3]{1,0}) tuple(%constant.41, %broadcast.17, %broadcast.16, %broadcast.16, %broadcast.15, /*index=5*/%broadcast.16, %broadcast.15, %broadcast_in_dim.9)
+  %while.15 = (s64[], f64[3,3]{1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,1]{2,1,0}, /*index=5*/f64[2794,3,3]{2,1,0}, f64[2794,3,1]{2,1,0}, f64[1,3]{1,0}) while(%while.14), condition=%region_2.5, body=%region_0.4
+  %constant.38 = f64[] constant(0)
+  %broadcast.14 = f64[3]{0} broadcast(%constant.38), dimensions={}
+  %add_any.9 = f64[3,3]{1,0} add(%while.15#1, %while.15#1)
+  %while.23 = (s64[], f64[3]{0}, f64[3,3]{1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,3]{2,1,0}, /*index=5*/f64[2794,3,1]{2,1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,1]{2,1,0}) tuple(%constant.41, %broadcast.14, %add_any.9, %while.15#2, %while.15#3, /*index=5*/%while.15#4, %while.15#5, %while.15#6)
+  %while.24 = (s64[], f64[3]{0}, f64[3,3]{1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,3]{2,1,0}, /*index=5*/f64[2794,3,1]{2,1,0}, f64[2794,3,3]{2,1,0}, f64[2794,3,1]{2,1,0}) while(%while.23), condition=%region_9.14, body=%region_4.13
+  %mul.29 = f64[3,3]{1,0} multiply(%while.15#1, %while.15#1)
+  %constant.42 = f64[] constant(0)
+  %reduce_sum.47 = f64[] reduce(%mul.29, %constant.42), dimensions={0,1}, to_apply=%region_3.6
+  ROOT %tuple.5 = (f64[], f64[3]{0}) tuple(%reduce_sum.47, %while.24#1)
+}
+

From a6d5fe5f93607e90ab6b066952268c3fb473faa9 Mon Sep 17 00:00:00 2001
From: Alejandro Pacheco <alepacheco99@gmail.com>
Date: Tue, 13 Oct 2026 02:20:42 +0200
Subject: [PATCH 0609/1000] Bound the copying pass of FastParseSequenceExample
 by the output sizes

FastParseSequenceExample counts the values of every requested feature
first (GetContextFeatureLengths / GetSequenceFeatureLengths), allocates
the outputs from those counts, and then copies the values into them in
Parse{Context,Sequence}{Dense,Sparse,Ragged}Features through raw
pointers. None of those writes was checked against the size of the
output; they relied on both passes reading the same bytes the same way.
Recent out-of-bounds writes in this file (16584652da, 65993abce) were
cases where they did not.

Make the copying pass unable to write past an output regardless of
whether the passes agree:

- Parse{Bytes,Float,Int64}Feature take the number of free slots in
  `out` and return -1 instead of writing past them.
- ParseFeature passes the room left in the output tensor.
- ParseSequenceDenseFeatures and ParseSequenceSparseFeatures go through
  ParseFeature instead of their own raw cursors, and dense padding goes
  through a bounded PadFeature helper.
- The ragged inner splits check the first pass's row count before each
  write.
- CopyTensorIntoTensor checks that a context default fits.

Output for valid input is unchanged; add a test covering every output
kind of FastParseSequenceExample.
---
 .../core/util/example_proto_fast_parsing.cc   | 224 +++++++++---------
 .../util/example_proto_fast_parsing_test.cc   |  82 +++++++
 2 files changed, 193 insertions(+), 113 deletions(-)

diff --git a/tensorflow/core/util/example_proto_fast_parsing.cc b/tensorflow/core/util/example_proto_fast_parsing.cc
index 40475158509..8b6d458c759 100644
--- a/tensorflow/core/util/example_proto_fast_parsing.cc
+++ b/tensorflow/core/util/example_proto_fast_parsing.cc
@@ -1884,10 +1884,17 @@ std::string ExampleName(const absl::Span<const tstring> example_names, int n) {
   return example_names.empty() ? "<unknown>" : example_names[n];
 }
 
+// The Parse*Feature functions below are used twice by FastParseSequenceExample:
+// first with a null `out` to count the values of a feature, and then, once the
+// output tensors have been allocated from those counts, to copy the values.
+// In the copying pass at most `out_size` values are written to `out`, and a
+// feature holding more values than that is reported as an error, so that a
+// disagreement between the two passes can't write past the end of an output.
+
 // Return the number of bytes elements parsed, or -1 on error. If out is null,
 // this method simply counts the number of elements without any copying.
 inline int ParseBytesFeature(protobuf::io::CodedInputStream* stream,
-                             tstring* out) {
+                             tstring* out, size_t out_size) {
   int num_elements = 0;
   uint32_t length;
   if (!stream->ExpectTag(kDelimitedTag(1)) || !stream->ReadVarint32(&length)) {
@@ -1906,6 +1913,9 @@ inline int ParseBytesFeature(protobuf::io::CodedInputStream* stream,
           return -1;
         }
       } else {
+        if (static_cast<size_t>(num_elements) >= out_size) {
+          return -1;
+        }
         if (static_cast<int64_t>(bytes_length) > stream->BytesUntilLimit()) {
           return -1;
         }
@@ -1922,22 +1932,10 @@ inline int ParseBytesFeature(protobuf::io::CodedInputStream* stream,
   return num_elements;
 }
 
-inline void PadFloatFeature(int num_to_pad, float* out) {
-  for (int i = 0; i < num_to_pad; i++) {
-    *out++ = 0.0;
-  }
-}
-
-inline void PadInt64Feature(int num_to_pad, int64_t* out) {
-  for (int i = 0; i < num_to_pad; i++) {
-    *out++ = 0;
-  }
-}
-
 // Return the number of float elements parsed, or -1 on error. If out is null,
 // this method simply counts the number of elements without any copying.
-inline int ParseFloatFeature(protobuf::io::CodedInputStream* stream,
-                             float* out) {
+inline int ParseFloatFeature(protobuf::io::CodedInputStream* stream, float* out,
+                             size_t out_size) {
   int num_elements = 0;
   uint32_t length;
   if (!stream->ExpectTag(kDelimitedTag(2)) || !stream->ReadVarint32(&length)) {
@@ -1959,6 +1957,9 @@ inline int ParseFloatFeature(protobuf::io::CodedInputStream* stream,
           return -1;
         }
         if (out != nullptr) {
+          if (static_cast<size_t>(num_elements) >= out_size) {
+            return -1;
+          }
           *out++ = absl::bit_cast<float>(buffer32);
         }
         num_elements++;
@@ -1972,6 +1973,9 @@ inline int ParseFloatFeature(protobuf::io::CodedInputStream* stream,
           return -1;
         }
         if (out != nullptr) {
+          if (static_cast<size_t>(num_elements) >= out_size) {
+            return -1;
+          }
           *out++ = absl::bit_cast<float>(buffer32);
         }
         num_elements++;
@@ -1988,7 +1992,7 @@ inline int ParseFloatFeature(protobuf::io::CodedInputStream* stream,
 // Return the number of int64 elements parsed, or -1 on error. If out is null,
 // this method simply counts the number of elements without any copying.
 inline int ParseInt64Feature(protobuf::io::CodedInputStream* stream,
-                             int64_t* out) {
+                             int64_t* out, size_t out_size) {
   int num_elements = 0;
   uint32_t length;
   if (!stream->ExpectTag(kDelimitedTag(3)) || !stream->ReadVarint32(&length)) {
@@ -2010,6 +2014,9 @@ inline int ParseInt64Feature(protobuf::io::CodedInputStream* stream,
           return -1;
         }
         if (out != nullptr) {
+          if (static_cast<size_t>(num_elements) >= out_size) {
+            return -1;
+          }
           *out++ = n;
         }
         num_elements++;
@@ -2022,6 +2029,9 @@ inline int ParseInt64Feature(protobuf::io::CodedInputStream* stream,
           return -1;
         }
         if (out != nullptr) {
+          if (static_cast<size_t>(num_elements) >= out_size) {
+            return -1;
+          }
           *out++ = n;
         }
         num_elements++;
@@ -2037,22 +2047,28 @@ inline int ParseInt64Feature(protobuf::io::CodedInputStream* stream,
 
 // Parses the next feature on `stream` into `out` starting at `out_offset`.
 // Updates `out_offset`, and returns the number of values added.
-// Returns -1 if the next feature on `stream` doesn't match `dtype`.
+// Returns -1 if the next feature on `stream` doesn't match `dtype`, or if it
+// has more values than `out` has room for after `out_offset`.
 inline int ParseFeature(DataType dtype, protobuf::io::CodedInputStream* stream,
                         Tensor* out, size_t* out_offset) {
+  const size_t out_elements = out->NumElements();
+  if (*out_offset > out_elements) {
+    return -1;
+  }
+  const size_t out_size = out_elements - *out_offset;
   int delta;
   switch (dtype) {
     case DT_STRING:
-      delta =
-          ParseBytesFeature(stream, out->flat<tstring>().data() + *out_offset);
+      delta = ParseBytesFeature(
+          stream, out->flat<tstring>().data() + *out_offset, out_size);
       break;
     case DT_FLOAT:
-      delta =
-          ParseFloatFeature(stream, out->flat<float>().data() + *out_offset);
+      delta = ParseFloatFeature(stream, out->flat<float>().data() + *out_offset,
+                                out_size);
       break;
     case DT_INT64:
-      delta =
-          ParseInt64Feature(stream, out->flat<int64_t>().data() + *out_offset);
+      delta = ParseInt64Feature(
+          stream, out->flat<int64_t>().data() + *out_offset, out_size);
       break;
     default:
       ReportUnexpectedDataType(dtype);
@@ -2064,17 +2080,47 @@ inline int ParseFeature(DataType dtype, protobuf::io::CodedInputStream* stream,
   return delta;
 }
 
+// Fills `num_to_pad` values of `out` starting at `out_offset` with zeros (or
+// empty strings), and updates `out_offset`. Returns false if `out` doesn't have
+// room for them.
+inline bool PadFeature(DataType dtype, size_t num_to_pad, Tensor* out,
+                       size_t* out_offset) {
+  const size_t out_elements = out->NumElements();
+  if (*out_offset > out_elements || num_to_pad > out_elements - *out_offset) {
+    return false;
+  }
+  switch (dtype) {
+    case DT_STRING:
+      // Newly allocated string tensors already hold empty strings.
+      break;
+    case DT_FLOAT: {
+      float* begin = out->flat<float>().data() + *out_offset;
+      std::fill(begin, begin + num_to_pad, 0.0f);
+      break;
+    }
+    case DT_INT64: {
+      int64_t* begin = out->flat<int64_t>().data() + *out_offset;
+      std::fill(begin, begin + num_to_pad, 0);
+      break;
+    }
+    default:
+      ReportUnexpectedDataType(dtype);
+  }
+  *out_offset += num_to_pad;
+  return true;
+}
+
 // Returns the length of the next feature on `stream`.
 // Returns -1 if the next feature on `stream` doesn't match `dtype`.
 inline int GetFeatureLength(DataType dtype,
                             protobuf::io::CodedInputStream* stream) {
   switch (dtype) {
     case DT_STRING:
-      return ParseBytesFeature(stream, nullptr);
+      return ParseBytesFeature(stream, nullptr, 0);
     case DT_FLOAT:
-      return ParseFloatFeature(stream, nullptr);
+      return ParseFloatFeature(stream, nullptr, 0);
     case DT_INT64:
-      return ParseInt64Feature(stream, nullptr);
+      return ParseInt64Feature(stream, nullptr, 0);
     default:
       ReportUnexpectedDataType(dtype);
       return -1;
@@ -2292,10 +2338,14 @@ absl::Status GetSequenceFeatureLengths(
 }
 
 // Copies src into dst[dst_offset:dst_offset+src.size], and then increments
-// dst_offset by src.size.
-void CopyTensorIntoTensor(DataType dtype, const Tensor& src, Tensor* dst,
+// dst_offset by src.size. Returns false if dst doesn't have room for src.
+bool CopyTensorIntoTensor(DataType dtype, const Tensor& src, Tensor* dst,
                           size_t* dst_offset) {
-  size_t src_size = src.NumElements();
+  const size_t src_size = src.NumElements();
+  const size_t dst_size = dst->NumElements();
+  if (*dst_offset > dst_size || src_size > dst_size - *dst_offset) {
+    return false;
+  }
   switch (dtype) {
     case DT_INT64: {
       auto src_t = src.flat<int64_t>().data();
@@ -2319,6 +2369,7 @@ void CopyTensorIntoTensor(DataType dtype, const Tensor& src, Tensor* dst,
       ReportUnexpectedDataType(dtype);
   }
   *dst_offset += src_size;
+  return true;
 }
 
 // Parses dense features in `context_features`, and writes their parsed
@@ -2365,7 +2416,11 @@ absl::Status ParseContextDenseFeatures(
                            " (data type: ", DataTypeString(c.dtype), ")",
                            " is required but could not be found."));
         }
-        CopyTensorIntoTensor(dtype, c.default_value, &out, &out_offset);
+        if (!CopyTensorIntoTensor(dtype, c.default_value, &out, &out_offset)) {
+          return absl::InvalidArgumentError(
+              absl::StrCat("Default value for context feature ", c.feature_name,
+                           " does not fit its output"));
+        }
         num_elements += c.default_value.NumElements();
       } else if (!feature_proto.empty()) {
         protobuf::io::CodedInputStream stream(
@@ -2416,6 +2471,8 @@ absl::Status ParseContextSparseFeatures(
         Tensor(allocator, DT_INT64, TensorShape({is_batch ? 2 : 1}));
     Tensor& out_values = context_result->sparse_values[t];
     size_t out_values_offset = 0;
+    // There is one row of indices per value, and ParseFeature never adds more
+    // values than out_values can hold, so these writes stay in bounds too.
     int64_t* out_indices =
         context_result->sparse_indices[t].flat<int64_t>().data();
     auto out_shape = context_result->sparse_shapes[t].vec<int64_t>();
@@ -2597,23 +2654,8 @@ absl::Status ParseSequenceDenseFeatures(
     (*dense_feature_lengths)[t] =
         Tensor(allocator, DT_INT64, dense_length_shape);
     int64_t* out_lengths = (*dense_feature_lengths)[t].flat<int64_t>().data();
-
-    tstring* out_bytes = nullptr;
-    float* out_float = nullptr;
-    int64_t* out_int64 = nullptr;
-    switch (dtype) {
-      case DT_STRING:
-        out_bytes = sequence_result->dense_values[t].flat<tstring>().data();
-        break;
-      case DT_FLOAT:
-        out_float = sequence_result->dense_values[t].flat<float>().data();
-        break;
-      case DT_INT64:
-        out_int64 = sequence_result->dense_values[t].flat<int64_t>().data();
-        break;
-      default:
-        ReportUnexpectedDataType(dtype);
-    }
+    Tensor& out = sequence_result->dense_values[t];
+    size_t out_offset = 0;
 
     // Fill in the values.
     for (int e = 0; e < num_examples; e++) {
@@ -2647,23 +2689,7 @@ absl::Status ParseSequenceDenseFeatures(
           auto limit = stream.PushLimit(feature_length);
           int num_added = 0;
           if (feature_length > 2) {
-            switch (dtype) {
-              case DT_STRING:
-                num_added = ParseBytesFeature(&stream, out_bytes);
-                if (num_added >= 0) out_bytes += num_added;
-                break;
-              case DT_FLOAT:
-                num_added = ParseFloatFeature(&stream, out_float);
-                if (num_added >= 0) out_float += num_added;
-                break;
-              case DT_INT64:
-                num_added = ParseInt64Feature(&stream, out_int64);
-                if (num_added >= 0) out_int64 += num_added;
-                break;
-              default:
-                ReportUnexpectedDataType(dtype);
-                num_added = 0;
-            }
+            num_added = ParseFeature(dtype, &stream, &out, &out_offset);
             if (num_added < 0) {
               // This should be unreachable -- we already scanned the feature in
               // GetSequenceFeatureLengths, and it hasn't changed since then.
@@ -2686,21 +2712,14 @@ absl::Status ParseSequenceDenseFeatures(
       }
       *out_lengths++ = num_rows;
       // Pad as necessary.
-      int num_to_pad = expected_max_elements - num_elements;
-      switch (dtype) {
-        case DT_STRING:
-          out_bytes += num_to_pad;
-          break;
-        case DT_FLOAT:
-          PadFloatFeature(num_to_pad, out_float);
-          out_float += num_to_pad;
-          break;
-        case DT_INT64:
-          PadInt64Feature(num_to_pad, out_int64);
-          out_int64 += num_to_pad;
-          break;
-        default:
-          ReportUnexpectedDataType(dtype);
+      if (num_elements > expected_max_elements ||
+          !PadFeature(dtype, expected_max_elements - num_elements, &out,
+                      &out_offset)) {
+        // This should be unreachable -- we already scanned the feature in
+        // GetSequenceFeatureLengths, and it hasn't changed since then.
+        return absl::InvalidArgumentError(
+            absl::StrCat("Error in sequence feature ", c.feature_name,
+                         " in example ", ExampleName(example_names, e)));
       }
     }
   }
@@ -2730,22 +2749,10 @@ absl::Status ParseSequenceSparseFeatures(
     sequence_result->sparse_shapes[t] =
         Tensor(allocator, DT_INT64, TensorShape({is_batch ? 3 : 2}));
 
-    tstring* out_bytes = nullptr;
-    float* out_float = nullptr;
-    int64_t* out_int64 = nullptr;
-    switch (dtype) {
-      case DT_STRING:
-        out_bytes = sequence_result->sparse_values[t].flat<tstring>().data();
-        break;
-      case DT_FLOAT:
-        out_float = sequence_result->sparse_values[t].flat<float>().data();
-        break;
-      case DT_INT64:
-        out_int64 = sequence_result->sparse_values[t].flat<int64_t>().data();
-        break;
-      default:
-        ReportUnexpectedDataType(dtype);
-    }
+    Tensor& out_values = sequence_result->sparse_values[t];
+    size_t out_values_offset = 0;
+    // There is one row of indices per value, and ParseFeature never adds more
+    // values than out_values can hold, so these writes stay in bounds too.
     int64_t* out_indices =
         sequence_result->sparse_indices[t].flat<int64_t>().data();
     auto out_shape = sequence_result->sparse_shapes[t].vec<int64_t>();
@@ -2774,24 +2781,8 @@ absl::Status ParseSequenceSparseFeatures(
         }
         if (feature_length > 2) {
           auto limit = stream.PushLimit(feature_length);
-          int num_added;
-          switch (dtype) {
-            case DT_STRING:
-              num_added = ParseBytesFeature(&stream, out_bytes);
-              if (num_added >= 0) out_bytes += num_added;
-              break;
-            case DT_FLOAT:
-              num_added = ParseFloatFeature(&stream, out_float);
-              if (num_added >= 0) out_float += num_added;
-              break;
-            case DT_INT64:
-              num_added = ParseInt64Feature(&stream, out_int64);
-              if (num_added >= 0) out_int64 += num_added;
-              break;
-            default:
-              ReportUnexpectedDataType(dtype);
-              num_added = 0;
-          }
+          int num_added =
+              ParseFeature(dtype, &stream, &out_values, &out_values_offset);
           if (num_added < 0) {
             // This should be unreachable -- we already scanned the feature in
             // GetSequenceFeatureLengths, and it hasn't changed since then.
@@ -2945,6 +2936,13 @@ absl::Status ParseSequenceRaggedFeatures(
                 absl::StrCat("Error in sequence feature ", c.feature_name,
                              " in example ", ExampleName(example_names, e)));
           }
+          // The inner splits were sized from the row count found by
+          // GetSequenceFeatureLengths; don't write past them if this pass
+          // finds more rows.
+          if (outer_split >= expected_num_rows) {
+            return absl::InvalidArgumentError(absl::StrCat(
+                "Unexpected number of rows for feature ", c.feature_name));
+          }
           if (int32_inner_splits) {
             *int32_inner_splits++ = inner_split;
           } else if (int64_inner_splits) {
diff --git a/tensorflow/core/util/example_proto_fast_parsing_test.cc b/tensorflow/core/util/example_proto_fast_parsing_test.cc
index a8180ec240f..9d4e9c16240 100644
--- a/tensorflow/core/util/example_proto_fast_parsing_test.cc
+++ b/tensorflow/core/util/example_proto_fast_parsing_test.cc
@@ -619,6 +619,88 @@ TEST(FastParse, OobWriteVulnerabilityMalformedSparseSequenceExample) {
   EXPECT_TRUE(absl::IsInvalidArgument(status));
 }
 
+TEST(FastParse, SequenceExampleFillsAllOutputKinds) {
+  // Two SequenceExamples with a different number of rows and values each, so
+  // that dense feature lists are padded and sparse/ragged ones are packed.
+  SequenceExample examples[2];
+  (*examples[0].mutable_context()->mutable_feature())["c"]
+      .mutable_int64_list()
+      ->add_value(11);
+  (*examples[1].mutable_context()->mutable_feature())["c"]
+      .mutable_int64_list()
+      ->add_value(12);
+  auto* lists0 = examples[0].mutable_feature_lists()->mutable_feature_list();
+  auto* lists1 = examples[1].mutable_feature_lists()->mutable_feature_list();
+  // Adds one row holding `values` to a float feature list.
+  auto add_float_row = [](FeatureList* list, std::vector<float> values) {
+    auto* float_list = list->add_feature()->mutable_float_list();
+    for (float v : values) float_list->add_value(v);
+  };
+  add_float_row(&(*lists0)["d"], {1, 2});
+  add_float_row(&(*lists0)["d"], {3, 4});
+  add_float_row(&(*lists1)["d"], {5, 6});
+  (*lists0)["s"].add_feature()->mutable_int64_list()->add_value(7);
+  auto* s01 = (*lists0)["s"].add_feature()->mutable_int64_list();
+  s01->add_value(8);
+  s01->add_value(9);
+  (*lists1)["s"].add_feature()->mutable_int64_list()->add_value(10);
+  (*lists0)["r"].add_feature()->mutable_bytes_list()->add_value("a");
+  auto* r01 = (*lists0)["r"].add_feature()->mutable_bytes_list();
+  r01->add_value("b");
+  r01->add_value("c");
+  const std::vector<tstring> serialized = {Serialize(examples[0]),
+                                           Serialize(examples[1])};
+
+  FastParseExampleConfig context_config;
+  context_config.dense.emplace_back("c", DT_INT64, PartialTensorShape({1}),
+                                    Tensor(DT_INT64, TensorShape({0})),
+                                    /*variable_length=*/false,
+                                    /*elements_per_stride=*/1);
+  FastParseExampleConfig sequence_config;
+  sequence_config.dense.emplace_back("d", DT_FLOAT, PartialTensorShape({2}),
+                                     Tensor(DT_FLOAT, TensorShape({0})),
+                                     /*variable_length=*/true,
+                                     /*elements_per_stride=*/2);
+  AddSparseFeature("s", DT_INT64, &sequence_config);
+  sequence_config.ragged.emplace_back("r", DT_STRING, DT_INT64);
+
+  Result context_result, sequence_result;
+  std::vector<Tensor> dense_lengths;
+  const absl::Status status = FastParseSequenceExample(
+      context_config, sequence_config, serialized, {}, nullptr, &context_result,
+      &sequence_result, &dense_lengths);
+  ASSERT_TRUE(status.ok()) << status;
+
+  auto expect_values = [](const Tensor& t, std::vector<int64_t> shape,
+                          auto expected) {
+    EXPECT_EQ(t.shape(), TensorShape(shape));
+    using T = typename decltype(expected)::value_type;
+    auto flat = t.flat<T>();
+    ASSERT_EQ(flat.size(), expected.size());
+    for (int i = 0; i < flat.size(); ++i) {
+      EXPECT_EQ(flat(i), expected[i]) << "i=" << i;
+    }
+  };
+  expect_values(context_result.dense_values[0], {2, 1},
+                std::vector<int64_t>{11, 12});
+  // The second example's missing row is padded with zeros.
+  expect_values(sequence_result.dense_values[0], {2, 2, 2},
+                std::vector<float>{1, 2, 3, 4, 5, 6, 0, 0});
+  expect_values(dense_lengths[0], {2}, std::vector<int64_t>{2, 1});
+  expect_values(sequence_result.sparse_values[0], {4},
+                std::vector<int64_t>{7, 8, 9, 10});
+  expect_values(sequence_result.sparse_indices[0], {4, 3},
+                std::vector<int64_t>{0, 0, 0, 0, 1, 0, 0, 1, 1, 1, 0, 0});
+  expect_values(sequence_result.sparse_shapes[0], {3},
+                std::vector<int64_t>{2, 2, 2});
+  expect_values(sequence_result.ragged_values[0], {3},
+                std::vector<tstring>{"a", "b", "c"});
+  expect_values(sequence_result.ragged_splits[0], {3},
+                std::vector<int64_t>{0, 1, 3});
+  expect_values(sequence_result.ragged_outer_splits[0], {3},
+                std::vector<int64_t>{0, 2, 2});
+}
+
 }  // namespace
 }  // namespace example
 }  // namespace tensorflow

From 5e82d4b5acd5ae84b42b7b6a0342944a4c523909 Mon Sep 17 00:00:00 2001
From: Alex Pivovarov <upwind@google.com>
Date: Mon, 12 Oct 2026 11:29:18 -0700
Subject: [PATCH 0610/1000] [XLA:GPU] Add batch=1 test cases to
 `topk_special_floats_test`

Previously, `TopKWithSpecialFloats` and `TopKWithSpecialBFloat16` in `xla/tests/topk_special_floats_test.cc` only tested `batch = 8`. On Blackwell GPUs (e.g. B200), when `k > 16` causes `TopkSpecializer` to fall back to `TopkDecomposer` (`hlo.sort` + `slice`), the 32-bit cost heuristic in `IsCubSortFasterOnBlackwell` does not rewrite `f32[8, n]` to CUB radix sort (`xla.gpu.ext.cub_sort_pairs`), so `batch = 8` only exercised the MLIR/LLVM bitonic sort fallback.

### Why this change is needed:
- Add `batch = 1` (and `n = 4096` in `TopKWithSpecialFloats`) in addition to `batch = 8` so the test suite covers both `batch = 1` and `batch = 8` across `TopK` lowering paths (`__gpu$TopK` for `k <= 16` and `hlo.sort` + `slice` for `k > 16`).
- Temporarily skip `(batch == 1 && n >= 4096 && k > 16)` on GPU in `TopKWithSpecialFloats` while https://github.com/openxla/xla/issues/50426 is being fixed (where `SortRewriter` lowers the fallback `hlo.sort` to `cub_sort_pairs`, which treats `-0.0` and `+0.0` as equal instead of `+0.0 > -0.0` under `order=TOTAL`).

PiperOrigin-RevId: 998045725
---
 third_party/xla/xla/tests/BUILD               |   1 +
 .../xla/xla/tests/topk_special_floats_test.cc | 172 ++++++++++--------
 2 files changed, 93 insertions(+), 80 deletions(-)

diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD
index b0609dcdbd8..7cb302036ee 100644
--- a/third_party/xla/xla/tests/BUILD
+++ b/third_party/xla/xla/tests/BUILD
@@ -2233,6 +2233,7 @@ xla_test(
         ":hlo_test_base",
         ":literal_test_util",
         ":xla_internal_test_main",
+        ":xla_test_backend_predicates",
         "//xla:array2d",
         "//xla:error_spec",
         "//xla:literal",
diff --git a/third_party/xla/xla/tests/topk_special_floats_test.cc b/third_party/xla/xla/tests/topk_special_floats_test.cc
index 66e602b8d45..30eaa81e35e 100644
--- a/third_party/xla/xla/tests/topk_special_floats_test.cc
+++ b/third_party/xla/xla/tests/topk_special_floats_test.cc
@@ -27,6 +27,7 @@ limitations under the License.
 #include "xla/tests/hlo_interpreter_reference_mixin.h"
 #include "xla/tests/hlo_test_base.h"
 #include "xla/tests/literal_test_util.h"
+#include "xla/tests/xla_test_backend_predicates.h"
 #include "xla/tsl/platform/test.h"
 #include "xla/types.h"
 
@@ -97,51 +98,59 @@ TEST_F(TopkTest, TopKWithSpecialFloats) {
       32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47,
       48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63};
 
-  for (int n : {512, 1024, 32768, 65536 + 4096}) {
-    Array2D<float> input(8, n, neg_smallest_qnan);
-    for (int i = 0; i < input_row_start.size(); ++i) {
-      for (int j = 0; j < 8; ++j) {
-        input(j, i) = input_row_start[i];
-      }
-    }
-
-    auto literal = LiteralUtil::CreateR2FromArray2D<float>(input);
-
-    for (int k : {2, 4, 8, 16, 32, 64}) {
-      const auto hlo_text_module = absl::StrFormat(R"(
-      HloModule topk
-
-      ENTRY TopK {
-        x = f32[8,%d]{1,0} parameter(0)
-        ROOT topk = (f32[8,%d]{1,0}, s32[8,%d]{1,0}) topk(x), k=%d, largest=true
-      }
-    )",
-                                                   n, k, k, k);
-
-      ASSERT_OK_AND_ASSIGN(auto module,
-                           ParseAndReturnVerifiedModule(hlo_text_module));
-
-      Array2D<float> expected_vals_arr(8, k);
-      Array2D<int32_t> expected_idxs_arr(8, k);
-      for (int i = 0; i < 8; ++i) {
-        for (int j = 0; j < k; ++j) {
-          expected_vals_arr(i, j) = expected_values_64[j];
-          expected_idxs_arr(i, j) = expected_indices_64[j];
+  for (int batch : {1, 8}) {
+    for (int n : {512, 1024, 4096, 32768, 65536 + 4096}) {
+      Array2D<float> input(batch, n, neg_smallest_qnan);
+      for (int i = 0; i < input_row_start.size(); ++i) {
+        for (int j = 0; j < batch; ++j) {
+          input(j, i) = input_row_start[i];
         }
       }
 
-      auto expected_vals_literal =
-          LiteralUtil::CreateR2FromArray2D<float>(expected_vals_arr);
-      auto expected_idxs_literal =
-          LiteralUtil::CreateR2FromArray2D<int32_t>(expected_idxs_arr);
-      auto expected_tuple = LiteralUtil::MakeTuple(
-          {&expected_vals_literal, &expected_idxs_literal});
+      auto literal = LiteralUtil::CreateR2FromArray2D<float>(input);
 
-      ASSERT_OK_AND_ASSIGN(Literal actual_tuple,
-                           Execute(std::move(module), {&literal}));
+      for (int k : {2, 4, 8, 16, 32, 64}) {
+        if (test::DeviceTypeIs(test::kGpu) && batch == 1 && n >= 4096 &&
+            k > 16) {
+          // TODO(b/572311831): [XLA:GPU] hlo.topk gives incorrect ordering
+          // for +0.0 vs -0.0 when lowered to CUB radix sort.
+          continue;
+        }
+        const auto hlo_text_module = absl::StrFormat(
+            R"(
+      HloModule topk
 
-      EXPECT_TRUE(LiteralTestUtil::Near(expected_tuple, actual_tuple,
-                                        xla::ErrorSpec{0, 0}));
+      ENTRY TopK {
+        x = f32[%d,%d]{1,0} parameter(0)
+        ROOT topk = (f32[%d,%d]{1,0}, s32[%d,%d]{1,0}) topk(x), k=%d, largest=true
+      }
+    )",
+            batch, n, batch, k, batch, k, k);
+
+        ASSERT_OK_AND_ASSIGN(auto module,
+                             ParseAndReturnVerifiedModule(hlo_text_module));
+
+        Array2D<float> expected_vals_arr(batch, k);
+        Array2D<int32_t> expected_idxs_arr(batch, k);
+        for (int i = 0; i < batch; ++i) {
+          for (int j = 0; j < k; ++j) {
+            expected_vals_arr(i, j) = expected_values_64[j];
+            expected_idxs_arr(i, j) = expected_indices_64[j];
+          }
+        }
+
+        auto expected_vals_literal =
+            LiteralUtil::CreateR2FromArray2D<float>(expected_vals_arr);
+        auto expected_idxs_literal =
+            LiteralUtil::CreateR2FromArray2D<int32_t>(expected_idxs_arr);
+        auto expected_tuple = LiteralUtil::MakeTuple(
+            {&expected_vals_literal, &expected_idxs_literal});
+
+        ASSERT_OK_AND_ASSIGN(Literal actual_tuple,
+                             Execute(std::move(module), {&literal}));
+        EXPECT_TRUE(LiteralTestUtil::Near(expected_tuple, actual_tuple,
+                                          xla::ErrorSpec{0, 0}));
+      }
     }
   }
 }
@@ -201,51 +210,54 @@ TEST_F(TopkTest, TopKWithSpecialBFloat16) {
       32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47,
       48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63};
 
-  for (int n : {512, 1024, 32768, 65536 + 4096}) {
-    Array2D<bfloat16> input(8, n, neg_smallest_qnan);
-    for (int i = 0; i < input_row_start.size(); ++i) {
-      for (int j = 0; j < 8; ++j) {
-        input(j, i) = input_row_start[i];
-      }
-    }
-
-    auto literal = LiteralUtil::CreateR2FromArray2D<bfloat16>(input);
-
-    for (int k : {2, 4, 8, 16, 32, 64}) {
-      const auto hlo_text_module = absl::StrFormat(R"(
-      HloModule topk
-
-      ENTRY TopK {
-        x = bf16[8,%d]{1,0} parameter(0)
-        ROOT topk = (bf16[8,%d]{1,0}, s32[8,%d]{1,0}) topk(x), k=%d, largest=true
-      }
-    )",
-                                                   n, k, k, k);
-
-      ASSERT_OK_AND_ASSIGN(auto module,
-                           ParseAndReturnVerifiedModule(hlo_text_module));
-
-      Array2D<bfloat16> expected_vals_arr(8, k);
-      Array2D<int32_t> expected_idxs_arr(8, k);
-      for (int i = 0; i < 8; ++i) {
-        for (int j = 0; j < k; ++j) {
-          expected_vals_arr(i, j) = expected_values_64[j];
-          expected_idxs_arr(i, j) = expected_indices_64[j];
+  for (int batch : {1, 8}) {
+    for (int n : {512, 1024, 4096, 32768, 65536 + 4096}) {
+      Array2D<bfloat16> input(batch, n, neg_smallest_qnan);
+      for (int i = 0; i < input_row_start.size(); ++i) {
+        for (int j = 0; j < batch; ++j) {
+          input(j, i) = input_row_start[i];
         }
       }
 
-      auto expected_vals_literal =
-          LiteralUtil::CreateR2FromArray2D<bfloat16>(expected_vals_arr);
-      auto expected_idxs_literal =
-          LiteralUtil::CreateR2FromArray2D<int32_t>(expected_idxs_arr);
-      auto expected_tuple = LiteralUtil::MakeTuple(
-          {&expected_vals_literal, &expected_idxs_literal});
+      auto literal = LiteralUtil::CreateR2FromArray2D<bfloat16>(input);
 
-      ASSERT_OK_AND_ASSIGN(Literal actual_tuple,
-                           Execute(std::move(module), {&literal}));
+      for (int k : {2, 4, 8, 16, 32, 64}) {
+        const auto hlo_text_module = absl::StrFormat(
+            R"(
+      HloModule topk
 
-      EXPECT_TRUE(LiteralTestUtil::Near(expected_tuple, actual_tuple,
-                                        xla::ErrorSpec{0, 0}));
+      ENTRY TopK {
+        x = bf16[%d,%d]{1,0} parameter(0)
+        ROOT topk = (bf16[%d,%d]{1,0}, s32[%d,%d]{1,0}) topk(x), k=%d, largest=true
+      }
+    )",
+            batch, n, batch, k, batch, k, k);
+
+        ASSERT_OK_AND_ASSIGN(auto module,
+                             ParseAndReturnVerifiedModule(hlo_text_module));
+
+        Array2D<bfloat16> expected_vals_arr(batch, k);
+        Array2D<int32_t> expected_idxs_arr(batch, k);
+        for (int i = 0; i < batch; ++i) {
+          for (int j = 0; j < k; ++j) {
+            expected_vals_arr(i, j) = expected_values_64[j];
+            expected_idxs_arr(i, j) = expected_indices_64[j];
+          }
+        }
+
+        auto expected_vals_literal =
+            LiteralUtil::CreateR2FromArray2D<bfloat16>(expected_vals_arr);
+        auto expected_idxs_literal =
+            LiteralUtil::CreateR2FromArray2D<int32_t>(expected_idxs_arr);
+        auto expected_tuple = LiteralUtil::MakeTuple(
+            {&expected_vals_literal, &expected_idxs_literal});
+
+        ASSERT_OK_AND_ASSIGN(Literal actual_tuple,
+                             Execute(std::move(module), {&literal}));
+
+        EXPECT_TRUE(LiteralTestUtil::Near(expected_tuple, actual_tuple,
+                                          xla::ErrorSpec{0, 0}));
+      }
     }
   }
 }

From eae482aea887251b5c9b054f4660feaccbc9bc83 Mon Sep 17 00:00:00 2001
From: Vlad Sytchenko <vsytch@google.com>
Date: Mon, 12 Oct 2026 11:30:41 -0700
Subject: [PATCH 0611/1000] [Mosaic] Mark tpu.scan as Pure

PiperOrigin-RevId: 998046885
---
 third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td
index f978b8ac117..4786a651980 100644
--- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td
+++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td
@@ -200,7 +200,7 @@ def TPU_ReduceOp : TPU_Op<"reduce", [Pure, SameOperandsAndResultRank]> {
 //  mask   : Elementwise vector mask. The scan operation starts from the
 //           lowest-indexed non-masked vector element (all previous elements
 //           have undefined values). Not taken for I1 input vectors.
-def TPU_ScanOp : TPU_Op<"scan"> {
+def TPU_ScanOp : TPU_Op<"scan", [Pure]> {
   let arguments = (ins
     VectorOfNonZeroRankOf<[I1, I16, I32, BF16, F32]>:$input,
     TPU_ReductionKindAttr:$kind,

From 2fd11341ad1ed636e9c8e4023235fe11833d0b37 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Mon, 12 Oct 2026 11:42:23 -0700
Subject: [PATCH 0612/1000] [XLA:GPU] Skip symmetric buffer test in
 p2p_ops_e2e_test on pre-Hopper GPUs

PiperOrigin-RevId: 998054813
---
 .../xla/xla/backends/gpu/tests/p2p_ops_e2e_test.cc        | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/third_party/xla/xla/backends/gpu/tests/p2p_ops_e2e_test.cc b/third_party/xla/xla/backends/gpu/tests/p2p_ops_e2e_test.cc
index 66b6229db35..7480ef8f314 100644
--- a/third_party/xla/xla/backends/gpu/tests/p2p_ops_e2e_test.cc
+++ b/third_party/xla/xla/backends/gpu/tests/p2p_ops_e2e_test.cc
@@ -39,6 +39,14 @@ class P2POps : public CollectiveOpsWithFlagsBase,
             /*enable_symmetric_buffer=*/GetParam(),
             /*memory_size=*/32 * kMB,
             /*collectives_memory_size=*/0) {}
+
+ protected:
+  void SetUp() override {
+    CollectiveOpsWithFlagsBase::SetUp();
+    if (!IsHopperAndHigher() && enable_symmetric_buffer_) {
+      GTEST_SKIP() << "Test requires Hopper or higher";
+    }
+  }
 };
 
 TEST_P(P2POps, CollectivePermute) {

From dd3d64c1617078b0df1bdf159a071296c4dbad5b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 12 Oct 2026 11:46:39 -0700
Subject: [PATCH 0613/1000] Integrate LLVM at llvm/llvm-project@2e0d9e12b36e

Updates LLVM usage to match
[2e0d9e12b36e](https://github.com/llvm/llvm-project/commit/2e0d9e12b36e853d9b108c4da09509cd54d1f3f9)

PiperOrigin-RevId: 998057483
---
 .../xla/third_party/llvm/generated.patch      |  11 -
 .../xla/third_party/llvm/workspace.bzl        |   4 +-
 .../xla/third_party/shardy/temporary.patch    | 826 ++++++++++++++++++
 .../xla/third_party/shardy/workspace.bzl      |   4 +-
 4 files changed, 830 insertions(+), 15 deletions(-)

diff --git a/third_party/xla/third_party/llvm/generated.patch b/third_party/xla/third_party/llvm/generated.patch
index e9c206cebe4..509398da979 100644
--- a/third_party/xla/third_party/llvm/generated.patch
+++ b/third_party/xla/third_party/llvm/generated.patch
@@ -1,12 +1 @@
 Auto generated patch. Do not edit or delete it, even if empty.
-diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll b/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll
---- a/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll
-+++ b/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll
-@@ -1,5 +1,7 @@
- ; Checks if conversion is only inserted for the spilled register
- ; instead of 2 conversions for the W register
-+
-+; REQUIRES: asserts
- ; RUN: llc -O2 -march=hexagon -mcpu=hexagonv81 -enable-xqf-gen=true \
- ; RUN: -hexagon-qfloat-mode=ieee -mattr=+hvxv81,+hvx-length128B \
- ; RUN: -enable-postra-xqf-check -debug-only=handle-qfp 2>&1 < %s -o /dev/null | FileCheck %s
diff --git a/third_party/xla/third_party/llvm/workspace.bzl b/third_party/xla/third_party/llvm/workspace.bzl
index 68a14533c79..e42d7dbdd78 100644
--- a/third_party/xla/third_party/llvm/workspace.bzl
+++ b/third_party/xla/third_party/llvm/workspace.bzl
@@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "b08e005995ed29a2f89aa9a3f599a9d45db65523"
-    LLVM_SHA256 = "3f4f1765816e22a21ac41dbf71d91d519b0b371b05d7e33ef2a0a862c838420c"
+    LLVM_COMMIT = "2e0d9e12b36e853d9b108c4da09509cd54d1f3f9"
+    LLVM_SHA256 = "8e4a851780af109d623fa6d6db9b1e62d7126272d17268fed1b149f0e737c1d7"
 
     tf_http_archive(
         name = name,
diff --git a/third_party/xla/third_party/shardy/temporary.patch b/third_party/xla/third_party/shardy/temporary.patch
index e69de29bb2d..6614bd1c1a7 100644
--- a/third_party/xla/third_party/shardy/temporary.patch
+++ b/third_party/xla/third_party/shardy/temporary.patch
@@ -0,0 +1,826 @@
+diff --git a/shardy/dialect/sdy/ir/canonicalization.cc b/shardy/dialect/sdy/ir/canonicalization.cc
+index 55f7ebcf..9af43e32 100644
+--- a/shardy/dialect/sdy/ir/canonicalization.cc
++++ b/shardy/dialect/sdy/ir/canonicalization.cc
+@@ -25,6 +25,7 @@ limitations under the License.
+ #include "llvm/ADT/SmallVector.h"
+ #include "llvm/Support/MathExtras.h"
+ #include "mlir/IR/BuiltinAttributes.h"
++#include "mlir/IR/BuiltinTypes.h"
+ #include "mlir/IR/Diagnostics.h"
+ #include "mlir/IR/MLIRContext.h"
+ #include "mlir/IR/OperationSupport.h"
+@@ -203,6 +204,42 @@ std::optional<ReduceScatterFusionInfo> matchAndComputeReduceScatterFusionAxes(
+   SmallVector<DimensionShardingAttr> reduceScatterOutDimShardings;
+   reduceScatterOutDimShardings.reserve(rank);
+ 
++  // TODO(b/553078138): This is a conservative workaround. The proper fix is to
++  // make a reduce-scatter with unreduced axes exportable for the non-divisible
++  // case (in `convertReduceScatter`); once that lands, this bail-out (which
++  // falls back to the less efficient `all_reduce` + `all_slice` form) can be
++  // removed so these cases fuse as well.
++  // Bail out of the fusion if it would create an `sdy.reduce_scatter` that
++  // scatters a dimension by a group size that does not evenly divide the local
++  // (already-sharded) dimension size. Exporting a reduce-scatter with unreduced
++  // axes (see `convertReduceScatter` in export_manual_reduction_collectives)
++  // requires the scattered dimension to be divisible by the group size, and
++  // would otherwise fail. Keeping the `all_reduce` + `all_slice` form preserves
++  // correct behavior for these non-divisible cases.
++  if (auto operandType =
++          dyn_cast<RankedTensorType>(allReduceOp.getTensor().getType())) {
++    // Local (per-device) shape of the all-reduce operand under its existing
++    // sharding. `allowNonDivisible=true` handles the case where the existing
++    // sharding does not evenly divide a dimension: the local size is rounded up
++    // (ceil), matching how Shardy pads such dimensions.
++    RankedTensorType localType = allReduceInSharding.getLocalTensorType(
++        operandType, mesh, /*allowNonDivisible=*/true);
++    for (int64_t dim = 0; dim < rank; ++dim) {
++      int64_t groupSize =
++          getTotalAxesSize(reduceScatterAxes[dim].getValue(), mesh);
++      if (groupSize <= 1) {
++        continue;
++      }
++      int64_t localDimSize = localType.getDimSize(dim);
++      if (ShapedType::isDynamic(localDimSize)) {
++        continue;
++      }
++      if (localDimSize % groupSize != 0) {
++        return std::nullopt;
++      }
++    }
++  }
++
+   for (auto [operandDimSharding, reduceScatterAxesPerDim] : llvm::zip_equal(
+            allReduceInSharding.getDimShardings(), reduceScatterAxes)) {
+     SmallVector<AxisRefAttr> expectedDimShardingAxes =
+diff --git a/shardy/dialect/sdy/ir/extensions/stablehlo_extensions.cc b/shardy/dialect/sdy/ir/extensions/stablehlo_extensions.cc
+index d4825b2e..75aac4e0 100644
+--- a/shardy/dialect/sdy/ir/extensions/stablehlo_extensions.cc
++++ b/shardy/dialect/sdy/ir/extensions/stablehlo_extensions.cc
+@@ -70,20 +70,20 @@ struct WhileShardableDataFlowOpInterface
+   }
+ };
+ 
+-template <typename OpTy>
+-struct RegionBranchShardableDataFlowOpInterface
++struct CaseShardableDataFlowOpInterface
+     : public ShardableDataFlowOpInterface::ExternalModel<
+-          RegionBranchShardableDataFlowOpInterface<OpTy>, OpTy> {
++          CaseShardableDataFlowOpInterface, stablehlo::CaseOp> {
+   ResultRange getOpResultEdgeOwners(Operation *op) const {
+     return op->getResults();
+   }
+ 
+   SmallVector<OpOperand *> getEdgeSources(Operation *op, Value owner) const {
++    auto caseOp = cast<stablehlo::CaseOp>(op);
+     auto opResult = dyn_cast<OpResult>(owner);
+     assert(opResult && opResult.getOwner() == op);
+     SmallVector<OpOperand *> sources;
+-    sources.reserve(op->getNumRegions());
+-    for (Region& branch : op->getRegions()) {
++    sources.reserve(caseOp.getBranches().size());
++    for (Region &branch : caseOp.getBranches()) {
+       sources.push_back(&branch.front().getTerminator()->getOpOperand(
+           opResult.getResultNumber()));
+     }
+@@ -132,10 +132,7 @@ void registerStablehloExtensions(MLIRContext *ctx) {
+   // Ensure dialect is loaded before attaching interfaces.
+   ctx->loadDialect<stablehlo::StablehloDialect>();
+   stablehlo::WhileOp::attachInterface<WhileShardableDataFlowOpInterface>(*ctx);
+-  stablehlo::CaseOp::attachInterface<
+-      RegionBranchShardableDataFlowOpInterface<stablehlo::CaseOp>>(*ctx);
+-  stablehlo::IfOp::attachInterface<
+-      RegionBranchShardableDataFlowOpInterface<stablehlo::IfOp>>(*ctx);
++  stablehlo::CaseOp::attachInterface<CaseShardableDataFlowOpInterface>(*ctx);
+   stablehlo::OptimizationBarrierOp::attachInterface<
+       OptBarrierShardableDataFlowOpInterface>(*ctx);
+ }
+diff --git a/shardy/dialect/sdy/ir/test/collective_canonicalization.mlir b/shardy/dialect/sdy/ir/test/collective_canonicalization.mlir
+index de616ede..14b0acf3 100644
+--- a/shardy/dialect/sdy/ir/test/collective_canonicalization.mlir
++++ b/shardy/dialect/sdy/ir/test/collective_canonicalization.mlir
+@@ -115,15 +115,33 @@ func.func @reduce_scatter_fusion4(%arg0 : tensor<64x8xf32> {sdy.sharding=#sdy.sh
+   return %1 : tensor<64x8xf32>
+ }
+ 
++// The local size of dim 0 is 2 / size("x"=2) = 1, which is not divisible by the
++// reduce-scatter group size (size("y"=2)). The fusion must be skipped and the
++// original `all_reduce` + `all_slice` preserved, otherwise the resulting
++// `reduce_scatter` cannot be exported (see `convertReduceScatter` in
++// export_manual_reduction_collectives).
+ // CHECK-LABEL: func @reduce_scatter_fusion_non_divisible
+ func.func @reduce_scatter_fusion_non_divisible(%arg0 : tensor<2x4xf32> {sdy.sharding=#sdy.sharding<@mesh, [{"x"}, {}]>}) -> tensor<2x4xf32> {
+-  // CHECK-NEXT: %0 = sdy.reduce_scatter [{"y"}, {}] %arg0 out_sharding=<@mesh, [{"x", "y"}, {}]> : tensor<2x4xf32>
+-  // CHECK-NEXT: return %0 : tensor<2x4xf32>
++  // CHECK-NEXT: %[[ALL_REDUCE:.*]] = sdy.all_reduce {"y"} %arg0 out_sharding=<@mesh, [{"x"}, {}]> : tensor<2x4xf32>
++  // CHECK-NEXT: %[[ALL_SLICE:.*]] = sdy.all_slice [{"y"}, {}] %[[ALL_REDUCE]] out_sharding=<@mesh, [{"x", "y"}, {}]> : tensor<2x4xf32>
++  // CHECK-NEXT: return %[[ALL_SLICE]] : tensor<2x4xf32>
+   %0 = sdy.all_reduce {"y"} %arg0 out_sharding=<@mesh, [{"x"}, {}]> : tensor<2x4xf32>
+   %1 = sdy.all_slice [{"y"}, {}] %0 out_sharding=<@mesh, [{"x", "y"}, {}]> : tensor<2x4xf32>
+   return %1 : tensor<2x4xf32>
+ }
+ 
++// The `all_reduce` operand is sharded indivisibly: dim 0 size 7 with
++// size("x"=2), so its local size is ceil(7 / 2) = 4, which is divisible by
++// the reduce-scatter group size (size("y"=2)). Therefore, fusion happens.
++// CHECK-LABEL: func @reduce_scatter_fusion_non_divisible_all_reduce
++func.func @reduce_scatter_fusion_non_divisible_all_reduce(%arg0 : tensor<7x4xf32> {sdy.sharding=#sdy.sharding<@mesh, [{"x"}, {}]>}) -> tensor<7x4xf32> {
++  // CHECK-NEXT: %0 = sdy.reduce_scatter [{"y"}, {}] %arg0 out_sharding=<@mesh, [{"x", "y"}, {}]> : tensor<7x4xf32>
++  // CHECK-NEXT: return %0 : tensor<7x4xf32>
++  %0 = sdy.all_reduce {"y"} %arg0 out_sharding=<@mesh, [{"x"}, {}]> : tensor<7x4xf32>
++  %1 = sdy.all_slice [{"y"}, {}] %0 out_sharding=<@mesh, [{"x", "y"}, {}]> : tensor<7x4xf32>
++  return %1 : tensor<7x4xf32>
++}
++
+ // CHECK-LABEL: func @reduce_scatter_fusion_missing_in_sharding
+ func.func @reduce_scatter_fusion_missing_in_sharding(%arg0 : tensor<16x8xf32>) -> tensor<16x8xf32> {
+   // CHECK-NEXT: %0 = sdy.reduce_scatter [{"x"}, {"y"}] %arg0 out_sharding=<@mesh2, [{"x"}, {"y"}]> : tensor<16x8xf32>
+diff --git a/shardy/dialect/sdy/ir/utils.cc b/shardy/dialect/sdy/ir/utils.cc
+index 590f7149..f6623986 100644
+--- a/shardy/dialect/sdy/ir/utils.cc
++++ b/shardy/dialect/sdy/ir/utils.cc
+@@ -247,13 +247,10 @@ TensorShardingPerValueAttr inlineMesh(
+                                          inlinedShardings);
+ }
+ 
+-namespace {
+-
+-template <typename SymbolTableOrOp>
+-std::pair<Attribute, MeshAttr> getCommonMeshAndRef(
+-    ArrayRef<TensorShardingAttr> operandShardings,
+-    ArrayRef<TensorShardingAttr> resultsShardings,
+-    SymbolTableOrOp symbolTableOrOp, const bool ignoreDeviceIds) {
++Attribute getCommonMeshOrRef(ArrayRef<TensorShardingAttr> operandShardings,
++                             ArrayRef<TensorShardingAttr> resultsShardings,
++                             const SymbolTable& symbolTable,
++                             const bool ignoreDeviceIds) {
+   Attribute meshOrRef;
+   MeshAttr mesh;
+   for (TensorShardingAttr sharding : llvm::concat<const TensorShardingAttr>(
+@@ -261,7 +258,7 @@ std::pair<Attribute, MeshAttr> getCommonMeshAndRef(
+     if (!sharding || sharding.getMeshOrRef() == meshOrRef) {
+       continue;
+     }
+-    MeshAttr otherMesh = sharding.getMesh(symbolTableOrOp);
++    MeshAttr otherMesh = sharding.getMesh(symbolTable);
+     if (!mesh || mesh.empty()) {
+       mesh = otherMesh;
+       meshOrRef = sharding.getMeshOrRef();
+@@ -272,7 +269,7 @@ std::pair<Attribute, MeshAttr> getCommonMeshAndRef(
+     }
+     if (!otherMesh.equals(mesh, ignoreDeviceIds)) {
+       // Found more than one mesh name.
+-      return {nullptr, nullptr};
++      return nullptr;
+     }
+     // Prefer iota device id over non-iota.
+     if (ignoreDeviceIds && otherMesh.getDeviceIds().empty()) {
+@@ -281,18 +278,7 @@ std::pair<Attribute, MeshAttr> getCommonMeshAndRef(
+     }
+   }
+ 
+-  return {meshOrRef, mesh};
+-}
+-
+-}  // namespace
+-
+-Attribute getCommonMeshOrRef(ArrayRef<TensorShardingAttr> operandShardings,
+-                             ArrayRef<TensorShardingAttr> resultsShardings,
+-                             const SymbolTable& symbolTable,
+-                             const bool ignoreDeviceIds) {
+-  return getCommonMeshAndRef(operandShardings, resultsShardings, symbolTable,
+-                             ignoreDeviceIds)
+-      .first;
++  return meshOrRef;
+ }
+ 
+ MeshAttr getCommonMesh(ArrayRef<TensorShardingAttr> shardings,
+@@ -300,24 +286,21 @@ MeshAttr getCommonMesh(ArrayRef<TensorShardingAttr> shardings,
+   return getCommonMesh(shardings, {}, symbolTable);
+ }
+ 
+-MeshAttr getCommonMesh(ArrayRef<TensorShardingAttr> shardings, Operation* op) {
+-  return getCommonMesh(shardings, {}, op);
+-}
+-
+ MeshAttr getCommonMesh(ArrayRef<TensorShardingAttr> operandShardings,
+                        ArrayRef<TensorShardingAttr> resultsShardings,
+                        const SymbolTable& symbolTable) {
+-  return getCommonMeshAndRef(operandShardings, resultsShardings, symbolTable,
+-                             /*ignoreDeviceIds=*/false)
+-      .second;
++  if (Attribute meshOrRef =
++          getCommonMeshOrRef(operandShardings, resultsShardings, symbolTable)) {
++    return getMeshOrLookup(symbolTable, meshOrRef);
++  }
++  return nullptr;
+ }
+ 
+ MeshAttr getCommonMesh(ArrayRef<TensorShardingAttr> operandShardings,
+                        ArrayRef<TensorShardingAttr> resultsShardings,
+                        Operation* op) {
+-  return getCommonMeshAndRef(operandShardings, resultsShardings, op,
+-                             /*ignoreDeviceIds=*/false)
+-      .second;
++  return getCommonMesh(operandShardings, resultsShardings,
++                       SymbolTable(op->getParentOfType<ModuleOp>()));
+ }
+ 
+ std::optional<StringRef> getCommonMeshName(
+diff --git a/shardy/dialect/sdy/ir/utils.h b/shardy/dialect/sdy/ir/utils.h
+index 3b56a553..22e4d9cc 100644
+--- a/shardy/dialect/sdy/ir/utils.h
++++ b/shardy/dialect/sdy/ir/utils.h
+@@ -195,12 +195,6 @@ Attribute getCommonMeshOrRef(ArrayRef<TensorShardingAttr> operandShardings,
+ MeshAttr getCommonMesh(ArrayRef<TensorShardingAttr> shardings,
+                        const SymbolTable& symbolTable);
+ 
+-// Returns the common `MeshAttr` bound by all the `TensorShardingAttr`s or
+-// nullptr if there is none.
+-//
+-// Ignores empty meshes unless all meshes are empty.
+-MeshAttr getCommonMesh(ArrayRef<TensorShardingAttr> shardings, Operation* op);
+-
+ // Returns the common `MeshAttr` bound by all the `TensorShardingAttr`s or
+ // nullptr if there is none.
+ //
+diff --git a/shardy/dialect/sdy/ir/utils_test.cc b/shardy/dialect/sdy/ir/utils_test.cc
+index f76d2b16..8dbaf619 100644
+--- a/shardy/dialect/sdy/ir/utils_test.cc
++++ b/shardy/dialect/sdy/ir/utils_test.cc
+@@ -81,27 +81,11 @@ class UtilsTest : public ShardyTestBase {
+     return symbolTableCollection.getSymbolTable(moduleOp.get());
+   }
+ 
+-  Operation* getModuleOp() { return moduleOp.get(); }
+-  Operation* getNestedOp() {
+-    return moduleOp->lookupSymbol<func::FuncOp>("main");
+-  }
+-
+  private:
+   OwningOpRef<ModuleOp> moduleOp;
+   mlir::SymbolTableCollection symbolTableCollection;
+ };
+ 
+-TEST_F(UtilsTest, GetCommonMesh_SingleListWithOperation) {
+-  EXPECT_NE(getCommonMesh({createTensorSharding("mesh_ab_23")}, getModuleOp()),
+-            nullptr);
+-  EXPECT_EQ(
+-      getCommonMesh({createTensorSharding("mesh_ab_23")}, getModuleOp()),
+-      getCommonMesh({createTensorSharding("mesh_ab_23")}, getSymbolTable()));
+-  EXPECT_EQ(
+-      getCommonMesh({createTensorSharding("mesh_ab_23")}, getNestedOp()),
+-      getCommonMesh({createTensorSharding("mesh_ab_23")}, getSymbolTable()));
+-}
+-
+ TEST_F(UtilsTest, GetCommonMeshName_AllSame) {
+   EXPECT_EQ(getCommonMeshName({createTensorSharding("mesh_ab_23")},
+                               {createTensorSharding("mesh_ab_23")},
+diff --git a/shardy/dialect/sdy/ir/verifiers.cc b/shardy/dialect/sdy/ir/verifiers.cc
+index 09f70f28..65c6212b 100644
+--- a/shardy/dialect/sdy/ir/verifiers.cc
++++ b/shardy/dialect/sdy/ir/verifiers.cc
+@@ -1186,7 +1186,8 @@ LogicalResult verifyPropagationEdgesShardingAttr(
+         "expected propagation edges attr to reference a sharding.");
+   }
+ 
+-  MeshAttr mesh = getCommonMesh(shardings, op);
++  MeshAttr mesh =
++      getCommonMesh(shardings, SymbolTable(op->getParentOfType<ModuleOp>()));
+   if (!mesh) {
+     return op->emitOpError(
+         "expected a common mesh for propagation edges attr.");
+@@ -1195,7 +1196,12 @@ LogicalResult verifyPropagationEdgesShardingAttr(
+   for (PropagationOneStepAttr propagationEdge : propagationEdges) {
+     for (AxisToPropagationDetailsAttr axisEntry :
+          propagationEdge.getAxisEntries()) {
+-      if (!mesh.hasAxis(axisEntry.getAxisName().getName())) {
++      bool meshContainsAxis = false;
++      if (mesh.hasAxis(axisEntry.getAxisName().getName())) {
++        meshContainsAxis = true;
++      }
++
++      if (!meshContainsAxis) {
+         return op->emitOpError("expected axis ref to be in one of the meshes");
+       }
+ 
+@@ -1284,8 +1290,8 @@ LogicalResult NamedComputationOp::verifySymbolUses(
+ 
+   // TODO(pxy): remove this once the `ShardableDataFlowOpInterface` is verified.
+   // Verify the in/out shardings.
+-  const SymbolTable& symbolTable = symbolTableCollection.getSymbolTable(
+-      getOperation()->getParentOfType<ModuleOp>());
++  const SymbolTable& symbolTable =
++      SymbolTable(getOperation()->getParentOfType<ModuleOp>());
+   if (inShardings && failed(verifyTensorShardingPerValueAttr(
+                          *inShardings, getOperandTypes(), *this, symbolTable,
+                          [this](StringRef msg) {
+diff --git a/shardy/dialect/sdy/transforms/export/test/insert_explicit_reshards/data_flow_ops.mlir b/shardy/dialect/sdy/transforms/export/test/insert_explicit_reshards/data_flow_ops.mlir
+index f287d6de..f2891d04 100644
+--- a/shardy/dialect/sdy/transforms/export/test/insert_explicit_reshards/data_flow_ops.mlir
++++ b/shardy/dialect/sdy/transforms/export/test/insert_explicit_reshards/data_flow_ops.mlir
+@@ -48,25 +48,6 @@ func.func @case_argument_used_outside_block(%arg0: tensor<210xf32> {sdy.sharding
+   return %3 : tensor<210xf32>
+ }
+ 
+-// CHECK-LABEL: func @if
+-func.func @if(%arg0: tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x":(1)2}]>}, %arg1: tensor<i1>) -> (tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}]>}) {
+-  %0 = "stablehlo.if"(%arg1) ({
+-    // CHECK: %[[RESHARD:.*]] = sdy.reshard %arg0 <@mesh, [{"x"}]> : tensor<210xf32>
+-    // CHECK-NEXT: stablehlo.abs %[[RESHARD]]
+-    %2 = stablehlo.abs %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x"}]>]>} : tensor<210xf32>
+-    // CHECK: %[[RESHARD:.*]] = sdy.reshard %{{.*}} <@mesh, [{"y"}]> : tensor<210xf32>
+-    // CHECK-NEXT: stablehlo.return %[[RESHARD]] : tensor<210xf32>
+-    stablehlo.return %2 : tensor<210xf32>
+-  }, {
+-    %2 = stablehlo.abs %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x":(1)2}]>]>} : tensor<210xf32>
+-    // CHECK: %[[RESHARD:.*]] = sdy.reshard %{{.*}} <@mesh, [{"y"}]> : tensor<210xf32>
+-    // CHECK-NEXT: stablehlo.return %[[RESHARD]] : tensor<210xf32>
+-    stablehlo.return %2 : tensor<210xf32>
+-  }) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>} : (tensor<i1>) -> tensor<210xf32>
+-  %1 = stablehlo.negate %0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>} : tensor<210xf32>
+-  return %1 : tensor<210xf32>
+-}
+-
+ // CHECK-LABEL: func @while
+ func.func @while(%arg0: tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}]>}) -> (tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}]>}) {
+   %c = stablehlo.constant dense<0> : tensor<i32>
+diff --git a/shardy/dialect/sdy/transforms/import/test/add_data_flow_edges.mlir b/shardy/dialect/sdy/transforms/import/test/add_data_flow_edges.mlir
+index 6b080b16..bd472047 100644
+--- a/shardy/dialect/sdy/transforms/import/test/add_data_flow_edges.mlir
++++ b/shardy/dialect/sdy/transforms/import/test/add_data_flow_edges.mlir
+@@ -55,24 +55,6 @@ func.func @case_token_result_skipped(
+   return %0#0, %0#1 : !stablehlo.token, tensor<4xi64>
+ }
+ 
+-// CHECK-LABEL: func @if
+-func.func @if(%arg0: tensor<i1>, %arg1: tensor<8xi64>, %arg2: tensor<8xi64>,
+-              %arg3: tensor<8xi64>, %arg4: tensor<8xi64>)
+-    -> (tensor<8xi64>, tensor<8xi64>) {
+-  // CHECK-NEXT: %[[IF:.*]]:2 = "stablehlo.if"(%arg0)
+-  // CHECK:      %[[EDGE_1:.*]] = sdy.data_flow_edge %[[IF]]#0 : tensor<8xi64>
+-  // CHECK-NEXT: %[[EDGE_2:.*]] = sdy.data_flow_edge %[[IF]]#1 : tensor<8xi64>
+-  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[EDGE_2]], %[[EDGE_2]]
+-  // CHECK-NEXT: return %[[EDGE_1]], %[[ADD]]
+-  %0:2 = "stablehlo.if"(%arg0) ({
+-    stablehlo.return %arg1, %arg2 : tensor<8xi64>, tensor<8xi64>
+-  }, {
+-    stablehlo.return %arg3, %arg4 : tensor<8xi64>, tensor<8xi64>
+-  }) : (tensor<i1>) -> (tensor<8xi64>, tensor<8xi64>)
+-  %1 = stablehlo.add %0#1, %0#1 : tensor<8xi64>
+-  return %0#0, %1 : tensor<8xi64>, tensor<8xi64>
+-}
+-
+ // CHECK-LABEL: func @optimization_barrier
+ func.func @optimization_barrier(%arg0: tensor<32x96xf32>, %arg1: tensor<32x96xf32>)
+     -> (tensor<32x96xf32>, tensor<32x96xf32>) {
+diff --git a/shardy/dialect/sdy/transforms/import/test/apply_sharding_constraints.mlir b/shardy/dialect/sdy/transforms/import/test/apply_sharding_constraints.mlir
+index 445f25fd..11d064d5 100644
+--- a/shardy/dialect/sdy/transforms/import/test/apply_sharding_constraints.mlir
++++ b/shardy/dialect/sdy/transforms/import/test/apply_sharding_constraints.mlir
+@@ -652,275 +652,6 @@ func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+ 
+ // -----
+ 
+-// A flat call graph (@main -> @foo) without func data flow edges, applying a
+-// single sharding constraint to the argument and call result. Without func data
+-// flow edges, closed shardings are attached directly to the function argument
+-// and call result and other uses are not replaced with the sharding constraint.
+-sdy.mesh @mesh = <["a"=2, "b"=2]>
+-
+-// CHECK-LABEL: func @foo(
+-// CHECK-SAME:    %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}, {}]>}) -> tensor<8x8xf32> {
+-func.func @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %arg0, %arg0 : tensor<8x8xf32>
+-  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %arg0 <@mesh, [{"a"}, {}]> {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=8, j=8}>} : tensor<8x8xf32>
+-  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
+-  %0 = stablehlo.add %arg0, %arg0 : tensor<8x8xf32>
+-  %1 = sdy.sharding_constraint %arg0 <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
+-  return %0 : tensor<8x8xf32>
+-}
+-
+-// CHECK-LABEL: func @main
+-func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  // CHECK-NEXT: %[[CALL:.*]] = call @foo(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"b"}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[CALL]], %[[CALL]] : tensor<8x8xf32>
+-  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %[[CALL]] <@mesh, [{}, {"b"}]> {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=8, j=8}>} : tensor<8x8xf32>
+-  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
+-  %0 = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  %1 = stablehlo.add %0, %0 : tensor<8x8xf32>
+-  %2 = sdy.sharding_constraint %0 <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
+-  return %1 : tensor<8x8xf32>
+-}
+-
+-// -----
+-
+-// A flat call graph (@main -> @foo) applying a single sharding constraint to
+-// the func data flow edges. The func data flow edges is left unsharded, the
+-// constraint is moved right after the edge, and other uses are replaced with
+-// the constraint.
+-sdy.mesh @mesh = <["a"=2, "b"=2]>
+-
+-// CHECK-LABEL: func @foo
+-func.func @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE:.*]] = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
+-  // CHECK-NOT: sdy.sharding
+-  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %[[FUNC_DATA_FLOW_EDGE]] <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
+-  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[WSC]], %[[WSC]] : tensor<8x8xf32>
+-  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
+-  %0 = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
+-  %1 = stablehlo.add %0, %0 : tensor<8x8xf32>
+-  %2 = sdy.sharding_constraint %0 <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
+-  return %1 : tensor<8x8xf32>
+-}
+-
+-// CHECK-LABEL: func @main
+-func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  // CHECK-NEXT: %[[CALL:.*]] = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE:.*]] = sdy.func_data_flow_edge %[[CALL]] : tensor<8x8xf32>
+-  // CHECK-NOT: sdy.sharding
+-  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %[[FUNC_DATA_FLOW_EDGE]] <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
+-  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[WSC]], %[[WSC]] : tensor<8x8xf32>
+-  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
+-  %0 = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  %1 = sdy.func_data_flow_edge %0 : tensor<8x8xf32>
+-  %2 = stablehlo.add %1, %1 : tensor<8x8xf32>
+-  %3 = sdy.sharding_constraint %1 <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
+-  return %2 : tensor<8x8xf32>
+-}
+-
+-// -----
+-
+-// A non-flat call graph (@main -> @foo -> @bar and @main -> @bar) without
+-// func data flow edges, applying a single sharding constraint to the argument
+-// and the call results. Without func data flow edges, closed shardings are
+-// attached directly to the argument and both results, and other uses are not
+-// replaced with the sharding constraint.
+-sdy.mesh @mesh = <["a"=2, "b"=2]>
+-
+-// CHECK-LABEL: func @bar(
+-// CHECK-SAME:    %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}, {}]>}) -> tensor<8x8xf32> {
+-func.func @bar(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %arg0, %arg0 : tensor<8x8xf32>
+-  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %arg0 <@mesh, [{"a"}, {}]> {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=8, j=8}>} : tensor<8x8xf32>
+-  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
+-  %0 = stablehlo.add %arg0, %arg0 : tensor<8x8xf32>
+-  %1 = sdy.sharding_constraint %arg0 <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
+-  return %0 : tensor<8x8xf32>
+-}
+-
+-// CHECK-LABEL: func @foo
+-func.func @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  // CHECK-NEXT: %[[CALL:.*]] = call @bar(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"b"}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[CALL]], %[[CALL]] : tensor<8x8xf32>
+-  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %[[CALL]] <@mesh, [{}, {"b"}]> {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=8, j=8}>} : tensor<8x8xf32>
+-  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
+-  %0 = call @bar(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  %1 = stablehlo.add %0, %0 : tensor<8x8xf32>
+-  %2 = sdy.sharding_constraint %0 <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
+-  return %1 : tensor<8x8xf32>
+-}
+-
+-// CHECK-LABEL: func @main
+-func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  // CHECK-NEXT: %[[CALL_0:.*]] = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  // CHECK-NEXT: %[[CALL_1:.*]] = call @bar(%[[CALL_0]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {"b"}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[CALL_1]], %[[CALL_1]] : tensor<8x8xf32>
+-  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %[[CALL_1]] <@mesh, [{"a"}, {"b"}]> {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=8, j=8}>} : tensor<8x8xf32>
+-  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
+-  %0 = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  %1 = call @bar(%0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  %2 = stablehlo.add %1, %1 : tensor<8x8xf32>
+-  %3 = sdy.sharding_constraint %1 <@mesh, [{"a"}, {"b"}]> : tensor<8x8xf32>
+-  return %2 : tensor<8x8xf32>
+-}
+-
+-// -----
+-
+-// A non-flat call graph (@main -> @foo -> @bar and @main -> @bar) applying a
+-// single sharding constraint to the func data flow edges on @bar's argument
+-// and the call results in @foo and @main. Each func data flow edges is left
+-// unsharded, the constraint is moved right after the edge, and other uses are
+-// replaced with the constraint.
+-sdy.mesh @mesh = <["a"=2, "b"=2]>
+-
+-// CHECK-LABEL: func @bar
+-func.func @bar(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE:.*]] = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
+-  // CHECK-NOT: sdy.sharding
+-  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %[[FUNC_DATA_FLOW_EDGE]] <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
+-  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[WSC]], %[[WSC]] : tensor<8x8xf32>
+-  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
+-  %0 = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
+-  %1 = stablehlo.add %0, %0 : tensor<8x8xf32>
+-  %2 = sdy.sharding_constraint %0 <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
+-  return %1 : tensor<8x8xf32>
+-}
+-
+-// CHECK-LABEL: func @foo
+-func.func @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE_0:.*]] = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
+-  // CHECK-NEXT: %[[CALL:.*]] = call @bar(%[[FUNC_DATA_FLOW_EDGE_0]])
+-  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE_1:.*]] = sdy.func_data_flow_edge %[[CALL]] : tensor<8x8xf32>
+-  // CHECK-NOT: sdy.sharding
+-  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %[[FUNC_DATA_FLOW_EDGE_1]] <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
+-  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[WSC]], %[[WSC]] : tensor<8x8xf32>
+-  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
+-  %0 = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
+-  %1 = call @bar(%0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  %2 = sdy.func_data_flow_edge %1 : tensor<8x8xf32>
+-  %3 = stablehlo.add %2, %2 : tensor<8x8xf32>
+-  %4 = sdy.sharding_constraint %2 <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
+-  return %3 : tensor<8x8xf32>
+-}
+-
+-// CHECK-LABEL: func @main
+-func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  // CHECK-NEXT: %[[CALL_0:.*]] = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE_0:.*]] = sdy.func_data_flow_edge %[[CALL_0]] : tensor<8x8xf32>
+-  // CHECK-NEXT: %[[CALL_1:.*]] = call @bar(%[[FUNC_DATA_FLOW_EDGE_0]]) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE_1:.*]] = sdy.func_data_flow_edge %[[CALL_1]] : tensor<8x8xf32>
+-  // CHECK-NOT: sdy.sharding
+-  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %[[FUNC_DATA_FLOW_EDGE_1]] <@mesh, [{"a"}, {"b"}]> : tensor<8x8xf32>
+-  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[WSC]], %[[WSC]] : tensor<8x8xf32>
+-  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
+-  %0 = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  %1 = sdy.func_data_flow_edge %0 : tensor<8x8xf32>
+-  %2 = call @bar(%1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  %3 = sdy.func_data_flow_edge %2 : tensor<8x8xf32>
+-  %4 = stablehlo.add %3, %3 : tensor<8x8xf32>
+-  %5 = sdy.sharding_constraint %3 <@mesh, [{"a"}, {"b"}]> : tensor<8x8xf32>
+-  return %4 : tensor<8x8xf32>
+-}
+-
+-// -----
+-
+-// Non-flat call graph (@main -> @foo -> @bar and @main -> @bar) without
+-// func data flow edges, applying a chain of two sharding constraints
+-// to stablehlo.add in @bar and @foo. Verifies that the head constraint's
+-// sharding is attached to `stablehlo.add` and uses after the chain (including
+-// the func returns in @bar and @foo) are replaced with the tail of the chain.
+-sdy.mesh @mesh = <["a"=2, "b"=2]>
+-
+-// CHECK-LABEL: func @bar
+-func.func @bar(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %arg0, %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}]>]>}
+-  // CHECK-NEXT: %[[WSC_0:.*]] = sdy.sharding_constraint %[[ADD]] <@mesh, [{"a"}, {}]>
+-  // CHECK-NEXT: %[[WSC_1:.*]] = sdy.sharding_constraint %[[WSC_0]] <@mesh, [{}, {"b"}]>
+-  // CHECK-NEXT: return %[[WSC_1]] : tensor<8x8xf32>
+-  %0 = stablehlo.add %arg0, %arg0 : tensor<8x8xf32>
+-  %1 = sdy.sharding_constraint %0 <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
+-  %2 = sdy.sharding_constraint %1 <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
+-  return %0 : tensor<8x8xf32>
+-}
+-
+-// CHECK-LABEL: func @foo
+-func.func @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  // CHECK-NEXT: %[[CALL:.*]] = call @bar(%arg0)
+-  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[CALL]], %[[CALL]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}]>]>}
+-  // CHECK-NEXT: %[[WSC_0:.*]] = sdy.sharding_constraint %[[ADD]] <@mesh, [{"a"}, {}]>
+-  // CHECK-NEXT: %[[WSC_1:.*]] = sdy.sharding_constraint %[[WSC_0]] <@mesh, [{}, {"b"}]>
+-  // CHECK-NEXT: return %[[WSC_1]] : tensor<8x8xf32>
+-  %0 = call @bar(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  %1 = stablehlo.add %0, %0 : tensor<8x8xf32>
+-  %2 = sdy.sharding_constraint %1 <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
+-  %3 = sdy.sharding_constraint %2 <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
+-  return %1 : tensor<8x8xf32>
+-}
+-
+-// CHECK-LABEL: func @main
+-func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  // CHECK-NEXT: %[[CALL_0:.*]] = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  // CHECK-NEXT: %[[CALL_1:.*]] = call @bar(%[[CALL_0]]) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  // CHECK-NEXT: return %[[CALL_1]] : tensor<8x8xf32>
+-  %0 = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  %1 = call @bar(%0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  return %1 : tensor<8x8xf32>
+-}
+-
+-// -----
+-
+-// A non-flat call graph (@main -> @foo -> @bar and @main -> @bar) with func
+-// data flow edges on function arguments and call results, applying a chain of
+-// two sharding constraints to stablehlo.add in @bar and @foo.
+-sdy.mesh @mesh = <["a"=2, "b"=2]>
+-
+-// CHECK-LABEL: func @bar
+-func.func @bar(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE:.*]] = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
+-  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[FUNC_DATA_FLOW_EDGE]], %[[FUNC_DATA_FLOW_EDGE]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}]>]>}
+-  // CHECK-NEXT: %[[WSC_0:.*]] = sdy.sharding_constraint %[[ADD]] <@mesh, [{"a"}, {}]>
+-  // CHECK-NEXT: %[[WSC_1:.*]] = sdy.sharding_constraint %[[WSC_0]] <@mesh, [{}, {"b"}]>
+-  // CHECK-NEXT: return %[[WSC_1]] : tensor<8x8xf32>
+-  %0 = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
+-  %1 = stablehlo.add %0, %0 : tensor<8x8xf32>
+-  %2 = sdy.sharding_constraint %1 <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
+-  %3 = sdy.sharding_constraint %2 <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
+-  return %1 : tensor<8x8xf32>
+-}
+-
+-// CHECK-LABEL: func @foo
+-func.func @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE_0:.*]] = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
+-  // CHECK-NEXT: %[[CALL:.*]] = call @bar(%[[FUNC_DATA_FLOW_EDGE_0]])
+-  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE_1:.*]] = sdy.func_data_flow_edge %[[CALL]] : tensor<8x8xf32>
+-  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[FUNC_DATA_FLOW_EDGE_1]], %[[FUNC_DATA_FLOW_EDGE_1]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}]>]>}
+-  // CHECK-NEXT: %[[WSC_0:.*]] = sdy.sharding_constraint %[[ADD]] <@mesh, [{"a"}, {}]>
+-  // CHECK-NEXT: %[[WSC_1:.*]] = sdy.sharding_constraint %[[WSC_0]] <@mesh, [{}, {"b"}]>
+-  // CHECK-NEXT: return %[[WSC_1]] : tensor<8x8xf32>
+-  %0 = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
+-  %1 = call @bar(%0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  %2 = sdy.func_data_flow_edge %1 : tensor<8x8xf32>
+-  %3 = stablehlo.add %2, %2 : tensor<8x8xf32>
+-  %4 = sdy.sharding_constraint %3 <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
+-  %5 = sdy.sharding_constraint %4 <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
+-  return %3 : tensor<8x8xf32>
+-}
+-
+-// CHECK-LABEL: func @main
+-func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
+-  // CHECK-NEXT: %[[CALL_0:.*]] = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE_0:.*]] = sdy.func_data_flow_edge %[[CALL_0]] : tensor<8x8xf32>
+-  // CHECK-NEXT: %[[CALL_1:.*]] = call @bar(%[[FUNC_DATA_FLOW_EDGE_0]]) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE_1:.*]] = sdy.func_data_flow_edge %[[CALL_1]] : tensor<8x8xf32>
+-  // CHECK-NEXT: return %[[FUNC_DATA_FLOW_EDGE_1]] : tensor<8x8xf32>
+-  %0 = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  %1 = sdy.func_data_flow_edge %0 : tensor<8x8xf32>
+-  %2 = call @bar(%1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
+-  %3 = sdy.func_data_flow_edge %2 : tensor<8x8xf32>
+-  return %3 : tensor<8x8xf32>
+-}
+-
+-// -----
+-
+ sdy.mesh @mesh = <["a"=2, "b"=2]>
+ 
+ // CHECK-LABEL: func @main
+diff --git a/shardy/dialect/sdy/transforms/propagation/op_sharding_rule_registry.cc b/shardy/dialect/sdy/transforms/propagation/op_sharding_rule_registry.cc
+index dc069471..e88aa109 100644
+--- a/shardy/dialect/sdy/transforms/propagation/op_sharding_rule_registry.cc
++++ b/shardy/dialect/sdy/transforms/propagation/op_sharding_rule_registry.cc
+@@ -1321,7 +1321,7 @@ OpShardingRuleAttr createOpShardingRule(Operation* op,
+             PropagationBarrierOp, ShardableDataFlowOpInterface, ShardingGroupOp,
+             ReshardOp, stablehlo::AfterAllOp, stablehlo::CaseOp,
+             stablehlo::ConstantOp, stablehlo::CreateTokenOp,
+-            stablehlo::GetTupleElementOp, stablehlo::IfOp, stablehlo::InfeedOp,
++            stablehlo::GetTupleElementOp, stablehlo::InfeedOp,
+             stablehlo::IotaOp, stablehlo::OutfeedOp,
+             stablehlo::OptimizationBarrierOp, stablehlo::PartitionIdOp,
+             stablehlo::RecvOp, stablehlo::SendOp, stablehlo::WhileOp>(
+diff --git a/shardy/dialect/sdy/transforms/propagation/sharding_group_map.cc b/shardy/dialect/sdy/transforms/propagation/sharding_group_map.cc
+index f61c0dd3..2035990f 100644
+--- a/shardy/dialect/sdy/transforms/propagation/sharding_group_map.cc
++++ b/shardy/dialect/sdy/transforms/propagation/sharding_group_map.cc
+@@ -15,10 +15,12 @@ limitations under the License.
+ 
+ #include "shardy/dialect/sdy/transforms/propagation/sharding_group_map.h"
+ 
++#include <algorithm>
++#include <cassert>
+ #include <cstdint>
+ 
+ #include "llvm/ADT/DenseMap.h"
+-#include "llvm/ADT/MapVector.h"
++#include "llvm/ADT/STLExtras.h"
+ #include "llvm/Support/ErrorHandling.h"
+ #include "mlir/IR/BuiltinOps.h"
+ #include "mlir/IR/Diagnostics.h"
+@@ -34,19 +36,21 @@ namespace sdy {
+ 
+ ShardingGroupMap::ShardingGroupMap(ModuleOp moduleOp) {
+   moduleOp.walk([&](ShardingGroupOp op) {
+-    int64_t groupId = op.getGroupId();
++    // After canonicalization all group ids will take distinct values in the
++    // range 0,1,...,N. Because of this we can directly index these group ids
++    // into the shardingGroupToValues vector (resizing when necessary).
++    shardingGroupToValues.resize(
++        std::max(op.getGroupId() + 1,
++                 static_cast<uint64_t>(shardingGroupToValues.size())));
+     // Each value can only map to one sharding group id after
+     // canonicalization.
+     auto [it, inserted] =
+-        valueToShardingGroup.try_emplace(op.getInput(), groupId);
+-    if (!inserted) {
+-      if (it->getSecond() != groupId) {
+-        llvm::report_fatal_error(
+-            "Value can only map to one sharding group id after import.");
+-      }
+-      return;
++        valueToShardingGroup.try_emplace(op.getInput(), op.getGroupId());
++    if (!inserted && it->getSecond() != op.getGroupId()) {
++      llvm::report_fatal_error(
++          "Value can only map to one sharding group id after import.");
+     }
+-    shardingGroupToValues[groupId].push_back(op.getInput());
++    shardingGroupToValues[op.getGroupId()].push_back(op.getInput());
+   });
+ }
+ 
+@@ -96,7 +100,7 @@ CommonSharding findCommonSharding(int64_t groupId, ValueRange groupMembers) {
+ }  // namespace
+ 
+ void ShardingGroupMap::syncGroupMemberShardings(ModuleOp module) {
+-  for (auto& [groupId, groupMembers] : shardingGroupToValues) {
++  for (auto [groupId, groupMembers] : llvm::enumerate(shardingGroupToValues)) {
+     auto [sharding, hasConflict] = findCommonSharding(groupId, groupMembers);
+     if (!sharding) {
+       continue;
+@@ -130,7 +134,8 @@ void ShardingGroupMap::syncGroupMemberShardings(ModuleOp module) {
+ ValueRange ShardingGroupMap::getGroupMembers(const Value& value) const {
+   if (auto it = valueToShardingGroup.find(value);
+       it != valueToShardingGroup.end()) {
+-    return shardingGroupToValues.find(it->getSecond())->second;
++    int64_t shardingGroupId = it->getSecond();
++    return shardingGroupToValues[shardingGroupId];
+   }
+   return {};
+ }
+diff --git a/shardy/dialect/sdy/transforms/propagation/sharding_group_map.h b/shardy/dialect/sdy/transforms/propagation/sharding_group_map.h
+index b924ca1a..90b38bd4 100644
+--- a/shardy/dialect/sdy/transforms/propagation/sharding_group_map.h
++++ b/shardy/dialect/sdy/transforms/propagation/sharding_group_map.h
+@@ -19,7 +19,6 @@ limitations under the License.
+ #include <cstdint>
+ 
+ #include "llvm/ADT/DenseMap.h"
+-#include "llvm/ADT/MapVector.h"
+ #include "mlir/IR/BuiltinOps.h"
+ #include "mlir/IR/Value.h"
+ #include "mlir/IR/ValueRange.h"
+@@ -46,7 +45,7 @@ class ShardingGroupMap {
+   ValueRange getGroupMembers(const Value& value) const;
+ 
+  private:
+-  llvm::MapVector<int64_t, SmallVector<Value>> shardingGroupToValues;
++  SmallVector<SmallVector<Value>> shardingGroupToValues;
+   llvm::SmallDenseMap<Value, int64_t> valueToShardingGroup;
+ };
+ 
+diff --git a/shardy/dialect/sdy/transforms/propagation/test/sharding_group_propagation.mlir b/shardy/dialect/sdy/transforms/propagation/test/sharding_group_propagation.mlir
+index 36ae5341..b9edf866 100644
+--- a/shardy/dialect/sdy/transforms/propagation/test/sharding_group_propagation.mlir
++++ b/shardy/dialect/sdy/transforms/propagation/test/sharding_group_propagation.mlir
+@@ -131,31 +131,6 @@ func.func @different_group_ids_greater_group_id_first(
+   return %arg0, %arg1, %arg2, %arg3 : tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>
+ }
+ 
+-// -----
+-sdy.mesh @mesh = <["a"=2, "b"=2]>
+-
+-// CHECK-LABEL: func @undensified_negative_and_large_group_ids
+-// CHECK-SAME:    %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a", ?}, {?}]>},
+-// CHECK-SAME:    %arg1: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a", ?}, {?}]>},
+-// CHECK-SAME:    %arg2: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{?}, {"b", ?}]>},
+-// CHECK-SAME:    %arg3: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{?}, {"b", ?}]>})
+-// CHECK-SAME:    -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a", ?}, {?}]>},
+-// CHECK-SAME:        tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a", ?}, {?}]>},
+-// CHECK-SAME:        tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{?}, {"b", ?}]>},
+-// CHECK-SAME:        tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{?}, {"b", ?}]>})
+-func.func @undensified_negative_and_large_group_ids(
+-  %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a", ?}, {?}]>},
+-  %arg1: tensor<8x8xf32>,
+-  %arg2: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{?}, {"b", ?}]>},
+-  %arg3: tensor<8x8xf32>)
+-   -> (tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>) {
+-  sdy.sharding_group %arg0 group_id=-1 : tensor<8x8xf32>
+-  sdy.sharding_group %arg1 group_id=-1 : tensor<8x8xf32>
+-  sdy.sharding_group %arg2 group_id=100000000000 : tensor<8x8xf32>
+-  sdy.sharding_group %arg3 group_id=100000000000 : tensor<8x8xf32>
+-  return %arg0, %arg1, %arg2, %arg3 : tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>
+-}
+-
+ // -----
+ 
+ // Tests for initial member sharding sync
+diff --git a/third_party/llvm/generated.patch b/third_party/llvm/generated.patch
+index e9c206ce..509398da 100644
+--- a/third_party/llvm/generated.patch
++++ b/third_party/llvm/generated.patch
+@@ -1,12 +1 @@
+ Auto generated patch. Do not edit or delete it, even if empty.
+-diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll b/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll
+---- a/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll
+-+++ b/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll
+-@@ -1,5 +1,7 @@
+- ; Checks if conversion is only inserted for the spilled register
+- ; instead of 2 conversions for the W register
+-+
+-+; REQUIRES: asserts
+- ; RUN: llc -O2 -march=hexagon -mcpu=hexagonv81 -enable-xqf-gen=true \
+- ; RUN: -hexagon-qfloat-mode=ieee -mattr=+hvxv81,+hvx-length128B \
+- ; RUN: -enable-postra-xqf-check -debug-only=handle-qfp 2>&1 < %s -o /dev/null | FileCheck %s
+diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
+index 68a14533..e42d7dbd 100644
+--- a/third_party/llvm/workspace.bzl
++++ b/third_party/llvm/workspace.bzl
+@@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
+ 
+ def repo(name):
+     """Imports LLVM."""
+-    LLVM_COMMIT = "b08e005995ed29a2f89aa9a3f599a9d45db65523"
+-    LLVM_SHA256 = "3f4f1765816e22a21ac41dbf71d91d519b0b371b05d7e33ef2a0a862c838420c"
++    LLVM_COMMIT = "2e0d9e12b36e853d9b108c4da09509cd54d1f3f9"
++    LLVM_SHA256 = "8e4a851780af109d623fa6d6db9b1e62d7126272d17268fed1b149f0e737c1d7"
+ 
+     tf_http_archive(
+         name = name,
diff --git a/third_party/xla/third_party/shardy/workspace.bzl b/third_party/xla/third_party/shardy/workspace.bzl
index 659b180a7ba..ba16a772a07 100644
--- a/third_party/xla/third_party/shardy/workspace.bzl
+++ b/third_party/xla/third_party/shardy/workspace.bzl
@@ -18,8 +18,8 @@
 load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 
 def repo():
-    SHARDY_COMMIT = "b5950641f8b5cbbebdf437e39197e1a170611587"
-    SHARDY_SHA256 = "ff228c5220c99e8de7fa8a480104ffaf1f6801c3b2641b05b6f04bc626779f12"
+    SHARDY_COMMIT = "a29d5efe295b7834e3b644d49e09211f1af0d28a"
+    SHARDY_SHA256 = "b28a57c907bd6b74b849a0110b8e4081358daca62336e72d418e0770895d1b7c"
 
     tf_http_archive(
         name = "shardy",

From 5824554f773f05a2ac664020163f98617b2aa5ac Mon Sep 17 00:00:00 2001
From: Jian Cai <jiancai@google.com>
Date: Mon, 12 Oct 2026 12:03:24 -0700
Subject: [PATCH 0614/1000] Add unit tests for XLA original value recovery
 across compiler passes.

Previously, hlo_original_value_reconstructor_test.cc only evaluated recovery computations from hand-written HLO recovery tables, leaving pass-level generation untested.

Add original_value_recovery_test.cc and hlo_original_value_analysis_test.cc to test recovery table generation and chaining across SpmdPartitioner, AlgebraicSimplifier, and CallInliner, and fix HloOriginalValueAnalysis::Create so recovery entries mapped from non-placeholder instructions are discovered. END_PUBLIC

Also add original_value_recovery_e2e_test.py to verify end-to-end JAX-on-TPU compilation, runtime HLO output logging, and offline tensor reconstruction against NumPy references, along with original_value_recovery_test_helper_nb (testonly = True) so Python tests can inspect full reconstructed tensor arrays and recovery table metadata (whereas comparison_tool_bridge only exposes scalar summary statistics).

PiperOrigin-RevId: 998067449
---
 .../hlo/utils/hlo_original_value_analysis.cc  |   28 +-
 .../utils/hlo_original_value_analysis_test.cc |   50 +
 third_party/xla/xla/service/BUILD             |   23 +
 .../service/original_value_recovery_test.cc   | 1198 +++++++++++++++++
 4 files changed, 1288 insertions(+), 11 deletions(-)
 create mode 100644 third_party/xla/xla/service/original_value_recovery_test.cc

diff --git a/third_party/xla/xla/hlo/utils/hlo_original_value_analysis.cc b/third_party/xla/xla/hlo/utils/hlo_original_value_analysis.cc
index a64d96466c8..0cb276de967 100644
--- a/third_party/xla/xla/hlo/utils/hlo_original_value_analysis.cc
+++ b/third_party/xla/xla/hlo/utils/hlo_original_value_analysis.cc
@@ -270,17 +270,7 @@ HloOriginalValueAnalysis::Create(const HloModule* optimized_module,
           continue;
         }
 
-        if (absl::StrContains(oa.instruction_name, "__ovp")) {
-          std::vector<OriginalTensorInfo> results;
-          absl::flat_hash_set<OriginalArray> visited_placeholders;
-          BuildTransformationChain(oa, placeholder_to_recoverables,
-                                   optimized_key, optimized_tensor_sharding,
-                                   visited_placeholders, results);
-          for (auto& info : results) {
-            original_tensor_by_optimized_tensor_key[optimized_key].push_back(
-                std::move(info));
-          }
-        } else {
+        if (!absl::StrContains(oa.instruction_name, "__ovp")) {
           auto relative_key = RelativeScopedTensorKey::FromString(
               oa.instruction_name, oa.shape_index);
           HloOriginalValueAnalysis::OriginalTensorInfo info;
@@ -294,6 +284,22 @@ HloOriginalValueAnalysis::Create(const HloModule* optimized_module,
           original_tensor_by_optimized_tensor_key[optimized_key].push_back(
               std::move(info));
         }
+        // Always traverse `BuildTransformationChain` even when `oa` is not a
+        // synthetic placeholder (no "__ovp" suffix). Passes such as
+        // `AlgebraicSimplifier` record recovery table entries whose source is
+        // an existing un-renamed instruction (e.g., `{"reshape.1"} : {"x.1"}`
+        // when `reshape.1 = reshape(x.1)` is folded into a consumer), so `x.1`
+        // both recovers `{"x.1"}` directly and serves as the source for
+        // recovering `{"reshape.1"}`.
+        std::vector<OriginalTensorInfo> results;
+        absl::flat_hash_set<OriginalArray> visited_placeholders;
+        BuildTransformationChain(oa, placeholder_to_recoverables, optimized_key,
+                                 optimized_tensor_sharding,
+                                 visited_placeholders, results);
+        for (auto& info : results) {
+          original_tensor_by_optimized_tensor_key[optimized_key].push_back(
+              std::move(info));
+        }
       }
     }
   }
diff --git a/third_party/xla/xla/hlo/utils/hlo_original_value_analysis_test.cc b/third_party/xla/xla/hlo/utils/hlo_original_value_analysis_test.cc
index 5c70a322685..563b942439e 100644
--- a/third_party/xla/xla/hlo/utils/hlo_original_value_analysis_test.cc
+++ b/third_party/xla/xla/hlo/utils/hlo_original_value_analysis_test.cc
@@ -132,6 +132,56 @@ ENTRY main {
                                  SizeIs(1)))))));
 }
 
+// Verifies that an optimized instruction with a non-placeholder OriginalArray
+// (e.g., {"x.1"} without "__ovp") both recovers {"x.1"} directly and triggers
+// recovery table entries whose source is {"x.1"} (such as {"reshape.1"} :
+// {"x.1"}).
+TEST_F(HloOriginalValueAnalysisTest, RecoveryFromNonPlaceholderOriginalArray) {
+  std::string hlo_text = R"hlo(
+HloModule module, origin_recovery_table={
+  {"reshape.1"}: {"x.1"}, "HloModule rec ENTRY %e { %p = f32[8] parameter(0) ROOT %r = f32[1,8,1] reshape(%p) }"
+}
+ENTRY main {
+  %x.1 = f32[8] parameter(0), origin={{"x.1"}}
+  ROOT %broadcast.1 = f32[1,2,3,8,1] broadcast(%x.1), dimensions={3}, origin={{"broadcast.1"}}
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text));
+  ASSERT_OK_AND_ASSIGN(auto analysis,
+                       HloOriginalValueAnalysis::Create(module.get()));
+
+  auto x_orig = RelativeScopedTensorKey::FromString("x.1", {});
+  auto reshape_orig = RelativeScopedTensorKey::FromString("reshape.1", {});
+  auto broadcast_orig = RelativeScopedTensorKey::FromString("broadcast.1", {});
+
+  EXPECT_THAT(
+      analysis->original_tensor_by_optimized_tensor_key(),
+      UnorderedElementsAre(
+          Pair(TensorKey{"x.1", {}},
+               UnorderedElementsAre(
+                   AllOf(Field(&HloOriginalValueAnalysis::OriginalTensorInfo::
+                                   original_scoped_tensor_key,
+                               x_orig),
+                         Field(&HloOriginalValueAnalysis::OriginalTensorInfo::
+                                   recovery_modules,
+                               IsEmpty())),
+                   AllOf(Field(&HloOriginalValueAnalysis::OriginalTensorInfo::
+                                   original_scoped_tensor_key,
+                               reshape_orig),
+                         Field(&HloOriginalValueAnalysis::OriginalTensorInfo::
+                                   recovery_modules,
+                               SizeIs(1))))),
+          Pair(TensorKey{"broadcast.1", {}},
+               ElementsAre(
+                   AllOf(Field(&HloOriginalValueAnalysis::OriginalTensorInfo::
+                                   original_scoped_tensor_key,
+                               broadcast_orig),
+                         Field(&HloOriginalValueAnalysis::OriginalTensorInfo::
+                                   recovery_modules,
+                               IsEmpty()))))));
+}
+
 TEST_F(HloOriginalValueAnalysisTest, ChainedRecoveryModules) {
   constexpr absl::string_view hlo_string = R"hlo(
 HloModule chained_module, entry_computation_layout={(f32[4,8]{1,0})->f32[4,8]{1,0}}, origin_recovery_table={
diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index 4958417f3bc..d60312dc516 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -6209,6 +6209,29 @@ xla_cc_test(
     ],
 )
 
+xla_cc_test(
+    name = "original_value_recovery_test",
+    srcs = ["original_value_recovery_test.cc"],
+    deps = [
+        ":call_inliner",
+        "//xla:literal",
+        "//xla:literal_util",
+        "//xla:shape_util",
+        "//xla:xla_data_proto_cc",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+        "//xla/hlo/transforms/simplifiers:algebraic_simplifier",
+        "//xla/hlo/utils:hlo_original_value_analysis",
+        "//xla/hlo/utils:hlo_original_value_analyzer_utils",
+        "//xla/hlo/utils:hlo_original_value_reconstructor",
+        "//xla/hlo/utils:hlo_sharding_reconstruction_util",
+        "//xla/service/spmd:spmd_partitioner",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "while_loop_pipeline_unroller",
     srcs = ["while_loop_pipeline_unroller.cc"],
diff --git a/third_party/xla/xla/service/original_value_recovery_test.cc b/third_party/xla/xla/service/original_value_recovery_test.cc
new file mode 100644
index 00000000000..101ad80c1a9
--- /dev/null
+++ b/third_party/xla/xla/service/original_value_recovery_test.cc
@@ -0,0 +1,1198 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/container/flat_hash_map.h"
+#include "absl/strings/string_view.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/ir/hlo_original_value.h"
+#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
+#include "xla/hlo/transforms/simplifiers/algebraic_simplifier.h"
+#include "xla/hlo/utils/hlo_original_value_analysis.h"
+#include "xla/hlo/utils/hlo_original_value_analyzer_utils.h"
+#include "xla/hlo/utils/hlo_original_value_reconstructor.h"
+#include "xla/hlo/utils/hlo_sharding_reconstruction_util.h"
+#include "xla/literal.h"
+#include "xla/literal_util.h"
+#include "xla/service/call_inliner.h"
+#include "xla/service/spmd/spmd_partitioner.h"
+#include "xla/shape_util.h"
+#include "xla/xla_data.pb.h"
+
+namespace xla {
+namespace {
+
+using ::testing::ElementsAre;
+
+// Tests original value recovery table generation, chaining, and reconstruction
+// across XLA compiler passes:
+// - SpmdPartitioner: 1D/2D sharding, transposed device assignment, uneven
+//   tiling with slice recovery, partial replication (last_tile_dim_replicate),
+//   manual subgroups (manual_axes), multi-leaf tuple sharding, and mixed
+//   sharded/replicated tuple elements with non-empty ShapeIndex.
+// - AlgebraicSimplifier & Chained Passes: reshape + broadcast folding,
+//   AlgebraicSimplifier followed by SpmdPartitioner (partitioned=false and
+//   partitioned=true), and CallInliner followed by AlgebraicSimplifier
+//   preserving call hierarchy scopes.
+// - HloOriginalValueReconstructor: partitioned=true per-shard callbacks, mixed
+//   partitioned=false and partitioned=true debug attributes, and multi-
+//   controller unaddressable device early abort.
+class OriginalValueRecoveryTest : public HloHardwareIndependentTestBase {
+ protected:
+  // Helper to locate the optimized tensor key in the analysis that corresponds
+  // to an original instruction name (optionally scoped) and shape index.
+  AbsoluteScopedTensorKey FindOptimizedTensorKey(
+      const HloOriginalValueAnalysis& analysis, absl::string_view original_name,
+      const ShapeIndex& original_shape_index = {}) {
+    RelativeScopedTensorKey rel_key = RelativeScopedTensorKey::FromString(
+        original_name, original_shape_index);
+    auto it =
+        analysis.original_to_optimized_tensor_map().find(rel_key.tensor_key);
+    if (it == analysis.original_to_optimized_tensor_map().end()) {
+      return {};
+    }
+    for (const auto& [opt_key, info] : it->second) {
+      if (info->original_scoped_tensor_key == rel_key) {
+        return AbsoluteScopedTensorKey::Create(opt_key, {});
+      }
+    }
+    return {};
+  }
+};
+
+TEST_F(OriginalValueRecoveryTest, SpmdPartitioning1DSharding) {
+  constexpr absl::string_view hlo_string = R"hlo(
+HloModule module, entry_computation_layout={(f32[4,8]{1,0})->f32[4,8]{1,0}}, num_partitions=2
+
+ENTRY entry (p: f32[4,8]) -> f32[4,8] {
+  %p = f32[4,8]{1,0} parameter(0), sharding={devices=[2,1]<=[2]}
+  ROOT %neg = f32[4,8]{1,0} negate(%p), sharding={devices=[2,1]<=[2]}, origin={{"neg_origin"}}
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_TRUE(module->original_value_recovery_table().empty());
+
+  spmd::SpmdPartitionerOptions options;
+  options.allow_module_signature_change = true;
+  spmd::SpmdPartitioner partitioner(2, 1, options);
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(partitioner, module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_FALSE(module->original_value_recovery_table().empty());
+
+  ASSERT_OK_AND_ASSIGN(auto analysis,
+                       HloOriginalValueAnalysis::Create(module.get()));
+  auto analysis_shared =
+      std::shared_ptr<const HloOriginalValueAnalysis>(std::move(analysis));
+
+  std::optional<Literal> recovered_literal;
+  auto callback =
+      [&](const AbsoluteScopedTensorKey& original_tensor_key,
+          const OriginalArray& original_tensor,
+          std::shared_ptr<Literal> recovered_data,
+          const std::vector<HloModule::DebugAttributes>& debug_attributes,
+          int64_t manual_shard_id) {
+        if (original_tensor.instruction_name == "neg_origin" &&
+            recovered_data != nullptr) {
+          recovered_literal = recovered_data->Clone();
+        }
+      };
+
+  HloOriginalValueReconstructor reconstructor(analysis_shared, callback,
+                                              std::nullopt, module.get());
+
+  AbsoluteScopedTensorKey opt_key =
+      FindOptimizedTensorKey(*analysis_shared, "neg_origin");
+  ASSERT_FALSE(opt_key.tensor_key.instruction_name.empty());
+
+  // Shard 0: [2, 8] filled with 1.0f
+  ASSERT_OK_AND_ASSIGN(Literal literal0,
+                       Literal::Make(ShapeUtil::MakeShape(F32, {2, 8})));
+  literal0.PopulateWithValue(1.0f);
+  ShardTensor shard0 = {
+      /*logical_shard_id=*/0,
+      /*data=*/std::make_shared<Literal>(std::move(literal0))};
+
+  // Shard 1: [2, 8] filled with 2.0f
+  ASSERT_OK_AND_ASSIGN(Literal literal1,
+                       Literal::Make(ShapeUtil::MakeShape(F32, {2, 8})));
+  literal1.PopulateWithValue(2.0f);
+  ShardTensor shard1 = {
+      /*logical_shard_id=*/1,
+      /*data=*/std::make_shared<Literal>(std::move(literal1))};
+
+  ASSERT_OK(reconstructor.ProcessShardTensor(opt_key, std::move(shard0)));
+  EXPECT_FALSE(recovered_literal.has_value());
+
+  ASSERT_OK(reconstructor.ProcessShardTensor(opt_key, std::move(shard1)));
+  ASSERT_TRUE(recovered_literal.has_value());
+
+  EXPECT_THAT(recovered_literal->shape().dimensions(), ElementsAre(4, 8));
+  EXPECT_EQ(recovered_literal->Get<float>({0, 0}), 1.0f);
+  EXPECT_EQ(recovered_literal->Get<float>({1, 7}), 1.0f);
+  EXPECT_EQ(recovered_literal->Get<float>({2, 0}), 2.0f);
+  EXPECT_EQ(recovered_literal->Get<float>({3, 7}), 2.0f);
+}
+
+TEST_F(OriginalValueRecoveryTest, SpmdPartitioning2DSharding) {
+  constexpr absl::string_view hlo_string = R"hlo(
+HloModule module, entry_computation_layout={(f32[4,6]{1,0}, f32[4,6]{1,0})->f32[4,6]{1,0}}, num_partitions=4
+
+ENTRY entry (p0: f32[4,6], p1: f32[4,6]) -> f32[4,6] {
+  %p0 = f32[4,6]{1,0} parameter(0), sharding={devices=[2,2]<=[4]}
+  %p1 = f32[4,6]{1,0} parameter(1), sharding={devices=[2,2]<=[4]}
+  ROOT %add = f32[4,6]{1,0} add(%p0, %p1), sharding={devices=[2,2]<=[4]}, origin={{"add_origin"}}
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_TRUE(module->original_value_recovery_table().empty());
+
+  spmd::SpmdPartitionerOptions options;
+  options.allow_module_signature_change = true;
+  spmd::SpmdPartitioner partitioner(4, 1, options);
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(partitioner, module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_FALSE(module->original_value_recovery_table().empty());
+
+  ASSERT_OK_AND_ASSIGN(auto analysis,
+                       HloOriginalValueAnalysis::Create(module.get()));
+  auto analysis_shared =
+      std::shared_ptr<const HloOriginalValueAnalysis>(std::move(analysis));
+
+  std::optional<Literal> recovered_literal;
+  auto callback =
+      [&](const AbsoluteScopedTensorKey& original_tensor_key,
+          const OriginalArray& original_tensor,
+          std::shared_ptr<Literal> recovered_data,
+          const std::vector<HloModule::DebugAttributes>& debug_attributes,
+          int64_t manual_shard_id) {
+        if (original_tensor.instruction_name == "add_origin" &&
+            recovered_data != nullptr) {
+          recovered_literal = recovered_data->Clone();
+        }
+      };
+
+  HloOriginalValueReconstructor reconstructor(analysis_shared, callback,
+                                              std::nullopt, module.get());
+
+  AbsoluteScopedTensorKey opt_key =
+      FindOptimizedTensorKey(*analysis_shared, "add_origin");
+  ASSERT_FALSE(opt_key.tensor_key.instruction_name.empty());
+
+  // Shard shape is [2, 3].
+  // Device grid [2, 2] <= [4]:
+  // logical_shard_id 0 -> tile (0, 0): top-left
+  // logical_shard_id 1 -> tile (0, 1): top-right
+  // logical_shard_id 2 -> tile (1, 0): bottom-left
+  // logical_shard_id 3 -> tile (1, 1): bottom-right
+  for (int64_t shard_id = 0; shard_id < 4; ++shard_id) {
+    ASSERT_OK_AND_ASSIGN(Literal shard_data,
+                         Literal::Make(ShapeUtil::MakeShape(F32, {2, 3})));
+    shard_data.PopulateWithValue(static_cast<float>(shard_id + 1));
+    ShardTensor shard = {shard_id,
+                         std::make_shared<Literal>(std::move(shard_data))};
+    ASSERT_OK(reconstructor.ProcessShardTensor(opt_key, std::move(shard)));
+  }
+
+  ASSERT_TRUE(recovered_literal.has_value());
+  EXPECT_THAT(recovered_literal->shape().dimensions(), ElementsAre(4, 6));
+
+  // Tile (0, 0) from shard 0 (1.0f)
+  EXPECT_EQ(recovered_literal->Get<float>({0, 0}), 1.0f);
+  EXPECT_EQ(recovered_literal->Get<float>({1, 2}), 1.0f);
+  // Tile (0, 1) from shard 1 (2.0f)
+  EXPECT_EQ(recovered_literal->Get<float>({0, 3}), 2.0f);
+  EXPECT_EQ(recovered_literal->Get<float>({1, 5}), 2.0f);
+  // Tile (1, 0) from shard 2 (3.0f)
+  EXPECT_EQ(recovered_literal->Get<float>({2, 0}), 3.0f);
+  EXPECT_EQ(recovered_literal->Get<float>({3, 2}), 3.0f);
+  // Tile (1, 1) from shard 3 (4.0f)
+  EXPECT_EQ(recovered_literal->Get<float>({2, 3}), 4.0f);
+  EXPECT_EQ(recovered_literal->Get<float>({3, 5}), 4.0f);
+}
+
+TEST_F(OriginalValueRecoveryTest, SpmdPartitioningTransposedSharding) {
+  constexpr absl::string_view hlo_string = R"hlo(
+HloModule module, entry_computation_layout={(f32[4,6]{1,0}, f32[4,6]{1,0})->f32[4,6]{1,0}}, num_partitions=4
+
+ENTRY entry (p0: f32[4,6], p1: f32[4,6]) -> f32[4,6] {
+  %p0 = f32[4,6]{1,0} parameter(0), sharding={devices=[2,2]<=[2,2]T(1,0)}
+  %p1 = f32[4,6]{1,0} parameter(1), sharding={devices=[2,2]<=[2,2]T(1,0)}
+  ROOT %add = f32[4,6]{1,0} add(%p0, %p1), sharding={devices=[2,2]<=[2,2]T(1,0)}, origin={{"transposed_origin"}}
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_TRUE(module->original_value_recovery_table().empty());
+
+  spmd::SpmdPartitionerOptions options;
+  options.allow_module_signature_change = true;
+  spmd::SpmdPartitioner partitioner(4, 1, options);
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(partitioner, module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_FALSE(module->original_value_recovery_table().empty());
+
+  ASSERT_OK_AND_ASSIGN(auto analysis,
+                       HloOriginalValueAnalysis::Create(module.get()));
+  auto analysis_shared =
+      std::shared_ptr<const HloOriginalValueAnalysis>(std::move(analysis));
+
+  std::optional<Literal> recovered_literal;
+  auto callback =
+      [&](const AbsoluteScopedTensorKey& original_tensor_key,
+          const OriginalArray& original_tensor,
+          std::shared_ptr<Literal> recovered_data,
+          const std::vector<HloModule::DebugAttributes>& debug_attributes,
+          int64_t manual_shard_id) {
+        if (original_tensor.instruction_name == "transposed_origin" &&
+            recovered_data != nullptr) {
+          recovered_literal = recovered_data->Clone();
+        }
+      };
+
+  HloOriginalValueReconstructor reconstructor(analysis_shared, callback,
+                                              std::nullopt, module.get());
+
+  AbsoluteScopedTensorKey opt_key =
+      FindOptimizedTensorKey(*analysis_shared, "transposed_origin");
+  ASSERT_FALSE(opt_key.tensor_key.instruction_name.empty());
+
+  // Shard shape is [2, 3].
+  // Device assignment is transposed T(1,0):
+  // logical_shard_id 0 -> tile (0, 0)
+  // logical_shard_id 1 -> tile (1, 0)
+  // logical_shard_id 2 -> tile (0, 1)
+  // logical_shard_id 3 -> tile (1, 1)
+  for (int64_t shard_id = 0; shard_id < 4; ++shard_id) {
+    ASSERT_OK_AND_ASSIGN(Literal shard_data,
+                         Literal::Make(ShapeUtil::MakeShape(F32, {2, 3})));
+    shard_data.PopulateWithValue(static_cast<float>(shard_id + 1));
+    ShardTensor shard = {shard_id,
+                         std::make_shared<Literal>(std::move(shard_data))};
+    ASSERT_OK(reconstructor.ProcessShardTensor(opt_key, std::move(shard)));
+  }
+
+  ASSERT_TRUE(recovered_literal.has_value());
+  EXPECT_THAT(recovered_literal->shape().dimensions(), ElementsAre(4, 6));
+
+  // Tile (0, 0) from shard 0 (1.0f)
+  EXPECT_EQ(recovered_literal->Get<float>({0, 0}), 1.0f);
+  EXPECT_EQ(recovered_literal->Get<float>({1, 2}), 1.0f);
+  // Tile (1, 0) from shard 1 (2.0f)
+  EXPECT_EQ(recovered_literal->Get<float>({2, 0}), 2.0f);
+  EXPECT_EQ(recovered_literal->Get<float>({3, 2}), 2.0f);
+  // Tile (0, 1) from shard 2 (3.0f)
+  EXPECT_EQ(recovered_literal->Get<float>({0, 3}), 3.0f);
+  EXPECT_EQ(recovered_literal->Get<float>({1, 5}), 3.0f);
+  // Tile (1, 1) from shard 3 (4.0f)
+  EXPECT_EQ(recovered_literal->Get<float>({2, 3}), 4.0f);
+  EXPECT_EQ(recovered_literal->Get<float>({3, 5}), 4.0f);
+}
+
+TEST_F(OriginalValueRecoveryTest, SpmdPartitioningUnevenTilingWithSlice) {
+  constexpr absl::string_view hlo_string = R"hlo(
+HloModule module, entry_computation_layout={(f32[3,4]{1,0})->f32[3,4]{1,0}}, num_partitions=2
+
+ENTRY entry (p: f32[3,4]) -> f32[3,4] {
+  %p = f32[3,4]{1,0} parameter(0), sharding={devices=[2,1]<=[2]}
+  ROOT %neg = f32[3,4]{1,0} negate(%p), sharding={devices=[2,1]<=[2]}, origin={{"uneven_origin"}}
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_TRUE(module->original_value_recovery_table().empty());
+
+  spmd::SpmdPartitionerOptions options;
+  options.allow_module_signature_change = true;
+  spmd::SpmdPartitioner partitioner(2, 1, options);
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(partitioner, module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_FALSE(module->original_value_recovery_table().empty());
+
+  ASSERT_OK_AND_ASSIGN(auto analysis,
+                       HloOriginalValueAnalysis::Create(module.get()));
+  auto analysis_shared =
+      std::shared_ptr<const HloOriginalValueAnalysis>(std::move(analysis));
+
+  std::optional<Literal> recovered_literal;
+  auto callback =
+      [&](const AbsoluteScopedTensorKey& original_tensor_key,
+          const OriginalArray& original_tensor,
+          std::shared_ptr<Literal> recovered_data,
+          const std::vector<HloModule::DebugAttributes>& debug_attributes,
+          int64_t manual_shard_id) {
+        if (original_tensor.instruction_name == "uneven_origin" &&
+            recovered_data != nullptr) {
+          recovered_literal = recovered_data->Clone();
+        }
+      };
+
+  HloOriginalValueReconstructor reconstructor(analysis_shared, callback,
+                                              std::nullopt, module.get());
+
+  AbsoluteScopedTensorKey opt_key =
+      FindOptimizedTensorKey(*analysis_shared, "uneven_origin");
+  ASSERT_FALSE(opt_key.tensor_key.instruction_name.empty());
+
+  // Shard shape is ceil(3/2) = 2, so shard shape is [2, 4].
+  ASSERT_OK_AND_ASSIGN(Literal literal0,
+                       Literal::Make(ShapeUtil::MakeShape(F32, {2, 4})));
+  literal0.PopulateWithValue(5.0f);
+  ShardTensor shard0 = {
+      /*logical_shard_id=*/0,
+      /*data=*/std::make_shared<Literal>(std::move(literal0))};
+
+  ASSERT_OK_AND_ASSIGN(Literal literal1,
+                       Literal::Make(ShapeUtil::MakeShape(F32, {2, 4})));
+  literal1.PopulateWithValue(9.0f);
+  ShardTensor shard1 = {
+      /*logical_shard_id=*/1,
+      /*data=*/std::make_shared<Literal>(std::move(literal1))};
+
+  ASSERT_OK(reconstructor.ProcessShardTensor(opt_key, std::move(shard0)));
+  ASSERT_OK(reconstructor.ProcessShardTensor(opt_key, std::move(shard1)));
+
+  ASSERT_TRUE(recovered_literal.has_value());
+  // The recovered literal should be sliced back to original shape [3, 4]
+  EXPECT_THAT(recovered_literal->shape().dimensions(), ElementsAre(3, 4));
+  EXPECT_EQ(recovered_literal->Get<float>({0, 0}), 5.0f);
+  EXPECT_EQ(recovered_literal->Get<float>({1, 3}), 5.0f);
+  EXPECT_EQ(recovered_literal->Get<float>({2, 0}), 9.0f);
+  EXPECT_EQ(recovered_literal->Get<float>({2, 3}), 9.0f);
+}
+
+TEST_F(OriginalValueRecoveryTest, SpmdPartitioningTupleSharding) {
+  constexpr absl::string_view hlo_string = R"hlo(
+HloModule module, entry_computation_layout={(f32[4,4]{1,0}, f32[4,8]{1,0})->(f32[4,4]{1,0}, f32[4,8]{1,0})}, num_partitions=2
+
+ENTRY entry (x: f32[4,4], y: f32[4,8]) -> (f32[4,4], f32[4,8]) {
+  %x = f32[4,4]{1,0} parameter(0), sharding={devices=[2,1]<=[2]}
+  %y = f32[4,8]{1,0} parameter(1), sharding={devices=[2,1]<=[2]}
+  ROOT %tuple = (f32[4,4]{1,0}, f32[4,8]{1,0}) tuple(%x, %y),
+    sharding={{devices=[2,1]<=[2]}, {devices=[2,1]<=[2]}},
+    origin={({"tuple_x"}, {"tuple_y"})}
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_TRUE(module->original_value_recovery_table().empty());
+
+  spmd::SpmdPartitionerOptions options;
+  options.allow_module_signature_change = true;
+  spmd::SpmdPartitioner partitioner(2, 1, options);
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(partitioner, module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_FALSE(module->original_value_recovery_table().empty());
+
+  ASSERT_OK_AND_ASSIGN(auto analysis,
+                       HloOriginalValueAnalysis::Create(module.get()));
+  auto analysis_shared =
+      std::shared_ptr<const HloOriginalValueAnalysis>(std::move(analysis));
+
+  std::optional<Literal> recovered_x;
+  std::optional<Literal> recovered_y;
+  auto callback =
+      [&](const AbsoluteScopedTensorKey& original_tensor_key,
+          const OriginalArray& original_tensor,
+          std::shared_ptr<Literal> recovered_data,
+          const std::vector<HloModule::DebugAttributes>& debug_attributes,
+          int64_t manual_shard_id) {
+        if (recovered_data != nullptr) {
+          if (original_tensor.instruction_name == "tuple_x") {
+            recovered_x = recovered_data->Clone();
+          } else if (original_tensor.instruction_name == "tuple_y") {
+            recovered_y = recovered_data->Clone();
+          }
+        }
+      };
+
+  HloOriginalValueReconstructor reconstructor(analysis_shared, callback,
+                                              std::nullopt, module.get());
+
+  AbsoluteScopedTensorKey opt_key_x =
+      FindOptimizedTensorKey(*analysis_shared, "tuple_x");
+  ASSERT_FALSE(opt_key_x.tensor_key.instruction_name.empty());
+
+  AbsoluteScopedTensorKey opt_key_y =
+      FindOptimizedTensorKey(*analysis_shared, "tuple_y");
+  ASSERT_FALSE(opt_key_y.tensor_key.instruction_name.empty());
+
+  // Shards for tuple_x: shape [2, 4]
+  ASSERT_OK_AND_ASSIGN(Literal x_shard0,
+                       Literal::Make(ShapeUtil::MakeShape(F32, {2, 4})));
+  x_shard0.PopulateWithValue(10.0f);
+  ASSERT_OK_AND_ASSIGN(Literal x_shard1,
+                       Literal::Make(ShapeUtil::MakeShape(F32, {2, 4})));
+  x_shard1.PopulateWithValue(20.0f);
+
+  ASSERT_OK(reconstructor.ProcessShardTensor(
+      opt_key_x, {0, std::make_shared<Literal>(std::move(x_shard0))}));
+  ASSERT_OK(reconstructor.ProcessShardTensor(
+      opt_key_x, {1, std::make_shared<Literal>(std::move(x_shard1))}));
+
+  ASSERT_TRUE(recovered_x.has_value());
+  EXPECT_THAT(recovered_x->shape().dimensions(), ElementsAre(4, 4));
+  EXPECT_EQ(recovered_x->Get<float>({0, 0}), 10.0f);
+  EXPECT_EQ(recovered_x->Get<float>({2, 0}), 20.0f);
+
+  // Shards for tuple_y: shape [2, 8]
+  ASSERT_OK_AND_ASSIGN(Literal y_shard0,
+                       Literal::Make(ShapeUtil::MakeShape(F32, {2, 8})));
+  y_shard0.PopulateWithValue(30.0f);
+  ASSERT_OK_AND_ASSIGN(Literal y_shard1,
+                       Literal::Make(ShapeUtil::MakeShape(F32, {2, 8})));
+  y_shard1.PopulateWithValue(40.0f);
+
+  ASSERT_OK(reconstructor.ProcessShardTensor(
+      opt_key_y, {0, std::make_shared<Literal>(std::move(y_shard0))}));
+  ASSERT_OK(reconstructor.ProcessShardTensor(
+      opt_key_y, {1, std::make_shared<Literal>(std::move(y_shard1))}));
+
+  ASSERT_TRUE(recovered_y.has_value());
+  EXPECT_THAT(recovered_y->shape().dimensions(), ElementsAre(4, 8));
+  EXPECT_EQ(recovered_y->Get<float>({0, 0}), 30.0f);
+  EXPECT_EQ(recovered_y->Get<float>({2, 0}), 40.0f);
+}
+
+TEST_F(OriginalValueRecoveryTest, SpmdPartitioningTupleWithReplicatedElement) {
+  constexpr absl::string_view hlo_string = R"hlo(
+HloModule module, entry_computation_layout={(f32[2]{0}, f32[4,4]{1,0})->(f32[2]{0}, f32[4,4]{1,0})}, num_partitions=2
+
+ENTRY entry (x: f32[2], y: f32[4,4]) -> (f32[2], f32[4,4]) {
+  %x = f32[2]{0} parameter(0), sharding={replicated}
+  %y = f32[4,4]{1,0} parameter(1), sharding={devices=[2,1]<=[2]}
+  ROOT %tuple = (f32[2]{0}, f32[4,4]{1,0}) tuple(%x, %y),
+    sharding={{replicated}, {devices=[2,1]<=[2]}},
+    origin={({"tuple_origin" {0}}, {"tuple_origin" {1}})}
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_TRUE(module->original_value_recovery_table().empty());
+
+  spmd::SpmdPartitionerOptions options;
+  options.allow_module_signature_change = true;
+  spmd::SpmdPartitioner partitioner(2, 1, options);
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(partitioner, module.get()));
+  EXPECT_TRUE(changed);
+  // Only the sharded element (index {1}) requires a recovery computation;
+  // the replicated element (index {0}) is directly propagated.
+  EXPECT_EQ(module->original_value_recovery_table().size(), 1);
+
+  ASSERT_OK_AND_ASSIGN(auto analysis,
+                       HloOriginalValueAnalysis::Create(module.get()));
+  auto analysis_shared =
+      std::shared_ptr<const HloOriginalValueAnalysis>(std::move(analysis));
+
+  std::optional<Literal> recovered_x;
+  std::optional<Literal> recovered_y;
+  auto callback =
+      [&](const AbsoluteScopedTensorKey& original_tensor_key,
+          const OriginalArray& original_tensor,
+          std::shared_ptr<Literal> recovered_data,
+          const std::vector<HloModule::DebugAttributes>& debug_attributes,
+          int64_t manual_shard_id) {
+        if (original_tensor.instruction_name == "tuple_origin" &&
+            recovered_data != nullptr) {
+          if (original_tensor.shape_index == ShapeIndex({0})) {
+            recovered_x = recovered_data->Clone();
+          } else if (original_tensor.shape_index == ShapeIndex({1})) {
+            recovered_y = recovered_data->Clone();
+          }
+        }
+      };
+
+  HloOriginalValueReconstructor reconstructor(analysis_shared, callback,
+                                              std::nullopt, module.get());
+
+  // Reconstruct the replicated tuple element {0} from a single shard.
+  AbsoluteScopedTensorKey opt_key_x =
+      FindOptimizedTensorKey(*analysis_shared, "tuple_origin", {0});
+  ASSERT_FALSE(opt_key_x.tensor_key.instruction_name.empty());
+
+  Literal x_shard0 = LiteralUtil::CreateR1<float>({3.0f, 6.0f});
+  ASSERT_OK(reconstructor.ProcessShardTensor(
+      opt_key_x, {0, std::make_shared<Literal>(std::move(x_shard0))}));
+  ASSERT_TRUE(recovered_x.has_value());
+  EXPECT_THAT(recovered_x->shape().dimensions(), ElementsAre(2));
+  EXPECT_EQ(recovered_x->Get<float>({0}), 3.0f);
+  EXPECT_EQ(recovered_x->Get<float>({1}), 6.0f);
+
+  // Reconstruct the sharded tuple element {1} from both shards.
+  AbsoluteScopedTensorKey opt_key_y =
+      FindOptimizedTensorKey(*analysis_shared, "tuple_origin", {1});
+  ASSERT_FALSE(opt_key_y.tensor_key.instruction_name.empty());
+
+  ASSERT_OK_AND_ASSIGN(Literal y_shard0,
+                       Literal::Make(ShapeUtil::MakeShape(F32, {2, 4})));
+  y_shard0.PopulateWithValue(7.0f);
+  ASSERT_OK_AND_ASSIGN(Literal y_shard1,
+                       Literal::Make(ShapeUtil::MakeShape(F32, {2, 4})));
+  y_shard1.PopulateWithValue(14.0f);
+
+  ASSERT_OK(reconstructor.ProcessShardTensor(
+      opt_key_y, {0, std::make_shared<Literal>(std::move(y_shard0))}));
+  ASSERT_OK(reconstructor.ProcessShardTensor(
+      opt_key_y, {1, std::make_shared<Literal>(std::move(y_shard1))}));
+
+  ASSERT_TRUE(recovered_y.has_value());
+  EXPECT_THAT(recovered_y->shape().dimensions(), ElementsAre(4, 4));
+  EXPECT_EQ(recovered_y->Get<float>({0, 0}), 7.0f);
+  EXPECT_EQ(recovered_y->Get<float>({2, 0}), 14.0f);
+}
+
+TEST_F(OriginalValueRecoveryTest, SpmdPartitioningPartiallyReplicatedSharding) {
+  constexpr absl::string_view hlo_string = R"hlo(
+HloModule module, entry_computation_layout={(f32[4,6]{1,0})->f32[4,6]{1,0}}, num_partitions=4
+
+ENTRY entry (p: f32[4,6]) -> f32[4,6] {
+  %p = f32[4,6]{1,0} parameter(0), sharding={devices=[2,1,2]<=[4] last_tile_dim_replicate}
+  ROOT %neg = f32[4,6]{1,0} negate(%p), sharding={devices=[2,1,2]<=[4] last_tile_dim_replicate}, origin={{"partial_rep_origin"}}
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_TRUE(module->original_value_recovery_table().empty());
+
+  spmd::SpmdPartitionerOptions options;
+  options.allow_module_signature_change = true;
+  spmd::SpmdPartitioner partitioner(4, 1, options);
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(partitioner, module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_FALSE(module->original_value_recovery_table().empty());
+
+  ASSERT_OK_AND_ASSIGN(auto analysis,
+                       HloOriginalValueAnalysis::Create(module.get()));
+  auto analysis_shared =
+      std::shared_ptr<const HloOriginalValueAnalysis>(std::move(analysis));
+
+  std::optional<Literal> recovered_literal;
+  auto callback =
+      [&](const AbsoluteScopedTensorKey& original_tensor_key,
+          const OriginalArray& original_tensor,
+          std::shared_ptr<Literal> recovered_data,
+          const std::vector<HloModule::DebugAttributes>& debug_attributes,
+          int64_t manual_shard_id) {
+        if (original_tensor.instruction_name == "partial_rep_origin" &&
+            recovered_data != nullptr) {
+          recovered_literal = recovered_data->Clone();
+        }
+      };
+
+  HloOriginalValueReconstructor reconstructor(analysis_shared, callback,
+                                              std::nullopt, module.get());
+
+  AbsoluteScopedTensorKey opt_key =
+      FindOptimizedTensorKey(*analysis_shared, "partial_rep_origin");
+  ASSERT_FALSE(opt_key.tensor_key.instruction_name.empty());
+
+  // Devices {0, 1} hold tile 0 (11.0f), devices {2, 3} hold tile 1 (22.0f).
+  for (int64_t shard_id = 0; shard_id < 4; ++shard_id) {
+    ASSERT_OK_AND_ASSIGN(Literal shard_data,
+                         Literal::Make(ShapeUtil::MakeShape(F32, {2, 6})));
+    shard_data.PopulateWithValue(shard_id < 2 ? 11.0f : 22.0f);
+    ASSERT_OK(reconstructor.ProcessShardTensor(
+        opt_key, {shard_id, std::make_shared<Literal>(std::move(shard_data))}));
+  }
+
+  ASSERT_TRUE(recovered_literal.has_value());
+  EXPECT_THAT(recovered_literal->shape().dimensions(), ElementsAre(4, 6));
+  EXPECT_EQ(recovered_literal->Get<float>({0, 0}), 11.0f);
+  EXPECT_EQ(recovered_literal->Get<float>({1, 5}), 11.0f);
+  EXPECT_EQ(recovered_literal->Get<float>({2, 0}), 22.0f);
+  EXPECT_EQ(recovered_literal->Get<float>({3, 5}), 22.0f);
+}
+
+TEST_F(OriginalValueRecoveryTest, SpmdPartitioningManualSubgroupSharding) {
+  constexpr absl::string_view hlo_string = R"hlo(
+HloModule module, entry_computation_layout={(f32[8,6]{1,0})->f32[8,6]{1,0}}, num_partitions=4
+
+ENTRY entry (p: f32[8,6]) -> f32[8,6] {
+  %p = f32[8,6]{1,0} parameter(0), sharding={mesh['x'=2,'y'=2], [{'y'}, {}], manual={'x'}}
+  ROOT %neg = f32[8,6]{1,0} negate(%p), sharding={mesh['x'=2,'y'=2], [{'y'}, {}], manual={'x'}}, origin={{"manual_subgroup_origin"}}
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_TRUE(module->original_value_recovery_table().empty());
+
+  spmd::SpmdPartitionerOptions options;
+  options.allow_module_signature_change = true;
+  spmd::SpmdPartitioner partitioner(4, 1, options);
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(partitioner, module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_FALSE(module->original_value_recovery_table().empty());
+
+  ASSERT_OK_AND_ASSIGN(auto analysis,
+                       HloOriginalValueAnalysis::Create(module.get()));
+  auto analysis_shared =
+      std::shared_ptr<const HloOriginalValueAnalysis>(std::move(analysis));
+
+  absl::flat_hash_map<int64_t, Literal> recovered_by_manual_id;
+  auto callback =
+      [&](const AbsoluteScopedTensorKey& original_tensor_key,
+          const OriginalArray& original_tensor,
+          std::shared_ptr<Literal> recovered_data,
+          const std::vector<HloModule::DebugAttributes>& debug_attributes,
+          int64_t manual_shard_id) {
+        if (original_tensor.instruction_name == "manual_subgroup_origin" &&
+            recovered_data != nullptr) {
+          recovered_by_manual_id[manual_shard_id] = recovered_data->Clone();
+        }
+      };
+
+  HloOriginalValueReconstructor reconstructor(analysis_shared, callback,
+                                              std::nullopt, module.get());
+
+  AbsoluteScopedTensorKey opt_key =
+      FindOptimizedTensorKey(*analysis_shared, "manual_subgroup_origin");
+  ASSERT_FALSE(opt_key.tensor_key.instruction_name.empty());
+
+  // Per-device shard shape is [2, 6] (8 divided by 2 tiled * 2 manual = 2).
+  // mesh['x'=2,'y'=2], [{'y'}, {}], manual={'x'}:
+  //   manual_id 0 (x=0): device 0 (y=0, tile 0) and device 1 (y=1, tile 1)
+  //   manual_id 1 (x=1): device 2 (y=0, tile 0) and device 3 (y=1, tile 1)
+  for (int64_t shard_id = 0; shard_id < 4; ++shard_id) {
+    ASSERT_OK_AND_ASSIGN(Literal shard_data,
+                         Literal::Make(ShapeUtil::MakeShape(F32, {2, 6})));
+    shard_data.PopulateWithValue(static_cast<float>((shard_id + 1) * 10));
+    ASSERT_OK(reconstructor.ProcessShardTensor(
+        opt_key, {shard_id, std::make_shared<Literal>(std::move(shard_data))}));
+  }
+
+  ASSERT_EQ(recovered_by_manual_id.size(), 2);
+  ASSERT_TRUE(recovered_by_manual_id.contains(0));
+  ASSERT_TRUE(recovered_by_manual_id.contains(1));
+
+  // Manual group 0 combines device 0 (10.0f) and device 1 (20.0f) -> [4, 6]
+  EXPECT_THAT(recovered_by_manual_id.at(0).shape().dimensions(),
+              ElementsAre(4, 6));
+  EXPECT_EQ(recovered_by_manual_id.at(0).Get<float>({0, 0}), 10.0f);
+  EXPECT_EQ(recovered_by_manual_id.at(0).Get<float>({1, 5}), 10.0f);
+  EXPECT_EQ(recovered_by_manual_id.at(0).Get<float>({2, 0}), 20.0f);
+  EXPECT_EQ(recovered_by_manual_id.at(0).Get<float>({3, 5}), 20.0f);
+
+  // Manual group 1 combines device 2 (30.0f) and device 3 (40.0f) -> [4, 6]
+  EXPECT_THAT(recovered_by_manual_id.at(1).shape().dimensions(),
+              ElementsAre(4, 6));
+  EXPECT_EQ(recovered_by_manual_id.at(1).Get<float>({0, 0}), 30.0f);
+  EXPECT_EQ(recovered_by_manual_id.at(1).Get<float>({1, 5}), 30.0f);
+  EXPECT_EQ(recovered_by_manual_id.at(1).Get<float>({2, 0}), 40.0f);
+  EXPECT_EQ(recovered_by_manual_id.at(1).Get<float>({3, 5}), 40.0f);
+}
+
+TEST_F(OriginalValueRecoveryTest,
+       AlgebraicSimplifierReshapeAndBroadcastMerged) {
+  constexpr absl::string_view hlo_string = R"hlo(
+HloModule module
+
+ENTRY entry (param0: f32[5]) -> f32[1,2,3,5,1] {
+  %param0 = f32[5]{0} parameter(0)
+  %reshape = f32[1,5,1]{2,1,0} reshape(%param0), origin={{"reshape"}}
+  ROOT %broadcast = f32[1,2,3,5,1]{4,3,2,1,0} broadcast(%reshape), dimensions={0,3,4}
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_TRUE(module->original_value_recovery_table().empty());
+
+  AlgebraicSimplifierOptions options;
+  AlgebraicSimplifier simplifier(options);
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(simplifier, module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_FALSE(module->original_value_recovery_table().empty());
+
+  ASSERT_OK_AND_ASSIGN(auto analysis,
+                       HloOriginalValueAnalysis::Create(module.get()));
+  auto analysis_shared =
+      std::shared_ptr<const HloOriginalValueAnalysis>(std::move(analysis));
+
+  std::optional<Literal> recovered_reshape;
+  auto callback =
+      [&](const AbsoluteScopedTensorKey& original_tensor_key,
+          const OriginalArray& original_tensor,
+          std::shared_ptr<Literal> recovered_data,
+          const std::vector<HloModule::DebugAttributes>& debug_attributes,
+          int64_t manual_shard_id) {
+        if (original_tensor.instruction_name == "reshape" &&
+            recovered_data != nullptr) {
+          recovered_reshape = recovered_data->Clone();
+        }
+      };
+
+  HloOriginalValueReconstructor reconstructor(analysis_shared, callback,
+                                              std::nullopt, module.get());
+
+  AbsoluteScopedTensorKey opt_key =
+      FindOptimizedTensorKey(*analysis_shared, "reshape");
+  ASSERT_FALSE(opt_key.tensor_key.instruction_name.empty());
+
+  Literal literal =
+      LiteralUtil::CreateR1<float>({1.0f, 2.0f, 3.0f, 4.0f, 5.0f});
+  ShardTensor shard = {/*logical_shard_id=*/0,
+                       /*data=*/std::make_shared<Literal>(std::move(literal))};
+
+  ASSERT_OK(reconstructor.ProcessShardTensor(opt_key, std::move(shard)));
+
+  ASSERT_TRUE(recovered_reshape.has_value());
+  EXPECT_THAT(recovered_reshape->shape().dimensions(), ElementsAre(1, 5, 1));
+  EXPECT_EQ(recovered_reshape->Get<float>({0, 0, 0}), 1.0f);
+  EXPECT_EQ(recovered_reshape->Get<float>({0, 1, 0}), 2.0f);
+  EXPECT_EQ(recovered_reshape->Get<float>({0, 2, 0}), 3.0f);
+  EXPECT_EQ(recovered_reshape->Get<float>({0, 3, 0}), 4.0f);
+  EXPECT_EQ(recovered_reshape->Get<float>({0, 4, 0}), 5.0f);
+}
+
+TEST_F(OriginalValueRecoveryTest,
+       ChainedPassesAlgebraicSimplifierThenSpmdPartitioner) {
+  constexpr absl::string_view hlo_string = R"hlo(
+HloModule module, entry_computation_layout={(f32[4]{0})->f32[1,2,3,4,1]{4,3,2,1,0}}, num_partitions=2
+
+ENTRY entry (param0: f32[4]) -> f32[1,2,3,4,1] {
+  %param0 = f32[4]{0} parameter(0), sharding={devices=[2]<=[2]}
+  %reshape = f32[1,4,1]{2,1,0} reshape(%param0), origin={{"reshape"}}
+  ROOT %broadcast = f32[1,2,3,4,1]{4,3,2,1,0} broadcast(%reshape), dimensions={0,3,4}
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_TRUE(module->original_value_recovery_table().empty());
+
+  // Pass 1: AlgebraicSimplifier merges reshape and broadcast.
+  AlgebraicSimplifierOptions simplifier_options;
+  AlgebraicSimplifier simplifier(simplifier_options);
+  ASSERT_OK_AND_ASSIGN(bool simplifier_changed,
+                       RunHloPass(simplifier, module.get()));
+  EXPECT_TRUE(simplifier_changed);
+  EXPECT_FALSE(module->original_value_recovery_table().empty());
+
+  // Pass 2: SpmdPartitioner partitions the remaining instructions.
+  spmd::SpmdPartitionerOptions spmd_options;
+  spmd_options.allow_module_signature_change = true;
+  spmd::SpmdPartitioner partitioner(2, 1, spmd_options);
+  ASSERT_OK_AND_ASSIGN(bool spmd_changed,
+                       RunHloPass(partitioner, module.get()));
+  EXPECT_TRUE(spmd_changed);
+
+  ASSERT_OK_AND_ASSIGN(auto analysis,
+                       HloOriginalValueAnalysis::Create(module.get()));
+  auto analysis_shared =
+      std::shared_ptr<const HloOriginalValueAnalysis>(std::move(analysis));
+
+  std::optional<Literal> recovered_reshape;
+  auto callback =
+      [&](const AbsoluteScopedTensorKey& original_tensor_key,
+          const OriginalArray& original_tensor,
+          std::shared_ptr<Literal> recovered_data,
+          const std::vector<HloModule::DebugAttributes>& debug_attributes,
+          int64_t manual_shard_id) {
+        if (original_tensor.instruction_name == "reshape" &&
+            recovered_data != nullptr) {
+          recovered_reshape = recovered_data->Clone();
+        }
+      };
+
+  HloOriginalValueReconstructor reconstructor(analysis_shared, callback,
+                                              std::nullopt, module.get());
+
+  AbsoluteScopedTensorKey opt_key =
+      FindOptimizedTensorKey(*analysis_shared, "reshape");
+  ASSERT_FALSE(opt_key.tensor_key.instruction_name.empty());
+
+  // Each shard has shape [2].
+  // Shard 0: {1.0f, 2.0f}
+  Literal shard0_lit = LiteralUtil::CreateR1<float>({1.0f, 2.0f});
+  // Shard 1: {3.0f, 4.0f}
+  Literal shard1_lit = LiteralUtil::CreateR1<float>({3.0f, 4.0f});
+
+  ASSERT_OK(reconstructor.ProcessShardTensor(
+      opt_key, {0, std::make_shared<Literal>(std::move(shard0_lit))}));
+  ASSERT_OK(reconstructor.ProcessShardTensor(
+      opt_key, {1, std::make_shared<Literal>(std::move(shard1_lit))}));
+
+  ASSERT_TRUE(recovered_reshape.has_value());
+  EXPECT_THAT(recovered_reshape->shape().dimensions(), ElementsAre(1, 4, 1));
+  EXPECT_EQ(recovered_reshape->Get<float>({0, 0, 0}), 1.0f);
+  EXPECT_EQ(recovered_reshape->Get<float>({0, 1, 0}), 2.0f);
+  EXPECT_EQ(recovered_reshape->Get<float>({0, 2, 0}), 3.0f);
+  EXPECT_EQ(recovered_reshape->Get<float>({0, 3, 0}), 4.0f);
+}
+
+TEST_F(OriginalValueRecoveryTest,
+       ChainedPassesPartitionedMiddleUnshardDropped) {
+  constexpr absl::string_view hlo_string = R"hlo(
+HloModule module, entry_computation_layout={(f32[4]{0})->f32[1,2,3,4,1]{4,3,2,1,0}}, num_partitions=2,
+debug_attributes={
+  {"reshape"}:({callback_id=123,partitioned=true})
+}
+
+ENTRY entry (param0: f32[4]) -> f32[1,2,3,4,1] {
+  %param0 = f32[4]{0} parameter(0), sharding={devices=[2]<=[2]}
+  %reshape = f32[1,4,1]{2,1,0} reshape(%param0), origin={{"reshape"}}
+  ROOT %broadcast = f32[1,2,3,4,1]{4,3,2,1,0} broadcast(%reshape), dimensions={0,3,4}
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_TRUE(module->original_value_recovery_table().empty());
+
+  AlgebraicSimplifierOptions simplifier_options;
+  AlgebraicSimplifier simplifier(simplifier_options);
+  ASSERT_OK_AND_ASSIGN(bool simplifier_changed,
+                       RunHloPass(simplifier, module.get()));
+  EXPECT_TRUE(simplifier_changed);
+
+  spmd::SpmdPartitionerOptions spmd_options;
+  spmd_options.allow_module_signature_change = true;
+  spmd::SpmdPartitioner partitioner(2, 1, spmd_options);
+  ASSERT_OK_AND_ASSIGN(bool spmd_changed,
+                       RunHloPass(partitioner, module.get()));
+  EXPECT_TRUE(spmd_changed);
+
+  ASSERT_OK_AND_ASSIGN(auto analysis,
+                       HloOriginalValueAnalysis::Create(module.get()));
+  auto analysis_shared =
+      std::shared_ptr<const HloOriginalValueAnalysis>(std::move(analysis));
+
+  int call_count = 0;
+  std::vector<std::optional<Literal>> received_literals;
+  auto callback =
+      [&](const AbsoluteScopedTensorKey& original_tensor_key,
+          const OriginalArray& original_tensor,
+          std::shared_ptr<Literal> recovered_data,
+          const std::vector<HloModule::DebugAttributes>& debug_attributes,
+          int64_t manual_shard_id) {
+        if (original_tensor.instruction_name == "reshape") {
+          ++call_count;
+          if (recovered_data != nullptr) {
+            received_literals.push_back(recovered_data->Clone());
+          } else {
+            received_literals.push_back(std::nullopt);
+          }
+        }
+      };
+
+  HloOriginalValueReconstructor reconstructor(analysis_shared, callback,
+                                              std::nullopt, module.get());
+
+  AbsoluteScopedTensorKey opt_key =
+      FindOptimizedTensorKey(*analysis_shared, "reshape");
+  ASSERT_FALSE(opt_key.tensor_key.instruction_name.empty());
+
+  Literal shard0_lit = LiteralUtil::CreateR1<float>({1.0f, 2.0f});
+  Literal shard1_lit = LiteralUtil::CreateR1<float>({3.0f, 4.0f});
+
+  ASSERT_OK(reconstructor.ProcessShardTensor(
+      opt_key, {0, std::make_shared<Literal>(std::move(shard0_lit))}));
+  ASSERT_OK(reconstructor.ProcessShardTensor(
+      opt_key, {1, std::make_shared<Literal>(std::move(shard1_lit))}));
+
+  // Because the unshard recovery module precedes the full-shape reshape module,
+  // per-shard recovery (partitioned=true) cannot be evaluated and reports
+  // nullptr for each shard.
+  ASSERT_EQ(call_count, 2);
+  EXPECT_FALSE(received_literals[0].has_value());
+  EXPECT_FALSE(received_literals[1].has_value());
+}
+
+TEST_F(OriginalValueRecoveryTest, PartitionedDebugAttributeReportsShards) {
+  constexpr absl::string_view hlo_string = R"hlo(
+HloModule module, entry_computation_layout={(f32[4,8]{1,0})->f32[4,8]{1,0}}, num_partitions=2,
+debug_attributes={
+  {"neg_origin"}:({callback_id=42,partitioned=true})
+}
+
+ENTRY entry (p: f32[4,8]) -> f32[4,8] {
+  %p = f32[4,8]{1,0} parameter(0), sharding={devices=[2,1]<=[2]}
+  ROOT %neg = f32[4,8]{1,0} negate(%p), sharding={devices=[2,1]<=[2]}, origin={{"neg_origin"}}
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_TRUE(module->original_value_recovery_table().empty());
+
+  spmd::SpmdPartitionerOptions options;
+  options.allow_module_signature_change = true;
+  spmd::SpmdPartitioner partitioner(2, 1, options);
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(partitioner, module.get()));
+  EXPECT_TRUE(changed);
+
+  ASSERT_OK_AND_ASSIGN(auto analysis,
+                       HloOriginalValueAnalysis::Create(module.get()));
+  auto analysis_shared =
+      std::shared_ptr<const HloOriginalValueAnalysis>(std::move(analysis));
+
+  std::vector<Literal> received_shards;
+  auto callback =
+      [&](const AbsoluteScopedTensorKey& original_tensor_key,
+          const OriginalArray& original_tensor,
+          std::shared_ptr<Literal> recovered_data,
+          const std::vector<HloModule::DebugAttributes>& debug_attributes,
+          int64_t manual_shard_id) {
+        if (original_tensor.instruction_name == "neg_origin" &&
+            recovered_data != nullptr) {
+          received_shards.push_back(recovered_data->Clone());
+        }
+      };
+
+  HloOriginalValueReconstructor reconstructor(analysis_shared, callback,
+                                              std::nullopt, module.get());
+
+  AbsoluteScopedTensorKey opt_key =
+      FindOptimizedTensorKey(*analysis_shared, "neg_origin");
+  ASSERT_FALSE(opt_key.tensor_key.instruction_name.empty());
+
+  ASSERT_OK_AND_ASSIGN(Literal literal0,
+                       Literal::Make(ShapeUtil::MakeShape(F32, {2, 8})));
+  literal0.PopulateWithValue(1.0f);
+  ASSERT_OK(reconstructor.ProcessShardTensor(
+      opt_key, {0, std::make_shared<Literal>(std::move(literal0))}));
+
+  ASSERT_OK_AND_ASSIGN(Literal literal1,
+                       Literal::Make(ShapeUtil::MakeShape(F32, {2, 8})));
+  literal1.PopulateWithValue(2.0f);
+  ASSERT_OK(reconstructor.ProcessShardTensor(
+      opt_key, {1, std::make_shared<Literal>(std::move(literal1))}));
+
+  // With partitioned=true, the callback is invoked for each shard immediately,
+  // skipping the final unshard!
+  ASSERT_EQ(received_shards.size(), 2);
+  EXPECT_THAT(received_shards[0].shape().dimensions(), ElementsAre(2, 8));
+  EXPECT_EQ(received_shards[0].Get<float>({0, 0}), 1.0f);
+  EXPECT_THAT(received_shards[1].shape().dimensions(), ElementsAre(2, 8));
+  EXPECT_EQ(received_shards[1].Get<float>({0, 0}), 2.0f);
+}
+
+TEST_F(OriginalValueRecoveryTest, MixedPartitionedDebugAttributes) {
+  constexpr absl::string_view hlo_string = R"hlo(
+HloModule module, entry_computation_layout={(f32[4,8]{1,0})->f32[4,8]{1,0}}, num_partitions=2,
+debug_attributes={
+  {"neg_origin"}:({callback_id=1,partitioned=false},{callback_id=2,partitioned=true})
+}
+
+ENTRY entry (p: f32[4,8]) -> f32[4,8] {
+  %p = f32[4,8]{1,0} parameter(0), sharding={devices=[2,1]<=[2]}
+  ROOT %neg = f32[4,8]{1,0} negate(%p), sharding={devices=[2,1]<=[2]}, origin={{"neg_origin"}}
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_TRUE(module->original_value_recovery_table().empty());
+
+  spmd::SpmdPartitionerOptions options;
+  options.allow_module_signature_change = true;
+  spmd::SpmdPartitioner partitioner(2, 1, options);
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(partitioner, module.get()));
+  EXPECT_TRUE(changed);
+
+  ASSERT_OK_AND_ASSIGN(auto analysis,
+                       HloOriginalValueAnalysis::Create(module.get()));
+  auto analysis_shared =
+      std::shared_ptr<const HloOriginalValueAnalysis>(std::move(analysis));
+
+  std::vector<Literal> recovered_unsharded;
+  std::vector<Literal> recovered_partitioned;
+  auto callback =
+      [&](const AbsoluteScopedTensorKey& original_tensor_key,
+          const OriginalArray& original_tensor,
+          std::shared_ptr<Literal> recovered_data,
+          const std::vector<HloModule::DebugAttributes>& debug_attributes,
+          int64_t manual_shard_id) {
+        if (original_tensor.instruction_name != "neg_origin" ||
+            recovered_data == nullptr) {
+          return;
+        }
+        for (const auto& attr : debug_attributes) {
+          if (attr.partitioned) {
+            recovered_partitioned.push_back(recovered_data->Clone());
+          } else {
+            recovered_unsharded.push_back(recovered_data->Clone());
+          }
+        }
+      };
+
+  HloOriginalValueReconstructor reconstructor(analysis_shared, callback,
+                                              std::nullopt, module.get());
+
+  AbsoluteScopedTensorKey opt_key =
+      FindOptimizedTensorKey(*analysis_shared, "neg_origin");
+  ASSERT_FALSE(opt_key.tensor_key.instruction_name.empty());
+
+  ASSERT_OK_AND_ASSIGN(Literal literal0,
+                       Literal::Make(ShapeUtil::MakeShape(F32, {2, 8})));
+  literal0.PopulateWithValue(1.0f);
+  ASSERT_OK(reconstructor.ProcessShardTensor(
+      opt_key, {0, std::make_shared<Literal>(std::move(literal0))}));
+
+  ASSERT_OK_AND_ASSIGN(Literal literal1,
+                       Literal::Make(ShapeUtil::MakeShape(F32, {2, 8})));
+  literal1.PopulateWithValue(2.0f);
+  ASSERT_OK(reconstructor.ProcessShardTensor(
+      opt_key, {1, std::make_shared<Literal>(std::move(literal1))}));
+
+  ASSERT_EQ(recovered_partitioned.size(), 2);
+  EXPECT_THAT(recovered_partitioned[0].shape().dimensions(), ElementsAre(2, 8));
+  EXPECT_EQ(recovered_partitioned[0].Get<float>({0, 0}), 1.0f);
+  EXPECT_THAT(recovered_partitioned[1].shape().dimensions(), ElementsAre(2, 8));
+  EXPECT_EQ(recovered_partitioned[1].Get<float>({0, 0}), 2.0f);
+
+  ASSERT_EQ(recovered_unsharded.size(), 1);
+  EXPECT_THAT(recovered_unsharded[0].shape().dimensions(), ElementsAre(4, 8));
+  EXPECT_EQ(recovered_unsharded[0].Get<float>({0, 0}), 1.0f);
+  EXPECT_EQ(recovered_unsharded[0].Get<float>({2, 0}), 2.0f);
+}
+
+TEST_F(OriginalValueRecoveryTest,
+       MultiControllerUnaddressableDeviceAbortsRecovery) {
+  constexpr absl::string_view hlo_string = R"hlo(
+HloModule module, entry_computation_layout={(f32[4,8]{1,0})->f32[4,8]{1,0}}, num_partitions=2
+
+ENTRY entry (p: f32[4,8]) -> f32[4,8] {
+  %p = f32[4,8]{1,0} parameter(0), sharding={devices=[2,1]<=[2]}
+  ROOT %neg = f32[4,8]{1,0} negate(%p), sharding={devices=[2,1]<=[2]}, origin={{"neg_origin"}}
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_TRUE(module->original_value_recovery_table().empty());
+
+  spmd::SpmdPartitionerOptions options;
+  options.allow_module_signature_change = true;
+  spmd::SpmdPartitioner partitioner(2, 1, options);
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(partitioner, module.get()));
+  EXPECT_TRUE(changed);
+
+  ASSERT_OK_AND_ASSIGN(auto analysis,
+                       HloOriginalValueAnalysis::Create(module.get()));
+  auto analysis_shared =
+      std::shared_ptr<const HloOriginalValueAnalysis>(std::move(analysis));
+
+  bool callback_called = false;
+  auto callback =
+      [&](const AbsoluteScopedTensorKey& original_tensor_key,
+          const OriginalArray& original_tensor,
+          std::shared_ptr<Literal> recovered_data,
+          const std::vector<HloModule::DebugAttributes>& debug_attributes,
+          int64_t manual_shard_id) { callback_called = true; };
+
+  // Only logical device 0 is addressable on this controller; device 1 is not.
+  auto logical_device_is_addressable = [](int64_t logical_device_id) -> bool {
+    return logical_device_id == 0;
+  };
+
+  HloOriginalValueReconstructor reconstructor(
+      analysis_shared, callback, logical_device_is_addressable, module.get());
+
+  AbsoluteScopedTensorKey opt_key =
+      FindOptimizedTensorKey(*analysis_shared, "neg_origin");
+  ASSERT_FALSE(opt_key.tensor_key.instruction_name.empty());
+
+  ASSERT_OK_AND_ASSIGN(Literal literal0,
+                       Literal::Make(ShapeUtil::MakeShape(F32, {2, 8})));
+  literal0.PopulateWithValue(1.0f);
+  ASSERT_OK(reconstructor.ProcessShardTensor(
+      opt_key, {0, std::make_shared<Literal>(std::move(literal0))}));
+
+  EXPECT_FALSE(callback_called);
+}
+
+TEST_F(OriginalValueRecoveryTest,
+       CallInlinerThenAlgebraicSimplifierPreservesScope) {
+  constexpr absl::string_view hlo_string = R"hlo(
+HloModule module
+
+callee (x: f32[5]) -> f32[1,2,3,5,1] {
+  %x = f32[5]{0} parameter(0)
+  %reshape = f32[1,5,1]{2,1,0} reshape(%x), origin={{"reshape"}}
+  ROOT %broadcast = f32[1,2,3,5,1]{4,3,2,1,0} broadcast(%reshape), dimensions={0,3,4}
+}
+
+ENTRY entry (param0: f32[5]) -> f32[1,2,3,5,1] {
+  %param0 = f32[5]{0} parameter(0)
+  ROOT %call.1 = f32[1,2,3,5,1]{4,3,2,1,0} call(%param0), to_apply=callee, origin={{"call.1"},["call.1"]}
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_TRUE(module->original_value_recovery_table().empty());
+
+  // Pass 1: CallInliner inlines callee and prefixes origin with "call.1/".
+  CallInliner inliner(/*single_call_site=*/false);
+  ASSERT_OK_AND_ASSIGN(bool inliner_changed, RunHloPass(inliner, module.get()));
+  EXPECT_TRUE(inliner_changed);
+
+  // Pass 2: AlgebraicSimplifier merges reshape and broadcast and records the
+  // recovery computation for {"call.1/reshape"}.
+  AlgebraicSimplifierOptions options;
+  AlgebraicSimplifier simplifier(options);
+  ASSERT_OK_AND_ASSIGN(bool simplifier_changed,
+                       RunHloPass(simplifier, module.get()));
+  EXPECT_TRUE(simplifier_changed);
+  EXPECT_FALSE(module->original_value_recovery_table().empty());
+
+  ASSERT_OK_AND_ASSIGN(auto analysis,
+                       HloOriginalValueAnalysis::Create(module.get()));
+  auto analysis_shared =
+      std::shared_ptr<const HloOriginalValueAnalysis>(std::move(analysis));
+
+  std::optional<AbsoluteScopedTensorKey> recovered_key;
+  std::optional<Literal> recovered_reshape;
+  auto callback =
+      [&](const AbsoluteScopedTensorKey& original_tensor_key,
+          const OriginalArray& original_tensor,
+          std::shared_ptr<Literal> recovered_data,
+          const std::vector<HloModule::DebugAttributes>& debug_attributes,
+          int64_t manual_shard_id) {
+        if (original_tensor_key.tensor_key.instruction_name == "reshape" &&
+            recovered_data != nullptr) {
+          recovered_key = original_tensor_key;
+          recovered_reshape = recovered_data->Clone();
+        }
+      };
+
+  HloOriginalValueReconstructor reconstructor(analysis_shared, callback,
+                                              std::nullopt, module.get());
+
+  AbsoluteScopedTensorKey opt_key =
+      FindOptimizedTensorKey(*analysis_shared, "call.1/reshape");
+  ASSERT_FALSE(opt_key.tensor_key.instruction_name.empty());
+
+  Literal literal =
+      LiteralUtil::CreateR1<float>({1.0f, 2.0f, 3.0f, 4.0f, 5.0f});
+  ASSERT_OK(reconstructor.ProcessShardTensor(
+      opt_key, {0, std::make_shared<Literal>(std::move(literal))}));
+
+  ASSERT_TRUE(recovered_key.has_value());
+  EXPECT_EQ(recovered_key->tensor_key.instruction_name, "reshape");
+  EXPECT_THAT(recovered_key->scope_instructions,
+              ElementsAre(ScopeInstruction::Create("call.1")));
+
+  ASSERT_TRUE(recovered_reshape.has_value());
+  EXPECT_THAT(recovered_reshape->shape().dimensions(), ElementsAre(1, 5, 1));
+  EXPECT_EQ(recovered_reshape->Get<float>({0, 0, 0}), 1.0f);
+  EXPECT_EQ(recovered_reshape->Get<float>({0, 4, 0}), 5.0f);
+}
+
+}  // namespace
+}  // namespace xla

From 042c49a45c210011ba85ee419a575bedf1b23577 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 12 Oct 2026 12:12:14 -0700
Subject: [PATCH 0615/1000] [PjRt-IFRT] Prefer
 `PjRtExecutable::GetOutputDimensions` for `OutputSpec::shard_shape`

When an executable returns manually partitioned outputs (`{manual}` sharding, e.g. from `shard_map(..., check_rep=False)`) or unevenly tiled outputs, `xla::hlo_sharding_util::TileShape` returns the unsharded global shape or fails its divisibility check rather than returning the per-device buffer shape, causing `PjRtExecutable::CommonMetadata::Serialize` to record the wrong `OutputSpec::shard_shape` or fail.

Query `pjrt_executable->GetOutputDimensions()` first when computing `OutputSpec::shard_shape`, falling back to `xla::hlo_sharding_util::TileShape` when `GetOutputDimensions()` is unavailable.

PiperOrigin-RevId: 998073278
---
 .../xla/python/pjrt_ifrt/pjrt_executable.cc   | 21 ++++---
 .../pjrt_ifrt/xla_executable_impl_test_lib.cc | 59 +++++++++++++++++++
 2 files changed, 73 insertions(+), 7 deletions(-)

diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc
index abc6d30dbbe..51bd21e0048 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc
@@ -523,6 +523,9 @@ absl::StatusOr<std::string> PjRtExecutable::CommonMetadata::Serialize(
 
   metadata.set_computation_name(pjrt_executable->name());
 
+  absl::StatusOr<std::vector<std::vector<xla::DimensionVector>>>
+      output_dimensions = pjrt_executable->GetOutputDimensions();
+
   // Encode output specs.
   for (int i = 0; i < output_dtypes.size(); ++i) {
     SerializedXlaExecutableMetadata::OutputSpec& output_spec =
@@ -545,13 +548,17 @@ absl::StatusOr<std::string> PjRtExecutable::CommonMetadata::Serialize(
     // Shape
     std::optional<Shape> shard_shape;
     {
-      // If the output sharding information is missing (e.g., serializing a
-      // previously deserialized executable) or the output dtype is token, the
-      // output shard shape is the same as the output shape. The former case
-      // will be removed once serialization/deserialization roundtrip preserves
-      // the output sharding information.
-      if (!output_hlo_shardings.has_value() ||
-          output_dtypes[i].kind() == DType::kToken) {
+      // Prefer the per-device output dimensions reported by the compiled
+      // executable. This handles manual SPMD shardings (`{manual}` from
+      // `shard_map`) and unevenly tiled shardings, where
+      // `hlo_sharding_util::TileShape` would otherwise return the unsharded
+      // global shape or fail its divisibility check.
+      if (output_dtypes[i].kind() == DType::kToken) {
+        shard_shape = output_shapes[i];
+      } else if (output_dimensions.ok() && !output_dimensions->empty() &&
+                 (*output_dimensions)[0].size() == output_dtypes.size()) {
+        shard_shape = Shape((*output_dimensions)[0][i]);
+      } else if (!output_hlo_shardings.has_value()) {
         shard_shape = output_shapes[i];
       } else {
         ABSL_ASSIGN_OR_RETURN(xla::PrimitiveType element_type,
diff --git a/third_party/xla/xla/python/pjrt_ifrt/xla_executable_impl_test_lib.cc b/third_party/xla/xla/python/pjrt_ifrt/xla_executable_impl_test_lib.cc
index 08923015c8c..7612fd2dcdf 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/xla_executable_impl_test_lib.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/xla_executable_impl_test_lib.cc
@@ -1298,6 +1298,65 @@ TEST(ExecutableTest, ExecutableSerialization) {
   }
 }
 
+TEST(ExecutableTest, ExecutableSerializationUsesExecutableOutputDimensions) {
+  static constexpr absl::string_view kUnevenShardingModule = R"mlir(
+    module @uneven_sharding attributes {
+      mhlo.num_replicas = 1 : i32,
+      mhlo.num_partitions = 2 : i32
+    } {
+      func.func @main(
+        %arg0: tensor<3x3xi32> {mhlo.sharding = "{devices=[2,1]<=[2]}"}
+      ) -> (
+        tensor<3x3xi32> {mhlo.sharding = "{devices=[2,1]<=[2]}"}
+      ) {
+        return %arg0 : tensor<3x3xi32>
+      }
+    }
+  )mlir";
+
+  ASSERT_OK_AND_ASSIGN(auto client, test_util::GetClient());
+  Compiler* compiler = client->GetDefaultCompiler();
+  absl::Span<Device* const> devices =
+      client->addressable_devices().subspan(0, 2);
+
+  ASSERT_OK_AND_ASSIGN(
+      LoadedExecutableRef loaded_executable,
+      CompileOnDevices(client.get(), compiler, kUnevenShardingModule, devices,
+                       /*replicated=*/false, /*serialize=*/false));
+
+  absl::StatusOr<std::string> serialized_executable =
+      loaded_executable->Serialize();
+  if (absl::IsUnimplemented(serialized_executable.status())) {
+    GTEST_SKIP() << "Serialization is not supported on this platform.";
+  }
+  ASSERT_OK(serialized_executable);
+
+  SerializedXlaExecutableMetadata metadata;
+  tsl::protobuf::io::ArrayInputStream input_stream(
+      serialized_executable->data(), serialized_executable->size());
+  ASSERT_TRUE(tsl::protobuf::util::ParseDelimitedFromZeroCopyStream(
+      &metadata, &input_stream, nullptr));
+
+  ASSERT_EQ(metadata.output_specs_size(), 1);
+  EXPECT_THAT(metadata.output_specs(0).shape().dims(), ElementsAre(3, 3));
+  // Padded per-device shard shape reported by
+  // `PjRtExecutable::GetOutputDimensions()` (whereas
+  // `hlo_sharding_util::TileShape` would fail `3 % 2 == 0`).
+  EXPECT_THAT(metadata.output_specs(0).shard_shape().dims(), ElementsAre(2, 3));
+
+  ASSERT_OK_AND_ASSIGN(DeviceListRef device_list,
+                       client->MakeDeviceList(devices));
+  auto options = std::make_unique<XlaDeserializeExecutableOptions>();
+  options->devices = std::move(device_list);
+  ASSERT_OK_AND_ASSIGN(
+      LoadedExecutableRef deserialized_executable,
+      compiler
+          ->DeserializeLoadedExecutable(absl::Cord(*serialized_executable),
+                                        std::move(options))
+          .Await());
+  EXPECT_EQ(deserialized_executable->name(), "uneven_sharding");
+}
+
 }  // namespace
 }  // namespace ifrt
 }  // namespace xla

From 2bc5e1fb6a34bcdccdd797f6600bcbb8a529e67c Mon Sep 17 00:00:00 2001
From: Pablo Zimmermann <karupayun@google.com>
Date: Mon, 12 Oct 2026 12:16:01 -0700
Subject: [PATCH 0616/1000] [XLA:GPU] Track dot-region instructions and report
 L2 bytes read for tiled dot fusions.

As part of migrating `gpu_dot_fusion_cost_model` into
`GpuIndexingPerformanceModel`, add an `in_dot_region` flag so `l2_bytes_read`
only counts operands read inside the dot loop and ignores epilogue operands.

Unlike `CalculateL2Bytes` in `gpu_dot_fusion_cost_model`, which only looks at
the `dot` instruction's direct inputs, this counts the actual fusion inputs in
the dot region. This is more accurate when there are fused prologue ops (like
casts or broadcasts) or tiles larger than the tensor, though it pads `K` to a
multiple of `block_k` and rounds sub-byte types up to 1 byte.

PiperOrigin-RevId: 998075752
---
 .../model/gpu_indexing_performance_model.cc   | 42 ++++++++++++-------
 .../gpu_indexing_performance_model_test.cc    |  1 +
 2 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc
index 92f82ea499b..198284fe031 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc
+++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc
@@ -234,8 +234,8 @@ int64_t GetNumBlocksForRegion(const TiledHloInstructionType* tiled_hlo,
 }
 
 // Visits all tiled HLOs in the computation (including those within regions)
-// using a visitor that receives the tiled HLO and the root num_blocks that
-// access that tiled HLO.
+// using a visitor that receives the tiled HLO, the root num_blocks that
+// access that tiled HLO, and whether the instruction is within a dot region.
 //
 // Note: Set num_blocks_at_root to 1 if the visitor doesn't use the block count.
 template <typename TiledHloComputationType>
@@ -243,27 +243,37 @@ void ForEachInstructionInTiledHloComputation(
     const TiledHloComputationType& tiled_hlo_computation,
     int64_t num_blocks_at_root,
     absl::FunctionRef<
-        void(const typename TiledHloComputationType::InstructionType*, int64_t)>
+        void(const typename TiledHloComputationType::InstructionType*, int64_t,
+             bool)>
         visitor) {
-  std::vector<std::pair<
-      const typename TiledHloComputationType::InstructionType*, int64_t>>
-      worklist;
+  struct WorkItem {
+    const typename TiledHloComputationType::InstructionType* instruction;
+    int64_t num_blocks_cur_hlo;
+    bool in_dot_region;
+  };
+  std::vector<WorkItem> worklist;
   for (const auto* instruction : tiled_hlo_computation.instructions()) {
-    worklist.push_back({instruction, num_blocks_at_root});
+    worklist.push_back(
+        {instruction, num_blocks_at_root, /*in_dot_region=*/false});
   }
 
   while (!worklist.empty()) {
-    auto [instruction, num_blocks_cur_hlo] = worklist.back();
+    auto [instruction, num_blocks_cur_hlo, in_dot_region] = worklist.back();
     worklist.pop_back();
 
-    visitor(instruction, num_blocks_cur_hlo);
+    visitor(instruction, num_blocks_cur_hlo, in_dot_region);
 
+    bool child_in_dot_region =
+        in_dot_region ||
+        HloPredicateIsOp<HloOpcode::kDot, HloOpcode::kScaledDot>(
+            instruction->hlo());
     for (const auto& [i, region] :
          llvm::enumerate(instruction->hlo_regions())) {
       int64_t num_blocks_cur_region =
           GetNumBlocksForRegion(instruction, num_blocks_cur_hlo, i);
       for (const auto& tiled_hlo : region.instructions()) {
-        worklist.push_back({&*tiled_hlo, num_blocks_cur_region});
+        worklist.push_back(
+            {&*tiled_hlo, num_blocks_cur_region, child_in_dot_region});
       }
     }
   }
@@ -293,7 +303,7 @@ bool DoesComputationFitInRegisters(
       tiled_hlo_computation,
       /*num_blocks_at_root=*/1,
       [&](const typename TiledHloComputationType::InstructionType* tiled_hlo,
-          int64_t unused_num_blocks_cur_hlo) {
+          int64_t /*num_blocks_cur_hlo*/, bool /*in_dot_region*/) {
         if (!fits) {
           return;
         }
@@ -465,6 +475,7 @@ absl::StatusOr<EstimateRunTimeData> EstimateRunTimeForTiledHloComputationImpl(
   int64_t dot_flops = 0;
   int64_t flops = 0;
   int64_t bytes_read = 0;
+  int64_t dot_l2_bytes_read = 0;
   int64_t num_blocks = tiled_hlo_computation.num_output_tiles();
 
   DotEstimates dot_estimates;
@@ -481,7 +492,7 @@ absl::StatusOr<EstimateRunTimeData> EstimateRunTimeForTiledHloComputationImpl(
   ForEachInstructionInTiledHloComputation(
       tiled_hlo_computation, num_blocks,
       [&](const typename TiledHloComputationType::InstructionType* tiled_hlo,
-          int64_t num_blocks_cur_hlo) {
+          int64_t num_blocks_cur_hlo, bool in_dot_region) {
         const HloInstruction* hlo = tiled_hlo->hlo();
         if (fusion_adaptor.ContainsInstruction(hlo)) {
           // Number of elements in the tile after padding.
@@ -543,6 +554,9 @@ absl::StatusOr<EstimateRunTimeData> EstimateRunTimeForTiledHloComputationImpl(
         int64_t tile_bytes_read = element_type_size * num_elements;
 
         bytes_read += tile_bytes_read;
+        if (in_dot_region) {
+          dot_l2_bytes_read += tile_bytes_read;
+        }
 
         double effective_bandwidth_utilization_rate =
             BandwidthUtilizationRateHeuristicForTiledMemoryAccess(*tiled_hlo,
@@ -597,7 +611,7 @@ absl::StatusOr<EstimateRunTimeData> EstimateRunTimeForTiledHloComputationImpl(
       /*write_time=*/write_time,
       /*compute_time=*/compute_time,
       /*exec_time=*/exec_time,
-      /*l2_bytes_read=*/0,
+      /*l2_bytes_read=*/dot_l2_bytes_read,
       /*shared_memory_per_block_bytes=*/
       dot_estimates.shared_memory_per_block_bytes,
       /*registers_per_thread=*/dot_estimates.registers_per_thread};
@@ -612,7 +626,7 @@ int64_t EstimateNumWarpsImpl(
   ForEachInstructionInTiledHloComputation(
       tiled_hlo_computation, /*num_blocks_at_root=*/1,
       [&](const typename TiledHloComputationType::InstructionType* tiled_hlo,
-          int64_t unused_num_blocks_cur_hlo) {
+          int64_t /*num_blocks_cur_hlo*/, bool /*in_dot_region*/) {
         largest_live_tile_size = std::max(
             largest_live_tile_size, GetPaddedTileSize(tiled_hlo->tile_sizes()));
       });
diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc
index 67195c26f6d..3740f0af54a 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc
+++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc
@@ -726,6 +726,7 @@ ENTRY main {
       4 * (64 * 32) * (num_blocks * CeilOfRatio(257, 64))        // ==> %p1
           + 4 * (64 * 32) * (num_blocks * CeilOfRatio(257, 64))  // ==> %p0
   );
+  EXPECT_EQ(result.l2_bytes_read, result.bytes_read);
   // (32 * 64 + 32 * 64) * 4 bytes (f32) * 1 stage = 16384 bytes.
   EXPECT_EQ(result.shared_memory_per_block_bytes, 2 * (32 * 64 * 4));
   // 1 accumulator register per thread (32 * 32 / 1024) plus base kernel-state

From 658ace0e092ac14cc058e2337316a86812ba70f8 Mon Sep 17 00:00:00 2001
From: Matthias Guenther <mrguenther@google.com>
Date: Mon, 12 Oct 2026 12:20:41 -0700
Subject: [PATCH 0617/1000] Remove the LINT directives around the StableHLO
 pin.

PiperOrigin-RevId: 998078463
---
 third_party/xla/third_party/stablehlo/workspace.bzl | 2 --
 1 file changed, 2 deletions(-)

diff --git a/third_party/xla/third_party/stablehlo/workspace.bzl b/third_party/xla/third_party/stablehlo/workspace.bzl
index 8764fbd0a3f..7c94cab1d31 100644
--- a/third_party/xla/third_party/stablehlo/workspace.bzl
+++ b/third_party/xla/third_party/stablehlo/workspace.bzl
@@ -18,10 +18,8 @@
 load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 
 def repo():
-    # LINT.IfChange
     STABLEHLO_COMMIT = "b9029ac2228a12e176ea20e56f7d3dcb6bf8a019"
     STABLEHLO_SHA256 = "16af6c8f0baacda15dd72d98dc295a213be455b8fa1899af46e144d6dae670bd"
-    # LINT.ThenChange(Google-internal path)
 
     tf_http_archive(
         name = "stablehlo",

From 39552e637c1450bd8094847aa6f402f31bef1bb7 Mon Sep 17 00:00:00 2001
From: Bill Varcho <varcho@google.com>
Date: Mon, 12 Oct 2026 12:26:28 -0700
Subject: [PATCH 0618/1000] [XLA:SPMD] Pass explicit enable-full-version=false
 in ragged_dot_insert_explicit_reshards.mlir.

Pin `-sdy-insert-explicit-reshards='enable-full-version=false'` in `ragged_dot_insert_explicit_reshards.mlir` so the test continues testing the non-full version when `enableFullVersion=true` becomes the default in Shardy.

PiperOrigin-RevId: 998081930
---
 .../spmd/shardy/test/ragged_dot_insert_explicit_reshards.mlir   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/spmd/shardy/test/ragged_dot_insert_explicit_reshards.mlir b/third_party/xla/xla/service/spmd/shardy/test/ragged_dot_insert_explicit_reshards.mlir
index 2d276e5130c..1728bb57e9d 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/ragged_dot_insert_explicit_reshards.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/ragged_dot_insert_explicit_reshards.mlir
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 // ==============================================================================
-// RUN: sdy_opt %s -allow-unregistered-dialect  -sdy-insert-explicit-reshards | FileCheck %s
+// RUN: sdy_opt %s -allow-unregistered-dialect  -sdy-insert-explicit-reshards='enable-full-version=false' | FileCheck %s
 
 sdy.mesh @mesh = <["x"=2, "y"=2, "z"=4]>
 sdy.mesh @other_mesh = <["x"=2, "y"=2]>

From f1a255c8a90dbe8e4056bd0ea65551307f3da66b Mon Sep 17 00:00:00 2001
From: Akhil Goel <akhil.goel@intel.com>
Date: Mon, 12 Oct 2026 12:39:10 -0700
Subject: [PATCH 0619/1000] PR #50560: [XLA:GPU][oneAPI] Align XPU Triton hook
 with upstream Triton

Imported from GitHub PR https://github.com/openxla/xla/pull/50560

This PR aligns the XPU Triton hook with https://github.com/triton-lang/triton/commit/0ccad974ee86942e704ac06dafdc0c3f214eb229
Copybara import of the project:

--
77847fc8dcebd04c63de06fa52c890ec28a3cd4c by Akhil Goel <akhil.goel@intel.com>:

Update XPU Triton hook

Merging this change closes #50560

PiperOrigin-RevId: 998089658
---
 .../triton/intel_xpu/intel_build.patch            | 15 +++++++++------
 .../third_party/triton/intel_xpu/workspace.bzl    |  4 ++--
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/third_party/triton/intel_xpu/intel_build.patch b/third_party/xla/third_party/triton/intel_xpu/intel_build.patch
index 4020f2426ef..29b24e67c35 100644
--- a/third_party/xla/third_party/triton/intel_xpu/intel_build.patch
+++ b/third_party/xla/third_party/triton/intel_xpu/intel_build.patch
@@ -15,21 +15,24 @@
 diff --git b/BUILD b/BUILD
 --- a/BUILD
 +++ b/BUILD
-@@ -5,3 +5,4 @@
+@@ -4,6 +4,7 @@
+ load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
  load("@rules_cc//cc:cc_library.bzl", "cc_library")
  load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
 +load("@local_config_sycl//sycl:build_defs.bzl", "if_sycl_is_configured")
  # copybara:uncomment load("//tools/build_defs/license:license.bzl", "license")
-@@ -545,7 +546,7 @@
-         "lib/Dialect/TritonGPU/Transforms/Utility.cpp",  # Avoid circular dependency.
+ 
+ # See go/triton-sla.
+@@ -566,7 +567,7 @@ cc_library(
          "lib/Dialect/TritonNvidiaGPU/Transforms/ClusterBarrierMbarAllocator.cpp",  # Avoid circular dependency.
+         "lib/Dialect/TritonNvidiaGPU/Transforms/MBarrierUtilities.cpp",  # Avoid circular dependency.
          "lib/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.cpp",  # Avoid circular dependency.
 -    ],
 +    ] + if_sycl_is_configured(["//third_party/intel:TritonIntelGPU_srcs"]),
      hdrs = glob([
          "include/triton/Dialect/Gluon/IR/*.h",
          "include/triton/Dialect/Triton/IR/*.h",
-@@ -565,7 +566,7 @@
+@@ -587,7 +588,7 @@ cc_library(
          # it would be better if upstream moved it there.
          "include/triton/Conversion/MLIRTypes.h",
          "include/triton/Dialect/TritonInstrument/Transforms/ConSanTargetHooks.h",
@@ -38,7 +41,7 @@ diff --git b/BUILD b/BUILD
      copts = select({
          ":compiler_is_msvc": [],
          "//conditions:default": [
-@@ -576,7 +577,10 @@
+@@ -598,7 +599,10 @@ cc_library(
              "-Wno-implicit-fallthrough",
          ],
      }),
@@ -50,7 +53,7 @@ diff --git b/BUILD b/BUILD
      deps = [
          ":Dump",
          ":GetEnv",  # for Utility
-@@ -633,7 +638,16 @@
+@@ -659,7 +663,16 @@ cc_library(
          "@llvm-project//mlir:UBDialect",
          "@triton//third_party/f2reduce",  # for Utility
          "@triton//third_party/nvidia:NVGPUDialect",
diff --git a/third_party/xla/third_party/triton/intel_xpu/workspace.bzl b/third_party/xla/third_party/triton/intel_xpu/workspace.bzl
index 1c3f756406b..594ae3e13ea 100644
--- a/third_party/xla/third_party/triton/intel_xpu/workspace.bzl
+++ b/third_party/xla/third_party/triton/intel_xpu/workspace.bzl
@@ -1,7 +1,7 @@
 """Intel XPU Triton archive metadata."""
 
-XPU_TRITON_COMMIT = "6046064c56a27193423475e0a6f412a08b13c6e6"
-XPU_TRITON_SHA256 = "158213c620f555aaf295613f48240c63a6b40dbe7db3805f97f4c6f2361b7f7d"
+XPU_TRITON_COMMIT = "616ef4c791a0feb79243689b985ca2a2afd9fa6b"
+XPU_TRITON_SHA256 = "e51b47dd8f6a59bd3574964337fc8b1deef04c34a2ff8fbe2a4ba2cba134f1cf"
 
 def use_xpu_triton(repository_ctx):
     return repository_ctx.getenv("ENABLE_INTEL_XPU_TRITON", "").strip() == "1"

From 8934dff38782949ff561feedbf69a2987206e0e7 Mon Sep 17 00:00:00 2001
From: Alex Pivovarov <upwind@google.com>
Date: Mon, 12 Oct 2026 12:56:24 -0700
Subject: [PATCH 0620/1000] [XLA:GPU] Add batch=1 and -0.0125 test cases to
 `approx_topk_special_floats_test`

Previously, `ApproxTopKWithSpecialFloats` and `ApproxTopKWithSpecialBFloat16` in `xla/tests/approx_topk_special_floats_test.cc` only tested `batch = 8`. On Blackwell GPUs (e.g. B200), the 32-bit cost heuristic in `IsCubSortFasterOnBlackwell` does not rewrite 2-operand `hlo.sort` on `f32[8, n]` to CUB radix sort (`xla.gpu.ext.cub_sort_pairs`), so `batch = 8` only exercised the bitonic sort path (`StableSortExpander` + bitonic sort).

### Why this change is needed:
- Add `batch = 1` in addition to `batch = 8` so `ApproxTopKWithSpecialFloats` exercises both the bitonic sort path (`n < 4096`) and the CUB radix sort path (`xla.gpu.ext.cub_sort_pairs` for `batch = 1, n >= 4096`). Unlike `hlo.topk` (`order=TOTAL`), `ApproxTopK` uses stable `PARTIAL` order where `+0.0 == -0.0`, which matches CUB radix sort's `ProcessFloatMinusZero` behavior.
- Add `-0.0125` (`neg_0_0125`) to the test inputs to provide symmetric coverage of negative normal values alongside `+0.0125` (`pos_0_0125`), negative subnormals, and `-64.5`.

PiperOrigin-RevId: 998099656
---
 .../tests/approx_topk_special_floats_test.cc  | 186 +++++++++---------
 1 file changed, 96 insertions(+), 90 deletions(-)

diff --git a/third_party/xla/xla/tests/approx_topk_special_floats_test.cc b/third_party/xla/xla/tests/approx_topk_special_floats_test.cc
index 23b7a7e9548..00cf1e14dc1 100644
--- a/third_party/xla/xla/tests/approx_topk_special_floats_test.cc
+++ b/third_party/xla/xla/tests/approx_topk_special_floats_test.cc
@@ -50,19 +50,20 @@ TEST_F(TopkTest, ApproxTopKWithSpecialFloats) {
 
   float pos_64_5 = absl::bit_cast<float>(0x42810000u);
   float pos_0_0125 = absl::bit_cast<float>(0x3C4d0000u);
+  float neg_0_0125 = absl::bit_cast<float>(0xBC4d0000u);
   float neg_64_5 = absl::bit_cast<float>(0xC2810000u);
 
   std::vector<float> input_row_start = {
-      neg_inf,       pos_inf,       neg_dnrm_1e39, pos_zero,     neg_zero,
-      neg_64_5,      pos_inf,       neg_64_5,      neg_zero,     pos_dnrm_1e40,
-      neg_dnrm_9e41, pos_zero,      neg_inf,       pos_64_5,     pos_0_0125,
-      pos_dnrm_1e38, neg_dnrm_1e39, pos_dnrm_1e40, neg_dnrm_9e41};
+      neg_inf,       pos_inf,       neg_dnrm_1e39, pos_zero,      neg_zero,
+      neg_64_5,      pos_inf,       neg_64_5,      neg_zero,      pos_dnrm_1e40,
+      neg_dnrm_9e41, pos_zero,      neg_inf,       pos_64_5,      pos_0_0125,
+      pos_dnrm_1e38, neg_dnrm_1e39, pos_dnrm_1e40, neg_dnrm_9e41, neg_0_0125};
 
   std::vector<float> expected_values_64 = {
       pos_inf,       pos_inf,       pos_64_5,      pos_0_0125,    pos_dnrm_1e38,
       pos_dnrm_1e40, pos_dnrm_1e40, pos_zero,      neg_zero,      neg_zero,
       pos_zero,      neg_dnrm_9e41, neg_dnrm_9e41, neg_dnrm_1e39, neg_dnrm_1e39,
-      neg_64_5,      neg_64_5,      neg_inf,       neg_inf,       neg_inf,
+      neg_0_0125,    neg_64_5,      neg_64_5,      neg_inf,       neg_inf,
       neg_inf,       neg_inf,       neg_inf,       neg_inf,       neg_inf,
       neg_inf,       neg_inf,       neg_inf,       neg_inf,       neg_inf,
       neg_inf,       neg_inf,       neg_inf,       neg_inf,       neg_inf,
@@ -74,23 +75,24 @@ TEST_F(TopkTest, ApproxTopKWithSpecialFloats) {
       neg_inf,       neg_inf,       neg_inf,       neg_inf};
 
   std::vector<int32_t> expected_indices_64 = {
-      1,  6,  13, 14, 15, 9,  17, 3,  4,  8,  11, 10, 18, 2,  16, 5,
-      7,  0,  12, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31,
+      1,  6,  13, 14, 15, 9,  17, 3,  4,  8,  11, 10, 18, 2,  16, 19,
+      5,  7,  0,  12, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31,
       32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47,
       48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63};
 
-  for (int n : {512, 1024, 32768, 65536 + 4096}) {
-    Array2D<float> input(8, n, neg_inf);
-    for (int i = 0; i < input_row_start.size(); ++i) {
-      for (int j = 0; j < 8; ++j) {
-        input(j, i) = input_row_start[i];
+  for (int batch : {1, 8}) {
+    for (int n : {512, 1024, 4096, 32768, 65536 + 4096}) {
+      Array2D<float> input(batch, n, neg_inf);
+      for (int i = 0; i < input_row_start.size(); ++i) {
+        for (int j = 0; j < batch; ++j) {
+          input(j, i) = input_row_start[i];
+        }
       }
-    }
 
-    auto literal = LiteralUtil::CreateR2FromArray2D<float>(input);
+      auto literal = LiteralUtil::CreateR2FromArray2D<float>(input);
 
-    for (int k : {2, 4, 8, 16, 32, 64}) {
-      const auto hlo_text_module = absl::Substitute(R"(
+      for (int k : {2, 4, 8, 16, 32, 64}) {
+        const auto hlo_text_module = absl::Substitute(R"(
       HloModule topk
 
       %top_k_gt_f32_comparator.4 (approx_top_k.35: f32[], approx_top_k.36: f32[], approx_top_k.37: s32[], approx_top_k.38: s32[]) -> pred[] {
@@ -102,41 +104,42 @@ TEST_F(TopkTest, ApproxTopKWithSpecialFloats) {
       }
 
       ENTRY TopK {
-        x = f32[8,$0]{1,0} parameter(0)
-        %approx_top_k.60 = s32[$0]{0} iota(), iota_dimension=0
-        %approx_top_k.61 = s32[8,$0]{1,0} broadcast(%approx_top_k.60), dimensions={1}
-        %approx_top_k.62 = (f32[8,$0]{1,0}, s32[8,$0]{1,0}) sort(x, %approx_top_k.61), dimensions={1}, is_stable=true, to_apply=%top_k_gt_f32_comparator.4
-        %approx_top_k.64 = f32[8,$1]{1,0} slice(%approx_top_k.62#0), slice={[0:8], [0:$1]}
-        %approx_top_k.66 = s32[8,$1]{1,0} slice(%approx_top_k.62#1), slice={[0:8], [0:$1]}
-        ROOT %approx_top_k.67 = (f32[8,$1]{1,0}, s32[8,$1]{1,0}) tuple(%approx_top_k.64, %approx_top_k.66)
+        x = f32[$0,$1]{1,0} parameter(0)
+        %approx_top_k.60 = s32[$1]{0} iota(), iota_dimension=0
+        %approx_top_k.61 = s32[$0,$1]{1,0} broadcast(%approx_top_k.60), dimensions={1}
+        %approx_top_k.62 = (f32[$0,$1]{1,0}, s32[$0,$1]{1,0}) sort(x, %approx_top_k.61), dimensions={1}, is_stable=true, to_apply=%top_k_gt_f32_comparator.4
+        %approx_top_k.64 = f32[$0,$2]{1,0} slice(%approx_top_k.62#0), slice={[0:$0], [0:$2]}
+        %approx_top_k.66 = s32[$0,$2]{1,0} slice(%approx_top_k.62#1), slice={[0:$0], [0:$2]}
+        ROOT %approx_top_k.67 = (f32[$0,$2]{1,0}, s32[$0,$2]{1,0}) tuple(%approx_top_k.64, %approx_top_k.66)
       }
     )",
-                                                    n, k);
+                                                      batch, n, k);
 
-      ASSERT_OK_AND_ASSIGN(auto module,
-                           ParseAndReturnVerifiedModule(hlo_text_module));
+        ASSERT_OK_AND_ASSIGN(auto module,
+                             ParseAndReturnVerifiedModule(hlo_text_module));
 
-      Array2D<float> expected_vals_arr(8, k);
-      Array2D<int32_t> expected_idxs_arr(8, k);
-      for (int i = 0; i < 8; ++i) {
-        for (int j = 0; j < k; ++j) {
-          expected_vals_arr(i, j) = expected_values_64[j];
-          expected_idxs_arr(i, j) = expected_indices_64[j];
+        Array2D<float> expected_vals_arr(batch, k);
+        Array2D<int32_t> expected_idxs_arr(batch, k);
+        for (int i = 0; i < batch; ++i) {
+          for (int j = 0; j < k; ++j) {
+            expected_vals_arr(i, j) = expected_values_64[j];
+            expected_idxs_arr(i, j) = expected_indices_64[j];
+          }
         }
+
+        auto expected_vals_literal =
+            LiteralUtil::CreateR2FromArray2D<float>(expected_vals_arr);
+        auto expected_idxs_literal =
+            LiteralUtil::CreateR2FromArray2D<int32_t>(expected_idxs_arr);
+        auto expected_tuple = LiteralUtil::MakeTuple(
+            {&expected_vals_literal, &expected_idxs_literal});
+
+        ASSERT_OK_AND_ASSIGN(Literal actual_tuple,
+                             Execute(std::move(module), {&literal}));
+
+        EXPECT_TRUE(LiteralTestUtil::Near(expected_tuple, actual_tuple,
+                                          xla::ErrorSpec{0, 0}));
       }
-
-      auto expected_vals_literal =
-          LiteralUtil::CreateR2FromArray2D<float>(expected_vals_arr);
-      auto expected_idxs_literal =
-          LiteralUtil::CreateR2FromArray2D<int32_t>(expected_idxs_arr);
-      auto expected_tuple = LiteralUtil::MakeTuple(
-          {&expected_vals_literal, &expected_idxs_literal});
-
-      ASSERT_OK_AND_ASSIGN(Literal actual_tuple,
-                           Execute(std::move(module), {&literal}));
-
-      EXPECT_TRUE(LiteralTestUtil::Near(expected_tuple, actual_tuple,
-                                        xla::ErrorSpec{0, 0}));
     }
   }
 }
@@ -158,19 +161,20 @@ TEST_F(TopkTest, ApproxTopKWithSpecialBFloat16) {
 
   bfloat16 pos_64_5 = ToBfloat16(0x4281u);
   bfloat16 pos_0_0125 = ToBfloat16(0x3C4du);
+  bfloat16 neg_0_0125 = ToBfloat16(0xBC4du);
   bfloat16 neg_64_5 = ToBfloat16(0xC281u);
 
   std::vector<bfloat16> input_row_start = {
-      neg_inf,       pos_inf,       neg_dnrm_1e39, pos_zero,     neg_zero,
-      neg_64_5,      pos_inf,       neg_64_5,      neg_zero,     pos_dnrm_1e40,
-      neg_dnrm_9e41, pos_zero,      neg_inf,       pos_64_5,     pos_0_0125,
-      pos_dnrm_1e38, neg_dnrm_1e39, pos_dnrm_1e40, neg_dnrm_9e41};
+      neg_inf,       pos_inf,       neg_dnrm_1e39, pos_zero,      neg_zero,
+      neg_64_5,      pos_inf,       neg_64_5,      neg_zero,      pos_dnrm_1e40,
+      neg_dnrm_9e41, pos_zero,      neg_inf,       pos_64_5,      pos_0_0125,
+      pos_dnrm_1e38, neg_dnrm_1e39, pos_dnrm_1e40, neg_dnrm_9e41, neg_0_0125};
 
   std::vector<bfloat16> expected_values_64 = {
       pos_inf,       pos_inf,       pos_64_5,      pos_0_0125,    pos_dnrm_1e38,
       pos_dnrm_1e40, pos_dnrm_1e40, pos_zero,      neg_zero,      neg_zero,
       pos_zero,      neg_dnrm_9e41, neg_dnrm_9e41, neg_dnrm_1e39, neg_dnrm_1e39,
-      neg_64_5,      neg_64_5,      neg_inf,       neg_inf,       neg_inf,
+      neg_0_0125,    neg_64_5,      neg_64_5,      neg_inf,       neg_inf,
       neg_inf,       neg_inf,       neg_inf,       neg_inf,       neg_inf,
       neg_inf,       neg_inf,       neg_inf,       neg_inf,       neg_inf,
       neg_inf,       neg_inf,       neg_inf,       neg_inf,       neg_inf,
@@ -182,23 +186,24 @@ TEST_F(TopkTest, ApproxTopKWithSpecialBFloat16) {
       neg_inf,       neg_inf,       neg_inf,       neg_inf};
 
   std::vector<int32_t> expected_indices_64 = {
-      1,  6,  13, 14, 15, 9,  17, 3,  4,  8,  11, 10, 18, 2,  16, 5,
-      7,  0,  12, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31,
+      1,  6,  13, 14, 15, 9,  17, 3,  4,  8,  11, 10, 18, 2,  16, 19,
+      5,  7,  0,  12, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31,
       32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47,
       48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63};
 
-  for (int n : {512, 1024, 32768, 65536 + 4096}) {
-    Array2D<bfloat16> input(8, n, neg_inf);
-    for (int i = 0; i < input_row_start.size(); ++i) {
-      for (int j = 0; j < 8; ++j) {
-        input(j, i) = input_row_start[i];
+  for (int batch : {1, 8}) {
+    for (int n : {512, 1024, 4096, 32768, 65536 + 4096}) {
+      Array2D<bfloat16> input(batch, n, neg_inf);
+      for (int i = 0; i < input_row_start.size(); ++i) {
+        for (int j = 0; j < batch; ++j) {
+          input(j, i) = input_row_start[i];
+        }
       }
-    }
 
-    auto literal = LiteralUtil::CreateR2FromArray2D<bfloat16>(input);
+      auto literal = LiteralUtil::CreateR2FromArray2D<bfloat16>(input);
 
-    for (int k : {2, 4, 8, 16, 32, 64}) {
-      const auto hlo_text_module = absl::Substitute(R"(
+      for (int k : {2, 4, 8, 16, 32, 64}) {
+        const auto hlo_text_module = absl::Substitute(R"(
       HloModule topk
 
       %top_k_gt_bf16_comparator.4 (approx_top_k.35: bf16[], approx_top_k.36: bf16[], approx_top_k.37: s32[], approx_top_k.38: s32[]) -> pred[] {
@@ -210,41 +215,42 @@ TEST_F(TopkTest, ApproxTopKWithSpecialBFloat16) {
       }
 
       ENTRY TopK {
-        x = bf16[8,$0]{1,0} parameter(0)
-        %approx_top_k.60 = s32[$0]{0} iota(), iota_dimension=0
-        %approx_top_k.61 = s32[8,$0]{1,0} broadcast(%approx_top_k.60), dimensions={1}
-        %approx_top_k.62 = (bf16[8,$0]{1,0}, s32[8,$0]{1,0}) sort(x, %approx_top_k.61), dimensions={1}, is_stable=true, to_apply=%top_k_gt_bf16_comparator.4
-        %approx_top_k.64 = bf16[8,$1]{1,0} slice(%approx_top_k.62#0), slice={[0:8], [0:$1]}
-        %approx_top_k.66 = s32[8,$1]{1,0} slice(%approx_top_k.62#1), slice={[0:8], [0:$1]}
-        ROOT %approx_top_k.67 = (bf16[8,$1]{1,0}, s32[8,$1]{1,0}) tuple(%approx_top_k.64, %approx_top_k.66)
+        x = bf16[$0,$1]{1,0} parameter(0)
+        %approx_top_k.60 = s32[$1]{0} iota(), iota_dimension=0
+        %approx_top_k.61 = s32[$0,$1]{1,0} broadcast(%approx_top_k.60), dimensions={1}
+        %approx_top_k.62 = (bf16[$0,$1]{1,0}, s32[$0,$1]{1,0}) sort(x, %approx_top_k.61), dimensions={1}, is_stable=true, to_apply=%top_k_gt_bf16_comparator.4
+        %approx_top_k.64 = bf16[$0,$2]{1,0} slice(%approx_top_k.62#0), slice={[0:$0], [0:$2]}
+        %approx_top_k.66 = s32[$0,$2]{1,0} slice(%approx_top_k.62#1), slice={[0:$0], [0:$2]}
+        ROOT %approx_top_k.67 = (bf16[$0,$2]{1,0}, s32[$0,$2]{1,0}) tuple(%approx_top_k.64, %approx_top_k.66)
       }
     )",
-                                                    n, k);
+                                                      batch, n, k);
 
-      ASSERT_OK_AND_ASSIGN(auto module,
-                           ParseAndReturnVerifiedModule(hlo_text_module));
+        ASSERT_OK_AND_ASSIGN(auto module,
+                             ParseAndReturnVerifiedModule(hlo_text_module));
 
-      Array2D<bfloat16> expected_vals_arr(8, k);
-      Array2D<int32_t> expected_idxs_arr(8, k);
-      for (int i = 0; i < 8; ++i) {
-        for (int j = 0; j < k; ++j) {
-          expected_vals_arr(i, j) = expected_values_64[j];
-          expected_idxs_arr(i, j) = expected_indices_64[j];
+        Array2D<bfloat16> expected_vals_arr(batch, k);
+        Array2D<int32_t> expected_idxs_arr(batch, k);
+        for (int i = 0; i < batch; ++i) {
+          for (int j = 0; j < k; ++j) {
+            expected_vals_arr(i, j) = expected_values_64[j];
+            expected_idxs_arr(i, j) = expected_indices_64[j];
+          }
         }
+
+        auto expected_vals_literal =
+            LiteralUtil::CreateR2FromArray2D<bfloat16>(expected_vals_arr);
+        auto expected_idxs_literal =
+            LiteralUtil::CreateR2FromArray2D<int32_t>(expected_idxs_arr);
+        auto expected_tuple = LiteralUtil::MakeTuple(
+            {&expected_vals_literal, &expected_idxs_literal});
+
+        ASSERT_OK_AND_ASSIGN(Literal actual_tuple,
+                             Execute(std::move(module), {&literal}));
+
+        EXPECT_TRUE(LiteralTestUtil::Near(expected_tuple, actual_tuple,
+                                          xla::ErrorSpec{0, 0}));
       }
-
-      auto expected_vals_literal =
-          LiteralUtil::CreateR2FromArray2D<bfloat16>(expected_vals_arr);
-      auto expected_idxs_literal =
-          LiteralUtil::CreateR2FromArray2D<int32_t>(expected_idxs_arr);
-      auto expected_tuple = LiteralUtil::MakeTuple(
-          {&expected_vals_literal, &expected_idxs_literal});
-
-      ASSERT_OK_AND_ASSIGN(Literal actual_tuple,
-                           Execute(std::move(module), {&literal}));
-
-      EXPECT_TRUE(LiteralTestUtil::Near(expected_tuple, actual_tuple,
-                                        xla::ErrorSpec{0, 0}));
     }
   }
 }

From b2ea5f1b69b57a7955931ef904573f50710590e7 Mon Sep 17 00:00:00 2001
From: Emilio Cota <ecg@google.com>
Date: Mon, 12 Oct 2026 13:18:46 -0700
Subject: [PATCH 0621/1000] [xla:xtile] xtile.extract: initialize the padding
 of partial tiles

When bufferizing xtile.extract for partial/boundary tiles, a full tile
is allocated via memref.alloc but only an in-bounds subview is
copied into it. The remaining (padding) lanes are left uninitialized.

This can result in UB in LLVM IR if the uninitialized (undef) values
are later operated on (e.g. in elementwise ops that operate on the
full tile before masking), and also would generate msan warnings
if it weren't for the InitializeAllocsPass, which under msan initializes
(and unpoisons) all tiles in full.

Fix this potential UB and prepare for the upcoming removal of
InitializeAllocsPass (its blanket unpoisonings will become actively
harmful once the upcoming msan support lands, because they might hide
actual bugs) by initializing the partial tile buffers after allocation,
i.e. before any subviews are copied into them, as follows:

1. A single 1D splat vector.store (vector<num_elements x element_type>) at
   [0, ..., 0] for standard integer/float element types (bitwidth >= 8),
   which lowers directly to a single flat LLVM vector store without emitting
   loops or basic blocks.
2. The existing `scf::buildLoopNest` fallback for `ComplexType` and sub-byte
   types.

This keeps the performance (runtime and compile time) impact of this
change to a negligible level.

PiperOrigin-RevId: 998113295
---
 third_party/xla/xla/codegen/xtile/ir/BUILD    |  2 +
 .../xtile/ir/transforms/tests/bufferize.mlir  | 82 +++++++++++++++++++
 .../codegen/xtile/ir/xtile_bufferization.cc   | 72 ++++++++++++++++
 .../xla/xla/codegen/xtile/ir/xtile_dialect.h  |  2 +
 .../xla/xla/codegen/xtile/ir/xtile_dialect.td |  2 +
 5 files changed, 160 insertions(+)

diff --git a/third_party/xla/xla/codegen/xtile/ir/BUILD b/third_party/xla/xla/codegen/xtile/ir/BUILD
index bd43d8f3e8e..6e3aa668f07 100644
--- a/third_party/xla/xla/codegen/xtile/ir/BUILD
+++ b/third_party/xla/xla/codegen/xtile/ir/BUILD
@@ -116,6 +116,7 @@ cc_library(
         "@llvm-project//mlir:BufferizationDialect",
         "@llvm-project//mlir:BufferizationInterfaces",
         "@llvm-project//mlir:BytecodeOpInterface",
+        "@llvm-project//mlir:ComplexDialect",
         "@llvm-project//mlir:DialectUtils",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:FunctionInterfaces",
@@ -127,6 +128,7 @@ cc_library(
         "@llvm-project//mlir:SideEffectInterfaces",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:VectorDialect",
         "@stablehlo//:stablehlo_ops",
     ],
 )
diff --git a/third_party/xla/xla/codegen/xtile/ir/transforms/tests/bufferize.mlir b/third_party/xla/xla/codegen/xtile/ir/transforms/tests/bufferize.mlir
index a9c4ead3dc3..97dcb28a138 100644
--- a/third_party/xla/xla/codegen/xtile/ir/transforms/tests/bufferize.mlir
+++ b/third_party/xla/xla/codegen/xtile/ir/transforms/tests/bufferize.mlir
@@ -17,6 +17,7 @@
 
 // CHECK: @extract_strided(%[[SOURCE:.*]]: memref<16xf32>, %[[OFFSET:.*]]: index)
 func.func @extract_strided(%source: memref<16xf32>, %tile_id: index) -> tensor<8xf32> {
+  // CHECK-DAG: %[[PAD:.*]] = arith.constant dense<0x7FC00000> : vector<8xf32>
   // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
   // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
   // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
@@ -39,6 +40,7 @@ func.func @extract_strided(%source: memref<16xf32>, %tile_id: index) -> tensor<8
     // CHECK-SAME: : memref<16xf32> to memref<?xf32, strided<[2], offset: ?>>
 
     // CHECK: %[[ALLOC_1:.*]] = memref.alloc() : memref<8xf32>
+    // CHECK: vector.store %[[PAD]], %[[ALLOC_1]][%{{.*}}] : memref<8xf32>, vector<8xf32>
 
     // CHECK: %[[ALLOC_1_SUBVIEW:.*]] = memref.subview %[[ALLOC_1]]
     // CHECK-SAME: [0] [%[[SIZE]]] [1] : memref<8xf32> to memref<?xf32, strided<[1]>>
@@ -114,6 +116,7 @@ func.func @extract_identity(%source: memref<16xf32>, %tile_id: index) -> tensor<
   // CHECK: } else {
   // CHECK:   %[[INPUT_SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[OFFSET]]] [%[[SIZE]]] [1] : memref<16xf32> to memref<?xf32, strided<[1], offset: ?>>
   // CHECK:   %[[ALLOC_ELSE:.*]] = memref.alloc() : memref<8xf32>
+  // CHECK:   vector.store %{{.*}}, %[[ALLOC_ELSE]][%{{.*}}] : memref<8xf32>, vector<8xf32>
   // CHECK:   %[[ALLOC_SUBVIEW:.*]] = memref.subview %[[ALLOC_ELSE]][0] [%[[SIZE]]] [1] : memref<8xf32> to memref<?xf32, strided<[1]>>
   // CHECK:   memref.copy %[[INPUT_SUBVIEW]], %[[ALLOC_SUBVIEW]] : memref<?xf32, strided<[1], offset: ?>> to memref<?xf32, strided<[1]>>
   // CHECK:   scf.yield %[[ALLOC_ELSE]] : memref<8xf32>
@@ -141,6 +144,7 @@ func.func @extract_5d(%source: memref<1x2x1x32768x256xf32>, %d1: index, %d3: ind
   // CHECK: } else {
   // CHECK:   %[[INPUT_SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[C_0]], %[[D1]], %[[C_0]], %[[D3]], %[[D4]]] [1, 1, 1, 1, %[[SIZE_4]]] [0, 1, 0, 1, 1] : memref<1x2x1x32768x256xf32> to memref<1x1x1x1x?xf32, strided<[0, 8388608, 0, 256, 1], offset: ?>>
   // CHECK:   %[[ALLOC_ELSE:.*]] = memref.alloc() : memref<1x1x1x1x16xf32>
+  // CHECK:   vector.store %{{.*}}, %[[ALLOC_ELSE]][%[[C_0]], %[[C_0]], %[[C_0]], %[[C_0]], %[[C_0]]] : memref<1x1x1x1x16xf32>, vector<16xf32>
   // CHECK:   %[[ALLOC_SUBVIEW:.*]] = memref.subview %[[ALLOC_ELSE]][0, 0, 0, 0, 0] [1, 1, 1, 1, %[[SIZE_4]]] [1, 1, 1, 1, 1] : memref<1x1x1x1x16xf32> to memref<1x1x1x1x?xf32, strided<[16, 16, 16, 16, 1]>>
   // CHECK:   memref.copy %[[INPUT_SUBVIEW]], %[[ALLOC_SUBVIEW]]
   // CHECK:   scf.yield %[[ALLOC_ELSE]] : memref<1x1x1x1x16xf32>
@@ -157,6 +161,7 @@ func.func @extract_5d(%source: memref<1x2x1x32768x256xf32>, %d1: index, %d3: ind
 // CHECK: @extract_static(%[[SOURCE:.*]]: memref<16xf32>)
 func.func @extract_static(%source: memref<16xf32>) -> tensor<8xf32> {
   // CHECK-NOT: scf.if
+  // CHECK-NOT: vector.store
   // CHECK: %[[SUBVIEW:.*]] = memref.subview %[[SOURCE]][0] [8] [1] : memref<16xf32> to memref<8xf32, strided<[1]>>
   // CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<8xf32>
   // CHECK: memref.copy %[[SUBVIEW]], %[[ALLOC]] : memref<8xf32, strided<[1]>> to memref<8xf32>
@@ -193,6 +198,7 @@ func.func @insert_in_bounds(%source: tensor<8xf32>, %destination: memref<16xf32>
 // CHECK: @extract_strided_in_bounds(%[[SOURCE:.*]]: memref<17xf32>, %[[TILE_ID:.*]]: index)
 func.func @extract_strided_in_bounds(%source: memref<17xf32>, %tile_id: index) -> tensor<8xf32> {
   // CHECK-NOT: scf.if
+  // CHECK-NOT: vector.store
   // CHECK: %[[OFFSET:.*]] = xla.apply_indexing #[[MAP]](%[[TILE_ID]])
   // CHECK: %[[SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[OFFSET]]] [8] [2] : memref<17xf32> to memref<8xf32, strided<[2], offset: ?>>
   // CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<8xf32>
@@ -237,6 +243,7 @@ func.func @extract_rank_reduced(%source: memref<2x8x16xf32>, %d0: index, %d1: in
   // CHECK: } else {
   // CHECK:   %[[INPUT_SUBVIEW:.*]] = memref.subview %[[SOURCE]][%[[D0]], %[[OFFSET_1]], %[[D2]]] [1, 4, %[[SIZE_2]]] [1, 1, 1] : memref<2x8x16xf32> to memref<4x?xf32, strided<[16, 1], offset: ?>>
   // CHECK:   %[[ALLOC_ELSE:.*]] = memref.alloc() : memref<4x8xf32>
+  // CHECK:   vector.store %{{.*}}, %[[ALLOC_ELSE]][%{{.*}}, %{{.*}}] : memref<4x8xf32>, vector<32xf32>
   // CHECK:   %[[ALLOC_SUBVIEW:.*]] = memref.subview %[[ALLOC_ELSE]][0, 0] [4, %[[SIZE_2]]] [1, 1] : memref<4x8xf32> to memref<4x?xf32, strided<[8, 1]>>
   // CHECK:   memref.copy %[[INPUT_SUBVIEW]], %[[ALLOC_SUBVIEW]]
   // CHECK:   scf.yield %[[ALLOC_ELSE]] : memref<4x8xf32>
@@ -276,3 +283,78 @@ func.func @insert_rank_reduced(%source: tensor<4x1x8xf32>, %destination: memref<
   return
 }
 
+// -----
+
+// CHECK-LABEL: @extract_int(
+func.func @extract_int(%source: memref<16xi32>, %tile_id: index) -> tensor<8xi32> {
+  // CHECK-DAG: %[[PAD:.*]] = arith.constant dense<-1> : vector<8xi32>
+  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+  // CHECK: scf.if {{.*}} -> (memref<8xi32>) {
+  // CHECK: } else {
+  // CHECK:   %[[ALLOC_ELSE:.*]] = memref.alloc() : memref<8xi32>
+  // CHECK:   vector.store %[[PAD]], %[[ALLOC_ELSE]][%[[C0]]] : memref<8xi32>, vector<8xi32>
+  // CHECK: }
+  %tile = xtile.extract %source[%tile_id][8][1] : memref<16xi32> -> tensor<8xi32>
+  return %tile : tensor<8xi32>
+}
+
+// -----
+
+// CHECK-LABEL: @extract_complex(
+func.func @extract_complex(%source: memref<16xcomplex<f32>>, %tile_id: index) -> tensor<8xcomplex<f32>> {
+  // CHECK-DAG: %[[NAN:.*]] = arith.constant 0x7FC00000 : f32
+  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+  // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+  // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
+  // CHECK: scf.if {{.*}} -> (memref<8xcomplex<f32>>) {
+  // CHECK: } else {
+  // CHECK:   %[[ALLOC_ELSE:.*]] = memref.alloc() : memref<8xcomplex<f32>>
+  // CHECK:   %[[PAD:.*]] = complex.create %[[NAN]], %[[NAN]] : complex<f32>
+  // CHECK:   scf.for %[[IV:.*]] = %[[C0]] to %[[C8]] step %[[C1]] {
+  // CHECK:     memref.store %[[PAD]], %[[ALLOC_ELSE]][%[[IV]]] : memref<8xcomplex<f32>>
+  // CHECK:   }
+  // CHECK: }
+  %tile = xtile.extract %source[%tile_id][8][1] : memref<16xcomplex<f32>> -> tensor<8xcomplex<f32>>
+  return %tile : tensor<8xcomplex<f32>>
+}
+
+// -----
+
+// CHECK-LABEL: @extract_subbyte_int(
+func.func @extract_subbyte_int(%source: memref<4x16xi4>, %d0: index, %d1: index) -> tensor<2x8xi4> {
+  // CHECK-DAG: %[[PAD:.*]] = arith.constant -1 : i4
+  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+  // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+  // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+  // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
+  // CHECK: scf.if {{.*}} -> (memref<2x8xi4>) {
+  // CHECK: } else {
+  // CHECK:   %[[ALLOC_ELSE:.*]] = memref.alloc() : memref<2x8xi4>
+  // CHECK:   scf.for %[[IV0:.*]] = %[[C0]] to %[[C2]] step %[[C1]] {
+  // CHECK:     scf.for %[[IV1:.*]] = %[[C0]] to %[[C8]] step %[[C1]] {
+  // CHECK:       memref.store %[[PAD]], %[[ALLOC_ELSE]][%[[IV0]], %[[IV1]]] : memref<2x8xi4>
+  // CHECK:     }
+  // CHECK:   }
+  // CHECK: }
+  %tile = xtile.extract %source[%d0, %d1][2, 8][1, 1] : memref<4x16xi4> -> tensor<2x8xi4>
+  return %tile : tensor<2x8xi4>
+}
+
+// -----
+
+// CHECK-LABEL: @extract_finite_only_float(
+func.func @extract_finite_only_float(%source: memref<16xf4E2M1FN>, %tile_id: index) -> tensor<8xf4E2M1FN> {
+  // CHECK-DAG: %[[PAD:.*]] = arith.constant 0.000000e+00 : f4E2M1FN
+  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+  // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+  // CHECK-DAG: %[[C8:.*]] = arith.constant 8 : index
+  // CHECK: scf.if {{.*}} -> (memref<8xf4E2M1FN>) {
+  // CHECK: } else {
+  // CHECK:   %[[ALLOC_ELSE:.*]] = memref.alloc() : memref<8xf4E2M1FN>
+  // CHECK:   scf.for %[[IV:.*]] = %[[C0]] to %[[C8]] step %[[C1]] {
+  // CHECK:     memref.store %[[PAD]], %[[ALLOC_ELSE]][%[[IV]]] : memref<8xf4E2M1FN>
+  // CHECK:   }
+  // CHECK: }
+  %tile = xtile.extract %source[%tile_id][8][1] : memref<16xf4E2M1FN> -> tensor<8xf4E2M1FN>
+  return %tile : tensor<8xf4E2M1FN>
+}
diff --git a/third_party/xla/xla/codegen/xtile/ir/xtile_bufferization.cc b/third_party/xla/xla/codegen/xtile/ir/xtile_bufferization.cc
index 72ca4fdd12c..72d30df898e 100644
--- a/third_party/xla/xla/codegen/xtile/ir/xtile_bufferization.cc
+++ b/third_party/xla/xla/codegen/xtile/ir/xtile_bufferization.cc
@@ -19,19 +19,25 @@ limitations under the License.
 #include <optional>
 
 #include "absl/log/check.h"
+#include "llvm/ADT/APFloat.h"
+#include "llvm/ADT/APInt.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/SmallVectorExtras.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/LogicalResult.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributeInterfaces.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypeInterfaces.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -215,6 +221,38 @@ static std::optional<mlir::Value> TileIsFullSize(
   return is_full_size;
 }
 
+// Returns a padding attribute chosen to surface unintended reads of
+// out-of-bounds lanes in tests.
+static mlir::TypedAttr GetPaddingAttr(mlir::ImplicitLocOpBuilder& builder,
+                                      mlir::Type element_type) {
+  if (auto float_type = mlir::dyn_cast<mlir::FloatType>(element_type)) {
+    const llvm::fltSemantics& semantics = float_type.getFloatSemantics();
+    return builder.getFloatAttr(float_type,
+                                llvm::APFloat::semanticsHasNaN(semantics)
+                                    ? llvm::APFloat::getNaN(semantics)
+                                    : llvm::APFloat::getZero(semantics));
+  }
+  if (auto int_type = mlir::dyn_cast<mlir::IntegerType>(element_type)) {
+    return builder.getIntegerAttr(
+        int_type, llvm::APInt::getAllOnes(int_type.getIntOrFloatBitWidth()));
+  }
+  return nullptr;
+}
+
+static mlir::Value GetPaddingValue(mlir::ImplicitLocOpBuilder& builder,
+                                   mlir::Type element_type) {
+  if (mlir::TypedAttr attr = GetPaddingAttr(builder, element_type)) {
+    return mlir::arith::ConstantOp::create(builder, element_type, attr);
+  }
+  if (auto complex_type = mlir::dyn_cast<mlir::ComplexType>(element_type)) {
+    mlir::Value component =
+        GetPaddingValue(builder, complex_type.getElementType());
+    return mlir::complex::CreateOp::create(builder, complex_type, component,
+                                           component);
+  }
+  llvm_unreachable("unhandled element type");
+}
+
 // Get a buffer copied from the original buffer that is padded to the full tile
 // size.
 static mlir::TypedValue<mlir::MemRefType> GetPaddedTileBuffer(
@@ -225,6 +263,40 @@ static mlir::TypedValue<mlir::MemRefType> GetPaddedTileBuffer(
       builder, GetStaticFoldResult(builder, tile_type.getShape()),
       tile_type.getElementType());
 
+  mlir::Type element_type = tile_type.getElementType();
+  llvm::SmallVector<mlir::Value> zeros(
+      tile_type.getRank(), mlir::arith::ConstantIndexOp::create(builder, 0));
+
+  if (element_type.isIntOrFloat() &&
+      element_type.getIntOrFloatBitWidth() >= 8) {
+    // A 1-D vector.store at [0, ..., 0] initializes the entire contiguous
+    // buffer in a single flat store (0-D tiles are always in-bounds and never
+    // reach here).
+    DCHECK_GE(tile_type.getRank(), 1);
+    auto vector_type =
+        mlir::VectorType::get({tile_type.getNumElements()}, element_type);
+    auto padding_vector = mlir::arith::ConstantOp::create(
+        builder, vector_type,
+        mlir::DenseElementsAttr::get(vector_type,
+                                     GetPaddingAttr(builder, element_type)));
+    mlir::vector::StoreOp::create(builder, padding_vector, buffer, zeros);
+  } else {
+    mlir::Value padding_value = GetPaddingValue(builder, element_type);
+    llvm::SmallVector<mlir::Value> ubs = llvm::map_to_vector(
+        tile_type.getShape(), [&builder](int64_t dim_size) -> mlir::Value {
+          return mlir::arith::ConstantIndexOp::create(builder, dim_size);
+        });
+    llvm::SmallVector<mlir::Value> steps(
+        tile_type.getRank(), mlir::arith::ConstantIndexOp::create(builder, 1));
+    mlir::scf::buildLoopNest(
+        builder, builder.getLoc(), zeros, ubs, steps,
+        [padding_value, &buffer](mlir::OpBuilder& nested_builder,
+                                 mlir::Location loc, mlir::ValueRange ivs) {
+          mlir::memref::StoreOp::create(nested_builder, loc, padding_value,
+                                        buffer, ivs);
+        });
+  }
+
   auto local_tile_size = GetClampedTileSize(builder, op, false);
   auto local_buffer_subview =
       GetLocalBufferSubview(builder, buffer, local_tile_size,
diff --git a/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.h b/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.h
index 3c5d7585278..07762e3f4ec 100644
--- a/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.h
+++ b/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.h
@@ -17,9 +17,11 @@ limitations under the License.
 #define XLA_CODEGEN_XTILE_IR_XTILE_DIALECT_H_
 
 #include "mlir/Dialect/Arith/IR/Arith.h"  // IWYU pragma: keep
+#include "mlir/Dialect/Complex/IR/Complex.h"  // IWYU pragma: keep
 #include "mlir/Dialect/MemRef/IR/MemRef.h"  // IWYU pragma: keep
 #include "mlir/Dialect/SCF/IR/SCF.h"  // IWYU pragma: keep
 #include "mlir/Dialect/Tensor/IR/Tensor.h"  // IWYU pragma: keep
+#include "mlir/Dialect/Vector/IR/VectorOps.h"  // IWYU pragma: keep
 #include "mlir/IR/Dialect.h"  // IWYU pragma: keep
 #include "stablehlo/dialect/StablehloOps.h"  // IWYU pragma: keep
 
diff --git a/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.td b/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.td
index 38c4bbf47a1..847d43e1f95 100644
--- a/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.td
+++ b/third_party/xla/xla/codegen/xtile/ir/xtile_dialect.td
@@ -27,10 +27,12 @@ def XTileDialect : Dialect {
 
   let dependentDialects = [
     "mlir::arith::ArithDialect",
+    "mlir::complex::ComplexDialect",
     "mlir::memref::MemRefDialect",
     "mlir::tensor::TensorDialect",
     "mlir::stablehlo::StablehloDialect",
     "mlir::scf::SCFDialect",
+    "mlir::vector::VectorDialect",
   ];
 
   let useDefaultAttributePrinterParser = 1;

From 4f96467fbe78b0a84382777ace8c30ed7bdc5843 Mon Sep 17 00:00:00 2001
From: Penporn Koanantakool <penporn@google.com>
Date: Mon, 12 Oct 2026 14:06:28 -0700
Subject: [PATCH 0622/1000] [XLA] Update XLA warnings.bazelrc

PiperOrigin-RevId: 998144170
---
 third_party/xla/warnings.bazelrc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/warnings.bazelrc b/third_party/xla/warnings.bazelrc
index 2ae6e25197d..efde9adcc80 100644
--- a/third_party/xla/warnings.bazelrc
+++ b/third_party/xla/warnings.bazelrc
@@ -22,6 +22,7 @@ build:warnings --copt=-Wno-ignored-qualifiers
 build:warnings --copt=-Wno-inconsistent-missing-override
 build:warnings --copt=-Wno-nontrivial-memaccess
 build:warnings --copt=-Wno-nullability-completeness
+build:warnings --copt=-Wno-pointer-bool-conversion
 build:warnings --copt=-Wno-potentially-evaluated-expression
 build:warnings --copt=-Wno-range-loop-analysis
 build:warnings --copt=-Wno-return-type-c-linkage
@@ -29,6 +30,7 @@ build:warnings --copt=-Wno-tautological-undefined-compare
 build:warnings --copt=-Wno-unused-but-set-variable
 build:warnings --copt=-Wno-unused-lambda-capture
 build:warnings --copt=-Wno-unused-local-typedef
+build:warnings --copt=-Wno-unused-template
 build:warnings --copt=-Wno-deprecated-volatile
 build:warnings --copt=-Wno-deprecated-anon-enum-enum-conversion
 build:warnings --copt=-Wno-deprecated-enum-compare
@@ -91,8 +93,6 @@ build:warnings --copt=-Wno-invalid-offsetof
 build:warnings --copt=-Wno-final-dtor-non-final-class
 build:warnings --copt=-Wnon-virtual-dtor
 build:warnings --copt=-Wimplicit-fallthrough
-build:warnings --copt=-Wno-pointer-bool-conversion
-build:warnings --copt=-Wno-unused-template
 build:warnings --copt=-Wno-builtin-macro-redefined
 build:warnings --copt=-Wno-macro-redefined
 build:warnings --copt=-Wno-deprecated-register

From 58519efe3c6d0d69e1543b6f06b3a2196e3d8ea9 Mon Sep 17 00:00:00 2001
From: Bill Varcho <varcho@google.com>
Date: Mon, 12 Oct 2026 14:09:57 -0700
Subject: [PATCH 0623/1000] [SDY][XLA][2/3] Enable useSingleDeviceMeshNames in
 XLA Shardy round-trip pipelines.

PiperOrigin-RevId: 998146235
---
 .../spmd/shardy/sdy_round_trip/pipelines.cc        |  8 ++++++--
 .../spmd/shardy/test/round_trip_pipeline.mlir      | 10 +++++-----
 .../test/sdy_round_trip_export_pipeline.mlir       |  6 +++---
 ...round_trip_import_pipeline_hlo_sharding_v3.mlir | 14 +++++++-------
 ...trip_import_pipeline_lift_and_dedup_meshes.mlir |  8 ++++----
 5 files changed, 25 insertions(+), 21 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/pipelines.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/pipelines.cc
index 1bdcafb886e..7db3eefc7bd 100644
--- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/pipelines.cc
+++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/pipelines.cc
@@ -56,7 +56,9 @@ void addSdyRoundTripExportPipeline(mlir::OpPassManager& pm,
   // Lift meshes before deduping, since the dedup meshes pass ignores inlined
   // meshes.
   if (!keepMeshesInlined) {
-    pm.addPass(mlir::sdy::createLiftInlinedMeshesPass());
+    pm.addPass(mlir::sdy::createLiftInlinedMeshesPass(
+        mlir::sdy::LiftInlinedMeshesPassOptions{
+            /*useSingleDeviceMeshNames=*/true}));
   }
   pm.addPass(createSdyRoundTripDedupMeshesPass());
   pm.addPass(createSdyRoundTripExportOpsPass());
@@ -110,7 +112,9 @@ void addSdyRoundTripImportPipeline(mlir::OpPassManager& pm,
     //
     // liftAndDedupMeshes is required here because of sdy shardings added
     // directly to hlo in tf2xla.
-    pm.addPass(mlir::sdy::createLiftInlinedMeshesPass());
+    pm.addPass(mlir::sdy::createLiftInlinedMeshesPass(
+        mlir::sdy::LiftInlinedMeshesPassOptions{
+            /*useSingleDeviceMeshNames=*/true}));
     pm.addPass(createSdyRoundTripDedupMeshesPass());
   }
   pm.addPass(createCanonicalizerPass(
diff --git a/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline.mlir b/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline.mlir
index ac9ec25c1ba..e6eed7176fe 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/round_trip_pipeline.mlir
@@ -87,18 +87,18 @@ func.func @main(
 // CHECK-NOT: xla.sdy.sharding
 
 // CHECK-V3: sdy.mesh @empty_mesh = <[]>
-// CHECK-V3: sdy.mesh @maximal_mesh_0 = <[], device_ids=[0]>
+// CHECK-V3: sdy.mesh @single_device_0 = <[], device_ids=[0]>
 sdy.mesh @empty_mesh = <[]>
-sdy.mesh @maximal_mesh_0 = <[], device_ids=[0]>
+sdy.mesh @single_device_0 = <[], device_ids=[0]>
 
 // CHECK-LABEL: func @main
 func.func @main(
   // CHECK-V2: %arg0: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@empty_mesh, [{?}, {}]>},
   // CHECK-V3: %arg0: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@empty_mesh, [{?}, {}]>},
-  // CHECK-V2: %arg1: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@maximal_mesh_0, []>})
-  // CHECK-V3: %arg1: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@maximal_mesh_0, []>})
+  // CHECK-V2: %arg1: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@single_device_0, []>})
+  // CHECK-V3: %arg1: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@single_device_0, []>})
   %arg0: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@empty_mesh, [{?}, {}]>},
-  %arg1: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@maximal_mesh_0, []>}
+  %arg1: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@single_device_0, []>}
   ) -> (tensor<8x16xf32>) {
   %0 = stablehlo.add %arg0, %arg1 : tensor<8x16xf32>
   return %0 : tensor<8x16xf32>
diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_export_pipeline.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_export_pipeline.mlir
index 8be9b934467..7dd40a61cfe 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_export_pipeline.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_export_pipeline.mlir
@@ -24,9 +24,9 @@ sdy.mesh @unused_mesh = <["axis_0"=16]>
 
 // CHECK-V2: module attributes {mhlo.frontend_attributes = {
 // CHECK-V2-SAME: xla.sdy.meshes = "{
-// CHECK-V2-SAME: maximal_mesh_5 = #sdy.mesh<[], device_ids=[5]>,
 // CHECK-V2-SAME: mesh = #sdy.mesh<[\22axis_0\22=2, \22axis_1\22=4, \22axis_2\22=4]>,
-// CHECK-V2-SAME: mesh_0 = #sdy.mesh<[\22a\22=2, \22b\22=2]>
+// CHECK-V2-SAME: mesh_0 = #sdy.mesh<[\22a\22=2, \22b\22=2]>,
+// CHECK-V2-SAME: single_device_5 = #sdy.mesh<[], device_ids=[5]>,
 // CHECK-V2-SAME: unused_mesh = #sdy.mesh<[\22axis_0\22=16]>}"}} {
 
 // CHECK-LABEL: func @multiple_shardings(
@@ -184,7 +184,7 @@ func.func @inlined_mesh(
   // CHECK-NEXT: %[[SHARDING:.*]] = stablehlo.custom_call @Sharding(%arg0)
   // CHECK-V2-SAME:   mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@mesh_0, [{\22a\22, \22b\22}]>]>"}, mhlo.sharding = "{devices=[4]<=[4]}"}
   // CHECK-V2-NEXT: %[[RESULT_SHARDING:.*]] = stablehlo.custom_call @xla.sdy.FuncResultSharding(%[[SHARDING]])
-  // CHECK-V2-SAME:   mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@maximal_mesh_5, []>]>"}
+  // CHECK-V2-SAME:   mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@single_device_5, []>]>"}
   // CHECK-V2-NEXT: return %[[RESULT_SHARDING]]
   // CHECK-V3-NOT:    mhlo.frontend_attributes
   // CHECK-V3-SAME:   mhlo.sharding = "{mesh['a'=2,'b'=2], [{'a', 'b'}]}"
diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir
index 1b0f94f6667..bb6f2d9f8a5 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir
@@ -18,8 +18,8 @@
 module @module_1 {
   // CHECK: sdy.mesh @mesh = <["a"=8, "b"=8, "c"=8]>
   // CHECK: sdy.mesh @mesh_0 = <["a"=2, "b"=2]>
-  // CHECK: sdy.mesh @maximal_mesh_5 = <[], device_ids=[5]>
-  // CHECK: sdy.mesh @maximal_mesh_0 = <[], device_ids=[0]>
+  // CHECK: sdy.mesh @single_device_5 = <[], device_ids=[5]>
+  // CHECK: sdy.mesh @single_device_0 = <[], device_ids=[0]>
 
   // CHECK-LABEL: func @results_with_sharding
   // CHECK-SAME:    %arg0: tensor<32xi32> {sdy.sharding = #sdy.sharding<@mesh, [{"b"}]>},
@@ -156,7 +156,7 @@ module @module_1 {
 
   // CHECK-LABEL: func @inlined_mesh(
   // CHECK-SAME: %arg0: tensor<32xi32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a"}]>},
-  // CHECK-SAME: %arg1: tensor<32xi32> {sdy.sharding = #sdy.sharding<@maximal_mesh_5, []>})
+  // CHECK-SAME: %arg1: tensor<32xi32> {sdy.sharding = #sdy.sharding<@single_device_5, []>})
   // CHECK-SAME: -> tensor<32xi32> {
   func.func @inlined_mesh(
     %arg0: tensor<32xi32> {mhlo.sharding = "{mesh['a'=2,'b'=2], [{'a'}]}"},
@@ -236,8 +236,8 @@ module @module_1 {
     %arg0: tensor<16x8xf32> {mhlo.sharding = "{mesh['a'=1,'b'=4,'c'=1], [{'b'},{?}]}"},
     %arg1: tensor<16x8xf32> {mhlo.frontend_attributes = {baz = 1 : i32, foo = "bar"}},
     %arg2: !stablehlo.token) -> tensor<16x8xf32> {
-    // CHECK-NEXT: %[[SEND:.*]] = "stablehlo.send"(%arg0, %arg2) <{channel_handle = #stablehlo.channel_handle<handle = 1, type = 2>, is_host_transfer = true}> {mhlo.frontend_attributes = {baz = 1 : i32}, sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh_0, []>]>} : (tensor<16x8xf32>, !stablehlo.token) -> !stablehlo.token
-    // CHECK-NEXT: %[[RECV:.*]]:2 = "stablehlo.recv"(%[[SEND]]) <{channel_handle = #stablehlo.channel_handle<handle = 1, type = 3>, is_host_transfer = true}> {mhlo.frontend_attributes = {baz = 1 : i32}, sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh_0, []>, <@maximal_mesh_0, []>]>} : (!stablehlo.token) -> (tensor<16x8xf32>, !stablehlo.token)
+    // CHECK-NEXT: %[[SEND:.*]] = "stablehlo.send"(%arg0, %arg2) <{channel_handle = #stablehlo.channel_handle<handle = 1, type = 2>, is_host_transfer = true}> {mhlo.frontend_attributes = {baz = 1 : i32}, sdy.sharding = #sdy.sharding_per_value<[<@single_device_0, []>]>} : (tensor<16x8xf32>, !stablehlo.token) -> !stablehlo.token
+    // CHECK-NEXT: %[[RECV:.*]]:2 = "stablehlo.recv"(%[[SEND]]) <{channel_handle = #stablehlo.channel_handle<handle = 1, type = 3>, is_host_transfer = true}> {mhlo.frontend_attributes = {baz = 1 : i32}, sdy.sharding = #sdy.sharding_per_value<[<@single_device_0, []>, <@single_device_0, []>]>} : (!stablehlo.token) -> (tensor<16x8xf32>, !stablehlo.token)
     // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[RECV]]#0, %arg1 : tensor<16x8xf32>
     // CHECK-NEXT: return %[[ADD]] : tensor<16x8xf32>
     %0 = "stablehlo.send"(%arg0, %arg2) {
@@ -265,7 +265,7 @@ module @maximal_sharding_module {
     // CHECK-NEXT: stablehlo.custom_call @xla_ffi_python_cpu_callback(%arg0) {
     // CHECK-SAME:   api_version = 4 : i32, backend_config = {descriptor = 126001424235520 : ui64},
     // CHECK-SAME:   has_side_effect = true, operand_layouts = [dense<0> : tensor<1xindex>], result_layouts = [],
-    // CHECK-SAME:   sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh_0, []>]>, xla_shape = "()"
+    // CHECK-SAME:   sdy.sharding = #sdy.sharding_per_value<[<@single_device_0, []>]>, xla_shape = "()"
     // CHECK-SAME: } : (tensor<2xi64>) -> ()
     // CHECK-NEXT: return %arg0 : tensor<2xi64>
     %2 = stablehlo.custom_call @xla_ffi_python_cpu_callback(%arg0) {
@@ -282,7 +282,7 @@ module @maximal_sharding_module {
     // CHECK-NEXT: stablehlo.custom_call @xla_ffi_python_cpu_callback(%arg0) {
     // CHECK-SAME:   api_version = 4 : i32, backend_config = {descriptor = 126001424235520 : ui64},
     // CHECK-SAME:   has_side_effect = true, operand_layouts = [dense<0> : tensor<1xindex>], result_layouts = [],
-    // CHECK-SAME:   sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh_0, []>]>, xla_shape = "()"
+    // CHECK-SAME:   sdy.sharding = #sdy.sharding_per_value<[<@single_device_0, []>]>, xla_shape = "()"
     // CHECK-SAME: } : (tensor<2xi64>) -> ()
     // CHECK-NEXT: return %arg0 : tensor<2xi64>
     %2 = stablehlo.custom_call @xla_ffi_python_cpu_callback(%arg0) {
diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_lift_and_dedup_meshes.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_lift_and_dedup_meshes.mlir
index 544c89db20e..27d98c6b6e6 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_lift_and_dedup_meshes.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_lift_and_dedup_meshes.mlir
@@ -17,12 +17,12 @@
 // CHECK-LABEL: module @multiple_inlined_mesh_shardings
 module @multiple_inlined_mesh_shardings {
   // CHECK: sdy.mesh @mesh = <["a"=2, "b"=2]>
-  // CHECK: sdy.mesh @maximal_mesh_0 = <[], device_ids=[0]>
-  // CHECK: sdy.mesh @maximal_mesh_2 = <[], device_ids=[2]>
+  // CHECK: sdy.mesh @single_device_0 = <[], device_ids=[0]>
+  // CHECK: sdy.mesh @single_device_2 = <[], device_ids=[2]>
 
   // CHECK-LABEL: func @inlined_mesh1(
   // CHECK-SAME: %arg0: tensor<32xi32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}]>})
-  // CHECK-SAME: -> (tensor<32xi32> {sdy.sharding = #sdy.sharding<@maximal_mesh_0, []>}) {
+  // CHECK-SAME: -> (tensor<32xi32> {sdy.sharding = #sdy.sharding<@single_device_0, []>}) {
   func.func @inlined_mesh1(
     %arg0: tensor<32xi32> {mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding<mesh<[\"a\"=2, \"b\"=2]>, [{\"a\"}]>"}}
   ) -> tensor<32xi32> {
@@ -35,7 +35,7 @@ module @multiple_inlined_mesh_shardings {
 
   // CHECK-LABEL: func @inlined_mesh2(
   // CHECK-SAME: %arg0: tensor<32xi32> {sdy.sharding = #sdy.sharding<@mesh, [{"a", "b"}]>})
-  // CHECK-SAME: -> (tensor<32xi32> {sdy.sharding = #sdy.sharding<@maximal_mesh_2, []>}) {
+  // CHECK-SAME: -> (tensor<32xi32> {sdy.sharding = #sdy.sharding<@single_device_2, []>}) {
   func.func @inlined_mesh2(
     %arg0: tensor<32xi32> {mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding<mesh<[\"a\"=4, \"b\"=1]>, [{\"a\"}]>"}}
   ) -> tensor<32xi32> {

From 97a39921e0ddf5c254450cb4ad4ced03c1426f00 Mon Sep 17 00:00:00 2001
From: Dmitri Latushko <latushko@google.com>
Date: Mon, 12 Oct 2026 14:25:56 -0700
Subject: [PATCH 0624/1000] Increase shard_count from 5 to 20 for
 forwardprop_test to avoid TSAN shard timeouts.

PiperOrigin-RevId: 998156158
---
 tensorflow/python/eager/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/eager/BUILD b/tensorflow/python/eager/BUILD
index c18803ac23f..3d1d87b4bcb 100644
--- a/tensorflow/python/eager/BUILD
+++ b/tensorflow/python/eager/BUILD
@@ -531,7 +531,7 @@ cuda_py_strict_test(
 cuda_py_strict_test(
     name = "forwardprop_test",
     srcs = ["forwardprop_test.py"],
-    shard_count = 5,
+    shard_count = 20,
     deps = [
         ":backprop",
         ":context",

From ed9697feb41a3b93cc4b9c3df980f8683d5b28b2 Mon Sep 17 00:00:00 2001
From: Changming Sun CA <changmingsun@google.com>
Date: Mon, 12 Oct 2026 15:07:05 -0700
Subject: [PATCH 0625/1000] Remove unused transitional uint8 depthwise
 convolution kernels.

PiperOrigin-RevId: 998180090
---
 tensorflow/lite/kernels/internal/BUILD        |    2 -
 .../internal/depthwiseconv_quantized_test.cc  |  171 +-
 .../depthwiseconv_uint8_transitional.h        | 8129 -----------------
 3 files changed, 4 insertions(+), 8298 deletions(-)
 delete mode 100644 tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_transitional.h

diff --git a/tensorflow/lite/kernels/internal/BUILD b/tensorflow/lite/kernels/internal/BUILD
index 253bd57a5c5..fe7bbf1bbf1 100644
--- a/tensorflow/lite/kernels/internal/BUILD
+++ b/tensorflow/lite/kernels/internal/BUILD
@@ -1099,12 +1099,10 @@ cc_test(
     name = "depthwiseconv_quantized_test",
     srcs = [
         "depthwiseconv_quantized_test.cc",
-        "optimized/depthwiseconv_uint8_transitional.h",
     ],
     copts = tflite_copts(),
     shard_count = 2,
     deps = [
-        ":compatibility",
         ":cpu_check",
         ":optimized_base",
         ":reference_base",
diff --git a/tensorflow/lite/kernels/internal/depthwiseconv_quantized_test.cc b/tensorflow/lite/kernels/internal/depthwiseconv_quantized_test.cc
index d5a2da2b9d5..1cc984c5a18 100644
--- a/tensorflow/lite/kernels/internal/depthwiseconv_quantized_test.cc
+++ b/tensorflow/lite/kernels/internal/depthwiseconv_quantized_test.cc
@@ -38,7 +38,6 @@ limitations under the License.
 #include "tensorflow/lite/kernels/internal/optimized/depthwiseconv_multithread.h"
 #include "tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8.h"
 #include "tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h"
-#include "tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_transitional.h"
 #include "tensorflow/lite/kernels/internal/reference/depthwiseconv_uint8.h"
 
 namespace tflite {
@@ -50,12 +49,6 @@ using optimized_ops::depthwise_conv::QuantizationTypeImpl;
 using ::testing::Bool;
 using ::testing::Values;
 
-#if defined(__aarch64__)
-static constexpr bool kLooseIntrinsicsTolerance = false;
-#else
-static constexpr bool kLooseIntrinsicsTolerance = true;
-#endif
-
 // Currently, this is used in place of a Boolean "is symmetric?".
 enum class ParamsSpecialization {
   kNone = 0,
@@ -239,83 +232,10 @@ inline void DispatchDepthwiseConvImpl(
 #endif
       break;
     }
-    case DepthwiseConvImplementation::kUseCModel3x3DotProduct: {
-      DotProduct3x3KernelType kernel_type =
-          optimized_ops::depthwise_conv::CategorizeDotProductKernel(
-              input_shape, filter_shape, output_shape, params);
-
-      ASSERT_TRUE(
-          kernel_type == DotProduct3x3KernelType::kPlain ||
-          kernel_type == DotProduct3x3KernelType::kStride2 ||
-          kernel_type ==
-              DotProduct3x3KernelType::kWithDepthMultiplicationStride1 ||
-          kernel_type ==
-              DotProduct3x3KernelType::kWithDepthMultiplicationStride2)
-          << "Kernel type = " << static_cast<int>(kernel_type)
-          << " depth_multiplier = " << params.depth_multiplier
-          << " pad_width = " << params.padding_values.width
-          << " pad_height = " << params.padding_values.height
-          << " stride_width = " << params.stride_width
-          << " stride_height = " << params.stride_height
-          << " input_width = " << input_shape.Dims(2)
-          << " input_height = " << input_shape.Dims(1)
-          << " output_width = " << output_shape.Dims(2)
-          << " output_height = " << output_shape.Dims(1)
-          << " depth = " << input_shape.Dims(3)
-          << " buffer need = " << input_shape.Dims(3) * input_shape.Dims(2) * 6
-          << " input_offset = " << params.input_offset;
-
-      optimized_ops::depthwise_conv::DepthwiseConvDotProduct3x3Impl<
-          DepthwiseConvImplementation::kUseCModel3x3DotProduct,
-          quantization_type>(
-          params, input_shape, input_data, filter_shape, filter_data,
-          bias_shape, bias_data, output_shape, output_data, /*thread_start=*/0,
-          /*thread_end=*/output_shape.Dims(1), /*thread_dim=*/1);
-      return;
-    }
-    case DepthwiseConvImplementation::kUseUnwound3x3DotProduct: {
-      DotProduct3x3KernelType kernel_type =
-          optimized_ops::depthwise_conv::CategorizeDotProductKernel(
-              input_shape, filter_shape, output_shape, params);
-      ASSERT_TRUE(
-          kernel_type == DotProduct3x3KernelType::kPlain ||
-          kernel_type == DotProduct3x3KernelType::kStride2 ||
-          kernel_type ==
-              DotProduct3x3KernelType::kWithDepthMultiplicationStride1 ||
-          kernel_type ==
-              DotProduct3x3KernelType::kWithDepthMultiplicationStride2);
-      optimized_ops::depthwise_conv::DepthwiseConvDotProduct3x3Impl<
-          DepthwiseConvImplementation::kUseUnwound3x3DotProduct,
-          quantization_type>(
-          params, input_shape, input_data, filter_shape, filter_data,
-          bias_shape, bias_data, output_shape, output_data, /*thread_start=*/0,
-          /*thread_end=*/output_shape.Dims(1), /*thread_dim=*/1);
-      return;
-    }
-    case DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct: {
-#if defined(USE_NEON)
-      DotProduct3x3KernelType kernel_type =
-          optimized_ops::depthwise_conv::CategorizeDotProductKernel(
-              input_shape, filter_shape, output_shape, params);
-
-      ASSERT_TRUE(
-          kernel_type == DotProduct3x3KernelType::kPlain ||
-          kernel_type == DotProduct3x3KernelType::kStride2 ||
-          kernel_type ==
-              DotProduct3x3KernelType::kWithDepthMultiplicationStride1 ||
-          kernel_type ==
-              DotProduct3x3KernelType::kWithDepthMultiplicationStride2);
-      optimized_ops::depthwise_conv::DepthwiseConvDotProduct3x3Impl<
-          DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct,
-          quantization_type>(
-          params, input_shape, input_data, filter_shape, filter_data,
-          bias_shape, bias_data, output_shape, output_data, /*thread_start=*/0,
-          /*thread_end=*/output_shape.Dims(1), /*thread_dim=*/1);
-      return;
-#else
+    case DepthwiseConvImplementation::kUseCModel3x3DotProduct:
+    case DepthwiseConvImplementation::kUseUnwound3x3DotProduct:
+    case DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct:
       break;
-#endif
-    }
     case DepthwiseConvImplementation::kUseGenericKernel: {
       DispatchDepthwiseConvGeneral<quantization_type>(
           params, input_shape, input_data, filter_shape, filter_data,
@@ -448,34 +368,8 @@ inline void DispatchDepthwiseConvImpl<QuantizationType::kPerChannelInt8>(
     }
     case DepthwiseConvImplementation::kUseCModel3x3DotProduct:
     case DepthwiseConvImplementation::kUseUnwound3x3DotProduct:
+    case DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct:
       break;
-    case DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct: {
-#if defined(USE_NEON)
-      DotProduct3x3KernelType kernel_type =
-          optimized_ops::depthwise_conv::CategorizeDotProductKernel<
-              QuantizationType::kPerChannelInt8>(
-              input_shape, filter_shape, output_shape, params,
-              params.output_shift_per_channel);
-
-      ASSERT_TRUE(
-          kernel_type == DotProduct3x3KernelType::kPlain ||
-          kernel_type == DotProduct3x3KernelType::kStride2 ||
-          kernel_type ==
-              DotProduct3x3KernelType::kWithDepthMultiplicationStride1 ||
-          kernel_type ==
-              DotProduct3x3KernelType::kWithDepthMultiplicationStride2);
-      optimized_ops::depthwise_conv::DepthwiseConvDotProduct3x3Impl<
-          DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct,
-          quantization_type>(
-          params, input_shape, input_data, filter_shape, filter_data,
-          bias_shape, bias_data, output_shape, output_data,
-          /*thread_start=*/0,
-          /*thread_end=*/output_shape.Dims(1), /*thread_dim=*/1);
-      return;
-#else
-      break;
-#endif
-    }
     case DepthwiseConvImplementation::kUseGenericKernel: {
       EXPECT_NE(params.output_multiplier_per_channel, nullptr);
       EXPECT_NE(params.output_shift_per_channel, nullptr);
@@ -1226,63 +1120,6 @@ INSTANTIATE_TEST_SUITE_P(
         ),
     TestParam::TestNameSuffix);
 
-INSTANTIATE_TEST_SUITE_P(
-    CModel, DepthwiseConvTest,
-    testing::Combine(
-        Values(DepthwiseConvImplementation::
-                   kUseCModel3x3DotProduct),            // forced_invocation
-        Values(200),                                    // tests_to_run
-        Values(QuantizationType::kNonPerChannelUint8),  // quantization_type
-        Bool(),                                         // test_stride
-        Bool(),                                         // test_pad
-        Bool(),                                         // test_depth_multiplier
-        Values(DepthwiseConvOutputRounding::kUpward),   // output_rounding
-        Values(1),                                      // num_threads
-        Values(false)                                   // loose_tolerance
-        ),
-    TestParam::TestNameSuffix);
-
-INSTANTIATE_TEST_SUITE_P(
-    Unwound, DepthwiseConvTest,
-    testing::Combine(
-        Values(DepthwiseConvImplementation::
-                   kUseUnwound3x3DotProduct),           // forced_invocation
-        Values(200),                                    // tests_to_run
-        Values(QuantizationType::kNonPerChannelUint8),  // quantization_type
-        Bool(),                                         // test_stride
-        Bool(),                                         // test_pad
-        Bool(),                                         // test_depth_multiplier
-        Values(DepthwiseConvOutputRounding::kUpward),   // output_rounding
-        Values(1),                                      // num_threads
-        Values(false)                                   // loose_tolerance
-        ),
-    TestParam::TestNameSuffix);
-
-#if defined(USE_NEON)
-
-// TODO(b/148145875): Remove this extra guard after checking that code runs
-// without lax vector conversions.
-#if defined(__aarch64__) && !defined(GOOGLE_L4T) && defined(__ANDROID__) && \
-    defined(__clang__)
-INSTANTIATE_TEST_SUITE_P(
-    IntrinsicsPerChannel, DepthwiseConvTest,
-    testing::Combine(
-        Values(DepthwiseConvImplementation::
-                   kUseIntrinsics3x3DotProduct),       // forced_invocation
-        Values(200),                                   // tests_to_run
-        Values(QuantizationType::kPerChannelInt8),     // quantization_type
-        Bool(),                                        // test_stride
-        Bool(),                                        // test_pad
-        Bool(),                                        // test_depth_multiplier
-        Values(DepthwiseConvOutputRounding::kUpward),  // output_rounding
-        Values(1),                                     // num_threads
-        Values(kLooseIntrinsicsTolerance)              // loose_tolerance
-        ),
-    TestParam::TestNameSuffix);
-#endif
-
-#endif
-
 #if defined(__aarch64__) && !defined(GOOGLE_L4T) && defined(__ANDROID__) && \
     defined(__clang__)
 INSTANTIATE_TEST_SUITE_P(
diff --git a/tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_transitional.h b/tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_transitional.h
deleted file mode 100644
index 63dca21cd87..00000000000
--- a/tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_transitional.h
+++ /dev/null
@@ -1,8129 +0,0 @@
-/* Copyright 2019 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.
-==============================================================================*/
-#ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_UINT8_TRANSITIONAL_H_
-#define TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_UINT8_TRANSITIONAL_H_
-
-// This file provides kernel implementations that are not used in shipped
-// inference code, but rather (a) show how model C++ code is designed and then
-// transformed into asm code, and (b) aid with maintenance and later development
-// of variations. Many projects (even including, say, the classic NAG libraries)
-// develop highly optimized code, but do not maintain intermediate versions.
-// Often the result is incomprehensible final-version code.
-
-#include <algorithm>
-
-#include "tensorflow/lite/kernels/internal/compatibility.h"
-#include "tensorflow/lite/kernels/internal/optimized/cpu_check.h"
-#include "tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8.h"
-#include "tensorflow/lite/kernels/internal/optimized/depthwiseconv_uint8_3x3_filter.h"
-#include "tensorflow/lite/kernels/internal/reference/depthwiseconv_uint8.h"
-#include "tensorflow/lite/kernels/internal/types.h"
-
-namespace tflite {
-namespace optimized_ops {
-namespace depthwise_conv {
-
-#ifdef USE_NEON
-
-inline void util_vst1_u8(uint8_t* data_addr, uint8x8_t reg) {
-  return vst1_u8(data_addr, reg);
-}
-inline void util_vst1_x8(uint8_t* data_addr, int8x8_t reg) {
-  return vst1_u8(data_addr, vreinterpret_u8_s8(reg));
-}
-inline void util_vst1_x8(int8_t* data_addr, int8x8_t reg) {
-  return vst1_s8(data_addr, reg);
-}
-
-// Lane operations are for clarity and convenience. We want to load and store
-// 4 8-bit lanes together. So these are treated much like 32-bit loads and
-// 32-bit stores. Stores require 32-bit alignment.
-
-#define vst1_lane_s8x4(dst, reg, lane_num)                                  \
-  TFLITE_DCHECK_EQ(reinterpret_cast<std::uintptr_t>(dst) % 4, 0);           \
-  vst1_lane_u32(reinterpret_cast<uint32_t*>(dst), vreinterpret_u32_s8(reg), \
-                lane_num)
-#define vst1_lane_u8x4(dst, reg, lane_num)                                  \
-  TFLITE_DCHECK_EQ(reinterpret_cast<std::uintptr_t>(dst) % 4, 0);           \
-  vst1_lane_u32(reinterpret_cast<uint32_t*>(dst), vreinterpret_u32_u8(reg), \
-                lane_num)
-#define vst1q_lane_s8x4(dst, reg, lane_num)                                   \
-  TFLITE_DCHECK_EQ(reinterpret_cast<std::uintptr_t>(dst) % 4, 0);             \
-  vst1q_lane_u32(reinterpret_cast<uint32_t*>(dst), vreinterpretq_u32_s8(reg), \
-                 lane_num)
-#define vst1q_lane_u8x4(dst, reg, lane_num)                                   \
-  TFLITE_DCHECK_EQ(reinterpret_cast<std::uintptr_t>(dst) % 4, 0);             \
-  vst1q_lane_u32(reinterpret_cast<uint32_t*>(dst), vreinterpretq_u32_u8(reg), \
-                 lane_num)
-
-// Important! Most compilation configurations will compile and run without
-// reinterpret_cast. Sanitizers may fail silently on lane-loading, with an
-// obscure bug or mis-feature probably in unhygienic macro expansion.
-#define vld1q_lane_s8x8(src, reg, lane_num)                                   \
-  vreinterpretq_s8_u64(vld1q_lane_u64(reinterpret_cast<const uint64_t*>(src), \
-                                      vreinterpretq_u64_s8(reg), lane_num))
-#define vld1_lane_8x4(src, reg, lane_num)                                \
-  vreinterpret_s8_s32(vld1_lane_s32(reinterpret_cast<const int32*>(src), \
-                                    vreinterpret_s32_s8(reg), lane_num))
-#define vld1q_lane_8x4(src, reg, lane_num)                                 \
-  vreinterpretq_s8_s32(vld1q_lane_s32(reinterpret_cast<const int32*>(src), \
-                                      vreinterpretq_s32_s8(reg), lane_num))
-#define vld1q_dup_s8x4(src) \
-  vreinterpretq_s8_s32(vld1q_dup_s32(reinterpret_cast<const int32*>(src)))
-
-#endif  // USE_NEON
-
-template <QuantizationType quantization_type>
-struct ProcessPerDepth<DepthwiseConvImplementation::kUseCModel3x3DotProduct,
-                       quantization_type> {
-  // Filter data is provided as filter_block[3][3][depth/8][2][4]: height 3,
-  // width 3,  sub-block 0 or 1, depth 4. Filter data is written as
-  // filter_bank[3][2][4][4]; height 3, sub-block, depth 4, width 4.
-  //
-  // Note that this rearrangement is much like that performed on input data when
-  // filling the workspace, and optimized versions will be similar.
-  static inline void FillFilterBank(int depth, const uint8_t* filter_block,
-                                    int8_t filter_bank[3][2][4][4]) {
-    constexpr int kSymmetricZeroPoint =
-        QuantizationTypeImpl<quantization_type>::kIntSymmetricZeroPoint;
-    // Load filter data in, 8-bytes down depth / sub-block at a time.
-    //
-    // loaded_filter has dimensions height 3, width 4, sub-block 0 or 1,
-    // depth 4.
-    uint8_t loaded_filter[3][4][2][4];
-    for (int y = 0; y < 3; ++y) {
-      for (int x = 0; x < 3; ++x) {
-        memcpy(loaded_filter[y][x][0], &filter_block[3 * y * depth + x * depth],
-               8);
-      }
-      // Pad the filter with symmetric representation of 0, so that the values
-      // become 0 when the zero-poing is added below. Thus these filter taps are
-      // effectively disregarded in later filtering.
-      memset(loaded_filter[y][3][0], kSymmetricZeroPoint, 8);
-    }
-    for (int y = 0; y < 3; ++y) {
-      for (int z = 0; z < 4; ++z) {
-        for (int x = 0; x < 4; ++x) {
-          filter_bank[y][0][z][x] =
-              loaded_filter[y][x][0][z] - kSymmetricZeroPoint;
-          filter_bank[y][1][z][x] =
-              loaded_filter[y][x][1][z] - kSymmetricZeroPoint;
-        }
-      }
-    }
-  }
-
-  // Adjust the bias (weights) data according to the input offset.
-  //
-  // The output calculation is
-  // out[h][w][d] = bias[d] + sum_ij (in[h+i][w+j][d] + in_offset) *
-  //                                 (filter[i][j][d] + filter_offset)
-  // (where offsets are expressed as differences from 128).
-  //
-  // Since we cannot efficiently handle varying offsets / bias across the image,
-  // we insist on filter_offset = 0.
-  //
-  // This function calculates
-  // adjusted_bias[d] = bias[d] + sum_ij in_offset * filter[i][j][d]
-  // which accounts for input offset. If the bias is constant over the depth,
-  // the adjusted bias will vary.
-  static inline void AdjustBias(int32_t input_offset,
-                                const int8_t filter_bank[3][2][4][4],
-                                const int32_t* bias_data,
-                                int32_t adjusted_bias_block[2][4]) {
-    constexpr int kSymmetricZeroPoint =
-        QuantizationTypeImpl<quantization_type>::kIntSymmetricZeroPoint;
-    TFLITE_DCHECK_GE(input_offset, -255);
-    TFLITE_DCHECK_LE(input_offset, 0);
-    // For instance, if input_offset == 128, no adjustment is needed.
-    const int32_t input_offset_difference = input_offset + kSymmetricZeroPoint;
-
-    for (int s = 0; s < 2; ++s) {
-      for (int z = 0; z < 4; ++z) {
-        adjusted_bias_block[s][z] = bias_data[4 * s + z];
-        for (int i = 0; i < 9; ++i) {
-          adjusted_bias_block[s][z] +=
-              input_offset_difference * filter_bank[i % 3][s][z][i / 3];
-        }
-      }
-    }
-  }
-
-  static void Run(const uint8_t* filter_data, const int32_t* bias_data,
-                  int8_t* shuffled_filter_data, int32_t* adjusted_bias_data,
-                  const DepthwiseConvDotProdParams* function_params) {
-    constexpr int shuffled_filter_increment = 2 * 3 * 4 * 4;
-    const int depth = function_params->output_depth;
-    const int depth_micro_repeats = function_params->depth_micro_repeats;
-    const int bias_increment = function_params->bias_increment;
-    const int32_t input_offset = function_params->input_offset;
-
-    int8_t filter_bank[3][2][4][4];
-    int32_t adjusted_bias_block[2][4];
-
-    for (int j_depth = 0; j_depth < depth_micro_repeats; ++j_depth) {
-      FillFilterBank(depth, filter_data + 8 * j_depth, filter_bank);
-      AdjustBias(input_offset, filter_bank,
-                 bias_data + 2 * bias_increment * j_depth, adjusted_bias_block);
-
-      memcpy(shuffled_filter_data, filter_bank[0][0][0],
-             shuffled_filter_increment);
-      shuffled_filter_data += shuffled_filter_increment;
-      memcpy(adjusted_bias_data, adjusted_bias_block[0],
-             8 * sizeof(adjusted_bias_block[0][0]));
-      adjusted_bias_data += 8;
-    }
-  }
-};
-
-template <QuantizationType quantization_type>
-struct ProcessPerDepth<DepthwiseConvImplementation::kUseUnwound3x3DotProduct,
-                       quantization_type> {
-  static inline void Run(const uint8_t* filter_data, const int32_t* bias_data,
-                         int8_t* shuffled_filter_data,
-                         int32_t* adjusted_bias_data,
-                         const DepthwiseConvDotProdParams* function_params) {
-    const int depth = function_params->output_depth;
-    const int depth_micro_repeats = function_params->depth_micro_repeats;
-    const int bias_increment = function_params->bias_increment;
-
-    // Simulate NEON-register transposition of subset of filter.
-    int8_t filter_bank_a_0[4][4];  // Depth 4, width 4.
-    int8_t filter_bank_a_1[4][4];
-    int8_t filter_bank_a_2[4][4];
-    int8_t filter_bank_b_0[4][4];
-    int8_t filter_bank_b_1[4][4];
-    int8_t filter_bank_b_2[4][4];
-
-    // Load filter data in, essentially dropping the [depth/8] dimension, which
-    // is equivalent to loading just the depth needed for one micro-block.
-    //
-    // loaded_filter has dimensions height 3, width 4, sub-block 0 or 1,
-    // depth 4.
-    uint8_t loaded_filter_0[4][2][4];
-    uint8_t loaded_filter_1[4][2][4];
-    uint8_t loaded_filter_2[4][2][4];
-
-    constexpr int kSymmetricZeroPoint =
-        QuantizationTypeImpl<quantization_type>::kIntSymmetricZeroPoint;
-    const int32_t input_offset = function_params->input_offset;
-    TFLITE_DCHECK_GE(input_offset, -255);
-    TFLITE_DCHECK_LE(input_offset, 0);
-    const int32_t input_offset_difference = input_offset + kSymmetricZeroPoint;
-
-    for (int j_depth = 0; j_depth < depth_micro_repeats; ++j_depth) {
-      const uint8_t* filter_block = filter_data + 8 * j_depth;
-
-      // Filter data is provided as filter_block[3][3][depth/8][2][4].
-      // height 3, width 3, micro-blocks, sub-block 0 or 1, depth 4.
-      // filter_bank[3][2][4][4]; Sub-block, height 3, depth 4, width 4.
-      for (int x = 0; x < 3; ++x) {
-        memcpy(loaded_filter_0[x][0], &filter_block[3 * 0 * depth + x * depth],
-               8);
-        memcpy(loaded_filter_1[x][0], &filter_block[3 * 1 * depth + x * depth],
-               8);
-        memcpy(loaded_filter_2[x][0], &filter_block[3 * 2 * depth + x * depth],
-               8);
-      }
-      // Pad the filter with -filter_offset, so that the values become 0 when
-      // the filter_offset is later added, and so the filter tap is effectively
-      // disregarded.
-      memset(loaded_filter_0[3][0], kSymmetricZeroPoint, 8);
-      memset(loaded_filter_1[3][0], kSymmetricZeroPoint, 8);
-      memset(loaded_filter_2[3][0], kSymmetricZeroPoint, 8);
-
-      for (int z = 0; z < 4; ++z) {
-        for (int x = 0; x < 4; ++x) {
-          filter_bank_a_0[z][x] =
-              loaded_filter_0[x][0][z] - kSymmetricZeroPoint;
-          filter_bank_b_0[z][x] =
-              loaded_filter_0[x][1][z] - kSymmetricZeroPoint;
-          filter_bank_a_1[z][x] =
-              loaded_filter_1[x][0][z] - kSymmetricZeroPoint;
-          filter_bank_b_1[z][x] =
-              loaded_filter_1[x][1][z] - kSymmetricZeroPoint;
-          filter_bank_a_2[z][x] =
-              loaded_filter_2[x][0][z] - kSymmetricZeroPoint;
-          filter_bank_b_2[z][x] =
-              loaded_filter_2[x][1][z] - kSymmetricZeroPoint;
-        }
-      }
-
-      memcpy(shuffled_filter_data, filter_bank_a_0, 16);
-      shuffled_filter_data += 16;
-      memcpy(shuffled_filter_data, filter_bank_b_0, 16);
-      shuffled_filter_data += 16;
-      memcpy(shuffled_filter_data, filter_bank_a_1, 16);
-      shuffled_filter_data += 16;
-      memcpy(shuffled_filter_data, filter_bank_b_1, 16);
-      shuffled_filter_data += 16;
-      memcpy(shuffled_filter_data, filter_bank_a_2, 16);
-      shuffled_filter_data += 16;
-      memcpy(shuffled_filter_data, filter_bank_b_2, 16);
-      shuffled_filter_data += 16;
-
-      int32_t adjusted_bias_data_0[4];
-      int32_t adjusted_bias_data_1[4];
-      // For instance, if input_offset == 128, no adjustment is needed.
-      for (int z = 0; z < 4; ++z) {
-        adjusted_bias_data_0[z] = bias_data[z];
-        adjusted_bias_data_1[z] = bias_data[4 + z];
-        for (int x = 0; x < 4; ++x) {
-          adjusted_bias_data_0[z] +=
-              input_offset_difference * filter_bank_a_0[z][x];
-          adjusted_bias_data_0[z] +=
-              input_offset_difference * filter_bank_a_1[z][x];
-          adjusted_bias_data_0[z] +=
-              input_offset_difference * filter_bank_a_2[z][x];
-          adjusted_bias_data_1[z] +=
-              input_offset_difference * filter_bank_b_0[z][x];
-          adjusted_bias_data_1[z] +=
-              input_offset_difference * filter_bank_b_1[z][x];
-          adjusted_bias_data_1[z] +=
-              input_offset_difference * filter_bank_b_2[z][x];
-
-          adjusted_bias_data[z] = adjusted_bias_data_0[z];
-          adjusted_bias_data[4 + z] = adjusted_bias_data_1[z];
-        }
-      }
-      bias_data += 2 * bias_increment;
-      adjusted_bias_data += 8;
-    }
-  }
-};
-
-#ifdef USE_NEON
-template <QuantizationType quantization_type>
-struct ProcessPerDepth<DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct,
-                       quantization_type> {
-  static void ProcessPerDepthIntrinsics(
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          filter_data,
-      const int32_t* bias_data, int8_t* shuffled_filter_data,
-      int32_t* adjusted_bias_data,
-      const DepthwiseConvDotProdParams* function_params) {
-    const int depth = function_params->output_depth;
-    const int depth_micro_repeats = function_params->depth_micro_repeats;
-    const int bias_increment = function_params->bias_increment;
-
-    constexpr int kSymmetricZeroPoint =
-        QuantizationTypeImpl<quantization_type>::kIntSymmetricZeroPoint;
-    constexpr uint8_t kSignBit =
-        QuantizationTypeImpl<quantization_type>::kUint8SignBit;
-    const int32_t input_offset = function_params->input_offset;
-    if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-      TFLITE_DCHECK_GE(input_offset, -255);
-      TFLITE_DCHECK_LE(input_offset, 0);
-    }
-    const int32_t input_offset_difference = input_offset + kSymmetricZeroPoint;
-    const int8x16_t ones_vector = vdupq_n_s8(1);
-
-    // Simulate NEON-register transposition of subset of filter.
-    int8x16_t input_0_a;
-    int8x16_t input_0_b;
-    int8x16_t input_0_c;
-    int8x16_t input_1_a;
-    int8x16_t input_1_b;
-    int8x16_t input_1_c;
-    int8x16_t input_2_a;
-    int8x16_t input_2_b;
-    int8x16_t input_2_c;
-
-    int8x16_t filter_0_a;
-    int8x16_t filter_0_b;
-    int8x16_t filter_1_a;
-    int8x16_t filter_1_b;
-    int8x16_t filter_2_a;
-    int8x16_t filter_2_b;
-
-    // For uint8, effect subtraction of zero-point = 128 by XOR of sign bit.
-    const int8x16_t sign_bit = vreinterpretq_s8_u8(vdupq_n_u8(kSignBit));
-
-    const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-        filter_block = filter_data;
-    for (int j_depth = 0; j_depth < depth_micro_repeats; ++j_depth) {
-      // Filter data is provided as filter_block[3][3][depth/8][2][4].
-      // height 3, width 3, micro-blocks, sub-block 0 or 1, depth 4.
-      // filter_bank[3][2][4][4]; Sub-block, height 3, depth 4, width 4.
-
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          filter_block_ptr = filter_block;
-      input_0_a = vld1q_lane_s8x8(filter_block_ptr, input_0_a, 0);
-      filter_block_ptr += depth;
-      input_0_b = vld1q_lane_s8x8(filter_block_ptr, input_0_b, 0);
-      filter_block_ptr += depth;
-      input_0_c = vld1q_lane_s8x8(filter_block_ptr, input_0_c, 0);
-      filter_block_ptr += depth;
-      input_1_a = vld1q_lane_s8x8(filter_block_ptr, input_1_a, 0);
-      filter_block_ptr += depth;
-      input_1_b = vld1q_lane_s8x8(filter_block_ptr, input_1_b, 0);
-      filter_block_ptr += depth;
-      input_1_c = vld1q_lane_s8x8(filter_block_ptr, input_1_c, 0);
-      filter_block_ptr += depth;
-      input_2_a = vld1q_lane_s8x8(filter_block_ptr, input_2_a, 0);
-      filter_block_ptr += depth;
-      input_2_b = vld1q_lane_s8x8(filter_block_ptr, input_2_b, 0);
-      filter_block_ptr += depth;
-      input_2_c = vld1q_lane_s8x8(filter_block_ptr, input_2_c, 0);
-
-      filter_0_a = vzip1q_s8(input_0_a, input_0_b);
-      filter_0_b = vzip1q_s8(input_0_c, sign_bit);
-      filter_1_a = vzip1q_s8(input_1_a, input_1_b);
-      filter_1_b = vzip1q_s8(input_1_c, sign_bit);
-      filter_2_a = vzip1q_s8(input_2_a, input_2_b);
-      filter_2_b = vzip1q_s8(input_2_c, sign_bit);
-      if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-        filter_0_a = veorq_s8(filter_0_a, sign_bit);
-        filter_0_b = veorq_s8(filter_0_b, sign_bit);
-        filter_1_a = veorq_s8(filter_1_a, sign_bit);
-        filter_1_b = veorq_s8(filter_1_b, sign_bit);
-        filter_2_a = veorq_s8(filter_2_a, sign_bit);
-        filter_2_b = veorq_s8(filter_2_b, sign_bit);
-      }
-      vzipq_s8x2_in_place(&filter_0_a, &filter_0_b);
-      vzipq_s8x2_in_place(&filter_1_a, &filter_1_b);
-      vzipq_s8x2_in_place(&filter_2_a, &filter_2_b);
-
-      vst1q_s8(shuffled_filter_data, filter_0_a);
-      shuffled_filter_data += 16;
-      vst1q_s8(shuffled_filter_data, filter_0_b);
-      shuffled_filter_data += 16;
-      vst1q_s8(shuffled_filter_data, filter_1_a);
-      shuffled_filter_data += 16;
-      vst1q_s8(shuffled_filter_data, filter_1_b);
-      shuffled_filter_data += 16;
-      vst1q_s8(shuffled_filter_data, filter_2_a);
-      shuffled_filter_data += 16;
-      vst1q_s8(shuffled_filter_data, filter_2_b);
-      shuffled_filter_data += 16;
-
-      int32x4_t adjusted_bias_data_a = vld1q_s32(bias_data);
-      bias_data += bias_increment;
-      int32x4_t adjusted_bias_data_b = vld1q_s32(bias_data);
-      bias_data += bias_increment;
-      // For instance, if input_offset is kIntSymmetricZeroPoint, no adjustment
-      // is needed.
-
-      int32x4_t filter_sum_a = vdupq_n_s32(0);
-      filter_sum_a = vdotq_s32(filter_sum_a, filter_0_a, ones_vector);
-      filter_sum_a = vdotq_s32(filter_sum_a, filter_1_a, ones_vector);
-      filter_sum_a = vdotq_s32(filter_sum_a, filter_2_a, ones_vector);
-      int32x4_t filter_sum_b = vdupq_n_s32(0);
-      filter_sum_b = vdotq_s32(filter_sum_b, filter_0_b, ones_vector);
-      filter_sum_b = vdotq_s32(filter_sum_b, filter_1_b, ones_vector);
-      filter_sum_b = vdotq_s32(filter_sum_b, filter_2_b, ones_vector);
-
-      adjusted_bias_data_a = vmlaq_n_s32(adjusted_bias_data_a, filter_sum_a,
-                                         input_offset_difference);
-      adjusted_bias_data_b = vmlaq_n_s32(adjusted_bias_data_b, filter_sum_b,
-                                         input_offset_difference);
-
-      vst1q_s32(adjusted_bias_data, adjusted_bias_data_a);
-      adjusted_bias_data += 4;
-      vst1q_s32(adjusted_bias_data, adjusted_bias_data_b);
-      adjusted_bias_data += 4;
-
-      filter_block += 8;
-    }
-  }
-
-  static inline void Run(const typename QuantizationTypeImpl<
-                             quantization_type>::ExternalType* filter_data,
-                         const int32_t* bias_data, int8_t* shuffled_filter_data,
-                         int32_t* adjusted_bias_data,
-                         const DepthwiseConvDotProdParams* function_params) {
-    ProcessPerDepthIntrinsics(filter_data, bias_data, shuffled_filter_data,
-                              adjusted_bias_data, function_params);
-  }
-};
-#endif
-
-template <QuantizationType quantization_type, int32_t max_padding>
-struct PackMacroBlock<
-    DepthwiseConvImplementation::kUseCModel3x3DotProduct, quantization_type,
-    DepthwiseConvDepthMultiplication::kNoMultiplication, max_padding> {
-  // A straight copy of a macro block of input data into a scratch buffer.
-  //
-  // Requirement: depth_micro_repeats > 0.
-  static inline void CopyMacroBlock(
-      int32_t height_block_number, int32_t width_block_number,
-      const DepthwiseConvDotProdParams& function_params,
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          input_block_data,
-      int8_t* scratch_block_data) {
-    TFLITE_DCHECK_LE(max_padding, 1);
-
-    // Strides.
-    // The input depth and count of micro blocks provide the width strides.
-    const int input_height_stride = function_params.input_height_stride;
-    const int workspace_height_stride = function_params.workspace_height_stride;
-    const int input_depth = function_params.input_depth;
-    const int depth_micro_repeats = function_params.depth_micro_repeats;
-    TFLITE_DCHECK_GT(depth_micro_repeats, 0);
-
-    // Remaining iteration and dimension parameters.
-    //
-    // If width_overall_micro_repeats = input_width_micro_repeats + 1, then the
-    // final micro block is incomplete.
-    const int width_overall_micro_repeats =
-        function_params.input_width_overall_micro_repeats;
-    int input_width_micro_repeats = function_params.input_width_micro_repeats;
-    const int residual_width = function_params.residual_width;
-    const int block_height = function_params.inbound_block_height;
-
-    const int padding_left = function_params.padding_left;
-    const int padding_right = function_params.padding_right;
-    const int padding_top = function_params.padding_top;
-    const int padding_bottom = function_params.padding_bottom;
-
-    const bool leading_width_padding =
-        padding_left > 0 && width_block_number == 0;
-    const bool trailing_width_padding =
-        padding_right > 0 &&
-        width_block_number == (function_params.width_macro_count - 1);
-    const bool leading_height_padding =
-        padding_top > 0 && height_block_number < 0;
-    const bool trailing_height_padding =
-        padding_bottom > 0 &&
-        height_block_number == (function_params.height_macro_count - 1);
-
-    // Modify the trailing case to reflect the input width.
-    int input_residual_width =
-        input_width_micro_repeats < width_overall_micro_repeats ? residual_width
-                                                                : 4;
-    if (trailing_width_padding) {
-      input_residual_width -= 1;
-      input_width_micro_repeats = width_overall_micro_repeats - 1;
-    }
-
-    constexpr int kSymmetricZeroPoint =
-        QuantizationTypeImpl<quantization_type>::kIntSymmetricZeroPoint;
-    const int32_t input_offset_difference =
-        function_params.input_offset + kSymmetricZeroPoint;
-
-    // We load data into a temporary buffer and then save, to match subsequent
-    // processing. This will make it easier to combine stages into one ASM
-    // routine.
-    int8_t tmp_load[4][2][4];
-
-    int copy_block_height = block_height;
-    if (leading_height_padding) {
-      memset(scratch_block_data, -input_offset_difference,
-             workspace_height_stride);
-      scratch_block_data += workspace_height_stride;
-      input_block_data += input_height_stride;
-      copy_block_height -= 1;
-    }
-    if (trailing_height_padding) {
-      copy_block_height -= 1;
-    }
-
-    // The outer 3 loops go through all the micro blocks in a macro block.
-    for (int k_height = 0; k_height < copy_block_height; ++k_height) {
-      for (int j_width = 0; j_width < width_overall_micro_repeats; ++j_width) {
-        // Figure out division of work (available input vs trailing padding).
-        int adjusted_residual_width =
-            j_width == input_width_micro_repeats ? input_residual_width : 4;
-
-        int start_width = 0;
-        if (leading_width_padding && j_width == 0) {
-          start_width = 1;
-          memset(tmp_load[0][0], -input_offset_difference, 8);
-        }
-        if (adjusted_residual_width < 4) {
-          for (int x = adjusted_residual_width; x < 4; ++x) {
-            memset(tmp_load[x][0], -input_offset_difference, 8);
-          }
-        }
-
-        for (int i_depth = 0; i_depth < depth_micro_repeats; ++i_depth) {
-          // The inner 3 loops go through the sub-block, depth and width within
-          // each micro block.
-
-          // Load, and apply symmetric offset.
-          int8_t* scratch_data =
-              scratch_block_data + k_height * workspace_height_stride +
-              j_width * 4 * 8 + i_depth * 4 * 8 * width_overall_micro_repeats;
-          const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-              input_data = input_block_data + k_height * input_height_stride +
-                           j_width * 4 * input_depth + i_depth * 8;
-          // Full-size macro blocks are 2*4*4 = 32 bytes.
-          for (int x = start_width; x < adjusted_residual_width; ++x) {
-            for (int s = 0; s < 2; ++s) {
-              for (int d = 0; d < 4; ++d) {
-                tmp_load[x][s][d] = input_data[x * input_depth + 4 * s + d] -
-                                    kSymmetricZeroPoint;
-              }
-            }
-          }
-
-          // Save results.
-          memcpy(&scratch_data[0], tmp_load[0][0], 8);
-          memcpy(&scratch_data[8], tmp_load[1][0], 8);
-          memcpy(&scratch_data[16], tmp_load[2][0], 8);
-          memcpy(&scratch_data[24], tmp_load[3][0], 8);
-        }
-      }
-    }
-
-    if (trailing_height_padding) {
-      memset(scratch_block_data + copy_block_height * workspace_height_stride,
-             -input_offset_difference, workspace_height_stride);
-    }
-  }
-
-  // Transpose 4x4 blocks within each sub-micro-block.
-  //
-  // Implemented somewhat like NEON register manipulation, so that we can see
-  // equivalence of the two approaches.
-  static inline void MicroTransposeBlocks(
-      const DepthwiseConvDotProdParams& function_params,
-      int8_t* scratch_block_data) {
-    const int workspace_height_stride = function_params.workspace_height_stride;
-    const int width_overall_micro_repeats =
-        function_params.input_width_overall_micro_repeats;
-    const int depth_micro_repeats = function_params.depth_micro_repeats;
-    const int block_height = function_params.inbound_block_height;
-
-    // Transpositions are 4x4, but doing 2 at a time is more efficient in the
-    // NEON code we are simulating.
-    int8_t tmp_load[4][2][4];         // [width][sub-block][depth]
-    int8_t tmp_transposed[4][2][4];   // [depth][sub-block][width]
-    int8_t tmp_interleaved[2][4][4];  // [sub-block][depth][width]
-
-    // The outer 3 loops go through all the micro blocks in a macro block.
-    for (int k_height = 0; k_height < block_height; ++k_height) {
-      for (int j_width = 0; j_width < width_overall_micro_repeats; ++j_width) {
-        for (int i_depth = 0; i_depth < depth_micro_repeats; ++i_depth) {
-          int8_t* scratch_data =
-              scratch_block_data + k_height * workspace_height_stride +
-              j_width * 4 * 8 + i_depth * 4 * 8 * width_overall_micro_repeats;
-          // A. Load data
-          memcpy(tmp_load[0][0], &scratch_data[0], 8);
-          memcpy(tmp_load[1][0], &scratch_data[8], 8);
-          memcpy(tmp_load[2][0], &scratch_data[16], 8);
-          memcpy(tmp_load[3][0], &scratch_data[24], 8);
-
-          // B. Simulate between-register transposition.
-          for (int x = 0; x < 4; ++x) {
-            for (int y = 0; y < 4; ++y) {
-              tmp_transposed[x][0][y] = tmp_load[y][0][x];
-              tmp_transposed[x][1][y] = tmp_load[y][1][x];
-            }
-          }
-
-          // C. Simulate between-register interleaving.
-          for (int x = 0; x < 4; ++x) {
-            for (int y = 0; y < 4; ++y) {
-              tmp_interleaved[0][x][y] = tmp_transposed[x][0][y];
-              tmp_interleaved[1][x][y] = tmp_transposed[x][1][y];
-            }
-          }
-          // D. Simulate mangled storage arrangement.
-          memcpy(&scratch_data[0], tmp_interleaved[0][0], 16);
-          memcpy(&scratch_data[16], tmp_interleaved[1][0], 16);
-        }
-      }
-    }
-  }
-
-  static inline void Run(
-      int32_t height_block_number, int32_t width_block_number,
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          input_block_data,
-      int8_t* scratch_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-    CopyMacroBlock(height_block_number, width_block_number, *function_params,
-                   input_block_data, scratch_block_data);
-    MicroTransposeBlocks(*function_params, scratch_block_data);
-  }
-};
-
-template <QuantizationType quantization_type, int32_t max_padding>
-struct PackMacroBlock<
-    DepthwiseConvImplementation::kUseCModel3x3DotProduct, quantization_type,
-    DepthwiseConvDepthMultiplication::kUnitInputDepth, max_padding> {
-  static inline void Run(
-      int32_t height_block_number, int32_t width_block_number,
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          input_block_data,
-      int8_t* scratch_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-    // Currently support for padding is limited to 1 on any side.
-    TFLITE_DCHECK_LE(max_padding, 1);
-
-    // Strides.
-    // The count of micro blocks (below) provides the width strides.
-    const int input_height_stride = function_params->input_height_stride;
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-
-    // Remaining iteration and dimension parameters.
-    //
-    // If width_overall_micro_repeats = input_width_micro_repeats + 1, then the
-    // final micro block is incomplete.
-    const int width_overall_micro_repeats =
-        function_params->input_width_overall_micro_repeats;
-    const int input_width_micro_repeats =
-        function_params->input_width_micro_repeats;
-    const int residual_width = function_params->residual_width;
-    const int block_height = function_params->inbound_block_height;
-    TFLITE_DCHECK_GE(workspace_height_stride, 4 * width_overall_micro_repeats);
-
-    const int padding_left = function_params->padding_left;
-    const int padding_right = function_params->padding_right;
-    const int padding_top = function_params->padding_top;
-    const int padding_bottom = function_params->padding_bottom;
-
-    const bool leading_width_padding =
-        padding_left > 0 && width_block_number == 0;
-    const bool trailing_width_padding =
-        padding_right > 0 &&
-        width_block_number == (function_params->width_macro_count - 1);
-    const bool leading_height_padding =
-        padding_top > 0 && height_block_number < 0;
-    const bool trailing_height_padding =
-        padding_bottom > 0 &&
-        height_block_number == (function_params->height_macro_count - 1);
-
-    constexpr int kSymmetricZeroPoint =
-        QuantizationTypeImpl<quantization_type>::kIntSymmetricZeroPoint;
-    const int32_t input_offset_difference =
-        function_params->input_offset + kSymmetricZeroPoint;
-
-    int copy_block_height = block_height;
-    if (leading_height_padding) {
-      memset(scratch_block_data, -input_offset_difference,
-             workspace_height_stride + kWorkspaceExtension);
-      scratch_block_data += workspace_height_stride;
-      input_block_data += input_height_stride;
-      copy_block_height -= 1;
-    }
-    if (trailing_height_padding) {
-      copy_block_height -= 1;
-    }
-
-    int adjusted_residual_width =
-        input_width_micro_repeats < width_overall_micro_repeats ? residual_width
-                                                                : 4;
-
-    if (trailing_width_padding) {
-      adjusted_residual_width -= 1;
-    }
-    int start_width = 0;
-    if (leading_width_padding) {
-      start_width = 1;
-      input_block_data += 1;
-    }
-
-    const int copy_size = (width_overall_micro_repeats - 1) * 4 +
-                          adjusted_residual_width - start_width;
-
-    TFLITE_DCHECK_LE(
-        copy_size,
-        input_height_stride - width_block_number * input_width_micro_repeats);
-    // We may drop up to stride-1 of trailing input.
-    TFLITE_DCHECK_GE(copy_size, input_height_stride - 1);
-
-    // When there is unit input depth, the micro-block iteration need only be
-    // through the height. The micro blocks are contiguous across the width.
-    for (int k_height = 0; k_height < copy_block_height; ++k_height) {
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          input_data = input_block_data + k_height * input_height_stride;
-      int8_t* scratch_data =
-          scratch_block_data + k_height * workspace_height_stride;
-
-      // Handle leading padding. This is overwritten if there is no padding.
-      scratch_data[0] = -input_offset_difference;
-
-      memcpy(&scratch_data[start_width], input_data, copy_size);
-      for (int i = 0; i < copy_size; ++i) {
-        scratch_data[start_width + i] += -kSymmetricZeroPoint;
-      }
-
-      // Handle trailing padding, and fill in remainder of micro block.
-      memset(&scratch_data[start_width + copy_size], -input_offset_difference,
-             4 - adjusted_residual_width + kWorkspaceExtension);
-    }
-
-    if (trailing_height_padding) {
-      memset(scratch_block_data + copy_block_height * workspace_height_stride,
-             -input_offset_difference,
-             workspace_height_stride + kWorkspaceExtension);
-    }
-  }
-};
-
-// Beginning of code section containing intermediate code transformation.
-//
-// This section is only compiled when kUseUnwound3x3DotProduct versions of
-// templated functions are selected.
-template <QuantizationType quantization_type>
-struct PackMacroBlock<DepthwiseConvImplementation::kUseUnwound3x3DotProduct,
-                      quantization_type,
-                      DepthwiseConvDepthMultiplication::kNoMultiplication,
-                      /*max_padding=*/0> {
-  static inline void Run(
-      int32_t height_block_number, int32_t width_block_number,
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          input_block_data,
-      int8_t* scratch_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-    const int width_overall_micro_repeats =
-        function_params->input_width_overall_micro_repeats;
-    const int input_width_micro_repeats =
-        function_params->input_width_micro_repeats;
-    const int depth_micro_repeats = function_params->depth_micro_repeats;
-    const int block_height = function_params->inbound_block_height;
-    const int residual_width = function_params->residual_width;
-    const int input_height_stride = function_params->input_height_stride;
-    const int input_depth = function_params->input_depth;
-
-    TFLITE_DCHECK_GE(depth_micro_repeats, 0);
-    constexpr int kSymmetricZeroPoint =
-        QuantizationTypeImpl<quantization_type>::kIntSymmetricZeroPoint;
-    const int micro_block_size = 4 * 8;
-    const int depth_advance = width_overall_micro_repeats * micro_block_size;
-    const int width_advance =
-        micro_block_size *
-        (1 - depth_micro_repeats * width_overall_micro_repeats);
-    const int height_advance = workspace_height_stride -
-                               width_overall_micro_repeats * micro_block_size;
-    const int input_depth_skip = 4 * input_depth - 8 * depth_micro_repeats;
-
-    // Transpositions are 4x4, but doing 2 at a time is more efficient in the
-    // NEON code we are simulating. Note the blocks of 4x4 are still interleaved
-    // down the depth.
-    int8_t tmp_load[4][2][4];
-    int8_t tmp_transposed[4][2][4];
-    int8_t tmp_interleaved[2][4][4];
-
-    // Work through one slice, by row, at a time.
-    int8_t* scratch_data = scratch_block_data;
-    for (int k_height = 0; k_height < block_height; ++k_height) {
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          input_data = input_block_data;
-      input_block_data += input_height_stride;
-
-      // Traverse the width one point at a time, but the depth in (micro) blocks
-      // of size 8.
-      //
-      // The depth and width margins, which are filled with "zeros", may be
-      // larger than is strictly needed to calculate output. This is because the
-      // conv calculation is performed across complete micro blocks.
-      for (int j_width = 0; j_width < input_width_micro_repeats; ++j_width) {
-        // Load, then zero.
-        for (int i_depth = 0; i_depth < depth_micro_repeats; ++i_depth) {
-          // A. Simulate register loading.
-          for (int x = 0; x < 4; ++x) {
-            for (int s = 0; s < 2; ++s) {
-              for (int d = 0; d < 4; ++d) {
-                tmp_load[x][s][d] = input_data[x * input_depth + 4 * s + d] -
-                                    kSymmetricZeroPoint;
-              }
-            }
-          }
-          // B. Simulate between-register transposition.
-          for (int x = 0; x < 4; ++x) {
-            for (int y = 0; y < 4; ++y) {
-              tmp_transposed[x][0][y] = tmp_load[y][0][x];
-              tmp_transposed[x][1][y] = tmp_load[y][1][x];
-            }
-          }
-
-          // C and D are to be performed together as 4-byte stores in NEON code.
-          // C. Simulate between-register interleaving.
-          for (int x = 0; x < 4; ++x) {
-            for (int y = 0; y < 4; ++y) {
-              tmp_interleaved[0][x][y] = tmp_transposed[x][0][y];
-              tmp_interleaved[1][x][y] = tmp_transposed[x][1][y];
-            }
-          }
-          // D. Simulate mangled storage arrangement.
-          memcpy(&scratch_data[0], tmp_interleaved[0][0], 8);
-          memcpy(&scratch_data[8], tmp_interleaved[0][2], 8);
-          memcpy(&scratch_data[16], tmp_interleaved[1][0], 8);
-          memcpy(&scratch_data[24], tmp_interleaved[1][2], 8);
-
-          scratch_data += depth_advance;
-          input_data += 8;
-        }
-        scratch_data += width_advance;
-        input_data += input_depth_skip;
-      }
-      if (width_overall_micro_repeats > input_width_micro_repeats) {
-        TFLITE_DCHECK_EQ(width_overall_micro_repeats,
-                         input_width_micro_repeats + 1);
-        TFLITE_DCHECK_GT(residual_width, 0);
-        // Figure out division of work (available input vs zero-ed).
-        const int adjusted_residual_width = residual_width;
-        // Load, then zero.
-        for (int i_depth = 0; i_depth < depth_micro_repeats; ++i_depth) {
-          // A. Simulate register loading.
-          for (int x = 0; x < adjusted_residual_width; ++x) {
-            for (int s = 0; s < 2; ++s) {
-              for (int d = 0; d < 4; ++d) {
-                tmp_load[x][s][d] = input_data[x * input_depth + 4 * s + d] -
-                                    kSymmetricZeroPoint;
-              }
-            }
-          }
-          for (int x = adjusted_residual_width; x < 4; ++x) {
-            for (int s = 0; s < 2; ++s) {
-              for (int d = 0; d < 4; ++d) {
-                tmp_load[x][s][d] = 0;
-              }
-            }
-          }
-          // B. Simulate between-register transposition.
-          for (int x = 0; x < 4; ++x) {
-            for (int y = 0; y < 4; ++y) {
-              tmp_transposed[x][0][y] = tmp_load[y][0][x];
-              tmp_transposed[x][1][y] = tmp_load[y][1][x];
-            }
-          }
-
-          // C and D are to be performed together as 4-byte stores in NEON code.
-          // C. Simulate between-register interleaving.
-          for (int x = 0; x < 4; ++x) {
-            for (int y = 0; y < 4; ++y) {
-              tmp_interleaved[0][x][y] = tmp_transposed[x][0][y];
-              tmp_interleaved[1][x][y] = tmp_transposed[x][1][y];
-            }
-          }
-          // D. Simulate mangled storage arrangement.
-          memcpy(&scratch_data[0], tmp_interleaved[0][0], 8);
-          memcpy(&scratch_data[8], tmp_interleaved[0][2], 8);
-          memcpy(&scratch_data[16], tmp_interleaved[1][0], 8);
-          memcpy(&scratch_data[24], tmp_interleaved[1][2], 8);
-
-          scratch_data += depth_advance;
-          input_data += 8;
-        }
-        scratch_data += width_advance;
-        input_data += input_depth_skip;
-      }
-      scratch_data += height_advance;
-    }
-
-    TFLITE_DCHECK_EQ(scratch_data, scratch_block_data +
-                                       block_height * workspace_height_stride);
-  }
-};
-
-template <QuantizationType quantization_type>
-struct PackMacroBlock<DepthwiseConvImplementation::kUseUnwound3x3DotProduct,
-                      quantization_type,
-                      DepthwiseConvDepthMultiplication::kNoMultiplication,
-                      /*max_padding=*/1> {
-  static inline void Run(
-      int32_t height_block_number, int32_t width_block_number,
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          input_block_data,
-      int8_t* scratch_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-    // Just use C model code for case of padding. Optimized versions merge the
-    // modifications therein to handle padding.
-    PackMacroBlock<DepthwiseConvImplementation::kUseCModel3x3DotProduct,
-                   quantization_type,
-                   DepthwiseConvDepthMultiplication::kNoMultiplication,
-                   /*max_padding=*/1>::Run(height_block_number,
-                                           width_block_number, input_block_data,
-                                           scratch_block_data, function_params);
-  }
-};
-
-template <QuantizationType quantization_type, int32_t max_padding>
-struct PackMacroBlock<
-    DepthwiseConvImplementation::kUseUnwound3x3DotProduct, quantization_type,
-    DepthwiseConvDepthMultiplication::kUnitInputDepth, max_padding> {
-  static inline void Run(
-      int32_t height_block_number, int32_t width_block_number,
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          input_block_data,
-      int8_t* scratch_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-    const int width_overall_micro_repeats =
-        function_params->input_width_overall_micro_repeats;
-    const int input_width_micro_repeats =
-        function_params->input_width_micro_repeats;
-    const int block_height = function_params->inbound_block_height;
-    const int residual_width = function_params->residual_width;
-    const int input_height_stride = function_params->input_height_stride;
-
-    const int padding_left = function_params->padding_left;
-    const int padding_right = function_params->padding_right;
-    const int padding_top = function_params->padding_top;
-    const int padding_bottom = function_params->padding_bottom;
-
-    constexpr int kSymmetricZeroPoint =
-        QuantizationTypeImpl<quantization_type>::kIntSymmetricZeroPoint;
-
-    TFLITE_DCHECK_GE(workspace_height_stride, 4 * width_overall_micro_repeats);
-
-    const bool leading_width_padding =
-        padding_left > 0 && width_block_number == 0;
-    const bool trailing_width_padding =
-        padding_right > 0 &&
-        width_block_number == (function_params->width_macro_count - 1);
-    const bool leading_height_padding =
-        padding_top > 0 && height_block_number < 0;
-    const bool trailing_height_padding =
-        padding_bottom > 0 &&
-        height_block_number == (function_params->height_macro_count - 1);
-
-    const int32_t input_offset = function_params->input_offset;
-    const int32_t input_offset_difference = input_offset + kSymmetricZeroPoint;
-
-    // Work through one slice, by row, at a time.
-    int8_t* scratch_data_base = scratch_block_data;
-
-    int copy_block_height = block_height;
-    if (leading_height_padding) {
-      copy_block_height -= 1;
-      memset(scratch_data_base, -input_offset_difference,
-             workspace_height_stride + kWorkspaceExtension);
-      scratch_data_base += workspace_height_stride;
-      input_block_data += input_height_stride;
-    }
-    if (trailing_height_padding) {
-      copy_block_height -= 1;
-    }
-
-    int adjusted_residual_width =
-        input_width_micro_repeats < width_overall_micro_repeats ? residual_width
-                                                                : 4;
-
-    if (trailing_width_padding) {
-      adjusted_residual_width -= 1;
-    }
-    int start_width = 0;
-    if (leading_width_padding) {
-      start_width = 1;
-      input_block_data += 1;
-    }
-
-    const int copy_size = (width_overall_micro_repeats - 1) * 4 +
-                          adjusted_residual_width - start_width;
-    // Adjusted so that later conditionals are simplified.
-    const int copy_size_adjusted =
-        trailing_width_padding ? copy_size + 1 : copy_size;
-
-    TFLITE_DCHECK_LE(
-        copy_size,
-        input_height_stride - width_block_number * input_width_micro_repeats);
-    // We may drop up to stride-1 of trailing input.
-    TFLITE_DCHECK_GE(copy_size, input_height_stride - 1);
-
-    // This is used to simulate what should happen in registers.
-    int8_t tmp_data[16];
-
-    int scratch_data_offset = 0;
-    int input_block_offset = 0;
-
-    if (copy_size >= 16) {
-      for (int k_height = 0; k_height < copy_block_height; ++k_height) {
-        // Work through one slice, by row, at a time.
-        int8_t* scratch_data = scratch_data_base + scratch_data_offset;
-
-        int copy_done = 0;
-
-        // The surrounding condition ensures that we always need at least one
-        // iteration of the main copy loop. In the case of leading width
-        // padding, we unroll this specially.
-        if (leading_width_padding) {
-          memcpy(tmp_data + 1, input_block_data + input_block_offset, 15);
-          for (int i = 0; i < 16; ++i) {
-            tmp_data[i] += -kSymmetricZeroPoint;
-          }
-          tmp_data[0] = -input_offset_difference;
-          memcpy(scratch_data, tmp_data, 16);
-          copy_done += 15;
-        }
-
-        // Main copy loop.
-        for (; (copy_done + 16) <= copy_size; copy_done += 16) {
-          memcpy(tmp_data, input_block_data + input_block_offset + copy_done,
-                 16);
-          for (int i = 0; i < 16; ++i) {
-            tmp_data[i] += -kSymmetricZeroPoint;
-          }
-          TFLITE_DCHECK_EQ((start_width + copy_done) % 16, 0);
-          memcpy(&scratch_data[start_width + copy_done], tmp_data, 16);
-        }
-
-        const int copy_remaining = copy_size - copy_done;
-        // Total amount
-        // = copy_size - copy_done + 4 - adjusted_residual_width
-        // = width_overall_micro_repeats * 4 - start_width - copy_done.
-        // Undone micro blocks
-        // = width_overall_micro_repeats - (start_width + copy_done) / 4.
-
-        // Conditional is (copy_remaining > 0 || trailing_width_padding).
-        if (copy_done < copy_size_adjusted) {
-          // Employ overlapping-load strategy in order to load full register,
-          // but use only part.
-          memcpy(tmp_data,
-                 input_block_data + input_block_offset + copy_done -
-                     (16 - copy_remaining),
-                 16);
-          // Shift to select the part that we need.
-          for (int i = 0; i < copy_remaining; ++i) {
-            tmp_data[i] = tmp_data[(16 - copy_remaining) + i];
-          }
-          for (int i = 0; i < 16; ++i) {
-            tmp_data[i] += -kSymmetricZeroPoint;
-          }
-          // Apply padding to remainder, some unnecessary but costless in regs.
-          for (int i = copy_remaining; i < 16; ++i) {
-            tmp_data[i] = -input_offset_difference;
-          }
-          const int final_repeats =
-              width_overall_micro_repeats - (start_width + copy_done) / 4;
-          for (int i = 0; i < final_repeats; ++i) {
-            memcpy(&scratch_data[start_width + copy_done], tmp_data + 4 * i, 4);
-            copy_done += 4;
-          }
-        }
-        memset(scratch_data + start_width + copy_done, -input_offset_difference,
-               kWorkspaceExtension);
-
-        scratch_data_offset += workspace_height_stride;
-        input_block_offset += input_height_stride;
-      }
-    } else if (copy_size >= 4) {
-      for (int k_height = 0; k_height < copy_block_height; ++k_height) {
-        // Work through one slice, by row, at a time.
-        int8_t* scratch_data = scratch_data_base + scratch_data_offset;
-
-        int copy_done = 0;
-
-        // The surrounding condition ensures that we always need at least one
-        // iteration of the main copy loop. In the case of leading width
-        // padding, we unroll this specially.
-        if (leading_width_padding) {
-          memcpy(tmp_data + 1, input_block_data + input_block_offset, 3);
-          for (int i = 0; i < 4; ++i) {
-            tmp_data[i] += -kSymmetricZeroPoint;
-          }
-          tmp_data[0] = -input_offset_difference;
-          memcpy(scratch_data, tmp_data, 4);
-          copy_done += 3;
-        }
-
-        for (; (copy_done + 4) <= copy_size; copy_done += 4) {
-          memcpy(tmp_data, input_block_data + input_block_offset + copy_done,
-                 4);
-          for (int i = 0; i < 4; ++i) {
-            tmp_data[i] += -kSymmetricZeroPoint;
-          }
-          // Perform as 4 int32 stores, because that is our alignment.
-          memcpy(&scratch_data[start_width + copy_done], tmp_data, 4);
-        }
-
-        // Total amount
-        // = copy_size - copy_done + 4 - adjusted_residual_width
-        // = width_overall_micro_repeats * 4 - start_width - copy_done.
-        // Undone micro blocks
-        // = width_overall_micro_repeats - (start_width + copy_done) / 4.
-        const int copy_remaining = copy_size - copy_done;
-        // Conditional is (copy_remaining > 0 || trailing_width_padding).
-        if (copy_done < copy_size_adjusted) {
-          TFLITE_DCHECK_LT(copy_remaining, 4);
-          // Employ overlapping-load strategy in order to load full register,
-          // but use only part.
-          memcpy(tmp_data,
-                 input_block_data + input_block_offset + copy_done -
-                     (4 - copy_remaining),
-                 4);
-          // Shift to select the part that we need.
-          for (int i = 0; i < copy_remaining; ++i) {
-            tmp_data[i] = tmp_data[(4 - copy_remaining) + i];
-          }
-          for (int i = 0; i < 4; ++i) {
-            tmp_data[i] += -kSymmetricZeroPoint;
-          }
-          // Apply padding to remainder, some unnecessary but costless in regs.
-          for (int i = copy_remaining; i < 4; ++i) {
-            tmp_data[i] = -input_offset_difference;
-          }
-          memcpy(&scratch_data[start_width + copy_done], tmp_data, 4);
-          copy_done += 4;
-        }
-        memset(scratch_data + start_width + copy_done, -input_offset_difference,
-               kWorkspaceExtension);
-
-        scratch_data_offset += workspace_height_stride;
-        input_block_offset += input_height_stride;
-      }
-    } else if (width_overall_micro_repeats == 2) {
-      for (int k_height = 0; k_height < copy_block_height; ++k_height) {
-        // Apply padding by quick fill of whole reg.
-        for (int i = 0; i < 8; ++i) {
-          tmp_data[i] = -input_offset;
-        }
-        for (int i = 0; i < copy_size; ++i) {
-          // Apply shift-left insert, tmp_data as both operands.
-          // The zero-index byte is left unchanged.
-          for (int i = 7; i > 0; --i) {
-            tmp_data[i] = tmp_data[i - 1];
-          }
-          tmp_data[1] =
-              input_block_data[input_block_offset + (copy_size - 1 - i)];
-        }
-        if (!leading_width_padding) {
-          // Remove leading padding, junking trailing byte, OK because max size
-          // is less than 8.
-          TFLITE_DCHECK_LT(copy_size_adjusted + start_width, 8);
-          for (int i = 0; i < 7; ++i) {
-            tmp_data[i] = tmp_data[i + 1];
-          }
-        }
-        for (int i = 0; i < 8; ++i) {
-          tmp_data[i] += -kSymmetricZeroPoint;
-        }
-        memcpy(scratch_data_base + scratch_data_offset, tmp_data, 8);
-        memset(scratch_data_base + scratch_data_offset + 8,
-               -input_offset_difference, kWorkspaceExtension);
-
-        scratch_data_offset += workspace_height_stride;
-        input_block_offset += input_height_stride;
-      }
-    } else {
-      TFLITE_DCHECK_EQ(width_overall_micro_repeats, 1);
-      // This path is basically the same as the preceding, 2-micro-block one,
-      // but here we simply store fewer bytes.
-      for (int k_height = 0; k_height < copy_block_height; ++k_height) {
-        // Apply padding by quick fill of whole reg.
-        for (int i = 0; i < 8; ++i) {
-          tmp_data[i] = -input_offset;
-        }
-        for (int i = 0; i < copy_size; ++i) {
-          // Apply shift-left insert, tmp_data as both operands.
-          // The zero-index byte is left unchanged.
-          for (int i = 7; i > 0; --i) {
-            tmp_data[i] = tmp_data[i - 1];
-          }
-          tmp_data[1] =
-              input_block_data[input_block_offset + (copy_size - 1 - i)];
-        }
-        if (!leading_width_padding) {
-          // Remove leading padding, junking trailing byte, OK because max size
-          // is less than 8.
-          TFLITE_DCHECK_LT(copy_size_adjusted + start_width, 8);
-          for (int i = 0; i < 7; ++i) {
-            tmp_data[i] = tmp_data[i + 1];
-          }
-        }
-        for (int i = 0; i < 8; ++i) {
-          tmp_data[i] += -kSymmetricZeroPoint;
-        }
-        memcpy(scratch_data_base + scratch_data_offset, tmp_data, 4);
-        memset(scratch_data_base + scratch_data_offset + 4,
-               -input_offset_difference, kWorkspaceExtension);
-
-        scratch_data_offset += workspace_height_stride;
-        input_block_offset += input_height_stride;
-      }
-    }
-
-    scratch_data_base += copy_block_height * workspace_height_stride;
-
-    if (trailing_height_padding) {
-      memset(scratch_data_base, -input_offset_difference,
-             workspace_height_stride + kWorkspaceExtension);
-      scratch_data_base += workspace_height_stride;
-    }
-
-    TFLITE_DCHECK_EQ(
-        scratch_data_base,
-        scratch_block_data + block_height * workspace_height_stride);
-  }
-};
-// The preceding section is only compiled when kUseUnwound3x3DotProduct versions
-// of templated functions are selected.
-//
-// End of code section containing intermediate code transformation.
-
-#ifdef USE_NEON
-template <QuantizationType quantization_type>
-struct PackMacroBlock<DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct,
-                      quantization_type,
-                      DepthwiseConvDepthMultiplication::kNoMultiplication,
-                      /*max_padding=*/0> {
-  static inline void PackMacroBlockIntrinsics(
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          input_block_data,
-      int8_t* scratch_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-    TFLITE_DCHECK_EQ(function_params->padding_bottom, 0);
-    TFLITE_DCHECK_EQ(function_params->padding_top, 0);
-    TFLITE_DCHECK_EQ(function_params->padding_left, 0);
-    TFLITE_DCHECK_EQ(function_params->padding_right, 0);
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-    const int width_overall_micro_repeats =
-        function_params->input_width_overall_micro_repeats;
-    const int input_width_micro_repeats =
-        function_params->input_width_micro_repeats;
-    const int depth_micro_repeats = function_params->depth_micro_repeats;
-    const int block_height = function_params->inbound_block_height;
-    const int residual_width = function_params->residual_width;
-    const int input_height_stride = function_params->input_height_stride;
-    const int input_depth = function_params->input_depth;
-
-    TFLITE_DCHECK_GE(depth_micro_repeats, 0);
-    constexpr uint8_t kSignBit =
-        QuantizationTypeImpl<quantization_type>::kUint8SignBit;
-    const int micro_block_size = 4 * 8;
-    const int depth_advance = width_overall_micro_repeats * micro_block_size;
-    const int width_advance =
-        micro_block_size *
-        (1 - depth_micro_repeats * width_overall_micro_repeats);
-    const int height_advance = workspace_height_stride -
-                               width_overall_micro_repeats * micro_block_size;
-    const int input_depth_skip = 4 * input_depth - 8 * depth_micro_repeats;
-
-    // Transpositions are 4x4, but doing 2 at a time is more efficient in NEON
-    // code. Note the blocks of 4x4 are still interleaved down the depth.
-    int8x16_t work_reg_a;
-    int8x16_t work_reg_b;
-
-    // Effect subtraction of zero-point = 128 by XOR of sign bit.
-    const int8x16_t sign_bit = vreinterpretq_s8_u8(vdupq_n_u8(kSignBit));
-
-    // Work through one slice, by row, at a time.
-    int8_t* scratch_data_0 = scratch_block_data;
-
-    for (int k_height = 0; k_height < block_height; ++k_height) {
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          input_data_0 = input_block_data;
-      int8x16_t input_data_a;
-      int8x16_t input_data_b;
-      int8x16_t input_data_c;
-      int8x16_t input_data_d;
-
-      // Traverse the width one point at a time, but the depth in (micro) blocks
-      // of size 8.
-      //
-      // The depth and width margins, which are filled with "zeros", may be
-      // larger than is strictly needed to calculate output. This is because the
-      // conv calculation is performed across complete micro blocks.
-      for (int j_width = 0; j_width < input_width_micro_repeats; ++j_width) {
-        int8x16_t work_reg_a_sp;
-        int8x16_t work_reg_b_sp;
-
-        int i_depth = 0;
-
-        if (depth_micro_repeats >= 2) {
-          i_depth += 2;
-
-          input_data_a = util_vld1q_x8(input_data_0);
-          input_data_b = util_vld1q_x8(input_data_0 + 1 * input_depth);
-          input_data_c = util_vld1q_x8(input_data_0 + 2 * input_depth);
-          input_data_d = util_vld1q_x8(input_data_0 + 3 * input_depth);
-          input_data_0 += 16;
-
-          for (; i_depth < depth_micro_repeats - 1; i_depth += 2) {
-            work_reg_a = vzip1q_s8(input_data_a, input_data_b);
-            work_reg_b = vzip1q_s8(input_data_c, input_data_d);
-            vzipq_s8x2_in_place(&work_reg_a, &work_reg_b);
-            if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-              work_reg_a = veorq_s8(work_reg_a, sign_bit);
-              work_reg_b = veorq_s8(work_reg_b, sign_bit);
-            }
-
-            work_reg_a_sp = vzip2q_s8(input_data_a, input_data_b);
-            work_reg_b_sp = vzip2q_s8(input_data_c, input_data_d);
-            vzipq_s8x2_in_place(&work_reg_a_sp, &work_reg_b_sp);
-
-            input_data_a = util_vld1q_x8(input_data_0);
-            input_data_b = util_vld1q_x8(input_data_0 + 1 * input_depth);
-            vst1q_s8(scratch_data_0, work_reg_a);
-            vst1q_s8(scratch_data_0 + 16, work_reg_b);
-
-            scratch_data_0 += depth_advance;
-
-            if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-              work_reg_a_sp = veorq_s8(work_reg_a_sp, sign_bit);
-              work_reg_b_sp = veorq_s8(work_reg_b_sp, sign_bit);
-            }
-
-            input_data_c = util_vld1q_x8(input_data_0 + 2 * input_depth);
-            input_data_d = util_vld1q_x8(input_data_0 + 3 * input_depth);
-            vst1q_s8(scratch_data_0, work_reg_a_sp);
-            vst1q_s8(scratch_data_0 + 16, work_reg_b_sp);
-
-            scratch_data_0 += depth_advance;
-            input_data_0 += 16;
-          }
-
-          work_reg_a = vzip1q_s8(input_data_a, input_data_b);
-          work_reg_b = vzip1q_s8(input_data_c, input_data_d);
-          vzipq_s8x2_in_place(&work_reg_a, &work_reg_b);
-          if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-            work_reg_a = veorq_s8(work_reg_a, sign_bit);
-            work_reg_b = veorq_s8(work_reg_b, sign_bit);
-          }
-          vst1q_s8(scratch_data_0, work_reg_a);
-          vst1q_s8(scratch_data_0 + 16, work_reg_b);
-
-          scratch_data_0 += depth_advance;
-
-          work_reg_a_sp = vzip2q_s8(input_data_a, input_data_b);
-          work_reg_b_sp = vzip2q_s8(input_data_c, input_data_d);
-          vzipq_s8x2_in_place(&work_reg_a_sp, &work_reg_b_sp);
-          if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-            work_reg_a_sp = veorq_s8(work_reg_a_sp, sign_bit);
-            work_reg_b_sp = veorq_s8(work_reg_b_sp, sign_bit);
-          }
-
-          vst1q_s8(scratch_data_0, work_reg_a_sp);
-          vst1q_s8(scratch_data_0 + 16, work_reg_b_sp);
-
-          scratch_data_0 += depth_advance;
-        }
-        for (; i_depth < depth_micro_repeats; ++i_depth) {
-          input_data_a = vld1q_lane_s8x8(input_data_0, input_data_a, 0);
-          input_data_b =
-              vld1q_lane_s8x8(input_data_0 + 1 * input_depth, input_data_b, 0);
-          input_data_c =
-              vld1q_lane_s8x8(input_data_0 + 2 * input_depth, input_data_c, 0);
-          input_data_d =
-              vld1q_lane_s8x8(input_data_0 + 3 * input_depth, input_data_d, 0);
-          work_reg_a = vzip1q_s8(input_data_a, input_data_b);
-          work_reg_b = vzip1q_s8(input_data_c, input_data_d);
-
-          input_data_0 += 8;
-
-          vzipq_s8x2_in_place(&work_reg_a, &work_reg_b);
-          if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-            work_reg_a = veorq_s8(work_reg_a, sign_bit);
-            work_reg_b = veorq_s8(work_reg_b, sign_bit);
-          }
-
-          vst1q_s8(scratch_data_0, work_reg_a);
-          vst1q_s8(scratch_data_0 + 16, work_reg_b);
-
-          scratch_data_0 += depth_advance;
-        }
-        scratch_data_0 += width_advance;
-        input_data_0 += input_depth_skip;
-      }
-      if (width_overall_micro_repeats > input_width_micro_repeats) {
-        TFLITE_DCHECK_EQ(width_overall_micro_repeats,
-                         input_width_micro_repeats + 1);
-        TFLITE_DCHECK_GT(residual_width, 0);
-        TFLITE_DCHECK_LT(residual_width, 4);
-        for (int i_depth = 0; i_depth < depth_micro_repeats; ++i_depth) {
-          input_data_c = vreinterpretq_s8_u8(vdupq_n_u8(kSignBit));
-          input_data_a = vld1q_lane_s8x8(input_data_0, input_data_a, 0);
-          input_data_d = vreinterpretq_s8_u8(vdupq_n_u8(kSignBit));
-          if (residual_width > 1) {
-            input_data_b =
-                vld1q_lane_s8x8(input_data_0 + input_depth, input_data_b, 0);
-            if (residual_width == 3) {
-              input_data_c = vld1q_lane_s8x8(input_data_0 + 2 * input_depth,
-                                             input_data_c, 0);
-            }
-          }
-          work_reg_a = vzip1q_s8(input_data_a, input_data_b);
-          work_reg_b = vzip1q_s8(input_data_c, input_data_d);
-
-          if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-            work_reg_a = veorq_s8(work_reg_a, sign_bit);
-            work_reg_b = veorq_s8(work_reg_b, sign_bit);
-          }
-          vzipq_s8x2_in_place(&work_reg_a, &work_reg_b);
-
-          vst1q_s8(scratch_data_0, work_reg_a);
-          vst1q_s8(scratch_data_0 + 16, work_reg_b);
-
-          scratch_data_0 += depth_advance;
-          input_data_0 += 8;
-        }
-        scratch_data_0 += width_advance;
-        input_data_0 += input_depth_skip;
-      }
-
-      scratch_data_0 += height_advance;
-      input_block_data += input_height_stride;
-    }
-    TFLITE_DCHECK_EQ(
-        scratch_data_0,
-        scratch_block_data + block_height * workspace_height_stride);
-  }
-
-  static inline void Run(
-      int32_t height_block_number, int32_t width_block_number,
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          input_block_data,
-      int8_t* scratch_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-#ifdef __aarch64__
-    PreloadInputBlock(input_block_data, function_params);
-#endif
-    PackMacroBlockIntrinsics(input_block_data, scratch_block_data,
-                             function_params);
-  }
-};
-
-template <QuantizationType quantization_type>
-struct PackMacroBlock<DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct,
-                      quantization_type,
-                      DepthwiseConvDepthMultiplication::kNoMultiplication,
-                      /*max_padding=*/1> {
-  static inline void PackMacroBlockIntrinsics(
-      int32_t height_block_number, int32_t width_block_number,
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          input_block_data,
-      int8_t* scratch_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-    constexpr uint8_t kSignBit =
-        QuantizationTypeImpl<quantization_type>::kUint8SignBit;
-
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-    const int width_overall_micro_repeats =
-        function_params->input_width_overall_micro_repeats;
-    const int input_width_micro_repeats =
-        function_params->input_width_micro_repeats;
-    const int depth_micro_repeats = function_params->depth_micro_repeats;
-    const int block_height = function_params->inbound_block_height;
-    const int residual_width = function_params->residual_width;
-    const int input_height_stride = function_params->input_height_stride;
-    const int input_depth = function_params->input_depth;
-
-    const int padding_left = function_params->padding_left;
-    const int padding_right = function_params->padding_right;
-    const int padding_top = function_params->padding_top;
-    const int padding_bottom = function_params->padding_bottom;
-
-    TFLITE_DCHECK_GT(depth_micro_repeats, 0);
-    constexpr int kSymmetricZeroPoint =
-        QuantizationTypeImpl<quantization_type>::kIntSymmetricZeroPoint;
-
-    const int micro_block_size = 4 * 8;
-    const int depth_advance = width_overall_micro_repeats * micro_block_size;
-    const int width_advance =
-        micro_block_size *
-        (1 - depth_micro_repeats * width_overall_micro_repeats);
-    const int height_advance = workspace_height_stride -
-                               width_overall_micro_repeats * micro_block_size;
-    const int input_depth_skip = 4 * input_depth - 8 * depth_micro_repeats;
-
-    const bool leading_width_padding =
-        padding_left > 0 && width_block_number == 0;
-    const bool trailing_width_padding =
-        padding_right > 0 &&
-        width_block_number == (function_params->width_macro_count - 1);
-    const bool leading_height_padding =
-        padding_top > 0 && height_block_number < 0;
-    const bool trailing_height_padding =
-        padding_bottom > 0 &&
-        height_block_number == (function_params->height_macro_count - 1);
-
-    const int32_t input_offset = function_params->input_offset;
-    const int32_t input_offset_difference = input_offset + kSymmetricZeroPoint;
-
-    // Transpositions are 4x4, but doing 2 at a time is more efficient in NEON
-    // code. Note the blocks of 4x4 are still interleaved down the depth.
-    int8x16_t work_reg_a;
-    int8x16_t work_reg_b;
-
-    // Effect subtraction of zero-point = 128 by XOR of sign bit.
-    const int8x16_t sign_bit = vreinterpretq_s8_u8(vdupq_n_u8(kSignBit));
-
-    // Work through one slice, by row, at a time.
-    int8_t* scratch_data_0 = scratch_block_data;
-
-    int copy_block_height = block_height;
-    if (leading_height_padding) {
-      copy_block_height -= 1;
-      memset(scratch_data_0, -input_offset_difference, workspace_height_stride);
-      scratch_data_0 += workspace_height_stride;
-      input_block_data += input_height_stride;
-    }
-    if (trailing_height_padding) {
-      copy_block_height -= 1;
-    }
-
-    for (int k_height = 0; k_height < copy_block_height; ++k_height) {
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          input_data_0 = input_block_data;
-      int8x16_t input_data_a;
-      int8x16_t input_data_b;
-      int8x16_t input_data_c;
-      int8x16_t input_data_d;
-
-      // Traverse the width one point at a time, but the depth in (micro) blocks
-      // of size 8.
-      //
-      // The depth and width margins, which are filled with "zeros", may be
-      // larger than is strictly needed to calculate output. This is because the
-      // conv calculation is performed across complete micro blocks.
-      for (int j_width = 0; j_width < width_overall_micro_repeats; ++j_width) {
-        // Figure out division of work (available input vs zero-ed).
-        int adjusted_residual_width =
-            j_width == (input_width_micro_repeats) ? residual_width : 4;
-
-        if (trailing_width_padding &&
-            j_width == (width_overall_micro_repeats - 1)) {
-          adjusted_residual_width -= 1;
-        }
-        int start_width = 0;
-        if (leading_width_padding && j_width == 0) {
-          start_width = 1;
-        }
-        if (start_width == 0) {
-          if (adjusted_residual_width == 4) {
-            int8x16_t work_reg_a_sp;
-            int8x16_t work_reg_b_sp;
-
-            int i_depth = 0;
-
-            if (depth_micro_repeats >= 2) {
-              i_depth += 2;
-
-              input_data_a = util_vld1q_x8(input_data_0);
-              input_data_b = util_vld1q_x8(input_data_0 + 1 * input_depth);
-              input_data_c = util_vld1q_x8(input_data_0 + 2 * input_depth);
-              input_data_d = util_vld1q_x8(input_data_0 + 3 * input_depth);
-              input_data_0 += 16;
-
-              for (; i_depth < depth_micro_repeats - 1; i_depth += 2) {
-                work_reg_a = vzip1q_s8(input_data_a, input_data_b);
-                work_reg_b = vzip1q_s8(input_data_c, input_data_d);
-                vzipq_s8x2_in_place(&work_reg_a, &work_reg_b);
-                if (quantization_type ==
-                    QuantizationType::kNonPerChannelUint8) {
-                  work_reg_a = veorq_s8(work_reg_a, sign_bit);
-                  work_reg_b = veorq_s8(work_reg_b, sign_bit);
-                }
-
-                work_reg_a_sp = vzip2q_s8(input_data_a, input_data_b);
-                work_reg_b_sp = vzip2q_s8(input_data_c, input_data_d);
-                vzipq_s8x2_in_place(&work_reg_a_sp, &work_reg_b_sp);
-
-                input_data_a = util_vld1q_x8(input_data_0);
-                input_data_b = util_vld1q_x8(input_data_0 + 1 * input_depth);
-                vst1q_s8(scratch_data_0, work_reg_a);
-                vst1q_s8(scratch_data_0 + 16, work_reg_b);
-
-                scratch_data_0 += depth_advance;
-
-                if (quantization_type ==
-                    QuantizationType::kNonPerChannelUint8) {
-                  work_reg_a_sp = veorq_s8(work_reg_a_sp, sign_bit);
-                  work_reg_b_sp = veorq_s8(work_reg_b_sp, sign_bit);
-                }
-
-                input_data_c = util_vld1q_x8(input_data_0 + 2 * input_depth);
-                input_data_d = util_vld1q_x8(input_data_0 + 3 * input_depth);
-                vst1q_s8(scratch_data_0, work_reg_a_sp);
-                vst1q_s8(scratch_data_0 + 16, work_reg_b_sp);
-
-                scratch_data_0 += depth_advance;
-                input_data_0 += 16;
-              }
-
-              work_reg_a = vzip1q_s8(input_data_a, input_data_b);
-              work_reg_b = vzip1q_s8(input_data_c, input_data_d);
-              vzipq_s8x2_in_place(&work_reg_a, &work_reg_b);
-              if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-                work_reg_a = veorq_s8(work_reg_a, sign_bit);
-                work_reg_b = veorq_s8(work_reg_b, sign_bit);
-              }
-              vst1q_s8(scratch_data_0, work_reg_a);
-              vst1q_s8(scratch_data_0 + 16, work_reg_b);
-
-              scratch_data_0 += depth_advance;
-
-              work_reg_a_sp = vzip2q_s8(input_data_a, input_data_b);
-              work_reg_b_sp = vzip2q_s8(input_data_c, input_data_d);
-              vzipq_s8x2_in_place(&work_reg_a_sp, &work_reg_b_sp);
-              if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-                work_reg_a_sp = veorq_s8(work_reg_a_sp, sign_bit);
-                work_reg_b_sp = veorq_s8(work_reg_b_sp, sign_bit);
-              }
-
-              vst1q_s8(scratch_data_0, work_reg_a_sp);
-              vst1q_s8(scratch_data_0 + 16, work_reg_b_sp);
-
-              scratch_data_0 += depth_advance;
-            }
-            for (; i_depth < depth_micro_repeats; ++i_depth) {
-              input_data_a = vld1q_lane_s8x8(input_data_0, input_data_a, 0);
-              input_data_b = vld1q_lane_s8x8(input_data_0 + 1 * input_depth,
-                                             input_data_b, 0);
-              input_data_c = vld1q_lane_s8x8(input_data_0 + 2 * input_depth,
-                                             input_data_c, 0);
-              input_data_d = vld1q_lane_s8x8(input_data_0 + 3 * input_depth,
-                                             input_data_d, 0);
-              work_reg_a = vzip1q_s8(input_data_a, input_data_b);
-              work_reg_b = vzip1q_s8(input_data_c, input_data_d);
-
-              input_data_0 += 8;
-
-              vzipq_s8x2_in_place(&work_reg_a, &work_reg_b);
-              if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-                work_reg_a = veorq_s8(work_reg_a, sign_bit);
-                work_reg_b = veorq_s8(work_reg_b, sign_bit);
-              }
-
-              vst1q_s8(scratch_data_0, work_reg_a);
-              vst1q_s8(scratch_data_0 + 16, work_reg_b);
-
-              scratch_data_0 += depth_advance;
-            }
-            scratch_data_0 += width_advance;
-            input_data_0 += input_depth_skip;
-          } else {
-            TFLITE_DCHECK_LT(adjusted_residual_width, 4);
-            for (int i_depth = 0; i_depth < depth_micro_repeats; ++i_depth) {
-              input_data_a = vreinterpretq_s8_u8(vdupq_n_u8(-input_offset));
-              input_data_b = vreinterpretq_s8_u8(vdupq_n_u8(-input_offset));
-              input_data_c = vreinterpretq_s8_u8(vdupq_n_u8(-input_offset));
-              input_data_d = vreinterpretq_s8_u8(vdupq_n_u8(-input_offset));
-              if (adjusted_residual_width > 0) {
-                input_data_a = vld1q_lane_s8x8(input_data_0, input_data_a, 0);
-                if (adjusted_residual_width > 1) {
-                  input_data_b = vld1q_lane_s8x8(input_data_0 + input_depth,
-                                                 input_data_b, 0);
-                  if (adjusted_residual_width == 3) {
-                    input_data_c = vld1q_lane_s8x8(
-                        input_data_0 + 2 * input_depth, input_data_c, 0);
-                  }
-                }
-              }
-              work_reg_a = vzip1q_s8(input_data_a, input_data_b);
-              work_reg_b = vzip1q_s8(input_data_c, input_data_d);
-
-              if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-                work_reg_a = veorq_s8(work_reg_a, sign_bit);
-                work_reg_b = veorq_s8(work_reg_b, sign_bit);
-              }
-              vzipq_s8x2_in_place(&work_reg_a, &work_reg_b);
-
-              vst1q_s8(scratch_data_0, work_reg_a);
-              vst1q_s8(scratch_data_0 + 16, work_reg_b);
-
-              scratch_data_0 += depth_advance;
-              input_data_0 += 8;
-            }
-            scratch_data_0 += width_advance;
-            input_data_0 += input_depth_skip;
-          }
-        } else {
-          if (adjusted_residual_width == 4) {
-            int8x16_t work_reg_a_sp;
-            int8x16_t work_reg_b_sp;
-
-            int i_depth = 0;
-
-            if (depth_micro_repeats >= 2) {
-              i_depth += 2;
-
-              input_data_a = vreinterpretq_s8_u8(vdupq_n_u8(-input_offset));
-              input_data_b = util_vld1q_x8(input_data_0 + 1 * input_depth);
-              input_data_c = util_vld1q_x8(input_data_0 + 2 * input_depth);
-              input_data_d = util_vld1q_x8(input_data_0 + 3 * input_depth);
-              input_data_0 += 16;
-
-              for (; i_depth < depth_micro_repeats - 1; i_depth += 2) {
-                work_reg_a = vzip1q_s8(input_data_a, input_data_b);
-                work_reg_b = vzip1q_s8(input_data_c, input_data_d);
-                vzipq_s8x2_in_place(&work_reg_a, &work_reg_b);
-                if (quantization_type ==
-                    QuantizationType::kNonPerChannelUint8) {
-                  work_reg_a = veorq_s8(work_reg_a, sign_bit);
-                  work_reg_b = veorq_s8(work_reg_b, sign_bit);
-                }
-
-                work_reg_a_sp = vzip2q_s8(input_data_a, input_data_b);
-                work_reg_b_sp = vzip2q_s8(input_data_c, input_data_d);
-                vzipq_s8x2_in_place(&work_reg_a_sp, &work_reg_b_sp);
-
-                input_data_a = vreinterpretq_s8_u8(vdupq_n_u8(-input_offset));
-                input_data_b = util_vld1q_x8(input_data_0 + 1 * input_depth);
-                vst1q_s8(scratch_data_0, work_reg_a);
-                vst1q_s8(scratch_data_0 + 16, work_reg_b);
-
-                scratch_data_0 += depth_advance;
-
-                if (quantization_type ==
-                    QuantizationType::kNonPerChannelUint8) {
-                  work_reg_a_sp = veorq_s8(work_reg_a_sp, sign_bit);
-                  work_reg_b_sp = veorq_s8(work_reg_b_sp, sign_bit);
-                }
-
-                input_data_c = util_vld1q_x8(input_data_0 + 2 * input_depth);
-                input_data_d = util_vld1q_x8(input_data_0 + 3 * input_depth);
-                vst1q_s8(scratch_data_0, work_reg_a_sp);
-                vst1q_s8(scratch_data_0 + 16, work_reg_b_sp);
-
-                scratch_data_0 += depth_advance;
-                input_data_0 += 16;
-              }
-
-              work_reg_a = vzip1q_s8(input_data_a, input_data_b);
-              work_reg_b = vzip1q_s8(input_data_c, input_data_d);
-              vzipq_s8x2_in_place(&work_reg_a, &work_reg_b);
-              if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-                work_reg_a = veorq_s8(work_reg_a, sign_bit);
-                work_reg_b = veorq_s8(work_reg_b, sign_bit);
-              }
-              vst1q_s8(scratch_data_0, work_reg_a);
-              vst1q_s8(scratch_data_0 + 16, work_reg_b);
-
-              scratch_data_0 += depth_advance;
-
-              work_reg_a_sp = vzip2q_s8(input_data_a, input_data_b);
-              work_reg_b_sp = vzip2q_s8(input_data_c, input_data_d);
-              vzipq_s8x2_in_place(&work_reg_a_sp, &work_reg_b_sp);
-              if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-                work_reg_a_sp = veorq_s8(work_reg_a_sp, sign_bit);
-                work_reg_b_sp = veorq_s8(work_reg_b_sp, sign_bit);
-              }
-
-              vst1q_s8(scratch_data_0, work_reg_a_sp);
-              vst1q_s8(scratch_data_0 + 16, work_reg_b_sp);
-
-              scratch_data_0 += depth_advance;
-            }
-            for (; i_depth < depth_micro_repeats; ++i_depth) {
-              input_data_a = vreinterpretq_s8_u8(vdupq_n_u8(-input_offset));
-              input_data_b = vld1q_lane_s8x8(input_data_0 + 1 * input_depth,
-                                             input_data_b, 0);
-              input_data_c = vld1q_lane_s8x8(input_data_0 + 2 * input_depth,
-                                             input_data_c, 0);
-              input_data_d = vld1q_lane_s8x8(input_data_0 + 3 * input_depth,
-                                             input_data_d, 0);
-              work_reg_a = vzip1q_s8(input_data_a, input_data_b);
-              work_reg_b = vzip1q_s8(input_data_c, input_data_d);
-
-              input_data_0 += 8;
-
-              vzipq_s8x2_in_place(&work_reg_a, &work_reg_b);
-              if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-                work_reg_a = veorq_s8(work_reg_a, sign_bit);
-                work_reg_b = veorq_s8(work_reg_b, sign_bit);
-              }
-
-              vst1q_s8(scratch_data_0, work_reg_a);
-              vst1q_s8(scratch_data_0 + 16, work_reg_b);
-
-              scratch_data_0 += depth_advance;
-            }
-            scratch_data_0 += width_advance;
-            input_data_0 += input_depth_skip;
-          } else {
-            TFLITE_DCHECK_LT(adjusted_residual_width, 4);
-
-            for (int i_depth = 0; i_depth < depth_micro_repeats; ++i_depth) {
-              input_data_a = vreinterpretq_s8_u8(vdupq_n_u8(-input_offset));
-              input_data_b = vreinterpretq_s8_u8(vdupq_n_u8(-input_offset));
-              input_data_c = vreinterpretq_s8_u8(vdupq_n_u8(-input_offset));
-              input_data_d = vreinterpretq_s8_u8(vdupq_n_u8(-input_offset));
-              // Skip loading first column.
-              if (adjusted_residual_width > 1) {
-                input_data_b = vld1q_lane_s8x8(input_data_0 + input_depth,
-                                               input_data_b, 0);
-                if (adjusted_residual_width == 3) {
-                  input_data_c = vld1q_lane_s8x8(input_data_0 + 2 * input_depth,
-                                                 input_data_c, 0);
-                }
-              }
-              work_reg_a = vzip1q_s8(input_data_a, input_data_b);
-              work_reg_b = vzip1q_s8(input_data_c, input_data_d);
-
-              if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-                work_reg_a = veorq_s8(work_reg_a, sign_bit);
-                work_reg_b = veorq_s8(work_reg_b, sign_bit);
-              }
-              vzipq_s8x2_in_place(&work_reg_a, &work_reg_b);
-
-              vst1q_s8(scratch_data_0, work_reg_a);
-              vst1q_s8(scratch_data_0 + 16, work_reg_b);
-
-              scratch_data_0 += depth_advance;
-              input_data_0 += 8;
-            }
-            scratch_data_0 += width_advance;
-            input_data_0 += input_depth_skip;
-          }
-        }
-      }
-      scratch_data_0 += height_advance;
-      input_block_data += input_height_stride;
-    }
-
-    if (trailing_height_padding) {
-      memset(scratch_data_0, -input_offset_difference, workspace_height_stride);
-      scratch_data_0 += workspace_height_stride;
-    }
-
-    TFLITE_DCHECK_EQ(
-        scratch_data_0,
-        scratch_block_data + block_height * workspace_height_stride);
-  }
-
-  static inline void Run(
-      int32_t height_block_number, int32_t width_block_number,
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          input_block_data,
-      int8_t* scratch_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-#ifdef __aarch64__
-    PreloadInputBlock(input_block_data, function_params);
-#endif
-
-    PackMacroBlockIntrinsics(height_block_number, width_block_number,
-                             input_block_data, scratch_block_data,
-                             function_params);
-  }
-};
-
-template <QuantizationType quantization_type>
-struct PackMacroBlock<DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct,
-                      quantization_type,
-                      DepthwiseConvDepthMultiplication::kUnitInputDepth,
-                      /*max_padding=*/1> {
-  static inline void PackMacroBlockIntrinsics(
-      int32_t height_block_number, int32_t width_block_number,
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          input_block_data,
-      int8_t* scratch_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-    const int width_overall_micro_repeats =
-        function_params->input_width_overall_micro_repeats;
-    const int input_width_micro_repeats =
-        function_params->input_width_micro_repeats;
-    const int block_height = function_params->inbound_block_height;
-    const int residual_width = function_params->residual_width;
-    const int input_height_stride = function_params->input_height_stride;
-
-    const int padding_left = function_params->padding_left;
-    const int padding_right = function_params->padding_right;
-    const int padding_top = function_params->padding_top;
-    const int padding_bottom = function_params->padding_bottom;
-
-    constexpr int kSymmetricZeroPoint =
-        QuantizationTypeImpl<quantization_type>::kIntSymmetricZeroPoint;
-
-    TFLITE_DCHECK_GE(workspace_height_stride, 4 * width_overall_micro_repeats);
-
-    const bool leading_width_padding =
-        padding_left > 0 && width_block_number == 0;
-    const bool trailing_width_padding =
-        padding_right > 0 &&
-        width_block_number == (function_params->width_macro_count - 1);
-    const bool leading_height_padding =
-        padding_top > 0 && height_block_number < 0;
-    const bool trailing_height_padding =
-        padding_bottom > 0 &&
-        height_block_number == (function_params->height_macro_count - 1);
-
-    const int32_t input_offset = function_params->input_offset;
-    const int32_t input_offset_difference = input_offset + kSymmetricZeroPoint;
-
-    // Work through one slice, by row, at a time.
-    int8_t* scratch_data_base = scratch_block_data;
-
-    int copy_block_height = block_height;
-    if (leading_height_padding) {
-      copy_block_height -= 1;
-      memset(scratch_data_base, -input_offset_difference,
-             workspace_height_stride + kWorkspaceExtension);
-      scratch_data_base += workspace_height_stride;
-      input_block_data += input_height_stride;
-    }
-    if (trailing_height_padding) {
-      copy_block_height -= 1;
-    }
-
-    int adjusted_residual_width =
-        input_width_micro_repeats < width_overall_micro_repeats ? residual_width
-                                                                : 4;
-
-    if (trailing_width_padding) {
-      adjusted_residual_width -= 1;
-    }
-    int start_width = 0;
-    if (leading_width_padding) {
-      start_width = 1;
-      input_block_data += 1;
-    }
-
-    const int copy_size = (width_overall_micro_repeats - 1) * 4 +
-                          adjusted_residual_width - start_width;
-    // Adjusted so that later conditionals are simplified.
-    const int copy_size_adjusted =
-        trailing_width_padding ? copy_size + 1 : copy_size;
-
-    TFLITE_DCHECK_LE(
-        copy_size,
-        input_height_stride - width_block_number * input_width_micro_repeats);
-    // We may drop up to stride-1 of trailing input.
-    TFLITE_DCHECK_GE(copy_size, input_height_stride - 1);
-
-    int scratch_data_offset = 0;
-    int input_block_offset = 0;
-
-    constexpr uint8_t kSignBit =
-        QuantizationTypeImpl<quantization_type>::kUint8SignBit;
-
-    // Transpositions are 4x4, but doing 2 at a time is more efficient in NEON
-    // code. Note the blocks of 4x4 are still interleaved down the depth.
-    int8x16_t work_reg;
-    int8x8_t half_work_reg;
-    int8x8_t padding_mask;
-
-    // Effect subtraction of zero-point = 128 by XOR of sign bit.
-    const int8x16_t sign_bit = vreinterpretq_s8_u8(vdupq_n_u8(kSignBit));
-    const int8x16_t padding_reg =
-        vreinterpretq_s8_u8(vdupq_n_u8(-input_offset));
-    padding_mask = vdup_n_s8(-1);
-    half_work_reg = vdup_n_s8(0);
-
-    if (copy_size >= 16) {
-      const int copy_remaining = (copy_size + start_width) & 0x7;
-      padding_mask = vreinterpret_s8_s64(vshl_s64(
-          vreinterpret_s64_s8(padding_mask), vdup_n_s64(8 * copy_remaining)));
-
-      for (int k_height = 0; k_height < copy_block_height; ++k_height) {
-        // Work through one slice, by row, at a time.
-        int8_t* scratch_data = scratch_data_base + scratch_data_offset;
-
-        int copy_done = 0;
-
-        // The surrounding condition ensures that we always need at least one
-        // iteration of the main copy loop. In the case of leading width
-        // padding, we unroll this specially.
-        if (leading_width_padding) {
-          work_reg = util_vld1q_x8(input_block_data + input_block_offset);
-          work_reg = vextq_s8(padding_reg, work_reg, 15);
-          if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-            work_reg = veorq_s8(work_reg, sign_bit);
-          }
-          vst1q_s8(scratch_data, work_reg);
-          copy_done += 15;
-        }
-
-        // Main copy loop.
-        for (; (copy_done + 16) <= copy_size; copy_done += 16) {
-          work_reg =
-              util_vld1q_x8(input_block_data + input_block_offset + copy_done);
-          if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-            work_reg = veorq_s8(work_reg, sign_bit);
-          }
-          TFLITE_DCHECK_EQ((start_width + copy_done) % 16, 0);
-          vst1q_s8(scratch_data + start_width + copy_done, work_reg);
-        }
-
-        if (copy_done + 8 <= copy_size) {
-          half_work_reg =
-              util_vld1_x8(input_block_data + input_block_offset + copy_done);
-          if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-            half_work_reg = veor_s8(half_work_reg, vget_low_s8(sign_bit));
-          }
-          TFLITE_DCHECK_EQ((start_width + copy_done) % 8, 0);
-          vst1_s8(scratch_data + start_width + copy_done, half_work_reg);
-          copy_done += 8;
-        }
-
-        TFLITE_DCHECK_EQ(copy_remaining, copy_size - copy_done);
-        // Total amount
-        // = copy_size - copy_done + 4 - adjusted_residual_width
-        // = width_overall_micro_repeats * 4 - start_width - copy_done.
-        // Undone micro blocks
-        // = width_overall_micro_repeats - (start_width + copy_done) / 4.
-
-        // Conditional is (copy_remaining > 0 || trailing_width_padding).
-        if (copy_done < copy_size_adjusted) {
-          // Employ overlapping-load strategy in order to load full register,
-          // but use only part.
-          // This has the advantage of resulting in zeros after shifting.
-          half_work_reg = util_vld1_x8(input_block_data + input_block_offset +
-                                       copy_size - 8);
-
-          half_work_reg = vreinterpret_s8_s64(
-              vshl_s64(vreinterpret_s64_s8(half_work_reg),
-                       vdup_n_s64(-8 * (8 - copy_remaining))));
-          half_work_reg = vbsl_s8(vreinterpret_u8_s8(padding_mask),
-                                  vget_low_s8(padding_reg), half_work_reg);
-
-          if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-            half_work_reg = veor_s8(half_work_reg, vget_low_s8(sign_bit));
-          }
-          TFLITE_DCHECK_EQ((start_width + copy_done) % 8, 0);
-          vst1_s8(scratch_data + start_width + copy_done, half_work_reg);
-        }
-
-        // Trailing guard.
-        vst1_s8(scratch_data + start_width + copy_done, half_work_reg);
-        vst1_s8(scratch_data + start_width + copy_done + 8, half_work_reg);
-
-        scratch_data_offset += workspace_height_stride;
-        input_block_offset += input_height_stride;
-      }
-    } else if (copy_size >= 4) {
-      const int copy_remaining = (copy_size + start_width) & 0x3;
-      padding_mask = vreinterpret_s8_s64(vshl_s64(
-          vreinterpret_s64_s8(padding_mask), vdup_n_s64(8 * copy_remaining)));
-
-      for (int k_height = 0; k_height < copy_block_height; ++k_height) {
-        // Work through one slice, by row, at a time.
-        int8_t* scratch_data = scratch_data_base + scratch_data_offset;
-
-        int copy_done = 0;
-
-        // The surrounding condition ensures that we always need at least one
-        // iteration of the main copy loop. In the case of leading width
-        // padding, we unroll this specially.
-        if (leading_width_padding) {
-          half_work_reg = vld1_lane_8x4(input_block_data + input_block_offset,
-                                        half_work_reg, 0);
-          half_work_reg = vext_s8(vget_low_s8(padding_reg), half_work_reg, 7);
-          if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-            half_work_reg = veor_s8(half_work_reg, vget_low_s8(sign_bit));
-          }
-          vst1_lane_s8x4(scratch_data, half_work_reg, 0);
-          copy_done += 3;
-        }
-
-        // Main copy loop.
-        for (; (copy_done + 4) <= copy_size; copy_done += 4) {
-          half_work_reg =
-              vld1_lane_8x4(input_block_data + input_block_offset + copy_done,
-                            half_work_reg, 0);
-          if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-            half_work_reg = veor_s8(half_work_reg, vget_low_s8(sign_bit));
-          }
-          TFLITE_DCHECK_EQ((start_width + copy_done) % 4, 0);
-          vst1_lane_s8x4(scratch_data + start_width + copy_done, half_work_reg,
-                         0);
-        }
-
-        TFLITE_DCHECK_EQ(copy_remaining, copy_size - copy_done);
-        // Total amount
-        // = copy_size - copy_done + 4 - adjusted_residual_width
-        // = width_overall_micro_repeats * 4 - start_width - copy_done.
-        // Undone micro blocks
-        // = width_overall_micro_repeats - (start_width + copy_done) / 4.
-
-        // Conditional is (copy_remaining > 0 || trailing_width_padding).
-        if (copy_done < copy_size_adjusted) {
-          TFLITE_DCHECK_LT(copy_remaining, 4);
-          // Employ overlapping-load strategy in order to load full register,
-          // but use only part.
-          // This has the advantage of resulting in zeros after shifting.
-          half_work_reg = vld1_lane_8x4(
-              input_block_data + input_block_offset + copy_size - 4,
-              half_work_reg, 0);
-
-          half_work_reg = vreinterpret_s8_s64(
-              vshl_s64(vreinterpret_s64_s8(half_work_reg),
-                       vdup_n_s64(-8 * (4 - copy_remaining))));
-          half_work_reg = vbsl_s8(vreinterpret_u8_s8(padding_mask),
-                                  vget_low_s8(padding_reg), half_work_reg);
-
-          if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-            half_work_reg = veor_s8(half_work_reg, vget_low_s8(sign_bit));
-          }
-          TFLITE_DCHECK_EQ((start_width + copy_done) % 4, 0);
-          vst1_lane_s8x4(scratch_data + start_width + copy_done, half_work_reg,
-                         0);
-          copy_done += 4;
-        }
-        // Trailing guard.
-        vst1_lane_s8x4(scratch_data + start_width + copy_done, half_work_reg,
-                       0);
-        vst1_lane_s8x4(scratch_data + start_width + copy_done + 4,
-                       half_work_reg, 0);
-        vst1_lane_s8x4(scratch_data + start_width + copy_done + 8,
-                       half_work_reg, 0);
-        vst1_lane_s8x4(scratch_data + start_width + copy_done + 12,
-                       half_work_reg, 0);
-
-        scratch_data_offset += workspace_height_stride;
-        input_block_offset += input_height_stride;
-      }
-    } else if (width_overall_micro_repeats == 2) {
-      // Special case of 1 + 3 + 1, padding + copy + padding.
-      // This is rarely executed in practice.
-      TFLITE_DCHECK_EQ(copy_size, 3);
-      TFLITE_DCHECK_EQ(start_width, 1);
-      TFLITE_DCHECK(leading_width_padding);
-      TFLITE_DCHECK(trailing_width_padding);
-
-      for (int k_height = 0; k_height < copy_block_height; ++k_height) {
-        half_work_reg = vreinterpret_s8_u8(vdup_n_u8(-input_offset));
-        half_work_reg = vld1_lane_s8(reinterpret_cast<const int8_t*>(
-                                         input_block_data + input_block_offset),
-                                     half_work_reg, 1);
-        half_work_reg =
-            vld1_lane_s8(reinterpret_cast<const int8_t*>(
-                             input_block_data + input_block_offset + 1),
-                         half_work_reg, 2);
-        half_work_reg =
-            vld1_lane_s8(reinterpret_cast<const int8_t*>(
-                             input_block_data + input_block_offset + 2),
-                         half_work_reg, 3);
-
-        if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-          half_work_reg = veor_s8(half_work_reg, vget_low_s8(sign_bit));
-        }
-        TFLITE_DCHECK_EQ(scratch_data_offset % 8, 0);
-        vst1_s8(scratch_data_base + scratch_data_offset, half_work_reg);
-
-        // Trailing guard.
-        vst1_lane_s8x4(scratch_data_base + scratch_data_offset + 4,
-                       half_work_reg, 0);
-        vst1_lane_s8x4(scratch_data_base + scratch_data_offset + 8,
-                       half_work_reg, 0);
-        vst1_lane_s8x4(scratch_data_base + scratch_data_offset + 12,
-                       half_work_reg, 0);
-        vst1_lane_s8x4(scratch_data_base + scratch_data_offset + 16,
-                       half_work_reg, 0);
-
-        scratch_data_offset += workspace_height_stride;
-        input_block_offset += input_height_stride;
-      }
-    } else {
-      TFLITE_DCHECK_EQ(width_overall_micro_repeats, 1);
-      const int copy_remaining = (copy_size + start_width) & 0x3;
-      padding_mask = vreinterpret_s8_s64(vshl_s64(
-          vreinterpret_s64_s8(padding_mask), vdup_n_s64(8 * copy_remaining)));
-      if (leading_width_padding) {
-        padding_mask = vreinterpret_s8_u8(
-            vset_lane_u8(255, vreinterpret_u8_s8(padding_mask), 0));
-      }
-
-      for (int k_height = 0; k_height < copy_block_height; ++k_height) {
-        for (int i = 0; i < copy_size; ++i) {
-          half_work_reg = vreinterpret_s8_s64(
-              vshl_n_s64(vreinterpret_s64_s8(half_work_reg), 8));
-          half_work_reg = vld1_lane_s8(
-              reinterpret_cast<const int8_t*>(
-                  input_block_data + input_block_offset + copy_size - 1 - i),
-              half_work_reg, 0);
-        }
-        if (leading_width_padding) {
-          half_work_reg = vreinterpret_s8_s64(
-              vshl_n_s64(vreinterpret_s64_s8(half_work_reg), 8));
-        }
-        half_work_reg = vbsl_s8(vreinterpret_u8_s8(padding_mask),
-                                vget_low_s8(padding_reg), half_work_reg);
-
-        if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-          half_work_reg = veor_s8(half_work_reg, vget_low_s8(sign_bit));
-        }
-        TFLITE_DCHECK_EQ(scratch_data_offset % 4, 0);
-        vst1_lane_s8x4(scratch_data_base + scratch_data_offset, half_work_reg,
-                       0);
-
-        // Trailing guard.
-        vst1_lane_s8x4(scratch_data_base + scratch_data_offset + 4,
-                       half_work_reg, 0);
-        vst1_lane_s8x4(scratch_data_base + scratch_data_offset + 8,
-                       half_work_reg, 0);
-        vst1_lane_s8x4(scratch_data_base + scratch_data_offset + 12,
-                       half_work_reg, 0);
-        vst1_lane_s8x4(scratch_data_base + scratch_data_offset + 16,
-                       half_work_reg, 0);
-
-        scratch_data_offset += workspace_height_stride;
-        input_block_offset += input_height_stride;
-      }
-    }
-
-    scratch_data_base += copy_block_height * workspace_height_stride;
-
-    if (trailing_height_padding) {
-      memset(scratch_data_base, -input_offset_difference,
-             workspace_height_stride + kWorkspaceExtension);
-      scratch_data_base += workspace_height_stride;
-    }
-
-    TFLITE_DCHECK_EQ(
-        scratch_data_base,
-        scratch_block_data + block_height * workspace_height_stride);
-  }
-
-  static inline void Run(
-      int32_t height_block_number, int32_t width_block_number,
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          input_block_data,
-      int8_t* scratch_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-#ifdef __aarch64__
-    PreloadInputBlock(input_block_data, function_params);
-#endif
-
-    PackMacroBlockIntrinsics(height_block_number, width_block_number,
-                             input_block_data, scratch_block_data,
-                             function_params);
-  }
-};
-
-template <QuantizationType quantization_type>
-struct PackMacroBlock<DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct,
-                      quantization_type,
-                      DepthwiseConvDepthMultiplication::kUnitInputDepth,
-                      /*max_padding=*/0> {
-  static inline void PackMacroBlockIntrinsics(
-      int32_t height_block_number, int32_t width_block_number,
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          input_block_data,
-      int8_t* scratch_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-    const int width_overall_micro_repeats =
-        function_params->input_width_overall_micro_repeats;
-    const int input_width_micro_repeats =
-        function_params->input_width_micro_repeats;
-    const int block_height = function_params->inbound_block_height;
-    const int residual_width = function_params->residual_width;
-    const int input_height_stride = function_params->input_height_stride;
-
-    TFLITE_DCHECK_EQ(function_params->padding_left, 0);
-    TFLITE_DCHECK_EQ(function_params->padding_right, 0);
-    TFLITE_DCHECK_EQ(function_params->padding_top, 0);
-    TFLITE_DCHECK_EQ(function_params->padding_bottom, 0);
-
-    TFLITE_DCHECK_GE(workspace_height_stride, 4 * width_overall_micro_repeats);
-
-    // Work through one slice, by row, at a time.
-    int8_t* scratch_data_base = scratch_block_data;
-
-    const int copy_block_height = block_height;
-
-    int adjusted_residual_width =
-        input_width_micro_repeats < width_overall_micro_repeats ? residual_width
-                                                                : 4;
-
-    const int copy_size =
-        (width_overall_micro_repeats - 1) * 4 + adjusted_residual_width;
-
-    TFLITE_DCHECK_LE(
-        copy_size,
-        input_height_stride - width_block_number * input_width_micro_repeats);
-    // We may drop up to stride-1 of trailing input.
-    TFLITE_DCHECK_GE(copy_size, input_height_stride - 1);
-
-    int scratch_data_offset = 0;
-    int input_block_offset = 0;
-
-    constexpr uint8_t kSignBit =
-        QuantizationTypeImpl<quantization_type>::kUint8SignBit;
-
-    // Transpositions are 4x4, but doing 2 at a time is more efficient in NEON
-    // code. Note the blocks of 4x4 are still interleaved down the depth.
-    int8x16_t work_reg;
-    int8x8_t half_work_reg;
-
-    // Effect subtraction of zero-point = 128 by XOR of sign bit.
-    const int8x16_t sign_bit = vreinterpretq_s8_u8(vdupq_n_u8(kSignBit));
-    half_work_reg = vdup_n_s8(0);
-
-    if (copy_size >= 16) {
-      const int copy_remaining = copy_size & 0x7;
-
-      for (int k_height = 0; k_height < copy_block_height; ++k_height) {
-        // Work through one slice, by row, at a time.
-        int8_t* scratch_data = scratch_data_base + scratch_data_offset;
-
-        int copy_done = 0;
-
-        // Main copy loop.
-        for (; (copy_done + 16) <= copy_size; copy_done += 16) {
-          work_reg =
-              util_vld1q_x8(input_block_data + input_block_offset + copy_done);
-          if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-            work_reg = veorq_s8(work_reg, sign_bit);
-          }
-          TFLITE_DCHECK_EQ(copy_done % 16, 0);
-          vst1q_s8(scratch_data + copy_done, work_reg);
-        }
-
-        if (copy_done + 8 <= copy_size) {
-          half_work_reg =
-              util_vld1_x8(input_block_data + input_block_offset + copy_done);
-          if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-            half_work_reg = veor_s8(half_work_reg, vget_low_s8(sign_bit));
-          }
-          TFLITE_DCHECK_EQ(copy_done % 8, 0);
-          vst1_s8(scratch_data + copy_done, half_work_reg);
-          copy_done += 8;
-        }
-
-        TFLITE_DCHECK_EQ(copy_remaining, copy_size - copy_done);
-        // Total amount
-        // = copy_size - copy_done + 4 - adjusted_residual_width
-        // = width_overall_micro_repeats * 4 - start_width - copy_done.
-        // Undone micro blocks
-        // = width_overall_micro_repeats - (start_width + copy_done) / 4.
-
-        // Conditional is (copy_remaining > 0 || trailing_width_padding).
-        if (copy_done < copy_size) {
-          // Employ overlapping-load strategy in order to load full register,
-          // but use only part.
-          // This has the advantage of resulting in zeros after shifting.
-          half_work_reg = util_vld1_x8(input_block_data + input_block_offset +
-                                       copy_size - 8);
-
-          half_work_reg = vreinterpret_s8_s64(
-              vshl_s64(vreinterpret_s64_s8(half_work_reg),
-                       vdup_n_s64(-8 * (8 - copy_remaining))));
-
-          if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-            half_work_reg = veor_s8(half_work_reg, vget_low_s8(sign_bit));
-          }
-          TFLITE_DCHECK_EQ(copy_done % 8, 0);
-          vst1_s8(scratch_data + copy_done, half_work_reg);
-          copy_done += 8;
-        }
-
-        // Trailing guard.
-        vst1_s8(scratch_data + copy_done, half_work_reg);
-        vst1_s8(scratch_data + copy_done + 8, half_work_reg);
-
-        scratch_data_offset += workspace_height_stride;
-        input_block_offset += input_height_stride;
-      }
-    } else if (copy_size >= 4) {
-      const int copy_remaining = copy_size & 0x3;
-
-      for (int k_height = 0; k_height < copy_block_height; ++k_height) {
-        // Work through one slice, by row, at a time.
-        int8_t* scratch_data = scratch_data_base + scratch_data_offset;
-
-        int copy_done = 0;
-
-        // Main copy loop.
-        for (; (copy_done + 4) <= copy_size; copy_done += 4) {
-          half_work_reg =
-              vld1_lane_8x4(input_block_data + input_block_offset + copy_done,
-                            half_work_reg, 0);
-          if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-            half_work_reg = veor_s8(half_work_reg, vget_low_s8(sign_bit));
-          }
-          TFLITE_DCHECK_EQ(copy_done % 4, 0);
-          vst1_lane_s8x4(scratch_data + copy_done, half_work_reg, 0);
-        }
-
-        TFLITE_DCHECK_EQ(copy_remaining, copy_size - copy_done);
-        // Total amount
-        // = copy_size - copy_done + 4 - adjusted_residual_width
-        // = width_overall_micro_repeats * 4 - start_width - copy_done.
-        // Undone micro blocks
-        // = width_overall_micro_repeats - (start_width + copy_done) / 4.
-
-        // Conditional is (copy_remaining > 0 || trailing_width_padding).
-        if (copy_done < copy_size) {
-          TFLITE_DCHECK_LT(copy_remaining, 4);
-          // Employ overlapping-load strategy in order to load full register,
-          // but use only part.
-          // This has the advantage of resulting in zeros after shifting.
-          half_work_reg = vld1_lane_8x4(
-              input_block_data + input_block_offset + copy_size - 4,
-              half_work_reg, 0);
-
-          half_work_reg = vreinterpret_s8_s64(
-              vshl_s64(vreinterpret_s64_s8(half_work_reg),
-                       vdup_n_s64(-8 * (4 - copy_remaining))));
-
-          if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-            half_work_reg = veor_s8(half_work_reg, vget_low_s8(sign_bit));
-          }
-          TFLITE_DCHECK_EQ(copy_done % 4, 0);
-          vst1_lane_s8x4(scratch_data + copy_done, half_work_reg, 0);
-          copy_done += 4;
-        }
-        // Trailing guard.
-        vst1_lane_s8x4(scratch_data + copy_done, half_work_reg, 0);
-        vst1_lane_s8x4(scratch_data + copy_done + 4, half_work_reg, 0);
-        vst1_lane_s8x4(scratch_data + copy_done + 8, half_work_reg, 0);
-        vst1_lane_s8x4(scratch_data + copy_done + 12, half_work_reg, 0);
-
-        scratch_data_offset += workspace_height_stride;
-        input_block_offset += input_height_stride;
-      }
-    } else {
-      TFLITE_DCHECK_EQ(width_overall_micro_repeats, 1);
-
-      for (int k_height = 0; k_height < copy_block_height; ++k_height) {
-        for (int i = 0; i < copy_size; ++i) {
-          half_work_reg = vreinterpret_s8_s64(
-              vshl_n_s64(vreinterpret_s64_s8(half_work_reg), 8));
-          half_work_reg = vld1_lane_s8(
-              reinterpret_cast<const int8_t*>(
-                  input_block_data + input_block_offset + copy_size - 1 - i),
-              half_work_reg, 0);
-        }
-
-        half_work_reg = veor_s8(half_work_reg, vget_low_s8(sign_bit));
-        TFLITE_DCHECK_EQ(scratch_data_offset % 4, 0);
-        vst1_lane_s8x4(scratch_data_base + scratch_data_offset, half_work_reg,
-                       0);
-
-        // Trailing guard.
-        vst1_lane_s8x4(scratch_data_base + scratch_data_offset + 4,
-                       half_work_reg, 0);
-        vst1_lane_s8x4(scratch_data_base + scratch_data_offset + 8,
-                       half_work_reg, 0);
-        vst1_lane_s8x4(scratch_data_base + scratch_data_offset + 12,
-                       half_work_reg, 0);
-        vst1_lane_s8x4(scratch_data_base + scratch_data_offset + 16,
-                       half_work_reg, 0);
-
-        scratch_data_offset += workspace_height_stride;
-        input_block_offset += input_height_stride;
-      }
-    }
-
-    scratch_data_base += copy_block_height * workspace_height_stride;
-
-    TFLITE_DCHECK_EQ(
-        scratch_data_base,
-        scratch_block_data + block_height * workspace_height_stride);
-  }
-
-  static inline void Run(
-      int32_t height_block_number, int32_t width_block_number,
-      const typename QuantizationTypeImpl<quantization_type>::ExternalType*
-          input_block_data,
-      int8_t* scratch_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-#ifdef __aarch64__
-    PreloadInputBlock(input_block_data, function_params);
-#endif
-
-    PackMacroBlockIntrinsics(height_block_number, width_block_number,
-                             input_block_data, scratch_block_data,
-                             function_params);
-  }
-};
-
-#endif  // ARM NEON
-
-// Apply filter to macro block of input data and store results.
-//
-// Requirement: depth_micro_repeats > 0 || residual_depth > 0.
-template <int32_t stride, QuantizationType quantization_type>
-struct KernelMacroBlock<
-    DepthwiseConvImplementation::kUseCModel3x3DotProduct, quantization_type,
-    DepthwiseConvDepthMultiplication::kNoMultiplication, stride> {
-  // Construct a width-shifted combination of two input sub-blocks, effectively
-  // concatenating them.
-  //
-  // The filter is applied using sub-blocks. These are in the needed form for
-  // the first (width) offset. For subsequent offsets, the filter is applied to
-  // shifted and combined data. The concatentation and shifting herein is fairly
-  // straightforward, but in the optimized code is an area of creativity in
-  // design because NEON instructions do not directly support the required
-  // between-register permutation.
-  //
-  // In NEON optimized code, input data is grouped in 4-byte blocks. In order to
-  // move along the width for each output point calculation, data is shifted, in
-  // essence between two such blocks.
-  //
-  // selected_data has format height 3, depth 4, width 4.
-  //
-  // When the micro block is trailing (the last across the macro-block width),
-  // it would be illegal to load the right (next) block, and the no_right_block
-  // indicates this scenario.
-  static inline void ConcatenateInputSubBlocks(int offset, int sub_block,
-                                               int workspace_height_stride,
-                                               int width_micro_stride,
-                                               bool no_right_block,
-                                               const int8_t* input_block,
-                                               int8_t selected_data[3][4][4]) {
-    TFLITE_DCHECK_GE(offset, 0);
-    TFLITE_DCHECK_LT(offset, 4);
-
-    // The input banks have same format as selected_data.
-    int8_t left_bank[3][4][4];
-    int8_t right_bank[3][4][4];
-
-    // Work through one slice, by row, at a time.
-    for (int k_height = 0; k_height < 3; ++k_height) {
-      // Simulate demangling of mangled storage arrangement.
-      const int8_t* left_input_block =
-          &input_block[k_height * workspace_height_stride + sub_block * 2 * 8];
-      memcpy(left_bank[k_height][0], left_input_block, 16);
-      if (no_right_block) {
-        memset(right_bank[k_height][0], 0, 16);
-      } else {
-        const int8_t* right_input_block =
-            &input_block[k_height * workspace_height_stride +
-                         sub_block * 2 * 8 + width_micro_stride];
-        memcpy(right_bank[k_height][0], right_input_block, 16);
-      }
-      for (int depth_index = 0; depth_index < 4; ++depth_index) {
-        memcpy(selected_data[k_height][depth_index],
-               &left_bank[k_height][depth_index][offset], 4 - offset);
-        memcpy(&selected_data[k_height][depth_index][4 - offset],
-               right_bank[k_height][depth_index], offset);
-      }
-    }
-  }
-
-  // Straight implementation of 3x3 filter within sub-micro block.
-  static inline void Calculate3x3FilterOutput(
-      const DepthwiseConvDotProdParams& params, int sub_block,
-      const int8_t selected_data[3][4][4], const int8_t filter_bank[3][2][4][4],
-      const int32_t* bias_data, uint8_t output_values[4]) {
-    const int32_t output_activation_min = params.quantized_activation_min;
-    const int32_t output_activation_max = params.quantized_activation_max;
-    const int32_t output_multiplier = params.output_multiplier;
-    const int32_t output_shift = params.output_shift;
-    const int32_t output_offset = params.output_offset;
-    for (int d = 0; d < 4; ++d) {
-      int32_t acc = 0;
-      for (int y = 0; y < 3; ++y) {
-        for (int x = 0; x < 4; ++x) {
-          int32_t input_val = selected_data[y][d][x];
-          int32_t filter_val = filter_bank[y][sub_block][d][x];
-          acc += filter_val * input_val;
-        }
-      }
-      acc += bias_data[d];
-      acc = reference_ops::depthwise_conv::DepthwiseConvRound<
-          DepthwiseConvOutputRounding::kUpward>(acc, output_multiplier,
-                                                output_shift);
-      acc += output_offset;
-      acc = std::max(acc, output_activation_min);
-      acc = std::min(acc, output_activation_max);
-      output_values[d] = static_cast<uint8_t>(acc);
-    }
-  }
-
-  static inline void Run(const int8_t* scratch_block_data,
-                         const int8_t* filter_workspace,
-                         const int32_t* bias_data, uint8_t* output_block_data,
-                         const DepthwiseConvDotProdParams* function_params) {
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-    const int input_width_overall_micro_repeats =
-        function_params->input_width_overall_micro_repeats;
-    const int output_width_micro_repeats =
-        function_params->output_width_micro_repeats;
-    const int depth_micro_repeats = function_params->depth_micro_repeats;
-    const int depth = function_params->input_depth;
-    const int stride_val = function_params->stride;
-    const int four_over_stride = function_params->four_over_stride;
-
-    const int output_width_overall_micro_repeats =
-        function_params->output_width_overall_micro_repeats;
-    const int block_height = function_params->outbound_block_height;
-    const int residual_width = function_params->output_residual_width;
-    const int output_height_stride = function_params->output_height_stride;
-    constexpr int bias_increment = 4;
-    TFLITE_DCHECK_EQ(function_params->bias_increment, bias_increment);
-
-    TFLITE_DCHECK(depth_micro_repeats > 0);
-    const int width_micro_stride = 4 * 8;
-    const int depth_micro_stride =
-        width_micro_stride * input_width_overall_micro_repeats;
-
-    constexpr int shuffled_filter_increment = 2 * 3 * 4 * 4;
-
-    // Simulate NEON-register transposition of subset of filter.
-    int8_t filter_bank[3][2][4][4];  // Height 3, sub-block,  depth 4, width 4.
-    // Simulate NEON-register input data concatenation + sub-selection.
-    int8_t sub_selected_input_data[3][4][4];  // Height 3, depth 4, width 4.
-    uint8_t output_values[4];                 // Depth 4.
-
-    // The outer 3 loops go through all the micro blocks in a macro block, and
-    // separately treat the two sub-blocks within each micro block.
-    for (int j_depth = 0; j_depth < depth_micro_repeats; ++j_depth) {
-      memcpy(filter_bank[0][0][0],
-             filter_workspace + j_depth * shuffled_filter_increment,
-             shuffled_filter_increment);
-
-      for (int s = 0; s < 2; ++s) {
-        for (int k_height = 0; k_height < block_height; ++k_height) {
-          const int8_t* scratch_data =
-              scratch_block_data +
-              workspace_height_stride * k_height * stride_val +
-              depth_micro_stride * j_depth;
-          uint8_t* output_data =
-              output_block_data + output_height_stride * k_height + 8 * j_depth;
-
-          for (int i_width = 0; i_width < output_width_overall_micro_repeats;
-               ++i_width) {
-            const int output_width = i_width == output_width_micro_repeats
-                                         ? residual_width
-                                         : four_over_stride;
-            const bool no_right_block = (output_width - 1) * stride_val < 2;
-            TFLITE_DCHECK_LE(output_width * stride_val, 4);
-            const int8_t* input_data =
-                scratch_data + width_micro_stride * i_width;
-            // Iterate over input width shifts within sub-micro blocks.
-            for (int x = 0; x < output_width; ++x) {
-              ConcatenateInputSubBlocks(x * stride_val, s,
-                                        workspace_height_stride,
-                                        width_micro_stride, no_right_block,
-                                        input_data, sub_selected_input_data);
-              Calculate3x3FilterOutput(
-                  *function_params, s, sub_selected_input_data, filter_bank,
-                  bias_data + (2 * j_depth + s) * bias_increment,
-                  output_values);
-              for (int d = 0; d < 4; ++d) {
-                output_data[depth * (four_over_stride * i_width + x) + 4 * s +
-                            d] = output_values[d];
-              }
-            }
-          }
-        }
-      }
-    }
-  }
-};
-
-// Apply filter to macro block of input data and store results.
-//
-// Parameters for repeats and residual sizes are in terms of outputs.
-//
-// Requirement: depth_micro_repeats > 0 || residual_depth > 0.
-template <int32_t stride, QuantizationType quantization_type>
-struct KernelMacroBlock<
-    DepthwiseConvImplementation::kUseCModel3x3DotProduct, quantization_type,
-    DepthwiseConvDepthMultiplication::kUnitInputDepth, stride> {
-  // Construct a width-shifted combination of two input sub-blocks, effectively
-  // concatenating them.
-  //
-  // The filter is applied using sub-blocks. These are in the needed form for
-  // the first (width) offset. For subsequent offsets, the filter is applied to
-  // shifted and combined data. The concatentation and shifting herein is fairly
-  // straightforward, but in the optimized code is an area of creativity in
-  // design because NEON instructions do not directly support the required
-  // between-register permutation.
-  //
-  // In NEON optimized code, input data is grouped in 4-byte blocks. In order to
-  // move along the width for each output point calculation, data is shifted, in
-  // essence between two such blocks.
-  //
-  // selected_data has format height 3, width 4.
-  //
-  // When the micro block is trailing (the last across the macro-block width),
-  // it would be illegal to load the right (next) block, and the no_right_block
-  // indicates this scenario.
-  static inline void ConcatenateInputSubBlocks(int offset,
-                                               int workspace_height_stride,
-                                               bool no_right_block,
-                                               const int8_t* input_block,
-                                               int8_t selected_data[3][4]) {
-    TFLITE_DCHECK_GE(offset, 0);
-    TFLITE_DCHECK_LT(offset, 4);
-    if (no_right_block) {
-      for (int k_height = 0; k_height < 3; ++k_height) {
-        memcpy(selected_data[k_height],
-               &input_block[k_height * workspace_height_stride + offset],
-               4 - offset);
-      }
-    } else {
-      for (int k_height = 0; k_height < 3; ++k_height) {
-        memcpy(selected_data[k_height],
-               &input_block[k_height * workspace_height_stride + offset], 4);
-      }
-    }
-  }
-
-  // Straight implementation of 3x3 filter within sub-micro block.
-  static inline void Calculate3x3FilterOutput(
-      const DepthwiseConvDotProdParams& function_params, int sub_block,
-      const int8_t selected_data[3][4], const int8_t filter_bank[3][2][4][4],
-      const int32_t* bias_data, uint8_t output_values[4]) {
-    const int32_t output_activation_min =
-        function_params.quantized_activation_min;
-    const int32_t output_activation_max =
-        function_params.quantized_activation_max;
-    const int32_t output_multiplier = function_params.output_multiplier;
-    const int32_t output_shift = function_params.output_shift;
-    const int32_t output_offset = function_params.output_offset;
-    for (int d = 0; d < 4; ++d) {
-      int32_t acc = 0;
-      for (int y = 0; y < 3; ++y) {
-        for (int x = 0; x < 4; ++x) {
-          int32_t input_val = selected_data[y][x];
-          int32_t filter_val = filter_bank[y][sub_block][d][x];
-          acc += filter_val * input_val;
-        }
-      }
-      acc += bias_data[d];
-      acc = reference_ops::depthwise_conv::DepthwiseConvRound<
-          DepthwiseConvOutputRounding::kUpward>(acc, output_multiplier,
-                                                output_shift);
-      acc += output_offset;
-      acc = std::max(acc, output_activation_min);
-      acc = std::min(acc, output_activation_max);
-      output_values[d] = static_cast<uint8_t>(acc);
-    }
-  }
-
-  static inline void Run(const int8_t* scratch_block_data,
-                         const int8_t* filter_workspace,
-                         const int32_t* bias_data, uint8_t* output_block_data,
-                         const DepthwiseConvDotProdParams* function_params) {
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-    const int output_width_micro_repeats =
-        function_params->output_width_micro_repeats;
-    const int depth_micro_repeats = function_params->depth_micro_repeats;
-    const int depth = function_params->output_depth;
-    const int stride_val = function_params->stride;
-    const int four_over_stride = function_params->four_over_stride;
-
-    const int workspace_width_micro_repeats =
-        function_params->workspace_width_micro_repeats;
-    const int output_width_overall_micro_repeats =
-        function_params->output_width_overall_micro_repeats;
-    const int block_height = function_params->outbound_block_height;
-    const int residual_width = function_params->output_residual_width;
-    const int output_height_stride = function_params->output_height_stride;
-    constexpr int bias_increment = 4;
-    TFLITE_DCHECK_EQ(function_params->bias_increment, bias_increment);
-
-    TFLITE_DCHECK(depth_micro_repeats > 0);
-
-    constexpr int shuffled_filter_increment = 2 * 3 * 4 * 4;
-
-    // Simulate NEON-register transposition of subset of filter.
-    int8_t filter_bank[3][2][4][4];  // Height 3, sub-block,  depth 4, width 4.
-    // Simulate NEON-register input data concatenation + sub-selection.
-    int8_t sub_selected_input_data[3][4];  // Height 3, depth 4, width 4.
-    uint8_t output_values[4];              // Depth 4.
-
-    // The outer 3 loops go through all the micro blocks in a macro block, and
-    // separately treat the two sub-blocks within each micro block.
-    for (int j_depth = 0; j_depth < depth_micro_repeats; ++j_depth) {
-      memcpy(filter_bank[0][0][0],
-             filter_workspace + j_depth * shuffled_filter_increment,
-             shuffled_filter_increment);
-
-      for (int s = 0; s < 2; ++s) {
-        for (int k_height = 0; k_height < block_height; ++k_height) {
-          const int8_t* scratch_data =
-              scratch_block_data +
-              workspace_height_stride * k_height * stride_val;
-          uint8_t* output_data =
-              output_block_data + output_height_stride * k_height + 8 * j_depth;
-
-          for (int i_width = 0; i_width < output_width_overall_micro_repeats;
-               ++i_width) {
-            const int output_width = i_width == output_width_micro_repeats
-                                         ? residual_width
-                                         : four_over_stride;
-            const bool no_right_block = i_width == output_width_micro_repeats &&
-                                        output_width_overall_micro_repeats ==
-                                            workspace_width_micro_repeats;
-            TFLITE_DCHECK_LE(output_width * stride_val, 4);
-            const int8_t* input_data = scratch_data + 4 * i_width;
-            // Iterate over input width shifts within 4x4 blocks.
-            for (int x = 0; x < output_width; ++x) {
-              ConcatenateInputSubBlocks(x * stride_val, workspace_height_stride,
-                                        no_right_block, input_data,
-                                        sub_selected_input_data);
-              Calculate3x3FilterOutput(
-                  *function_params, s, sub_selected_input_data, filter_bank,
-                  bias_data + (2 * j_depth + s) * bias_increment,
-                  output_values);
-              for (int d = 0; d < 4; ++d) {
-                output_data[depth * (four_over_stride * i_width + x) + 4 * s +
-                            d] = output_values[d];
-              }
-            }
-          }
-        }
-      }
-    }
-  }
-};
-
-// Beginning of code section containing intermediate code transformation.
-//
-// This section is only compiled when kUseUnwound3x3DotProduct versions of
-// templated functions are selected.
-template <int32_t stride, QuantizationType quantization_type>
-struct KernelMacroBlock<
-    DepthwiseConvImplementation::kUseUnwound3x3DotProduct, quantization_type,
-    DepthwiseConvDepthMultiplication::kNoMultiplication, stride> {
-  static inline void Run(const int8_t* scratch_block_data,
-                         const int8_t* filter_workspace,
-                         const int32_t* bias_data, uint8_t* output_block_data,
-                         const DepthwiseConvDotProdParams* function_params) {
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-    const int input_width_overall_micro_repeats =
-        function_params->input_width_overall_micro_repeats;
-    const int output_width_micro_repeats =
-        function_params->output_width_micro_repeats;
-    const int depth_micro_repeats = function_params->depth_micro_repeats;
-    const int depth = function_params->input_depth;
-    const int stride_val = function_params->stride;
-    const int four_over_stride = function_params->four_over_stride;
-
-    const int output_width_overall_micro_repeats =
-        function_params->output_width_overall_micro_repeats;
-    const int block_height = function_params->outbound_block_height;
-    const int residual_width = function_params->output_residual_width;
-    const int output_height_stride = function_params->output_height_stride;
-    const int bias_increment = function_params->bias_increment;
-
-    TFLITE_DCHECK(depth_micro_repeats > 0);
-    const int width_micro_stride = 4 * 8;
-    const int depth_micro_stride =
-        width_micro_stride * input_width_overall_micro_repeats;
-
-    const int32_t output_activation_min =
-        function_params->quantized_activation_min;
-    const int32_t output_activation_max =
-        function_params->quantized_activation_max;
-    const int32_t output_multiplier = function_params->output_multiplier;
-    const int32_t output_shift = function_params->output_shift;
-    const int32_t output_offset = function_params->output_offset;
-
-    // Simulate NEON-register transposition of subset of filter.
-    int8_t filter_bank_a_0[4][4];  // Depth 4, width 4.
-    int8_t filter_bank_a_1[4][4];
-    int8_t filter_bank_a_2[4][4];
-    int8_t filter_bank_b_0[4][4];
-    int8_t filter_bank_b_1[4][4];
-    int8_t filter_bank_b_2[4][4];
-    // Simulate NEON-register input data concatenation + sub-selection.
-    // Also sub-block, height 3, depth 4, width 4.
-    uint8_t output_values[4];  // Sub-block, depth 4.
-    // selected_data has format Depth 4, width 4.
-    int8_t left_bank_0[4][4];
-    int8_t left_bank_1[4][4];
-    int8_t left_bank_2[4][4];
-    int8_t right_bank_0[4][4];
-    int8_t right_bank_1[4][4];
-    int8_t right_bank_2[4][4];
-    memset(right_bank_0[0], 0, 16);
-    memset(right_bank_1[0], 0, 16);
-    memset(right_bank_2[0], 0, 16);
-
-    constexpr int shuffled_filter_increment = 2 * 3 * 4 * 4;
-
-    for (int j_depth = 0; j_depth < depth_micro_repeats; ++j_depth) {
-      const int8_t* filter_block =
-          filter_workspace + shuffled_filter_increment * j_depth;
-
-      memcpy(filter_bank_a_0, filter_block, 16);
-      memcpy(filter_bank_b_0, filter_block + 16, 16);
-      memcpy(filter_bank_a_1, filter_block + 32, 16);
-      memcpy(filter_bank_b_1, filter_block + 48, 16);
-      memcpy(filter_bank_a_2, filter_block + 64, 16);
-      memcpy(filter_bank_b_2, filter_block + 80, 16);
-
-      for (int s = 0; s < 2; ++s) {
-        // Work through one slice, by row, at a time.
-        for (int k_height = 0; k_height < block_height; ++k_height) {
-          const int8_t* scratch_data =
-              scratch_block_data +
-              workspace_height_stride * k_height * stride_val +
-              depth_micro_stride * j_depth;
-          uint8_t* output_data =
-              output_block_data + output_height_stride * k_height + 8 * j_depth;
-          const int8_t* input_data_0 = scratch_data + s * 2 * 8;
-
-          // Load first sub-micro block of data into operational banks.
-          memcpy(left_bank_0[0], input_data_0, 16);
-          memcpy(left_bank_1[0], input_data_0 + workspace_height_stride, 16);
-          memcpy(left_bank_2[0], input_data_0 + 2 * workspace_height_stride,
-                 16);
-
-          for (int i_width = 0; i_width < output_width_overall_micro_repeats;
-               ++i_width) {
-            const int output_width = i_width == output_width_micro_repeats
-                                         ? residual_width
-                                         : four_over_stride;
-            TFLITE_DCHECK_LE(output_width * stride_val, 4);
-            const int8_t* input_data =
-                input_data_0 + width_micro_stride * i_width;
-            const bool no_right_block = (output_width - 1) * stride_val < 2;
-
-            // Load next sub-micro block of data.
-            if (!no_right_block) {
-              memcpy(right_bank_0[0], input_data + width_micro_stride, 16);
-              memcpy(right_bank_1[0],
-                     input_data + workspace_height_stride + width_micro_stride,
-                     16);
-              memcpy(
-                  right_bank_2[0],
-                  input_data + 2 * workspace_height_stride + width_micro_stride,
-                  16);
-            }
-
-            // Iterate over input width shifts within 4x4 blocks.
-            for (int x = 0; x < output_width; ++x) {
-              // Operate on depth of 4 in batches.
-              for (int d = 0; d < 4; ++d) {
-                int32_t acc = 0;
-                for (int x = 0; x < 4; ++x) {
-                  int32_t input_val = left_bank_0[d][x];
-                  int32_t filter_val = filter_bank_a_0[d][x];
-                  acc += filter_val * input_val;
-                }
-                for (int x = 0; x < 4; ++x) {
-                  int32_t input_val = left_bank_1[d][x];
-                  int32_t filter_val = filter_bank_a_1[d][x];
-                  acc += filter_val * input_val;
-                }
-                for (int x = 0; x < 4; ++x) {
-                  int32_t input_val = left_bank_2[d][x];
-                  int32_t filter_val = filter_bank_a_2[d][x];
-                  acc += filter_val * input_val;
-                }
-                acc += bias_data[d];
-                acc = reference_ops::depthwise_conv::DepthwiseConvRound<
-                    DepthwiseConvOutputRounding::kUpward>(
-                    acc, output_multiplier, output_shift);
-                acc += output_offset;
-                acc = std::max(acc, output_activation_min);
-                acc = std::min(acc, output_activation_max);
-                output_values[d] = static_cast<uint8_t>(acc);
-              }
-
-              for (int d = 0; d < 4; ++d) {
-                output_data[depth * (four_over_stride * i_width + x) + 4 * s +
-                            d] = output_values[d];
-              }
-
-              // Simulate shifting instructions.
-              if (stride_val == 1) {
-                for (int depth_index = 0; depth_index < 4; ++depth_index) {
-                  for (int z = 0; z < 3; ++z) {
-                    left_bank_0[depth_index][z] =
-                        left_bank_0[depth_index][z + 1];
-                    left_bank_1[depth_index][z] =
-                        left_bank_1[depth_index][z + 1];
-                    left_bank_2[depth_index][z] =
-                        left_bank_2[depth_index][z + 1];
-                  }
-                  left_bank_0[depth_index][3] = right_bank_0[depth_index][0];
-                  left_bank_1[depth_index][3] = right_bank_1[depth_index][0];
-                  left_bank_2[depth_index][3] = right_bank_2[depth_index][0];
-                  for (int z = 0; z < 3; ++z) {
-                    right_bank_0[depth_index][z] =
-                        right_bank_0[depth_index][z + 1];
-                    right_bank_1[depth_index][z] =
-                        right_bank_1[depth_index][z + 1];
-                    right_bank_2[depth_index][z] =
-                        right_bank_2[depth_index][z + 1];
-                  }
-                }
-              } else {
-                for (int depth_index = 0; depth_index < 4; ++depth_index) {
-                  for (int z = 0; z < 2; ++z) {
-                    left_bank_0[depth_index][z] =
-                        left_bank_0[depth_index][z + 2];
-                    left_bank_1[depth_index][z] =
-                        left_bank_1[depth_index][z + 2];
-                    left_bank_2[depth_index][z] =
-                        left_bank_2[depth_index][z + 2];
-                  }
-                  left_bank_0[depth_index][2] = right_bank_0[depth_index][0];
-                  left_bank_1[depth_index][2] = right_bank_1[depth_index][0];
-                  left_bank_2[depth_index][2] = right_bank_2[depth_index][0];
-                  left_bank_0[depth_index][3] = right_bank_0[depth_index][1];
-                  left_bank_1[depth_index][3] = right_bank_1[depth_index][1];
-                  left_bank_2[depth_index][3] = right_bank_2[depth_index][1];
-                  for (int z = 0; z < 2; ++z) {
-                    right_bank_0[depth_index][z] =
-                        right_bank_0[depth_index][z + 2];
-                    right_bank_1[depth_index][z] =
-                        right_bank_1[depth_index][z + 2];
-                    right_bank_2[depth_index][z] =
-                        right_bank_2[depth_index][z + 2];
-                  }
-                }
-              }
-            }
-          }
-        }
-        bias_data += bias_increment;
-
-        // Move filter for second sub-block into operational filter.
-        for (int z = 0; z < 4; ++z) {
-          for (int x = 0; x < 4; ++x) {
-            filter_bank_a_0[z][x] = filter_bank_b_0[z][x];
-            filter_bank_a_1[z][x] = filter_bank_b_1[z][x];
-            filter_bank_a_2[z][x] = filter_bank_b_2[z][x];
-          }
-        }
-      }
-    }
-  }
-};
-
-template <int32_t stride, QuantizationType quantization_type>
-struct KernelMacroBlock<
-    DepthwiseConvImplementation::kUseUnwound3x3DotProduct, quantization_type,
-    DepthwiseConvDepthMultiplication::kUnitInputDepth, stride> {
-  static inline void Run(const int8_t* scratch_block_data,
-                         const int8_t* filter_workspace,
-                         const int32_t* bias_data, uint8_t* output_block_data,
-                         const DepthwiseConvDotProdParams* function_params) {
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-    const int output_width_micro_repeats =
-        function_params->output_width_micro_repeats;
-    const int depth_micro_repeats = function_params->depth_micro_repeats;
-    const int output_depth = function_params->output_depth;
-    const int stride_val = function_params->stride;
-    const int four_over_stride = function_params->four_over_stride;
-
-    const int output_width_overall_micro_repeats =
-        function_params->output_width_overall_micro_repeats;
-    const int block_height = function_params->outbound_block_height;
-    const int residual_width = function_params->output_residual_width;
-    const int output_height_stride = function_params->output_height_stride;
-    const int bias_increment = function_params->bias_increment;
-
-    const int32_t output_activation_min =
-        function_params->quantized_activation_min;
-    const int32_t output_activation_max =
-        function_params->quantized_activation_max;
-    const int32_t output_multiplier = function_params->output_multiplier;
-    const int32_t output_shift = function_params->output_shift;
-    const int32_t output_offset = function_params->output_offset;
-
-    TFLITE_DCHECK(depth_micro_repeats > 0);
-
-    TFLITE_DCHECK_EQ(bias_increment, 4);
-
-    constexpr int shuffled_filter_increment = 2 * 3 * 4 * 4;
-
-    // Simulate NEON-register transposition of subset of filter.
-    int8_t filter_bank_a_0[4][4];  // Depth 4, width 4.
-    int8_t filter_bank_a_1[4][4];
-    int8_t filter_bank_a_2[4][4];
-    int8_t filter_bank_b_0[4][4];
-    int8_t filter_bank_b_1[4][4];
-    int8_t filter_bank_b_2[4][4];
-    // Simulate NEON-register input data concatenation + sub-selection.
-    // Also sub-block, height 3, depth 4, width 4.
-
-    int8_t input_bank_0[8];
-    int8_t input_bank_1[8];
-    int8_t input_bank_2[8];
-
-    TFLITE_DCHECK_GE(depth_micro_repeats, 1);
-
-    uint8_t output_values[2][4];  // Sub-block, depth 4.
-
-    for (int j_depth = 0; j_depth < depth_micro_repeats; ++j_depth) {
-      memcpy(filter_bank_a_0, filter_workspace, 16);
-      memcpy(filter_bank_b_0, filter_workspace + 16, 16);
-      memcpy(filter_bank_a_1, filter_workspace + 32, 16);
-      memcpy(filter_bank_b_1, filter_workspace + 48, 16);
-      memcpy(filter_bank_a_2, filter_workspace + 64, 16);
-      memcpy(filter_bank_b_2, filter_workspace + 80, 16);
-
-      // Work through one slice, by row, at a time.
-      for (int k_height = 0; k_height < block_height; ++k_height) {
-        const int8_t* scratch_data =
-            scratch_block_data +
-            workspace_height_stride * k_height * stride_val;
-        uint8_t* output_data =
-            output_block_data + output_height_stride * k_height + 8 * j_depth;
-
-        memcpy(input_bank_0, scratch_data, 4);
-        memcpy(input_bank_1, scratch_data + workspace_height_stride, 4);
-        memcpy(input_bank_2, scratch_data + 2 * workspace_height_stride, 4);
-
-        for (int i_width = 0; i_width < output_width_overall_micro_repeats;
-             ++i_width) {
-          const int output_width = i_width == output_width_micro_repeats
-                                       ? residual_width
-                                       : four_over_stride;
-
-          TFLITE_DCHECK_LE(output_width * stride_val, 4);
-          const int8_t* input_data = scratch_data + 4 * i_width;
-
-          memcpy(input_bank_0 + 4, input_data + 4, 4);
-          memcpy(input_bank_1 + 4, input_data + workspace_height_stride + 4, 4);
-          memcpy(input_bank_2 + 4, input_data + 2 * workspace_height_stride + 4,
-                 4);
-
-          // Iterate over input width shifts within 4x4 blocks.
-          for (int w = 0; w < output_width; ++w) {
-            constexpr int offset =
-                0;  // Shift input instead of offset in multiply-accumulate.
-
-            {
-              const int s = 0;
-              for (int d = 0; d < 4; ++d) {
-                int32_t acc = bias_data[s * 4 + d];
-                for (int x = 0; x < 4; ++x) {
-                  int32_t input_val_0 = input_bank_0[offset + x];
-                  int32_t filter_val_0 = filter_bank_a_0[d][x];
-                  acc += filter_val_0 * input_val_0;
-                  int32_t input_val_1 = input_bank_1[offset + x];
-                  int32_t filter_val_1 = filter_bank_a_1[d][x];
-                  acc += filter_val_1 * input_val_1;
-                  int32_t input_val_2 = input_bank_2[offset + x];
-                  int32_t filter_val_2 = filter_bank_a_2[d][x];
-                  acc += filter_val_2 * input_val_2;
-                }
-                acc = reference_ops::depthwise_conv::DepthwiseConvRound<
-                    DepthwiseConvOutputRounding::kUpward>(
-                    acc, output_multiplier, output_shift);
-                acc += output_offset;
-                acc = std::max(acc, output_activation_min);
-                acc = std::min(acc, output_activation_max);
-                output_values[s][d] = static_cast<uint8_t>(acc);
-
-                output_data[s * 4 + d] = output_values[s][d];
-              }
-            }
-            {
-              const int s = 1;
-              for (int d = 0; d < 4; ++d) {
-                int32_t acc = bias_data[s * 4 + d];
-                for (int x = 0; x < 4; ++x) {
-                  int32_t input_val_0 = input_bank_0[offset + x];
-                  int32_t filter_val_0 = filter_bank_b_0[d][x];
-                  acc += filter_val_0 * input_val_0;
-                  int32_t input_val_1 = input_bank_1[offset + x];
-                  int32_t filter_val_1 = filter_bank_b_1[d][x];
-                  acc += filter_val_1 * input_val_1;
-                  int32_t input_val_2 = input_bank_2[offset + x];
-                  int32_t filter_val_2 = filter_bank_b_2[d][x];
-                  acc += filter_val_2 * input_val_2;
-                }
-                acc = reference_ops::depthwise_conv::DepthwiseConvRound<
-                    DepthwiseConvOutputRounding::kUpward>(
-                    acc, output_multiplier, output_shift);
-                acc += output_offset;
-                acc = std::max(acc, output_activation_min);
-                acc = std::min(acc, output_activation_max);
-                output_values[s][d] = static_cast<uint8_t>(acc);
-
-                output_data[s * 4 + d] = output_values[s][d];
-              }
-            }
-
-            // Simulate register shifts.
-            for (int i = 0; i < (8 - stride_val); ++i) {
-              input_bank_0[i] = input_bank_0[i + stride_val];
-              input_bank_1[i] = input_bank_1[i + stride_val];
-              input_bank_2[i] = input_bank_2[i + stride_val];
-            }
-
-            output_data += output_depth;
-          }
-        }
-      }
-      bias_data += 2 * bias_increment;
-      filter_workspace += shuffled_filter_increment;
-    }
-  }
-};
-// The preceding section is only compiled when kUseUnwound3x3DotProduct versions
-// of templated functions are selected.
-//
-// End of code section containing intermediate code transformation.
-
-#ifdef USE_NEON
-template <>
-struct KernelMacroBlock<
-    DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct,
-    QuantizationType::kNonPerChannelUint8,
-    DepthwiseConvDepthMultiplication::kNoMultiplication,
-    /*stride=*/1> {
-  static inline uint8x8_t vqmovxn_s16(int16x8_t x) { return vqmovun_s16(x); }
-  static inline uint8x8_t util_vmin_x8(uint8x8_t a, uint8x8_t b) {
-    return vmin_u8(a, b);
-  }
-  static inline uint8x8_t util_vmax_x8(uint8x8_t a, uint8x8_t b) {
-    return vmax_u8(a, b);
-  }
-  static inline uint8x16_t util_vminq_x8(uint8x16_t a, uint8x16_t b) {
-    return vminq_u8(a, b);
-  }
-  static inline uint8x16_t util_vmaxq_x8(uint8x16_t a, uint8x16_t b) {
-    return vmaxq_u8(a, b);
-  }
-
-  static inline void KernelMacroBlockIntrinsics(
-      const int8_t* scratch_block_data, const int8_t* filter_workspace,
-      const int32_t* bias_data, uint8_t* output_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-    static constexpr QuantizationType quantization_type =
-        QuantizationType::kNonPerChannelUint8;
-
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-    const int input_width_overall_micro_repeats =
-        function_params->input_width_overall_micro_repeats;
-    const int output_width_micro_repeats =
-        function_params->output_width_micro_repeats;
-    const int depth_micro_repeats = function_params->depth_micro_repeats;
-    const int depth = function_params->input_depth;
-
-    const int output_width_overall_micro_repeats =
-        function_params->output_width_overall_micro_repeats;
-    const int block_height = function_params->outbound_block_height;
-    const int residual_width = function_params->output_residual_width;
-    const int output_height_stride = function_params->output_height_stride;
-    constexpr int kBiasIncrement = 4;
-
-    TFLITE_DCHECK(depth_micro_repeats > 0);
-    const int width_micro_stride = 4 * 8;
-    const int depth_micro_stride =
-        width_micro_stride * input_width_overall_micro_repeats;
-
-    const int32_t output_activation_min =
-        function_params->quantized_activation_min;
-    const int32_t output_activation_max =
-        function_params->quantized_activation_max;
-    const int32_t output_multiplier = function_params->output_multiplier;
-    const int32_t output_shift = function_params->output_shift;
-    const int32_t output_offset = function_params->output_offset;
-    if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-      TFLITE_DCHECK_GE(output_activation_min, 0);
-      TFLITE_DCHECK_LT(output_activation_min, 256);
-      TFLITE_DCHECK_GE(output_activation_max, 0);
-      TFLITE_DCHECK_LT(output_activation_max, 256);
-    } else {
-      TFLITE_DCHECK_GE(output_activation_min, -128);
-      TFLITE_DCHECK_LT(output_activation_min, 128);
-      TFLITE_DCHECK_GE(output_activation_max, -128);
-      TFLITE_DCHECK_LT(output_activation_max, 128);
-    }
-    TFLITE_DCHECK_GE(output_offset, -32878);
-    TFLITE_DCHECK_LT(output_offset, 32768);
-
-    const int16x8_t output_offset_vec =
-        vdupq_n_s16(static_cast<int16_t>(output_offset));
-    const uint8x16_t output_activation_min_vec =
-        vdupq_n_u8(static_cast<uint8_t>(output_activation_min));
-    const uint8x16_t output_activation_max_vec =
-        vdupq_n_u8(static_cast<uint8_t>(output_activation_max));
-
-    const int8_t* input_data_depthwise = scratch_block_data;
-    typename QuantizationTypeImpl<quantization_type>::ExternalType*
-        output_data_depthwise = output_block_data;
-    for (int j_depth = 0; j_depth < depth_micro_repeats; ++j_depth) {
-      // Simulate NEON-register transposition of subset of filter.
-      int8x16_t filter_reg_0_a;
-      int8x16_t filter_reg_0_b;
-      int8x16_t filter_reg_1_a;
-      int8x16_t filter_reg_1_b;
-      int8x16_t filter_reg_2_a;
-      int8x16_t filter_reg_2_b;
-      int8x16_t filter_reg_0_a_shifted;
-      int8x16_t filter_reg_1_a_shifted;
-      int8x16_t filter_reg_2_a_shifted;
-
-      filter_reg_0_a = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_0_b = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_1_a = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_1_b = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_2_a = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_2_b = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-
-      filter_reg_0_a_shifted = vreinterpretq_s8_u32(
-          vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_0_a), 8));
-      filter_reg_1_a_shifted = vreinterpretq_s8_u32(
-          vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_1_a), 8));
-      filter_reg_2_a_shifted = vreinterpretq_s8_u32(
-          vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_2_a), 8));
-
-      if (block_height == 4) {
-        for (int s = 0; s < 2; ++s) {
-          // Work through one slice, by row, at a time.
-          const int8_t* input_data_base = input_data_depthwise + 2 * 8 * s;
-          typename QuantizationTypeImpl<quantization_type>::ExternalType*
-              output_data_base = output_data_depthwise + 4 * s;
-
-          const int8_t* next_input_data = input_data_base;
-          typename QuantizationTypeImpl<quantization_type>::ExternalType*
-              output_data = output_data_base;
-
-          const int32x4_t adjusted_bias_data = vld1q_s32(bias_data);
-          bias_data += kBiasIncrement;
-
-          // Load first sub-micro block of data into operational banks.
-          int8x16_t left_bank_0_reg = vld1q_s8(next_input_data);
-          int8x16_t left_bank_1_reg =
-              vld1q_s8(next_input_data + workspace_height_stride);
-          int8x16_t left_bank_2_reg =
-              vld1q_s8(next_input_data + 2 * workspace_height_stride);
-          int8x16_t left_bank_3_reg =
-              vld1q_s8(next_input_data + 3 * workspace_height_stride);
-          int8x16_t left_bank_4_reg =
-              vld1q_s8(next_input_data + 4 * workspace_height_stride);
-          int8x16_t left_bank_5_reg =
-              vld1q_s8(next_input_data + 5 * workspace_height_stride);
-
-          int32x4_t acc0;
-          int32x4_t acc1;
-          int32x4_t acc2;
-          int32x4_t acc3;
-
-          acc0 = adjusted_bias_data;
-          acc1 = adjusted_bias_data;
-          acc2 = adjusted_bias_data;
-          acc3 = adjusted_bias_data;
-
-          acc0 = vdotq_s32(acc0, filter_reg_2_a, left_bank_2_reg);
-          acc1 = vdotq_s32(acc1, filter_reg_1_a, left_bank_2_reg);
-          acc2 = vdotq_s32(acc2, filter_reg_0_a, left_bank_2_reg);
-          acc3 = vdotq_s32(acc3, filter_reg_0_a, left_bank_3_reg);
-
-          for (int i_width = 0; i_width < output_width_micro_repeats;
-               ++i_width) {
-            next_input_data += width_micro_stride;
-
-            // Iterate over input width shifts within 4x4 blocks.
-            {
-              acc0 = vdotq_s32(acc0, filter_reg_0_a, left_bank_0_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_1_a, left_bank_1_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_0_a, left_bank_1_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_2_a, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_1_a, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_2_a, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_1_a, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_2_a, left_bank_5_reg);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc0, -output_shift);
-              acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc1, -output_shift);
-              acc2 = vqrdmulhq_n_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc2, -output_shift);
-              acc3 = vqrdmulhq_n_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc3, -output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              uint8x16_t acc_u8_all = vcombine_u8(vqmovxn_s16(acc_s16_0_1),
-                                                  vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_u8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_u8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_u8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_u8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              output_data += depth;
-            }
-
-            // Load next sub-micro block of data.
-            int8x16_t right_bank_0_reg;
-            int8x16_t right_bank_1_reg;
-            int8x16_t right_bank_2_reg;
-            int8x16_t right_bank_3_reg;
-            int8x16_t right_bank_4_reg;
-            int8x16_t right_bank_5_reg;
-
-            // Loading of next block always valid.
-            right_bank_0_reg = vld1q_s8(next_input_data);
-            right_bank_1_reg =
-                vld1q_s8(next_input_data + workspace_height_stride);
-            right_bank_2_reg =
-                vld1q_s8(next_input_data + 2 * workspace_height_stride);
-            right_bank_3_reg =
-                vld1q_s8(next_input_data + 3 * workspace_height_stride);
-            right_bank_4_reg =
-                vld1q_s8(next_input_data + 4 * workspace_height_stride);
-            right_bank_5_reg =
-                vld1q_s8(next_input_data + 5 * workspace_height_stride);
-
-            {
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_s32(acc0, filter_reg_0_a_shifted, left_bank_0_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_1_a_shifted, left_bank_1_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_2_a_shifted, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_0_a_shifted, left_bank_1_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_1_a_shifted, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_2_a_shifted, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_0_a_shifted, left_bank_2_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_1_a_shifted, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_2_a_shifted, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_0_a_shifted, left_bank_3_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_1_a_shifted, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_2_a_shifted, left_bank_5_reg);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc0, -output_shift);
-              acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc1, -output_shift);
-              acc2 = vqrdmulhq_n_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc2, -output_shift);
-              acc3 = vqrdmulhq_n_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc3, -output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              uint8x16_t acc_u8_all = vcombine_u8(vqmovxn_s16(acc_s16_0_1),
-                                                  vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_u8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_u8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_u8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_u8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              left_bank_0_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_0_reg)));
-              left_bank_1_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_1_reg)));
-              left_bank_2_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_2_reg)));
-              left_bank_3_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_3_reg)));
-              left_bank_4_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_4_reg)));
-              left_bank_5_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_5_reg)));
-              vtrn1_s8x2_in_place(&left_bank_0_reg, &right_bank_0_reg);
-              vtrn1_s8x2_in_place(&left_bank_1_reg, &right_bank_1_reg);
-              vtrn1_s8x2_in_place(&left_bank_2_reg, &right_bank_2_reg);
-              vtrn1_s8x2_in_place(&left_bank_3_reg, &right_bank_3_reg);
-              vtrn1_s8x2_in_place(&left_bank_4_reg, &right_bank_4_reg);
-              vtrn1_s8x2_in_place(&left_bank_5_reg, &right_bank_5_reg);
-
-              output_data += depth;
-            }
-
-            {
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_s32(acc0, filter_reg_0_a, left_bank_0_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_1_a, left_bank_1_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_2_a, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_0_a, left_bank_1_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_1_a, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_2_a, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_0_a, left_bank_2_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_1_a, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_2_a, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_0_a, left_bank_3_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_1_a, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_2_a, left_bank_5_reg);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc0, -output_shift);
-              acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc1, -output_shift);
-              acc2 = vqrdmulhq_n_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc2, -output_shift);
-              acc3 = vqrdmulhq_n_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc3, -output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              uint8x16_t acc_u8_all = vcombine_u8(vqmovxn_s16(acc_s16_0_1),
-                                                  vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_u8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_u8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_u8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_u8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              output_data += depth;
-            }
-
-            {
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_s32(acc0, filter_reg_0_a_shifted, left_bank_0_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_1_a_shifted, left_bank_1_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_2_a_shifted, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_0_a_shifted, left_bank_1_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_1_a_shifted, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_2_a_shifted, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_0_a_shifted, left_bank_2_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_1_a_shifted, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_2_a_shifted, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_0_a_shifted, left_bank_3_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_1_a_shifted, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_2_a_shifted, left_bank_5_reg);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc0, -output_shift);
-              acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc1, -output_shift);
-              acc2 = vqrdmulhq_n_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc2, -output_shift);
-              acc3 = vqrdmulhq_n_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc3, -output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              uint8x16_t acc_u8_all = vcombine_u8(vqmovxn_s16(acc_s16_0_1),
-                                                  vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_u8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_u8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_u8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_u8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              left_bank_0_reg = right_bank_0_reg;
-              left_bank_1_reg = right_bank_1_reg;
-              left_bank_2_reg = right_bank_2_reg;
-              left_bank_3_reg = right_bank_3_reg;
-              left_bank_4_reg = right_bank_4_reg;
-              left_bank_5_reg = right_bank_5_reg;
-
-              output_data += depth;
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_s32(acc0, filter_reg_2_a, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_1_a, left_bank_2_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_0_a, left_bank_2_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_0_a, left_bank_3_reg);
-            }
-          }
-
-          if (residual_width > 0) {
-            next_input_data += width_micro_stride;
-            const int output_width = residual_width;
-
-            // Load next sub-micro block of data.
-            int8x16_t right_bank_0_reg;
-            int8x16_t right_bank_1_reg;
-            int8x16_t right_bank_2_reg;
-            int8x16_t right_bank_3_reg;
-            int8x16_t right_bank_4_reg;
-            int8x16_t right_bank_5_reg;
-            // Logic: (output_width - 1) * stride_val < 2.
-            const bool no_right_block = output_width < 3;
-
-            if (no_right_block) {
-              // Only needed for sanitizer checks.
-              right_bank_0_reg = vdupq_n_s8(0);
-              right_bank_1_reg = vdupq_n_s8(0);
-              right_bank_2_reg = vdupq_n_s8(0);
-              right_bank_3_reg = vdupq_n_s8(0);
-              right_bank_4_reg = vdupq_n_s8(0);
-              right_bank_5_reg = vdupq_n_s8(0);
-            } else {
-              right_bank_0_reg = vld1q_s8(next_input_data);
-              right_bank_1_reg =
-                  vld1q_s8(next_input_data + workspace_height_stride);
-              right_bank_2_reg =
-                  vld1q_s8(next_input_data + 2 * workspace_height_stride);
-              right_bank_3_reg =
-                  vld1q_s8(next_input_data + 3 * workspace_height_stride);
-              right_bank_4_reg =
-                  vld1q_s8(next_input_data + 4 * workspace_height_stride);
-              right_bank_5_reg =
-                  vld1q_s8(next_input_data + 5 * workspace_height_stride);
-            }
-
-            // Iterate over input width shifts within 4x4 blocks.
-            for (int x = 0; x < output_width; ++x) {
-              acc0 = vdotq_s32(acc0, filter_reg_0_a, left_bank_0_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_1_a, left_bank_1_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_0_a, left_bank_1_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_2_a, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_1_a, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_2_a, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_1_a, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_2_a, left_bank_5_reg);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc0, -output_shift);
-              acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc1, -output_shift);
-              acc2 = vqrdmulhq_n_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc2, -output_shift);
-              acc3 = vqrdmulhq_n_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc3, -output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              uint8x16_t acc_u8_all = vcombine_u8(vqmovxn_s16(acc_s16_0_1),
-                                                  vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_u8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_u8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_u8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_u8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              biregister_rotate_8(&left_bank_0_reg, &right_bank_0_reg);
-              biregister_rotate_8(&left_bank_1_reg, &right_bank_1_reg);
-              biregister_rotate_8(&left_bank_2_reg, &right_bank_2_reg);
-              biregister_rotate_8(&left_bank_3_reg, &right_bank_3_reg);
-              biregister_rotate_8(&left_bank_4_reg, &right_bank_4_reg);
-              biregister_rotate_8(&left_bank_5_reg, &right_bank_5_reg);
-
-              output_data += depth;
-
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_s32(acc0, filter_reg_2_a, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_1_a, left_bank_2_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_0_a, left_bank_2_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_0_a, left_bank_3_reg);
-            }
-          }
-          input_data_base += 4 * workspace_height_stride;
-          output_data_base += 4 * output_height_stride;
-
-          // Move to next sub-block: advance to second set of filters, to new
-          // bias.
-          filter_reg_0_a = filter_reg_0_b;
-          filter_reg_1_a = filter_reg_1_b;
-          filter_reg_2_a = filter_reg_2_b;
-          filter_reg_0_a_shifted = vreinterpretq_s8_u32(
-              vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_0_a), 8));
-          filter_reg_1_a_shifted = vreinterpretq_s8_u32(
-              vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_1_a), 8));
-          filter_reg_2_a_shifted = vreinterpretq_s8_u32(
-              vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_2_a), 8));
-        }
-      } else {
-        const int8_t* input_data_base = input_data_depthwise;
-        typename QuantizationTypeImpl<quantization_type>::ExternalType*
-            output_data_base = output_data_depthwise;
-
-        const int32x4_t adjusted_bias_data_a = vld1q_s32(bias_data);
-        bias_data += kBiasIncrement;
-        const int32x4_t adjusted_bias_data_b = vld1q_s32(bias_data);
-        bias_data += kBiasIncrement;
-
-        for (int k_height = 0; k_height < block_height; ++k_height) {
-          const int8_t* next_input_data = input_data_base;
-          typename QuantizationTypeImpl<quantization_type>::ExternalType*
-              output_data = output_data_base;
-
-          // Load first sub-micro block of data into operational banks.
-          int8x16_t left_bank_0_reg_a = vld1q_s8(next_input_data);
-          int8x16_t left_bank_1_reg_a =
-              vld1q_s8(next_input_data + workspace_height_stride);
-          int8x16_t left_bank_2_reg_a =
-              vld1q_s8(next_input_data + 2 * workspace_height_stride);
-          int8x16_t left_bank_0_reg_b = vld1q_s8(next_input_data + 16);
-          int8x16_t left_bank_1_reg_b =
-              vld1q_s8(next_input_data + workspace_height_stride + 16);
-          int8x16_t left_bank_2_reg_b =
-              vld1q_s8(next_input_data + 2 * workspace_height_stride + 16);
-
-          for (int i_width = 0; i_width < output_width_overall_micro_repeats;
-               ++i_width) {
-            next_input_data += width_micro_stride;
-            const int output_width =
-                i_width == output_width_micro_repeats ? residual_width : 4;
-
-            int8x16_t right_bank_0_reg_a;
-            int8x16_t right_bank_1_reg_a;
-            int8x16_t right_bank_2_reg_a;
-            int8x16_t right_bank_0_reg_b;
-            int8x16_t right_bank_1_reg_b;
-            int8x16_t right_bank_2_reg_b;
-            // Logic: (output_width - 1) * stride_val < 2.
-            const bool no_right_block = output_width < 3;
-
-            // Load next sub-micro block of data.
-            if (no_right_block) {
-              // Only needed for sanitizer checks.
-              right_bank_0_reg_a = vdupq_n_s8(0);
-              right_bank_1_reg_a = vdupq_n_s8(0);
-              right_bank_2_reg_a = vdupq_n_s8(0);
-              right_bank_0_reg_b = vdupq_n_s8(0);
-              right_bank_1_reg_b = vdupq_n_s8(0);
-              right_bank_2_reg_b = vdupq_n_s8(0);
-            } else {
-              right_bank_0_reg_a = vld1q_s8(next_input_data);
-              right_bank_1_reg_a =
-                  vld1q_s8(next_input_data + workspace_height_stride);
-              right_bank_2_reg_a =
-                  vld1q_s8(next_input_data + 2 * workspace_height_stride);
-              right_bank_0_reg_b = vld1q_s8(next_input_data + 16);
-              right_bank_1_reg_b =
-                  vld1q_s8(next_input_data + workspace_height_stride + 16);
-              right_bank_2_reg_b =
-                  vld1q_s8(next_input_data + 2 * workspace_height_stride + 16);
-            }
-
-            // Iterate over input width shifts within 4x4 blocks.
-            for (int x = 0; x < output_width; ++x) {
-              int32x4_t acc_a = adjusted_bias_data_a;
-              int32x4_t acc_b = adjusted_bias_data_b;
-              acc_a = vdotq_s32(acc_a, filter_reg_0_a, left_bank_0_reg_a);
-              acc_a = vdotq_s32(acc_a, filter_reg_1_a, left_bank_1_reg_a);
-              acc_a = vdotq_s32(acc_a, filter_reg_2_a, left_bank_2_reg_a);
-              acc_b = vdotq_s32(acc_b, filter_reg_0_b, left_bank_0_reg_b);
-              acc_b = vdotq_s32(acc_b, filter_reg_1_b, left_bank_1_reg_b);
-              acc_b = vdotq_s32(acc_b, filter_reg_2_b, left_bank_2_reg_b);
-
-              // Fixed-point multiplication.
-              acc_a = vqrdmulhq_n_s32(acc_a, output_multiplier);
-              acc_b = vqrdmulhq_n_s32(acc_b, output_multiplier);
-              acc_a = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc_a, -output_shift);
-              acc_b = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc_b, -output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_0 =
-                  vcombine_s16(vqmovn_s32(acc_a), vqmovn_s32(acc_b));
-              acc_s16_0_0 = vqaddq_s16(acc_s16_0_0, output_offset_vec);
-              // Apply the activation function.
-              uint8x8_t acc_u8_0_0 = vqmovxn_s16(acc_s16_0_0);
-              acc_u8_0_0 = util_vmax_x8(acc_u8_0_0,
-                                        vget_low_u8(output_activation_min_vec));
-              acc_u8_0_0 = util_vmin_x8(acc_u8_0_0,
-                                        vget_low_u8(output_activation_max_vec));
-
-              util_vst1_u8(output_data, acc_u8_0_0);
-
-              biregister_rotate_8(&left_bank_0_reg_a, &right_bank_0_reg_a);
-              biregister_rotate_8(&left_bank_1_reg_a, &right_bank_1_reg_a);
-              biregister_rotate_8(&left_bank_2_reg_a, &right_bank_2_reg_a);
-              biregister_rotate_8(&left_bank_0_reg_b, &right_bank_0_reg_b);
-              biregister_rotate_8(&left_bank_1_reg_b, &right_bank_1_reg_b);
-              biregister_rotate_8(&left_bank_2_reg_b, &right_bank_2_reg_b);
-
-              output_data += depth;
-            }
-          }
-          input_data_base += workspace_height_stride;
-          output_data_base += output_height_stride;
-        }
-      }
-      input_data_depthwise += depth_micro_stride;
-      output_data_depthwise += 8;
-    }
-  }  // NOLINT(readability/fn_size) Manually unrolled.
-
-  static inline void Run(const int8_t* scratch_block_data,
-                         const int8_t* filter_workspace,
-                         const int32_t* bias_data, uint8_t* output_block_data,
-                         const DepthwiseConvDotProdParams* function_params) {
-    KernelMacroBlockIntrinsics(scratch_block_data, filter_workspace, bias_data,
-                               output_block_data, function_params);
-  }
-};
-
-template <>
-struct KernelMacroBlock<
-    DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct,
-    QuantizationType::kNonPerChannelUint8,
-    DepthwiseConvDepthMultiplication::kNoMultiplication,
-    /*stride=*/2> {
-  static inline uint8x8_t vqmovxn_s16(int16x8_t x) { return vqmovun_s16(x); }
-  static inline uint8x8_t util_vmin_x8(uint8x8_t a, uint8x8_t b) {
-    return vmin_u8(a, b);
-  }
-  static inline uint8x8_t util_vmax_x8(uint8x8_t a, uint8x8_t b) {
-    return vmax_u8(a, b);
-  }
-
-  static inline void KernelMacroBlockIntrinsics(
-      const int8_t* scratch_block_data, const int8_t* filter_workspace,
-      const int32_t* bias_data, uint8_t* output_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-    static constexpr QuantizationType quantization_type =
-        QuantizationType::kNonPerChannelUint8;
-
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-    const int input_width_overall_micro_repeats =
-        function_params->input_width_overall_micro_repeats;
-    const int output_width_micro_repeats =
-        function_params->output_width_micro_repeats;
-    const int depth_micro_repeats = function_params->depth_micro_repeats;
-    const int depth = function_params->input_depth;
-    constexpr int kStrideVal = 2;
-    constexpr int kFourOverStride = 2;
-    TFLITE_DCHECK_EQ(function_params->stride, kStrideVal);
-    TFLITE_DCHECK_EQ(function_params->four_over_stride, kFourOverStride);
-
-    const int workspace_width_micro_repeats =
-        function_params->workspace_width_micro_repeats;
-    const int output_width_overall_micro_repeats =
-        function_params->output_width_overall_micro_repeats;
-    const int block_height = function_params->outbound_block_height;
-    const int residual_width = function_params->output_residual_width;
-    const int output_height_stride = function_params->output_height_stride;
-    constexpr int kBiasIncrement = 4;
-
-    TFLITE_DCHECK(depth_micro_repeats > 0);
-    const int width_micro_stride = 4 * 8;
-    const int depth_micro_stride =
-        width_micro_stride * input_width_overall_micro_repeats;
-
-    const int32_t output_activation_min =
-        function_params->quantized_activation_min;
-    const int32_t output_activation_max =
-        function_params->quantized_activation_max;
-    const int32_t output_multiplier = function_params->output_multiplier;
-    const int32_t output_shift = function_params->output_shift;
-    const int32_t output_offset = function_params->output_offset;
-    if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-      TFLITE_DCHECK_GE(output_activation_min, 0);
-      TFLITE_DCHECK_LT(output_activation_min, 256);
-      TFLITE_DCHECK_GE(output_activation_max, 0);
-      TFLITE_DCHECK_LT(output_activation_max, 256);
-    } else {
-      TFLITE_DCHECK_GE(output_activation_min, -128);
-      TFLITE_DCHECK_LT(output_activation_min, 128);
-      TFLITE_DCHECK_GE(output_activation_max, -128);
-      TFLITE_DCHECK_LT(output_activation_max, 128);
-    }
-    TFLITE_DCHECK_GE(output_offset, -32878);
-    TFLITE_DCHECK_LT(output_offset, 32768);
-
-    // This version only does min/max on 64 bits.
-    const int16x8_t output_offset_vec =
-        vdupq_n_s16(static_cast<int16_t>(output_offset));
-    const uint8x8_t output_activation_min_vec =
-        vdup_n_u8(static_cast<uint8_t>(output_activation_min));
-    const uint8x8_t output_activation_max_vec =
-        vdup_n_u8(static_cast<uint8_t>(output_activation_max));
-
-    constexpr int shuffled_filter_increment = 2 * 3 * 4 * 4;
-
-    TFLITE_DCHECK_LE(block_height, 2);
-
-    for (int j_depth = 0; j_depth < depth_micro_repeats; ++j_depth) {
-      const int8_t* filter_block =
-          filter_workspace + shuffled_filter_increment * j_depth;
-
-      if (block_height == 2) {
-        for (int s = 0; s < 2; ++s) {
-          // Simulate NEON-register transposition of subset of filter.
-          int8x16_t filter_reg_0_a;
-          int8x16_t filter_reg_1_a;
-          int8x16_t filter_reg_2_a;
-
-          filter_reg_0_a = vld1q_s8(filter_block + s * 16);
-          filter_reg_1_a = vld1q_s8(filter_block + s * 16 + 32);
-          filter_reg_2_a = vld1q_s8(filter_block + s * 16 + 64);
-
-          const int8_t* scratch_data =
-              scratch_block_data + depth_micro_stride * j_depth;
-          typename QuantizationTypeImpl<quantization_type>::ExternalType*
-              output_data = output_block_data + 8 * j_depth;
-          const int8_t* input_data_0 = scratch_data + s * 2 * 8;
-
-          const int32x4_t adjusted_bias_data = vld1q_s32(bias_data);
-
-          // Load first sub-micro block of data into operational banks.
-          int8x16_t left_bank_0_reg = vld1q_s8(input_data_0);
-          int8x16_t left_bank_1_reg =
-              vld1q_s8(input_data_0 + workspace_height_stride);
-          int8x16_t left_bank_2_reg =
-              vld1q_s8(input_data_0 + 2 * workspace_height_stride);
-          int8x16_t left_bank_3_reg =
-              vld1q_s8(input_data_0 + 3 * workspace_height_stride);
-          int8x16_t left_bank_4_reg =
-              vld1q_s8(input_data_0 + 4 * workspace_height_stride);
-
-          int8x16_t right_bank_0_reg;
-          int8x16_t right_bank_1_reg;
-          int8x16_t right_bank_2_reg;
-          int8x16_t right_bank_3_reg;
-          int8x16_t right_bank_4_reg;
-
-          int32x4_t acc0;
-          int32x4_t acc1;
-          int16x8_t acc_s16_0_1;
-          uint8x8_t acc_u8;
-
-          int i_width = 0;
-
-          // When output_width_micro_repeats <
-          // output_width_overall_micro_repeats, 0 < residual_width <= 2, and so
-          // residual_width == 1 is then true iff residual_width < 2.
-          const int adjusted_width_micro_repeats =
-              (output_width_micro_repeats <
-               output_width_overall_micro_repeats) &&
-                      (residual_width == 1)
-                  ? output_width_micro_repeats
-                  : output_width_overall_micro_repeats;
-
-          for (; i_width < adjusted_width_micro_repeats; ++i_width) {
-            const int output_width = kFourOverStride;
-            TFLITE_DCHECK_LE(output_width * kStrideVal, 4);
-            const int8_t* input_data =
-                input_data_0 + width_micro_stride * i_width;
-            acc0 = adjusted_bias_data;
-            acc1 = adjusted_bias_data;
-            right_bank_0_reg = vld1q_s8(input_data + width_micro_stride);
-            right_bank_1_reg = vld1q_s8(input_data + width_micro_stride +
-                                        workspace_height_stride);
-
-            acc0 = vdotq_s32(acc0, filter_reg_0_a, left_bank_0_reg);
-            acc1 = vdotq_s32(acc1, filter_reg_0_a, left_bank_2_reg);
-            typename QuantizationTypeImpl<quantization_type>::ExternalType*
-                output_data_base = output_data + depth * 2 * i_width + 4 * s;
-
-            right_bank_2_reg = vld1q_s8(input_data + width_micro_stride +
-                                        2 * workspace_height_stride);
-            right_bank_3_reg = vld1q_s8(input_data + width_micro_stride +
-                                        3 * workspace_height_stride);
-            acc0 = vdotq_s32(acc0, filter_reg_1_a, left_bank_1_reg);
-            acc0 = vdotq_s32(acc0, filter_reg_2_a, left_bank_2_reg);
-            acc1 = vdotq_s32(acc1, filter_reg_1_a, left_bank_3_reg);
-            acc1 = vdotq_s32(acc1, filter_reg_2_a, left_bank_4_reg);
-            right_bank_4_reg = vld1q_s8(input_data + width_micro_stride +
-                                        4 * workspace_height_stride);
-
-            // Fixed-point multiplication.
-            acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-            acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc0, -output_shift);
-            acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-            acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc1, -output_shift);
-            // Add the output offset.
-            acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            acc_u8 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8 = util_vmax_x8(acc_u8, output_activation_min_vec);
-            acc_u8 = util_vmin_x8(acc_u8, output_activation_max_vec);
-
-            left_bank_0_reg = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_0_reg)));
-            left_bank_1_reg = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_1_reg)));
-            left_bank_2_reg = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_2_reg)));
-            left_bank_3_reg = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_3_reg)));
-            left_bank_4_reg = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_4_reg)));
-            acc0 = adjusted_bias_data;
-            acc1 = adjusted_bias_data;
-            vtrn1_s8x2_in_place(&left_bank_0_reg, &right_bank_0_reg);
-            vtrn1_s8x2_in_place(&left_bank_1_reg, &right_bank_1_reg);
-            vtrn1_s8x2_in_place(&left_bank_2_reg, &right_bank_2_reg);
-            vst1_lane_u8x4(output_data_base, acc_u8, 0);
-            vst1_lane_u8x4(output_data_base + output_height_stride, acc_u8, 1);
-
-            vtrn1_s8x2_in_place(&left_bank_3_reg, &right_bank_3_reg);
-            vtrn1_s8x2_in_place(&left_bank_4_reg, &right_bank_4_reg);
-
-            acc0 = vdotq_s32(acc0, filter_reg_0_a, left_bank_0_reg);
-            acc1 = vdotq_s32(acc1, filter_reg_0_a, left_bank_2_reg);
-            acc0 = vdotq_s32(acc0, filter_reg_1_a, left_bank_1_reg);
-            acc1 = vdotq_s32(acc1, filter_reg_1_a, left_bank_3_reg);
-            acc0 = vdotq_s32(acc0, filter_reg_2_a, left_bank_2_reg);
-            acc1 = vdotq_s32(acc1, filter_reg_2_a, left_bank_4_reg);
-
-            // Fixed-point multiplication.
-            acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-            acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc0, -output_shift);
-            acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-            acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc1, -output_shift);
-            // Add the output offset.
-            acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            acc_u8 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8 = util_vmax_x8(acc_u8, output_activation_min_vec);
-            acc_u8 = util_vmin_x8(acc_u8, output_activation_max_vec);
-
-            vst1_lane_u8x4(output_data_base + depth, acc_u8, 0);
-            vst1_lane_u8x4(output_data_base + depth + output_height_stride,
-                           acc_u8, 1);
-
-            left_bank_0_reg = right_bank_0_reg;
-            left_bank_1_reg = right_bank_1_reg;
-            left_bank_2_reg = right_bank_2_reg;
-            left_bank_3_reg = right_bank_3_reg;
-            left_bank_4_reg = right_bank_4_reg;
-          }
-          for (; i_width < output_width_overall_micro_repeats; ++i_width) {
-            TFLITE_DCHECK_NE(residual_width, kFourOverStride);
-
-            // No need to load next ("right") block of data.
-
-            typename QuantizationTypeImpl<quantization_type>::ExternalType*
-                output_data_base = output_data + depth * 2 * i_width + 4 * s;
-
-            // Iterate over input width shifts within 4x4 blocks.
-            {
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-
-              acc0 = vdotq_s32(acc0, filter_reg_0_a, left_bank_0_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_1_a, left_bank_1_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_2_a, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_0_a, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_1_a, left_bank_3_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_2_a, left_bank_4_reg);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc0, -output_shift);
-              acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc1, -output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              // Apply the activation function.
-              uint8x8_t acc_u8 = vqmovxn_s16(acc_s16_0_1);
-              acc_u8 = util_vmax_x8(acc_u8, output_activation_min_vec);
-              acc_u8 = util_vmin_x8(acc_u8, output_activation_max_vec);
-
-              vst1_lane_u8x4(output_data_base, acc_u8, 0);
-              vst1_lane_u8x4(output_data_base + output_height_stride, acc_u8,
-                             1);
-
-              left_bank_0_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_0_reg)));
-              left_bank_1_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_1_reg)));
-              left_bank_2_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_2_reg)));
-              left_bank_3_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_3_reg)));
-              left_bank_4_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_4_reg)));
-              vtrn1_s8x2_in_place(&left_bank_0_reg, &right_bank_0_reg);
-              vtrn1_s8x2_in_place(&left_bank_1_reg, &right_bank_1_reg);
-              vtrn1_s8x2_in_place(&left_bank_2_reg, &right_bank_2_reg);
-              vtrn1_s8x2_in_place(&left_bank_3_reg, &right_bank_3_reg);
-              vtrn1_s8x2_in_place(&left_bank_4_reg, &right_bank_4_reg);
-            }
-          }
-          bias_data += kBiasIncrement;
-        }
-      } else {
-        // block_height == 1.
-        int8x16_t filter_reg_0_a;
-        int8x16_t filter_reg_1_a;
-        int8x16_t filter_reg_2_a;
-        int8x16_t filter_reg_0_b;
-        int8x16_t filter_reg_1_b;
-        int8x16_t filter_reg_2_b;
-
-        filter_reg_0_a = vld1q_s8(filter_block);
-        filter_reg_1_a = vld1q_s8(filter_block + 32);
-        filter_reg_2_a = vld1q_s8(filter_block + 64);
-        filter_reg_0_b = vld1q_s8(filter_block + 16);
-        filter_reg_1_b = vld1q_s8(filter_block + 16 + 32);
-        filter_reg_2_b = vld1q_s8(filter_block + 16 + 64);
-
-        const int8_t* scratch_data =
-            scratch_block_data + depth_micro_stride * j_depth;
-        typename QuantizationTypeImpl<quantization_type>::ExternalType*
-            output_data = output_block_data + 8 * j_depth;
-        const int8_t* input_data_0 = scratch_data;
-
-        const int32x4_t adjusted_bias_data_a = vld1q_s32(bias_data);
-        bias_data += kBiasIncrement;
-        const int32x4_t adjusted_bias_data_b = vld1q_s32(bias_data);
-        bias_data += kBiasIncrement;
-
-        // Load first sub-micro block of data into operational banks.
-        int8x16_t left_bank_0_reg_a = vld1q_s8(input_data_0);
-        int8x16_t left_bank_1_reg_a =
-            vld1q_s8(input_data_0 + workspace_height_stride);
-        int8x16_t left_bank_2_reg_a =
-            vld1q_s8(input_data_0 + 2 * workspace_height_stride);
-        int8x16_t left_bank_0_reg_b = vld1q_s8(input_data_0 + 16);
-        int8x16_t left_bank_1_reg_b =
-            vld1q_s8(input_data_0 + workspace_height_stride + 16);
-        int8x16_t left_bank_2_reg_b =
-            vld1q_s8(input_data_0 + 2 * workspace_height_stride + 16);
-
-        int8x16_t right_bank_0_reg_a;
-        int8x16_t right_bank_1_reg_a;
-        int8x16_t right_bank_2_reg_a;
-        int8x16_t right_bank_0_reg_b;
-        int8x16_t right_bank_1_reg_b;
-        int8x16_t right_bank_2_reg_b;
-
-        int32x4_t acc0_a;
-        int32x4_t acc0_b;
-
-        for (int i_width = 0; i_width < output_width_overall_micro_repeats;
-             ++i_width) {
-          const int output_width = i_width == output_width_micro_repeats
-                                       ? residual_width
-                                       : kFourOverStride;
-          TFLITE_DCHECK_LE(output_width * kStrideVal, 4);
-          const int8_t* input_data =
-              input_data_0 + width_micro_stride * i_width;
-          const bool no_right_block = i_width == output_width_micro_repeats &&
-                                      output_width_overall_micro_repeats ==
-                                          workspace_width_micro_repeats;
-
-          if (!no_right_block) {
-            // Load next sub-micro block of data.
-            right_bank_0_reg_a = vld1q_s8(input_data + width_micro_stride);
-            right_bank_1_reg_a = vld1q_s8(input_data + width_micro_stride +
-                                          workspace_height_stride);
-            right_bank_2_reg_a = vld1q_s8(input_data + width_micro_stride +
-                                          2 * workspace_height_stride);
-            right_bank_0_reg_b = vld1q_s8(input_data + width_micro_stride + 16);
-            right_bank_1_reg_b = vld1q_s8(input_data + width_micro_stride +
-                                          workspace_height_stride + 16);
-            right_bank_2_reg_b = vld1q_s8(input_data + width_micro_stride +
-                                          2 * workspace_height_stride + 16);
-          }
-
-          typename QuantizationTypeImpl<quantization_type>::ExternalType*
-              output_data_base = output_data + depth * 2 * i_width;
-
-          // Iterate over input width shifts within 4x4 blocks.
-          {
-            acc0_a = adjusted_bias_data_a;
-            acc0_b = adjusted_bias_data_b;
-
-            acc0_a = vdotq_s32(acc0_a, filter_reg_0_a, left_bank_0_reg_a);
-            acc0_a = vdotq_s32(acc0_a, filter_reg_1_a, left_bank_1_reg_a);
-            acc0_a = vdotq_s32(acc0_a, filter_reg_2_a, left_bank_2_reg_a);
-            acc0_b = vdotq_s32(acc0_b, filter_reg_0_b, left_bank_0_reg_b);
-            acc0_b = vdotq_s32(acc0_b, filter_reg_1_b, left_bank_1_reg_b);
-            acc0_b = vdotq_s32(acc0_b, filter_reg_2_b, left_bank_2_reg_b);
-
-            // Fixed-point multiplication.
-            acc0_a = vqrdmulhq_n_s32(acc0_a, output_multiplier);
-            acc0_b = vqrdmulhq_n_s32(acc0_b, output_multiplier);
-            acc0_a = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc0_a, -output_shift);
-            acc0_b = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc0_b, -output_shift);
-            // Add the output offset.
-            int16x8_t acc_s16_0_1 =
-                vcombine_s16(vqmovn_s32(acc0_a), vqmovn_s32(acc0_b));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            uint8x8_t acc_u8 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8 = util_vmax_x8(acc_u8, output_activation_min_vec);
-            acc_u8 = util_vmin_x8(acc_u8, output_activation_max_vec);
-
-            util_vst1_u8(output_data_base, acc_u8);
-
-            left_bank_0_reg_a = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_0_reg_a)));
-            left_bank_1_reg_a = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_1_reg_a)));
-            left_bank_2_reg_a = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_2_reg_a)));
-            left_bank_0_reg_b = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_0_reg_b)));
-            left_bank_1_reg_b = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_1_reg_b)));
-            left_bank_2_reg_b = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_2_reg_b)));
-            vtrn1_s8x2_in_place(&left_bank_0_reg_a, &right_bank_0_reg_a);
-            vtrn1_s8x2_in_place(&left_bank_1_reg_a, &right_bank_1_reg_a);
-            vtrn1_s8x2_in_place(&left_bank_2_reg_a, &right_bank_2_reg_a);
-            vtrn1_s8x2_in_place(&left_bank_0_reg_b, &right_bank_0_reg_b);
-            vtrn1_s8x2_in_place(&left_bank_1_reg_b, &right_bank_1_reg_b);
-            vtrn1_s8x2_in_place(&left_bank_2_reg_b, &right_bank_2_reg_b);
-          }
-
-          if (output_width > 1) {
-            acc0_a = adjusted_bias_data_a;
-            acc0_b = adjusted_bias_data_b;
-
-            acc0_a = vdotq_s32(acc0_a, filter_reg_0_a, left_bank_0_reg_a);
-            acc0_a = vdotq_s32(acc0_a, filter_reg_1_a, left_bank_1_reg_a);
-            acc0_a = vdotq_s32(acc0_a, filter_reg_2_a, left_bank_2_reg_a);
-            acc0_b = vdotq_s32(acc0_b, filter_reg_0_b, left_bank_0_reg_b);
-            acc0_b = vdotq_s32(acc0_b, filter_reg_1_b, left_bank_1_reg_b);
-            acc0_b = vdotq_s32(acc0_b, filter_reg_2_b, left_bank_2_reg_b);
-
-            // Fixed-point multiplication.
-            acc0_a = vqrdmulhq_n_s32(acc0_a, output_multiplier);
-            acc0_b = vqrdmulhq_n_s32(acc0_b, output_multiplier);
-            acc0_a = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc0_a, -output_shift);
-            acc0_b = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc0_b, -output_shift);
-            // Add the output offset.
-            int16x8_t acc_s16_0_1 =
-                vcombine_s16(vqmovn_s32(acc0_a), vqmovn_s32(acc0_b));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            uint8x8_t acc_u8 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8 = util_vmax_x8(acc_u8, output_activation_min_vec);
-            acc_u8 = util_vmin_x8(acc_u8, output_activation_max_vec);
-
-            util_vst1_u8(output_data_base + depth, acc_u8);
-
-            left_bank_0_reg_a = right_bank_0_reg_a;
-            left_bank_1_reg_a = right_bank_1_reg_a;
-            left_bank_2_reg_a = right_bank_2_reg_a;
-            left_bank_0_reg_b = right_bank_0_reg_b;
-            left_bank_1_reg_b = right_bank_1_reg_b;
-            left_bank_2_reg_b = right_bank_2_reg_b;
-          }
-        }
-      }
-    }
-  }  // NOLINT(readability/fn_size) Manually unrolled.
-
-  static inline void Run(const int8_t* scratch_block_data,
-                         const int8_t* filter_workspace,
-                         const int32_t* bias_data, uint8_t* output_block_data,
-                         const DepthwiseConvDotProdParams* function_params) {
-    KernelMacroBlockIntrinsics(scratch_block_data, filter_workspace, bias_data,
-                               output_block_data, function_params);
-  }
-};
-
-template <>
-struct KernelMacroBlock<
-    DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct,
-    QuantizationType::kNonPerChannelUint8,
-    DepthwiseConvDepthMultiplication::kUnitInputDepth,
-    /*stride=*/1> {
-  static inline uint8x8_t vqmovxn_s16(int16x8_t x) { return vqmovun_s16(x); }
-  static inline uint8x8_t util_vmin_x8(uint8x8_t a, uint8x8_t b) {
-    return vmin_u8(a, b);
-  }
-  static inline uint8x8_t util_vmax_x8(uint8x8_t a, uint8x8_t b) {
-    return vmax_u8(a, b);
-  }
-  static inline uint8x16_t util_vminq_x8(uint8x16_t a, uint8x16_t b) {
-    return vminq_u8(a, b);
-  }
-  static inline uint8x16_t util_vmaxq_x8(uint8x16_t a, uint8x16_t b) {
-    return vmaxq_u8(a, b);
-  }
-
-  static inline void KernelMacroBlockIntrinsics(
-      const int8_t* scratch_block_data, const int8_t* filter_workspace,
-      const int32_t* bias_data, uint8_t* output_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-    static constexpr QuantizationType quantization_type =
-        QuantizationType::kNonPerChannelUint8;
-
-    TFLITE_DCHECK_EQ(function_params->stride, 1);
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-    const int output_width_micro_repeats =
-        function_params->output_width_micro_repeats;
-    const int depth_micro_repeats = function_params->depth_micro_repeats;
-    const int output_depth = function_params->output_depth;
-
-    const int output_width_overall_micro_repeats =
-        function_params->output_width_overall_micro_repeats;
-    const int block_height = function_params->outbound_block_height;
-    const int residual_width = function_params->output_residual_width;
-    const int output_height_stride = function_params->output_height_stride;
-    constexpr int kBiasIncrement = 4;
-
-    TFLITE_DCHECK(depth_micro_repeats > 0);
-
-    const int32_t output_activation_min =
-        function_params->quantized_activation_min;
-    const int32_t output_activation_max =
-        function_params->quantized_activation_max;
-    const int32_t output_multiplier = function_params->output_multiplier;
-    const int32_t output_shift = function_params->output_shift;
-    const int32_t output_offset = function_params->output_offset;
-    if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-      TFLITE_DCHECK_GE(output_activation_min, 0);
-      TFLITE_DCHECK_LT(output_activation_min, 256);
-      TFLITE_DCHECK_GE(output_activation_max, 0);
-      TFLITE_DCHECK_LT(output_activation_max, 256);
-    } else {
-      TFLITE_DCHECK_GE(output_activation_min, -128);
-      TFLITE_DCHECK_LT(output_activation_min, 128);
-      TFLITE_DCHECK_GE(output_activation_max, -128);
-      TFLITE_DCHECK_LT(output_activation_max, 128);
-    }
-    TFLITE_DCHECK_GE(output_offset, -32878);
-    TFLITE_DCHECK_LT(output_offset, 32768);
-
-    const int16x8_t output_offset_vec =
-        vdupq_n_s16(static_cast<int16_t>(output_offset));
-    const uint8x16_t output_activation_min_vec =
-        vdupq_n_u8(static_cast<uint8_t>(output_activation_min));
-    const uint8x16_t output_activation_max_vec =
-        vdupq_n_u8(static_cast<uint8_t>(output_activation_max));
-
-    typename QuantizationTypeImpl<quantization_type>::ExternalType*
-        output_data_depthwise = output_block_data;
-    for (int j_depth = 0; j_depth < depth_micro_repeats; ++j_depth) {
-      // Simulate NEON-register transposition of subset of filter.
-      int8x16_t filter_reg_0_a;
-      int8x16_t filter_reg_0_b;
-      int8x16_t filter_reg_1_a;
-      int8x16_t filter_reg_1_b;
-      int8x16_t filter_reg_2_a;
-      int8x16_t filter_reg_2_b;
-      int8x16_t filter_reg_0_a_shifted;
-      int8x16_t filter_reg_1_a_shifted;
-      int8x16_t filter_reg_2_a_shifted;
-
-      filter_reg_0_a = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_0_b = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_1_a = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_1_b = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_2_a = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_2_b = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-
-      filter_reg_0_a_shifted = vreinterpretq_s8_u32(
-          vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_0_a), 8));
-      filter_reg_1_a_shifted = vreinterpretq_s8_u32(
-          vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_1_a), 8));
-      filter_reg_2_a_shifted = vreinterpretq_s8_u32(
-          vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_2_a), 8));
-
-      // When output_width_micro_repeats < output_width_overall_micro_repeats,
-      // 0 < residual_width <= 2, and so residual_width == 1 is then true iff
-      // residual_width < 2.
-      const int adjusted_width_micro_repeats =
-          (output_width_micro_repeats < output_width_overall_micro_repeats) &&
-                  (residual_width < 4)
-              ? output_width_micro_repeats
-              : output_width_overall_micro_repeats;
-
-      if (block_height == 4) {
-        for (int s = 0; s < 2; ++s) {
-          // Work through one slice, by row, at a time.
-          typename QuantizationTypeImpl<quantization_type>::ExternalType*
-              output_data_base = output_data_depthwise + 4 * s;
-
-          const int8_t* next_input_data = scratch_block_data;
-          typename QuantizationTypeImpl<quantization_type>::ExternalType*
-              output_data = output_data_base;
-
-          const int32x4_t adjusted_bias_data = vld1q_s32(bias_data);
-          bias_data += kBiasIncrement;
-
-          int8x16_t input_bank_a_reg;  //  left 0, right 0, left 1, right 1.
-          int8x16_t input_bank_b_reg;  //  left 2, right 2, left 3, right 3.
-          int8x16_t input_bank_c_reg;  //  left 4, right 4, left 5, right 5.
-
-          // Load first sub-micro block of data into operational banks.
-          input_bank_a_reg =
-              vld1q_dup_s8x4(next_input_data);  // Load lane 0, avoiding
-                                                // uninitialized variable.
-          input_bank_a_reg = vld1q_lane_8x4(
-              next_input_data + workspace_height_stride, input_bank_a_reg, 2);
-          input_bank_b_reg = vld1q_dup_s8x4(
-              next_input_data +
-              2 * workspace_height_stride);  // Load lane 0, avoiding
-                                             // uninitialized variable.
-          input_bank_b_reg =
-              vld1q_lane_8x4(next_input_data + 3 * workspace_height_stride,
-                             input_bank_b_reg, 2);
-          input_bank_c_reg = vld1q_dup_s8x4(
-              next_input_data +
-              4 * workspace_height_stride);  // Load lane 0, avoiding
-                                             // uninitialized variable.
-          input_bank_c_reg =
-              vld1q_lane_8x4(next_input_data + 5 * workspace_height_stride,
-                             input_bank_c_reg, 2);
-
-          int32x4_t acc0;
-          int32x4_t acc1;
-          int32x4_t acc2;
-          int32x4_t acc3;
-
-          acc0 = adjusted_bias_data;
-          acc1 = adjusted_bias_data;
-          acc2 = adjusted_bias_data;
-          acc3 = adjusted_bias_data;
-
-          acc0 = vdotq_four_lane_s32(acc0, filter_reg_2_a, input_bank_b_reg, 0);
-          acc1 = vdotq_four_lane_s32(acc1, filter_reg_1_a, input_bank_b_reg, 0);
-          acc2 = vdotq_four_lane_s32(acc2, filter_reg_0_a, input_bank_b_reg, 0);
-          acc3 = vdotq_four_lane_s32(acc3, filter_reg_0_a, input_bank_b_reg, 2);
-
-          int i_width = 0;
-          for (; i_width < adjusted_width_micro_repeats; ++i_width) {
-            next_input_data += 4;
-
-            // Iterate over input width shifts within 4x4 blocks.
-            {
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_0_a, input_bank_a_reg,
-                                         0);
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_1_a, input_bank_a_reg,
-                                         2);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_0_a, input_bank_a_reg,
-                                         2);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_2_a, input_bank_b_reg,
-                                         2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_1_a, input_bank_b_reg,
-                                         2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_2_a, input_bank_c_reg,
-                                         0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_1_a, input_bank_c_reg,
-                                         0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_2_a, input_bank_c_reg,
-                                         2);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc0, -output_shift);
-              acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc1, -output_shift);
-              acc2 = vqrdmulhq_n_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc2, -output_shift);
-              acc3 = vqrdmulhq_n_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc3, -output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              uint8x16_t acc_u8_all = vcombine_u8(vqmovxn_s16(acc_s16_0_1),
-                                                  vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_u8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_u8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_u8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_u8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              output_data += output_depth;
-            }
-            // Load next sub-micro block of data.
-            input_bank_a_reg =
-                vld1q_lane_8x4(next_input_data, input_bank_a_reg, 1);
-            input_bank_a_reg = vld1q_lane_8x4(
-                next_input_data + workspace_height_stride, input_bank_a_reg, 3);
-            input_bank_b_reg =
-                vld1q_lane_8x4(next_input_data + 2 * workspace_height_stride,
-                               input_bank_b_reg, 1);
-            input_bank_b_reg =
-                vld1q_lane_8x4(next_input_data + 3 * workspace_height_stride,
-                               input_bank_b_reg, 3);
-            input_bank_c_reg =
-                vld1q_lane_8x4(next_input_data + 4 * workspace_height_stride,
-                               input_bank_c_reg, 1);
-            input_bank_c_reg =
-                vld1q_lane_8x4(next_input_data + 5 * workspace_height_stride,
-                               input_bank_c_reg, 3);
-
-            {
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_0_a_shifted,
-                                         input_bank_a_reg, 0);
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_1_a_shifted,
-                                         input_bank_a_reg, 2);
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_2_a_shifted,
-                                         input_bank_b_reg, 0);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_0_a_shifted,
-                                         input_bank_a_reg, 2);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_1_a_shifted,
-                                         input_bank_b_reg, 0);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_2_a_shifted,
-                                         input_bank_b_reg, 2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_0_a_shifted,
-                                         input_bank_b_reg, 0);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_1_a_shifted,
-                                         input_bank_b_reg, 2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_2_a_shifted,
-                                         input_bank_c_reg, 0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_0_a_shifted,
-                                         input_bank_b_reg, 2);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_1_a_shifted,
-                                         input_bank_c_reg, 0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_2_a_shifted,
-                                         input_bank_c_reg, 2);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc0, -output_shift);
-              acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc1, -output_shift);
-              acc2 = vqrdmulhq_n_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc2, -output_shift);
-              acc3 = vqrdmulhq_n_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc3, -output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              uint8x16_t acc_u8_all = vcombine_u8(vqmovxn_s16(acc_s16_0_1),
-                                                  vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_u8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_u8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_u8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_u8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              input_bank_a_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_a_reg), 16));
-              input_bank_b_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_b_reg), 16));
-              input_bank_c_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_c_reg), 16));
-
-              output_data += output_depth;
-            }
-
-            {
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_0_a, input_bank_a_reg,
-                                         0);
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_1_a, input_bank_a_reg,
-                                         2);
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_2_a, input_bank_b_reg,
-                                         0);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_0_a, input_bank_a_reg,
-                                         2);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_1_a, input_bank_b_reg,
-                                         0);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_2_a, input_bank_b_reg,
-                                         2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_0_a, input_bank_b_reg,
-                                         0);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_1_a, input_bank_b_reg,
-                                         2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_2_a, input_bank_c_reg,
-                                         0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_0_a, input_bank_b_reg,
-                                         2);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_1_a, input_bank_c_reg,
-                                         0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_2_a, input_bank_c_reg,
-                                         2);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc0, -output_shift);
-              acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc1, -output_shift);
-              acc2 = vqrdmulhq_n_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc2, -output_shift);
-              acc3 = vqrdmulhq_n_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc3, -output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              uint8x16_t acc_u8_all = vcombine_u8(vqmovxn_s16(acc_s16_0_1),
-                                                  vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_u8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_u8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_u8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_u8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              output_data += output_depth;
-            }
-
-            {
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_0_a_shifted,
-                                         input_bank_a_reg, 0);
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_1_a_shifted,
-                                         input_bank_a_reg, 2);
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_2_a_shifted,
-                                         input_bank_b_reg, 0);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_0_a_shifted,
-                                         input_bank_a_reg, 2);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_1_a_shifted,
-                                         input_bank_b_reg, 0);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_2_a_shifted,
-                                         input_bank_b_reg, 2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_0_a_shifted,
-                                         input_bank_b_reg, 0);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_1_a_shifted,
-                                         input_bank_b_reg, 2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_2_a_shifted,
-                                         input_bank_c_reg, 0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_0_a_shifted,
-                                         input_bank_b_reg, 2);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_1_a_shifted,
-                                         input_bank_c_reg, 0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_2_a_shifted,
-                                         input_bank_c_reg, 2);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc0, -output_shift);
-              acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc1, -output_shift);
-              acc2 = vqrdmulhq_n_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc2, -output_shift);
-              acc3 = vqrdmulhq_n_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc3, -output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              uint8x16_t acc_u8_all = vcombine_u8(vqmovxn_s16(acc_s16_0_1),
-                                                  vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_u8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_u8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_u8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_u8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              input_bank_a_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_a_reg), 16));
-              input_bank_b_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_b_reg), 16));
-              input_bank_c_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_c_reg), 16));
-
-              output_data += output_depth;
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_2_a, input_bank_b_reg,
-                                         0);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_1_a, input_bank_b_reg,
-                                         0);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_0_a, input_bank_b_reg,
-                                         0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_0_a, input_bank_b_reg,
-                                         2);
-            }
-          }
-
-          if (i_width < output_width_overall_micro_repeats) {
-            next_input_data += 4;
-            const int output_width = residual_width;
-
-            // Load next sub-micro block of data.
-            input_bank_a_reg =
-                vld1q_lane_8x4(next_input_data, input_bank_a_reg, 1);
-            input_bank_a_reg = vld1q_lane_8x4(
-                next_input_data + workspace_height_stride, input_bank_a_reg, 3);
-            input_bank_b_reg =
-                vld1q_lane_8x4(next_input_data + 2 * workspace_height_stride,
-                               input_bank_b_reg, 1);
-            input_bank_b_reg =
-                vld1q_lane_8x4(next_input_data + 3 * workspace_height_stride,
-                               input_bank_b_reg, 3);
-            input_bank_c_reg =
-                vld1q_lane_8x4(next_input_data + 4 * workspace_height_stride,
-                               input_bank_c_reg, 1);
-            input_bank_c_reg =
-                vld1q_lane_8x4(next_input_data + 5 * workspace_height_stride,
-                               input_bank_c_reg, 3);
-
-            // Iterate over input width shifts within 4x4 blocks.
-            for (int x = 0; x < output_width; ++x) {
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_0_a, input_bank_a_reg,
-                                         0);
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_1_a, input_bank_a_reg,
-                                         2);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_0_a, input_bank_a_reg,
-                                         2);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_2_a, input_bank_b_reg,
-                                         2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_1_a, input_bank_b_reg,
-                                         2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_2_a, input_bank_c_reg,
-                                         0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_1_a, input_bank_c_reg,
-                                         0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_2_a, input_bank_c_reg,
-                                         2);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc0, -output_shift);
-              acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc1, -output_shift);
-              acc2 = vqrdmulhq_n_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc2, -output_shift);
-              acc3 = vqrdmulhq_n_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc3, -output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              uint8x16_t acc_u8_all = vcombine_u8(vqmovxn_s16(acc_s16_0_1),
-                                                  vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_u8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_u8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_u8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_u8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              input_bank_a_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_a_reg), 8));
-              input_bank_b_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_b_reg), 8));
-              input_bank_c_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_c_reg), 8));
-
-              output_data += output_depth;
-
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_2_a, input_bank_b_reg,
-                                         0);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_1_a, input_bank_b_reg,
-                                         0);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_0_a, input_bank_b_reg,
-                                         0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_0_a, input_bank_b_reg,
-                                         2);
-            }
-          }
-          // scratch_block_data += 4 * workspace_height_stride;
-          output_data_base += 4 * output_height_stride;
-
-          // Move to next sub-block: advance to second set of filters, to new
-          // bias.
-          filter_reg_0_a = filter_reg_0_b;
-          filter_reg_1_a = filter_reg_1_b;
-          filter_reg_2_a = filter_reg_2_b;
-          filter_reg_0_a_shifted = vreinterpretq_s8_u32(
-              vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_0_a), 8));
-          filter_reg_1_a_shifted = vreinterpretq_s8_u32(
-              vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_1_a), 8));
-          filter_reg_2_a_shifted = vreinterpretq_s8_u32(
-              vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_2_a), 8));
-        }
-      } else {
-        // Block height < 4.
-        typename QuantizationTypeImpl<quantization_type>::ExternalType*
-            output_data_base = output_data_depthwise;
-
-        const int32x4_t adjusted_bias_data_a = vld1q_s32(bias_data);
-        bias_data += kBiasIncrement;
-        const int32x4_t adjusted_bias_data_b = vld1q_s32(bias_data);
-        bias_data += kBiasIncrement;
-
-        for (int k_height = 0; k_height < block_height; ++k_height) {
-          const int8_t* next_input_data =
-              scratch_block_data + k_height * workspace_height_stride;
-          typename QuantizationTypeImpl<quantization_type>::ExternalType*
-              output_data = output_data_base;
-
-          int8x16_t input_bank_p_reg;  //  left 0, right 0, left 1, right 1.
-          int8x16_t input_bank_q_reg;  //  left 2, right 2, left 3, right 3.
-
-          // Load first sub-micro block of data into operational banks.
-          input_bank_p_reg =
-              vld1q_dup_s8x4(next_input_data);  // Load lane 0, avoiding
-                                                // uninitialized variable.
-          input_bank_p_reg = vld1q_lane_8x4(
-              next_input_data + workspace_height_stride, input_bank_p_reg, 2);
-          input_bank_q_reg = vld1q_dup_s8x4(
-              next_input_data +
-              2 * workspace_height_stride);  // Load lane 0, avoiding
-                                             // uninitialized variable.
-
-          for (int i_width = 0; i_width < output_width_overall_micro_repeats;
-               ++i_width) {
-            next_input_data += 4;
-            const int output_width =
-                i_width == output_width_micro_repeats ? residual_width : 4;
-
-            // Load next sub-micro block of data.
-            input_bank_p_reg =
-                vld1q_lane_8x4(next_input_data, input_bank_p_reg, 1);
-            input_bank_p_reg = vld1q_lane_8x4(
-                next_input_data + workspace_height_stride, input_bank_p_reg, 3);
-            input_bank_q_reg =
-                vld1q_lane_8x4(next_input_data + 2 * workspace_height_stride,
-                               input_bank_q_reg, 1);
-            // Iterate over input width shifts within 4x4 blocks.
-            for (int x = 0; x < output_width; ++x) {
-              int32x4_t acc_a = adjusted_bias_data_a;
-              int32x4_t acc_b = adjusted_bias_data_b;
-              acc_a = vdotq_four_lane_s32(acc_a, filter_reg_0_a,
-                                          input_bank_p_reg, 0);
-              acc_a = vdotq_four_lane_s32(acc_a, filter_reg_1_a,
-                                          input_bank_p_reg, 2);
-              acc_a = vdotq_four_lane_s32(acc_a, filter_reg_2_a,
-                                          input_bank_q_reg, 0);
-              acc_b = vdotq_four_lane_s32(acc_b, filter_reg_0_b,
-                                          input_bank_p_reg, 0);
-              acc_b = vdotq_four_lane_s32(acc_b, filter_reg_1_b,
-                                          input_bank_p_reg, 2);
-              acc_b = vdotq_four_lane_s32(acc_b, filter_reg_2_b,
-                                          input_bank_q_reg, 0);
-
-              // Fixed-point multiplication.
-              acc_a = vqrdmulhq_n_s32(acc_a, output_multiplier);
-              acc_b = vqrdmulhq_n_s32(acc_b, output_multiplier);
-              acc_a = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc_a, -output_shift);
-              acc_b = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                  acc_b, -output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_0 =
-                  vcombine_s16(vqmovn_s32(acc_a), vqmovn_s32(acc_b));
-              acc_s16_0_0 = vqaddq_s16(acc_s16_0_0, output_offset_vec);
-              // Apply the activation function.
-              uint8x8_t acc_u8_0_0 = vqmovxn_s16(acc_s16_0_0);
-              acc_u8_0_0 = util_vmax_x8(acc_u8_0_0,
-                                        vget_low_u8(output_activation_min_vec));
-              acc_u8_0_0 = util_vmin_x8(acc_u8_0_0,
-                                        vget_low_u8(output_activation_max_vec));
-
-              util_vst1_u8(output_data, acc_u8_0_0);
-
-              input_bank_p_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_p_reg), 8));
-              input_bank_q_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_q_reg), 8));
-
-              output_data += output_depth;
-            }
-          }
-          output_data_base += output_height_stride;
-        }
-      }
-      output_data_depthwise += 8;
-    }
-  }  // NOLINT(readability/fn_size) Manually unrolled.
-
-  static inline void Run(const int8_t* scratch_block_data,
-                         const int8_t* filter_workspace,
-                         const int32_t* bias_data, uint8_t* output_block_data,
-                         const DepthwiseConvDotProdParams* function_params) {
-    KernelMacroBlockIntrinsics(scratch_block_data, filter_workspace, bias_data,
-                               output_block_data, function_params);
-  }
-};
-
-template <>
-struct KernelMacroBlock<
-    DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct,
-    QuantizationType::kNonPerChannelUint8,
-    DepthwiseConvDepthMultiplication::kUnitInputDepth,
-    /*stride=*/2> {
-  static inline uint8x8_t vqmovxn_s16(int16x8_t x) { return vqmovun_s16(x); }
-  static inline uint8x8_t util_vmin_x8(uint8x8_t a, uint8x8_t b) {
-    return vmin_u8(a, b);
-  }
-  static inline uint8x8_t util_vmax_x8(uint8x8_t a, uint8x8_t b) {
-    return vmax_u8(a, b);
-  }
-
-  static inline void KernelMacroBlockIntrinsics(
-      const int8_t* scratch_block_data, const int8_t* filter_workspace,
-      const int32_t* bias_data, uint8_t* output_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-    static constexpr QuantizationType quantization_type =
-        QuantizationType::kNonPerChannelUint8;
-
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-    const int output_width_micro_repeats =
-        function_params->output_width_micro_repeats;
-    const int depth_micro_repeats = function_params->depth_micro_repeats;
-    const int output_depth = function_params->output_depth;
-    constexpr int kStrideVal = 2;
-    TFLITE_DCHECK_EQ(function_params->stride, kStrideVal);
-
-    const int output_width_overall_micro_repeats =
-        function_params->output_width_overall_micro_repeats;
-    const int block_height = function_params->outbound_block_height;
-    const int residual_width = function_params->output_residual_width;
-    const int output_height_stride = function_params->output_height_stride;
-    constexpr int kBiasIncrement = 4;
-
-    const int32_t output_activation_min =
-        function_params->quantized_activation_min;
-    const int32_t output_activation_max =
-        function_params->quantized_activation_max;
-    const int32_t output_multiplier = function_params->output_multiplier;
-    const int32_t output_shift = function_params->output_shift;
-    const int32_t output_offset = function_params->output_offset;
-    if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-      TFLITE_DCHECK_GE(output_activation_min, 0);
-      TFLITE_DCHECK_LT(output_activation_min, 256);
-      TFLITE_DCHECK_GE(output_activation_max, 0);
-      TFLITE_DCHECK_LT(output_activation_max, 256);
-    } else {
-      TFLITE_DCHECK_GE(output_activation_min, -128);
-      TFLITE_DCHECK_LT(output_activation_min, 128);
-      TFLITE_DCHECK_GE(output_activation_max, -128);
-      TFLITE_DCHECK_LT(output_activation_max, 128);
-    }
-    TFLITE_DCHECK_GE(output_offset, -32878);
-    TFLITE_DCHECK_LT(output_offset, 32768);
-
-    TFLITE_DCHECK_GE(depth_micro_repeats, 1);
-
-    const int16x8_t output_offset_vec =
-        vdupq_n_s16(static_cast<int16_t>(output_offset));
-    const uint8x16_t output_activation_min_vec =
-        vdupq_n_u8(static_cast<uint8_t>(output_activation_min));
-    const uint8x16_t output_activation_max_vec =
-        vdupq_n_u8(static_cast<uint8_t>(output_activation_max));
-
-    for (int j_depth = 0; j_depth < (depth_micro_repeats * 1 + 0); ++j_depth) {
-      int8x16_t filter_reg_0_a;
-      int8x16_t filter_reg_0_b;
-      int8x16_t filter_reg_1_a;
-      int8x16_t filter_reg_1_b;
-      int8x16_t filter_reg_2_a;
-      int8x16_t filter_reg_2_b;
-
-      filter_reg_0_a = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_0_b = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_1_a = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_1_b = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_2_a = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_2_b = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-
-      const int32x4_t adjusted_bias_data_s_0 = vld1q_s32(bias_data);
-      bias_data += kBiasIncrement;
-      const int32x4_t adjusted_bias_data_s_1 = vld1q_s32(bias_data);
-      bias_data += kBiasIncrement;
-
-      if (block_height == 2) {
-        const int8_t* scratch_data = scratch_block_data;
-        typename QuantizationTypeImpl<quantization_type>::ExternalType*
-            output_data = output_block_data + 8 * j_depth;
-
-        int8x16_t input_bank_a_reg;  //  left 0, right 0, left 1, right 1.
-        int8x16_t input_bank_b_reg;  //  left 2, right 2, left 3, right 3.
-        int8x16_t input_bank_c_reg;  //  left 4, right 4, xxx, xxx.
-
-        // Load first sub-micro block of data into operational banks.
-        input_bank_a_reg =
-            vld1q_dup_s8x4(scratch_data);  // Load lane 0, avoiding
-                                           // uninitialized variable.
-        input_bank_a_reg = vld1q_lane_8x4(
-            scratch_data + workspace_height_stride, input_bank_a_reg, 2);
-        input_bank_b_reg = vld1q_dup_s8x4(
-            scratch_data +
-            2 * workspace_height_stride);  // Load lane 0, avoiding
-                                           // uninitialized variable.
-        input_bank_b_reg = vld1q_lane_8x4(
-            scratch_data + 3 * workspace_height_stride, input_bank_b_reg, 2);
-        input_bank_c_reg = vld1q_dup_s8x4(
-            scratch_data +
-            4 * workspace_height_stride);  // Load lane 0, avoiding
-                                           // uninitialized variable.
-
-        int32x4_t acc0;
-        int32x4_t acc1;
-
-        // When output_width_micro_repeats < output_width_overall_micro_repeats,
-        // 0 < residual_width <= 2, and so residual_width == 1 is then true iff
-        // residual_width < 2.
-        const int adjusted_width_micro_repeats =
-            (output_width_micro_repeats < output_width_overall_micro_repeats) &&
-                    (residual_width < 2)
-                ? output_width_micro_repeats
-                : output_width_overall_micro_repeats;
-
-        int i_width = 0;
-        for (; i_width < adjusted_width_micro_repeats; ++i_width) {
-          const int8_t* input_data = scratch_data + 4 + 4 * i_width;
-
-          // Load next sub-micro block of data.
-          input_bank_a_reg = vld1q_lane_8x4(input_data, input_bank_a_reg, 1);
-          input_bank_a_reg = vld1q_lane_8x4(
-              input_data + workspace_height_stride, input_bank_a_reg, 3);
-          input_bank_b_reg = vld1q_lane_8x4(
-              input_data + 2 * workspace_height_stride, input_bank_b_reg, 1);
-          input_bank_b_reg = vld1q_lane_8x4(
-              input_data + 3 * workspace_height_stride, input_bank_b_reg, 3);
-          input_bank_c_reg = vld1q_lane_8x4(
-              input_data + 4 * workspace_height_stride, input_bank_c_reg, 1);
-
-          int16x8_t acc_s16_0_1;
-          uint8x8_t acc_u8_0_1;
-          // Iterate over input width shifts within 4x4 blocks.
-          {
-            acc0 = adjusted_bias_data_s_0;
-            acc1 = adjusted_bias_data_s_0;
-
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_0_a, input_bank_a_reg, 0);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_1_a, input_bank_a_reg, 2);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_2_a, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_0_a, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_1_a, input_bank_b_reg, 2);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_2_a, input_bank_c_reg, 0);
-
-            // Fixed-point multiplication.
-            acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-            acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc0, -output_shift);
-            acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-            acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc1, -output_shift);
-            // Add the output offset.
-            acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            acc_u8_0_1 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8_0_1 = util_vmax_x8(acc_u8_0_1,
-                                      vget_low_u8(output_activation_min_vec));
-            acc_u8_0_1 = util_vmin_x8(acc_u8_0_1,
-                                      vget_low_u8(output_activation_max_vec));
-
-            vst1_lane_u8x4(output_data, acc_u8_0_1, 0);
-            vst1_lane_u8x4(output_data + output_height_stride, acc_u8_0_1, 1);
-
-            acc0 = adjusted_bias_data_s_1;
-            acc1 = adjusted_bias_data_s_1;
-
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_0_b, input_bank_a_reg, 0);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_1_b, input_bank_a_reg, 2);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_2_b, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_0_b, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_1_b, input_bank_b_reg, 2);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_2_b, input_bank_c_reg, 0);
-
-            // Fixed-point multiplication.
-            acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-            acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc0, -output_shift);
-            acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-            acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc1, -output_shift);
-            // Add the output offset.
-            acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            acc_u8_0_1 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8_0_1 = util_vmax_x8(acc_u8_0_1,
-                                      vget_low_u8(output_activation_min_vec));
-            acc_u8_0_1 = util_vmin_x8(acc_u8_0_1,
-                                      vget_low_u8(output_activation_max_vec));
-
-            vst1_lane_u8x4(output_data + 4, acc_u8_0_1, 0);
-            vst1_lane_u8x4(output_data + 4 + output_height_stride, acc_u8_0_1,
-                           1);
-
-            input_bank_a_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_a_reg), 16));
-            input_bank_b_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_b_reg), 16));
-            input_bank_c_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_c_reg), 16));
-
-            output_data += output_depth;
-          }
-
-          // output_width == four_over_stride.
-          acc0 = adjusted_bias_data_s_0;
-          acc1 = adjusted_bias_data_s_0;
-
-          acc0 = vdotq_four_lane_s32(acc0, filter_reg_0_a, input_bank_a_reg, 0);
-          acc0 = vdotq_four_lane_s32(acc0, filter_reg_1_a, input_bank_a_reg, 2);
-          acc0 = vdotq_four_lane_s32(acc0, filter_reg_2_a, input_bank_b_reg, 0);
-          acc1 = vdotq_four_lane_s32(acc1, filter_reg_0_a, input_bank_b_reg, 0);
-          acc1 = vdotq_four_lane_s32(acc1, filter_reg_1_a, input_bank_b_reg, 2);
-          acc1 = vdotq_four_lane_s32(acc1, filter_reg_2_a, input_bank_c_reg, 0);
-
-          // Fixed-point multiplication.
-          acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-          acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-              acc0, -output_shift);
-          acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-          acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-              acc1, -output_shift);
-          // Add the output offset.
-          acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-          acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-          // Apply the activation function.
-          acc_u8_0_1 = vqmovxn_s16(acc_s16_0_1);
-          acc_u8_0_1 =
-              util_vmax_x8(acc_u8_0_1, vget_low_u8(output_activation_min_vec));
-          acc_u8_0_1 =
-              util_vmin_x8(acc_u8_0_1, vget_low_u8(output_activation_max_vec));
-
-          vst1_lane_u8x4(output_data, acc_u8_0_1, 0);
-          vst1_lane_u8x4(output_data + output_height_stride, acc_u8_0_1, 1);
-
-          acc0 = adjusted_bias_data_s_1;
-          acc1 = adjusted_bias_data_s_1;
-
-          acc0 = vdotq_four_lane_s32(acc0, filter_reg_0_b, input_bank_a_reg, 0);
-          acc0 = vdotq_four_lane_s32(acc0, filter_reg_1_b, input_bank_a_reg, 2);
-          acc0 = vdotq_four_lane_s32(acc0, filter_reg_2_b, input_bank_b_reg, 0);
-          acc1 = vdotq_four_lane_s32(acc1, filter_reg_0_b, input_bank_b_reg, 0);
-          acc1 = vdotq_four_lane_s32(acc1, filter_reg_1_b, input_bank_b_reg, 2);
-          acc1 = vdotq_four_lane_s32(acc1, filter_reg_2_b, input_bank_c_reg, 0);
-
-          // Fixed-point multiplication.
-          acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-          acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-              acc0, -output_shift);
-          acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-          acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-              acc1, -output_shift);
-          // Add the output offset.
-          acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-          acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-          // Apply the activation function.
-          acc_u8_0_1 = vqmovxn_s16(acc_s16_0_1);
-          acc_u8_0_1 =
-              util_vmax_x8(acc_u8_0_1, vget_low_u8(output_activation_min_vec));
-          acc_u8_0_1 =
-              util_vmin_x8(acc_u8_0_1, vget_low_u8(output_activation_max_vec));
-
-          vst1_lane_u8x4(output_data + 4, acc_u8_0_1, 0);
-          vst1_lane_u8x4(output_data + 4 + output_height_stride, acc_u8_0_1, 1);
-
-          input_bank_a_reg = vreinterpretq_s8_u64(
-              vshrq_n_u64(vreinterpretq_u64_s8(input_bank_a_reg), 8));
-          input_bank_b_reg = vreinterpretq_s8_u64(
-              vshrq_n_u64(vreinterpretq_u64_s8(input_bank_b_reg), 8));
-          input_bank_c_reg = vreinterpretq_s8_u64(
-              vshrq_n_u64(vreinterpretq_u64_s8(input_bank_c_reg), 8));
-
-          output_data += output_depth;
-        }
-        for (; i_width < output_width_overall_micro_repeats; ++i_width) {
-          // output_width == 1.
-          const int8_t* input_data = scratch_data + 4 + 4 * i_width;
-
-          // Load next sub-micro block of data.
-          input_bank_a_reg = vld1q_lane_8x4(input_data, input_bank_a_reg, 1);
-          input_bank_a_reg = vld1q_lane_8x4(
-              input_data + workspace_height_stride, input_bank_a_reg, 3);
-          input_bank_b_reg = vld1q_lane_8x4(
-              input_data + 2 * workspace_height_stride, input_bank_b_reg, 1);
-          input_bank_b_reg = vld1q_lane_8x4(
-              input_data + 3 * workspace_height_stride, input_bank_b_reg, 3);
-          input_bank_c_reg = vld1q_lane_8x4(
-              input_data + 4 * workspace_height_stride, input_bank_c_reg, 1);
-
-          int16x8_t acc_s16_0_1;
-          uint8x8_t acc_u8_0_1;
-          // Iterate over input width shifts within 4x4 blocks.
-          {
-            acc0 = adjusted_bias_data_s_0;
-            acc1 = adjusted_bias_data_s_0;
-
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_0_a, input_bank_a_reg, 0);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_1_a, input_bank_a_reg, 2);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_2_a, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_0_a, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_1_a, input_bank_b_reg, 2);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_2_a, input_bank_c_reg, 0);
-
-            // Fixed-point multiplication.
-            acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-            acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc0, -output_shift);
-            acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-            acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc1, -output_shift);
-            // Add the output offset.
-            acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            acc_u8_0_1 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8_0_1 = util_vmax_x8(acc_u8_0_1,
-                                      vget_low_u8(output_activation_min_vec));
-            acc_u8_0_1 = util_vmin_x8(acc_u8_0_1,
-                                      vget_low_u8(output_activation_max_vec));
-
-            vst1_lane_u8x4(output_data, acc_u8_0_1, 0);
-            vst1_lane_u8x4(output_data + output_height_stride, acc_u8_0_1, 1);
-
-            acc0 = adjusted_bias_data_s_1;
-            acc1 = adjusted_bias_data_s_1;
-
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_0_b, input_bank_a_reg, 0);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_1_b, input_bank_a_reg, 2);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_2_b, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_0_b, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_1_b, input_bank_b_reg, 2);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_2_b, input_bank_c_reg, 0);
-
-            // Fixed-point multiplication.
-            acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-            acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc0, -output_shift);
-            acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-            acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc1, -output_shift);
-            // Add the output offset.
-            acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            acc_u8_0_1 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8_0_1 = util_vmax_x8(acc_u8_0_1,
-                                      vget_low_u8(output_activation_min_vec));
-            acc_u8_0_1 = util_vmin_x8(acc_u8_0_1,
-                                      vget_low_u8(output_activation_max_vec));
-
-            vst1_lane_u8x4(output_data + 4, acc_u8_0_1, 0);
-            vst1_lane_u8x4(output_data + 4 + output_height_stride, acc_u8_0_1,
-                           1);
-
-            input_bank_a_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_a_reg), 16));
-            input_bank_b_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_b_reg), 16));
-            input_bank_c_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_c_reg), 16));
-
-            output_data += output_depth;
-          }
-        }
-      } else {
-        TFLITE_DCHECK_EQ(block_height, 1);
-        // Work through one slice, by row, at a time.
-        const int8_t* scratch_data = scratch_block_data;
-        typename QuantizationTypeImpl<quantization_type>::ExternalType*
-            output_data = output_block_data + 8 * j_depth;
-
-        int8x16_t input_bank_a_reg;  //  left 0, right 0, left 1, right 1.
-        int8x16_t input_bank_b_reg;  //  left 2, right 2, xxx, xxx.
-
-        // Load first sub-micro block of data into operational banks.
-        input_bank_a_reg =
-            vld1q_dup_s8x4(scratch_data);  // Load lane 0, avoiding
-                                           // uninitialized variable.
-        input_bank_a_reg = vld1q_lane_8x4(
-            scratch_data + workspace_height_stride, input_bank_a_reg, 2);
-        input_bank_b_reg = vld1q_dup_s8x4(
-            scratch_data +
-            2 * workspace_height_stride);  // Load lane 0, avoiding
-                                           // uninitialized variable.
-
-        int32x4_t acc0;
-        int32x4_t acc1;
-
-        for (int i_width = 0; i_width < output_width_overall_micro_repeats;
-             ++i_width) {
-          const int output_width =
-              i_width == output_width_micro_repeats ? residual_width : 2;
-
-          TFLITE_DCHECK_LE(output_width, 2);
-          TFLITE_DCHECK_GE(output_width, 1);
-          TFLITE_DCHECK_LE(output_width * kStrideVal, 4);
-          const int8_t* input_data = scratch_data + 4 + 4 * i_width;
-
-          // Load next sub-micro block of data.
-          input_bank_a_reg = vld1q_lane_8x4(input_data, input_bank_a_reg, 1);
-          input_bank_a_reg = vld1q_lane_8x4(
-              input_data + workspace_height_stride, input_bank_a_reg, 3);
-          input_bank_b_reg = vld1q_lane_8x4(
-              input_data + 2 * workspace_height_stride, input_bank_b_reg, 1);
-
-          int16x8_t acc_s16_0_1;
-          uint8x8_t acc_u8_0_1;
-
-          // Iterate over input width shifts within 4x4 blocks.
-          {
-            acc0 = adjusted_bias_data_s_0;
-
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_2_a, input_bank_b_reg, 0);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_0_a, input_bank_a_reg, 0);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_1_a, input_bank_a_reg, 2);
-
-            acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-            acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc0, -output_shift);
-
-            // Second sub-block accumulation.
-            acc1 = adjusted_bias_data_s_1;
-
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_2_b, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_0_b, input_bank_a_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_1_b, input_bank_a_reg, 2);
-
-            acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-            acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc1, -output_shift);
-
-            // Add the output offset.
-            acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            acc_u8_0_1 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8_0_1 = util_vmax_x8(acc_u8_0_1,
-                                      vget_low_u8(output_activation_min_vec));
-            acc_u8_0_1 = util_vmin_x8(acc_u8_0_1,
-                                      vget_low_u8(output_activation_max_vec));
-
-            // This stores the results for both sub-blocks together.
-            util_vst1_u8(output_data, acc_u8_0_1);
-
-            input_bank_a_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_a_reg), 16));
-            input_bank_b_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_b_reg), 16));
-
-            output_data += output_depth;
-          }
-          if (output_width == 2) {
-            acc0 = adjusted_bias_data_s_0;
-
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_2_a, input_bank_b_reg, 0);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_0_a, input_bank_a_reg, 0);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_1_a, input_bank_a_reg, 2);
-
-            acc0 = vqrdmulhq_n_s32(acc0, output_multiplier);
-            acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc0, -output_shift);
-
-            // Second sub-block accumulation.
-            acc1 = adjusted_bias_data_s_1;
-
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_2_b, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_0_b, input_bank_a_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_1_b, input_bank_a_reg, 2);
-
-            acc1 = vqrdmulhq_n_s32(acc1, output_multiplier);
-            acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::Run(
-                acc1, -output_shift);
-
-            // Add the output offset.
-            acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            acc_u8_0_1 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8_0_1 = util_vmax_x8(acc_u8_0_1,
-                                      vget_low_u8(output_activation_min_vec));
-            acc_u8_0_1 = util_vmin_x8(acc_u8_0_1,
-                                      vget_low_u8(output_activation_max_vec));
-
-            // This stores the results for both sub-blocks together.
-            util_vst1_u8(output_data, acc_u8_0_1);
-
-            input_bank_a_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_a_reg), 16));
-            input_bank_b_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_b_reg), 16));
-
-            output_data += output_depth;
-          }
-        }
-      }
-    }
-  }
-
-  static inline void Run(const int8_t* scratch_block_data,
-                         const int8_t* filter_workspace,
-                         const int32_t* bias_data, uint8_t* output_block_data,
-                         const DepthwiseConvDotProdParams* function_params) {
-    KernelMacroBlockIntrinsics(scratch_block_data, filter_workspace, bias_data,
-                               output_block_data, function_params);
-  }
-};
-
-template <>
-struct KernelMacroBlock<
-    DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct,
-    QuantizationType::kPerChannelInt8,
-    DepthwiseConvDepthMultiplication::kNoMultiplication,
-    /*stride=*/1> {
-  static inline int8x8_t vqmovxn_s16(int16x8_t x) { return vqmovn_s16(x); }
-  static inline int8x8_t util_vmin_x8(int8x8_t a, int8x8_t b) {
-    return vmin_s8(a, b);
-  }
-  static inline int8x8_t util_vmax_x8(int8x8_t a, int8x8_t b) {
-    return vmax_s8(a, b);
-  }
-  static inline int8x16_t util_vminq_x8(int8x16_t a, int8x16_t b) {
-    return vminq_s8(a, b);
-  }
-  static inline int8x16_t util_vmaxq_x8(int8x16_t a, int8x16_t b) {
-    return vmaxq_s8(a, b);
-  }
-
-  static inline void KernelMacroBlockIntrinsics(
-      const int8_t* scratch_block_data, const int8_t* filter_workspace,
-      const int32_t* bias_data, int8_t* output_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-    static constexpr QuantizationType quantization_type =
-        QuantizationType::kPerChannelInt8;
-
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-    const int input_width_overall_micro_repeats =
-        function_params->input_width_overall_micro_repeats;
-    const int output_width_micro_repeats =
-        function_params->output_width_micro_repeats;
-    const int depth_micro_repeats = function_params->depth_micro_repeats;
-    const int depth = function_params->input_depth;
-
-    const int output_width_overall_micro_repeats =
-        function_params->output_width_overall_micro_repeats;
-    const int block_height = function_params->outbound_block_height;
-    const int residual_width = function_params->output_residual_width;
-    const int output_height_stride = function_params->output_height_stride;
-    constexpr int kBiasIncrement = 4;
-
-    TFLITE_DCHECK(depth_micro_repeats > 0);
-    const int width_micro_stride = 4 * 8;
-    const int depth_micro_stride =
-        width_micro_stride * input_width_overall_micro_repeats;
-
-    const int32_t output_activation_min =
-        function_params->quantized_activation_min;
-    const int32_t output_activation_max =
-        function_params->quantized_activation_max;
-    const int32_t output_offset = function_params->output_offset;
-    const int32_t* output_shift_per_channel =
-        function_params->output_shift_per_channel;
-    const int32_t* output_multiplier_per_channel =
-        function_params->output_multiplier_per_channel;
-    if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-      TFLITE_DCHECK_GE(output_activation_min, 0);
-      TFLITE_DCHECK_LT(output_activation_min, 256);
-      TFLITE_DCHECK_GE(output_activation_max, 0);
-      TFLITE_DCHECK_LT(output_activation_max, 256);
-    } else {
-      TFLITE_DCHECK_GE(output_activation_min, -128);
-      TFLITE_DCHECK_LT(output_activation_min, 128);
-      TFLITE_DCHECK_GE(output_activation_max, -128);
-      TFLITE_DCHECK_LT(output_activation_max, 128);
-      TFLITE_DCHECK_NE(output_shift_per_channel, nullptr);
-      TFLITE_DCHECK_NE(output_multiplier_per_channel, nullptr);
-    }
-    TFLITE_DCHECK_GE(output_offset, -32878);
-    TFLITE_DCHECK_LT(output_offset, 32768);
-
-    const int16x8_t output_offset_vec =
-        vdupq_n_s16(static_cast<int16_t>(output_offset));
-    const int8x16_t output_activation_min_vec =
-        vdupq_n_s8(static_cast<int8_t>(output_activation_min));
-    const int8x16_t output_activation_max_vec =
-        vdupq_n_s8(static_cast<int8_t>(output_activation_max));
-
-    const int8_t* input_data_depthwise = scratch_block_data;
-    typename QuantizationTypeImpl<quantization_type>::ExternalType*
-        output_data_depthwise = output_block_data;
-    for (int j_depth = 0; j_depth < depth_micro_repeats; ++j_depth) {
-      // Simulate NEON-register transposition of subset of filter.
-      int8x16_t filter_reg_0_a;
-      int8x16_t filter_reg_0_b;
-      int8x16_t filter_reg_1_a;
-      int8x16_t filter_reg_1_b;
-      int8x16_t filter_reg_2_a;
-      int8x16_t filter_reg_2_b;
-      int8x16_t filter_reg_0_a_shifted;
-      int8x16_t filter_reg_1_a_shifted;
-      int8x16_t filter_reg_2_a_shifted;
-
-      filter_reg_0_a = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_0_b = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_1_a = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_1_b = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_2_a = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_2_b = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-
-      filter_reg_0_a_shifted = vreinterpretq_s8_u32(
-          vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_0_a), 8));
-      filter_reg_1_a_shifted = vreinterpretq_s8_u32(
-          vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_1_a), 8));
-      filter_reg_2_a_shifted = vreinterpretq_s8_u32(
-          vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_2_a), 8));
-
-      if (block_height == 4) {
-        for (int s = 0; s < 2; ++s) {
-          // Work through one slice, by row, at a time.
-          const int8_t* input_data_base = input_data_depthwise + 2 * 8 * s;
-          typename QuantizationTypeImpl<quantization_type>::ExternalType*
-              output_data_base = output_data_depthwise + 4 * s;
-
-          const int8_t* next_input_data = input_data_base;
-          typename QuantizationTypeImpl<quantization_type>::ExternalType*
-              output_data = output_data_base;
-
-          const int32x4_t adjusted_bias_data = vld1q_s32(bias_data);
-          bias_data += kBiasIncrement;
-
-          const int32x4_t output_shift =
-              vld1q_s32(output_shift_per_channel + j_depth * 8 + 4 * s);
-          const int32x4_t output_multiplier =
-              vld1q_s32(output_multiplier_per_channel + j_depth * 8 + 4 * s);
-
-          // Load first sub-micro block of data into operational banks.
-          int8x16_t left_bank_0_reg = vld1q_s8(next_input_data);
-          int8x16_t left_bank_1_reg =
-              vld1q_s8(next_input_data + workspace_height_stride);
-          int8x16_t left_bank_2_reg =
-              vld1q_s8(next_input_data + 2 * workspace_height_stride);
-          int8x16_t left_bank_3_reg =
-              vld1q_s8(next_input_data + 3 * workspace_height_stride);
-          int8x16_t left_bank_4_reg =
-              vld1q_s8(next_input_data + 4 * workspace_height_stride);
-          int8x16_t left_bank_5_reg =
-              vld1q_s8(next_input_data + 5 * workspace_height_stride);
-
-          int32x4_t acc0;
-          int32x4_t acc1;
-          int32x4_t acc2;
-          int32x4_t acc3;
-
-          acc0 = adjusted_bias_data;
-          acc1 = adjusted_bias_data;
-          acc2 = adjusted_bias_data;
-          acc3 = adjusted_bias_data;
-
-          acc0 = vdotq_s32(acc0, filter_reg_2_a, left_bank_2_reg);
-          acc1 = vdotq_s32(acc1, filter_reg_1_a, left_bank_2_reg);
-          acc2 = vdotq_s32(acc2, filter_reg_0_a, left_bank_2_reg);
-          acc3 = vdotq_s32(acc3, filter_reg_0_a, left_bank_3_reg);
-
-          for (int i_width = 0; i_width < output_width_micro_repeats;
-               ++i_width) {
-            next_input_data += width_micro_stride;
-
-            // Iterate over input width shifts within 4x4 blocks.
-            {
-              acc0 = vdotq_s32(acc0, filter_reg_0_a, left_bank_0_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_1_a, left_bank_1_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_0_a, left_bank_1_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_2_a, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_1_a, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_2_a, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_1_a, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_2_a, left_bank_5_reg);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc0, output_shift);
-              acc1 = vqrdmulhq_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc1, output_shift);
-              acc2 = vqrdmulhq_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc2, output_shift);
-              acc3 = vqrdmulhq_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc3, output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              int8x16_t acc_u8_all = vcombine_s8(vqmovxn_s16(acc_s16_0_1),
-                                                 vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_s8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_s8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_s8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_s8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              output_data += depth;
-            }
-
-            // Load next sub-micro block of data.
-            int8x16_t right_bank_0_reg;
-            int8x16_t right_bank_1_reg;
-            int8x16_t right_bank_2_reg;
-            int8x16_t right_bank_3_reg;
-            int8x16_t right_bank_4_reg;
-            int8x16_t right_bank_5_reg;
-
-            // Loading of next block always valid.
-            right_bank_0_reg = vld1q_s8(next_input_data);
-            right_bank_1_reg =
-                vld1q_s8(next_input_data + workspace_height_stride);
-            right_bank_2_reg =
-                vld1q_s8(next_input_data + 2 * workspace_height_stride);
-            right_bank_3_reg =
-                vld1q_s8(next_input_data + 3 * workspace_height_stride);
-            right_bank_4_reg =
-                vld1q_s8(next_input_data + 4 * workspace_height_stride);
-            right_bank_5_reg =
-                vld1q_s8(next_input_data + 5 * workspace_height_stride);
-
-            {
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_s32(acc0, filter_reg_0_a_shifted, left_bank_0_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_1_a_shifted, left_bank_1_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_2_a_shifted, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_0_a_shifted, left_bank_1_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_1_a_shifted, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_2_a_shifted, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_0_a_shifted, left_bank_2_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_1_a_shifted, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_2_a_shifted, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_0_a_shifted, left_bank_3_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_1_a_shifted, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_2_a_shifted, left_bank_5_reg);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc0, output_shift);
-              acc1 = vqrdmulhq_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc1, output_shift);
-              acc2 = vqrdmulhq_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc2, output_shift);
-              acc3 = vqrdmulhq_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc3, output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              int8x16_t acc_u8_all = vcombine_s8(vqmovxn_s16(acc_s16_0_1),
-                                                 vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_s8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_s8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_s8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_s8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              left_bank_0_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_0_reg)));
-              left_bank_1_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_1_reg)));
-              left_bank_2_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_2_reg)));
-              left_bank_3_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_3_reg)));
-              left_bank_4_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_4_reg)));
-              left_bank_5_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_5_reg)));
-              vtrn1_s8x2_in_place(&left_bank_0_reg, &right_bank_0_reg);
-              vtrn1_s8x2_in_place(&left_bank_1_reg, &right_bank_1_reg);
-              vtrn1_s8x2_in_place(&left_bank_2_reg, &right_bank_2_reg);
-              vtrn1_s8x2_in_place(&left_bank_3_reg, &right_bank_3_reg);
-              vtrn1_s8x2_in_place(&left_bank_4_reg, &right_bank_4_reg);
-              vtrn1_s8x2_in_place(&left_bank_5_reg, &right_bank_5_reg);
-
-              output_data += depth;
-            }
-
-            {
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_s32(acc0, filter_reg_0_a, left_bank_0_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_1_a, left_bank_1_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_2_a, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_0_a, left_bank_1_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_1_a, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_2_a, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_0_a, left_bank_2_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_1_a, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_2_a, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_0_a, left_bank_3_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_1_a, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_2_a, left_bank_5_reg);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc0, output_shift);
-              acc1 = vqrdmulhq_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc1, output_shift);
-              acc2 = vqrdmulhq_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc2, output_shift);
-              acc3 = vqrdmulhq_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc3, output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              int8x16_t acc_u8_all = vcombine_s8(vqmovxn_s16(acc_s16_0_1),
-                                                 vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_s8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_s8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_s8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_s8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              output_data += depth;
-            }
-
-            {
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_s32(acc0, filter_reg_0_a_shifted, left_bank_0_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_1_a_shifted, left_bank_1_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_2_a_shifted, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_0_a_shifted, left_bank_1_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_1_a_shifted, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_2_a_shifted, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_0_a_shifted, left_bank_2_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_1_a_shifted, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_2_a_shifted, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_0_a_shifted, left_bank_3_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_1_a_shifted, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_2_a_shifted, left_bank_5_reg);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc0, output_shift);
-              acc1 = vqrdmulhq_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc1, output_shift);
-              acc2 = vqrdmulhq_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc2, output_shift);
-              acc3 = vqrdmulhq_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc3, output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              int8x16_t acc_u8_all = vcombine_s8(vqmovxn_s16(acc_s16_0_1),
-                                                 vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_s8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_s8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_s8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_s8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              left_bank_0_reg = right_bank_0_reg;
-              left_bank_1_reg = right_bank_1_reg;
-              left_bank_2_reg = right_bank_2_reg;
-              left_bank_3_reg = right_bank_3_reg;
-              left_bank_4_reg = right_bank_4_reg;
-              left_bank_5_reg = right_bank_5_reg;
-
-              output_data += depth;
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_s32(acc0, filter_reg_2_a, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_1_a, left_bank_2_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_0_a, left_bank_2_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_0_a, left_bank_3_reg);
-            }
-          }
-
-          if (residual_width > 0) {
-            next_input_data += width_micro_stride;
-            const int output_width = residual_width;
-
-            // Load next sub-micro block of data.
-            int8x16_t right_bank_0_reg;
-            int8x16_t right_bank_1_reg;
-            int8x16_t right_bank_2_reg;
-            int8x16_t right_bank_3_reg;
-            int8x16_t right_bank_4_reg;
-            int8x16_t right_bank_5_reg;
-            // Logic: (output_width - 1) * stride_val < 2.
-            const bool no_right_block = output_width < 3;
-
-            if (no_right_block) {
-              // Only needed for sanitizer checks.
-              right_bank_0_reg = vdupq_n_s8(0);
-              right_bank_1_reg = vdupq_n_s8(0);
-              right_bank_2_reg = vdupq_n_s8(0);
-              right_bank_3_reg = vdupq_n_s8(0);
-              right_bank_4_reg = vdupq_n_s8(0);
-              right_bank_5_reg = vdupq_n_s8(0);
-            } else {
-              right_bank_0_reg = vld1q_s8(next_input_data);
-              right_bank_1_reg =
-                  vld1q_s8(next_input_data + workspace_height_stride);
-              right_bank_2_reg =
-                  vld1q_s8(next_input_data + 2 * workspace_height_stride);
-              right_bank_3_reg =
-                  vld1q_s8(next_input_data + 3 * workspace_height_stride);
-              right_bank_4_reg =
-                  vld1q_s8(next_input_data + 4 * workspace_height_stride);
-              right_bank_5_reg =
-                  vld1q_s8(next_input_data + 5 * workspace_height_stride);
-            }
-
-            // Iterate over input width shifts within 4x4 blocks.
-            for (int x = 0; x < output_width; ++x) {
-              acc0 = vdotq_s32(acc0, filter_reg_0_a, left_bank_0_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_1_a, left_bank_1_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_0_a, left_bank_1_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_2_a, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_1_a, left_bank_3_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_2_a, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_1_a, left_bank_4_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_2_a, left_bank_5_reg);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc0, output_shift);
-              acc1 = vqrdmulhq_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc1, output_shift);
-              acc2 = vqrdmulhq_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc2, output_shift);
-              acc3 = vqrdmulhq_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc3, output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              int8x16_t acc_u8_all = vcombine_s8(vqmovxn_s16(acc_s16_0_1),
-                                                 vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_s8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_s8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_s8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_s8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              biregister_rotate_8(&left_bank_0_reg, &right_bank_0_reg);
-              biregister_rotate_8(&left_bank_1_reg, &right_bank_1_reg);
-              biregister_rotate_8(&left_bank_2_reg, &right_bank_2_reg);
-              biregister_rotate_8(&left_bank_3_reg, &right_bank_3_reg);
-              biregister_rotate_8(&left_bank_4_reg, &right_bank_4_reg);
-              biregister_rotate_8(&left_bank_5_reg, &right_bank_5_reg);
-
-              output_data += depth;
-
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_s32(acc0, filter_reg_2_a, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_1_a, left_bank_2_reg);
-              acc2 = vdotq_s32(acc2, filter_reg_0_a, left_bank_2_reg);
-              acc3 = vdotq_s32(acc3, filter_reg_0_a, left_bank_3_reg);
-            }
-          }
-          input_data_base += 4 * workspace_height_stride;
-          output_data_base += 4 * output_height_stride;
-
-          // Move to next sub-block: advance to second set of filters, to new
-          // bias.
-          filter_reg_0_a = filter_reg_0_b;
-          filter_reg_1_a = filter_reg_1_b;
-          filter_reg_2_a = filter_reg_2_b;
-          filter_reg_0_a_shifted = vreinterpretq_s8_u32(
-              vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_0_a), 8));
-          filter_reg_1_a_shifted = vreinterpretq_s8_u32(
-              vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_1_a), 8));
-          filter_reg_2_a_shifted = vreinterpretq_s8_u32(
-              vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_2_a), 8));
-        }
-      } else {
-        const int8_t* input_data_base = input_data_depthwise;
-        typename QuantizationTypeImpl<quantization_type>::ExternalType*
-            output_data_base = output_data_depthwise;
-
-        const int32x4_t adjusted_bias_data_a = vld1q_s32(bias_data);
-        bias_data += kBiasIncrement;
-        const int32x4_t adjusted_bias_data_b = vld1q_s32(bias_data);
-        bias_data += kBiasIncrement;
-
-        const int32x4_t output_shift_a =
-            vld1q_s32(output_shift_per_channel + j_depth * 8);
-        const int32x4_t output_multiplier_a =
-            vld1q_s32(output_multiplier_per_channel + j_depth * 8);
-        const int32x4_t output_shift_b =
-            vld1q_s32(output_shift_per_channel + j_depth * 8 + 4);
-        const int32x4_t output_multiplier_b =
-            vld1q_s32(output_multiplier_per_channel + j_depth * 8 + 4);
-
-        for (int k_height = 0; k_height < block_height; ++k_height) {
-          const int8_t* next_input_data = input_data_base;
-          typename QuantizationTypeImpl<quantization_type>::ExternalType*
-              output_data = output_data_base;
-
-          // Load first sub-micro block of data into operational banks.
-          int8x16_t left_bank_0_reg_a = vld1q_s8(next_input_data);
-          int8x16_t left_bank_1_reg_a =
-              vld1q_s8(next_input_data + workspace_height_stride);
-          int8x16_t left_bank_2_reg_a =
-              vld1q_s8(next_input_data + 2 * workspace_height_stride);
-          int8x16_t left_bank_0_reg_b = vld1q_s8(next_input_data + 16);
-          int8x16_t left_bank_1_reg_b =
-              vld1q_s8(next_input_data + workspace_height_stride + 16);
-          int8x16_t left_bank_2_reg_b =
-              vld1q_s8(next_input_data + 2 * workspace_height_stride + 16);
-
-          for (int i_width = 0; i_width < output_width_overall_micro_repeats;
-               ++i_width) {
-            next_input_data += width_micro_stride;
-            const int output_width =
-                i_width == output_width_micro_repeats ? residual_width : 4;
-
-            int8x16_t right_bank_0_reg_a;
-            int8x16_t right_bank_1_reg_a;
-            int8x16_t right_bank_2_reg_a;
-            int8x16_t right_bank_0_reg_b;
-            int8x16_t right_bank_1_reg_b;
-            int8x16_t right_bank_2_reg_b;
-            // Logic: (output_width - 1) * stride_val < 2.
-            const bool no_right_block = output_width < 3;
-
-            // Load next sub-micro block of data.
-            if (no_right_block) {
-              // Only needed for sanitizer checks.
-              right_bank_0_reg_a = vdupq_n_s8(0);
-              right_bank_1_reg_a = vdupq_n_s8(0);
-              right_bank_2_reg_a = vdupq_n_s8(0);
-              right_bank_0_reg_b = vdupq_n_s8(0);
-              right_bank_1_reg_b = vdupq_n_s8(0);
-              right_bank_2_reg_b = vdupq_n_s8(0);
-            } else {
-              right_bank_0_reg_a = vld1q_s8(next_input_data);
-              right_bank_1_reg_a =
-                  vld1q_s8(next_input_data + workspace_height_stride);
-              right_bank_2_reg_a =
-                  vld1q_s8(next_input_data + 2 * workspace_height_stride);
-              right_bank_0_reg_b = vld1q_s8(next_input_data + 16);
-              right_bank_1_reg_b =
-                  vld1q_s8(next_input_data + workspace_height_stride + 16);
-              right_bank_2_reg_b =
-                  vld1q_s8(next_input_data + 2 * workspace_height_stride + 16);
-            }
-
-            // Iterate over input width shifts within 4x4 blocks.
-            for (int x = 0; x < output_width; ++x) {
-              int32x4_t acc_a = adjusted_bias_data_a;
-              int32x4_t acc_b = adjusted_bias_data_b;
-              acc_a = vdotq_s32(acc_a, filter_reg_0_a, left_bank_0_reg_a);
-              acc_a = vdotq_s32(acc_a, filter_reg_1_a, left_bank_1_reg_a);
-              acc_a = vdotq_s32(acc_a, filter_reg_2_a, left_bank_2_reg_a);
-              acc_b = vdotq_s32(acc_b, filter_reg_0_b, left_bank_0_reg_b);
-              acc_b = vdotq_s32(acc_b, filter_reg_1_b, left_bank_1_reg_b);
-              acc_b = vdotq_s32(acc_b, filter_reg_2_b, left_bank_2_reg_b);
-
-              // Fixed-point multiplication.
-              acc_a = vqrdmulhq_s32(acc_a, output_multiplier_a);
-              acc_b = vqrdmulhq_s32(acc_b, output_multiplier_b);
-              acc_a =
-                  DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                      acc_a, output_shift_a);
-              acc_b =
-                  DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                      acc_b, output_shift_b);
-              // Add the output offset.
-              int16x8_t acc_s16_0_0 =
-                  vcombine_s16(vqmovn_s32(acc_a), vqmovn_s32(acc_b));
-              acc_s16_0_0 = vqaddq_s16(acc_s16_0_0, output_offset_vec);
-              // Apply the activation function.
-              int8x8_t acc_u8_0_0 = vqmovxn_s16(acc_s16_0_0);
-              acc_u8_0_0 = util_vmax_x8(acc_u8_0_0,
-                                        vget_low_s8(output_activation_min_vec));
-              acc_u8_0_0 = util_vmin_x8(acc_u8_0_0,
-                                        vget_low_s8(output_activation_max_vec));
-
-              vst1_s8(output_data, acc_u8_0_0);
-
-              biregister_rotate_8(&left_bank_0_reg_a, &right_bank_0_reg_a);
-              biregister_rotate_8(&left_bank_1_reg_a, &right_bank_1_reg_a);
-              biregister_rotate_8(&left_bank_2_reg_a, &right_bank_2_reg_a);
-              biregister_rotate_8(&left_bank_0_reg_b, &right_bank_0_reg_b);
-              biregister_rotate_8(&left_bank_1_reg_b, &right_bank_1_reg_b);
-              biregister_rotate_8(&left_bank_2_reg_b, &right_bank_2_reg_b);
-
-              output_data += depth;
-            }
-          }
-          input_data_base += workspace_height_stride;
-          output_data_base += output_height_stride;
-        }
-      }
-      input_data_depthwise += depth_micro_stride;
-      output_data_depthwise += 8;
-    }
-  }  // NOLINT(readability/fn_size) Manually unrolled.
-
-  static inline void Run(const int8_t* scratch_block_data,
-                         const int8_t* filter_workspace,
-                         const int32_t* bias_data, int8_t* output_block_data,
-                         const DepthwiseConvDotProdParams* function_params) {
-    KernelMacroBlockIntrinsics(scratch_block_data, filter_workspace, bias_data,
-                               output_block_data, function_params);
-  }
-};
-
-template <>
-struct KernelMacroBlock<
-    DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct,
-    QuantizationType::kPerChannelInt8,
-    DepthwiseConvDepthMultiplication::kNoMultiplication,
-    /*stride=*/2> {
-  static inline int8x8_t vqmovxn_s16(int16x8_t x) { return vqmovn_s16(x); }
-  static inline int8x8_t util_vmin_x8(int8x8_t a, int8x8_t b) {
-    return vmin_s8(a, b);
-  }
-  static inline int8x8_t util_vmax_x8(int8x8_t a, int8x8_t b) {
-    return vmax_s8(a, b);
-  }
-
-  static inline void KernelMacroBlockIntrinsics(
-      const int8_t* scratch_block_data, const int8_t* filter_workspace,
-      const int32_t* bias_data, int8_t* output_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-    static constexpr QuantizationType quantization_type =
-        QuantizationType::kPerChannelInt8;
-
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-    const int input_width_overall_micro_repeats =
-        function_params->input_width_overall_micro_repeats;
-    const int output_width_micro_repeats =
-        function_params->output_width_micro_repeats;
-    const int depth_micro_repeats = function_params->depth_micro_repeats;
-    const int depth = function_params->input_depth;
-    constexpr int kStrideVal = 2;
-    constexpr int kFourOverStride = 2;
-    TFLITE_DCHECK_EQ(function_params->stride, kStrideVal);
-    TFLITE_DCHECK_EQ(function_params->four_over_stride, kFourOverStride);
-
-    const int workspace_width_micro_repeats =
-        function_params->workspace_width_micro_repeats;
-    const int output_width_overall_micro_repeats =
-        function_params->output_width_overall_micro_repeats;
-    const int block_height = function_params->outbound_block_height;
-    const int residual_width = function_params->output_residual_width;
-    const int output_height_stride = function_params->output_height_stride;
-    constexpr int kBiasIncrement = 4;
-
-    TFLITE_DCHECK(depth_micro_repeats > 0);
-    const int width_micro_stride = 4 * 8;
-    const int depth_micro_stride =
-        width_micro_stride * input_width_overall_micro_repeats;
-
-    const int32_t output_activation_min =
-        function_params->quantized_activation_min;
-    const int32_t output_activation_max =
-        function_params->quantized_activation_max;
-    const int32_t output_offset = function_params->output_offset;
-    const int32_t* output_shift_per_channel =
-        function_params->output_shift_per_channel;
-    const int32_t* output_multiplier_per_channel =
-        function_params->output_multiplier_per_channel;
-    if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-      TFLITE_DCHECK_GE(output_activation_min, 0);
-      TFLITE_DCHECK_LT(output_activation_min, 256);
-      TFLITE_DCHECK_GE(output_activation_max, 0);
-      TFLITE_DCHECK_LT(output_activation_max, 256);
-    } else {
-      TFLITE_DCHECK_GE(output_activation_min, -128);
-      TFLITE_DCHECK_LT(output_activation_min, 128);
-      TFLITE_DCHECK_GE(output_activation_max, -128);
-      TFLITE_DCHECK_LT(output_activation_max, 128);
-      TFLITE_DCHECK_NE(output_shift_per_channel, nullptr);
-      TFLITE_DCHECK_NE(output_multiplier_per_channel, nullptr);
-    }
-    TFLITE_DCHECK_GE(output_offset, -32878);
-    TFLITE_DCHECK_LT(output_offset, 32768);
-
-    // This version only does min/max on 64 bits.
-    const int16x8_t output_offset_vec =
-        vdupq_n_s16(static_cast<int16_t>(output_offset));
-    const int8x8_t output_activation_min_vec =
-        vdup_n_s8(static_cast<int8_t>(output_activation_min));
-    const int8x8_t output_activation_max_vec =
-        vdup_n_s8(static_cast<int8_t>(output_activation_max));
-
-    constexpr int shuffled_filter_increment = 2 * 3 * 4 * 4;
-
-    TFLITE_DCHECK_LE(block_height, 2);
-
-    for (int j_depth = 0; j_depth < depth_micro_repeats; ++j_depth) {
-      const int8_t* filter_block =
-          filter_workspace + shuffled_filter_increment * j_depth;
-
-      if (block_height == 2) {
-        for (int s = 0; s < 2; ++s) {
-          // Simulate NEON-register transposition of subset of filter.
-          int8x16_t filter_reg_0_a;
-          int8x16_t filter_reg_1_a;
-          int8x16_t filter_reg_2_a;
-
-          filter_reg_0_a = vld1q_s8(filter_block + s * 16);
-          filter_reg_1_a = vld1q_s8(filter_block + s * 16 + 32);
-          filter_reg_2_a = vld1q_s8(filter_block + s * 16 + 64);
-
-          const int8_t* scratch_data =
-              scratch_block_data + depth_micro_stride * j_depth;
-          typename QuantizationTypeImpl<quantization_type>::ExternalType*
-              output_data = output_block_data + 8 * j_depth;
-          const int8_t* input_data_0 = scratch_data + s * 2 * 8;
-
-          const int32x4_t adjusted_bias_data = vld1q_s32(bias_data);
-
-          const int32x4_t output_shift =
-              vld1q_s32(output_shift_per_channel + j_depth * 8 + 4 * s);
-          const int32x4_t output_multiplier =
-              vld1q_s32(output_multiplier_per_channel + j_depth * 8 + 4 * s);
-
-          // Load first sub-micro block of data into operational banks.
-          int8x16_t left_bank_0_reg = vld1q_s8(input_data_0);
-          int8x16_t left_bank_1_reg =
-              vld1q_s8(input_data_0 + workspace_height_stride);
-          int8x16_t left_bank_2_reg =
-              vld1q_s8(input_data_0 + 2 * workspace_height_stride);
-          int8x16_t left_bank_3_reg =
-              vld1q_s8(input_data_0 + 3 * workspace_height_stride);
-          int8x16_t left_bank_4_reg =
-              vld1q_s8(input_data_0 + 4 * workspace_height_stride);
-
-          int8x16_t right_bank_0_reg;
-          int8x16_t right_bank_1_reg;
-          int8x16_t right_bank_2_reg;
-          int8x16_t right_bank_3_reg;
-          int8x16_t right_bank_4_reg;
-
-          int32x4_t acc0;
-          int32x4_t acc1;
-          int16x8_t acc_s16_0_1;
-          int8x8_t acc_u8;
-
-          int i_width = 0;
-
-          // When output_width_micro_repeats <
-          // output_width_overall_micro_repeats, 0 < residual_width <= 2, and so
-          // residual_width == 1 is then true iff residual_width < 2.
-          const int adjusted_width_micro_repeats =
-              (output_width_micro_repeats <
-               output_width_overall_micro_repeats) &&
-                      (residual_width == 1)
-                  ? output_width_micro_repeats
-                  : output_width_overall_micro_repeats;
-
-          for (; i_width < adjusted_width_micro_repeats; ++i_width) {
-            const int output_width = kFourOverStride;
-            TFLITE_DCHECK_LE(output_width * kStrideVal, 4);
-            const int8_t* input_data =
-                input_data_0 + width_micro_stride * i_width;
-            acc0 = adjusted_bias_data;
-            acc1 = adjusted_bias_data;
-            right_bank_0_reg = vld1q_s8(input_data + width_micro_stride);
-            right_bank_1_reg = vld1q_s8(input_data + width_micro_stride +
-                                        workspace_height_stride);
-
-            acc0 = vdotq_s32(acc0, filter_reg_0_a, left_bank_0_reg);
-            acc1 = vdotq_s32(acc1, filter_reg_0_a, left_bank_2_reg);
-            typename QuantizationTypeImpl<quantization_type>::ExternalType*
-                output_data_base = output_data + depth * 2 * i_width + 4 * s;
-
-            right_bank_2_reg = vld1q_s8(input_data + width_micro_stride +
-                                        2 * workspace_height_stride);
-            right_bank_3_reg = vld1q_s8(input_data + width_micro_stride +
-                                        3 * workspace_height_stride);
-            acc0 = vdotq_s32(acc0, filter_reg_1_a, left_bank_1_reg);
-            acc0 = vdotq_s32(acc0, filter_reg_2_a, left_bank_2_reg);
-            acc1 = vdotq_s32(acc1, filter_reg_1_a, left_bank_3_reg);
-            acc1 = vdotq_s32(acc1, filter_reg_2_a, left_bank_4_reg);
-            right_bank_4_reg = vld1q_s8(input_data + width_micro_stride +
-                                        4 * workspace_height_stride);
-
-            // Fixed-point multiplication.
-            acc0 = vqrdmulhq_s32(acc0, output_multiplier);
-            acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc0, output_shift);
-            acc1 = vqrdmulhq_s32(acc1, output_multiplier);
-            acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc1, output_shift);
-            // Add the output offset.
-            acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            acc_u8 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8 = util_vmax_x8(acc_u8, output_activation_min_vec);
-            acc_u8 = util_vmin_x8(acc_u8, output_activation_max_vec);
-
-            left_bank_0_reg = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_0_reg)));
-            left_bank_1_reg = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_1_reg)));
-            left_bank_2_reg = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_2_reg)));
-            left_bank_3_reg = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_3_reg)));
-            left_bank_4_reg = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_4_reg)));
-            acc0 = adjusted_bias_data;
-            acc1 = adjusted_bias_data;
-            vtrn1_s8x2_in_place(&left_bank_0_reg, &right_bank_0_reg);
-            vtrn1_s8x2_in_place(&left_bank_1_reg, &right_bank_1_reg);
-            vtrn1_s8x2_in_place(&left_bank_2_reg, &right_bank_2_reg);
-            vst1_lane_s8x4(output_data_base, acc_u8, 0);
-            vst1_lane_s8x4(output_data_base + output_height_stride, acc_u8, 1);
-
-            vtrn1_s8x2_in_place(&left_bank_3_reg, &right_bank_3_reg);
-            vtrn1_s8x2_in_place(&left_bank_4_reg, &right_bank_4_reg);
-
-            acc0 = vdotq_s32(acc0, filter_reg_0_a, left_bank_0_reg);
-            acc1 = vdotq_s32(acc1, filter_reg_0_a, left_bank_2_reg);
-            acc0 = vdotq_s32(acc0, filter_reg_1_a, left_bank_1_reg);
-            acc1 = vdotq_s32(acc1, filter_reg_1_a, left_bank_3_reg);
-            acc0 = vdotq_s32(acc0, filter_reg_2_a, left_bank_2_reg);
-            acc1 = vdotq_s32(acc1, filter_reg_2_a, left_bank_4_reg);
-
-            // Fixed-point multiplication.
-            acc0 = vqrdmulhq_s32(acc0, output_multiplier);
-            acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc0, output_shift);
-            acc1 = vqrdmulhq_s32(acc1, output_multiplier);
-            acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc1, output_shift);
-            // Add the output offset.
-            acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            acc_u8 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8 = util_vmax_x8(acc_u8, output_activation_min_vec);
-            acc_u8 = util_vmin_x8(acc_u8, output_activation_max_vec);
-
-            vst1_lane_s8x4(output_data_base + depth, acc_u8, 0);
-            vst1_lane_s8x4(output_data_base + depth + output_height_stride,
-                           acc_u8, 1);
-
-            left_bank_0_reg = right_bank_0_reg;
-            left_bank_1_reg = right_bank_1_reg;
-            left_bank_2_reg = right_bank_2_reg;
-            left_bank_3_reg = right_bank_3_reg;
-            left_bank_4_reg = right_bank_4_reg;
-          }
-          for (; i_width < output_width_overall_micro_repeats; ++i_width) {
-            TFLITE_DCHECK_NE(residual_width, kFourOverStride);
-
-            // No need to load next ("right") block of data.
-
-            typename QuantizationTypeImpl<quantization_type>::ExternalType*
-                output_data_base = output_data + depth * 2 * i_width + 4 * s;
-
-            // Iterate over input width shifts within 4x4 blocks.
-            {
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-
-              acc0 = vdotq_s32(acc0, filter_reg_0_a, left_bank_0_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_1_a, left_bank_1_reg);
-              acc0 = vdotq_s32(acc0, filter_reg_2_a, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_0_a, left_bank_2_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_1_a, left_bank_3_reg);
-              acc1 = vdotq_s32(acc1, filter_reg_2_a, left_bank_4_reg);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc0, output_shift);
-              acc1 = vqrdmulhq_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc1, output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              // Apply the activation function.
-              int8x8_t acc_u8 = vqmovxn_s16(acc_s16_0_1);
-              acc_u8 = util_vmax_x8(acc_u8, output_activation_min_vec);
-              acc_u8 = util_vmin_x8(acc_u8, output_activation_max_vec);
-
-              vst1_lane_s8x4(output_data_base, acc_u8, 0);
-              vst1_lane_s8x4(output_data_base + output_height_stride, acc_u8,
-                             1);
-
-              left_bank_0_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_0_reg)));
-              left_bank_1_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_1_reg)));
-              left_bank_2_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_2_reg)));
-              left_bank_3_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_3_reg)));
-              left_bank_4_reg = vreinterpretq_s8_u16(
-                  vrev32q_u16(vreinterpretq_u16_s8(left_bank_4_reg)));
-              vtrn1_s8x2_in_place(&left_bank_0_reg, &right_bank_0_reg);
-              vtrn1_s8x2_in_place(&left_bank_1_reg, &right_bank_1_reg);
-              vtrn1_s8x2_in_place(&left_bank_2_reg, &right_bank_2_reg);
-              vtrn1_s8x2_in_place(&left_bank_3_reg, &right_bank_3_reg);
-              vtrn1_s8x2_in_place(&left_bank_4_reg, &right_bank_4_reg);
-            }
-          }
-          bias_data += kBiasIncrement;
-        }
-      } else {
-        // block_height == 1.
-        int8x16_t filter_reg_0_a;
-        int8x16_t filter_reg_1_a;
-        int8x16_t filter_reg_2_a;
-        int8x16_t filter_reg_0_b;
-        int8x16_t filter_reg_1_b;
-        int8x16_t filter_reg_2_b;
-
-        filter_reg_0_a = vld1q_s8(filter_block);
-        filter_reg_1_a = vld1q_s8(filter_block + 32);
-        filter_reg_2_a = vld1q_s8(filter_block + 64);
-        filter_reg_0_b = vld1q_s8(filter_block + 16);
-        filter_reg_1_b = vld1q_s8(filter_block + 16 + 32);
-        filter_reg_2_b = vld1q_s8(filter_block + 16 + 64);
-
-        const int8_t* scratch_data =
-            scratch_block_data + depth_micro_stride * j_depth;
-        typename QuantizationTypeImpl<quantization_type>::ExternalType*
-            output_data = output_block_data + 8 * j_depth;
-        const int8_t* input_data_0 = scratch_data;
-
-        const int32x4_t adjusted_bias_data_a = vld1q_s32(bias_data);
-        bias_data += kBiasIncrement;
-        const int32x4_t adjusted_bias_data_b = vld1q_s32(bias_data);
-        bias_data += kBiasIncrement;
-
-        const int32x4_t output_shift_a =
-            vld1q_s32(output_shift_per_channel + j_depth * 8);
-        const int32x4_t output_multiplier_a =
-            vld1q_s32(output_multiplier_per_channel + j_depth * 8);
-        const int32x4_t output_shift_b =
-            vld1q_s32(output_shift_per_channel + j_depth * 8 + 4);
-        const int32x4_t output_multiplier_b =
-            vld1q_s32(output_multiplier_per_channel + j_depth * 8 + 4);
-
-        // Load first sub-micro block of data into operational banks.
-        int8x16_t left_bank_0_reg_a = vld1q_s8(input_data_0);
-        int8x16_t left_bank_1_reg_a =
-            vld1q_s8(input_data_0 + workspace_height_stride);
-        int8x16_t left_bank_2_reg_a =
-            vld1q_s8(input_data_0 + 2 * workspace_height_stride);
-        int8x16_t left_bank_0_reg_b = vld1q_s8(input_data_0 + 16);
-        int8x16_t left_bank_1_reg_b =
-            vld1q_s8(input_data_0 + workspace_height_stride + 16);
-        int8x16_t left_bank_2_reg_b =
-            vld1q_s8(input_data_0 + 2 * workspace_height_stride + 16);
-
-        int8x16_t right_bank_0_reg_a;
-        int8x16_t right_bank_1_reg_a;
-        int8x16_t right_bank_2_reg_a;
-        int8x16_t right_bank_0_reg_b;
-        int8x16_t right_bank_1_reg_b;
-        int8x16_t right_bank_2_reg_b;
-
-        int32x4_t acc0_a;
-        int32x4_t acc0_b;
-
-        for (int i_width = 0; i_width < output_width_overall_micro_repeats;
-             ++i_width) {
-          const int output_width = i_width == output_width_micro_repeats
-                                       ? residual_width
-                                       : kFourOverStride;
-          TFLITE_DCHECK_LE(output_width * kStrideVal, 4);
-          const int8_t* input_data =
-              input_data_0 + width_micro_stride * i_width;
-          const bool no_right_block = i_width == output_width_micro_repeats &&
-                                      output_width_overall_micro_repeats ==
-                                          workspace_width_micro_repeats;
-
-          if (!no_right_block) {
-            // Load next sub-micro block of data.
-            right_bank_0_reg_a = vld1q_s8(input_data + width_micro_stride);
-            right_bank_1_reg_a = vld1q_s8(input_data + width_micro_stride +
-                                          workspace_height_stride);
-            right_bank_2_reg_a = vld1q_s8(input_data + width_micro_stride +
-                                          2 * workspace_height_stride);
-            right_bank_0_reg_b = vld1q_s8(input_data + width_micro_stride + 16);
-            right_bank_1_reg_b = vld1q_s8(input_data + width_micro_stride +
-                                          workspace_height_stride + 16);
-            right_bank_2_reg_b = vld1q_s8(input_data + width_micro_stride +
-                                          2 * workspace_height_stride + 16);
-          }
-
-          typename QuantizationTypeImpl<quantization_type>::ExternalType*
-              output_data_base = output_data + depth * 2 * i_width;
-
-          // Iterate over input width shifts within 4x4 blocks.
-          {
-            acc0_a = adjusted_bias_data_a;
-            acc0_b = adjusted_bias_data_b;
-
-            acc0_a = vdotq_s32(acc0_a, filter_reg_0_a, left_bank_0_reg_a);
-            acc0_a = vdotq_s32(acc0_a, filter_reg_1_a, left_bank_1_reg_a);
-            acc0_a = vdotq_s32(acc0_a, filter_reg_2_a, left_bank_2_reg_a);
-            acc0_b = vdotq_s32(acc0_b, filter_reg_0_b, left_bank_0_reg_b);
-            acc0_b = vdotq_s32(acc0_b, filter_reg_1_b, left_bank_1_reg_b);
-            acc0_b = vdotq_s32(acc0_b, filter_reg_2_b, left_bank_2_reg_b);
-
-            // Fixed-point multiplication.
-            acc0_a = vqrdmulhq_s32(acc0_a, output_multiplier_a);
-            acc0_b = vqrdmulhq_s32(acc0_b, output_multiplier_b);
-            acc0_a = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc0_a, output_shift_a);
-            acc0_b = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc0_b, output_shift_b);
-            // Add the output offset.
-            int16x8_t acc_s16_0_1 =
-                vcombine_s16(vqmovn_s32(acc0_a), vqmovn_s32(acc0_b));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            int8x8_t acc_u8 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8 = util_vmax_x8(acc_u8, output_activation_min_vec);
-            acc_u8 = util_vmin_x8(acc_u8, output_activation_max_vec);
-
-            vst1_s8(output_data_base, acc_u8);
-
-            left_bank_0_reg_a = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_0_reg_a)));
-            left_bank_1_reg_a = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_1_reg_a)));
-            left_bank_2_reg_a = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_2_reg_a)));
-            left_bank_0_reg_b = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_0_reg_b)));
-            left_bank_1_reg_b = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_1_reg_b)));
-            left_bank_2_reg_b = vreinterpretq_s8_u16(
-                vrev32q_u16(vreinterpretq_u16_s8(left_bank_2_reg_b)));
-            vtrn1_s8x2_in_place(&left_bank_0_reg_a, &right_bank_0_reg_a);
-            vtrn1_s8x2_in_place(&left_bank_1_reg_a, &right_bank_1_reg_a);
-            vtrn1_s8x2_in_place(&left_bank_2_reg_a, &right_bank_2_reg_a);
-            vtrn1_s8x2_in_place(&left_bank_0_reg_b, &right_bank_0_reg_b);
-            vtrn1_s8x2_in_place(&left_bank_1_reg_b, &right_bank_1_reg_b);
-            vtrn1_s8x2_in_place(&left_bank_2_reg_b, &right_bank_2_reg_b);
-          }
-
-          if (output_width > 1) {
-            acc0_a = adjusted_bias_data_a;
-            acc0_b = adjusted_bias_data_b;
-
-            acc0_a = vdotq_s32(acc0_a, filter_reg_0_a, left_bank_0_reg_a);
-            acc0_a = vdotq_s32(acc0_a, filter_reg_1_a, left_bank_1_reg_a);
-            acc0_a = vdotq_s32(acc0_a, filter_reg_2_a, left_bank_2_reg_a);
-            acc0_b = vdotq_s32(acc0_b, filter_reg_0_b, left_bank_0_reg_b);
-            acc0_b = vdotq_s32(acc0_b, filter_reg_1_b, left_bank_1_reg_b);
-            acc0_b = vdotq_s32(acc0_b, filter_reg_2_b, left_bank_2_reg_b);
-
-            // Fixed-point multiplication.
-            acc0_a = vqrdmulhq_s32(acc0_a, output_multiplier_a);
-            acc0_b = vqrdmulhq_s32(acc0_b, output_multiplier_b);
-            acc0_a = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc0_a, output_shift_a);
-            acc0_b = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc0_b, output_shift_b);
-            // Add the output offset.
-            int16x8_t acc_s16_0_1 =
-                vcombine_s16(vqmovn_s32(acc0_a), vqmovn_s32(acc0_b));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            int8x8_t acc_u8 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8 = util_vmax_x8(acc_u8, output_activation_min_vec);
-            acc_u8 = util_vmin_x8(acc_u8, output_activation_max_vec);
-
-            vst1_s8(output_data_base + depth, acc_u8);
-
-            left_bank_0_reg_a = right_bank_0_reg_a;
-            left_bank_1_reg_a = right_bank_1_reg_a;
-            left_bank_2_reg_a = right_bank_2_reg_a;
-            left_bank_0_reg_b = right_bank_0_reg_b;
-            left_bank_1_reg_b = right_bank_1_reg_b;
-            left_bank_2_reg_b = right_bank_2_reg_b;
-          }
-        }
-      }
-    }
-  }  // NOLINT(readability/fn_size) Manually unrolled.
-
-  static inline void Run(const int8_t* scratch_block_data,
-                         const int8_t* filter_workspace,
-                         const int32_t* bias_data, int8_t* output_block_data,
-                         const DepthwiseConvDotProdParams* function_params) {
-    KernelMacroBlockIntrinsics(scratch_block_data, filter_workspace, bias_data,
-                               output_block_data, function_params);
-  }
-};
-
-template <>
-struct KernelMacroBlock<
-    DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct,
-    QuantizationType::kPerChannelInt8,
-    DepthwiseConvDepthMultiplication::kUnitInputDepth,
-    /*stride=*/1> {
-  static inline int8x8_t vqmovxn_s16(int16x8_t x) { return vqmovn_s16(x); }
-  static inline int8x8_t util_vmin_x8(int8x8_t a, int8x8_t b) {
-    return vmin_s8(a, b);
-  }
-  static inline int8x8_t util_vmax_x8(int8x8_t a, int8x8_t b) {
-    return vmax_s8(a, b);
-  }
-  static inline int8x16_t util_vminq_x8(int8x16_t a, int8x16_t b) {
-    return vminq_s8(a, b);
-  }
-  static inline int8x16_t util_vmaxq_x8(int8x16_t a, int8x16_t b) {
-    return vmaxq_s8(a, b);
-  }
-
-  static inline void KernelMacroBlockIntrinsics(
-      const int8_t* scratch_block_data, const int8_t* filter_workspace,
-      const int32_t* bias_data, int8_t* output_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-    static constexpr QuantizationType quantization_type =
-        QuantizationType::kPerChannelInt8;
-
-    TFLITE_DCHECK_EQ(function_params->stride, 1);
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-    const int output_width_micro_repeats =
-        function_params->output_width_micro_repeats;
-    const int depth_micro_repeats = function_params->depth_micro_repeats;
-    const int output_depth = function_params->output_depth;
-
-    const int output_width_overall_micro_repeats =
-        function_params->output_width_overall_micro_repeats;
-    const int block_height = function_params->outbound_block_height;
-    const int residual_width = function_params->output_residual_width;
-    const int output_height_stride = function_params->output_height_stride;
-    constexpr int kBiasIncrement = 4;
-
-    TFLITE_DCHECK(depth_micro_repeats > 0);
-
-    const int32_t output_activation_min =
-        function_params->quantized_activation_min;
-    const int32_t output_activation_max =
-        function_params->quantized_activation_max;
-    const int32_t output_offset = function_params->output_offset;
-    const int32_t* output_shift_per_channel =
-        function_params->output_shift_per_channel;
-    const int32_t* output_multiplier_per_channel =
-        function_params->output_multiplier_per_channel;
-    if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-      TFLITE_DCHECK_GE(output_activation_min, 0);
-      TFLITE_DCHECK_LT(output_activation_min, 256);
-      TFLITE_DCHECK_GE(output_activation_max, 0);
-      TFLITE_DCHECK_LT(output_activation_max, 256);
-    } else {
-      TFLITE_DCHECK_GE(output_activation_min, -128);
-      TFLITE_DCHECK_LT(output_activation_min, 128);
-      TFLITE_DCHECK_GE(output_activation_max, -128);
-      TFLITE_DCHECK_LT(output_activation_max, 128);
-      TFLITE_DCHECK_NE(output_shift_per_channel, nullptr);
-      TFLITE_DCHECK_NE(output_multiplier_per_channel, nullptr);
-    }
-    TFLITE_DCHECK_GE(output_offset, -32878);
-    TFLITE_DCHECK_LT(output_offset, 32768);
-
-    const int16x8_t output_offset_vec =
-        vdupq_n_s16(static_cast<int16_t>(output_offset));
-    const int8x16_t output_activation_min_vec =
-        vdupq_n_s8(static_cast<int8_t>(output_activation_min));
-    const int8x16_t output_activation_max_vec =
-        vdupq_n_s8(static_cast<int8_t>(output_activation_max));
-
-    typename QuantizationTypeImpl<quantization_type>::ExternalType*
-        output_data_depthwise = output_block_data;
-    for (int j_depth = 0; j_depth < depth_micro_repeats; ++j_depth) {
-      // Simulate NEON-register transposition of subset of filter.
-      int8x16_t filter_reg_0_a;
-      int8x16_t filter_reg_0_b;
-      int8x16_t filter_reg_1_a;
-      int8x16_t filter_reg_1_b;
-      int8x16_t filter_reg_2_a;
-      int8x16_t filter_reg_2_b;
-      int8x16_t filter_reg_0_a_shifted;
-      int8x16_t filter_reg_1_a_shifted;
-      int8x16_t filter_reg_2_a_shifted;
-
-      filter_reg_0_a = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_0_b = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_1_a = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_1_b = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_2_a = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_2_b = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-
-      filter_reg_0_a_shifted = vreinterpretq_s8_u32(
-          vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_0_a), 8));
-      filter_reg_1_a_shifted = vreinterpretq_s8_u32(
-          vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_1_a), 8));
-      filter_reg_2_a_shifted = vreinterpretq_s8_u32(
-          vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_2_a), 8));
-
-      // When output_width_micro_repeats < output_width_overall_micro_repeats,
-      // 0 < residual_width <= 2, and so residual_width == 1 is then true iff
-      // residual_width < 2.
-      const int adjusted_width_micro_repeats =
-          (output_width_micro_repeats < output_width_overall_micro_repeats) &&
-                  (residual_width < 4)
-              ? output_width_micro_repeats
-              : output_width_overall_micro_repeats;
-
-      if (block_height == 4) {
-        for (int s = 0; s < 2; ++s) {
-          // Work through one slice, by row, at a time.
-          typename QuantizationTypeImpl<quantization_type>::ExternalType*
-              output_data_base = output_data_depthwise + 4 * s;
-
-          const int8_t* next_input_data = scratch_block_data;
-          typename QuantizationTypeImpl<quantization_type>::ExternalType*
-              output_data = output_data_base;
-
-          const int32x4_t adjusted_bias_data = vld1q_s32(bias_data);
-          bias_data += kBiasIncrement;
-
-          const int32x4_t output_shift =
-              vld1q_s32(output_shift_per_channel + j_depth * 8 + 4 * s);
-          const int32x4_t output_multiplier =
-              vld1q_s32(output_multiplier_per_channel + j_depth * 8 + 4 * s);
-
-          int8x16_t input_bank_a_reg;  //  left 0, right 0, left 1, right 1.
-          int8x16_t input_bank_b_reg;  //  left 2, right 2, left 3, right 3.
-          int8x16_t input_bank_c_reg;  //  left 4, right 4, left 5, right 5.
-
-          // Load first sub-micro block of data into operational banks.
-          input_bank_a_reg =
-              vld1q_dup_s8x4(next_input_data);  // Load lane 0, avoiding
-                                                // uninitialized variable.
-          input_bank_a_reg = vld1q_lane_8x4(
-              next_input_data + workspace_height_stride, input_bank_a_reg, 2);
-          input_bank_b_reg = vld1q_dup_s8x4(
-              next_input_data +
-              2 * workspace_height_stride);  // Load lane 0, avoiding
-                                             // uninitialized variable.
-          input_bank_b_reg =
-              vld1q_lane_8x4(next_input_data + 3 * workspace_height_stride,
-                             input_bank_b_reg, 2);
-          input_bank_c_reg = vld1q_dup_s8x4(
-              next_input_data +
-              4 * workspace_height_stride);  // Load lane 0, avoiding
-                                             // uninitialized variable.
-          input_bank_c_reg =
-              vld1q_lane_8x4(next_input_data + 5 * workspace_height_stride,
-                             input_bank_c_reg, 2);
-
-          int32x4_t acc0;
-          int32x4_t acc1;
-          int32x4_t acc2;
-          int32x4_t acc3;
-
-          acc0 = adjusted_bias_data;
-          acc1 = adjusted_bias_data;
-          acc2 = adjusted_bias_data;
-          acc3 = adjusted_bias_data;
-
-          acc0 = vdotq_four_lane_s32(acc0, filter_reg_2_a, input_bank_b_reg, 0);
-          acc1 = vdotq_four_lane_s32(acc1, filter_reg_1_a, input_bank_b_reg, 0);
-          acc2 = vdotq_four_lane_s32(acc2, filter_reg_0_a, input_bank_b_reg, 0);
-          acc3 = vdotq_four_lane_s32(acc3, filter_reg_0_a, input_bank_b_reg, 2);
-
-          int i_width = 0;
-          for (; i_width < adjusted_width_micro_repeats; ++i_width) {
-            next_input_data += 4;
-
-            // Iterate over input width shifts within 4x4 blocks.
-            {
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_0_a, input_bank_a_reg,
-                                         0);
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_1_a, input_bank_a_reg,
-                                         2);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_0_a, input_bank_a_reg,
-                                         2);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_2_a, input_bank_b_reg,
-                                         2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_1_a, input_bank_b_reg,
-                                         2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_2_a, input_bank_c_reg,
-                                         0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_1_a, input_bank_c_reg,
-                                         0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_2_a, input_bank_c_reg,
-                                         2);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc0, output_shift);
-              acc1 = vqrdmulhq_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc1, output_shift);
-              acc2 = vqrdmulhq_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc2, output_shift);
-              acc3 = vqrdmulhq_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc3, output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              int8x16_t acc_u8_all = vcombine_s8(vqmovxn_s16(acc_s16_0_1),
-                                                 vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_s8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_s8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_s8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_s8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              output_data += output_depth;
-            }
-            // Load next sub-micro block of data.
-            input_bank_a_reg =
-                vld1q_lane_8x4(next_input_data, input_bank_a_reg, 1);
-            input_bank_a_reg = vld1q_lane_8x4(
-                next_input_data + workspace_height_stride, input_bank_a_reg, 3);
-            input_bank_b_reg =
-                vld1q_lane_8x4(next_input_data + 2 * workspace_height_stride,
-                               input_bank_b_reg, 1);
-            input_bank_b_reg =
-                vld1q_lane_8x4(next_input_data + 3 * workspace_height_stride,
-                               input_bank_b_reg, 3);
-            input_bank_c_reg =
-                vld1q_lane_8x4(next_input_data + 4 * workspace_height_stride,
-                               input_bank_c_reg, 1);
-            input_bank_c_reg =
-                vld1q_lane_8x4(next_input_data + 5 * workspace_height_stride,
-                               input_bank_c_reg, 3);
-
-            {
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_0_a_shifted,
-                                         input_bank_a_reg, 0);
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_1_a_shifted,
-                                         input_bank_a_reg, 2);
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_2_a_shifted,
-                                         input_bank_b_reg, 0);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_0_a_shifted,
-                                         input_bank_a_reg, 2);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_1_a_shifted,
-                                         input_bank_b_reg, 0);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_2_a_shifted,
-                                         input_bank_b_reg, 2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_0_a_shifted,
-                                         input_bank_b_reg, 0);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_1_a_shifted,
-                                         input_bank_b_reg, 2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_2_a_shifted,
-                                         input_bank_c_reg, 0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_0_a_shifted,
-                                         input_bank_b_reg, 2);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_1_a_shifted,
-                                         input_bank_c_reg, 0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_2_a_shifted,
-                                         input_bank_c_reg, 2);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc0, output_shift);
-              acc1 = vqrdmulhq_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc1, output_shift);
-              acc2 = vqrdmulhq_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc2, output_shift);
-              acc3 = vqrdmulhq_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc3, output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              int8x16_t acc_u8_all = vcombine_s8(vqmovxn_s16(acc_s16_0_1),
-                                                 vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_s8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_s8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_s8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_s8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              input_bank_a_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_a_reg), 16));
-              input_bank_b_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_b_reg), 16));
-              input_bank_c_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_c_reg), 16));
-
-              output_data += output_depth;
-            }
-
-            {
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_0_a, input_bank_a_reg,
-                                         0);
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_1_a, input_bank_a_reg,
-                                         2);
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_2_a, input_bank_b_reg,
-                                         0);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_0_a, input_bank_a_reg,
-                                         2);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_1_a, input_bank_b_reg,
-                                         0);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_2_a, input_bank_b_reg,
-                                         2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_0_a, input_bank_b_reg,
-                                         0);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_1_a, input_bank_b_reg,
-                                         2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_2_a, input_bank_c_reg,
-                                         0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_0_a, input_bank_b_reg,
-                                         2);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_1_a, input_bank_c_reg,
-                                         0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_2_a, input_bank_c_reg,
-                                         2);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc0, output_shift);
-              acc1 = vqrdmulhq_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc1, output_shift);
-              acc2 = vqrdmulhq_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc2, output_shift);
-              acc3 = vqrdmulhq_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc3, output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              int8x16_t acc_u8_all = vcombine_s8(vqmovxn_s16(acc_s16_0_1),
-                                                 vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_s8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_s8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_s8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_s8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              output_data += output_depth;
-            }
-
-            {
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_0_a_shifted,
-                                         input_bank_a_reg, 0);
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_1_a_shifted,
-                                         input_bank_a_reg, 2);
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_2_a_shifted,
-                                         input_bank_b_reg, 0);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_0_a_shifted,
-                                         input_bank_a_reg, 2);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_1_a_shifted,
-                                         input_bank_b_reg, 0);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_2_a_shifted,
-                                         input_bank_b_reg, 2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_0_a_shifted,
-                                         input_bank_b_reg, 0);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_1_a_shifted,
-                                         input_bank_b_reg, 2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_2_a_shifted,
-                                         input_bank_c_reg, 0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_0_a_shifted,
-                                         input_bank_b_reg, 2);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_1_a_shifted,
-                                         input_bank_c_reg, 0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_2_a_shifted,
-                                         input_bank_c_reg, 2);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc0, output_shift);
-              acc1 = vqrdmulhq_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc1, output_shift);
-              acc2 = vqrdmulhq_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc2, output_shift);
-              acc3 = vqrdmulhq_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc3, output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              int8x16_t acc_u8_all = vcombine_s8(vqmovxn_s16(acc_s16_0_1),
-                                                 vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_s8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_s8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_s8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_s8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              input_bank_a_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_a_reg), 16));
-              input_bank_b_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_b_reg), 16));
-              input_bank_c_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_c_reg), 16));
-
-              output_data += output_depth;
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_2_a, input_bank_b_reg,
-                                         0);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_1_a, input_bank_b_reg,
-                                         0);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_0_a, input_bank_b_reg,
-                                         0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_0_a, input_bank_b_reg,
-                                         2);
-            }
-          }
-
-          if (i_width < output_width_overall_micro_repeats) {
-            next_input_data += 4;
-            const int output_width = residual_width;
-
-            // Load next sub-micro block of data.
-            input_bank_a_reg =
-                vld1q_lane_8x4(next_input_data, input_bank_a_reg, 1);
-            input_bank_a_reg = vld1q_lane_8x4(
-                next_input_data + workspace_height_stride, input_bank_a_reg, 3);
-            input_bank_b_reg =
-                vld1q_lane_8x4(next_input_data + 2 * workspace_height_stride,
-                               input_bank_b_reg, 1);
-            input_bank_b_reg =
-                vld1q_lane_8x4(next_input_data + 3 * workspace_height_stride,
-                               input_bank_b_reg, 3);
-            input_bank_c_reg =
-                vld1q_lane_8x4(next_input_data + 4 * workspace_height_stride,
-                               input_bank_c_reg, 1);
-            input_bank_c_reg =
-                vld1q_lane_8x4(next_input_data + 5 * workspace_height_stride,
-                               input_bank_c_reg, 3);
-
-            // Iterate over input width shifts within 4x4 blocks.
-            for (int x = 0; x < output_width; ++x) {
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_0_a, input_bank_a_reg,
-                                         0);
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_1_a, input_bank_a_reg,
-                                         2);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_0_a, input_bank_a_reg,
-                                         2);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_2_a, input_bank_b_reg,
-                                         2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_1_a, input_bank_b_reg,
-                                         2);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_2_a, input_bank_c_reg,
-                                         0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_1_a, input_bank_c_reg,
-                                         0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_2_a, input_bank_c_reg,
-                                         2);
-
-              // Fixed-point multiplication.
-              acc0 = vqrdmulhq_s32(acc0, output_multiplier);
-              acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc0, output_shift);
-              acc1 = vqrdmulhq_s32(acc1, output_multiplier);
-              acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc1, output_shift);
-              acc2 = vqrdmulhq_s32(acc2, output_multiplier);
-              acc2 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc2, output_shift);
-              acc3 = vqrdmulhq_s32(acc3, output_multiplier);
-              acc3 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                  acc3, output_shift);
-              // Add the output offset.
-              int16x8_t acc_s16_0_1 =
-                  vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-              int16x8_t acc_s16_2_3 =
-                  vcombine_s16(vqmovn_s32(acc2), vqmovn_s32(acc3));
-              acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-              acc_s16_2_3 = vqaddq_s16(acc_s16_2_3, output_offset_vec);
-              // Apply the activation function.
-              int8x16_t acc_u8_all = vcombine_s8(vqmovxn_s16(acc_s16_0_1),
-                                                 vqmovxn_s16(acc_s16_2_3));
-              acc_u8_all = util_vmaxq_x8(acc_u8_all, output_activation_min_vec);
-              acc_u8_all = util_vminq_x8(acc_u8_all, output_activation_max_vec);
-
-              vst1q_lane_s8x4(output_data, acc_u8_all, 0);
-              vst1q_lane_s8x4(output_data + output_height_stride, acc_u8_all,
-                              1);
-              vst1q_lane_s8x4(output_data + 2 * output_height_stride,
-                              acc_u8_all, 2);
-              vst1q_lane_s8x4(output_data + 3 * output_height_stride,
-                              acc_u8_all, 3);
-
-              input_bank_a_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_a_reg), 8));
-              input_bank_b_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_b_reg), 8));
-              input_bank_c_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_c_reg), 8));
-
-              output_data += output_depth;
-
-              acc0 = adjusted_bias_data;
-              acc1 = adjusted_bias_data;
-              acc2 = adjusted_bias_data;
-              acc3 = adjusted_bias_data;
-
-              acc0 = vdotq_four_lane_s32(acc0, filter_reg_2_a, input_bank_b_reg,
-                                         0);
-              acc1 = vdotq_four_lane_s32(acc1, filter_reg_1_a, input_bank_b_reg,
-                                         0);
-              acc2 = vdotq_four_lane_s32(acc2, filter_reg_0_a, input_bank_b_reg,
-                                         0);
-              acc3 = vdotq_four_lane_s32(acc3, filter_reg_0_a, input_bank_b_reg,
-                                         2);
-            }
-          }
-          // scratch_block_data += 4 * workspace_height_stride;
-          output_data_base += 4 * output_height_stride;
-
-          // Move to next sub-block: advance to second set of filters, to new
-          // bias.
-          filter_reg_0_a = filter_reg_0_b;
-          filter_reg_1_a = filter_reg_1_b;
-          filter_reg_2_a = filter_reg_2_b;
-          filter_reg_0_a_shifted = vreinterpretq_s8_u32(
-              vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_0_a), 8));
-          filter_reg_1_a_shifted = vreinterpretq_s8_u32(
-              vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_1_a), 8));
-          filter_reg_2_a_shifted = vreinterpretq_s8_u32(
-              vshlq_n_u32(vreinterpretq_u32_s8(filter_reg_2_a), 8));
-        }
-      } else {
-        // Block height < 4.
-        typename QuantizationTypeImpl<quantization_type>::ExternalType*
-            output_data_base = output_data_depthwise;
-
-        const int32x4_t adjusted_bias_data_a = vld1q_s32(bias_data);
-        bias_data += kBiasIncrement;
-        const int32x4_t adjusted_bias_data_b = vld1q_s32(bias_data);
-        bias_data += kBiasIncrement;
-
-        const int32x4_t output_shift_a =
-            vld1q_s32(output_shift_per_channel + j_depth * 8);
-        const int32x4_t output_multiplier_a =
-            vld1q_s32(output_multiplier_per_channel + j_depth * 8);
-        const int32x4_t output_shift_b =
-            vld1q_s32(output_shift_per_channel + j_depth * 8 + 4);
-        const int32x4_t output_multiplier_b =
-            vld1q_s32(output_multiplier_per_channel + j_depth * 8 + 4);
-
-        for (int k_height = 0; k_height < block_height; ++k_height) {
-          const int8_t* next_input_data =
-              scratch_block_data + k_height * workspace_height_stride;
-          typename QuantizationTypeImpl<quantization_type>::ExternalType*
-              output_data = output_data_base;
-
-          int8x16_t input_bank_p_reg;  //  left 0, right 0, left 1, right 1.
-          int8x16_t input_bank_q_reg;  //  left 2, right 2, left 3, right 3.
-
-          // Load first sub-micro block of data into operational banks.
-          input_bank_p_reg =
-              vld1q_dup_s8x4(next_input_data);  // Load lane 0, avoiding
-                                                // uninitialized variable.
-          input_bank_p_reg = vld1q_lane_8x4(
-              next_input_data + workspace_height_stride, input_bank_p_reg, 2);
-          input_bank_q_reg = vld1q_dup_s8x4(
-              next_input_data +
-              2 * workspace_height_stride);  // Load lane 0, avoiding
-                                             // uninitialized variable.
-
-          for (int i_width = 0; i_width < output_width_overall_micro_repeats;
-               ++i_width) {
-            next_input_data += 4;
-            const int output_width =
-                i_width == output_width_micro_repeats ? residual_width : 4;
-
-            // Load next sub-micro block of data.
-            input_bank_p_reg =
-                vld1q_lane_8x4(next_input_data, input_bank_p_reg, 1);
-            input_bank_p_reg = vld1q_lane_8x4(
-                next_input_data + workspace_height_stride, input_bank_p_reg, 3);
-            input_bank_q_reg =
-                vld1q_lane_8x4(next_input_data + 2 * workspace_height_stride,
-                               input_bank_q_reg, 1);
-            // Iterate over input width shifts within 4x4 blocks.
-            for (int x = 0; x < output_width; ++x) {
-              int32x4_t acc_a = adjusted_bias_data_a;
-              int32x4_t acc_b = adjusted_bias_data_b;
-              acc_a = vdotq_four_lane_s32(acc_a, filter_reg_0_a,
-                                          input_bank_p_reg, 0);
-              acc_a = vdotq_four_lane_s32(acc_a, filter_reg_1_a,
-                                          input_bank_p_reg, 2);
-              acc_a = vdotq_four_lane_s32(acc_a, filter_reg_2_a,
-                                          input_bank_q_reg, 0);
-              acc_b = vdotq_four_lane_s32(acc_b, filter_reg_0_b,
-                                          input_bank_p_reg, 0);
-              acc_b = vdotq_four_lane_s32(acc_b, filter_reg_1_b,
-                                          input_bank_p_reg, 2);
-              acc_b = vdotq_four_lane_s32(acc_b, filter_reg_2_b,
-                                          input_bank_q_reg, 0);
-
-              // Fixed-point multiplication.
-              acc_a = vqrdmulhq_s32(acc_a, output_multiplier_a);
-              acc_b = vqrdmulhq_s32(acc_b, output_multiplier_b);
-              acc_a =
-                  DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                      acc_a, output_shift_a);
-              acc_b =
-                  DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                      acc_b, output_shift_b);
-              // Add the output offset.
-              int16x8_t acc_s16_0_0 =
-                  vcombine_s16(vqmovn_s32(acc_a), vqmovn_s32(acc_b));
-              acc_s16_0_0 = vqaddq_s16(acc_s16_0_0, output_offset_vec);
-              // Apply the activation function.
-              int8x8_t acc_u8_0_0 = vqmovxn_s16(acc_s16_0_0);
-              acc_u8_0_0 = util_vmax_x8(acc_u8_0_0,
-                                        vget_low_s8(output_activation_min_vec));
-              acc_u8_0_0 = util_vmin_x8(acc_u8_0_0,
-                                        vget_low_s8(output_activation_max_vec));
-
-              vst1_s8(output_data, acc_u8_0_0);
-
-              input_bank_p_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_p_reg), 8));
-              input_bank_q_reg = vreinterpretq_s8_u64(
-                  vshrq_n_u64(vreinterpretq_u64_s8(input_bank_q_reg), 8));
-
-              output_data += output_depth;
-            }
-          }
-          output_data_base += output_height_stride;
-        }
-      }
-      output_data_depthwise += 8;
-    }
-  }  // NOLINT(readability/fn_size) Manually unrolled.
-
-  static inline void Run(const int8_t* scratch_block_data,
-                         const int8_t* filter_workspace,
-                         const int32_t* bias_data, int8_t* output_block_data,
-                         const DepthwiseConvDotProdParams* function_params) {
-    KernelMacroBlockIntrinsics(scratch_block_data, filter_workspace, bias_data,
-                               output_block_data, function_params);
-  }
-};
-
-template <>
-struct KernelMacroBlock<
-    DepthwiseConvImplementation::kUseIntrinsics3x3DotProduct,
-    QuantizationType::kPerChannelInt8,
-    DepthwiseConvDepthMultiplication::kUnitInputDepth,
-    /*stride=*/2> {
-  static inline int8x8_t vqmovxn_s16(int16x8_t x) { return vqmovn_s16(x); }
-  static inline int8x8_t util_vmin_x8(int8x8_t a, int8x8_t b) {
-    return vmin_s8(a, b);
-  }
-  static inline int8x8_t util_vmax_x8(int8x8_t a, int8x8_t b) {
-    return vmax_s8(a, b);
-  }
-
-  static inline void KernelMacroBlockIntrinsics(
-      const int8_t* scratch_block_data, const int8_t* filter_workspace,
-      const int32_t* bias_data, int8_t* output_block_data,
-      const DepthwiseConvDotProdParams* function_params) {
-    static constexpr QuantizationType quantization_type =
-        QuantizationType::kPerChannelInt8;
-
-    const int workspace_height_stride =
-        function_params->workspace_height_stride;
-    const int output_width_micro_repeats =
-        function_params->output_width_micro_repeats;
-    const int depth_micro_repeats = function_params->depth_micro_repeats;
-    const int output_depth = function_params->output_depth;
-    constexpr int kStrideVal = 2;
-    TFLITE_DCHECK_EQ(function_params->stride, kStrideVal);
-
-    const int output_width_overall_micro_repeats =
-        function_params->output_width_overall_micro_repeats;
-    const int block_height = function_params->outbound_block_height;
-    const int residual_width = function_params->output_residual_width;
-    const int output_height_stride = function_params->output_height_stride;
-    constexpr int kBiasIncrement = 4;
-
-    const int32_t output_activation_min =
-        function_params->quantized_activation_min;
-    const int32_t output_activation_max =
-        function_params->quantized_activation_max;
-    const int32_t output_offset = function_params->output_offset;
-    const int32_t* output_shift_per_channel =
-        function_params->output_shift_per_channel;
-    const int32_t* output_multiplier_per_channel =
-        function_params->output_multiplier_per_channel;
-    if (quantization_type == QuantizationType::kNonPerChannelUint8) {
-      TFLITE_DCHECK_GE(output_activation_min, 0);
-      TFLITE_DCHECK_LT(output_activation_min, 256);
-      TFLITE_DCHECK_GE(output_activation_max, 0);
-      TFLITE_DCHECK_LT(output_activation_max, 256);
-    } else {
-      TFLITE_DCHECK_GE(output_activation_min, -128);
-      TFLITE_DCHECK_LT(output_activation_min, 128);
-      TFLITE_DCHECK_GE(output_activation_max, -128);
-      TFLITE_DCHECK_LT(output_activation_max, 128);
-      TFLITE_DCHECK_NE(output_shift_per_channel, nullptr);
-      TFLITE_DCHECK_NE(output_multiplier_per_channel, nullptr);
-    }
-    TFLITE_DCHECK_GE(output_offset, -32878);
-    TFLITE_DCHECK_LT(output_offset, 32768);
-
-    TFLITE_DCHECK_GE(depth_micro_repeats, 1);
-
-    const int16x8_t output_offset_vec =
-        vdupq_n_s16(static_cast<int16_t>(output_offset));
-    const int8x16_t output_activation_min_vec =
-        vdupq_n_s8(static_cast<int8_t>(output_activation_min));
-    const int8x16_t output_activation_max_vec =
-        vdupq_n_s8(static_cast<int8_t>(output_activation_max));
-
-    for (int j_depth = 0; j_depth < (depth_micro_repeats * 1 + 0); ++j_depth) {
-      int8x16_t filter_reg_0_a;
-      int8x16_t filter_reg_0_b;
-      int8x16_t filter_reg_1_a;
-      int8x16_t filter_reg_1_b;
-      int8x16_t filter_reg_2_a;
-      int8x16_t filter_reg_2_b;
-
-      filter_reg_0_a = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_0_b = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_1_a = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_1_b = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_2_a = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-      filter_reg_2_b = vld1q_s8(filter_workspace);
-      filter_workspace += 16;
-
-      const int32x4_t adjusted_bias_data_s_0 = vld1q_s32(bias_data);
-      bias_data += kBiasIncrement;
-      const int32x4_t adjusted_bias_data_s_1 = vld1q_s32(bias_data);
-      bias_data += kBiasIncrement;
-
-      const int32x4_t output_shift_s_0 =
-          vld1q_s32(output_shift_per_channel + j_depth * 8);
-      const int32x4_t output_multiplier_s_0 =
-          vld1q_s32(output_multiplier_per_channel + j_depth * 8);
-      const int32x4_t output_shift_s_1 =
-          vld1q_s32(output_shift_per_channel + j_depth * 8 + 4);
-      const int32x4_t output_multiplier_s_1 =
-          vld1q_s32(output_multiplier_per_channel + j_depth * 8 + 4);
-
-      if (block_height == 2) {
-        const int8_t* scratch_data = scratch_block_data;
-        typename QuantizationTypeImpl<quantization_type>::ExternalType*
-            output_data = output_block_data + 8 * j_depth;
-
-        int8x16_t input_bank_a_reg;  //  left 0, right 0, left 1, right 1.
-        int8x16_t input_bank_b_reg;  //  left 2, right 2, left 3, right 3.
-        int8x16_t input_bank_c_reg;  //  left 4, right 4, xxx, xxx.
-
-        // Load first sub-micro block of data into operational banks.
-        input_bank_a_reg =
-            vld1q_dup_s8x4(scratch_data);  // Load lane 0, avoiding
-                                           // uninitialized variable.
-        input_bank_a_reg = vld1q_lane_8x4(
-            scratch_data + workspace_height_stride, input_bank_a_reg, 2);
-        input_bank_b_reg = vld1q_dup_s8x4(
-            scratch_data +
-            2 * workspace_height_stride);  // Load lane 0, avoiding
-                                           // uninitialized variable.
-        input_bank_b_reg = vld1q_lane_8x4(
-            scratch_data + 3 * workspace_height_stride, input_bank_b_reg, 2);
-        input_bank_c_reg = vld1q_dup_s8x4(
-            scratch_data +
-            4 * workspace_height_stride);  // Load lane 0, avoiding
-                                           // uninitialized variable.
-
-        int32x4_t acc0;
-        int32x4_t acc1;
-
-        // When output_width_micro_repeats < output_width_overall_micro_repeats,
-        // 0 < residual_width <= 2, and so residual_width == 1 is then true iff
-        // residual_width < 2.
-        const int adjusted_width_micro_repeats =
-            (output_width_micro_repeats < output_width_overall_micro_repeats) &&
-                    (residual_width < 2)
-                ? output_width_micro_repeats
-                : output_width_overall_micro_repeats;
-
-        int i_width = 0;
-        for (; i_width < adjusted_width_micro_repeats; ++i_width) {
-          const int8_t* input_data = scratch_data + 4 + 4 * i_width;
-
-          // Load next sub-micro block of data.
-          input_bank_a_reg = vld1q_lane_8x4(input_data, input_bank_a_reg, 1);
-          input_bank_a_reg = vld1q_lane_8x4(
-              input_data + workspace_height_stride, input_bank_a_reg, 3);
-          input_bank_b_reg = vld1q_lane_8x4(
-              input_data + 2 * workspace_height_stride, input_bank_b_reg, 1);
-          input_bank_b_reg = vld1q_lane_8x4(
-              input_data + 3 * workspace_height_stride, input_bank_b_reg, 3);
-          input_bank_c_reg = vld1q_lane_8x4(
-              input_data + 4 * workspace_height_stride, input_bank_c_reg, 1);
-
-          int16x8_t acc_s16_0_1;
-          int8x8_t acc_u8_0_1;
-          // Iterate over input width shifts within 4x4 blocks.
-          {
-            acc0 = adjusted_bias_data_s_0;
-            acc1 = adjusted_bias_data_s_0;
-
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_0_a, input_bank_a_reg, 0);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_1_a, input_bank_a_reg, 2);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_2_a, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_0_a, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_1_a, input_bank_b_reg, 2);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_2_a, input_bank_c_reg, 0);
-
-            // Fixed-point multiplication.
-            acc0 = vqrdmulhq_s32(acc0, output_multiplier_s_0);
-            acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc0, output_shift_s_0);
-            acc1 = vqrdmulhq_s32(acc1, output_multiplier_s_0);
-            acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc1, output_shift_s_0);
-            // Add the output offset.
-            acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            acc_u8_0_1 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8_0_1 = util_vmax_x8(acc_u8_0_1,
-                                      vget_low_s8(output_activation_min_vec));
-            acc_u8_0_1 = util_vmin_x8(acc_u8_0_1,
-                                      vget_low_s8(output_activation_max_vec));
-
-            vst1_lane_s8x4(output_data, acc_u8_0_1, 0);
-            vst1_lane_s8x4(output_data + output_height_stride, acc_u8_0_1, 1);
-
-            acc0 = adjusted_bias_data_s_1;
-            acc1 = adjusted_bias_data_s_1;
-
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_0_b, input_bank_a_reg, 0);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_1_b, input_bank_a_reg, 2);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_2_b, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_0_b, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_1_b, input_bank_b_reg, 2);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_2_b, input_bank_c_reg, 0);
-
-            // Fixed-point multiplication.
-            acc0 = vqrdmulhq_s32(acc0, output_multiplier_s_1);
-            acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc0, output_shift_s_1);
-            acc1 = vqrdmulhq_s32(acc1, output_multiplier_s_1);
-            acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc1, output_shift_s_1);
-            // Add the output offset.
-            acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            acc_u8_0_1 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8_0_1 = util_vmax_x8(acc_u8_0_1,
-                                      vget_low_s8(output_activation_min_vec));
-            acc_u8_0_1 = util_vmin_x8(acc_u8_0_1,
-                                      vget_low_s8(output_activation_max_vec));
-
-            vst1_lane_s8x4(output_data + 4, acc_u8_0_1, 0);
-            vst1_lane_s8x4(output_data + 4 + output_height_stride, acc_u8_0_1,
-                           1);
-
-            input_bank_a_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_a_reg), 16));
-            input_bank_b_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_b_reg), 16));
-            input_bank_c_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_c_reg), 16));
-
-            output_data += output_depth;
-          }
-
-          // output_width == four_over_stride.
-          acc0 = adjusted_bias_data_s_0;
-          acc1 = adjusted_bias_data_s_0;
-
-          acc0 = vdotq_four_lane_s32(acc0, filter_reg_0_a, input_bank_a_reg, 0);
-          acc0 = vdotq_four_lane_s32(acc0, filter_reg_1_a, input_bank_a_reg, 2);
-          acc0 = vdotq_four_lane_s32(acc0, filter_reg_2_a, input_bank_b_reg, 0);
-          acc1 = vdotq_four_lane_s32(acc1, filter_reg_0_a, input_bank_b_reg, 0);
-          acc1 = vdotq_four_lane_s32(acc1, filter_reg_1_a, input_bank_b_reg, 2);
-          acc1 = vdotq_four_lane_s32(acc1, filter_reg_2_a, input_bank_c_reg, 0);
-
-          // Fixed-point multiplication.
-          acc0 = vqrdmulhq_s32(acc0, output_multiplier_s_0);
-          acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-              acc0, output_shift_s_0);
-          acc1 = vqrdmulhq_s32(acc1, output_multiplier_s_0);
-          acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-              acc1, output_shift_s_0);
-          // Add the output offset.
-          acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-          acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-          // Apply the activation function.
-          acc_u8_0_1 = vqmovxn_s16(acc_s16_0_1);
-          acc_u8_0_1 =
-              util_vmax_x8(acc_u8_0_1, vget_low_s8(output_activation_min_vec));
-          acc_u8_0_1 =
-              util_vmin_x8(acc_u8_0_1, vget_low_s8(output_activation_max_vec));
-
-          vst1_lane_s8x4(output_data, acc_u8_0_1, 0);
-          vst1_lane_s8x4(output_data + output_height_stride, acc_u8_0_1, 1);
-
-          acc0 = adjusted_bias_data_s_1;
-          acc1 = adjusted_bias_data_s_1;
-
-          acc0 = vdotq_four_lane_s32(acc0, filter_reg_0_b, input_bank_a_reg, 0);
-          acc0 = vdotq_four_lane_s32(acc0, filter_reg_1_b, input_bank_a_reg, 2);
-          acc0 = vdotq_four_lane_s32(acc0, filter_reg_2_b, input_bank_b_reg, 0);
-          acc1 = vdotq_four_lane_s32(acc1, filter_reg_0_b, input_bank_b_reg, 0);
-          acc1 = vdotq_four_lane_s32(acc1, filter_reg_1_b, input_bank_b_reg, 2);
-          acc1 = vdotq_four_lane_s32(acc1, filter_reg_2_b, input_bank_c_reg, 0);
-
-          // Fixed-point multiplication.
-          acc0 = vqrdmulhq_s32(acc0, output_multiplier_s_1);
-          acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-              acc0, output_shift_s_1);
-          acc1 = vqrdmulhq_s32(acc1, output_multiplier_s_1);
-          acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-              acc1, output_shift_s_1);
-          // Add the output offset.
-          acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-          acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-          // Apply the activation function.
-          acc_u8_0_1 = vqmovxn_s16(acc_s16_0_1);
-          acc_u8_0_1 =
-              util_vmax_x8(acc_u8_0_1, vget_low_s8(output_activation_min_vec));
-          acc_u8_0_1 =
-              util_vmin_x8(acc_u8_0_1, vget_low_s8(output_activation_max_vec));
-
-          vst1_lane_s8x4(output_data + 4, acc_u8_0_1, 0);
-          vst1_lane_s8x4(output_data + 4 + output_height_stride, acc_u8_0_1, 1);
-
-          input_bank_a_reg = vreinterpretq_s8_u64(
-              vshrq_n_u64(vreinterpretq_u64_s8(input_bank_a_reg), 16));
-          input_bank_b_reg = vreinterpretq_s8_u64(
-              vshrq_n_u64(vreinterpretq_u64_s8(input_bank_b_reg), 16));
-          input_bank_c_reg = vreinterpretq_s8_u64(
-              vshrq_n_u64(vreinterpretq_u64_s8(input_bank_c_reg), 16));
-
-          output_data += output_depth;
-        }
-        for (; i_width < output_width_overall_micro_repeats; ++i_width) {
-          // output_width == 1.
-          const int8_t* input_data = scratch_data + 4 + 4 * i_width;
-
-          // Load next sub-micro block of data.
-          input_bank_a_reg = vld1q_lane_8x4(input_data, input_bank_a_reg, 1);
-          input_bank_a_reg = vld1q_lane_8x4(
-              input_data + workspace_height_stride, input_bank_a_reg, 3);
-          input_bank_b_reg = vld1q_lane_8x4(
-              input_data + 2 * workspace_height_stride, input_bank_b_reg, 1);
-          input_bank_b_reg = vld1q_lane_8x4(
-              input_data + 3 * workspace_height_stride, input_bank_b_reg, 3);
-          input_bank_c_reg = vld1q_lane_8x4(
-              input_data + 4 * workspace_height_stride, input_bank_c_reg, 1);
-
-          int16x8_t acc_s16_0_1;
-          int8x8_t acc_u8_0_1;
-          // Iterate over input width shifts within 4x4 blocks.
-          {
-            acc0 = adjusted_bias_data_s_0;
-            acc1 = adjusted_bias_data_s_0;
-
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_0_a, input_bank_a_reg, 0);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_1_a, input_bank_a_reg, 2);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_2_a, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_0_a, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_1_a, input_bank_b_reg, 2);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_2_a, input_bank_c_reg, 0);
-
-            // Fixed-point multiplication.
-            acc0 = vqrdmulhq_s32(acc0, output_multiplier_s_0);
-            acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc0, output_shift_s_0);
-            acc1 = vqrdmulhq_s32(acc1, output_multiplier_s_0);
-            acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc1, output_shift_s_0);
-            // Add the output offset.
-            acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            acc_u8_0_1 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8_0_1 = util_vmax_x8(acc_u8_0_1,
-                                      vget_low_s8(output_activation_min_vec));
-            acc_u8_0_1 = util_vmin_x8(acc_u8_0_1,
-                                      vget_low_s8(output_activation_max_vec));
-
-            vst1_lane_s8x4(output_data, acc_u8_0_1, 0);
-            vst1_lane_s8x4(output_data + output_height_stride, acc_u8_0_1, 1);
-
-            acc0 = adjusted_bias_data_s_1;
-            acc1 = adjusted_bias_data_s_1;
-
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_0_b, input_bank_a_reg, 0);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_1_b, input_bank_a_reg, 2);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_2_b, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_0_b, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_1_b, input_bank_b_reg, 2);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_2_b, input_bank_c_reg, 0);
-
-            // Fixed-point multiplication.
-            acc0 = vqrdmulhq_s32(acc0, output_multiplier_s_1);
-            acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc0, output_shift_s_1);
-            acc1 = vqrdmulhq_s32(acc1, output_multiplier_s_1);
-            acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc1, output_shift_s_1);
-            // Add the output offset.
-            acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            acc_u8_0_1 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8_0_1 = util_vmax_x8(acc_u8_0_1,
-                                      vget_low_s8(output_activation_min_vec));
-            acc_u8_0_1 = util_vmin_x8(acc_u8_0_1,
-                                      vget_low_s8(output_activation_max_vec));
-
-            vst1_lane_s8x4(output_data + 4, acc_u8_0_1, 0);
-            vst1_lane_s8x4(output_data + 4 + output_height_stride, acc_u8_0_1,
-                           1);
-
-            input_bank_a_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_a_reg), 16));
-            input_bank_b_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_b_reg), 16));
-            input_bank_c_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_c_reg), 16));
-
-            output_data += output_depth;
-          }
-        }
-      } else {
-        TFLITE_DCHECK_EQ(block_height, 1);
-        // Work through one slice, by row, at a time.
-        const int8_t* scratch_data = scratch_block_data;
-        typename QuantizationTypeImpl<quantization_type>::ExternalType*
-            output_data = output_block_data + 8 * j_depth;
-
-        int8x16_t input_bank_a_reg;  //  left 0, right 0, left 1, right 1.
-        int8x16_t input_bank_b_reg;  //  left 2, right 2, xxx, xxx.
-
-        // Load first sub-micro block of data into operational banks.
-        input_bank_a_reg =
-            vld1q_dup_s8x4(scratch_data);  // Load lane 0, avoiding
-                                           // uninitialized variable.
-        input_bank_a_reg = vld1q_lane_8x4(
-            scratch_data + workspace_height_stride, input_bank_a_reg, 2);
-        input_bank_b_reg = vld1q_dup_s8x4(
-            scratch_data +
-            2 * workspace_height_stride);  // Load lane 0, avoiding
-                                           // uninitialized variable.
-
-        int32x4_t acc0;
-        int32x4_t acc1;
-
-        for (int i_width = 0; i_width < output_width_overall_micro_repeats;
-             ++i_width) {
-          const int output_width =
-              i_width == output_width_micro_repeats ? residual_width : 2;
-
-          TFLITE_DCHECK_LE(output_width, 2);
-          TFLITE_DCHECK_GE(output_width, 1);
-          TFLITE_DCHECK_LE(output_width * kStrideVal, 4);
-          const int8_t* input_data = scratch_data + 4 + 4 * i_width;
-
-          // Load next sub-micro block of data.
-          input_bank_a_reg = vld1q_lane_8x4(input_data, input_bank_a_reg, 1);
-          input_bank_a_reg = vld1q_lane_8x4(
-              input_data + workspace_height_stride, input_bank_a_reg, 3);
-          input_bank_b_reg = vld1q_lane_8x4(
-              input_data + 2 * workspace_height_stride, input_bank_b_reg, 1);
-
-          int16x8_t acc_s16_0_1;
-          int8x8_t acc_u8_0_1;
-
-          // Iterate over input width shifts within 4x4 blocks.
-          {
-            acc0 = adjusted_bias_data_s_0;
-
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_2_a, input_bank_b_reg, 0);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_0_a, input_bank_a_reg, 0);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_1_a, input_bank_a_reg, 2);
-
-            acc0 = vqrdmulhq_s32(acc0, output_multiplier_s_0);
-            acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc0, output_shift_s_0);
-
-            // Second sub-block accumulation.
-            acc1 = adjusted_bias_data_s_1;
-
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_2_b, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_0_b, input_bank_a_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_1_b, input_bank_a_reg, 2);
-
-            acc1 = vqrdmulhq_s32(acc1, output_multiplier_s_1);
-            acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc1, output_shift_s_1);
-
-            // Add the output offset.
-            acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            acc_u8_0_1 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8_0_1 = util_vmax_x8(acc_u8_0_1,
-                                      vget_low_s8(output_activation_min_vec));
-            acc_u8_0_1 = util_vmin_x8(acc_u8_0_1,
-                                      vget_low_s8(output_activation_max_vec));
-
-            // This stores the results for both sub-blocks together.
-            vst1_s8(output_data, acc_u8_0_1);
-
-            input_bank_a_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_a_reg), 16));
-            input_bank_b_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_b_reg), 16));
-
-            output_data += output_depth;
-          }
-          if (output_width == 2) {
-            acc0 = adjusted_bias_data_s_0;
-
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_2_a, input_bank_b_reg, 0);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_0_a, input_bank_a_reg, 0);
-            acc0 =
-                vdotq_four_lane_s32(acc0, filter_reg_1_a, input_bank_a_reg, 2);
-
-            acc0 = vqrdmulhq_s32(acc0, output_multiplier_s_0);
-            acc0 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc0, output_shift_s_0);
-
-            // Second sub-block accumulation.
-            acc1 = adjusted_bias_data_s_1;
-
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_2_b, input_bank_b_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_0_b, input_bank_a_reg, 0);
-            acc1 =
-                vdotq_four_lane_s32(acc1, filter_reg_1_b, input_bank_a_reg, 2);
-
-            acc1 = vqrdmulhq_s32(acc1, output_multiplier_s_1);
-            acc1 = DivideByPOT<DepthwiseConvOutputRounding::kUpward>::RunMult(
-                acc1, output_shift_s_1);
-
-            // Add the output offset.
-            acc_s16_0_1 = vcombine_s16(vqmovn_s32(acc0), vqmovn_s32(acc1));
-            acc_s16_0_1 = vqaddq_s16(acc_s16_0_1, output_offset_vec);
-            // Apply the activation function.
-            acc_u8_0_1 = vqmovxn_s16(acc_s16_0_1);
-            acc_u8_0_1 = util_vmax_x8(acc_u8_0_1,
-                                      vget_low_s8(output_activation_min_vec));
-            acc_u8_0_1 = util_vmin_x8(acc_u8_0_1,
-                                      vget_low_s8(output_activation_max_vec));
-
-            // This stores the results for both sub-blocks together.
-            vst1_s8(output_data, acc_u8_0_1);
-
-            input_bank_a_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_a_reg), 16));
-            input_bank_b_reg = vreinterpretq_s8_u64(
-                vshrq_n_u64(vreinterpretq_u64_s8(input_bank_b_reg), 16));
-
-            output_data += output_depth;
-          }
-        }
-      }
-    }
-  }
-
-  static inline void Run(const int8_t* scratch_block_data,
-                         const int8_t* filter_workspace,
-                         const int32_t* bias_data, int8_t* output_block_data,
-                         const DepthwiseConvDotProdParams* function_params) {
-    KernelMacroBlockIntrinsics(scratch_block_data, filter_workspace, bias_data,
-                               output_block_data, function_params);
-  }
-};
-
-#undef vst1_lane_s8x4
-#undef vst1_lane_u8x4
-#undef vst1q_lane_s8x4
-#undef vst1q_lane_u8x4
-#undef vld1q_lane_s8x8
-#undef vld1_lane_8x4
-#undef vld1q_lane_8x4
-#undef vld1q_dup_s8x4
-
-#endif  //  USE_NEON
-
-}  // namespace depthwise_conv
-}  // namespace optimized_ops
-}  // namespace tflite
-
-#endif  // TENSORFLOW_LITE_KERNELS_INTERNAL_OPTIMIZED_DEPTHWISECONV_UINT8_TRANSITIONAL_H_

From 47bb5724f97e1385a5ec9711474faf5446b1df34 Mon Sep 17 00:00:00 2001
From: Bill Varcho <varcho@google.com>
Date: Mon, 12 Oct 2026 15:11:19 -0700
Subject: [PATCH 0626/1000] [SDY][Refactor][3/3] Support single_device HLO V3
 format and terminology.

Update HLO V3 mesh parsing, printing, and Shardy round-trip import/export pipeline to use single_device instead of maximal_mesh.

PiperOrigin-RevId: 998182497
---
 third_party/xla/xla/hlo/ir/mesh_and_axis.cc   | 10 +++----
 third_party/xla/xla/hlo/ir/mesh_and_axis.h    |  6 ++--
 .../xla/xla/hlo/ir/mesh_and_axis_test.cc      | 25 +++++++++--------
 .../xla/xla/hlo/ir/named_sharding_test.cc     |  6 ++--
 third_party/xla/xla/hlo/parser/hlo_parser.cc  | 23 +++++++++------
 .../xla/xla/hlo/parser/hlo_parser_test.cc     | 22 +++++++++++----
 .../stablehlo_round_trip/stablehlo_import.cc  | 11 ++++----
 .../spmd/shardy/test/import_shardings.mlir    | 28 +++++++++----------
 .../test/sdy_round_trip_export_pipeline.mlir  |  2 +-
 ..._trip_import_pipeline_hlo_sharding_v3.mlir |  2 +-
 .../test/stablehlo_export_pipeline.mlir       | 14 +++++-----
 ...o_export_pipeline_simplify_replicated.mlir |  2 +-
 .../xla/xla/service/spmd/shardy/utils.h       |  3 +-
 third_party/xla/xla/xla_data.proto            |  8 +++---
 14 files changed, 91 insertions(+), 71 deletions(-)

diff --git a/third_party/xla/xla/hlo/ir/mesh_and_axis.cc b/third_party/xla/xla/hlo/ir/mesh_and_axis.cc
index b7ad1da346b..89b67853593 100644
--- a/third_party/xla/xla/hlo/ir/mesh_and_axis.cc
+++ b/third_party/xla/xla/hlo/ir/mesh_and_axis.cc
@@ -46,12 +46,12 @@ namespace xla {
 
 absl::Status Mesh::Validate() {
   if (device_assignment_.num_dimensions() == 0) {
-    // Empty mesh or maximal mesh.
+    // Empty mesh or single-device mesh.
     if (device_assignment_.num_elements() <= 1) {
       return absl::OkStatus();
     }
     return absl::InvalidArgumentError(absl::StrCat(
-        "Non-maximal mesh must have exactly 1 device id. Number of "
+        "Single-device mesh must have at most 1 device id. Number of "
         "device ids: ",
         device_assignment_.num_elements()));
   }
@@ -109,9 +109,9 @@ Mesh::Mesh(TileAssignment device_assignment,
 }
 
 std::string Mesh::ToString() const {
-  if (IsMaximal()) {
+  if (IsSingleDevice()) {
     return absl::StrCat(
-        "maximal_mesh[device_id=", device_assignment_.array()(0), "]");
+        "single_device[device_id=", device_assignment_.array()(0), "]");
   }
 
   std::string mesh_str = "mesh";
@@ -143,7 +143,7 @@ MeshProto Mesh::ToProto() const {
     if (device_assignment_.num_elements() == 0) {
       return MeshProto();
     }
-    // Maximal mesh
+    // Single-device mesh
     // TODO(b/454008727): Validate device_ids_size is 1.
     proto.add_device_ids(*device_assignment_.array().begin());
     return proto;
diff --git a/third_party/xla/xla/hlo/ir/mesh_and_axis.h b/third_party/xla/xla/hlo/ir/mesh_and_axis.h
index 21da41633ca..8d18c6e85da 100644
--- a/third_party/xla/xla/hlo/ir/mesh_and_axis.h
+++ b/third_party/xla/xla/hlo/ir/mesh_and_axis.h
@@ -44,8 +44,8 @@ class AxisRef;
 // - device_assignment_.dimensions() represents the axis sizes.
 // - device_assignment_.array() represents the list of device IDs.
 //
-// For maximal mesh, axes_names is empty and device_assignment_ contains the
-// single device id.
+// For single-device mesh, axes_names is empty and device_assignment_ contains
+// the single device id.
 //
 // Example: device_assignment {{3, 0, 2}, {1, 4, 5}} with axes names
 // {"data", "model"} represents the mesh ["data"=2, "model"=3].
@@ -54,7 +54,7 @@ class Mesh {
   // Empty mesh
   explicit Mesh() = default;
 
-  // Maximal Mesh
+  // Single-device Mesh
   explicit Mesh(int64_t device_id) : device_assignment_(device_id) {}
 
   // Constructs an iota device assignment mesh with given axes sizes and names.
diff --git a/third_party/xla/xla/hlo/ir/mesh_and_axis_test.cc b/third_party/xla/xla/hlo/ir/mesh_and_axis_test.cc
index e086ae10726..0ab1571d1b4 100644
--- a/third_party/xla/xla/hlo/ir/mesh_and_axis_test.cc
+++ b/third_party/xla/xla/hlo/ir/mesh_and_axis_test.cc
@@ -272,8 +272,8 @@ TEST(MeshAndAxisTest, MeshToString) {
   EXPECT_EQ(mesh_ooo.ToString(),
             "mesh['ooo'=10], device_ids=(8,3,7,5,4,2,6,0,1,9)");
 
-  Mesh maximal_mesh(5);
-  EXPECT_EQ(maximal_mesh.ToString(), "maximal_mesh[device_id=5]");
+  Mesh single_device_mesh(5);
+  EXPECT_EQ(single_device_mesh.ToString(), "single_device[device_id=5]");
 
   Mesh mesh_with_special_characters({10, 12, 15},
                                     {"abc def", "xyz.pqr", "<axis> def"});
@@ -364,26 +364,29 @@ TEST(MeshAndAxisTest, EmptyMesh) {
   EXPECT_EQ(empty_mesh, Mesh::FromProto(empty_mesh.ToProto()));
 }
 
-TEST(MeshAndAxisTest, MaximalMesh) {
-  Mesh maximal_mesh(5);
-  EXPECT_TRUE(maximal_mesh.IsMaximal());
-  Mesh non_maximal_mesh({2, 3}, {"a", "b"});
-  EXPECT_FALSE(non_maximal_mesh.IsMaximal());
+TEST(MeshAndAxisTest, SingleDeviceMesh) {
+  Mesh single_device_mesh(5);
+  EXPECT_TRUE(single_device_mesh.IsSingleDevice());
+  EXPECT_TRUE(single_device_mesh.IsMaximal());
+  Mesh non_single_device_mesh({2, 3}, {"a", "b"});
+  EXPECT_FALSE(non_single_device_mesh.IsSingleDevice());
+  EXPECT_FALSE(non_single_device_mesh.IsMaximal());
   Mesh mesh_single_axis({1}, {"a"});
+  EXPECT_FALSE(mesh_single_axis.IsSingleDevice());
   EXPECT_FALSE(mesh_single_axis.IsMaximal());
 
-  EXPECT_EQ(maximal_mesh, Mesh(5));
-  EXPECT_NE(maximal_mesh, Mesh(6));
+  EXPECT_EQ(single_device_mesh, Mesh(5));
+  EXPECT_NE(single_device_mesh, Mesh(6));
 
   MeshProto expected_proto;
   expected_proto.add_device_ids(5);
-  EXPECT_THAT(maximal_mesh.ToProto(), EqualsProto(expected_proto));
+  EXPECT_THAT(single_device_mesh.ToProto(), EqualsProto(expected_proto));
 
   MeshProto from_proto;
   from_proto.add_device_ids(7);
   EXPECT_EQ(Mesh(7), Mesh::FromProto(from_proto));
 
-  EXPECT_EQ(maximal_mesh, Mesh::FromProto(maximal_mesh.ToProto()));
+  EXPECT_EQ(single_device_mesh, Mesh::FromProto(single_device_mesh.ToProto()));
 }
 
 TEST(MeshAndAxisTest, AxisRefSize) {
diff --git a/third_party/xla/xla/hlo/ir/named_sharding_test.cc b/third_party/xla/xla/hlo/ir/named_sharding_test.cc
index 42bdc1ae935..464935986a6 100644
--- a/third_party/xla/xla/hlo/ir/named_sharding_test.cc
+++ b/third_party/xla/xla/hlo/ir/named_sharding_test.cc
@@ -199,9 +199,9 @@ TEST(NamedShardingTest, ToString) {
   EXPECT_EQ(sharding_dim.ToString(),
             "{mesh['a'=2,'b'=4,'c'=3,'d'=8], [{'c'}, {'a', 'b':(2)2, ?}]}");
 
-  Mesh maximal_mesh(5);
-  NamedSharding single_device_sharding(maximal_mesh);
-  EXPECT_EQ(single_device_sharding.ToString(), "{maximal_mesh[device_id=5]}");
+  Mesh single_device_mesh(5);
+  NamedSharding single_device_sharding(single_device_mesh);
+  EXPECT_EQ(single_device_sharding.ToString(), "{single_device[device_id=5]}");
 
   Mesh non_iota_mesh(
       TileAssignment(/*dims=*/{2, 4, 4, 2}, /*reshape_dims=*/{1, 4, 1, 16},
diff --git a/third_party/xla/xla/hlo/parser/hlo_parser.cc b/third_party/xla/xla/hlo/parser/hlo_parser.cc
index 97d92958954..f157d074e29 100644
--- a/third_party/xla/xla/hlo/parser/hlo_parser.cc
+++ b/third_party/xla/xla/hlo/parser/hlo_parser.cc
@@ -4156,10 +4156,11 @@ bool HloParserImpl::ParseCollectiveDeviceListBase(
     return true;
   }
 
-  // If the first token is 'mesh' or 'maximal_mesh', then we are parsing a (v3)
-  // MeshAxesReplicaGroupList.
+  // If the first token is 'mesh', 'single_device', or 'maximal_mesh', then we
+  // are parsing a (v3) MeshAxesReplicaGroupList.
   if (lexer_.GetKind() == TokKind::kIdent &&
-      (lexer_.GetStrVal() == "mesh" || lexer_.GetStrVal() == "maximal_mesh")) {
+      (lexer_.GetStrVal() == "mesh" || lexer_.GetStrVal() == "single_device" ||
+       lexer_.GetStrVal() == "maximal_mesh")) {
     std::optional<Mesh> mesh;
     if (!ParseMesh(mesh)) {
       return false;
@@ -4335,19 +4336,22 @@ bool HloParserImpl::ParseMesh(std::optional<Mesh>& mesh) {
   std::string kind = lexer_.GetStrVal();
   lexer_.Lex();
 
-  if (kind == "maximal_mesh") {
-    if (!ParseToken(TokKind::kLsquare, "expected '[' after maximal_mesh")) {
+  if (kind == "single_device" || kind == "maximal_mesh") {
+    if (!ParseToken(TokKind::kLsquare,
+                    absl::StrCat("expected '[' after ", kind))) {
       return false;
     }
     std::string device_ids_str;
     if (!ParseAttributeName(&device_ids_str) || device_ids_str != "device_id") {
-      return TokenError("expected device_id= attribute in maximal_mesh");
+      return TokenError(
+          absl::StrCat("expected device_id= attribute in ", kind));
     }
     int64_t device_id;
     if (!ParseInt64(&device_id)) {
       return false;
     }
-    if (!ParseToken(TokKind::kRsquare, "expected ']' to end maximal_mesh")) {
+    if (!ParseToken(TokKind::kRsquare,
+                    absl::StrCat("expected ']' to end ", kind))) {
       return false;
     }
     mesh.emplace(device_id);
@@ -4355,7 +4359,7 @@ bool HloParserImpl::ParseMesh(std::optional<Mesh>& mesh) {
   }
 
   if (kind != "mesh") {
-    return TokenError("Expected 'mesh' or 'maximal_mesh'");
+    return TokenError("Expected 'mesh', 'single_device', or 'maximal_mesh'");
   }
 
   // Parse @mesh[...]
@@ -4892,7 +4896,8 @@ bool HloParserImpl::ParseSingleSharding(std::optional<HloSharding>& sharding,
   }
 
   if (lexer_.GetKind() == TokKind::kIdent &&
-      (lexer_.GetStrVal() == "mesh" || lexer_.GetStrVal() == "maximal_mesh")) {
+      (lexer_.GetStrVal() == "mesh" || lexer_.GetStrVal() == "single_device" ||
+       lexer_.GetStrVal() == "maximal_mesh")) {
     std::optional<NamedSharding> named_sharding;
     if (ParseSingleSharding(named_sharding)) {
       sharding = HloSharding(std::move(*named_sharding));
diff --git a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc
index 6578a8cc02e..8f6c7233512 100644
--- a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc
+++ b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc
@@ -4934,10 +4934,16 @@ TEST_F(HloParserTest, ParseNamedShardingReplicatedAxes) {
   EXPECT_EQ(sharding.ToString(/*include_metadata=*/true), original);
 }
 
-TEST_F(HloParserTest, ParseNamedShardingMaximal) {
-  const std::string original = "{maximal_mesh[device_id=5]}";
+TEST_F(HloParserTest, ParseNamedShardingSingleDevice) {
+  const std::string original = "{single_device[device_id=5]}";
   ASSERT_OK_AND_ASSIGN(HloSharding sharding, ParseSharding(original));
   EXPECT_EQ(sharding.ToString(/*include_metadata=*/true), original);
+
+  // Backward compatibility: parsing maximal_mesh
+  ASSERT_OK_AND_ASSIGN(HloSharding sharding_maximal,
+                       ParseSharding("{maximal_mesh[device_id=5]}"));
+  EXPECT_EQ(sharding_maximal.ToString(/*include_metadata=*/true),
+            "{single_device[device_id=5]}");
 }
 
 TEST_F(HloParserTest, ParseNamedShardingWithSpecialCharacters) {
@@ -5006,7 +5012,7 @@ TEST_F(HloParserTest, ParseNamedShardingTuple) {
 TEST_F(HloParserTest, ParseMixedShardingTuple1) {
   const std::string original =
       "{{replicated}, {mesh['a'=2,'b'=4], replicated}, {maximal device=5}, "
-      "{maximal_mesh[device_id=5]}}";
+      "{single_device[device_id=5]}}";
   ASSERT_OK_AND_ASSIGN(HloSharding sharding, ParseSharding(original));
 
   EXPECT_EQ(sharding.ToString(/*include_metadata=*/true), original);
@@ -5155,12 +5161,16 @@ TEST_F(HloParserTest, ParseReplicaGroupsV3) {
   EXPECT_EQ("{{0,1},{2,3}}", ReplicaGroupsToString(replica_groups));
 }
 
-TEST_F(HloParserTest, ParseReplicaGroupsMaximalMesh) {
-  const std::string original = "maximal_mesh[device_id=1] {}";
-  auto status = ParseReplicaGroupsOnly(original).status();
+TEST_F(HloParserTest, ParseReplicaGroupsSingleDevice) {
+  const std::string original = "single_device[device_id=1] {}";
   ASSERT_OK_AND_ASSIGN(std::vector<ReplicaGroup> replica_groups,
                        ParseReplicaGroupsOnly(original));
   EXPECT_EQ("{{1}}", ReplicaGroupsToString(replica_groups));
+
+  // Backward compatibility: parsing maximal_mesh
+  ASSERT_OK_AND_ASSIGN(std::vector<ReplicaGroup> replica_groups_maximal,
+                       ParseReplicaGroupsOnly("maximal_mesh[device_id=1] {}"));
+  EXPECT_EQ("{{1}}", ReplicaGroupsToString(replica_groups_maximal));
 }
 
 TEST_F(HloParserTest, ParseDynamicReshapeMissingOperands) {
diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.cc b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.cc
index 629abcd5149..f321b8d56a1 100644
--- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.cc
+++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.cc
@@ -687,15 +687,16 @@ class ImportShardingsPass
     symbolTable.insert(MeshOp::create(opBuilder, moduleOp.getLoc(),
                                       kGlobalMeshName, globalMesh));
 
-    SmallDenseMap<int64_t, StringRef> deviceIdToMaximalMeshName;
+    SmallDenseMap<int64_t, StringRef> deviceIdToSingleDeviceMeshName;
     for (int64_t deviceId : deviceIdsForMaximalMesh) {
-      // Create a mesh name with its deviceId as a suffix for each maximal mesh.
-      std::string meshName = absl::StrCat("maximal_mesh_", deviceId);
+      // Create a mesh name with its deviceId as a suffix for each single-device
+      // mesh.
+      std::string meshName = absl::StrCat("single_device_", deviceId);
       auto meshOp = MeshOp::create(
           opBuilder, moduleOp.getLoc(), meshName,
           MeshAttr::getSingleDevice(moduleOp.getContext(), deviceId));
       symbolTable.insert(meshOp);
-      deviceIdToMaximalMeshName[deviceId] = meshOp.getSymName();
+      deviceIdToSingleDeviceMeshName[deviceId] = meshOp.getSymName();
     }
 
     for (FuncOp funcOp : moduleOp.getOps<FuncOp>()) {
@@ -703,7 +704,7 @@ class ImportShardingsPass
       // only to the main/entry function.
       bool isMain = funcOp.getSymName() == "main";
       if (mlir::failed(importShardings(
-              funcOp, globalMesh, deviceIdToMaximalMeshName,
+              funcOp, globalMesh, deviceIdToSingleDeviceMeshName,
               isMain ? allowPropagationToArgs : ArrayRef<bool>(),
               isMain ? allowPropagationToResults : ArrayRef<bool>(),
               inlineMesh))) {
diff --git a/third_party/xla/xla/service/spmd/shardy/test/import_shardings.mlir b/third_party/xla/xla/service/spmd/shardy/test/import_shardings.mlir
index 61e3f7b182c..c579b8e38f0 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/import_shardings.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/import_shardings.mlir
@@ -141,10 +141,10 @@ func.func @unknown_sharding(%arg0: tensor<8x8xf32> {mhlo.sharding = "{devices=[4
 // -----
 
 // CHECK-LABEL: sdy.mesh @mesh = <[]>
-// CHECK-LABEL: sdy.mesh @maximal_mesh_0 = <[], device_ids=[0]>
+// CHECK-LABEL: sdy.mesh @single_device_0 = <[], device_ids=[0]>
 
 // CHECK-LABEL: func @one_maximal_mesh(
-// CHECK-SAME:      %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@maximal_mesh_0, []>}
+// CHECK-SAME:      %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@single_device_0, []>}
 func.func @one_maximal_mesh(%arg0: tensor<8x8xf32> {mhlo.sharding = "{maximal device=0}"},
                             %arg1: tensor<8x8xf32>) -> tensor<8x8xf32> {
   %0 = stablehlo.add %arg0, %arg1 : tensor<8x8xf32>
@@ -153,12 +153,12 @@ func.func @one_maximal_mesh(%arg0: tensor<8x8xf32> {mhlo.sharding = "{maximal de
 
 // -----
 
-// CHECK-LABEL: sdy.mesh @maximal_mesh_0 = <[], device_ids=[0]>
-// CHECK-LABEL: sdy.mesh @maximal_mesh_4 = <[], device_ids=[4]>
+// CHECK-LABEL: sdy.mesh @single_device_0 = <[], device_ids=[0]>
+// CHECK-LABEL: sdy.mesh @single_device_4 = <[], device_ids=[4]>
 
 // CHECK-LABEL: func @two_maximal_shardings_should_be_sorted(
-// CHECK-SAME:      %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@maximal_mesh_4, []>},
-// CHECK-SAME:      %arg1: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@maximal_mesh_0, []>})
+// CHECK-SAME:      %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@single_device_4, []>},
+// CHECK-SAME:      %arg1: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@single_device_0, []>})
 func.func @two_maximal_shardings_should_be_sorted(%arg0: tensor<8x8xf32> {mhlo.sharding = "{maximal device=4}"},
                             %arg1: tensor<8x8xf32> {mhlo.sharding = "{maximal device=0}"}) -> tensor<8x8xf32> {
   %0 = stablehlo.add %arg0, %arg1 : tensor<8x8xf32>
@@ -166,11 +166,11 @@ func.func @two_maximal_shardings_should_be_sorted(%arg0: tensor<8x8xf32> {mhlo.s
 }
 
 // -----
-// CHECK-COUNT-1: sdy.mesh @maximal_mesh_0 = <[], device_ids=[0]>
+// CHECK-COUNT-1: sdy.mesh @single_device_0 = <[], device_ids=[0]>
 
 // CHECK-LABEL: func @duplicate_maximal_sharding_should_be_deduped(
-// CHECK-SAME:      %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@maximal_mesh_0, []>},
-// CHECK-SAME:      %arg1: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@maximal_mesh_0, []>})
+// CHECK-SAME:      %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@single_device_0, []>},
+// CHECK-SAME:      %arg1: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@single_device_0, []>})
 func.func @duplicate_maximal_sharding_should_be_deduped(%arg0: tensor<8x8xf32> {mhlo.sharding = "{maximal device=0}"},
                             %arg1: tensor<8x8xf32> {mhlo.sharding = "{maximal device=0}"}) -> tensor<8x8xf32> {
   %0 = stablehlo.add %arg0, %arg1 : tensor<8x8xf32>
@@ -180,11 +180,11 @@ func.func @duplicate_maximal_sharding_should_be_deduped(%arg0: tensor<8x8xf32> {
 // -----
 
 // CHECK-LABEL: sdy.mesh @mesh = <["_axis_0"=8, "_axis_1"=4]>
-// CHECK-LABEL: sdy.mesh @maximal_mesh_0 = <[], device_ids=[0]>
+// CHECK-LABEL: sdy.mesh @single_device_0 = <[], device_ids=[0]>
 
 // CHECK-LABEL: func @two_meshes(
 // CHECK-SAME:      %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"_axis_1"}, {}]>},
-// CHECK-SAME:      %arg1: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@maximal_mesh_0, []>}, %arg2: tensor<8x16xf32>)
+// CHECK-SAME:      %arg1: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@single_device_0, []>}, %arg2: tensor<8x16xf32>)
 func.func @two_meshes(%arg0: tensor<8x8xf32> {mhlo.sharding = "{devices=[4,1,8]<=[8,4]T(1,0) last_tile_dim_replicate}"},
                             %arg1: tensor<8x8xf32> {mhlo.sharding = "{maximal device=0}"},
                             %arg2: tensor<8x16xf32>) -> tensor<8x16xf32> {
@@ -195,7 +195,7 @@ func.func @two_meshes(%arg0: tensor<8x8xf32> {mhlo.sharding = "{devices=[4,1,8]<
 
 // -----
 // CHECK-LABEL: sdy.mesh @mesh = <["_axis_0"=8, "_axis_1"=4]>
-// CHECK-LABEL: sdy.mesh @maximal_mesh_0 = <[], device_ids=[0]>
+// CHECK-LABEL: sdy.mesh @single_device_0 = <[], device_ids=[0]>
 
 // CHECK-LABEL: func @maximal_sharding_on_op(
 // CHECK-SAME:      %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"_axis_1"}, {}]>},
@@ -204,9 +204,9 @@ func.func @two_meshes(%arg0: tensor<8x8xf32> {mhlo.sharding = "{devices=[4,1,8]<
 func.func @maximal_sharding_on_op(%arg0: tensor<8x8xf32> {mhlo.sharding = "{devices=[4,1,8]<=[8,4]T(1,0) last_tile_dim_replicate}"},
                             %arg1: tensor<8x8xf32>) -> tensor<8x8xf32> {
 // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %arg0, %arg1
-// CHECK-SAME{LITERAL}: {sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh_4, []>]>}
+// CHECK-SAME{LITERAL}: {sdy.sharding = #sdy.sharding_per_value<[<@single_device_4, []>]>}
 // CHECK-NEXT: %[[MULTIPLY:.*]] = stablehlo.multiply %[[ADD]], %[[ADD]]
-// CHECK-SAME{LITERAL}: {sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh_0, []>]>}
+// CHECK-SAME{LITERAL}: {sdy.sharding = #sdy.sharding_per_value<[<@single_device_0, []>]>}
   %0 = stablehlo.add %arg0, %arg1 {mhlo.sharding = "{maximal device=4}"} : tensor<8x8xf32>
   %1 = stablehlo.multiply %0, %0 {mhlo.sharding = "{maximal device=0}"} : tensor<8x8xf32>
   return %1 : tensor<8x8xf32>
diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_export_pipeline.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_export_pipeline.mlir
index 7dd40a61cfe..a8715aba66c 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_export_pipeline.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_export_pipeline.mlir
@@ -177,7 +177,7 @@ func.func @constant() -> tensor<i32> {
 // CHECK-V2-SAME: -> (tensor<32xi32> {mhlo.sharding = "{maximal device=5}"}) {
 // CHECK-V3-NOT:    mhlo.frontend_attributes
 // CHECK-V3-SAME: mhlo.sharding = "{mesh['a'=2,'b'=2], [{'a'}]}"
-// CHECK-V3-SAME: -> (tensor<32xi32> {mhlo.sharding = "{maximal_mesh[device_id=5]}"}) {
+// CHECK-V3-SAME: -> (tensor<32xi32> {mhlo.sharding = "{single_device[device_id=5]}"}) {
 func.func @inlined_mesh(
   %arg0: tensor<32xi32> {sdy.sharding = #sdy.sharding<mesh<["a"=2, "b"=2]>, [{"a"}]>}
 ) -> (tensor<32xi32> {sdy.sharding = #sdy.sharding<mesh<[], device_ids=[5]>, []>}) {
diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir
index bb6f2d9f8a5..eeadf165681 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir
@@ -160,7 +160,7 @@ module @module_1 {
   // CHECK-SAME: -> tensor<32xi32> {
   func.func @inlined_mesh(
     %arg0: tensor<32xi32> {mhlo.sharding = "{mesh['a'=2,'b'=2], [{'a'}]}"},
-    %arg1: tensor<32xi32> {mhlo.sharding = "{maximal_mesh[device_id=5]}"}
+    %arg1: tensor<32xi32> {mhlo.sharding = "{single_device[device_id=5]}"}
   ) -> tensor<32xi32> {
     // CHECK-NEXT: %[[SHARDING:.*]] = sdy.sharding_constraint %arg0 <@mesh_0, [{"a", "b"}]> : tensor<32xi32>
     // CHECK-NEXT: return %[[SHARDING]]
diff --git a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_pipeline.mlir b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_pipeline.mlir
index 6954c558e09..63d2f308741 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_pipeline.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_pipeline.mlir
@@ -261,12 +261,12 @@ func.func @sharding_in_manual_computation_body(%arg0: tensor<8x16xf32> {sdy.shar
 }
 
 // CHECK-V2-LABEL: func @mesh_with_device_id_should_be_converted_to_maximal_sharding(%arg0: tensor<8x8xf32> {mhlo.sharding = "{maximal device=0}"}, %arg1: tensor<8x8xf32>)
-// CHECK-V3-LABEL: func @mesh_with_device_id_should_be_converted_to_maximal_sharding(%arg0: tensor<8x8xf32> {mhlo.sharding = "{maximal_mesh[device_id=0]}"}, %arg1: tensor<8x8xf32>)
+// CHECK-V3-LABEL: func @mesh_with_device_id_should_be_converted_to_maximal_sharding(%arg0: tensor<8x8xf32> {mhlo.sharding = "{single_device[device_id=0]}"}, %arg1: tensor<8x8xf32>)
 func.func @mesh_with_device_id_should_be_converted_to_maximal_sharding(%arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@maximal_mesh_0, []>}, %arg1: tensor<8x8xf32>) -> tensor<8x8xf32> {
     // CHECK: %[[ADD:.*]] = stablehlo.add %arg0, %arg1
     %0 = stablehlo.add %arg0, %arg1 : tensor<8x8xf32>
     // CHECK-V2: %[[ADD_WITH_SHARDING:.*]] = stablehlo.add %[[ADD]], %[[ADD]] {mhlo.sharding = "{maximal device=1}"}
-    // CHECK-V3: %[[ADD_WITH_SHARDING:.*]] = stablehlo.add %[[ADD]], %[[ADD]] {mhlo.sharding = "{maximal_mesh[device_id=1]}"}
+    // CHECK-V3: %[[ADD_WITH_SHARDING:.*]] = stablehlo.add %[[ADD]], %[[ADD]] {mhlo.sharding = "{single_device[device_id=1]}"}
     %1 = stablehlo.add %0, %0 {sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh_1, []>]>} : tensor<8x8xf32>
     return %1 : tensor<8x8xf32>
 }
@@ -422,7 +422,7 @@ func.func public @callback_no_result(%arg0: tensor<f64>) {
   // CHECK-NEXT: stablehlo.custom_call @xla_python_cpu_callback(%[[C]], %arg0) {
   // CHECK-SAME:   api_version = 2 : i32, backend_config = "56238273106176",
   // CHECK-V2-SAME:   has_side_effect = true,  mhlo.sharding = "{maximal device=0}",
-  // CHECK-V3-SAME:   has_side_effect = true,  mhlo.sharding = "{maximal_mesh[device_id=0]}",
+  // CHECK-V3-SAME:   has_side_effect = true,  mhlo.sharding = "{single_device[device_id=0]}",
   // CHECK-SAME:   operand_layouts = [dense<> : tensor<0xindex>, dense<> : tensor<0xindex>],
   // CHECK-SAME: } : (tensor<i64>, tensor<f64>) -> ()
   %c = stablehlo.constant dense<56238273106176> : tensor<i64>
@@ -437,7 +437,7 @@ func.func public @callback_tuple_result_token_used(%arg0: !stablehlo.token, %arg
   // CHECK-NEXT: %[[CALLBACK:.*]] = stablehlo.custom_call @xla_python_cpu_callback(%[[C]], %arg0, %arg1) {
   // CHECK-SAME:   api_version = 2 : i32, backend_config = "56238119409280",
   // CHECK-V2-SAME:   has_side_effect = true, mhlo.sharding = "{maximal device=0}",
-  // CHECK-V3-SAME:   has_side_effect = true, mhlo.sharding = "{maximal_mesh[device_id=0]}",
+  // CHECK-V3-SAME:   has_side_effect = true, mhlo.sharding = "{single_device[device_id=0]}",
   // CHECK-SAME:   operand_layouts = [dense<> : tensor<0xindex>, dense<> : tensor<0xindex>, dense<0> : tensor<1xindex>],
   // CHECK-SAME:   result_layouts = [dense<> : tensor<0xindex>]
   // CHECK-SAME: } : (tensor<i64>, !stablehlo.token, tensor<2xi64>) -> tuple<!stablehlo.token>
@@ -464,7 +464,7 @@ func.func @callback_no_tuple_result_used(%arg0: tensor<2xf64>) -> tensor<2xf64>
 // CHECK-SAME:      (%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
 func.func @maximal_sharding_no_results(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
   // CHECK-V2-NEXT: stablehlo.custom_call @foo(%arg0) {has_side_effect = true, mhlo.sharding = "{maximal device=0}"} : (tensor<8x8xf32>) -> ()
-  // CHECK-V3-NEXT: stablehlo.custom_call @foo(%arg0) {has_side_effect = true, mhlo.sharding = "{maximal_mesh[device_id=0]}"} : (tensor<8x8xf32>) -> ()
+  // CHECK-V3-NEXT: stablehlo.custom_call @foo(%arg0) {has_side_effect = true, mhlo.sharding = "{single_device[device_id=0]}"} : (tensor<8x8xf32>) -> ()
   // CHECK-NEXT: return %arg0 : tensor<8x8xf32>
   stablehlo.custom_call @foo(%arg0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh_0, []>]>} : (tensor<8x8xf32>) -> ()
   return %arg0 : tensor<8x8xf32>
@@ -477,10 +477,10 @@ func.func @reshard_maximal_sharding(%arg0: tensor<8x8xf32>) -> (tensor<8x8xf32>,
   // CHECK-V2-NEXT: %[[COPY_0:.*]] = mhlo.copy %[[CALL]] {mhlo.sharding = "{devices=[8,1,4]<=[32] last_tile_dim_replicate}"} : tensor<8x8xf32>
   // CHECK-V2-NEXT: %[[COPY_1:.*]] = mhlo.copy %[[CALL]] {mhlo.sharding = "{devices=[4,1,8]<=[8,4]T(1,0) last_tile_dim_replicate}"} : tensor<8x8xf32>
   // CHECK-V2-NEXT: %[[COPY_2:.*]] = mhlo.copy %[[CALL]] {mhlo.sharding = "{maximal device=1}"} : tensor<8x8xf32>
-  // CHECK-V3-NEXT: %[[CALL:.*]] = stablehlo.custom_call @foo(%arg0) {has_side_effect = true, mhlo.sharding = "{maximal_mesh[device_id=0]}"} : (tensor<8x8xf32>) -> tensor<8x8xf32>
+  // CHECK-V3-NEXT: %[[CALL:.*]] = stablehlo.custom_call @foo(%arg0) {has_side_effect = true, mhlo.sharding = "{single_device[device_id=0]}"} : (tensor<8x8xf32>) -> tensor<8x8xf32>
   // CHECK-V3-NEXT: %[[COPY_0:.*]] = mhlo.copy %[[CALL]] {mhlo.sharding = "{mesh['x'=8,'y'=4], [{'x'}, {}]}"} : tensor<8x8xf32>
   // CHECK-V3-NEXT: %[[COPY_1:.*]] = mhlo.copy %[[CALL]] {mhlo.sharding = "{mesh['x'=8,'y'=4], [{'y'}, {}]}"} : tensor<8x8xf32>
-  // CHECK-V3-NEXT: %[[COPY_2:.*]] = mhlo.copy %[[CALL]] {mhlo.sharding = "{maximal_mesh[device_id=1]}"} : tensor<8x8xf32>
+  // CHECK-V3-NEXT: %[[COPY_2:.*]] = mhlo.copy %[[CALL]] {mhlo.sharding = "{single_device[device_id=1]}"} : tensor<8x8xf32>
   // CHECK-NEXT: return %[[CALL]], %[[COPY_0]], %[[COPY_1]], %[[COPY_2]] : tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>
   %0 = sdy.reshard %arg0 <@maximal_mesh_0, []> : tensor<8x8xf32>
   %1 = stablehlo.custom_call @foo(%0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh_0, []>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32>
diff --git a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_pipeline_simplify_replicated.mlir b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_pipeline_simplify_replicated.mlir
index 994d29a7492..1f716b800cd 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_pipeline_simplify_replicated.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_pipeline_simplify_replicated.mlir
@@ -55,7 +55,7 @@ sdy.mesh @maximal_mesh = <[], device_ids=[0]>
 // CHECK-LABEL: func @simplify_replicated_sharding_empty_mesh(
 // CHECK-SAME:      %arg0: tensor<8x8xf32> {mhlo.sharding = "{mesh[], replicated}"},
 // CHECK-SAME:      %arg1: tensor<8x8xf32> {mhlo.sharding = "{mesh[], [{?}, {}]}"},
-// CHECK-SAME:      %arg2: tensor<8x8xf32> {mhlo.sharding = "{maximal_mesh[device_id=0]}"})
+// CHECK-SAME:      %arg2: tensor<8x8xf32> {mhlo.sharding = "{single_device[device_id=0]}"})
 func.func @simplify_replicated_sharding_empty_mesh(%arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@empty_mesh, [{}, {}]>},
                                                    %arg1: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@empty_mesh, [{?}, {}]>},
                                                    %arg2: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@maximal_mesh, []>}) -> tensor<8x8xf32> {
diff --git a/third_party/xla/xla/service/spmd/shardy/utils.h b/third_party/xla/xla/service/spmd/shardy/utils.h
index 55bcc364b9c..d145879cb30 100644
--- a/third_party/xla/xla/service/spmd/shardy/utils.h
+++ b/third_party/xla/xla/service/spmd/shardy/utils.h
@@ -181,7 +181,8 @@ bool hasFrontendMeshes(mlir::ModuleOp module);
 
 // Check if the module has any sort of Shardy mesh:
 // - `mesh`
-// - `maximal_mesh_{X}`
+// - `single_device_{X}`
+// - `maximal_mesh_{X}` (for backward compatibility)
 // - `empty_mesh`
 // TODO(b/420837831): delete this once we don't fall back to GSPMD.
 bool hasShardyMesh(mlir::ModuleOp module);
diff --git a/third_party/xla/xla/xla_data.proto b/third_party/xla/xla/xla_data.proto
index 604b7fa3f12..0f14dc62de0 100644
--- a/third_party/xla/xla/xla_data.proto
+++ b/third_party/xla/xla/xla_data.proto
@@ -1066,7 +1066,7 @@ message StatisticsViz {
 // If the list of axes is empty
 //   - If the `device_ids` is not provided, it is an empty mesh.
 //   - If the `device_ids` is provided, it must be a single non-negative
-//     integer, we call it a **maximal-sharding mesh**.
+//     integer, we call it a **single-device mesh**.
 
 // If the list of axes is provided
 //   - If a device ID list is specified, the product of the axis sizes must
@@ -1074,14 +1074,14 @@ message StatisticsViz {
 //   - If a device ID list is not specified, then the device ID list is
 //     implicitly `iota(product(axis_sizes))` with an optional transform
 //     in the `iota_transform` field applied.
-//   - It is not a maximal-sharding mesh even if the total size of axes is 1.
+//   - It is not a single-device mesh even if the total size of axes is 1.
 
 // Here are some examples of meshes:
 
 // - An empty mesh represents a placeholder mesh that can be replaced during
 //   propagation: <[]>
 // - A mesh without axes list and a single non-negative device ID, which is a
-//   maximal-sharding mesh: <[], device_ids=[3]>
+//   single-device mesh: <[], device_ids=[3]>
 // - A mesh with two axes and device IDs defined as iota(6) with no
 //   transformation applied: <["a"=2, "b"=3]>
 // - A mesh with two axes and device IDs defined as iota(6) with a transpose
@@ -1092,7 +1092,7 @@ message StatisticsViz {
 // **Constraints:**
 // - Elements in `device_ids` must be non-negative.
 // - If `axes` is empty, the size of `device_ids` can be 0 (empty mesh) or 1
-//   (maximal-sharding mesh).
+//   (single-device mesh).
 // - If `axes` is not empty,
 //     - Elements in `axes` must not have duplicate names.
 //     - If `device_ids` is specified, the original `device_ids` is not

From 97e039b134000e5894ef24ebd87b8b30521f1a6c Mon Sep 17 00:00:00 2001
From: Dmitri Latushko <latushko@google.com>
Date: Mon, 12 Oct 2026 15:40:57 -0700
Subject: [PATCH 0627/1000] Prefer non-derived root-cause errors over earlier
 derived cancellation errors in ExecutorState::NodeDone.

PiperOrigin-RevId: 998198613
---
 tensorflow/core/common_runtime/executor.cc    |  15 ++-
 .../core/common_runtime/executor_test.cc      | 120 ++++++++++++++++++
 2 files changed, 133 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/common_runtime/executor.cc b/tensorflow/core/common_runtime/executor.cc
index 615c58f47ab..b9c932ff906 100644
--- a/tensorflow/core/common_runtime/executor.cc
+++ b/tensorflow/core/common_runtime/executor.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/memory/memory.h"
+#include "absl/status/status.h"
 #include "absl/strings/ascii.h"
 #include "absl/strings/numbers.h"
 #include "absl/strings/str_join.h"
@@ -1253,6 +1254,11 @@ bool ExecutorState<PropagatorStateType>::NodeDone(
     // Some error happened. This thread of computation is done.
     {
       mutex_lock l(mu_);
+      const bool is_derived = StatusGroup::IsDerived(s) ||
+                              ((absl::IsCancelled(s) || absl::IsAborted(s)) &&
+                               cancellation_manager_ &&
+                               (cancellation_manager_->IsCancelled() ||
+                                cancellation_manager_->IsCancelling()));
       if (status_.ok()) {
         // If this is the first node to fail in this run, we are responsible for
         // aborting all other execution in the step.
@@ -1262,13 +1268,18 @@ bool ExecutorState<PropagatorStateType>::NodeDone(
         // being derived. Note that the original node that fails might also
         // trigger cancellation, and here we make sure the original error is
         // exposed to users and not buried as a derived error.
-        if (cancellation_manager_ && cancellation_manager_->IsCancelled() &&
-            (absl::IsCancelled(s) || absl::IsAborted(s))) {
+        if (is_derived && !StatusGroup::IsDerived(s)) {
           status_ = StatusGroup::MakeDerived(s);
           maybe_derived_s = status_;
         } else {
           status_ = s;
         }
+      } else if (StatusGroup::IsDerived(status_) && !is_derived) {
+        // If a previous node in this executor reported a derived cancellation
+        // error (for example, a sibling _Recv cancelled via a shared
+        // CancellationManager before a nested function finished reporting its
+        // original error), prefer the non-derived root-cause error.
+        status_ = s;
       }
     }
 
diff --git a/tensorflow/core/common_runtime/executor_test.cc b/tensorflow/core/common_runtime/executor_test.cc
index 2097278fbbb..1a6df72f3f1 100644
--- a/tensorflow/core/common_runtime/executor_test.cc
+++ b/tensorflow/core/common_runtime/executor_test.cc
@@ -31,6 +31,7 @@ limitations under the License.
 #include "tensorflow/core/common_runtime/process_util.h"
 #include "tensorflow/core/common_runtime/step_stats_collector.h"
 #include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/cancellation.h"
 #include "tensorflow/core/framework/local_rendezvous.h"
 #include "tensorflow/core/framework/op.h"
 #include "tensorflow/core/framework/rendezvous.h"
@@ -39,8 +40,10 @@ limitations under the License.
 #include "tensorflow/core/framework/versions.pb.h"
 #include "tensorflow/core/graph/algorithm.h"
 #include "tensorflow/core/graph/testlib.h"
+#include "tensorflow/core/lib/core/status.h"
 #include "tensorflow/core/lib/core/status_test_util.h"
 #include "tensorflow/core/lib/random/simple_philox.h"
+#include "tensorflow/core/lib/strings/str_util.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/platform/logging.h"
 #include "tensorflow/core/platform/strcat.h"
@@ -498,6 +501,123 @@ TEST_F(ExecutorTest, NoInputTensors) {
   TF_ASSERT_OK(Run(rendez_));
 }
 
+namespace {
+
+// A test Rendezvous that waits until all three _Recv nodes are in-flight,
+// cancels the shared CancellationManager (simulating a nested function
+// cancelling the shared manager when an op inside it fails), and then completes
+// the _Recv callbacks in deterministic order:
+//   1. "derived_cancelled": CancelledError while CancellationManager is
+//      cancelled (marked derived by ExecutorState::NodeDone).
+//   2. "root_cause_error": non-derived InvalidArgumentError (simulating the
+//      nested function completing with its original root-cause error).
+//   3. "later_derived_cancelled": derived CancelledError arriving after the
+//      non-derived error has already been recorded.
+class OrderedFailureRendezvous : public Rendezvous {
+ public:
+  absl::Status Send(const ParsedKey& key, const Args& send_args,
+                    const Tensor& val, const bool is_dead) override {
+    return absl::OkStatus();
+  }
+
+  void RecvAsync(const ParsedKey& key, const Args& recv_args,
+                 DoneCallback done) override {
+    Args derived_args;
+    DoneCallback derived_done;
+    Args root_cause_args;
+    DoneCallback root_cause_done;
+    Args later_derived_args;
+    DoneCallback later_derived_done;
+    {
+      mutex_lock l(mu_);
+      if (key.edge_name == "derived_cancelled") {
+        derived_args_ = recv_args;
+        derived_done_ = std::move(done);
+      } else if (key.edge_name == "root_cause_error") {
+        root_cause_args_ = recv_args;
+        root_cause_done_ = std::move(done);
+      } else if (key.edge_name == "later_derived_cancelled") {
+        later_derived_args_ = recv_args;
+        later_derived_done_ = std::move(done);
+      }
+      if (!derived_done_ || !root_cause_done_ || !later_derived_done_) {
+        return;
+      }
+      derived_args = derived_args_;
+      derived_done = std::move(derived_done_);
+      root_cause_args = root_cause_args_;
+      root_cause_done = std::move(root_cause_done_);
+      later_derived_args = later_derived_args_;
+      later_derived_done = std::move(later_derived_done_);
+    }
+
+    if (derived_args.cancellation_manager != nullptr) {
+      CancellationManager* cm = derived_args.cancellation_manager;
+      CancellationToken token = cm->get_cancellation_token();
+      bool registered = cm->RegisterCallback(token, [&]() {
+        derived_done(absl::CancelledError("RecvFromRemoteAsync is cancelled."),
+                     Rendezvous::Args(), derived_args, Tensor(),
+                     /*is_dead=*/false);
+      });
+      CHECK(registered);
+      cm->StartCancel();
+    } else {
+      derived_done(absl::CancelledError("RecvFromRemoteAsync is cancelled."),
+                   Rendezvous::Args(), derived_args, Tensor(),
+                   /*is_dead=*/false);
+    }
+    root_cause_done(
+        absl::InvalidArgumentError("Nan in summary histogram for: tag"),
+        Rendezvous::Args(), root_cause_args, Tensor(), /*is_dead=*/false);
+    later_derived_done(
+        StatusGroup::MakeDerived(absl::CancelledError("Later recv cancelled.")),
+        Rendezvous::Args(), later_derived_args, Tensor(), /*is_dead=*/false);
+  }
+
+  void StartAbort(const absl::Status& status) override {}
+
+ private:
+  mutex mu_;
+  Args derived_args_ TF_GUARDED_BY(mu_);
+  DoneCallback derived_done_ TF_GUARDED_BY(mu_);
+  Args root_cause_args_ TF_GUARDED_BY(mu_);
+  DoneCallback root_cause_done_ TF_GUARDED_BY(mu_);
+  Args later_derived_args_ TF_GUARDED_BY(mu_);
+  DoneCallback later_derived_done_ TF_GUARDED_BY(mu_);
+};
+
+}  // namespace
+
+TEST_F(ExecutorTest, NonDerivedErrorOverridesDerivedCancellationError) {
+  auto g = std::make_unique<Graph>(OpRegistry::Global());
+  auto in0 =
+      test::graph::Recv(g.get(), "derived_cancelled", "float", ALICE, 1, BOB);
+  auto in1 =
+      test::graph::Recv(g.get(), "root_cause_error", "float", ALICE, 1, BOB);
+  auto in2 = test::graph::Recv(g.get(), "later_derived_cancelled", "float",
+                               ALICE, 1, BOB);
+  auto add0 = test::graph::Add(g.get(), in0, in1);
+  auto add1 = test::graph::Add(g.get(), add0, in2);
+  test::graph::Send(g.get(), add1, "out", BOB, 1, ALICE);
+  Create(std::move(g));
+
+  CancellationManager cm;
+  auto* rendez = new OrderedFailureRendezvous();
+  Executor::Args args;
+  args.rendezvous = rendez;
+  args.cancellation_manager = &cm;
+  args.stats_collector = &step_stats_collector_;
+  args.runner = runner_;
+
+  absl::Status s = exec_->Run(args);
+  EXPECT_TRUE(absl::IsInvalidArgument(s)) << s;
+  EXPECT_FALSE(StatusGroup::IsDerived(s)) << s;
+  EXPECT_TRUE(
+      absl::StrContains(s.message(), "Nan in summary histogram for: tag"))
+      << s;
+  rendez->Unref();
+}
+
 // Create a graph that is 'depth' deep. At each level, fan-in and fan-out a
 // maximum of 'width' nodes. All nodes are no-ops and all dependencies are
 // control dependencies.

From 4ce902b79b5fb0a18cdadcb417d56bbce24a915b Mon Sep 17 00:00:00 2001
From: Vlad Sytchenko <vsytch@google.com>
Date: Mon, 12 Oct 2026 16:02:38 -0700
Subject: [PATCH 0628/1000] [XLA] Only default to custom interference if there
 is no HLO interference

PiperOrigin-RevId: 998210421
---
 .../xla/xla/service/buffer_assignment.cc      | 77 +++++++++----------
 .../xla/xla/service/buffer_assignment.h       |  5 +-
 .../xla/xla/service/buffer_assignment_test.cc | 12 ++-
 3 files changed, 49 insertions(+), 45 deletions(-)

diff --git a/third_party/xla/xla/service/buffer_assignment.cc b/third_party/xla/xla/service/buffer_assignment.cc
index e3f1da7179a..4d460e4bdf9 100644
--- a/third_party/xla/xla/service/buffer_assignment.cc
+++ b/third_party/xla/xla/service/buffer_assignment.cc
@@ -1926,20 +1926,6 @@ bool BufferAssigner::LiveRangeInterferes(
     const HloValue* buffer1, const HloLiveRange::LiveRangeBounds& live_range1,
     const HloValue* buffer2, const HloLiveRange::LiveRangeBounds& live_range2,
     BufferAssignment* assignment) {
-  if (opts_.live_range_interference_options.has_value()) {
-    const auto& options = *opts_.live_range_interference_options;
-    if (options.has_custom_interference(assignment->alias_analysis(),
-                                        *buffer1) &&
-        options.has_custom_interference(assignment->alias_analysis(),
-                                        *buffer2)) {
-      std::optional<bool> custom_interference =
-          options.interferes(assignment->alias_analysis(), *buffer1, *buffer2);
-      if (custom_interference.has_value()) {
-        return *custom_interference;
-      }
-    }
-  }
-
   CHECK((assignment->hlo_live_range().total_order_scheduled()));
 
   // An HloValue can hold multiple instruction positions during its lifetime
@@ -1969,6 +1955,7 @@ bool BufferAssigner::LiveRangeInterferes(
         return false;
       };
 
+  bool interferes = false;
   if (!(live_range1.start > live_range2.end ||
         live_range2.start > live_range1.end)) {
     if (live_range1.end == live_range2.start) {
@@ -1978,7 +1965,7 @@ bool BufferAssigner::LiveRangeInterferes(
         VLOG(4) << "End of live range of " << buffer1->ToShortString()
                 << " is equal to the start of live range of "
                 << buffer2->ToShortString() << ", buffer cannot be shared.";
-        return true;
+        interferes = true;
       }
     } else if (live_range2.end == live_range1.start) {
       auto operand_value = buffer2;
@@ -1987,7 +1974,7 @@ bool BufferAssigner::LiveRangeInterferes(
         VLOG(4) << "End of live range of " << buffer2->ToShortString()
                 << " is equal to the start of live range of "
                 << buffer1->ToShortString() << ", buffer cannot be shared.";
-        return true;
+        interferes = true;
       }
     } else {
       VLOG(4) << "Can't assign: assignee " << *buffer1 << " may interfere with "
@@ -1996,10 +1983,25 @@ bool BufferAssigner::LiveRangeInterferes(
       VLOG(4) << "assigned_buffer.end: " << live_range1.end;
       VLOG(4) << "live_range2.start" << live_range2.start;
       VLOG(4) << "live_range2.end" << live_range2.end;
-      return true;
+      interferes = true;
     }
   }
-  return false;
+
+  if (interferes && opts_.live_range_interference_options.has_value()) {
+    const auto& options = *opts_.live_range_interference_options;
+    if (options.has_custom_interference(assignment->alias_analysis(),
+                                        *buffer1) &&
+        options.has_custom_interference(assignment->alias_analysis(),
+                                        *buffer2)) {
+      std::optional<bool> custom_interference =
+          options.interferes(assignment->alias_analysis(), *buffer1, *buffer2);
+      if (custom_interference.has_value()) {
+        return *custom_interference;
+      }
+    }
+  }
+
+  return interferes;
 }
 
 absl::StatusOr<bool> BufferAssigner::MaybeAssignBuffer(
@@ -2129,32 +2131,27 @@ absl::StatusOr<bool> BufferAssigner::MaybeAssignBuffer(
           return false;
         }
       } else {
-        std::optional<bool> custom_interference;
-        if (opts_.live_range_interference_options.has_value()) {
-          const auto& options = *opts_.live_range_interference_options;
-          if (options.has_custom_interference(assignment->alias_analysis(),
-                                              *new_value) &&
-              options.has_custom_interference(assignment->alias_analysis(),
-                                              assigned_buffer)) {
-            custom_interference = options.interferes(
-                assignment->alias_analysis(), *new_value, assigned_buffer);
+        if (assignment->hlo_ordering().MayInterfere(
+                assigned_buffer, *new_value, assignment->dataflow_analysis(),
+                alias_info_)) {
+          std::optional<bool> custom_interference;
+          if (opts_.live_range_interference_options.has_value()) {
+            const auto& options = *opts_.live_range_interference_options;
+            if (options.has_custom_interference(assignment->alias_analysis(),
+                                                *new_value) &&
+                options.has_custom_interference(assignment->alias_analysis(),
+                                                assigned_buffer)) {
+              custom_interference = options.interferes(
+                  assignment->alias_analysis(), *new_value, assigned_buffer);
+            }
           }
-        }
-        if (custom_interference.has_value()) {
-          if (*custom_interference) {
+          if (custom_interference.value_or(true)) {
+            // Fallback to partial order based interference detection (slower)
+            // when we don't have a total order scheduled module.
             VLOG(4) << "Can't assign: assignee " << assigned_buffer
-                    << " custom interference with "
-                    << new_value->ToShortString();
+                    << " may interfere with " << new_value->ToShortString();
             return false;
           }
-        } else if (assignment->hlo_ordering().MayInterfere(
-                       assigned_buffer, *new_value,
-                       assignment->dataflow_analysis(), alias_info_)) {
-          // Fallback to partial order based interference detection (slower)
-          // when we don't have a total order scheduled module.
-          VLOG(4) << "Can't assign: assignee " << assigned_buffer
-                  << " may interfere with " << new_value->ToShortString();
-          return false;
         }
       }
 
diff --git a/third_party/xla/xla/service/buffer_assignment.h b/third_party/xla/xla/service/buffer_assignment.h
index 38a89223cd4..464754cdf69 100644
--- a/third_party/xla/xla/service/buffer_assignment.h
+++ b/third_party/xla/xla/service/buffer_assignment.h
@@ -496,9 +496,12 @@ class BufferAssignment {
         has_custom_interference;
 
     // Symmetric binary predicate: returns true/false to override interference
-    // between lhs and rhs, or std::nullopt to fall back to default analysis.
+    // between lhs and rhs when their HLO live ranges overlap, or std::nullopt
+    // to fall back to default analysis.
     // Must be symmetric: interferes(alias_analysis, a, b) ==
     // interferes(alias_analysis, b, a).
+    // Only consulted when standard HLO live range analysis indicates that lhs
+    // and rhs overlap; ignored when their HLO live ranges do not overlap.
     // Returning true forces interference (preventing allocation reuse), while
     // returning false forces non-interference (enabling allocation reuse).
     //
diff --git a/third_party/xla/xla/service/buffer_assignment_test.cc b/third_party/xla/xla/service/buffer_assignment_test.cc
index ca87aadb1b5..c7b8dc19a58 100644
--- a/third_party/xla/xla/service/buffer_assignment_test.cc
+++ b/third_party/xla/xla/service/buffer_assignment_test.cc
@@ -5971,15 +5971,17 @@ ENTRY main {
   EXPECT_EQ(buffer0_slice.offset(), slice0.offset());
 }
 
-TEST_F(BufferAssignmentTest, CustomInterferenceForcesInterference) {
+TEST_F(BufferAssignmentTest,
+       CustomInterferenceIgnoredWhenLiveRangesDoNotOverlap) {
   constexpr absl::string_view kHlo = R"hlo(
 HloModule main, is_scheduled=true
 
 ENTRY main {
   buffer.0 = s32[4096]{0} custom-call(), custom_call_target="AllocateBuffer"
-  copy.0 = s32[4096]{0} copy(buffer.0)
+  use.0 = (s32[4096]{0}) tuple(buffer.0)
   buffer.1 = s32[4096]{0} custom-call(), custom_call_target="AllocateBuffer"
-  ROOT copy.1 = s32[4096]{0} copy(buffer.1)
+  use.1 = (s32[4096]{0}) tuple(buffer.1)
+  ROOT root = tuple()
 }
 )hlo";
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
@@ -6022,7 +6024,9 @@ ENTRY main {
   ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice buffer1_slice,
                        buffers->GetUniqueTopLevelSlice(buffer1));
 
-  EXPECT_NE(buffer0_slice.index(), buffer1_slice.index());
+  // buffer.0 and buffer.1 have non-overlapping HLO live ranges, so custom
+  // interference returning true is ignored and they still share an allocation.
+  EXPECT_EQ(buffer0_slice, buffer1_slice);
 }
 
 TEST_F(BufferAssignmentTest, HasLiveRangeInterferenceOptionFastMerge) {

From c1e67201498a4294ecef88bc222335fb6b186908 Mon Sep 17 00:00:00 2001
From: Christian Sigg <csigg@google.com>
Date: Mon, 12 Oct 2026 16:25:00 -0700
Subject: [PATCH 0629/1000] [XLA:GPU] Add the TensorIR fusion emitter and
 autotuner backend behind a flag.

This imports NVIDIA's TensorIR fusion emitter (github.com/sergey-kozub/xla
branch experimental, commits f7e1b8db43 and 2f0c163c4a, rebased onto
openxla/xla main 05783d744f).

The emitter replaces the hand-written HLO->TensorIR converter
(conversion.{h,cc}, which built nv_tensor_ir ops directly with an
OpBuilder) with an MLIR lowering pipeline: HloFunctionImporter emits
StableHLO and --legalize-stablehlo-to-tensor-ir rewrites it to an
nv_tensor_ir.graph; hlo_to_tensor_ir::ImportAndLegalizeComputation is
shared between the emitter and the autotuner. Layouts become boundary
stride and alignment attributes. Kernels are assembled ahead of time
with tileiras and launched through CustomKernelThunk, which makes them
serializable and command-buffer capturable.

Autotuning uses the TensorIR tiling cost model and picks a spread of
tile sizes within the register budget. Compiled cubins go through the
kernel reuse cache so identical fusions share one compilation. Clamp and
select with scalar operands are rejected (operand shapes must be
uniform).

The backend is guarded by --xla_gpu_experimental_enable_tensor_ir (off
by default): the CUDA factory always registers TensorIrBackend and
ConfigAssignerPass::GetEnabledBackends prunes it unless the flag is set,
so the emitter is inert by default. It needs sm_90+ and a CUDA 13.3+
toolkit (tileiras) at runtime.

Changes relative to NVIDIA's commits:
- Dropped the superseded conversion.{h,cc}.
- llvm::stable_sort for equal-cost tilings in the autotuner.
- The kernel-cache continuation captures the fusion name by value.
- Removed unused includes and deps in fusion.cc/BUILD; added the deps
  the layering check requires.
- Include-what-you-use cleanup (direct includes and deps for every used
  symbol); TensorIrFusion drops its unused HloFusionAnalysis reference;
  factory_test drops two deps it never included (flagged by the OSS dwyu
  check).
- Modules are created with xla::llvm_ir::CreateMlirModuleOp (XLA forbids
  mlir::ModuleOp::create); ApplyConfig rejects unsupported instructions
  like the other backends; CreateFloatSplat takes an integral literal,
  the rank-0 result reshape uses an explicit getResult(), the tiling
  groups live in an absl::btree_map, and the autotuner test checks the
  actual cap of 10 configs.
- fusion and custom_kernel are cuda-only (the tensor_ir runtime and
  compiler include cuda.h); builds without CUDA link a fusion_stub whose
  Emit returns Unimplemented, selected in codegen:fusions through
  if_cuda_is_configured, so the OSS CPU, ROCm and oneAPI builds work.

XLA's WORKSPACE builds (--noenable_bzlmod, the "XLA Linux x86 CPU
WORKSPACE" CI job) had no flatbuffers repository, which @tensor_ir's
Artifact schema needs now that the emitter links it:
third_party/flatbuffers/{workspace.bzl,flatbuffers.BUILD} define it
(C++ targets only, the 25.9.23 that TensorFlow pins) and MODULE.bazel
moves the BCR module to the same version.

TensorFlow's WORKSPACE build initializes XLA's third-party repositories
one by one in tensorflow/workspace2.bzl; it now also initializes
cuda_tile and tensor_ir, which the GPU compiler depends on through the
emitter.

Co-authored-by: Sergey Kozub <skozub@nvidia.com>
PiperOrigin-RevId: 998222565
---
 tensorflow/workspace2.bzl                     |    4 +
 third_party/xla/MODULE.bazel                  |    2 +-
 third_party/xla/third_party/flatbuffers/BUILD |    1 +
 .../third_party/flatbuffers/flatbuffers.BUILD |  151 ++
 .../xla/third_party/flatbuffers/workspace.bzl |   17 +
 third_party/xla/workspace2.bzl                |    4 +
 .../backends/autotuner/backend_config.proto   |    1 +
 .../xla/xla/backends/autotuner/backends.proto |    1 +
 .../xla/xla/backends/gpu/autotuner/BUILD      |   77 +-
 .../backends/gpu/autotuner/factory_cuda.cc    |    3 +
 .../backends/gpu/autotuner/factory_test.cc    |    2 +-
 .../xla/backends/gpu/autotuner/tensor_ir.cc   |  267 +++
 .../xla/backends/gpu/autotuner/tensor_ir.h    |   65 +
 .../backends/gpu/autotuner/tensor_ir_test.cc  |  285 +++
 .../xla/xla/backends/gpu/codegen/BUILD        |    9 +-
 .../xla/xla/backends/gpu/codegen/fusions.cc   |    3 +
 .../xla/backends/gpu/codegen/tensor_ir/BUILD  |  247 ++-
 .../codegen/tensor_ir/compilation_pipeline.cc |   95 +
 .../codegen/tensor_ir/compilation_pipeline.h  |   56 +
 .../gpu/codegen/tensor_ir/conversion.cc       | 1041 ---------
 .../gpu/codegen/tensor_ir/conversion.h        |   40 -
 .../gpu/codegen/tensor_ir/custom_kernel.cc    |   99 +
 .../gpu/codegen/tensor_ir/custom_kernel.h     |   39 +
 .../backends/gpu/codegen/tensor_ir/fusion.cc  |  286 +++
 .../backends/gpu/codegen/tensor_ir/fusion.h   |   38 +
 .../gpu/codegen/tensor_ir/fusion_stub.cc      |   33 +
 .../gpu/codegen/tensor_ir/hlo_to_tensor_ir.cc |  140 ++
 .../gpu/codegen/tensor_ir/hlo_to_tensor_ir.h  |   50 +
 .../tensor_ir/hlo_to_tensor_ir_test.cc        |  289 +++
 .../gpu/codegen/tensor_ir/layout_utils.cc     |  158 ++
 .../gpu/codegen/tensor_ir/layout_utils.h      |   71 +
 .../codegen/tensor_ir/layout_utils_test.cc    |  414 ++++
 .../backends/gpu/codegen/tensor_ir/support.cc |   32 +
 .../backends/gpu/codegen/tensor_ir/support.h  |    9 +
 .../gpu/codegen/tensor_ir/support_test.cc     |   58 +
 .../gpu/codegen/tensor_ir/tests/BUILD         |   93 +
 .../tensor_ir/tests/elementwise_convert.hlo   |   11 +-
 .../tests/fusion_emitter_device_test.cc       |  594 ++++++
 .../tests/fusion_emitter_thunk_test.cc        |  182 ++
 .../tensor_ir/tests/layout_bitcast.hlo        |    7 +
 .../tensor_ir/tests/layout_broadcast.hlo      |   16 +-
 .../codegen/tensor_ir/tests/reduce_float.hlo  |   19 +-
 .../tensor_ir/tests/reduce_int_signed.hlo     |   19 +-
 .../tensor_ir/tests/reduce_int_unsigned.hlo   |   19 +-
 .../gpu/codegen/tensor_ir/transforms/BUILD    |   70 +
 .../legalize_stablehlo_to_tensor_ir.cc        | 1877 +++++++++++++++++
 .../gpu/codegen/tensor_ir/transforms/passes.h |   35 +
 .../codegen/tensor_ir/transforms/passes.td    |   39 +
 .../codegen/tensor_ir/transforms/tests/BUILD  |   39 +
 .../legalize_stablehlo_to_tensor_ir.mlir      |  468 ++++
 ...galize_stablehlo_to_tensor_ir_bitcast.mlir |   74 +
 ...galize_stablehlo_to_tensor_ir_compare.mlir |  136 ++
 ...galize_stablehlo_to_tensor_ir_invalid.mlir |  172 ++
 ...egalize_stablehlo_to_tensor_ir_matmul.mlir |  241 +++
 ...egalize_stablehlo_to_tensor_ir_reduce.mlir |  283 +++
 ...ize_stablehlo_to_tensor_ir_structural.mlir |  223 ++
 .../tensor_ir/transforms/type_conversion.cc   |   65 +
 .../tensor_ir/transforms/type_conversion.h    |   30 +
 .../xla/xla/backends/gpu/codegen/tools/BUILD  |    4 +-
 .../gpu/codegen/tools/fusion_to_tensorir.cc   |   56 +-
 third_party/xla/xla/backends/gpu/tests/BUILD  |    3 +
 .../xla/xla/backends/gpu/tests/xla-opt.cc     |   13 +-
 .../transforms/fusion_block_level_rewriter.cc |   10 +
 .../gpu/transforms/priority_fusion.cc         |    1 +
 third_party/xla/xla/debug_options_flags.cc    |    8 +
 .../gpu/autotuning/config_assigner_pass.cc    |   63 +-
 .../autotuning/config_assigner_pass_test.cc   |   47 +
 .../xla/xla/service/gpu/backend_configs.proto |   17 +-
 .../xla/service/gpu/hlo_fusion_analysis.cc    |    4 +
 .../xla/xla/service/gpu/hlo_fusion_analysis.h |    1 +
 .../xla/xla/service/gpu/ir_emission_utils.h   |    4 +
 third_party/xla/xla/xla.proto                 |    7 +-
 72 files changed, 7832 insertions(+), 1205 deletions(-)
 create mode 100644 third_party/xla/third_party/flatbuffers/BUILD
 create mode 100644 third_party/xla/third_party/flatbuffers/flatbuffers.BUILD
 create mode 100644 third_party/xla/third_party/flatbuffers/workspace.bzl
 create mode 100644 third_party/xla/xla/backends/gpu/autotuner/tensor_ir.cc
 create mode 100644 third_party/xla/xla/backends/gpu/autotuner/tensor_ir.h
 create mode 100644 third_party/xla/xla/backends/gpu/autotuner/tensor_ir_test.cc
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/compilation_pipeline.cc
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/compilation_pipeline.h
 delete mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/conversion.cc
 delete mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/conversion.h
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/custom_kernel.cc
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/custom_kernel.h
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/fusion.cc
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/fusion.h
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/fusion_stub.cc
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir.cc
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir.h
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir_test.cc
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils.cc
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils.h
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils_test.cc
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/support_test.cc
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/fusion_emitter_device_test.cc
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/fusion_emitter_thunk_test.cc
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/BUILD
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/legalize_stablehlo_to_tensor_ir.cc
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/passes.h
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/passes.td
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/BUILD
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir.mlir
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_bitcast.mlir
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_compare.mlir
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_invalid.mlir
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_matmul.mlir
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_reduce.mlir
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_structural.mlir
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/type_conversion.cc
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/type_conversion.h

diff --git a/tensorflow/workspace2.bzl b/tensorflow/workspace2.bzl
index bcfe4415614..da028662d57 100644
--- a/tensorflow/workspace2.bzl
+++ b/tensorflow/workspace2.bzl
@@ -28,6 +28,7 @@ load("@xla//third_party/benchmark:workspace.bzl", benchmark = "repo")
 load("@xla//third_party/clang_toolchain:cc_configure_clang.bzl", "cc_download_clang_toolchain")
 load("@xla//third_party/compute_library:workspace.bzl", compute_library = "repo")
 load("@xla//third_party/cpuinfo:workspace.bzl", cpuinfo = "repo")
+load("@xla//third_party/cuda_tile:workspace.bzl", cuda_tile = "repo")
 load("@xla//third_party/cudnn_frontend:workspace.bzl", cudnn_frontend = "repo")
 load("@xla//third_party/cutlass:workspace.bzl", cutlass = "repo")
 load("@xla//third_party/dlpack:workspace.bzl", dlpack = "repo")
@@ -63,6 +64,7 @@ load("@xla//third_party/shardy:workspace.bzl", shardy = "repo")
 load("@xla//third_party/slinky:workspace.bzl", slinky = "repo")
 load("@xla//third_party/spdlog:workspace.bzl", spdlog = "repo")
 load("@xla//third_party/stablehlo:workspace.bzl", stablehlo = "repo")
+load("@xla//third_party/tensor_ir:workspace.bzl", tensor_ir = "repo")
 load("@xla//third_party/tensorrt:tensorrt_configure.bzl", "tensorrt_configure")
 load("@xla//third_party/tensorrt:workspace.bzl", tensorrt = "repo")
 load("@xla//third_party/triton:workspace.bzl", triton = "repo")
@@ -126,6 +128,7 @@ def _initialize_third_party():
     bazel_features_deps()
     benchmark()
     com_google_highway()
+    cuda_tile()
     ducc()
     dlpack()
     eigen3()
@@ -163,6 +166,7 @@ def _initialize_third_party():
     spdlog()
     sobol_data()
     stablehlo()
+    tensor_ir()
     vulkan_headers()
     tensorrt()
     nvshmem()
diff --git a/third_party/xla/MODULE.bazel b/third_party/xla/MODULE.bazel
index c5e23f8c1ba..293c64ee973 100644
--- a/third_party/xla/MODULE.bazel
+++ b/third_party/xla/MODULE.bazel
@@ -11,7 +11,7 @@ bazel_dep(name = "bazel_features", version = "1.43.0")
 bazel_dep(name = "bazel_skylib", version = "1.9.0")
 bazel_dep(name = "boringssl", version = "0.20250818.0")
 bazel_dep(name = "curl", version = "8.22.0")
-bazel_dep(name = "flatbuffers", version = "24.3.25.bcr.1")
+bazel_dep(name = "flatbuffers", version = "25.9.23")
 bazel_dep(name = "google_benchmark", version = "1.9.4", repo_name = "com_google_benchmark")
 bazel_dep(name = "googletest", version = "1.17.0", repo_name = "com_google_googletest")
 bazel_dep(name = "grpc", version = "1.81.0")
diff --git a/third_party/xla/third_party/flatbuffers/BUILD b/third_party/xla/third_party/flatbuffers/BUILD
new file mode 100644
index 00000000000..cc558415f2b
--- /dev/null
+++ b/third_party/xla/third_party/flatbuffers/BUILD
@@ -0,0 +1 @@
+# copybara:uncomment package(default_applicable_licenses = ["//third_party/tensorflow:license"])
diff --git a/third_party/xla/third_party/flatbuffers/flatbuffers.BUILD b/third_party/xla/third_party/flatbuffers/flatbuffers.BUILD
new file mode 100644
index 00000000000..75df1d25f19
--- /dev/null
+++ b/third_party/xla/third_party/flatbuffers/flatbuffers.BUILD
@@ -0,0 +1,151 @@
+# Copyright 2026 The OpenXLA 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.
+# ==============================================================================
+
+load("@rules_cc//cc:defs.bzl", "cc_binary", "cc_library")
+
+package(default_visibility = ["//visibility:public"])
+
+licenses(["notice"])  # Apache 2.0
+
+exports_files(["LICENSE"])
+
+config_setting(
+    name = "platform_freebsd",
+    values = {"cpu": "freebsd"},
+)
+
+config_setting(
+    name = "platform_openbsd",
+    values = {"cpu": "openbsd"},
+)
+
+config_setting(
+    name = "windows",
+    values = {"cpu": "x64_windows"},
+)
+
+# Public flatc library to compile flatbuffer files at runtime.
+cc_library(
+    name = "flatbuffers",
+    hdrs = ["//:public_headers"],
+    linkstatic = 1,
+    strip_include_prefix = "/include",
+    visibility = ["//visibility:public"],
+    deps = ["//src:flatbuffers"],
+)
+
+# Public C++ headers for the Flatbuffers library.
+filegroup(
+    name = "public_headers",
+    srcs = [
+        "include/flatbuffers/allocator.h",
+        "include/flatbuffers/array.h",
+        "include/flatbuffers/base.h",
+        "include/flatbuffers/buffer.h",
+        "include/flatbuffers/buffer_ref.h",
+        "include/flatbuffers/code_generator.h",
+        "include/flatbuffers/code_generators.h",
+        "include/flatbuffers/default_allocator.h",
+        "include/flatbuffers/detached_buffer.h",
+        "include/flatbuffers/file_manager.h",
+        "include/flatbuffers/flatbuffer_builder.h",
+        "include/flatbuffers/flatbuffers.h",
+        "include/flatbuffers/flex_flat_util.h",
+        "include/flatbuffers/flexbuffers.h",
+        "include/flatbuffers/grpc.h",
+        "include/flatbuffers/hash.h",
+        "include/flatbuffers/idl.h",
+        "include/flatbuffers/minireflect.h",
+        "include/flatbuffers/reflection.h",
+        "include/flatbuffers/reflection_generated.h",
+        "include/flatbuffers/registry.h",
+        "include/flatbuffers/stl_emulation.h",
+        "include/flatbuffers/string.h",
+        "include/flatbuffers/struct.h",
+        "include/flatbuffers/table.h",
+        "include/flatbuffers/util.h",
+        "include/flatbuffers/vector.h",
+        "include/flatbuffers/vector_downward.h",
+        "include/flatbuffers/verifier.h",
+    ],
+    visibility = ["//visibility:public"],
+)
+
+# Public flatc compiler library.
+cc_library(
+    name = "flatc_library",
+    linkstatic = 1,
+    visibility = ["//visibility:public"],
+    deps = [
+        "@flatbuffers//src:flatc_library",
+    ],
+)
+
+# Public flatc compiler.
+cc_binary(
+    name = "flatc",
+    linkopts = select({
+        ":platform_freebsd": [
+            "-lm",
+        ],
+        # If Visual Studio 2022 developers facing linking errors,
+        # change the line below as ":windows": ["/DEFAULTLIB:msvcrt.lib"],
+        ":windows": [],
+        "//conditions:default": [
+            "-lm",
+            "-ldl",
+        ],
+    }),
+    visibility = ["//visibility:public"],
+    deps = [
+        "@flatbuffers//src:flatc",
+    ],
+)
+
+filegroup(
+    name = "flatc_headers",
+    srcs = [
+        "include/flatbuffers/flatc.h",
+    ],
+    visibility = ["//visibility:public"],
+)
+
+# Library used by flatbuffer_cc_library rules.
+cc_library(
+    name = "runtime_cc",
+    hdrs = [
+        "include/flatbuffers/allocator.h",
+        "include/flatbuffers/array.h",
+        "include/flatbuffers/base.h",
+        "include/flatbuffers/buffer.h",
+        "include/flatbuffers/buffer_ref.h",
+        "include/flatbuffers/default_allocator.h",
+        "include/flatbuffers/detached_buffer.h",
+        "include/flatbuffers/flatbuffer_builder.h",
+        "include/flatbuffers/flatbuffers.h",
+        "include/flatbuffers/flexbuffers.h",
+        "include/flatbuffers/stl_emulation.h",
+        "include/flatbuffers/string.h",
+        "include/flatbuffers/struct.h",
+        "include/flatbuffers/table.h",
+        "include/flatbuffers/util.h",
+        "include/flatbuffers/vector.h",
+        "include/flatbuffers/vector_downward.h",
+        "include/flatbuffers/verifier.h",
+    ],
+    linkstatic = 1,
+    strip_include_prefix = "/include",
+    visibility = ["//visibility:public"],
+)
diff --git a/third_party/xla/third_party/flatbuffers/workspace.bzl b/third_party/xla/third_party/flatbuffers/workspace.bzl
new file mode 100644
index 00000000000..2e21c905637
--- /dev/null
+++ b/third_party/xla/third_party/flatbuffers/workspace.bzl
@@ -0,0 +1,17 @@
+"""Loads the Flatbuffers library (WORKSPACE builds only; bzlmod builds use the BCR module)."""
+
+load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
+
+_FLATBUFFERS_VERSION = "25.9.23"
+
+# curl -L https://github.com/google/flatbuffers/archive/v<_FLATBUFFERS_VERSION>.tar.gz | shasum -a 256
+_FLATBUFFERS_SHA256 = "9102253214dea6ae10c2ac966ea1ed2155d22202390b532d1dea64935c518ada"
+
+def repo():
+    tf_http_archive(
+        name = "flatbuffers",
+        strip_prefix = "flatbuffers-%s" % _FLATBUFFERS_VERSION,
+        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",
+    )
diff --git a/third_party/xla/workspace2.bzl b/third_party/xla/workspace2.bzl
index 2c4dca151d4..9b0c4fb098b 100644
--- a/third_party/xla/workspace2.bzl
+++ b/third_party/xla/workspace2.bzl
@@ -39,6 +39,7 @@ load("//third_party/dlpack:workspace.bzl", dlpack = "repo")
 load("//third_party/ducc:workspace.bzl", ducc = "repo")
 load("//third_party/eigen3:workspace.bzl", eigen3 = "repo")
 load("//third_party/farmhash:workspace.bzl", farmhash = "repo")
+load("//third_party/flatbuffers:workspace.bzl", flatbuffers = "repo")
 load("//third_party/fmt:workspace.bzl", fmt = "repo")
 load("//third_party/FP16:workspace.bzl", FP16 = "repo")
 load("//third_party/fxdiv:workspace.bzl", fxdiv = "repo")
@@ -109,6 +110,9 @@ def _initialize_third_party():
     ducc()
     eigen3()
     farmhash()
+
+    # Only needed in WORKSPACE builds; bzlmod uses the BCR module.
+    flatbuffers()
     fmt()
     fxdiv()
     riegeli()
diff --git a/third_party/xla/xla/backends/autotuner/backend_config.proto b/third_party/xla/xla/backends/autotuner/backend_config.proto
index d2ae907a3d5..d4bda1d623d 100644
--- a/third_party/xla/xla/backends/autotuner/backend_config.proto
+++ b/third_party/xla/xla/backends/autotuner/backend_config.proto
@@ -36,5 +36,6 @@ message BackendConfig {
     xla.cpu.LlvmKernelOptions llvm_kernel = 5;
     xla.gpu.NativeEmitterBackendConfig native_emitter = 6;
     xla.cpu.YnnFusionOptions ynn_fusion = 7;
+    xla.gpu.TensorIrFusionConfig tensor_ir = 8;
   }
 }
diff --git a/third_party/xla/xla/backends/autotuner/backends.proto b/third_party/xla/xla/backends/autotuner/backends.proto
index dafa1687822..359c53b8819 100644
--- a/third_party/xla/xla/backends/autotuner/backends.proto
+++ b/third_party/xla/xla/backends/autotuner/backends.proto
@@ -41,4 +41,5 @@ enum Backend {
   HIPBLASLT_FISSION = 16;
   YNNPACK = 17;
   BLOCK_LEVEL_EMITTER_CPU = 18;
+  TENSOR_IR = 19;
 }
diff --git a/third_party/xla/xla/backends/gpu/autotuner/BUILD b/third_party/xla/xla/backends/gpu/autotuner/BUILD
index 4fa20a5570a..c614e7ec79e 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/BUILD
+++ b/third_party/xla/xla/backends/gpu/autotuner/BUILD
@@ -146,6 +146,80 @@ xla_test(
     ],
 )
 
+cc_library(
+    name = "tensor_ir",
+    srcs = ["tensor_ir.cc"],
+    hdrs = ["tensor_ir.h"],
+    tags = ["gpu"],
+    deps = [
+        ":gpu_codegen_backend",
+        "//xla:xla_proto_cc",
+        "//xla/backends/autotuner:backends_proto_cc",
+        "//xla/backends/autotuner:codegen_backend",
+        "//xla/backends/gpu/codegen/tensor_ir:hlo_to_tensor_ir",
+        "//xla/backends/gpu/codegen/tensor_ir:support",
+        "//xla/hlo/ir:hlo",
+        "//xla/service:compiler",
+        "//xla/service/gpu:backend_configs_cc",
+        "//xla/service/gpu:gpu_constants",
+        "//xla/service/gpu:ir_emission_utils",
+        "//xla/service/llvm_ir:llvm_util",
+        "//xla/stream_executor/cuda:cuda_compute_capability",
+        "@com_google_absl//absl/container:btree",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@tensor_ir//:NVTensorIRDialect",
+        "@tensor_ir//:NVTensorIRTilingAnalysis",
+        "@tensor_ir//:NVTensorIRTransform",
+        "@tensor_ir//:NVTensorIRUtils",
+    ],
+)
+
+xla_test(
+    name = "tensor_ir_test",
+    srcs = ["tensor_ir_test.cc"],
+    backend_tags = {
+        "b200": ["broken"],
+    },
+    backends = [
+        "a100",
+        "h100",
+        "b200",
+    ],
+    tags = [
+        "cuda-only",
+        "no_mac",
+    ],
+    use_legacy_runtime = True,
+    deps = [
+        ":tensor_ir",
+        "//xla:autotuning_proto_cc",
+        "//xla:xla_proto_cc",
+        "//xla/backends/autotuner:backend_config_proto_cc",
+        "//xla/backends/autotuner:codegen_backend",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+        "//xla/service:compiler",
+        "//xla/service:platform_util",
+        "//xla/service/gpu:backend_configs_cc",
+        "//xla/service/gpu:ir_emission_utils",
+        "//xla/stream_executor:platform",
+        "//xla/stream_executor:stream_executor_h",
+        "//xla/stream_executor/cuda:cuda_compute_capability",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "cublaslt",
     srcs = ["cublaslt.cc"],
@@ -553,6 +627,7 @@ cc_library(
         ":factory",
         ":fission_backend",
         ":native_emitter",
+        ":tensor_ir",
         ":triton",
         "//xla:xla_proto_cc",
         "//xla/backends/autotuner:backends_proto_cc",
@@ -1151,9 +1226,7 @@ xla_test(
         "//xla/stream_executor:platform_manager",
         "//xla/stream_executor:stream_executor_address_allocator",
         "//xla/stream_executor:stream_executor_h",
-        "//xla/stream_executor:stream_executor_memory_allocator",
         "//xla/stream_executor/platform:platform_object_registry",
-        "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
diff --git a/third_party/xla/xla/backends/gpu/autotuner/factory_cuda.cc b/third_party/xla/xla/backends/gpu/autotuner/factory_cuda.cc
index e181373ff5f..e9861cf63d8 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/factory_cuda.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/factory_cuda.cc
@@ -31,6 +31,7 @@ limitations under the License.
 #include "xla/backends/gpu/autotuner/factory.h"
 #include "xla/backends/gpu/autotuner/fission_backend.h"
 #include "xla/backends/gpu/autotuner/native_emitter.h"
+#include "xla/backends/gpu/autotuner/tensor_ir.h"
 #include "xla/backends/gpu/autotuner/triton.h"
 #include "xla/backends/gpu/transforms/dot_algorithm_rewriter.h"
 #include "xla/backends/gpu/transforms/gemm_rewriter.h"
@@ -102,6 +103,8 @@ std::vector<std::unique_ptr<CodegenBackend>> GetCodegenBackendsForCuda(
   backends.push_back(std::make_unique<BlockLevelEmitterBackend>(
       debug_options, compiler, shape_size_fn, target_config, thread_pool,
       mlir_context_pool));
+  backends.push_back(std::make_unique<TensorIrBackend>(debug_options, compiler,
+                                                       target_config));
 
   if (!backend_allowlist.empty()) {
     backends.erase(
diff --git a/third_party/xla/xla/backends/gpu/autotuner/factory_test.cc b/third_party/xla/xla/backends/gpu/autotuner/factory_test.cc
index 0b2c889100b..e20551f7c92 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/factory_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/factory_test.cc
@@ -132,7 +132,7 @@ TEST_P(FactoryTest, GetCodegenBackends) {
 INSTANTIATE_TEST_SUITE_P(
     All, FactoryTest,
     ::testing::Values(
-        FactoryTestParams{{}, 6, /*run_on_cuda=*/true, /*run_on_rocm=*/false},
+        FactoryTestParams{{}, 7, /*run_on_cuda=*/true, /*run_on_rocm=*/false},
         FactoryTestParams{{}, 6, /*run_on_cuda=*/false, /*run_on_rocm=*/true},
         FactoryTestParams{{Backend::TRITON}, 1},
         FactoryTestParams{{Backend::TRITON, Backend::CUBLASLT},
diff --git a/third_party/xla/xla/backends/gpu/autotuner/tensor_ir.cc b/third_party/xla/xla/backends/gpu/autotuner/tensor_ir.cc
new file mode 100644
index 00000000000..c528f37393b
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/autotuner/tensor_ir.cc
@@ -0,0 +1,267 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/gpu/autotuner/tensor_ir.h"
+
+#include <algorithm>
+#include <cstdint>
+#include <iterator>
+#include <memory>
+#include <utility>
+#include <vector>
+
+#include "tensor_ir/Analysis/Tiling/Arch.h"
+#include "tensor_ir/Analysis/Tiling/Enumerate.h"
+#include "tensor_ir/Analysis/Tiling/Evaluate.h"
+#include "tensor_ir/Dialect/TensorIR.h"
+#include "tensor_ir/Transform/Passes.h"
+#include "tensor_ir/Utils/ComputeCapability.h"
+#include "absl/container/btree_map.h"
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/LogicalResult.h"
+#include "llvm/Support/MathExtras.h"
+#include "mlir/Dialect/Arith/IR/ArithDialect.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OwningOpRef.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Support/LogicalResult.h"
+#include "xla/backends/autotuner/codegen_backend.h"
+#include "xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir.h"
+#include "xla/backends/gpu/codegen/tensor_ir/support.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_opcode.h"
+#include "xla/service/gpu/backend_configs.pb.h"
+#include "xla/service/gpu/gpu_constants.h"
+#include "xla/service/gpu/ir_emission_utils.h"
+#include "xla/service/llvm_ir/llvm_util.h"
+#include "xla/stream_executor/cuda/cuda_compute_capability.h"
+
+namespace xla {
+namespace gpu {
+
+namespace {
+
+using ::mlir::nv_tensor_ir::tiling_analysis::TilingConfig;
+using ::mlir::nv_tensor_ir::tiling_analysis::TilingEvaluation;
+
+std::unique_ptr<BackendConfig> Pack(const TilingConfig& tiling_config) {
+  auto config = std::make_unique<BackendConfig>();
+  TensorIrFusionConfig* tensor_ir_config = config->mutable_tensor_ir();
+  for (int32_t dim : tiling_config.getTileShape()) {
+    tensor_ir_config->add_tile_size(dim);
+  }
+  tensor_ir_config->set_reduction_tile_size(
+      tiling_config.getReductionTileSize());
+  return config;
+}
+
+// Enumerates up to `num_tilings_to_generate` candidate tilings for `graph`
+// and returns up to `num_tilings_to_select` of them that fit the register
+// budget. The selection is spread across tile sizes (ascending) and each
+// tile-size group is ranked by the tiling evaluator's cost model (best first),
+// so the first returned tiling is the best one of the smallest tile size.
+absl::StatusOr<llvm::SmallVector<TilingConfig>> SelectBestTilings(
+    mlir::nv_tensor_ir::GraphOp graph, int64_t num_tilings_to_generate,
+    int64_t num_tilings_to_select,
+    const mlir::nv_tensor_ir::tiling_analysis::ArchInfo& arch_info) {
+  mlir::FailureOr<mlir::nv_tensor_ir::tiling_analysis::TilingEvaluator>
+      evaluator = mlir::nv_tensor_ir::tiling_analysis::TilingEvaluator::create(
+          graph, arch_info);
+  if (llvm::failed(evaluator)) {
+    return absl::InvalidArgumentError(
+        "TensorIrBackend: failed to create the tiling evaluator");
+  }
+  llvm::SmallVector<TilingConfig> tilings =
+      mlir::nv_tensor_ir::tiling_analysis::enumerateTilings(
+          *evaluator, num_tilings_to_generate);
+
+  // Use a storage size threshold to filter out tilings that are too large.
+  constexpr int64_t kStorageSizeThreshold =
+      ::mlir::nv_tensor_ir::tiling_analysis::kMaxRegistersPerSM *
+      ::mlir::nv_tensor_ir::tiling_analysis::kRegisterSize;
+
+  // Evaluate each tiling and keep those that fit in the register budget,
+  // grouped by the product of the tile shape. The next step spreads the
+  // selection across these groups.
+  using EvaluatedTiling = std::pair<TilingConfig, TilingEvaluation>;
+  absl::btree_map<int64_t, llvm::SmallVector<EvaluatedTiling>>
+      tilings_grouped_by_tile_size;
+  for (const TilingConfig& tiling_config : tilings) {
+    auto eval = evaluator->evaluate(tiling_config);
+    if (llvm::succeeded(eval) &&
+        eval->registerStorageBytes < kStorageSizeThreshold) {
+      int64_t tile_size =
+          llvm::product_of(tiling_config.getTileShape(), int64_t{1});
+      tilings_grouped_by_tile_size[tile_size].emplace_back(tiling_config,
+                                                           *eval);
+    }
+  }
+
+  // Pick `num_tilings_to_select` tilings spread across tile sizes. Each group
+  // is sorted by cost (best first) and takes an equal share of the remaining
+  // budget, rounded up. A group smaller than its share leaves the unused slots
+  // for later groups.
+  llvm::SmallVector<EvaluatedTiling> selected;
+  int64_t budget = num_tilings_to_select;
+  int64_t remaining_groups = tilings_grouped_by_tile_size.size();
+  for (auto& [tile_size, group] : tilings_grouped_by_tile_size) {
+    if (budget == 0) {
+      break;
+    }
+    llvm::stable_sort(group, [](const auto& a, const auto& b) {
+      return a.second < b.second;
+    });
+    int64_t take = std::min<int64_t>(
+        group.size(), llvm::divideCeil(budget, remaining_groups));
+    budget -= take;
+    --remaining_groups;
+    selected.append(std::make_move_iterator(group.begin()),
+                    std::make_move_iterator(group.begin() + take));
+  }
+
+  // Return the best tilings.
+  llvm::SmallVector<TilingConfig> result;
+  result.reserve(selected.size());
+  for (auto& [tiling_config, unused_eval] : selected) {
+    result.push_back(std::move(tiling_config));
+  }
+  return result;
+}
+
+}  // namespace
+
+absl::StatusOr<std::vector<std::unique_ptr<BackendConfig>>>
+TensorIrBackend::GetSupportedConfigs(const HloInstruction& instr) {
+  if (!IsSupported(instr)) {
+    return std::vector<std::unique_ptr<BackendConfig>>();
+  }
+  const auto* fusion = Cast<HloFusionInstruction>(&instr);
+  const HloComputation& computation = *fusion->fused_instructions_computation();
+
+  mlir::MLIRContext context;
+  context.loadDialect<mlir::nv_tensor_ir::TensorIRDialect,
+                      mlir::arith::ArithDialect>();
+  mlir::OwningOpRef<mlir::ModuleOp> module =
+      llvm_ir::CreateMlirModuleOp(mlir::UnknownLoc::get(&context));
+  ABSL_ASSIGN_OR_RETURN(
+      mlir::nv_tensor_ir::GraphOp graph_op,
+      tensor_ir::ImportAndLegalizeComputation(computation, *module));
+
+  // `enumerateTilings` reads the "iteration_space" attribute that the
+  // layout-propagation passes attach to the graph; run them first.
+  mlir::PassManager pass_manager(&context);
+  pass_manager.addNestedPass<mlir::nv_tensor_ir::GraphOp>(
+      mlir::nv_tensor_ir::createLayoutPropagationAnnotationPass());
+  pass_manager.addNestedPass<mlir::nv_tensor_ir::GraphOp>(
+      mlir::nv_tensor_ir::createLayoutPropagationNormalizationPass());
+  if (llvm::failed(pass_manager.run(*module))) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("TensorIrBackend: layout propagation failed for "
+                     "fusion: ",
+                     instr.ToString()));
+  }
+
+  // Get the compute capability and target-specific architecture information.
+  const auto* cuda_cc = target_config()
+                            .device_description.gpu_compute_capability()
+                            .cuda_compute_capability();
+  if (cuda_cc == nullptr) {
+    return absl::InvalidArgumentError(
+        "TensorIrBackend: requires a CUDA compute capability");
+  }
+  mlir::FailureOr<mlir::nv_tensor_ir::SmTarget> sm_target =
+      mlir::nv_tensor_ir::SmTarget::fromCc(cuda_cc->major * 10 +
+                                           cuda_cc->minor);
+  if (llvm::failed(sm_target)) {
+    return absl::InvalidArgumentError(
+        "TensorIrBackend: unsupported compute capability");
+  }
+  const mlir::nv_tensor_ir::tiling_analysis::ArchInfo arch_info =
+      mlir::nv_tensor_ir::tiling_analysis::getArchInfo(
+          sm_target->getComputeCapability(),
+          target_config().device_description.core_count());
+
+  // Select the best tilings.
+  ABSL_ASSIGN_OR_RETURN(llvm::SmallVector<TilingConfig> tilings,
+                   SelectBestTilings(graph_op, /*num_tilings_to_generate=*/1000,
+                                     /*num_tilings_to_select=*/10, arch_info));
+
+  std::vector<std::unique_ptr<BackendConfig>> configs;
+  configs.reserve(tilings.size());
+  for (const TilingConfig& tiling_config : tilings) {
+    configs.push_back(Pack(tiling_config));
+  }
+  return configs;
+}
+
+absl::StatusOr<std::unique_ptr<BackendConfig>>
+TensorIrBackend::GetDefaultConfig(const HloInstruction& instr) {
+  ABSL_ASSIGN_OR_RETURN(auto configs, GetSupportedConfigs(instr));
+  if (configs.empty()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("TensorIrBackend: no supported tiling configs for "
+                     "instruction: ",
+                     instr.ToString()));
+  }
+  return std::move(configs.front());
+}
+
+absl::Status TensorIrBackend::ApplyConfig(HloInstruction& instr,
+                                          const BackendConfig& config) {
+  if (!IsSupported(instr)) {
+    return absl::InvalidArgumentError(
+        "TensorIrBackend does not support this instruction.");
+  }
+  if (!config.has_tensor_ir()) {
+    return absl::InvalidArgumentError("Expected TensorIrFusionConfig.");
+  }
+  ABSL_ASSIGN_OR_RETURN(GpuBackendConfig gpu_backend_config,
+                   instr.backend_config<GpuBackendConfig>());
+  FusionBackendConfig& backend_config =
+      *gpu_backend_config.mutable_fusion_backend_config();
+  backend_config.set_kind(kTensorIrFusionKind);
+  *backend_config.mutable_tensor_ir_fusion_config() = config.tensor_ir();
+  ABSL_RETURN_IF_ERROR(instr.set_backend_config(std::move(gpu_backend_config)));
+  instr.set_fusion_kind(HloInstruction::FusionKind::kCustom);
+  return absl::OkStatus();
+}
+
+bool TensorIrBackend::IsSupported(const HloInstruction& instr) {
+  if (instr.opcode() != HloOpcode::kFusion) {
+    return false;
+  }
+  if (!tensor_ir::IsSupportedComputeCapability(
+           target_config().device_description.gpu_compute_capability())
+           .IsAllowed()) {
+    return false;
+  }
+  const auto* fusion = Cast<HloFusionInstruction>(&instr);
+  return tensor_ir::IsSupportedFusionComputation(
+             *fusion->fused_instructions_computation())
+      .IsAllowed();
+}
+
+}  // namespace gpu
+}  // namespace xla
diff --git a/third_party/xla/xla/backends/gpu/autotuner/tensor_ir.h b/third_party/xla/xla/backends/gpu/autotuner/tensor_ir.h
new file mode 100644
index 00000000000..e4ba7f5137c
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/autotuner/tensor_ir.h
@@ -0,0 +1,65 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_BACKENDS_GPU_AUTOTUNER_TENSOR_IR_H_
+#define XLA_BACKENDS_GPU_AUTOTUNER_TENSOR_IR_H_
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "xla/backends/autotuner/backends.pb.h"
+#include "xla/backends/autotuner/codegen_backend.h"
+#include "xla/backends/gpu/autotuner/gpu_codegen_backend.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/service/compiler.h"
+#include "xla/xla.pb.h"
+
+namespace xla {
+namespace gpu {
+
+// Codegen backend for fusions emitted through the TensorIR/CudaTile
+// pipeline (see xla/backends/gpu/codegen/tensor_ir).
+//
+// This backend enumerates candidate tiling configurations for a fusion
+// (using the TensorIR tiling-analysis heuristics) and applies the chosen
+// one as a TensorIrFusionConfig on the fusion's FusionBackendConfig.
+class TensorIrBackend : public GpuCodegenBackend {
+ public:
+  TensorIrBackend(const DebugOptions* debug_options, Compiler* compiler,
+                  const Compiler::GpuTargetConfig* target_config)
+      : GpuCodegenBackend(autotuner::Backend::TENSOR_IR, debug_options,
+                          compiler, target_config) {}
+
+  absl::StatusOr<std::vector<std::unique_ptr<BackendConfig>>>
+  GetSupportedConfigs(const HloInstruction& instr) override;
+
+  absl::StatusOr<std::unique_ptr<BackendConfig>> GetDefaultConfig(
+      const HloInstruction& instr) override;
+
+  absl::Status ApplyConfig(HloInstruction& instr,
+                           const BackendConfig& config) override;
+
+  bool IsSupported(const HloInstruction& instr) override;
+
+  std::string version() const override { return "1"; }
+};
+
+}  // namespace gpu
+}  // namespace xla
+
+#endif  // XLA_BACKENDS_GPU_AUTOTUNER_TENSOR_IR_H_
diff --git a/third_party/xla/xla/backends/gpu/autotuner/tensor_ir_test.cc b/third_party/xla/xla/backends/gpu/autotuner/tensor_ir_test.cc
new file mode 100644
index 00000000000..a20c3b257cc
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/autotuner/tensor_ir_test.cc
@@ -0,0 +1,285 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/gpu/autotuner/tensor_ir.h"
+
+#include <memory>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "absl/status/statusor.h"
+#include "xla/autotuning.pb.h"
+#include "xla/backends/autotuner/backend_config.pb.h"
+#include "xla/backends/autotuner/codegen_backend.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
+#include "xla/service/compiler.h"
+#include "xla/service/gpu/backend_configs.pb.h"
+#include "xla/service/gpu/ir_emission_utils.h"
+#include "xla/service/platform_util.h"
+#include "xla/stream_executor/cuda/cuda_compute_capability.h"
+#include "xla/stream_executor/platform.h"
+#include "xla/stream_executor/stream_executor.h"
+#include "xla/xla.pb.h"
+
+namespace xla::gpu {
+namespace {
+
+using ::absl_testing::IsOk;
+using ::absl_testing::StatusIs;
+using ::testing::ElementsAre;
+using ::testing::IsEmpty;
+using ::testing::Not;
+
+constexpr char kElementwiseFusionHlo[] = R"(
+HloModule m
+
+%fused_add (p0: f32[32,16], p1: f32[32,16]) -> f32[32,16] {
+  %p0 = f32[32,16]{1,0} parameter(0)
+  %p1 = f32[32,16]{1,0} parameter(1)
+  ROOT %add = f32[32,16]{1,0} add(%p0, %p1)
+}
+
+ENTRY %entry_computation (p0: f32[32,16], p1: f32[32,16]) -> f32[32,16] {
+  %p0 = f32[32,16]{1,0} parameter(0)
+  %p1 = f32[32,16]{1,0} parameter(1)
+  ROOT %fusion = f32[32,16]{1,0} fusion(%p0, %p1), kind=kLoop, calls=%fused_add
+})";
+
+constexpr char kReductionFusionHlo[] = R"(
+HloModule m
+
+%func (lhs: f32[], rhs: f32[]) -> f32[] {
+  %lhs = f32[] parameter(0)
+  %rhs = f32[] parameter(1)
+  ROOT %sum = f32[] add(%lhs, %rhs)
+}
+
+%fused_reduce (p0: f32[32,4096]) -> f32[32] {
+  %p0 = f32[32,4096]{1,0} parameter(0)
+  %c0 = f32[] constant(0)
+  ROOT %reduce = f32[32]{0} reduce(%p0, %c0), dimensions={1}, to_apply=%func
+}
+
+ENTRY %entry_computation (p0: f32[32,4096]) -> f32[32] {
+  %p0 = f32[32,4096]{1,0} parameter(0)
+  ROOT %fusion = f32[32]{0} fusion(%p0), kind=kInput, calls=%fused_reduce
+})";
+
+// Integer fusions exercise the signedness bridge: StableHLO uses signless
+// `i32`, while nv_tensor_ir requires signed `si32`.
+constexpr char kElementwiseS32FusionHlo[] = R"(
+HloModule m
+
+%fused_mul (p0: s32[32,16], p1: s32[32,16]) -> s32[32,16] {
+  %p0 = s32[32,16]{1,0} parameter(0)
+  %p1 = s32[32,16]{1,0} parameter(1)
+  ROOT %mul = s32[32,16]{1,0} multiply(%p0, %p1)
+}
+
+ENTRY %entry_computation (p0: s32[32,16], p1: s32[32,16]) -> s32[32,16] {
+  %p0 = s32[32,16]{1,0} parameter(0)
+  %p1 = s32[32,16]{1,0} parameter(1)
+  ROOT %fusion = s32[32,16]{1,0} fusion(%p0, %p1), kind=kLoop, calls=%fused_mul
+})";
+
+// `pad` isn't part of the opcode set the TensorIR emitter supports.
+constexpr char kUnsupportedOpFusionHlo[] = R"(
+HloModule m
+
+%fused_pad (p0: f32[4]) -> f32[6] {
+  %p0 = f32[4]{0} parameter(0)
+  %c0 = f32[] constant(0)
+  ROOT %pad = f32[6]{0} pad(%p0, %c0), padding=1_1
+}
+
+ENTRY %entry_computation (p0: f32[4]) -> f32[6] {
+  %p0 = f32[4]{0} parameter(0)
+  ROOT %fusion = f32[6]{0} fusion(%p0), kind=kLoop, calls=%fused_pad
+})";
+
+class TensorIrBackendTest : public HloHardwareIndependentTestBase {
+ protected:
+  TensorIrBackendTest()
+      : platform_(PlatformUtil::GetDefaultPlatform().value()),
+        stream_executor_(platform_->ExecutorForDevice(0).value()),
+        target_config_(stream_executor_),
+        compiler_(Compiler::GetForPlatform(platform_->id()).value()),
+        backend_(&debug_options_, compiler_.get(), &target_config_) {
+    // The backend only supports Hopper and newer. Pin the target so that these
+    // tests do not depend on which GPU the test machine happens to have;
+    // `IsSupportedReturnsFalseForPreHopper` covers the other side.
+    target_config_.device_description.set_cuda_compute_capability(
+        se::CudaComputeCapability::Hopper());
+  }
+
+  DebugOptions debug_options_;
+  se::Platform* platform_;
+  se::StreamExecutor* stream_executor_;
+  Compiler::GpuTargetConfig target_config_;
+  std::unique_ptr<Compiler> compiler_;
+  TensorIrBackend backend_;
+};
+
+TEST_F(TensorIrBackendTest, IsSupportedReturnsTrueForSupportedFusion) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kElementwiseFusionHlo));
+  EXPECT_TRUE(
+      backend_.IsSupported(*module->entry_computation()->root_instruction()));
+}
+
+TEST_F(TensorIrBackendTest, IsSupportedReturnsFalseForNonFusion) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(R"(
+HloModule m
+ENTRY %entry_computation (p0: f32[4], p1: f32[4]) -> f32[4] {
+  %p0 = f32[4]{0} parameter(0)
+  %p1 = f32[4]{0} parameter(1)
+  ROOT %add = f32[4]{0} add(%p0, %p1)
+})"));
+  EXPECT_FALSE(
+      backend_.IsSupported(*module->entry_computation()->root_instruction()));
+}
+
+TEST_F(TensorIrBackendTest, IsSupportedReturnsFalseForPreHopper) {
+  target_config_.device_description.set_cuda_compute_capability(
+      se::CudaComputeCapability::Ampere());
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kElementwiseFusionHlo));
+  EXPECT_FALSE(
+      backend_.IsSupported(*module->entry_computation()->root_instruction()));
+}
+
+TEST_F(TensorIrBackendTest, IsSupportedReturnsFalseForUnsupportedOpcode) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kUnsupportedOpFusionHlo));
+  EXPECT_FALSE(
+      backend_.IsSupported(*module->entry_computation()->root_instruction()));
+}
+
+TEST_F(TensorIrBackendTest, GetDefaultConfigReturnsTensorIrConfig) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kElementwiseFusionHlo));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<BackendConfig> config,
+                       backend_.GetDefaultConfig(
+                           *module->entry_computation()->root_instruction()));
+
+  // GetDefaultConfig returns the top-ranked config from GetSupportedConfigs.
+  ASSERT_TRUE(config->has_tensor_ir());
+  EXPECT_THAT(config->tensor_ir().tile_size(), Not(IsEmpty()));
+}
+
+TEST_F(TensorIrBackendTest, GetDefaultConfigFailsForUnsupportedInstruction) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kUnsupportedOpFusionHlo));
+  EXPECT_THAT(backend_.GetDefaultConfig(
+                  *module->entry_computation()->root_instruction()),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+TEST_F(TensorIrBackendTest, GetSupportedConfigsReturnsRankedTilings) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kReductionFusionHlo));
+  ASSERT_OK_AND_ASSIGN(std::vector<std::unique_ptr<BackendConfig>> configs,
+                       backend_.GetSupportedConfigs(
+                           *module->entry_computation()->root_instruction()));
+
+  // GetSupportedConfigs asks SelectBestTilings for at most 10 configs.
+  EXPECT_THAT(configs, Not(IsEmpty()));
+  EXPECT_LE(configs.size(), 10);
+  for (const std::unique_ptr<BackendConfig>& config : configs) {
+    ASSERT_TRUE(config->has_tensor_ir());
+    EXPECT_THAT(config->tensor_ir().tile_size(), Not(IsEmpty()));
+  }
+}
+
+// If the signless `i32` to signed `si32` bridge regressed, legalization would
+// fail and `GetSupportedConfigs` would return an error instead of configs.
+TEST_F(TensorIrBackendTest, GetSupportedConfigsHandlesSignedIntegers) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kElementwiseS32FusionHlo));
+  ASSERT_OK_AND_ASSIGN(std::vector<std::unique_ptr<BackendConfig>> configs,
+                       backend_.GetSupportedConfigs(
+                           *module->entry_computation()->root_instruction()));
+
+  EXPECT_THAT(configs, Not(IsEmpty()));
+  for (const std::unique_ptr<BackendConfig>& config : configs) {
+    ASSERT_TRUE(config->has_tensor_ir());
+    EXPECT_THAT(config->tensor_ir().tile_size(), Not(IsEmpty()));
+  }
+}
+
+TEST_F(TensorIrBackendTest, GetSupportedConfigsReturnsEmptyForUnsupported) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kUnsupportedOpFusionHlo));
+  ASSERT_OK_AND_ASSIGN(std::vector<std::unique_ptr<BackendConfig>> configs,
+                       backend_.GetSupportedConfigs(
+                           *module->entry_computation()->root_instruction()));
+  EXPECT_THAT(configs, IsEmpty());
+}
+
+TEST_F(TensorIrBackendTest, ApplyConfigSetsFusionBackendConfig) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kElementwiseFusionHlo));
+  HloInstruction* fusion = module->entry_computation()->root_instruction();
+
+  BackendConfig config;
+  config.mutable_tensor_ir()->add_tile_size(8);
+  config.mutable_tensor_ir()->add_tile_size(16);
+  config.mutable_tensor_ir()->set_reduction_tile_size(64);
+
+  ASSERT_THAT(backend_.ApplyConfig(*fusion, config), IsOk());
+
+  ASSERT_OK_AND_ASSIGN(GpuBackendConfig gpu_backend_config,
+                       fusion->backend_config<GpuBackendConfig>());
+  const FusionBackendConfig& backend_config =
+      gpu_backend_config.fusion_backend_config();
+  EXPECT_EQ(backend_config.kind(), kTensorIrFusionKind);
+  EXPECT_THAT(backend_config.tensor_ir_fusion_config().tile_size(),
+              ElementsAre(8, 16));
+  EXPECT_EQ(backend_config.tensor_ir_fusion_config().reduction_tile_size(), 64);
+  EXPECT_EQ(fusion->fusion_kind(), HloInstruction::FusionKind::kCustom);
+}
+
+TEST_F(TensorIrBackendTest, ApplyConfigFailsForWrongConfigType) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kElementwiseFusionHlo));
+  HloInstruction* fusion = module->entry_computation()->root_instruction();
+
+  BackendConfig config;
+  config.mutable_gemm()->set_algorithm(1);
+
+  EXPECT_THAT(backend_.ApplyConfig(*fusion, config),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+TEST_F(TensorIrBackendTest, ApplyConfigFailsForUnsupportedInstruction) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kUnsupportedOpFusionHlo));
+  HloInstruction* fusion = module->entry_computation()->root_instruction();
+
+  BackendConfig config;
+  config.mutable_tensor_ir()->add_tile_size(8);
+
+  EXPECT_THAT(backend_.ApplyConfig(*fusion, config),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+}  // namespace
+}  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/codegen/BUILD b/third_party/xla/xla/backends/gpu/codegen/BUILD
index d1d5ba1dd25..aeee005c995 100644
--- a/third_party/xla/xla/backends/gpu/codegen/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/BUILD
@@ -16,6 +16,10 @@
 load("@rules_cc//cc:cc_library.bzl", "cc_library")
 load("//xla:xla.default.bzl", "xla_cc_test")
 load("//xla/tests:build_defs.bzl", "xla_test")
+load(
+    "//xla/tsl/platform/default:cuda_build_defs.bzl",
+    "if_cuda_is_configured",
+)
 
 package(
     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
@@ -240,7 +244,10 @@ cc_library(
         "//xla/service/gpu:hlo_fusion_analysis",
         "//xla/service/gpu:ir_emission_utils",
         "@com_google_absl//absl/status:statusor",
-    ],
+    ] + if_cuda_is_configured(
+        ["//xla/backends/gpu/codegen/tensor_ir:fusion"],
+        no_cuda = ["//xla/backends/gpu/codegen/tensor_ir:fusion_stub"],
+    ),
 )
 
 cc_library(
diff --git a/third_party/xla/xla/backends/gpu/codegen/fusions.cc b/third_party/xla/xla/backends/gpu/codegen/fusions.cc
index 48ce3af58af..7c82e266121 100644
--- a/third_party/xla/xla/backends/gpu/codegen/fusions.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/fusions.cc
@@ -30,6 +30,7 @@ limitations under the License.
 #include "xla/backends/gpu/codegen/emitters/transpose.h"
 #include "xla/backends/gpu/codegen/fusion_emitter.h"
 #include "xla/backends/gpu/codegen/sort.h"
+#include "xla/backends/gpu/codegen/tensor_ir/fusion.h"
 #include "xla/backends/gpu/codegen/triton/fusion.h"
 #include "xla/codegen/ir_emission_utils.h"
 #include "xla/hlo/ir/hlo_instruction.h"
@@ -106,6 +107,8 @@ std::unique_ptr<FusionInterface> GetFusionEmitter(
       return std::make_unique<TritonFusion>(analysis);
     case HloFusionAnalysis::EmitterFusionKind::kCuDnn:
       return std::make_unique<CuDnnFusion>(analysis);
+    case HloFusionAnalysis::EmitterFusionKind::kTensorIr:
+      return std::make_unique<TensorIrFusion>();
   }
 }
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD
index b340d8ccf06..d8230292686 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD
@@ -1,4 +1,6 @@
 load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//xla:xla.default.bzl", "xla_cc_test")
+load("//xla/tsl:tsl.bzl", "if_google")
 
 package(
     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
@@ -13,36 +15,6 @@ package_group(
     ],
 )
 
-cc_library(
-    name = "conversion",
-    srcs = ["conversion.cc"],
-    hdrs = ["conversion.h"],
-    deps = [
-        "//xla:comparison_util",
-        "//xla:literal",
-        "//xla:literal_util",
-        "//xla:shape_util",
-        "//xla:util",
-        "//xla:xla_data_proto_cc",
-        "//xla/hlo/ir:hlo",
-        "//xla/service/llvm_ir:llvm_util",
-        "@com_google_absl//absl/algorithm:container",
-        "@com_google_absl//absl/cleanup",
-        "@com_google_absl//absl/log",
-        "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:status_macros",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/types:span",
-        "@llvm-project//llvm:Support",
-        "@llvm-project//mlir:ArithDialect",
-        "@llvm-project//mlir:IR",
-        "@llvm-project//mlir:SideEffectInterfaces",
-        "@llvm-project//mlir:Support",
-        "@tensor_ir//:NVTensorIRDialect",
-    ],
-)
-
 cc_library(
     name = "support",
     srcs = ["support.cc"],
@@ -52,6 +24,221 @@ cc_library(
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/service:decision",
+        "//xla/stream_executor:device_description",
+        "//xla/stream_executor/cuda:cuda_compute_capability",
         "@com_google_absl//absl/strings",
     ],
 )
+
+xla_cc_test(
+    name = "support_test",
+    srcs = ["support_test.cc"],
+    deps = [
+        ":support",
+        "//xla/stream_executor:device_description",
+        "//xla/stream_executor/cuda:cuda_compute_capability",
+        "//xla/stream_executor/rocm:rocm_compute_capability",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+cc_library(
+    name = "hlo_to_tensor_ir",
+    srcs = ["hlo_to_tensor_ir.cc"],
+    hdrs = ["hlo_to_tensor_ir.h"],
+    deps = [
+        ":layout_utils",
+        "//xla/backends/gpu/codegen/tensor_ir/transforms:passes",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/translate/hlo_to_mhlo:hlo_function_importer",
+        "//xla/mlir/utils:error_util",
+        "//xla/service/llvm_ir:llvm_util",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@stablehlo//:stablehlo_ops",
+        "@tensor_ir//:NVTensorIRDialect",
+    ],
+)
+
+cc_library(
+    name = "compilation_pipeline",
+    srcs = ["compilation_pipeline.cc"],
+    hdrs = ["compilation_pipeline.h"],
+    # tileiras, ptxas and libnvvm staged from the CUDA 13.3+ toolkit;
+    # FindCudaExecutable("tileiras") picks them up from runfiles.
+    data = if_google(["@local_config_cuda//cuda:runtime_tileiras"]),
+    deps = [
+        "//xla/stream_executor:device_description",
+        "//xla/stream_executor/cuda:cuda_compute_capability",
+        "//xla/stream_executor/cuda:subprocess_compilation",
+        "@com_google_absl//absl/base",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
+        "@tensor_ir//:NVTensorIRConversionOptions",
+        "@tensor_ir//:NVTensorIRCudaTilePipelines",
+        "@tensor_ir//:NVTensorIRTileIRAssembly",
+        "@tsl//tsl/platform:path",
+    ],
+)
+
+cc_library(
+    name = "fusion",
+    srcs = ["fusion.cc"],
+    hdrs = ["fusion.h"],
+    # The tensor_ir compiler and runtime libraries include cuda.h; builds
+    # without CUDA link :fusion_stub instead.
+    tags = [
+        "cuda-only",
+        "gpu",
+    ],
+    visibility = [
+        "//xla/backends/gpu/codegen:__subpackages__",
+        "//xla/service/gpu:__subpackages__",
+    ],
+    deps = [
+        ":compilation_pipeline",
+        ":custom_kernel",
+        ":hlo_to_tensor_ir",
+        ":layout_utils",
+        ":support",
+        "//xla:future",
+        "//xla/backends/gpu/codegen:fusion_emitter",
+        "//xla/backends/gpu/codegen:kernel_compiler",
+        "//xla/backends/gpu/codegen/kernels:custom_kernel",
+        "//xla/backends/gpu/codegen/kernels:ptx_custom_kernel",
+        "//xla/backends/gpu/runtime:custom_kernel_thunk",
+        "//xla/backends/gpu/runtime:thunk",
+        "//xla/codegen/emitters:kernel_arguments",
+        "//xla/hlo/ir:hlo",
+        "//xla/mlir/utils:error_util",
+        "//xla/service/gpu:backend_configs_cc",
+        "//xla/service/gpu:gpu_constants",
+        "//xla/service/gpu:ir_emission_utils",
+        "//xla/service/gpu:ir_emitter_context",
+        "//xla/service/gpu:kernel_reuse_cache",
+        "//xla/service/llvm_ir:llvm_util",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@cuda_tile//:CudaTileDialect",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@tensor_ir//:NVTensorIRCompiler",
+        "@tensor_ir//:NVTensorIRCompilerOptions",
+        "@tensor_ir//:NVTensorIRConversionOptions",
+        "@tensor_ir//:NVTensorIRDialect",
+        "@tensor_ir//:NVTensorIRUtils",
+    ],
+)
+
+cc_library(
+    name = "layout_utils",
+    srcs = ["layout_utils.cc"],
+    hdrs = ["layout_utils.h"],
+    deps = [
+        "//xla:shape_util",
+        "//xla/codegen/emitters:kernel_arguments",
+        "//xla/hlo/ir:hlo",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
+        "@llvm-project//mlir:IR",
+        "@tensor_ir//:NVTensorIRDialect",
+    ],
+)
+
+cc_library(
+    name = "custom_kernel",
+    srcs = ["custom_kernel.cc"],
+    hdrs = ["custom_kernel.h"],
+    # The tensor_ir runtime library includes cuda.h.
+    tags = [
+        "cuda-only",
+        "gpu",
+    ],
+    deps = [
+        "//xla/service/gpu:kernel_reuse_cache",
+        "//xla/service/gpu:launch_dimensions",
+        "//xla/stream_executor:launch_dim",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Support",
+        "@tensor_ir//:NVTensorIRRuntime",
+    ],
+)
+
+cc_library(
+    name = "fusion_stub",
+    srcs = ["fusion_stub.cc"],
+    hdrs = ["fusion.h"],
+    visibility = [
+        "//xla/backends/gpu/codegen:__subpackages__",
+    ],
+    deps = [
+        "//xla/backends/gpu/codegen:fusion_emitter",
+        "//xla/backends/gpu/runtime:thunk",
+        "//xla/hlo/ir:hlo",
+        "//xla/service/gpu:ir_emitter_context",
+        "@com_google_absl//absl/status",
+    ],
+)
+
+xla_cc_test(
+    name = "layout_utils_test",
+    srcs = ["layout_utils_test.cc"],
+    deps = [
+        ":layout_utils",
+        "//xla:shape_util",
+        "//xla:xla_data_proto_cc",
+        "//xla/codegen/emitters:kernel_arguments",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@tensor_ir//:NVTensorIRDialect",
+        "@tensor_ir//:NVTensorIRSupport",
+    ],
+)
+
+xla_cc_test(
+    name = "hlo_to_tensor_ir_test",
+    srcs = ["hlo_to_tensor_ir_test.cc"],
+    deps = [
+        ":hlo_to_tensor_ir",
+        "//xla/hlo/parser:hlo_parser",
+        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+        "//xla/service/llvm_ir:llvm_util",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@stablehlo//:stablehlo_ops",
+        "@tensor_ir//:NVTensorIRDialect",
+    ],
+)
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/compilation_pipeline.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/compilation_pipeline.cc
new file mode 100644
index 00000000000..842bc84801c
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/compilation_pipeline.cc
@@ -0,0 +1,95 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/gpu/codegen/tensor_ir/compilation_pipeline.h"
+
+#include <cstdint>
+#include <cstdlib>
+#include <string>
+#include <vector>
+
+#include "tensor_ir/Compiler/CudaTile/Pipelines.h"
+#include "tensor_ir/Compiler/CudaTile/TileIRAssembly.h"
+#include "tensor_ir/Conversion/TensorToCudaTile/Options.h"
+#include "absl/base/call_once.h"
+#include "absl/log/log.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "mlir/Pass/PassManager.h"
+#include "xla/stream_executor/cuda/cuda_compute_capability.h"
+#include "xla/stream_executor/cuda/subprocess_compilation.h"
+#include "xla/stream_executor/device_description.h"
+#include "tsl/platform/path.h"
+
+namespace xla::gpu::tensor_ir {
+
+mlir::nv_tensor_ir::TensorToCudaTilePipelineOptions GetPipelineOptions(
+    const stream_executor::GpuComputeCapability& gpu_cc,
+    llvm::ArrayRef<int32_t> tile_size, int64_t reduction_tile_size) {
+  mlir::nv_tensor_ir::TensorToCudaTilePipelineOptions options;
+  options.tileSize.assign(tile_size.begin(), tile_size.end());
+  options.reductionTileSize = reduction_tile_size;
+  // Skip the TileAnalyzer candidate search when an explicit tile shape is
+  // already provided (e.g. by the autotuner).
+  options.maxCandidates = tile_size.empty() ? 1 : 0;
+
+  if (const auto* cuda_cc = gpu_cc.cuda_compute_capability()) {
+    options.computeCapability = cuda_cc->major * 10 + cuda_cc->minor;
+  }
+  return options;
+}
+
+void CreateTensorIrPipeline(
+    mlir::OpPassManager* pm,
+    const mlir::nv_tensor_ir::TensorToCudaTilePipelineOptions& options) {
+  mlir::nv_tensor_ir::buildTensorToCudaTileConversionPipeline(*pm, options);
+}
+
+void SetUpTileIrAssembler() {
+  static absl::once_flag once;
+  absl::call_once(once, [] {
+    // The same search that finds `ptxas`: the CUDA toolkit roots first, `$PATH`
+    // last. `tileiras` ships next to `ptxas` from CUDA 13.1 on; TensorIR itself
+    // needs the 13.3 release or newer.
+    absl::StatusOr<std::string> path = stream_executor::FindCudaExecutable(
+        "tileiras", /*preferred_cuda_dir=*/"");
+    if (!path.ok()) {
+      LOG(WARNING) << "TensorIR: no Tile IR assembler (CUDA 13.3 or newer "
+                      "ships one), falling back to Tile IR bytecode for the "
+                      "driver to JIT: "
+                   << path.status();
+      return;
+    }
+    // `tileiras` finds `libnvvm` and `ptxas` relative to the resolved location
+    // of its own executable (`/proc/self/exe`). When the toolkit is staged
+    // through symlinks (bazel runfiles, hermetic CUDA) that resolves to a
+    // location with no neighbours. `CUDA_HOME` anchors both lookups on the
+    // staged layout instead: `<root>/bin/tileiras`,
+    // `<root>/nvvm/lib64/libnvvm.so` and `<root>/bin/ptxas`. The rest of the
+    // environment is a hermetic minimum.
+    std::string root(tsl::io::Dirname(tsl::io::Dirname(*path)));
+    std::vector<std::string> environment = {absl::StrCat("CUDA_HOME=", root)};
+    for (const char* name : {"PATH", "TMPDIR"}) {
+      if (const char* value = std::getenv(name); value != nullptr) {
+        environment.push_back(absl::StrCat(name, "=", value));
+      }
+    }
+    mlir::nv_tensor_ir::backend::cuda_tile::setTileIRAssembler(*path,
+                                                               environment);
+  });
+}
+
+}  // namespace xla::gpu::tensor_ir
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/compilation_pipeline.h b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/compilation_pipeline.h
new file mode 100644
index 00000000000..07918d4196d
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/compilation_pipeline.h
@@ -0,0 +1,56 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_COMPILATION_PIPELINE_H_
+#define XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_COMPILATION_PIPELINE_H_
+
+#include <cstdint>
+
+#include "tensor_ir/Conversion/TensorToCudaTile/Options.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Support/LLVM.h"
+#include "xla/stream_executor/device_description.h"
+
+namespace xla::gpu::tensor_ir {
+
+// Returns the TensorIR-to-CudaTile pipeline options for the given target
+// and tiling configuration. `tile_size` is the output tile shape and
+// `reduction_tile_size` is the tile size for the contracting dimension (only
+// relevant for fusions that contain a reduction).
+mlir::nv_tensor_ir::TensorToCudaTilePipelineOptions GetPipelineOptions(
+    const stream_executor::GpuComputeCapability& gpu_cc,
+    llvm::ArrayRef<int32_t> tile_size,
+    int64_t reduction_tile_size =
+        mlir::nv_tensor_ir::kDefaultReductionTileSize);
+
+// Adds the TensorIR-to-CudaTile analysis and conversion pipeline to `pm`.
+void CreateTensorIrPipeline(
+    mlir::OpPassManager* pm,
+    const mlir::nv_tensor_ir::TensorToCudaTilePipelineOptions& options);
+
+// Points TensorIR's Tile IR assembler at the `tileiras` binary staged in the
+// runfiles. TensorIR only looks on `$PATH` by default, where a Bazel-staged
+// binary never is, and quietly falls back to emitting Tile IR bytecode for the
+// driver to JIT.
+//
+// Runs the lookup once and is safe to call from several threads. Missing an
+// assembler is not an error: CUDA before 13.1 does not ship one, and the
+// bytecode fallback still produces a working kernel wherever the driver can
+// JIT it.
+void SetUpTileIrAssembler();
+
+}  // namespace xla::gpu::tensor_ir
+
+#endif  // XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_COMPILATION_PIPELINE_H_
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/conversion.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/conversion.cc
deleted file mode 100644
index 22a3372183d..00000000000
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/conversion.cc
+++ /dev/null
@@ -1,1041 +0,0 @@
-/* Copyright 2026 The OpenXLA Authors.
-
-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.
-==============================================================================*/
-
-#include "xla/backends/gpu/codegen/tensor_ir/conversion.h"
-
-#include <cstdint>
-#include <numeric>
-#include <optional>
-#include <utility>
-
-#include "tensor_ir/Dialect/TensorIR.h"
-#include "tensor_ir/Dialect/TensorIRAttrs.h"
-#include "absl/algorithm/container.h"
-#include "absl/cleanup/cleanup.h"
-#include "absl/log/log.h"
-#include "absl/status/status.h"
-#include "absl/status/status_macros.h"
-#include "absl/status/statusor.h"
-#include "absl/strings/str_cat.h"
-#include "absl/strings/str_join.h"
-#include "absl/types/span.h"
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/Casting.h"
-#include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/IR/Attributes.h"
-#include "mlir/IR/Block.h"
-#include "mlir/IR/Builders.h"
-#include "mlir/IR/BuiltinAttributeInterfaces.h"
-#include "mlir/IR/BuiltinAttributes.h"
-#include "mlir/IR/BuiltinOps.h"
-#include "mlir/IR/BuiltinTypeInterfaces.h"
-#include "mlir/IR/Location.h"
-#include "mlir/IR/MLIRContext.h"
-#include "mlir/IR/OwningOpRef.h"
-#include "mlir/IR/Value.h"
-#include "mlir/IR/ValueRange.h"
-#include "mlir/IR/Verifier.h"
-#include "mlir/Interfaces/SideEffectInterfaces.h"
-#include "mlir/Support/LLVM.h"
-#include "xla/comparison_util.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_opcode.h"
-#include "xla/layout_util.h"
-#include "xla/literal.h"
-#include "xla/literal_util.h"
-#include "xla/primitive_util.h"
-#include "xla/service/llvm_ir/llvm_util.h"
-#include "xla/shape.h"
-#include "xla/shape_util.h"
-#include "xla/util.h"
-#include "xla/xla_data.pb.h"
-
-// NOLINTBEGIN(clang-diagnostic-pre-c++20-compat)
-
-namespace xla::gpu::tensor_ir {
-namespace {
-
-namespace arith = ::mlir::arith;
-namespace tir = ::mlir::nv_tensor_ir;
-
-absl::StatusOr<mlir::Value> ConvertReductionInstruction(
-    const HloInstruction& source, mlir::ValueRange operands,
-    mlir::Block& target);
-
-// Creates MLIR location from HLO instruction metadata.
-mlir::Location GetLocationFromInstruction(const HloInstruction& source,
-                                          mlir::MLIRContext* context) {
-  const OpMetadata& metadata = source.metadata();
-
-  llvm::SmallVector<mlir::Location> locations;
-  if (!metadata.op_name().empty()) {
-    locations.push_back(
-        mlir::NameLoc::get(mlir::StringAttr::get(context, metadata.op_name())));
-  }
-  if (!metadata.source_file().empty()) {
-    locations.push_back(mlir::FileLineColRange::get(
-        context, metadata.source_file(), metadata.source_line(),
-        metadata.source_column(), metadata.source_end_line(),
-        metadata.source_end_column()));
-  }
-
-  if (locations.empty()) {
-    return mlir::UnknownLoc::get(context);
-  }
-  if (locations.size() == 1) {
-    return locations.front();
-  }
-  return mlir::FusedLoc::get(context, locations);
-}
-
-// Creates MLIR type from HLO primitive type.
-absl::StatusOr<mlir::Type> GetElementType(PrimitiveType type,
-                                          mlir::Builder& builder) {
-  switch (type) {
-    case PrimitiveType::PRED:
-      return builder.getI1Type();
-    case PrimitiveType::S8:
-    case PrimitiveType::S16:
-    case PrimitiveType::S32:
-    case PrimitiveType::S64:
-      return builder.getIntegerType(primitive_util::BitWidth(type),
-                                    /*isSigned=*/true);
-    case PrimitiveType::U8:
-    case PrimitiveType::U16:
-    case PrimitiveType::U32:
-    case PrimitiveType::U64:
-      return builder.getIntegerType(primitive_util::BitWidth(type),
-                                    /*isSigned=*/false);
-    case PrimitiveType::F16:
-      return builder.getF16Type();
-    case PrimitiveType::BF16:
-      return builder.getBF16Type();
-    case PrimitiveType::F32:
-      return builder.getF32Type();
-    case PrimitiveType::F64:
-      return builder.getF64Type();
-    default:
-      return absl::InvalidArgumentError(
-          absl::StrCat("Unsupported primitive type: ", type));
-  }
-}
-
-// Converts a scalar literal to MLIR attribute.
-absl::StatusOr<mlir::TypedAttr> GetScalarAttribute(
-    const LiteralBase& literal, mlir::Builder& builder,
-    bool use_signless_integer_type) {
-  ABSL_ASSIGN_OR_RETURN(auto element_type,
-                   GetElementType(literal.shape().element_type(), builder));
-  if (element_type.isFloat()) {
-    if (auto cst = literal.GetAsDouble({})) {
-      return mlir::FloatAttr::get(element_type, *cst);
-    }
-  } else {
-    if (auto cst = literal.GetIntegralAsS64({})) {
-      if (use_signless_integer_type) {
-        element_type =
-            builder.getIntegerType(element_type.getIntOrFloatBitWidth());
-      }
-      return mlir::IntegerAttr::get(element_type, *cst);
-    }
-  }
-  return absl::InvalidArgumentError(
-      absl::StrCat("Unsupported scalar literal: ", literal.ToString()));
-}
-
-// Creates MLIR tensor type from HLO shape.
-absl::StatusOr<mlir::RankedTensorType> GetTensorType(const Shape& shape,
-                                                     mlir::Builder& builder) {
-  ABSL_ASSIGN_OR_RETURN(auto element_type,
-                   GetElementType(shape.element_type(), builder));
-  // Promote 0-D scalar shapes to 1-D tensors of size 1 for TensorIR.
-  if (shape.dimensions().empty()) {
-    return mlir::RankedTensorType::get({1}, element_type);
-  }
-  return mlir::RankedTensorType::get(shape.dimensions(), element_type);
-}
-
-// Builds a constant operation with given value and shape.
-absl::StatusOr<mlir::Value> BuildConstant(const LiteralBase& literal,
-                                          const Shape& shape,
-                                          mlir::ImplicitLocOpBuilder& builder) {
-  ABSL_ASSIGN_OR_RETURN(auto tensor_type, GetTensorType(shape, builder));
-  auto dense_attr = primitive_util::ArrayTypeSwitch(
-      [&](auto type) -> mlir::DenseElementsAttr {
-        using NativeT = primitive_util::NativeTypeOf<type>;
-        return mlir::DenseElementsAttr::get(tensor_type,
-                                            literal.GetFirstElement<NativeT>());
-      },
-      literal.shape().element_type());
-  return tir::ConstantOp::create(builder, dense_attr);
-}
-
-absl::StatusOr<mlir::Value> BuildFloatConstant(
-    double value, const Shape& shape, mlir::ImplicitLocOpBuilder& builder) {
-  return BuildConstant(LiteralUtil::CreateR0(shape.element_type(), value),
-                       shape, builder);
-}
-
-// Builds a reshape operation from HLO bitcast instruction.
-absl::StatusOr<mlir::Value> BuildBitcast(const HloInstruction& source,
-                                         mlir::Value operand,
-                                         mlir::ImplicitLocOpBuilder& builder) {
-  mlir::Value result = operand;
-
-  // Calculate the operand/result shapes with default layout.
-  Shape operand_normal_shape =
-      ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout(
-          source.operand(0)->shape());
-  Shape result_normal_shape =
-      ShapeUtil::MakeShapeWithDescendingLayoutAndSamePhysicalLayout(
-          source.shape());
-
-  // If the operand doesn't have a default layout, transpose it.
-  // Example: [A,B,C]{2,0,1} -> [B,A,C]{2,1,0} uses permutation (1,0,2)
-  if (source.operand(0)->shape() != operand_normal_shape) {
-    auto permutation = llvm::to_vector(
-        llvm::reverse(source.operand(0)->shape().layout().minor_to_major()));
-    ABSL_ASSIGN_OR_RETURN(auto transpose_type,
-                     GetTensorType(operand_normal_shape, builder));
-    VLOG(3) << "Transposing operand: "
-            << llvm_ir::DumpToString(result.getType()) << " to "
-            << llvm_ir::DumpToString(transpose_type);
-    result =
-        tir::TransposeOp::create(builder, transpose_type, result, permutation);
-  }
-
-  // If the normalized shapes are not equal, create a reshape operation.
-  if (operand_normal_shape != result_normal_shape) {
-    ABSL_ASSIGN_OR_RETURN(auto reshape_type,
-                     GetTensorType(result_normal_shape, builder));
-    VLOG(3) << "Reshaping operand: " << llvm_ir::DumpToString(result.getType())
-            << " to " << llvm_ir::DumpToString(reshape_type);
-    result = tir::ReshapeOp::create(builder, reshape_type, result);
-  }
-
-  // If the result doesn't have a default layout, transpose it.
-  // Example: [A,B,C]{2,1,0} -> [C,A,B]{0,2,1} uses permutation (2,0,1)
-  if (source.shape() != result_normal_shape) {
-    llvm::SmallVector<int64_t> permutation(source.shape().dimensions().size());
-    for (auto [idx, pos] : llvm::enumerate(
-             llvm::reverse(source.shape().layout().minor_to_major()))) {
-      permutation[pos] = idx;
-    }
-    ABSL_ASSIGN_OR_RETURN(auto transpose_type,
-                     GetTensorType(source.shape(), builder));
-    VLOG(3) << "Transposing operand: "
-            << llvm_ir::DumpToString(result.getType()) << " to "
-            << llvm_ir::DumpToString(transpose_type);
-    result =
-        tir::TransposeOp::create(builder, transpose_type, result, permutation);
-  }
-
-  return result;
-}
-
-// Builds a broadcast operation from HLO broadcast instruction.
-absl::StatusOr<mlir::Value> BuildBroadcast(
-    const HloBroadcastInstruction& source, mlir::Value operand,
-    mlir::ImplicitLocOpBuilder& builder) {
-  auto dimensions = source.dimensions();
-
-  // If the operand is a scalar constant, return shaped constant.
-  if (const auto* constant = DynCast<HloConstantInstruction>(source.operand(0));
-      constant != nullptr && constant->shape().dimensions().empty()) {
-    VLOG(3) << "Broadcasting scalar constant: "
-            << constant->literal().ToString() << " to " << source.shape();
-    return BuildConstant(constant->literal(), source.shape(), builder);
-  }
-
-  // If the source dimensions are not sorted, a transpose is needed.
-  // https://openxla.org/stablehlo/spec#broadcast_in_dim
-  if (!absl::c_is_sorted(dimensions)) {
-    // Get the transpose permutation by sorting the dimensions.
-    llvm::SmallVector<int64_t> permutation(dimensions.size());
-    std::iota(permutation.begin(), permutation.end(), 0);
-    absl::c_sort(permutation, [&](int64_t a, int64_t b) {
-      return dimensions[a] < dimensions[b];
-    });
-
-    // Create the transpose operation.
-    llvm::SmallVector<int64_t> permuted_dims;
-    permuted_dims.reserve(dimensions.size());
-    for (int64_t i : permutation) {
-      permuted_dims.push_back(source.operand(0)->shape().dimensions(i));
-    }
-    mlir::ShapedType transpose_type =
-        llvm::cast<mlir::ShapedType>(operand.getType()).clone(permuted_dims);
-    VLOG(3) << "Transposing operand: "
-            << llvm_ir::DumpToString(operand.getType()) << " to "
-            << llvm_ir::DumpToString(transpose_type);
-    operand =
-        tir::TransposeOp::create(builder, transpose_type, operand, permutation);
-  }
-
-  // Reshape operand, add unit dimensions for the broadcast.
-  llvm::SmallVector<int64_t> pre_broadcast_dims;
-  for (int i = 0, n = source.shape().dimensions().size(); i < n; ++i) {
-    auto it = absl::c_find(dimensions, i);
-    if (it != dimensions.end()) {
-      pre_broadcast_dims.push_back(
-          source.operand(0)->shape().dimensions(it - dimensions.begin()));
-    } else {
-      pre_broadcast_dims.push_back(1);
-    }
-  }
-  mlir::ShapedType reshape_type =
-      llvm::cast<mlir::ShapedType>(operand.getType()).clone(pre_broadcast_dims);
-  VLOG(3) << "Reshaping operand: " << llvm_ir::DumpToString(operand.getType())
-          << " to " << llvm_ir::DumpToString(reshape_type);
-  operand = tir::ReshapeOp::create(builder, reshape_type, operand);
-
-  // Create the broadcast operation.
-  mlir::ShapedType broadcast_type =
-      reshape_type.clone(source.shape().dimensions());
-  VLOG(3) << "Broadcasting operand: "
-          << llvm_ir::DumpToString(operand.getType()) << " to "
-          << llvm_ir::DumpToString(broadcast_type);
-  return tir::BroadcastOp::create(builder, broadcast_type, operand);
-}
-
-// Builds a dot operation from HLO dot instruction.
-absl::StatusOr<mlir::Value> BuildDot(const HloDotInstruction& source,
-                                     mlir::Value lhs, mlir::Value rhs,
-                                     mlir::ImplicitLocOpBuilder& builder) {
-  // Calculate LHS transpose permutation.
-  int lhs_rank = source.operand(0)->shape().dimensions().size();
-  const auto& lhs_batch = source.dot_dimension_numbers().lhs_batch_dimensions();
-  const auto& lhs_contracting =
-      source.dot_dimension_numbers().lhs_contracting_dimensions();
-  auto lhs_non_contracting =
-      GetNonContractingDims(lhs_rank, lhs_contracting, lhs_batch);
-
-  llvm::SmallVector<int64_t> lhs_permutation(lhs_batch.begin(),
-                                             lhs_batch.end());
-  lhs_permutation.append(lhs_non_contracting.begin(),
-                         lhs_non_contracting.end());
-  lhs_permutation.append(lhs_contracting.begin(), lhs_contracting.end());
-
-  // Transpose LHS, if needed.
-  llvm::SmallVector<int64_t> lhs_permuted_dims;
-  for (int i = 0; i < lhs_rank; ++i) {
-    lhs_permuted_dims.push_back(
-        source.operand(0)->shape().dimensions(lhs_permutation[i]));
-  }
-
-  if (!absl::c_is_sorted(lhs_permutation)) {
-    mlir::ShapedType transpose_type =
-        llvm::cast<mlir::ShapedType>(lhs.getType()).clone(lhs_permuted_dims);
-    VLOG(3) << "Transposing LHS operand: "
-            << llvm_ir::DumpToString(lhs.getType()) << " to "
-            << llvm_ir::DumpToString(transpose_type);
-    lhs =
-        tir::TransposeOp::create(builder, transpose_type, lhs, lhs_permutation);
-  }
-
-  // Calculate LHS dot operand shape.
-  int64_t batch_size =
-      Product(absl::MakeSpan(lhs_permuted_dims).first(lhs_batch.size()));
-  int64_t contracting_size =
-      Product(absl::MakeSpan(lhs_permuted_dims).last(lhs_contracting.size()));
-  int64_t lhs_non_contracting_size =
-      Product(absl::MakeSpan(lhs_permuted_dims)
-                  .subspan(lhs_batch.size(), lhs_non_contracting.size()));
-
-  llvm::SmallVector<int64_t> lhs_expected_dims;
-  if (!lhs_batch.empty()) {
-    lhs_expected_dims.push_back(batch_size);
-  }
-  lhs_expected_dims.push_back(lhs_non_contracting_size);
-  lhs_expected_dims.push_back(contracting_size);
-
-  // Reshape LHS, if needed.
-  auto lhs_reshape_type =
-      llvm::cast<mlir::ShapedType>(lhs.getType()).clone(lhs_expected_dims);
-  if (lhs.getType() != lhs_reshape_type) {
-    VLOG(3) << "Reshaping LHS operand: " << llvm_ir::DumpToString(lhs.getType())
-            << " to " << llvm_ir::DumpToString(lhs_reshape_type);
-    lhs = tir::ReshapeOp::create(builder, lhs_reshape_type, lhs);
-  }
-
-  // Calculate RHS transpose permutation.
-  int rhs_rank = source.operand(1)->shape().dimensions().size();
-  const auto& rhs_batch = source.dot_dimension_numbers().rhs_batch_dimensions();
-  const auto& rhs_contracting =
-      source.dot_dimension_numbers().rhs_contracting_dimensions();
-  auto rhs_non_contracting =
-      GetNonContractingDims(rhs_rank, rhs_contracting, rhs_batch);
-
-  llvm::SmallVector<int64_t> rhs_permutation(rhs_batch.begin(),
-                                             rhs_batch.end());
-  rhs_permutation.append(rhs_contracting.begin(), rhs_contracting.end());
-  rhs_permutation.append(rhs_non_contracting.begin(),
-                         rhs_non_contracting.end());
-
-  // Transpose RHS, if needed.
-  llvm::SmallVector<int64_t> rhs_permuted_dims;
-  for (int i = 0; i < rhs_rank; ++i) {
-    rhs_permuted_dims.push_back(
-        source.operand(1)->shape().dimensions(rhs_permutation[i]));
-  }
-
-  if (!absl::c_is_sorted(rhs_permutation)) {
-    mlir::ShapedType transpose_type =
-        llvm::cast<mlir::ShapedType>(rhs.getType()).clone(rhs_permuted_dims);
-    VLOG(3) << "Transposing RHS operand: "
-            << llvm_ir::DumpToString(rhs.getType()) << " to "
-            << llvm_ir::DumpToString(transpose_type);
-    rhs =
-        tir::TransposeOp::create(builder, transpose_type, rhs, rhs_permutation);
-  }
-
-  // Calculate RHS dot operand shape.
-  int64_t rhs_non_contracting_size = Product(
-      absl::MakeSpan(rhs_permuted_dims).last(rhs_non_contracting.size()));
-
-  llvm::SmallVector<int64_t> rhs_expected_dims;
-  if (!rhs_batch.empty()) {
-    rhs_expected_dims.push_back(batch_size);
-  }
-  rhs_expected_dims.push_back(contracting_size);
-  rhs_expected_dims.push_back(rhs_non_contracting_size);
-
-  // Reshape RHS, if needed.
-  auto rhs_reshape_type =
-      llvm::cast<mlir::ShapedType>(rhs.getType()).clone(rhs_expected_dims);
-  if (rhs.getType() != rhs_reshape_type) {
-    VLOG(3) << "Reshaping RHS operand: " << llvm_ir::DumpToString(rhs.getType())
-            << " to " << llvm_ir::DumpToString(rhs_reshape_type);
-    rhs = tir::ReshapeOp::create(builder, rhs_reshape_type, rhs);
-  }
-
-  // Calculate matmul result shape.
-  ABSL_ASSIGN_OR_RETURN(auto result_type, GetTensorType(source.shape(), builder));
-
-  llvm::SmallVector<int64_t> matmul_shape;
-  if (!lhs_batch.empty()) {
-    matmul_shape.push_back(batch_size);
-  }
-  matmul_shape.push_back(lhs_non_contracting_size);
-  matmul_shape.push_back(rhs_non_contracting_size);
-  mlir::ShapedType matmul_type = result_type.clone(matmul_shape);
-
-  // Create the matmul operation.
-  VLOG(3) << "Creating matmul operation: " << source.name();
-  auto matmul_op = tir::MatmulOp::create(builder, matmul_type, lhs, rhs);
-  mlir::Value result = matmul_op.getResult();
-
-  // Align the result to the source shape.
-  if (result_type != matmul_type) {
-    VLOG(3) << "Reshaping result: " << llvm_ir::DumpToString(matmul_type)
-            << " to " << llvm_ir::DumpToString(result_type);
-    result = tir::ReshapeOp::create(builder, result_type, result);
-  }
-  return result;
-}
-
-// Returns the shaped type for reduction output with unit dimensions for reduced
-// axes.
-mlir::ShapedType GetReduceOutputType(mlir::RankedTensorType result_type,
-                                     const Shape& operand_shape,
-                                     const Shape& result_shape,
-                                     absl::Span<const int32_t> reduce_dims) {
-  llvm::SmallVector<int64_t> output_shape;
-  output_shape.reserve(operand_shape.dimensions().size());
-  for (int i = 0, p = 0; i < operand_shape.dimensions().size(); ++i) {
-    if (!absl::c_contains(reduce_dims, i)) {
-      output_shape.push_back(result_shape.dimensions(p++));
-    } else {
-      output_shape.push_back(1);
-    }
-  }
-  return result_type.clone(output_shape);
-}
-
-// Builds a reduction operation from HLO reduce instruction.
-absl::StatusOr<mlir::Value> BuildReduce(const HloReduceInstruction& source,
-                                        mlir::ValueRange operands,
-                                        mlir::ImplicitLocOpBuilder& builder) {
-  llvm::SmallVector<mlir::Attribute> initial_values;
-  for (const HloInstruction* init_instr : source.init_values()) {
-    const auto* constant = Cast<HloConstantInstruction>(init_instr);
-    ABSL_ASSIGN_OR_RETURN(auto initial_value,
-                     GetScalarAttribute(constant->literal(), builder,
-                                        /*use_signless_integer_type=*/false));
-    initial_values.push_back(initial_value);
-  }
-
-  llvm::SmallVector<int32_t> reduce_dims(source.dimensions().begin(),
-                                         source.dimensions().end());
-  ABSL_ASSIGN_OR_RETURN(auto result_type, GetTensorType(source.shape(), builder));
-  mlir::ShapedType output_type = GetReduceOutputType(
-      result_type, source.operand(0)->shape(), source.shape(), reduce_dims);
-
-  VLOG(3) << "Creating reduction operation: " << source.name();
-  auto reduce_op = tir::ReduceUDOp::create(
-      builder, output_type, operands.take_front(source.input_count()),
-      reduce_dims, builder.getArrayAttr(initial_values));
-
-  mlir::Block* body = builder.createBlock(&reduce_op.getRegion());
-  for (int i = 0; i < 2; ++i) {
-    for (mlir::Attribute init : initial_values) {
-      auto arg_type = llvm::cast<mlir::TypedAttr>(init).getType();
-      if (arg_type.isInteger()) {
-        arg_type = builder.getIntegerType(arg_type.getIntOrFloatBitWidth());
-      }
-      body->addArgument(arg_type, reduce_op.getLoc());
-    }
-  }
-
-  const HloComputation* computation = source.to_apply();
-  VLOG(3) << "Converting reduction HLO computation: " << computation->name();
-
-  llvm::DenseMap<const HloInstruction*, mlir::Value> converted;
-  for (const HloInstruction* instruction :
-       computation->MakeInstructionPostOrder()) {
-    llvm::SmallVector<mlir::Value> new_operands;
-    for (const auto& operand : instruction->operands()) {
-      new_operands.push_back(converted[operand]);
-    }
-    ABSL_ASSIGN_OR_RETURN(
-        converted[instruction],
-        ConvertReductionInstruction(*instruction, new_operands, *body));
-  }
-
-  mlir::Value result = converted[computation->root_instruction()];
-  tir::YieldOp::create(builder, result);
-  builder.setInsertionPointAfter(reduce_op);
-
-  VLOG(3) << "Reshaping result: " << llvm_ir::DumpToString(output_type)
-          << " to " << llvm_ir::DumpToString(result_type);
-  return tir::ReshapeOp::create(builder, result_type, reduce_op.getResult(0));
-}
-
-// Returns the reduction type, if the body is a supported single-op reduction.
-std::optional<tir::ReductionMode> GetReductionType(
-    const HloReduceInstruction& reduce) {
-  // Reduction must have exactly two parameters.
-  const HloComputation* comp = reduce.to_apply();
-  if (comp->num_parameters() != 2) {
-    return std::nullopt;
-  }
-  const HloInstruction* param0 = comp->parameter_instruction(0);
-  const HloInstruction* param1 = comp->parameter_instruction(1);
-
-  // Reduction root must use the parameters as operands.
-  const HloInstruction* root = comp->root_instruction();
-  if (root->operand_count() != 2 || root->operand(0) != param0 ||
-      root->operand(1) != param1) {
-    return std::nullopt;
-  }
-
-  // Reduction initial value must be a constant.
-  const auto* init_value =
-      DynCast<HloConstantInstruction>(reduce.init_values().front());
-  if (init_value == nullptr) {
-    return std::nullopt;
-  }
-
-  if (root->opcode() == HloOpcode::kAdd && init_value->literal().IsAll(0)) {
-    return tir::ReductionMode::add;
-  }
-
-  if (root->opcode() == HloOpcode::kMultiply &&
-      init_value->literal().IsAll(1)) {
-    return tir::ReductionMode::mul;
-  }
-
-  return std::nullopt;
-}
-
-// Builds a reduction operation for a reduction type.
-absl::StatusOr<mlir::Value> BuildSimpleReduce(
-    const HloReduceInstruction& source, tir::ReductionMode mode,
-    mlir::Value operand, mlir::ImplicitLocOpBuilder& builder) {
-  llvm::SmallVector<int32_t> reduce_dims(source.dimensions().begin(),
-                                         source.dimensions().end());
-  ABSL_ASSIGN_OR_RETURN(auto result_type, GetTensorType(source.shape(), builder));
-  mlir::ShapedType output_type = GetReduceOutputType(
-      result_type, source.operand(0)->shape(), source.shape(), reduce_dims);
-
-  VLOG(3) << "Creating reduction operation: " << source.name();
-  auto reduce_op =
-      tir::ReduceOp::create(builder, output_type, operand, reduce_dims, mode);
-
-  VLOG(3) << "Reshaping result: " << llvm_ir::DumpToString(output_type)
-          << " to " << llvm_ir::DumpToString(result_type);
-  return tir::ReshapeOp::create(builder, result_type, reduce_op.getResult());
-}
-
-// Maps HLO compare instruction to TensorIR comparator.
-tir::Comparator GetComparator(const HloCompareInstruction& source) {
-  bool is_float = primitive_util::IsFloatingPointType(
-      source.operand(0)->shape().element_type());
-  if (is_float) {
-    switch (source.direction()) {
-      case ComparisonDirection::kEq:
-        return tir::Comparator::oeq;
-      case ComparisonDirection::kNe:
-        return tir::Comparator::une;
-      case ComparisonDirection::kGt:
-        return tir::Comparator::ogt;
-      case ComparisonDirection::kGe:
-        return tir::Comparator::oge;
-      case ComparisonDirection::kLt:
-        return tir::Comparator::olt;
-      case ComparisonDirection::kLe:
-        return tir::Comparator::ole;
-    }
-  } else {
-    switch (source.direction()) {
-      case ComparisonDirection::kEq:
-        return tir::Comparator::eq;
-      case ComparisonDirection::kNe:
-        return tir::Comparator::neq;
-      case ComparisonDirection::kGt:
-        return tir::Comparator::gt;
-      case ComparisonDirection::kGe:
-        return tir::Comparator::ge;
-      case ComparisonDirection::kLt:
-        return tir::Comparator::lt;
-      case ComparisonDirection::kLe:
-        return tir::Comparator::le;
-    }
-  }
-}
-
-// Maps HLO compare instruction to `arith` comparator (floating-point).
-arith::CmpFPredicate GetArithFloatComparator(
-    const HloCompareInstruction& source) {
-  switch (source.direction()) {
-    case ComparisonDirection::kEq:
-      return arith::CmpFPredicate::OEQ;
-    case ComparisonDirection::kNe:
-      return arith::CmpFPredicate::UNE;
-    case ComparisonDirection::kGt:
-      return arith::CmpFPredicate::OGT;
-    case ComparisonDirection::kGe:
-      return arith::CmpFPredicate::OGE;
-    case ComparisonDirection::kLt:
-      return arith::CmpFPredicate::OLT;
-    case ComparisonDirection::kLe:
-      return arith::CmpFPredicate::OLE;
-  }
-}
-
-// Maps HLO compare instruction to `arith` comparator (integer).
-arith::CmpIPredicate GetArithIntComparator(
-    const HloCompareInstruction& source) {
-  bool is_signed = primitive_util::IsSignedIntegralType(
-      source.operand(0)->shape().element_type());
-  switch (source.direction()) {
-    case ComparisonDirection::kEq:
-      return arith::CmpIPredicate::eq;
-    case ComparisonDirection::kNe:
-      return arith::CmpIPredicate::ne;
-    case ComparisonDirection::kGt:
-      return is_signed ? arith::CmpIPredicate::sgt : arith::CmpIPredicate::ugt;
-    case ComparisonDirection::kGe:
-      return is_signed ? arith::CmpIPredicate::sge : arith::CmpIPredicate::uge;
-    case ComparisonDirection::kLt:
-      return is_signed ? arith::CmpIPredicate::slt : arith::CmpIPredicate::ult;
-    case ComparisonDirection::kLe:
-      return is_signed ? arith::CmpIPredicate::sle : arith::CmpIPredicate::ule;
-  }
-}
-
-// Template for creating `arith` operations.
-template <typename Op, typename... Args>
-mlir::Value CreateArithOp(Args&&... args) {
-  return Op::create(std::forward<Args>(args)...).getResult();
-}
-
-// Creates strides attribute (helper).
-mlir::NamedAttribute CreateStridesAttributeFromShape(mlir::MLIRContext* context,
-                                                     const Shape& shape) {
-  // Calculate strides from the shape and the layout.
-  llvm::SmallVector<int64_t> strides(shape.dimensions().size());
-  int64_t stride = 1;
-  for (int dim : shape.layout().minor_to_major()) {
-    strides[dim] = stride;
-    stride *= shape.dimensions(dim);
-  }
-
-  // Create the attribute with strides formatted as a string.
-  auto strides_fmt = absl::StrCat("(", absl::StrJoin(strides, ","), ")");
-  auto strides_attr = mlir::StringAttr::get(context, strides_fmt);
-  return {tir::TensorIRDialect::getStrideAttrName(), strides_attr};
-}
-
-mlir::DictionaryAttr CreateStridesDictionaryAttribute(
-    mlir::MLIRContext* context, const Shape& shape) {
-  if (LayoutUtil::IsMonotonicWithDim0Major(shape.layout())) {
-    return mlir::DictionaryAttr::get(context);
-  }
-  auto attr = CreateStridesAttributeFromShape(context, shape);
-  return mlir::DictionaryAttr::get(context, {attr});
-}
-
-// Converts HLO instruction `source` with already-converted `operands` values
-// to TensorIR ops and insert them at the end of `target` block.
-absl::StatusOr<mlir::Value> ConvertFusionInstruction(
-    const HloInstruction& source, mlir::ValueRange operands,
-    mlir::Block& target) {
-  mlir::MLIRContext* context = target.getParent()->getContext();
-  VLOG(3) << "Converting HLO instruction: " << source.ToString();
-
-  mlir::Location location = GetLocationFromInstruction(source, context);
-  mlir::ImplicitLocOpBuilder builder(location, context);
-  builder.setInsertionPointToEnd(&target);
-
-  switch (source.opcode()) {
-    // Unary elementwise operations.
-    case HloOpcode::kAbs:
-      return tir::AbsOp::create(builder, operands);
-    case HloOpcode::kCeil:
-      return tir::CeilOp::create(builder, operands);
-    case HloOpcode::kConvert: {
-      ABSL_ASSIGN_OR_RETURN(auto convert_type,
-                       GetTensorType(source.shape(), builder));
-      return tir::ConvertOp::create(builder, convert_type, operands[0]);
-    }
-    case HloOpcode::kCos:
-      return tir::CosOp::create(builder, operands);
-    case HloOpcode::kErf:
-      return tir::ErfOp::create(builder, operands);
-    case HloOpcode::kExp:
-      return tir::ExpOp::create(builder, operands);
-    case HloOpcode::kExpm1: {
-      // NOTE: Decomposing Expm1(x) to Exp(x) - 1.0 can result in a severe loss
-      // of precision for values of x close to 0.
-      ABSL_ASSIGN_OR_RETURN(auto one,
-                       BuildFloatConstant(1.0, source.shape(), builder));
-      return tir::SubOp::create(builder,
-                                tir::ExpOp::create(builder, operands[0]), one);
-    }
-    case HloOpcode::kFloor:
-      return tir::FloorOp::create(builder, operands);
-    case HloOpcode::kLog:
-      return tir::LogOp::create(builder, operands);
-    case HloOpcode::kLog1p: {
-      // NOTE: Decomposing Log1p(x) to Log(x + 1.0) can result in a severe loss
-      // of precision for values of x close to 0.
-      ABSL_ASSIGN_OR_RETURN(auto one,
-                       BuildFloatConstant(1.0, source.shape(), builder));
-      return tir::LogOp::create(builder,
-                                tir::AddOp::create(builder, operands[0], one));
-    }
-    case HloOpcode::kNot:
-      return tir::LogicalNotOp::create(builder, operands);
-    case HloOpcode::kNegate:
-      return tir::NegOp::create(builder, operands);
-    case HloOpcode::kRsqrt:
-      return tir::RsqrtOp::create(builder, operands);
-    case HloOpcode::kSin:
-      return tir::SinOp::create(builder, operands);
-    case HloOpcode::kSqrt:
-      return tir::SqrtOp::create(builder, operands);
-    case HloOpcode::kTan:
-      return tir::TanOp::create(builder, operands);
-    case HloOpcode::kTanh:
-      return tir::TanhFwdOp::create(builder, operands);
-
-    // Binary elementwise operations.
-    case HloOpcode::kAdd:
-      return tir::AddOp::create(builder, operands);
-    case HloOpcode::kAtan2:
-      return tir::Atan2Op::create(builder, operands);
-    case HloOpcode::kCompare: {
-      tir::Comparator comparator =
-          GetComparator(*Cast<HloCompareInstruction>(&source));
-      return tir::CmpOp::create(builder, comparator, operands[0], operands[1]);
-    }
-    case HloOpcode::kDivide:
-      return tir::DivOp::create(builder, operands);
-    case HloOpcode::kMaximum:
-      return tir::MaxOp::create(builder, operands);
-    case HloOpcode::kMinimum:
-      return tir::MinOp::create(builder, operands);
-    case HloOpcode::kMultiply:
-      return tir::MulOp::create(builder, operands);
-    case HloOpcode::kPower:
-      return tir::PowOp::create(builder, operands);
-    case HloOpcode::kRemainder:
-      return tir::RemOp::create(builder, operands);
-    case HloOpcode::kSubtract:
-      return tir::SubOp::create(builder, operands);
-    case HloOpcode::kAnd:
-      return tir::LogicalAndOp::create(builder, operands);
-    case HloOpcode::kOr:
-      return tir::LogicalOrOp::create(builder, operands);
-
-    // Ternary elementwise operations.
-    case HloOpcode::kSelect:
-      return tir::BinarySelectOp::create(builder, operands);
-    case HloOpcode::kClamp:
-      return tir::MinOp::create(
-          builder, operands[2],
-          tir::MaxOp::create(builder, operands[0], operands[1]));
-
-    // Layout modification operations.
-    case HloOpcode::kReshape:
-    case HloOpcode::kBitcast:
-      return BuildBitcast(source, operands[0], builder);
-    case HloOpcode::kBroadcast:
-      return BuildBroadcast(*Cast<HloBroadcastInstruction>(&source),
-                            operands[0], builder);
-    case HloOpcode::kSlice: {
-      ABSL_ASSIGN_OR_RETURN(auto slice_type, GetTensorType(source.shape(), builder));
-      return tir::SliceOp::create(builder, slice_type, operands[0],
-                                  source.slice_starts(), source.slice_limits(),
-                                  source.slice_strides());
-    }
-    case HloOpcode::kTranspose: {
-      auto transpose = Cast<HloTransposeInstruction>(&source);
-      ABSL_ASSIGN_OR_RETURN(auto transpose_type,
-                       GetTensorType(source.shape(), builder));
-      return tir::TransposeOp::create(builder, transpose_type, operands[0],
-                                      transpose->dimensions());
-    }
-
-    // Reduction operations.
-    case HloOpcode::kDot:
-      return BuildDot(*Cast<HloDotInstruction>(&source), operands[0],
-                      operands[1], builder);
-    case HloOpcode::kReduce: {
-      auto reduce = Cast<HloReduceInstruction>(&source);
-      auto type = GetReductionType(*reduce);
-      return type.has_value()
-                 ? BuildSimpleReduce(*reduce, *type, operands[0], builder)
-                 : BuildReduce(*reduce, operands, builder);
-    }
-
-    // Miscellaneous operations.
-    case HloOpcode::kParameter:
-      return target.getArgument(source.parameter_number());
-    case HloOpcode::kConstant:
-      return BuildConstant(source.literal(), source.shape(), builder);
-    case HloOpcode::kIota: {
-      auto iota = Cast<HloIotaInstruction>(&source);
-      ABSL_ASSIGN_OR_RETURN(auto iota_type, GetTensorType(source.shape(), builder));
-      return tir::IotaOp::create(builder, iota_type, iota->iota_dimension(),
-                                 /*dynamic_sizes=*/{});
-    }
-    case HloOpcode::kConcatenate: {
-      ABSL_ASSIGN_OR_RETURN(auto concatenate_type,
-                       GetTensorType(source.shape(), builder));
-      return tir::ConcatenateOp::create(builder, concatenate_type, operands,
-                                        source.concatenate_dimension());
-    }
-
-    default:
-      return absl::UnimplementedError(absl::StrCat(
-          "Unsupported instruction: ", HloOpcodeString(source.opcode())));
-  }
-}
-
-// Converts HLO instruction `source` in the context of a reduction computation
-// with already-converted `operands` values to TensorIR ops and insert them at
-// the end of `target` block.
-absl::StatusOr<mlir::Value> ConvertReductionInstruction(
-    const HloInstruction& source, mlir::ValueRange operands,
-    mlir::Block& target) {
-  mlir::MLIRContext* context = target.getParent()->getContext();
-  VLOG(3) << "Converting HLO instruction: " << source.ToString();
-
-  mlir::Location location = GetLocationFromInstruction(source, context);
-  mlir::ImplicitLocOpBuilder builder(location, context);
-  builder.setInsertionPointToEnd(&target);
-
-  // `arith` dialect (used in reduction body) has separate instructions for
-  // floating-point and signed/unsigned integer types.
-  const HloInstruction* type_source =
-      source.operand_count() != 0 ? source.operand(0) : &source;
-  PrimitiveType element_type = type_source->shape().element_type();
-  bool is_float = primitive_util::IsFloatingPointType(element_type);
-  bool is_signed = primitive_util::IsSignedIntegralType(element_type);
-
-  switch (source.opcode()) {
-    // Binary elementwise operations.
-    case HloOpcode::kAdd:
-      return is_float ? CreateArithOp<arith::AddFOp>(builder, operands)
-                      : CreateArithOp<arith::AddIOp>(builder, operands);
-    case HloOpcode::kCompare: {
-      const auto& compare = *Cast<HloCompareInstruction>(&source);
-      return is_float ? CreateArithOp<arith::CmpFOp>(
-                            builder, GetArithFloatComparator(compare),
-                            operands[0], operands[1])
-                      : CreateArithOp<arith::CmpIOp>(
-                            builder, GetArithIntComparator(compare),
-                            operands[0], operands[1]);
-    }
-    case HloOpcode::kMaximum:
-      return is_float    ? CreateArithOp<arith::MaximumFOp>(builder, operands)
-             : is_signed ? CreateArithOp<arith::MaxSIOp>(builder, operands)
-                         : CreateArithOp<arith::MaxUIOp>(builder, operands);
-    case HloOpcode::kMinimum:
-      return is_float    ? CreateArithOp<arith::MinimumFOp>(builder, operands)
-             : is_signed ? CreateArithOp<arith::MinSIOp>(builder, operands)
-                         : CreateArithOp<arith::MinUIOp>(builder, operands);
-    case HloOpcode::kMultiply:
-      return is_float ? CreateArithOp<arith::MulFOp>(builder, operands)
-                      : CreateArithOp<arith::MulIOp>(builder, operands);
-    case HloOpcode::kAnd:
-      return CreateArithOp<arith::AndIOp>(builder, operands);
-    case HloOpcode::kOr:
-      return CreateArithOp<arith::OrIOp>(builder, operands);
-    case HloOpcode::kXor:
-      return CreateArithOp<arith::XOrIOp>(builder, operands);
-
-    // Ternary elementwise operations.
-    case HloOpcode::kSelect:
-      return CreateArithOp<arith::SelectOp>(builder, operands);
-    case HloOpcode::kClamp:
-      return is_float    ? CreateArithOp<arith::MinimumFOp>(
-                               builder, operands[2],
-                               CreateArithOp<arith::MaximumFOp>(
-                                   builder, operands[0], operands[1]))
-             : is_signed ? CreateArithOp<arith::MinSIOp>(
-                               builder, operands[2],
-                               CreateArithOp<arith::MaxSIOp>(
-                                   builder, operands[0], operands[1]))
-                         : CreateArithOp<arith::MinUIOp>(
-                               builder, operands[2],
-                               CreateArithOp<arith::MaxUIOp>(
-                                   builder, operands[0], operands[1]));
-
-    // Miscellaneous operations.
-    case HloOpcode::kParameter:
-      return target.getArgument(source.parameter_number());
-    case HloOpcode::kConstant: {
-      const auto& constant = *Cast<HloConstantInstruction>(&source);
-      ABSL_ASSIGN_OR_RETURN(auto value,
-                       GetScalarAttribute(constant.literal(), builder,
-                                          /*use_signless_integer_type=*/true));
-      return arith::ConstantOp::create(builder, value);
-    }
-
-    default:
-      return absl::UnimplementedError(absl::StrCat(
-          "Unsupported instruction: ", HloOpcodeString(source.opcode())));
-  }
-}
-
-}  // namespace
-
-// Converts an HLO fusion computation into a TensorIR `GraphOp` appended
-// to the body of `target`. If conversion fails, `target` is unmodified.
-absl::StatusOr<mlir::nv_tensor_ir::GraphOp> ConvertFusionComputation(
-    const HloComputation& source, mlir::ModuleOp target) {
-  mlir::MLIRContext* context = target.getContext();
-  VLOG(3) << "Converting HLO computation: " << source.name();
-
-  mlir::ImplicitLocOpBuilder builder(mlir::UnknownLoc::get(context), context);
-  builder.setInsertionPointToEnd(target.getBody());
-
-  llvm::SmallVector<mlir::Type> input_types;
-  llvm::SmallVector<mlir::Attribute> input_attrs;
-  for (const auto& parameter : source.parameter_instructions()) {
-    ABSL_ASSIGN_OR_RETURN(auto parameter_type,
-                     GetTensorType(parameter->shape(), builder));
-    input_types.push_back(parameter_type);
-    input_attrs.push_back(
-        CreateStridesDictionaryAttribute(context, parameter->shape()));
-  }
-
-  mlir::ArrayAttr arg_attrs;
-  if (llvm::any_of(input_attrs, [](mlir::Attribute attr) {
-        return !llvm::cast<mlir::DictionaryAttr>(attr).empty();
-      })) {
-    arg_attrs = mlir::ArrayAttr::get(context, input_attrs);
-  }
-
-  llvm::SmallVector<mlir::Type> output_types;
-  ABSL_ASSIGN_OR_RETURN(auto output_type,
-                   GetTensorType(source.root_instruction()->shape(), builder));
-  output_types.push_back(output_type);
-
-  mlir::ArrayAttr res_attrs;
-  auto output_attr = CreateStridesDictionaryAttribute(
-      context, source.root_instruction()->shape());
-  if (!output_attr.empty()) {
-    res_attrs = mlir::ArrayAttr::get(context, {output_attr});
-  }
-
-  auto function_type =
-      mlir::FunctionType::get(context, input_types, output_types);
-  VLOG(3) << "Function type: " << llvm_ir::DumpToString(function_type);
-
-  auto graph_op = tir::GraphOp::create(builder, source.name(),
-                                       /*sym_visibility=*/nullptr,
-                                       function_type, arg_attrs, res_attrs);
-
-  absl::Cleanup cleanup = [&] { graph_op.erase(); };
-
-  mlir::Block* body = builder.createBlock(&graph_op.getRegion());
-  for (mlir::Type type : input_types) {
-    body->addArgument(type, target.getLoc());
-  }
-
-  llvm::DenseMap<const HloInstruction*, mlir::Value> converted;
-  for (const HloInstruction* instruction : source.MakeInstructionPostOrder()) {
-    llvm::SmallVector<mlir::Value> operands;
-    operands.reserve(instruction->operand_count());
-    for (const auto& operand : instruction->operands()) {
-      operands.push_back(converted[operand]);
-    }
-    ABSL_ASSIGN_OR_RETURN(
-        converted[instruction],
-        ConvertFusionInstruction(*instruction, operands, *graph_op.getBody()));
-  }
-
-  mlir::Value result = converted[source.root_instruction()];
-  tir::ResultsOp::create(builder, result);
-
-  // Remove dead operations (scalar constants).
-  for (mlir::Operation& op : llvm::make_early_inc_range(graph_op.getOps())) {
-    if (mlir::isOpTriviallyDead(&op)) {
-      op.erase();
-    }
-  }
-
-  if (mlir::failed(mlir::verify(graph_op))) {
-    return absl::InternalError(absl::StrCat(
-        "Verification failed for TensorIR graph: ", source.name()));
-  }
-
-  std::move(cleanup).Cancel();
-  return graph_op;
-}
-
-// Creates a new MLIR module containing the converted TensorIR graph for
-// `source`.
-absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ConvertFusionComputation(
-    const HloComputation& source, mlir::MLIRContext* context) {
-  mlir::Location location = mlir::UnknownLoc::get(context);
-  mlir::OwningOpRef<mlir::ModuleOp> module =
-      llvm_ir::CreateMlirModuleOp(location);
-  ABSL_RETURN_IF_ERROR(ConvertFusionComputation(source, *module).status());
-  return module;
-}
-
-}  // namespace xla::gpu::tensor_ir
-
-// NOLINTEND(clang-diagnostic-pre-c++20-compat)
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/conversion.h b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/conversion.h
deleted file mode 100644
index eb0ffc1f2f4..00000000000
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/conversion.h
+++ /dev/null
@@ -1,40 +0,0 @@
-/* Copyright 2026 The OpenXLA Authors.
-
-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.
-==============================================================================*/
-
-#ifndef XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_CONVERSION_H_
-#define XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_CONVERSION_H_
-
-#include "tensor_ir/Dialect/TensorIR.h"
-#include "absl/status/statusor.h"
-#include "mlir/IR/BuiltinOps.h"
-#include "mlir/IR/MLIRContext.h"
-#include "mlir/IR/OwningOpRef.h"
-#include "xla/hlo/ir/hlo_computation.h"
-
-namespace xla::gpu::tensor_ir {
-
-// Converts an HLO fusion computation into a TensorIR `GraphOp` appended
-// to the body of `target`. If conversion fails, `target` is unmodified.
-absl::StatusOr<mlir::nv_tensor_ir::GraphOp> ConvertFusionComputation(
-    const HloComputation& source, mlir::ModuleOp target);
-
-// Creates a new MLIR module containing the converted TensorIR graph for
-// `source`.
-absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ConvertFusionComputation(
-    const HloComputation& source, mlir::MLIRContext* context);
-
-}  // namespace xla::gpu::tensor_ir
-
-#endif  // XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_CONVERSION_H_
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/custom_kernel.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/custom_kernel.cc
new file mode 100644
index 00000000000..7a4b5830e23
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/custom_kernel.cc
@@ -0,0 +1,99 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/gpu/codegen/tensor_ir/custom_kernel.h"
+
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <vector>
+
+#include "tensor_ir/Runtime/CudaTile/CudaTileRuntimeKernel.h"
+#include "tensor_ir/Runtime/CudaTile/KernelLaunchHelpers.h"
+#include "tensor_ir/Runtime/IRuntimeKernel.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/Support/Casting.h"
+#include "xla/service/gpu/kernel_reuse_cache.h"
+#include "xla/service/gpu/launch_dimensions.h"
+#include "xla/stream_executor/launch_dim.h"
+
+namespace xla::gpu::tensor_ir {
+
+absl::StatusOr<KernelReuseCache::Entry> MakeKernelCacheEntry(
+    const ::tensor_ir::rt::IRuntimeKernel& kernel) {
+  const auto* cuda_tile_kernel_ptr =
+      llvm::dyn_cast<::tensor_ir::rt::CudaTileRuntimeKernel>(&kernel);
+  if (cuda_tile_kernel_ptr == nullptr) {
+    return absl::InternalError(
+        absl::StrCat("TensorIR kernel '", kernel.name(),
+                     "' is not a CudaTile runtime kernel"));
+  }
+  const ::tensor_ir::rt::CudaTileRuntimeKernel& cuda_tile_kernel =
+      *cuda_tile_kernel_ptr;
+
+  if (!cuda_tile_kernel.hasDeviceCode()) {
+    return absl::InternalError(absl::StrCat(
+        "TensorIR kernel '", cuda_tile_kernel.name(), "' has no device code"));
+  }
+  if (cuda_tile_kernel.hasTileIRBytecode()) {
+    // The buffer holds TileIR bytecode rather than an assembled cubin. That is
+    // fine as far as loading goes: `cuModuleLoadData`, which StreamExecutor
+    // calls to load a cubin, documents "Tile IR data" alongside cubin, fatbin
+    // and PTX as an accepted image, and JITs it. But that JIT lives in the
+    // driver and needs r610 or newer, so say so once per process -- this is a
+    // property of the installation, not of any one kernel.
+    LOG_FIRST_N(WARNING, 1)
+        << "TensorIR kernel '" << cuda_tile_kernel.name()
+        << "' is TileIR bytecode rather than an assembled cubin, because no "
+           "'tileiras' assembler was found; CUDA 13.3 or newer ships one. The "
+           "CUDA driver will JIT the bytecode when the kernel is loaded, which "
+           "needs a native r610 or newer driver.";
+  }
+
+  // The launch grid is fixed at compile time for everything XLA emits: the
+  // compiler only installs a runtime grid computer for dynamic shapes, and a
+  // fusion never has any.
+  std::optional<::tensor_ir::rt::GridDims> grid = cuda_tile_kernel.staticGrid();
+  if (!grid.has_value()) {
+    return absl::InternalError(
+        absl::StrCat("TensorIR kernel '", cuda_tile_kernel.name(),
+                     "' has a launch grid that is only known at run time"));
+  }
+
+  // CudaTile's launch ABI, from `CudaTileRuntimeKernel::launch`: one thread per
+  // block, no shared memory, and -- since XLA only produces static shapes, for
+  // which the compiler installs a `PointerOnlyArgPacker` -- one device pointer
+  // per argument, in order. That is exactly what
+  // `CreateSharedCubinCustomKernel` builds.
+  //
+  // The one thing lost is the `CU_CLUSTER_SCHEDULING_POLICY_SPREAD` launch
+  // attribute that `launch` sets. It is a scheduling hint for thread block
+  // clusters, and these kernels are launched without a cluster dimension.
+  llvm::ArrayRef<char> device_code = cuda_tile_kernel.deviceCode();
+  KernelReuseCache::Entry entry;
+  entry.kernel_name = cuda_tile_kernel.funcName();
+  entry.launch_dimensions = LaunchDimensions(
+      se::BlockDim(grid->x, grid->y, grid->z), se::ThreadDim(1, 1, 1));
+  entry.shmem_bytes = 0;
+  entry.binary = std::make_shared<const std::vector<uint8_t>>(
+      device_code.begin(), device_code.end());
+  return entry;
+}
+
+}  // namespace xla::gpu::tensor_ir
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/custom_kernel.h b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/custom_kernel.h
new file mode 100644
index 00000000000..8206c460a3a
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/custom_kernel.h
@@ -0,0 +1,39 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_CUSTOM_KERNEL_H_
+#define XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_CUSTOM_KERNEL_H_
+
+#include "tensor_ir/Runtime/IRuntimeKernel.h"
+#include "absl/status/statusor.h"
+#include "xla/service/gpu/kernel_reuse_cache.h"
+
+namespace xla::gpu::tensor_ir {
+
+// Extracts what is needed to launch a kernel compiled by the TensorIR CudaTile
+// backend with a `CustomKernelThunk`: the cubin, the entry point and the launch
+// grid. Going through the generic thunk rather than a bespoke one is what lets
+// XLA serialize the kernel and record it into a command buffer, and returning a
+// `KernelReuseCache::Entry` lets identical fusions share one compiled kernel.
+//
+// `kernel` must have come from a compiler created with
+// `CompilerBackend::CudaTile`, which only ever returns a
+// `CudaTileRuntimeKernel`.
+absl::StatusOr<KernelReuseCache::Entry> MakeKernelCacheEntry(
+    const ::tensor_ir::rt::IRuntimeKernel& kernel);
+
+}  // namespace xla::gpu::tensor_ir
+
+#endif  // XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_CUSTOM_KERNEL_H_
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/fusion.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/fusion.cc
new file mode 100644
index 00000000000..9316072c212
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/fusion.cc
@@ -0,0 +1,286 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/gpu/codegen/tensor_ir/fusion.h"
+
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "cuda_tile/Dialect/CudaTile/IR/Dialect.h"
+#include "tensor_ir/Compiler/Compiler.h"
+#include "tensor_ir/Compiler/CudaTile/CudaTileCompiler.h"
+#include "tensor_ir/Conversion/TensorToCudaTile/Options.h"
+#include "tensor_ir/Dialect/TensorIR.h"
+#include "tensor_ir/Dialect/TensorIRAttrs.h"
+#include "tensor_ir/Options/OptionsEnums.h"
+#include "tensor_ir/Utils/ComputeCapability.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/LogicalResult.h"
+#include "mlir/Dialect/Arith/IR/ArithDialect.h"
+#include "mlir/IR/Block.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OwningOpRef.h"
+#include "mlir/IR/Region.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Support/LogicalResult.h"
+#include "xla/backends/gpu/codegen/kernel_compiler.h"
+#include "xla/backends/gpu/codegen/kernels/custom_kernel.h"
+#include "xla/backends/gpu/codegen/kernels/ptx_custom_kernel.h"
+#include "xla/backends/gpu/codegen/tensor_ir/compilation_pipeline.h"
+#include "xla/backends/gpu/codegen/tensor_ir/custom_kernel.h"
+#include "xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir.h"
+#include "xla/backends/gpu/codegen/tensor_ir/layout_utils.h"
+#include "xla/backends/gpu/codegen/tensor_ir/support.h"
+#include "xla/backends/gpu/runtime/custom_kernel_thunk.h"
+#include "xla/backends/gpu/runtime/thunk.h"
+#include "xla/codegen/emitters/kernel_arguments.h"
+#include "xla/future.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/mlir/utils/error_util.h"
+#include "xla/service/gpu/backend_configs.pb.h"
+#include "xla/service/gpu/gpu_constants.h"
+#include "xla/service/gpu/ir_emission_utils.h"
+#include "xla/service/gpu/ir_emitter_context.h"
+#include "xla/service/gpu/kernel_reuse_cache.h"
+#include "xla/service/llvm_ir/llvm_util.h"
+
+namespace xla::gpu {
+
+namespace {
+
+// Imports, lowers and compiles the fusion to a cubin.
+absl::StatusOr<KernelReuseCache::Entry> CompileFusion(
+    IrEmitterContext& ir_emitter_context, const HloFusionInstruction& fusion,
+    const TensorIrFusionConfig& tensor_ir_config,
+    const emitters::KernelArguments& kernel_arguments) {
+  const HloComputation* computation = fusion.fused_instructions_computation();
+
+  // Create the MLIR context and module.
+  BorrowedMlirContext borrowed_context = ir_emitter_context.BorrowMlirContext();
+  mlir::MLIRContext& context = **borrowed_context;
+  context.loadDialect<mlir::nv_tensor_ir::TensorIRDialect,
+                      mlir::cuda_tile::CudaTileDialect,
+                      mlir::arith::ArithDialect>();
+
+  // `OwningOpRef` so that the module is destroyed when we leave this function;
+  // the MLIR context is borrowed from a pool and outlives it.
+  mlir::OwningOpRef<mlir::ModuleOp> module_ref =
+      llvm_ir::CreateMlirModuleOp(mlir::UnknownLoc::get(&context));
+  mlir::ModuleOp module = *module_ref;
+  ABSL_ASSIGN_OR_RETURN(
+      mlir::nv_tensor_ir::GraphOp graph_op,
+      tensor_ir::ImportAndLegalizeComputation(*computation, module));
+  {
+    mlir::BaseScopedDiagnosticHandler diagnostic_handler(&context);
+    if (llvm::failed(graph_op.verify())) {
+      return diagnostic_handler.Combine(absl::InternalError(
+          absl::StrCat("TensorIrFusion: invalid TensorIR graph for fusion: ",
+                       fusion.ToString(), ": ")));
+    }
+  }
+
+  const std::vector<emitters::KernelArgument>& kernel_args =
+      kernel_arguments.args();
+  // The graph has one argument per fused parameter, which the HLO verifier
+  // ties to the fusion's operands. A fusion is permitted to have operands that
+  // no fused parameter reads, though, in which case the counts disagree; fail
+  // loudly here rather than letting the alignments be attached off-by-one.
+  int64_t num_inputs = computation->num_parameters();
+  if (kernel_args.size() != num_inputs + 1) {
+    return absl::InternalError(absl::StrCat(
+        "TensorIrFusion: expected ", num_inputs + 1,
+        " kernel arguments (one per fused parameter plus a single output), "
+        "got ",
+        kernel_args.size(), " for fusion: ", fusion.ToString()));
+  }
+  // Layout strides were already attached by `ImportAndLegalizeComputation`;
+  // alignments need buffer assignment and so can only be added here.
+  ABSL_RETURN_IF_ERROR(tensor_ir::AttachBufferAlignments(graph_op, kernel_args));
+
+  // Run the conversion pipeline.
+  llvm::SmallVector<int32_t> tile_size(tensor_ir_config.tile_size().begin(),
+                                       tensor_ir_config.tile_size().end());
+  // An unset `reduction_tile_size` is 0, which the pipeline rejects ("must be
+  // a positive power of two") even for fusions without a reduction. Fall back
+  // to the pipeline's own default instead.
+  int64_t reduction_tile_size = tensor_ir_config.reduction_tile_size();
+  if (reduction_tile_size == 0) {
+    reduction_tile_size = mlir::nv_tensor_ir::kDefaultReductionTileSize;
+  }
+  mlir::nv_tensor_ir::TensorToCudaTilePipelineOptions pipeline_options =
+      tensor_ir::GetPipelineOptions(
+          ir_emitter_context.gpu_device_info().gpu_compute_capability(),
+          tile_size, reduction_tile_size);
+
+  // Note: the graph is *not* lowered here. `ICompiler::compile()` runs the
+  // TensorIR-to-CudaTile pipeline itself and requires the module to still
+  // contain exactly one (unlowered) `nv_tensor_ir.graph`; pre-lowering it here
+  // makes the compiler reject the module.
+
+  // The CudaTile bytecode writer only accepts `DILocAttr`, `FileLineColLoc`,
+  // `UnknownLoc` and `CallSiteLoc` locations, while the HLO importer attaches
+  // `NameLoc`s derived from HLO instruction and parameter names. Drop them,
+  // otherwise writing the kernel's TileIR bytecode fails. Block arguments need
+  // the same treatment: the lowering propagates their locations onto the ops
+  // it creates.
+  mlir::Location unknown_loc = mlir::UnknownLoc::get(&context);
+  module.walk([&](mlir::Operation* op) {
+    op->setLoc(unknown_loc);
+    for (mlir::Region& region : op->getRegions()) {
+      for (mlir::Block& block : region) {
+        for (mlir::BlockArgument arg : block.getArguments()) {
+          arg.setLoc(unknown_loc);
+        }
+      }
+    }
+  });
+
+  // Compile the kernel.
+  using ::mlir::nv_tensor_ir::ArchPortability;
+  using ::mlir::nv_tensor_ir::backend::cuda_tile::CudaTileArtifactKind;
+  using ::mlir::nv_tensor_ir::backend::cuda_tile::CudaTileCompileOptions;
+  mlir::FailureOr<mlir::nv_tensor_ir::SmTarget> sm_target =
+      mlir::nv_tensor_ir::SmTarget::fromCc(pipeline_options.computeCapability,
+                                           ArchPortability::arch_conditional);
+  if (llvm::failed(sm_target)) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("TensorIrFusion: unsupported compute capability: ",
+                     pipeline_options.computeCapability));
+  }
+  CudaTileCompileOptions compile_options(*sm_target, pipeline_options.numCTAs,
+                                         pipeline_options.numWarps, tile_size);
+  // Ask the compiler for device code rather than TileIR bytecode. Cubin
+  // assembly requires an arch-conditional target, which `SmTarget::fromCc`
+  // only grants from sm_90 on; below that it silently clamps the portability
+  // down and `CudaTileCompileOptions::validate()` then rejects the request.
+  // `IsSupportedComputeCapability` above has already ruled those out.
+  if (sm_target->getPortability() != ArchPortability::arch_conditional) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("TensorIrFusion: unsupported compute capability: ",
+                     pipeline_options.computeCapability));
+  }
+  compile_options.artifactKind = CudaTileArtifactKind::Cubin;
+  // `ICompiler::compile()` rebuilds the pipeline options from the compile
+  // options (`makePipelineOptions`), so the reduction tile size has to travel
+  // through here as well; `pipeline_options.reductionTileSize` alone would be
+  // dropped. The field comes from third_party/tensor_ir/patches/hotfixes.patch.
+  compile_options.reductionTileSize = reduction_tile_size;
+
+  tensor_ir::SetUpTileIrAssembler();
+  std::unique_ptr<mlir::nv_tensor_ir::ICompiler> compiler =
+      mlir::nv_tensor_ir::ICompiler::create(
+          mlir::nv_tensor_ir::CompilerBackend::CudaTile);
+  auto kernel_or = compiler->compile(module, compile_options);
+  if (!kernel_or.ok()) {
+    return absl::InternalError(
+        absl::StrCat("TensorIrFusion: failed to compile fusion ", fusion.name(),
+                     ": ", kernel_or.status().message()));
+  }
+
+  return tensor_ir::MakeKernelCacheEntry(**kernel_or);
+}
+
+}  // namespace
+
+AsyncThunkSequence TensorIrFusion::Emit(
+    IrEmitterContext& ir_emitter_context,
+    const HloFusionInstruction& fusion) const {
+  // Verify the fusion is supported.
+  ABSL_ASSIGN_OR_RETURN(GpuBackendConfig gpu_backend_config,
+                   fusion.backend_config<GpuBackendConfig>());
+  const FusionBackendConfig& backend_config =
+      gpu_backend_config.fusion_backend_config();
+  if (backend_config.kind() != kTensorIrFusionKind) {
+    return absl::InternalError(absl::StrCat(
+        "TensorIrFusion: unsupported fusion kind: ", backend_config.kind()));
+  }
+  if (!backend_config.has_tensor_ir_fusion_config()) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "TensorIrFusion: missing tensor_ir_fusion_config for fusion: ",
+        fusion.ToString()));
+  }
+  const TensorIrFusionConfig& tensor_ir_config =
+      backend_config.tensor_ir_fusion_config();
+
+  if (auto decision = tensor_ir::IsSupportedComputeCapability(
+          ir_emitter_context.gpu_device_info().gpu_compute_capability());
+      !decision.IsAllowed()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("TensorIrFusion: ", decision.Explain()));
+  }
+
+  const HloComputation* computation = fusion.fused_instructions_computation();
+  if (auto decision = tensor_ir::IsSupportedFusionComputation(*computation);
+      !decision.IsAllowed()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("TensorIrFusion: ", decision.Explain()));
+  }
+
+  ABSL_ASSIGN_OR_RETURN(
+      emitters::KernelArguments kernel_arguments,
+      emitters::KernelArguments::Create(ir_emitter_context.buffer_assignment(),
+                                        GetDefaultBufferAlignment(), &fusion));
+
+  // Kernels are shared between identical fusions; the tiling is part of the
+  // key because the same computation can be compiled with different tilings.
+  std::string discriminator = absl::StrCat(
+      "TensorIrFusion:", absl::StrJoin(tensor_ir_config.tile_size(), ","), ":",
+      tensor_ir_config.reduction_tile_size());
+  auto [future_entry, cached] = ir_emitter_context.kernel_cache().GetWithStatus(
+      computation, kernel_arguments.args(), discriminator,
+      [&]() -> xla::Future<KernelReuseCache::Entry> {
+        return CompileFusion(ir_emitter_context, fusion, tensor_ir_config,
+                             kernel_arguments);
+      });
+  Thunk::ThunkInfo thunk_info = Thunk::ThunkInfo::WithProfileAnnotation(
+      &fusion, ir_emitter_context.GetNextThunkId());
+  return future_entry.Map(
+      [fusion_name = std::string(fusion.name()),
+       thunk_info = std::move(thunk_info),
+       kernel_arguments = std::move(kernel_arguments), cached = cached,
+       devices_in_process =
+           ir_emitter_context.gpu_topology().num_devices_per_process()](
+          const KernelReuseCache::Entry& entry) mutable
+          -> absl::StatusOr<ThunkSequence> {
+        if (cached) {
+          VLOG(3) << "Reuse: " << fusion_name << " -> " << entry.kernel_name;
+        }
+        ABSL_ASSIGN_OR_RETURN(
+            CustomKernel custom_kernel,
+            kernel::CreateSharedCubinCustomKernel(
+                entry.kernel_name, entry.binary, kernel_arguments.args().size(),
+                entry.launch_dimensions.block_counts(),
+                entry.launch_dimensions.thread_counts_per_block(),
+                entry.shmem_bytes));
+        return ThunkSequence::Of<CustomKernelThunk>(
+            thunk_info, std::move(custom_kernel), kernel_arguments,
+            devices_in_process);
+      });
+}
+
+}  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/fusion.h b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/fusion.h
new file mode 100644
index 00000000000..b75c3cddeee
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/fusion.h
@@ -0,0 +1,38 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_FUSION_H_
+#define XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_FUSION_H_
+
+#include "xla/backends/gpu/codegen/fusion_emitter.h"
+#include "xla/backends/gpu/runtime/thunk.h"
+#include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/service/gpu/ir_emitter_context.h"
+
+namespace xla::gpu {
+
+// Emits fusions with FusionBackendConfig.kind == kTensorIrFusionKind through
+// the TensorIR/CudaTile compilation pipeline.
+class TensorIrFusion : public FusionInterface {
+ public:
+  TensorIrFusion() = default;
+
+  AsyncThunkSequence Emit(IrEmitterContext& ir_emitter_context,
+                          const HloFusionInstruction& fusion) const final;
+};
+
+}  // namespace xla::gpu
+
+#endif  // XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_FUSION_H_
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/fusion_stub.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/fusion_stub.cc
new file mode 100644
index 00000000000..02a34529fd6
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/fusion_stub.cc
@@ -0,0 +1,33 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "absl/status/status.h"
+#include "xla/backends/gpu/codegen/tensor_ir/fusion.h"
+#include "xla/backends/gpu/runtime/thunk.h"
+#include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/service/gpu/ir_emitter_context.h"
+
+namespace xla::gpu {
+
+// Stands in for the TensorIR emitter in builds without CUDA: the tensor_ir
+// runtime and compiler libraries include cuda.h.
+AsyncThunkSequence TensorIrFusion::Emit(
+    IrEmitterContext& ir_emitter_context,
+    const HloFusionInstruction& fusion) const {
+  return absl::UnimplementedError(
+      "The TensorIR fusion emitter is only available in CUDA builds.");
+}
+
+}  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir.cc
new file mode 100644
index 00000000000..ef544713f62
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir.cc
@@ -0,0 +1,140 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir.h"
+
+#include <string>
+#include <unordered_map>
+
+#include "tensor_ir/Dialect/TensorIR.h"
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/LogicalResult.h"
+#include "mlir/Dialect/Arith/IR/ArithDialect.h"
+#include "mlir/Dialect/Func/IR/FuncDialect.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OwningOpRef.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Pass/PassManager.h"
+#include "stablehlo/dialect/StablehloOps.h"
+#include "xla/backends/gpu/codegen/tensor_ir/layout_utils.h"
+#include "xla/backends/gpu/codegen/tensor_ir/transforms/passes.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.h"
+#include "xla/mlir/utils/error_util.h"
+#include "xla/service/llvm_ir/llvm_util.h"
+
+namespace xla::gpu::tensor_ir {
+
+absl::StatusOr<mlir::nv_tensor_ir::GraphOp> ImportAndLegalizeComputation(
+    const HloComputation& computation, mlir::ModuleOp module) {
+  if (module == nullptr) {
+    return absl::InvalidArgumentError(
+        "ImportAndLegalizeComputation: module is null");
+  }
+  mlir::MLIRContext* context = module.getContext();
+  context->loadDialect<mlir::nv_tensor_ir::TensorIRDialect,
+                       mlir::stablehlo::StablehloDialect,
+                       mlir::func::FuncDialect, mlir::arith::ArithDialect>();
+
+  // Collect MLIR diagnostics so that they end up in the returned status rather
+  // than on stderr.
+  mlir::BaseScopedDiagnosticHandler diagnostic_handler(context);
+
+  mlir::SymbolTable symbol_table(module);
+  mlir::Builder builder(context);
+  std::unordered_map<const HloComputation*, mlir::func::FuncOp> function_map;
+
+  absl::StatusOr<mlir::func::FuncOp> func_op =
+      HloFunctionImporter::ImportAsFunc(computation, symbol_table,
+                                        &function_map, &builder,
+                                        /*is_main=*/true,
+                                        /*flatten_computation_args_result=*/
+                                        false);
+  if (!func_op.ok()) {
+    return diagnostic_handler.Combine(absl::Status(
+        func_op.status().code(),
+        absl::StrCat("Failed to import HLO computation '", computation.name(),
+                     "' as StableHLO: ", func_op.status().message())));
+  }
+
+  mlir::PassManager pass_manager(context);
+  pass_manager.addPass(
+      mlir::nv_tensor_ir::xla::createLegalizeStablehloToTensorIrPass());
+  if (llvm::failed(pass_manager.run(module))) {
+    return diagnostic_handler.Combine(absl::InternalError(
+        absl::StrCat("Failed to legalize StableHLO to TensorIR for "
+                     "computation '",
+                     computation.name(), "': ")));
+  }
+
+  llvm::SmallVector<mlir::nv_tensor_ir::GraphOp> graph_ops(
+      module.getOps<mlir::nv_tensor_ir::GraphOp>());
+  if (graph_ops.size() != 1) {
+    return diagnostic_handler.Combine(absl::InternalError(
+        absl::StrCat("Expected exactly one nv_tensor_ir.graph after legalizing "
+                     "computation '",
+                     computation.name(), "', got ", graph_ops.size())));
+  }
+
+  // Surface any error diagnostic that was reported without the pass manager
+  // signalling failure.
+  absl::Status diagnostics = diagnostic_handler.ConsumeStatus();
+  if (!diagnostics.ok()) {
+    return diagnostics;
+  }
+
+  // The computation is imported as the module's main function, so both the
+  // function and the graph legalized from it are named `main`. Rename the
+  // graph after the fact: the name ends up as the compiled kernel's name, and
+  // calling every kernel `main` makes profiles and traces unattributable.
+  mlir::nv_tensor_ir::GraphOp graph_op = graph_ops.front();
+  graph_op.setSymName(
+      llvm_ir::SanitizeFunctionName(std::string(computation.name())));
+
+  // Non-default HLO layouts survive as `nv_tensor_ir.stride` attributes on the
+  // graph boundary. This has to happen here rather than in the emitter: the
+  // layout-propagation passes read these attributes and silently assume
+  // row-major when they are absent, so a graph without them tiles differently
+  // from one with them.
+  ABSL_RETURN_IF_ERROR(AttachLayoutStrides(graph_op, computation));
+  return graph_op;
+}
+
+absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ImportAndLegalizeComputation(
+    const HloComputation& computation, mlir::MLIRContext* context) {
+  if (context == nullptr) {
+    return absl::InvalidArgumentError(
+        "ImportAndLegalizeComputation: context is null");
+  }
+  mlir::OwningOpRef<mlir::ModuleOp> module =
+      llvm_ir::CreateMlirModuleOp(mlir::UnknownLoc::get(context));
+  // All the work, including diagnostic capture, happens in the overload above.
+  absl::StatusOr<mlir::nv_tensor_ir::GraphOp> graph_op =
+      ImportAndLegalizeComputation(computation, *module);
+  if (!graph_op.ok()) {
+    return graph_op.status();
+  }
+  return module;
+}
+
+}  // namespace xla::gpu::tensor_ir
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir.h b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir.h
new file mode 100644
index 00000000000..d2651486e31
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir.h
@@ -0,0 +1,50 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_HLO_TO_TENSOR_IR_H_
+#define XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_HLO_TO_TENSOR_IR_H_
+
+#include "tensor_ir/Dialect/TensorIR.h"
+#include "absl/status/statusor.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OwningOpRef.h"
+#include "xla/hlo/ir/hlo_computation.h"
+
+namespace xla::gpu::tensor_ir {
+
+// Imports `computation` into `module` as a StableHLO `func.func` (via
+// `HloFunctionImporter`) and legalizes it to a single `nv_tensor_ir.graph`
+// (via the `legalize-stablehlo-to-tensor-ir` pass). Returns the resulting
+// graph, which is owned by `module`.
+//
+// `module` must be empty of graphs; the caller is expected to pass a freshly
+// created module. Any MLIR diagnostic emitted during import or legalization is
+// folded into the returned status.
+//
+// Note: the computation is imported as the module's main function, but the
+// resulting graph is renamed after `computation` (sanitized for use as a
+// symbol name), so that the compiled kernel is attributable in profiles.
+absl::StatusOr<mlir::nv_tensor_ir::GraphOp> ImportAndLegalizeComputation(
+    const HloComputation& computation, mlir::ModuleOp module);
+
+// Same as above, but creates the containing module in `context` and returns
+// it. The returned module contains exactly one `nv_tensor_ir.graph`.
+absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> ImportAndLegalizeComputation(
+    const HloComputation& computation, mlir::MLIRContext* context);
+
+}  // namespace xla::gpu::tensor_ir
+
+#endif  // XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_HLO_TO_TENSOR_IR_H_
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir_test.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir_test.cc
new file mode 100644
index 00000000000..0793fb7198b
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir_test.cc
@@ -0,0 +1,289 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir.h"
+
+#include <cstdint>
+
+#include "tensor_ir/Dialect/TensorIR.h"
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Dialect/Func/IR/FuncDialect.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OwningOpRef.h"
+#include "mlir/IR/Types.h"
+#include "mlir/IR/Value.h"
+#include "mlir/IR/Visitors.h"
+#include "mlir/Support/LLVM.h"
+#include "stablehlo/dialect/StablehloOps.h"
+#include "xla/hlo/parser/hlo_parser.h"
+#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
+#include "xla/service/llvm_ir/llvm_util.h"
+
+namespace xla::gpu::tensor_ir {
+namespace {
+
+using ::absl_testing::IsOk;
+using ::absl_testing::StatusIs;
+using ::testing::HasSubstr;
+using ::testing::Not;
+
+namespace tir = ::mlir::nv_tensor_ir;
+
+class HloToTensorIrTest : public HloHardwareIndependentTestBase {
+ protected:
+  HloToTensorIrTest() {
+    context_.loadDialect<tir::TensorIRDialect, mlir::func::FuncDialect,
+                         mlir::stablehlo::StablehloDialect>();
+    module_ = llvm_ir::CreateMlirModuleOp(mlir::UnknownLoc::get(&context_));
+  }
+
+  mlir::MLIRContext context_;
+  mlir::OwningOpRef<mlir::ModuleOp> module_;
+};
+
+// Returns the number of ops of type `OpTy` nested inside `graph`.
+template <typename OpTy>
+int64_t CountOps(tir::GraphOp graph) {
+  int64_t count = 0;
+  graph.walk([&](OpTy op) { ++count; });
+  return count;
+}
+
+TEST_F(HloToTensorIrTest, ElementwiseAddF32) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule test_module
+
+    ENTRY main {
+      p0 = f32[4,8]{1,0} parameter(0)
+      p1 = f32[4,8]{1,0} parameter(1)
+      ROOT add = f32[4,8]{1,0} add(p0, p1)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto hlo_module,
+                       ParseAndReturnUnverifiedModule(kHloText));
+
+  ASSERT_OK_AND_ASSIGN(
+      tir::GraphOp graph,
+      ImportAndLegalizeComputation(*hlo_module->entry_computation(), *module_));
+
+  // The func.func has been rewritten into a single nv_tensor_ir.graph.
+  EXPECT_TRUE(module_->getOps<mlir::func::FuncOp>().empty());
+
+  auto graph_type = graph.getFunctionType();
+  ASSERT_EQ(graph_type.getNumInputs(), 2);
+  ASSERT_EQ(graph_type.getNumResults(), 1);
+  auto arg_type =
+      mlir::dyn_cast<mlir::RankedTensorType>(graph_type.getInput(0));
+  ASSERT_TRUE(arg_type != nullptr);
+  EXPECT_EQ(arg_type.getShape(), llvm::ArrayRef<int64_t>({4, 8}));
+  EXPECT_TRUE(arg_type.getElementType().isF32());
+
+  EXPECT_EQ(CountOps<tir::AddOp>(graph), 1);
+  EXPECT_EQ(CountOps<tir::ResultsOp>(graph), 1);
+}
+
+TEST_F(HloToTensorIrTest, ElementwiseMulS32UsesSignedIntegers) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule test_module
+
+    ENTRY main {
+      p0 = s32[4,8]{1,0} parameter(0)
+      p1 = s32[4,8]{1,0} parameter(1)
+      ROOT mul = s32[4,8]{1,0} multiply(p0, p1)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto hlo_module,
+                       ParseAndReturnUnverifiedModule(kHloText));
+
+  ASSERT_OK_AND_ASSIGN(
+      tir::GraphOp graph,
+      ImportAndLegalizeComputation(*hlo_module->entry_computation(), *module_));
+
+  auto graph_type = graph.getFunctionType();
+  ASSERT_EQ(graph_type.getNumInputs(), 2);
+  ASSERT_EQ(graph_type.getNumResults(), 1);
+
+  // The signedness bridge: StableHLO uses signless i32, nv_tensor_ir requires
+  // signed si32.
+  for (mlir::Type type : graph_type.getInputs()) {
+    auto tensor_type = mlir::dyn_cast<mlir::RankedTensorType>(type);
+    ASSERT_TRUE(tensor_type != nullptr);
+    EXPECT_TRUE(tensor_type.getElementType().isSignedInteger(32))
+        << "expected si32 element type, got signless/unsigned integer";
+  }
+  auto result_type =
+      mlir::dyn_cast<mlir::RankedTensorType>(graph_type.getResult(0));
+  ASSERT_TRUE(result_type != nullptr);
+  EXPECT_TRUE(result_type.getElementType().isSignedInteger(32));
+
+  // The block arguments must be converted as well.
+  for (mlir::BlockArgument arg : graph.getGraphBody().front().getArguments()) {
+    auto tensor_type = mlir::dyn_cast<mlir::RankedTensorType>(arg.getType());
+    ASSERT_TRUE(tensor_type != nullptr);
+    EXPECT_TRUE(tensor_type.getElementType().isSignedInteger(32));
+  }
+
+  EXPECT_EQ(CountOps<tir::MulOp>(graph), 1);
+}
+
+TEST_F(HloToTensorIrTest, GraphIsNamedAfterTheComputation) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule test_module
+
+    ENTRY tensor_ir_fusion {
+      p0 = f32[4]{0} parameter(0)
+      ROOT n = f32[4]{0} negate(p0)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto hlo_module,
+                       ParseAndReturnUnverifiedModule(kHloText));
+
+  ASSERT_OK_AND_ASSIGN(
+      tir::GraphOp graph,
+      ImportAndLegalizeComputation(*hlo_module->entry_computation(), *module_));
+
+  EXPECT_EQ(graph.getSymName().str(), "tensor_ir_fusion");
+}
+
+TEST_F(HloToTensorIrTest, GraphNameIsSanitized) {
+  // HLO names may contain '.' and '-', which are not valid in a kernel symbol.
+  constexpr absl::string_view kHloText = R"(
+    HloModule test_module
+
+    ENTRY fused-computation.1 {
+      p0 = f32[4]{0} parameter(0)
+      ROOT n = f32[4]{0} negate(p0)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto hlo_module,
+                       ParseAndReturnUnverifiedModule(kHloText));
+
+  ASSERT_OK_AND_ASSIGN(
+      tir::GraphOp graph,
+      ImportAndLegalizeComputation(*hlo_module->entry_computation(), *module_));
+
+  EXPECT_EQ(graph.getSymName().str(), "fused_computation_1");
+}
+
+TEST_F(HloToTensorIrTest, UnsupportedOpFailsAndNamesTheOp) {
+  // stablehlo.xor has no nv_tensor_ir equivalent.
+  constexpr absl::string_view kHloText = R"(
+    HloModule test_module
+
+    ENTRY main {
+      p0 = s32[4,8]{1,0} parameter(0)
+      p1 = s32[4,8]{1,0} parameter(1)
+      ROOT x = s32[4,8]{1,0} xor(p0, p1)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto hlo_module,
+                       ParseAndReturnUnverifiedModule(kHloText));
+
+  auto graph =
+      ImportAndLegalizeComputation(*hlo_module->entry_computation(), *module_);
+  EXPECT_THAT(graph, Not(IsOk()));
+  EXPECT_THAT(graph.status().message(), HasSubstr("stablehlo.xor"));
+}
+
+TEST_F(HloToTensorIrTest, TupleRootFailsCleanly) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule test_module
+
+    ENTRY main {
+      p0 = f32[4]{0} parameter(0)
+      p1 = f32[4]{0} parameter(1)
+      a = f32[4]{0} add(p0, p1)
+      m = f32[4]{0} multiply(p0, p1)
+      ROOT t = (f32[4]{0}, f32[4]{0}) tuple(a, m)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto hlo_module,
+                       ParseAndReturnUnverifiedModule(kHloText));
+
+  auto graph =
+      ImportAndLegalizeComputation(*hlo_module->entry_computation(), *module_);
+  EXPECT_THAT(graph, Not(IsOk()));
+  EXPECT_THAT(graph.status().message(), HasSubstr("stablehlo.tuple"));
+}
+
+TEST_F(HloToTensorIrTest, NullModuleIsRejected) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule test_module
+
+    ENTRY main {
+      p0 = f32[4]{0} parameter(0)
+      ROOT n = f32[4]{0} negate(p0)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto hlo_module,
+                       ParseAndReturnUnverifiedModule(kHloText));
+
+  EXPECT_THAT(ImportAndLegalizeComputation(*hlo_module->entry_computation(),
+                                           mlir::ModuleOp()),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+TEST_F(HloToTensorIrTest, CreatesAndOwnsTheModule) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule test_module
+
+    ENTRY main {
+      p0 = f32[4]{0} parameter(0)
+      ROOT n = f32[4]{0} negate(p0)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto hlo_module,
+                       ParseAndReturnUnverifiedModule(kHloText));
+
+  ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef<mlir::ModuleOp> module,
+                       ImportAndLegalizeComputation(
+                           *hlo_module->entry_computation(), &context_));
+
+  ASSERT_TRUE(module);
+  llvm::SmallVector<tir::GraphOp> graphs(module->getOps<tir::GraphOp>());
+  ASSERT_EQ(graphs.size(), 1);
+  EXPECT_EQ(CountOps<tir::NegOp>(graphs.front()), 1);
+  EXPECT_TRUE(module->getOps<mlir::func::FuncOp>().empty());
+}
+
+TEST_F(HloToTensorIrTest, NullContextIsRejected) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule test_module
+
+    ENTRY main {
+      p0 = f32[4]{0} parameter(0)
+      ROOT n = f32[4]{0} negate(p0)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto hlo_module,
+                       ParseAndReturnUnverifiedModule(kHloText));
+
+  EXPECT_THAT(ImportAndLegalizeComputation(*hlo_module->entry_computation(),
+                                           /*context=*/nullptr),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+}  // namespace
+}  // namespace xla::gpu::tensor_ir
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils.cc
new file mode 100644
index 00000000000..714ffac1a60
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils.cc
@@ -0,0 +1,158 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/gpu/codegen/tensor_ir/layout_utils.h"
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "tensor_ir/Dialect/TensorIR.h"
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_join.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/MLIRContext.h"
+#include "xla/codegen/emitters/kernel_arguments.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/layout_util.h"
+#include "xla/shape.h"
+
+namespace xla::gpu::tensor_ir {
+
+std::string ComputeStrideString(const Shape& shape) {
+  int64_t rank = shape.dimensions().size();
+  if (rank == 0) {
+    return "()";
+  }
+  std::vector<int64_t> strides(rank, 1);
+  int64_t stride = 1;
+  for (int64_t dim : shape.layout().minor_to_major()) {
+    strides[dim] = stride;
+    stride *= shape.dimensions(dim);
+  }
+  return absl::StrCat("(", absl::StrJoin(strides, ","), ")");
+}
+
+bool HasDefaultLayout(const Shape& shape) {
+  if (!shape.IsArray()) {
+    return false;
+  }
+  // Scalars have only one layout, and a shape without an explicit layout is
+  // treated as having the default one (what `LayoutUtil::SetToDefaultLayout`
+  // would assign); `shape.layout()` must not be called on such a shape.
+  if (shape.dimensions().empty() || !shape.has_layout()) {
+    return true;
+  }
+  return LayoutUtil::IsMonotonicWithDim0Major(shape.layout());
+}
+
+namespace {
+
+// Both attachment functions below target the single-result graphs that
+// `ImportAndLegalizeComputation` produces.
+absl::Status CheckGraphIsSingleResult(mlir::nv_tensor_ir::GraphOp graph,
+                                      absl::string_view caller) {
+  if (graph.getNumResults() != 1) {
+    return absl::InternalError(
+        absl::StrCat(caller, ": graph '", graph.getSymName().str(), "' has ",
+                     graph.getNumResults(), " results, expected exactly 1"));
+  }
+  return absl::OkStatus();
+}
+
+}  // namespace
+
+absl::Status AttachLayoutStrides(mlir::nv_tensor_ir::GraphOp graph,
+                                 const HloComputation& computation) {
+  ABSL_RETURN_IF_ERROR(CheckGraphIsSingleResult(graph, "AttachLayoutStrides"));
+
+  int64_t num_parameters = computation.num_parameters();
+  if (graph.getNumArguments() != num_parameters) {
+    return absl::InternalError(absl::StrCat(
+        "AttachLayoutStrides: graph '", graph.getSymName().str(), "' has ",
+        graph.getNumArguments(), " arguments, expected ", num_parameters,
+        " for computation '", computation.name(), "'"));
+  }
+
+  mlir::MLIRContext* context = graph.getContext();
+  mlir::StringAttr stride_attr_name = mlir::StringAttr::get(
+      context, mlir::nv_tensor_ir::TensorIRDialect::getStrideAttrName());
+
+  for (int64_t i = 0; i < num_parameters; ++i) {
+    const Shape& shape = computation.parameter_instruction(i)->shape();
+    if (!shape.IsArray()) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "AttachLayoutStrides: non-array parameter at index ", i,
+          " of computation '", computation.name(), "': ", shape.ToString()));
+    }
+    if (!HasDefaultLayout(shape)) {
+      graph.setArgAttr(
+          i, stride_attr_name,
+          mlir::StringAttr::get(context, ComputeStrideString(shape)));
+    }
+  }
+
+  const Shape& root_shape = computation.root_instruction()->shape();
+  if (!root_shape.IsArray()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("AttachLayoutStrides: non-array root of computation '",
+                     computation.name(), "': ", root_shape.ToString()));
+  }
+  if (!HasDefaultLayout(root_shape)) {
+    graph.setResultAttr(
+        0, stride_attr_name,
+        mlir::StringAttr::get(context, ComputeStrideString(root_shape)));
+  }
+
+  return absl::OkStatus();
+}
+
+absl::Status AttachBufferAlignments(
+    mlir::nv_tensor_ir::GraphOp graph,
+    absl::Span<const emitters::KernelArgument> kernel_args) {
+  ABSL_RETURN_IF_ERROR(CheckGraphIsSingleResult(graph, "AttachBufferAlignments"));
+
+  int64_t num_inputs = graph.getNumArguments();
+  if (kernel_args.size() != num_inputs + 1) {
+    return absl::InternalError(absl::StrCat(
+        "AttachBufferAlignments: expected ", num_inputs + 1,
+        " kernel arguments for graph '", graph.getSymName().str(), "' with ",
+        num_inputs, " arguments, but got ", kernel_args.size()));
+  }
+
+  mlir::MLIRContext* context = graph.getContext();
+  mlir::StringAttr alignment_attr_name = mlir::StringAttr::get(
+      context, mlir::nv_tensor_ir::TensorIRDialect::getAlignmentAttrName());
+  mlir::IntegerType i64_type = mlir::IntegerType::get(context, 64);
+
+  for (int64_t i = 0; i < num_inputs; ++i) {
+    graph.setArgAttr(
+        i, alignment_attr_name,
+        mlir::IntegerAttr::get(i64_type, kernel_args[i].alignment()));
+  }
+  graph.setResultAttr(
+      0, alignment_attr_name,
+      mlir::IntegerAttr::get(i64_type, kernel_args[num_inputs].alignment()));
+
+  return absl::OkStatus();
+}
+
+}  // namespace xla::gpu::tensor_ir
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils.h b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils.h
new file mode 100644
index 00000000000..25652cff8e8
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils.h
@@ -0,0 +1,71 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_LAYOUT_UTILS_H_
+#define XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_LAYOUT_UTILS_H_
+
+#include <string>
+
+#include "tensor_ir/Dialect/TensorIR.h"
+#include "absl/status/status.h"
+#include "absl/types/span.h"
+#include "xla/codegen/emitters/kernel_arguments.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/shape.h"
+
+namespace xla::gpu::tensor_ir {
+
+// Returns the CuTe stride tuple for `shape`'s layout, e.g. "(16,1)" for a
+// row-major 2-D shape with 16 columns. Strides are in dimension order
+// (dimension 0 first) and counted in elements, not bytes.
+// For rank-0 scalar shapes, this returns "()".
+std::string ComputeStrideString(const Shape& shape);
+
+// True when `shape` is an array with the default descending minor-to-major
+// layout (or with no explicit layout, which is treated as the default one), in
+// which case no stride attribute needs to be attached.
+bool HasDefaultLayout(const Shape& shape);
+
+// Attaches `nv_tensor_ir.stride` to every argument and to result 0 of `graph`
+// whose corresponding shape does not have the default layout.
+//
+// Strides are intrinsic to `computation`: they are derived purely from the
+// layouts of its parameters and its root, with no buffer assignment involved.
+// This is therefore part of importing a computation, and is called by
+// `ImportAndLegalizeComputation`, so that everyone who imports a computation
+// (the emitter, the autotuner, tools) sees the same graph.
+//
+// `graph` must have exactly one argument per computation parameter and exactly
+// one result; this mirrors what the importer produces.
+absl::Status AttachLayoutStrides(mlir::nv_tensor_ir::GraphOp graph,
+                                 const HloComputation& computation);
+
+// Attaches `nv_tensor_ir.alignment` to every argument and to result 0 of
+// `graph`.
+//
+// Alignments come from buffer assignment, so unlike strides they are a
+// property of the buffers the computation happens to be assigned rather than
+// of the computation itself. This can only run in the emitter, after kernel
+// arguments exist.
+//
+// `kernel_args` must hold one argument per graph argument plus one for the
+// single result.
+absl::Status AttachBufferAlignments(
+    mlir::nv_tensor_ir::GraphOp graph,
+    absl::Span<const emitters::KernelArgument> kernel_args);
+
+}  // namespace xla::gpu::tensor_ir
+
+#endif  // XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_LAYOUT_UTILS_H_
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils_test.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils_test.cc
new file mode 100644
index 00000000000..3eaa0c8301b
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils_test.cc
@@ -0,0 +1,414 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/gpu/codegen/tensor_ir/layout_utils.h"
+
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include "tensor_ir/Dialect/TensorIR.h"
+#include "tensor_ir/Support/TCutegen.h"
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/MLIRContext.h"
+#include "xla/codegen/emitters/kernel_arguments.h"
+#include "xla/hlo/ir/hlo_casting_utils.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/xla_data.pb.h"
+
+namespace xla::gpu::tensor_ir {
+namespace {
+
+using ::absl_testing::IsOk;
+using ::absl_testing::StatusIs;
+namespace tcg = ::mlir::nv_tensor_ir::tcutegen;
+
+TEST(LayoutUtilsTest, ComputeStrideStringRowMajor2D) {
+  Shape shape = ShapeUtil::MakeShapeWithDenseLayout(F32, {4, 16}, {1, 0});
+  EXPECT_EQ(ComputeStrideString(shape), "(16,1)");
+
+  auto stride = tcg::Stride::fromString(ComputeStrideString(shape));
+  ASSERT_TRUE(stride.has_value());
+  EXPECT_EQ(stride->toString(), "(16,1)");
+}
+
+TEST(LayoutUtilsTest, ComputeStrideStringColumnMajor2D) {
+  Shape shape = ShapeUtil::MakeShapeWithDenseLayout(F32, {4, 16}, {0, 1});
+  EXPECT_EQ(ComputeStrideString(shape), "(1,4)");
+
+  auto stride = tcg::Stride::fromString(ComputeStrideString(shape));
+  ASSERT_TRUE(stride.has_value());
+  EXPECT_EQ(stride->toString(), "(1,4)");
+}
+
+TEST(LayoutUtilsTest, ComputeStrideStringRank1) {
+  Shape shape = ShapeUtil::MakeShapeWithDenseLayout(F32, {8}, {0});
+  EXPECT_EQ(ComputeStrideString(shape), "(1)");
+
+  auto stride = tcg::Stride::fromString(ComputeStrideString(shape));
+  ASSERT_TRUE(stride.has_value());
+  EXPECT_EQ(stride->toString(), "(1)");
+}
+
+TEST(LayoutUtilsTest, ComputeStrideStringRank0Scalar) {
+  Shape shape = ShapeUtil::MakeShape(F32, {});
+  EXPECT_EQ(ComputeStrideString(shape), "()");
+
+  // Verify that "()" is parsed by tcg::Stride::fromString.
+  auto stride = tcg::Stride::fromString(ComputeStrideString(shape));
+  ASSERT_TRUE(stride.has_value());
+  EXPECT_EQ(stride->toString(), "()");
+}
+
+TEST(LayoutUtilsTest, ComputeStrideStringTransposed3D) {
+  // f32[2,3,4] with minor_to_major {0,2,1}
+  // dim 0 is most minor: stride = 1, product = 2
+  // dim 2 is next: stride = 2, product = 2 * 4 = 8
+  // dim 1 is most major: stride = 8, product = 8 * 3 = 24
+  // Strides in dim order (0, 1, 2): (1, 8, 2)
+  Shape shape = ShapeUtil::MakeShapeWithDenseLayout(F32, {2, 3, 4}, {0, 2, 1});
+  EXPECT_EQ(ComputeStrideString(shape), "(1,8,2)");
+
+  auto stride = tcg::Stride::fromString(ComputeStrideString(shape));
+  ASSERT_TRUE(stride.has_value());
+  EXPECT_EQ(stride->toString(), "(1,8,2)");
+}
+
+TEST(LayoutUtilsTest, ComputeStrideStringDegenerateDimOfSize1) {
+  // f32[1, 16] with minor_to_major {1, 0}: dim 1 stride = 1, dim 0 stride = 16
+  // -> "(16,1)"
+  Shape row_major = ShapeUtil::MakeShapeWithDenseLayout(F32, {1, 16}, {1, 0});
+  EXPECT_EQ(ComputeStrideString(row_major), "(16,1)");
+
+  // f32[1, 16] with minor_to_major {0, 1}: dim 0 stride = 1, dim 1 stride = 1
+  // -> "(1,1)"
+  Shape col_major = ShapeUtil::MakeShapeWithDenseLayout(F32, {1, 16}, {0, 1});
+  EXPECT_EQ(ComputeStrideString(col_major), "(1,1)");
+
+  auto stride = tcg::Stride::fromString(ComputeStrideString(col_major));
+  ASSERT_TRUE(stride.has_value());
+  EXPECT_EQ(stride->toString(), "(1,1)");
+}
+
+TEST(LayoutUtilsTest, HasDefaultLayoutTests) {
+  EXPECT_TRUE(HasDefaultLayout(ShapeUtil::MakeShape(F32, {})));
+  EXPECT_TRUE(
+      HasDefaultLayout(ShapeUtil::MakeShapeWithDenseLayout(F32, {8}, {0})));
+  EXPECT_TRUE(HasDefaultLayout(
+      ShapeUtil::MakeShapeWithDenseLayout(F32, {4, 16}, {1, 0})));
+  EXPECT_FALSE(HasDefaultLayout(
+      ShapeUtil::MakeShapeWithDenseLayout(F32, {4, 16}, {0, 1})));
+  EXPECT_TRUE(HasDefaultLayout(
+      ShapeUtil::MakeShapeWithDenseLayout(F32, {2, 3, 4}, {2, 1, 0})));
+  EXPECT_FALSE(HasDefaultLayout(
+      ShapeUtil::MakeShapeWithDenseLayout(F32, {2, 3, 4}, {0, 2, 1})));
+  EXPECT_FALSE(HasDefaultLayout(ShapeUtil::MakeTupleShape({})));
+
+  Shape no_layout = ShapeUtil::MakeShape(F32, {4, 16});
+  no_layout.clear_layout();
+  EXPECT_TRUE(HasDefaultLayout(no_layout));
+}
+
+// Builds a bare `nv_tensor_ir.graph` with the given argument/result types. The
+// body is left empty: these tests only look at boundary attributes.
+mlir::nv_tensor_ir::GraphOp MakeGraph(mlir::OpBuilder& b,
+                                      llvm::ArrayRef<mlir::Type> arg_types,
+                                      llvm::ArrayRef<mlir::Type> result_types) {
+  auto graph = mlir::nv_tensor_ir::GraphOp::create(
+      b, mlir::UnknownLoc::get(b.getContext()), "test_graph",
+      /*sym_visibility=*/nullptr, b.getFunctionType(arg_types, result_types),
+      /*arg_attrs=*/nullptr, /*res_attrs=*/nullptr);
+  graph.addEntryBlock();
+  return graph;
+}
+
+class AttachLayoutStridesTest : public HloHardwareIndependentTestBase {};
+
+TEST_F(AttachLayoutStridesTest, DefaultAndNonDefaultLayouts) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule test_module
+
+    fused_computation {
+      p0 = f32[4,16]{1,0} parameter(0)
+      p1 = f32[4,16]{0,1} parameter(1)
+      p2 = f32[] parameter(2)
+      b2 = f32[4,16]{0,1} broadcast(p2), dimensions={}
+      add0 = f32[4,16]{0,1} add(p0, p1)
+      ROOT add = f32[4,16]{0,1} add(add0, b2)
+    }
+
+    ENTRY entry {
+      p0 = f32[4,16]{1,0} parameter(0)
+      p1 = f32[4,16]{0,1} parameter(1)
+      p2 = f32[] parameter(2)
+      ROOT fusion = f32[4,16]{0,1} fusion(p0, p1, p2), kind=kCustom, calls=fused_computation
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto hlo_module, ParseAndReturnVerifiedModule(kHloText));
+  const HloComputation& computation =
+      *Cast<HloFusionInstruction>(
+           hlo_module->entry_computation()->root_instruction())
+           ->fused_instructions_computation();
+
+  mlir::MLIRContext context;
+  context.loadDialect<mlir::nv_tensor_ir::TensorIRDialect>();
+  mlir::OpBuilder b(&context);
+
+  auto t = mlir::RankedTensorType::get({4, 16}, b.getF32Type());
+  auto scalar_t = mlir::RankedTensorType::get({}, b.getF32Type());
+  auto graph = MakeGraph(b, {t, t, scalar_t}, {t});
+
+  ASSERT_THAT(AttachLayoutStrides(graph, computation), IsOk());
+
+  auto stride_attr_name =
+      mlir::nv_tensor_ir::TensorIRDialect::getStrideAttrName();
+
+  // Arg 0 is row-major, i.e. the default layout -> no stride attribute.
+  EXPECT_EQ(graph.getArgAttr(0, stride_attr_name), nullptr);
+
+  // Arg 1 is column-major -> stride "(1,4)".
+  auto a1_stride =
+      graph.getArgAttrOfType<mlir::StringAttr>(1, stride_attr_name);
+  ASSERT_TRUE(a1_stride != nullptr);
+  EXPECT_EQ(a1_stride.getValue().str(), "(1,4)");
+
+  // Arg 2 is a rank-0 scalar, which always counts as default -> no attribute.
+  EXPECT_EQ(graph.getArgAttr(2, stride_attr_name), nullptr);
+
+  // The root is column-major -> stride "(1,4)".
+  auto res_stride =
+      graph.getResultAttrOfType<mlir::StringAttr>(0, stride_attr_name);
+  ASSERT_TRUE(res_stride != nullptr);
+  EXPECT_EQ(res_stride.getValue().str(), "(1,4)");
+
+  graph.erase();
+}
+
+// The whole point of splitting strides out of the emitter is that a plain
+// computation, with no fusion instruction and no buffer assignment, is enough
+// to produce them.
+TEST_F(AttachLayoutStridesTest, WorksOnAComputationWithoutAFusion) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule test_module
+
+    ENTRY entry {
+      p0 = f32[8,16,32]{2,0,1} parameter(0)
+      ROOT out = f32[16,8,32]{2,1,0} bitcast(p0)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto hlo_module, ParseAndReturnVerifiedModule(kHloText));
+  const HloComputation& computation = *hlo_module->entry_computation();
+
+  mlir::MLIRContext context;
+  context.loadDialect<mlir::nv_tensor_ir::TensorIRDialect>();
+  mlir::OpBuilder b(&context);
+
+  auto in_t = mlir::RankedTensorType::get({8, 16, 32}, b.getF32Type());
+  auto out_t = mlir::RankedTensorType::get({16, 8, 32}, b.getF32Type());
+  auto graph = MakeGraph(b, {in_t}, {out_t});
+
+  ASSERT_THAT(AttachLayoutStrides(graph, computation), IsOk());
+
+  auto stride_attr_name =
+      mlir::nv_tensor_ir::TensorIRDialect::getStrideAttrName();
+  auto a0_stride =
+      graph.getArgAttrOfType<mlir::StringAttr>(0, stride_attr_name);
+  ASSERT_TRUE(a0_stride != nullptr);
+  EXPECT_EQ(a0_stride.getValue().str(), "(32,256,1)");
+  // The root is row-major -> no stride attribute.
+  EXPECT_EQ(graph.getResultAttr(0, stride_attr_name), nullptr);
+
+  graph.erase();
+}
+
+TEST_F(AttachLayoutStridesTest, ErrorPathGraphNumArgumentsMismatch) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule test_module
+
+    ENTRY entry {
+      p0 = f32[4,16]{1,0} parameter(0)
+      ROOT id = f32[4,16]{1,0} copy(p0)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto hlo_module, ParseAndReturnVerifiedModule(kHloText));
+
+  mlir::MLIRContext context;
+  context.loadDialect<mlir::nv_tensor_ir::TensorIRDialect>();
+  mlir::OpBuilder b(&context);
+
+  // Graph has 2 arguments, but the computation has 1 parameter.
+  auto t0 = mlir::RankedTensorType::get({4, 16}, b.getF32Type());
+  auto graph = MakeGraph(b, {t0, t0}, {t0});
+
+  EXPECT_THAT(AttachLayoutStrides(graph, *hlo_module->entry_computation()),
+              StatusIs(absl::StatusCode::kInternal));
+
+  graph.erase();
+}
+
+TEST_F(AttachLayoutStridesTest, ErrorPathTupleParameter) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule test_module
+
+    ENTRY entry {
+      p0 = (f32[4], f32[4]) parameter(0)
+      gte0 = f32[4] get-tuple-element(p0), index=0
+      ROOT id = f32[4] copy(gte0)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto hlo_module, ParseAndReturnVerifiedModule(kHloText));
+
+  mlir::MLIRContext context;
+  context.loadDialect<mlir::nv_tensor_ir::TensorIRDialect>();
+  mlir::OpBuilder b(&context);
+
+  auto t0 = mlir::RankedTensorType::get({4}, b.getF32Type());
+  auto graph = MakeGraph(b, {t0}, {t0});
+
+  EXPECT_THAT(AttachLayoutStrides(graph, *hlo_module->entry_computation()),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+
+  graph.erase();
+}
+
+class AttachBufferAlignmentsTest : public HloHardwareIndependentTestBase {};
+
+TEST_F(AttachBufferAlignmentsTest, AttachesAlignmentToArgsAndResult) {
+  mlir::MLIRContext context;
+  context.loadDialect<mlir::nv_tensor_ir::TensorIRDialect>();
+  mlir::OpBuilder b(&context);
+
+  auto t = mlir::RankedTensorType::get({4, 16}, b.getF32Type());
+  auto graph = MakeGraph(b, {t, t}, {t});
+
+  Shape shape = ShapeUtil::MakeShapeWithDenseLayout(F32, {4, 16}, {1, 0});
+  std::vector<emitters::KernelArgument> kernel_args;
+  kernel_args.reserve(3);
+  for (int64_t alignment : {128, 64, 256}) {
+    kernel_args.emplace_back(shape);
+    kernel_args.back().set_alignment(alignment);
+  }
+
+  ASSERT_THAT(AttachBufferAlignments(graph, kernel_args), IsOk());
+
+  auto alignment_attr_name =
+      mlir::nv_tensor_ir::TensorIRDialect::getAlignmentAttrName();
+  auto a0 = graph.getArgAttrOfType<mlir::IntegerAttr>(0, alignment_attr_name);
+  ASSERT_TRUE(a0 != nullptr);
+  EXPECT_EQ(a0.getInt(), 128);
+  auto a1 = graph.getArgAttrOfType<mlir::IntegerAttr>(1, alignment_attr_name);
+  ASSERT_TRUE(a1 != nullptr);
+  EXPECT_EQ(a1.getInt(), 64);
+  auto res =
+      graph.getResultAttrOfType<mlir::IntegerAttr>(0, alignment_attr_name);
+  ASSERT_TRUE(res != nullptr);
+  EXPECT_EQ(res.getInt(), 256);
+
+  graph.erase();
+}
+
+TEST_F(AttachBufferAlignmentsTest, ErrorPathKernelArgsCountMismatch) {
+  mlir::MLIRContext context;
+  context.loadDialect<mlir::nv_tensor_ir::TensorIRDialect>();
+  mlir::OpBuilder b(&context);
+
+  auto t0 = mlir::RankedTensorType::get({4, 16}, b.getF32Type());
+  auto graph = MakeGraph(b, {t0}, {t0});
+
+  // Pass 1 argument instead of 2 (1 input + 1 output).
+  std::vector<emitters::KernelArgument> kernel_args;
+  kernel_args.emplace_back(
+      ShapeUtil::MakeShapeWithDenseLayout(F32, {4, 16}, {1, 0}));
+
+  EXPECT_THAT(AttachBufferAlignments(graph, kernel_args),
+              StatusIs(absl::StatusCode::kInternal));
+
+  graph.erase();
+}
+
+// Strides are attached during import and alignments later, in the emitter.
+// Both write into the same per-argument attribute dictionary, so verify that
+// the second call does not drop what the first one wrote.
+TEST_F(AttachBufferAlignmentsTest, ComposesWithLayoutStrides) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule test_module
+
+    ENTRY entry {
+      p0 = f32[4,16]{0,1} parameter(0)
+      ROOT id = f32[4,16]{0,1} copy(p0)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto hlo_module, ParseAndReturnVerifiedModule(kHloText));
+
+  mlir::MLIRContext context;
+  context.loadDialect<mlir::nv_tensor_ir::TensorIRDialect>();
+  mlir::OpBuilder b(&context);
+
+  auto t = mlir::RankedTensorType::get({4, 16}, b.getF32Type());
+  auto graph = MakeGraph(b, {t}, {t});
+
+  Shape shape = ShapeUtil::MakeShapeWithDenseLayout(F32, {4, 16}, {0, 1});
+  std::vector<emitters::KernelArgument> kernel_args;
+  kernel_args.reserve(2);
+  for (int64_t alignment : {128, 256}) {
+    kernel_args.emplace_back(shape);
+    kernel_args.back().set_alignment(alignment);
+  }
+
+  ASSERT_THAT(AttachLayoutStrides(graph, *hlo_module->entry_computation()),
+              IsOk());
+  ASSERT_THAT(AttachBufferAlignments(graph, kernel_args), IsOk());
+
+  auto alignment_attr_name =
+      mlir::nv_tensor_ir::TensorIRDialect::getAlignmentAttrName();
+  auto stride_attr_name =
+      mlir::nv_tensor_ir::TensorIRDialect::getStrideAttrName();
+
+  auto a0_stride =
+      graph.getArgAttrOfType<mlir::StringAttr>(0, stride_attr_name);
+  ASSERT_TRUE(a0_stride != nullptr);
+  EXPECT_EQ(a0_stride.getValue().str(), "(1,4)");
+  auto a0_align =
+      graph.getArgAttrOfType<mlir::IntegerAttr>(0, alignment_attr_name);
+  ASSERT_TRUE(a0_align != nullptr);
+  EXPECT_EQ(a0_align.getInt(), 128);
+
+  auto res_stride =
+      graph.getResultAttrOfType<mlir::StringAttr>(0, stride_attr_name);
+  ASSERT_TRUE(res_stride != nullptr);
+  EXPECT_EQ(res_stride.getValue().str(), "(1,4)");
+  auto res_align =
+      graph.getResultAttrOfType<mlir::IntegerAttr>(0, alignment_attr_name);
+  ASSERT_TRUE(res_align != nullptr);
+  EXPECT_EQ(res_align.getInt(), 256);
+
+  graph.erase();
+}
+
+}  // namespace
+}  // namespace xla::gpu::tensor_ir
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.cc
index 7e26284faae..be878071930 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.cc
@@ -23,6 +23,8 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/primitive_util.h"
+#include "xla/stream_executor/cuda/cuda_compute_capability.h"
+#include "xla/stream_executor/device_description.h"
 #include "xla/xla_data.pb.h"
 
 namespace xla::gpu::tensor_ir {
@@ -264,6 +266,21 @@ CodegenDecision IsInstructionSupportedForFusion(const HloInstruction& instr) {
       break;
     }
 
+    // HLO allows scalar `min`/`max` bounds for clamp and a scalar predicate
+    // for select, but the TensorIR ops require all operands to have the same
+    // shape.
+    case HloOpcode::kClamp:
+    case HloOpcode::kSelect: {
+      for (const HloInstruction* operand : instr.operands()) {
+        if (operand->shape().dimensions() != instr.shape().dimensions()) {
+          return CodegenDecision::Forbid(
+              absl::StrCat("Unsupported scalar operand of ternary operation: ",
+                           instr.name()));
+        }
+      }
+      break;
+    }
+
     default: {
       break;
     }
@@ -271,4 +288,19 @@ CodegenDecision IsInstructionSupportedForFusion(const HloInstruction& instr) {
   return CodegenDecision::Allow();
 }
 
+CodegenDecision IsSupportedComputeCapability(
+    const se::GpuComputeCapability& cc) {
+  const auto* cuda_cc = cc.cuda_compute_capability();
+  if (cuda_cc == nullptr) {
+    return CodegenDecision::Forbid("TensorIR requires a CUDA GPU");
+  }
+  if (!cuda_cc->IsAtLeastHopper()) {
+    return CodegenDecision::Forbid(
+        absl::StrCat("TensorIR requires compute capability ",
+                     se::CudaComputeCapability::Hopper().ToString(),
+                     " or newer, got ", cuda_cc->ToString()));
+  }
+  return CodegenDecision::Allow();
+}
+
 }  // namespace xla::gpu::tensor_ir
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.h b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.h
index 88f3b25dc92..192b6dbba27 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.h
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.h
@@ -19,6 +19,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/service/decision.h"
+#include "xla/stream_executor/device_description.h"
 
 namespace xla::gpu::tensor_ir {
 
@@ -33,6 +34,14 @@ CodegenDecision IsSupportedFusionComputation(const HloComputation& comp);
 // fused computation.
 CodegenDecision IsInstructionSupportedForFusion(const HloInstruction& instr);
 
+// Returns `Decision::Allow` if `cc` is an architecture the TensorIR fusion
+// emitter can target, which means CUDA and Hopper or newer. CudaTile only
+// assembles a cubin for an arch-conditional target and refuses to make one
+// arch-conditional below sm_90; all the emitter could produce there is Tile IR
+// bytecode for the driver to JIT, which is a path we neither want nor test.
+CodegenDecision IsSupportedComputeCapability(
+    const se::GpuComputeCapability& cc);
+
 }  // namespace xla::gpu::tensor_ir
 
 #endif  // XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_SUPPORT_H_
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support_test.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support_test.cc
new file mode 100644
index 00000000000..2e1f5e99907
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support_test.cc
@@ -0,0 +1,58 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/gpu/codegen/tensor_ir/support.h"
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "xla/stream_executor/cuda/cuda_compute_capability.h"
+#include "xla/stream_executor/device_description.h"
+#include "xla/stream_executor/rocm/rocm_compute_capability.h"
+
+namespace xla::gpu::tensor_ir {
+namespace {
+
+using ::testing::HasSubstr;
+
+TEST(IsSupportedComputeCapabilityTest, AllowsHopper) {
+  EXPECT_TRUE(IsSupportedComputeCapability(
+                  se::GpuComputeCapability(se::CudaComputeCapability::Hopper()))
+                  .IsAllowed());
+}
+
+TEST(IsSupportedComputeCapabilityTest, AllowsNewerThanHopper) {
+  EXPECT_TRUE(
+      IsSupportedComputeCapability(
+          se::GpuComputeCapability(se::CudaComputeCapability::Blackwell()))
+          .IsAllowed());
+}
+
+TEST(IsSupportedComputeCapabilityTest, ForbidsAmpere) {
+  CodegenDecision decision = IsSupportedComputeCapability(
+      se::GpuComputeCapability(se::CudaComputeCapability::Ampere()));
+  EXPECT_FALSE(decision.IsAllowed());
+  EXPECT_THAT(decision.Explain(), HasSubstr("9.0"));
+  EXPECT_THAT(decision.Explain(), HasSubstr("8.0"));
+}
+
+TEST(IsSupportedComputeCapabilityTest, ForbidsNonCuda) {
+  CodegenDecision decision = IsSupportedComputeCapability(
+      se::GpuComputeCapability(se::RocmComputeCapability("gfx942")));
+  EXPECT_FALSE(decision.IsAllowed());
+  EXPECT_THAT(decision.Explain(), HasSubstr("CUDA"));
+}
+
+}  // namespace
+}  // namespace xla::gpu::tensor_ir
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/BUILD b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/BUILD
index 01bef89604c..2964d0d1666 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/BUILD
@@ -1,4 +1,6 @@
 load("//xla:lit.bzl", "lit_test_suite")
+load("//xla/tests:build_defs.bzl", "xla_test")
+load("//xla/tsl:tsl.bzl", "if_google")
 
 package(
     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
@@ -22,3 +24,94 @@ lit_test_suite(
         "@llvm-project//llvm:FileCheck",
     ],
 )
+
+# Checks the shape of what the TensorIR emitter hands to the runtime. Compiles
+# only, never launches, so unlike the numerics tests below it does not need a
+# driver that can load a TensorIR kernel.
+xla_test(
+    name = "fusion_emitter_thunk_test",
+    srcs = ["fusion_emitter_thunk_test.cc"],
+    backend_tags = {
+        # The emitter compiles for whatever device the test runs on, and
+        # `IsSupportedComputeCapability` rejects anything before Hopper.
+        "gpu": ["requires-gpu-sm90"],
+    },
+    backends = ["gpu"],
+    tags = [
+        "cuda-only",
+        "no_mac",
+    ],
+    deps = [
+        "//xla:xla_proto_cc",
+        "//xla/backends/gpu/runtime:custom_kernel_thunk",
+        "//xla/backends/gpu/runtime:thunk",
+        "//xla/backends/gpu/runtime:thunk_proto_cc",
+        "//xla/backends/gpu/tests:gpu_pjrt_codegen_test",
+        "//xla/hlo/testlib:verified_hlo_module",
+        "//xla/service:executable",
+        "//xla/service/gpu:gpu_executable",
+        "//xla/service/gpu:launch_dimensions",
+        "//xla/stream_executor:launch_dim",
+        "//xla/tests:xla_internal_test_main",  # fixdeps: keep
+        "@com_google_absl//absl/base",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest",
+    ],
+)
+
+# End-to-end numerics tests for the TensorIR fusion emitter.
+#
+# This target is `manual` because it needs a newer toolchain than CI has, not
+# because it fails. It passes on an H100 with a CUDA 13.3 toolkit.
+#
+# Why CUDA 13.3. The emitter asks the compiler for a cubin, which it assembles
+# by running `tileiras` (found next to the toolkit's `nvcc`, in
+# `--xla_gpu_cuda_data_dir`, or on `PATH`). CUDA 13.0 ships no `tileiras`, and
+# the 13.1 and 13.2 releases lack the `--list-versions` probe TensorIR uses to
+# check bytecode compatibility, so 13.3 is the first release the compiler can
+# assemble with. Without an assembler the compiler hands back Tile IR bytecode
+# for the driver to JIT. TensorIR emits the CUDA 13.3 bytecode format by
+# default (`max(compatibility, 13.3)`, see the tensor_ir README), which needs
+# the matching r610 or newer driver; the r590 (CUDA 13.1) JIT only takes the
+# `compatibility` format, which the emitter does not request.
+#
+# Why Hopper. The compiler refuses to assemble a cubin below sm_90
+# (`CudaTileCompileOptions::validateDerived`).
+#
+# Why the driver matters even for an assembled cubin. Loading a Tile IR cubin
+# makes the driver dlopen() libnvidia-gpucomp.so; a driver install without it
+# fails with CUDA_ERROR_JIT_COMPILER_NOT_FOUND.
+#
+# TODO: drop the `manual` tag once CI builds against CUDA 13.3 or newer.
+xla_test(
+    name = "fusion_emitter_device_test",
+    srcs = ["fusion_emitter_device_test.cc"],
+    backend_tags = {
+        "gpu": ["requires-gpu-sm90"],
+    },
+    backends = ["gpu"],
+    tags = [
+        "cuda-only",
+        "manual",
+        "no_mac",
+    ] + if_google(["notap"]),
+    deps = [
+        "//xla:error_spec",
+        "//xla:literal",
+        "//xla:literal_util",
+        "//xla:xla_proto_cc",
+        "//xla/backends/gpu/tests:gpu_pjrt_codegen_test",
+        "//xla/hlo/testlib:verified_hlo_module",
+        "//xla/tests:hlo_interpreter_reference_mixin",
+        "//xla/tests:xla_internal_test_main",  # fixdeps: keep
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest",
+    ] + if_google([
+        # Forge drivers lack libnvidia-gpucomp.so and the compat library supplies it.
+        "@local_config_cuda//cuda:cuda_compat_tileir",
+    ]),
+)
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_convert.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_convert.hlo
index 813fbe9c3ef..68da5afe042 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_convert.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/elementwise_convert.hlo
@@ -40,8 +40,11 @@ test_float_to_unsigned {
 
 // -----
 
+// Converting to PRED is decomposed by the HLO importer into a comparison
+// against zero, so no `convert` op reaches TensorIR.
 // CHECK-LABEL: @test_float_to_pred
-// CHECK: convert %{{.*}} : tensor<8x16xf32> -> tensor<8x16xi1>
+// CHECK: %[[ZERO:.*]] = constant dense<0.000000e+00> : tensor<8x16xf32>
+// CHECK: cmp %{{.*}} une %[[ZERO]] : tensor<8x16xf32>
 
 test_float_to_pred {
   %p0 = f32[8,16] parameter(0)
@@ -81,7 +84,8 @@ test_signed_to_float {
 // -----
 
 // CHECK-LABEL: @test_signed_to_pred
-// CHECK: convert %{{.*}} : tensor<8x16xsi16> -> tensor<8x16xi1>
+// CHECK: %[[ZERO:.*]] = constant dense<0> : tensor<8x16xsi16>
+// CHECK: cmp %{{.*}} neq %[[ZERO]] : tensor<8x16xsi16>
 
 test_signed_to_pred {
   %p0 = s16[8,16] parameter(0)
@@ -121,7 +125,8 @@ test_unsigned_to_float {
 // -----
 
 // CHECK-LABEL: @test_unsigned_to_pred
-// CHECK: convert %{{.*}} : tensor<8x16xui16> -> tensor<8x16xi1>
+// CHECK: %[[ZERO:.*]] = constant dense<0> : tensor<8x16xui16>
+// CHECK: cmp %{{.*}} neq %[[ZERO]] : tensor<8x16xui16>
 
 test_unsigned_to_pred {
   %p0 = u16[8,16] parameter(0)
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/fusion_emitter_device_test.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/fusion_emitter_device_test.cc
new file mode 100644
index 00000000000..c6eaa3fd034
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/fusion_emitter_device_test.cc
@@ -0,0 +1,594 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include <cstdint>
+#include <limits>
+#include <memory>
+#include <utility>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/strings/string_view.h"
+#include "xla/backends/gpu/tests/gpu_pjrt_codegen_test.h"
+#include "xla/error_spec.h"
+#include "xla/hlo/testlib/verified_hlo_module.h"
+#include "xla/literal.h"
+#include "xla/literal_util.h"
+#include "xla/tests/hlo_interpreter_reference_mixin.h"
+#include "xla/xla.pb.h"
+
+namespace xla::gpu {
+namespace {
+
+constexpr ErrorSpec kExactMatch{/*aabs=*/0, /*arel=*/0};
+
+// End-to-end numerics tests for the TensorIR fusion emitter.
+//
+// A fusion is routed to the TensorIR emitter iff it is a `kCustom` fusion
+// whose `fusion_backend_config.kind` is `__tensorir` and which carries a
+// `tensor_ir_fusion_config` (see `HloFusionAnalysis::GetEmitterFusionKind` and
+// `TensorIrFusion::Emit`). The tile sizes inside that config are optional; if
+// they are omitted, the pipeline's own tile analysis picks them, which is what
+// most tests below rely on.
+//
+// All tests run without HLO passes so that the optimizer does not rewrite the
+// hand-written fusions and route them to a different emitter.
+class TensorIrEmitterTest
+    : public HloInterpreterReferenceMixin<GpuPjRtCodegenTest> {};
+
+TEST_F(TensorIrEmitterTest, ElementwiseAddF32) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = f32[8,16] parameter(0)
+  p1 = f32[8,16] parameter(1)
+  ROOT add = f32[8,16] add(p0, p1)
+}
+
+ENTRY main {
+  p0 = f32[8,16] parameter(0)
+  p1 = f32[8,16] parameter(1)
+  ROOT fusion = f32[8,16] fusion(p0, p1), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+// The tile size in the backend config is expressed in *iteration space*
+// terms, not in output-shape terms: layout propagation flattens the iteration
+// space of an elementwise fusion to rank 1, so a rank-2 tile is rejected with
+// "tile rank 2 vs iteration space rank 1".
+TEST_F(TensorIrEmitterTest, ElementwiseAddF32WithExplicitTileSize) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = f32[64,128] parameter(0)
+  p1 = f32[64,128] parameter(1)
+  ROOT add = f32[64,128] add(p0, p1)
+}
+
+ENTRY main {
+  p0 = f32[64,128] parameter(0)
+  p1 = f32[64,128] parameter(1)
+  ROOT fusion = f32[64,128] fusion(p0, p1), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{"tile_size":[256]}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+TEST_F(TensorIrEmitterTest, ElementwiseMulMaxF32) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = f32[8,16] parameter(0)
+  p1 = f32[8,16] parameter(1)
+  mul = f32[8,16] multiply(p0, p1)
+  ROOT max = f32[8,16] maximum(mul, p0)
+}
+
+ENTRY main {
+  p0 = f32[8,16] parameter(0)
+  p1 = f32[8,16] parameter(1)
+  ROOT fusion = f32[8,16] fusion(p0, p1), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+TEST_F(TensorIrEmitterTest, UnaryElementwiseF32) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = f32[8,16] parameter(0)
+  abs = f32[8,16] abs(p0)
+  ROOT neg = f32[8,16] negate(abs)
+}
+
+ENTRY main {
+  p0 = f32[8,16] parameter(0)
+  ROOT fusion = f32[8,16] fusion(p0), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+TEST_F(TensorIrEmitterTest, ConvertS32ToF32) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = s32[8,16] parameter(0)
+  ROOT convert = f32[8,16] convert(p0)
+}
+
+ENTRY main {
+  p0 = s32[8,16] parameter(0)
+  ROOT fusion = f32[8,16] fusion(p0), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+TEST_F(TensorIrEmitterTest, ConvertF32ToF16) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = f32[8,16] parameter(0)
+  ROOT convert = f16[8,16] convert(p0)
+}
+
+ENTRY main {
+  p0 = f32[8,16] parameter(0)
+  ROOT fusion = f16[8,16] fusion(p0), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+// Integer elementwise arithmetic exercises the signless -> signed bridge in
+// the TypeConverter: StableHLO uses signless `i32` while nv_tensor_ir requires
+// signed `si32`.
+TEST_F(TensorIrEmitterTest, ElementwiseS32) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = s32[8,16] parameter(0)
+  p1 = s32[8,16] parameter(1)
+  mul = s32[8,16] multiply(p0, p1)
+  ROOT sub = s32[8,16] subtract(mul, p1)
+}
+
+ENTRY main {
+  p0 = s32[8,16] parameter(0)
+  p1 = s32[8,16] parameter(1)
+  ROOT fusion = s32[8,16] fusion(p0, p1), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+// Signed division and remainder must not be lowered to their unsigned
+// counterparts, which is only observable for negative operands.
+TEST_F(TensorIrEmitterTest, SignedDivideAndRemainderS32) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = s32[2,4] parameter(0)
+  p1 = s32[2,4] parameter(1)
+  div = s32[2,4] divide(p0, p1)
+  rem = s32[2,4] remainder(p0, p1)
+  ROOT add = s32[2,4] add(div, rem)
+}
+
+ENTRY main {
+  p0 = s32[2,4] parameter(0)
+  p1 = s32[2,4] parameter(1)
+  ROOT fusion = s32[2,4] fusion(p0, p1), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(kHloText));
+  Literal lhs = LiteralUtil::CreateR2<int32_t>(
+      {{-7, -1, 0, 1}, {7, -2147483648, 13, -13}});
+  Literal rhs =
+      LiteralUtil::CreateR2<int32_t>({{2, -3, 5, -1}, {-2, 3, -4, 4}});
+  EXPECT_TRUE(
+      RunAndCompareNoHloPasses(std::move(module), {&lhs, &rhs}, kExactMatch));
+}
+
+TEST_F(TensorIrEmitterTest, SelectAndClampF32) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = f32[8,16] parameter(0)
+  p1 = f32[8,16] parameter(1)
+  gt = pred[8,16] compare(p0, p1), direction=GT
+  sel = f32[8,16] select(gt, p0, p1)
+  ROOT res = f32[8,16] maximum(sel, p1)
+}
+
+ENTRY main {
+  p0 = f32[8,16] parameter(0)
+  p1 = f32[8,16] parameter(1)
+  ROOT fusion = f32[8,16] fusion(p0, p1), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+// This case is load-bearing, not decoration: float `NE` must lower to the
+// *unordered* comparator `une`, because IEEE-754 makes `NaN != x` true for
+// every `x`, including `NaN != NaN`. The emitter used to lower it to the
+// ordered `one`, which returns false for NaN operands; that was a real bug we
+// fixed. Feeding an actual NaN (plus +/-Inf and -0.0) through the kernel and
+// comparing exactly is the only thing that catches a regression here; random
+// inputs never produce a NaN.
+TEST_F(TensorIrEmitterTest, CompareNeWithNaN) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = f32[2,4] parameter(0)
+  p1 = f32[2,4] parameter(1)
+  ROOT ne = pred[2,4] compare(p0, p1), direction=NE
+}
+
+ENTRY main {
+  p0 = f32[2,4] parameter(0)
+  p1 = f32[2,4] parameter(1)
+  ROOT fusion = pred[2,4] fusion(p0, p1), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(kHloText));
+  const float kNaN = std::numeric_limits<float>::quiet_NaN();
+  const float kInf = std::numeric_limits<float>::infinity();
+  // Row 0 pairs a NaN against itself, a number, an infinity and -0.0; all four
+  // comparisons must be true. Row 1 covers the ordered cases.
+  Literal lhs = LiteralUtil::CreateR2<float>(
+      {{kNaN, kNaN, kNaN, kNaN}, {1.0f, 1.0f, -0.0f, kInf}});
+  Literal rhs = LiteralUtil::CreateR2<float>(
+      {{kNaN, 1.0f, kInf, -0.0f}, {1.0f, 2.0f, 0.0f, kInf}});
+  EXPECT_TRUE(
+      RunAndCompareNoHloPasses(std::move(module), {&lhs, &rhs}, kExactMatch));
+}
+
+// The other float comparison directions must stay *ordered*: every comparison
+// against a NaN other than `NE` is false.
+TEST_F(TensorIrEmitterTest, CompareOrderedDirectionsWithNaN) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = f32[2,4] parameter(0)
+  p1 = f32[2,4] parameter(1)
+  eq = pred[2,4] compare(p0, p1), direction=EQ
+  lt = pred[2,4] compare(p0, p1), direction=LT
+  ge = pred[2,4] compare(p0, p1), direction=GE
+  or0 = pred[2,4] or(eq, lt)
+  ROOT res = pred[2,4] or(or0, ge)
+}
+
+ENTRY main {
+  p0 = f32[2,4] parameter(0)
+  p1 = f32[2,4] parameter(1)
+  ROOT fusion = pred[2,4] fusion(p0, p1), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(kHloText));
+  const float kNaN = std::numeric_limits<float>::quiet_NaN();
+  Literal lhs = LiteralUtil::CreateR2<float>(
+      {{kNaN, kNaN, 1.0f, 2.0f}, {-1.0f, 0.0f, 1.0f, 2.0f}});
+  Literal rhs = LiteralUtil::CreateR2<float>(
+      {{kNaN, 1.0f, kNaN, 2.0f}, {1.0f, -0.0f, 1.0f, 3.0f}});
+  EXPECT_TRUE(
+      RunAndCompareNoHloPasses(std::move(module), {&lhs, &rhs}, kExactMatch));
+}
+
+// Reductions whose body is a plain `add`/`max` lower to TensorIR's built-in
+// reduction modes.
+TEST_F(TensorIrEmitterTest, ReduceAddF32) {
+  constexpr absl::string_view kHloText = R"(
+add_f32 {
+  lhs = f32[] parameter(0)
+  rhs = f32[] parameter(1)
+  ROOT add = f32[] add(lhs, rhs)
+}
+
+fused_computation {
+  p0 = f32[16,64] parameter(0)
+  c0 = f32[] constant(0)
+  ROOT reduce = f32[16] reduce(p0, c0), dimensions={1}, to_apply=add_f32
+}
+
+ENTRY main {
+  p0 = f32[16,64] parameter(0)
+  ROOT fusion = f32[16] fusion(p0), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  // The reference interpreter sums sequentially while the kernel sums in a
+  // tree, so the results differ in the last few ulps.
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, ErrorSpec{1e-5, 1e-5}));
+}
+
+TEST_F(TensorIrEmitterTest, ReduceMaxF32) {
+  constexpr absl::string_view kHloText = R"(
+max_f32 {
+  lhs = f32[] parameter(0)
+  rhs = f32[] parameter(1)
+  ROOT max = f32[] maximum(lhs, rhs)
+}
+
+fused_computation {
+  p0 = f32[16,64] parameter(0)
+  c0 = f32[] constant(-inf)
+  ROOT reduce = f32[16] reduce(p0, c0), dimensions={1}, to_apply=max_f32
+}
+
+ENTRY main {
+  p0 = f32[16,64] parameter(0)
+  ROOT fusion = f32[16] fusion(p0), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+// A reduction body that is not a bare built-in mode (here: operands flipped)
+// lowers to `reduce_ud`, i.e. a user-defined reduction region.
+TEST_F(TensorIrEmitterTest, ReduceUserDefinedF32) {
+  constexpr absl::string_view kHloText = R"(
+add_flipped {
+  lhs = f32[] parameter(0)
+  rhs = f32[] parameter(1)
+  ROOT add = f32[] add(rhs, lhs)
+}
+
+fused_computation {
+  p0 = f32[16,64] parameter(0)
+  c0 = f32[] constant(0)
+  ROOT reduce = f32[16] reduce(p0, c0), dimensions={1}, to_apply=add_flipped
+}
+
+ENTRY main {
+  p0 = f32[16,64] parameter(0)
+  ROOT fusion = f32[16] fusion(p0), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, ErrorSpec{1e-5, 1e-5}));
+}
+
+// Integer reductions always lower to `reduce_ud`, and additionally exercise
+// the signless -> signed bridge inside the reduction region.
+TEST_F(TensorIrEmitterTest, ReduceUserDefinedS32) {
+  constexpr absl::string_view kHloText = R"(
+add_s32 {
+  lhs = s32[] parameter(0)
+  rhs = s32[] parameter(1)
+  ROOT add = s32[] add(lhs, rhs)
+}
+
+fused_computation {
+  p0 = s32[16,64] parameter(0)
+  c0 = s32[] constant(0)
+  ROOT reduce = s32[16] reduce(p0, c0), dimensions={1}, to_apply=add_s32
+}
+
+ENTRY main {
+  p0 = s32[16,64] parameter(0)
+  ROOT fusion = s32[16] fusion(p0), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+TEST_F(TensorIrEmitterTest, DotF32) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = f32[64,32] parameter(0)
+  p1 = f32[32,16] parameter(1)
+  ROOT dot = f32[64,16] dot(p0, p1),
+    lhs_contracting_dims={1}, rhs_contracting_dims={0}
+}
+
+ENTRY main {
+  p0 = f32[64,32] parameter(0)
+  p1 = f32[32,16] parameter(1)
+  ROOT fusion = f32[64,16] fusion(p0, p1), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, ErrorSpec{1e-3, 1e-3}));
+}
+
+TEST_F(TensorIrEmitterTest, DotWithBatchF32) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = f32[4,16,32] parameter(0)
+  p1 = f32[4,32,16] parameter(1)
+  ROOT dot = f32[4,16,16] dot(p0, p1),
+    lhs_batch_dims={0}, rhs_batch_dims={0},
+    lhs_contracting_dims={2}, rhs_contracting_dims={1}
+}
+
+ENTRY main {
+  p0 = f32[4,16,32] parameter(0)
+  p1 = f32[4,32,16] parameter(1)
+  ROOT fusion = f32[4,16,16] fusion(p0, p1), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, ErrorSpec{1e-3, 1e-3}));
+}
+
+TEST_F(TensorIrEmitterTest, BroadcastTransposeSlice) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = f32[8,16] parameter(0)
+  transpose = f32[16,8] transpose(p0), dimensions={1,0}
+  slice = f32[8,4] slice(transpose), slice={[0:8], [0:4]}
+  ROOT reshape = f32[32] reshape(slice)
+}
+
+ENTRY main {
+  p0 = f32[8,16] parameter(0)
+  ROOT fusion = f32[32] fusion(p0), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+// Non-default layouts must reach the kernel as stride attributes; see
+// `tensor_ir::AttachBoundaryAttributes`. The first operand is column major.
+TEST_F(TensorIrEmitterTest, NonDefaultOperandLayout) {
+  constexpr absl::string_view kHloText = R"(
+HloModule m, entry_computation_layout={(f32[8,16]{0,1},f32[8,16]{1,0})->f32[8,16]{1,0}}
+
+fused_computation {
+  p0 = f32[8,16]{0,1} parameter(0)
+  p1 = f32[8,16]{1,0} parameter(1)
+  ROOT add = f32[8,16]{1,0} add(p0, p1)
+}
+
+ENTRY main {
+  p0 = f32[8,16]{0,1} parameter(0)
+  p1 = f32[8,16]{1,0} parameter(1)
+  ROOT fusion = f32[8,16]{1,0} fusion(p0, p1), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+// Same, but the fusion output is column major as well.
+TEST_F(TensorIrEmitterTest, NonDefaultOutputLayout) {
+  constexpr absl::string_view kHloText = R"(
+HloModule m, entry_computation_layout={(f32[8,16]{1,0},f32[8,16]{1,0})->f32[8,16]{0,1}}
+
+fused_computation {
+  p0 = f32[8,16]{1,0} parameter(0)
+  p1 = f32[8,16]{1,0} parameter(1)
+  ROOT add = f32[8,16]{0,1} add(p0, p1)
+}
+
+ENTRY main {
+  p0 = f32[8,16]{1,0} parameter(0)
+  p1 = f32[8,16]{1,0} parameter(1)
+  ROOT fusion = f32[8,16]{0,1} fusion(p0, p1), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+// Rank-0 tensors need care: CudaTile's layout propagation rejects a rank-0
+// graph argument, graph result or constant, so the legalization pass rewrites
+// all three to rank 1. See `PromoteRank0Tensors`.
+TEST_F(TensorIrEmitterTest, ScalarParameterBroadcast) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = f32[256] parameter(0)
+  p1 = f32[] parameter(1)
+  b = f32[256] broadcast(p1), dimensions={}
+  ROOT a = f32[256] add(p0, b)
+}
+
+ENTRY main {
+  p0 = f32[256] parameter(0)
+  p1 = f32[] parameter(1)
+  ROOT fusion = f32[256] fusion(p0, p1), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+TEST_F(TensorIrEmitterTest, ScalarConstantBroadcast) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = f32[256] parameter(0)
+  c = f32[] constant(2)
+  b = f32[256] broadcast(c), dimensions={}
+  ROOT m = f32[256] multiply(p0, b)
+}
+
+ENTRY main {
+  p0 = f32[256] parameter(0)
+  ROOT fusion = f32[256] fusion(p0), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+// A full reduce, so the fusion root is rank 0.
+TEST_F(TensorIrEmitterTest, ReduceToScalar) {
+  constexpr absl::string_view kHloText = R"(
+add {
+  a = f32[] parameter(0)
+  b = f32[] parameter(1)
+  ROOT s = f32[] add(a, b)
+}
+
+fused_computation {
+  p0 = f32[256] parameter(0)
+  c0 = f32[] constant(0)
+  ROOT r = f32[] reduce(p0, c0), dimensions={0}, to_apply=add
+}
+
+ENTRY main {
+  p0 = f32[256] parameter(0)
+  ROOT fusion = f32[] fusion(p0), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, ErrorSpec{1e-5, 1e-5}));
+}
+
+// Acceptance criterion 4: the same kernels, executed from inside a command
+// buffer rather than launched directly. The thunk test next door proves the
+// conversion happens; this proves the recorded graph computes the same thing.
+class TensorIrCommandBufferTest : public TensorIrEmitterTest {
+ protected:
+  DebugOptions GetDebugOptionsForTest() const override {
+    DebugOptions debug_options = TensorIrEmitterTest::GetDebugOptionsForTest();
+    debug_options.clear_xla_gpu_enable_command_buffer();
+    debug_options.add_xla_gpu_enable_command_buffer(DebugOptions::FUSION);
+    // Otherwise a single-kernel graph is below the size threshold.
+    debug_options.set_xla_gpu_graph_min_graph_size(1);
+    return debug_options;
+  }
+};
+
+TEST_F(TensorIrCommandBufferTest, ElementwiseAddF32) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = f32[8,16] parameter(0)
+  p1 = f32[8,16] parameter(1)
+  ROOT add = f32[8,16] add(p0, p1)
+}
+
+ENTRY main {
+  p0 = f32[8,16] parameter(0)
+  p1 = f32[8,16] parameter(1)
+  ROOT fusion = f32[8,16] fusion(p0, p1), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+TEST_F(TensorIrCommandBufferTest, ReduceToScalar) {
+  constexpr absl::string_view kHloText = R"(
+add {
+  a = f32[] parameter(0)
+  b = f32[] parameter(1)
+  ROOT s = f32[] add(a, b)
+}
+
+fused_computation {
+  p0 = f32[256] parameter(0)
+  c0 = f32[] constant(0)
+  ROOT r = f32[] reduce(p0, c0), dimensions={0}, to_apply=add
+}
+
+ENTRY main {
+  p0 = f32[256] parameter(0)
+  ROOT fusion = f32[] fusion(p0), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, ErrorSpec{1e-5, 1e-5}));
+}
+
+}  // namespace
+}  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/fusion_emitter_thunk_test.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/fusion_emitter_thunk_test.cc
new file mode 100644
index 00000000000..9b54ee50ee5
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/fusion_emitter_thunk_test.cc
@@ -0,0 +1,182 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include <memory>
+#include <utility>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/base/casts.h"
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "xla/backends/gpu/runtime/custom_kernel_thunk.h"
+#include "xla/backends/gpu/runtime/thunk.h"
+#include "xla/backends/gpu/runtime/thunk.pb.h"
+#include "xla/backends/gpu/tests/gpu_pjrt_codegen_test.h"
+#include "xla/hlo/testlib/verified_hlo_module.h"
+#include "xla/service/executable.h"
+#include "xla/service/gpu/gpu_executable.h"
+#include "xla/service/gpu/launch_dimensions.h"
+#include "xla/stream_executor/launch_dim.h"
+#include "xla/xla.pb.h"
+
+namespace xla::gpu {
+namespace {
+
+using ::testing::ElementsAre;
+
+constexpr absl::string_view kAddF32 = R"(
+fused_computation {
+  p0 = f32[8,16] parameter(0)
+  p1 = f32[8,16] parameter(1)
+  ROOT add = f32[8,16] add(p0, p1)
+}
+
+ENTRY main {
+  p0 = f32[8,16] parameter(0)
+  p1 = f32[8,16] parameter(1)
+  ROOT fusion = f32[8,16] fusion(p0, p1), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+
+// Checks the shape of what the TensorIR emitter hands to the runtime. Unlike
+// the numerics tests next door these only compile and never launch, so they
+// need a GPU to be visible but not one whose driver can load the kernel.
+class TensorIrThunkTest : public GpuPjRtCodegenTest {
+ protected:
+  // Compiles without the HLO optimization passes, which would otherwise
+  // rewrite the hand-written fusion and route it to a different emitter.
+  absl::StatusOr<std::unique_ptr<Executable>> Compile(
+      absl::string_view hlo_text) {
+    ABSL_ASSIGN_OR_RETURN(std::unique_ptr<VerifiedHloModule> module,
+                     ParseAndReturnVerifiedModule(hlo_text));
+    return CompileToExecutable(std::move(module),
+                               /*run_optimization_passes=*/false);
+  }
+
+  // Compiles `hlo_text` and returns its only thunk.
+  absl::StatusOr<const Thunk*> CompileToSingleThunk(
+      absl::string_view hlo_text, std::unique_ptr<Executable>* executable) {
+    ABSL_ASSIGN_OR_RETURN(*executable, Compile(hlo_text));
+    auto* gpu_executable = absl::down_cast<GpuExecutable*>(executable->get());
+    const ThunkSequence& thunks = gpu_executable->thunk_executor().thunks();
+    if (thunks.size() != 1) {
+      return absl::InternalError(
+          absl::StrCat("expected exactly one thunk, got ", thunks.size()));
+    }
+    return thunks.front().get();
+  }
+
+  absl::StatusOr<const CustomKernelThunk*> CompileToSingleCustomKernelThunk(
+      absl::string_view hlo_text, std::unique_ptr<Executable>* executable) {
+    ABSL_ASSIGN_OR_RETURN(const Thunk* thunk,
+                     CompileToSingleThunk(hlo_text, executable));
+    const auto* custom_kernel_thunk =
+        dynamic_cast<const CustomKernelThunk*>(thunk);
+    if (custom_kernel_thunk == nullptr) {
+      return absl::InternalError(
+          absl::StrCat("expected a CustomKernelThunk, got ",
+                       Thunk::KindToString(thunk->kind())));
+    }
+    return custom_kernel_thunk;
+  }
+};
+
+TEST_F(TensorIrThunkTest, EmitsACustomKernelThunk) {
+  std::unique_ptr<Executable> executable;
+  ASSERT_OK_AND_ASSIGN(const CustomKernelThunk* thunk,
+                       CompileToSingleCustomKernelThunk(kAddF32, &executable));
+  EXPECT_EQ(thunk->kind(), Thunk::Kind::kCustomKernel);
+}
+
+// CudaTile's launch ABI: the grid comes from the compiler's tiling, every
+// block is a single thread, and nothing uses shared memory.
+TEST_F(TensorIrThunkTest, UsesTheCudaTileLaunchAbi) {
+  std::unique_ptr<Executable> executable;
+  ASSERT_OK_AND_ASSIGN(const CustomKernelThunk* thunk,
+                       CompileToSingleCustomKernelThunk(kAddF32, &executable));
+
+  LaunchDimensions launch = thunk->launch_dimensions();
+  EXPECT_EQ(launch.thread_counts_per_block().x, 1);
+  EXPECT_EQ(launch.thread_counts_per_block().y, 1);
+  EXPECT_EQ(launch.thread_counts_per_block().z, 1);
+  EXPECT_GE(launch.block_counts().x, 1);
+  EXPECT_EQ(thunk->shmem_bytes(), 0);
+}
+
+// The kernel takes one device pointer per buffer: the two inputs, then the
+// output.
+TEST_F(TensorIrThunkTest, PassesInputsThenOutput) {
+  std::unique_ptr<Executable> executable;
+  ASSERT_OK_AND_ASSIGN(const CustomKernelThunk* thunk,
+                       CompileToSingleCustomKernelThunk(kAddF32, &executable));
+
+  EXPECT_EQ(thunk->arguments().size(), 3);
+  EXPECT_THAT(thunk->written(), ElementsAre(false, false, true));
+}
+
+// The reason for using `CustomKernelThunk` rather than a bespoke thunk: the
+// kernel, its device code and its launch dimensions all round-trip through a
+// proto, which ahead-of-time compilation requires. A bespoke thunk would have
+// to reimplement all of this.
+TEST_F(TensorIrThunkTest, RoundTripsThroughAProto) {
+  std::unique_ptr<Executable> executable;
+  ASSERT_OK_AND_ASSIGN(const CustomKernelThunk* thunk,
+                       CompileToSingleCustomKernelThunk(kAddF32, &executable));
+
+  ASSERT_OK_AND_ASSIGN(ThunkProto proto, thunk->ToProto());
+  ASSERT_TRUE(proto.has_custom_kernel_thunk());
+
+  auto* gpu_executable = absl::down_cast<GpuExecutable*>(executable.get());
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<CustomKernelThunk> restored,
+                       CustomKernelThunk::FromProto(
+                           thunk->thunk_info(), proto.custom_kernel_thunk(),
+                           gpu_executable->allocations(),
+                           /*devices_in_process=*/1));
+
+  EXPECT_EQ(restored->custom_kernel_name(), thunk->custom_kernel_name());
+  EXPECT_EQ(restored->launch_dimensions(), thunk->launch_dimensions());
+  EXPECT_EQ(restored->shmem_bytes(), thunk->shmem_bytes());
+  EXPECT_EQ(restored->arguments().size(), thunk->arguments().size());
+  EXPECT_EQ(restored->written(), thunk->written());
+}
+
+// Using `CustomKernelThunk` also makes the kernel eligible for command
+// buffers. The bespoke thunk it replaced was not: the conversion pass only
+// knows the thunk kinds it enumerates.
+class TensorIrCommandBufferTest : public TensorIrThunkTest {
+ protected:
+  DebugOptions GetDebugOptionsForTest() const override {
+    DebugOptions debug_options = TensorIrThunkTest::GetDebugOptionsForTest();
+    debug_options.clear_xla_gpu_enable_command_buffer();
+    debug_options.add_xla_gpu_enable_command_buffer(DebugOptions::FUSION);
+    // Otherwise a single-kernel graph is below the size threshold.
+    debug_options.set_xla_gpu_graph_min_graph_size(1);
+    return debug_options;
+  }
+};
+
+TEST_F(TensorIrCommandBufferTest, RecordsIntoACommandBuffer) {
+  std::unique_ptr<Executable> executable;
+  ASSERT_OK_AND_ASSIGN(const Thunk* thunk,
+                       CompileToSingleThunk(kAddF32, &executable));
+  EXPECT_EQ(thunk->kind(), Thunk::Kind::kCommandBuffer);
+}
+
+}  // namespace
+}  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_bitcast.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_bitcast.hlo
index 079ed0680b1..7b36866ebc1 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_bitcast.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_bitcast.hlo
@@ -1,5 +1,12 @@
 // RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s
 
+// These tests check how non-default HLO layouts are turned into explicit
+// `transpose`/`reshape` ops, and into `nv_tensor_ir.stride` attributes on the
+// graph boundary. The strides are attached by `AttachLayoutStrides` from
+// inside `ImportAndLegalizeComputation`, so they are part of what this tool
+// prints. The matching `nv_tensor_ir.alignment` attributes need buffer
+// assignment and are added later, by the emitter only.
+
 // CHECK-LABEL: @test_reshape_split
 // CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>
 // CHECK: %[[OUT:.*]] = reshape %[[ARG0]] : tensor<8x16xf32> -> tensor<8x4x4xf32>
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_broadcast.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_broadcast.hlo
index 874d349509c..4c2084192d8 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_broadcast.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_broadcast.hlo
@@ -1,10 +1,15 @@
 // RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s
 
-// CHECK-LABEL: @test_brodcast_constant
-// CHECK: %[[OUT:.*]] = constant dense<1.000000e+00> : tensor<8x16xf32>
+// A rank-0 constant is emitted at rank 1: CudaTile's layout propagation
+// rejects a rank-0 constant, so `PromoteRank0Tensors` rebuilds it as `[1]`.
+// The broadcast is still materialised rather than folded into the constant.
+// CHECK-LABEL: @test_broadcast_constant
+// CHECK: %[[CST:.*]] = constant dense<1.000000e+00> : tensor<1xf32>
+// CHECK: %[[TEMP:.*]] = reshape %[[CST]] : tensor<1xf32> -> tensor<1x1xf32>
+// CHECK: %[[OUT:.*]] = broadcast %[[TEMP]] : tensor<1x1xf32> -> tensor<8x16xf32>
 // CHECK: results %[[OUT]]
 
-test_brodcast_constant {
+test_broadcast_constant {
   %c0 = f32[] constant(1)
   %out = f32[8,16] broadcast(%c0), dimensions={}
 }
@@ -38,10 +43,11 @@ test_broadcast_transpose {
 
 // -----
 
+// A rank-0 parameter is declared as `[1]` on the boundary, which leaves the
+// broadcast reading it directly.
 // CHECK-LABEL: @test_broadcast_scalar
 // CHECK-SAME: %[[ARG0:.*]]: tensor<1xf32>
-// CHECK: %[[TEMP:.*]] = reshape %[[ARG0]] : tensor<1xf32> -> tensor<1xf32>
-// CHECK: %[[OUT:.*]] = broadcast %[[TEMP]] : tensor<1xf32> -> tensor<32xf32>
+// CHECK: %[[OUT:.*]] = broadcast %[[ARG0]] : tensor<1xf32> -> tensor<32xf32>
 // CHECK: results %[[OUT]]
 
 test_broadcast_scalar {
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_float.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_float.hlo
index 552d91b3d00..5b7a019b52c 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_float.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_float.hlo
@@ -40,11 +40,11 @@ test_reduce_mul {
 
 // -----
 
+// The reduction body adds its operands in the flipped order; addition is
+// commutative, so this still matches the built-in `add` reduction mode.
 // CHECK-LABEL: @test_reduce_add_flip
 // CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xf32>
-// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) <dimensions = [1], identity = [0.000000e+00 : f32]> (%[[LHS:.*]]: f32, %[[RHS:.*]]: f32) {
-// CHECK: %[[RESULT:.*]] = arith.addf %[[RHS]], %[[LHS]] : f32
-// CHECK: yield %[[RESULT]]
+// CHECK: %[[REDUCE:.*]] = reduce(%[[ARG0]]) <dimensions = [1], reduction_mode = <add>> : tensor<16x64xf32> -> tensor<16x1xf32>
 // CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xf32> -> tensor<16xf32>
 // CHECK: results %[[OUT]]
 
@@ -62,11 +62,10 @@ test_reduce_add_flip {
 
 // -----
 
+// As above, but for the commutative `mul` reduction mode.
 // CHECK-LABEL: @test_reduce_mul_flip
 // CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xf32>
-// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) <dimensions = [1], identity = [1.000000e+00 : f32]> (%[[LHS:.*]]: f32, %[[RHS:.*]]: f32) {
-// CHECK: %[[RESULT:.*]] = arith.mulf %[[RHS]], %[[LHS]] : f32
-// CHECK: yield %[[RESULT]]
+// CHECK: %[[REDUCE:.*]] = reduce(%[[ARG0]]) <dimensions = [1], reduction_mode = <mul>> : tensor<16x64xf32> -> tensor<16x1xf32>
 // CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xf32> -> tensor<16xf32>
 // CHECK: results %[[OUT]]
 
@@ -86,9 +85,7 @@ test_reduce_mul_flip {
 
 // CHECK-LABEL: @test_reduce_min
 // CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xf32>
-// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) <dimensions = [1], identity = [0x7F800000 : f32]> (%[[LHS:.*]]: f32, %[[RHS:.*]]: f32) {
-// CHECK: %[[RESULT:.*]] = arith.minimumf %[[LHS]], %[[RHS]] : f32
-// CHECK: yield %[[RESULT]]
+// CHECK: %[[REDUCE:.*]] = reduce(%[[ARG0]]) <dimensions = [1], reduction_mode = <min>> : tensor<16x64xf32> -> tensor<16x1xf32>
 // CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xf32> -> tensor<16xf32>
 // CHECK: results %[[OUT]]
 
@@ -108,9 +105,7 @@ test_reduce_min {
 
 // CHECK-LABEL: @test_reduce_max
 // CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xf32>
-// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) <dimensions = [1], identity = [0xFF800000 : f32]> (%[[LHS:.*]]: f32, %[[RHS:.*]]: f32) {
-// CHECK: %[[RESULT:.*]] = arith.maximumf %[[LHS]], %[[RHS]] : f32
-// CHECK: yield %[[RESULT]]
+// CHECK: %[[REDUCE:.*]] = reduce(%[[ARG0]]) <dimensions = [1], reduction_mode = <max>> : tensor<16x64xf32> -> tensor<16x1xf32>
 // CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xf32> -> tensor<16xf32>
 // CHECK: results %[[OUT]]
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_int_signed.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_int_signed.hlo
index 8f47188340e..3e88f5321bd 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_int_signed.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_int_signed.hlo
@@ -1,10 +1,10 @@
 // RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s
 
+// The reduction body adds its operands in the flipped order; addition is
+// commutative, so this still matches the built-in `add` reduction mode.
 // CHECK-LABEL: @test_reduce_add_flip
 // CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xsi32>
-// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) <dimensions = [1], identity = [0 : si32]> (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) {
-// CHECK: %[[RESULT:.*]] = arith.addi %[[RHS]], %[[LHS]] : i32
-// CHECK: yield %[[RESULT]]
+// CHECK: %[[REDUCE:.*]] = reduce(%[[ARG0]]) <dimensions = [1], reduction_mode = <add>> : tensor<16x64xsi32> -> tensor<16x1xsi32>
 // CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xsi32> -> tensor<16xsi32>
 // CHECK: results %[[OUT]]
 
@@ -22,11 +22,10 @@ test_reduce_add_flip {
 
 // -----
 
+// As above, but for the commutative `mul` reduction mode.
 // CHECK-LABEL: @test_reduce_mul_flip
 // CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xsi32>
-// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) <dimensions = [1], identity = [1 : si32]> (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) {
-// CHECK: %[[RESULT:.*]] = arith.muli %[[RHS]], %[[LHS]] : i32
-// CHECK: yield %[[RESULT]]
+// CHECK: %[[REDUCE:.*]] = reduce(%[[ARG0]]) <dimensions = [1], reduction_mode = <mul>> : tensor<16x64xsi32> -> tensor<16x1xsi32>
 // CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xsi32> -> tensor<16xsi32>
 // CHECK: results %[[OUT]]
 
@@ -46,9 +45,7 @@ test_reduce_mul_flip {
 
 // CHECK-LABEL: @test_reduce_min
 // CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xsi32>
-// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) <dimensions = [1], identity = [2147483647 : si32]> (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) {
-// CHECK: %[[RESULT:.*]] = arith.minsi %[[LHS]], %[[RHS]] : i32
-// CHECK: yield %[[RESULT]]
+// CHECK: %[[REDUCE:.*]] = reduce(%[[ARG0]]) <dimensions = [1], reduction_mode = <min>> : tensor<16x64xsi32> -> tensor<16x1xsi32>
 // CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xsi32> -> tensor<16xsi32>
 // CHECK: results %[[OUT]]
 
@@ -68,9 +65,7 @@ test_reduce_min {
 
 // CHECK-LABEL: @test_reduce_max
 // CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xsi32>
-// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) <dimensions = [1], identity = [-2147483648 : si32]> (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) {
-// CHECK: %[[RESULT:.*]] = arith.maxsi %[[LHS]], %[[RHS]] : i32
-// CHECK: yield %[[RESULT]]
+// CHECK: %[[REDUCE:.*]] = reduce(%[[ARG0]]) <dimensions = [1], reduction_mode = <max>> : tensor<16x64xsi32> -> tensor<16x1xsi32>
 // CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xsi32> -> tensor<16xsi32>
 // CHECK: results %[[OUT]]
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_int_unsigned.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_int_unsigned.hlo
index dc97e8f4074..36b28917b6b 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_int_unsigned.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/reduce_int_unsigned.hlo
@@ -1,10 +1,10 @@
 // RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s
 
+// The reduction body adds its operands in the flipped order; addition is
+// commutative, so this still matches the built-in `add` reduction mode.
 // CHECK-LABEL: @test_reduce_add_flip
 // CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xui32>
-// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) <dimensions = [1], identity = [0 : ui32]> (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) {
-// CHECK: %[[RESULT:.*]] = arith.addi %[[RHS]], %[[LHS]] : i32
-// CHECK: yield %[[RESULT]]
+// CHECK: %[[REDUCE:.*]] = reduce(%[[ARG0]]) <dimensions = [1], reduction_mode = <add>> : tensor<16x64xui32> -> tensor<16x1xui32>
 // CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xui32> -> tensor<16xui32>
 // CHECK: results %[[OUT]]
 
@@ -22,11 +22,10 @@ test_reduce_add_flip {
 
 // -----
 
+// As above, but for the commutative `mul` reduction mode.
 // CHECK-LABEL: @test_reduce_mul_flip
 // CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xui32>
-// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) <dimensions = [1], identity = [1 : ui32]> (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) {
-// CHECK: %[[RESULT:.*]] = arith.muli %[[RHS]], %[[LHS]] : i32
-// CHECK: yield %[[RESULT]]
+// CHECK: %[[REDUCE:.*]] = reduce(%[[ARG0]]) <dimensions = [1], reduction_mode = <mul>> : tensor<16x64xui32> -> tensor<16x1xui32>
 // CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xui32> -> tensor<16xui32>
 // CHECK: results %[[OUT]]
 
@@ -46,9 +45,7 @@ test_reduce_mul_flip {
 
 // CHECK-LABEL: @test_reduce_min
 // CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xui32>
-// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) <dimensions = [1], identity = [4294967295 : ui32]> (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) {
-// CHECK: %[[RESULT:.*]] = arith.minui %[[LHS]], %[[RHS]] : i32
-// CHECK: yield %[[RESULT]]
+// CHECK: %[[REDUCE:.*]] = reduce(%[[ARG0]]) <dimensions = [1], reduction_mode = <min>> : tensor<16x64xui32> -> tensor<16x1xui32>
 // CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xui32> -> tensor<16xui32>
 // CHECK: results %[[OUT]]
 
@@ -68,9 +65,7 @@ test_reduce_min {
 
 // CHECK-LABEL: @test_reduce_max
 // CHECK-SAME: %[[ARG0:.*]]: tensor<16x64xui32>
-// CHECK: %[[REDUCE:.*]] = reduce_ud(%[[ARG0]]) <dimensions = [1], identity = [0 : ui32]> (%[[LHS:.*]]: i32, %[[RHS:.*]]: i32) {
-// CHECK: %[[RESULT:.*]] = arith.maxui %[[LHS]], %[[RHS]] : i32
-// CHECK: yield %[[RESULT]]
+// CHECK: %[[REDUCE:.*]] = reduce(%[[ARG0]]) <dimensions = [1], reduction_mode = <max>> : tensor<16x64xui32> -> tensor<16x1xui32>
 // CHECK: %[[OUT:.*]] = reshape %[[REDUCE]] : tensor<16x1xui32> -> tensor<16xui32>
 // CHECK: results %[[OUT]]
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/BUILD b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/BUILD
new file mode 100644
index 00000000000..ce95b3e1e9e
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/BUILD
@@ -0,0 +1,70 @@
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [":friends"],
+    licenses = ["notice"],
+)
+
+package_group(
+    name = "friends",
+    includes = [
+        "//xla:friends",
+    ],
+)
+
+gentbl_cc_library(
+    name = "passes_inc_gen",
+    tbl_outs = {
+        "passes.h.inc": [
+            "-gen-pass-decls",
+            "-name=TensorIrTransforms",
+        ],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "passes.td",
+    visibility = ["//visibility:private"],
+    deps = ["@llvm-project//mlir:PassBaseTdFiles"],
+)
+
+cc_library(
+    name = "type_conversion",
+    srcs = [
+        "type_conversion.cc",
+    ],
+    hdrs = [
+        "type_conversion.h",
+    ],
+    deps = [
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+    ],
+)
+
+cc_library(
+    name = "passes",
+    srcs = [
+        "legalize_stablehlo_to_tensor_ir.cc",
+    ],
+    hdrs = [
+        "passes.h",
+    ],
+    deps = [
+        ":passes_inc_gen",
+        ":type_conversion",
+        "//xla/mlir_hlo",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@stablehlo//:stablehlo_ops",
+        "@tensor_ir//:NVTensorIRDialect",
+    ],
+)
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/legalize_stablehlo_to_tensor_ir.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/legalize_stablehlo_to_tensor_ir.cc
new file mode 100644
index 00000000000..6052d8b3eed
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/legalize_stablehlo_to_tensor_ir.cc
@@ -0,0 +1,1877 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include <cstdint>
+#include <optional>
+#include <utility>
+
+#include "tensor_ir/Dialect/TensorIR.h"
+#include "tensor_ir/Dialect/TensorIRAttrs.h"
+#include "llvm/ADT/APFloat.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/Sequence.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/IR/ArithAttributes.h"
+#include "mlir/Dialect/Arith/IR/ArithDialect.h"
+#include "mlir/Dialect/Func/IR/FuncDialect.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Block.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributeInterfaces.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Types.h"
+#include "mlir/IR/Value.h"
+#include "mlir/IR/ValueRange.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "mlir/Pass/Pass.h"  // IWYU pragma: keep
+#include "mlir/Support/LLVM.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "stablehlo/dialect/StablehloOps.h"
+#include "xla/backends/gpu/codegen/tensor_ir/transforms/type_conversion.h"
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+
+namespace mlir::nv_tensor_ir::xla {
+#define GEN_PASS_DEF_LEGALIZESTABLEHLOTOTENSORIRPASS
+#include "xla/backends/gpu/codegen/tensor_ir/transforms/passes.h.inc"
+
+namespace {
+template <typename SourceOp, typename TargetOp>
+struct ElementwiseUnaryOpConversion
+    : public mlir::OpConversionPattern<SourceOp> {
+  using mlir::OpConversionPattern<SourceOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      SourceOp op, typename SourceOp::Adaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    mlir::Type result_type =
+        this->getTypeConverter()->convertType(op.getType());
+    if (!result_type) {
+      return rewriter.notifyMatchFailure(op, "failed to convert result type");
+    }
+    rewriter.replaceOpWithNewOp<TargetOp>(op, result_type,
+                                          adaptor.getOperand());
+    return mlir::success();
+  }
+};
+
+template <typename SourceOp, typename TargetOp>
+struct ElementwiseBinaryOpConversion
+    : public mlir::OpConversionPattern<SourceOp> {
+  using mlir::OpConversionPattern<SourceOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      SourceOp op, typename SourceOp::Adaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    mlir::Type result_type =
+        this->getTypeConverter()->convertType(op.getType());
+    if (!result_type) {
+      return rewriter.notifyMatchFailure(op, "failed to convert result type");
+    }
+    rewriter.replaceOpWithNewOp<TargetOp>(op, result_type, adaptor.getLhs(),
+                                          adaptor.getRhs());
+    return mlir::success();
+  }
+};
+
+struct SelectOpConversion
+    : public mlir::OpConversionPattern<mlir::stablehlo::SelectOp> {
+  using mlir::OpConversionPattern<
+      mlir::stablehlo::SelectOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::stablehlo::SelectOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    mlir::Type result_type =
+        this->getTypeConverter()->convertType(op.getType());
+    if (!result_type) {
+      return rewriter.notifyMatchFailure(op, "failed to convert result type");
+    }
+    rewriter.replaceOpWithNewOp<mlir::nv_tensor_ir::BinarySelectOp>(
+        op, result_type, adaptor.getPred(), adaptor.getOnTrue(),
+        adaptor.getOnFalse());
+    return mlir::success();
+  }
+};
+
+struct CompareOpConversion
+    : public mlir::OpConversionPattern<mlir::stablehlo::CompareOp> {
+  using mlir::OpConversionPattern<
+      mlir::stablehlo::CompareOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::stablehlo::CompareOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    if (op.getCompareType().has_value() &&
+        *op.getCompareType() == mlir::stablehlo::ComparisonType::TOTALORDER) {
+      return rewriter.notifyMatchFailure(
+          op, "TOTALORDER comparison is not supported in nv_tensor_ir");
+    }
+
+    mlir::Type converted_res_type =
+        this->getTypeConverter()->convertType(op.getType());
+    if (!converted_res_type) {
+      return rewriter.notifyMatchFailure(op, "failed to convert result type");
+    }
+
+    bool is_float = false;
+    if (op.getCompareType().has_value() &&
+        *op.getCompareType() != mlir::stablehlo::ComparisonType::NOTYPE) {
+      is_float =
+          (*op.getCompareType() == mlir::stablehlo::ComparisonType::FLOAT);
+    } else {
+      mlir::Type elem_type = op.getLhs().getType().getElementType();
+      if (mlir::isa<mlir::FloatType>(elem_type)) {
+        is_float = true;
+      } else if (mlir::isa<mlir::IntegerType>(elem_type)) {
+        is_float = false;
+      } else {
+        return rewriter.notifyMatchFailure(
+            op, "unsupported operand element type for comparison");
+      }
+    }
+
+    mlir::nv_tensor_ir::Comparator comparator;
+    if (is_float) {
+      // Every float direction maps to its ordered comparator except NE, which
+      // maps to the *unordered* `une`. This asymmetry is deliberate and is not
+      // a typo: IEEE 754 defines `NaN != x` as true, so not-equal is the only
+      // direction that must hold when either operand is NaN. Ordered `one`
+      // returns false in that case and would silently change the result of any
+      // fusion that can produce a NaN.
+      switch (op.getComparisonDirection()) {
+        case mlir::stablehlo::ComparisonDirection::EQ:
+          comparator = mlir::nv_tensor_ir::Comparator::oeq;
+          break;
+        case mlir::stablehlo::ComparisonDirection::NE:
+          comparator = mlir::nv_tensor_ir::Comparator::une;
+          break;
+        case mlir::stablehlo::ComparisonDirection::GE:
+          comparator = mlir::nv_tensor_ir::Comparator::oge;
+          break;
+        case mlir::stablehlo::ComparisonDirection::GT:
+          comparator = mlir::nv_tensor_ir::Comparator::ogt;
+          break;
+        case mlir::stablehlo::ComparisonDirection::LE:
+          comparator = mlir::nv_tensor_ir::Comparator::ole;
+          break;
+        case mlir::stablehlo::ComparisonDirection::LT:
+          comparator = mlir::nv_tensor_ir::Comparator::olt;
+          break;
+      }
+    } else {
+      switch (op.getComparisonDirection()) {
+        case mlir::stablehlo::ComparisonDirection::EQ:
+          comparator = mlir::nv_tensor_ir::Comparator::eq;
+          break;
+        case mlir::stablehlo::ComparisonDirection::NE:
+          comparator = mlir::nv_tensor_ir::Comparator::neq;
+          break;
+        case mlir::stablehlo::ComparisonDirection::GE:
+          comparator = mlir::nv_tensor_ir::Comparator::ge;
+          break;
+        case mlir::stablehlo::ComparisonDirection::GT:
+          comparator = mlir::nv_tensor_ir::Comparator::gt;
+          break;
+        case mlir::stablehlo::ComparisonDirection::LE:
+          comparator = mlir::nv_tensor_ir::Comparator::le;
+          break;
+        case mlir::stablehlo::ComparisonDirection::LT:
+          comparator = mlir::nv_tensor_ir::Comparator::lt;
+          break;
+      }
+    }
+
+    rewriter.replaceOpWithNewOp<mlir::nv_tensor_ir::CmpOp>(
+        op, converted_res_type, comparator, adaptor.getLhs(), adaptor.getRhs());
+    return mlir::success();
+  }
+};
+
+struct ConvertOpConversion
+    : public mlir::OpConversionPattern<mlir::stablehlo::ConvertOp> {
+  using mlir::OpConversionPattern<
+      mlir::stablehlo::ConvertOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::stablehlo::ConvertOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    mlir::Type result_type =
+        this->getTypeConverter()->convertType(op.getType());
+    if (!result_type) {
+      return rewriter.notifyMatchFailure(op, "failed to convert result type");
+    }
+    rewriter.replaceOpWithNewOp<mlir::nv_tensor_ir::ConvertOp>(
+        op, result_type, adaptor.getOperand());
+    return mlir::success();
+  }
+};
+
+// stablehlo.clamp(min, operand, max) has no nv_tensor_ir counterpart and is
+// expressed as min(max_value, max(min_value, operand)). The nesting order and
+// the operand order within each of the two operations are chosen to match the
+// string-building converter this pass replaces, so that the two produce
+// byte-identical output for the same fusion.
+struct ClampOpConversion
+    : public mlir::OpConversionPattern<mlir::stablehlo::ClampOp> {
+  using mlir::OpConversionPattern<
+      mlir::stablehlo::ClampOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::stablehlo::ClampOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    mlir::Type result_type =
+        this->getTypeConverter()->convertType(op.getType());
+    if (!result_type) {
+      return rewriter.notifyMatchFailure(op, "failed to convert result type");
+    }
+    // StableHLO additionally allows rank-0 `min`/`max` operands that are
+    // broadcast against `operand`. nv_tensor_ir.max and nv_tensor_ir.min
+    // require all operands to have the same shape, and silently pairing a
+    // scalar with a shaped operand would not verify, so reject that form
+    // explicitly rather than emitting something that cannot be built. XLA's
+    // HLO importer never produces it, because HLO clamp is already
+    // shape-matched by the time a fusion is formed.
+    if (adaptor.getMin().getType() != result_type ||
+        adaptor.getMax().getType() != result_type) {
+      return rewriter.notifyMatchFailure(
+          op,
+          "scalar clamp bounds are not supported; min and max must have "
+          "the same shape as the operand");
+    }
+    mlir::Value lower_bounded = mlir::nv_tensor_ir::MaxOp::create(
+        rewriter, op.getLoc(), result_type, adaptor.getMin(),
+        adaptor.getOperand());
+    rewriter.replaceOpWithNewOp<mlir::nv_tensor_ir::MinOp>(
+        op, result_type, adaptor.getMax(), lower_bounded);
+    return mlir::success();
+  }
+};
+
+// Returns a splat nv_tensor_ir.constant of the integral value `value` over
+// `type`, which must have a floating-point element type.
+mlir::Value CreateFloatSplat(mlir::ConversionPatternRewriter& rewriter,
+                             mlir::Location loc, mlir::RankedTensorType type,
+                             int64_t value) {
+  auto element_type = mlir::cast<mlir::FloatType>(type.getElementType());
+  llvm::APFloat splat(element_type.getFloatSemantics(), value);
+  return mlir::nv_tensor_ir::ConstantOp::create(
+      rewriter, loc, mlir::DenseElementsAttr::get(type, splat));
+}
+
+// NOTE: Decomposing Expm1(x) to Exp(x) - 1.0 can result in a severe loss of
+// precision for values of x close to 0. nv_tensor_ir has no expm1, so this is
+// the only available lowering; the caveat is reproduced from the converter
+// this pass replaces because it is invisible in the emitted IR.
+struct Expm1OpConversion
+    : public mlir::OpConversionPattern<mlir::stablehlo::Expm1Op> {
+  using mlir::OpConversionPattern<
+      mlir::stablehlo::Expm1Op>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::stablehlo::Expm1Op op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    auto result_type = mlir::dyn_cast_or_null<mlir::RankedTensorType>(
+        this->getTypeConverter()->convertType(op.getType()));
+    if (!result_type ||
+        !mlir::isa<mlir::FloatType>(result_type.getElementType())) {
+      return rewriter.notifyMatchFailure(
+          op, "expected a ranked floating-point result");
+    }
+    mlir::Location loc = op.getLoc();
+    mlir::Value one = CreateFloatSplat(rewriter, loc, result_type, /*value=*/1);
+    mlir::Value exp = mlir::nv_tensor_ir::ExpOp::create(
+        rewriter, loc, result_type, adaptor.getOperand());
+    rewriter.replaceOpWithNewOp<mlir::nv_tensor_ir::SubOp>(op, result_type, exp,
+                                                           one);
+    return mlir::success();
+  }
+};
+
+// NOTE: Decomposing Log1p(x) to Log(x + 1.0) can result in a severe loss of
+// precision for values of x close to 0. See the note on Expm1OpConversion.
+struct Log1pOpConversion
+    : public mlir::OpConversionPattern<mlir::stablehlo::Log1pOp> {
+  using mlir::OpConversionPattern<
+      mlir::stablehlo::Log1pOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::stablehlo::Log1pOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    auto result_type = mlir::dyn_cast_or_null<mlir::RankedTensorType>(
+        this->getTypeConverter()->convertType(op.getType()));
+    if (!result_type ||
+        !mlir::isa<mlir::FloatType>(result_type.getElementType())) {
+      return rewriter.notifyMatchFailure(
+          op, "expected a ranked floating-point result");
+    }
+    mlir::Location loc = op.getLoc();
+    mlir::Value one = CreateFloatSplat(rewriter, loc, result_type, /*value=*/1);
+    mlir::Value sum = mlir::nv_tensor_ir::AddOp::create(
+        rewriter, loc, result_type, adaptor.getOperand(), one);
+    rewriter.replaceOpWithNewOp<mlir::nv_tensor_ir::LogOp>(op, result_type,
+                                                           sum);
+    return mlir::success();
+  }
+};
+
+//===----------------------------------------------------------------------===//
+// Structural operations.
+//===----------------------------------------------------------------------===//
+
+struct ReshapeOpConversion
+    : public mlir::OpConversionPattern<mlir::stablehlo::ReshapeOp> {
+  using mlir::OpConversionPattern<
+      mlir::stablehlo::ReshapeOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::stablehlo::ReshapeOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    mlir::Type result_type =
+        this->getTypeConverter()->convertType(op.getType());
+    if (!result_type) {
+      return rewriter.notifyMatchFailure(op, "failed to convert result type");
+    }
+    rewriter.replaceOpWithNewOp<mlir::nv_tensor_ir::ReshapeOp>(
+        op, result_type, adaptor.getOperand(),
+        /*dynamic_sizes=*/mlir::ValueRange{});
+    return mlir::success();
+  }
+};
+
+struct TransposeOpConversion
+    : public mlir::OpConversionPattern<mlir::stablehlo::TransposeOp> {
+  using mlir::OpConversionPattern<
+      mlir::stablehlo::TransposeOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::stablehlo::TransposeOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    mlir::Type result_type =
+        this->getTypeConverter()->convertType(op.getType());
+    if (!result_type) {
+      return rewriter.notifyMatchFailure(op, "failed to convert result type");
+    }
+    // Both dialects define the permutation as output[i] = input[perm[i]], so
+    // it can be forwarded as is.
+    rewriter.replaceOpWithNewOp<mlir::nv_tensor_ir::TransposeOp>(
+        op, result_type, adaptor.getOperand(), op.getPermutationAttr());
+    return mlir::success();
+  }
+};
+
+struct SliceOpConversion
+    : public mlir::OpConversionPattern<mlir::stablehlo::SliceOp> {
+  using mlir::OpConversionPattern<
+      mlir::stablehlo::SliceOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::stablehlo::SliceOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    mlir::Type result_type =
+        this->getTypeConverter()->convertType(op.getType());
+    if (!result_type) {
+      return rewriter.notifyMatchFailure(op, "failed to convert result type");
+    }
+    rewriter.replaceOpWithNewOp<mlir::nv_tensor_ir::SliceOp>(
+        op, result_type, adaptor.getOperand(), op.getStartIndicesAttr(),
+        op.getLimitIndicesAttr(), op.getStridesAttr());
+    return mlir::success();
+  }
+};
+
+struct ConcatenateOpConversion
+    : public mlir::OpConversionPattern<mlir::stablehlo::ConcatenateOp> {
+  using mlir::OpConversionPattern<
+      mlir::stablehlo::ConcatenateOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::stablehlo::ConcatenateOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    mlir::Type result_type =
+        this->getTypeConverter()->convertType(op.getType());
+    if (!result_type) {
+      return rewriter.notifyMatchFailure(op, "failed to convert result type");
+    }
+    rewriter.replaceOpWithNewOp<mlir::nv_tensor_ir::ConcatenateOp>(
+        op, result_type, adaptor.getInputs(), op.getDimension());
+    return mlir::success();
+  }
+};
+
+struct IotaOpConversion
+    : public mlir::OpConversionPattern<mlir::stablehlo::IotaOp> {
+  using mlir::OpConversionPattern<mlir::stablehlo::IotaOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::stablehlo::IotaOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    mlir::Type result_type =
+        this->getTypeConverter()->convertType(op.getType());
+    if (!result_type) {
+      return rewriter.notifyMatchFailure(op, "failed to convert result type");
+    }
+    rewriter.replaceOpWithNewOp<mlir::nv_tensor_ir::IotaOp>(
+        op, result_type, op.getIotaDimension(),
+        /*dynamic_sizes=*/mlir::ValueRange{});
+    return mlir::success();
+  }
+};
+
+struct ConstantOpConversion
+    : public mlir::OpConversionPattern<mlir::stablehlo::ConstantOp> {
+  using mlir::OpConversionPattern<
+      mlir::stablehlo::ConstantOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::stablehlo::ConstantOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    auto result_type = mlir::dyn_cast_or_null<mlir::RankedTensorType>(
+        this->getTypeConverter()->convertType(op.getType()));
+    if (!result_type) {
+      return rewriter.notifyMatchFailure(op, "failed to convert result type");
+    }
+    auto value = mlir::dyn_cast<mlir::DenseElementsAttr>(op.getValue());
+    if (!value) {
+      return rewriter.notifyMatchFailure(op,
+                                         "expected a dense elements attribute");
+    }
+    // nv_tensor_ir.constant infers its result type from the value attribute
+    // (InferTypeOpAdaptor), so the attribute itself has to be rebuilt over the
+    // converted element type. The raw storage is bit-identical because the
+    // type conversion only turns signless integers into signed integers of the
+    // same width.
+    if (value.getType() != result_type) {
+      if (!mlir::DenseElementsAttr::isValidRawBuffer(result_type,
+                                                     value.getRawData())) {
+        return rewriter.notifyMatchFailure(
+            op,
+            "constant value cannot be reinterpreted over the converted type");
+      }
+      value = mlir::DenseElementsAttr::getFromRawBuffer(result_type,
+                                                        value.getRawData());
+    }
+    rewriter.replaceOpWithNewOp<mlir::nv_tensor_ir::ConstantOp>(
+        op, mlir::cast<mlir::TypedAttr>(value));
+    return mlir::success();
+  }
+};
+
+struct BroadcastInDimOpConversion
+    : public mlir::OpConversionPattern<mlir::stablehlo::BroadcastInDimOp> {
+  using mlir::OpConversionPattern<
+      mlir::stablehlo::BroadcastInDimOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::stablehlo::BroadcastInDimOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    auto result_type = mlir::dyn_cast_or_null<mlir::RankedTensorType>(
+        this->getTypeConverter()->convertType(op.getType()));
+    if (!result_type) {
+      return rewriter.notifyMatchFailure(op, "failed to convert result type");
+    }
+    auto input_type =
+        mlir::dyn_cast<mlir::RankedTensorType>(adaptor.getOperand().getType());
+    if (!input_type) {
+      return rewriter.notifyMatchFailure(op, "expected a ranked tensor input");
+    }
+    llvm::ArrayRef<int64_t> broadcast_dimensions = op.getBroadcastDimensions();
+    if (static_cast<int64_t>(broadcast_dimensions.size()) !=
+        input_type.getRank()) {
+      return rewriter.notifyMatchFailure(
+          op, "expected one broadcast dimension per input dimension");
+    }
+
+    mlir::Type element_type = result_type.getElementType();
+    mlir::Location loc = op.getLoc();
+    mlir::Value value = adaptor.getOperand();
+
+    // nv_tensor_ir.broadcast preserves the rank and can only expand literal
+    // unit dimensions, so the input is first reshaped to the output rank. That
+    // reshape relinearizes row-major, which only produces the intended result
+    // if the input dimensions already appear in output order. If they do not,
+    // sort them with a transpose first.
+    if (!llvm::is_sorted(broadcast_dimensions)) {
+      llvm::SmallVector<int64_t> permutation =
+          llvm::to_vector(llvm::seq<int64_t>(0, input_type.getRank()));
+      llvm::sort(permutation, [&](int64_t lhs, int64_t rhs) {
+        return broadcast_dimensions[lhs] < broadcast_dimensions[rhs];
+      });
+      llvm::SmallVector<int64_t> transposed_shape;
+      transposed_shape.reserve(permutation.size());
+      for (int64_t dimension : permutation) {
+        transposed_shape.push_back(input_type.getDimSize(dimension));
+      }
+      value = mlir::nv_tensor_ir::TransposeOp::create(
+          rewriter, loc,
+          mlir::RankedTensorType::get(transposed_shape, element_type), value,
+          rewriter.getDenseI64ArrayAttr(permutation));
+    }
+
+    llvm::SmallVector<int64_t> expanded_shape(result_type.getRank(), 1);
+    for (auto [dimension, output_dimension] :
+         llvm::enumerate(broadcast_dimensions)) {
+      if (output_dimension < 0 || output_dimension >= result_type.getRank()) {
+        return rewriter.notifyMatchFailure(
+            op, "broadcast dimension is out of bounds");
+      }
+      expanded_shape[output_dimension] = input_type.getDimSize(dimension);
+    }
+    auto expanded_type =
+        mlir::RankedTensorType::get(expanded_shape, element_type);
+    if (expanded_type != value.getType()) {
+      value = mlir::nv_tensor_ir::ReshapeOp::create(
+          rewriter, loc, expanded_type, value,
+          /*dynamic_sizes=*/mlir::ValueRange{});
+    }
+    // A broadcast that changes no dimension is rejected by the verifier.
+    if (expanded_type != result_type) {
+      value = mlir::nv_tensor_ir::BroadcastOp::create(
+          rewriter, loc, result_type, value,
+          /*dynamic_sizes=*/mlir::ValueRange{});
+    }
+    rewriter.replaceOp(op, value);
+    return mlir::success();
+  }
+};
+
+//===----------------------------------------------------------------------===//
+// Bitcast.
+//===----------------------------------------------------------------------===//
+
+// Returns true if `minor_to_major` is the default descending layout, i.e. the
+// tensor is already stored in row-major order.
+bool IsDefaultLayout(llvm::ArrayRef<int64_t> minor_to_major) {
+  int64_t rank = minor_to_major.size();
+  for (int64_t dimension : llvm::seq<int64_t>(0, rank)) {
+    if (minor_to_major[dimension] != rank - 1 - dimension) {
+      return false;
+    }
+  }
+  return true;
+}
+
+// Reads a minor-to-major layout that HloFunctionImporter attached to `op` as a
+// discardable `dense<[...]> : tensor<Nxindex>` attribute. Returns std::nullopt
+// if the attribute is absent or is not a permutation of [0, rank).
+std::optional<llvm::SmallVector<int64_t>> GetLayoutAttribute(
+    mlir::Operation* op, llvm::StringRef name, int64_t rank) {
+  auto attr = op->getAttrOfType<mlir::DenseIntElementsAttr>(name);
+  if (!attr || attr.getNumElements() != rank) {
+    return std::nullopt;
+  }
+  llvm::SmallVector<int64_t> minor_to_major;
+  minor_to_major.reserve(rank);
+  for (const llvm::APInt& value : attr.getValues<llvm::APInt>()) {
+    minor_to_major.push_back(value.getSExtValue());
+  }
+  llvm::SmallVector<bool> seen(rank, false);
+  for (int64_t dimension : minor_to_major) {
+    if (dimension < 0 || dimension >= rank || seen[dimension]) {
+      return std::nullopt;
+    }
+    seen[dimension] = true;
+  }
+  return minor_to_major;
+}
+
+// Returns the dimensions of a tensor of shape `shape` and layout
+// `minor_to_major`, reordered from most-major to most-minor. This is the shape
+// the same buffer would have if it were relabelled to carry the default
+// descending layout, i.e. xla::ShapeUtil's
+// MakeShapeWithDescendingLayoutAndSamePhysicalLayout.
+llvm::SmallVector<int64_t> GetPhysicalShape(
+    llvm::ArrayRef<int64_t> shape, llvm::ArrayRef<int64_t> minor_to_major) {
+  llvm::SmallVector<int64_t> physical_shape;
+  physical_shape.reserve(shape.size());
+  for (int64_t dimension : llvm::reverse(minor_to_major)) {
+    physical_shape.push_back(shape[dimension]);
+  }
+  return physical_shape;
+}
+
+// Lowers mhlo.bitcast, which reinterprets a buffer under a new shape and
+// layout without moving any element.
+//
+// nv_tensor_ir tensors, like StableHLO tensors, have no layout: every tensor
+// is row-major. A bitcast between two non-row-major shapes therefore cannot be
+// a bare reshape, because a reshape relinearizes in row-major order. The
+// lowering normalizes both sides to row-major and relinearizes in between:
+//
+//   1. transpose the operand from its logical order into physical
+//      (major-to-minor) order, if its layout is not already the default;
+//   2. reshape, if the two physical shapes differ;
+//   3. transpose from the result's physical order back into its logical
+//      order, if the result layout is not the default.
+//
+// Each step is skipped when it would be the identity, so a bitcast between two
+// default-layout shapes lowers to a single reshape.
+//
+// The layouts are read from the `source_layout` and `result_layout` attributes
+// that HloFunctionImporter attaches to every mhlo.bitcast it creates. They are
+// discardable (mhlo.bitcast declares no attributes in ODS), so they are the
+// only record of the layout anywhere in the imported module -- the MLIR tensor
+// type cannot carry it. Without them the physical arrangement is unknowable,
+// so their absence is a match failure rather than an assumption that the
+// layout is the default one: guessing would silently miscompile exactly the
+// non-default-layout bitcasts this code exists to handle.
+struct BitcastOpConversion
+    : public mlir::OpConversionPattern<mlir::mhlo::BitcastOp> {
+  using mlir::OpConversionPattern<mlir::mhlo::BitcastOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::mhlo::BitcastOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    auto result_type = mlir::dyn_cast_or_null<mlir::RankedTensorType>(
+        this->getTypeConverter()->convertType(op.getType()));
+    if (!result_type) {
+      return rewriter.notifyMatchFailure(op, "failed to convert result type");
+    }
+    auto input_type =
+        mlir::dyn_cast<mlir::RankedTensorType>(adaptor.getOperand().getType());
+    if (!input_type) {
+      return rewriter.notifyMatchFailure(op, "expected a ranked tensor input");
+    }
+    // A bitcast reinterprets the buffer, it does not convert elements.
+    if (input_type.getElementType() != result_type.getElementType()) {
+      return rewriter.notifyMatchFailure(
+          op, "expected the input and result element types to match");
+    }
+
+    std::optional<llvm::SmallVector<int64_t>> source_layout =
+        GetLayoutAttribute(op, "source_layout", input_type.getRank());
+    std::optional<llvm::SmallVector<int64_t>> result_layout =
+        GetLayoutAttribute(op, "result_layout", result_type.getRank());
+    if (!source_layout.has_value() || !result_layout.has_value()) {
+      return rewriter.notifyMatchFailure(
+          op,
+          "mhlo.bitcast is missing the source_layout or result_layout "
+          "attribute, so the physical element order is unknown");
+    }
+
+    mlir::Location loc = op.getLoc();
+    mlir::Type element_type = result_type.getElementType();
+    mlir::Value value = adaptor.getOperand();
+
+    // 1. Logical -> physical order for the operand.
+    if (!IsDefaultLayout(*source_layout)) {
+      llvm::SmallVector<int64_t> permutation =
+          llvm::to_vector(llvm::reverse(*source_layout));
+      value = mlir::nv_tensor_ir::TransposeOp::create(
+          rewriter, loc,
+          mlir::RankedTensorType::get(
+              GetPhysicalShape(input_type.getShape(), *source_layout),
+              element_type),
+          value, rewriter.getDenseI64ArrayAttr(permutation));
+    }
+
+    // 2. Relinearize, both sides now being row-major.
+    auto physical_result_type = mlir::RankedTensorType::get(
+        GetPhysicalShape(result_type.getShape(), *result_layout), element_type);
+    if (physical_result_type != value.getType()) {
+      value = mlir::nv_tensor_ir::ReshapeOp::create(
+          rewriter, loc, physical_result_type, value,
+          /*dynamic_sizes=*/mlir::ValueRange{});
+    }
+
+    // 3. Physical -> logical order for the result. The permutation is the
+    // inverse of the one used in step 1, because it undoes the reordering
+    // rather than applying it.
+    if (!IsDefaultLayout(*result_layout)) {
+      llvm::SmallVector<int64_t> permutation(result_type.getRank());
+      for (auto [physical, logical] :
+           llvm::enumerate(llvm::reverse(*result_layout))) {
+        permutation[logical] = physical;
+      }
+      value = mlir::nv_tensor_ir::TransposeOp::create(
+          rewriter, loc, result_type, value,
+          rewriter.getDenseI64ArrayAttr(permutation));
+    }
+
+    rewriter.replaceOp(op, value);
+    return mlir::success();
+  }
+};
+
+//===----------------------------------------------------------------------===//
+// Reductions.
+//===----------------------------------------------------------------------===//
+
+// Returns the built-in reduction mode implemented by the combiner `op`, or
+// std::nullopt if `op` is not a recognized combiner. The nv_tensor_ir variants
+// are accepted as well because the elementwise patterns may already have
+// rewritten the combiner inside the reduction body.
+std::optional<mlir::nv_tensor_ir::ReductionMode> GetReductionMode(
+    mlir::Operation& op) {
+  if (mlir::isa<mlir::stablehlo::AddOp, mlir::nv_tensor_ir::AddOp>(op)) {
+    return mlir::nv_tensor_ir::ReductionMode::add;
+  }
+  if (mlir::isa<mlir::stablehlo::MulOp, mlir::nv_tensor_ir::MulOp>(op)) {
+    return mlir::nv_tensor_ir::ReductionMode::mul;
+  }
+  if (mlir::isa<mlir::stablehlo::MaxOp, mlir::nv_tensor_ir::MaxOp>(op)) {
+    return mlir::nv_tensor_ir::ReductionMode::max;
+  }
+  if (mlir::isa<mlir::stablehlo::MinOp, mlir::nv_tensor_ir::MinOp>(op)) {
+    return mlir::nv_tensor_ir::ReductionMode::min;
+  }
+  return std::nullopt;
+}
+
+// Returns true if `value` is the identity element of `mode`.
+// nv_tensor_ir.reduce has no init operand and implicitly uses the identity of
+// its reduction mode, so a stablehlo.reduce with any other init value must not
+// be lowered onto it.
+bool IsReductionIdentity(mlir::DenseElementsAttr value,
+                         mlir::nv_tensor_ir::ReductionMode mode) {
+  if (!value.isSplat()) {
+    return false;
+  }
+  mlir::Type element_type = value.getElementType();
+  if (mlir::isa<mlir::FloatType>(element_type)) {
+    llvm::APFloat splat = value.getSplatValue<llvm::APFloat>();
+    switch (mode) {
+      case mlir::nv_tensor_ir::ReductionMode::add:
+        return splat.isZero();
+      case mlir::nv_tensor_ir::ReductionMode::mul:
+        return splat == llvm::APFloat(splat.getSemantics(), 1);
+      case mlir::nv_tensor_ir::ReductionMode::max:
+        return splat.isInfinity() && splat.isNegative();
+      case mlir::nv_tensor_ir::ReductionMode::min:
+        return splat.isInfinity() && !splat.isNegative();
+      default:
+        return false;
+    }
+  }
+  auto integer_type = mlir::dyn_cast<mlir::IntegerType>(element_type);
+  if (!integer_type) {
+    return false;
+  }
+  llvm::APInt splat = value.getSplatValue<llvm::APInt>();
+  unsigned width = integer_type.getWidth();
+  bool is_unsigned = integer_type.isUnsigned();
+  switch (mode) {
+    case mlir::nv_tensor_ir::ReductionMode::add:
+      return splat.isZero();
+    case mlir::nv_tensor_ir::ReductionMode::mul:
+      return splat == llvm::APInt(width, 1);
+    case mlir::nv_tensor_ir::ReductionMode::max:
+      return splat == (is_unsigned ? llvm::APInt::getMinValue(width)
+                                   : llvm::APInt::getSignedMinValue(width));
+    case mlir::nv_tensor_ir::ReductionMode::min:
+      return splat == (is_unsigned ? llvm::APInt::getMaxValue(width)
+                                   : llvm::APInt::getSignedMaxValue(width));
+    default:
+      return false;
+  }
+}
+
+// Returns the signless equivalent of `type` if it is a signed or unsigned
+// integer, and `type` itself otherwise. The reduce_ud region verifier
+// (TensorOps.cpp, ReduceUDOp::verifyRegions) compares the block argument and
+// yield types against the *signless* form of the identity's type, so the whole
+// region is built over signless integers even though the identity and the
+// tensor operands keep their signedness. This deliberately bypasses
+// TensorIrTypeConverter, which maps signless integers to signed.
+mlir::Type ToSignlessType(mlir::Type type) {
+  if (auto integer_type = mlir::dyn_cast<mlir::IntegerType>(type)) {
+    if (!integer_type.isSignless()) {
+      return mlir::IntegerType::get(type.getContext(), integer_type.getWidth());
+    }
+  }
+  return type;
+}
+
+// Returns a scalar TypedAttr holding the single value of the splat `value`,
+// retyped to `element_type`. Used both for the reduce_ud identity, which keeps
+// the signed/unsigned element type, and for constants in the combiner region,
+// which are signless.
+mlir::TypedAttr GetScalarAttribute(mlir::DenseElementsAttr value,
+                                   mlir::Type element_type) {
+  if (mlir::isa<mlir::FloatType>(element_type)) {
+    return mlir::FloatAttr::get(element_type,
+                                value.getSplatValue<llvm::APFloat>());
+  }
+  return mlir::IntegerAttr::get(element_type,
+                                value.getSplatValue<llvm::APInt>());
+}
+
+// Translates one operation of a stablehlo.reduce body into the arith
+// operations that implement it on signless scalars, and appends them at the
+// builder's insertion point. `operands` are the already-translated scalar
+// operands. Returns a null Value if the operation has no translation.
+//
+// The body of stablehlo.reduce operates on rank-0 tensors, whereas the
+// reduce_ud region operates on plain scalars, so this is a translation and not
+// a clone. Signedness is read from the *stablehlo* operand type, which still
+// carries it (`ui32` for unsigned, signless `i32` for signed), because the
+// scalars themselves are uniformly signless.
+mlir::Value ConvertReductionBodyOp(mlir::Operation& op,
+                                   mlir::ValueRange operands,
+                                   mlir::OpBuilder& builder) {
+  mlir::Location loc = op.getLoc();
+
+  if (auto constant = mlir::dyn_cast<mlir::stablehlo::ConstantOp>(op)) {
+    auto value = mlir::dyn_cast<mlir::DenseElementsAttr>(constant.getValue());
+    if (!value || !value.isSplat()) {
+      return nullptr;
+    }
+    return mlir::arith::ConstantOp::create(
+        builder, loc,
+        GetScalarAttribute(value, ToSignlessType(value.getElementType())));
+  }
+
+  // Everything below is an elementwise operation, so the element type of the
+  // first operand decides which arith variant applies.
+  auto operand_type = mlir::dyn_cast<mlir::RankedTensorType>(
+      op.getNumOperands() > 0 ? op.getOperand(0).getType()
+                              : op.getResult(0).getType());
+  if (!operand_type) {
+    return nullptr;
+  }
+  mlir::Type element_type = operand_type.getElementType();
+  bool is_float = mlir::isa<mlir::FloatType>(element_type);
+  // StableHLO spells unsigned integers `ui<N>` and signed ones as signless
+  // `i<N>`. `i1` is XLA's PRED, which XLA treats as unsigned.
+  bool is_unsigned =
+      element_type.isUnsignedInteger() || element_type.isInteger(/*width=*/1);
+
+  if (mlir::isa<mlir::stablehlo::AddOp>(op)) {
+    if (is_float) {
+      return mlir::arith::AddFOp::create(builder, loc, operands[0],
+                                         operands[1]);
+    }
+    return mlir::arith::AddIOp::create(builder, loc, operands[0], operands[1]);
+  }
+  if (mlir::isa<mlir::stablehlo::MulOp>(op)) {
+    if (is_float) {
+      return mlir::arith::MulFOp::create(builder, loc, operands[0],
+                                         operands[1]);
+    }
+    return mlir::arith::MulIOp::create(builder, loc, operands[0], operands[1]);
+  }
+  if (mlir::isa<mlir::stablehlo::MaxOp>(op)) {
+    if (is_float) {
+      return mlir::arith::MaximumFOp::create(builder, loc, operands[0],
+                                             operands[1]);
+    }
+    if (is_unsigned) {
+      return mlir::arith::MaxUIOp::create(builder, loc, operands[0],
+                                          operands[1]);
+    }
+    return mlir::arith::MaxSIOp::create(builder, loc, operands[0], operands[1]);
+  }
+  if (mlir::isa<mlir::stablehlo::MinOp>(op)) {
+    if (is_float) {
+      return mlir::arith::MinimumFOp::create(builder, loc, operands[0],
+                                             operands[1]);
+    }
+    if (is_unsigned) {
+      return mlir::arith::MinUIOp::create(builder, loc, operands[0],
+                                          operands[1]);
+    }
+    return mlir::arith::MinSIOp::create(builder, loc, operands[0], operands[1]);
+  }
+  if (mlir::isa<mlir::stablehlo::AndOp>(op)) {
+    return mlir::arith::AndIOp::create(builder, loc, operands[0], operands[1]);
+  }
+  if (mlir::isa<mlir::stablehlo::OrOp>(op)) {
+    return mlir::arith::OrIOp::create(builder, loc, operands[0], operands[1]);
+  }
+  if (mlir::isa<mlir::stablehlo::XorOp>(op)) {
+    return mlir::arith::XOrIOp::create(builder, loc, operands[0], operands[1]);
+  }
+  if (mlir::isa<mlir::stablehlo::SelectOp>(op)) {
+    return mlir::arith::SelectOp::create(builder, loc, operands[0], operands[1],
+                                         operands[2]);
+  }
+  if (mlir::isa<mlir::stablehlo::ClampOp>(op)) {
+    // clamp(min, x, max) == min(max, max(min, x)); see ClampOpConversion.
+    if (is_float) {
+      mlir::Value lower = mlir::arith::MaximumFOp::create(
+          builder, loc, operands[0], operands[1]);
+      return mlir::arith::MinimumFOp::create(builder, loc, operands[2], lower);
+    }
+    if (is_unsigned) {
+      mlir::Value lower =
+          mlir::arith::MaxUIOp::create(builder, loc, operands[0], operands[1]);
+      return mlir::arith::MinUIOp::create(builder, loc, operands[2], lower);
+    }
+    mlir::Value lower =
+        mlir::arith::MaxSIOp::create(builder, loc, operands[0], operands[1]);
+    return mlir::arith::MinSIOp::create(builder, loc, operands[2], lower);
+  }
+  if (auto compare = mlir::dyn_cast<mlir::stablehlo::CompareOp>(op)) {
+    // arith has no total-order comparison; see CompareOpConversion.
+    if (compare.getCompareType().has_value() &&
+        *compare.getCompareType() ==
+            mlir::stablehlo::ComparisonType::TOTALORDER) {
+      return nullptr;
+    }
+    if (is_float) {
+      mlir::arith::CmpFPredicate predicate;
+      switch (compare.getComparisonDirection()) {
+        case mlir::stablehlo::ComparisonDirection::EQ:
+          predicate = mlir::arith::CmpFPredicate::OEQ;
+          break;
+        // See CompareOpConversion for why not-equal alone is unordered.
+        case mlir::stablehlo::ComparisonDirection::NE:
+          predicate = mlir::arith::CmpFPredicate::UNE;
+          break;
+        case mlir::stablehlo::ComparisonDirection::GE:
+          predicate = mlir::arith::CmpFPredicate::OGE;
+          break;
+        case mlir::stablehlo::ComparisonDirection::GT:
+          predicate = mlir::arith::CmpFPredicate::OGT;
+          break;
+        case mlir::stablehlo::ComparisonDirection::LE:
+          predicate = mlir::arith::CmpFPredicate::OLE;
+          break;
+        case mlir::stablehlo::ComparisonDirection::LT:
+          predicate = mlir::arith::CmpFPredicate::OLT;
+          break;
+      }
+      return mlir::arith::CmpFOp::create(builder, loc, predicate, operands[0],
+                                         operands[1]);
+    }
+    mlir::arith::CmpIPredicate predicate;
+    switch (compare.getComparisonDirection()) {
+      case mlir::stablehlo::ComparisonDirection::EQ:
+        predicate = mlir::arith::CmpIPredicate::eq;
+        break;
+      case mlir::stablehlo::ComparisonDirection::NE:
+        predicate = mlir::arith::CmpIPredicate::ne;
+        break;
+      case mlir::stablehlo::ComparisonDirection::GE:
+        predicate = is_unsigned ? mlir::arith::CmpIPredicate::uge
+                                : mlir::arith::CmpIPredicate::sge;
+        break;
+      case mlir::stablehlo::ComparisonDirection::GT:
+        predicate = is_unsigned ? mlir::arith::CmpIPredicate::ugt
+                                : mlir::arith::CmpIPredicate::sgt;
+        break;
+      case mlir::stablehlo::ComparisonDirection::LE:
+        predicate = is_unsigned ? mlir::arith::CmpIPredicate::ule
+                                : mlir::arith::CmpIPredicate::sle;
+        break;
+      case mlir::stablehlo::ComparisonDirection::LT:
+        predicate = is_unsigned ? mlir::arith::CmpIPredicate::ult
+                                : mlir::arith::CmpIPredicate::slt;
+        break;
+    }
+    return mlir::arith::CmpIOp::create(builder, loc, predicate, operands[0],
+                                       operands[1]);
+  }
+
+  return nullptr;
+}
+
+struct ReduceOpConversion
+    : public mlir::OpConversionPattern<mlir::stablehlo::ReduceOp> {
+  // Runs before the elementwise patterns so that the combiner in the reduction
+  // body has not been rewritten yet when the body is inspected.
+  ReduceOpConversion(const mlir::TypeConverter& type_converter,
+                     mlir::MLIRContext* context)
+      : mlir::OpConversionPattern<mlir::stablehlo::ReduceOp>(
+            type_converter, context, /*benefit=*/2) {}
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::stablehlo::ReduceOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    if (adaptor.getInputs().size() != 1 || op->getNumResults() != 1) {
+      // nv_tensor_ir.reduce_ud is variadic and could express a multi-output
+      // reduction, but nothing upstream produces one yet, so the extra
+      // bookkeeping is not carried here.
+      return rewriter.notifyMatchFailure(
+          op, "only single-input, single-result reductions are supported");
+    }
+    auto result_type = mlir::dyn_cast_or_null<mlir::RankedTensorType>(
+        this->getTypeConverter()->convertType(op->getResult(0).getType()));
+    if (!result_type) {
+      return rewriter.notifyMatchFailure(op, "failed to convert result type");
+    }
+    auto input_type = mlir::dyn_cast<mlir::RankedTensorType>(
+        adaptor.getInputs()[0].getType());
+    if (!input_type) {
+      return rewriter.notifyMatchFailure(op, "expected a ranked tensor input");
+    }
+    if (input_type.getElementType() != result_type.getElementType()) {
+      return rewriter.notifyMatchFailure(
+          op, "expected the input and result element types to match");
+    }
+
+    mlir::DenseElementsAttr init_value;
+    if (!mlir::matchPattern(adaptor.getInitValues()[0],
+                            mlir::m_Constant(&init_value)) ||
+        !init_value.isSplat()) {
+      return rewriter.notifyMatchFailure(
+          op, "expected a splat constant reduction init value");
+    }
+
+    llvm::SmallVector<int32_t> dimensions;
+    dimensions.reserve(op.getDimensions().size());
+    for (int64_t dimension : op.getDimensions()) {
+      // nv_tensor_ir.reduce takes an i32 dense array. Bounding the dimension
+      // by the tensor rank also makes the narrowing below lossless.
+      if (dimension < 0 || dimension >= input_type.getRank()) {
+        return rewriter.notifyMatchFailure(
+            op, "reduction dimension is out of range");
+      }
+      dimensions.push_back(static_cast<int32_t>(dimension));
+    }
+    if (dimensions.empty()) {
+      return rewriter.notifyMatchFailure(
+          op, "expected at least one reduction dimension");
+    }
+
+    // Both nv_tensor_ir.reduce and nv_tensor_ir.reduce_ud preserve the rank
+    // and contract the reduced dimensions to 1, while stablehlo.reduce drops
+    // them, so a reshape is needed to recover the stablehlo result shape.
+    llvm::SmallVector<int64_t> reduced_shape(input_type.getShape());
+    for (int32_t dimension : dimensions) {
+      reduced_shape[dimension] = 1;
+    }
+    auto reduced_type =
+        mlir::RankedTensorType::get(reduced_shape, input_type.getElementType());
+
+    mlir::Value reduced = BuildReduce(op, adaptor, reduced_type, dimensions,
+                                      init_value, rewriter);
+    if (!reduced) {
+      return rewriter.notifyMatchFailure(
+          op, "reduction body has no nv_tensor_ir equivalent");
+    }
+    if (reduced_type != result_type) {
+      reduced = mlir::nv_tensor_ir::ReshapeOp::create(
+          rewriter, op.getLoc(), result_type, reduced,
+          /*dynamic_sizes=*/mlir::ValueRange{});
+    }
+    rewriter.replaceOp(op, reduced);
+    return mlir::success();
+  }
+
+ private:
+  // Emits the reduction itself, preferring the built-in nv_tensor_ir.reduce
+  // over nv_tensor_ir.reduce_ud because the tiling pipeline understands the
+  // built-in modes best. Returns a null Value if neither applies.
+  static mlir::Value BuildReduce(mlir::stablehlo::ReduceOp op,
+                                 OpAdaptor adaptor,
+                                 mlir::RankedTensorType reduced_type,
+                                 llvm::ArrayRef<int32_t> dimensions,
+                                 mlir::DenseElementsAttr init_value,
+                                 mlir::ConversionPatternRewriter& rewriter) {
+    if (std::optional<mlir::nv_tensor_ir::ReductionMode> mode =
+            GetBuiltinReductionMode(op, init_value)) {
+      return mlir::nv_tensor_ir::ReduceOp::create(
+          rewriter, op.getLoc(), reduced_type, adaptor.getInputs()[0],
+          rewriter.getDenseI32ArrayAttr(dimensions), *mode);
+    }
+    return BuildUserDefinedReduce(op, adaptor, reduced_type, dimensions,
+                                  init_value, rewriter);
+  }
+
+  // Returns the built-in reduction mode this reduction can use, if any. The
+  // combiner has to be a single recognized operation over the two block
+  // arguments, and the init value has to be that mode's implicit identity,
+  // because nv_tensor_ir.reduce has no init operand.
+  static std::optional<mlir::nv_tensor_ir::ReductionMode>
+  GetBuiltinReductionMode(mlir::stablehlo::ReduceOp op,
+                          mlir::DenseElementsAttr init_value) {
+    mlir::Block& body = op.getBody().front();
+    if (!llvm::hasNItems(body.without_terminator(), 1)) {
+      return std::nullopt;
+    }
+    mlir::Operation& combiner = *body.begin();
+    mlir::Operation* terminator = body.getTerminator();
+    if (terminator->getNumOperands() != 1 || combiner.getNumResults() != 1 ||
+        terminator->getOperand(0) != combiner.getResult(0)) {
+      return std::nullopt;
+    }
+    // The combiner has to consume both block arguments and nothing else.
+    // add/mul/max/min are all commutative, so the order does not matter, but
+    // a body such as `add(%arg0, %arg0)` is not a plain reduction and must not
+    // be mistaken for one.
+    if (body.getNumArguments() != 2 || combiner.getNumOperands() != 2) {
+      return std::nullopt;
+    }
+    mlir::Value lhs = combiner.getOperand(0);
+    mlir::Value rhs = combiner.getOperand(1);
+    mlir::Value arg0 = body.getArgument(0);
+    mlir::Value arg1 = body.getArgument(1);
+    if (!((lhs == arg0 && rhs == arg1) || (lhs == arg1 && rhs == arg0))) {
+      return std::nullopt;
+    }
+    std::optional<mlir::nv_tensor_ir::ReductionMode> mode =
+        GetReductionMode(combiner);
+    if (!mode.has_value() || !IsReductionIdentity(init_value, *mode)) {
+      return std::nullopt;
+    }
+    return mode;
+  }
+
+  // Emits nv_tensor_ir.reduce_ud, whose explicit identity and combiner region
+  // can express any reduction body, and translates the stablehlo body into
+  // that region. Returns a null Value if the body cannot be translated.
+  static mlir::Value BuildUserDefinedReduce(
+      mlir::stablehlo::ReduceOp op, OpAdaptor adaptor,
+      mlir::RankedTensorType reduced_type, llvm::ArrayRef<int32_t> dimensions,
+      mlir::DenseElementsAttr init_value,
+      mlir::ConversionPatternRewriter& rewriter) {
+    mlir::Location loc = op.getLoc();
+    // The identity keeps the converted (signed or unsigned) element type: the
+    // op verifier compares it against the result element type as is.
+    mlir::TypedAttr identity =
+        GetScalarAttribute(init_value, reduced_type.getElementType());
+    auto reduce_op = mlir::nv_tensor_ir::ReduceUDOp::create(
+        rewriter, loc, mlir::TypeRange{reduced_type}, adaptor.getInputs(),
+        rewriter.getDenseI32ArrayAttr(dimensions),
+        rewriter.getArrayAttr({identity}));
+
+    // The region takes (prev_result, curr_operand) over *signless* scalars.
+    mlir::Type scalar_type = ToSignlessType(reduced_type.getElementType());
+    mlir::Block* body = rewriter.createBlock(&reduce_op.getBody());
+    body->addArgument(scalar_type, loc);
+    body->addArgument(scalar_type, loc);
+
+    mlir::Block& source_body = op.getBody().front();
+    if (source_body.getNumArguments() != body->getNumArguments()) {
+      rewriter.eraseOp(reduce_op);
+      return nullptr;
+    }
+    mlir::IRMapping mapping;
+    mapping.map(source_body.getArguments(), body->getArguments());
+    for (mlir::Operation& source_op : source_body.without_terminator()) {
+      llvm::SmallVector<mlir::Value> operands;
+      operands.reserve(source_op.getNumOperands());
+      for (mlir::Value operand : source_op.getOperands()) {
+        operands.push_back(mapping.lookupOrNull(operand));
+      }
+      if (llvm::is_contained(operands, mlir::Value())) {
+        rewriter.eraseOp(reduce_op);
+        return nullptr;
+      }
+      mlir::Value result =
+          ConvertReductionBodyOp(source_op, operands, rewriter);
+      if (!result || source_op.getNumResults() != 1) {
+        rewriter.eraseOp(reduce_op);
+        return nullptr;
+      }
+      mapping.map(source_op.getResult(0), result);
+    }
+
+    llvm::SmallVector<mlir::Value> yielded;
+    for (mlir::Value operand : source_body.getTerminator()->getOperands()) {
+      yielded.push_back(mapping.lookupOrNull(operand));
+    }
+    if (yielded.size() != 1 || llvm::is_contained(yielded, mlir::Value())) {
+      rewriter.eraseOp(reduce_op);
+      return nullptr;
+    }
+    mlir::nv_tensor_ir::YieldOp::create(rewriter, loc, yielded);
+
+    rewriter.setInsertionPointAfter(reduce_op);
+    return reduce_op.getResult(0);
+  }
+};
+
+//===----------------------------------------------------------------------===//
+// Matmul.
+//===----------------------------------------------------------------------===//
+
+// Returns the dimensions of a rank-`rank` operand that are neither batch nor
+// contracting dimensions ("free" dimensions), in increasing order. Keeping
+// them in the operand's original relative order is what makes the final
+// reshape back to the stablehlo result shape a pure relinearization.
+llvm::SmallVector<int64_t> GetFreeDimensions(
+    int64_t rank, llvm::ArrayRef<int64_t> batch_dimensions,
+    llvm::ArrayRef<int64_t> contracting_dimensions) {
+  llvm::SmallVector<int64_t> free_dimensions;
+  for (int64_t dimension : llvm::seq<int64_t>(0, rank)) {
+    if (llvm::is_contained(batch_dimensions, dimension) ||
+        llvm::is_contained(contracting_dimensions, dimension)) {
+      continue;
+    }
+    free_dimensions.push_back(dimension);
+  }
+  return free_dimensions;
+}
+
+// Returns the product of the extents of `dimensions` in `type`, i.e. the size
+// of the single dimension they collapse into. The product of no dimensions is
+// 1, which is the right neutral extent for a missing batch or free group.
+int64_t GetDimensionProduct(mlir::RankedTensorType type,
+                            llvm::ArrayRef<int64_t> dimensions) {
+  int64_t product = 1;
+  for (int64_t dimension : dimensions) {
+    product *= type.getDimSize(dimension);
+  }
+  return product;
+}
+
+// Returns `value` permuted by `permutation`, or `value` unchanged if the
+// permutation is the identity. Both dialects define the permutation as
+// output[i] = input[permutation[i]].
+mlir::Value MaybeTranspose(mlir::ConversionPatternRewriter& rewriter,
+                           mlir::Location loc, mlir::Value value,
+                           llvm::ArrayRef<int64_t> permutation) {
+  // A permutation of [0, rank) is the identity exactly when it is sorted.
+  if (llvm::is_sorted(permutation)) {
+    return value;
+  }
+  auto type = mlir::cast<mlir::RankedTensorType>(value.getType());
+  llvm::SmallVector<int64_t> shape;
+  shape.reserve(permutation.size());
+  for (int64_t dimension : permutation) {
+    shape.push_back(type.getDimSize(dimension));
+  }
+  return mlir::nv_tensor_ir::TransposeOp::create(
+      rewriter, loc, type.clone(shape), value,
+      rewriter.getDenseI64ArrayAttr(permutation));
+}
+
+// Returns `value` reshaped to `shape`, or `value` unchanged if it already has
+// exactly that type. Comparing types rather than shapes keeps an identity
+// reshape from showing up in the output.
+mlir::Value MaybeReshape(mlir::ConversionPatternRewriter& rewriter,
+                         mlir::Location loc, mlir::Value value,
+                         llvm::ArrayRef<int64_t> shape) {
+  auto type = mlir::cast<mlir::RankedTensorType>(value.getType());
+  mlir::RankedTensorType reshaped_type = type.clone(shape);
+  if (reshaped_type == type) {
+    return value;
+  }
+  return mlir::nv_tensor_ir::ReshapeOp::create(
+      rewriter, loc, reshaped_type, value,
+      /*dynamic_sizes=*/mlir::ValueRange{});
+}
+
+// Returns failure if `precision_config` requests anything other than DEFAULT.
+//
+// nv_tensor_ir.matmul carries neither a precision nor an algorithm attribute,
+// so precision_config and algorithm have nowhere to go. Rather than silently
+// dropping a request that changes the numerics, reject it. An all-DEFAULT
+// precision_config asks for nothing in particular and is accepted (and
+// dropped), which is what almost every dot carries.
+mlir::LogicalResult CheckDefaultPrecision(
+    mlir::Operation* op, std::optional<mlir::ArrayAttr> precision_config,
+    mlir::ConversionPatternRewriter& rewriter) {
+  if (!precision_config.has_value()) {
+    return mlir::success();
+  }
+  for (mlir::Attribute precision : *precision_config) {
+    auto precision_attr =
+        mlir::dyn_cast<mlir::stablehlo::PrecisionAttr>(precision);
+    if (!precision_attr ||
+        precision_attr.getValue() != mlir::stablehlo::Precision::DEFAULT) {
+      return rewriter.notifyMatchFailure(
+          op, "nv_tensor_ir.matmul cannot express a non-default precision");
+    }
+  }
+  return mlir::success();
+}
+
+// Lowers a dot onto nv_tensor_ir.matmul by reducing it to a canonical rank-3
+// (batched) or rank-2 (unbatched) matmul:
+//
+//   1. lhs: transpose to [batch..., free..., contracting...], then reshape to
+//      [B, M, K] with B = prod(batch), M = prod(free), K = prod(contracting).
+//   2. rhs: transpose to [batch..., contracting..., free...], then reshape to
+//      [B, K, N].
+//   3. matmul -> [B, M, N].
+//   4. reshape to the stablehlo result shape, which is
+//      [batch..., lhs_free..., rhs_free...] in exactly that order.
+//
+// Each transpose and reshape is only emitted when it is not the identity, so
+// the common already-canonical dot lowers to a bare matmul. On success `op` is
+// replaced.
+mlir::LogicalResult LowerDotToMatmul(
+    mlir::Operation* op, mlir::Value lhs_value, mlir::Value rhs_value,
+    llvm::ArrayRef<int64_t> lhs_batch, llvm::ArrayRef<int64_t> rhs_batch,
+    llvm::ArrayRef<int64_t> lhs_contract, llvm::ArrayRef<int64_t> rhs_contract,
+    mlir::RankedTensorType result_type,
+    mlir::ConversionPatternRewriter& rewriter) {
+  auto lhs_type = mlir::dyn_cast<mlir::RankedTensorType>(lhs_value.getType());
+  auto rhs_type = mlir::dyn_cast<mlir::RankedTensorType>(rhs_value.getType());
+  if (!lhs_type || !rhs_type) {
+    return rewriter.notifyMatchFailure(op, "expected ranked tensor operands");
+  }
+
+  if (lhs_batch.size() != rhs_batch.size()) {
+    return rewriter.notifyMatchFailure(
+        op, "expected the same number of batch dimensions on both operands");
+  }
+  if (lhs_contract.size() != rhs_contract.size()) {
+    return rewriter.notifyMatchFailure(
+        op,
+        "expected the same number of contracting dimensions on both "
+        "operands");
+  }
+  auto in_range = [](llvm::ArrayRef<int64_t> dimensions, int64_t rank) {
+    return llvm::all_of(dimensions, [rank](int64_t dimension) {
+      return dimension >= 0 && dimension < rank;
+    });
+  };
+  if (!in_range(lhs_batch, lhs_type.getRank()) ||
+      !in_range(lhs_contract, lhs_type.getRank()) ||
+      !in_range(rhs_batch, rhs_type.getRank()) ||
+      !in_range(rhs_contract, rhs_type.getRank())) {
+    return rewriter.notifyMatchFailure(op, "dot dimension is out of range");
+  }
+  // The batch and contracting lists pair up positionally between the two
+  // operands, so the paired extents have to agree for the collapsed B and K
+  // to be the same on both sides.
+  for (auto [lhs_dimension, rhs_dimension] : llvm::zip(lhs_batch, rhs_batch)) {
+    if (lhs_type.getDimSize(lhs_dimension) !=
+        rhs_type.getDimSize(rhs_dimension)) {
+      return rewriter.notifyMatchFailure(
+          op, "paired batch dimensions have different extents");
+    }
+  }
+  for (auto [lhs_dimension, rhs_dimension] :
+       llvm::zip(lhs_contract, rhs_contract)) {
+    if (lhs_type.getDimSize(lhs_dimension) !=
+        rhs_type.getDimSize(rhs_dimension)) {
+      return rewriter.notifyMatchFailure(
+          op, "paired contracting dimensions have different extents");
+    }
+  }
+
+  llvm::SmallVector<int64_t> lhs_free =
+      GetFreeDimensions(lhs_type.getRank(), lhs_batch, lhs_contract);
+  llvm::SmallVector<int64_t> rhs_free =
+      GetFreeDimensions(rhs_type.getRank(), rhs_batch, rhs_contract);
+
+  int64_t batch_size = GetDimensionProduct(lhs_type, lhs_batch);
+  int64_t contracting_size = GetDimensionProduct(lhs_type, lhs_contract);
+  int64_t lhs_free_size = GetDimensionProduct(lhs_type, lhs_free);
+  int64_t rhs_free_size = GetDimensionProduct(rhs_type, rhs_free);
+  // Without batch dimensions the canonical form is the plain rank-2 matmul;
+  // introducing a size-1 batch dimension would only add two reshapes.
+  bool has_batch = !lhs_batch.empty();
+
+  mlir::Location loc = op->getLoc();
+
+  llvm::SmallVector<int64_t> lhs_permutation = llvm::to_vector(lhs_batch);
+  llvm::append_range(lhs_permutation, lhs_free);
+  llvm::append_range(lhs_permutation, lhs_contract);
+  llvm::SmallVector<int64_t> lhs_shape;
+  if (has_batch) {
+    lhs_shape.push_back(batch_size);
+  }
+  lhs_shape.push_back(lhs_free_size);
+  lhs_shape.push_back(contracting_size);
+  mlir::Value lhs = MaybeTranspose(rewriter, loc, lhs_value, lhs_permutation);
+  lhs = MaybeReshape(rewriter, loc, lhs, lhs_shape);
+
+  llvm::SmallVector<int64_t> rhs_permutation = llvm::to_vector(rhs_batch);
+  llvm::append_range(rhs_permutation, rhs_contract);
+  llvm::append_range(rhs_permutation, rhs_free);
+  llvm::SmallVector<int64_t> rhs_shape;
+  if (has_batch) {
+    rhs_shape.push_back(batch_size);
+  }
+  rhs_shape.push_back(contracting_size);
+  rhs_shape.push_back(rhs_free_size);
+  mlir::Value rhs = MaybeTranspose(rewriter, loc, rhs_value, rhs_permutation);
+  rhs = MaybeReshape(rewriter, loc, rhs, rhs_shape);
+
+  llvm::SmallVector<int64_t> matmul_shape;
+  if (has_batch) {
+    matmul_shape.push_back(batch_size);
+  }
+  matmul_shape.push_back(lhs_free_size);
+  matmul_shape.push_back(rhs_free_size);
+  mlir::Value matmul = mlir::nv_tensor_ir::MatmulOp::create(
+      rewriter, loc, result_type.clone(matmul_shape), lhs, rhs);
+
+  rewriter.replaceOp(
+      op, MaybeReshape(rewriter, loc, matmul, result_type.getShape()));
+  return mlir::success();
+}
+
+struct DotGeneralOpConversion
+    : public mlir::OpConversionPattern<mlir::stablehlo::DotGeneralOp> {
+  using mlir::OpConversionPattern<
+      mlir::stablehlo::DotGeneralOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::stablehlo::DotGeneralOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    auto result_type = mlir::dyn_cast_or_null<mlir::RankedTensorType>(
+        this->getTypeConverter()->convertType(op.getType()));
+    if (!result_type) {
+      return rewriter.notifyMatchFailure(op, "failed to convert result type");
+    }
+    if (op.getAlgorithm().has_value()) {
+      return rewriter.notifyMatchFailure(
+          op, "nv_tensor_ir.matmul cannot express a dot algorithm");
+    }
+    if (mlir::failed(
+            CheckDefaultPrecision(op, op.getPrecisionConfig(), rewriter))) {
+      return mlir::failure();
+    }
+
+    mlir::stablehlo::DotDimensionNumbersAttr dimension_numbers =
+        op.getDotDimensionNumbers();
+    return LowerDotToMatmul(op, adaptor.getLhs(), adaptor.getRhs(),
+                            dimension_numbers.getLhsBatchingDimensions(),
+                            dimension_numbers.getRhsBatchingDimensions(),
+                            dimension_numbers.getLhsContractingDimensions(),
+                            dimension_numbers.getRhsContractingDimensions(),
+                            result_type, rewriter);
+  }
+};
+
+// stablehlo.dot is what the HLO importer emits for a plain two-operand dot.
+// It has no dimension numbers: it always contracts the last dimension of the
+// lhs with the first dimension of the rhs and has no batch dimensions, so it
+// is just a dot_general with those numbers filled in.
+struct DotOpConversion
+    : public mlir::OpConversionPattern<mlir::stablehlo::DotOp> {
+  using mlir::OpConversionPattern<mlir::stablehlo::DotOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::stablehlo::DotOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    auto result_type = mlir::dyn_cast_or_null<mlir::RankedTensorType>(
+        this->getTypeConverter()->convertType(op.getType()));
+    if (!result_type) {
+      return rewriter.notifyMatchFailure(op, "failed to convert result type");
+    }
+    if (mlir::failed(
+            CheckDefaultPrecision(op, op.getPrecisionConfig(), rewriter))) {
+      return mlir::failure();
+    }
+    auto lhs_type =
+        mlir::dyn_cast<mlir::RankedTensorType>(adaptor.getLhs().getType());
+    if (!lhs_type || lhs_type.getRank() == 0) {
+      return rewriter.notifyMatchFailure(
+          op, "expected a ranked lhs of rank at least one");
+    }
+    int64_t lhs_contract = lhs_type.getRank() - 1;
+    int64_t rhs_contract = 0;
+    return LowerDotToMatmul(op, adaptor.getLhs(), adaptor.getRhs(),
+                            /*lhs_batch=*/{}, /*rhs_batch=*/{},
+                            llvm::ArrayRef<int64_t>(lhs_contract),
+                            llvm::ArrayRef<int64_t>(rhs_contract), result_type,
+                            rewriter);
+  }
+};
+
+struct ResultsOpConversion
+    : public mlir::OpConversionPattern<mlir::nv_tensor_ir::ResultsOp> {
+  using mlir::OpConversionPattern<
+      mlir::nv_tensor_ir::ResultsOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::nv_tensor_ir::ResultsOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    rewriter.replaceOpWithNewOp<mlir::nv_tensor_ir::ResultsOp>(
+        op, adaptor.getOperands());
+    return mlir::success();
+  }
+};
+
+void populateLegalizeStablehloToTensorIrPatterns(
+    mlir::RewritePatternSet& patterns, mlir::TypeConverter& type_converter,
+    mlir::MLIRContext* context) {
+  // Unary elementwise operations.
+  patterns.add<ElementwiseUnaryOpConversion<mlir::stablehlo::AbsOp,
+                                            mlir::nv_tensor_ir::AbsOp>,
+               ElementwiseUnaryOpConversion<mlir::stablehlo::CeilOp,
+                                            mlir::nv_tensor_ir::CeilOp>,
+               ElementwiseUnaryOpConversion<mlir::stablehlo::FloorOp,
+                                            mlir::nv_tensor_ir::FloorOp>,
+               ElementwiseUnaryOpConversion<mlir::stablehlo::NegOp,
+                                            mlir::nv_tensor_ir::NegOp>,
+               ElementwiseUnaryOpConversion<mlir::stablehlo::CosineOp,
+                                            mlir::nv_tensor_ir::CosOp>,
+               ElementwiseUnaryOpConversion<mlir::stablehlo::SineOp,
+                                            mlir::nv_tensor_ir::SinOp>,
+               ElementwiseUnaryOpConversion<mlir::stablehlo::TanOp,
+                                            mlir::nv_tensor_ir::TanOp>,
+               ElementwiseUnaryOpConversion<mlir::stablehlo::ExpOp,
+                                            mlir::nv_tensor_ir::ExpOp>,
+               ElementwiseUnaryOpConversion<mlir::stablehlo::LogOp,
+                                            mlir::nv_tensor_ir::LogOp>,
+               ElementwiseUnaryOpConversion<mlir::stablehlo::RsqrtOp,
+                                            mlir::nv_tensor_ir::RsqrtOp>,
+               ElementwiseUnaryOpConversion<mlir::stablehlo::SqrtOp,
+                                            mlir::nv_tensor_ir::SqrtOp>,
+               ElementwiseUnaryOpConversion<mlir::stablehlo::TanhOp,
+                                            mlir::nv_tensor_ir::TanhFwdOp>,
+               ElementwiseUnaryOpConversion<mlir::stablehlo::LogisticOp,
+                                            mlir::nv_tensor_ir::SigmoidFwdOp>,
+               ElementwiseUnaryOpConversion<mlir::stablehlo::NotOp,
+                                            mlir::nv_tensor_ir::LogicalNotOp>,
+               // The HLO importer has no StableHLO op to emit for Erf -- there
+               // is none -- so it emits mhlo.erf.
+               ElementwiseUnaryOpConversion<mlir::mhlo::ErfOp,
+                                            mlir::nv_tensor_ir::ErfOp>>(
+      type_converter, context);
+
+  // Binary elementwise operations.
+  // Note: stablehlo.remainder follows the sign of the dividend (truncated
+  // remainder), matching nv_tensor_ir.rem, NOT nv_tensor_ir.mod which follows
+  // floored-remainder semantics.
+  patterns.add<ElementwiseBinaryOpConversion<mlir::stablehlo::AddOp,
+                                             mlir::nv_tensor_ir::AddOp>,
+               ElementwiseBinaryOpConversion<mlir::stablehlo::SubtractOp,
+                                             mlir::nv_tensor_ir::SubOp>,
+               ElementwiseBinaryOpConversion<mlir::stablehlo::MulOp,
+                                             mlir::nv_tensor_ir::MulOp>,
+               ElementwiseBinaryOpConversion<mlir::stablehlo::DivOp,
+                                             mlir::nv_tensor_ir::DivOp>,
+               ElementwiseBinaryOpConversion<mlir::stablehlo::MaxOp,
+                                             mlir::nv_tensor_ir::MaxOp>,
+               ElementwiseBinaryOpConversion<mlir::stablehlo::MinOp,
+                                             mlir::nv_tensor_ir::MinOp>,
+               ElementwiseBinaryOpConversion<mlir::stablehlo::PowOp,
+                                             mlir::nv_tensor_ir::PowOp>,
+               ElementwiseBinaryOpConversion<mlir::stablehlo::RemOp,
+                                             mlir::nv_tensor_ir::RemOp>,
+               ElementwiseBinaryOpConversion<mlir::stablehlo::Atan2Op,
+                                             mlir::nv_tensor_ir::Atan2Op>,
+               ElementwiseBinaryOpConversion<mlir::stablehlo::AndOp,
+                                             mlir::nv_tensor_ir::LogicalAndOp>,
+               ElementwiseBinaryOpConversion<mlir::stablehlo::OrOp,
+                                             mlir::nv_tensor_ir::LogicalOrOp>>(
+      type_converter, context);
+
+  // Select, Compare, Convert, and Results.
+  // Note: stablehlo.xor has no nv_tensor_ir equivalent at tensor granularity
+  // and is intentionally left unhandled so that the conversion target and
+  // final check reject it. Inside a reduction body it is expressible, and
+  // ConvertReductionBodyOp does handle it.
+  patterns.add<SelectOpConversion, CompareOpConversion, ConvertOpConversion,
+               ResultsOpConversion>(type_converter, context);
+
+  // Operations without a direct nv_tensor_ir counterpart, expressed in terms
+  // of the operations that do exist.
+  patterns.add<ClampOpConversion, Expm1OpConversion, Log1pOpConversion>(
+      type_converter, context);
+
+  // Structural operations, constants and matmul.
+  // Note: pad, gather, scatter, dynamic_slice, sort, reverse, convolution and
+  // the control flow operations have no nv_tensor_ir equivalent and are
+  // intentionally left unhandled.
+  patterns.add<ReshapeOpConversion, TransposeOpConversion, SliceOpConversion,
+               ConcatenateOpConversion, IotaOpConversion, ConstantOpConversion,
+               BroadcastInDimOpConversion, BitcastOpConversion,
+               DotGeneralOpConversion, DotOpConversion>(type_converter,
+                                                        context);
+
+  // Reductions, registered with a higher benefit than the elementwise
+  // patterns so that the combiner in the reduction body can be inspected.
+  patterns.add<ReduceOpConversion>(type_converter, context);
+}
+
+// Rewrites the rank-0 tensors CudaTile cannot handle to rank 1, reshaping on
+// the inside so that the rest of the body is unchanged.
+//
+// CudaTile's layout propagation rejects a rank-0 tensor in three places, each
+// with an unhelpful message: a rank-0 graph argument and a rank-0 constant
+// both fail with "failed to compute layout", and a rank-0 graph result with
+// "failed to normalize layout". A rank-0 value produced by a `reshape` is
+// accepted, which is what makes this rewrite possible at all. XLA, meanwhile,
+// produces all three routinely: a scalar fusion parameter, a scalar constant
+// feeding a broadcast, and a full reduce.
+//
+// An XLA buffer for a rank-0 shape and for the corresponding `[1]` shape have
+// the same size and alignment, and both have a (trivially) default layout. So
+// declaring the boundary as `[1]` leaves the kernel ABI untouched, and leaves
+// `AttachLayoutStrides`, which only emits an attribute for a non-default
+// layout, with nothing to do in either case.
+void PromoteRank0Tensors(mlir::nv_tensor_ir::GraphOp graph) {
+  static constexpr int64_t kRank1Shape[] = {1};
+
+  mlir::Block& body = graph.getGraphBody().front();
+  auto results_op =
+      mlir::dyn_cast<mlir::nv_tensor_ir::ResultsOp>(body.getTerminator());
+  if (results_op == nullptr) {
+    return;
+  }
+
+  mlir::OpBuilder builder(graph.getContext());
+
+  // `create` inserts before the insertion point without moving it, so
+  // successive reshapes land in argument order rather than reversed.
+  builder.setInsertionPointToStart(&body);
+  llvm::SmallVector<mlir::Type> input_types;
+  input_types.reserve(body.getNumArguments());
+  for (mlir::BlockArgument arg : body.getArguments()) {
+    auto type = mlir::dyn_cast<mlir::RankedTensorType>(arg.getType());
+    if (type == nullptr || type.getRank() != 0) {
+      input_types.push_back(arg.getType());
+      continue;
+    }
+    arg.setType(type.clone(kRank1Shape));
+    auto reshape = mlir::nv_tensor_ir::ReshapeOp::create(
+        builder, arg.getLoc(), type, arg,
+        /*dynamic_sizes=*/mlir::ValueRange{});
+    arg.replaceAllUsesExcept(reshape.getResult(), reshape);
+    input_types.push_back(arg.getType());
+  }
+
+  builder.setInsertionPoint(results_op);
+  llvm::SmallVector<mlir::Type> result_types;
+  result_types.reserve(results_op->getNumOperands());
+  for (mlir::OpOperand& operand : results_op->getOpOperands()) {
+    auto type = mlir::dyn_cast<mlir::RankedTensorType>(operand.get().getType());
+    if (type == nullptr || type.getRank() != 0) {
+      result_types.push_back(operand.get().getType());
+      continue;
+    }
+    mlir::RankedTensorType promoted = type.clone(kRank1Shape);
+    auto reshape = mlir::nv_tensor_ir::ReshapeOp::create(
+        builder, results_op.getLoc(), promoted, operand.get(),
+        /*dynamic_sizes=*/mlir::ValueRange{});
+    operand.set(reshape.getResult());
+    result_types.push_back(promoted);
+  }
+
+  graph.setFunctionType(
+      mlir::FunctionType::get(graph.getContext(), input_types, result_types));
+
+  // A rank-0 constant is rejected even though it is not on the boundary, so
+  // rebuild it at rank 1 and reshape back down.
+  llvm::SmallVector<mlir::nv_tensor_ir::ConstantOp> rank0_constants;
+  graph.walk([&](mlir::nv_tensor_ir::ConstantOp constant) {
+    auto type =
+        mlir::dyn_cast<mlir::RankedTensorType>(constant.getResult().getType());
+    if (type != nullptr && type.getRank() == 0) {
+      rank0_constants.push_back(constant);
+    }
+  });
+  for (mlir::nv_tensor_ir::ConstantOp constant : rank0_constants) {
+    auto value = mlir::dyn_cast<mlir::DenseElementsAttr>(constant.getValue());
+    if (value == nullptr) {
+      continue;
+    }
+    auto type =
+        mlir::cast<mlir::RankedTensorType>(constant.getResult().getType());
+    builder.setInsertionPoint(constant);
+    auto promoted = mlir::nv_tensor_ir::ConstantOp::create(
+        builder, constant.getLoc(),
+        mlir::cast<mlir::TypedAttr>(value.reshape(type.clone(kRank1Shape))));
+    auto reshape = mlir::nv_tensor_ir::ReshapeOp::create(
+        builder, constant.getLoc(), type, promoted.getResult(),
+        /*dynamic_sizes=*/mlir::ValueRange{});
+    constant.getResult().replaceAllUsesWith(reshape.getResult());
+    constant.erase();
+  }
+
+  // A reshape is a pure shape reinterpretation, so reshape(reshape(x)) is just
+  // reshape(x), and where the two cancel out it is x itself. Promotion creates
+  // such chains wherever it meets a reshape the importer already emitted: a
+  // full reduce, for example, is `reduce -> tensor<1xT>` followed by
+  // `reshape -> tensor<T>`, and promoting the result reshapes it straight back
+  // up. Collapse them rather than leaving them in the graph.
+  //
+  // `walk` is post-order, so the inner reshape of a chain is always rewritten
+  // before the outer one and one pass suffices. Erasing is left to a second
+  // pass, running outermost first, because erasing during the first would
+  // leave dangling entries in `reshapes`.
+  llvm::SmallVector<mlir::nv_tensor_ir::ReshapeOp> reshapes;
+  graph.walk([&](mlir::nv_tensor_ir::ReshapeOp reshape) {
+    // More than one operand means dynamic result sizes, which XLA never
+    // produces and which would not be safe to re-point.
+    if (reshape->getNumOperands() == 1) {
+      reshapes.push_back(reshape);
+    }
+  });
+  for (mlir::nv_tensor_ir::ReshapeOp reshape : reshapes) {
+    auto inner =
+        reshape.getOperand(0).getDefiningOp<mlir::nv_tensor_ir::ReshapeOp>();
+    if (inner != nullptr && inner->getNumOperands() == 1) {
+      reshape.setOperand(0, inner.getOperand(0));
+    }
+  }
+  for (mlir::nv_tensor_ir::ReshapeOp reshape : llvm::reverse(reshapes)) {
+    if (reshape.getOperand(0).getType() == reshape.getResult().getType()) {
+      reshape.getResult().replaceAllUsesWith(reshape.getOperand(0));
+    }
+    if (reshape.use_empty()) {
+      reshape.erase();
+    }
+  }
+}
+
+class LegalizeStablehloToTensorIrPass
+    : public impl::LegalizeStablehloToTensorIrPassBase<
+          LegalizeStablehloToTensorIrPass> {
+ public:
+  using LegalizeStablehloToTensorIrPassBase::
+      LegalizeStablehloToTensorIrPassBase;
+
+  void runOnOperation() override {
+    mlir::ModuleOp module = getOperation();
+    mlir::MLIRContext* context = &getContext();
+
+    TensorIrTypeConverter type_converter;
+
+    // Collect func.func operations to process.
+    llvm::SmallVector<mlir::func::FuncOp> func_ops;
+    module.walk(
+        [&](mlir::func::FuncOp func_op) { func_ops.push_back(func_op); });
+
+    for (mlir::func::FuncOp func_op : func_ops) {
+      mlir::OpBuilder builder(func_op);
+
+      // Convert function signature types.
+      auto func_type = func_op.getFunctionType();
+      llvm::SmallVector<mlir::Type> new_inputs;
+      llvm::SmallVector<mlir::Type> new_results;
+      if (failed(
+              type_converter.convertTypes(func_type.getInputs(), new_inputs)) ||
+          failed(type_converter.convertTypes(func_type.getResults(),
+                                             new_results))) {
+        func_op.emitOpError("failed to convert function signature types");
+        signalPassFailure();
+        return;
+      }
+      auto converted_func_type =
+          builder.getFunctionType(new_inputs, new_results);
+
+      auto graph_op = mlir::nv_tensor_ir::GraphOp::create(
+          builder, func_op.getLoc(), func_op.getName(),
+          /*sym_visibility=*/nullptr, converted_func_type,
+          func_op.getArgAttrsAttr(), func_op.getResAttrsAttr());
+
+      // Move the function's body region into the graph's graphBody region.
+      graph_op.getGraphBody().takeBody(func_op.getBody());
+
+      // Convert block argument types.
+      mlir::Block& entry_block = graph_op.getGraphBody().front();
+      for (auto [arg, new_type] :
+           llvm::zip(entry_block.getArguments(), new_inputs)) {
+        arg.setType(new_type);
+      }
+
+      // Replace the terminating mlir::func::ReturnOp with
+      // mlir::nv_tensor_ir::ResultsOp carrying the same operands.
+      if (auto return_op = mlir::dyn_cast<mlir::func::ReturnOp>(
+              entry_block.getTerminator())) {
+        builder.setInsertionPoint(return_op);
+        mlir::nv_tensor_ir::ResultsOp::create(builder, return_op.getLoc(),
+                                              return_op.getOperands());
+        return_op.erase();
+      }
+
+      func_op.erase();
+    }
+
+    mlir::ConversionTarget target(*context);
+    // Illegal by default. Only the dialects that legitimately survive into the
+    // output are legalized. StableHLO is not the only source dialect this pass
+    // sees: HloFunctionImporter emits mhlo ops (e.g. mhlo.bitcast, mhlo.erf)
+    // for constructs StableHLO has no equivalent for. If the target were only
+    // "stablehlo is illegal", those would be legal by default, survive inside
+    // the nv_tensor_ir.graph body, still verify and print, and only fail much
+    // later when nothing can lower them.
+    target.markUnknownOpDynamicallyLegal(
+        [](mlir::Operation*) { return false; });
+    target
+        .addLegalDialect<mlir::nv_tensor_ir::TensorIRDialect,
+                         mlir::arith::ArithDialect, mlir::func::FuncDialect>();
+    // The module is the operation being converted; the builtin dialect is
+    // otherwise illegal so that a leftover unrealized_conversion_cast is
+    // flagged too.
+    target.addLegalOp<mlir::ModuleOp>();
+    target.addDynamicallyLegalOp<mlir::nv_tensor_ir::ResultsOp>(
+        [&](mlir::nv_tensor_ir::ResultsOp op) {
+          return type_converter.isLegal(op);
+        });
+
+    mlir::RewritePatternSet patterns(context);
+    populateLegalizeStablehloToTensorIrPatterns(patterns, type_converter,
+                                                context);
+
+    if (mlir::failed(mlir::applyPartialConversion(module, target,
+                                                  std::move(patterns)))) {
+      // Deliberately no early return: the conversion rolled back, and the walk
+      // below names every op that was left illegal.
+      signalPassFailure();
+    }
+
+    // Some operands of the source IR only survive as attributes on the
+    // nv_tensor_ir side: a stablehlo.reduce init value becomes the `identity`
+    // attribute of reduce/reduce_ud, for example. Their (already converted)
+    // defining ops are then dead, and would otherwise be printed as unused
+    // rank-0 constants in the graph body. The dialect conversion driver does
+    // not clean those up, so do it here.
+    //
+    // This has to run to a fixpoint: deadness cascades, and a single walk only
+    // catches the first layer. `walk` is post-order, so an op is always visited
+    // before the op that consumes it, and a constant feeding a dead convert
+    // still looks live on the pass that removes the convert. Post-order also
+    // guarantees a nested op is erased before its parent, so erasing in
+    // collection order can never leave a dangling pointer.
+    bool erased_any = true;
+    while (erased_any) {
+      erased_any = false;
+      llvm::SmallVector<mlir::Operation*> dead_ops;
+      module.walk([&](mlir::Operation* op) {
+        if (op != module.getOperation() && mlir::isOpTriviallyDead(op)) {
+          dead_ops.push_back(op);
+        }
+      });
+      for (mlir::Operation* op : dead_ops) {
+        op->erase();
+        erased_any = true;
+      }
+    }
+
+    // applyPartialConversion only reports ops its patterns were asked to
+    // convert, so walk the module and flag anything the target considers
+    // illegal, naming the offending op.
+    bool found_illegal_op = false;
+    module.walk([&](mlir::Operation* op) {
+      if (op == module.getOperation()) {
+        return;
+      }
+      if (!target.isLegal(op)) {
+        op->emitOpError() << "operation was not legalized to nv_tensor_ir "
+                             "dialect";
+        found_illegal_op = true;
+      }
+    });
+
+    if (found_illegal_op) {
+      signalPassFailure();
+      return;
+    }
+
+    // Runs last, on a fully legalized graph, so that it only ever sees
+    // nv_tensor_ir types and can insert nv_tensor_ir.reshape directly.
+    for (auto graph : module.getOps<mlir::nv_tensor_ir::GraphOp>()) {
+      PromoteRank0Tensors(graph);
+    }
+  }
+};
+
+}  // namespace
+
+}  // namespace mlir::nv_tensor_ir::xla
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/passes.h b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/passes.h
new file mode 100644
index 00000000000..ab8781aaf22
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/passes.h
@@ -0,0 +1,35 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_TRANSFORMS_PASSES_H_
+#define XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_TRANSFORMS_PASSES_H_
+
+#include "tensor_ir/Dialect/TensorIR.h"  // IWYU pragma: keep
+#include "mlir/Dialect/Arith/IR/Arith.h"  // IWYU pragma: keep
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // IWYU pragma: keep
+#include "mlir/Pass/Pass.h"  // IWYU pragma: keep
+#include "stablehlo/dialect/StablehloOps.h"  // IWYU pragma: keep
+
+namespace mlir::nv_tensor_ir::xla {
+
+#define GEN_PASS_DECL
+#include "xla/backends/gpu/codegen/tensor_ir/transforms/passes.h.inc"
+
+#define GEN_PASS_REGISTRATION
+#include "xla/backends/gpu/codegen/tensor_ir/transforms/passes.h.inc"
+
+}  // namespace mlir::nv_tensor_ir::xla
+
+#endif  // XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_TRANSFORMS_PASSES_H_
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/passes.td b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/passes.td
new file mode 100644
index 00000000000..808b4e1c27b
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/passes.td
@@ -0,0 +1,39 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_TRANSFORMS_PASSES_TD_
+#define XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_TRANSFORMS_PASSES_TD_
+
+include "mlir/Pass/PassBase.td"
+
+def LegalizeStablehloToTensorIrPass
+    : Pass<"legalize-stablehlo-to-tensor-ir", "mlir::ModuleOp"> {
+  let summary = "Legalize StableHLO ops to the nv_tensor_ir dialect.";
+  let description = [{
+    Converts func.func functions to nv_tensor_ir.graph operations and legalizes
+    contained StableHLO operations to nv_tensor_ir operations.
+  }];
+  let dependentDialects = [
+    "::mlir::nv_tensor_ir::TensorIRDialect",
+    "::mlir::stablehlo::StablehloDialect",
+    // The HLO importer emits mhlo ops (erf, bitcast) for constructs StableHLO
+    // has no equivalent for, so the pass matches on those too.
+    "::mlir::mhlo::MhloDialect",
+    "::mlir::arith::ArithDialect",
+    "::mlir::func::FuncDialect",
+  ];
+}
+
+#endif  // XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_TRANSFORMS_PASSES_TD_
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/BUILD b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/BUILD
new file mode 100644
index 00000000000..aff7e1e5e26
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/BUILD
@@ -0,0 +1,39 @@
+# Copyright 2026 The OpenXLA 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.
+# ==============================================================================
+
+load("//xla:lit.bzl", "lit_test_suite")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [":friends"],
+    licenses = ["notice"],
+)
+
+package_group(
+    name = "friends",
+    includes = [
+        "//xla:friends",
+    ],
+)
+
+lit_test_suite(
+    name = "mlir_lit_tests",
+    srcs = glob(["*.mlir"]),
+    cfg = "//xla:lit.cfg.py",
+    tools = [
+        "//xla/backends/gpu/tests:xla-opt",
+        "@llvm-project//llvm:FileCheck",
+    ],
+)
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir.mlir b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir.mlir
new file mode 100644
index 00000000000..35dc33bfb1d
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir.mlir
@@ -0,0 +1,468 @@
+// RUN: xla-opt %s --legalize-stablehlo-to-tensor-ir --split-input-file | FileCheck %s
+
+// CHECK-LABEL: nv_tensor_ir.graph @identity
+// CHECK-SAME: (%[[IN:.*]]: tensor<8x16xf32>) -> tensor<8x16xf32>
+// CHECK-NEXT: results %[[IN]] : tensor<8x16xf32>
+func.func @identity(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
+  return %arg0 : tensor<8x16xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @attributes_preserved
+// CHECK-SAME: (%[[IN:.*]]: tensor<8x16xf32> {nv_tensor_ir.alignment = 16 : i64, nv_tensor_ir.stride = "(16,1)"})
+// CHECK-SAME: -> (tensor<8x16xf32> {nv_tensor_ir.alignment = 32 : i64, nv_tensor_ir.stride = "(16,1)"})
+// CHECK-NEXT: results %[[IN]] : tensor<8x16xf32>
+func.func @attributes_preserved(
+    %arg0: tensor<8x16xf32> {nv_tensor_ir.alignment = 16 : i64, nv_tensor_ir.stride = "(16,1)"})
+    -> (tensor<8x16xf32> {nv_tensor_ir.alignment = 32 : i64, nv_tensor_ir.stride = "(16,1)"}) {
+  return %arg0 : tensor<8x16xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @multiple_results
+// CHECK-SAME: (%[[IN0:.*]]: tensor<8x16xf32>, %[[IN1:.*]]: tensor<4xf32>)
+// CHECK-SAME: -> (tensor<8x16xf32>, tensor<4xf32>)
+// CHECK-NEXT: results %[[IN0]], %[[IN1]] : tensor<8x16xf32>, tensor<4xf32>
+func.func @multiple_results(%arg0: tensor<8x16xf32>, %arg1: tensor<4xf32>) -> (tensor<8x16xf32>, tensor<4xf32>) {
+  return %arg0, %arg1 : tensor<8x16xf32>, tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @envelope_i32
+// CHECK-SAME: (%[[IN:.*]]: tensor<8x16xsi32>) -> tensor<8x16xsi32>
+// CHECK-NEXT: results %[[IN]] : tensor<8x16xsi32>
+func.func @envelope_i32(%arg0: tensor<8x16xi32>) -> tensor<8x16xi32> {
+  return %arg0 : tensor<8x16xi32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @envelope_ui32
+// CHECK-SAME: (%[[IN:.*]]: tensor<8x16xui32>) -> tensor<8x16xui32>
+// CHECK-NEXT: results %[[IN]] : tensor<8x16xui32>
+func.func @envelope_ui32(%arg0: tensor<8x16xui32>) -> tensor<8x16xui32> {
+  return %arg0 : tensor<8x16xui32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @envelope_i1
+// CHECK-SAME: (%[[IN:.*]]: tensor<8x16xi1>) -> tensor<8x16xi1>
+// CHECK-NEXT: results %[[IN]] : tensor<8x16xi1>
+func.func @envelope_i1(%arg0: tensor<8x16xi1>) -> tensor<8x16xi1> {
+  return %arg0 : tensor<8x16xi1>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @abs
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = abs %[[IN]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @abs(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.abs %arg0 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @ceil
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = ceil %[[IN]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @ceil(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.ceil %arg0 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @floor
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = floor %[[IN]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @floor(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.floor %arg0 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @negate
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = neg %[[IN]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @negate(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.negate %arg0 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @cosine
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = cos %[[IN]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @cosine(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.cosine %arg0 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @sine
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = sin %[[IN]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @sine(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.sine %arg0 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @tan
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = tan %[[IN]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @tan(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.tan %arg0 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @exponential
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = exp %[[IN]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @exponential(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.exponential %arg0 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @log
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = log %[[IN]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @log(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.log %arg0 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @rsqrt
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = rsqrt %[[IN]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @rsqrt(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.rsqrt %arg0 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @sqrt
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = sqrt %[[IN]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @sqrt(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.sqrt %arg0 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @tanh
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = tanh_fwd %[[IN]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @tanh(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.tanh %arg0 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @logistic
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = sigmoid_fwd %[[IN]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @logistic(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.logistic %arg0 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @not
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xi1>) -> tensor<4xi1>
+// CHECK-NEXT: %[[RES:.*]] = not %[[IN]] : tensor<4xi1>
+// CHECK-NEXT: results %[[RES]] : tensor<4xi1>
+func.func @not(%arg0: tensor<4xi1>) -> tensor<4xi1> {
+  %0 = stablehlo.not %arg0 : tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+// StableHLO has no erf operation, so the HLO importer emits the mhlo one and
+// this pass has to accept it.
+// CHECK-LABEL: nv_tensor_ir.graph @erf
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = erf %[[IN]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @erf(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = mhlo.erf %arg0 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// nv_tensor_ir has no expm1, so it is decomposed into exp(x) - 1.
+// CHECK-LABEL: nv_tensor_ir.graph @exponential_minus_one
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[ONE:.*]] = constant dense<1.000000e+00> : tensor<4xf32>
+// CHECK-NEXT: %[[EXP:.*]] = exp %[[IN]] : tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = sub %[[EXP]], %[[ONE]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @exponential_minus_one(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.exponential_minus_one %arg0 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// Likewise, log1p is decomposed into log(x + 1).
+// CHECK-LABEL: nv_tensor_ir.graph @log_plus_one
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[ONE:.*]] = constant dense<1.000000e+00> : tensor<4xf32>
+// CHECK-NEXT: %[[ADD:.*]] = add %[[IN]], %[[ONE]] : tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = log %[[ADD]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @log_plus_one(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.log_plus_one %arg0 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @add
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xf32>, %[[RHS:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = add %[[LHS]], %[[RHS]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @add(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.add %arg0, %arg1 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @subtract
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xf32>, %[[RHS:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = sub %[[LHS]], %[[RHS]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @subtract(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.subtract %arg0, %arg1 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @multiply
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xf32>, %[[RHS:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = mul %[[LHS]], %[[RHS]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @multiply(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.multiply %arg0, %arg1 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @divide
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xf32>, %[[RHS:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = div %[[LHS]], %[[RHS]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @divide(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.divide %arg0, %arg1 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @maximum
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xf32>, %[[RHS:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = max %[[LHS]], %[[RHS]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @maximum(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.maximum %arg0, %arg1 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @minimum
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xf32>, %[[RHS:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = min %[[LHS]], %[[RHS]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @minimum(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.minimum %arg0, %arg1 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @power
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xf32>, %[[RHS:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = pow %[[LHS]], %[[RHS]] : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @power(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.power %arg0, %arg1 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @remainder
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xsi32>, %[[RHS:.*]]: tensor<4xsi32>) -> tensor<4xsi32>
+// CHECK-NEXT: %[[RES:.*]] = rem %[[LHS]], %[[RHS]] : tensor<4xsi32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xsi32>
+func.func @remainder(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi32> {
+  %0 = stablehlo.remainder %arg0, %arg1 : tensor<4xi32>
+  return %0 : tensor<4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @atan2
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xf32>, %[[RHS:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = atan2 %[[LHS]], %[[RHS]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @atan2(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.atan2 %arg0, %arg1 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @and
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xi1>, %[[RHS:.*]]: tensor<4xi1>) -> tensor<4xi1>
+// CHECK-NEXT: %[[RES:.*]] = and %[[LHS]], %[[RHS]] : tensor<4xi1>
+// CHECK-NEXT: results %[[RES]] : tensor<4xi1>
+func.func @and(%arg0: tensor<4xi1>, %arg1: tensor<4xi1>) -> tensor<4xi1> {
+  %0 = stablehlo.and %arg0, %arg1 : tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @or
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xi1>, %[[RHS:.*]]: tensor<4xi1>) -> tensor<4xi1>
+// CHECK-NEXT: %[[RES:.*]] = or %[[LHS]], %[[RHS]] : tensor<4xi1>
+// CHECK-NEXT: results %[[RES]] : tensor<4xi1>
+func.func @or(%arg0: tensor<4xi1>, %arg1: tensor<4xi1>) -> tensor<4xi1> {
+  %0 = stablehlo.or %arg0, %arg1 : tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @select
+// CHECK-SAME: (%[[PRED:.*]]: tensor<4xi1>, %[[ON_TRUE:.*]]: tensor<4xf32>, %[[ON_FALSE:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = binary_select %[[PRED]], %[[ON_TRUE]], %[[ON_FALSE]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @select(%arg0: tensor<4xi1>, %arg1: tensor<4xf32>, %arg2: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.select %arg0, %arg1, %arg2 : tensor<4xi1>, tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// nv_tensor_ir has no clamp, so it is expressed as min(max, max(min, x)). The
+// nesting and operand order match the string-building converter this pass
+// replaces.
+// CHECK-LABEL: nv_tensor_ir.graph @clamp
+// CHECK-SAME: (%[[MIN:.*]]: tensor<4xf32>, %[[IN:.*]]: tensor<4xf32>, %[[MAX:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[LOWER:.*]] = max %[[MIN]], %[[IN]] : tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = min %[[MAX]], %[[LOWER]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @clamp(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>, %arg2: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.clamp %arg0, %arg1, %arg2 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// nv_tensor_ir.max and nv_tensor_ir.min pick the signed comparison from the
+// signed element type, so no extra conversion is needed.
+// CHECK-LABEL: nv_tensor_ir.graph @clamp_signed
+// CHECK-SAME: (%[[MIN:.*]]: tensor<4xsi32>, %[[IN:.*]]: tensor<4xsi32>, %[[MAX:.*]]: tensor<4xsi32>) -> tensor<4xsi32>
+// CHECK-NEXT: %[[LOWER:.*]] = max %[[MIN]], %[[IN]] : tensor<4xsi32>
+// CHECK-NEXT: %[[RES:.*]] = min %[[MAX]], %[[LOWER]] : tensor<4xsi32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xsi32>
+func.func @clamp_signed(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>, %arg2: tensor<4xi32>) -> tensor<4xi32> {
+  %0 = stablehlo.clamp %arg0, %arg1, %arg2 : tensor<4xi32>
+  return %0 : tensor<4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @compare_float
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xf32>, %[[RHS:.*]]: tensor<4xf32>) -> tensor<4xi1>
+// CHECK-NEXT: %[[RES:.*]] = cmp %[[LHS]] olt %[[RHS]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xi1>
+func.func @compare_float(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xi1> {
+  %0 = stablehlo.compare LT, %arg0, %arg1, FLOAT : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @compare_signed
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xsi32>, %[[RHS:.*]]: tensor<4xsi32>) -> tensor<4xi1>
+// CHECK-NEXT: %[[RES:.*]] = cmp %[[LHS]] ge %[[RHS]] : tensor<4xsi32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xi1>
+func.func @compare_signed(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi1> {
+  %0 = stablehlo.compare GE, %arg0, %arg1, SIGNED : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @compare_unsigned
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xui32>, %[[RHS:.*]]: tensor<4xui32>) -> tensor<4xi1>
+// CHECK-NEXT: %[[RES:.*]] = cmp %[[LHS]] eq %[[RHS]] : tensor<4xui32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xi1>
+func.func @compare_unsigned(%arg0: tensor<4xui32>, %arg1: tensor<4xui32>) -> tensor<4xi1> {
+  %0 = stablehlo.compare EQ, %arg0, %arg1, UNSIGNED : (tensor<4xui32>, tensor<4xui32>) -> tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @convert_f32_to_i32
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xf32>) -> tensor<4xsi32>
+// CHECK-NEXT: %[[RES:.*]] = convert %[[IN]] : tensor<4xf32> -> tensor<4xsi32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xsi32>
+func.func @convert_f32_to_i32(%arg0: tensor<4xf32>) -> tensor<4xi32> {
+  %0 = stablehlo.convert %arg0 : (tensor<4xf32>) -> tensor<4xi32>
+  return %0 : tensor<4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @convert_i32_to_f32
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xsi32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = convert %[[IN]] : tensor<4xsi32> -> tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @convert_i32_to_f32(%arg0: tensor<4xi32>) -> tensor<4xf32> {
+  %0 = stablehlo.convert %arg0 : (tensor<4xi32>) -> tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_bitcast.mlir b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_bitcast.mlir
new file mode 100644
index 00000000000..35d1e941fad
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_bitcast.mlir
@@ -0,0 +1,74 @@
+// RUN: xla-opt %s --legalize-stablehlo-to-tensor-ir --split-input-file | FileCheck %s
+
+// mhlo.bitcast reinterprets a buffer under a new shape and layout without
+// moving any element. nv_tensor_ir tensors are always row-major, so the
+// lowering normalizes both sides to row-major and relinearizes in between:
+// transpose into physical order, reshape, transpose back into logical order.
+// Steps that would be the identity are skipped.
+//
+// The layouts come from the `source_layout` and `result_layout` attributes the
+// HLO importer attaches to every mhlo.bitcast. They are discardable attributes
+// (mhlo.bitcast declares none in ODS), so they are the only record of the
+// physical element order in the imported module.
+
+// Both sides carry a non-default layout, and the two physical shapes agree, so
+// the relinearizing reshape is skipped and only the two transposes remain.
+// CHECK-LABEL: nv_tensor_ir.graph @bitcast_both_non_default
+// CHECK-SAME: (%[[IN:.*]]: tensor<8x16x32xf32>) -> tensor<16x8x32xf32>
+// CHECK-NEXT: %[[PHYSICAL:.*]] = transpose %[[IN]] permutation = [2, 1, 0] : tensor<8x16x32xf32> -> tensor<32x16x8xf32>
+// CHECK-NEXT: %[[RES:.*]] = transpose %[[PHYSICAL]] permutation = [1, 2, 0] : tensor<32x16x8xf32> -> tensor<16x8x32xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<16x8x32xf32>
+func.func @bitcast_both_non_default(%arg0: tensor<8x16x32xf32>) -> tensor<16x8x32xf32> {
+  %0 = mhlo.bitcast %arg0 {result_layout = dense<[1, 0, 2]> : tensor<3xindex>, source_layout = dense<[0, 1, 2]> : tensor<3xindex>} : (tensor<8x16x32xf32>) -> tensor<16x8x32xf32>
+  return %0 : tensor<16x8x32xf32>
+}
+
+// -----
+
+// Both sides are row-major already, so the bitcast is a plain reshape.
+// CHECK-LABEL: nv_tensor_ir.graph @bitcast_default_layouts
+// CHECK-SAME: (%[[IN:.*]]: tensor<8x16xf32>) -> tensor<8x4x4xf32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[IN]] : tensor<8x16xf32> -> tensor<8x4x4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<8x4x4xf32>
+func.func @bitcast_default_layouts(%arg0: tensor<8x16xf32>) -> tensor<8x4x4xf32> {
+  %0 = mhlo.bitcast %arg0 {result_layout = dense<[2, 1, 0]> : tensor<3xindex>, source_layout = dense<[1, 0]> : tensor<2xindex>} : (tensor<8x16xf32>) -> tensor<8x4x4xf32>
+  return %0 : tensor<8x4x4xf32>
+}
+
+// -----
+
+// Only the operand is stored transposed, so only the first transpose is
+// emitted and the trailing one is skipped.
+// CHECK-LABEL: nv_tensor_ir.graph @bitcast_source_non_default
+// CHECK-SAME: (%[[IN:.*]]: tensor<8x16xf32>) -> tensor<4x4x8xf32>
+// CHECK-NEXT: %[[PHYSICAL:.*]] = transpose %[[IN]] permutation = [1, 0] : tensor<8x16xf32> -> tensor<16x8xf32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[PHYSICAL]] : tensor<16x8xf32> -> tensor<4x4x8xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4x4x8xf32>
+func.func @bitcast_source_non_default(%arg0: tensor<8x16xf32>) -> tensor<4x4x8xf32> {
+  %0 = mhlo.bitcast %arg0 {result_layout = dense<[2, 1, 0]> : tensor<3xindex>, source_layout = dense<[0, 1]> : tensor<2xindex>} : (tensor<8x16xf32>) -> tensor<4x4x8xf32>
+  return %0 : tensor<4x4x8xf32>
+}
+
+// -----
+
+// Only the result is stored transposed.
+// CHECK-LABEL: nv_tensor_ir.graph @bitcast_result_non_default
+// CHECK-SAME: (%[[IN:.*]]: tensor<8x16xf32>) -> tensor<16x8xf32>
+// CHECK-NEXT: %[[RES:.*]] = transpose %[[IN]] permutation = [1, 0] : tensor<8x16xf32> -> tensor<16x8xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<16x8xf32>
+func.func @bitcast_result_non_default(%arg0: tensor<8x16xf32>) -> tensor<16x8xf32> {
+  %0 = mhlo.bitcast %arg0 {result_layout = dense<[0, 1]> : tensor<2xindex>, source_layout = dense<[1, 0]> : tensor<2xindex>} : (tensor<8x16xf32>) -> tensor<16x8xf32>
+  return %0 : tensor<16x8xf32>
+}
+
+// -----
+
+// Same shape, same default layout: every step is the identity, so the bitcast
+// disappears entirely.
+// CHECK-LABEL: nv_tensor_ir.graph @bitcast_identity
+// CHECK-SAME: (%[[IN:.*]]: tensor<8x16xf32>) -> tensor<8x16xf32>
+// CHECK-NEXT: results %[[IN]] : tensor<8x16xf32>
+func.func @bitcast_identity(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
+  %0 = mhlo.bitcast %arg0 {result_layout = dense<[1, 0]> : tensor<2xindex>, source_layout = dense<[1, 0]> : tensor<2xindex>} : (tensor<8x16xf32>) -> tensor<8x16xf32>
+  return %0 : tensor<8x16xf32>
+}
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_compare.mlir b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_compare.mlir
new file mode 100644
index 00000000000..386765ec672
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_compare.mlir
@@ -0,0 +1,136 @@
+// RUN: xla-opt %s --legalize-stablehlo-to-tensor-ir --split-input-file | FileCheck %s
+
+// The twelve rows of the comparison-direction mapping. Float directions map to
+// the ordered comparators except NE, which maps to the unordered `une` because
+// IEEE 754 defines `NaN != x` as true.
+
+// CHECK-LABEL: nv_tensor_ir.graph @compare_float_eq
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xf32>, %[[RHS:.*]]: tensor<4xf32>) -> tensor<4xi1>
+// CHECK-NEXT: %[[RES:.*]] = cmp %[[LHS]] oeq %[[RHS]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xi1>
+func.func @compare_float_eq(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xi1> {
+  %0 = stablehlo.compare EQ, %arg0, %arg1 : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+// NE is the one asymmetric row: `une`, not `one`.
+// CHECK-LABEL: nv_tensor_ir.graph @compare_float_ne
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xf32>, %[[RHS:.*]]: tensor<4xf32>) -> tensor<4xi1>
+// CHECK-NEXT: %[[RES:.*]] = cmp %[[LHS]] une %[[RHS]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xi1>
+func.func @compare_float_ne(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xi1> {
+  %0 = stablehlo.compare NE, %arg0, %arg1 : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @compare_float_gt
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xf32>, %[[RHS:.*]]: tensor<4xf32>) -> tensor<4xi1>
+// CHECK-NEXT: %[[RES:.*]] = cmp %[[LHS]] ogt %[[RHS]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xi1>
+func.func @compare_float_gt(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xi1> {
+  %0 = stablehlo.compare GT, %arg0, %arg1 : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @compare_float_ge
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xf32>, %[[RHS:.*]]: tensor<4xf32>) -> tensor<4xi1>
+// CHECK-NEXT: %[[RES:.*]] = cmp %[[LHS]] oge %[[RHS]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xi1>
+func.func @compare_float_ge(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xi1> {
+  %0 = stablehlo.compare GE, %arg0, %arg1 : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @compare_float_lt
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xf32>, %[[RHS:.*]]: tensor<4xf32>) -> tensor<4xi1>
+// CHECK-NEXT: %[[RES:.*]] = cmp %[[LHS]] olt %[[RHS]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xi1>
+func.func @compare_float_lt(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xi1> {
+  %0 = stablehlo.compare LT, %arg0, %arg1 : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @compare_float_le
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xf32>, %[[RHS:.*]]: tensor<4xf32>) -> tensor<4xi1>
+// CHECK-NEXT: %[[RES:.*]] = cmp %[[LHS]] ole %[[RHS]] : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xi1>
+func.func @compare_float_le(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xi1> {
+  %0 = stablehlo.compare LE, %arg0, %arg1 : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @compare_integer_eq
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xsi32>, %[[RHS:.*]]: tensor<4xsi32>) -> tensor<4xi1>
+// CHECK-NEXT: %[[RES:.*]] = cmp %[[LHS]] eq %[[RHS]] : tensor<4xsi32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xi1>
+func.func @compare_integer_eq(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi1> {
+  %0 = stablehlo.compare EQ, %arg0, %arg1 : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @compare_integer_ne
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xsi32>, %[[RHS:.*]]: tensor<4xsi32>) -> tensor<4xi1>
+// CHECK-NEXT: %[[RES:.*]] = cmp %[[LHS]] neq %[[RHS]] : tensor<4xsi32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xi1>
+func.func @compare_integer_ne(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi1> {
+  %0 = stablehlo.compare NE, %arg0, %arg1 : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @compare_integer_gt
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xsi32>, %[[RHS:.*]]: tensor<4xsi32>) -> tensor<4xi1>
+// CHECK-NEXT: %[[RES:.*]] = cmp %[[LHS]] gt %[[RHS]] : tensor<4xsi32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xi1>
+func.func @compare_integer_gt(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi1> {
+  %0 = stablehlo.compare GT, %arg0, %arg1 : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @compare_integer_ge
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xsi32>, %[[RHS:.*]]: tensor<4xsi32>) -> tensor<4xi1>
+// CHECK-NEXT: %[[RES:.*]] = cmp %[[LHS]] ge %[[RHS]] : tensor<4xsi32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xi1>
+func.func @compare_integer_ge(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi1> {
+  %0 = stablehlo.compare GE, %arg0, %arg1 : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @compare_integer_lt
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xsi32>, %[[RHS:.*]]: tensor<4xsi32>) -> tensor<4xi1>
+// CHECK-NEXT: %[[RES:.*]] = cmp %[[LHS]] lt %[[RHS]] : tensor<4xsi32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xi1>
+func.func @compare_integer_lt(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi1> {
+  %0 = stablehlo.compare LT, %arg0, %arg1 : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @compare_integer_le
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4xsi32>, %[[RHS:.*]]: tensor<4xsi32>) -> tensor<4xi1>
+// CHECK-NEXT: %[[RES:.*]] = cmp %[[LHS]] le %[[RHS]] : tensor<4xsi32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xi1>
+func.func @compare_integer_le(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi1> {
+  %0 = stablehlo.compare LE, %arg0, %arg1 : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_invalid.mlir b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_invalid.mlir
new file mode 100644
index 00000000000..9a0adf0e3f0
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_invalid.mlir
@@ -0,0 +1,172 @@
+// RUN: xla-opt %s --legalize-stablehlo-to-tensor-ir --verify-diagnostics --split-input-file
+
+func.func @unsupported_xor(%arg0: tensor<4xi1>, %arg1: tensor<4xi1>) -> tensor<4xi1> {
+  // expected-error @+2 {{failed to legalize operation 'stablehlo.xor'}}
+  // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+  %0 = stablehlo.xor %arg0, %arg1 : tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+func.func @unsupported_totalorder(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xi1> {
+  // expected-error @+2 {{failed to legalize operation 'stablehlo.compare'}}
+  // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+  %0 = stablehlo.compare LT, %arg0, %arg1, TOTALORDER : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xi1>
+  return %0 : tensor<4xi1>
+}
+
+// -----
+
+// nv_tensor_ir.matmul has nowhere to record a precision request, so anything
+// other than DEFAULT is rejected instead of being silently dropped.
+func.func @dot_general_non_default_precision(%arg0: tensor<8x16xf32>, %arg1: tensor<16x32xf32>) -> tensor<8x32xf32> {
+  // expected-error @+2 {{failed to legalize operation 'stablehlo.dot_general'}}
+  // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+  %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [HIGHEST, HIGHEST] : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32>
+  return %0 : tensor<8x32xf32>
+}
+
+// -----
+
+// Likewise for an explicit dot algorithm.
+func.func @dot_general_algorithm(%arg0: tensor<8x16xf32>, %arg1: tensor<16x32xf32>) -> tensor<8x32xf32> {
+  // expected-error @+2 {{failed to legalize operation 'stablehlo.dot_general'}}
+  // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+  %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], algorithm = <lhs_precision_type = tf32, rhs_precision_type = tf32, accumulation_type = f32, lhs_component_count = 1, rhs_component_count = 1, num_primitive_operations = 1, allow_imprecise_accumulation = false> : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32>
+  return %0 : tensor<8x32xf32>
+}
+
+// -----
+
+func.func @reduce_non_constant_init(%arg0: tensor<4x8xf32>, %arg1: tensor<f32>) -> tensor<4xf32> {
+  // expected-error @+4 {{failed to legalize operation 'stablehlo.reduce'}}
+  // expected-error @+3 {{'stablehlo.add' op operation was not legalized to nv_tensor_ir dialect}}
+  // expected-error @+2 {{'stablehlo.return' op operation was not legalized to nv_tensor_ir dialect}}
+  // expected-error @+1 {{'stablehlo.reduce' op operation was not legalized to nv_tensor_ir dialect}}
+  %0 = stablehlo.reduce(%arg0 init: %arg1) applies stablehlo.add across dimensions = [1] : (tensor<4x8xf32>, tensor<f32>) -> tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// stablehlo.subtract is not one of the built-in nv_tensor_ir reduction modes.
+func.func @reduce_unsupported_combiner(%arg0: tensor<4x8xf32>) -> tensor<4xf32> {
+  // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+  %init = stablehlo.constant dense<0.000000e+00> : tensor<f32>
+  // expected-error @+2 {{failed to legalize operation 'stablehlo.reduce'}}
+  // expected-error @+1 {{'stablehlo.reduce' op operation was not legalized to nv_tensor_ir dialect}}
+  %0 = stablehlo.reduce(%arg0 init: %init) across dimensions = [1] : (tensor<4x8xf32>, tensor<f32>) -> tensor<4xf32>
+   reducer(%lhs: tensor<f32>, %rhs: tensor<f32>) {
+    // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+    %1 = stablehlo.subtract %lhs, %rhs : tensor<f32>
+    // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+    stablehlo.return %1 : tensor<f32>
+  }
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// arith has no total-order comparison, in a reduction body just as much as in
+// an elementwise position.
+func.func @reduce_totalorder_combiner(%arg0: tensor<4x8xf32>) -> tensor<4xf32> {
+  // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+  %init = stablehlo.constant dense<0.000000e+00> : tensor<f32>
+  // expected-error @+2 {{failed to legalize operation 'stablehlo.reduce'}}
+  // expected-error @+1 {{'stablehlo.reduce' op operation was not legalized to nv_tensor_ir dialect}}
+  %0 = stablehlo.reduce(%arg0 init: %init) across dimensions = [1] : (tensor<4x8xf32>, tensor<f32>) -> tensor<4xf32>
+   reducer(%lhs: tensor<f32>, %rhs: tensor<f32>) {
+    // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+    %1 = stablehlo.compare LT, %lhs, %rhs, TOTALORDER : (tensor<f32>, tensor<f32>) -> tensor<i1>
+    // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+    %2 = stablehlo.select %1, %lhs, %rhs : tensor<i1>, tensor<f32>
+    // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+    stablehlo.return %2 : tensor<f32>
+  }
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// Only single-input, single-result reductions are supported.
+func.func @reduce_variadic(%arg0: tensor<4x8xf32>, %arg1: tensor<4x8xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
+  // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+  %init = stablehlo.constant dense<0.000000e+00> : tensor<f32>
+  // expected-error @+2 {{failed to legalize operation 'stablehlo.reduce'}}
+  // expected-error @+1 {{'stablehlo.reduce' op operation was not legalized to nv_tensor_ir dialect}}
+  %0:2 = stablehlo.reduce(%arg0 init: %init), (%arg1 init: %init) across dimensions = [1] : (tensor<4x8xf32>, tensor<4x8xf32>, tensor<f32>, tensor<f32>) -> (tensor<4xf32>, tensor<4xf32>)
+   reducer(%lhs0: tensor<f32>, %rhs0: tensor<f32>) (%lhs1: tensor<f32>, %rhs1: tensor<f32>) {
+    // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+    %1 = stablehlo.add %lhs0, %rhs0 : tensor<f32>
+    // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+    %2 = stablehlo.add %lhs1, %rhs1 : tensor<f32>
+    // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+    stablehlo.return %1, %2 : tensor<f32>, tensor<f32>
+  }
+  return %0#0, %0#1 : tensor<4xf32>, tensor<4xf32>
+}
+
+// -----
+
+// The nv_tensor_ir dialect has no pad operation.
+func.func @unsupported_pad(%arg0: tensor<4xf32>, %arg1: tensor<f32>) -> tensor<6xf32> {
+  // expected-error @+2 {{failed to legalize operation 'stablehlo.pad'}}
+  // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+  %0 = stablehlo.pad %arg0, %arg1, low = [1], high = [1], interior = [0] : (tensor<4xf32>, tensor<f32>) -> tensor<6xf32>
+  return %0 : tensor<6xf32>
+}
+
+// -----
+
+// StableHLO allows rank-0 clamp bounds that are broadcast against the operand,
+// but nv_tensor_ir.max and nv_tensor_ir.min require matching shapes. The form
+// is rejected rather than lowered into something that would not verify.
+func.func @clamp_scalar_bounds(%arg0: tensor<f32>, %arg1: tensor<4xf32>, %arg2: tensor<f32>) -> tensor<4xf32> {
+  // expected-error @+2 {{failed to legalize operation 'stablehlo.clamp'}}
+  // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+  %0 = stablehlo.clamp %arg0, %arg1, %arg2 : (tensor<f32>, tensor<4xf32>, tensor<f32>) -> tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// source_layout and result_layout are discardable attributes, so they can be
+// dropped by any rewrite that does not copy them. Without them the physical
+// element order is unknowable and assuming the default layout would silently
+// miscompile, so the bitcast is rejected instead.
+func.func @bitcast_without_layouts(%arg0: tensor<8x16xf32>) -> tensor<8x4x4xf32> {
+  // expected-error @+2 {{failed to legalize operation 'mhlo.bitcast'}}
+  // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+  %0 = mhlo.bitcast %arg0 : (tensor<8x16xf32>) -> tensor<8x4x4xf32>
+  return %0 : tensor<8x4x4xf32>
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// Dialect leakage. The conversion target is illegal by default, so an op from
+// any dialect other than nv_tensor_ir/arith/func fails the pass instead of
+// surviving inside the nv_tensor_ir.graph body.
+//
+// The motivating producers are the mhlo ops that HloFunctionImporter emits for
+// constructs StableHLO cannot express. mhlo.bitcast and mhlo.erf are handled
+// by this pass now, so two unhandled ops with the same shape stand in for the
+// general case here.
+//===----------------------------------------------------------------------===//
+
+func.func @unsupported_math_dialect(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  // expected-error @+2 {{failed to legalize operation 'math.absf'}}
+  // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+  %0 = math.absf %arg0 : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+func.func @unsupported_tensor_dialect(%arg0: tensor<4xi32>) -> tensor<4xf32> {
+  // expected-error @+2 {{failed to legalize operation 'tensor.bitcast'}}
+  // expected-error @+1 {{operation was not legalized to nv_tensor_ir dialect}}
+  %0 = tensor.bitcast %arg0 : tensor<4xi32> to tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_matmul.mlir b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_matmul.mlir
new file mode 100644
index 00000000000..70406bc6380
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_matmul.mlir
@@ -0,0 +1,241 @@
+// RUN: xla-opt %s --legalize-stablehlo-to-tensor-ir --split-input-file | FileCheck %s
+
+// CHECK-LABEL: nv_tensor_ir.graph @matmul
+// CHECK-SAME: (%[[LHS:.*]]: tensor<8x16xf32>, %[[RHS:.*]]: tensor<16x32xf32>) -> tensor<8x32xf32>
+// CHECK-NEXT: %[[RES:.*]] = matmul(%[[LHS]], %[[RHS]]) : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<8x32xf32>
+func.func @matmul(%arg0: tensor<8x16xf32>, %arg1: tensor<16x32xf32>) -> tensor<8x32xf32> {
+  %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32>
+  return %0 : tensor<8x32xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @matmul_i32
+// CHECK-SAME: (%[[LHS:.*]]: tensor<8x16xsi32>, %[[RHS:.*]]: tensor<16x32xsi32>) -> tensor<8x32xsi32>
+// CHECK-NEXT: %[[RES:.*]] = matmul(%[[LHS]], %[[RHS]]) : (tensor<8x16xsi32>, tensor<16x32xsi32>) -> tensor<8x32xsi32>
+// CHECK-NEXT: results %[[RES]] : tensor<8x32xsi32>
+func.func @matmul_i32(%arg0: tensor<8x16xi32>, %arg1: tensor<16x32xi32>) -> tensor<8x32xi32> {
+  %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<8x16xi32>, tensor<16x32xi32>) -> tensor<8x32xi32>
+  return %0 : tensor<8x32xi32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @matmul_batched
+// CHECK-SAME: (%[[LHS:.*]]: tensor<8x16x32xf32>, %[[RHS:.*]]: tensor<8x32x48xf32>) -> tensor<8x16x48xf32>
+// CHECK-NEXT: %[[RES:.*]] = matmul(%[[LHS]], %[[RHS]]) : (tensor<8x16x32xf32>, tensor<8x32x48xf32>) -> tensor<8x16x48xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<8x16x48xf32>
+func.func @matmul_batched(%arg0: tensor<8x16x32xf32>, %arg1: tensor<8x32x48xf32>) -> tensor<8x16x48xf32> {
+  %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<8x16x32xf32>, tensor<8x32x48xf32>) -> tensor<8x16x48xf32>
+  return %0 : tensor<8x16x48xf32>
+}
+
+// -----
+
+// The lhs contracts along its leading dimension, so it is transposed into the
+// (M, K) layout that nv_tensor_ir.matmul requires.
+// CHECK-LABEL: nv_tensor_ir.graph @matmul_lhs_transposed
+// CHECK-SAME: (%[[LHS:.*]]: tensor<16x8xf32>, %[[RHS:.*]]: tensor<16x32xf32>) -> tensor<8x32xf32>
+// CHECK-NEXT: %[[LHS_T:.*]] = transpose %[[LHS]] permutation = [1, 0] : tensor<16x8xf32> -> tensor<8x16xf32>
+// CHECK-NEXT: %[[RES:.*]] = matmul(%[[LHS_T]], %[[RHS]]) : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<8x32xf32>
+func.func @matmul_lhs_transposed(%arg0: tensor<16x8xf32>, %arg1: tensor<16x32xf32>) -> tensor<8x32xf32> {
+  %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [0] x [0] : (tensor<16x8xf32>, tensor<16x32xf32>) -> tensor<8x32xf32>
+  return %0 : tensor<8x32xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @matmul_rhs_transposed
+// CHECK-SAME: (%[[LHS:.*]]: tensor<8x16xf32>, %[[RHS:.*]]: tensor<32x16xf32>) -> tensor<8x32xf32>
+// CHECK-NEXT: %[[RHS_T:.*]] = transpose %[[RHS]] permutation = [1, 0] : tensor<32x16xf32> -> tensor<16x32xf32>
+// CHECK-NEXT: %[[RES:.*]] = matmul(%[[LHS]], %[[RHS_T]]) : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<8x32xf32>
+func.func @matmul_rhs_transposed(%arg0: tensor<8x16xf32>, %arg1: tensor<32x16xf32>) -> tensor<8x32xf32> {
+  %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [1] : (tensor<8x16xf32>, tensor<32x16xf32>) -> tensor<8x32xf32>
+  return %0 : tensor<8x32xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @matmul_batched_both_transposed
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4x32x16xf32>, %[[RHS:.*]]: tensor<4x48x32xf32>) -> tensor<4x16x48xf32>
+// CHECK-NEXT: %[[LHS_T:.*]] = transpose %[[LHS]] permutation = [0, 2, 1] : tensor<4x32x16xf32> -> tensor<4x16x32xf32>
+// CHECK-NEXT: %[[RHS_T:.*]] = transpose %[[RHS]] permutation = [0, 2, 1] : tensor<4x48x32xf32> -> tensor<4x32x48xf32>
+// CHECK-NEXT: %[[RES:.*]] = matmul(%[[LHS_T]], %[[RHS_T]]) : (tensor<4x16x32xf32>, tensor<4x32x48xf32>) -> tensor<4x16x48xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4x16x48xf32>
+func.func @matmul_batched_both_transposed(%arg0: tensor<4x32x16xf32>, %arg1: tensor<4x48x32xf32>) -> tensor<4x16x48xf32> {
+  %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [1] x [2] : (tensor<4x32x16xf32>, tensor<4x48x32xf32>) -> tensor<4x16x48xf32>
+  return %0 : tensor<4x16x48xf32>
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// The nine cases of the old string-building converter's specification,
+// .../codegen/tensor_ir/tests/dot.hlo, ported to StableHLO. The op sequence
+// below must keep matching what that file pins.
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: nv_tensor_ir.graph @test_dot_no_batch
+// CHECK-SAME: (%[[LHS:.*]]: tensor<8x32xf32>, %[[RHS:.*]]: tensor<32x16xf32>) -> tensor<8x16xf32>
+// CHECK-NEXT: %[[RES:.*]] = matmul(%[[LHS]], %[[RHS]]) : (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<8x16xf32>
+func.func @test_dot_no_batch(%arg0: tensor<8x32xf32>, %arg1: tensor<32x16xf32>) -> tensor<8x16xf32> {
+  %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<8x32xf32>, tensor<32x16xf32>) -> tensor<8x16xf32>
+  return %0 : tensor<8x16xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @test_dot_with_batch
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4x8x32xf32>, %[[RHS:.*]]: tensor<4x32x16xf32>) -> tensor<4x8x16xf32>
+// CHECK-NEXT: %[[RES:.*]] = matmul(%[[LHS]], %[[RHS]]) : (tensor<4x8x32xf32>, tensor<4x32x16xf32>) -> tensor<4x8x16xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4x8x16xf32>
+func.func @test_dot_with_batch(%arg0: tensor<4x8x32xf32>, %arg1: tensor<4x32x16xf32>) -> tensor<4x8x16xf32> {
+  %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<4x8x32xf32>, tensor<4x32x16xf32>) -> tensor<4x8x16xf32>
+  return %0 : tensor<4x8x16xf32>
+}
+
+// -----
+
+// A batch dimension that is not leading.
+// CHECK-LABEL: nv_tensor_ir.graph @test_dot_transpose_lhs
+// CHECK-SAME: (%[[LHS:.*]]: tensor<32x4x8xf32>, %[[RHS:.*]]: tensor<4x32x16xf32>) -> tensor<4x8x16xf32>
+// CHECK-NEXT: %[[LHS_T:.*]] = transpose %[[LHS]] permutation = [1, 2, 0] : tensor<32x4x8xf32> -> tensor<4x8x32xf32>
+// CHECK-NEXT: %[[RES:.*]] = matmul(%[[LHS_T]], %[[RHS]]) : (tensor<4x8x32xf32>, tensor<4x32x16xf32>) -> tensor<4x8x16xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4x8x16xf32>
+func.func @test_dot_transpose_lhs(%arg0: tensor<32x4x8xf32>, %arg1: tensor<4x32x16xf32>) -> tensor<4x8x16xf32> {
+  %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [1] x [0], contracting_dims = [0] x [1] : (tensor<32x4x8xf32>, tensor<4x32x16xf32>) -> tensor<4x8x16xf32>
+  return %0 : tensor<4x8x16xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @test_dot_transpose_rhs
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4x8x32xf32>, %[[RHS:.*]]: tensor<32x4x16xf32>) -> tensor<4x8x16xf32>
+// CHECK-NEXT: %[[RHS_T:.*]] = transpose %[[RHS]] permutation = [1, 0, 2] : tensor<32x4x16xf32> -> tensor<4x32x16xf32>
+// CHECK-NEXT: %[[RES:.*]] = matmul(%[[LHS]], %[[RHS_T]]) : (tensor<4x8x32xf32>, tensor<4x32x16xf32>) -> tensor<4x8x16xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4x8x16xf32>
+func.func @test_dot_transpose_rhs(%arg0: tensor<4x8x32xf32>, %arg1: tensor<32x4x16xf32>) -> tensor<4x8x16xf32> {
+  %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [1], contracting_dims = [2] x [0] : (tensor<4x8x32xf32>, tensor<32x4x16xf32>) -> tensor<4x8x16xf32>
+  return %0 : tensor<4x8x16xf32>
+}
+
+// -----
+
+// Multiple batch dimensions, collapsed into one.
+// CHECK-LABEL: nv_tensor_ir.graph @test_dot_reshape_batch
+// CHECK-SAME: (%[[LHS:.*]]: tensor<2x4x8x32xf32>, %[[RHS:.*]]: tensor<2x4x32x16xf32>) -> tensor<2x4x8x16xf32>
+// CHECK-NEXT: %[[LHS_R:.*]] = reshape %[[LHS]] : tensor<2x4x8x32xf32> -> tensor<8x8x32xf32>
+// CHECK-NEXT: %[[RHS_R:.*]] = reshape %[[RHS]] : tensor<2x4x32x16xf32> -> tensor<8x32x16xf32>
+// CHECK-NEXT: %[[MATMUL:.*]] = matmul(%[[LHS_R]], %[[RHS_R]]) : (tensor<8x8x32xf32>, tensor<8x32x16xf32>) -> tensor<8x8x16xf32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[MATMUL]] : tensor<8x8x16xf32> -> tensor<2x4x8x16xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<2x4x8x16xf32>
+func.func @test_dot_reshape_batch(%arg0: tensor<2x4x8x32xf32>, %arg1: tensor<2x4x32x16xf32>) -> tensor<2x4x8x16xf32> {
+  %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0, 1] x [0, 1], contracting_dims = [3] x [2] : (tensor<2x4x8x32xf32>, tensor<2x4x32x16xf32>) -> tensor<2x4x8x16xf32>
+  return %0 : tensor<2x4x8x16xf32>
+}
+
+// -----
+
+// Multiple contracting dimensions, merged into one.
+// CHECK-LABEL: nv_tensor_ir.graph @test_dot_reshape_contracting
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4x8x32x2xf32>, %[[RHS:.*]]: tensor<4x32x2x16xf32>) -> tensor<4x8x16xf32>
+// CHECK-NEXT: %[[LHS_R:.*]] = reshape %[[LHS]] : tensor<4x8x32x2xf32> -> tensor<4x8x64xf32>
+// CHECK-NEXT: %[[RHS_R:.*]] = reshape %[[RHS]] : tensor<4x32x2x16xf32> -> tensor<4x64x16xf32>
+// CHECK-NEXT: %[[RES:.*]] = matmul(%[[LHS_R]], %[[RHS_R]]) : (tensor<4x8x64xf32>, tensor<4x64x16xf32>) -> tensor<4x8x16xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4x8x16xf32>
+func.func @test_dot_reshape_contracting(%arg0: tensor<4x8x32x2xf32>, %arg1: tensor<4x32x2x16xf32>) -> tensor<4x8x16xf32> {
+  %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [2, 3] x [1, 2] : (tensor<4x8x32x2xf32>, tensor<4x32x2x16xf32>) -> tensor<4x8x16xf32>
+  return %0 : tensor<4x8x16xf32>
+}
+
+// -----
+
+// Two free dimensions on the lhs, merged into one.
+// CHECK-LABEL: nv_tensor_ir.graph @test_dot_reshape_lhs_noncontracting
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4x8x8x32xf32>, %[[RHS:.*]]: tensor<4x32x16xf32>) -> tensor<4x8x8x16xf32>
+// CHECK-NEXT: %[[LHS_R:.*]] = reshape %[[LHS]] : tensor<4x8x8x32xf32> -> tensor<4x64x32xf32>
+// CHECK-NEXT: %[[MATMUL:.*]] = matmul(%[[LHS_R]], %[[RHS]]) : (tensor<4x64x32xf32>, tensor<4x32x16xf32>) -> tensor<4x64x16xf32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[MATMUL]] : tensor<4x64x16xf32> -> tensor<4x8x8x16xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4x8x8x16xf32>
+func.func @test_dot_reshape_lhs_noncontracting(%arg0: tensor<4x8x8x32xf32>, %arg1: tensor<4x32x16xf32>) -> tensor<4x8x8x16xf32> {
+  %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [3] x [1] : (tensor<4x8x8x32xf32>, tensor<4x32x16xf32>) -> tensor<4x8x8x16xf32>
+  return %0 : tensor<4x8x8x16xf32>
+}
+
+// -----
+
+// Two free dimensions on the rhs, merged into one.
+// CHECK-LABEL: nv_tensor_ir.graph @test_dot_reshape_rhs_noncontracting
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4x8x32xf32>, %[[RHS:.*]]: tensor<4x32x16x16xf32>) -> tensor<4x8x16x16xf32>
+// CHECK-NEXT: %[[RHS_R:.*]] = reshape %[[RHS]] : tensor<4x32x16x16xf32> -> tensor<4x32x256xf32>
+// CHECK-NEXT: %[[MATMUL:.*]] = matmul(%[[LHS]], %[[RHS_R]]) : (tensor<4x8x32xf32>, tensor<4x32x256xf32>) -> tensor<4x8x256xf32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[MATMUL]] : tensor<4x8x256xf32> -> tensor<4x8x16x16xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4x8x16x16xf32>
+func.func @test_dot_reshape_rhs_noncontracting(%arg0: tensor<4x8x32xf32>, %arg1: tensor<4x32x16x16xf32>) -> tensor<4x8x16x16xf32> {
+  %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [0] x [0], contracting_dims = [2] x [1] : (tensor<4x8x32xf32>, tensor<4x32x16x16xf32>) -> tensor<4x8x16x16xf32>
+  return %0 : tensor<4x8x16x16xf32>
+}
+
+// -----
+
+// Non-leading and reordered batch dimensions, reordered multi-dimensional
+// contraction, and multiple free dimensions on both operands, all at once.
+// CHECK-LABEL: nv_tensor_ir.graph @test_dot_transpose_reshape_all
+// CHECK-SAME: (%[[LHS:.*]]: tensor<9x8x5x4x3x2xf32>, %[[RHS:.*]]: tensor<2x6x8x3x7x9xf32>) -> tensor<2x3x5x4x6x7xf32>
+// CHECK-NEXT: %[[LHS_T:.*]] = transpose %[[LHS]] permutation = [5, 4, 2, 3, 1, 0] : tensor<9x8x5x4x3x2xf32> -> tensor<2x3x5x4x8x9xf32>
+// CHECK-NEXT: %[[LHS_R:.*]] = reshape %[[LHS_T]] : tensor<2x3x5x4x8x9xf32> -> tensor<6x20x72xf32>
+// CHECK-NEXT: %[[RHS_T:.*]] = transpose %[[RHS]] permutation = [0, 3, 2, 5, 1, 4] : tensor<2x6x8x3x7x9xf32> -> tensor<2x3x8x9x6x7xf32>
+// CHECK-NEXT: %[[RHS_R:.*]] = reshape %[[RHS_T]] : tensor<2x3x8x9x6x7xf32> -> tensor<6x72x42xf32>
+// CHECK-NEXT: %[[MATMUL:.*]] = matmul(%[[LHS_R]], %[[RHS_R]]) : (tensor<6x20x72xf32>, tensor<6x72x42xf32>) -> tensor<6x20x42xf32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[MATMUL]] : tensor<6x20x42xf32> -> tensor<2x3x5x4x6x7xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<2x3x5x4x6x7xf32>
+func.func @test_dot_transpose_reshape_all(%arg0: tensor<9x8x5x4x3x2xf32>, %arg1: tensor<2x6x8x3x7x9xf32>) -> tensor<2x3x5x4x6x7xf32> {
+  %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [5, 4] x [0, 3], contracting_dims = [1, 0] x [2, 5] : (tensor<9x8x5x4x3x2xf32>, tensor<2x6x8x3x7x9xf32>) -> tensor<2x3x5x4x6x7xf32>
+  return %0 : tensor<2x3x5x4x6x7xf32>
+}
+
+// -----
+
+// Multiple contracting dimensions without any batch dimension: the collapsed
+// form is the plain rank-2 matmul, not a size-1 batched one.
+// CHECK-LABEL: nv_tensor_ir.graph @matmul_multiple_contracting_dims
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4x8x16xf32>, %[[RHS:.*]]: tensor<8x16x32xf32>) -> tensor<4x32xf32>
+// CHECK-NEXT: %[[LHS_R:.*]] = reshape %[[LHS]] : tensor<4x8x16xf32> -> tensor<4x128xf32>
+// CHECK-NEXT: %[[RHS_R:.*]] = reshape %[[RHS]] : tensor<8x16x32xf32> -> tensor<128x32xf32>
+// CHECK-NEXT: %[[RES:.*]] = matmul(%[[LHS_R]], %[[RHS_R]]) : (tensor<4x128xf32>, tensor<128x32xf32>) -> tensor<4x32xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4x32xf32>
+func.func @matmul_multiple_contracting_dims(%arg0: tensor<4x8x16xf32>, %arg1: tensor<8x16x32xf32>) -> tensor<4x32xf32> {
+  %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1, 2] x [0, 1] : (tensor<4x8x16xf32>, tensor<8x16x32xf32>) -> tensor<4x32xf32>
+  return %0 : tensor<4x32xf32>
+}
+
+// -----
+
+// A trailing batch dimension on both operands.
+// CHECK-LABEL: nv_tensor_ir.graph @matmul_trailing_batch_dim
+// CHECK-SAME: (%[[LHS:.*]]: tensor<4x8x16xf32>, %[[RHS:.*]]: tensor<16x8x32xf32>) -> tensor<8x4x32xf32>
+// CHECK-NEXT: %[[LHS_T:.*]] = transpose %[[LHS]] permutation = [1, 0, 2] : tensor<4x8x16xf32> -> tensor<8x4x16xf32>
+// CHECK-NEXT: %[[RHS_T:.*]] = transpose %[[RHS]] permutation = [1, 0, 2] : tensor<16x8x32xf32> -> tensor<8x16x32xf32>
+// CHECK-NEXT: %[[RES:.*]] = matmul(%[[LHS_T]], %[[RHS_T]]) : (tensor<8x4x16xf32>, tensor<8x16x32xf32>) -> tensor<8x4x32xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<8x4x32xf32>
+func.func @matmul_trailing_batch_dim(%arg0: tensor<4x8x16xf32>, %arg1: tensor<16x8x32xf32>) -> tensor<8x4x32xf32> {
+  %0 = stablehlo.dot_general %arg0, %arg1, batching_dims = [1] x [1], contracting_dims = [2] x [0] : (tensor<4x8x16xf32>, tensor<16x8x32xf32>) -> tensor<8x4x32xf32>
+  return %0 : tensor<8x4x32xf32>
+}
+
+// -----
+
+// An all-DEFAULT precision_config asks for nothing in particular, so it is
+// accepted and dropped. Anything else is rejected, see the invalid tests.
+// CHECK-LABEL: nv_tensor_ir.graph @matmul_default_precision
+// CHECK-SAME: (%[[LHS:.*]]: tensor<8x16xf32>, %[[RHS:.*]]: tensor<16x32xf32>) -> tensor<8x32xf32>
+// CHECK-NEXT: %[[RES:.*]] = matmul(%[[LHS]], %[[RHS]]) : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<8x32xf32>
+func.func @matmul_default_precision(%arg0: tensor<8x16xf32>, %arg1: tensor<16x32xf32>) -> tensor<8x32xf32> {
+  %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<8x16xf32>, tensor<16x32xf32>) -> tensor<8x32xf32>
+  return %0 : tensor<8x32xf32>
+}
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_reduce.mlir b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_reduce.mlir
new file mode 100644
index 00000000000..a4633bf0a8c
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_reduce.mlir
@@ -0,0 +1,283 @@
+// RUN: xla-opt %s --legalize-stablehlo-to-tensor-ir --split-input-file | FileCheck %s
+
+// nv_tensor_ir.reduce preserves the rank, so the reduced dimensions are
+// contracted to 1 and a reshape recovers the stablehlo result shape.
+// CHECK-LABEL: nv_tensor_ir.graph @reduce_add
+// CHECK-SAME: (%[[IN:.*]]: tensor<16x32x64xf32>) -> tensor<16x64xf32>
+// CHECK: %[[REDUCE:.*]] = reduce(%[[IN]]) <dimensions = [1], reduction_mode = <add>> : tensor<16x32x64xf32> -> tensor<16x1x64xf32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[REDUCE]] : tensor<16x1x64xf32> -> tensor<16x64xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<16x64xf32>
+func.func @reduce_add(%arg0: tensor<16x32x64xf32>) -> tensor<16x64xf32> {
+  %init = stablehlo.constant dense<0.000000e+00> : tensor<f32>
+  %0 = stablehlo.reduce(%arg0 init: %init) applies stablehlo.add across dimensions = [1] : (tensor<16x32x64xf32>, tensor<f32>) -> tensor<16x64xf32>
+  return %0 : tensor<16x64xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @reduce_mul
+// CHECK-SAME: (%[[IN:.*]]: tensor<4x8xf32>) -> tensor<4xf32>
+// CHECK: %[[REDUCE:.*]] = reduce(%[[IN]]) <dimensions = [1], reduction_mode = <mul>> : tensor<4x8xf32> -> tensor<4x1xf32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[REDUCE]] : tensor<4x1xf32> -> tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @reduce_mul(%arg0: tensor<4x8xf32>) -> tensor<4xf32> {
+  %init = stablehlo.constant dense<1.000000e+00> : tensor<f32>
+  %0 = stablehlo.reduce(%arg0 init: %init) applies stablehlo.multiply across dimensions = [1] : (tensor<4x8xf32>, tensor<f32>) -> tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// The init value has to be -inf, the identity of the max reduction.
+// CHECK-LABEL: nv_tensor_ir.graph @reduce_max
+// CHECK-SAME: (%[[IN:.*]]: tensor<4x8xf32>) -> tensor<4xf32>
+// CHECK: %[[REDUCE:.*]] = reduce(%[[IN]]) <dimensions = [1], reduction_mode = <max>> : tensor<4x8xf32> -> tensor<4x1xf32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[REDUCE]] : tensor<4x1xf32> -> tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @reduce_max(%arg0: tensor<4x8xf32>) -> tensor<4xf32> {
+  %init = stablehlo.constant dense<0xFF800000> : tensor<f32>
+  %0 = stablehlo.reduce(%arg0 init: %init) applies stablehlo.maximum across dimensions = [1] : (tensor<4x8xf32>, tensor<f32>) -> tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @reduce_min
+// CHECK-SAME: (%[[IN:.*]]: tensor<4x8xf32>) -> tensor<4xf32>
+// CHECK: %[[REDUCE:.*]] = reduce(%[[IN]]) <dimensions = [1], reduction_mode = <min>> : tensor<4x8xf32> -> tensor<4x1xf32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[REDUCE]] : tensor<4x1xf32> -> tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @reduce_min(%arg0: tensor<4x8xf32>) -> tensor<4xf32> {
+  %init = stablehlo.constant dense<0x7F800000> : tensor<f32>
+  %0 = stablehlo.reduce(%arg0 init: %init) applies stablehlo.minimum across dimensions = [1] : (tensor<4x8xf32>, tensor<f32>) -> tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @reduce_i32
+// CHECK-SAME: (%[[IN:.*]]: tensor<4x8xsi32>) -> tensor<4xsi32>
+// CHECK: %[[REDUCE:.*]] = reduce(%[[IN]]) <dimensions = [1], reduction_mode = <add>> : tensor<4x8xsi32> -> tensor<4x1xsi32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[REDUCE]] : tensor<4x1xsi32> -> tensor<4xsi32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xsi32>
+func.func @reduce_i32(%arg0: tensor<4x8xi32>) -> tensor<4xi32> {
+  %init = stablehlo.constant dense<0> : tensor<i32>
+  %0 = stablehlo.reduce(%arg0 init: %init) applies stablehlo.add across dimensions = [1] : (tensor<4x8xi32>, tensor<i32>) -> tensor<4xi32>
+  return %0 : tensor<4xi32>
+}
+
+// -----
+
+// The dimensions are narrowed to an i32 dense array and are not sorted.
+// CHECK-LABEL: nv_tensor_ir.graph @reduce_multiple_dimensions
+// CHECK-SAME: (%[[IN:.*]]: tensor<4x8x16xf32>) -> tensor<8xf32>
+// CHECK: %[[REDUCE:.*]] = reduce(%[[IN]]) <dimensions = [2, 0], reduction_mode = <add>> : tensor<4x8x16xf32> -> tensor<1x8x1xf32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[REDUCE]] : tensor<1x8x1xf32> -> tensor<8xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<8xf32>
+func.func @reduce_multiple_dimensions(%arg0: tensor<4x8x16xf32>) -> tensor<8xf32> {
+  %init = stablehlo.constant dense<0.000000e+00> : tensor<f32>
+  %0 = stablehlo.reduce(%arg0 init: %init) applies stablehlo.add across dimensions = [2, 0] : (tensor<4x8x16xf32>, tensor<f32>) -> tensor<8xf32>
+  return %0 : tensor<8xf32>
+}
+
+// -----
+
+// A rank-0 graph result is declared as `[1]`, which is exactly the shape the
+// reduce already produces, so no reshape survives. See `PromoteRank0Tensors`.
+// CHECK-LABEL: nv_tensor_ir.graph @reduce_to_scalar
+// CHECK-SAME: (%[[IN:.*]]: tensor<8xf32>) -> tensor<1xf32>
+// CHECK: %[[REDUCE:.*]] = reduce(%[[IN]]) <dimensions = [0], reduction_mode = <add>> : tensor<8xf32> -> tensor<1xf32>
+// CHECK-NEXT: results %[[REDUCE]] : tensor<1xf32>
+func.func @reduce_to_scalar(%arg0: tensor<8xf32>) -> tensor<f32> {
+  %init = stablehlo.constant dense<0.000000e+00> : tensor<f32>
+  %0 = stablehlo.reduce(%arg0 init: %init) applies stablehlo.add across dimensions = [0] : (tensor<8xf32>, tensor<f32>) -> tensor<f32>
+  return %0 : tensor<f32>
+}
+
+// -----
+
+// Both reductions are matched, i.e. the elementwise patterns do not rewrite the
+// stablehlo.add inside the first body before the second reduce is visited.
+// CHECK-LABEL: nv_tensor_ir.graph @two_reductions
+// CHECK-SAME: (%[[IN:.*]]: tensor<4x8xf32>) -> (tensor<4xf32>, tensor<4xf32>)
+// CHECK: %[[REDUCE0:.*]] = reduce(%[[IN]]) <dimensions = [1], reduction_mode = <add>> : tensor<4x8xf32> -> tensor<4x1xf32>
+// CHECK-NEXT: %[[RES0:.*]] = reshape %[[REDUCE0]] : tensor<4x1xf32> -> tensor<4xf32>
+// CHECK-NEXT: %[[REDUCE1:.*]] = reduce(%[[IN]]) <dimensions = [1], reduction_mode = <max>> : tensor<4x8xf32> -> tensor<4x1xf32>
+// CHECK-NEXT: %[[RES1:.*]] = reshape %[[REDUCE1]] : tensor<4x1xf32> -> tensor<4xf32>
+// CHECK-NEXT: results %[[RES0]], %[[RES1]] : tensor<4xf32>, tensor<4xf32>
+func.func @two_reductions(%arg0: tensor<4x8xf32>) -> (tensor<4xf32>, tensor<4xf32>) {
+  %zero = stablehlo.constant dense<0.000000e+00> : tensor<f32>
+  %neg_inf = stablehlo.constant dense<0xFF800000> : tensor<f32>
+  %0 = stablehlo.reduce(%arg0 init: %zero) applies stablehlo.add across dimensions = [1] : (tensor<4x8xf32>, tensor<f32>) -> tensor<4xf32>
+  %1 = stablehlo.reduce(%arg0 init: %neg_inf) applies stablehlo.maximum across dimensions = [1] : (tensor<4x8xf32>, tensor<f32>) -> tensor<4xf32>
+  return %0, %1 : tensor<4xf32>, tensor<4xf32>
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// nv_tensor_ir.reduce_ud. Reductions whose combiner is not one of the four
+// built-in modes carry an explicit identity and a region holding the combiner.
+//
+// The region operates on *signless* scalars even when the reduction itself is
+// over a signed or unsigned element type: the op verifier drops the signedness
+// of the identity before comparing it against the block argument types. The
+// signedness is therefore carried by the identity attribute and the operand
+// type only, and the arith operation inside the region is the one that picks
+// the signed or unsigned interpretation.
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: nv_tensor_ir.graph @reduce_and
+// CHECK-SAME: (%[[IN:.*]]: tensor<4x8xsi32>) -> tensor<4xsi32>
+// CHECK-NEXT: %[[REDUCE:.*]] = reduce_ud(%[[IN]]) <dimensions = [1], identity = [-1 : si32]> (%[[PREV:.*]]: i32, %[[CURR:.*]]: i32) {
+// CHECK-NEXT:   %[[AND:.*]] = arith.andi %[[PREV]], %[[CURR]] : i32
+// CHECK-NEXT:   nv_tensor_ir.yield %[[AND]] : i32
+// CHECK-NEXT: } : tensor<4x8xsi32> -> tensor<4x1xsi32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[REDUCE]] : tensor<4x1xsi32> -> tensor<4xsi32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xsi32>
+func.func @reduce_and(%arg0: tensor<4x8xi32>) -> tensor<4xi32> {
+  %init = stablehlo.constant dense<-1> : tensor<i32>
+  %0 = stablehlo.reduce(%arg0 init: %init) across dimensions = [1] : (tensor<4x8xi32>, tensor<i32>) -> tensor<4xi32>
+   reducer(%lhs: tensor<i32>, %rhs: tensor<i32>) {
+    %1 = stablehlo.and %lhs, %rhs : tensor<i32>
+    stablehlo.return %1 : tensor<i32>
+  }
+  return %0 : tensor<4xi32>
+}
+
+// -----
+
+// The identity keeps the unsigned element type, the region arguments do not.
+// CHECK-LABEL: nv_tensor_ir.graph @reduce_or_unsigned
+// CHECK-SAME: (%[[IN:.*]]: tensor<4x8xui32>) -> tensor<4xui32>
+// CHECK-NEXT: %[[REDUCE:.*]] = reduce_ud(%[[IN]]) <dimensions = [1], identity = [0 : ui32]> (%[[PREV:.*]]: i32, %[[CURR:.*]]: i32) {
+// CHECK-NEXT:   %[[OR:.*]] = arith.ori %[[PREV]], %[[CURR]] : i32
+// CHECK-NEXT:   nv_tensor_ir.yield %[[OR]] : i32
+// CHECK-NEXT: } : tensor<4x8xui32> -> tensor<4x1xui32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[REDUCE]] : tensor<4x1xui32> -> tensor<4xui32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xui32>
+func.func @reduce_or_unsigned(%arg0: tensor<4x8xui32>) -> tensor<4xui32> {
+  %init = stablehlo.constant dense<0> : tensor<ui32>
+  %0 = stablehlo.reduce(%arg0 init: %init) across dimensions = [1] : (tensor<4x8xui32>, tensor<ui32>) -> tensor<4xui32>
+   reducer(%lhs: tensor<ui32>, %rhs: tensor<ui32>) {
+    %1 = stablehlo.or %lhs, %rhs : tensor<ui32>
+    stablehlo.return %1 : tensor<ui32>
+  }
+  return %0 : tensor<4xui32>
+}
+
+// -----
+
+// A multi-operation body is translated operation by operation, including the
+// i1 intermediate produced by the comparison.
+// CHECK-LABEL: nv_tensor_ir.graph @reduce_compare_select
+// CHECK-SAME: (%[[IN:.*]]: tensor<4x8xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[REDUCE:.*]] = reduce_ud(%[[IN]]) <dimensions = [1], identity = [0.000000e+00 : f32]> (%[[PREV:.*]]: f32, %[[CURR:.*]]: f32) {
+// CHECK-NEXT:   %[[CMP:.*]] = arith.cmpf olt, %[[PREV]], %[[CURR]] : f32
+// CHECK-NEXT:   %[[SEL:.*]] = arith.select %[[CMP]], %[[PREV]], %[[CURR]] : f32
+// CHECK-NEXT:   nv_tensor_ir.yield %[[SEL]] : f32
+// CHECK-NEXT: } : tensor<4x8xf32> -> tensor<4x1xf32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[REDUCE]] : tensor<4x1xf32> -> tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @reduce_compare_select(%arg0: tensor<4x8xf32>) -> tensor<4xf32> {
+  %init = stablehlo.constant dense<0.000000e+00> : tensor<f32>
+  %0 = stablehlo.reduce(%arg0 init: %init) across dimensions = [1] : (tensor<4x8xf32>, tensor<f32>) -> tensor<4xf32>
+   reducer(%lhs: tensor<f32>, %rhs: tensor<f32>) {
+    %1 = stablehlo.compare LT, %lhs, %rhs : (tensor<f32>, tensor<f32>) -> tensor<i1>
+    %2 = stablehlo.select %1, %lhs, %rhs : tensor<i1>, tensor<f32>
+    stablehlo.return %2 : tensor<f32>
+  }
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// Constants inside the body become scalar arith.constant, and clamp is
+// decomposed the same way as its elementwise counterpart.
+// CHECK-LABEL: nv_tensor_ir.graph @reduce_clamp
+// CHECK-SAME: (%[[IN:.*]]: tensor<4x8xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[REDUCE:.*]] = reduce_ud(%[[IN]]) <dimensions = [1], identity = [0.000000e+00 : f32]> (%[[PREV:.*]]: f32, %[[CURR:.*]]: f32) {
+// CHECK-NEXT:   %[[LO:.*]] = arith.constant -1.000000e+06 : f32
+// CHECK-NEXT:   %[[HI:.*]] = arith.constant 1.000000e+06 : f32
+// CHECK-NEXT:   %[[SUM:.*]] = arith.addf %[[PREV]], %[[CURR]] : f32
+// CHECK-NEXT:   %[[LOWER:.*]] = arith.maximumf %[[LO]], %[[SUM]] : f32
+// CHECK-NEXT:   %[[CLAMPED:.*]] = arith.minimumf %[[HI]], %[[LOWER]] : f32
+// CHECK-NEXT:   nv_tensor_ir.yield %[[CLAMPED]] : f32
+// CHECK-NEXT: } : tensor<4x8xf32> -> tensor<4x1xf32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[REDUCE]] : tensor<4x1xf32> -> tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @reduce_clamp(%arg0: tensor<4x8xf32>) -> tensor<4xf32> {
+  %init = stablehlo.constant dense<0.000000e+00> : tensor<f32>
+  %0 = stablehlo.reduce(%arg0 init: %init) across dimensions = [1] : (tensor<4x8xf32>, tensor<f32>) -> tensor<4xf32>
+   reducer(%lhs: tensor<f32>, %rhs: tensor<f32>) {
+    %lo = stablehlo.constant dense<-1.000000e+06> : tensor<f32>
+    %hi = stablehlo.constant dense<1.000000e+06> : tensor<f32>
+    %1 = stablehlo.add %lhs, %rhs : tensor<f32>
+    %2 = stablehlo.clamp %lo, %1, %hi : tensor<f32>
+    stablehlo.return %2 : tensor<f32>
+  }
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// nv_tensor_ir.reduce has no init operand and implicitly uses the identity of
+// its reduction mode, so an add reduction seeded with 1.0 cannot use it and
+// falls back to reduce_ud, whose identity is explicit.
+// CHECK-LABEL: nv_tensor_ir.graph @reduce_non_identity_init
+// CHECK-SAME: (%[[IN:.*]]: tensor<4x8xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[REDUCE:.*]] = reduce_ud(%[[IN]]) <dimensions = [1], identity = [1.000000e+00 : f32]> (%[[PREV:.*]]: f32, %[[CURR:.*]]: f32) {
+// CHECK-NEXT:   %[[SUM:.*]] = arith.addf %[[PREV]], %[[CURR]] : f32
+// CHECK-NEXT:   nv_tensor_ir.yield %[[SUM]] : f32
+// CHECK-NEXT: } : tensor<4x8xf32> -> tensor<4x1xf32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[REDUCE]] : tensor<4x1xf32> -> tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @reduce_non_identity_init(%arg0: tensor<4x8xf32>) -> tensor<4xf32> {
+  %init = stablehlo.constant dense<1.000000e+00> : tensor<f32>
+  %0 = stablehlo.reduce(%arg0 init: %init) applies stablehlo.add across dimensions = [1] : (tensor<4x8xf32>, tensor<f32>) -> tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// add(%prev, %prev) ignores the incoming element, so it is not the built-in
+// add reduction even though the combiner and the identity both look like one.
+// CHECK-LABEL: nv_tensor_ir.graph @reduce_repeated_argument
+// CHECK-SAME: (%[[IN:.*]]: tensor<4x8xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[REDUCE:.*]] = reduce_ud(%[[IN]]) <dimensions = [1], identity = [0.000000e+00 : f32]> (%[[PREV:.*]]: f32, %[[CURR:.*]]: f32) {
+// CHECK-NEXT:   %[[SUM:.*]] = arith.addf %[[PREV]], %[[PREV]] : f32
+// CHECK-NEXT:   nv_tensor_ir.yield %[[SUM]] : f32
+// CHECK-NEXT: } : tensor<4x8xf32> -> tensor<4x1xf32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[REDUCE]] : tensor<4x1xf32> -> tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @reduce_repeated_argument(%arg0: tensor<4x8xf32>) -> tensor<4xf32> {
+  %init = stablehlo.constant dense<0.000000e+00> : tensor<f32>
+  %0 = stablehlo.reduce(%arg0 init: %init) across dimensions = [1] : (tensor<4x8xf32>, tensor<f32>) -> tensor<4xf32>
+   reducer(%lhs: tensor<f32>, %rhs: tensor<f32>) {
+    %1 = stablehlo.add %lhs, %lhs : tensor<f32>
+    stablehlo.return %1 : tensor<f32>
+  }
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// A flipped built-in combiner still uses the fast path: add and multiply are
+// commutative, so the operand order carries no information. This is a
+// deliberate divergence from the string-building converter, which required
+// add(%prev, %curr) exactly and fell back to reduce_ud otherwise.
+// CHECK-LABEL: nv_tensor_ir.graph @reduce_add_flipped
+// CHECK-SAME: (%[[IN:.*]]: tensor<4x8xf32>) -> tensor<4xf32>
+// CHECK-NEXT: %[[REDUCE:.*]] = reduce(%[[IN]]) <dimensions = [1], reduction_mode = <add>> : tensor<4x8xf32> -> tensor<4x1xf32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[REDUCE]] : tensor<4x1xf32> -> tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @reduce_add_flipped(%arg0: tensor<4x8xf32>) -> tensor<4xf32> {
+  %init = stablehlo.constant dense<0.000000e+00> : tensor<f32>
+  %0 = stablehlo.reduce(%arg0 init: %init) across dimensions = [1] : (tensor<4x8xf32>, tensor<f32>) -> tensor<4xf32>
+   reducer(%lhs: tensor<f32>, %rhs: tensor<f32>) {
+    %1 = stablehlo.add %rhs, %lhs : tensor<f32>
+    stablehlo.return %1 : tensor<f32>
+  }
+  return %0 : tensor<4xf32>
+}
+
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_structural.mlir b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_structural.mlir
new file mode 100644
index 00000000000..02b3ffcac9c
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_structural.mlir
@@ -0,0 +1,223 @@
+// RUN: xla-opt %s --legalize-stablehlo-to-tensor-ir --split-input-file | FileCheck %s
+
+// CHECK-LABEL: nv_tensor_ir.graph @reshape
+// CHECK-SAME: (%[[IN:.*]]: tensor<2x3xf32>) -> tensor<6x1xf32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[IN]] : tensor<2x3xf32> -> tensor<6x1xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<6x1xf32>
+func.func @reshape(%arg0: tensor<2x3xf32>) -> tensor<6x1xf32> {
+  %0 = stablehlo.reshape %arg0 : (tensor<2x3xf32>) -> tensor<6x1xf32>
+  return %0 : tensor<6x1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @reshape_i32
+// CHECK-SAME: (%[[IN:.*]]: tensor<2x3xsi32>) -> tensor<3x2xsi32>
+// CHECK-NEXT: %[[RES:.*]] = reshape %[[IN]] : tensor<2x3xsi32> -> tensor<3x2xsi32>
+// CHECK-NEXT: results %[[RES]] : tensor<3x2xsi32>
+func.func @reshape_i32(%arg0: tensor<2x3xi32>) -> tensor<3x2xi32> {
+  %0 = stablehlo.reshape %arg0 : (tensor<2x3xi32>) -> tensor<3x2xi32>
+  return %0 : tensor<3x2xi32>
+}
+
+// -----
+
+// The permutation is forwarded as is. With the opposite convention the result
+// shape would be 3x4x2 and the nv_tensor_ir.transpose verifier would reject it.
+// CHECK-LABEL: nv_tensor_ir.graph @transpose
+// CHECK-SAME: (%[[IN:.*]]: tensor<2x3x4xf32>) -> tensor<4x2x3xf32>
+// CHECK-NEXT: %[[RES:.*]] = transpose %[[IN]] permutation = [2, 0, 1] : tensor<2x3x4xf32> -> tensor<4x2x3xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4x2x3xf32>
+func.func @transpose(%arg0: tensor<2x3x4xf32>) -> tensor<4x2x3xf32> {
+  %0 = stablehlo.transpose %arg0, dims = [2, 0, 1] : (tensor<2x3x4xf32>) -> tensor<4x2x3xf32>
+  return %0 : tensor<4x2x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @slice
+// CHECK-SAME: (%[[IN:.*]]: tensor<8x8xf32>) -> tensor<4x2xf32>
+// CHECK-NEXT: %[[RES:.*]] = slice %[[IN]] starts = [0, 1] limits = [8, 7] strides = [2, 3] : tensor<8x8xf32> -> tensor<4x2xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4x2xf32>
+func.func @slice(%arg0: tensor<8x8xf32>) -> tensor<4x2xf32> {
+  %0 = stablehlo.slice %arg0 [0:8:2, 1:7:3] : (tensor<8x8xf32>) -> tensor<4x2xf32>
+  return %0 : tensor<4x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @concatenate
+// CHECK-SAME: (%[[LHS:.*]]: tensor<2x2xf32>, %[[RHS:.*]]: tensor<2x3xf32>) -> tensor<2x5xf32>
+// CHECK-NEXT: %[[RES:.*]] = concatenate %[[LHS]], %[[RHS]] dimension = 1 : (tensor<2x2xf32>, tensor<2x3xf32>) -> tensor<2x5xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<2x5xf32>
+func.func @concatenate(%arg0: tensor<2x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<2x5xf32> {
+  %0 = stablehlo.concatenate %arg0, %arg1, dim = 1 : (tensor<2x2xf32>, tensor<2x3xf32>) -> tensor<2x5xf32>
+  return %0 : tensor<2x5xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @iota
+// CHECK-SAME: () -> tensor<2x3xf32>
+// CHECK-NEXT: %[[RES:.*]] = iota dimension = 1 : tensor<2x3xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<2x3xf32>
+func.func @iota() -> tensor<2x3xf32> {
+  %0 = stablehlo.iota dim = 1 : tensor<2x3xf32>
+  return %0 : tensor<2x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @iota_i32
+// CHECK-SAME: () -> tensor<4x5xsi32>
+// CHECK-NEXT: %[[RES:.*]] = iota dimension = 0 : tensor<4x5xsi32>
+// CHECK-NEXT: results %[[RES]] : tensor<4x5xsi32>
+func.func @iota_i32() -> tensor<4x5xi32> {
+  %0 = stablehlo.iota dim = 0 : tensor<4x5xi32>
+  return %0 : tensor<4x5xi32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @constant_f32
+// CHECK-SAME: () -> tensor<4xf32>
+// CHECK-NEXT: %[[RES:.*]] = constant dense<1.000000e+00> : tensor<4xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xf32>
+func.func @constant_f32() -> tensor<4xf32> {
+  %0 = stablehlo.constant dense<1.000000e+00> : tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// The value attribute is rebuilt over the converted element type, because
+// nv_tensor_ir.constant infers its result type from the attribute.
+// CHECK-LABEL: nv_tensor_ir.graph @constant_i32
+// CHECK-SAME: () -> tensor<4xsi32>
+// CHECK-NEXT: %[[RES:.*]] = constant dense<[1, 2, 3, 4]> : tensor<4xsi32>
+// CHECK-NEXT: results %[[RES]] : tensor<4xsi32>
+func.func @constant_i32() -> tensor<4xi32> {
+  %0 = stablehlo.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
+  return %0 : tensor<4xi32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @constant_i32_splat
+// CHECK-SAME: () -> tensor<2x2xsi32>
+// CHECK-NEXT: %[[RES:.*]] = constant dense<7> : tensor<2x2xsi32>
+// CHECK-NEXT: results %[[RES]] : tensor<2x2xsi32>
+func.func @constant_i32_splat() -> tensor<2x2xi32> {
+  %0 = stablehlo.constant dense<7> : tensor<2x2xi32>
+  return %0 : tensor<2x2xi32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @constant_i64
+// CHECK-SAME: () -> tensor<2xsi64>
+// CHECK-NEXT: %[[RES:.*]] = constant dense<[-1, 9223372036854775807]> : tensor<2xsi64>
+// CHECK-NEXT: results %[[RES]] : tensor<2xsi64>
+func.func @constant_i64() -> tensor<2xi64> {
+  %0 = stablehlo.constant dense<[-1, 9223372036854775807]> : tensor<2xi64>
+  return %0 : tensor<2xi64>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @constant_ui32
+// CHECK-SAME: () -> tensor<2xui32>
+// CHECK-NEXT: %[[RES:.*]] = constant dense<[1, 4294967295]> : tensor<2xui32>
+// CHECK-NEXT: results %[[RES]] : tensor<2xui32>
+func.func @constant_ui32() -> tensor<2xui32> {
+  %0 = stablehlo.constant dense<[1, 4294967295]> : tensor<2xui32>
+  return %0 : tensor<2xui32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @constant_i1
+// CHECK-SAME: () -> tensor<2xi1>
+// CHECK-NEXT: %[[RES:.*]] = constant dense<[true, false]> : tensor<2xi1>
+// CHECK-NEXT: results %[[RES]] : tensor<2xi1>
+func.func @constant_i1() -> tensor<2xi1> {
+  %0 = stablehlo.constant dense<[true, false]> : tensor<2xi1>
+  return %0 : tensor<2xi1>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @broadcast_leading_dim
+// CHECK-SAME: (%[[IN:.*]]: tensor<3xf32>) -> tensor<2x3xf32>
+// CHECK-NEXT: %[[EXPANDED:.*]] = reshape %[[IN]] : tensor<3xf32> -> tensor<1x3xf32>
+// CHECK-NEXT: %[[RES:.*]] = broadcast %[[EXPANDED]] : tensor<1x3xf32> -> tensor<2x3xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<2x3xf32>
+func.func @broadcast_leading_dim(%arg0: tensor<3xf32>) -> tensor<2x3xf32> {
+  %0 = stablehlo.broadcast_in_dim %arg0, dims = [1] : (tensor<3xf32>) -> tensor<2x3xf32>
+  return %0 : tensor<2x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: nv_tensor_ir.graph @broadcast_trailing_dim
+// CHECK-SAME: (%[[IN:.*]]: tensor<3xf32>) -> tensor<3x2xf32>
+// CHECK-NEXT: %[[EXPANDED:.*]] = reshape %[[IN]] : tensor<3xf32> -> tensor<3x1xf32>
+// CHECK-NEXT: %[[RES:.*]] = broadcast %[[EXPANDED]] : tensor<3x1xf32> -> tensor<3x2xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<3x2xf32>
+func.func @broadcast_trailing_dim(%arg0: tensor<3xf32>) -> tensor<3x2xf32> {
+  %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<3xf32>) -> tensor<3x2xf32>
+  return %0 : tensor<3x2xf32>
+}
+
+// -----
+
+// The input already has the output rank, so no reshape is needed.
+// CHECK-LABEL: nv_tensor_ir.graph @broadcast_unit_dim
+// CHECK-SAME: (%[[IN:.*]]: tensor<16x1x64xf32>) -> tensor<16x32x64xf32>
+// CHECK-NEXT: %[[RES:.*]] = broadcast %[[IN]] : tensor<16x1x64xf32> -> tensor<16x32x64xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<16x32x64xf32>
+func.func @broadcast_unit_dim(%arg0: tensor<16x1x64xf32>) -> tensor<16x32x64xf32> {
+  %0 = stablehlo.broadcast_in_dim %arg0, dims = [0, 1, 2] : (tensor<16x1x64xf32>) -> tensor<16x32x64xf32>
+  return %0 : tensor<16x32x64xf32>
+}
+
+// -----
+
+// A no-op broadcast is rejected by the nv_tensor_ir.broadcast verifier, so no
+// operation is emitted at all.
+// CHECK-LABEL: nv_tensor_ir.graph @broadcast_noop
+// CHECK-SAME: (%[[IN:.*]]: tensor<4xf32>) -> tensor<4xf32>
+// CHECK-NEXT: results %[[IN]] : tensor<4xf32>
+func.func @broadcast_noop(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+  %0 = stablehlo.broadcast_in_dim %arg0, dims = [0] : (tensor<4xf32>) -> tensor<4xf32>
+  return %0 : tensor<4xf32>
+}
+
+// -----
+
+// A rank-0 argument is declared as `[1]` on the boundary, so the reshape that
+// expands it for the broadcast starts from rank 1. See `PromoteRank0Tensors`.
+// CHECK-LABEL: nv_tensor_ir.graph @broadcast_scalar
+// CHECK-SAME: (%[[IN:.*]]: tensor<1xf32>) -> tensor<2x3xf32>
+// CHECK-NEXT: %[[EXPANDED:.*]] = reshape %[[IN]] : tensor<1xf32> -> tensor<1x1xf32>
+// CHECK-NEXT: %[[RES:.*]] = broadcast %[[EXPANDED]] : tensor<1x1xf32> -> tensor<2x3xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<2x3xf32>
+func.func @broadcast_scalar(%arg0: tensor<f32>) -> tensor<2x3xf32> {
+  %0 = stablehlo.broadcast_in_dim %arg0, dims = [] : (tensor<f32>) -> tensor<2x3xf32>
+  return %0 : tensor<2x3xf32>
+}
+
+// -----
+
+// The broadcast dimensions are not ascending, so the input dimensions have to
+// be sorted with a transpose before the reshape.
+// CHECK-LABEL: nv_tensor_ir.graph @broadcast_permuted
+// CHECK-SAME: (%[[IN:.*]]: tensor<2x3xf32>) -> tensor<4x3x2xf32>
+// CHECK-NEXT: %[[SORTED:.*]] = transpose %[[IN]] permutation = [1, 0] : tensor<2x3xf32> -> tensor<3x2xf32>
+// CHECK-NEXT: %[[EXPANDED:.*]] = reshape %[[SORTED]] : tensor<3x2xf32> -> tensor<1x3x2xf32>
+// CHECK-NEXT: %[[RES:.*]] = broadcast %[[EXPANDED]] : tensor<1x3x2xf32> -> tensor<4x3x2xf32>
+// CHECK-NEXT: results %[[RES]] : tensor<4x3x2xf32>
+func.func @broadcast_permuted(%arg0: tensor<2x3xf32>) -> tensor<4x3x2xf32> {
+  %0 = stablehlo.broadcast_in_dim %arg0, dims = [2, 1] : (tensor<2x3xf32>) -> tensor<4x3x2xf32>
+  return %0 : tensor<4x3x2xf32>
+}
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/type_conversion.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/type_conversion.cc
new file mode 100644
index 00000000000..3bd94b09081
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/type_conversion.cc
@@ -0,0 +1,65 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/gpu/codegen/tensor_ir/transforms/type_conversion.h"
+
+#include <optional>
+
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Types.h"
+#include "mlir/Support/LogicalResult.h"
+
+namespace mlir::nv_tensor_ir::xla {
+
+TensorIrTypeConverter::TensorIrTypeConverter() {
+  // Pass-through fallback for types that do not need conversion (e.g. floats).
+  addConversion([](mlir::Type type) { return type; });
+
+  // Map signless IntegerType of width 8/16/32/64 to signed IntegerType.
+  // i1 remains unchanged. Unsigned integers remain unchanged.
+  addConversion([](mlir::IntegerType type) -> mlir::Type {
+    if (type.isSignless() && (type.getWidth() == 8 || type.getWidth() == 16 ||
+                              type.getWidth() == 32 || type.getWidth() == 64)) {
+      return mlir::IntegerType::get(type.getContext(), type.getWidth(),
+                                    mlir::IntegerType::Signed);
+    }
+    return type;
+  });
+
+  // Map RankedTensorType recursively: same shape, converted element type,
+  // encoding dropped (nv_tensor_ir rejects non-null encodings).
+  addConversion(
+      [this](mlir::RankedTensorType type) -> std::optional<mlir::Type> {
+        mlir::Type elem_type = convertType(type.getElementType());
+        if (!elem_type) {
+          return std::nullopt;
+        }
+        return mlir::RankedTensorType::get(type.getShape(), elem_type);
+      });
+
+  // Convert FunctionType: convert all input and result types.
+  addConversion([this](mlir::FunctionType type) -> std::optional<mlir::Type> {
+    llvm::SmallVector<mlir::Type> inputs;
+    llvm::SmallVector<mlir::Type> results;
+    if (failed(convertTypes(type.getInputs(), inputs)) ||
+        failed(convertTypes(type.getResults(), results))) {
+      return std::nullopt;
+    }
+    return mlir::FunctionType::get(type.getContext(), inputs, results);
+  });
+}
+
+}  // namespace mlir::nv_tensor_ir::xla
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/type_conversion.h b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/type_conversion.h
new file mode 100644
index 00000000000..6fd4b251cdd
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/type_conversion.h
@@ -0,0 +1,30 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_TRANSFORMS_TYPE_CONVERSION_H_
+#define XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_TRANSFORMS_TYPE_CONVERSION_H_
+
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir::nv_tensor_ir::xla {
+
+class TensorIrTypeConverter : public mlir::TypeConverter {
+ public:
+  TensorIrTypeConverter();
+};
+
+}  // namespace mlir::nv_tensor_ir::xla
+
+#endif  // XLA_BACKENDS_GPU_CODEGEN_TENSOR_IR_TRANSFORMS_TYPE_CONVERSION_H_
diff --git a/third_party/xla/xla/backends/gpu/codegen/tools/BUILD b/third_party/xla/xla/backends/gpu/codegen/tools/BUILD
index 51e8db082d5..259f996c02c 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tools/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/tools/BUILD
@@ -96,10 +96,12 @@ xla_cc_binary(
     ],
     visibility = [":codegen_tests"],
     deps = [
-        "//xla/backends/gpu/codegen/tensor_ir:conversion",
+        "//xla/backends/gpu/codegen/tensor_ir:hlo_to_tensor_ir",
         "//xla/backends/gpu/codegen/tensor_ir:support",
         "//xla/hlo/ir:hlo",
+        "//xla/service/llvm_ir:llvm_util",
         "//xla/tools:hlo_module_loader",
+        "@com_google_absl//absl/status:statusor",
         "@cuda_tile//:CudaTileDialect",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
diff --git a/third_party/xla/xla/backends/gpu/codegen/tools/fusion_to_tensorir.cc b/third_party/xla/xla/backends/gpu/codegen/tools/fusion_to_tensorir.cc
index 50249e81c5b..5896e0e19ea 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tools/fusion_to_tensorir.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/tools/fusion_to_tensorir.cc
@@ -13,18 +13,19 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include <utility>
+#include <cstdint>
 
 #include "cuda_tile/Dialect/CudaTile/IR/Dialect.h"
-#include "mlir/Tools/mlir-translate/MlirTranslateMain.h"
-#include "mlir/Tools/mlir-translate/Translation.h"
 #include "tensor_ir/Compiler/CudaTile/Pipelines.h"
 #include "tensor_ir/Conversion/TensorToCudaTile/Options.h"
 #include "tensor_ir/Dialect/TensorIR.h"
+#include "absl/status/statusor.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/LogicalResult.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/DialectRegistry.h"
 #include "mlir/IR/Location.h"
@@ -32,11 +33,14 @@ limitations under the License.
 #include "mlir/IR/OwningOpRef.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/LLVM.h"
-#include "xla/backends/gpu/codegen/tensor_ir/conversion.h"
+#include "mlir/Tools/mlir-translate/MlirTranslateMain.h"
+#include "mlir/Tools/mlir-translate/Translation.h"
+#include "xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir.h"
 #include "xla/backends/gpu/codegen/tensor_ir/support.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/service/llvm_ir/llvm_util.h"
 #include "xla/tools/hlo_module_loader.h"
 
 namespace xla::gpu::tensor_ir {
@@ -47,6 +51,20 @@ llvm::cl::opt<bool> compile_flag("compile",
                                  llvm::cl::desc("Compile to CudaTile dialect."),
                                  llvm::cl::init(false));
 
+// NOLINTNEXTLINE
+llvm::cl::opt<int64_t> alignment_flag(
+    "alignment",
+    llvm::cl::desc(
+        "Alignment (in bytes) to use for the graph's arguments and result."),
+    llvm::cl::init(1));
+
+// NOLINTNEXTLINE
+llvm::cl::list<int32_t> tile_shape_flag(
+    "tile_shape",
+    llvm::cl::desc("Tile shape to use for the CudaTile conversion pipeline "
+                   "(only used with --compile)."),
+    llvm::cl::CommaSeparated);
+
 mlir::OwningOpRef<mlir::ModuleOp> HloToTensorIRTranslate(
     llvm::StringRef input, mlir::MLIRContext* context) {
   context->loadAllAvailableDialects();
@@ -70,17 +88,39 @@ mlir::OwningOpRef<mlir::ModuleOp> HloToTensorIRTranslate(
     return nullptr;
   }
 
-  auto module_or = ConvertFusionComputation(*comp, context);
-  if (!module_or.ok()) {
+  mlir::OwningOpRef<mlir::ModuleOp> module =
+      llvm_ir::CreateMlirModuleOp(mlir::UnknownLoc::get(context));
+  absl::StatusOr<mlir::nv_tensor_ir::GraphOp> graph_op_or =
+      ImportAndLegalizeComputation(*comp, *module);
+  if (!graph_op_or.ok()) {
     mlir::emitError(mlir::UnknownLoc::get(context))
-        << module_or.status().message();
+        << graph_op_or.status().message();
     return nullptr;
   }
-  mlir::OwningOpRef<mlir::ModuleOp> module = *std::move(module_or);
+  mlir::nv_tensor_ir::GraphOp graph_op = *graph_op_or;
+
+  // Annotate the graph's arguments and result with the alignment attribute.
+  if (alignment_flag > 1) {
+    mlir::StringAttr alignment_attr_name = mlir::StringAttr::get(
+        context, mlir::nv_tensor_ir::TensorIRDialect::getAlignmentAttrName());
+    mlir::IntegerAttr alignment_attr = mlir::IntegerAttr::get(
+        mlir::IntegerType::get(context, 64), alignment_flag.getValue());
+    for (int64_t i = 0; i < graph_op.getNumArguments(); ++i) {
+      graph_op.setArgAttr(i, alignment_attr_name, alignment_attr);
+    }
+    for (int64_t i = 0; i < graph_op.getNumResults(); ++i) {
+      graph_op.setResultAttr(i, alignment_attr_name, alignment_attr);
+    }
+  }
 
   if (compile_flag) {
     mlir::PassManager pass_manager(context);
     mlir::nv_tensor_ir::TensorToCudaTilePipelineOptions options;
+    if (!tile_shape_flag.empty()) {
+      options.tileSize.assign(tile_shape_flag.begin(), tile_shape_flag.end());
+      options.reductionTileSize = options.tileSize.back();
+      options.tileSize.pop_back();
+    }
     mlir::nv_tensor_ir::buildTensorToCudaTileConversionPipeline(pass_manager,
                                                                 options);
     if (llvm::failed(pass_manager.run(*module))) {
diff --git a/third_party/xla/xla/backends/gpu/tests/BUILD b/third_party/xla/xla/backends/gpu/tests/BUILD
index a951c40b4e4..f4409759688 100644
--- a/third_party/xla/xla/backends/gpu/tests/BUILD
+++ b/third_party/xla/xla/backends/gpu/tests/BUILD
@@ -802,12 +802,14 @@ cc_binary(
     srcs = ["xla-opt.cc"],
     deps = [
         "//xla/backends/gpu/codegen/emitters/transforms:passes",
+        "//xla/backends/gpu/codegen/tensor_ir/transforms:passes",
         "//xla/backends/gpu/codegen/triton:compilation_pipeline",
         "//xla/backends/gpu/codegen/triton/ir:triton_xla",
         "//xla/backends/gpu/codegen/triton/transforms:passes",
         "//xla/codegen/emitters/ir:xla",
         "//xla/codegen/emitters/transforms:passes",
         "//xla/codegen/xtile/ir:xtile",
+        "//xla/mlir_hlo",
         "//xla/stream_executor:device_description",
         "//xla/stream_executor/cuda:cuda_compute_capability",
         "//xla/stream_executor/rocm:rocm_compute_capability",
@@ -825,6 +827,7 @@ cc_binary(
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TensorDialect",
         "@stablehlo//:stablehlo_ops",
+        "@tensor_ir//:NVTensorIRDialect",
         "@triton//:AllPassesAndDialects",
         "@triton//third_party/amd:TestAMDAnalysis",  # buildcleaner: keep
     ],
diff --git a/third_party/xla/xla/backends/gpu/tests/xla-opt.cc b/third_party/xla/xla/backends/gpu/tests/xla-opt.cc
index dafec9a8b6a..a1fdb9026ae 100644
--- a/third_party/xla/xla/backends/gpu/tests/xla-opt.cc
+++ b/third_party/xla/xla/backends/gpu/tests/xla-opt.cc
@@ -16,12 +16,13 @@ limitations under the License.
 #include <string>
 
 #include "bin/RegisterTritonDialects.h"
+#include "tensor_ir/Dialect/TensorIR.h"
 #include "llvm/Support/CommandLine.h"
 #include "mlir/Dialect/Func/Extensions/InlinerExtension.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Func/IR/FuncDialect.h"
 #include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h"
-#include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/MemRef/IR/MemRefDialect.h"
+#include "mlir/Dialect/Tensor/IR/TensorDialect.h"
 #include "mlir/InitAllExtensions.h"
 #include "mlir/Pass/PassOptions.h"
 #include "mlir/Pass/PassRegistry.h"
@@ -31,12 +32,14 @@ limitations under the License.
 #include "mlir/Tools/mlir-opt/MlirOptMain.h"
 #include "stablehlo/dialect/StablehloOps.h"
 #include "xla/backends/gpu/codegen/emitters/transforms/passes.h"
+#include "xla/backends/gpu/codegen/tensor_ir/transforms/passes.h"
 #include "xla/backends/gpu/codegen/triton/compilation_pipeline.h"
 #include "xla/backends/gpu/codegen/triton/ir/triton_xla_ops.h"
 #include "xla/backends/gpu/codegen/triton/transforms/passes.h"
 #include "xla/codegen/emitters/ir/xla_dialect.h"
 #include "xla/codegen/emitters/transforms/passes.h"
 #include "xla/codegen/xtile/ir/xtile_dialect.h"
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/rocm/rocm_compute_capability.h"
@@ -94,8 +97,10 @@ int main(int argc, char** argv) {
   registry.insert<mlir::func::FuncDialect, mlir::tensor::TensorDialect,
                   mlir::triton::xla::XlaTritonDialect, xla::XlaDialect,
                   xla::xtile::XTileDialect, mlir::stablehlo::StablehloDialect,
-                  mlir::memref::MemRefDialect>();
+                  mlir::memref::MemRefDialect, mlir::mhlo::MhloDialect,
+                  mlir::nv_tensor_ir::TensorIRDialect>();
   mlir::triton::xla::registerTritonXlaTransformsPasses();
+  mlir::nv_tensor_ir::xla::registerTensorIrTransformsPasses();
   xla::emitters::registerTransformsPasses();
   xla::gpu::registerGpuFusionTransformsPasses();
 
diff --git a/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc
index 0d79ddb6e16..86f941bb4b6 100644
--- a/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc
@@ -253,6 +253,16 @@ absl::StatusOr<bool> ProcessFusionInstruction(
   ABSL_ASSIGN_OR_RETURN(auto backend_config,
                    fusion_instruction->backend_config<GpuBackendConfig>());
 
+  if (backend_config.has_fusion_backend_config() &&
+      backend_config.fusion_backend_config().kind() == kTensorIrFusionKind) {
+    // The fusion was explicitly assigned to the TensorIR emitter (by the
+    // autotuner or by hand). Rewriting it into a Triton block-level fusion
+    // would silently discard that assignment.
+    VLOG(2) << "Not rewriting fusion " << fusion_instruction->name()
+            << " because it is assigned to the TensorIR emitter.";
+    return false;
+  }
+
   if (backend_config.has_fusion_backend_config() &&
       backend_config.fusion_backend_config().has_block_level_fusion_config()) {
     // Fusion is already block-level! Skip.
diff --git a/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc b/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc
index 2b95620a76d..d0ec57af11c 100644
--- a/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc
@@ -1526,6 +1526,7 @@ HloInstruction::FusionKind PriorityFusion::ChooseKind(
     case HloFusionAnalysis::EmitterFusionKind::kCustomFusion:
     case HloFusionAnalysis::EmitterFusionKind::kCuDnn:
     case HloFusionAnalysis::EmitterFusionKind::kSort:
+    case HloFusionAnalysis::EmitterFusionKind::kTensorIr:
       return HloInstruction::FusionKind::kCustom;
     case HloFusionAnalysis::EmitterFusionKind::kConcatenate:
     case HloFusionAnalysis::EmitterFusionKind::kReduction:
diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index 6f063f4cf77..6d70cc36c80 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -609,6 +609,7 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() {
   opts.set_xla_gpu_gxl_scratch_size_bytes(64 * 1024 * 1024);
   opts.set_xla_gpu_enable_persistent_symmetric_memory(false);
   opts.set_xla_gpu_experimental_enable_raft_for_stable_topk(false);
+  opts.set_xla_gpu_experimental_enable_tensor_ir(false);
   opts.set_xla_gpu_async_copy_min_bytes(-1);
 
   // Disable float checks.
@@ -3737,6 +3738,13 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
           &DebugOptions::set_xla_gpu_experimental_enable_raft_for_stable_topk),
       debug_options->xla_gpu_experimental_enable_raft_for_stable_topk(),
       "If true, enables RAFT for stable TopK."));
+  flag_list->push_back(tsl::Flag(
+      "xla_gpu_experimental_enable_tensor_ir",
+      bool_setter_for(&DebugOptions::set_xla_gpu_experimental_enable_tensor_ir),
+      debug_options->xla_gpu_experimental_enable_tensor_ir(),
+      "If true, the GPU autotuner may assign fusions to the experimental "
+      "TensorIR (CUDA Tile) emitter. Requires sm_90+ and CUDA 13.3+ "
+      "(tileiras)."));
   flag_list->push_back(tsl::Flag(
       "xla_gpu_experimental_ragged_all_to_all_use_device_kernel",
       bool_setter_for(
diff --git a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc
index 177debd715b..89822e5e713 100644
--- a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc
+++ b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc
@@ -373,18 +373,18 @@ InstructionFilterFn GetShouldAssignConfigToInstructionFn(
       !debug_options.xla_gpu_deterministic_ops() &&
       debug_options.xla_gpu_experimental_enable_fusion_autotuner();
 
-  return [do_not_autotune_cublas, do_not_autotune_cudnn,
-          enable_fusion_autotuner,
-          gpu_version](const HloInstruction& instruction) -> bool {
-    AutotuneDecision decision = ShouldAssignConfigToInstruction(
-        do_not_autotune_cublas, do_not_autotune_cudnn, enable_fusion_autotuner,
-        instruction);
-    if (!decision) {
-      VLOG(3) << "Not assigning configs to " << instruction.name() << ": "
-              << decision.Explain();
-    }
-    return decision.IsAllowed();
-  };
+  return
+      [do_not_autotune_cublas, do_not_autotune_cudnn, enable_fusion_autotuner,
+       gpu_version](const HloInstruction& instruction) -> bool {
+        AutotuneDecision decision = ShouldAssignConfigToInstruction(
+            do_not_autotune_cublas, do_not_autotune_cudnn,
+            enable_fusion_autotuner, instruction);
+        if (!decision) {
+          VLOG(3) << "Not assigning configs to " << instruction.name() << ": "
+                  << decision.Explain();
+        }
+        return decision.IsAllowed();
+      };
 }
 
 absl::StatusOr<std::vector<std::unique_ptr<CodegenBackend>>>
@@ -427,6 +427,10 @@ ConfigAssignerPass::GetEnabledBackends(
     disabled_autotune_backends.push_back(autotuner::Backend::TRITON);
   }
 
+  if (!debug_options.xla_gpu_experimental_enable_tensor_ir()) {
+    disabled_autotune_backends.push_back(autotuner::Backend::TENSOR_IR);
+  }
+
   autotune_backends.erase(
       std::remove_if(autotune_backends.begin(), autotune_backends.end(),
                      [&](autotuner::Backend backend) {
@@ -476,26 +480,25 @@ absl::StatusOr<std::unique_ptr<ConfigAssignerPass>> ConfigAssignerPass::Create(
 
   std::unique_ptr<Autotuner> autotuner = nullptr;
   if (!is_deviceless) {
-      // TODO(intel-tf): Enable buffer checking for SYCL once
-      // BufferComparatorKernel and RedzoneAllocatorKernel are registered for
-      // SYCL platform.
-      bool is_buffer_check_supported = stream_executor->GetPlatform()->id() !=
-                                       stream_executor::sycl::kSyclPlatformId;
-      std::unique_ptr<Profiler> profiler = GpuProfiler::Create(
-          stream_executor,
-          GetProfileOptions(debug_options, is_buffer_check_supported),
-          allocator);
-      Autotuner::Options autotuner_options =
-          GetAutotunerOptions(debug_options, is_buffer_check_supported);
-      autotuner_options.cache_context = AutotuneCacheContext::Create(
-          target_config->device_description, orchestrator->codegen_backends());
+    // TODO(intel-tf): Enable buffer checking for SYCL once
+    // BufferComparatorKernel and RedzoneAllocatorKernel are registered for
+    // SYCL platform.
+    bool is_buffer_check_supported = stream_executor->GetPlatform()->id() !=
+                                     stream_executor::sycl::kSyclPlatformId;
+    std::unique_ptr<Profiler> profiler = GpuProfiler::Create(
+        stream_executor,
+        GetProfileOptions(debug_options, is_buffer_check_supported), allocator);
+    Autotuner::Options autotuner_options =
+        GetAutotunerOptions(debug_options, is_buffer_check_supported);
+    autotuner_options.cache_context = AutotuneCacheContext::Create(
+        target_config->device_description, orchestrator->codegen_backends());
 
-      std::vector<std::unique_ptr<Profiler>> profilers;
-      profilers.push_back(std::move(profiler));
+    std::vector<std::unique_ptr<Profiler>> profilers;
+    profilers.push_back(std::move(profiler));
 
-      ABSL_ASSIGN_OR_RETURN(autotuner,
-                       Autotuner::Create(*orchestrator, std::move(profilers),
-                                         autotuner_options, thread_pool));
+    ABSL_ASSIGN_OR_RETURN(autotuner,
+                     Autotuner::Create(*orchestrator, std::move(profilers),
+                                       autotuner_options, thread_pool));
   }
 
   VLOG(1) << "ConfigAssigner options: " << assigner_options.ToString();
diff --git a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc
index 8a9ee807834..e63c8f02985 100644
--- a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc
+++ b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc
@@ -733,6 +733,53 @@ TEST_F(AutotunerFlagsTest, GetEnabledBackendsRespectsDeterminism) {
   EXPECT_TRUE(has_triton);
 }
 
+TEST_F(AutotunerFlagsTest, GetEnabledBackendsRespectsTensorIrFlag) {
+  GpuCompiler::GpuTargetConfig target_config(stream_executor_);
+  GpuAliasInfo alias_info(stream_executor_->GetDeviceDescription());
+  mlir::MLIRContext mlir_context;
+  RegisterSymbolicExprStorage(&mlir_context);
+  MlirContextPool mlir_context_pool(CreateMlirContext);
+
+  {
+    DebugOptions debug_options = GetDebugOptionsForTest();
+    debug_options.set_xla_gpu_experimental_enable_tensor_ir(false);
+
+    ASSERT_OK_AND_ASSIGN(std::vector<std::unique_ptr<CodegenBackend>> backends,
+                         ConfigAssignerPass::GetEnabledBackends(
+                             stream_executor_, allocator_.get(), &target_config,
+                             &alias_info, debug_options, &mlir_context,
+                             /*shape_size_fn=*/[](const Shape&) { return 0; },
+                             &compiler_, stream_executor_->GetPlatform()->id(),
+                             /*thread_pool=*/nullptr, &mlir_context_pool));
+
+    for (const auto& backend : backends) {
+      EXPECT_NE(backend->backend(), autotuner::Backend::TENSOR_IR);
+    }
+  }
+
+  {
+    DebugOptions debug_options = GetDebugOptionsForTest();
+    debug_options.set_xla_gpu_experimental_enable_tensor_ir(true);
+
+    ASSERT_OK_AND_ASSIGN(std::vector<std::unique_ptr<CodegenBackend>> backends,
+                         ConfigAssignerPass::GetEnabledBackends(
+                             stream_executor_, allocator_.get(), &target_config,
+                             &alias_info, debug_options, &mlir_context,
+                             /*shape_size_fn=*/[](const Shape&) { return 0; },
+                             &compiler_, stream_executor_->GetPlatform()->id(),
+                             /*thread_pool=*/nullptr, &mlir_context_pool));
+
+    bool has_tensor_ir = false;
+    for (const auto& backend : backends) {
+      if (backend->backend() == autotuner::Backend::TENSOR_IR) {
+        has_tensor_ir = true;
+        break;
+      }
+    }
+    EXPECT_TRUE(has_tensor_ir);
+  }
+}
+
 TEST_F(ConfigAssignerPassTest, CublasLtSelectFirstConfig) {
   absl::SetVLogLevel("config_assigner*", 10);
   AutotunerCache::ClearAutotuneResults();
diff --git a/third_party/xla/xla/service/gpu/backend_configs.proto b/third_party/xla/xla/service/gpu/backend_configs.proto
index d41d52e154c..9de21d89f22 100644
--- a/third_party/xla/xla/service/gpu/backend_configs.proto
+++ b/third_party/xla/xla/service/gpu/backend_configs.proto
@@ -284,7 +284,7 @@ message DynamicSliceConfig {
   }
 }
 
-// Next id: 8
+// Next id: 9
 message FusionBackendConfig {
   // kLoop, kInput, or kOutput (from HloInstruction::FusionKind), or your own
   // custom string.
@@ -312,10 +312,25 @@ message FusionBackendConfig {
 
   CuDnnFusionConfig cudnn_fusion_config = 5;
 
+  // Only valid when kind == "__tensorir". Code generation of such fusions
+  // will fail if this field is not set.
+  TensorIrFusionConfig tensor_ir_fusion_config = 8;
+
   reserved 3, 7;
   reserved "dynamic_memcpy_config";
 }
 
+// Backend config for fusions codegen'd through the TensorIR/CudaTile
+// pipeline (kind == "__tensorir").
+message TensorIrFusionConfig {
+  // Tile size for each dimension of the fusion's output.
+  repeated int32 tile_size = 1;
+
+  // Tile size for the reduction (contracting) dimension. Ignored if the
+  // fusion doesn't contain a reduction.
+  int64 reduction_tile_size = 2;
+}
+
 // Backed config for norm executed by cuDNN.
 message CudnnNormBackendConfig {
   // Epsilon parameter.
diff --git a/third_party/xla/xla/service/gpu/hlo_fusion_analysis.cc b/third_party/xla/xla/service/gpu/hlo_fusion_analysis.cc
index b821481aa03..0003e06c35f 100644
--- a/third_party/xla/xla/service/gpu/hlo_fusion_analysis.cc
+++ b/third_party/xla/xla/service/gpu/hlo_fusion_analysis.cc
@@ -131,6 +131,10 @@ HloFusionAnalysis::EmitterFusionKind GetEmitterFusionKind(
     return HloFusionAnalysis::EmitterFusionKind::kCuDnn;
   }
 
+  if (fusion_backend_config.kind() == kTensorIrFusionKind) {
+    return HloFusionAnalysis::EmitterFusionKind::kTensorIr;
+  }
+
   std::optional<HloInstructionAdaptor> first_reduce_hero;
   for (auto [root, hero] : llvm::zip(fusion_roots, fusion_heroes)) {
     if (IsRealReductionHero(root.instruction(), hero.instruction(),
diff --git a/third_party/xla/xla/service/gpu/hlo_fusion_analysis.h b/third_party/xla/xla/service/gpu/hlo_fusion_analysis.h
index 24c279ff24d..c33bf5bf702 100644
--- a/third_party/xla/xla/service/gpu/hlo_fusion_analysis.h
+++ b/third_party/xla/xla/service/gpu/hlo_fusion_analysis.h
@@ -52,6 +52,7 @@ class HloFusionAnalysis {
     kScatter,
     kCuDnn,
     kSort,
+    kTensorIr,
   };
 
   // Precomputed information about inputs (arguments) and outputs (roots) of the
diff --git a/third_party/xla/xla/service/gpu/ir_emission_utils.h b/third_party/xla/xla/service/gpu/ir_emission_utils.h
index 1d1fd800180..fe5aa99990f 100644
--- a/third_party/xla/xla/service/gpu/ir_emission_utils.h
+++ b/third_party/xla/xla/service/gpu/ir_emission_utils.h
@@ -130,6 +130,10 @@ inline constexpr absl::string_view kTritonNestedGemmFusionKind =
 // Fusions that use cuDNN have FusionBackendConfig.kind equal to this string.
 inline constexpr absl::string_view kCuDnnFusionKind = "__cudnn$fusion";
 
+// Fusions codegen'd through the TensorIR/CudaTile pipeline have
+// FusionBackendConfig.kind equal to this string.
+inline constexpr absl::string_view kTensorIrFusionKind = "__tensorir";
+
 inline constexpr absl::string_view kUncompilableFusion =
     "__uncompilable_fusion";
 
diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index aea36a7d7e3..d696df8c9d1 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -1942,7 +1942,12 @@ message DebugOptions {
   // Enables Raft for stable TopK.
   optional bool xla_gpu_experimental_enable_raft_for_stable_topk = 546;
 
-  // Next id: 557
+  // Lets the GPU autotuner consider the experimental TensorIR (CUDA Tile)
+  // fusion emitter as a codegen backend, so fusions may be assigned to it.
+  // Requires sm_90+ and CUDA 13.3+ (tileiras) at runtime. Off by default.
+  optional bool xla_gpu_experimental_enable_tensor_ir = 557;
+
+  // Next id: 558
 
   // Extra options to pass to the compilation backend (e.g. LLVM); specific
   // interpretation of these values is left to the backend.

From 50d6afb1809af1d20431470dd1d9a15ff68dee29 Mon Sep 17 00:00:00 2001
From: Ce Zheng CA <zce@google.com>
Date: Mon, 12 Oct 2026 16:46:44 -0700
Subject: [PATCH 0630/1000] [XLA:CPU] Deduplicate Eigen device instantiations
 and remove redundant TensorReshapingOp wrappers in convolution_lib.h

In `xla/backends/cpu/runtime/convolution_lib.h`:
1. Previously, `EigenTransposedConv2D`, `EigenGenericConv2D`, and `EigenConv3D` branched on `if (device != nullptr) { out.device(*device, ...) = expr; } else { out = expr; }`. Evaluating `out = expr` without `.device(...)` instantiated a second copy of the entire `Eigen::TensorEvaluator`, `gemm_pack_rhs`, `gemm_pack_lhs`, `general_matrix_vector_product`, `TensorImagePatchOp`, and `TensorVolumePatchOp` template hierarchy for `Eigen::DefaultDevice` alongside `Eigen::ThreadPoolDevice`. Evaluating on `Eigen::ThreadPoolDevice(/*pool=*/nullptr, /*num_cores=*/1)` when `device == nullptr` takes the `num_threads == 1` fast path in `TensorContractionThreadPool.h` (`evalProductSequential`) and `TensorDeviceThreadPool.h` (`parallelForAsync`) synchronously on the calling thread without touching `pool_`, eliminating the duplicate `Eigen::DefaultDevice` instantiations.
2. Map `C`, `input`, `kernel`, and `output` directly with their target `Eigen::TensorMap` ranks instead of constructing 4D/5D `TensorMap`s and wrapping `input`, `kernel`, `output`, and `.contract(...)` in `Eigen::TensorReshapingOp` (`input_reshaped_dims`, `output_reshaped_dims`, `kernel_reshaped_dims`, `kernel_dims`, and `post_contract_dims`).

Runtime and memory impact:
- CPU / Memory: Zero overhead; single-threaded execution still runs `evalProductSequential` directly on the calling thread, and removing the intermediate `TensorReshapingOp` wrappers simplifies the Eigen expression tree by 68 net lines of code.
- Binary size: Reduces `libconvolution_lib.a` (`-c opt --copt=-gmlt`) by `-1,664,986 B` (`-24.73%`, from `6,733,274 B` to `5,068,288 B`).

PiperOrigin-RevId: 998234110
---
 .../backends/cpu/runtime/convolution_lib.h    | 289 +++++++-----------
 1 file changed, 102 insertions(+), 187 deletions(-)

diff --git a/third_party/xla/xla/backends/cpu/runtime/convolution_lib.h b/third_party/xla/xla/backends/cpu/runtime/convolution_lib.h
index cc42cfd2c07..ff1c82b0dda 100644
--- a/third_party/xla/xla/backends/cpu/runtime/convolution_lib.h
+++ b/third_party/xla/xla/backends/cpu/runtime/convolution_lib.h
@@ -181,8 +181,8 @@ void EigenTransposedConv2D(
   // Grouped convolutions are not supported yet.
   CHECK(kernel_channels == input_channels);
 
-  using TensorMap2D =
-      Eigen::TensorMap<Eigen::Tensor<ScalarType, 2, Eigen::RowMajor>,
+  using TensorMap3D =
+      Eigen::TensorMap<Eigen::Tensor<ScalarType, 3, Eigen::RowMajor>,
                        Eigen::Unaligned>;
   using ConstTensorMap3D =
       Eigen::TensorMap<Eigen::Tensor<const ScalarType, 3, Eigen::RowMajor>,
@@ -215,8 +215,8 @@ void EigenTransposedConv2D(
       Eigen::IndexPair<Eigen::DenseIndex>(1, 1)};
 
   // Compute intermediate results (convolution matrix) into conv_matrix.
-  TensorMap2D C(conv_matrix_data, input_batch * input_image_size,
-                kernel_total_size);
+  TensorMap3D C(conv_matrix_data, input_batch * input_image_size,
+                kernel_x * kernel_y, kernel_filters);
 
   ConstTensorMap2D A(lhs, input_batch * input_image_size, input_channels);
   ConstTensorMap3D B(rhs, kernel_x * kernel_y, kernel_channels, kernel_filters);
@@ -249,22 +249,10 @@ void EigenTransposedConv2D(
     done();
   };
 
-  // Molds the output of the contraction into the shape expected by packing
-  // algorithm:
-  // - the minor dimension (dims[1]): the patch values to be packed; contiguous
-  //   in memory
-  // - the major dimension (dims[0]): everything else
-  Eigen::DSizes<Eigen::Index, 2> post_contract_dims;
-  post_contract_dims[0] = input_batch * input_image_size;
-  post_contract_dims[1] = kernel_total_size;
-
-  if (device != nullptr) {
-    C.device(*device, std::move(pack_patches)) =
-        A.contract(B, contract_dims).reshape(post_contract_dims);
-  } else {
-    C = A.contract(B, contract_dims).reshape(post_contract_dims);
-    pack_patches();
-  }
+  Eigen::ThreadPoolDevice single_thread_device(/*pool=*/nullptr,
+                                               /*num_cores=*/1);
+  C.device(device != nullptr ? *device : single_thread_device,
+           std::move(pack_patches)) = A.contract(B, contract_dims);
 }
 
 // Algorithm that works for all types of 2D convolutions. Even though it works
@@ -283,41 +271,9 @@ void EigenGenericConv2D(
     Eigen::Index lhs_y_dilation, Eigen::Index rhs_x_dilation,
     Eigen::Index rhs_y_dilation, Eigen::Index feature_group_count,
     DoneCallback done) {
-  // For non-grouped convolutions, we can optimize Eigen expressions and avoid
-  // introducing an extra dimension of size `1`.
-  if constexpr (!is_grouped) {
-    DCHECK_EQ(feature_group_count, 1) << "Expected feature group count to be 1";
-  }
-
-  const Eigen::TensorMap<Eigen::Tensor<const ScalarType, 4, Eigen::RowMajor>,
-                         Eigen::Aligned>
-      input(lhs, input_batch, input_x, input_y, input_channels);
-
-  const Eigen::TensorMap<Eigen::Tensor<const ScalarType, 4, Eigen::RowMajor>,
-                         Eigen::Aligned>
-      kernel(rhs, kernel_x, kernel_y, kernel_channels, kernel_filters);
-
-  Eigen::TensorMap<Eigen::Tensor<ScalarType, 4, Eigen::RowMajor>,
-                   Eigen::Aligned>
-      output(out, input_batch, output_x, output_y, kernel_filters);
-
   Eigen::array<Eigen::IndexPair<Eigen::Index>, 1> contract_dims;
   contract_dims[0] = Eigen::IndexPair<Eigen::Index>(1, 0);
 
-  Eigen::DSizes<Eigen::Index, 5> input_reshaped_dims;
-  input_reshaped_dims[0] = input_batch;
-  input_reshaped_dims[1] = input_x;
-  input_reshaped_dims[2] = input_y;
-  input_reshaped_dims[3] = feature_group_count;
-  input_reshaped_dims[4] = input_channels / feature_group_count;
-
-  Eigen::DSizes<Eigen::Index, 5> output_reshaped_dims;
-  output_reshaped_dims[0] = input_batch;
-  output_reshaped_dims[1] = output_x;
-  output_reshaped_dims[2] = output_y;
-  output_reshaped_dims[3] = feature_group_count;
-  output_reshaped_dims[4] = kernel_filters / feature_group_count;
-
   // Molds the output of the patch extraction code into a 2d tensor:
   // - the first dimension (dims[0]): the patch values to be multiplied with the
   //   kernels
@@ -326,47 +282,24 @@ void EigenGenericConv2D(
   pre_contract_dims[0] = output_y * output_x * input_batch;
   pre_contract_dims[1] = kernel_channels * kernel_y * kernel_x;
 
-  // Molds the output of the contraction into the shape expected by the user:
-  Eigen::DSizes<Eigen::Index, 4> post_contract_dims;
-  post_contract_dims[0] = input_batch;
-  post_contract_dims[1] = output_x;
-  post_contract_dims[2] = output_y;
-  post_contract_dims[3] = kernel_filters / feature_group_count;
+  // IMPORTANT: Below, the row and column dimensions must be flipped when passed
+  // to Eigen.
 
-  // Kernel dimensions for non-grouped convolution.
-  Eigen::DSizes<Eigen::Index, 2> kernel_dims;
-  kernel_dims[0] = kernel_channels * kernel_y * kernel_x;
-  kernel_dims[1] = kernel_filters;
+  // For non-grouped convolutions, we can optimize Eigen expressions and avoid
+  // introducing an extra dimension of size `1`.
+  if constexpr (!is_grouped) {
+    DCHECK_EQ(feature_group_count, 1) << "Expected feature group count to be 1";
 
-  // Kernel dimensions for grouped convolution.
-  Eigen::DSizes<Eigen::Index, 3> kernel_reshaped_dims;
-  kernel_reshaped_dims[0] = kernel_channels * kernel_y * kernel_x;
-  kernel_reshaped_dims[1] = feature_group_count;
-  kernel_reshaped_dims[2] = kernel_filters / feature_group_count;
+    const Eigen::TensorMap<Eigen::Tensor<const ScalarType, 4, Eigen::RowMajor>,
+                           Eigen::Aligned>
+        input(lhs, input_batch, input_x, input_y, input_channels);
+    const Eigen::TensorMap<Eigen::Tensor<const ScalarType, 2, Eigen::RowMajor>,
+                           Eigen::Aligned>
+        kernel(rhs, pre_contract_dims[1], kernel_filters);
+    Eigen::TensorMap<Eigen::Tensor<ScalarType, 2, Eigen::RowMajor>,
+                     Eigen::Aligned>
+        output(out, pre_contract_dims[0], kernel_filters);
 
-  // IMPORTANT: Below in `convolve_group` and `convolve` lambdas, the row and
-  // column dimensions must be flipped when passed to Eigen.
-
-  // Constructs convolution and output expressions for a given group index.
-  auto convolve_group = [=](int64_t i) {
-    auto convolved =
-        input.reshape(input_reshaped_dims)
-            .chip(i, 3)
-            .extract_image_patches(
-                kernel_y, kernel_x, y_stride, x_stride, rhs_y_dilation,
-                rhs_x_dilation, lhs_y_dilation, lhs_x_dilation,
-                padding_y_before, padding_y_after, padding_x_before,
-                padding_x_after, static_cast<ScalarType>(0.0f))
-            .reshape(pre_contract_dims)
-            .contract(kernel.reshape(kernel_reshaped_dims).chip(i, 1),
-                      contract_dims)
-            .reshape(post_contract_dims);
-    auto output_reshaped = output.reshape(output_reshaped_dims).chip(i, 3);
-    return std::make_pair(output_reshaped, convolved);
-  };
-
-  // Constructs convolution and output expressions for full input.
-  auto convolve = [=] {
     auto convolved =
         input
             .extract_image_patches(
@@ -375,57 +308,70 @@ void EigenGenericConv2D(
                 padding_y_before, padding_y_after, padding_x_before,
                 padding_x_after, static_cast<ScalarType>(0.0f))
             .reshape(pre_contract_dims)
-            .contract(kernel.reshape(kernel_dims), contract_dims)
-            .reshape(post_contract_dims);
-    return std::make_pair(output, convolved);
-  };
+            .contract(kernel, contract_dims);
 
-  // For non-grouped convolutions, we need to execute only one Eigen expression.
-  if constexpr (!is_grouped) {
-    auto [output, convolved] = convolve();
-
-    if (device != nullptr) {
-      output.device(*device, std::move(done)) = convolved;
-    } else {
-      output = convolved;
-      done();
-    }
-
-    return;
-  }
-
-  // For grouped convolutions, we need to execute multiple Eigen expressions and
-  // we might use thread pool to run them concurrently.
-  if (device != nullptr) {
-    // Although we schedule at most one tasks for each thread, individual
-    // convolution might also schedule more tasks into the same thread pool.
-    auto max_tasks = static_cast<Eigen::Index>(device->numThreads());
-    auto task_size = Eigen::numext::div_ceil(feature_group_count, max_tasks);
-    auto num_tasks = Eigen::numext::div_ceil(feature_group_count, task_size);
+    Eigen::ThreadPoolDevice single_thread_device(/*pool=*/nullptr,
+                                                 /*num_cores=*/1);
+    output.device(device != nullptr ? *device : single_thread_device,
+                  std::move(done)) = convolved;
+  } else {
+    const Eigen::TensorMap<Eigen::Tensor<const ScalarType, 5, Eigen::RowMajor>,
+                           Eigen::Aligned>
+        input(lhs, input_batch, input_x, input_y, feature_group_count,
+              input_channels / feature_group_count);
+    const Eigen::TensorMap<Eigen::Tensor<const ScalarType, 3, Eigen::RowMajor>,
+                           Eigen::Aligned>
+        kernel(rhs, pre_contract_dims[1], feature_group_count,
+               kernel_filters / feature_group_count);
 
     // Signal done callback when all feature groups are done.
     tsl::CountDownAsyncValueRef<tsl::Chain> count_down(feature_group_count);
     count_down.AsPtr().AndThen([done = std::move(done)]() mutable { done(); });
 
-    Worker::Parallelize(
-        device->getPool(), /*num_workers=*/num_tasks, num_tasks,
-        [=](Eigen::Index task_index) {
-          Eigen::Index start = task_index * task_size;
-          Eigen::Index end = std::min(start + task_size, feature_group_count);
-          for (Eigen::Index i = start; i < end; ++i) {
-            auto on_done = [count_down]() mutable { count_down.CountDown(); };
-            auto [output, convolved] = convolve_group(i);
-            output.device(*device, std::move(on_done)) = convolved;
-          }
-        });
+    auto convolve_range = [=](const Eigen::ThreadPoolDevice& dev,
+                              Eigen::Index start, Eigen::Index end) {
+      Eigen::TensorMap<Eigen::Tensor<ScalarType, 3, Eigen::RowMajor>,
+                       Eigen::Aligned>
+          output(out, pre_contract_dims[0], feature_group_count,
+                 kernel_filters / feature_group_count);
+      for (Eigen::Index i = start; i < end; ++i) {
+        auto on_done = [count_down]() mutable { count_down.CountDown(); };
+        auto convolved =
+            input.chip(i, 3)
+                .extract_image_patches(
+                    kernel_y, kernel_x, y_stride, x_stride, rhs_y_dilation,
+                    rhs_x_dilation, lhs_y_dilation, lhs_x_dilation,
+                    padding_y_before, padding_y_after, padding_x_before,
+                    padding_x_after, static_cast<ScalarType>(0.0f))
+                .reshape(pre_contract_dims)
+                .contract(kernel.chip(i, 1), contract_dims);
+        output.chip(i, 1).device(dev, std::move(on_done)) = convolved;
+      }
+    };
 
-  } else {
-    // Convolve all feature groups sequentially in the caller thread.
-    for (Eigen::Index i = 0; i < feature_group_count; ++i) {
-      auto [output, convolved] = convolve_group(i);
-      output = convolved;
+    // For grouped convolutions, we need to execute multiple Eigen expressions
+    // and we might use thread pool to run them concurrently.
+    if (device != nullptr) {
+      // Although we schedule at most one tasks for each thread, individual
+      // convolution might also schedule more tasks into the same thread pool.
+      auto max_tasks = static_cast<Eigen::Index>(device->numThreads());
+      auto task_size = Eigen::numext::div_ceil(feature_group_count, max_tasks);
+      auto num_tasks = Eigen::numext::div_ceil(feature_group_count, task_size);
+
+      Worker::Parallelize(device->getPool(), /*num_workers=*/num_tasks,
+                          num_tasks, [=](Eigen::Index task_index) {
+                            Eigen::Index start = task_index * task_size;
+                            Eigen::Index end = std::min(start + task_size,
+                                                        feature_group_count);
+                            convolve_range(*device, start, end);
+                          });
+
+    } else {
+      // Convolve all feature groups sequentially in the caller thread.
+      Eigen::ThreadPoolDevice single_thread_device(/*pool=*/nullptr,
+                                                   /*num_cores=*/1);
+      convolve_range(single_thread_device, 0, feature_group_count);
     }
-    done();
   }
 }
 
@@ -496,35 +442,6 @@ void EigenConv3D(const Eigen::ThreadPoolDevice* device, ScalarType* out,
                  Eigen::Index lhs_z_dilation, Eigen::Index rhs_x_dilation,
                  Eigen::Index rhs_y_dilation, Eigen::Index rhs_z_dilation,
                  Eigen::Index feature_group_count, DoneCallback done) {
-  using ConstTType =
-      Eigen::TensorMap<Eigen::Tensor<const ScalarType, 5, Eigen::RowMajor>,
-                       Eigen::Aligned>;
-  const ConstTType input(lhs, input_batch, input_x, input_y, input_z,
-                         input_channels);
-
-  const ConstTType kernel(rhs, kernel_x, kernel_y, kernel_z, kernel_channels,
-                          kernel_filters);
-
-  Eigen::TensorMap<Eigen::Tensor<ScalarType, 5, Eigen::RowMajor>,
-                   Eigen::Aligned>
-      output(out, input_batch, output_x, output_y, output_z, kernel_filters);
-
-  Eigen::DSizes<Eigen::Index, 6> input_reshaped_dims;
-  input_reshaped_dims[0] = input_batch;
-  input_reshaped_dims[1] = input_x;
-  input_reshaped_dims[2] = input_y;
-  input_reshaped_dims[3] = input_z;
-  input_reshaped_dims[4] = feature_group_count;
-  input_reshaped_dims[5] = input_channels / feature_group_count;
-
-  Eigen::DSizes<Eigen::Index, 6> output_reshaped_dims;
-  output_reshaped_dims[0] = input_batch;
-  output_reshaped_dims[1] = output_x;
-  output_reshaped_dims[2] = output_y;
-  output_reshaped_dims[3] = output_z;
-  output_reshaped_dims[4] = feature_group_count;
-  output_reshaped_dims[5] = kernel_filters / feature_group_count;
-
   Eigen::array<Eigen::IndexPair<Eigen::Index>, 1> contract_dims;
   contract_dims[0] = Eigen::IndexPair<Eigen::Index>(1, 0);
 
@@ -536,26 +453,33 @@ void EigenConv3D(const Eigen::ThreadPoolDevice* device, ScalarType* out,
   pre_contract_dims[0] = output_x * output_y * output_z * input_batch;
   pre_contract_dims[1] = kernel_channels * kernel_x * kernel_y * kernel_z;
 
-  // Molds the output of the contraction into the shape expected by the user:
-  Eigen::DSizes<Eigen::Index, 5> post_contract_dims;
-  post_contract_dims[0] = input_batch;
-  post_contract_dims[1] = output_x;
-  post_contract_dims[2] = output_y;
-  post_contract_dims[3] = output_z;
-  post_contract_dims[4] = kernel_filters / feature_group_count;
+  const Eigen::TensorMap<Eigen::Tensor<const ScalarType, 6, Eigen::RowMajor>,
+                         Eigen::Aligned>
+      input(lhs, input_batch, input_x, input_y, input_z, feature_group_count,
+            input_channels / feature_group_count);
 
-  Eigen::DSizes<Eigen::Index, 3> kernel_dims;
-  kernel_dims[0] = kernel_channels * kernel_x * kernel_y * kernel_z;
-  kernel_dims[1] = feature_group_count;
-  kernel_dims[2] = kernel_filters / feature_group_count;
+  const Eigen::TensorMap<Eigen::Tensor<const ScalarType, 3, Eigen::RowMajor>,
+                         Eigen::Aligned>
+      kernel(rhs, pre_contract_dims[1], feature_group_count,
+             kernel_filters / feature_group_count);
+
+  Eigen::TensorMap<Eigen::Tensor<ScalarType, 3, Eigen::RowMajor>,
+                   Eigen::Aligned>
+      output(out, pre_contract_dims[0], feature_group_count,
+             kernel_filters / feature_group_count);
 
   // Signal done callback when all feature groups are done.
   tsl::CountDownAsyncValueRef<tsl::Chain> count_down(feature_group_count);
   count_down.AsPtr().AndThen([done = std::move(done)]() mutable { done(); });
 
+  Eigen::ThreadPoolDevice single_thread_device(/*pool=*/nullptr,
+                                               /*num_cores=*/1);
+  const Eigen::ThreadPoolDevice& dev =
+      device != nullptr ? *device : single_thread_device;
+
   for (Eigen::Index i = 0; i < feature_group_count; ++i) {
     // The dimension order must be flipped when passed to Eigen.
-    auto input_chip = input.reshape(input_reshaped_dims).chip(i, 4);
+    auto input_chip = input.chip(i, 4);
     auto patches =
         Eigen::TensorVolumePatchOp<Eigen::Dynamic, Eigen::Dynamic,
                                    Eigen::Dynamic, decltype(input_chip)>(
@@ -565,20 +489,11 @@ void EigenConv3D(const Eigen::ThreadPoolDevice* device, ScalarType* out,
             padding_z_after, padding_y_before, padding_y_after,
             padding_x_before, padding_x_after, static_cast<ScalarType>(0.0f));
 
-    auto convolved =
-        patches.reshape(pre_contract_dims)
-            .contract(kernel.reshape(kernel_dims).chip(i, 1), contract_dims)
-            .reshape(post_contract_dims);
+    auto convolved = patches.reshape(pre_contract_dims)
+                         .contract(kernel.chip(i, 1), contract_dims);
 
-    auto output_reshaped = output.reshape(output_reshaped_dims).chip(i, 4);
-
-    if (device != nullptr) {
-      auto on_done = [count_down]() mutable { count_down.CountDown(); };
-      output_reshaped.device(*device, std::move(on_done)) = convolved;
-    } else {
-      output_reshaped = convolved;
-      count_down.CountDown();
-    }
+    auto on_done = [count_down]() mutable { count_down.CountDown(); };
+    output.chip(i, 1).device(dev, std::move(on_done)) = convolved;
   }
 }
 

From e01806cec204c9860b9b88c3478288bc00cfdae8 Mon Sep 17 00:00:00 2001
From: Enver Kayaaslan <enver@google.com>
Date: Mon, 12 Oct 2026 16:49:31 -0700
Subject: [PATCH 0631/1000] Factor out the predicate for the pattern match on
 allgather(broadcast).

PiperOrigin-RevId: 998235484
---
 .../all_gather_broadcast_reorder.cc           | 78 ++++++++++++-------
 1 file changed, 48 insertions(+), 30 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/collectives/all_gather_broadcast_reorder.cc b/third_party/xla/xla/hlo/transforms/collectives/all_gather_broadcast_reorder.cc
index 96f4a003c19..0c89f4184a3 100644
--- a/third_party/xla/xla/hlo/transforms/collectives/all_gather_broadcast_reorder.cc
+++ b/third_party/xla/xla/hlo/transforms/collectives/all_gather_broadcast_reorder.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "xla/hlo/transforms/collectives/all_gather_broadcast_reorder.h"
 
 #include <cstdint>
+#include <optional>
 #include <vector>
 
 #include "absl/container/flat_hash_set.h"
@@ -25,6 +26,7 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "absl/types/span.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_opcode.h"
@@ -35,6 +37,48 @@ limitations under the License.
 #include "tsl/platform/errors.h"
 
 namespace xla {
+namespace {
+
+// Returns whether `inst` matches the all-gather(broadcast) reorder pattern,
+// and if so, whether the all-gather is along a uniform dimension.
+std::optional<bool> MatchAllGatherBroadcastReorder(const HloInstruction* inst) {
+  // Check for all-gather with a broadcast operand.
+  if (inst->opcode() != HloOpcode::kAllGather || !inst->shape().IsArray() ||
+      inst->operand(0)->opcode() != HloOpcode::kBroadcast) {
+    return std::nullopt;
+  }
+
+  const auto* ag = Cast<HloAllGatherInstruction>(inst);
+  const auto* bcast = Cast<HloBroadcastInstruction>(inst->operand(0));
+
+  // We categorize each dimension of the all-gather result as either
+  // uniform (same data along that dimension) or not. The all_gather
+  // dimension is not uniform since we don't know uniformity of data across
+  // the all-gather participants. In addition, the set of dimensions
+  // for the broadcast instruction feeding into all-gather are also not
+  // uniform. If there are any remaining uniform dims and their size > 1,
+  // then doing the broadcast after the all-gather is beneficial as the
+  // all-gather will be operating on smaller data.
+  absl::flat_hash_set<int64_t> non_uniform_dims(bcast->dimensions().begin(),
+                                                bcast->dimensions().end());
+  const bool all_gather_along_uniform_dim =
+      non_uniform_dims.insert(ag->all_gather_dimension()).second;
+
+  // Find the product of the size of uniform dims.
+  int64_t uniform_dim_size = 1;
+  for (int64_t i = 0; i < ag->shape().dimensions().size(); ++i) {
+    if (!non_uniform_dims.contains(i)) {
+      uniform_dim_size *= ag->shape().dimensions(i);
+    }
+  }
+
+  if (uniform_dim_size == 1) {
+    return std::nullopt;
+  }
+  return all_gather_along_uniform_dim;
+}
+
+}  // namespace
 
 absl::StatusOr<bool> AllGatherBroadcastReorder::RunImpl(
     HloModule* module,
@@ -51,9 +95,9 @@ absl::StatusOr<bool> AllGatherBroadcastReorder::RunImpl(
   bool changed = false;
   for (auto computation : module->computations(execution_threads)) {
     for (HloInstruction *inst : computation->MakeInstructionPostOrder()) {
-      // Check for all-gather with a broadcast operand.
-      if (inst->opcode() != HloOpcode::kAllGather || !inst->shape().IsArray() ||
-          inst->operand(0)->opcode() != HloOpcode::kBroadcast) {
+      std::optional<bool> all_gather_along_uniform_dim =
+          MatchAllGatherBroadcastReorder(inst);
+      if (!all_gather_along_uniform_dim.has_value()) {
         continue;
       }
 
@@ -61,37 +105,11 @@ absl::StatusOr<bool> AllGatherBroadcastReorder::RunImpl(
       HloBroadcastInstruction *bcast =
           Cast<HloBroadcastInstruction>(inst->mutable_operand(0));
 
-      // We categorize each dimension of the all-gather result as either
-      // uniform (same data along that dimension) or not. The all_gather
-      // dimension is not uniform since we don't know uniformity of data across
-      // the all-gather participants. In addition, the set of dimensions
-      // for the broadcast instruction feeding into all-gather are also not
-      // uniform. If there are any remaining uniform dims and their size > 1,
-      // then doing the broadcast after the all-gather is beneficial as the
-      // all-gather will be operating on smaller data.
-      absl::flat_hash_set<int64_t> non_uniform_dims;
-      non_uniform_dims.insert(bcast->dimensions().begin(),
-                              bcast->dimensions().end());
-      const bool all_gather_along_uniform_dim =
-          non_uniform_dims.insert(ag->all_gather_dimension()).second;
-
-      // Find the product of the size of uniform dims.
-      int64_t uniform_dim_size = 1;
-      for (int64_t i = 0; i < ag->shape().dimensions().size(); ++i) {
-        if (non_uniform_dims.count(i) == 0) {
-          uniform_dim_size *= ag->shape().dimensions(i);
-        }
-      }
-
-      if (uniform_dim_size == 1) {
-        continue;
-      }
-
       HloInstruction *replacement;
       const int64_t ag_dim = ag->all_gather_dimension();
       // Transform the all-gather(broadcast(x)) to broadcast(all-gather(x)).
       // There are 2 cases here:
-      if (!all_gather_along_uniform_dim) {
+      if (!*all_gather_along_uniform_dim) {
         // If the all-gather happens along one of the non-uniform dimensions of
         // the broadcast, then issue all-gather(x) and then a broadcast.
         // Example:

From 74b42410d03bcceb35190ba321e01e25ba2dc37b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 12 Oct 2026 17:04:48 -0700
Subject: [PATCH 0632/1000] [XLA] Add sample boilerplate for PjRtRelocatable

This change adds an implementation of PjRtRelocatable to the sample
phase compiler implementation and adds tests for it. It also adds an
implementation of CompileRelocatable.

It omits Link because we don't have a good sample PjRtExecutable we can
produce. This is good enough to produce some tests of the
PjRtRelocatable interface to demonstrate how it's supposed to be used.

PiperOrigin-RevId: 998243383
---
 third_party/xla/xla/pjrt/BUILD                |  35 +++
 .../pjrt/pjrt_phase_compile_sample_plugin.cc  | 250 +++++++++++++++++-
 .../pjrt/pjrt_phase_compile_sample_plugin.h   |  84 ++++++
 .../xla/xla/pjrt/pjrt_phase_compiler_test.cc  | 101 +++++++
 .../xla/xla/pjrt/pjrt_relocatable_test.cc     | 227 +++++++++++++++-
 5 files changed, 691 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/xla/pjrt/BUILD b/third_party/xla/xla/pjrt/BUILD
index 1cfa8417c3f..2c8cb2366ed 100644
--- a/third_party/xla/xla/pjrt/BUILD
+++ b/third_party/xla/xla/pjrt/BUILD
@@ -1317,7 +1317,15 @@ cc_library(
         ":pjrt_client",
         ":pjrt_compiler",
         ":pjrt_executable",
+        ":pjrt_layout",
+        ":pjrt_relocatable",
+        "//xla:shape_util",
+        "//xla:status_macros",
+        "//xla:util",
+        "//xla:xla_data_proto_cc",
         "//xla/hlo/builder:xla_computation",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/translate:stablehlo",
         "//xla/pjrt/c:pjrt_c_api_hdrs",
         "//xla/pjrt/c:pjrt_c_api_helpers",
         "//xla/pjrt/c:pjrt_c_api_phase_compile_extension_hdrs",
@@ -1330,16 +1338,23 @@ cc_library(
         "//xla/tsl/framework/mlir:status_scoped_diagnostic_handler",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:Transforms",
+        "@riegeli//riegeli/bytes:string_reader",
+        "@riegeli//riegeli/bytes:string_writer",
+        "@riegeli//riegeli/messages:parse_message",
+        "@riegeli//riegeli/messages:serialize_message",
         "@stablehlo//:stablehlo_passes_optimization",
         "@stablehlo//:stablehlo_portable_api",
         "@stablehlo//:stablehlo_serialization",
+        "@tsl//tsl/platform:fingerprint",
     ],
 )
 
@@ -1363,6 +1378,7 @@ xla_cc_test(
         ":pjrt_compiler",
         ":pjrt_executable",
         ":pjrt_phase_compile_sample_plugin",
+        ":pjrt_relocatable",
         "//xla:debug_options_flags",
         "//xla/backends/cpu:target_machine_options",
         "//xla/hlo/builder:xla_computation",
@@ -1373,6 +1389,7 @@ xla_cc_test(
         "//xla/pjrt/proto:pjrt_partial_program_proto_cc",
         "//xla/tsl/lib/core:status_test_util",
         "//xla/tsl/util/proto:proto_matchers",
+        "@com_google_absl//absl/base",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
@@ -1564,7 +1581,25 @@ xla_cc_test(
     name = "pjrt_relocatable_test",
     srcs = ["pjrt_relocatable_test.cc"],
     deps = [
+        ":maybe_owning_mlir_module",
+        ":pjrt_compiler",
+        ":pjrt_layout",
+        ":pjrt_phase_compile_sample_plugin",
         ":pjrt_relocatable",
+        "//xla:debug_options_flags",
+        "//xla:shape_util",
+        "//xla:util",
+        "//xla:xla_data_proto_cc",
+        "//xla/backends/cpu:target_machine_options",
+        "//xla/pjrt/plugin/xla_cpu:cpu_topology",
+        "//xla/pjrt/plugin/xla_cpu:cpu_topology_description",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest_main",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:Support",
+        "@stablehlo//:reference_api",
     ],
 )
diff --git a/third_party/xla/xla/pjrt/pjrt_phase_compile_sample_plugin.cc b/third_party/xla/xla/pjrt/pjrt_phase_compile_sample_plugin.cc
index 41c9e552ffd..9b9dd7ebf1c 100644
--- a/third_party/xla/xla/pjrt/pjrt_phase_compile_sample_plugin.cc
+++ b/third_party/xla/xla/pjrt/pjrt_phase_compile_sample_plugin.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/pjrt/pjrt_phase_compile_sample_plugin.h"
 
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <utility>
@@ -22,9 +23,11 @@ limitations under the License.
 
 #include "absl/log/log.h"
 #include "absl/status/status.h"
+#include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
+#include "absl/types/span.h"
 #include "llvm/Support/LogicalResult.h"
 #include "llvm/Support/raw_ostream.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
@@ -33,10 +36,20 @@ limitations under the License.
 #include "mlir/IR/OwningOpRef.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
+#include "riegeli/bytes/string_reader.h"
+#include "riegeli/bytes/string_writer.h"
+#include "riegeli/messages/parse_message.h"
+#include "riegeli/messages/serialize_message.h"
 #include "stablehlo/api/PortableApi.h"
 #include "stablehlo/dialect/Serialization.h"
 #include "stablehlo/transforms/optimization/Passes.h"
 #include "xla/hlo/builder/xla_computation.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/ir/hlo_input_output_alias_config.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/translate/stablehlo.h"
+#include "xla/layout_util.h"
 #include "xla/pjrt/c/pjrt_c_api.h"
 #include "xla/pjrt/c/pjrt_c_api_helpers.h"
 #include "xla/pjrt/c/pjrt_c_api_phase_compile_extension.h"
@@ -47,10 +60,18 @@ limitations under the License.
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/pjrt/pjrt_compiler.h"
 #include "xla/pjrt/pjrt_executable.h"
+#include "xla/pjrt/pjrt_layout.h"
+#include "xla/pjrt/pjrt_relocatable.h"
 #include "xla/pjrt/plugin/xla_cpu/cpu_client_options.h"
 #include "xla/pjrt/plugin/xla_cpu/xla_cpu_pjrt_client.h"
 #include "xla/pjrt/proto/pjrt_partial_program.pb.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+#include "xla/status_macros.h"
 #include "xla/tsl/framework/mlir/status_scoped_diagnostic_handler.h"
+#include "xla/util.h"
+#include "xla/xla_data.pb.h"
+#include "tsl/platform/fingerprint.h"
 
 namespace pjrt {
 namespace phase_compile_sample_plugin {
@@ -77,11 +98,15 @@ StablehloTypeSerialization::Deserialize(const std::string& program,
 }
 
 absl::StatusOr<std::string> StablehloTypeSerialization::Serialize(
-    const mlir::OwningOpRef<mlir::ModuleOp>& module_op) {
+    mlir::ModuleOp module_op) {
+  if (!module_op) {
+    return absl::InvalidArgumentError(
+        "SHLOFormat serialization failed: MLIR module is null");
+  }
   auto version = mlir::stablehlo::getCurrentVersion();
   std::string bytecode;
   llvm::raw_string_ostream os(bytecode);
-  if (failed(mlir::stablehlo::serializePortableArtifact(*module_op, version, os,
+  if (failed(mlir::stablehlo::serializePortableArtifact(module_op, version, os,
                                                         true))) {
     return absl::InvalidArgumentError(
         "SHLOFormat serialization failed: Could not serialize MLIR module");
@@ -89,6 +114,11 @@ absl::StatusOr<std::string> StablehloTypeSerialization::Serialize(
   return bytecode;
 }
 
+absl::StatusOr<std::string> StablehloTypeSerialization::Serialize(
+    const mlir::OwningOpRef<mlir::ModuleOp>& module_op) {
+  return Serialize(*module_op);
+}
+
 namespace {
 
 constexpr absl::string_view kStablehloBytecodeFormat = "bytecode";
@@ -158,6 +188,9 @@ absl::StatusOr<std::vector<xla::PjRtPartialProgramProto>> PhaseCompiler(
     serialized_output_object.set_producer_phase(kPhaseName);
     serialized_output_object.add_consumer_phases({std::string(kNextPhaseName)});
     serialized_output_object.set_version("1.0");
+    if (!input_program.program_name().empty()) {
+      serialized_output_object.set_program_name(input_program.program_name());
+    }
 
     serialized_output_objects.push_back(std::move(serialized_output_object));
   }
@@ -166,6 +199,190 @@ absl::StatusOr<std::vector<xla::PjRtPartialProgramProto>> PhaseCompiler(
 
 }  // namespace
 
+/*static*/
+absl::StatusOr<std::unique_ptr<SampleRelocatable>> SampleRelocatable::Create(
+    std::vector<xla::PjRtPartialProgramProto> partial_programs) {
+  if (partial_programs.empty()) {
+    return absl::InvalidArgumentError(
+        "SampleRelocatable requires at least one partial program.");
+  }
+
+  mlir::MLIRContext context;
+  ABSL_ASSIGN_OR_RETURN(mlir::OwningOpRef<mlir::ModuleOp> module_op,
+                   StablehloTypeSerialization::Deserialize(
+                       partial_programs[0].program(), context));
+  ABSL_ASSIGN_OR_RETURN(std::unique_ptr<xla::HloModule> hlo_module,
+                   xla::ConvertStablehloToHlo(*module_op));
+
+  xla::ProgramShape shape =
+      hlo_module->entry_computation()->ComputeProgramShape();
+  if (!xla::LayoutUtil::HasLayout(shape)) {
+    xla::LayoutUtil::SetToDefaultLayout(&shape);
+  }
+
+  std::vector<std::pair<int64_t, int64_t>> aliases;
+  ABSL_RETURN_IF_ERROR(
+      hlo_module->input_output_alias_config().ForEachAliasWithStatus(
+          [&](const xla::ShapeIndex& output_index,
+              const xla::HloInputOutputAliasConfig::Alias& alias)
+              -> absl::Status {
+            TF_RET_CHECK(output_index.size() <= 1)
+                << "Nested tuple outputs are unsupported.";
+            TF_RET_CHECK(alias.parameter_index.empty())
+                << "Tuple inputs are unsupported.";
+            int64_t out_idx = output_index.empty() ? 0 : output_index[0];
+            aliases.emplace_back(alias.parameter_number, out_idx);
+            return absl::OkStatus();
+          }));
+
+  bool has_side_effects = false;
+  for (const xla::HloComputation* computation : hlo_module->computations()) {
+    has_side_effects = has_side_effects || computation->HasSideEffect();
+    if (has_side_effects) {
+      break;
+    }
+  }
+
+  uint64_t fp = 0;
+  for (const xla::PjRtPartialProgramProto& prog : partial_programs) {
+    fp = tsl::FingerprintCat64(fp, tsl::Fingerprint64(prog.program()));
+  }
+  std::string name = partial_programs[0].program_name();
+  std::string build_id = absl::StrCat(absl::Hex(fp, absl::kZeroPad16));
+  std::string call_backend_config =
+      absl::StrCat("sample_backend_config:", build_id);
+
+  return std::make_unique<SampleRelocatable>(
+      name, build_id, std::move(shape), call_backend_config, std::move(aliases),
+      has_side_effects, std::move(partial_programs));
+}
+
+absl::StatusOr<std::string> SampleRelocatable::Serialize() const {
+  std::string serialized;
+  riegeli::StringWriter<> writer(&serialized);
+  for (const xla::PjRtPartialProgramProto& program : partial_programs_) {
+    ABSL_RETURN_IF_ERROR(riegeli::SerializeLengthPrefixedMessage(
+        program, writer,
+        riegeli::SerializeMessageOptions().set_deterministic(true)));
+  }
+  if (!writer.Close()) {
+    return writer.status();
+  }
+  return serialized;
+}
+
+/*static*/
+absl::StatusOr<std::unique_ptr<SampleRelocatable>>
+SampleRelocatable::Deserialize(absl::string_view serialized) {
+  if (serialized.empty()) {
+    return absl::InvalidArgumentError(
+        "Serialized SampleRelocatable cannot be empty.");
+  }
+  riegeli::StringReader<> reader(serialized);
+  std::vector<xla::PjRtPartialProgramProto> partial_programs;
+  while (reader.Pull()) {
+    xla::PjRtPartialProgramProto program;
+    ABSL_RETURN_IF_ERROR(riegeli::ParseLengthPrefixedMessage(reader, program));
+    partial_programs.push_back(std::move(program));
+  }
+  if (!reader.Close()) {
+    return reader.status();
+  }
+  return Create(std::move(partial_programs));
+}
+
+absl::StatusOr<std::vector<xla::PrimitiveType>>
+SampleRelocatable::GetParameterElementTypes() const {
+  std::vector<xla::PrimitiveType> element_types;
+  element_types.reserve(shape_.parameters().size());
+  for (const xla::Shape& shape : shape_.parameters()) {
+    TF_RET_CHECK(!shape.IsTuple())
+        << "Relocatable has unflattened program shape.";
+    element_types.push_back(shape.element_type());
+  }
+  return element_types;
+}
+
+absl::StatusOr<std::vector<xla::PrimitiveType>>
+SampleRelocatable::GetOutputElementTypes() const {
+  std::vector<xla::PrimitiveType> element_types;
+  if (!shape_.result().IsTuple()) {
+    element_types.push_back(shape_.result().element_type());
+    return element_types;
+  }
+  element_types.reserve(shape_.result().tuple_shapes().size());
+  for (const xla::Shape& shape : shape_.result().tuple_shapes()) {
+    TF_RET_CHECK(!shape.IsTuple())
+        << "Relocatable has unflattened program shape.";
+    element_types.push_back(shape.element_type());
+  }
+  return element_types;
+}
+
+absl::StatusOr<std::vector<xla::DimensionVector>>
+SampleRelocatable::GetParameterDimensions() const {
+  std::vector<xla::DimensionVector> dimensions;
+  dimensions.reserve(shape_.parameters().size());
+  for (const xla::Shape& shape : shape_.parameters()) {
+    TF_RET_CHECK(!shape.IsTuple())
+        << "Relocatable has unflattened program shape.";
+    absl::Span<const int64_t> shape_dims = shape.dimensions();
+    dimensions.emplace_back(shape_dims.begin(), shape_dims.end());
+  }
+  return dimensions;
+}
+
+absl::StatusOr<std::vector<xla::DimensionVector>>
+SampleRelocatable::GetOutputDimensions() const {
+  std::vector<xla::DimensionVector> dimensions;
+  if (!shape_.result().IsTuple()) {
+    absl::Span<const int64_t> shape_dims = shape_.result().dimensions();
+    dimensions.emplace_back(shape_dims.begin(), shape_dims.end());
+    return dimensions;
+  }
+  dimensions.reserve(shape_.result().tuple_shapes().size());
+  for (const xla::Shape& shape : shape_.result().tuple_shapes()) {
+    TF_RET_CHECK(!shape.IsTuple())
+        << "Relocatable has unflattened program shape.";
+    absl::Span<const int64_t> shape_dims = shape.dimensions();
+    dimensions.emplace_back(shape_dims.begin(), shape_dims.end());
+  }
+  return dimensions;
+}
+
+absl::StatusOr<std::vector<std::shared_ptr<const xla::PjRtLayout>>>
+SampleRelocatable::GetParameterLayouts() const {
+  std::vector<std::shared_ptr<const xla::PjRtLayout>> layouts;
+  layouts.reserve(shape_.parameters().size());
+  for (const xla::Shape& shape : shape_.parameters()) {
+    TF_RET_CHECK(!shape.IsTuple())
+        << "Relocatable has unflattened program shape.";
+    TF_RET_CHECK(shape.has_layout()) << "Relocatable is missing layouts.";
+    layouts.push_back(std::make_shared<const xla::PjRtLayout>(shape.layout()));
+  }
+  return layouts;
+}
+
+absl::StatusOr<std::vector<std::shared_ptr<const xla::PjRtLayout>>>
+SampleRelocatable::GetOutputLayouts() const {
+  std::vector<std::shared_ptr<const xla::PjRtLayout>> layouts;
+  if (!shape_.result().IsTuple()) {
+    TF_RET_CHECK(shape_.result().has_layout())
+        << "Relocatable is missing layouts.";
+    layouts.push_back(
+        std::make_shared<const xla::PjRtLayout>(shape_.result().layout()));
+    return layouts;
+  }
+  layouts.reserve(shape_.result().tuple_shapes().size());
+  for (const xla::Shape& shape : shape_.result().tuple_shapes()) {
+    TF_RET_CHECK(!shape.IsTuple())
+        << "Relocatable has unflattened program shape.";
+    TF_RET_CHECK(shape.has_layout()) << "Relocatable is missing layouts.";
+    layouts.push_back(std::make_shared<const xla::PjRtLayout>(shape.layout()));
+  }
+  return layouts;
+}
+
 absl::Status SamplePhaseCompiler::RegisterAllPhases() {
   xla::CompilationPhaseFunctions phase_functions;
   phase_functions.compiler = PhaseCompiler;
@@ -193,6 +410,35 @@ SamplePhaseCompiler::Compile(xla::CompileOptions options,
       "compiler.");
 }
 
+absl::StatusOr<std::unique_ptr<xla::PjRtRelocatable>>
+SamplePhaseCompiler::CompileToRelocatable(
+    xla::CompileOptions options, xla::MaybeOwningMlirModule module,
+    const xla::PjRtTopologyDescription& topology, absl::string_view name,
+    bool is_entrypoint, xla::PjRtClient* client) {
+  ABSL_ASSIGN_OR_RETURN(std::string serialized_module,
+                   StablehloTypeSerialization::Serialize(module.mlir_module()));
+
+  xla::PjRtPartialProgramProto partial_program;
+  partial_program.set_program(std::move(serialized_module));
+  partial_program.set_program_format(kStablehloBytecodeFormat);
+  partial_program.set_producer_phase("n/a");
+  partial_program.add_consumer_phases({std::string(kPhaseName)});
+  partial_program.set_version("1.0");
+  partial_program.set_program_name(name);
+
+  ABSL_ASSIGN_OR_RETURN(std::vector<xla::PjRtPartialProgramProto> programs,
+                   RunPhases(std::move(options), {std::move(partial_program)},
+                             topology, {std::string(kPhaseName)}));
+
+  return SampleRelocatable::Create(std::move(programs));
+}
+
+absl::StatusOr<std::unique_ptr<xla::PjRtRelocatable>>
+SamplePhaseCompiler::DeserializeRelocatable(
+    absl::string_view serialized) const {
+  return SampleRelocatable::Deserialize(serialized);
+}
+
 PJRT_Error* PJRT_PhaseCompile_Get_Compiler(
     PJRT_PhaseCompile_Get_Compiler_Args* args) {
   PJRT_RETURN_IF_ERROR(ActualStructSizeIsGreaterOrEqual(
diff --git a/third_party/xla/xla/pjrt/pjrt_phase_compile_sample_plugin.h b/third_party/xla/xla/pjrt/pjrt_phase_compile_sample_plugin.h
index ef25add9b78..03d6198c6d7 100644
--- a/third_party/xla/xla/pjrt/pjrt_phase_compile_sample_plugin.h
+++ b/third_party/xla/xla/pjrt/pjrt_phase_compile_sample_plugin.h
@@ -16,8 +16,11 @@ limitations under the License.
 #ifndef XLA_PJRT_PJRT_PHASE_COMPILE_SAMPLE_PLUGIN_H_
 #define XLA_PJRT_PJRT_PHASE_COMPILE_SAMPLE_PLUGIN_H_
 
+#include <cstdint>
 #include <memory>
 #include <string>
+#include <utility>
+#include <vector>
 
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
@@ -31,6 +34,12 @@ limitations under the License.
 #include "xla/pjrt/maybe_owning_mlir_module.h"
 #include "xla/pjrt/pjrt_compiler.h"
 #include "xla/pjrt/pjrt_executable.h"
+#include "xla/pjrt/pjrt_layout.h"
+#include "xla/pjrt/pjrt_relocatable.h"
+#include "xla/pjrt/proto/pjrt_partial_program.pb.h"
+#include "xla/shape.h"
+#include "xla/util.h"
+#include "xla/xla_data.pb.h"
 
 namespace pjrt {
 namespace phase_compile_sample_plugin {
@@ -53,6 +62,7 @@ class StablehloTypeSerialization {
 
   // Serializes an MLIR ModuleOp into a StableHLO bytecode string.
   // Returns an error if serialization fails.
+  static absl::StatusOr<std::string> Serialize(mlir::ModuleOp module_op);
   static absl::StatusOr<std::string> Serialize(
       const mlir::OwningOpRef<mlir::ModuleOp>& module_op);
 
@@ -65,6 +75,72 @@ class StablehloTypeSerialization {
 // phase in the `PjRtPartialProgramProto` objects.
 constexpr absl::string_view kPhaseName = "stablehlo_to_optimized_stablehlo";
 
+// Sample implementation of `PjRtRelocatable` that contains a mock-optimized
+// StableHLO program inside its set of partial programs, along with any
+// dependent StableHLO programs.
+class SampleRelocatable : public xla::PjRtRelocatable {
+ public:
+  SampleRelocatable(
+      absl::string_view name, absl::string_view build_id,
+      xla::ProgramShape shape, absl::string_view call_backend_config,
+      std::vector<std::pair<int64_t, int64_t>> param_output_aliases,
+      bool has_side_effects,
+      std::vector<xla::PjRtPartialProgramProto> partial_programs)
+      : name_(name),
+        build_id_(build_id),
+        shape_(std::move(shape)),
+        call_backend_config_(call_backend_config),
+        param_output_aliases_(std::move(param_output_aliases)),
+        has_side_effects_(has_side_effects),
+        partial_programs_(std::move(partial_programs)) {}
+
+  ~SampleRelocatable() override = default;
+
+  static absl::StatusOr<std::unique_ptr<SampleRelocatable>> Create(
+      std::vector<xla::PjRtPartialProgramProto> partial_programs);
+
+  absl::string_view name() const override { return name_; }
+  absl::string_view build_id() const override { return build_id_; }
+  absl::string_view call_backend_config() const override {
+    return call_backend_config_;
+  }
+  bool has_side_effects() const override { return has_side_effects_; }
+  const xla::ProgramShape& shape() const { return shape_; }
+  const std::vector<xla::PjRtPartialProgramProto>& partial_programs() const {
+    return partial_programs_;
+  }
+
+  absl::StatusOr<std::string> Serialize() const override;
+  static absl::StatusOr<std::unique_ptr<SampleRelocatable>> Deserialize(
+      absl::string_view serialized);
+
+  absl::StatusOr<std::vector<xla::PrimitiveType>> GetParameterElementTypes()
+      const override;
+  absl::StatusOr<std::vector<xla::PrimitiveType>> GetOutputElementTypes()
+      const override;
+  absl::StatusOr<std::vector<xla::DimensionVector>> GetParameterDimensions()
+      const override;
+  absl::StatusOr<std::vector<xla::DimensionVector>> GetOutputDimensions()
+      const override;
+  absl::StatusOr<std::vector<std::shared_ptr<const xla::PjRtLayout>>>
+  GetParameterLayouts() const override;
+  absl::StatusOr<std::vector<std::shared_ptr<const xla::PjRtLayout>>>
+  GetOutputLayouts() const override;
+  absl::StatusOr<std::vector<std::pair<int64_t, int64_t>>>
+  GetParameterOutputAliases() const override {
+    return param_output_aliases_;
+  }
+
+ private:
+  std::string name_;
+  std::string build_id_;
+  xla::ProgramShape shape_;
+  std::string call_backend_config_;
+  std::vector<std::pair<int64_t, int64_t>> param_output_aliases_;
+  bool has_side_effects_;
+  std::vector<xla::PjRtPartialProgramProto> partial_programs_;
+};
+
 // This class demonstrates an example phase compiler that the plugin developer
 // needs to implement.
 class SamplePhaseCompiler : public xla::PjRtPhaseCompiler {
@@ -80,6 +156,14 @@ class SamplePhaseCompiler : public xla::PjRtPhaseCompiler {
       xla::CompileOptions options, xla::MaybeOwningMlirModule module,
       const xla::PjRtTopologyDescription& topology,
       xla::PjRtClient* client) override;
+
+  absl::StatusOr<std::unique_ptr<xla::PjRtRelocatable>> CompileToRelocatable(
+      xla::CompileOptions options, xla::MaybeOwningMlirModule module,
+      const xla::PjRtTopologyDescription& topology, absl::string_view name,
+      bool is_entrypoint, xla::PjRtClient* client) override;
+
+  absl::StatusOr<std::unique_ptr<xla::PjRtRelocatable>> DeserializeRelocatable(
+      absl::string_view serialized) const override;
 };
 
 // Creates a phase compile extension for the sample plugin.
diff --git a/third_party/xla/xla/pjrt/pjrt_phase_compiler_test.cc b/third_party/xla/xla/pjrt/pjrt_phase_compiler_test.cc
index 91883b28118..6902adfab3c 100644
--- a/third_party/xla/xla/pjrt/pjrt_phase_compiler_test.cc
+++ b/third_party/xla/xla/pjrt/pjrt_phase_compiler_test.cc
@@ -22,6 +22,7 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/base/casts.h"
 #include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
@@ -39,6 +40,7 @@ limitations under the License.
 #include "xla/pjrt/pjrt_compiler.h"
 #include "xla/pjrt/pjrt_executable.h"
 #include "xla/pjrt/pjrt_phase_compile_sample_plugin.h"
+#include "xla/pjrt/pjrt_relocatable.h"
 #include "xla/pjrt/plugin/xla_cpu/cpu_topology.h"
 #include "xla/pjrt/plugin/xla_cpu/cpu_topology_description.h"
 #include "xla/pjrt/proto/pjrt_partial_program.pb.h"
@@ -50,6 +52,7 @@ namespace {
 
 using ::testing::ElementsAre;
 using ::testing::HasSubstr;
+using ::testing::SizeIs;
 using ::tsl::proto_testing::EqualsProto;
 
 constexpr absl::string_view kStablehloModuleStr = R"(
@@ -105,6 +108,17 @@ class SamplePhaseCompilerTest : public ::testing::Test {
             std::vector<xla::CpuTopology::CpuDevice>(),
             xla::cpu::TargetMachineOptions(xla::GetDebugOptionsFromFlags())));
   }
+
+  xla::MaybeOwningMlirModule ParseStablehloModule(
+      absl::string_view module_str = kStablehloModuleStr) {
+    auto context = std::make_shared<mlir::MLIRContext>();
+    mlir::FailureOr<mlir::OwningOpRef<mlir::ModuleOp>> stablehlo_module =
+        mlir::stablehlo::parseStablehloModule(std::string(module_str),
+                                              *context);
+    CHECK(mlir::succeeded(stablehlo_module));
+    return xla::MaybeOwningMlirModule(std::move(context),
+                                      std::move(*stablehlo_module));
+  }
 };
 
 // Test that the sample phase compiler's RegisterAllPhases method fails when
@@ -257,5 +271,92 @@ TEST_F(SamplePhaseCompilerTest, TestSamplePhaseCompilerGetPhaseNames) {
                   ElementsAre(phase_compile_sample_plugin::kPhaseName)));
 }
 
+TEST_F(SamplePhaseCompilerTest, TestSamplePhaseCompilerCompileToRelocatable) {
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<xla::PjRtRelocatable> relocatable,
+      phase_compiler_->CompileToRelocatable(
+          xla::CompileOptions(), ParseStablehloModule(), *topology_description_,
+          /*name=*/"my_relocatable",
+          /*is_entrypoint=*/false, /*client=*/nullptr));
+
+  EXPECT_EQ(relocatable->name(), "my_relocatable");
+  const auto* sample_relocatable =
+      absl::down_cast<const phase_compile_sample_plugin::SampleRelocatable*>(
+          relocatable.get());
+  ASSERT_THAT(sample_relocatable->partial_programs(), SizeIs(1));
+  EXPECT_EQ(sample_relocatable->partial_programs()[0].program_name(),
+            "my_relocatable");
+  EXPECT_EQ(sample_relocatable->partial_programs()[0].producer_phase(),
+            phase_compile_sample_plugin::kPhaseName);
+
+  mlir::MLIRContext context;
+  EXPECT_THAT(
+      phase_compile_sample_plugin::StablehloTypeSerialization::Deserialize(
+          sample_relocatable->partial_programs()[0].program(), context),
+      absl_testing::IsOk());
+}
+
+TEST_F(SamplePhaseCompilerTest,
+       TestSamplePhaseCompilerCompileToRelocatableWithNullModule) {
+  mlir::ModuleOp null_module;
+  EXPECT_THAT(
+      phase_compiler_->CompileToRelocatable(
+          xla::CompileOptions(), xla::MaybeOwningMlirModule(null_module),
+          *topology_description_, /*name=*/"bad",
+          /*is_entrypoint=*/false, /*client=*/nullptr),
+      absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+TEST_F(SamplePhaseCompilerTest,
+       TestSamplePhaseCompilerSerializeAndDeserializeRelocatable) {
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<xla::PjRtRelocatable> relocatable,
+      phase_compiler_->CompileToRelocatable(
+          xla::CompileOptions(), ParseStablehloModule(), *topology_description_,
+          /*name=*/"base",
+          /*is_entrypoint=*/false, /*client=*/nullptr));
+
+  ASSERT_OK_AND_ASSIGN(std::string serialized, relocatable->Serialize());
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<xla::PjRtRelocatable> deserialized,
+                       phase_compiler_->DeserializeRelocatable(serialized));
+
+  EXPECT_EQ(deserialized->name(), relocatable->name());
+  EXPECT_EQ(deserialized->build_id(), relocatable->build_id());
+  EXPECT_EQ(deserialized->call_backend_config(),
+            relocatable->call_backend_config());
+
+  const auto* sample_relocatable =
+      absl::down_cast<const phase_compile_sample_plugin::SampleRelocatable*>(
+          relocatable.get());
+  const auto* sample_deserialized =
+      absl::down_cast<const phase_compile_sample_plugin::SampleRelocatable*>(
+          deserialized.get());
+  ASSERT_EQ(sample_deserialized->partial_programs().size(),
+            sample_relocatable->partial_programs().size());
+  for (size_t i = 0; i < sample_relocatable->partial_programs().size(); ++i) {
+    EXPECT_THAT(sample_deserialized->partial_programs()[i],
+                EqualsProto(sample_relocatable->partial_programs()[i]));
+  }
+}
+
+TEST_F(SamplePhaseCompilerTest,
+       TestSamplePhaseCompilerDeserializeRelocatableWithInvalidInput) {
+  EXPECT_THAT(phase_compiler_->DeserializeRelocatable(""),
+              absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
+  EXPECT_THAT(phase_compiler_->DeserializeRelocatable("not_a_valid_proto"),
+              absl_testing::StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+TEST_F(SamplePhaseCompilerTest,
+       TestSamplePhaseCompilerLinkRelocatablesUnimplemented) {
+  // We don't have a sample for LinkRelocatables since we don't have an easy
+  // way to create a sample PjRtExecutable that actually works.
+  EXPECT_THAT(phase_compiler_->LinkRelocatables(xla::CompileOptions(),
+                                                *topology_description_,
+                                                /*entrypoint=*/nullptr,
+                                                /*dependencies=*/{}),
+              absl_testing::StatusIs(absl::StatusCode::kUnimplemented));
+}
+
 }  // namespace
 }  // namespace pjrt
diff --git a/third_party/xla/xla/pjrt/pjrt_relocatable_test.cc b/third_party/xla/xla/pjrt/pjrt_relocatable_test.cc
index cf9cd3554cf..77da7ab6ab2 100644
--- a/third_party/xla/xla/pjrt/pjrt_relocatable_test.cc
+++ b/third_party/xla/xla/pjrt/pjrt_relocatable_test.cc
@@ -12,15 +12,234 @@ 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.
 ==============================================================================*/
-#include "xla/pjrt/pjrt_relocatable.h"  // IWYU pragma: keep
+#include "xla/pjrt/pjrt_relocatable.h"
 
-#include <gmock/gmock.h>  // IWYU pragma: keep
-#include <gtest/gtest.h>  // IWYU pragma: keep
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/log/check.h"
+#include "absl/status/status_matchers.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OwningOpRef.h"
+#include "mlir/Support/LLVM.h"
+#include "stablehlo/reference/Api.h"
+#include "xla/backends/cpu/target_machine_options.h"
+#include "xla/debug_options_flags.h"
+#include "xla/layout_util.h"
+#include "xla/pjrt/maybe_owning_mlir_module.h"
+#include "xla/pjrt/pjrt_compiler.h"
+#include "xla/pjrt/pjrt_layout.h"
+#include "xla/pjrt/pjrt_phase_compile_sample_plugin.h"
+#include "xla/pjrt/plugin/xla_cpu/cpu_topology.h"
+#include "xla/pjrt/plugin/xla_cpu/cpu_topology_description.h"
+#include "xla/util.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla {
 namespace {
 
-TEST(PjRtRelocatableTest, Empty) {}
+using ::testing::ElementsAre;
+using ::testing::IsEmpty;
+using ::testing::Not;
+using ::testing::Pair;
+using ::testing::SizeIs;
+
+class PjRtRelocatableTest : public ::testing::Test {
+ protected:
+  PjRtRelocatableTest() {
+    phase_compiler_ = std::make_unique<
+        pjrt::phase_compile_sample_plugin::SamplePhaseCompiler>();
+    CHECK_OK(phase_compiler_->RegisterAllPhases());
+
+    topology_description_ = std::make_unique<CpuTopologyDescription>(
+        CpuId(), CpuName(), "<unknown>",
+        CpuTopology(std::vector<CpuTopology::CpuDevice>(),
+                    cpu::TargetMachineOptions(GetDebugOptionsFromFlags())));
+  }
+
+  absl::StatusOr<std::unique_ptr<PjRtRelocatable>> CompileToRelocatable(
+      absl::string_view module_str, absl::string_view name = "test_relocatable",
+      bool is_entrypoint = false) {
+    auto context = std::make_shared<mlir::MLIRContext>();
+    mlir::FailureOr<mlir::OwningOpRef<mlir::ModuleOp>> stablehlo_module =
+        mlir::stablehlo::parseStablehloModule(std::string(module_str),
+                                              *context);
+    CHECK(mlir::succeeded(stablehlo_module));
+    return phase_compiler_->CompileToRelocatable(
+        CompileOptions(),
+        MaybeOwningMlirModule(std::move(context), std::move(*stablehlo_module)),
+        *topology_description_, name, is_entrypoint,
+        /*client=*/nullptr);
+  }
+
+  std::unique_ptr<pjrt::phase_compile_sample_plugin::SamplePhaseCompiler>
+      phase_compiler_;
+  std::unique_ptr<PjRtTopologyDescription> topology_description_;
+};
+
+TEST_F(PjRtRelocatableTest, BasicMetadataAndSingleOutput) {
+  constexpr absl::string_view kModuleStr = R"(
+    module {
+      func.func @main(%arg0: tensor<4x8xf32>, %arg1: tensor<2xi32>) -> tensor<4x8xf32> {
+        %0 = stablehlo.constant dense<0.0> : tensor<4x8xf32>
+        %1 = stablehlo.add %arg0, %0 : tensor<4x8xf32>
+        func.return %1 : tensor<4x8xf32>
+      }
+    }
+  )";
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtRelocatable> relocatable,
+                       CompileToRelocatable(kModuleStr, "add_op"));
+
+  EXPECT_EQ(relocatable->name(), "add_op");
+  EXPECT_THAT(relocatable->build_id(), Not(IsEmpty()));
+  EXPECT_THAT(relocatable->call_backend_config(), Not(IsEmpty()));
+  EXPECT_FALSE(relocatable->has_side_effects());
+
+  EXPECT_THAT(relocatable->GetParameterElementTypes(),
+              absl_testing::IsOkAndHolds(ElementsAre(F32, S32)));
+  EXPECT_THAT(relocatable->GetOutputElementTypes(),
+              absl_testing::IsOkAndHolds(ElementsAre(F32)));
+
+  EXPECT_THAT(relocatable->GetParameterDimensions(),
+              absl_testing::IsOkAndHolds(
+                  ElementsAre(DimensionVector{4, 8}, DimensionVector{2})));
+  EXPECT_THAT(relocatable->GetOutputDimensions(),
+              absl_testing::IsOkAndHolds(ElementsAre(DimensionVector{4, 8})));
+
+  ASSERT_OK_AND_ASSIGN(
+      std::vector<std::shared_ptr<const PjRtLayout>> param_layouts,
+      relocatable->GetParameterLayouts());
+  ASSERT_THAT(param_layouts, SizeIs(2));
+  EXPECT_EQ(*param_layouts[0], PjRtLayout(LayoutUtil::MakeDescendingLayout(2)));
+  EXPECT_EQ(*param_layouts[1], PjRtLayout(LayoutUtil::MakeDescendingLayout(1)));
+
+  ASSERT_OK_AND_ASSIGN(
+      std::vector<std::shared_ptr<const PjRtLayout>> output_layouts,
+      relocatable->GetOutputLayouts());
+  ASSERT_THAT(output_layouts, SizeIs(1));
+  EXPECT_EQ(*output_layouts[0],
+            PjRtLayout(LayoutUtil::MakeDescendingLayout(2)));
+
+  EXPECT_THAT(relocatable->GetParameterOutputAliases(),
+              absl_testing::IsOkAndHolds(IsEmpty()));
+}
+
+TEST_F(PjRtRelocatableTest, MultipleOutputsAndAliases) {
+  constexpr absl::string_view kModuleStr = R"(
+    module {
+      func.func @main(%arg0: tensor<4xf32> {tf.aliasing_output = 1 : i32},
+                      %arg1: tensor<2x3xi64>) -> (tensor<2x3xi64>, tensor<4xf32>) {
+        func.return %arg1, %arg0 : tensor<2x3xi64>, tensor<4xf32>
+      }
+    }
+  )";
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtRelocatable> relocatable,
+                       CompileToRelocatable(kModuleStr, "swap_op"));
+
+  EXPECT_THAT(relocatable->GetParameterElementTypes(),
+              absl_testing::IsOkAndHolds(ElementsAre(F32, S64)));
+  EXPECT_THAT(relocatable->GetOutputElementTypes(),
+              absl_testing::IsOkAndHolds(ElementsAre(S64, F32)));
+
+  EXPECT_THAT(relocatable->GetParameterDimensions(),
+              absl_testing::IsOkAndHolds(
+                  ElementsAre(DimensionVector{4}, DimensionVector{2, 3})));
+  EXPECT_THAT(relocatable->GetOutputDimensions(),
+              absl_testing::IsOkAndHolds(
+                  ElementsAre(DimensionVector{2, 3}, DimensionVector{4})));
+
+  ASSERT_OK_AND_ASSIGN(
+      std::vector<std::shared_ptr<const PjRtLayout>> output_layouts,
+      relocatable->GetOutputLayouts());
+  ASSERT_THAT(output_layouts, SizeIs(2));
+  EXPECT_EQ(*output_layouts[0],
+            PjRtLayout(LayoutUtil::MakeDescendingLayout(2)));
+  EXPECT_EQ(*output_layouts[1],
+            PjRtLayout(LayoutUtil::MakeDescendingLayout(1)));
+
+  EXPECT_THAT(relocatable->GetParameterOutputAliases(),
+              absl_testing::IsOkAndHolds(ElementsAre(Pair(0, 1))));
+}
+
+TEST_F(PjRtRelocatableTest, HasSideEffects) {
+  constexpr absl::string_view kSideEffectModuleStr = R"(
+    module {
+      func.func @main(%arg0: tensor<4xi32>) -> tensor<4xi32> {
+        %0 = stablehlo.custom_call @side_effecting_op(%arg0) {has_side_effect = true} : (tensor<4xi32>) -> tensor<4xi32>
+        func.return %0 : tensor<4xi32>
+      }
+    }
+  )";
+
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<PjRtRelocatable> relocatable,
+      CompileToRelocatable(kSideEffectModuleStr, "side_effect_op"));
+  EXPECT_TRUE(relocatable->has_side_effects());
+}
+
+TEST_F(PjRtRelocatableTest, SerializeAndDeserializeRoundTrip) {
+  constexpr absl::string_view kModuleStr = R"(
+    module {
+      func.func @main(%arg0: tensor<4x2xf32> {tf.aliasing_output = 0 : i32}) -> tensor<4x2xf32> {
+        %0 = stablehlo.custom_call @side_effecting_op(%arg0) {has_side_effect = true} : (tensor<4x2xf32>) -> tensor<4x2xf32>
+        func.return %0 : tensor<4x2xf32>
+      }
+    }
+  )";
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtRelocatable> original,
+                       CompileToRelocatable(kModuleStr, "round_trip_op"));
+  ASSERT_OK_AND_ASSIGN(std::string serialized, original->Serialize());
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtRelocatable> deserialized,
+                       phase_compiler_->DeserializeRelocatable(serialized));
+
+  EXPECT_EQ(deserialized->name(), original->name());
+  EXPECT_EQ(deserialized->build_id(), original->build_id());
+  EXPECT_EQ(deserialized->call_backend_config(),
+            original->call_backend_config());
+  EXPECT_EQ(deserialized->has_side_effects(), original->has_side_effects());
+  EXPECT_EQ(deserialized->GetParameterElementTypes(),
+            original->GetParameterElementTypes());
+  EXPECT_EQ(deserialized->GetOutputElementTypes(),
+            original->GetOutputElementTypes());
+  EXPECT_EQ(deserialized->GetParameterDimensions(),
+            original->GetParameterDimensions());
+  EXPECT_EQ(deserialized->GetOutputDimensions(),
+            original->GetOutputDimensions());
+  EXPECT_EQ(deserialized->GetParameterOutputAliases(),
+            original->GetParameterOutputAliases());
+
+  ASSERT_OK_AND_ASSIGN(
+      std::vector<std::shared_ptr<const PjRtLayout>> orig_param_layouts,
+      original->GetParameterLayouts());
+  ASSERT_OK_AND_ASSIGN(
+      std::vector<std::shared_ptr<const PjRtLayout>> deser_param_layouts,
+      deserialized->GetParameterLayouts());
+  ASSERT_EQ(deser_param_layouts.size(), orig_param_layouts.size());
+  for (size_t i = 0; i < orig_param_layouts.size(); ++i) {
+    EXPECT_EQ(*deser_param_layouts[i], *orig_param_layouts[i]);
+  }
+
+  ASSERT_OK_AND_ASSIGN(
+      std::vector<std::shared_ptr<const PjRtLayout>> orig_out_layouts,
+      original->GetOutputLayouts());
+  ASSERT_OK_AND_ASSIGN(
+      std::vector<std::shared_ptr<const PjRtLayout>> deser_out_layouts,
+      deserialized->GetOutputLayouts());
+  ASSERT_EQ(deser_out_layouts.size(), orig_out_layouts.size());
+  for (size_t i = 0; i < orig_out_layouts.size(); ++i) {
+    EXPECT_EQ(*deser_out_layouts[i], *orig_out_layouts[i]);
+  }
+}
 
 }  // namespace
 }  // namespace xla

From 2c686468751fa7a6a0945050f24436c041feb132 Mon Sep 17 00:00:00 2001
From: Michael Wong <michaelcw@google.com>
Date: Mon, 12 Oct 2026 17:52:04 -0700
Subject: [PATCH 0633/1000] Avoid using Status for control flow in
 HeapSimulator.

PiperOrigin-RevId: 998266399
---
 .../service/heap_simulator/heap_simulator.cc  | 85 +++++++++----------
 .../service/heap_simulator/heap_simulator.h   | 14 +--
 2 files changed, 42 insertions(+), 57 deletions(-)

diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
index 7d7b1015f44..89fe04a9dde 100644
--- a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
+++ b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
@@ -2380,53 +2380,46 @@ GlobalDecreasingSizeBestFitHeap<
 }
 
 template <typename BufferType>
-absl::Status GlobalDecreasingSizeBestFitHeap<BufferType>::
-    SlicedAllocationFinder::DoesPermutationFit(
-        absl::Span<const int64_t> permutation_of_slice_times,
-        const FreeChunkRoot& root, int64_t offset) const {
-  absl::Status result =
-      DoesPermutationFitImpl(permutation_of_slice_times, root, offset);
-  VLOG(3) << "SlicedAllocationFinder::DoesPermutationFit\n"
-          << "  permutation of slice times: [ "
-          << absl::StrJoin(permutation_of_slice_times, ",") << " ]\n"
-          << "  offset: " << offset << "\n"
-          << "  root: " << root.ToString() << "\n"
-          << "  -> " << result;
-  return result;
-}
+bool GlobalDecreasingSizeBestFitHeap<BufferType>::SlicedAllocationFinder::
+    DoesPermutationFit(absl::Span<const int64_t> permutation_of_slice_times,
+                       const FreeChunkRoot& root, int64_t offset) const {
+  auto vlog_prefix = [&]() {
+    return absl::StrCat("SlicedAllocationFinder::DoesPermutationFit\n",
+                        "  permutation of slice times: [ ",
+                        absl::StrJoin(permutation_of_slice_times, ","), " ]\n",
+                        "  offset: ", offset, "\n", "  root: ", root.ToString(),
+                        "\n  -> ");
+  };
 
-template <typename BufferType>
-absl::Status GlobalDecreasingSizeBestFitHeap<BufferType>::
-    SlicedAllocationFinder::DoesPermutationFitImpl(
-        absl::Span<const int64_t> permutation_of_slice_times,
-        const FreeChunkRoot& root, int64_t offset) const {
   if (permutation_of_slice_times.size() != sorted_slice_sizes_.size()) {
-    return InvalidArgumentStrCat(
-        sorted_slice_sizes_.size(), " slices times expected in permutation. ",
-        permutation_of_slice_times.size(), " specified.");
+    VLOG(3) << vlog_prefix() << sorted_slice_sizes_.size()
+            << " slices times expected in permutation. "
+            << permutation_of_slice_times.size() << " specified.";
+    return false;
   }
   if (offset >= root.chunk.chunk_end()) {
-    return FailedPrecondition(
-        "%s", absl::StrCat("Free chunk root ", root.chunk.ToString(),
-                           " does not overlap with offset ", offset, "."));
+    VLOG(3) << vlog_prefix() << "Free chunk root " << root.chunk.ToString()
+            << " does not overlap with offset " << offset << ".";
+    return false;
   }
   if (offset + max_colocation_size_ > root.chunk.chunk_end()) {
-    return FailedPrecondition(
-        "%s", absl::StrCat("Not enough space to fit enitre allocation [",
-                           offset, ", ", offset + max_colocation_size_,
-                           ") in free chunk root ", root.chunk.ToString()));
+    VLOG(3) << vlog_prefix() << "Not enough space to fit entire allocation ["
+            << offset << ", " << offset + max_colocation_size_
+            << ") in free chunk root " << root.chunk.ToString();
+    return false;
   }
   if (!is_offset_allowed_(offset)) {
-    return FailedPrecondition(
-        "%s", absl::StrCat("We are not permitted to place an allocation at ",
-                           "offset ", offset, "."));
+    VLOG(3) << vlog_prefix()
+            << "We are not permitted to place an allocation at offset "
+            << offset << ".";
+    return false;
   }
 
   auto piece_fwd_it = root.pieces.lower_bound(offset);
   if (piece_fwd_it == root.pieces.end()) {
-    return FailedPrecondition(
-        "%s", absl::StrCat("Offset ", offset, " comes before free chunk root ",
-                           root.chunk.ToString()));
+    VLOG(3) << vlog_prefix() << "Offset " << offset
+            << " comes before free chunk root " << root.chunk.ToString();
+    return false;
   }
   ++piece_fwd_it;
   auto piece_reverse_it = std::make_reverse_iterator(piece_fwd_it);
@@ -2455,11 +2448,10 @@ absl::Status GlobalDecreasingSizeBestFitHeap<BufferType>::
     if (current_piece_time > current_slice_time) {
       // The current piece is not free far enough back in time to support the
       // current slice.
-      return FailedPrecondition(
-          "%s",
-          absl::StrCat("At slice time t", current_slice_time, ", slice ",
-                       slice_index, " does not fit at offset ", current_offset,
-                       " in root ", root.chunk.ToString()));
+      VLOG(3) << vlog_prefix() << "At slice time t" << current_slice_time
+              << ", slice " << slice_index << " does not fit at offset "
+              << current_offset << " in root " << root.chunk.ToString();
+      return false;
     }
 
     if (remaining_in_slice >= remaining_in_piece) {
@@ -2475,12 +2467,14 @@ absl::Status GlobalDecreasingSizeBestFitHeap<BufferType>::
   }
 
   if (!out_of_slices(slice_index)) {
-    return InternalStrCat("Ran out of space in root ", root.chunk.ToString(),
-                          " to fit slice permutation; however, we should "
-                          "have caught such a condition earlier.");
+    VLOG(3) << vlog_prefix() << "Ran out of space in root "
+            << root.chunk.ToString()
+            << " to fit slice permutation; however, we should "
+               "have caught such a condition earlier.";
+    return false;
   }
 
-  return absl::OkStatus();
+  return true;
 }
 
 // Future opportunities:
@@ -2518,8 +2512,7 @@ GlobalDecreasingSizeBestFitHeap<BufferType>::SlicedAllocationFinder::FindInRoot(
          !slice_time_permutation_iterator_->Done();
          slice_time_permutation_iterator_->Next()) {
       if (DoesPermutationFit(slice_time_permutation_iterator_->Get(), root,
-                             offset)
-              .ok()) {
+                             offset)) {
         return PermutationToChunks(slice_time_permutation_iterator_->Get(),
                                    offset);
       }
diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator.h b/third_party/xla/xla/service/heap_simulator/heap_simulator.h
index 2d7771dd834..edacc1ca4e7 100644
--- a/third_party/xla/xla/service/heap_simulator/heap_simulator.h
+++ b/third_party/xla/xla/service/heap_simulator/heap_simulator.h
@@ -872,22 +872,14 @@ class GlobalDecreasingSizeBestFitHeap : public HeapAlgorithm<BufferType> {
     // The latest slice time for the specified sliced allocation request.
     int64_t LatestSliceTime() const { return sorted_slice_sizes_.size() - 1; }
 
-    // Returns ok if the given permutation of slice times results in an
-    // allocation of free space in root, at the specified offset. Otherwise,
-    // returns the reason such an allocation would not fit.
+    // Returns whether the given permutation of slice times results in an
+    // allocation of free space in root, at the specified offset.
     //
     // permutation_of_slice_times[i] is the slice time that the ith slice
     // (spatially) should be allocated. Such a slice has size
     // sorted_slice_sizes_[i] and would be allocated at offset +
     // sum(sorted_slice_sizes[j], for j in [0, i-1]).
-    absl::Status DoesPermutationFit(
-        absl::Span<const int64_t> permutation_of_slice_times,
-        const FreeChunkRoot& root, int64_t offset) const;
-
-    // Only DoesSlicedPermutationFit() should call this method directly. Other
-    // callers should call DoesSlicedPermutationFit(), which contains some
-    // wrapper VLOGGING.
-    absl::Status DoesPermutationFitImpl(
+    bool DoesPermutationFit(
         absl::Span<const int64_t> permutation_of_slice_times,
         const FreeChunkRoot& root, int64_t offset) const;
 

From 5ec4bb2fb4ebb0d983121c52827fdb0af326932b Mon Sep 17 00:00:00 2001
From: Dmitri Latushko <latushko@google.com>
Date: Mon, 12 Oct 2026 18:15:48 -0700
Subject: [PATCH 0634/1000] Shard assert_cardinality_test across 4 shards to
 prevent ASAN and TSAN timeouts

PiperOrigin-RevId: 998276271
---
 tensorflow/python/data/experimental/kernel_tests/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/python/data/experimental/kernel_tests/BUILD b/tensorflow/python/data/experimental/kernel_tests/BUILD
index 0dfdc5460a4..d3b2b3eca75 100644
--- a/tensorflow/python/data/experimental/kernel_tests/BUILD
+++ b/tensorflow/python/data/experimental/kernel_tests/BUILD
@@ -10,6 +10,7 @@ tf_py_strict_test(
     name = "assert_cardinality_test",
     size = "small",
     srcs = ["assert_cardinality_test.py"],
+    shard_count = 4,
     deps = [
         "//tensorflow/python/data/experimental/ops:cardinality",
         "//tensorflow/python/data/experimental/ops:global_shuffle_op",

From c122332f7644cd4fd092c753156f22c360951d7c Mon Sep 17 00:00:00 2001
From: Javid Khan <dxbjavid@gmail.com>
Date: Tue, 13 Oct 2026 11:16:36 +0530
Subject: [PATCH 0635/1000] widen ragged bincount loop indices to int64 and
 cover sparse test in graph mode

---
 tensorflow/core/kernels/bincount_op.cc                   | 9 +++++----
 .../python/kernel_tests/math_ops/bincount_op_test.py     | 1 +
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/kernels/bincount_op.cc b/tensorflow/core/kernels/bincount_op.cc
index e660801addc..7691ed75812 100644
--- a/tensorflow/core/kernels/bincount_op.cc
+++ b/tensorflow/core/kernels/bincount_op.cc
@@ -515,13 +515,14 @@ class RaggedBincountOp : public OpKernel {
                 absl::InvalidArgumentError(
                     absl::StrCat("size (", size, ") must be non-negative")));
 
-    int num_rows = splits.size() - 1;
-    int num_values = values.size();
-    int batch_idx = 0;
+    int64_t num_values = values.size();
+    int64_t batch_idx = 0;
 
     OP_REQUIRES(ctx, splits.size() > 0,
                 absl::InvalidArgumentError("Splits must be non-empty"));
 
+    int64_t num_rows = splits.size() - 1;
+
     OP_REQUIRES(ctx, splits(0) == 0,
                 absl::InvalidArgumentError(absl::StrCat(
                     "Splits must start with 0, not with ", splits(0))));
@@ -549,7 +550,7 @@ class RaggedBincountOp : public OpKernel {
     fill(ctx->eigen_device<Device>(), out_t->flat<T>());
     const auto out = out_t->matrix<T>();
 
-    for (int idx = 0; idx < num_values; ++idx) {
+    for (int64_t idx = 0; idx < num_values; ++idx) {
       while (idx >= splits(batch_idx)) {
         batch_idx++;
       }
diff --git a/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py b/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py
index 97f6ba9bc6b..b25ed3bb528 100644
--- a/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py
+++ b/tensorflow/python/kernel_tests/math_ops/bincount_op_test.py
@@ -780,6 +780,7 @@ class SparseBincountOpTest(test_util.TensorFlowTestCase,
           )
       )
 
+  @test_util.run_in_graph_and_eager_modes
   def test_sparse_bincount_invalid_weights_length(self):
     # Regression test: the 2-D SparseBincount path reads `weights(i)` with the
     # value index `i`, so a `weights` shorter than the values would be read out

From 38777b9dd5d44a3ae6788a72836b8a1bf72cbac6 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 12 Oct 2026 18:43:31 -0700
Subject: [PATCH 0636/1000] [XLA][BufferAssignment] Allow forcing allocations
 at fixed offsets

To accommodate code generated without relocations in mind, this change allows
forcing Buffer Assignment to place specific buffers at fixed offsets within
their allocations:
1. Introduces `BufferAssigner::Options::fixed_offset` to specify a required
   fixed byte offset for a buffer position.
2. Extends `BufferAssignment::NewAllocation` with an `offset` parameter,
   sizing the allocation to `offset + size` and placing the buffer at
   `offset`.
3. Updates `MaybeAssignBuffer` and `AssignSingleHloBuffer` to validate fixed
   offsets and bounds across aliased `HloValue`s, require matching fixed
   offsets when reusing an allocation, and prevent fixed-offset and non-fixed
   buffers from reusing each other's allocations.
4. Sorts fixed-offset buffers first (largest-first) and routes them to the
   default allocation manager in `AssignBuffersForGlobalComputations` (bypassing
   `FAST_MERGE` and `HeapSimulator` delay), and excludes fixed-offset buffers
   and allocations from cross-color temporary heap reuse in
   `ReuseCompatibleTempHeaps`.
5. Updates `CombineTempAllocations` to prioritize zero-fixed-offset temporary
   allocations at base offset 0 while keeping interfering zero-fixed-offset
   and non-zero fixed-offset allocations standalone.
6. Adds unit tests covering zero and non-zero fixed offsets, interfering and
   non-overlapping zero-offset buffers, conflicting fixed offsets, and
   non-fixed-offset allocation isolation in `buffer_assignment_test.cc`.

PiperOrigin-RevId: 998287616
---
 .../xla/xla/service/buffer_assignment.cc      | 223 +++++++++++-
 .../xla/xla/service/buffer_assignment.h       |  14 +-
 .../xla/xla/service/buffer_assignment_test.cc | 344 ++++++++++++++++++
 3 files changed, 569 insertions(+), 12 deletions(-)

diff --git a/third_party/xla/xla/service/buffer_assignment.cc b/third_party/xla/xla/service/buffer_assignment.cc
index 4d460e4bdf9..dbd46c8dd91 100644
--- a/third_party/xla/xla/service/buffer_assignment.cc
+++ b/third_party/xla/xla/service/buffer_assignment.cc
@@ -1117,10 +1117,10 @@ BufferAllocation* BufferAssignment::NewEmptyAllocation(
 }
 
 absl::StatusOr<BufferAllocation*> BufferAssignment::NewAllocation(
-    const HloBuffer& buffer, int64_t size) {
+    const HloBuffer& buffer, int64_t size, int64_t offset) {
   ABSL_ASSIGN_OR_RETURN(auto color, buffer.color());
-  BufferAllocation* allocation = NewEmptyAllocation(size, color);
-  ABSL_RETURN_IF_ERROR(AddAssignment(allocation, buffer, /*offset=*/0, size));
+  BufferAllocation* allocation = NewEmptyAllocation(offset + size, color);
+  ABSL_RETURN_IF_ERROR(AddAssignment(allocation, buffer, offset, size));
   allocation->peak_buffers_.push_back(buffer.values()[0]);
   return allocation;
 }
@@ -1165,7 +1165,10 @@ absl::Status BufferAssignment::AddAssignment(BufferAllocation* allocation,
 // Combines allocations of temporary buffers of the same color into one big
 // BufferAllocation.
 absl::Status BufferAssignment::CombineTempAllocations(
-    std::optional<BufferValue::Color> temp_buffer_color) {
+    std::optional<BufferValue::Color> temp_buffer_color,
+    std::optional<std::function<std::optional<int64_t>(
+        const HloAliasAnalysis&, const HloInstruction*, const ShapeIndex&)>>
+        fixed_offset) {
   VLOG(1) << "CombineTempAllocations()";
 
   // Move all temp allocations into a single run at the end of the allocations
@@ -1179,6 +1182,36 @@ absl::Status BufferAssignment::CombineTempAllocations(
   std::vector<BufferAllocation> combined_allocations;
   // Reserve the max possible size we'd need to ensure pointer stability.
   combined_allocations.reserve(allocations_.end() - first_temp_it);
+  // Prioritize temp allocations requiring a fixed offset of zero so that they
+  // are placed first at offset 0 within the combined allocation, allowing other
+  // temp allocations of the same color to be placed on top.
+  auto has_zero_fixed_offset_fn = [&](const BufferAllocation& alloc) {
+    if (!fixed_offset.has_value()) {
+      return false;
+    }
+    // NOLINTNEXTLINE(*-custom-deterministic-iteration-order)
+    for (const auto& buffer_offset_size : alloc.assigned_buffers()) {
+      const HloValue* val = buffer_offset_size.first;
+      std::optional<int64_t> req_offset =
+          (*fixed_offset)(alias_analysis(), val->instruction(), val->index());
+      if (req_offset.has_value() && *req_offset == 0 &&
+          buffer_offset_size.second.offset == 0) {
+        return true;
+      }
+    }
+    return false;
+  };
+
+  std::stable_sort(first_temp_it, allocations_.end(),
+                   [&](const BufferAllocation& a, const BufferAllocation& b) {
+                     bool a_zero = has_zero_fixed_offset_fn(a);
+                     bool b_zero = has_zero_fixed_offset_fn(b);
+                     if (a_zero != b_zero) {
+                       return a_zero > b_zero;
+                     }
+                     return false;
+                   });
+
   // Holds the pointer to a combined allocation of each (color, page), if any.
   absl::flat_hash_map<std::pair<BufferValue::Color, int64_t>, BufferAllocation*>
       combined_allocation_map;
@@ -1189,6 +1222,28 @@ absl::Status BufferAssignment::CombineTempAllocations(
     BufferAllocation& temp_allocation = *it;
     BufferValue::Color color = temp_allocation.color();
     int64_t page = temp_allocation.page_id();
+    bool has_fixed_offset = false;
+    bool all_fixed_offsets_are_zero = true;
+    if (fixed_offset.has_value()) {
+      // NOLINTNEXTLINE(*-custom-deterministic-iteration-order)
+      for (const auto& [val, slice] : temp_allocation.assigned_buffers()) {
+        std::optional<int64_t> offset =
+            (*fixed_offset)(alias_analysis(), val->instruction(), val->index());
+        if (offset.has_value()) {
+          has_fixed_offset = true;
+          if (*offset != 0 || slice.offset != 0) {
+            all_fixed_offsets_are_zero = false;
+          }
+        }
+      }
+    }
+    // Allocations with non-zero fixed offset requirements cannot be safely
+    // merged into a combined allocation at base 0.
+    if (has_fixed_offset && !all_fixed_offsets_are_zero) {
+      combined_allocations.push_back(std::move(temp_allocation));
+      continue;
+    }
+
     auto combined_it = combined_allocation_map.find({color, page});
     if (combined_it == combined_allocation_map.end()) {
       // We have found the first temp allocation of this color. Collect
@@ -1201,6 +1256,12 @@ absl::Status BufferAssignment::CombineTempAllocations(
                                       &combined_allocations.back());
       continue;
     }
+    // Any subsequent fixed-offset-0 allocation interfered with the first one
+    // during MaybeAssignBuffer and cannot be shifted to a non-zero base offset.
+    if (has_fixed_offset) {
+      combined_allocations.push_back(std::move(temp_allocation));
+      continue;
+    }
     if (combined_it->second->size() + temp_allocation.size() >=
         multiheap_size_constraint_per_heap_) {
       // We cannot put more into the current combined_it. So, appoint a new
@@ -2171,6 +2232,69 @@ absl::StatusOr<bool> BufferAssigner::MaybeAssignBuffer(
     }
   }
 
+  int64_t target_offset = 0;
+  if (opts_.fixed_offset.has_value()) {
+    std::optional<int64_t> req_offset;
+    for (const HloValue* val : hlo_buffer.values()) {
+      std::optional<int64_t> offset = (*opts_.fixed_offset)(
+          assignment->alias_analysis(), val->instruction(), val->index());
+      if (offset.has_value()) {
+        if (req_offset.has_value() && *req_offset != *offset) {
+          return Internal(
+              "Conflicting fixed offset requirements for buffer %s: %d vs %d",
+              hlo_buffer.ToString(), *req_offset, *offset);
+        }
+        req_offset = offset;
+      }
+    }
+    if (req_offset.has_value()) {
+      target_offset = *req_offset;
+      if (target_offset < 0 ||
+          target_offset + assignment->HloBufferSize(hlo_buffer) >
+              allocation->size()) {
+        return false;
+      }
+      // NOLINTNEXTLINE(*-custom-deterministic-iteration-order)
+      for (const auto& buffer_offset_size : allocation->assigned_buffers()) {
+        const HloValue* val = buffer_offset_size.first;
+        const HloBuffer& exist_buffer =
+            assignment->alias_analysis().GetBufferContainingValue(*val);
+        bool exist_has_fixed_offset = absl::c_any_of(
+            exist_buffer.values(), [&](const HloValue* exist_val) {
+              return (*opts_.fixed_offset)(assignment->alias_analysis(),
+                                           exist_val->instruction(),
+                                           exist_val->index())
+                  .has_value();
+            });
+        if (!exist_has_fixed_offset) {
+          VLOG(4) << "Can't assign: allocation contains non-fixed offset "
+                  << "buffer, fixed offset buffer " << hlo_buffer
+                  << " cannot reuse it.";
+          return false;
+        }
+        int64_t exist_offset = buffer_offset_size.second.offset;
+        if (exist_offset != target_offset) {
+          VLOG(4) << "Can't assign: allocation has buffer at offset "
+                  << exist_offset << ", but required fixed offset is "
+                  << target_offset;
+          return false;
+        }
+      }
+    } else {
+      // NOLINTNEXTLINE(*-custom-deterministic-iteration-order)
+      for (const auto& buffer_offset_size : allocation->assigned_buffers()) {
+        const HloValue* val = buffer_offset_size.first;
+        if ((*opts_.fixed_offset)(assignment->alias_analysis(),
+                                  val->instruction(), val->index())
+                .has_value()) {
+          VLOG(4) << "Can't assign: allocation contains fixed offset buffer, "
+                  << "non-fixed buffer " << hlo_buffer << " cannot reuse it.";
+          return false;
+        }
+      }
+    }
+  }
+
   // If the buffer is live out of the computation then it should only be
   // assigned a buffer which exactly fits the result to avoid wasting memory
   // (result buffers can have arbitrary lifetimes).
@@ -2182,7 +2306,7 @@ absl::StatusOr<bool> BufferAssigner::MaybeAssignBuffer(
   }
 
   ABSL_RETURN_IF_ERROR(
-      assignment->AddAssignment(allocation, hlo_buffer, /*offset=*/0,
+      assignment->AddAssignment(allocation, hlo_buffer, target_offset,
                                 assignment->HloBufferSize(hlo_buffer)));
   return true;
 }
@@ -2298,6 +2422,20 @@ bool BufferAssigner::DelayTemporaryBufferAssignment(
       })) {
     return false;
   }
+
+  // Buffers requiring a fixed offset must not be delayed to heap simulation
+  // because the heap simulator packs temporary buffers into dynamic offsets
+  // based on live-range intervals without honoring fixed offset requirements.
+  if (opts_.fixed_offset.has_value()) {
+    for (const HloValue* hlo_value : hlo_buffer->values()) {
+      if ((*opts_.fixed_offset)(assignment->alias_analysis(),
+                                hlo_value->instruction(), hlo_value->index())
+              .has_value()) {
+        return false;
+      }
+    }
+  }
+
   bool all_computations_have_sequential_order = true;
   for (const HloValue* hlo_value : hlo_buffer->values()) {
     HloComputation* computation = hlo_value->instruction()->parent();
@@ -2371,8 +2509,28 @@ absl::Status BufferAssigner::AssignSingleHloBuffer(
   }
 
   if (!assignment->HasAllocation(*hlo_buffer)) {
-    ABSL_ASSIGN_OR_RETURN(BufferAllocation * allocation,
-                     assignment->NewAllocation(*hlo_buffer, buffer_size));
+    int64_t target_offset = 0;
+    if (opts_.fixed_offset.has_value()) {
+      std::optional<int64_t> req_offset;
+      for (const HloValue* val : hlo_buffer->values()) {
+        std::optional<int64_t> offset = (*opts_.fixed_offset)(
+            assignment->alias_analysis(), val->instruction(), val->index());
+        if (offset.has_value()) {
+          if (req_offset.has_value() && *req_offset != *offset) {
+            return Internal(
+                "Conflicting fixed offset requirements for buffer %s: %d vs %d",
+                hlo_buffer->ToString(), *req_offset, *offset);
+          }
+          req_offset = offset;
+        }
+      }
+      if (req_offset.has_value()) {
+        target_offset = *req_offset;
+      }
+    }
+    ABSL_ASSIGN_OR_RETURN(
+        BufferAllocation * allocation,
+        assignment->NewAllocation(*hlo_buffer, buffer_size, target_offset));
     allocation_manager->RegisterNewAllocation(hlo_buffer, allocation->index());
     VLOG(3) << "New allocation #" << allocation->index()
             << " for: " << *hlo_buffer;
@@ -2507,7 +2665,32 @@ absl::Status BufferAssigner::AssignBuffersForGlobalComputations(
     buffer_order = BufferOrder::kLiveRangeStart;
   }
 
+  auto has_fixed_offset = [&](const HloBuffer* buffer) {
+    if (!opts_.fixed_offset.has_value()) {
+      return false;
+    }
+    return absl::c_any_of(buffer->values(), [&](const HloValue* val) {
+      return (*opts_.fixed_offset)(alias_analysis, val->instruction(),
+                                   val->index())
+          .has_value();
+    });
+  };
+
   std::vector<Comparator> comparators;
+  if (opts_.fixed_offset.has_value()) {
+    comparators.push_back(
+        [&](const HloBuffer* a, const HloBuffer* b) -> std::optional<bool> {
+          bool a_fixed = has_fixed_offset(a);
+          bool b_fixed = has_fixed_offset(b);
+          if (a_fixed != b_fixed) {
+            return a_fixed > b_fixed;
+          }
+          if (a_fixed) {
+            return CompareSize(size_of, a, b);
+          }
+          return std::nullopt;
+        });
+  }
   switch (buffer_order) {
     case BufferOrder::kBiggestFirst:
       comparators.push_back(std::move(compare_size));
@@ -2541,8 +2724,11 @@ absl::Status BufferAssigner::AssignBuffersForGlobalComputations(
   DefaultBufferAllocationsManagerForComputationsWithoutOrdering default_manager(
       assignment, this);
 
-  auto get_manager = [&](LogicalBuffer::Color color)
+  auto get_manager = [&](const HloBuffer* buffer, LogicalBuffer::Color color)
       -> BufferAllocationsManagerForComputationsWithoutOrdering* {
+    if (has_fixed_offset(buffer)) {
+      return &default_manager;
+    }
     if (algorithm ==
         buffer_assignment::
             AssignmentAlgorithmForComputationsWithoutOrderingProto::
@@ -2559,7 +2745,7 @@ absl::Status BufferAssigner::AssignBuffersForGlobalComputations(
     VLOG(3) << "Assigning buffer for " << *buffer;
     ABSL_ASSIGN_OR_RETURN(auto color, buffer->color());
     BufferAllocationsManagerForComputationsWithoutOrdering* allocation_manager =
-        get_manager(color);
+        get_manager(buffer, color);
 
     ABSL_ASSIGN_OR_RETURN(bool special, AssignSpecialHloBuffer(
                                        buffer, allocation_manager, assignment));
@@ -2694,6 +2880,15 @@ absl::StatusOr<int64_t> BufferAssigner::ReuseCompatibleTempHeaps(
     if (allocation.color() != buffer_color &&
         allocation.IsPreallocatedTempBuffer() &&
         opts_.can_use_allocation(buffer_color, allocation.color())) {
+      if (opts_.fixed_offset.has_value() &&
+          absl::c_any_of(allocation.assigned_buffers(), [&](const auto& entry) {
+            const HloValue* val = entry.first;
+            return (*opts_.fixed_offset)(assignment->alias_analysis(),
+                                         val->instruction(), val->index())
+                .has_value();
+          })) {
+        continue;
+      }
       destinations.push_back(&allocation);
     }
   }
@@ -2730,6 +2925,13 @@ absl::StatusOr<int64_t> BufferAssigner::ReuseCompatibleTempHeaps(
               assignment->alias_analysis(), *buffer_value)) {
         return false;
       }
+      if (opts_.fixed_offset.has_value() &&
+          (*opts_.fixed_offset)(assignment->alias_analysis(),
+                                buffer_value->instruction(),
+                                buffer_value->index())
+              .has_value()) {
+        return false;
+      }
       return values->contains(buffer_value) &&
              buffer_value->color() == buffer_color &&
              live_ranges.contains(buffer_value);
@@ -3527,7 +3729,8 @@ absl::Status BufferAssigner::RunAssignBuffers(
   // performed after all buffers have been assigned, and after maybe_live_out
   // is marked, since it is used to determine whether an allocation contains
   // temporary buffers or not.
-  ABSL_RETURN_IF_ERROR(assignment->CombineTempAllocations(opts_.temp_buffer_color));
+  ABSL_RETURN_IF_ERROR(assignment->CombineTempAllocations(opts_.temp_buffer_color,
+                                                     opts_.fixed_offset));
   return absl::OkStatus();
 }
 
diff --git a/third_party/xla/xla/service/buffer_assignment.h b/third_party/xla/xla/service/buffer_assignment.h
index 464754cdf69..a36cdcb7132 100644
--- a/third_party/xla/xla/service/buffer_assignment.h
+++ b/third_party/xla/xla/service/buffer_assignment.h
@@ -763,7 +763,8 @@ class BufferAssignment {
   // Helper that calls NewEmptyAllocation and AddAssignment in one call,
   // creating an allocation containing a single LogicalBuffer.
   absl::StatusOr<BufferAllocation*> NewAllocation(const HloBuffer& buffer,
-                                                  int64_t size);
+                                                  int64_t size,
+                                                  int64_t offset = 0);
 
   // Adds a LogicalBuffer to the set assigned to the given allocation.
   absl::Status AddAssignment(BufferAllocation* allocation,
@@ -780,7 +781,10 @@ class BufferAssignment {
 
   // Combines allocations of temporary buffers into one big BufferAllocation.
   absl::Status CombineTempAllocations(
-      std::optional<BufferValue::Color> temp_buffer_color);
+      std::optional<BufferValue::Color> temp_buffer_color,
+      std::optional<std::function<std::optional<int64_t>(
+          const HloAliasAnalysis&, const HloInstruction*, const ShapeIndex&)>>
+          fixed_offset = std::nullopt);
 
   // Computes stats for the assignment, to be retrieved by GetStats.
   void ComputeSummaryStats();
@@ -900,6 +904,8 @@ class BufferAssigner {
       const HloAliasAnalysis&, const HloInstruction*, const ShapeIndex&)>;
   using LiveRangeInterferenceOptions =
       BufferAssignment::LiveRangeInterferenceOptions;
+  using FixedOffset = std::function<std::optional<int64_t>(
+      const HloAliasAnalysis&, const HloInstruction*, const ShapeIndex&)>;
 
   // The order in which to process buffers during buffer assignment.
   enum class BufferOrder {
@@ -927,6 +933,10 @@ class BufferAssigner {
     // Optional callbacks to override live-range interference analysis between
     // two buffer values (e.g. for target-specific concurrency or barriers).
     std::optional<LiveRangeInterferenceOptions> live_range_interference_options;
+
+    // Optional callback returning a required fixed offset for a buffer
+    // position.
+    std::optional<FixedOffset> fixed_offset;
     // Description of any buffer offsets that are already set by an earlier
     // pass.
     std::unique_ptr<memory_space_assignment::PresetAssignments>
diff --git a/third_party/xla/xla/service/buffer_assignment_test.cc b/third_party/xla/xla/service/buffer_assignment_test.cc
index c7b8dc19a58..19606072b9f 100644
--- a/third_party/xla/xla/service/buffer_assignment_test.cc
+++ b/third_party/xla/xla/service/buffer_assignment_test.cc
@@ -6107,5 +6107,349 @@ ENTRY main {
   EXPECT_EQ(buffer0_slice.offset(), slice0.offset());
 }
 
+TEST_F(BufferAssignmentTest, FixedOffsetOption) {
+  constexpr absl::string_view kHlo = R"hlo(
+HloModule main, is_scheduled=true
+
+ENTRY main {
+  buffer.0 = s32[4096]{0} custom-call(), custom_call_target="AllocateBuffer"
+  call-start.0 = ((), (), (s32[4096]{0})) call-start(), to_apply={
+    ROOT tuple.0 = tuple()
+  }
+  call-done.0 = () call-done(call-start.0)
+  ROOT copy.0 = s32[4096]{0} copy(buffer.0)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kHlo));
+
+  BufferAssigner::Options opts;
+  opts.allocate_buffers_for_constants = true;
+  opts.must_not_live_out = [](const HloAliasAnalysis& alias_analysis,
+                              const HloInstruction* instruction,
+                              const ShapeIndex& index) -> bool {
+    return absl::StartsWith(instruction->name(), "call-start");
+  };
+  opts.fixed_offset = [](const HloAliasAnalysis& alias_analysis,
+                         const HloInstruction* instruction,
+                         const ShapeIndex& index) -> std::optional<int64_t> {
+    if (absl::StartsWith(instruction->name(), "call-start") &&
+        index == ShapeIndex{2, 0}) {
+      return 0;
+    }
+    return std::nullopt;
+  };
+
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<BufferAssignment> buffers,
+      BufferAssigner::Run(
+          module.get(),
+          std::make_unique<SequentialHloOrdering>(module->schedule()),
+          &BufferSizeBytes, &alias_info_,
+          [](LogicalBuffer::Color) { return 1; }, std::move(opts)));
+
+  const HloInstruction* call_start0 =
+      module->entry_computation()->GetInstructionWithName("call-start.0");
+  const HloInstruction* buffer0 =
+      module->entry_computation()->GetInstructionWithName("buffer.0");
+
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice call_start_slice,
+                       buffers->GetUniqueSlice(call_start0, ShapeIndex{2, 0}));
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice buffer0_slice,
+                       buffers->GetUniqueTopLevelSlice(buffer0));
+
+  EXPECT_EQ(call_start_slice.offset(), 0);
+  EXPECT_NE(buffer0_slice.offset(), 0);
+  EXPECT_EQ(call_start_slice.index(), buffer0_slice.index());
+}
+
+TEST_F(BufferAssignmentTest, FixedNonZeroOffset) {
+  constexpr absl::string_view kHlo = R"hlo(
+HloModule main, is_scheduled=true
+
+ENTRY main {
+  buffer.0 = s32[4096]{0} custom-call(), custom_call_target="AllocateBuffer"
+  call-start.0 = ((), (), (s32[4096]{0})) call-start(), to_apply={
+    ROOT tuple.0 = tuple()
+  }
+  call-done.0 = () call-done(call-start.0)
+  ROOT copy.0 = s32[4096]{0} copy(buffer.0)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kHlo));
+
+  constexpr int64_t kFixedOffset = 1024;
+  BufferAssigner::Options opts;
+  opts.allocate_buffers_for_constants = true;
+  opts.must_not_live_out = [](const HloAliasAnalysis& alias_analysis,
+                              const HloInstruction* instruction,
+                              const ShapeIndex& index) -> bool {
+    return absl::StartsWith(instruction->name(), "call-start");
+  };
+  opts.fixed_offset = [kFixedOffset](
+                          const HloAliasAnalysis& alias_analysis,
+                          const HloInstruction* instruction,
+                          const ShapeIndex& index) -> std::optional<int64_t> {
+    if (absl::StartsWith(instruction->name(), "call-start") &&
+        index == ShapeIndex{2, 0}) {
+      return kFixedOffset;
+    }
+    return std::nullopt;
+  };
+
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<BufferAssignment> buffers,
+      BufferAssigner::Run(
+          module.get(),
+          std::make_unique<SequentialHloOrdering>(module->schedule()),
+          &BufferSizeBytes, &alias_info_,
+          [](LogicalBuffer::Color) { return 1; }, std::move(opts)));
+
+  const HloInstruction* call_start0 =
+      module->entry_computation()->GetInstructionWithName("call-start.0");
+  const HloInstruction* buffer0 =
+      module->entry_computation()->GetInstructionWithName("buffer.0");
+
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice call_start_slice,
+                       buffers->GetUniqueSlice(call_start0, ShapeIndex{2, 0}));
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice buffer0_slice,
+                       buffers->GetUniqueTopLevelSlice(buffer0));
+
+  EXPECT_EQ(call_start_slice.offset(), kFixedOffset);
+  EXPECT_NE(buffer0_slice.offset(), kFixedOffset);
+  const BufferAllocation& call_start_alloc =
+      buffers->GetAllocation(call_start_slice.index());
+  EXPECT_GE(call_start_alloc.size(), kFixedOffset + call_start_slice.size());
+}
+
+TEST_F(BufferAssignmentTest,
+       FixedZeroOffsetInterferingBuffersGetDistinctAllocations) {
+  constexpr absl::string_view kHlo = R"hlo(
+HloModule main, is_scheduled=true
+
+ENTRY main {
+  call-start.0 = ((), (), (s32[4096]{0})) call-start(), to_apply={
+    ROOT tuple.0 = tuple()
+  }
+  call-start.1 = ((), (), (s32[4096]{0})) call-start(), to_apply={
+    ROOT tuple.0 = tuple()
+  }
+  call-done.0 = () call-done(call-start.0)
+  call-done.1 = () call-done(call-start.1)
+  buffer.0 = s32[4096]{0} custom-call(), custom_call_target="AllocateBuffer"
+  ROOT copy.0 = s32[4096]{0} copy(buffer.0)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kHlo));
+
+  BufferAssigner::Options opts;
+  opts.allocate_buffers_for_constants = true;
+  opts.must_not_live_out = [](const HloAliasAnalysis& alias_analysis,
+                              const HloInstruction* instruction,
+                              const ShapeIndex& index) -> bool {
+    return absl::StartsWith(instruction->name(), "call-start");
+  };
+  opts.fixed_offset = [](const HloAliasAnalysis& alias_analysis,
+                         const HloInstruction* instruction,
+                         const ShapeIndex& index) -> std::optional<int64_t> {
+    if (absl::StartsWith(instruction->name(), "call-start") &&
+        index == ShapeIndex{2, 0}) {
+      return 0;
+    }
+    return std::nullopt;
+  };
+
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<BufferAssignment> buffers,
+      BufferAssigner::Run(
+          module.get(),
+          std::make_unique<SequentialHloOrdering>(module->schedule()),
+          &BufferSizeBytes, &alias_info_,
+          [](LogicalBuffer::Color) { return 1; }, std::move(opts)));
+
+  const HloInstruction* call_start0 =
+      module->entry_computation()->GetInstructionWithName("call-start.0");
+  const HloInstruction* call_start1 =
+      module->entry_computation()->GetInstructionWithName("call-start.1");
+
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice slice0,
+                       buffers->GetUniqueSlice(call_start0, ShapeIndex{2, 0}));
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice slice1,
+                       buffers->GetUniqueSlice(call_start1, ShapeIndex{2, 0}));
+
+  EXPECT_EQ(slice0.offset(), 0);
+  EXPECT_EQ(slice1.offset(), 0);
+  EXPECT_NE(slice0.index(), slice1.index());
+}
+
+TEST_F(BufferAssignmentTest,
+       FixedZeroOffsetNonOverlappingBuffersReuseAllocation) {
+  constexpr absl::string_view kHlo = R"hlo(
+HloModule main, is_scheduled=true
+
+ENTRY main {
+  call-start.0 = ((), (), (s32[4096]{0})) call-start(), to_apply={
+    ROOT tuple.0 = tuple()
+  }
+  call-done.0 = () call-done(call-start.0)
+  call-start.1 = ((), (), (s32[8192]{0})) call-start(), to_apply={
+    ROOT tuple.0 = tuple()
+  }
+  call-done.1 = () call-done(call-start.1)
+  buffer.0 = s32[4096]{0} custom-call(), custom_call_target="AllocateBuffer"
+  ROOT copy.0 = s32[4096]{0} copy(buffer.0)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kHlo));
+
+  BufferAssigner::Options opts;
+  opts.allocate_buffers_for_constants = true;
+  opts.assignment_algorithm_for_computations_without_ordering =
+      buffer_assignment::
+          AssignmentAlgorithmForComputationsWithoutOrderingProto::FAST_MERGE;
+  opts.color_memory_limit = [](LogicalBuffer::Color) { return 1024 * 1024; };
+  opts.must_not_live_out = [](const HloAliasAnalysis& alias_analysis,
+                              const HloInstruction* instruction,
+                              const ShapeIndex& index) -> bool {
+    return absl::StartsWith(instruction->name(), "call-start");
+  };
+  opts.fixed_offset = [](const HloAliasAnalysis& alias_analysis,
+                         const HloInstruction* instruction,
+                         const ShapeIndex& index) -> std::optional<int64_t> {
+    if (absl::StartsWith(instruction->name(), "call-start") &&
+        index == ShapeIndex{2, 0}) {
+      return 0;
+    }
+    return std::nullopt;
+  };
+
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<BufferAssignment> buffers,
+      BufferAssigner::Run(
+          module.get(),
+          std::make_unique<SequentialHloOrdering>(module->schedule()),
+          &BufferSizeBytes, &alias_info_,
+          [](LogicalBuffer::Color) { return 1; }, std::move(opts)));
+
+  const HloInstruction* call_start0 =
+      module->entry_computation()->GetInstructionWithName("call-start.0");
+  const HloInstruction* call_start1 =
+      module->entry_computation()->GetInstructionWithName("call-start.1");
+  const HloInstruction* buffer0 =
+      module->entry_computation()->GetInstructionWithName("buffer.0");
+
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice slice0,
+                       buffers->GetUniqueSlice(call_start0, ShapeIndex{2, 0}));
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice slice1,
+                       buffers->GetUniqueSlice(call_start1, ShapeIndex{2, 0}));
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice buffer0_slice,
+                       buffers->GetUniqueTopLevelSlice(buffer0));
+
+  EXPECT_EQ(slice0.offset(), 0);
+  EXPECT_EQ(slice1.offset(), 0);
+  EXPECT_EQ(slice0.index(), slice1.index());
+  EXPECT_EQ(buffer0_slice.index(), slice0.index());
+  EXPECT_NE(buffer0_slice.offset(), 0);
+}
+
+TEST_F(BufferAssignmentTest, ConflictingFixedOffsetsReturnsError) {
+  constexpr absl::string_view kHlo = R"hlo(
+HloModule main, is_scheduled=true
+
+ENTRY main {
+  buffer.0 = s32[4096]{0} custom-call(), custom_call_target="AllocateBuffer"
+  alias.0 = s32[4096]{0} custom-call(buffer.0),
+    custom_call_target="InPlaceOp",
+    output_to_operand_aliasing={{}: (0, {})}
+  ROOT copy.0 = s32[4096]{0} copy(alias.0)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kHlo));
+
+  BufferAssigner::Options opts;
+  opts.allocate_buffers_for_constants = true;
+  opts.fixed_offset = [](const HloAliasAnalysis& alias_analysis,
+                         const HloInstruction* instruction,
+                         const ShapeIndex& index) -> std::optional<int64_t> {
+    if (instruction->name() == "buffer.0") {
+      return 0;
+    }
+    if (instruction->name() == "alias.0") {
+      return 1024;
+    }
+    return std::nullopt;
+  };
+
+  EXPECT_THAT(BufferAssigner::Run(
+                  module.get(),
+                  std::make_unique<SequentialHloOrdering>(module->schedule()),
+                  &BufferSizeBytes, &alias_info_,
+                  [](LogicalBuffer::Color) { return 1; }, std::move(opts)),
+              absl_testing::StatusIs(
+                  absl::StatusCode::kInternal,
+                  HasSubstr("Conflicting fixed offset requirements")));
+}
+
+TEST_F(BufferAssignmentTest,
+       FixedZeroOffsetDoesNotReuseNonFixedOffsetAllocation) {
+  constexpr absl::string_view kHlo = R"hlo(
+HloModule main, is_scheduled=true
+
+ENTRY main {
+  buffer.0 = s32[8192]{0} custom-call(), custom_call_target="AllocateBuffer"
+  use.0 = () custom-call(buffer.0), custom_call_target="UseBuffer"
+  call-start.0 = ((), (), (s32[4096]{0})) call-start(), to_apply={
+    ROOT tuple.0 = tuple()
+  }
+  call-done.0 = () call-done(call-start.0)
+  ROOT root_tuple = () tuple()
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kHlo));
+
+  BufferAssigner::Options opts;
+  opts.allocate_buffers_for_constants = true;
+  opts.must_not_live_out = [](const HloAliasAnalysis& alias_analysis,
+                              const HloInstruction* instruction,
+                              const ShapeIndex& index) -> bool {
+    return absl::StartsWith(instruction->name(), "call-start");
+  };
+  opts.fixed_offset = [](const HloAliasAnalysis& alias_analysis,
+                         const HloInstruction* instruction,
+                         const ShapeIndex& index) -> std::optional<int64_t> {
+    if (absl::StartsWith(instruction->name(), "call-start") &&
+        index == ShapeIndex{2, 0}) {
+      return 0;
+    }
+    return std::nullopt;
+  };
+
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<BufferAssignment> buffers,
+      BufferAssigner::Run(
+          module.get(),
+          std::make_unique<SequentialHloOrdering>(module->schedule()),
+          &BufferSizeBytes, &alias_info_,
+          [](LogicalBuffer::Color) { return 1; }, std::move(opts)));
+
+  const HloInstruction* call_start0 =
+      module->entry_computation()->GetInstructionWithName("call-start.0");
+  const HloInstruction* buffer0 =
+      module->entry_computation()->GetInstructionWithName("buffer.0");
+
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice slice0,
+                       buffers->GetUniqueSlice(call_start0, ShapeIndex{2, 0}));
+  ASSERT_OK_AND_ASSIGN(const BufferAllocation::Slice buffer0_slice,
+                       buffers->GetUniqueTopLevelSlice(buffer0));
+
+  EXPECT_EQ(slice0.index(), buffer0_slice.index());
+  EXPECT_EQ(slice0.offset(), 0);
+  EXPECT_GE(buffer0_slice.offset(), slice0.size());
+}
 }  // namespace
 }  // namespace xla

From ae03a9cbabf67d326d6704a787af126e5db6b753 Mon Sep 17 00:00:00 2001
From: Parker Schuh <parkers@google.com>
Date: Mon, 12 Oct 2026 19:13:18 -0700
Subject: [PATCH 0637/1000] Remove addressable_device_logical_ids because it is
 unused. Because of the ability to reload an executable at runtime at
 different offsets when running multipod, we should prefer tracking the device
 information at the higher level (like ifrt) based off of information we pull
 from the PjRtExecutable. (see GetDeviceListFromDeviceAssignment, but note
 that it also doesn't really handle subslice loading properly).

PiperOrigin-RevId: 998299049
---
 .../xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc | 22 ++----------
 .../xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.h  |  1 -
 .../pjrt/c_api_client/pjrt_c_api_client.cc    | 35 -------------------
 .../xla/pjrt/c_api_client/pjrt_c_api_client.h |  7 ----
 .../c_api_client/pjrt_c_api_client_test.cc    | 30 ----------------
 third_party/xla/xla/pjrt/common_pjrt_client.h |  5 ---
 .../pjrt/interpreter/interpreter_client.cc    | 12 +------
 .../xla/pjrt/interpreter/interpreter_client.h |  7 ----
 third_party/xla/xla/pjrt/pjrt_client.h        |  9 -----
 third_party/xla/xla/pjrt/tf_pjrt_client.h     |  4 ---
 .../functional_hlo_runner.cc                  | 31 ++++------------
 11 files changed, 9 insertions(+), 154 deletions(-)

diff --git a/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc b/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc
index 18cb6bcff36..84268b334f1 100644
--- a/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc
+++ b/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc
@@ -1919,11 +1919,8 @@ PJRT_Error* PJRT_LoadedExecutable_AddressableDeviceLogicalIds(
       PJRT_LoadedExecutable_AddressableDeviceLogicalIds_Args_STRUCT_SIZE,
       args->struct_size));
 
-  args->num_addressable_device_logical_ids =
-      args->executable->addressable_device_logical_ids.size();
-  args->addressable_device_logical_ids =
-      args->executable->addressable_device_logical_ids.data();
-  return nullptr;
+  return StatusToPjRtError(xla::Unimplemented(
+      "PJRT_LoadedExecutable_AddressableDeviceLogicalIds is not implemented."));
 }
 
 PJRT_Error* PJRT_Executable_NumOutputs(PJRT_Executable_NumOutputs_Args* args) {
@@ -3750,25 +3747,10 @@ PJRT_Executable::PJRT_Executable(xla::PjRtExecutable* unowned_executable)
     : executable(unowned_executable),
       fingerprint(executable->FingerprintExecutable()) {}
 
-static void PopulatePjrtExecutableAddressableDeviceLogicalIds(
-    PJRT_LoadedExecutable* executable) {
-  CHECK(executable->client != nullptr) << ": client was null";
-  absl::Span<const xla::PjRtLoadedExecutable::LogicalDeviceIds> cpp_ids =
-      executable->get()->addressable_device_logical_ids();
-
-  std::vector<PJRT_LogicalDeviceIds>& exec_ids =
-      executable->addressable_device_logical_ids;
-  exec_ids.reserve(cpp_ids.size());
-  for (const auto& id : cpp_ids) {
-    exec_ids.push_back(PJRT_LogicalDeviceIds{id.replica, id.partition});
-  }
-}
-
 PJRT_LoadedExecutable::PJRT_LoadedExecutable(
     std::shared_ptr<xla::PjRtLoadedExecutable> executable, PJRT_Client* client)
     : executable(std::move(executable)), client(client) {
   pjrt::PopulatePjrtExecutableAddressableDevices(this);
-  PopulatePjrtExecutableAddressableDeviceLogicalIds(this);
 }
 
 namespace pjrt {
diff --git a/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.h b/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.h
index 6ef2e0d9b4b..956d362463b 100644
--- a/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.h
+++ b/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.h
@@ -189,7 +189,6 @@ struct PJRT_LoadedExecutable {
   // addressed by the compiled executable program. `client` owns the objects
   // these point to.
   std::vector<PJRT_Device*> addressable_devices;
-  std::vector<PJRT_LogicalDeviceIds> addressable_device_logical_ids;
 
   PJRT_LoadedExecutable(std::shared_ptr<xla::PjRtLoadedExecutable> executable,
                         PJRT_Client* client);
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 273c6acd9d6..e432478a05f 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
@@ -2832,44 +2832,9 @@ PjRtCApiLoadedExecutable::PjRtCApiLoadedExecutable(
   executable_ =
       std::make_unique<PjRtCApiExecutable>(pjrt_c_api(), args.executable);
   InitDevices();
-  InitAddressableDeviceLogicalIds();
   InitDeviceAssignment();
 }
 
-void PjRtCApiLoadedExecutable::InitAddressableDeviceLogicalIds() {
-  if (pjrt_c_api()->pjrt_api_version.major_version == 0 &&
-      pjrt_c_api()->pjrt_api_version.minor_version < 96) {
-    return;
-  }
-  PJRT_LoadedExecutable_AddressableDeviceLogicalIds_Args args;
-  args.struct_size =
-      PJRT_LoadedExecutable_AddressableDeviceLogicalIds_Args_STRUCT_SIZE;
-  args.extension_start = nullptr;
-  args.executable = c_loaded_executable();
-  args.addressable_device_logical_ids = nullptr;
-  args.num_addressable_device_logical_ids = 0;
-
-  const PJRT_Api* api = pjrt_c_api();
-  if (api->PJRT_LoadedExecutable_AddressableDeviceLogicalIds == nullptr) {
-    return;
-  }
-  pjrt::LogFatalIfPjrtError(
-      api->PJRT_LoadedExecutable_AddressableDeviceLogicalIds(&args), api);
-
-  if (args.num_addressable_device_logical_ids == 0 ||
-      args.addressable_device_logical_ids == nullptr) {
-    return;
-  }
-  addressable_device_logical_ids_.reserve(
-      args.num_addressable_device_logical_ids);
-
-  for (size_t i = 0; i < args.num_addressable_device_logical_ids; ++i) {
-    PJRT_LogicalDeviceIds id = args.addressable_device_logical_ids[i];
-    addressable_device_logical_ids_.push_back(
-        LogicalDeviceIds{id.replica, id.partition});
-  }
-}
-
 void PjRtCApiLoadedExecutable::InitDevices() {
   PJRT_LoadedExecutable_AddressableDevices_Args args;
   args.struct_size = PJRT_LoadedExecutable_AddressableDevices_Args_STRUCT_SIZE;
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 51052421965..81098544246 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
@@ -823,11 +823,6 @@ class PjRtCApiLoadedExecutable : public PjRtLoadedExecutable {
     return *device_assignment_;
   }
 
-  absl::Span<const LogicalDeviceIds> addressable_device_logical_ids()
-      const override {
-    return addressable_device_logical_ids_;
-  }
-
   absl::Span<PjRtDevice* const> addressable_devices() const override {
     return addressable_devices_;
   }
@@ -934,7 +929,6 @@ class PjRtCApiLoadedExecutable : public PjRtLoadedExecutable {
   }
 
  private:
-  std::vector<LogicalDeviceIds> addressable_device_logical_ids_;
   // Groups data needed to support send/recv execution callbacks.
   struct HloOutputCallbackState {
     // Key is {replica_id, partition_id}.
@@ -1013,7 +1007,6 @@ class PjRtCApiLoadedExecutable : public PjRtLoadedExecutable {
   std::vector<PjRtDevice*> addressable_devices_;
   std::unique_ptr<const DeviceAssignment> device_assignment_;
 
-  void InitAddressableDeviceLogicalIds();
   void InitDevices();
   void InitDeviceAssignment();
 };
diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc
index d5d72885516..202bdc7abe9 100644
--- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc
+++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc
@@ -1060,36 +1060,6 @@ ENTRY Identity() -> f32[2, 2] {
   EXPECT_THAT(poison_result, IsOkAndHolds(false));
 }
 
-TEST(PjRtCApiClientTest, AddressableDeviceLogicalIds) {
-  SetUpCpuPjRtApi();
-  ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtClient> client,
-                       GetCApiClient("cpu"));
-  ASSERT_GT(client->addressable_devices().size(), 1);
-
-  XlaBuilder builder("Identity");
-  Shape shape = ShapeUtil::MakeShape(S32, {2, 3});
-  XlaOp input = Parameter(&builder, 0, shape, "input");
-  XlaComputation computation = builder.Build(input).value();
-
-  DeviceAssignment device_assignment(1, 2);
-  device_assignment(0, 0) = 0;
-  device_assignment(0, 1) = 1;
-
-  CompileOptions options;
-  options.executable_build_options.set_device_assignment(device_assignment);
-
-  ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtLoadedExecutable> executable,
-                       client->CompileAndLoad(computation, options));
-
-  absl::Span<const PjRtLoadedExecutable::LogicalDeviceIds> logical_ids =
-      executable->addressable_device_logical_ids();
-  ASSERT_EQ(logical_ids.size(), 2);
-  EXPECT_EQ(logical_ids[0].replica, 0);
-  EXPECT_EQ(logical_ids[0].partition, 0);
-  EXPECT_EQ(logical_ids[1].replica, 0);
-  EXPECT_EQ(logical_ids[1].partition, 1);
-}
-
 TEST(PjRtCApiClientTest, Bitcast) {
   SetUpCpuPjRtApi();
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtClient> client,
diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.h b/third_party/xla/xla/pjrt/common_pjrt_client.h
index 2eeffdb8254..261ad93d8ba 100644
--- a/third_party/xla/xla/pjrt/common_pjrt_client.h
+++ b/third_party/xla/xla/pjrt/common_pjrt_client.h
@@ -654,11 +654,6 @@ class CommonPjRtLoadedExecutable : public PjRtLoadedExecutable {
     return *device_assignment_;
   }
 
-  absl::Span<const LogicalDeviceIds> addressable_device_logical_ids()
-      const override {
-    return addressable_device_logical_ids_;
-  }
-
   void Delete() override {
     if (load_state_ != nullptr) {
       load_state_->Delete();
diff --git a/third_party/xla/xla/pjrt/interpreter/interpreter_client.cc b/third_party/xla/xla/pjrt/interpreter/interpreter_client.cc
index e8f7ca7e7b8..4a6a5b61c98 100644
--- a/third_party/xla/xla/pjrt/interpreter/interpreter_client.cc
+++ b/third_party/xla/xla/pjrt/interpreter/interpreter_client.cc
@@ -162,15 +162,12 @@ InterpreterLoadedExecutable::InterpreterLoadedExecutable(
     std::shared_ptr<InterpreterExecutable> executable,
     std::unique_ptr<HloEvaluatorInterface> hlo_evaluator,
     std::shared_ptr<DeviceAssignment> device_assignment,
-    std::vector<LogicalDeviceIds> addressable_device_logical_ids,
     std::vector<PjRtDevice*> addressable_devices)
     : client_(ABSL_DIE_IF_NULL(client)),
       name_(ABSL_DIE_IF_NULL(executable)->name()),
       executable_(std::move(executable)),
       hlo_evaluator_(std::move(hlo_evaluator)),
       device_assignment_(std::move(device_assignment)),
-      addressable_device_logical_ids_(
-          std::move(addressable_device_logical_ids)),
       addressable_devices_(std::move(addressable_devices)) {
   if (executable_ && executable_->dynamic_dimension_inference().has_value()) {
     hlo_evaluator_->set_dynamic_dimension_inference(
@@ -529,19 +526,12 @@ absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> InterpreterClient::Load(
                         "num_replicas: %d, num_partitions: %d",
                         num_replicas, num_partitions));
   }
-  std::vector<PjRtLoadedExecutable::LogicalDeviceIds>
-      addressable_device_logical_ids;
   std::vector<PjRtDevice*> addressable_devices;
-  PjRtLoadedExecutable::LogicalDeviceIds logical_device_ids;
-  logical_device_ids.replica = 0;
-  logical_device_ids.partition = 0;
-  addressable_device_logical_ids.push_back(std::move(logical_device_ids));
   addressable_devices.push_back(&interpreter_device_);
 
   return std::make_unique<InterpreterLoadedExecutable>(
       this, std::move(shared_interpreter_executable), std::move(evaluator),
-      std::move(device_assignment), std::move(addressable_device_logical_ids),
-      std::move(addressable_devices));
+      std::move(device_assignment), std::move(addressable_devices));
 }
 
 absl::StatusOr<std::unique_ptr<PjRtExecutable>>
diff --git a/third_party/xla/xla/pjrt/interpreter/interpreter_client.h b/third_party/xla/xla/pjrt/interpreter/interpreter_client.h
index 2e9bda4b9e9..5a6ecf597db 100644
--- a/third_party/xla/xla/pjrt/interpreter/interpreter_client.h
+++ b/third_party/xla/xla/pjrt/interpreter/interpreter_client.h
@@ -297,7 +297,6 @@ class InterpreterLoadedExecutable final : public PjRtLoadedExecutable {
       std::shared_ptr<InterpreterExecutable> executable,
       std::unique_ptr<HloEvaluatorInterface> hlo_evaluator,
       std::shared_ptr<DeviceAssignment> device_assignment,
-      std::vector<LogicalDeviceIds> addressable_device_logical_ids,
       std::vector<PjRtDevice*> addressable_devices);
 
   InterpreterExecutable* GetExecutable() const override;
@@ -321,11 +320,6 @@ class InterpreterLoadedExecutable final : public PjRtLoadedExecutable {
     return *device_assignment_;
   }
 
-  absl::Span<const LogicalDeviceIds> addressable_device_logical_ids()
-      const override {
-    return addressable_device_logical_ids_;
-  }
-
   absl::Span<PjRtDevice* const> addressable_devices() const override {
     return addressable_devices_;
   }
@@ -361,7 +355,6 @@ class InterpreterLoadedExecutable final : public PjRtLoadedExecutable {
   std::shared_ptr<InterpreterExecutable> executable_ ABSL_GUARDED_BY(mutex_);
   std::unique_ptr<HloEvaluatorInterface> hlo_evaluator_ ABSL_GUARDED_BY(mutex_);
   std::shared_ptr<DeviceAssignment> device_assignment_;
-  std::vector<LogicalDeviceIds> addressable_device_logical_ids_;
   std::vector<PjRtDevice*> addressable_devices_;
 };
 
diff --git a/third_party/xla/xla/pjrt/pjrt_client.h b/third_party/xla/xla/pjrt/pjrt_client.h
index 8ad889abe2d..63aa3f45d5d 100644
--- a/third_party/xla/xla/pjrt/pjrt_client.h
+++ b/third_party/xla/xla/pjrt/pjrt_client.h
@@ -1411,21 +1411,12 @@ class PjRtLoadedExecutable {
   virtual absl::StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
   GetCostAnalysis() const;
 
-  // The replica and partition indices of device_assignment to be run by this
-  // client. On single-host platforms without partitioning, this is all replicas
-  // (i.e. addressable_device_logical_ids_[i] = (i, 0)), but this may not be the
-  // case on multi-host platforms. If there are 4 replicas and 2 partitions on a
-  // single host platform, size of addressable_device_logical_ids_ is 4*2 = 8.
   struct LogicalDeviceIds {
     int replica;
     int partition;
   };
-  virtual absl::Span<const LogicalDeviceIds> addressable_device_logical_ids()
-      const = 0;
 
   // An addressable_device is one which the client can issue commands to.
-  // addressable_devices()[i] is the Device to which
-  // addressable_device_logical_ids()[i] is assigned.
   virtual absl::Span<PjRtDevice* const> addressable_devices() const = 0;
 
   // Donation Semantics:
diff --git a/third_party/xla/xla/pjrt/tf_pjrt_client.h b/third_party/xla/xla/pjrt/tf_pjrt_client.h
index 0943e22bddc..a934a74f8f6 100644
--- a/third_party/xla/xla/pjrt/tf_pjrt_client.h
+++ b/third_party/xla/xla/pjrt/tf_pjrt_client.h
@@ -144,10 +144,6 @@ class TfPjRtExecutable : public PjRtLoadedExecutable {
   const DeviceAssignment& device_assignment() const override {
     return wrapped_->device_assignment();
   }
-  absl::Span<const LogicalDeviceIds> addressable_device_logical_ids()
-      const override {
-    return wrapped_->addressable_device_logical_ids();
-  }
   absl::Span<PjRtDevice* const> addressable_devices() const override {
     return wrapped_->addressable_devices();
   }
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
index ac8cce570ff..9cffeee12ef 100644
--- a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
@@ -756,11 +756,10 @@ CopyArgumentsToDevice(PjRtClient& client,
     return device_buffers;
   };
 
-  absl::Span<const PjRtLoadedExecutable::LogicalDeviceIds>
-      addressable_device_logical_ids =
-          executable->addressable_device_logical_ids();
   ABSL_ASSIGN_OR_RETURN(std::vector<std::shared_ptr<HloModule>> hlo_modules,
                    executable->GetHloModules());
+  CHECK_EQ(hlo_modules.size(), 1);
+  HloModule* module = hlo_modules.front().get();
 
   for (int i = 0; i < num_addressable_devices; ++i) {
     PjRtDevice* curr_device = addressable_devices[i];
@@ -779,11 +778,6 @@ CopyArgumentsToDevice(PjRtClient& client,
     const std::vector<Literal>& curr_device_arguments =
         arguments.at(source_device_id);
 
-    int executable_idx = hlo_modules.size() == 1
-                             ? 0
-                             : addressable_device_logical_ids[i].partition;
-    HloModule* module = hlo_modules[executable_idx].get();
-
     argument_buffers[i].reserve(curr_device_arguments.size());
     for (int arg_i = 0; arg_i < curr_device_arguments.size(); ++arg_i) {
       const Literal& literal = curr_device_arguments[arg_i];
@@ -813,26 +807,18 @@ CreateUninitializedArgumentsOnDevice(PjRtClient& client,
                                      bool flatten_arguments = false) {
   absl::Span<PjRtDevice* const> addressable_devices =
       executable->addressable_devices();
-  absl::Span<const PjRtLoadedExecutable::LogicalDeviceIds>
-      addressable_device_logical_ids =
-          executable->addressable_device_logical_ids();
   ABSL_ASSIGN_OR_RETURN(std::vector<std::shared_ptr<HloModule>> hlo_modules,
                    executable->GetHloModules());
   VLOG(1) << "FunctionalHloRunner: local_executable count = "
           << hlo_modules.size();
+  CHECK_EQ(hlo_modules.size(), 1);
+  const HloModule& hlo_module = *hlo_modules.front();
 
   LOG(INFO) << "Starting argument buffer shape calculation.";
   PerDeviceShapeVecType argument_shapes_per_device;
-  // This must be true, based on the comment on
-  // PjRtLoadedExecutable::addressable_devices().
-  CHECK_EQ(addressable_devices.size(), addressable_device_logical_ids.size());
   for (int i = 0; i < static_cast<int>(addressable_devices.size()); ++i) {
     VLOG(3) << "Calculating fake argument shapes for device " << i;
     PjRtDevice* device = addressable_devices[i];
-    int executable_idx = hlo_modules.size() == 1
-                             ? 0
-                             : addressable_device_logical_ids[i].partition;
-    const HloModule& hlo_module = *hlo_modules[executable_idx];
 
     std::vector<Shape> argument_shapes;
     if (flatten_arguments) {
@@ -916,13 +902,12 @@ CreateArgumentsOnDevice(PjRtClient& client,
   size_t num_addressable_devices = addressable_devices.size();
 
   PerDeviceLiteralVecType per_device_argument_literals;
-  absl::Span<const PjRtLoadedExecutable::LogicalDeviceIds>
-      addressable_device_logical_ids =
-          executable->addressable_device_logical_ids();
   ABSL_ASSIGN_OR_RETURN(std::vector<std::shared_ptr<HloModule>> hlo_modules,
                    executable->GetHloModules());
   VLOG(1) << "FunctionalHloRunner: local_executable count = "
           << hlo_modules.size();
+  CHECK_EQ(hlo_modules.size(), 1);
+  HloModule* my_hlo_module = hlo_modules.front().get();
 
   const bool kUseRandomInputs = running_options.module_argument_mode ==
                                     ModuleArgumentMode::kUseRandomInputs ||
@@ -949,10 +934,6 @@ CreateArgumentsOnDevice(PjRtClient& client,
     VLOG(3) << "Creating fake arguments for device " << i;
     LiteralVec& argument_literals =
         per_device_argument_literals[addressable_devices[i]->id()];
-    int executable_idx = hlo_modules.size() == 1
-                             ? 0
-                             : addressable_device_logical_ids[i].partition;
-    HloModule* my_hlo_module = hlo_modules[executable_idx].get();
     if (flatten_arguments) {
       ABSL_RETURN_IF_ERROR(EnsureSingleTupleForFlattening(*my_hlo_module));
     }

From 7fc70a240d246eb2501026bee5a87dd34111c50d Mon Sep 17 00:00:00 2001
From: Ricky Zhou <rickyz@google.com>
Date: Mon, 12 Oct 2026 19:13:31 -0700
Subject: [PATCH 0638/1000] Add a fast path for comparing a PjRtLayout against
 itself.

PiperOrigin-RevId: 998299116
---
 third_party/xla/xla/pjrt/pjrt_layout.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/xla/xla/pjrt/pjrt_layout.h b/third_party/xla/xla/pjrt/pjrt_layout.h
index e0367f8387f..f276bece0d7 100644
--- a/third_party/xla/xla/pjrt/pjrt_layout.h
+++ b/third_party/xla/xla/pjrt/pjrt_layout.h
@@ -70,7 +70,7 @@ class PjRtLayout {
   std::string ToString() const { return xla_layout_.ToString(); }
 
   bool operator==(const PjRtLayout& other) const {
-    return xla_layout_ == other.xla_layout_;
+    return this == &other || xla_layout_ == other.xla_layout_;
   }
 
   bool operator!=(const PjRtLayout& other) const { return !(*this == other); }

From 00f3eb230cb37abfa91f4cedbb64eebc434119b6 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 12 Oct 2026 19:45:09 -0700
Subject: [PATCH 0639/1000] Remove addressable_device_logical_ids because it is
 unused. Because of the ability to reload an executable at runtime at
 different offsets when running multipod, we should prefer tracking the device
 information at the higher level (like ifrt) based off of information we pull
 from the PjRtExecutable. (see GetDeviceListFromDeviceAssignment, but note
 that it also doesn't really handle subslice loading properly).

Reverts ae03a9cbabf67d326d6704a787af126e5db6b753

PiperOrigin-RevId: 998311746
---
 .../xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc | 22 ++++++++++--
 .../xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.h  |  1 +
 .../pjrt/c_api_client/pjrt_c_api_client.cc    | 35 +++++++++++++++++++
 .../xla/pjrt/c_api_client/pjrt_c_api_client.h |  7 ++++
 .../c_api_client/pjrt_c_api_client_test.cc    | 30 ++++++++++++++++
 third_party/xla/xla/pjrt/common_pjrt_client.h |  5 +++
 .../pjrt/interpreter/interpreter_client.cc    | 12 ++++++-
 .../xla/pjrt/interpreter/interpreter_client.h |  7 ++++
 third_party/xla/xla/pjrt/pjrt_client.h        |  9 +++++
 third_party/xla/xla/pjrt/tf_pjrt_client.h     |  4 +++
 .../functional_hlo_runner.cc                  | 31 ++++++++++++----
 11 files changed, 154 insertions(+), 9 deletions(-)

diff --git a/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc b/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc
index 84268b334f1..18cb6bcff36 100644
--- a/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc
+++ b/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.cc
@@ -1919,8 +1919,11 @@ PJRT_Error* PJRT_LoadedExecutable_AddressableDeviceLogicalIds(
       PJRT_LoadedExecutable_AddressableDeviceLogicalIds_Args_STRUCT_SIZE,
       args->struct_size));
 
-  return StatusToPjRtError(xla::Unimplemented(
-      "PJRT_LoadedExecutable_AddressableDeviceLogicalIds is not implemented."));
+  args->num_addressable_device_logical_ids =
+      args->executable->addressable_device_logical_ids.size();
+  args->addressable_device_logical_ids =
+      args->executable->addressable_device_logical_ids.data();
+  return nullptr;
 }
 
 PJRT_Error* PJRT_Executable_NumOutputs(PJRT_Executable_NumOutputs_Args* args) {
@@ -3747,10 +3750,25 @@ PJRT_Executable::PJRT_Executable(xla::PjRtExecutable* unowned_executable)
     : executable(unowned_executable),
       fingerprint(executable->FingerprintExecutable()) {}
 
+static void PopulatePjrtExecutableAddressableDeviceLogicalIds(
+    PJRT_LoadedExecutable* executable) {
+  CHECK(executable->client != nullptr) << ": client was null";
+  absl::Span<const xla::PjRtLoadedExecutable::LogicalDeviceIds> cpp_ids =
+      executable->get()->addressable_device_logical_ids();
+
+  std::vector<PJRT_LogicalDeviceIds>& exec_ids =
+      executable->addressable_device_logical_ids;
+  exec_ids.reserve(cpp_ids.size());
+  for (const auto& id : cpp_ids) {
+    exec_ids.push_back(PJRT_LogicalDeviceIds{id.replica, id.partition});
+  }
+}
+
 PJRT_LoadedExecutable::PJRT_LoadedExecutable(
     std::shared_ptr<xla::PjRtLoadedExecutable> executable, PJRT_Client* client)
     : executable(std::move(executable)), client(client) {
   pjrt::PopulatePjrtExecutableAddressableDevices(this);
+  PopulatePjrtExecutableAddressableDeviceLogicalIds(this);
 }
 
 namespace pjrt {
diff --git a/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.h b/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.h
index 956d362463b..6ef2e0d9b4b 100644
--- a/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.h
+++ b/third_party/xla/xla/pjrt/c/pjrt_c_api_wrapper_impl.h
@@ -189,6 +189,7 @@ struct PJRT_LoadedExecutable {
   // addressed by the compiled executable program. `client` owns the objects
   // these point to.
   std::vector<PJRT_Device*> addressable_devices;
+  std::vector<PJRT_LogicalDeviceIds> addressable_device_logical_ids;
 
   PJRT_LoadedExecutable(std::shared_ptr<xla::PjRtLoadedExecutable> executable,
                         PJRT_Client* client);
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 e432478a05f..273c6acd9d6 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
@@ -2832,9 +2832,44 @@ PjRtCApiLoadedExecutable::PjRtCApiLoadedExecutable(
   executable_ =
       std::make_unique<PjRtCApiExecutable>(pjrt_c_api(), args.executable);
   InitDevices();
+  InitAddressableDeviceLogicalIds();
   InitDeviceAssignment();
 }
 
+void PjRtCApiLoadedExecutable::InitAddressableDeviceLogicalIds() {
+  if (pjrt_c_api()->pjrt_api_version.major_version == 0 &&
+      pjrt_c_api()->pjrt_api_version.minor_version < 96) {
+    return;
+  }
+  PJRT_LoadedExecutable_AddressableDeviceLogicalIds_Args args;
+  args.struct_size =
+      PJRT_LoadedExecutable_AddressableDeviceLogicalIds_Args_STRUCT_SIZE;
+  args.extension_start = nullptr;
+  args.executable = c_loaded_executable();
+  args.addressable_device_logical_ids = nullptr;
+  args.num_addressable_device_logical_ids = 0;
+
+  const PJRT_Api* api = pjrt_c_api();
+  if (api->PJRT_LoadedExecutable_AddressableDeviceLogicalIds == nullptr) {
+    return;
+  }
+  pjrt::LogFatalIfPjrtError(
+      api->PJRT_LoadedExecutable_AddressableDeviceLogicalIds(&args), api);
+
+  if (args.num_addressable_device_logical_ids == 0 ||
+      args.addressable_device_logical_ids == nullptr) {
+    return;
+  }
+  addressable_device_logical_ids_.reserve(
+      args.num_addressable_device_logical_ids);
+
+  for (size_t i = 0; i < args.num_addressable_device_logical_ids; ++i) {
+    PJRT_LogicalDeviceIds id = args.addressable_device_logical_ids[i];
+    addressable_device_logical_ids_.push_back(
+        LogicalDeviceIds{id.replica, id.partition});
+  }
+}
+
 void PjRtCApiLoadedExecutable::InitDevices() {
   PJRT_LoadedExecutable_AddressableDevices_Args args;
   args.struct_size = PJRT_LoadedExecutable_AddressableDevices_Args_STRUCT_SIZE;
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 81098544246..51052421965 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
@@ -823,6 +823,11 @@ class PjRtCApiLoadedExecutable : public PjRtLoadedExecutable {
     return *device_assignment_;
   }
 
+  absl::Span<const LogicalDeviceIds> addressable_device_logical_ids()
+      const override {
+    return addressable_device_logical_ids_;
+  }
+
   absl::Span<PjRtDevice* const> addressable_devices() const override {
     return addressable_devices_;
   }
@@ -929,6 +934,7 @@ class PjRtCApiLoadedExecutable : public PjRtLoadedExecutable {
   }
 
  private:
+  std::vector<LogicalDeviceIds> addressable_device_logical_ids_;
   // Groups data needed to support send/recv execution callbacks.
   struct HloOutputCallbackState {
     // Key is {replica_id, partition_id}.
@@ -1007,6 +1013,7 @@ class PjRtCApiLoadedExecutable : public PjRtLoadedExecutable {
   std::vector<PjRtDevice*> addressable_devices_;
   std::unique_ptr<const DeviceAssignment> device_assignment_;
 
+  void InitAddressableDeviceLogicalIds();
   void InitDevices();
   void InitDeviceAssignment();
 };
diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc
index 202bdc7abe9..d5d72885516 100644
--- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc
+++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client_test.cc
@@ -1060,6 +1060,36 @@ ENTRY Identity() -> f32[2, 2] {
   EXPECT_THAT(poison_result, IsOkAndHolds(false));
 }
 
+TEST(PjRtCApiClientTest, AddressableDeviceLogicalIds) {
+  SetUpCpuPjRtApi();
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtClient> client,
+                       GetCApiClient("cpu"));
+  ASSERT_GT(client->addressable_devices().size(), 1);
+
+  XlaBuilder builder("Identity");
+  Shape shape = ShapeUtil::MakeShape(S32, {2, 3});
+  XlaOp input = Parameter(&builder, 0, shape, "input");
+  XlaComputation computation = builder.Build(input).value();
+
+  DeviceAssignment device_assignment(1, 2);
+  device_assignment(0, 0) = 0;
+  device_assignment(0, 1) = 1;
+
+  CompileOptions options;
+  options.executable_build_options.set_device_assignment(device_assignment);
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtLoadedExecutable> executable,
+                       client->CompileAndLoad(computation, options));
+
+  absl::Span<const PjRtLoadedExecutable::LogicalDeviceIds> logical_ids =
+      executable->addressable_device_logical_ids();
+  ASSERT_EQ(logical_ids.size(), 2);
+  EXPECT_EQ(logical_ids[0].replica, 0);
+  EXPECT_EQ(logical_ids[0].partition, 0);
+  EXPECT_EQ(logical_ids[1].replica, 0);
+  EXPECT_EQ(logical_ids[1].partition, 1);
+}
+
 TEST(PjRtCApiClientTest, Bitcast) {
   SetUpCpuPjRtApi();
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtClient> client,
diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.h b/third_party/xla/xla/pjrt/common_pjrt_client.h
index 261ad93d8ba..2eeffdb8254 100644
--- a/third_party/xla/xla/pjrt/common_pjrt_client.h
+++ b/third_party/xla/xla/pjrt/common_pjrt_client.h
@@ -654,6 +654,11 @@ class CommonPjRtLoadedExecutable : public PjRtLoadedExecutable {
     return *device_assignment_;
   }
 
+  absl::Span<const LogicalDeviceIds> addressable_device_logical_ids()
+      const override {
+    return addressable_device_logical_ids_;
+  }
+
   void Delete() override {
     if (load_state_ != nullptr) {
       load_state_->Delete();
diff --git a/third_party/xla/xla/pjrt/interpreter/interpreter_client.cc b/third_party/xla/xla/pjrt/interpreter/interpreter_client.cc
index 4a6a5b61c98..e8f7ca7e7b8 100644
--- a/third_party/xla/xla/pjrt/interpreter/interpreter_client.cc
+++ b/third_party/xla/xla/pjrt/interpreter/interpreter_client.cc
@@ -162,12 +162,15 @@ InterpreterLoadedExecutable::InterpreterLoadedExecutable(
     std::shared_ptr<InterpreterExecutable> executable,
     std::unique_ptr<HloEvaluatorInterface> hlo_evaluator,
     std::shared_ptr<DeviceAssignment> device_assignment,
+    std::vector<LogicalDeviceIds> addressable_device_logical_ids,
     std::vector<PjRtDevice*> addressable_devices)
     : client_(ABSL_DIE_IF_NULL(client)),
       name_(ABSL_DIE_IF_NULL(executable)->name()),
       executable_(std::move(executable)),
       hlo_evaluator_(std::move(hlo_evaluator)),
       device_assignment_(std::move(device_assignment)),
+      addressable_device_logical_ids_(
+          std::move(addressable_device_logical_ids)),
       addressable_devices_(std::move(addressable_devices)) {
   if (executable_ && executable_->dynamic_dimension_inference().has_value()) {
     hlo_evaluator_->set_dynamic_dimension_inference(
@@ -526,12 +529,19 @@ absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> InterpreterClient::Load(
                         "num_replicas: %d, num_partitions: %d",
                         num_replicas, num_partitions));
   }
+  std::vector<PjRtLoadedExecutable::LogicalDeviceIds>
+      addressable_device_logical_ids;
   std::vector<PjRtDevice*> addressable_devices;
+  PjRtLoadedExecutable::LogicalDeviceIds logical_device_ids;
+  logical_device_ids.replica = 0;
+  logical_device_ids.partition = 0;
+  addressable_device_logical_ids.push_back(std::move(logical_device_ids));
   addressable_devices.push_back(&interpreter_device_);
 
   return std::make_unique<InterpreterLoadedExecutable>(
       this, std::move(shared_interpreter_executable), std::move(evaluator),
-      std::move(device_assignment), std::move(addressable_devices));
+      std::move(device_assignment), std::move(addressable_device_logical_ids),
+      std::move(addressable_devices));
 }
 
 absl::StatusOr<std::unique_ptr<PjRtExecutable>>
diff --git a/third_party/xla/xla/pjrt/interpreter/interpreter_client.h b/third_party/xla/xla/pjrt/interpreter/interpreter_client.h
index 5a6ecf597db..2e9bda4b9e9 100644
--- a/third_party/xla/xla/pjrt/interpreter/interpreter_client.h
+++ b/third_party/xla/xla/pjrt/interpreter/interpreter_client.h
@@ -297,6 +297,7 @@ class InterpreterLoadedExecutable final : public PjRtLoadedExecutable {
       std::shared_ptr<InterpreterExecutable> executable,
       std::unique_ptr<HloEvaluatorInterface> hlo_evaluator,
       std::shared_ptr<DeviceAssignment> device_assignment,
+      std::vector<LogicalDeviceIds> addressable_device_logical_ids,
       std::vector<PjRtDevice*> addressable_devices);
 
   InterpreterExecutable* GetExecutable() const override;
@@ -320,6 +321,11 @@ class InterpreterLoadedExecutable final : public PjRtLoadedExecutable {
     return *device_assignment_;
   }
 
+  absl::Span<const LogicalDeviceIds> addressable_device_logical_ids()
+      const override {
+    return addressable_device_logical_ids_;
+  }
+
   absl::Span<PjRtDevice* const> addressable_devices() const override {
     return addressable_devices_;
   }
@@ -355,6 +361,7 @@ class InterpreterLoadedExecutable final : public PjRtLoadedExecutable {
   std::shared_ptr<InterpreterExecutable> executable_ ABSL_GUARDED_BY(mutex_);
   std::unique_ptr<HloEvaluatorInterface> hlo_evaluator_ ABSL_GUARDED_BY(mutex_);
   std::shared_ptr<DeviceAssignment> device_assignment_;
+  std::vector<LogicalDeviceIds> addressable_device_logical_ids_;
   std::vector<PjRtDevice*> addressable_devices_;
 };
 
diff --git a/third_party/xla/xla/pjrt/pjrt_client.h b/third_party/xla/xla/pjrt/pjrt_client.h
index 63aa3f45d5d..8ad889abe2d 100644
--- a/third_party/xla/xla/pjrt/pjrt_client.h
+++ b/third_party/xla/xla/pjrt/pjrt_client.h
@@ -1411,12 +1411,21 @@ class PjRtLoadedExecutable {
   virtual absl::StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
   GetCostAnalysis() const;
 
+  // The replica and partition indices of device_assignment to be run by this
+  // client. On single-host platforms without partitioning, this is all replicas
+  // (i.e. addressable_device_logical_ids_[i] = (i, 0)), but this may not be the
+  // case on multi-host platforms. If there are 4 replicas and 2 partitions on a
+  // single host platform, size of addressable_device_logical_ids_ is 4*2 = 8.
   struct LogicalDeviceIds {
     int replica;
     int partition;
   };
+  virtual absl::Span<const LogicalDeviceIds> addressable_device_logical_ids()
+      const = 0;
 
   // An addressable_device is one which the client can issue commands to.
+  // addressable_devices()[i] is the Device to which
+  // addressable_device_logical_ids()[i] is assigned.
   virtual absl::Span<PjRtDevice* const> addressable_devices() const = 0;
 
   // Donation Semantics:
diff --git a/third_party/xla/xla/pjrt/tf_pjrt_client.h b/third_party/xla/xla/pjrt/tf_pjrt_client.h
index a934a74f8f6..0943e22bddc 100644
--- a/third_party/xla/xla/pjrt/tf_pjrt_client.h
+++ b/third_party/xla/xla/pjrt/tf_pjrt_client.h
@@ -144,6 +144,10 @@ class TfPjRtExecutable : public PjRtLoadedExecutable {
   const DeviceAssignment& device_assignment() const override {
     return wrapped_->device_assignment();
   }
+  absl::Span<const LogicalDeviceIds> addressable_device_logical_ids()
+      const override {
+    return wrapped_->addressable_device_logical_ids();
+  }
   absl::Span<PjRtDevice* const> addressable_devices() const override {
     return wrapped_->addressable_devices();
   }
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
index 9cffeee12ef..ac8cce570ff 100644
--- a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
@@ -756,10 +756,11 @@ CopyArgumentsToDevice(PjRtClient& client,
     return device_buffers;
   };
 
+  absl::Span<const PjRtLoadedExecutable::LogicalDeviceIds>
+      addressable_device_logical_ids =
+          executable->addressable_device_logical_ids();
   ABSL_ASSIGN_OR_RETURN(std::vector<std::shared_ptr<HloModule>> hlo_modules,
                    executable->GetHloModules());
-  CHECK_EQ(hlo_modules.size(), 1);
-  HloModule* module = hlo_modules.front().get();
 
   for (int i = 0; i < num_addressable_devices; ++i) {
     PjRtDevice* curr_device = addressable_devices[i];
@@ -778,6 +779,11 @@ CopyArgumentsToDevice(PjRtClient& client,
     const std::vector<Literal>& curr_device_arguments =
         arguments.at(source_device_id);
 
+    int executable_idx = hlo_modules.size() == 1
+                             ? 0
+                             : addressable_device_logical_ids[i].partition;
+    HloModule* module = hlo_modules[executable_idx].get();
+
     argument_buffers[i].reserve(curr_device_arguments.size());
     for (int arg_i = 0; arg_i < curr_device_arguments.size(); ++arg_i) {
       const Literal& literal = curr_device_arguments[arg_i];
@@ -807,18 +813,26 @@ CreateUninitializedArgumentsOnDevice(PjRtClient& client,
                                      bool flatten_arguments = false) {
   absl::Span<PjRtDevice* const> addressable_devices =
       executable->addressable_devices();
+  absl::Span<const PjRtLoadedExecutable::LogicalDeviceIds>
+      addressable_device_logical_ids =
+          executable->addressable_device_logical_ids();
   ABSL_ASSIGN_OR_RETURN(std::vector<std::shared_ptr<HloModule>> hlo_modules,
                    executable->GetHloModules());
   VLOG(1) << "FunctionalHloRunner: local_executable count = "
           << hlo_modules.size();
-  CHECK_EQ(hlo_modules.size(), 1);
-  const HloModule& hlo_module = *hlo_modules.front();
 
   LOG(INFO) << "Starting argument buffer shape calculation.";
   PerDeviceShapeVecType argument_shapes_per_device;
+  // This must be true, based on the comment on
+  // PjRtLoadedExecutable::addressable_devices().
+  CHECK_EQ(addressable_devices.size(), addressable_device_logical_ids.size());
   for (int i = 0; i < static_cast<int>(addressable_devices.size()); ++i) {
     VLOG(3) << "Calculating fake argument shapes for device " << i;
     PjRtDevice* device = addressable_devices[i];
+    int executable_idx = hlo_modules.size() == 1
+                             ? 0
+                             : addressable_device_logical_ids[i].partition;
+    const HloModule& hlo_module = *hlo_modules[executable_idx];
 
     std::vector<Shape> argument_shapes;
     if (flatten_arguments) {
@@ -902,12 +916,13 @@ CreateArgumentsOnDevice(PjRtClient& client,
   size_t num_addressable_devices = addressable_devices.size();
 
   PerDeviceLiteralVecType per_device_argument_literals;
+  absl::Span<const PjRtLoadedExecutable::LogicalDeviceIds>
+      addressable_device_logical_ids =
+          executable->addressable_device_logical_ids();
   ABSL_ASSIGN_OR_RETURN(std::vector<std::shared_ptr<HloModule>> hlo_modules,
                    executable->GetHloModules());
   VLOG(1) << "FunctionalHloRunner: local_executable count = "
           << hlo_modules.size();
-  CHECK_EQ(hlo_modules.size(), 1);
-  HloModule* my_hlo_module = hlo_modules.front().get();
 
   const bool kUseRandomInputs = running_options.module_argument_mode ==
                                     ModuleArgumentMode::kUseRandomInputs ||
@@ -934,6 +949,10 @@ CreateArgumentsOnDevice(PjRtClient& client,
     VLOG(3) << "Creating fake arguments for device " << i;
     LiteralVec& argument_literals =
         per_device_argument_literals[addressable_devices[i]->id()];
+    int executable_idx = hlo_modules.size() == 1
+                             ? 0
+                             : addressable_device_logical_ids[i].partition;
+    HloModule* my_hlo_module = hlo_modules[executable_idx].get();
     if (flatten_arguments) {
       ABSL_RETURN_IF_ERROR(EnsureSingleTupleForFlattening(*my_hlo_module));
     }

From 676d37492dfde59b34aed09bdabed9bf1903056e Mon Sep 17 00:00:00 2001
From: Penporn Koanantakool <penporn@google.com>
Date: Mon, 12 Oct 2026 20:05:09 -0700
Subject: [PATCH 0640/1000] [XLA:HLO] Update `hlo_dump_ui_bin_sanitized.js`
 golden file.

PiperOrigin-RevId: 998320462
---
 .../xla/tools/hlo_dump/assets/hlo_dump_ui_bin_sanitized.js    | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/tools/hlo_dump/assets/hlo_dump_ui_bin_sanitized.js b/third_party/xla/xla/tools/hlo_dump/assets/hlo_dump_ui_bin_sanitized.js
index 8f5c133233c..696882a7344 100644
--- a/third_party/xla/xla/tools/hlo_dump/assets/hlo_dump_ui_bin_sanitized.js
+++ b/third_party/xla/xla/tools/hlo_dump/assets/hlo_dump_ui_bin_sanitized.js
@@ -17,8 +17,8 @@ limitations under the License.
 /* eslint-disable */
 var $jscomp=$jscomp||{};$jscomp.scope={};$jscomp.ASSUME_ES2020=!0;$jscomp.ASSUME_ES6=!0;$jscomp.ASSUME_ES5=!0;$jscomp.ASSUME_NO_NATIVE_MAP=!1;$jscomp.ASSUME_NO_NATIVE_SET=!1;$jscomp.ISOLATE_POLYFILLS=!1;$jscomp.FORCE_POLYFILL_PROMISE=!1;$jscomp.FORCE_POLYFILL_PROMISE_WHEN_NO_UNHANDLED_REJECTION=!1;$jscomp.INSTRUMENT_ASYNC_CONTEXT=!0;
 $jscomp.defineProperty=$jscomp.ASSUME_ES5||typeof Object.defineProperties=="function"?Object.defineProperty:function(a,b,c){if(a==Array.prototype||a==Object.prototype)return a;a[b]=c.value;return a};$jscomp.getGlobal=function(a){a=["object"==typeof globalThis&&globalThis,a,"object"==typeof window&&window,"object"==typeof self&&self,"object"==typeof global&&global];for(var b=0;b<a.length;++b){var c=a[b];if(c&&c.Math==Math)return c}throw Error("Cannot find global object");};
-$jscomp.global=$jscomp.ASSUME_ES2020?globalThis:$jscomp.getGlobal(this);$jscomp.IS_SYMBOL_NATIVE=typeof Symbol==="function"&&typeof Symbol("x")==="symbol";$jscomp.TRUST_ES6_POLYFILLS=!$jscomp.ISOLATE_POLYFILLS||$jscomp.IS_SYMBOL_NATIVE;$jscomp.polyfills={};$jscomp.propertyToPolyfillSymbol={};$jscomp.POLYFILL_PREFIX="$jscp$";var JSCompiler_inline_result$jscomp$0,classes$jscomp$inline_20="Int8 Uint8 Uint8Clamped Int16 Uint16 Int32 Uint32 Float32 Float64".split(" ");
-$jscomp.global.BigInt64Array&&(classes$jscomp$inline_20.push("BigInt64"),classes$jscomp$inline_20.push("BigUint64"));JSCompiler_inline_result$jscomp$0=classes$jscomp$inline_20;$jscomp.TYPED_ARRAY_CLASSES=JSCompiler_inline_result$jscomp$0;$jscomp.polyfillTypedArrayMethod=function(a,b,c,d){if(b)for(var e=0;e<$jscomp.TYPED_ARRAY_CLASSES.length;e++){var f=$jscomp.TYPED_ARRAY_CLASSES[e]+"Array.prototype."+a;$jscomp.ISOLATE_POLYFILLS?$jscomp.polyfillIsolated(f,b,c,d):$jscomp.polyfillUnisolated(f,b,c,d)}};
+$jscomp.global=$jscomp.ASSUME_ES2020?globalThis:$jscomp.getGlobal(this);$jscomp.IS_SYMBOL_NATIVE=typeof Symbol==="function"&&typeof Symbol("x")==="symbol";$jscomp.TRUST_ES6_POLYFILLS=!$jscomp.ISOLATE_POLYFILLS||$jscomp.IS_SYMBOL_NATIVE;$jscomp.polyfills={};$jscomp.propertyToPolyfillSymbol={};$jscomp.POLYFILL_PREFIX="$jscp$";var JSCompiler_inline_result$jscomp$0,classes$jscomp$inline_17="Int8 Uint8 Uint8Clamped Int16 Uint16 Int32 Uint32 Float32 Float64".split(" ");
+$jscomp.global.BigInt64Array&&(classes$jscomp$inline_17.push("BigInt64"),classes$jscomp$inline_17.push("BigUint64"));JSCompiler_inline_result$jscomp$0=classes$jscomp$inline_17;$jscomp.TYPED_ARRAY_CLASSES=JSCompiler_inline_result$jscomp$0;$jscomp.polyfillTypedArrayMethod=function(a,b,c,d){if(b)for(var e=0;e<$jscomp.TYPED_ARRAY_CLASSES.length;e++){var f=$jscomp.TYPED_ARRAY_CLASSES[e]+"Array.prototype."+a;$jscomp.ISOLATE_POLYFILLS?$jscomp.polyfillIsolated(f,b,c,d):$jscomp.polyfillUnisolated(f,b,c,d)}};
 $jscomp.polyfill=function(a,b,c,d){b&&($jscomp.ISOLATE_POLYFILLS?$jscomp.polyfillIsolated(a,b,c,d):$jscomp.polyfillUnisolated(a,b,c,d))};$jscomp.polyfillUnisolated=function(a,b){var c=$jscomp.global;a=a.split(".");for(var d=0;d<a.length-1;d++){var e=a[d];if(!(e in c))return;c=c[e]}a=a[a.length-1];d=c[a];b=b(d);b!=d&&b!=null&&$jscomp.defineProperty(c,a,{configurable:!0,writable:!0,value:b})};
 $jscomp.polyfillIsolated=function(a,b,c){var d=a.split("."),e=d.length===1;a=d[0];a=!e&&a in $jscomp.polyfills?$jscomp.polyfills:$jscomp.global;for(var f=0;f<d.length-1;f++){var g=d[f];if(!(g in a))return;a=a[g]}d=d[d.length-1];c=$jscomp.IS_SYMBOL_NATIVE&&c==="es6"?a[d]:null;b=b(c);b!=null&&(e?$jscomp.defineProperty($jscomp.polyfills,d,{configurable:!0,writable:!0,value:b}):b!==c&&($jscomp.propertyToPolyfillSymbol[d]===void 0&&(e=Math.random()*1E9>>>0,$jscomp.propertyToPolyfillSymbol[d]=$jscomp.IS_SYMBOL_NATIVE?
 $jscomp.global.Symbol(d):$jscomp.POLYFILL_PREFIX+e+"$"+d),e=$jscomp.propertyToPolyfillSymbol[d],$jscomp.defineProperty(a,e,{configurable:!0,writable:!0,value:b})))};$jscomp.asyncContextEnter=function(){return $jscomp.asyncContextIdentity};$jscomp.asyncContextStart=function(){return $jscomp.asyncContextEnter};$jscomp.asyncContextIdentity=function(a){return a};

From 486b5049728e633c0904a3008a654c8ea8bf911f Mon Sep 17 00:00:00 2001
From: Sean Talts <talts@google.com>
Date: Mon, 12 Oct 2026 20:12:44 -0700
Subject: [PATCH 0641/1000] [XLA:CPU] Raise default
 xla_cpu_small_while_loop_byte_threshold to 1000000000

While loops consisting of many small operations incur high per-thunk dispatch
and concurrency overheads when not outlined by SmallWhileLoopHoistingPass.
Loops containing custom calls, collectives, scatter, sort, FFT, or custom
fusions are already excluded by SmallWhileLoopHoistingPass.

Raise the default byte threshold from 1024 to 1000000000.

PiperOrigin-RevId: 998323502
---
 third_party/xla/xla/service/cpu/cpu_options.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/cpu/cpu_options.cc b/third_party/xla/xla/service/cpu/cpu_options.cc
index 1ae6afe2561..b2f91fd2368 100644
--- a/third_party/xla/xla/service/cpu/cpu_options.cc
+++ b/third_party/xla/xla/service/cpu/cpu_options.cc
@@ -102,7 +102,7 @@ absl::StatusOr<int64_t> SmallWhileLoopByteThreshold(
 
   auto itr = extra_options_map.find(kSmallWhileLoopByteThreshold);
   if (itr == extra_options_map.end()) {
-    return 1024;  // Default value.
+    return 1000000000;  // Default value.
   }
 
   int64_t byte_threshold;

From c35f7e21cd2d30b5df29752d5df50ea52d45079f Mon Sep 17 00:00:00 2001
From: Sean Talts <talts@google.com>
Date: Mon, 12 Oct 2026 20:19:16 -0700
Subject: [PATCH 0642/1000] [XLA:PJRT] Generalize interleave/deinterleave
 transpose kernels to small dimensions (K = 2..7)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Follow-up to the initial K=3 interleave/deinterleave transpose kernels addressing review comments.

### Why

The previous change introduced specialized interleave/deinterleave kernels for a hardcoded small stride-1 dimension size of K=3. However, small non-power-of-2 dimensions frequently arise in tensor reshapes, channels-last conversions, and packed layouts where the default block transpose cannot vectorize efficiently.

Based on empirical benchmarking across K in [2, 16], element sizes 1–16 B, and multiple architectures (x86_64 Haswell AVX2 baseline, AMD Milan, Intel Sapphire Rapids, and Arm Neoverse V2):
* Small stride-1 loops with K in [2, 7] see massive performance improvements from vector interleaving/deinterleaving—winning configurations are typically ~5x faster (and up to 20x / -95% latency on small-element K=2, 5, 6).
* Configurations where the default kernel is already competitive or where vectorization falls off (e.g. K >= 8 or 16-byte deinterleave) are left on the default path and remain unaffected.
* Dimension coalescing: K is now resolved from the small side's coalesced loop dimension, correctly catching multi-dimensional small strides like `{2, 3, W} -> {W, 2, 3}` as K = 6.

### Design Choices

* **Binary size control via `WithInterleaveSize`:** Templating the outer `Transpose<>` on K would bloat the binary by +233 KB (+68%), while a naive runtime switch inside the macro-kernel loop caused significant latency regressions (+34%) from per-block lambda allocation and store-forwarding stalls. We resolve this by using a helper method (`WithInterleaveSize`) that dispatches to compile-time K outside the block loop, once per innermost loop. Together with only instantiating kernels for winning (K, element size) pairs, `transpose.o` grows by only +38 KB (+11.2%).
* **Zero regression on existing paths:** The default kernel and memcpy paths are untouched, and `TransposePlan` struct layout preserves original field offsets (storing the interleave size in an existing padding hole).
* **Selection clarity:** Kernel selection logic in `transpose.cc` is factored into clear named conditions, with verified invariants enforced via `DCHECK`.

<details>

<summary>Benchmark Results</summary>

Tested (`--runs 6 --config=benchmark`) on Milan, Sapphire Rapids, Haswell (-march=haswell), and Arm Neoverse V2 (`--cpu=arm`).

#### Target Geomeans (cpu/op; negative is faster)

| Target | AMD Milan | Milan (Haswell) | Intel Sapphire Rapids | SPR (Haswell) | Arm Neoverse V2 |
|---|---|---|---|---|---|
| `transpose_test` | -14.50% | -16.29% | -10.83% | -13.87% | -12.19% |
| `literal_linearizer_benchmark` | -0.10% | +1.08% | -0.17% | +0.22% | -0.36% |

In `literal_linearizer_benchmark`, the 23 cases whose plans use interleave kernels are 15-28% faster across all platforms. Inactive cases on Haswell show minor noise (0.4-1.7%) attributable to code layout shifts since default transpose code paths are byte-identical.

#### Representative Microbenchmark Wins (AMD Milan, default build)

* **K = 2 interleave (uint8):** 76.5 µs -> 3.6 µs (**-95.3%**, ~21x faster)
* **K = 2 deinterleave (uint8):** 48.8 µs -> 20.8 µs (**-57.5%**, ~2.4x faster)
* **K = 5 interleave (uint8):** 157.1 µs -> 15.4 µs (**-90.2%**, ~10x faster)
* **K = 6 interleave (uint8):** 175.8 µs -> 15.1 µs (**-91.4%**, ~11.6x faster)
* **K = 7 interleave (uint8):** 194.6 µs -> 32.7 µs (**-83.2%**, ~6x faster)
* **Coalesced {2,3} interleave (uint8):** 248.3 µs -> 15.2 µs (**-93.9%**, ~16x faster)
* **K = 2 interleave (uint32):** 73.5 µs -> 14.0 µs (**-81.0%**, ~5.3x faster)
* **K = 5 interleave (uint32):** 183.9 µs -> 78.8 µs (**-57.1%**, ~2.3x faster)
* **K = 6 interleave (uint32):** 204.9 µs -> 92.9 µs (**-54.7%**, ~2.2x faster)
* **K = 7 interleave (uint32):** 230.2 µs -> 113.4 µs (**-50.8%**, ~2.0x faster)
* **CHW 3x256x4 interleave (uint8):** 4.05 µs -> 0.11 µs (**-97.2%**, ~35x faster)

</details>

PiperOrigin-RevId: 998326217
---
 third_party/xla/xla/pjrt/transpose.cc        | 217 +++++++++++++++----
 third_party/xla/xla/pjrt/transpose.h         |  14 +-
 third_party/xla/xla/pjrt/transpose_kernels.h |   6 +-
 third_party/xla/xla/pjrt/transpose_test.cc   | 153 +++++++++++--
 4 files changed, 321 insertions(+), 69 deletions(-)

diff --git a/third_party/xla/xla/pjrt/transpose.cc b/third_party/xla/xla/pjrt/transpose.cc
index c3f57b33ee9..9868c7817e9 100644
--- a/third_party/xla/xla/pjrt/transpose.cc
+++ b/third_party/xla/xla/pjrt/transpose.cc
@@ -187,12 +187,88 @@ struct TransposePlan::Node {
   bool is_inner_dim_in_b = false;
 };
 
-// Size of the stride-1 dimension handled by the interleave kernels.
-constexpr int kInterleaveSize = 3;
+// Returns whether to use the `kind` interleave kernel for K = `k` and
+// `elem_size`-byte elements.
+//
+// Measured against the default kernel on 2026-10-11 with the
+// BM_TransposeSmallDim_* benchmarks in transpose_test.cc on AMD Milan, Intel
+// Sapphire Rapids and Arm Neoverse V2. K >= 8 was much slower on some of these.
+// K = 4 was faster only for interleave of 1- and 2-byte elements. Deinterleave
+// beyond the per-size K limits was slower on at least one platform. Among the
+// selected cases, the only slowdown was 16-byte interleave at K = 6 and 7 on
+// Sapphire Rapids (up to 6%).
+constexpr bool InterleaveKernelIsFaster(TransposePlan::InnerKernelKind kind,
+                                        int64_t k, int64_t elem_size) {
+  if (kind == TransposePlan::InnerKernelKind::kInterleave) {
+    return k >= 2 && k < 8 && (k != 4 || elem_size <= 2);
+  }
+  int64_t max_k = 0;
+  switch (elem_size) {
+    case 1:
+    case 2:
+      max_k = 7;
+      break;
+    case 4:
+      max_k = 5;
+      break;
+    case 8:
+      max_k = 3;
+      break;
+  }
+  return k >= 2 && k != 4 && k <= max_k;
+}
+
+// Maps a runtime K to a compile-time constant for `fn`, instantiating `fn` only
+// for the K that InterleaveKernelIsFaster can select for `kind` and T, so
+// kernels that are never used are not compiled.
+template <typename T, TransposePlan::InnerKernelKind kind, typename Fn>
+ABSL_ATTRIBUTE_ALWAYS_INLINE void WithInterleaveSize(int64_t k, Fn fn) {
+  const auto call = [&](auto const_k) {
+    if constexpr (InterleaveKernelIsFaster(kind, const_k, sizeof(T))) {
+      fn(const_k);
+    } else {
+      LOG(FATAL) << "Unreachable: interleave kernel with K=" << k
+                 << " sizeof(T)=" << sizeof(T);
+    }
+  };
+  switch (k) {
+    case 2:
+      call(std::integral_constant<int, 2>{});
+      break;
+    case 3:
+      call(std::integral_constant<int, 3>{});
+      break;
+    case 4:
+      call(std::integral_constant<int, 4>{});
+      break;
+    case 5:
+      call(std::integral_constant<int, 5>{});
+      break;
+    case 6:
+      call(std::integral_constant<int, 6>{});
+      break;
+    case 7:
+      call(std::integral_constant<int, 7>{});
+      break;
+    default:
+      LOG(FATAL) << "Invalid interleave size " << k;
+  }
+}
+
+// The value of MacroKernel's `k` when K is read from the stride at run time.
+constexpr int kRuntimeInterleaveSize = -1;
+
+// Returns K for the `kind` interleave kernel given the sentinel strides: the
+// small side's K-groups are back to back, so its stride is K elements.
+template <typename T, TransposePlan::InnerKernelKind kind>
+int64_t InterleaveSizeFromStrides(int64_t lda, int64_t ldb) {
+  return (kind == TransposePlan::InnerKernelKind::kInterleave ? ldb : lda) /
+         int64_t{sizeof(T)};
+}
 
 template <typename T, int inner_bs,
           TransposePlan::Transformation transformation,
-          TransposePlan::InnerKernelKind kind>
+          TransposePlan::InnerKernelKind kind, int k = kRuntimeInterleaveSize>
 ABSL_ATTRIBUTE_FUNC_ALIGN(64)
 void MacroKernel(const char* __restrict a, int64_t lda, int outer_bs_a,
                  char* __restrict b, int64_t ldb, int outer_bs_b,
@@ -201,11 +277,28 @@ void MacroKernel(const char* __restrict a, int64_t lda, int outer_bs_a,
             << " outer_bs_a=" << outer_bs_a << " outer_bs_b=" << outer_bs_b
             << " inner_bs=" << inner_bs;
 
+  // `k == kRuntimeInterleaveSize` means the caller has not dispatched on K, so
+  // dispatch here.
   if constexpr (kind == TransposePlan::InnerKernelKind::kInterleave) {
-    InterleaveKernel<T, kInterleaveSize>(a, lda, b, outer_bs_a * inner_bs);
+    if constexpr (k != kRuntimeInterleaveSize) {
+      InterleaveKernel<T, k>(a, lda, b, outer_bs_a * inner_bs);
+    } else {
+      WithInterleaveSize<T, kind>(
+          InterleaveSizeFromStrides<T, kind>(lda, ldb), [&](auto const_k) {
+            InterleaveKernel<T, const_k>(a, lda, b, outer_bs_a * inner_bs);
+          });
+    }
     return;
-  } else if constexpr (kind == TransposePlan::InnerKernelKind::kDeinterleave) {
-    DeinterleaveKernel<T, kInterleaveSize>(a, b, ldb, outer_bs_b * inner_bs);
+  }
+  if constexpr (kind == TransposePlan::InnerKernelKind::kDeinterleave) {
+    if constexpr (k != kRuntimeInterleaveSize) {
+      DeinterleaveKernel<T, k>(a, b, ldb, outer_bs_b * inner_bs);
+    } else {
+      WithInterleaveSize<T, kind>(
+          InterleaveSizeFromStrides<T, kind>(lda, ldb), [&](auto const_k) {
+            DeinterleaveKernel<T, const_k>(a, b, ldb, outer_bs_b * inner_bs);
+          });
+    }
     return;
   }
 
@@ -303,10 +396,25 @@ void Transpose(const char* __restrict a, int outer_bs_a, char* __restrict b,
     const int64_t lda_block = next_node->lda;
     const int64_t ldb_block = next_node->ldb;
     int64_t i;
-    for (i = 0; i < stop; i += inc) {
-      MacroKernel<T, inner_bs, transformation, kind>(
-          a_offset(i), lda_block, outer_bs_a, b_offset(i), ldb_block,
-          outer_bs_b, scratch, bits_per_element);
+    if constexpr (kind == TransposePlan::InnerKernelKind::kInterleave ||
+                  kind == TransposePlan::InnerKernelKind::kDeinterleave) {
+      // Dispatch K once per loop rather than per block; a per-block dispatch
+      // was measurably slower.
+      WithInterleaveSize<T, kind>(
+          InterleaveSizeFromStrides<T, kind>(lda_block, ldb_block),
+          [&](auto const_k) {
+            for (i = 0; i < stop; i += inc) {
+              MacroKernel<T, inner_bs, transformation, kind, const_k>(
+                  a_offset(i), lda_block, outer_bs_a, b_offset(i), ldb_block,
+                  outer_bs_b, scratch, bits_per_element);
+            }
+          });
+    } else {
+      for (i = 0; i < stop; i += inc) {
+        MacroKernel<T, inner_bs, transformation, kind>(
+            a_offset(i), lda_block, outer_bs_a, b_offset(i), ldb_block,
+            outer_bs_b, scratch, bits_per_element);
+      }
     }
     // Handle trailing elements that didn't fit in a complete macrokernel.
     // Only the innermost dimensions have non-trivial outer_bs blocking.
@@ -560,7 +668,7 @@ void TransposePlan::ExecuteTyped(const char* a, char* b,
       break;
     }
     case InnerKernelKind::kInterleave:
-      if constexpr (transformation == Transformation::kNone && sizeof(T) <= 8) {
+      if constexpr (transformation == Transformation::kNone) {
         DCHECK_EQ(inner_block_elems_, 1);
         handle_inner_block_elems(
             std::integral_constant<int, 1>{},
@@ -568,12 +676,11 @@ void TransposePlan::ExecuteTyped(const char* a, char* b,
                                    InnerKernelKind::kInterleave>{});
       } else {
         LOG(FATAL) << "Unreachable: kInterleave with transformation="
-                   << static_cast<int>(transformation)
-                   << " sizeof(T)=" << sizeof(T);
+                   << static_cast<int>(transformation);
       }
       break;
     case InnerKernelKind::kDeinterleave:
-      if constexpr (transformation == Transformation::kNone && sizeof(T) <= 8) {
+      if constexpr (transformation == Transformation::kNone) {
         DCHECK_EQ(inner_block_elems_, 1);
         handle_inner_block_elems(
             std::integral_constant<int, 1>{},
@@ -581,8 +688,7 @@ void TransposePlan::ExecuteTyped(const char* a, char* b,
                                    InnerKernelKind::kDeinterleave>{});
       } else {
         LOG(FATAL) << "Unreachable: kDeinterleave with transformation="
-                   << static_cast<int>(transformation)
-                   << " sizeof(T)=" << sizeof(T);
+                   << static_cast<int>(transformation);
       }
       break;
     case InnerKernelKind::kMemcpy:
@@ -1360,28 +1466,55 @@ absl::Status TransposePlan::Initialize() {
     b_stride1_size = std::min(b_stride1_size, b_dims_[pos_stride1b_in_b]);
   }
 
-  if (!inner_kernel_is_memcpy() && transformation_ == Transformation::kNone &&
-      (elem_size_in_bytes_ == 1 || elem_size_in_bytes_ == 2 ||
-       elem_size_in_bytes_ == 4 || elem_size_in_bytes_ == 8)) {
-    auto loop_a = absl::c_find_if(
+  // For kInterleave and kDeinterleave, the trip count of the large side's
+  // stride-1 loop.
+  int64_t large_loop_size = 0;
+  if (!inner_kernel_is_memcpy() && transformation_ == Transformation::kNone) {
+    const auto loop_a_it = absl::c_find_if(
         loop_order, [](const Loop& l) { return l.is_inner_dim_in_a; });
-    auto loop_b = absl::c_find_if(
+    const auto loop_b_it = absl::c_find_if(
         loop_order, [](const Loop& l) { return l.is_inner_dim_in_b; });
-    if (loop_a != loop_order.end() && loop_b != loop_order.end() &&
-        loop_a != loop_b && loop_a->lda == elem_size_in_bytes_ &&
-        loop_a->ldb == sentinel_ldb_ && loop_b->lda == sentinel_lda_ &&
-        loop_b->ldb == elem_size_in_bytes_) {
-      if (!loop_b->tile_interior && loop_b->tile_size == 1 &&
-          loop_b->dim_size == kInterleaveSize &&
-          b_stride1_size == kInterleaveSize && a_stride1_size >= 8 &&
-          sentinel_ldb_ == kInterleaveSize * elem_size_in_bytes_) {
-        inner_kernel_kind_ = InnerKernelKind::kInterleave;
-      } else if (!loop_a->tile_interior && loop_a->tile_size == 1 &&
-                 loop_a->dim_size == kInterleaveSize &&
-                 a_stride1_size == kInterleaveSize && b_stride1_size >= 8 &&
-                 sentinel_lda_ == kInterleaveSize * elem_size_in_bytes_) {
-        inner_kernel_kind_ = InnerKernelKind::kDeinterleave;
-      }
+    DCHECK(loop_a_it != loop_order.end());
+    DCHECK(loop_b_it != loop_order.end());
+    // loop_a and loop_b iterate over A's and B's stride-1 dimensions (or their
+    // tile interiors).
+    const Loop& loop_a = *loop_a_it;
+    const Loop& loop_b = *loop_b_it;
+    DCHECK_EQ(loop_a.lda, elem_size_in_bytes_);
+    DCHECK_EQ(loop_b.ldb, elem_size_in_bytes_);
+    DCHECK_EQ(loop_a.ldb, sentinel_ldb_);
+    DCHECK_EQ(loop_b.lda, sentinel_lda_);
+    const auto trip_count = [](const Loop& l) {
+      return l.tile_interior ? std::min(l.tile_size, l.dim_size) : l.dim_size;
+    };
+    const int64_t loop_a_size = trip_count(loop_a);
+    const int64_t loop_b_size = trip_count(loop_b);
+    // Returns whether to use the `kind` kernel with K = small.dim_size, where
+    // `small` is the stride-1 loop of the small side.
+    const auto use_kernel = [&](InnerKernelKind kind, const Loop& small,
+                                int64_t group_stride, int64_t large_size) {
+      DCHECK_EQ(small.tile_size == 1, !small.tile_interior);
+      const int64_t k = small.dim_size;
+      // With a partial tile, the trip count is not a constant K.
+      const bool no_partial_tile = !small.tile_interior;
+      // The K-groups are back to back, so the small side is one contiguous run.
+      const bool contiguous = group_stride == k * elem_size_in_bytes_;
+      // Faster than the default kernel for this K and element size.
+      const bool faster =
+          InterleaveKernelIsFaster(kind, k, elem_size_in_bytes_);
+      const bool large_side_is_long = large_size >= 8;
+      return no_partial_tile && contiguous && faster && large_side_is_long;
+    };
+    if (use_kernel(InnerKernelKind::kInterleave, loop_b, sentinel_ldb_,
+                   loop_a_size)) {
+      inner_kernel_kind_ = InnerKernelKind::kInterleave;
+      interleave_size_ = loop_b.dim_size;
+      large_loop_size = loop_a_size;
+    } else if (use_kernel(InnerKernelKind::kDeinterleave, loop_a, sentinel_lda_,
+                          loop_b_size)) {
+      inner_kernel_kind_ = InnerKernelKind::kDeinterleave;
+      interleave_size_ = loop_a.dim_size;
+      large_loop_size = loop_b_size;
     }
   }
 
@@ -1413,13 +1546,13 @@ absl::Status TransposePlan::Initialize() {
     if (inner_kernel_kind_ == InnerKernelKind::kInterleave) {
       inner_block_elems_ = 1;
       outer_block_elems_a_ = std::min<int64_t>(
-          a_stride1_size, max_inner_block_elems * kMaxOuterBlockElems);
+          large_loop_size, max_inner_block_elems * kMaxOuterBlockElems);
       outer_block_elems_b_ = 1;
     } else if (inner_kernel_kind_ == InnerKernelKind::kDeinterleave) {
       inner_block_elems_ = 1;
       outer_block_elems_a_ = 1;
       outer_block_elems_b_ = std::min<int64_t>(
-          b_stride1_size, max_inner_block_elems * kMaxOuterBlockElems);
+          large_loop_size, max_inner_block_elems * kMaxOuterBlockElems);
     } else {
       inner_block_elems_ = max_inner_block_elems;
       while (inner_block_elems_ > std::min(a_stride1_size, b_stride1_size)) {
@@ -1634,7 +1767,7 @@ int TransposePlan::ChooseParallelizationStrategy(
   int64_t acc = elem_size_in_bytes_;
   if (inner_kernel_kind_ == InnerKernelKind::kInterleave ||
       inner_kernel_kind_ == InnerKernelKind::kDeinterleave) {
-    acc *= kInterleaveSize * inner_block_elems_ * outer_block_elems_a_ *
+    acc *= interleave_size_ * inner_block_elems_ * outer_block_elems_a_ *
            outer_block_elems_b_;
   } else if (!inner_kernel_is_memcpy()) {
     acc *= inner_block_elems_ * inner_block_elems_ * outer_block_elems_a_ *
@@ -1907,7 +2040,7 @@ std::string TransposePlan::ToString() const {
       transformation_str = absl::StrCat("pack(", bits_per_element_, ")");
       break;
   }
-  absl::string_view inner_kernel_str;
+  std::string inner_kernel_str;
   switch (inner_kernel_kind_) {
     case InnerKernelKind::kDefault:
       inner_kernel_str = "default";
@@ -1916,10 +2049,10 @@ std::string TransposePlan::ToString() const {
       inner_kernel_str = "memcpy";
       break;
     case InnerKernelKind::kInterleave:
-      inner_kernel_str = "interleave";
+      inner_kernel_str = absl::StrCat("interleave(", interleave_size_, ")");
       break;
     case InnerKernelKind::kDeinterleave:
-      inner_kernel_str = "deinterleave";
+      inner_kernel_str = absl::StrCat("deinterleave(", interleave_size_, ")");
       break;
   }
   return absl::StrFormat(
diff --git a/third_party/xla/xla/pjrt/transpose.h b/third_party/xla/xla/pjrt/transpose.h
index 5e46dbfde70..9ffca077675 100644
--- a/third_party/xla/xla/pjrt/transpose.h
+++ b/third_party/xla/xla/pjrt/transpose.h
@@ -103,13 +103,13 @@ class TransposePlan {
     kDefault,
     // A and B have the same stride-1 dimension, so the inner loop is a memcpy.
     kMemcpy,
-    // B's stride-1 dimension has size 3 and A's stride-1 dimension is large:
-    // three contiguous rows of A, `lda` bytes apart, are merged element by
-    // element into one contiguous run of B, i.e. b[3 * i + k] = a[k][i].
+    // B's stride-1 dimension has a small size K and A's is large, each after
+    // being coalesced with any dimensions contiguous with it in both A and B:
+    // K rows of A, `lda` bytes apart, are merged element by element into one
+    // contiguous run of B, i.e. b[K * i + k] = a[k][i].
     kInterleave,
-    // The reverse of kInterleave: A's stride-1 dimension has size 3 and one
-    // contiguous run of A is split into three rows of B, `ldb` bytes apart,
-    // i.e. b[k][i] = a[3 * i + k].
+    // The reverse of kInterleave: one contiguous run of A is split into K rows
+    // of B, `ldb` bytes apart, i.e. b[k][i] = a[K * i + k].
     kDeinterleave,
   };
 
@@ -376,6 +376,8 @@ class TransposePlan {
   bool a_is_tiled_ = false;
   bool b_is_tiled_ = false;
   bool is_identity_permutation_ = false;
+  // K for kInterleave and kDeinterleave. Placed here to fill padding.
+  int interleave_size_ = 0;
 
   // Per-chunk loop nests. Each loop nest has its own start/end bounds
   // representing one chunk of the work.
diff --git a/third_party/xla/xla/pjrt/transpose_kernels.h b/third_party/xla/xla/pjrt/transpose_kernels.h
index b12193e6efc..b4016543d26 100644
--- a/third_party/xla/xla/pjrt/transpose_kernels.h
+++ b/third_party/xla/xla/pjrt/transpose_kernels.h
@@ -1118,9 +1118,9 @@ struct TransposeMicroKernel {
 };
 
 // Kernels for InnerKernelKind::kInterleave and kDeinterleave over `n` elements
-// of each of the K rows. They are written as scalar loops because, with K a
-// compile-time constant, clang vectorizes them into structured loads/stores
-// (ld3/st3 on NEON, shuffle sequences on x86).
+// of each of the K rows. They are written as scalar loops with a compile-time
+// K for clang to vectorize, e.g. with NEON structured loads/stores (ld2/ld3,
+// st2/st3) for K = 2, 3 and shuffles on x86.
 template <typename T, int K>
 void InterleaveKernel(const char* __restrict a, int64_t lda, char* __restrict b,
                       int64_t n) {
diff --git a/third_party/xla/xla/pjrt/transpose_test.cc b/third_party/xla/xla/pjrt/transpose_test.cc
index 0efc1cedd41..520b450ffb8 100644
--- a/third_party/xla/xla/pjrt/transpose_test.cc
+++ b/third_party/xla/xla/pjrt/transpose_test.cc
@@ -720,7 +720,12 @@ std::vector<TransposeTestCase> GetTransposeTestCases() {
       TransposeTestCase(/*dims=*/{257, 255, 3}, /*permutation=*/{2, 0, 1}),
       TransposeTestCase(/*dims=*/{5, 3, 64, 64}, /*permutation=*/{2, 3, 0, 1}),
       TransposeTestCase(/*dims=*/{5, 3, 64, 64}, /*permutation=*/{0, 2, 3, 1}),
-      TransposeTestCase(/*dims=*/{5, 64, 64, 3}, /*permutation=*/{0, 3, 1, 2})};
+      TransposeTestCase(/*dims=*/{5, 64, 64, 3}, /*permutation=*/{0, 3, 1, 2}),
+
+      // Other small stride-1 sizes, including one coalesced from 2 x 3.
+      TransposeTestCase(/*dims=*/{5, 256, 256}, /*permutation=*/{1, 2, 0}),
+      TransposeTestCase(/*dims=*/{256, 256, 6}, /*permutation=*/{2, 0, 1}),
+      TransposeTestCase(/*dims=*/{2, 3, 64, 64}, /*permutation=*/{2, 3, 0, 1})};
   return cases;
 }
 
@@ -1330,6 +1335,51 @@ static void* benchmarks = []() {
   return nullptr;
 }();
 
+// Sweeps the size K of a small stride-1 dimension on either side.
+static void* small_dim_benchmarks = []() {
+  using BenchmarkFn =
+      void (*)(const TransposeTestCase&, int, testing::benchmark::State&);
+  const std::vector<std::pair<std::string, BenchmarkFn>> variants = {
+      {"u8", BM_Transpose<uint8_t>},         {"u16", BM_Transpose<uint16_t>},
+      {"u32", BM_Transpose<uint32_t>},       {"u64", BM_Transpose<uint64_t>},
+      {"u128", BM_Transpose<absl::uint128>},
+  };
+  std::vector<std::pair<std::string, TransposeTestCase>> cases;
+  for (int64_t k = 2; k <= 16; ++k) {
+    cases.push_back({absl::StrCat("K_", k, "_interleave"),
+                     TransposeTestCase(/*dims=*/{k, 256, 256},
+                                       /*permutation=*/{1, 2, 0})});
+    cases.push_back({absl::StrCat("K_", k, "_deinterleave"),
+                     TransposeTestCase(/*dims=*/{256, 256, k},
+                                       /*permutation=*/{2, 0, 1})});
+  }
+  // The 2 and 3 dimensions are coalesced into K = 6.
+  cases.push_back(
+      {"K_2x3_interleave", TransposeTestCase(/*dims=*/{2, 3, 256, 256},
+                                             /*permutation=*/{2, 3, 0, 1})});
+  cases.push_back(
+      {"K_2x3_deinterleave", TransposeTestCase(/*dims=*/{256, 256, 2, 3},
+                                               /*permutation=*/{2, 3, 0, 1})});
+  // The small W is coalesced with H on the large side.
+  cases.push_back({"HWC_256x4x3_deinterleave",
+                   TransposeTestCase(/*dims=*/{256, 4, 3},
+                                     /*permutation=*/{2, 0, 1})});
+  cases.push_back(
+      {"CHW_3x256x4_interleave", TransposeTestCase(/*dims=*/{3, 256, 4},
+                                                   /*permutation=*/{1, 2, 0})});
+  for (const auto& [type_name, fn] : variants) {
+    for (const auto& [suffix, testcase] : cases) {
+      const std::string name =
+          absl::StrCat("BM_TransposeSmallDim_", type_name, "_", suffix);
+      benchmark::RegisterBenchmark(name.c_str(), [fn = fn, testcase = testcase](
+                                                     benchmark::State& state) {
+        fn(testcase, /*parallelism=*/1, state);
+      });
+    }
+  }
+  return nullptr;
+}();
+
 enum class StrideMode {
   kTight,          // lda = ldb <= 64B (exercises SseSquare / 128-bit kernel)
   kWideSymmetric,  // lda = ldb > 64B (exercises AvxRectangular gather)
@@ -1842,30 +1892,57 @@ TEST(TransposeTest, InterleaveAndDeinterleaveSelection) {
     std::string expected_kernel;
   };
   const SelectionCase cases[] = {
-      {{3, 8}, {1, 0}, 1, {}, {}, {}, "inner_kernel=interleave"},
-      {{3, 8}, {1, 0}, 2, {}, {}, {}, "inner_kernel=interleave"},
-      {{3, 8}, {1, 0}, 4, {}, {}, {}, "inner_kernel=interleave"},
-      {{3, 8}, {1, 0}, 8, {}, {}, {}, "inner_kernel=interleave"},
-      {{3, 8}, {1, 0}, 16, {}, {}, {}, "inner_kernel=default"},
+      {{3, 8}, {1, 0}, 1, {}, {}, {}, "inner_kernel=interleave(3)"},
+      {{3, 8}, {1, 0}, 2, {}, {}, {}, "inner_kernel=interleave(3)"},
+      {{3, 8}, {1, 0}, 4, {}, {}, {}, "inner_kernel=interleave(3)"},
+      {{3, 8}, {1, 0}, 8, {}, {}, {}, "inner_kernel=interleave(3)"},
+      {{3, 8}, {1, 0}, 16, {}, {}, {}, "inner_kernel=interleave(3)"},
       {{3, 7}, {1, 0}, 4, {}, {}, {}, "inner_kernel=default"},
-      {{8, 3}, {1, 0}, 1, {}, {}, {}, "inner_kernel=deinterleave"},
-      {{8, 3}, {1, 0}, 2, {}, {}, {}, "inner_kernel=deinterleave"},
-      {{8, 3}, {1, 0}, 4, {}, {}, {}, "inner_kernel=deinterleave"},
-      {{8, 3}, {1, 0}, 8, {}, {}, {}, "inner_kernel=deinterleave"},
+      {{8, 3}, {1, 0}, 1, {}, {}, {}, "inner_kernel=deinterleave(3)"},
+      {{8, 3}, {1, 0}, 2, {}, {}, {}, "inner_kernel=deinterleave(3)"},
+      {{8, 3}, {1, 0}, 4, {}, {}, {}, "inner_kernel=deinterleave(3)"},
+      {{8, 3}, {1, 0}, 8, {}, {}, {}, "inner_kernel=deinterleave(3)"},
       {{8, 3}, {1, 0}, 16, {}, {}, {}, "inner_kernel=default"},
       {{7, 3}, {1, 0}, 4, {}, {}, {}, "inner_kernel=default"},
+      {{2, 8}, {1, 0}, 4, {}, {}, {}, "inner_kernel=interleave(2)"},
+      {{5, 8}, {1, 0}, 4, {}, {}, {}, "inner_kernel=interleave(5)"},
+      {{6, 8}, {1, 0}, 4, {}, {}, {}, "inner_kernel=interleave(6)"},
+      {{7, 8}, {1, 0}, 4, {}, {}, {}, "inner_kernel=interleave(7)"},
+      {{7, 8}, {1, 0}, 16, {}, {}, {}, "inner_kernel=interleave(7)"},
+      {{4, 8}, {1, 0}, 1, {}, {}, {}, "inner_kernel=interleave(4)"},
+      {{4, 8}, {1, 0}, 2, {}, {}, {}, "inner_kernel=interleave(4)"},
+      {{4, 8}, {1, 0}, 4, {}, {}, {}, "inner_kernel=default"},
+      {{8, 16}, {1, 0}, 4, {}, {}, {}, "inner_kernel=default"},
+      {{8, 2}, {1, 0}, 4, {}, {}, {}, "inner_kernel=deinterleave(2)"},
+      {{8, 2}, {1, 0}, 16, {}, {}, {}, "inner_kernel=default"},
+      {{8, 5}, {1, 0}, 4, {}, {}, {}, "inner_kernel=deinterleave(5)"},
+      {{8, 6}, {1, 0}, 2, {}, {}, {}, "inner_kernel=deinterleave(6)"},
+      {{8, 6}, {1, 0}, 4, {}, {}, {}, "inner_kernel=default"},
+      {{8, 7}, {1, 0}, 2, {}, {}, {}, "inner_kernel=deinterleave(7)"},
+      {{8, 7}, {1, 0}, 4, {}, {}, {}, "inner_kernel=default"},
+      {{8, 4}, {1, 0}, 1, {}, {}, {}, "inner_kernel=default"},
+      {{8, 4}, {1, 0}, 4, {}, {}, {}, "inner_kernel=default"},
+      {{16, 8}, {1, 0}, 4, {}, {}, {}, "inner_kernel=default"},
+      {{2, 3, 8}, {2, 0, 1}, 4, {}, {}, {}, "inner_kernel=interleave(6)"},
+      {{8, 2, 3}, {1, 2, 0}, 2, {}, {}, {}, "inner_kernel=deinterleave(6)"},
+      {{2, 2, 8}, {2, 0, 1}, 1, {}, {}, {}, "inner_kernel=interleave(4)"},
+      {{2, 2, 8}, {2, 0, 1}, 4, {}, {}, {}, "inner_kernel=default"},
+      {{2, 4, 8}, {2, 0, 1}, 4, {}, {}, {}, "inner_kernel=default"},
+      {{3, 2, 4}, {1, 2, 0}, 4, {}, {}, {}, "inner_kernel=interleave(3)"},
+      {{3, 256, 4}, {1, 2, 0}, 4, {}, {}, {}, "inner_kernel=interleave(3)"},
+      {{256, 4, 3}, {2, 0, 1}, 4, {}, {}, {}, "inner_kernel=deinterleave(3)"},
       {{3, 4, 16}, {2, 1, 0}, 4, {}, {}, {}, "inner_kernel=default"},
       {{16, 4, 3}, {2, 1, 0}, 4, {}, {}, {}, "inner_kernel=default"},
-      {{3, 8}, {1, 0}, 4, {32, 4}, {}, {}, "inner_kernel=interleave"},
-      {{3, 8}, {1, 0}, 4, {-32, 4}, {}, {}, "inner_kernel=interleave"},
-      {{8, 3}, {1, 0}, 4, {12, 4}, {}, {}, "inner_kernel=deinterleave"},
+      {{3, 8}, {1, 0}, 4, {32, 4}, {}, {}, "inner_kernel=interleave(3)"},
+      {{3, 8}, {1, 0}, 4, {-32, 4}, {}, {}, "inner_kernel=interleave(3)"},
+      {{8, 3}, {1, 0}, 4, {12, 4}, {}, {}, "inner_kernel=deinterleave(3)"},
       {{2, 8, 3},
        {0, 2, 1},
        4,
        {-128, 12, 4},
        {},
        {},
-       "inner_kernel=deinterleave"},
+       "inner_kernel=deinterleave(3)"},
       {{8, 3}, {1, 0}, 4, {32, 4}, {}, {}, "inner_kernel=default"},
       {{8, 3}, {1, 0}, 4, {-32, 4}, {}, {}, "inner_kernel=default"},
       {{32, 224, 224, 3},
@@ -1874,16 +1951,18 @@ TEST(TransposeTest, InterleaveAndDeinterleaveSelection) {
        {},
        {},
        {8, 128},
-       "inner_kernel=deinterleave"},
+       "inner_kernel=deinterleave(3)"},
       {{32, 3, 224, 224},
        {0, 2, 3, 1},
        4,
        {},
        {8, 128},
        {},
-       "inner_kernel=interleave"},
-      {{6, 8}, {1, 0}, 4, {}, {}, {3}, "inner_kernel=default"},
-      {{8, 6}, {1, 0}, 4, {}, {3}, {}, "inner_kernel=default"},
+       "inner_kernel=interleave(3)"},
+      {{6, 8}, {1, 0}, 4, {}, {}, {3}, "inner_kernel=interleave(6)"},
+      {{8, 6}, {1, 0}, 2, {}, {3}, {}, "inner_kernel=deinterleave(6)"},
+      {{7, 8}, {1, 0}, 4, {}, {}, {3}, "inner_kernel=default"},
+      {{8, 7}, {1, 0}, 4, {}, {3}, {}, "inner_kernel=default"},
   };
   for (const auto& tc : cases) {
     TransposePlan::Options o;
@@ -1908,4 +1987,42 @@ TEST(TransposeTest, InterleaveAndDeinterleaveSelection) {
   }
 }
 
+TEST(TransposeTest, InterleaveAndDeinterleaveAllSizes) {
+  for (int e : {1, 2, 4, 8, 16}) {
+    std::vector<TransposeTestCase> cases = {
+        TransposeTestCase({6, 8}, {1, 0}, {}, /*output_tiling=*/{3}),
+        TransposeTestCase({8, 6}, {1, 0}, /*input_tiling=*/{3}),
+        TransposeTestCase({3, 2, 4}, {1, 2, 0}),
+        TransposeTestCase({16, 4, 3}, {2, 0, 1}),
+        TransposeTestCase({3, 16, 4}, {1, 2, 0}),
+    };
+    for (int64_t h : {1, 3}) {
+      for (int64_t w : {8, 9, 37}) {
+        for (int64_t k = 2; k <= 16; ++k) {
+          cases.push_back(TransposeTestCase({k, h, w}, {1, 2, 0}));
+          cases.push_back(TransposeTestCase({h, w, k}, {2, 0, 1}));
+        }
+        cases.push_back(TransposeTestCase({2, 3, h, w}, {2, 3, 0, 1}));
+        cases.push_back(TransposeTestCase({h, w, 2, 3}, {2, 3, 0, 1}));
+      }
+    }
+    for (int64_t k : {2, 5, 7}) {
+      cases.push_back(TransposeTestCase({k, 9}, {1, 0}, {}, {},
+                                        /*input_striding=*/{-9 * e, e}));
+      cases.push_back(TransposeTestCase({9, k}, {1, 0}, {}, {},
+                                        /*input_striding=*/{k * e, -e}));
+      cases.push_back(TransposeTestCase({k, 37}, {1, 0}, /*input_tiling=*/{8}));
+      cases.push_back(
+          TransposeTestCase({37, k}, {1, 0}, {}, /*output_tiling=*/{8}));
+    }
+    for (const TransposeTestCase& test : cases) {
+      SCOPED_TRACE(absl::StrCat(test.ToString(), ",elem_size=", e));
+      ExecuteTest(test, e, /*parallelism=*/1,
+                  TransposePlan::Transformation::kNone,
+                  TransposePlan::ChunkContiguity::kNone,
+                  /*use_execute_chunk=*/false);
+    }
+  }
+}
+
 }  // namespace xla

From 661555332c7061a4b8d4c0af88c959a113fd80bb Mon Sep 17 00:00:00 2001
From: Ce Zheng CA <zce@google.com>
Date: Mon, 12 Oct 2026 20:26:02 -0700
Subject: [PATCH 0643/1000] [XLA:CPU] Remove redundant Eigen::Aligned16
 contraction instantiation in dot_lib.h

`TypedMatMul` in `xla/backends/cpu/runtime/dot_lib.h` branches at runtime on 16-byte pointer alignment and instantiates both `MatMul<..., Eigen::Aligned16>` and `MatMul<..., Eigen::Unaligned>`. However, `Eigen::TensorContractionEvaluatorBase` (`evalTo`, `evalToAsync`, `evalGemmPartial`), `EvalParallelContext`, and `general_matrix_vector_product` unconditionally wrap the output buffer in `TensorMap<..., Eigen::Unaligned>` and instantiate `TensorContractionInputMapper<..., Eigen::Unaligned>` and `LhsAlignment = Eigen::Unaligned` for all contraction inputs.

Register `Eigen::Unaligned` template specializations of `DirectColMajorAccess` alongside `Eigen::Aligned` in `xla/tsl/framework/contraction/eigen_contraction_kernel.h` so the oneDNN direct column-major access fast path is used for `Eigen::Unaligned` `TensorMap` inputs as well, and use `Eigen::Unaligned` unconditionally in `MatMul`, removing the redundant `Eigen::AlignmentType` template parameter and runtime alignment branch.

```diff
-template <typename LhsType, typename RhsType, typename OutType,
-          Eigen::AlignmentType alignment>
+template <typename LhsType, typename RhsType, typename OutType>
 void MatMul(const Eigen::ThreadPoolDevice* device, OutType* out, LhsType* lhs,
             RhsType* rhs, int64_t m, int64_t n, int64_t k,
             int32_t transpose_lhs, int32_t transpose_rhs, DoneCallback done) {
...
-  const Eigen::TensorMap<Eigen::Tensor<const LhsType, 2>, alignment> a(
+  const Eigen::TensorMap<Eigen::Tensor<const LhsType, 2>, Eigen::Unaligned> a(
       lhs, lhs_rows, lhs_cols);
-  const Eigen::TensorMap<Eigen::Tensor<const RhsType, 2>, alignment> b(
+  const Eigen::TensorMap<Eigen::Tensor<const RhsType, 2>, Eigen::Unaligned> b(
       rhs, rhs_rows, rhs_cols);
-  Eigen::TensorMap<Eigen::Tensor<OutType, 2>, alignment> c(out, m, n);
+  Eigen::TensorMap<Eigen::Tensor<OutType, 2>, Eigen::Unaligned> c(out, m, n);
```

Performance & Memory Impact:
- CPU runtime: Zero regression (all generated contraction and GEMV functions in `dot_lib` emit identical instruction sequences and use the same oneDNN direct column-major access fast path as `Eigen::Aligned16`, while avoiding a 3-pointer alignment branch on entry to `TypedMatMul`).
- Memory usage: Unchanged at runtime; reduces instruction cache footprint and binary size of `dot_lib_onednn` + `dot_thunk` by `-3.57 MB` (`-47.0%`, from `7.60 MB` to `4.03 MB`).

PiperOrigin-RevId: 998328529
---
 .../xla/xla/backends/cpu/runtime/BUILD        |  2 -
 .../xla/xla/backends/cpu/runtime/dot_lib.h    | 40 ++++-------
 .../contraction/eigen_contraction_kernel.h    | 37 +++++-----
 .../eigen_contraction_kernel_test.cc          | 69 +++++++++++++++++++
 4 files changed, 103 insertions(+), 45 deletions(-)

diff --git a/third_party/xla/xla/backends/cpu/runtime/BUILD b/third_party/xla/xla/backends/cpu/runtime/BUILD
index c10be593330..d1eaa64f95d 100644
--- a/third_party/xla/xla/backends/cpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/cpu/runtime/BUILD
@@ -802,7 +802,6 @@ cc_library(
         "ignore_for_dep=third_party/tensorflow/compiler/xla/tsl/framework/contraction/eigen_contraction_kernel.h",
     ]),
     deps = [
-        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/functional:any_invocable",
         "@eigen_archive//:eigen3",
     ],
@@ -825,7 +824,6 @@ cc_library(
     hdrs = ["dot_lib.h"],
     deps = [
         "//xla/tsl/framework/contraction:eigen_contraction_kernel",
-        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/functional:any_invocable",
         "@eigen_archive//:eigen3",
     ],
diff --git a/third_party/xla/xla/backends/cpu/runtime/dot_lib.h b/third_party/xla/xla/backends/cpu/runtime/dot_lib.h
index 386a6ad7a62..461e7a558d7 100644
--- a/third_party/xla/xla/backends/cpu/runtime/dot_lib.h
+++ b/third_party/xla/xla/backends/cpu/runtime/dot_lib.h
@@ -21,7 +21,6 @@ limitations under the License.
 #include <type_traits>
 #include <utility>
 
-#include "absl/base/optimization.h"
 #include "absl/functional/any_invocable.h"
 
 #define EIGEN_USE_THREADS
@@ -34,8 +33,7 @@ namespace xla::cpu::internal {
 using DoneCallback = absl::AnyInvocable<void()>;
 
 // Col-major x Col-major MatMul implementation as Eigen contraction.
-template <typename LhsType, typename RhsType, typename OutType,
-          Eigen::AlignmentType alignment>
+template <typename LhsType, typename RhsType, typename OutType>
 void MatMul(const Eigen::ThreadPoolDevice* device, OutType* out, LhsType* lhs,
             RhsType* rhs, int64_t m, int64_t n, int64_t k,
             int32_t transpose_lhs, int32_t transpose_rhs, DoneCallback done);
@@ -50,8 +48,7 @@ void TypedMatMul(const Eigen::ThreadPoolDevice* device, void* out, void* lhs,
 // TypedMatMul/MatMul implementation details.
 //===----------------------------------------------------------------------===//
 
-template <typename LhsType, typename RhsType, typename OutType,
-          Eigen::AlignmentType alignment>
+template <typename LhsType, typename RhsType, typename OutType>
 void MatMul(const Eigen::ThreadPoolDevice* device, OutType* out, LhsType* lhs,
             RhsType* rhs, int64_t m, int64_t n, int64_t k,
             int32_t transpose_lhs, int32_t transpose_rhs, DoneCallback done) {
@@ -67,11 +64,16 @@ void MatMul(const Eigen::ThreadPoolDevice* device, OutType* out, LhsType* lhs,
     std::swap(rhs_rows, rhs_cols);
   }
 
-  const Eigen::TensorMap<Eigen::Tensor<const LhsType, 2>, alignment> a(
+  // Both Eigen's TensorContractionEvaluatorBase and DirectColMajorAccess in
+  // eigen_contraction_kernel.h support Eigen::Unaligned TensorMap inputs with
+  // the same generated code and oneDNN direct col-major access as
+  // Eigen::Aligned16. Using Eigen::Unaligned avoids instantiating a duplicate
+  // set of contraction kernels for Eigen::Aligned16.
+  const Eigen::TensorMap<Eigen::Tensor<const LhsType, 2>, Eigen::Unaligned> a(
       lhs, lhs_rows, lhs_cols);
-  const Eigen::TensorMap<Eigen::Tensor<const RhsType, 2>, alignment> b(
+  const Eigen::TensorMap<Eigen::Tensor<const RhsType, 2>, Eigen::Unaligned> b(
       rhs, rhs_rows, rhs_cols);
-  Eigen::TensorMap<Eigen::Tensor<OutType, 2>, alignment> c(out, m, n);
+  Eigen::TensorMap<Eigen::Tensor<OutType, 2>, Eigen::Unaligned> c(out, m, n);
 
   typedef typename Eigen::Tensor<LhsType, 2>::DimensionPair DimPair;
   int lhs_contract_dim = transpose_lhs ? 0 : 1;
@@ -108,24 +110,10 @@ template <typename LhsType, typename RhsType, typename OutType>
 void TypedMatMul(const Eigen::ThreadPoolDevice* device, void* out, void* lhs,
                  void* rhs, int64_t m, int64_t n, int64_t k, bool transpose_lhs,
                  bool transpose_rhs, DoneCallback done) {
-  auto is_16_byte_aligned = [](void* ptr) {
-    return reinterpret_cast<uintptr_t>(ptr) % 16 == 0;
-  };
-
-  bool is_aligned = is_16_byte_aligned(lhs) && is_16_byte_aligned(rhs) &&
-                    is_16_byte_aligned(out);
-
-  if (ABSL_PREDICT_TRUE(is_aligned)) {
-    MatMul<LhsType, RhsType, OutType, Eigen::Aligned16>(
-        device, static_cast<OutType*>(out), static_cast<LhsType*>(lhs),
-        static_cast<RhsType*>(rhs), m, n, k, transpose_lhs, transpose_rhs,
-        std::move(done));
-  } else {
-    MatMul<LhsType, RhsType, OutType, Eigen::Unaligned>(
-        device, static_cast<OutType*>(out), static_cast<LhsType*>(lhs),
-        static_cast<RhsType*>(rhs), m, n, k, transpose_lhs, transpose_rhs,
-        std::move(done));
-  }
+  MatMul<LhsType, RhsType, OutType>(
+      device, static_cast<OutType*>(out), static_cast<LhsType*>(lhs),
+      static_cast<RhsType*>(rhs), m, n, k, transpose_lhs, transpose_rhs,
+      std::move(done));
 }
 
 // Declare TypedMatMul template for all supported data types to enable
diff --git a/third_party/xla/xla/tsl/framework/contraction/eigen_contraction_kernel.h b/third_party/xla/xla/tsl/framework/contraction/eigen_contraction_kernel.h
index f8bdc27fc93..57aaab76da1 100644
--- a/third_party/xla/xla/tsl/framework/contraction/eigen_contraction_kernel.h
+++ b/third_party/xla/xla/tsl/framework/contraction/eigen_contraction_kernel.h
@@ -462,35 +462,38 @@ struct DirectColMajorAccess {
 
 #define SIMPLE_TENSOR const Tensor<Scalar, 2, Eigen::ColMajor, StorageIndex>
 
-#define TENSOR_MAP_ROWMAJOR                                               \
+#define TENSOR_MAP_ROWMAJOR(ALIGNMENT)                                    \
   const TensorMap<Tensor<const Scalar, 2, Eigen::RowMajor, StorageIndex>, \
-                  Eigen::Aligned>
+                  ALIGNMENT>
 
-#define TENSOR_MAP_COLMAJOR                                               \
+#define TENSOR_MAP_COLMAJOR(ALIGNMENT)                                    \
   const TensorMap<Tensor<const Scalar, 2, Eigen::ColMajor, StorageIndex>, \
-                  Eigen::Aligned>
+                  ALIGNMENT>
 
-#define TENSOR_MAP_CONST_ROWMAJOR                                   \
-  const TensorMap<Tensor<Scalar, 2, Eigen::RowMajor, StorageIndex>, \
-                  Eigen::Aligned>
+#define TENSOR_MAP_CONST_ROWMAJOR(ALIGNMENT) \
+  const TensorMap<Tensor<Scalar, 2, Eigen::RowMajor, StorageIndex>, ALIGNMENT>
 
-#define TENSOR_MAP_CONST_COLMAJOR                                   \
-  const TensorMap<Tensor<Scalar, 2, Eigen::ColMajor, StorageIndex>, \
-                  Eigen::Aligned>
+#define TENSOR_MAP_CONST_COLMAJOR(ALIGNMENT) \
+  const TensorMap<Tensor<Scalar, 2, Eigen::ColMajor, StorageIndex>, ALIGNMENT>
 
 // This is reshaped convolution filter from `eigen_spatial_convolutions.h`.
-#define TENSOR_RESHAPE                                                        \
+#define TENSOR_RESHAPE(ALIGNMENT)                                             \
   const TensorReshapingOp<                                                    \
       const Eigen::DSizes<StorageIndex, 2>,                                   \
       const TensorMap<Tensor<const Scalar, 4, Eigen::RowMajor, StorageIndex>, \
-                      Eigen::Aligned>>
+                      ALIGNMENT>>
 
 REGISTER_DIRECT_COL_MAJOR_ACCESS(SIMPLE_TENSOR);
-REGISTER_DIRECT_COL_MAJOR_ACCESS(TENSOR_MAP_ROWMAJOR);
-REGISTER_DIRECT_COL_MAJOR_ACCESS(TENSOR_MAP_COLMAJOR);
-REGISTER_DIRECT_COL_MAJOR_ACCESS(TENSOR_MAP_CONST_ROWMAJOR);
-REGISTER_DIRECT_COL_MAJOR_ACCESS(TENSOR_MAP_CONST_COLMAJOR);
-REGISTER_DIRECT_COL_MAJOR_ACCESS(TENSOR_RESHAPE);
+REGISTER_DIRECT_COL_MAJOR_ACCESS(TENSOR_MAP_ROWMAJOR(Eigen::Aligned));
+REGISTER_DIRECT_COL_MAJOR_ACCESS(TENSOR_MAP_COLMAJOR(Eigen::Aligned));
+REGISTER_DIRECT_COL_MAJOR_ACCESS(TENSOR_MAP_CONST_ROWMAJOR(Eigen::Aligned));
+REGISTER_DIRECT_COL_MAJOR_ACCESS(TENSOR_MAP_CONST_COLMAJOR(Eigen::Aligned));
+REGISTER_DIRECT_COL_MAJOR_ACCESS(TENSOR_RESHAPE(Eigen::Aligned));
+REGISTER_DIRECT_COL_MAJOR_ACCESS(TENSOR_MAP_ROWMAJOR(Eigen::Unaligned));
+REGISTER_DIRECT_COL_MAJOR_ACCESS(TENSOR_MAP_COLMAJOR(Eigen::Unaligned));
+REGISTER_DIRECT_COL_MAJOR_ACCESS(TENSOR_MAP_CONST_ROWMAJOR(Eigen::Unaligned));
+REGISTER_DIRECT_COL_MAJOR_ACCESS(TENSOR_MAP_CONST_COLMAJOR(Eigen::Unaligned));
+REGISTER_DIRECT_COL_MAJOR_ACCESS(TENSOR_RESHAPE(Eigen::Unaligned));
 
 #undef SIMPLE_TENSOR
 #undef TENSOR_MAP_ROWMAJOR
diff --git a/third_party/xla/xla/tsl/framework/contraction/eigen_contraction_kernel_test.cc b/third_party/xla/xla/tsl/framework/contraction/eigen_contraction_kernel_test.cc
index 107ad0238a0..77c4f445f0a 100644
--- a/third_party/xla/xla/tsl/framework/contraction/eigen_contraction_kernel_test.cc
+++ b/third_party/xla/xla/tsl/framework/contraction/eigen_contraction_kernel_test.cc
@@ -56,6 +56,38 @@ void RunEigenMatMul(int m, int k, int n) {
   ExpectClose<OutType>(out, expected);
 }
 
+template <typename LhsType, typename RhsType, typename OutType, int Alignment>
+void RunEigenTensorMapMatMul(int m, int k, int n) {
+  Eigen::Tensor<LhsType, 2> lhs_storage(m, k);
+  Eigen::Tensor<RhsType, 2> rhs_storage(k, n);
+  Eigen::Tensor<OutType, 2> out_storage(m, n);
+
+  lhs_storage.setRandom();
+  rhs_storage.setRandom();
+  out_storage.setZero();
+
+  const Eigen::TensorMap<Eigen::Tensor<const LhsType, 2>, Alignment> lhs(
+      lhs_storage.data(), m, k);
+  const Eigen::TensorMap<Eigen::Tensor<const RhsType, 2>, Alignment> rhs(
+      rhs_storage.data(), k, n);
+  Eigen::TensorMap<Eigen::Tensor<OutType, 2>, Alignment> out(out_storage.data(),
+                                                             m, n);
+
+  using DimPair = typename Eigen::Tensor<LhsType, 2>::DimensionPair;
+  std::array<DimPair, 1> dims({DimPair(1, 0)});
+  out = lhs.contract(rhs, dims);
+
+  const Eigen::TensorMap<Eigen::Tensor<const LhsType, 2>, Eigen::Aligned>
+      lhs_aligned(lhs_storage.data(), m, k);
+  const Eigen::TensorMap<Eigen::Tensor<const RhsType, 2>, Eigen::Aligned>
+      rhs_aligned(rhs_storage.data(), k, n);
+  Eigen::Tensor<OutType, 2> expected =
+      lhs_aligned.template cast<OutType>().contract(
+          rhs_aligned.template cast<OutType>(), dims);
+
+  ExpectClose<OutType>(out_storage, expected);
+}
+
 struct EigenContractionKernelTestParams {
   int m;
   int k;
@@ -81,6 +113,43 @@ TEST_P(EigenContractionKernelTest, Float) {
   RunEigenMatMul<float, float, float>(param.m, param.k, param.n);
 }
 
+TEST_P(EigenContractionKernelTest, FloatTensorMapUnaligned) {
+  EigenContractionKernelTestParams param = GetParam();
+  RunEigenTensorMapMatMul<float, float, float, Eigen::Unaligned>(
+      param.m, param.k, param.n);
+}
+
+#if defined(TENSORFLOW_USE_MKLDNN_CONTRACTION_KERNEL)
+TEST(DirectColMajorAccessTest, SupportsAlignedAndUnalignedTensorMap) {
+  using Scalar = float;
+  using StorageIndex = Eigen::Index;
+  using Device = Eigen::DefaultDevice;
+  using ContractT = Eigen::array<StorageIndex, 1>;
+  constexpr int kPacketSize = Eigen::internal::packet_traits<Scalar>::size;
+
+  using AlignedTensorMap = const Eigen::TensorMap<
+      Eigen::Tensor<const Scalar, 2, Eigen::ColMajor, StorageIndex>,
+      Eigen::Aligned>;
+  using UnalignedTensorMap = const Eigen::TensorMap<
+      Eigen::Tensor<const Scalar, 2, Eigen::ColMajor, StorageIndex>,
+      Eigen::Unaligned>;
+
+  using AlignedMapper = Eigen::internal::TensorContractionInputMapper<
+      Scalar, StorageIndex, Eigen::internal::Lhs,
+      Eigen::TensorEvaluator<AlignedTensorMap, Device>, ContractT, ContractT,
+      kPacketSize, /*inner_dim_contiguous=*/true,
+      /*inner_dim_reordered=*/false, Eigen::Unaligned>;
+  using UnalignedMapper = Eigen::internal::TensorContractionInputMapper<
+      Scalar, StorageIndex, Eigen::internal::Lhs,
+      Eigen::TensorEvaluator<UnalignedTensorMap, Device>, ContractT, ContractT,
+      kPacketSize, /*inner_dim_contiguous=*/true,
+      /*inner_dim_reordered=*/false, Eigen::Unaligned>;
+
+  EXPECT_TRUE(Eigen::internal::DirectColMajorAccess<AlignedMapper>::value);
+  EXPECT_TRUE(Eigen::internal::DirectColMajorAccess<UnalignedMapper>::value);
+}
+#endif
+
 INSTANTIATE_TEST_SUITE_P(
     EigenContractionKernelTestSuite, EigenContractionKernelTest,
     testing::ValuesIn<EigenContractionKernelTestParams>({{10, 10, 10},

From 877319c8f5740f5d68ceddcb654d8034e05082c0 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 12 Oct 2026 21:53:55 -0700
Subject: [PATCH 0644/1000] Automated Code Change

PiperOrigin-RevId: 998360322
---
 third_party/xla/xla/backends/cpu/codegen/BUILD              | 1 +
 third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc | 1 +
 2 files changed, 2 insertions(+)

diff --git a/third_party/xla/xla/backends/cpu/codegen/BUILD b/third_party/xla/xla/backends/cpu/codegen/BUILD
index 11fee55599e..34f52742488 100644
--- a/third_party/xla/xla/backends/cpu/codegen/BUILD
+++ b/third_party/xla/xla/backends/cpu/codegen/BUILD
@@ -241,6 +241,7 @@ cc_library(
         "@com_google_absl//absl/base:config",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:vlog_is_on",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
diff --git a/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc b/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc
index e54d0a319cd..02c5ec0f759 100644
--- a/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc
+++ b/third_party/xla/xla/backends/cpu/codegen/fusion_compiler.cc
@@ -23,6 +23,7 @@ limitations under the License.
 
 #include "absl/base/config.h"  // IWYU pragma: keep
 #include "absl/log/log.h"
+#include "absl/log/vlog_is_on.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"

From 0b28b0e5aa5f170df8e8c2b28c4bf0c41a1c1d61 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 12 Oct 2026 22:55:47 -0700
Subject: [PATCH 0645/1000] Automated Code Change

PiperOrigin-RevId: 998388955
---
 third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc
index de18d0234a0..7355893546f 100644
--- a/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <algorithm>
 #include <cstdint>
+#include <functional>
 #include <memory>
 #include <optional>
 #include <string>

From 8bc2b6baec0ee7d251a7deef3147e4c21c1e44a3 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 12 Oct 2026 23:20:36 -0700
Subject: [PATCH 0646/1000] Automated Code Change

PiperOrigin-RevId: 998399243
---
 tensorflow/core/framework/BUILD             | 1 +
 tensorflow/core/framework/op_def_builder.cc | 4 ++--
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensorflow/core/framework/BUILD b/tensorflow/core/framework/BUILD
index 81e3b1ba576..13f93988d5f 100644
--- a/tensorflow/core/framework/BUILD
+++ b/tensorflow/core/framework/BUILD
@@ -1114,6 +1114,7 @@ cc_library(
         "//tensorflow/core/platform:errors",
         "//tensorflow/core/platform:macros",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
     ],
 )
 
diff --git a/tensorflow/core/framework/op_def_builder.cc b/tensorflow/core/framework/op_def_builder.cc
index e0ee1cd81e0..728f0537d0d 100644
--- a/tensorflow/core/framework/op_def_builder.cc
+++ b/tensorflow/core/framework/op_def_builder.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/strings/escaping.h"
+#include "absl/types/span.h"
 #include "tensorflow/core/framework/attr_value.pb.h"
 #include "tensorflow/core/framework/attr_value_util.h"
 #include "tensorflow/core/framework/op_def_util.h"
@@ -495,8 +496,7 @@ void FinalizeDoc(const std::string& text, OpDef* op_def,
   // Trim trailing blank lines from the description.
   while (start_l < end_l && lines[end_l - 1].empty()) --end_l;
   std::string desc = absl::StrJoin(
-      absl::Span<const std::string>(lines.data() + start_l, end_l - start_l),
-      "\n");
+      absl::MakeConstSpan(lines).subspan(start_l, end_l - start_l), "\n");
   if (!desc.empty()) op_def->set_description(desc);
 
   // name: description

From de9f07d8795ca3bf56fc40c1fb04852fc1fc2989 Mon Sep 17 00:00:00 2001
From: Volodymyr Kysenko <vksnk@google.com>
Date: Mon, 12 Oct 2026 23:30:55 -0700
Subject: [PATCH 0647/1000] Support per-tensor INT8 quantized KV cache in SDPA.

Add `DefineSdpaDot` to handle both float and per-tensor `INT8` quantized `KV` operands in `DefineSdpaNode` by dynamically quantizing the LHS to `int8`.

PiperOrigin-RevId: 998403184
---
 .../lite/delegates/ynnpack/attention.cc       | 117 ++++++++++++++----
 .../lite/delegates/ynnpack/attention_model.cc |  17 ++-
 .../lite/delegates/ynnpack/attention_model.h  |   4 +-
 .../lite/delegates/ynnpack/attention_test.cc  |  84 +++++++++++++
 4 files changed, 191 insertions(+), 31 deletions(-)

diff --git a/tensorflow/lite/delegates/ynnpack/attention.cc b/tensorflow/lite/delegates/ynnpack/attention.cc
index 28e4d605ecd..ee48fda8fce 100644
--- a/tensorflow/lite/delegates/ynnpack/attention.cc
+++ b/tensorflow/lite/delegates/ynnpack/attention.cc
@@ -72,6 +72,66 @@ SdpaInputs GetSdpaInputs(TfLiteContext* context, const NodeInfo& node) {
   return inputs;
 }
 
+// Computes `output = a @ b` where `a` is float [..., M, K] and `b` is float or
+// per-tensor quantized int8 [..., K, N]. When `b` is quantized, `a` is
+// dynamically quantized to int8 along its reduction axis (-1).
+TfLiteStatus DefineSdpaDot(TfLiteContext* context, ynn_subgraph_t subgraph,
+                           uint32_t a_id, uint32_t b_id,
+                           const TfLiteTensor& b_tensor, ynn_type output_type,
+                           uint32_t* output_id) {
+  if (!IsQuantized(b_tensor)) {
+    if (*output_id == YNN_INVALID_VALUE_ID) {
+      TF_LITE_ENSURE_YNN_STATUS(ynn_define_tensor(
+          subgraph, output_type, 0, nullptr, nullptr, 0, output_id));
+    }
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_dot(subgraph, /*num_k_dims=*/1, a_id,
+                                             b_id, YNN_INVALID_VALUE_ID,
+                                             output_id, 0));
+    return kTfLiteOk;
+  }
+
+  uint32_t b_scale_id = YNN_INVALID_VALUE_ID;
+  uint32_t b_zp_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_STATUS(DefineQuantizationParams(context, subgraph, b_tensor,
+                                                 &b_scale_id, &b_zp_id));
+
+  int32_t reduce_axis = -1;
+  uint32_t min_max_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_YNN_STATUS(ynn_define_reduce(
+      subgraph, ynn_reduce_min_max, 1, &reduce_axis, a_id, YNN_INVALID_VALUE_ID,
+      &min_max_id, YNN_NODE_FLAG_KEEP_DIMS));
+
+  uint32_t a_zp_id = YNN_INVALID_VALUE_ID;
+  uint32_t a_scale_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_YNN_STATUS(ynn_define_dynamic_quantization(
+      subgraph, min_max_id, ynn_type_int8, &a_zp_id, &a_scale_id, 0));
+
+  uint32_t q_a_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_YNN_STATUS(ynn_define_quantize(
+      subgraph, a_id, ynn_type_int8, a_zp_id, a_scale_id, &q_a_id, 0));
+
+  uint32_t dot_zp_id = YNN_INVALID_VALUE_ID;
+  uint32_t dot_scale_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_YNN_STATUS(ynn::define_dot_quantization(
+      subgraph, /*num_k_dims=*/1, q_a_id, a_zp_id, a_scale_id, b_id, b_zp_id,
+      b_scale_id, dot_zp_id, dot_scale_id));
+
+  uint32_t accum_init_id = YNN_INVALID_VALUE_ID;
+  if (dot_zp_id != YNN_INVALID_VALUE_ID) {
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_unary(subgraph, ynn_unary_negate,
+                                               dot_zp_id, &accum_init_id, 0));
+  }
+
+  uint32_t accum_id = YNN_INVALID_VALUE_ID;
+  TF_LITE_ENSURE_YNN_STATUS(ynn_define_dot(subgraph, /*num_k_dims=*/1, q_a_id,
+                                           b_id, accum_init_id, &accum_id, 0));
+
+  TF_LITE_ENSURE_YNN_STATUS(
+      ynn_define_dequantize(subgraph, accum_id, YNN_INVALID_VALUE_ID,
+                            dot_scale_id, output_type, output_id, 0));
+  return kTfLiteOk;
+}
+
 }  // namespace
 
 bool IsSdpa(const TfLiteRegistration* registration, const TfLiteNode* node) {
@@ -116,13 +176,16 @@ TfLiteStatus IsSdpaSupported(const TfLiteRegistration* registration,
     return type == kTfLiteFloat32 || type == kTfLiteFloat16 ||
            type == kTfLiteBFloat16;
   };
+  auto is_quantized_int8 = [](const TfLiteTensor& tensor) {
+    return tensor.type == kTfLiteInt8 && IsQuantized(tensor);
+  };
 
   TF_LITE_ENSURE(context, IsTensorSupported(q));
   TF_LITE_ENSURE(context, is_float_type(q.type));
   TF_LITE_ENSURE(context, IsTensorSupported(k));
-  TF_LITE_ENSURE(context, is_float_type(k.type));
+  TF_LITE_ENSURE(context, k.type == q.type || is_quantized_int8(k));
   TF_LITE_ENSURE(context, IsTensorSupported(v));
-  TF_LITE_ENSURE(context, is_float_type(v.type));
+  TF_LITE_ENSURE(context, is_float_type(v.type) || is_quantized_int8(v));
   TF_LITE_ENSURE(context, IsTensorSupported(output));
   TF_LITE_ENSURE(context, is_float_type(output.type));
 
@@ -333,11 +396,6 @@ TfLiteStatus DefineSdpaNode(TfLiteContext* context, ynn_subgraph_t subgraph,
     k_trans_id = current_k_val_id;
   }
 
-  uint32_t scale_const_id = YNN_INVALID_VALUE_ID;
-  TF_LITE_ENSURE_YNN_STATUS(ynn_define_tensor(
-      subgraph, GetYnnType(q_tensor.type), 0, nullptr, &scale_val,
-      YNN_VALUE_FLAG_COPY_DATA_FP32, &scale_const_id));
-
   const int q_head_dim = is_seq_major ? 2 : 1;
   const int k_head_dim = is_seq_major ? 2 : 1;
   const int n_q = q_tensor.dims->data[q_head_dim];
@@ -347,11 +405,13 @@ TfLiteStatus DefineSdpaNode(TfLiteContext* context, ynn_subgraph_t subgraph,
 
   const int q_seq_dim = is_seq_major ? 1 : 2;
   const int q_seq = q_tensor.dims->data[q_seq_dim];
-  const bool use_decode1 =
+  const bool is_small_q =
       !is_seq_major && (q_seq * static_cast<int>(g_heads_per_kv) <= 8);
+  const bool use_decode1 =
+      is_small_q && !IsQuantized(k_tensor) && !IsQuantized(v_tensor);
   // Prefill with grouped heads: keep the head axis as [n_kv, g] and let the dot
   // broadcast K/V over g, instead of folding g into the row axis.
-  const bool gqa_batch = g_heads_per_kv > 1 && !is_seq_major && !use_decode1;
+  const bool gqa_batch = g_heads_per_kv > 1 && !is_seq_major && !is_small_q;
   const bool gqa_fold = g_heads_per_kv > 1 && !gqa_batch;
 
   if (g_heads_per_kv > 1) {
@@ -384,10 +444,17 @@ TfLiteStatus DefineSdpaNode(TfLiteContext* context, ynn_subgraph_t subgraph,
   }
 
   // Scale Q.
-  uint32_t q_scaled_id = YNN_INVALID_VALUE_ID;
-  TF_LITE_ENSURE_YNN_STATUS(ynn_define_binary(subgraph, ynn_binary_multiply,
-                                              q_trans_id, scale_const_id,
-                                              &q_scaled_id, 0));
+  uint32_t q_scaled_id = q_trans_id;
+  if (scale_val != 1.0f) {
+    uint32_t scale_const_id = YNN_INVALID_VALUE_ID;
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_tensor(
+        subgraph, GetYnnType(q_tensor.type), 0, nullptr, &scale_val,
+        YNN_VALUE_FLAG_COPY_DATA_FP32, &scale_const_id));
+    q_scaled_id = YNN_INVALID_VALUE_ID;
+    TF_LITE_ENSURE_YNN_STATUS(ynn_define_binary(subgraph, ynn_binary_multiply,
+                                                q_trans_id, scale_const_id,
+                                                &q_scaled_id, 0));
+  }
 
   // Scores: S = Q @ K^T, [B, H, Q, S].
   const int32_t swap_last_two_perm[] = {-1, -2};
@@ -413,9 +480,9 @@ TfLiteStatus DefineSdpaNode(TfLiteContext* context, ynn_subgraph_t subgraph,
         ynn_define_static_transpose(subgraph, 2, swap_last_two_perm, k_trans_id,
                                     &k_trans_t_id, YNN_NODE_FLAG_KEEP_DIMS));
 
-    TF_LITE_ENSURE_YNN_STATUS(
-        ynn_define_dot(subgraph, /*num_k_dims=*/1, q_scaled_id, k_trans_t_id,
-                       YNN_INVALID_VALUE_ID, &scores_id, 0));
+    TF_LITE_ENSURE_STATUS(DefineSdpaDot(context, subgraph, q_scaled_id,
+                                        k_trans_t_id, k_tensor, ynn_type_fp32,
+                                        &scores_id));
   }
 
   uint32_t logits_id = scores_id;
@@ -493,8 +560,11 @@ TfLiteStatus DefineSdpaNode(TfLiteContext* context, ynn_subgraph_t subgraph,
   }
 
   uint32_t probs_id = YNN_INVALID_VALUE_ID;
-  TF_LITE_ENSURE_YNN_STATUS(ynn_define_tensor(
-      subgraph, GetYnnType(v_tensor.type), 0, nullptr, nullptr, 0, &probs_id));
+  if (!IsQuantized(v_tensor)) {
+    TF_LITE_ENSURE_YNN_STATUS(
+        ynn_define_tensor(subgraph, GetYnnType(v_tensor.type), 0, nullptr,
+                          nullptr, 0, &probs_id));
+  }
   TF_LITE_ENSURE_YNN_STATUS(
       ynn::define_softmax(subgraph, masked_logits_id, 1.0f, probs_id));
 
@@ -551,14 +621,9 @@ TfLiteStatus DefineSdpaNode(TfLiteContext* context, ynn_subgraph_t subgraph,
         subgraph, 2, is_seq_major ? seq_major_v_perm : swap_last_two_perm,
         current_v_val_id, &v_trans_id, YNN_NODE_FLAG_KEEP_DIMS));
 
-    if (*post_bmm_ptr == YNN_INVALID_VALUE_ID) {
-      TF_LITE_ENSURE_YNN_STATUS(
-          ynn_define_tensor(subgraph, GetYnnType(output_tensor.type), 0,
-                            nullptr, nullptr, 0, post_bmm_ptr));
-    }
-    TF_LITE_ENSURE_YNN_STATUS(
-        ynn_define_dot(subgraph, /*num_k_dims=*/1, probs_id, v_trans_id,
-                       YNN_INVALID_VALUE_ID, post_bmm_ptr, 0));
+    TF_LITE_ENSURE_STATUS(
+        DefineSdpaDot(context, subgraph, probs_id, v_trans_id, v_tensor,
+                      GetYnnType(output_tensor.type), post_bmm_ptr));
   }
 
   uint32_t post_trans_id = *post_bmm_ptr;
diff --git a/tensorflow/lite/delegates/ynnpack/attention_model.cc b/tensorflow/lite/delegates/ynnpack/attention_model.cc
index ee1add8c6b5..f7825d51a4d 100644
--- a/tensorflow/lite/delegates/ynnpack/attention_model.cc
+++ b/tensorflow/lite/delegates/ynnpack/attention_model.cc
@@ -209,7 +209,8 @@ TfLiteRegistration* Register_DummySDPA() {
 AttentionModel::AttentionModel(
     int b, int t, int s, int h, int n, float scale, bool transpose_io,
     bool use_delegate, const TfLiteYNNPackDelegateOptions& delegate_options,
-    AttentionImpl impl, int n_kv) {
+    AttentionImpl impl, int n_kv, bool quantize_kv, bool bool_mask,
+    int32_t kv_zero_point) {
   if (n_kv <= 0) {
     n_kv = n;
   }
@@ -237,14 +238,22 @@ AttentionModel::AttentionModel(
   }
 
   query_id_ = AddInput({TensorType_FLOAT32, query_shape});
-  key_id_ = AddInput({TensorType_FLOAT32, key_shape});
-  value_id_ = AddInput({TensorType_FLOAT32, value_shape});
+  if (quantize_kv) {
+    key_id_ = AddInput(
+        {TensorType_INT8, key_shape, 0.0f, 0.0f, 0.05f, kv_zero_point});
+    value_id_ = AddInput(
+        {TensorType_INT8, value_shape, 0.0f, 0.0f, 0.05f, kv_zero_point});
+  } else {
+    key_id_ = AddInput({TensorType_FLOAT32, key_shape});
+    value_id_ = AddInput({TensorType_FLOAT32, value_shape});
+  }
   if (impl != AttentionImpl::kFullSequence) {
     runtime_bmm_params_id_ = AddInput({TensorType_INT32, {1}});
   } else {
     runtime_bmm_params_id_ = -1;
   }
-  mask_id_ = AddInput({TensorType_FLOAT32, mask_shape});
+  mask_id_ =
+      AddInput({bool_mask ? TensorType_BOOL : TensorType_FLOAT32, mask_shape});
 
   std::vector<int> output_shape = transpose_io ? std::vector<int>{b, t, n, h}
                                                : std::vector<int>{b, n, t, h};
diff --git a/tensorflow/lite/delegates/ynnpack/attention_model.h b/tensorflow/lite/delegates/ynnpack/attention_model.h
index af5c13c851e..38157f63b42 100644
--- a/tensorflow/lite/delegates/ynnpack/attention_model.h
+++ b/tensorflow/lite/delegates/ynnpack/attention_model.h
@@ -16,6 +16,7 @@ limitations under the License.
 #ifndef TENSORFLOW_LITE_DELEGATES_YNNPACK_ATTENTION_MODEL_H_
 #define TENSORFLOW_LITE_DELEGATES_YNNPACK_ATTENTION_MODEL_H_
 
+#include <cstdint>
 #include <vector>
 
 #include "tensorflow/lite/c/common.h"
@@ -44,7 +45,8 @@ class AttentionModel : public MultiOpModel {
   AttentionModel(int b, int t, int s, int h, int n, float scale,
                  bool transpose_io, bool use_delegate,
                  const TfLiteYNNPackDelegateOptions& delegate_options,
-                 AttentionImpl impl, int n_kv = -1);
+                 AttentionImpl impl, int n_kv = -1, bool quantize_kv = false,
+                 bool bool_mask = false, int32_t kv_zero_point = 0);
 
   int query() const { return query_id_; }
   int key() const { return key_id_; }
diff --git a/tensorflow/lite/delegates/ynnpack/attention_test.cc b/tensorflow/lite/delegates/ynnpack/attention_test.cc
index 13ef8845c05..5f9cb3bce3b 100644
--- a/tensorflow/lite/delegates/ynnpack/attention_test.cc
+++ b/tensorflow/lite/delegates/ynnpack/attention_test.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include <cmath>
 #include <cstddef>
+#include <cstdint>
 #include <limits>
 #include <string>
 #include <vector>
@@ -431,6 +432,89 @@ TEST(AttentionGqaTest, OdmlSdpaTransposedGqaDecodeAndPrefill) {
   }
 }
 
+TEST(AttentionGqaTest, OdmlSdpaTransposedQuantizedKvAndBoolMask) {
+  for (int t : {1, 4}) {
+    for (int n_kv : {1, 2}) {
+      for (int32_t kv_zp : {0, -5}) {
+        const int b = 1;
+        const int s = 24;
+        const int h = 16;
+        const int n_q = 4;
+        const int s_active = 11;
+        const float scale = 1.0f / std::sqrt(static_cast<float>(h));
+        const float kv_scale = 0.05f;
+
+        TfLiteYNNPackDelegateOptions options =
+            TfLiteYNNPackDelegateOptionsDefault();
+        options.num_threads = 1;
+        options.static_shape = true;
+
+        std::vector<float> q_data(b * n_q * t * h);
+        std::vector<int8_t> k_qdata(b * n_kv * s * h);
+        std::vector<float> k_fdata(b * n_kv * s * h);
+        std::vector<int8_t> v_qdata(b * n_kv * h * s);
+        std::vector<float> v_fdata(b * n_kv * h * s);
+        std::vector<bool> mask_bdata(b * 1 * t * s);
+        std::vector<float> mask_fdata(b * 1 * t * s);
+
+        for (size_t i = 0; i < q_data.size(); ++i) {
+          q_data[i] = 0.5f * std::sin(static_cast<float>(i + 1));
+        }
+        for (size_t i = 0; i < k_qdata.size(); ++i) {
+          int val = static_cast<int>(
+              std::round(20.0f * std::cos(static_cast<float>(i + 2))));
+          k_qdata[i] = static_cast<int8_t>(val);
+          k_fdata[i] = static_cast<float>(k_qdata[i] - kv_zp) * kv_scale;
+        }
+        for (size_t i = 0; i < v_qdata.size(); ++i) {
+          int val = static_cast<int>(
+              std::round(20.0f * std::sin(static_cast<float>(i + 3))));
+          v_qdata[i] = static_cast<int8_t>(val);
+          v_fdata[i] = static_cast<float>(v_qdata[i] - kv_zp) * kv_scale;
+        }
+        for (int i = 0; i < b * t; ++i) {
+          for (int j = 0; j < s; ++j) {
+            bool allow = (j < s_active) && (j <= i + 4);
+            mask_bdata[i * s + j] = allow;
+            mask_fdata[i * s + j] = allow ? 0.0f : -10000.0f;
+          }
+        }
+
+        AttentionModel model_ref(b, t, s, h, n_q, scale, /*transpose_io=*/false,
+                                 /*use_delegate=*/false, options,
+                                 AttentionImpl::kOdmlRuntimeBmm, n_kv);
+        model_ref.PopulateTensor(model_ref.query(), q_data);
+        model_ref.PopulateTensor(model_ref.key(), k_fdata);
+        model_ref.PopulateTensor(model_ref.value(), v_fdata);
+        model_ref.PopulateTensor(model_ref.runtime_bmm_params(), {s_active});
+        model_ref.PopulateTensor(model_ref.mask(), mask_fdata);
+        ASSERT_EQ(model_ref.Invoke(), kTfLiteOk);
+
+        AttentionModel model(b, t, s, h, n_q, scale, /*transpose_io=*/false,
+                             /*use_delegate=*/true, options,
+                             AttentionImpl::kOdmlSdpa, n_kv,
+                             /*quantize_kv=*/true, /*bool_mask=*/true, kv_zp);
+        model.PopulateTensor(model.query(), q_data);
+        model.PopulateTensor(model.key(), k_qdata);
+        model.PopulateTensor(model.value(), v_qdata);
+        model.PopulateTensor(model.runtime_bmm_params(), {s_active});
+        model.PopulateTensor(model.mask(), mask_bdata);
+        ASSERT_EQ(model.Invoke(), kTfLiteOk);
+
+        auto out_delegate = model.ExtractVector<float>(model.output());
+        auto out_ref = model_ref.ExtractVector<float>(model_ref.output());
+        ASSERT_EQ(out_delegate.size(), out_ref.size());
+
+        for (size_t i = 0; i < out_delegate.size(); ++i) {
+          EXPECT_FALSE(std::isnan(out_delegate[i]));
+          EXPECT_FALSE(std::isnan(out_ref[i]));
+          EXPECT_NEAR(out_delegate[i], out_ref[i], 2e-2f);
+        }
+      }
+    }
+  }
+}
+
 INSTANTIATE_TEST_SUITE_P(AttentionTest, AttentionTest,
                          ::testing::Values(AttentionImpl::kFullSequence,
                                            AttentionImpl::kOdmlRuntimeBmm,

From ef1025c86256bddeb8f0df8798852e506de88bfb Mon Sep 17 00:00:00 2001
From: Saurabh Singh <geniusthegreat0963@gmail.com>
Date: Tue, 13 Oct 2026 15:08:57 +0530
Subject: [PATCH 0648/1000] Add symmetric test coverage for complex infinities
 and allclose

- Add reverse test pair (inf+0j, inf+infj) for complex isclose
- Add [inf], [-inf] and [-inf], [-inf] cases for allclose
---
 .../python/ops/numpy_ops/np_math_ops_test.py     | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

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 9a5bf467bfd..9505b44a6c3 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
@@ -412,8 +412,12 @@ class MathTest(test.TestCase, parameterized.TestCase):
     a_vals = [np.inf, -np.inf, -np.inf, np.inf, 1.0, 1.0, np.inf]
     b_vals = [np.inf, -np.inf, np.inf, -np.inf, 1.0, np.inf, 1.0]
     if np.issubdtype(dtype, np.complexfloating):
-      a_vals.extend([complex(np.inf, np.inf), complex(np.inf, np.inf)])
-      b_vals.extend([complex(np.inf, 0), complex(np.inf, np.inf)])
+      a_vals.extend([
+          complex(np.inf, np.inf), complex(np.inf, np.inf), complex(np.inf, 0)
+      ])
+      b_vals.extend([
+          complex(np.inf, 0), complex(np.inf, np.inf), complex(np.inf, np.inf)
+      ])
 
     a = np.array(a_vals, dtype)
     b = np.array(b_vals, dtype)
@@ -429,6 +433,14 @@ class MathTest(test.TestCase, parameterized.TestCase):
         np_math_ops.allclose([dtype(-np.inf)], [dtype(np.inf)]),
         np.allclose([dtype(-np.inf)], [dtype(np.inf)]),
     )
+    self.match(
+        np_math_ops.allclose([dtype(np.inf)], [dtype(-np.inf)]),
+        np.allclose([dtype(np.inf)], [dtype(-np.inf)]),
+    )
+    self.match(
+        np_math_ops.allclose([dtype(-np.inf)], [dtype(-np.inf)]),
+        np.allclose([dtype(-np.inf)], [dtype(-np.inf)]),
+    )
 
   @parameterized.parameters([np.int32, np.int64])
   def testIsCloseIntegerTolerances(self, dtype):

From bc42aef95a4607e4d2ef3099fa9ae869ad5d8d26 Mon Sep 17 00:00:00 2001
From: Praveen Batra <praveenbatra@google.com>
Date: Mon, 12 Oct 2026 23:32:47 -0700
Subject: [PATCH 0649/1000] Add some test skips

PiperOrigin-RevId: 998403933
---
 third_party/xla/xla/tests/reduce_window_test.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/xla/xla/tests/reduce_window_test.cc b/third_party/xla/xla/tests/reduce_window_test.cc
index 0008fa95012..95cbe7ea3d2 100644
--- a/third_party/xla/xla/tests/reduce_window_test.cc
+++ b/third_party/xla/xla/tests/reduce_window_test.cc
@@ -2072,7 +2072,7 @@ ENTRY %SyncTensorsGraph.43 (p0.1: f32[], p1.7: pred[3,3]) -> (pred[]) {
 }
 
 TEST_F(ReduceWindowHloTest, VariadicWithNonTrivialWindows) {
-  if (test::DeviceTypeIs(test::kGpu)) {
+  if (test::DeviceTypeIs(test::kGpu) || test::HasModifiers({test::kGrm})) {
     GTEST_SKIP();
   }
   const char* const hlo_string = R"(

From e3ba48047f660a0f1436934d5a3d2d53a75b1ed8 Mon Sep 17 00:00:00 2001
From: Alejandro Pacheco <alepacheco99@gmail.com>
Date: Tue, 13 Oct 2026 12:04:09 +0200
Subject: [PATCH 0650/1000] Handle zero-size outputs in ParseFeature and test
 the copy-pass bounds

A tensor with no elements can have no buffer, and a null output pointer
makes the Parse*Feature functions count values instead of storing them,
skipping the bound. ParseFeature now uses GetFeatureLength when there is
no room left, and fails if the feature has any values.

Add tests for ParseFeature (values that don't fit, and an empty output),
PadFeature and CopyTensorIntoTensor. These checks can't fail through
FastParseSequenceExample since the counting pass sizes every output, so
the helpers are exported for tests like TestFastParse.
---
 .../core/util/example_proto_fast_parsing.cc   | 58 ++++++++++-----
 .../core/util/example_proto_fast_parsing.h    | 15 ++++
 .../util/example_proto_fast_parsing_test.cc   | 72 +++++++++++++++++++
 3 files changed, 128 insertions(+), 17 deletions(-)

diff --git a/tensorflow/core/util/example_proto_fast_parsing.cc b/tensorflow/core/util/example_proto_fast_parsing.cc
index 8b6d458c759..19ae723ec32 100644
--- a/tensorflow/core/util/example_proto_fast_parsing.cc
+++ b/tensorflow/core/util/example_proto_fast_parsing.cc
@@ -2045,6 +2045,23 @@ inline int ParseInt64Feature(protobuf::io::CodedInputStream* stream,
   return num_elements;
 }
 
+// Returns the length of the next feature on `stream`.
+// Returns -1 if the next feature on `stream` doesn't match `dtype`.
+inline int GetFeatureLength(DataType dtype,
+                            protobuf::io::CodedInputStream* stream) {
+  switch (dtype) {
+    case DT_STRING:
+      return ParseBytesFeature(stream, nullptr, 0);
+    case DT_FLOAT:
+      return ParseFloatFeature(stream, nullptr, 0);
+    case DT_INT64:
+      return ParseInt64Feature(stream, nullptr, 0);
+    default:
+      ReportUnexpectedDataType(dtype);
+      return -1;
+  }
+}
+
 // Parses the next feature on `stream` into `out` starting at `out_offset`.
 // Updates `out_offset`, and returns the number of values added.
 // Returns -1 if the next feature on `stream` doesn't match `dtype`, or if it
@@ -2056,6 +2073,11 @@ inline int ParseFeature(DataType dtype, protobuf::io::CodedInputStream* stream,
     return -1;
   }
   const size_t out_size = out_elements - *out_offset;
+  if (out_size == 0) {
+    // `out` may have no buffer at all, and a null pointer would make the
+    // Parse*Feature functions count the values instead of storing them.
+    return GetFeatureLength(dtype, stream) == 0 ? 0 : -1;
+  }
   int delta;
   switch (dtype) {
     case DT_STRING:
@@ -2110,23 +2132,6 @@ inline bool PadFeature(DataType dtype, size_t num_to_pad, Tensor* out,
   return true;
 }
 
-// Returns the length of the next feature on `stream`.
-// Returns -1 if the next feature on `stream` doesn't match `dtype`.
-inline int GetFeatureLength(DataType dtype,
-                            protobuf::io::CodedInputStream* stream) {
-  switch (dtype) {
-    case DT_STRING:
-      return ParseBytesFeature(stream, nullptr, 0);
-    case DT_FLOAT:
-      return ParseFloatFeature(stream, nullptr, 0);
-    case DT_INT64:
-      return ParseInt64Feature(stream, nullptr, 0);
-    default:
-      ReportUnexpectedDataType(dtype);
-      return -1;
-  }
-}
-
 inline DataType ParseDataType(protobuf::io::CodedInputStream* stream) {
   uint8_t peek_tag = PeekTag(stream);
   switch (peek_tag) {
@@ -2994,6 +2999,25 @@ absl::Status ParseSequenceRaggedFeatures(
 
 }  // namespace
 
+int TestParseFeature(DataType dtype, const std::string& serialized_feature,
+                     Tensor* out, size_t* out_offset) {
+  protobuf::io::CodedInputStream stream(
+      reinterpret_cast<const uint8_t*>(serialized_feature.data()),
+      serialized_feature.size());
+  EnableAliasing(&stream);
+  return ParseFeature(dtype, &stream, out, out_offset);
+}
+
+bool TestPadFeature(DataType dtype, size_t num_to_pad, Tensor* out,
+                    size_t* out_offset) {
+  return PadFeature(dtype, num_to_pad, out, out_offset);
+}
+
+bool TestCopyTensorIntoTensor(DataType dtype, const Tensor& src, Tensor* dst,
+                              size_t* dst_offset) {
+  return CopyTensorIntoTensor(dtype, src, dst, dst_offset);
+}
+
 // TODO(sundberg): Use the threadpool to parallelize example parsing.
 // TODO(b/111553342): Support extracting feature statistics from the examples.
 absl::Status FastParseSequenceExample(
diff --git a/tensorflow/core/util/example_proto_fast_parsing.h b/tensorflow/core/util/example_proto_fast_parsing.h
index 583a4238737..ab68f052033 100644
--- a/tensorflow/core/util/example_proto_fast_parsing.h
+++ b/tensorflow/core/util/example_proto_fast_parsing.h
@@ -166,6 +166,21 @@ absl::Status FastParseSequenceExample(
 // It is exported here as a convenient API to test parser part separately.
 bool TestFastParse(const std::string& serialized, Example* example);
 
+// These run the helpers that FastParseSequenceExample uses to fill its outputs,
+// so that their bounds checks can be tested on their own. The counting pass
+// sizes every output to fit, so the checks can't fail through
+// FastParseSequenceExample.
+// TestParseFeature parses the serialized Feature into `out` starting at
+// `*out_offset` and returns the number of values added, or -1 if they don't
+// fit. TestPadFeature and TestCopyTensorIntoTensor return false if the values
+// don't fit.
+int TestParseFeature(DataType dtype, const std::string& serialized_feature,
+                     Tensor* out, size_t* out_offset);
+bool TestPadFeature(DataType dtype, size_t num_to_pad, Tensor* out,
+                    size_t* out_offset);
+bool TestCopyTensorIntoTensor(DataType dtype, const Tensor& src, Tensor* dst,
+                              size_t* dst_offset);
+
 }  // namespace example
 }  // namespace tensorflow
 
diff --git a/tensorflow/core/util/example_proto_fast_parsing_test.cc b/tensorflow/core/util/example_proto_fast_parsing_test.cc
index 9d4e9c16240..00a6753011a 100644
--- a/tensorflow/core/util/example_proto_fast_parsing_test.cc
+++ b/tensorflow/core/util/example_proto_fast_parsing_test.cc
@@ -701,6 +701,78 @@ TEST(FastParse, SequenceExampleFillsAllOutputKinds) {
                 std::vector<int64_t>{0, 2, 2});
 }
 
+TEST(FastParse, ParseFeatureRejectsValuesThatDoNotFit) {
+  Feature floats, ints, bytes;
+  for (float v : {1, 2, 3}) floats.mutable_float_list()->add_value(v);
+  for (int64_t v : {1, 2, 3}) ints.mutable_int64_list()->add_value(v);
+  for (const char* v : {"a", "b", "c"})
+    bytes.mutable_bytes_list()->add_value(v);
+  for (const auto& [dtype, feature] :
+       {std::pair{DT_FLOAT, &floats}, std::pair{DT_INT64, &ints},
+        std::pair{DT_STRING, &bytes}}) {
+    // Room for two of the three values after the offset.
+    Tensor out(dtype, TensorShape({3}));
+    size_t out_offset = 1;
+    EXPECT_EQ(TestParseFeature(dtype, feature->SerializeAsString(), &out,
+                               &out_offset),
+              -1)
+        << DataTypeString(dtype);
+    EXPECT_EQ(out_offset, 1);
+
+    Tensor fits(dtype, TensorShape({4}));
+    out_offset = 1;
+    EXPECT_EQ(TestParseFeature(dtype, feature->SerializeAsString(), &fits,
+                               &out_offset),
+              3)
+        << DataTypeString(dtype);
+    EXPECT_EQ(out_offset, 4);
+  }
+}
+
+TEST(FastParse, ParseFeatureIntoEmptyTensor) {
+  Feature feature;
+  feature.mutable_float_list()->add_value(1);
+  Tensor out(DT_FLOAT, TensorShape({0}));
+  size_t out_offset = 0;
+  EXPECT_EQ(TestParseFeature(DT_FLOAT, feature.SerializeAsString(), &out,
+                             &out_offset),
+            -1);
+  EXPECT_EQ(out_offset, 0);
+
+  // A feature without values still fits.
+  feature.mutable_float_list()->clear_value();
+  EXPECT_EQ(TestParseFeature(DT_FLOAT, feature.SerializeAsString(), &out,
+                             &out_offset),
+            0);
+  EXPECT_EQ(out_offset, 0);
+}
+
+TEST(FastParse, PadFeatureRejectsPaddingThatDoesNotFit) {
+  Tensor out(DT_INT64, TensorShape({3}));
+  size_t out_offset = 1;
+  EXPECT_FALSE(TestPadFeature(DT_INT64, 3, &out, &out_offset));
+  EXPECT_EQ(out_offset, 1);
+  EXPECT_TRUE(TestPadFeature(DT_INT64, 2, &out, &out_offset));
+  EXPECT_EQ(out_offset, 3);
+}
+
+TEST(FastParse, CopyTensorIntoTensorRejectsSourceThatDoesNotFit) {
+  // A context dense default with more values than the output has room for.
+  Tensor default_value(DT_INT64, TensorShape({3}));
+  default_value.flat<int64_t>().setConstant(7);
+  Tensor out(DT_INT64, TensorShape({2}));
+  size_t out_offset = 0;
+  EXPECT_FALSE(
+      TestCopyTensorIntoTensor(DT_INT64, default_value, &out, &out_offset));
+  EXPECT_EQ(out_offset, 0);
+
+  Tensor fits(DT_INT64, TensorShape({4}));
+  out_offset = 1;
+  EXPECT_TRUE(
+      TestCopyTensorIntoTensor(DT_INT64, default_value, &fits, &out_offset));
+  EXPECT_EQ(out_offset, 4);
+}
+
 }  // namespace
 }  // namespace example
 }  // namespace tensorflow

From 96d81bebd606fa9322e25c9cd284f8f77f46b52d Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Tue, 13 Oct 2026 15:38:06 +0530
Subject: [PATCH 0651/1000] log null VariantTensorData skip at WARNING

Signed-off-by: mohammed adib <adib@bugqore.com>
---
 tensorflow/core/data/serialization_utils.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/data/serialization_utils.cc b/tensorflow/core/data/serialization_utils.cc
index 551652fcf78..7cb62a52e0c 100644
--- a/tensorflow/core/data/serialization_utils.cc
+++ b/tensorflow/core/data/serialization_utils.cc
@@ -186,7 +186,7 @@ VariantTensorDataReader::VariantTensorDataReader(
     const std::vector<const tensorflow::VariantTensorData*>& data) {
   for (const auto& d : data) {
     if (d == nullptr) {
-      VLOG(1) << "Found null VariantTensorData pointer; skipping.";
+      LOG(WARNING) << "Found null VariantTensorData pointer; skipping.";
       continue;
     }
     std::string metadata;

From 292e407794e3002d7f909a36b74165c7a71dfb20 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 13 Oct 2026 00:49:14 -0700
Subject: [PATCH 0652/1000] Reverts e0205e3191d84ceda5485a6a27a5ddc9ed2a974d

PiperOrigin-RevId: 998440810
---
 .../concat_prune_inactive_small_tile.hlo      |  9 ++++-
 .../codegen/tiling/symbolic_tile_analysis.cc  |  9 ++++-
 .../tiling/symbolic_tile_analysis_test.cc     | 39 -------------------
 .../tiling/tiled_hlo_computation_test.cc      |  4 --
 .../codegen/tiling/tiled_hlo_schedule_test.cc |  3 --
 .../tiling/tiling_specification_test.cc       |  4 --
 .../codegen/xtile/codegen/fusion_emitter.cc   |  3 +-
 .../codegen/tiled_emitter_constraints_test.cc |  4 --
 .../model/triton_emitter_constraints_test.cc  |  4 --
 9 files changed, 15 insertions(+), 64 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive_small_tile.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive_small_tile.hlo
index a85a1167b3a..dba7de98b75 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive_small_tile.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive_small_tile.hlo
@@ -1,5 +1,8 @@
-// RUN: fusion_to_triton %s 2>&1 | FileCheck %s
-// RUN: triton_test_correctness %s
+// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s 2>&1 | FileCheck %s
+// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
+
+// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s 2>&1 | FileCheck %s --check-prefixes=CHECK-TILING
+// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
 
 // This is a fork of concat_prune_inactive.hlo that uses a smaller concat-dim
 // tile size (16 instead of 64). With the sign-agnostic boundary-alignment
@@ -14,6 +17,8 @@
 
 // CHECK:      xtile.entry_func @triton_fn
 
+// CHECK-TILING:      xtile.entry_func @triton_fn
+
 HloModule min_repro, is_scheduled=true, entry_computation_layout={(f32[64,10]{1,0}, f32[1,16,10]{2,1,0}, f32[1,16,10]{2,1,0})->f32[64,10]{1,0}}
 
 %fused_computation (param_0: f32[64,10], param_1: f32[1,16,10], param_2: f32[1,16,10]) -> f32[64,10] {
diff --git a/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis.cc b/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis.cc
index ddd77bc6e4d..38e915a268b 100644
--- a/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis.cc
+++ b/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis.cc
@@ -1398,8 +1398,13 @@ SymbolicTileAnalysis::AnalyzeFromInstruction(
 /*static*/ SymbolicTileAnalysisOrError SymbolicTileAnalysis::AnalyzeFusion(
     const HloFusionAdaptor& fusion, MLIRContext* mlir_context,
     EmitterSpecificConstraintsBuilder emitter_specific_constraints_builder) {
-  LOG(FATAL) << "Calling into symbolic tile analysis is disallowed. "
-             << "Please use experimental tiling propagation instead.";
+  CHECK(!fusion.GetRoots()
+             .front()
+             .instruction()
+             .GetModule()
+             ->config()
+             .debug_options()
+             .xla_gpu_experimental_enable_tiling_propagation());
   RegisterSymbolicExprStorage(mlir_context);
   auto real_root_index_or = GetRealRootIndex(fusion.GetRoots());
   if (!real_root_index_or.ok()) {
diff --git a/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis_test.cc b/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis_test.cc
index ae4e5bbc456..d7dc6c536ac 100644
--- a/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis_test.cc
+++ b/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis_test.cc
@@ -177,49 +177,10 @@ class FakeEmitterSpecificConstraints : public EmitterSpecificConstraints {
   int64_t dim0_tile_size_;
 };
 
-using SymbolicTileAnalysisDeathTest = HloHardwareIndependentTestBase;
-
-TEST_F(SymbolicTileAnalysisDeathTest, CallingSymbolicTileAnalysisFailsCheck) {
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
-HloModule m
-
-fused_computation {
-  p0 = f32[10,10] parameter(0)
-  ROOT r = f32[10,10] abs(p0)
-}
-
-ENTRY main {
-  p0 = f32[10,10] parameter(0)
-  ROOT fusion = f32[10,10] fusion(p0), kind=kCustom, calls=fused_computation
-}
-)"));
-  mlir::MLIRContext mlir_context;
-  EXPECT_DEATH(SymbolicTileAnalysis::AnalyzeComputation(
-                   *module->entry_computation()
-                        ->root_instruction()
-                        ->fused_instructions_computation(),
-                   &mlir_context),
-               "Calling into symbolic tile analysis is disallowed");
-
-  module->mutable_config()
-      .mutable_debug_options()
-      .set_xla_gpu_experimental_enable_tiling_propagation(false);
-  EXPECT_DEATH(SymbolicTileAnalysis::AnalyzeComputation(
-                   *module->entry_computation()
-                        ->root_instruction()
-                        ->fused_instructions_computation(),
-                   &mlir_context),
-               "Calling into symbolic tile analysis is disallowed");
-}
-
 class SymbolicTileAnalysisTest : public HloHardwareIndependentTestBase {
  public:
   SymbolicTileAnalysisTest() = default;
 
-  void SetUp() override {
-    GTEST_SKIP() << "Legacy symbolic tile analysis is disallowed.";
-  }
-
   std::optional<SymbolicTileAnalysis> TryAnalyzeModule(
       HloModule* module,
       EmitterSpecificConstraintsBuilder emitter_specific_constraints_builder =
diff --git a/third_party/xla/xla/codegen/tiling/tiled_hlo_computation_test.cc b/third_party/xla/xla/codegen/tiling/tiled_hlo_computation_test.cc
index d2e3c2e9736..853b69b13c5 100644
--- a/third_party/xla/xla/codegen/tiling/tiled_hlo_computation_test.cc
+++ b/third_party/xla/xla/codegen/tiling/tiled_hlo_computation_test.cc
@@ -45,10 +45,6 @@ class TiledHloComputationTest : public HloHardwareIndependentTestBase {
  public:
   TiledHloComputationTest() { RegisterSymbolicExprStorage(&mlir_context_); }
 
-  void SetUp() override {
-    GTEST_SKIP() << "Legacy symbolic tile analysis is disallowed.";
-  }
-
  protected:
   DebugOptions GetDebugOptionsForTest() const override {
     DebugOptions debug_options =
diff --git a/third_party/xla/xla/codegen/tiling/tiled_hlo_schedule_test.cc b/third_party/xla/xla/codegen/tiling/tiled_hlo_schedule_test.cc
index 22a324ab1ce..a58d9ef83f4 100644
--- a/third_party/xla/xla/codegen/tiling/tiled_hlo_schedule_test.cc
+++ b/third_party/xla/xla/codegen/tiling/tiled_hlo_schedule_test.cc
@@ -52,9 +52,6 @@ using ::testing::HasSubstr;
 class TiledHloScheduleTest : public HloHardwareIndependentTestBase {
  protected:
   TiledHloScheduleTest() { RegisterSymbolicExprStorage(&mlir_context_); }
-  void SetUp() override {
-    GTEST_SKIP() << "Legacy symbolic tile analysis is disallowed.";
-  }
   DebugOptions GetDebugOptionsForTest() const override {
     DebugOptions debug_options =
         HloHardwareIndependentTestBase::GetDebugOptionsForTest();
diff --git a/third_party/xla/xla/codegen/tiling/tiling_specification_test.cc b/third_party/xla/xla/codegen/tiling/tiling_specification_test.cc
index 1d72d1c88ce..6a302a41de5 100644
--- a/third_party/xla/xla/codegen/tiling/tiling_specification_test.cc
+++ b/third_party/xla/xla/codegen/tiling/tiling_specification_test.cc
@@ -51,10 +51,6 @@ class TilingSpecificationTest : public HloHardwareIndependentTestBase {
  public:
   TilingSpecificationTest() { RegisterSymbolicExprStorage(&mlir_context_); }
 
-  void SetUp() override {
-    GTEST_SKIP() << "Legacy symbolic tile analysis is disallowed.";
-  }
-
  protected:
   DebugOptions GetDebugOptionsForTest() const override {
     DebugOptions debug_options =
diff --git a/third_party/xla/xla/codegen/xtile/codegen/fusion_emitter.cc b/third_party/xla/xla/codegen/xtile/codegen/fusion_emitter.cc
index b7c339d07c8..83b44deb93b 100644
--- a/third_party/xla/xla/codegen/xtile/codegen/fusion_emitter.cc
+++ b/third_party/xla/xla/codegen/xtile/codegen/fusion_emitter.cc
@@ -1256,8 +1256,7 @@ absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> EmitXTileModule(
     return absl::InvalidArgumentError("Multi-output fusion is disabled.");
   }
 
-  LOG(FATAL) << "Calling into legacy fusion emitter is disallowed. "
-             << "Please use experimental tiling propagation instead.";
+  CHECK(!debug_options.xla_gpu_experimental_enable_tiling_propagation());
 
   const HloComputation* hlo_computation =
       fusion.fused_instructions_computation();
diff --git a/third_party/xla/xla/codegen/xtile/codegen/tiled_emitter_constraints_test.cc b/third_party/xla/xla/codegen/xtile/codegen/tiled_emitter_constraints_test.cc
index ef96001e82e..863b83b4af1 100644
--- a/third_party/xla/xla/codegen/xtile/codegen/tiled_emitter_constraints_test.cc
+++ b/third_party/xla/xla/codegen/xtile/codegen/tiled_emitter_constraints_test.cc
@@ -53,10 +53,6 @@ class TiledEmitterConstraintsTest : public HloHardwareIndependentTestBase {
 
  public:
   TiledEmitterConstraintsTest() { RegisterSymbolicExprStorage(&mlir_context_); }
-
-  void SetUp() override {
-    GTEST_SKIP() << "Legacy symbolic tile analysis is disallowed.";
-  }
   std::optional<SymbolicTileAnalysis> TryAnalyzeModule(
       HloModule* module, bool with_tiled_emitter_specific_constraints = true) {
     EmitterSpecificConstraintsBuilder constraints_builder = nullptr;
diff --git a/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc b/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc
index 95b1abb12a7..fe4d40c7fed 100644
--- a/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc
+++ b/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc
@@ -58,10 +58,6 @@ using ::testing::HasSubstr;
 class TritonEmitterConstraintsTest : public HloHardwareIndependentTestBase {
  public:
   TritonEmitterConstraintsTest() = default;
-
-  void SetUp() override {
-    GTEST_SKIP() << "Legacy symbolic tile analysis is disallowed.";
-  }
   std::optional<SymbolicTileAnalysis> TryAnalyzeModule(
       HloModule* module, bool with_triton_emitter_specific_constraints = true) {
     EmitterSpecificConstraintsBuilder constraints_builder = nullptr;

From 688f13655f99c1b6af05663294c9e29549df2428 Mon Sep 17 00:00:00 2001
From: nishad shabbir <nishadshabbir28@gmail.com>
Date: Tue, 13 Oct 2026 16:06:59 +0530
Subject: [PATCH 0653/1000] Cast the name_scope values back to lists

gradients_util._AsList() returns its argument unchanged when it is already
a list or a tuple, so ys_list and xs are not necessarily both lists and
concatenating them raises TypeError whenever the two differ. That happens
for a tensor ys with a tuple xs, and for tuple/list either way round.

Restore the list() calls and drive the test over the container
combinations, since a tensor for both ys and xs leaves them both as lists
and cannot catch this.
---
 tensorflow/python/ops/gradients_impl.py |  3 +-
 tensorflow/python/ops/gradients_test.py | 38 +++++++++++++++++--------
 2 files changed, 28 insertions(+), 13 deletions(-)

diff --git a/tensorflow/python/ops/gradients_impl.py b/tensorflow/python/ops/gradients_impl.py
index 50eea59c192..2cb1314eb99 100644
--- a/tensorflow/python/ops/gradients_impl.py
+++ b/tensorflow/python/ops/gradients_impl.py
@@ -428,7 +428,8 @@ def hessians(ys,
   }
   # Compute first-order derivatives and iterate for each x in xs.
   hessians = []
-  with ops.name_scope(name, "hessians", ys_list + xs):
+  # _AsList() passes tuples through unchanged, so cast before concatenating.
+  with ops.name_scope(name, "hessians", list(ys_list) + list(xs)):
     _gradients = gradients(ys, xs, **kwargs)
     for gradient, x in zip(_gradients, xs):
       # change shape to one-dimension without graph branching
diff --git a/tensorflow/python/ops/gradients_test.py b/tensorflow/python/ops/gradients_test.py
index 62b2d051d85..50c27d245fe 100644
--- a/tensorflow/python/ops/gradients_test.py
+++ b/tensorflow/python/ops/gradients_test.py
@@ -866,19 +866,33 @@ class HessianTest(test_util.TensorFlowTestCase):
     # `name` is documented as grouping all the gradient ops together, so the
     # ops hessians() adds must land under it rather than in the enclosing
     # scope.
+    # _AsList() leaves tuples alone, so the ys and xs containers also have to
+    # survive being mixed.
+    containers = [
+        lambda y, x: (y, x),
+        lambda y, x: (y, (x,)),
+        lambda y, x: ((y,), [x]),
+        lambda y, x: ([y], (x,)),
+        lambda y, x: ([y], [x]),
+        lambda y, x: ((y,), (x,)),
+    ]
     for name, expected in [(None, "hessians"), ("my_hessians", "my_hessians")]:
-      with ops.Graph().as_default() as g:
-        x = constant_op.constant([1.0, 2.0, 3.0])
-        y = math_ops.reduce_sum(x * x * x)
-        existing_ops = set(g.get_operations())
-        if name is None:
-          gradients.hessians(y, x)
-        else:
-          gradients.hessians(y, x, name=name)
-        added = [op.name for op in g.get_operations() if op not in existing_ops]
-      self.assertNotEmpty(added)
-      for op_name in added:
-        self.assertEqual(expected, op_name.split("/")[0])
+      for wrap in containers:
+        with ops.Graph().as_default() as g:
+          x = constant_op.constant([1.0, 2.0, 3.0])
+          y = math_ops.reduce_sum(x * x * x)
+          existing_ops = set(g.get_operations())
+          ys, xs = wrap(y, x)
+          if name is None:
+            gradients.hessians(ys, xs)
+          else:
+            gradients.hessians(ys, xs, name=name)
+          added = [
+              op.name for op in g.get_operations() if op not in existing_ops
+          ]
+        self.assertNotEmpty(added)
+        for op_name in added:
+          self.assertEqual(expected, op_name.split("/")[0])
 
 
 class IndexedSlicesToTensorTest(test_util.TensorFlowTestCase):

From 0762f279ebb8e8b66cdcf38b396d4173fe9d4cc8 Mon Sep 17 00:00:00 2001
From: Alejandro Pacheco <alepacheco99@gmail.com>
Date: Tue, 13 Oct 2026 12:44:28 +0200
Subject: [PATCH 0654/1000] Fail the copy-pass helpers on an unexpected dtype

ReportUnexpectedDataType() only DCHECKs, so in optimized builds the
default branches of PadFeature and CopyTensorIntoTensor reported success
and advanced the offset without writing anything. Return false there, and
-1 from ParseFeature, so an unexpected dtype becomes an error.
---
 .../core/util/example_proto_fast_parsing.cc   | 14 +++++++----
 .../util/example_proto_fast_parsing_test.cc   | 24 +++++++++++++++++++
 2 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/tensorflow/core/util/example_proto_fast_parsing.cc b/tensorflow/core/util/example_proto_fast_parsing.cc
index 19ae723ec32..a3453e73ddd 100644
--- a/tensorflow/core/util/example_proto_fast_parsing.cc
+++ b/tensorflow/core/util/example_proto_fast_parsing.cc
@@ -2064,8 +2064,9 @@ inline int GetFeatureLength(DataType dtype,
 
 // Parses the next feature on `stream` into `out` starting at `out_offset`.
 // Updates `out_offset`, and returns the number of values added.
-// Returns -1 if the next feature on `stream` doesn't match `dtype`, or if it
-// has more values than `out` has room for after `out_offset`.
+// Returns -1 if the next feature on `stream` doesn't match `dtype`, if it has
+// more values than `out` has room for after `out_offset`, or if `dtype` isn't
+// supported.
 inline int ParseFeature(DataType dtype, protobuf::io::CodedInputStream* stream,
                         Tensor* out, size_t* out_offset) {
   const size_t out_elements = out->NumElements();
@@ -2094,7 +2095,7 @@ inline int ParseFeature(DataType dtype, protobuf::io::CodedInputStream* stream,
       break;
     default:
       ReportUnexpectedDataType(dtype);
-      delta = 0;
+      return -1;
   }
   if (delta > 0) {
     *out_offset += delta;
@@ -2104,7 +2105,7 @@ inline int ParseFeature(DataType dtype, protobuf::io::CodedInputStream* stream,
 
 // Fills `num_to_pad` values of `out` starting at `out_offset` with zeros (or
 // empty strings), and updates `out_offset`. Returns false if `out` doesn't have
-// room for them.
+// room for them or if `dtype` isn't supported.
 inline bool PadFeature(DataType dtype, size_t num_to_pad, Tensor* out,
                        size_t* out_offset) {
   const size_t out_elements = out->NumElements();
@@ -2127,6 +2128,7 @@ inline bool PadFeature(DataType dtype, size_t num_to_pad, Tensor* out,
     }
     default:
       ReportUnexpectedDataType(dtype);
+      return false;
   }
   *out_offset += num_to_pad;
   return true;
@@ -2343,7 +2345,8 @@ absl::Status GetSequenceFeatureLengths(
 }
 
 // Copies src into dst[dst_offset:dst_offset+src.size], and then increments
-// dst_offset by src.size. Returns false if dst doesn't have room for src.
+// dst_offset by src.size. Returns false if dst doesn't have room for src or if
+// dtype isn't supported.
 bool CopyTensorIntoTensor(DataType dtype, const Tensor& src, Tensor* dst,
                           size_t* dst_offset) {
   const size_t src_size = src.NumElements();
@@ -2372,6 +2375,7 @@ bool CopyTensorIntoTensor(DataType dtype, const Tensor& src, Tensor* dst,
     }
     default:
       ReportUnexpectedDataType(dtype);
+      return false;
   }
   *dst_offset += src_size;
   return true;
diff --git a/tensorflow/core/util/example_proto_fast_parsing_test.cc b/tensorflow/core/util/example_proto_fast_parsing_test.cc
index 00a6753011a..a1fc2758870 100644
--- a/tensorflow/core/util/example_proto_fast_parsing_test.cc
+++ b/tensorflow/core/util/example_proto_fast_parsing_test.cc
@@ -756,6 +756,30 @@ TEST(FastParse, PadFeatureRejectsPaddingThatDoesNotFit) {
   EXPECT_EQ(out_offset, 3);
 }
 
+TEST(FastParse, HelpersRejectUnexpectedDataType) {
+  Feature feature;
+  feature.mutable_int64_list()->add_value(1);
+  const std::string serialized = feature.SerializeAsString();
+  Tensor src(DT_BOOL, TensorShape({2}));
+  Tensor out(DT_BOOL, TensorShape({4}));
+  size_t out_offset = 1;
+#ifdef NDEBUG
+  // ReportUnexpectedDataType() only DCHECKs, so the helpers must fail on
+  // their own and leave the offset alone.
+  EXPECT_EQ(TestParseFeature(DT_BOOL, serialized, &out, &out_offset), -1);
+  EXPECT_FALSE(TestPadFeature(DT_BOOL, 2, &out, &out_offset));
+  EXPECT_FALSE(TestCopyTensorIntoTensor(DT_BOOL, src, &out, &out_offset));
+  EXPECT_EQ(out_offset, 1);
+#else
+  EXPECT_DEATH(TestParseFeature(DT_BOOL, serialized, &out, &out_offset),
+               "unexpected DataType");
+  EXPECT_DEATH(TestPadFeature(DT_BOOL, 2, &out, &out_offset),
+               "unexpected DataType");
+  EXPECT_DEATH(TestCopyTensorIntoTensor(DT_BOOL, src, &out, &out_offset),
+               "unexpected DataType");
+#endif
+}
+
 TEST(FastParse, CopyTensorIntoTensorRejectsSourceThatDoesNotFit) {
   // A context dense default with more values than the output has room for.
   Tensor default_value(DT_INT64, TensorShape({3}));

From 8d11fb06ed6fcedf909e4c9ecbd40a9dc96bc370 Mon Sep 17 00:00:00 2001
From: Aditya Gaurav <mr.x.l.r.8.pride@gmail.com>
Date: Sun, 27 Sep 2026 07:38:57 +0000
Subject: [PATCH 0655/1000] Validate constant num_bits in XLA
 QuantizeAndDequantizeV3

The tf2xla lowering of QuantizeAndDequantizeV3 reads num_bits from
input 3 and uses it directly to compute the quantization bounds. The
range check that the CPU kernel applies is missing, so with
jit_compile=True a num_bits of 0 or a negative value returns a result
instead of raising InvalidArgumentError.

When input 3 folds to a constant, apply the same range check as the
CPU kernel. Dynamic num_bits is left unchanged.

Fixes #127914
---
 tensorflow/compiler/tests/BUILD               |  2 ++
 tensorflow/compiler/tests/unary_ops_test.py   | 20 +++++++++++++++++++
 .../kernels/quantize_and_dequantize_op.cc     |  7 +++++++
 3 files changed, 29 insertions(+)

diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD
index deb8ce4f21f..5f40fbd0a52 100644
--- a/tensorflow/compiler/tests/BUILD
+++ b/tensorflow/compiler/tests/BUILD
@@ -1896,7 +1896,9 @@ tf_xla_py_strict_test(
     ],
     deps = [
         ":xla_test",
+        "//tensorflow/python/eager:def_function",
         "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
         "//tensorflow/python/ops:array_ops",
         "//tensorflow/python/ops:array_ops_stack",
         "//tensorflow/python/ops:bitwise_ops",
diff --git a/tensorflow/compiler/tests/unary_ops_test.py b/tensorflow/compiler/tests/unary_ops_test.py
index 01341eb7929..a691785f52f 100644
--- a/tensorflow/compiler/tests/unary_ops_test.py
+++ b/tensorflow/compiler/tests/unary_ops_test.py
@@ -20,7 +20,9 @@ import numpy as np
 import six
 
 from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import def_function
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import array_ops_stack  # pylint: disable=g-direct-tensorflow-import
 from tensorflow.python.ops import bitwise_ops
@@ -335,6 +337,24 @@ class UnaryOpsTest(xla_test.XLATestCase):
           ),
       )
 
+  def testQuantizeAndDequantizeV3InvalidNumBits(self):
+    for num_bits in (0, -1):
+
+      @def_function.function(jit_compile=True)
+      def quantize_and_dequantize_v3(x):
+        return array_ops.quantize_and_dequantize_v3(
+            x, -1.0, 1.0, num_bits=num_bits, signed_input=True
+        )
+
+      with self.session() as session:
+        with self.test_scope():
+          x = array_ops.placeholder(dtypes.float32, [4], name="a")
+          output = quantize_and_dequantize_v3(x)
+        with self.assertRaisesRegex(
+            errors.InvalidArgumentError, "num_bits is out of range"
+        ):
+          session.run(output, {x: np.array([-1, -0.5, 0, 0.3], np.float32)})
+
   def testComplexOps(self):
     for dtype in self.complex_types:
 
diff --git a/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc b/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc
index 4f34224d79e..62c63d92797 100644
--- a/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/quantize_and_dequantize_op.cc
@@ -112,6 +112,13 @@ class QuantizeAndDequantizeOp : public XlaOpKernel {
       OP_REQUIRES(
           ctx, ctx->num_inputs() == 4,
           errors::Internal("Expected 4 inputs to QuantizeAndDequantize"));
+      int64_t num_bits_val;
+      if (ctx->ConstantInputAsIntScalar(3, &num_bits_val).ok()) {
+        OP_REQUIRES(
+            ctx, num_bits_val > 0 && num_bits_val < (signed_input_ ? 62 : 63),
+            errors::InvalidArgument("num_bits is out of range: ", num_bits_val,
+                                    " with signed_input_ ", signed_input_));
+      }
       num_bits = ctx->Input(3);
     } else {
       num_bits = xla::ConstantR0<int32_t>(b, num_bits_);

From ead4b3e19fdaf6f051f98089339d3164c65f278d Mon Sep 17 00:00:00 2001
From: Aditya Gaurav <mr.x.l.r.8.pride@gmail.com>
Date: Sun, 27 Sep 2026 07:44:57 +0000
Subject: [PATCH 0656/1000] Use self.evaluate in QuantizeAndDequantizeV3
 num_bits test

---
 tensorflow/compiler/tests/unary_ops_test.py | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/tensorflow/compiler/tests/unary_ops_test.py b/tensorflow/compiler/tests/unary_ops_test.py
index a691785f52f..bf4f44b4971 100644
--- a/tensorflow/compiler/tests/unary_ops_test.py
+++ b/tensorflow/compiler/tests/unary_ops_test.py
@@ -346,14 +346,12 @@ class UnaryOpsTest(xla_test.XLATestCase):
             x, -1.0, 1.0, num_bits=num_bits, signed_input=True
         )
 
-      with self.session() as session:
-        with self.test_scope():
-          x = array_ops.placeholder(dtypes.float32, [4], name="a")
-          output = quantize_and_dequantize_v3(x)
+      with self.test_scope():
+        x = np.array([-1, -0.5, 0, 0.3], dtype=np.float32)
         with self.assertRaisesRegex(
             errors.InvalidArgumentError, "num_bits is out of range"
         ):
-          session.run(output, {x: np.array([-1, -0.5, 0, 0.3], np.float32)})
+          self.evaluate(quantize_and_dequantize_v3(x))
 
   def testComplexOps(self):
     for dtype in self.complex_types:

From c0d2cc05947d15a93be96ec5dda21094b8876ca1 Mon Sep 17 00:00:00 2001
From: Aditya Gaurav <mr.x.l.r.8.pride@gmail.com>
Date: Mon, 28 Sep 2026 13:47:26 +0000
Subject: [PATCH 0657/1000] Test QuantizeAndDequantizeV3 num_bits upper bounds
 with an explicit session

---
 tensorflow/compiler/tests/unary_ops_test.py | 34 +++++++++++++--------
 1 file changed, 22 insertions(+), 12 deletions(-)

diff --git a/tensorflow/compiler/tests/unary_ops_test.py b/tensorflow/compiler/tests/unary_ops_test.py
index bf4f44b4971..8081411771c 100644
--- a/tensorflow/compiler/tests/unary_ops_test.py
+++ b/tensorflow/compiler/tests/unary_ops_test.py
@@ -338,20 +338,30 @@ class UnaryOpsTest(xla_test.XLATestCase):
       )
 
   def testQuantizeAndDequantizeV3InvalidNumBits(self):
-    for num_bits in (0, -1):
+    for signed_input, invalid_num_bits in (
+        (True, (0, -1, 62, 63)),
+        (False, (0, -1, 63, 64)),
+    ):
+      for num_bits in invalid_num_bits:
 
-      @def_function.function(jit_compile=True)
-      def quantize_and_dequantize_v3(x):
-        return array_ops.quantize_and_dequantize_v3(
-            x, -1.0, 1.0, num_bits=num_bits, signed_input=True
-        )
-
-      with self.test_scope():
-        x = np.array([-1, -0.5, 0, 0.3], dtype=np.float32)
-        with self.assertRaisesRegex(
-            errors.InvalidArgumentError, "num_bits is out of range"
+        @def_function.function(jit_compile=True)
+        def quantize_and_dequantize_v3(
+            x, num_bits=num_bits, signed_input=signed_input
         ):
-          self.evaluate(quantize_and_dequantize_v3(x))
+          return array_ops.quantize_and_dequantize_v3(
+              x, -1.0, 1.0, num_bits=num_bits, signed_input=signed_input
+          )
+
+        with self.session() as sess:
+          with self.test_scope():
+            x = array_ops.placeholder(dtypes.float32)
+            output = quantize_and_dequantize_v3(x)
+          with self.assertRaisesRegex(
+              errors.InvalidArgumentError, "num_bits is out of range"
+          ):
+            sess.run(
+                output, {x: np.array([-1, -0.5, 0, 0.3], dtype=np.float32)}
+            )
 
   def testComplexOps(self):
     for dtype in self.complex_types:

From 765d6e4de75c8e015c01a645e43bd34a08bc9d7a Mon Sep 17 00:00:00 2001
From: Aditya Gaurav <mr.x.l.r.8.pride@gmail.com>
Date: Sat, 3 Oct 2026 06:10:36 +0000
Subject: [PATCH 0658/1000] Test valid QuantizeAndDequantizeV3 num_bits bounds
 and use device_scope

---
 tensorflow/compiler/tests/unary_ops_test.py | 24 ++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/tensorflow/compiler/tests/unary_ops_test.py b/tensorflow/compiler/tests/unary_ops_test.py
index 8081411771c..626af157c7c 100644
--- a/tensorflow/compiler/tests/unary_ops_test.py
+++ b/tensorflow/compiler/tests/unary_ops_test.py
@@ -353,7 +353,7 @@ class UnaryOpsTest(xla_test.XLATestCase):
           )
 
         with self.session() as sess:
-          with self.test_scope():
+          with self.device_scope():
             x = array_ops.placeholder(dtypes.float32)
             output = quantize_and_dequantize_v3(x)
           with self.assertRaisesRegex(
@@ -363,6 +363,28 @@ class UnaryOpsTest(xla_test.XLATestCase):
                 output, {x: np.array([-1, -0.5, 0, 0.3], dtype=np.float32)}
             )
 
+    for signed_input, valid_num_bits in (
+        (True, (1, 61)),
+        (False, (1, 62)),
+    ):
+      for num_bits in valid_num_bits:
+
+        @def_function.function(jit_compile=True)
+        def quantize_and_dequantize_v3_valid(
+            x, num_bits=num_bits, signed_input=signed_input
+        ):
+          return array_ops.quantize_and_dequantize_v3(
+              x, -1.0, 1.0, num_bits=num_bits, signed_input=signed_input
+          )
+
+        with self.session() as sess:
+          with self.device_scope():
+            x = array_ops.placeholder(dtypes.float32)
+            output = quantize_and_dequantize_v3_valid(x)
+          sess.run(
+              output, {x: np.array([-1, -0.5, 0, 0.3], dtype=np.float32)}
+          )
+
   def testComplexOps(self):
     for dtype in self.complex_types:
 

From 30c975e46ffe11cb4659aed8b725099b906d1cce Mon Sep 17 00:00:00 2001
From: Aditya Gaurav <mr.x.l.r.8.pride@gmail.com>
Date: Sat, 3 Oct 2026 10:42:10 +0000
Subject: [PATCH 0659/1000] Rename QuantizeAndDequantizeV3 num_bits test to
 NumBitsBoundaries

---
 tensorflow/compiler/tests/unary_ops_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/compiler/tests/unary_ops_test.py b/tensorflow/compiler/tests/unary_ops_test.py
index 626af157c7c..2a13259f0a7 100644
--- a/tensorflow/compiler/tests/unary_ops_test.py
+++ b/tensorflow/compiler/tests/unary_ops_test.py
@@ -337,7 +337,7 @@ class UnaryOpsTest(xla_test.XLATestCase):
           ),
       )
 
-  def testQuantizeAndDequantizeV3InvalidNumBits(self):
+  def testQuantizeAndDequantizeV3NumBitsBoundaries(self):
     for signed_input, invalid_num_bits in (
         (True, (0, -1, 62, 63)),
         (False, (0, -1, 63, 64)),

From 24326ec17d142a2eb0c8806f2c0ec8e0c9ffeb17 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Tue, 13 Oct 2026 01:45:34 -0700
Subject: [PATCH 0660/1000] [XLA:GPU]: Support per-output tiled tuple producers
 in TiledHlo.

PiperOrigin-RevId: 998467622
---
 .../xla/xla/codegen/tiling/experimental/BUILD |   1 +
 .../codegen/tiling/experimental/tiled_hlo.cc  | 107 ++++++-
 .../tiling/experimental/tiled_hlo_test.cc     | 295 ++++++++++++++++++
 3 files changed, 389 insertions(+), 14 deletions(-)

diff --git a/third_party/xla/xla/codegen/tiling/experimental/BUILD b/third_party/xla/xla/codegen/tiling/experimental/BUILD
index ce30f449e7f..321028eed61 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/BUILD
+++ b/third_party/xla/xla/codegen/tiling/experimental/BUILD
@@ -280,6 +280,7 @@ cc_library(
         "@com_google_absl//absl/hash",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo.cc b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo.cc
index 48e31f49b29..69da047a424 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo.cc
@@ -35,12 +35,15 @@ limitations under the License.
 #include "absl/hash/hash.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
+#include "absl/status/status.h"
 #include "absl/status/status_macros.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SmallVector.h"
@@ -137,7 +140,77 @@ class OrderedTiledHloPtrSet {
                   llvm::SmallVector<experimental::Tile, 2>{std::move(tile)});
   }
 
-  // Same as above for an instruction with one tile per result.
+  // Inserts an operand of `user` into the set. For tuple-producing operands
+  // whose outputs do not share a single tile, buffers each output's tile across
+  // its `get-tuple-element` users and only marks the producer as newly inserted
+  // once all outputs have been tiled.
+  absl::StatusOr<std::pair<TiledHloInstruction*, bool>> InsertOperand(
+      const HloInstruction& user, const HloInstructionAdaptor& operand_adaptor,
+      Tile tile) {
+    const HloInstruction* operand_hlo = &operand_adaptor.instruction();
+    if (!operand_hlo->shape().IsTuple() ||
+        !tile.tiling_space().HasPerOutputTiles(operand_adaptor)) {
+      return Insert(operand_hlo, std::move(tile));
+    }
+    if (user.opcode() != HloOpcode::kGetTupleElement) {
+      return absl::UnimplementedError(absl::StrCat(
+          "Tuple-producing instruction ", operand_hlo->name(),
+          " must be consumed through get-tuple-element, but is used by ",
+          user.ToString()));
+    }
+
+    PendingTupleProducer& pending = pending_tuple_producers_[operand_hlo];
+    const int64_t num_outputs = operand_hlo->shape().tuple_shapes().size();
+    if (pending.tiled_hlo == nullptr) {
+      pending.tiles.resize(num_outputs);
+      pending.tiled_hlo = &instructions_.emplace_back(
+          operand_hlo, llvm::SmallVector<experimental::Tile, 2>{});
+    }
+
+    const int64_t index = user.tuple_index();
+    if (pending.tiles[index].has_value()) {
+      if (*pending.tiles[index] != tile) {
+        return absl::UnimplementedError(
+            absl::StrCat("Output ", index, " of ", operand_hlo->name(),
+                         " is tiled with multiple incompatible tiles."));
+      }
+      return std::make_pair(pending.tiled_hlo, false);
+    }
+    pending.tiles[index] = std::move(tile);
+    // Keep the instruction pending until all outputs have a tile, then
+    // materialize it below.
+    if (++pending.num_assigned < num_outputs) {
+      return std::make_pair(pending.tiled_hlo, false);
+    }
+
+    llvm::SmallVector<experimental::Tile, 2> tiles;
+    tiles.reserve(num_outputs);
+    for (const std::optional<experimental::Tile>& t : pending.tiles) {
+      tiles.push_back(*t);
+    }
+    pending.tiled_hlo->set_tiles(std::move(tiles));
+    hash_set_.insert(pending.tiled_hlo);
+    insertion_order_.push_back(pending.tiled_hlo);
+    return std::make_pair(pending.tiled_hlo, true);
+  }
+
+  // Consumes the set, returning pointers to its elements in insertion order.
+  absl::StatusOr<std::vector<TiledHloInstruction*>>
+  ConsumeInInsertionOrder() && {
+    // Validate that all pending tuple producers have been inserted.
+    for (const auto& [tuple_producer, pending] : pending_tuple_producers_) {
+      for (int64_t k = 0; k < pending.tiles.size(); ++k) {
+        if (!pending.tiles[k].has_value()) {
+          return absl::UnimplementedError(
+              absl::StrCat("Output ", k, " of ", tuple_producer->name(),
+                           " is not consumed in the fusion."));
+        }
+      }
+    }
+    return std::move(insertion_order_);
+  }
+
+ private:
   std::pair<TiledHloInstruction*, bool> Insert(
       const HloInstruction* hlo,
       llvm::SmallVector<experimental::Tile, 2> tiles) {
@@ -152,12 +225,6 @@ class OrderedTiledHloPtrSet {
     return {&candidate, true};
   }
 
-  // Consumes the set, returning pointers to its elements in insertion order.
-  std::vector<TiledHloInstruction*> ConsumeInInsertionOrder() && {
-    return std::move(insertion_order_);
-  }
-
- private:
   struct PtrHash {
     size_t operator()(const TiledHloInstruction* v) const {
       return absl::HashOf(v->hlo(), v->tiles());
@@ -172,6 +239,12 @@ class OrderedTiledHloPtrSet {
     }
   };
 
+  struct PendingTupleProducer {
+    TiledHloInstruction* tiled_hlo = nullptr;
+    llvm::SmallVector<std::optional<experimental::Tile>, 2> tiles;
+    int64_t num_assigned = 0;
+  };
+
   // Stores non-owning pointers to the elements in the set. Elements are
   // compared by the value behind the pointer, not the pointer itself.
   absl::flat_hash_set<TiledHloInstruction*, PtrHash, PtrEqual> hash_set_;
@@ -181,6 +254,11 @@ class OrderedTiledHloPtrSet {
 
   // Pointers to the elements of this set, in insertion order.
   std::vector<TiledHloInstruction*> insertion_order_;
+
+  // Tuple-producing instructions whose per-output tiles are being accumulated
+  // across their `get-tuple-element` users before being finalized in the set.
+  llvm::MapVector<const HloInstruction*, PendingTupleProducer>
+      pending_tuple_producers_;
 };
 
 // Sorts tiled hlo instructions in def-before-use order, starting from
@@ -510,7 +588,7 @@ absl::StatusOr<TiledHloRegion> CreateHloRegion(
 
     ABSL_ASSIGN_OR_RETURN(
         auto operands_tiles,
-        PropagateTileToInput(tiling_space, *hlo, tiled_hlo->tile(), 0));
+        PropagateTilesToInputs(tiling_space, *hlo, tiled_hlo->tiles()));
 
     RegionSchema spec = GetRegionSchema(*tiled_hlo);
 
@@ -546,9 +624,9 @@ absl::StatusOr<TiledHloRegion> CreateHloRegion(
       // Hoists and deduplicates loop-invariant instructions.
       RegionContext* target_region =
           FindTargetRegion(&current_context, operands_tiles[id]);
-      auto [operand_tiled_hlo, inserted] =
-          target_region->instructions_set->Insert(
-              &operands[id].instruction(), std::move(operands_tiles[id]));
+      ABSL_ASSIGN_OR_RETURN((auto [operand_tiled_hlo, inserted]),
+                       target_region->instructions_set->InsertOperand(
+                           *hlo, operands[id], std::move(operands_tiles[id])));
       resolved_operands[id] = operand_tiled_hlo;
       if (inserted) {
         target_region->worklist->push_back(operand_tiled_hlo);
@@ -569,9 +647,10 @@ absl::StatusOr<TiledHloRegion> CreateHloRegion(
     }
   }
 
-  return TiledHloRegion{
-      std::move(tiled_hlo_instructions_set).ConsumeInInsertionOrder(),
-      std::move(canonical_roots)};
+  ABSL_ASSIGN_OR_RETURN(
+      std::vector<TiledHloInstruction*> instructions,
+      std::move(tiled_hlo_instructions_set).ConsumeInInsertionOrder());
+  return TiledHloRegion{std::move(instructions), std::move(canonical_roots)};
 }
 
 }  // namespace
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc
index 7edc81b2bda..2aab60bab0e 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc
@@ -1188,6 +1188,301 @@ TEST_P(TileAnalysisTest, ScanWithLoop) {
   )"));
 }
 
+// Tests for index-wise variadic instructions (currently: multi-operand
+// all-gather). Such an instruction is consumed through one `get-tuple-element`
+// per output, and is tiled as a single `TiledHloInstruction` holding one tile
+// per output.
+class IndexWiseVariadicTileAnalysisTest : public TileAnalysisTestBase {
+ protected:
+  static std::vector<const TiledHloInstruction*> FindAll(
+      const TiledHloComputation& tiled_computation, HloOpcode opcode) {
+    std::vector<const TiledHloInstruction*> result;
+    for (const TiledHloInstruction* instruction :
+         tiled_computation.instructions()) {
+      if (instruction->hlo()->opcode() == opcode) {
+        result.push_back(instruction);
+      }
+    }
+    return result;
+  }
+
+  static int64_t IndexOf(const TiledHloComputation& tiled_computation,
+                         const TiledHloInstruction* instruction) {
+    absl::Span<const TiledHloInstruction* const> instructions =
+        tiled_computation.instructions();
+    return absl::c_find(instructions, instruction) - instructions.begin();
+  }
+};
+
+TEST_F(IndexWiseVariadicTileAnalysisTest,
+       VariadicAllGatherIsTiledAsOneNodeWithPerOutputTiles) {
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"hlo(
+    f {
+      p0 = f32[1,2,8,16] parameter(0)
+      p1 = s32[1,2] parameter(1)
+      ag = (f32[4,2,8,16], s32[4,2]) all-gather(p0, p1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      gte0 = f32[4,2,8,16] get-tuple-element(ag), index=0
+      gte1 = s32[4,2] get-tuple-element(ag), index=1
+      ROOT t = (f32[4,2,8,16], s32[4,2]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      arg0 = f32[1,2,8,16] parameter(0)
+      arg1 = s32[1,2] parameter(1)
+      ROOT fusion = (f32[4,2,8,16], s32[4,2]) fusion(arg0, arg1),
+        kind=kCustom, calls=f
+    })hlo",
+                                    {1, 1, 4, 16, 1, 1}));
+
+  // The dimensions of the outputs are laid out back to back in the ordered
+  // dimension list of the all-gather: output 0 at positions 0-3, output 1 at
+  // 4-5.
+  EXPECT_THAT(tiled_computation, MatchString(R"(
+    Dimensions:
+      0 type: parallel size: 4 tile size: 1 dim ID:0 hlo: %gte0 = f32[4,2,8,16]{3,2,1,0} get-tuple-element(%ag), index=0
+      1 type: parallel size: 2 tile size: 1 dim ID:1 hlo: %gte0 = f32[4,2,8,16]{3,2,1,0} get-tuple-element(%ag), index=0
+      2 type: parallel size: 8 tile size: 4 dim ID:2 hlo: %gte0 = f32[4,2,8,16]{3,2,1,0} get-tuple-element(%ag), index=0
+      3 type: parallel size: 16 tile size: 16 dim ID:3 hlo: %gte0 = f32[4,2,8,16]{3,2,1,0} get-tuple-element(%ag), index=0
+      4 type: parallel size: 4 tile size: 1 dim ID:0 hlo: %gte1 = s32[4,2]{1,0} get-tuple-element(%ag), index=1
+      5 type: parallel size: 2 tile size: 1 dim ID:1 hlo: %gte1 = s32[4,2]{1,0} get-tuple-element(%ag), index=1
+    Root tiles:
+      0 root tile:  offsets [tid_0, tid_1, tid_2 * 4, 0] sizes [1, 1, 4, 16] strides [1, 1, 1, 1] upper bounds [4, 2, 8, 16]
+      1 root tile:  offsets [tid_4, tid_5] sizes [1, 1] strides [1, 1] upper bounds [4, 2]
+
+    Tiled HLO:
+      arg0.tile_0 = parameter(0)  offsets [0, tid_1, tid_2 * 4, 0] sizes [1, 1, 4, 16] strides [1, 1, 1, 1] upper bounds [1, 2, 8, 16]
+                                replica ids { offsets [tid_0] sizes [1] strides [1] upper bounds [4] }
+      arg1.tile_0 = parameter(1)  offsets [0, tid_5] sizes [1, 1] strides [1, 1] upper bounds [1, 2]
+                                replica ids { offsets [tid_4] sizes [1] strides [1] upper bounds [4] }
+      ag.tile_0 = all-gather(arg0.tile_0, arg1.tile_0)
+        #0  offsets [tid_0, tid_1, tid_2 * 4, 0] sizes [1, 1, 4, 16] strides [1, 1, 1, 1] upper bounds [4, 2, 8, 16]
+        #1  offsets [tid_4, tid_5] sizes [1, 1] strides [1, 1] upper bounds [4, 2]
+      gte0.tile_0 = get-tuple-element(ag.tile_0)  offsets [tid_0, tid_1, tid_2 * 4, 0] sizes [1, 1, 4, 16] strides [1, 1, 1, 1] upper bounds [4, 2, 8, 16]
+      gte1.tile_0 = get-tuple-element(ag.tile_0)  offsets [tid_4, tid_5] sizes [1, 1] strides [1, 1] upper bounds [4, 2]
+  )"));
+
+  std::vector<const TiledHloInstruction*> all_gathers =
+      FindAll(tiled_computation, HloOpcode::kAllGather);
+  ASSERT_EQ(all_gathers.size(), 1);
+  const TiledHloInstruction* ag = all_gathers[0];
+  EXPECT_THAT(ag,
+              IsHloWithOperands(HloOpcode::kAllGather,
+                                std::vector<HloOpcode>{HloOpcode::kParameter,
+                                                       HloOpcode::kParameter}));
+  ASSERT_EQ(ag->operands().size(), 2);
+  ASSERT_EQ(ag->tiles().size(), 2);
+  EXPECT_EQ(ag->operand(0)->hlo()->name(), "arg0");
+  EXPECT_EQ(ag->operand(1)->hlo()->name(), "arg1");
+  // Each operand carries one replica id, derived from its own output's tile
+  // id. The scheduler is responsible for mapping both to the same peer.
+  EXPECT_EQ(ag->operand(0)->tile().num_replica_ids(), 1);
+  EXPECT_EQ(ag->operand(1)->tile().num_replica_ids(), 1);
+
+  // Ordering: both operands precede the all-gather, which precedes both GTEs.
+  const int64_t ag_index = IndexOf(tiled_computation, ag);
+  EXPECT_LT(IndexOf(tiled_computation, ag->operand(0)), ag_index);
+  EXPECT_LT(IndexOf(tiled_computation, ag->operand(1)), ag_index);
+  EXPECT_GT(IndexOf(tiled_computation, tiled_computation.roots()[0]), ag_index);
+  EXPECT_GT(IndexOf(tiled_computation, tiled_computation.roots()[1]), ag_index);
+
+  // One GTE root per output, both consuming the same tiled all-gather.
+  ASSERT_EQ(tiled_computation.roots().size(), 2);
+  EXPECT_EQ(tiled_computation.roots()[0]->operand(0), ag);
+  EXPECT_EQ(tiled_computation.roots()[1]->operand(0), ag);
+  EXPECT_EQ(tiled_computation.roots()[0]->tile(), ag->tile(0));
+  EXPECT_EQ(tiled_computation.roots()[1]->tile(), ag->tile(1));
+}
+
+TEST_F(IndexWiseVariadicTileAnalysisTest,
+       VariadicAllGatherConsumedByGetTupleElementAndConvert) {
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"hlo(
+    f {
+      p0 = f32[2,16] parameter(0)
+      p1 = s32[4] parameter(1)
+      ag = (f32[8,16], s32[16]) all-gather(p0, p1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      gte0 = f32[8,16] get-tuple-element(ag), index=0
+      gte1 = s32[16] get-tuple-element(ag), index=1
+      c0 = bf16[8,16] convert(gte0)
+      c1 = f32[16] convert(gte1)
+      ROOT t = (bf16[8,16], f32[16]) tuple(c0, c1)
+    }
+
+    ENTRY e {
+      arg0 = f32[2,16] parameter(0)
+      arg1 = s32[4] parameter(1)
+      ROOT fusion = (bf16[8,16], f32[16]) fusion(arg0, arg1),
+        kind=kCustom, calls=f
+    })hlo",
+                                    {2, 8, 2}));
+
+  std::vector<const TiledHloInstruction*> all_gathers =
+      FindAll(tiled_computation, HloOpcode::kAllGather);
+  ASSERT_EQ(all_gathers.size(), 1);
+  const TiledHloInstruction* ag = all_gathers[0];
+  ASSERT_EQ(ag->operands().size(), 2);
+  ASSERT_EQ(ag->tiles().size(), 2);
+  EXPECT_EQ(ag->operand(0)->tile().num_replica_ids(), 1);
+  EXPECT_EQ(ag->operand(1)->tile().num_replica_ids(), 1);
+
+  std::vector<const TiledHloInstruction*> gtes =
+      FindAll(tiled_computation, HloOpcode::kGetTupleElement);
+  ASSERT_EQ(gtes.size(), 2);
+  for (const TiledHloInstruction* gte : gtes) {
+    ASSERT_EQ(gte->operands().size(), 1);
+    EXPECT_EQ(gte->operand(0), ag);
+  }
+}
+
+TEST_F(IndexWiseVariadicTileAnalysisTest,
+       VariadicAllGatherWithIdenticalOperandShapes) {
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"hlo(
+    f {
+      p0 = s32[1,2] parameter(0)
+      p1 = s32[1,2] parameter(1)
+      p2 = s32[1,2] parameter(2)
+      ag = (s32[4,2], s32[4,2], s32[4,2]) all-gather(p0, p1, p2),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      gte0 = s32[4,2] get-tuple-element(ag), index=0
+      gte1 = s32[4,2] get-tuple-element(ag), index=1
+      gte2 = s32[4,2] get-tuple-element(ag), index=2
+      ROOT t = (s32[4,2], s32[4,2], s32[4,2]) tuple(gte0, gte1, gte2)
+    }
+
+    ENTRY e {
+      arg0 = s32[1,2] parameter(0)
+      arg1 = s32[1,2] parameter(1)
+      arg2 = s32[1,2] parameter(2)
+      ROOT fusion = (s32[4,2], s32[4,2], s32[4,2]) fusion(arg0, arg1, arg2),
+        kind=kCustom, calls=f
+    })hlo",
+                                    {1, 1, 1, 1, 1, 1}));
+
+  std::vector<const TiledHloInstruction*> all_gathers =
+      FindAll(tiled_computation, HloOpcode::kAllGather);
+  ASSERT_EQ(all_gathers.size(), 1);
+  const TiledHloInstruction* ag = all_gathers[0];
+  ASSERT_EQ(ag->operands().size(), 3);
+  ASSERT_EQ(ag->tiles().size(), 3);
+  ASSERT_EQ(tiled_computation.roots().size(), 3);
+  for (int64_t k = 0; k < 3; ++k) {
+    EXPECT_EQ(ag->operand(k)->hlo()->parameter_number(), k);
+    EXPECT_EQ(tiled_computation.roots()[k]->operand(0), ag);
+    EXPECT_EQ(tiled_computation.roots()[k]->hlo()->tuple_index(), k);
+    EXPECT_EQ(ag->tile(k), tiled_computation.roots()[k]->tile());
+  }
+  EXPECT_EQ(FindAll(tiled_computation, HloOpcode::kParameter).size(), 3);
+}
+
+TEST_F(IndexWiseVariadicTileAnalysisTest,
+       VariadicAllGatherWithGeneralOperandExtentAndDotProducers) {
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"hlo(
+    f {
+      lhs0 = f32[2,8] parameter(0)
+      rhs0 = f32[8,16] parameter(1)
+      lhs1 = f32[4,8] parameter(2)
+      rhs1 = f32[8,32] parameter(3)
+      dot0 = f32[2,16] dot(lhs0, rhs0),
+        lhs_contracting_dims={1}, rhs_contracting_dims={0}
+      dot1 = f32[4,32] dot(lhs1, rhs1),
+        lhs_contracting_dims={1}, rhs_contracting_dims={0}
+      ag = (f32[8,16], f32[16,32]) all-gather(dot0, dot1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      gte0 = f32[8,16] get-tuple-element(ag), index=0
+      gte1 = f32[16,32] get-tuple-element(ag), index=1
+      ROOT t = (f32[8,16], f32[16,32]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      lhs0 = f32[2,8] parameter(0)
+      rhs0 = f32[8,16] parameter(1)
+      lhs1 = f32[4,8] parameter(2)
+      rhs1 = f32[8,32] parameter(3)
+      ROOT fusion = (f32[8,16], f32[16,32]) fusion(lhs0, rhs0, lhs1, rhs1),
+        kind=kCustom, calls=f
+    })hlo",
+                                    {2, 8, 2, 16, 8, 8}));
+
+  std::vector<const TiledHloInstruction*> all_gathers =
+      FindAll(tiled_computation, HloOpcode::kAllGather);
+  ASSERT_EQ(all_gathers.size(), 1);
+  EXPECT_EQ(FindAll(tiled_computation, HloOpcode::kDot).size(), 2);
+}
+
+TEST_F(IndexWiseVariadicTileAnalysisTest,
+       VariadicAllGatherWithUnconsumedOutputIsRejected) {
+  EXPECT_THAT(
+      ParseAndTile(R"hlo(
+    f {
+      p0 = f32[1,8] parameter(0)
+      p1 = s32[1,2] parameter(1)
+      ag = (f32[4,8], s32[4,2]) all-gather(p0, p1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      ROOT gte0 = f32[4,8] get-tuple-element(ag), index=0
+    }
+
+    ENTRY e {
+      arg0 = f32[1,8] parameter(0)
+      arg1 = s32[1,2] parameter(1)
+      ROOT fusion = f32[4,8] fusion(arg0, arg1), kind=kCustom, calls=f
+    })hlo",
+                   {1, 4}),
+      StatusIs(absl::StatusCode::kUnimplemented,
+               HasSubstr("Output 1 of ag is not consumed in the fusion.")));
+}
+
+TEST_F(SameShapeMultiOutputFusionTileAnalysisTest,
+       VariadicReduceIsTiledAsOneNode) {
+  // A variadic reduce is not index-wise: every output depends on every input.
+  // Its GTE roots share one tile and must resolve to a single tiled node.
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"hlo(
+    add_pair {
+      a = f32[] parameter(0)
+      b = s32[] parameter(1)
+      c = f32[] parameter(2)
+      d = s32[] parameter(3)
+      add_f = f32[] add(a, c)
+      add_s = s32[] add(b, d)
+      ROOT t = (f32[], s32[]) tuple(add_f, add_s)
+    }
+
+    f {
+      p0 = f32[4,8] parameter(0)
+      p1 = s32[4,8] parameter(1)
+      c0 = f32[] constant(0)
+      c1 = s32[] constant(0)
+      reduce = (f32[4], s32[4]) reduce(p0, p1, c0, c1), dimensions={1},
+        to_apply=add_pair
+      gte0 = f32[4] get-tuple-element(reduce), index=0
+      gte1 = s32[4] get-tuple-element(reduce), index=1
+      ROOT t = (f32[4], s32[4]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      arg0 = f32[4,8] parameter(0)
+      arg1 = s32[4,8] parameter(1)
+      ROOT fusion = (f32[4], s32[4]) fusion(arg0, arg1), kind=kLoop, calls=f
+    })hlo",
+                                    {2, 8}));
+
+  std::vector<const TiledHloInstruction*> reduces;
+  for (const TiledHloInstruction* instruction :
+       tiled_computation.instructions()) {
+    if (instruction->hlo()->opcode() == HloOpcode::kReduce) {
+      reduces.push_back(instruction);
+    }
+  }
+  ASSERT_EQ(reduces.size(), 1);
+  EXPECT_EQ(reduces[0]->tiles().size(), 1);
+  EXPECT_EQ(reduces[0]->operands().size(), 4);
+}
+
 // TODO(b/422676780): Port the remaining tests.
 
 }  // namespace

From 1730485b8fbfc212e1f3ca8932c6bf0904836028 Mon Sep 17 00:00:00 2001
From: Levon Ter-Grigoryan <patrios@google.com>
Date: Tue, 13 Oct 2026 01:51:07 -0700
Subject: [PATCH 0661/1000] [XLA:GPU] Speed up
 CollectivePipelineParallelismTest.JaxExampleWithDecomposedCycle by shrinking
 input shapes.

`JaxExampleWithDecomposedCycle` previously used `f32[4,4096,4096]` and `f32[4,5,4096,8192]` input tensors (2.95 GB total), which caused single-threaded CPU random input generation in `MakeFakeArguments`, synchronous unsharded host-to-device transfers across 4 GPUs, and 13 loop iterations of large GEMMs and collectives to take over 100 seconds on 8xH100 CI.

None of the pipeline parallelism or collective-permute decomposition passes exercised by this test depend on the 4096 and 8192 contracting and non-contracting dimensions (`xla_gpu_collective_permute_decomposer_threshold` is set to 0, and SPMD partitioning shards only along dimension 0). Shrinking those dimensions to 16 and 32 to match the other tests in the suite reduces the tensor volume by 65,536x (to 45 KB) and cuts `MakeFakeArguments` time from 8.3s to under 1ms while preserving full test coverage.

PiperOrigin-RevId: 998470491
---
 .../collective_pipeline_parallelism_test.cc   | 206 +++++++++---------
 1 file changed, 103 insertions(+), 103 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/tests/collective_pipeline_parallelism_test.cc b/third_party/xla/xla/backends/gpu/tests/collective_pipeline_parallelism_test.cc
index ee59c1e17bc..d7e1589b712 100644
--- a/third_party/xla/xla/backends/gpu/tests/collective_pipeline_parallelism_test.cc
+++ b/third_party/xla/xla/backends/gpu/tests/collective_pipeline_parallelism_test.cc
@@ -1689,8 +1689,8 @@ TEST_P(CollectivePipelineParallelismTest,
 TEST_P(CollectivePipelineParallelismTest, JaxExampleWithDecomposedCycle) {
   constexpr char kModuleStr[] = R"(
 HloModule jit_entry_computation, entry_computation_layout={
-    (f32[4,4096,4096]{2,1,0}, f32[4,5,4096,8192]{3,2,1,0})->
-    f32[4,5,4096,8192]{3,2,1,0}},
+    (f32[4,16,16]{2,1,0}, f32[4,5,16,32]{3,2,1,0})->
+    f32[4,5,16,32]{3,2,1,0}},
     allow_spmd_sharding_propagation_to_parameters={false,false},
     allow_spmd_sharding_propagation_to_output={true}, num_partitions=4
 
@@ -1719,31 +1719,31 @@ HloModule jit_entry_computation, entry_computation_layout={
   ROOT %select.29 = s32[] select(%and.27, %add.28, %remainder.22)
 }
 
-%_pad.30 (Arg_0.31: f32[4,1,4096,8192], Arg_1.32: s32[]) -> f32[5,1,4096,8192] {
-  %Arg_0.31 = f32[4,1,4096,8192]{3,2,1,0} parameter(0)
+%_pad.30 (Arg_0.31: f32[4,1,16,32], Arg_1.32: s32[]) -> f32[5,1,16,32] {
+  %Arg_0.31 = f32[4,1,16,32]{3,2,1,0} parameter(0)
   %Arg_1.32 = s32[] parameter(1)
   %convert.33 = f32[] convert(%Arg_1.32)
-  ROOT %pad.34 = f32[5,1,4096,8192]{3,2,1,0} pad(%Arg_0.31, %convert.33),
+  ROOT %pad.34 = f32[5,1,16,32]{3,2,1,0} pad(%Arg_0.31, %convert.33),
       padding=1_0x0_0x0_0x0_0
 }
 
-%_where_0.35 (Arg_0.36: pred[], Arg_1.37: f32[4,1,4096,8192],
-    Arg_2.38: f32[4,1,4096,8192]) -> f32[4,1,4096,8192] {
+%_where_0.35 (Arg_0.36: pred[], Arg_1.37: f32[4,1,16,32],
+    Arg_2.38: f32[4,1,16,32]) -> f32[4,1,16,32] {
   %Arg_0.36 = pred[] parameter(0)
-  %broadcast.39 = pred[4,1,4096,8192]{3,2,1,0} broadcast(%Arg_0.36),
+  %broadcast.39 = pred[4,1,16,32]{3,2,1,0} broadcast(%Arg_0.36),
       dimensions={}
-  %Arg_1.37 = f32[4,1,4096,8192]{3,2,1,0} parameter(1)
-  %Arg_2.38 = f32[4,1,4096,8192]{3,2,1,0} parameter(2)
-  ROOT %select.40 = f32[4,1,4096,8192]{3,2,1,0} select(%broadcast.39, %Arg_1.37,
+  %Arg_1.37 = f32[4,1,16,32]{3,2,1,0} parameter(1)
+  %Arg_2.38 = f32[4,1,16,32]{3,2,1,0} parameter(2)
+  ROOT %select.40 = f32[4,1,16,32]{3,2,1,0} select(%broadcast.39, %Arg_1.37,
       %Arg_2.38)
 }
 
-%_where_1.41 (Arg_0.42: pred[4,1,4096,8192], Arg_1.43: f32[4,1,4096,8192],
-    Arg_2.44: f32[4,1,4096,8192]) -> f32[4,1,4096,8192] {
-  %Arg_0.42 = pred[4,1,4096,8192]{3,2,1,0} parameter(0)
-  %Arg_1.43 = f32[4,1,4096,8192]{3,2,1,0} parameter(1)
-  %Arg_2.44 = f32[4,1,4096,8192]{3,2,1,0} parameter(2)
-  ROOT %select.45 = f32[4,1,4096,8192]{3,2,1,0} select(%Arg_0.42, %Arg_1.43,
+%_where_1.41 (Arg_0.42: pred[4,1,16,32], Arg_1.43: f32[4,1,16,32],
+    Arg_2.44: f32[4,1,16,32]) -> f32[4,1,16,32] {
+  %Arg_0.42 = pred[4,1,16,32]{3,2,1,0} parameter(0)
+  %Arg_1.43 = f32[4,1,16,32]{3,2,1,0} parameter(1)
+  %Arg_2.44 = f32[4,1,16,32]{3,2,1,0} parameter(2)
+  ROOT %select.45 = f32[4,1,16,32]{3,2,1,0} select(%Arg_0.42, %Arg_1.43,
       %Arg_2.44)
 }
 
@@ -1772,12 +1772,12 @@ HloModule jit_entry_computation, entry_computation_layout={
   ROOT %select.65 = s32[] select(%and.63, %add.64, %remainder.58)
 }
 
-%_pad_2.66 (Arg_0.67: f32[4,1,4096,8192], Arg_1.68: s32[])
-    -> f32[7,1,4096,8192] {
-  %Arg_0.67 = f32[4,1,4096,8192]{3,2,1,0} parameter(0)
+%_pad_2.66 (Arg_0.67: f32[4,1,16,32], Arg_1.68: s32[])
+    -> f32[7,1,16,32] {
+  %Arg_0.67 = f32[4,1,16,32]{3,2,1,0} parameter(0)
   %Arg_1.68 = s32[] parameter(1)
   %convert.69 = f32[] convert(%Arg_1.68)
-  ROOT %pad.70 = f32[7,1,4096,8192]{3,2,1,0} pad(%Arg_0.67, %convert.69),
+  ROOT %pad.70 = f32[7,1,16,32]{3,2,1,0} pad(%Arg_0.67, %convert.69),
       padding=0_3x0_0x0_0x0_0
 }
 
@@ -1806,21 +1806,21 @@ HloModule jit_entry_computation, entry_computation_layout={
   ROOT %select.90 = s32[] select(%and.88, %add.89, %remainder.83)
 }
 
-%None.91 (Arg_0.92: f32[4,4096,4096], Arg_1.93: f32[4,5,4096,8192],
-    Arg_2.94: f32[4,5,4096,8192], Arg_3.95: f32[4,1,4096,8192],
-    Arg_4.96: f32[4,1,4096,8192], Arg_5.97: s32[])
-    -> (f32[4,4096,4096], f32[4,5,4096,8192], f32[4,5,4096,8192],
-    f32[4,1,4096,8192], f32[4,1,4096,8192]) {
-  %Arg_0.92 = f32[4,4096,4096]{2,1,0} parameter(0)
-  %Arg_1.93 = f32[4,5,4096,8192]{3,2,1,0} parameter(1)
-  %Arg_2.94 = f32[4,5,4096,8192]{3,2,1,0} parameter(2)
+%None.91 (Arg_0.92: f32[4,16,16], Arg_1.93: f32[4,5,16,32],
+    Arg_2.94: f32[4,5,16,32], Arg_3.95: f32[4,1,16,32],
+    Arg_4.96: f32[4,1,16,32], Arg_5.97: s32[])
+    -> (f32[4,16,16], f32[4,5,16,32], f32[4,5,16,32],
+    f32[4,1,16,32], f32[4,1,16,32]) {
+  %Arg_0.92 = f32[4,16,16]{2,1,0} parameter(0)
+  %Arg_1.93 = f32[4,5,16,32]{3,2,1,0} parameter(1)
+  %Arg_2.94 = f32[4,5,16,32]{3,2,1,0} parameter(2)
   %iota.113 = s32[4]{0} iota(), iota_dimension=0
-  %broadcast.114 = s32[4,1,4096,8192]{3,2,1,0} broadcast(%iota.113),
+  %broadcast.114 = s32[4,1,16,32]{3,2,1,0} broadcast(%iota.113),
       dimensions={0}
   %constant.98 = s32[] constant(0)
-  %broadcast.99 = s32[4,1,4096,8192]{3,2,1,0} broadcast(%constant.98),
+  %broadcast.99 = s32[4,1,16,32]{3,2,1,0} broadcast(%constant.98),
       dimensions={}
-  %compare.115 = pred[4,1,4096,8192]{3,2,1,0} compare(%broadcast.114,
+  %compare.115 = pred[4,1,16,32]{3,2,1,0} compare(%broadcast.114,
       %broadcast.99), direction=EQ
   %Arg_5.97 = s32[] parameter(5)
   %constant.102 = s32[] constant(5)
@@ -1830,31 +1830,31 @@ HloModule jit_entry_computation, entry_computation_layout={
   %compare.105 = pred[] compare(%call.104, %constant.103), direction=LT
   %add.106 = s32[] add(%call.104, %constant.102)
   %select.107 = s32[] select(%compare.105, %add.106, %call.104)
-  %dynamic-slice.108 = f32[4,1,4096,8192]{3,2,1,0} dynamic-slice(%Arg_1.93,
+  %dynamic-slice.108 = f32[4,1,16,32]{3,2,1,0} dynamic-slice(%Arg_1.93,
       %constant.103, %select.107, %constant.103, %constant.103),
-      dynamic_slice_sizes={4,1,4096,8192}
-  %Arg_4.96 = f32[4,1,4096,8192]{3,2,1,0} parameter(4)
-  %call.112 = f32[4,1,4096,8192]{3,2,1,0} call(%compare.111, %dynamic-slice.108,
+      dynamic_slice_sizes={4,1,16,32}
+  %Arg_4.96 = f32[4,1,16,32]{3,2,1,0} parameter(4)
+  %call.112 = f32[4,1,16,32]{3,2,1,0} call(%compare.111, %dynamic-slice.108,
       %Arg_4.96), to_apply=%_where_0.35
-  %Arg_3.95 = f32[4,1,4096,8192]{3,2,1,0} parameter(3)
-  %call.109 = f32[5,1,4096,8192]{3,2,1,0} call(%Arg_3.95, %constant.103),
+  %Arg_3.95 = f32[4,1,16,32]{3,2,1,0} parameter(3)
+  %call.109 = f32[5,1,16,32]{3,2,1,0} call(%Arg_3.95, %constant.103),
       to_apply=%_pad.30
-  %slice.110 = f32[4,1,4096,8192]{3,2,1,0} slice(%call.109),
-      slice={[0:4], [0:1], [0:4096], [0:8192]}
-  %call.116 = f32[4,1,4096,8192]{3,2,1,0} call(%compare.115, %call.112,
+  %slice.110 = f32[4,1,16,32]{3,2,1,0} slice(%call.109),
+      slice={[0:4], [0:1], [0:16], [0:32]}
+  %call.116 = f32[4,1,16,32]{3,2,1,0} call(%compare.115, %call.112,
       %slice.110), to_apply=%_where_1.41
-  %reshape.117 = f32[4,4096,8192]{2,1,0} reshape(%call.116)
-  %dot.118 = f32[4,4096,8192]{2,1,0} dot(%Arg_0.92, %reshape.117),
+  %reshape.117 = f32[4,16,32]{2,1,0} reshape(%call.116)
+  %dot.118 = f32[4,16,32]{2,1,0} dot(%Arg_0.92, %reshape.117),
       lhs_batch_dims={0}, lhs_contracting_dims={2}, rhs_batch_dims={0},
       rhs_contracting_dims={1}
-  %reshape.150 = f32[4,1,4096,8192]{3,2,1,0} reshape(%dot.118)
+  %reshape.150 = f32[4,1,16,32]{3,2,1,0} reshape(%dot.118)
   %constant.100 = s32[] constant(2)
   %add.159 = s32[] add(%Arg_5.97, %constant.100)
   %call.160 = s32[] call(%add.159, %constant.102), to_apply=%remainder.76
   %compare.161 = pred[] compare(%call.160, %constant.103), direction=LT
   %add.162 = s32[] add(%call.160, %constant.102)
   %select.163 = s32[] select(%compare.161, %add.162, %call.160)
-  %dynamic-update-slice.164 = f32[4,5,4096,8192]{3,2,1,0}
+  %dynamic-update-slice.164 = f32[4,5,16,32]{3,2,1,0}
       dynamic-update-slice(%Arg_2.94, %reshape.150, %constant.103, %select.163,
       %constant.103, /*index=5*/%constant.103)
   %constant.101 = s32[] constant(1)
@@ -1863,82 +1863,82 @@ HloModule jit_entry_computation, entry_computation_layout={
   %compare.153 = pred[] compare(%call.152, %constant.103), direction=LT
   %add.154 = s32[] add(%call.152, %constant.102)
   %select.155 = s32[] select(%compare.153, %add.154, %call.152)
-  %dynamic-slice.156 = f32[4,1,4096,8192]{3,2,1,0} dynamic-slice(%Arg_2.94,
+  %dynamic-slice.156 = f32[4,1,16,32]{3,2,1,0} dynamic-slice(%Arg_2.94,
       %constant.103, %select.155, %constant.103, %constant.103),
-      dynamic_slice_sizes={4,1,4096,8192}
-  %call.157 = f32[7,1,4096,8192]{3,2,1,0} call(%dynamic-slice.156,
+      dynamic_slice_sizes={4,1,16,32}
+  %call.157 = f32[7,1,16,32]{3,2,1,0} call(%dynamic-slice.156,
       %constant.103), to_apply=%_pad_2.66
-  %slice.158 = f32[4,1,4096,8192]{3,2,1,0} slice(%call.157),
-      slice={[3:7], [0:1], [0:4096], [0:8192]}
-  ROOT %tuple.165 = (f32[4,4096,4096]{2,1,0}, f32[4,5,4096,8192]{3,2,1,0},
-      f32[4,5,4096,8192]{3,2,1,0}, f32[4,1,4096,8192]{3,2,1,0},
-      f32[4,1,4096,8192]{3,2,1,0}) tuple(%Arg_0.92, %Arg_1.93,
+  %slice.158 = f32[4,1,16,32]{3,2,1,0} slice(%call.157),
+      slice={[3:7], [0:1], [0:16], [0:32]}
+  ROOT %tuple.165 = (f32[4,16,16]{2,1,0}, f32[4,5,16,32]{3,2,1,0},
+      f32[4,5,16,32]{3,2,1,0}, f32[4,1,16,32]{3,2,1,0},
+      f32[4,1,16,32]{3,2,1,0}) tuple(%Arg_0.92, %Arg_1.93,
       %dynamic-update-slice.164, %reshape.150, %slice.158)
 }
 
-%region_0.166 (arg_tuple.167: (s32[], f32[4,4096,4096], f32[4,5,4096,8192],
-    f32[4,5,4096,8192], f32[4,1,4096,8192], /*index=5*/f32[4,1,4096,8192],
-    s32[13])) -> (s32[], f32[4,4096,4096], f32[4,5,4096,8192],
-    f32[4,5,4096,8192], f32[4,1,4096,8192], /*index=5*/f32[4,1,4096,8192],
+%region_0.166 (arg_tuple.167: (s32[], f32[4,16,16], f32[4,5,16,32],
+    f32[4,5,16,32], f32[4,1,16,32], /*index=5*/f32[4,1,16,32],
+    s32[13])) -> (s32[], f32[4,16,16], f32[4,5,16,32],
+    f32[4,5,16,32], f32[4,1,16,32], /*index=5*/f32[4,1,16,32],
     s32[13]) {
-  %arg_tuple.167 = (s32[], f32[4,4096,4096]{2,1,0}, f32[4,5,4096,8192]{3,2,1,0},
-      f32[4,5,4096,8192]{3,2,1,0}, f32[4,1,4096,8192]{3,2,1,0},
-      /*index=5*/f32[4,1,4096,8192]{3,2,1,0}, s32[13]{0}) parameter(0)
+  %arg_tuple.167 = (s32[], f32[4,16,16]{2,1,0}, f32[4,5,16,32]{3,2,1,0},
+      f32[4,5,16,32]{3,2,1,0}, f32[4,1,16,32]{3,2,1,0},
+      /*index=5*/f32[4,1,16,32]{3,2,1,0}, s32[13]{0}) parameter(0)
   %get-tuple-element.168 = s32[] get-tuple-element(%arg_tuple.167), index=0
   %constant.175 = s32[] constant(1)
   %add.184 = s32[] add(%get-tuple-element.168, %constant.175)
-  %get-tuple-element.169 = f32[4,4096,4096]{2,1,0}
+  %get-tuple-element.169 = f32[4,16,16]{2,1,0}
       get-tuple-element(%arg_tuple.167), index=1
-  %get-tuple-element.170 = f32[4,5,4096,8192]{3,2,1,0}
+  %get-tuple-element.170 = f32[4,5,16,32]{3,2,1,0}
       get-tuple-element(%arg_tuple.167), index=2
-  %get-tuple-element.171 = f32[4,5,4096,8192]{3,2,1,0}
+  %get-tuple-element.171 = f32[4,5,16,32]{3,2,1,0}
       get-tuple-element(%arg_tuple.167), index=3
-  %get-tuple-element.172 = f32[4,1,4096,8192]{3,2,1,0}
+  %get-tuple-element.172 = f32[4,1,16,32]{3,2,1,0}
       get-tuple-element(%arg_tuple.167), index=4
-  %get-tuple-element.173 = f32[4,1,4096,8192]{3,2,1,0}
+  %get-tuple-element.173 = f32[4,1,16,32]{3,2,1,0}
       get-tuple-element(%arg_tuple.167), index=5
   %get-tuple-element.174 = s32[13]{0} get-tuple-element(%arg_tuple.167), index=6
   %dynamic-slice.176 = s32[1]{0} dynamic-slice(%get-tuple-element.174,
       %get-tuple-element.168), dynamic_slice_sizes={1}
   %reshape.177 = s32[] reshape(%dynamic-slice.176)
-  %call.178 = (f32[4,4096,4096]{2,1,0}, f32[4,5,4096,8192]{3,2,1,0},
-      f32[4,5,4096,8192]{3,2,1,0}, f32[4,1,4096,8192]{3,2,1,0},
-      f32[4,1,4096,8192]{3,2,1,0}) call(%get-tuple-element.169,
+  %call.178 = (f32[4,16,16]{2,1,0}, f32[4,5,16,32]{3,2,1,0},
+      f32[4,5,16,32]{3,2,1,0}, f32[4,1,16,32]{3,2,1,0},
+      f32[4,1,16,32]{3,2,1,0}) call(%get-tuple-element.169,
       %get-tuple-element.170, %get-tuple-element.171, %get-tuple-element.172,
       %get-tuple-element.173, /*index=5*/%reshape.177), to_apply=%None.91
-  %get-tuple-element.179 = f32[4,4096,4096]{2,1,0} get-tuple-element(%call.178),
+  %get-tuple-element.179 = f32[4,16,16]{2,1,0} get-tuple-element(%call.178),
       index=0
-  %get-tuple-element.180 = f32[4,5,4096,8192]{3,2,1,0}
+  %get-tuple-element.180 = f32[4,5,16,32]{3,2,1,0}
       get-tuple-element(%call.178), index=1
-  %get-tuple-element.181 = f32[4,5,4096,8192]{3,2,1,0}
+  %get-tuple-element.181 = f32[4,5,16,32]{3,2,1,0}
       get-tuple-element(%call.178), index=2
-  %get-tuple-element.182 = f32[4,1,4096,8192]{3,2,1,0}
+  %get-tuple-element.182 = f32[4,1,16,32]{3,2,1,0}
       get-tuple-element(%call.178), index=3
-  %get-tuple-element.183 = f32[4,1,4096,8192]{3,2,1,0}
+  %get-tuple-element.183 = f32[4,1,16,32]{3,2,1,0}
       get-tuple-element(%call.178), index=4
-  ROOT %tuple.185 = (s32[], f32[4,4096,4096]{2,1,0},
-      f32[4,5,4096,8192]{3,2,1,0}, f32[4,5,4096,8192]{3,2,1,0},
-      f32[4,1,4096,8192]{3,2,1,0}, /*index=5*/f32[4,1,4096,8192]{3,2,1,0},
+  ROOT %tuple.185 = (s32[], f32[4,16,16]{2,1,0},
+      f32[4,5,16,32]{3,2,1,0}, f32[4,5,16,32]{3,2,1,0},
+      f32[4,1,16,32]{3,2,1,0}, /*index=5*/f32[4,1,16,32]{3,2,1,0},
       s32[13]{0}) tuple(%add.184, %get-tuple-element.179,
       %get-tuple-element.180, %get-tuple-element.181, %get-tuple-element.182,
       /*index=5*/%get-tuple-element.183, %get-tuple-element.174)
 }
 
-%region_1.186 (arg_tuple.187: (s32[], f32[4,4096,4096], f32[4,5,4096,8192],
-    f32[4,5,4096,8192], f32[4,1,4096,8192], /*index=5*/f32[4,1,4096,8192],
+%region_1.186 (arg_tuple.187: (s32[], f32[4,16,16], f32[4,5,16,32],
+    f32[4,5,16,32], f32[4,1,16,32], /*index=5*/f32[4,1,16,32],
     s32[13])) -> pred[] {
-  %arg_tuple.187 = (s32[], f32[4,4096,4096]{2,1,0}, f32[4,5,4096,8192]{3,2,1,0},
-      f32[4,5,4096,8192]{3,2,1,0}, f32[4,1,4096,8192]{3,2,1,0},
-      /*index=5*/f32[4,1,4096,8192]{3,2,1,0}, s32[13]{0}) parameter(0)
-  %get-tuple-element.189 = f32[4,4096,4096]{2,1,0}
+  %arg_tuple.187 = (s32[], f32[4,16,16]{2,1,0}, f32[4,5,16,32]{3,2,1,0},
+      f32[4,5,16,32]{3,2,1,0}, f32[4,1,16,32]{3,2,1,0},
+      /*index=5*/f32[4,1,16,32]{3,2,1,0}, s32[13]{0}) parameter(0)
+  %get-tuple-element.189 = f32[4,16,16]{2,1,0}
       get-tuple-element(%arg_tuple.187), index=1
-  %get-tuple-element.190 = f32[4,5,4096,8192]{3,2,1,0}
+  %get-tuple-element.190 = f32[4,5,16,32]{3,2,1,0}
       get-tuple-element(%arg_tuple.187), index=2
-  %get-tuple-element.191 = f32[4,5,4096,8192]{3,2,1,0}
+  %get-tuple-element.191 = f32[4,5,16,32]{3,2,1,0}
       get-tuple-element(%arg_tuple.187), index=3
-  %get-tuple-element.192 = f32[4,1,4096,8192]{3,2,1,0}
+  %get-tuple-element.192 = f32[4,1,16,32]{3,2,1,0}
       get-tuple-element(%arg_tuple.187), index=4
-  %get-tuple-element.193 = f32[4,1,4096,8192]{3,2,1,0}
+  %get-tuple-element.193 = f32[4,1,16,32]{3,2,1,0}
       get-tuple-element(%arg_tuple.187), index=5
   %get-tuple-element.194 = s32[13]{0} get-tuple-element(%arg_tuple.187), index=6
   %get-tuple-element.188 = s32[] get-tuple-element(%arg_tuple.187), index=0
@@ -1947,39 +1947,39 @@ HloModule jit_entry_computation, entry_computation_layout={
       direction=LT
 }
 
-ENTRY %main.204 (Arg_0.1: f32[4,4096,4096], Arg_1.2: f32[4,5,4096,8192])
-    -> f32[4,5,4096,8192] {
+ENTRY %main.204 (Arg_0.1: f32[4,16,16], Arg_1.2: f32[4,5,16,32])
+    -> f32[4,5,16,32] {
   %constant.3 = s32[] constant(0)
-  %Arg_0.1 = f32[4,4096,4096]{2,1,0} parameter(0),
+  %Arg_0.1 = f32[4,16,16]{2,1,0} parameter(0),
       sharding={devices=[4,1,1]<=[4]}
-  %Arg_1.2 = f32[4,5,4096,8192]{3,2,1,0} parameter(1),
+  %Arg_1.2 = f32[4,5,16,32]{3,2,1,0} parameter(1),
       sharding={devices=[4,1,1,1]<=[4]}
   %constant.4 = f32[] constant(0)
-  %broadcast.5 = f32[4,5,4096,8192]{3,2,1,0} broadcast(%constant.4),
+  %broadcast.5 = f32[4,5,16,32]{3,2,1,0} broadcast(%constant.4),
       dimensions={}
   %constant.6 = f32[] constant(0)
-  %broadcast.7 = f32[4,1,4096,8192]{3,2,1,0} broadcast(%constant.6),
+  %broadcast.7 = f32[4,1,16,32]{3,2,1,0} broadcast(%constant.6),
       dimensions={}
   %iota.8 = s32[13]{0} iota(), iota_dimension=0
-  %tuple.9 = (s32[], f32[4,4096,4096]{2,1,0}, f32[4,5,4096,8192]{3,2,1,0},
-      f32[4,5,4096,8192]{3,2,1,0}, f32[4,1,4096,8192]{3,2,1,0},
-      /*index=5*/f32[4,1,4096,8192]{3,2,1,0}, s32[13]{0}) tuple(%constant.3,
+  %tuple.9 = (s32[], f32[4,16,16]{2,1,0}, f32[4,5,16,32]{3,2,1,0},
+      f32[4,5,16,32]{3,2,1,0}, f32[4,1,16,32]{3,2,1,0},
+      /*index=5*/f32[4,1,16,32]{3,2,1,0}, s32[13]{0}) tuple(%constant.3,
       %Arg_0.1, %Arg_1.2, %broadcast.5, %broadcast.7, /*index=5*/%broadcast.7,
       %iota.8)
-  %while.197 = (s32[], f32[4,4096,4096]{2,1,0}, f32[4,5,4096,8192]{3,2,1,0},
-      f32[4,5,4096,8192]{3,2,1,0}, f32[4,1,4096,8192]{3,2,1,0},
-      /*index=5*/f32[4,1,4096,8192]{3,2,1,0}, s32[13]{0}) while(%tuple.9),
+  %while.197 = (s32[], f32[4,16,16]{2,1,0}, f32[4,5,16,32]{3,2,1,0},
+      f32[4,5,16,32]{3,2,1,0}, f32[4,1,16,32]{3,2,1,0},
+      /*index=5*/f32[4,1,16,32]{3,2,1,0}, s32[13]{0}) while(%tuple.9),
       condition=%region_1.186, body=%region_0.166
   %get-tuple-element.198 = s32[] get-tuple-element(%while.197), index=0
-  %get-tuple-element.199 = f32[4,4096,4096]{2,1,0}
+  %get-tuple-element.199 = f32[4,16,16]{2,1,0}
       get-tuple-element(%while.197), index=1
-  %get-tuple-element.200 = f32[4,5,4096,8192]{3,2,1,0}
+  %get-tuple-element.200 = f32[4,5,16,32]{3,2,1,0}
       get-tuple-element(%while.197), index=2
-  ROOT %get-tuple-element.201 = f32[4,5,4096,8192]{3,2,1,0}
+  ROOT %get-tuple-element.201 = f32[4,5,16,32]{3,2,1,0}
       get-tuple-element(%while.197), index=3
-  %get-tuple-element.202 = f32[4,1,4096,8192]{3,2,1,0}
+  %get-tuple-element.202 = f32[4,1,16,32]{3,2,1,0}
       get-tuple-element(%while.197), index=4
-  %get-tuple-element.203 = f32[4,1,4096,8192]{3,2,1,0}
+  %get-tuple-element.203 = f32[4,1,16,32]{3,2,1,0}
       get-tuple-element(%while.197), index=5
 }
   )";

From 00c4189d90f71665ddaff46f8d53403be33d8918 Mon Sep 17 00:00:00 2001
From: Stanislav Bardyuk <sbardyuk@google.com>
Date: Tue, 13 Oct 2026 01:52:15 -0700
Subject: [PATCH 0662/1000] PR #50167: [XLA] Return NaN from Igammac for
 out-of-domain and NaN inputs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50167

## 📝 Summary of Changes

`xla::Igammac(a, x)` returned 1 for every input with `x <= 0 || a <= 0`. That is right for `x == 0`, but `x < 0` and `a <= 0` are outside the domain of the function and the references TF compares against (Eigen, which is what eager uses, and scipy) give NaN there. This change splits the mask the same way `Igamma` in the same file already does: `x == 0` keeps returning 1, `x < 0`, `a <= 0` and NaN inputs return NaN, `x == inf` keeps returning 0. Valid inputs take exactly the same code path as before. The only other value that changes is `Igammac(NaN, inf)`, which was 0 and is now NaN.

## 🎯 Justification

Fixes #49175. `tf.math.igammac` under `jit_compile=True` gave `[0.336429, 1.0, 1.0]` for `[0.6374078, -1.6019852, -0.24963307]` while eager execution (Eigen) gave `[0.33642888, nan, nan]`. Igammac was the only one of the incomplete gamma helpers in math.cc without the NaN handling.

## 🚀 Kind of Contribution

🐛 Bug Fix

## 🧪 Unit Tests:

New `MathTest.IgammacSpecialValues` in xla/hlo/builder/lib/math_test.cc covering NaN inputs, the values from the issue, negative x, a == 0, a < 0, x == 0, x == -0.0 and x == inf. It fails before the change (7 of 14 lanes are wrong) and passes after it.

```
bazel test //xla/hlo/builder/lib:math_test_cpu
```

## 🧪 Execution Tests:

Same test, it runs the computation on the CPU backend. Also checked that a 40x40 grid of valid (a, x) pairs in f32 and f64 gives bit-identical results before and after the change.
Copybara import of the project:

--
d2e9082b33631d3df127bf5b1251e4182a6c5a75 by Stanislav Bardyuk <sbardyuk@google.com>:

[XLA] Return NaN from Igammac for out-of-domain and NaN inputs

Igammac mapped the whole region x <= 0 || a <= 0 to 1. That is only right
for x == 0; x < 0 and a <= 0 are outside the domain and Eigen and scipy
return NaN there. Split the mask the same way Igamma does: x == 0 stays 1,
x == inf stays 0, domain errors and NaN inputs give NaN. Valid inputs are
unchanged.

Fixes #49175.

Merging this change closes #50167

PiperOrigin-RevId: 998471153
---
 third_party/xla/xla/hlo/builder/lib/math.cc   | 16 +++++++++-----
 .../xla/xla/hlo/builder/lib/math_test.cc      | 21 +++++++++++++++++++
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/third_party/xla/xla/hlo/builder/lib/math.cc b/third_party/xla/xla/hlo/builder/lib/math.cc
index be9ba4fd357..fe7b94d54f7 100644
--- a/third_party/xla/xla/hlo/builder/lib/math.cc
+++ b/third_party/xla/xla/hlo/builder/lib/math.cc
@@ -1109,22 +1109,28 @@ XlaOp RandomGammaGrad(XlaOp a, XlaOp x) {
 XlaOp Igammac(XlaOp a, XlaOp x) {
   auto& b = *a.builder();
   auto doit = [&b](XlaOp a, XlaOp x, PrimitiveType type) -> XlaOp {
-    XlaOp out_of_range = Or(Le(x, ScalarLike(x, 0)), Le(a, ScalarLike(a, 0)));
+    XlaOp is_nan = Or(IsNan(a), IsNan(x));
+    XlaOp x_is_zero = Eq(x, ScalarLike(x, 0));
+    XlaOp x_is_infinity =
+        Eq(x, ScalarLike(x, std::numeric_limits<float>::infinity()));
+    // Same domain as Igamma: x < 0 or a <= 0 is a domain error (NaN), while
+    // x == 0 is a valid boundary where igammac(a, 0) = 1.
+    XlaOp domain_error = Or(Lt(x, ScalarLike(x, 0)), Le(a, ScalarLike(a, 0)));
     XlaOp use_igamma = Or(Lt(x, ScalarLike(x, 1)), Lt(x, a));
     XlaOp ax = a * Log(x) - x - Lgamma(a);
     XlaOp underflow = Lt(ax, -Log(MaxFiniteValue(&b, type)));
-    XlaOp enabled = Not(Or(out_of_range, underflow));
+    XlaOp enabled = Not(Or(Or(Or(x_is_zero, domain_error), underflow), is_nan));
     ax = Exp(ax);
+    const double nan = std::numeric_limits<double>::quiet_NaN();
     XlaOp result =
         Select(use_igamma,
                ScalarLike(a, 1) - IgammaSeries<VALUE>(
                                       ax, x, a, And(enabled, use_igamma), type),
                IgammacContinuedFraction<VALUE>(
                    ax, x, a, And(enabled, Not(use_igamma)), type));
-    XlaOp x_is_infinity =
-        Eq(x, ScalarLike(x, std::numeric_limits<float>::infinity()));
+    result = Select(x_is_zero, FullLike(result, 1), result);
     result = Select(x_is_infinity, ZerosLike(result), result);
-    return Select(out_of_range, FullLike(a, 1), result);
+    return Select(Or(domain_error, is_nan), FullLike(a, nan), result);
   };
   return b.ReportErrorOrReturn([&]() -> absl::StatusOr<XlaOp> {
     ABSL_ASSIGN_OR_RETURN(auto a_shape, b.GetShape(a));
diff --git a/third_party/xla/xla/hlo/builder/lib/math_test.cc b/third_party/xla/xla/hlo/builder/lib/math_test.cc
index dce946ee200..0b126329390 100644
--- a/third_party/xla/xla/hlo/builder/lib/math_test.cc
+++ b/third_party/xla/xla/hlo/builder/lib/math_test.cc
@@ -631,6 +631,27 @@ TEST_F(MathTest, Igammac) {
   ComputeAndCompareR3<float>(&builder, expected, {}, kErrorSpec);
 }
 
+TEST_F(MathTest, IgammacSpecialValues) {
+  SetFastMathDisabled(true);
+  XlaBuilder builder(TestName());
+  const float nan = std::numeric_limits<float>::quiet_NaN();
+  const float inf = std::numeric_limits<float>::infinity();
+  auto a = ConstantR1<float>(
+      &builder,
+      {nan, nan, 0.53327996, nan, 0.6374078, -1.6019852, -0.24963307,
+       0.53327996, 0.0, -2.0, -2.0, 0.53327996, 0.53327996, 0.53327996});
+  auto x = ConstantR1<float>(
+      &builder, {nan, 8.97671773, nan, 0.0, 0.6374078, -1.6019852, -0.24963307,
+                 -1.0, 2.0, 0.0, inf, 0.0, -0.0, inf});
+
+  Igammac(a, x);
+  // Domain errors (x < 0, a <= 0) and NaN inputs give NaN as in Eigen and
+  // Igamma; igammac(a, 0) = 1 and igammac(a, inf) = 0 for a > 0.
+  std::vector<float> expected = {nan, nan, nan, nan, 0.33642888, nan, nan,
+                                 nan, nan, nan, nan, 1.0,        1.0, 0.0};
+  ComputeAndCompareR1<float>(&builder, expected, {}, kErrorSpec);
+}
+
 TEST_F(MathTest, IgammacF16) {
   SetFastMathDisabled(true);
 

From bf06958241dcedd0da8481de1de0199725d2030c Mon Sep 17 00:00:00 2001
From: nishad shabbir <nishadshabbir28@gmail.com>
Date: Tue, 13 Oct 2026 18:21:39 +0530
Subject: [PATCH 0663/1000] Drive the name scope test over every container
 pairing

Wrapping each argument independently covers all nine combinations of
tensor, list and tuple instead of the six I picked by hand. The one I had
left out, a tuple ys with a tensor xs, is among the four that raise
TypeError without the list() casts.
---
 tensorflow/python/ops/gradients_test.py | 47 ++++++++++++-------------
 1 file changed, 22 insertions(+), 25 deletions(-)

diff --git a/tensorflow/python/ops/gradients_test.py b/tensorflow/python/ops/gradients_test.py
index 50c27d245fe..3c7a87cb37f 100644
--- a/tensorflow/python/ops/gradients_test.py
+++ b/tensorflow/python/ops/gradients_test.py
@@ -866,33 +866,30 @@ class HessianTest(test_util.TensorFlowTestCase):
     # `name` is documented as grouping all the gradient ops together, so the
     # ops hessians() adds must land under it rather than in the enclosing
     # scope.
-    # _AsList() leaves tuples alone, so the ys and xs containers also have to
-    # survive being mixed.
-    containers = [
-        lambda y, x: (y, x),
-        lambda y, x: (y, (x,)),
-        lambda y, x: ((y,), [x]),
-        lambda y, x: ([y], (x,)),
-        lambda y, x: ([y], [x]),
-        lambda y, x: ((y,), (x,)),
+    # _AsList() leaves tuples alone, so every pairing of the containers the
+    # two arguments accept has to survive being concatenated.
+    wrappers = [
+        lambda v: v,
+        lambda v: (v,),
+        lambda v: [v],
     ]
     for name, expected in [(None, "hessians"), ("my_hessians", "my_hessians")]:
-      for wrap in containers:
-        with ops.Graph().as_default() as g:
-          x = constant_op.constant([1.0, 2.0, 3.0])
-          y = math_ops.reduce_sum(x * x * x)
-          existing_ops = set(g.get_operations())
-          ys, xs = wrap(y, x)
-          if name is None:
-            gradients.hessians(ys, xs)
-          else:
-            gradients.hessians(ys, xs, name=name)
-          added = [
-              op.name for op in g.get_operations() if op not in existing_ops
-          ]
-        self.assertNotEmpty(added)
-        for op_name in added:
-          self.assertEqual(expected, op_name.split("/")[0])
+      for wrap_ys in wrappers:
+        for wrap_xs in wrappers:
+          with ops.Graph().as_default() as g:
+            x = constant_op.constant([1.0, 2.0, 3.0])
+            y = math_ops.reduce_sum(x * x * x)
+            existing_ops = set(g.get_operations())
+            if name is None:
+              gradients.hessians(wrap_ys(y), wrap_xs(x))
+            else:
+              gradients.hessians(wrap_ys(y), wrap_xs(x), name=name)
+            added = [
+                op.name for op in g.get_operations() if op not in existing_ops
+            ]
+          self.assertNotEmpty(added)
+          for op_name in added:
+            self.assertEqual(expected, op_name.split("/")[0])
 
 
 class IndexedSlicesToTensorTest(test_util.TensorFlowTestCase):

From fbce2ae83af5e54b7f149a89aa003247988ff4bd Mon Sep 17 00:00:00 2001
From: Alexandros Theodoridis <atheodor@amd.com>
Date: Tue, 13 Oct 2026 02:00:57 -0700
Subject: [PATCH 0664/1000] PR #50568: [ROCm] Add rocm rbe dev setup
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50568

📝 Summary of Changes
Add dev setup for rocm rbe

🎯 Justification
By default rocm rbe cluster uses internal network
setup which is internal to the cloude provider
dev machines can't access it directly.
Removing internal prefix allowes dev machines
to use rbe cluster.

🚀 Kind of Contribution
Please remove what does not apply: ♻️ Cleanup

📊 Benchmark (for Performance Improvements)
Not relevant

🧪 Unit Tests:
CI

🧪 Execution Tests:
CI

Copybara import of the project:

--
4f5fe6baf16af4a4c5b224c5d2c5d9aa5f726815 by Alexandros Theodoridis <atheodor@amd.com>:

Add rocm rbe dev setup

Merging this change closes #50568

PiperOrigin-RevId: 998475351
---
 .../xla/build_tools/rocm/rocm_xla.bazelrc      | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/build_tools/rocm/rocm_xla.bazelrc b/third_party/xla/build_tools/rocm/rocm_xla.bazelrc
index e0c18369bc4..f69f7bbd880 100644
--- a/third_party/xla/build_tools/rocm/rocm_xla.bazelrc
+++ b/third_party/xla/build_tools/rocm/rocm_xla.bazelrc
@@ -1,12 +1,10 @@
 # Test-related settings.
 
-build:rocm_dev --remote_upload_local_results=false
-build:rocm_dev --remote_cache="grpcs://internal.actinolite.cluster.engflow.com"
-
 common:rocm_base --test_verbose_timeout_warnings
 common:rocm_base --test_timeout=920,2400,7200,9600
 
 build:rocm_rbe --remote_cache="grpcs://internal.actinolite.cluster.engflow.com"
+build:rocm_rbe --remote_executor="grpcs://internal.actinolite.cluster.engflow.com"
 build:rocm_rbe --repo_env=REMOTE_GPU_TESTING=1
 build:rocm_rbe --bes_backend="grpcs://internal.actinolite.cluster.engflow.com"
 build:rocm_rbe --bes_results_url="https://internal.actinolite.cluster.engflow.com/invocation/"
@@ -25,11 +23,15 @@ build:rocm_rbe --incompatible_strict_action_env
 build:rocm_rbe --experimental_strict_repo_env
 
 
-test:rocm_rbe --jobs=50
-test:rocm_rbe --remote_executor=grpcs://internal.actinolite.cluster.engflow.com
-test:rocm_rbe --remote_timeout=3600
-test:rocm_rbe --strategy=TestRunner=remote,local
-test:rocm_rbe --test_sharding_strategy=explicit
+build:rocm_rbe --jobs=50
+build:rocm_rbe --remote_timeout=3600
+build:rocm_rbe --strategy=TestRunner=remote,local
+build:rocm_rbe --test_sharding_strategy=explicit
+
+build:rocm_rbe_dev --config=rocm_rbe
+build:rocm_rbe_dev --remote_cache="grpcs://actinolite.cluster.engflow.com"
+build:rocm_rbe_dev --remote_executor="grpcs://actinolite.cluster.engflow.com"
+build:rocm_rbe_dev --repo_env=REMOTE_GPU_TESTING=1
 
 # Dynamic execution config for ROCm RBE - races remote vs local with remote head start
 build:rocm_rbe_dynamic --config=rocm_rbe

From c4e00871abb14b0cad622f8f6fca45e5dc55453e Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 13 Oct 2026 02:06:10 -0700
Subject: [PATCH 0665/1000] Reverts a814257ee0c331ed289e2a867bfd77524d5ef1f1

PiperOrigin-RevId: 998478175
---
 .../simplifiers/algebraic_simplifier.cc       | 189 +-----------------
 .../simplifiers/algebraic_simplifier_test.cc  | 132 ------------
 2 files changed, 6 insertions(+), 315 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
index 5a52d6f43c2..69e32bffd23 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
@@ -934,179 +934,6 @@ absl::Status AlgebraicSimplifierVisitor::HandleAbs(HloInstruction* abs) {
   return absl::OkStatus();
 }
 
-namespace {
-
-// Returns the value of a constant integral scalar, looking through converts
-// (e.g. `convert(u32 constant)` left behind by X64 elimination) and applying
-// the conversions so type semantics are respected.
-std::optional<int64_t> GetConstantIntegralScalar(const HloInstruction* hlo) {
-  if (!ShapeUtil::IsScalar(hlo->shape()) ||
-      !primitive_util::IsIntegralType(hlo->shape().element_type())) {
-    return std::nullopt;
-  }
-  const HloInstruction* current = hlo;
-  absl::InlinedVector<PrimitiveType, 2> convert_types;
-  while (current->opcode() == HloOpcode::kConvert) {
-    convert_types.push_back(current->shape().element_type());
-    current = current->operand(0);
-  }
-  if (!Match(current, m::ConstantScalar())) {
-    return std::nullopt;
-  }
-  Literal lit = current->literal().Clone();
-  for (auto it = convert_types.rbegin(); it != convert_types.rend(); ++it) {
-    absl::StatusOr<Literal> converted = lit.Convert(*it);
-    if (!converted.ok()) {
-      return std::nullopt;
-    }
-    lit = *std::move(converted);
-  }
-  return lit.GetFirstInteger();
-}
-
-// Returns `pad(operand, zero, config)`, or `operand` if `config` is a no-op.
-absl::StatusOr<HloInstruction*> MakePadIfNeeded(
-    HloInstruction* operand, HloInstruction* zero, const PaddingConfig& config,
-    const OpMetadata* metadata = nullptr) {
-  if (absl::c_all_of(config.dimensions(), [](const auto& dim) {
-        return dim.edge_padding_low() == 0 && dim.edge_padding_high() == 0 &&
-               dim.interior_padding() == 0;
-      })) {
-    return operand;
-  }
-  return MakePadHlo(operand, zero, config, metadata);
-}
-
-// An array that is zero everywhere except for an axis-aligned box, which holds
-// `source`. `low(d)` / `high(d)` are the number of zeros before / after the box
-// in dimension `d`.
-class ZeroPaddedBox {
- public:
-  ZeroPaddedBox(HloInstruction* source, std::vector<int64_t> low,
-                std::vector<int64_t> high)
-      : source_(source), low_(std::move(low)), high_(std::move(high)) {
-    CHECK_EQ(low_.size(), source_->shape().dimensions().size());
-    CHECK_EQ(high_.size(), source_->shape().dimensions().size());
-  }
-
-  HloInstruction* source() const { return source_; }
-  int64_t low(int64_t dim) const { return low_[dim]; }
-  int64_t high(int64_t dim) const { return high_[dim]; }
-  int64_t rank() const { return low_.size(); }
-
- private:
-  HloInstruction* source_;
-  std::vector<int64_t> low_;
-  std::vector<int64_t> high_;
-};
-
-// Matches `pad(source, 0)` without interior or negative padding.
-std::optional<ZeroPaddedBox> MatchZeroPaddedBox(HloInstruction* hlo) {
-  if (!hlo->shape().IsArray() || hlo->opcode() != HloOpcode::kPad ||
-      !IsAll(hlo->operand(1), 0)) {
-    return std::nullopt;
-  }
-  std::vector<int64_t> low;
-  std::vector<int64_t> high;
-  for (const auto& dim : hlo->padding_config().dimensions()) {
-    if (dim.interior_padding() != 0 || dim.edge_padding_low() < 0 ||
-        dim.edge_padding_high() < 0) {
-      return std::nullopt;
-    }
-    low.push_back(dim.edge_padding_low());
-    high.push_back(dim.edge_padding_high());
-  }
-  return ZeroPaddedBox(hlo->mutable_operand(0), std::move(low),
-                       std::move(high));
-}
-
-// Returns the only dimension in which `a` and `b` are placed differently, or
-// std::nullopt if they are placed identically or differ in several dimensions.
-std::optional<int64_t> FindConcatDim(const ZeroPaddedBox& a,
-                                     const ZeroPaddedBox& b) {
-  std::optional<int64_t> concat_dim;
-  for (int64_t d = 0; d < a.rank(); ++d) {
-    if (a.low(d) == b.low(d) && a.high(d) == b.high(d)) {
-      continue;
-    }
-    // The boxes may be placed differently in only one dimension, the one we
-    // concatenate along.
-    if (concat_dim.has_value()) {
-      return std::nullopt;
-    }
-    concat_dim = d;
-  }
-  return concat_dim;
-}
-
-// Returns the operands of `concatenate(first.source, zeros, second.source)`
-// along `concat_dim`, where the zeros fill a gap of size `gap` (omitted if 0).
-std::vector<HloInstruction*> MakeConcatParts(const ZeroPaddedBox& first,
-                                             const ZeroPaddedBox& second,
-                                             int64_t concat_dim, int64_t gap,
-                                             HloInstruction* zero) {
-  std::vector<HloInstruction*> parts = {first.source()};
-  if (gap > 0) {
-    std::vector<int64_t> gap_dims(first.source()->shape().dimensions().begin(),
-                                  first.source()->shape().dimensions().end());
-    gap_dims[concat_dim] = gap;
-    parts.push_back(MakeBroadcastHlo(zero, {}, gap_dims));
-  }
-  parts.push_back(second.source());
-  return parts;
-}
-
-// Folds `add(A, B)` where A and B are zero-padded boxes that coincide in all
-// but one dimension and are disjoint in that one into
-// `pad(concatenate(A.source, zeros, B.source), 0)`.
-//
-// Autograd emits such sums for the gradient of a tensor that is split into
-// pieces, e.g. `q, k, v = x[0], x[1], x[2]`. The result is again a zero-padded
-// box, so n-ary sums collapse into nested concatenates, which HandleConcatenate
-// flattens.
-absl::StatusOr<HloInstruction*> TryFoldAddOfDisjointZeroPaddedBoxes(
-    HloInstruction* add) {
-  std::optional<ZeroPaddedBox> a = MatchZeroPaddedBox(add->mutable_operand(0));
-  if (!a.has_value()) {
-    return nullptr;
-  }
-  std::optional<ZeroPaddedBox> b = MatchZeroPaddedBox(add->mutable_operand(1));
-  if (!b.has_value()) {
-    return nullptr;
-  }
-  std::optional<int64_t> maybe_concat_dim = FindConcatDim(*a, *b);
-  if (!maybe_concat_dim.has_value()) {
-    return nullptr;
-  }
-  const int64_t concat_dim = *maybe_concat_dim;
-  if (a->low(concat_dim) > b->low(concat_dim)) {
-    std::swap(a, b);
-  }
-  // Both boxes lie inside the same array, so with `a` starting first they are
-  // disjoint iff `b` starts at or after the end of `a`. `b`'s end, and hence
-  // its high padding, is then determined by its size; it needs no check.
-  const int64_t a_end =
-      add->shape().dimensions(concat_dim) - a->high(concat_dim);
-  const int64_t gap = b->low(concat_dim) - a_end;
-  if (gap < 0) {
-    return nullptr;
-  }
-  HloInstruction* zero = add->AddInstruction(HloInstruction::CreateConstant(
-      LiteralUtil::Zero(add->shape().element_type())));
-  ABSL_ASSIGN_OR_RETURN(HloInstruction * concat,
-                   MakeConcatHlo(MakeConcatParts(*a, *b, concat_dim, gap, zero),
-                                 concat_dim));
-  PaddingConfig padding_config = MakeNoPaddingConfig(a->rank());
-  for (int64_t d = 0; d < a->rank(); ++d) {
-    auto* pad_dim = padding_config.mutable_dimensions(d);
-    pad_dim->set_edge_padding_low(a->low(d));
-    pad_dim->set_edge_padding_high(d == concat_dim ? b->high(d) : a->high(d));
-  }
-  return MakePadIfNeeded(concat, zero, padding_config, &add->metadata());
-}
-
-}  // namespace
-
 absl::Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) {
   HloInstruction *lhs, *rhs;
   CHECK(Match(add, m::Add(m::Op(&lhs), m::Op(&rhs))));
@@ -1188,15 +1015,6 @@ absl::Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) {
     }
   }
 
-  // Sum of zero-padded blocks that tile a dimension => pad(concatenate(...)).
-  if (!options_.is_layout_sensitive()) {
-    ABSL_ASSIGN_OR_RETURN(HloInstruction * folded,
-                     TryFoldAddOfDisjointZeroPaddedBoxes(add));
-    if (folded != nullptr) {
-      return ReplaceInstruction(add, folded);
-    }
-  }
-
   // Convert add with fullshape into add with partial shape when a
   // portion of add is effective:
   //             zero (fullshape)   rhs (partialshape)
@@ -8514,8 +8332,13 @@ absl::Status AlgebraicSimplifierVisitor::HandleDynamicUpdateSlice(
       for (int64_t dim = 0; dim < updated_shape.dimensions().size(); ++dim) {
         auto padding_config_dim = padding_config.add_dimensions();
         auto slice_dim_start = update_start_indx->operand(dim + offset);
+        if (!Match(slice_dim_start, m::ConstantScalar())) {
+          compatible = false;
+          break;
+        }
+        VLOG(2) << "slice: " << slice_dim_start->ToString();
         std::optional<int64_t> start =
-            GetConstantIntegralScalar(slice_dim_start);
+            slice_dim_start->literal().GetFirstInteger();
         if (!start) {
           compatible = false;
           break;
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
index c9f2b454c59..2725280e898 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
@@ -9678,56 +9678,6 @@ ENTRY AddBroadcastZeroWithDynamicSlice {
   EXPECT_THAT(root->operand(1)->opcode(), HloOpcode::kPad);
 }
 
-TEST_F(AlgebraicSimplifierTest, DynamicUpdateSliceOfBroadcastToPadWithConvert) {
-  constexpr absl::string_view hlo_string = R"(
-HloModule DusWithConvert
-
-ENTRY DusWithConvert {
-  constant = f32[] constant(0)
-  broadcast = f32[10,12] broadcast(constant), dimensions={}
-  param = f32[2,12] parameter(0)
-  u_idx = u32[] constant(3)
-  s_idx = s32[] convert(u_idx)
-  zero = s32[] constant(0)
-  ROOT dus = f32[10,12] dynamic-update-slice(broadcast, param, s_idx, zero)
-}
-)";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
-  AlgebraicSimplifier simplifier(default_options_);
-  ASSERT_TRUE(simplifier.Run(module.get()).value());
-  auto root = module->entry_computation()->root_instruction();
-  EXPECT_THAT(root->opcode(), HloOpcode::kPad);
-  EXPECT_EQ(root->padding_config().dimensions(0).edge_padding_low(), 3);
-  EXPECT_EQ(root->padding_config().dimensions(0).edge_padding_high(), 5);
-}
-
-TEST_F(AlgebraicSimplifierTest,
-       DynamicUpdateSliceOfBroadcastToPadWithTruncatingConvert) {
-  constexpr absl::string_view hlo_string = R"(
-HloModule DusWithTruncatingConvert
-
-ENTRY DusWithTruncatingConvert {
-  constant = f32[] constant(0)
-  broadcast = f32[10,12] broadcast(constant), dimensions={}
-  param = f32[2,12] parameter(0)
-  big_idx = s32[] constant(257)
-  trunc_idx = s8[] convert(big_idx)
-  s_idx = s32[] convert(trunc_idx)
-  zero = s32[] constant(0)
-  ROOT dus = f32[10,12] dynamic-update-slice(broadcast, param, s_idx, zero)
-}
-)";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
-  AlgebraicSimplifier simplifier(default_options_);
-  ASSERT_TRUE(simplifier.Run(module.get()).value());
-  auto root = module->entry_computation()->root_instruction();
-  EXPECT_THAT(root->opcode(), HloOpcode::kPad);
-  // 257 as s8 is 1, so low padding must be 1 (if convert was ignored, it would
-  // be 257).
-  EXPECT_EQ(root->padding_config().dimensions(0).edge_padding_low(), 1);
-  EXPECT_EQ(root->padding_config().dimensions(0).edge_padding_high(), 7);
-}
-
 // Test that dynamic-update-slice with a scalar broadcast does not become a pad
 // if the dynamic-update-slice is for host memory offload.
 TEST_F(AlgebraicSimplifierTest, DynamicUpdateSliceOfBroadcastToPadHostOffload) {
@@ -10015,88 +9965,6 @@ ENTRY AddDynamicUpdateSliceToAddSlice {
                         m::Parameter(2), m::Constant(), m::Constant())));
 }
 
-TEST_F(AlgebraicSimplifierTest, AddOfTilingDusIntoZerosBecomesConcat) {
-  // Gradient of `q, k, v = x[0], x[1], x[2]`, with indices as left behind by
-  // X64 elimination.
-  constexpr absl::string_view kHlo = R"(
-HloModule m
-
-ENTRY main {
-  pq = f32[1,8,16] parameter(0)
-  pk = f32[1,8,16] parameter(1)
-  pv = f32[1,8,16] parameter(2)
-  zero = f32[] constant(0)
-  zeros = f32[3,8,16] broadcast(zero), dimensions={}
-  ua = u32[] constant(0)
-  ub = u32[] constant(1)
-  uc = u32[] constant(2)
-  ja = s32[] convert(ua)
-  jb = s32[] convert(ub)
-  jc = s32[] convert(uc)
-  dq = f32[3,8,16] dynamic-update-slice(zeros, pq, ja, ja, ja)
-  dk = f32[3,8,16] dynamic-update-slice(zeros, pk, jb, ja, ja)
-  dv = f32[3,8,16] dynamic-update-slice(zeros, pv, jc, ja, ja)
-  sum = f32[3,8,16] add(dq, dk)
-  ROOT out = f32[3,8,16] add(sum, dv)
-}
-)";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo));
-  HloPassFix<AlgebraicSimplifier> simplifier(default_options_);
-  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&simplifier, module.get()));
-  EXPECT_TRUE(changed);
-  EXPECT_THAT(module->entry_computation()->root_instruction(),
-              GmockMatch(m::Concatenate(m::Parameter(0), m::Parameter(1),
-                                        m::Parameter(2))));
-}
-
-TEST_F(AlgebraicSimplifierTest, AddOfDisjointPadsWithGapBecomesPadOfConcat) {
-  constexpr absl::string_view kHlo = R"(
-HloModule m
-
-ENTRY main {
-  pa = f32[4,2] parameter(0)
-  pb = f32[4,2] parameter(1)
-  zero = f32[] constant(0)
-  a = f32[4,10] pad(pa, zero), padding=0_0x0_8
-  b = f32[4,10] pad(pb, zero), padding=0_0x6_2
-  ROOT out = f32[4,10] add(a, b)
-}
-)";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo));
-  AlgebraicSimplifier simplifier(default_options_);
-  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&simplifier, module.get()));
-  EXPECT_TRUE(changed);
-  const HloInstruction* root = module->entry_computation()->root_instruction();
-  EXPECT_THAT(
-      root, GmockMatch(m::Pad(m::Concatenate(m::Parameter(0),
-                                             m::Broadcast(m::ConstantScalar(0)),
-                                             m::Parameter(1)),
-                              m::ConstantScalar(0))));
-  EXPECT_EQ(root->padding_config().dimensions(1).edge_padding_low(), 0);
-  EXPECT_EQ(root->padding_config().dimensions(1).edge_padding_high(), 2);
-  EXPECT_EQ(root->operand(0)->operand(1)->shape().dimensions(1), 4);
-}
-
-TEST_F(AlgebraicSimplifierTest, AddOfOverlappingPadsIsNotFolded) {
-  constexpr absl::string_view kHlo = R"(
-HloModule m
-
-ENTRY main {
-  pa = f32[4,4] parameter(0)
-  pb = f32[4,4] parameter(1)
-  zero = f32[] constant(0)
-  a = f32[4,10] pad(pa, zero), padding=0_0x0_6
-  b = f32[4,10] pad(pb, zero), padding=0_0x2_4
-  ROOT out = f32[4,10] add(a, b)
-}
-)";
-  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHlo));
-  AlgebraicSimplifier simplifier(default_options_);
-  ASSERT_OK(RunHloPass(&simplifier, module.get()).status());
-  EXPECT_THAT(module->entry_computation()->root_instruction(),
-              GmockMatch(m::Add(m::Pad(), m::Pad())));
-}
-
 TEST_F(AlgebraicSimplifierTest, ScalarMultiplyReduction) {
   constexpr absl::string_view hlo_string = R"(
 HloModule ConstScalarMultiply

From 8756e1304c62d0aed5b8e7246d353785d4360eec Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Sat, 3 Oct 2026 15:13:00 -0700
Subject: [PATCH 0666/1000] Validate axis bounds in experimental.numpy roll

NumPy raises AxisError for out-of-bounds axes, but roll passed
them straight to the backend kernel, which failed with a confusing
runtime error. When the input rank is statically known, scalar int
axes and sequences of int axes (including range and ndarray, as
accepted by manip_ops.roll) are now normalized and bounds-checked
with a clear ValueError, matching np.roll. Behavior for valid axes
is unchanged.
---
 .../python/ops/numpy_ops/np_array_ops.py      | 28 +++++++++++++++
 .../python/ops/numpy_ops/np_array_ops_test.py | 35 +++++++++++++++++++
 2 files changed, 63 insertions(+)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index 727fd5f4605..9ec374111c7 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -16,6 +16,7 @@
 # pylint: disable=g-direct-tensorflow-import
 
 import builtins
+import collections.abc
 import enum
 import functools
 import math
@@ -1725,6 +1726,33 @@ def roll(a, shift, axis=None):  # pylint: disable=missing-docstring
   a = asarray(a)
 
   if axis is not None:
+    # NumPy raises AxisError for out-of-bounds axes instead of letting the
+    # backend kernel fail with a confusing error.
+    maybe_rank = a.shape.rank
+    if maybe_rank is not None:
+      if isinstance(axis, (int, np.integer)) and not isinstance(
+          axis, (bool, np.bool_)
+      ):
+        normalized = axis + maybe_rank if axis < 0 else axis
+        if normalized < 0 or normalized >= maybe_rank:
+          raise ValueError(
+              f'Argument `axis` (received axis={axis}) is out of bounds '
+              f'for input of rank {maybe_rank}.'
+          )
+      elif isinstance(axis, collections.abc.Iterable):
+        axes = builtins.list(axis)
+        if all(
+            isinstance(x, (int, np.integer))
+            and not isinstance(x, (bool, np.bool_))
+            for x in axes
+        ):
+          for x in axes:
+            normalized = x + maybe_rank if x < 0 else x
+            if normalized < 0 or normalized >= maybe_rank:
+              raise ValueError(
+                  f'Argument `axis` (received axis={x}) is out of bounds '
+                  f'for input of rank {maybe_rank}.'
+              )
     return manip_ops.roll(a, shift, axis)
 
   # If axis is None, the roll happens as a 1-d tensor.
diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index a08d791a183..13ea35fb9d6 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1659,6 +1659,41 @@ class ArrayMethodsTest(test.TestCase):
     with self.assertRaisesRegex(errors_impl.InvalidArgumentError, kernel_err):
       self.evaluate(flip_unknown_rank_invalid_pos(a))
 
+  def testRoll(self):
+    np.random.seed(0)
+    random_seed.set_seed(0)
+
+    def _test(*args, **kwargs):
+      expected = np.roll(*args, **kwargs)
+      raw_ans = np_array_ops.roll(*args, **kwargs)
+
+      self.assertAllEqual(expected, raw_ans)
+
+    a = np.random.rand(2, 3, 4)
+
+    # No axis rolls the flattened array.
+    _test(a, 1)
+    # A scalar shift on a single axis, including negative values.
+    _test(a, 1, axis=0)
+    _test(a, -1, axis=2)
+    _test(a, 2, axis=-1)
+    # A sequence of shifts across multiple axes.
+    _test(a, (1, 2, 3), axis=(0, 1, 2))
+    _test(a, [1, -2], axis=[-3, 2])
+    _test(a, np.array([0, 1]), axis=np.array([-1, 1]))
+
+    # Out-of-bounds axes raise, matching np.roll's AxisError.
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_array_ops.roll(a, 1, axis=3)
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_array_ops.roll(a, 1, axis=-4)
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_array_ops.roll(a, (1, 2), axis=(0, 3))
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_array_ops.roll(a, (1, -1), axis=[0, -4])
+    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+      np_array_ops.roll(a, range(2), axis=range(2, 4))
+
   def testNdim(self):
     self.assertAllEqual(0, np_array_ops.ndim(0.5))
     self.assertAllEqual(1, np_array_ops.ndim([1, 2]))

From 994c81ea5cc2ddd9262a3fc553f0dda59d31bf46 Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Mon, 5 Oct 2026 21:10:13 -0700
Subject: [PATCH 0667/1000] Address review: rely on backend kernel axis check,
 improve kernel error message

---
 .../python/ops/numpy_ops/np_array_ops.py      | 28 -------------------
 .../python/ops/numpy_ops/np_array_ops_test.py | 22 +++++++++++----
 2 files changed, 16 insertions(+), 34 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops.py b/tensorflow/python/ops/numpy_ops/np_array_ops.py
index 9ec374111c7..727fd5f4605 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops.py
@@ -16,7 +16,6 @@
 # pylint: disable=g-direct-tensorflow-import
 
 import builtins
-import collections.abc
 import enum
 import functools
 import math
@@ -1726,33 +1725,6 @@ def roll(a, shift, axis=None):  # pylint: disable=missing-docstring
   a = asarray(a)
 
   if axis is not None:
-    # NumPy raises AxisError for out-of-bounds axes instead of letting the
-    # backend kernel fail with a confusing error.
-    maybe_rank = a.shape.rank
-    if maybe_rank is not None:
-      if isinstance(axis, (int, np.integer)) and not isinstance(
-          axis, (bool, np.bool_)
-      ):
-        normalized = axis + maybe_rank if axis < 0 else axis
-        if normalized < 0 or normalized >= maybe_rank:
-          raise ValueError(
-              f'Argument `axis` (received axis={axis}) is out of bounds '
-              f'for input of rank {maybe_rank}.'
-          )
-      elif isinstance(axis, collections.abc.Iterable):
-        axes = builtins.list(axis)
-        if all(
-            isinstance(x, (int, np.integer))
-            and not isinstance(x, (bool, np.bool_))
-            for x in axes
-        ):
-          for x in axes:
-            normalized = x + maybe_rank if x < 0 else x
-            if normalized < 0 or normalized >= maybe_rank:
-              raise ValueError(
-                  f'Argument `axis` (received axis={x}) is out of bounds '
-                  f'for input of rank {maybe_rank}.'
-              )
     return manip_ops.roll(a, shift, axis)
 
   # If axis is None, the roll happens as a 1-d tensor.
diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index 13ea35fb9d6..79e3e79b895 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1682,16 +1682,26 @@ class ArrayMethodsTest(test.TestCase):
     _test(a, [1, -2], axis=[-3, 2])
     _test(a, np.array([0, 1]), axis=np.array([-1, 1]))
 
-    # Out-of-bounds axes raise, matching np.roll's AxisError.
-    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+    # Out-of-bounds axes raise, checked by the backend kernel.
+    with self.assertRaisesRegex(
+        (ValueError, errors_impl.InvalidArgumentError), 'out of bounds'
+    ):
       np_array_ops.roll(a, 1, axis=3)
-    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+    with self.assertRaisesRegex(
+        (ValueError, errors_impl.InvalidArgumentError), 'out of bounds'
+    ):
       np_array_ops.roll(a, 1, axis=-4)
-    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+    with self.assertRaisesRegex(
+        (ValueError, errors_impl.InvalidArgumentError), 'out of bounds'
+    ):
       np_array_ops.roll(a, (1, 2), axis=(0, 3))
-    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+    with self.assertRaisesRegex(
+        (ValueError, errors_impl.InvalidArgumentError), 'out of bounds'
+    ):
       np_array_ops.roll(a, (1, -1), axis=[0, -4])
-    with self.assertRaisesRegex(ValueError, 'out of bounds'):
+    with self.assertRaisesRegex(
+        (ValueError, errors_impl.InvalidArgumentError), 'out of bounds'
+    ):
       np_array_ops.roll(a, range(2), axis=range(2, 4))
 
   def testNdim(self):

From 9976b3d4f915bf556bfd29138c905b8ee7e2ef24 Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Mon, 5 Oct 2026 22:36:11 -0700
Subject: [PATCH 0668/1000] Use int64 axis/shift arithmetic and restore 'is out
 of range' message

---
 .../python/ops/numpy_ops/np_array_ops_test.py | 20 ++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index 79e3e79b895..57c21f5cf78 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1684,26 +1684,36 @@ class ArrayMethodsTest(test.TestCase):
 
     # Out-of-bounds axes raise, checked by the backend kernel.
     with self.assertRaisesRegex(
-        (ValueError, errors_impl.InvalidArgumentError), 'out of bounds'
+        (ValueError, errors_impl.InvalidArgumentError), 'out of range'
     ):
       np_array_ops.roll(a, 1, axis=3)
     with self.assertRaisesRegex(
-        (ValueError, errors_impl.InvalidArgumentError), 'out of bounds'
+        (ValueError, errors_impl.InvalidArgumentError), 'out of range'
     ):
       np_array_ops.roll(a, 1, axis=-4)
     with self.assertRaisesRegex(
-        (ValueError, errors_impl.InvalidArgumentError), 'out of bounds'
+        (ValueError, errors_impl.InvalidArgumentError), 'out of range'
     ):
       np_array_ops.roll(a, (1, 2), axis=(0, 3))
     with self.assertRaisesRegex(
-        (ValueError, errors_impl.InvalidArgumentError), 'out of bounds'
+        (ValueError, errors_impl.InvalidArgumentError), 'out of range'
     ):
       np_array_ops.roll(a, (1, -1), axis=[0, -4])
     with self.assertRaisesRegex(
-        (ValueError, errors_impl.InvalidArgumentError), 'out of bounds'
+        (ValueError, errors_impl.InvalidArgumentError), 'out of range'
     ):
       np_array_ops.roll(a, range(2), axis=range(2, 4))
 
+    # Graph-mode tracing of roll with concrete axes.
+    @def_function.function
+    def f(x, s, ax):
+      return np_array_ops.roll(x, s, axis=ax)
+
+    self.assertAllEqual(
+        np.roll(a, 1, axis=0),
+        f(a, 1, 0),
+    )
+
   def testNdim(self):
     self.assertAllEqual(0, np_array_ops.ndim(0.5))
     self.assertAllEqual(1, np_array_ops.ndim([1, 2]))

From 926095e92447ea19f9bbdd3de2234f6e680dcd35 Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Tue, 6 Oct 2026 15:08:24 -0700
Subject: [PATCH 0669/1000] Widen roll shift arithmetic to int64 and cap
 dimension sizes at int32

---
 .../python/ops/numpy_ops/np_array_ops_test.py | 22 ++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index 57c21f5cf78..8c73ba153c2 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1684,25 +1684,29 @@ class ArrayMethodsTest(test.TestCase):
 
     # Out-of-bounds axes raise, checked by the backend kernel.
     with self.assertRaisesRegex(
-        (ValueError, errors_impl.InvalidArgumentError), 'out of range'
+        errors_impl.InvalidArgumentError, 'out of range'
     ):
       np_array_ops.roll(a, 1, axis=3)
     with self.assertRaisesRegex(
-        (ValueError, errors_impl.InvalidArgumentError), 'out of range'
+        errors_impl.InvalidArgumentError, 'out of range'
     ):
       np_array_ops.roll(a, 1, axis=-4)
     with self.assertRaisesRegex(
-        (ValueError, errors_impl.InvalidArgumentError), 'out of range'
+        errors_impl.InvalidArgumentError, 'out of range'
     ):
       np_array_ops.roll(a, (1, 2), axis=(0, 3))
     with self.assertRaisesRegex(
-        (ValueError, errors_impl.InvalidArgumentError), 'out of range'
+        errors_impl.InvalidArgumentError, 'out of range'
     ):
       np_array_ops.roll(a, (1, -1), axis=[0, -4])
     with self.assertRaisesRegex(
-        (ValueError, errors_impl.InvalidArgumentError), 'out of range'
+        errors_impl.InvalidArgumentError, 'out of range'
     ):
       np_array_ops.roll(a, range(2), axis=range(2, 4))
+    with self.assertRaisesRegex(
+        errors_impl.InvalidArgumentError, 'out of range'
+    ):
+      np_array_ops.roll(a, 1, axis=1 << 33)
 
     # Graph-mode tracing of roll with concrete axes.
     @def_function.function
@@ -1713,6 +1717,14 @@ class ArrayMethodsTest(test.TestCase):
         np.roll(a, 1, axis=0),
         f(a, 1, 0),
     )
+    self.assertAllEqual(
+        np.roll(a, 1, axis=0),
+        f(a, constant_op.constant(1), constant_op.constant(0)),
+    )
+    with self.assertRaisesRegex(
+        errors_impl.InvalidArgumentError, 'out of range'
+    ):
+      f(a, 1, 3)
 
   def testNdim(self):
     self.assertAllEqual(0, np_array_ops.ndim(0.5))

From a344fe7fbe6288c491ceb63555623943511ffab5 Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Wed, 7 Oct 2026 07:14:52 -0700
Subject: [PATCH 0670/1000] Add zero-size fast exit before Roll functor
 dispatch

---
 tensorflow/core/kernels/roll_op.cc | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/core/kernels/roll_op.cc b/tensorflow/core/kernels/roll_op.cc
index bd12d5d914d..7a0629eeb5d 100644
--- a/tensorflow/core/kernels/roll_op.cc
+++ b/tensorflow/core/kernels/roll_op.cc
@@ -108,6 +108,7 @@ class RollOp : public OpKernel {
     Tensor* output = nullptr;
     OP_REQUIRES_OK(context,
                    context->allocate_output(0, input.shape(), &output));
+    if (num_elements == 0) return;
     auto input_flat = input.flat<T>().data();
     auto output_flat = output->flat<T>().data();
 

From b720c5dffcf2931971cf4981b7dab7ef20700821 Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Wed, 7 Oct 2026 15:23:01 -0700
Subject: [PATCH 0671/1000] Add empty-array and scalar roll test cases

---
 tensorflow/python/ops/numpy_ops/np_array_ops_test.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index 8c73ba153c2..3a09a5c0a33 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1682,6 +1682,11 @@ class ArrayMethodsTest(test.TestCase):
     _test(a, [1, -2], axis=[-3, 2])
     _test(a, np.array([0, 1]), axis=np.array([-1, 1]))
 
+    # Empty and scalar arrays.
+    _test(np.empty((0, 3, 4)), 1, axis=1)
+    _test(np.empty((3, 4, 0)), 1, axis=1)
+    _test(np.array(42.0), 1)
+
     # Out-of-bounds axes raise, checked by the backend kernel.
     with self.assertRaisesRegex(
         errors_impl.InvalidArgumentError, 'out of range'

From 51be0efa8dc351c962dc8d4227a74b0b62194569 Mon Sep 17 00:00:00 2001
From: Christian Sigg <csigg@google.com>
Date: Tue, 13 Oct 2026 02:26:07 -0700
Subject: [PATCH 0672/1000] [xla][gpu] Emit ignore_denormal_mode on
 floating-point atomic additions

Set the target-independent `ignore_denormal_mode` unit attribute on the
`llvm.atomicrmw fadd` ops created by `emitNVidiaAtomicFAdd` and
`emitAMDAtomicFAdd`. It translates to `!atomic.ignore.denormal.mode` in LLVM IR
and replaces the ROCDL-only `rocdl.ignore_denormal_mode` attribute, so NVPTX
also keeps lowering to native `atom.add` once `-nvptx-allow-ftz-atomics`
defaults back to false. The AMD path still sets the ROCDL `no_remote_memory`
and `no_fine_grained_memory` attributes.

`gpu_atomic_test` now checks for the metadata on f32, f16, bf16 and f64 atomic
adds on both CUDA and ROCm.

Also drop the redundant `mlir:DialectUtils` dep from `:passes`, flagged by
build_cleaner. `StaticValueUtils.h` is already provided by `mlir:IR`.

Depends on llvm/llvm-project#217588 (commit 1491508743104), which adds
`ignore_denormal_mode` to `LLVM::AtomicRMWOp`.

PiperOrigin-RevId: 998488034
---
 .../xla/backends/gpu/tests/gpu_atomic_test.cc | 88 +++++++++++++++++--
 .../xla/xla/codegen/emitters/transforms/BUILD |  1 -
 .../emitters/transforms/lower_tensors.cc      |  9 +-
 .../transforms/tests/lower_tensors.mlir       | 32 +++----
 4 files changed, 104 insertions(+), 26 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/tests/gpu_atomic_test.cc b/third_party/xla/xla/backends/gpu/tests/gpu_atomic_test.cc
index f22ba6a6579..9994bc3d1e1 100644
--- a/third_party/xla/xla/backends/gpu/tests/gpu_atomic_test.cc
+++ b/third_party/xla/xla/backends/gpu/tests/gpu_atomic_test.cc
@@ -102,16 +102,91 @@ TEST_F(GpuAtomicTest, TestAddAtomicF32) {
 )";
 
   ASSERT_OK(CompileAndVerifyIr(hlo_string, IsBuiltWithRocm() ? R"(
-CHECK: atomicrmw fadd ptr %[[ADDR:.*]], float %[[VALUE:.*]] syncscope("agent-one-as") monotonic
+CHECK: atomicrmw fadd ptr %[[ADDR:.*]], float %[[VALUE:.*]] syncscope("agent-one-as") monotonic{{.*}}!atomic.ignore.denormal.mode
 )"
                                                              : R"(
-CHECK: atomicrmw fadd ptr %[[ADDR:.*]], float %[[VALUE:.*]] monotonic
+CHECK: atomicrmw fadd ptr %[[ADDR:.*]], float %[[VALUE:.*]] monotonic{{.*}}!atomic.ignore.denormal.mode
+)"));
+}
+
+TEST_F(GpuAtomicTest, TestAddAtomicF16) {
+  if (!IsBuiltWithRocm() &&
+      !device_description().cuda_compute_capability().SupportsAllFeaturesOf(
+          stream_executor::CudaComputeCapability::Volta())) {
+    return;
+  }
+  const char* hlo_string = R"(
+    HloModule TensorFlowScatterV1
+
+    update_f16 (lhs: f16[], rhs: f16[]) -> f16[] {
+      lhs = f16[] parameter(0)
+      rhs = f16[] parameter(1)
+      ROOT add = f16[] add(lhs, rhs)
+    }
+
+    ENTRY main {
+      operand = f16[3,3] parameter(0)
+      indices = s32[2] parameter(1)
+      updates = f16[2,3] parameter(2)
+      ROOT scatter = f16[3,3] scatter(operand, indices, updates),
+          to_apply=update_f16,
+          update_window_dims={1},
+          inserted_window_dims={0},
+          scatter_dims_to_operand_dims={0},
+          index_vector_dim=1, unique_indices=false
+    }
+)";
+
+  // ROCm GPUs with packed fp16 atomics add a <2 x half> to a global pointer.
+  ASSERT_OK(CompileAndVerifyIr(hlo_string, IsBuiltWithRocm() ? R"(
+CHECK: atomicrmw fadd ptr {{.*}}, {{(<2 x half>|half)}} %[[VALUE:.*]] syncscope("agent-one-as") monotonic{{.*}}!atomic.ignore.denormal.mode
+)"
+                                                             : R"(
+CHECK: atomicrmw fadd ptr %[[ADDR:.*]], half %[[VALUE:.*]] monotonic{{.*}}!atomic.ignore.denormal.mode
+)"));
+}
+
+TEST_F(GpuAtomicTest, TestAddAtomicBF16) {
+  if (!IsBuiltWithRocm() &&
+      !device_description().cuda_compute_capability().SupportsAllFeaturesOf(
+          stream_executor::CudaComputeCapability::Hopper())) {
+    return;
+  }
+  const char* hlo_string = R"(
+    HloModule TensorFlowScatterV1
+
+    update_bf16 (lhs: bf16[], rhs: bf16[]) -> bf16[] {
+      lhs = bf16[] parameter(0)
+      rhs = bf16[] parameter(1)
+      ROOT add = bf16[] add(lhs, rhs)
+    }
+
+    ENTRY main {
+      operand = bf16[3,3] parameter(0)
+      indices = s32[2] parameter(1)
+      updates = bf16[2,3] parameter(2)
+      ROOT scatter = bf16[3,3] scatter(operand, indices, updates),
+          to_apply=update_bf16,
+          update_window_dims={1},
+          inserted_window_dims={0},
+          scatter_dims_to_operand_dims={0},
+          index_vector_dim=1, unique_indices=false
+    }
+)";
+
+  // ROCm GPUs with packed bf16 atomics add a <2 x bfloat> to a global pointer.
+  ASSERT_OK(CompileAndVerifyIr(hlo_string, IsBuiltWithRocm() ? R"(
+CHECK: atomicrmw fadd ptr {{.*}}, {{(<2 x bfloat>|bfloat)}} %[[VALUE:.*]] syncscope("agent-one-as") monotonic{{.*}}!atomic.ignore.denormal.mode
+)"
+                                                             : R"(
+CHECK: atomicrmw fadd ptr %[[ADDR:.*]], bfloat %[[VALUE:.*]] monotonic{{.*}}!atomic.ignore.denormal.mode
 )"));
 }
 
 TEST_F(GpuAtomicTest, TestAddAtomicF64) {
   // Atomic add required sm_60 or above.
-  if (device_description().cuda_compute_capability().SupportsAllFeaturesOf(
+  if (!IsBuiltWithRocm() &&
+      !device_description().cuda_compute_capability().SupportsAllFeaturesOf(
           stream_executor::CudaComputeCapability::Pascal())) {
     return;
   }
@@ -138,8 +213,11 @@ TEST_F(GpuAtomicTest, TestAddAtomicF64) {
     }
 )";
 
-  ASSERT_OK(CompileAndVerifyIr(hlo_string, R"(
-CHECK: atomicrmw fadd ptr %[[ADDR:.*]], double %[[VALUE:.*]] monotonic
+  ASSERT_OK(CompileAndVerifyIr(hlo_string, IsBuiltWithRocm() ? R"(
+CHECK: atomicrmw fadd ptr %[[ADDR:.*]], double %[[VALUE:.*]] syncscope("agent-one-as") monotonic{{.*}}!atomic.ignore.denormal.mode
+)"
+                                                             : R"(
+CHECK: atomicrmw fadd ptr %[[ADDR:.*]], double %[[VALUE:.*]] monotonic{{.*}}!atomic.ignore.denormal.mode
 )"));
 }
 
diff --git a/third_party/xla/xla/codegen/emitters/transforms/BUILD b/third_party/xla/xla/codegen/emitters/transforms/BUILD
index 33c393a9b7b..683f8a3b84e 100644
--- a/third_party/xla/xla/codegen/emitters/transforms/BUILD
+++ b/third_party/xla/xla/codegen/emitters/transforms/BUILD
@@ -251,7 +251,6 @@ cc_library(
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:ComplexDialect",
         "@llvm-project//mlir:DataLayoutInterfaces",
-        "@llvm-project//mlir:DialectUtils",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:FuncUtils",
         "@llvm-project//mlir:GPUDialect",
diff --git a/third_party/xla/xla/codegen/emitters/transforms/lower_tensors.cc b/third_party/xla/xla/codegen/emitters/transforms/lower_tensors.cc
index b64f626f12c..0bc4fed72cd 100644
--- a/third_party/xla/xla/codegen/emitters/transforms/lower_tensors.cc
+++ b/third_party/xla/xla/codegen/emitters/transforms/lower_tensors.cc
@@ -1015,8 +1015,10 @@ class RewriteAtomicRMW : public OpRewritePattern<AtomicRMWOp> {
       return failure();
     }
 
-    ml::AtomicRMWOp::create(b, loc, ml::AtomicBinOp::fadd, addr, modifier_arg,
-                            ml::AtomicOrdering::monotonic);
+    auto rmw_op =
+        ml::AtomicRMWOp::create(b, loc, ml::AtomicBinOp::fadd, addr,
+                                modifier_arg, ml::AtomicOrdering::monotonic);
+    rmw_op.setIgnoreDenormalMode(true);
     return success();
   }
 
@@ -1125,17 +1127,16 @@ class RewriteAtomicRMW : public OpRewritePattern<AtomicRMWOp> {
     auto op = ml::AtomicRMWOp::create(
         b, loc, ml::AtomicBinOp::fadd, addr, modifier_arg,
         ml::AtomicOrdering::monotonic, "agent-one-as");
+    op.setIgnoreDenormalMode(true);
 
     auto unitAttr = b.getUnitAttr();
     auto* rocdl =
         op->getContext()->getOrLoadDialect<mlir::ROCDL::ROCDLDialect>();
     auto noRemoteMemHelper = rocdl->getNoRemoteMemoryAttrHelper();
     auto noFineMemHelper = rocdl->getNoFineGrainedMemoryAttrHelper();
-    auto ignoreDenormalModeHelper = rocdl->getIgnoreDenormalModeAttrHelper();
 
     noRemoteMemHelper.setAttr(op, unitAttr);
     noFineMemHelper.setAttr(op, unitAttr);
-    ignoreDenormalModeHelper.setAttr(op, unitAttr);
 
     return success();
   }
diff --git a/third_party/xla/xla/codegen/emitters/transforms/tests/lower_tensors.mlir b/third_party/xla/xla/codegen/emitters/transforms/tests/lower_tensors.mlir
index 97dcdd8c6bd..7a58d6c303a 100644
--- a/third_party/xla/xla/codegen/emitters/transforms/tests/lower_tensors.mlir
+++ b/third_party/xla/xla/codegen/emitters/transforms/tests/lower_tensors.mlir
@@ -725,27 +725,27 @@ func.func @direct_atomic_rmw_fadd_f32(%in: tensor<8xf32>,
 // CHECK-PASCAL-LABEL: @direct_atomic_rmw_fadd_f32
 // CHECK-PASCAL: %[[C2:.*]] = arith.constant 2
 // CHECK-PASCAL: %[[ADDR:.*]] = llvm.getelementptr
-// CHECK-PASCAL: llvm.atomicrmw fadd %[[ADDR]], %[[C2]] monotonic
+// CHECK-PASCAL: llvm.atomicrmw ignore_denormal_mode fadd %[[ADDR]], %[[C2]] monotonic
 
 // CHECK-VOLTA-LABEL: @direct_atomic_rmw_fadd_f32
 // CHECK-VOLTA: %[[C2:.*]] = arith.constant 2
 // CHECK-VOLTA: %[[ADDR:.*]] = llvm.getelementptr
-// CHECK-VOLTA: llvm.atomicrmw fadd %[[ADDR]], %[[C2]] monotonic
+// CHECK-VOLTA: llvm.atomicrmw ignore_denormal_mode fadd %[[ADDR]], %[[C2]] monotonic
 
 // CHECK-AMPERE-LABEL: @direct_atomic_rmw_fadd_f32
 // CHECK-AMPERE: %[[C2:.*]] = arith.constant 2
 // CHECK-AMPERE: %[[ADDR:.*]] = llvm.getelementptr
-// CHECK-AMPERE: llvm.atomicrmw fadd %[[ADDR]], %[[C2]] monotonic
+// CHECK-AMPERE: llvm.atomicrmw ignore_denormal_mode fadd %[[ADDR]], %[[C2]] monotonic
 
 // CHECK-GFX908-MI100-LABEL: @direct_atomic_rmw_fadd_f32
 // CHECK-GFX908-MI100: %[[C2:.*]] = arith.constant 2
 // CHECK-GFX908-MI100: %[[ADDR:.*]] = llvm.getelementptr
-// CHECK-GFX908-MI100: llvm.atomicrmw fadd %[[ADDR]], %[[C2]] syncscope("agent-one-as") monotonic
+// CHECK-GFX908-MI100: llvm.atomicrmw ignore_denormal_mode fadd %[[ADDR]], %[[C2]] syncscope("agent-one-as") monotonic {rocdl.no_fine_grained_memory, rocdl.no_remote_memory}
 
 // CHECK-GFX90A-MI200-LABEL: @direct_atomic_rmw_fadd_f32
 // CHECK-GFX90A-MI200: %[[C2:.*]] = arith.constant 2
 // CHECK-GFX90A-MI200: %[[ADDR:.*]] = llvm.getelementptr
-// CHECK-GFX90A-MI200: llvm.atomicrmw fadd %[[ADDR]], %[[C2]] syncscope("agent-one-as") monotonic
+// CHECK-GFX90A-MI200: llvm.atomicrmw ignore_denormal_mode fadd %[[ADDR]], %[[C2]] syncscope("agent-one-as") monotonic {rocdl.no_fine_grained_memory, rocdl.no_remote_memory}
 
 // -----
 
@@ -765,12 +765,12 @@ func.func @direct_atomic_rmw_fadd_f16(%in: tensor<8xf16>,
 // CHECK-VOLTA-LABEL: @direct_atomic_rmw_fadd_f16
 // CHECK-VOLTA: %[[C2:.*]] = arith.constant 2
 // CHECK-VOLTA: %[[ADDR:.*]] = llvm.getelementptr
-// CHECK-VOLTA: llvm.atomicrmw fadd %[[ADDR]], %[[C2]] monotonic
+// CHECK-VOLTA: llvm.atomicrmw ignore_denormal_mode fadd %[[ADDR]], %[[C2]] monotonic
 
 // CHECK-AMPERE-LABEL: @direct_atomic_rmw_fadd_f16
 // CHECK-AMPERE: %[[C2:.*]] = arith.constant 2
 // CHECK-AMPERE: %[[ADDR:.*]] = llvm.getelementptr
-// CHECK-AMPERE: llvm.atomicrmw fadd %[[ADDR]], %[[C2]] monotonic
+// CHECK-AMPERE: llvm.atomicrmw ignore_denormal_mode fadd %[[ADDR]], %[[C2]] monotonic
 
 // CHECK-GFX908-MI100-LABEL: @direct_atomic_rmw_fadd_f16
 // CHECK-GFX908-MI100: %[[CST:.*]] = arith.constant 2
@@ -788,7 +788,7 @@ func.func @direct_atomic_rmw_fadd_f16(%in: tensor<8xf16>,
 // CHECK-GFX908-MI100: %[[VAL_SHIFT:.*]] = llvm.shl %[[VAL_WIDE]], %[[SHIFT]]
 // CHECK-GFX908-MI100: %[[ADDR:.*]] = llvm.inttoptr %[[ADDR_MASKED]]
 // CHECK-GFX908-MI100: %[[VAL:.*]] = llvm.bitcast %[[VAL_SHIFT]] : i32 to vector<2xf16>
-// CHECK-GFX908-MI100: llvm.atomicrmw fadd %[[ADDR]], %[[VAL]] syncscope("agent-one-as") monotonic
+// CHECK-GFX908-MI100: llvm.atomicrmw ignore_denormal_mode fadd %[[ADDR]], %[[VAL]] syncscope("agent-one-as") monotonic {rocdl.no_fine_grained_memory, rocdl.no_remote_memory}
 
 
 // CHECK-GFX90A-MI200-LABEL: @direct_atomic_rmw_fadd_f16
@@ -807,7 +807,7 @@ func.func @direct_atomic_rmw_fadd_f16(%in: tensor<8xf16>,
 // CHECK-GFX90A-MI200: %[[VAL_SHIFT:.*]] = llvm.shl %[[VAL_WIDE]], %[[SHIFT]]
 // CHECK-GFX90A-MI200: %[[ADDR:.*]] = llvm.inttoptr %[[ADDR_MASKED]]
 // CHECK-GFX90A-MI200: %[[VAL:.*]] = llvm.bitcast %[[VAL_SHIFT]] : i32 to vector<2xf16>
-// CHECK-GFX90A-MI200: llvm.atomicrmw fadd %[[ADDR]], %[[VAL]] syncscope("agent-one-as") monotonic
+// CHECK-GFX90A-MI200: llvm.atomicrmw ignore_denormal_mode fadd %[[ADDR]], %[[VAL]] syncscope("agent-one-as") monotonic {rocdl.no_fine_grained_memory, rocdl.no_remote_memory}
 
 // -----
 
@@ -856,7 +856,7 @@ func.func @direct_atomic_rmw_fadd_bf16(%in: tensor<8xbf16>,
 // CHECK-HOPPER-LABEL: @direct_atomic_rmw_fadd_bf16
 // CHECK-HOPPER: %[[C2:.*]] = arith.constant 2
 // CHECK-HOPPER: %[[ADDR:.*]] = llvm.getelementptr
-// CHECK-HOPPER: llvm.atomicrmw fadd %[[ADDR]], %[[C2]] monotonic
+// CHECK-HOPPER: llvm.atomicrmw ignore_denormal_mode fadd %[[ADDR]], %[[C2]] monotonic
 
 // -----
 
@@ -892,7 +892,7 @@ func.func @direct_atomic_rmw_fadd_bf16_widened(%in: tensor<8xbf16>,
 // CHECK-GFX942-MI300: %[[VAL_SHIFT:.*]] = llvm.shl %[[VAL_WIDE]], %[[SHIFT]]
 // CHECK-GFX942-MI300: %[[ADDR_PTR:.*]] = llvm.inttoptr %[[ADDR_MASKED]]
 // CHECK-GFX942-MI300: %[[VAL:.*]] = llvm.bitcast %[[VAL_SHIFT]] : i32 to vector<2xbf16>
-// CHECK-GFX942-MI300: llvm.atomicrmw fadd %[[ADDR_PTR]], %[[VAL]] syncscope("agent-one-as") monotonic
+// CHECK-GFX942-MI300: llvm.atomicrmw ignore_denormal_mode fadd %[[ADDR_PTR]], %[[VAL]] syncscope("agent-one-as") monotonic {rocdl.no_fine_grained_memory, rocdl.no_remote_memory}
 
 // -----
 
@@ -909,28 +909,28 @@ func.func @direct_atomic_rmw_fadd_f64(%in: tensor<8xf64>,
 // CHECK-PASCAL-LABEL: @direct_atomic_rmw_fadd_f64
 // CHECK-PASCAL: %[[C2:.*]] = arith.constant 2
 // CHECK-PASCAL: %[[ADDR:.*]] = llvm.getelementptr
-// CHECK-PASCAL: llvm.atomicrmw fadd %[[ADDR]], %[[C2]] monotonic
+// CHECK-PASCAL: llvm.atomicrmw ignore_denormal_mode fadd %[[ADDR]], %[[C2]] monotonic
 
 // CHECK-VOLTA-LABEL: @direct_atomic_rmw_fadd_f64
 // CHECK-VOLTA: %[[C2:.*]] = arith.constant 2
 // CHECK-VOLTA: %[[ADDR:.*]] = llvm.getelementptr
-// CHECK-VOLTA: llvm.atomicrmw fadd %[[ADDR]], %[[C2]] monotonic
+// CHECK-VOLTA: llvm.atomicrmw ignore_denormal_mode fadd %[[ADDR]], %[[C2]] monotonic
 
 // CHECK-AMPERE-LABEL: @direct_atomic_rmw_fadd_f64
 // CHECK-AMPERE: %[[C2:.*]] = arith.constant 2
 // CHECK-AMPERE: %[[ADDR:.*]] = llvm.getelementptr
-// CHECK-AMPERE: llvm.atomicrmw fadd %[[ADDR]], %[[C2]] monotonic
+// CHECK-AMPERE: llvm.atomicrmw ignore_denormal_mode fadd %[[ADDR]], %[[C2]] monotonic
 
 // CHECK-GFX908-MI100-LABEL: @direct_atomic_rmw_fadd_f64
 // CHECK-GFX908-MI100: %[[C2:.*]] = arith.constant 2
 // CHECK-GFX908-MI100: %[[ADDR:.*]] = llvm.getelementptr
-// CHECK-GFX908-MI100: llvm.atomicrmw fadd %[[ADDR]], %[[C2]] syncscope("agent-one-as") monotonic
+// CHECK-GFX908-MI100: llvm.atomicrmw ignore_denormal_mode fadd %[[ADDR]], %[[C2]] syncscope("agent-one-as") monotonic {rocdl.no_fine_grained_memory, rocdl.no_remote_memory}
 
 
 // CHECK-GFX90A-MI200-LABEL: @direct_atomic_rmw_fadd_f64
 // CHECK-GFX90A-MI200: %[[C2:.*]] = arith.constant 2
 // CHECK-GFX90A-MI200: %[[ADDR:.*]] = llvm.getelementptr
-// CHECK-GFX90A-MI200: llvm.atomicrmw fadd %[[ADDR]], %[[C2]] syncscope("agent-one-as") monotonic
+// CHECK-GFX90A-MI200: llvm.atomicrmw ignore_denormal_mode fadd %[[ADDR]], %[[C2]] syncscope("agent-one-as") monotonic {rocdl.no_fine_grained_memory, rocdl.no_remote_memory}
 
 
 // -----

From cd9b84683456472bb3273ad6f3e090dd9b397da9 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Tue, 13 Oct 2026 02:31:41 -0700
Subject: [PATCH 0673/1000] [XLA:GPU] Use absl::call_once in
 StreamExecutorCommon::GetDeviceDescription.

PiperOrigin-RevId: 998490694
---
 third_party/xla/xla/stream_executor/BUILD     |  6 ++--
 .../stream_executor/stream_executor_common.cc | 15 +++++-----
 .../stream_executor/stream_executor_common.h  | 18 ++++--------
 .../stream_executor/stream_executor_test.cc   | 29 +++++++++++++++++--
 4 files changed, 44 insertions(+), 24 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/BUILD b/third_party/xla/xla/stream_executor/BUILD
index 07847c0a9eb..b56cad132a8 100644
--- a/third_party/xla/xla/stream_executor/BUILD
+++ b/third_party/xla/xla/stream_executor/BUILD
@@ -1144,11 +1144,10 @@ cc_library(
         ":platform",
         ":stream_executor_h",
         "//xla/tsl/util:env_var",
-        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/base",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/synchronization",
     ],
 )
 
@@ -1341,11 +1340,12 @@ xla_cc_test(
     name = "stream_executor_test",
     srcs = ["stream_executor_test.cc"],
     deps = [
+        ":device_description",
         ":platform",
         ":platform_manager",
         ":stream_executor_h",
         "//xla/stream_executor/host:host_platform",
-        "//xla/tsl/platform:statusor",
+        "//xla/tsl/platform:env",
         "//xla/tsl/platform:test",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
diff --git a/third_party/xla/xla/stream_executor/stream_executor_common.cc b/third_party/xla/xla/stream_executor/stream_executor_common.cc
index ebeff36f1de..fd25ca014a8 100644
--- a/third_party/xla/xla/stream_executor/stream_executor_common.cc
+++ b/third_party/xla/xla/stream_executor/stream_executor_common.cc
@@ -17,11 +17,12 @@ limitations under the License.
 
 #include <cstdint>
 #include <memory>
+#include <utility>
 
+#include "absl/base/call_once.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/statusor.h"
-#include "absl/synchronization/mutex.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/platform.h"
 #include "xla/tsl/util/env_var.h"
@@ -42,12 +43,12 @@ StreamExecutorCommon::StreamExecutorCommon(const Platform* platform)
       memory_limit_bytes_(GetMemoryLimitBytesFromEnvironmentVariable()) {}
 
 const DeviceDescription& StreamExecutorCommon::GetDeviceDescription() const {
-  absl::MutexLock lock(mu_);
-  if (device_description_ != nullptr) {
-    return *device_description_;
-  }
-
-  device_description_ = CreateDeviceDescription().value();
+  absl::call_once(device_description_once_, [this]() {
+    absl::StatusOr<std::unique_ptr<DeviceDescription>> desc =
+        CreateDeviceDescription();
+    CHECK_OK(desc);
+    device_description_ = *std::move(desc);
+  });
   return *device_description_;
 }
 
diff --git a/third_party/xla/xla/stream_executor/stream_executor_common.h b/third_party/xla/xla/stream_executor/stream_executor_common.h
index 1e755d637b0..3b19542210c 100644
--- a/third_party/xla/xla/stream_executor/stream_executor_common.h
+++ b/third_party/xla/xla/stream_executor/stream_executor_common.h
@@ -1,5 +1,3 @@
-#include "xla/stream_executor/activate_context.h"
-#include "xla/stream_executor/device_description.h"
 /* Copyright 2015 The OpenXLA Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -21,8 +19,9 @@ limitations under the License.
 #include <cstdint>
 #include <memory>
 
-#include "absl/base/thread_annotations.h"
-#include "absl/synchronization/mutex.h"
+#include "absl/base/call_once.h"
+#include "xla/stream_executor/activate_context.h"
+#include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/stream_executor.h"
 
@@ -54,19 +53,14 @@ class StreamExecutorCommon : public StreamExecutor {
   int64_t GetMemoryLimitBytes() const override { return memory_limit_bytes_; }
 
  private:
-  // Reader/writer lock for mutable data structures on this StreamExecutor.
-  //
-  // Mutable so that caching functions (like DeviceDescription, AsBlas, etc.)
-  // can acquire the lock on their first (mutating) call as well.
-  mutable absl::Mutex mu_;
-
   // Reference to the platform that created this executor.
   const Platform* platform_;
 
+  mutable absl::once_flag device_description_once_;
+
   // Slot to cache the owned DeviceDescription for the underlying device
   // once it has been queried from DeviceDescription().
-  mutable std::unique_ptr<DeviceDescription> device_description_
-      ABSL_GUARDED_BY(mu_);
+  mutable std::unique_ptr<DeviceDescription> device_description_;
 
   // Memory limit in bytes. Value less or equal to 0 indicates there is no
   // limit.
diff --git a/third_party/xla/xla/stream_executor/stream_executor_test.cc b/third_party/xla/xla/stream_executor/stream_executor_test.cc
index d582e34f21d..a3c3364b708 100644
--- a/third_party/xla/xla/stream_executor/stream_executor_test.cc
+++ b/third_party/xla/xla/stream_executor/stream_executor_test.cc
@@ -17,16 +17,20 @@ limitations under the License.
 
 #include <cstdint>
 #include <memory>
+#include <vector>
 
 #include <gtest/gtest.h>
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
+#include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/platform_manager.h"
-#include "xla/tsl/platform/statusor.h"
+#include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/test.h"
+#include "xla/tsl/platform/threadpool.h"
 
 namespace stream_executor {
+namespace {
 
 struct TestResource : public StreamExecutor::Resource {
   TestResource() = default;
@@ -34,7 +38,7 @@ struct TestResource : public StreamExecutor::Resource {
   int32_t value = 0;
 };
 
-static absl::StatusOr<StreamExecutor*> NewStreamExecutor() {
+absl::StatusOr<StreamExecutor*> NewStreamExecutor() {
   ABSL_ASSIGN_OR_RETURN(auto platform, PlatformManager::PlatformWithName("Host"));
   ABSL_ASSIGN_OR_RETURN(auto stream_exec, platform->ExecutorForDevice(0));
   return stream_exec;
@@ -60,4 +64,25 @@ TEST(StreamExecutorTest, GetOrCreateResource) {
   EXPECT_EQ(executor->GetOrNullResource<TestResource>(), resource);
 }
 
+TEST(StreamExecutorTest, GetDeviceDescriptionConcurrent) {
+  ASSERT_OK_AND_ASSIGN(StreamExecutor * executor, NewStreamExecutor());
+
+  constexpr int kNumThreads = 8;
+  std::vector<const DeviceDescription*> descriptions(kNumThreads, nullptr);
+  {
+    tsl::thread::ThreadPool pool(tsl::Env::Default(), "test_pool", kNumThreads);
+    for (int i = 0; i < kNumThreads; ++i) {
+      pool.Schedule([executor, &descriptions, i]() {
+        descriptions[i] = &executor->GetDeviceDescription();
+      });
+    }
+  }
+
+  ASSERT_NE(descriptions[0], nullptr);
+  for (int i = 1; i < kNumThreads; ++i) {
+    EXPECT_EQ(descriptions[i], descriptions[0]);
+  }
+}
+
+}  // namespace
 }  // namespace stream_executor

From 92725d3768fd8b1b123cd1c7cfb06db2d5ecb2e6 Mon Sep 17 00:00:00 2001
From: Sai Ganesh Muthuraman <saiganeshm@google.com>
Date: Tue, 13 Oct 2026 02:37:27 -0700
Subject: [PATCH 0674/1000] Support bracketed IPv6 host-port pairs in TSL
 profiler capture.

PiperOrigin-RevId: 998493498
---
 .../xla/xla/tsl/profiler/rpc/client/BUILD     |   8 +-
 .../profiler/rpc/client/capture_profile.cc    |  30 ++++-
 .../remote_profiler_session_manager_test.cc   |  69 ++++++++++
 third_party/xla/xla/tsl/profiler/utils/BUILD  |   3 +
 .../xla/tsl/profiler/utils/session_manager.cc |  43 +++++-
 .../profiler/utils/session_manager_test.cc    | 122 +++++++++++++++++-
 6 files changed, 260 insertions(+), 15 deletions(-)

diff --git a/third_party/xla/xla/tsl/profiler/rpc/client/BUILD b/third_party/xla/xla/tsl/profiler/rpc/client/BUILD
index ed680d9157b..8678b750fac 100644
--- a/third_party/xla/xla/tsl/profiler/rpc/client/BUILD
+++ b/third_party/xla/xla/tsl/profiler/rpc/client/BUILD
@@ -1,4 +1,4 @@
-load("//xla/tsl:tsl.bzl", "internal_visibility")
+load("//xla/tsl:tsl.bzl", "if_oss", "internal_visibility")
 load(
     "//xla/tsl/platform:build_config.bzl",
     "tf_protos_profiler_service",
@@ -289,10 +289,14 @@ tsl_cc_test(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/time",
         "@com_google_googletest//:gtest_main",
+        "@grpc//:grpc++",
         "@tsl//tsl/platform:platform_port",
         "@tsl//tsl/profiler/lib:profiler_factory_impl",
         "@tsl//tsl/profiler/lib:profiler_session",
         "@tsl//tsl/profiler/lib:profiler_session_impl",
         "@tsl//tsl/profiler/protobuf:profiler_options_proto_cc",
-    ] + tf_protos_profiler_service(),
+    ] + if_oss([
+        "@tsl//tsl/profiler/protobuf:profiler_service_cc_grpc_proto",
+        "@tsl//tsl/profiler/protobuf:profiler_service_proto_cc",
+    ]) + tf_protos_profiler_service(),
 )
diff --git a/third_party/xla/xla/tsl/profiler/rpc/client/capture_profile.cc b/third_party/xla/xla/tsl/profiler/rpc/client/capture_profile.cc
index 850ec4261c2..ec27f2d8ff5 100644
--- a/third_party/xla/xla/tsl/profiler/rpc/client/capture_profile.cc
+++ b/third_party/xla/xla/tsl/profiler/rpc/client/capture_profile.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 #include "xla/tsl/profiler/rpc/client/capture_profile.h"
 
+#include <cstddef>
 #include <cstdint>
 #include <iostream>
 #include <memory>
@@ -26,9 +27,10 @@ limitations under the License.
 #include "absl/log/log.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
+#include "absl/strings/match.h"
 #include "absl/strings/str_join.h"
-#include "absl/strings/str_split.h"
 #include "absl/strings/string_view.h"
+#include "absl/strings/strip.h"
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
 #include "xla/tsl/platform/errors.h"
@@ -104,15 +106,31 @@ NewProfileSessionRequest PopulateNewProfileSessionRequest(
     absl::string_view repository_root, absl::string_view session_id,
     const RemoteProfilerSessionManagerOptions& opts) {
   NewProfileSessionRequest request;
-  std::vector<absl::string_view> parts =
-      absl::StrSplit(opts.service_addresses(0), ':');
-  DCHECK(!parts.empty());
+  absl::string_view host_port = opts.service_addresses(0);
+  absl::string_view host = host_port;
+  const size_t colon_pos = host_port.rfind(':');
+  const size_t close_bracket_pos = host_port.rfind(']');
+  const bool has_unbracketed_port =
+      close_bracket_pos == absl::string_view::npos &&
+      host_port.find(':') == colon_pos;
+  const bool has_bracketed_port =
+      close_bracket_pos != absl::string_view::npos &&
+      colon_pos > close_bracket_pos;
+  if (colon_pos != absl::string_view::npos &&
+      (has_unbracketed_port || has_bracketed_port)) {
+    host = host_port.substr(0, colon_pos);
+  }
+  if (absl::StartsWith(host, "[") && absl::EndsWith(host, "]")) {
+    host = absl::StripPrefix(host, "[");
+    host = absl::StripSuffix(host, "]");
+  }
+  DCHECK(!host.empty());
 
   *request.mutable_request() =
-      PopulateProfileRequest(repository_root, session_id, parts[0], opts);
+      PopulateProfileRequest(repository_root, session_id, host, opts);
   request.set_repository_root(repository_root.data(), repository_root.size());
   request.set_session_id(session_id.data(), session_id.size());
-  for (const auto& hostname : opts.service_addresses()) {
+  for (const std::string& hostname : opts.service_addresses()) {
     request.add_hosts(hostname);
   }
   return request;
diff --git a/third_party/xla/xla/tsl/profiler/rpc/client/remote_profiler_session_manager_test.cc b/third_party/xla/xla/tsl/profiler/rpc/client/remote_profiler_session_manager_test.cc
index a7607639517..2392400a2f9 100644
--- a/third_party/xla/xla/tsl/profiler/rpc/client/remote_profiler_session_manager_test.cc
+++ b/third_party/xla/xla/tsl/profiler/rpc/client/remote_profiler_session_manager_test.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 #include "xla/tsl/profiler/rpc/client/remote_profiler_session_manager.h"
 
+#include <cstddef>
 #include <cstdint>
 #include <memory>
 #include <string>
@@ -25,15 +26,19 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
+#include "grpcpp/security/server_credentials.h"
+#include "grpcpp/server_builder.h"
 #include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/status_matchers.h"  // IWYU pragma: keep
 #include "xla/tsl/platform/test.h"
 #include "xla/tsl/profiler/rpc/client/profiler_client_test_util.h"
 #include "xla/tsl/profiler/rpc/profiler_server.h"
+#include "xla/tsl/profiler/rpc/profiler_service_impl.h"
 #include "xla/tsl/profiler/utils/file_system_utils.h"
 #include "tsl/platform/host_info.h"
 #include "tsl/profiler/lib/profiler_session.h"
 #include "tsl/profiler/protobuf/profiler_options.pb.h"
+#include "tsl/profiler/protobuf/profiler_service.grpc.pb.h"
 #include "tsl/profiler/protobuf/profiler_service.pb.h"
 
 namespace tsl {
@@ -271,6 +276,70 @@ TEST(RemoteProfilerSessionManagerTest, OverrideHostnamesMismatch) {
   EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument);
 }
 
+// Returns true if the test environment supports binding an IPv6 loopback
+// socket ([::1]).
+bool IsIpv6LoopbackAvailable() {
+  std::unique_ptr<tensorflow::grpc::ProfilerService::Service> service =
+      CreateProfilerService();
+  ::grpc::ServerBuilder builder;
+  int selected_port = 0;
+  builder.AddListeningPort("[::1]:0", ::grpc::InsecureServerCredentials(),
+                           &selected_port);
+  builder.RegisterService(service.get());
+  std::unique_ptr<::grpc::Server> server = builder.BuildAndStart();
+  if (server != nullptr) {
+    server->Shutdown();
+  }
+  return selected_port != 0;
+}
+
+TEST(RemoteProfilerSessionManagerTest, Ipv6Address) {
+  if (!IsIpv6LoopbackAvailable()) {
+    GTEST_SKIP() << "IPv6 loopback [::1] is unavailable in this test "
+                    "environment.";
+  }
+  absl::Duration duration = absl::Milliseconds(30);
+  RemoteProfilerSessionManagerOptions options;
+  *options.mutable_profiler_options() = tsl::ProfilerSession::DefaultOptions();
+  options.mutable_profiler_options()->set_duration_ms(
+      absl::ToInt64Milliseconds(duration));
+
+  std::string service_address;
+  std::unique_ptr<ProfilerServer> server =
+      StartServer(duration, &service_address);
+  const size_t colon_pos = service_address.rfind(':');
+  ASSERT_NE(colon_pos, std::string::npos);
+  std::string ipv6_service_address =
+      absl::StrCat("[::1]", service_address.substr(colon_pos));
+  options.add_service_addresses(ipv6_service_address);
+
+  absl::Time approx_start = absl::Now();
+  absl::Duration grace = absl::Seconds(kGracePeriodSeconds);
+  absl::Duration max_duration = duration + grace;
+  options.set_max_session_duration_ms(absl::ToInt64Milliseconds(max_duration));
+  options.set_session_creation_timestamp_ns(absl::ToUnixNanos(approx_start));
+
+  ProfileRequest request = PopulateProfileRequest(
+      TmpDir(), "session_id", ipv6_service_address, options);
+  absl::Status status;
+  std::unique_ptr<RemoteProfilerSessionManager> sessions =
+      RemoteProfilerSessionManager::Create(options, request, status);
+  ASSERT_OK(status);
+  std::vector<Response> responses = sessions->WaitForCompletion();
+  absl::Duration elapsed = absl::Now() - approx_start;
+  ASSERT_EQ(responses.size(), 1);
+  EXPECT_EQ(responses.back().service_address, ipv6_service_address);
+  EXPECT_OK(responses.back().status);
+  EXPECT_TRUE(responses.back().profile_response->empty_trace());
+  EXPECT_EQ(responses.back().profile_response->tool_data_size(), 0);
+  EXPECT_THAT(elapsed, DurationApproxLess(max_duration));
+
+  std::string expected_filename = absl::StrCat(
+      "[__1]_", service_address.substr(colon_pos + 1), ".xplane.pb");
+  EXPECT_OK(Env::Default()->FileExists(ProfilerJoinPath(
+      request.repository_root(), request.session_id(), expected_filename)));
+}
+
 }  // namespace
 }  // namespace profiler
 }  // namespace tsl
diff --git a/third_party/xla/xla/tsl/profiler/utils/BUILD b/third_party/xla/xla/tsl/profiler/utils/BUILD
index 5b8060d5737..c6340f1d445 100644
--- a/third_party/xla/xla/tsl/profiler/utils/BUILD
+++ b/third_party/xla/xla/tsl/profiler/utils/BUILD
@@ -533,8 +533,11 @@ tsl_cc_test(
         ":session_manager",
         "//xla/tsl/platform:test",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest_main",
+        "@tsl//tsl/profiler/protobuf:profiler_options_proto_cc",
     ],
 )
 
diff --git a/third_party/xla/xla/tsl/profiler/utils/session_manager.cc b/third_party/xla/xla/tsl/profiler/utils/session_manager.cc
index f98dfbb71c9..77bef813bb3 100644
--- a/third_party/xla/xla/tsl/profiler/utils/session_manager.cc
+++ b/third_party/xla/xla/tsl/profiler/utils/session_manager.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <algorithm>
 #include <climits>
+#include <cstddef>
 #include <cstdint>
 #include <string>
 #include <variant>
@@ -308,12 +309,42 @@ absl::Status ValidateRemoteProfilerSessionManagerOptions(
 }
 
 absl::Status ValidateHostPortPair(absl::string_view host_port) {
-  uint32_t port;
-  std::vector<absl::string_view> parts = absl::StrSplit(host_port, ':');
-  // Must be host:port, port must be a number, host must not contain a '/',
-  // host also must not be empty.
-  if (parts.size() != 2 || !absl::SimpleAtoi(parts[1], &port) ||
-      absl::StrContains(parts[0], "/") || parts[0].empty()) {
+  absl::string_view host;
+  absl::string_view port_str;
+  if (absl::StartsWith(host_port, "[")) {
+    const size_t close_bracket = host_port.find(']');
+    if (close_bracket == absl::string_view::npos ||
+        close_bracket + 1 >= host_port.size() ||
+        host_port[close_bracket + 1] != ':') {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Could not interpret \"", host_port, "\" as a host-port pair."));
+    }
+    host = host_port.substr(1, close_bracket - 1);
+    port_str = host_port.substr(close_bracket + 2);
+    // Per RFC 3986 and gRPC's dns:/// resolver, bracketed hosts must be IPv6
+    // literals (contain ':') and must not contain nested brackets.
+    if (host.empty() || !absl::StrContains(host, ':') ||
+        absl::StrContains(host, '[')) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Could not interpret \"", host_port, "\" as a host-port pair."));
+    }
+  } else {
+    if (absl::StrContains(host_port, '[') ||
+        absl::StrContains(host_port, ']')) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Could not interpret \"", host_port, "\" as a host-port pair."));
+    }
+    std::vector<absl::string_view> parts = absl::StrSplit(host_port, ':');
+    if (parts.size() != 2 || parts[0].empty()) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Could not interpret \"", host_port, "\" as a host-port pair."));
+    }
+    host = parts[0];
+    port_str = parts[1];
+  }
+  uint32_t port = 0;
+  if (!absl::SimpleAtoi(port_str, &port) || port > 65535 ||
+      absl::StrContains(host, '/')) {
     return absl::InvalidArgumentError(absl::StrCat(
         "Could not interpret \"", host_port, "\" as a host-port pair."));
   }
diff --git a/third_party/xla/xla/tsl/profiler/utils/session_manager_test.cc b/third_party/xla/xla/tsl/profiler/utils/session_manager_test.cc
index 5347ffbe9ab..053cb5b43ec 100644
--- a/third_party/xla/xla/tsl/profiler/utils/session_manager_test.cc
+++ b/third_party/xla/xla/tsl/profiler/utils/session_manager_test.cc
@@ -18,15 +18,24 @@ limitations under the License.
 #include <string>
 #include <variant>
 
+#include <gmock/gmock.h>
 #include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
 #include "absl/strings/string_view.h"
 #include "xla/tsl/platform/test.h"
+#include "tsl/profiler/protobuf/profiler_options.pb.h"
 
 namespace tsl {
 namespace profiler {
 namespace {
 
-using tensorflow::RemoteProfilerSessionManagerOptions;
+using ::absl_testing::IsOk;
+using ::absl_testing::StatusIs;
+using ::tensorflow::RemoteProfilerSessionManagerOptions;
+using ::testing::ElementsAre;
+using ::testing::Eq;
+using ::testing::HasSubstr;
 
 TEST(SessionManagerTest, OptionsWithSessionIdTest) {
   absl::string_view logdir = "/tmp/logdir";
@@ -84,6 +93,117 @@ TEST(SessionManagerTest, UseSystemHostnameFalseTest) {
   EXPECT_FALSE(it->second.bool_value());
 }
 
+TEST(SessionManagerTest, ValidateHostPortPairValidAddresses) {
+  EXPECT_THAT(ValidateHostPortPair("localhost:8466"), IsOk());
+  EXPECT_THAT(ValidateHostPortPair("127.0.0.1:8466"), IsOk());
+  EXPECT_THAT(ValidateHostPortPair("[::1]:8466"), IsOk());
+  EXPECT_THAT(ValidateHostPortPair("[::]:8466"), IsOk());
+  EXPECT_THAT(ValidateHostPortPair("[2001:db8:21da:7::b]:8466"), IsOk());
+  EXPECT_THAT(ValidateHostPortPair("[2001:db8:85a3::8a2e:370:7334]:8466"),
+              IsOk());
+  EXPECT_THAT(
+      ValidateHostPortPair("[2001:0db8:85a3:0000:0000:8a2e:0370:7334]:8466"),
+      IsOk());
+  EXPECT_THAT(ValidateHostPortPair("[::ffff:192.0.2.128]:8466"), IsOk());
+  EXPECT_THAT(ValidateHostPortPair("[fe80::1%eth0]:8466"), IsOk());
+}
+
+TEST(SessionManagerTest, ValidateHostPortPairInvalidUnbracketedIpv6) {
+  EXPECT_THAT(ValidateHostPortPair("2001:db8::1:8466"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("::1:8466"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+}
+
+TEST(SessionManagerTest, ValidateHostPortPairInvalidBracketAndPortForms) {
+  EXPECT_THAT(ValidateHostPortPair("[2001:db8::1]"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("[2001:db8::1]:"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("[2001:db8::1]:abc"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("[2001:db8::1]:-1"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("[2001:db8::1]:70000"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("[2001:db8::1]8466"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("[]"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("[]:8466"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("[localhost]:8466"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("[127.0.0.1]:8466"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("[[::1]:8466"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("[::1]]:8466"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("[[::1]]:8466"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("host[0]:8466"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("host]:8466"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("[2001:db8::1/64]:8466"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("host/path:8466"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("localhost"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("8466"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair(":8466"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair("localhost:"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+  EXPECT_THAT(ValidateHostPortPair(""),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("Could not interpret")));
+}
+
+TEST(SessionManagerTest, MultiHostIpv6OptionsValidation) {
+  absl::string_view service_addresses = "[2001:db8::1]:8466,[2001:db8::2]:8466";
+  absl::string_view logdir = "/tmp/logdir";
+  absl::flat_hash_map<std::string, std::variant<bool, int, std::string>> opts;
+  bool is_cloud_tpu_session = false;
+
+  RemoteProfilerSessionManagerOptions options =
+      GetRemoteSessionManagerOptionsLocked(service_addresses, logdir,
+                                           /*worker_list=*/"",
+                                           /*include_dataset_ops=*/false,
+                                           /*duration_ms=*/100, opts,
+                                           &is_cloud_tpu_session);
+  EXPECT_FALSE(is_cloud_tpu_session);
+  EXPECT_THAT(options.service_addresses(),
+              ElementsAre(Eq("[2001:db8::1]:8466"), Eq("[2001:db8::2]:8466")));
+  EXPECT_THAT(ValidateRemoteProfilerSessionManagerOptions(options), IsOk());
+}
+
 }  // namespace
 }  // namespace profiler
 }  // namespace tsl

From 1ec71d73bf94513f5842cfec4b11343c554c625d Mon Sep 17 00:00:00 2001
From: abhirupone <abhirup.karmakar00@gmail.com>
Date: Tue, 13 Oct 2026 19:56:00 +0530
Subject: [PATCH 0675/1000] Address review feedback on RaggedTensor static
 validation and monotonicity checks

- Replace np.diff with direct array comparison in RowPartition monotonicity checks
- Reuse const_rowids in RowPartition.from_value_rowids instead of querying constant_value twice
- Use np.int64 accumulation in np.sum for _static_nvals_or_constant
- Add comment explaining validate=False usage in ragged_squeeze_op
- Add jit_compile test cases for dynamic first dim and static validation of row_starts
---
 .../python/ops/ragged/ragged_squeeze_op.py    |  3 ++
 .../python/ops/ragged/ragged_tensor_test.py   | 23 +++++++++++
 tensorflow/python/ops/ragged/row_partition.py | 39 +++++++++----------
 3 files changed, 44 insertions(+), 21 deletions(-)

diff --git a/tensorflow/python/ops/ragged/ragged_squeeze_op.py b/tensorflow/python/ops/ragged/ragged_squeeze_op.py
index 8755bee8b5a..afd2256892a 100644
--- a/tensorflow/python/ops/ragged/ragged_squeeze_op.py
+++ b/tensorflow/python/ops/ragged/ragged_squeeze_op.py
@@ -112,6 +112,9 @@ def squeeze(input: ragged_tensor.Ragged, axis=None, name=None):  # pylint: disab
     if remaining_row_splits and 0 in ragged_dims:
       remaining_row_splits.pop(0)
 
+    # Pass validate=False because remaining_row_splits come from an
+    # already-valid input RaggedTensor and are gated by the squeezability
+    # Assert checks above.
     squeezed_rt = RaggedTensor.from_nested_row_splits(
         squeezed_rt, remaining_row_splits, validate=False)
 
diff --git a/tensorflow/python/ops/ragged/ragged_tensor_test.py b/tensorflow/python/ops/ragged/ragged_tensor_test.py
index 6d87498243e..07a587a3b1d 100644
--- a/tensorflow/python/ops/ragged/ragged_tensor_test.py
+++ b/tensorflow/python/ops/ragged/ragged_tensor_test.py
@@ -29,6 +29,7 @@ from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import tensor as tensor_lib
 from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_spec
 from tensorflow.python.framework import test_util
 from tensorflow.python.framework import type_spec
 from tensorflow.python.framework.type_utils import fulltypes_for_flat_tensors
@@ -370,6 +371,28 @@ class RaggedTensorTest(test_util.TensorFlowTestCase, parameterized.TestCase):
         ValueError, r"Shapes \(2,\) and \(4,\) are incompatible"):
       fn(constant_op.constant([4.0, 5.0, 6.0, 7.0]))
 
+  def testFromRowSplitsDynamicFirstDimUnderJit(self):
+    @def_function.function(
+        jit_compile=True,
+        input_signature=[
+            tensor_spec.TensorSpec(shape=[None], dtype=dtypes.float32)
+        ])
+    def fn(x):
+      return RaggedTensor.from_row_splits(x, [0, 2, 3])
+
+    result = fn(constant_op.constant([1.0, 2.0, 3.0]))
+    self.assertAllEqual(result, [[1.0, 2.0], [3.0]])
+
+  def testFromRowStartsStaticValidationUnderJit(self):
+    @def_function.function(jit_compile=True)
+    def fn(x):
+      values = array_ops.reshape(x, [-1, 2])
+      return RaggedTensor.from_row_starts(values, [0, 3, 2])
+
+    with self.assertRaisesRegex(
+        ValueError, r"row_starts must be monotonic increasing"):
+      fn(constant_op.constant([4.0, 5.0, 6.0, 7.0]))
+
   def testFromRowStarts(self):
     values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
     row_starts = constant_op.constant([0, 2, 2, 5, 6], dtypes.int64)
diff --git a/tensorflow/python/ops/ragged/row_partition.py b/tensorflow/python/ops/ragged/row_partition.py
index 89b3a33c8da..716032f5ba5 100644
--- a/tensorflow/python/ops/ragged/row_partition.py
+++ b/tensorflow/python/ops/ragged/row_partition.py
@@ -251,8 +251,8 @@ class RowPartition(composite_tensor.CompositeTensor):
                         [value_rowids, nrows]):
       value_rowids = cls._convert_row_partition(
           value_rowids, "value_rowids", dtype_hint=dtype_hint, dtype=dtype)
+      const_rowids = tensor_util.constant_value(value_rowids)
       if nrows is None:
-        const_rowids = tensor_util.constant_value(value_rowids)
         if const_rowids is None:
           nrows = array_ops.concat([value_rowids[-1:], [-1]], axis=0)[0] + 1
           const_nrows = None
@@ -266,7 +266,6 @@ class RowPartition(composite_tensor.CompositeTensor):
         if const_nrows is not None:
           if const_nrows < 0:
             raise ValueError("Expected nrows >= 0; got %d" % const_nrows)
-          const_rowids = tensor_util.constant_value(value_rowids)
           if const_rowids is not None and const_rowids.size > 0:
             if not const_nrows >= const_rowids[-1] + 1:
               raise ValueError(
@@ -277,23 +276,21 @@ class RowPartition(composite_tensor.CompositeTensor):
       nrows.shape.assert_has_rank(0)
 
       if validate:
-        value_rowids_const = tensor_util.constant_value(value_rowids)
-        if value_rowids_const is not None:
-          if len(value_rowids_const) > 0 and value_rowids_const[0] < 0:
+        if const_rowids is not None:
+          if len(const_rowids) > 0 and const_rowids[0] < 0:
             raise ValueError(
                 "Arguments to from_value_rowids do not form a valid "
                 "RowPartition: value_rowids must be non-negative, "
-                f"got {value_rowids_const[0]}")
-          diff = np.diff(value_rowids_const)
-          if np.any(diff < 0):
-            bad_idx = int(np.argmax(diff < 0))
+                f"got {const_rowids[0]}")
+          if np.any(const_rowids[1:] < const_rowids[:-1]):
+            bad_idx = int(np.argmax(const_rowids[1:] < const_rowids[:-1]))
             raise ValueError(
                 "Arguments to from_value_rowids do not form a valid "
                 "RowPartition: value_rowids must be monotonic increasing, "
                 f"but value_rowids[{bad_idx}] "
-                f"({value_rowids_const[bad_idx]}) > "
+                f"({const_rowids[bad_idx]}) > "
                 f"value_rowids[{bad_idx + 1}] "
-                f"({value_rowids_const[bad_idx + 1]})")
+                f"({const_rowids[bad_idx + 1]})")
         msg = ("Arguments to from_value_rowids do not form a valid "
                "RowPartition")
         checks = [
@@ -390,9 +387,9 @@ class RowPartition(composite_tensor.CompositeTensor):
                 "Arguments to from_row_splits do not form a valid "
                 f"RaggedTensor: row_splits[0] must be zero, "
                 f"got {row_splits_const[0]}")
-          diff = np.diff(row_splits_const)
-          if np.any(diff < 0):
-            bad_idx = int(np.argmax(diff < 0))
+          if np.any(row_splits_const[1:] < row_splits_const[:-1]):
+            bad_idx = int(
+                np.argmax(row_splits_const[1:] < row_splits_const[:-1]))
             raise ValueError(
                 "Arguments to from_row_splits do not form a valid "
                 "RaggedTensor: row_splits must be monotonic increasing, "
@@ -517,9 +514,9 @@ class RowPartition(composite_tensor.CompositeTensor):
                 "Arguments to from_row_starts do not form a valid "
                 f"RaggedTensor: row_starts[0] must be zero, "
                 f"got {row_starts_const[0]}")
-          diff = np.diff(row_starts_const)
-          if np.any(diff < 0):
-            bad_idx = int(np.argmax(diff < 0))
+          if np.any(row_starts_const[1:] < row_starts_const[:-1]):
+            bad_idx = int(
+                np.argmax(row_starts_const[1:] < row_starts_const[:-1]))
             raise ValueError(
                 "Arguments to from_row_starts do not form a valid "
                 "RaggedTensor: row_starts must be monotonic increasing, "
@@ -591,9 +588,9 @@ class RowPartition(composite_tensor.CompositeTensor):
                 "Arguments to from_row_limits do not form a valid "
                 "RaggedTensor: row_limits must be non-negative, "
                 f"got {row_limits_const[0]}")
-          diff = np.diff(row_limits_const)
-          if np.any(diff < 0):
-            bad_idx = int(np.argmax(diff < 0))
+          if np.any(row_limits_const[1:] < row_limits_const[:-1]):
+            bad_idx = int(
+                np.argmax(row_limits_const[1:] < row_limits_const[:-1]))
             raise ValueError(
                 "Arguments to from_row_limits do not form a valid "
                 "RaggedTensor: row_limits must be monotonic increasing, "
@@ -990,7 +987,7 @@ class RowPartition(composite_tensor.CompositeTensor):
     if self._row_lengths is not None:
       row_lengths_const = tensor_util.constant_value(self._row_lengths)
       if row_lengths_const is not None:
-        return int(row_lengths_const.sum())
+        return int(np.sum(row_lengths_const, dtype=np.int64))
     if self._uniform_row_length is not None and self._nrows is not None:
       nrows = self.static_nrows
       row_len = self.static_uniform_row_length

From b102894ee56a55de3fb48e78f302b5af7c726ec7 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Tue, 13 Oct 2026 03:26:20 -0700
Subject: [PATCH 0676/1000] [XLA:GPU]: Scheduling for variadics and xtile
 lowering.

PiperOrigin-RevId: 998516116
---
 .../xla/xla/codegen/tiling/experimental/BUILD |   4 +
 .../codegen/tiling/experimental/scheduling.cc | 297 +++++++++++++++++-
 .../tiling/experimental/scheduling_test.cc    | 285 +++++++++++++++++
 .../codegen/xtile/codegen/emitter_helpers.h   |  36 ++-
 .../codegen/experimental_fusion_emitter.cc    |  75 +++--
 5 files changed, 662 insertions(+), 35 deletions(-)

diff --git a/third_party/xla/xla/codegen/tiling/experimental/BUILD b/third_party/xla/xla/codegen/tiling/experimental/BUILD
index 321028eed61..df27fcac371 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/BUILD
+++ b/third_party/xla/xla/codegen/tiling/experimental/BUILD
@@ -218,6 +218,7 @@ cc_library(
         ":tiled_hlo",
         "//xla:permutation_util",
         "//xla:shape_util",
+        "//xla:status_macros",
         "//xla:util",
         "//xla:xla_data_proto_cc",
         "//xla/hlo/analysis:indexing_analysis",
@@ -227,6 +228,8 @@ cc_library(
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
@@ -250,6 +253,7 @@ xla_cc_test(
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
         "//xla/hlo/utils:hlo_traversal",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/strings:string_view",
diff --git a/third_party/xla/xla/codegen/tiling/experimental/scheduling.cc b/third_party/xla/xla/codegen/tiling/experimental/scheduling.cc
index 1a096bf75df..e9d8af7bfd9 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/scheduling.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/scheduling.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include <algorithm>
 #include <cstdint>
 #include <numeric>
+#include <optional>
 #include <set>
 #include <string>
 #include <utility>
@@ -26,12 +27,20 @@ limitations under the License.
 #include "absl/algorithm/container.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/str_join.h"
 #include "absl/types/span.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SmallVector.h"
 #include "mlir/IR/MLIRContext.h"
+#include "xla/codegen/tiling/experimental/tile.h"
 #include "xla/codegen/tiling/experimental/tiled_hlo.h"
 #include "xla/codegen/tiling/experimental/tiling_space.h"
 #include "xla/hlo/analysis/indexing_analysis.h"
@@ -39,11 +48,13 @@ limitations under the License.
 #include "xla/hlo/analysis/symbolic_expr.h"
 #include "xla/hlo/analysis/symbolic_map.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
+#include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/permutation_util.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
+#include "xla/status_macros.h"
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
 
@@ -270,6 +281,271 @@ llvm::SmallVector<int64_t> GetParallelDimensionsPermutation(
   return {};
 }
 
+// What the scheduler needs to know about the index-wise variadic instructions
+// of a computation.
+struct IndexWiseVariadicInfo {
+  // The index-wise variadic instructions, in any region of the computation.
+  llvm::SmallVector<const TiledHloInstruction*, 2> instructions;
+  // Maps each all-gather dimension ID to how many tiles each peer owns along
+  // that dimension.
+  //
+  // Example: `f32[4, 16] -> f32[16, 16] all-gather` across 4 GPUs on dim 0
+  // (`d0`) with tile size [2, 16]:
+  //   - `d0` has 16 / 2 = 8 tiles in total across all 4 GPUs.
+  //   - Each GPU owns a 4-row slice, which is 4 / 2 = 2 tiles along `d0`.
+  //   - So `tiles_per_peer_by_dim[d0] = 2`.
+  llvm::DenseMap<int64_t, int64_t> tiles_per_peer_by_dim;
+  // Number of peers the outputs are read from. 1 if `tiles_per_peer_by_dim`
+  // is empty.
+  int64_t num_peers = 1;
+};
+
+// Collects the index-wise variadic instructions of `region` and its nested
+// regions.
+void CollectIndexWiseVariadic(
+    const TiledHloRegion& region,
+    llvm::SmallVectorImpl<const TiledHloInstruction*>& result) {
+  for (const TiledHloInstruction* instruction : region.instructions()) {
+    if (IsIndexWiseVariadic(*instruction->hlo())) {
+      result.push_back(instruction);
+    }
+    for (const TiledHloRegion& nested : instruction->hlo_regions()) {
+      CollectIndexWiseVariadic(nested, result);
+    }
+  }
+}
+
+// Finds the index-wise variadic instructions of the computation and, for the
+// operand tiles that carry a replica id, derives which parallel tile dimension
+// carries the peer index and how many tiles per peer it spans.
+absl::StatusOr<IndexWiseVariadicInfo> BuildIndexWiseVariadicInfo(
+    const TiledHloComputation& tiled_computation) {
+  IndexWiseVariadicInfo info;
+  CollectIndexWiseVariadic(tiled_computation.tiled_root_region(),
+                           info.instructions);
+  if (info.instructions.empty()) {
+    return info;
+  }
+
+  const TilingSpace& tiling_space = tiled_computation.tiling_space();
+  TF_RET_CHECK(!tiling_space.IsSymbolic())
+      << "Expected tile sizes to be assigned at this point.";
+  llvm::SmallVector<TilingSpace::DimensionInfo, 4> dims =
+      tiling_space.dimensions();
+  const int64_t num_dims = dims.size();
+
+  llvm::SmallVector<int64_t, 4> tile_sizes;
+  tile_sizes.reserve(num_dims);
+  for (const TilingSpace::DimensionInfo& dim : dims) {
+    tile_sizes.push_back(*dim.tile_size);
+  }
+  // Variables are laid out as [tile ids, tile sizes, runtime variables].
+  const int64_t num_vars = 2 * num_dims + tiling_space.num_rt_vars();
+
+  std::optional<int64_t> num_peers;
+  for (const TiledHloInstruction* instruction : info.instructions) {
+    const HloInstruction& hlo = *instruction->hlo();
+    for (const auto& [output_index, operand] :
+         llvm::enumerate(instruction->operands())) {
+      const Tile& tile = operand->tile();
+      if (tile.num_replica_ids() == 0) {
+        // The output tile is not associated with a single peer; nothing to
+        // factor out of its tile ids.
+        continue;
+      }
+      TF_RET_CHECK(tile.num_replica_ids() == 1)
+          << "Only one replica dimension is supported.";
+      const DimTile& replica_dim_tile = tile.replica_ids().back();
+      const int64_t num_replicas = replica_dim_tile.upper_bound.GetValue();
+      SymbolicExpr peer = replica_dim_tile.offset;
+      llvm::DenseSet<VariableID> used_vars;
+      peer.GetUsedVariables(used_vars);
+      if (used_vars.size() != 1 || *used_vars.begin() >= num_dims ||
+          dims[*used_vars.begin()].type !=
+              TilingSpace::DimensionSemantics::kParallel) {
+        return absl::UnimplementedError(absl::StrFormat(
+            "Peer index %s of output %d of %s must depend on exactly one "
+            "parallel tile id.",
+            peer.ToString(num_dims), output_index, hlo.name()));
+      }
+      const int64_t dim_id = *used_vars.begin();
+      // Number of blocks across all replicas.
+      const int64_t num_blocks =
+          CeilOfRatio(dims[dim_id].dimension_size, tile_sizes[dim_id]);
+      if (num_blocks < num_replicas || num_blocks % num_replicas != 0) {
+        return absl::UnimplementedError(absl::StrFormat(
+            "Dimension %d of output %d of %s has %d blocks, which is not a "
+            "multiple of the number of peers %d.",
+            dim_id, output_index, hlo.name(), num_blocks, num_replicas));
+      }
+      const int64_t tiles_per_peer = num_blocks / num_replicas;
+      std::vector<int64_t> values(num_vars, 0);
+      auto peer_at = [&](int64_t tile_id) {
+        values[dim_id] = tile_id;
+        return peer.Evaluate(values);
+      };
+      if (!llvm::all_of(llvm::seq<int64_t>(0, num_blocks), [&](int64_t t) {
+            return peer_at(t) == t / tiles_per_peer;
+          })) {
+        return absl::UnimplementedError(absl::StrFormat(
+            "Peer index %s of output %d of %s must be tile id %d divided by "
+            "%d.",
+            peer.ToString(num_dims), output_index, hlo.name(), dim_id,
+            tiles_per_peer));
+      }
+
+      if (num_peers.has_value() && *num_peers != num_replicas) {
+        return absl::InvalidArgumentError(absl::StrFormat(
+            "Output %d of %s reads from %d peers, but other outputs read "
+            "from %d peers.",
+            output_index, hlo.name(), num_replicas, *num_peers));
+      }
+      num_peers = num_replicas;
+      info.tiles_per_peer_by_dim[dim_id] = tiles_per_peer;
+    }
+  }
+  info.num_peers = num_peers.value_or(1);
+  return info;
+}
+
+// Internal struct for holding the parallel dimensions of a root.
+struct RootDims {
+  // Parallel output dimension whose tile ID encodes the peer index
+  // (all-gather).
+  std::optional<int64_t> peer_dim_id;
+  int64_t tiles_per_peer = 1;
+  llvm::SmallVector<int64_t, 4> non_peer_dim_ids;
+  llvm::SmallVector<int64_t, 4> non_peer_block_counts;
+  int64_t num_tiles = 1;
+};
+
+absl::StatusOr<llvm::SmallVector<RootDims, 2>> BuildRootsForIndexWiseVariadic(
+    const TiledHloComputation& tiled_computation,
+    const IndexWiseVariadicInfo& info) {
+  const TilingSpace& tiling_space = tiled_computation.tiling_space();
+  const bool has_peer_dim = !info.tiles_per_peer_by_dim.empty();
+  // Group the parallel dimensions by the root (output) that defines them.
+  llvm::SmallVector<RootDims, 2> roots;
+  int64_t num_grouped_dims = 0;
+  for (const auto& [root_index, tiled_root] :
+       llvm::enumerate(tiled_computation.roots())) {
+    const HloInstruction& hlo = *tiled_root->hlo();
+    RootDims& root = roots.emplace_back();
+    for (int64_t i = 0; i < GetFirstShape(&hlo).dimensions().size(); ++i) {
+      const TilingSpace::DimensionInfo* dim =
+          &tiling_space.GetDimensionInfo(hlo, i);
+      if (dim->type != TilingSpace::DimensionSemantics::kParallel) {
+        continue;
+      }
+      ++num_grouped_dims;
+      const int64_t dim_id = dim->id.value();
+      auto it = info.tiles_per_peer_by_dim.find(dim_id);
+      if (it == info.tiles_per_peer_by_dim.end()) {
+        root.non_peer_dim_ids.push_back(dim_id);
+        root.non_peer_block_counts.push_back(
+            CeilOfRatio(dim->dimension_size, *dim->tile_size));
+        continue;
+      }
+      if (root.peer_dim_id.has_value()) {
+        return absl::UnimplementedError(absl::StrFormat(
+            "Output %d of %s has more than one dimension carrying the peer "
+            "index.",
+            root_index, hlo.name()));
+      }
+      root.peer_dim_id = dim_id;
+      root.tiles_per_peer = it->second;
+    }
+    if (has_peer_dim && !root.peer_dim_id.has_value()) {
+      return absl::UnimplementedError(
+          absl::StrFormat("Output %d of %s does not depend on the peer index.",
+                          root_index, hlo.name()));
+    }
+    root.num_tiles = Product(root.non_peer_block_counts) * root.tiles_per_peer;
+  }
+  if (num_grouped_dims != tiling_space.num_parallel_dimensions()) {
+    return absl::UnimplementedError(
+        "Not all parallel dimensions are defined by the roots.");
+  }
+  return roots;
+}
+
+absl::StatusOr<Schedule> BuildIndexWiseVariadicSchedule(
+    mlir::MLIRContext* ctx, const IndexWiseVariadicInfo& info,
+    llvm::ArrayRef<RootDims> roots, int64_t num_tiles_per_pid,
+    int64_t max_tiles, bool has_peer_dim, SymbolicExpr global_tile_id) {
+  Schedule schedule;
+  SymbolicExpr peer = global_tile_id.floorDiv(max_tiles);
+  SymbolicExpr intra_peer_tile_id =
+      has_peer_dim ? global_tile_id % max_tiles : global_tile_id;
+  for (const RootDims& root : roots) {
+    llvm::SmallVector<int64_t, 4> block_counts = root.non_peer_block_counts;
+    if (root.tiles_per_peer > 1) {
+      block_counts.push_back(root.tiles_per_peer);
+    }
+    // `DelinearizeIndex` does not wrap its outermost component, so for a padded
+    // tile (`intra_peer_tile_id >= num_tiles`) that component exceeds its
+    // block count and the tile is masked by its upper bound.
+    llvm::SmallVector<SymbolicExpr, 4> delinearized =
+        DelinearizeIndex(block_counts, intra_peer_tile_id, ctx);
+    for (int64_t i = 0; i < root.non_peer_dim_ids.size(); ++i) {
+      schedule.dim_id_to_pid_expr[root.non_peer_dim_ids[i]] = delinearized[i];
+    }
+    if (!root.peer_dim_id.has_value()) {
+      continue;
+    }
+    SymbolicExpr peer_tile_id = peer * root.tiles_per_peer;
+    if (root.tiles_per_peer > 1) {
+      peer_tile_id = peer_tile_id + delinearized.back();
+    }
+    if (root.num_tiles < max_tiles &&
+        (root.non_peer_block_counts.empty() ||
+         root.non_peer_block_counts.front() == 1)) {
+      // No non-trivial non-peer dimension can absorb the padding: a non-peer
+      // dimension covered by a single block folds its tile offset to 0, and
+      // an unpadded `tiles_per_peer` offset would step into the next peer
+      // instead of out of bounds. Push the peer tile id past the last peer.
+      peer_tile_id =
+          peer_tile_id + intra_peer_tile_id.floorDiv(root.num_tiles) *
+                             (info.num_peers * root.tiles_per_peer);
+    }
+    schedule.dim_id_to_pid_expr[*root.peer_dim_id] = peer_tile_id;
+  }
+  schedule.num_tiles = info.num_peers * max_tiles;
+  schedule.num_pids = CeilOfRatio(schedule.num_tiles, num_tiles_per_pid);
+  return schedule;
+}
+
+// Schedules a computation that contains index-wise variadic instructions.
+//
+// All roots are "zipped": one program handles tile `i` of every root, so
+// that each variadic instruction is executed once per program for all of
+// its outputs. If the outputs are associated with a peer the peer index is
+// the outermost tile index and is shared by all roots; below it every root
+// gets its own intra-peer tile index over its remaining ("non-peer")
+// parallel dimensions and any intra-peer tiles along the peer dimension
+// itself. Roots with fewer tiles than the largest root are padded: their
+// padded tiles are mapped out of bounds so that the emitter masks them.
+absl::StatusOr<Schedule> GetIndexWiseVariadicSchedule(
+    const TiledHloComputation& tiled_computation,
+    const IndexWiseVariadicInfo& info, SymbolicExpr global_tile_id,
+    int64_t num_tiles_per_pid) {
+  mlir::MLIRContext* ctx = tiled_computation.GetMLIRContext();
+  const bool has_peer_dim = !info.tiles_per_peer_by_dim.empty();
+  using Roots = llvm::SmallVector<RootDims, 2>;
+  ABSL_ASSIGN_OR_RETURN(Roots roots,
+                   BuildRootsForIndexWiseVariadic(tiled_computation, info));
+
+  // Find the maximum number of tiles across all roots.
+  int64_t max_tiles = 1;
+  for (const RootDims& root : roots) {
+    max_tiles = std::max(max_tiles, root.num_tiles);
+  }
+  // Now build the scheduled rooted on the output with the max number of tiles.
+  return BuildIndexWiseVariadicSchedule(ctx, info, roots, num_tiles_per_pid,
+                                        max_tiles, has_peer_dim,
+                                        global_tile_id);
+}
+
 }  // namespace
 
 std::string Schedule::ToString() const {
@@ -290,6 +566,20 @@ std::string Schedule::ToString() const {
 
 absl::StatusOr<Schedule> GetSchedule(
     const TiledHloComputation& tiled_computation, int64_t num_tiles_per_pid) {
+  mlir::MLIRContext* ctx = tiled_computation.GetMLIRContext();
+  SymbolicExpr program_id = CreateDimExpr(0, ctx);
+  SymbolicExpr global_tile_id = program_id;
+  if (num_tiles_per_pid > 1) {
+    SymbolicExpr tile_id = CreateDimExpr(1, ctx);
+    global_tile_id = program_id * num_tiles_per_pid + tile_id;
+  }
+  ABSL_ASSIGN_OR_RETURN(IndexWiseVariadicInfo variadic_info,
+                   BuildIndexWiseVariadicInfo(tiled_computation));
+  if (!variadic_info.instructions.empty()) {
+    return GetIndexWiseVariadicSchedule(tiled_computation, variadic_info,
+                                        global_tile_id, num_tiles_per_pid);
+  }
+
   // Compute the block counts for each parallel dimension.
   llvm::SmallVector<int64_t, 4> parallel_dim_block_counts;
   llvm::SmallVector<int64_t, 4> parallel_dim_ids;
@@ -312,13 +602,6 @@ absl::StatusOr<Schedule> GetSchedule(
         llvm::to_vector<4>(xla::Permute(parallel_dim_ids, permutation));
   }
 
-  mlir::MLIRContext* ctx = tiled_computation.GetMLIRContext();
-  SymbolicExpr program_id = CreateDimExpr(0, ctx);
-  SymbolicExpr global_tile_id = program_id;
-  if (num_tiles_per_pid > 1) {
-    SymbolicExpr tile_id = CreateDimExpr(1, ctx);
-    global_tile_id = program_id * num_tiles_per_pid + tile_id;
-  }
   llvm::SmallVector<SymbolicExpr, 4> delinearized_pid =
       DelinearizeIndex(parallel_dim_block_counts, global_tile_id, ctx);
   Schedule schedule;
diff --git a/third_party/xla/xla/codegen/tiling/experimental/scheduling_test.cc b/third_party/xla/xla/codegen/tiling/experimental/scheduling_test.cc
index 58c4821731c..238f0367d86 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/scheduling_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/scheduling_test.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/log/check.h"
+#include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/status_matchers.h"
 #include "absl/strings/string_view.h"
@@ -41,7 +42,9 @@ namespace xla::gpu::experimental {
 namespace {
 
 using absl_testing::IsOkAndHolds;
+using absl_testing::StatusIs;
 using mlir::MLIRContext;
+using testing::HasSubstr;
 
 }  // namespace
 
@@ -271,4 +274,286 @@ TEST_F(SchedulingTest, GetRaggedDotNonContractingPermutationNoSwapWhenMGeN) {
           "d0 -> pid / 2, d1 -> pid mod 2, num_pids=32, num_tiles=32")));
 }
 
+// d0 and d4 share the peer index (pid / 4); output 1 is padded on d5.
+TEST_F(SchedulingTest, VariadicAllGatherSharesPeerIndex) {
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"(
+    f {
+      p0 = f32[1,2,8,16] parameter(0)
+      p1 = s32[1,2] parameter(1)
+      ag = (f32[4,2,8,16], s32[4,2]) all-gather(p0, p1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      gte0 = f32[4,2,8,16] get-tuple-element(ag), index=0
+      gte1 = s32[4,2] get-tuple-element(ag), index=1
+      ROOT t = (f32[4,2,8,16], s32[4,2]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      arg0 = f32[1,2,8,16] parameter(0)
+      arg1 = s32[1,2] parameter(1)
+      ROOT fusion = (f32[4,2,8,16], s32[4,2]) fusion(arg0, arg1),
+        kind=kCustom, calls=f
+    })",
+                                    {1, 1, 4, 16, 1, 1}));
+  EXPECT_THAT(
+      GetSchedule(tiled_computation),
+      IsOkAndHolds(MatchSchedule("d0 -> pid / 4, d1 -> (pid mod 4) / 2, "
+                                 "d2 -> pid mod 4 mod 2, d3 -> 0, "
+                                 "d4 -> pid / 4, d5 -> pid mod 4, "
+                                 "num_pids=16, num_tiles=16")));
+}
+
+TEST_F(SchedulingTest, VariadicAllGatherNonZeroGatherDim) {
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"(
+    f {
+      p0 = f32[2,1,16] parameter(0)
+      p1 = s32[2,1] parameter(1)
+      ag = (f32[2,4,16], s32[2,4]) all-gather(p0, p1),
+        replica_groups={{0,1,2,3}}, dimensions={1}
+      gte0 = f32[2,4,16] get-tuple-element(ag), index=0
+      gte1 = s32[2,4] get-tuple-element(ag), index=1
+      ROOT t = (f32[2,4,16], s32[2,4]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      arg0 = f32[2,1,16] parameter(0)
+      arg1 = s32[2,1] parameter(1)
+      ROOT fusion = (f32[2,4,16], s32[2,4]) fusion(arg0, arg1),
+        kind=kCustom, calls=f
+    })",
+                                    {1, 1, 8, 1, 1}));
+  EXPECT_THAT(
+      GetSchedule(tiled_computation),
+      IsOkAndHolds(MatchSchedule("d0 -> (pid mod 4) / 2, d1 -> pid / 4, "
+                                 "d2 -> pid mod 4 mod 2, d3 -> pid mod 4, "
+                                 "d4 -> pid / 4, "
+                                 "num_pids=16, num_tiles=16")));
+}
+
+// Output 1 has no private dims; padded tiles push d2 past the last peer.
+TEST_F(SchedulingTest, VariadicAllGatherOutputWithoutPrivateDimsIsPadded) {
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"(
+    f {
+      p0 = s32[1,2] parameter(0)
+      p1 = s32[1] parameter(1)
+      ag = (s32[4,2], s32[4]) all-gather(p0, p1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      gte0 = s32[4,2] get-tuple-element(ag), index=0
+      gte1 = s32[4] get-tuple-element(ag), index=1
+      ROOT t = (s32[4,2], s32[4]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      arg0 = s32[1,2] parameter(0)
+      arg1 = s32[1] parameter(1)
+      ROOT fusion = (s32[4,2], s32[4]) fusion(arg0, arg1), kind=kCustom,
+        calls=f
+    })",
+                                    {1, 1, 1}));
+  EXPECT_THAT(GetSchedule(tiled_computation),
+              IsOkAndHolds(MatchSchedule("d0 -> pid / 2, d1 -> pid mod 2, "
+                                         "d2 -> pid / 2 + (pid mod 2) * 4, "
+                                         "num_pids=8, num_tiles=8")));
+}
+
+// Output 1's single-tile private dim (d3) cannot mask; d2 is pushed out of
+// bounds.
+TEST_F(SchedulingTest, VariadicAllGatherOutputWithSingleTileIsPadded) {
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"(
+    f {
+      p0 = s32[1,2] parameter(0)
+      p1 = s32[1,1] parameter(1)
+      ag = (s32[4,2], s32[4,1]) all-gather(p0, p1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      gte0 = s32[4,2] get-tuple-element(ag), index=0
+      gte1 = s32[4,1] get-tuple-element(ag), index=1
+      ROOT t = (s32[4,2], s32[4,1]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      arg0 = s32[1,2] parameter(0)
+      arg1 = s32[1,1] parameter(1)
+      ROOT fusion = (s32[4,2], s32[4,1]) fusion(arg0, arg1), kind=kCustom,
+        calls=f
+    })",
+                                    {1, 1, 1, 1}));
+  EXPECT_THAT(GetSchedule(tiled_computation),
+              IsOkAndHolds(MatchSchedule("d0 -> pid / 2, d1 -> pid mod 2, "
+                                         "d2 -> pid / 2 + (pid mod 2) * 4, "
+                                         "d3 -> pid mod 2, "
+                                         "num_pids=8, num_tiles=8")));
+}
+
+// Six outputs with 2..10 tiles/peer share peer index `pid / 10`.
+TEST_F(SchedulingTest, VariadicAllGatherWithMixedTileCounts) {
+  ASSERT_OK_AND_ASSIGN(
+      const TiledHloComputation tiled_computation,
+      ParseAndTile(R"(
+    f {
+      p0 = f32[1,163840] parameter(0)
+      p1 = pred[1,256] parameter(1)
+      p2 = s32[1,512] parameter(2)
+      p3 = pred[1,2] parameter(3)
+      p4 = pred[1,2] parameter(4)
+      p5 = s32[1,16] parameter(5)
+      ag = (f32[4,163840], pred[4,256], s32[4,512], pred[4,2], pred[4,2],
+        s32[4,16]) all-gather(p0, p1, p2, p3, p4, p5),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      gte0 = f32[4,163840] get-tuple-element(ag), index=0
+      gte1 = pred[4,256] get-tuple-element(ag), index=1
+      gte2 = s32[4,512] get-tuple-element(ag), index=2
+      gte3 = pred[4,2] get-tuple-element(ag), index=3
+      gte4 = pred[4,2] get-tuple-element(ag), index=4
+      gte5 = s32[4,16] get-tuple-element(ag), index=5
+      ROOT t = (f32[4,163840], pred[4,256], s32[4,512], pred[4,2], pred[4,2],
+        s32[4,16]) tuple(gte0, gte1, gte2, gte3, gte4, gte5)
+    }
+
+    ENTRY e {
+      p0 = f32[1,163840] parameter(0)
+      p1 = pred[1,256] parameter(1)
+      p2 = s32[1,512] parameter(2)
+      p3 = pred[1,2] parameter(3)
+      p4 = pred[1,2] parameter(4)
+      p5 = s32[1,16] parameter(5)
+      ROOT fusion = (f32[4,163840], pred[4,256], s32[4,512], pred[4,2],
+        pred[4,2], s32[4,16]) fusion(p0, p1, p2, p3, p4, p5), kind=kCustom,
+        calls=f
+    })",
+                   {1, 16384, 1, 64, 1, 128, 1, 1, 1, 1, 1, 4}));
+  EXPECT_THAT(GetSchedule(tiled_computation),
+              IsOkAndHolds(MatchSchedule("d0 -> pid / 10, d1 -> pid mod 10, "
+                                         "d2 -> pid / 10, d3 -> pid mod 10, "
+                                         "d4 -> pid / 10, d5 -> pid mod 10, "
+                                         "d6 -> pid / 10, d7 -> pid mod 10, "
+                                         "d8 -> pid / 10, d9 -> pid mod 10, "
+                                         "d10 -> pid / 10, d11 -> pid mod 10, "
+                                         "num_pids=40, num_tiles=40")));
+}
+
+TEST_F(SchedulingTest, VariadicAllGatherTileSpanningPeersIsRejected) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    f {
+      p0 = f32[1,8,16] parameter(0)
+      p1 = s32[1,8] parameter(1)
+      ag = (f32[4,8,16], s32[4,8]) all-gather(p0, p1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      gte0 = f32[4,8,16] get-tuple-element(ag), index=0
+      gte1 = s32[4,8] get-tuple-element(ag), index=1
+      ROOT t = (f32[4,8,16], s32[4,8]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      arg0 = f32[1,8,16] parameter(0)
+      arg1 = s32[1,8] parameter(1)
+      ROOT fusion = (f32[4,8,16], s32[4,8]) fusion(arg0, arg1),
+        kind=kCustom, calls=f
+    })");
+  auto fusion_adaptor = HloFusionAdaptor::ForInstruction(root);
+  ASSERT_OK_AND_ASSIGN(auto tiling_space,
+                       TilingSpace::Create(*fusion_adaptor, &mlir_context_));
+  ASSERT_OK(tiling_space->AssignTileSizes({2, 8, 16, 1, 8}));
+  EXPECT_THAT(
+      TiledHloComputation::Tile(*fusion_adaptor, std::move(tiling_space)),
+      StatusIs(absl::StatusCode::kUnimplemented,
+               HasSubstr("expected tile size to divide per-rank size")));
+}
+
+// Both outputs have 2 gather tiles/peer; output 1 is padded past the last peer.
+TEST_F(SchedulingTest, VariadicAllGatherMultipleTilesPerPeerOnGatherDim) {
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"(
+    f {
+      p0 = f32[4,16] parameter(0)
+      p1 = s32[2] parameter(1)
+      ag = (f32[16,16], s32[8]) all-gather(p0, p1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      gte0 = f32[16,16] get-tuple-element(ag), index=0
+      gte1 = s32[8] get-tuple-element(ag), index=1
+      ROOT t = (f32[16,16], s32[8]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      arg0 = f32[4,16] parameter(0)
+      arg1 = s32[2] parameter(1)
+      ROOT fusion = (f32[16,16], s32[8]) fusion(arg0, arg1),
+        kind=kCustom, calls=f
+    })",
+                                    {2, 8, 1}));
+  EXPECT_THAT(
+      GetSchedule(tiled_computation),
+      IsOkAndHolds(MatchSchedule(
+          "d0 -> (pid / 4) * 2 + pid mod 4 mod 2, d1 -> (pid mod 4) / 2, "
+          "d2 -> (pid / 4) * 2 + pid mod 4 + ((pid mod 4) / 2) * 8, "
+          "num_pids=16, num_tiles=16")));
+}
+
+TEST_F(SchedulingTest,
+       InteriorVariadicAllGatherConsumedByElementwiseSharesPeerIndex) {
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"(
+    f {
+      p0 = f32[2,16] parameter(0)
+      p1 = s32[4] parameter(1)
+      ag = (f32[8,16], s32[16]) all-gather(p0, p1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      gte0 = f32[8,16] get-tuple-element(ag), index=0
+      gte1 = s32[16] get-tuple-element(ag), index=1
+      c0 = bf16[8,16] convert(gte0)
+      c1 = f32[16] convert(gte1)
+      ROOT t = (bf16[8,16], f32[16]) tuple(c0, c1)
+    }
+
+    ENTRY e {
+      arg0 = f32[2,16] parameter(0)
+      arg1 = s32[4] parameter(1)
+      ROOT fusion = (bf16[8,16], f32[16]) fusion(arg0, arg1),
+        kind=kCustom, calls=f
+    })",
+                                    {2, 8, 2}));
+  EXPECT_THAT(GetSchedule(tiled_computation),
+              IsOkAndHolds(MatchSchedule("d0 -> pid / 2, d1 -> pid mod 2, "
+                                         "d2 -> (pid / 2) * 2 + pid mod 2, "
+                                         "num_pids=8, num_tiles=8")));
+}
+
+TEST_F(SchedulingTest, VariadicAllGatherWithDotProducersSharesPeerIndex) {
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"(
+    f {
+      lhs0 = f32[2,8] parameter(0)
+      rhs0 = f32[8,16] parameter(1)
+      lhs1 = f32[4,8] parameter(2)
+      rhs1 = f32[8,32] parameter(3)
+      dot0 = f32[2,16] dot(lhs0, rhs0),
+        lhs_contracting_dims={1}, rhs_contracting_dims={0}
+      dot1 = f32[4,32] dot(lhs1, rhs1),
+        lhs_contracting_dims={1}, rhs_contracting_dims={0}
+      ag = (f32[8,16], f32[16,32]) all-gather(dot0, dot1),
+        replica_groups={{0,1,2,3}}, dimensions={0}
+      gte0 = f32[8,16] get-tuple-element(ag), index=0
+      gte1 = f32[16,32] get-tuple-element(ag), index=1
+      ROOT t = (f32[8,16], f32[16,32]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      lhs0 = f32[2,8] parameter(0)
+      rhs0 = f32[8,16] parameter(1)
+      lhs1 = f32[4,8] parameter(2)
+      rhs1 = f32[8,32] parameter(3)
+      ROOT fusion = (f32[8,16], f32[16,32]) fusion(lhs0, rhs0, lhs1, rhs1),
+        kind=kCustom, calls=f
+    })",
+                                    {2, 8, 2, 16, 8, 8}));
+  EXPECT_THAT(
+      GetSchedule(tiled_computation),
+      IsOkAndHolds(MatchSchedule(
+          "d0 -> pid / 4, d1 -> pid mod 4, "
+          "d2 -> (pid / 4) * 2 + pid mod 4 mod 2, d3 -> (pid mod 4) / 2, "
+          "num_pids=16, num_tiles=16")));
+}
+
 }  // namespace xla::gpu::experimental
diff --git a/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.h b/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.h
index ffdf21b9152..f8333680506 100644
--- a/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.h
+++ b/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.h
@@ -28,6 +28,7 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/types/span.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
@@ -117,20 +118,45 @@ class EmitterContext {
     return tiled_computation_;
   }
 
+  // Returns the single result of `tiled_hlo`. Multi-result instructions (e.g.
+  // variadic collectives) must be addressed with an explicit result index.
   TensorValue TiledHloToTensorValue(
       const gpu::experimental::TiledHloInstruction& tiled_hlo) const {
-    return tiled_hlo_to_tensor_.at(&tiled_hlo);
+    const auto& results = tiled_hlo_to_tensors_.at(&tiled_hlo);
+    CHECK_EQ(results.size(), 1)
+        << "Expected a single result for " << tiled_hlo.hlo()->ToString();
+    return results.front();
+  }
+
+  // Returns result `result_index` of `tiled_hlo`.
+  TensorValue TiledHloToTensorValue(
+      const gpu::experimental::TiledHloInstruction& tiled_hlo,
+      int64_t result_index) const {
+    const auto& results = tiled_hlo_to_tensors_.at(&tiled_hlo);
+    CHECK_GE(result_index, 0);
+    CHECK_LT(result_index, results.size())
+        << "Result index out of range for " << tiled_hlo.hlo()->ToString();
+    return results[result_index];
+  }
+
+  // Maps `tiled_hlo` to its results. Returns false if it was already mapped.
+  bool MapTiledHloToTensorValues(
+      const gpu::experimental::TiledHloInstruction* tiled_hlo,
+      llvm::ArrayRef<TensorValue> values) {
+    return tiled_hlo_to_tensors_
+        .try_emplace(tiled_hlo, values.begin(), values.end())
+        .second;
   }
 
   bool IsEmitted(
       const gpu::experimental::TiledHloInstruction& tiled_hlo) const {
-    return tiled_hlo_to_tensor_.contains(&tiled_hlo);
+    return tiled_hlo_to_tensors_.contains(&tiled_hlo);
   }
 
   bool MapTiledHloToTensorValue(
       const gpu::experimental::TiledHloInstruction* tiled_hlo,
       TensorValue value) {
-    return tiled_hlo_to_tensor_.insert(std::make_pair(tiled_hlo, value)).second;
+    return MapTiledHloToTensorValues(tiled_hlo, {value});
   }
 
   std::pair<mlir::Value, Interval> GetSequentialDimValue(
@@ -167,8 +193,8 @@ class EmitterContext {
   mlir::Value pid_;
   mlir::Value tid_;
   absl::flat_hash_map<const gpu::experimental::TiledHloInstruction*,
-                      TensorValue>
-      tiled_hlo_to_tensor_;
+                      llvm::SmallVector<TensorValue, 1>>
+      tiled_hlo_to_tensors_;
   gpu::experimental::Schedule schedule_;
   const HloFusionInstruction* fusion_ = nullptr;
   xtile::EntryFuncOp entry_func_;
diff --git a/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc b/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc
index e8b4bb74d0a..4facf30eec6 100644
--- a/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc
+++ b/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include <optional>
 #include <string>
 #include <utility>
+#include <variant>
 #include <vector>
 
 #include "absl/algorithm/container.h"
@@ -113,11 +114,14 @@ namespace arith = ::mlir::arith;
 namespace stablehlo = ::mlir::stablehlo;
 namespace ge = ::xla::gpu::experimental;
 
+// One tensor per result of a tiled HLO instruction.
+using EmittedResults = SmallVector<TensorValue, 1>;
+
 absl::StatusOr<std::vector<TensorValue>> EmitTiledComputation(
     EmitterContext& emitter_ctx, const ge::TiledHloRegion& region,
     absl::Span<const ge::TiledHloInstruction* const> roots);
 
-absl::StatusOr<TensorValue> EmitTiledHloInstruction(
+absl::StatusOr<EmittedResults> EmitTiledHloInstruction(
     EmitterContext& emitter_ctx, const ge::TiledHloInstruction& tiled_hlo);
 
 Value MakeIndex(mlir::ImplicitLocOpBuilder& b, int64_t value) {
@@ -144,7 +148,8 @@ TensorValue ReshapeTo(mlir::ImplicitLocOpBuilder& b, TensorValue input,
       stablehlo::ReshapeOp::create(b, result_type, input).getResult());
 }
 
-absl::StatusOr<TensorValue> EmitAllGather(
+// Emits a (possibly variadic) all-gather and returns one tensor per output.
+absl::StatusOr<EmittedResults> EmitAllGather(
     EmitterContext& emitter_ctx, const HloAllGatherInstruction* all_gather,
     const ge::TiledHloInstruction& tiled_all_gather, ValueRange operands) {
   if (all_gather->device_list()->replica_groups().empty()) {
@@ -163,13 +168,25 @@ absl::StatusOr<TensorValue> EmitAllGather(
   bool use_global_device_ids = all_gather->use_global_device_ids();
 
   ImplicitLocOpBuilder& b = emitter_ctx.b();
-  ABSL_ASSIGN_OR_RETURN(
-      auto output_element_type,
-      xtile::PrimitiveTypeToMlirType(b, all_gather->shape().element_type()));
-  ABSL_ASSIGN_OR_RETURN(SmallVector<int64_t> tile_sizes,
-                   tiled_all_gather.tile().GetStaticTileSizes());
-  auto output_type =
-      mlir::RankedTensorType::get(tile_sizes, output_element_type);
+  // One result per gathered operand. For a variadic all-gather the HLO shape
+  // is a tuple and the tiled instruction holds one tile per output.
+  const int64_t num_outputs = all_gather->shape().IsTuple()
+                                  ? all_gather->shape().tuple_shapes().size()
+                                  : 1;
+  SmallVector<Type> output_types;
+  output_types.reserve(num_outputs);
+  for (int64_t k = 0; k < num_outputs; ++k) {
+    const Shape& output_shape = all_gather->shape().IsTuple()
+                                    ? all_gather->shape().tuple_shapes(k)
+                                    : all_gather->shape();
+    ABSL_ASSIGN_OR_RETURN(
+        auto output_element_type,
+        xtile::PrimitiveTypeToMlirType(b, output_shape.element_type()));
+    ABSL_ASSIGN_OR_RETURN(SmallVector<int64_t> tile_sizes,
+                     tiled_all_gather.tile(k).GetStaticTileSizes());
+    output_types.push_back(
+        mlir::RankedTensorType::get(tile_sizes, output_element_type));
+  }
 
   auto replica_groups_type = mlir::RankedTensorType::get(
       {static_cast<int64_t>(all_gather->replica_groups().size()),
@@ -184,10 +201,15 @@ absl::StatusOr<TensorValue> EmitAllGather(
                                                                /*type=*/0)
                      : nullptr;
   auto all_gather_op = mlir::stablehlo::AllGatherOp::create(
-      b, mlir::TypeRange{output_type}, operands,
+      b, mlir::TypeRange{output_types}, operands,
       all_gather->all_gather_dimension(), replica_groups_attr,
       channel_handle_attr, use_global_device_ids);
-  return mlir::cast<TensorValue>(all_gather_op.getResult(0));
+  EmittedResults results;
+  results.reserve(num_outputs);
+  for (Value result : all_gather_op.getResults()) {
+    results.push_back(mlir::cast<TensorValue>(result));
+  }
+  return results;
 }
 
 absl::StatusOr<TensorValue> EmitAllReduce(
@@ -700,9 +722,9 @@ absl::StatusOr<TensorValue> EmitRegionInstructionWithDependencies(
     if (!visited.contains(instr)) {
       continue;
     }
-    ABSL_ASSIGN_OR_RETURN(TensorValue value,
+    ABSL_ASSIGN_OR_RETURN(EmittedResults values,
                      EmitTiledHloInstruction(emitter_ctx, *instr));
-    TF_RET_CHECK(emitter_ctx.MapTiledHloToTensorValue(instr, value))
+    TF_RET_CHECK(emitter_ctx.MapTiledHloToTensorValues(instr, values))
         << instr->hlo()->ToString();
   }
   TF_RET_CHECK(emitter_ctx.IsEmitted(*root))
@@ -2274,8 +2296,9 @@ absl::StatusOr<TensorValue> EmitReduce(
              : EmitReduceWithRegion(emitter_ctx, tiled_hlo);
 }
 
-absl::StatusOr<TensorValue> EmitTiledHloInstruction(
-    EmitterContext& emitter_ctx, const ge::TiledHloInstruction& tiled_hlo) {
+absl::StatusOr<std::variant<TensorValue, EmittedResults>>
+EmitTiledHloInstructionImpl(EmitterContext& emitter_ctx,
+                            const ge::TiledHloInstruction& tiled_hlo) {
   auto& b = emitter_ctx.b();
   const HloInstruction* hlo = tiled_hlo.hlo();
   VLOG(4) << "EmitTiledHloInstruction: " << hlo->ToString();
@@ -2344,12 +2367,8 @@ absl::StatusOr<TensorValue> EmitTiledHloInstruction(
     return result.front();
   }
   if (hlo->opcode() == HloOpcode::kGetTupleElement) {
-    int64_t index = hlo->tuple_index();
-    if (index == 0) {
-      return emitter_ctx.TiledHloToTensorValue(*tiled_hlo.operand(0));
-    }
-    return absl::UnimplementedError(
-        absl::StrCat("Unsupported get-tuple-element index ", index));
+    return emitter_ctx.TiledHloToTensorValue(*tiled_hlo.operand(0),
+                                             hlo->tuple_index());
   }
   std::vector<Value> operands;
   operands.reserve(hlo->operands().size());
@@ -2431,15 +2450,25 @@ absl::StatusOr<TensorValue> EmitTiledHloInstruction(
       absl::StrCat("Unsupported operation ", hlo->ToString()));
 }
 
+absl::StatusOr<EmittedResults> EmitTiledHloInstruction(
+    EmitterContext& emitter_ctx, const ge::TiledHloInstruction& tiled_hlo) {
+  ABSL_ASSIGN_OR_RETURN(auto result,
+                   EmitTiledHloInstructionImpl(emitter_ctx, tiled_hlo));
+  if (auto* single = std::get_if<TensorValue>(&result)) {
+    return EmittedResults{*single};
+  }
+  return std::move(std::get<EmittedResults>(result));
+}
+
 absl::StatusOr<std::vector<TensorValue>> EmitTiledComputation(
     EmitterContext& emitter_ctx, const ge::TiledHloRegion& region,
     absl::Span<const ge::TiledHloInstruction* const> roots) {
   for (const auto* tiled_hlo : region.instructions()) {
     const HloInstruction* hlo = tiled_hlo->hlo();
     VLOG(8) << "Emitting " << hlo->ToString(HloPrintOptions::ShortParsable());
-    ABSL_ASSIGN_OR_RETURN(TensorValue result,
+    ABSL_ASSIGN_OR_RETURN(EmittedResults results,
                      EmitTiledHloInstruction(emitter_ctx, *tiled_hlo));
-    TF_RET_CHECK(emitter_ctx.MapTiledHloToTensorValue(tiled_hlo, result))
+    TF_RET_CHECK(emitter_ctx.MapTiledHloToTensorValues(tiled_hlo, results))
         << hlo->ToString();
   }
   std::vector<TensorValue> results;

From fcb146d1780ac754ca435fca708c76fdad743cd5 Mon Sep 17 00:00:00 2001
From: Harsha H S <hsharsha@users.noreply.github.com>
Date: Tue, 13 Oct 2026 03:49:52 -0700
Subject: [PATCH 0677/1000] PR #50379: [ROCm] Bump up to theRock10.1.0
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50379

📝 Summary of Changes
Use the latest released theRock 10.1.0 tar ball

Copybara import of the project:

--
1af23a00b50803758dba87f8b3caeb4ef5a21716 by Harsha Havanur Shamsundara <Harsha.HavanurShamsundara@amd.com>:

Bump up to theRock10.1.0

Merging this change closes #50379

PiperOrigin-RevId: 998525612
---
 third_party/xla/.github/workflows/rocm_ci.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/.github/workflows/rocm_ci.yml b/third_party/xla/.github/workflows/rocm_ci.yml
index 43aa1903d48..e1b7b177371 100644
--- a/third_party/xla/.github/workflows/rocm_ci.yml
+++ b/third_party/xla/.github/workflows/rocm_ci.yml
@@ -38,8 +38,8 @@ concurrency:
   cancel-in-progress: ${{ github.ref != 'main' }}
 
 env:
-  ROCM_DISTRO_URL: https://stable.repo.amd.com/rocm/core/tarball/therock-dist-linux-multiarch-10.0.0.tar.gz
-  ROCM_DISTRO_HASH: 1c5e807875d26a2470ecc7323daa5b5b9009208a55c3290ac255a909cde15fc6
+  ROCM_DISTRO_URL: https://stable.repo.amd.com/rocm/core/tarball/therock-dist-linux-multiarch-10.1.0.tar.gz
+  ROCM_DISTRO_HASH: a95c72f0052168f2262365dff71c2259331721ad0a65da7957dc7d96ccb77021
   ROCM_AMD_GPU_TARGETS: gfx950
   ROCM_SINGLE_GPU_RBE_POOL: linux_x64_gpu
 

From ec1b954897024c0616aa74284c3aae9a018aed78 Mon Sep 17 00:00:00 2001
From: Ilia Sergachev <isergachev@nvidia.com>
Date: Tue, 13 Oct 2026 03:52:07 -0700
Subject: [PATCH 0678/1000] PR #50288: [GPU] Fix while thunk unroll race.
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50288

📝 Summary of Changes

is_unrolled_loop_ used to be concurrently flipped by threads of
different GPUs in thunk Initialize() causing compilation mismatches and hangs.

🚀 Kind of Contribution

🐛 Bug Fix

🧪 Unit Tests:

added
Copybara import of the project:

--
c40ff3bcfbd050063bac307f4b1da1dbe2b6b4ae by Ilia Sergachev <isergachev@nvidia.com>:

[GPU] Fix while thunk unroll race.

is_unrolled_loop_ used to be concurrently flipped by threads of
different GPUs in Initialize() causing mismatches and hangs.

Merging this change closes #50288

PiperOrigin-RevId: 998526489
---
 .../xla/xla/backends/gpu/runtime/BUILD        |  1 +
 .../xla/backends/gpu/runtime/while_thunk.cc   | 38 +++++------
 .../backends/gpu/runtime/while_thunk_test.cc  | 67 +++++++++++++++++++
 3 files changed, 85 insertions(+), 21 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index f2d11fbc751..3656013960b 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -3897,6 +3897,7 @@ xla_test(
         "//xla/tsl/platform:test",
         "//xla/tsl/util/proto:parse_text_proto",
         "//xla/tsl/util/proto:proto_matchers",
+        "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
diff --git a/third_party/xla/xla/backends/gpu/runtime/while_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/while_thunk.cc
index e0eccd023ca..ec7d75ee892 100644
--- a/third_party/xla/xla/backends/gpu/runtime/while_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/while_thunk.cc
@@ -134,26 +134,6 @@ absl::Status WhileThunk::Initialize(const InitializeParams& params) {
     ABSL_RETURN_IF_ERROR(command_body_executor_->Initialize(params));
   }
 
-  is_unrolled_loop_ = false;
-  if (command_condition_executor_.has_value() &&
-      command_body_executor_.has_value() && enable_loop_unroll_ &&
-      command_body_executor_->support_loop_unroll() &&
-      command_condition_executor_->support_loop_unroll() &&
-      trip_count_.has_value()) {
-    is_unrolled_loop_ = true;
-  }
-  VLOG(3) << "WhileThunk::Initialize command buffer: enable_loop_unroll_="
-          << enable_loop_unroll_ << ", body_support="
-          << (command_body_executor_.has_value()
-                  ? command_body_executor_->support_loop_unroll()
-                  : false)
-          << ", cond_support="
-          << (command_condition_executor_.has_value()
-                  ? command_condition_executor_->support_loop_unroll()
-                  : false)
-          << ", trip_count=" << trip_count_.value_or(-1)
-          << ", is_unrolled_loop_=" << is_unrolled_loop_;
-
   return host_memory_pools_.GetOrCreateAndInitialize(
       params.executor->device_ordinal(), [&](PoolState* state) -> absl::Status {
         ABSL_ASSIGN_OR_RETURN(
@@ -270,7 +250,23 @@ absl::Status WhileThunk::SetOrUpdateCommandBufferExecutors(
   command_condition_executor_ = std::move(condition_executor);
   command_body_executor_ = std::move(body_executor);
   enable_loop_unroll_ = enable_loop_unroll;
-  is_unrolled_loop_ = false;
+  is_unrolled_loop_ = command_condition_executor_.has_value() &&
+                      command_body_executor_.has_value() &&
+                      enable_loop_unroll_ &&
+                      command_body_executor_->support_loop_unroll() &&
+                      command_condition_executor_->support_loop_unroll() &&
+                      trip_count_.has_value();
+  VLOG(3) << "WhileThunk command buffer: enable_loop_unroll_="
+          << enable_loop_unroll_ << ", body_support="
+          << (command_body_executor_.has_value()
+                  ? command_body_executor_->support_loop_unroll()
+                  : false)
+          << ", cond_support="
+          << (command_condition_executor_.has_value()
+                  ? command_condition_executor_->support_loop_unroll()
+                  : false)
+          << ", trip_count=" << trip_count_.value_or(-1)
+          << ", is_unrolled_loop_=" << is_unrolled_loop_;
   return absl::OkStatus();
 }
 
diff --git a/third_party/xla/xla/backends/gpu/runtime/while_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/while_thunk_test.cc
index ea2a95c0ae7..6deb42c57b3 100644
--- a/third_party/xla/xla/backends/gpu/runtime/while_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/while_thunk_test.cc
@@ -26,6 +26,7 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/functional/any_invocable.h"
 #include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
@@ -559,5 +560,71 @@ TEST(WhileThunkTest, TransformNested) {
               Kind::kCustomCall);
 }
 
+TEST(WhileThunkTest, UnrollDecisionIsReadyBeforeDeviceInitialization) {
+  constexpr int kTripCount = 7;
+  BufferAllocation pred_alloc(/*index=*/0, /*size=*/sizeof(bool), /*color=*/0);
+  BufferAllocation::Slice pred_slice(&pred_alloc, /*offset=*/0,
+                                     /*size=*/sizeof(bool));
+  WhileThunk thunk(Thunk::ThunkInfo(), pred_slice, ThunkSequence(),
+                   ThunkSequence(), kTripCount, /*devices_per_host=*/1);
+  CommandRecordCounts cond_counts;
+  CommandRecordCounts body_counts;
+  ASSERT_OK_AND_ASSIGN(CommandExecutor cond_executor,
+                       MakeCommandExecutor(&cond_counts));
+  ASSERT_OK_AND_ASSIGN(CommandExecutor body_executor,
+                       MakeCommandExecutor(&body_counts));
+  ASSERT_OK(thunk.SetOrUpdateCommandBufferExecutors(
+      std::move(cond_executor), std::move(body_executor),
+      /*enable_loop_unroll=*/true));
+
+  bool pred = true;
+  std::vector<se::DeviceAddressBase> buffers = {
+      se::DeviceAddressBase(&pred, sizeof(pred))};
+  BufferAllocations allocations(buffers, /*device_ordinal=*/0,
+                                /*memory_allocator=*/nullptr);
+  ServiceExecutableRunOptions run_options;
+  Thunk::ExecuteParams execute_params = Thunk::ExecuteParams::Create(
+      run_options, allocations, /*stream=*/nullptr,
+      /*command_buffer_trace_stream=*/nullptr, /*collective_params=*/nullptr,
+      /*collective_cliques=*/nullptr, /*collective_memory=*/nullptr);
+  CommandStateManager state_manager;
+  Command::RecordParams record_params{state_manager};
+  ::testing::NiceMock<se::MockCommandBuffer> command_buffer;
+  NestedCommandBuffer child;
+  ConfigureNestedCommandBuffer(&child);
+  FakeSeCommand child_command;
+  using RecordChild = absl::AnyInvocable<absl::Status(se::CommandBuffer*)>;
+  EXPECT_CALL(command_buffer,
+              CreateChildCommand(::testing::A<RecordChild>(), ::testing::_))
+      .WillOnce([&](RecordChild record,
+                    absl::Span<const se::CommandBuffer::Command* const>)
+                    -> absl::StatusOr<const se::CommandBuffer::Command*> {
+        ABSL_RETURN_IF_ERROR(record(child.command_buffer.get()));
+        ABSL_RETURN_IF_ERROR(child.command_buffer->Finalize());
+        return &child_command;
+      });
+  EXPECT_CALL(command_buffer,
+              CreateWhile(::testing::A<se::DeviceAddress<bool>>(), ::testing::_,
+                          ::testing::_, ::testing::_))
+      .Times(0);
+  ASSERT_OK_AND_ASSIGN(const se::CommandBuffer::Command* recorded,
+                       thunk.Record(execute_params, record_params,
+                                    Command::RecordCreate{}, &command_buffer));
+  EXPECT_EQ(recorded, &child_command);
+  EXPECT_EQ(cond_counts.creates, kTripCount);
+  EXPECT_EQ(body_counts.creates, kTripCount);
+
+  ASSERT_OK(child.command_buffer->Update());
+  EXPECT_CALL(command_buffer,
+              UpdateChildCommand(&child_command, ::testing::A<RecordChild>()))
+      .WillOnce([&](const se::CommandBuffer::Command*, RecordChild update) {
+        return update(child.command_buffer.get());
+      });
+  ASSERT_OK(thunk.Record(execute_params, record_params,
+                         Command::RecordUpdate{recorded}, &command_buffer));
+  EXPECT_EQ(cond_counts.updates, kTripCount);
+  EXPECT_EQ(body_counts.updates, kTripCount);
+}
+
 }  // namespace
 }  // namespace xla::gpu

From cecf9f7999867bebab54de5266603d0b5d8a164f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 13 Oct 2026 03:53:35 -0700
Subject: [PATCH 0679/1000] Previously, if any buffer color exceeded its memory
 limit during fast buffer assignment, fallback cleared all allocations and ran
 buffer assignment for every color from scratch.

This CL separates the fallback mechanic for different colors.

PiperOrigin-RevId: 998527152
---
 .../xla/xla/service/buffer_assignment.cc      | 139 ++++++++++++------
 .../xla/xla/service/buffer_assignment.h       |  27 +++-
 .../xla/xla/service/buffer_assignment_test.cc |  83 +++++++++++
 3 files changed, 202 insertions(+), 47 deletions(-)

diff --git a/third_party/xla/xla/service/buffer_assignment.cc b/third_party/xla/xla/service/buffer_assignment.cc
index dbd46c8dd91..32f0bbd1e9f 100644
--- a/third_party/xla/xla/service/buffer_assignment.cc
+++ b/third_party/xla/xla/service/buffer_assignment.cc
@@ -1339,6 +1339,31 @@ absl::Status BufferAssignment::CombineTempAllocations(
   return absl::OkStatus();
 }
 
+void BufferAssignment::ClearAllocationsForColors(
+    const absl::flat_hash_set<LogicalBuffer::Color>& colors) {
+  allocations_.erase(std::remove_if(allocations_.begin(), allocations_.end(),
+                                    [&](const BufferAllocation& alloc) {
+                                      return colors.contains(alloc.color());
+                                    }),
+                     allocations_.end());
+
+  temp_allocation_total_size_ = 0;
+  allocation_index_for_value_.clear();
+  absl::erase_if(fragmentation_bytes_by_color_, [&](const auto& entry) {
+    return colors.contains(entry.first);
+  });
+  stats_ = Stats();
+  for (BufferAllocation::Index index = 0; index < allocations_.size();
+       ++index) {
+    BufferAllocation& allocation = allocations_[index];
+    allocation.set_index(index);
+    // NOLINTNEXTLINE : the order of the loop is not important.
+    for (const auto& [value, offset_size] : allocation.assigned_buffers_) {
+      allocation_index_for_value_[value] = index;
+    }
+  }
+}
+
 void BufferAssignment::ComputeSummaryStats() {
   for (auto& allocation : Allocations()) {
     if (allocation.is_entry_computation_parameter()) {
@@ -1360,6 +1385,12 @@ void BufferAssignment::ComputeSummaryStats() {
     stats_.total_allocation_count++;
     stats_.total_allocation_bytes += allocation.size();
   }
+  if (!fragmentation_bytes_by_color_.empty()) {
+    stats_.preallocated_temp_fragmentation_bytes = 0;
+    for (const auto& [color, bytes] : fragmentation_bytes_by_color_) {
+      stats_.preallocated_temp_fragmentation_bytes += bytes;
+    }
+  }
 }
 
 absl::StatusOr<int64_t> BufferAssignment::ComputeTotalFragmentationBytes(
@@ -2581,7 +2612,8 @@ absl::Status BufferAssigner::AssignBuffersForGlobalComputations(
         buffers_to_assign_sequentially,
     BufferAssignment* assignment,
     buffer_assignment::AssignmentAlgorithmForComputationsWithoutOrderingProto::
-        Value algorithm) {
+        Value algorithm,
+    const absl::flat_hash_set<LogicalBuffer::Color>* colors_to_assign) {
   if (computations.empty()) {
     return absl::OkStatus();
   }
@@ -2591,7 +2623,8 @@ absl::Status BufferAssigner::AssignBuffersForGlobalComputations(
   // First assign the preset allocations.
   absl::flat_hash_set<const HloBuffer*> preset_assigned_buffers;
 
-  ABSL_RETURN_IF_ERROR(AssignPresetBuffers(&preset_assigned_buffers, assignment));
+  ABSL_RETURN_IF_ERROR(AssignPresetBuffers(&preset_assigned_buffers, assignment,
+                                      colors_to_assign));
 
   const HloAliasAnalysis& alias_analysis = assignment->alias_analysis();
 
@@ -2607,10 +2640,31 @@ absl::Status BufferAssigner::AssignBuffersForGlobalComputations(
     TF_RET_CHECK(!buffer.values().empty());
     const HloComputation* comp = buffer.values()[0]->instruction()->parent();
     if (computations_set.contains(comp)) {
+      ABSL_ASSIGN_OR_RETURN(LogicalBuffer::Color color, buffer.color());
+      if (colors_to_assign != nullptr && !colors_to_assign->contains(color)) {
+        continue;
+      }
       sorted_buffers.push_back(&buffer);
     }
   }
 
+  for (const HloComputation* computation : computations) {
+    const HloInstructionSequence* instruction_sequence =
+        assignment->hlo_ordering().SequentialOrder(*computation);
+    const bool has_sequential_order = instruction_sequence != nullptr;
+    if (has_sequential_order && buffers_to_assign_sequentially != nullptr) {
+      // Every sequential computation must get an entry in the
+      // buffers_to_assign_sequentially map, even if we end up with an empty
+      // set of buffers. This ensures we can correctly determine whether to
+      // run whole-module heap simulation.
+      buffers_to_assign_sequentially->emplace(computation,
+                                              flat_hash_set<const HloValue*>());
+    }
+  }
+  if (sorted_buffers.empty()) {
+    return absl::OkStatus();
+  }
+
   // Generate a post order sort of instructions for sorting of the
   // HloBuffers.
   flat_hash_map<const HloInstruction*, int> post_order_position;
@@ -2632,20 +2686,6 @@ absl::Status BufferAssigner::AssignBuffersForGlobalComputations(
     }
   }
 
-  for (const HloComputation* computation : computations) {
-    const HloInstructionSequence* instruction_sequence =
-        assignment->hlo_ordering().SequentialOrder(*computation);
-    const bool has_sequential_order = instruction_sequence != nullptr;
-    if (has_sequential_order) {
-      // Every sequential computation must get an entry in the
-      // buffers_to_assign_sequentially map, even if we end up with an empty
-      // set of buffers. This ensures we can correctly determine whether to
-      // run whole-module heap simulation.
-      buffers_to_assign_sequentially->emplace(computation,
-                                              flat_hash_set<const HloValue*>());
-    }
-  }
-
   using Comparator = absl::AnyInvocable<std::optional<bool>(
       const HloBuffer* a, const HloBuffer* b)>;
   auto size_of = absl::bind_front(&BufferAssignment::HloBufferSize, assignment);
@@ -3021,7 +3061,8 @@ absl::StatusOr<int64_t> BufferAssigner::ReuseCompatibleTempHeaps(
 
 absl::Status BufferAssigner::AssignPresetBuffers(
     absl::flat_hash_set<const HloBuffer*>* assigned_buffers,
-    BufferAssignment* assignment) {
+    BufferAssignment* assignment,
+    const absl::flat_hash_set<LogicalBuffer::Color>* colors_to_assign) {
   if (!opts_.preset_assignments) {
     return absl::OkStatus();
   }
@@ -3032,6 +3073,9 @@ absl::Status BufferAssigner::AssignPresetBuffers(
   for (auto& color_and_info :
        opts_.preset_assignments->assignment_informations()) {
     LogicalBuffer::Color color(color_and_info.first);
+    if (colors_to_assign != nullptr && !colors_to_assign->contains(color)) {
+      continue;
+    }
     auto inserted = preset_allocations.emplace(
         color,
         assignment->NewEmptyAllocation(color_and_info.second.size, color));
@@ -3050,6 +3094,10 @@ absl::Status BufferAssigner::AssignPresetBuffers(
     const HloPosition& defining_position = position_and_chunk.first;
     const HloBuffer& buffer = alias_analysis.GetUniqueBufferAt(
         defining_position.instruction, defining_position.index);
+    ABSL_ASSIGN_OR_RETURN(LogicalBuffer::Color color, buffer.color());
+    if (colors_to_assign != nullptr && !colors_to_assign->contains(color)) {
+      continue;
+    }
     for (const HloValue* value : buffer.values()) {
       VLOG(3) << "Preset allocation for value: " << value->ToShortString();
       const HeapSimulator::Chunk& chunk = position_and_chunk.second;
@@ -3455,11 +3503,8 @@ absl::Status BufferAssigner::AssignBuffersFromHeapSimulator(
     BufferValue::Color color,
     std::optional<BufferAssignment::BufferIsolationOptions> isolation_options) {
   IsolateHeapBuffers(isolation_options, assignment, color, result);
-  if (assignment->stats_.preallocated_temp_fragmentation_bytes == -1) {
-    assignment->stats_.preallocated_temp_fragmentation_bytes =
-        result.fragmentation_size;
-  } else {
-    assignment->stats_.preallocated_temp_fragmentation_bytes +=
+  if (result.fragmentation_size >= 0) {
+    assignment->fragmentation_bytes_by_color_[color] +=
         result.fragmentation_size;
   }
   VLOG(1) << "Result size from heap simulator: " << result.heap_size;
@@ -3590,20 +3635,23 @@ absl::Status BufferAssigner::RunAssignBuffersWithFallback(
   }
 
   // 2. Primary Run: Execute buffer assignment using the primary configured
-  // algorithms.
+  // algorithms across all colors.
   absl::Status primary_status = RunAssignBuffers(
       module, global_computations, thread_local_computations, assignment,
       opts_.buffer_assignment_algorithm,
       opts_.assignment_algorithm_for_computations_without_ordering);
 
-  // 3. Check if the primary run succeeded and met memory limits.
-  // We calculate the total allocated memory per color and compare it against
-  // the adjusted memory limit (subtracting a safety margin of 2.5 GiB).
-  bool need_fallback = false;
+  // 3. Check if the primary run succeeded and met memory limits per color.
+  // Collect only the colors that exceeded their adjusted memory limit
+  // (subtracting a safety margin of 2.5 GiB).
+  absl::flat_hash_set<LogicalBuffer::Color> fallback_colors;
   if (absl::IsResourceExhausted(primary_status)) {
-    need_fallback = true;
     VLOG(1) << "Primary BufferAssignment failed with ResourceExhausted. "
-            << "Triggering in-place fallback.";
+            << "Triggering in-place fallback for all colors.";
+    for (const HloBuffer& buffer : assignment->alias_analysis().buffers()) {
+      ABSL_ASSIGN_OR_RETURN(LogicalBuffer::Color color, buffer.color());
+      fallback_colors.insert(color);
+    }
   } else if (primary_status.ok()) {
     // Aggregate total allocated bytes across all allocations by color.
     // Ensure we account for alignment fragmentation exactly the way
@@ -3632,34 +3680,38 @@ absl::Status BufferAssigner::RunAssignBuffersWithFallback(
             std::max<int64_t>(0, memory_limit - kTwoPointFiveGiB);
 
         if (total_allocated_bytes > adjusted_limit) {
-          need_fallback = true;
+          fallback_colors.insert(color);
           VLOG(1) << "Primary BufferAssignment exceeded memory limit for color "
                   << color << " (" << total_allocated_bytes
                   << " bytes vs limit of " << adjusted_limit
                   << " bytes). Triggering in-place fallback to DEFAULT.";
-          break;
+        } else {
+          VLOG(1) << "Primary BufferAssignment did not exceed memory limit for "
+                     "color "
+                  << color << " (" << total_allocated_bytes
+                  << " bytes vs limit of " << adjusted_limit << " bytes).";
         }
+      } else {
         VLOG(1)
-            << "Primary BufferAssignment did not exceed memory limit for color "
-            << color << " (" << total_allocated_bytes << " bytes vs limit of "
-            << adjusted_limit << " bytes).";
+            << "Primary BufferAssignment did not set memory limit for color "
+            << color << ". Fallback not triggered.";
       }
     }
+  } else {
+    return primary_status;
   }
 
-  // 4. Execute In-Place Fallback if needed.
-  // We reset the BufferAssignment object allocations in-place, restore the
-  // backup preset assignments, and execute the assignment passes using the
-  // DEFAULT fallback algorithms.
-  if (need_fallback) {
-    assignment->ClearAllocations();
+  // 4. Execute In-Place Fallback only for the colors that exceeded their limit.
+  if (!fallback_colors.empty()) {
+    assignment->ClearAllocationsForColors(fallback_colors);
     if (backup_presets) {
       opts_.preset_assignments = backup_presets->ClonePresetAssignments();
     }
     ABSL_RETURN_IF_ERROR(RunAssignBuffers(
         module, global_computations, thread_local_computations, assignment,
         opts_.fallback_algorithm,
-        opts_.fallback_algorithm_for_computations_without_ordering));
+        opts_.fallback_algorithm_for_computations_without_ordering,
+        &fallback_colors));
   }
   return absl::OkStatus();
 }
@@ -3672,7 +3724,8 @@ absl::Status BufferAssigner::RunAssignBuffers(
     buffer_assignment::BufferAssignmentAlgorithmProto::Value
         sequential_algorithm,
     buffer_assignment::AssignmentAlgorithmForComputationsWithoutOrderingProto::
-        Value non_sequential_algorithm) {
+        Value non_sequential_algorithm,
+    const absl::flat_hash_set<LogicalBuffer::Color>* colors_to_assign) {
   // First assign buffers for global computations. Temporary buffers for
   // sequential computations are collected in
   // 'buffers_to_assign_sequentially'.
@@ -3680,7 +3733,7 @@ absl::Status BufferAssigner::RunAssignBuffers(
       buffers_to_assign_sequentially;
   ABSL_RETURN_IF_ERROR(AssignBuffersForGlobalComputations(
       global_computations, &buffers_to_assign_sequentially, assignment,
-      non_sequential_algorithm));
+      non_sequential_algorithm, colors_to_assign));
   // Assign buffers with sequential ordering, if any. If all global
   // computations are sequential, we can run heap simulation on the whole
   // module, which reduces memory usage.
diff --git a/third_party/xla/xla/service/buffer_assignment.h b/third_party/xla/xla/service/buffer_assignment.h
index a36cdcb7132..72ac832d24f 100644
--- a/third_party/xla/xla/service/buffer_assignment.h
+++ b/third_party/xla/xla/service/buffer_assignment.h
@@ -28,6 +28,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/container/btree_map.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/functional/any_invocable.h"
@@ -402,7 +403,12 @@ class BufferAllocation {
   absl::Status AddAssignment(const HloValue& buffer, int64_t offset,
                              int64_t size);
 
-  void set_index(Index index) { index_ = index; }
+  void set_index(Index index) {
+    index_ = index;
+    for (HeapSimulatorTrace& trace : heap_traces_) {
+      trace.set_buffer_allocation_index(index);
+    }
+  }
   void set_size(int64_t size) { size_ = size; }
 
   // The index of the allocation in the BufferAssignment.
@@ -718,9 +724,14 @@ class BufferAssignment {
     allocations_.clear();
     temp_allocation_total_size_ = 0;
     allocation_index_for_value_.clear();
+    fragmentation_bytes_by_color_.clear();
     stats_ = Stats();
   }
 
+  // Clears only the buffer allocations belonging to `colors`.
+  void ClearAllocationsForColors(
+      const absl::flat_hash_set<LogicalBuffer::Color>& colors);
+
   // Creates and returns a new BufferAllocation, with no assigned
   // LogicalBuffers. Ownership is maintained internally.
   BufferAllocation* NewEmptyAllocation(int64_t size,
@@ -823,6 +834,8 @@ class BufferAssignment {
 
   Stats stats_;
 
+  absl::btree_map<LogicalBuffer::Color, int64_t> fragmentation_bytes_by_color_;
+
   absl::flat_hash_map<HloBuffer::Id, int64_t> cached_buffer_sizes_;
 
   BufferAssignment(const BufferAssignment&) = delete;
@@ -1073,7 +1086,9 @@ class BufferAssigner {
           sequential_algorithm,
       buffer_assignment::
           AssignmentAlgorithmForComputationsWithoutOrderingProto::Value
-              non_sequential_algorithm);
+              non_sequential_algorithm,
+      const absl::flat_hash_set<LogicalBuffer::Color>* colors_to_assign =
+          nullptr);
 
   absl::Status RunAssignBuffersWithFallback(
       const HloModule* module,
@@ -1093,7 +1108,9 @@ class BufferAssigner {
       BufferAssignment* assignment,
       buffer_assignment::
           AssignmentAlgorithmForComputationsWithoutOrderingProto::Value
-              algorithm);
+              algorithm,
+      const absl::flat_hash_set<LogicalBuffer::Color>* colors_to_assign =
+          nullptr);
 
   // Assigns dedicated thread-local buffers for the given thread-local
   // computations. `assignment` is modified to reflect the new buffer
@@ -1113,7 +1130,9 @@ class BufferAssigner {
   // to assigned_buffers and skip buffer allocation.
   absl::Status AssignPresetBuffers(
       absl::flat_hash_set<const HloBuffer*>* assigned_buffers,
-      BufferAssignment* assignment);
+      BufferAssignment* assignment,
+      const absl::flat_hash_set<LogicalBuffer::Color>* colors_to_assign =
+          nullptr);
 
   // Handles HloBuffers that represent DUS views or constants. Returns true if
   // the buffer was handled (either allocated as a constant or skipped).
diff --git a/third_party/xla/xla/service/buffer_assignment_test.cc b/third_party/xla/xla/service/buffer_assignment_test.cc
index 19606072b9f..bba5c535cdb 100644
--- a/third_party/xla/xla/service/buffer_assignment_test.cc
+++ b/third_party/xla/xla/service/buffer_assignment_test.cc
@@ -866,6 +866,89 @@ TEST_F(BufferAssignmentTest, OOMFallbackIgnoresThreadLocalAllocations) {
   EXPECT_EQ(other_bytes_fallback, 288);
 }
 
+TEST_F(BufferAssignmentTest, OOMFallbackOnlyForExceededColor) {
+  auto builder = HloComputation::Builder(TestName());
+  auto param = builder.AddInstruction(
+      HloInstruction::CreateParameter(0, ShapeUtil::MakeShape(F32, {}), "p1"));
+
+  auto with_space = [](Shape s, int64_t space) {
+    s.mutable_layout()->set_memory_space(space);
+    return s;
+  };
+  Shape s10 = ShapeUtil::MakeShape(F32, {10});  // 40 bytes
+  Shape s20 = ShapeUtil::MakeShape(F32, {20});  // 80 bytes
+  Shape s30 = ShapeUtil::MakeShape(F32, {30});  // 120 bytes
+  Shape s40 = ShapeUtil::MakeShape(F32, {40});  // 160 bytes
+  Shape s1 = ShapeUtil::MakeShape(F32, {1});    // 4 bytes
+
+  // Chain in memory space 0 (exceeds limit -> falls back to DEFAULT: 288B).
+  auto a0 = builder.AddInstruction(
+      HloInstruction::CreateCustomCall(with_space(s10, 0), {param}, "dummy"));
+  auto b0 = builder.AddInstruction(
+      HloInstruction::CreateCustomCall(with_space(s20, 0), {a0}, "dummy"));
+  auto c0 = builder.AddInstruction(
+      HloInstruction::CreateCustomCall(with_space(s30, 0), {b0}, "dummy"));
+  auto d0 = builder.AddInstruction(
+      HloInstruction::CreateCustomCall(with_space(s40, 0), {c0}, "dummy"));
+
+  // Identical chain in memory space 2 (within limit -> stays FAST_MERGE: 404B).
+  auto a2 = builder.AddInstruction(
+      HloInstruction::CreateCustomCall(with_space(s10, 2), {param}, "dummy"));
+  auto b2 = builder.AddInstruction(
+      HloInstruction::CreateCustomCall(with_space(s20, 2), {a2}, "dummy"));
+  auto c2 = builder.AddInstruction(
+      HloInstruction::CreateCustomCall(with_space(s30, 2), {b2}, "dummy"));
+  auto d2 = builder.AddInstruction(
+      HloInstruction::CreateCustomCall(with_space(s40, 2), {c2}, "dummy"));
+
+  auto root = builder.AddInstruction(
+      HloInstruction::CreateCustomCall(with_space(s1, 0), {d0, d2}, "dummy"));
+
+  auto module = CreateNewVerifiedModule();
+  module->AddEntryComputation(builder.Build());
+
+  HloSchedule schedule(module.get());
+  schedule.set_sequence(module->entry_computation(),
+                        {param, a0, b0, c0, d0, a2, b2, c2, d2, root});
+  CHECK_OK(module->set_schedule(schedule));
+
+  BufferAssigner::Options opts;
+  opts.assignment_algorithm_for_computations_without_ordering =
+      buffer_assignment::
+          AssignmentAlgorithmForComputationsWithoutOrderingProto::FAST_MERGE;
+  opts.buffer_assignment_algorithm = buffer_assignment::
+      BufferAssignmentAlgorithmProto::FAST_MERGE_WITH_FALLBACK;
+  opts.fallback_algorithm =
+      buffer_assignment::BufferAssignmentAlgorithmProto::DEFAULT;
+
+  constexpr int64_t kSafetyMargin = int64_t{5} << 29;
+  opts.color_memory_limit = [&](LogicalBuffer::Color color) {
+    // Color 0 exceeds limit (tight limit); Color 2 fits comfortably.
+    return color == 0 ? (kSafetyMargin + 10) : (kSafetyMargin + 10000000);
+  };
+
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<BufferAssignment> assignment,
+      BufferAssigner::Run(
+          module.get(), std::make_unique<SequentialHloOrdering>(schedule),
+          &BufferSizeBytes, &alias_info_,
+          [](LogicalBuffer::Color) { return 1; }, std::move(opts)));
+
+  int64_t color0_bytes = 0;
+  int64_t color2_bytes = 0;
+  for (const BufferAllocation& alloc : assignment->Allocations()) {
+    if (alloc.color() == 0) {
+      color0_bytes += alloc.size();
+    } else if (alloc.color() == 2) {
+      color2_bytes += alloc.size();
+    }
+  }
+  // Color 0 fell back to DEFAULT (288 bytes), while Color 2 kept FAST_MERGE
+  // (400 bytes of temps without param/root).
+  EXPECT_EQ(color0_bytes, 288);
+  EXPECT_EQ(color2_bytes, 400);
+}
+
 MATCHER(IdEq, "") {
   auto* actual = std::get<0>(arg);
   auto* expected = std::get<1>(arg);

From c46cb7c8bae511581604ded09b40caeeaa4ebce6 Mon Sep 17 00:00:00 2001
From: Marco Minutoli <marco.minutoli@amd.com>
Date: Tue, 13 Oct 2026 03:53:42 -0700
Subject: [PATCH 0680/1000] PR #49829: [ROCm] Fuse FP8 output scaling and
 conversion into hipBLASLt GEMMs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/49829

## Motivation

cuBLASLt fuses FP8 output scaling and type conversion into GEMM custom calls, but the GemmRewriter blocked this path on ROCm. This enables the same fusion on hipBLASLt. On gfx950, hipBLASLt has D-scaled FP8-output kernels only for E4M3 x E4M3 -> E4M3, so a D scale is passed only when it cannot be folded away.

Part of a series: depends on #50068, which adds `GemmBackendConfig.has_d_scale` (its commit is included here until it lands). #50069 adds an autotuner fallback for missing FP8-output kernels and is independent of this PR.

## Details

- Plain FP8 dots, and patterns with no runtime D scale, pass no D scale. With the DEFAULT epilogue, a runtime D scale is folded into the A scale, which is exact for any sign. These need no D-scaled kernels and cover all OCP/FNUZ type combinations.
- RELU, BIAS, GELU and SILU epilogues with a runtime D scale use hipBLASLt's D-scaled kernels when the types have them (OCP: E4M3 x E4M3 -> E4M3; FNUZ: all supported combinations). Otherwise the scaling and conversion stay unfused. RELU/SILU are not fused onto a GEMM that already carries a D scale, since hipBLASLt applies the activation before the scale.
- Set `GemmBackendConfig.has_d_scale` (#50068) on GEMMs that keep a D scale, so hipBLASLt queries D-scaled algorithms only for them.
- Guard FP8 output fusion against non-F16 vector biases on ROCm; hipBLASLt defaults the bias data type for FP8 matmuls, which misinterprets BF16 biases.
- If a GPU's hipBLASLt lacks a kernel this assumes, compilation fails with "No supported configs found"; #50069 adds an autotuner fallback for that case.
- Pre-existing, not changed here: `F8ScaleD` folds an output scale into `a_scale` for the RELU epilogue without checking its sign, on CUDA too. cuBLASLt and hipBLASLt apply the activation before the D scale, so this is wrong for a negative runtime scale (relu(x)·s ≠ relu(x·s) for s < 0). Options are to require a constant positive scale for the RELU fold or to drop RELU from `F8ScaleD`; I'd like the reviewer's call.

## Benchmark

MI355X, measured on a prototype of this change (to be re-measured on the final version), `--xla_gpu_enable_triton_gemm=false`, bf16 dequant, fused vs. unfused: 8192×8192×4096 1.07×, 4096×16384×4096 1.08×, DeepSeek-V3 down 16384×7168×2048 1.17×, MoE 16×16384 (K 2048→N 7168 / K 7168→N 2048) 1.29× / 1.12×. Mixed E5M2×E4M3 inputs: 1.08–1.35×. With default flags, the f32-dequant forms of these GEMMs go through hipBLASLt instead of Triton: 1.21–1.63×.

## Test Plan

On gfx950, ran `gemm_rewriter_fp8_test`, `dot_algorithm_support_test`, `gemm_rewriter_test`, `fission_backend_test`, `hipblaslt_test`, `gpu_blas_lt_test`, `rocm_blas_lt_test` and `gpublas_lt_matmul_thunk_test`. New tests cover GELU with a runtime D scale, the cases that stay unfused on ROCm, and the fission path for mixed inputs; `dot_algorithm_support_test` covers the mixed-input FP8 dots.

## Test Result

All 8 targets pass on gfx950, including the `dot_algorithm_support_test` mixed-input FP8-output cases.

Copybara import of the project:

--
abcfa3d5fa9becf9d80b3fa6ec67c7e720510e04 by Marco Minutoli <marco.minutoli@amd.com>:

[ROCm] Record in the GEMM config whether an FP8 output has a D scale

Motivation:
hipBLASLt applies a D scale only with algorithms selected while
D_SCALE_POINTER is set, and on gfx950 it has D-scaled FP8-output kernels
only for E4M3 x E4M3 -> E4M3. The algorithm query must therefore know
whether a GEMM has a D scale, rather than infer it from the output type
or from operand positions.

Details:
- Add GemmBackendConfig.has_d_scale and carry it into GemmConfig and its
  proto.
- hipBLASLt sets D_SCALE_POINTER before the algorithm query only when the
  flag is set, and ExecuteOnStream rejects a D scale that does not match
  the plan.
- On ROCm, the thunk passes the last operand as the D scale only when the
  flag is set. CUDA is unchanged.

Nothing sets the flag yet. It is used by the FP8 output fusion and by the
autotuner fallback for missing FP8-output kernels.

Test Plan:
gpu_blas_lt_test, rocm_blas_lt_test, gpublas_lt_matmul_thunk_test,
gemm_rewriter_fp8_test and hipblaslt_test pass on gfx950.

--
a5e4003302db83a367a8aa707d8364c23144337c by Marco Minutoli <marco.minutoli@amd.com>:

[ROCm] Fuse FP8 output scaling and conversion into hipBLASLt GEMMs

Motivation:
cuBLASLt fuses FP8 output scaling and type conversion into GEMM custom
calls, but the GemmRewriter blocked this path on ROCm. hipBLASLt supports
FP8 outputs, but on gfx950 it has D-scaled FP8-output kernels only for
E4M3 x E4M3 -> E4M3, so a D scale is passed only when it cannot be folded.

Details:
- Enable F8ConvertD on ROCm, including the FNUZ clamp bounds.
- Pass no D scale for plain FP8 dots, and fold a runtime D scale into the
  A scale for the DEFAULT epilogue.
- Use D-scaled kernels for RELU/BIAS/GELU/SILU epilogues only when the
  types have them; otherwise leave the scaling and conversion unfused, as
  with a matrix bias, an amax reduction or a vector bias that is not F16.
- Set GemmBackendConfig.has_d_scale on GEMMs that keep a D scale.

Test Plan:
gemm_rewriter_fp8_test, dot_algorithm_support_test, gemm_rewriter_test,
fission_backend_test, gpu_blas_lt_test, rocm_blas_lt_test and
gpublas_lt_matmul_thunk_test pass on gfx950.

--
17133d056740899daa4a173463f537c7fee08092 by Marco Minutoli <marco.minutoli@amd.com>:

[ROCm] Remove unused deps of fission_backend_test

Flagged by DWYU now that this PR touches the target.

--
7e72f2e51017a8c22485105071391c91e0b92550 by Marco Minutoli <marco.minutoli@amd.com>:

[ROCm] Remove unused deps of gpu_blas_lt

Flagged by DWYU now that this PR touches the target.

Merging this change closes #49829

PiperOrigin-RevId: 998527191
---
 .../xla/xla/backends/gpu/autotuner/BUILD      |   2 -
 .../gpu/autotuner/fission_backend_test.cc     |  49 +++++
 .../xla/xla/backends/gpu/transforms/BUILD     |   3 +
 .../backends/gpu/transforms/gemm_rewriter.cc  | 171 +++++++++++++++-
 .../gpu/transforms/gemm_rewriter_fp8_test.cc  | 188 ++++++++++++++++--
 5 files changed, 391 insertions(+), 22 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/autotuner/BUILD b/third_party/xla/xla/backends/gpu/autotuner/BUILD
index c614e7ec79e..6f74a53ae44 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/BUILD
+++ b/third_party/xla/xla/backends/gpu/autotuner/BUILD
@@ -930,7 +930,6 @@ xla_test(
         "//xla:autotune_results_proto_cc",
         "//xla:autotuning_proto_cc",
         "//xla:xla_proto_cc",
-        "//xla/backends/autotuner:backends_proto_cc",
         "//xla/backends/autotuner:codegen_backend",
         "//xla/backends/gpu/codegen/emitters:mlir_kernel_emitter",
         "//xla/backends/gpu/transforms:dot_algorithm_rewriter",
@@ -944,7 +943,6 @@ xla_test(
         "//xla/service:platform_util",
         "//xla/service/gpu:alias_info",
         "//xla/service/gpu:mlir_context_pool",
-        "//xla/service/gpu/model:gpu_performance_model_base",
         "//xla/stream_executor:device_description",
         "//xla/stream_executor:platform",
         "//xla/stream_executor:stream_executor_h",
diff --git a/third_party/xla/xla/backends/gpu/autotuner/fission_backend_test.cc b/third_party/xla/xla/backends/gpu/autotuner/fission_backend_test.cc
index 3f217f38d44..aeb7da4d8e4 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/fission_backend_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/fission_backend_test.cc
@@ -752,6 +752,55 @@ TEST_F(CublasFissionBackendTest,
   EXPECT_THAT(configs, testing::Not(testing::IsEmpty()));
 }
 
+// A scaled E5M2 x E4M3 dot with an E4M3 output needs no D-scaled kernels, so
+// hipBLASLt fuses its FP8 output instead of leaving a much slower Triton GEMM.
+TEST_F(CublasFissionBackendTest, HipblasLtConfigsForMixedFp8OutputFusion) {
+  const auto* rocm_cc =
+      device_description_.gpu_compute_capability().rocm_compute_capability();
+  if (rocm_cc == nullptr || !rocm_cc->has_ocp_fp8_support()) {
+    GTEST_SKIP() << "Requires a ROCm GPU with OCP FP8 support.";
+  }
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(R"(
+  computation {
+    x = f8e5m2[64,128]{1,0} parameter(0)
+    y = f8e4m3fn[128,64]{1,0} parameter(1)
+    x_scale = f32[] parameter(2)
+    y_scale = f32[] parameter(3)
+    z_scale = f32[] parameter(4)
+    x_f32 = f32[64,128]{1,0} convert(x)
+    y_f32 = f32[128,64]{1,0} convert(y)
+    x_scale_bcast = f32[64,128]{1,0} broadcast(x_scale), dimensions={}
+    y_scale_bcast = f32[128,64]{1,0} broadcast(y_scale), dimensions={}
+    x_unscaled = f32[64,128]{1,0} multiply(x_f32, x_scale_bcast)
+    y_unscaled = f32[128,64]{1,0} multiply(y_f32, y_scale_bcast)
+    dot = f32[64,64]{1,0} dot(x_unscaled, y_unscaled), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    z_scale_bcast = f32[64,64]{1,0} broadcast(z_scale), dimensions={}
+    dot_scaled = f32[64,64]{1,0} divide(dot, z_scale_bcast)
+    c1 = f32[] constant(-448)
+    c1_bcast = f32[64,64]{1,0} broadcast(c1), dimensions={}
+    c2 = f32[] constant(448)
+    c2_bcast = f32[64,64]{1,0} broadcast(c2), dimensions={}
+    dot_clamped = f32[64,64]{1,0} clamp(c1_bcast, dot_scaled, c2_bcast)
+    ROOT out = f8e4m3fn[64,64]{1,0} convert(dot_clamped)
+  }
+
+  ENTRY main {
+    x = f8e5m2[64,128]{1,0} parameter(0)
+    y = f8e4m3fn[128,64]{1,0} parameter(1)
+    x_scale = f32[] parameter(2)
+    y_scale = f32[] parameter(3)
+    z_scale = f32[] parameter(4)
+    ROOT fusion = f8e4m3fn[64,64]{1,0} fusion(x, y, x_scale, y_scale, z_scale),
+      kind=kCustom, calls=computation,
+      backend_config={"fusion_backend_config":{"kind":"__triton_gemm"}}
+  })"));
+  ASSERT_OK_AND_ASSIGN(std::vector<std::unique_ptr<BackendConfig>> configs,
+                       fission_backend_->GetSupportedConfigs(
+                           *module->entry_computation()->root_instruction()));
+  EXPECT_THAT(configs, Not(IsEmpty()));
+}
+
 TEST_F(CublasFissionBackendTest,
        GetSupportedConfigsWithEstimatesComposesPrologueEpilogueAndGemm) {
   if (IsRocm(stream_executor_)) {
diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD
index 648e57ec637..6aa224982d8 100644
--- a/third_party/xla/xla/backends/gpu/transforms/BUILD
+++ b/third_party/xla/xla/backends/gpu/transforms/BUILD
@@ -2280,6 +2280,8 @@ xla_test(
         ":gemm_rewriter",
         ":gemm_rewriter_test_lib",
         "//xla:error_spec",
+        "//xla:literal",
+        "//xla:literal_util",
         "//xla:xla_data_proto_cc",
         "//xla:xla_proto_cc",
         "//xla/hlo/ir:hlo",
@@ -2293,6 +2295,7 @@ xla_test(
         "//xla/stream_executor:device_description",
         "//xla/stream_executor:semantic_version",
         "//xla/tests:hlo_pjrt_interpreter_reference_mixin",
+        "//xla/tests:test_utils",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
diff --git a/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter.cc
index 6d376b5b658..d0aefdf752f 100644
--- a/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter.cc
@@ -1610,14 +1610,57 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
     return absl::OkStatus();
   }
 
+  // Returns whether the BLAS library supports an FP8 output of type `d_type`
+  // for the input types of `gemm`. On ROCm, C has the type of D (see
+  // matmul_utils.cc), and hipBLASLt supports E4M3 x E4M3 -> E4M3 and mixed
+  // E4M3/E5M2 inputs -> E5M2, plus mixed inputs -> E4M3 for OCP types.
+  absl::StatusOr<bool> IsSupportedF8OutputType(const HloInstruction& gemm,
+                                               PrimitiveType d_type) {
+    if (gpu_version_.IsCuda()) {
+      return d_type == F8E4M3FN || d_type == F8E5M2;
+    }
+    if (!gpu_version_.IsRocm()) {
+      return false;
+    }
+    ABSL_ASSIGN_OR_RETURN(auto rocm_compute_capability,
+                     GetRocmComputeCapability(gpu_version_));
+    PrimitiveType a_type = gemm.operand(0)->shape().element_type();
+    PrimitiveType b_type = gemm.operand(1)->shape().element_type();
+    bool has_e5m2_input = a_type == F8E5M2 || b_type == F8E5M2 ||
+                          a_type == F8E5M2FNUZ || b_type == F8E5M2FNUZ;
+    switch (d_type) {
+      case F8E4M3FN:
+        return rocm_compute_capability.has_ocp_fp8_support();
+      case F8E5M2:
+        return rocm_compute_capability.has_ocp_fp8_support() && has_e5m2_input;
+      case F8E4M3FNUZ:
+        return rocm_compute_capability.has_nanoo_fp8_support() &&
+               !has_e5m2_input;
+      case F8E5M2FNUZ:
+        return rocm_compute_capability.has_nanoo_fp8_support() &&
+               has_e5m2_input;
+      default:
+        return false;
+    }
+  }
+
+  // Returns whether hipBLASLt can apply a D scale to an FP8 output of type
+  // `d_type` for the input types of `gemm`: only E4M3 x E4M3 -> E4M3 for OCP
+  // types, and every IsSupportedF8OutputType combination for FNUZ types.
+  bool SupportsF8OutputDScale(const HloInstruction& gemm,
+                              PrimitiveType d_type) const {
+    if (d_type != F8E4M3FN && d_type != F8E5M2) {
+      return true;
+    }
+    return gemm.operand(0)->shape().element_type() == F8E4M3FN &&
+           gemm.operand(1)->shape().element_type() == F8E4M3FN &&
+           d_type == F8E4M3FN;
+  }
+
   absl::Status F8ConvertD(HloInstruction* instr, HloInstruction* existing_gemm,
                           HloInstruction* d_scale, HloInstruction* clamp_lower,
                           HloInstruction* clamp_upper,
                           bool mult_scale = false) {
-    // TODO: add ROCm support to this fusion pattern
-    if (gpu_version_.IsRocm()) {
-      return absl::OkStatus();
-    }
     // Verify the data types and the operands of clamp.
     if (instr->shape().element_type() == F8E4M3FN) {
       if (!clamp_lower->literal().IsAllFloat(static_cast<float>(
@@ -1633,10 +1676,36 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
               std::numeric_limits<tsl::float8_e5m2>::max()))) {
         return absl::OkStatus();
       }
+    } else if (instr->shape().element_type() == F8E4M3FNUZ) {
+      if (!clamp_lower->literal().IsAllFloat(static_cast<float>(
+              std::numeric_limits<tsl::float8_e4m3fnuz>::lowest())) ||
+          !clamp_upper->literal().IsAllFloat(static_cast<float>(
+              std::numeric_limits<tsl::float8_e4m3fnuz>::max()))) {
+        return absl::OkStatus();
+      }
+    } else if (instr->shape().element_type() == F8E5M2FNUZ) {
+      if (!clamp_lower->literal().IsAllFloat(static_cast<float>(
+              std::numeric_limits<tsl::float8_e5m2fnuz>::lowest())) ||
+          !clamp_upper->literal().IsAllFloat(static_cast<float>(
+              std::numeric_limits<tsl::float8_e5m2fnuz>::max()))) {
+        return absl::OkStatus();
+      }
     } else {
       return absl::OkStatus();
     }
 
+    ABSL_ASSIGN_OR_RETURN(
+        bool is_supported_output_type,
+        IsSupportedF8OutputType(*existing_gemm, instr->shape().element_type()));
+    if (!is_supported_output_type) {
+      VLOG(1) << "The conversion of the result of "
+              << existing_gemm->ToShortString() << " to "
+              << PrimitiveType_Name(instr->shape().element_type())
+              << " is not fused into the FP8 Custom Call because the output "
+                 "type is not supported for its input types.";
+      return absl::OkStatus();
+    }
+
     if (d_scale && !ShapeUtil::IsScalar(d_scale->shape())) {
       return absl::OkStatus();
     }
@@ -1688,6 +1757,16 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
       if (!reduce_damax) {
         return absl::OkStatus();
       }
+      // hipBLASLt ships no tuned AMAX_D kernels on current ROCm releases (none
+      // for gfx950; slow ones for gfx942), so keep the amax reduction and the
+      // FP8 conversion unfused.
+      if (gpu_version_.IsRocm()) {
+        VLOG(1) << "The scaling and conversion of the result of "
+                << existing_gemm->ToShortString()
+                << " is not fused into the FP8 Custom Call because hipBLASLt "
+                   "does not support the amax epilogue on this platform.";
+        return absl::OkStatus();
+      }
     } else if (gemm_users.size() > 2) {
       return absl::OkStatus();
     }
@@ -1707,12 +1786,82 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
                    "matrix bias with element type other than BF16 or F16.";
         return absl::OkStatus();
       }
+      if (gpu_version_.IsRocm()) {
+        VLOG(1) << "The scaling and conversion of the result of "
+                << existing_gemm->ToShortString()
+                << " is not fused into the FP8 Custom Call because "
+                   "hipBLASLt does not support matrix bias with FP8 output.";
+        return absl::OkStatus();
+      }
       // Turn off the output to operand aliasing, since the fp8 output and
       // bf16/fp16 bias have different sizes.
       xla::Cast<HloCustomCallInstruction>(existing_gemm)
           ->set_output_to_operand_aliasing({});
     }
 
+    // hipBLASLt reads the vector bias of a GEMM with FP8 output as F16.
+    if (gpu_version_.IsRocm()) {
+      ABSL_ASSIGN_OR_RETURN(
+          bool has_vector_bias,
+          gpublas_lt::EpilogueAddsVectorBias(gemm_backend_config.epilogue()));
+      if (has_vector_bias &&
+          existing_gemm->operands().back()->shape().element_type() != F16) {
+        VLOG(1) << "The scaling and conversion of the result of "
+                << existing_gemm->ToShortString()
+                << " is not fused into the FP8 Custom Call because hipBLASLt "
+                   "only supports an F16 vector bias with FP8 output.";
+        return absl::OkStatus();
+      }
+    }
+
+    // On ROCm, hipBLASLt applies the A and B scales with alpha before the
+    // epilogue and the D scale after it, so for the DEFAULT epilogue (beta is 0
+    // here) the D scale is folded into the A scale. Other epilogues need
+    // kernels that apply a D scale; without them, the output stays unfused.
+    if (gpu_version_.IsRocm()) {
+      if (d_scale && d_scale->IsConstant() &&
+          d_scale->literal().IsAllFloat(1.0f)) {
+        d_scale = nullptr;
+      }
+      std::vector<HloInstruction*> new_operands(
+          existing_gemm->operands().begin(), existing_gemm->operands().end());
+      GpuBackendConfig new_gpu_backend_config = gpu_backend_config;
+      if (d_scale) {
+        bool fold_d_scale =
+            gemm_backend_config.epilogue() == GemmBackendConfig::DEFAULT;
+        if (!fold_d_scale &&
+            !SupportsF8OutputDScale(*existing_gemm,
+                                    instr->shape().element_type())) {
+          VLOG(1) << "The scaling and conversion of the result of "
+                  << existing_gemm->ToShortString()
+                  << " is not fused into the FP8 Custom Call because hipBLASLt "
+                     "does not support a D scale for its types.";
+          return absl::OkStatus();
+        }
+        ABSL_ASSIGN_OR_RETURN(HloInstruction * d_scale_f32,
+                         InvertAndConvertScalar(d_scale, !mult_scale));
+        if (!fold_d_scale) {
+          new_operands.push_back(d_scale_f32);
+          new_gpu_backend_config.mutable_gemm_backend_config()->set_has_d_scale(
+              true);
+        } else if (new_operands[2]->IsConstant() &&
+                   new_operands[2]->literal().IsAllFloat(1.0f)) {
+          new_operands[2] = d_scale_f32;
+        } else {
+          ABSL_ASSIGN_OR_RETURN(new_operands[2],
+                           MakeBinaryHlo(HloOpcode::kMultiply, new_operands[2],
+                                         d_scale_f32, &d_scale->metadata()));
+        }
+      }
+      std::unique_ptr<HloInstruction> new_gemm =
+          existing_gemm->CloneWithNewOperands(instr->shape(), new_operands);
+      ABSL_RETURN_IF_ERROR(new_gemm->set_backend_config(new_gpu_backend_config));
+      ABSL_RETURN_IF_ERROR(ReplaceWithNewInstruction(instr, std::move(new_gemm)));
+      VLOG(1) << "Conversion" << (d_scale ? " and scaling" : "")
+              << " fused into FP8 GEMM.";
+      return absl::OkStatus();
+    }
+
     // If necessary, invert the scaling factor of D and convert to F32. When no
     // scaling factor was captured, set the factor to one.
     if (d_scale) {
@@ -1976,6 +2125,12 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
     if (!SupportsEpilogueFusion(gemm->shape().element_type())) {
       return false;
     }
+    // On ROCm, hipBLASLt reads the vector bias of a GEMM with FP8 output as
+    // F16, and a D scale operand must stay last.
+    if (gpu_version_.IsRocm() &&
+        primitive_util::IsF8Type(gemm->shape().element_type())) {
+      return false;
+    }
 
     HloInstruction* bias = broadcast->mutable_operand(0);
 
@@ -2119,6 +2274,10 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
 
     ABSL_ASSIGN_OR_RETURN(auto gpu_config, gemm->backend_config<GpuBackendConfig>());
     GemmBackendConfig& config = GetMutableGemmBackendConfig(gpu_config);
+    // The activation follows the D scale, but hipBLASLt applies it before.
+    if (config.has_d_scale()) {
+      return absl::OkStatus();
+    }
     if (config.epilogue() == GemmBackendConfig::DEFAULT) {
       config.set_epilogue(GemmBackendConfig::RELU);
     } else if (config.epilogue() == GemmBackendConfig::BIAS) {
@@ -2258,6 +2417,10 @@ class GemmRewriterVisitor : public DfsHloRewriteVisitor {
 
     ABSL_ASSIGN_OR_RETURN(auto gpu_config, gemm->backend_config<GpuBackendConfig>());
     GemmBackendConfig& config = GetMutableGemmBackendConfig(gpu_config);
+    // The activation follows the D scale, but hipBLASLt applies it before.
+    if (config.has_d_scale()) {
+      return absl::OkStatus();
+    }
 
     if (config.epilogue() == GemmBackendConfig::DEFAULT) {
       config.set_epilogue(GemmBackendConfig::SILU);
diff --git a/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter_fp8_test.cc b/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter_fp8_test.cc
index 9a8d9ada204..1faf3edf36f 100644
--- a/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter_fp8_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/gemm_rewriter_fp8_test.cc
@@ -39,12 +39,15 @@ limitations under the License.
 #include "xla/hlo/testlib/pattern_matcher_gmock.h"
 #include "xla/hlo/testlib/test.h"
 #include "xla/hlo/testlib/verified_hlo_module.h"
+#include "xla/literal.h"
+#include "xla/literal_util.h"
 #include "xla/service/device_assignment.h"
 #include "xla/service/hlo_module_config.h"
 #include "xla/service/pattern_matcher.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/semantic_version.h"
 #include "xla/tests/hlo_pjrt_interpreter_reference_mixin.h"
+#include "xla/tests/test_utils.h"
 #include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
 
@@ -122,6 +125,20 @@ class ParameterizedFp8GemmRewriteTest
     EXPECT_EQ(call->custom_call_target(), "__cublas$lt$matmul$f8");
   }
 
+  // Runs `hlo_text` with the F32 scales in parameters 2-4 set to 4, 8 and 0.25.
+  void RunAndCompareWithScales(absl::string_view hlo_text) {
+    ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text));
+    ASSERT_OK_AND_ASSIGN(std::vector<Literal> args,
+                         MakeFakeArguments(module.get()));
+    args[2] = LiteralUtil::CreateR0<float>(4.0f);
+    args[3] = LiteralUtil::CreateR0<float>(8.0f);
+    args[4] = LiteralUtil::CreateR0<float>(0.25f);
+    // The relative tolerance allows one E4M3 ULP (2^-3).
+    EXPECT_TRUE(RunAndCompare(std::move(module),
+                              LiteralUtil::MakePointers(args),
+                              ErrorSpec{1e-2, 1.3e-1}));
+  }
+
   void MatchOptimizedHlo(absl::string_view hlo, const absl::string_view pattern,
                          bool print_operand_shape = false) {
     GemmRewriteTestBase::MatchOptimizedHlo(
@@ -1431,9 +1448,9 @@ TEST_F(ParameterizedFp8GemmRewriteTest, UnscaledABScaledDF8) {
       z_scale_bcast = f32[16,16] broadcast(z_scale), dimensions={}
       dot_a = f32[16,16] dot(x, y), lhs_contracting_dims={1}, rhs_contracting_dims={0}
       dot_a_scaled = f32[16,16] divide(dot_a, z_scale_bcast)
-      c1 = f32[] constant(-448.)
+      c1 = f32[] constant(-<<F8E4M3_AMAX>>)
       c1_bcast = f32[16,16] broadcast(c1), dimensions={}
-      c2 = f32[] constant(448.)
+      c2 = f32[] constant(<<F8E4M3_AMAX>>)
       c2_bcast = f32[16,16] broadcast(c2), dimensions={}
       dot_a_clamped = f32[16,16] clamp(c1_bcast, dot_a_scaled, c2_bcast)
       ROOT dot_a_f8 = <<F8E4M3>>[16,16] convert(dot_a_clamped)
@@ -1457,7 +1474,7 @@ TEST_F(ParameterizedFp8GemmRewriteTest, UnscaledABScaledDF8) {
 ; CHECK-NEXT:    [[C1:%[^ ]+]] = f32[] constant(1)
 ; CHECK-PTX-NEXT:    [[C2:%[^ ]+]] = f32[] constant(1)
 ; CHECK-PTX-NEXT:    [[OUT:%[^ ]+]] = (<<F8E4M3>>[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[P2_INV]], [[C1]], [[C2]]),
-; CHECK-GCN-NEXT:    [[OUT:%[^ ]+]] = (f32[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[P2_INV]], [[C1]]),
+; CHECK-GCN-NEXT:    [[OUT:%[^ ]+]] = (<<F8E4M3>>[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[P2_INV]], [[C1]]),
 ; CHECK:           custom_call_target="__cublas$lt$matmul$f8",
 ; CHECK:           backend_config={
 ; CHECK-DAG:         "alpha_real":1
@@ -1668,12 +1685,17 @@ TEST_F(ParameterizedFp8GemmRewriteTest, ScaledABScaledDF8) {
 ; CHECK-NEXT:    [[P1:%[^ ]+]] = <<F8E4M3>>[32,16]{1,0} parameter(1)
 ; CHECK-NEXT:    [[P1_TRANSPOSE:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} transpose([[P1]]), dimensions={1,0}
 ; CHECK-NEXT:    [[P2:%[^ ]+]] = f32[] parameter(2)
-; CHECK-NEXT:    [[P3:%[^ ]+]] = f32[] parameter(3)
+; CHECK-PTX-NEXT:    [[P3:%[^ ]+]] = f32[] parameter(3)
 ; CHECK-PTX-NEXT:    [[C2:%[^ ]+]] = f32[] constant(1)
 ; CHECK-PTX-NEXT:    [[P4:%[^ ]+]] = f32[] parameter(4)
 ; CHECK-PTX-NEXT:    [[P4_INV:%[^ ]+]] = f32[] divide([[C2]], [[P4]])
 ; CHECK-PTX-NEXT:    [[OUT:%[^ ]+]] = (<<F8E4M3>>[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[P2]], [[P3]], [[P4_INV]]),
-; CHECK-GCN-NEXT:    [[OUT:%[^ ]+]] = (f32[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[P2]], [[P3]]),
+; CHECK-GCN-NEXT:    [[C2:%[^ ]+]] = f32[] constant(1)
+; CHECK-GCN-NEXT:    [[P4:%[^ ]+]] = f32[] parameter(4)
+; CHECK-GCN-NEXT:    [[P4_INV:%[^ ]+]] = f32[] divide([[C2]], [[P4]])
+; CHECK-GCN-NEXT:    [[A_SCALE:%[^ ]+]] = f32[] multiply([[P2]], [[P4_INV]])
+; CHECK-GCN-NEXT:    [[P3:%[^ ]+]] = f32[] parameter(3)
+; CHECK-GCN-NEXT:    [[OUT:%[^ ]+]] = (<<F8E4M3>>[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[A_SCALE]], [[P3]]),
 ; CHECK:           custom_call_target="__cublas$lt$matmul$f8",
 ; CHECK:           backend_config={
 ; CHECK-DAG:         "alpha_real":1
@@ -1689,6 +1711,7 @@ TEST_F(ParameterizedFp8GemmRewriteTest, ScaledABScaledDF8) {
 ; CHECK-DAG:           "operand_precision":["DEFAULT","DEFAULT"]
 ; CHECK-DAG:         }
 ; CHECK-DAG:         "epilogue":"DEFAULT"
+; CHECK-GCN-NOT:     "has_d_scale"
 ; CHECK:           }
       )");
 }
@@ -1778,11 +1801,10 @@ TEST_F(ParameterizedFp8GemmRewriteTest, ScaledABScaledDReluActivationF8) {
 ; CHECK-NEXT:    [[P1_TRANSPOSE:%[^ ]+]] = <<F8E4M3>>[16,32]{1,0} transpose([[P1]]), dimensions={1,0}
 ; CHECK-NEXT:    [[P2:%[^ ]+]] = f32[] parameter(2)
 ; CHECK-NEXT:    [[P3:%[^ ]+]] = f32[] parameter(3)
-; CHECK-PTX-NEXT:    [[C2:%[^ ]+]] = f32[] constant(1)
-; CHECK-PTX-NEXT:    [[P4:%[^ ]+]] = f32[] parameter(4)
-; CHECK-PTX-NEXT:    [[P4_INV:%[^ ]+]] = f32[] divide([[C2]], [[P4]])
-; CHECK-PTX-NEXT:    [[OUT:%[^ ]+]] = (<<F8E4M3>>[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[P2]], [[P3]], [[P4_INV]]),
-; CHECK-GCN-NEXT:    [[OUT:%[^ ]+]] = (f32[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[P2]], [[P3]]),
+; CHECK-NEXT:    [[C2:%[^ ]+]] = f32[] constant(1)
+; CHECK-NEXT:    [[P4:%[^ ]+]] = f32[] parameter(4)
+; CHECK-NEXT:    [[P4_INV:%[^ ]+]] = f32[] divide([[C2]], [[P4]])
+; CHECK-NEXT:    [[OUT:%[^ ]+]] = (<<F8E4M3>>[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[P2]], [[P3]], [[P4_INV]]),
 ; CHECK:           custom_call_target="__cublas$lt$matmul$f8",
 ; CHECK:           backend_config={
 ; CHECK-DAG:         "alpha_real":1
@@ -1798,6 +1820,7 @@ TEST_F(ParameterizedFp8GemmRewriteTest, ScaledABScaledDReluActivationF8) {
 ; CHECK-DAG:           "operand_precision":["DEFAULT","DEFAULT"]
 ; CHECK-DAG:         }
 ; CHECK-DAG:         "epilogue":"RELU"
+; CHECK-GCN-DAG:     "has_d_scale":true
 ; CHECK:           }
       )");
 }
@@ -1930,12 +1953,11 @@ TEST_F(ParameterizedFp8GemmRewriteTest, ScaledABScaledDVectorBiasF8) {
 ; CHECK-NEXT:    [[P3:%[^ ]+]] = f16[] parameter(4)
 ; CHECK-NEXT:    [[CV1:%[^ ]+]] = f32[] convert([[P3]])
 ; CHECK-NEXT:    [[VB:%[^ ]+]] = f16[16]{0} parameter(2)
-; CHECK-PTX-NEXT:    [[C2:%[^ ]+]] = f16[] constant(1)
-; CHECK-PTX-NEXT:    [[P4:%[^ ]+]] = f16[] parameter(5)
-; CHECK-PTX-NEXT:    [[DV:%[^ ]+]] = f16[] divide([[C2]], [[P4]])
-; CHECK-PTX-NEXT:    [[CV2:%[^ ]+]] = f32[] convert([[DV]])
-; CHECK-PTX:     [[OUT:%[^ ]+]] = (<<F8E4M3>>[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[CV]], [[CV1]], [[VB]], /*index=5*/[[CV2]]),
-; CHECK-GCN:     [[OUT:%[^ ]+]] = (f16[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[CV]], [[CV1]], [[VB]]),
+; CHECK-NEXT:    [[C2:%[^ ]+]] = f16[] constant(1)
+; CHECK-NEXT:    [[P4:%[^ ]+]] = f16[] parameter(5)
+; CHECK-NEXT:    [[DV:%[^ ]+]] = f16[] divide([[C2]], [[P4]])
+; CHECK-NEXT:    [[CV2:%[^ ]+]] = f32[] convert([[DV]])
+; CHECK:     [[OUT:%[^ ]+]] = (<<F8E4M3>>[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call([[P0]], [[P1_TRANSPOSE]], [[CV]], [[CV1]], [[VB]], /*index=5*/[[CV2]]),
 ; CHECK:           custom_call_target="__cublas$lt$matmul$f8",
 ; CHECK:           backend_config={
 ; CHECK-DAG:         "alpha_real":1
@@ -1951,6 +1973,7 @@ TEST_F(ParameterizedFp8GemmRewriteTest, ScaledABScaledDVectorBiasF8) {
 ; CHECK-DAG:           "operand_precision":["DEFAULT","DEFAULT"]
 ; CHECK-DAG:         }
 ; CHECK-DAG:         "epilogue":"BIAS"
+; CHECK-GCN-DAG:     "has_d_scale":true
 ; CHECK:           }
       )");
 }
@@ -2522,6 +2545,139 @@ TEST_F(ParameterizedFp8GemmRewriteTest,
       )");
 }
 
+// On ROCm, hipBLASLt reads the vector bias of a GEMM with FP8 output as F16,
+// so an FP8 vector bias is not fused into it.
+TEST_F(ParameterizedFp8GemmRewriteTest,
+       UnscaledABUnscaledDVectorBiasF8OutputNotFused) {
+  if (!IsRocm()) {
+    GTEST_SKIP() << "ROCm-specific.";
+  }
+  const char* hlo_text = R"(
+    HloModule test
+
+    ENTRY test {
+      x = <<F8E4M3>>[16,32] parameter(0)
+      y = <<F8E4M3>>[32,16] parameter(1)
+      b = <<F8E4M3>>[16] parameter(2)
+      b_bcast = <<F8E4M3>>[16,16] broadcast(b), dimensions={1}
+      dot_a = <<F8E4M3>>[16,16] dot(x, y), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+      ROOT out = <<F8E4M3>>[16,16] add(dot_a, b_bcast)
+    }
+)";
+
+  RunAndFilecheckHloRewrite(
+      hlo_text,
+      GemmRewriter(CudaHopperOrRocmCapability(), GetToolkitVersion(),
+                   GemmRewriterOptions{GemmRewriterOptions::DType::kFp8Only}),
+      R"(
+; CHECK:         custom-call({{[^,]+}}, {{[^,]+}}, {{[^,]+}}, {{[^,)]+}}), custom_call_target="__cublas$lt$matmul$f8"
+; CHECK-SAME:      "epilogue":"DEFAULT"
+; CHECK:         add(
+      )");
+}
+
+// A scaled FP8 GEMM, where <<ACTIVATION>> computes <<ACT>> from `dot`.
+constexpr absl::string_view kScaledABScaledDActivationF8Hlo = R"(
+    ENTRY test {
+      x = <<F8E4M3>>[16,32] parameter(0)
+      y = <<F8E4M3>>[32,16] parameter(1)
+      x_f32 = f32[16,32] convert(x)
+      y_f32 = f32[32,16] convert(y)
+      x_scale = f32[] parameter(2)
+      y_scale = f32[] parameter(3)
+      z_scale = f32[] parameter(4)
+      x_scale_bcast = f32[16,32] broadcast(x_scale), dimensions={}
+      y_scale_bcast = f32[32,16] broadcast(y_scale), dimensions={}
+      z_scale_bcast = f32[16,16] broadcast(z_scale), dimensions={}
+      x_unscaled = f32[16,32] multiply(x_f32, x_scale_bcast)
+      y_unscaled = f32[32,16] multiply(y_f32, y_scale_bcast)
+      dot = f32[16,16] dot(x_unscaled, y_unscaled), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+      <<ACTIVATION>>
+      act_scaled = f32[16,16] divide(<<ACT>>, z_scale_bcast)
+      c1 = f32[] constant(-<<F8E4M3_AMAX>>)
+      c1_bcast = f32[16,16] broadcast(c1), dimensions={}
+      c2 = f32[] constant(<<F8E4M3_AMAX>>)
+      c2_bcast = f32[16,16] broadcast(c2), dimensions={}
+      act_clamped = f32[16,16] clamp(c1_bcast, act_scaled, c2_bcast)
+      ROOT out = <<F8E4M3>>[16,16] convert(act_clamped)
+    }
+)";
+
+constexpr absl::string_view kGelu = R"(
+      mul.0 = f32[16,16] multiply(dot, dot)
+      mul.1 = f32[16,16] multiply(dot, mul.0)
+      const.0 = f32[] constant(0.044715)
+      bcast.0 = f32[16,16] broadcast(const.0), dimensions={}
+      mul.2 = f32[16,16] multiply(mul.1, bcast.0)
+      add.0 = f32[16,16] add(dot, mul.2)
+      const.1 = f32[] constant(0.797884583)
+      bcast.1 = f32[16,16] broadcast(const.1), dimensions={}
+      mul.3 = f32[16,16] multiply(add.0, bcast.1)
+      tanh = f32[16,16] tanh(mul.3)
+      const.2 = f32[] constant(1)
+      bcast.2 = f32[16,16] broadcast(const.2), dimensions={}
+      add.2 = f32[16,16] add(tanh, bcast.2)
+      const.3 = f32[] constant(0.5)
+      bcast.3 = f32[16,16] broadcast(const.3), dimensions={}
+      mul.4 = f32[16,16] multiply(add.2, bcast.3)
+      act = f32[16,16] multiply(dot, mul.4))";
+
+// On ROCm, the D scale of a GELU epilogue cannot be folded into the A scale,
+// so it is passed to a hipBLASLt kernel that applies it.
+TEST_F(ParameterizedFp8GemmRewriteTest, ScaledABScaledDGeluF8) {
+  if (!IsRocm()) {
+    GTEST_SKIP() << "ROCm-specific.";
+  }
+  std::string hlo_text =
+      absl::StrReplaceAll(kScaledABScaledDActivationF8Hlo,
+                          {{"<<ACTIVATION>>", kGelu}, {"<<ACT>>", "act"}});
+  RunAndCompareWithScales(hlo_text);
+  MatchOptimizedHlo(hlo_text, R"(
+; CHECK:         (<<F8E4M3>>[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call({{[^,]+}}, {{[^,]+}}, {{[^,]+}}, {{[^,]+}}, {{[^,)]+}}),
+; CHECK-SAME:      custom_call_target="__cublas$lt$matmul$f8"
+; CHECK-SAME:      "epilogue":"GELU"
+; CHECK-SAME:      "has_d_scale":true
+      )");
+}
+
+// On ROCm, the output scaling and conversion stay unfused when hipBLASLt has no
+// kernel for them: with a matrix bias, a vector bias that is not F16, or an
+// E5M2 output of E4M3 inputs.
+TEST_F(ParameterizedFp8GemmRewriteTest, ScaledABScaledDNotFusedOnRocm) {
+  if (!IsRocm()) {
+    GTEST_SKIP() << "ROCm-specific.";
+  }
+  constexpr absl::string_view kMatrixBias = R"(
+      b = f32[16,16] parameter(5)
+      act = f32[16,16] add(dot, b))";
+  constexpr absl::string_view kF32VectorBias = R"(
+      b = f32[16] parameter(5)
+      b_bcast = f32[16,16] broadcast(b), dimensions={1}
+      act = f32[16,16] add(dot, b_bcast))";
+  struct Case {
+    absl::string_view activation, act, root, amax;
+  };
+  for (const Case& c :
+       {Case{kMatrixBias, "act", "ROOT out = <<F8E4M3>>", "<<F8E4M3_AMAX>>"},
+        Case{kF32VectorBias, "act", "ROOT out = <<F8E4M3>>", "<<F8E4M3_AMAX>>"},
+        Case{"", "dot", "ROOT out = <<F8E5M2>>", "57344."}}) {
+    std::string hlo_text = absl::StrReplaceAll(
+        kScaledABScaledDActivationF8Hlo, {{"<<ACTIVATION>>", c.activation},
+                                          {"<<ACT>>", c.act},
+                                          {"ROOT out = <<F8E4M3>>", c.root},
+                                          {"<<F8E4M3_AMAX>>", c.amax}});
+    SCOPED_TRACE(hlo_text);
+    RunAndFilecheckHloRewrite(
+        hlo_text,
+        GemmRewriter(CudaHopperOrRocmCapability(), GetToolkitVersion(),
+                     GemmRewriterOptions{GemmRewriterOptions::DType::kFp8Only}),
+        R"(
+; CHECK:         = (f32[16,16]{1,0}, s8[{{[0-9]+}}]{0}) custom-call(
+; CHECK-SAME:      custom_call_target="__cublas$lt$matmul$f8"
+      )");
+  }
+}
+
 TEST_F(ParameterizedFp8GemmRewriteTest, ScaledABScaledDWithDAmaxF8) {
   const char* hlo_text = R"(
     HloModule test

From 10e97716c9749a6f8ea464c90d2085dc0bdf1553 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Tue, 13 Oct 2026 20:56:18 +0530
Subject: [PATCH 0681/1000] Address review: explicit is not None check and
 exception-restore test

- _experimental_default_mesh: use `if previous_default is not None` instead
  of relying on Mesh's implicit truthiness, for consistency with
  _default_layout's `is None` check.
- Add test_default_mesh_restore_on_exception to APITest, verifying the
  default mesh is correctly restored after an exception inside
  api.default_mesh(), including the nested-scope case.
---
 tensorflow/dtensor/python/dtensor_device.py |  2 +-
 tensorflow/dtensor/python/tests/api_test.py | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/tensorflow/dtensor/python/dtensor_device.py b/tensorflow/dtensor/python/dtensor_device.py
index 5ee21f263c9..26de0af2900 100644
--- a/tensorflow/dtensor/python/dtensor_device.py
+++ b/tensorflow/dtensor/python/dtensor_device.py
@@ -368,7 +368,7 @@ class DTensorDevice(object):
       # mesh and `_current_default_mesh` are left stuck on `mesh` for the
       # rest of the program.
       _pywrap_dtensor_device.ExperimentalClearDefaultMesh(self._device_info)
-      if previous_default:
+      if previous_default is not None:
         _pywrap_dtensor_device.ExperimentalSetDefaultMesh(
             self._device_info,
             previous_default.to_string().encode("utf-8"))
diff --git a/tensorflow/dtensor/python/tests/api_test.py b/tensorflow/dtensor/python/tests/api_test.py
index 72310866514..d676f60e9a3 100644
--- a/tensorflow/dtensor/python/tests/api_test.py
+++ b/tensorflow/dtensor/python/tests/api_test.py
@@ -300,6 +300,24 @@ class APITest(test_util.DTensorBaseTest):
 
     self.assertIsNone(api.get_default_mesh())
 
+  def test_default_mesh_restore_on_exception(self):
+    self.assertIsNone(api.get_default_mesh())
+    with self.assertRaisesRegex(RuntimeError, "simulated error"):
+      with api.default_mesh(self.mesh):
+        self.assertEqual(api.get_default_mesh(), self.mesh)
+        raise RuntimeError("simulated error")
+    self.assertIsNone(api.get_default_mesh())
+
+    # Verify nested scope restoration on exception
+    with api.default_mesh(self.mesh):
+      self.assertEqual(api.get_default_mesh(), self.mesh)
+      with self.assertRaisesRegex(RuntimeError, "inner simulated error"):
+        with api.default_mesh(self.mesh.host_mesh()):
+          self.assertEqual(api.get_default_mesh(), self.mesh.host_mesh())
+          raise RuntimeError("inner simulated error")
+      self.assertEqual(api.get_default_mesh(), self.mesh)
+    self.assertIsNone(api.get_default_mesh())
+
 
 if __name__ == '__main__':
   test.main()

From 0f466e055fb6a1faa15d01fb65421ff06245cc49 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Tue, 13 Oct 2026 20:59:01 +0530
Subject: [PATCH 0682/1000] Fix index_matrix batch-offset bug; remove
 accidental file duplication

The previous push on this branch accidentally duplicated the entire file
starting mid-line in element_spec's return statement, shadowing the real
fix with the original buggy definitions further down and producing a
fatal SyntaxError. Reverted to the clean upstream file and reapplied only
the intended minimal change:

- _index_matrix: index elem_spec.shape[layout_idx + 1] instead of
  elem_spec.shape[layout_idx], since elem_spec's shape includes the
  leading batch dimension at index 0 while layout_idx enumerates the
  non-batch sharding specs starting from 0.
- DTensorDataset.__init__: build _index_matrices from
  self._batched_dataset.element_spec (always correctly batched/rank-
  aligned for both dataset_already_batched=True and False) instead of
  the unbatched dataset.element_spec.

Also adds testIndexMatrixForBatchedElementSpec to InputUtilHelpersTest,
directly exercising _index_matrix with a batched element spec against
two non-batch sharded layouts, to cover the dataset_already_batched=True
path that previously had no test coverage.
---
 tensorflow/dtensor/python/input_util.py       | 604 +-----------------
 .../dtensor/python/tests/input_util_test.py   |  35 +
 2 files changed, 38 insertions(+), 601 deletions(-)

diff --git a/tensorflow/dtensor/python/input_util.py b/tensorflow/dtensor/python/input_util.py
index 488870c2169..fc89c239aef 100644
--- a/tensorflow/dtensor/python/input_util.py
+++ b/tensorflow/dtensor/python/input_util.py
@@ -552,11 +552,9 @@ class DTensorDataset(dataset_ops.UnaryUnchangedStructureDataset):
         _shard_counts(layout, batch_dim) for layout in flattened_layouts
     ]
     # `_index_matrix` requires the *batched* element spec (shape rank must
-    # match `layout.rank`, including the leading batch dim). `flattened_elem_spec`
-    # is `dataset.element_spec`, which only has that rank when the dataset was
-    # already batched; use `self._batched_dataset.element_spec` instead so this
-    # is correct for both `dataset_already_batched` values.
-    flattened_batched_elem_spec = nest.flatten(self._batched_dataset.element_spec)
+    # match `layout.rank`, including the leading batch dim).
+    flattened_batched_elem_spec = nest.flatten(
+        self._batched_dataset.element_spec)
     self._index_matrices = [
         _index_matrix(layout, elem_spec)
         for layout, elem_spec in zip(flattened_layouts,
@@ -675,602 +673,6 @@ class DTensorDataset(dataset_ops.UnaryUnchangedStructureDataset):
     partitioned_dataset = enumerated_dataset.map(slice_batch)
     return partitioned_dataset
 
-  @property
-  def element_spec(self):
-    return self._global_element_specfrom tensorflow.python.data.experimental.ops import distribute
-from tensorflow.python.data.ops import dataset_ops
-from tensorflow.python.data.ops import iterator_ops
-from tensorflow.python.eager import context
-from tensorflow.python.framework import constant_op
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import errors
-from tensorflow.python.framework import ops
-from tensorflow.python.framework import tensor
-from tensorflow.python.framework import tensor_shape
-from tensorflow.python.framework import tensor_spec
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import math_ops
-from tensorflow.python.types import data as data_types
-from tensorflow.python.util import nest
-from tensorflow.python.util.tf_export import tf_export
-
-
-@dataclasses.dataclass
-class TFDataServiceConfig:
-  """Specifies the tf.data service configuration to use.
-
-  Attributes:
-    dispatcher_address: a string specifying the address of the tf.data service
-      dispatcher server.
-    job_name: a non-empty string identifying the shared job that will be created
-      on tf.data service to process this dataset.
-  """
-  dispatcher_address: str
-  job_name: str
-
-
-# TODO(b/223275517): Add support for get_next_as_optional().
-class _DTensorIterator(iterator_ops.OwnedIterator):
-  """An iterator for a tf.data.Dataset distributed using DTensor.
-
-  DTensorIterator encapsulates multiple underlying dataset iterators. It handles
-  retrieving the tensors to be placed on each underlying device and then uses
-  the 'pack' operation to create and return a DTensor. Thus users need only
-  interact with a single DTensorIterator to automatically distribute dataset
-  tensors onto devices.
-  """
-
-  def __init__(
-      self,
-      dtensor_components: Tuple[tensor.Tensor],
-      global_element_spec: tensor_spec.TensorSpec,
-      layouts: Any):
-    """Initializes a distributed iterator for DTensor datasets.
-
-    This iterator encapsulates tf.data iterators for the underlying devices, and
-    treats it as a packed DTensor of iterator resource tensors.
-
-    Args:
-      dtensor_components: a tuple containing the underlying iterator resources
-        packed into a DTensor. This is expected to be a tuple with a single
-        element.
-      global_element_spec: the underlying dataset's element spec from a global
-        view.
-      layouts: a structure of DTensor layouts to be applied to the elements
-        returned by the underlying iterators. This can be a single layout or
-        (possibly nested) tuples or dictionaries of layouts, and the structure
-        must match the structure of the iterator elements.
-    """
-    # dtensor_components is expected to be a single-element tuple.
-    [self._iterator_resource_dtensor] = dtensor_components
-    self._global_element_spec = global_element_spec
-    self._layouts = layouts
-    self._layouts_str = nest.map_structure(
-        lambda layout: layout.to_string(), layouts)
-
-    super().__init__(
-        components=dtensor_components, element_spec=global_element_spec)
-
-  def __next__(self):
-    try:
-      # IteratorGetNext will return a DTensor on the host, so move it to the
-      # device mesh. If the dataset layouts are on the host mesh itself, this
-      # is handled by DTensor as a no-op.
-      host_elem = self._next_internal()
-      context.async_wait()
-      device_elem = nest.map_structure(
-          api.copy_to_mesh, host_elem, self._layouts)
-      context.async_wait()
-      return device_elem
-    except errors.OutOfRangeError as e:
-      # Match TF2 eager executor behavior by raising StopIteration when iterator
-      # is out of range.
-      if context.executing_eagerly():
-        raise StopIteration from e
-      else:
-        raise e
-
-  @property
-  def _type_spec(self):
-    return _DTensorIteratorSpec(self._global_element_spec, self._layouts_str)
-
-
-class _DTensorIteratorSpec(iterator_ops.IteratorSpec):
-  """Type specification for `_DTensorIterator`."""
-
-  __slots__ = ['_global_element_spec', '_layouts_str']
-
-  def __init__(
-      self, global_element_spec: tensor_spec.TensorSpec, layouts_str: Any):
-    super().__init__(global_element_spec)
-    self._global_element_spec = global_element_spec
-    self._layouts_str = layouts_str
-
-  @property
-  def value_type(self):
-    return _DTensorIterator
-
-  def _serialize(self):
-    return (self._global_element_spec, self._layouts_str)
-
-  @property
-  def _component_specs(self):
-    return (tensor_spec.TensorSpec([], dtypes.resource),)
-
-  def _to_components(self, value):
-    return (value._iterator_resource_dtensor,)  # pylint: disable=protected-access
-
-  def _from_components(self, components):
-    layouts = nest.map_structure(
-        layout_lib.Layout.from_string, self._layouts_str)
-    return _DTensorIterator(
-        dtensor_components=components,
-        global_element_spec=self._global_element_spec,
-        layouts=layouts)
-
-  @classmethod
-  def from_value(cls, value):
-    return cls(value._global_element_spec, value._layouts_str)  # pylint: disable=protected-access
-
-
-def _validate_input(flattened_layouts: Sequence[layout_lib.Layout],
-                    flattened_elem_spec: Sequence[tensor_spec.TensorSpec],
-                    dataset_already_batched: bool):
-  """Checks that the dataset's layouts and element specs are compatible.
-
-  Args:
-    flattened_layouts: the flattened list of layouts used to distribute the
-      dataset.
-    flattened_elem_spec: the flattened list of element specs used in the
-      dataset's components.
-    dataset_already_batched: whether the dataset to be validated is already
-      batched.
-
-  Raises:
-    ValueError: if the dataset's inputs are incompatible.
-  """
-  if not flattened_elem_spec:
-    raise ValueError(
-        'Expected input element spec of at least one element, was empty.')
-
-  first_elem_shape = flattened_elem_spec[0].shape
-
-  for layout, elem_spec in zip(flattened_layouts, flattened_elem_spec):
-    if elem_spec.shape.rank is None:
-      raise ValueError(
-          'Dataset element shape must have a valid rank, got spec %s.' %
-          elem_spec)
-
-    # Check that layout's rank matches the element's rank. If dataset is not yet
-    # batched, then the layout's rank must be one greater than the element's
-    # rank.
-    expected_rank = elem_spec.shape.rank
-    if not dataset_already_batched:
-      expected_rank += 1
-    if layout.rank != expected_rank:
-      raise ValueError(
-          ('Expected layout with rank %d for element spec %s, got layout %s. '
-           'Check that the dataset is not batched before passing to '
-           'DTensorDataset.') %
-          (expected_rank, elem_spec, layout.sharding_specs))
-
-    if dataset_already_batched:
-      # Check that the batch dimension size of all dataset elements match.
-      batch_dim_size = first_elem_shape.as_list()[0]
-      if batch_dim_size is None:
-        raise ValueError(
-            ('Size of batch dimension of element spec %s is None. Ensure '
-             'drop_remainder=True when batching the dataset.') % elem_spec)
-
-      if elem_spec.shape.as_list()[0] != batch_dim_size:
-        raise ValueError(
-            ('Size of batch dimension of element spec %s does not match '
-             'expected size %d.') % (elem_spec, batch_dim_size))
-
-
-def _shard_counts(layout: layout_lib.Layout,
-                  batch_dim: Optional[str] = None) -> List[int]:
-  """Computes a list of the number of shards in each dimension of the layout.
-
-  The shard counts are used to slice each dataset element. The batch dimension's
-  count is overridden to 1 since we only consider how many shards to make
-  locally (within each local replica). Sharding across clients is handled by
-  either tf.data.Dataset's shard transformation (in the single-client case) or
-  tf.data service's distribute function (in the multi-client case).
-
-  Args:
-    layout: the layout to compute the shard counts for.
-    batch_dim: the name of the batch dimension of the layout, if present.
-
-  Returns:
-    A list of shard counts, one element per dimension of the layout.
-  """
-  shard_counts = []
-  for spec in layout.sharding_specs:
-    if spec in (batch_dim, layout_lib.UNSHARDED):
-      shard_counts.append(1)
-    else:
-      shard_counts.append(layout.mesh.dim_size(spec))
-  return shard_counts
-
-
-def _index_matrix(layout: layout_lib.Layout,
-                  elem_spec: tensor_spec.TensorSpec) -> tensor.Tensor:
-  """Computes a utility matrix to derive device-based slice offsets.
-
-  This function builds a matrix of shape `[mesh.rank, layout.rank]` for each
-  dataset element. This matrix can be used to slice the DTensor components
-  returned by the iterator according to the local device that component is to be
-  placed on. This can be done by multiplying the device offsets of shape
-  `[1, mesh.rank]` with this index matrix to get a `[1, layout.rank]` shape
-  tensor containing the slice offsets.
-
-  Note: the index on the batch dim is always 0 since sharding on the batch
-  dimension is handled by either tf.data.Dataset's shard transformation (in the
-  single-client case) or tf.data service's distribute function (in the
-  multi-client case). If there is no sharding on the batch dimension (or any
-  other dimension), the slice index remains 0.
-
-  Args:
-    layout: the layout of the dataset element.
-    elem_spec: the spec of the dataset element.
-
-  Returns:
-    The index matrix as a tensor.
-  """
-  matrix = []
-  for dim in layout.mesh.dim_names:
-    row = [0]
-    for layout_idx, spec in enumerate(layout.sharding_specs[1:]):
-      if spec == layout_lib.UNSHARDED or spec != dim:
-        row.append(0)
-      else:
-        row.append(elem_spec.shape[layout_idx] // layout.mesh.dim_size(dim))
-    matrix.append(row)
-
-  return constant_op.constant(matrix, dtype=dtypes.int32)
-
-
-def _pack_iterator_resource_dtensor(
-    datasets: List[Tuple[int, data_types.DatasetV2]],
-    layouts: Any,
-    mesh: layout_lib.Mesh,
-    num_local_devices_per_replica: int):
-  """Creates a DTensor iterator resource for the per-replica datasets.
-
-  Given a list of replica ID to tf.data.Dataset mappings, this function creates
-  iterators for each device and then packs the underlying iterator resource
-  tensors into a single DTensor. This resource tensor is used by the
-  IteratorGetNext op to retrieve the next element in the dataset.
-
-  Args:
-    datasets: a list of tuples of each unique local replica ID to the dataset
-      object whose elements will be placed on the devices corresponding to that
-      replica.
-    layouts: a structure of DTensor layouts to be applied to the elements
-      returned by the underlying iterators. This can be a single layout or
-      (possibly nested) tuples or dictionaries of layouts, and the structure
-      must match the structure of the iterator elements.
-    mesh: the DTensor mesh to place the iterator batches on.
-    num_local_devices_per_replica: the number of devices in each data-parallel
-      replica.
-
-  Returns:
-    A DTensor of the underlying iterator resource tensors.
-  """
-  host_mesh_devices = mesh.host_mesh().local_devices()
-  device_idx = 0
-
-  iterators = []
-  for _, dataset in datasets:
-    for idx in range(num_local_devices_per_replica):
-      with ops.device_v2(host_mesh_devices[device_idx]):
-        device_dataset = dataset.shard(
-            num_shards=num_local_devices_per_replica, index=idx)
-        iterators.append(iter(device_dataset))
-      device_idx += 1
-
-  if device_idx != len(host_mesh_devices):
-    raise ValueError(
-        'The `datasets` argument does not have the correct number of'
-        f' underlying datasets, found {device_idx} but expected'
-        f' {len(host_mesh_devices)}.')
-
-  host_layouts = nest.map_structure(
-      lambda l: layout_lib.Layout(l.sharding_specs, mesh.host_mesh()), layouts)
-
-  # Pack the iterator resource tensors into a replicated 0-dimensional DTensor
-  # and set the element layouts.
-  iterator_resources = [it._iterator_resource for it in iterators]  # pylint: disable=protected-access
-  d_iterator_resource = api.pack(
-      iterator_resources,
-      layout_lib.Layout.replicated(mesh=mesh.host_mesh(), rank=0))
-  api._dtensor_device().set_iterator_element_layouts(  # pylint: disable=protected-access
-      d_iterator_resource, nest.flatten(host_layouts))
-
-  return d_iterator_resource
-
-
-@tf_export('experimental.dtensor.DTensorDataset', v1=[])
-class DTensorDataset(dataset_ops.UnaryUnchangedStructureDataset):
-  """A dataset of DTensors.
-
-  DTensorDataset encapsulates a `tf.data.Dataset` whose elements are
-  automatically packed and returned as DTensors based on a given mesh and
-  layouts.
-  """
-
-  def __init__(self,
-               dataset: data_types.DatasetV2,
-               *,
-               mesh: layout_lib.Mesh,
-               layouts: Any,
-               global_batch_size: int,
-               dataset_already_batched: bool = False,
-               batch_dim: Optional[str] = None,
-               prefetch: Optional[int] = None,
-               tf_data_service_config: Optional[TFDataServiceConfig] = None):
-    """Creates a DTensorDataset.
-
-    DTensorDataset automatically handles distribution of the dataset elements to
-    each client's devices. It can be used to create an iterator that returns
-    DTensors of the input data on each iteration.
-
-    DTensorDataset works best with unbatched datasets. It takes the mesh and the
-    provided layouts to automatically calculate how to batch the input locally
-    for each replica.
-
-    If the provided dataset is already batched according to the per-replica
-    batch size, then `dataset_already_batched` must be set and DTensorDataset
-    will check that the batch size is consistent with the intended
-    `global_batch_size` using the layout information. Each replica receives a
-    separate slice of the global batch, thus the per-replica batch size can be
-    computed as the global batch size divided by the number of model replicas.
-    For a DTensor mesh, the number of replicas is equal to the size of the
-    mesh's batch dimension.
-
-    Note: `tf.experimental.dtensor.DTensorDataset` instances do *not* implement
-    the full interface of `tf.data.Dataset`. It only supports two usages we will
-    mention below: iteration and `element_spec`. We don't support any other APIs
-    to transform or inspect the dataset.
-
-    TODO(b/223275517): add support for input datasets that are already batched
-    to the global batch size.
-
-    Args:
-      dataset: a `tf.data.Dataset` object.
-      mesh: the DTensor mesh to place the dataset batches on.
-      layouts: a structure of DTensor layouts to be applied to the input dataset
-        values. This can be a single layout or (possibly nested) tuples or
-        dictionaries of layouts, and the structure must match the structure of
-        the dataset. Either all or none of the layouts should be sharded on the
-        batch dimension; having only a subset of layouts batch sharded will not
-        work and raises a ValueError.
-      global_batch_size: the desired global batch size.
-      dataset_already_batched: must be set only if the dataset is already
-        batched to the per-replica batch size. The batched dataset must have
-        `drop_remainder=True` set since DTensor requires static shapes for
-        slicing the input tensors.
-      batch_dim: the mesh dimension on which the input's batch dimension is
-        sharded. Set to None if the input layouts do not shard on the batch
-        dimension.
-      prefetch: number of batches to prefetch using Dataset.prefetch.
-      tf_data_service_config: if operating in multi-client mode, this config
-        specifies the tf.data service configuration to use.
-
-    Raises:
-      ValueError: on any of the following situations,
-        1. if the structures and ranks of layouts and the dataset do not match.
-        2. if the shapes in the dataset's spec are not fully defined.
-        3. if batch_dim is specified and all layouts are not batch-sharded.
-        4. if per_replica_batch_size is specified for an already batched Dataset
-           but it does not match the expected per-replica size based on the
-           provided mesh.
-      TypeError: if type of structures of layouts and the dataset do not match.
-    """
-    super().__init__(dataset, dataset_ops.to_variant(dataset))
-
-    # TODO(b/271162918): fix multi-client use case.
-    if tf_data_service_config is not None:
-      raise NotImplementedError(
-          'Multi-client DTensorDataset is currently not supported.'
-          ' Check b/271162918.')
-
-    self._mesh = mesh
-    self._layouts = layouts
-    self._batch_dim = batch_dim
-    self._prefetch = prefetch
-    self._tf_data_service_config = tf_data_service_config
-
-    nest.assert_same_structure(dataset.element_spec, layouts)
-    flattened_layouts = nest.flatten(layouts)
-    flattened_elem_spec = nest.flatten(dataset.element_spec)
-
-    if batch_dim:
-      self.num_global_replicas = mesh.dim_size(batch_dim)
-      self._local_replica_ids = list(
-          dict.fromkeys(
-              [loc[batch_dim] for loc in mesh.local_device_locations()]))
-
-      for layout in flattened_layouts:
-        if batch_dim != layout.sharding_specs[0]:
-          raise ValueError(
-              ('batch_dim %s was specified but at least one layout did not '
-               'contain it: %s') % (batch_dim, layout))
-    else:
-      # Only one replica since there is no sharding on the batch dimension.
-      self.num_global_replicas = 1
-      self._local_replica_ids = [0]
-
-    # Validate layout and element spec compatibility, and raise ValueError if
-    # invalid.
-    _validate_input(
-        flattened_layouts,
-        flattened_elem_spec,
-        dataset_already_batched=dataset_already_batched)
-
-    expected_batch_size = global_batch_size // self.num_global_replicas
-    if not dataset_already_batched:
-      self._batched_dataset = dataset.batch(
-          expected_batch_size, drop_remainder=True)
-    else:
-      per_replica_batch_size = flattened_elem_spec[0].shape.as_list()[0]
-      if per_replica_batch_size != expected_batch_size:
-        raise ValueError(
-            ('per_replica_batch_size does not matched expected size based on '
-             'the mesh, got %d but expected %d.') %
-            (per_replica_batch_size, expected_batch_size))
-      self._batched_dataset = dataset
-
-    # Construct a global element spec of the dataset.
-    flattened_global_elem_spec = []
-    batch_tensor_shape = tensor_shape.as_shape([global_batch_size])
-    for elem_spec in nest.flatten(self._batched_dataset.element_spec):
-      new_elem_spec = tensor_spec.TensorSpec(
-          shape=operator.concat(batch_tensor_shape, elem_spec.shape[1:]),
-          dtype=elem_spec.dtype,
-          name=elem_spec.name)
-      flattened_global_elem_spec.append(new_elem_spec)
-    self._global_element_spec = nest.pack_sequence_as(
-        dataset.element_spec, flattened_global_elem_spec)
-
-    num_global_devices_per_replica = config.num_global_devices(
-        mesh.device_type()) // self.num_global_replicas
-    self._num_local_replicas = len(self._local_replica_ids)
-    self._num_local_devices_per_replica = mesh.num_local_devices(
-    ) // self._num_local_replicas
-    # The number of clients each replica is split over.
-    self._num_clients_per_replica = (
-        num_global_devices_per_replica // self._num_local_devices_per_replica)
-    # In the case where a replica is split across multiple clients, an offset
-    # needs to be added to the index used by the partitioning logic such that
-    # the local devices on that client can be correctly matched to slices of the
-    # input tensor(s). If replicas are wholly contained within a client, then
-    # this offset is always 0.
-    self._partition_offset = (config.client_id() % self._num_clients_per_replica
-                             ) * self._num_local_devices_per_replica
-
-    # Helper data structures used in partitioning the dataset tensors.
-    self._all_shard_counts = [
-        _shard_counts(layout, batch_dim) for layout in flattened_layouts
-    ]
-    self._index_matrices = [
-        _index_matrix(layout, elem_spec)
-        for layout, elem_spec in zip(flattened_layouts, flattened_elem_spec)
-    ]
-
-  def __iter__(self):
-    datasets: List[Tuple[int, data_types.DatasetV2]] = []
-
-    # Start with the batched the dataset.
-    local_dataset = self._batched_dataset
-
-    if self._batch_dim is not None:
-      if self._num_clients_per_replica > 1:
-        # If a replica is split over multiple clients then each batch needs to
-        # be repeated before distribution as many times as there are clients
-        # corresponding to that replica.
-        local_dataset = self._repeat_batch(local_dataset,
-                                           self._num_clients_per_replica)
-        sharding_policy = data_service_ops.ShardingPolicy.DATA
-      else:
-        # Replicas are unique to each client, so FILE based sharding can be used
-        # which is more performant since each worker does not need to read the
-        # entire dataset.
-        sharding_policy = data_service_ops.ShardingPolicy.FILE
-    else:
-      # No batch dimension sharding specified so disable dataset sharding during
-      # the distribute step.
-      sharding_policy = data_service_ops.ShardingPolicy.OFF
-
-    # Apply distribution here (if specified) so all remaining transformations
-    # are executed locally.
-    if self._tf_data_service_config is not None:
-      local_dataset = local_dataset.apply(
-          data_service_ops.distribute(
-              processing_mode=sharding_policy,
-              service=self._tf_data_service_config.dispatcher_address,
-              job_name=f'{self._tf_data_service_config.job_name}_{config.client_id()}',
-              target_workers='LOCAL'))
-
-    for local_replica_idx, replica_id in enumerate(self._local_replica_ids):
-      # Select the shard for the corresponding replica.
-      dataset = distribute._AutoShardDataset(
-          local_dataset,
-          num_workers=self._num_local_replicas,
-          index=local_replica_idx,
-          num_replicas=self.num_global_replicas)
-
-      # Repeat each batch for each local device in the replica.
-      dataset = self._repeat_batch(dataset, self._num_local_devices_per_replica)
-
-      # Slice each shard further for all non-batch dim shards. If there is no
-      # non-batch dim sharding, this slice is essentially a no-op.
-      dataset = self._partition(dataset)
-
-      # Apply prefetch as the last step. Since each batch is repeated, the
-      # number of elements to prefetch has to be scaled by the same size.
-      if self._prefetch is not None:
-        dataset = dataset.prefetch(
-            self._prefetch * self._num_local_devices_per_replica)
-
-      datasets.append((replica_id, dataset))
-
-    # Convert the datasets into iterators placed on the host.
-    d_iterator_resource = _pack_iterator_resource_dtensor(
-        datasets=datasets,
-        layouts=self._layouts,
-        mesh=self._mesh,
-        num_local_devices_per_replica=self._num_local_devices_per_replica)
-
-    return _DTensorIterator(
-        dtensor_components=(d_iterator_resource,),
-        global_element_spec=self._global_element_spec,
-        layouts=self._layouts)
-
-  def _repeat_batch(self, dataset, repeats):
-    if repeats == 1:
-      # Remove this shortcut if tf.data can optimize this away.
-      return dataset
-
-    def repeat(*x):
-      return dataset_ops.DatasetV2.from_tensors(x).repeat(repeats)
-
-    return dataset.flat_map(repeat)
-
-  def _partition(self, dataset):
-    """Slices each dataset element on any sharded non-batch dimension."""
-    if self._num_local_devices_per_replica == 1 and self._partition_offset == 0:
-      # Remove this shortcut if tf.data can optimize this away.
-      return dataset
-
-    # TODO(b/223275517): decouple from self and make testable.
-    def slice_batch(index, batch):
-      flattened_batch = nest.flatten(batch)
-      flattened_output = []
-
-      norm_index = math_ops.cast(
-          index % self._num_local_devices_per_replica, dtype=dtypes.int32)
-      norm_index += self._partition_offset
-      coords = self._mesh.coords(norm_index)
-      coords = array_ops.reshape(coords, (1, -1))
-
-      for element, shard_counts, idx_matrix in zip(flattened_batch,
-                                                   self._all_shard_counts,
-                                                   self._index_matrices):
-        indexes = math_ops.matmul(coords, idx_matrix)
-        start = array_ops.reshape(indexes, (-1,))
-        size = array_ops.shape_v2(
-            element, out_type=dtypes.int32) // shard_counts
-        flattened_output.append(
-            array_ops.slice(element, begin=start, size=size))
-
-      return nest.pack_sequence_as(batch, flattened_output)
-
-    enumerated_dataset = dataset.enumerate()
-    partitioned_dataset = enumerated_dataset.map(slice_batch)
-    return partitioned_dataset
-
   @property
   def element_spec(self):
     return self._global_element_spec
diff --git a/tensorflow/dtensor/python/tests/input_util_test.py b/tensorflow/dtensor/python/tests/input_util_test.py
index aa1d386e339..306eb03359c 100644
--- a/tensorflow/dtensor/python/tests/input_util_test.py
+++ b/tensorflow/dtensor/python/tests/input_util_test.py
@@ -604,6 +604,41 @@ class InputUtilHelpersTest(test_util.DTensorBaseTest):
 
     self.assertEqual(input_util._shard_counts(layout, batch_dim), counts)
 
+  @parameterized.parameters(
+      {
+          'mesh_dims': [(MESH_DIM_BATCH, 2),
+                        (MESH_DIM_HEIGHT, 4),
+                        (MESH_DIM_WIDTH, 2)],
+          'layout_specs': [MESH_DIM_BATCH, MESH_DIM_HEIGHT, UNSHARDED],
+          'elem_shape': [8, 16, 3],
+          'expected_matrix': [[0, 0, 0], [0, 4, 0], [0, 0, 0]],
+      }, {
+          'mesh_dims': [(MESH_DIM_BATCH, 2),
+                        (MESH_DIM_HEIGHT, 4),
+                        (MESH_DIM_WIDTH, 2)],
+          'layout_specs': [MESH_DIM_BATCH, MESH_DIM_WIDTH, MESH_DIM_HEIGHT],
+          'elem_shape': [8, 6, 16],
+          'expected_matrix': [[0, 0, 0], [0, 0, 4], [0, 3, 0]],
+      })
+  def testIndexMatrixForBatchedElementSpec(
+      self, mesh_dims, layout_specs, elem_shape, expected_matrix):
+    """Regression test for the `dataset_already_batched=True` offset bug.
+
+    `_index_matrix` must be called with the *batched* element spec (shape
+    rank matching `layout.rank`, including the leading batch dim) and must
+    index it at `layout_idx + 1` to skip that leading batch dim. Passing an
+    unbatched spec, or omitting the `+ 1` offset, silently computes the wrong
+    slice offsets for sharded, non-batch dimensions.
+    """
+    num_devices = np.prod([size for _, size in mesh_dims])
+    mesh = mesh_util.create_mesh(
+        mesh_dims=mesh_dims, devices=['CPU:%d' % i for i in range(num_devices)])
+    layout = Layout(layout_specs, mesh)
+    elem_spec = tensor_spec.TensorSpec(shape=elem_shape, dtype=dtypes.float32)
+
+    matrix = input_util._index_matrix(layout, elem_spec)
+    self.assertAllEqual(matrix, expected_matrix)
+
 
 class DTensorIteratorSpecTest(test_util.DTensorBaseTest):
 

From f68c587c7c29c3e278b12fcfd4a531b791eda480 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Tue, 13 Oct 2026 21:00:55 +0530
Subject: [PATCH 0683/1000] Address review: guard against length mismatch; add
 regression tests

- _enumerate_core_locations and create_tpu_mesh: add explicit
  len(ring_bounds) != len(host_bounds) / len(topology_shape) checks
  before the element-wise any(...) comparisons, since zip() silently
  truncates to the shorter sequence and could let an out-of-bounds
  value in a trailing dimension pass validation undetected.
- Add testCreateMeshWithInvalidRingBounds and
  testEnumerateCoreLocationsWithSmallRingBounds to DeviceAssignmentTest,
  covering both modified validation branches with the exact cases that
  a lexicographical list comparison would have missed.
---
 .../tests/tpu_device_assignment_test.py       | 33 +++++++++++++++++++
 tensorflow/dtensor/python/tpu_util.py         |  8 +++--
 2 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/tensorflow/dtensor/python/tests/tpu_device_assignment_test.py b/tensorflow/dtensor/python/tests/tpu_device_assignment_test.py
index 08ece48382e..b4537e51012 100644
--- a/tensorflow/dtensor/python/tests/tpu_device_assignment_test.py
+++ b/tensorflow/dtensor/python/tests/tpu_device_assignment_test.py
@@ -802,6 +802,39 @@ class DeviceAssignmentTest(test_util.DTensorBaseTest):
     with self.assertRaises(ValueError):
       tpu_util.create_tpu_mesh(['x'], [1], 'mesh_with_wrong_shape')
 
+  def testCreateMeshWithInvalidRingBounds(self):
+    topology_shape = list(tpu_util._tpu_topology.mesh_shape)
+    # A bound that exceeds the topology in a non-first dimension
+    # would pass a lexicographical check but must fail element-wise.
+    invalid_ring_bounds = [
+        topology_shape[0],
+        topology_shape[1] + 1,
+        topology_shape[2],
+        topology_shape[3],
+    ]
+    with self.assertRaisesRegex(
+        ValueError, r"ring_bounds .* should be <= topology sizes"
+    ):
+      tpu_util.create_tpu_mesh(
+          ["x", "y"],
+          [2, 1],
+          "mesh_invalid_ring_bounds",
+          ring_bounds=invalid_ring_bounds,
+      )
+
+  def testEnumerateCoreLocationsWithSmallRingBounds(self):
+    # A ring bound smaller than the host bound in a non-first dimension.
+    with self.assertRaisesRegex(
+        ValueError, r"ring_bounds .* should be >= host_bounds"
+    ):
+      tpu_util._enumerate_core_locations(
+          bounds=[4, 4, 1, 2],
+          ring_bounds=[4, 1, 1, 2],
+          axes=["core", "x", "y", "z"],
+          can_split_host_across_rings=False,
+          ring_size=8,
+      )
+
   # Build rings for the batch dimension.
   def testCreateMeshWithPositiveRingDims(self):
     mesh = tpu_util.create_tpu_mesh(['x', 'y'], [2, 1],
diff --git a/tensorflow/dtensor/python/tpu_util.py b/tensorflow/dtensor/python/tpu_util.py
index 24bea30bccb..25ccbe56ccd 100644
--- a/tensorflow/dtensor/python/tpu_util.py
+++ b/tensorflow/dtensor/python/tpu_util.py
@@ -445,7 +445,9 @@ def _enumerate_core_locations(
 
   # Reorder ring_bounds and validate it's element-wise >= host_bounds.
   ring_bounds = [ring_bounds[i] for i in axes]  # pyrefly: ignore[bad-index]
-  if any(r < h for r, h in zip(ring_bounds, host_bounds)):
+  if len(ring_bounds) != len(host_bounds) or any(
+      r < h for r, h in zip(ring_bounds, host_bounds)
+  ):
     raise ValueError(
         "ring_bounds %s should be >= host_bounds %s"
         % (ring_bounds, host_bounds)
@@ -713,7 +715,9 @@ def create_tpu_mesh(
     ring_bounds = topology_shape
   elif len(ring_bounds) != 4:
     raise ValueError("Expected 4 elements in ring_bounds, got %s" % ring_bounds)
-  elif any(r > t for r, t in zip(ring_bounds, topology_shape)):
+  elif len(ring_bounds) != len(topology_shape) or any(
+      r > t for r, t in zip(ring_bounds, topology_shape)
+  ):
     raise ValueError(
         "ring_bounds %s should be <= topology sizes %s"
         % (ring_bounds, topology_shape)

From cea5ac3cf15ca0d1fb130d5d6afba707d9971f8f Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Tue, 13 Oct 2026 21:02:18 +0530
Subject: [PATCH 0684/1000] Address review: defensive empty-stack check and
 regression test suite

- clear_load_options: raise a descriptive RuntimeError instead of letting
  an empty stack pop raise a generic IndexError.
- Add load_context_test.py covering default state, single and nested
  contexts, nested context with None options, thread isolation, and
  stack unwinding on an exception inside a nested context -- verified
  each case manually against the actual fix before committing.
- Register the new test target in python/distribute/BUILD.

Left the legacy duplicate in
python/keras/saving/saved_model/load_context.py untouched, since it's
a separate, frozen copy and the reviewer noted fixing it is not
required for this PR.
---
 tensorflow/python/distribute/BUILD            |  9 ++
 tensorflow/python/distribute/load_context.py  |  4 +
 .../python/distribute/load_context_test.py    | 96 +++++++++++++++++++
 3 files changed, 109 insertions(+)
 create mode 100644 tensorflow/python/distribute/load_context_test.py

diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD
index 2c0bc7c8fda..b5e63d595ad 100644
--- a/tensorflow/python/distribute/BUILD
+++ b/tensorflow/python/distribute/BUILD
@@ -2692,6 +2692,15 @@ pytype_strict_library(
     deps = ["//tensorflow/python/util:tf_export"],
 )
 
+tf_py_strict_test(
+    name = "load_context_test",
+    srcs = ["load_context_test.py"],
+    deps = [
+        ":load_context",
+        "//tensorflow/python/eager:test",
+    ],
+)
+
 py_library(
     name = "merge_call_interim",
     srcs = [
diff --git a/tensorflow/python/distribute/load_context.py b/tensorflow/python/distribute/load_context.py
index 6950c240770..26676200783 100644
--- a/tensorflow/python/distribute/load_context.py
+++ b/tensorflow/python/distribute/load_context.py
@@ -35,6 +35,10 @@ class LoadContext(threading.local):
     self._load_options_stack.append(load_options)
 
   def clear_load_options(self):
+    if not self._load_options_stack:
+      raise RuntimeError(
+          "clear_load_options() called when load context stack is empty."
+      )
     self._load_options_stack.pop()
 
   def load_options(self):
diff --git a/tensorflow/python/distribute/load_context_test.py b/tensorflow/python/distribute/load_context_test.py
new file mode 100644
index 00000000000..119f9759ad3
--- /dev/null
+++ b/tensorflow/python/distribute/load_context_test.py
@@ -0,0 +1,96 @@
+# Copyright 2023 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 load_context."""
+
+import threading
+
+from tensorflow.python.distribute import load_context
+from tensorflow.python.eager import test
+
+
+class LoadContextTest(test.TestCase):
+
+  def test_default_state(self):
+    self.assertFalse(load_context.in_load_context())
+    self.assertIsNone(load_context.get_load_options())
+
+  def test_single_context(self):
+    options = object()
+    with load_context.load_context(options):
+      self.assertTrue(load_context.in_load_context())
+      self.assertIs(load_context.get_load_options(), options)
+    self.assertFalse(load_context.in_load_context())
+    self.assertIsNone(load_context.get_load_options())
+
+  def test_nested_context(self):
+    outer_options = object()
+    inner_options = object()
+    with load_context.load_context(outer_options):
+      self.assertTrue(load_context.in_load_context())
+      self.assertIs(load_context.get_load_options(), outer_options)
+      with load_context.load_context(inner_options):
+        self.assertTrue(load_context.in_load_context())
+        self.assertIs(load_context.get_load_options(), inner_options)
+      self.assertTrue(load_context.in_load_context())
+      self.assertIs(load_context.get_load_options(), outer_options)
+    self.assertFalse(load_context.in_load_context())
+    self.assertIsNone(load_context.get_load_options())
+
+  def test_nested_context_with_none(self):
+    outer_options = object()
+    with load_context.load_context(outer_options):
+      with load_context.load_context(None):
+        self.assertTrue(load_context.in_load_context())
+        self.assertIsNone(load_context.get_load_options())
+      self.assertTrue(load_context.in_load_context())
+      self.assertIs(load_context.get_load_options(), outer_options)
+
+  def test_thread_isolation(self):
+    options = object()
+    errors = []
+
+    def thread_fn():
+      try:
+        if load_context.in_load_context():
+          errors.append("Expected not in_load_context in background thread")
+        if load_context.get_load_options() is not None:
+          errors.append("Expected None load_options in background thread")
+      except Exception as e:  # pylint: disable=broad-except
+        errors.append(str(e))
+
+    with load_context.load_context(options):
+      t = threading.Thread(target=thread_fn)
+      t.start()
+      t.join()
+
+    self.assertEmpty(errors)
+
+  def test_exception_unwinds_stack(self):
+    outer_options = object()
+    inner_options = object()
+    with load_context.load_context(outer_options):
+      try:
+        with load_context.load_context(inner_options):
+          raise ValueError("Simulated failure")
+      except ValueError:
+        pass
+      self.assertTrue(load_context.in_load_context())
+      self.assertIs(load_context.get_load_options(), outer_options)
+    self.assertFalse(load_context.in_load_context())
+    self.assertIsNone(load_context.get_load_options())
+
+
+if __name__ == "__main__":
+  test.main()

From 59c992fb77944b4184b5f49f1a520f35050bc81b Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Tue, 13 Oct 2026 21:04:22 +0530
Subject: [PATCH 0685/1000] Address review: add regression test for
 experimental_replication_mode fix

Add test_central_storage_strategy_input_options to
CentralStorageStrategyTest, covering both branches of the fixed typo:
InputReplicationMode.PER_REPLICA correctly raises NotImplementedError
(previously masked by the AttributeError from the misspelled
attribute), and InputReplicationMode.PER_WORKER options propagate
without error.
---
 .../parameter_server_strategy_test.py         | 24 +++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/tensorflow/python/distribute/parameter_server_strategy_test.py b/tensorflow/python/distribute/parameter_server_strategy_test.py
index fb59cd754bb..979feb91086 100644
--- a/tensorflow/python/distribute/parameter_server_strategy_test.py
+++ b/tensorflow/python/distribute/parameter_server_strategy_test.py
@@ -898,6 +898,30 @@ class CentralStorageStrategyTest(strategy_test_lib.DistributionTestBase,
     strategy, _, _ = create_test_objects(num_gpus=0)
     self.assertFalse(strategy.extended._in_multi_worker_mode())
 
+  @combinations.generate(combinations.combine(mode=['eager']))
+  def test_central_storage_strategy_input_options(self):
+    strategy = central_storage_strategy.CentralStorageStrategy._from_num_gpus(
+        0)
+    dataset = dataset_ops.Dataset.range(10).batch(2)
+
+    # PER_REPLICA mode raises NotImplementedError.
+    options = distribute_lib.InputOptions(
+        experimental_replication_mode=distribute_lib.InputReplicationMode
+        .PER_REPLICA)
+    with self.assertRaisesRegex(
+        NotImplementedError,
+        r'InputReplicationMode\.PER_REPLICA is only supported in '
+        r'`experimental_distribute_datasets_from_function`\.'):
+      strategy.experimental_distribute_dataset(dataset, options=options)
+
+    # Valid options propagate without raising AttributeError.
+    options = distribute_lib.InputOptions(
+        experimental_replication_mode=distribute_lib.InputReplicationMode
+        .PER_WORKER)
+    dist_dataset = strategy.experimental_distribute_dataset(
+        dataset, options=options)
+    self.assertIsNotNone(dist_dataset)
+
 
 if __name__ == '__main__':
   test.main()

From 6dd4e406b21bbeac51c30d23a1cca81d7c3a069f Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Tue, 13 Oct 2026 21:05:15 +0530
Subject: [PATCH 0686/1000] Wrap mock.patch.object calls to fix PyLint
 line-length failures

Lines 161 and 182 in mesh_util_test.py exceeded the 80-character limit,
failing the CI PyLint check. Wrapped both
'with mock.patch.object(...)' calls onto multiple lines.
---
 tensorflow/dtensor/python/tests/mesh_util_test.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tensorflow/dtensor/python/tests/mesh_util_test.py b/tensorflow/dtensor/python/tests/mesh_util_test.py
index 07503c6352f..18c81fab6d5 100644
--- a/tensorflow/dtensor/python/tests/mesh_util_test.py
+++ b/tensorflow/dtensor/python/tests/mesh_util_test.py
@@ -158,7 +158,9 @@ class MeshUtilTest(test_util.DTensorBaseTest):
 
   def test_create_distributed_mesh_tpu_with_local_devices_raises_error(self):
     self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
-    with mock.patch.object(accelerator_util, 'is_initialized', return_value=True):
+    with mock.patch.object(
+        accelerator_util, 'is_initialized', return_value=True
+    ):
       with self.assertRaisesRegex(
           ValueError, 'Do not specify devices for TPU meshes'
       ):
@@ -179,7 +181,9 @@ class MeshUtilTest(test_util.DTensorBaseTest):
 
   def test_create_distributed_mesh_empty_local_devices(self):
     self.skipForDeviceType(['GPU', 'TPU'], reason='Test is intended for CPUs')
-    with mock.patch.object(accelerator_util, 'is_initialized', return_value=True):
+    with mock.patch.object(
+        accelerator_util, 'is_initialized', return_value=True
+    ):
       with self.assertRaisesRegex(
           ValueError, 'Received an empty `devices` list'
       ):

From 3b091579e3b7eeb64b5284adf16bcde7a93f1239 Mon Sep 17 00:00:00 2001
From: Andrei Ivanov <anivanov@nvidia.com>
Date: Tue, 13 Oct 2026 06:12:13 -0700
Subject: [PATCH 0687/1000] PR #48199: [XLA:GPU] Fix profiling of multi-kernel
 matmuls
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/48199

📝 Summary of Changes

- Extract KernelExecutionTimer to enable testing matmul perf table profiler time aggregation.
- Change matmul perf table profiler to compute median of kernel sums instead of median of all kernels. The summation is made over kernels belonging to a single matmul.

🎯 Justification

The matmul perf table profiler previously took the median over a flat list of kernel durations, which is incorrect when one logical matmul lowers to multiple kernels (e.g., padding, split-K).

🚀 Kind of Contribution

🐛 Bug Fix, ♻️ Cleanup, 🧪 Tests

🧪 Unit Tests:

Added KernelExecutionTimerTest.MeasuresCompleteMultiKernelExecutions: testing correct aggregation of multi-kernel matmul runtimes.
Copybara import of the project:

--
67fde136bb41ab719ebddac06a55d146dff31e15 by Andrei Ivanov <anivanov@nvidia.com>:

[XLA:GPU] Extract matmul kernel execution timing

--
d296393d9f47a6c87291277135910b9b322ab96a by Andrei Ivanov <anivanov@nvidia.com>:

[XLA:GPU] Fix profiling of multi-kernel matmuls

--
d2e6b565ad688ad217538146df96f0de269662c6 by Andrei Ivanov <anivanov@nvidia.com>:

Add profiler dependency to matmul perf test

--
27b7fa8e6d08eec82fc93d24be173568e10d03ac by Andrei Ivanov <anivanov@nvidia.com>:

[XLA:GPU] Disable ASan for matmul perf table test

Merging this change closes #48199

PiperOrigin-RevId: 998584949
---
 .../xla/service/gpu/model/hlo_op_profiler.h   |  1 +
 .../service/gpu/model/hlo_op_profiler_cuda.cc |  7 +++
 .../service/gpu/model/hlo_op_profiler_rocm.cc |  7 +++
 third_party/xla/xla/tools/BUILD               |  8 +++
 .../xla/xla/tools/matmul_perf_table_gen.cc    | 50 +++++++++++++------
 .../xla/xla/tools/matmul_perf_table_gen.h     | 23 +++++++++
 .../xla/tools/matmul_perf_table_gen_test.cc   | 49 ++++++++++++++++++
 7 files changed, 130 insertions(+), 15 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/hlo_op_profiler.h b/third_party/xla/xla/service/gpu/model/hlo_op_profiler.h
index c4aa0b468c5..db9b52c86b7 100644
--- a/third_party/xla/xla/service/gpu/model/hlo_op_profiler.h
+++ b/third_party/xla/xla/service/gpu/model/hlo_op_profiler.h
@@ -40,6 +40,7 @@ class HloOpProfiler {
    public:
     virtual ~KernelTracer() = default;
     virtual uint64_t getMedianKernelTimeNs() && = 0;
+    virtual uint64_t getSumKernelTimeNs() && = 0;
   };
 
   HloOpProfiler(HloRunnerInterface* absl_nonnull runner,
diff --git a/third_party/xla/xla/service/gpu/model/hlo_op_profiler_cuda.cc b/third_party/xla/xla/service/gpu/model/hlo_op_profiler_cuda.cc
index ea35bd3bf1b..2c050f8d068 100644
--- a/third_party/xla/xla/service/gpu/model/hlo_op_profiler_cuda.cc
+++ b/third_party/xla/xla/service/gpu/model/hlo_op_profiler_cuda.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include <cstdint>
 #include <memory>
+#include <numeric>
 #include <string>
 #include <vector>
 
@@ -60,6 +61,12 @@ class CuptiKernelTracer : public HloOpProfiler::KernelTracer,
     return (kernel_times_ns_[i - 1] + kernel_times_ns_[i] + 1) / 2;
   }
 
+  uint64_t getSumKernelTimeNs() && override {
+    cupti_tracer_->Disable();  // Also flushes buffer.
+    return std::accumulate(kernel_times_ns_.begin(), kernel_times_ns_.end(),
+                           uint64_t{0});
+  }
+
  private:
   // CuptiTraceCollector
   void AddEvent(profiler::CuptiTracerEvent&& event) override {
diff --git a/third_party/xla/xla/service/gpu/model/hlo_op_profiler_rocm.cc b/third_party/xla/xla/service/gpu/model/hlo_op_profiler_rocm.cc
index c682dfc00cc..83979f37f29 100644
--- a/third_party/xla/xla/service/gpu/model/hlo_op_profiler_rocm.cc
+++ b/third_party/xla/xla/service/gpu/model/hlo_op_profiler_rocm.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include <algorithm>
 #include <cstdint>
 #include <memory>
+#include <numeric>
 #include <string>
 #include <vector>
 
@@ -59,6 +60,12 @@ class RocmKernelTracer : public HloOpProfiler::KernelTracer,
     return (kernel_times_ns_[i - 1] + kernel_times_ns_[i] + 1) / 2;
   }
 
+  uint64_t getSumKernelTimeNs() && override {
+    rocm_tracer_->Disable();  // Also flushes buffer.
+    return std::accumulate(kernel_times_ns_.begin(), kernel_times_ns_.end(),
+                           uint64_t{0});
+  }
+
  private:
   static profiler::RocmTraceCollectorOptions MakeCollectorOptions(
       uint32_t num_gpus) {
diff --git a/third_party/xla/xla/tools/BUILD b/third_party/xla/xla/tools/BUILD
index f1962bdab5d..43b7aa8ca76 100644
--- a/third_party/xla/xla/tools/BUILD
+++ b/third_party/xla/xla/tools/BUILD
@@ -764,6 +764,8 @@ cc_library(
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/functional:any_invocable",
+        "@com_google_absl//absl/functional:function_ref",
         "@com_google_absl//absl/hash",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
@@ -781,15 +783,21 @@ cc_library(
 xla_test(
     name = "matmul_perf_table_gen_test",
     srcs = ["matmul_perf_table_gen_test.cc"],
+    # TODO(b/372714955): Remove once default CUDA ver is >= 13.4 (CUPTI session leak fixed in 13.4).
+    args = if_google(["--heap_check="]),
     backends = ["gpu"],
     local_defines = if_cuda(["GOOGLE_CUDA"]),
+    # TODO(b/372714955): Remove once default CUDA ver is >= 13.4 (CUPTI session leak fixed in 13.4).
+    tags = if_google(["noasan"]),
     deps = [
         ":matmul_perf_table_gen",
         "//xla:xla_data_proto_cc",
         "//xla/backends/gpu/tests:hlo_pjrt_gpu_test_base",
         "//xla/service/gpu/model:hlo_op_profile_proto_cc",
+        "//xla/service/gpu/model:hlo_op_profiler_lib",
         "//xla/tsl/platform:statusor",
         "//xla/tsl/util/proto:proto_matchers",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest_main",
         "@com_google_protobuf//:protobuf",
diff --git a/third_party/xla/xla/tools/matmul_perf_table_gen.cc b/third_party/xla/xla/tools/matmul_perf_table_gen.cc
index 530c137195f..25ca1df0206 100644
--- a/third_party/xla/xla/tools/matmul_perf_table_gen.cc
+++ b/third_party/xla/xla/tools/matmul_perf_table_gen.cc
@@ -29,6 +29,7 @@ limitations under the License.
 
 #include "absl/container/btree_map.h"
 #include "absl/container/flat_hash_set.h"
+#include "absl/functional/function_ref.h"
 #include "absl/hash/hash.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
@@ -194,13 +195,6 @@ std::unique_ptr<HloModule> GetModule(absl::string_view lhs_dtype,
   return *std::move(parsed);
 }
 
-void Measure(HloRunnerInterface& runner, OpaqueExecutable* executable,
-             const std::vector<Literal>& args_small,
-             const std::vector<Literal>& args_large) {
-  CHECK_OK(runner.ExecuteWithExecutable(executable, args_small).status());
-  CHECK_OK(runner.ExecuteWithExecutable(executable, args_large).status());
-}
-
 void AddDotsFromStaticSpec(const MatmulPerfTableGen::Config& config,
                            std::vector<EntrySpec>& specs) {
   auto inc = [](uint32_t i, const MatmulPerfTableGen::StepSpec& spec) {
@@ -383,6 +377,29 @@ int64_t GetFlops(const HloDotInstruction& dot) {
 
 }  // namespace
 
+internal::KernelExecutionTimer::KernelExecutionTimer(
+    KernelTracerFactory tracer_factory)
+    : tracer_factory_(std::move(tracer_factory)) {}
+
+uint64_t internal::KernelExecutionTimer::MeasureRepeated(
+    int num_executions,
+    absl::FunctionRef<absl::Status(int execution)> execute) {
+  std::vector<uint64_t> execution_times_ns;
+  execution_times_ns.reserve(num_executions);
+  for (int execution = 0; execution < num_executions; ++execution) {
+    std::unique_ptr<HloOpProfiler::KernelTracer> tracer = tracer_factory_();
+    CHECK_OK(execute(execution));
+    execution_times_ns.push_back(std::move(*tracer).getSumKernelTimeNs());
+  }
+
+  std::sort(execution_times_ns.begin(), execution_times_ns.end());
+  size_t i = execution_times_ns.size() / 2;
+  if (execution_times_ns.size() % 2 != 0) {
+    return execution_times_ns[i];
+  }
+  return (execution_times_ns[i - 1] + execution_times_ns[i] + 1) / 2;
+}
+
 std::unique_ptr<OpaqueExecutable> MatmulPerfTableGen::Compile(
     std::unique_ptr<HloModule> module) {
   auto compiled =
@@ -419,14 +436,17 @@ absl::Duration MatmulPerfTableGen::Profile(std::unique_ptr<HloModule> module) {
   // First run to warm up stuff.
   CHECK_OK(runner_.ExecuteWithExecutable(compiled.get(), *args_small).status());
 
-  // Trace `kNumProfilingRuns` times to get decent measurement.
-  std::unique_ptr<HloOpProfiler::KernelTracer> tracer =
-      HloOpProfiler::GetKernelTracer();
-  for (int i = 0; i < kNumProfilingRuns; i++) {
-    Measure(runner_, compiled.get(), *args_small, *args_large);
-  }
-
-  return absl::Nanoseconds(std::move(*tracer).getMedianKernelTimeNs());
+  internal::KernelExecutionTimer timer(
+      [] { return HloOpProfiler::GetKernelTracer(); });
+  return absl::Nanoseconds(
+      timer.MeasureRepeated(2 * kNumProfilingRuns, [&](int execution) {
+        if (execution % 2 == 0) {
+          return runner_.ExecuteWithExecutable(compiled.get(), *args_small)
+              .status();
+        }
+        return runner_.ExecuteWithExecutable(compiled.get(), *args_large)
+            .status();
+      }));
 }
 
 absl::StatusOr<DeviceHloInstructionProfiles> MatmulPerfTableGen::Merge(
diff --git a/third_party/xla/xla/tools/matmul_perf_table_gen.h b/third_party/xla/xla/tools/matmul_perf_table_gen.h
index e783e9553e3..ab5572c9fe2 100644
--- a/third_party/xla/xla/tools/matmul_perf_table_gen.h
+++ b/third_party/xla/xla/tools/matmul_perf_table_gen.h
@@ -16,12 +16,15 @@ limitations under the License.
 #ifndef XLA_TOOLS_MATMUL_PERF_TABLE_GEN_H_
 #define XLA_TOOLS_MATMUL_PERF_TABLE_GEN_H_
 
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <utility>
 #include <vector>
 
 #include "absl/base/nullability.h"
+#include "absl/functional/any_invocable.h"
+#include "absl/functional/function_ref.h"
 #include "absl/log/die_if_null.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
@@ -29,12 +32,32 @@ limitations under the License.
 #include "absl/time/time.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/service/gpu/model/hlo_op_profile.pb.h"
+#include "xla/service/gpu/model/hlo_op_profiler.h"
 #include "xla/service/hlo_runner_interface.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/xla_data.pb.h"
 
 namespace xla::gpu {
 
+namespace internal {
+
+class KernelExecutionTimer {
+ public:
+  using KernelTracerFactory =
+      absl::AnyInvocable<std::unique_ptr<HloOpProfiler::KernelTracer>()>;
+
+  explicit KernelExecutionTimer(KernelTracerFactory tracer_factory);
+
+  uint64_t MeasureRepeated(
+      int num_executions,
+      absl::FunctionRef<absl::Status(int execution)> execute);
+
+ private:
+  KernelTracerFactory tracer_factory_;
+};
+
+}  // namespace internal
+
 class MatmulPerfTableGen {
  public:
   struct DataTypeSpec {
diff --git a/third_party/xla/xla/tools/matmul_perf_table_gen_test.cc b/third_party/xla/xla/tools/matmul_perf_table_gen_test.cc
index d070e536232..72d01e06df3 100644
--- a/third_party/xla/xla/tools/matmul_perf_table_gen_test.cc
+++ b/third_party/xla/xla/tools/matmul_perf_table_gen_test.cc
@@ -15,14 +15,21 @@ limitations under the License.
 
 #include "xla/tools/matmul_perf_table_gen.h"
 
+#include <algorithm>
+#include <cstddef>
 #include <cstdint>
+#include <memory>
+#include <numeric>
+#include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/status/status.h"
 #include "absl/strings/string_view.h"
 #include "google/protobuf/text_format.h"
 #include "xla/backends/gpu/tests/hlo_pjrt_gpu_test_base.h"
 #include "xla/service/gpu/model/hlo_op_profile.pb.h"
+#include "xla/service/gpu/model/hlo_op_profiler.h"
 #include "xla/tsl/platform/statusor.h"
 #include "xla/tsl/util/proto/proto_matchers.h"
 #include "xla/xla_data.pb.h"
@@ -30,6 +37,48 @@ limitations under the License.
 namespace xla::gpu {
 namespace {
 
+class FakeKernelTracer final : public HloOpProfiler::KernelTracer {
+ public:
+  void AddKernelTimeNs(uint64_t kernel_time_ns) {
+    kernel_times_ns_.push_back(kernel_time_ns);
+  }
+
+  uint64_t getMedianKernelTimeNs() && override {
+    std::sort(kernel_times_ns_.begin(), kernel_times_ns_.end());
+    size_t i = kernel_times_ns_.size() / 2;
+    if (kernel_times_ns_.size() % 2 != 0) {
+      return kernel_times_ns_[i];
+    }
+    return (kernel_times_ns_[i - 1] + kernel_times_ns_[i] + 1) / 2;
+  }
+
+  uint64_t getSumKernelTimeNs() && override {
+    return std::accumulate(kernel_times_ns_.begin(), kernel_times_ns_.end(),
+                           uint64_t{0});
+  }
+
+ private:
+  std::vector<uint64_t> kernel_times_ns_;
+};
+
+TEST(KernelExecutionTimerTest, MeasuresCompleteMultiKernelExecutions) {
+  FakeKernelTracer* active_tracer;
+  internal::KernelExecutionTimer timer(
+      [&]() -> std::unique_ptr<HloOpProfiler::KernelTracer> {
+        auto tracer = std::make_unique<FakeKernelTracer>();
+        active_tracer = tracer.get();
+        return tracer;
+      });
+
+  uint64_t measured_time_ns = timer.MeasureRepeated(10, [&](int) {
+    active_tracer->AddKernelTimeNs(25);
+    active_tracer->AddKernelTimeNs(75);
+    return absl::OkStatus();
+  });
+
+  EXPECT_EQ(measured_time_ns, 100);
+}
+
 class MatmulPerfTableGenTest : public HloPjRtGpuTestBase {
   void SetUp() override {
     auto compute_capability = device_description().gpu_compute_capability();

From 6f3191031a5649513b9e49abdfa4184f5dc58db0 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Tue, 13 Oct 2026 21:34:44 +0530
Subject: [PATCH 0688/1000] Cover both graph and eager modes in
 replication_mode regression test

Per review feedback, CentralStorageStrategy.experimental_distribute_dataset
is used in both graph and eager mode, so the regression test added for the
experimental_replication_mode typo fix must cover both to catch graph-mode
regressions.
---
 tensorflow/python/distribute/parameter_server_strategy_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/distribute/parameter_server_strategy_test.py b/tensorflow/python/distribute/parameter_server_strategy_test.py
index 979feb91086..a8352089ffd 100644
--- a/tensorflow/python/distribute/parameter_server_strategy_test.py
+++ b/tensorflow/python/distribute/parameter_server_strategy_test.py
@@ -898,7 +898,7 @@ class CentralStorageStrategyTest(strategy_test_lib.DistributionTestBase,
     strategy, _, _ = create_test_objects(num_gpus=0)
     self.assertFalse(strategy.extended._in_multi_worker_mode())
 
-  @combinations.generate(combinations.combine(mode=['eager']))
+  @combinations.generate(combinations.combine(mode=['graph', 'eager']))
   def test_central_storage_strategy_input_options(self):
     strategy = central_storage_strategy.CentralStorageStrategy._from_num_gpus(
         0)

From fe976e3f8d649d7aa9dcd079dc2f335b3de63537 Mon Sep 17 00:00:00 2001
From: Pablo Zimmermann <karupayun@google.com>
Date: Tue, 13 Oct 2026 06:14:08 -0700
Subject: [PATCH 0689/1000] [XLA:GPU] Apply HBM bandwidth derating and exclude
 cache re-reads for dot fusions.

As part of migrating `gpu_dot_fusion_cost_model` into
`GpuIndexingPerformanceModel`, tiled dot fusions need `read_time` and
`write_time` to reflect pure DRAM transfer times. Scale HBM bandwidth by the
transfer-size derate and exclude K-loop cache re-reads from DRAM read time,
since L2 traffic is modeled separately.

Unlike `CalculateHbmTime` in `gpu_dot_fusion_cost_model`, which assumes fully
coalesced accesses on the `dot` instruction's direct inputs and output, this
computes HBM times across the actual fusion parameters and roots while
accounting for per-operand coalescing utilization.

PiperOrigin-RevId: 998585834
---
 .../model/gpu_indexing_performance_model.cc   | 52 ++++++++++++++++---
 .../gpu_indexing_performance_model_test.cc    |  8 +++
 2 files changed, 52 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc
index 198284fe031..a15bd65296f 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc
+++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc
@@ -90,6 +90,8 @@ using ::xla::xtile::TilingFromAnnotatedFusion;
 
 namespace {
 
+namespace dot_model = ::xla::gpu::gpu_dot_fusion_cost_model::detail;
+
 // Information about an operand read.
 struct OperandReadInfo {
   // Total number of bytes read from the operand.
@@ -120,25 +122,44 @@ void RecordOperandRead(const HloInstruction* hlo, int64_t tile_bytes_read,
                bandwidth_utilization_rate);
 }
 
-// Computes the memory read time across all operands in `read_map`.
+// Returns the unique bytes read from DRAM for an operand, capped at the
+// operand's total size since re-reads across tiles are assumed to hit cache.
+int64_t GetDramBytesRead(const HloInstruction* hlo, const OperandReadInfo& info,
+                         const HloCostAnalysis::ShapeSizeFunction& shape_size) {
+  return std::min(shape_size(hlo->shape()), info.total_bytes_read);
+}
+
+// Computes total net HBM bytes read across all unique operands.
+int64_t GetNetBytesRead(const OperandReadMap& read_map,
+                        const HloCostAnalysis::ShapeSizeFunction& shape_size) {
+  int64_t net_bytes = 0;
+  // NOLINTNEXTLINE(*-custom-deterministic-iteration-order)
+  for (const auto& [hlo, info] : read_map) {
+    net_bytes += GetDramBytesRead(hlo, info, shape_size);
+  }
+  return net_bytes;
+}
+
+// Computes the memory read time across all operands.
 absl::Duration ComputeReadTime(
     const OperandReadMap& read_map, const se::DeviceDescription& device_info,
-    int64_t num_blocks, const HloCostAnalysis::ShapeSizeFunction& shape_size) {
+    int64_t num_blocks, const HloCostAnalysis::ShapeSizeFunction& shape_size,
+    bool include_cache_rereads = true, double bandwidth_derate = 1.0) {
   absl::Duration read_time = absl::ZeroDuration();
   // `absl::Duration` addition uses exact fixed-point integer arithmetic, so
   // map iteration order does not affect the accumulated result.
   // NOLINTNEXTLINE(*-custom-deterministic-iteration-order)
   for (const auto& [hlo, operand_read_info] : read_map) {
-    int64_t operand_size = shape_size(hlo->shape());
-    int64_t n_bytes_net =
-        std::min(operand_size, operand_read_info.total_bytes_read);
+    int64_t dram_bytes = GetDramBytesRead(hlo, operand_read_info, shape_size);
+    int64_t total_bytes =
+        include_cache_rereads ? operand_read_info.total_bytes_read : dram_bytes;
 
     read_time += GpuPerformanceModelBase::ReadTimeWithDRAMHeuristic(
-        device_info, num_blocks, n_bytes_net,
-        operand_read_info.total_bytes_read,
+        device_info, num_blocks, /*n_bytes_net=*/dram_bytes,
+        /*n_bytes_total=*/total_bytes,
         /*element_type=*/hlo->shape().element_type(),
         /*hbm_bandwidth_utilization_rate=*/
-        operand_read_info.read_bandwidth_utilization_rate);
+        operand_read_info.read_bandwidth_utilization_rate * bandwidth_derate);
   }
   return read_time;
 }
@@ -603,6 +624,21 @@ absl::StatusOr<EstimateRunTimeData> EstimateRunTimeForTiledHloComputationImpl(
   // significantly more complex for that.
   exec_time = std::max(exec_time, dot_estimates.exec_time);
 
+  if (dot_estimates.exec_time > absl::ZeroDuration()) {
+    int64_t total_net_hbm_bytes =
+        GetNetBytesRead(n_bytes_total_map, shape_size) + bytes_written;
+    float effective_hbm_bandwidth =
+        dot_model::GetEffectiveHbmBandwidth(total_net_hbm_bytes, device_info);
+    double hbm_bandwidth_derate = static_cast<double>(effective_hbm_bandwidth) /
+                                  device_info.memory_bandwidth();
+    // Exclude K-loop cache re-reads because `CalculateL2Time` models L2
+    // traffic separately.
+    read_time =
+        ComputeReadTime(n_bytes_total_map, device_info, num_blocks, shape_size,
+                        /*include_cache_rereads=*/false, hbm_bandwidth_derate);
+    write_time /= hbm_bandwidth_derate;
+  }
+
   return EstimateRunTimeData{
       /*flops=*/flops + dot_flops,
       /*bytes_read=*/bytes_read,
diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc
index 3740f0af54a..aee1cad361a 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc
+++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc
@@ -727,6 +727,14 @@ ENTRY main {
           + 4 * (64 * 32) * (num_blocks * CeilOfRatio(257, 64))  // ==> %p0
   );
   EXPECT_EQ(result.l2_bytes_read, result.bytes_read);
+  const HloDotInstruction* dot = Cast<HloDotInstruction>(
+      module->entry_computation()->root_instruction()->fused_expression_root());
+  dot_model::HbmEstimates expected_hbm = dot_model::CalculateHbmTime(
+      dot_model::DotProblemInfo(*dot), device_info_);
+  EXPECT_NEAR(absl::ToDoubleNanoseconds(result.read_time),
+              absl::ToDoubleNanoseconds(expected_hbm.read_time), 1.0);
+  EXPECT_NEAR(absl::ToDoubleNanoseconds(result.write_time),
+              absl::ToDoubleNanoseconds(expected_hbm.write_time), 15.0);
   // (32 * 64 + 32 * 64) * 4 bytes (f32) * 1 stage = 16384 bytes.
   EXPECT_EQ(result.shared_memory_per_block_bytes, 2 * (32 * 64 * 4));
   // 1 accumulator register per thread (32 * 32 / 1024) plus base kernel-state

From d75529487757c4912edb0b9c8d63114cc709faf2 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 13 Oct 2026 06:27:29 -0700
Subject: [PATCH 0690/1000] Integrate LLVM at llvm/llvm-project@4a0a7e3ddc3b

Updates LLVM usage to match
[4a0a7e3ddc3b](https://github.com/llvm/llvm-project/commit/4a0a7e3ddc3bdc4fde7c5b9b9e6890fb4dbf9153)

PiperOrigin-RevId: 998591845
---
 .../xla/third_party/llvm/workspace.bzl        |   4 +-
 .../xla/third_party/shardy/temporary.patch    | 903 ++----------------
 .../xla/third_party/shardy/workspace.bzl      |   4 +-
 .../triton/common/llvm_cl998243985.patch      |  14 +
 .../xla/third_party/triton/common/series      |   1 +
 .../codegen/triton/extern_function_helper.cc  |  11 +-
 6 files changed, 118 insertions(+), 819 deletions(-)
 create mode 100755 third_party/xla/third_party/triton/common/llvm_cl998243985.patch

diff --git a/third_party/xla/third_party/llvm/workspace.bzl b/third_party/xla/third_party/llvm/workspace.bzl
index e42d7dbdd78..5906818eabd 100644
--- a/third_party/xla/third_party/llvm/workspace.bzl
+++ b/third_party/xla/third_party/llvm/workspace.bzl
@@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "2e0d9e12b36e853d9b108c4da09509cd54d1f3f9"
-    LLVM_SHA256 = "8e4a851780af109d623fa6d6db9b1e62d7126272d17268fed1b149f0e737c1d7"
+    LLVM_COMMIT = "4a0a7e3ddc3bdc4fde7c5b9b9e6890fb4dbf9153"
+    LLVM_SHA256 = "1453c2abaa0fdb957a6108a1802d282bd01010dcde904288c698d44f7539c488"
 
     tf_http_archive(
         name = name,
diff --git a/third_party/xla/third_party/shardy/temporary.patch b/third_party/xla/third_party/shardy/temporary.patch
index 6614bd1c1a7..4d53e31bb73 100644
--- a/third_party/xla/third_party/shardy/temporary.patch
+++ b/third_party/xla/third_party/shardy/temporary.patch
@@ -1,826 +1,109 @@
-diff --git a/shardy/dialect/sdy/ir/canonicalization.cc b/shardy/dialect/sdy/ir/canonicalization.cc
-index 55f7ebcf..9af43e32 100644
---- a/shardy/dialect/sdy/ir/canonicalization.cc
-+++ b/shardy/dialect/sdy/ir/canonicalization.cc
-@@ -25,6 +25,7 @@ limitations under the License.
- #include "llvm/ADT/SmallVector.h"
- #include "llvm/Support/MathExtras.h"
- #include "mlir/IR/BuiltinAttributes.h"
-+#include "mlir/IR/BuiltinTypes.h"
- #include "mlir/IR/Diagnostics.h"
- #include "mlir/IR/MLIRContext.h"
- #include "mlir/IR/OperationSupport.h"
-@@ -203,6 +204,42 @@ std::optional<ReduceScatterFusionInfo> matchAndComputeReduceScatterFusionAxes(
-   SmallVector<DimensionShardingAttr> reduceScatterOutDimShardings;
-   reduceScatterOutDimShardings.reserve(rank);
+diff --git a/shardy/dialect/sdy/ir/dialect.cc b/shardy/dialect/sdy/ir/dialect.cc
+index 719686a5..53082307 100644
+--- a/shardy/dialect/sdy/ir/dialect.cc
++++ b/shardy/dialect/sdy/ir/dialect.cc
+@@ -969,9 +969,7 @@ Type TensorShardingAttr::getLocalType(Type globalType, MeshAttr mesh,
+ RankedTensorType TensorShardingAttr::getLocalTensorType(
+     RankedTensorType globalTensorType, MeshAttr mesh,
+     bool allowNonDivisible) const {
+-  // getLocalType should always return something castable to RankedTensorType if
+-  // a RankedTensorType is passed in.
+-  return cast<RankedTensorType>(
++  return dyn_cast_or_null<RankedTensorType>(
+       getLocalType(globalTensorType, mesh, allowNonDivisible));
+ }
  
-+  // TODO(b/553078138): This is a conservative workaround. The proper fix is to
-+  // make a reduce-scatter with unreduced axes exportable for the non-divisible
-+  // case (in `convertReduceScatter`); once that lands, this bail-out (which
-+  // falls back to the less efficient `all_reduce` + `all_slice` form) can be
-+  // removed so these cases fuse as well.
-+  // Bail out of the fusion if it would create an `sdy.reduce_scatter` that
-+  // scatters a dimension by a group size that does not evenly divide the local
-+  // (already-sharded) dimension size. Exporting a reduce-scatter with unreduced
-+  // axes (see `convertReduceScatter` in export_manual_reduction_collectives)
-+  // requires the scattered dimension to be divisible by the group size, and
-+  // would otherwise fail. Keeping the `all_reduce` + `all_slice` form preserves
-+  // correct behavior for these non-divisible cases.
-+  if (auto operandType =
-+          dyn_cast<RankedTensorType>(allReduceOp.getTensor().getType())) {
-+    // Local (per-device) shape of the all-reduce operand under its existing
-+    // sharding. `allowNonDivisible=true` handles the case where the existing
-+    // sharding does not evenly divide a dimension: the local size is rounded up
-+    // (ceil), matching how Shardy pads such dimensions.
-+    RankedTensorType localType = allReduceInSharding.getLocalTensorType(
-+        operandType, mesh, /*allowNonDivisible=*/true);
-+    for (int64_t dim = 0; dim < rank; ++dim) {
-+      int64_t groupSize =
-+          getTotalAxesSize(reduceScatterAxes[dim].getValue(), mesh);
-+      if (groupSize <= 1) {
-+        continue;
-+      }
-+      int64_t localDimSize = localType.getDimSize(dim);
-+      if (ShapedType::isDynamic(localDimSize)) {
-+        continue;
-+      }
-+      if (localDimSize % groupSize != 0) {
-+        return std::nullopt;
-+      }
-+    }
-+  }
+diff --git a/shardy/dialect/sdy/ir/dialect_test.cc b/shardy/dialect/sdy/ir/dialect_test.cc
+index 1ea9160e..8f149153 100644
+--- a/shardy/dialect/sdy/ir/dialect_test.cc
++++ b/shardy/dialect/sdy/ir/dialect_test.cc
+@@ -624,12 +624,25 @@ TEST_F(DialectTest, TensorShardingAttrGetLocalTensorType) {
+       {createDimSharding({createAxis("y")}), createDimSharding({})});
+   EXPECT_EQ(sharding3.getLocalTensorType(globalType, mesh),
+             RankedTensorType::get({3, 4}, builder.getI32Type()));
++  EXPECT_EQ(sharding3.getLocalTensorType(globalType, mesh,
++                                         /*allowNonDivisible=*/false),
++            nullptr);
+ 
+   TensorShardingAttr sharding4 = createTensorSharding(
+       {createDimSharding({createAxis("x"), createAxis("z")}),
+        createDimSharding({createAxis("y")})});
+   EXPECT_EQ(sharding4.getLocalTensorType(globalType, mesh),
+             RankedTensorType::get({1, 2}, builder.getI32Type()));
++  EXPECT_EQ(sharding4.getLocalTensorType(globalType, mesh,
++                                         /*allowNonDivisible=*/false),
++            nullptr);
 +
-   for (auto [operandDimSharding, reduceScatterAxesPerDim] : llvm::zip_equal(
-            allReduceInSharding.getDimShardings(), reduceScatterAxes)) {
-     SmallVector<AxisRefAttr> expectedDimShardingAxes =
-diff --git a/shardy/dialect/sdy/ir/extensions/stablehlo_extensions.cc b/shardy/dialect/sdy/ir/extensions/stablehlo_extensions.cc
-index d4825b2e..75aac4e0 100644
---- a/shardy/dialect/sdy/ir/extensions/stablehlo_extensions.cc
-+++ b/shardy/dialect/sdy/ir/extensions/stablehlo_extensions.cc
-@@ -70,20 +70,20 @@ struct WhileShardableDataFlowOpInterface
-   }
- };
- 
--template <typename OpTy>
--struct RegionBranchShardableDataFlowOpInterface
-+struct CaseShardableDataFlowOpInterface
-     : public ShardableDataFlowOpInterface::ExternalModel<
--          RegionBranchShardableDataFlowOpInterface<OpTy>, OpTy> {
-+          CaseShardableDataFlowOpInterface, stablehlo::CaseOp> {
-   ResultRange getOpResultEdgeOwners(Operation *op) const {
-     return op->getResults();
-   }
- 
-   SmallVector<OpOperand *> getEdgeSources(Operation *op, Value owner) const {
-+    auto caseOp = cast<stablehlo::CaseOp>(op);
-     auto opResult = dyn_cast<OpResult>(owner);
-     assert(opResult && opResult.getOwner() == op);
-     SmallVector<OpOperand *> sources;
--    sources.reserve(op->getNumRegions());
--    for (Region& branch : op->getRegions()) {
-+    sources.reserve(caseOp.getBranches().size());
-+    for (Region &branch : caseOp.getBranches()) {
-       sources.push_back(&branch.front().getTerminator()->getOpOperand(
-           opResult.getResultNumber()));
-     }
-@@ -132,10 +132,7 @@ void registerStablehloExtensions(MLIRContext *ctx) {
-   // Ensure dialect is loaded before attaching interfaces.
-   ctx->loadDialect<stablehlo::StablehloDialect>();
-   stablehlo::WhileOp::attachInterface<WhileShardableDataFlowOpInterface>(*ctx);
--  stablehlo::CaseOp::attachInterface<
--      RegionBranchShardableDataFlowOpInterface<stablehlo::CaseOp>>(*ctx);
--  stablehlo::IfOp::attachInterface<
--      RegionBranchShardableDataFlowOpInterface<stablehlo::IfOp>>(*ctx);
-+  stablehlo::CaseOp::attachInterface<CaseShardableDataFlowOpInterface>(*ctx);
-   stablehlo::OptimizationBarrierOp::attachInterface<
-       OptBarrierShardableDataFlowOpInterface>(*ctx);
- }
-diff --git a/shardy/dialect/sdy/ir/test/collective_canonicalization.mlir b/shardy/dialect/sdy/ir/test/collective_canonicalization.mlir
-index de616ede..14b0acf3 100644
---- a/shardy/dialect/sdy/ir/test/collective_canonicalization.mlir
-+++ b/shardy/dialect/sdy/ir/test/collective_canonicalization.mlir
-@@ -115,15 +115,33 @@ func.func @reduce_scatter_fusion4(%arg0 : tensor<64x8xf32> {sdy.sharding=#sdy.sh
-   return %1 : tensor<64x8xf32>
++  TensorShardingAttr sharding5 = createTensorSharding(
++      {createDimSharding({createAxis("x"), createAxis("z")}),
++       createDimSharding({createAxis("x")})});
++  EXPECT_EQ(sharding5.getLocalTensorType(globalType, mesh,
++                                         /*allowNonDivisible=*/false),
++            RankedTensorType::get({1, 2}, builder.getI32Type()));
  }
  
-+// The local size of dim 0 is 2 / size("x"=2) = 1, which is not divisible by the
-+// reduce-scatter group size (size("y"=2)). The fusion must be skipped and the
-+// original `all_reduce` + `all_slice` preserved, otherwise the resulting
-+// `reduce_scatter` cannot be exported (see `convertReduceScatter` in
-+// export_manual_reduction_collectives).
- // CHECK-LABEL: func @reduce_scatter_fusion_non_divisible
- func.func @reduce_scatter_fusion_non_divisible(%arg0 : tensor<2x4xf32> {sdy.sharding=#sdy.sharding<@mesh, [{"x"}, {}]>}) -> tensor<2x4xf32> {
--  // CHECK-NEXT: %0 = sdy.reduce_scatter [{"y"}, {}] %arg0 out_sharding=<@mesh, [{"x", "y"}, {}]> : tensor<2x4xf32>
--  // CHECK-NEXT: return %0 : tensor<2x4xf32>
-+  // CHECK-NEXT: %[[ALL_REDUCE:.*]] = sdy.all_reduce {"y"} %arg0 out_sharding=<@mesh, [{"x"}, {}]> : tensor<2x4xf32>
-+  // CHECK-NEXT: %[[ALL_SLICE:.*]] = sdy.all_slice [{"y"}, {}] %[[ALL_REDUCE]] out_sharding=<@mesh, [{"x", "y"}, {}]> : tensor<2x4xf32>
-+  // CHECK-NEXT: return %[[ALL_SLICE]] : tensor<2x4xf32>
-   %0 = sdy.all_reduce {"y"} %arg0 out_sharding=<@mesh, [{"x"}, {}]> : tensor<2x4xf32>
-   %1 = sdy.all_slice [{"y"}, {}] %0 out_sharding=<@mesh, [{"x", "y"}, {}]> : tensor<2x4xf32>
-   return %1 : tensor<2x4xf32>
+ TEST_F(DialectTest, DimensionShardingAttrGetSharded) {
+diff --git a/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc b/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc
+index 8417874d..48217e35 100644
+--- a/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc
++++ b/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc
+@@ -776,7 +776,9 @@ bool allValidShapes(ModuleOp moduleOp) {
+                   }
+                   if (auto tupleType = dyn_cast<TupleType>(type)) {
+                     op->emitError("Shardy propagation doesn't support tuples: ")
+-                        << tupleType;
++                        << tupleType
++                        << ". Please flatten or canonicalize tuples before "
++                           "running Shardy propagation.";
+                     return WalkResult::interrupt();
+                   }
+                 }
+diff --git a/shardy/dialect/sdy/transforms/propagation/test/basic_propagation.mlir b/shardy/dialect/sdy/transforms/propagation/test/basic_propagation.mlir
+index c5cd8022..6d9f32c2 100644
+--- a/shardy/dialect/sdy/transforms/propagation/test/basic_propagation.mlir
++++ b/shardy/dialect/sdy/transforms/propagation/test/basic_propagation.mlir
+@@ -1104,3 +1104,14 @@ func.func private @foo(%arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mes
+   %0 = stablehlo.abs %arg0 : tensor<8x8xf32>
+   return %0: tensor<8x8xf32>
  }
- 
-+// The `all_reduce` operand is sharded indivisibly: dim 0 size 7 with
-+// size("x"=2), so its local size is ceil(7 / 2) = 4, which is divisible by
-+// the reduce-scatter group size (size("y"=2)). Therefore, fusion happens.
-+// CHECK-LABEL: func @reduce_scatter_fusion_non_divisible_all_reduce
-+func.func @reduce_scatter_fusion_non_divisible_all_reduce(%arg0 : tensor<7x4xf32> {sdy.sharding=#sdy.sharding<@mesh, [{"x"}, {}]>}) -> tensor<7x4xf32> {
-+  // CHECK-NEXT: %0 = sdy.reduce_scatter [{"y"}, {}] %arg0 out_sharding=<@mesh, [{"x", "y"}, {}]> : tensor<7x4xf32>
-+  // CHECK-NEXT: return %0 : tensor<7x4xf32>
-+  %0 = sdy.all_reduce {"y"} %arg0 out_sharding=<@mesh, [{"x"}, {}]> : tensor<7x4xf32>
-+  %1 = sdy.all_slice [{"y"}, {}] %0 out_sharding=<@mesh, [{"x", "y"}, {}]> : tensor<7x4xf32>
-+  return %1 : tensor<7x4xf32>
++
++// -----
++
++sdy.mesh @mesh = <["a"=2]>
++
++func.func @tuple_not_supported(%arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}, {}]>}) -> tensor<8x8xf32> {
++  // expected-error @+1 {{Shardy propagation doesn't support tuples: 'tuple<tensor<8x8xf32>>'. Please flatten or canonicalize tuples before running Shardy propagation.}}
++  %0 = stablehlo.tuple %arg0 : tuple<tensor<8x8xf32>>
++  %1 = stablehlo.get_tuple_element %0[0] : (tuple<tensor<8x8xf32>>) -> tensor<8x8xf32>
++  return %1 : tensor<8x8xf32>
 +}
-+
- // CHECK-LABEL: func @reduce_scatter_fusion_missing_in_sharding
- func.func @reduce_scatter_fusion_missing_in_sharding(%arg0 : tensor<16x8xf32>) -> tensor<16x8xf32> {
-   // CHECK-NEXT: %0 = sdy.reduce_scatter [{"x"}, {"y"}] %arg0 out_sharding=<@mesh2, [{"x"}, {"y"}]> : tensor<16x8xf32>
-diff --git a/shardy/dialect/sdy/ir/utils.cc b/shardy/dialect/sdy/ir/utils.cc
-index 590f7149..f6623986 100644
---- a/shardy/dialect/sdy/ir/utils.cc
-+++ b/shardy/dialect/sdy/ir/utils.cc
-@@ -247,13 +247,10 @@ TensorShardingPerValueAttr inlineMesh(
-                                          inlinedShardings);
- }
- 
--namespace {
--
--template <typename SymbolTableOrOp>
--std::pair<Attribute, MeshAttr> getCommonMeshAndRef(
--    ArrayRef<TensorShardingAttr> operandShardings,
--    ArrayRef<TensorShardingAttr> resultsShardings,
--    SymbolTableOrOp symbolTableOrOp, const bool ignoreDeviceIds) {
-+Attribute getCommonMeshOrRef(ArrayRef<TensorShardingAttr> operandShardings,
-+                             ArrayRef<TensorShardingAttr> resultsShardings,
-+                             const SymbolTable& symbolTable,
-+                             const bool ignoreDeviceIds) {
-   Attribute meshOrRef;
-   MeshAttr mesh;
-   for (TensorShardingAttr sharding : llvm::concat<const TensorShardingAttr>(
-@@ -261,7 +258,7 @@ std::pair<Attribute, MeshAttr> getCommonMeshAndRef(
-     if (!sharding || sharding.getMeshOrRef() == meshOrRef) {
-       continue;
-     }
--    MeshAttr otherMesh = sharding.getMesh(symbolTableOrOp);
-+    MeshAttr otherMesh = sharding.getMesh(symbolTable);
-     if (!mesh || mesh.empty()) {
-       mesh = otherMesh;
-       meshOrRef = sharding.getMeshOrRef();
-@@ -272,7 +269,7 @@ std::pair<Attribute, MeshAttr> getCommonMeshAndRef(
-     }
-     if (!otherMesh.equals(mesh, ignoreDeviceIds)) {
-       // Found more than one mesh name.
--      return {nullptr, nullptr};
-+      return nullptr;
-     }
-     // Prefer iota device id over non-iota.
-     if (ignoreDeviceIds && otherMesh.getDeviceIds().empty()) {
-@@ -281,18 +278,7 @@ std::pair<Attribute, MeshAttr> getCommonMeshAndRef(
-     }
-   }
- 
--  return {meshOrRef, mesh};
--}
--
--}  // namespace
--
--Attribute getCommonMeshOrRef(ArrayRef<TensorShardingAttr> operandShardings,
--                             ArrayRef<TensorShardingAttr> resultsShardings,
--                             const SymbolTable& symbolTable,
--                             const bool ignoreDeviceIds) {
--  return getCommonMeshAndRef(operandShardings, resultsShardings, symbolTable,
--                             ignoreDeviceIds)
--      .first;
-+  return meshOrRef;
- }
- 
- MeshAttr getCommonMesh(ArrayRef<TensorShardingAttr> shardings,
-@@ -300,24 +286,21 @@ MeshAttr getCommonMesh(ArrayRef<TensorShardingAttr> shardings,
-   return getCommonMesh(shardings, {}, symbolTable);
- }
- 
--MeshAttr getCommonMesh(ArrayRef<TensorShardingAttr> shardings, Operation* op) {
--  return getCommonMesh(shardings, {}, op);
--}
--
- MeshAttr getCommonMesh(ArrayRef<TensorShardingAttr> operandShardings,
-                        ArrayRef<TensorShardingAttr> resultsShardings,
-                        const SymbolTable& symbolTable) {
--  return getCommonMeshAndRef(operandShardings, resultsShardings, symbolTable,
--                             /*ignoreDeviceIds=*/false)
--      .second;
-+  if (Attribute meshOrRef =
-+          getCommonMeshOrRef(operandShardings, resultsShardings, symbolTable)) {
-+    return getMeshOrLookup(symbolTable, meshOrRef);
-+  }
-+  return nullptr;
- }
- 
- MeshAttr getCommonMesh(ArrayRef<TensorShardingAttr> operandShardings,
-                        ArrayRef<TensorShardingAttr> resultsShardings,
-                        Operation* op) {
--  return getCommonMeshAndRef(operandShardings, resultsShardings, op,
--                             /*ignoreDeviceIds=*/false)
--      .second;
-+  return getCommonMesh(operandShardings, resultsShardings,
-+                       SymbolTable(op->getParentOfType<ModuleOp>()));
- }
- 
- std::optional<StringRef> getCommonMeshName(
-diff --git a/shardy/dialect/sdy/ir/utils.h b/shardy/dialect/sdy/ir/utils.h
-index 3b56a553..22e4d9cc 100644
---- a/shardy/dialect/sdy/ir/utils.h
-+++ b/shardy/dialect/sdy/ir/utils.h
-@@ -195,12 +195,6 @@ Attribute getCommonMeshOrRef(ArrayRef<TensorShardingAttr> operandShardings,
- MeshAttr getCommonMesh(ArrayRef<TensorShardingAttr> shardings,
-                        const SymbolTable& symbolTable);
- 
--// Returns the common `MeshAttr` bound by all the `TensorShardingAttr`s or
--// nullptr if there is none.
--//
--// Ignores empty meshes unless all meshes are empty.
--MeshAttr getCommonMesh(ArrayRef<TensorShardingAttr> shardings, Operation* op);
--
- // Returns the common `MeshAttr` bound by all the `TensorShardingAttr`s or
- // nullptr if there is none.
- //
-diff --git a/shardy/dialect/sdy/ir/utils_test.cc b/shardy/dialect/sdy/ir/utils_test.cc
-index f76d2b16..8dbaf619 100644
---- a/shardy/dialect/sdy/ir/utils_test.cc
-+++ b/shardy/dialect/sdy/ir/utils_test.cc
-@@ -81,27 +81,11 @@ class UtilsTest : public ShardyTestBase {
-     return symbolTableCollection.getSymbolTable(moduleOp.get());
-   }
- 
--  Operation* getModuleOp() { return moduleOp.get(); }
--  Operation* getNestedOp() {
--    return moduleOp->lookupSymbol<func::FuncOp>("main");
--  }
--
-  private:
-   OwningOpRef<ModuleOp> moduleOp;
-   mlir::SymbolTableCollection symbolTableCollection;
- };
- 
--TEST_F(UtilsTest, GetCommonMesh_SingleListWithOperation) {
--  EXPECT_NE(getCommonMesh({createTensorSharding("mesh_ab_23")}, getModuleOp()),
--            nullptr);
--  EXPECT_EQ(
--      getCommonMesh({createTensorSharding("mesh_ab_23")}, getModuleOp()),
--      getCommonMesh({createTensorSharding("mesh_ab_23")}, getSymbolTable()));
--  EXPECT_EQ(
--      getCommonMesh({createTensorSharding("mesh_ab_23")}, getNestedOp()),
--      getCommonMesh({createTensorSharding("mesh_ab_23")}, getSymbolTable()));
--}
--
- TEST_F(UtilsTest, GetCommonMeshName_AllSame) {
-   EXPECT_EQ(getCommonMeshName({createTensorSharding("mesh_ab_23")},
-                               {createTensorSharding("mesh_ab_23")},
-diff --git a/shardy/dialect/sdy/ir/verifiers.cc b/shardy/dialect/sdy/ir/verifiers.cc
-index 09f70f28..65c6212b 100644
---- a/shardy/dialect/sdy/ir/verifiers.cc
-+++ b/shardy/dialect/sdy/ir/verifiers.cc
-@@ -1186,7 +1186,8 @@ LogicalResult verifyPropagationEdgesShardingAttr(
-         "expected propagation edges attr to reference a sharding.");
-   }
- 
--  MeshAttr mesh = getCommonMesh(shardings, op);
-+  MeshAttr mesh =
-+      getCommonMesh(shardings, SymbolTable(op->getParentOfType<ModuleOp>()));
-   if (!mesh) {
-     return op->emitOpError(
-         "expected a common mesh for propagation edges attr.");
-@@ -1195,7 +1196,12 @@ LogicalResult verifyPropagationEdgesShardingAttr(
-   for (PropagationOneStepAttr propagationEdge : propagationEdges) {
-     for (AxisToPropagationDetailsAttr axisEntry :
-          propagationEdge.getAxisEntries()) {
--      if (!mesh.hasAxis(axisEntry.getAxisName().getName())) {
-+      bool meshContainsAxis = false;
-+      if (mesh.hasAxis(axisEntry.getAxisName().getName())) {
-+        meshContainsAxis = true;
-+      }
-+
-+      if (!meshContainsAxis) {
-         return op->emitOpError("expected axis ref to be in one of the meshes");
-       }
- 
-@@ -1284,8 +1290,8 @@ LogicalResult NamedComputationOp::verifySymbolUses(
- 
-   // TODO(pxy): remove this once the `ShardableDataFlowOpInterface` is verified.
-   // Verify the in/out shardings.
--  const SymbolTable& symbolTable = symbolTableCollection.getSymbolTable(
--      getOperation()->getParentOfType<ModuleOp>());
-+  const SymbolTable& symbolTable =
-+      SymbolTable(getOperation()->getParentOfType<ModuleOp>());
-   if (inShardings && failed(verifyTensorShardingPerValueAttr(
-                          *inShardings, getOperandTypes(), *this, symbolTable,
-                          [this](StringRef msg) {
-diff --git a/shardy/dialect/sdy/transforms/export/test/insert_explicit_reshards/data_flow_ops.mlir b/shardy/dialect/sdy/transforms/export/test/insert_explicit_reshards/data_flow_ops.mlir
-index f287d6de..f2891d04 100644
---- a/shardy/dialect/sdy/transforms/export/test/insert_explicit_reshards/data_flow_ops.mlir
-+++ b/shardy/dialect/sdy/transforms/export/test/insert_explicit_reshards/data_flow_ops.mlir
-@@ -48,25 +48,6 @@ func.func @case_argument_used_outside_block(%arg0: tensor<210xf32> {sdy.sharding
-   return %3 : tensor<210xf32>
- }
- 
--// CHECK-LABEL: func @if
--func.func @if(%arg0: tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x":(1)2}]>}, %arg1: tensor<i1>) -> (tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}]>}) {
--  %0 = "stablehlo.if"(%arg1) ({
--    // CHECK: %[[RESHARD:.*]] = sdy.reshard %arg0 <@mesh, [{"x"}]> : tensor<210xf32>
--    // CHECK-NEXT: stablehlo.abs %[[RESHARD]]
--    %2 = stablehlo.abs %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x"}]>]>} : tensor<210xf32>
--    // CHECK: %[[RESHARD:.*]] = sdy.reshard %{{.*}} <@mesh, [{"y"}]> : tensor<210xf32>
--    // CHECK-NEXT: stablehlo.return %[[RESHARD]] : tensor<210xf32>
--    stablehlo.return %2 : tensor<210xf32>
--  }, {
--    %2 = stablehlo.abs %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x":(1)2}]>]>} : tensor<210xf32>
--    // CHECK: %[[RESHARD:.*]] = sdy.reshard %{{.*}} <@mesh, [{"y"}]> : tensor<210xf32>
--    // CHECK-NEXT: stablehlo.return %[[RESHARD]] : tensor<210xf32>
--    stablehlo.return %2 : tensor<210xf32>
--  }) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>} : (tensor<i1>) -> tensor<210xf32>
--  %1 = stablehlo.negate %0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"y"}]>]>} : tensor<210xf32>
--  return %1 : tensor<210xf32>
--}
--
- // CHECK-LABEL: func @while
- func.func @while(%arg0: tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}]>}) -> (tensor<210xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}]>}) {
-   %c = stablehlo.constant dense<0> : tensor<i32>
-diff --git a/shardy/dialect/sdy/transforms/import/test/add_data_flow_edges.mlir b/shardy/dialect/sdy/transforms/import/test/add_data_flow_edges.mlir
-index 6b080b16..bd472047 100644
---- a/shardy/dialect/sdy/transforms/import/test/add_data_flow_edges.mlir
-+++ b/shardy/dialect/sdy/transforms/import/test/add_data_flow_edges.mlir
-@@ -55,24 +55,6 @@ func.func @case_token_result_skipped(
-   return %0#0, %0#1 : !stablehlo.token, tensor<4xi64>
- }
- 
--// CHECK-LABEL: func @if
--func.func @if(%arg0: tensor<i1>, %arg1: tensor<8xi64>, %arg2: tensor<8xi64>,
--              %arg3: tensor<8xi64>, %arg4: tensor<8xi64>)
--    -> (tensor<8xi64>, tensor<8xi64>) {
--  // CHECK-NEXT: %[[IF:.*]]:2 = "stablehlo.if"(%arg0)
--  // CHECK:      %[[EDGE_1:.*]] = sdy.data_flow_edge %[[IF]]#0 : tensor<8xi64>
--  // CHECK-NEXT: %[[EDGE_2:.*]] = sdy.data_flow_edge %[[IF]]#1 : tensor<8xi64>
--  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[EDGE_2]], %[[EDGE_2]]
--  // CHECK-NEXT: return %[[EDGE_1]], %[[ADD]]
--  %0:2 = "stablehlo.if"(%arg0) ({
--    stablehlo.return %arg1, %arg2 : tensor<8xi64>, tensor<8xi64>
--  }, {
--    stablehlo.return %arg3, %arg4 : tensor<8xi64>, tensor<8xi64>
--  }) : (tensor<i1>) -> (tensor<8xi64>, tensor<8xi64>)
--  %1 = stablehlo.add %0#1, %0#1 : tensor<8xi64>
--  return %0#0, %1 : tensor<8xi64>, tensor<8xi64>
--}
--
- // CHECK-LABEL: func @optimization_barrier
- func.func @optimization_barrier(%arg0: tensor<32x96xf32>, %arg1: tensor<32x96xf32>)
-     -> (tensor<32x96xf32>, tensor<32x96xf32>) {
-diff --git a/shardy/dialect/sdy/transforms/import/test/apply_sharding_constraints.mlir b/shardy/dialect/sdy/transforms/import/test/apply_sharding_constraints.mlir
-index 445f25fd..11d064d5 100644
---- a/shardy/dialect/sdy/transforms/import/test/apply_sharding_constraints.mlir
-+++ b/shardy/dialect/sdy/transforms/import/test/apply_sharding_constraints.mlir
-@@ -652,275 +652,6 @@ func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
- 
- // -----
- 
--// A flat call graph (@main -> @foo) without func data flow edges, applying a
--// single sharding constraint to the argument and call result. Without func data
--// flow edges, closed shardings are attached directly to the function argument
--// and call result and other uses are not replaced with the sharding constraint.
--sdy.mesh @mesh = <["a"=2, "b"=2]>
--
--// CHECK-LABEL: func @foo(
--// CHECK-SAME:    %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}, {}]>}) -> tensor<8x8xf32> {
--func.func @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %arg0, %arg0 : tensor<8x8xf32>
--  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %arg0 <@mesh, [{"a"}, {}]> {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=8, j=8}>} : tensor<8x8xf32>
--  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
--  %0 = stablehlo.add %arg0, %arg0 : tensor<8x8xf32>
--  %1 = sdy.sharding_constraint %arg0 <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
--  return %0 : tensor<8x8xf32>
--}
--
--// CHECK-LABEL: func @main
--func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  // CHECK-NEXT: %[[CALL:.*]] = call @foo(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"b"}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[CALL]], %[[CALL]] : tensor<8x8xf32>
--  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %[[CALL]] <@mesh, [{}, {"b"}]> {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=8, j=8}>} : tensor<8x8xf32>
--  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
--  %0 = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  %1 = stablehlo.add %0, %0 : tensor<8x8xf32>
--  %2 = sdy.sharding_constraint %0 <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
--  return %1 : tensor<8x8xf32>
--}
--
--// -----
--
--// A flat call graph (@main -> @foo) applying a single sharding constraint to
--// the func data flow edges. The func data flow edges is left unsharded, the
--// constraint is moved right after the edge, and other uses are replaced with
--// the constraint.
--sdy.mesh @mesh = <["a"=2, "b"=2]>
--
--// CHECK-LABEL: func @foo
--func.func @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE:.*]] = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
--  // CHECK-NOT: sdy.sharding
--  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %[[FUNC_DATA_FLOW_EDGE]] <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
--  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[WSC]], %[[WSC]] : tensor<8x8xf32>
--  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
--  %0 = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
--  %1 = stablehlo.add %0, %0 : tensor<8x8xf32>
--  %2 = sdy.sharding_constraint %0 <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
--  return %1 : tensor<8x8xf32>
--}
--
--// CHECK-LABEL: func @main
--func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  // CHECK-NEXT: %[[CALL:.*]] = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE:.*]] = sdy.func_data_flow_edge %[[CALL]] : tensor<8x8xf32>
--  // CHECK-NOT: sdy.sharding
--  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %[[FUNC_DATA_FLOW_EDGE]] <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
--  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[WSC]], %[[WSC]] : tensor<8x8xf32>
--  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
--  %0 = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  %1 = sdy.func_data_flow_edge %0 : tensor<8x8xf32>
--  %2 = stablehlo.add %1, %1 : tensor<8x8xf32>
--  %3 = sdy.sharding_constraint %1 <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
--  return %2 : tensor<8x8xf32>
--}
--
--// -----
--
--// A non-flat call graph (@main -> @foo -> @bar and @main -> @bar) without
--// func data flow edges, applying a single sharding constraint to the argument
--// and the call results. Without func data flow edges, closed shardings are
--// attached directly to the argument and both results, and other uses are not
--// replaced with the sharding constraint.
--sdy.mesh @mesh = <["a"=2, "b"=2]>
--
--// CHECK-LABEL: func @bar(
--// CHECK-SAME:    %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}, {}]>}) -> tensor<8x8xf32> {
--func.func @bar(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %arg0, %arg0 : tensor<8x8xf32>
--  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %arg0 <@mesh, [{"a"}, {}]> {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=8, j=8}>} : tensor<8x8xf32>
--  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
--  %0 = stablehlo.add %arg0, %arg0 : tensor<8x8xf32>
--  %1 = sdy.sharding_constraint %arg0 <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
--  return %0 : tensor<8x8xf32>
--}
--
--// CHECK-LABEL: func @foo
--func.func @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  // CHECK-NEXT: %[[CALL:.*]] = call @bar(%arg0) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"b"}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[CALL]], %[[CALL]] : tensor<8x8xf32>
--  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %[[CALL]] <@mesh, [{}, {"b"}]> {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=8, j=8}>} : tensor<8x8xf32>
--  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
--  %0 = call @bar(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  %1 = stablehlo.add %0, %0 : tensor<8x8xf32>
--  %2 = sdy.sharding_constraint %0 <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
--  return %1 : tensor<8x8xf32>
--}
--
--// CHECK-LABEL: func @main
--func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  // CHECK-NEXT: %[[CALL_0:.*]] = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  // CHECK-NEXT: %[[CALL_1:.*]] = call @bar(%[[CALL_0]]) {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {"b"}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[CALL_1]], %[[CALL_1]] : tensor<8x8xf32>
--  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %[[CALL_1]] <@mesh, [{"a"}, {"b"}]> {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=8, j=8}>} : tensor<8x8xf32>
--  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
--  %0 = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  %1 = call @bar(%0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  %2 = stablehlo.add %1, %1 : tensor<8x8xf32>
--  %3 = sdy.sharding_constraint %1 <@mesh, [{"a"}, {"b"}]> : tensor<8x8xf32>
--  return %2 : tensor<8x8xf32>
--}
--
--// -----
--
--// A non-flat call graph (@main -> @foo -> @bar and @main -> @bar) applying a
--// single sharding constraint to the func data flow edges on @bar's argument
--// and the call results in @foo and @main. Each func data flow edges is left
--// unsharded, the constraint is moved right after the edge, and other uses are
--// replaced with the constraint.
--sdy.mesh @mesh = <["a"=2, "b"=2]>
--
--// CHECK-LABEL: func @bar
--func.func @bar(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE:.*]] = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
--  // CHECK-NOT: sdy.sharding
--  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %[[FUNC_DATA_FLOW_EDGE]] <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
--  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[WSC]], %[[WSC]] : tensor<8x8xf32>
--  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
--  %0 = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
--  %1 = stablehlo.add %0, %0 : tensor<8x8xf32>
--  %2 = sdy.sharding_constraint %0 <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
--  return %1 : tensor<8x8xf32>
--}
--
--// CHECK-LABEL: func @foo
--func.func @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE_0:.*]] = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
--  // CHECK-NEXT: %[[CALL:.*]] = call @bar(%[[FUNC_DATA_FLOW_EDGE_0]])
--  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE_1:.*]] = sdy.func_data_flow_edge %[[CALL]] : tensor<8x8xf32>
--  // CHECK-NOT: sdy.sharding
--  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %[[FUNC_DATA_FLOW_EDGE_1]] <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
--  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[WSC]], %[[WSC]] : tensor<8x8xf32>
--  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
--  %0 = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
--  %1 = call @bar(%0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  %2 = sdy.func_data_flow_edge %1 : tensor<8x8xf32>
--  %3 = stablehlo.add %2, %2 : tensor<8x8xf32>
--  %4 = sdy.sharding_constraint %2 <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
--  return %3 : tensor<8x8xf32>
--}
--
--// CHECK-LABEL: func @main
--func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  // CHECK-NEXT: %[[CALL_0:.*]] = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE_0:.*]] = sdy.func_data_flow_edge %[[CALL_0]] : tensor<8x8xf32>
--  // CHECK-NEXT: %[[CALL_1:.*]] = call @bar(%[[FUNC_DATA_FLOW_EDGE_0]]) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE_1:.*]] = sdy.func_data_flow_edge %[[CALL_1]] : tensor<8x8xf32>
--  // CHECK-NOT: sdy.sharding
--  // CHECK-NEXT: %[[WSC:.*]] = sdy.sharding_constraint %[[FUNC_DATA_FLOW_EDGE_1]] <@mesh, [{"a"}, {"b"}]> : tensor<8x8xf32>
--  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[WSC]], %[[WSC]] : tensor<8x8xf32>
--  // CHECK-NEXT: return %[[ADD]] : tensor<8x8xf32>
--  %0 = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  %1 = sdy.func_data_flow_edge %0 : tensor<8x8xf32>
--  %2 = call @bar(%1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  %3 = sdy.func_data_flow_edge %2 : tensor<8x8xf32>
--  %4 = stablehlo.add %3, %3 : tensor<8x8xf32>
--  %5 = sdy.sharding_constraint %3 <@mesh, [{"a"}, {"b"}]> : tensor<8x8xf32>
--  return %4 : tensor<8x8xf32>
--}
--
--// -----
--
--// Non-flat call graph (@main -> @foo -> @bar and @main -> @bar) without
--// func data flow edges, applying a chain of two sharding constraints
--// to stablehlo.add in @bar and @foo. Verifies that the head constraint's
--// sharding is attached to `stablehlo.add` and uses after the chain (including
--// the func returns in @bar and @foo) are replaced with the tail of the chain.
--sdy.mesh @mesh = <["a"=2, "b"=2]>
--
--// CHECK-LABEL: func @bar
--func.func @bar(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %arg0, %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}]>]>}
--  // CHECK-NEXT: %[[WSC_0:.*]] = sdy.sharding_constraint %[[ADD]] <@mesh, [{"a"}, {}]>
--  // CHECK-NEXT: %[[WSC_1:.*]] = sdy.sharding_constraint %[[WSC_0]] <@mesh, [{}, {"b"}]>
--  // CHECK-NEXT: return %[[WSC_1]] : tensor<8x8xf32>
--  %0 = stablehlo.add %arg0, %arg0 : tensor<8x8xf32>
--  %1 = sdy.sharding_constraint %0 <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
--  %2 = sdy.sharding_constraint %1 <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
--  return %0 : tensor<8x8xf32>
--}
--
--// CHECK-LABEL: func @foo
--func.func @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  // CHECK-NEXT: %[[CALL:.*]] = call @bar(%arg0)
--  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[CALL]], %[[CALL]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}]>]>}
--  // CHECK-NEXT: %[[WSC_0:.*]] = sdy.sharding_constraint %[[ADD]] <@mesh, [{"a"}, {}]>
--  // CHECK-NEXT: %[[WSC_1:.*]] = sdy.sharding_constraint %[[WSC_0]] <@mesh, [{}, {"b"}]>
--  // CHECK-NEXT: return %[[WSC_1]] : tensor<8x8xf32>
--  %0 = call @bar(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  %1 = stablehlo.add %0, %0 : tensor<8x8xf32>
--  %2 = sdy.sharding_constraint %1 <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
--  %3 = sdy.sharding_constraint %2 <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
--  return %1 : tensor<8x8xf32>
--}
--
--// CHECK-LABEL: func @main
--func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  // CHECK-NEXT: %[[CALL_0:.*]] = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  // CHECK-NEXT: %[[CALL_1:.*]] = call @bar(%[[CALL_0]]) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  // CHECK-NEXT: return %[[CALL_1]] : tensor<8x8xf32>
--  %0 = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  %1 = call @bar(%0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  return %1 : tensor<8x8xf32>
--}
--
--// -----
--
--// A non-flat call graph (@main -> @foo -> @bar and @main -> @bar) with func
--// data flow edges on function arguments and call results, applying a chain of
--// two sharding constraints to stablehlo.add in @bar and @foo.
--sdy.mesh @mesh = <["a"=2, "b"=2]>
--
--// CHECK-LABEL: func @bar
--func.func @bar(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE:.*]] = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
--  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[FUNC_DATA_FLOW_EDGE]], %[[FUNC_DATA_FLOW_EDGE]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}]>]>}
--  // CHECK-NEXT: %[[WSC_0:.*]] = sdy.sharding_constraint %[[ADD]] <@mesh, [{"a"}, {}]>
--  // CHECK-NEXT: %[[WSC_1:.*]] = sdy.sharding_constraint %[[WSC_0]] <@mesh, [{}, {"b"}]>
--  // CHECK-NEXT: return %[[WSC_1]] : tensor<8x8xf32>
--  %0 = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
--  %1 = stablehlo.add %0, %0 : tensor<8x8xf32>
--  %2 = sdy.sharding_constraint %1 <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
--  %3 = sdy.sharding_constraint %2 <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
--  return %1 : tensor<8x8xf32>
--}
--
--// CHECK-LABEL: func @foo
--func.func @foo(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE_0:.*]] = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
--  // CHECK-NEXT: %[[CALL:.*]] = call @bar(%[[FUNC_DATA_FLOW_EDGE_0]])
--  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE_1:.*]] = sdy.func_data_flow_edge %[[CALL]] : tensor<8x8xf32>
--  // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[FUNC_DATA_FLOW_EDGE_1]], %[[FUNC_DATA_FLOW_EDGE_1]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}, {}]>]>}
--  // CHECK-NEXT: %[[WSC_0:.*]] = sdy.sharding_constraint %[[ADD]] <@mesh, [{"a"}, {}]>
--  // CHECK-NEXT: %[[WSC_1:.*]] = sdy.sharding_constraint %[[WSC_0]] <@mesh, [{}, {"b"}]>
--  // CHECK-NEXT: return %[[WSC_1]] : tensor<8x8xf32>
--  %0 = sdy.func_data_flow_edge %arg0 : tensor<8x8xf32>
--  %1 = call @bar(%0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  %2 = sdy.func_data_flow_edge %1 : tensor<8x8xf32>
--  %3 = stablehlo.add %2, %2 : tensor<8x8xf32>
--  %4 = sdy.sharding_constraint %3 <@mesh, [{"a"}, {}]> : tensor<8x8xf32>
--  %5 = sdy.sharding_constraint %4 <@mesh, [{}, {"b"}]> : tensor<8x8xf32>
--  return %3 : tensor<8x8xf32>
--}
--
--// CHECK-LABEL: func @main
--func.func @main(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
--  // CHECK-NEXT: %[[CALL_0:.*]] = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE_0:.*]] = sdy.func_data_flow_edge %[[CALL_0]] : tensor<8x8xf32>
--  // CHECK-NEXT: %[[CALL_1:.*]] = call @bar(%[[FUNC_DATA_FLOW_EDGE_0]]) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  // CHECK-NEXT: %[[FUNC_DATA_FLOW_EDGE_1:.*]] = sdy.func_data_flow_edge %[[CALL_1]] : tensor<8x8xf32>
--  // CHECK-NEXT: return %[[FUNC_DATA_FLOW_EDGE_1]] : tensor<8x8xf32>
--  %0 = call @foo(%arg0) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  %1 = sdy.func_data_flow_edge %0 : tensor<8x8xf32>
--  %2 = call @bar(%1) : (tensor<8x8xf32>) -> tensor<8x8xf32>
--  %3 = sdy.func_data_flow_edge %2 : tensor<8x8xf32>
--  return %3 : tensor<8x8xf32>
--}
--
--// -----
--
- sdy.mesh @mesh = <["a"=2, "b"=2]>
- 
- // CHECK-LABEL: func @main
-diff --git a/shardy/dialect/sdy/transforms/propagation/op_sharding_rule_registry.cc b/shardy/dialect/sdy/transforms/propagation/op_sharding_rule_registry.cc
-index dc069471..e88aa109 100644
---- a/shardy/dialect/sdy/transforms/propagation/op_sharding_rule_registry.cc
-+++ b/shardy/dialect/sdy/transforms/propagation/op_sharding_rule_registry.cc
-@@ -1321,7 +1321,7 @@ OpShardingRuleAttr createOpShardingRule(Operation* op,
-             PropagationBarrierOp, ShardableDataFlowOpInterface, ShardingGroupOp,
-             ReshardOp, stablehlo::AfterAllOp, stablehlo::CaseOp,
-             stablehlo::ConstantOp, stablehlo::CreateTokenOp,
--            stablehlo::GetTupleElementOp, stablehlo::IfOp, stablehlo::InfeedOp,
-+            stablehlo::GetTupleElementOp, stablehlo::InfeedOp,
-             stablehlo::IotaOp, stablehlo::OutfeedOp,
-             stablehlo::OptimizationBarrierOp, stablehlo::PartitionIdOp,
-             stablehlo::RecvOp, stablehlo::SendOp, stablehlo::WhileOp>(
-diff --git a/shardy/dialect/sdy/transforms/propagation/sharding_group_map.cc b/shardy/dialect/sdy/transforms/propagation/sharding_group_map.cc
-index f61c0dd3..2035990f 100644
---- a/shardy/dialect/sdy/transforms/propagation/sharding_group_map.cc
-+++ b/shardy/dialect/sdy/transforms/propagation/sharding_group_map.cc
-@@ -15,10 +15,12 @@ limitations under the License.
- 
- #include "shardy/dialect/sdy/transforms/propagation/sharding_group_map.h"
- 
-+#include <algorithm>
-+#include <cassert>
- #include <cstdint>
- 
- #include "llvm/ADT/DenseMap.h"
--#include "llvm/ADT/MapVector.h"
-+#include "llvm/ADT/STLExtras.h"
- #include "llvm/Support/ErrorHandling.h"
- #include "mlir/IR/BuiltinOps.h"
- #include "mlir/IR/Diagnostics.h"
-@@ -34,19 +36,21 @@ namespace sdy {
- 
- ShardingGroupMap::ShardingGroupMap(ModuleOp moduleOp) {
-   moduleOp.walk([&](ShardingGroupOp op) {
--    int64_t groupId = op.getGroupId();
-+    // After canonicalization all group ids will take distinct values in the
-+    // range 0,1,...,N. Because of this we can directly index these group ids
-+    // into the shardingGroupToValues vector (resizing when necessary).
-+    shardingGroupToValues.resize(
-+        std::max(op.getGroupId() + 1,
-+                 static_cast<uint64_t>(shardingGroupToValues.size())));
-     // Each value can only map to one sharding group id after
-     // canonicalization.
-     auto [it, inserted] =
--        valueToShardingGroup.try_emplace(op.getInput(), groupId);
--    if (!inserted) {
--      if (it->getSecond() != groupId) {
--        llvm::report_fatal_error(
--            "Value can only map to one sharding group id after import.");
--      }
--      return;
-+        valueToShardingGroup.try_emplace(op.getInput(), op.getGroupId());
-+    if (!inserted && it->getSecond() != op.getGroupId()) {
-+      llvm::report_fatal_error(
-+          "Value can only map to one sharding group id after import.");
-     }
--    shardingGroupToValues[groupId].push_back(op.getInput());
-+    shardingGroupToValues[op.getGroupId()].push_back(op.getInput());
-   });
- }
- 
-@@ -96,7 +100,7 @@ CommonSharding findCommonSharding(int64_t groupId, ValueRange groupMembers) {
- }  // namespace
- 
- void ShardingGroupMap::syncGroupMemberShardings(ModuleOp module) {
--  for (auto& [groupId, groupMembers] : shardingGroupToValues) {
-+  for (auto [groupId, groupMembers] : llvm::enumerate(shardingGroupToValues)) {
-     auto [sharding, hasConflict] = findCommonSharding(groupId, groupMembers);
-     if (!sharding) {
-       continue;
-@@ -130,7 +134,8 @@ void ShardingGroupMap::syncGroupMemberShardings(ModuleOp module) {
- ValueRange ShardingGroupMap::getGroupMembers(const Value& value) const {
-   if (auto it = valueToShardingGroup.find(value);
-       it != valueToShardingGroup.end()) {
--    return shardingGroupToValues.find(it->getSecond())->second;
-+    int64_t shardingGroupId = it->getSecond();
-+    return shardingGroupToValues[shardingGroupId];
-   }
-   return {};
- }
-diff --git a/shardy/dialect/sdy/transforms/propagation/sharding_group_map.h b/shardy/dialect/sdy/transforms/propagation/sharding_group_map.h
-index b924ca1a..90b38bd4 100644
---- a/shardy/dialect/sdy/transforms/propagation/sharding_group_map.h
-+++ b/shardy/dialect/sdy/transforms/propagation/sharding_group_map.h
-@@ -19,7 +19,6 @@ limitations under the License.
- #include <cstdint>
- 
- #include "llvm/ADT/DenseMap.h"
--#include "llvm/ADT/MapVector.h"
- #include "mlir/IR/BuiltinOps.h"
- #include "mlir/IR/Value.h"
- #include "mlir/IR/ValueRange.h"
-@@ -46,7 +45,7 @@ class ShardingGroupMap {
-   ValueRange getGroupMembers(const Value& value) const;
- 
-  private:
--  llvm::MapVector<int64_t, SmallVector<Value>> shardingGroupToValues;
-+  SmallVector<SmallVector<Value>> shardingGroupToValues;
-   llvm::SmallDenseMap<Value, int64_t> valueToShardingGroup;
- };
- 
-diff --git a/shardy/dialect/sdy/transforms/propagation/test/sharding_group_propagation.mlir b/shardy/dialect/sdy/transforms/propagation/test/sharding_group_propagation.mlir
-index 36ae5341..b9edf866 100644
---- a/shardy/dialect/sdy/transforms/propagation/test/sharding_group_propagation.mlir
-+++ b/shardy/dialect/sdy/transforms/propagation/test/sharding_group_propagation.mlir
-@@ -131,31 +131,6 @@ func.func @different_group_ids_greater_group_id_first(
-   return %arg0, %arg1, %arg2, %arg3 : tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>
- }
- 
--// -----
--sdy.mesh @mesh = <["a"=2, "b"=2]>
--
--// CHECK-LABEL: func @undensified_negative_and_large_group_ids
--// CHECK-SAME:    %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a", ?}, {?}]>},
--// CHECK-SAME:    %arg1: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a", ?}, {?}]>},
--// CHECK-SAME:    %arg2: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{?}, {"b", ?}]>},
--// CHECK-SAME:    %arg3: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{?}, {"b", ?}]>})
--// CHECK-SAME:    -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a", ?}, {?}]>},
--// CHECK-SAME:        tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a", ?}, {?}]>},
--// CHECK-SAME:        tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{?}, {"b", ?}]>},
--// CHECK-SAME:        tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{?}, {"b", ?}]>})
--func.func @undensified_negative_and_large_group_ids(
--  %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a", ?}, {?}]>},
--  %arg1: tensor<8x8xf32>,
--  %arg2: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{?}, {"b", ?}]>},
--  %arg3: tensor<8x8xf32>)
--   -> (tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>) {
--  sdy.sharding_group %arg0 group_id=-1 : tensor<8x8xf32>
--  sdy.sharding_group %arg1 group_id=-1 : tensor<8x8xf32>
--  sdy.sharding_group %arg2 group_id=100000000000 : tensor<8x8xf32>
--  sdy.sharding_group %arg3 group_id=100000000000 : tensor<8x8xf32>
--  return %arg0, %arg1, %arg2, %arg3 : tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>
--}
--
- // -----
- 
- // Tests for initial member sharding sync
-diff --git a/third_party/llvm/generated.patch b/third_party/llvm/generated.patch
-index e9c206ce..509398da 100644
---- a/third_party/llvm/generated.patch
-+++ b/third_party/llvm/generated.patch
-@@ -1,12 +1 @@
- Auto generated patch. Do not edit or delete it, even if empty.
--diff -ruN --strip-trailing-cr a/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll b/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll
----- a/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll
--+++ b/llvm/test/CodeGen/Hexagon/autohvx/xqf-postra-conv-double2.ll
--@@ -1,5 +1,7 @@
-- ; Checks if conversion is only inserted for the spilled register
-- ; instead of 2 conversions for the W register
--+
--+; REQUIRES: asserts
-- ; RUN: llc -O2 -march=hexagon -mcpu=hexagonv81 -enable-xqf-gen=true \
-- ; RUN: -hexagon-qfloat-mode=ieee -mattr=+hvxv81,+hvx-length128B \
-- ; RUN: -enable-postra-xqf-check -debug-only=handle-qfp 2>&1 < %s -o /dev/null | FileCheck %s
 diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
-index 68a14533..e42d7dbd 100644
+index e42d7dbd..5906818e 100644
 --- a/third_party/llvm/workspace.bzl
 +++ b/third_party/llvm/workspace.bzl
 @@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
  
  def repo(name):
      """Imports LLVM."""
--    LLVM_COMMIT = "b08e005995ed29a2f89aa9a3f599a9d45db65523"
--    LLVM_SHA256 = "3f4f1765816e22a21ac41dbf71d91d519b0b371b05d7e33ef2a0a862c838420c"
-+    LLVM_COMMIT = "2e0d9e12b36e853d9b108c4da09509cd54d1f3f9"
-+    LLVM_SHA256 = "8e4a851780af109d623fa6d6db9b1e62d7126272d17268fed1b149f0e737c1d7"
+-    LLVM_COMMIT = "2e0d9e12b36e853d9b108c4da09509cd54d1f3f9"
+-    LLVM_SHA256 = "8e4a851780af109d623fa6d6db9b1e62d7126272d17268fed1b149f0e737c1d7"
++    LLVM_COMMIT = "4a0a7e3ddc3bdc4fde7c5b9b9e6890fb4dbf9153"
++    LLVM_SHA256 = "1453c2abaa0fdb957a6108a1802d282bd01010dcde904288c698d44f7539c488"
  
      tf_http_archive(
          name = name,
+diff --git a/third_party/stablehlo/workspace.bzl b/third_party/stablehlo/workspace.bzl
+index 5c603385..7c94cab1 100644
+--- a/third_party/stablehlo/workspace.bzl
++++ b/third_party/stablehlo/workspace.bzl
+@@ -18,10 +18,8 @@
+ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
+ 
+ def repo():
+-    #
+     STABLEHLO_COMMIT = "b9029ac2228a12e176ea20e56f7d3dcb6bf8a019"
+     STABLEHLO_SHA256 = "16af6c8f0baacda15dd72d98dc295a213be455b8fa1899af46e144d6dae670bd"
+-    #
+ 
+     tf_http_archive(
+         name = "stablehlo",
diff --git a/third_party/xla/third_party/shardy/workspace.bzl b/third_party/xla/third_party/shardy/workspace.bzl
index ba16a772a07..152e8c5bffc 100644
--- a/third_party/xla/third_party/shardy/workspace.bzl
+++ b/third_party/xla/third_party/shardy/workspace.bzl
@@ -18,8 +18,8 @@
 load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 
 def repo():
-    SHARDY_COMMIT = "a29d5efe295b7834e3b644d49e09211f1af0d28a"
-    SHARDY_SHA256 = "b28a57c907bd6b74b849a0110b8e4081358daca62336e72d418e0770895d1b7c"
+    SHARDY_COMMIT = "38ab99f4d252c59d491899a694fe8cbadfeccec3"
+    SHARDY_SHA256 = "f8a14581ac86bf170c1aa6814daa000112f9c3a93afa677c28315433e7b2a33f"
 
     tf_http_archive(
         name = "shardy",
diff --git a/third_party/xla/third_party/triton/common/llvm_cl998243985.patch b/third_party/xla/third_party/triton/common/llvm_cl998243985.patch
new file mode 100755
index 00000000000..470de51ad17
--- /dev/null
+++ b/third_party/xla/third_party/triton/common/llvm_cl998243985.patch
@@ -0,0 +1,14 @@
+--- a/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp
++++ b/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp
+@@ -217,8 +217,8 @@
+         Value loaded = LLVM::LoadOp::create(
+             rewriter, loc, valueTy, ptr, alignment, /*isVolatile=*/false,
+             /*isNonTemporal=*/false, /*isInvariant=*/false,
+-            /*isInvariantGroup=*/false, LLVM::AtomicOrdering::monotonic,
+-            getAtomicSyncScope(scope));
++            /*isInvariantGroup=*/false, /*isNoundef=*/false,
++            LLVM::AtomicOrdering::monotonic, getAtomicSyncScope(scope));
+         return SmallVector<Value>{loaded};
+       });
+   return results.front();
+
diff --git a/third_party/xla/third_party/triton/common/series b/third_party/xla/third_party/triton/common/series
index 126e281f788..7661facc3a2 100644
--- a/third_party/xla/third_party/triton/common/series
+++ b/third_party/xla/third_party/triton/common/series
@@ -37,3 +37,4 @@ llvm_cl991819890.patch
 fix_load_ops.patch
 cxx17_support.patch
 llvm_cl996124567.patch
+llvm_cl998243985.patch
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/extern_function_helper.cc b/third_party/xla/xla/backends/gpu/codegen/triton/extern_function_helper.cc
index f809923c028..93bcf9bd2b1 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/extern_function_helper.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/extern_function_helper.cc
@@ -347,11 +347,12 @@ mlir::Value CreateAtomicSpinWaitOps(
 
   // Loop: atomic load + compare + conditional branch
   builder.setInsertionPointToStart(loop_block);
-  LLVM::LoadOp loaded = LLVM::LoadOp::create(
-      builder, params.loc, i32_type, addr, /*alignment=*/4,
-      /*isVolatile=*/false,
-      /*isNonTemporal=*/false, /*isInvariant=*/false,
-      /*isInvariantGroup=*/false, ordering, builder.getStringAttr(syncscope));
+  LLVM::LoadOp loaded =
+      LLVM::LoadOp::create(builder, params.loc, i32_type, addr, /*alignment=*/4,
+                           /*isVolatile=*/false,
+                           /*isNonTemporal=*/false, /*isInvariant=*/false,
+                           /*isInvariantGroup=*/false, /*isNoundef=*/false,
+                           ordering, builder.getStringAttr(syncscope));
   LLVM::ICmpPredicate predicate =
       ComparatorToICmpPredicate(instruction.comparator);
   LLVM::ICmpOp condition =

From b2ac52bcfe6f6c61fa65bd0c63c397ef0721a71d Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Tue, 13 Oct 2026 22:04:04 +0530
Subject: [PATCH 0691/1000] Address review: complete nested-none assertions and
 cover empty-stack guard

- test_nested_context_with_none now also asserts the fully-unwound state
  (not in_load_context, get_load_options() is None) after the outer
  context exits, matching the pattern in test_nested_context.
- Add test_clear_load_options_empty_stack_raises to cover the
  clear_load_options() RuntimeError guard on an empty stack.
---
 tensorflow/python/distribute/load_context_test.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/tensorflow/python/distribute/load_context_test.py b/tensorflow/python/distribute/load_context_test.py
index 119f9759ad3..22069b1adf1 100644
--- a/tensorflow/python/distribute/load_context_test.py
+++ b/tensorflow/python/distribute/load_context_test.py
@@ -56,6 +56,15 @@ class LoadContextTest(test.TestCase):
         self.assertIsNone(load_context.get_load_options())
       self.assertTrue(load_context.in_load_context())
       self.assertIs(load_context.get_load_options(), outer_options)
+    self.assertFalse(load_context.in_load_context())
+    self.assertIsNone(load_context.get_load_options())
+
+  def test_clear_load_options_empty_stack_raises(self):
+    with self.assertRaisesRegex(
+        RuntimeError,
+        r"clear_load_options\(\) called when load context stack is empty\.",
+    ):
+      load_context._load_context.clear_load_options()
 
   def test_thread_isolation(self):
     options = object()

From c968117514519dd6427bbb4764542767359c5f8e Mon Sep 17 00:00:00 2001
From: Oleg Shyshkov <shyshkov@google.com>
Date: Tue, 13 Oct 2026 06:33:41 -0700
Subject: [PATCH 0692/1000] [XLA:GPU] Update SolLatencyEstimator to support
 Triton collective custom fusions.

Identifies Triton collectives by checking for custom fusions with the right kind rather than relying on CollectiveBackendConfig kernel strategy flags. The old code didn't work with fusion at all.

PiperOrigin-RevId: 998594544
---
 third_party/xla/xla/service/gpu/model/BUILD   |   1 +
 .../gpu/model/sol_latency_estimator.cc        | 127 ++++++-----
 .../gpu/model/sol_latency_estimator_test.cc   | 207 ++++++++++++------
 3 files changed, 206 insertions(+), 129 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/BUILD b/third_party/xla/xla/service/gpu/model/BUILD
index 4b5ba59841e..f132263c0d6 100644
--- a/third_party/xla/xla/service/gpu/model/BUILD
+++ b/third_party/xla/xla/service/gpu/model/BUILD
@@ -71,6 +71,7 @@ cc_library(
         "//xla/service:latency_hiding_scheduler",
         "//xla/service/gpu:backend_configs_cc",
         "//xla/service/gpu:flag_utils",
+        "//xla/service/gpu:ir_emission_utils",
         "//xla/stream_executor:device_description",
         "//xla/tsl/platform:env",
         "@com_google_absl//absl/log",
diff --git a/third_party/xla/xla/service/gpu/model/sol_latency_estimator.cc b/third_party/xla/xla/service/gpu/model/sol_latency_estimator.cc
index 9805455775b..3679913b5c5 100644
--- a/third_party/xla/xla/service/gpu/model/sol_latency_estimator.cc
+++ b/third_party/xla/xla/service/gpu/model/sol_latency_estimator.cc
@@ -37,6 +37,7 @@ limitations under the License.
 #include "xla/service/collective_ops_utils.h"
 #include "xla/service/gpu/backend_configs.pb.h"
 #include "xla/service/gpu/flag_utils.h"
+#include "xla/service/gpu/ir_emission_utils.h"
 #include "xla/service/gpu/model/collective_interpolator.h"
 #include "xla/service/gpu/model/gpu_hlo_cost_analysis.h"
 #include "xla/service/gpu/model/gpu_performance_model.h"
@@ -64,6 +65,26 @@ bool IsSupportedCollectiveOp(const HloInstruction& instr) {
       &instr);
 }
 
+bool IsTritonCollectiveFusion(const HloInstruction& instr) {
+  if (!instr.IsCustomFusion()) {
+    return false;
+  }
+  absl::StatusOr<GpuBackendConfig> gpu_config =
+      instr.backend_config<GpuBackendConfig>();
+  return gpu_config.ok() && gpu_config->fusion_backend_config().kind() ==
+                                kTritonCollectiveFusionKind;
+}
+
+bool IsAsyncTritonCollectiveFusionStart(const HloInstruction& instr) {
+  return instr.opcode() == HloOpcode::kAsyncStart &&
+         IsTritonCollectiveFusion(*instr.async_wrapped_instruction());
+}
+
+bool IsAsyncTritonCollectiveFusionDone(const HloInstruction& instr) {
+  return instr.opcode() == HloOpcode::kAsyncDone &&
+         IsTritonCollectiveFusion(*instr.async_wrapped_instruction());
+}
+
 bool IsHostOffloaded(const HloInstruction& instr) {
   auto backend_config = instr.backend_config<GpuBackendConfig>();
   return backend_config.ok() &&
@@ -251,27 +272,22 @@ absl::StatusOr<absl::Duration> DispatchEstimation(
       // If an intra-node collective kernel will be used (Triton codegen kernel
       // — uses the NVLink-based cost formula), apply the NVLink model instead
       // of the NCCL-calibrated CollectiveInterpolator.
-      absl::StatusOr<GpuBackendConfig> backend_cfg =
-          instr.backend_config<GpuBackendConfig>();
-      if (backend_cfg.ok()) {
-        const auto ks =
-            backend_cfg->collective_backend_config().kernel_strategy();
-        if (IsTritonCollectiveKernel(ks)) {
-          SolGPUCostModel sol_model(sol_flags);
-          const int num_gpus = num_groups_and_devices->second;
-          const int active_links =
-              gpu_device_info.device_interconnect_info().active_links;
-          const int64_t size_bytes =
-              static_cast<int64_t>(analysis.BytesTransferred(instr));
-          // Add HBM time (reading/writing local buffers) on top of NVLink time.
-          absl::Duration hbm_time =
-              absl::Seconds(1.0f * analysis.bytes_accessed(instr) /
-                            gpu_device_info.memory_bandwidth());
-          ABSL_ASSIGN_OR_RETURN(absl::Duration nvlink_time,
-                           sol_model.IntraNodeAllReduceLatency(
-                               size_bytes, num_gpus, active_links));
-          return hbm_time + nvlink_time;
-        }
+      if (instr.IsFused() &&
+          IsTritonCollectiveFusion(*instr.parent()->FusionInstruction())) {
+        SolGPUCostModel sol_model(sol_flags);
+        const int num_gpus = num_groups_and_devices->second;
+        const int active_links =
+            gpu_device_info.device_interconnect_info().active_links;
+        const int64_t size_bytes =
+            static_cast<int64_t>(analysis.BytesTransferred(instr));
+        // Add HBM time (reading/writing local buffers) on top of NVLink time.
+        absl::Duration hbm_time =
+            absl::Seconds(1.0f * analysis.bytes_accessed(instr) /
+                          gpu_device_info.memory_bandwidth());
+        ABSL_ASSIGN_OR_RETURN(absl::Duration nvlink_time,
+                         sol_model.IntraNodeAllReduceLatency(
+                             size_bytes, num_gpus, active_links));
+        return hbm_time + nvlink_time;
       }
       // NCCL intra-node collective: use empirical CollectiveInterpolator.
       if (collective_interpolator == nullptr) {
@@ -358,6 +374,9 @@ SolLatencyEstimator::ComputeCollectiveTime(
 
   const HloInstruction* collective =
       instr.IsAsynchronous() ? instr.async_wrapped_instruction() : &instr;
+  if (IsTritonCollectiveFusion(*collective)) {
+    collective = collective->fused_expression_root();
+  }
   if (const auto* cp = DynCast<HloCollectivePermuteInstruction>(collective)) {
     // Handles the collective-permute ops.
     int64_t partition_size = GetPartitionSize(*cp, sol_flags);
@@ -472,23 +491,17 @@ LatencyEstimator::TimeCost SolLatencyEstimator::GetLatencyBetween(
     return kLowLatency;
   }
 
-  // For sync intra-node AllReduce (Triton codegen kernel annotated as
-  // TRITON_ONE_SHOT / TRITON_TWO_SHOT) the kernel runs on the compute stream
-  // and fully blocks further compute — no overlap is possible between start
-  // and done.  NodeCost() already reports the full blocking latency on the
-  // critical path, so return kLowLatency here to avoid double-counting.
-  {
-    absl::StatusOr<GpuBackendConfig> cfg =
-        from.GetInstr().backend_config<GpuBackendConfig>();
-    if (cfg.ok() && cfg->collective_backend_config().is_sync()) {
-      const auto ks = cfg->collective_backend_config().kernel_strategy();
-      if (IsTritonCollectiveKernel(ks)) {
-        VLOG(10) << "GetLatencyBetween: Returning kLowLatency for sync "
-                    "intra-node AllReduce (cost already in NodeCost): "
-                 << from.GetInstr().name();
-        return kLowLatency;
-      }
-    }
+  // For sync intra-node AllReduce (Triton collective fusion) the kernel runs on
+  // the compute stream and fully blocks further compute — no overlap is
+  // possible between start and done. NodeCost() already reports the full
+  // blocking latency on the critical path, so return kLowLatency here to avoid
+  // double-counting.
+  if (IsAsyncTritonCollectiveFusionStart(from.GetInstr()) &&
+      IsGPUSyncCollective(from.GetInstr())) {
+    VLOG(10) << "GetLatencyBetween: Returning kLowLatency for sync "
+                "intra-node AllReduce (cost already in NodeCost): "
+             << from.GetInstr().name();
+    return kLowLatency;
   }
 
   if (!IsAsyncPair(from, target) && !IsSupportedCollectiveOp(from.GetInstr())) {
@@ -523,31 +536,27 @@ LatencyEstimator::TimeCost SolLatencyEstimator::NodeCost(
   if (const std::optional<TimeCost> latency = GetLatencyFromMetadata(*instr)) {
     return *latency;
   }
-  // Sync intra-node AllReduce kernels (Triton codegen, annotated as
-  // TRITON_ONE_SHOT / TRITON_TWO_SHOT) run on the compute stream and block
-  // compute entirely.  Their full latency is on the critical path, so we must
-  // report it here as NodeCost rather than as a hidden GetLatencyBetween.
-  if (hlo_query::IsAsyncCollectiveStartOp(instr, /*include_send_recv=*/true)) {
-    absl::StatusOr<GpuBackendConfig> cfg =
-        instr->backend_config<GpuBackendConfig>();
-    if (cfg.ok() && cfg->collective_backend_config().is_sync()) {
-      const auto ks = cfg->collective_backend_config().kernel_strategy();
-      if (IsTritonCollectiveKernel(ks)) {
-        absl::StatusOr<absl::Duration> t = ComputeCollectiveTime(
-            *instr, gpu_info_, shape_size_function_, sol_flags_,
-            *cost_analysis_, collective_interpolator_.get());
-        if (t.ok()) {
-          VLOG(10) << "NodeCost: Sync intra-node AllReduce cost for "
-                   << instr->name() << ": " << absl::ToDoubleMicroseconds(*t)
-                   << " us";
-          return absl::ToDoubleMicroseconds(*t);
-        }
-      }
+  // Sync intra-node AllReduce kernels (Triton collective fusions) run on the
+  // compute stream and block compute entirely. Their full latency is on the
+  // critical path, so we must report it here as NodeCost rather than as a
+  // hidden GetLatencyBetween.
+  if (IsAsyncTritonCollectiveFusionStart(*instr) &&
+      IsGPUSyncCollective(*instr)) {
+    absl::StatusOr<absl::Duration> t = ComputeCollectiveTime(
+        *instr, gpu_info_, shape_size_function_, sol_flags_, *cost_analysis_,
+        collective_interpolator_.get());
+    if (t.ok()) {
+      VLOG(10) << "NodeCost: Sync intra-node AllReduce cost for "
+               << instr->name() << ": " << absl::ToDoubleMicroseconds(*t)
+               << " us";
+      return absl::ToDoubleMicroseconds(*t);
     }
   }
 
   if (hlo_query::IsAsyncCollectiveStartOp(instr, /*include_send_recv=*/true) ||
-      hlo_query::IsAsyncCollectiveDoneOp(instr, /*include_send_recv=*/true)) {
+      hlo_query::IsAsyncCollectiveDoneOp(instr, /*include_send_recv=*/true) ||
+      IsAsyncTritonCollectiveFusionStart(*instr) ||
+      IsAsyncTritonCollectiveFusionDone(*instr)) {
     VLOG(10) << "NodeCost: Returning kLowCost for async start/done op "
              << instr->name();
     return kLowCost;
diff --git a/third_party/xla/xla/service/gpu/model/sol_latency_estimator_test.cc b/third_party/xla/xla/service/gpu/model/sol_latency_estimator_test.cc
index f52179e018a..afe1bc297d8 100644
--- a/third_party/xla/xla/service/gpu/model/sol_latency_estimator_test.cc
+++ b/third_party/xla/xla/service/gpu/model/sol_latency_estimator_test.cc
@@ -1058,14 +1058,18 @@ TEST_F(IsSolLatencyEstimatorEnabledTest, CreatesEstimatorWithGfx942Profiles) {
 // ---- Triton collective scheduler integration tests -----------------------
 //
 // These tests verify that SolLatencyEstimator correctly uses the NVLink-based
-// cost formula when an AllReduceStart is annotated with a Triton kernel
-// strategy, and that the sync case is correctly handled by NodeCost.
+// cost formula when an async-start wraps a __triton_collective custom fusion,
+// and that both async and sync cases are handled properly by NodeCost and
+// GetLatencyBetween.
 
 class TritonAllReduceSchedulerTest : public HloHardwareIndependentTestBase {
  protected:
   // 32768 F32 elements = 128 KB → one-shot strategy (< 256 KB threshold).
   // All 8 replicas in a single group → SINGLE_PARTITION communication type.
-  static constexpr absl::string_view kOneShotAllReduceHlo = R"(
+  // Mirrors the HLO produced by CollectivesScheduleLinearizer /
+  // FuseSupportedCollectives: async-start calls an async computation whose root
+  // is a kCustom fusion with kind="__triton_collective".
+  static constexpr absl::string_view kOneShotTritonAllReduceHlo = R"(
 HloModule m, num_partitions=8
 
 add {
@@ -1074,12 +1078,49 @@ add {
   ROOT r = f32[] add(p0, p1)
 }
 
-ENTRY e {
-  p = f32[32768] parameter(0)
-  ar-start = f32[32768] all-reduce-start(p),
+triton_collective_computation {
+  param_0 = f32[32768] parameter(0)
+  ROOT all-reduce = f32[32768] all-reduce(param_0),
       replica_groups={{0,1,2,3,4,5,6,7}},
       to_apply=add
-  ROOT ar-done = f32[32768] all-reduce-done(ar-start)
+}
+
+async_computation {
+  param_0 = f32[32768] parameter(0)
+  ROOT custom-fusion = f32[32768] fusion(param_0),
+      kind=kCustom, calls=triton_collective_computation,
+      backend_config={"fusion_backend_config":{"kind":"__triton_collective"}}
+}
+
+ENTRY e {
+  p = f32[32768] parameter(0)
+  ar-start = ((f32[32768]), f32[32768]) async-start(p),
+      calls=async_computation
+  ROOT ar-done = f32[32768] async-done(ar-start)
+})";
+
+  // Same 128 KB all-reduce without Triton fusion (NCCL path).
+  static constexpr absl::string_view kOneShotNcclAllReduceHlo = R"(
+HloModule m, num_partitions=8
+
+add {
+  p0 = f32[] parameter(0)
+  p1 = f32[] parameter(1)
+  ROOT r = f32[] add(p0, p1)
+}
+
+async_computation {
+  param_0 = f32[32768] parameter(0)
+  ROOT all-reduce = f32[32768] all-reduce(param_0),
+      replica_groups={{0,1,2,3,4,5,6,7}},
+      to_apply=add
+}
+
+ENTRY e {
+  p = f32[32768] parameter(0)
+  ar-start = ((f32[32768]), f32[32768]) async-start(p),
+      calls=async_computation
+  ROOT ar-done = f32[32768] async-done(ar-start)
 })";
 
   TritonAllReduceSchedulerTest() {
@@ -1105,15 +1146,6 @@ ENTRY e {
     };
   }
 
-  // Annotates `instr` with the given Triton kernel strategy.
-  static void SetKernelStrategy(
-      HloInstruction* instr,
-      CollectiveBackendConfig::CollectiveKernelStrategy strategy) {
-    GpuBackendConfig cfg = instr->backend_config<GpuBackendConfig>().value();
-    cfg.mutable_collective_backend_config()->set_kernel_strategy(strategy);
-    CHECK_OK(instr->set_backend_config(cfg));
-  }
-
   // Marks `instr` as synchronous (runs on compute stream, no async overlap).
   static void SetSync(HloInstruction* instr) {
     GpuBackendConfig cfg = instr->backend_config<GpuBackendConfig>().value();
@@ -1139,30 +1171,40 @@ ENTRY e {
         static_cast<int64_t>(estimator->NodeCost(&instr)));
   }
 
+  // Creates a SolLatencyEstimator and calls GetLatencyBetween.
+  absl::Duration GetLatencyBetween(const HloInstruction& from,
+                                   const HloInstruction& target,
+                                   const HloComputation* computation) {
+    auto estimator = *SolLatencyEstimator::Create(
+        scheduler_config_, std::make_unique<DummyLatencyEstimator>(),
+        gpu_device_info_, HloCostAnalysis::DefaultShapeSize, computation);
+    return absl::Microseconds(static_cast<int64_t>(estimator->GetLatencyBetween(
+        HloGraphNode(&from, /*original_position=*/-1),
+        HloGraphNode(&target, /*original_position=*/-1))));
+  }
+
   se::DeviceDescription gpu_device_info_;
   SolGPUCostModel::Config sol_flags_;
   SchedulerConfig scheduler_config_;
 };
 
 // Test 1: async Triton ONE_SHOT AllReduce.
-// With the KERNEL_STRATEGY_TRITON_ONE_SHOT annotation and 18 NVLink
-// links at 20 GB/s, ComputeCollectiveTime must return the NVLink formula
-// result:
+// When async-start wraps a __triton_collective custom fusion and the device has
+// 18 NVLink links at 20 GB/s, ComputeCollectiveTime and GetLatencyBetween must
+// return the NVLink formula result:
 //   transfer = 7 × 128 KB / (18 × 20 GB/s) ≈ 2.5 µs
 //   total ≈ launch(1µs) + 2.5µs + barrier(1.5µs) ≈ 5 µs
-// This must be different from the NCCL ring estimate (which uses nic_speed_gbps
-// and has a much larger RTT term).
+// This must be different from the NCCL estimate on an unfused all-reduce.
 TEST_F(TritonAllReduceSchedulerTest, AsyncTritonOneShotUsesNvlinkFormula) {
-  ASSERT_OK_AND_ASSIGN(auto module,
-                       ParseAndReturnVerifiedModule(kOneShotAllReduceHlo));
+  ASSERT_OK_AND_ASSIGN(auto triton_module, ParseAndReturnVerifiedModule(
+                                               kOneShotTritonAllReduceHlo));
 
-  HloInstruction* ar_start = hlo_query::FindInstruction(
-      module->entry_computation(), HloOpcode::kAllReduceStart);
-  ASSERT_NE(ar_start, nullptr);
-
-  // Annotate with Triton one-shot strategy (async: is_sync is false by default)
-  SetKernelStrategy(ar_start,
-                    CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT);
+  HloInstruction* triton_start = hlo_query::FindInstruction(
+      triton_module->entry_computation(), HloOpcode::kAsyncStart);
+  HloInstruction* triton_done = hlo_query::FindInstruction(
+      triton_module->entry_computation(), HloOpcode::kAsyncDone);
+  ASSERT_NE(triton_start, nullptr);
+  ASSERT_NE(triton_done, nullptr);
 
   // A real interpolator is passed so the NCCL fallback path doesn't crash.
   // If the Triton dispatch fires correctly, the interpolator is not called.
@@ -1172,99 +1214,124 @@ TEST_F(TritonAllReduceSchedulerTest, AsyncTritonOneShotUsesNvlinkFormula) {
   ASSERT_OK_AND_ASSIGN(
       absl::Duration triton_time,
       SolLatencyEstimator::ComputeCollectiveTime(
-          *ar_start, gpu_device_info_, HloCostAnalysis::DefaultShapeSize,
+          *triton_start, gpu_device_info_, HloCostAnalysis::DefaultShapeSize,
           sol_flags_, interpolator.get()));
   triton_time = absl::Trunc(triton_time, absl::Microseconds(1));
 
-  // The NVLink formula gives ~5 µs.
+  // The NVLink formula gives ~5 µs, and async-done has 0 collective time.
   EXPECT_EQ(triton_time, absl::Microseconds(5));
+  ASSERT_OK_AND_ASSIGN(absl::Duration done_time,
+                       ComputeCollectiveTime(*triton_done));
+  EXPECT_EQ(done_time, absl::ZeroDuration());
 
-  // Verify this is different from the NCCL ring estimate: remove annotation.
-  SetKernelStrategy(ar_start, CollectiveBackendConfig::KERNEL_STRATEGY_DEFAULT);
+  // For async Triton AllReduce, GetLatencyBetween(start, done) also returns the
+  // NVLink collective time (~5 µs).
+  EXPECT_EQ(GetLatencyBetween(*triton_start, *triton_done,
+                              triton_module->entry_computation()),
+            absl::Microseconds(5));
+
+  // Verify this is different from the NCCL estimate on an unfused all-reduce.
+  ASSERT_OK_AND_ASSIGN(auto nccl_module,
+                       ParseAndReturnVerifiedModule(kOneShotNcclAllReduceHlo));
+  HloInstruction* nccl_start = hlo_query::FindInstruction(
+      nccl_module->entry_computation(), HloOpcode::kAsyncStart);
+  ASSERT_NE(nccl_start, nullptr);
   ASSERT_OK_AND_ASSIGN(
       absl::Duration nccl_time,
       SolLatencyEstimator::ComputeCollectiveTime(
-          *ar_start, gpu_device_info_, HloCostAnalysis::DefaultShapeSize,
+          *nccl_start, gpu_device_info_, HloCostAnalysis::DefaultShapeSize,
           sol_flags_, interpolator.get()));
-  // NCCL ring with nic=100 GB/s and rtt=100 µs gives a much larger value.
   EXPECT_NE(absl::Trunc(nccl_time, absl::Microseconds(1)), triton_time);
 }
 
-// Test 2: sync Triton ONE_SHOT AllReduce - NodeCost must NOT return kLowCost.
-// When is_sync=true and kernel_strategy=TRITON_ONE_SHOT, NodeCost() must
-// delegate to ComputeCollectiveTime so the latency is on the critical path.
-// We verify this by calling ComputeCollectiveTime directly (which uses our
-// NVLink flags) and checking it returns > kLowCost (1.0 µs).
-TEST_F(TritonAllReduceSchedulerTest,
-       SyncTritonOneShotComputeCollectiveTimeExceedsLowCost) {
-  ASSERT_OK_AND_ASSIGN(auto module,
-                       ParseAndReturnVerifiedModule(kOneShotAllReduceHlo));
+// Test 2: async Triton AllReduce start and done have kLowCost = 1.0 µs NodeCost
+// because their latency is modelled on the edge via GetLatencyBetween.
+TEST_F(TritonAllReduceSchedulerTest, AsyncTritonAllReduceNodeCostIsLow) {
+  ASSERT_OK_AND_ASSIGN(
+      auto module, ParseAndReturnVerifiedModule(kOneShotTritonAllReduceHlo));
 
   HloInstruction* ar_start = hlo_query::FindInstruction(
-      module->entry_computation(), HloOpcode::kAllReduceStart);
+      module->entry_computation(), HloOpcode::kAsyncStart);
+  HloInstruction* ar_done = hlo_query::FindInstruction(
+      module->entry_computation(), HloOpcode::kAsyncDone);
+  ASSERT_NE(ar_start, nullptr);
+  ASSERT_NE(ar_done, nullptr);
+
+  EXPECT_EQ(NodeCost(*ar_start, module->entry_computation()),
+            absl::Microseconds(1));
+  EXPECT_EQ(NodeCost(*ar_done, module->entry_computation()),
+            absl::Microseconds(1));
+}
+
+// Test 3: sync Triton ONE_SHOT AllReduce - ComputeCollectiveTime returns ~5 µs.
+TEST_F(TritonAllReduceSchedulerTest,
+       SyncTritonOneShotComputeCollectiveTimeExceedsLowCost) {
+  ASSERT_OK_AND_ASSIGN(
+      auto module, ParseAndReturnVerifiedModule(kOneShotTritonAllReduceHlo));
+
+  HloInstruction* ar_start = hlo_query::FindInstruction(
+      module->entry_computation(), HloOpcode::kAsyncStart);
   ASSERT_NE(ar_start, nullptr);
 
-  // Mark as synchronous Triton one-shot.
   SetSync(ar_start);
-  SetKernelStrategy(ar_start,
-                    CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT);
 
   ASSERT_OK_AND_ASSIGN(absl::Duration cost, ComputeCollectiveTime(*ar_start));
 
   // The NVLink formula must return significantly more than kLowCost = 1.0 µs.
-  // We expect ~5 µs.  Any value > 2 µs confirms we are not returning kLowCost.
+  // We expect ~5 µs. Any value > 2 µs confirms we are not returning kLowCost.
   EXPECT_GT(cost, absl::Microseconds(2));
   EXPECT_EQ(absl::Trunc(cost, absl::Microseconds(1)), absl::Microseconds(5));
 }
 
-// Test 3: sync Triton ONE_SHOT AllReduce exercises the actual NodeCost
-// branching logic end-to-end via SolLatencyEstimator::Create.
+// Test 4: sync Triton ONE_SHOT AllReduce exercises NodeCost and
+// GetLatencyBetween end-to-end via SolLatencyEstimator::Create.
 // SolLatencyEstimator::Create derives sol_flags from GetConfig(module,
 // device_info), which calls GetIciBandwidthPerLaneGbps on the device.
 // For RTXA6000 with SM9.0 (Hopper) and 18 active NVLink links:
 //   nvlink_bw_per_lane_gbps = kSm90NvlinkBandwidth = 20.0 GB/s
 //   nvlink_barrier_latency  = 1.5 µs  (kUnknownKey table entry)
-// → NodeCost returns ~5 µs via ComputeCollectiveTime, confirming that the
-//   IsGPUSyncCollective + TRITON_ONE_SHOT branch in NodeCost fires correctly
-//   with production-derived (not test-injected) sol_flags.
+// → NodeCost(ar-start) returns ~5 µs on the critical path, while
+//   GetLatencyBetween(ar-start, ar-done) returns kLowLatency = 1 µs to avoid
+//   double-counting.
 TEST_F(TritonAllReduceSchedulerTest,
        SyncTritonOneShotNodeCostExceedsLowCostEndToEnd) {
-  ASSERT_OK_AND_ASSIGN(auto module,
-                       ParseAndReturnVerifiedModule(kOneShotAllReduceHlo));
+  ASSERT_OK_AND_ASSIGN(
+      auto module, ParseAndReturnVerifiedModule(kOneShotTritonAllReduceHlo));
 
   HloInstruction* ar_start = hlo_query::FindInstruction(
-      module->entry_computation(), HloOpcode::kAllReduceStart);
+      module->entry_computation(), HloOpcode::kAsyncStart);
+  HloInstruction* ar_done = hlo_query::FindInstruction(
+      module->entry_computation(), HloOpcode::kAsyncDone);
   ASSERT_NE(ar_start, nullptr);
+  ASSERT_NE(ar_done, nullptr);
 
   SetSync(ar_start);
-  SetKernelStrategy(ar_start,
-                    CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT);
 
-  // NodeCost() internally calls SolLatencyEstimator::Create which reads
-  // sol_flags from GetConfig(module, device_info).  For SM9.0 + 18 links
-  // the derived config gives nvlink_bw=20.0 GB/s and barrier=1.5 µs, so
-  // the Triton one-shot formula produces ~5 µs — well above kLowCost=1 µs.
   absl::Duration node_cost = NodeCost(*ar_start, module->entry_computation());
 
   // The Triton one-shot formula: launch(1µs) + 7×128KB/(18×20GB/s) + 1.5µs
-  // ≈ 5 µs.  Verify the NodeCost branch fires (> kLowCost) and matches.
+  // ≈ 5 µs. Verify the NodeCost branch fires (> kLowCost) and matches.
   EXPECT_GT(node_cost, absl::Microseconds(1))
       << "NodeCost for sync Triton AllReduce should exceed kLowCost=1µs; "
-         "verify that the TRITON_ONE_SHOT branch in NodeCost fired";
+         "verify that the sync Triton collective fusion branch in NodeCost "
+         "fired";
   EXPECT_EQ(node_cost, absl::Microseconds(5));
+  EXPECT_EQ(NodeCost(*ar_done, module->entry_computation()),
+            absl::Microseconds(1));
+  EXPECT_EQ(GetLatencyBetween(*ar_start, *ar_done, module->entry_computation()),
+            absl::Microseconds(1));
 }
 
-// Test 4: without the annotation (NCCL path), async AllReduceStart NodeCost
+// Test 5: without Triton fusion (NCCL path), async AllReduceStart NodeCost
 // returns SolLatencyEstimator::kLowCost = 1.0 (latency is hidden, not on
 // critical path).
 TEST_F(TritonAllReduceSchedulerTest, AsyncNcclAllReduceNodeCostIsLow) {
   ASSERT_OK_AND_ASSIGN(auto module,
-                       ParseAndReturnVerifiedModule(kOneShotAllReduceHlo));
+                       ParseAndReturnVerifiedModule(kOneShotNcclAllReduceHlo));
 
   HloInstruction* ar_start = hlo_query::FindInstruction(
-      module->entry_computation(), HloOpcode::kAllReduceStart);
+      module->entry_computation(), HloOpcode::kAsyncStart);
   ASSERT_NE(ar_start, nullptr);
-  // No annotation: is_sync=false (default), kernel_strategy=DEFAULT (default).
 
   absl::Duration node_cost = NodeCost(*ar_start, module->entry_computation());
 

From 19ff5604c8bbf59e708c7b790b3cfa909205491e Mon Sep 17 00:00:00 2001
From: abhirupone <abhirup.karmakar00@gmail.com>
Date: Tue, 13 Oct 2026 22:07:25 +0530
Subject: [PATCH 0693/1000] Add xla_cpu_jit dependency to ragged_tensor_test
 target in BUILD

- Add //tensorflow/compiler/jit:xla_cpu_jit to deps for ragged_tensor_test so that the XLA Host CPU compiler is linked and registered for jit_compile=True test execution
---
 tensorflow/python/ops/ragged/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/python/ops/ragged/BUILD b/tensorflow/python/ops/ragged/BUILD
index 4c8416eb018..15eb004800f 100644
--- a/tensorflow/python/ops/ragged/BUILD
+++ b/tensorflow/python/ops/ragged/BUILD
@@ -798,6 +798,7 @@ py_test(
         "//tensorflow/python/framework:tensor",
         "//tensorflow/python/framework:tensor_shape",
         "//tensorflow/python/framework:test_lib",
+        "//tensorflow/compiler/jit:xla_cpu_jit",
         "//tensorflow/python/framework:type_spec",
         "//tensorflow/python/framework:type_utils",
         "//tensorflow/python/ops:array_grad",

From 42c89751883c4b281ac4534458ad97b08cc1d1c8 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Tue, 13 Oct 2026 06:52:18 -0700
Subject: [PATCH 0694/1000] [XLA:GPU] Eliminate double calls to Thunk::Prepare
 and Thunk::Initialize

Currently during command buffer conversion we call Prepare/Initialize twice.
1. With ThunkExecutor.
2. With CommandExecutor.

This change avoids that. Additionally some cleanups in collective thunk initialization.

PiperOrigin-RevId: 998602508
---
 .../backends/gpu/runtime/all_to_all_thunk.cc  |  1 -
 .../gpu/runtime/collective_reduce_thunk.cc    |  3 +-
 .../gpu/runtime/collective_reduce_thunk.h     |  5 +-
 .../backends/gpu/runtime/collective_thunk.cc  |  8 --
 .../backends/gpu/runtime/collective_thunk.h   | 10 ---
 .../gpu/runtime/command_buffer_thunk.cc       | 39 ++++------
 .../backends/gpu/runtime/conditional_thunk.cc | 22 +++---
 .../backends/gpu/runtime/custom_call_thunk.cc | 45 +++++++-----
 .../runtime/dynamic_slice_fusion_v2_thunk.cc  | 24 +++---
 .../gpu/runtime/ragged_all_to_all_thunk.cc    |  2 -
 .../xla/backends/gpu/runtime/recv_thunk.cc    |  5 --
 .../xla/xla/backends/gpu/runtime/recv_thunk.h |  1 -
 .../xla/backends/gpu/runtime/send_thunk.cc    |  5 --
 .../xla/xla/backends/gpu/runtime/send_thunk.h |  2 -
 .../xla/backends/gpu/runtime/while_thunk.cc   | 20 +++--
 .../backends/gpu/runtime/while_thunk_test.cc  | 69 ++++++++++++++++++
 .../collective_ops_command_buffer_test.cc     | 73 +++++++++++++++++++
 17 files changed, 220 insertions(+), 114 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk.cc
index d25bc5653a3..fb9457e6f38 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/all_to_all_thunk.cc
@@ -125,7 +125,6 @@ AllToAllThunk::AllToAllThunk(ThunkInfo thunk_info,
 }
 
 absl::Status AllToAllThunk::Initialize(const InitializeParams& params) {
-  ABSL_RETURN_IF_ERROR(CollectiveThunk::Initialize(params));
   CHECK_GT(params.local_device_count, 0);
   XLA_VLOG_DEVICE(5, params.executor->device_ordinal())
       << "Local device count : " << params.local_device_count;
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.cc
index 8f40bfe7ce3..46965663c48 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.cc
@@ -102,8 +102,7 @@ CollectiveReduceThunk::CollectiveReduceThunk(
   return GetAllReduceConfigInst(inst).config.group_mode;
 }
 
-absl::Status CollectiveReduceThunk::InitializeCollective(
-    const InitializeParams& params, const GpuCliqueKey& clique_key) {
+absl::Status CollectiveReduceThunk::Initialize(const InitializeParams& params) {
   if (!has_dynamic_root_) {
     return absl::OkStatus();
   }
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.h b/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.h
index 0d59fa7a066..efcbc36d667 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.h
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_reduce_thunk.h
@@ -78,12 +78,11 @@ class CollectiveReduceThunk : public AllReduceReduceScatterThunkBase {
 
   absl::StatusOr<ThunkProto> ToProto() const override;
 
+  absl::Status Initialize(const InitializeParams& params) override;
+
  protected:
   bool RequiresRendezvous() const override { return true; }
 
-  absl::Status InitializeCollective(const InitializeParams& params,
-                                    const GpuCliqueKey& clique_key) override;
-
   absl::Status RunCollective(const ExecuteParams& params,
                              const GpuCliqueKey& clique_key, se::Stream& stream,
                              Communicator& comm) override;
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/collective_thunk.cc
index e0762966398..5fc3ccf9850 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_thunk.cc
@@ -420,14 +420,6 @@ absl::Status CollectiveThunk::Prepare(const PrepareParams& params) {
   return PrepareCollective(params, clique_key);
 }
 
-absl::Status CollectiveThunk::Initialize(const InitializeParams& params) {
-  ABSL_ASSIGN_OR_RETURN(
-      GpuCliqueKey clique_key,
-      GetGpuCliqueKey(*params.collective_params, config().replica_groups,
-                      config().group_mode, communication_id_));
-  return InitializeCollective(params, clique_key);
-}
-
 absl::Status CollectiveThunk::FirstCallRendezvous(
     const ExecuteParams& params, const GpuCliqueKey& clique_key,
     absl::string_view label, RendezvousFlag& flag) {
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_thunk.h b/third_party/xla/xla/backends/gpu/runtime/collective_thunk.h
index 5b1352883db..87b6e464ac4 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_thunk.h
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_thunk.h
@@ -130,7 +130,6 @@ class CollectiveThunk : public Command {
   bool requires_warmup() const override { return true; }
 
   absl::Status Prepare(const PrepareParams& params) override;
-  absl::Status Initialize(const InitializeParams& params) override;
   absl::Status ExecuteOnStream(const ExecuteParams& params) override;
   absl::StatusOr<const se::CommandBuffer::Command*> Record(
       const ExecuteParams& execute_params, const RecordParams& record_params,
@@ -171,15 +170,6 @@ class CollectiveThunk : public Command {
     return absl::OkStatus();
   }
 
-  // Initializes collective operation for execution.
-  //
-  // At this stage it is possible to resolve buffer slices from a buffer
-  // assignment, but the content of all buffers is undefined.
-  virtual absl::Status InitializeCollective(const InitializeParams& params,
-                                            const GpuCliqueKey& clique_key) {
-    return absl::OkStatus();
-  }
-
   // Run collective operation on a given stream.
   //
   // A collective thunk is normally an independent operation in a sense that
diff --git a/third_party/xla/xla/backends/gpu/runtime/command_buffer_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/command_buffer_thunk.cc
index a85f056a213..a0d2dc88feb 100644
--- a/third_party/xla/xla/backends/gpu/runtime/command_buffer_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/command_buffer_thunk.cc
@@ -142,25 +142,13 @@ absl::Status CommandBufferThunk::Prepare(const PrepareParams& params) {
     return absl::OkStatus();
   }
 
-  // Always prepare thunks if they are present so we are ready to fall back
-  // on them if we detect profiling activity.
+  // `commands_` borrows pointers to the same Thunk objects owned by `thunks_`,
+  // so preparing `thunks_` (when present) already prepares all commands.
   if (thunks_) {
-    ABSL_RETURN_IF_ERROR(thunks_->Prepare(params));
+    return thunks_->Prepare(params);
   }
 
-  // TODO(b/290773547): Disabled CUDA graphs when profiling is active because of
-  // memory corruption.
-  if (tsl::profiler::ProfilerLock::HasActiveSession() && thunks_ &&
-      !enable_command_buffers_during_profiling_) {
-    VLOG(1) << "Prepare command buffer thunk as a regular thunk sequence "
-               "because we detected active profiling session";
-    TraceMe trace("WARNING: CommandBuffer disabled when profiling");
-    return absl::OkStatus();
-  }
-
-  ABSL_RETURN_IF_ERROR(commands_.Prepare(params));
-
-  return absl::OkStatus();
+  return commands_.Prepare(params);
 }
 
 absl::Status CommandBufferThunk::Initialize(const InitializeParams& params) {
@@ -170,13 +158,14 @@ absl::Status CommandBufferThunk::Initialize(const InitializeParams& params) {
     return absl::OkStatus();
   }
 
-  // Initialize commands.
-  ABSL_RETURN_IF_ERROR(commands_.Initialize(params));
-
-  // Always initialize thunks if they are present so we are ready to fall back
-  // on them if we detect profiling activity.
+  // Always initialize fallback thunks (or commands when there are no fallback
+  // thunks). `commands_` borrows pointers to the same Thunk objects owned by
+  // `thunks_`, so initializing `thunks_` (when present) already initializes all
+  // commands.
   if (thunks_) {
     ABSL_RETURN_IF_ERROR(thunks_->Initialize(params));
+  } else {
+    ABSL_RETURN_IF_ERROR(commands_.Initialize(params));
   }
 
   // TODO(b/290773547): Disabled CUDA graphs when profiling is active because of
@@ -209,6 +198,10 @@ absl::Status CommandBufferThunk::Initialize(const InitializeParams& params) {
     cmd_buffer->warmup_done = true;
   }
 
+  if (!cmd_buffer->warmup_done) {
+    return absl::OkStatus();
+  }
+
   // Construct ExecuteParams with empty fields for everything that is not needed
   // for recording commands.
   Thunk::ExecuteParams execute_params(
@@ -224,10 +217,6 @@ absl::Status CommandBufferThunk::Initialize(const InitializeParams& params) {
       /*mock_collectives=*/false, /*execution_id=*/0,
       /*rng_seed=*/0, params.persistent_alloc_indices);
 
-  if (!cmd_buffer->warmup_done) {
-    return absl::OkStatus();
-  }
-
   // If command buffer is in `kCreate` state it means that command buffer
   // sequence was never recorded into it. We initialize all command buffers
   // before execution, because command buffers when instantiated will allocate
diff --git a/third_party/xla/xla/backends/gpu/runtime/conditional_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/conditional_thunk.cc
index cad0060c570..b160edee413 100644
--- a/third_party/xla/xla/backends/gpu/runtime/conditional_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/conditional_thunk.cc
@@ -146,12 +146,11 @@ absl::Status ConditionalThunk::Prepare(const PrepareParams& params) {
   } else {
     TF_RET_CHECK(!branch_executors_.empty());
   }
-  for (auto& branch_executor : branch_executors_) {
-    ABSL_RETURN_IF_ERROR(branch_executor.Prepare(params));
-  }
-  if (command_branch_executors_.has_value()) {
-    for (CommandExecutor& branch_executor : *command_branch_executors_) {
-      ABSL_RETURN_IF_ERROR(branch_executor.Prepare(params));
+  for (size_t i = 0; i < branch_executors_.size(); ++i) {
+    if (!branch_executors_[i].thunks().empty()) {
+      ABSL_RETURN_IF_ERROR(branch_executors_[i].Prepare(params));
+    } else if (command_branch_executors_.has_value()) {
+      ABSL_RETURN_IF_ERROR((*command_branch_executors_)[i].Prepare(params));
     }
   }
   return absl::OkStatus();
@@ -163,12 +162,11 @@ absl::Status ConditionalThunk::Initialize(const InitializeParams& params) {
   } else {
     TF_RET_CHECK(!branch_executors_.empty());
   }
-  for (auto& branch_executor : branch_executors_) {
-    ABSL_RETURN_IF_ERROR(branch_executor.Initialize(params));
-  }
-  if (command_branch_executors_.has_value()) {
-    for (CommandExecutor& branch_executor : *command_branch_executors_) {
-      ABSL_RETURN_IF_ERROR(branch_executor.Initialize(params));
+  for (size_t i = 0; i < branch_executors_.size(); ++i) {
+    if (!branch_executors_[i].thunks().empty()) {
+      ABSL_RETURN_IF_ERROR(branch_executors_[i].Initialize(params));
+    } else if (command_branch_executors_.has_value()) {
+      ABSL_RETURN_IF_ERROR((*command_branch_executors_)[i].Initialize(params));
     }
   }
 
diff --git a/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc
index 7355893546f..8a7a3762b5c 100644
--- a/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc
@@ -370,10 +370,13 @@ InvokeContext CustomCallThunk::BuildInvokeContext(
   ffi::ExecutionState* initialize_state = nullptr;
 
   if (execution_scoped_state) {
-    PrepareAndInitState& prepare_and_init = tsl::any_cast<PrepareAndInitState>(
-        execution_scoped_state->at(this->thunk_info().thunk_id));
-    prepare_state = &prepare_and_init.prepare;
-    initialize_state = &prepare_and_init.init;
+    if (auto it = execution_scoped_state->find(this->thunk_info().thunk_id);
+        it != execution_scoped_state->end()) {
+      PrepareAndInitState& prepare_and_init =
+          tsl::any_cast<PrepareAndInitState>(it->second);
+      prepare_state = &prepare_and_init.prepare;
+      initialize_state = &prepare_and_init.init;
+    }
   }
 
   // `called_computation_` is forwarded to the FFI handler both for direct
@@ -471,9 +474,14 @@ absl::Status CustomCallThunk::ExecuteFfiHandler(
 }
 
 absl::Status CustomCallThunk::Prepare(const PrepareParams& params) {
-  // Run the prepare stage at most once per execution. When lowered into a
-  // command buffer, CommandBufferThunk drives this same object via both its
-  // command sequence and its fallback thunk sequence; see PrepareAndInitState.
+  const auto* c_bundle = std::get_if<XLA_FFI_Handler_Bundle>(&bundle_);
+  const auto* owned_bundle = std::get_if<OwnedHandlerBundle>(&bundle_);
+  if ((c_bundle == nullptr || c_bundle->prepare == nullptr) &&
+      (owned_bundle == nullptr || owned_bundle->prepare == nullptr)) {
+    return absl::OkStatus();
+  }
+
+  // Run the prepare stage at most once per execution.
   if (params.execution_scoped_state != nullptr) {
     auto [it, _] = params.execution_scoped_state->try_emplace(
         thunk_info().thunk_id, std::in_place_type<PrepareAndInitState>);
@@ -487,8 +495,7 @@ absl::Status CustomCallThunk::Prepare(const PrepareParams& params) {
   const RunId run_id =
       params.collective_params ? params.collective_params->run_id : RunId{-1};
 
-  if (const auto* c_bundle = std::get_if<XLA_FFI_Handler_Bundle>(&bundle_);
-      c_bundle && c_bundle->prepare) {
+  if (c_bundle && c_bundle->prepare) {
     return ExecuteFfiHandler(
         run_id, c_bundle->prepare, XLA_FFI_ExecutionStage_PREPARE,
         /*stream=*/nullptr,
@@ -502,8 +509,7 @@ absl::Status CustomCallThunk::Prepare(const PrepareParams& params) {
         /*collective_memory=*/nullptr,
         /*computation_streams=*/{});
   }
-  if (const auto* owned_bundle = std::get_if<OwnedHandlerBundle>(&bundle_);
-      owned_bundle && owned_bundle->prepare) {
+  if (owned_bundle && owned_bundle->prepare) {
     return ExecuteFfiHandler(
         run_id, *owned_bundle->prepare, xla::ffi::ExecutionStage::kPrepare,
         /*stream=*/nullptr,
@@ -522,9 +528,14 @@ absl::Status CustomCallThunk::Prepare(const PrepareParams& params) {
 }
 
 absl::Status CustomCallThunk::Initialize(const InitializeParams& params) {
-  // Run the initialize stage at most once per execution. When lowered into a
-  // command buffer, CommandBufferThunk drives this same object via both its
-  // command sequence and its fallback thunk sequence; see PrepareAndInitState.
+  const auto* c_bundle = std::get_if<XLA_FFI_Handler_Bundle>(&bundle_);
+  const auto* owned_bundle = std::get_if<OwnedHandlerBundle>(&bundle_);
+  if ((c_bundle == nullptr || c_bundle->initialize == nullptr) &&
+      (owned_bundle == nullptr || owned_bundle->initialize == nullptr)) {
+    return absl::OkStatus();
+  }
+
+  // Run the initialize stage at most once per execution.
   if (params.execution_scoped_state != nullptr) {
     auto [it, _] = params.execution_scoped_state->try_emplace(
         thunk_info().thunk_id, std::in_place_type<PrepareAndInitState>);
@@ -538,8 +549,7 @@ absl::Status CustomCallThunk::Initialize(const InitializeParams& params) {
   const RunId run_id =
       params.collective_params ? params.collective_params->run_id : RunId{-1};
 
-  if (const auto* c_bundle = std::get_if<XLA_FFI_Handler_Bundle>(&bundle_);
-      c_bundle && c_bundle->initialize) {
+  if (c_bundle && c_bundle->initialize) {
     return ExecuteFfiHandler(
         run_id, *c_bundle->initialize, XLA_FFI_ExecutionStage_INITIALIZE,
         params.stream, params.execution_scoped_state,
@@ -550,8 +560,7 @@ absl::Status CustomCallThunk::Initialize(const InitializeParams& params) {
         params.collective_memory,
         /*computation_streams=*/{});
   }
-  if (const auto* owned_bundle = std::get_if<OwnedHandlerBundle>(&bundle_);
-      owned_bundle && owned_bundle->initialize) {
+  if (owned_bundle && owned_bundle->initialize) {
     return ExecuteFfiHandler(
         run_id, *owned_bundle->initialize,
         xla::ffi::ExecutionStage::kInitialize, params.stream,
diff --git a/third_party/xla/xla/backends/gpu/runtime/dynamic_slice_fusion_v2_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/dynamic_slice_fusion_v2_thunk.cc
index aab37e21641..de8feeb0a09 100644
--- a/third_party/xla/xla/backends/gpu/runtime/dynamic_slice_fusion_v2_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/dynamic_slice_fusion_v2_thunk.cc
@@ -176,10 +176,10 @@ std::string DynamicSliceFusionV2Thunk::ToString(int indent) const {
 
 absl::Status DynamicSliceFusionV2Thunk::Prepare(const PrepareParams& params) {
   if (!params.buffer_allocations) {
-    if (command_executor_.has_value()) {
-      ABSL_RETURN_IF_ERROR(command_executor_->Prepare(params));
+    if (!executor_.thunks().empty() || !command_executor_.has_value()) {
+      return executor_.Prepare(params);
     }
-    return executor_.Prepare(params);
+    return command_executor_->Prepare(params);
   }
 
   // Embedded thunks and commands use synthetic embedded allocation indices
@@ -193,19 +193,19 @@ absl::Status DynamicSliceFusionV2Thunk::Prepare(const PrepareParams& params) {
   PrepareParams embedded_params = params;
   embedded_params.buffer_allocations = &embedded_allocs;
 
-  if (command_executor_.has_value()) {
-    ABSL_RETURN_IF_ERROR(command_executor_->Prepare(embedded_params));
+  if (!executor_.thunks().empty() || !command_executor_.has_value()) {
+    return executor_.Prepare(embedded_params);
   }
-  return executor_.Prepare(embedded_params);
+  return command_executor_->Prepare(embedded_params);
 }
 
 absl::Status DynamicSliceFusionV2Thunk::Initialize(
     const InitializeParams& params) {
   if (!params.buffer_allocations) {
-    if (command_executor_.has_value()) {
-      ABSL_RETURN_IF_ERROR(command_executor_->Initialize(params));
+    if (!executor_.thunks().empty() || !command_executor_.has_value()) {
+      return executor_.Initialize(params);
     }
-    return executor_.Initialize(params);
+    return command_executor_->Initialize(params);
   }
 
   // Embedded thunks and commands use synthetic embedded allocation indices
@@ -219,10 +219,10 @@ absl::Status DynamicSliceFusionV2Thunk::Initialize(
   InitializeParams embedded_params = params;
   embedded_params.buffer_allocations = &embedded_allocs;
 
-  if (command_executor_.has_value()) {
-    ABSL_RETURN_IF_ERROR(command_executor_->Initialize(embedded_params));
+  if (!executor_.thunks().empty() || !command_executor_.has_value()) {
+    return executor_.Initialize(embedded_params);
   }
-  return executor_.Initialize(embedded_params);
+  return command_executor_->Initialize(embedded_params);
 }
 
 absl::Status DynamicSliceFusionV2Thunk::VerifyBufferAssignment(
diff --git a/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc
index 4c3570d6262..575c04db091 100644
--- a/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/ragged_all_to_all_thunk.cc
@@ -613,8 +613,6 @@ RaggedAllToAllThunk::GetCliqueRequirements(const GpuCliqueKey& clique_key,
 }
 
 absl::Status RaggedAllToAllThunk::Initialize(const InitializeParams& params) {
-  ABSL_RETURN_IF_ERROR(CollectiveThunk::Initialize(params));
-
   se::StreamExecutor* executor = params.executor;
   ABSL_RETURN_IF_ERROR(per_device_states_.GetOrCreateAndInitialize(
       executor->device_ordinal(),
diff --git a/third_party/xla/xla/backends/gpu/runtime/recv_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/recv_thunk.cc
index 165f9778db6..ac7eae2562f 100644
--- a/third_party/xla/xla/backends/gpu/runtime/recv_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/recv_thunk.cc
@@ -63,11 +63,6 @@ RecvThunk::RecvThunk(ThunkInfo thunk_info, const P2PConfig& config,
       config_(config),
       hlo_name_(instr_name) {}
 
-absl::Status RecvThunk::Initialize(const InitializeParams& params) {
-  ABSL_RETURN_IF_ERROR(CollectiveThunk::Initialize(params));
-  return absl::OkStatus();
-}
-
 absl::StatusOr<std::unique_ptr<RecvThunk>> RecvThunk::FromProto(
     ThunkInfo thunk_info, const RecvThunkProto& thunk_proto,
     absl::Span<const BufferAllocation> buffer_allocations) {
diff --git a/third_party/xla/xla/backends/gpu/runtime/recv_thunk.h b/third_party/xla/xla/backends/gpu/runtime/recv_thunk.h
index 0e3170f1df2..9c54eff13e7 100644
--- a/third_party/xla/xla/backends/gpu/runtime/recv_thunk.h
+++ b/third_party/xla/xla/backends/gpu/runtime/recv_thunk.h
@@ -46,7 +46,6 @@ class RecvThunk : public CollectiveThunk {
             const Buffer& buffer);
   RecvThunk(ThunkInfo thunk_info, const P2PConfig& config, const Buffer& buffer,
             absl::string_view instr_name);
-  absl::Status Initialize(const InitializeParams& params) override;
 
   static absl::StatusOr<std::unique_ptr<RecvThunk>> FromProto(
       ThunkInfo thunk_info, const RecvThunkProto& thunk_proto,
diff --git a/third_party/xla/xla/backends/gpu/runtime/send_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/send_thunk.cc
index 4d49962ba73..f6a69301643 100644
--- a/third_party/xla/xla/backends/gpu/runtime/send_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/send_thunk.cc
@@ -63,11 +63,6 @@ SendThunk::SendThunk(ThunkInfo thunk_info, const P2PConfig& config,
       config_(config),
       hlo_name_(instr_name) {}
 
-absl::Status SendThunk::Initialize(const InitializeParams& params) {
-  ABSL_RETURN_IF_ERROR(CollectiveThunk::Initialize(params));
-  return absl::OkStatus();
-}
-
 absl::StatusOr<const se::CommandBuffer::Command*> SendThunk::Record(
     const ExecuteParams& execute_params, const RecordParams& record_params,
     RecordAction record_action, se::CommandBuffer* command_buffer) {
diff --git a/third_party/xla/xla/backends/gpu/runtime/send_thunk.h b/third_party/xla/xla/backends/gpu/runtime/send_thunk.h
index ac624afaa35..0e42db9148f 100644
--- a/third_party/xla/xla/backends/gpu/runtime/send_thunk.h
+++ b/third_party/xla/xla/backends/gpu/runtime/send_thunk.h
@@ -47,8 +47,6 @@ class SendThunk : public CollectiveThunk {
   SendThunk(ThunkInfo thunk_info, const P2PConfig& config, const Buffer& buffer,
             absl::string_view instr_name);
 
-  absl::Status Initialize(const InitializeParams& params) override;
-
   absl::StatusOr<const se::CommandBuffer::Command*> Record(
       const ExecuteParams& execute_params, const RecordParams& record_params,
       RecordAction record_action, se::CommandBuffer* command_buffer) override;
diff --git a/third_party/xla/xla/backends/gpu/runtime/while_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/while_thunk.cc
index ec7d75ee892..e1fa1aea58e 100644
--- a/third_party/xla/xla/backends/gpu/runtime/while_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/while_thunk.cc
@@ -113,24 +113,28 @@ WhileThunk::WhileThunk(
       host_memory_pools_(devices_per_host) {}
 
 absl::Status WhileThunk::Prepare(const PrepareParams& params) {
-  ABSL_RETURN_IF_ERROR(condition_executor_.Prepare(params));
-  ABSL_RETURN_IF_ERROR(body_executor_.Prepare(params));
-  if (command_condition_executor_.has_value()) {
+  if (!condition_executor_.thunks().empty()) {
+    ABSL_RETURN_IF_ERROR(condition_executor_.Prepare(params));
+  } else if (command_condition_executor_.has_value()) {
     ABSL_RETURN_IF_ERROR(command_condition_executor_->Prepare(params));
   }
-  if (command_body_executor_.has_value()) {
+  if (!body_executor_.thunks().empty()) {
+    ABSL_RETURN_IF_ERROR(body_executor_.Prepare(params));
+  } else if (command_body_executor_.has_value()) {
     ABSL_RETURN_IF_ERROR(command_body_executor_->Prepare(params));
   }
   return absl::OkStatus();
 }
 
 absl::Status WhileThunk::Initialize(const InitializeParams& params) {
-  ABSL_RETURN_IF_ERROR(condition_executor_.Initialize(params));
-  ABSL_RETURN_IF_ERROR(body_executor_.Initialize(params));
-  if (command_condition_executor_.has_value()) {
+  if (!condition_executor_.thunks().empty()) {
+    ABSL_RETURN_IF_ERROR(condition_executor_.Initialize(params));
+  } else if (command_condition_executor_.has_value()) {
     ABSL_RETURN_IF_ERROR(command_condition_executor_->Initialize(params));
   }
-  if (command_body_executor_.has_value()) {
+  if (!body_executor_.thunks().empty()) {
+    ABSL_RETURN_IF_ERROR(body_executor_.Initialize(params));
+  } else if (command_body_executor_.has_value()) {
     ABSL_RETURN_IF_ERROR(command_body_executor_->Initialize(params));
   }
 
diff --git a/third_party/xla/xla/backends/gpu/runtime/while_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/while_thunk_test.cc
index 6deb42c57b3..e8a2f9f5309 100644
--- a/third_party/xla/xla/backends/gpu/runtime/while_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/while_thunk_test.cc
@@ -102,6 +102,7 @@ WhileThunk CreateWhileThunk(
 
 struct CommandRecordCounts {
   int prepares = 0;
+  int initializes = 0;
   int creates = 0;
   int updates = 0;
 };
@@ -118,6 +119,11 @@ class RecordingCommand : public Command {
     return absl::OkStatus();
   }
 
+  absl::Status Initialize(const InitializeParams&) override {
+    ++counts_->initializes;
+    return absl::OkStatus();
+  }
+
   absl::StatusOr<const se::CommandBuffer::Command*> Record(
       const Thunk::ExecuteParams&, const RecordParams&, RecordAction action,
       se::CommandBuffer*) override {
@@ -321,6 +327,69 @@ TEST(WhileThunkTest, PreparePropagatesToCommandBufferExecutors) {
   EXPECT_EQ(body_counts.prepares, 1);
 }
 
+TEST(WhileThunkTest, PrepareAndInitializeDoNotDuplicateSharedCommands) {
+  BufferAllocation pred_alloc(/*index=*/0, /*size=*/sizeof(bool), /*color=*/0);
+  BufferAllocation::Slice pred_slice(&pred_alloc, /*offset=*/0,
+                                     /*size=*/sizeof(bool));
+  CommandRecordCounts cond_counts;
+  CommandRecordCounts body_counts;
+  auto cond_cmd = std::make_unique<RecordingCommand>(&cond_counts);
+  auto body_cmd = std::make_unique<RecordingCommand>(&body_counts);
+
+  CommandSequence cond_commands;
+  cond_commands.Append(cond_cmd.get());
+  CommandSequence body_commands;
+  body_commands.Append(body_cmd.get());
+
+  ThunkSequence cond_thunks;
+  cond_thunks.push_back(std::move(cond_cmd));
+  ThunkSequence body_thunks;
+  body_thunks.push_back(std::move(body_cmd));
+
+  WhileThunk thunk(Thunk::ThunkInfo(), pred_slice, std::move(cond_thunks),
+                   std::move(body_thunks), /*trip_count=*/std::nullopt,
+                   /*devices_per_host=*/1);
+
+  ASSERT_OK_AND_ASSIGN(CommandExecutor cond_executor,
+                       CommandExecutor::Create(
+                           std::move(cond_commands),
+                           CommandExecutor::SynchronizationMode::kSerialize));
+  ASSERT_OK_AND_ASSIGN(CommandExecutor body_executor,
+                       CommandExecutor::Create(
+                           std::move(body_commands),
+                           CommandExecutor::SynchronizationMode::kSerialize));
+  ASSERT_OK(thunk.SetOrUpdateCommandBufferExecutors(
+      std::move(cond_executor), std::move(body_executor),
+      /*enable_loop_unroll=*/false));
+
+  ASSERT_OK_AND_ASSIGN(std::string platform_name,
+                       PlatformUtil::CanonicalPlatformName("gpu"));
+  ASSERT_OK_AND_ASSIGN(se::Platform * platform,
+                       se::PlatformManager::PlatformWithName(platform_name));
+  ASSERT_OK_AND_ASSIGN(se::StreamExecutor * executor,
+                       platform->ExecutorForDevice(0));
+
+  std::vector<se::DeviceAddressBase> buffers;
+  BufferAllocations allocations(buffers, /*device_ordinal=*/0,
+                                /*memory_allocator=*/nullptr);
+
+  Thunk::PrepareParams prepare_params{/*collective_params=*/nullptr,
+                                      /*collective_clique_requests=*/nullptr,
+                                      /*collective_memory_requests=*/nullptr,
+                                      /*executor=*/executor,
+                                      /*buffer_allocations=*/&allocations};
+  ASSERT_OK(thunk.Prepare(prepare_params));
+  EXPECT_EQ(cond_counts.prepares, 1);
+  EXPECT_EQ(body_counts.prepares, 1);
+
+  Thunk::InitializeParams init_params;
+  init_params.executor = executor;
+  init_params.buffer_allocations = &allocations;
+  ASSERT_OK(thunk.Initialize(init_params));
+  EXPECT_EQ(cond_counts.initializes, 1);
+  EXPECT_EQ(body_counts.initializes, 1);
+}
+
 TEST(WhileThunkTest, RecordCreatesAndUpdatesCommandBufferWhile) {
   BufferAllocation pred_alloc(/*index=*/0, /*size=*/sizeof(bool), /*color=*/0);
   BufferAllocation::Slice pred_slice(&pred_alloc, /*offset=*/0,
diff --git a/third_party/xla/xla/backends/gpu/tests/collective_ops_command_buffer_test.cc b/third_party/xla/xla/backends/gpu/tests/collective_ops_command_buffer_test.cc
index 6fad10528ba..c1dab0bdcf6 100644
--- a/third_party/xla/xla/backends/gpu/tests/collective_ops_command_buffer_test.cc
+++ b/third_party/xla/xla/backends/gpu/tests/collective_ops_command_buffer_test.cc
@@ -401,5 +401,78 @@ TEST_F(CollectiveOpsCommandBufferPeerAccessTest, FilterRaggedAllToAllOnly) {
       << dump_content;
 }
 
+TEST_F(CollectiveOpsCommandBufferTest, WhileWithAsyncAllReduce) {
+  constexpr absl::string_view kHloText = R"(
+  HloModule module, replica_count=2
+
+  add {
+    lhs = f32[] parameter(0)
+    rhs = f32[] parameter(1)
+    ROOT add = f32[] add(lhs, rhs)
+  }
+
+  cond {
+    param = (s32[], f32[2]) parameter(0)
+    iter = s32[] get-tuple-element(param), index=0
+    limit = s32[] constant(2)
+    ROOT cmp = pred[] compare(iter, limit), direction=LT
+  }
+
+  body {
+    param = (s32[], f32[2]) parameter(0)
+    iter = s32[] get-tuple-element(param), index=0
+    data = f32[2] get-tuple-element(param), index=1
+    one = s32[] constant(1)
+    next_iter = s32[] add(iter, one)
+    ar = f32[2] all-reduce(data), replica_groups={{0,1}}, to_apply=add
+    ROOT out = (s32[], f32[2]) tuple(next_iter, ar)
+  }
+
+  ENTRY main {
+    c0 = s32[] parameter(0)
+    data = f32[2] parameter(1)
+    init = (s32[], f32[2]) tuple(c0, data)
+    while_res = (s32[], f32[2]) while(init), condition=cond, body=body
+    ROOT res = f32[2] get-tuple-element(while_res), index=1
+  }
+  )";
+
+  constexpr int64_t kNumReplicas = 2;
+  if (device_count() < kNumReplicas) {
+    GTEST_SKIP() << "Test requires at least " << kNumReplicas << " devices ("
+                 << device_count() << " available)";
+  }
+
+  HloModuleConfig config =
+      GetModuleConfigForTest(/*replica_count=*/kNumReplicas);
+  config.mutable_debug_options().clear_xla_gpu_enable_command_buffer();
+  config.mutable_debug_options().add_xla_gpu_enable_command_buffer(
+      DebugOptions::FUSION);
+  config.mutable_debug_options().add_xla_gpu_enable_command_buffer(
+      DebugOptions::WHILE);
+  config.mutable_debug_options().add_xla_gpu_enable_command_buffer(
+      DebugOptions::COLLECTIVES);
+  config.mutable_debug_options().set_xla_gpu_graph_min_graph_size(1);
+  config.mutable_debug_options().add_xla_disable_hlo_passes(
+      "collective-fusion");
+
+  ASSERT_OK_AND_ASSIGN(auto module,
+                       ParseAndReturnVerifiedModule(kHloText, config));
+
+  Literal iter_init = LiteralUtil::CreateR0<int32_t>(0);
+  Literal input = LiteralUtil::CreateR1<float>({1.0f, 2.0f});
+  ASSERT_OK_AND_ASSIGN(
+      ExecutionResult execution_result,
+      ExecuteReplicated(std::move(module),
+                        std::vector<Literal*>{&iter_init, &input},
+                        /*run_hlo_passes=*/true));
+  const std::vector<Literal>& results = execution_result.results;
+  ASSERT_EQ(results.size(), kNumReplicas);
+  EXPECT_TRUE(LiteralTestUtil::Equal(LiteralUtil::CreateR1<float>({4.0f, 8.0f}),
+                                     results[0]));
+  EXPECT_TRUE(LiteralTestUtil::Equal(LiteralUtil::CreateR1<float>({4.0f, 8.0f}),
+                                     results[1]));
+}
+
 }  // namespace
 }  // namespace xla

From 1d2a71940e2f337f1aaf10f2849004330bf12c5a Mon Sep 17 00:00:00 2001
From: TJ Xu <tjx@nvidia.com>
Date: Tue, 13 Oct 2026 06:56:56 -0700
Subject: [PATCH 0695/1000] PR #50604: Enable dedicated host memcpy streams by
 default
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50604

📝 Summary of Changes

It was turned off by default due to the forward compatibility policy from google side.
Turning it on by default now.
Copybara import of the project:

--
682dfc67ba321d69a9a471c84c987854500787b9 by TJ Xu <tjx@nvidia.com>:

enable dedcated host memcpy streams by default

Merging this change closes #50604

PiperOrigin-RevId: 998604476
---
 third_party/xla/xla/service/gpu/compile_module_to_llvm_ir.cc | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/compile_module_to_llvm_ir.cc b/third_party/xla/xla/service/gpu/compile_module_to_llvm_ir.cc
index a89fb1bfb35..1326b4aacc6 100644
--- a/third_party/xla/xla/service/gpu/compile_module_to_llvm_ir.cc
+++ b/third_party/xla/xla/service/gpu/compile_module_to_llvm_ir.cc
@@ -103,9 +103,7 @@ CompileModuleResults InitializeResults(const HloModule* hlo_module) {
                       .xla_gpu_experimental_enable_collective_multi_streaming()
                   ? kDefaultNumCommunicationStreams
                   : 1,
-              // Keep disabled by default until the 14-day AOT
-              // forward-compatibility window has elapsed (9th of October 2026).
-              /*enable_dedicated_memcpy_streams=*/false,
+              /*enable_dedicated_memcpy_streams=*/true,
           });
   return results;
 }

From ff42e9551fb9d50d6d054d4bc42edc22d1bb80db Mon Sep 17 00:00:00 2001
From: Michael Kuperstein <mkuper@google.com>
Date: Tue, 13 Oct 2026 07:17:09 -0700
Subject: [PATCH 0696/1000] Clarify ReduceWindow init_values requirements in
 HLO operation semantics.

The `ReduceWindow` section in operation_semantics.md previously deferred the details of `init_values` to the `Reduce` specification rather than stating its requirements directly.

State explicitly in the `ReduceWindow` specification that `init_values` is a sequence of `N` scalars matching the operand element types and that all initial values must form an identity under `computation`, removing the cross-reference to `Reduce`.

PiperOrigin-RevId: 998614262
---
 third_party/xla/docs/operation_semantics.md | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/docs/operation_semantics.md b/third_party/xla/docs/operation_semantics.md
index 4fb025f39b4..2bedf6b5c75 100644
--- a/third_party/xla/docs/operation_semantics.md
+++ b/third_party/xla/docs/operation_semantics.md
@@ -4141,10 +4141,11 @@ window_strides, padding)`**
 :                     :                     : types `T_0,..., T_{N-1}`, each   :
 :                     :                     : representing the base area on    :
 :                     :                     : which the window is placed.      :
-| `init_values`       | `N XlaOps`          | The N starting values for the    |
+| `init_values`       | `N XlaOps`          | A sequence of N scalars of types |
+:                     :                     : `T_0,..., T_{N-1}`, representing :
+:                     :                     : the starting values for the      :
 :                     :                     : reduction, one for each of the N :
-:                     :                     : operands. See [Reduce](#reduce)  :
-:                     :                     : for details.                     :
+:                     :                     : operands.                        :
 | `computation`       | `XlaComputation`    | Reduction function of type `T_0, |
 :                     :                     : ..., T_{N-1}, T_0, ..., T_{N-1}  :
 :                     :                     : -> Collate(T_0, ..., T_{N-1})`,  :
@@ -4171,6 +4172,7 @@ Where:
 
 *   N is required to be greater or equal to 1.
 *   All input arrays must have the same dimensions.
+*   All initial values have to form an identity under `computation`.
 *   If `N = 1`, `Collate(T)` is `T`.
 *   If `N > 1`, `Collate(T_0, ..., T_{N-1})` is a tuple of `N` elements of type
     `(T0,...T{N-1})`.

From 083ff59e4cc08546606e60959cca577957f362f3 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Tue, 13 Oct 2026 22:58:58 +0530
Subject: [PATCH 0697/1000] Address review: integration test for batched
 offsets, is_fully_defined check

- DTensorDataset.__init__ / _validate_input: replace the rank-only check
  with is_fully_defined() so a dataset element spec with a known rank but
  an unknown non-batch dimension size is rejected up front with a clear
  ValueError, instead of silently passing validation and later hitting a
  TypeError in _index_matrix.
- Add testDatasetAlreadyBatchedOffsets, an end-to-end integration test
  that instantiates DTensorDataset with dataset_already_batched=True and
  a non-batch sharded layout, verifying the previously-fixed
  _index_matrices construction is exercised at runtime, not just in the
  _index_matrix unit test.
- Update testIterableFailsWithUnknownShapeDatasetSpec's expected message
  to match the new is_fully_defined() error text.
---
 tensorflow/dtensor/python/input_util.py       |  4 +--
 .../dtensor/python/tests/input_util_test.py   | 30 ++++++++++++++++++-
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/tensorflow/dtensor/python/input_util.py b/tensorflow/dtensor/python/input_util.py
index fc89c239aef..033c1541b3e 100644
--- a/tensorflow/dtensor/python/input_util.py
+++ b/tensorflow/dtensor/python/input_util.py
@@ -226,9 +226,9 @@ def _validate_input(flattened_layouts: Sequence[layout_lib.Layout],
   first_elem_shape = flattened_elem_spec[0].shape
 
   for layout, elem_spec in zip(flattened_layouts, flattened_elem_spec):
-    if elem_spec.shape.rank is None:
+    if not elem_spec.shape.is_fully_defined():
       raise ValueError(
-          'Dataset element shape must have a valid rank, got spec %s.' %
+          'Dataset element shape must be fully defined, got spec %s.' %
           elem_spec)
 
     # Check that layout's rank matches the element's rank. If dataset is not yet
diff --git a/tensorflow/dtensor/python/tests/input_util_test.py b/tensorflow/dtensor/python/tests/input_util_test.py
index 306eb03359c..a062dbbba70 100644
--- a/tensorflow/dtensor/python/tests/input_util_test.py
+++ b/tensorflow/dtensor/python/tests/input_util_test.py
@@ -83,7 +83,7 @@ class DTensorDatasetTest(test_util.DTensorBaseTest):
             tensor_shape.TensorShape(None), dtype=dtypes.int32))
 
     with self.assertRaisesRegex(
-        ValueError, 'Dataset element shape must have a valid rank'):
+        ValueError, 'Dataset element shape must be fully defined'):
       input_util.DTensorDataset(
           dataset=dataset,
           global_batch_size=8,
@@ -105,6 +105,34 @@ class DTensorDatasetTest(test_util.DTensorBaseTest):
           layouts=images_layout,
           batch_dim=MESH_DIM_BATCH)
 
+  def testDatasetAlreadyBatchedOffsets(self):
+    layout = Layout(
+        [MESH_DIM_BATCH, MESH_DIM_HEIGHT, MESH_DIM_WIDTH], self.mesh)
+
+    # Global batch size 4. Non-batch sharded on height and width, so this
+    # exercises the `dataset_already_batched=True` + non-batch-sharded-layout
+    # path fixed in `_index_matrix` / `DTensorDataset.__init__` (previously,
+    # the un-batched `element_spec` was used to build `_index_matrices`,
+    # which shifted every non-batch axis by one and produced wrong slice
+    # offsets whenever a non-batch sharded axis's size differed from the
+    # batch size).
+    images = self._images([4, 4, 6])
+    dataset = dataset_ops.DatasetV2.from_tensor_slices(images).batch(
+        4, drop_remainder=True)
+
+    d_dataset = input_util.DTensorDataset(
+        dataset=dataset,
+        global_batch_size=4,
+        mesh=self.mesh,
+        layouts=layout,
+        batch_dim=MESH_DIM_BATCH,
+        dataset_already_batched=True)
+
+    iterator = iter(d_dataset)
+    d_tensor = next(iterator)
+
+    self.assertDTensorEqual(images, layout, d_tensor)
+
   @parameterized.named_parameters(('Eager', False), ('Graph', True))
   def testRangeIteration(self, is_graph):
     batch_size = 8

From 3b61493f6692f89cd2b6aabdcbc1ee32d328ccab Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Tue, 13 Oct 2026 23:02:56 +0530
Subject: [PATCH 0698/1000] Address review: fix flawed negative-test boundary
 in ring_bounds test

testCreateMeshWithInvalidRingBounds previously matched topology_shape's
first element and only incremented the second, which would have raised
under the old lexicographical '>' comparison too (first tied, second
larger => list comparison already True), so it didn't actually exercise
the bug being fixed. Decrementing the first element and incrementing the
second makes the lexicographical comparison evaluate False (old, buggy
behavior) while the element-wise comparison correctly evaluates True
(fixed behavior), verified empirically.
---
 .../dtensor/python/tests/tpu_device_assignment_test.py     | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/tensorflow/dtensor/python/tests/tpu_device_assignment_test.py b/tensorflow/dtensor/python/tests/tpu_device_assignment_test.py
index b4537e51012..b5f79cbf0ec 100644
--- a/tensorflow/dtensor/python/tests/tpu_device_assignment_test.py
+++ b/tensorflow/dtensor/python/tests/tpu_device_assignment_test.py
@@ -804,10 +804,11 @@ class DeviceAssignmentTest(test_util.DTensorBaseTest):
 
   def testCreateMeshWithInvalidRingBounds(self):
     topology_shape = list(tpu_util._tpu_topology.mesh_shape)
-    # A bound that exceeds the topology in a non-first dimension
-    # would pass a lexicographical check but must fail element-wise.
+    # A bound that exceeds the topology in a non-first dimension but is
+    # smaller in the first dimension would pass a lexicographical check
+    # but must fail element-wise.
     invalid_ring_bounds = [
-        topology_shape[0],
+        topology_shape[0] - 1,
         topology_shape[1] + 1,
         topology_shape[2],
         topology_shape[3],

From 4e6c3c924bb0f173356537cbcaf52fe8b552e8e5 Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Tue, 13 Oct 2026 10:36:00 -0700
Subject: [PATCH 0699/1000] Use a local RandomState and test empty inputs in
 testHighRank

np.random.seed changes NumPy's global state for the rest of the
process, so draw the features from a local RandomState instead. Also
run each rank with a zero-sized leading dimension, which the oneDNN
kernel collapses to a 2-D input with zero rows.
---
 .../python/kernel_tests/nn_ops/softmax_op_test.py    | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py b/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py
index 4f6223ea635..cf6d29da872 100644
--- a/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py
@@ -279,11 +279,15 @@ class SoftmaxTest(test.TestCase):
   def testHighRank(self):
     # Regression test for GitHub issue 128814: the oneDNN kernel rejected
     # inputs of more than 5 dimensions.
-    np.random.seed(0)
+    rng = np.random.RandomState(0)
     for rank in (6, 7):
-      features = np.random.rand(*([2] * (rank - 1) + [3])).astype(np.float32)
-      self._testSoftmax(features)
-      self._testSoftmax(features, dim=1)
+      for empty in (False, True):
+        shape = [2] * (rank - 1) + [3]
+        if empty:
+          shape[2] = 0
+        features = rng.rand(*shape).astype(np.float32)
+        self._testSoftmax(features)
+        self._testSoftmax(features, dim=1)
 
   def testShapeInference(self):
     op = nn_ops.softmax([[[1., 1., 1., 1.], [1., 2., 3., 4.]],

From b3598b0012307242d81a909dc87b656a82791e91 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 13 Oct 2026 08:11:23 -0700
Subject: [PATCH 0700/1000] Remove redundant LIBTPU_ON_GCE preprocessor guard
 from tf_record_dataset_op.cc, preserving the PLATFORM_CLOUD_TPU and
 TPU_GCS_FS checks.

PiperOrigin-RevId: 998645139
---
 tensorflow/core/kernels/data/tf_record_dataset_op.cc | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tensorflow/core/kernels/data/tf_record_dataset_op.cc b/tensorflow/core/kernels/data/tf_record_dataset_op.cc
index b6fdc67e874..f98ef9a4986 100644
--- a/tensorflow/core/kernels/data/tf_record_dataset_op.cc
+++ b/tensorflow/core/kernels/data/tf_record_dataset_op.cc
@@ -55,9 +55,7 @@ constexpr int64_t kCloudTpuBlockSize = 127LL << 20;  // 127MB.
 constexpr int64_t kS3BlockSize = kCloudTpuBlockSize;
 
 bool is_cloud_tpu_gcs_fs() {
-#if defined(LIBTPU_ON_GCE)
-  return true;
-#elif defined(PLATFORM_CLOUD_TPU) && defined(TPU_GCS_FS)
+#if defined(PLATFORM_CLOUD_TPU) && defined(TPU_GCS_FS)
   return true;
 #else
   return false;

From a9df6faef3d65475c71119e4d3d7badad56ff1a3 Mon Sep 17 00:00:00 2001
From: greypilgrim-083 <prashantchaudhary7353@gmail.com>
Date: Tue, 13 Oct 2026 23:19:09 +0530
Subject: [PATCH 0701/1000] Address review: fix gtest_main dep conflict and add
 green-path tests to sliding_window_dataset_op_test

- Remove @com_google_googletest//:gtest_main from sliding_window_dataset_op_test
  deps in experimental/BUILD to fix multiple main() definition linker error
  (//tensorflow/core:test_main already provides main())
- Add SizeOneWithExtremeStrideDoesNotOverflow test: verifies size=1 with
  extreme stride initializes successfully (no overflow)
- Add SignedInt64OverflowingTargetBufferSize test: verifies size=3 with
  stride=2^62 returns InvalidArgument due to overflow
---
 .../core/kernels/data/experimental/BUILD      |  1 -
 .../sliding_window_dataset_op_test.cc         | 33 +++++++++++++++++++
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/data/experimental/BUILD b/tensorflow/core/kernels/data/experimental/BUILD
index 84eae61f1a4..33051088c08 100644
--- a/tensorflow/core/kernels/data/experimental/BUILD
+++ b/tensorflow/core/kernels/data/experimental/BUILD
@@ -868,7 +868,6 @@ tf_cc_test(
         "//tensorflow/core/kernels/data:range_dataset_op",
         "//tensorflow/core/platform:status",
         "@com_google_absl//absl/status",
-        "@com_google_googletest//:gtest_main",
         "@xla//xla/tsl/lib/core:status_test_util",
         "@xla//xla/tsl/platform:status_matchers",
     ],
diff --git a/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op_test.cc b/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op_test.cc
index b41c1a9e687..5725d28b632 100644
--- a/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op_test.cc
+++ b/tensorflow/core/kernels/data/experimental/sliding_window_dataset_op_test.cc
@@ -141,6 +141,39 @@ TEST_F(SlidingWindowDatasetOpTest, AddOverflowingTargetBufferSize) {
                              ::testing::HasSubstr("overflow")));
 }
 
+TEST_F(SlidingWindowDatasetOpTest, SizeOneWithExtremeStrideDoesNotOverflow) {
+  auto dataset_params = SlidingWindowDatasetParams(
+      RangeDatasetParams(0, 3, 1),
+      /*size=*/1,
+      /*shift=*/1,
+      /*stride=*/std::numeric_limits<int64_t>::max(),
+      /*drop_remainder=*/true,
+      /*output_dtypes=*/{DT_VARIANT},
+      /*output_shapes=*/{PartialTensorShape({})},
+      /*node_name=*/kNodeName);
+  TF_EXPECT_OK(Initialize(dataset_params));
+}
+
+SlidingWindowDatasetParams SlidingWindowDatasetParamsWithSignedInt64Overflow() {
+  return SlidingWindowDatasetParams(
+      RangeDatasetParams(0, 3, 1),
+      /*size=*/3,
+      /*shift=*/1,
+      /*stride=*/static_cast<int64_t>(1) << 62,
+      /*drop_remainder=*/true,
+      /*output_dtypes=*/{DT_VARIANT},
+      /*output_shapes=*/{PartialTensorShape({})},
+      /*node_name=*/kNodeName);
+}
+
+TEST_F(SlidingWindowDatasetOpTest, SignedInt64OverflowingTargetBufferSize) {
+  auto dataset_params = SlidingWindowDatasetParamsWithSignedInt64Overflow();
+  EXPECT_THAT(
+      Initialize(dataset_params),
+      tsl::testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                             ::testing::HasSubstr("overflow")));
+}
+
 }  // namespace
 }  // namespace experimental
 }  // namespace data

From 842b77fce53771df2d98c51a77ff60441c5a457f Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Tue, 13 Oct 2026 08:51:06 -0700
Subject: [PATCH 0702/1000] [XLA:GPU] Cut host overhead of CustomCallThunk FFI
 stages.

PiperOrigin-RevId: 998666879
---
 .../xla/xla/backends/gpu/runtime/BUILD        |  3 +-
 .../backends/gpu/runtime/custom_call_thunk.cc | 98 ++++++++++++-------
 .../backends/gpu/runtime/custom_call_thunk.h  | 68 +++++++++++--
 3 files changed, 126 insertions(+), 43 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index 3656013960b..394bad04d18 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -1037,6 +1037,7 @@ cc_library(
         ":collective_params",
         ":command",
         ":command_state",
+        ":per_device_state",
         ":record_ffi",
         ":thunk",
         ":thunk_proto_cc",
@@ -1060,7 +1061,6 @@ cc_library(
         "//xla/ffi/api:record_api",
         "//xla/hlo/ir:hlo",
         "//xla/runtime:buffer_use",
-        "//xla/runtime:object_pool",
         "//xla/service:buffer_assignment",
         "//xla/service:hlo_proto_cc",
         "//xla/service:shaped_slice",
@@ -5503,7 +5503,6 @@ cc_library(
     name = "per_device_state",
     srcs = ["per_device_state.cc"],
     hdrs = ["per_device_state.h"],
-    visibility = ["//visibility:private"],
     deps = [
         ":cow_storage",
         ":device_slot",
diff --git a/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc b/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc
index 8a7a3762b5c..f350f394b15 100644
--- a/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "xla/backends/gpu/runtime/custom_call_thunk.h"
 
 #include <algorithm>
+#include <atomic>
 #include <cstdint>
 #include <functional>
 #include <memory>
@@ -64,7 +65,6 @@ limitations under the License.
 #include "xla/ffi/record_ffi.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/primitive_util.h"
-#include "xla/runtime/object_pool.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/gpu/buffer_allocations.h"
 #include "xla/service/hlo.pb.h"
@@ -242,7 +242,7 @@ absl::StatusOr<std::unique_ptr<CustomCallThunk>> CustomCallThunk::Create(
       thunk_info, std::move(target_name), std::move(bundle),
       std::move(operands), std::move(results), std::move(call_frame),
       std::move(attributes), std::move(execution_state), called_computation,
-      cpu_target_machine_options, use_pdl));
+      gpu_compute_capability, cpu_target_machine_options, use_pdl));
 }
 
 absl::StatusOr<std::unique_ptr<CustomCallThunk>> CustomCallThunk::Create(
@@ -282,7 +282,22 @@ absl::StatusOr<std::unique_ptr<CustomCallThunk>> CustomCallThunk::Create(
       thunk_info, std::move(target_name), std::move(bundle),
       std::move(operands), std::move(results), std::move(call_frame),
       std::move(attributes), std::move(execution_state), called_computation,
-      cpu_target_machine_options));
+      gpu_compute_capability, cpu_target_machine_options));
+}
+
+static std::vector<CustomCallThunk::BufferArg> ToBufferArgs(
+    absl::Span<const NullableShapedSlice> slices) {
+  std::vector<CustomCallThunk::BufferArg> args;
+  args.reserve(slices.size());
+  for (const NullableShapedSlice& slice : slices) {
+    if (slice.has_value()) {
+      args.push_back({slice->slice.index(), slice->slice.offset(),
+                      slice->slice.size(), /*present=*/true});
+    } else {
+      args.push_back(CustomCallThunk::BufferArg{});
+    }
+  }
+  return args;
 }
 
 CustomCallThunk::CustomCallThunk(
@@ -293,52 +308,70 @@ CustomCallThunk::CustomCallThunk(
     ffi::AttributesMap attributes,
     std::unique_ptr<ffi::ExecutionState> execution_state,
     const HloComputation* called_computation,
+    const se::GpuComputeCapability& gpu_compute_capability,
     std::optional<xla::cpu::TargetMachineOptions> cpu_target_machine_options,
     bool use_pdl)
     : TracedCommand(Thunk::kCustomCall, thunk_info),
       target_name_(std::move(target_name)),
       operands_(std::move(operands)),
       results_(std::move(results)),
+      operand_args_(ToBufferArgs(operands_)),
+      result_args_(ToBufferArgs(results_)),
       bundle_(std::move(bundle)),
       use_pdl_(use_pdl),
       attributes_(std::move(attributes)),
       call_frame_(std::move(call_frame)),
-      call_frames_([this] { return call_frame_->Copy(); }),
       execution_state_(std::move(execution_state)),
       called_computation_(called_computation),
+      gpu_compute_capability_(gpu_compute_capability),
       cpu_target_machine_options_(std::move(cpu_target_machine_options)) {}
 
-absl::StatusOr<ObjectPool<CallFrame>::BorrowedObject>
-CustomCallThunk::BuildCallFrame(
+absl::StatusOr<CustomCallThunk::CallFrameLease> CustomCallThunk::BuildCallFrame(
     const BufferAllocations* absl_nullable buffer_allocations) {
-  auto device_memory = [&](BufferAllocation::Slice slice) {
-    return buffer_allocations ? buffer_allocations->GetDeviceAddress(slice)
-                              : se::DeviceAddressBase{};
+  if (!call_frame_.has_value()) {
+    return Internal("Call frame prototype is not initialized");
+  }
+
+  auto device_memory = [&](const BufferArg& arg) -> se::DeviceAddressBase {
+    if (!arg.present || buffer_allocations == nullptr) {
+      return se::DeviceAddressBase{};
+    }
+    se::DeviceAddressBase base =
+        buffer_allocations->GetDeviceAddress(arg.index);
+    DCHECK_LE(arg.offset + arg.size, base.size());
+    return base.GetByteSlice(arg.offset, arg.size);
   };
 
-  absl::InlinedVector<se::DeviceAddressBase, 8> arguments;
-  arguments.reserve(operands_.size());
-  for (auto& operand : operands_) {
-    if (!operand.has_value()) {
-      arguments.push_back(se::DeviceAddressBase{});
-    } else {
-      arguments.push_back(device_memory(operand->slice));
-    }
+  absl::InlinedVector<se::DeviceAddressBase, 32> arguments;
+  arguments.reserve(operand_args_.size());
+  for (const BufferArg& arg : operand_args_) {
+    arguments.push_back(device_memory(arg));
   }
 
-  absl::InlinedVector<se::DeviceAddressBase, 4> results;
-  results.reserve(results_.size());
-  for (auto& result : results_) {
-    if (!result.has_value()) {
-      results.push_back(se::DeviceAddressBase{});
-    } else {
-      results.push_back(device_memory(result->slice));
-    }
+  absl::InlinedVector<se::DeviceAddressBase, 8> results;
+  results.reserve(result_args_.size());
+  for (const BufferArg& arg : result_args_) {
+    results.push_back(device_memory(arg));
   }
 
-  ABSL_ASSIGN_OR_RETURN(auto call_frame, call_frames_->GetOrCreate());
-  ABSL_RETURN_IF_ERROR(call_frame->UpdateWithBuffers(arguments, results));
-  return call_frame;
+  // Fast path: the per-device call frame, if no other execution on this
+  // device currently holds it. Otherwise fall back to a private copy.
+  int device_ordinal =
+      buffer_allocations ? buffer_allocations->device_ordinal() : 0;
+  ABSL_ASSIGN_OR_RETURN(CallFrameSlot * slot,
+                   call_frames_.GetOrCreate(device_ordinal));
+  std::optional<CallFrameLease> lease;
+  if (!slot->busy.exchange(true, std::memory_order_acquire)) {
+    if (!slot->frame.has_value()) {
+      slot->frame.emplace(call_frame_->Copy());
+    }
+    lease.emplace(slot);
+  } else {
+    lease.emplace(call_frame_->Copy());
+  }
+
+  ABSL_RETURN_IF_ERROR((*lease)->UpdateWithBuffers(arguments, results));
+  return std::move(*lease);
 }
 
 InvokeContext CustomCallThunk::BuildInvokeContext(
@@ -360,11 +393,8 @@ InvokeContext CustomCallThunk::BuildInvokeContext(
     allocator = buffer_allocations->memory_allocator();
   }
 
-  const se::GpuComputeCapability* gpu_compute_capability = nullptr;
-  if (stream != nullptr) {
-    gpu_compute_capability =
-        &stream->parent()->GetDeviceDescription().gpu_compute_capability();
-  }
+  const se::GpuComputeCapability* gpu_compute_capability =
+      &gpu_compute_capability_;
 
   ffi::ExecutionState* prepare_state = nullptr;
   ffi::ExecutionState* initialize_state = nullptr;
@@ -613,8 +643,6 @@ absl::StatusOr<const se::CommandBuffer::Command*> CustomCallThunk::Record(
     const RecordParams& record_params, RecordAction record_action,
     se::CommandBuffer* command_buffer) {
   se::StreamExecutor* executor = execute_params.stream->parent();
-  ABSL_ASSIGN_OR_RETURN(auto call_frame,
-                   BuildCallFrame(execute_params.buffer_allocations));
 
   // Retrieve or create the state for recording.
   auto* state = record_params.state.GetOrCreate<CustomCallRecordState>(
diff --git a/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.h b/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.h
index 0e6f8d3c7f8..4f57d273e05 100644
--- a/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.h
+++ b/third_party/xla/xla/backends/gpu/runtime/custom_call_thunk.h
@@ -16,9 +16,12 @@ limitations under the License.
 #ifndef XLA_BACKENDS_GPU_RUNTIME_CUSTOM_CALL_THUNK_H_
 #define XLA_BACKENDS_GPU_RUNTIME_CUSTOM_CALL_THUNK_H_
 
+#include <atomic>
+#include <cstdint>
 #include <memory>
 #include <optional>
 #include <string>
+#include <utility>
 #include <variant>
 #include <vector>
 
@@ -31,6 +34,7 @@ limitations under the License.
 #include "xla/backends/gpu/runtime/collective_cliques.h"
 #include "xla/backends/gpu/runtime/collective_memory.h"
 #include "xla/backends/gpu/runtime/command.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk.pb.h"
 #include "xla/backends/gpu/runtime/traced_command.h"
@@ -44,7 +48,6 @@ limitations under the License.
 #include "xla/ffi/invoke.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/runtime/buffer_use.h"
-#include "xla/runtime/object_pool.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/gpu/buffer_allocations.h"
 #include "xla/service/shaped_slice.h"
@@ -192,7 +195,50 @@ class CustomCallThunk : public TracedCommand {
       const se::GpuComputeCapability& gpu_compute_capability,
       std::optional<xla::cpu::TargetMachineOptions> cpu_target_machine_options);
 
+  // A buffer argument resolved at construction time. `present == false`
+  // represents a null pointer argument.
+  struct BufferArg {
+    BufferAllocation::Index index = -1;
+    int64_t offset = 0;
+    int64_t size = 0;
+    bool present = false;
+  };
+
  private:
+  // Per-device call frame reused across executions. `busy` is set while a
+  // `CallFrameLease` refers to `frame`; concurrent executions on the same
+  // device fall back to a private copy of the reference call frame.
+  struct CallFrameSlot {
+    std::optional<ffi::CallFrame> frame;
+    std::atomic<bool> busy{false};
+  };
+
+  // RAII handle to a call frame: either the per-device slot (released on
+  // destruction) or an owned copy.
+  class CallFrameLease {
+   public:
+    explicit CallFrameLease(CallFrameSlot* slot) : slot_(slot) {}
+    explicit CallFrameLease(ffi::CallFrame frame) : owned_(std::move(frame)) {}
+    ~CallFrameLease() {
+      if (slot_ != nullptr) {
+        slot_->busy.store(false, std::memory_order_release);
+      }
+    }
+    CallFrameLease(CallFrameLease&& other) noexcept
+        : slot_(std::exchange(other.slot_, nullptr)),
+          owned_(std::move(other.owned_)) {}
+    CallFrameLease& operator=(CallFrameLease&&) = delete;
+    CallFrameLease(const CallFrameLease&) = delete;
+    CallFrameLease& operator=(const CallFrameLease&) = delete;
+
+    ffi::CallFrame& operator*() { return slot_ ? *slot_->frame : *owned_; }
+    ffi::CallFrame* operator->() { return &**this; }
+
+   private:
+    CallFrameSlot* slot_ = nullptr;
+    std::optional<ffi::CallFrame> owned_;
+  };
+
   CustomCallThunk(
       ThunkInfo thunk_info, std::string target_name,
       std::variant<XLA_FFI_Handler_Bundle, OwnedHandlerBundle> bundle,
@@ -201,11 +247,12 @@ class CustomCallThunk : public TracedCommand {
       xla::ffi::AttributesMap attributes,
       std::unique_ptr<ffi::ExecutionState> execution_state,
       const HloComputation* called_computation,
+      const se::GpuComputeCapability& gpu_compute_capability,
       std::optional<xla::cpu::TargetMachineOptions> cpu_target_machine_options,
       bool use_pdl = false);
 
-  absl::StatusOr<ObjectPool<xla::ffi::CallFrame>::BorrowedObject>
-  BuildCallFrame(const BufferAllocations* absl_nullable buffer_allocations);
+  absl::StatusOr<CallFrameLease> BuildCallFrame(
+      const BufferAllocations* absl_nullable buffer_allocations);
 
   xla::ffi::InvokeContext BuildInvokeContext(
       RunId run_id, se::Stream* absl_nullable stream,
@@ -254,6 +301,12 @@ class CustomCallThunk : public TracedCommand {
   std::vector<NullableShapedSlice> operands_;
   std::vector<NullableShapedSlice> results_;
 
+  // Compact copies of the operand/result slices iterated on every execution
+  // (`NullableShapedSlice` carries a `Shape` and is too large to stream over
+  // for thousands of thunks per step).
+  std::vector<BufferArg> operand_args_;
+  std::vector<BufferArg> result_args_;
+
   // registry) or an owned bundle (from xla::ffi::Bind()).
   std::variant<XLA_FFI_Handler_Bundle, OwnedHandlerBundle> bundle_;
   bool use_pdl_ = false;
@@ -262,9 +315,9 @@ class CustomCallThunk : public TracedCommand {
   // Reference call frame pre-initialized at construction time.
   std::optional<ffi::CallFrame> call_frame_;
 
-  // A pool of call frames used at run time. Newly created call frames are
-  // copied from the reference call frame and updated with buffer addresses.
-  std::optional<ObjectPool<ffi::CallFrame>> call_frames_;
+  // Per-device call frames used at run time, copied from the reference call
+  // frame on first use and updated with buffer addresses on every call.
+  PerDeviceState<CallFrameSlot> call_frames_;
 
   // Execution state bound to the FFI handler instance. Optional.
   std::shared_ptr<ffi::ExecutionState> execution_state_;
@@ -279,6 +332,9 @@ class CustomCallThunk : public TracedCommand {
   // custom calls that access called computation can only be linked statically.
   const HloComputation* called_computation_ = nullptr;
 
+  // Compute capability of the target device, forwarded to FFI handlers.
+  se::GpuComputeCapability gpu_compute_capability_;
+
   std::optional<xla::cpu::TargetMachineOptions> cpu_target_machine_options_;
 };
 

From f9649b6d550d151767cbb15f4a433962160458e7 Mon Sep 17 00:00:00 2001
From: Sean Talts <talts@google.com>
Date: Tue, 13 Oct 2026 08:58:43 -0700
Subject: [PATCH 0703/1000] [XLA:CPU] Run while loop bodies with buffers <= 64
 KiB sequentially in ThunkExecutor

Concurrency overheads from the DAG executor (per-iteration ExecuteState
allocation, atomic counter countdowns, thread pool work stealing) dominate
execution time across while loop iterations unless buffer sizes are large.
Configure WhileThunk's nested condition and body executors with a 64 KiB
execute_sequential_buffer_threshold (up from the 512 B default), so small-buffer
loops execute sequentially while large-buffer loops retain multi-core DAG
concurrency.

PiperOrigin-RevId: 998671411
---
 .../xla/backends/cpu/runtime/while_thunk.cc   | 15 +++++--
 .../backends/cpu/runtime/while_thunk_test.cc  | 41 +++++++++++++++++++
 2 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/backends/cpu/runtime/while_thunk.cc b/third_party/xla/xla/backends/cpu/runtime/while_thunk.cc
index 7c863b0eaf2..18d2be3804c 100644
--- a/third_party/xla/xla/backends/cpu/runtime/while_thunk.cc
+++ b/third_party/xla/xla/backends/cpu/runtime/while_thunk.cc
@@ -50,10 +50,17 @@ namespace xla::cpu {
 absl::StatusOr<std::unique_ptr<WhileThunk>> WhileThunk::Create(
     Info info, BufferAllocation::Slice cond_buffer, ThunkSequence cond_sequence,
     ThunkSequence body_sequence, std::optional<int64_t> trip_count) {
-  ABSL_ASSIGN_OR_RETURN(ThunkExecutor cond_executor,
-                   ThunkExecutor::Create(std::move(cond_sequence)));
-  ABSL_ASSIGN_OR_RETURN(ThunkExecutor body_executor,
-                   ThunkExecutor::Create(std::move(body_sequence)));
+  ThunkExecutor::Options executor_options;
+  // Concurrency overheads across while loop iterations (ExecuteState
+  // allocation, ready queue synchronization, thread pool work stealing)
+  // dominate execution time unless buffer sizes are large.
+  executor_options.execute_sequential_buffer_threshold = 64 * 1024;
+  ABSL_ASSIGN_OR_RETURN(
+      ThunkExecutor cond_executor,
+      ThunkExecutor::Create(std::move(cond_sequence), executor_options));
+  ABSL_ASSIGN_OR_RETURN(
+      ThunkExecutor body_executor,
+      ThunkExecutor::Create(std::move(body_sequence), executor_options));
 
   if (cond_buffer.size() != sizeof(bool)) {
     return Internal("Unsupported cond buffer size %d", cond_buffer.size());
diff --git a/third_party/xla/xla/backends/cpu/runtime/while_thunk_test.cc b/third_party/xla/xla/backends/cpu/runtime/while_thunk_test.cc
index f68b9fd68e8..f1c2d1b5fcd 100644
--- a/third_party/xla/xla/backends/cpu/runtime/while_thunk_test.cc
+++ b/third_party/xla/xla/backends/cpu/runtime/while_thunk_test.cc
@@ -244,5 +244,46 @@ TEST(WhileThunkTest, NonBlockingExecuteWithTripCount) {
   EXPECT_EQ(counter, LiteralUtil::CreateR0<int32_t>(kNumIterations));
 }
 
+TEST(WhileThunkTest, SequentialExecution) {
+  BufferAllocation alloc(0, 2 * 1024 * 1024, 0);
+  BufferAllocation::Slice pred_slice(&alloc, 0, sizeof(bool));
+
+  ThunkSequence cond_sequence;
+  cond_sequence.push_back(std::make_unique<BufferUseThunk>(
+      BufferUse::Read(pred_slice, ShapeUtil::MakeShape(PRED, {1}))));
+
+  ThunkSequence body_sequence;
+  for (int i = 0; i < 10; ++i) {
+    BufferAllocation::Slice slice(&alloc, 1024 + i * 2048, 2048);
+    body_sequence.push_back(std::make_unique<BufferUseThunk>(
+        BufferUse::Write(slice, ShapeUtil::MakeShape(U8, {2048}))));
+  }
+
+  ASSERT_OK_AND_ASSIGN(
+      auto thunk,
+      WhileThunk::Create({"while"}, pred_slice, std::move(cond_sequence),
+                         std::move(body_sequence)));
+  EXPECT_TRUE(thunk->body_executor().is_sequential());
+
+  ThunkSequence large_cond_sequence;
+  large_cond_sequence.push_back(std::make_unique<BufferUseThunk>(
+      BufferUse::Read(pred_slice, ShapeUtil::MakeShape(PRED, {1}))));
+
+  constexpr int64_t kLargeSliceSize = 128 * 1024;
+  ThunkSequence large_body_sequence;
+  for (int i = 0; i < 10; ++i) {
+    BufferAllocation::Slice slice(&alloc, 1024 + i * kLargeSliceSize,
+                                  kLargeSliceSize);
+    large_body_sequence.push_back(std::make_unique<BufferUseThunk>(
+        BufferUse::Write(slice, ShapeUtil::MakeShape(U8, {kLargeSliceSize}))));
+  }
+
+  ASSERT_OK_AND_ASSIGN(auto large_thunk,
+                       WhileThunk::Create({"while_large"}, pred_slice,
+                                          std::move(large_cond_sequence),
+                                          std::move(large_body_sequence)));
+  EXPECT_FALSE(large_thunk->body_executor().is_sequential());
+}
+
 }  // namespace
 }  // namespace xla::cpu

From 4ffb660634f26a4781b59f4417891892bbc166ab Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Tue, 13 Oct 2026 11:15:56 -0700
Subject: [PATCH 0704/1000] Run testHighRank eagerly too and test an empty last
 dimension

Eager softmax also reaches the oneDNN kernel, through the graph passes
of the function that wraps the op, so run the test in both modes. Also
test an empty last dimension, comparing it directly since NumPy can't
take the maximum along an empty dimension.
---
 .../python/kernel_tests/nn_ops/softmax_op_test.py | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py b/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py
index cf6d29da872..91fb22161c1 100644
--- a/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py
@@ -276,17 +276,24 @@ class SoftmaxTest(test.TestCase):
         use_gpu=False)
     self._testOverflow(use_gpu=False)
 
+  @test_util.run_in_graph_and_eager_modes(use_gpu=False)
   def testHighRank(self):
     # Regression test for GitHub issue 128814: the oneDNN kernel rejected
     # inputs of more than 5 dimensions.
     rng = np.random.RandomState(0)
     for rank in (6, 7):
-      for empty in (False, True):
+      for empty_dim in (None, 2, -1):
         shape = [2] * (rank - 1) + [3]
-        if empty:
-          shape[2] = 0
+        if empty_dim is not None:
+          shape[empty_dim] = 0
         features = rng.rand(*shape).astype(np.float32)
-        self._testSoftmax(features)
+        if empty_dim == -1:
+          # NumPy can't take the maximum along an empty dimension.
+          self.assertAllEqual(
+              np.zeros(shape, dtype=np.float32),
+              self.evaluate(nn_ops.softmax(features)))
+        else:
+          self._testSoftmax(features)
         self._testSoftmax(features, dim=1)
 
   def testShapeInference(self):

From e4c589c6c1cebe070f17131fa4a7d70098d05cc9 Mon Sep 17 00:00:00 2001
From: Tori Baker <vwbaker@google.com>
Date: Tue, 13 Oct 2026 10:39:39 -0700
Subject: [PATCH 0705/1000] Don't fuse epilogue transposes the user can absorb
 in GemmFusion V2.

As with prologue transposes, fusing an epilogue transpose that reorders a non-contracting dimension (M or N) without moving the minor-most dimension makes the output multi-strided, constraining Triton tiling and requiring a separate transpose kernel under cuBLAS. Leave such a transpose to its user when it has a single elementwise or concatenate user that can absorb it at no extra memory cost.

PiperOrigin-RevId: 998736698
---
 .../backends/gpu/transforms/gemm_fusion.cc    | 27 +++++++++-
 .../gpu/transforms/gemm_fusion_test.cc        | 51 +++++++++++++++++++
 2 files changed, 77 insertions(+), 1 deletion(-)

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 7ab4fc1a7ed..d28df71b899 100644
--- a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc
@@ -1650,7 +1650,22 @@ std::optional<TrackerInfo> ComputeCandidateTracker(
   return std::nullopt;
 }
 
+// Returns true if the only user of `transpose`, looking through bitcasts and
+// reshapes, can absorb it at no extra memory cost. `transpose` must be from
+// the original module so that all its users are visible.
+bool UserCanAbsorbTranspose(const HloInstruction& transpose) {
+  const HloInstruction* user = &transpose;
+  do {
+    if (user->user_count() != 1) {
+      return false;
+    }
+    user = user->users()[0];
+  } while (HloPredicateIsOp<HloOpcode::kBitcast, HloOpcode::kReshape>(user));
+  return user->IsElementwise() || user->opcode() == HloOpcode::kConcatenate;
+}
+
 FusionDecision ShouldFuseUserTranspose(const HloInstruction& transpose,
+                                       const HloInstruction& original_transpose,
                                        const HloInstruction& fusion,
                                        const ShapeTracker& tracker) {
   const HloInstruction* dot = hlo_query::FindInstruction(
@@ -1692,6 +1707,16 @@ FusionDecision ShouldFuseUserTranspose(const HloInstruction& transpose,
     }
   }
 
+  // As in ShouldFuseTranspose, leave a transpose that splits M or N without
+  // moving the minor dimension to its user if the user can absorb it.
+  if ((!transpose_tracker.MapsToOneStride(lhs_nc_dims) ||
+       !transpose_tracker.MapsToOneStride(rhs_nc_dims)) &&
+      !TransposesMinorDimension(&original_transpose) &&
+      UserCanAbsorbTranspose(original_transpose)) {
+    return FusionDecision::Forbid(
+        "Transpose within a non-contracting dimension can be fused into its "
+        "user instead.");
+  }
   return FusionDecision::Allow();
 }
 
@@ -1723,7 +1748,7 @@ FusionDecision ShouldFuseUser(const HloInstruction* user,
         return FusionDecision::Forbid(
             "No shape tracker found for transpose user.");
       }
-      return ShouldFuseUserTranspose(*user, *fusion, *tracker);
+      return ShouldFuseUserTranspose(*user, original_user, *fusion, *tracker);
     case HloOpcode::kConcatenate:
       return FusionDecision::Forbid("Not fusing concatenate into epilogue.");
     default:
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 65a8047fc74..4d3eb4bfe45 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
@@ -2719,6 +2719,57 @@ ENTRY e {
                                    GmockMatch(m::Parameter())));
 }
 
+TEST_F(GemmFusionTestV2, DoNotFuseEpilogueTransposeIfUserCanAbsorbIt) {
+  // The transpose only reorders the sub-dimensions of the LHS non-contracting
+  // dimension, so it is left for its user (a residual add) to absorb, even
+  // though the user could be fused into the GEMM otherwise. The bias is still
+  // fused.
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
+HloModule m
+
+ENTRY e {
+  p0 = f16[3072,64]{1,0} parameter(0)
+  p1 = f16[64,32]{1,0} parameter(1)
+  dot = f16[3072,32]{1,0} dot(p0, p1),
+    lhs_contracting_dims={1}, rhs_contracting_dims={0}
+  p2 = f16[32]{0} parameter(2)
+  bias = f16[3072,32]{1,0} broadcast(p2), dimensions={1}
+  add = f16[3072,32]{1,0} add(dot, bias)
+  bitcast = f16[3,16,64,32]{3,2,1,0} bitcast(add)
+  transpose = f16[16,64,3,32]{3,2,1,0} transpose(bitcast),
+    dimensions={1,2,0,3}
+  p3 = f16[16,64,3,32]{3,2,1,0} parameter(3)
+  ROOT sum = f16[16,64,3,32]{3,2,1,0} add(p3, transpose)
+})"));
+  ASSERT_THAT(GemmFusion(device_info_).Run(module.get()), IsOkAndHolds(true));
+  const HloInstruction* fusion = nullptr;
+  ASSERT_THAT(module->entry_computation()->root_instruction(),
+              GmockMatch(m::Add(m::Parameter(), TransposeOrTransposeBitcast(
+                                                    m::Fusion(&fusion)))));
+  EXPECT_THAT(fusion->fused_expression_root(), GmockMatch(m::Add()));
+}
+
+TEST_F(GemmFusionTestV2, FuseEpilogueTransposeIfRoot) {
+  // There is no user kernel that could absorb the transpose.
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
+HloModule m
+
+ENTRY e {
+  p0 = f16[3072,64]{1,0} parameter(0)
+  p1 = f16[64,32]{1,0} parameter(1)
+  dot = f16[3072,32]{1,0} dot(p0, p1),
+    lhs_contracting_dims={1}, rhs_contracting_dims={0}
+  bitcast = f16[3,16,64,32]{3,2,1,0} bitcast(dot)
+  ROOT transpose = f16[16,64,3,32]{3,2,1,0} transpose(bitcast),
+    dimensions={1,2,0,3}
+})"));
+  ASSERT_THAT(GemmFusion(device_info_).Run(module.get()), IsOkAndHolds(true));
+  const HloInstruction* fusion =
+      module->entry_computation()->root_instruction();
+  ASSERT_THAT(fusion, GmockMatch(m::Fusion()));
+  EXPECT_THAT(fusion->fused_expression_root(), GmockMatch(m::Transpose()));
+}
+
 TEST_F(GemmFusionTestV2, ConcatResetTrackerCrash) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m

From 8c53108d204c6eaaf4ad2331098b340d295c89bc Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Tue, 13 Oct 2026 11:48:47 -0700
Subject: [PATCH 0706/1000] Guard XLA Range sizes against overflow and
 out-of-order bounds

Compute the static size from |limit - start| and |delta| in uint64_t,
as the CPU kernel does, so that a range whose span doesn't fit in its
type, such as int64 from -2**63 to 2**63 - 1, gets its full size
instead of overflowing.

For dynamic bounds, compute the size in the unsigned type of the same
width, and give runtime values in the wrong order for delta's sign,
which the CPU kernel rejects, size 0. Their size was past the output's
static size, so the result was read from past the output's buffer.
Skip Abs for unsigned deltas, since XLA's Abs rejects unsigned types,
so that unsigned ranges with dynamic bounds compile.

Run testRangeDynamicBounds for int32, int64 and float32, with
out-of-order runtime bounds and a uint32 range, and test the int64
range in testRange.
---
 tensorflow/compiler/tests/ternary_ops_test.py | 35 ++++++++---
 .../compiler/tf2xla/kernels/sequence_ops.cc   | 62 ++++++++++++++-----
 2 files changed, 75 insertions(+), 22 deletions(-)

diff --git a/tensorflow/compiler/tests/ternary_ops_test.py b/tensorflow/compiler/tests/ternary_ops_test.py
index 1d97c18df6f..afe3b355839 100644
--- a/tensorflow/compiler/tests/ternary_ops_test.py
+++ b/tensorflow/compiler/tests/ternary_ops_test.py
@@ -76,15 +76,24 @@ class TernaryOpsTest(xla_test.XLATestCase, parameterized.TestCase):
           dtype(2),
           expected=np.array([1, 3, 5], dtype=dtype),
       )
+    # The distance from start to limit doesn't fit in int64.
+    self._testTernary(
+        math_ops.range,
+        np.int64(-2**63),
+        np.int64(2**63 - 1),
+        np.int64(2**62),
+        expected=np.array([-2**63, -2**62, 0, 2**62], dtype=np.int64),
+    )
 
-  def _rangeOfMaskedSize(self, args):
+  def _rangeOfMaskedSize(self, args, dtype=np.int32):
     # n is the number of true entries in mask, a size only known at run time,
     # and args maps it to the range's start, limit and delta.
     @def_function.function(jit_compile=True)
     def f(mask):
       indices = array_ops.reshape(array_ops.where(mask), [-1])
       n = array_ops.size(array_ops.gather(math_ops.range(6), indices))
-      return math_ops.range(*args(n))
+      return gen_math_ops._range(
+          *[math_ops.cast(a, dtype) for a in args(math_ops.cast(n, dtype))])
 
     with self.session() as session:
       with self.device_scope():
@@ -97,13 +106,23 @@ class TernaryOpsTest(xla_test.XLATestCase, parameterized.TestCase):
     # Regression test for GitHub issue 128950: with a negative delta, the
     # output was sized from the wrong bounds of start and limit, and with a
     # dynamic start, its values began at start's bound. Here n is 3, bound 6.
+    for dtype in (np.int32, np.int64, np.float32):
+      self.assertAllEqual(
+          self._rangeOfMaskedSize(lambda n: (n - 1, -1, -1), dtype), [2, 1, 0])
+      self.assertAllEqual(
+          self._rangeOfMaskedSize(lambda n: (6, n, -1), dtype), [6, 5, 4])
+      self.assertAllEqual(
+          self._rangeOfMaskedSize(lambda n: (n, 5, 1), dtype), [3, 4])
+      self.assertAllEqual(
+          self._rangeOfMaskedSize(lambda n: (n, -n, -2), dtype), [3, 1, -1])
+      # The CPU kernel rejects runtime values in the wrong order for delta's
+      # sign. XLA can't report that at run time, and returns an empty range.
+      self.assertAllEqual(
+          self._rangeOfMaskedSize(lambda n: (n, 1, 1), dtype), [])
+      self.assertAllEqual(
+          self._rangeOfMaskedSize(lambda n: (n, 4, -1), dtype), [])
     self.assertAllEqual(
-        self._rangeOfMaskedSize(lambda n: (n - 1, -1, -1)), [2, 1, 0])
-    self.assertAllEqual(
-        self._rangeOfMaskedSize(lambda n: (6, n, -1)), [6, 5, 4])
-    self.assertAllEqual(self._rangeOfMaskedSize(lambda n: (n, 5, 1)), [3, 4])
-    self.assertAllEqual(
-        self._rangeOfMaskedSize(lambda n: (n, -n, -2)), [3, 1, -1])
+        self._rangeOfMaskedSize(lambda n: (n, 10, 2), np.uint32), [3, 5, 7, 9])
 
   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 a9998e90e0a..d101a659a0a 100644
--- a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include <cmath>
 #include <cstdint>
 #include <cstdlib>
+#include <limits>
 #include <type_traits>
 
 #include "absl/status/statusor.h"
@@ -67,12 +68,21 @@ absl::StatusOr<xla::XlaOp> CreateRangeTensor(
   }
   int64_t size;
   if constexpr (std::is_integral_v<T>) {
-    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);
+    // Take |limit - start| and |delta| in uint64_t, which holds them for every
+    // T, as the CPU kernel does, so that a range wider than T can hold doesn't
+    // overflow.
+    const bool count_up = delta > static_cast<T>(0);
+    const uint64_t range =
+        count_up ? static_cast<uint64_t>(limit) - static_cast<uint64_t>(start)
+                 : static_cast<uint64_t>(start) - static_cast<uint64_t>(limit);
+    const uint64_t abs_delta = count_up ? static_cast<uint64_t>(delta)
+                                        : -static_cast<uint64_t>(delta);
+    const uint64_t size_unsigned = range == 0 ? 0 : (range - 1) / abs_delta + 1;
+    if (size_unsigned > std::numeric_limits<int64_t>::max()) {
+      return errors::InvalidArgument("Requires ((limit - start) / delta) <= ",
+                                     std::numeric_limits<int64_t>::max());
+    }
+    size = static_cast<int64_t>(size_unsigned);
   } else {
     double start_f = static_cast<double>(start);
     double limit_f = static_cast<double>(limit);
@@ -187,17 +197,41 @@ class RangeOp : public XlaOpKernel {
       xla::XlaOp delta = ctx->Input(2);
       xla::XlaOp limit = ctx->Input(1);
       xla::XlaOp start = ctx->Input(0);
+      // The runtime values can be in the wrong order for delta's sign, which
+      // the CPU kernel rejects. Give such a range size 0 rather than a size
+      // past the output's static size.
+      const xla::XlaOp is_empty =
+          count_down ? xla::Le(start, limit) : xla::Ge(start, limit);
+      xla::XlaOp dynamic_size;
       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);
-        dynamic_size = xla::ConvertElementType(dynamic_size, xla::S32);
-        output = xla::SetDimensionSize(output.value(), dynamic_size, 0);
+        // Take |limit - start| and |delta| in the unsigned type of the same
+        // width, which holds them. XLA's Abs doesn't take unsigned types, whose
+        // delta is positive anyway.
+        const xla::PrimitiveType xla_type = ctx->input_xla_type(0);
+        const xla::PrimitiveType unsigned_type =
+            xla::primitive_util::UnsignedIntegralTypeForBitWidth(
+                xla::primitive_util::BitWidth(xla_type));
+        const auto as_unsigned = [&](xla::XlaOp x) {
+          return xla::BitcastConvertType(x, unsigned_type);
+        };
+        const xla::XlaOp range = count_down
+                                     ? as_unsigned(start) - as_unsigned(limit)
+                                     : as_unsigned(limit) - as_unsigned(start);
+        const xla::XlaOp abs_delta =
+            as_unsigned(xla::primitive_util::IsSignedIntegralType(xla_type)
+                            ? xla::Abs(delta)
+                            : delta);
+        const xla::XlaOp one = xla::One(ctx->builder(), unsigned_type);
+        dynamic_size =
+            xla::Select(is_empty, xla::Zero(ctx->builder(), unsigned_type),
+                        (range - one) / abs_delta + one);
       } else {
-        auto dynamic_size = (xla::Ceil(xla::Abs((limit - start) / delta)));
-        dynamic_size = xla::ConvertElementType(dynamic_size, xla::S32);
-        output = xla::SetDimensionSize(output.value(), dynamic_size, 0);
+        dynamic_size = xla::Select(
+            is_empty, xla::Zero(ctx->builder(), ctx->input_xla_type(0)),
+            xla::Ceil(xla::Abs((limit - start) / delta)));
       }
+      dynamic_size = xla::ConvertElementType(dynamic_size, xla::S32);
+      output = xla::SetDimensionSize(output.value(), dynamic_size, 0);
     }
 
     ctx->SetOutput(0, output.value());

From 7996ed6c15ce7b0e3fd58dc20fa3d9f54e64cc40 Mon Sep 17 00:00:00 2001
From: Ce Zheng CA <zce@google.com>
Date: Tue, 13 Oct 2026 11:12:32 -0700
Subject: [PATCH 0707/1000] [XLA:CodeGen] Remove unused ROCDL/NVVM dialect
 dependencies from LowerTensorsPass and CPU fusion emitter

`LowerTensorsPass` runs on both CPU (`target_type="cpu"`, via `FusionCompiler` in `xla/backends/cpu/codegen/emitters/cpu_fusion_emitter.cc`) and GPU, yet declared `mlir::ROCDL::ROCDLDialect` in `dependentDialects` (`xla/codegen/emitters/transforms/passes.td`) and called `getOrLoadDialect<mlir::ROCDL::ROCDLDialect>()` in `emitAMDAtomicFAdd()` (`xla/codegen/emitters/transforms/lower_tensors.cc`) only to set the `"rocdl.no_remote_memory"` and `"rocdl.no_fine_grained_memory"` discardable attributes.

Set those `"rocdl.*"` discardable attributes directly via `op->setDiscardableAttr(mlir::ROCDL::ROCDLDialect::*AttrHelper::getNameStr(), unitAttr)` (`MLIRContext` automatically links `"rocdl.*"` `StringAttr`s to `ROCDLDialect` when `ROCDLDialect` is loaded on AMDGPU targets via `MlirKernelEmitter` and `LowerToLLVMGPUPass`). Also remove unused `NVVMDialect`, `ROCDLDialect`, and `DialectUtils` headers and `BUILD` dependencies from `xla/codegen/emitters/transforms/passes.h`, `xla/codegen/emitters/transforms/BUILD`, `xla/backends/cpu/codegen/emitters/cpu_fusion_emitter.cc`, and `xla/backends/cpu/codegen/emitters/BUILD`.

**Performance & Memory Impact:**
- **CPU Runtime & Memory:** Slightly positive on CPU/TPU compilations, as `LowerTensorsPass` no longer loads and initializes `ROCDLDialect` (`ROCDLOps`) into every CPU `MLIRContext`. Unchanged on AMDGPU, where `ROCDLDialect` is already registered and loaded by `MlirKernelEmitter` and `LowerToLLVMGPUPass`.

PiperOrigin-RevId: 998762325
---
 .../xla/xla/backends/cpu/codegen/emitters/BUILD   |  2 --
 .../cpu/codegen/emitters/cpu_fusion_emitter.cc    |  2 --
 .../xla/xla/codegen/emitters/transforms/BUILD     |  6 ------
 .../codegen/emitters/transforms/lower_tensors.cc  | 15 +++++++--------
 .../xla/xla/codegen/emitters/transforms/passes.h  |  2 --
 .../xla/xla/codegen/emitters/transforms/passes.td |  1 -
 6 files changed, 7 insertions(+), 21 deletions(-)

diff --git a/third_party/xla/xla/backends/cpu/codegen/emitters/BUILD b/third_party/xla/xla/backends/cpu/codegen/emitters/BUILD
index 0999c8517e5..4aed914c45c 100644
--- a/third_party/xla/xla/backends/cpu/codegen/emitters/BUILD
+++ b/third_party/xla/xla/backends/cpu/codegen/emitters/BUILD
@@ -102,9 +102,7 @@ cc_library(
         "@llvm-project//mlir:LLVMDialect",
         "@llvm-project//mlir:LLVMToLLVMIRTranslation",
         "@llvm-project//mlir:MemRefTransforms",
-        "@llvm-project//mlir:NVVMToLLVMIRTranslation",
         "@llvm-project//mlir:Pass",
-        "@llvm-project//mlir:ROCDLToLLVMIRTranslation",
         "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TensorDialect",
diff --git a/third_party/xla/xla/backends/cpu/codegen/emitters/cpu_fusion_emitter.cc b/third_party/xla/xla/backends/cpu/codegen/emitters/cpu_fusion_emitter.cc
index ab7414bd433..bc42a8844c2 100644
--- a/third_party/xla/xla/backends/cpu/codegen/emitters/cpu_fusion_emitter.cc
+++ b/third_party/xla/xla/backends/cpu/codegen/emitters/cpu_fusion_emitter.cc
@@ -53,8 +53,6 @@ limitations under the License.
 #include "mlir/Support/LLVM.h"
 #include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"
-#include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"
-#include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h"
 #include "mlir/Target/LLVMIR/Export.h"
 #include "xla/backends/cpu/alignment.h"
 #include "xla/backends/cpu/codegen/kernel_api_ir_builder.h"
diff --git a/third_party/xla/xla/codegen/emitters/transforms/BUILD b/third_party/xla/xla/codegen/emitters/transforms/BUILD
index 683f8a3b84e..56f7e808d29 100644
--- a/third_party/xla/xla/codegen/emitters/transforms/BUILD
+++ b/third_party/xla/xla/codegen/emitters/transforms/BUILD
@@ -285,9 +285,7 @@ cc_library(
         "//xla/codegen/emitters/ir:xla",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
-        "@llvm-project//mlir:NVVMDialect",
         "@llvm-project//mlir:Pass",
-        "@llvm-project//mlir:ROCDLDialect",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TransformUtils",
     ],
@@ -313,9 +311,7 @@ cc_library(
         "@llvm-project//mlir:DataLayoutInterfaces",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
-        "@llvm-project//mlir:NVVMDialect",
         "@llvm-project//mlir:Pass",
-        "@llvm-project//mlir:ROCDLDialect",
         "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TransformUtils",
@@ -339,9 +335,7 @@ cc_library(
         "@llvm-project//mlir:GPUDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:MathDialect",
-        "@llvm-project//mlir:NVVMDialect",
         "@llvm-project//mlir:Pass",
-        "@llvm-project//mlir:ROCDLDialect",
         "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TransformUtils",
diff --git a/third_party/xla/xla/codegen/emitters/transforms/lower_tensors.cc b/third_party/xla/xla/codegen/emitters/transforms/lower_tensors.cc
index 0bc4fed72cd..3ba72fb9408 100644
--- a/third_party/xla/xla/codegen/emitters/transforms/lower_tensors.cc
+++ b/third_party/xla/xla/codegen/emitters/transforms/lower_tensors.cc
@@ -37,7 +37,7 @@ limitations under the License.
 #include "mlir/Dialect/LLVMIR/LLVMAttrs.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
-#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
+#include "mlir/Dialect/LLVMIR/ROCDLDialectDecl.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
@@ -1130,13 +1130,12 @@ class RewriteAtomicRMW : public OpRewritePattern<AtomicRMWOp> {
     op.setIgnoreDenormalMode(true);
 
     auto unitAttr = b.getUnitAttr();
-    auto* rocdl =
-        op->getContext()->getOrLoadDialect<mlir::ROCDL::ROCDLDialect>();
-    auto noRemoteMemHelper = rocdl->getNoRemoteMemoryAttrHelper();
-    auto noFineMemHelper = rocdl->getNoFineGrainedMemoryAttrHelper();
-
-    noRemoteMemHelper.setAttr(op, unitAttr);
-    noFineMemHelper.setAttr(op, unitAttr);
+    op->setDiscardableAttr(
+        mlir::ROCDL::ROCDLDialect::NoRemoteMemoryAttrHelper::getNameStr(),
+        unitAttr);
+    op->setDiscardableAttr(
+        mlir::ROCDL::ROCDLDialect::NoFineGrainedMemoryAttrHelper::getNameStr(),
+        unitAttr);
 
     return success();
   }
diff --git a/third_party/xla/xla/codegen/emitters/transforms/passes.h b/third_party/xla/xla/codegen/emitters/transforms/passes.h
index 7d396204e0c..5c8f11643d5 100644
--- a/third_party/xla/xla/codegen/emitters/transforms/passes.h
+++ b/third_party/xla/xla/codegen/emitters/transforms/passes.h
@@ -19,8 +19,6 @@ limitations under the License.
 #include <memory>
 #include <string>
 
-#include "mlir/Dialect/LLVMIR/NVVMDialect.h"  // IWYU pragma: keep
-#include "mlir/Dialect/LLVMIR/ROCDLDialect.h"  // IWYU pragma: keep
 #include "mlir/Pass/Pass.h"
 
 namespace stream_executor {
diff --git a/third_party/xla/xla/codegen/emitters/transforms/passes.td b/third_party/xla/xla/codegen/emitters/transforms/passes.td
index 9743cdc91ee..f3615481118 100644
--- a/third_party/xla/xla/codegen/emitters/transforms/passes.td
+++ b/third_party/xla/xla/codegen/emitters/transforms/passes.td
@@ -118,7 +118,6 @@ def LowerTensorsPass : Pass<"xla-lower-tensors", "mlir::ModuleOp"> {
     "xla::gpu::XlaGpuDialect",
     "xla::XlaDialect",
     "mlir::vector::VectorDialect",
-    "mlir::ROCDL::ROCDLDialect",
   ];
   let options = [
     Option<"gpu_device_info_", "gpu_device_info", "std::string", /*default=*/"",

From ff46c8fa9148563a5ff4d95ed6ac0e5f15e1b5e9 Mon Sep 17 00:00:00 2001
From: Alex Pivovarov <upwind@google.com>
Date: Tue, 13 Oct 2026 11:28:31 -0700
Subject: [PATCH 0708/1000] [XLA:GPU] Support stable PARTIAL order in
 TopkSpecializer U64 RAFT adapter.

#### Why:
On GPU, `approx_top_k` lowers to a stable `PARTIAL`-order sort + slice (`backend_config={is_stable = true, order = "PARTIAL"}`), where `+0.0` and `-0.0` compare equal (`!(+0.0 > -0.0)` and `!(-0.0 > +0.0)`) and ties between `+0.0` and `-0.0` must be broken strictly by their original index (`iota`). Previously, `ShouldRewriteStableTopKToUint64` only supported `ComparisonOrder::kTotal` (where `+0.0 > -0.0`), causing stable `PARTIAL`-order `TopK` operations to skip RAFT `select_k` and fall back to full `sort + slice`.

#### Changes:
- Add `BuildPackF32ToU64PartialOrder` and `BuildUnpackU64ToF32PartialOrder` in `xla/backends/gpu/transforms/topk_specializer.cc` while keeping `BuildPackF32ToU64TotalOrder` and `BuildUnpackU64ToF32TotalOrder` unchanged for `TOTAL` order:
  - **Bits `[63..32]` (`partial_radix_key`)**: Canonicalize `-0.0f`'s radix key (`0x7FFFFFFF`) to `+0.0f`'s radix key (`0x80000000`) by adding `is_neg_zero` (`1` for `-0.0f`, `0` otherwise), so `+0.0f` and `-0.0f` tie on the upper 32 bits.
  - **Bits `[31..1]` (`(0x7FFFFFFF - iota) << 1`)**: Store the 31-bit reversed `s32` index to break ties in ascending index order.
  - **Bit `[0]` (`is_neg_zero`)**: Preserve whether a zero was originally `-0.0f` without affecting element ordering (since bits `[31..1]` are strictly unique per index). During unpacking, subtract `is_neg_zero` from `partial_radix_key` before unflipping so `-0.0f` and `+0.0f` are restored bit-identically.
- Update `ShouldRewriteStableTopKToUint64` and `RewriteStableTopKToUint64` to support `ComparisonOrder::kPartial` in addition to `ComparisonOrder::kTotal`.
- Add unit and execution tests in `xla/backends/gpu/transforms/topk_specializer_test.cc` (`RewriteStablePartialOrderTopKF32ToUint64`, `RewriteStablePartialOrderTopKBF16ToUint64`, `SkipsSmallPartialOrderTopK`, and `RewriteStablePartialOrderTopKProducesCorrectResult`).

PiperOrigin-RevId: 998774308
---
 .../gpu/transforms/topk_specializer.cc        | 177 ++++++++++++-
 .../gpu/transforms/topk_specializer_test.cc   | 239 +++++++++++++++++-
 2 files changed, 403 insertions(+), 13 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/transforms/topk_specializer.cc b/third_party/xla/xla/backends/gpu/transforms/topk_specializer.cc
index 377bc32f768..74894ce0ef6 100644
--- a/third_party/xla/xla/backends/gpu/transforms/topk_specializer.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/topk_specializer.cc
@@ -73,8 +73,10 @@ HloInstruction* BroadcastU64(HloComputation* comp, const Shape& target_shape,
       ShapeUtil::ChangeElementType(target_shape, U64), constant, {}));
 }
 
-// Emits HLO to pack F32 values and Iota indices into U64 priority keys.
-HloInstruction* BuildPackF32ToU64(HloInstruction* data, HloComputation* comp) {
+// Emits HLO to pack F32 values and Iota indices into U64 priority keys for
+// TOTAL order (+0.0 > -0.0).
+HloInstruction* BuildPackF32ToU64TotalOrder(HloInstruction* data,
+                                            HloComputation* comp) {
   const Shape& shape = data->shape();
   int64_t iota_dim = shape.dimensions().size() - 1;
 
@@ -131,9 +133,9 @@ HloInstruction* BuildPackF32ToU64(HloInstruction* data, HloComputation* comp) {
       u64_shape, HloOpcode::kOr, val_u64_top, iota_neg_u64));
 }
 
-// Emits HLO to unpack U64 priority keys back to F32 values.
-HloInstruction* BuildUnpackU64ToF32(HloInstruction* u64_values,
-                                    HloComputation* comp) {
+// Emits HLO to unpack TOTAL-order U64 priority keys back to F32 values.
+HloInstruction* BuildUnpackU64ToF32TotalOrder(HloInstruction* u64_values,
+                                              HloComputation* comp) {
   const Shape& shape = u64_values->shape();
   Shape u32_shape = ShapeUtil::ChangeElementType(shape, U32);
   Shape f32_shape = ShapeUtil::ChangeElementType(shape, F32);
@@ -176,10 +178,156 @@ HloInstruction* BuildUnpackU64ToF32(HloInstruction* u64_values,
       HloInstruction::CreateBitcastConvert(f32_shape, original_u32));
 }
 
+// Emits HLO to pack F32 values and Iota indices into U64 priority keys for
+// PARTIAL order (+0.0 == -0.0).
+//
+// Key layout:
+//   [63..32]: partial_radix_key (where -0.0f's radix key 0x7FFFFFFF is
+//             canonicalized to +0.0f's radix key 0x80000000 so that +0.0f and
+//             -0.0f tie on the upper 32 bits)
+//   [31..1]:  0x7FFFFFFF - iota (31-bit reversed index, breaking ties in
+//             ascending index order)
+//   [0]:      is_neg_zero (1 if original value was -0.0f, 0 otherwise; never
+//             affects ordering since bits [31..1] are strictly unique per
+//             index, and allows unpacking to restore -0.0f bit-identically)
+HloInstruction* BuildPackF32ToU64PartialOrder(HloInstruction* data,
+                                              HloComputation* comp) {
+  const Shape& shape = data->shape();
+  int64_t iota_dim = shape.dimensions().size() - 1;
+
+  Shape pred_shape = ShapeUtil::ChangeElementType(shape, PRED);
+  Shape u32_shape = ShapeUtil::ChangeElementType(shape, U32);
+  Shape s32_shape = ShapeUtil::ChangeElementType(shape, S32);
+  Shape u64_shape = ShapeUtil::ChangeElementType(shape, U64);
+
+  // 1. Generate 31-bit reversed indices shifted left by 1:
+  //    (0x7FFFFFFF - iota) << 1
+  HloInstruction* iota_u32 =
+      comp->AddInstruction(HloInstruction::CreateIota(u32_shape, iota_dim));
+  HloInstruction* broadcast_7f = BroadcastU32(comp, shape, 0x7FFFFFFF);
+  HloInstruction* iota_neg = comp->AddInstruction(HloInstruction::CreateBinary(
+      u32_shape, HloOpcode::kSubtract, broadcast_7f, iota_u32));
+  HloInstruction* broadcast_1 = BroadcastU32(comp, shape, 1);
+  HloInstruction* iota_shifted =
+      comp->AddInstruction(HloInstruction::CreateBinary(
+          u32_shape, HloOpcode::kShiftLeft, iota_neg, broadcast_1));
+
+  // 2. Pure Bitwise Radix Float Flip (F32 -> U32)
+  // s32_val = bitcast(data)
+  HloInstruction* s32_val = comp->AddInstruction(
+      HloInstruction::CreateBitcastConvert(s32_shape, data));
+
+  // sign_smeared = s32_val >> 31 (arithmetic shift)
+  HloInstruction* const_31_s32 = comp->AddInstruction(
+      HloInstruction::CreateConstant(LiteralUtil::CreateR0<int32_t>(31)));
+  HloInstruction* broadcast_31_s32 = comp->AddInstruction(
+      HloInstruction::CreateBroadcast(s32_shape, const_31_s32, {}));
+  HloInstruction* sign_smeared = comp->AddInstruction(
+      HloInstruction::CreateBinary(s32_shape, HloOpcode::kShiftRightArithmetic,
+                                   s32_val, broadcast_31_s32));
+
+  // mask = bitcast(sign_smeared) | 0x80000000
+  HloInstruction* sign_smeared_u32 = comp->AddInstruction(
+      HloInstruction::CreateBitcastConvert(u32_shape, sign_smeared));
+  HloInstruction* broadcast_8 = BroadcastU32(comp, shape, 0x80000000);
+  HloInstruction* mask = comp->AddInstruction(HloInstruction::CreateBinary(
+      u32_shape, HloOpcode::kOr, sign_smeared_u32, broadcast_8));
+
+  // radix_key = bitcast(data) ^ mask
+  HloInstruction* u32_val = comp->AddInstruction(
+      HloInstruction::CreateBitcastConvert(u32_shape, data));
+  HloInstruction* radix_key = comp->AddInstruction(
+      HloInstruction::CreateBinary(u32_shape, HloOpcode::kXor, u32_val, mask));
+
+  // 3. Canonicalize -0.0f (u32_val == 0x80000000, radix_key == 0x7FFFFFFF) to
+  //    +0.0f's radix_key (0x80000000) by adding is_neg_zero (1 for -0.0f, 0
+  //    otherwise), and save is_neg_zero in bit 0 of the lower 32 bits.
+  HloInstruction* is_neg_zero_pred =
+      comp->AddInstruction(HloInstruction::CreateCompare(
+          pred_shape, u32_val, broadcast_8, ComparisonDirection::kEq));
+  HloInstruction* is_neg_zero = comp->AddInstruction(
+      HloInstruction::CreateConvert(u32_shape, is_neg_zero_pred));
+  HloInstruction* partial_radix_key =
+      comp->AddInstruction(HloInstruction::CreateBinary(
+          u32_shape, HloOpcode::kAdd, radix_key, is_neg_zero));
+  HloInstruction* low_u32 = comp->AddInstruction(HloInstruction::CreateBinary(
+      u32_shape, HloOpcode::kOr, iota_shifted, is_neg_zero));
+
+  // 4. Pack into U64: (partial_radix_key << 32) | low_u32
+  HloInstruction* val_u64 = comp->AddInstruction(
+      HloInstruction::CreateConvert(u64_shape, partial_radix_key));
+  HloInstruction* broadcast_32_u64 = BroadcastU64(comp, shape, 32);
+  HloInstruction* val_u64_top =
+      comp->AddInstruction(HloInstruction::CreateBinary(
+          u64_shape, HloOpcode::kShiftLeft, val_u64, broadcast_32_u64));
+
+  HloInstruction* low_u64 =
+      comp->AddInstruction(HloInstruction::CreateConvert(u64_shape, low_u32));
+
+  return comp->AddInstruction(HloInstruction::CreateBinary(
+      u64_shape, HloOpcode::kOr, val_u64_top, low_u64));
+}
+
+// Emits HLO to unpack PARTIAL-order U64 priority keys back to F32 values.
+HloInstruction* BuildUnpackU64ToF32PartialOrder(HloInstruction* u64_values,
+                                                HloComputation* comp) {
+  const Shape& shape = u64_values->shape();
+  Shape u32_shape = ShapeUtil::ChangeElementType(shape, U32);
+  Shape f32_shape = ShapeUtil::ChangeElementType(shape, F32);
+
+  // 1. Shift right by 32 and cast to U32 to extract partial_radix_key, and
+  //    subtract bit 0 (is_neg_zero) to restore the original radix_key
+  //    (0x80000000 - 1 = 0x7FFFFFFF for -0.0f).
+  HloInstruction* broadcast_32_u64 = BroadcastU64(comp, shape, 32);
+  HloInstruction* rshift = comp->AddInstruction(HloInstruction::CreateBinary(
+      shape, HloOpcode::kShiftRightLogical, u64_values, broadcast_32_u64));
+  HloInstruction* partial_radix_key =
+      comp->AddInstruction(HloInstruction::CreateConvert(u32_shape, rshift));
+
+  HloInstruction* low_u32 = comp->AddInstruction(
+      HloInstruction::CreateConvert(u32_shape, u64_values));
+  HloInstruction* broadcast_1 = BroadcastU32(comp, shape, 1);
+  HloInstruction* is_neg_zero =
+      comp->AddInstruction(HloInstruction::CreateBinary(
+          u32_shape, HloOpcode::kAnd, low_u32, broadcast_1));
+  HloInstruction* radix_key = comp->AddInstruction(HloInstruction::CreateBinary(
+      u32_shape, HloOpcode::kSubtract, partial_radix_key, is_neg_zero));
+
+  // 2. Isolate the MSB: msb = radix_key >> 31 (Logical shift)
+  HloInstruction* const_31_u32 = comp->AddInstruction(
+      HloInstruction::CreateConstant(LiteralUtil::CreateR0<uint32_t>(31)));
+  HloInstruction* broadcast_31_u32 = comp->AddInstruction(
+      HloInstruction::CreateBroadcast(u32_shape, const_31_u32, {}));
+  HloInstruction* msb = comp->AddInstruction(HloInstruction::CreateBinary(
+      u32_shape, HloOpcode::kShiftRightLogical, radix_key, broadcast_31_u32));
+
+  // 3. Reconstruct mask: msb_minus_one = msb - 1
+  // If MSB was 1: 1 - 1 = 0x00000000.
+  // If MSB was 0: 0 - 1 = 0xFFFFFFFF (unsigned underflow).
+  HloInstruction* msb_minus_one =
+      comp->AddInstruction(HloInstruction::CreateBinary(
+          u32_shape, HloOpcode::kSubtract, msb, broadcast_1));
+
+  // 4. unmask = msb_minus_one | 0x80000000
+  HloInstruction* broadcast_8 = BroadcastU32(comp, shape, 0x80000000);
+  HloInstruction* unmask = comp->AddInstruction(HloInstruction::CreateBinary(
+      u32_shape, HloOpcode::kOr, msb_minus_one, broadcast_8));
+
+  // 5. Unflip the bits: original_u32 = radix_key ^ unmask
+  HloInstruction* original_u32 =
+      comp->AddInstruction(HloInstruction::CreateBinary(
+          u32_shape, HloOpcode::kXor, radix_key, unmask));
+
+  // 6. Bitcast back to F32
+  return comp->AddInstruction(
+      HloInstruction::CreateBitcastConvert(f32_shape, original_u32));
+}
+
 // Checks if we can safely route stable TopK to RAFT using the Uint64 adapter.
 bool ShouldRewriteStableTopKToUint64(HloCustomCallInstruction* topk) {
-  // The Uint64 bit-packing adapter implements TOTAL ordering.
-  if (hlo_instruction_utils::GetTopKOrder(topk) != ComparisonOrder::kTotal) {
+  // The Uint64 bit-packing adapter supports TOTAL and PARTIAL ordering.
+  ComparisonOrder order = hlo_instruction_utils::GetTopKOrder(topk);
+  if (order != ComparisonOrder::kTotal && order != ComparisonOrder::kPartial) {
     return false;
   }
   if (!hlo_instruction_utils::IsTopKStable(topk)) {
@@ -202,6 +350,11 @@ bool ShouldRewriteStableTopKToUint64(HloCustomCallInstruction* topk) {
     return false;
   }
 
+  // The PARTIAL-order U64 adapter packs the reversed index into 31 bits.
+  if (order == ComparisonOrder::kPartial && n > (size_t{1} << 31)) {
+    return false;
+  }
+
   double ratio = static_cast<double>(k) / n;
   if (ratio >= 0.85) {
     return false;
@@ -248,7 +401,10 @@ absl::StatusOr<HloInstruction*> RewriteStableTopKToUint64(
     data = comp->AddInstruction(HloInstruction::CreateConvert(
         ShapeUtil::ChangeElementType(data->shape(), F32), data));
   }
-  HloInstruction* packed_u64 = BuildPackF32ToU64(data, comp);
+  ComparisonOrder order = hlo_instruction_utils::GetTopKOrder(topk);
+  HloInstruction* packed_u64 = order == ComparisonOrder::kPartial
+                                   ? BuildPackF32ToU64PartialOrder(data, comp)
+                                   : BuildPackF32ToU64TotalOrder(data, comp);
 
   // 2. Create the specialized __gpu$TopK custom call
   Shape k_shape =
@@ -277,7 +433,10 @@ absl::StatusOr<HloInstruction*> RewriteStableTopKToUint64(
       HloInstruction::CreateGetTupleElement(idx_shape, new_topk, 1));
 
   // Unpack to F32.
-  HloInstruction* unpacked_data = BuildUnpackU64ToF32(u64_vals, comp);
+  HloInstruction* unpacked_data =
+      order == ComparisonOrder::kPartial
+          ? BuildUnpackU64ToF32PartialOrder(u64_vals, comp)
+          : BuildUnpackU64ToF32TotalOrder(u64_vals, comp);
 
   // If the original input was BF16, downcast the F32 result back to BF16.
   if (original_type == BF16) {
diff --git a/third_party/xla/xla/backends/gpu/transforms/topk_specializer_test.cc b/third_party/xla/xla/backends/gpu/transforms/topk_specializer_test.cc
index 17e5bd5f00c..bf0e0277ac8 100644
--- a/third_party/xla/xla/backends/gpu/transforms/topk_specializer_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/topk_specializer_test.cc
@@ -355,7 +355,7 @@ TEST_F(TopkTest, RewriteStableTopKDisabledByDefault) {
   EXPECT_FALSE(changed);
 }
 
-TEST_F(TopkTest, SkipsPartialOrderTopK) {
+TEST_F(TopkTest, RewriteStablePartialOrderTopKF32ToUint64) {
   const char* hlo = R"(
     HloModule m
 
@@ -369,9 +369,183 @@ TEST_F(TopkTest, SkipsPartialOrderTopK) {
 
     ENTRY top_k {
       arg = f32[8,1024] parameter(0)
-      small_k = (f32[8,8], s32[8,8]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = true, order = "PARTIAL"}
-      raft_k = (f32[8,32], s32[8,32]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = true, order = "PARTIAL"}
-      ROOT result = tuple(small_k, raft_k)
+      ROOT result = (f32[8,32], s32[8,32]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = true, order = "PARTIAL"}
+    }
+  )";
+
+  HloModuleConfig config = GetModuleConfigForTest();
+  config.mutable_debug_options()
+      .set_xla_gpu_experimental_enable_raft_for_stable_topk(true);
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(hlo, config));
+
+  if (!device_description().gpu_compute_capability().IsCuda()) {
+    GTEST_SKIP() << "RAFT is CUDA-only.";
+  }
+
+  ASSERT_OK_AND_ASSIGN(
+      bool changed,
+      TopkSpecializer(device_description().gpu_compute_capability())
+          .Run(module.get()));
+  ASSERT_TRUE(changed);
+
+  const char* check_pattern = R"(
+// CHECK-LABEL: ENTRY %top_k
+// CHECK: %[[ARG:[^ ]+]] = f32[8,1024]{{.*}} parameter(0)
+
+// 1. Bitcast to U32 and S32
+// CHECK: %[[VAL_U32:[^ ]+]] = u32[8,1024]{{.*}} bitcast-convert(%[[ARG]])
+// CHECK: %[[VAL_S32:[^ ]+]] = s32[8,1024]{{.*}} bitcast-convert(%[[ARG]])
+
+// 2. Pure Bitwise Radix Flip (F32 -> U32)
+// CHECK: %[[SRA:[^ ]+]] = s32[8,1024]{{.*}} shift-right-arithmetic(%[[VAL_S32]], {{.*}})
+// CHECK: %[[SRA_U32:[^ ]+]] = u32[8,1024]{{.*}} bitcast-convert(%[[SRA]])
+// CHECK: %[[MASK:[^ ]+]] = u32[8,1024]{{.*}} or(%[[SRA_U32]], {{.*}})
+// CHECK: %[[RADIX_KEY:[^ ]+]] = u32[8,1024]{{.*}} xor(%[[VAL_U32]], %[[MASK]])
+
+// 3. Canonicalize -0.0f and pack into U64 with 31-bit reversed iota and is_neg_zero bit
+// CHECK: %[[IS_NEG_ZERO_PRED:[^ ]+]] = pred[8,1024]{{.*}} compare(%[[VAL_U32]], {{.*}}), direction=EQ
+// CHECK: %[[IS_NEG_ZERO:[^ ]+]] = u32[8,1024]{{.*}} convert(%[[IS_NEG_ZERO_PRED]])
+// CHECK: %[[PARTIAL_KEY:[^ ]+]] = u32[8,1024]{{.*}} add(%[[RADIX_KEY]], %[[IS_NEG_ZERO]])
+// CHECK: %[[KEY_U64:[^ ]+]] = u64[8,1024]{{.*}} convert(%[[PARTIAL_KEY]])
+// CHECK: %[[SHIFT_LEFT:[^ ]+]] = u64[8,1024]{{.*}} shift-left(%[[KEY_U64]], {{.*}})
+// CHECK: %[[IOTA:[^ ]+]] = u32[8,1024]{{.*}} iota(), iota_dimension=1
+// CHECK: %[[IOTA_NEG:[^ ]+]] = u32[8,1024]{{.*}} subtract({{.*}}, %[[IOTA]])
+// CHECK: %[[IOTA_SHIFTED:[^ ]+]] = u32[8,1024]{{.*}} shift-left(%[[IOTA_NEG]], {{.*}})
+// CHECK: %[[LOW_U32:[^ ]+]] = u32[8,1024]{{.*}} or(%[[IOTA_SHIFTED]], %[[IS_NEG_ZERO]])
+// CHECK: %[[LOW_U64:[^ ]+]] = u64[8,1024]{{.*}} convert(%[[LOW_U32]])
+// CHECK: %[[PACKED:[^ ]+]] = u64[8,1024]{{.*}} or(%[[SHIFT_LEFT]], %[[LOW_U64]])
+
+// 4. CustomCall (__gpu$TopK)
+// CHECK: %[[CUSTOM_CALL:[^ ]+]] = (u64[8,32]{{.*}}, s32[8,32]{{.*}}, u8[33554432]{{.*}}) custom-call(%[[PACKED]]), custom_call_target="__gpu$TopK", api_version=API_VERSION_TYPED_FFI, backend_config={is_stable = false, order = "PARTIAL"}
+
+// 5. Unpack U64 -> U32 and restore -0.0f radix key
+// CHECK: %[[SRL:[^ ]+]] = u64[8,32]{{.*}} shift-right-logical(%[[CUSTOM_CALL]]#0, {{.*}})
+// CHECK: %[[PARTIAL_UNPACK_U32:[^ ]+]] = u32[8,32]{{.*}} convert(%[[SRL]])
+// CHECK: %[[UNPACK_LOW_U32:[^ ]+]] = u32[8,32]{{.*}} convert(%[[CUSTOM_CALL]]#0)
+// CHECK: %[[UNPACK_IS_NEG_ZERO:[^ ]+]] = u32[8,32]{{.*}} and(%[[UNPACK_LOW_U32]], {{.*}})
+// CHECK: %[[UNPACK_U32:[^ ]+]] = u32[8,32]{{.*}} subtract(%[[PARTIAL_UNPACK_U32]], %[[UNPACK_IS_NEG_ZERO]])
+
+// 6. Pure Bitwise Radix Unflip (U32 -> F32)
+// CHECK: %[[MSB:[^ ]+]] = u32[8,32]{{.*}} shift-right-logical(%[[UNPACK_U32]], {{.*}})
+// CHECK: %[[MSB_MINUS_ONE:[^ ]+]] = u32[8,32]{{.*}} subtract(%[[MSB]], {{.*}})
+// CHECK: %[[UNMASK:[^ ]+]] = u32[8,32]{{.*}} or(%[[MSB_MINUS_ONE]], {{.*}})
+// CHECK: %[[UNFLIPPED:[^ ]+]] = u32[8,32]{{.*}} xor(%[[UNPACK_U32]], %[[UNMASK]])
+// CHECK: %[[F32_OUT:[^ ]+]] = f32[8,32]{{.*}} bitcast-convert(%[[UNFLIPPED]])
+
+// 7. Tuple output
+// CHECK: ROOT {{.*}} tuple({{.*}}%[[F32_OUT]], {{.*}}%[[CUSTOM_CALL]]#1)
+  )";
+
+  ASSERT_OK_AND_ASSIGN(bool filecheck_matched,
+                       RunFileCheck(module->ToString(), check_pattern));
+  EXPECT_TRUE(filecheck_matched);
+}
+
+TEST_F(TopkTest, RewriteStablePartialOrderTopKBF16ToUint64) {
+  const char* hlo = R"(
+    HloModule m
+
+    %compare-gt.1 {
+      p.1.lhs = s32[] parameter(2)
+      p.1.rhs = s32[] parameter(3)
+      p.0.lhs = bf16[] parameter(0)
+      p.0.rhs = bf16[] parameter(1)
+      ROOT compare = pred[] compare(p.0.lhs, p.0.rhs), direction=GT
+    }
+
+    ENTRY top_k {
+      arg = bf16[8,65540] parameter(0)
+      ROOT result = (bf16[8,32], s32[8,32]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = true, order = "PARTIAL"}
+    }
+  )";
+
+  HloModuleConfig config = GetModuleConfigForTest();
+  config.mutable_debug_options()
+      .set_xla_gpu_experimental_enable_raft_for_stable_topk(true);
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(hlo, config));
+
+  if (!device_description().gpu_compute_capability().IsCuda()) {
+    GTEST_SKIP() << "RAFT is CUDA-only.";
+  }
+
+  ASSERT_OK_AND_ASSIGN(
+      bool changed,
+      TopkSpecializer(device_description().gpu_compute_capability())
+          .Run(module.get()));
+  ASSERT_TRUE(changed);
+
+  const char* check_pattern = R"(
+// CHECK-LABEL: ENTRY %top_k
+// CHECK: %[[ARG:[^ ]+]] = bf16[8,65540]{{.*}} parameter(0)
+
+// 1. Convert BF16 -> F32, then Bitcast to U32 and S32
+// CHECK: %[[F32_VAL:[^ ]+]] = f32[8,65540]{{.*}} convert(%[[ARG]])
+// CHECK: %[[VAL_U32:[^ ]+]] = u32[8,65540]{{.*}} bitcast-convert(%[[F32_VAL]])
+// CHECK: %[[VAL_S32:[^ ]+]] = s32[8,65540]{{.*}} bitcast-convert(%[[F32_VAL]])
+
+// 2. Pure Bitwise Radix Flip (F32 -> U32)
+// CHECK: %[[SRA:[^ ]+]] = s32[8,65540]{{.*}} shift-right-arithmetic(%[[VAL_S32]], {{.*}})
+// CHECK: %[[SRA_U32:[^ ]+]] = u32[8,65540]{{.*}} bitcast-convert(%[[SRA]])
+// CHECK: %[[MASK:[^ ]+]] = u32[8,65540]{{.*}} or(%[[SRA_U32]], {{.*}})
+// CHECK: %[[RADIX_KEY:[^ ]+]] = u32[8,65540]{{.*}} xor(%[[VAL_U32]], %[[MASK]])
+
+// 3. Canonicalize -0.0f and pack into U64 with 31-bit reversed iota and is_neg_zero bit
+// CHECK: %[[IS_NEG_ZERO_PRED:[^ ]+]] = pred[8,65540]{{.*}} compare(%[[VAL_U32]], {{.*}}), direction=EQ
+// CHECK: %[[IS_NEG_ZERO:[^ ]+]] = u32[8,65540]{{.*}} convert(%[[IS_NEG_ZERO_PRED]])
+// CHECK: %[[PARTIAL_KEY:[^ ]+]] = u32[8,65540]{{.*}} add(%[[RADIX_KEY]], %[[IS_NEG_ZERO]])
+// CHECK: %[[KEY_U64:[^ ]+]] = u64[8,65540]{{.*}} convert(%[[PARTIAL_KEY]])
+// CHECK: %[[SHIFT_LEFT:[^ ]+]] = u64[8,65540]{{.*}} shift-left(%[[KEY_U64]], {{.*}})
+// CHECK: %[[IOTA:[^ ]+]] = u32[8,65540]{{.*}} iota(), iota_dimension=1
+// CHECK: %[[IOTA_NEG:[^ ]+]] = u32[8,65540]{{.*}} subtract({{.*}}, %[[IOTA]])
+// CHECK: %[[IOTA_SHIFTED:[^ ]+]] = u32[8,65540]{{.*}} shift-left(%[[IOTA_NEG]], {{.*}})
+// CHECK: %[[LOW_U32:[^ ]+]] = u32[8,65540]{{.*}} or(%[[IOTA_SHIFTED]], %[[IS_NEG_ZERO]])
+// CHECK: %[[LOW_U64:[^ ]+]] = u64[8,65540]{{.*}} convert(%[[LOW_U32]])
+// CHECK: %[[PACKED:[^ ]+]] = u64[8,65540]{{.*}} or(%[[SHIFT_LEFT]], %[[LOW_U64]])
+
+// 4. CustomCall (__gpu$TopK)
+// CHECK: %[[CUSTOM_CALL:[^ ]+]] = (u64[8,32]{{.*}}, s32[8,32]{{.*}}, u8[33554432]{{.*}}) custom-call(%[[PACKED]]), custom_call_target="__gpu$TopK", api_version=API_VERSION_TYPED_FFI, backend_config={is_stable = false, order = "PARTIAL"}
+
+// 5. Unpack U64 -> U32 and restore -0.0f radix key
+// CHECK: %[[SRL:[^ ]+]] = u64[8,32]{{.*}} shift-right-logical(%[[CUSTOM_CALL]]#0, {{.*}})
+// CHECK: %[[PARTIAL_UNPACK_U32:[^ ]+]] = u32[8,32]{{.*}} convert(%[[SRL]])
+// CHECK: %[[UNPACK_LOW_U32:[^ ]+]] = u32[8,32]{{.*}} convert(%[[CUSTOM_CALL]]#0)
+// CHECK: %[[UNPACK_IS_NEG_ZERO:[^ ]+]] = u32[8,32]{{.*}} and(%[[UNPACK_LOW_U32]], {{.*}})
+// CHECK: %[[UNPACK_U32:[^ ]+]] = u32[8,32]{{.*}} subtract(%[[PARTIAL_UNPACK_U32]], %[[UNPACK_IS_NEG_ZERO]])
+
+// 6. Pure Bitwise Radix Unflip (U32 -> F32)
+// CHECK: %[[MSB:[^ ]+]] = u32[8,32]{{.*}} shift-right-logical(%[[UNPACK_U32]], {{.*}})
+// CHECK: %[[MSB_MINUS_ONE:[^ ]+]] = u32[8,32]{{.*}} subtract(%[[MSB]], {{.*}})
+// CHECK: %[[UNMASK:[^ ]+]] = u32[8,32]{{.*}} or(%[[MSB_MINUS_ONE]], {{.*}})
+// CHECK: %[[UNFLIPPED:[^ ]+]] = u32[8,32]{{.*}} xor(%[[UNPACK_U32]], %[[UNMASK]])
+// CHECK: %[[F32_OUT:[^ ]+]] = f32[8,32]{{.*}} bitcast-convert(%[[UNFLIPPED]])
+
+// 7. Convert F32 -> BF16 and Tuple output
+// CHECK: %[[BF16_OUT:[^ ]+]] = bf16[8,32]{{.*}} convert(%[[F32_OUT]])
+// CHECK: ROOT {{.*}} tuple({{.*}}%[[BF16_OUT]], {{.*}}%[[CUSTOM_CALL]]#1)
+  )";
+
+  ASSERT_OK_AND_ASSIGN(bool filecheck_matched,
+                       RunFileCheck(module->ToString(), check_pattern));
+  EXPECT_TRUE(filecheck_matched);
+}
+
+TEST_F(TopkTest, SkipsSmallPartialOrderTopK) {
+  const char* hlo = R"(
+    HloModule m
+
+    %compare-gt.1 {
+      p.1.lhs = s32[] parameter(2)
+      p.1.rhs = s32[] parameter(3)
+      p.0.lhs = f32[] parameter(0)
+      p.0.rhs = f32[] parameter(1)
+      ROOT compare = pred[] compare(p.0.lhs, p.0.rhs), direction=GT
+    }
+
+    ENTRY top_k {
+      arg = f32[8,1024] parameter(0)
+      ROOT small_k = (f32[8,8], s32[8,8]) custom-call(arg), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = true, order = "PARTIAL"}
     }
   )";
 
@@ -388,5 +562,62 @@ TEST_F(TopkTest, SkipsPartialOrderTopK) {
   EXPECT_FALSE(changed);
 }
 
+TEST_F(TopkTest, RewriteStablePartialOrderTopKProducesCorrectResult) {
+  if (!device_description().gpu_compute_capability().IsCuda()) {
+    GTEST_SKIP() << "RAFT is CUDA-only.";
+  }
+
+  const char* hlo = R"(
+    HloModule m
+
+    %compare-gt.1 {
+      p.1.lhs = s32[] parameter(2)
+      p.1.rhs = s32[] parameter(3)
+      p.0.lhs = f32[] parameter(0)
+      p.0.rhs = f32[] parameter(1)
+      ROOT compare = pred[] compare(p.0.lhs, p.0.rhs), direction=GT
+    }
+
+    ENTRY top_k {
+      arg = f32[4,1024] parameter(0)
+      iota = s32[4,1024] iota(), iota_dimension=1
+      c2 = s32[] constant(2)
+      b2 = s32[4,1024] broadcast(c2), dimensions={}
+      rem = s32[4,1024] remainder(iota, b2)
+      c0 = s32[] constant(0)
+      b0 = s32[4,1024] broadcast(c0), dimensions={}
+      is_even = pred[4,1024] compare(rem, b0), direction=EQ
+      pos_zero_c = f32[] constant(0)
+      neg_zero_c = f32[] constant(-0)
+      pos_zero = f32[4,1024] broadcast(pos_zero_c), dimensions={}
+      neg_zero = f32[4,1024] broadcast(neg_zero_c), dimensions={}
+      signed_zeros = f32[4,1024] select(is_even, pos_zero, neg_zero)
+      ROOT result = (f32[4,32], s32[4,32]) custom-call(signed_zeros), custom_call_target="TopK", called_computations={%compare-gt.1}, backend_config={is_stable = true, order = "PARTIAL"}
+    }
+  )";
+
+  HloModuleConfig config = GetModuleConfigForTest();
+  config.mutable_debug_options()
+      .set_xla_gpu_experimental_enable_raft_for_stable_topk(true);
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> topk_module,
+                       ParseAndReturnVerifiedModule(hlo, config));
+  std::unique_ptr<HloModule> reference_module = topk_module->Clone("reference");
+
+  ASSERT_OK_AND_ASSIGN(
+      bool changed,
+      TopkSpecializer(device_description().gpu_compute_capability())
+          .Run(topk_module.get()));
+  ASSERT_TRUE(changed);
+
+  ASSERT_OK_AND_ASSIGN(changed, TopkDecomposer().Run(reference_module.get()));
+  ASSERT_TRUE(changed);
+  reference_module->mutable_config()
+      .mutable_debug_options()
+      .add_xla_disable_hlo_passes("topk-rewriter");
+
+  EXPECT_TRUE(RunAndCompareTwoModules(
+      std::move(topk_module), std::move(reference_module), std::nullopt));
+}
+
 }  // namespace
 }  // namespace xla::gpu

From a0a3f512aaf63410fd202c5ba8099ad0b1d043ca Mon Sep 17 00:00:00 2001
From: Blake Hechtman <blakehechtman@google.com>
Date: Tue, 13 Oct 2026 11:48:14 -0700
Subject: [PATCH 0709/1000] Add Tile::kMatchDimension for tile dimensions that
 match the shape dimension

Introduce Tile::kMatchDimension (printed and parsed as d in T(d, 128)) as a
separate sentinel from Tile::kCombineDimension (*) across XLA Layout,
HloParser, LayoutUtil, ShapeUtil, AlgebraicSimplifier, and PjRtLayout.

PiperOrigin-RevId: 998788337
---
 third_party/xla/xla/hlo/parser/hlo_parser.cc  |  7 +++-
 .../xla/xla/hlo/parser/hlo_parser_test.cc     | 10 +++++
 .../simplifiers/algebraic_simplifier.cc       | 12 +++---
 third_party/xla/xla/layout.cc                 |  2 +
 third_party/xla/xla/layout.h                  | 16 +++++---
 third_party/xla/xla/layout_test.cc            |  6 +++
 third_party/xla/xla/layout_util.cc            | 37 ++++++++++---------
 third_party/xla/xla/layout_util_test.cc       |  6 +--
 .../xla/pjrt/interpreter/interpreter_client.h |  2 +-
 third_party/xla/xla/pjrt/pjrt_layout.h        | 10 +++--
 third_party/xla/xla/shape_util.cc             |  2 +-
 11 files changed, 70 insertions(+), 40 deletions(-)

diff --git a/third_party/xla/xla/hlo/parser/hlo_parser.cc b/third_party/xla/xla/hlo/parser/hlo_parser.cc
index f157d074e29..757db3152d6 100644
--- a/third_party/xla/xla/hlo/parser/hlo_parser.cc
+++ b/third_party/xla/xla/hlo/parser/hlo_parser.cc
@@ -7311,7 +7311,7 @@ bool HloParserImpl::ParseDimLevelTypes(
 //   ::= 'T' ('(' dim_list ')')+
 // dim_list
 //   ::= /*empty*/
-//   ::= (int64_t | '*') (',' (int64_t | '*'))*
+//   ::= (int64_t | '*' | 'd') (',' (int64_t | '*' | 'd'))*
 bool HloParserImpl::ParseTiles(std::vector<Tile>* tiles) {
   auto parse_and_add_tile_dimension = [&]() {
     if (lexer_.GetKind() == TokKind::kAsterisk) {
@@ -7319,6 +7319,11 @@ bool HloParserImpl::ParseTiles(std::vector<Tile>* tiles) {
       lexer_.Lex();
       return true;
     }
+    if (lexer_.GetKind() == TokKind::kIdent && lexer_.GetStrVal() == "d") {
+      tiles->back().add_dimensions(Tile::kMatchDimension);
+      lexer_.Lex();
+      return true;
+    }
     int64_t i;
     if (ParseInt64(&i)) {
       tiles->back().add_dimensions(i);
diff --git a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc
index 8f6c7233512..9fc0944aba3 100644
--- a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc
+++ b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc
@@ -5850,6 +5850,16 @@ TEST_F(HloParserTest, ParseShapeStringWithTilingLayout) {
       << "expected: " << ShapeUtil::HumanStringWithLayout(expected)
       << "actual:   " << ShapeUtil::HumanStringWithLayout(actual);
 
+  // Tile with matching dimension ('d').
+  shape_string = "bf16[123,456]{1,0:T(d,128)(2,1)}";
+  ASSERT_OK_AND_ASSIGN(actual, ParseShape(shape_string));
+  expected = ShapeUtil::MakeShapeWithDenseLayout(
+      BF16, {123, 456}, {1, 0},
+      {Tile({Tile::kMatchDimension, 128}), Tile({2, 1})});
+  EXPECT_EQ(expected, actual)
+      << "expected: " << ShapeUtil::HumanStringWithLayout(expected)
+      << "actual:   " << ShapeUtil::HumanStringWithLayout(actual);
+
   // Wrong minor_to_major.
   shape_string = "f32[123,456,789]{1:T(2, * , 128)}";
   auto result = ParseShape(shape_string);
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
index 69e32bffd23..543737216a5 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
@@ -295,12 +295,12 @@ bool IsAllFpConstantPowerOf2(const HloInstruction* op) {
   return mantissa == 0.5 || mantissa == -0.5;
 }
 
-bool HasCombineDimensionTile(const Shape& shape) {
+bool HasMatchDimensionTile(const Shape& shape) {
   if (!shape.has_layout()) {
     return false;
   }
   for (const Tile& tile : shape.layout().tiles()) {
-    if (absl::c_linear_search(tile.dimensions(), Tile::kCombineDimension)) {
+    if (absl::c_linear_search(tile.dimensions(), Tile::kMatchDimension)) {
       return true;
     }
   }
@@ -316,8 +316,8 @@ bool TransposeIsBitcast(const HloInstruction* transpose) {
                                      transpose->dimensions())) {
     return false;
   }
-  if (!HasCombineDimensionTile(operand->shape()) &&
-      !HasCombineDimensionTile(transpose->shape())) {
+  if (!HasMatchDimensionTile(operand->shape()) &&
+      !HasMatchDimensionTile(transpose->shape())) {
     return true;
   }
   return absl::c_equal(operand->shape().layout().tiles(),
@@ -1896,8 +1896,8 @@ absl::Status AlgebraicSimplifierVisitor::HandleCopy(HloInstruction* copy) {
       copy->operand(0)->user_count() == 1 &&
       ShapeUtil::ReshapeIsBitcast(copy->operand(0)->shape(), copy->shape()) &&
       (!options_.is_layout_sensitive() ||
-       (!HasCombineDimensionTile(copy->operand(0)->shape()) &&
-        !HasCombineDimensionTile(copy->shape())) ||
+       (!HasMatchDimensionTile(copy->operand(0)->shape()) &&
+        !HasMatchDimensionTile(copy->shape())) ||
        absl::c_equal(copy->operand(0)->shape().layout().tiles(),
                      copy->shape().layout().tiles()))) {
     return ReplaceWithNewInstruction(
diff --git a/third_party/xla/xla/layout.cc b/third_party/xla/xla/layout.cc
index d882a5b05d9..60973d26d5b 100644
--- a/third_party/xla/xla/layout.cc
+++ b/third_party/xla/xla/layout.cc
@@ -102,6 +102,8 @@ void Tile::Print(Printer* printer) const {
     } else {
       if (dim == kCombineDimension) {
         printer->Append("*");
+      } else if (dim == kMatchDimension) {
+        printer->Append("d");
       } else {
         printer->Append("Invalid value ");
         printer->Append(dim);
diff --git a/third_party/xla/xla/layout.h b/third_party/xla/xla/layout.h
index 89248d9f96e..233ddb7b5f8 100644
--- a/third_party/xla/xla/layout.h
+++ b/third_party/xla/xla/layout.h
@@ -52,7 +52,8 @@ class Tile {
     Tile tile;
     tile.dimensions_.reserve(tile_proto.dimensions_size());
     for (int64_t dimension : tile_proto.dimensions()) {
-      TF_RET_CHECK(dimension >= 0 || dimension == kCombineDimension);
+      TF_RET_CHECK(dimension >= 0 || dimension == kCombineDimension ||
+                   dimension == kMatchDimension);
       tile.add_dimensions(dimension);
     }
     return tile;
@@ -74,7 +75,8 @@ class Tile {
   // For example, (*,*,2,*,8) means that the tile has 2 dimensions, and the
   // dimension sizes are 2 and 8. '*' means the corresponding dimension in
   // the shape should be combined with the next more minor dimension before
-  // tiling.
+  // tiling. 'd' means the tile dimension size matches the corresponding
+  // dimension in the shape.
   void Print(Printer* printer) const;
   std::string ToString() const;
 
@@ -95,13 +97,15 @@ class Tile {
   }
 
   // This dimension size means the corresponding dimension in the shape is
-  // combined with the next minor dimension before tiling is applied. In a top
-  // level tile T(*, N), it indicates an unpadded dynamic dimension whose
-  // extent equals the logical dimension rounded up to the subtile packing
-  // factor.
+  // combined with the next minor dimension before tiling is applied.
   static constexpr int64_t kCombineDimension =
       std::numeric_limits<int64_t>::min();
 
+  // This dimension size means the tile dimension size matches the corresponding
+  // dimension in the shape.
+  static constexpr int64_t kMatchDimension =
+      std::numeric_limits<int64_t>::min() + 1;
+
   template <typename H>
   friend H AbslHashValue(H h, const Tile& t) {
     return H::combine(std::move(h), t.dimensions_);
diff --git a/third_party/xla/xla/layout_test.cc b/third_party/xla/xla/layout_test.cc
index 8b61b9c9943..3923dcb370e 100644
--- a/third_party/xla/xla/layout_test.cc
+++ b/third_party/xla/xla/layout_test.cc
@@ -51,6 +51,12 @@ TEST(Layout, ToStringForTileWithCombinedDimensions) {
             "{3,2,1,0:T(*,*,42,123)}");
 }
 
+TEST(Layout, ToStringForTileWithMatchedDimensions) {
+  EXPECT_EQ(
+      Layout({3, 2, 1, 0}, {Tile({Tile::kMatchDimension, 128})}).ToString(),
+      "{3,2,1,0:T(d,128)}");
+}
+
 TEST(Layout, ToStringForTailPaddingAlignment) {
   EXPECT_EQ(Layout({3, 2, 1, 0})
                 .set_tail_padding_alignment_in_elements(100)
diff --git a/third_party/xla/xla/layout_util.cc b/third_party/xla/xla/layout_util.cc
index f112fbe50a3..2397d7ea315 100644
--- a/third_party/xla/xla/layout_util.cc
+++ b/third_party/xla/xla/layout_util.cc
@@ -73,11 +73,11 @@ void SetDefaultLayoutToContainer(T* minor_to_major) {
   }
   for (const Tile& tile : tiles) {
     for (int64_t dim : tile.dimensions()) {
-      if (dim < 0 && dim != Tile::kCombineDimension) {
-        LOG(FATAL)
-            << "Tile dimension size needs to be minimum int64_t value if "
-               "it's negative. Value is "
-            << dim;
+      if (dim < 0 && dim != Tile::kCombineDimension &&
+          dim != Tile::kMatchDimension) {
+        LOG(FATAL) << "Tile dimension size needs to be kCombineDimension or "
+                      "kMatchDimension if it's negative. Value is "
+                   << dim;
       }
     }
     *layout.add_tiles() = tile;
@@ -287,10 +287,11 @@ Layout CreateDefaultLayoutForRank(int64_t num_dims) {
       return InvalidArgument("layout has invalid tiles: %s", shape.ToString());
     }
     for (int64_t dim : tile.dimensions()) {
-      if (dim <= 0 && dim != Tile::kCombineDimension) {
+      if (dim <= 0 && dim != Tile::kCombineDimension &&
+          dim != Tile::kMatchDimension) {
         return InvalidArgument(
-            "layout has invalid tiles: tile dimension %d must be positive or "
-            "kCombineDimension: %s",
+            "layout has invalid tiles: tile dimension %d must be positive, "
+            "kCombineDimension, or kMatchDimension: %s",
             dim, shape.ToString());
       }
     }
@@ -618,7 +619,7 @@ Layout LayoutUtil::MoveDimToMinor(const Layout& layout, const int64_t dim) {
     tile = shape.layout().tiles()[0];
   }
 
-  auto resolve_combine_dim = [&](int64_t i, int64_t dim_size) {
+  auto resolve_match_dim = [&](int64_t i, int64_t dim_size) {
     absl::Span<const int64_t> tile_dims = tile.dimensions();
     absl::Span<const int64_t> sub_tile_dims =
         shape.layout().tiles().size() > 1 ? shape.layout().tiles(1).dimensions()
@@ -641,13 +642,13 @@ Layout LayoutUtil::MoveDimToMinor(const Layout& layout, const int64_t dim) {
   // Initialize to number of elements in a tile.
   for (int64_t d_idx = 0; d_idx < tile.dimensions().size(); ++d_idx) {
     int64_t tile_dim_size = tile.dimensions()[d_idx];
-    if (tile_dim_size == Tile::kCombineDimension) {
+    if (tile_dim_size == Tile::kMatchDimension) {
       int64_t minor = tile.dimensions().size() - 1 - d_idx;
       if (minor >= shape.layout().minor_to_major().size()) {
         continue;
       }
       int64_t logical_dim = Minor(shape.layout(), minor);
-      tile_dim_size = resolve_combine_dim(d_idx, shape.dimensions(logical_dim));
+      tile_dim_size = resolve_match_dim(d_idx, shape.dimensions(logical_dim));
     }
     tile_multiplier *= tile_dim_size;
   }
@@ -662,8 +663,8 @@ Layout LayoutUtil::MoveDimToMinor(const Layout& layout, const int64_t dim) {
     if (minor < tile.dimensions().size()) {
       int64_t d_idx = tile.dimensions().size() - 1 - minor;
       int64_t tile_dim_size = tile.dimensions()[d_idx];
-      if (tile_dim_size == Tile::kCombineDimension) {
-        tile_dim_size = resolve_combine_dim(d_idx, shape_dim_size);
+      if (tile_dim_size == Tile::kMatchDimension) {
+        tile_dim_size = resolve_match_dim(d_idx, shape_dim_size);
       }
       linear_index += tile_multiplier * (index / tile_dim_size) +
                       within_tile_multiplier * (index % tile_dim_size);
@@ -733,7 +734,7 @@ Layout LayoutUtil::MoveDimToMinor(const Layout& layout, const int64_t dim) {
       int64_t d = current_shape[suffix_start + i];
       int64_t e = current_indices[suffix_start + i];
       int64_t t = tile.dimension(i);
-      if (t == Tile::kCombineDimension) {
+      if (t == Tile::kMatchDimension) {
         t = d == 0 ? 1 : d;
       }
       next_shape.push_back(CeilOfRatio(d, t));
@@ -746,7 +747,7 @@ Layout LayoutUtil::MoveDimToMinor(const Layout& layout, const int64_t dim) {
       int64_t d = current_shape[suffix_start + i];
       int64_t e = current_indices[suffix_start + i];
       int64_t t = tile.dimension(i);
-      if (t == Tile::kCombineDimension) {
+      if (t == Tile::kMatchDimension) {
         t = d == 0 ? 1 : d;
       }
       next_shape.push_back(t);
@@ -819,7 +820,7 @@ Layout LayoutUtil::MoveDimToMinor(const Layout& layout, const int64_t dim) {
     for (int i = 0; i < tile_rank; ++i) {
       int64_t d = current_shape[suffix_start + i];
       int64_t t = tile.dimension(i);
-      if (t == Tile::kCombineDimension) {
+      if (t == Tile::kMatchDimension) {
         t = d == 0 ? 1 : d;
       }
       next_shape.push_back(CeilOfRatio(d, t));
@@ -828,7 +829,7 @@ Layout LayoutUtil::MoveDimToMinor(const Layout& layout, const int64_t dim) {
     for (int i = 0; i < tile_rank; ++i) {
       int64_t d = current_shape[suffix_start + i];
       int64_t t = tile.dimension(i);
-      if (t == Tile::kCombineDimension) {
+      if (t == Tile::kMatchDimension) {
         t = d == 0 ? 1 : d;
       }
       next_shape.push_back(t);
@@ -866,7 +867,7 @@ Layout LayoutUtil::MoveDimToMinor(const Layout& layout, const int64_t dim) {
       int64_t outer_idx = current_indices[suffix_start + i];
       int64_t inner_idx = current_indices[suffix_start + tile_rank + i];
       int64_t tile_dim = step.tile.dimension(i);
-      if (tile_dim == Tile::kCombineDimension) {
+      if (tile_dim == Tile::kMatchDimension) {
         int64_t d = step.shape_before[suffix_start + i];
         tile_dim = d == 0 ? 1 : d;
       }
diff --git a/third_party/xla/xla/layout_util_test.cc b/third_party/xla/xla/layout_util_test.cc
index c13c915eb9d..9d15bca9a9c 100644
--- a/third_party/xla/xla/layout_util_test.cc
+++ b/third_party/xla/xla/layout_util_test.cc
@@ -355,13 +355,13 @@ TEST_F(LayoutUtilTest, HumanStringWithTiling) {
             "bf16[8,2,3,1004]{3,2,1,0:T(2,*,*,128)}");
 }
 
-TEST_F(LayoutUtilTest, LinearIndexWithCombineDimension) {
+TEST_F(LayoutUtilTest, LinearIndexWithMatchDimension) {
   Shape shape = ShapeUtil::MakeShapeWithDenseLayout(F32, {16, 256}, {1, 0});
   Tile* tile = shape.mutable_layout()->add_tiles();
-  tile->add_dimensions(Tile::kCombineDimension);
+  tile->add_dimensions(Tile::kMatchDimension);
   tile->add_dimensions(128);
 
-  // Tile (*, 128) on [16, 256] means each tile has size (16, 128).
+  // Tile (d, 128) on [16, 256] means each tile has size (16, 128).
   // (0, 0) -> 0
   EXPECT_EQ(LayoutUtil::LinearIndex(shape, {0, 0}), 0);
   EXPECT_EQ(LayoutUtil::LinearIndexForNestedTiling(shape, {0, 0}), 0);
diff --git a/third_party/xla/xla/pjrt/interpreter/interpreter_client.h b/third_party/xla/xla/pjrt/interpreter/interpreter_client.h
index 2e9bda4b9e9..0c93fd53372 100644
--- a/third_party/xla/xla/pjrt/interpreter/interpreter_client.h
+++ b/third_party/xla/xla/pjrt/interpreter/interpreter_client.h
@@ -164,7 +164,7 @@ class InterpreterLiteralWrapperBuffer final : public PjRtBuffer {
     Layout layout = on_device_shape().layout();
     for (const Tile& tile : layout.tiles()) {
       for (int64_t d : tile.dimensions()) {
-        if (d == Tile::kCombineDimension) {
+        if (d == Tile::kCombineDimension || d == Tile::kMatchDimension) {
           layout.clear_tiles();
           return std::make_shared<PjRtLayout>(std::move(layout));
         }
diff --git a/third_party/xla/xla/pjrt/pjrt_layout.h b/third_party/xla/xla/pjrt/pjrt_layout.h
index f276bece0d7..06dc96c838e 100644
--- a/third_party/xla/xla/pjrt/pjrt_layout.h
+++ b/third_party/xla/xla/pjrt/pjrt_layout.h
@@ -43,11 +43,13 @@ class PjRtLayout {
     xla_layout_.set_memory_space(xla::Layout::kDefaultMemorySpace);
 
     for (const auto& tile : xla_layout_.tiles()) {
-      CHECK(absl::c_none_of(
-          tile.dimensions(),
-          [](int64_t d) { return d == xla::Tile::kCombineDimension; }))
+      CHECK(absl::c_none_of(tile.dimensions(),
+                            [](int64_t d) {
+                              return d == xla::Tile::kCombineDimension ||
+                                     d == xla::Tile::kMatchDimension;
+                            }))
           << "Runtime buffers must not have layouts with `kCombineDimension` "
-             "in tile dimensions, but got "
+             "or `kMatchDimension` in tile dimensions, but got "
           << xla_layout_;
     }
   }
diff --git a/third_party/xla/xla/shape_util.cc b/third_party/xla/xla/shape_util.cc
index db1b53073a3..db6379cd74d 100644
--- a/third_party/xla/xla/shape_util.cc
+++ b/third_party/xla/xla/shape_util.cc
@@ -2544,7 +2544,7 @@ ShapeUtil::ByteStrides(const Shape& shape) {
                                                   shape, minor_to_major[dim])
                                             : 1;
     int64_t tile_dim = tile_dimensions[tile_dim_size - dim - 1];
-    if (tile_dim == Tile::kCombineDimension) {
+    if (tile_dim == Tile::kMatchDimension) {
       int64_t packing = 1;
       if (shape.layout().tiles().size() > 1) {
         const Tile& sub_tile = shape.layout().tiles(1);

From f7ed51baa8c6b716dcd705d63f274607c16994db Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Tue, 13 Oct 2026 12:00:40 -0700
Subject: [PATCH 0710/1000] [XLA:CPU] Improve float32 and float64 expm1
 precision from 5-6 ULP to < 1 ULP.

Previously, `ElementalIrEmitter::EmitExpm1` and `RewriteExpm1Pattern` in
`expand_float_ops.cc` lowered `expm1(x)` as a piecewise expression:
- `|x| > 0.5`: `exp(x) - 1`
- `|x| <= 0.5`: `tanh(x / 2) * (exp(x) + 1)`

This achieved a worst-case error of 6 ULP in float32 and > 5 ULP in float64.

Replace this piecewise lowering with a dedicated intrinsic
(`xla::codegen::intrinsics::Expm1`) for float32 and float64. The new
implementation reduces `x = n * ln(2) + r` into `[-ln(2)/2, ln(2)/2]`,
approximates `e^r - 1` with a Sollya minimax polynomial (degree 7 for `float32`,
degree 12 for `float64`), and reconstructs
`expm1(x) = (2^n - 1) + 2^n * (e^r - 1)` without catastrophic cancellation.
FMA instructions are used to preserve low-order bits across the quadratic term
(`0.5 * r_hi^2`) and intermediate sums, preventing double rounding across
power-of-two spacing boundaries.

This bounds worst-case error to < 0.83 ULP for `float32` (verified exhaustively
across all 2^32 inputs) and < 1 ULP for `float64`, while using fewer
instructions than evaluating both `tanh(x / 2)` and `exp(x)`.

Tighten the CPU `Expm1` error bound in `exhaustive_unary_test_functions.cc` to
1 ULP (`distance_err(1)`).

PiperOrigin-RevId: 998797194
---
 .../cpu/codegen/emitters/transforms/BUILD     |   1 -
 .../emitters/transforms/expand_float_ops.cc   |  58 +--
 .../transforms/tests/expand_float_ops.mlir    |  20 -
 .../transforms/lower_xla_intrinsic_lib.cc     |   1 +
 .../tests/lower_xla_intrinsic_lib.mlir        |  70 ++++
 .../intrinsic/accuracy/accuracy_budget.h      |  24 +-
 .../accuracy/intrinsic_accuracy_test.cc       |   5 +
 .../xla/xla/codegen/intrinsic/cpp/BUILD       |  68 ++++
 .../intrinsic/cpp/cpp_gen_intrinsics.cc       |  40 ++
 .../intrinsic/cpp/cpp_gen_intrinsics.h        |   8 +-
 .../xla/xla/codegen/intrinsic/cpp/expm1.cc    | 373 ++++++++++++++++++
 .../xla/xla/codegen/intrinsic/cpp/expm1.h     |  35 ++
 .../xla/codegen/intrinsic/cpp/expm1_test.cc   | 291 ++++++++++++++
 .../intrinsic/cpp/intrinsic_declarations.h    |  28 ++
 .../xla/codegen/intrinsic/cpp/vector_ops.h    | 133 ++++++-
 third_party/xla/xla/codegen/intrinsic_lib.cc  |  17 +-
 third_party/xla/xla/codegen/intrinsic_lib.h   |   2 +-
 .../xla/xla/codegen/intrinsic_lib_test.cc     |  38 +-
 third_party/xla/xla/service/BUILD             |   3 +-
 .../cpu/tests/cpu_vectorization_test.cc       |  50 +++
 .../xla/xla/service/elemental_ir_emitter.cc   |  39 +-
 .../exhaustive_unary_test_functions.cc        |   8 +-
 22 files changed, 1184 insertions(+), 128 deletions(-)
 create mode 100644 third_party/xla/xla/codegen/intrinsic/cpp/expm1.cc
 create mode 100644 third_party/xla/xla/codegen/intrinsic/cpp/expm1.h
 create mode 100644 third_party/xla/xla/codegen/intrinsic/cpp/expm1_test.cc

diff --git a/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/BUILD b/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/BUILD
index 6342d760f84..6026d1178f1 100644
--- a/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/BUILD
+++ b/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/BUILD
@@ -66,7 +66,6 @@ cc_library(
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:DataLayoutInterfaces",
-        "@llvm-project//mlir:DialectUtils",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:LLVMDialect",
diff --git a/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/expand_float_ops.cc b/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/expand_float_ops.cc
index 1d08be217fd..0f135558882 100644
--- a/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/expand_float_ops.cc
+++ b/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/expand_float_ops.cc
@@ -14,12 +14,12 @@ limitations under the License.
 ==============================================================================*/
 
 #include <cassert>
-#include <memory>
+#include <memory>  // IWYU pragma: keep
 #include <utility>
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/Math/IR/Math.h"
+#include "mlir/Dialect/Math/IR/Math.h"  // IWYU pragma: keep
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
 #include "mlir/IR/BuiltinAttributes.h"
@@ -31,7 +31,7 @@ limitations under the License.
 #include "mlir/IR/Types.h"
 #include "mlir/IR/Value.h"
 #include "mlir/IR/ValueRange.h"
-#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/Pass.h"  // IWYU pragma: keep
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -100,56 +100,6 @@ struct RewriteExtFPattern : public mlir::OpRewritePattern<ma::ExtFOp> {
   }
 };
 
-// Use a more numerically stable implementation of expm1(x).
-// |x| > 0.5: exp(x) - 1
-// |x| < 0.5: tanh(x/2) * (exp(x)+1)
-class RewriteExpm1Pattern : public mlir::OpRewritePattern<mlir::math::ExpM1Op> {
- public:
-  using OpRewritePattern::OpRewritePattern;
-
-  mlir::LogicalResult matchAndRewrite(
-      mlir::math::ExpM1Op op, mlir::PatternRewriter& rewriter) const override {
-    mlir::ImplicitLocOpBuilder b(op.getLoc(), rewriter);
-
-    mlir::Type type = op.getType();
-    mlir::Type element_type = mlir::getElementTypeOrSelf(type);
-    mlir::Value one = GetConst(b, type, b.getFloatAttr(element_type, 1.0));
-    mlir::Value half = GetConst(b, type, b.getFloatAttr(element_type, 0.5));
-    mlir::Value zero = GetConst(b, type, b.getFloatAttr(element_type, 0.0));
-    mlir::Value x = op.getOperand();
-
-    mlir::arith::FastMathFlagsAttr fastmath = op.getFastmathAttr();
-
-    mlir::Value exp_x = b.create<mlir::math::ExpOp>(x, fastmath);
-
-    mlir::Value exp_x_minus_1 =
-        b.create<mlir::arith::SubFOp>(exp_x, one, fastmath);
-
-    mlir::Value half_x = b.create<mlir::arith::MulFOp>(x, half, fastmath);
-    mlir::Value tanh_half_x = b.create<mlir::math::TanhOp>(half_x, fastmath);
-    mlir::Value exp_x_plus_1 =
-        b.create<mlir::arith::AddFOp>(exp_x, one, fastmath);
-    mlir::Value small_result =
-        b.create<mlir::arith::MulFOp>(tanh_half_x, exp_x_plus_1, fastmath);
-
-    mlir::Value abs_x = b.create<mlir::math::AbsFOp>(x, fastmath);
-    mlir::Value x_is_large = b.create<mlir::arith::CmpFOp>(
-        mlir::arith::CmpFPredicate::OGT, abs_x, half);
-    mlir::Value normal_result = b.create<mlir::arith::SelectOp>(
-        x_is_large, exp_x_minus_1, small_result);
-
-    // half_x can underflow resulting in zero.
-    // TODO(willfroom): Do we actually need this check? tanh(0) == 0.
-    mlir::Value half_x_is_zero = b.create<mlir::arith::CmpFOp>(
-        mlir::arith::CmpFPredicate::OEQ, half_x, zero);
-    mlir::Value result =
-        b.create<mlir::arith::SelectOp>(half_x_is_zero, x, normal_result);
-
-    rewriter.replaceOp(op, result);
-    return mlir::success();
-  }
-};
-
 class ExpandFloatOpsPass
     : public impl::ExpandFloatOpsPassBase<ExpandFloatOpsPass> {
  public:
@@ -157,7 +107,7 @@ class ExpandFloatOpsPass
 
   void runOnOperation() override {
     mlir::RewritePatternSet patterns(&getContext());
-    patterns.add<RewriteExtFPattern, RewriteExpm1Pattern>(&getContext());
+    patterns.add<RewriteExtFPattern>(&getContext());
 
     if (mlir::failed(
             mlir::applyPatternsGreedily(getOperation(), std::move(patterns)))) {
diff --git a/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/tests/expand_float_ops.mlir b/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/tests/expand_float_ops.mlir
index 91104684749..979a77d665e 100644
--- a/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/tests/expand_float_ops.mlir
+++ b/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/tests/expand_float_ops.mlir
@@ -27,23 +27,3 @@ func.func @extend_vector(%input: vector<8xbf16>) -> vector<8xf32> {
   %truncated = arith.extf %input : vector<8xbf16> to vector<8xf32>
   func.return %truncated : vector<8xf32>
 }
-
-// -----
-
-func.func @expm1(%arg0: f64) -> f64 {
-  %ret = math.expm1 %arg0 : f64
-  return %ret : f64
-}
-
-// CHECK-LABEL: @expm1
-// CHECK-NOT: math.expm1
-
-// -----
-
-func.func @expm1_vector(%arg0: vector<4xf64>) -> vector<4xf64> {
-  %ret = math.expm1 %arg0 : vector<4xf64>
-  return %ret : vector<4xf64>
-}
-
-// CHECK-LABEL: @expm1_vector
-// CHECK-NOT: math.expm1
diff --git a/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc b/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc
index 0eec56a66c6..a88e4357252 100644
--- a/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc
+++ b/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc
@@ -336,6 +336,7 @@ class LowerXlaIntrinsicLibPass
     mlir::RewritePatternSet patterns(context);
     absl::string_view cpu_features = cpu_features_.getValue();
     patterns.add<LowerIntrinsicPattern<ci::Exp, mm::ExpOp>,
+                 LowerIntrinsicPattern<ci::Expm1, mm::ExpM1Op>,
                  LowerIntrinsicPattern<ci::Log1p, mm::Log1pOp>,
                  LowerIntrinsicPattern<ci::Rsqrt, mm::RsqrtOp>,
                  LowerIntrinsicPattern<ci::Tanh, mm::TanhOp>,
diff --git a/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir b/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir
index 7eedc95ced8..d0df27f8185 100644
--- a/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir
+++ b/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir
@@ -365,6 +365,76 @@ func.func @rsqrt_f32_vector_16(%arg0: vector<16xf32>) -> vector<16xf32> {
 // AVX512: %[[CALL:.*]] = call @xla.rsqrt.v16f32(%arg0) : (vector<16xf32>) -> vector<16xf32>
 // AVX512: return %[[CALL]]
 
+// -----
 
+module {
+  func.func @expm1_f32(%arg0: f32) -> f32 {
+    %ret = math.expm1 %arg0 : f32
+    return %ret : f32
+  }
+}
 
+// CHECK-LABEL: @expm1_f32
+// CHECK-NOT: math.expm1
+// CHECK: %[[RESULT:.*]] = call @xla.expm1.f32(%arg0) : (f32) -> f32
+// CHECK: return %[[RESULT]] : f32
 
+// -----
+
+module {
+  func.func @expm1_f32_vector(%arg0: vector<4xf32>) -> vector<4xf32> {
+    %ret = math.expm1 %arg0 : vector<4xf32>
+    return %ret : vector<4xf32>
+  }
+}
+
+// CHECK-LABEL: @expm1_f32_vector
+// CHECK-NOT: math.expm1
+// CHECK: %[[RESULT:.*]] = call @xla.expm1.v4f32(%arg0) : (vector<4xf32>) -> vector<4xf32>
+// CHECK: return %[[RESULT]] : vector<4xf32>
+
+// -----
+
+module {
+  func.func @expm1_f64(%arg0: f64) -> f64 {
+    %ret = math.expm1 %arg0 : f64
+    return %ret : f64
+  }
+}
+
+// CHECK-LABEL: @expm1_f64
+// CHECK-NOT: math.expm1
+// CHECK: %[[RESULT:.*]] = call @xla.expm1.f64(%arg0) : (f64) -> f64
+// CHECK: return %[[RESULT]] : f64
+
+// -----
+
+module {
+  func.func @expm1_f64_vector(%arg0: vector<4xf64>) -> vector<4xf64> {
+    %ret = math.expm1 %arg0 : vector<4xf64>
+    return %ret : vector<4xf64>
+  }
+}
+
+// CHECK-LABEL: @expm1_f64_vector
+// CHECK-NOT: math.expm1
+// CHECK: %[[RESULT:.*]] = call @xla.expm1.v4f64(%arg0) : (vector<4xf64>) -> vector<4xf64>
+// CHECK: return %[[RESULT]] : vector<4xf64>
+
+// -----
+
+// expm1 vector_16 variant is gated on CPU features (AVX512).
+func.func @expm1_f32_vector_16(%arg0: vector<16xf32>) -> vector<16xf32> {
+  %ret = math.expm1 %arg0 : vector<16xf32>
+  return %ret : vector<16xf32>
+}
+// CHECK-LABEL: @expm1_f32_vector_16
+// CHECK: %[[S0:.*]] = vector.extract_strided_slice %arg0 offsets = [0], sizes = [8], strides = [1] : vector<16xf32> to vector<8xf32>
+// CHECK: %[[C0:.*]] = call @xla.expm1.v8f32(%[[S0]])
+// CHECK: %[[S1:.*]] = vector.extract_strided_slice %arg0 offsets = [8], sizes = [8], strides = [1] : vector<16xf32> to vector<8xf32>
+// CHECK: %[[C1:.*]] = call @xla.expm1.v8f32(%[[S1]])
+// CHECK: return
+
+// AVX512-LABEL: @expm1_f32_vector_16
+// AVX512: %[[CALL:.*]] = call @xla.expm1.v16f32(%arg0) : (vector<16xf32>) -> vector<16xf32>
+// AVX512: return %[[CALL]]
diff --git a/third_party/xla/xla/codegen/intrinsic/accuracy/accuracy_budget.h b/third_party/xla/xla/codegen/intrinsic/accuracy/accuracy_budget.h
index 8eb3273018a..2d0479caedf 100644
--- a/third_party/xla/xla/codegen/intrinsic/accuracy/accuracy_budget.h
+++ b/third_party/xla/xla/codegen/intrinsic/accuracy/accuracy_budget.h
@@ -64,11 +64,27 @@ constexpr AccuracyBudget kExpF32Budget = {
      /*subnormal=*/0},
 };
 
-constexpr AccuracyBudget kExpF64Budget = {
+constexpr AccuracyBudget kExpF64Budget = {/*cpu=*/{/*regular=*/1,
+                                                   /*subnormal=*/41132809365},
+                                          /*gpu=*/{/*regular=*/1,
+                                                   /*subnormal=*/0}};
+
+// Expm1
+constexpr AccuracyBudget kExpm1F32Budget = {
     /*cpu=*/{/*regular=*/1,
-             /*subnormal=*/41132809365},
-    /*gpu=*/{/*regular=*/1,
-             /*subnormal=*/0}};
+             /*subnormal=*/0},
+    /*gpu=*/
+    {/*regular=*/1,
+     /*subnormal=*/0},
+};
+
+constexpr AccuracyBudget kExpm1F64Budget = {
+    /*cpu=*/{/*regular=*/1,
+             /*subnormal=*/0},
+    /*gpu=*/
+    {/*regular=*/1,
+     /*subnormal=*/0},
+};
 
 // Log1p
 constexpr AccuracyBudget kLog1pF32Budget = {
diff --git a/third_party/xla/xla/codegen/intrinsic/accuracy/intrinsic_accuracy_test.cc b/third_party/xla/xla/codegen/intrinsic/accuracy/intrinsic_accuracy_test.cc
index 6f64b784aa6..16317626605 100644
--- a/third_party/xla/xla/codegen/intrinsic/accuracy/intrinsic_accuracy_test.cc
+++ b/third_party/xla/xla/codegen/intrinsic/accuracy/intrinsic_accuracy_test.cc
@@ -250,6 +250,11 @@ std::vector<IntrinsicAccuracyTestParam> GetAccuracyTestParams() {
       {"exponential", F32, accuracy::kGoldenExp, accuracy::kExpF32Budget},
       {"exponential", F64, accuracy::kGoldenExp, accuracy::kExpF64Budget},
 
+      {"exponential-minus-one", F32, accuracy::kGoldenExpm1,
+       accuracy::kExpm1F32Budget},
+      {"exponential-minus-one", F64, accuracy::kGoldenExpm1,
+       accuracy::kExpm1F64Budget},
+
       {"log-plus-one", F32, accuracy::kGoldenLog1p, accuracy::kLog1pF32Budget},
       {"log-plus-one", F64, accuracy::kGoldenLog1p, accuracy::kLog1pF64Budget},
 
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/BUILD b/third_party/xla/xla/codegen/intrinsic/cpp/BUILD
index a3be0cd9a5b..f0ef230b757 100644
--- a/third_party/xla/xla/codegen/intrinsic/cpp/BUILD
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/BUILD
@@ -235,6 +235,71 @@ xla_cc_test(
     ],
 )
 
+cc_library(
+    name = "expm1",
+    srcs = ["expm1.cc"],
+    hdrs = ["expm1.h"],
+    copts = [
+        "-Wno-psabi",
+    ],
+    deps = [
+        ":vector_ops",
+    ],
+)
+
+cc_ir_header(
+    name = "expm1_32_ll",
+    src = "expm1.cc",
+    copts = [
+        "-mprefer-vector-width=256",
+    ],
+    deps = [
+        ":cpp_context_provider",
+        ":expm1",
+        ":vector_ops",
+    ],
+)
+
+cc_ir_header(
+    name = "expm1_64_ll",
+    src = "expm1.cc",
+    copts = [
+        "-mprefer-vector-width=512",
+    ] + if_oss(
+        [],
+        select({
+            ":x86_64": ["-mavx512f"],
+            "//conditions:default": [],
+        }),
+    ),
+    deps = [
+        ":cpp_context_provider",
+        ":expm1",
+        ":vector_ops",
+    ],
+)
+
+xla_cc_test(
+    name = "expm1_test",
+    srcs = ["expm1_test.cc"],
+    copts = [
+        "-Wno-psabi",
+    ],
+    deps = [
+        ":_cpp_gen_intrinsics",
+        ":expm1",
+        ":expm1_32_ll",
+        ":expm1_64_ll",
+        ":vector_ops",
+        "//xla:types",
+        "@com_google_absl//absl/base",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Core",
+        "@llvm-project//llvm:Support",
+    ],
+)
+
 cc_binary(
     name = "dump_llvm_ir",
     srcs = ["dump_llvm_ir.cc"],
@@ -255,6 +320,8 @@ cc_library(
         ":eigen_unary_16_ll",
         ":eigen_unary_32_ll",
         ":eigen_unary_64_ll",
+        ":expm1_32_ll",
+        ":expm1_64_ll",
         "//xla/codegen/intrinsic",
         "//xla/service/llvm_ir:llvm_util",
         "@com_google_absl//absl/log",
@@ -279,6 +346,7 @@ cc_library(
         "//xla/codegen/intrinsic:type",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
         "@llvm-project//llvm:ir_headers",
     ],
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/cpp_gen_intrinsics.cc b/third_party/xla/xla/codegen/intrinsic/cpp/cpp_gen_intrinsics.cc
index 5ee1581de33..af0108bbc1f 100644
--- a/third_party/xla/xla/codegen/intrinsic/cpp/cpp_gen_intrinsics.cc
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/cpp_gen_intrinsics.cc
@@ -47,11 +47,27 @@ limitations under the License.
 #include "xla/codegen/intrinsic/cpp/eigen_unary_16_ll.h"
 #include "xla/codegen/intrinsic/cpp/eigen_unary_32_ll.h"
 #include "xla/codegen/intrinsic/cpp/eigen_unary_64_ll.h"
+#include "xla/codegen/intrinsic/cpp/expm1_32_ll.h"
+#include "xla/codegen/intrinsic/cpp/expm1_64_ll.h"
 #include "xla/codegen/intrinsic/intrinsic.h"
 #include "xla/service/llvm_ir/llvm_util.h"
 
 namespace xla::codegen {
 
+namespace {
+
+const std::string& SelectIrString(const intrinsics::IntrinsicOptions& options,
+                                  const std::string& ir_32,
+                                  const std::string& ir_64) {
+  if (options.Contains("+avx512f") && (options.prefer_vector_width >= 512 ||
+                                       options.prefer_vector_width == 0)) {
+    return ir_64;
+  }
+  return ir_32;
+}
+
+}  // namespace
+
 const std::string& GetCppGenIrString(
     const intrinsics::IntrinsicOptions& options) {
   if (options.Contains("+avx512f") && (options.prefer_vector_width >= 512 ||
@@ -72,6 +88,30 @@ bool AreEigenIntrinsicsAvailable(absl::string_view features) {
   return !GetCppGenIrString(options).empty();
 }
 
+bool AreCppGenIntrinsicsAvailable() {
+  return !GetCppGenIrString(intrinsics::IntrinsicOptions()).empty();
+}
+
+std::vector<CppGenIntrinsicLibrary> GetCppGenLibraries(
+    const intrinsics::IntrinsicOptions& options) {
+  std::vector<CppGenIntrinsicLibrary> libraries;
+  if (options.device_type != intrinsics::DeviceType::kIntelCpu &&
+      options.device_type != intrinsics::DeviceType::kAmdCpu &&
+      options.device_type != intrinsics::DeviceType::kArmCpu) {
+    return libraries;
+  }
+  const std::string& eigen_ir = GetCppGenIrString(options);
+  if (!eigen_ir.empty()) {
+    libraries.emplace_back(eigen_ir, "eigen");
+  }
+  const std::string& expm1_ir =
+      SelectIrString(options, ::llvm_ir::kExpm132LlIr, ::llvm_ir::kExpm164LlIr);
+  if (!expm1_ir.empty()) {
+    libraries.emplace_back(expm1_ir, "expm1");
+  }
+  return libraries;
+}
+
 namespace {
 
 void PrepareForInlining(llvm::Function* func) {
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/cpp_gen_intrinsics.h b/third_party/xla/xla/codegen/intrinsic/cpp/cpp_gen_intrinsics.h
index 2fd151299d6..4c93bbdb447 100644
--- a/third_party/xla/xla/codegen/intrinsic/cpp/cpp_gen_intrinsics.h
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/cpp_gen_intrinsics.h
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include <memory>
 #include <string>
+#include <vector>
 
 #include "absl/strings/string_view.h"
 #include "llvm/IR/Attributes.h"
@@ -42,9 +43,14 @@ std::unique_ptr<llvm::Module> ParseEmbeddedBitcode(
     llvm::LLVMContext& context, const std::string& bitcode,
     absl::string_view source_name = "embedded_module");
 
-// Returns true if the Eigen C++ intrinsics were compiled and are available.
+// Returns true if the C++ bitcode intrinsics were compiled and are available.
 // If the compiler does not support vector extensions, this will return false.
 bool AreEigenIntrinsicsAvailable(absl::string_view features = "");
+bool AreCppGenIntrinsicsAvailable();
+
+class CppGenIntrinsicLibrary;
+std::vector<CppGenIntrinsicLibrary> GetCppGenLibraries(
+    const intrinsics::IntrinsicOptions& options);
 
 // Helper for Intrinsic<T> classes that use CppGen backend for some types.
 // Looks up a function by name in the module (assuming it was linked from
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/expm1.cc b/third_party/xla/xla/codegen/intrinsic/cpp/expm1.cc
new file mode 100644
index 00000000000..8233375811d
--- /dev/null
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/expm1.cc
@@ -0,0 +1,373 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#if defined(__FLT16_MANT_DIG__) && defined(__has_attribute) &&    \
+    __has_attribute(ext_vector_type) && defined(__has_builtin) && \
+    __has_builtin(__builtin_vectorelements)
+
+#include "xla/codegen/intrinsic/cpp/expm1.h"
+
+#include <cstdint>
+#include <limits>
+
+#include "xla/codegen/intrinsic/cpp/vector_ops.h"
+
+namespace xla::codegen {
+namespace {
+
+// Computes expm1(x) = e^x - 1 for float32 within 0.72 ULP across all inputs
+// (assuming hardware FMA; without FMA, `@llvm.fmuladd` falls back to unfused
+// multiply/add and accuracy degrades by a few ULPs).
+// See https://en.cppreference.com/w/cpp/numeric/math/expm1 for background.
+//
+// Algorithm overview:
+// 1. Range reduction: We write x = n * ln(2) + (r_hi + r_lo), where
+//    n = round(x / ln(2)) and (r_hi, r_lo) = FastTwoSum(x - n * kLn2Hi,
+//    -n * kLn2Lo) in [-ln(2)/2, ln(2)/2]. Then:
+//      expm1(x) = e^x - 1 = 2^n * e^(r_hi + r_lo) - 1.
+// 2. Polynomial approximation: On [-ln(2)/2, ln(2)/2], we approximate the
+//    cubic-and-higher terms with a degree-4 Sollya minimax polynomial q(r).
+//    Since |r_lo| <= 0.5 * ulp(r_hi), we can drop r_hi * r_lo, r_lo^2, and
+//    higher-order r_lo terms:
+//      e^(r_hi + r_lo) - 1
+//        ~ r_hi + r_lo + (r_hi + r_lo)^2 / 2 + (r_hi + r_lo)^3 * q(r_hi + r_lo)
+//        ~ r_hi + r_lo + r_hi^2 / 2 + r_hi^3 * q(r_hi)
+//        = r_hi + tail
+//    where tail = r_lo + 0.5 * r_hi^2 + r_hi^3 * q(r_hi).
+// 3. Reconstruction: Substituting e^(r_hi + r_lo) ~ 1 + r_hi + tail into
+//    expm1(x) = 2^n * e^(r_hi + r_lo) - 1 gives:
+//      expm1(x) ~ 2^n * (1 + r_hi + tail) - 1
+//               = (2^n - 1) + 2^n * (r_hi + tail).
+//    We evaluate this by summing the two largest terms first
+//    (`head = (2^n - 1) + 2^n * r_hi`) and accumulating the lower-order
+//    rounding bits of `2^n - 1` and `head` into `tail` before the final
+//    addition (see Steps 3 and 5).
+template <typename VecT>
+inline VecT Expm1F32Impl(VecT input_x) {
+  constexpr float kLog2e = 0x1.715476p+0f;  // 1 / ln(2)
+  constexpr float kHalf = 0.5f;
+  constexpr float kOne = 1.0f;
+  constexpr float kNegOne = -1.0f;
+
+  // Step 1: Thresholds and input clamping.
+  // For x > 88.72283935546875f (0x1.62e42ep+6f), expm1(x) overflows float32
+  // to +infinity. For x < -18.0f, e^x < 2^-25.95, so expm1(x) rounds to -1.0f
+  // in round-to-nearest-even.
+  constexpr float kExpHi = 0x1.62e42ep+6f;
+  constexpr float kExpLo = -18.0f;
+
+  // We clamp x to [kExpLo, kExpHi] before computing n = round(x / ln(2)) so
+  // that n stays within [-26, 128] and integer conversions never overflow
+  // int32.
+  VecT x_clamped = Clamp(input_x, kExpLo, kExpHi);
+
+  // Step 2: Cody-Waite range reduction (x = n * ln(2) + r_hi + r_lo).
+  // Compute n = round(x / ln(2)) using round-to-nearest-even so |r_hi| <=
+  // ln(2)/2.
+  VecT n = __builtin_elementwise_rint(x_clamped * kLog2e);
+
+  // We split ln(2) into a 15-bit high part (kLn2Hi) and a 24-bit low part
+  // (kLn2Lo) to represent ln(2) to ~39 bits of precision. Because ln(2) is
+  // irrational, a single float32 constant only holds 24 bits of ln(2). When x
+  // is close to n * ln(2), their leading bits cancel, so subtracting a single
+  // 24-bit constant leaves a remainder missing the lower bits of ln(2).
+  // Choosing kLn2Hi (0x1.62e4p-1f) with 9 trailing zero bits ensures that for
+  // |n| <= 128 (at most 7 bits), the product n * kLn2Hi requires at most
+  // 15 + 7 = 22 <= 24 significand bits and is exact in float32. By Sterbenz's
+  // lemma, x - n * kLn2Hi is also exact, and FastTwoSum normalizes
+  // (x - n * kLn2Hi) and (-n * kLn2Lo) into a non-overlapping pair
+  // (r_hi, r_lo).
+  constexpr float kLn2Hi = 0x1.62e4p-1f;
+  constexpr float kLn2Lo = 0x1.7f7d1cp-20f;
+
+  VecT neg_n = -n;
+  auto [r_hi, r_lo] = FastTwoSum(Fmuladd(neg_n, Splat<VecT>(kLn2Hi), x_clamped),
+                                 neg_n * kLn2Lo);
+
+  // Step 3: Polynomial approximation for e^(r_hi + r_lo) - 1 on
+  // [-ln(2)/2, ln(2)/2].
+  //
+  // We approximate q(r) = (e^r - 1 - r - 0.5 * r^2) / r^3 using a degree-4
+  // Sollya minimax polynomial (so e^r - 1 is degree 7 overall):
+  //   q(r) ~ C3 + C4 * r + C5 * r^2 + C6 * r^3 + C7 * r^4
+  //
+  // Sollya script to derive C3..C7 and certify the relative ULP bound:
+  //   prec = 165!;
+  //   dom = [-log(2)/2; log(2)/2];
+  //   f = (expm1(x) - x - 0.5*x^2) / x^3;
+  //   p = fpminimax(f, 4, [|SG...|], dom);
+  //   display = hexadecimal!;
+  //   print(p);
+  //   display = decimal!;
+  //   print("Max relative error (in 2^-24 ULPs):",
+  //         sup(supnorm(x + 0.5*x^2 + x^3*p, expm1(x), dom, relative, 1b-60))
+  //         * 2^24);
+  // Certified supnorm relative error bound: 0.0197 ULPs (in units of 2^-24).
+  VecT q = HornerPoly(r_hi,
+                      0x1.a2ffc4p-13f,  // C7
+                      0x1.6d0ffap-10f,  // C6
+                      0x1.110fbp-7f,    // C5
+                      0x1.55551ap-5f,   // C4
+                      0x1.555556p-3f);  // C3
+
+  // Because `|r_lo| <= 0.5 * ulp(r_hi)` from FastTwoSum, the cross-term
+  // `r_hi * r_lo` and `r_lo^2` in `(r_hi + r_lo)^2 / 2` are `<= 2^-24 *
+  // r_hi^2` and can be dropped:
+  //   e^(r_hi + r_lo) - 1
+  //     ~ r_hi + r_lo + (r_hi + r_lo)^2 / 2 + (r_hi + r_lo)^3 * q(r_hi + r_lo)
+  //     ~ r_hi + r_lo + r_hi^2 / 2 + r_hi^3 * q(r_hi)
+  //     = r_hi + (0.5 * r_hi^2 + (r_hi^3 * q(r_hi) + r_lo))
+  //     = r_hi + tail.
+  VecT r2 = r_hi * r_hi;
+  VecT r_lo_adj = Fmuladd(r2, r_hi * q, r_lo);
+  VecT half_r_hi = kHalf * r_hi;
+  VecT tail = Fmuladd(half_r_hi, r_hi, r_lo_adj);
+
+  // Step 4: Construct 2^n via IEEE-754 exponent bit manipulation.
+  // We split n = n1 + n2 into two scale factors s1 = 2^n1 and s2 = 2^n2 with
+  // n1 = min(n, 127) and n2 = n - n1.
+  // Near the upper threshold x ~ 88.72, n reaches 128. Because the maximum
+  // finite float32 exponent is 127, constructing 2^128 directly would overflow
+  // to +infinity. Clamping n1 to 127 ensures both s1 = 2^n1 and s2 = 2^n2 stay
+  // finite (with s2 = 1 almost everywhere, and s2 = 2 only when n == 128).
+  constexpr int32_t kMaxExp = 127;
+  constexpr int32_t kMantissaBits = 23;
+
+  auto n_int = ToSignedInt(n);
+  auto max_exp = Splat<decltype(n_int)>(kMaxExp);
+  auto n1 = n_int < max_exp ? n_int : max_exp;
+  auto n2 = n_int - n1;
+
+  auto exp1_bits = (n1 + max_exp) << kMantissaBits;
+  auto exp2_bits = (n2 + max_exp) << kMantissaBits;
+  VecT s1 = __builtin_bit_cast(VecT, exp1_bits);
+  VecT s2 = __builtin_bit_cast(VecT, exp2_bits);
+
+  // Step 5: Reconstruct expm1(x) from 2^n = s1 * s2 (Step 4) and
+  // e^(r_hi + r_lo) ~ 1 + r_hi + tail (Step 3), where
+  // `tail = 0.5 * r_hi^2 + r_hi^3 * q(r_hi) + r_lo`:
+  //   expm1(x) = 2^n * e^(r_hi + r_lo) - 1
+  //            ~ s1 * s2 * (1 + r_hi + tail) - 1
+  //            = ((s1 - 1/s2) + s1 * (r_hi + tail)) * s2
+  //            ~ ((s1 - 1) + s1 * (r_hi + tail)) * s2
+  //            = (((s1 - 1) + s1 * r_hi) + s1 * tail) * s2.
+  // Note that for all valid n, s1 - 1/s2 rounds to the same float32 value as
+  // s1 - 1 (when n <= 127, s2 = 1; when n == 128, s1 = 2^127 and both
+  // 2^127 - 0.5 and 2^127 - 1 round to 2^127).
+  //
+  // First we compute `head = (s1 - 1) + s1 * r_hi`, summing the two largest
+  // terms so that any leading cancellation (e.g., when s1 = 0.5, where
+  // `s1 - 1 = -0.5` and `s1 * r_hi > 0`) happens before adding the smaller
+  // `s1 * tail` term.
+  //
+  // Next, we compute the low-order bits discarded when forming `s1 - 1` and
+  // `head`, and add them into `tail_total = s1 * tail + head_err` before
+  // the final addition `head + tail_total`:
+  // - `FastTwoSum(s1, -1)` recovers the `-1` lost when `s1 >= 2^25` (where
+  //   float32 spacing is >= 2.0, so `s1 - 1` rounds to `s1`).
+  // - Because `s1_minus_one` and `head` are close in magnitude, their
+  //   difference `d = s1_minus_one - head` is exact. `Fmuladd(s1, r_hi, d)`
+  //   computes `(s1 * r_hi + s1_minus_one) - head` before rounding—recovering
+  //   the exact low-order bits that were rounded off when `head` was formed.
+  // Keeping these low-order bits prevents double-rounding across power-of-two
+  // spacing boundaries (e.g., at n = 24, where `head` temporarily rounds above
+  // 2^24 into the 2.0-spacing range before negative `s1 * tail` brings the
+  // result back below 2^24).
+  auto [s1_minus_one, s1_m1_err] = FastTwoSum(s1, Splat<VecT>(-kOne));
+  VecT head = Fmuladd(s1, r_hi, s1_minus_one);
+  VecT head_err = Fmuladd(s1, r_hi, s1_minus_one - head) + s1_m1_err;
+  VecT tail_total = Fmuladd(s1, tail, head_err);
+  VecT unscaled_result = head + tail_total;
+  VecT result = unscaled_result * s2;
+
+  // Step 6: Handle special values and boundary conditions:
+  // - If x == +/-0.0, return input_x directly to preserve the sign of -0.0.
+  // - If x < kExpLo, return -1.0 (asymptote as x -> -infinity).
+  // - If x > kExpHi, return +infinity.
+  // - If x is NaN, `ToSignedInt` guards `n_int` to 0 (so `s1 = s2 = 1.0`),
+  //   while `r_hi` and `tail` are quiet NaN, so `result` evaluates to quiet
+  //   NaN and all ordered comparisons below evaluate to false.
+  constexpr float kInf = std::numeric_limits<float>::infinity();
+  result = input_x == 0.0f ? input_x : result;
+  result = input_x < kExpLo ? Splat<VecT>(kNegOne) : result;
+  result = input_x > kExpHi ? Splat<VecT>(kInf) : result;
+  return result;
+}
+
+// Computes expm1(x) = e^x - 1 for float64 within 1 ULP across all inputs.
+// Follows the same algorithm as Expm1F32Impl above (Cody-Waite range reduction,
+// degree-12 Sollya minimax polynomial approximation on [-ln(2)/2, ln(2)/2],
+// and two-step reconstruction preserving low-order bits).
+template <typename VecT>
+inline VecT Expm1F64Impl(VecT input_x) {
+  constexpr double kLog2e = 0x1.71547652b82fep+0;  // 1 / ln(2)
+  constexpr double kHalf = 0.5;
+  constexpr double kOne = 1.0;
+  constexpr double kNegOne = -1.0;
+
+  // Step 1: Thresholds and input clamping.
+  // For x > 709.78271289338397 (0x1.62e42fefa39efp+9), expm1(x) overflows
+  // float64 to +infinity. For x < -0x1.2b708872320e2p+5 (~ -37.42), e^x <
+  // 2^-54, so expm1(x) rounds to -1.0 in round-to-nearest-even.
+  constexpr double kExpHi = 0x1.62e42fefa39efp+9;
+  constexpr double kExpLo = -0x1.2b708872320e2p+5;
+
+  // We clamp x to [kExpLo, kExpHi] before computing n = round(x / ln(2)) so
+  // that n stays within [-55, 1024] and integer conversions never overflow
+  // int64.
+  VecT x_clamped = Clamp(input_x, kExpLo, kExpHi);
+
+  // Step 2: Cody-Waite range reduction (x = n * ln(2) + r_hi + r_lo).
+  // Compute n = round(x / ln(2)) using round-to-nearest-even so |r_hi| <=
+  // ln(2)/2.
+  VecT n = __builtin_elementwise_rint(x_clamped * kLog2e);
+
+  // We split ln(2) into a 36-bit high part (kLn2Hi) and a 53-bit low part
+  // (kLn2Lo). Since |n| <= 1024 takes at most 10 bits, n * kLn2Hi requires at
+  // most 36 + 10 = 46 <= 53 bits and is exact in float64, making
+  // x - n * kLn2Hi exact by Sterbenz's lemma.
+  constexpr double kLn2Hi = 0x1.62e42fefa0000p-1;
+  constexpr double kLn2Lo = 0x1.cf79abc9e3b3ap-40;
+
+  VecT neg_n = -n;
+  auto [r_hi, r_lo] = FastTwoSum(Fmuladd(neg_n, Splat<VecT>(kLn2Hi), x_clamped),
+                                 neg_n * kLn2Lo);
+
+  // Step 3: Polynomial approximation for e^(r_hi + r_lo) - 1 on
+  // [-ln(2)/2, ln(2)/2].
+  //
+  // We approximate q(r) = (e^r - 1 - r - 0.5 * r^2) / r^3 using a degree-9
+  // Sollya minimax polynomial (so e^r - 1 is degree 12 overall):
+  //   q(r) ~ C3 + C4 * r + ... + C12 * r^9
+  //
+  // Sollya script to derive C3..C12 and certify the relative ULP bound:
+  //   prec = 256!;
+  //   dom = [-log(2)/2; log(2)/2];
+  //   p = fpminimax(expm1(x), [|3, 4, 5, 6, 7, 8, 9, 10, 11, 12|], [|D...|],
+  //                 dom, relative, floating, x + 0.5*x^2);
+  //   display = hexadecimal!;
+  //   print(p);
+  //   display = decimal!;
+  //   print("Max relative error (in 2^-53 ULPs):",
+  //         sup(supnorm(p, expm1(x), dom, relative, 1b-80)) * 2^53);
+  // Certified supnorm relative error bound: 0.0028 ULPs (in units of 2^-53).
+  VecT q = HornerPoly(r_hi,
+                      0x1.1f0413a6fc15ap-29,  // C12
+                      0x1.af5f1ac764956p-26,  // C11
+                      0x1.27e521cebe596p-22,  // C10
+                      0x1.71ddf816ea1afp-19,  // C9
+                      0x1.a01a01872c144p-16,  // C8
+                      0x1.a01a01b001948p-13,  // C7
+                      0x1.6c16c16c1c0d1p-10,  // C6
+                      0x1.111111110f657p-7,   // C5
+                      0x1.555555555554ap-5,   // C4
+                      0x1.5555555555559p-3);  // C3
+
+  VecT r2 = r_hi * r_hi;
+  VecT r_lo_adj = Fmuladd(r2, r_hi * q, r_lo);
+  VecT half_r_hi = kHalf * r_hi;
+  VecT tail = Fmuladd(half_r_hi, r_hi, r_lo_adj);
+
+  // Step 4: Construct 2^n via IEEE-754 exponent bit manipulation.
+  // We split n = n1 + n2 into two scale factors s1 = 2^n1 and s2 = 2^n2 with
+  // n1 = min(n, 1023) and n2 = n - n1.
+  // Near the upper threshold x ~ 709.78, n reaches 1024. Because the maximum
+  // finite float64 exponent is 1023, constructing 2^1024 directly would
+  // overflow to +infinity. Clamping n1 to 1023 ensures both s1 = 2^n1 and
+  // s2 = 2^n2 stay finite (with s2 = 1 almost everywhere, and s2 = 2 only when
+  // n == 1024).
+  constexpr int64_t kMaxExp = 1023;
+  constexpr int64_t kMantissaBits = 52;
+
+  auto n_int = ToSignedInt(n);
+  auto max_exp = Splat<decltype(n_int)>(kMaxExp);
+  auto n1 = n_int < max_exp ? n_int : max_exp;
+  auto n2 = n_int - n1;
+
+  auto exp1_bits = (n1 + max_exp) << kMantissaBits;
+  auto exp2_bits = (n2 + max_exp) << kMantissaBits;
+  VecT s1 = __builtin_bit_cast(VecT, exp1_bits);
+  VecT s2 = __builtin_bit_cast(VecT, exp2_bits);
+
+  // Step 5: Reconstruct expm1(x) from 2^n = s1 * s2 (Step 4) and
+  // e^(r_hi + r_lo) ~ 1 + r_hi + tail (Step 3), where
+  // `tail = 0.5 * r_hi^2 + r_hi^3 * q(r_hi) + r_lo`:
+  //   expm1(x) = 2^n * e^(r_hi + r_lo) - 1
+  //            ~ s1 * s2 * (1 + r_hi + tail) - 1
+  //            = ((s1 - 1/s2) + s1 * (r_hi + tail)) * s2
+  //            ~ ((s1 - 1) + s1 * (r_hi + tail)) * s2
+  //            = (((s1 - 1) + s1 * r_hi) + s1 * tail) * s2.
+  // Note that for all valid n, s1 - 1/s2 rounds to the same float64 value as
+  // s1 - 1 (when n <= 1023, s2 = 1; when n == 1024, s1 = 2^1023 and both
+  // 2^1023 - 0.5 and 2^1023 - 1 round to 2^1023).
+  //
+  // First we compute `head = (s1 - 1) + s1 * r_hi`, summing the two largest
+  // terms so that any leading cancellation (e.g., when s1 = 0.5, where
+  // `s1 - 1 = -0.5` and `s1 * r_hi > 0`) happens before adding the smaller
+  // `s1 * tail` term.
+  //
+  // Next, we compute the low-order bits discarded when forming `s1 - 1` and
+  // `head`, and add them into `tail_total = s1 * tail + head_err` before
+  // the final addition `head + tail_total`:
+  // - `FastTwoSum(s1, -1)` recovers the `-1` lost when `s1 >= 2^54` (where
+  //   float64 spacing is >= 2.0, so `s1 - 1` rounds to `s1`).
+  // - Because `s1_minus_one` and `head` are close in magnitude, their
+  //   difference `d = s1_minus_one - head` is exact. `Fmuladd(s1, r_hi, d)`
+  //   computes `(s1 * r_hi + s1_minus_one) - head` before rounding—recovering
+  //   the exact low-order bits that were rounded off when `head` was formed.
+  // Keeping these low-order bits prevents double-rounding across power-of-two
+  // spacing boundaries (e.g., at n = 53, where `head` temporarily rounds above
+  // 2^53 into the 2.0-spacing range before negative `s1 * tail` brings the
+  // result back below 2^53).
+  auto [s1_minus_one, s1_m1_err] = FastTwoSum(s1, Splat<VecT>(-kOne));
+  VecT head = Fmuladd(s1, r_hi, s1_minus_one);
+  VecT head_err = Fmuladd(s1, r_hi, s1_minus_one - head) + s1_m1_err;
+  VecT tail_total = Fmuladd(s1, tail, head_err);
+  VecT unscaled_result = head + tail_total;
+  VecT result = unscaled_result * s2;
+
+  // Step 6: Handle special values and boundary conditions:
+  // - If x == +/-0.0, return input_x directly to preserve the sign of -0.0.
+  // - If x < kExpLo, return -1.0 (asymptote as x -> -infinity).
+  // - If x > kExpHi, return +infinity.
+  // - If x is NaN, `ToSignedInt` guards `n_int` to 0 (so `s1 = s2 = 1.0`),
+  //   while `r_hi` and `tail` are quiet NaN, so `result` evaluates to quiet
+  //   NaN and all ordered comparisons below evaluate to false.
+  constexpr double kInf = std::numeric_limits<double>::infinity();
+  result = input_x == 0.0 ? input_x : result;
+  result = input_x < kExpLo ? Splat<VecT>(kNegOne) : result;
+  result = input_x > kExpHi ? Splat<VecT>(kInf) : result;
+  return result;
+}
+
+}  // namespace
+
+//===--------------------------------------------------------------------===//
+// XLA entry points, renamed with asm in header file.
+//===--------------------------------------------------------------------===//
+
+// Single precision
+float expm1_f32(float x) { return Expm1F32Impl(x); }
+Vec4f expm1_v4f32(Vec4f x) { return Expm1F32Impl(x); }
+Vec8f expm1_v8f32(Vec8f x) { return Expm1F32Impl(x); }
+Vec16f expm1_v16f32(Vec16f x) { return Expm1F32Impl(x); }
+
+// Double precision
+double expm1_f64(double x) { return Expm1F64Impl(x); }
+Vec4d expm1_v4f64(Vec4d x) { return Expm1F64Impl(x); }
+Vec8d expm1_v8f64(Vec8d x) { return Expm1F64Impl(x); }
+
+}  // namespace xla::codegen
+#endif  // defined(__FLT16_MANT_DIG__) && defined(__has_attribute) && ...
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/expm1.h b/third_party/xla/xla/codegen/intrinsic/cpp/expm1.h
new file mode 100644
index 00000000000..45d1914c431
--- /dev/null
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/expm1.h
@@ -0,0 +1,35 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+#ifndef XLA_CODEGEN_INTRINSIC_CPP_EXPM1_H_
+#define XLA_CODEGEN_INTRINSIC_CPP_EXPM1_H_
+
+#include "xla/codegen/intrinsic/cpp/vector_ops.h"
+
+namespace xla::codegen {
+
+// Single precision expm1 entry points.
+float expm1_f32(float x) asm("xla.expm1.f32");
+Vec4f expm1_v4f32(Vec4f x) asm("xla.expm1.v4f32");
+Vec8f expm1_v8f32(Vec8f x) asm("xla.expm1.v8f32");
+Vec16f expm1_v16f32(Vec16f x) asm("xla.expm1.v16f32");
+
+// Double precision expm1 entry points.
+double expm1_f64(double x) asm("xla.expm1.f64");
+Vec4d expm1_v4f64(Vec4d x) asm("xla.expm1.v4f64");
+Vec8d expm1_v8f64(Vec8d x) asm("xla.expm1.v8f64");
+
+}  // namespace xla::codegen
+
+#endif  // XLA_CODEGEN_INTRINSIC_CPP_EXPM1_H_
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/expm1_test.cc b/third_party/xla/xla/codegen/intrinsic/cpp/expm1_test.cc
new file mode 100644
index 00000000000..3c5f27be03b
--- /dev/null
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/expm1_test.cc
@@ -0,0 +1,291 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/codegen/intrinsic/cpp/expm1.h"
+
+#include <cmath>
+#include <cstdint>
+#include <limits>
+#include <memory>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/base/casts.h"
+#include "absl/strings/string_view.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/raw_ostream.h"
+#include "xla/codegen/intrinsic/cpp/cpp_gen_intrinsics.h"
+#include "xla/codegen/intrinsic/cpp/expm1_32_ll.h"
+#include "xla/codegen/intrinsic/cpp/expm1_64_ll.h"
+#include "xla/codegen/intrinsic/cpp/vector_ops.h"
+#include "xla/types.h"
+
+namespace xla::codegen::intrinsics {
+namespace {
+
+using ::testing::AllOf;
+using ::testing::ContainsRegex;
+using ::testing::Not;
+
+TEST(Expm1Test, BitcodeIsVectorizedAndPortable) {
+  for (const std::string& bitcode :
+       {llvm_ir::kExpm132LlIr, llvm_ir::kExpm164LlIr}) {
+    llvm::LLVMContext context;
+    std::unique_ptr<llvm::Module> module =
+        ParseEmbeddedBitcode(context, bitcode);
+    std::string ir;
+    llvm::raw_string_ostream stream(ir);
+    module->print(stream, nullptr);
+
+    EXPECT_THAT(
+        ir,
+        AllOf(ContainsRegex("define .*@xla.expm1.f32\\(float"),
+              ContainsRegex("define .*@xla.expm1.v4f32\\(<4 x float>"),
+              ContainsRegex("define .*@xla.expm1.v8f32\\(.*<8 x float>"),
+              ContainsRegex("define .*@xla.expm1.v16f32\\(.*<16 x float>"),
+              ContainsRegex("define .*@xla.expm1.f64\\(double"),
+              ContainsRegex("define .*@xla.expm1.v4f64\\(.*<4 x double>"),
+              ContainsRegex("define .*@xla.expm1.v8f64\\(.*<8 x double>"),
+              ContainsRegex("fmul <4 x float>"),
+              ContainsRegex("fmul <4 x double>"), ContainsRegex("llvm.fmuladd"),
+              Not(ContainsRegex("llvm\\.fma")),
+              Not(ContainsRegex("llvm\\.x86")),
+              Not(ContainsRegex("llvm\\.aarch64"))));
+  }
+}
+
+float ComputeUlpErrorF32(float actual, double expected) {
+  if (std::isnan(expected)) {
+    return std::isnan(actual) ? 0.0f : std::numeric_limits<float>::infinity();
+  }
+  float expected_f32 = static_cast<float>(expected);
+  if (std::isinf(expected_f32)) {
+    return (std::isinf(actual) &&
+            std::signbit(actual) == std::signbit(expected_f32))
+               ? 0.0f
+               : std::numeric_limits<float>::infinity();
+  }
+  if (expected == 0.0) {
+    return (actual == 0.0f && std::signbit(actual) == std::signbit(expected))
+               ? 0.0f
+               : std::numeric_limits<float>::infinity();
+  }
+  float direction = (static_cast<double>(actual) < expected)
+                        ? -std::numeric_limits<float>::infinity()
+                        : std::numeric_limits<float>::infinity();
+  float next = std::nextafter(expected_f32, direction);
+  double ulp =
+      std::abs(static_cast<double>(next) - static_cast<double>(expected_f32));
+  return static_cast<float>(std::abs(static_cast<double>(actual) - expected) /
+                            ulp);
+}
+
+double ComputeUlpErrorF64(double actual, double expected) {
+  if (std::isnan(expected)) {
+    return std::isnan(actual) ? 0.0 : std::numeric_limits<double>::infinity();
+  }
+  if (std::isinf(expected)) {
+    return (std::isinf(actual) &&
+            std::signbit(actual) == std::signbit(expected))
+               ? 0.0
+               : std::numeric_limits<double>::infinity();
+  }
+  if (expected == 0.0) {
+    return (actual == 0.0 && std::signbit(actual) == std::signbit(expected))
+               ? 0.0
+               : std::numeric_limits<double>::infinity();
+  }
+  double direction = (actual < expected)
+                         ? -std::numeric_limits<double>::infinity()
+                         : std::numeric_limits<double>::infinity();
+  double next = std::nextafter(expected, direction);
+  double ulp = std::abs(next - expected);
+  return std::abs(actual - expected) / ulp;
+}
+
+TEST(Expm1Test, ScalarF32AccuracyAndSpecialCases) {
+  const float special_inputs[] = {
+      0.0f,
+      -0.0f,
+      std::numeric_limits<float>::infinity(),
+      -std::numeric_limits<float>::infinity(),
+      std::numeric_limits<float>::quiet_NaN(),
+      std::numeric_limits<float>::signaling_NaN(),
+      std::numeric_limits<float>::denorm_min(),
+      -std::numeric_limits<float>::denorm_min(),
+      std::numeric_limits<float>::min(),
+      -std::numeric_limits<float>::min(),
+      88.0f,
+      88.722839f,
+      89.0f,
+      -17.32868f,
+      -20.0f,
+      -100.0f,
+      -0.69314718f,  // -ln(2) where expm1(x) = -0.5
+      -0.34657359f,  // -0.5 * ln(2)
+      0.34657359f,   // +0.5 * ln(2)
+      0.69314718f,   // +ln(2)
+      1e-4f,
+      -1e-4f,
+      1e-7f,
+      -1e-7f,
+  };
+  for (float x : special_inputs) {
+    float actual = expm1_f32(x);
+    double expected = std::expm1(static_cast<double>(x));
+    EXPECT_LE(ComputeUlpErrorF32(actual, expected), 1.0f)
+        << "x = " << x << ", actual = " << actual
+        << ", expected = " << static_cast<float>(expected);
+  }
+
+  for (int i = -10000; i <= 10000; ++i) {
+    float x = static_cast<float>(i) * 0.01f;
+    float actual = expm1_f32(x);
+    double expected = std::expm1(static_cast<double>(x));
+    EXPECT_LE(ComputeUlpErrorF32(actual, expected), 1.0f)
+        << "x = " << x << ", actual = " << actual
+        << ", expected = " << static_cast<float>(expected);
+  }
+}
+
+TEST(Expm1Test, VectorF32MatchesScalar) {
+  Vec4f v4 = {0.0f, -0.69314718f, 1.5f, 89.0f};
+  Vec4f r4 = expm1_v4f32(v4);
+  for (int i = 0; i < 4; ++i) {
+    EXPECT_LE(ComputeUlpErrorF32(r4[i], std::expm1(static_cast<double>(v4[i]))),
+              1.0f);
+  }
+
+  Vec8f v8 = {0.0f,        -0.0f, -0.69314718f,
+              0.69314718f, 1e-5f, std::numeric_limits<float>::quiet_NaN(),
+              -20.0f,      89.0f};
+  Vec8f r8 = expm1_v8f32(v8);
+  for (int i = 0; i < 8; ++i) {
+    EXPECT_LE(ComputeUlpErrorF32(r8[i], std::expm1(static_cast<double>(v8[i]))),
+              1.0f);
+  }
+}
+
+TEST(Expm1Test, ScalarF64AccuracyAndSpecialCases) {
+  const double special_inputs[] = {
+      0.0,
+      -0.0,
+      std::numeric_limits<double>::infinity(),
+      -std::numeric_limits<double>::infinity(),
+      std::numeric_limits<double>::quiet_NaN(),
+      std::numeric_limits<double>::signaling_NaN(),
+      std::numeric_limits<double>::denorm_min(),
+      -std::numeric_limits<double>::denorm_min(),
+      std::numeric_limits<double>::min(),
+      -std::numeric_limits<double>::min(),
+      709.0,
+      709.78271289338397,
+      710.0,
+      -37.42994775023705,
+      -40.0,
+      -100.0,
+      -0.6931471805599453,  // -ln(2)
+      -0.3465735902799726,  // -0.5 * ln(2)
+      0.3465735902799726,   // +0.5 * ln(2)
+      0.6931471805599453,   // +ln(2)
+      1e-8,
+      -1e-8,
+      1e-15,
+      -1e-15,
+  };
+  for (double x : special_inputs) {
+    double actual = expm1_f64(x);
+    double expected = std::expm1(x);
+    EXPECT_LE(ComputeUlpErrorF64(actual, expected), 1.0)
+        << "x = " << x << ", actual = " << actual
+        << ", expected = " << expected;
+  }
+
+  for (int i = -10000; i <= 10000; ++i) {
+    double x = static_cast<double>(i) * 0.07;
+    double actual = expm1_f64(x);
+    double expected = std::expm1(x);
+    EXPECT_LE(ComputeUlpErrorF64(actual, expected), 1.0)
+        << "x = " << x << ", actual = " << actual
+        << ", expected = " << expected;
+  }
+}
+
+TEST(Expm1Test, VectorF64MatchesScalar) {
+  Vec4d v4 = {0.0, -0.6931471805599453, 1.5, 710.0};
+  Vec4d r4 = expm1_v4f64(v4);
+  for (int i = 0; i < 4; ++i) {
+    EXPECT_LE(ComputeUlpErrorF64(r4[i], std::expm1(v4[i])), 1.0);
+  }
+
+  Vec8d v8 = {0.0,
+              -0.0,
+              -0.6931471805599453,
+              0.6931471805599453,
+              1e-9,
+              std::numeric_limits<double>::quiet_NaN(),
+              -40.0,
+              710.0};
+  Vec8d r8 = expm1_v8f64(v8);
+  for (int i = 0; i < 8; ++i) {
+    EXPECT_LE(ComputeUlpErrorF64(r8[i], std::expm1(v8[i])), 1.0);
+  }
+}
+
+TEST(Expm1Test, ExhaustiveF16AndBF16ZeroUlpViaF32) {
+  for (uint32_t bits = 0; bits <= 0xFFFF; ++bits) {
+    xla::half h = absl::bit_cast<xla::half>(static_cast<uint16_t>(bits));
+    float xf = static_cast<float>(h);
+    xla::half actual = static_cast<xla::half>(expm1_f32(xf));
+    if (std::isnan(xf)) {
+      EXPECT_TRUE(std::isnan(static_cast<float>(actual)));
+      continue;
+    }
+    xla::half expected =
+        static_cast<xla::half>(std::expm1(static_cast<double>(xf)));
+    if (std::isnan(static_cast<float>(expected))) {
+      EXPECT_TRUE(std::isnan(static_cast<float>(actual)));
+    } else {
+      EXPECT_EQ(absl::bit_cast<uint16_t>(actual),
+                absl::bit_cast<uint16_t>(expected))
+          << "F16 mismatch at x = " << xf;
+    }
+  }
+
+  for (uint32_t bits = 0; bits <= 0xFFFF; ++bits) {
+    xla::bfloat16 bf =
+        absl::bit_cast<xla::bfloat16>(static_cast<uint16_t>(bits));
+    float xbf = static_cast<float>(bf);
+    xla::bfloat16 actual_bf = static_cast<xla::bfloat16>(expm1_f32(xbf));
+    if (std::isnan(xbf)) {
+      EXPECT_TRUE(std::isnan(static_cast<float>(actual_bf)));
+      continue;
+    }
+    xla::bfloat16 expected_bf =
+        static_cast<xla::bfloat16>(std::expm1(static_cast<double>(xbf)));
+    if (std::isnan(static_cast<float>(expected_bf))) {
+      EXPECT_TRUE(std::isnan(static_cast<float>(actual_bf)));
+    } else {
+      EXPECT_EQ(absl::bit_cast<uint16_t>(actual_bf),
+                absl::bit_cast<uint16_t>(expected_bf))
+          << "BF16 mismatch at x = " << xbf;
+    }
+  }
+}
+
+}  // namespace
+}  // namespace xla::codegen::intrinsics
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/intrinsic_declarations.h b/third_party/xla/xla/codegen/intrinsic/cpp/intrinsic_declarations.h
index b5090e67966..ac6abb08950 100644
--- a/third_party/xla/xla/codegen/intrinsic/cpp/intrinsic_declarations.h
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/intrinsic_declarations.h
@@ -20,6 +20,7 @@ limitations under the License.
 
 #include "absl/log/check.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/match.h"
 #include "absl/strings/string_view.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/DerivedTypes.h"
@@ -83,6 +84,33 @@ class EigenAtan : public Intrinsic<EigenAtan> {
                              UnaryFunctionType(module, type));
   }
 };
+
+class Expm1 : public Intrinsic<Expm1> {
+ public:
+  static constexpr absl::string_view kName = "expm1";
+
+  static std::vector<std::vector<Type>> SupportedVectorTypes(
+      absl::string_view features = "") {
+    if (!AreCppGenIntrinsicsAvailable()) {
+      return {};
+    }
+    std::vector<std::vector<Type>> types = {
+        {Type::S(xla::F32)}, {Type::V(xla::F32, 4)}, {Type::V(xla::F32, 8)},
+        {Type::S(xla::F64)}, {Type::V(xla::F64, 4)},
+    };
+    if (absl::StrContains(features, "+avx512f")) {
+      types.push_back({Type::V(xla::F32, 16)});
+      types.push_back({Type::V(xla::F64, 8)});
+    }
+    return types;
+  }
+
+  static absl::StatusOr<llvm::Function*> CreateDefinition(
+      llvm::Module* module, const IntrinsicOptions& options, Type type) {
+    return GetCppGenFunction(module, Name(type),
+                             UnaryFunctionType(module, type));
+  }
+};
 }  // namespace xla::codegen::intrinsics
 
 #endif  // XLA_CODEGEN_INTRINSIC_CPP_INTRINSIC_DECLARATIONS_H_
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/vector_ops.h b/third_party/xla/xla/codegen/intrinsic/cpp/vector_ops.h
index 9003ca03897..62cb244487c 100644
--- a/third_party/xla/xla/codegen/intrinsic/cpp/vector_ops.h
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/vector_ops.h
@@ -22,6 +22,7 @@ limitations under the License.
 #include <cstddef>
 #include <cstdint>
 #include <type_traits>
+#include <utility>
 
 #include "Eigen/Core"
 
@@ -43,25 +44,38 @@ typedef double Vec2d __attribute__((vector_size(16)));
 typedef double Vec4d __attribute__((vector_size(32)));
 typedef double Vec8d __attribute__((vector_size(64)));
 
-// Corresponding integer types
-typedef uint32_t Vec4i __attribute__((vector_size(16)));
-typedef uint32_t Vec8i __attribute__((vector_size(32)));
+// Corresponding unsigned integer types
+typedef uint32_t Vec4u __attribute__((vector_size(16)));
+typedef uint32_t Vec8u __attribute__((vector_size(32)));
+typedef uint32_t Vec16u __attribute__((vector_size(64)));
 typedef uint64_t Vec2q __attribute__((vector_size(16)));
 typedef uint64_t Vec4q __attribute__((vector_size(32)));
 typedef uint64_t Vec8q __attribute__((vector_size(64)));
 
+// Corresponding signed integer types
+typedef int32_t Vec4i __attribute__((vector_size(16)));
+typedef int32_t Vec8i __attribute__((vector_size(32)));
+typedef int32_t Vec16i __attribute__((vector_size(64)));
+typedef int64_t Vec2l __attribute__((vector_size(16)));
+typedef int64_t Vec4l __attribute__((vector_size(32)));
+typedef int64_t Vec8l __attribute__((vector_size(64)));
+
 namespace internal {
-// Helper type to select the corresponding integer vector type.
+// Helper type to select the corresponding unsigned integer vector type.
 template <typename ScalarInt, size_t Width>
 struct MakeIntVec;
 
 template <>
 struct MakeIntVec<uint32_t, 4> {
-  using type = Vec4i;
+  using type = Vec4u;
 };
 template <>
 struct MakeIntVec<uint32_t, 8> {
-  using type = Vec8i;
+  using type = Vec8u;
+};
+template <>
+struct MakeIntVec<uint32_t, 16> {
+  using type = Vec16u;
 };
 template <>
 struct MakeIntVec<uint64_t, 2> {
@@ -76,7 +90,37 @@ struct MakeIntVec<uint64_t, 8> {
   using type = Vec8q;
 };
 
-// This trait takes a float vector and provides its integer vector partner.
+// Helper type to select the corresponding signed integer vector type.
+template <typename ScalarInt, size_t Width>
+struct MakeSignedIntVec;
+
+template <>
+struct MakeSignedIntVec<int32_t, 4> {
+  using type = Vec4i;
+};
+template <>
+struct MakeSignedIntVec<int32_t, 8> {
+  using type = Vec8i;
+};
+template <>
+struct MakeSignedIntVec<int32_t, 16> {
+  using type = Vec16i;
+};
+template <>
+struct MakeSignedIntVec<int64_t, 2> {
+  using type = Vec2l;
+};
+template <>
+struct MakeSignedIntVec<int64_t, 4> {
+  using type = Vec4l;
+};
+template <>
+struct MakeSignedIntVec<int64_t, 8> {
+  using type = Vec8l;
+};
+
+// This trait takes a float vector and provides its unsigned integer vector
+// partner.
 template <typename FloatVec>
 struct CorrespondingIntVector {
  private:
@@ -88,6 +132,20 @@ struct CorrespondingIntVector {
  public:
   using type = typename MakeIntVec<ScalarInt, kWidth>::type;
 };
+
+// This trait takes a float vector and provides its signed integer vector
+// partner.
+template <typename FloatVec>
+struct CorrespondingSignedIntVector {
+ private:
+  using ScalarFloat = decltype(FloatVec{}[0]);
+  static constexpr size_t kWidth = sizeof(FloatVec) / sizeof(ScalarFloat);
+  using ScalarInt =
+      std::conditional_t<sizeof(ScalarFloat) == 4, int32_t, int64_t>;
+
+ public:
+  using type = typename MakeSignedIntVec<ScalarInt, kWidth>::type;
+};
 }  // namespace internal
 
 // ===--------------------------------------------------------------------===//
@@ -133,6 +191,23 @@ struct ArrayMap<Vec8d> {
   using type = Eigen::Array<double, 8, 1>;
 };
 
+namespace internal {
+template <typename T, typename S, size_t... Is>
+constexpr T SplatImpl(S val, std::index_sequence<Is...>) {
+  return T{(static_cast<void>(Is), static_cast<decltype(T{}[0])>(val))...};
+}
+}  // namespace internal
+
+template <typename T, typename S>
+constexpr T Splat(S val) {
+  if constexpr (std::is_fundamental_v<T>) {
+    return static_cast<T>(val);
+  } else {
+    constexpr size_t kWidth = sizeof(T) / sizeof(decltype(T{}[0]));
+    return internal::SplatImpl<T>(val, std::make_index_sequence<kWidth>{});
+  }
+}
+
 // Computes the absolute value of a vector using bitwise operations.
 // FloatVec: The floating-point vector type (e.g., Vec4f).
 // x: The input vector.
@@ -140,7 +215,7 @@ struct ArrayMap<Vec8d> {
 template <typename FloatVec>
 FloatVec BitwiseAbs(FloatVec x) {
   using IntVec = typename internal::CorrespondingIntVector<FloatVec>::type;
-  // Get the underlying scalar integer type (e.g., int from Vec4i).
+  // Get the underlying scalar integer type (e.g., uint32_t from Vec4u).
   using ScalarInt = decltype(IntVec{}[0]);
 
   // Create a mask to clear the sign bit (e.g., 0x7FFFFFFF for int).
@@ -174,6 +249,48 @@ template <typename Vec, typename Scalar>
 Vec Clamp(Vec x, Scalar min, Scalar max) {
   return x < min ? min : x > max ? max : x;
 }
+
+// Evaluates a polynomial c_0 * x^n + c_1 * x^(n-1) + ... + c_n using Horner's
+// method, where `coeffs` are given in descending degree order.
+template <typename T, typename S, typename... Coeffs>
+constexpr T HornerPoly(T x, S c0, Coeffs... coeffs) {
+  T acc = Splat<T>(c0);
+  ((acc = acc * x + Splat<T>(coeffs)), ...);
+  return acc;
+}
+
+// Emits `@llvm.fmuladd` to compute `a * b + c`. When the target CPU has
+// hardware FMA, this lowers to a single fused multiply-add instruction;
+// without hardware FMA, `@llvm.fmuladd` falls back to unfused multiply/add.
+template <typename VecT>
+inline VecT Fmuladd(VecT a, VecT b, VecT c) {
+#pragma clang fp contract(on)
+  return a * b + c;
+}
+
+// Computes the error-free transformation `s + e = a + b` assuming `|a| >= |b|`
+// (Dekker's FastTwoSum), returning `{s, e}` where `s = RN(a + b)` and `e` is
+// the rounding error.
+template <typename VecT>
+inline std::pair<VecT, VecT> FastTwoSum(VecT a, VecT b) {
+  VecT s = a + b;
+  VecT e = (a - s) + b;
+  return {s, e};
+}
+
+template <typename VecT>
+inline auto ToSignedInt(VecT x) {
+  // Guard NaN lanes to 0 before float-to-int conversion (`fptosi` on NaN is UB
+  // and produces `poison`).
+  VecT safe_x = x == x ? x : VecT{};
+  if constexpr (std::is_arithmetic_v<VecT>) {
+    using IntT = std::conditional_t<sizeof(VecT) == 4, int32_t, int64_t>;
+    return static_cast<IntT>(safe_x);
+  } else {
+    using IntVec = typename internal::CorrespondingSignedIntVector<VecT>::type;
+    return __builtin_convertvector(safe_x, IntVec);
+  }
+}
 }  // namespace codegen
 }  // namespace xla
 
diff --git a/third_party/xla/xla/codegen/intrinsic_lib.cc b/third_party/xla/xla/codegen/intrinsic_lib.cc
index ffcf7315200..9a61ed2f7c3 100644
--- a/third_party/xla/xla/codegen/intrinsic_lib.cc
+++ b/third_party/xla/xla/codegen/intrinsic_lib.cc
@@ -124,22 +124,13 @@ class IntrinsicAdapter : public IntrinsicFunction {
 };
 
 IntrinsicFunctionLib::IntrinsicFunctionLib(const IntrinsicOptions& options)
-    : options_(options) {
-  if (options.device_type == intrinsics::DeviceType::kIntelCpu ||
-      options.device_type == intrinsics::DeviceType::kAmdCpu ||
-      options.device_type == intrinsics::DeviceType::kArmCpu) {
-    const std::string& ir_string = GetCppGenIrString(options);
-    if (!ir_string.empty()) {
-      auto eigen_lib =
-          std::make_unique<CppGenIntrinsicLibrary>(ir_string, "eigen");
-      ir_libraries_.push_back(std::move(eigen_lib));
-    }
-  }
-
+    : options_(options), ir_libraries_(GetCppGenLibraries(options)) {
   intrinsic_functions_.push_back(
       std::make_unique<IntrinsicAdapter<intrinsics::Ldexp>>());
   intrinsic_functions_.push_back(
       std::make_unique<IntrinsicAdapter<intrinsics::Exp>>());
+  intrinsic_functions_.push_back(
+      std::make_unique<IntrinsicAdapter<intrinsics::Expm1>>());
   intrinsic_functions_.push_back(
       std::make_unique<IntrinsicAdapter<intrinsics::FpTrunc>>());
   intrinsic_functions_.push_back(
@@ -292,7 +283,7 @@ IntrinsicFunctionLib::DefineIntrinsicFunctions(llvm::Module& module) {
         for (const auto& signature : signatures) {
           if (!ir_libraries_linked) {
             for (const auto& lib : ir_libraries_) {
-              lib->LinkIntoModule(module);
+              lib.LinkIntoModule(module);
             }
             ir_libraries_linked = true;
           }
diff --git a/third_party/xla/xla/codegen/intrinsic_lib.h b/third_party/xla/xla/codegen/intrinsic_lib.h
index 2c69dc29a31..7dde5ed4da3 100644
--- a/third_party/xla/xla/codegen/intrinsic_lib.h
+++ b/third_party/xla/xla/codegen/intrinsic_lib.h
@@ -63,7 +63,7 @@ class IntrinsicFunctionLib {
   std::vector<std::unique_ptr<IntrinsicFunction>> intrinsic_functions_;
   absl::flat_hash_map<absl::string_view, absl::string_view> targets_;
   const intrinsics::IntrinsicOptions options_;
-  std::vector<std::unique_ptr<CppGenIntrinsicLibrary>> ir_libraries_;
+  std::vector<CppGenIntrinsicLibrary> ir_libraries_;
 };
 
 }  // namespace xla::codegen
diff --git a/third_party/xla/xla/codegen/intrinsic_lib_test.cc b/third_party/xla/xla/codegen/intrinsic_lib_test.cc
index 18002be54c9..b2242c4ec49 100644
--- a/third_party/xla/xla/codegen/intrinsic_lib_test.cc
+++ b/third_party/xla/xla/codegen/intrinsic_lib_test.cc
@@ -56,7 +56,7 @@ TEST(IntrinsicLibTest, ExpVectorizations) {
   std::vector<llvm::VecDesc> vec_descs = lib.Vectorizations();
   std::vector<std::string> vec_descs_str;
   for (const auto& vec_desc : vec_descs) {
-    if (vec_desc.getScalarFnName().starts_with("xla.exp")) {
+    if (vec_desc.getScalarFnName().starts_with("xla.exp.")) {
       vec_descs_str.push_back(ToString(vec_desc));
     }
   }
@@ -134,6 +134,42 @@ TEST(IntrinsicLibTest, DefinesWideEigenAtan) {
   }
 }
 
+TEST(IntrinsicLibTest, Expm1Vectorizations) {
+  IntrinsicOptions options;
+  auto lib = IntrinsicFunctionLib(options);
+  std::vector<llvm::VecDesc> vec_descs = lib.Vectorizations();
+  std::vector<std::string> vec_descs_str;
+  for (const auto& vec_desc : vec_descs) {
+    if (vec_desc.getScalarFnName().starts_with("xla.expm1")) {
+      vec_descs_str.push_back(ToString(vec_desc));
+    }
+  }
+
+  EXPECT_THAT(
+      vec_descs_str,
+      UnorderedElementsAre("xla.expm1.f32:xla.expm1.v4f32:4:_ZGV_LLVM_N4v",
+                           "xla.expm1.f32:xla.expm1.v8f32:8:_ZGV_LLVM_N8v",
+                           "xla.expm1.f64:xla.expm1.v4f64:4:_ZGV_LLVM_N4v"));
+
+  IntrinsicOptions options_avx512;
+  options_avx512.features = "+avx512f";
+  auto lib_avx512 = IntrinsicFunctionLib(options_avx512);
+  std::vector<llvm::VecDesc> vec_descs_512 = lib_avx512.Vectorizations();
+  std::vector<std::string> vec_descs_512_str;
+  for (const auto& vec_desc : vec_descs_512) {
+    if (vec_desc.getScalarFnName().starts_with("xla.expm1")) {
+      vec_descs_512_str.push_back(ToString(vec_desc));
+    }
+  }
+
+  EXPECT_THAT(
+      vec_descs_512_str,
+      UnorderedElementsAre("xla.expm1.f32:xla.expm1.v4f32:4:_ZGV_LLVM_N4v",
+                           "xla.expm1.f32:xla.expm1.v8f32:8:_ZGV_LLVM_N8v",
+                           "xla.expm1.f32:xla.expm1.v16f32:16:_ZGV_LLVM_N16v",
+                           "xla.expm1.f64:xla.expm1.v4f64:4:_ZGV_LLVM_N4v",
+                           "xla.expm1.f64:xla.expm1.v8f64:8:_ZGV_LLVM_N8v"));
+}
 TEST(IntrinsicLibTest, CppGenIntrinsicLibraryPreservesNoInline) {
   llvm::LLVMContext context;
   llvm::Module dst_module("dst_module", context);
diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index d60312dc516..5f26fbc2818 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -4158,9 +4158,10 @@ cc_library(
         "//xla:util",
         "//xla:window_util",
         "//xla:xla_data_proto_cc",
-        "//xla/codegen/intrinsic",
         "//xla/codegen/intrinsic:fptrunc",
         "//xla/codegen/intrinsic:log1p",
+        "//xla/codegen/intrinsic:type",
+        "//xla/codegen/intrinsic/cpp:_intrinsic_declarations",
         "//xla/hlo/ir:hlo",
         "//xla/service/llvm_ir:ir_array",
         "//xla/service/llvm_ir:ir_builder_mixin",
diff --git a/third_party/xla/xla/service/cpu/tests/cpu_vectorization_test.cc b/third_party/xla/xla/service/cpu/tests/cpu_vectorization_test.cc
index f6a9a862a24..5489da59fb8 100644
--- a/third_party/xla/xla/service/cpu/tests/cpu_vectorization_test.cc
+++ b/third_party/xla/xla/service/cpu/tests/cpu_vectorization_test.cc
@@ -431,6 +431,56 @@ TEST_F(AtanJitVectorizationTest, AtanF32) {
                                /*match_optimized_ir=*/true));
 }
 
+class Expm1JitVectorizationTest : public HloHardwareIndependentTestBase {
+ protected:
+  DebugOptions GetDebugOptionsForTest() const override {
+    DebugOptions debug_options =
+        HloHardwareIndependentTestBase::GetDebugOptionsForTest();
+    debug_options.set_xla_cpu_max_isa("AVX512");
+    debug_options.set_xla_cpu_prefer_vector_width(512);
+    return debug_options;
+  }
+};
+
+TEST_F(Expm1JitVectorizationTest, Expm1F32) {
+  if (!tsl::port::IsX86CPU()) {
+    GTEST_SKIP() << "This feature only works for x86 CPUs.";
+  }
+
+  const absl::string_view hlo_text = R"(
+    HloModule Expm1F32
+
+    ENTRY Expm1F32 {
+      p0 = f32[1024] parameter(0)
+      ROOT expm1 = f32[1024] exponential-minus-one(p0)
+    }
+  )";
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> hlo_module,
+                       ParseAndReturnVerifiedModule(hlo_text));
+
+  auto compiler = GetCpuCompiler();
+  auto llvm_compiler = absl::down_cast<LLVMCompiler*>(compiler.get());
+  Compiler::CompileOptions compile_options;
+  compile_options.device_allocator = nullptr;
+
+  int num_elements = 16;
+  if (!tsl::port::TestCPUFeature(tsl::port::CPUFeature::AVX512F)) {
+    if (tsl::port::TestCPUFeature(tsl::port::CPUFeature::AVX2)) {
+      num_elements = 8;
+    } else {
+      num_elements = 4;
+    }
+  }
+
+  std::string check_lines =
+      absl::StrFormat("CHECK: fmul contract <%d x float>", num_elements);
+
+  ASSERT_OK(CompileAndVerifyIr(llvm_compiler, compile_options,
+                               std::move(hlo_module), check_lines,
+                               /*match_optimized_ir=*/true));
+}
+
 }  // namespace
 }  // namespace cpu
 }  // namespace xla
diff --git a/third_party/xla/xla/service/elemental_ir_emitter.cc b/third_party/xla/xla/service/elemental_ir_emitter.cc
index f5bface7646..3307adef641 100644
--- a/third_party/xla/xla/service/elemental_ir_emitter.cc
+++ b/third_party/xla/xla/service/elemental_ir_emitter.cc
@@ -51,9 +51,10 @@ limitations under the License.
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include "xla/codegen/intrinsic/cpp/intrinsic_declarations.h"
 #include "xla/codegen/intrinsic/fptrunc.h"
-#include "xla/codegen/intrinsic/intrinsic.h"
 #include "xla/codegen/intrinsic/log1p.h"
+#include "xla/codegen/intrinsic/type.h"
 #include "xla/comparison_util.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_computation.h"
@@ -2273,28 +2274,20 @@ absl::StatusOr<llvm::Value*> ElementalIrEmitter::EmitExp2(
 
 absl::StatusOr<llvm::Value*> ElementalIrEmitter::EmitExpm1(
     PrimitiveType prim_type, llvm::Value* value) {
-  auto x = value;
-  auto type = llvm_ir::PrimitiveTypeToIrType(prim_type, module_->getContext());
-  auto one = llvm::ConstantFP::get(type, 1.0);
-  auto half = llvm::ConstantFP::get(type, 0.5);
-  auto zero = llvm::ConstantFP::get(type, 0.0);
-
-  // expm1(x) == tanh(x/2)*(exp(x)+1)
-  // x/2 can underflow, if it does we approximate expm1 with x.
-  auto x_over_two = FMul(x, half);
-  auto x_over_two_is_zero = FCmpOEQ(x_over_two, zero);
-  auto abs_x =
-      llvm_ir::EmitCallToIntrinsic(llvm::Intrinsic::fabs, {x}, {type}, b_);
-  // Use a naive exp(x)-1 calculation if |x| is > 0.5
-  auto x_magnitude_is_large = FCmpOGT(abs_x, half);
-  ABSL_ASSIGN_OR_RETURN(auto tanh_of_x_over_two, EmitTanh(prim_type, x_over_two));
-  ABSL_ASSIGN_OR_RETURN(auto exp_of_x, EmitExp(prim_type, x, ""));
-  auto exp_of_x_plus_one = FAdd(exp_of_x, one);
-  auto exp_of_x_minus_one = FSub(exp_of_x, one);
-  auto expm1_of_x = FMul(tanh_of_x_over_two, exp_of_x_plus_one);
-  expm1_of_x = Select(x_magnitude_is_large, exp_of_x_minus_one, expm1_of_x);
-  expm1_of_x = Select(x_over_two_is_zero, x, expm1_of_x);
-  return expm1_of_x;
+  if (prim_type == F32 || prim_type == F64) {
+    llvm::Function* expm1 = codegen::intrinsics::Expm1::GetOrInsertDeclaration(
+        module_, IntrinsicType::S(prim_type));
+    return b_->CreateCall(expm1, {value});
+  }
+  if (prim_type == F16 || prim_type == BF16) {
+    llvm::Type* f32_type = b_->getFloatTy();
+    llvm::Value* f32_value = b_->CreateFPExt(value, f32_type);
+    llvm::Function* expm1 = codegen::intrinsics::Expm1::GetOrInsertDeclaration(
+        module_, IntrinsicType::S(F32));
+    llvm::Value* f32_result = b_->CreateCall(expm1, {f32_value});
+    return b_->CreateFPTrunc(f32_result, value->getType());
+  }
+  return Unimplemented("expm1");
 }
 
 absl::StatusOr<llvm::Value*> ElementalIrEmitter::EmitPow(
diff --git a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_functions.cc b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_functions.cc
index d3b1f731f92..9f8eef54404 100644
--- a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_functions.cc
+++ b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_functions.cc
@@ -81,7 +81,13 @@ UNARY_TEST(Exp, {
       })
       .Run();
 })
-UNARY_TEST(Expm1, { Expm1Op<kT>(this).Error(GetDefaultSpecGenerator()).Run(); })
+UNARY_TEST(Expm1, {
+  Expm1Op<kT>(this)
+      .CpuError(
+          +[](NativeT) { return ErrorSpec::Builder().distance_err(1).build(); })
+      .Error(GetDefaultSpecGenerator())
+      .Run();
+})
 UNARY_TEST(Exp2, { Exp2Op<kT>(this).Error(GetDefaultSpecGenerator()).Run(); })
 
 UNARY_TEST(Logistic, {

From 91ba0b5f2508a6492053d9577baa473ff0d0086e Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 13 Oct 2026 12:10:47 -0700
Subject: [PATCH 0711/1000] [XLA:GPU] Guard compilation_pipeline
 runtime_tileiras with if_cuda.

- Guard the compilation_pipeline runtime_tileiras data dependency behind if_cuda.

Ensures builds not configured to run with cuda don't pull in cuda dependencies.

PiperOrigin-RevId: 998804331
---
 third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD
index d8230292686..b6c6daed396 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD
@@ -1,3 +1,4 @@
+load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda")
 load("@rules_cc//cc:cc_library.bzl", "cc_library")
 load("//xla:xla.default.bzl", "xla_cc_test")
 load("//xla/tsl:tsl.bzl", "if_google")
@@ -73,7 +74,7 @@ cc_library(
     hdrs = ["compilation_pipeline.h"],
     # tileiras, ptxas and libnvvm staged from the CUDA 13.3+ toolkit;
     # FindCudaExecutable("tileiras") picks them up from runfiles.
-    data = if_google(["@local_config_cuda//cuda:runtime_tileiras"]),
+    data = if_google(if_cuda(["@local_config_cuda//cuda:runtime_tileiras"])),
     deps = [
         "//xla/stream_executor:device_description",
         "//xla/stream_executor/cuda:cuda_compute_capability",

From b7798e3e187514c7c3f59c852208444343cec2f2 Mon Sep 17 00:00:00 2001
From: Bixia Zheng <bixia@google.com>
Date: Tue, 13 Oct 2026 12:26:01 -0700
Subject: [PATCH 0712/1000] Canonicalize layouts after Shardy device-local code
 generation.

When converting global code to device-local code , invoke CanonicalizeLayoutAfterShardingPropagation to update the entry computation layout. Parameter and output layouts are canonicalized if their shapes have been changed.

PiperOrigin-RevId: 998814023
---
 third_party/xla/xla/service/spmd/shardy/BUILD |  2 +-
 .../service/spmd/shardy/shardy_xla_pass.cc    | 62 +++++++++++++++----
 .../spmd/shardy/shardy_xla_pass_test.cc       | 60 ++++++++++++++++++
 3 files changed, 110 insertions(+), 14 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/shardy/BUILD b/third_party/xla/xla/service/spmd/shardy/BUILD
index c387d7b6729..7dce0f393c5 100644
--- a/third_party/xla/xla/service/spmd/shardy/BUILD
+++ b/third_party/xla/xla/service/spmd/shardy/BUILD
@@ -58,7 +58,6 @@ cc_library(
         "//xla/service/spmd/shardy/stablehlo_round_trip:stablehlo_import",
         "//xla/tsl/framework/mlir:status_scoped_diagnostic_handler",
         "//xla/tsl/platform:env",
-        "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
@@ -148,6 +147,7 @@ xla_cc_test(
         "//xla/service:computation_layout",
         "//xla/service/spmd:sharding_format_picker",
         "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_googletest//:gtest_main",
         "@shardy//shardy/dialect/sdy/transforms/common:partitioner_stage",
         "@shardy//shardy/dialect/sdy/transforms/common:propagation_options",
diff --git a/third_party/xla/xla/service/spmd/shardy/shardy_xla_pass.cc b/third_party/xla/xla/service/spmd/shardy/shardy_xla_pass.cc
index f1ed9e0fc12..7b12619d15f 100644
--- a/third_party/xla/xla/service/spmd/shardy/shardy_xla_pass.cc
+++ b/third_party/xla/xla/service/spmd/shardy/shardy_xla_pass.cc
@@ -74,7 +74,6 @@ limitations under the License.
 #include "xla/status_macros.h"
 #include "xla/tsl/framework/mlir/status_scoped_diagnostic_handler.h"
 #include "xla/tsl/platform/env.h"
-#include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"
 #include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
@@ -223,6 +222,53 @@ ComputationLayout getFlattenedComputationLayout(
   return flattenedComputationLayout;
 }
 
+// Updates the entry computation layout of `hloModule` after Shardy device-local
+// code generation by copying layouts from `oldLayout` to the new local program
+// shape and canonicalizing layouts for any parameter or output whose shape
+// changed during partitioning.
+absl::Status updateEntryComputationLayout(HloModule* hloModule,
+                                          const ComputationLayout& oldLayout) {
+  ProgramShape newProgramShape =
+      hloModule->entry_computation()->ComputeProgramShape();
+
+  // Record which parameters changed shape and copy the original parameter
+  // layouts onto the new local parameter shapes.
+  mlir::SmallVector<bool> updateParametersLayout(
+      newProgramShape.parameters_size());
+  for (int64_t i = 0; i < newProgramShape.parameters_size(); ++i) {
+    const Shape& oldParamShape = oldLayout.parameter_shape(i);
+    Shape* newParamShape = newProgramShape.mutable_parameters(i);
+    updateParametersLayout[i] =
+        !Shape::Equal().IgnoreLayout()(oldParamShape, *newParamShape);
+    ABSL_RETURN_IF_ERROR(
+        LayoutUtil::CopyLayoutBetweenShapes(oldParamShape, newParamShape));
+  }
+
+  // Record which outputs changed shape and copy the original result layout onto
+  // the new local result shape.
+  const Shape& oldResultShape = oldLayout.result_shape();
+  Shape* newResultShape = newProgramShape.mutable_result();
+  mlir::SmallVector<bool> updateOutputLayout;
+  if (newResultShape->IsTuple()) {
+    updateOutputLayout.reserve(newResultShape->tuple_shapes().size());
+    for (int64_t i = 0; i < newResultShape->tuple_shapes().size(); ++i) {
+      updateOutputLayout.push_back(!Shape::Equal().IgnoreLayout()(
+          oldResultShape.tuple_shapes(i), newResultShape->tuple_shapes(i)));
+    }
+  } else {
+    updateOutputLayout.push_back(
+        !Shape::Equal().IgnoreLayout()(oldResultShape, *newResultShape));
+  }
+  ABSL_RETURN_IF_ERROR(
+      LayoutUtil::CopyLayoutBetweenShapes(oldResultShape, newResultShape));
+  *hloModule->mutable_entry_computation_layout() =
+      ComputationLayout(newProgramShape, /*ignore_layouts=*/false);
+
+  // Give the changed parameter and output shapes canonical layouts.
+  return hlo_sharding_util::CanonicalizeLayoutAfterShardingPropagation(
+      hloModule, updateOutputLayout, updateParametersLayout);
+}
+
 // Returns the corresponding parameter number and index dependent on whether the
 // returned HLO module will have a single tuple level tuple.
 // - If `useTupleArgs` is true, then the param number will be 0 since there is
@@ -579,18 +625,8 @@ absl::StatusOr<bool> ShardyXLA::RunImpl(
       hloModule->set_spmd_parameters_shardings(std::move(param_shardings));
     }
 
-    ProgramShape newProgramShape =
-        hloModule->entry_computation()->ComputeProgramShape();
-    for (int64_t i = 0; i < newProgramShape.parameters_size(); ++i) {
-      ABSL_RETURN_IF_ERROR(LayoutUtil::CopyLayoutBetweenShapes(
-          flattenedEntryComputationLayout.parameter_shape(i),
-          newProgramShape.mutable_parameters(i)));
-    }
-    ABSL_RETURN_IF_ERROR(LayoutUtil::CopyLayoutBetweenShapes(
-        flattenedEntryComputationLayout.result_shape(),
-        newProgramShape.mutable_result()));
-    *hloModule->mutable_entry_computation_layout() =
-        ComputationLayout(newProgramShape, /*ignore_layouts=*/false);
+    ABSL_RETURN_IF_ERROR(updateEntryComputationLayout(
+        hloModule, flattenedEntryComputationLayout));
   } else {
     // Restore entry computation layout.
     *hloModule->mutable_entry_computation_layout() =
diff --git a/third_party/xla/xla/service/spmd/shardy/shardy_xla_pass_test.cc b/third_party/xla/xla/service/spmd/shardy/shardy_xla_pass_test.cc
index a3a1e2b01db..d72aec0a28f 100644
--- a/third_party/xla/xla/service/spmd/shardy/shardy_xla_pass_test.cc
+++ b/third_party/xla/xla/service/spmd/shardy/shardy_xla_pass_test.cc
@@ -17,14 +17,18 @@ limitations under the License.
 
 #include <memory>
 #include <string>
+#include <utility>
+#include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/log/log.h"
+#include "absl/status/statusor.h"
 #include "shardy/dialect/sdy/transforms/common/partitioner_stage.h"
 #include "shardy/dialect/sdy/transforms/common/propagation_options.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/ir/hlo_print_options.h"
 #include "xla/hlo/ir/hlo_sharding.h"
@@ -32,6 +36,8 @@ limitations under the License.
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/hlo/testlib/verified_hlo_module.h"
 #include "xla/hlo/utils/hlo_matchers.h"
+#include "xla/layout.h"
+#include "xla/layout_util.h"
 #include "xla/service/computation_layout.h"
 #include "xla/service/spmd/sharding_format_picker.h"
 #include "xla/service/spmd/shardy/constants.h"
@@ -1108,6 +1114,60 @@ TEST_P(ShardyXLATest, ShardyGenerateDeviceLocalCodeUseTupleArgs) {
             "f32[4,16]{1,0:T(8,128)}))->f32[4,4]{1,0:T(8,128)}");
 }
 
+TEST_P(ShardyXLATest, ShardyGenerateDeviceLocalCodeCanonicalizeLayout) {
+  const char* const hloString = R"(
+    HloModule pjit_f,
+      entry_computation_layout={(f32[8,16]{1,0:T(8,128)}, f32[8,16]{0,1:T(8,128)}, f32[8,16]{0,1:T(8,128)})->(f32[8,16]{1,0:T(8,128)}, f32[8,16]{0,1:T(8,128)})},
+      allow_spmd_sharding_propagation_to_parameters={false,false,true},
+      allow_spmd_sharding_propagation_to_output={false,false},
+      num_partitions=8
+
+    ENTRY %main (Arg_0: f32[8,16], Arg_1: f32[8,16], Arg_2: f32[8,16]) -> (f32[8,16], f32[8,16]) {
+      %Arg_0 = f32[8,16]{1,0} parameter(0), sharding={devices=[2,4]<=[8]}
+      %Arg_1 = f32[8,16]{0,1} parameter(1)
+      %Arg_2 = f32[8,16]{0,1} parameter(2)
+      %add_0 = f32[8,16]{1,0} add(%Arg_0, %Arg_2), sharding={devices=[2,4]<=[8]}
+      ROOT %tuple = (f32[8,16]{1,0}, f32[8,16]{0,1}) tuple(%add_0, %Arg_1)
+    })";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(hloString));
+  module->set_layout_canonicalization_callback(
+      [](const HloModule& m)
+          -> absl::StatusOr<std::pair<std::vector<Shape>, Shape>> {
+        auto chooseCompactLayout = [](Shape shape) {
+          *shape.mutable_layout() = LayoutUtil::MakeLayout(
+              {1, 0}, {Tile({shape.dimensions(0), shape.dimensions(1)})});
+          return shape;
+        };
+        std::vector<Shape> paramShapes;
+        for (const HloInstruction* param :
+             m.entry_computation()->parameter_instructions()) {
+          paramShapes.push_back(chooseCompactLayout(param->shape()));
+        }
+        Shape resultShape = m.entry_computation()->root_instruction()->shape();
+        for (Shape& subshape : *resultShape.mutable_tuple_shapes()) {
+          subshape = chooseCompactLayout(subshape);
+        }
+        return std::make_pair(std::move(paramShapes), std::move(resultShape));
+      });
+
+  mlir::sdy::PropagationOptions propagationOptions;
+  propagationOptions.partitionerStage =
+      mlir::sdy::PartitionerStage::kConvertGlobalToLocal;
+  runShardyWithStablehloImport(module.get(),
+                               /*runSdyShardingPropagation=*/true,
+                               /*expectChanged=*/true, propagationOptions);
+
+  // Arg_0, Arg_2 (via propagation), and output element 0 shapes changed from
+  // f32[8,16] to f32[4,4], so their layouts are canonicalized to {1,0:T(4,4)}.
+  // Arg_1 and output element 1 shapes did not change, so their preset
+  // {0,1:T(8,128)} layouts are preserved.
+  EXPECT_EQ(module->entry_computation_layout().ToString(),
+            "(f32[4,4]{1,0:T(4,4)}, f32[8,16]{0,1:T(8,128)}, "
+            "f32[4,4]{1,0:T(4,4)})->(f32[4,4]{1,0:T(4,4)}, "
+            "f32[8,16]{0,1:T(8,128)})");
+}
+
 TEST_P(ShardyXLATest, TestMaximalShardingNoResults) {
   const char* const hloString = R"(
 HloModule maximal_sharding_module, entry_computation_layout={(s64[2]{0})->s64[2]{0}},

From 254e5e0aa3de3d60df295a10c75fcb5c0c9a6d91 Mon Sep 17 00:00:00 2001
From: Sean Talts <talts@google.com>
Date: Tue, 13 Oct 2026 12:39:00 -0700
Subject: [PATCH 0713/1000] [XLA:CPU] Don't hoist while loops containing dots
 or convolutions that need Eigen or unsupported types

This cl just plumbs through target machine features a little more and reduces which while loops we hoist now that we set xla_cpu_small_while_loop_byte_threshold to 1e9 in a recent CL. This fixes tests that broke internally and the net performance effects should be ~neutral on our benchmarks, since those models this affects would just fail outright.

Context
SmallWhileLoopHoistingPass outlines while loops into calls that are compiled
as a single kernel with runtime calls disabled. Dots and convolutions that
would otherwise call Eigen then fall back to a naive loop nest, and those with
operand types that IrEmitter::HandleDot/HandleConvolution reject (e.g. BF16)
fail to compile. With the higher default byte threshold this affects most loops
containing large GEMMs or convolutions.

Treat such instructions as unavailable for hoisting, using the same Eigen
eligibility checks as the emitters.

PiperOrigin-RevId: 998822551
---
 third_party/xla/xla/service/cpu/BUILD         |  6 ++
 .../xla/xla/service/cpu/cpu_compiler.cc       |  3 +-
 .../cpu/small_while_loop_hoisting_pass.cc     | 46 ++++++++--
 .../cpu/small_while_loop_hoisting_pass.h      |  6 +-
 .../small_while_loop_hoisting_pass_test.cc    | 92 ++++++++++++++++++-
 5 files changed, 143 insertions(+), 10 deletions(-)

diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD
index d3f9f52b2d7..cbc6da2e898 100644
--- a/third_party/xla/xla/service/cpu/BUILD
+++ b/third_party/xla/xla/service/cpu/BUILD
@@ -1693,7 +1693,10 @@ cc_library(
     hdrs = ["small_while_loop_hoisting_pass.h"],
     deps = [
         ":cpu_executable",
+        ":dot_op_emitter",
+        ":ir_emission_utils",
         "//xla:util",
+        "//xla/backends/cpu/codegen:target_machine_features",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass",
         "//xla/service:collective_ops_utils",
@@ -1796,10 +1799,13 @@ xla_cc_test(
     deps = [
         ":backend_config_proto_cc",
         ":small_while_loop_hoisting_pass",
+        ":target_machine_features_stub",
+        "//xla/backends/cpu/codegen:target_machine_features",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
         "//xla/hlo/testlib:test",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest_main",
     ],
diff --git a/third_party/xla/xla/service/cpu/cpu_compiler.cc b/third_party/xla/xla/service/cpu/cpu_compiler.cc
index a1acec04a1d..dccc220b678 100644
--- a/third_party/xla/xla/service/cpu/cpu_compiler.cc
+++ b/third_party/xla/xla/service/cpu/cpu_compiler.cc
@@ -1207,7 +1207,8 @@ absl::Status CpuCompiler::RunHloPassesAfterLayoutAssn(
     ABSL_ASSIGN_OR_RETURN(
         int64_t byte_threshold,
         xla::cpu::options::SmallWhileLoopByteThreshold(module->config()));
-    pipeline.AddPass<SmallWhileLoopHoistingPass>(byte_threshold);
+    pipeline.AddPass<SmallWhileLoopHoistingPass>(byte_threshold,
+                                                 target_machine_features);
   }
 
   pipeline.AddPass<PropagateCallMetadata>();
diff --git a/third_party/xla/xla/service/cpu/small_while_loop_hoisting_pass.cc b/third_party/xla/xla/service/cpu/small_while_loop_hoisting_pass.cc
index bba1058e3b4..d53988bdfa4 100644
--- a/third_party/xla/xla/service/cpu/small_while_loop_hoisting_pass.cc
+++ b/third_party/xla/xla/service/cpu/small_while_loop_hoisting_pass.cc
@@ -26,12 +26,15 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
+#include "xla/backends/cpu/codegen/target_machine_features.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/service/collective_ops_utils.h"
 #include "xla/service/cpu/cpu_executable.h"
+#include "xla/service/cpu/dot_op_emitter.h"
+#include "xla/service/cpu/ir_emission_utils.h"
 #include "xla/service/hlo_cost_analysis.h"
 #include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/statusor.h"
@@ -39,7 +42,21 @@ limitations under the License.
 
 namespace xla::cpu {
 
-static bool InstructionIsUnavailable(const HloInstruction* instr) {
+// The call thunk emitter compiles the loop with runtime calls disabled, so dots
+// and convolutions that would otherwise call Eigen fall back to a naive loop
+// nest. IrEmitter::HandleDot and HandleConvolution also reject operand types
+// that the thunk emitter supports (e.g. BF16).
+static bool OperandTypesUnsupported(const HloInstruction* instr) {
+  constexpr PrimitiveType kSupportedTypes[] = {
+      PRED, S8, U8, S16, U16, S32, U32, S64, U64, F16, F32, F64, C64, C128};
+  PrimitiveType lhs_type = instr->operand(0)->shape().element_type();
+  return lhs_type != instr->operand(1)->shape().element_type() ||
+         !absl::c_linear_search(kSupportedTypes, lhs_type);
+}
+
+static bool InstructionIsUnavailable(
+    const HloInstruction* instr,
+    const TargetMachineFeatures& target_machine_features) {
   switch (instr->opcode()) {
     // The following instructions are not currently supported by the call thunk
     // emitter due to how the legacy & thunk emitters interact; specifically,
@@ -58,6 +75,18 @@ static bool InstructionIsUnavailable(const HloInstruction* instr) {
     case HloOpcode::kFusion:
       return instr->fusion_kind() == HloInstruction::FusionKind::kCustom;
 
+    case HloOpcode::kDot:
+      return OperandTypesUnsupported(instr) ||
+             GetDotImplementationStrategy(instr->GetModule()->config(), *instr,
+                                          target_machine_features,
+                                          /*allow_runtime_calls=*/true) ==
+                 DotImplementationStrategy::kEigen;
+
+    case HloOpcode::kConvolution:
+      return OperandTypesUnsupported(instr) ||
+             PotentiallyImplementedAsEigenConvolution(*instr,
+                                                      target_machine_features);
+
     default:
       return IsCollective(instr);
   }
@@ -67,19 +96,21 @@ static bool InstructionIsUnavailable(const HloInstruction* instr) {
 // "unavailable" for the call thunk emitter.
 static bool ContainsUnavailableInstruction(
     const HloInstruction* instr,
+    const TargetMachineFeatures& target_machine_features,
     absl::flat_hash_map<const HloInstruction*, bool>& has_unavailable_instr) {
   if (const auto itr = has_unavailable_instr.find(instr);
       itr != has_unavailable_instr.end()) {
     return itr->second;
   }
 
-  if (InstructionIsUnavailable(instr)) {
+  if (InstructionIsUnavailable(instr, target_machine_features)) {
     return has_unavailable_instr.insert({instr, true}).first->second;
   }
 
   for (const HloComputation* called_comp : instr->called_computations()) {
     for (const HloInstruction* sub_instr : called_comp->instructions()) {
-      if (ContainsUnavailableInstruction(sub_instr, has_unavailable_instr)) {
+      if (ContainsUnavailableInstruction(sub_instr, target_machine_features,
+                                         has_unavailable_instr)) {
         return has_unavailable_instr.insert({instr, true}).first->second;
       }
     }
@@ -89,8 +120,10 @@ static bool ContainsUnavailableInstruction(
 }
 
 SmallWhileLoopHoistingPass::SmallWhileLoopHoistingPass(
-    int64_t small_buffer_access_size)
-    : small_buffer_access_size_(small_buffer_access_size) {}
+    int64_t small_buffer_access_size,
+    const TargetMachineFeatures* target_machine_features)
+    : small_buffer_access_size_(small_buffer_access_size),
+      target_machine_features_(target_machine_features) {}
 
 absl::StatusOr<bool> SmallWhileLoopHoistingPass::RunImpl(
     HloModule* module,
@@ -104,7 +137,8 @@ absl::StatusOr<bool> SmallWhileLoopHoistingPass::RunImpl(
   bool changed = false;
   absl::flat_hash_map<const HloInstruction*, bool> has_unavailable_instr;
   for (HloInstruction* while_instr : while_instrs) {
-    if (ContainsUnavailableInstruction(while_instr, has_unavailable_instr)) {
+    if (ContainsUnavailableInstruction(while_instr, *target_machine_features_,
+                                       has_unavailable_instr)) {
       continue;
     }
 
diff --git a/third_party/xla/xla/service/cpu/small_while_loop_hoisting_pass.h b/third_party/xla/xla/service/cpu/small_while_loop_hoisting_pass.h
index 5ee2e37fa02..f20487db8e0 100644
--- a/third_party/xla/xla/service/cpu/small_while_loop_hoisting_pass.h
+++ b/third_party/xla/xla/service/cpu/small_while_loop_hoisting_pass.h
@@ -21,6 +21,7 @@ limitations under the License.
 #include "absl/container/flat_hash_set.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
+#include "xla/backends/cpu/codegen/target_machine_features.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/pass/hlo_pass_interface.h"
@@ -31,7 +32,9 @@ namespace xla::cpu {
 // kernel instead of a series of thunks.
 class SmallWhileLoopHoistingPass final : public HloModulePass {
  public:
-  explicit SmallWhileLoopHoistingPass(int64_t small_buffer_access_size);
+  SmallWhileLoopHoistingPass(
+      int64_t small_buffer_access_size,
+      const TargetMachineFeatures* target_machine_features);
 
   absl::string_view name() const final { return "small-while-loop-hoisting"; }
 
@@ -45,6 +48,7 @@ class SmallWhileLoopHoistingPass final : public HloModulePass {
 
  private:
   int64_t small_buffer_access_size_;
+  const TargetMachineFeatures* target_machine_features_;
 };
 }  // namespace xla::cpu
 
diff --git a/third_party/xla/xla/service/cpu/small_while_loop_hoisting_pass_test.cc b/third_party/xla/xla/service/cpu/small_while_loop_hoisting_pass_test.cc
index cb237be4ac2..5a8519bcfb6 100644
--- a/third_party/xla/xla/service/cpu/small_while_loop_hoisting_pass_test.cc
+++ b/third_party/xla/xla/service/cpu/small_while_loop_hoisting_pass_test.cc
@@ -15,28 +15,39 @@ limitations under the License.
 
 #include "xla/service/cpu/small_while_loop_hoisting_pass.h"
 
+#include <cstdint>
 #include <memory>
 #include <optional>
 #include <string>
 
 #include <gtest/gtest.h>
 #include "absl/status/statusor.h"
+#include "absl/strings/str_replace.h"
 #include "absl/strings/string_view.h"
+#include "xla/backends/cpu/codegen/target_machine_features.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/hlo/testlib/test.h"
 #include "xla/service/cpu/backend_config.pb.h"
+#include "xla/service/cpu/target_machine_features_stub.h"
 
 namespace xla {
 namespace {
 
 class SmallWhileLoopHoistingPassTest : public HloHardwareIndependentTestBase {
  protected:
-  absl::StatusOr<bool> RunSmallWhileLoopHoistingPass(HloModule* module) {
-    return cpu::SmallWhileLoopHoistingPass(1024).Run(module);
+  absl::StatusOr<bool> RunSmallWhileLoopHoistingPass(
+      HloModule* module, int64_t small_buffer_access_size = 1024) {
+    return cpu::SmallWhileLoopHoistingPass(small_buffer_access_size,
+                                           &target_machine_features_)
+        .Run(module);
   }
+
+  cpu::TargetMachineFeaturesStub target_machine_features_{[](int64_t) {
+    return cpu::TargetMachineFeatures::kEigenExpectedTensorAlignment;
+  }};
 };
 
 TEST_F(SmallWhileLoopHoistingPassTest, SmallWhileLoopHoisting) {
@@ -263,5 +274,82 @@ TEST_F(SmallWhileLoopHoistingPassTest, NoYnnWhileLoopHoisting) {
   EXPECT_FALSE(changed);
 }
 
+TEST_F(SmallWhileLoopHoistingPassTest, DotAndConvolutionWhileLoopHoisting) {
+  constexpr absl::string_view kHloTemplate = R"(
+    HloModule while_loop
+
+    while_body {
+      state = (s32[], $lhs, $rhs) parameter(0)
+      counter = s32[] get-tuple-element(state), index=0
+      one = s32[] constant(1)
+      next = s32[] add(counter, one)
+      lhs = $lhs get-tuple-element(state), index=1
+      rhs = $rhs get-tuple-element(state), index=2
+      op = $op
+      ROOT tuple = (s32[], $lhs, $rhs) tuple(next, lhs, rhs)
+    }
+
+    while_condition {
+      state = (s32[], $lhs, $rhs) parameter(0)
+      counter = s32[] get-tuple-element(state), index=0
+      limit = s32[] constant(10)
+      ROOT less_than = pred[] compare(counter, limit), direction=LT
+    }
+
+    ENTRY main {
+      zero = s32[] constant(0)
+      lhs = $lhs parameter(0)
+      rhs = $rhs parameter(1)
+      init = (s32[], $lhs, $rhs) tuple(zero, lhs, rhs)
+      ROOT while_loop = (s32[], $lhs, $rhs) while(init),
+          condition=while_condition, body=while_body
+    }
+    )";
+
+  struct TestCase {
+    absl::string_view lhs;
+    absl::string_view rhs;
+    absl::string_view op;
+    bool expect_hoisted;
+  };
+  // Only the matrix-vector product is emitted inline without Eigen and in a
+  // type the legacy emitter supports.
+  const TestCase kTestCases[] = {
+      {"f32[256,256]", "f32[256,1]",
+       "f32[256,1] dot(lhs, rhs), lhs_contracting_dims={1}, "
+       "rhs_contracting_dims={0}",
+       true},
+      {"f32[512,256]", "f32[256,256]",
+       "f32[512,256] dot(lhs, rhs), lhs_contracting_dims={1}, "
+       "rhs_contracting_dims={0}",
+       false},
+      {"bf16[256,256]", "bf16[256,1]",
+       "bf16[256,1] dot(lhs, rhs), lhs_contracting_dims={1}, "
+       "rhs_contracting_dims={0}",
+       false},
+      {"f32[8,512,256]", "f32[8,256,256]",
+       "f32[8,512,256] dot(lhs, rhs), lhs_batch_dims={0}, "
+       "lhs_contracting_dims={2}, rhs_batch_dims={0}, rhs_contracting_dims={1}",
+       false},
+      {"f32[1,32,32,16]", "f32[3,3,16,16]",
+       "f32[1,32,32,16] convolution(lhs, rhs), window={size=3x3 pad=1_1x1_1}, "
+       "dim_labels=b01f_01io->b01f",
+       false},
+  };
+  for (const TestCase& test_case : kTestCases) {
+    SCOPED_TRACE(test_case.op);
+    std::string hlo_string =
+        absl::StrReplaceAll(kHloTemplate, {{"$lhs", test_case.lhs},
+                                           {"$rhs", test_case.rhs},
+                                           {"$op", test_case.op}});
+    ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> m,
+                         ParseAndReturnVerifiedModule(hlo_string));
+    ASSERT_OK_AND_ASSIGN(bool changed,
+                         RunSmallWhileLoopHoistingPass(
+                             m.get(), /*small_buffer_access_size=*/1000000000));
+    EXPECT_EQ(changed, test_case.expect_hoisted);
+  }
+}
+
 }  // namespace
 }  // namespace xla

From 2ccdc423bc76e9c8b178cdbacd4a68b5aa6f279a Mon Sep 17 00:00:00 2001
From: greypilgrim-083 <prashantchaudhary7353@gmail.com>
Date: Wed, 14 Oct 2026 02:10:54 +0530
Subject: [PATCH 0714/1000] Address review: restore gtest_main and add direct
 header dependencies to sliding_window_dataset_op_test

- Re-add @com_google_googletest//:gtest_main to deps
- Add //tensorflow/core/framework:types_proto_cc (for types.pb.h)
- Add //tensorflow/core/platform:types (for platform/types.h)
---
 tensorflow/core/kernels/data/experimental/BUILD | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensorflow/core/kernels/data/experimental/BUILD b/tensorflow/core/kernels/data/experimental/BUILD
index 33051088c08..54c6e5b0e82 100644
--- a/tensorflow/core/kernels/data/experimental/BUILD
+++ b/tensorflow/core/kernels/data/experimental/BUILD
@@ -865,9 +865,12 @@ tf_cc_test(
         "//tensorflow/core:test_main",
         "//tensorflow/core/data:dataset_test_base",
         "//tensorflow/core/data:name_utils",
+        "//tensorflow/core/framework:types_proto_cc",
         "//tensorflow/core/kernels/data:range_dataset_op",
         "//tensorflow/core/platform:status",
+        "//tensorflow/core/platform:types",
         "@com_google_absl//absl/status",
+        "@com_google_googletest//:gtest_main",
         "@xla//xla/tsl/lib/core:status_test_util",
         "@xla//xla/tsl/platform:status_matchers",
     ],

From 5d0718098a02fad818408b60cd69a2e891766b68 Mon Sep 17 00:00:00 2001
From: Changming Sun CA <changmingsun@google.com>
Date: Tue, 13 Oct 2026 13:16:17 -0700
Subject: [PATCH 0715/1000] Update XNNPACK in TensorFlow/XLA

PiperOrigin-RevId: 998846435
---
 tensorflow/lite/tools/cmake/modules/xnnpack.cmake  | 2 +-
 tensorflow/workspace2.bzl                          | 6 +++---
 third_party/xla/third_party/kleidiai/workspace.bzl | 6 +++---
 third_party/xla/third_party/slinky/workspace.bzl   | 6 +++---
 third_party/xla/third_party/xnnpack/workspace.bzl  | 6 +++---
 5 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/tensorflow/lite/tools/cmake/modules/xnnpack.cmake b/tensorflow/lite/tools/cmake/modules/xnnpack.cmake
index d7cf1bd6fcd..444ddf1af6b 100644
--- a/tensorflow/lite/tools/cmake/modules/xnnpack.cmake
+++ b/tensorflow/lite/tools/cmake/modules/xnnpack.cmake
@@ -23,7 +23,7 @@ OverridableFetchContent_Declare(
   xnnpack
   GIT_REPOSITORY https://github.com/google/XNNPACK
   # Sync with tensorflow/workspace2.bzl
-  GIT_TAG d89ef6669a14db203b3b7935b1b3862cb63fb6df
+  GIT_TAG 0a66ebb5817fce48df605451444dcfc27a5e4ee4
   GIT_PROGRESS TRUE
   PREFIX "${CMAKE_BINARY_DIR}"
   SOURCE_DIR "${CMAKE_BINARY_DIR}/xnnpack"
diff --git a/tensorflow/workspace2.bzl b/tensorflow/workspace2.bzl
index da028662d57..1b078dba9c1 100644
--- a/tensorflow/workspace2.bzl
+++ b/tensorflow/workspace2.bzl
@@ -265,9 +265,9 @@ def _tf_repositories():
     # XNNPack dependency.
     tf_http_archive(
         name = "KleidiAI",
-        sha256 = "652ee046ebf99bee5ed78aaf8961ea1df6ac09b6ebbd3a33106d9ce255efba69",
-        strip_prefix = "kleidiai-4bc7bd457930de119f8b826ecfdebfbd9c08cf8c",
-        urls = tf_mirror_urls("https://gitlab.arm.com/kleidi/kleidiai/-/archive/4bc7bd457930de119f8b826ecfdebfbd9c08cf8c/kleidiai-4bc7bd457930de119f8b826ecfdebfbd9c08cf8c.zip"),
+        sha256 = "a424e81535ebb6d0117e9282fba45d70aee0bedb481f603df55e1805bed0e21b",
+        strip_prefix = "kleidiai-5147983e14600ce7f8a6583dc04ba07753ac54c4",
+        urls = tf_mirror_urls("https://gitlab.arm.com/kleidi/kleidiai/-/archive/5147983e14600ce7f8a6583dc04ba07753ac54c4/kleidiai-5147983e14600ce7f8a6583dc04ba07753ac54c4.zip"),
     )
 
     FXdiv()
diff --git a/third_party/xla/third_party/kleidiai/workspace.bzl b/third_party/xla/third_party/kleidiai/workspace.bzl
index c6a1c0e517e..c9ed642a952 100644
--- a/third_party/xla/third_party/kleidiai/workspace.bzl
+++ b/third_party/xla/third_party/kleidiai/workspace.bzl
@@ -20,7 +20,7 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 def repo():
     tf_http_archive(
         name = "KleidiAI",
-        sha256 = "6b3e6630be314a28f6ea28fed14f7109b0b7c472f1e06d2dba17ffccda3b9466",
-        strip_prefix = "kleidiai-dce86647385ab2638aa5abebcb652f3e4271970d",
-        urls = tf_mirror_urls("https://gitlab.arm.com/kleidi/kleidiai/-/archive/dce86647385ab2638aa5abebcb652f3e4271970d/kleidiai-dce86647385ab2638aa5abebcb652f3e4271970d.zip"),
+        sha256 = "a424e81535ebb6d0117e9282fba45d70aee0bedb481f603df55e1805bed0e21b",
+        strip_prefix = "kleidiai-5147983e14600ce7f8a6583dc04ba07753ac54c4",
+        urls = tf_mirror_urls("https://gitlab.arm.com/kleidi/kleidiai/-/archive/5147983e14600ce7f8a6583dc04ba07753ac54c4/kleidiai-5147983e14600ce7f8a6583dc04ba07753ac54c4.zip"),
     )
diff --git a/third_party/xla/third_party/slinky/workspace.bzl b/third_party/xla/third_party/slinky/workspace.bzl
index f9952ab6b80..0d87a2ad678 100644
--- a/third_party/xla/third_party/slinky/workspace.bzl
+++ b/third_party/xla/third_party/slinky/workspace.bzl
@@ -20,7 +20,7 @@ load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 def repo():
     tf_http_archive(
         name = "slinky",
-        sha256 = "6ba811c39fd400149d58a42a1a5ad33e7fe2554c9a45990788779e0657039405",
-        strip_prefix = "slinky-2c9e129fc82215530af851cfd8d3094de0519241",
-        urls = tf_mirror_urls("https://github.com/dsharlet/slinky/archive/2c9e129fc82215530af851cfd8d3094de0519241.zip"),
+        sha256 = "cca25390873ad01a6e461edd39d220fb16c912beba2ebe02ae3b897ca63f30ef",
+        strip_prefix = "slinky-a3eec6fcfe3093be04d6c3d7dcab7d4e70bdd0dc",
+        urls = tf_mirror_urls("https://github.com/dsharlet/slinky/archive/a3eec6fcfe3093be04d6c3d7dcab7d4e70bdd0dc.zip"),
     )
diff --git a/third_party/xla/third_party/xnnpack/workspace.bzl b/third_party/xla/third_party/xnnpack/workspace.bzl
index ca6374f1cae..d56ec2b0a8b 100644
--- a/third_party/xla/third_party/xnnpack/workspace.bzl
+++ b/third_party/xla/third_party/xnnpack/workspace.bzl
@@ -21,8 +21,8 @@ def repo():
     # LINT.IfChange
     tf_http_archive(
         name = "XNNPACK",
-        sha256 = "aab0ed2b3972125ca057b8e2b16127fbf68f8128ea07dd7efe56cdb92f2c1917",
-        strip_prefix = "XNNPACK-d89ef6669a14db203b3b7935b1b3862cb63fb6df",
-        urls = tf_mirror_urls("https://github.com/google/XNNPACK/archive/d89ef6669a14db203b3b7935b1b3862cb63fb6df.zip"),
+        sha256 = "cb485a0fe76dec243568ec6ce6c0b75e18909ed7df5459921b46eaada332fe8f",
+        strip_prefix = "XNNPACK-0a66ebb5817fce48df605451444dcfc27a5e4ee4",
+        urls = tf_mirror_urls("https://github.com/google/XNNPACK/archive/0a66ebb5817fce48df605451444dcfc27a5e4ee4.zip"),
     )
     # LINT.ThenChange(//tensorflow/lite/tools/cmake/modules/xnnpack.cmake)

From 86e07deecf64b493f0c6747a8dada74e55f33253 Mon Sep 17 00:00:00 2001
From: Parker Schuh <parkers@google.com>
Date: Tue, 13 Oct 2026 14:23:31 -0700
Subject: [PATCH 0716/1000] Add
 CommonPjRtClient::MakeLoadStateAndResolveAddressableDevices as a temporary
 override to gather together support for loading TPU executables using
 CommonPjRtClient APIs.

PiperOrigin-RevId: 998890832
---
 .../xla/xla/pjrt/common_pjrt_client.cc        | 172 +++++++++++++-----
 third_party/xla/xla/pjrt/common_pjrt_client.h |  17 ++
 2 files changed, 140 insertions(+), 49 deletions(-)

diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.cc b/third_party/xla/xla/pjrt/common_pjrt_client.cc
index 0ed4110ee4d..578aad5f869 100644
--- a/third_party/xla/xla/pjrt/common_pjrt_client.cc
+++ b/third_party/xla/xla/pjrt/common_pjrt_client.cc
@@ -554,52 +554,14 @@ absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> CommonPjRtClient::Load(
   return loaded_executable;
 }
 
-absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
-CommonPjRtClient::LoadInternal(std::shared_ptr<PjRtExecutable> executable,
-                               const LoadOptions& load_options, bool dump) {
-  tsl::profiler::TraceMe traceme("CommonPjRtClient::Load");
-  VLOG(1) << "CommonPjRtClient::Load";
-  ABSL_ASSIGN_OR_RETURN(const PjRtTopologyDescription* topology,
-                   GetTopologyDescription());
-  ABSL_ASSIGN_OR_RETURN(auto hlo_module, executable->GetHloModule());
-
-  ABSL_ASSIGN_OR_RETURN(CompileOptions compile_options,
-                   executable->GetCompileOptions());
-  ABSL_RETURN_IF_ERROR(compile_options.ApplyAllOptionOverrides());
-  if (IsEarlyExitCompilation(compile_options)) {
-    return InvalidArgument(
-        "Executable compiled with xla_early_exit_with_layouts cannot be "
-        "loaded.");
-  }
-  if (!IsGpuId(platform_id())) {
-    absl::StatusOr<std::unique_ptr<PjRtRuntimeAbiVersion>> runtime_abi_version =
-        RuntimeAbiVersion();
-    if (!absl::IsUnimplemented(runtime_abi_version.status())) {
-      ABSL_RETURN_IF_ERROR(runtime_abi_version.status());
-      ABSL_ASSIGN_OR_RETURN(
-          std::unique_ptr<PjRtExecutableAbiVersion> executable_abi_version,
-          executable->GetAbiVersion());
-      ABSL_RETURN_IF_ERROR(
-          (*runtime_abi_version)->IsCompatibleWith(*executable_abi_version));
-    }
-  }
-  std::vector<PjRtLoadedExecutable::LogicalDeviceIds>
-      addressable_device_logical_ids;
-  std::vector<PjRtDevice*> addressable_devices;
-  int num_replicas;
-  int num_partitions;
-  std::shared_ptr<DeviceAssignment> device_assignment;
-  ABSL_RETURN_IF_ERROR(ParseDeviceAssignmentCompileOptions(
-      compile_options.compile_portable_executable,
-      &compile_options.executable_build_options,
-      [this, topology, &compile_options](int num_replicas, int num_partitions) {
-        return topology->GetDefaultDeviceAssignment(
-            process_index(), num_replicas,
-            /*num_replicas_per_slice=*/std::nullopt, num_partitions,
-            compile_options.multi_slice_config);
-      },
-      &num_replicas, &num_partitions, &device_assignment));
-
+absl::StatusOr<tsl::RCReference<PjRtExecutableLoadState>>
+CommonPjRtClient::MakeLoadStateAndResolveAddressableDevices(
+    PjRtExecutable* executable, CompileOptions* compile_options,
+    std::vector<PjRtLoadedExecutable::LogicalDeviceIds>&
+        addressable_device_logical_ids,
+    std::vector<PjRtDevice*>& addressable_devices,
+    std::shared_ptr<DeviceAssignment>& device_assignment,
+    const LoadOptions& load_options) {
   // Find devices that are addressable by this client/task.
   if (device_assignment != nullptr) {
     int num_replicas = device_assignment->replica_count();
@@ -630,6 +592,73 @@ CommonPjRtClient::LoadInternal(std::shared_ptr<PjRtExecutable> executable,
     }
   }
 
+  return raw_client()->MakeLoadState();
+}
+
+absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
+CommonPjRtClient::LoadInternal(std::shared_ptr<PjRtExecutable> executable,
+                               const LoadOptions& load_options, bool dump) {
+  ABSL_ASSIGN_OR_RETURN(CompileOptions compile_options,
+                   executable->GetCompileOptions());
+  ABSL_RETURN_IF_ERROR(compile_options.ApplyAllOptionOverrides());
+  return LoadInternal(executable, std::move(compile_options), nullptr,
+                      load_options, dump);
+}
+
+absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
+CommonPjRtClient::LoadInternal(
+    std::shared_ptr<PjRtExecutable> executable,
+    CompileOptions&& compile_options,
+    std::shared_ptr<DeviceAssignment> device_assignment,
+    const LoadOptions& load_options, bool dump) {
+  tsl::profiler::TraceMe traceme("CommonPjRtClient::Load");
+  VLOG(1) << "CommonPjRtClient::Load";
+
+  ABSL_ASSIGN_OR_RETURN(const PjRtTopologyDescription* topology,
+                   GetTopologyDescription());
+  ABSL_ASSIGN_OR_RETURN(auto hlo_module, executable->GetHloModule());
+  if (IsEarlyExitCompilation(compile_options)) {
+    return InvalidArgument(
+        "Executable compiled with xla_early_exit_with_layouts cannot be "
+        "loaded.");
+  }
+  if (!IsGpuId(platform_id())) {
+    absl::StatusOr<std::unique_ptr<PjRtRuntimeAbiVersion>> runtime_abi_version =
+        RuntimeAbiVersion();
+    if (!absl::IsUnimplemented(runtime_abi_version.status())) {
+      ABSL_RETURN_IF_ERROR(runtime_abi_version.status());
+      ABSL_ASSIGN_OR_RETURN(
+          std::unique_ptr<PjRtExecutableAbiVersion> executable_abi_version,
+          executable->GetAbiVersion());
+      ABSL_RETURN_IF_ERROR(
+          (*runtime_abi_version)->IsCompatibleWith(*executable_abi_version));
+    }
+  }
+  std::vector<PjRtLoadedExecutable::LogicalDeviceIds>
+      addressable_device_logical_ids;
+  std::vector<PjRtDevice*> addressable_devices;
+  if (device_assignment == nullptr) {
+    int num_replicas;
+    int num_partitions;
+    ABSL_RETURN_IF_ERROR(ParseDeviceAssignmentCompileOptions(
+        compile_options.compile_portable_executable,
+        &compile_options.executable_build_options,
+        [this, topology, &compile_options](int num_replicas,
+                                           int num_partitions) {
+          return topology->GetDefaultDeviceAssignment(
+              process_index(), num_replicas,
+              /*num_replicas_per_slice=*/std::nullopt, num_partitions,
+              compile_options.multi_slice_config);
+        },
+        &num_replicas, &num_partitions, &device_assignment));
+  }
+
+  ABSL_ASSIGN_OR_RETURN(
+      auto load_state,
+      MakeLoadStateAndResolveAddressableDevices(
+          executable.get(), &compile_options, addressable_device_logical_ids,
+          addressable_devices, device_assignment, load_options));
+  ABSL_RETURN_IF_ERROR(load_state->Preload(executable.get()));
   const auto& ex_options = compile_options.executable_build_options;
   const bool xla_dump_hlo_unoptimized_snapshots =
       ex_options.has_debug_options() &&
@@ -675,8 +704,55 @@ CommonPjRtClient::LoadInternal(std::shared_ptr<PjRtExecutable> executable,
     }
   }
 
+  // TODO(parkers): get to the bottom of this...
   auto parameter_shapes =
-      GetParameterShapes(hlo_module->compute_computation_layout());
+      IsTpuId(platform_id())
+          ? GetParameterShapes(hlo_module->entry_computation_layout())
+          : GetParameterShapes(hlo_module->compute_computation_layout());
+  if (IsTpuId(platform_id())) {
+    auto canonicalize_shapes =
+        [&](absl::Span<xla::Shape> shapes) -> absl::Status {
+      for (xla::Shape& shape : shapes) {
+        ABSL_ASSIGN_OR_RETURN(int kind_id, GetMemorySpaceKindForShape(shape));
+        ABSL_ASSIGN_OR_RETURN(shape,
+                         topology->MakeCanonicalShapeForMemorySpace(
+                             kind_id, shape,
+                             shape.has_layout() ? &shape.layout() : nullptr));
+      }
+      return absl::OkStatus();
+    };
+    ABSL_RETURN_IF_ERROR(canonicalize_shapes(
+        result_shape.IsTuple()
+            ? absl::MakeSpan(*result_shape.mutable_tuple_shapes())
+            : absl::MakeSpan(&result_shape, 1)));
+    ABSL_RETURN_IF_ERROR(canonicalize_shapes(absl::MakeSpan(parameter_shapes)));
+    // canonicalize_shapes calls MakeCanonicalShapeForMemorySpace, which by
+    // default sets dynamic_shape_metadata_prefix_bytes to 1024 for all dynamic
+    // shapes. However, custom call PadRealToStatic consumes dynamic inputs
+    // without the 1024-byte metadata prefix. Find parameters feeding
+    // PadRealToStatic custom calls and explicitly set their prefix bytes to 0.
+    const HloComputation* entry_comp = hlo_module->entry_computation();
+    for (const HloInstruction* instruction : entry_comp->instructions()) {
+      if (instruction->IsCustomCall("PadRealToStatic") &&
+          instruction->operand_count() > 0) {
+        const HloInstruction* operand = instruction->operand(0);
+        while (operand->opcode() == HloOpcode::kCopy ||
+               operand->opcode() == HloOpcode::kBitcast) {
+          operand = operand->operand(0);
+        }
+        CHECK(operand->opcode() == HloOpcode::kParameter);
+        CHECK(operand->shape().IsArray());
+        int64_t param_no = operand->parameter_number();
+        DCHECK_LT(param_no, parameter_shapes.size());
+        DCHECK(parameter_shapes[param_no].IsArray());
+        if (param_no < parameter_shapes.size()) {
+          parameter_shapes[param_no]
+              .mutable_layout()
+              ->set_dynamic_shape_metadata_prefix_bytes(0);
+        }
+      }
+    }
+  }
   using InputHloSnapshotBits =
       CommonPjRtLoadedExecutable::DispatchInfo::InputHloSnapshotBits;
   std::unique_ptr<InputHloSnapshotBits> input_hlo_snapshot_bits;
@@ -705,8 +781,6 @@ CommonPjRtClient::LoadInternal(std::shared_ptr<PjRtExecutable> executable,
   if (IsCpuId(platform_id())) {
     dispatch_info.parameter_memory_space_kind_ids.clear();
   }
-  auto load_state = raw_client()->MakeLoadState();
-  ABSL_RETURN_IF_ERROR(load_state->Preload(executable.get()));
   auto loaded_executable = std::make_unique<CommonPjRtLoadedExecutable>(
       this, raw_client()->ToAsyncExecutable(std::move(executable)),
       std::move(dispatch_info), std::move(load_state));
diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.h b/third_party/xla/xla/pjrt/common_pjrt_client.h
index 2eeffdb8254..1c615e4486f 100644
--- a/third_party/xla/xla/pjrt/common_pjrt_client.h
+++ b/third_party/xla/xla/pjrt/common_pjrt_client.h
@@ -592,7 +592,24 @@ class CommonPjRtClient : public PjRtClient {
       std::shared_ptr<PjRtExecutable> executable,
       const LoadOptions& load_options) override;
 
+  // TODO(parkers): Remove the need for this temporary function once more
+  // backend differences are resolved.
+  virtual absl::StatusOr<tsl::RCReference<PjRtExecutableLoadState>>
+  MakeLoadStateAndResolveAddressableDevices(
+      PjRtExecutable* executable, CompileOptions* compile_options,
+      std::vector<PjRtLoadedExecutable::LogicalDeviceIds>&
+          addressable_device_logical_ids,
+      std::vector<PjRtDevice*>& addressable_devices,
+      std::shared_ptr<DeviceAssignment>& device_assignment,
+      const LoadOptions& load_options);
+
  protected:
+  absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> LoadInternal(
+      std::shared_ptr<PjRtExecutable> executable,
+      CompileOptions&& compile_options,
+      std::shared_ptr<DeviceAssignment> device_assignment,
+      const LoadOptions& load_options, bool dump);
+
   absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> LoadInternal(
       std::shared_ptr<PjRtExecutable> executable,
       const LoadOptions& load_options, bool dump);

From 0d1623270eeae3b6efa56e66b1153ec89e88bfa4 Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Tue, 13 Oct 2026 14:37:16 -0700
Subject: [PATCH 0717/1000] Take diagonals of XLA MatrixDiagPart inputs with
 dynamic dimensions

MatrixDiagPartOp sized its output from the input's static shape, which
for a dynamic dimension is its bound, so the diagonals came out as long
as at the bound, with entries from outside the actual matrix. With
dynamic dimensions, take the diagonals of the matrices at their bound
sizes, place each diagonal of a band after the padding it has at the
actual sizes, with padding_value past its actual length, and give the
output its actual sizes.

MatrixDiagPartV3, which tf.linalg.diag_part uses, is lowered through
MLIR, whose lowering only takes static shapes. Send its inputs with
dynamic dimensions to MatrixDiagPartOp, and keep the MLIR lowering for
the others.
---
 tensorflow/compiler/tests/BUILD               |   2 +
 .../compiler/tests/matrix_diag_ops_test.py    |  65 +++++++
 .../tf2xla/kernels/matrix_diag_ops.cc         | 165 +++++++++++++++---
 3 files changed, 204 insertions(+), 28 deletions(-)

diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD
index deb8ce4f21f..ba636fa0f38 100644
--- a/tensorflow/compiler/tests/BUILD
+++ b/tensorflow/compiler/tests/BUILD
@@ -1213,6 +1213,8 @@ tf_xla_py_strict_test(
     ],
     deps = [
         ":xla_test",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/ops:array_ops",
         "//tensorflow/python/ops:array_ops_gen",
         "//tensorflow/python/platform:test",
diff --git a/tensorflow/compiler/tests/matrix_diag_ops_test.py b/tensorflow/compiler/tests/matrix_diag_ops_test.py
index 0c265a11ccb..e37d4a11f7c 100644
--- a/tensorflow/compiler/tests/matrix_diag_ops_test.py
+++ b/tensorflow/compiler/tests/matrix_diag_ops_test.py
@@ -17,6 +17,8 @@
 import numpy as np
 
 from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
 from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import gen_array_ops
 from tensorflow.python.platform import googletest
@@ -767,6 +769,69 @@ class MatrixDiagPartTest(xla_test.XLATestCase):
                   "align": align
               }, solution)
 
+  def _diagPartOfSelected(self, op, x, masks, **params):
+    # masks maps axes of x to masks of the entries to keep along them, which
+    # gives those dimensions a size only known at run time.
+    axes = sorted(masks)
+
+    @def_function.function(jit_compile=True)
+    def f(x, *x_masks):
+      for axis, mask in zip(axes, x_masks):
+        x = array_ops.gather(
+            x, array_ops.reshape(array_ops.where(mask), [-1]), axis=axis)
+      return op(x, **params)
+
+    x = np.array(x, dtype=np.int32)
+    with self.session() as session:
+      with self.device_scope():
+        x_ph = array_ops.placeholder(dtypes.int32, x.shape)
+        mask_phs = [
+            array_ops.placeholder(dtypes.bool, [x.shape[axis]])
+            for axis in axes
+        ]
+        output = f(x_ph, *mask_phs)
+      feed_dict = {x_ph: x}
+      feed_dict.update({ph: masks[axis] for ph, axis in zip(mask_phs, axes)})
+      return session.run(output, feed_dict)
+
+  def testDynamicDimensions(self):
+    # Regression test for GitHub issue 129053: with a dynamic dimension, the
+    # diagonals had the length they have at its bound. The mask keeps 2 of 5
+    # entries, so the dimensions it selects have size 2 and bound 5.
+    keep = [True, False, True, False, False]
+    x = np.arange(1, 16).reshape(5, 3)  # Keeps [[1, 2, 3], [7, 8, 9]].
+    for op in (array_ops.matrix_diag_part, gen_array_ops.matrix_diag_part):
+      self.assertAllEqual(self._diagPartOfSelected(op, x, {-2: keep}), [1, 8])
+    self.assertAllEqual(
+        self._diagPartOfSelected(
+            array_ops.matrix_diag_part,
+            x, {-2: keep},
+            k=(-1, 1),
+            padding_value=9), [[2, 9], [1, 8], [7, 9]])
+    self.assertAllEqual(
+        self._diagPartOfSelected(
+            array_ops.matrix_diag_part,
+            x, {-2: keep},
+            k=(-1, 1),
+            padding_value=9,
+            align="LEFT_RIGHT"), [[2, 9], [1, 8], [9, 7]])
+    self.assertAllEqual(
+        self._diagPartOfSelected(
+            gen_array_ops.matrix_diag_part_v2,
+            x, {-2: keep},
+            k=[-1, 1],
+            padding_value=9), [[2, 9], [1, 8], [7, 9]])
+    self.assertAllEqual(
+        self._diagPartOfSelected(array_ops.matrix_diag_part,
+                                 np.arange(1, 16).reshape(3, 5), {-1: keep}),
+        [1, 8])
+    self.assertAllEqual(
+        self._diagPartOfSelected(array_ops.matrix_diag_part,
+                                 np.arange(1, 51).reshape(2, 5, 5), {
+                                     -2: keep,
+                                     -1: keep
+                                 }), [[1, 13], [26, 38]])
+
 
 if __name__ == "__main__":
   googletest.main()
diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc b/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc
index 8dc866ee430..0f08021416f 100644
--- a/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc
@@ -31,6 +31,7 @@ limitations under the License.
 #include "xla/hlo/builder/lib/constants.h"
 #include "xla/hlo/builder/lib/matrix.h"
 #include "xla/hlo/builder/xla_builder.h"
+#include "xla/shape.h"
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
 #include "tensorflow/core/framework/op_kernel.h"
@@ -416,42 +417,125 @@ class MatrixDiagPartOp : public XlaOpKernel {
 
     // Computes output.
     xla::XlaOp input = context->Input(0);
+
+    // With dynamic dimensions, take the diagonals of the matrices at their
+    // bound sizes, which hold the actual matrices in the top left, and give
+    // the output its actual sizes at the end.
+    xla::Shape input_xla_shape;
+    OP_REQUIRES_VALUE(input_xla_shape, context, context->InputXlaShape(0));
+    std::vector<xla::XlaOp> input_dim_sizes(input_rank);
+    for (int i = 0; i < input_rank; ++i) {
+      if (input_xla_shape.is_dynamic_dimension(i)) {
+        input_dim_sizes[i] = xla::GetDimensionSize(input, i);
+        input = xla::RemoveDynamicDimension(input, i);
+      }
+    }
+    const bool dynamic_matrix =
+        input_xla_shape.is_dynamic_dimension(input_rank - 2) ||
+        input_xla_shape.is_dynamic_dimension(input_rank - 1);
+    const auto actual_size = [&](int dim) {
+      return input_xla_shape.is_dynamic_dimension(dim)
+                 ? input_dim_sizes[dim]
+                 : xla::ConstantR0<int32_t>(
+                       context->builder(),
+                       static_cast<int32_t>(input_shape.dim_size(dim)));
+    };
+    // The length of the longest diagonal from lower to upper at the actual
+    // sizes.
+    const auto actual_len = [&](int64_t lower, int64_t upper) {
+      const xla::XlaOp actual_num_rows = actual_size(input_rank - 2);
+      const xla::XlaOp actual_num_cols = actual_size(input_rank - 1);
+      return xla::Min(
+          xla::Add(actual_num_rows,
+                   xla::ScalarLike(actual_num_rows,
+                                   std::min(upper, int64_t{0}))),
+          xla::Sub(actual_num_cols,
+                   xla::ScalarLike(actual_num_cols,
+                                   std::max(lower, int64_t{0}))));
+    };
+    xla::XlaOp actual_max_diag_len;
+    if (dynamic_matrix) {
+      actual_max_diag_len = actual_len(lower_diag_index, upper_diag_index);
+    }
+
+    xla::XlaOp output;
     std::vector<xla::XlaOp> diag_list;
     xla::PaddingConfig padding_config =
         xla::MakeNoPaddingConfig(input_rank - 1);
     if (num_diags == 1) {
-      context->SetOutput(
-          0, is_gpu_ ? xla::GetMatrixDiagonalViaGather(input, upper_diag_index)
-                     : xla::GetMatrixDiagonal(input, upper_diag_index));
-      return;
-    }
-    for (int diag_index = upper_diag_index; diag_index >= lower_diag_index;
-         --diag_index) {
-      xla::XlaOp single_diag =
-          is_gpu_ ? xla::GetMatrixDiagonalViaGather(input, diag_index)
-                  : xla::GetMatrixDiagonal(input, diag_index);
-      const int64_t diag_len = ComputeDiagLen(diag_index, num_rows, num_cols);
-      const int64_t padding_len = max_diag_len - diag_len;
-      if (padding_len > 0) {
-        if (IsLeftAligned(diag_index, left_align_superdiagonal_,
-                          left_align_subdiagonal_)) {
+      output = is_gpu_
+                   ? xla::GetMatrixDiagonalViaGather(input, upper_diag_index)
+                   : xla::GetMatrixDiagonal(input, upper_diag_index);
+    } else {
+      auto slice_sizes = input_shape.dim_sizes();
+      slice_sizes.pop_back();
+      slice_sizes.back() = max_diag_len;
+      for (int diag_index = upper_diag_index; diag_index >= lower_diag_index;
+           --diag_index) {
+        xla::XlaOp single_diag =
+            is_gpu_ ? xla::GetMatrixDiagonalViaGather(input, diag_index)
+                    : xla::GetMatrixDiagonal(input, diag_index);
+        const int64_t diag_len =
+            ComputeDiagLen(diag_index, num_rows, num_cols);
+        const int64_t padding_len = max_diag_len - diag_len;
+        const bool left_aligned = IsLeftAligned(
+            diag_index, left_align_superdiagonal_, left_align_subdiagonal_);
+        if (dynamic_matrix) {
+          // Place the diagonal after the padding it has at the actual sizes,
+          // and pad its entries past its actual length, which come from
+          // outside the actual matrix, with padding_value too.
+          const xla::XlaOp len = actual_len(diag_index, diag_index);
+          const xla::XlaOp offset =
+              left_aligned ? xla::ScalarLike(len, 0)
+                           : xla::Sub(actual_max_diag_len, len);
           padding_config.mutable_dimensions(input_rank - 2)
-              ->set_edge_padding_low(0);
+              ->set_edge_padding_low(max_diag_len);
           padding_config.mutable_dimensions(input_rank - 2)
-              ->set_edge_padding_high(padding_len);
-        } else {
-          padding_config.mutable_dimensions(input_rank - 2)
-              ->set_edge_padding_low(padding_len);
-          padding_config.mutable_dimensions(input_rank - 2)
-              ->set_edge_padding_high(0);
+              ->set_edge_padding_high(max_diag_len);
+          single_diag = xla::Pad(single_diag, padding_value, padding_config);
+          std::vector<xla::XlaOp> start_indices(input_rank - 1,
+                                                xla::ScalarLike(len, 0));
+          start_indices.back() =
+              xla::Sub(xla::ScalarLike(len, max_diag_len), offset);
+          single_diag =
+              xla::DynamicSlice(single_diag, start_indices, slice_sizes);
+          const xla::XlaOp in_diagonal = xla::Lt(
+              xla::Iota(context->builder(), xla::S32, max_diag_len),
+              xla::Add(offset, len));
+          single_diag = xla::Select(
+              xla::BroadcastInDim(in_diagonal, slice_sizes, {input_rank - 2}),
+              single_diag, xla::Broadcast(padding_value, slice_sizes));
+        } else if (padding_len > 0) {
+          if (left_aligned) {
+            padding_config.mutable_dimensions(input_rank - 2)
+                ->set_edge_padding_low(0);
+            padding_config.mutable_dimensions(input_rank - 2)
+                ->set_edge_padding_high(padding_len);
+          } else {
+            padding_config.mutable_dimensions(input_rank - 2)
+                ->set_edge_padding_low(padding_len);
+            padding_config.mutable_dimensions(input_rank - 2)
+                ->set_edge_padding_high(0);
+          }
+          single_diag = xla::Pad(single_diag, padding_value, padding_config);
         }
-        single_diag = xla::Pad(single_diag, padding_value, padding_config);
+        diag_list.emplace_back(single_diag);
       }
-      diag_list.emplace_back(single_diag);
+      auto concat =
+          xla::ConcatInDim(context->builder(), diag_list, input_rank - 2);
+      output = xla::Reshape(concat, output_shape.dim_sizes());
     }
-    auto concat =
-        xla::ConcatInDim(context->builder(), diag_list, input_rank - 2);
-    context->SetOutput(0, xla::Reshape(concat, output_shape.dim_sizes()));
+
+    if (dynamic_matrix) {
+      output = xla::SetDimensionSize(output, actual_max_diag_len,
+                                     output_shape.dims() - 1);
+    }
+    for (int i = 0; i < input_rank - 2; ++i) {
+      if (input_xla_shape.is_dynamic_dimension(i)) {
+        output = xla::SetDimensionSize(output, input_dim_sizes[i], i);
+      }
+    }
+    context->SetOutput(0, output);
   }
 
  private:
@@ -466,10 +550,35 @@ REGISTER_XLA_OP(Name("MatrixDiagPartV2")
                     .CompileTimeConstantInput("k")
                     .CompileTimeConstantInput("padding_value"),
                 MatrixDiagPartOp);
+
+// MatrixDiagPartV3 is lowered through MLIR, whose lowering only takes static
+// shapes, so an input with a dynamic dimension goes through MatrixDiagPartOp.
+class MatrixDiagPartV3Op : public XlaOpKernel {
+ public:
+  explicit MatrixDiagPartV3Op(OpKernelConstruction* context)
+      : XlaOpKernel(context),
+        mlir_kernel_(context),
+        dynamic_shape_kernel_(context) {}
+
+  void Compile(XlaOpKernelContext* context) override {
+    xla::Shape input_xla_shape;
+    OP_REQUIRES_VALUE(input_xla_shape, context, context->InputXlaShape(0));
+    XlaOpKernel& kernel =
+        input_xla_shape.is_dynamic()
+            ? static_cast<XlaOpKernel&>(dynamic_shape_kernel_)
+            : static_cast<XlaOpKernel&>(mlir_kernel_);
+    kernel.Compile(context);
+  }
+
+ private:
+  MlirXlaOpKernel mlir_kernel_;
+  MatrixDiagPartOp dynamic_shape_kernel_;
+};
+
 REGISTER_XLA_OP(Name("MatrixDiagPartV3")
                     .CompileTimeConstantInput("k")
                     .CompileTimeConstantInput("padding_value"),
-                MlirXlaOpKernel);
+                MatrixDiagPartV3Op);
 
 class MatrixSetDiagOp : public XlaOpKernel {
  public:

From a6e14b01fb5715f09fef82363a0ec3f32c055ccb Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Tue, 13 Oct 2026 14:30:45 -0700
Subject: [PATCH 0718/1000] [XLA] Guard add and multiply constant-reassociation
 rules in AlgebraicSimplifier on fast math.

Floating-point and complex addition and multiplication are not associative, so reassociating expressions like `(A + C1) + C2 => A + (C1 + C2)` or `(A * C1) * C2 => A * (C1 * C2)` can change rounding or introduce/mask overflow, underflow, and NaNs. Restrict these reassociation rewrites in `HandleAdd` and `HandleMultiply` to non-floating-point, non-complex types unless `enable_fast_math` is enabled.

PiperOrigin-RevId: 998895281
---
 .../simplifiers/algebraic_simplifier.cc       |  74 +++++---
 .../simplifiers/algebraic_simplifier_test.cc  | 165 +++++++++++++++---
 2 files changed, 191 insertions(+), 48 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
index 543737216a5..98c101a5671 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
@@ -966,11 +966,14 @@ absl::Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) {
   //
   VLOG(10) << "trying transform [(A + C1) + C2 => A + (C1 + C2)]";
   HloInstruction *a, *c1, *c2;
-  if (Match(add, m::Add(m::Add(m::NonConstant(&a), m::Constant(&c1)),
-                        m::Constant(&c2))) ||
-      Match(add, m::Add(m::Add(m::NonConstant(&a),
-                               m::Broadcast(m::ConstantScalar(&c1))),
-                        m::Broadcast(m::ConstantScalar(&c2))))) {
+  if (((!ShapeUtil::ElementIsFloating(add->shape()) &&
+        !ShapeUtil::ElementIsComplex(add->shape())) ||
+       options_.enable_fast_math()) &&
+      (Match(add, m::Add(m::Add(m::NonConstant(&a), m::Constant(&c1)),
+                         m::Constant(&c2))) ||
+       Match(add, m::Add(m::Add(m::NonConstant(&a),
+                                m::Broadcast(m::ConstantScalar(&c1))),
+                         m::Broadcast(m::ConstantScalar(&c2)))))) {
     ABSL_ASSIGN_OR_RETURN(auto* sum_of_constants,
                      MakeBinaryHlo(HloOpcode::kAdd, c1, c2));
     if (ShapeUtil::IsScalar(sum_of_constants->shape()) &&
@@ -984,11 +987,14 @@ absl::Status AlgebraicSimplifierVisitor::HandleAdd(HloInstruction* add) {
   }
 
   VLOG(10) << "trying transform [(C1 - A) + C2 => (C1 + C2) - A]";
-  if (Match(add, m::Add(m::Subtract(m::Constant(&c1), m::NonConstant(&a)),
-                        m::Constant(&c2))) ||
-      Match(add, m::Add(m::Subtract(m::Broadcast(m::ConstantScalar(&c1)),
-                                    m::NonConstant(&a)),
-                        m::Broadcast(m::ConstantScalar(&c2))))) {
+  if (((!ShapeUtil::ElementIsFloating(add->shape()) &&
+        !ShapeUtil::ElementIsComplex(add->shape())) ||
+       options_.enable_fast_math()) &&
+      (Match(add, m::Add(m::Subtract(m::Constant(&c1), m::NonConstant(&a)),
+                         m::Constant(&c2))) ||
+       Match(add, m::Add(m::Subtract(m::Broadcast(m::ConstantScalar(&c1)),
+                                     m::NonConstant(&a)),
+                         m::Broadcast(m::ConstantScalar(&c2)))))) {
     ABSL_ASSIGN_OR_RETURN(HloInstruction * sum_of_constants,
                      MakeBinaryHlo(HloOpcode::kAdd, c1, c2));
     if (ShapeUtil::IsScalar(sum_of_constants->shape()) &&
@@ -5110,7 +5116,10 @@ absl::Status AlgebraicSimplifierVisitor::HandleMultiply(
     HloInstruction *a, *b, *c1, *c2;
     // Mul(Mul(x, constant1), Mul(y, constant2)) => Mul(Mul(x, y),
     // constant1*constant2)
-    if (Match(multiply,
+    if (((!ShapeUtil::ElementIsFloating(multiply->shape()) &&
+          !ShapeUtil::ElementIsComplex(multiply->shape())) ||
+         options_.enable_fast_math()) &&
+        Match(multiply,
               m::MultiplyAnyOrder(
                   m::MultiplyAnyOrder(m::NonConstant(&a), m::Constant(&c1)),
                   m::MultiplyAnyOrder(m::NonConstant(&b), m::Constant(&c2))))) {
@@ -5135,7 +5144,10 @@ absl::Status AlgebraicSimplifierVisitor::HandleMultiply(
   {
     HloInstruction *a, *c1, *c2;
     // Mul(Mul(a, constant1), constant2) => Mul(a, constant1*constant2)
-    if (Match(multiply,
+    if (((!ShapeUtil::ElementIsFloating(multiply->shape()) &&
+          !ShapeUtil::ElementIsComplex(multiply->shape())) ||
+         options_.enable_fast_math()) &&
+        Match(multiply,
               m::MultiplyAnyOrder(
                   m::MultiplyAnyOrder(m::NonConstant(&a), m::Constant(&c1)),
                   m::Constant(&c2)))) {
@@ -5159,15 +5171,18 @@ absl::Status AlgebraicSimplifierVisitor::HandleMultiply(
     HloInstruction *a, *b, *constant, *op;
     // Mul(Mul(a, constant1), Broadcast(b)) =>
     // Mul(Broadcast(Mul(b, constant1), a))
-    if (Match(multiply,
-              m::MultiplyAnyOrder(m::MultiplyAnyOrder(m::NonConstant(&a),
-                                                      m::Constant(&constant)),
-                                  m::Op(&op))) ||
-        Match(multiply,
-              m::MultiplyAnyOrder(
-                  m::MultiplyAnyOrder(m::NonConstant(&a),
-                                      m::Broadcast(m::Constant(&constant))),
-                  m::Op(&op)))) {
+    if (((!ShapeUtil::ElementIsFloating(multiply->shape()) &&
+          !ShapeUtil::ElementIsComplex(multiply->shape())) ||
+         options_.enable_fast_math()) &&
+        (Match(multiply,
+               m::MultiplyAnyOrder(m::MultiplyAnyOrder(m::NonConstant(&a),
+                                                       m::Constant(&constant)),
+                                   m::Op(&op))) ||
+         Match(multiply,
+               m::MultiplyAnyOrder(
+                   m::MultiplyAnyOrder(m::NonConstant(&a),
+                                       m::Broadcast(m::Constant(&constant))),
+                   m::Op(&op))))) {
       // Check that the other side was a broadcast, and not of a constant.
       if (ShapeUtil::IsScalar(constant->shape()) &&
           Match(op, m::Broadcast(m::NonConstant()))) {
@@ -5193,13 +5208,16 @@ absl::Status AlgebraicSimplifierVisitor::HandleMultiply(
 
   VLOG(10) << "trying transform [(A * C1) * C2 => A * (C1 * C2)]";
   HloInstruction *a, *c1, *c2;
-  if (Match(multiply,
-            m::Multiply(m::Multiply(m::NonConstant(&a), m::Constant(&c1)),
-                        m::Constant(&c2))) ||
-      Match(multiply,
-            m::Multiply(
-                m::Multiply(m::Op(&a), m::Broadcast(m::ConstantScalar(&c1))),
-                m::Broadcast(m::ConstantScalar(&c2))))) {
+  if (((!ShapeUtil::ElementIsFloating(multiply->shape()) &&
+        !ShapeUtil::ElementIsComplex(multiply->shape())) ||
+       options_.enable_fast_math()) &&
+      (Match(multiply,
+             m::Multiply(m::Multiply(m::NonConstant(&a), m::Constant(&c1)),
+                         m::Constant(&c2))) ||
+       Match(multiply,
+             m::Multiply(
+                 m::Multiply(m::Op(&a), m::Broadcast(m::ConstantScalar(&c1))),
+                 m::Broadcast(m::ConstantScalar(&c2)))))) {
     ABSL_ASSIGN_OR_RETURN(auto* product_of_constants,
                      MakeBinaryHlo(HloOpcode::kMultiply, c1, c2));
     if (ShapeUtil::IsScalar(product_of_constants->shape()) &&
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
index 2725280e898..3b7af4365b8 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
@@ -672,8 +672,14 @@ TEST_F(AlgebraicSimplifierTest, MultiplyChain) {
       ROOT z = f32[] multiply(x, y)
     }
   )";
-  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
-  ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).value());
+  ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  EXPECT_THAT(AlgebraicSimplifier(default_options_).Run(m.get()),
+              absl_testing::IsOkAndHolds(false));
+
+  AlgebraicSimplifierOptions options = default_options_;
+  options.set_enable_fast_math(true);
+  ASSERT_THAT(AlgebraicSimplifier(options).Run(m.get()),
+              absl_testing::IsOkAndHolds(true));
   EXPECT_THAT(
       m->entry_computation()->root_instruction(),
       GmockMatch(m::MultiplyAnyOrder(
@@ -693,8 +699,14 @@ TEST_F(AlgebraicSimplifierTest, MultiplyChain2) {
       ROOT y = f32[] multiply(c, b)
     }
   )";
-  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
-  ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).value());
+  ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  EXPECT_THAT(AlgebraicSimplifier(default_options_).Run(m.get()),
+              absl_testing::IsOkAndHolds(false));
+
+  AlgebraicSimplifierOptions options = default_options_;
+  options.set_enable_fast_math(true);
+  ASSERT_THAT(AlgebraicSimplifier(options).Run(m.get()),
+              absl_testing::IsOkAndHolds(true));
   EXPECT_THAT(m->entry_computation()->root_instruction(),
               GmockMatch(m::MultiplyAnyOrder(
                   m::Parameter(0), m::MultiplyAnyOrder(m::ConstantScalar(2),
@@ -716,8 +728,14 @@ TEST_F(AlgebraicSimplifierTest, MultiplyBroadcastReassoc) {
       ROOT z = f32[2,2] multiply(y, x)
     }
   )";
-  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
-  ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).value());
+  ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  EXPECT_THAT(AlgebraicSimplifier(default_options_).Run(m.get()),
+              absl_testing::IsOkAndHolds(false));
+
+  AlgebraicSimplifierOptions options = default_options_;
+  options.set_enable_fast_math(true);
+  ASSERT_THAT(AlgebraicSimplifier(options).Run(m.get()),
+              absl_testing::IsOkAndHolds(true));
   EXPECT_THAT(m->entry_computation()->root_instruction(),
               GmockMatch(m::MultiplyAnyOrder(
                   m::Parameter(0), m::Broadcast(m::MultiplyAnyOrder(
@@ -1011,8 +1029,14 @@ TEST_F(AlgebraicSimplifierTest, MultiplyReassociateMergeConstants) {
       ROOT multiply1 = f32[] multiply(multiply0, c1)
     }
   )";
-  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
-  ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).value());
+  ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  EXPECT_THAT(AlgebraicSimplifier(default_options_).Run(m.get()),
+              absl_testing::IsOkAndHolds(false));
+
+  AlgebraicSimplifierOptions options = default_options_;
+  options.set_enable_fast_math(true);
+  ASSERT_THAT(AlgebraicSimplifier(options).Run(m.get()),
+              absl_testing::IsOkAndHolds(true));
   EXPECT_THAT(m->entry_computation()->root_instruction(),
               GmockMatch(m::Multiply(m::Parameter(0),
                                      m::Multiply(m::ConstantScalar(2.0),
@@ -1032,8 +1056,14 @@ TEST_F(AlgebraicSimplifierTest, MultiplyReassociateMergeBroadcastedConstants) {
       ROOT multiply1 = f32[4] multiply(multiply0, b1)
     }
   )";
-  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
-  ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).value());
+  ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  EXPECT_THAT(AlgebraicSimplifier(default_options_).Run(m.get()),
+              absl_testing::IsOkAndHolds(false));
+
+  AlgebraicSimplifierOptions options = default_options_;
+  options.set_enable_fast_math(true);
+  ASSERT_THAT(AlgebraicSimplifier(options).Run(m.get()),
+              absl_testing::IsOkAndHolds(true));
   EXPECT_THAT(
       m->entry_computation()->root_instruction(),
       GmockMatch(m::Multiply(
@@ -1128,8 +1158,14 @@ TEST_F(AlgebraicSimplifierTest,
       ROOT multiply1 = f32[4] multiply(multiply0, b1)
     }
   )";
-  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
-  ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).value());
+  ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  EXPECT_THAT(AlgebraicSimplifier(default_options_).Run(m.get()),
+              absl_testing::IsOkAndHolds(false));
+
+  AlgebraicSimplifierOptions options = default_options_;
+  options.set_enable_fast_math(true);
+  ASSERT_THAT(AlgebraicSimplifier(options).Run(m.get()),
+              absl_testing::IsOkAndHolds(true));
   EXPECT_THAT(
       m->entry_computation()->root_instruction(),
       GmockMatch(m::Multiply(
@@ -2141,14 +2177,63 @@ TEST_F(AlgebraicSimplifierTest, AddReassociateMergeConstants) {
   auto computation = m->AddEntryComputationWithLayouts(builder.Build());
   HloInstruction* root = computation->root_instruction();
   EXPECT_EQ(root->opcode(), HloOpcode::kAdd);
-  AlgebraicSimplifier simplifier(default_options_);
-  ASSERT_TRUE(simplifier.Run(m.get()).value());
+  EXPECT_THAT(AlgebraicSimplifier(default_options_).Run(m.get()),
+              absl_testing::IsOkAndHolds(false));
+
+  AlgebraicSimplifierOptions options = default_options_;
+  options.set_enable_fast_math(true);
+  ASSERT_THAT(AlgebraicSimplifier(options).Run(m.get()),
+              absl_testing::IsOkAndHolds(true));
   root = computation->root_instruction();
   EXPECT_THAT(root, GmockMatch(m::Add(
                         m::Op().Is(param0),
                         m::Add(m::Op().Is(constant1), m::Op().Is(constant2)))));
 }
 
+TEST_F(AlgebraicSimplifierTest, AddReassociateMergeConstantsComplex) {
+  constexpr absl::string_view kModuleStr = R"(
+    HloModule m
+    test {
+      p0 = c64[] parameter(0)
+      c0 = c64[] constant((1.0, 2.0))
+      c1 = c64[] constant((3.0, 4.0))
+      add0 = c64[] add(p0, c0)
+      ROOT add1 = c64[] add(add0, c1)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  EXPECT_THAT(AlgebraicSimplifier(default_options_).Run(m.get()),
+              absl_testing::IsOkAndHolds(false));
+
+  AlgebraicSimplifierOptions options = default_options_;
+  options.set_enable_fast_math(true);
+  ASSERT_THAT(AlgebraicSimplifier(options).Run(m.get()),
+              absl_testing::IsOkAndHolds(true));
+  EXPECT_THAT(m->entry_computation()->root_instruction(),
+              GmockMatch(m::Add(m::Parameter(0),
+                                m::Add(m::Constant(), m::Constant()))));
+}
+
+TEST_F(AlgebraicSimplifierTest, AddReassociateMergeConstantsInt) {
+  constexpr absl::string_view kModuleStr = R"(
+    HloModule m
+    test {
+      p0 = s32[] parameter(0)
+      c0 = s32[] constant(42)
+      c1 = s32[] constant(3)
+      add0 = s32[] add(p0, c0)
+      ROOT add1 = s32[] add(add0, c1)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  ASSERT_THAT(AlgebraicSimplifier(default_options_).Run(m.get()),
+              absl_testing::IsOkAndHolds(true));
+  EXPECT_THAT(
+      m->entry_computation()->root_instruction(),
+      GmockMatch(m::Add(m::Parameter(0),
+                        m::Add(m::ConstantScalar(42), m::ConstantScalar(3)))));
+}
+
 TEST_F(AlgebraicSimplifierTest, AddReassociateMergeBroadcastedConstants) {
   constexpr absl::string_view kModuleStr = R"(
     HloModule m
@@ -2162,14 +2247,42 @@ TEST_F(AlgebraicSimplifierTest, AddReassociateMergeBroadcastedConstants) {
       ROOT add1 = f32[4] add(add0, b1)
     }
   )";
-  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
-  ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).value());
+  ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  EXPECT_THAT(AlgebraicSimplifier(default_options_).Run(m.get()),
+              absl_testing::IsOkAndHolds(false));
+
+  AlgebraicSimplifierOptions options = default_options_;
+  options.set_enable_fast_math(true);
+  ASSERT_THAT(AlgebraicSimplifier(options).Run(m.get()),
+              absl_testing::IsOkAndHolds(true));
   EXPECT_THAT(m->entry_computation()->root_instruction(),
               GmockMatch(m::Add(m::Parameter(0),
                                 m::Broadcast(m::Add(m::ConstantScalar(1.0),
                                                     m::ConstantScalar(2.0))))));
 }
 
+TEST_F(AlgebraicSimplifierTest, AddReassociateMergeBroadcastedConstantsInt) {
+  constexpr absl::string_view kModuleStr = R"(
+    HloModule m
+    test {
+      p0 = s32[4] parameter(0)
+      c0 = s32[] constant(1)
+      c1 = s32[] constant(2)
+      b0 = s32[4] broadcast(c0), dimensions={}
+      b1 = s32[4] broadcast(c1), dimensions={}
+      add0 = s32[4] add(p0, b0)
+      ROOT add1 = s32[4] add(add0, b1)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  ASSERT_THAT(AlgebraicSimplifier(default_options_).Run(m.get()),
+              absl_testing::IsOkAndHolds(true));
+  EXPECT_THAT(m->entry_computation()->root_instruction(),
+              GmockMatch(m::Add(m::Parameter(0),
+                                m::Broadcast(m::Add(m::ConstantScalar(1),
+                                                    m::ConstantScalar(2))))));
+}
+
 TEST_F(AlgebraicSimplifierTest, ReplaceSubtractOfEqualOperandsWithZero) {
   constexpr absl::string_view kModuleStr = R"(
     HloModule m
@@ -2216,8 +2329,14 @@ TEST_F(AlgebraicSimplifierTest, SubAddReassociateMergeConstants) {
       ROOT add = f32[] add(sub, c1)
     }
   )";
-  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
-  ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).value());
+  ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  EXPECT_THAT(AlgebraicSimplifier(default_options_).Run(m.get()),
+              absl_testing::IsOkAndHolds(false));
+
+  AlgebraicSimplifierOptions options = default_options_;
+  options.set_enable_fast_math(true);
+  ASSERT_THAT(AlgebraicSimplifier(options).Run(m.get()),
+              absl_testing::IsOkAndHolds(true));
   EXPECT_THAT(m->entry_computation()->root_instruction(),
               GmockMatch(m::Subtract(
                   m::Add(m::ConstantScalar(1.0), m::ConstantScalar(2.0)),
@@ -2254,8 +2373,14 @@ TEST_F(AlgebraicSimplifierTest, SubAddReassociateMergeBroadcastedConstants) {
       ROOT add = f32[4] add(sub, b1)
     }
   )";
-  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
-  ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).value());
+  ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  EXPECT_THAT(AlgebraicSimplifier(default_options_).Run(m.get()),
+              absl_testing::IsOkAndHolds(false));
+
+  AlgebraicSimplifierOptions options = default_options_;
+  options.set_enable_fast_math(true);
+  ASSERT_THAT(AlgebraicSimplifier(options).Run(m.get()),
+              absl_testing::IsOkAndHolds(true));
   EXPECT_THAT(
       m->entry_computation()->root_instruction(),
       GmockMatch(m::Subtract(

From 51570f43cfb93ecee025a653cf8eb0a03583a01e Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Tue, 13 Oct 2026 14:56:15 -0700
Subject: [PATCH 0719/1000] Build the actual matrix sizes once in XLA
 MatrixDiagPart

Build the actual row and column counts once, instead of again for
each diagonal of a band, and still only when a matrix dimension is
dynamic.
---
 tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc b/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc
index 0f08021416f..19518533c02 100644
--- a/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc
@@ -440,11 +440,13 @@ class MatrixDiagPartOp : public XlaOpKernel {
                        context->builder(),
                        static_cast<int32_t>(input_shape.dim_size(dim)));
     };
+    // Only built with dynamic matrix dimensions, so that static inputs build
+    // the same HLO as before.
+    xla::XlaOp actual_num_rows;
+    xla::XlaOp actual_num_cols;
     // The length of the longest diagonal from lower to upper at the actual
     // sizes.
     const auto actual_len = [&](int64_t lower, int64_t upper) {
-      const xla::XlaOp actual_num_rows = actual_size(input_rank - 2);
-      const xla::XlaOp actual_num_cols = actual_size(input_rank - 1);
       return xla::Min(
           xla::Add(actual_num_rows,
                    xla::ScalarLike(actual_num_rows,
@@ -455,6 +457,8 @@ class MatrixDiagPartOp : public XlaOpKernel {
     };
     xla::XlaOp actual_max_diag_len;
     if (dynamic_matrix) {
+      actual_num_rows = actual_size(input_rank - 2);
+      actual_num_cols = actual_size(input_rank - 1);
       actual_max_diag_len = actual_len(lower_diag_index, upper_diag_index);
     }
 

From aeeec7987d75f415d109b26cba325ee5e793b277 Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Tue, 13 Oct 2026 15:01:47 -0700
Subject: [PATCH 0720/1000] Stop exporting the bundled curl from
 libtensorflow_framework

libtensorflow_framework exports the functions of the curl it links
statically: 94 public curl_* ones, and 663 Curl_* and 18 curlx_*
internal ones. Nothing outside the library calls them, and exported,
they could stand in for another libcurl's, such as the system's, in
the same process, with TensorFlow's BoringSSL behind them. Make them
local in the version script, as the LLVM C API and the hwloc names
already are, and add them to the macOS list of private symbols.

BoringSSL's symbols stay exported: libtensorflow_cc imports some of
them from this library.
---
 tensorflow/tf_framework_version_script.lds | 9 +++++++++
 tensorflow/tf_private_symbols.lds          | 3 +++
 2 files changed, 12 insertions(+)

diff --git a/tensorflow/tf_framework_version_script.lds b/tensorflow/tf_framework_version_script.lds
index 1d2a8351f41..6137f6b6e81 100644
--- a/tensorflow/tf_framework_version_script.lds
+++ b/tensorflow/tf_framework_version_script.lds
@@ -21,4 +21,13 @@ tensorflow {
        and can still be shared across TF DSOs (tensorflow#125854). */
     hwloc*;
     LLVM*;
+
+    # The statically linked curl: its public curl_* functions and its
+    # internal Curl_* and curlx_* ones. Nothing outside this library calls
+    # them, and exported, they could stand in for another libcurl's, such as
+    # the system's, in the same process, with TensorFlow's BoringSSL behind
+    # them (tensorflow#129050).
+    curl_*;
+    Curl_*;
+    curlx_*;
 };
diff --git a/tensorflow/tf_private_symbols.lds b/tensorflow/tf_private_symbols.lds
index 264a68e0d7f..7484e64155a 100644
--- a/tensorflow/tf_private_symbols.lds
+++ b/tensorflow/tf_private_symbols.lds
@@ -8,3 +8,6 @@ _jsimd_*
 _hwloc_*
 _LLVM*
 *mlir*
+_curl_*
+_Curl_*
+_curlx_*

From cc26da544f2a718fe7940a508d538f2c9117e59b Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Tue, 13 Oct 2026 15:22:32 -0700
Subject: [PATCH 0721/1000] Restore 64-bit axis handling and complete test edge
 coverage

The rebase onto master restored the 32-bit `int axis` declaration that
truncates 64-bit axis values (e.g. `1 << 33` wrapping to 0), and
dropped the error wording/logging of the originally supplied axis plus
the input rank. Re-apply the int64 axis handling and message.

Also add the rolls on the empty dimension itself (axis=0 / axis=2),
the negative 64-bit boundary `-(1 << 33)`, and a graph-mode
tensor-axis error case to the NumPy-level test suite.
---
 tensorflow/core/kernels/roll_op.cc                  | 10 ++++++----
 .../python/ops/numpy_ops/np_array_ops_test.py       | 13 ++++++++++++-
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/tensorflow/core/kernels/roll_op.cc b/tensorflow/core/kernels/roll_op.cc
index 7a0629eeb5d..7bf92351d38 100644
--- a/tensorflow/core/kernels/roll_op.cc
+++ b/tensorflow/core/kernels/roll_op.cc
@@ -71,13 +71,15 @@ class RollOp : public OpKernel {
     // total modulo sum of shifts for each dimension
     absl::InlinedVector<int64_t, 4> shift_mod_sum(num_dims, 0);
     for (int i = 0; i < num_shifts; i++) {
-      int axis = axis_flat(i);
+      int64_t axis = axis_flat(i);
       if (axis < 0) {
         axis += num_dims;
       }
-      OP_REQUIRES(context, FastBoundsCheck(axis, num_dims),
-                  absl::InvalidArgumentError(
-                      absl::StrCat("axis ", axis, " is out of range")));
+      OP_REQUIRES(
+          context, FastBoundsCheck(axis, num_dims),
+          absl::InvalidArgumentError(absl::StrCat(
+              "axis ", axis_flat(i), " is out of range for input of rank ",
+              num_dims, ".")));
       const int64_t ds = std::max<int64_t>(input.dim_size(axis), 1);
       // reduce the shift first so the sum cannot overflow int64_t
       const int64_t sum =
diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index 3a09a5c0a33..a2ce3355c99 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1682,9 +1682,12 @@ class ArrayMethodsTest(test.TestCase):
     _test(a, [1, -2], axis=[-3, 2])
     _test(a, np.array([0, 1]), axis=np.array([-1, 1]))
 
-    # Empty and scalar arrays.
+    # Empty and scalar arrays, including rolls on the empty dimension itself
+    # (exercising the max(dim_size, 1) safeguard).
+    _test(np.empty((0, 3, 4)), 1, axis=0)
     _test(np.empty((0, 3, 4)), 1, axis=1)
     _test(np.empty((3, 4, 0)), 1, axis=1)
+    _test(np.empty((3, 4, 0)), 1, axis=2)
     _test(np.array(42.0), 1)
 
     # Out-of-bounds axes raise, checked by the backend kernel.
@@ -1712,6 +1715,10 @@ class ArrayMethodsTest(test.TestCase):
         errors_impl.InvalidArgumentError, 'out of range'
     ):
       np_array_ops.roll(a, 1, axis=1 << 33)
+    with self.assertRaisesRegex(
+        errors_impl.InvalidArgumentError, 'out of range'
+    ):
+      np_array_ops.roll(a, 1, axis=-(1 << 33))
 
     # Graph-mode tracing of roll with concrete axes.
     @def_function.function
@@ -1730,6 +1737,10 @@ class ArrayMethodsTest(test.TestCase):
         errors_impl.InvalidArgumentError, 'out of range'
     ):
       f(a, 1, 3)
+    with self.assertRaisesRegex(
+        errors_impl.InvalidArgumentError, 'out of range'
+    ):
+      f(a, constant_op.constant(1), constant_op.constant(3))
 
   def testNdim(self):
     self.assertAllEqual(0, np_array_ops.ndim(0.5))

From 5edc4bcc98ddaa67a8a28297c2afb354b655c24e Mon Sep 17 00:00:00 2001
From: Bill Varcho <varcho@google.com>
Date: Tue, 13 Oct 2026 15:26:42 -0700
Subject: [PATCH 0722/1000] [SDY][optimization] Consolidate multiple module
 walks in ExportStablehloShardingsPass into a single pass.

PiperOrigin-RevId: 998929769
---
 .../stablehlo_round_trip/export_shardings.cc  | 65 +++++++++----------
 1 file changed, 30 insertions(+), 35 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc
index df36194d8a0..3540a6ae657 100644
--- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc
+++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc
@@ -32,6 +32,7 @@ limitations under the License.
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/AffineMap.h"
@@ -310,32 +311,9 @@ class ExportStablehloShardingsPass
                  clearReverseOpSharding);
     }
 
-    moduleOp.walk([&](stablehlo::CustomCallOp customCall) {
-      // StableHLO doesn't have an equivalent of `erf` and `topk` ops.
-      // If they have a sharding annotation, we need to move it into
-      // `mhlo.attributes`, which StableHLO->MHLO conversion would lift back up.
-      StringRef callTargetName = customCall.getCallTargetName();
-      if (callTargetName != "mhlo.erf" && callTargetName != "mhlo.topk") {
-        return;
-      }
-      // TODO(bartchr): refactor `addFrontendAttribute` to take a key for the
-      // dictionary attribute. Then can re-use the logic instead of duplicating
-      // it here for `kMhloAttributesAttr`.
-      if (auto sdySharding =
-              customCall->getAttrOfType<StringAttr>(kXlaShardingAttr)) {
-        customCall->removeAttr(kXlaShardingAttr);
-        SmallVector<mlir::NamedAttribute> newAttributes(
-            customCall->getAttrOfType<DictionaryAttr>(kMhloAttributesAttr)
-                .getValue());
-        newAttributes.push_back(
-            builder.getNamedAttr(kXlaShardingAttr, sdySharding));
-        customCall->setAttr(kMhloAttributesAttr,
-                            builder.getDictionaryAttr(newAttributes));
-      }
-    });
     // Collect all used mesh names from StableHLO attributes
     llvm::DenseSet<StringRef> usedMeshNames;
-    auto collectMeshesFromReplicaGroups = [&](auto op) {
+    auto collectMeshesFromReplicaGroups = [&](Operation* op) {
       if (auto attr = op->getAttr("replica_groups")) {
         if (auto rgv3 =
                 mlir::dyn_cast<stablehlo::ReplicaGroupMeshAxesAttr>(attr)) {
@@ -350,17 +328,34 @@ class ExportStablehloShardingsPass
       }
     };
 
-    moduleOp.walk(
-        [&](stablehlo::AllGatherOp op) { collectMeshesFromReplicaGroups(op); });
-    moduleOp.walk(
-        [&](stablehlo::AllReduceOp op) { collectMeshesFromReplicaGroups(op); });
-    moduleOp.walk([&](stablehlo::ReduceScatterOp op) {
-      collectMeshesFromReplicaGroups(op);
-    });
-    moduleOp.walk(
-        [&](stablehlo::AllToAllOp op) { collectMeshesFromReplicaGroups(op); });
-    moduleOp.walk([&](stablehlo::CollectiveBroadcastOp op) {
-      collectMeshesFromReplicaGroups(op);
+    moduleOp.walk([&](Operation* op) {
+      if (auto customCall = dyn_cast<stablehlo::CustomCallOp>(op)) {
+        // StableHLO doesn't have an equivalent of `erf` and `topk` ops.
+        // If they have a sharding annotation, we need to move it into
+        // `mhlo.attributes`, which StableHLO->MHLO conversion would lift back
+        // up.
+        StringRef callTargetName = customCall.getCallTargetName();
+        if (callTargetName == "mhlo.erf" || callTargetName == "mhlo.topk") {
+          // TODO(bartchr): refactor `addFrontendAttribute` to take a key for
+          // the dictionary attribute. Then can re-use the logic instead of
+          // duplicating it here for `kMhloAttributesAttr`.
+          if (auto sdySharding =
+                  customCall->getAttrOfType<StringAttr>(kXlaShardingAttr)) {
+            customCall->removeAttr(kXlaShardingAttr);
+            SmallVector<mlir::NamedAttribute> newAttributes(
+                customCall->getAttrOfType<DictionaryAttr>(kMhloAttributesAttr)
+                    .getValue());
+            newAttributes.push_back(
+                builder.getNamedAttr(kXlaShardingAttr, sdySharding));
+            customCall->setAttr(kMhloAttributesAttr,
+                                builder.getDictionaryAttr(newAttributes));
+          }
+        }
+      } else if (isa<stablehlo::AllGatherOp, stablehlo::AllReduceOp,
+                     stablehlo::ReduceScatterOp, stablehlo::AllToAllOp,
+                     stablehlo::CollectiveBroadcastOp>(op)) {
+        collectMeshesFromReplicaGroups(op);
+      }
     });
 
     // Remove all mesh symbols that are not used

From de52d9c15ecb24c2f3af36df27864bc4f6c36d7b Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Tue, 13 Oct 2026 15:33:51 -0700
Subject: [PATCH 0723/1000] Use a C-style comment for the curl symbols in the
 version script

It works in every linker's version scripts and matches the hwloc
comment above it.
---
 tensorflow/tf_framework_version_script.lds | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tensorflow/tf_framework_version_script.lds b/tensorflow/tf_framework_version_script.lds
index 6137f6b6e81..7f4bb14e22f 100644
--- a/tensorflow/tf_framework_version_script.lds
+++ b/tensorflow/tf_framework_version_script.lds
@@ -22,11 +22,11 @@ tensorflow {
     hwloc*;
     LLVM*;
 
-    # The statically linked curl: its public curl_* functions and its
-    # internal Curl_* and curlx_* ones. Nothing outside this library calls
-    # them, and exported, they could stand in for another libcurl's, such as
-    # the system's, in the same process, with TensorFlow's BoringSSL behind
-    # them (tensorflow#129050).
+    /* The statically linked curl: its public curl_* functions and its
+       internal Curl_* and curlx_* ones. Nothing outside this library calls
+       them, and exported, they could stand in for another libcurl's, such as
+       the system's, in the same process, with TensorFlow's BoringSSL behind
+       them (tensorflow#129050). */
     curl_*;
     Curl_*;
     curlx_*;

From 64be59dca5cc6acdbb37be26a573df8270b54b64 Mon Sep 17 00:00:00 2001
From: Ibrahim Umit Akgun <umitakgun@google.com>
Date: Tue, 13 Oct 2026 15:31:03 -0700
Subject: [PATCH 0724/1000] [XLA:GPU] Avoid revisiting HLO cost analysis for
 unmodified operands in PriorityFusion.

In `PriorityFusionQueue::OnFusingInstruction`, `to_update_priority_` collects both the newly created or modified `fusion` instruction and all of its fusible operands whose priorities need to be recomputed because their consumer changed.

Previously, `PriorityFusionQueue::UpdatePriorities` called `cost_analysis_->RevisitInstruction` and `UpdatePerformanceModelCache` on every instruction in `to_update_priority_`. However, unmodified operands only had their user lists updated; neither `GpuHloCostAnalysis::RevisitInstruction` nor `CombinedGpuPerformanceModel::EstimateRunTimeForInstruction` depends on an instruction's users, and `GpuPerformanceModelCache::Invalidate` preserves the standalone runtime cache entries of unmodified operands. Revisiting unmodified operands (which are frequently large fusions themselves) repeatedly re-traversed their fused subcomputations at every downstream fusion step.

Track the newly created or modified fusions in `updated_fusions_` and only run `cost_analysis_->RevisitInstruction` and `UpdatePerformanceModelCache` on `updated_fusions_` before recomputing priorities for `to_update_priority_`. All unit tests pass, and XLA:GPU H100 benchmarks show no runtime performance degradation (1.00x geomean across device and wall time).

PiperOrigin-RevId: 998932267
---
 .../xla/xla/backends/gpu/transforms/BUILD     |  3 +
 .../gpu/transforms/priority_fusion.cc         | 20 +++--
 .../gpu/transforms/priority_fusion_test.cc    | 82 +++++++++++++++++++
 3 files changed, 99 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD
index 6aa224982d8..48b69fe4626 100644
--- a/third_party/xla/xla/backends/gpu/transforms/BUILD
+++ b/third_party/xla/xla/backends/gpu/transforms/BUILD
@@ -2611,6 +2611,7 @@ cc_library(
         "//xla/service/gpu:gpu_fusible",
         "//xla/service/gpu:hlo_fusion_analysis",
         "//xla/service/gpu:ir_emission_utils",
+        "//xla/service/gpu:mlir_context_pool",
         "//xla/service/gpu/model:combined_gpu_performance_model",
         "//xla/service/gpu/model:fusion_analysis_cache",
         "//xla/service/gpu/model:gpu_hlo_cost_analysis",
@@ -2623,6 +2624,7 @@ cc_library(
         "//xla/tsl/concurrency:future",
         "//xla/tsl/platform:env",
         "//xla/tsl/platform:logging",
+        "//xla/tsl/util:sorted_range",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/base:nullability",
@@ -2647,6 +2649,7 @@ xla_cc_test(
     srcs = ["priority_fusion_test.cc"],
     deps = [
         ":priority_fusion",
+        "//xla:shape_util",
         "//xla:xla_proto_cc",
         "//xla/backends/gpu/codegen/emitters:mlir_kernel_emitter",
         "//xla/codegen/xtile:xtile_config_proto_cc",
diff --git a/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc b/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc
index d0ec57af11c..8a6f2334e64 100644
--- a/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc
@@ -63,6 +63,7 @@ limitations under the License.
 #include "xla/service/gpu/gpu_fusible.h"
 #include "xla/service/gpu/hlo_fusion_analysis.h"
 #include "xla/service/gpu/ir_emission_utils.h"
+#include "xla/service/gpu/mlir_context_pool.h"
 #include "xla/service/gpu/model/combined_gpu_performance_model.h"
 #include "xla/service/gpu/model/fusion_analysis_cache.h"
 #include "xla/service/gpu/model/gpu_hlo_cost_analysis.h"
@@ -79,6 +80,7 @@ limitations under the License.
 #include "xla/tsl/concurrency/future.h"
 #include "xla/tsl/platform/logging.h"
 #include "xla/tsl/platform/threadpool.h"
+#include "xla/tsl/util/sorted_range.h"
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
 
@@ -402,19 +404,19 @@ class PriorityFusionQueue {
 
   // Update priorities of all affected ops.
   absl::Status UpdatePriorities() {
-    // Revisit costs of all updated ops. It's important to update cost analysis
-    // before recalculating priorities.
-    for (auto instruction : to_update_priority_) {
+    // Revisit costs of all updated fusion ops. Unmodified operands in
+    // `to_update_priority_` only had their consumers updated; their own HLO
+    // costs and standalone runtime estimates have not changed.
+    for (HloInstruction* instruction : tsl::SortedRange(updated_fusions_)) {
       ABSL_RETURN_IF_ERROR(cost_analysis_->RevisitInstruction(instruction));
-    }
-    for (auto producer : to_update_priority_) {
-      ABSL_RETURN_IF_ERROR(UpdatePerformanceModelCache(producer));
+      ABSL_RETURN_IF_ERROR(UpdatePerformanceModelCache(instruction));
     }
 
     ABSL_RETURN_IF_ERROR(ComputeAndSetPriorities(std::vector<HloInstruction*>{
         to_update_priority_.begin(), to_update_priority_.end()}));
 
     to_update_priority_.clear();
+    updated_fusions_.clear();
     operands_to_new_consumers_.clear();
     operands_to_removed_consumers_runtimes_.clear();
     return absl::OkStatus();
@@ -583,11 +585,13 @@ class PriorityFusionQueue {
     // fuse. For now, as we only allow multi-output fusion if there is just a
     // single fusible consumer, this is not needed.
     to_update_priority_.insert(fusion);
+    updated_fusions_.insert(fusion);
   }
 
   // Removes data for the instruction.
   void RemoveInstruction(HloInstruction* instruction) {
     to_update_priority_.erase(instruction);
+    updated_fusions_.erase(instruction);
 
     auto reverse_it = reverse_map_.find(instruction);
     if (reverse_it == reverse_map_.end()) {
@@ -1252,6 +1256,10 @@ class PriorityFusionQueue {
   // avoid recomputing priorities multiple times before we dequeue a new
   // producer.
   absl::flat_hash_set<HloInstruction*> to_update_priority_;
+  // The subset of `to_update_priority_` consisting of newly created or modified
+  // fusion instructions whose HLO cost analysis and standalone performance
+  // model cache entries need to be updated before recomputing priorities.
+  absl::flat_hash_set<HloInstruction*> updated_fusions_;
   absl::flat_hash_map<HloInstruction*, std::vector<HloInstruction*>>
       operands_to_new_consumers_;
   absl::flat_hash_map<HloInstruction*, GpuPerformanceModel::RunTimes>
diff --git a/third_party/xla/xla/backends/gpu/transforms/priority_fusion_test.cc b/third_party/xla/xla/backends/gpu/transforms/priority_fusion_test.cc
index c3006bca1d9..c73af235746 100644
--- a/third_party/xla/xla/backends/gpu/transforms/priority_fusion_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/priority_fusion_test.cc
@@ -46,6 +46,8 @@ limitations under the License.
 #include "xla/service/hlo_cost_analysis.h"
 #include "xla/service/hlo_module_config.h"
 #include "xla/service/pattern_matcher.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/threadpool.h"
@@ -1764,5 +1766,85 @@ CHECK: ROOT %[[EPILOGUE_FUSION:.*]] = f32[100]{0} fusion(%[[SCAN_FUSION]]), kind
       )");
 }
 
+TEST_F(PriorityFusionTest, DoesNotRevisitCostAnalysisForUnmodifiedOperands) {
+  // Construct an HLO graph where `unfused_producer` is a kInput reduction
+  // fusion with a distinct internal broadcast shape `f32[17,33]`. It feeds a
+  // chain of downstream consumers (`consumer_red` -> `add1` -> `mul1` -> `add2`
+  // -> `exp1`) that get fused together over multiple PriorityFusion steps.
+  // Because both `unfused_producer` and the consumer chain contain reductions,
+  // `unfused_producer` cannot fuse into the consumer fusion and remains an
+  // unmodified operand across every consumer fusion step.
+  constexpr absl::string_view kHlo = R"(
+    HloModule test_module
+
+    add_scalar {
+      a = f32[] parameter(0)
+      b = f32[] parameter(1)
+      ROOT sum = f32[] add(a, b)
+    }
+
+    producer_R {
+      p = f32[17] parameter(0)
+      bcast = f32[17,33] broadcast(p), dimensions={0}
+      z = f32[] constant(0)
+      ROOT red = f32[17] reduce(bcast, z), dimensions={1}, to_apply=add_scalar
+    }
+
+    ENTRY main {
+      p0 = f32[17] parameter(0)
+      p1 = f32[17,32] parameter(1)
+      p2 = f32[17] parameter(2)
+      z = f32[] constant(0)
+      unfused_producer = f32[17] fusion(p0), kind=kInput, calls=producer_R
+      consumer_red = f32[17] reduce(p1, z), dimensions={1}, to_apply=add_scalar
+      add1 = f32[17] add(consumer_red, unfused_producer)
+      mul1 = f32[17] multiply(add1, p2)
+      add2 = f32[17] add(mul1, unfused_producer)
+      ROOT exp1 = f32[17] exponential(add2)
+    })";
+
+  // Count how many times `GpuHloCostAnalysis` queries the byte size of the
+  // internal shape `f32[17,33]` inside `producer_R`. This shape only exists
+  // inside `unfused_producer`'s fused computation, so any call to
+  // `cost_analysis_->RevisitInstruction(unfused_producer)` re-traverses
+  // `producer_R` and increments `internal_shape_size_calls`.
+  int internal_shape_size_calls = 0;
+  GpuHloCostAnalysis::Options options;
+  options.count_multiple_input_accesses = true;
+  options.shape_size = [&](const Shape& shape) {
+    if (shape.IsArray() && shape.dimensions().size() == 2 &&
+        shape.dimensions(0) == 17 && shape.dimensions(1) == 33) {
+      ++internal_shape_size_calls;
+    }
+    return ShapeUtil::ByteSizeOf(shape, /*pointer_size=*/8);
+  };
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kHlo));
+
+  // Measure how many times a single initial GpuHloCostAnalysis pass queries the
+  // internal shape `f32[17,33]` of `unfused_producer`.
+  GpuHloCostAnalysis baseline_analysis(options, device_info_);
+  ASSERT_THAT(module->entry_computation()->Accept(&baseline_analysis),
+              absl_testing::IsOk());
+  const int initial_calls = internal_shape_size_calls;
+  ASSERT_GT(initial_calls, 0);
+
+  internal_shape_size_calls = 0;
+  PriorityFusion priority_fusion(/*thread_pool=*/nullptr, device_info_,
+                                 &alias_info_, options, &mlir_context_);
+  EXPECT_THAT(priority_fusion.Run(module.get()),
+              absl_testing::IsOkAndHolds(true));
+
+  // `unfused_producer` is an operand of the consumer fusion that gets updated
+  // across multiple fusion steps (`consumer_red`, `add1`, `mul1`, `add2`,
+  // `exp1`), so its priority is recomputed at each step
+  // (`to_update_priority_`). However, `unfused_producer` itself is never
+  // modified (`updated_fusions_`), so its HLO cost analysis should only be
+  // visited once during initial queue creation and never revisited during
+  // `UpdatePriorities()`.
+  EXPECT_EQ(internal_shape_size_calls, initial_calls);
+}
+
 }  // namespace gpu
 }  // namespace xla

From 37ef9c93ee3f18df032a9b35bcdfbb925afe4ff9 Mon Sep 17 00:00:00 2001
From: Bill Varcho <varcho@google.com>
Date: Tue, 13 Oct 2026 15:43:04 -0700
Subject: [PATCH 0725/1000] [SDY][perf] Batch argument and result attribute
 updates in convertShardyAttrsWithHloShardingV3.

Each `setArgAttr`/`removeArgAttr` (and result equivalent) rebuilds the function's whole arg/result attribute array, so updating every argument one at a time is O(N^2) in the number of arguments. Build the dictionaries with `NamedAttrList` and set them once with `setAllArgAttrs`/`setAllResultAttrs`, matching `convertShardyAttrsWithoutHloShardingV3` in the same file.

PiperOrigin-RevId: 998939402
---
 .../sdy_round_trip/import_shardy_attrs.cc     | 25 ++++++++++++++-----
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc
index 4b3f5b085e2..58083e905f3 100644
--- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc
+++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc
@@ -242,32 +242,45 @@ ManualAxesAttr extractManualAxes(const xla::HloSharding& hloSharding,
 
 // The sharding information is in the `kXlaShardingAttr` attribute.
 void convertShardyAttrsWithHloShardingV3(FuncOp funcOp) {
+  // Replace `kXlaShardingAttr` with `kShardingAttr` on all args, then set all
+  // arg attrs at once. Per-arg `setArgAttr` would be O(N^2).
+  llvm::SmallVector<DictionaryAttr> funcArgAttrs;
+  funcArgAttrs.reserve(funcOp.getNumArguments());
   for (auto [argNum, argType] : llvm::enumerate(funcOp.getArgumentTypes())) {
+    mlir::NamedAttrList attrs(funcOp.getArgAttrDict(argNum));
     if (auto oldSharding =
-            funcOp.getArgAttrOfType<StringAttr>(argNum, kXlaShardingAttr)) {
+            mlir::dyn_cast_or_null<StringAttr>(attrs.get(kXlaShardingAttr))) {
       if (auto sdySharding = convertToSdyShardingAttr(
               parseShardingFromString(oldSharding),
               mlir::sdy::getTensorRank(argType), funcOp.getContext())) {
-        funcOp.setArgAttr(argNum, kShardingAttr, sdySharding);
+        attrs.set(kShardingAttr, sdySharding);
       }
     }
-    funcOp.removeArgAttr(argNum, kXlaShardingAttr);
+    attrs.erase(kXlaShardingAttr);
+    funcArgAttrs.push_back(attrs.getDictionary(funcOp.getContext()));
   }
+  funcOp.setAllArgAttrs(funcArgAttrs);
 
+  // Same for results, but single-device shardings are dropped.
+  llvm::SmallVector<DictionaryAttr> newResultAttrs;
+  newResultAttrs.reserve(funcOp.getNumResults());
   for (auto [resNum, resType] : llvm::enumerate(funcOp.getResultTypes())) {
+    mlir::NamedAttrList attrs(funcOp.getResultAttrDict(resNum));
     if (auto oldSharding =
-            funcOp.getResultAttrOfType<StringAttr>(resNum, kXlaShardingAttr)) {
+            mlir::dyn_cast_or_null<StringAttr>(attrs.get(kXlaShardingAttr))) {
       HloSharding hloSharding = parseShardingFromString(oldSharding);
       if (!hloSharding.IsSingleDevice()) {
         if (auto sdySharding = convertToSdyShardingAttr(
                 hloSharding, mlir::sdy::getTensorRank(resType),
                 funcOp.getContext())) {
-          funcOp.setResultAttr(resNum, kShardingAttr, sdySharding);
+          attrs.set(kShardingAttr, sdySharding);
         }
       }
     }
-    funcOp.removeResultAttr(resNum, kXlaShardingAttr);
+    attrs.erase(kXlaShardingAttr);
+    newResultAttrs.push_back(attrs.getDictionary(funcOp.getContext()));
   }
+  funcOp.setAllResultAttrs(newResultAttrs);
 
   if (funcOp.isExternal()) {
     return;

From 11881c6ffabf66ea0a23f784509681e25daab7ce Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Tue, 13 Oct 2026 16:28:03 -0700
Subject: [PATCH 0726/1000] [XLA] Fix integer power lowering for even bases
 with large exponents

`ElementalIrEmitter::EmitIntegerPow` and `mapMhloOpToStdScalarOp<mhlo::PowOp>` run six iterations of exponentiation by squaring, shifting `exponent` right by one bit on each iteration. When `exponent >= 64` and `lhs` is even, the result modulo two to the bit width must be zero. However, when the low six bits of `exponent` are zero, the loop never multiplies `accumulator` by `base` and returns `1`. When the lowest bit is one and bits one through five are zero, it returns `lhs` instead of `0`.

Select `0` for the accumulator after the loop whenever the shifted `exponent` is nonzero and the lowest bit of the original base is zero. For bases `1` and `-1`, `base` to the 64th power is `1`, so the loop result remains exact.

Existing programs that compute integer `power` with an even base and an exponent of at least 64 now receive `0`. Other inputs are unchanged.

Tests:
- `//third_party/tensorflow/compiler/xla/service:elemental_ir_emitter_test`: `IntegerPowEvenBaseExponentGeBitWidth` verifies that integer `power` returns `0` for even bases `{0, 2, -2}` with large exponents.
PiperOrigin-RevId: 998964788
---
 .../mhlo/transforms/map_mhlo_to_scalar_op.h   | 49 +++++++++++--------
 .../xla/xla/service/elemental_ir_emitter.cc   |  5 ++
 .../xla/service/elemental_ir_emitter_test.cc  | 17 +++++++
 3 files changed, 50 insertions(+), 21 deletions(-)

diff --git a/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_mhlo_to_scalar_op.h b/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_mhlo_to_scalar_op.h
index 1d170bc7faa..bc16e71d045 100644
--- a/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_mhlo_to_scalar_op.h
+++ b/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_mhlo_to_scalar_op.h
@@ -1198,28 +1198,35 @@ inline Value mapMhloOpToStdScalarOp<mhlo::PowOp>(
   auto originalBase = adaptor.getLhs();
   auto originalExponent = adaptor.getRhs();
 
-  Value accum =
-      scf::ForOp::create(
-          lb, lowerBound, upperBound, step,
-          SmallVector<Value>({one, originalBase, originalExponent}),
-          [&](OpBuilder& b, Location, Value /*v*/, ValueRange iters) {
-            Value accum = iters[0];
-            Value base = iters[1];
-            Value exponent = iters[2];
+  auto forOp = scf::ForOp::create(
+      lb, lowerBound, upperBound, step,
+      SmallVector<Value>({one, originalBase, originalExponent}),
+      [&](OpBuilder& b, Location, Value /*v*/, ValueRange iters) {
+        Value accum = iters[0];
+        Value base = iters[1];
+        Value exponent = iters[2];
 
-            Value condition = arith::CmpIOp::create(
-                b, loc, arith::CmpIPredicate::eq,
-                ::mlir::arith::AndIOp::create(b, loc, exponent, one), one);
-            Value multiplied =
-                ::mlir::arith::MulIOp::create(b, loc, accum, base);
-            accum = ::mlir::arith::SelectOp::create(b, loc, condition,
-                                                    multiplied, accum);
-            base = ::mlir::arith::MulIOp::create(b, loc, base, base);
-            exponent = ::mlir::arith::ShRUIOp::create(b, loc, exponent, one);
-            scf::YieldOp::create(b, loc,
-                                 SmallVector<Value>({accum, base, exponent}));
-          })
-          .getResult(0);
+        Value condition = arith::CmpIOp::create(
+            b, loc, arith::CmpIPredicate::eq,
+            ::mlir::arith::AndIOp::create(b, loc, exponent, one), one);
+        Value multiplied = ::mlir::arith::MulIOp::create(b, loc, accum, base);
+        accum = ::mlir::arith::SelectOp::create(b, loc, condition, multiplied,
+                                                accum);
+        base = ::mlir::arith::MulIOp::create(b, loc, base, base);
+        exponent = ::mlir::arith::ShRUIOp::create(b, loc, exponent, one);
+        scf::YieldOp::create(b, loc,
+                             SmallVector<Value>({accum, base, exponent}));
+      });
+  Value accum = forOp.getResult(0);
+  Value remainingExponent = forOp.getResult(2);
+  Value exponentNonZero = arith::CmpIOp::create(lb, arith::CmpIPredicate::ne,
+                                                remainingExponent, zero);
+  Value baseIsEven = arith::CmpIOp::create(
+      lb, arith::CmpIPredicate::eq,
+      ::mlir::arith::AndIOp::create(lb, originalBase, one), zero);
+  accum = ::mlir::arith::SelectOp::create(
+      lb, ::mlir::arith::AndIOp::create(lb, exponentNonZero, baseIsEven), zero,
+      accum);
 
   if (IsUnsignedIntegerType{}(getElementTypeOrSelf(argTypes.front()))) {
     return accum;
diff --git a/third_party/xla/xla/service/elemental_ir_emitter.cc b/third_party/xla/xla/service/elemental_ir_emitter.cc
index 3307adef641..b76c9f2d6ac 100644
--- a/third_party/xla/xla/service/elemental_ir_emitter.cc
+++ b/third_party/xla/xla/service/elemental_ir_emitter.cc
@@ -2511,6 +2511,11 @@ llvm::Value* ElementalIrEmitter::EmitIntegerPow(llvm::Value* lhs,
     base = b_->CreateMul(base, base);
     exponent = b_->CreateLShr(exponent, 1);
   }
+  accumulator = b_->CreateSelect(
+      b_->CreateAnd(
+          b_->CreateICmpNE(exponent, exp_zero),
+          b_->CreateICmpEQ(b_->CreateAnd(original_base, base_one), base_zero)),
+      base_zero, accumulator);
   if (!is_signed) {
     return accumulator;
   }
diff --git a/third_party/xla/xla/service/elemental_ir_emitter_test.cc b/third_party/xla/xla/service/elemental_ir_emitter_test.cc
index a07a8f63556..e667173f5ef 100644
--- a/third_party/xla/xla/service/elemental_ir_emitter_test.cc
+++ b/third_party/xla/xla/service/elemental_ir_emitter_test.cc
@@ -670,5 +670,22 @@ ENTRY main {
   RunTest(hlo_text, {&base, &exponent});
 }
 
+TEST_F(ElementalIrEmitterExecutionTest, IntegerPowEvenBaseExponentGeBitWidth) {
+  const std::string hlo_text = R"(
+HloModule IntegerPowEvenBaseExponentGeBitWidth
+
+ENTRY main {
+  base = s32[8]{0} parameter(0)
+  exponent = s32[8]{0} parameter(1)
+  ROOT power = s32[8]{0} power(base, exponent)
+}
+)";
+
+  Literal base = LiteralUtil::CreateR1<int32_t>({0, 0, 2, -2, 0, 2, -2, 0});
+  Literal exponent =
+      LiteralUtil::CreateR1<int32_t>({1, 32, 32, 33, 64, 64, 65, 1});
+  RunTest(hlo_text, {&base, &exponent});
+}
+
 }  // namespace
 }  // namespace xla

From 97c224626eeff18ad8e207b935874ec9632d0f96 Mon Sep 17 00:00:00 2001
From: Alex Pivovarov <upwind@google.com>
Date: Tue, 13 Oct 2026 16:39:46 -0700
Subject: [PATCH 0727/1000] [XLA:GPU] Use reshape instead of broadcast for
 batch=1 iota in approx_topk_special_floats_test.

In `xla/tests/approx_topk_special_floats_test.cc`, the input HLO previously used `broadcast(1D iota)` for both `batch = 1` and `batch = 8`.

However JAX emits different HLOs around iota for batch 1 and 8.

Match the unoptimized HLO emitted by JAX (`reshape(1D iota)` when `batch == 1`, and `broadcast(1D iota)` when `batch > 1`) so that the single `GpuAlgebraicSimplifier` pass folds `reshape(1D iota)` into a 2D `iota` before `TopkRewriter` runs, exercising `TopK` lowering for both `batch = 1` and `batch = 8`.

PiperOrigin-RevId: 998971277
---
 .../tests/approx_topk_special_floats_test.cc   | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/xla/tests/approx_topk_special_floats_test.cc b/third_party/xla/xla/tests/approx_topk_special_floats_test.cc
index 00cf1e14dc1..4fb23a53c9c 100644
--- a/third_party/xla/xla/tests/approx_topk_special_floats_test.cc
+++ b/third_party/xla/xla/tests/approx_topk_special_floats_test.cc
@@ -92,7 +92,8 @@ TEST_F(TopkTest, ApproxTopKWithSpecialFloats) {
       auto literal = LiteralUtil::CreateR2FromArray2D<float>(input);
 
       for (int k : {2, 4, 8, 16, 32, 64}) {
-        const auto hlo_text_module = absl::Substitute(R"(
+        const auto hlo_text_module = absl::Substitute(
+            R"(
       HloModule topk
 
       %top_k_gt_f32_comparator.4 (approx_top_k.35: f32[], approx_top_k.36: f32[], approx_top_k.37: s32[], approx_top_k.38: s32[]) -> pred[] {
@@ -106,14 +107,16 @@ TEST_F(TopkTest, ApproxTopKWithSpecialFloats) {
       ENTRY TopK {
         x = f32[$0,$1]{1,0} parameter(0)
         %approx_top_k.60 = s32[$1]{0} iota(), iota_dimension=0
-        %approx_top_k.61 = s32[$0,$1]{1,0} broadcast(%approx_top_k.60), dimensions={1}
+        %approx_top_k.61 = s32[$0,$1]{1,0} $3
         %approx_top_k.62 = (f32[$0,$1]{1,0}, s32[$0,$1]{1,0}) sort(x, %approx_top_k.61), dimensions={1}, is_stable=true, to_apply=%top_k_gt_f32_comparator.4
         %approx_top_k.64 = f32[$0,$2]{1,0} slice(%approx_top_k.62#0), slice={[0:$0], [0:$2]}
         %approx_top_k.66 = s32[$0,$2]{1,0} slice(%approx_top_k.62#1), slice={[0:$0], [0:$2]}
         ROOT %approx_top_k.67 = (f32[$0,$2]{1,0}, s32[$0,$2]{1,0}) tuple(%approx_top_k.64, %approx_top_k.66)
       }
     )",
-                                                      batch, n, k);
+            batch, n, k,
+            batch == 1 ? "reshape(%approx_top_k.60)"
+                       : "broadcast(%approx_top_k.60), dimensions={1}");
 
         ASSERT_OK_AND_ASSIGN(auto module,
                              ParseAndReturnVerifiedModule(hlo_text_module));
@@ -203,7 +206,8 @@ TEST_F(TopkTest, ApproxTopKWithSpecialBFloat16) {
       auto literal = LiteralUtil::CreateR2FromArray2D<bfloat16>(input);
 
       for (int k : {2, 4, 8, 16, 32, 64}) {
-        const auto hlo_text_module = absl::Substitute(R"(
+        const auto hlo_text_module = absl::Substitute(
+            R"(
       HloModule topk
 
       %top_k_gt_bf16_comparator.4 (approx_top_k.35: bf16[], approx_top_k.36: bf16[], approx_top_k.37: s32[], approx_top_k.38: s32[]) -> pred[] {
@@ -217,14 +221,16 @@ TEST_F(TopkTest, ApproxTopKWithSpecialBFloat16) {
       ENTRY TopK {
         x = bf16[$0,$1]{1,0} parameter(0)
         %approx_top_k.60 = s32[$1]{0} iota(), iota_dimension=0
-        %approx_top_k.61 = s32[$0,$1]{1,0} broadcast(%approx_top_k.60), dimensions={1}
+        %approx_top_k.61 = s32[$0,$1]{1,0} $3
         %approx_top_k.62 = (bf16[$0,$1]{1,0}, s32[$0,$1]{1,0}) sort(x, %approx_top_k.61), dimensions={1}, is_stable=true, to_apply=%top_k_gt_bf16_comparator.4
         %approx_top_k.64 = bf16[$0,$2]{1,0} slice(%approx_top_k.62#0), slice={[0:$0], [0:$2]}
         %approx_top_k.66 = s32[$0,$2]{1,0} slice(%approx_top_k.62#1), slice={[0:$0], [0:$2]}
         ROOT %approx_top_k.67 = (bf16[$0,$2]{1,0}, s32[$0,$2]{1,0}) tuple(%approx_top_k.64, %approx_top_k.66)
       }
     )",
-                                                      batch, n, k);
+            batch, n, k,
+            batch == 1 ? "reshape(%approx_top_k.60)"
+                       : "broadcast(%approx_top_k.60), dimensions={1}");
 
         ASSERT_OK_AND_ASSIGN(auto module,
                              ParseAndReturnVerifiedModule(hlo_text_module));

From d8623b220d1891d779b0436ff020cc46c2a09403 Mon Sep 17 00:00:00 2001
From: Esun Kim <veblush@google.com>
Date: Tue, 13 Oct 2026 17:35:03 -0700
Subject: [PATCH 0728/1000] This corrects a numerical precision issue
 identified in several TFLite reference kernels. The issue stems from a
 "narrowing" effect where intermediate variables intended for double precision
 calculations were being evaluated in float precision due to the placement of
 type casts.

PiperOrigin-RevId: 999000488
---
 tensorflow/lite/core/c/common.h               | 13 ++++++----
 tensorflow/lite/kernels/activations.cc        | 24 ++++++++++++-------
 tensorflow/lite/kernels/add.cc                | 11 +++++----
 tensorflow/lite/kernels/div.cc                |  6 ++---
 tensorflow/lite/kernels/elementwise.cc        |  7 ++++--
 tensorflow/lite/kernels/kernel_util.cc        |  4 ++--
 tensorflow/lite/kernels/lstm.cc               |  3 ++-
 tensorflow/lite/kernels/mul.cc                |  5 ++--
 tensorflow/lite/kernels/squared_difference.cc | 12 ++++++----
 tensorflow/lite/kernels/sub.cc                | 13 ++++++----
 tensorflow/lite/kernels/svdf.cc               | 14 ++++++-----
 11 files changed, 67 insertions(+), 45 deletions(-)

diff --git a/tensorflow/lite/core/c/common.h b/tensorflow/lite/core/c/common.h
index 66c8bf6fad6..ffc910d986c 100644
--- a/tensorflow/lite/core/c/common.h
+++ b/tensorflow/lite/core/c/common.h
@@ -269,11 +269,14 @@ void TfLiteFloatArrayFree(TfLiteFloatArray* a);
 
 #define TF_LITE_ENSURE_NEAR(context, a, b, epsilon)                          \
   do {                                                                       \
-    auto delta = ((a) > (b)) ? ((a) - (b)) : ((b) - (a));                    \
-    if (delta > epsilon) {                                                   \
-      TF_LITE_KERNEL_LOG((context), "%s:%d %s not near %s (%f != %f)",       \
-                         __FILE__, __LINE__, #a, #b, static_cast<double>(a), \
-                         static_cast<double>(b));                            \
+    const auto _val_a = (a);                                                 \
+    const auto _val_b = (b);                                                 \
+    const auto _eps = (epsilon);                                             \
+    const auto _delta = _val_a - _val_b;                                     \
+    if (_delta > _eps || _delta < -_eps) {                                   \
+      TF_LITE_KERNEL_LOG(                                                    \
+          (context), "%s:%d %s not near %s (%f != %f)", __FILE__, __LINE__,  \
+          #a, #b, static_cast<double>(_val_a), static_cast<double>(_val_b)); \
       return kTfLiteError;                                                   \
     }                                                                        \
   } while (0)
diff --git a/tensorflow/lite/kernels/activations.cc b/tensorflow/lite/kernels/activations.cc
index 20ffc244a31..b6d7c4cdbf8 100644
--- a/tensorflow/lite/kernels/activations.cc
+++ b/tensorflow/lite/kernels/activations.cc
@@ -231,7 +231,8 @@ TfLiteStatus ReluPrepare(TfLiteContext* context, TfLiteNode* node) {
 
   if (input->type == kTfLiteInt8 || input->type == kTfLiteUInt8 ||
       input->type == kTfLiteInt16) {
-    double real_multiplier = input->params.scale / output->params.scale;
+    double real_multiplier = static_cast<double>(input->params.scale) /
+                             static_cast<double>(output->params.scale);
     QuantizeMultiplier(real_multiplier, &data->output_multiplier,
                        &data->output_shift);
   }
@@ -315,11 +316,13 @@ TfLiteStatus LeakyReluPrepare(TfLiteContext* context, TfLiteNode* node) {
     const auto* params =
         reinterpret_cast<TfLiteLeakyReluParams*>(node->builtin_data);
 
-    double alpha_multiplier =
-        input->params.scale * params->alpha / output->params.scale;
+    double alpha_multiplier = static_cast<double>(input->params.scale) *
+                              static_cast<double>(params->alpha) /
+                              static_cast<double>(output->params.scale);
     QuantizeMultiplier(alpha_multiplier, &data->output_multiplier_alpha,
                        &data->output_shift_alpha);
-    double identity_multiplier = input->params.scale / output->params.scale;
+    double identity_multiplier = static_cast<double>(input->params.scale) /
+                                 static_cast<double>(output->params.scale);
     QuantizeMultiplier(identity_multiplier, &data->output_multiplier_identity,
                        &data->output_shift_identity);
   }
@@ -641,8 +644,9 @@ TfLiteStatus SoftmaxPrepare(TfLiteContext* context, TfLiteNode* node) {
     data->params.zero_point = output->params.zero_point;
     data->params.scale = output->params.scale;
 
-    double input_scale_beta_rescale =
-        input->params.scale * params->beta /
+    const double input_scale_beta_rescale =
+        static_cast<double>(input->params.scale) *
+        static_cast<double>(params->beta) /
         (10.0 / 65535.0);  // scale the input_diff such that [-65535, 0]
                            // correspond to [-10.0, 0.0]
     QuantizeMultiplier(input_scale_beta_rescale, &data->params.input_multiplier,
@@ -736,9 +740,11 @@ TfLiteStatus PreluPrepare(TfLiteContext* context, TfLiteNode* node) {
     // output real multiplier 1 is input_scale / output_scale;
     // for input_q - input_zp < 0:
     // output real multiplier 2 is input_scale  * alpha_scale/ output_scale.
-    double real_multiplier_1 = input->params.scale / output->params.scale;
-    double real_multiplier_2 =
-        input->params.scale * alpha->params.scale / output->params.scale;
+    double real_multiplier_1 = static_cast<double>(input->params.scale) /
+                               static_cast<double>(output->params.scale);
+    double real_multiplier_2 = static_cast<double>(input->params.scale) *
+                               static_cast<double>(alpha->params.scale) /
+                               static_cast<double>(output->params.scale);
     QuantizeMultiplier(real_multiplier_1, &data->output_multiplier_1,
                        &data->output_shift_1);
     QuantizeMultiplier(real_multiplier_2, &data->output_multiplier_2,
diff --git a/tensorflow/lite/kernels/add.cc b/tensorflow/lite/kernels/add.cc
index 9e5dfcd228c..ee61fd93de3 100644
--- a/tensorflow/lite/kernels/add.cc
+++ b/tensorflow/lite/kernels/add.cc
@@ -181,14 +181,15 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     // therefore the addition will still fit in a 32 bit accumulator.
     data->left_shift = general_scale_int16 ? 15 : 20;
     const double twice_max_input_scale =
-        2 * std::max(input1->params.scale, input2->params.scale);
+        2 * static_cast<double>(
+                std::max(input1->params.scale, input2->params.scale));
     const double real_input1_multiplier =
-        input1->params.scale / twice_max_input_scale;
+        static_cast<double>(input1->params.scale) / twice_max_input_scale;
     const double real_input2_multiplier =
-        input2->params.scale / twice_max_input_scale;
+        static_cast<double>(input2->params.scale) / twice_max_input_scale;
     const double real_output_multiplier =
-        twice_max_input_scale /
-        ((1 << data->left_shift) * output->params.scale);
+        twice_max_input_scale / (static_cast<double>(1 << data->left_shift) *
+                                 static_cast<double>(output->params.scale));
 
     QuantizeMultiplierSmallerThanOneExp(
         real_input1_multiplier, &data->input1_multiplier, &data->input1_shift);
diff --git a/tensorflow/lite/kernels/div.cc b/tensorflow/lite/kernels/div.cc
index d7cebed68f9..c5a59cd2494 100644
--- a/tensorflow/lite/kernels/div.cc
+++ b/tensorflow/lite/kernels/div.cc
@@ -104,8 +104,9 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     TF_LITE_ENSURE_STATUS(CalculateActivationRangeQuantized(
         context, params->activation, output, &data->output_activation_min,
         &data->output_activation_max));
-    const double real_multiplier =
-        input1->params.scale / (input2->params.scale * output->params.scale);
+    const double real_multiplier = static_cast<double>(input1->params.scale) /
+                                   (static_cast<double>(input2->params.scale) *
+                                    static_cast<double>(output->params.scale));
     QuantizeMultiplier(real_multiplier, &data->output_multiplier,
                        &data->output_shift);
   }
@@ -279,7 +280,6 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_OK(context,
                     GetOutputSafe(context, node, kOutputTensor, &output));
 
-
   if (output->type == kTfLiteFloat32) {
     // Div by zero seems ok in this case, we don't do a check at this point.
     // However, unlike in TF where infinities are returned, here we return an
diff --git a/tensorflow/lite/kernels/elementwise.cc b/tensorflow/lite/kernels/elementwise.cc
index 0bff391ffb3..cead70011b2 100644
--- a/tensorflow/lite/kernels/elementwise.cc
+++ b/tensorflow/lite/kernels/elementwise.cc
@@ -86,13 +86,16 @@ bool IsLogSupportedType(const TfLiteType type) {
 inline void SetAbsOutputMultiplier(const float input_scale,
                                    const float output_scale,
                                    int32_t* multiplier, int32_t* shift) {
-  QuantizeMultiplier(input_scale / output_scale, multiplier, shift);
+  QuantizeMultiplier(
+      static_cast<double>(input_scale) / static_cast<double>(output_scale),
+      multiplier, shift);
 }
 
 inline void SetRsqrtOutputMultiplier(const float input_scale,
                                      const float output_scale,
                                      int32_t* multiplier, int32_t* shift) {
-  const double scale = 1. / (std::sqrt(input_scale) * output_scale);
+  const double scale = 1.0 / (std::sqrt(static_cast<double>(input_scale)) *
+                              static_cast<double>(output_scale));
   QuantizeMultiplier(scale, multiplier, shift);
 }
 
diff --git a/tensorflow/lite/kernels/kernel_util.cc b/tensorflow/lite/kernels/kernel_util.cc
index 36ad5178158..3344dc2b13d 100644
--- a/tensorflow/lite/kernels/kernel_util.cc
+++ b/tensorflow/lite/kernels/kernel_util.cc
@@ -330,8 +330,8 @@ TfLiteStatus GetQuantizedConvolutionMultipler(TfLiteContext* context,
                                               const TfLiteTensor* filter,
                                               TfLiteTensor* output,
                                               double* multiplier) {
-  const double input_product_scale =
-      static_cast<double>(input->params.scale * filter->params.scale);
+  const double input_product_scale = static_cast<double>(input->params.scale) *
+                                     static_cast<double>(filter->params.scale);
   TF_LITE_ENSURE(context, input_product_scale >= 0);
   *multiplier = input_product_scale / static_cast<double>(output->params.scale);
 
diff --git a/tensorflow/lite/kernels/lstm.cc b/tensorflow/lite/kernels/lstm.cc
index 10dde7f4987..2c1114d48ce 100644
--- a/tensorflow/lite/kernels/lstm.cc
+++ b/tensorflow/lite/kernels/lstm.cc
@@ -2325,7 +2325,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
       return kTfLiteError;
     }
 
-    double real_accum_multiplier = 4096 * bias->params.scale;
+    double real_accum_multiplier =
+        4096.0 * static_cast<double>(bias->params.scale);
     int32 accum_multiplier;
     int accum_shift;
     tflite::QuantizeMultiplier(real_accum_multiplier, &accum_multiplier,
diff --git a/tensorflow/lite/kernels/mul.cc b/tensorflow/lite/kernels/mul.cc
index 73ffd3488f9..e2d516ac733 100644
--- a/tensorflow/lite/kernels/mul.cc
+++ b/tensorflow/lite/kernels/mul.cc
@@ -122,8 +122,9 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     TF_LITE_ENSURE_STATUS(CalculateActivationRangeQuantized(
         context, params->activation, output, &data->output_activation_min,
         &data->output_activation_max));
-    double real_multiplier =
-        input1->params.scale * input2->params.scale / output->params.scale;
+    double real_multiplier = static_cast<double>(input1->params.scale) *
+                             static_cast<double>(input2->params.scale) /
+                             static_cast<double>(output->params.scale);
     QuantizeMultiplier(real_multiplier, &data->output_multiplier,
                        &data->output_shift);
   }
diff --git a/tensorflow/lite/kernels/squared_difference.cc b/tensorflow/lite/kernels/squared_difference.cc
index 57bf92cc37c..21597d0f424 100644
--- a/tensorflow/lite/kernels/squared_difference.cc
+++ b/tensorflow/lite/kernels/squared_difference.cc
@@ -106,16 +106,18 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     // shift to make integer for scales.
     data->arithmetic_params.left_shift = 7;
     const double twice_max_input_scale =
-        2 * std::max(input1_quantization_params.scale,
-                     input2_quantization_params.scale);
+        2 * static_cast<double>(std::max(input1_quantization_params.scale,
+                                         input2_quantization_params.scale));
     const double real_input1_multiplier =
-        input1_quantization_params.scale / twice_max_input_scale;
+        static_cast<double>(input1_quantization_params.scale) /
+        twice_max_input_scale;
     double real_input2_multiplier =
-        input2_quantization_params.scale / twice_max_input_scale;
+        static_cast<double>(input2_quantization_params.scale) /
+        twice_max_input_scale;
     const double real_output_multiplier =
         (twice_max_input_scale * twice_max_input_scale) /
         ((1 << data->arithmetic_params.left_shift * 2) *
-         output_quantization_params.scale);
+         static_cast<double>(output_quantization_params.scale));
     tflite::QuantizeMultiplierSmallerThanOneExp(
         real_input1_multiplier, &data->arithmetic_params.input1_multiplier,
         &data->arithmetic_params.input1_shift);
diff --git a/tensorflow/lite/kernels/sub.cc b/tensorflow/lite/kernels/sub.cc
index 0136e94ab21..7cb5508227a 100644
--- a/tensorflow/lite/kernels/sub.cc
+++ b/tensorflow/lite/kernels/sub.cc
@@ -136,15 +136,18 @@ TfLiteStatus PrepareGeneralSubOp(TfLiteContext* context,
   // 31, therefore the addition will still fit in a 32 bit accumulator.
   op_params->left_shift = output->type == kTfLiteInt16 ? 15 : 20;
   const double twice_max_input_scale =
-      2 * std::max(input1_quantization_params.scale,
-                   input2_quantization_params.scale);
+      2 * static_cast<double>(std::max(input1_quantization_params.scale,
+                                       input2_quantization_params.scale));
   const double real_input1_multiplier =
-      input1_quantization_params.scale / twice_max_input_scale;
+      static_cast<double>(input1_quantization_params.scale) /
+      twice_max_input_scale;
   const double real_input2_multiplier =
-      input2_quantization_params.scale / twice_max_input_scale;
+      static_cast<double>(input2_quantization_params.scale) /
+      twice_max_input_scale;
   const double real_output_multiplier =
       twice_max_input_scale /
-      ((1 << op_params->left_shift) * output_quantization_params.scale);
+      ((1 << op_params->left_shift) *
+       static_cast<double>(output_quantization_params.scale));
 
   tflite::QuantizeMultiplierSmallerThanOneExp(real_input1_multiplier,
                                               &op_params->input1_multiplier,
diff --git a/tensorflow/lite/kernels/svdf.cc b/tensorflow/lite/kernels/svdf.cc
index 74a38153631..5173322f260 100644
--- a/tensorflow/lite/kernels/svdf.cc
+++ b/tensorflow/lite/kernels/svdf.cc
@@ -275,12 +275,14 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     TF_LITE_ENSURE(context, output->quantization.type != kTfLiteNoQuantization);
     auto* output_params = reinterpret_cast<TfLiteAffineQuantization*>(
         output->quantization.params);
-    const double effective_scale_1 = input_params->scale->data[0] *
-                                     weights_feature_params->scale->data[0] /
-                                     state_params->scale->data[0];
-    const double effective_scale_2 = state_params->scale->data[0] *
-                                     weight_time_params->scale->data[0] /
-                                     output_params->scale->data[0];
+    const double effective_scale_1 =
+        static_cast<double>(input_params->scale->data[0]) *
+        static_cast<double>(weights_feature_params->scale->data[0]) /
+        static_cast<double>(state_params->scale->data[0]);
+    const double effective_scale_2 =
+        static_cast<double>(state_params->scale->data[0]) *
+        static_cast<double>(weight_time_params->scale->data[0]) /
+        static_cast<double>(output_params->scale->data[0]);
     QuantizeMultiplier(effective_scale_1, &op_data->effective_scale_1_a,
                        &op_data->effective_scale_1_b);
     QuantizeMultiplier(effective_scale_2, &op_data->effective_scale_2_a,

From 7c5599d76019a675eb739ee72da6d1c87bbf53ba Mon Sep 17 00:00:00 2001
From: Amir Saadat <amirsaadat@google.com>
Date: Tue, 13 Oct 2026 17:44:52 -0700
Subject: [PATCH 0729/1000] Add sample coordinates and decoupled diff score
 support to HLO mismatch HTML dumps.

In interactive HLO mismatch HTML graph dumps (DumpHloModuleMismatchWithGraphData):
- Annotate Actual, Expected, and Rel Error tooltip lines with sample coordinates (e.g. `(at max diff, index {x, y})` or `(mean)`).
- Allow callers to provide an explicit `diff_score` override in MismatchDetails to decouple node heatmap color thresholds from the sampled element relative error.

PiperOrigin-RevId: 999005678
---
 third_party/xla/xla/tools/hlo_dump/BUILD      |   1 -
 .../xla/xla/tools/hlo_dump/hlo_dump_utils.cc  | 134 ++++-
 .../xla/xla/tools/hlo_dump/hlo_dump_utils.h   |   7 +-
 .../xla/tools/hlo_dump/hlo_dump_utils_test.cc | 496 +++++++++++++++++-
 4 files changed, 615 insertions(+), 23 deletions(-)

diff --git a/third_party/xla/xla/tools/hlo_dump/BUILD b/third_party/xla/xla/tools/hlo_dump/BUILD
index c3c551b3795..1f598e1d2b5 100644
--- a/third_party/xla/xla/tools/hlo_dump/BUILD
+++ b/third_party/xla/xla/tools/hlo_dump/BUILD
@@ -92,7 +92,6 @@ xla_cc_test(
         "//xla/hlo/parser:hlo_parser",
         "//xla/service:hlo_proto_cc",
         "//xla/tsl/platform:env",
-        "//xla/tsl/platform:statusor",
         "//xla/tsl/platform:test",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/strings",
diff --git a/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils.cc b/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils.cc
index 96ef805b27f..95caad0262c 100644
--- a/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils.cc
+++ b/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils.cc
@@ -1341,7 +1341,39 @@ PopulateTensorVisualizations(const HloModule& module,
           }
           ApplyBoundingBoxTail(bbox, mismatch->rel_error, info);
         } else {
-          info.mismatch_count = 1;
+          int64_t mismatch_count = 1;
+          if (!mismatch->is_summary_mean &&
+              mismatch->percentage_of_elems_exceeding_both_errors.has_value() &&
+              info.total_elements > 0) {
+            mismatch_count = std::max<int64_t>(
+                1, std::llround(
+                       (*mismatch->percentage_of_elems_exceeding_both_errors /
+                        100.0) *
+                       static_cast<double>(info.total_elements)));
+          }
+          info.mismatch_count = mismatch_count;
+          if (!mismatch->is_summary_mean &&
+              mismatch->sample_index.has_value() &&
+              mismatch->sample_index->size() == info.shape.size()) {
+            if (!info.shape.empty()) {
+              const double top_rel_error = (std::isnan(mismatch->rel_error) ||
+                                            std::isinf(mismatch->rel_error))
+                                               ? 1.0
+                                               : mismatch->rel_error;
+              info.top_mismatches.push_back(
+                  {*mismatch->sample_index, top_rel_error});
+            }
+            if (mismatch_count == 1) {
+              info.box_min = *mismatch->sample_index;
+              info.box_max = *mismatch->sample_index;
+            } else {
+              info.box_min.assign(info.shape.size(), 0);
+              info.box_max.resize(info.shape.size());
+              for (size_t d = 0; d < info.shape.size(); ++d) {
+                info.box_max[d] = std::max<int64_t>(0, info.shape[d] - 1);
+              }
+            }
+          }
         }
       }
 
@@ -1486,6 +1518,61 @@ std::string ConvertHloToHtml(
     }
   }
 
+  hlo_dump_ui_js = absl::StrReplaceAll(
+      hlo_dump_ui_js,
+      {{"vec4 getColorForScore(float score){\\n"
+        "if (score==-1.0) return vec4(0.702,0.533,1.0,1.0);\\n"
+        "if (score<0.0) return vec4(0.82,0.82,0.82,1.0);\\n"
+        "if (score>=100.0) return vec4(1.0,0.090,0.090,1.0);\\n"
+        "if (score<1.0){\\n"
+        "return mix(vec4(0.6,1.0,0.6,1.0),vec4(0.753,0.961,0.502,1.0),score);"
+        "\\n"
+        "}else if (score<5.0){\\n"
+        "return mix(vec4(0.753,0.961,0.502,1.0),vec4(0.878,0.933,0.251,1.0),"
+        "(score-1.0)/4.0);\\n"
+        "}else if (score<10.0){\\n"
+        "return mix(vec4(0.878,0.933,0.251,1.0),vec4(0.933,0.933,0.0,1.0),"
+        "(score-5.0)/5.0);\\n"
+        "}else if (score<30.0){\\n"
+        "return mix(vec4(0.933,0.933,0.0,1.0),vec4(1.0,0.753,0.0,1.0),"
+        "(score-10.0)/20.0);\\n"
+        "}else if (score<60.0){\\n"
+        "return mix(vec4(1.0,0.753,0.0,1.0),vec4(1.0,0.502,0.0,1.0),"
+        "(score-30.0)/30.0);\\n"
+        "}else{\\n"
+        "return mix(vec4(1.0,0.502,0.0,1.0),vec4(1.0,0.090,0.090,1.0),"
+        "(score-60.0)/40.0);\\n"
+        "}\\n"
+        "}",
+        "vec4 getColorForScore(float score){\\n"
+        "if (score==-1.0) return vec4(0.702,0.533,1.0,1.0);\\n"
+        "if (score<0.0) return vec4(0.82,0.82,0.82,1.0);\\n"
+        "if (score==0.0) return vec4(0.6,1.0,0.6,1.0);\\n"
+        "return vec4(1.0,0.090,0.090,1.0);\\n"
+        "}"},
+       {"float diffFactor=aDiffScore==-1.0 ? 5.0 : "
+        "clamp(aDiffScore/20.0,0.0,5.0);",
+        "float diffFactor=5.0;"},
+       {"if (aDiffScore==-1.0){\\n"
+        "depth=-1.0;\\n"
+        "}else{\\n"
+        "depth=-clamp((aDiffScore+1.0)/101.0,0.0,1.0);\\n"
+        "}",
+        "if (aDiffScore==-1.0){\\n"
+        "depth=-1.0;\\n"
+        "}else if (aDiffScore>0.0){\\n"
+        "depth=-0.5-clamp(aDiffScore/200.0,0.0,0.49);\\n"
+        "}else if (aDiffScore==0.0){\\n"
+        "depth=-0.01;\\n"
+        "}else{\\n"
+        "depth=0.0;\\n"
+        "}"},
+       {"var b=`Diff Score: ${a.diffScore}`;a.diffScore===-1&&(b=\"Diff "
+        "Score: NaN/Inf Mismatch\");",
+        "var b=a.diffScore===-1?\"Diff Score: NaN/Inf Mismatch\":a.diffScore<0?"
+        "\"Diff Score: Not Evaluated\":`Diff Score: "
+        "${Number(a.diffScore.toPrecision(6))}`;"}});
+
   std::string background_styles = GenerateBackgroundStyles(annotations);
 
   std::string processed_hlo_text(hlo_text);
@@ -1632,18 +1719,38 @@ absl::flat_hash_map<TensorKey, TensorAnnotation> PopulateMismatchAnnotations(
 
     TensorAnnotation ann;
     ann.anchor_id = absl::StrCat("step", target_instr->unique_id());
-    ann.background_color = "pink";
+    const bool is_finite_vals =
+        std::isfinite(mismatch.actual) && std::isfinite(mismatch.expected);
+    const bool is_zero_both_errors =
+        is_finite_vals &&
+        ((mismatch.percentage_of_elems_exceeding_both_errors.has_value() &&
+          *mismatch.percentage_of_elems_exceeding_both_errors == 0.0) ||
+         (mismatch.diff_score.has_value() && *mismatch.diff_score == 0.0));
+    ann.background_color = is_zero_both_errors ? "lightgreen" : "pink";
 
     std::vector<std::string> tooltip_parts;
     tooltip_parts.push_back("<b>Numeric Mismatch:</b>");
     if (mismatch.custom_description.has_value()) {
       tooltip_parts.push_back(*mismatch.custom_description);
     } else {
-      tooltip_parts.push_back(absl::StrFormat("Actual: %g", mismatch.actual));
+      std::string sample_suffix;
+      if (mismatch.is_summary_mean) {
+        sample_suffix = " (mean)";
+      } else if (mismatch.sample_index.has_value()) {
+        if (!mismatch.sample_index->empty()) {
+          sample_suffix =
+              absl::StrFormat(" (at max diff, index {%s})",
+                              absl::StrJoin(*mismatch.sample_index, ","));
+        } else {
+          sample_suffix = " (at max diff)";
+        }
+      }
       tooltip_parts.push_back(
-          absl::StrFormat("Expected: %g", mismatch.expected));
+          absl::StrFormat("Actual: %g%s", mismatch.actual, sample_suffix));
       tooltip_parts.push_back(
-          absl::StrFormat("Rel Error: %g", mismatch.rel_error));
+          absl::StrFormat("Expected: %g%s", mismatch.expected, sample_suffix));
+      tooltip_parts.push_back(absl::StrFormat(
+          "Rel Error: %g%s", mismatch.rel_error, sample_suffix));
       if (mismatch.percentage_of_elems_exceeding_abs_error.has_value()) {
         tooltip_parts.push_back(
             absl::StrFormat("Elems exceeding abs error: %.2f%%",
@@ -1675,7 +1782,8 @@ absl::flat_hash_map<TensorKey, TensorAnnotation> PopulateMismatchAnnotations(
 }
 
 GraphData PopulateMismatchGraphData(
-    const HloModule& module, absl::Span<const MismatchDetails> mismatches) {
+    const HloModule& module, absl::Span<const MismatchDetails> mismatches,
+    double unverified_diff_score) {
   GraphData graph_data;
 
   absl::flat_hash_map<const HloComputation*, const HloInstruction*>
@@ -1717,9 +1825,14 @@ GraphData PopulateMismatchGraphData(
 
   for (const MismatchDetails& m : mismatches) {
     double score = 100.0;
-    if (std::isnan(m.actual) || std::isinf(m.actual) ||
-        std::isnan(m.expected) || std::isinf(m.expected)) {
+    if (m.diff_score.has_value()) {
+      score = *m.diff_score;
+    } else if (std::isnan(m.actual) || std::isinf(m.actual) ||
+               std::isnan(m.expected) || std::isinf(m.expected)) {
       score = kNanInfMismatchDiffScore;
+    } else if (m.percentage_of_elems_exceeding_both_errors.has_value() &&
+               *m.percentage_of_elems_exceeding_both_errors == 0.0) {
+      score = 0.0;
     } else if (m.rel_error > 0.0) {
       score = m.rel_error * 100.0;
     }
@@ -1779,7 +1892,7 @@ GraphData PopulateMismatchGraphData(
     double y = next_y;
     next_y += 2.0;
 
-    double score = 0.0;
+    double score = unverified_diff_score;
     if (auto it = instr_to_score.find(instr); it != instr_to_score.end()) {
       score = it->second;
     }
@@ -1830,7 +1943,8 @@ absl::StatusOr<std::string> DumpHloModuleMismatchWithGraphData(
   }
   absl::flat_hash_map<TensorKey, TensorAnnotation> annotations =
       PopulateMismatchAnnotations(module, mismatches);
-  GraphData graph_data = PopulateMismatchGraphData(module, mismatches);
+  GraphData graph_data =
+      PopulateMismatchGraphData(module, mismatches, kUnverifiedDiffScore);
   auto tensor_visualizations = PopulateTensorVisualizations(module, mismatches);
 
   xla::StackFrameIndexProto stack_frame_index = module.stack_frames().proto();
diff --git a/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils.h b/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils.h
index db1f0565df1..f1d2221ac81 100644
--- a/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils.h
+++ b/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils.h
@@ -34,6 +34,7 @@ limitations under the License.
 namespace xla::numerics::debug_info {
 
 inline constexpr double kNanInfMismatchDiffScore = -1.0;
+inline constexpr double kUnverifiedDiffScore = -2.0;
 
 struct SliceBoundingBox {
   int64_t slice_index = 0;
@@ -81,6 +82,8 @@ struct MismatchDetails {
   double actual = 0.0;
   double expected = 0.0;
   double rel_error = 0.0;
+  bool is_summary_mean = false;
+  std::optional<std::vector<int64_t>> sample_index;
   std::optional<double> percentage_of_elems_exceeding_abs_error;
   std::optional<double> percentage_of_elems_exceeding_rel_error;
   std::optional<double> percentage_of_elems_exceeding_both_errors;
@@ -93,6 +96,7 @@ struct MismatchDetails {
   // falling back to a fabricated single-element mismatch at the origin. HLO
   // highlighting, the graph view, and tooltips are unaffected.
   bool has_element_level_data = true;
+  std::optional<double> diff_score;
 };
 
 struct TensorKey {
@@ -185,7 +189,8 @@ absl::flat_hash_map<TensorKey, TensorAnnotation> PopulateMismatchAnnotations(
     const HloModule& module, absl::Span<const MismatchDetails> mismatches);
 
 GraphData PopulateMismatchGraphData(
-    const HloModule& module, absl::Span<const MismatchDetails> mismatches);
+    const HloModule& module, absl::Span<const MismatchDetails> mismatches,
+    double unverified_diff_score = 0.0);
 
 // Generates an interactive HTML dump for an HLO module and writes it to
 // `output_dir`/`output_filename`, creating `output_dir` if needed.
diff --git a/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils_test.cc b/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils_test.cc
index f3648af4e6c..ca21b41e251 100644
--- a/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils_test.cc
+++ b/third_party/xla/xla/tools/hlo_dump/hlo_dump_utils_test.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <array>
 #include <cstdint>
+#include <cstdlib>
 #include <fstream>
 #include <iterator>
 #include <limits>
@@ -38,7 +39,6 @@ limitations under the License.
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/tsl/platform/env.h"
-#include "xla/tsl/platform/statusor.h"
 #include "xla/tsl/platform/test.h"
 #include "xla/xla_data.pb.h"
 #include "tsl/platform/path.h"
@@ -436,8 +436,8 @@ ENTRY main {
   ROOT add = f32[10] add(p0, p0)
 }
 )hlo";
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          ParseAndReturnUnverifiedModule(hlo_string));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnUnverifiedModule(hlo_string));
 
   MismatchDetails details;
   details.target_instruction_name = "add";
@@ -473,6 +473,88 @@ ENTRY main {
                                 "Result of reduce: False"));
 }
 
+TEST(HloDumpUtilsTest, PopulateMismatchAnnotations_SampleIndexAndSummaryMean) {
+  const absl::string_view hlo_string = R"hlo(
+HloModule test_module
+ENTRY main {
+  p0 = f32[2,4] parameter(0)
+  add = f32[2,4] add(p0, p0)
+  ROOT sub = f32[2,4] subtract(add, p0)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnUnverifiedModule(hlo_string));
+
+  MismatchDetails per_elem_details;
+  per_elem_details.target_instruction_name = "add";
+  per_elem_details.actual = 14.0;
+  per_elem_details.expected = 10.0;
+  per_elem_details.rel_error = 0.4;
+  per_elem_details.sample_index = std::vector<int64_t>{1, 3};
+
+  MismatchDetails mean_details;
+  mean_details.target_instruction_name = "sub";
+  mean_details.actual = 7.5;
+  mean_details.expected = -2.5;
+  mean_details.rel_error = 4.0;
+  mean_details.is_summary_mean = true;
+
+  MismatchDetails scalar_details;
+  scalar_details.target_instruction_name = "p0";
+  scalar_details.actual = 3.0;
+  scalar_details.expected = 2.0;
+  scalar_details.rel_error = 0.5;
+  scalar_details.sample_index = std::vector<int64_t>{};
+
+  MismatchDetails unknown_details;
+  unknown_details.target_instruction_name = "nonexistent_op";
+  unknown_details.actual = 1.0;
+  unknown_details.expected = 0.0;
+  unknown_details.rel_error = 1.0;
+
+  auto annotations = PopulateMismatchAnnotations(
+      *module,
+      {per_elem_details, mean_details, scalar_details, unknown_details});
+  auto add_key = TensorKey::Create("add", ShapeIndex{});
+  ASSERT_TRUE(annotations.contains(add_key));
+  ASSERT_TRUE(annotations[add_key].tooltip_data.has_value());
+  EXPECT_TRUE(absl::StrContains(*annotations[add_key].tooltip_data,
+                                "Actual: 14 (at max diff, index {1,3})"));
+  EXPECT_TRUE(absl::StrContains(*annotations[add_key].tooltip_data,
+                                "Expected: 10 (at max diff, index {1,3})"));
+  EXPECT_TRUE(absl::StrContains(*annotations[add_key].tooltip_data,
+                                "Rel Error: 0.4 (at max diff, index {1,3})"));
+
+  auto sub_key = TensorKey::Create("sub", ShapeIndex{});
+  ASSERT_TRUE(annotations.contains(sub_key));
+  ASSERT_TRUE(annotations[sub_key].tooltip_data.has_value());
+  EXPECT_TRUE(absl::StrContains(*annotations[sub_key].tooltip_data,
+                                "Actual: 7.5 (mean)"));
+  EXPECT_TRUE(absl::StrContains(*annotations[sub_key].tooltip_data,
+                                "Expected: -2.5 (mean)"));
+  EXPECT_TRUE(absl::StrContains(*annotations[sub_key].tooltip_data,
+                                "Rel Error: 4 (mean)"));
+
+  auto p0_key = TensorKey::Create("p0", ShapeIndex{});
+  ASSERT_TRUE(annotations.contains(p0_key));
+  ASSERT_TRUE(annotations[p0_key].tooltip_data.has_value());
+  EXPECT_TRUE(absl::StrContains(*annotations[p0_key].tooltip_data,
+                                "Actual: 3 (at max diff)"));
+  EXPECT_FALSE(
+      absl::StrContains(*annotations[p0_key].tooltip_data, "index {}"));
+  EXPECT_FALSE(
+      annotations.contains(TensorKey::Create("nonexistent_op", ShapeIndex{})));
+
+  MismatchDetails custom_desc_details = mean_details;
+  custom_desc_details.custom_description = "Summary mean comparison";
+  auto custom_annotations =
+      PopulateMismatchAnnotations(*module, {custom_desc_details});
+  ASSERT_TRUE(custom_annotations.contains(sub_key));
+  ASSERT_TRUE(custom_annotations[sub_key].tooltip_data.has_value());
+  EXPECT_TRUE(absl::StrContains(*custom_annotations[sub_key].tooltip_data,
+                                "Summary mean comparison"));
+}
+
 TEST(HloDumpUtilsTest, PopulateMismatchGraphData_Basic) {
   const absl::string_view hlo_string = R"hlo(
 HloModule test_module
@@ -481,8 +563,8 @@ ENTRY main {
   ROOT add = f32[10] add(p0, p0)
 }
 )hlo";
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          ParseAndReturnUnverifiedModule(hlo_string));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnUnverifiedModule(hlo_string));
 
   MismatchDetails details;
   details.target_instruction_name = "add";
@@ -529,8 +611,8 @@ ENTRY main {
   ROOT tuple = (f32[10], f32[]) tuple(add, reduce)
 }
 )hlo";
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          ParseAndReturnUnverifiedModule(hlo_string));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnUnverifiedModule(hlo_string));
 
   MismatchDetails details;
   details.target_instruction_name = "tuple";
@@ -565,8 +647,8 @@ ENTRY main {
   ROOT fusion = f32[10] fusion(p0), kind=kLoop, calls=fused_comp
 }
 )hlo";
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          ParseAndReturnUnverifiedModule(hlo_string));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnUnverifiedModule(hlo_string));
 
   MismatchDetails details;
   details.target_instruction_name = "fusion";
@@ -727,8 +809,8 @@ ENTRY main {
   ROOT add = f32[10] add(p0, p0)
 }
 )hlo";
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                          ParseAndReturnUnverifiedModule(hlo_string));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnUnverifiedModule(hlo_string));
 
   MismatchDetails details;
   details.target_instruction_name = "add";
@@ -749,6 +831,119 @@ ENTRY main {
   EXPECT_TRUE(found_add);
 }
 
+TEST(HloDumpUtilsTest, PopulateMismatchGraphData_ExplicitDiffScore) {
+  const absl::string_view hlo_string = R"hlo(
+HloModule test_module
+ENTRY main {
+  p0 = f32[10] parameter(0)
+  ROOT add = f32[10] add(p0, p0)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnUnverifiedModule(hlo_string));
+
+  MismatchDetails details;
+  details.target_instruction_name = "add";
+  details.actual = 1.000243;
+  details.expected = 1.0;
+  // Sampled element rel_error is small (0.0245%), but overall tensor diff_score
+  // is 100.0 (red).
+  details.rel_error = 0.000245;
+  details.diff_score = 100.0;
+
+  auto graph_data = PopulateMismatchGraphData(*module, {details});
+  bool found_add = false;
+  for (const auto& node : graph_data.nodes) {
+    if (node.key == "add") {
+      found_add = true;
+      EXPECT_EQ(node.diff_score, 100.0);
+    }
+  }
+  EXPECT_TRUE(found_add);
+}
+
+TEST(HloDumpUtilsTest, PopulateMismatchGraphData_ExplicitDiffScoreNanInf) {
+  const absl::string_view hlo_string = R"hlo(
+HloModule test_module
+ENTRY main {
+  p0 = f32[10] parameter(0)
+  add = f32[10] add(p0, p0)
+  sub = f32[10] subtract(add, p0)
+  ROOT mul = f32[10] multiply(sub, p0)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnUnverifiedModule(hlo_string));
+
+  MismatchDetails p0_low;
+  p0_low.target_instruction_name = "p0";
+  p0_low.actual = 1.0;
+  p0_low.expected = 2.0;
+  p0_low.rel_error = 0.25;
+  p0_low.diff_score = 25.0;
+
+  MismatchDetails p0_high = p0_low;
+  p0_high.diff_score = 75.0;
+
+  MismatchDetails add_finite;
+  add_finite.target_instruction_name = "add";
+  add_finite.actual = 1.0;
+  add_finite.expected = 2.0;
+  add_finite.rel_error = 0.5;
+  add_finite.diff_score = 50.0;
+
+  MismatchDetails details = add_finite;
+  details.diff_score = kNanInfMismatchDiffScore;
+
+  MismatchDetails add_after_nan = add_finite;
+  add_after_nan.diff_score = 80.0;
+
+  MismatchDetails expected_nan;
+  expected_nan.target_instruction_name = "sub";
+  expected_nan.actual = 1.0;
+  expected_nan.expected = std::numeric_limits<double>::quiet_NaN();
+  expected_nan.rel_error = 0.5;
+
+  MismatchDetails expected_inf;
+  expected_inf.target_instruction_name = "mul";
+  expected_inf.actual = 1.0;
+  expected_inf.expected = std::numeric_limits<double>::infinity();
+  expected_inf.rel_error = 0.5;
+
+  MismatchDetails unknown_target;
+  unknown_target.target_instruction_name = "nonexistent_op";
+  unknown_target.actual = 1.0;
+  unknown_target.expected = 2.0;
+  unknown_target.rel_error = 0.5;
+
+  auto graph_data = PopulateMismatchGraphData(
+      *module, {p0_low, p0_high, add_finite, details, add_after_nan,
+                expected_nan, expected_inf, unknown_target});
+  bool found_add = false;
+  bool found_p0 = false;
+  bool found_sub = false;
+  bool found_mul = false;
+  for (const auto& node : graph_data.nodes) {
+    if (node.key == "add") {
+      found_add = true;
+      EXPECT_EQ(node.diff_score, kNanInfMismatchDiffScore);
+    } else if (node.key == "p0") {
+      found_p0 = true;
+      EXPECT_EQ(node.diff_score, 75.0);
+    } else if (node.key == "sub") {
+      found_sub = true;
+      EXPECT_EQ(node.diff_score, kNanInfMismatchDiffScore);
+    } else if (node.key == "mul") {
+      found_mul = true;
+      EXPECT_EQ(node.diff_score, kNanInfMismatchDiffScore);
+    }
+  }
+  EXPECT_TRUE(found_add);
+  EXPECT_TRUE(found_p0);
+  EXPECT_TRUE(found_sub);
+  EXPECT_TRUE(found_mul);
+}
+
 TEST(HloDumpUtilsTest, ConvertHloToHtmlCompactGte) {
   std::string hlo_text =
       "ENTRY %test {\n"
@@ -1823,5 +2018,284 @@ ENTRY main {
   EXPECT_TRUE(tsl::Env::Default()->FileExists(second).ok());
 }
 
+TEST(HloDumpUtilsTest,
+     PopulateTensorVisualizations_TupleSubshapeAndFallbackTarget) {
+  const absl::string_view hlo_string = R"hlo(
+HloModule tuple_vis_module
+ENTRY main {
+  p0 = f32[4,8] parameter(0)
+  p1 = f32[16] parameter(1)
+  tok = token[] after-all()
+  empty_tuple = () tuple()
+  inner_tuple = (f32[4,8]) tuple(p0)
+  nested_tuple = ((f32[4,8]), f32[16]) tuple(inner_tuple, p1)
+  tuple_oob = (f32[4,8], f32[16]) tuple(p0, p1)
+  ROOT tuple = (f32[4,8], f32[16]) tuple(p0, p1)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnUnverifiedModule(hlo_string));
+
+  MismatchDetails tuple_mismatch;
+  tuple_mismatch.target_instruction_name = "tuple";
+  tuple_mismatch.output_shape_index = 1;
+  tuple_mismatch.actual = 2.0;
+  tuple_mismatch.expected = 1.0;
+  tuple_mismatch.rel_error = 1.0;
+
+  MismatchDetails inner_tuple_mismatch;
+  inner_tuple_mismatch.target_instruction_name = "inner_tuple";
+  inner_tuple_mismatch.actual = 2.0;
+  inner_tuple_mismatch.expected = 1.0;
+  inner_tuple_mismatch.rel_error = 1.0;
+
+  MismatchDetails tuple_oob_mismatch;
+  tuple_oob_mismatch.target_instruction_name = "tuple_oob";
+  tuple_oob_mismatch.output_shape_index = 99;
+  tuple_oob_mismatch.actual = 2.0;
+  tuple_oob_mismatch.expected = 1.0;
+  tuple_oob_mismatch.rel_error = 1.0;
+
+  MismatchDetails p0_mismatched_box;
+  p0_mismatched_box.target_instruction_name = "p0";
+  p0_mismatched_box.actual = 2.0;
+  p0_mismatched_box.expected = 1.0;
+  p0_mismatched_box.rel_error = 1.0;
+  MismatchBoundingBox mismatched_rank_box;
+  mismatched_rank_box.box_min = {1};
+  mismatched_rank_box.box_max = {3};
+  mismatched_rank_box.mismatch_count = 3;
+  mismatched_rank_box.total_elements = 32;
+  mismatched_rank_box.pattern = "SINGLE_SLICE";
+  p0_mismatched_box.bounding_box = mismatched_rank_box;
+
+  MismatchDetails fallback_with_box;
+  fallback_with_box.target_instruction_name = "external_op_with_box";
+  fallback_with_box.actual = 3.0;
+  fallback_with_box.expected = 1.0;
+  fallback_with_box.rel_error = 2.0;
+  MismatchBoundingBox box;
+  box.box_min = {1, 2};
+  box.box_max = {4, 6};
+  box.mismatch_count = 12;
+  box.total_elements = 128;
+  box.pattern = "DENSE_BLOCK";
+  fallback_with_box.bounding_box = box;
+
+  MismatchDetails fallback_without_box;
+  fallback_without_box.target_instruction_name = "external_op_no_box";
+  fallback_without_box.actual = 4.0;
+  fallback_without_box.expected = 2.0;
+  fallback_without_box.rel_error = 1.0;
+
+  auto visualizations = PopulateTensorVisualizations(
+      *module, {tuple_mismatch, inner_tuple_mismatch, tuple_oob_mismatch,
+                p0_mismatched_box, fallback_with_box, fallback_without_box,
+                fallback_with_box});
+
+  ASSERT_TRUE(visualizations.contains("tuple"));
+  const auto& tuple_vis = visualizations.at("tuple");
+  EXPECT_TRUE(tuple_vis.has_mismatch);
+  EXPECT_EQ(tuple_vis.shape, (std::vector<int64_t>{16}));
+  EXPECT_EQ(tuple_vis.total_elements, 16);
+  EXPECT_EQ(tuple_vis.mismatch_count, 1);
+
+  ASSERT_TRUE(visualizations.contains("inner_tuple"));
+  EXPECT_EQ(visualizations.at("inner_tuple").shape,
+            (std::vector<int64_t>{4, 8}));
+
+  ASSERT_TRUE(visualizations.contains("tuple_oob"));
+  EXPECT_EQ(visualizations.at("tuple_oob").shape, (std::vector<int64_t>{4, 8}));
+
+  ASSERT_TRUE(visualizations.contains("nested_tuple"));
+  EXPECT_TRUE(visualizations.at("nested_tuple").shape.empty());
+
+  ASSERT_TRUE(visualizations.contains("p0"));
+  const auto& p0_vis = visualizations.at("p0");
+  EXPECT_EQ(p0_vis.shape, (std::vector<int64_t>{4, 8}));
+  EXPECT_EQ(p0_vis.box_min, (std::vector<int64_t>{0, 0}));
+  EXPECT_EQ(p0_vis.box_max, (std::vector<int64_t>{0, 0}));
+
+  ASSERT_TRUE(visualizations.contains("external_op_with_box"));
+  const auto& ext_box_vis = visualizations.at("external_op_with_box");
+  EXPECT_TRUE(ext_box_vis.has_mismatch);
+  EXPECT_EQ(ext_box_vis.box_min, (std::vector<int64_t>{1, 2}));
+  EXPECT_EQ(ext_box_vis.box_max, (std::vector<int64_t>{4, 6}));
+  EXPECT_EQ(ext_box_vis.mismatch_count, 12);
+  EXPECT_EQ(ext_box_vis.total_elements, 128);
+  EXPECT_EQ(ext_box_vis.pattern, "DENSE_BLOCK");
+
+  ASSERT_TRUE(visualizations.contains("external_op_no_box"));
+  const auto& ext_no_box_vis = visualizations.at("external_op_no_box");
+  EXPECT_TRUE(ext_no_box_vis.has_mismatch);
+  EXPECT_EQ(ext_no_box_vis.mismatch_count, 1);
+
+  // Verify ComputeBoundingBoxFromLiteralMask edge cases: non-array tuple,
+  // zero-element array, and size-1 leading dimension.
+  Literal tuple_lit = LiteralUtil::MakeTuple({});
+  MismatchBoundingBox tuple_bbox = ComputeBoundingBoxFromLiteralMask(tuple_lit);
+  EXPECT_EQ(tuple_bbox.mismatch_count, 0);
+  EXPECT_TRUE(tuple_bbox.tensor_shape.empty());
+
+  Literal empty_lit = LiteralUtil::CreateR1<bool>({});
+  MismatchBoundingBox empty_bbox = ComputeBoundingBoxFromLiteralMask(empty_lit);
+  EXPECT_EQ(empty_bbox.total_elements, 0);
+  EXPECT_EQ(empty_bbox.mismatch_count, 0);
+
+  Literal dim1_lit = LiteralUtil::CreateR2<bool>({{true, true, false, false}});
+  MismatchBoundingBox dim1_bbox = ComputeBoundingBoxFromLiteralMask(dim1_lit);
+  EXPECT_EQ(dim1_bbox.mismatch_count, 2);
+  EXPECT_EQ(dim1_bbox.pattern, "DENSE_BLOCK");
+
+  // Verify DumpHloModuleMismatchWithGraphData write-error and temp-dir paths.
+  auto bad_path_or = DumpHloModuleMismatchWithGraphData(
+      *module, {tuple_mismatch}, "nonexistent_subdir_xyz/mismatch.html");
+  EXPECT_FALSE(bad_path_or.ok());
+
+  const char* orig_env = std::getenv("TEST_UNDECLARED_OUTPUTS_DIR");
+  std::optional<std::string> saved_env;
+  if (orig_env != nullptr) {
+    saved_env = std::string(orig_env);
+  }
+  tsl::setenv("TEST_UNDECLARED_OUTPUTS_DIR", "", 1);
+  auto empty_env_or = DumpHloModuleMismatchWithGraphData(
+      *module, {tuple_mismatch},
+      tsl::io::JoinPath(::testing::TempDir(), "temp_mismatch_empty_env.html"));
+  EXPECT_TRUE(empty_env_or.ok());
+
+  tsl::unsetenv("TEST_UNDECLARED_OUTPUTS_DIR");
+  auto unset_env_or = DumpHloModuleMismatchWithGraphData(
+      *module, {tuple_mismatch}, "temp_mismatch_unset_env.html");
+  EXPECT_TRUE(unset_env_or.ok());
+
+  if (saved_env.has_value()) {
+    tsl::setenv("TEST_UNDECLARED_OUTPUTS_DIR", saved_env->c_str(), 1);
+  }
+}
+
+TEST(HloDumpUtilsTest,
+     PopulateMismatchGraphDataAndVisualizations_UnverifiedGrayAndActualScore) {
+  const absl::string_view hlo_string = R"hlo(
+HloModule hv_module
+ENTRY main {
+  p0 = f32[4,8] parameter(0)
+  p1 = f32[4,8] parameter(1)
+  sub = f32[4,8] subtract(p0, p1)
+  ROOT add = f32[4,8] add(sub, p1)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnUnverifiedModule(hlo_string));
+
+  MismatchDetails sub_mismatch;
+  sub_mismatch.target_instruction_name = "sub";
+  sub_mismatch.actual = -6.99993;
+  sub_mismatch.expected = -6.99822;
+  sub_mismatch.rel_error = 0.000244621;
+  sub_mismatch.sample_index = {1, 5};
+  sub_mismatch.percentage_of_elems_exceeding_both_errors = 0.0;
+
+  MismatchDetails add_mismatch;
+  add_mismatch.target_instruction_name = "add";
+  add_mismatch.actual = 10.0;
+  add_mismatch.expected = 8.0;
+  add_mismatch.rel_error = 0.25;
+  add_mismatch.sample_index = {2, 3};
+  add_mismatch.percentage_of_elems_exceeding_both_errors = 25.0;
+
+  auto annotations =
+      PopulateMismatchAnnotations(*module, {sub_mismatch, add_mismatch});
+  EXPECT_EQ(annotations[TensorKey::Create("sub", {})].background_color,
+            "lightgreen");
+  EXPECT_EQ(annotations[TensorKey::Create("add", {})].background_color, "pink");
+
+  GraphData graph_data = PopulateMismatchGraphData(
+      *module, {sub_mismatch, add_mismatch}, kUnverifiedDiffScore);
+  absl::flat_hash_map<std::string, double> scores;
+  for (const auto& node : graph_data.nodes) {
+    scores[node.key] = node.diff_score;
+  }
+  EXPECT_DOUBLE_EQ(scores["p0"], kUnverifiedDiffScore);
+  EXPECT_DOUBLE_EQ(scores["p1"], kUnverifiedDiffScore);
+  EXPECT_DOUBLE_EQ(scores["sub"], 0.0);
+  EXPECT_DOUBLE_EQ(scores["add"], 25.0);
+
+  auto visualizations =
+      PopulateTensorVisualizations(*module, {sub_mismatch, add_mismatch});
+  ASSERT_TRUE(visualizations.contains("sub"));
+  const auto& sub_vis = visualizations.at("sub");
+  EXPECT_TRUE(sub_vis.has_mismatch);
+  EXPECT_EQ(sub_vis.mismatch_count, 1);
+  EXPECT_EQ(sub_vis.box_min, (std::vector<int64_t>{1, 5}));
+  EXPECT_EQ(sub_vis.box_max, (std::vector<int64_t>{1, 5}));
+  ASSERT_EQ(sub_vis.top_mismatches.size(), 1);
+  EXPECT_EQ(sub_vis.top_mismatches[0].coord, (std::vector<int64_t>{1, 5}));
+  EXPECT_NEAR(sub_vis.top_mismatches[0].rel_error, 0.000244621, 1e-9);
+
+  ASSERT_TRUE(visualizations.contains("add"));
+  const auto& add_vis = visualizations.at("add");
+  EXPECT_TRUE(add_vis.has_mismatch);
+  EXPECT_EQ(add_vis.mismatch_count, 8);
+  EXPECT_EQ(add_vis.box_min, (std::vector<int64_t>{0, 0}));
+  EXPECT_EQ(add_vis.box_max, (std::vector<int64_t>{3, 7}));
+  ASSERT_EQ(add_vis.top_mismatches.size(), 1);
+  EXPECT_EQ(add_vis.top_mismatches[0].coord, (std::vector<int64_t>{2, 3}));
+  EXPECT_DOUBLE_EQ(add_vis.top_mismatches[0].rel_error, 0.25);
+
+  const std::string dir =
+      tsl::io::JoinPath(::testing::TempDir(), "hv_unverified_gray_test");
+  ASSERT_OK_AND_ASSIGN(
+      const std::string html_path,
+      DumpHloModuleMismatchWithGraphData(*module, {sub_mismatch, add_mismatch},
+                                         "hv_report.html", dir));
+  std::string html_content;
+  ASSERT_OK(
+      tsl::ReadFileToString(tsl::Env::Default(), html_path, &html_content));
+  EXPECT_TRUE(absl::StrContains(html_content, "Diff Score: Not Evaluated"));
+  EXPECT_TRUE(absl::StrContains(
+      html_content, "if (score==0.0) return vec4(0.6,1.0,0.6,1.0);"));
+  EXPECT_TRUE(absl::StrContains(html_content, "bg-lightgreen"));
+  EXPECT_TRUE(absl::StrContains(html_content, "bg-pink"));
+}
+
+TEST(HloDumpUtilsTest,
+     PopulateMismatchGraphData_FiniteZeroCrossingDoesNotTriggerNanInfScore) {
+  const absl::string_view hlo_string = R"hlo(
+HloModule zero_crossing_module
+ENTRY main {
+  p0 = f32[4] parameter(0)
+  sub = f32[4] negate(p0)
+  ROOT add = f32[4] add(sub, p0)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnUnverifiedModule(hlo_string));
+
+  MismatchDetails benign_zero_crossing;
+  benign_zero_crossing.target_instruction_name = "sub";
+  benign_zero_crossing.actual = 1e-7;
+  benign_zero_crossing.expected = 0.0;
+  benign_zero_crossing.rel_error = std::numeric_limits<double>::infinity();
+  benign_zero_crossing.percentage_of_elems_exceeding_both_errors = 0.0;
+
+  MismatchDetails true_zero_crossing;
+  true_zero_crossing.target_instruction_name = "add";
+  true_zero_crossing.actual = 1.0;
+  true_zero_crossing.expected = 0.0;
+  true_zero_crossing.rel_error = std::numeric_limits<double>::infinity();
+  true_zero_crossing.percentage_of_elems_exceeding_both_errors = 25.0;
+
+  GraphData graph_data = PopulateMismatchGraphData(
+      *module, {benign_zero_crossing, true_zero_crossing},
+      kUnverifiedDiffScore);
+  absl::flat_hash_map<std::string, double> scores;
+  for (const auto& node : graph_data.nodes) {
+    scores[node.key] = node.diff_score;
+  }
+  EXPECT_DOUBLE_EQ(scores["sub"], 0.0);
+  EXPECT_GT(scores["add"], 0.0);
+  EXPECT_NE(scores["add"], kNanInfMismatchDiffScore);
+}
+
 }  // namespace
 }  // namespace xla::numerics::debug_info

From 0e0775e140f803f094a2965a1de2de75c23f5c71 Mon Sep 17 00:00:00 2001
From: Hyeontaek Lim <hyeontaek@google.com>
Date: Tue, 13 Oct 2026 18:24:45 -0700
Subject: [PATCH 0730/1000] Qualify `mlir::dyn_cast` and `mlir::isa` in
 `ExportStablehloShardingsPass`

Consolidating the module walks in `ExportStablehloShardingsPass` introduced unqualified calls to `dyn_cast` and `isa` with explicit template arguments. Because `mlir::dyn_cast` and `mlir::isa` are not brought into the file namespace via `using` declarations, unqualified template calls rely on C++20 argument-dependent lookup and fail to compile under `-Werror,-Wc++20-extensions`:

```
xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc:332:29: error: use of function template name with no prior declaration in function call with explicit template arguments is a C++20 extension [-Werror,-Wc++20-extensions]
```

Explicitly qualifying these calls as `mlir::dyn_cast` and `mlir::isa` matches the rest of `export_shardings.cc` and resolves the C++20 extension error in C++17 builds.

PiperOrigin-RevId: 999024271
---
 .../spmd/shardy/stablehlo_round_trip/export_shardings.cc  | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc
index 3540a6ae657..4a4925565d5 100644
--- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc
+++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc
@@ -329,7 +329,7 @@ class ExportStablehloShardingsPass
     };
 
     moduleOp.walk([&](Operation* op) {
-      if (auto customCall = dyn_cast<stablehlo::CustomCallOp>(op)) {
+      if (auto customCall = mlir::dyn_cast<stablehlo::CustomCallOp>(op)) {
         // StableHLO doesn't have an equivalent of `erf` and `topk` ops.
         // If they have a sharding annotation, we need to move it into
         // `mhlo.attributes`, which StableHLO->MHLO conversion would lift back
@@ -351,9 +351,9 @@ class ExportStablehloShardingsPass
                                 builder.getDictionaryAttr(newAttributes));
           }
         }
-      } else if (isa<stablehlo::AllGatherOp, stablehlo::AllReduceOp,
-                     stablehlo::ReduceScatterOp, stablehlo::AllToAllOp,
-                     stablehlo::CollectiveBroadcastOp>(op)) {
+      } else if (mlir::isa<stablehlo::AllGatherOp, stablehlo::AllReduceOp,
+                           stablehlo::ReduceScatterOp, stablehlo::AllToAllOp,
+                           stablehlo::CollectiveBroadcastOp>(op)) {
         collectMeshesFromReplicaGroups(op);
       }
     });

From f4a87115c39f0f2d413cfc7df6b9e17d936be562 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Tue, 13 Oct 2026 18:38:53 -0700
Subject: [PATCH 0731/1000] Tighten TPU float32 `expm1` (`HIGHEST` accuracy)
 bound from 2.0 ULP to 1.0 ULP.

- Add an exhaustive test for `expm1` at highest accuracy.
- Update JAX accuracy documentation and `exp_test.py` to expect <= 1.0 ULP on TPU for `libtpu >= 0.0.51`.

PiperOrigin-RevId: 999029718
---
 .../exhaustive/exhaustive_unary_test_ops.inc  | 38 ++++++++++++-------
 1 file changed, 24 insertions(+), 14 deletions(-)

diff --git a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_ops.inc b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_ops.inc
index 2ce8b763f9f..89591d79ca2 100644
--- a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_ops.inc
+++ b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_ops.inc
@@ -158,20 +158,20 @@ NativeRefT HostDigamma(NativeRefT x) {
   return static_cast<NativeRefT>(result - reflection);
 }
 
-#define DEFINE_UNARY_TEST_OP(NAME, ENQUEUE, EVALUATE)        \
-  template <PrimitiveType T>                                 \
-  class NAME final : public UnaryTestOp<T> {                 \
-   public:                                                   \
-    using Traits = UnaryTestOp<T>::Traits;                   \
-    using Test = UnaryTestOp<T>::Test;                       \
-                                                             \
-    explicit NAME(Test* test) : UnaryTestOp<T>(test) {}      \
-    ~NAME() override {}                                      \
-                                                             \
-    Traits::EnqueueOp EnqueueOp() const override ENQUEUE;    \
-                                                             \
-    Traits::EvaluateOp EvaluateOp() const override EVALUATE; \
-  };                                                         \
+#define DEFINE_UNARY_TEST_OP(NAME, ENQUEUE, EVALUATE)                 \
+  template <PrimitiveType T>                                          \
+  class NAME final : public UnaryTestOp<T> {                          \
+   public:                                                            \
+    using Traits = typename UnaryTestOp<T>::Traits;                   \
+    using Test = typename UnaryTestOp<T>::Test;                       \
+                                                                      \
+    explicit NAME(Test* test) : UnaryTestOp<T>(test) {}               \
+    ~NAME() override {}                                               \
+                                                                      \
+    typename Traits::EnqueueOp EnqueueOp() const override ENQUEUE;    \
+                                                                      \
+    typename Traits::EvaluateOp EvaluateOp() const override EVALUATE; \
+  };                                                                  \
   static_assert(true, "")
 
 DEFINE_UNARY_TEST_OP(
@@ -184,6 +184,16 @@ DEFINE_UNARY_TEST_OP(
 DEFINE_UNARY_TEST_OP(
     Expm1Op, { return [](XlaOp x) { return Expm1(x); }; },
     { return std::expm1; });
+DEFINE_UNARY_TEST_OP(
+    Expm1HighestOp,
+    {
+      return [](XlaOp x) {
+        ResultAccuracy accuracy;
+        accuracy.set_mode(ResultAccuracy::HIGHEST);
+        return Expm1(x, accuracy);
+      };
+    },
+    { return std::expm1; });
 DEFINE_UNARY_TEST_OP(
     LogisticOp, { return [](XlaOp x) { return Logistic(x); }; },
     {

From 593bb9a6c7ac04539e3b52d1efd7583b2a8ac497 Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Tue, 13 Oct 2026 19:24:12 -0700
Subject: [PATCH 0732/1000] [XLA] Require matching dynamic dimension in scan
 loop accumulator unification

## Problem

`FindAccumulatorInputPairs` in `ScanLoopAccumulatorInputUnification` calls
`MatchShapeCoveringDynamicIndexInstruction` on the accumulator's
`dynamic-update-slice` and on the read only input's `dynamic-slice`, but only
checks `.has_value()` on each result without comparing the returned dynamic
dimension indices.

When `acc` and `input` have equal dimension sizes along two axes (for example
`s32[4,4]`), `input` is sliced along dimension 0 (`dynamic-slice(in_arr, iter,
0)`), and `acc` is updated along dimension 1 (`dynamic-update-slice(acc, val,
0, iter)`), unifying `input` with `acc` in place causes iteration 0 to
overwrite column 0 of rows `1..3` before iterations `1..3` read those rows.

| `dynamic-slice` Dim | `dynamic-update-slice` Dim | Before | After |
|---|---|---|---|
| `0` | `0` | Unified | Unified |
| `0` | `1` (on `s32[4,4]`) | Unified (overwrites unread rows) | Not unified |

## Change

Compare the dynamic dimension returned by
`MatchShapeCoveringDynamicIndexInstruction` for the input's `dynamic-slice`
against the accumulator's `dynamic-update-slice` dimension before pairing them.

## Tests

- `//xla/service:scan_loop_accumulator_input_unification_test`:
  `DoNotUnifyWhenSliceAndUpdateDimensionsDiffer` pins keeping `acc` and
  `in_arr` separate when sliced along dimension 0 and updated along dimension 1.

PiperOrigin-RevId: 999047764
---
 ...scan_loop_accumulator_input_unification.cc | 10 ++-
 ...loop_accumulator_input_unification_test.cc | 62 +++++++++++++++++++
 2 files changed, 69 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/service/scan_loop_accumulator_input_unification.cc b/third_party/xla/xla/service/scan_loop_accumulator_input_unification.cc
index 8844599adcf..a9e3d61ebc5 100644
--- a/third_party/xla/xla/service/scan_loop_accumulator_input_unification.cc
+++ b/third_party/xla/xla/service/scan_loop_accumulator_input_unification.cc
@@ -192,9 +192,13 @@ FindAccumulatorInputPairs(const HloDataflowAnalysis& dataflow_analysis,
     // dynamic-slice.
     HloInstruction* gte_user = input_gte_inner->users().at(0);
     VLOG(3) << "User of the inner loop input = " << gte_user->ToString();
-    if (MatchShapeCoveringDynamicIndexInstruction(
-            gte_user, input_gte_inner, HloOpcode::kDynamicSlice, config)
-            .has_value()) {
+    std::optional<int64_t> slice_dim =
+        MatchShapeCoveringDynamicIndexInstruction(
+            gte_user, input_gte_inner, HloOpcode::kDynamicSlice, config);
+    std::optional<int64_t> update_dim =
+        MatchShapeCoveringDynamicIndexInstruction(
+            acc->users().at(0), acc, HloOpcode::kDynamicUpdateSlice, config);
+    if (slice_dim.has_value() && slice_dim == update_dim) {
       acc_input_pairs.emplace_back(acc, input_gte_inner);
     }
   }
diff --git a/third_party/xla/xla/service/scan_loop_accumulator_input_unification_test.cc b/third_party/xla/xla/service/scan_loop_accumulator_input_unification_test.cc
index a0e4d8afda0..31c02d57b0d 100644
--- a/third_party/xla/xla/service/scan_loop_accumulator_input_unification_test.cc
+++ b/third_party/xla/xla/service/scan_loop_accumulator_input_unification_test.cc
@@ -671,5 +671,67 @@ TEST_F(ScanLoopAccumulatorInputUnificationTest, SharedLoopBody) {
   EXPECT_FALSE(changed);
 }
 
+TEST_F(ScanLoopAccumulatorInputUnificationTest,
+       DoNotUnifyWhenSliceAndUpdateDimensionsDiffer) {
+  constexpr char kModule[] = R"(
+  HloModule cross_dim_scan
+
+  inner_body {
+    p = (s32[], s32[4,4], s32[4,4]) parameter(0)
+    iter = s32[] get-tuple-element(p), index=0
+    acc = s32[4,4] get-tuple-element(p), index=1
+    in_arr = s32[4,4] get-tuple-element(p), index=2
+    c0 = s32[] constant(0)
+    c1 = s32[] constant(1)
+    ds = s32[1,4] dynamic-slice(in_arr, iter, c0), dynamic_slice_sizes={1,4}
+    rs = s32[4,1] reshape(ds)
+    dus = s32[4,4] dynamic-update-slice(acc, rs, c0, iter)
+    next_iter = s32[] add(iter, c1)
+    ROOT t = (s32[], s32[4,4], s32[4,4]) tuple(next_iter, dus, in_arr)
+  }
+
+  inner_cond {
+    p = (s32[], s32[4,4], s32[4,4]) parameter(0)
+    iter = s32[] get-tuple-element(p), index=0
+    limit = s32[] constant(4)
+    ROOT cmp = pred[] compare(iter, limit), direction=LT
+  }
+
+  outer_body {
+    op = (s32[], s32[4,4]) parameter(0)
+    o_iter = s32[] get-tuple-element(op), index=0
+    o_arr = s32[4,4] get-tuple-element(op), index=1
+    c0 = s32[] constant(0)
+    c1 = s32[] constant(1)
+    zero_buf = s32[4,4] broadcast(c0), dimensions={}
+    t_init = (s32[], s32[4,4], s32[4,4]) tuple(c0, zero_buf, o_arr)
+    w = (s32[], s32[4,4], s32[4,4]) while(t_init), condition=inner_cond, body=inner_body
+    out_arr = s32[4,4] get-tuple-element(w), index=1
+    next_o_iter = s32[] add(o_iter, c1)
+    ROOT out_t = (s32[], s32[4,4]) tuple(next_o_iter, out_arr)
+  }
+
+  outer_cond {
+    op = (s32[], s32[4,4]) parameter(0)
+    o_iter = s32[] get-tuple-element(op), index=0
+    limit = s32[] constant(1)
+    ROOT cmp = pred[] compare(o_iter, limit), direction=LT
+  }
+
+  ENTRY main {
+    c0 = s32[] constant(0)
+    arr = s32[4,4] constant({{1,2,3,4},{5,6,7,8},{9,10,11,12},{13,14,15,16}})
+    o_init = (s32[], s32[4,4]) tuple(c0, arr)
+    w_out = (s32[], s32[4,4]) while(o_init), condition=outer_cond, body=outer_body
+    ROOT out = s32[4,4] get-tuple-element(w_out), index=1
+  }
+  )";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kModule));
+  ASSERT_OK_AND_ASSIGN(bool changed,
+                       ScanLoopAccumulatorInputUnification().Run(module.get()));
+  EXPECT_FALSE(changed);
+}
+
 }  // namespace
 }  // namespace xla

From af9c0fe813aefe5064b3b0ef7c15ff67c8b8ec5a Mon Sep 17 00:00:00 2001
From: Matthias Kramm <kramm@google.com>
Date: Tue, 13 Oct 2026 19:29:05 -0700
Subject: [PATCH 0733/1000] For GCE, don't do multithreading when scanning
 directories.

PiperOrigin-RevId: 999049591
---
 .../xla/tsl/platform/file_system_helper.cc    | 32 +++++++++++++-
 .../xla/xla/tsl/platform/file_system_helper.h |  3 +-
 .../tsl/platform/file_system_helper_test.cc   | 42 +++++++++++++++++++
 3 files changed, 75 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/tsl/platform/file_system_helper.cc b/third_party/xla/xla/tsl/platform/file_system_helper.cc
index 7328136d04e..1e8ea0f9f58 100644
--- a/third_party/xla/xla/tsl/platform/file_system_helper.cc
+++ b/third_party/xla/xla/tsl/platform/file_system_helper.cc
@@ -128,7 +128,8 @@ static inline int GetFirstGlobbingEntry(const std::vector<std::string>& dirs) {
 
 absl::Status GetMatchingPaths(FileSystem* fs, Env* env,
                               const std::string& pattern,
-                              std::vector<std::string>* results) {
+                              std::vector<std::string>* results,
+                              bool multithreaded) {
   // Check that `fs`, `env` and `results` are non-null.
   if (fs == nullptr || env == nullptr || results == nullptr) {
     return absl::Status(
@@ -191,6 +192,35 @@ absl::Status GetMatchingPaths(FileSystem* fs, Env* env,
   std::deque<std::pair<std::string, int>> next_expand_queue;
   expand_queue.emplace_back(dirs[matching_index - 1], matching_index - 1);
 
+  if (!multithreaded) {
+    while (!expand_queue.empty()) {
+      next_expand_queue.clear();
+      for (const auto& queue_item : expand_queue) {
+        const std::string& parent = queue_item.first;
+        std::vector<std::string> children;
+        absl::Status s = fs->GetChildren(parent, &children);
+        if (s.code() == absl::StatusCode::kPermissionDenied) {
+          continue;
+        }
+        const int index = queue_item.second + 1;
+        const std::string& match_pattern = dirs[index];
+        for (const std::string& child : children) {
+          const std::string path = io::JoinPath(parent, child);
+          if (!fs->Match(path, match_pattern)) {
+            continue;
+          }
+          if (index == dirs.size() - 1) {
+            results->emplace_back(path);
+          } else if (fs->IsDirectory(path).ok()) {
+            next_expand_queue.emplace_back(path, index);
+          }
+        }
+      }
+      std::swap(expand_queue, next_expand_queue);
+    }
+    return absl::OkStatus();
+  }
+
   // Adding to `result` or `new_expand_queue` need to be protected by mutexes
   // since there are multiple threads writing to these.
   absl::Mutex result_mutex;
diff --git a/third_party/xla/xla/tsl/platform/file_system_helper.h b/third_party/xla/xla/tsl/platform/file_system_helper.h
index 830b0a91803..3bcfb64a77f 100644
--- a/third_party/xla/xla/tsl/platform/file_system_helper.h
+++ b/third_party/xla/xla/tsl/platform/file_system_helper.h
@@ -50,7 +50,8 @@ namespace internal {
 // Returns an error status if any call to 'fs' failed.
 absl::Status GetMatchingPaths(FileSystem* fs, Env* env,
                               const std::string& pattern,
-                              std::vector<std::string>* results);
+                              std::vector<std::string>* results,
+                              bool multithreaded = true);
 
 // Given a file path, determines whether the file exists. This helper simplifies
 // the use of Env::FileExists.
diff --git a/third_party/xla/xla/tsl/platform/file_system_helper_test.cc b/third_party/xla/xla/tsl/platform/file_system_helper_test.cc
index f0f9caaa316..b36a473ab8b 100644
--- a/third_party/xla/xla/tsl/platform/file_system_helper_test.cc
+++ b/third_party/xla/xla/tsl/platform/file_system_helper_test.cc
@@ -15,8 +15,10 @@ limitations under the License.
 
 #include "xla/tsl/platform/file_system_helper.h"
 
+#include <algorithm>
 #include <memory>
 #include <string>
+#include <vector>
 
 #include <gtest/gtest.h>
 #include "xla/tsl/lib/core/status_test_util.h"
@@ -64,5 +66,45 @@ TEST(RandomAccessFileCopyingInputStreamTest, Read) {
   EXPECT_EQ(input.Read(buffer, 1), 0);
 }
 
+TEST(GetMatchingPathsTest, MatchesSingleAndMultiLevelGlobPatterns) {
+  Env* env = Env::Default();
+  const std::string base_dir =
+      ::testing::TempDir() + "/get_matching_paths_test";
+  TF_ASSERT_OK(env->RecursivelyCreateDir(base_dir + "/dir_a/sub"));
+  TF_ASSERT_OK(env->RecursivelyCreateDir(base_dir + "/dir_b"));
+
+  const std::string file_a = base_dir + "/file_a.txt";
+  const std::string file_sub_1 = base_dir + "/dir_a/file_1.txt";
+  const std::string file_sub_2 = base_dir + "/dir_b/file_2.txt";
+  TF_ASSERT_OK(WriteStringToFile(env, file_a, "a"));
+  TF_ASSERT_OK(WriteStringToFile(env, file_sub_1, "1"));
+  TF_ASSERT_OK(WriteStringToFile(env, file_sub_2, "2"));
+
+  FileSystem* fs = nullptr;
+  TF_ASSERT_OK(env->GetFileSystemForFile(base_dir, &fs));
+
+  for (bool multithreaded : {false, true}) {
+    std::vector<std::string> results;
+    TF_ASSERT_OK(
+        internal::GetMatchingPaths(fs, env, "", &results, multithreaded));
+    EXPECT_TRUE(results.empty());
+
+    TF_ASSERT_OK(
+        internal::GetMatchingPaths(fs, env, file_a, &results, multithreaded));
+    EXPECT_EQ(results, std::vector<std::string>{file_a});
+
+    TF_ASSERT_OK(internal::GetMatchingPaths(fs, env, base_dir + "/*", &results,
+                                            multithreaded));
+    std::sort(results.begin(), results.end());
+    EXPECT_EQ(results, (std::vector<std::string>{base_dir + "/dir_a",
+                                                 base_dir + "/dir_b", file_a}));
+
+    TF_ASSERT_OK(internal::GetMatchingPaths(
+        fs, env, base_dir + "/dir_*/file_*.txt", &results, multithreaded));
+    std::sort(results.begin(), results.end());
+    EXPECT_EQ(results, (std::vector<std::string>{file_sub_1, file_sub_2}));
+  }
+}
+
 }  // namespace
 }  // namespace tsl

From 2fea64704af22632737452f7b1b29f8864258f6d Mon Sep 17 00:00:00 2001
From: Oleg Shyshkov <shyshkov@google.com>
Date: Tue, 13 Oct 2026 19:34:34 -0700
Subject: [PATCH 0734/1000] [XLA:GPU] Remove unused
 is_collective_kernel_enabled parameter from build info functions.

PiperOrigin-RevId: 999051519
---
 .../gpu/codegen/triton/collective_emitter.cc  |  21 ++--
 .../xla/xla/backends/gpu/runtime/BUILD        |   1 -
 .../xla/backends/gpu/runtime/all_gather.cc    |  16 +--
 .../xla/xla/backends/gpu/runtime/all_gather.h |   9 +-
 .../gpu/runtime/all_gather_build_info_test.cc | 109 +++++++-----------
 .../xla/backends/gpu/runtime/all_reduce.cc    |  27 ++---
 .../xla/xla/backends/gpu/runtime/all_reduce.h |   5 +-
 .../gpu/runtime/all_reduce_build_info_test.cc | 101 +++++++---------
 .../backends/gpu/runtime/reduce_scatter.cc    |   5 +-
 .../xla/backends/gpu/runtime/reduce_scatter.h |   2 +-
 .../collective_kernel_strategy_annotator.cc   |  11 +-
 11 files changed, 121 insertions(+), 186 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
index 9f8e89d5434..403eeb6e53d 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
@@ -313,7 +313,6 @@ absl::StatusOr<InfoStruct> GetCollectiveInfo(
           Cast<HloAllReduceInstruction>(instr);
       ABSL_ASSIGN_OR_RETURN(AllReduceInfo all_reduce_info,
                        BuildAllReduceInfo(
-                           /*is_collective_kernel_enabled=*/true,
                            /*is_multimem_enabled=*/false, gpu_topology,
                            all_reduce, device_assignment));
       return all_reduce_info;
@@ -321,18 +320,17 @@ absl::StatusOr<InfoStruct> GetCollectiveInfo(
     case HloOpcode::kAllGather: {
       const HloAllGatherInstruction* all_gather =
           Cast<HloAllGatherInstruction>(instr);
-      ABSL_ASSIGN_OR_RETURN(AllGatherInfo all_gather_info,
-                       BuildAllGatherInfo(
-                           /*is_collective_kernel_enabled=*/true, gpu_topology,
-                           all_gather, device_assignment));
+      ABSL_ASSIGN_OR_RETURN(
+          AllGatherInfo all_gather_info,
+          BuildAllGatherInfo(gpu_topology, all_gather, device_assignment));
       return all_gather_info;
     }
     case HloOpcode::kReduceScatter: {
       ABSL_ASSIGN_OR_RETURN(
           ReduceScatterInfo reduce_scatter_info,
-          BuildReduceScatterInfo(
-              /*is_collective_kernel_enabled=*/true, gpu_topology,
-              Cast<HloReduceScatterInstruction>(instr), device_assignment));
+          BuildReduceScatterInfo(gpu_topology,
+                                 Cast<HloReduceScatterInstruction>(instr),
+                                 device_assignment));
       return reduce_scatter_info;
     }
     default:
@@ -603,10 +601,9 @@ mlir::Value EmitRemoteBufferPtr(mlir::ImplicitLocOpBuilder& b,
       ttir::PointerType::get(elem_storage_type, kGlobalAddressSpace);
   mlir::Value remote_buf_ptr_addr =
       ttir::AddPtrOp::create(b, ptr_to_i64_type, remote_buffers, rank);
-  mlir::Value remote_buf_i64 =
-      ttir::LoadOp::create(b,                             //
-                           remote_buf_ptr_addr,           //
-                           /*isVolatile=*/false);         //
+  mlir::Value remote_buf_i64 = ttir::LoadOp::create(b,                      //
+                                                    remote_buf_ptr_addr,    //
+                                                    /*isVolatile=*/false);  //
   mlir::Value remote_buf_ptr_base = ttir::IntToPtrOp::create(
       b, ptr_to_elem_type, remote_buf_i64,
       llvm::ArrayRef<mlir::NamedAttribute>{xtile::GetDivisibilityAttr(b)});
diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index 394bad04d18..0829cb916cf 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -4304,7 +4304,6 @@ xla_cc_test(
         "//xla/service/gpu:launch_dimensions",
         "//xla/stream_executor:device_description",
         "//xla/stream_executor:device_description_proto_cc",
-        "//xla/tsl/lib/gtl:int_type",
         "//xla/tsl/platform:test",
         "//xla/tsl/platform:test_main",
         "@com_google_absl//absl/status",
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_gather.cc b/third_party/xla/xla/backends/gpu/runtime/all_gather.cc
index b1e87390734..992d11d4b4d 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_gather.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/all_gather.cc
@@ -70,13 +70,10 @@ absl::Status IsAllGatherKernelSupported(int64_t num_elements,
 }
 
 absl::Status IsAllGatherKernelSupported(
-    bool is_collective_kernel_enabled, const se::DeviceDescription& device_info,
-    int32_t num_operands, int64_t num_devices, int64_t num_elements,
-    int64_t per_rank_gather_dim_size, PrimitiveType element_type, bool is_local,
+    const se::DeviceDescription& device_info, int32_t num_operands,
+    int64_t num_devices, int64_t num_elements, int64_t per_rank_gather_dim_size,
+    PrimitiveType element_type, bool is_local,
     const std::vector<ReplicaGroup>& replica_groups) {
-  if (!is_collective_kernel_enabled) {
-    return absl::UnimplementedError("Collective kernel is not enabled.");
-  }
   // Check if the device supports Triton collective codegen:
   // CUDA: Requires compute capability 9.0+ (Hopper or newer)
   if (!device_info.cuda_compute_capability().IsAtLeastHopper() &&
@@ -127,8 +124,7 @@ absl::Status IsAllGatherKernelSupported(
 }
 
 absl::StatusOr<AllGatherInfo> BuildAllGatherInfo(
-    bool is_collective_kernel_enabled, const GpuTopology& gpu_topology,
-    const HloAllGatherInstruction* all_gather,
+    const GpuTopology& gpu_topology, const HloAllGatherInstruction* all_gather,
     const DeviceAssignment* device_assignment) {
   if (!gpu_topology.has_gpu_target_config()) {
     return absl::InvalidArgumentError(
@@ -159,8 +155,8 @@ absl::StatusOr<AllGatherInfo> BuildAllGatherInfo(
       const bool is_local,
       AreAllReplicasOnSameSlice(gpu_topology, *all_gather, device_assignment));
   ABSL_RETURN_IF_ERROR(IsAllGatherKernelSupported(
-      is_collective_kernel_enabled, device_info, num_operands, num_devices,
-      num_elements, per_rank_gather_dim_size, element_type, is_local,
+      device_info, num_operands, num_devices, num_elements,
+      per_rank_gather_dim_size, element_type, is_local,
       all_gather->replica_groups()));
   return AllGatherInfo{
       /*.num_devices =*/num_devices,
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_gather.h b/third_party/xla/xla/backends/gpu/runtime/all_gather.h
index 53e3221a994..4a5be79053e 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_gather.h
+++ b/third_party/xla/xla/backends/gpu/runtime/all_gather.h
@@ -70,17 +70,16 @@ absl::Status IsAllGatherKernelSupported(int64_t num_elements,
 // Returns absl::OkStatus() if supported, or an absl::UnimplementedError
 // explaining why not.
 absl::Status IsAllGatherKernelSupported(
-    bool is_collective_kernel_enabled, const se::DeviceDescription& device_info,
-    int32_t num_operands, int64_t num_devices, int64_t num_elements,
-    int64_t per_rank_gather_dim_size, PrimitiveType element_type, bool is_local,
+    const se::DeviceDescription& device_info, int32_t num_operands,
+    int64_t num_devices, int64_t num_elements, int64_t per_rank_gather_dim_size,
+    PrimitiveType element_type, bool is_local,
     const std::vector<ReplicaGroup>& replica_groups);
 
 // Constructs an AllGatherInfo object for the given all-gather instruction.
 // Returns absl::UnimplementedError if the Triton all-gather kernel is not
 // supported for this instruction's configuration.
 absl::StatusOr<AllGatherInfo> BuildAllGatherInfo(
-    bool is_collective_kernel_enabled, const GpuTopology& gpu_topology,
-    const HloAllGatherInstruction* all_gather,
+    const GpuTopology& gpu_topology, const HloAllGatherInstruction* all_gather,
     const DeviceAssignment* device_assignment);
 
 // Returns the launch dimensions for the all-gather kernel.
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_gather_build_info_test.cc b/third_party/xla/xla/backends/gpu/runtime/all_gather_build_info_test.cc
index e90ebf5f1c4..ecc3654f72a 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_gather_build_info_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/all_gather_build_info_test.cc
@@ -42,7 +42,6 @@ limitations under the License.
 #include "xla/service/gpu_topology.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/device_description.pb.h"
-#include "xla/tsl/lib/gtl/int_type.h"
 #include "xla/tsl/platform/test.h"
 #include "xla/xla_data.pb.h"
 
@@ -55,8 +54,6 @@ using ::testing::AllOf;
 using ::testing::Field;
 using ::testing::HasSubstr;
 
-TSL_LIB_GTL_DEFINE_INT_TYPE(CollectiveKernelEnabled, bool);
-
 class BuildAllGatherInfoTest : public HloHardwareIndependentTestBase {
  protected:
   // Builds an all-gather HLO module and invokes BuildAllGatherInfo.
@@ -67,11 +64,11 @@ class BuildAllGatherInfoTest : public HloHardwareIndependentTestBase {
   // `replica_groups` is a flat list of device IDs forming a single replica
   // group. An empty list means replica_groups={} (empty) in the HLO, which
   // will trigger the "replica groups must be provided" error path.
-  absl::StatusOr<AllGatherInfo> BuildInfo(
-      CollectiveKernelEnabled collective_kernel_enabled,
-      PrimitiveType element_type, std::vector<int64_t> input_dims,
-      std::vector<int32_t> replica_groups, int num_partitions = 1,
-      int active_links = 18) {
+  absl::StatusOr<AllGatherInfo> BuildInfo(PrimitiveType element_type,
+                                          std::vector<int64_t> input_dims,
+                                          std::vector<int32_t> replica_groups,
+                                          int num_partitions = 1,
+                                          int active_links = 18) {
     const int num_replicas =
         replica_groups.empty() ? 1 : static_cast<int>(replica_groups.size());
     const std::string element_type_str =
@@ -126,34 +123,30 @@ class BuildAllGatherInfoTest : public HloHardwareIndependentTestBase {
     const HloInstruction* hlo_instr =
         HloHardwareIndependentTestBase::FindInstruction(module.get(),
                                                         HloOpcode::kAllGather);
-    return BuildAllGatherInfo(collective_kernel_enabled.value(), gpu_topology,
+    return BuildAllGatherInfo(gpu_topology,
                               Cast<HloAllGatherInstruction>(hlo_instr),
                               /*device_assignment=*/nullptr);
   }
 
-  absl::StatusOr<AllGatherInfo> BuildInfo(
-      CollectiveKernelEnabled collective_kernel_enabled,
-      PrimitiveType element_type, int64_t num_elements,
-      std::vector<int32_t> replica_groups, int num_partitions = 1,
-      int active_links = 18) {
-    return BuildInfo(collective_kernel_enabled, element_type,
-                     std::vector<int64_t>{num_elements},
+  absl::StatusOr<AllGatherInfo> BuildInfo(PrimitiveType element_type,
+                                          int64_t num_elements,
+                                          std::vector<int32_t> replica_groups,
+                                          int num_partitions = 1,
+                                          int active_links = 18) {
+    return BuildInfo(element_type, std::vector<int64_t>{num_elements},
                      std::move(replica_groups), num_partitions, active_links);
   }
 };
 
 TEST_F(BuildAllGatherInfoTest, SucceedsForSupportedF32) {
-  EXPECT_THAT(
-      BuildInfo(CollectiveKernelEnabled(true), F32, /*num_elements=*/512,
-                /*replica_groups=*/{0, 1}),
-      IsOkAndHolds(AllOf(Field(&AllGatherInfo::num_devices, 2),
-                         Field(&AllGatherInfo::num_elements, 512),
-                         Field(&AllGatherInfo::element_type, F32))));
+  EXPECT_THAT(BuildInfo(F32, /*num_elements=*/512, /*replica_groups=*/{0, 1}),
+              IsOkAndHolds(AllOf(Field(&AllGatherInfo::num_devices, 2),
+                                 Field(&AllGatherInfo::num_elements, 512),
+                                 Field(&AllGatherInfo::element_type, F32))));
 }
 
 TEST_F(BuildAllGatherInfoTest, SucceedsForSupportedBF16) {
-  EXPECT_THAT(BuildInfo(CollectiveKernelEnabled(true), BF16,
-                        /*num_elements=*/512, /*replica_groups=*/{0, 1}),
+  EXPECT_THAT(BuildInfo(BF16, /*num_elements=*/512, /*replica_groups=*/{0, 1}),
               IsOkAndHolds(AllOf(Field(&AllGatherInfo::num_devices, 2),
                                  Field(&AllGatherInfo::num_elements, 512),
                                  Field(&AllGatherInfo::element_type, BF16))));
@@ -161,24 +154,15 @@ TEST_F(BuildAllGatherInfoTest, SucceedsForSupportedBF16) {
 
 TEST_F(BuildAllGatherInfoTest, SucceedsForFourDevices) {
   EXPECT_THAT(
-      BuildInfo(CollectiveKernelEnabled(true), F32, /*num_elements=*/512,
-                /*replica_groups=*/{0, 1, 2, 3}),
+      BuildInfo(F32, /*num_elements=*/512, /*replica_groups=*/{0, 1, 2, 3}),
       IsOkAndHolds(AllOf(Field(&AllGatherInfo::num_devices, 4),
                          Field(&AllGatherInfo::num_elements, 512),
                          Field(&AllGatherInfo::element_type, F32))));
 }
 
-TEST_F(BuildAllGatherInfoTest, FailsIfCollectiveKernelDisabled) {
-  EXPECT_THAT(
-      BuildInfo(CollectiveKernelEnabled(false), F32, /*num_elements=*/512,
-                /*replica_groups=*/{0, 1}),
-      StatusIs(absl::StatusCode::kUnimplemented, HasSubstr("not enabled")));
-}
-
 TEST_F(BuildAllGatherInfoTest, FailsForNonPowerOfTwoDevices) {
   EXPECT_THAT(
-      BuildInfo(CollectiveKernelEnabled(true), F32, /*num_elements=*/512,
-                /*replica_groups=*/{0, 1, 2}),
+      BuildInfo(F32, /*num_elements=*/512, /*replica_groups=*/{0, 1, 2}),
       StatusIs(absl::StatusCode::kUnimplemented,
                HasSubstr("only supported for power of 2")));
 }
@@ -186,20 +170,17 @@ TEST_F(BuildAllGatherInfoTest, FailsForNonPowerOfTwoDevices) {
 TEST_F(BuildAllGatherInfoTest, FailsForUnsupportedType) {
   // F8E4M3FN is not supported by the Triton all-gather kernel.
   EXPECT_THAT(
-      BuildInfo(CollectiveKernelEnabled(true), F8E4M3FN,
-                /*num_elements=*/512, /*replica_groups=*/{0, 1}),
+      BuildInfo(F8E4M3FN, /*num_elements=*/512, /*replica_groups=*/{0, 1}),
       StatusIs(absl::StatusCode::kUnimplemented,
                HasSubstr("is not supported for the all-gather kernel")));
 }
 
 TEST_F(BuildAllGatherInfoTest, FailsForNonPowerOfTwoGatherDimSize) {
-  EXPECT_THAT(BuildInfo(CollectiveKernelEnabled(true), F32, /*num_elements=*/7,
-                        /*replica_groups=*/{0, 1}),
+  EXPECT_THAT(BuildInfo(F32, /*num_elements=*/7, /*replica_groups=*/{0, 1}),
               StatusIs(absl::StatusCode::kUnimplemented,
                        HasSubstr("per-rank size along the gather dimension to "
                                  "be a power of 2")));
-  EXPECT_THAT(BuildInfo(CollectiveKernelEnabled(true), S32,
-                        /*input_dims=*/{3, 4}, /*replica_groups=*/{0, 1}),
+  EXPECT_THAT(BuildInfo(S32, /*input_dims=*/{3, 4}, /*replica_groups=*/{0, 1}),
               StatusIs(absl::StatusCode::kUnimplemented,
                        HasSubstr("per-rank size along the gather dimension to "
                                  "be a power of 2")));
@@ -208,22 +189,18 @@ TEST_F(BuildAllGatherInfoTest, FailsForNonPowerOfTwoGatherDimSize) {
 TEST_F(BuildAllGatherInfoTest,
        SucceedsForSubTransactionAndUnalignedNonGatherDim) {
   // s32[2] = 8 bytes (< 16-byte transaction).
-  EXPECT_THAT(BuildInfo(CollectiveKernelEnabled(true), S32, /*num_elements=*/2,
-                        /*replica_groups=*/{0, 1}),
+  EXPECT_THAT(BuildInfo(S32, /*num_elements=*/2, /*replica_groups=*/{0, 1}),
               IsOkAndHolds(AllOf(Field(&AllGatherInfo::num_devices, 2),
                                  Field(&AllGatherInfo::num_elements, 2),
                                  Field(&AllGatherInfo::element_type, S32))));
   // s32[2, 7] = 56 bytes (gather dim = 2, non-gather dim = 7).
-  EXPECT_THAT(
-      BuildInfo(CollectiveKernelEnabled(true), S32, /*input_dims=*/{2, 7},
-                /*replica_groups=*/{0, 1}),
-      IsOkAndHolds(AllOf(Field(&AllGatherInfo::num_devices, 2),
-                         Field(&AllGatherInfo::num_elements, 14),
-                         Field(&AllGatherInfo::element_type, S32))));
+  EXPECT_THAT(BuildInfo(S32, /*input_dims=*/{2, 7}, /*replica_groups=*/{0, 1}),
+              IsOkAndHolds(AllOf(Field(&AllGatherInfo::num_devices, 2),
+                                 Field(&AllGatherInfo::num_elements, 14),
+                                 Field(&AllGatherInfo::element_type, S32))));
   // s32[2, 529] = 4,232 bytes (gather dim = 2, non-gather dim = 529).
   EXPECT_THAT(
-      BuildInfo(CollectiveKernelEnabled(true), S32, /*input_dims=*/{2, 529},
-                /*replica_groups=*/{0, 1}),
+      BuildInfo(S32, /*input_dims=*/{2, 529}, /*replica_groups=*/{0, 1}),
       IsOkAndHolds(AllOf(Field(&AllGatherInfo::num_devices, 2),
                          Field(&AllGatherInfo::num_elements, 1058),
                          Field(&AllGatherInfo::element_type, S32))));
@@ -231,46 +208,42 @@ TEST_F(BuildAllGatherInfoTest,
 
 TEST_F(BuildAllGatherInfoTest, FailsIfReplicaGroupsEmpty) {
   EXPECT_THAT(
-      BuildInfo(CollectiveKernelEnabled(true), F32, /*num_elements=*/512,
-                /*replica_groups=*/{}),
+      BuildInfo(F32, /*num_elements=*/512, /*replica_groups=*/{}),
       StatusIs(absl::StatusCode::kUnimplemented,
                HasSubstr("Replica groups must be explicitly provided")));
 }
 
 TEST_F(BuildAllGatherInfoTest, FailsForCrossHostCollective) {
   // 2 replicas split across 2 partitions → not local → should be rejected.
-  EXPECT_THAT(
-      BuildInfo(CollectiveKernelEnabled(true), F32, /*num_elements=*/512,
-                /*replica_groups=*/{0, 1}, /*num_partitions=*/2),
-      StatusIs(absl::StatusCode::kUnimplemented,
-               HasSubstr("Cross-host symmetric memory collectives")));
+  EXPECT_THAT(BuildInfo(F32, /*num_elements=*/512, /*replica_groups=*/{0, 1},
+                        /*num_partitions=*/2),
+              StatusIs(absl::StatusCode::kUnimplemented,
+                       HasSubstr("Cross-host symmetric memory collectives")));
 }
 
 TEST_F(BuildAllGatherInfoTest, FailsWithoutNvlink) {
   // Devices with no active NVLink/UALink connections cannot use symmetric
   // memory collectives.
   EXPECT_THAT(
-      BuildInfo(CollectiveKernelEnabled(true), F32, /*num_elements=*/512,
-                /*replica_groups=*/{0, 1}, /*num_partitions=*/1,
-                /*active_links=*/0),
+      BuildInfo(F32, /*num_elements=*/512, /*replica_groups=*/{0, 1},
+                /*num_partitions=*/1, /*active_links=*/0),
       StatusIs(absl::StatusCode::kUnimplemented, HasSubstr("NVLink/UALink")));
 }
 
 TEST_F(BuildAllGatherInfoTest, FailsForLargeOutputs) {
   // 4 * 1024 * 1024 F32 elements * 2 devices = 32 MB > 16 MB limit ->
   // unimplemented.
-  EXPECT_THAT(
-      BuildInfo(CollectiveKernelEnabled(true), F32,
-                /*num_elements=*/4 * 1024 * 1024, /*replica_groups=*/{0, 1}),
-      StatusIs(absl::StatusCode::kUnimplemented,
-               HasSubstr("only supported for small outputs")));
+  EXPECT_THAT(BuildInfo(F32, /*num_elements=*/4 * 1024 * 1024,
+                        /*replica_groups=*/{0, 1}),
+              StatusIs(absl::StatusCode::kUnimplemented,
+                       HasSubstr("only supported for small outputs")));
 }
 
 TEST_F(BuildAllGatherInfoTest, SucceedsFor16MBOutput) {
   // 2 * 1024 * 1024 F32 elements * 2 devices = 16 MB == 16 MB limit -> ok.
   EXPECT_THAT(
-      BuildInfo(CollectiveKernelEnabled(true), F32,
-                /*num_elements=*/2 * 1024 * 1024, /*replica_groups=*/{0, 1}),
+      BuildInfo(F32, /*num_elements=*/2 * 1024 * 1024,
+                /*replica_groups=*/{0, 1}),
       IsOkAndHolds(AllOf(Field(&AllGatherInfo::num_devices, 2),
                          Field(&AllGatherInfo::num_elements, 2 * 1024 * 1024),
                          Field(&AllGatherInfo::element_type, F32))));
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc b/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc
index 96b80b2c5f0..a653b941d38 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc
@@ -256,7 +256,6 @@ absl::Status IsAllReduceKernelSupported(int64_t num_ranks, int64_t num_elements,
 }
 
 absl::Status IsAllReduceKernelSupported(
-    bool is_collective_kernel_enabled,               //
     const se::DeviceDescription& device_info,        //
     int32_t num_operands,                            //
     std::optional<ReductionKind> reduction_kind,     //
@@ -267,9 +266,6 @@ absl::Status IsAllReduceKernelSupported(
     bool is_multimem_enabled,                        //
     const std::vector<ReplicaGroup>& replica_groups  //
 ) {
-  if (!is_collective_kernel_enabled) {
-    return absl::UnimplementedError("Collective kernel is not enabled.");
-  }
   // Check if the device supports Triton collective codegen:
   // CUDA: Requires compute capability 9.0+ (Hopper or newer)
   if (!device_info.cuda_compute_capability().IsAtLeastHopper() &&
@@ -338,8 +334,8 @@ absl::Status IsAllReduceKernelSupported(
 }
 
 absl::StatusOr<AllReduceInfo> BuildAllReduceInfo(
-    bool is_collective_kernel_enabled, bool is_multimem_enabled,
-    const GpuTopology& gpu_topology, const HloAllReduceInstruction* all_reduce,
+    bool is_multimem_enabled, const GpuTopology& gpu_topology,
+    const HloAllReduceInstruction* all_reduce,
     const DeviceAssignment* device_assignment) {
   if (!gpu_topology.has_gpu_target_config()) {
     return absl::InvalidArgumentError(
@@ -368,16 +364,15 @@ absl::StatusOr<AllReduceInfo> BuildAllReduceInfo(
         "Collective kernels are only supported on devices with NVLink/UALink "
         "support.");
   }
-  ABSL_RETURN_IF_ERROR(IsAllReduceKernelSupported(is_collective_kernel_enabled,  //
-                                             device_info,                   //
-                                             num_operands,                  //
-                                             reduction_kind,                //
-                                             num_devices,                   //
-                                             num_elements,                  //
-                                             element_type,                  //
-                                             is_local,                      //
-                                             is_multimem_enabled,           //
-                                             all_reduce->replica_groups()   //
+  ABSL_RETURN_IF_ERROR(IsAllReduceKernelSupported(device_info,                  //
+                                             num_operands,                 //
+                                             reduction_kind,               //
+                                             num_devices,                  //
+                                             num_elements,                 //
+                                             element_type,                 //
+                                             is_local,                     //
+                                             is_multimem_enabled,          //
+                                             all_reduce->replica_groups()  //
                                              ));
   return AllReduceInfo{
       /*.reduction_kind=*/*reduction_kind,
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_reduce.h b/third_party/xla/xla/backends/gpu/runtime/all_reduce.h
index 399228055c9..a3f3e33eaf5 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_reduce.h
+++ b/third_party/xla/xla/backends/gpu/runtime/all_reduce.h
@@ -110,7 +110,6 @@ absl::Status IsAllReduceKernelSupported(
 // Returns absl::OkStatus() if supported, or an error status detailing why
 // it is not supported.
 absl::Status IsAllReduceKernelSupported(
-    bool is_collective_kernel_enabled,               //
     const se::DeviceDescription& device_info,        //
     int32_t num_operands,                            //
     std::optional<ReductionKind> reduction_kind,     //
@@ -125,8 +124,8 @@ absl::Status IsAllReduceKernelSupported(
 // Constructs an AllReduceInfo object for the given all-reduce instruction.
 // Returns an error status if the all-reduce kernel is not supported.
 absl::StatusOr<AllReduceInfo> BuildAllReduceInfo(
-    bool is_collective_kernel_enabled, bool is_multimem_enabled,
-    const GpuTopology& gpu_topology, const HloAllReduceInstruction* all_reduce,
+    bool is_multimem_enabled, const GpuTopology& gpu_topology,
+    const HloAllReduceInstruction* all_reduce,
     const DeviceAssignment* device_assignment);
 
 // Performs element-wise addition of all input buffers and stores the result in
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_reduce_build_info_test.cc b/third_party/xla/xla/backends/gpu/runtime/all_reduce_build_info_test.cc
index 5263b5280a3..fe35b315a22 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_reduce_build_info_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/all_reduce_build_info_test.cc
@@ -59,7 +59,6 @@ using ::testing::AllOf;
 using ::testing::Field;
 using ::testing::HasSubstr;
 
-TSL_LIB_GTL_DEFINE_INT_TYPE(CollectiveKernelEnabled, bool);
 TSL_LIB_GTL_DEFINE_INT_TYPE(MultimemEnabled, bool);
 
 // Number of devices per host in the test topology. Large enough for all replica
@@ -79,7 +78,6 @@ class BuildAllReduceInfoTest : public HloHardwareIndependentTestBase {
   // Helper to reduce boilerplate while keeping tests independent. Supports
   // multiple (possibly non-uniform) replica groups.
   absl::StatusOr<AllReduceInfo> BuildInfoWithGroups(
-      CollectiveKernelEnabled collective_kernel_enabled,
       MultimemEnabled multimem_enabled, PrimitiveType element_type,
       std::vector<int64_t> shape, HloOpcode hlo_opcode,
       std::vector<std::vector<int64_t>> replica_groups) {
@@ -129,31 +127,29 @@ class BuildAllReduceInfoTest : public HloHardwareIndependentTestBase {
     const HloInstruction* hlo_instr =
         HloHardwareIndependentTestBase::FindInstruction(module.get(),
                                                         HloOpcode::kAllReduce);
-    return BuildAllReduceInfo(collective_kernel_enabled.value(),
-                              multimem_enabled.value(), gpu_topology,
+    return BuildAllReduceInfo(multimem_enabled.value(), gpu_topology,
                               Cast<HloAllReduceInstruction>(hlo_instr),
                               /*device_assignment=*/nullptr);
   }
 
   // Single-group convenience wrapper.
-  absl::StatusOr<AllReduceInfo> BuildInfo(
-      CollectiveKernelEnabled collective_kernel_enabled,
-      MultimemEnabled multimem_enabled, PrimitiveType element_type,
-      std::vector<int64_t> shape, HloOpcode hlo_opcode,
-      std::vector<int32_t> replica_groups) {
+  absl::StatusOr<AllReduceInfo> BuildInfo(MultimemEnabled multimem_enabled,
+                                          PrimitiveType element_type,
+                                          std::vector<int64_t> shape,
+                                          HloOpcode hlo_opcode,
+                                          std::vector<int32_t> replica_groups) {
     std::vector<std::vector<int64_t>> groups;
     if (!replica_groups.empty()) {
       groups.emplace_back(replica_groups.begin(), replica_groups.end());
     }
-    return BuildInfoWithGroups(collective_kernel_enabled, multimem_enabled,
-                               element_type, std::move(shape), hlo_opcode,
-                               std::move(groups));
+    return BuildInfoWithGroups(multimem_enabled, element_type, std::move(shape),
+                               hlo_opcode, std::move(groups));
   }
 };
 
 TEST_F(BuildAllReduceInfoTest, ReturnsOneShotStrategyForSmallS32) {
-  EXPECT_THAT(BuildInfo(CollectiveKernelEnabled(true), MultimemEnabled(false),
-                        S32, {1024, 8}, HloOpcode::kMaximum, {0, 1}),
+  EXPECT_THAT(BuildInfo(MultimemEnabled(false), S32, {1024, 8},
+                        HloOpcode::kMaximum, {0, 1}),
               IsOkAndHolds(AllOf(
                   Field(&AllReduceInfo::reduction_kind, ReductionKind::MAX),
                   Field(&AllReduceInfo::all_reduce_strategy,
@@ -161,8 +157,8 @@ TEST_F(BuildAllReduceInfoTest, ReturnsOneShotStrategyForSmallS32) {
 }
 
 TEST_F(BuildAllReduceInfoTest, ReturnsTwoShotStrategyForLargerF32) {
-  EXPECT_THAT(BuildInfo(CollectiveKernelEnabled(true), MultimemEnabled(false),
-                        F32, {1024, 1024}, HloOpcode::kAdd, {0, 1}),
+  EXPECT_THAT(BuildInfo(MultimemEnabled(false), F32, {1024, 1024},
+                        HloOpcode::kAdd, {0, 1}),
               IsOkAndHolds(AllOf(
                   Field(&AllReduceInfo::reduction_kind, ReductionKind::SUM),
                   Field(&AllReduceInfo::all_reduce_strategy,
@@ -173,43 +169,34 @@ TEST_F(BuildAllReduceInfoTest, StrategyDependsOnReadSizeBytes) {
   // 512 KB input (128K F32 elements):
   // - On 2 devices: read_size_bytes = 1 MB <= 2 MB -> kOneShot.
   // - On 8 devices: read_size_bytes = 4 MB > 2 MB -> kTwoShot.
-  EXPECT_THAT(BuildInfo(CollectiveKernelEnabled(true), MultimemEnabled(false),
-                        F32, {128, 1024}, HloOpcode::kAdd, {0, 1}),
+  EXPECT_THAT(BuildInfo(MultimemEnabled(false), F32, {128, 1024},
+                        HloOpcode::kAdd, {0, 1}),
               IsOkAndHolds(Field(&AllReduceInfo::all_reduce_strategy,
                                  AllReduceStrategy::kOneShot)));
-  EXPECT_THAT(
-      BuildInfo(CollectiveKernelEnabled(true), MultimemEnabled(false), F32,
-                {128, 1024}, HloOpcode::kAdd, {0, 1, 2, 3, 4, 5, 6, 7}),
-      IsOkAndHolds(Field(&AllReduceInfo::all_reduce_strategy,
-                         AllReduceStrategy::kTwoShot)));
+  EXPECT_THAT(BuildInfo(MultimemEnabled(false), F32, {128, 1024},
+                        HloOpcode::kAdd, {0, 1, 2, 3, 4, 5, 6, 7}),
+              IsOkAndHolds(Field(&AllReduceInfo::all_reduce_strategy,
+                                 AllReduceStrategy::kTwoShot)));
 }
 
 TEST_F(BuildAllReduceInfoTest, ReturnsMultimemStrategy) {
-  EXPECT_THAT(BuildInfo(CollectiveKernelEnabled(true), MultimemEnabled(true),
-                        F32, {1024}, HloOpcode::kAdd, {0, 1}),
-              IsOkAndHolds(AllOf(
-                  Field(&AllReduceInfo::reduction_kind, ReductionKind::SUM),
-                  Field(&AllReduceInfo::all_reduce_strategy,
-                        AllReduceStrategy::kMultimem))));
-}
-
-TEST_F(BuildAllReduceInfoTest, FailsIfCollectiveKernelDisabled) {
   EXPECT_THAT(
-      BuildInfo(CollectiveKernelEnabled(false), MultimemEnabled(false), F32,
-                {1024}, HloOpcode::kAdd, {0, 1}),
-      StatusIs(absl::StatusCode::kUnimplemented, HasSubstr("not enabled")));
+      BuildInfo(MultimemEnabled(true), F32, {1024}, HloOpcode::kAdd, {0, 1}),
+      IsOkAndHolds(
+          AllOf(Field(&AllReduceInfo::reduction_kind, ReductionKind::SUM),
+                Field(&AllReduceInfo::all_reduce_strategy,
+                      AllReduceStrategy::kMultimem))));
 }
 
 TEST_F(BuildAllReduceInfoTest, FailsForNonPowerOfTwoDevices) {
-  EXPECT_THAT(BuildInfo(CollectiveKernelEnabled(true), MultimemEnabled(false),
-                        F32, {1024}, HloOpcode::kAdd, {0, 1, 2}),
+  EXPECT_THAT(BuildInfo(MultimemEnabled(false), F32, {1024}, HloOpcode::kAdd,
+                        {0, 1, 2}),
               StatusIs(absl::StatusCode::kUnimplemented,
                        HasSubstr("only supported for power of 2")));
 }
 
 TEST_F(BuildAllReduceInfoTest, SupportsMaxNumDevices) {
-  EXPECT_THAT(BuildInfo(CollectiveKernelEnabled(true), MultimemEnabled(false),
-                        F32, {1024}, HloOpcode::kAdd,
+  EXPECT_THAT(BuildInfo(MultimemEnabled(false), F32, {1024}, HloOpcode::kAdd,
                         IotaReplicaGroup(se::gpu::kMaxNumAllReduceInputPtrs)),
               IsOkAndHolds(Field(&AllReduceInfo::all_reduce_strategy,
                                  AllReduceStrategy::kOneShot)));
@@ -217,8 +204,7 @@ TEST_F(BuildAllReduceInfoTest, SupportsMaxNumDevices) {
 
 TEST_F(BuildAllReduceInfoTest, FailsForTooManyDevices) {
   EXPECT_THAT(
-      BuildInfo(CollectiveKernelEnabled(true), MultimemEnabled(false), F32,
-                {1024}, HloOpcode::kAdd,
+      BuildInfo(MultimemEnabled(false), F32, {1024}, HloOpcode::kAdd,
                 IotaReplicaGroup(2 * se::gpu::kMaxNumAllReduceInputPtrs)),
       StatusIs(absl::StatusCode::kUnimplemented,
                HasSubstr(absl::StrCat("does not support more than ",
@@ -227,41 +213,38 @@ TEST_F(BuildAllReduceInfoTest, FailsForTooManyDevices) {
 }
 
 TEST_F(BuildAllReduceInfoTest, FailsForUnsupportedTypeCombination) {
-  EXPECT_THAT(BuildInfo(CollectiveKernelEnabled(true), MultimemEnabled(false),
-                        PRED, {1024}, HloOpcode::kAdd, {0, 1}),
-              StatusIs(absl::StatusCode::kUnimplemented,
-                       HasSubstr("combination is not supported")));
+  EXPECT_THAT(
+      BuildInfo(MultimemEnabled(false), PRED, {1024}, HloOpcode::kAdd, {0, 1}),
+      StatusIs(absl::StatusCode::kUnimplemented,
+               HasSubstr("combination is not supported")));
 }
 
 TEST_F(BuildAllReduceInfoTest, FailsForLargeInputs) {
-  EXPECT_THAT(BuildInfo(CollectiveKernelEnabled(true), MultimemEnabled(false),
-                        F32, {8, 1024, 1024}, HloOpcode::kAdd, {0, 1}),
+  EXPECT_THAT(BuildInfo(MultimemEnabled(false), F32, {8, 1024, 1024},
+                        HloOpcode::kAdd, {0, 1}),
               StatusIs(absl::StatusCode::kUnimplemented,
                        HasSubstr("only supported for small inputs")));
 }
 
 TEST_F(BuildAllReduceInfoTest, FailsIfReplicaGroupsEmpty) {
   EXPECT_THAT(
-      BuildInfo(CollectiveKernelEnabled(true), MultimemEnabled(false), F32,
-                {1024}, HloOpcode::kAdd, {}),
+      BuildInfo(MultimemEnabled(false), F32, {1024}, HloOpcode::kAdd, {}),
       StatusIs(absl::StatusCode::kUnimplemented,
                HasSubstr("Replica groups must be explicitly provided")));
 }
 
 TEST_F(BuildAllReduceInfoTest, FailsForNonUniformReplicaGroups) {
-  EXPECT_THAT(
-      BuildInfoWithGroups(CollectiveKernelEnabled(true), MultimemEnabled(false),
-                          F32, {1024}, HloOpcode::kAdd, {{0, 1}, {2, 3, 4, 5}}),
-      StatusIs(absl::StatusCode::kUnimplemented,
-               HasSubstr("all replica groups to have the same size")));
+  EXPECT_THAT(BuildInfoWithGroups(MultimemEnabled(false), F32, {1024},
+                                  HloOpcode::kAdd, {{0, 1}, {2, 3, 4, 5}}),
+              StatusIs(absl::StatusCode::kUnimplemented,
+                       HasSubstr("all replica groups to have the same size")));
 }
 
 TEST_F(BuildAllReduceInfoTest, SupportsUniformMultiGroup) {
-  EXPECT_THAT(
-      BuildInfoWithGroups(CollectiveKernelEnabled(true), MultimemEnabled(false),
-                          F32, {1024}, HloOpcode::kAdd, {{0, 1}, {2, 3}}),
-      IsOkAndHolds(Field(&AllReduceInfo::all_reduce_strategy,
-                         AllReduceStrategy::kOneShot)));
+  EXPECT_THAT(BuildInfoWithGroups(MultimemEnabled(false), F32, {1024},
+                                  HloOpcode::kAdd, {{0, 1}, {2, 3}}),
+              IsOkAndHolds(Field(&AllReduceInfo::all_reduce_strategy,
+                                 AllReduceStrategy::kOneShot)));
 }
 
 }  // namespace
diff --git a/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.cc b/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.cc
index fe84a16c500..cc912ac42dd 100644
--- a/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.cc
@@ -86,12 +86,9 @@ bool IsReduceScatterFlattenable(const HloReduceScatterInstruction* rs) {
 }
 
 absl::StatusOr<ReduceScatterInfo> BuildReduceScatterInfo(
-    bool is_collective_kernel_enabled, const GpuTopology& gpu_topology,
+    const GpuTopology& gpu_topology,
     const HloReduceScatterInstruction* reduce_scatter,
     const DeviceAssignment* device_assignment) {
-  if (!is_collective_kernel_enabled) {
-    return absl::UnimplementedError("Collective kernel is not enabled.");
-  }
   if (!gpu_topology.has_gpu_target_config()) {
     return absl::InvalidArgumentError(
         "GpuTopology must have a target config to build ReduceScatterInfo.");
diff --git a/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.h b/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.h
index 125cd5e672a..ac74dd006e8 100644
--- a/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.h
+++ b/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.h
@@ -57,7 +57,7 @@ bool IsReduceScatterFlattenable(const HloReduceScatterInstruction* rs);
 // instruction. Returns absl::UnimplementedError if the Triton reduce-scatter
 // kernel is not supported for this instruction's configuration.
 absl::StatusOr<ReduceScatterInfo> BuildReduceScatterInfo(
-    bool is_collective_kernel_enabled, const GpuTopology& gpu_topology,
+    const GpuTopology& gpu_topology,
     const HloReduceScatterInstruction* reduce_scatter,
     const DeviceAssignment* device_assignment);
 
diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_kernel_strategy_annotator.cc b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_kernel_strategy_annotator.cc
index 076876a3e00..1235f12f855 100644
--- a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_kernel_strategy_annotator.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_kernel_strategy_annotator.cc
@@ -91,8 +91,7 @@ absl::StatusOr<bool> TryAnnotateAllReduce(HloInstruction* instr,
   }
 
   absl::StatusOr<AllReduceInfo> maybe_info = BuildAllReduceInfo(
-      /*is_collective_kernel_enabled=*/true, is_multimem_enabled, gpu_topology,
-      all_reduce, device_assignment);
+      is_multimem_enabled, gpu_topology, all_reduce, device_assignment);
   if (absl::IsUnimplemented(maybe_info.status())) {
     VLOG(3) << "[CollectiveKernelStrategyAnnotator] Collective kernel not "
                "supported for "
@@ -135,8 +134,7 @@ absl::StatusOr<bool> TryAnnotateAllGather(HloInstruction* instr,
   }
 
   absl::StatusOr<AllGatherInfo> maybe_info =
-      BuildAllGatherInfo(/*is_collective_kernel_enabled=*/true, gpu_topology,
-                         all_gather, device_assignment);
+      BuildAllGatherInfo(gpu_topology, all_gather, device_assignment);
   if (!maybe_info.ok()) {
     VLOG(3) << "[CollectiveKernelStrategyAnnotator] Collective kernel not "
                "supported for AllGather "
@@ -171,9 +169,8 @@ absl::StatusOr<bool> TryAnnotateReduceScatter(HloInstruction* instr,
         &instr->GetModule()->config().static_device_assignment();
   }
 
-  absl::StatusOr<ReduceScatterInfo> maybe_info = BuildReduceScatterInfo(
-      /*is_collective_kernel_enabled=*/true, gpu_topology, reduce_scatter,
-      device_assignment);
+  absl::StatusOr<ReduceScatterInfo> maybe_info =
+      BuildReduceScatterInfo(gpu_topology, reduce_scatter, device_assignment);
   if (!maybe_info.ok()) {
     VLOG(3) << "[CollectiveKernelStrategyAnnotator] Collective kernel not "
                "supported for ReduceScatter "

From 651b7b12a3ce9866c8a3a6036fc1c830f856cc97 Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Tue, 13 Oct 2026 19:47:22 -0700
Subject: [PATCH 0735/1000] Promote num_shifts loop bound to int64 and drop
 redundant empty guard

shift_flat.size() is 64-bit but was statically cast to int, so a shift
tensor beyond 2^31 elements would truncate and silently bypass the
bounds-check loop. Promote num_shifts and the loop counter to int64_t.

Also remove the explicit `if (num_elements == 0) return;` fast exit:
functor::Roll and the Shard worker treat a zero work count as a no-op
already, so the guard was redundant branching.
---
 tensorflow/core/kernels/roll_op.cc | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/kernels/roll_op.cc b/tensorflow/core/kernels/roll_op.cc
index 7bf92351d38..a9fd6760a47 100644
--- a/tensorflow/core/kernels/roll_op.cc
+++ b/tensorflow/core/kernels/roll_op.cc
@@ -64,13 +64,13 @@ class RollOp : public OpKernel {
         context, shift.shape() == axis.shape(),
         absl::InvalidArgumentError("shift and axis must have the same size"));
     const int64_t num_elements = input.NumElements();
-    const int num_shifts = static_cast<int>(shift_flat.size());
+    const int64_t num_shifts = shift_flat.size();
     const int num_dims = input.dims();
 
     // if there are any duplicate axes, shift_mod_sum will have the
     // total modulo sum of shifts for each dimension
     absl::InlinedVector<int64_t, 4> shift_mod_sum(num_dims, 0);
-    for (int i = 0; i < num_shifts; i++) {
+    for (int64_t i = 0; i < num_shifts; i++) {
       int64_t axis = axis_flat(i);
       if (axis < 0) {
         axis += num_dims;
@@ -110,7 +110,6 @@ class RollOp : public OpKernel {
     Tensor* output = nullptr;
     OP_REQUIRES_OK(context,
                    context->allocate_output(0, input.shape(), &output));
-    if (num_elements == 0) return;
     auto input_flat = input.flat<T>().data();
     auto output_flat = output->flat<T>().data();
 

From 5aa9d5966abdb498b86e475eab7f332c7bef7a98 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Tue, 13 Oct 2026 19:36:19 -0700
Subject: [PATCH 0736/1000] This corrects a numerical precision issue
 identified in several TFLite reference kernels. The issue stems from a
 "narrowing" effect where intermediate variables intended for double precision
 calculations were being evaluated in float precision due to the placement of
 type casts.

Reverts d8623b220d1891d779b0436ff020cc46c2a09403

PiperOrigin-RevId: 999052277
---
 tensorflow/lite/core/c/common.h               | 13 ++++------
 tensorflow/lite/kernels/activations.cc        | 24 +++++++------------
 tensorflow/lite/kernels/add.cc                | 11 ++++-----
 tensorflow/lite/kernels/div.cc                |  6 ++---
 tensorflow/lite/kernels/elementwise.cc        |  7 ++----
 tensorflow/lite/kernels/kernel_util.cc        |  4 ++--
 tensorflow/lite/kernels/lstm.cc               |  3 +--
 tensorflow/lite/kernels/mul.cc                |  5 ++--
 tensorflow/lite/kernels/squared_difference.cc | 12 ++++------
 tensorflow/lite/kernels/sub.cc                | 13 ++++------
 tensorflow/lite/kernels/svdf.cc               | 14 +++++------
 11 files changed, 45 insertions(+), 67 deletions(-)

diff --git a/tensorflow/lite/core/c/common.h b/tensorflow/lite/core/c/common.h
index ffc910d986c..66c8bf6fad6 100644
--- a/tensorflow/lite/core/c/common.h
+++ b/tensorflow/lite/core/c/common.h
@@ -269,14 +269,11 @@ void TfLiteFloatArrayFree(TfLiteFloatArray* a);
 
 #define TF_LITE_ENSURE_NEAR(context, a, b, epsilon)                          \
   do {                                                                       \
-    const auto _val_a = (a);                                                 \
-    const auto _val_b = (b);                                                 \
-    const auto _eps = (epsilon);                                             \
-    const auto _delta = _val_a - _val_b;                                     \
-    if (_delta > _eps || _delta < -_eps) {                                   \
-      TF_LITE_KERNEL_LOG(                                                    \
-          (context), "%s:%d %s not near %s (%f != %f)", __FILE__, __LINE__,  \
-          #a, #b, static_cast<double>(_val_a), static_cast<double>(_val_b)); \
+    auto delta = ((a) > (b)) ? ((a) - (b)) : ((b) - (a));                    \
+    if (delta > epsilon) {                                                   \
+      TF_LITE_KERNEL_LOG((context), "%s:%d %s not near %s (%f != %f)",       \
+                         __FILE__, __LINE__, #a, #b, static_cast<double>(a), \
+                         static_cast<double>(b));                            \
       return kTfLiteError;                                                   \
     }                                                                        \
   } while (0)
diff --git a/tensorflow/lite/kernels/activations.cc b/tensorflow/lite/kernels/activations.cc
index b6d7c4cdbf8..20ffc244a31 100644
--- a/tensorflow/lite/kernels/activations.cc
+++ b/tensorflow/lite/kernels/activations.cc
@@ -231,8 +231,7 @@ TfLiteStatus ReluPrepare(TfLiteContext* context, TfLiteNode* node) {
 
   if (input->type == kTfLiteInt8 || input->type == kTfLiteUInt8 ||
       input->type == kTfLiteInt16) {
-    double real_multiplier = static_cast<double>(input->params.scale) /
-                             static_cast<double>(output->params.scale);
+    double real_multiplier = input->params.scale / output->params.scale;
     QuantizeMultiplier(real_multiplier, &data->output_multiplier,
                        &data->output_shift);
   }
@@ -316,13 +315,11 @@ TfLiteStatus LeakyReluPrepare(TfLiteContext* context, TfLiteNode* node) {
     const auto* params =
         reinterpret_cast<TfLiteLeakyReluParams*>(node->builtin_data);
 
-    double alpha_multiplier = static_cast<double>(input->params.scale) *
-                              static_cast<double>(params->alpha) /
-                              static_cast<double>(output->params.scale);
+    double alpha_multiplier =
+        input->params.scale * params->alpha / output->params.scale;
     QuantizeMultiplier(alpha_multiplier, &data->output_multiplier_alpha,
                        &data->output_shift_alpha);
-    double identity_multiplier = static_cast<double>(input->params.scale) /
-                                 static_cast<double>(output->params.scale);
+    double identity_multiplier = input->params.scale / output->params.scale;
     QuantizeMultiplier(identity_multiplier, &data->output_multiplier_identity,
                        &data->output_shift_identity);
   }
@@ -644,9 +641,8 @@ TfLiteStatus SoftmaxPrepare(TfLiteContext* context, TfLiteNode* node) {
     data->params.zero_point = output->params.zero_point;
     data->params.scale = output->params.scale;
 
-    const double input_scale_beta_rescale =
-        static_cast<double>(input->params.scale) *
-        static_cast<double>(params->beta) /
+    double input_scale_beta_rescale =
+        input->params.scale * params->beta /
         (10.0 / 65535.0);  // scale the input_diff such that [-65535, 0]
                            // correspond to [-10.0, 0.0]
     QuantizeMultiplier(input_scale_beta_rescale, &data->params.input_multiplier,
@@ -740,11 +736,9 @@ TfLiteStatus PreluPrepare(TfLiteContext* context, TfLiteNode* node) {
     // output real multiplier 1 is input_scale / output_scale;
     // for input_q - input_zp < 0:
     // output real multiplier 2 is input_scale  * alpha_scale/ output_scale.
-    double real_multiplier_1 = static_cast<double>(input->params.scale) /
-                               static_cast<double>(output->params.scale);
-    double real_multiplier_2 = static_cast<double>(input->params.scale) *
-                               static_cast<double>(alpha->params.scale) /
-                               static_cast<double>(output->params.scale);
+    double real_multiplier_1 = input->params.scale / output->params.scale;
+    double real_multiplier_2 =
+        input->params.scale * alpha->params.scale / output->params.scale;
     QuantizeMultiplier(real_multiplier_1, &data->output_multiplier_1,
                        &data->output_shift_1);
     QuantizeMultiplier(real_multiplier_2, &data->output_multiplier_2,
diff --git a/tensorflow/lite/kernels/add.cc b/tensorflow/lite/kernels/add.cc
index ee61fd93de3..9e5dfcd228c 100644
--- a/tensorflow/lite/kernels/add.cc
+++ b/tensorflow/lite/kernels/add.cc
@@ -181,15 +181,14 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     // therefore the addition will still fit in a 32 bit accumulator.
     data->left_shift = general_scale_int16 ? 15 : 20;
     const double twice_max_input_scale =
-        2 * static_cast<double>(
-                std::max(input1->params.scale, input2->params.scale));
+        2 * std::max(input1->params.scale, input2->params.scale);
     const double real_input1_multiplier =
-        static_cast<double>(input1->params.scale) / twice_max_input_scale;
+        input1->params.scale / twice_max_input_scale;
     const double real_input2_multiplier =
-        static_cast<double>(input2->params.scale) / twice_max_input_scale;
+        input2->params.scale / twice_max_input_scale;
     const double real_output_multiplier =
-        twice_max_input_scale / (static_cast<double>(1 << data->left_shift) *
-                                 static_cast<double>(output->params.scale));
+        twice_max_input_scale /
+        ((1 << data->left_shift) * output->params.scale);
 
     QuantizeMultiplierSmallerThanOneExp(
         real_input1_multiplier, &data->input1_multiplier, &data->input1_shift);
diff --git a/tensorflow/lite/kernels/div.cc b/tensorflow/lite/kernels/div.cc
index c5a59cd2494..d7cebed68f9 100644
--- a/tensorflow/lite/kernels/div.cc
+++ b/tensorflow/lite/kernels/div.cc
@@ -104,9 +104,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     TF_LITE_ENSURE_STATUS(CalculateActivationRangeQuantized(
         context, params->activation, output, &data->output_activation_min,
         &data->output_activation_max));
-    const double real_multiplier = static_cast<double>(input1->params.scale) /
-                                   (static_cast<double>(input2->params.scale) *
-                                    static_cast<double>(output->params.scale));
+    const double real_multiplier =
+        input1->params.scale / (input2->params.scale * output->params.scale);
     QuantizeMultiplier(real_multiplier, &data->output_multiplier,
                        &data->output_shift);
   }
@@ -280,6 +279,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_OK(context,
                     GetOutputSafe(context, node, kOutputTensor, &output));
 
+
   if (output->type == kTfLiteFloat32) {
     // Div by zero seems ok in this case, we don't do a check at this point.
     // However, unlike in TF where infinities are returned, here we return an
diff --git a/tensorflow/lite/kernels/elementwise.cc b/tensorflow/lite/kernels/elementwise.cc
index cead70011b2..0bff391ffb3 100644
--- a/tensorflow/lite/kernels/elementwise.cc
+++ b/tensorflow/lite/kernels/elementwise.cc
@@ -86,16 +86,13 @@ bool IsLogSupportedType(const TfLiteType type) {
 inline void SetAbsOutputMultiplier(const float input_scale,
                                    const float output_scale,
                                    int32_t* multiplier, int32_t* shift) {
-  QuantizeMultiplier(
-      static_cast<double>(input_scale) / static_cast<double>(output_scale),
-      multiplier, shift);
+  QuantizeMultiplier(input_scale / output_scale, multiplier, shift);
 }
 
 inline void SetRsqrtOutputMultiplier(const float input_scale,
                                      const float output_scale,
                                      int32_t* multiplier, int32_t* shift) {
-  const double scale = 1.0 / (std::sqrt(static_cast<double>(input_scale)) *
-                              static_cast<double>(output_scale));
+  const double scale = 1. / (std::sqrt(input_scale) * output_scale);
   QuantizeMultiplier(scale, multiplier, shift);
 }
 
diff --git a/tensorflow/lite/kernels/kernel_util.cc b/tensorflow/lite/kernels/kernel_util.cc
index 3344dc2b13d..36ad5178158 100644
--- a/tensorflow/lite/kernels/kernel_util.cc
+++ b/tensorflow/lite/kernels/kernel_util.cc
@@ -330,8 +330,8 @@ TfLiteStatus GetQuantizedConvolutionMultipler(TfLiteContext* context,
                                               const TfLiteTensor* filter,
                                               TfLiteTensor* output,
                                               double* multiplier) {
-  const double input_product_scale = static_cast<double>(input->params.scale) *
-                                     static_cast<double>(filter->params.scale);
+  const double input_product_scale =
+      static_cast<double>(input->params.scale * filter->params.scale);
   TF_LITE_ENSURE(context, input_product_scale >= 0);
   *multiplier = input_product_scale / static_cast<double>(output->params.scale);
 
diff --git a/tensorflow/lite/kernels/lstm.cc b/tensorflow/lite/kernels/lstm.cc
index 2c1114d48ce..10dde7f4987 100644
--- a/tensorflow/lite/kernels/lstm.cc
+++ b/tensorflow/lite/kernels/lstm.cc
@@ -2325,8 +2325,7 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
       return kTfLiteError;
     }
 
-    double real_accum_multiplier =
-        4096.0 * static_cast<double>(bias->params.scale);
+    double real_accum_multiplier = 4096 * bias->params.scale;
     int32 accum_multiplier;
     int accum_shift;
     tflite::QuantizeMultiplier(real_accum_multiplier, &accum_multiplier,
diff --git a/tensorflow/lite/kernels/mul.cc b/tensorflow/lite/kernels/mul.cc
index e2d516ac733..73ffd3488f9 100644
--- a/tensorflow/lite/kernels/mul.cc
+++ b/tensorflow/lite/kernels/mul.cc
@@ -122,9 +122,8 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     TF_LITE_ENSURE_STATUS(CalculateActivationRangeQuantized(
         context, params->activation, output, &data->output_activation_min,
         &data->output_activation_max));
-    double real_multiplier = static_cast<double>(input1->params.scale) *
-                             static_cast<double>(input2->params.scale) /
-                             static_cast<double>(output->params.scale);
+    double real_multiplier =
+        input1->params.scale * input2->params.scale / output->params.scale;
     QuantizeMultiplier(real_multiplier, &data->output_multiplier,
                        &data->output_shift);
   }
diff --git a/tensorflow/lite/kernels/squared_difference.cc b/tensorflow/lite/kernels/squared_difference.cc
index 21597d0f424..57bf92cc37c 100644
--- a/tensorflow/lite/kernels/squared_difference.cc
+++ b/tensorflow/lite/kernels/squared_difference.cc
@@ -106,18 +106,16 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     // shift to make integer for scales.
     data->arithmetic_params.left_shift = 7;
     const double twice_max_input_scale =
-        2 * static_cast<double>(std::max(input1_quantization_params.scale,
-                                         input2_quantization_params.scale));
+        2 * std::max(input1_quantization_params.scale,
+                     input2_quantization_params.scale);
     const double real_input1_multiplier =
-        static_cast<double>(input1_quantization_params.scale) /
-        twice_max_input_scale;
+        input1_quantization_params.scale / twice_max_input_scale;
     double real_input2_multiplier =
-        static_cast<double>(input2_quantization_params.scale) /
-        twice_max_input_scale;
+        input2_quantization_params.scale / twice_max_input_scale;
     const double real_output_multiplier =
         (twice_max_input_scale * twice_max_input_scale) /
         ((1 << data->arithmetic_params.left_shift * 2) *
-         static_cast<double>(output_quantization_params.scale));
+         output_quantization_params.scale);
     tflite::QuantizeMultiplierSmallerThanOneExp(
         real_input1_multiplier, &data->arithmetic_params.input1_multiplier,
         &data->arithmetic_params.input1_shift);
diff --git a/tensorflow/lite/kernels/sub.cc b/tensorflow/lite/kernels/sub.cc
index 7cb5508227a..0136e94ab21 100644
--- a/tensorflow/lite/kernels/sub.cc
+++ b/tensorflow/lite/kernels/sub.cc
@@ -136,18 +136,15 @@ TfLiteStatus PrepareGeneralSubOp(TfLiteContext* context,
   // 31, therefore the addition will still fit in a 32 bit accumulator.
   op_params->left_shift = output->type == kTfLiteInt16 ? 15 : 20;
   const double twice_max_input_scale =
-      2 * static_cast<double>(std::max(input1_quantization_params.scale,
-                                       input2_quantization_params.scale));
+      2 * std::max(input1_quantization_params.scale,
+                   input2_quantization_params.scale);
   const double real_input1_multiplier =
-      static_cast<double>(input1_quantization_params.scale) /
-      twice_max_input_scale;
+      input1_quantization_params.scale / twice_max_input_scale;
   const double real_input2_multiplier =
-      static_cast<double>(input2_quantization_params.scale) /
-      twice_max_input_scale;
+      input2_quantization_params.scale / twice_max_input_scale;
   const double real_output_multiplier =
       twice_max_input_scale /
-      ((1 << op_params->left_shift) *
-       static_cast<double>(output_quantization_params.scale));
+      ((1 << op_params->left_shift) * output_quantization_params.scale);
 
   tflite::QuantizeMultiplierSmallerThanOneExp(real_input1_multiplier,
                                               &op_params->input1_multiplier,
diff --git a/tensorflow/lite/kernels/svdf.cc b/tensorflow/lite/kernels/svdf.cc
index 5173322f260..74a38153631 100644
--- a/tensorflow/lite/kernels/svdf.cc
+++ b/tensorflow/lite/kernels/svdf.cc
@@ -275,14 +275,12 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     TF_LITE_ENSURE(context, output->quantization.type != kTfLiteNoQuantization);
     auto* output_params = reinterpret_cast<TfLiteAffineQuantization*>(
         output->quantization.params);
-    const double effective_scale_1 =
-        static_cast<double>(input_params->scale->data[0]) *
-        static_cast<double>(weights_feature_params->scale->data[0]) /
-        static_cast<double>(state_params->scale->data[0]);
-    const double effective_scale_2 =
-        static_cast<double>(state_params->scale->data[0]) *
-        static_cast<double>(weight_time_params->scale->data[0]) /
-        static_cast<double>(output_params->scale->data[0]);
+    const double effective_scale_1 = input_params->scale->data[0] *
+                                     weights_feature_params->scale->data[0] /
+                                     state_params->scale->data[0];
+    const double effective_scale_2 = state_params->scale->data[0] *
+                                     weight_time_params->scale->data[0] /
+                                     output_params->scale->data[0];
     QuantizeMultiplier(effective_scale_1, &op_data->effective_scale_1_a,
                        &op_data->effective_scale_1_b);
     QuantizeMultiplier(effective_scale_2, &op_data->effective_scale_2_a,

From 9a0f5f21d19b07b635423a65be8b084311eccbaa Mon Sep 17 00:00:00 2001
From: Hyeontaek Lim <hyeontaek@google.com>
Date: Tue, 13 Oct 2026 20:37:08 -0700
Subject: [PATCH 0737/1000] [PjRt-IFRT] Implement `PjRtClient::Load`

`xla::ifrt::Client::Load` API provides a capability to load compiled (and not loaded) `xla::ifrt::Executable`s with load options to obtain `xla::ifrt::LoadedExecutable`s.

Implement `xla::ifrt::PjRtClient::Load` to load compiled `xla::ifrt::PjRtExecutable` batches onto the underlying `xla::PjRtClient`, validate target devices against the compiled device assignment (or bind portable executables to an optional single target device at load time), and forward `loaded_host_callbacks` and `outputs_bundle_slice_sizes` (defaulting to the compile-time `outputs_bundle_slice_sizes` when unspecified in `xla::ifrt::LoadOptions`).

PiperOrigin-RevId: 999076724
---
 third_party/xla/xla/python/pjrt_ifrt/BUILD    |  13 +
 .../xla/xla/python/pjrt_ifrt/pjrt_client.cc   | 116 +++-
 .../xla/xla/python/pjrt_ifrt/pjrt_compiler.cc |  23 +-
 .../xla/python/pjrt_ifrt/pjrt_executable.cc   |  51 +-
 .../xla/python/pjrt_ifrt/pjrt_executable.h    |  28 +-
 .../pjrt_executable_impl_test_lib.cc          | 568 ++++++++++++++++++
 6 files changed, 765 insertions(+), 34 deletions(-)

diff --git a/third_party/xla/xla/python/pjrt_ifrt/BUILD b/third_party/xla/xla/python/pjrt_ifrt/BUILD
index f5270da9053..f5b3660b795 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/BUILD
+++ b/third_party/xla/xla/python/pjrt_ifrt/BUILD
@@ -333,15 +333,28 @@ cc_library(
     srcs = ["pjrt_executable_impl_test_lib.cc"],
     deps = [
         ":pjrt_ifrt",
+        ":xla_ifrt",
         "//xla:literal",
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",
+        "//xla/client:executable_build_options",
+        "//xla/pjrt:mlir_to_hlo",
         "//xla/pjrt:pjrt_executable",
+        "//xla/python/ifrt",
+        "//xla/python/ifrt:basic_device_list",
         "//xla/python/ifrt:test_util",
+        "//xla/python/ifrt/hlo:hlo_program",
+        "//xla/service:device_assignment",
+        "//xla/tsl/concurrency:future",
         "//xla/tsl/concurrency:ref_count",
         "//xla/tsl/platform:statusor",
         "//xla/tsl/platform:test",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/types:span",
+        "@llvm-project//mlir:IR",
     ],
     alwayslink = True,
 )
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_client.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_client.cc
index 97d28d14351..3c52a073bf5 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_client.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_client.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "xla/python/pjrt_ifrt/pjrt_client.h"
 
 #include <algorithm>
+#include <cstddef>
 #include <cstdint>
 #include <functional>
 #include <memory>
@@ -59,6 +60,7 @@ limitations under the License.
 #include "xla/pjrt/pjrt_common.h"
 #include "xla/pjrt/pjrt_compiler.h"
 #include "xla/pjrt/pjrt_device_description.h"
+#include "xla/pjrt/pjrt_executable.h"
 #include "xla/pjrt/pjrt_layout.h"
 #include "xla/python/ifrt/array.h"
 #include "xla/python/ifrt/array_spec.h"
@@ -73,6 +75,7 @@ limitations under the License.
 #include "xla/python/ifrt/device_list.h"
 #include "xla/python/ifrt/dtype.h"
 #include "xla/python/ifrt/executable.h"
+#include "xla/python/ifrt/host_callback.h"
 #include "xla/python/ifrt/index_domain.h"
 #include "xla/python/ifrt/layout.h"
 #include "xla/python/ifrt/memory.h"
@@ -89,13 +92,17 @@ limitations under the License.
 #include "xla/python/pjrt_ifrt/pjrt_buffer_hash_util.h"
 #include "xla/python/pjrt_ifrt/pjrt_device.h"
 #include "xla/python/pjrt_ifrt/pjrt_dtype.h"
+#include "xla/python/pjrt_ifrt/pjrt_executable.h"
 #include "xla/python/pjrt_ifrt/pjrt_layout.h"
 #include "xla/python/pjrt_ifrt/pjrt_memory.h"
 #include "xla/python/pjrt_ifrt/pjrt_remap.h"
 #include "xla/python/pjrt_ifrt/pjrt_topology.h"
 #include "xla/python/pjrt_ifrt/pjrt_tuple.h"
+#include "xla/python/pjrt_ifrt/xla_client.h"
+#include "xla/python/pjrt_ifrt/xla_compiler.h"
 #include "xla/python/pjrt_ifrt/xla_sharding.h"
 #include "xla/runtime/device_id.h"
+#include "xla/service/device_assignment.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/status_macros.h"
@@ -1961,7 +1968,114 @@ absl::StatusOr<BundleRef> PjRtClient::ConcatBundles(
 absl::StatusOr<std::vector<tsl::Future<LoadedExecutableRef>>> PjRtClient::Load(
     absl::Span<const ExecutableRef> executables,
     absl::Span<std::unique_ptr<LoadOptions>> options) {
-  return absl::UnimplementedError("Load is not implemented in PjRtClient.");
+  if (executables.size() != options.size()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("executables and options must have the same size, got ",
+                     executables.size(), " and ", options.size()));
+  }
+
+  std::vector<tsl::Future<LoadedExecutableRef>> loaded_executables;
+  loaded_executables.reserve(executables.size());
+
+  absl::flat_hash_map<xla::GlobalDeviceId, DeviceId>
+      pjrt_global_device_id_to_ifrt_device_id;
+  pjrt_global_device_id_to_ifrt_device_id.reserve(
+      ifrt_device_id_to_pjrt_global_device_id_.size());
+  for (const auto& [ifrt_device_id, pjrt_global_device_id] :
+       ifrt_device_id_to_pjrt_global_device_id_) {
+    pjrt_global_device_id_to_ifrt_device_id.insert(
+        {pjrt_global_device_id, ifrt_device_id});
+  }
+
+  for (size_t i = 0; i < executables.size(); ++i) {
+    auto* pjrt_executable =
+        dyn_cast_or_null<PjRtExecutable>(executables[i].get());
+    if (pjrt_executable == nullptr) {
+      return absl::InvalidArgumentError(
+          "PjRtClient::Load requires PjRtExecutable");
+    }
+    if (options[i] == nullptr) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("LoadOptions at index ", i, " must not be null"));
+    }
+    if (options[i]->devices != nullptr && options[i]->devices->empty()) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "LoadOptions::devices at index ", i, " must be non-empty"));
+    }
+
+    DeviceListRef device_list = std::move(options[i]->devices);
+    if (pjrt_executable->common_metadata().is_portable &&
+        device_list != nullptr && device_list->size() != 1) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Only single-shard portable execution is supported, got ",
+          device_list->size(), " devices at index ", i));
+    }
+    std::optional<std::vector<int>> outputs_bundle_slice_sizes =
+        std::move(options[i]->outputs_bundle_slice_sizes);
+    std::vector<tsl::RCReference<LoadedHostCallback>> loaded_host_callbacks;
+    if (auto* xla_load_options = dyn_cast<XlaLoadOptions>(options[i].get());
+        xla_load_options != nullptr) {
+      loaded_host_callbacks =
+          std::move(xla_load_options->loaded_host_callbacks);
+    }
+
+    ABSL_ASSIGN_OR_RETURN(
+        std::unique_ptr<xla::PjRtLoadedExecutable> pjrt_loaded_executable,
+        pjrt_client_->Load(pjrt_executable->shared_ptr_pjrt_executable(),
+                           xla::LoadOptions()));
+
+    if (!pjrt_executable->common_metadata().is_portable) {
+      xla::DeviceAssignment device_assignment =
+          pjrt_loaded_executable->device_assignment();
+      absl::Status result;
+      device_assignment.Each(
+          [&](int64_t replica, int64_t computation, int64_t* device_id) {
+            if (!result.ok()) {
+              return;
+            }
+            auto it = pjrt_global_device_id_to_ifrt_device_id.find(
+                xla::GlobalDeviceId(*device_id));
+            if (it != pjrt_global_device_id_to_ifrt_device_id.end()) {
+              *device_id = it->second.value();
+            } else {
+              result.Update(absl::InvalidArgumentError(absl::StrCat(
+                  "Unknown IFRT device ID for PjRt global device ID ",
+                  *device_id)));
+            }
+          });
+      ABSL_RETURN_IF_ERROR(result);
+      ABSL_ASSIGN_OR_RETURN(
+          DeviceListRef compiled_device_list,
+          GetDeviceListFromDeviceAssignment(this, device_assignment));
+      if (device_list == nullptr) {
+        device_list = std::move(compiled_device_list);
+      } else if (*device_list != *compiled_device_list) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "LoadOptions::devices at index ", i,
+            " does not match the executable's compiled device assignment: ",
+            *device_list, " vs. ", *compiled_device_list));
+      }
+    }
+
+    PjRtExecutable::CommonMetadata common_metadata =
+        pjrt_executable->common_metadata();
+    if (outputs_bundle_slice_sizes.has_value()) {
+      common_metadata.outputs_bundle_slice_sizes =
+          std::move(outputs_bundle_slice_sizes);
+    }
+
+    ABSL_ASSIGN_OR_RETURN(LoadedExecutableRef loaded_executable,
+                     PjRtLoadedExecutable::Create(
+                         this,
+                         std::shared_ptr<xla::PjRtLoadedExecutable>(
+                             std::move(pjrt_loaded_executable)),
+                         std::move(loaded_host_callbacks),
+                         std::move(device_list), std::move(common_metadata)));
+    loaded_executables.push_back(
+        tsl::Future<LoadedExecutableRef>(std::move(loaded_executable)));
+  }
+
+  return loaded_executables;
 }
 
 absl::StatusOr<std::shared_ptr<Topology>> PjRtClient::GetTopologyForDevices(
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_compiler.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_compiler.cc
index d47cbdef862..4dc613d73b5 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_compiler.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_compiler.cc
@@ -124,6 +124,11 @@ tsl::Future<ExecutableRef> PjRtCompiler::Compile(
       cast<HloProgram>(std::move(program));
   ABSL_ASSIGN_OR_RETURN(auto xla_compile_options,
                    GetXlaCompileOptions(std::move(options)));
+  if (!xla_compile_options->loaded_host_callbacks.empty()) {
+    return absl::InvalidArgumentError(
+        "loaded_host_callbacks cannot be passed to Compiler::Compile; pass "
+        "them via XlaLoadOptions to Client::Load instead.");
+  }
   const auto* pjrt_topology = dyn_cast<PjRtTopology>(&topology);
   if (pjrt_topology == nullptr) {
     return absl::InvalidArgumentError("PjRtCompiler requires a PjRtTopology");
@@ -144,7 +149,8 @@ tsl::Future<ExecutableRef> PjRtCompiler::Compile(
   return PjRtExecutable::Create(
       std::move(*xla_program).ToMaybeOwningMlirModule(),
       std::move(xla_compile_options->compile_options),
-      *pjrt_topology->description(), compile_client);
+      *pjrt_topology->description(), compile_client,
+      std::move(xla_compile_options->outputs_bundle_slice_sizes));
 }
 
 tsl::Future<LoadedExecutableRef> PjRtCompiler::DeserializeLoadedExecutable(
@@ -178,12 +184,15 @@ tsl::Future<LoadedExecutableRef> PjRtCompiler::DeserializeLoadedExecutable(
   // through from Australis and are always present in the
   // DeserializeExecutableOptions.
   DeviceListRef device_list;
-  if (xla_deserialize_options->devices.has_value()) {
-    device_list = std::move(xla_deserialize_options->devices.value());
-  } else {
-    ABSL_ASSIGN_OR_RETURN(device_list,
-                     GetDeviceListFromDeviceAssignment(
-                         client_, pjrt_loaded_executable->device_assignment()));
+  if (!common_metadata_and_serialized_pjrt_executable.first.is_portable) {
+    if (xla_deserialize_options->devices.has_value()) {
+      device_list = *std::move(xla_deserialize_options->devices);
+    } else {
+      ABSL_ASSIGN_OR_RETURN(
+          device_list,
+          GetDeviceListFromDeviceAssignment(
+              client_, pjrt_loaded_executable->device_assignment()));
+    }
   }
   return PjRtLoadedExecutable::Create(
       client_,
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc
index 51bd21e0048..15e2d11e67f 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.cc
@@ -443,7 +443,8 @@ char PjRtLoadedExecutable::ID = 0;
 absl::StatusOr<ExecutableRef> PjRtExecutable::Create(
     xla::MaybeOwningMlirModule module, xla::CompileOptions compile_options,
     const xla::PjRtTopologyDescription& topology,
-    xla::PjRtClient* compile_client) {
+    xla::PjRtClient* compile_client,
+    std::optional<std::vector<int>> outputs_bundle_slice_sizes) {
   const bool is_portable = compile_options.compile_portable_executable;
 
   // We have to do process the MLIR before the compile call, since the latter
@@ -486,6 +487,8 @@ absl::StatusOr<ExecutableRef> PjRtExecutable::Create(
   common_metadata.output_hlo_shardings = std::move(output_hlo_shardings);
   common_metadata.output_memory_kinds = std::move(output_memory_kinds);
   common_metadata.output_layouts = std::move(output_layouts);
+  common_metadata.outputs_bundle_slice_sizes =
+      std::move(outputs_bundle_slice_sizes);
 
   return ExecutableRef(new PjRtExecutable(std::move(pjrt_executable),
                                           std::move(common_metadata)));
@@ -843,9 +846,6 @@ absl::StatusOr<LoadedExecutableRef> PjRtLoadedExecutable::Create(
       std::vector<MemoryKind> output_memory_kinds,
       GetMemoryKinds(pjrt_loaded_executable->GetOutputMemoryKinds(),
                      output_dtypes));
-  std::vector<ShardingRef> output_shardings =
-      MakeShardings(output_shapes, output_hlo_shardings, output_memory_kinds,
-                    executable_devices);
   ABSL_ASSIGN_OR_RETURN(
       std::optional<std::vector<std::shared_ptr<const xla::PjRtLayout>>>
           output_layouts,
@@ -864,7 +864,8 @@ absl::StatusOr<LoadedExecutableRef> PjRtLoadedExecutable::Create(
       std::move(outputs_bundle_slice_sizes);
 
   return LoadedExecutableRef(new PjRtLoadedExecutable(
-      client, std::move(pjrt_loaded_executable), std::move(executable_devices),
+      client, std::move(pjrt_loaded_executable),
+      is_portable ? DeviceListRef() : std::move(executable_devices),
       std::move(loaded_host_callbacks), std::move(common_metadata)));
 }
 
@@ -877,7 +878,9 @@ PjRtLoadedExecutable::PjRtLoadedExecutable(
     : client_(client),
       pjrt_loaded_executable_(std::move(pjrt_loaded_executable)),
       devices_(std::move(devices)),
-      addressable_devices_(devices_->AddressableDeviceList()->devices()),
+      addressable_devices_(devices_ != nullptr
+                               ? devices_->AddressableDeviceList()->devices()
+                               : absl::Span<Device* const>()),
       all_loaded_host_callbacks_(
           std::make_shared<std::vector<tsl::RCReference<LoadedHostCallback>>>(
               std::move(all_loaded_host_callbacks))),
@@ -887,9 +890,13 @@ PjRtLoadedExecutable::PjRtLoadedExecutable(
           GatherHloOutputCallbacks(*all_loaded_host_callbacks_)),
       common_metadata_(std::move(common_metadata)),
       user_context_(UserContextScope::current()) {
-  output_shardings_ = MakeShardings(
-      common_metadata_.output_shapes, common_metadata_.output_hlo_shardings,
-      common_metadata_.output_memory_kinds, devices_);
+  if (devices_ != nullptr) {
+    output_shardings_ = MakeShardings(
+        common_metadata_.output_shapes,
+        common_metadata_.is_portable ? std::nullopt
+                                     : common_metadata_.output_hlo_shardings,
+        common_metadata_.output_memory_kinds, devices_);
+  }
   pjrt_output_memory_kinds_.emplace_back();
   for (const MemoryKind& memory_kind : common_metadata_.output_memory_kinds) {
     pjrt_output_memory_kinds_.back().push_back(memory_kind.value());
@@ -911,19 +918,29 @@ PjRtLoadedExecutable::Execute(absl::Span<ArrayRef> args,
   std::vector<std::unique_ptr<PjRtBuffer>> owned_buffers;
 
   int num_computations;
-  const bool portable_execution = devices.has_value();
+  const bool portable_execution = common_metadata_.is_portable;
   PjRtCompatibleDevice* portable_execution_device = nullptr;
+  if (devices_ != nullptr && devices.has_value() &&
+      (*devices == nullptr || **devices != *devices_)) {
+    return InvalidArgument(
+        "Cannot execute with devices that differ from the bound devices");
+  }
   if (portable_execution) {
-    if ((*devices)->size() != 1) {
+    const DeviceListRef& execution_devices =
+        devices.has_value() ? *devices : devices_;
+    if (execution_devices == nullptr || execution_devices->empty()) {
+      return InvalidArgument("No devices provided for portable executable");
+    }
+    if (execution_devices->size() != 1) {
       return InvalidArgument(
           "Only single-shard portable execution is supported");
     }
-    num_computations = 1;
     portable_execution_device =
-        static_cast<PjRtDevice*>((*devices)->devices().front());
+        static_cast<PjRtDevice*>(execution_devices->devices().front());
+    num_computations = 1;
   } else {
-    if (devices_->devices().empty()) {
-      return InvalidArgument("No devices provided for portable executable");
+    if (devices_ == nullptr || devices_->empty()) {
+      return InvalidArgument("No devices provided for executable");
     }
     num_computations = addressable_devices_.size();
   }
@@ -1144,7 +1161,7 @@ PjRtLoadedExecutable::Execute(absl::Span<ArrayRef> args,
   for (int i = 0; i < num_outputs; ++i) {
     PjRtArray::PjRtBuffers buffers;
     buffers.reserve(num_computations);
-    const MemoryKind dst_memory_kind = output_shardings_[i]->memory_kind();
+    const MemoryKind dst_memory_kind = common_metadata_.output_memory_kinds[i];
 
     for (int j = 0; j < num_computations; ++j) {
       if (j > 0) {
@@ -1161,7 +1178,7 @@ PjRtLoadedExecutable::Execute(absl::Span<ArrayRef> args,
           std::shared_ptr<PjRtBuffer>(pjrt_outputs[j][i].release()));
     }
     std::optional<ShardingRef> sharding;
-    if (portable_execution) {
+    if (output_shardings_.empty()) {
       if (auto it = single_device_shardings.find(dst_memory_kind);
           it == single_device_shardings.end()) {
         sharding = single_device_shardings
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.h b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.h
index 0761ae98dbd..46b87d70e0f 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.h
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable.h
@@ -70,6 +70,7 @@ class PjRtCompatibleExecutable
   // APIs that allow direct access to `xla::PjRtExecutable` for PjRt-only
   // operations.
   virtual xla::PjRtExecutable* pjrt_executable() = 0;
+  virtual std::shared_ptr<xla::PjRtExecutable> shared_ptr_pjrt_executable() = 0;
 
   static char ID;  // NOLINT
 };
@@ -101,7 +102,8 @@ class PjRtExecutable final
   static absl::StatusOr<ExecutableRef> Create(
       xla::MaybeOwningMlirModule module, xla::CompileOptions compile_options,
       const xla::PjRtTopologyDescription& topology,
-      xla::PjRtClient* compile_client = nullptr);
+      xla::PjRtClient* compile_client,
+      std::optional<std::vector<int>> outputs_bundle_slice_sizes);
 
   // PjRtCompatibleExecutable implementation.
 
@@ -109,6 +111,10 @@ class PjRtExecutable final
     DCHECK(this);
     return pjrt_executable_.get();
   }
+  std::shared_ptr<xla::PjRtExecutable> shared_ptr_pjrt_executable() override {
+    DCHECK(this);
+    return pjrt_executable_;
+  }
 
   // Executable implementation.
 
@@ -216,6 +222,11 @@ class PjRtExecutable final
             xla_deserialize_executable_options);
   };
 
+  const CommonMetadata& common_metadata() const {
+    DCHECK(this);
+    return common_metadata_;
+  }
+
  protected:
   PjRtExecutable(std::shared_ptr<xla::PjRtExecutable> pjrt_executable,
                  CommonMetadata common_metadata);
@@ -355,8 +366,8 @@ class PjRtLoadedExecutable final
       absl::Span<BundleRef> args, const ExecuteOptions& options) override;
 
   std::optional<DeviceListRef> devices() const override {
-    if (pjrt_loaded_executable_->addressable_devices().empty()) {
-      // Portable executable.
+    DCHECK(this);
+    if (devices_ == nullptr) {
       return std::nullopt;
     }
     return devices_;
@@ -382,11 +393,11 @@ class PjRtLoadedExecutable final
 
   PjRtClient* client_;
   std::shared_ptr<xla::PjRtLoadedExecutable> pjrt_loaded_executable_;
-  // Devices that `pjrt_loaded_executable_` runs on. Empty if the executable is
-  // portable.
+  // Devices that `pjrt_loaded_executable_` runs on. `nullptr` if the
+  // executable is portable and not bound to a device at load time.
   DeviceListRef devices_;
   // Addressable devices. The underlying device list is owned by
-  // `devices_->AddressableDeviceList()`.
+  // `devices_->AddressableDeviceList()`. Empty if `devices_` is `nullptr`.
   absl::Span<Device* const> addressable_devices_;
   std::shared_ptr<std::vector<tsl::RCReference<LoadedHostCallback>>>
       all_loaded_host_callbacks_;
@@ -394,9 +405,8 @@ class PjRtLoadedExecutable final
   std::vector<PjRtHloOutputLoadedHostCallback*> host_hlo_output_callbacks_;
 
   PjRtExecutable::CommonMetadata common_metadata_;
-  // If the executable is portable, shardings in `output_shardings_` will use an
-  // arbitrary addressable device, and will be overridden by a
-  // `SingleDeviceSharding` generated on the fly at execution time.
+  // Output shardings. Empty if `devices_` is `nullptr`; unbound portable
+  // execution constructs `SingleDeviceSharding` on the fly at execution time.
   std::vector<ShardingRef> output_shardings_;
   // PjRt-style memory kinds. Used only for `GetOutputMemoryKinds()`.
   std::vector<std::vector<absl::string_view>> pjrt_output_memory_kinds_;
diff --git a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable_impl_test_lib.cc b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable_impl_test_lib.cc
index 9124e665d95..f3f9fe64c6f 100644
--- a/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable_impl_test_lib.cc
+++ b/third_party/xla/xla/python/pjrt_ifrt/pjrt_executable_impl_test_lib.cc
@@ -15,14 +15,47 @@ limitations under the License.
 
 #include <cstdint>
 #include <memory>
+#include <optional>
+#include <tuple>
 #include <utility>
+#include <vector>
 
+#include "absl/algorithm/container.h"
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "absl/strings/string_view.h"
 #include "absl/types/span.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/OwningOpRef.h"
+#include "xla/client/executable_build_options.h"
 #include "xla/literal.h"
+#include "xla/pjrt/mlir_to_hlo.h"
 #include "xla/pjrt/pjrt_executable.h"
+#include "xla/python/ifrt/array.h"
+#include "xla/python/ifrt/basic_device_list.h"
+#include "xla/python/ifrt/bundle.h"
+#include "xla/python/ifrt/client.h"
+#include "xla/python/ifrt/compiler.h"
+#include "xla/python/ifrt/device.h"
+#include "xla/python/ifrt/device_list.h"
+#include "xla/python/ifrt/dtype.h"
+#include "xla/python/ifrt/executable.h"
+#include "xla/python/ifrt/hlo/hlo_program.h"
+#include "xla/python/ifrt/host_callback.h"
+#include "xla/python/ifrt/memory.h"
+#include "xla/python/ifrt/rtti.h"
+#include "xla/python/ifrt/shape.h"
+#include "xla/python/ifrt/sharding.h"
 #include "xla/python/ifrt/test_util.h"
+#include "xla/python/ifrt/topology.h"
+#include "xla/python/ifrt/value.h"
 #include "xla/python/pjrt_ifrt/pjrt_host_callback.h"
+#include "xla/python/pjrt_ifrt/xla_client.h"
+#include "xla/python/pjrt_ifrt/xla_compiler.h"
+#include "xla/service/device_assignment.h"
 #include "xla/shape_util.h"
+#include "xla/tsl/concurrency/future.h"
 #include "xla/tsl/concurrency/ref_count.h"
 #include "xla/tsl/platform/statusor.h"
 #include "xla/tsl/platform/test.h"
@@ -32,6 +65,32 @@ namespace xla {
 namespace ifrt {
 namespace {
 
+using ::absl_testing::IsOk;
+using ::absl_testing::StatusIs;
+using ::testing::ElementsAre;
+using ::testing::ElementsAreArray;
+using ::testing::SizeIs;
+
+constexpr absl::string_view kModuleAddOne = R"mlir(
+module @add_one {
+  func.func @main(%arg0: tensor<2x3xf32>) -> tensor<2x3xf32> {
+    %0 = stablehlo.constant dense<1.000000e+00> : tensor<f32>
+    %1 = "stablehlo.broadcast_in_dim"(%0) {broadcast_dimensions = array<i64>} : (tensor<f32>) -> tensor<2x3xf32>
+    %2 = stablehlo.add %arg0, %1 : tensor<2x3xf32>
+    return %2 : tensor<2x3xf32>
+  }
+})mlir";
+
+constexpr absl::string_view kModuleTwoOutputs = R"mlir(
+module @two_outputs {
+  func.func @main(%arg0: tensor<2x3xf32>) -> (tensor<2x3xf32>, tensor<2x3xf32>) {
+    %0 = stablehlo.constant dense<1.000000e+00> : tensor<f32>
+    %1 = "stablehlo.broadcast_in_dim"(%0) {broadcast_dimensions = array<i64>} : (tensor<f32>) -> tensor<2x3xf32>
+    %2 = stablehlo.add %arg0, %1 : tensor<2x3xf32>
+    return %arg0, %2 : tensor<2x3xf32>, tensor<2x3xf32>
+  }
+})mlir";
+
 TEST(PjRtExecutableImplTest, HloOutputCallbackWrapper) {
   TF_ASSERT_OK_AND_ASSIGN(auto client, test_util::GetClient());
 
@@ -66,6 +125,515 @@ TEST(PjRtExecutableImplTest, HloOutputCallbackWrapper) {
   EXPECT_TRUE(invoked);
 }
 
+TEST(PjRtExecutableImplTest, CompileAndLoadSeparately) {
+  ASSERT_OK_AND_ASSIGN(std::shared_ptr<Client> client, test_util::GetClient());
+  ASSERT_GE(client->addressable_devices().size(), 2);
+  Compiler* compiler = client->GetDefaultCompiler();
+
+  mlir::MLIRContext context;
+  ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef<mlir::ModuleOp> module,
+                       xla::ParseMlirModuleString(kModuleAddOne, context));
+
+  Device* device0 = client->addressable_devices().at(0);
+  Device* device1 = client->addressable_devices().at(1);
+  ASSERT_OK_AND_ASSIGN(DeviceListRef device_list,
+                       client->MakeDeviceList({device0, device1}));
+  ASSERT_OK_AND_ASSIGN(DeviceListRef single_device_list,
+                       client->MakeDeviceList({device0}));
+  ASSERT_OK_AND_ASSIGN(std::shared_ptr<Topology> topology,
+                       client->GetTopologyForDevices(device_list));
+
+  xla::CompileOptions compile_options;
+  ExecutableBuildOptions& build_options =
+      compile_options.executable_build_options;
+  build_options.set_num_replicas(2);
+  build_options.set_num_partitions(1);
+  build_options.set_use_spmd_partitioning(true);
+  DeviceAssignment device_assignment(2, 1);
+  device_assignment(0, 0) = device0->Id().value();
+  device_assignment(1, 0) = device1->Id().value();
+  build_options.set_device_assignment(device_assignment);
+
+  auto xla_compile_options =
+      std::make_unique<XlaCompileOptions>(compile_options, device_list);
+  ASSERT_OK_AND_ASSIGN(ExecutableRef executable,
+                       compiler
+                           ->Compile(std::make_unique<HloProgram>(*module),
+                                     *topology, std::move(xla_compile_options))
+                           .Await());
+  ASSERT_NE(executable, nullptr);
+  EXPECT_EQ(executable->name(), "add_one");
+  EXPECT_EQ(executable->num_devices(), 2);
+
+  std::vector<ExecutableRef> executables = {executable, executable};
+  std::vector<std::unique_ptr<LoadOptions>> load_options;
+  load_options.push_back(std::make_unique<XlaLoadOptions>(
+      device_list, /*loaded_host_callbacks=*/
+      std::vector<tsl::RCReference<LoadedHostCallback>>(),
+      /*outputs_bundle_slice_sizes=*/std::nullopt));
+  load_options.push_back(std::make_unique<XlaLoadOptions>(
+      /*devices=*/DeviceListRef(), /*loaded_host_callbacks=*/
+      std::vector<tsl::RCReference<LoadedHostCallback>>(),
+      /*outputs_bundle_slice_sizes=*/std::nullopt));
+  ASSERT_OK_AND_ASSIGN(
+      std::vector<tsl::Future<LoadedExecutableRef>> loaded_executable_futures,
+      client->Load(executables, absl::MakeSpan(load_options)));
+  ASSERT_THAT(loaded_executable_futures, SizeIs(2));
+  ASSERT_OK_AND_ASSIGN(LoadedExecutableRef loaded_executable,
+                       loaded_executable_futures[0].Await());
+  ASSERT_NE(loaded_executable, nullptr);
+  EXPECT_EQ(loaded_executable->name(), "add_one");
+  ASSERT_TRUE(loaded_executable->devices().has_value());
+  EXPECT_THAT((*loaded_executable->devices())->devices(),
+              ElementsAre(device0, device1));
+
+  ASSERT_OK_AND_ASSIGN(LoadedExecutableRef loaded_executable_default_devices,
+                       loaded_executable_futures[1].Await());
+  ASSERT_NE(loaded_executable_default_devices, nullptr);
+  ASSERT_TRUE(loaded_executable_default_devices->devices().has_value());
+  EXPECT_THAT((*loaded_executable_default_devices->devices())->devices(),
+              ElementsAre(device0, device1));
+
+  DType dtype(DType::kF32);
+  Shape shape({2, 3});
+  std::vector<float> data(6);
+  absl::c_iota(data, 0);
+  std::vector<float> expected_out_data(6);
+  absl::c_iota(expected_out_data, 1);
+  ShardingRef sharding =
+      ConcreteEvenSharding::Create(device_list, MemoryKind(), shape, shape,
+                                   /*is_fully_replicated=*/true);
+  ASSERT_OK_AND_ASSIGN(
+      ArrayRef input,
+      client->MakeArrayFromHostBuffer(
+          data.data(), dtype, shape, /*byte_strides=*/std::nullopt, sharding,
+          /*layout=*/nullptr,
+          Client::HostBufferSemantics::kImmutableOnlyDuringCall,
+          /*on_done_with_host_buffer=*/{}));
+
+  LoadedExecutable::ExecuteOptions execute_options;
+  execute_options.fill_status = true;
+  for (const std::optional<DeviceListRef>& execute_devices :
+       {std::optional<DeviceListRef>(),
+        std::optional<DeviceListRef>(device_list)}) {
+    ASSERT_OK_AND_ASSIGN(
+        LoadedExecutable::ExecuteResult result,
+        loaded_executable->Execute(absl::MakeSpan(&input, 1), execute_options,
+                                   execute_devices));
+    EXPECT_THAT(result.status.Await(), IsOk());
+    ASSERT_THAT(result.outputs, SizeIs(1));
+
+    ASSERT_OK_AND_ASSIGN(
+        std::vector<ArrayRef> output_shards,
+        result.outputs.front()->DisassembleIntoSingleDeviceArrays(
+            ArrayCopySemantics::kDonateInput,
+            SingleDeviceShardSemantics::kAddressableShards));
+    ASSERT_THAT(output_shards, SizeIs(2));
+    for (const ArrayRef& output_shard : output_shards) {
+      std::vector<float> out_data(6);
+      EXPECT_THAT(output_shard
+                      ->CopyToHostBuffer(out_data.data(),
+                                         /*byte_strides=*/std::nullopt,
+                                         ArrayCopySemantics::kAlwaysCopy)
+                      .Await(),
+                  IsOk());
+      EXPECT_THAT(out_data, ElementsAreArray(expected_out_data));
+    }
+  }
+
+  EXPECT_THAT(loaded_executable->Execute(absl::MakeSpan(&input, 1),
+                                         execute_options, single_device_list),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+TEST(PjRtExecutableImplTest, CompileAndLoadWithOutputsBundleSliceSizes) {
+  ASSERT_OK_AND_ASSIGN(std::shared_ptr<Client> client, test_util::GetClient());
+  Compiler* compiler = client->GetDefaultCompiler();
+
+  mlir::MLIRContext context;
+  ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef<mlir::ModuleOp> module,
+                       xla::ParseMlirModuleString(kModuleTwoOutputs, context));
+
+  Device* device = client->addressable_devices().at(0);
+  ASSERT_OK_AND_ASSIGN(DeviceListRef device_list,
+                       client->MakeDeviceList({device}));
+  ASSERT_OK_AND_ASSIGN(std::shared_ptr<Topology> topology,
+                       client->GetTopologyForDevices(device_list));
+
+  xla::CompileOptions compile_options;
+  ExecutableBuildOptions& build_options =
+      compile_options.executable_build_options;
+  build_options.set_device_ordinal(device->Id().value());
+  build_options.set_num_replicas(1);
+  build_options.set_num_partitions(1);
+  build_options.set_use_spmd_partitioning(true);
+  DeviceAssignment device_assignment(1, 1);
+  device_assignment(0, 0) = device->Id().value();
+  build_options.set_device_assignment(device_assignment);
+
+  auto xla_compile_options =
+      std::make_unique<XlaCompileOptions>(compile_options, device_list);
+  xla_compile_options->outputs_bundle_slice_sizes = std::vector<int>{1, 1};
+  ASSERT_OK_AND_ASSIGN(ExecutableRef executable,
+                       compiler
+                           ->Compile(std::make_unique<HloProgram>(*module),
+                                     *topology, std::move(xla_compile_options))
+                           .Await());
+
+  std::vector<ExecutableRef> executables = {executable, executable};
+  std::vector<std::unique_ptr<LoadOptions>> load_options;
+  load_options.push_back(std::make_unique<XlaLoadOptions>(
+      device_list, /*loaded_host_callbacks=*/
+      std::vector<tsl::RCReference<LoadedHostCallback>>(),
+      /*outputs_bundle_slice_sizes=*/std::nullopt));
+  load_options.push_back(std::make_unique<XlaLoadOptions>(
+      device_list, /*loaded_host_callbacks=*/
+      std::vector<tsl::RCReference<LoadedHostCallback>>(),
+      /*outputs_bundle_slice_sizes=*/std::vector<int>{2}));
+  ASSERT_OK_AND_ASSIGN(
+      std::vector<tsl::Future<LoadedExecutableRef>> loaded_executable_futures,
+      client->Load(executables, absl::MakeSpan(load_options)));
+  ASSERT_THAT(loaded_executable_futures, SizeIs(2));
+
+  ASSERT_OK_AND_ASSIGN(LoadedExecutableRef loaded_preserved,
+                       loaded_executable_futures[0].Await());
+  ASSERT_OK_AND_ASSIGN(LoadedExecutableRef loaded_overridden,
+                       loaded_executable_futures[1].Await());
+
+  DType dtype(DType::kF32);
+  Shape shape({2, 3});
+  std::vector<float> data(6);
+  absl::c_iota(data, 0);
+  std::vector<float> expected_out_data(6);
+  absl::c_iota(expected_out_data, 1);
+  ShardingRef sharding = SingleDeviceSharding::Create(device, MemoryKind());
+  ASSERT_OK_AND_ASSIGN(
+      ArrayRef input,
+      client->MakeArrayFromHostBuffer(
+          data.data(), dtype, shape, /*byte_strides=*/std::nullopt, sharding,
+          /*layout=*/nullptr,
+          Client::HostBufferSemantics::kImmutableOnlyDuringCall,
+          /*on_done_with_host_buffer=*/{}));
+  std::vector<ValueRef> input_values = {input};
+  ASSERT_OK_AND_ASSIGN(BundleRef input_bundle,
+                       client->Bundle(absl::MakeSpan(input_values),
+                                      ArrayCopySemantics::kReuseInput));
+
+  LoadedExecutable::ExecuteOptions execute_options;
+  execute_options.fill_status = true;
+
+  // Verify compile-time outputs_bundle_slice_sizes ({1, 1}) is preserved when
+  // LoadOptions::outputs_bundle_slice_sizes is std::nullopt.
+  ASSERT_OK_AND_ASSIGN(LoadedExecutable::ExecuteBundleResult result_preserved,
+                       loaded_preserved->ExecuteBundle(
+                           absl::MakeSpan(&input_bundle, 1), execute_options));
+  EXPECT_THAT(result_preserved.status.Await(), IsOk());
+  ASSERT_THAT(result_preserved.outputs, SizeIs(2));
+  EXPECT_EQ(result_preserved.outputs[0]->num_values(), 1);
+  EXPECT_EQ(result_preserved.outputs[1]->num_values(), 1);
+
+  for (int i = 0; i < 2; ++i) {
+    ASSERT_OK_AND_ASSIGN(std::vector<ValueRef> values,
+                         result_preserved.outputs[i]->GetValues(
+                             ArrayCopySemantics::kReuseInput));
+    ASSERT_THAT(values, SizeIs(1));
+    auto* out_array = xla::ifrt::dyn_cast<Array>(values.front().get());
+    ASSERT_NE(out_array, nullptr);
+    std::vector<float> out_data(6);
+    EXPECT_THAT(out_array
+                    ->CopyToHostBuffer(out_data.data(),
+                                       /*byte_strides=*/std::nullopt,
+                                       ArrayCopySemantics::kAlwaysCopy)
+                    .Await(),
+                IsOk());
+    EXPECT_THAT(out_data, ElementsAreArray(i == 0 ? data : expected_out_data));
+  }
+
+  // Verify LoadOptions::outputs_bundle_slice_sizes ({2}) overrides compile-time
+  // outputs_bundle_slice_sizes.
+  ASSERT_OK_AND_ASSIGN(LoadedExecutable::ExecuteBundleResult result_overridden,
+                       loaded_overridden->ExecuteBundle(
+                           absl::MakeSpan(&input_bundle, 1), execute_options));
+  EXPECT_THAT(result_overridden.status.Await(), IsOk());
+  ASSERT_THAT(result_overridden.outputs, SizeIs(1));
+  EXPECT_EQ(result_overridden.outputs.front()->num_values(), 2);
+
+  ASSERT_OK_AND_ASSIGN(std::vector<ValueRef> overridden_values,
+                       result_overridden.outputs.front()->GetValues(
+                           ArrayCopySemantics::kReuseInput));
+  ASSERT_THAT(overridden_values, SizeIs(2));
+  for (int i = 0; i < 2; ++i) {
+    auto* out_array = xla::ifrt::dyn_cast<Array>(overridden_values[i].get());
+    ASSERT_NE(out_array, nullptr);
+    std::vector<float> out_data(6);
+    EXPECT_THAT(out_array
+                    ->CopyToHostBuffer(out_data.data(),
+                                       /*byte_strides=*/std::nullopt,
+                                       ArrayCopySemantics::kAlwaysCopy)
+                    .Await(),
+                IsOk());
+    EXPECT_THAT(out_data, ElementsAreArray(i == 0 ? data : expected_out_data));
+  }
+}
+
+TEST(PjRtExecutableImplTest, LoadBatchWithDuplicateExecutables) {
+  ASSERT_OK_AND_ASSIGN(std::shared_ptr<Client> client, test_util::GetClient());
+  ASSERT_GE(client->addressable_devices().size(), 2);
+  Compiler* compiler = client->GetDefaultCompiler();
+
+  mlir::MLIRContext context;
+  ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef<mlir::ModuleOp> module,
+                       xla::ParseMlirModuleString(kModuleAddOne, context));
+
+  Device* device0 = client->addressable_devices().at(0);
+  Device* device1 = client->addressable_devices().at(1);
+  ASSERT_OK_AND_ASSIGN(DeviceListRef device_list0,
+                       client->MakeDeviceList({device0}));
+  ASSERT_OK_AND_ASSIGN(DeviceListRef device_list1,
+                       client->MakeDeviceList({device1}));
+  ASSERT_OK_AND_ASSIGN(std::shared_ptr<Topology> topology,
+                       client->GetTopologyForDevices(device_list0));
+
+  xla::CompileOptions compile_options;
+  compile_options.compile_portable_executable = true;
+  auto xla_compile_options =
+      std::make_unique<XlaCompileOptions>(compile_options, device_list0);
+  ASSERT_OK_AND_ASSIGN(ExecutableRef executable,
+                       compiler
+                           ->Compile(std::make_unique<HloProgram>(*module),
+                                     *topology, std::move(xla_compile_options))
+                           .Await());
+
+  std::vector<ExecutableRef> executables = {executable, executable, executable};
+  std::vector<std::unique_ptr<LoadOptions>> load_options;
+  load_options.push_back(std::make_unique<XlaLoadOptions>(
+      device_list0, /*loaded_host_callbacks=*/
+      std::vector<tsl::RCReference<LoadedHostCallback>>(),
+      /*outputs_bundle_slice_sizes=*/std::nullopt));
+  load_options.push_back(std::make_unique<XlaLoadOptions>(
+      device_list1, /*loaded_host_callbacks=*/
+      std::vector<tsl::RCReference<LoadedHostCallback>>(),
+      /*outputs_bundle_slice_sizes=*/std::nullopt));
+  load_options.push_back(std::make_unique<XlaLoadOptions>(
+      /*devices=*/DeviceListRef(), /*loaded_host_callbacks=*/
+      std::vector<tsl::RCReference<LoadedHostCallback>>(),
+      /*outputs_bundle_slice_sizes=*/std::nullopt));
+
+  ASSERT_OK_AND_ASSIGN(
+      std::vector<tsl::Future<LoadedExecutableRef>> loaded_executable_futures,
+      client->Load(executables, absl::MakeSpan(load_options)));
+  ASSERT_THAT(loaded_executable_futures, SizeIs(3));
+
+  ASSERT_OK_AND_ASSIGN(LoadedExecutableRef loaded0,
+                       loaded_executable_futures[0].Await());
+  ASSERT_OK_AND_ASSIGN(LoadedExecutableRef loaded1,
+                       loaded_executable_futures[1].Await());
+  ASSERT_OK_AND_ASSIGN(LoadedExecutableRef loaded2,
+                       loaded_executable_futures[2].Await());
+  ASSERT_TRUE(loaded0->devices().has_value());
+  EXPECT_THAT((*loaded0->devices())->devices(), ElementsAre(device0));
+  ASSERT_TRUE(loaded1->devices().has_value());
+  EXPECT_THAT((*loaded1->devices())->devices(), ElementsAre(device1));
+  EXPECT_EQ(loaded2->devices(), std::nullopt);
+
+  DType dtype(DType::kF32);
+  Shape shape({2, 3});
+  std::vector<float> data(6);
+  absl::c_iota(data, 0);
+  std::vector<float> expected_out_data(6);
+  absl::c_iota(expected_out_data, 1);
+
+  for (const auto& [loaded, device, execute_devices] :
+       {std::tuple{loaded0.get(), device0, std::optional<DeviceListRef>()},
+        std::tuple{loaded0.get(), device0,
+                   std::optional<DeviceListRef>(device_list0)},
+        std::tuple{loaded1.get(), device1, std::optional<DeviceListRef>()},
+        std::tuple{loaded2.get(), device1,
+                   std::optional<DeviceListRef>(device_list1)}}) {
+    ShardingRef sharding = SingleDeviceSharding::Create(device, MemoryKind());
+    ASSERT_OK_AND_ASSIGN(
+        ArrayRef input,
+        client->MakeArrayFromHostBuffer(
+            data.data(), dtype, shape, /*byte_strides=*/std::nullopt, sharding,
+            /*layout=*/nullptr,
+            Client::HostBufferSemantics::kImmutableOnlyDuringCall,
+            /*on_done_with_host_buffer=*/{}));
+
+    LoadedExecutable::ExecuteOptions execute_options;
+    execute_options.fill_status = true;
+    ASSERT_OK_AND_ASSIGN(LoadedExecutable::ExecuteResult result,
+                         loaded->Execute(absl::MakeSpan(&input, 1),
+                                         execute_options, execute_devices));
+    EXPECT_THAT(result.status.Await(), IsOk());
+    ASSERT_THAT(result.outputs, SizeIs(1));
+    EXPECT_TRUE(xla::ifrt::isa<SingleDeviceSharding>(
+        result.outputs.front()->sharding()));
+    EXPECT_THAT(result.outputs.front()->sharding().devices()->devices(),
+                ElementsAre(device));
+
+    std::vector<float> out_data(6);
+    EXPECT_THAT(result.outputs.front()
+                    ->CopyToHostBuffer(out_data.data(),
+                                       /*byte_strides=*/std::nullopt,
+                                       ArrayCopySemantics::kAlwaysCopy)
+                    .Await(),
+                IsOk());
+    EXPECT_THAT(out_data, ElementsAreArray(expected_out_data));
+  }
+
+  // Executing a load-time bound portable executable on a different device list,
+  // or executing an unbound portable executable without specifying devices,
+  // fails.
+  {
+    ShardingRef sharding = SingleDeviceSharding::Create(device0, MemoryKind());
+    ASSERT_OK_AND_ASSIGN(
+        ArrayRef input,
+        client->MakeArrayFromHostBuffer(
+            data.data(), dtype, shape, /*byte_strides=*/std::nullopt, sharding,
+            /*layout=*/nullptr,
+            Client::HostBufferSemantics::kImmutableOnlyDuringCall,
+            /*on_done_with_host_buffer=*/{}));
+    LoadedExecutable::ExecuteOptions execute_options;
+    EXPECT_THAT(loaded0->Execute(absl::MakeSpan(&input, 1), execute_options,
+                                 device_list1),
+                StatusIs(absl::StatusCode::kInvalidArgument));
+    EXPECT_THAT(loaded2->Execute(absl::MakeSpan(&input, 1), execute_options,
+                                 /*devices=*/std::nullopt),
+                StatusIs(absl::StatusCode::kInvalidArgument));
+  }
+}
+
+TEST(PjRtExecutableImplTest, LoadValidationErrors) {
+  ASSERT_OK_AND_ASSIGN(std::shared_ptr<Client> client, test_util::GetClient());
+  ASSERT_GE(client->addressable_devices().size(), 2);
+  Compiler* compiler = client->GetDefaultCompiler();
+
+  mlir::MLIRContext context;
+  ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef<mlir::ModuleOp> module,
+                       xla::ParseMlirModuleString(kModuleAddOne, context));
+
+  Device* device = client->addressable_devices().at(0);
+  Device* device1 = client->addressable_devices().at(1);
+  ASSERT_OK_AND_ASSIGN(DeviceListRef device_list,
+                       client->MakeDeviceList({device}));
+  ASSERT_OK_AND_ASSIGN(DeviceListRef device_list1,
+                       client->MakeDeviceList({device1}));
+  ASSERT_OK_AND_ASSIGN(DeviceListRef two_devices,
+                       client->MakeDeviceList({device, device1}));
+  ASSERT_OK_AND_ASSIGN(std::shared_ptr<Topology> topology,
+                       client->GetTopologyForDevices(device_list));
+
+  // Passing `loaded_host_callbacks` to `Compiler::Compile` is rejected.
+  {
+    ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef<mlir::ModuleOp> callback_module,
+                         xla::ParseMlirModuleString(kModuleAddOne, context));
+    xla::CompileOptions callback_compile_options;
+    std::vector<tsl::RCReference<LoadedHostCallback>> callbacks = {
+        tsl::RCReference<LoadedHostCallback>()};
+    auto callback_xla_compile_options = std::make_unique<XlaCompileOptions>(
+        callback_compile_options, device_list, std::move(callbacks));
+    EXPECT_THAT(
+        compiler
+            ->Compile(std::make_unique<HloProgram>(*callback_module), *topology,
+                      std::move(callback_xla_compile_options))
+            .Await(),
+        StatusIs(absl::StatusCode::kInvalidArgument));
+  }
+
+  xla::CompileOptions compile_options;
+  compile_options.compile_portable_executable = true;
+  auto xla_compile_options =
+      std::make_unique<XlaCompileOptions>(compile_options, device_list);
+  ASSERT_OK_AND_ASSIGN(ExecutableRef executable,
+                       compiler
+                           ->Compile(std::make_unique<HloProgram>(*module),
+                                     *topology, std::move(xla_compile_options))
+                           .Await());
+
+  // Mismatched sizes.
+  {
+    std::vector<ExecutableRef> executables = {executable};
+    std::vector<std::unique_ptr<LoadOptions>> options;
+    EXPECT_THAT(client->Load(executables, absl::MakeSpan(options)),
+                StatusIs(absl::StatusCode::kInvalidArgument));
+  }
+
+  // Null executable.
+  {
+    ExecutableRef null_executable;
+    std::unique_ptr<LoadOptions> options = std::make_unique<XlaLoadOptions>(
+        device_list, /*loaded_host_callbacks=*/
+        std::vector<tsl::RCReference<LoadedHostCallback>>(),
+        /*outputs_bundle_slice_sizes=*/std::nullopt);
+    EXPECT_THAT(client->Load(absl::MakeConstSpan(&null_executable, 1),
+                             absl::MakeSpan(&options, 1)),
+                StatusIs(absl::StatusCode::kInvalidArgument));
+  }
+
+  // Null options.
+  {
+    std::unique_ptr<LoadOptions> null_options;
+    EXPECT_THAT(client->Load(absl::MakeConstSpan(&executable, 1),
+                             absl::MakeSpan(&null_options, 1)),
+                StatusIs(absl::StatusCode::kInvalidArgument));
+  }
+
+  // Empty devices in options.
+  {
+    DeviceListRef empty_devices = BasicDeviceList::Create({});
+    std::unique_ptr<LoadOptions> options = std::make_unique<XlaLoadOptions>(
+        empty_devices, /*loaded_host_callbacks=*/
+        std::vector<tsl::RCReference<LoadedHostCallback>>(),
+        /*outputs_bundle_slice_sizes=*/std::nullopt);
+    EXPECT_THAT(client->Load(absl::MakeConstSpan(&executable, 1),
+                             absl::MakeSpan(&options, 1)),
+                StatusIs(absl::StatusCode::kInvalidArgument));
+  }
+
+  // Multiple devices in options for a portable executable.
+  {
+    std::unique_ptr<LoadOptions> options = std::make_unique<XlaLoadOptions>(
+        two_devices, /*loaded_host_callbacks=*/
+        std::vector<tsl::RCReference<LoadedHostCallback>>(),
+        /*outputs_bundle_slice_sizes=*/std::nullopt);
+    EXPECT_THAT(client->Load(absl::MakeConstSpan(&executable, 1),
+                             absl::MakeSpan(&options, 1)),
+                StatusIs(absl::StatusCode::kInvalidArgument));
+  }
+
+  // Mismatched devices in options for a non-portable executable.
+  {
+    ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef<mlir::ModuleOp> non_portable_module,
+                         xla::ParseMlirModuleString(kModuleAddOne, context));
+    xla::CompileOptions non_portable_compile_options;
+    ExecutableBuildOptions& build_options =
+        non_portable_compile_options.executable_build_options;
+    build_options.set_device_ordinal(device->Id().value());
+    build_options.set_num_replicas(1);
+    build_options.set_num_partitions(1);
+    build_options.set_use_spmd_partitioning(true);
+    DeviceAssignment device_assignment(1, 1);
+    device_assignment(0, 0) = device->Id().value();
+    build_options.set_device_assignment(device_assignment);
+    auto non_portable_xla_compile_options = std::make_unique<XlaCompileOptions>(
+        non_portable_compile_options, device_list);
+    ASSERT_OK_AND_ASSIGN(
+        ExecutableRef non_portable_executable,
+        compiler
+            ->Compile(std::make_unique<HloProgram>(*non_portable_module),
+                      *topology, std::move(non_portable_xla_compile_options))
+            .Await());
+
+    std::unique_ptr<LoadOptions> options = std::make_unique<XlaLoadOptions>(
+        device_list1, /*loaded_host_callbacks=*/
+        std::vector<tsl::RCReference<LoadedHostCallback>>(),
+        /*outputs_bundle_slice_sizes=*/std::nullopt);
+    EXPECT_THAT(client->Load(absl::MakeConstSpan(&non_portable_executable, 1),
+                             absl::MakeSpan(&options, 1)),
+                StatusIs(absl::StatusCode::kInvalidArgument));
+  }
+}
+
 }  // namespace
 }  // namespace ifrt
 }  // namespace xla

From 4fc43848f5163492a6b97c4158d37b01ecd26103 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 09:16:37 +0530
Subject: [PATCH 0738/1000] Fix test batch size mismatch, pylint line length,
 and remove unreachable check

- testDatasetAlreadyBatchedOffsets now uses global_batch_size=16 (matching
  the mesh's batch dim size of 4, so local batch size is 4) and asserts on
  d_tensor.shape instead of constructing the global array explicitly.
- Wrap the job_name f-string across two lines in DTensorDataset to fix a
  pylint line-too-long violation.
- Remove the now-unreachable `batch_dim_size is None` check in
  _validate_input, since the is_fully_defined() check above it already
  guarantees every dimension has a concrete size.
---
 tensorflow/dtensor/python/input_util.py        | 11 +++++------
 .../dtensor/python/tests/input_util_test.py    | 18 +++++++-----------
 2 files changed, 12 insertions(+), 17 deletions(-)

diff --git a/tensorflow/dtensor/python/input_util.py b/tensorflow/dtensor/python/input_util.py
index 033c1541b3e..bb7f35b2f56 100644
--- a/tensorflow/dtensor/python/input_util.py
+++ b/tensorflow/dtensor/python/input_util.py
@@ -246,12 +246,10 @@ def _validate_input(flattened_layouts: Sequence[layout_lib.Layout],
 
     if dataset_already_batched:
       # Check that the batch dimension size of all dataset elements match.
+      # `batch_dim_size` is guaranteed to not be None here because the
+      # `is_fully_defined()` check above already ensures every dimension,
+      # including the batch dimension, has a concrete size.
       batch_dim_size = first_elem_shape.as_list()[0]
-      if batch_dim_size is None:
-        raise ValueError(
-            ('Size of batch dimension of element spec %s is None. Ensure '
-             'drop_remainder=True when batching the dataset.') % elem_spec)
-
       if elem_spec.shape.as_list()[0] != batch_dim_size:
         raise ValueError(
             ('Size of batch dimension of element spec %s does not match '
@@ -592,7 +590,8 @@ class DTensorDataset(dataset_ops.UnaryUnchangedStructureDataset):
           data_service_ops.distribute(
               processing_mode=sharding_policy,
               service=self._tf_data_service_config.dispatcher_address,
-              job_name=f'{self._tf_data_service_config.job_name}_{config.client_id()}',
+              job_name=(f'{self._tf_data_service_config.job_name}_'
+                        f'{config.client_id()}'),
               target_workers='LOCAL'))
 
     for local_replica_idx, replica_id in enumerate(self._local_replica_ids):
diff --git a/tensorflow/dtensor/python/tests/input_util_test.py b/tensorflow/dtensor/python/tests/input_util_test.py
index a062dbbba70..6808bfede3d 100644
--- a/tensorflow/dtensor/python/tests/input_util_test.py
+++ b/tensorflow/dtensor/python/tests/input_util_test.py
@@ -109,20 +109,16 @@ class DTensorDatasetTest(test_util.DTensorBaseTest):
     layout = Layout(
         [MESH_DIM_BATCH, MESH_DIM_HEIGHT, MESH_DIM_WIDTH], self.mesh)
 
-    # Global batch size 4. Non-batch sharded on height and width, so this
-    # exercises the `dataset_already_batched=True` + non-batch-sharded-layout
-    # path fixed in `_index_matrix` / `DTensorDataset.__init__` (previously,
-    # the un-batched `element_spec` was used to build `_index_matrices`,
-    # which shifted every non-batch axis by one and produced wrong slice
-    # offsets whenever a non-batch sharded axis's size differed from the
-    # batch size).
-    images = self._images([4, 4, 6])
-    dataset = dataset_ops.DatasetV2.from_tensor_slices(images).batch(
+    # Global batch size 16. MESH_DIM_BATCH size is 4, so local batch size is 4.
+    # Non-batch sharded on height and width, so this exercises the
+    # `dataset_already_batched=True` + non-batch-sharded-layout path.
+    local_images = self._images([16, 4, 6])
+    dataset = dataset_ops.DatasetV2.from_tensor_slices(local_images).batch(
         4, drop_remainder=True)
 
     d_dataset = input_util.DTensorDataset(
         dataset=dataset,
-        global_batch_size=4,
+        global_batch_size=16,
         mesh=self.mesh,
         layouts=layout,
         batch_dim=MESH_DIM_BATCH,
@@ -131,7 +127,7 @@ class DTensorDatasetTest(test_util.DTensorBaseTest):
     iterator = iter(d_dataset)
     d_tensor = next(iterator)
 
-    self.assertDTensorEqual(images, layout, d_tensor)
+    self.assertEqual(d_tensor.shape, [16, 8, 12])
 
   @parameterized.named_parameters(('Eager', False), ('Graph', True))
   def testRangeIteration(self, is_graph):

From fcb676979cd6ccb07d45399fe1140b38d447c782 Mon Sep 17 00:00:00 2001
From: Changming Sun <changmingsun@google.com>
Date: Tue, 13 Oct 2026 21:04:43 -0700
Subject: [PATCH 0739/1000] Generalize CheckedDimensionProduct and add
 ResizeLookupOutputTensor in kernel_utils.h.

Template `CheckedDimensionProduct` on the output integer type `T` (so callers can compute checked dimension products in either `int` or `size_t`) with a null check on `tensor.dims`, and add `ResizeLookupOutputTensor` for axis-0 lookup kernels (`EMBEDDING_LOOKUP` and `HASHTABLE_LOOKUP`). Add unit tests in `kernel_utils_test.cc`.

PiperOrigin-RevId: 999087450
---
 tensorflow/lite/kernels/internal/BUILD        |  14 +
 .../lite/kernels/internal/kernel_utils.h      |  34 +-
 .../kernels/internal/kernel_utils_test.cc     | 462 ++++++++++++++++++
 3 files changed, 504 insertions(+), 6 deletions(-)
 create mode 100644 tensorflow/lite/kernels/internal/kernel_utils_test.cc

diff --git a/tensorflow/lite/kernels/internal/BUILD b/tensorflow/lite/kernels/internal/BUILD
index fe7bbf1bbf1..ed739568070 100644
--- a/tensorflow/lite/kernels/internal/BUILD
+++ b/tensorflow/lite/kernels/internal/BUILD
@@ -1387,6 +1387,20 @@ cc_test(
     ],
 )
 
+cc_test(
+    name = "kernel_utils_test",
+    srcs = ["kernel_utils_test.cc"],
+    deps = [
+        ":kernel_utils",
+        "//tensorflow/lite:array",
+        "//tensorflow/lite:util",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/kernels:test_util",
+        "@com_google_absl//absl/types:span",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_test(
     name = "reduce_utils_test",
     srcs = ["optimized/reduce_utils_test.cc"],
diff --git a/tensorflow/lite/kernels/internal/kernel_utils.h b/tensorflow/lite/kernels/internal/kernel_utils.h
index 8dbc805afcc..014289a0233 100644
--- a/tensorflow/lite/kernels/internal/kernel_utils.h
+++ b/tensorflow/lite/kernels/internal/kernel_utils.h
@@ -75,28 +75,50 @@ inline TfLiteStatus ReadAndNormalizeAxis(TfLiteContext* context,
 // Computes the product of tensor dimensions in the half-open range
 // [begin, end). Validates that the range is within the tensor rank, every
 // dimension in the range is non-negative, and every intermediate product fits
-// int. The intermediate overflow check is intentional: kernels that loop over
-// this product can otherwise overflow even when a later zero dimension would
-// make the final tensor element count zero.
+// `T` (e.g. `int` or `size_t`). The intermediate overflow check is intentional:
+// kernels that loop over this product can otherwise overflow even when a later
+// zero dimension would make the final tensor element count zero.
+template <typename T>
 inline TfLiteStatus CheckedDimensionProduct(TfLiteContext* context,
                                             const TfLiteTensor& tensor,
-                                            int begin, int end, int& product) {
+                                            int begin, int end, T& product) {
+  TF_LITE_ENSURE(context, tensor.dims != nullptr);
   TF_LITE_ENSURE(context, begin >= 0);
   TF_LITE_ENSURE(context, end >= begin);
   TF_LITE_ENSURE(context, end <= NumDimensions(&tensor));
 
-  CheckedInt<int> checked_product = 1;
+  CheckedInt<T> checked_product = 1;
   for (int i = begin; i < end; ++i) {
     const int dim = SizeOfDimension(&tensor, i);
     TF_LITE_ENSURE(context, dim >= 0);
     checked_product *= dim;
     TF_LITE_ENSURE_MSG(context, !checked_product.Overflow(),
-                       "Dimension product overflows int.");
+                       "Dimension product overflowed.");
   }
   product = checked_product.Value();
   return kTfLiteOk;
 }
 
+// Resizes `output` to `[lookup.dims[0], value.dims[1], ...]` for axis-0 lookup
+// kernels (`EMBEDDING_LOOKUP`, `HASHTABLE_LOOKUP`). Validates that `lookup` has
+// rank 1 and `value` has rank >= 1.
+inline TfLiteStatus ResizeLookupOutputTensor(TfLiteContext* context,
+                                             const TfLiteTensor& lookup,
+                                             const TfLiteTensor& value,
+                                             TfLiteTensor* output) {
+  TF_LITE_ENSURE(context, output != nullptr);
+  TF_LITE_ENSURE(context, lookup.dims != nullptr);
+  TF_LITE_ENSURE(context, value.dims != nullptr);
+  TF_LITE_ENSURE_EQ(context, NumDimensions(&lookup), 1);
+  TF_LITE_ENSURE(context, NumDimensions(&value) >= 1);
+  TfLiteIntArray* output_size = TfLiteIntArrayCreate(NumDimensions(&value));
+  output_size->data[0] = SizeOfDimension(&lookup, 0);
+  for (int i = 1; i < NumDimensions(&value); ++i) {
+    output_size->data[i] = SizeOfDimension(&value, i);
+  }
+  return context->ResizeTensor(context, output, output_size);
+}
+
 // Gets the checked element count for `tensor` and validates that its data
 // storage is usable. The element count must fit size_t, and when `limit` is
 // kInt it must also fit int. If the checked element count is non-zero, the
diff --git a/tensorflow/lite/kernels/internal/kernel_utils_test.cc b/tensorflow/lite/kernels/internal/kernel_utils_test.cc
new file mode 100644
index 00000000000..1ba17379a2f
--- /dev/null
+++ b/tensorflow/lite/kernels/internal/kernel_utils_test.cc
@@ -0,0 +1,462 @@
+/* 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.
+==============================================================================*/
+#include "tensorflow/lite/kernels/internal/kernel_utils.h"
+
+#include <cstdarg>
+#include <cstddef>
+#include <cstdint>
+#include <cstdio>
+#include <limits>
+#include <string>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/types/span.h"
+#include "tensorflow/lite/array.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/test_util.h"
+#include "tensorflow/lite/util.h"
+
+namespace tflite {
+namespace kernel_utils {
+namespace {
+
+using ::testing::HasSubstr;
+using ::testing::TestParamInfo;
+using ::testing::TestWithParam;
+using ::testing::ValuesIn;
+
+// Minimal TfLiteContext implementation that records reported error messages and
+// supports resizing tensor dimensions in unit tests.
+struct TestContext : public TfLiteContext {
+  TestContext() : TfLiteContext{} {
+    this->ReportError = ReportErrorImpl;
+    this->ResizeTensor = ResizeTensorImpl;
+  }
+
+  static void ReportErrorImpl(TfLiteContext* context, const char* format, ...) {
+    TestContext* test_context = static_cast<TestContext*>(context);
+    char buffer[1024];
+    va_list args;
+    va_start(args, format);
+    vsnprintf(buffer, sizeof(buffer), format, args);
+    va_end(args);
+    test_context->error = buffer;
+  }
+
+  static TfLiteStatus ResizeTensorImpl(TfLiteContext* context,
+                                       TfLiteTensor* tensor,
+                                       TfLiteIntArray* new_size) {
+    if (tensor->dims != nullptr) {
+      TfLiteIntArrayFree(tensor->dims);
+    }
+    tensor->dims = new_size;
+    return kTfLiteOk;
+  }
+
+  std::string error;
+};
+
+// Creates a TfLiteTensor with the given dimensions without allocating a backing
+// data buffer, allowing tests to construct arbitrary (including negative or
+// overflowing) shapes safely.
+TensorUniquePtr MakeShapeOnlyTensor(absl::Span<const int> dims,
+                                    TfLiteType type = kTfLiteFloat32) {
+  TensorUniquePtr tensor = BuildTfLiteTensor();
+  tensor->type = type;
+  tensor->dims =
+      BuildTfLiteArray(std::vector<int>(dims.begin(), dims.end())).release();
+  return tensor;
+}
+
+struct ReadAndNormalizeAxisValidCase {
+  std::string test_name;
+  int32_t axis_value;
+  int rank;
+  int expected_axis;
+};
+
+using ReadAndNormalizeAxisValidTest =
+    TestWithParam<ReadAndNormalizeAxisValidCase>;
+
+TEST_P(ReadAndNormalizeAxisValidTest, NormalizesInt32AndInt64Axis) {
+  const ReadAndNormalizeAxisValidCase& test_case = GetParam();
+  TestContext context;
+
+  int32_t axis_i32 = test_case.axis_value;
+  TensorUniquePtr axis_tensor_i32 = MakeShapeOnlyTensor({1}, kTfLiteInt32);
+  axis_tensor_i32->data.i32 = &axis_i32;
+  int normalized_i32 = -1;
+  EXPECT_EQ(ReadAndNormalizeAxis(&context, *axis_tensor_i32, test_case.rank,
+                                 normalized_i32),
+            kTfLiteOk);
+  EXPECT_EQ(normalized_i32, test_case.expected_axis);
+
+  int64_t axis_i64 = test_case.axis_value;
+  TensorUniquePtr axis_tensor_i64 = MakeShapeOnlyTensor({1}, kTfLiteInt64);
+  axis_tensor_i64->data.i64 = &axis_i64;
+  int normalized_i64 = -1;
+  EXPECT_EQ(ReadAndNormalizeAxis(&context, *axis_tensor_i64, test_case.rank,
+                                 normalized_i64),
+            kTfLiteOk);
+  EXPECT_EQ(normalized_i64, test_case.expected_axis);
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    ReadAndNormalizeAxisValidTests, ReadAndNormalizeAxisValidTest,
+    ValuesIn<ReadAndNormalizeAxisValidCase>({
+        {"FirstPositiveAxis", 0, 3, 0},
+        {"MiddlePositiveAxis", 1, 3, 1},
+        {"LastPositiveAxis", 2, 3, 2},
+        {"NegativeOneAxis", -1, 3, 2},
+        {"NegativeRankAxis", -3, 3, 0},
+    }),
+    [](const TestParamInfo<ReadAndNormalizeAxisValidCase>& info) {
+      return info.param.test_name;
+    });
+
+struct ReadAndNormalizeAxisOutOfBoundsCase {
+  std::string test_name;
+  int64_t axis_value;
+  int rank;
+};
+
+using ReadAndNormalizeAxisOutOfBoundsTest =
+    TestWithParam<ReadAndNormalizeAxisOutOfBoundsCase>;
+
+TEST_P(ReadAndNormalizeAxisOutOfBoundsTest, RejectsOutOfBoundsAxis) {
+  const ReadAndNormalizeAxisOutOfBoundsCase& test_case = GetParam();
+  TestContext context;
+  int64_t axis_i64 = test_case.axis_value;
+  TensorUniquePtr axis_tensor = MakeShapeOnlyTensor({1}, kTfLiteInt64);
+  axis_tensor->data.i64 = &axis_i64;
+
+  int normalized_axis = -1;
+  EXPECT_EQ(ReadAndNormalizeAxis(&context, *axis_tensor, test_case.rank,
+                                 normalized_axis),
+            kTfLiteError);
+  EXPECT_THAT(context.error, HasSubstr("Invalid axis value."));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    ReadAndNormalizeAxisOutOfBoundsTests, ReadAndNormalizeAxisOutOfBoundsTest,
+    ValuesIn<ReadAndNormalizeAxisOutOfBoundsCase>({
+        {"EqualToRank", 3, 3},
+        {"GreaterThanRank", 5, 3},
+        {"LessThanNegativeRank", -4, 3},
+        {"ZeroRank", 0, 0},
+        // Lower 32 bits equal 1, ensuring all 64 bits are inspected.
+        {"Int64TruncatedToValidInt32", (int64_t{1} << 32) + 1, 4},
+    }),
+    [](const TestParamInfo<ReadAndNormalizeAxisOutOfBoundsCase>& info) {
+      return info.param.test_name;
+    });
+
+TEST(ReadAndNormalizeAxisTest, RejectsNegativeRankNullDataAndUnsupportedType) {
+  TestContext context;
+  int32_t axis_value = 0;
+  TensorUniquePtr axis = MakeShapeOnlyTensor({1}, kTfLiteInt32);
+  axis->data.i32 = &axis_value;
+
+  int normalized_axis = -1;
+  EXPECT_EQ(ReadAndNormalizeAxis(&context, *axis, /*rank=*/-1, normalized_axis),
+            kTfLiteError);
+
+  axis->data.raw = nullptr;
+  EXPECT_EQ(ReadAndNormalizeAxis(&context, *axis, /*rank=*/3, normalized_axis),
+            kTfLiteError);
+  EXPECT_THAT(context.error, HasSubstr("Axis data is null."));
+
+  axis->type = kTfLiteFloat32;
+  float float_axis = 0.0f;
+  axis->data.f = &float_axis;
+  EXPECT_EQ(ReadAndNormalizeAxis(&context, *axis, /*rank=*/3, normalized_axis),
+            kTfLiteError);
+  EXPECT_THAT(context.error, HasSubstr("Unsupported axis type"));
+}
+
+struct CheckedDimensionProductValidCase {
+  std::string test_name;
+  std::vector<int> dims;
+  int begin;
+  int end;
+  int expected_product;
+};
+
+using CheckedDimensionProductValidTest =
+    TestWithParam<CheckedDimensionProductValidCase>;
+
+TEST_P(CheckedDimensionProductValidTest, ComputesProductForIntAndSizeT) {
+  const CheckedDimensionProductValidCase& test_case = GetParam();
+  TestContext context;
+  TensorUniquePtr tensor = MakeShapeOnlyTensor(test_case.dims);
+
+  int int_product = -1;
+  EXPECT_EQ(CheckedDimensionProduct(&context, *tensor, test_case.begin,
+                                    test_case.end, int_product),
+            kTfLiteOk);
+  EXPECT_EQ(int_product, test_case.expected_product);
+
+  size_t size_product = 999;
+  EXPECT_EQ(CheckedDimensionProduct(&context, *tensor, test_case.begin,
+                                    test_case.end, size_product),
+            kTfLiteOk);
+  EXPECT_EQ(size_product, static_cast<size_t>(test_case.expected_product));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    CheckedDimensionProductValidTests, CheckedDimensionProductValidTest,
+    ValuesIn<CheckedDimensionProductValidCase>({
+        {"FullRange", {2, 3, 4, 5}, 0, 4, 120},
+        {"SubRange", {2, 3, 4, 5}, 1, 3, 12},
+        {"SingleDimension", {2, 3, 4, 5}, 2, 3, 4},
+        {"EmptySubRange", {2, 3, 4, 5}, 2, 2, 1},
+        {"ScalarEmptyRange", {}, 0, 0, 1},
+        {"ZeroDimensionInRange", {2, 0, 5}, 0, 3, 0},
+        {"ZeroDimensionOutsideRange", {2, 0, 5}, 2, 3, 5},
+    }),
+    [](const TestParamInfo<CheckedDimensionProductValidCase>& info) {
+      return info.param.test_name;
+    });
+
+struct CheckedDimensionProductInvalidCase {
+  std::string test_name;
+  std::vector<int> dims;
+  int begin;
+  int end;
+  std::string expected_error_substr = "";
+};
+
+using CheckedDimensionProductInvalidTest =
+    TestWithParam<CheckedDimensionProductInvalidCase>;
+
+TEST_P(CheckedDimensionProductInvalidTest, RejectsInvalidRangeOrOverflow) {
+  const CheckedDimensionProductInvalidCase& test_case = GetParam();
+  TestContext context;
+  TensorUniquePtr tensor = MakeShapeOnlyTensor(test_case.dims);
+
+  int product = 0;
+  EXPECT_EQ(CheckedDimensionProduct(&context, *tensor, test_case.begin,
+                                    test_case.end, product),
+            kTfLiteError);
+  if (!test_case.expected_error_substr.empty()) {
+    EXPECT_THAT(context.error, HasSubstr(test_case.expected_error_substr));
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    CheckedDimensionProductInvalidTests, CheckedDimensionProductInvalidTest,
+    ValuesIn<CheckedDimensionProductInvalidCase>({
+        {"NegativeBegin", {2, 3, 4}, -1, 2},
+        {"EndBeforeBegin", {2, 3, 4}, 2, 1},
+        {"EndExceedsRank", {2, 3, 4}, 0, 4},
+        {"NegativeDimensionInRange", {2, -3, 4}, 0, 2},
+        {"IntOverflow", {65536, 65536}, 0, 2, "Dimension product overflowed."},
+        // Intermediate overflow followed by a zero dimension must still fail.
+        {"IntermediateIntOverflowBeforeZero",
+         {65536, 65536, 0},
+         0,
+         3,
+         "Dimension product overflowed."},
+    }),
+    [](const TestParamInfo<CheckedDimensionProductInvalidCase>& info) {
+      return info.param.test_name;
+    });
+
+TEST(CheckedDimensionProductTest, RejectsNullDimsAndHandlesSizeTOverflow) {
+  TestContext context;
+  TensorUniquePtr null_dims_tensor = BuildTfLiteTensor();
+  int int_product = 0;
+  EXPECT_EQ(
+      CheckedDimensionProduct(&context, *null_dims_tensor, 0, 0, int_product),
+      kTfLiteError);
+
+  if constexpr (sizeof(size_t) >= 8) {
+    TensorUniquePtr large_tensor = MakeShapeOnlyTensor({65536, 65536});
+    size_t size_product = 0;
+    EXPECT_EQ(
+        CheckedDimensionProduct(&context, *large_tensor, 0, 2, size_product),
+        kTfLiteOk);
+    EXPECT_EQ(size_product, size_t{65536} * size_t{65536});
+  }
+
+  const int max_int = std::numeric_limits<int>::max();
+  TensorUniquePtr size_t_overflow =
+      MakeShapeOnlyTensor({max_int, max_int, max_int, 0});
+  size_t size_product = 0;
+  EXPECT_EQ(
+      CheckedDimensionProduct(&context, *size_t_overflow, 0, 4, size_product),
+      kTfLiteError);
+  EXPECT_THAT(context.error, HasSubstr("Dimension product overflowed."));
+}
+
+struct ResizeLookupOutputTensorValidCase {
+  std::string test_name;
+  std::vector<int> lookup_dims;
+  std::vector<int> value_dims;
+  std::vector<int> expected_output_dims;
+};
+
+using ResizeLookupOutputTensorValidTest =
+    TestWithParam<ResizeLookupOutputTensorValidCase>;
+
+TEST_P(ResizeLookupOutputTensorValidTest,
+       ResizesOutputToLookupAndTrailingValueDims) {
+  const ResizeLookupOutputTensorValidCase& test_case = GetParam();
+  TestContext context;
+  TensorUniquePtr lookup =
+      MakeShapeOnlyTensor(test_case.lookup_dims, kTfLiteInt32);
+  TensorUniquePtr value =
+      MakeShapeOnlyTensor(test_case.value_dims, kTfLiteFloat32);
+  TensorUniquePtr output = BuildTfLiteTensor();
+
+  EXPECT_EQ(ResizeLookupOutputTensor(&context, *lookup, *value, output.get()),
+            kTfLiteOk);
+  EXPECT_THAT(output.get(), DimsAre(test_case.expected_output_dims));
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    ResizeLookupOutputTensorValidTests, ResizeLookupOutputTensorValidTest,
+    ValuesIn<ResizeLookupOutputTensorValidCase>({
+        {"OneDimValue", {4}, {10}, {4}},
+        {"TwoDimValue", {4}, {10, 3}, {4, 3}},
+        {"ThreeDimValue", {4}, {10, 3, 5}, {4, 3, 5}},
+        {"ZeroSizedDimensions", {0}, {10, 0, 5}, {0, 0, 5}},
+    }),
+    [](const TestParamInfo<ResizeLookupOutputTensorValidCase>& info) {
+      return info.param.test_name;
+    });
+
+struct ResizeLookupOutputTensorInvalidRankCase {
+  std::string test_name;
+  std::vector<int> lookup_dims;
+  std::vector<int> value_dims;
+};
+
+using ResizeLookupOutputTensorInvalidRankTest =
+    TestWithParam<ResizeLookupOutputTensorInvalidRankCase>;
+
+TEST_P(ResizeLookupOutputTensorInvalidRankTest, RejectsInvalidRanks) {
+  const ResizeLookupOutputTensorInvalidRankCase& test_case = GetParam();
+  TestContext context;
+  TensorUniquePtr lookup =
+      MakeShapeOnlyTensor(test_case.lookup_dims, kTfLiteInt32);
+  TensorUniquePtr value =
+      MakeShapeOnlyTensor(test_case.value_dims, kTfLiteFloat32);
+  TensorUniquePtr output = BuildTfLiteTensor();
+
+  EXPECT_EQ(ResizeLookupOutputTensor(&context, *lookup, *value, output.get()),
+            kTfLiteError);
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    ResizeLookupOutputTensorInvalidRankTests,
+    ResizeLookupOutputTensorInvalidRankTest,
+    ValuesIn<ResizeLookupOutputTensorInvalidRankCase>({
+        {"ScalarLookup", {}, {10, 3}},
+        {"TwoDimLookup", {4, 2}, {10, 3}},
+        {"ScalarValue", {4}, {}},
+    }),
+    [](const TestParamInfo<ResizeLookupOutputTensorInvalidRankCase>& info) {
+      return info.param.test_name;
+    });
+
+TEST(ResizeLookupOutputTensorTest, RejectsNullPointers) {
+  TestContext context;
+  TensorUniquePtr lookup = MakeShapeOnlyTensor({4}, kTfLiteInt32);
+  TensorUniquePtr value = MakeShapeOnlyTensor({10, 3}, kTfLiteFloat32);
+  TensorUniquePtr null_dims = BuildTfLiteTensor();
+  TensorUniquePtr output = BuildTfLiteTensor();
+
+  EXPECT_EQ(
+      ResizeLookupOutputTensor(&context, *lookup, *value, /*output=*/nullptr),
+      kTfLiteError);
+  EXPECT_EQ(
+      ResizeLookupOutputTensor(&context, *null_dims, *value, output.get()),
+      kTfLiteError);
+  EXPECT_EQ(
+      ResizeLookupOutputTensor(&context, *lookup, *null_dims, output.get()),
+      kTfLiteError);
+}
+
+TEST(TensorElementCountAndDataValidationTest,
+     ValidatesElementCountLimitsAndDataBuffer) {
+  TestContext context;
+  int32_t dummy_storage = 0;
+
+  TensorUniquePtr valid_tensor = MakeShapeOnlyTensor({2, 3}, kTfLiteInt32);
+  valid_tensor->data.i32 = &dummy_storage;
+
+  size_t size_count = 0;
+  EXPECT_EQ(GetTensorElementCountAndValidateData(
+                &context, *valid_tensor, "Input", ElementCountLimit::kSizeT,
+                size_count),
+            kTfLiteOk);
+  EXPECT_EQ(size_count, 6u);
+
+  int int_count = 0;
+  EXPECT_EQ(
+      CheckedTensorElementsAndData(&context, *valid_tensor, "Input", int_count),
+      kTfLiteOk);
+  EXPECT_EQ(int_count, 6);
+  EXPECT_EQ(ValidateTensorElementsAndData(&context, *valid_tensor, "Input"),
+            kTfLiteOk);
+
+  // Empty tensors are allowed to have a null data buffer.
+  TensorUniquePtr empty_tensor = MakeShapeOnlyTensor({2, 0, 3}, kTfLiteInt32);
+  EXPECT_EQ(
+      CheckedTensorElementsAndData(&context, *empty_tensor, "Empty", int_count),
+      kTfLiteOk);
+  EXPECT_EQ(int_count, 0);
+
+  // Non-empty tensors with null data must be rejected.
+  valid_tensor->data.raw = nullptr;
+  EXPECT_EQ(ValidateTensorElementsAndData(&context, *valid_tensor, "Missing"),
+            kTfLiteError);
+  EXPECT_THAT(context.error, HasSubstr("Missing data is null."));
+
+  // Invalid shapes (negative dimension) must be rejected.
+  TensorUniquePtr negative_dim_tensor =
+      MakeShapeOnlyTensor({2, -1}, kTfLiteInt32);
+  negative_dim_tensor->data.i32 = &dummy_storage;
+  EXPECT_EQ(
+      ValidateTensorElementsAndData(&context, *negative_dim_tensor, "BadShape"),
+      kTfLiteError);
+  EXPECT_THAT(
+      context.error,
+      HasSubstr("BadShape tensor shape is invalid or size overflowed."));
+
+  // Element counts exceeding INT_MAX pass with kSizeT (on 64-bit platforms)
+  // and fail with kInt / CheckedTensorElementsAndData.
+  if constexpr (sizeof(size_t) >= 8) {
+    TensorUniquePtr large_tensor =
+        MakeShapeOnlyTensor({65536, 65536}, kTfLiteInt32);
+    large_tensor->data.i32 = &dummy_storage;
+    EXPECT_EQ(ValidateTensorElementsAndData(&context, *large_tensor, "Large",
+                                            ElementCountLimit::kSizeT),
+              kTfLiteOk);
+    EXPECT_EQ(ValidateTensorElementsAndData(&context, *large_tensor, "Large",
+                                            ElementCountLimit::kInt),
+              kTfLiteError);
+    EXPECT_EQ(CheckedTensorElementsAndData(&context, *large_tensor, "Large",
+                                           int_count),
+              kTfLiteError);
+  }
+}
+
+}  // namespace
+}  // namespace kernel_utils
+}  // namespace tflite

From ab0c8f6bad438c45812773e88132ad84c2362d13 Mon Sep 17 00:00:00 2001
From: Ce Zheng CA <zce@google.com>
Date: Tue, 13 Oct 2026 22:12:31 -0700
Subject: [PATCH 0740/1000] [XLA:CPU] Gate PowerPC/RISC-V LLVM backends and
 ir_compiler 32-bit ARM target init on target/cross config settings

Update `if_llvm_powerpc_available` and `if_llvm_riscv_available` in `xla/tsl/platform/google/build_config_root.bzl` to select on `xla/tsl:ppc64le_or_cross` and `xla/tsl:riscv64_or_cross` (matching `xla/tsl/platform/default/build_config_root.bzl`), and gate `XLA_LLVM_AARCH32_AVAILABLE=1` in `xla/backends/cpu/codegen:ir_compiler` on `xla/tsl:aarch32_or_cross`.

Keep `if_llvm_aarch32_available` returning `then` in `xla/tsl/platform/google/build_config_root.bzl` so host AOT compilation tools (`tfcompile` and `xla_compile_lib`) built in `cfg = "exec"` without `--define=with_cross_compiler_support=true` continue to link `ARMCodeGen`/`ARMAsmParser` and register the 32-bit ARM target in `llvm::TargetRegistry` before `IrCompiler::InferTargetMachine` is invoked when cross-compiling `armv7` targets.

Meanwhile, gating `XLA_LLVM_AARCH32_AVAILABLE=1` on `aarch32_or_cross` in `ir_compiler` (matching `jit_compiler` and `cpu_compiler_pure`, neither of which links `ARMCodeGen` on non-ARM hosts) and gating PowerPC/RISC-V on their target/cross config settings avoids pulling ~6.54 MB of unused 32-bit ARM, PowerPC, and RISC-V LLVM codegen backends into x86_64/AArch64 XLA:CPU binaries, with zero runtime CPU or memory impact on active target architectures.

PiperOrigin-RevId: 999114222
---
 third_party/xla/xla/backends/cpu/codegen/BUILD | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/backends/cpu/codegen/BUILD b/third_party/xla/xla/backends/cpu/codegen/BUILD
index 34f52742488..a3759df9e0c 100644
--- a/third_party/xla/xla/backends/cpu/codegen/BUILD
+++ b/third_party/xla/xla/backends/cpu/codegen/BUILD
@@ -18,7 +18,6 @@ load("//xla:xla.default.bzl", "xla_cc_test", "xla_internal")
 load("//xla/tsl:tsl.bzl", "internal_visibility", "tsl_copts")
 load(
     "//xla/tsl/platform:build_config_root.bzl",
-    "if_llvm_aarch32_available",
     "if_llvm_aarch64_available",
     "if_llvm_powerpc_available",
     "if_llvm_riscv_available",
@@ -118,7 +117,14 @@ cc_library(
     name = "ir_compiler",
     srcs = ["ir_compiler.cc"],
     hdrs = ["ir_compiler.h"],
-    defines = if_llvm_aarch32_available(["XLA_LLVM_AARCH32_AVAILABLE=1"]) + if_llvm_aarch64_available([
+    # TODO(xla-cpu): Have users who need a specific target architecture include
+    # and initialize that LLVM backend directly.
+    defines = select({
+        "//xla/tsl:aarch32_or_cross": [
+            "XLA_LLVM_AARCH32_AVAILABLE=1",
+        ],
+        "//conditions:default": [],
+    }) + if_llvm_aarch64_available([
         "XLA_LLVM_AARCH64_AVAILABLE=1",
     ]) + if_llvm_powerpc_available([
         "XLA_LLVM_POWERPC_AVAILABLE=1",

From 580ce342f9abba346d0bd1b6eafac7b4ade6d934 Mon Sep 17 00:00:00 2001
From: Changming Sun <changmingsun@google.com>
Date: Tue, 13 Oct 2026 22:14:51 -0700
Subject: [PATCH 0741/1000] Harden sub-byte tensor byte calculation in
 tflite::BytesRequired.

PiperOrigin-RevId: 999115144
---
 tensorflow/lite/util.cc      | 61 +++++++++++++++++++++++++-----------
 tensorflow/lite/util_test.cc | 61 ++++++++++++++++++++++++++++++++++++
 2 files changed, 103 insertions(+), 19 deletions(-)

diff --git a/tensorflow/lite/util.cc b/tensorflow/lite/util.cc
index 28b78d92984..8b54b891f18 100644
--- a/tensorflow/lite/util.cc
+++ b/tensorflow/lite/util.cc
@@ -256,27 +256,50 @@ TfLiteStatus BytesRequired(TfLiteType type, const int* dims, size_t dims_size,
       context_, MultiplyAndCheckOverflow(type_size, count, bytes) == kTfLiteOk,
       "BytesRequired number of bytes overflowed.\n");
 
-  // GetSizeOfType doesn't work for kTfLiteInt4 due to it having 2 values packed
-  // into 1 byte so the output of GetSizeOfType is the same as int8 aka 1 byte.
-  // Thus the required bytes must be divided by half after everything for int4.
-  if (type == kTfLiteInt4 || type == kTfLiteUInt4) {
-    *bytes = (*bytes + 1) / 2;
-  } else if (type == kTfLiteInt2) {
-    // For kTfLiteInt2, 4 elements are packed into a single byte.
-    // The '*bytes' variable at this point holds the total number of elements,
-    // because GetSizeOfType returns sizeof(int8_t) for each Int2 element.
-    // To get the actual number of bytes needed for the packed representation,
-    // we need to divide the total number of elements by 4.
-    // The expression `(*bytes + 3) / 4` implements integer division with
-    // ceiling, ensuring that we allocate enough bytes to store all elements.
-    // For example:
-    // 1 element: (1 + 3) / 4 = 1 byte
-    // 4 elements: (4 + 3) / 4 = 1 byte
-    // 5 elements: (5 + 3) / 4 = 2 bytes
-    *bytes = (*bytes + 3) / 4;
+  switch (type) {
+    case kTfLiteInt4:
+    case kTfLiteUInt4:
+      // GetSizeOfType doesn't work for kTfLiteInt4/kTfLiteUInt4 due to 2 values
+      // being packed into 1 byte, so the output of GetSizeOfType is the same as
+      // int8 (1 byte). Divide by 2 with ceiling rounding without overflowing
+      // size_t when '*bytes' is near SIZE_MAX.
+      *bytes = (*bytes / 2) + (*bytes % 2);
+      return kTfLiteOk;
+    case kTfLiteInt2:
+      // For kTfLiteInt2, 4 elements are packed into a single byte.
+      // The '*bytes' variable at this point holds the total number of elements,
+      // because GetSizeOfType returns sizeof(int8_t) for each Int2 element.
+      // To get the actual number of bytes needed for the packed representation,
+      // we need to divide the total number of elements by 4 with ceiling
+      // rounding without overflowing size_t when '*bytes' is near SIZE_MAX.
+      *bytes = (*bytes / 4) + (*bytes % 4 != 0 ? 1 : 0);
+      return kTfLiteOk;
+    case kTfLiteFloat32:
+    case kTfLiteInt32:
+    case kTfLiteUInt32:
+    case kTfLiteUInt8:
+    case kTfLiteInt64:
+    case kTfLiteUInt64:
+    case kTfLiteBool:
+    case kTfLiteComplex64:
+    case kTfLiteComplex128:
+    case kTfLiteUInt16:
+    case kTfLiteInt16:
+    case kTfLiteInt8:
+    case kTfLiteFloat16:
+    case kTfLiteBFloat16:
+    case kTfLiteFloat64:
+    case kTfLiteFloat8E4M3FN:
+    case kTfLiteFloat8E5M2:
+      return kTfLiteOk;
+    case kTfLiteNoType:
+    case kTfLiteString:
+    case kTfLiteResource:
+    case kTfLiteVariant:
+      return kTfLiteError;
   }
 
-  return kTfLiteOk;
+  return kTfLiteError;
 }
 
 #ifndef TF_LITE_STATIC_MEMORY
diff --git a/tensorflow/lite/util_test.cc b/tensorflow/lite/util_test.cc
index d0716b0b979..80e116cbe0e 100644
--- a/tensorflow/lite/util_test.cc
+++ b/tensorflow/lite/util_test.cc
@@ -213,6 +213,67 @@ TEST(FourBitTest, BytesRequiredOdd) {
   ASSERT_EQ(required_bytes_four_bit, 3);
 }
 
+void SilentReportError(TfLiteContext*, const char*, ...) {}
+
+TEST(FourBitTest, BytesRequiredSizeMaxBoundaryDoesNotWrapToZero) {
+  TfLiteContext context{};
+  context.ReportError = SilentReportError;
+
+  // Fermat prime factorization of SIZE_MAX on 64-bit (2^64 - 1) and 32-bit
+  // (2^32 - 1) platforms. Every dimension fits in a positive 32-bit signed int.
+#if SIZE_MAX == UINT64_MAX
+  constexpr int kSizeMaxDims[] = {3, 5, 17, 257, 641, 65537, 6700417};
+  // 2^64 - 2 = 2 * 7 * 7 * 73 * 127 * 337 * 92737 * 649657.
+  constexpr int kSizeMaxMinusOneDims[] = {14, 7, 73, 127, 337, 92737, 649657};
+#else
+  constexpr int kSizeMaxDims[] = {3, 5, 17, 257, 65537};
+  // 2^32 - 2 = 2 * (2^31 - 1) = 2 * 2147483647.
+  constexpr int kSizeMaxMinusOneDims[] = {2, 2147483647};
+#endif
+  constexpr size_t kNumDims = sizeof(kSizeMaxDims) / sizeof(kSizeMaxDims[0]);
+
+  for (TfLiteType type : {kTfLiteInt4, kTfLiteUInt4}) {
+    size_t required_bytes = 0;
+    ASSERT_EQ(tflite::BytesRequired(type, kSizeMaxDims, kNumDims,
+                                    &required_bytes, &context),
+              kTfLiteOk);
+    EXPECT_EQ(required_bytes, (std::numeric_limits<size_t>::max() / 2) + 1);
+  }
+
+  {
+    size_t required_bytes = 0;
+    ASSERT_EQ(tflite::BytesRequired(kTfLiteInt2, kSizeMaxDims, kNumDims,
+                                    &required_bytes, &context),
+              kTfLiteOk);
+    EXPECT_EQ(required_bytes, (std::numeric_limits<size_t>::max() / 4) + 1);
+  }
+
+  {
+    size_t required_bytes = 0;
+    constexpr size_t kNumMinusOneDims =
+        sizeof(kSizeMaxMinusOneDims) / sizeof(kSizeMaxMinusOneDims[0]);
+    ASSERT_EQ(
+        tflite::BytesRequired(kTfLiteInt2, kSizeMaxMinusOneDims,
+                              kNumMinusOneDims, &required_bytes, &context),
+        kTfLiteOk);
+    EXPECT_EQ(required_bytes,
+              ((std::numeric_limits<size_t>::max() - 1) / 4) + 1);
+  }
+}
+
+TEST(BytesRequiredTest, RejectsUnsupportedTypes) {
+  TfLiteContext context{};
+  context.ReportError = SilentReportError;
+  constexpr int kDims[] = {2, 4};
+
+  for (TfLiteType type :
+       {kTfLiteNoType, kTfLiteString, kTfLiteResource, kTfLiteVariant}) {
+    size_t required_bytes = 0;
+    EXPECT_EQ(tflite::BytesRequired(type, kDims, 2, &required_bytes, &context),
+              kTfLiteError);
+  }
+}
+
 TEST(TestMakeUniqueTensor, Valid) {
   TensorUniquePtr t = BuildTfLiteTensor(kTfLiteInt32, {2, 3}, kTfLiteDynamic);
   ASSERT_NE(t.get(), nullptr);

From ead98bdbe84fb434dfaaa6aa10d9caba6bc4db07 Mon Sep 17 00:00:00 2001
From: mohammed adib <adib@bugqore.com>
Date: Wed, 14 Oct 2026 11:32:33 +0530
Subject: [PATCH 0742/1000] log empty metadata skip at WARNING

Signed-off-by: mohammed adib <adib@bugqore.com>
---
 tensorflow/core/data/serialization_utils.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/data/serialization_utils.cc b/tensorflow/core/data/serialization_utils.cc
index 7cb62a52e0c..44710634f66 100644
--- a/tensorflow/core/data/serialization_utils.cc
+++ b/tensorflow/core/data/serialization_utils.cc
@@ -193,7 +193,7 @@ VariantTensorDataReader::VariantTensorDataReader(
     d->get_metadata(&metadata);
     auto keys = str_util::Split(metadata, kDelimiter, str_util::SkipEmpty());
     if (keys.empty()) {
-      VLOG(1) << "Found empty metadata in VariantTensorData; skipping.";
+      LOG(WARNING) << "Found empty metadata in VariantTensorData; skipping.";
       continue;
     }
     const std::string name = keys[0];

From 5518e4d54569d3f1b0dd49a9bb47dbf48d7f6870 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 11:37:02 +0530
Subject: [PATCH 0743/1000] Fix testDatasetAlreadyBatchedOffsets assertion per
 review correction

The previous global_batch_size=16 fix was right, but the follow-up
shape assertion (d_tensor.shape == [16, 8, 12]) was wrong: the
dataset's non-batch dims (height/width) represent the global spatial
shape, not per-replica slice shapes, so the global DTensor is
actually [16, 4, 6]. More importantly, a bare shape check doesn't
verify the values are sharded at the correct offsets, which is what
this test exists to catch. Use assertDTensorEqual(local_images,
layout, d_tensor) instead, per review.
---
 tensorflow/dtensor/python/tests/input_util_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/dtensor/python/tests/input_util_test.py b/tensorflow/dtensor/python/tests/input_util_test.py
index 6808bfede3d..72428872e9e 100644
--- a/tensorflow/dtensor/python/tests/input_util_test.py
+++ b/tensorflow/dtensor/python/tests/input_util_test.py
@@ -127,7 +127,7 @@ class DTensorDatasetTest(test_util.DTensorBaseTest):
     iterator = iter(d_dataset)
     d_tensor = next(iterator)
 
-    self.assertEqual(d_tensor.shape, [16, 8, 12])
+    self.assertDTensorEqual(local_images, layout, d_tensor)
 
   @parameterized.named_parameters(('Eager', False), ('Graph', True))
   def testRangeIteration(self, is_graph):

From 8c5344aa5018c5ec2e75786c92607011ddd1727b Mon Sep 17 00:00:00 2001
From: Alex Pivovarov <upwind@google.com>
Date: Tue, 13 Oct 2026 23:29:55 -0700
Subject: [PATCH 0744/1000] [XLA:GPU] Print dtype in the select_k_exec_raft
 VLOG statement.

Include the template data type (`float`, `bfloat16`, or `uint64_t`) in the `select_k_exec_raft` `VLOG(3)` trace to make it easier to see which RAFT `select_k` specialization is invoked at runtime.

PiperOrigin-RevId: 999146551
---
 .../xla/xla/backends/gpu/runtime/BUILD        |  4 +---
 .../gpu/runtime/select_k_exec_raft_impl.h     | 24 ++++++++++++++-----
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index 0829cb916cf..c77f636dc5c 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -1553,12 +1553,10 @@ cuda_library(
         "//xla/stream_executor:device_address_allocator",
         "//xla/stream_executor:stream",
         "@com_google_absl//absl/cleanup",
-        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log",
-        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
         "@local_config_cuda//cuda:cuda_headers",
         "@raft//:raft_matrix",
     ],
diff --git a/third_party/xla/xla/backends/gpu/runtime/select_k_exec_raft_impl.h b/third_party/xla/xla/backends/gpu/runtime/select_k_exec_raft_impl.h
index 339e3cd0a4d..0d7877c65c1 100644
--- a/third_party/xla/xla/backends/gpu/runtime/select_k_exec_raft_impl.h
+++ b/third_party/xla/xla/backends/gpu/runtime/select_k_exec_raft_impl.h
@@ -18,17 +18,16 @@ limitations under the License.
 
 #include <cstddef>
 #include <cstdint>
+#include <exception>
 #include <memory>
 #include <optional>
-#include <utility>
+#include <type_traits>
 
 #include "absl/cleanup/cleanup.h"
-#include "absl/container/flat_hash_map.h"
-#include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
-#include "absl/status/status_macros.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
 #include "third_party/gpus/cuda/include/cuda_bf16.h"
 #include "raft/core/device_mdspan.hpp"
 #include "raft/core/mdspan_types.hpp"
@@ -136,6 +135,19 @@ struct RaftStreamResource : public se::Stream::Resource {
 template <typename T>
 SelectAlgo choose_select_k_algorithm(uint32_t rows, uint32_t cols, uint32_t k);
 
+template <typename T>
+constexpr absl::string_view SelectKTypeName() {
+  if constexpr (std::is_same_v<T, float>) {
+    return "float";
+  } else if constexpr (std::is_same_v<T, nv_bfloat16>) {
+    return "bfloat16";
+  } else if constexpr (std::is_same_v<T, std::uint64_t>) {
+    return "uint64_t";
+  } else {
+    return "unknown";
+  }
+}
+
 }  // namespace raft_internal
 
 // Host-side entry point for raft select_k
@@ -150,8 +162,8 @@ absl::Status select_k_exec(int device_ordinal,
                            se::DeviceAddressBase scratch_buffer) {
   // Pick the most suitable algorithm
   SelectAlgo algo = raft_internal::choose_select_k_algorithm<T>(batch, n, k);
-  VLOG(3) << "select_k_exec_raft: "
-          << "device_ordinal: " << device_ordinal << ", "
+  VLOG(3) << "select_k_exec_raft<" << raft_internal::SelectKTypeName<T>()
+          << ">: device_ordinal: " << device_ordinal << ", "
           << "allocator: " << allocator << ", "
           << "stream: " << stream << ", "
           << "data_in: " << data_in.opaque() << " (" << data_in.size() << "B)"

From cc588e79f0566322e8ac93a343a04922c1e29081 Mon Sep 17 00:00:00 2001
From: Kasper Nielsen <kasper0406@gmail.com>
Date: Tue, 13 Oct 2026 23:48:02 -0700
Subject: [PATCH 0745/1000] PR #47670: Preserve chlo.erfc as a
 stablehlo.composite
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/47670

#ai-generated

`chlo.erf` and `chlo.top_k` already survive StableHLO serialization as `stablehlo.composite` ops. `chlo.erfc` does not, so it is expanded into its full polynomial approximation at export time and downstream StableHLO consumers cannot recognize it.

Where I hit this: exporting `jax.nn.gelu(x, approximate=False)` (which lowers through `chlo.erfc`) via `jax.export`. With jax/jaxlib 0.9.1:

```
jax.scipy.special.erf          -> stablehlo.composite "chlo.erf" ...
jax.nn.gelu(approximate=False) -> ~90 StableHLO ops, no composite
```

This blocks the stablehlo-coreml converter (and any other StableHLO consumer) from recognizing erfc; it has to match a ~70-op polynomial or give up.

This adds `chlo::ErfcOp` to `ChloPreserveHighLevelOpsPass` and `ChloRecomposeOpsPass` so erfc round-trips exactly like erf. `chlo.erfc` has no attributes, so the existing builtin-attribute passthrough covers it.

On the compile path there is no `MHLO_ErfcOp`, so the recomposed `chlo.erfc` is left alone by `ChloLegalizeToHighLevelMhloPass` and expanded by `ChloLegalizeToStablehloPass` — the same path `chlo.acosh`/`cosh`/`sinh` already take. Verified end to end: the composite lowers to the identical 92-op expansion produced today, so numerics are unchanged.

Tests: erfc cases added to `chlo_preserve_high_level_ops.mlir` and `chlo_recompose_ops.mlir`, mirroring the existing erf cases. The preserve test's `preserve | recompose | diff` RUN line covers the round trip.

Copybara import of the project:

--
a081467164c0a3d8089bd0f78cdbe36eb25278dc by Kasper Nielsen <kasper0406@gmail.com>:

Preserve chlo.erfc as a stablehlo.composite

chlo.erf and chlo.top_k already survive StableHLO serialization as
stablehlo.composite ops, letting downstream consumers pattern-match them
instead of a large expanded polynomial. chlo.erfc did not, so exporting
jax.nn.gelu(x, approximate=False) through jax.export produced ~90
StableHLO ops (an exp, two polynomial ladders and four compare/select
pairs) with no way to recover the original erfc.

Add chlo::ErfcOp to ChloPreserveHighLevelOpsPass and ChloRecomposeOpsPass
so it round-trips the same way chlo.erf does. chlo.erfc has no
attributes, so the existing builtin attribute passthrough covers it and
no serialization logic is needed.

On the compile path the recomposed chlo.erfc has no high-level MHLO
counterpart (there is no MHLO_ErfcOp), so it is left alone by
ChloLegalizeToHighLevelMhloPass and expanded by
ChloLegalizeToStablehloPass. This is the path chlo.acosh, chlo.cosh and
chlo.sinh already take, since their direct MHLO lowerings are disabled by
default. Numerics are unchanged.

Merging this change closes #47670

PiperOrigin-RevId: 999154306
---
 .../transforms/chlo_preserve_high_level_ops.cpp   |  1 +
 .../transforms/chlo_recompose_ops.cpp             |  1 +
 .../chlo_preserve_high_level_ops.mlir             |  9 +++++++++
 .../tests/stablehlo_ext/chlo_recompose_ops.mlir   | 15 +++++++++++++++
 4 files changed, 26 insertions(+)

diff --git a/third_party/xla/xla/mlir_hlo/stablehlo_ext/transforms/chlo_preserve_high_level_ops.cpp b/third_party/xla/xla/mlir_hlo/stablehlo_ext/transforms/chlo_preserve_high_level_ops.cpp
index 2bb82cc8321..bd2ceaf82da 100644
--- a/third_party/xla/xla/mlir_hlo/stablehlo_ext/transforms/chlo_preserve_high_level_ops.cpp
+++ b/third_party/xla/xla/mlir_hlo/stablehlo_ext/transforms/chlo_preserve_high_level_ops.cpp
@@ -235,6 +235,7 @@ struct ChloPreserveHighLevelOpsPass
                  ChloOpToCompositePattern<chlo::CoshOp>,
                  ChloOpToCompositePattern<chlo::SinhOp>,
                  ChloOpToCompositePattern<chlo::ErfOp>,
+                 ChloOpToCompositePattern<chlo::ErfcOp>,
                  ChloOpToCompositePattern<chlo::MulhiOp>,
                  ChloOpToCompositePattern<chlo::RaggedDotOp>,
                  ChloOpToCompositePattern<chlo::TopKOp>,
diff --git a/third_party/xla/xla/mlir_hlo/stablehlo_ext/transforms/chlo_recompose_ops.cpp b/third_party/xla/xla/mlir_hlo/stablehlo_ext/transforms/chlo_recompose_ops.cpp
index f418e0f11b8..18e6ea0119f 100644
--- a/third_party/xla/xla/mlir_hlo/stablehlo_ext/transforms/chlo_recompose_ops.cpp
+++ b/third_party/xla/xla/mlir_hlo/stablehlo_ext/transforms/chlo_recompose_ops.cpp
@@ -483,6 +483,7 @@ struct ChloRecomposeOpsPass
       ChloOpRecomposePattern<chlo::CoshOp>,
       ChloOpRecomposePattern<chlo::SinhOp>,
       ChloOpRecomposePattern<chlo::ErfOp>,
+      ChloOpRecomposePattern<chlo::ErfcOp>,
       ChloOpRecomposePattern<chlo::MulhiOp>,
       ChloOpRecomposePattern<chlo::RaggedDotOp>,
       ChloOpRecomposePattern<chlo::ScanOp>,
diff --git a/third_party/xla/xla/mlir_hlo/tests/stablehlo_ext/chlo_preserve_high_level_ops.mlir b/third_party/xla/xla/mlir_hlo/tests/stablehlo_ext/chlo_preserve_high_level_ops.mlir
index 3beb0510f8a..feb6987b9c3 100644
--- a/third_party/xla/xla/mlir_hlo/tests/stablehlo_ext/chlo_preserve_high_level_ops.mlir
+++ b/third_party/xla/xla/mlir_hlo/tests/stablehlo_ext/chlo_preserve_high_level_ops.mlir
@@ -80,6 +80,15 @@ func.func @erf_preserve(%arg0: tensor<3x20x20xbf16>) -> tensor<?x20x20xbf16> {
 
 // -----
 
+// CHECK-LABEL: func @erfc_preserve
+func.func @erfc_preserve(%arg0: tensor<3x20x20xbf16>) -> tensor<?x20x20xbf16> {
+  // CHECK: stablehlo.composite "chlo.erfc" %arg0 {decomposition = @chlo.erfc.impl, version = 1 : i32}
+  %0 = chlo.erfc %arg0 : tensor<3x20x20xbf16> -> tensor<?x20x20xbf16>
+  return %0 : tensor<?x20x20xbf16>
+}
+
+// -----
+
 // CHECK-LABEL: func @acosh_preserve
 func.func @acosh_preserve(%arg0: tensor<3x20x20xbf16>) -> tensor<?x20x20xbf16> {
   // CHECK: stablehlo.composite "chlo.acosh" %arg0 {decomposition = @chlo.acosh.impl, version = 1 : i32}
diff --git a/third_party/xla/xla/mlir_hlo/tests/stablehlo_ext/chlo_recompose_ops.mlir b/third_party/xla/xla/mlir_hlo/tests/stablehlo_ext/chlo_recompose_ops.mlir
index 5c83ca5e9a4..b9971b057e2 100644
--- a/third_party/xla/xla/mlir_hlo/tests/stablehlo_ext/chlo_recompose_ops.mlir
+++ b/third_party/xla/xla/mlir_hlo/tests/stablehlo_ext/chlo_recompose_ops.mlir
@@ -32,6 +32,21 @@ func.func private @chlo.erf.impl(%arg0: tensor<3x20x20xbf16>) -> tensor<?x20x20x
 
 // -----
 
+// CHECK-LABEL: func @erfc_recompose_composite
+func.func @erfc_recompose_composite(%arg0: tensor<3x20x20xbf16>) -> tensor<?x20x20xbf16> {
+  // CHECK-NEXT: chlo.erfc
+  // CHECK-NOT: stablehlo.composite
+  %0 = stablehlo.composite "chlo.erfc" %arg0 {decomposition = @chlo.erfc.impl, version = 1 : i32} : (tensor<3x20x20xbf16>) -> tensor<?x20x20xbf16>
+  return %0 : tensor<?x20x20xbf16>
+}
+// CHECK-NOT: @chlo.erfc.impl
+func.func private @chlo.erfc.impl(%arg0: tensor<3x20x20xbf16>) -> tensor<?x20x20xbf16> {
+  %0 = chlo.erfc %arg0 : tensor<3x20x20xbf16> -> tensor<?x20x20xbf16>
+  return %0 : tensor<?x20x20xbf16>
+}
+
+// -----
+
 // CHECK-LABEL: func @acosh_recompose_composite
 func.func @acosh_recompose_composite(%arg0: tensor<3x20x20xbf16>) -> tensor<?x20x20xbf16> {
   // CHECK-NEXT: chlo.acosh

From df8091a9a02f96f5941a32b6fede2a9f994732f9 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 12:56:27 +0530
Subject: [PATCH 0746/1000] Reorder _validate_input checks to fix
 drop_remainder regression
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Moving is_fully_defined() to the start of the per-element validation
loop (in an earlier round) short-circuited the dataset_already_batched
branch below it for datasets batched with drop_remainder=False, whose
batch dim is None — breaking
testIterOnBatchedDatasetFailsNoDropLastBatch, which expects the more
specific "Ensure drop_remainder=True..." message.

Restore the original elem_spec.shape.rank is None check and the
batch_dim_size is None / drop_remainder error at their original
positions, and run is_fully_defined() after the dataset_already_batched
block instead, per review.
---
 tensorflow/dtensor/python/input_util.py | 17 ++++++++++++-----
 1 file changed, 12 insertions(+), 5 deletions(-)

diff --git a/tensorflow/dtensor/python/input_util.py b/tensorflow/dtensor/python/input_util.py
index bb7f35b2f56..da0c6e02e02 100644
--- a/tensorflow/dtensor/python/input_util.py
+++ b/tensorflow/dtensor/python/input_util.py
@@ -226,9 +226,9 @@ def _validate_input(flattened_layouts: Sequence[layout_lib.Layout],
   first_elem_shape = flattened_elem_spec[0].shape
 
   for layout, elem_spec in zip(flattened_layouts, flattened_elem_spec):
-    if not elem_spec.shape.is_fully_defined():
+    if elem_spec.shape.rank is None:
       raise ValueError(
-          'Dataset element shape must be fully defined, got spec %s.' %
+          'Dataset element shape must have a valid rank, got spec %s.' %
           elem_spec)
 
     # Check that layout's rank matches the element's rank. If dataset is not yet
@@ -246,15 +246,22 @@ def _validate_input(flattened_layouts: Sequence[layout_lib.Layout],
 
     if dataset_already_batched:
       # Check that the batch dimension size of all dataset elements match.
-      # `batch_dim_size` is guaranteed to not be None here because the
-      # `is_fully_defined()` check above already ensures every dimension,
-      # including the batch dimension, has a concrete size.
       batch_dim_size = first_elem_shape.as_list()[0]
+      if batch_dim_size is None:
+        raise ValueError(
+            ('Size of batch dimension of element spec %s is None. Ensure '
+             'drop_remainder=True when batching the dataset.') % elem_spec)
+
       if elem_spec.shape.as_list()[0] != batch_dim_size:
         raise ValueError(
             ('Size of batch dimension of element spec %s does not match '
              'expected size %d.') % (elem_spec, batch_dim_size))
 
+    if not elem_spec.shape.is_fully_defined():
+      raise ValueError(
+          'Dataset element shape must be fully defined, got spec %s.' %
+          elem_spec)
+
 
 def _shard_counts(layout: layout_lib.Layout,
                   batch_dim: Optional[str] = None) -> List[int]:

From 7b4ca340cfa8d42b4ff8a60b87c5061e20d4cca2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Eusebio=20Dur=C3=A1n=20Monta=C3=B1a?= <eusebiodm@google.com>
Date: Wed, 14 Oct 2026 00:58:55 -0700
Subject: [PATCH 0747/1000] Lazily initialize `YnnThreadpool` to decouple
 `Thunk` from YNNPACK

`Thunk::ExecuteParams` eagerly created a `YnnThreadpool`, which forced `:thunk` and `nanort_executable` to depend on YNNPACK, which also pulls in Perfetto, even when no `YnnFusionThunk` is used.

Instead, we forward-declare `ynn_threadpool` in `thunk.h` and store it inside `Thunk::YnnParams` using a `std::unique_ptr` with a custom deleter (`void (*)(ynn_threadpool*)`) and an `absl::once_flag`. `YnnFusionThunk::Execute` now lazily creates the threadpool on the first YNN fusion call in an execution, avoiding the overhead and binary size cost when YNNPACK isn't used without needing a global factory registry.

Also change `CreateYnnThreadpool` to return `YnnThreadpool` directly instead of `absl::StatusOr<YnnThreadpool>`, since it can't fail and this avoids a `CHECK_OK` during lazy initialization.

PiperOrigin-RevId: 999184995
---
 third_party/xla/xla/backends/cpu/nanort/BUILD |  2 --
 .../backends/cpu/nanort/nanort_executable.cc  | 23 ++-----------
 .../xla/xla/backends/cpu/runtime/BUILD        |  4 +--
 .../xla/xla/backends/cpu/runtime/thunk.cc     | 12 -------
 .../xla/xla/backends/cpu/runtime/thunk.h      | 32 +++++++++++++------
 .../xla/backends/cpu/runtime/ynnpack/BUILD    |  3 +-
 .../cpu/runtime/ynnpack/ynn_fusion_thunk.cc   |  9 ++++--
 .../runtime/ynnpack/ynn_fusion_thunk_test.cc  | 24 ++++++--------
 .../cpu/runtime/ynnpack/ynn_interop.cc        |  2 +-
 .../cpu/runtime/ynnpack/ynn_interop.h         |  8 ++---
 .../cpu/runtime/ynnpack/ynn_threadpool.cc     | 14 +++-----
 .../cpu/runtime/ynnpack/ynn_threadpool.h      |  7 ++--
 third_party/xla/xla/pjrt/cpu/cpu_client.cc    |  6 ----
 .../xla/xla/service/cpu/cpu_executable.cc     |  7 ----
 14 files changed, 56 insertions(+), 97 deletions(-)

diff --git a/third_party/xla/xla/backends/cpu/nanort/BUILD b/third_party/xla/xla/backends/cpu/nanort/BUILD
index f879f33711a..10f96d816e3 100644
--- a/third_party/xla/xla/backends/cpu/nanort/BUILD
+++ b/third_party/xla/xla/backends/cpu/nanort/BUILD
@@ -124,8 +124,6 @@ cc_library(
         "//xla/backends/cpu/runtime:function_library",
         "//xla/backends/cpu/runtime:thread_pool_task_runner",
         "//xla/backends/cpu/runtime:thunk",
-        "//xla/backends/cpu/runtime/ynnpack:ynn_interop",
-        "//xla/backends/cpu/runtime/ynnpack:ynn_threadpool",
         "//xla/ffi:execution_context",
         "//xla/hlo/ir:hlo",
         "//xla/runtime:device_id",
diff --git a/third_party/xla/xla/backends/cpu/nanort/nanort_executable.cc b/third_party/xla/xla/backends/cpu/nanort/nanort_executable.cc
index 01591e370cb..e1cf35472a7 100644
--- a/third_party/xla/xla/backends/cpu/nanort/nanort_executable.cc
+++ b/third_party/xla/xla/backends/cpu/nanort/nanort_executable.cc
@@ -34,8 +34,6 @@ limitations under the License.
 #include "xla/backends/cpu/runtime/function_library.h"
 #include "xla/backends/cpu/runtime/thread_pool_task_runner.h"
 #include "xla/backends/cpu/runtime/thunk.h"
-#include "xla/backends/cpu/runtime/ynnpack/ynn_interop.h"
-#include "xla/backends/cpu/runtime/ynnpack/ynn_threadpool.h"
 #include "xla/executable_run_options.h"
 #include "xla/ffi/execution_context.h"
 #include "xla/hlo/ir/hlo_module.h"
@@ -406,8 +404,7 @@ tsl::AsyncValueRef<NanoRtExecutable::ExecuteEvent> NanoRtExecutable::Execute(
   struct ExecutionContext {
     ExecutionContext(cpu::BufferAllocations::Buffers buffers,
                      FunctionLibrary* function_library,
-                     const ExecuteOptions& options,
-                     std::optional<Thunk::YnnParams> ynn_params)
+                     const ExecuteOptions& options)
         : allocations(std::move(buffers)),
           execute_params(Thunk::ExecuteParams{function_library, &allocations,
                                               /*xfeed=*/nullptr,
@@ -420,19 +417,15 @@ tsl::AsyncValueRef<NanoRtExecutable::ExecuteEvent> NanoRtExecutable::Execute(
               options.device_assignment(), /*collectives=*/nullptr),
           custom_call_execute_params(
               RunId(options.launch_id()), options.local_device_id().value(),
-              options.intra_op_thread_pool(), options.ffi_context()),
-          ynn_params(std::move(ynn_params)) {
+              options.intra_op_thread_pool(), options.ffi_context()) {
       execute_params.collective_params = &collective_execute_params;
       execute_params.custom_call_params = &custom_call_execute_params;
-      execute_params.ynn_params =
-          this->ynn_params.has_value() ? &*this->ynn_params : nullptr;
     }
 
     cpu::BufferAllocations allocations;
     Thunk::ExecuteParams execute_params;
     Thunk::CollectiveExecuteParams collective_execute_params;
     Thunk::CustomCallExecuteParams custom_call_execute_params;
-    std::optional<Thunk::YnnParams> ynn_params;
   };
 
   // Do a heap allocation if we're running with a thread pool, using
@@ -441,18 +434,8 @@ tsl::AsyncValueRef<NanoRtExecutable::ExecuteEvent> NanoRtExecutable::Execute(
   // allocation when it is not required.
   if (options.intra_op_thread_pool() || options.ffi_context() ||
       options.device_assignment()) {
-    // Prepare for executing YNNPACK fusions. Without a YNN threadpool, YNNPACK
-    // fusions run single-threaded in the caller thread.
-    std::optional<Thunk::YnnParams> ynn_params;
-    if (executable->has_ynn_fusions() && options.intra_op_thread_pool()) {
-      ABSL_ASSIGN_OR_RETURN(YnnThreadpool ynn_threadpool,
-                       CreateYnnThreadpool(options.intra_op_thread_pool()));
-      ynn_params.emplace(std::move(ynn_threadpool));
-    }
-
     auto execution_context = std::make_unique<ExecutionContext>(
-        std::move(buffers), executable->function_library(), options,
-        std::move(ynn_params));
+        std::move(buffers), executable->function_library(), options);
 
     auto execute_event =
         executable->thunks().Execute(execution_context->execute_params);
diff --git a/third_party/xla/xla/backends/cpu/runtime/BUILD b/third_party/xla/xla/backends/cpu/runtime/BUILD
index d1eaa64f95d..14d1f559594 100644
--- a/third_party/xla/xla/backends/cpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/cpu/runtime/BUILD
@@ -193,8 +193,6 @@ cc_library(
         "//xla:executable_run_options",
         "//xla/backends/cpu/collectives:cpu_collectives",
         "//xla/backends/cpu/collectives:in_process_collectives",
-        "//xla/backends/cpu/runtime/ynnpack:ynn_interop",
-        "//xla/backends/cpu/runtime/ynnpack:ynn_threadpool",
         "//xla/ffi:execution_context",
         "//xla/runtime:buffer_use",
         "//xla/runtime:device_id",
@@ -206,6 +204,7 @@ cc_library(
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:logging",
         "//xla/tsl/platform:statusor",
+        "@com_google_absl//absl/base",
         "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/functional:function_ref",
@@ -1376,7 +1375,6 @@ cc_library(
         ":while_thunk",
         "//xla:shape_util",
         "//xla:util",
-        "//xla/backends/cpu/runtime/ynnpack:ynn_interop",
         "//xla/runtime:work_group",
         "//xla/service:collective_ops_utils",
         "//xla/service:shaped_slice",
diff --git a/third_party/xla/xla/backends/cpu/runtime/thunk.cc b/third_party/xla/xla/backends/cpu/runtime/thunk.cc
index 581f5ba4749..d01cbd3ba7c 100644
--- a/third_party/xla/xla/backends/cpu/runtime/thunk.cc
+++ b/third_party/xla/xla/backends/cpu/runtime/thunk.cc
@@ -30,8 +30,6 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "xla/backends/cpu/collectives/cpu_collectives.h"
 #include "xla/backends/cpu/collectives/in_process_collectives.h"
-#include "xla/backends/cpu/runtime/ynnpack/ynn_interop.h"
-#include "xla/backends/cpu/runtime/ynnpack/ynn_threadpool.h"
 #include "xla/executable_run_options.h"
 #include "xla/runtime/device_id.h"
 #include "xla/service/cpu/cpu_executable_run_options.h"
@@ -161,16 +159,6 @@ Thunk::CustomCallExecuteParams::CustomCallExecuteParams(
       intra_op_thread_pool(intra_op_thread_pool),
       ffi_execution_context(ffi_execution_context) {}
 
-absl::StatusOr<Thunk::YnnParams> Thunk::YnnParams::Create(
-    const ExecutableRunOptions* run_options) {
-  ABSL_ASSIGN_OR_RETURN(YnnThreadpool threadpool,
-                   CreateYnnThreadpool(run_options->intra_op_thread_pool()));
-  return YnnParams(std::move(threadpool));
-}
-
-Thunk::YnnParams::YnnParams(YnnThreadpool threadpool)
-    : threadpool(std::move(threadpool)) {}
-
 Thunk::ExecuteSession::ExecuteSession(int64_t max_workers,
                                       int64_t split_threshold)
     : lock_(std::make_shared<std::nullopt_t>(std::nullopt)),
diff --git a/third_party/xla/xla/backends/cpu/runtime/thunk.h b/third_party/xla/xla/backends/cpu/runtime/thunk.h
index a8789ab9139..649c06eba50 100644
--- a/third_party/xla/xla/backends/cpu/runtime/thunk.h
+++ b/third_party/xla/xla/backends/cpu/runtime/thunk.h
@@ -26,6 +26,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/base/call_once.h"
 #include "absl/base/no_destructor.h"
 #include "absl/container/inlined_vector.h"
 #include "absl/functional/function_ref.h"
@@ -36,8 +37,6 @@ limitations under the License.
 #include "xla/backends/cpu/runtime/buffer_allocations.h"
 #include "xla/backends/cpu/runtime/function_library.h"
 #include "xla/backends/cpu/runtime/xfeed_manager.h"
-#include "xla/backends/cpu/runtime/ynnpack/ynn_interop.h"
-#include "xla/backends/cpu/runtime/ynnpack/ynn_threadpool.h"
 #include "xla/executable_run_options.h"
 #include "xla/ffi/execution_context.h"
 #include "xla/runtime/buffer_use.h"
@@ -49,6 +48,8 @@ limitations under the License.
 #include "xla/tsl/platform/statusor.h"
 #include "tsl/platform/platform.h"
 
+struct ynn_threadpool;
+
 namespace Eigen {
 struct ThreadPoolDevice;
 }  // namespace Eigen
@@ -259,14 +260,26 @@ class Thunk {
   // YnnParams
   //===--------------------------------------------------------------------===//
 
-  // Parameters capturing all the details required for running XNNPACK fusions.
-  struct YnnParams {
-    static absl::StatusOr<YnnParams> Create(
-        const ExecutableRunOptions* run_options);
+  // Parameters capturing all the details required for running YNNPACK fusions.
+  // `ynn_threadpool` is forward-declared and lazily initialized with a
+  // function-pointer deleter so that `:thunk` doesn't need to depend on YNNPACK
+  // (avoiding binary size bloat for targets that don't use YNNPACK).
+  class YnnParams {
+   public:
+    using Threadpool =
+        std::unique_ptr<ynn_threadpool, void (*)(ynn_threadpool*)>;
 
-    YnnThreadpool threadpool = nullptr;
+    // Lazily initializes and returns the threadpool. `factory` is invoked at
+    // most once per execution.
+    const Threadpool& GetOrCreateThreadpool(
+        absl::FunctionRef<Threadpool()> factory) const {
+      absl::call_once(once_, [&] { threadpool_ = factory(); });
+      return threadpool_;
+    }
 
-    explicit YnnParams(YnnThreadpool threadpool);
+   private:
+    mutable absl::once_flag once_;
+    mutable Threadpool threadpool_ = {nullptr, nullptr};
   };
 
   //===--------------------------------------------------------------------===//
@@ -283,7 +296,6 @@ class Thunk {
     TaskRunner* task_runner = nullptr;
     CollectiveExecuteParams* collective_params = nullptr;
     CustomCallExecuteParams* custom_call_params = nullptr;
-    YnnParams* ynn_params = nullptr;
     int64_t run_id = -1;          // -1 means no run id is set.
     int64_t device_ordinal = -1;  // -1 means no device ordinal is set.
     ExecuteSession session = ExecuteSession(ExecuteSession::kMaxWorkers,
@@ -292,6 +304,8 @@ class Thunk {
 
     // Per-execution custom options.
     const CustomOptions* custom_options = nullptr;
+
+    YnnParams ynn_params;
   };
 
   // An execute event that becomes ready when all tasks are completed.
diff --git a/third_party/xla/xla/backends/cpu/runtime/ynnpack/BUILD b/third_party/xla/xla/backends/cpu/runtime/ynnpack/BUILD
index bd13aefd33a..001303db065 100644
--- a/third_party/xla/xla/backends/cpu/runtime/ynnpack/BUILD
+++ b/third_party/xla/xla/backends/cpu/runtime/ynnpack/BUILD
@@ -86,7 +86,6 @@ cc_library(
         "@XNNPACK//ynnpack:ynnpack_h",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
-        "@com_google_absl//absl/status:statusor",
         "@eigen_archive//:eigen3",
     ],
 )
@@ -97,6 +96,7 @@ cc_library(
     hdrs = ["ynn_fusion_thunk.h"],
     deps = [
         ":ynn_interop",
+        ":ynn_threadpool",
         "//xla:shape_util",
         "//xla/backends/cpu/runtime:thunk",
         "//xla/hlo/ir:hlo",
@@ -132,7 +132,6 @@ xla_cc_test(
     deps = [
         ":ynn_fusion_thunk",
         ":ynn_interop",
-        ":ynn_threadpool",
         "//xla:literal_util",
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",
diff --git a/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_fusion_thunk.cc b/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_fusion_thunk.cc
index f7408e90faf..8145368fccc 100644
--- a/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_fusion_thunk.cc
+++ b/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_fusion_thunk.cc
@@ -37,6 +37,7 @@ limitations under the License.
 #include "absl/types/span.h"
 #include "xla/backends/cpu/runtime/thunk.h"
 #include "xla/backends/cpu/runtime/ynnpack/ynn_interop.h"
+#include "xla/backends/cpu/runtime/ynnpack/ynn_threadpool.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/runtime/buffer_use.h"
 #include "xla/stream_executor/device_address.h"
@@ -284,8 +285,12 @@ YnnFusionThunk::BufferUses YnnFusionThunk::buffer_uses() const {
 }
 
 const YnnThreadpool& GetYnnThreadpool(const Thunk::ExecuteParams& params) {
-  static absl::NoDestructor<YnnThreadpool> no_threadpool(nullptr);
-  return params.ynn_params ? params.ynn_params->threadpool : *no_threadpool;
+  static absl::NoDestructor<YnnThreadpool> no_threadpool(nullptr, nullptr);
+  if (params.intra_op_threadpool == nullptr) {
+    return *no_threadpool;
+  }
+  return params.ynn_params.GetOrCreateThreadpool(
+      [&] { return CreateYnnThreadpool(params.intra_op_threadpool); });
 }
 
 tsl::AsyncValueRef<YnnFusionThunk::ExecuteEvent> YnnFusionThunk::Execute(
diff --git a/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_fusion_thunk_test.cc b/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_fusion_thunk_test.cc
index 04f901f8491..1577457818b 100644
--- a/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_fusion_thunk_test.cc
+++ b/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_fusion_thunk_test.cc
@@ -30,7 +30,6 @@ limitations under the License.
 #include "xla/backends/cpu/runtime/thunk.h"
 #include "xla/backends/cpu/runtime/thunk_testlib.h"
 #include "xla/backends/cpu/runtime/ynnpack/ynn_interop.h"
-#include "xla/backends/cpu/runtime/ynnpack/ynn_threadpool.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/literal_util.h"
 #include "xla/service/buffer_assignment.h"
@@ -160,22 +159,24 @@ TEST_P(YnnFusionThunkTest, ElementwiseAdd) {
                       reinterpret_cast<HloInstruction*>(0xDEADBEEF),
                       {lhs_arg, rhs_arg}, {out_res}, &BuildBinaryAddSubgraph));
 
-  YnnThreadpool threadpool;
-  if (use_threadpool()) {
-    TF_ASSERT_OK_AND_ASSIGN(threadpool, CreateYnnThreadpool(&device));
-  }
-  Thunk::YnnParams ynn_params(std::move(threadpool));
-
   Thunk::ExecuteParams params;
   params.buffer_allocations = &allocations;
   params.intra_op_threadpool = use_threadpool() ? &device : nullptr;
-  params.ynn_params = &ynn_params;
 
   auto execute_event = thunk->Execute(params);
   tsl::BlockUntilReady(execute_event);
   ASSERT_FALSE(execute_event.IsError()) << execute_event.GetError();
 
   EXPECT_EQ(out, LiteralUtil::CreateR1<float>({5.0, 5.0, 5.0, 5.0}));
+
+  bool factory_called = false;
+  const YnnThreadpool& threadpool =
+      params.ynn_params.GetOrCreateThreadpool([&] {
+        factory_called = true;
+        return YnnThreadpool(nullptr, nullptr);
+      });
+  EXPECT_EQ(factory_called, !use_threadpool());
+  EXPECT_EQ(threadpool != nullptr, use_threadpool());
 }
 
 TEST_P(YnnFusionThunkTest, Iota) {
@@ -200,16 +201,9 @@ TEST_P(YnnFusionThunkTest, Iota) {
                       reinterpret_cast<HloInstruction*>(0xDEADBEEF), {},
                       {out_res}, &BuildIotaSubgraph));
 
-  YnnThreadpool threadpool;
-  if (use_threadpool()) {
-    TF_ASSERT_OK_AND_ASSIGN(threadpool, CreateYnnThreadpool(&device));
-  }
-  Thunk::YnnParams ynn_params(std::move(threadpool));
-
   Thunk::ExecuteParams params;
   params.buffer_allocations = &allocations;
   params.intra_op_threadpool = use_threadpool() ? &device : nullptr;
-  params.ynn_params = &ynn_params;
 
   auto execute_event = thunk->Execute(params);
   tsl::BlockUntilReady(execute_event);
diff --git a/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_interop.cc b/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_interop.cc
index 693e4ca0cbe..0b1211623be 100644
--- a/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_interop.cc
+++ b/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_interop.cc
@@ -42,7 +42,7 @@ absl::StatusOr<YnnThreadpool> CreateYnnThreadpool(
     absl::FunctionRef<ynn_status(ynn_threadpool_t*)> builder) {
   ynn_threadpool_t threadpool = nullptr;
   YNN_RETURN_IF_ERROR(builder(&threadpool));
-  return YnnThreadpool(threadpool);
+  return YnnThreadpool(threadpool, ynn_delete_threadpool);
 }
 
 absl::StatusOr<ynn_type> YnnType(const PrimitiveType& type) {
diff --git a/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_interop.h b/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_interop.h
index 372b833e054..9d58f639981 100644
--- a/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_interop.h
+++ b/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_interop.h
@@ -87,16 +87,16 @@ namespace internal {
 struct YnnDeleter {
   void operator()(ynn_subgraph* subgraph) { ynn_delete_subgraph(subgraph); }
   void operator()(ynn_runtime* runtime) { ynn_delete_runtime(runtime); }
-  void operator()(ynn_threadpool* threadpool) {
-    ynn_delete_threadpool(threadpool);
-  }
 };
 
 }  // namespace internal
 
 using YnnSubgraph = std::unique_ptr<ynn_subgraph, internal::YnnDeleter>;
 using YnnRuntime = std::unique_ptr<ynn_runtime, internal::YnnDeleter>;
-using YnnThreadpool = std::unique_ptr<ynn_threadpool, internal::YnnDeleter>;
+// Uses a function-pointer deleter instead of `internal::YnnDeleter` to match
+// `Thunk::YnnParams::Threadpool`, which cannot depend on YNNPACK headers.
+using YnnThreadpool =
+    std::unique_ptr<ynn_threadpool, void (*)(ynn_threadpool*)>;
 
 absl::StatusOr<YnnSubgraph> CreateYnnSubgraph(
     absl::FunctionRef<ynn_status(ynn_subgraph_t*)> builder);
diff --git a/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_threadpool.cc b/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_threadpool.cc
index 11027e4b90a..fbabb55dbd4 100644
--- a/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_threadpool.cc
+++ b/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_threadpool.cc
@@ -29,17 +29,13 @@ limitations under the License.
 
 namespace xla::cpu {
 
-absl::StatusOr<YnnThreadpool> CreateYnnThreadpool(
-    Eigen::ThreadPoolInterface* threadpool) {
-  return CreateYnnThreadpool([&](ynn_threadpool_t* ynn_threadpool) {
-    *ynn_threadpool =
-        reinterpret_cast<ynn_threadpool_t>(new SlinkyThreadPool(threadpool));
-    return ynn_status_success;
-  });
+YnnThreadpool CreateYnnThreadpool(Eigen::ThreadPoolInterface* threadpool) {
+  return YnnThreadpool(
+      reinterpret_cast<ynn_threadpool_t>(new SlinkyThreadPool(threadpool)),
+      ynn_delete_threadpool);
 }
 
-absl::StatusOr<YnnThreadpool> CreateYnnThreadpool(
-    const Eigen::ThreadPoolDevice* device) {
+YnnThreadpool CreateYnnThreadpool(const Eigen::ThreadPoolDevice* device) {
   return CreateYnnThreadpool(device->getPool());
 }
 
diff --git a/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_threadpool.h b/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_threadpool.h
index 9fe8186c34a..ca0c366268d 100644
--- a/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_threadpool.h
+++ b/third_party/xla/xla/backends/cpu/runtime/ynnpack/ynn_threadpool.h
@@ -16,7 +16,6 @@ limitations under the License.
 #ifndef XLA_BACKENDS_CPU_RUNTIME_YNNPACK_YNN_THREADPOOL_H_
 #define XLA_BACKENDS_CPU_RUNTIME_YNNPACK_YNN_THREADPOOL_H_
 
-#include "absl/status/statusor.h"
 #include "xla/backends/cpu/runtime/ynnpack/ynn_interop.h"
 
 namespace Eigen {
@@ -27,12 +26,10 @@ class ThreadPoolInterface;
 namespace xla::cpu {
 
 // Creates an YNNPACK threadpool from an Eigen threadpool.
-absl::StatusOr<YnnThreadpool> CreateYnnThreadpool(
-    Eigen::ThreadPoolInterface* threadpool);
+YnnThreadpool CreateYnnThreadpool(Eigen::ThreadPoolInterface* threadpool);
 
 // Creates an YNNPACK threadpool from an Eigen ThreadPoolDevice.
-absl::StatusOr<YnnThreadpool> CreateYnnThreadpool(
-    const Eigen::ThreadPoolDevice* device);
+YnnThreadpool CreateYnnThreadpool(const Eigen::ThreadPoolDevice* device);
 
 }  // namespace xla::cpu
 
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
index 9483a4c67b8..f6d284af426 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
@@ -1616,11 +1616,6 @@ PjRtRawLoadedExecutable::RawExecuteResult CpuPjRtRawLoadedExecutable::Execute(
         cpu::Thunk::CustomCallExecuteParams custom_call_execute_params,
         cpu::Thunk::CustomCallExecuteParams::Create(&run_options));
 
-    std::optional<cpu::Thunk::YnnParams> ynn_params;
-    if (cpu_executable->has_ynn_fusions()) {
-      ABSL_ASSIGN_OR_RETURN(ynn_params, cpu::Thunk::YnnParams::Create(&run_options));
-    }
-
     cpu::ThreadPoolTaskRunner task_runner(
         run_options.intra_op_thread_pool()->getPool());
 
@@ -1632,7 +1627,6 @@ PjRtRawLoadedExecutable::RawExecuteResult CpuPjRtRawLoadedExecutable::Execute(
         &task_runner,
         &collective_params,
         &custom_call_execute_params,
-        ynn_params ? &*ynn_params : nullptr,
         run_options.run_id().ToInt(),
         run_options.device_ordinal(),
         cpu::Thunk::ExecuteSession(cpu::Thunk::ExecuteSession::kMaxWorkers,
diff --git a/third_party/xla/xla/service/cpu/cpu_executable.cc b/third_party/xla/xla/service/cpu/cpu_executable.cc
index 298a2891ee5..744536feb18 100644
--- a/third_party/xla/xla/service/cpu/cpu_executable.cc
+++ b/third_party/xla/xla/service/cpu/cpu_executable.cc
@@ -268,12 +268,6 @@ absl::Status CpuExecutable::ExecuteThunks(
   ABSL_ASSIGN_OR_RETURN(Thunk::CustomCallExecuteParams custom_call_execute_params,
                    Thunk::CustomCallExecuteParams::Create(run_options));
 
-  // Prepare for executing YNNPACK fusions.
-  std::optional<Thunk::YnnParams> ynn_params;
-  if (has_ynn_fusions()) {
-    ABSL_ASSIGN_OR_RETURN(ynn_params, Thunk::YnnParams::Create(run_options));
-  }
-
   // Use the intra-op thread pool to offload thunk executor tasks.
   auto* intra_op_thread_pool =
       run_options ? run_options->intra_op_thread_pool() : nullptr;
@@ -284,7 +278,6 @@ absl::Status CpuExecutable::ExecuteThunks(
       &*function_library_, &allocations,
       GetXfeedManager(GetDeviceOrdinal(run_options)), intra_op_thread_pool,
       &task_runner, &collective_execute_params, &custom_call_execute_params,
-      ynn_params ? &*ynn_params : nullptr,
       /*run_id=*/-1,
       /*device_ordinal=*/GetDeviceOrdinal(run_options),
       /*session=*/

From 8b7adbc139398657da80bf95b6db413b9cccf029 Mon Sep 17 00:00:00 2001
From: Shaogang Wang <shawnw@nvidia.com>
Date: Wed, 14 Oct 2026 01:02:29 -0700
Subject: [PATCH 0748/1000] PR #50179: [XLA:GPU] PJRT: typed MemFraction policy
 with opt-in growth
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50179

📝 Summary of Changes

Add an optional typed memory-fraction policy while preserving the existing GPU allocator configuration interface. `GpuAllocatorConfig::memory_fraction` remains a `double` with a fixed default of 0.75. When set, `memory_fraction_policy` takes precedence through `GetMemoryFraction()`.

| String      | Policy                         | Meaning                                |
|-------------|--------------------------------|----------------------------------------|
| `0.75`      | `FixedMemFraction{0.75}`        | Fixed cap at 75%.                       |
| `0.75+`     | `FlexMemFraction{0.75, 1.0}`    | Start at 75%, with a growth cap of 100%. |
| `0.75-0.85` | `FlexMemFraction{0.75, 0.85}`   | Start at 75%, with a growth cap of 85%.  |
| `0.75-0.75` | `FixedMemFraction{0.75}`        | Equal endpoints retain a fixed cap.    |

Move the implementation into `xla/pjrt/gpu/allocator_config.{h,cc}`. Keep the legacy plugin header and Bazel target as forwarding interfaces, so existing C++ callers and bindings retain their numeric field and include path. The PJRT C API keeps the numeric `memory_fraction` option and adds the string `memory_fraction_policy` option with the same precedence.

Bare fractions above 1 preserve unified-memory oversubscription. Parsing and canonical serialization support scientific notation. Explicit C++17 equality operators simplify policy comparisons in tests. Retain the GPU platform, CUPTI, and NVML dependencies used by the GPU tools, with keep annotations.

This PR is based on main following the landing of #50176, #50177, and #50178. Allocators use the resolved start fraction here; #50180 enables growth for flexible policies.

🎯 Justification

Introduce explicit opt-in growth policies without breaking existing memory limits, numeric C++ access, bindings, or downstream include paths. Keep the policy implementation in `pjrt/gpu` without requiring a plugin dependency.

🚀 Kind of Contribution

✨ New Feature, ♻️ Cleanup, 🧪 Tests

🧪 Unit Tests:

`//xla/pjrt/gpu:allocator_config_test` covers parsing, invalid inputs, scientific notation, serialization round trips, fixed defaults, legacy numeric values, policy precedence, and clearing an override.

`//xla/pjrt/plugin/xla_gpu:xla_gpu_allocator_config_test` exercises the forwarding header and target, numeric reads and writes, and `absl::SimpleAtod` with the legacy `double*` member.

🧪 Execution Tests:

`//xla/pjrt/c:pjrt_c_api_gpu_test` covers C API option acceptance and rejection, including client creation with fixed, shorthand, and range policies on one GPU. Existing `//xla/pjrt/gpu:se_gpu_pjrt_client_test` coverage checks shared-pool allocation behavior.

Validation on Prenyx B200 (Slurm job `2540205`, at most two visible GPUs): all 19 focused tests pass with no failures or skips across the four targets above (13 policy/configuration, one legacy compatibility, four C API allocator, and one shared-pool allocation test). The final source snapshot matches commit `13eebf0cd313219c46c40ed4a27d7904e1018121`.

`//xla/tools:matmul_perf_table_gen_main` and `//xla/tools/multihost_hlo_runner:hlo_runner_main` both build successfully under CUDA. A minimal `matmul_perf_table_gen_main --dry_run=true` invocation passes on B200, exercising GPU client creation and platform registration.

Clang Format, Buildifier, [clang-tidy-cuda](https://github.com/openxla/xla/actions/runs/37926048763/job/113805132965), and [DWYU](https://github.com/openxla/xla/actions/runs/37926048793/job/113805133029) pass. The unused NUMA include flagged by clang-tidy and corresponding stale dependencies were removed.

The oneAPI CI job fails in upstream Triton (`TritonGPUToLLVM.cpp:224`, `Chipset` to `ROCDL::TargetInfo` conversion); the [same failure occurs on the upstream main baseline](https://github.com/openxla/xla/actions/runs/37921869954/job/113791950244).

Copybara import of the project:

--
eff4530dd7261366f2d28792986f4adf2752b692 by Shaogang Wang <shawnw@nvidia.com>:

[XLA:GPU] PJRT: typed MemFraction memory policy with START-CAP grammar

GpuAllocatorConfig::memory_fraction becomes a MemFraction, a variant of
FixedMemFraction{fraction} (a hard cap, today's semantics) and
FlexMemFraction{start, cap} (preallocate start, allow growth up to cap).
ParseMemFraction accepts the XLA_PYTHON_CLIENT_MEM_FRACTION grammar: "0.75"
is Flex{0.75, 1.0}, "0.75-0.85" is Flex{0.75, 0.85}, "0.75-0.75" is
Fixed{0.75}; a bare fraction at or above 1 is fixed. The PJRT C API gains
the string option "memory_fraction_policy" carrying the full grammar; the
float option "memory_fraction" keeps working and means a bare fraction.
Callers that only need a number use MemFractionStart.

This change is API shape only: CreateBFCAllocator accepts the policy but
still builds the same pool for every value, so there is no behavior
change. Growth for Flex policies and the XLA flag that overrides the
policy follow in the next change. Fourth of five stacked changes (see
#48985 for the discussion that led to the variant).

--
14b6a558315fa2440e1bc604615912b2f794a02c by Shaogang Wang <shawnw@nvidia.com>:

[XLA:GPU] Preserve fixed memory fractions and make growth opt-in

--
035a1068638e14107a810ac313a2376c3fb65a3e by Shaogang Wang <shawnw@nvidia.com>:

[XLA:GPU] Clean up unused dependencies after allocator config move

--
09e3c08df4e08bdb3fd022125b5c2ae3519ac579 by Shaogang Wang <shawnw@nvidia.com>:

[XLA:GPU] Preserve allocator config compatibility and runtime registrations

Keep the legacy allocator header, Bazel target, and numeric memory_fraction member. Resolve an optional typed policy consistently across allocator paths, retain runtime registration dependencies, and simplify policy tests with C++17 equality operators.

--
1b5bb5b509640057292ee285545a62dc1c224efe by Shaogang Wang <shawnw@nvidia.com>:

[XLA:GPU] Remove unused NUMA include from the GPU client header

--
13eebf0cd313219c46c40ed4a27d7904e1018121 by Shaogang Wang <shawnw@nvidia.com>:

[XLA:GPU] Clean up GPU client dependencies

Merging this change closes #50179

PiperOrigin-RevId: 999186606
---
 third_party/xla/xla/pjrt/c/BUILD              |   2 +-
 .../xla/xla/pjrt/c/pjrt_c_api_gpu_internal.cc |  15 +-
 .../xla/xla/pjrt/c/pjrt_c_api_gpu_test.cc     |  63 ++++++++
 third_party/xla/xla/pjrt/gpu/BUILD            |  35 +++-
 .../xla/xla/pjrt/gpu/allocator_config.cc      | 122 ++++++++++++++
 .../xla/xla/pjrt/gpu/allocator_config.h       | 146 +++++++++++++++++
 .../xla/xla/pjrt/gpu/allocator_config_test.cc | 153 ++++++++++++++++++
 third_party/xla/xla/pjrt/gpu/gpu_helpers.cc   |  13 +-
 third_party/xla/xla/pjrt/gpu/gpu_helpers.h    |  10 +-
 .../xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc    |  15 +-
 .../xla/xla/pjrt/gpu/se_gpu_pjrt_client.h     |   3 +-
 third_party/xla/xla/pjrt/plugin/xla_gpu/BUILD |  15 +-
 .../plugin/xla_gpu/xla_gpu_allocator_config.h |  61 +------
 .../xla_gpu/xla_gpu_allocator_config_test.cc  |  35 ++++
 .../plugin/xla_gpu/xla_gpu_client_options.h   |   2 +-
 third_party/xla/xla/tools/BUILD               |   4 +-
 .../xla/xla/tools/multihost_hlo_runner/BUILD  |  10 +-
 17 files changed, 611 insertions(+), 93 deletions(-)
 create mode 100644 third_party/xla/xla/pjrt/gpu/allocator_config.cc
 create mode 100644 third_party/xla/xla/pjrt/gpu/allocator_config.h
 create mode 100644 third_party/xla/xla/pjrt/gpu/allocator_config_test.cc
 create mode 100644 third_party/xla/xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config_test.cc

diff --git a/third_party/xla/xla/pjrt/c/BUILD b/third_party/xla/xla/pjrt/c/BUILD
index d046c487831..514cb384a9a 100644
--- a/third_party/xla/xla/pjrt/c/BUILD
+++ b/third_party/xla/xla/pjrt/c/BUILD
@@ -667,11 +667,11 @@ cc_library(
         "//xla/pjrt:pjrt_executable",
         "//xla/pjrt/extensions/abi_version:gpu_abi_version_extension",
         "//xla/pjrt/extensions/cross_host_transfers:pjrt_c_api_cross_host_transfer_extension",
+        "//xla/pjrt/gpu:allocator_config",
         "//xla/pjrt/gpu:gpu_helpers",
         "//xla/pjrt/gpu:se_gpu_pjrt_client",
         "//xla/pjrt/gpu:se_gpu_pjrt_compiler",  # buildcleaner: keep to register GPU AOT compiler
         "//xla/pjrt/gpu:se_gpu_topology_description",
-        "//xla/pjrt/plugin/xla_gpu:xla_gpu_allocator_config",
         "//xla/pjrt/plugin/xla_gpu:xla_gpu_client_options",
         "//xla/python:custom_call_batch_partitioner",
         "//xla/python:custom_partition_callback",
diff --git a/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_internal.cc b/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_internal.cc
index b91ad680b2e..49798c4efab 100644
--- a/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_internal.cc
+++ b/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_internal.cc
@@ -56,6 +56,7 @@ limitations under the License.
 #include "xla/pjrt/c/pjrt_c_api_xla_transform_internal.h"
 #include "xla/pjrt/extensions/abi_version/gpu_abi_version_extension.h"
 #include "xla/pjrt/extensions/cross_host_transfers/pjrt_c_api_cross_host_transfers_extension.h"
+#include "xla/pjrt/gpu/allocator_config.h"
 #include "xla/pjrt/gpu/gpu_helpers.h"
 #include "xla/pjrt/gpu/se_gpu_pjrt_client.h"
 #include "xla/pjrt/gpu/se_gpu_topology_description.h"
@@ -64,7 +65,6 @@ limitations under the License.
 #include "xla/pjrt/pjrt_compiler.h"
 #include "xla/pjrt/pjrt_device_description.h"
 #include "xla/pjrt/pjrt_executable.h"
-#include "xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config.h"
 #include "xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h"
 #include "xla/python/custom_call_batch_partitioner.h"
 #include "xla/python/custom_partition_callback.h"
@@ -108,6 +108,8 @@ PJRT_Error* PJRT_Client_Create(PJRT_Client_Create_Args* args) {
           {"allocator", PJRT_NamedValue_Type::PJRT_NamedValue_kString},
           {"memory_fraction", PJRT_NamedValue_Type::PJRT_NamedValue_kFloat},
           {"preallocate", PJRT_NamedValue_Type::PJRT_NamedValue_kBool},
+          {"memory_fraction_policy",
+           PJRT_NamedValue_Type::PJRT_NamedValue_kString},
           {"collective_memory_size",
            PJRT_NamedValue_Type::PJRT_NamedValue_kInt64},
           {"visible_devices", PJRT_NamedValue_Type::PJRT_NamedValue_kInt64List},
@@ -168,6 +170,17 @@ PJRT_Error* PJRT_Client_Create(PJRT_Client_Create_Args* args) {
       it != create_options.end()) {
     allocator_config.collective_memory_size = std::get<int64_t>(it->second);
   }
+  if (auto it = create_options.find("memory_fraction_policy");
+      it != create_options.end()) {
+    // Full grammar ("0.75", "0.75+", "0.75-0.85"); wins over the float.
+    absl::StatusOr<xla::MemFraction> fraction =
+        xla::ParseMemFraction(std::get<std::string>(it->second));
+    if (!fraction.ok()) {
+      return StatusToPjRtError(absl::InvalidArgumentError(absl::StrFormat(
+          "memory_fraction_policy: %s", fraction.status().message())));
+    }
+    allocator_config.memory_fraction_policy = *fraction;
+  }
   std::optional<std::set<int>> visible_devices;
   if (auto it = create_options.find("visible_devices");
       it != create_options.end()) {
diff --git a/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_test.cc b/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_test.cc
index d9fc2514ea4..7088a8fb60e 100644
--- a/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_test.cc
+++ b/third_party/xla/xla/pjrt/c/pjrt_c_api_gpu_test.cc
@@ -712,6 +712,9 @@ TEST(PjrtCApiGpuAllocatorTest, ValidOptionsParsing) {
     if (allocator_option == "bfc" || allocator_option == "cuda_async") {
       options["memory_fraction"] = 0.5f;
     }
+    if (allocator_option == "bfc") {
+      options["memory_fraction_policy"] = static_cast<std::string>("0.5-0.9");
+    }
     if (allocator_option == "cuda_async") {
       options["preallocate"] = true;
     }
@@ -786,6 +789,66 @@ TEST(PjrtCApiGpuAllocatorTest, InvalidAllocatorOptionsParsing) {
   api->PJRT_Error_Destroy(&error_destroy_args);
 }
 
+TEST(PjrtCApiGpuAllocatorTest, MemoryFractionPolicies) {
+  const PJRT_Api* api = GetPjrtApi();
+  for (const char* policy : {"0.01", "0.01+", "0.01-0.02", "0.01-0.01"}) {
+    SCOPED_TRACE(policy);
+    absl::flat_hash_map<std::string, xla::PjRtValueType> options = {
+        {"allocator", std::string("bfc")},
+        {"visible_devices", std::vector<int64_t>{0}},
+        {"memory_fraction_policy", std::string(policy)},
+    };
+    ASSERT_OK_AND_ASSIGN(std::vector<PJRT_NamedValue> c_options,
+                         ConvertToPjRtNamedValueList(options));
+    PJRT_Client_Create_Args create_args = {};
+    create_args.struct_size = PJRT_Client_Create_Args_STRUCT_SIZE;
+    create_args.create_options = c_options.data();
+    create_args.num_options = c_options.size();
+    std::unique_ptr<PJRT_Error, PJRT_ErrorDeleter> error(
+        api->PJRT_Client_Create(&create_args), MakeErrorDeleter(api));
+    std::unique_ptr<PJRT_Client, PJRT_ClientDeleter> client(
+        create_args.client, MakeClientDeleter(api));
+    ASSERT_EQ(error, nullptr) << GetErrorMessage(error.get(), api);
+    ASSERT_NE(client, nullptr);
+  }
+}
+
+TEST(PjrtCApiGpuAllocatorTest, InvalidMemoryFractionPolicyRejected) {
+  auto api = GetPjrtApi();
+  absl::flat_hash_map<std::string, xla::PjRtValueType> options = {
+      {"allocator", static_cast<std::string>("bfc")},
+      {"memory_fraction_policy", static_cast<std::string>("0.9-0.5")},
+  };
+  ASSERT_OK_AND_ASSIGN(std::vector<PJRT_NamedValue> c_options,
+                       ::pjrt::ConvertToPjRtNamedValueList(options));
+  PJRT_Client_Create_Args create_arg;
+  create_arg.struct_size = PJRT_Client_Create_Args_STRUCT_SIZE;
+  create_arg.extension_start = nullptr;
+  create_arg.client = nullptr;
+  create_arg.create_options = c_options.data();
+  create_arg.num_options = c_options.size();
+  create_arg.kv_get_callback = nullptr;
+  create_arg.kv_get_user_arg = nullptr;
+  create_arg.kv_try_get_callback = nullptr;
+  create_arg.kv_try_get_user_arg = nullptr;
+  create_arg.kv_put_callback = nullptr;
+  create_arg.kv_put_user_arg = nullptr;
+  PJRT_Error* error = api->PJRT_Client_Create(&create_arg);
+  EXPECT_NE(error, nullptr);
+  EXPECT_THAT(::pjrt::PjrtErrorToStatus(error, api),
+              absl_testing::StatusIs(
+                  absl::StatusCode::kInvalidArgument,
+                  HasSubstr("memory_fraction_policy: Invalid memory fraction "
+                            "\"0.9-0.5\"")));
+
+  PJRT_Error_Destroy_Args error_destroy_args;
+  error_destroy_args.struct_size = PJRT_Error_Destroy_Args_STRUCT_SIZE;
+  error_destroy_args.extension_start = nullptr;
+  error_destroy_args.error = error;
+
+  api->PJRT_Error_Destroy(&error_destroy_args);
+}
+
 TEST(PjrtCApiGpuMaxInflightComputationsTest, ValidOptionsParsing) {
   auto api = GetPjrtApi();
   absl::flat_hash_map<std::string, xla::PjRtValueType> options = {
diff --git a/third_party/xla/xla/pjrt/gpu/BUILD b/third_party/xla/xla/pjrt/gpu/BUILD
index 01ce7dfc471..44ffae3ae0e 100644
--- a/third_party/xla/xla/pjrt/gpu/BUILD
+++ b/third_party/xla/xla/pjrt/gpu/BUILD
@@ -37,18 +37,44 @@ package(
 
 xla_gpu_internal_packages()
 
+cc_library(
+    name = "allocator_config",
+    srcs = ["allocator_config.cc"],
+    hdrs = ["allocator_config.h"],
+    visibility = internal_visibility([":legacy_gpu_client_users"]),
+    deps = [
+        "//xla/tsl/framework:allocator",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+)
+
+xla_cc_test(
+    name = "allocator_config_test",
+    srcs = ["allocator_config_test.cc"],
+    deps = [
+        ":allocator_config",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "gpu_helpers",
     srcs = ["gpu_helpers.cc"],
     hdrs = ["gpu_helpers.h"],
     visibility = internal_visibility(["//xla/pjrt:friends"]),
     deps = [
+        ":allocator_config",
         "//xla:types",
         "//xla:util",
         "//xla/client:client_library",
         "//xla/client:local_client",
-        "//xla/pjrt/plugin/xla_gpu:xla_gpu_allocator_config",
         "//xla/service:platform_util",
+        "//xla/stream_executor:memory_space",
         "//xla/stream_executor:platform",
         "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor/integrations:device_mem_allocator",
@@ -78,6 +104,7 @@ cc_library(
         "TENSORFLOW_USE_SYCL=1",
     ]),
     deps = [
+        ":allocator_config",
         ":gpu_helpers",
         ":gpu_metrics",
         ":se_gpu_pjrt_cross_host_transfer_proto_cc",
@@ -132,7 +159,6 @@ cc_library(
         "//xla/pjrt/distributed:protocol_proto_cc",
         "//xla/pjrt/distributed:topology_util",
         "//xla/pjrt/distributed/coordination:coordination_service_proto_cc",
-        "//xla/pjrt/plugin/xla_gpu:xla_gpu_allocator_config",
         "//xla/pjrt/plugin/xla_gpu:xla_gpu_client_options",
         "//xla/pjrt/proto:compile_options_proto_cc",
         "//xla/pjrt/proto:topology_description_proto_cc",
@@ -238,7 +264,6 @@ cc_library(
     ]) + if_rocm([
         # keep sorted
         "//xla/stream_executor/rocm:rocm_device_address_vmm_allocator",
-        "@local_config_rocm//rocm:rocm_headers",
     ]) + if_sycl([
         # keep sorted
         "//xla/backends/gpu/runtime:thunk_executor",
@@ -255,6 +280,7 @@ cc_library(
     hdrs = ["se_gpu_pjrt_client.h"],
     visibility = internal_visibility(["//xla/pjrt/gpu:legacy_gpu_client_users"]),
     deps = [
+        ":allocator_config",
         ":se_gpu_pjrt_client_impl",
         ":se_gpu_pjrt_compiler_cuda_registration",
         ":se_gpu_topology_description",
@@ -277,12 +303,10 @@ cc_library(
         "//xla/pjrt:raw_buffer",
         "//xla/pjrt/distributed:key_value_store_interface",
         "//xla/pjrt/distributed/coordination:coordination_service_proto_cc",
-        "//xla/pjrt/plugin/xla_gpu:xla_gpu_allocator_config",
         "//xla/pjrt/plugin/xla_gpu:xla_gpu_client_options",
         "//xla/pjrt/se:local_device_state",
         "//xla/pjrt/se:pjrt_stream_executor_client",
         "//xla/runtime:device_id",
-        "//xla/service:device_assignment",
         "//xla/service:gpu_topology",
         "//xla/service:gpu_topology_proto_cc",
         "//xla/service/gpu:gpu_executable_run_options",
@@ -301,7 +325,6 @@ cc_library(
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
         "@com_google_absl//absl/types:span",
-        "@tsl//tsl/platform:platform_port",
     ] + if_rocm([
         ":se_gpu_pjrt_compiler_rocm_registration",
     ]) + if_sycl([
diff --git a/third_party/xla/xla/pjrt/gpu/allocator_config.cc b/third_party/xla/xla/pjrt/gpu/allocator_config.cc
new file mode 100644
index 00000000000..a3f85c2fba8
--- /dev/null
+++ b/third_party/xla/xla/pjrt/gpu/allocator_config.cc
@@ -0,0 +1,122 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/pjrt/gpu/allocator_config.h"
+
+#include <cmath>
+#include <cstddef>
+#include <string>
+#include <variant>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/ascii.h"
+#include "absl/strings/numbers.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+
+namespace xla {
+namespace {
+
+constexpr absl::string_view kGrammar =
+    "expected START, START+, or START-CAP as fractions of total device memory, "
+    "for example \"0.75\" (fixed at 75%), \"0.75+\" (preallocate 75%, grow to "
+    "all memory), or \"0.75-0.85\" (grow to at most 85%)";
+
+absl::Status InvalidSpec(absl::string_view spec, absl::string_view reason) {
+  return absl::InvalidArgumentError(absl::StrCat(
+      "Invalid memory fraction \"", spec, "\": ", reason, "; ", kGrammar));
+}
+
+absl::StatusOr<double> ParseFraction(absl::string_view text,
+                                     absl::string_view spec) {
+  double value = 0;
+  if (!absl::SimpleAtod(absl::StripAsciiWhitespace(text), &value) ||
+      !std::isfinite(value) || value <= 0) {
+    return InvalidSpec(spec, "fractions must be positive numbers");
+  }
+  return value;
+}
+
+}  // namespace
+
+absl::StatusOr<MemFraction> ParseMemFraction(absl::string_view spec) {
+  absl::string_view text = absl::StripAsciiWhitespace(spec);
+  const bool grow_to_full_memory = !text.empty() && text.back() == '+';
+  size_t separator = absl::string_view::npos;
+  if (grow_to_full_memory) {
+    text.remove_suffix(1);
+  } else {
+    // A minus sign in a numeric exponent is not a range separator.
+    separator = text.find('-');
+    while (separator != absl::string_view::npos && separator > 0 &&
+           (text[separator - 1] == 'e' || text[separator - 1] == 'E')) {
+      separator = text.find('-', separator + 1);
+    }
+  }
+  absl::StatusOr<double> start = ParseFraction(text.substr(0, separator), spec);
+  if (!start.ok()) {
+    return start.status();
+  }
+  if (!grow_to_full_memory && separator == absl::string_view::npos) {
+    return MemFractionFromFraction(*start);
+  }
+  absl::StatusOr<double> cap =
+      grow_to_full_memory ? absl::StatusOr<double>(1.0)
+                          : ParseFraction(text.substr(separator + 1), spec);
+  if (!cap.ok()) {
+    return cap.status();
+  }
+  if (*cap < *start) {
+    return InvalidSpec(spec, "the cap must be at least the start");
+  }
+  if (*cap == *start) {
+    return MemFraction(FixedMemFraction{*start});
+  }
+  if (*cap > 1) {
+    return InvalidSpec(spec,
+                       "a growth cap cannot exceed 1 (all of device memory)");
+  }
+  return MemFraction(FlexMemFraction{*start, *cap});
+}
+
+MemFraction MemFractionFromFraction(double fraction) {
+  return FixedMemFraction{fraction};
+}
+
+double MemFractionStart(const MemFraction& fraction) {
+  if (const auto* fixed = std::get_if<FixedMemFraction>(&fraction)) {
+    return fixed->fraction;
+  }
+  return std::get<FlexMemFraction>(fraction).start;
+}
+
+std::string MemFractionToString(const MemFraction& fraction) {
+  if (const auto* fixed = std::get_if<FixedMemFraction>(&fraction)) {
+    return absl::StrCat(fixed->fraction);
+  }
+  const auto& flex = std::get<FlexMemFraction>(fraction);
+  if (flex.cap == 1) {
+    return absl::StrCat(flex.start, "+");
+  }
+  return absl::StrCat(flex.start, "-", flex.cap);
+}
+
+MemFraction GpuAllocatorConfig::GetMemoryFraction() const {
+  return memory_fraction_policy.value_or(
+      MemFractionFromFraction(memory_fraction));
+}
+
+}  // namespace xla
diff --git a/third_party/xla/xla/pjrt/gpu/allocator_config.h b/third_party/xla/xla/pjrt/gpu/allocator_config.h
new file mode 100644
index 00000000000..de5637c8343
--- /dev/null
+++ b/third_party/xla/xla/pjrt/gpu/allocator_config.h
@@ -0,0 +1,146 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_PJRT_GPU_ALLOCATOR_CONFIG_H_
+#define XLA_PJRT_GPU_ALLOCATOR_CONFIG_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <variant>
+#include <vector>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "xla/tsl/framework/allocator.h"
+
+namespace xla {
+
+// How much of a device's memory the BFC allocator may use, as fractions of
+// total device memory. ParseMemFraction accepts the following grammar:
+//
+//   "0.75"       FixedMemFraction{0.75}: hard cap at 75%
+//   "0.75+"      FlexMemFraction{0.75, 1.0}: preallocate 75%, grow to all
+//   "0.75-0.85"  FlexMemFraction{0.75, 0.85}: preallocate 75%, grow to 85%
+//   "0.75-0.75"  FixedMemFraction{0.75}: hard cap, the pool never grows
+//
+// Only the shared spatial BFC pool (preallocate=true with
+// --xla_gpu_enable_allocator_spatial_partitioning, on device memory) grows.
+// Every other allocator kind and mode uses just the start fraction.
+
+// The pool never grows; `fraction` is a hard cap. With unified memory the
+// fraction may exceed 1 to oversubscribe device memory.
+struct FixedMemFraction {
+  double fraction = 0.75;
+
+  bool operator==(const FixedMemFraction& other) const {
+    return fraction == other.fraction;
+  }
+};
+
+// Preallocate `start` and let default-memory allocations grow the pool up to
+// `cap`. Invariant: 0 < start < cap <= 1.
+struct FlexMemFraction {
+  double start = 0.75;
+  double cap = 1.0;
+
+  bool operator==(const FlexMemFraction& other) const {
+    return start == other.start && cap == other.cap;
+  }
+};
+
+using MemFraction = std::variant<FixedMemFraction, FlexMemFraction>;
+
+// Parses "START", "START+", or "START-CAP" as described above. "START+" is
+// shorthand for "START-1.0". Returns InvalidArgument for malformed input,
+// a cap below the start, or a growth cap above 1.
+absl::StatusOr<MemFraction> ParseMemFraction(absl::string_view spec);
+
+// Policy denoted by a bare fraction: a fixed cap, preserving the semantics of
+// the numeric PJRT C API memory_fraction option.
+MemFraction MemFractionFromFraction(double fraction);
+
+// Fraction allocated up front: FixedMemFraction::fraction or
+// FlexMemFraction::start.
+double MemFractionStart(const MemFraction& fraction);
+
+// Canonical spelling accepted by ParseMemFraction.
+std::string MemFractionToString(const MemFraction& fraction);
+
+struct GpuAllocatorConfig {
+  enum class Kind {
+    kDefault,   // Client picks the best option for the platform.
+    kPlatform,  // Synchronous passthrough allocator that calls the
+                // StreamExecutor Allocate/Deallocate APIs directly via a
+                // MultiDeviceAdapter wrapping StreamExecutorAllocator
+                // instances, with no BFC caching or pooled growth.
+    kBFC,  // Allocator using a "Best-Fit with Coalescing" algorithm. Currently
+           // only available for GPU.
+    kCudaAsync,  // Use the CUDA async allocator.
+    kVmm,  // Use Virtual Memory Management (VMM) allocator. This allocator
+           // uses CUDA VMM APIs to manage virtual address space separately from
+           // physical memory, enabling features like memory oversubscription
+           // and fine-grained memory mapping control.
+  };
+  Kind kind = Kind::kDefault;
+
+  // The fixed fraction of total device memory the allocator may use when
+  // memory_fraction_policy is unset. This is the default value of
+  // XLA_CLIENT_MEM_FRACTION.
+  //
+  // If gpu_system_memory_size is set, it determines the initial allocation.
+  double memory_fraction = 0.75;
+
+  // Only used if kind == kBFC. The absolute size of reserved memory space for
+  // GPU system in bytes.
+  //
+  // If null, the start fraction from GetMemoryFraction() will be used.
+  std::optional<int64_t> gpu_system_memory_size = std::nullopt;
+
+  // Only used if kind == kBFC. If true, the allocator will immediately allocate
+  // the initial amount specified by GetMemoryFraction(). This reduces
+  // fragmentation, allowing more of the total memory to be used. If false, the
+  // allocator will allocate more memory as allocations are requested.
+  bool preallocate = true;
+
+  // Amount of collective memory (ncclMemAlloc) to preallocate. If this value is
+  // 0, collective memory space will be grown as needed to fit the application's
+  // usage, with the drawback of potentially higher fragmentation. If set,
+  // should be set to a multiple of 512MB to avoid wasting memory due to
+  // granularity requirements.
+  size_t collective_memory_size = 0;
+
+  // Callbacks that get called when the underlying suballocator allocates or
+  // deallocates memory. See `SubAllocator::Visitor` for more details.
+  std::vector<tsl::SubAllocator::Visitor> sub_allocator_alloc_visitors;
+  std::vector<tsl::SubAllocator::Visitor> sub_allocator_free_visitors;
+
+  // Overrides memory_fraction when set. Growth requires an explicit flexible
+  // policy; other allocator kinds and modes use only its start fraction.
+  // If gpu_system_memory_size is set, it replaces the initial allocation,
+  // while the growth cap still applies.
+  //
+  // PJRT C API create options: "memory_fraction" (float, a fixed fraction) and
+  // "memory_fraction_policy" (string, the full grammar; takes precedence).
+  std::optional<MemFraction> memory_fraction_policy = std::nullopt;
+
+  // Resolves the optional policy, falling back to the legacy fixed fraction.
+  MemFraction GetMemoryFraction() const;
+};
+
+}  // namespace xla
+
+#endif  // XLA_PJRT_GPU_ALLOCATOR_CONFIG_H_
diff --git a/third_party/xla/xla/pjrt/gpu/allocator_config_test.cc b/third_party/xla/xla/pjrt/gpu/allocator_config_test.cc
new file mode 100644
index 00000000000..a0b97437293
--- /dev/null
+++ b/third_party/xla/xla/pjrt/gpu/allocator_config_test.cc
@@ -0,0 +1,153 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/pjrt/gpu/allocator_config.h"
+
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+
+namespace xla {
+namespace {
+
+using ::absl_testing::IsOkAndHolds;
+using ::absl_testing::StatusIs;
+using ::testing::HasSubstr;
+
+TEST(GpuAllocatorConfigTest, DefaultHasFixedCap) {
+  GpuAllocatorConfig config;
+  EXPECT_EQ(config.memory_fraction, 0.75);
+  EXPECT_FALSE(config.memory_fraction_policy.has_value());
+  EXPECT_THAT(config.GetMemoryFraction(), FixedMemFraction{0.75});
+}
+
+TEST(GpuAllocatorConfigTest, NumericFractionHasFixedCap) {
+  GpuAllocatorConfig config;
+  for (double fraction : {0.5, 0.9, 1.5}) {
+    config.memory_fraction = fraction;
+    EXPECT_THAT(config.GetMemoryFraction(), FixedMemFraction{fraction});
+  }
+}
+
+TEST(GpuAllocatorConfigTest, ExplicitPolicyOverridesNumericFraction) {
+  GpuAllocatorConfig config;
+  config.memory_fraction = 0.25;
+  config.memory_fraction_policy = FlexMemFraction{0.5, 0.9};
+  EXPECT_THAT(config.GetMemoryFraction(), (FlexMemFraction{0.5, 0.9}));
+  config.memory_fraction_policy = FixedMemFraction{0.6};
+  EXPECT_THAT(config.GetMemoryFraction(), FixedMemFraction{0.6});
+  EXPECT_EQ(config.memory_fraction, 0.25);
+  config.memory_fraction_policy.reset();
+  EXPECT_THAT(config.GetMemoryFraction(), FixedMemFraction{0.25});
+}
+
+TEST(MemFractionTest, BareFractionHasFixedCap) {
+  ASSERT_OK_AND_ASSIGN(MemFraction fraction, ParseMemFraction("0.75"));
+  EXPECT_THAT(fraction, FixedMemFraction{0.75});
+  EXPECT_THAT(ParseMemFraction(" 0.5 "), IsOkAndHolds(FixedMemFraction{0.5}));
+}
+
+TEST(MemFractionTest, PlusSuffixGrowsToAllDeviceMemory) {
+  EXPECT_THAT(ParseMemFraction("0.75+"),
+              IsOkAndHolds(FlexMemFraction{0.75, 1.0}));
+  EXPECT_THAT(ParseMemFraction(" 0.5+ "),
+              IsOkAndHolds(FlexMemFraction{0.5, 1.0}));
+  EXPECT_THAT(ParseMemFraction("0.75-1.0"),
+              IsOkAndHolds(FlexMemFraction{0.75, 1.0}));
+  // The shorthand has the same semantics as an explicit range ending at 1.
+  EXPECT_THAT(ParseMemFraction("1+"), IsOkAndHolds(FixedMemFraction{1.0}));
+  EXPECT_THAT(ParseMemFraction("1.5+"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("cap must be at least the start")));
+}
+
+TEST(MemFractionTest, RangeSetsStartAndCap) {
+  EXPECT_THAT(ParseMemFraction("0.75-0.85"),
+              IsOkAndHolds(FlexMemFraction{0.75, 0.85}));
+  EXPECT_THAT(ParseMemFraction("0.5-1.0"),
+              IsOkAndHolds(FlexMemFraction{0.5, 1.0}));
+}
+
+TEST(MemFractionTest, EqualEndsAreFixed) {
+  EXPECT_THAT(ParseMemFraction("0.75-0.75"),
+              IsOkAndHolds(FixedMemFraction{0.75}));
+  // Oversubscription (unified memory) is only possible without growth.
+  EXPECT_THAT(ParseMemFraction("1.5-1.5"), IsOkAndHolds(FixedMemFraction{1.5}));
+}
+
+TEST(MemFractionTest, BareFractionAtOrAboveOneIsFixed) {
+  EXPECT_THAT(ParseMemFraction("1"), IsOkAndHolds(FixedMemFraction{1.0}));
+  EXPECT_THAT(ParseMemFraction("1.5"), IsOkAndHolds(FixedMemFraction{1.5}));
+}
+
+TEST(MemFractionTest, NumericFractionHasFixedCap) {
+  EXPECT_THAT(MemFractionFromFraction(0.75), FixedMemFraction{0.75});
+  EXPECT_THAT(MemFractionFromFraction(0.9), FixedMemFraction{0.9});
+  EXPECT_THAT(MemFractionFromFraction(1.0), FixedMemFraction{1.0});
+  EXPECT_THAT(MemFractionFromFraction(1.5), FixedMemFraction{1.5});
+}
+
+TEST(MemFractionTest, ScientificNotation) {
+  EXPECT_THAT(ParseMemFraction("1e-3"), IsOkAndHolds(FixedMemFraction{0.001}));
+  EXPECT_THAT(ParseMemFraction("1e-3+"),
+              IsOkAndHolds(FlexMemFraction{0.001, 1.0}));
+  EXPECT_THAT(ParseMemFraction("1e-3-2e-3"),
+              IsOkAndHolds(FlexMemFraction{0.001, 0.002}));
+  EXPECT_THAT(ParseMemFraction("1E-3-1E+0"),
+              IsOkAndHolds(FlexMemFraction{0.001, 1.0}));
+}
+
+TEST(MemFractionTest, RejectsMalformedInput) {
+  for (const char* spec :
+       {"", "abc", "0", "-0.5", "nan", "inf", "0.5-", "0.5-abc", "0.5-0.7-0.9",
+        "+", "0+", "-0.5+", "nan+", "inf+", "0.5++", "0.5+-1", "0.5-0.9+"}) {
+    EXPECT_THAT(
+        ParseMemFraction(spec),
+        StatusIs(absl::StatusCode::kInvalidArgument, HasSubstr("START-CAP")))
+        << spec;
+  }
+  EXPECT_THAT(ParseMemFraction("0.9-0.5"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("cap must be at least the start")));
+  EXPECT_THAT(ParseMemFraction("0.5-1.5"),
+              StatusIs(absl::StatusCode::kInvalidArgument,
+                       HasSubstr("cannot exceed 1")));
+}
+
+TEST(MemFractionTest, StartIsThePreallocatedFraction) {
+  EXPECT_EQ(MemFractionStart(FlexMemFraction{0.6, 0.9}), 0.6);
+  EXPECT_EQ(MemFractionStart(FixedMemFraction{0.3}), 0.3);
+}
+
+TEST(MemFractionTest, ToStringRoundTrips) {
+  for (const char* spec :
+       {"0.75", "0.75+", "0.75-1.0", "0.75-0.85", "0.75-0.75", "0.5", "1.5",
+        "1+", "1e-8", "1e-8+", "1e-8-2e-8"}) {
+    ASSERT_OK_AND_ASSIGN(MemFraction fraction, ParseMemFraction(spec));
+    const std::string text = MemFractionToString(fraction);
+    ASSERT_OK_AND_ASSIGN(MemFraction reparsed, ParseMemFraction(text));
+    EXPECT_EQ(MemFractionToString(reparsed), text) << spec;
+    EXPECT_EQ(reparsed, fraction) << spec;
+  }
+  EXPECT_EQ(MemFractionToString(FlexMemFraction{0.75, 1.0}), "0.75+");
+  EXPECT_EQ(MemFractionToString(FlexMemFraction{0.75, 0.85}), "0.75-0.85");
+  EXPECT_EQ(MemFractionToString(FixedMemFraction{0.75}), "0.75");
+}
+
+}  // namespace
+}  // namespace xla
diff --git a/third_party/xla/xla/pjrt/gpu/gpu_helpers.cc b/third_party/xla/xla/pjrt/gpu/gpu_helpers.cc
index f8ef800f473..dae520e47f0 100644
--- a/third_party/xla/xla/pjrt/gpu/gpu_helpers.cc
+++ b/third_party/xla/xla/pjrt/gpu/gpu_helpers.cc
@@ -36,9 +36,11 @@ limitations under the License.
 #include "absl/types/span.h"
 #include "xla/client/client_library.h"
 #include "xla/client/local_client.h"
+#include "xla/pjrt/gpu/allocator_config.h"
 #include "xla/service/platform_util.h"
 #include "xla/stream_executor/integrations/device_mem_allocator.h"
 #include "xla/stream_executor/integrations/stream_executor_allocator.h"
+#include "xla/stream_executor/memory_space.h"
 #include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/framework/allocator.h"
@@ -100,8 +102,8 @@ void EnablePeerAccess(absl::Span<se::StreamExecutor* const> executors) {
 
 // Builds a BFCAllocator for all local GPUs.
 absl::StatusOr<std::shared_ptr<tsl::BFCAllocator>> CreateBFCAllocator(
-    se::StreamExecutor* executor, double memory_fraction, bool preallocate,
-    std::optional<int64_t> gpu_system_memory_size,
+    se::StreamExecutor* executor, const MemFraction& memory_fraction,
+    bool preallocate, std::optional<int64_t> gpu_system_memory_size,
     const std::vector<tsl::SubAllocator::Visitor>& sub_allocator_alloc_visitors,
     const std::vector<tsl::SubAllocator::Visitor>& sub_allocator_free_visitors,
     bool enable_spatial_partitioning) {
@@ -144,9 +146,12 @@ absl::StatusOr<std::shared_ptr<tsl::BFCAllocator>> CreateBFCAllocator(
   // unified memory.
   // When unified memory is enabled, allow GPU memory oversubscription by
   // setting memory_fraction > 1.
+  // Growth for FlexMemFraction policies is wired up by a follow-up change;
+  // until then only the start fraction matters.
+  const double start_fraction = MemFractionStart(memory_fraction);
   size_t allocator_memory = enable_unified_memory
-                                ? total_memory * fmax(1.0, memory_fraction)
-                                : total_memory * memory_fraction;
+                                ? total_memory * fmax(1.0, start_fraction)
+                                : total_memory * start_fraction;
   // If gpu_system_memory_size is set, use it instead of default value.
   if (gpu_system_memory_size.has_value()) {
     allocator_memory = gpu_system_memory_size.value();
diff --git a/third_party/xla/xla/pjrt/gpu/gpu_helpers.h b/third_party/xla/xla/pjrt/gpu/gpu_helpers.h
index 7a53ba33d1d..78d9930e6af 100644
--- a/third_party/xla/xla/pjrt/gpu/gpu_helpers.h
+++ b/third_party/xla/xla/pjrt/gpu/gpu_helpers.h
@@ -28,7 +28,7 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "xla/client/local_client.h"
-#include "xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config.h"
+#include "xla/pjrt/gpu/allocator_config.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/framework/allocator.h"
 #include "xla/tsl/framework/bfc_allocator.h"
@@ -63,9 +63,13 @@ absl::StatusOr<std::unique_ptr<tsl::BFCAllocator>> GetGpuHostAllocator(
 // heuristic, for both owned-hole reuse and central-gap carves.
 // Equal-size holes prefer lower addresses for collective memory and higher
 // addresses for default memory.
+//
+// memory_fraction describes how much of device memory the pool may use (see
+// MemFraction). Only its start fraction is used here; growth for
+// FlexMemFraction policies is wired up by a follow-up change.
 absl::StatusOr<std::shared_ptr<tsl::BFCAllocator>> CreateBFCAllocator(
-    se::StreamExecutor* executor, double memory_fraction, bool preallocate,
-    std::optional<int64_t> gpu_system_memory_size,
+    se::StreamExecutor* executor, const MemFraction& memory_fraction,
+    bool preallocate, std::optional<int64_t> gpu_system_memory_size,
     const std::vector<tsl::SubAllocator::Visitor>& sub_allocator_alloc_visitors,
     const std::vector<tsl::SubAllocator::Visitor>& sub_allocator_free_visitors,
     bool enable_spatial_partitioning = false);
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
index 84f1583b17c..ec881f12410 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
@@ -79,6 +79,7 @@ limitations under the License.
 #include "xla/pjrt/distributed/key_value_store_interface.h"
 #include "xla/pjrt/distributed/protocol.pb.h"
 #include "xla/pjrt/distributed/topology_util.h"
+#include "xla/pjrt/gpu/allocator_config.h"
 #include "xla/pjrt/gpu/gpu_helpers.h"
 #include "xla/pjrt/gpu/se_gpu_pjrt_cross_host_transfer.pb.h"
 #include "xla/pjrt/gpu/se_gpu_pjrt_runtime_abi_version.h"
@@ -93,7 +94,6 @@ limitations under the License.
 #include "xla/pjrt/pjrt_compiler.h"
 #include "xla/pjrt/pjrt_device_description.h"
 #include "xla/pjrt/pjrt_executable.h"
-#include "xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config.h"
 #include "xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h"
 #include "xla/pjrt/proto/topology_description.pb.h"
 #include "xla/pjrt/raw_buffer.h"
@@ -168,7 +168,6 @@ limitations under the License.
 #include "xla/stream_executor/cuda/cuda_device_address_vmm_allocator.h"
 #include "xla/stream_executor/gpu/gpu_cudamallocasync_allocator.h"
 #elif TENSORFLOW_USE_ROCM
-#include "rocm/rocm_config.h"
 #include "xla/stream_executor/rocm/rocm_device_address_vmm_allocator.h"
 #endif
 
@@ -1297,6 +1296,7 @@ GetStreamExecutorGpuDeviceAllocator(
     se::Platform* platform, const GpuAllocatorConfig& allocator_config,
     const std::map<int, std::unique_ptr<LocalDeviceState>>& addressable_devices,
     bool preallocate_host_memory) {
+  const MemFraction memory_fraction = allocator_config.GetMemoryFraction();
   std::vector<se::MultiDeviceAdapter::AllocatorInfo> allocators;
   const DebugOptions& debug_options = xla::GetDebugOptionsFromFlags();
   GpuAllocatorConfig::Kind effective_kind = allocator_config.kind;
@@ -1318,7 +1318,7 @@ GetStreamExecutorGpuDeviceAllocator(
         ABSL_ASSIGN_OR_RETURN(
             auto async_allocator,
             CreateCudaAsyncAllocator(
-                *(ordinal_and_device.second), allocator_config.memory_fraction,
+                *(ordinal_and_device.second), MemFractionStart(memory_fraction),
                 allocator_config.preallocate, false, false, true));
         allocators.push_back(
             {std::move(async_allocator),
@@ -1353,8 +1353,7 @@ GetStreamExecutorGpuDeviceAllocator(
         ABSL_ASSIGN_OR_RETURN(
             auto bfc_allocator,
             CreateBFCAllocator(ordinal_and_device.second->executor(),
-                               allocator_config.memory_fraction,
-                               allocator_config.preallocate,
+                               memory_fraction, allocator_config.preallocate,
                                allocator_config.gpu_system_memory_size,
                                allocator_config.sub_allocator_alloc_visitors,
                                allocator_config.sub_allocator_free_visitors,
@@ -1450,7 +1449,7 @@ GetStreamExecutorGpuDeviceAllocator(
             {device->executor(), device->compute_stream()});
       }
       return se::gpu::CudaDeviceAddressVmmAllocator::Create(
-          platform, allocator_config.memory_fraction,
+          platform, MemFractionStart(memory_fraction),
           allocator_config.gpu_system_memory_size, executor_streams,
           /*reclaim_exempt_memory_space=*/
           static_cast<int64_t>(gpu::MemorySpaceColor::kCollective));
@@ -1462,7 +1461,7 @@ GetStreamExecutorGpuDeviceAllocator(
             {device->executor(), device->compute_stream()});
       }
       return se::gpu::RocmDeviceAddressVmmAllocator::Create(
-          platform, allocator_config.memory_fraction,
+          platform, MemFractionStart(memory_fraction),
           allocator_config.gpu_system_memory_size, executor_streams,
           /*reclaim_exempt_memory_space=*/
           static_cast<int64_t>(gpu::MemorySpaceColor::kCollective));
@@ -1482,7 +1481,7 @@ GetStreamExecutorGpuDeviceAllocator(
           auto collective_bfc_allocator,
           CreateCollectiveBFCAllocator(
               ordinal_and_device.second->executor(),
-              /*memory_fraction=*/1.0 - allocator_config.memory_fraction,
+              /*memory_fraction=*/1.0 - MemFractionStart(memory_fraction),
               allocator_config.collective_memory_size));
       allocators.push_back(
           {std::move(collective_bfc_allocator),
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h
index 25620fb9a75..9859a752970 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.h
@@ -43,6 +43,7 @@ limitations under the License.
 #include "xla/pjrt/device_event.h"
 #include "xla/pjrt/distributed/coordination/coordination_service.pb.h"
 #include "xla/pjrt/distributed/key_value_store_interface.h"
+#include "xla/pjrt/gpu/allocator_config.h"
 #include "xla/pjrt/gpu/se_gpu_topology_description.h"
 #include "xla/pjrt/host_memory_allocator.h"
 #include "xla/pjrt/maybe_owning_mlir_module.h"
@@ -50,7 +51,6 @@ limitations under the License.
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/pjrt/pjrt_compiler.h"
 #include "xla/pjrt/pjrt_executable.h"
-#include "xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config.h"
 #include "xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h"
 #include "xla/pjrt/raw_buffer.h"
 #include "xla/pjrt/se/buffer_sequencing_event.h"
@@ -70,7 +70,6 @@ limitations under the License.
 #include "xla/tsl/concurrency/ref_count.h"
 #include "xla/tsl/framework/allocator.h"
 #include "xla/xla_data.pb.h"
-#include "tsl/platform/numa.h"
 
 namespace xla {
 
diff --git a/third_party/xla/xla/pjrt/plugin/xla_gpu/BUILD b/third_party/xla/xla/pjrt/plugin/xla_gpu/BUILD
index 5b7aa93d96f..19af2982e1b 100644
--- a/third_party/xla/xla/pjrt/plugin/xla_gpu/BUILD
+++ b/third_party/xla/xla/pjrt/plugin/xla_gpu/BUILD
@@ -13,6 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
+load("//xla:xla.default.bzl", "xla_cc_test")
 load("//xla/tests:build_defs.bzl", "xla_test")
 load("//xla/tsl:tsl.bzl", "if_google")
 load("//xla/tsl/platform:rules_cc.bzl", "cc_library")
@@ -42,9 +43,9 @@ cc_library(
     srcs = [],
     hdrs = ["xla_gpu_client_options.h"],
     deps = [
-        ":xla_gpu_allocator_config",
         "//xla/pjrt:host_memory_allocator",
         "//xla/pjrt/distributed:key_value_store_interface",
+        "//xla/pjrt/gpu:allocator_config",
         "@com_google_absl//absl/time",
     ],
 )
@@ -54,7 +55,17 @@ cc_library(
     srcs = [],
     hdrs = ["xla_gpu_allocator_config.h"],
     deps = [
-        "//xla/tsl/framework:allocator",
+        "//xla/pjrt/gpu:allocator_config",
+    ],
+)
+
+xla_cc_test(
+    name = "xla_gpu_allocator_config_test",
+    srcs = ["xla_gpu_allocator_config_test.cc"],
+    deps = [
+        ":xla_gpu_allocator_config",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest_main",
     ],
 )
 
diff --git a/third_party/xla/xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config.h b/third_party/xla/xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config.h
index 13c252af9e8..7b68f350a0e 100644
--- a/third_party/xla/xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config.h
+++ b/third_party/xla/xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config.h
@@ -16,64 +16,7 @@ limitations under the License.
 #ifndef XLA_PJRT_PLUGIN_XLA_GPU_XLA_GPU_ALLOCATOR_CONFIG_H_
 #define XLA_PJRT_PLUGIN_XLA_GPU_XLA_GPU_ALLOCATOR_CONFIG_H_
 
-#include <cstddef>
-#include <cstdint>
-#include <optional>
-#include <vector>
-
-#include "xla/tsl/framework/allocator.h"
-
-namespace xla {
-
-struct GpuAllocatorConfig {
-  enum class Kind {
-    kDefault,   // Client picks the best option for the platform.
-    kPlatform,  // Synchronous passthrough allocator that calls the
-                // StreamExecutor Allocate/Deallocate APIs directly via a
-                // MultiDeviceAdapter wrapping StreamExecutorAllocator
-                // instances, with no BFC caching or pooled growth.
-    kBFC,  // Allocator using a "Best-Fit with Coalescing" algorithm. Currently
-           // only available for GPU.
-    kCudaAsync,  // Use the CUDA async allocator.
-    kVmm,  // Use Virtual Memory Management (VMM) allocator. This allocator
-           // uses CUDA VMM APIs to manage virtual address space separately from
-           // physical memory, enabling features like memory oversubscription
-           // and fine-grained memory mapping control.
-  };
-  Kind kind = Kind::kDefault;
-
-  // Only used if kind == kBFC. The maximum fraction of available memory to
-  // allocate. This is the default value of XLA_CLIENT_MEM_FRACTION.
-  //
-  // If `gpu_system_memory_size` is set, it determines memory allocation.
-  // `memory_fraction` won't be used in this case.
-  double memory_fraction = 0.75;
-
-  // Only used if kind == kBFC. The absolute size of reserved memory space for
-  // GPU system in bytes.
-  //
-  // If null, the default value `memory_fraction` will be used.
-  std::optional<int64_t> gpu_system_memory_size = std::nullopt;
-
-  // Only used if kind == kBFC. If true, the allocator will immediately allocate
-  // the maximum amount allowed by `memory_fraction`. This reduces
-  // fragmentation, allowing more of the total memory to be used. If false, the
-  // allocator will allocate more memory as allocations are requested.
-  bool preallocate = true;
-
-  // Amount of collective memory (ncclMemAlloc) to preallocate. If this value is
-  // 0, collective memory space will be grown as needed to fit the application's
-  // usage, with the drawback of potentially higher fragmentation. If set,
-  // should be set to a multiple of 512MB to avoid wasting memory due to
-  // granularity requirements.
-  size_t collective_memory_size = 0;
-
-  // Callbacks that get called when the underlying suballocator allocates or
-  // deallocates memory. See `SubAllocator::Visitor` for more details.
-  std::vector<tsl::SubAllocator::Visitor> sub_allocator_alloc_visitors;
-  std::vector<tsl::SubAllocator::Visitor> sub_allocator_free_visitors;
-};
-
-}  // namespace xla
+// Compatibility header. New users should include allocator_config.h directly.
+#include "xla/pjrt/gpu/allocator_config.h"  // IWYU pragma: export
 
 #endif  // XLA_PJRT_PLUGIN_XLA_GPU_XLA_GPU_ALLOCATOR_CONFIG_H_
diff --git a/third_party/xla/xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config_test.cc b/third_party/xla/xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config_test.cc
new file mode 100644
index 00000000000..305f6008919
--- /dev/null
+++ b/third_party/xla/xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config_test.cc
@@ -0,0 +1,35 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config.h"
+
+#include <gtest/gtest.h>
+#include "absl/strings/numbers.h"
+
+namespace xla {
+namespace {
+
+TEST(GpuAllocatorConfigCompatibilityTest, NumericMemoryFraction) {
+  GpuAllocatorConfig config;
+  double fraction = config.memory_fraction;
+  EXPECT_EQ(fraction, 0.75);
+  config.memory_fraction = 0.5;
+  EXPECT_TRUE(absl::SimpleAtod("0.25", &config.memory_fraction));
+  EXPECT_EQ(config.memory_fraction, 0.25);
+  EXPECT_EQ(MemFractionStart(config.GetMemoryFraction()), 0.25);
+}
+
+}  // namespace
+}  // namespace xla
diff --git a/third_party/xla/xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h b/third_party/xla/xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h
index b85aeae17e5..0aa1b48519c 100644
--- a/third_party/xla/xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h
+++ b/third_party/xla/xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h
@@ -23,8 +23,8 @@ limitations under the License.
 
 #include "absl/time/time.h"
 #include "xla/pjrt/distributed/key_value_store_interface.h"
+#include "xla/pjrt/gpu/allocator_config.h"
 #include "xla/pjrt/host_memory_allocator.h"
-#include "xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config.h"
 
 namespace xla {
 
diff --git a/third_party/xla/xla/tools/BUILD b/third_party/xla/xla/tools/BUILD
index 43b7aa8ca76..8e5d13d41ad 100644
--- a/third_party/xla/xla/tools/BUILD
+++ b/third_party/xla/xla/tools/BUILD
@@ -837,9 +837,9 @@ cc_library(
         "@com_google_protobuf//:protobuf",
         "@tsl//tsl/platform:platform_port",
     ] + if_cuda([
-        "//xla/stream_executor:cuda_platform",
+        "//xla/stream_executor:cuda_platform",  # buildcleaner: keep
     ]) + if_rocm_is_configured([
-        "//xla/stream_executor:rocm_platform",
+        "//xla/stream_executor:rocm_platform",  # buildcleaner: keep
     ]),
 )
 
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/BUILD b/third_party/xla/xla/tools/multihost_hlo_runner/BUILD
index 55ee59268d0..6703a79dc9b 100644
--- a/third_party/xla/xla/tools/multihost_hlo_runner/BUILD
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/BUILD
@@ -96,14 +96,14 @@ cc_library(
         "//xla/service:gpu_plugin",
         "//xla/backends/profiler/gpu:device_tracer",
     ]) + if_cuda([
-        "//xla/stream_executor:cuda_platform",
-        "//xla/backends/profiler/gpu:cupti_tracer",
+        "//xla/stream_executor:cuda_platform",  # buildcleaner: keep
+        "//xla/backends/profiler/gpu:cupti_tracer",  # buildcleaner: keep
     ] + if_google(
         [
             "//third_party/py/jax/jaxlib/cuda:cuda_gpu_kernels",  # fixdeps: keep
         ],
     )) + if_rocm([
-        "//xla/stream_executor:rocm_platform",
+        "//xla/stream_executor:rocm_platform",  # buildcleaner: keep
     ]),
     alwayslink = True,
 )
@@ -337,7 +337,9 @@ xla_test(
         "@tsl//tsl/platform",
         "@tsl//tsl/platform:path",
         "@tsl//tsl/platform:protobuf",
-    ] + if_cuda(["//xla/tsl/cuda:nvml"]),
+    ] + if_cuda([
+        "//xla/tsl/cuda:nvml",  # buildcleaner: keep
+    ]),
 )
 
 tsl_pybind_extension(

From 1a0a8951c77cdeffd9401beac9819170fa43f43e Mon Sep 17 00:00:00 2001
From: Aliia Khasanova <aliia@google.com>
Date: Wed, 14 Oct 2026 01:07:28 -0700
Subject: [PATCH 0749/1000] Add a list of the xla::DebugOptions fields that can
 change a loaded executable

Tools such as AOT golden verification need to know which DebugOptions fields
can change how a loaded executable behaves. Today that knowledge lives in
hand-kept lists that drift as flags are added.

This is a plain C++ list rather than a field option in xla.proto: an extension
on FieldOptions is registered again by every shared object that embeds its own
copy of xla.proto, and protobuf aborts on duplicate extension registrations.

PiperOrigin-RevId: 999189097
---
 third_party/xla/xla/BUILD                     |  23 ++++
 .../xla/xla/debug_options_runtime_fields.cc   | 118 ++++++++++++++++++
 .../xla/xla/debug_options_runtime_fields.h    |  50 ++++++++
 .../xla/debug_options_runtime_fields_test.cc  |  56 +++++++++
 third_party/xla/xla/xla.proto                 |   4 +
 5 files changed, 251 insertions(+)
 create mode 100644 third_party/xla/xla/debug_options_runtime_fields.cc
 create mode 100644 third_party/xla/xla/debug_options_runtime_fields.h
 create mode 100644 third_party/xla/xla/debug_options_runtime_fields_test.cc

diff --git a/third_party/xla/xla/BUILD b/third_party/xla/xla/BUILD
index b34a5be0a64..edfa0401a22 100644
--- a/third_party/xla/xla/BUILD
+++ b/third_party/xla/xla/BUILD
@@ -1379,6 +1379,29 @@ cc_library(
         ],
 )
 
+cc_library(
+    name = "debug_options_runtime_fields",
+    srcs = ["debug_options_runtime_fields.cc"],
+    hdrs = ["debug_options_runtime_fields.h"],
+    visibility = internal_visibility([":friends"]),
+    deps = [
+        ":xla_proto_cc",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+xla_cc_test(
+    name = "debug_options_runtime_fields_test",
+    srcs = ["debug_options_runtime_fields_test.cc"],
+    deps = [
+        ":debug_options_runtime_fields",
+        ":xla_proto_cc",
+        "//xla/tsl/platform:test_main",
+        "@com_google_googletest//:gtest",
+    ],
+)
+
 xla_cc_test(
     name = "debug_options_flags_test",
     srcs = ["debug_options_flags_test.cc"],
diff --git a/third_party/xla/xla/debug_options_runtime_fields.cc b/third_party/xla/xla/debug_options_runtime_fields.cc
new file mode 100644
index 00000000000..0e9fdb07ebd
--- /dev/null
+++ b/third_party/xla/xla/debug_options_runtime_fields.cc
@@ -0,0 +1,118 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/debug_options_runtime_fields.h"
+
+#include <cstddef>
+
+#include "absl/algorithm/container.h"
+#include "absl/types/span.h"
+#include "xla/xla.pb.h"
+
+namespace xla {
+namespace {
+
+// Sorted by field number; the number is in the comment.
+constexpr int kRuntimeFields[] = {
+    DebugOptions::kXlaForceHostPlatformDeviceCountFieldNumber,            // 102
+    DebugOptions::kXlaHloEvaluatorUseFastPathFieldNumber,                 // 106
+    DebugOptions::kXlaDumpToFieldNumber,                                  // 109
+    DebugOptions::kXlaGpuNcclTerminationTimeoutSecondsFieldNumber,        // 163
+    DebugOptions::kXlaGpuGraphMinGraphSizeFieldNumber,                    // 208
+    DebugOptions::kXlaGpuEnableHighestPriorityAsyncStreamFieldNumber,     // 216
+    DebugOptions::kXlaGpuEnableCommandBufferFieldNumber,                  // 258
+    DebugOptions::kXlaGpuEnableNcclCommSplittingFieldNumber,              // 272
+    DebugOptions::kXlaGpuNcclCollectiveMaxNchannelsFieldNumber,           // 273
+    DebugOptions::kXlaGpuNcclP2PMaxNchannelsFieldNumber,                  // 274
+    DebugOptions::kXlaGpuNcclInitMaxRankPerRootRatioFieldNumber,          // 277
+    DebugOptions::kXlaGpuNcclTerminateOnErrorFieldNumber,                 // 301
+    DebugOptions::kXlaCmdBufferTraceCacheSizeFieldNumber,                 // 311
+    DebugOptions::kXlaGpuTempBufferUseSeparateColorFieldNumber,           // 312
+    DebugOptions::kXlaEnableCommandBuffersDuringProfilingFieldNumber,     // 317
+    DebugOptions::kXlaGpuExecutableWarnStuckTimeoutSecondsFieldNumber,    // 327
+    DebugOptions::kXlaGpuExecutableTerminateTimeoutSecondsFieldNumber,    // 328
+    DebugOptions::                                                        //
+    kXlaGpuExperimentalParallelCollectiveOverlapLimitFieldNumber,         // 336
+    DebugOptions::kXlaGpuRequireExclusiveLockFieldNumber,                 // 347
+    DebugOptions::kXlaGpuNcclBlockingCommunicatorsFieldNumber,            // 390
+    DebugOptions::                                                        //
+    kXlaGpuFirstCollectiveCallWarnStuckTimeoutSecondsFieldNumber,         // 391
+    DebugOptions::                                                        //
+    kXlaGpuFirstCollectiveCallTerminateTimeoutSecondsFieldNumber,         // 392
+    DebugOptions::kXlaGpuNcclAsyncExecutionFieldNumber,                   // 393
+    DebugOptions::kXlaEarlyExitWithLayoutsFieldNumber,                    // 397
+    DebugOptions::kXlaGpuCommandBufferSchedulingModeFieldNumber,          // 404
+    DebugOptions::                                                        //
+    kXlaGpuExperimentalEnableNcclSymmetricBuffersFieldNumber,             // 406
+    DebugOptions::kXlaGpuCommandBufferUnrollLoopsFieldNumber,             // 411
+    DebugOptions::                                                        //
+    kXlaGpuExperimentalEnableChecksumTracingOnThunksFieldNumber,          // 414
+    DebugOptions::                                                        //
+    kXlaCpuCollectiveCallTerminateTimeoutSecondsFieldNumber,              // 417
+    DebugOptions::kXlaCpuCollectiveCallWarnStuckSecondsFieldNumber,       // 418
+    DebugOptions::kXlaGpuExperimentalThunkBufferDebugFilterFieldNumber,   // 424
+    DebugOptions::kXlaGpuDetectNanFieldNumber,                            // 426
+    DebugOptions::kXlaGpuDetectInfFieldNumber,                            // 428
+    DebugOptions::                                                        //
+    kXlaGpuExperimentalEnableBufferSaverOnThunksFieldNumber,              // 431
+    DebugOptions::kXlaGpuExecutableEmbedDebugInfoFieldNumber,             // 437
+    DebugOptions::kXlaCpuCollectiveTimeoutSecondsFieldNumber,             // 438
+    DebugOptions::kXlaGpuExecutionTerminateTimeoutFieldNumber,            // 451
+    DebugOptions::kXlaGpuLogMinmaxFieldNumber,                            // 453
+    DebugOptions::kXlaGpuExecutionProgressTrackingFieldNumber,            // 457
+    DebugOptions::kXlaGpuCollectivesImplementationFieldNumber,            // 468
+    DebugOptions::kXlaGpuCommandBufferUpdateModeFieldNumber,              // 469
+    DebugOptions::kXlaGpuExecutableNumComputeStreamsFieldNumber,          // 486
+    DebugOptions::kXlaGpuExecutableNumCommunicationStreamsFieldNumber,    // 487
+    DebugOptions::                                                        //
+    kXlaGpuExperimentalThunkBufferDebugModuleOutputsFieldNumber,          // 492
+    DebugOptions::kXlaGpuEnableAllocatorSpatialPartitioningFieldNumber,   // 494
+    DebugOptions::                                                        //
+    kXlaGpuEnableCollectivesCommandBufferFilterFieldNumber,               // 497
+    DebugOptions::                                                        //
+    kXlaEnableNcclSymmetricBuffersForCollectivesFieldNumber,              // 503
+    DebugOptions::kXlaGpuEnableNcclUserBuffersInDefaultSpaceFieldNumber,  // 506
+    DebugOptions::kXlaGpuGxlScratchSizeBytesFieldNumber,                  // 514
+    DebugOptions::kXlaGpuExperimentalEarlyExitFieldNumber,                // 549
+    DebugOptions::kXlaGpuExperimentalVmmDisabledFieldNumber,              // 551
+    DebugOptions::kXlaGpuDeviceExecutionTerminateTimeoutFieldNumber,      // 553
+};
+
+// Hand-written because std::is_sorted and std::adjacent_find are not
+// constexpr in C++17, which the OSS build uses.
+constexpr bool IsStrictlyIncreasing(absl::Span<const int> fields) {
+  for (size_t i = 1; i < fields.size(); ++i) {
+    if (fields[i - 1] >= fields[i]) {
+      return false;
+    }
+  }
+  return true;
+}
+
+static_assert(IsStrictlyIncreasing(kRuntimeFields),
+              "kRuntimeFields must be sorted by field number, without "
+              "duplicates.");
+
+}  // namespace
+
+absl::Span<const int> RuntimeDebugOptionsFieldNumbers() {
+  return kRuntimeFields;
+}
+
+bool IsRuntimeDebugOptionsField(int field_number) {
+  return absl::c_binary_search(kRuntimeFields, field_number);
+}
+
+}  // namespace xla
diff --git a/third_party/xla/xla/debug_options_runtime_fields.h b/third_party/xla/xla/debug_options_runtime_fields.h
new file mode 100644
index 00000000000..48e6f52ed24
--- /dev/null
+++ b/third_party/xla/xla/debug_options_runtime_fields.h
@@ -0,0 +1,50 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_DEBUG_OPTIONS_RUNTIME_FIELDS_H_
+#define XLA_DEBUG_OPTIONS_RUNTIME_FIELDS_H_
+
+#include "absl/types/span.h"
+
+namespace xla {
+
+// Returns the field numbers of the runtime `DebugOptions` flags, sorted. Used
+// by tools such as AOT golden verification.
+//
+// A runtime flag is read after compilation (i.e. when loading or running the
+// executable) and can change execution: results, which kernels or thunks run,
+// or whether loading or running fails. A flag whose value is copied into the
+// HLO, a backend config or a thunk while compiling isn't a runtime flag, since
+// the executable carries the copy.
+//
+// Flags that only affect diagnostics (dump formats and filters, NVTX
+// annotations, profiler setup) are left out even when read after compilation.
+// The exception is `xla_dump_to`, because whether it's empty changes
+// execution: with the NaN/Inf checker in DUMP mode an empty value makes loading
+// fail, and with the buffer saver enabled it skips the saver thunks. Tools
+// should compare `empty()` rather than the path.
+//
+// If you're adding a new runtime flag, add it to `kRuntimeFields` in
+// debug_options_runtime_fields.cc, since nothing enforces this list. If unsure,
+// add it: a missing flag makes tools silently ignore a real difference, an
+// extra one only makes them report a harmless one.
+absl::Span<const int> RuntimeDebugOptionsFieldNumbers();
+
+// Returns whether `field_number` is in `RuntimeDebugOptionsFieldNumbers()`.
+bool IsRuntimeDebugOptionsField(int field_number);
+
+}  // namespace xla
+
+#endif  // XLA_DEBUG_OPTIONS_RUNTIME_FIELDS_H_
diff --git a/third_party/xla/xla/debug_options_runtime_fields_test.cc b/third_party/xla/xla/debug_options_runtime_fields_test.cc
new file mode 100644
index 00000000000..1a7fdff1fe2
--- /dev/null
+++ b/third_party/xla/xla/debug_options_runtime_fields_test.cc
@@ -0,0 +1,56 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/debug_options_runtime_fields.h"
+
+#include <gtest/gtest.h>
+#include "xla/xla.pb.h"
+
+namespace xla {
+namespace {
+
+TEST(DebugOptionsRuntimeFieldsTest, ListsFieldsThatChangeExecution) {
+  // Read by the NaN checker thunk pass when the executable is loaded.
+  EXPECT_TRUE(
+      IsRuntimeDebugOptionsField(DebugOptions::kXlaGpuDetectNanFieldNumber));
+  // Read from the process's flags when the client is created.
+  EXPECT_TRUE(IsRuntimeDebugOptionsField(
+      DebugOptions::kXlaForceHostPlatformDeviceCountFieldNumber));
+  // Read by the watchdog while the executable runs.
+  EXPECT_TRUE(IsRuntimeDebugOptionsField(
+      DebugOptions::kXlaGpuDeviceExecutionTerminateTimeoutFieldNumber));
+  // An empty value makes loading fail with the NaN checker in DUMP mode.
+  EXPECT_TRUE(IsRuntimeDebugOptionsField(DebugOptions::kXlaDumpToFieldNumber));
+}
+
+TEST(DebugOptionsRuntimeFieldsTest, OmitsCompileOnlyAndDiagnosticsFields) {
+  // Read only while compiling.
+  EXPECT_FALSE(IsRuntimeDebugOptionsField(
+      DebugOptions::kXlaDebugBufferAssignmentShowMaxFieldNumber));
+  EXPECT_FALSE(IsRuntimeDebugOptionsField(
+      DebugOptions::
+          kXlaGpuExperimentalEnableDynamicSliceExtendedOffsetsFieldNumber));
+  // Read after compilation, but only changes which dump files are written.
+  EXPECT_FALSE(
+      IsRuntimeDebugOptionsField(DebugOptions::kXlaDumpHloAsTextFieldNumber));
+  // Read after compilation, but only affects profiling.
+  EXPECT_FALSE(IsRuntimeDebugOptionsField(
+      DebugOptions::kXlaGpuTraceAnnotationLevelFieldNumber));
+  EXPECT_FALSE(IsRuntimeDebugOptionsField(
+      DebugOptions::kXlaGpuEnableCuptiMultiSubscriberFieldNumber));
+}
+
+}  // namespace
+}  // namespace xla
diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index d696df8c9d1..0a71b05df2e 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -61,6 +61,10 @@ message ThunkBufferDebugFilter {
 // All singular scalar fields must be explicitly defined as `optional` so that
 // field presence is available to support merging between command-line flags and
 // stored instances. This is enforced via a unit test.
+//
+// Fields that can change how a loaded executable behaves must be added to
+// `kRuntimeFields` in `debug_options_runtime_fields.cc`. This is not
+// enforced; see `debug_options_runtime_fields.h` for the criteria.
 message DebugOptions {
   // Enum to define all collective ops
   // that xla supports.

From 596eca059dd448a67178f25e31e7ab34222b1aaa Mon Sep 17 00:00:00 2001
From: Mikhail Goncharov <goncharov@google.com>
Date: Wed, 14 Oct 2026 02:01:46 -0700
Subject: [PATCH 0750/1000] [XLA:GPU] Deprecate
 --xla_gpu_experimental_enable_tiling_propagation flag and proto field

flag is mute and new tiling is always on

PiperOrigin-RevId: 999214406
---
 third_party/xla/xla/debug_options_flags.cc | 10 ++++------
 third_party/xla/xla/xla.proto              |  4 +++-
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index 6d70cc36c80..d709e4bf7a8 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -3901,12 +3901,9 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
       debug_options->xla_gpu_rocm_max_trace_events(),
       "Maximum number of ROCm trace events (applies to callback/activity/"
       "annotation). Set as high as memory allows; up to 1e9."));
-  flag_list->push_back(tsl::Flag(
-      "xla_gpu_experimental_enable_tiling_propagation",
-      bool_setter_for(
-          &DebugOptions::set_xla_gpu_experimental_enable_tiling_propagation),
-      debug_options->xla_gpu_experimental_enable_tiling_propagation(),
-      "If true, enable experimental tiling propagation."));
+  flag_list->push_back(
+      tsl::Flag("xla_gpu_experimental_enable_tiling_propagation",
+                noop_flag_setter<bool>, true, "[Deprecated, do not use]."));
 
   auto setter_for_xla_gpu_detect_nan =
       [debug_options, detection_mode](const std::string& value) {
@@ -4117,6 +4114,7 @@ DebugOptions GetDebugOptionsFromProtoAndFlags(
     DebugOptions* proto_debug_options) {
   if (proto_debug_options != nullptr) {
     DebugOptions debug_options = *proto_debug_options;
+    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(true);
     // Apply overrides from XLA_FLAGS environment variable.
     std::vector<tsl::Flag> flag_list;
     MakeDebugOptionsFlags(&flag_list, &debug_options);
diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index 0a71b05df2e..72459561fae 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -1882,7 +1882,9 @@ message DebugOptions {
   optional int64 xla_gpu_async_copy_min_bytes = 522;
 
   // If true, enable experimental tiling propagation.
-  optional bool xla_gpu_experimental_enable_tiling_propagation = 456;
+  // TODO(b/514293537): drop the flag.
+  optional bool xla_gpu_experimental_enable_tiling_propagation = 456
+      [deprecated = true];
 
   // Enables and controls expetimental cost model-driven tiling choices for GEMM
   // fusions.

From 18dbca5e64a8b2bf2db3caf18632257d6624f0cf Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Wed, 14 Oct 2026 02:36:05 -0700
Subject: [PATCH 0751/1000] [XLA:GPU]: Add a kReplica dim type for narrowing
 collectives.

For narrowing collectives we don't have a way to increase the degree of
dimension as we go through them. This introduces a new dim type to have that degree of freedom when going through reduce-scatter.

PiperOrigin-RevId: 999230759
---
 .../codegen/tiling/experimental/scheduling.cc |  31 +++--
 .../tiling/experimental/scheduling_test.cc    |  26 ++++
 .../tiling/experimental/tile_propagation.cc   |  21 +++-
 .../experimental/tile_propagation_test.cc     |  40 +++---
 .../tiling/experimental/tiled_hlo_test.cc     |  13 +-
 .../tiling/experimental/tiling_space.cc       | 116 ++++++++++++++----
 .../tiling/experimental/tiling_space.h        |  22 +++-
 .../tiling/experimental/tiling_space_test.cc  |  42 +++++++
 .../validate_reshape_propagation.cc           |   4 +-
 .../codegen/xtile/codegen/emitter_helpers.cc  |   1 +
 .../xtile/tiling_from_block_parameters.cc     |   2 +
 11 files changed, 245 insertions(+), 73 deletions(-)

diff --git a/third_party/xla/xla/codegen/tiling/experimental/scheduling.cc b/third_party/xla/xla/codegen/tiling/experimental/scheduling.cc
index e9d8af7bfd9..7c0bdc0debd 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/scheduling.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/scheduling.cc
@@ -580,17 +580,23 @@ absl::StatusOr<Schedule> GetSchedule(
                                         global_tile_id, num_tiles_per_pid);
   }
 
-  // Compute the block counts for each parallel dimension.
+  // Compute the block counts for each replica and parallel dimension.
+  // Replica dimensions are placed outermost so each contiguous group of
+  // parallel blocks targets the same peer.
+  const TilingSpace& tiling_space = tiled_computation.tiling_space();
+  llvm::SmallVector<int64_t, 4> block_counts;
+  llvm::SmallVector<int64_t, 4> dim_ids;
   llvm::SmallVector<int64_t, 4> parallel_dim_block_counts;
   llvm::SmallVector<int64_t, 4> parallel_dim_ids;
   for (const auto& [dim_id, dimension] :
-       llvm::enumerate(tiled_computation.tiling_space().dimensions())) {
-    if (dimension.type != TilingSpace::DimensionSemantics::kParallel) {
-      continue;
+       llvm::enumerate(tiling_space.dimensions())) {
+    if (dimension.type == TilingSpace::DimensionSemantics::kReplica) {
+      block_counts.push_back(tiling_space.GetNumTiles(dimension));
+      dim_ids.push_back(dim_id);
+    } else if (dimension.type == TilingSpace::DimensionSemantics::kParallel) {
+      parallel_dim_block_counts.push_back(tiling_space.GetNumTiles(dimension));
+      parallel_dim_ids.push_back(dim_id);
     }
-    parallel_dim_block_counts.push_back(
-        CeilOfRatio(dimension.dimension_size, *dimension.tile_size));
-    parallel_dim_ids.push_back(dim_id);
   }
 
   llvm::SmallVector<int64_t> permutation =
@@ -601,15 +607,16 @@ absl::StatusOr<Schedule> GetSchedule(
     parallel_dim_ids =
         llvm::to_vector<4>(xla::Permute(parallel_dim_ids, permutation));
   }
+  block_counts.append(parallel_dim_block_counts);
+  dim_ids.append(parallel_dim_ids);
 
   llvm::SmallVector<SymbolicExpr, 4> delinearized_pid =
-      DelinearizeIndex(parallel_dim_block_counts, global_tile_id, ctx);
+      DelinearizeIndex(block_counts, global_tile_id, ctx);
   Schedule schedule;
-  for (const auto& [parallel_dim_id, expr] :
-       llvm::zip(parallel_dim_ids, delinearized_pid)) {
-    schedule.dim_id_to_pid_expr[parallel_dim_id] = expr;
+  for (const auto& [dim_id, expr] : llvm::zip(dim_ids, delinearized_pid)) {
+    schedule.dim_id_to_pid_expr[dim_id] = expr;
   }
-  schedule.num_tiles = Product(parallel_dim_block_counts);
+  schedule.num_tiles = Product(block_counts);
   schedule.num_pids = CeilOfRatio(schedule.num_tiles, num_tiles_per_pid);
   return schedule;
 }
diff --git a/third_party/xla/xla/codegen/tiling/experimental/scheduling_test.cc b/third_party/xla/xla/codegen/tiling/experimental/scheduling_test.cc
index 238f0367d86..bb4b43876b2 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/scheduling_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/scheduling_test.cc
@@ -556,4 +556,30 @@ TEST_F(SchedulingTest, VariadicAllGatherWithDotProducersSharesPeerIndex) {
           "num_pids=16, num_tiles=16")));
 }
 
+TEST_F(SchedulingTest, SingleOperandReduceScatterReplicaDimension) {
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"(
+    add {
+      lhs = f32[] parameter(0)
+      rhs = f32[] parameter(1)
+      ROOT add = f32[] add(lhs, rhs)
+    }
+
+    f {
+      p0 = f32[8,64] parameter(0)
+      ROOT rs = f32[4,64] reduce-scatter(p0),
+        replica_groups={{0,1}}, dimensions={0}, to_apply=add
+    }
+
+    ENTRY e {
+      arg0 = f32[8,64] parameter(0)
+      ROOT fusion = f32[4,64] fusion(arg0), kind=kCustom, calls=f
+    })",
+                                    {1, 32}));
+  EXPECT_THAT(GetSchedule(tiled_computation),
+              IsOkAndHolds(
+                  MatchSchedule("d0 -> (pid mod 4) / 2, d1 -> pid mod 4 mod 2, "
+                                "d2 -> pid / 4, num_pids=8, num_tiles=8")));
+}
+
 }  // namespace xla::gpu::experimental
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
index 1eddb4f135b..6996fd03093 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
@@ -1599,7 +1599,7 @@ absl::StatusOr<Tiles> PropagateTileToInputForReduceScatterOp(
   const int64_t scatter_dim = reduce_scatter.scatter_dimension();
   const int64_t shard_size = reduce_scatter.shape().dimensions(scatter_dim);
   const int64_t num_replicas =
-      reduce_scatter.operand(0)->shape().dimensions(scatter_dim) / shard_size;
+      reduce_scatter.device_list()->num_devices_per_group();
 
   SmallVector<DimTile> dim_tiles(output_tile.dim_tiles());
   DimTile& scatter_tile = dim_tiles[scatter_dim];
@@ -1608,14 +1608,25 @@ absl::StatusOr<Tiles> PropagateTileToInputForReduceScatterOp(
         absl::StrCat("Non-unit stride on the scatter dimension of ",
                      reduce_scatter.ToString()));
   }
-  scatter_tile.offset = scatter_tile.offset * num_replicas;
-  scatter_tile.size = scatter_tile.size * num_replicas;
-  scatter_tile.upper_bound = scatter_tile.upper_bound * num_replicas;
+  if (num_replicas <= 1) {
+    return Tiles{output_tile};
+  }
+  const TilingSpace::DimensionInfo& replica_dim_info =
+      tiling_space.GetDimensionInfo(reduce_scatter,
+                                    reduce_scatter.shape().dimensions().size());
+  SymbolicExpr replica_id =
+      CreateDimExpr(replica_dim_info.id.value(), tiling_space.mlir_context());
+  scatter_tile.offset = (scatter_tile.offset + replica_id * shard_size -
+                         replica_id * scatter_tile.size)
+                            .Canonicalize();
+  scatter_tile.size = (scatter_tile.size * num_replicas).Canonicalize();
+  scatter_tile.upper_bound =
+      (scatter_tile.upper_bound * num_replicas).Canonicalize();
 
   if (!tiling_space.IsSymbolic()) {
     std::optional<int64_t> size = TryGetConstantValue(scatter_tile.size);
     if (!size.has_value() || shard_size % *size != 0 ||
-        !scatter_tile.offset.Canonicalize().IsMultipleOf(*size)) {
+        !scatter_tile.offset.IsMultipleOf(*size)) {
       return absl::FailedPreconditionError(absl::StrCat(
           "Tiling propagation rejected for ", reduce_scatter.ToString(),
           ": operand tile ", scatter_tile.size.ToString(), " at ",
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
index 1cc8b0550ae..f4b6d815421 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
@@ -456,16 +456,14 @@ TEST_F(TilePropagationTest, CanPropagateToInputsOfReduceScatterOp) {
       std::unique_ptr<TilingSpace> tiling_space,
       TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
                           &mlir_context_));
-  ASSERT_OK_AND_ASSIGN(
-      auto tiled_operands,
-      PropagateTileToInput(
-          *tiling_space, *root,
-          GetTestTile(*tiling_space, root->shape().dimensions()), 0));
+  ASSERT_OK_AND_ASSIGN(auto tiled_operands,
+                       PropagateTileToInput(*tiling_space, *root,
+                                            tiling_space->tiled_roots()[0], 0));
   EXPECT_THAT(tiled_operands, MatchToString(R"(
-    0) (tid_0, tid_1)
-      -> offsets [tid_0 * ts_0 * 2, tid_1 * ts_1]
+    0) (tid_0, tid_1, tid_2)
+      -> offsets [tid_0 * ts_0 * 2 + tid_2 * 64, tid_1 * ts_1]
          sizes [ts_0 * 2, ts_1]
-         strides [1, 2]
+         strides [1, 1]
          upper bounds [128, 256]
   )"));
 }
@@ -495,24 +493,22 @@ TEST_F(TilePropagationTest, CanPropagateFromReduceScatterToDotOp) {
       std::unique_ptr<TilingSpace> tiling_space,
       TilingSpace::Create(*HloFusionAdaptor::ForInstruction(fusion),
                           &mlir_context_));
-  ASSERT_OK_AND_ASSIGN(
-      auto rs_operands,
-      PropagateTileToInput(
-          *tiling_space, *root,
-          GetTestTile(*tiling_space, root->shape().dimensions()), 0));
+  ASSERT_OK_AND_ASSIGN(auto rs_operands,
+                       PropagateTileToInput(*tiling_space, *root,
+                                            tiling_space->tiled_roots()[0], 0));
   ASSERT_OK_AND_ASSIGN(auto dot_operands,
                        PropagateTileToInput(*tiling_space, *root->operand(0),
                                             rs_operands[0], 0));
   EXPECT_THAT(dot_operands, MatchToString(R"(
-    0) (tid_0, tid_1, tid_2)
-      -> offsets [tid_0 * ts_0 * 2, tid_2 * ts_2]
-         sizes [ts_0 * 2, ts_2]
+    0) (tid_0, tid_1, tid_2, tid_3)
+      -> offsets [tid_0 * ts_0 * 2 + tid_2 * 64, tid_3 * ts_3]
+         sizes [ts_0 * 2, ts_3]
          strides [1, 1]
          upper bounds [128, 256]
-    1) (tid_0, tid_1, tid_2)
-      -> offsets [tid_2 * ts_2, tid_1 * ts_1]
-         sizes [ts_2, ts_1]
-         strides [1, 2]
+    1) (tid_0, tid_1, tid_2, tid_3)
+      -> offsets [tid_3 * ts_3, tid_1 * ts_1]
+         sizes [ts_3, ts_1]
+         strides [1, 1]
          upper bounds [256, 512]
   )"));
 }
@@ -538,8 +534,8 @@ TEST_F(TilePropagationTest, CanPropagateToInputsOfReduceScatterOpScatterDim1) {
                        PropagateTileToInput(*tiling_space, *root,
                                             tiling_space->tiled_roots()[0], 0));
   EXPECT_THAT(tiled_operands, MatchToString(R"(
-    0) (tid_0, tid_1)
-      -> offsets [tid_0 * ts_0, tid_1 * ts_1 * 2]
+    0) (tid_0, tid_1, tid_2)
+      -> offsets [tid_0 * ts_0, tid_1 * ts_1 * 2 + tid_2 * 128]
          sizes [ts_0, ts_1 * 2]
          strides [1, 1]
          upper bounds [128, 256]
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc
index 2aab60bab0e..372e0c9243a 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc
@@ -921,17 +921,18 @@ TEST_P(TileAnalysisTest, CollectiveDotReduceScatter) {
     Dimensions:
       0 type: parallel size: 64 tile size: 8 dim ID:0 hlo: %rs = f32[64,512]{1,0} reduce-scatter(%dot), replica_groups={{0,1}}, dimensions={0}, to_apply=%add
       1 type: parallel size: 512 tile size: 32 dim ID:1 hlo: %rs = f32[64,512]{1,0} reduce-scatter(%dot), replica_groups={{0,1}}, dimensions={0}, to_apply=%add
-      2 type: sequential size: 256 tile size: 32 dim ID:2 hlo: %dot = f32[128,512]{1,0} dot(%p0, %p1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+      2 type: replica size: 2 tile size: 1 dim ID:2 hlo: %rs = f32[64,512]{1,0} reduce-scatter(%dot), replica_groups={{0,1}}, dimensions={0}, to_apply=%add
+      3 type: sequential size: 256 tile size: 32 dim ID:2 hlo: %dot = f32[128,512]{1,0} dot(%p0, %p1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
     Root tiles:
-      0 root tile:  offsets [tid_0 * 8, tid_1 * 32] sizes [8, 32] strides [1, 1] upper bounds [64, 512]
+      0 root tile:  offsets [tid_0 * 16 + tid_2 * 8, tid_1 * 32] sizes [8, 32] strides [1, 1] upper bounds [64, 512]
 
     Tiled HLO:
-      dot.tile_0 = dot(p0.1.tile_0, p1.1.tile_0)  offsets [tid_0 * 16, tid_1 * 32] sizes [16, 32] strides [1, 1] upper bounds [128, 512]
+      dot.tile_0 = dot(p0.1.tile_0, p1.1.tile_0)  offsets [tid_0 * 16 + tid_2 * 64, tid_1 * 32] sizes [16, 32] strides [1, 1] upper bounds [128, 512]
       region #0 {
-        p0.1.tile_0 = parameter(0)  offsets [tid_0 * 16, tid_2 * 32] sizes [16, 32] strides [1, 1] upper bounds [128, 256]
-        p1.1.tile_0 = parameter(1)  offsets [tid_2 * 32, tid_1 * 32] sizes [32, 32] strides [1, 1] upper bounds [256, 512]
+        p0.1.tile_0 = parameter(0)  offsets [tid_0 * 16 + tid_2 * 64, tid_3 * 32] sizes [16, 32] strides [1, 1] upper bounds [128, 256]
+        p1.1.tile_0 = parameter(1)  offsets [tid_3 * 32, tid_1 * 32] sizes [32, 32] strides [1, 1] upper bounds [256, 512]
       }
-      rs.tile_0 = reduce-scatter(dot.tile_0)  offsets [tid_0 * 8, tid_1 * 32] sizes [8, 32] strides [1, 1] upper bounds [64, 512]
+      rs.tile_0 = reduce-scatter(dot.tile_0)  offsets [tid_0 * 16 + tid_2 * 8, tid_1 * 32] sizes [8, 32] strides [1, 1] upper bounds [64, 512]
   )"));
 }
 
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc
index 0524dfce615..a82a13aa60e 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc
@@ -31,6 +31,7 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/str_join.h"
+#include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
@@ -96,16 +97,22 @@ bool IsIndexWiseVariadic(const HloInstruction& hlo) {
 }
 
 llvm::DenseMap<SymbolicExpr, SymbolicExpr> GetTileSizeReplacementMap(
-    const TilingSpace& tiling_space, absl::Span<const int64_t> tile_sizes) {
-  CHECK_EQ(tile_sizes.size(), tiling_space.dimensions().size());
+    const TilingSpace& tiling_space) {
   mlir::MLIRContext* ctx = tiling_space.mlir_context();
+  llvm::SmallVector<TilingSpace::DimensionInfo, 4> dims =
+      tiling_space.dimensions();
   llvm::DenseMap<SymbolicExpr, SymbolicExpr> replacement_map;
-  for (const auto& [index, dim] : llvm::enumerate(tiling_space.dimensions())) {
-    replacement_map[CreateSymbolExpr(dim.id.value(), tile_sizes.size(), ctx)] =
-        CreateSymbolicConstant(tile_sizes[index], ctx);
+  for (const TilingSpace::DimensionInfo& dim : dims) {
+    const int64_t tile_size = *dim.tile_size;
+    replacement_map[CreateSymbolExpr(dim.id.value(), dims.size(), ctx)] =
+        CreateSymbolicConstant(tile_size, ctx);
+    int64_t scale = 1;
+    if (dim.replica_dim.has_value()) {
+      scale = dims[dim.replica_dim->value()].dimension_size;
+    }
     // If the tile size is greater than or equal to the dimension size, then
     // the dimension is trivial and can be replaced with 0.
-    if (dim.dimension_size <= tile_sizes[index]) {
+    if (dim.dimension_size <= tile_size * scale) {
       replacement_map[CreateDimExpr(dim.id.value(), ctx)] =
           CreateSymbolicConstant(0, ctx);
     }
@@ -115,9 +122,13 @@ llvm::DenseMap<SymbolicExpr, SymbolicExpr> GetTileSizeReplacementMap(
 
 std::string TilingSpace::DimensionInfo::ToString() const {
   std::stringstream ss;
-  ss << id << " type: "
-     << (type == DimensionSemantics::kParallel ? "parallel" : "sequential")
-     << " size: " << dimension_size;
+  absl::string_view type_str = "sequential";
+  if (type == DimensionSemantics::kParallel) {
+    type_str = "parallel";
+  } else if (type == DimensionSemantics::kReplica) {
+    type_str = "replica";
+  }
+  ss << id << " type: " << type_str << " size: " << dimension_size;
   if (tile_size.has_value()) {
     ss << " tile size: " << *tile_size;
   }
@@ -153,6 +164,9 @@ void TilingSpace::ProcessInstruction(const HloInstruction& hlo) {
     case HloOpcode::kReduce:
       ProcessReduce(hlo);
       break;
+    case HloOpcode::kReduceScatter:
+      ProcessReduceScatter(hlo);
+      break;
     case HloOpcode::kScan:
       ProcessScan(hlo);
       break;
@@ -210,6 +224,29 @@ void TilingSpace::ProcessReduce(const HloInstruction& hlo) {
   }
 }
 
+// Adds a `kReplica` dimension for `reduce-scatter` and links it to the scatter
+// dimension.
+void TilingSpace::ProcessReduceScatter(const HloInstruction& hlo) {
+  if (hlo.shape().IsTuple()) {
+    return;
+  }
+  const auto* reduce_scatter = Cast<HloReduceScatterInstruction>(&hlo);
+  const int64_t scatter_dim = reduce_scatter->scatter_dimension();
+  auto it = hlo_to_dimension_.find(std::make_pair(&hlo, scatter_dim));
+  if (it == hlo_to_dimension_.end()) {
+    return;
+  }
+  const TiledDimId scatter_dim_id = it->second->id;
+  const int64_t num_replicas =
+      reduce_scatter->device_list()->num_devices_per_group();
+  if (num_replicas <= 1) {
+    return;
+  }
+  AppendDimension(&hlo, hlo.shape().dimensions().size(), num_replicas,
+                  DimensionSemantics::kReplica);
+  dimensions_[scatter_dim_id.value()].replica_dim = dimensions_.back().id;
+}
+
 // Register sequential dimensions and RTVars for a kRaggedDot instruction.
 //
 // --- kRaggedNonContracting ---
@@ -437,15 +474,21 @@ absl::Status TilingSpace::AssignTileSizes(
     return absl::InternalError(
         "Tile sizes have already been assigned to this tiling space.");
   }
-  CHECK_EQ(tile_sizes.size(), dimensions_.size());
   is_symbolic_ = false;
 
-  for (const auto& [index, size] : llvm::enumerate(tile_sizes)) {
-    dimensions_[index].tile_size = size;
+  int64_t tile_size_index = 0;
+  for (DimensionInfo& dim : dimensions_) {
+    if (dim.type == DimensionSemantics::kReplica) {
+      dim.tile_size = 1;
+    } else {
+      CHECK_LT(tile_size_index, tile_sizes.size());
+      dim.tile_size = tile_sizes[tile_size_index++];
+    }
   }
+  CHECK_EQ(tile_size_index, tile_sizes.size());
 
   llvm::DenseMap<SymbolicExpr, SymbolicExpr> replacement_map =
-      GetTileSizeReplacementMap(*this, tile_sizes);
+      GetTileSizeReplacementMap(*this);
 
   if (!constraint_.IsSatisfiedBy(tile_sizes)) {
     return absl::InvalidArgumentError(absl::StrFormat(
@@ -594,8 +637,7 @@ absl::StatusOr<std::unique_ptr<TilingSpace>> TilingSpace::Create(
     for (auto [index, dim] : llvm::enumerate(dims)) {
       const DimensionInfo& dim_info =
           tiling_space->GetDimensionInfo(root.instruction(), index);
-      dim_tiles.push_back(
-          tiling_space->GetDefaultRootDimTile(dim_info.id, dim));
+      dim_tiles.push_back(tiling_space->GetDefaultRootDimTile(dim_info, dim));
     }
     Tile root_tile{*tiling_space, std::move(dim_tiles)};
     const Shape& root_shape = root.shape();
@@ -664,7 +706,7 @@ std::unique_ptr<TilingSpace> TilingSpace::Clone(
     dim_tiles.reserve(root_tile.dim_tiles().size());
     for (const DimTile& dim_tile : root_tile.dim_tiles()) {
       dim_tiles.push_back(cloned->GetDefaultRootDimTile(
-          TiledDimId(dim_tile.size.GetValue() - num_dimensions()),
+          cloned->dimensions_[dim_tile.size.GetValue() - num_dimensions()],
           dim_tile.upper_bound.GetValue()));
     }
     cloned->tiled_roots_.push_back(Tile{*cloned, std::move(dim_tiles)});
@@ -677,11 +719,33 @@ std::unique_ptr<TilingSpace> TilingSpace::Clone(
   return cloned;
 }
 
-DimTile TilingSpace::GetDefaultRootDimTile(TiledDimId id,
+DimTile TilingSpace::GetDefaultRootDimTile(const DimensionInfo& dim_info,
                                            int64_t dim_size) const {
-  return GetDefaultDimTile(
-      id, CreateSymbolExpr(id.value(), num_dimensions(), mlir_context_),
-      dim_size);
+  SymbolicExpr tile_size =
+      CreateSymbolExpr(dim_info.id.value(), num_dimensions(), mlir_context_);
+  if (!dim_info.replica_dim.has_value()) {
+    return GetDefaultDimTile(dim_info.id, tile_size, dim_size);
+  }
+  const int64_t num_replicas =
+      dimensions_[dim_info.replica_dim->value()].dimension_size;
+  SymbolicExpr tile_id = CreateDimExpr(dim_info.id.value(), mlir_context_);
+  SymbolicExpr replica_id =
+      CreateDimExpr(dim_info.replica_dim->value(), mlir_context_);
+  return DimTile{(tile_id * (tile_size * num_replicas) + replica_id * tile_size)
+                     .Canonicalize(),
+                 tile_size, CreateSymbolicConstant(1, mlir_context_),
+                 CreateSymbolicConstant(dim_size, mlir_context_)};
+}
+
+int64_t TilingSpace::GetNumTiles(const DimensionInfo& dim) const {
+  CHECK(!is_symbolic_);
+  CHECK(dim.tile_size.has_value());
+  const int64_t step =
+      *dim.tile_size *
+      (dim.replica_dim.has_value()
+           ? dimensions_[dim.replica_dim->value()].dimension_size
+           : 1);
+  return llvm::divideCeil(dim.dimension_size, step);
 }
 
 int64_t TilingSpace::num_parallel_dimensions() const {
@@ -703,7 +767,7 @@ void TilingSpace::InitSimplificationIndexing() {
   for (const DimensionInfo& dim_info : dimensions_) {
     int64_t tile_size = dim_info.tile_size.value();
     CHECK_GT(tile_size, 0);
-    int64_t upper_bound = llvm::divideCeil(dim_info.dimension_size, tile_size);
+    int64_t upper_bound = GetNumTiles(dim_info);
     dim_vars_indexing_.push_back(IndexingMap::Variable{0, upper_bound - 1});
     // Even though ts_X must already be replaced with constants right now, we
     // initialize their bounds to [tile_size, tile_size] for completeness.
@@ -769,6 +833,9 @@ TilingSpace::GetValidTilings() {
   // limitation of `GetFlatTilingsForInputSpace`, we set the tile size to 1 here
   // and later replace with the actual dimension size.
   for (const auto& dim : dimensions_) {
+    if (dim.type == DimensionSemantics::kReplica) {
+      continue;
+    }
     if (dim.type == DimensionSemantics::kSequential &&
         dim.hlo->opcode() == HloOpcode::kReduce) {
       input_space.push_back(1);
@@ -780,11 +847,16 @@ TilingSpace::GetValidTilings() {
   ABSL_ASSIGN_OR_RETURN(auto flat_tilings, GetFlatTilingsForInputSpace(input_space));
 
   for (auto& flat_tiling : flat_tilings) {
-    for (const auto& [idx, dim] : llvm::enumerate(dimensions_)) {
+    int64_t idx = 0;
+    for (const auto& dim : dimensions_) {
+      if (dim.type == DimensionSemantics::kReplica) {
+        continue;
+      }
       if (dim.type == DimensionSemantics::kSequential &&
           dim.hlo->opcode() == HloOpcode::kReduce) {
         flat_tiling[idx] = dim.dimension_size;
       }
+      ++idx;
     }
   }
 
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h
index db7dc9b02c4..acd57f97058 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h
@@ -63,7 +63,7 @@ class TilingSpace {
   TilingSpace(const TilingSpace&) = delete;
   TilingSpace& operator=(const TilingSpace&) = delete;
 
-  enum class DimensionSemantics { kParallel, kSequential };
+  enum class DimensionSemantics { kParallel, kSequential, kReplica };
   struct DimensionInfo {
     // Unique ID for the dimension within the tiling space.
     TiledDimId id;
@@ -91,6 +91,13 @@ class TilingSpace {
     // Tile size for the dimension.
     std::optional<int64_t> tile_size;
 
+    // If this dimension is partitioned across peers (e.g. the scatter
+    // dimension of a reduce-scatter), the ID of the associated `kReplica`
+    // dimension. When set, the root tile offset along this dimension is split
+    // into an intra-shard tile index (`id`, with `dimension_size / (tile_size
+    // * R)` tiles) and a peer index (`*replica_dim`, with `R` values).
+    std::optional<TiledDimId> replica_dim;
+
     std::string ToString() const;
   };
 
@@ -160,6 +167,11 @@ class TilingSpace {
   const DimensionInfo& GetDimensionInfo(const HloInstruction& hlo,
                                         int64_t dim_position) const;
 
+  // Returns the number of tiles along `dim` after tile sizes have been
+  // assigned. For a dimension associated with a `replica_dim`, this is the
+  // number of intra-peer tiles (`CeilOfRatio(dimension_size, tile_size * R)`).
+  int64_t GetNumTiles(const DimensionInfo& dim) const;
+
   // Assigns tile sizes to the dimensions and checks if the constraints derived
   // from the operations are satisfied. That does NOT include backend-specific
   // constraints.
@@ -242,6 +254,7 @@ class TilingSpace {
 
   void ProcessDotLike(const HloInstruction& hlo);
   void ProcessReduce(const HloInstruction& hlo);
+  void ProcessReduceScatter(const HloInstruction& hlo);
   void ProcessScan(const HloInstruction& hlo);
   void ProcessDynamicSlice(const HloInstruction& hlo);
   void ProcessGetTupleElement(const HloInstruction& hlo);
@@ -256,8 +269,9 @@ class TilingSpace {
   void InitSimplificationIndexing();
 
   // Returns the default symbolic tile, in this space's context, for the root
-  // dimension `id` of size `dim_size`.
-  DimTile GetDefaultRootDimTile(TiledDimId id, int64_t dim_size) const;
+  // dimension `dim_info` of size `dim_size`.
+  DimTile GetDefaultRootDimTile(const DimensionInfo& dim_info,
+                                int64_t dim_size) const;
 
   // Maps from (hlo, dim_position) to the dimension info.
   absl::flat_hash_map<std::pair<const HloInstruction*, int64_t>,
@@ -308,7 +322,7 @@ bool IsIndexWiseVariadic(const HloInstruction& hlo);
 
 // Returns a symbol replacement map to set concrete tile sizes.
 llvm::DenseMap<SymbolicExpr, SymbolicExpr> GetTileSizeReplacementMap(
-    const TilingSpace& tiling_space, absl::Span<const int64_t> tile_sizes);
+    const TilingSpace& tiling_space);
 
 }  // namespace xla::gpu::experimental
 
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc
index 5c95e5f95c0..bfc37ddbaad 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc
@@ -569,6 +569,48 @@ TEST_F(TilingSpaceTest,
   EXPECT_EQ(tiling_space->num_parallel_dimensions(), 4);
 }
 
+TEST_F(TilingSpaceTest,
+       ReduceScatterRootGetsReplicaDimensionAndReplicaSplitRootTile) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    add {
+      lhs = f32[] parameter(0)
+      rhs = f32[] parameter(1)
+      ROOT add = f32[] add(lhs, rhs)
+    }
+    f {
+      p0 = f32[32,16] parameter(0)
+      ROOT rs = f32[8,16] reduce-scatter(p0),
+        replica_groups={{0,1,2,3}}, dimensions={0}, to_apply=add
+    }
+
+    ENTRY e {
+      p0 = f32[32,16] parameter(0)
+      ROOT fusion = f32[8,16] fusion(p0), kind=kCustom, calls=f
+    }
+  )");
+  auto fusion_adaptor = HloFusionAdaptor::ForInstruction(root);
+  ASSERT_OK_AND_ASSIGN(auto tiling_space,
+                       TilingSpace::Create(*fusion_adaptor, &mlir_context_));
+  EXPECT_EQ(tiling_space->num_dimensions(), 3);
+  EXPECT_EQ(tiling_space->num_parallel_dimensions(), 2);
+  EXPECT_THAT(*tiling_space, MatchString(R"(
+    Dimensions:
+        0 type: parallel size: 8 dim ID:0
+          hlo: %rs = f32[8,16]{1,0} reduce-scatter(%p0), replica_groups={{0,1,2,3}}, dimensions={0}, to_apply=%add
+        1 type: parallel size: 16 dim ID:1
+          hlo: %rs = f32[8,16]{1,0} reduce-scatter(%p0), replica_groups={{0,1,2,3}}, dimensions={0}, to_apply=%add
+        2 type: replica size: 4 dim ID:2
+          hlo: %rs = f32[8,16]{1,0} reduce-scatter(%p0), replica_groups={{0,1,2,3}}, dimensions={0}, to_apply=%add
+    Root tiles:
+      0 root tile:
+           offsets [tid_0 * ts_0 * 4 + tid_2 * ts_0, tid_1 * ts_1]
+           sizes [ts_0, ts_1]
+           strides [1, 1]
+           upper bounds [8, 16]
+  )"));
+}
+
 class TilingSpaceSameShapeMultiOutputTest : public TilingSpaceTest {
  protected:
   DebugOptions GetDebugOptionsForTest() const override {
diff --git a/third_party/xla/xla/codegen/tiling/experimental/validate_reshape_propagation.cc b/third_party/xla/xla/codegen/tiling/experimental/validate_reshape_propagation.cc
index 3d22acfa7b0..aa2d3f03f69 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/validate_reshape_propagation.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/validate_reshape_propagation.cc
@@ -163,9 +163,9 @@ absl::Status ValidateConfig(const Shape& from_shape, const Shape& to_shape,
       PropagateTileToInput(*tiling_space, *reshape, output_tile, 0));
   CHECK_EQ(input_tiles.size(), 1);
   Tile input_tile = input_tiles[0];
-  llvm::DenseMap<SymbolicExpr, SymbolicExpr> replacement_map =
-      GetTileSizeReplacementMap(*tiling_space, config.sizes);
   ABSL_RETURN_IF_ERROR(tiling_space->AssignTileSizes(config.sizes));
+  llvm::DenseMap<SymbolicExpr, SymbolicExpr> replacement_map =
+      GetTileSizeReplacementMap(*tiling_space);
 
   input_tile.Replace(replacement_map);
   input_tile.Simplify();
diff --git a/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.cc b/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.cc
index 62abc2cbb41..067430edc09 100644
--- a/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.cc
+++ b/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.cc
@@ -446,6 +446,7 @@ absl::StatusOr<SmallVector<Value>> EmitterContext::EvaluateTilingParameters(
   for (const int64_t dim : dim_ids) {
     switch (tiling_space.dimensions()[dim].type) {
       case ge::TilingSpace::DimensionSemantics::kParallel:
+      case ge::TilingSpace::DimensionSemantics::kReplica:
         dim_replacements[dim] = schedule_.GetPidExpr(dim);
         break;
       case ge::TilingSpace::DimensionSemantics::kSequential: {
diff --git a/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.cc b/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.cc
index ef6d000f647..bdddb6a2001 100644
--- a/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.cc
+++ b/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.cc
@@ -186,6 +186,8 @@ absl::StatusOr<llvm::SmallVector<int64_t>> GetTilingSpaceConcreteSizes(
         tile_sizes.push_back(parallel_tile_sizes[parallel_dim_count]);
         parallel_dim_count++;
         break;
+      case DimensionSemantics::kReplica:
+        break;
       case DimensionSemantics::kSequential: {
         if (dim.hlo->has_backend_config()) {
           ABSL_ASSIGN_OR_RETURN(xla::xtile::Tile config,

From b4f66794650303430bcb082eb4d44f00675a82b7 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Wed, 14 Oct 2026 02:48:25 -0700
Subject: [PATCH 0752/1000] [XLA:GPU] Migrate GpuModuleGlobals and GetGpuMutex
 to PerDeviceState.

PiperOrigin-RevId: 999236374
---
 .../xla/xla/backends/gpu/runtime/BUILD        |   4 +
 third_party/xla/xla/service/gpu/BUILD         |   9 ++
 .../xla/xla/service/gpu/gpu_executable.cc     |   4 +-
 .../xla/service/gpu/gpu_executable_test.cc    |  91 ++++++++++++++
 .../xla/xla/service/gpu/gpu_module_globals.cc | 111 +++++++++---------
 .../xla/xla/service/gpu/gpu_module_globals.h  |  31 +++--
 .../xla/service/gpu/stream_executor_util.cc   |  34 +++---
 .../service/gpu/stream_executor_util_test.cc  |  65 ++++++++++
 8 files changed, 258 insertions(+), 91 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index c77f636dc5c..c460bcfebcc 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -5435,6 +5435,7 @@ xla_test(
 cc_library(
     name = "device_slot",
     hdrs = ["device_slot.h"],
+    compatible_with = get_compatible_with_portable(),
     visibility = ["//visibility:private"],
     deps = [
         "@com_google_absl//absl/base",
@@ -5448,6 +5449,7 @@ cc_library(
     name = "vector_storage",
     srcs = ["vector_storage.cc"],
     hdrs = ["vector_storage.h"],
+    compatible_with = get_compatible_with_portable(),
     visibility = ["//visibility:private"],
     deps = [
         ":device_slot",
@@ -5471,6 +5473,7 @@ cc_library(
     name = "cow_storage",
     srcs = ["cow_storage.cc"],
     hdrs = ["cow_storage.h"],
+    compatible_with = get_compatible_with_portable(),
     visibility = ["//visibility:private"],
     deps = [
         ":device_slot",
@@ -5500,6 +5503,7 @@ cc_library(
     name = "per_device_state",
     srcs = ["per_device_state.cc"],
     hdrs = ["per_device_state.h"],
+    compatible_with = get_compatible_with_portable(),
     deps = [
         ":cow_storage",
         ":device_slot",
diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 98221e9f0be..4784424d0c9 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -850,6 +850,7 @@ cc_library(
         "//xla/backends/gpu/runtime:collective_params",
         "//xla/backends/gpu/runtime:command_buffer_conversion_pass",
         "//xla/backends/gpu/runtime:execution_stream_id",
+        "//xla/backends/gpu/runtime:per_device_state",
         "//xla/backends/gpu/runtime:sequential_thunk",
         "//xla/backends/gpu/runtime:thunk",
         "//xla/backends/gpu/runtime:thunk_buffer_debug_pass",
@@ -947,6 +948,7 @@ xla_test(
     backends = ["gpu"],
     deps = [
         ":buffer_allocations",
+        ":dense_data_intermediate",
         ":gpu_executable",
         ":gpu_executable_buffer_allocator",
         ":gpu_executable_proto_cc",
@@ -988,6 +990,7 @@ xla_test(
         "//xla/stream_executor:device_description",
         "//xla/stream_executor:mock_stream",
         "//xla/stream_executor:mock_stream_executor",
+        "//xla/stream_executor:module_spec",
         "//xla/stream_executor:semantic_version",
         "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor/abi:executable_abi_version",
@@ -1006,6 +1009,7 @@ xla_test(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest_main",
         "@riegeli//riegeli/bytes:cfile_reader",
@@ -3056,6 +3060,7 @@ cc_library(
         "//xla:shape_util",
         "//xla:util",
         "//xla:xla_data_proto_cc",
+        "//xla/backends/gpu/runtime:per_device_state",
         "//xla/hlo/ir:hlo",
         "//xla/service:hlo_module_config",
         "//xla/stream_executor:data_type",
@@ -3076,6 +3081,7 @@ cc_library(
         "//xla/tsl/util/proto:proto_utils",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
@@ -3100,11 +3106,14 @@ xla_cc_test(
         ":stream_executor_util",
         "//xla:autotuning_proto_cc",
         "//xla/service:hlo_module_config",
+        "//xla/stream_executor:mock_stream_executor",
+        "//xla/tsl/platform:env",
         "//xla/tsl/protobuf:dnn_proto_cc",
         "//xla/tsl/util/proto:proto_utils",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
         "@com_google_googletest//:gtest_main",
     ],
diff --git a/third_party/xla/xla/service/gpu/gpu_executable.cc b/third_party/xla/xla/service/gpu/gpu_executable.cc
index 3b7f75a9bef..a5f2eaf1efa 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable.cc
@@ -510,7 +510,9 @@ GpuExecutable::GpuExecutable(
       debug_buffer_assignment_show_max_(
           debug_options.xla_debug_buffer_assignment_show_max()),
       constants_(std::move(constants)),
-      module_globals_(std::make_unique<GpuModuleGlobals>(binary_, constants_)),
+      module_globals_(std::make_unique<GpuModuleGlobals>(
+          binary_, constants_,
+          gpu_topology.has_value() ? gpu_topology->num_devices_per_host() : 0)),
       output_info_(std::move(output_info)),
       enable_debug_info_manager_(enable_debug_info_manager),
       thunk_sequence_proto_(std::move(thunk_sequence_proto)),
diff --git a/third_party/xla/xla/service/gpu/gpu_executable_test.cc b/third_party/xla/xla/service/gpu/gpu_executable_test.cc
index dcefa19792c..fdbc3615056 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable_test.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable_test.cc
@@ -31,6 +31,8 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
+#include "absl/synchronization/blocking_counter.h"
+#include "absl/synchronization/notification.h"
 #include "absl/types/span.h"
 #include "riegeli/bytes/cfile_reader.h"
 #include "riegeli/bytes/string_reader.h"
@@ -60,6 +62,7 @@ limitations under the License.
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/buffer_value.h"
 #include "xla/service/gpu/buffer_allocations.h"
+#include "xla/service/gpu/dense_data_intermediate.h"
 #include "xla/service/gpu/gpu_executable.pb.h"
 #include "xla/service/gpu/gpu_executable_buffer_allocator.h"
 #include "xla/service/gpu/gpu_module_globals.h"
@@ -81,6 +84,7 @@ limitations under the License.
 #include "xla/stream_executor/gpu/tma_metadata.h"
 #include "xla/stream_executor/mock_stream.h"
 #include "xla/stream_executor/mock_stream_executor.h"
+#include "xla/stream_executor/module_spec.h"
 #include "xla/stream_executor/semantic_version.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/lib/core/status_test_util.h"
@@ -141,6 +145,7 @@ TEST(GpuModuleGlobalsTest, EmptyBinarySkipsModuleLoading) {
   GpuModuleGlobals module_globals(binary, constants);
 
   EXPECT_CALL(stream, parent()).WillOnce(Return(&executor));
+  EXPECT_CALL(executor, device_ordinal()).WillRepeatedly(Return(0));
   EXPECT_CALL(executor, GetPlatform()).Times(0);
   EXPECT_CALL(executor, LoadModule(::testing::_)).Times(0);
 
@@ -149,6 +154,92 @@ TEST(GpuModuleGlobalsTest, EmptyBinarySkipsModuleLoading) {
   EXPECT_TRUE(resolved->empty());
 }
 
+TEST(GpuModuleGlobalsTest, ConcurrentResolveAcrossDevices) {
+  constexpr int kNumDevices = 8;
+  constexpr int kThreadsPerDevice = 4;
+
+  std::vector<uint8_t> binary = {1, 2, 3, 4};
+  std::vector<GpuModuleGlobals::ConstantInfo> constants = {
+      GpuModuleGlobals::ConstantInfo{
+          /*symbol_name=*/"const_0",
+          /*content=*/DenseDataIntermediate::Own({10, 20, 30, 40}),
+          /*allocation_index=*/0}};
+
+  std::vector<std::unique_ptr<se::MockStreamExecutor>> executors;
+  std::vector<std::unique_ptr<se::MockStream>> streams;
+  for (int d = 0; d < kNumDevices; ++d) {
+    auto executor = std::make_unique<se::MockStreamExecutor>();
+    auto stream = std::make_unique<se::MockStream>();
+    const void* handle_id =
+        reinterpret_cast<const void*>(static_cast<uintptr_t>(0x1000 + d));
+    se::ModuleHandle module_handle(handle_id);
+    se::DeviceAddressBase symbol_addr(
+        reinterpret_cast<void*>(static_cast<uintptr_t>(0x2000 + d * 16)), 4);
+
+    EXPECT_CALL(*stream, parent()).WillRepeatedly(Return(executor.get()));
+    EXPECT_CALL(*executor, device_ordinal()).WillRepeatedly(Return(d));
+    EXPECT_CALL(*executor, LoadModule(::testing::_))
+        .Times(1)
+        .WillOnce(Return(module_handle));
+    EXPECT_CALL(*executor, GetSymbol("const_0", module_handle))
+        .Times(1)
+        .WillOnce(Return(symbol_addr));
+    EXPECT_CALL(*stream, Memcpy(::testing::A<se::DeviceAddressBase*>(),
+                                ::testing::A<const void*>(), 4))
+        .Times(1)
+        .WillOnce(Return(absl::OkStatus()));
+    EXPECT_CALL(*stream, BlockHostUntilDone())
+        .Times(1)
+        .WillOnce(Return(absl::OkStatus()));
+    EXPECT_CALL(*executor, UnloadModule(module_handle))
+        .Times(1)
+        .WillOnce(Return(true));
+
+    executors.push_back(std::move(executor));
+    streams.push_back(std::move(stream));
+  }
+
+  GpuModuleGlobals module_globals(binary, constants,
+                                  /*devices_per_host=*/kNumDevices);
+
+  tsl::thread::ThreadPool pool(tsl::Env::Default(), "resolve_test",
+                               kNumDevices * kThreadsPerDevice);
+  absl::BlockingCounter ready(kNumDevices * kThreadsPerDevice);
+  absl::Notification go;
+  absl::BlockingCounter done(kNumDevices * kThreadsPerDevice);
+
+  std::vector<const GpuModuleGlobals::BufferAllocToDeviceMemoryMap*> results(
+      kNumDevices * kThreadsPerDevice, nullptr);
+
+  for (int d = 0; d < kNumDevices; ++d) {
+    for (int t = 0; t < kThreadsPerDevice; ++t) {
+      int idx = d * kThreadsPerDevice + t;
+      pool.Schedule([&, d, idx]() {
+        ready.DecrementCount();
+        go.WaitForNotification();
+        ASSERT_OK_AND_ASSIGN(results[idx],
+                             module_globals.Resolve(streams[d].get()));
+        done.DecrementCount();
+      });
+    }
+  }
+
+  ready.Wait();
+  go.Notify();
+  done.Wait();
+
+  for (int d = 0; d < kNumDevices; ++d) {
+    const auto* base = results[d * kThreadsPerDevice];
+    ASSERT_NE(base, nullptr);
+    ASSERT_TRUE(base->contains(0));
+    EXPECT_EQ(base->at(0).opaque(),
+              reinterpret_cast<void*>(static_cast<uintptr_t>(0x2000 + d * 16)));
+    for (int t = 1; t < kThreadsPerDevice; ++t) {
+      EXPECT_EQ(results[d * kThreadsPerDevice + t], base);
+    }
+  }
+}
+
 TEST_F(GpuExecutableTest, ExecuteWithWatchdogs) {
   constexpr absl::string_view kHlo = R"(
     HloModule watchdog_test
diff --git a/third_party/xla/xla/service/gpu/gpu_module_globals.cc b/third_party/xla/xla/service/gpu/gpu_module_globals.cc
index 0a53d039ffa..df2a7299342 100644
--- a/third_party/xla/xla/service/gpu/gpu_module_globals.cc
+++ b/third_party/xla/xla/service/gpu/gpu_module_globals.cc
@@ -15,9 +15,8 @@ limitations under the License.
 
 #include "xla/service/gpu/gpu_module_globals.h"
 
-#include <memory>
+#include <optional>
 #include <string>
-#include <utility>
 
 #include "absl/base/nullability.h"
 #include "absl/container/flat_hash_map.h"
@@ -27,8 +26,8 @@ limitations under the License.
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
-#include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/literal.h"
 #include "xla/map_util.h"
@@ -81,70 +80,70 @@ GpuModuleGlobals::ConstantInfo::FromProto(
 absl::StatusOr<const GpuModuleGlobals::BufferAllocToDeviceMemoryMap*>
 GpuModuleGlobals::Resolve(se::Stream* stream) {
   se::StreamExecutor* executor = stream->parent();
+  int device_ordinal = executor->device_ordinal();
 
-  absl::MutexLock lock(mutex_);
-  auto it = globals_.find(executor);
-  if (it != globals_.end()) {
-    return it->second.get();
-  }
+  ABSL_RETURN_IF_ERROR(device_states_.GetOrCreateAndInitialize(
+      device_ordinal, [&](DeviceState* state) -> absl::Status {
+        se::MultiModuleLoaderSpec module_spec;
+        if (!binary_.empty()) {
+          module_spec.AddCudaCubinInMemory(binary_);
+        }
 
-  se::MultiModuleLoaderSpec module_spec;
-  if (!binary_.empty()) {
-    module_spec.AddCudaCubinInMemory(binary_);
-  }
+        se::ModuleHandle module_handle;
+        // There is no module to load when constants compilation produced no
+        // binary.
+        if (binary_.empty()) {
+          TF_RET_CHECK(constants_.empty())
+              << "Constants metadata is present without a constants binary";
+        } else {
+          ABSL_ASSIGN_OR_RETURN(module_handle, executor->LoadModule(module_spec));
+        }
 
-  auto globals = std::make_unique<BufferAllocToDeviceMemoryMap>();
-  se::ModuleHandle module_handle;
-  // There is no module to load when constants compilation produced no binary.
-  if (binary_.empty()) {
-    TF_RET_CHECK(constants_.empty())
-        << "Constants metadata is present without a constants binary";
-  } else {
-    ABSL_ASSIGN_OR_RETURN(module_handle, executor->LoadModule(module_spec));
-  }
+        // A flag signalling if constant initialization submitted memcpy
+        // operations to the `stream`.
+        int submitted_mem_copies = 0;
 
-  // A flag signalling if constant initialization submitted memcpy operations
-  // to the `stream`.
-  int submitted_mem_copies = 0;
+        for (const ConstantInfo& info : constants_) {
+          absl::StatusOr<se::DeviceAddressBase> global_status;
+          if (static_cast<bool>(module_handle)) {
+            global_status =
+                executor->GetSymbol(info.symbol_name, module_handle);
+          }
 
-  for (const ConstantInfo& info : constants_) {
-    absl::StatusOr<se::DeviceAddressBase> global_status;
-    if (static_cast<bool>(module_handle)) {
-      global_status = executor->GetSymbol(info.symbol_name, module_handle);
-    }
+          se::DeviceAddressBase global;
 
-    se::DeviceAddressBase global;
+          CHECK(static_cast<bool>(module_handle) && global_status.ok());
+          // The constant was defined in the PTX and has been allocated by the
+          // CUDA driver.
+          global = *global_status;
+          XLA_VLOG_DEVICE(3, device_ordinal) << absl::StreamFormat(
+              "Resolved global %s to %p", info.symbol_name, global.opaque());
 
-    CHECK(static_cast<bool>(module_handle) && global_status.ok());
-    // The constant was defined in the PTX and has been allocated by the CUDA
-    // driver.
-    global = *global_status;
-    XLA_VLOG_DEVICE(3, executor->device_ordinal()) << absl::StreamFormat(
-        "Resolved global %s to %p", info.symbol_name, global.opaque());
+          if (!info.content.span().empty()) {
+            // This means the constant did not have an initializer in the PTX
+            // and therefore must be initialized by XLA here.
+            ABSL_RETURN_IF_ERROR(stream->Memcpy(&global, info.content.span().data(),
+                                           info.content.span().size()));
+            submitted_mem_copies = true;
+          }
 
-    if (!info.content.span().empty()) {
-      // This means the constant did not have an initializer in the PTX and
-      // therefore must be initialized by XLA here.
-      ABSL_RETURN_IF_ERROR(stream->Memcpy(&global, info.content.span().data(),
-                                     info.content.span().size()));
-      submitted_mem_copies = true;
-    }
+          if (info.allocation_index != -1) {
+            InsertOrDie(&state->globals, info.allocation_index, global);
+          }
+        }
 
-    if (info.allocation_index != -1) {
-      InsertOrDie(globals.get(), info.allocation_index, global);
-    }
-  }
+        // Wait for the completion of all host->device transfers, to guarantee
+        // that destructor will not race with any operations in flight
+        // (deallocate xla::Literal owned by the HLO module).
+        if (submitted_mem_copies) {
+          CHECK_OK(stream->BlockHostUntilDone());
+        }
 
-  // Wait for the completion of all host->device transfers, to guarantee that
-  // destructor will not race with any operations in flight (deallocate
-  // xla::Literal owned by the HLO module).
-  if (submitted_mem_copies) {
-    CHECK_OK(stream->BlockHostUntilDone());
-  }
+        state->module_handle.emplace(executor, module_handle);
+        return absl::OkStatus();
+      }));
 
-  module_handles_.emplace(executor,
-                          se::ScopedModuleHandle(executor, module_handle));
-  return globals_.emplace(executor, std::move(globals)).first->second.get();
+  return &device_states_.Find(device_ordinal)->globals;
 }
 
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/service/gpu/gpu_module_globals.h b/third_party/xla/xla/service/gpu/gpu_module_globals.h
index c4c366c585a..e67eba9d379 100644
--- a/third_party/xla/xla/service/gpu/gpu_module_globals.h
+++ b/third_party/xla/xla/service/gpu/gpu_module_globals.h
@@ -17,15 +17,14 @@ limitations under the License.
 #define XLA_SERVICE_GPU_GPU_MODULE_GLOBALS_H_
 
 #include <cstdint>
-#include <memory>
+#include <optional>
 #include <string>
 #include <vector>
 
 #include "absl/base/nullability.h"
-#include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/status/statusor.h"
-#include "absl/synchronization/mutex.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/gpu/dense_data_intermediate.h"
 #include "xla/service/gpu/gpu_executable.pb.h"
@@ -61,27 +60,27 @@ class GpuModuleGlobals {
   };
 
   GpuModuleGlobals(const std::vector<uint8_t>& binary,
-                   const std::vector<ConstantInfo>& constants)
-      : binary_(binary), constants_(constants) {}
+                   const std::vector<ConstantInfo>& constants,
+                   int devices_per_host = 0)
+      : binary_(binary),
+        constants_(constants),
+        device_states_(devices_per_host) {}
 
   // Loads the executable module for `stream` and initializes constant globals.
-  // Loaded modules and resolved global addresses are cached per executor.
+  // Loaded modules and resolved global addresses are cached per device ordinal.
   absl::StatusOr<const BufferAllocToDeviceMemoryMap*> Resolve(
       se::Stream* stream);
 
  private:
+  struct DeviceState {
+    // Required to keep the loaded module alive until this helper is destroyed.
+    std::optional<se::ScopedModuleHandle> module_handle;
+    BufferAllocToDeviceMemoryMap globals;
+  };
+
   const std::vector<uint8_t>& binary_;
   const std::vector<ConstantInfo>& constants_;
-
-  absl::Mutex mutex_;
-  // Cache of module handles. Required to keep loaded modules alive until this
-  // helper is destroyed.
-  absl::flat_hash_map<se::StreamExecutor*, se::ScopedModuleHandle>
-      module_handles_ ABSL_GUARDED_BY(mutex_);
-  // Cache of constant buffer allocation maps used by `Resolve`.
-  absl::flat_hash_map<se::StreamExecutor*,
-                      std::unique_ptr<BufferAllocToDeviceMemoryMap>>
-      globals_ ABSL_GUARDED_BY(mutex_);
+  PerDeviceState<DeviceState> device_states_;
 };
 
 }  // namespace gpu
diff --git a/third_party/xla/xla/service/gpu/stream_executor_util.cc b/third_party/xla/xla/service/gpu/stream_executor_util.cc
index e67c505f303..83a14d69de4 100644
--- a/third_party/xla/xla/service/gpu/stream_executor_util.cc
+++ b/third_party/xla/xla/service/gpu/stream_executor_util.cc
@@ -19,7 +19,6 @@ limitations under the License.
 #include <cstdint>
 #include <iterator>
 #include <limits>
-#include <map>
 #include <memory>
 #include <optional>
 #include <random>
@@ -31,7 +30,8 @@ limitations under the License.
 #include <vector>
 
 #include "absl/algorithm/container.h"
-#include "absl/base/const_init.h"
+#include "absl/base/no_destructor.h"
+#include "absl/base/optimization.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
@@ -43,6 +43,7 @@ limitations under the License.
 #include "absl/types/span.h"
 #include "Eigen/Core"
 #include "xla/autotuning.pb.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/layout.h"
 #include "xla/layout_util.h"
@@ -80,7 +81,6 @@ using tsl::profiler::TraceMeLevel;
 namespace xla {
 namespace gpu {
 
-
 namespace {
 
 using se::dnn::DataLayout;
@@ -343,21 +343,19 @@ FindVectorizedFeatureDims(const ConvolutionDimensionNumbers& dnums,
 
 // Returns a mutex that can be used to lock the given stream executor.
 absl::Mutex& GetGpuMutex(const se::StreamExecutor* stream_exec) {
-  static absl::Mutex mu(absl::kConstInit);
-  // se::Platform*s are global singletons guaranteed to live forever.
-  static auto* mutexes =
-      new std::map<std::pair<const se::Platform*, /*device_ordinal*/ int64_t>,
-                   absl::Mutex>();
-
-  absl::MutexLock global_lock(mu);
-  auto it = mutexes
-                ->emplace(std::piecewise_construct,
-                          std::make_tuple(stream_exec->GetPlatform(),
-                                          stream_exec->device_ordinal()),
-                          std::make_tuple())
-                .first;
-
-  return it->second;
+  // Pre-allocate 16 slots to cover standard multi-GPU host topologies
+  // (typically 8 or 16 GPUs per host) with O(1) flat vector indexing. Ordinals
+  // >= 16 fall back to PerDeviceState's dynamic copy-on-write storage.
+  static absl::NoDestructor<PerDeviceState<absl::Mutex>> mutexes(
+      /*num_devices=*/16);
+  int device_ordinal = stream_exec->device_ordinal();
+  if (ABSL_PREDICT_FALSE(device_ordinal < 0)) {
+    static absl::NoDestructor<absl::Mutex> fallback_mu;
+    return *fallback_mu;
+  }
+  absl::StatusOr<absl::Mutex*> mu = mutexes->GetOrCreate(device_ordinal);
+  CHECK_OK(mu.status());
+  return **mu;
 }
 
 absl::StatusOr<std::unique_ptr<se::Kernel>> CreateKernel(
diff --git a/third_party/xla/xla/service/gpu/stream_executor_util_test.cc b/third_party/xla/xla/service/gpu/stream_executor_util_test.cc
index cf22f97fea8..0c4ae1f91e4 100644
--- a/third_party/xla/xla/service/gpu/stream_executor_util_test.cc
+++ b/third_party/xla/xla/service/gpu/stream_executor_util_test.cc
@@ -16,6 +16,8 @@ limitations under the License.
 #include "xla/service/gpu/stream_executor_util.h"
 
 #include <cstdint>
+#include <memory>
+#include <utility>
 #include <vector>
 
 #include <gmock/gmock.h>
@@ -23,10 +25,16 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
 #include "absl/status/statusor.h"
+#include "absl/synchronization/blocking_counter.h"
+#include "absl/synchronization/mutex.h"
+#include "absl/synchronization/notification.h"
 #include "absl/time/time.h"
 #include "xla/autotuning.pb.h"
 #include "xla/service/gpu/cublas_cudnn.h"
 #include "xla/service/hlo_module_config.h"
+#include "xla/stream_executor/mock_stream_executor.h"
+#include "xla/tsl/platform/env.h"
+#include "xla/tsl/platform/threadpool.h"
 #include "xla/tsl/protobuf/dnn.pb.h"
 #include "xla/tsl/util/proto/proto_utils.h"
 
@@ -111,6 +119,63 @@ TEST(StreamExecutorUtilTest, CudnnConvKindFromProto) {
               StatusIs(absl::StatusCode::kInternal));
 }
 
+TEST(StreamExecutorUtilTest, GetGpuMutexConcurrentPerDevice) {
+  constexpr int kNumDevices = 8;
+  constexpr int kThreadsPerDevice = 4;
+
+  std::vector<std::unique_ptr<se::MockStreamExecutor>> executors;
+  for (int d = 0; d < kNumDevices; ++d) {
+    auto exec = std::make_unique<se::MockStreamExecutor>();
+    EXPECT_CALL(*exec, device_ordinal()).WillRepeatedly(::testing::Return(d));
+    executors.push_back(std::move(exec));
+  }
+
+  tsl::thread::ThreadPool pool(tsl::Env::Default(), "gpu_mutex_test",
+                               kNumDevices * kThreadsPerDevice);
+  absl::BlockingCounter ready(kNumDevices * kThreadsPerDevice);
+  absl::Notification go;
+  absl::BlockingCounter done(kNumDevices * kThreadsPerDevice);
+
+  std::vector<absl::Mutex*> mutex_ptrs(kNumDevices * kThreadsPerDevice,
+                                       nullptr);
+
+  for (int d = 0; d < kNumDevices; ++d) {
+    for (int t = 0; t < kThreadsPerDevice; ++t) {
+      int idx = d * kThreadsPerDevice + t;
+      pool.Schedule([&, d, idx]() {
+        ready.DecrementCount();
+        go.WaitForNotification();
+        absl::Mutex& mu = GetGpuMutex(executors[d].get());
+        {
+          absl::MutexLock lock(mu);
+          mutex_ptrs[idx] = &mu;
+        }
+        done.DecrementCount();
+      });
+    }
+  }
+
+  ready.Wait();
+  go.Notify();
+  done.Wait();
+
+  for (int d = 0; d < kNumDevices; ++d) {
+    absl::Mutex* base = mutex_ptrs[d * kThreadsPerDevice];
+    ASSERT_NE(base, nullptr);
+    for (int t = 1; t < kThreadsPerDevice; ++t) {
+      EXPECT_EQ(mutex_ptrs[d * kThreadsPerDevice + t], base);
+    }
+    for (int other_d = d + 1; other_d < kNumDevices; ++other_d) {
+      EXPECT_NE(base, mutex_ptrs[other_d * kThreadsPerDevice]);
+    }
+  }
+
+  se::MockStreamExecutor neg_exec;
+  EXPECT_CALL(neg_exec, device_ordinal()).WillRepeatedly(::testing::Return(-1));
+  absl::Mutex& fallback_mu = GetGpuMutex(&neg_exec);
+  EXPECT_EQ(&fallback_mu, &GetGpuMutex(&neg_exec));
+}
+
 }  // namespace
 
 }  // namespace xla::gpu

From 94d0618cf02ce82ed8e27ca823ff5cfa4d31de8c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2026 02:48:58 -0700
Subject: [PATCH 0753/1000] Automated Code Change

PiperOrigin-RevId: 999236629
---
 third_party/xla/xla/service/gpu/model/BUILD                  | 5 ++++-
 .../xla/xla/service/gpu/model/gpu_dot_fusion_cost_model.cc   | 2 +-
 .../xla/service/gpu/model/gpu_dot_fusion_cost_model_test.cc  | 1 +
 .../xla/xla/service/gpu/model/gpu_hlo_cost_analysis_test.cc  | 1 +
 .../service/gpu/model/gpu_indexing_performance_model_test.cc | 1 +
 .../xla/xla/service/gpu/model/sol_latency_estimator_test.cc  | 1 +
 .../xla/service/gpu/model/triton_emitter_constraints_test.cc | 1 +
 7 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/BUILD b/third_party/xla/xla/service/gpu/model/BUILD
index f132263c0d6..4897fb4064c 100644
--- a/third_party/xla/xla/service/gpu/model/BUILD
+++ b/third_party/xla/xla/service/gpu/model/BUILD
@@ -109,6 +109,7 @@ xla_cc_test(
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/time",
         "@com_google_googletest//:gtest",
     ],
@@ -499,6 +500,7 @@ cc_library(
         "//xla:util",
         "//xla:xla_data_proto_cc",
         "//xla/codegen/xtile:block_level_parameters",
+        "//xla/codegen/xtile:xtile_config_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/service/gpu:backend_configs_cc",
         "//xla/stream_executor:device_description",
@@ -506,7 +508,6 @@ cc_library(
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/container:flat_hash_set",
-        "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
@@ -524,6 +525,7 @@ xla_cc_test(
     deps = [
         ":gpu_dot_fusion_cost_model",
         ":gpu_performance_model_base",
+        "//xla:xla_data_proto_cc",
         "//xla/codegen/xtile:block_level_parameters",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
@@ -696,6 +698,7 @@ xla_cc_test(
     srcs = ["triton_emitter_constraints_test.cc"],
     deps = [
         ":triton_emitter_constraints",
+        "//xla:xla_proto_cc",
         "//xla/codegen/tiling:symbolic_tile_analysis",
         "//xla/codegen/tiling:tiling_specification",
         "//xla/codegen/tiling/experimental:tile",
diff --git a/third_party/xla/xla/service/gpu/model/gpu_dot_fusion_cost_model.cc b/third_party/xla/xla/service/gpu/model/gpu_dot_fusion_cost_model.cc
index b5302bcb434..5f8770a27d7 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_dot_fusion_cost_model.cc
+++ b/third_party/xla/xla/service/gpu/model/gpu_dot_fusion_cost_model.cc
@@ -25,7 +25,6 @@ limitations under the License.
 #include <vector>
 
 #include "absl/container/flat_hash_set.h"
-#include "absl/container/inlined_vector.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
@@ -35,6 +34,7 @@ limitations under the License.
 #include "absl/time/time.h"
 #include "absl/types/span.h"
 #include "xla/codegen/xtile/block_level_parameters.h"
+#include "xla/codegen/xtile/xtile_config.pb.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/primitive_util.h"
diff --git a/third_party/xla/xla/service/gpu/model/gpu_dot_fusion_cost_model_test.cc b/third_party/xla/xla/service/gpu/model/gpu_dot_fusion_cost_model_test.cc
index 36c56cb7f6c..d8464ef5797 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_dot_fusion_cost_model_test.cc
+++ b/third_party/xla/xla/service/gpu/model/gpu_dot_fusion_cost_model_test.cc
@@ -34,6 +34,7 @@ limitations under the License.
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/tsl/platform/statusor.h"
+#include "xla/xla_data.pb.h"
 
 namespace xla {
 namespace gpu {
diff --git a/third_party/xla/xla/service/gpu/model/gpu_hlo_cost_analysis_test.cc b/third_party/xla/xla/service/gpu/model/gpu_hlo_cost_analysis_test.cc
index 99ebf183ee2..c3c6d74b91d 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_hlo_cost_analysis_test.cc
+++ b/third_party/xla/xla/service/gpu/model/gpu_hlo_cost_analysis_test.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "xla/service/gpu/model/gpu_hlo_cost_analysis.h"
 
 #include <cstdint>
+#include <string>
 #include <vector>
 
 #include <gmock/gmock.h>
diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc
index aee1cad361a..74cda88982d 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc
+++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model_test.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <string>
 #include <utility>
 #include <variant>
 
diff --git a/third_party/xla/xla/service/gpu/model/sol_latency_estimator_test.cc b/third_party/xla/xla/service/gpu/model/sol_latency_estimator_test.cc
index afe1bc297d8..5a75dfa6284 100644
--- a/third_party/xla/xla/service/gpu/model/sol_latency_estimator_test.cc
+++ b/third_party/xla/xla/service/gpu/model/sol_latency_estimator_test.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/status_macros.h"
+#include "absl/strings/string_view.h"
 #include "absl/time/time.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_opcode.h"
diff --git a/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc b/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc
index fe4d40c7fed..96d32c61608 100644
--- a/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc
+++ b/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc
@@ -45,6 +45,7 @@ limitations under the License.
 #include "xla/service/instruction_fusion.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/tsl/platform/test.h"
+#include "xla/xla.pb.h"
 
 namespace xla {
 namespace gpu {

From efd1e0e64a642be925df64413049b50b93f6456a Mon Sep 17 00:00:00 2001
From: Sergey Kozub <skozub@nvidia.com>
Date: Wed, 14 Oct 2026 02:57:35 -0700
Subject: [PATCH 0754/1000] [XLA:GPU] TensorIR: support multi-output fusions,
 copies and FP8; refine the tiling cost model.

This imports NVIDIA's commit bff9155343 (github.com/sergey-kozub/xla
branch experimental) on top of the TensorIR fusion emitter.

Multi-output fusions: the importer flattens the root tuple, so the
nv_tensor_ir.graph has one result per tuple element. AttachLayoutStrides
and AttachBufferAlignments attach their attributes to every result, the
emitter expects one kernel argument per fused parameter plus one per
output, and the support check admits a flat root tuple.

Copies: mhlo.copy only changes the layout, which lives on the graph
boundary as stride attributes, so it legalizes to its operand. F8E4M3FN
and F8E5M2 are accepted as element types.

Tiling cost model (third_party/tensor_ir/patches/tiling_analysis.patch):
memory cost is counted in 32-byte sectors instead of 128-byte cache
lines; stores that only partially cover a sector count twice
(read-modify-write), and vectorizable store dimensions weigh twice as
much when ordering dimensions; for reductions the tile-size lower bound
drops below one warp while the grid still runs in a single wave, based
on the occupancy estimated from the register storage; the search-space
reduction stops once the per-dimension limits sum to at most the maximum
tile size; element sizes round sub-byte types up.

Changes relative to NVIDIA's commit:
- The kernel-cache continuation keeps capturing the fusion name by
  value.
- support depends on shape_util directly (layering check).

Co-authored-by: Christian Sigg <csigg@google.com>
PiperOrigin-RevId: 999240785
---
 .../tensor_ir/patches/tiling_analysis.patch   | 190 +++++++++++-------
 .../xla/backends/gpu/codegen/tensor_ir/BUILD  |   5 +
 .../backends/gpu/codegen/tensor_ir/fusion.cc  |  16 +-
 .../gpu/codegen/tensor_ir/hlo_to_tensor_ir.cc |   2 +-
 .../tensor_ir/hlo_to_tensor_ir_test.cc        |   7 +-
 .../gpu/codegen/tensor_ir/layout_utils.cc     |  69 ++++---
 .../gpu/codegen/tensor_ir/layout_utils.h      |  16 +-
 .../codegen/tensor_ir/layout_utils_test.cc    |  91 +++++++++
 .../backends/gpu/codegen/tensor_ir/support.cc |  40 ++--
 .../gpu/codegen/tensor_ir/support_test.cc     |  63 ++++++
 .../tests/fusion_emitter_device_test.cc       |  76 +++++++
 .../tests/fusion_emitter_thunk_test.cc        |  27 +++
 .../codegen/tensor_ir/tests/layout_other.hlo  |  14 ++
 .../codegen/tensor_ir/tests/multi_output.hlo  |  43 ++++
 .../legalize_stablehlo_to_tensor_ir.cc        |  36 +++-
 ...ize_stablehlo_to_tensor_ir_structural.mlir |  14 ++
 16 files changed, 575 insertions(+), 134 deletions(-)
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/multi_output.hlo

diff --git a/third_party/xla/third_party/tensor_ir/patches/tiling_analysis.patch b/third_party/xla/third_party/tensor_ir/patches/tiling_analysis.patch
index 9b00b444d29..2f539d701fe 100644
--- a/third_party/xla/third_party/tensor_ir/patches/tiling_analysis.patch
+++ b/third_party/xla/third_party/tensor_ir/patches/tiling_analysis.patch
@@ -3,7 +3,7 @@ model evaluation) imported from dynamic-kernel-generator/tensor_ir.
 
 diff --git a/include/tensor_ir/Analysis/Tiling/Arch.h b/include/tensor_ir/Analysis/Tiling/Arch.h
 new file mode 100644
-index 0000000..8d4e3e6
+index 0000000..4becf41
 --- /dev/null
 +++ b/include/tensor_ir/Analysis/Tiling/Arch.h
 @@ -0,0 +1,57 @@
@@ -21,8 +21,8 @@ index 0000000..8d4e3e6
 +/// Common parameters for all GPU architectures.
 +/// Number of threads per warp.
 +constexpr int kWarpSize = 32;
-+/// Size of a cache line in bytes.
-+constexpr int kCacheLineSize = 128;
++/// Size of a memory sector in bytes.
++constexpr int kMemorySectorSize = 32;
 +/// Maximum number of registers per thread.
 +constexpr int kMaxRegistersPerThread = 255;
 +/// Maximum number of registers per SM.
@@ -97,7 +97,7 @@ index 0000000..ade5641
 +#endif // TENSOR_IR_ANALYSIS_TILING_ENUMERATE_H_
 diff --git a/include/tensor_ir/Analysis/Tiling/Evaluate.h b/include/tensor_ir/Analysis/Tiling/Evaluate.h
 new file mode 100644
-index 0000000..350dfa8
+index 0000000..403592b
 --- /dev/null
 +++ b/include/tensor_ir/Analysis/Tiling/Evaluate.h
 @@ -0,0 +1,191 @@
@@ -145,7 +145,7 @@ index 0000000..350dfa8
 +/// - The register and shared memory storage, as the peak over the iteration
 +///   spaces. The register storage assumes that every live value of a tile is
 +///   kept in registers.
-+/// - The number of memory accesses and accessed cache lines (loads and stores),
++/// - The number of memory accesses and accessed sectors (loads and stores),
 +///   the number of loop iterations over the reduction tiles, and the number of
 +///   tensor core operations per MMA generation.
 +///
@@ -209,7 +209,7 @@ index 0000000..350dfa8
 +  /// Number of memory accesses (loads/stores).
 +  int64_t memoryAccessCount = 0;
 +
-+  /// Estimate cost of memory accesses (~number of cache lines).
++  /// Estimate cost of memory accesses (~number of sectors).
 +  int64_t memoryAccessCost = 0;
 +
 +  /// Number of loop iterations.
@@ -428,10 +428,10 @@ index 0000000..d11d8cb
 +} // namespace mlir::nv_tensor_ir
 diff --git a/lib/Analysis/Tiling/Enumerate.cpp b/lib/Analysis/Tiling/Enumerate.cpp
 new file mode 100644
-index 0000000..74441bd
+index 0000000..e4bbed9
 --- /dev/null
 +++ b/lib/Analysis/Tiling/Enumerate.cpp
-@@ -0,0 +1,869 @@
+@@ -0,0 +1,910 @@
 +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 +
@@ -482,6 +482,8 @@ index 0000000..74441bd
 +/// - Tensor sources count the tensors where the dimension is contiguous
 +///   (stride 1, the loads and stores can be vectorized), and where it is
 +///   broadcast (stride 0, the loaded values are reused across the tile).
++///   Stores are counted twice, as storing a partial sector needs a
++///   read-modify-write.
 +/// - Concat sources mark their concatenation dimension as non-tileable, as the
 +///   evaluator requires a unit tile size for it.
 +/// - Matmul sources mark the dimensions that map to the M and N dimensions of
@@ -527,6 +529,11 @@ index 0000000..74441bd
 +public:
 +  explicit DimensionLayoutIterator(size_t rank);
 +
++  LogicalResult start(LayoutSourceAttrInterface layout, bool writes) {
++    this->writes = writes;
++    return dispatch(layout);
++  }
++
 +  LogicalResult processTensor(TensorSourceAttr tensor) override;
 +  LogicalResult processConcat(ConcatSourceAttr concat) override;
 +  LogicalResult processReduction(ReductionSourceAttr reduction) override;
@@ -538,6 +545,7 @@ index 0000000..74441bd
 +private:
 +  SmallVector<DimensionInfo> dimensions;
 +  SmallVector<DimensionInfo *> pointers;
++  bool writes = false;
 +};
 +
 +DimensionLayoutIterator::DimensionLayoutIterator(size_t rank) {
@@ -554,7 +562,8 @@ index 0000000..74441bd
 +  for (auto [idx, stride] : llvm::enumerate(cuteLayout.stride().getValues())) {
 +    if (stride.isLeaf() && stride.isStatic()) {
 +      if (stride.as_int() == 1) {
-+        pointers[idx]->vectorizableCount++;
++        // Increased weights for write operations (read-modify-write).
++        pointers[idx]->vectorizableCount += writes ? 2 : 1;
 +      } else if (stride.as_int() == 0) {
 +        pointers[idx]->broadcastedCount++;
 +      }
@@ -806,6 +815,9 @@ index 0000000..74441bd
 +/// Generates the combinations of bits assigned to slots within the budget. If
 +/// the search space is too large, it is first reduced by lowering the limits of
 +/// the least important dimensions, so that the budget covers the rest evenly.
++/// The lowering stops once the search space fits the budget, or once the limits
++/// sum to at most `maxBits`: lowering them further would make even the largest
++/// combination smaller than the upper bound.
 +///
 +/// The generator assigns the bits depth-first, starting from the largest values
 +/// of the first slots. When the search space exceeds the budget, the output is
@@ -818,14 +830,18 @@ index 0000000..74441bd
 +                             ArrayRef<DimensionInfo> dimensions) {
 +  // If the number of combinations is within the budget, then generate all
 +  // combinations (fast path).
-+  if (countCombinations(minBits, maxBits, limits) <= maxResults) {
++  auto fitsBudget = [&](ArrayRef<int> limits) {
++    return llvm::sum_of(limits) <= maxBits ||
++           countCombinations(minBits, maxBits, limits) <= maxResults;
++  };
++  if (fitsBudget(limits)) {
 +    return generateCombinationsWithinRange(minBits, maxBits, limits, maxResults,
 +                                           descending);
 +  }
 +
 +  // Conservative upper bound for the tile size of the vectorizable dimensions,
 +  // assuming small element size (one byte).
-+  int vectorizableCap = llvm::Log2_64(kCacheLineSize);
++  int vectorizableCap = llvm::Log2_64(kMemorySectorSize);
 +
 +  // Try reducing the search space by reducing the limits starting from the
 +  // less important dimensions. Reduce the limit on vectorizable dimensions and
@@ -835,7 +851,7 @@ index 0000000..74441bd
 +    int cap = dimensions[dim].vectorizableCount > 0 ? vectorizableCap : 0;
 +    if (limits[dim] > cap) {
 +      newLimits[dim] = cap;
-+      if (countCombinations(minBits, maxBits, newLimits) <= maxResults) {
++      if (fitsBudget(newLimits)) {
 +        return generateCombinationsWithinRange(minBits, maxBits, newLimits,
 +                                               maxResults, descending);
 +      }
@@ -843,15 +859,16 @@ index 0000000..74441bd
 +  }
 +
 +  // The search space is still too large, so lower the cap on all slots by up
-+  // to kReductionFactor bits (halving the tile size each step). As above, stop
-+  // at the first cap that fits the budget. If no cap fits, the lowest one is
-+  // used and the output is truncated.
++  // to kReductionExponentFactorMax bits (halving the tile size each step).
++  // As above, stop at the first cap that fits the budget or brings the limits
++  // within `maxBits`. Otherwise, the lowest cap is used and the output is
++  // truncated.
 +  constexpr int kReductionExponentFactorMax = 2;
 +  for (int i = 1; i <= kReductionExponentFactorMax; ++i) {
 +    for (int &limit : newLimits) {
 +      limit = std::min(limit, vectorizableCap - i);
 +    }
-+    if (countCombinations(minBits, maxBits, newLimits) <= maxResults) {
++    if (fitsBudget(newLimits)) {
 +      break;
 +    }
 +  }
@@ -879,11 +896,12 @@ index 0000000..74441bd
 +/// - The upper bound on the number of tile elements is the number of elements
 +///   per SM if the grid runs in a single wave. Larger tiles would leave SMs
 +///   idle, but the bound is kept above the minimum tile size (one warp).
-+/// - The main tile combined with the reduction tile is limited to 2^15
-+///   elements, so larger reduction tiles leave room for smaller main tiles.
-+/// - The evaluator is used to lower the upper bound while the estimated
-+///   register storage exceeds the register file, as such tiles are likely to
-+///   spill.
++/// - For each reduction tile size, the evaluator is used to lower the
++///   per-dimension limits, so that the estimated register storage fits the
++///   register file, as larger tiles are likely to spill.
++/// - The lower bound is the minimum tile size (one warp). With reductions, it
++///   is lowered further when several blocks fit on one SM, as smaller tiles
++///   still run in a single wave.
 +/// - Non-tileable dimensions are forced to a unit tile size. They are still
 +///   counted in the single wave bound, as they multiply the number of blocks.
 +///
@@ -933,9 +951,8 @@ index 0000000..74441bd
 +          *TilingConfig::get(tileShape, 1 << reductionTileSizeBits));
 +      if (succeeded(evalResult)) {
 +        // Proportionally lower the limit to avoid spilling.
-+        int64_t ratio = llvm::divideCeil(evalResult->registerStorageBytes,
-+                                         kMaxRegistersPerSM * kRegisterSize);
-+        assert(ratio > 0);
++        int64_t ratio = 1 + evalResult->registerStorageBytes /
++                                (kMaxRegistersPerSM * kRegisterSize);
 +        int newLimit = kMaxTileSizeBits - llvm::Log2_64_Ceil(ratio);
 +        result[idx] = std::min(limits[idx], std::max(newLimit, 0));
 +      }
@@ -1013,14 +1030,38 @@ index 0000000..74441bd
 +        llvm::divideCeil(maxResults - result.size(),
 +                         reductionTileSizeBits - minReductionTileSizeBits + 1);
 +
-+    // Update the limits for the tile dimensions.
++    // Lower the minimum tile size (32) down to the smallest tile for which the
++    // grid still runs in a single wave (reductions only).
++    // The register storage of a unit main tile estimates how many blocks can
++    // be resident on one SM (capped by the hardware limit). The single wave
++    // bound assumes one block per SM; if N blocks fit on each SM, the grid can
++    // have N times more blocks, so the tile can be N times smaller and still
++    // run in a single wave.
++    int lowerBound = kMinTileSizeBits;
++    if (reductionTileSizeBits > 0) {
++      auto evalResult = evaluator.evaluate(
++          *TilingConfig::get(SmallVector<int32_t>(dimensions.size(), 1),
++                             1 << reductionTileSizeBits));
++      if (succeeded(evalResult)) {
++        int occupancy = std::clamp<int>((kMaxRegistersPerSM * kRegisterSize) /
++                                            evalResult->registerStorageBytes,
++                                        1, archInfo.maxBlocksPerSM);
++        int bitsRemaining =
++            bitsPerProcessorSingleWave - llvm::Log2_32_Ceil(occupancy);
++        lowerBound = std::min(lowerBound, std::max(bitsRemaining, 0));
++      }
++    }
++
++    // Lower the per-dimension limits so that a tile with this reduction tile
++    // size fits into the register file. The tile cannot be larger than all
++    // dimensions at their limits, nor than the maximum tile size.
 +    SmallVector<int> limits = getMaxSizeBits(dimensions, maxTileSizeBits,
 +                                             reductionTileSizeBits, evaluator);
-+    int localUpperBound = std::min(llvm::sum_of(limits), globalUpperBound);
++    int upperBound = std::min(kMaxTileSizeBits, llvm::sum_of(limits));
 +
-+    // Calculate the bounds for the tile size.
-+    int upperBound = std::min(kMaxTileSizeBits, localUpperBound);
-+    int lowerBound = std::min(kMinTileSizeBits, upperBound);
++    // Set the bounds for the tile size.
++    upperBound = std::min(upperBound, globalUpperBound);
++    lowerBound = std::min(lowerBound, upperBound);
 +
 +    LLVM_DEBUG({
 +      llvm::dbgs() << "Reduction tile size: " << (1 << reductionTileSizeBits)
@@ -1271,9 +1312,9 @@ index 0000000..74441bd
 +
 +  // Collect the dimensions of the main iteration space.
 +  DimensionLayoutIterator iter(shape.size());
-+  auto inputResult = iter.dispatch(evaluator.getInputLayout());
++  auto inputResult = iter.start(evaluator.getInputLayout(), /*writes=*/false);
 +  (void)inputResult;
-+  auto outputResult = iter.dispatch(evaluator.getOutputLayout());
++  auto outputResult = iter.start(evaluator.getOutputLayout(), /*writes=*/true);
 +  (void)outputResult;
 +
 +  // Order the dimensions using the importance heuristic.
@@ -1303,10 +1344,10 @@ index 0000000..74441bd
 +} // namespace mlir::nv_tensor_ir
 diff --git a/lib/Analysis/Tiling/Evaluate.cpp b/lib/Analysis/Tiling/Evaluate.cpp
 new file mode 100644
-index 0000000..365b1d3
+index 0000000..e86311b
 --- /dev/null
 +++ b/lib/Analysis/Tiling/Evaluate.cpp
-@@ -0,0 +1,1388 @@
+@@ -0,0 +1,1399 @@
 +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 +
@@ -1670,10 +1711,8 @@ index 0000000..365b1d3
 +
 +/// Calculate the storage needed for the element type.
 +int getSizeBytes(Type type) {
-+  if (auto tensorType = dyn_cast<TensorType>(type)) {
-+    type = tensorType.getElementType();
-+  }
-+  return std::max(type.getIntOrFloatBitWidth() / 8, 1u);
++  return llvm::divideCeil(
++      mlir::getElementTypeOrSelf(type).getIntOrFloatBitWidth(), 8);
 +}
 +
 +/// Calculate the storage needed for one tile element in the iteration space.
@@ -1996,13 +2035,13 @@ index 0000000..365b1d3
 +/// with their combination.
 +///
 +/// Tensor sources are the memory accesses. Each one counts as one access, and
-+/// the number of cache lines is estimated from the CuTe layout clamped to the
++/// the number of sectors is estimated from the CuTe layout clamped to the
 +/// tile shape:
 +///
 +/// - Broadcast dimensions (stride 0) are skipped.
-+/// - Dimensions with a stride below the cache line size are densely packed.
++/// - Dimensions with a stride below the sector size are densely packed.
 +///   Together they form one contiguous run per combination of the other
-+///   dimensions, which may straddle one extra cache line if it is unaligned.
++///   dimensions, which may straddle one extra sector if it is unaligned.
 +/// - Every step along a larger (or dynamic) stride starts a new run.
 +///
 +/// Tensor sources without a graph input or output (e.g. constants) cost
@@ -2033,6 +2072,11 @@ index 0000000..365b1d3
 +  EvaluationLayoutIterator(TilingEvaluator &parent, TilingConfig config)
 +      : TilingLayoutIterator(std::move(config)), parent(parent) {}
 +
++  LogicalResult start(LayoutSourceAttrInterface layout, bool writes) {
++    this->writes = writes;
++    return dispatch(layout);
++  }
++
 +  LogicalResult processTensor(TensorSourceAttr tensor) override;
 +  LogicalResult processComposite(CompositeSourceAttr composite) override;
 +  LogicalResult processConcat(ConcatSourceAttr concat) override;
@@ -2045,6 +2089,7 @@ index 0000000..365b1d3
 +private:
 +  TilingEvaluator &parent;
 +  SmallVector<TilingEvaluation> evals;
++  bool writes = false;
 +};
 +
 +LogicalResult EvaluationLayoutIterator::processTensor(TensorSourceAttr tensor) {
@@ -2072,16 +2117,15 @@ index 0000000..365b1d3
 +
 +  // Get the element type size (bytes).
 +  int elementSizeBytes = getSizeBytes(elementType);
-+  int elementsPerCacheLine = kCacheLineSize / elementSizeBytes;
++  int elementsPerSector = kMemorySectorSize / elementSizeBytes;
 +
-+  // If the offset is not aligned to the cache line boundary, an additional
-+  // cache line access may be needed.
-+  bool isUnalignedAccess =
-+      (tensor.getOffset() * elementSizeBytes) % kCacheLineSize != 0;
++  // If the offset is not aligned to the sector boundary, an additional sector
++  // access may be needed.
++  bool isUnalignedAccess = tensor.getOffset() % elementsPerSector != 0;
 +
-+  // Calculate the number of cache line accessed for the current tile.
-+  int64_t cacheLineAccesses = 1;
-+  int64_t innerSpan = 0;
++  // Calculate the number of sector accessed for the current tile.
++  int64_t sectorAccesses = 1;
++  int64_t innerSpan = 0, innerSize = 1;
 +  for (auto [tileSize, dimLayout] :
 +       llvm::zip_equal(config.tileShape, cuteLayout)) {
 +    auto cgShape = dimLayout.shape();
@@ -2093,45 +2137,53 @@ index 0000000..365b1d3
 +    int64_t actualSize = std::min<int64_t>(tileSize, maxSize);
 +
 +    // Assume large stride value for dynamic strides.
-+    int64_t cacheLinesInDim = actualSize;
 +    if (tcutegen::is_static(cgStride)) {
 +      // Skip broadcasted dimensions.
 +      int64_t stride = cgStride.as_int();
 +      if (stride == 0) {
 +        continue;
 +      }
-+      if (stride < elementsPerCacheLine) {
-+        // If the accesses are not aligned to the cache line boundary, an
-+        // additional cache line access may be needed.
++      if (stride < elementsPerSector) {
++        // If the accesses are not aligned to the sector boundary, an additional
++        // sector access may be needed.
 +        int64_t accessSize = actualSize * stride;
-+        isUnalignedAccess |= accessSize % elementsPerCacheLine != 0 &&
-+                             elementsPerCacheLine % accessSize != 0;
++        isUnalignedAccess |= accessSize % elementsPerSector != 0 &&
++                             elementsPerSector % accessSize != 0;
 +        // Extend the contiguous span by this dimension's extent.
 +        innerSpan += accessSize - stride;
-+        cacheLinesInDim = 1;
-+      } else {
-+        // Each step along a large-stride dimension starts a new contiguous run.
-+        // Unless the stride is a multiple of the cache line size, these runs
-+        // start mid-line and may straddle an extra cache line.
-+        isUnalignedAccess |= stride % elementsPerCacheLine != 0;
++        innerSize *= actualSize;
++        continue;
 +      }
++      // Each step along a large-stride dimension starts a new contiguous run.
++      // Unless the stride is a multiple of the sector size, these runs start
++      // mid-sector and may straddle an extra sector.
++      isUnalignedAccess |= stride % elementsPerSector != 0;
 +    }
-+    cacheLineAccesses *= cacheLinesInDim;
++    sectorAccesses *= actualSize;
 +  }
 +
-+  // Densely packed dimensions (stride below a cache line size) form a single
++  // Densely packed dimensions (stride below a sector size) form a single
 +  // contiguous run per combination of the remaining dimensions. The run covers
-+  // `innerSpan` elements and may straddle one extra cache line when the access
++  // `innerSpan` elements and may straddle one extra sector when the access
 +  // is unaligned.
 +  if (innerSpan > 0) {
-+    cacheLineAccesses *= llvm::divideCeil(innerSpan + 1, elementsPerCacheLine) +
-+                         isUnalignedAccess;
++    sectorAccesses *=
++        llvm::divideCeil(innerSpan + 1, elementsPerSector) + isUnalignedAccess;
++
++    // Stores that only partially cover a sector need a read-modify-write: the
++    // sector is read before the written elements are merged into it. This is
++    // the case when the run is not a whole number of sectors, or when it has
++    // gaps (the span is larger than the number of written elements).
++    if (writes &&
++        (innerSize % elementsPerSector != 0 || innerSpan + 1 != innerSize)) {
++      sectorAccesses *= 2;
++    }
 +  }
 +
 +  // Add the evaluation result.
 +  evals.push_back(TilingEvaluation{
 +      .memoryAccessCount = 1,
-+      .memoryAccessCost = cacheLineAccesses,
++      .memoryAccessCost = sectorAccesses,
 +  });
 +  return success();
 +}
@@ -2456,7 +2508,7 @@ index 0000000..365b1d3
 +}
 +
 +bool TilingEvaluation::operator<(const TilingEvaluation &other) const {
-+  // Total number of memory cache lines accessed is the main discriminator.
++  // Total number of memory sectors accessed is the main discriminator.
 +  int64_t score = blockCount * memoryAccessCost;
 +  int64_t otherScore = other.blockCount * other.memoryAccessCost;
 +  if (score != otherScore) {
@@ -2608,13 +2660,13 @@ index 0000000..365b1d3
 +
 +  // Get the evaluation result for the input layout.
 +  EvaluationLayoutIterator iterator(*this, config);
-+  if (failed(iterator.dispatch(inputLayout))) {
++  if (failed(iterator.start(inputLayout, /*writes=*/false))) {
 +    return failure();
 +  }
 +  TilingEvaluation result = iterator.getResult();
 +
 +  // Get the evaluation result for the output layout.
-+  if (failed(iterator.dispatch(outputLayout))) {
++  if (failed(iterator.start(outputLayout, /*writes=*/true))) {
 +    return failure();
 +  }
 +  TilingEvaluation output = iterator.getResult();
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD
index b6c6daed396..c19c9016151 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD
@@ -22,6 +22,7 @@ cc_library(
     hdrs = ["support.h"],
     deps = [
         "//xla:comparison_util",
+        "//xla:shape_util",
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/service:decision",
@@ -36,9 +37,12 @@ xla_cc_test(
     srcs = ["support_test.cc"],
     deps = [
         ":support",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
         "//xla/stream_executor:device_description",
         "//xla/stream_executor/cuda:cuda_compute_capability",
         "//xla/stream_executor/rocm:rocm_compute_capability",
+        "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest_main",
     ],
 )
@@ -157,6 +161,7 @@ cc_library(
         "//xla/hlo/ir:hlo",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
         "@llvm-project//mlir:IR",
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/fusion.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/fusion.cc
index 9316072c212..9fb103ac655 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/fusion.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/fusion.cc
@@ -106,15 +106,17 @@ absl::StatusOr<KernelReuseCache::Entry> CompileFusion(
 
   const std::vector<emitters::KernelArgument>& kernel_args =
       kernel_arguments.args();
-  // The graph has one argument per fused parameter, which the HLO verifier
-  // ties to the fusion's operands. A fusion is permitted to have operands that
-  // no fused parameter reads, though, in which case the counts disagree; fail
-  // loudly here rather than letting the alignments be attached off-by-one.
+  // The graph has one argument per fused parameter and one result per fusion
+  // output, which the HLO verifier ties to the fusion's operands and result. A
+  // fusion is permitted to have operands that no fused parameter reads, though,
+  // in which case the counts disagree; fail loudly here rather than letting the
+  // alignments be attached off-by-one.
   int64_t num_inputs = computation->num_parameters();
-  if (kernel_args.size() != num_inputs + 1) {
+  int64_t num_outputs = graph_op.getNumResults();
+  if (kernel_args.size() != num_inputs + num_outputs) {
     return absl::InternalError(absl::StrCat(
-        "TensorIrFusion: expected ", num_inputs + 1,
-        " kernel arguments (one per fused parameter plus a single output), "
+        "TensorIrFusion: expected ", num_inputs + num_outputs,
+        " kernel arguments (one per fused parameter plus one per output), "
         "got ",
         kernel_args.size(), " for fusion: ", fusion.ToString()));
   }
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir.cc
index ef544713f62..47f64276e4d 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir.cc
@@ -69,7 +69,7 @@ absl::StatusOr<mlir::nv_tensor_ir::GraphOp> ImportAndLegalizeComputation(
                                         &function_map, &builder,
                                         /*is_main=*/true,
                                         /*flatten_computation_args_result=*/
-                                        false);
+                                        true);
   if (!func_op.ok()) {
     return diagnostic_handler.Combine(absl::Status(
         func_op.status().code(),
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir_test.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir_test.cc
index 0793fb7198b..02acac00dd0 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/hlo_to_tensor_ir_test.cc
@@ -207,7 +207,7 @@ TEST_F(HloToTensorIrTest, UnsupportedOpFailsAndNamesTheOp) {
   EXPECT_THAT(graph.status().message(), HasSubstr("stablehlo.xor"));
 }
 
-TEST_F(HloToTensorIrTest, TupleRootFailsCleanly) {
+TEST_F(HloToTensorIrTest, TupleRootBecomesMultipleResults) {
   constexpr absl::string_view kHloText = R"(
     HloModule test_module
 
@@ -224,8 +224,9 @@ TEST_F(HloToTensorIrTest, TupleRootFailsCleanly) {
 
   auto graph =
       ImportAndLegalizeComputation(*hlo_module->entry_computation(), *module_);
-  EXPECT_THAT(graph, Not(IsOk()));
-  EXPECT_THAT(graph.status().message(), HasSubstr("stablehlo.tuple"));
+  ASSERT_THAT(graph, IsOk());
+  EXPECT_EQ(graph->getNumResults(), 2);
+  EXPECT_EQ(graph->getNumArguments(), 2);
 }
 
 TEST_F(HloToTensorIrTest, NullModuleIsRejected) {
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils.cc
index 714ffac1a60..a59d7bdec0d 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils.cc
@@ -22,9 +22,9 @@ limitations under the License.
 #include "tensor_ir/Dialect/TensorIR.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
-#include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -66,24 +66,29 @@ bool HasDefaultLayout(const Shape& shape) {
 
 namespace {
 
-// Both attachment functions below target the single-result graphs that
-// `ImportAndLegalizeComputation` produces.
-absl::Status CheckGraphIsSingleResult(mlir::nv_tensor_ir::GraphOp graph,
-                                      absl::string_view caller) {
-  if (graph.getNumResults() != 1) {
-    return absl::InternalError(
-        absl::StrCat(caller, ": graph '", graph.getSymName().str(), "' has ",
-                     graph.getNumResults(), " results, expected exactly 1"));
+// Returns the shapes of the graph results `ImportAndLegalizeComputation`
+// produces for `computation`: the elements of a root tuple, or the root shape
+// itself otherwise.
+absl::StatusOr<absl::Span<const Shape>> GetResultShapes(
+    const HloComputation& computation) {
+  const Shape& root_shape = computation.root_instruction()->shape();
+  absl::Span<const Shape> shapes =
+      root_shape.IsTuple() ? absl::MakeConstSpan(root_shape.tuple_shapes())
+                           : absl::MakeConstSpan(&root_shape, 1);
+  for (const Shape& shape : shapes) {
+    if (!shape.IsArray()) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("AttachLayoutStrides: non-array result of computation '",
+                       computation.name(), "': ", shape.ToString()));
+    }
   }
-  return absl::OkStatus();
+  return shapes;
 }
 
 }  // namespace
 
 absl::Status AttachLayoutStrides(mlir::nv_tensor_ir::GraphOp graph,
                                  const HloComputation& computation) {
-  ABSL_RETURN_IF_ERROR(CheckGraphIsSingleResult(graph, "AttachLayoutStrides"));
-
   int64_t num_parameters = computation.num_parameters();
   if (graph.getNumArguments() != num_parameters) {
     return absl::InternalError(absl::StrCat(
@@ -110,16 +115,20 @@ absl::Status AttachLayoutStrides(mlir::nv_tensor_ir::GraphOp graph,
     }
   }
 
-  const Shape& root_shape = computation.root_instruction()->shape();
-  if (!root_shape.IsArray()) {
-    return absl::InvalidArgumentError(
-        absl::StrCat("AttachLayoutStrides: non-array root of computation '",
-                     computation.name(), "': ", root_shape.ToString()));
+  ABSL_ASSIGN_OR_RETURN(absl::Span<const Shape> result_shapes,
+                   GetResultShapes(computation));
+  if (graph.getNumResults() != result_shapes.size()) {
+    return absl::InternalError(absl::StrCat(
+        "AttachLayoutStrides: graph '", graph.getSymName().str(), "' has ",
+        graph.getNumResults(), " results, expected ", result_shapes.size(),
+        " for computation '", computation.name(), "'"));
   }
-  if (!HasDefaultLayout(root_shape)) {
-    graph.setResultAttr(
-        0, stride_attr_name,
-        mlir::StringAttr::get(context, ComputeStrideString(root_shape)));
+  for (int64_t i = 0; i < result_shapes.size(); ++i) {
+    if (!HasDefaultLayout(result_shapes[i])) {
+      graph.setResultAttr(i, stride_attr_name,
+                          mlir::StringAttr::get(
+                              context, ComputeStrideString(result_shapes[i])));
+    }
   }
 
   return absl::OkStatus();
@@ -128,14 +137,14 @@ absl::Status AttachLayoutStrides(mlir::nv_tensor_ir::GraphOp graph,
 absl::Status AttachBufferAlignments(
     mlir::nv_tensor_ir::GraphOp graph,
     absl::Span<const emitters::KernelArgument> kernel_args) {
-  ABSL_RETURN_IF_ERROR(CheckGraphIsSingleResult(graph, "AttachBufferAlignments"));
-
   int64_t num_inputs = graph.getNumArguments();
-  if (kernel_args.size() != num_inputs + 1) {
+  int64_t num_results = graph.getNumResults();
+  if (kernel_args.size() != num_inputs + num_results) {
     return absl::InternalError(absl::StrCat(
-        "AttachBufferAlignments: expected ", num_inputs + 1,
+        "AttachBufferAlignments: expected ", num_inputs + num_results,
         " kernel arguments for graph '", graph.getSymName().str(), "' with ",
-        num_inputs, " arguments, but got ", kernel_args.size()));
+        num_inputs, " arguments and ", num_results, " results, but got ",
+        kernel_args.size()));
   }
 
   mlir::MLIRContext* context = graph.getContext();
@@ -148,9 +157,11 @@ absl::Status AttachBufferAlignments(
         i, alignment_attr_name,
         mlir::IntegerAttr::get(i64_type, kernel_args[i].alignment()));
   }
-  graph.setResultAttr(
-      0, alignment_attr_name,
-      mlir::IntegerAttr::get(i64_type, kernel_args[num_inputs].alignment()));
+  for (int64_t i = 0; i < num_results; ++i) {
+    graph.setResultAttr(i, alignment_attr_name,
+                        mlir::IntegerAttr::get(
+                            i64_type, kernel_args[num_inputs + i].alignment()));
+  }
 
   return absl::OkStatus();
 }
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils.h b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils.h
index 25652cff8e8..9c5b5b7cd37 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils.h
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils.h
@@ -38,8 +38,9 @@ std::string ComputeStrideString(const Shape& shape);
 // which case no stride attribute needs to be attached.
 bool HasDefaultLayout(const Shape& shape);
 
-// Attaches `nv_tensor_ir.stride` to every argument and to result 0 of `graph`
-// whose corresponding shape does not have the default layout.
+// Attaches `nv_tensor_ir.stride` to every argument and result of `graph` whose
+// corresponding shape does not have the default layout. Results correspond to
+// the elements of the root tuple, or to the root itself if it is not a tuple.
 //
 // Strides are intrinsic to `computation`: they are derived purely from the
 // layouts of its parameters and its root, with no buffer assignment involved.
@@ -47,21 +48,20 @@ bool HasDefaultLayout(const Shape& shape);
 // `ImportAndLegalizeComputation`, so that everyone who imports a computation
 // (the emitter, the autotuner, tools) sees the same graph.
 //
-// `graph` must have exactly one argument per computation parameter and exactly
-// one result; this mirrors what the importer produces.
+// `graph` must have exactly one argument per computation parameter and one
+// result per computation result; this mirrors what the importer produces.
 absl::Status AttachLayoutStrides(mlir::nv_tensor_ir::GraphOp graph,
                                  const HloComputation& computation);
 
-// Attaches `nv_tensor_ir.alignment` to every argument and to result 0 of
-// `graph`.
+// Attaches `nv_tensor_ir.alignment` to every argument and result of `graph`.
 //
 // Alignments come from buffer assignment, so unlike strides they are a
 // property of the buffers the computation happens to be assigned rather than
 // of the computation itself. This can only run in the emitter, after kernel
 // arguments exist.
 //
-// `kernel_args` must hold one argument per graph argument plus one for the
-// single result.
+// `kernel_args` must hold one argument per graph argument followed by one per
+// graph result.
 absl::Status AttachBufferAlignments(
     mlir::nv_tensor_ir::GraphOp graph,
     absl::Span<const emitters::KernelArgument> kernel_args);
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils_test.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils_test.cc
index 3eaa0c8301b..e0d95f049d3 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/layout_utils_test.cc
@@ -296,6 +296,65 @@ TEST_F(AttachLayoutStridesTest, ErrorPathTupleParameter) {
   graph.erase();
 }
 
+TEST_F(AttachLayoutStridesTest, MultipleResultsFromRootTuple) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule test_module
+
+    ENTRY entry {
+      p0 = f32[4,16]{1,0} parameter(0)
+      c0 = f32[4,16]{0,1} copy(p0)
+      c1 = f32[4,16]{1,0} copy(p0)
+      ROOT t = (f32[4,16]{0,1}, f32[4,16]{1,0}) tuple(c0, c1)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto hlo_module, ParseAndReturnVerifiedModule(kHloText));
+
+  mlir::MLIRContext context;
+  context.loadDialect<mlir::nv_tensor_ir::TensorIRDialect>();
+  mlir::OpBuilder b(&context);
+
+  auto t = mlir::RankedTensorType::get({4, 16}, b.getF32Type());
+  auto graph = MakeGraph(b, {t}, {t, t});
+
+  ASSERT_THAT(AttachLayoutStrides(graph, *hlo_module->entry_computation()),
+              IsOk());
+
+  auto stride_attr_name =
+      mlir::nv_tensor_ir::TensorIRDialect::getStrideAttrName();
+  auto res0 = graph.getResultAttrOfType<mlir::StringAttr>(0, stride_attr_name);
+  ASSERT_TRUE(res0 != nullptr);
+  EXPECT_EQ(res0.getValue().str(), "(1,4)");
+  EXPECT_TRUE(graph.getResultAttrOfType<mlir::StringAttr>(
+                  1, stride_attr_name) == nullptr);
+
+  graph.erase();
+}
+
+TEST_F(AttachLayoutStridesTest, ErrorPathResultCountMismatch) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule test_module
+
+    ENTRY entry {
+      p0 = f32[4,16]{1,0} parameter(0)
+      ROOT id = f32[4,16]{1,0} copy(p0)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto hlo_module, ParseAndReturnVerifiedModule(kHloText));
+
+  mlir::MLIRContext context;
+  context.loadDialect<mlir::nv_tensor_ir::TensorIRDialect>();
+  mlir::OpBuilder b(&context);
+
+  // Graph has 2 results, but the computation has 1.
+  auto t = mlir::RankedTensorType::get({4, 16}, b.getF32Type());
+  auto graph = MakeGraph(b, {t}, {t, t});
+
+  EXPECT_THAT(AttachLayoutStrides(graph, *hlo_module->entry_computation()),
+              StatusIs(absl::StatusCode::kInternal));
+
+  graph.erase();
+}
+
 class AttachBufferAlignmentsTest : public HloHardwareIndependentTestBase {};
 
 TEST_F(AttachBufferAlignmentsTest, AttachesAlignmentToArgsAndResult) {
@@ -332,6 +391,38 @@ TEST_F(AttachBufferAlignmentsTest, AttachesAlignmentToArgsAndResult) {
   graph.erase();
 }
 
+TEST_F(AttachBufferAlignmentsTest, AttachesAlignmentToMultipleResults) {
+  mlir::MLIRContext context;
+  context.loadDialect<mlir::nv_tensor_ir::TensorIRDialect>();
+  mlir::OpBuilder b(&context);
+
+  auto t = mlir::RankedTensorType::get({4, 16}, b.getF32Type());
+  auto graph = MakeGraph(b, {t}, {t, t});
+
+  Shape shape = ShapeUtil::MakeShapeWithDenseLayout(F32, {4, 16}, {1, 0});
+  std::vector<emitters::KernelArgument> kernel_args;
+  kernel_args.reserve(3);
+  for (int64_t alignment : {128, 64, 256}) {
+    kernel_args.emplace_back(shape);
+    kernel_args.back().set_alignment(alignment);
+  }
+
+  ASSERT_THAT(AttachBufferAlignments(graph, kernel_args), IsOk());
+
+  auto alignment_attr_name =
+      mlir::nv_tensor_ir::TensorIRDialect::getAlignmentAttrName();
+  auto r0 =
+      graph.getResultAttrOfType<mlir::IntegerAttr>(0, alignment_attr_name);
+  ASSERT_TRUE(r0 != nullptr);
+  EXPECT_EQ(r0.getInt(), 64);
+  auto r1 =
+      graph.getResultAttrOfType<mlir::IntegerAttr>(1, alignment_attr_name);
+  ASSERT_TRUE(r1 != nullptr);
+  EXPECT_EQ(r1.getInt(), 256);
+
+  graph.erase();
+}
+
 TEST_F(AttachBufferAlignmentsTest, ErrorPathKernelArgsCountMismatch) {
   mlir::MLIRContext context;
   context.loadDialect<mlir::nv_tensor_ir::TensorIRDialect>();
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.cc
index be878071930..f1753f0fb32 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support.cc
@@ -23,6 +23,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/primitive_util.h"
+#include "xla/shape_util.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/xla_data.pb.h"
@@ -44,7 +45,9 @@ bool IsSupportedPrimitiveType(PrimitiveType type) {
     case PrimitiveType::F16:
     case PrimitiveType::BF16:
     case PrimitiveType::F32:
-    case PrimitiveType::F64: {
+    case PrimitiveType::F64:
+    case PrimitiveType::F8E4M3FN:
+    case PrimitiveType::F8E5M2: {
       return true;
     }
 
@@ -108,6 +111,7 @@ bool IsSupportedFusionOpcode(HloOpcode opcode) {
     case HloOpcode::kParameter:
     case HloOpcode::kConstant:
     case HloOpcode::kIota:
+    case HloOpcode::kCopy:
     case HloOpcode::kConcatenate: {
       return true;
     }
@@ -200,20 +204,26 @@ CodegenDecision IsInstructionSupportedForFusion(const HloInstruction& instr) {
         *instr.fused_instructions_computation());
   }
 
-  if (!instr.shape().IsArray()) {
-    return CodegenDecision::Forbid(absl::StrCat("Unsupported non-array shape: ",
-                                                instr.shape().ToString()));
-  }
-
-  if (!IsSupportedPrimitiveType(instr.shape().element_type())) {
-    return CodegenDecision::Forbid(
-        absl::StrCat("Unsupported element type: ",
-                     primitive_util::LowercasePrimitiveTypeName(
-                         instr.shape().element_type())));
-  }
-  if (!IsSupportedFusionOpcode(instr.opcode())) {
-    return CodegenDecision::Forbid(absl::StrCat(
-        "Unsupported instruction: ", HloOpcodeString(instr.opcode())));
+  // The root of a multi-output fusion is a flat tuple of arrays. Only `kTuple`
+  // gets this exemption: the tuple's operands are checked as instructions of
+  // the computation, while any other tuple-shaped opcode is unsupported.
+  bool is_root_tuple = instr.opcode() == HloOpcode::kTuple && instr.IsRoot() &&
+                       !ShapeUtil::IsNestedTuple(instr.shape());
+  if (!is_root_tuple) {
+    if (!instr.shape().IsArray()) {
+      return CodegenDecision::Forbid(absl::StrCat(
+          "Unsupported non-array shape: ", instr.shape().ToString()));
+    }
+    if (!IsSupportedPrimitiveType(instr.shape().element_type())) {
+      return CodegenDecision::Forbid(
+          absl::StrCat("Unsupported element type: ",
+                       primitive_util::LowercasePrimitiveTypeName(
+                           instr.shape().element_type())));
+    }
+    if (!IsSupportedFusionOpcode(instr.opcode())) {
+      return CodegenDecision::Forbid(absl::StrCat(
+          "Unsupported instruction: ", HloOpcodeString(instr.opcode())));
+    }
   }
 
   switch (instr.opcode()) {
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support_test.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support_test.cc
index 2e1f5e99907..8bf4bcffe36 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/support_test.cc
@@ -15,8 +15,13 @@ limitations under the License.
 
 #include "xla/backends/gpu/codegen/tensor_ir/support.h"
 
+#include <memory>
+
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/strings/string_view.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/rocm/rocm_compute_capability.h"
@@ -54,5 +59,63 @@ TEST(IsSupportedComputeCapabilityTest, ForbidsNonCuda) {
   EXPECT_THAT(decision.Explain(), HasSubstr("CUDA"));
 }
 
+class IsSupportedFusionComputationTest : public HloHardwareIndependentTestBase {
+};
+
+TEST_F(IsSupportedFusionComputationTest, AllowsMultiOutputTupleRoot) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule m
+
+    fused_computation {
+      p0 = f32[8] parameter(0)
+      neg = f32[8] negate(p0)
+      abs = f32[8] abs(p0)
+      ROOT tuple = (f32[8], f32[8]) tuple(neg, abs)
+    }
+
+    ENTRY entry {
+      p0 = f32[8] parameter(0)
+      ROOT fusion = (f32[8], f32[8]) fusion(p0), kind=kLoop,
+          calls=fused_computation
+    })";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kHloText));
+  EXPECT_TRUE(IsSupportedFusionComputation(
+                  *module->GetComputationWithName("fused_computation"))
+                  .IsAllowed());
+}
+
+TEST_F(IsSupportedFusionComputationTest, ForbidsTupleShapedNonTupleRoot) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule m
+
+    compare {
+      p0.lhs = f32[] parameter(0)
+      p0.rhs = f32[] parameter(1)
+      p1.lhs = f32[] parameter(2)
+      p1.rhs = f32[] parameter(3)
+      ROOT lt = pred[] compare(p0.lhs, p0.rhs), direction=LT
+    }
+
+    fused_computation {
+      p0 = f32[8] parameter(0)
+      p1 = f32[8] parameter(1)
+      ROOT sort = (f32[8], f32[8]) sort(p0, p1), dimensions={0}, to_apply=compare
+    }
+
+    ENTRY entry {
+      p0 = f32[8] parameter(0)
+      p1 = f32[8] parameter(1)
+      ROOT fusion = (f32[8], f32[8]) fusion(p0, p1), kind=kLoop,
+          calls=fused_computation
+    })";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kHloText));
+  CodegenDecision decision = IsSupportedFusionComputation(
+      *module->GetComputationWithName("fused_computation"));
+  EXPECT_FALSE(decision.IsAllowed());
+  EXPECT_THAT(decision.Explain(), HasSubstr("non-array shape"));
+}
+
 }  // namespace
 }  // namespace xla::gpu::tensor_ir
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/fusion_emitter_device_test.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/fusion_emitter_device_test.cc
index c6eaa3fd034..8d431ce0c15 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/fusion_emitter_device_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/fusion_emitter_device_test.cc
@@ -86,6 +86,82 @@ ENTRY main {
   EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
 }
 
+TEST_F(TensorIrEmitterTest, CopyChangesLayout) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = f32[8,16]{1,0} parameter(0)
+  neg = f32[8,16]{1,0} negate(p0)
+  ROOT copy = f32[8,16]{0,1} copy(neg)
+}
+
+ENTRY main {
+  p0 = f32[8,16]{1,0} parameter(0)
+  ROOT fusion = f32[8,16]{0,1} fusion(p0), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+TEST_F(TensorIrEmitterTest, MultiOutputElementwiseF32) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = f32[8,16] parameter(0)
+  p1 = f32[8,16] parameter(1)
+  add = f32[8,16] add(p0, p1)
+  mul = f32[8,16] multiply(p0, p1)
+  ROOT tuple = (f32[8,16], f32[8,16]) tuple(add, mul)
+}
+
+ENTRY main {
+  p0 = f32[8,16] parameter(0)
+  p1 = f32[8,16] parameter(1)
+  ROOT fusion = (f32[8,16], f32[8,16]) fusion(p0, p1), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+TEST_F(TensorIrEmitterTest, MultiOutputWithDifferentTypes) {
+  constexpr absl::string_view kHloText = R"(
+fused_computation {
+  p0 = f32[8,16] parameter(0)
+  neg = f32[8,16] negate(p0)
+  cvt = s32[8,16] convert(p0)
+  ROOT tuple = (f32[8,16], s32[8,16]) tuple(neg, cvt)
+}
+
+ENTRY main {
+  p0 = f32[8,16] parameter(0)
+  ROOT fusion = (f32[8,16], s32[8,16]) fusion(p0), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, kExactMatch));
+}
+
+TEST_F(TensorIrEmitterTest, MultiOutputElementwiseAndReduceF32) {
+  constexpr absl::string_view kHloText = R"(
+add_f32 {
+  a = f32[] parameter(0)
+  b = f32[] parameter(1)
+  ROOT sum = f32[] add(a, b)
+}
+
+fused_computation {
+  p0 = f32[8,16] parameter(0)
+  neg = f32[8,16] negate(p0)
+  zero = f32[] constant(0)
+  sum = f32[8] reduce(p0, zero), dimensions={1}, to_apply=add_f32
+  ROOT tuple = (f32[8,16], f32[8]) tuple(neg, sum)
+}
+
+ENTRY main {
+  p0 = f32[8,16] parameter(0)
+  ROOT fusion = (f32[8,16], f32[8]) fusion(p0), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+  EXPECT_TRUE(RunAndCompareNoHloPasses(kHloText, ErrorSpec{1e-5, 1e-5}));
+}
+
 TEST_F(TensorIrEmitterTest, ElementwiseMulMaxF32) {
   constexpr absl::string_view kHloText = R"(
 fused_computation {
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/fusion_emitter_thunk_test.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/fusion_emitter_thunk_test.cc
index 9b54ee50ee5..612f3b49380 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/fusion_emitter_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/fusion_emitter_thunk_test.cc
@@ -54,6 +54,22 @@ ENTRY main {
     backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
 })";
 
+constexpr absl::string_view kTwoOutputsF32 = R"(
+fused_computation {
+  p0 = f32[8,16] parameter(0)
+  p1 = f32[8,16] parameter(1)
+  add = f32[8,16] add(p0, p1)
+  mul = f32[8,16] multiply(p0, p1)
+  ROOT tuple = (f32[8,16], f32[8,16]) tuple(add, mul)
+}
+
+ENTRY main {
+  p0 = f32[8,16] parameter(0)
+  p1 = f32[8,16] parameter(1)
+  ROOT fusion = (f32[8,16], f32[8,16]) fusion(p0, p1), kind=kCustom, calls=fused_computation,
+    backend_config={"fusion_backend_config":{"kind":"__tensorir","tensor_ir_fusion_config":{}}}
+})";
+
 // Checks the shape of what the TensorIR emitter hands to the runtime. Unlike
 // the numerics tests next door these only compile and never launch, so they
 // need a GPU to be visible but not one whose driver can load the kernel.
@@ -130,6 +146,17 @@ TEST_F(TensorIrThunkTest, PassesInputsThenOutput) {
   EXPECT_THAT(thunk->written(), ElementsAre(false, false, true));
 }
 
+// A fusion with several outputs gets one pointer per output after the inputs.
+TEST_F(TensorIrThunkTest, PassesInputsThenAllOutputs) {
+  std::unique_ptr<Executable> executable;
+  ASSERT_OK_AND_ASSIGN(
+      const CustomKernelThunk* thunk,
+      CompileToSingleCustomKernelThunk(kTwoOutputsF32, &executable));
+
+  EXPECT_EQ(thunk->arguments().size(), 4);
+  EXPECT_THAT(thunk->written(), ElementsAre(false, false, true, true));
+}
+
 // The reason for using `CustomKernelThunk` rather than a bespoke thunk: the
 // kernel, its device code and its launch dimensions all round-trip through a
 // proto, which ahead-of-time compilation requires. A bespoke thunk would have
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_other.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_other.hlo
index 8dcaaf3e344..afda84d9ff0 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_other.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/layout_other.hlo
@@ -21,3 +21,17 @@ test_transpose {
   %p0 = f32[2,4,8,16] parameter(0)
   %out = f32[8,2,4,16] transpose(%p0), dimensions={2,0,1,3}
 }
+
+// -----
+
+// CHECK-LABEL: @test_copy_changes_layout
+// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>
+// CHECK-SAME: -> (tensor<8x16xf32> {nv_tensor_ir.stride = "(1,8)"})
+// CHECK: %[[NEG:.*]] = neg %[[ARG0]]
+// CHECK: results %[[NEG]]
+
+test_copy_changes_layout {
+  %p0 = f32[8,16]{1,0} parameter(0)
+  %neg = f32[8,16]{1,0} negate(%p0)
+  ROOT %out = f32[8,16]{0,1} copy(%neg)
+}
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/multi_output.hlo b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/multi_output.hlo
new file mode 100644
index 00000000000..272f225ce3c
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/multi_output.hlo
@@ -0,0 +1,43 @@
+// RUN: fusion_to_tensorir --hlo-to-tensorir --split-input-file %s | FileCheck %s
+
+// CHECK-LABEL: @test_two_outputs
+// CHECK-SAME: %[[ARG0:.*]]: tensor<8x16xf32>, %[[ARG1:.*]]: tensor<8x16xf32>
+// CHECK-SAME: -> (tensor<8x16xf32>, tensor<8x16xf32>)
+// CHECK: %[[ADD:.*]] = add %[[ARG0]], %[[ARG1]]
+// CHECK: %[[MUL:.*]] = mul %[[ARG0]], %[[ARG1]]
+// CHECK: results %[[ADD]], %[[MUL]]
+
+test_two_outputs {
+  %p0 = f32[8,16] parameter(0)
+  %p1 = f32[8,16] parameter(1)
+  %add = f32[8,16] add(%p0, %p1)
+  %mul = f32[8,16] multiply(%p0, %p1)
+  ROOT %out = (f32[8,16], f32[8,16]) tuple(%add, %mul)
+}
+
+// -----
+
+// CHECK-LABEL: @test_different_types
+// CHECK-SAME: -> (tensor<8x16xf32>, tensor<8x16xsi32>)
+// CHECK: results
+
+test_different_types {
+  %p0 = f32[8,16] parameter(0)
+  %neg = f32[8,16] negate(%p0)
+  %cvt = s32[8,16] convert(%p0)
+  ROOT %out = (f32[8,16], s32[8,16]) tuple(%neg, %cvt)
+}
+
+// -----
+
+// CHECK-LABEL: @test_non_default_layout
+// CHECK-SAME: nv_tensor_ir.stride = "(1,8)"
+// CHECK-SAME: nv_tensor_ir.stride = "(1,8)"
+// CHECK: results
+
+test_non_default_layout {
+  %p0 = f32[8,16]{0,1} parameter(0)
+  %neg = f32[8,16]{0,1} negate(%p0)
+  %abs = f32[8,16]{0,1} abs(%p0)
+  ROOT %out = (f32[8,16]{0,1}, f32[8,16]{0,1}) tuple(%neg, %abs)
+}
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/legalize_stablehlo_to_tensor_ir.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/legalize_stablehlo_to_tensor_ir.cc
index 6052d8b3eed..be9f47c9e40 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/legalize_stablehlo_to_tensor_ir.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/legalize_stablehlo_to_tensor_ir.cc
@@ -716,6 +716,38 @@ struct BitcastOpConversion
   }
 };
 
+// Legalizes mhlo.copy to its operand.
+//
+// nv_tensor_ir values are logical tensors, and the physical layout only exists
+// on the graph boundary, where `AttachLayoutStrides` records it as
+// `nv_tensor_ir.stride` attributes derived from the HLO shapes. A copy only
+// changes the layout (the `result_layout` attribute HloFunctionImporter
+// attaches), never an element, so as far as the graph is concerned it is the
+// identity: a copy to a different layout is expressed by the strides of the
+// value it feeds.
+//
+// A `cross_program_prefetch_index` on the copy is a hint for the runtime and
+// does not affect the values, so it is ignored.
+struct CopyOpConversion : public mlir::OpConversionPattern<mlir::mhlo::CopyOp> {
+  using mlir::OpConversionPattern<mlir::mhlo::CopyOp>::OpConversionPattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::mhlo::CopyOp op, OpAdaptor adaptor,
+      mlir::ConversionPatternRewriter& rewriter) const override {
+    auto result_type = mlir::dyn_cast_or_null<mlir::RankedTensorType>(
+        this->getTypeConverter()->convertType(op.getType()));
+    if (!result_type) {
+      return rewriter.notifyMatchFailure(op, "failed to convert result type");
+    }
+    if (adaptor.getOperand().getType() != result_type) {
+      return rewriter.notifyMatchFailure(
+          op, "expected the operand and result types to match");
+    }
+    rewriter.replaceOp(op, adaptor.getOperand());
+    return mlir::success();
+  }
+};
+
 //===----------------------------------------------------------------------===//
 // Reductions.
 //===----------------------------------------------------------------------===//
@@ -1580,8 +1612,8 @@ void populateLegalizeStablehloToTensorIrPatterns(
   patterns.add<ReshapeOpConversion, TransposeOpConversion, SliceOpConversion,
                ConcatenateOpConversion, IotaOpConversion, ConstantOpConversion,
                BroadcastInDimOpConversion, BitcastOpConversion,
-               DotGeneralOpConversion, DotOpConversion>(type_converter,
-                                                        context);
+               CopyOpConversion, DotGeneralOpConversion, DotOpConversion>(
+      type_converter, context);
 
   // Reductions, registered with a higher benefit than the elementwise
   // patterns so that the combiner in the reduction body can be inspected.
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_structural.mlir b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_structural.mlir
index 02b3ffcac9c..442b2651e49 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_structural.mlir
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/transforms/tests/legalize_stablehlo_to_tensor_ir_structural.mlir
@@ -221,3 +221,17 @@ func.func @broadcast_permuted(%arg0: tensor<2x3xf32>) -> tensor<4x3x2xf32> {
   %0 = stablehlo.broadcast_in_dim %arg0, dims = [2, 1] : (tensor<2x3xf32>) -> tensor<4x3x2xf32>
   return %0 : tensor<4x3x2xf32>
 }
+
+// -----
+
+// mhlo.copy only changes the layout, which lives on the graph boundary rather
+// than in the body, so it legalizes to its operand.
+// CHECK-LABEL: nv_tensor_ir.graph @copy_is_identity
+// CHECK-SAME: (%[[IN:.*]]: tensor<8x16xf32>) -> tensor<8x16xf32>
+// CHECK-NEXT: %[[NEG:.*]] = neg %[[IN]] : tensor<8x16xf32>
+// CHECK-NEXT: results %[[NEG]] : tensor<8x16xf32>
+func.func @copy_is_identity(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
+  %0 = stablehlo.negate %arg0 : tensor<8x16xf32>
+  %1 = "mhlo.copy"(%0) {result_layout = dense<[0, 1]> : tensor<2xindex>, xla_shape = "f32[8,16]{0,1}"} : (tensor<8x16xf32>) -> tensor<8x16xf32>
+  return %1 : tensor<8x16xf32>
+}

From 50c8c58e595b88a7c29924ac1606db8bd282790f Mon Sep 17 00:00:00 2001
From: Alexander Lyashuk <crem@google.com>
Date: Wed, 14 Oct 2026 03:15:43 -0700
Subject: [PATCH 0755/1000] [XLA:GPU] Avoid combining reshapes by default in
 ShapeTracker::GetSteps.

`ShapeTracker::GetSteps()` previously returned the raw transformation sequence from internal projections, emitting `reshape -> transpose -> reshape` whenever adjacent transposed dimensions were merged in the `BufferView`. Meanwhile, `ToInstructionChain()` and `DebugString()` separately invoked `OptimizeSteps()` on the result of `GetSteps()` to expand and drop redundant reshapes.

Apply `OptimizeSteps()` directly inside `ShapeTracker::GetSteps()` when `avoid_combining_reshapes` is true (the default), and forward the parameter from `ToInstructionChain()` and `DebugString()`.

PiperOrigin-RevId: 999250054
---
 .../xla/xla/hlo/analysis/shape_tracker.cc     | 20 ++----
 .../xla/xla/hlo/analysis/shape_tracker.h      |  2 +-
 .../xla/hlo/analysis/shape_tracker_test.cc    | 66 +++++++++++--------
 3 files changed, 45 insertions(+), 43 deletions(-)

diff --git a/third_party/xla/xla/hlo/analysis/shape_tracker.cc b/third_party/xla/xla/hlo/analysis/shape_tracker.cc
index e183e2d7a7e..f1380ea33b9 100644
--- a/third_party/xla/xla/hlo/analysis/shape_tracker.cc
+++ b/third_party/xla/xla/hlo/analysis/shape_tracker.cc
@@ -783,7 +783,8 @@ absl::Status ShapeTracker::PrependBitcast(const xla::Shape& src_shape,
   return absl::OkStatus();
 }
 
-std::vector<ShapeTracker::Step> ShapeTracker::GetSteps() const {
+std::vector<ShapeTracker::Step> ShapeTracker::GetSteps(
+    bool avoid_combining_reshapes) const {
   if (ShapeUtil::ElementsIn(input_shape_) == 1) {
     if (input_shape_.dimensions() == output_shape_.dimensions()) {
       return {};
@@ -837,6 +838,9 @@ std::vector<ShapeTracker::Step> ShapeTracker::GetSteps() const {
                                           output_shape_.dimensions().end())});
   }
 
+  if (avoid_combining_reshapes) {
+    return OptimizeSteps(steps, input_shape_, output_shape_);
+  }
   return steps;
 }
 
@@ -981,12 +985,7 @@ absl::StatusOr<HloInstruction*> ShapeTracker::ToInstructionChain(
   HloComputation* computation = inst->parent();
   HloInstruction* current_inst = inst;
 
-  std::vector<Step> steps = GetSteps();
-  if (avoid_combining_reshapes) {
-    steps = OptimizeSteps(steps, input_shape_, output_shape_);
-  }
-
-  for (const auto& step : steps) {
+  for (const auto& step : GetSteps(avoid_combining_reshapes)) {
     switch (step.type) {
       case Step::Type::kReshape: {
         Shape new_shape = ShapeUtil::MakeShape(
@@ -1015,12 +1014,7 @@ std::string ShapeTracker::DebugString(bool avoid_combining_reshapes) const {
   std::vector<int64_t> current_dims(input_shape_.dimensions().begin(),
                                     input_shape_.dimensions().end());
 
-  std::vector<Step> steps = GetSteps();
-  if (avoid_combining_reshapes) {
-    steps = OptimizeSteps(steps, input_shape_, output_shape_);
-  }
-
-  for (const auto& step : steps) {
+  for (const auto& step : GetSteps(avoid_combining_reshapes)) {
     switch (step.type) {
       case Step::Type::kReshape: {
         absl::StrAppend(&result, " -> R[", absl::StrJoin(step.dimensions, ","),
diff --git a/third_party/xla/xla/hlo/analysis/shape_tracker.h b/third_party/xla/xla/hlo/analysis/shape_tracker.h
index f99794f4a02..545ba266803 100644
--- a/third_party/xla/xla/hlo/analysis/shape_tracker.h
+++ b/third_party/xla/xla/hlo/analysis/shape_tracker.h
@@ -171,7 +171,7 @@ class ShapeTracker {
   };
 
   // Returns the sequence of simplified steps represented by this tracker.
-  std::vector<Step> GetSteps() const;
+  std::vector<Step> GetSteps(bool avoid_combining_reshapes = true) const;
 
   // Returns a debug string representation of the tracked operations.
   std::string DebugString(bool avoid_combining_reshapes = true) const;
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 bba5b2cfec9..0d5e628d191 100644
--- a/third_party/xla/xla/hlo/analysis/shape_tracker_test.cc
+++ b/third_party/xla/xla/hlo/analysis/shape_tracker_test.cc
@@ -238,6 +238,17 @@ TEST(ShapeTrackerTest, GetStepsSingleTranspose) {
   EXPECT_EQ(steps[0].dimensions, (std::vector<int64_t>{1, 0}));
 }
 
+TEST(ShapeTrackerTest, GetStepsSingleTransposeMergingDimensions) {
+  Shape shape = ShapeUtil::MakeShape(F32, {4, 8, 16, 32});
+  ShapeTracker tracker(shape);
+  ASSERT_TRUE(tracker.AppendTranspose({0, 2, 3, 1}).ok());
+
+  std::vector<ShapeTracker::Step> steps = tracker.GetSteps();
+  ASSERT_EQ(steps.size(), 1);
+  EXPECT_EQ(steps[0].type, ShapeTracker::Step::Type::kTranspose);
+  EXPECT_EQ(steps[0].dimensions, (std::vector<int64_t>{0, 2, 3, 1}));
+}
+
 TEST(ShapeTrackerTest, GetInvertedDoubleIdentity) {
   Shape shape = ShapeUtil::MakeShape(F32, {2, 3, 4});
   ShapeTracker tracker(shape);
@@ -335,13 +346,11 @@ TEST(ShapeTrackerTest, GetStepsAndDebugString) {
   ASSERT_TRUE(tracker.AppendTranspose({0, 2, 1}).ok());
 
   std::vector<ShapeTracker::Step> steps = tracker.GetSteps();
-  ASSERT_EQ(steps.size(), 3);
-  EXPECT_EQ(steps[0].type, ShapeTracker::Step::Type::kReshape);
-  EXPECT_EQ(steps[0].dimensions, (std::vector<int64_t>{2, 3}));
-  EXPECT_EQ(steps[1].type, ShapeTracker::Step::Type::kTranspose);
-  EXPECT_EQ(steps[1].dimensions, (std::vector<int64_t>{1, 0}));
-  EXPECT_EQ(steps[2].type, ShapeTracker::Step::Type::kReshape);
-  EXPECT_EQ(steps[2].dimensions, (std::vector<int64_t>{1, 3, 2}));
+  ASSERT_EQ(steps.size(), 2);
+  EXPECT_EQ(steps[0].type, ShapeTracker::Step::Type::kTranspose);
+  EXPECT_EQ(steps[0].dimensions, (std::vector<int64_t>{2, 0, 1}));
+  EXPECT_EQ(steps[1].type, ShapeTracker::Step::Type::kReshape);
+  EXPECT_EQ(steps[1].dimensions, (std::vector<int64_t>{1, 3, 2}));
 
   EXPECT_EQ(tracker.DebugString(), "[1,2,3] -> T[3,1,2] -> R[1,3,2]");
 }
@@ -356,13 +365,11 @@ TEST(ShapeTrackerTest, GetStepsFallback) {
   EXPECT_EQ(tracker.output_shape().dimensions(), (std::vector<int64_t>{8, 3}));
 
   std::vector<ShapeTracker::Step> steps = tracker.GetSteps();
-  ASSERT_EQ(steps.size(), 3);
-  EXPECT_EQ(steps[0].type, ShapeTracker::Step::Type::kReshape);
-  EXPECT_EQ(steps[0].dimensions, (std::vector<int64_t>{6, 4}));
-  EXPECT_EQ(steps[1].type, ShapeTracker::Step::Type::kTranspose);
-  EXPECT_EQ(steps[1].dimensions, (std::vector<int64_t>{1, 0}));
-  EXPECT_EQ(steps[2].type, ShapeTracker::Step::Type::kReshape);
-  EXPECT_EQ(steps[2].dimensions, (std::vector<int64_t>{8, 3}));
+  ASSERT_EQ(steps.size(), 2);
+  EXPECT_EQ(steps[0].type, ShapeTracker::Step::Type::kTranspose);
+  EXPECT_EQ(steps[0].dimensions, (std::vector<int64_t>{2, 0, 1}));
+  EXPECT_EQ(steps[1].type, ShapeTracker::Step::Type::kReshape);
+  EXPECT_EQ(steps[1].dimensions, (std::vector<int64_t>{8, 3}));
 
   EXPECT_EQ(tracker.DebugString(), "[2,3,4] -> T[4,2,3] -> R[8,3]");
 }
@@ -395,13 +402,9 @@ TEST(ShapeTrackerTest, PrependTransposeNonSymmetric) {
             (std::vector<int64_t>{4, 2, 3}));
 
   std::vector<ShapeTracker::Step> steps = tracker.GetSteps();
-  ASSERT_EQ(steps.size(), 3);
-  EXPECT_EQ(steps[0].type, ShapeTracker::Step::Type::kReshape);
-  EXPECT_EQ(steps[0].dimensions, (std::vector<int64_t>{4, 6}));
-  EXPECT_EQ(steps[1].type, ShapeTracker::Step::Type::kTranspose);
-  EXPECT_EQ(steps[1].dimensions, (std::vector<int64_t>{1, 0}));
-  EXPECT_EQ(steps[2].type, ShapeTracker::Step::Type::kReshape);
-  EXPECT_EQ(steps[2].dimensions, (std::vector<int64_t>{2, 3, 4}));
+  ASSERT_EQ(steps.size(), 1);
+  EXPECT_EQ(steps[0].type, ShapeTracker::Step::Type::kTranspose);
+  EXPECT_EQ(steps[0].dimensions, (std::vector<int64_t>{1, 2, 0}));
 
   EXPECT_EQ(tracker.DebugString(), "[4,2,3] -> T[2,3,4]");
 }
@@ -719,7 +722,7 @@ TEST(ShapeTrackerTest, SimplifyConsecutiveReshapes) {
   ASSERT_TRUE(tracker.AppendTranspose({2, 1, 0}).ok());
   ASSERT_TRUE(tracker.AppendReshape({1, 512, 5376}).ok());
 
-  std::vector<ShapeTracker::Step> steps = tracker.GetSteps();
+  std::vector<ShapeTracker::Step> steps = tracker.GetSteps(false);
 
   // The intermediate reshape to [512, 5376] is skipped, and we go directly
   // from [5376, 512] to [1, 512, 5376] via a single reshape.
@@ -730,6 +733,13 @@ TEST(ShapeTrackerTest, SimplifyConsecutiveReshapes) {
   EXPECT_EQ(steps[1].dimensions, (std::vector<int64_t>{1, 0}));
   EXPECT_EQ(steps[2].type, ShapeTracker::Step::Type::kReshape);
   EXPECT_EQ(steps[2].dimensions, (std::vector<int64_t>{1, 512, 5376}));
+
+  std::vector<ShapeTracker::Step> opt_steps = tracker.GetSteps();
+  ASSERT_EQ(opt_steps.size(), 2);
+  EXPECT_EQ(opt_steps[0].type, ShapeTracker::Step::Type::kTranspose);
+  EXPECT_EQ(opt_steps[0].dimensions, (std::vector<int64_t>{2, 0, 1}));
+  EXPECT_EQ(opt_steps[1].type, ShapeTracker::Step::Type::kReshape);
+  EXPECT_EQ(opt_steps[1].dimensions, (std::vector<int64_t>{1, 512, 5376}));
 }
 
 TEST(ShapeTrackerTest, AppendBitcastOneSixToSix) {
@@ -1456,13 +1466,11 @@ TEST(ShapeTrackerTest, ZipWithOneElementShape) {
             (std::vector<int64_t>{1, 3, 2}));
 
   std::vector<ShapeTracker::Step> steps = zipped.GetSteps();
-  ASSERT_EQ(steps.size(), 3);
-  EXPECT_EQ(steps[0].type, ShapeTracker::Step::Type::kReshape);
-  EXPECT_EQ(steps[0].dimensions, (std::vector<int64_t>{2, 3}));
-  EXPECT_EQ(steps[1].type, ShapeTracker::Step::Type::kTranspose);
-  EXPECT_EQ(steps[1].dimensions, (std::vector<int64_t>{1, 0}));
-  EXPECT_EQ(steps[2].type, ShapeTracker::Step::Type::kReshape);
-  EXPECT_EQ(steps[2].dimensions, (std::vector<int64_t>{1, 3, 2}));
+  ASSERT_EQ(steps.size(), 2);
+  EXPECT_EQ(steps[0].type, ShapeTracker::Step::Type::kTranspose);
+  EXPECT_EQ(steps[0].dimensions, (std::vector<int64_t>{2, 0, 1}));
+  EXPECT_EQ(steps[1].type, ShapeTracker::Step::Type::kReshape);
+  EXPECT_EQ(steps[1].dimensions, (std::vector<int64_t>{1, 3, 2}));
 
   EXPECT_EQ(zipped.DebugString(), "[1,2,3] -> T[3,1,2] -> R[1,3,2]");
 }

From 7a793778649067d397e96bbb52859c02a40813f3 Mon Sep 17 00:00:00 2001
From: Aliia Khasanova <aliia@google.com>
Date: Wed, 14 Oct 2026 03:41:35 -0700
Subject: [PATCH 0756/1000] Compare only runtime DebugOptions in GPU AOT golden
 verification

Golden verification compared every DebugOptions field except a hand-kept
denylist, so any compile-only flag change broke the goldens, although it cannot
change a loaded executable. The collective ops AOT tests fail this way today,
on five such flags.

Compare only the fields in the shared runtime list instead. Ignore
`xla_dump_to`, because its path is host-specific.

Known limits:

-   New DebugOptions fields are ignored unless they are added to the runtime
    list.
-   Compile-only flags that change only machine code, e.g. `xla_gpu_ftz`, go
    undetected.
-   An empty vs non-empty `xla_dump_to` goes undetected.

PiperOrigin-RevId: 999261092
---
 third_party/xla/xla/tests/BUILD               |   2 +
 .../xla/xla/tests/aot_compatibility/BUILD     |   1 +
 .../tests/aot_compatibility/test_lib_test.cc  | 334 ++++++++++++++++--
 .../xla/tests/aot_interception_pjrt_client.cc | 141 ++------
 4 files changed, 336 insertions(+), 142 deletions(-)

diff --git a/third_party/xla/xla/tests/BUILD b/third_party/xla/xla/tests/BUILD
index 7cb302036ee..2cb37351da6 100644
--- a/third_party/xla/xla/tests/BUILD
+++ b/third_party/xla/xla/tests/BUILD
@@ -3713,6 +3713,7 @@ cc_library(
     srcs = ["aot_interception_pjrt_client.cc"],
     hdrs = ["aot_interception_pjrt_client.h"],
     deps = [
+        "//xla:debug_options_runtime_fields",
         "//xla:literal",
         "//xla:shape_util",
         "//xla:status_macros",
@@ -3735,6 +3736,7 @@ cc_library(
         "//xla/util/split_proto:split_executable_and_options_writer",
         "//xla/util/split_proto:split_gpu_executable_writer",
         "//xla/util/split_proto:split_proto_reader",
+        "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
diff --git a/third_party/xla/xla/tests/aot_compatibility/BUILD b/third_party/xla/xla/tests/aot_compatibility/BUILD
index dff85198bf7..d74abfbc7bc 100644
--- a/third_party/xla/xla/tests/aot_compatibility/BUILD
+++ b/third_party/xla/xla/tests/aot_compatibility/BUILD
@@ -41,6 +41,7 @@ xla_cc_test(
     tags = ["no_oss"],
     deps = [
         ":test_lib",
+        "//xla:debug_options_runtime_fields",
         "//xla:xla_proto_cc",
         "//xla/pjrt/proto:compile_options_proto_cc",
         "//xla/service/cpu:executable_proto_cc",
diff --git a/third_party/xla/xla/tests/aot_compatibility/test_lib_test.cc b/third_party/xla/xla/tests/aot_compatibility/test_lib_test.cc
index 1d28217ffaf..4b003aeacb8 100644
--- a/third_party/xla/xla/tests/aot_compatibility/test_lib_test.cc
+++ b/third_party/xla/xla/tests/aot_compatibility/test_lib_test.cc
@@ -19,14 +19,18 @@ limitations under the License.
 
 #include <cstdint>
 #include <cstdlib>
+#include <optional>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
 #include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
+#include "xla/debug_options_runtime_fields.h"
 #include "xla/pjrt/proto/compile_options.pb.h"
 #include "xla/service/cpu/executable.pb.h"
 #include "xla/service/gpu/gpu_executable.pb.h"
@@ -287,68 +291,306 @@ HumanReadableAotExecutable GpuExecutableWithDebugOptions(
   return executable;
 }
 
+// Returns an enum value of `field` other than its default, or null if the enum
+// has no other value.
+const tsl::protobuf::EnumValueDescriptor* NonDefaultEnumValue(
+    const tsl::protobuf::FieldDescriptor* field) {
+  const tsl::protobuf::EnumValueDescriptor* default_value =
+      field->default_value_enum();
+  const tsl::protobuf::EnumDescriptor* type = field->enum_type();
+  for (int i = 0; i < type->value_count(); ++i) {
+    if (default_value == nullptr ||
+        type->value(i)->number() != default_value->number()) {
+      return type->value(i);
+    }
+  }
+  return nullptr;
+}
+
+// Sets `field` of `message` to a value that differs from its default: a
+// singular field gets a value other than its default, a repeated field gets one
+// such element, and a message gets its first field set the same way. Returns
+// false if it cannot, e.g. for a type it does not know.
+bool SetToNonDefault(const tsl::protobuf::FieldDescriptor* field,
+                     tsl::protobuf::Message* message, int depth = 0) {
+  using FieldDescriptor = tsl::protobuf::FieldDescriptor;
+  const tsl::protobuf::Reflection* reflection = message->GetReflection();
+  const bool repeated = field->is_repeated();
+  switch (field->cpp_type()) {
+    case FieldDescriptor::CPPTYPE_BOOL: {
+      const bool value = !field->default_value_bool();
+      if (repeated) {
+        reflection->AddBool(message, field, value);
+      } else {
+        reflection->SetBool(message, field, value);
+      }
+      return true;
+    }
+    case FieldDescriptor::CPPTYPE_INT32: {
+      const int32_t value = field->default_value_int32() + 1;
+      if (repeated) {
+        reflection->AddInt32(message, field, value);
+      } else {
+        reflection->SetInt32(message, field, value);
+      }
+      return true;
+    }
+    case FieldDescriptor::CPPTYPE_INT64: {
+      const int64_t value = field->default_value_int64() + 1;
+      if (repeated) {
+        reflection->AddInt64(message, field, value);
+      } else {
+        reflection->SetInt64(message, field, value);
+      }
+      return true;
+    }
+    case FieldDescriptor::CPPTYPE_UINT32: {
+      const uint32_t value = field->default_value_uint32() + 1;
+      if (repeated) {
+        reflection->AddUInt32(message, field, value);
+      } else {
+        reflection->SetUInt32(message, field, value);
+      }
+      return true;
+    }
+    case FieldDescriptor::CPPTYPE_UINT64: {
+      const uint64_t value = field->default_value_uint64() + 1;
+      if (repeated) {
+        reflection->AddUInt64(message, field, value);
+      } else {
+        reflection->SetUInt64(message, field, value);
+      }
+      return true;
+    }
+    case FieldDescriptor::CPPTYPE_FLOAT: {
+      const float value = field->default_value_float() + 1;
+      if (repeated) {
+        reflection->AddFloat(message, field, value);
+      } else {
+        reflection->SetFloat(message, field, value);
+      }
+      return true;
+    }
+    case FieldDescriptor::CPPTYPE_DOUBLE: {
+      const double value = field->default_value_double() + 1;
+      if (repeated) {
+        reflection->AddDouble(message, field, value);
+      } else {
+        reflection->SetDouble(message, field, value);
+      }
+      return true;
+    }
+    case FieldDescriptor::CPPTYPE_STRING: {
+      std::string value =
+          absl::StrCat(field->default_value_string(), "_non_default");
+      if (repeated) {
+        reflection->AddString(message, field, std::move(value));
+      } else {
+        reflection->SetString(message, field, std::move(value));
+      }
+      return true;
+    }
+    case FieldDescriptor::CPPTYPE_ENUM: {
+      const tsl::protobuf::EnumValueDescriptor* value =
+          NonDefaultEnumValue(field);
+      if (value == nullptr) {
+        return false;
+      }
+      if (repeated) {
+        reflection->AddEnum(message, field, value);
+      } else {
+        reflection->SetEnum(message, field, value);
+      }
+      return true;
+    }
+    case FieldDescriptor::CPPTYPE_MESSAGE: {
+      // An empty submessage equals an unset one in an EQUIVALENT comparison,
+      // so one of its fields must be set as well. The depth limit stops a
+      // message type that contains itself.
+      constexpr int kMaxDepth = 4;
+      const tsl::protobuf::Descriptor* type = field->message_type();
+      if (depth >= kMaxDepth || type->field_count() == 0) {
+        return false;
+      }
+      tsl::protobuf::Message* child =
+          repeated ? reflection->AddMessage(message, field)
+                   : reflection->MutableMessage(message, field);
+      return SetToNonDefault(type->field(0), child, depth + 1);
+    }
+  }
+  return false;
+}
+
+// Returns `DebugOptions` with only `field` set, to a non-default value. Fails
+// the test and returns nullopt if `SetToNonDefault` cannot set it.
+std::optional<DebugOptions> DebugOptionsWithNonDefault(
+    const tsl::protobuf::FieldDescriptor* field) {
+  DebugOptions debug_options;
+  if (!SetToNonDefault(field, &debug_options)) {
+    ADD_FAILURE() << "Cannot set DebugOptions." << field->name() << " of type "
+                  << field->type_name()
+                  << " to a non-default value; teach SetToNonDefault to.";
+    return std::nullopt;
+  }
+  // Guards against a helper bug that would make every "ignores" check pass.
+  EXPECT_FALSE(tsl::protobuf::util::MessageDifferencer::Equivalent(
+      debug_options, DebugOptions()))
+      << field->name();
+  return debug_options;
+}
+
 class CompareGPUExecutablesDebugOptionsTest
-    : public ::testing::TestWithParam<DebugOptionsCopy> {};
+    : public ::testing::TestWithParam<DebugOptionsCopy> {
+ protected:
+  // Compares two GPU executables that differ only in the `GetParam()` copy of
+  // their debug options.
+  static absl::Status Compare(const DebugOptions& fresh,
+                              const DebugOptions& golden) {
+    return AOTInterceptionPjrtClient::CompareGPUExecutables(
+        GpuExecutableWithDebugOptions(GetParam(), fresh),
+        GpuExecutableWithDebugOptions(GetParam(), golden));
+  }
+};
 
-// Dump paths and host-specific paths cannot change the compiled executable, so
-// a difference in them must not cause a mismatch.
-TEST_P(CompareGPUExecutablesDebugOptionsTest, IgnoresDenylistedFields) {
-  DebugOptions golden_options;
-  golden_options.set_xla_dump_to("/tmp/golden_dir");
-  golden_options.set_xla_gpu_cuda_data_dir("/host_a/cuda");
+// The fields whose defaults changed since the v3 goldens of the GPU collective
+// AOT tests. None can change a loaded executable, so a difference in them must
+// not cause a mismatch.
+TEST_P(CompareGPUExecutablesDebugOptionsTest, IgnoresCompileOnlyFields) {
+  constexpr int kCompileOnlyFields[] = {
+      DebugOptions::kXlaGpuExperimentalAutotuneBackendsFieldNumber,
+      DebugOptions::kXlaGpuExperimentalDeterministicOpsTritonFieldNumber,
+      DebugOptions::kXlaGpuExperimentalEnableConvFusionFieldNumber,
+      DebugOptions::kXlaGpuExperimentalGemmFusionV2FieldNumber,
+      DebugOptions::kXlaGpuUnsupportedUseCrossHostOneShotKernelFieldNumber,
+  };
   DebugOptions fresh_options;
-  fresh_options.set_xla_dump_to("/tmp/fresh_dir");
-  fresh_options.set_xla_gpu_cuda_data_dir("/host_b/cuda");
+  for (const int number : kCompileOnlyFields) {
+    ASSERT_FALSE(IsRuntimeDebugOptionsField(number)) << number;
+    const tsl::protobuf::FieldDescriptor* field =
+        DebugOptions::descriptor()->FindFieldByNumber(number);
+    ASSERT_NE(field, nullptr) << number;
+    ASSERT_TRUE(SetToNonDefault(field, &fresh_options)) << field->name();
+  }
 
-  EXPECT_OK(AOTInterceptionPjrtClient::CompareGPUExecutables(
-      GpuExecutableWithDebugOptions(GetParam(), fresh_options),
-      GpuExecutableWithDebugOptions(GetParam(), golden_options)));
+  EXPECT_OK(Compare(fresh_options, DebugOptions()));
 }
 
 TEST_P(CompareGPUExecutablesDebugOptionsTest, DetectsNumericsFlagDifference) {
+  ASSERT_TRUE(
+      IsRuntimeDebugOptionsField(DebugOptions::kXlaGpuDetectNanFieldNumber));
   DebugOptions golden_options;
-  golden_options.set_xla_gpu_enable_fast_min_max(true);
+  golden_options.set_xla_gpu_detect_nan(DebugOptions::DETECTION_MODE_WARNING);
   DebugOptions fresh_options;
-  fresh_options.set_xla_gpu_enable_fast_min_max(false);
+  fresh_options.set_xla_gpu_detect_nan(DebugOptions::DETECTION_MODE_FAIL);
 
-  EXPECT_THAT(AOTInterceptionPjrtClient::CompareGPUExecutables(
-                  GpuExecutableWithDebugOptions(GetParam(), fresh_options),
-                  GpuExecutableWithDebugOptions(GetParam(), golden_options)),
-              StatusIs(absl::StatusCode::kInternal,
-                       HasSubstr("xla_gpu_enable_fast_min_max")));
+  EXPECT_THAT(
+      Compare(fresh_options, golden_options),
+      StatusIs(absl::StatusCode::kInternal, HasSubstr("xla_gpu_detect_nan")));
 }
 
 TEST_P(CompareGPUExecutablesDebugOptionsTest, DetectsCollectiveFlagDifference) {
+  ASSERT_TRUE(IsRuntimeDebugOptionsField(
+      DebugOptions::kXlaGpuNcclTerminationTimeoutSecondsFieldNumber));
   DebugOptions golden_options;
-  golden_options.set_xla_gpu_all_reduce_combine_threshold_bytes(31457287);
+  golden_options.set_xla_gpu_nccl_termination_timeout_seconds(600);
   DebugOptions fresh_options;
-  fresh_options.set_xla_gpu_all_reduce_combine_threshold_bytes(1024);
+  fresh_options.set_xla_gpu_nccl_termination_timeout_seconds(30);
 
-  EXPECT_THAT(
-      AOTInterceptionPjrtClient::CompareGPUExecutables(
-          GpuExecutableWithDebugOptions(GetParam(), fresh_options),
-          GpuExecutableWithDebugOptions(GetParam(), golden_options)),
-      StatusIs(absl::StatusCode::kInternal,
-               HasSubstr("xla_gpu_all_reduce_combine_threshold_bytes")));
+  EXPECT_THAT(Compare(fresh_options, golden_options),
+              StatusIs(absl::StatusCode::kInternal,
+                       HasSubstr("xla_gpu_nccl_termination_timeout_seconds")));
 }
 
-// Ignoring a denylisted field must not hide a compared field that differs
-// alongside it.
+// Ignoring fields, `xla_dump_to` included, must not hide a runtime field that
+// differs alongside them, and must keep the ignored ones out of the report.
 TEST_P(CompareGPUExecutablesDebugOptionsTest,
-       DetectsComparedFlagDifferenceNextToDenylistedOne) {
+       DetectsRuntimeFlagDifferenceNextToIgnoredOnes) {
+  ASSERT_TRUE(
+      IsRuntimeDebugOptionsField(DebugOptions::kXlaGpuDetectInfFieldNumber));
+  ASSERT_FALSE(IsRuntimeDebugOptionsField(
+      DebugOptions::kXlaGpuExperimentalEnableConvFusionFieldNumber));
   DebugOptions golden_options;
   golden_options.set_xla_dump_to("/tmp/golden_dir");
-  golden_options.set_xla_gpu_enable_fast_min_max(true);
+  golden_options.set_xla_gpu_experimental_enable_conv_fusion(false);
+  golden_options.set_xla_gpu_detect_inf(DebugOptions::DETECTION_MODE_NONE);
   DebugOptions fresh_options;
   fresh_options.set_xla_dump_to("/tmp/fresh_dir");
-  fresh_options.set_xla_gpu_enable_fast_min_max(false);
+  fresh_options.set_xla_gpu_experimental_enable_conv_fusion(true);
+  fresh_options.set_xla_gpu_detect_inf(DebugOptions::DETECTION_MODE_FAIL);
 
-  EXPECT_THAT(AOTInterceptionPjrtClient::CompareGPUExecutables(
-                  GpuExecutableWithDebugOptions(GetParam(), fresh_options),
-                  GpuExecutableWithDebugOptions(GetParam(), golden_options)),
-              StatusIs(absl::StatusCode::kInternal,
-                       AllOf(HasSubstr("xla_gpu_enable_fast_min_max"),
-                             Not(HasSubstr("xla_dump_to")))));
+  EXPECT_THAT(
+      Compare(fresh_options, golden_options),
+      StatusIs(
+          absl::StatusCode::kInternal,
+          AllOf(HasSubstr("xla_gpu_detect_inf"), Not(HasSubstr("xla_dump_to")),
+                Not(HasSubstr("xla_gpu_experimental_enable_conv_fusion")))));
+}
+
+// `xla_dump_to` is ignored although it is a runtime field: the path is
+// host-specific, and the loading process sets its own. So two different dump
+// directories must not cause a mismatch.
+TEST_P(CompareGPUExecutablesDebugOptionsTest, IgnoresXlaDumpToPath) {
+  DebugOptions golden_options;
+  golden_options.set_xla_dump_to("/tmp/golden_dir");
+  DebugOptions fresh_options;
+  fresh_options.set_xla_dump_to("/tmp/fresh_dir");
+
+  EXPECT_OK(Compare(fresh_options, golden_options));
+}
+
+// Nor must an empty `xla_dump_to` vs a non-empty one, in either direction.
+TEST_P(CompareGPUExecutablesDebugOptionsTest, IgnoresXlaDumpToEmptiness) {
+  ASSERT_TRUE(IsRuntimeDebugOptionsField(DebugOptions::kXlaDumpToFieldNumber));
+  DebugOptions empty_options;
+  DebugOptions set_options;
+  set_options.set_xla_dump_to("/tmp/dump_dir");
+
+  EXPECT_OK(Compare(/*fresh=*/set_options, /*golden=*/empty_options));
+  EXPECT_OK(Compare(/*fresh=*/empty_options, /*golden=*/set_options));
+}
+
+// Every runtime field must be compared, whatever its type, except
+// `xla_dump_to`, which is ignored and has its own tests.
+TEST_P(CompareGPUExecutablesDebugOptionsTest,
+       DetectsEveryRuntimeFieldDifference) {
+  for (const int number : RuntimeDebugOptionsFieldNumbers()) {
+    if (number == DebugOptions::kXlaDumpToFieldNumber) {
+      continue;
+    }
+    const tsl::protobuf::FieldDescriptor* field =
+        DebugOptions::descriptor()->FindFieldByNumber(number);
+    ASSERT_NE(field, nullptr) << "No DebugOptions field number " << number;
+    SCOPED_TRACE(field->name());
+    std::optional<DebugOptions> fresh_options =
+        DebugOptionsWithNonDefault(field);
+    if (!fresh_options.has_value()) {
+      continue;
+    }
+
+    EXPECT_THAT(Compare(*fresh_options, DebugOptions()),
+                StatusIs(absl::StatusCode::kInternal,
+                         HasSubstr(std::string(field->name()))));
+  }
+}
+
+// Every field outside the runtime list must be ignored, whatever its type.
+TEST_P(CompareGPUExecutablesDebugOptionsTest,
+       IgnoresEveryNonRuntimeFieldDifference) {
+  const tsl::protobuf::Descriptor* descriptor = DebugOptions::descriptor();
+  for (int i = 0; i < descriptor->field_count(); ++i) {
+    const tsl::protobuf::FieldDescriptor* field = descriptor->field(i);
+    if (IsRuntimeDebugOptionsField(field->number())) {
+      continue;
+    }
+    SCOPED_TRACE(field->name());
+    std::optional<DebugOptions> fresh_options =
+        DebugOptionsWithNonDefault(field);
+    if (!fresh_options.has_value()) {
+      continue;
+    }
+
+    EXPECT_OK(Compare(*fresh_options, DebugOptions()));
+  }
 }
 
 INSTANTIATE_TEST_SUITE_P(
@@ -376,9 +618,23 @@ TEST(TestLibTest, CompareGPUExecutables_FailsOnGenuineDifferences) {
                      HasSubstr("different_module"))));
 }
 
+// An executable whose payload is not a `gpu_executable` must still be compared
+// in full, not treated as equal.
+TEST(TestLibTest, CompareGPUExecutables_DetectsDifferencesInNonGpuPayload) {
+  HumanReadableAotExecutable golden;
+  golden.mutable_cpu_executable()->set_entry_function_name("golden_entry");
+  HumanReadableAotExecutable fresh;
+  fresh.mutable_cpu_executable()->set_entry_function_name("fresh_entry");
+
+  EXPECT_THAT(AOTInterceptionPjrtClient::CompareGPUExecutables(fresh, golden),
+              StatusIs(absl::StatusCode::kInternal,
+                       AllOf(HasSubstr("entry_function_name"),
+                             HasSubstr("fresh_entry"))));
+}
+
 // The CPU comparator must ignore both copies of debug_options as well as the
 // compiled machine-code fields (object_files, target_machine_options,
-// data_layout). See b/528258781.
+// data_layout). See b/567825028.
 TEST(TestLibTest, CompareGoldenCPUExecutable_IgnoresDebugOptionsAndReturnsOk) {
   HumanReadableAotExecutable golden;
   golden.mutable_cpu_executable()->set_entry_function_name("main");
diff --git a/third_party/xla/xla/tests/aot_interception_pjrt_client.cc b/third_party/xla/xla/tests/aot_interception_pjrt_client.cc
index e34b2b51699..42e2a65d720 100644
--- a/third_party/xla/xla/tests/aot_interception_pjrt_client.cc
+++ b/third_party/xla/xla/tests/aot_interception_pjrt_client.cc
@@ -24,17 +24,18 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/base/no_destructor.h"
 #include "absl/functional/any_invocable.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
-#include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "riegeli/base/maker.h"
 #include "riegeli/bytes/string_reader.h"
 #include "riegeli/bytes/string_writer.h"
+#include "xla/debug_options_runtime_fields.h"
 #include "xla/hlo/builder/xla_computation.h"
 #include "xla/layout.h"
 #include "xla/literal.h"
@@ -85,94 +86,26 @@ const tsl::protobuf::FieldDescriptor* FieldByPath(
   return field;
 }
 
-// DebugOptions fields that cannot change the compiled program. Every other
-// field, including fields added in future, is compared. That includes
-// runtime-only knobs such as watchdog timeouts: GpuExecutable reads them from
-// the serialized options, so the artifact carries them into every process that
-// loads it. A field belongs here only if every reader of it is in one of the
-// groups below; when in doubt, leave it out. Keep each group sorted.
-// TODO(b/567825028): Replace this hand-kept list once DebugOptions separates
-// compiler options from runtime, debug and environment options.
-constexpr absl::string_view kIgnoredDebugOptionsFields[] = {
-    // Dump controls not covered by `kIgnoredDebugOptionsFieldPrefixes`.
-    "xla_enable_dumping",
-    "xla_gpu_experimental_dump_fdo_profiles",
-    "xla_gpu_experimental_dump_gpu_executable",
-    "xla_hlo_graph_addresses",
-    "xla_hlo_graph_sharding_color",
-    // HLO text rendering: read only by HloModule::ToString(), which formats
-    // dumps, logs and error messages.
-    "xla_hlo_print_inline_stack_frames",
-    "xla_syntax_sugar_async_ops",
-    // Logging, profiling and tracing: log text, timers, profiler and NVTX
-    // payloads only.
-    "xla_debug_buffer_assignment_show_max",
-    "xla_detailed_logging",
-    "xla_enable_hlo_modules_upload",
-    "xla_enable_scoped_logging_timers",
-    "xla_gpu_enable_cupti_multi_subscriber",
-    "xla_gpu_print_compilation_stats",
-    "xla_gpu_rocm_max_trace_events",
-    "xla_gpu_trace_annotation_level",
-    // Host-specific: where the toolchain and compilation caches live and how
-    // many threads compile, not what is compiled.
-    "xla_gpu_cuda_data_dir",
-    "xla_gpu_experimental_autotuner_cache_dir",
-    "xla_gpu_force_compilation_parallelism",
-    "xla_gpu_kernel_cache_file",
-    "xla_gpu_per_fusion_autotune_cache_dir",
-    "xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found",
-    // Verification only: can abort compilation or log, but never changes a
-    // successfully compiled executable.
-    "xla_gpu_crash_on_verification_failures",
-    "xla_gpu_llvm_verification_level",
-    "xla_hlo_pass_fix_detect_cycles",
-    "xla_unsupported_crash_on_hlo_pass_fix_max_iterations",
-    "xla_unsupported_crash_on_hlo_pass_noop_change",
-    "xla_unsupported_crash_on_hlo_pass_silent_hlo_change",
-    // Not read by the GPU compiler or runtime from the stored options: read by
-    // the CPU or TPU compilers only, or only from process-global flags.
-    "xla_embed_ir_in_executable",
-    "xla_flags_reset",
-    "xla_force_host_platform_device_count",
-    "xla_tpu_detect_inf",
-    "xla_tpu_detect_nan",
-    // Test harness only: never read by the compiler or runtime.
-    "xla_test_add_command_buffer_mode",
-    "xla_test_all_input_layouts",
-    "xla_test_all_output_layouts",
-};
-
-// Every DebugOptions field with one of these prefixes is ignored, including
-// fields added in future.
-constexpr absl::string_view kIgnoredDebugOptionsFieldPrefixes[] = {
-    // Read by the CPU backend only.
-    "xla_cpu_",
-    // Dump controls: where, whether and in which format to write debug dumps.
-    // xla_dump_to is also overridden per process at GpuExecutable
-    // deserialization.
-    "xla_dump_",
-    "xla_gpu_dump_",
-};
-
-// Resolves `kIgnoredDebugOptionsFields` plus every field matching
-// `kIgnoredDebugOptionsFieldPrefixes`.
-std::vector<const tsl::protobuf::FieldDescriptor*> IgnoredDebugOptionsFields() {
-  const tsl::protobuf::Descriptor* descriptor = DebugOptions::descriptor();
-  std::vector<const tsl::protobuf::FieldDescriptor*> fields;
-  for (const absl::string_view name : kIgnoredDebugOptionsFields) {
-    fields.push_back(FieldByPath(descriptor, {name}));
-  }
-  for (int i = 0; i < descriptor->field_count(); ++i) {
-    const tsl::protobuf::FieldDescriptor* field = descriptor->field(i);
-    for (const absl::string_view prefix : kIgnoredDebugOptionsFieldPrefixes) {
-      if (absl::StartsWith(field->name(), prefix)) {
-        fields.push_back(field);
-        break;
-      }
-    }
-  }
-  return fields;
+// Returns every DebugOptions field outside `RuntimeDebugOptionsFieldNumbers()`,
+// including fields added in future: fields that cannot change how a loaded
+// executable behaves, so a difference in them is not a compatibility break.
+absl::Span<const tsl::protobuf::FieldDescriptor* const>
+NonRuntimeDebugOptionsFields() {
+  static const absl::NoDestructor<
+      std::vector<const tsl::protobuf::FieldDescriptor*>>
+      kFields([] {
+        const tsl::protobuf::Descriptor* descriptor =
+            DebugOptions::descriptor();
+        std::vector<const tsl::protobuf::FieldDescriptor*> fields;
+        for (int i = 0; i < descriptor->field_count(); ++i) {
+          const tsl::protobuf::FieldDescriptor* field = descriptor->field(i);
+          if (!IsRuntimeDebugOptionsField(field->number())) {
+            fields.push_back(field);
+          }
+        }
+        return fields;
+      }());
+  return *kFields;
 }
 
 // Runs the structural comparison shared by all backends. `extra_ignored_fields`
@@ -197,14 +130,17 @@ absl::Status CompareStructurally(
   ignored_fields.insert(ignored_fields.end(), extra_ignored_fields.begin(),
                         extra_ignored_fields.end());
   // Both debug_options copies (compile options and HLO module config) are
-  // `DebugOptions`, so ignoring its fields applies to both. The list was
-  // classified for GPU; it is a no-op on CPU, which still ignores
-  // debug_options wholesale.
-  const std::vector<const tsl::protobuf::FieldDescriptor*>
-      ignored_debug_options_fields = IgnoredDebugOptionsFields();
+  // `DebugOptions`, so only their runtime fields are compared. This is a no-op
+  // on CPU, which still ignores debug_options entirely.
+  const absl::Span<const tsl::protobuf::FieldDescriptor* const>
+      non_runtime_debug_options_fields = NonRuntimeDebugOptionsFields();
   ignored_fields.insert(ignored_fields.end(),
-                        ignored_debug_options_fields.begin(),
-                        ignored_debug_options_fields.end());
+                        non_runtime_debug_options_fields.begin(),
+                        non_runtime_debug_options_fields.end());
+  // `xla_dump_to` is a runtime field, but its path is host-specific, so this
+  // check ignores the field entirely.
+  ignored_fields.push_back(DebugOptions::descriptor()->FindFieldByNumber(
+      DebugOptions::kXlaDumpToFieldNumber));
   for (const auto* field : ignored_fields) {
     CHECK(field != nullptr)
         << "AOTInterceptionPjrtClient: a proto field descriptor to ignore was "
@@ -333,8 +269,8 @@ absl::Status AOTInterceptionPjrtClient::CompareGPUExecutables(
     const HumanReadableAotExecutable& fresh,
     const HumanReadableAotExecutable& golden) {
   // The ignored fields hold backend machine code and device-specific details.
-  // Both debug_options copies are compared, except for the fields in
-  // `kIgnoredDebugOptionsFields`.
+  // Both debug_options copies are compared on their runtime fields only,
+  // except `xla_dump_to`.
   return CompareStructurally(
       fresh, golden,
       {
@@ -383,11 +319,10 @@ absl::Status AOTInterceptionPjrtClient::CompareGoldenCPUExecutable(
               "target_machine_options"),
           cpu::CompilationResultProto::descriptor()->FindFieldByName(
               "data_layout"),
-          // TODO(b/528258781): Debug options are still ignored wholesale on
-          // CPU: `kIgnoredDebugOptionsFields` was classified for GPU only (it
-          // ignores every xla_cpu_* field), and the CPU golden predates
-          // current DebugOptions defaults. Compare them here too once CPU flags
-          // are classified and the golden regenerated.
+          // TODO(b/567825028): Debug options are still ignored entirely on CPU,
+          // and the CPU golden predates current DebugOptions defaults. Compare
+          // the runtime fields here too once the runtime field list is audited
+          // for CPU completeness and the golden is regenerated.
           ExecutableBuildOptionsProto::descriptor()->FindFieldByName(
               "debug_options"),
           HloModuleConfigProto::descriptor()->FindFieldByName("debug_options"),

From d2f2f22565d433d8f54076d28a45b19502b6f471 Mon Sep 17 00:00:00 2001
From: bodapatisaikrishna <saikrishnabodapati@gmail.com>
Date: Wed, 14 Oct 2026 17:17:39 +0530
Subject: [PATCH 0757/1000] Move in-function imports to module level in
 training_ops_test

---
 tensorflow/python/training/training_ops_test.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/training/training_ops_test.py b/tensorflow/python/training/training_ops_test.py
index 83e56c512dc..d8d98753979 100644
--- a/tensorflow/python/training/training_ops_test.py
+++ b/tensorflow/python/training/training_ops_test.py
@@ -19,6 +19,7 @@ import threading
 
 import numpy as np
 
+from tensorflow.python.eager import context
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
@@ -27,6 +28,7 @@ from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.framework.test_util import TensorFlowTestCase
 # Import resource_variable_ops for the variables-to-tensor implicit conversion.
+from tensorflow.python.ops import gen_resource_variable_ops
 from tensorflow.python.ops import gen_training_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import resource_variable_ops  # pylint: disable=unused-import
@@ -750,9 +752,6 @@ class TrainingOpsTest(TensorFlowTestCase):
 
   @test_util.run_v2_only
   def testResourceVariableUninitialized(self):
-    from tensorflow.python.eager import context
-    from tensorflow.python.ops import gen_resource_variable_ops
-
     shared_name = (
         context.anonymous_name() if context.executing_eagerly() else ""
     )

From 3f1379fa29a46c316eaa8ba31f1cd4112c824a1a Mon Sep 17 00:00:00 2001
From: Sevin Fide Varoglu <svaroglu@nvidia.com>
Date: Wed, 14 Oct 2026 03:55:26 -0700
Subject: [PATCH 0758/1000] PR #47889: [XLA:GPU] Support MultiGpuBarrierKernel
 with more than 32 peers
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/47889

📝 Summary of Changes
Launch kMaxPeers (128) threads per CTA instead of one warp (32), giving each peer its own thread for parallel signal/wait. Fix UB in SyncRemoteBlocks, move signal buffers to __shared__, and test both 1-thread and kMaxPeers-thread paths.

🎯 Justification
Fixes the >32-peer CUDA_ERROR_ILLEGAL_ADDRESS crash (#47283)

🚀 Kind of Contribution
🐛 Bug Fix

📊 Benchmark (for Performance Improvements)
N/A

🧪 Unit Tests:
N/A

🧪 Execution Tests:
Tested a simple benchmark on 64 ranks.

Copybara import of the project:

--
54aa7b6654021e7c7ff45f416d3883e2f23eeff8 by Sevin F. Varoglu <svaroglu@nvidia.com>:

[XLA:GPU] Support MultiGpuBarrierKernel with more than 32 peers

--
6557ee29d43a0553f8bb6e838790e6b82b3d3111 by Sevin F. Varoglu <svaroglu@nvidia.com>:

Address review feedback

--
286b2010ac5320055be306f6deffb39f51fc0793 by Sevin F. Varoglu <svaroglu@nvidia.com>:

Test: launch MultiGpuBarrier with ThreadDim(1) to exercise strided loop

--
d006a40da1700187e8ac421fe95a7c5974ec89b5 by Sevin F. Varoglu <svaroglu@nvidia.com>:

Add review feedback

--
6655f6ff9591800278d9fa15be2581c89d52a6b1 by Sevin F. Varoglu <svaroglu@nvidia.com>:

Add review feedback

Merging this change closes #47889

PiperOrigin-RevId: 999266712
---
 .../backends/gpu/runtime/multi_gpu_barrier.cc |   5 +-
 .../xla/xla/stream_executor/cuda/BUILD        |   1 +
 .../cuda/multi_gpu_barrier_kernel_cuda.cu.cc  |  17 +-
 .../cuda/multi_gpu_barrier_test.cc            | 202 +++++++++---------
 .../gpu/collective_signal.cu.h                |  21 +-
 .../gpu/multi_gpu_barrier_kernel.cu.h         |  36 ++--
 .../gpu/multi_gpu_barrier_kernel.h            |  10 +-
 7 files changed, 146 insertions(+), 146 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.cc b/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.cc
index c80846081d0..a478e23f367 100644
--- a/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.cc
@@ -129,6 +129,9 @@ absl::Status LaunchMultiGpuBarrierWithNccl(
 
   TF_RET_CHECK(symmetric_memory != nullptr)
       << "Symmetric memory is required on rank " << rank.value();
+  TF_RET_CHECK(num_devices <= MultiGpuBarrierWithNcclKernel::kMaxPeers)
+      << "Number of participants exceeds "
+         "MultiGpuBarrierWithNcclKernel::kMaxPeers";
 
   ABSL_ASSIGN_OR_RETURN(
       MultiGpuBarrierWithNcclKernel::KernelType * kernel,
@@ -191,7 +194,7 @@ absl::StatusOr<const XLA_FFI_Command*> RecordMultiGpuBarrierWithNccl(
 }
 
 size_t GetMultiGpuBarrierSignalBufferSize() {
-  return stream_executor::gpu::MultiGpuBarrierKernel::kMaxPeers *
+  return stream_executor::gpu::MultiGpuBarrierWithNcclKernel::kMaxPeers *
          sizeof(uint32_t);
 }
 
diff --git a/third_party/xla/xla/stream_executor/cuda/BUILD b/third_party/xla/xla/stream_executor/cuda/BUILD
index dbd2b0af4b9..cf50493f7e4 100644
--- a/third_party/xla/xla/stream_executor/cuda/BUILD
+++ b/third_party/xla/xla/stream_executor/cuda/BUILD
@@ -3003,6 +3003,7 @@ xla_test(
         "//xla/tsl/platform:env",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest_main",
     ],
 )
diff --git a/third_party/xla/xla/stream_executor/cuda/multi_gpu_barrier_kernel_cuda.cu.cc b/third_party/xla/xla/stream_executor/cuda/multi_gpu_barrier_kernel_cuda.cu.cc
index 18ab5fab658..6c7ada59f23 100644
--- a/third_party/xla/xla/stream_executor/cuda/multi_gpu_barrier_kernel_cuda.cu.cc
+++ b/third_party/xla/xla/stream_executor/cuda/multi_gpu_barrier_kernel_cuda.cu.cc
@@ -33,19 +33,18 @@ namespace stream_executor::gpu {
 __global__ void MultiGpuBarrierWithNcclKernelImpl(
     int64_t rank, int64_t num_ranks, ncclWindow_t signal_buffers_handle,
     uint32_t* sync_counter) {
-  // 1. Get individual signal buffers pointers.
-  std::array<uint32_t* __restrict__, MultiGpuBarrierWithNcclKernel::kMaxPeers>
+  __shared__ std::array<uint32_t* __restrict__,
+                        MultiGpuBarrierWithNcclKernel::kMaxPeers>
       signal_buffers;
 
-#pragma unroll
-  for (int64_t i = 0; i < MultiGpuBarrierWithNcclKernel::kMaxPeers; ++i) {
-    if (i < num_ranks) {
-      signal_buffers[i] = reinterpret_cast<uint32_t*>(
-          ncclGetLsaPointer(signal_buffers_handle, 0, i));
-    }
+  for (int64_t i = threadIdx.x; i < num_ranks; i += blockDim.x) {
+    signal_buffers[i] = reinterpret_cast<uint32_t*>(
+        ncclGetLsaPointer(signal_buffers_handle, 0, i));
   }
+  __syncthreads();
 
-  SyncRemoteBlocksAndUpdateCounter<PlatformType::kCuda>(
+  SyncRemoteBlocksAndUpdateCounter<PlatformType::kCuda,
+                                   MultiGpuBarrierWithNcclKernel::kMaxPeers>(
       rank, num_ranks, signal_buffers, sync_counter);
 }
 
diff --git a/third_party/xla/xla/stream_executor/cuda/multi_gpu_barrier_test.cc b/third_party/xla/xla/stream_executor/cuda/multi_gpu_barrier_test.cc
index 6bb2db06cfb..c634c8d86f5 100644
--- a/third_party/xla/xla/stream_executor/cuda/multi_gpu_barrier_test.cc
+++ b/third_party/xla/xla/stream_executor/cuda/multi_gpu_barrier_test.cc
@@ -24,6 +24,7 @@ limitations under the License.
 #include <gtest/gtest.h>
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
 #include "xla/backends/gpu/collectives/nccl_symmetric_memory.h"
 #include "xla/core/collectives/symmetric_memory.h"
 #include "xla/stream_executor/cuda/cuda_platform_id.h"
@@ -62,7 +63,7 @@ class MultiGpuBarrierTest : public ::testing::Test {
                    << visible_device_count;
     }
 
-    // Limit to MultiGpuBarrierKernel::kMaxPeers (32)
+    // Limit to MultiGpuBarrierKernel::kMaxPeers
     num_devices_ =
         std::min<int>(visible_device_count, MultiGpuBarrierKernel::kMaxPeers);
 
@@ -153,76 +154,68 @@ absl::StatusOr<std::vector<T>> CopyToHost(Stream* stream,
 }
 
 TEST_F(MultiGpuBarrierTest, BarrierSynchronization) {
-  // 1. Allocate Signal Buffers on each device
   std::vector<DeviceAddress<uint32_t>> signal_buffers;
-
   for (int i = 0; i < num_devices_; ++i) {
-    auto alloc = executors_[i]->AllocateArray<uint32_t>(num_devices_);
-    ASSERT_OK(streams_[i]->MemZero(&alloc, num_devices_ * sizeof(uint32_t)));
-    signal_buffers.push_back(alloc);
+    signal_buffers.push_back(
+        executors_[i]->AllocateArray<uint32_t>(num_devices_));
   }
 
-  // 2. Allocate Counters on each device
   std::vector<DeviceAddress<uint32_t>> counters;
-  std::vector<void*> counter_ptrs;
   for (int i = 0; i < num_devices_; ++i) {
-    auto c = executors_[i]->AllocateArray<uint32_t>(1);
-    // It is ok to initialize the counter to 0,
-    // as the kernel pre-increments signal_value.
-    ASSERT_OK(streams_[i]->MemZero(&c, sizeof(uint32_t)));
-    counters.push_back(c);
-    counter_ptrs.push_back(c.opaque());
+    counters.push_back(executors_[i]->AllocateArray<uint32_t>(1));
   }
 
-  // 3. Prepare Kernel Arguments
-  std::array<void*, MultiGpuBarrierKernel::kMaxPeers> kernel_arg_ptrs;
+  std::array<void*, MultiGpuBarrierKernel::kMaxPeers> kernel_arg_ptrs = {};
   for (int i = 0; i < num_devices_; ++i) {
     kernel_arg_ptrs[i] = signal_buffers[i].opaque();
   }
 
-  // 4. Launch Kernel REPEATEDLY 8 times to verify auto-increment
-  for (int step = 0; step < 8; ++step) {
+  for (int64_t threads_per_block :
+       {int64_t{1}, MultiGpuBarrierKernel::kMaxPeers}) {
+    SCOPED_TRACE(absl::StrCat("threads_per_block=", threads_per_block));
+
+    for (int i = 0; i < num_devices_; ++i) {
+      ASSERT_OK(streams_[i]->MemZero(&signal_buffers[i],
+                                     num_devices_ * sizeof(uint32_t)));
+      ASSERT_OK(streams_[i]->MemZero(&counters[i], sizeof(uint32_t)));
+    }
+
+    for (int step = 0; step < 8; ++step) {
+      for (int i = 0; i < num_devices_; ++i) {
+        ASSERT_OK_AND_ASSIGN(
+            auto kernel,
+            (GpuKernelRegistry::GetGlobalRegistry()
+                 .LoadKernel<MultiGpuBarrierKernel>(executors_[i])));
+
+        ASSERT_OK(kernel.Launch(
+            ThreadDim(threads_per_block, 1, 1), BlockDim(1, 1, 1),
+            streams_[i].get(), static_cast<int64_t>(i),
+            static_cast<int64_t>(num_devices_), kernel_arg_ptrs, counters[i]));
+      }
+    }
+
+    for (int i = 0; i < num_devices_; ++i) {
+      ASSERT_OK(streams_[i]->BlockHostUntilDone());
+    }
+
+    for (int i = 0; i < num_devices_; ++i) {
+      uint32_t val;
+      ASSERT_OK(streams_[i]->Memcpy(&val, counters[i], sizeof(uint32_t)));
+      ASSERT_OK(streams_[i]->BlockHostUntilDone());
+      EXPECT_EQ(val, 8) << "Counter on device " << i << " failed to increment";
+    }
+
     for (int i = 0; i < num_devices_; ++i) {
       ASSERT_OK_AND_ASSIGN(
-          auto kernel, (GpuKernelRegistry::GetGlobalRegistry()
-                            .LoadKernel<MultiGpuBarrierKernel>(executors_[i])));
-
-      ASSERT_OK(kernel.Launch(ThreadDim(num_devices_, 1, 1), BlockDim(1, 1, 1),
-                              streams_[i].get(), static_cast<int64_t>(i),
-                              static_cast<int64_t>(num_devices_),
-                              kernel_arg_ptrs, counters[i]));
-    }
-  }
-
-  for (int i = 0; i < num_devices_; ++i) {
-    ASSERT_OK(streams_[i]->BlockHostUntilDone());
-  }
-
-  // 5. Verify Counters
-  // After 8 runs, counters should be 8.
-  for (int i = 0; i < num_devices_; ++i) {
-    uint32_t val;
-    ASSERT_OK(streams_[i]->Memcpy(&val, counters[i], sizeof(uint32_t)));
-    ASSERT_OK(streams_[i]->BlockHostUntilDone());
-    EXPECT_EQ(val, 8) << "Counter on device " << i << " failed to increment";
-  }
-
-  // 6. Verify Signal Buffers
-  // Each device's signal buffer is an array of size 'num_devices'.
-  // By the end of step 8, every device should have written the value '8'
-  // into its designated slot on every peer's buffer.
-  for (int i = 0; i < num_devices_; ++i) {
-    // Copy the device's signal buffer back to the host
-    ASSERT_OK_AND_ASSIGN(std::vector<uint32_t> host_buffer,
-                         CopyToHost<uint32_t>(streams_[i].get(),
-                                              signal_buffers[i], num_devices_));
-
-    // Verify every slot contains the final step value (8)
-    for (int j = 0; j < num_devices_; ++j) {
-      EXPECT_EQ(host_buffer[j], 8)
-          << "Signal buffer on Device " << i << " at slot " << j
-          << " (which belongs to Peer " << j << ")"
-          << " has incorrect value.";
+          std::vector<uint32_t> host_buffer,
+          CopyToHost<uint32_t>(streams_[i].get(), signal_buffers[i],
+                               num_devices_));
+      for (int j = 0; j < num_devices_; ++j) {
+        EXPECT_EQ(host_buffer[j], 8)
+            << "Signal buffer on Device " << i << " at slot " << j
+            << " (which belongs to Peer " << j << ")"
+            << " has incorrect value.";
+      }
     }
   }
 }
@@ -236,7 +229,6 @@ TEST_F(MultiGpuBarrierTest, BarrierSynchronizationWithNccl) {
     collective_allocators.push_back(std::move(allocator));
   }
 
-  // Allocate Signal Buffers on each device.
   ASSERT_OK_AND_ASSIGN(
       std::vector<std::unique_ptr<MemoryAllocation>> signal_buffers,
       AllocateBuffers(collective_allocators, num_devices_,
@@ -256,57 +248,61 @@ TEST_F(MultiGpuBarrierTest, BarrierSynchronizationWithNccl) {
           signal_buffer_symmetric_memory,
       CreateSymmetricMemory(exec, comms, signal_buffers, executors_));
 
-  // Allocate Counters on each device.
   ASSERT_OK_AND_ASSIGN(
       std::vector<std::unique_ptr<MemoryAllocation>> counters,
       AllocateBuffers(collective_allocators, num_devices_, sizeof(uint32_t)));
 
-  // Launch Kernel REPEATEDLY 8 times to verify auto-increment.
-  for (int step = 0; step < 8; ++step) {
+  for (int64_t threads_per_block :
+       {int64_t{1}, MultiGpuBarrierWithNcclKernel::kMaxPeers}) {
+    SCOPED_TRACE(absl::StrCat("threads_per_block=", threads_per_block));
+
+    for (int i = 0; i < num_devices_; ++i) {
+      auto sig_addr = signal_buffers[i]->address();
+      ASSERT_OK(
+          streams_[i]->MemZero(&sig_addr, num_devices_ * sizeof(uint32_t)));
+      auto ctr_addr = counters[i]->address();
+      ASSERT_OK(streams_[i]->MemZero(&ctr_addr, sizeof(uint32_t)));
+    }
+
+    for (int step = 0; step < 8; ++step) {
+      for (int i = 0; i < num_devices_; ++i) {
+        ASSERT_OK_AND_ASSIGN(
+            auto kernel,
+            (GpuKernelRegistry::GetGlobalRegistry()
+                 .LoadKernel<MultiGpuBarrierWithNcclKernel>(executors_[i])));
+
+        ASSERT_OK(
+            kernel.Launch(ThreadDim(threads_per_block, 1, 1), BlockDim(1, 1, 1),
+                          streams_[i].get(), static_cast<int64_t>(i),
+                          static_cast<int64_t>(num_devices_),
+                          signal_buffer_symmetric_memory[i].get(),
+                          DeviceAddress<uint32_t>(counters[i]->address())));
+      }
+    }
+
+    for (int i = 0; i < num_devices_; ++i) {
+      ASSERT_OK(streams_[i]->BlockHostUntilDone());
+    }
+
+    for (int i = 0; i < num_devices_; ++i) {
+      uint32_t val;
+      ASSERT_OK(
+          streams_[i]->Memcpy(&val, counters[i]->address(), sizeof(uint32_t)));
+      ASSERT_OK(streams_[i]->BlockHostUntilDone());
+      EXPECT_EQ(val, 8) << "Counter on device " << i << " failed to increment";
+    }
+
     for (int i = 0; i < num_devices_; ++i) {
       ASSERT_OK_AND_ASSIGN(
-          auto kernel,
-          (GpuKernelRegistry::GetGlobalRegistry()
-               .LoadKernel<MultiGpuBarrierWithNcclKernel>(executors_[i])));
-
-      ASSERT_OK(kernel.Launch(ThreadDim(num_devices_, 1, 1), BlockDim(1, 1, 1),
-                              streams_[i].get(), static_cast<int64_t>(i),
-                              static_cast<int64_t>(num_devices_),
-                              signal_buffer_symmetric_memory[i].get(),
-                              DeviceAddress<uint32_t>(counters[i]->address())));
-    }
-  }
-
-  for (int i = 0; i < num_devices_; ++i) {
-    ASSERT_OK(streams_[i]->BlockHostUntilDone());
-  }
-
-  // Verify counters. After 8 runs, counters should be 8.
-  for (int i = 0; i < num_devices_; ++i) {
-    uint32_t val;
-    ASSERT_OK(
-        streams_[i]->Memcpy(&val, counters[i]->address(), sizeof(uint32_t)));
-    ASSERT_OK(streams_[i]->BlockHostUntilDone());
-    EXPECT_EQ(val, 8) << "Counter on device " << i << " failed to increment";
-  }
-
-  // Verify Signal Buffers.
-  // Each device's signal buffer is an array of size 'num_devices'.
-  // By the end of step 8, every device should have written the value '8'
-  // into its designated slot on every peer's buffer.
-  for (int i = 0; i < num_devices_; ++i) {
-    // Copy the device's signal buffer back to the host
-    ASSERT_OK_AND_ASSIGN(
-        std::vector<uint32_t> host_buffer,
-        CopyToHost<uint32_t>(streams_[i].get(), signal_buffers[i]->address(),
-                             num_devices_));
-
-    // Verify every slot contains the final step value (8)
-    for (int j = 0; j < num_devices_; ++j) {
-      EXPECT_EQ(host_buffer[j], 8)
-          << "Signal buffer on Device " << i << " at slot " << j
-          << " (which belongs to Peer " << j << ")"
-          << " has incorrect value.";
+          std::vector<uint32_t> host_buffer,
+          CopyToHost<uint32_t>(streams_[i].get(), signal_buffers[i]->address(),
+                               num_devices_));
+      for (int j = 0; j < num_devices_; ++j) {
+        EXPECT_EQ(host_buffer[j], 8)
+            << "Signal buffer on Device " << i << " at slot " << j
+            << " (which belongs to Peer " << j << ")"
+            << " has incorrect value.";
+      }
     }
   }
 }
diff --git a/third_party/xla/xla/stream_executor/gpu/collective_signal.cu.h b/third_party/xla/xla/stream_executor/gpu/collective_signal.cu.h
index fbb41709e36..423eeb6ddc2 100644
--- a/third_party/xla/xla/stream_executor/gpu/collective_signal.cu.h
+++ b/third_party/xla/xla/stream_executor/gpu/collective_signal.cu.h
@@ -37,16 +37,27 @@ __device__ void WaitSignalFlag(uint32_t* addr, uint32_t expected);
 
 // Generic Barrier across blocks.
 // MaxPeers is templated to allow reuse between different kernels.
+//
+// Each thread of the block strides over the peer ranks starting from
+// `threadIdx.x` with a step of `blockDim.x`. This lets the caller synchronize
+// with `num_ranks` peers using any number of threads (including fewer threads
+// than peers), as long as `num_ranks <= MaxPeers`.
+//
+// Puts and waits run in two separate phases (all puts, then all waits) so
+// that peers observe this rank's signals as early as possible, instead of
+// serializing each put/wait pair.
 template <PlatformType T, int64_t MaxPeers>
 __device__ __forceinline__ void SyncRemoteBlocks(
-    // Use raw pointer with __restrict__ directly here
-    std::array<uint32_t* __restrict__, MaxPeers> signal_pad_ptrs, int64_t rank,
-    int64_t num_ranks, uint32_t signal_value) {
-  if (threadIdx.x < num_ranks) {
-    auto target_rank = threadIdx.x;
+    const std::array<uint32_t* __restrict__, MaxPeers>& signal_pad_ptrs,
+    int64_t rank, int64_t num_ranks, uint32_t signal_value) {
+  for (int64_t target_rank = threadIdx.x; target_rank < num_ranks;
+       target_rank += blockDim.x) {
     PutSignalFlag<T>(
         signal_pad_ptrs[target_rank] + blockIdx.x * num_ranks + rank,
         signal_value);
+  }
+  for (int64_t target_rank = threadIdx.x; target_rank < num_ranks;
+       target_rank += blockDim.x) {
     WaitSignalFlag<T>(
         signal_pad_ptrs[rank] + blockIdx.x * num_ranks + target_rank,
         signal_value);
diff --git a/third_party/xla/xla/stream_executor/gpu/multi_gpu_barrier_kernel.cu.h b/third_party/xla/xla/stream_executor/gpu/multi_gpu_barrier_kernel.cu.h
index 88be205d934..c1409cc6d78 100644
--- a/third_party/xla/xla/stream_executor/gpu/multi_gpu_barrier_kernel.cu.h
+++ b/third_party/xla/xla/stream_executor/gpu/multi_gpu_barrier_kernel.cu.h
@@ -25,30 +25,20 @@ limitations under the License.
 
 namespace stream_executor::gpu {
 
-template <PlatformType PlatformT>
+template <PlatformType PlatformT, int64_t MaxPeers>
 __device__ __forceinline__ void SyncRemoteBlocksAndUpdateCounter(
     int64_t rank, int64_t num_ranks,
-    const std::array<uint32_t* __restrict__, MultiGpuBarrierKernel::kMaxPeers>&
-        signal_buffers,
+    const std::array<uint32_t* __restrict__, MaxPeers>& signal_buffers,
     uint32_t* sync_counter) {
-  // 2. Read State & Pre-Increment
-  // Pre-increment allows to initialize the counter to 0. We sync on 1.
-  // Every rank maintains its own counter, so this read is local and fast.
-  // Since we launch 1 block, a standard load is sufficient, as all threads see
-  // the same value.
   uint32_t signal_value = *sync_counter + 1;
 
-  // 3. Barrier
-  SyncRemoteBlocks<PlatformT, MultiGpuBarrierKernel::kMaxPeers>(
-      signal_buffers, rank, num_ranks, signal_value);
+  SyncRemoteBlocks<PlatformT, MaxPeers>(signal_buffers, rank, num_ranks,
+                                        signal_value);
 
-  // Explicit synchronization.
-  // Adding __syncthreads() ensures that all threads have finished accessing
-  // global memory/peers before we update the state for the next iteration.
+  // CTA-wide barrier: all threads must finish their remote puts/waits before
+  // thread 0 updates sync_counter.
   __syncthreads();
 
-  // 4. Update State
-  // Only the first thread needs to update the counter for the NEXT execution.
   if (threadIdx.x == 0 && blockIdx.x == 0) {
     *sync_counter = signal_value;
   }
@@ -59,17 +49,17 @@ __global__ void MultiGpuBarrierKernelImpl(
     int64_t rank, int64_t num_ranks,
     std::array<void*, MultiGpuBarrierKernel::kMaxPeers> signal_buffers_void,
     uint32_t* sync_counter) {
-  // 1. Cast void* pointers
-  std::array<uint32_t* __restrict__, MultiGpuBarrierKernel::kMaxPeers>
-      signal_buffers;
+  __shared__
+      std::array<uint32_t* __restrict__, MultiGpuBarrierKernel::kMaxPeers>
+          signal_buffers;
 
-#pragma unroll
-  for (int64_t i = 0; i < MultiGpuBarrierKernel::kMaxPeers; ++i) {
+  for (int64_t i = threadIdx.x; i < num_ranks; i += blockDim.x) {
     signal_buffers[i] = reinterpret_cast<uint32_t*>(signal_buffers_void[i]);
   }
+  __syncthreads();
 
-  SyncRemoteBlocksAndUpdateCounter<PlatformT>(rank, num_ranks, signal_buffers,
-                                              sync_counter);
+  SyncRemoteBlocksAndUpdateCounter<PlatformT, MultiGpuBarrierKernel::kMaxPeers>(
+      rank, num_ranks, signal_buffers, sync_counter);
 }
 
 }  // namespace stream_executor::gpu
diff --git a/third_party/xla/xla/stream_executor/gpu/multi_gpu_barrier_kernel.h b/third_party/xla/xla/stream_executor/gpu/multi_gpu_barrier_kernel.h
index 71c494b980c..8635eb641a2 100644
--- a/third_party/xla/xla/stream_executor/gpu/multi_gpu_barrier_kernel.h
+++ b/third_party/xla/xla/stream_executor/gpu/multi_gpu_barrier_kernel.h
@@ -47,8 +47,8 @@ namespace stream_executor::gpu {
 //     Zero-initialized memory correctly sets up the wait condition for the
 //     first barrier launch.
 struct MultiGpuBarrierKernel {
-  // Maximum number of peers supported by the barrier.
-  // Can be extended to support larger GPU clusters in the future.
+  // Maximum number of peers for the P2P (non-NCCL) barrier. Sized to the
+  // kernel argument array; use MultiGpuBarrierWithNcclKernel for >32 peers.
   static constexpr int64_t kMaxPeers = 32;
 
   using KernelType =
@@ -64,9 +64,9 @@ struct MultiGpuBarrierKernel {
 //  pointer storage symmetric memory at rank's location. This replaces the need
 //  to do host-side rendezvous to exchange pointers.
 struct MultiGpuBarrierWithNcclKernel {
-  // Maximum number of peers supported by the barrier.
-  // Can be extended to support larger GPU clusters in the future.
-  static constexpr int64_t kMaxPeers = 32;
+  // Maximum number of peers for the NCCL barrier. No kernel argument array,
+  // so this can be larger than MultiGpuBarrierKernel::kMaxPeers.
+  static constexpr int64_t kMaxPeers = 128;
 
   using KernelType = stream_executor::TypedKernel<
       int64_t, int64_t, xla::SymmetricMemory*,  // signal buffers

From 10d9e4089d5c7125f6a0b1217998e84536961d93 Mon Sep 17 00:00:00 2001
From: Levon Ter-Grigoryan <patrios@google.com>
Date: Wed, 14 Oct 2026 04:01:30 -0700
Subject: [PATCH 0759/1000] Add `COLLECTIVE_KERNEL_ALL_TO_ALL` to
 `DebugOptions::CollectiveKernelEnabled` and update flag parsing in
 `debug_options_flags.cc` to support
 `--xla_gpu_experimental_use_collective_kernels=all_to_all`.

PiperOrigin-RevId: 999268804
---
 third_party/xla/xla/debug_options_flags.cc      | 2 +-
 third_party/xla/xla/debug_options_flags_test.cc | 6 ++++++
 third_party/xla/xla/xla.proto                   | 1 +
 3 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index d709e4bf7a8..cab2e9bdae0 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -3384,7 +3384,7 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
           debug_options->xla_gpu_experimental_use_collective_kernels()),
       "Experimental: comma-separated filter of collective ops that should use "
       "custom kernels (e.g. Triton one-shot / two-shot) instead of NCCL. "
-      "Accepted values: ALL_REDUCE, ALL_GATHER, REDUCE_SCATTER, "
+      "Accepted values: ALL_REDUCE, ALL_GATHER, REDUCE_SCATTER, ALL_TO_ALL, "
       "ALL_COLLECTIVES (case-insensitive; the "
       "COLLECTIVE_KERNEL_ prefix may be omitted). Supports +/- "
       "incremental modifiers (e.g. +ALL_REDUCE,-ALL_GATHER). The deprecated "
diff --git a/third_party/xla/xla/debug_options_flags_test.cc b/third_party/xla/xla/debug_options_flags_test.cc
index 68064de02b6..86eba8bf450 100644
--- a/third_party/xla/xla/debug_options_flags_test.cc
+++ b/third_party/xla/xla/debug_options_flags_test.cc
@@ -244,6 +244,12 @@ TEST(DebugOptions, CollectiveKernelsFlagUnderscoreAllGather) {
               ElementsAre(DebugOptions::COLLECTIVE_KERNEL_ALL_GATHER));
 }
 
+TEST(DebugOptions, CollectiveKernelsFlagAllToAll) {
+  DebugOptions opts = ParseCollectiveKernelsFlag("ALL_TO_ALL");
+  EXPECT_THAT(opts.xla_gpu_experimental_use_collective_kernels(),
+              ElementsAre(DebugOptions::COLLECTIVE_KERNEL_ALL_TO_ALL));
+}
+
 TEST(DebugOptions, CollectiveKernelsFlagBothHyphenated) {
   DebugOptions opts = ParseCollectiveKernelsFlag("ALL_REDUCE,ALL_GATHER");
   EXPECT_THAT(opts.xla_gpu_experimental_use_collective_kernels(),
diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index 72459561fae..1d0936e2bb6 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -112,6 +112,7 @@ message DebugOptions {
     COLLECTIVE_KERNEL_ALL_REDUCE = 1;
     COLLECTIVE_KERNEL_ALL_GATHER = 2;
     COLLECTIVE_KERNEL_REDUCE_SCATTER = 3;
+    COLLECTIVE_KERNEL_ALL_TO_ALL = 4;
     COLLECTIVE_KERNEL_ALL_COLLECTIVES = 8;
   }
 

From 7bef9ab9883dab63e365d1d232323570bb1419c8 Mon Sep 17 00:00:00 2001
From: Oleg Shyshkov <shyshkov@google.com>
Date: Wed, 14 Oct 2026 04:07:36 -0700
Subject: [PATCH 0760/1000] [XLA:GPU] Remove RunAllReduceKernel and its
 associated unit tests.

PiperOrigin-RevId: 999271905
---
 .../xla/xla/backends/gpu/runtime/BUILD        |  43 --
 .../xla/backends/gpu/runtime/all_reduce.cc    | 150 -------
 .../xla/xla/backends/gpu/runtime/all_reduce.h |  49 --
 .../backends/gpu/runtime/all_reduce_test.cc   | 425 ------------------
 4 files changed, 667 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index c460bcfebcc..3ccc68e99d3 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -4221,7 +4221,6 @@ cc_library(
         "//xla:util",
         "//xla:xla_data_proto_cc",
         "//xla/backends/gpu/transforms/collectives:collective_ops_utils",
-        "//xla/core/collectives:rank_id",
         "//xla/core/collectives:reduction_kind",
         "//xla/hlo/ir:hlo",
         "//xla/service:collective_ops_utils",
@@ -4230,15 +4229,9 @@ cc_library(
         "//xla/service/gpu:gpu_constants",
         "//xla/service/gpu:ir_emission_utils",
         "//xla/service/gpu:launch_dimensions",
-        "//xla/stream_executor:device_address",
         "//xla/stream_executor:device_description",
-        "//xla/stream_executor:stream",
         "//xla/stream_executor/gpu:all_reduce_kernel",
-        "//xla/stream_executor/gpu:collective_kernel_metadata",
-        "//xla/stream_executor/gpu:gpu_kernel_registry",
-        "//xla/tsl/util:safe_reinterpret_cast",
         "@com_google_absl//absl/algorithm:container",
-        "@com_google_absl//absl/base",
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
@@ -4319,56 +4312,20 @@ xla_test(
     srcs = ["all_reduce_test.cc"],
     backend_tags = {
         "gpu": [
-            "multi_gpu",
-            "xla_amdgpu_any",
             "no_oss",
         ],
-        "gb200": ["broken"],  # TODO(b/505963049): re-enable when fixed.
-        "gb300": ["broken"],  # TODO(b/505963049): re-enable when fixed.
     },
     backends = ["gpu"],
     disabled_backends = [],
     deps = [
-        ":all_reduce",
-        "//xla:array",
         "//xla:literal",
         "//xla:literal_util",
-        "//xla:shape_util",
-        "//xla:status_macros",
-        "//xla:types",
-        "//xla:util",
-        "//xla:xla_data_proto_cc",
-        "//xla/core/collectives:rank_id",
-        "//xla/core/collectives:reduction_kind",
         "//xla/hlo/ir:hlo",
-        "//xla/service/gpu:gpu_constants",
-        "//xla/service/gpu:launch_dimensions",
-        "//xla/stream_executor:device_address",
-        "//xla/stream_executor:memory_space",
-        "//xla/stream_executor:platform",
-        "//xla/stream_executor:platform_manager",
-        "//xla/stream_executor:stream",
-        "//xla/stream_executor:stream_executor_h",
-        "//xla/stream_executor/gpu:all_reduce_kernel",
-        "//xla/stream_executor/gpu:collective_kernel_metadata",
-        "//xla/stream_executor/gpu:gpu_executor_header",
-        "//xla/stream_executor/gpu:gpu_init",
-        "//xla/stream_executor/gpu:multicast_memory",
-        "//xla/stream_executor/host:host_platform",
         "//xla/tests:hlo_pjrt_test_base",
-        "//xla/tests:literal_test_util",
-        "//xla/tsl/platform:errors",
         "//xla/tsl/platform:test",
         "//xla/tsl/platform:test_main",
-        "@com_google_absl//absl/log",
-        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:status_matchers",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/strings:str_format",
-        "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest",
     ],
 )
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc b/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc
index a653b941d38..9d614471938 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc
@@ -21,7 +21,6 @@ limitations under the License.
 #include <vector>
 
 #include "absl/algorithm/container.h"
-#include "absl/base/casts.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
@@ -34,7 +33,6 @@ limitations under the License.
 #include "llvm/ADT/bit.h"
 #include "xla/backends/gpu/runtime/collective_params.h"
 #include "xla/backends/gpu/transforms/collectives/collective_ops_utils.h"
-#include "xla/core/collectives/rank_id.h"
 #include "xla/core/collectives/reduction_kind.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
@@ -48,13 +46,8 @@ limitations under the License.
 #include "xla/service/gpu/launch_dimensions.h"
 #include "xla/service/gpu_topology.h"
 #include "xla/shape_util.h"
-#include "xla/stream_executor/device_address.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/gpu/all_reduce_kernel.h"
-#include "xla/stream_executor/gpu/collective_kernel_metadata.h"
-#include "xla/stream_executor/gpu/gpu_kernel_registry.h"
-#include "xla/stream_executor/stream.h"
-#include "xla/tsl/util/safe_reinterpret_cast.h"
 #include "xla/types.h"
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
@@ -69,102 +62,12 @@ static constexpr int64_t kMB = kKB * 1024;
 static constexpr int64_t kMaxOneShotAllReduceSizeBytes = 2 * kMB;
 static constexpr int64_t kMaxTwoShotAllReduceSizeBytes = 32 * kMB;
 
-template <typename T, ReductionKind kReductionKindV>
-class TagRegistry {
- private:
-  template <AllReduceStrategy kAllReduceStrategyV>
-  struct Impl {
-    using ElementType = T;
-    static constexpr ReductionKind kReductionKind = kReductionKindV;
-    static constexpr AllReduceStrategy kAllReduceStrategy = kAllReduceStrategyV;
-  };
-
- public:
-  static constexpr auto kOneShot = Impl<AllReduceStrategy::kOneShot>{};
-  static constexpr auto kTwoShot = Impl<AllReduceStrategy::kTwoShot>{};
-  static constexpr auto kMultimem = Impl<AllReduceStrategy::kMultimem>{};
-};
-
-// Static set of supported kernel tags.
-static constexpr auto kAddF32Tags = TagRegistry<float, ReductionKind::SUM>{};
-static constexpr auto kAddBF16Tags =
-    TagRegistry<bfloat16, ReductionKind::SUM>{};
-static constexpr auto kOrPredTags = TagRegistry<bool, ReductionKind::MAX>{};
 // Heuristic maxima after some benchmarking.
 // It's nicer to have this as a power of 2 because we consume blocks in powers
 // of 2 from a higher dimension to a lower dimension.
 static constexpr int64_t kMaxBlocksPerGrid = 32;
 static constexpr uint64_t kMaxThreadsPerBlock = 512;
 
-template <typename TagType>
-absl::Status LaunchTypedKernel(
-    TagType, se::Stream* stream, const LaunchDimensions& launch_dimensions,
-    se::DeviceAddressBase symmetric_input_buffer,
-    se::DeviceAddressBase local_input_buffer,
-    se::DeviceAddressBase output_buffer, int64_t rank, int64_t num_ranks,
-    int64_t num_elements, se::DeviceAddressBase symmetric_signal_buffer,
-    uint32_t signal_value, se::DeviceAddressBase metadata) {
-  using ElementType = typename TagType::ElementType;
-  static constexpr bool kIsTwoShot =
-      TagType::kAllReduceStrategy == AllReduceStrategy::kTwoShot;
-
-  ABSL_ASSIGN_OR_RETURN(
-      auto kernel,
-      (se::gpu::GpuKernelRegistry::GetGlobalRegistry()
-           .LoadKernel<
-               se::gpu::AllReduceKernel<ElementType, TagType::kReductionKind,
-                                        TagType::kAllReduceStrategy>>(
-               stream->parent())));
-
-  se::gpu::AllReduceKernelParams<ElementType> params{};
-  params.input_buffer =
-      tsl::safe_reinterpret_cast<ElementType*>(local_input_buffer.opaque());
-  params.output_buffer =
-      tsl::safe_reinterpret_cast<ElementType*>(output_buffer.opaque());
-  params.symmetric_input_ptrs =
-      tsl::safe_reinterpret_cast<ElementType*>(symmetric_input_buffer.opaque());
-  params.symmetric_signal_ptrs =
-      tsl::safe_reinterpret_cast<uint32_t*>(symmetric_signal_buffer.opaque());
-  params.rank = rank;
-  params.num_ranks = num_ranks;
-  params.num_elements = num_elements;
-  params.num_elements_per_rank = num_elements / (kIsTwoShot ? num_ranks : 1);
-  // NB: num_elements_per_block can be bigger or smaller than blockDim.x.
-  // If its smaller, then the block stride loop will run just once.
-  // If its bigger, then the block stride loop will run multiple times.
-  params.num_elements_per_block = RoundUpTo(
-      CeilOfRatio(params.num_elements_per_rank,
-                  absl::implicit_cast<int64_t>(launch_dimensions.num_blocks())),
-      se::gpu::kNumElementsPerThread);
-  params.rank_offset =
-      kIsTwoShot ? params.rank * params.num_elements_per_rank : 0;
-  for (int i = 0; i < params.num_ranks; ++i) {
-    params.rotated_ranks[i] = (i + rank) % params.num_ranks;
-  }
-  params.signal_value = signal_value;
-  params.metadata =
-      tsl::safe_reinterpret_cast<CollectiveKernelMetadata*>(metadata.opaque());
-
-  VLOG(3) << "Launching all-reduce kernel with params: " << "strategy: "
-          << absl::StrFormat("%v", TagType::kAllReduceStrategy)
-          << ", rank: " << params.rank << ", num_ranks: " << params.num_ranks
-          << ", num_elements: " << params.num_elements
-          << ", num_elements_per_rank: " << params.num_elements_per_rank
-          << ", num_elements_per_block: " << params.num_elements_per_block
-          << ", num_threads_per_block: "
-          << launch_dimensions.num_threads_per_block()
-          << ", num_blocks_per_grid: " << launch_dimensions.num_blocks()
-          << ", rank_offset: {" << params.rank_offset << ", rotated_ranks: "
-          << absl::StrJoin(
-                 absl::MakeSpan(params.rotated_ranks.data(), params.num_ranks),
-                 ", ")
-          << "}, signal_value: " << params.signal_value;
-
-  return kernel.Launch(launch_dimensions.thread_counts_per_block(),
-                       launch_dimensions.block_counts(), stream,
-                       std::move(params));
-}
-
 bool IsElementReductionSupported(PrimitiveType element_type,
                                  ReductionKind reduction_kind) {
   absl::Span<const HloOpcode> supported_ops =
@@ -383,59 +286,6 @@ absl::StatusOr<AllReduceInfo> BuildAllReduceInfo(
   };
 }
 
-absl::Status RunAllReduceKernel(
-    se::Stream* stream,                             //
-    const LaunchDimensions& launch_dimensions,      //
-    PrimitiveType element_type,                     //
-    ReductionKind reduction_kind,                   //
-    AllReduceStrategy all_reduce_strategy,          //
-    se::DeviceAddressBase symmetric_input_buffer,   //
-    se::DeviceAddressBase local_input_buffer,       //
-    se::DeviceAddressBase output_buffer,            //
-    RankId rank,                                    //
-    int64_t num_ranks,                              //
-    int64_t num_elements,                           //
-    se::DeviceAddressBase symmetric_signal_buffer,  //
-    uint32_t signal_value,                          //
-    se::DeviceAddressBase metadata) {
-  ABSL_RETURN_IF_ERROR(IsAllReduceKernelSupported(num_ranks, num_elements,
-                                             element_type, reduction_kind,
-                                             all_reduce_strategy));
-  const auto launch_kernel_impl = [&](auto tag) -> absl::Status {
-    return LaunchTypedKernel(
-        tag, stream, launch_dimensions, symmetric_input_buffer,
-        local_input_buffer, output_buffer, rank.value(), num_ranks,
-        num_elements, symmetric_signal_buffer, signal_value, metadata);
-  };
-  const auto launch_kernel = [&](auto tag_registry,
-                                 AllReduceStrategy strategy) -> absl::Status {
-    switch (strategy) {
-      case AllReduceStrategy::kOneShot:
-        return launch_kernel_impl(tag_registry.kOneShot);
-      case AllReduceStrategy::kTwoShot:
-        return launch_kernel_impl(tag_registry.kTwoShot);
-      case AllReduceStrategy::kMultimem:
-        return launch_kernel_impl(tag_registry.kMultimem);
-    }
-  };
-
-  if (element_type == F32 && reduction_kind == ReductionKind::SUM) {
-    return launch_kernel(kAddF32Tags, all_reduce_strategy);
-  }
-
-  if (element_type == BF16 && reduction_kind == ReductionKind::SUM) {
-    return launch_kernel(kAddBF16Tags, all_reduce_strategy);
-  }
-
-  if (element_type == PRED && reduction_kind == ReductionKind::MAX) {
-    return launch_kernel(kOrPredTags, all_reduce_strategy);
-  }
-
-  return absl::InvalidArgumentError(
-      "Unsupported AllReduce kernel. This line should never be reached if the "
-      "result of `IsAllReduceKernelSupported` is correct.");
-}
-
 absl::StatusOr<CollectiveKernelSpec> CreateAllReduceKernelSpec(
     const HloInstruction* instr, const LaunchDimensions& launch_dimensions,
     SymmetricMemoryType scratch_memory_type) {
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_reduce.h b/third_party/xla/xla/backends/gpu/runtime/all_reduce.h
index a3f3e33eaf5..77a99efe838 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_reduce.h
+++ b/third_party/xla/xla/backends/gpu/runtime/all_reduce.h
@@ -27,16 +27,13 @@ limitations under the License.
 #include "absl/types/span.h"
 #include "llvm/ADT/STLExtras.h"
 #include "xla/backends/gpu/runtime/collective_params.h"
-#include "xla/core/collectives/rank_id.h"
 #include "xla/core/collectives/reduction_kind.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/service/gpu/launch_dimensions.h"
-#include "xla/stream_executor/device_address.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/gpu/all_reduce_kernel.h"
-#include "xla/stream_executor/stream.h"
 #include "xla/types.h"  // IWYU pragma: keep
 #include "xla/xla_data.pb.h"
 
@@ -128,52 +125,6 @@ absl::StatusOr<AllReduceInfo> BuildAllReduceInfo(
     const HloAllReduceInstruction* all_reduce,
     const DeviceAssignment* device_assignment);
 
-// Performs element-wise addition of all input buffers and stores the result in
-// the output buffer.
-// The kernel is intended to be used for all-reduce operations in environment
-// where direct peer memory access is available. Input buffers can point to
-// memory on different devices. The caller is responsible to gather pointers
-// from different devices.
-//
-// The kernel copies data from local input buffer to remote input buffer of the
-// current rank at the start of the kernel.
-//
-// The kernel performs synchronization across devices at the start and the end
-// of the kernel. The synchronization happens between blocks with the same id.
-//
-// Input arguments:
-//  - remove_input_buffers: A list of buffers with inputs on other devices.
-//    The data in the buffers maybe not be initialized until blocks on different
-//    devices are synchronized.
-//  - local_input_buffer: The buffer with local input. Can be the same as
-//    the output buffer.
-//  - output_buffer: The buffer to store the result.
-//  - rank: Identifier of the device.
-//  - num_ranks: The number of devices participating in the operation.
-//  - num_elements: The number of elements in each buffer.
-//  - signal_flags_buffers: A list of buffers with signal flags that are used to
-//    synchronize blocks on different devices. The size of each signal buffer
-//    should be equal to the `num_ranks * num_blocks`.
-//  - signal_value: The value that is written to the signal flags. Should be
-//    different for different invocations of the kernel with the same signal
-//    buffer.
-absl::Status RunAllReduceKernel(
-    se::Stream* stream,                              //
-    const LaunchDimensions& launch_dimensions,       //
-    PrimitiveType element_type,                      //
-    ReductionKind reduction_kind,                    //
-    se::gpu::AllReduceStrategy all_reduce_strategy,  //
-    se::DeviceAddressBase symmetric_input_buffer,    //
-    se::DeviceAddressBase local_input_buffer,        //
-    se::DeviceAddressBase output_buffer,             //
-    RankId rank,                                     //
-    int64_t num_ranks,                               //
-    int64_t num_elements,                            //
-    se::DeviceAddressBase symmetric_signal_buffer,   //
-    uint32_t signal_value,                           //
-    se::DeviceAddressBase metadata                   //
-);
-
 }  // namespace xla::gpu
 
 #endif  // XLA_BACKENDS_GPU_RUNTIME_ALL_REDUCE_H_
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_reduce_test.cc b/third_party/xla/xla/backends/gpu/runtime/all_reduce_test.cc
index 898ba919812..cd46386ea65 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_reduce_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/all_reduce_test.cc
@@ -13,451 +13,26 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "xla/backends/gpu/runtime/all_reduce.h"
-
-#include <algorithm>
-#include <cstddef>
-#include <cstdint>
 #include <memory>
-#include <string>
-#include <tuple>
 #include <utility>
 #include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
-#include "absl/log/check.h"
-#include "absl/log/log.h"
 #include "absl/status/status.h"
-#include "absl/status/status_macros.h"
 #include "absl/status/status_matchers.h"
-#include "absl/status/statusor.h"
-#include "absl/strings/str_cat.h"
-#include "absl/strings/str_format.h"
-#include "absl/strings/str_join.h"
-#include "absl/types/span.h"
-#include "xla/array.h"
-#include "xla/core/collectives/rank_id.h"
-#include "xla/core/collectives/reduction_kind.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/literal.h"
 #include "xla/literal_util.h"
-#include "xla/primitive_util.h"
-#include "xla/service/gpu/gpu_constants.h"
-#include "xla/service/gpu/launch_dimensions.h"
-#include "xla/status_macros.h"
-#include "xla/stream_executor/device_address.h"
-#include "xla/stream_executor/gpu/all_reduce_kernel.h"
-#include "xla/stream_executor/gpu/collective_kernel_metadata.h"
-#include "xla/stream_executor/gpu/gpu_executor.h"
-#include "xla/stream_executor/gpu/gpu_init.h"
-#include "xla/stream_executor/gpu/multicast_memory.h"
-#include "xla/stream_executor/memory_space.h"
-#include "xla/stream_executor/platform.h"
-#include "xla/stream_executor/platform_manager.h"
-#include "xla/stream_executor/stream.h"
-#include "xla/stream_executor/stream_executor.h"
 #include "xla/tests/hlo_pjrt_test_base.h"
-#include "xla/tests/literal_test_util.h"
-#include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/test.h"
-#include "xla/types.h"
-#include "xla/util.h"
-#include "xla/xla_data.pb.h"
 
 namespace xla::gpu {
 namespace {
 
 using ::absl_testing::StatusIs;
-using ::stream_executor::gpu::AllReduceStrategy;
 using ::testing::HasSubstr;
 
-se::StreamExecutor* GetGpuExecutor(int64_t device_ordinal) {
-  auto* platform =
-      se::PlatformManager::PlatformWithName(se::GpuPlatformName()).value();
-  return platform->ExecutorForDevice(device_ordinal).value();
-}
-
-struct TestParams {
-  AllReduceStrategy all_reduce_strategy;
-  int64_t num_elements;
-};
-
-struct PtrFormatter {
-  void operator()(std::string* out, const void* ptr) const {
-    absl::StrAppend(out, absl::StrFormat("%p", ptr));
-  }
-};
-
-class AllReduceKernelTest : public ::testing::Test,
-                            public ::testing::WithParamInterface<TestParams> {
- public:
-  AllReduceKernelTest() : params_(GetParam()) {}
-
-  template <typename T>
-  absl::StatusOr<std::vector<Array<T>>> RunKernel(
-      const std::vector<se::StreamExecutor*>& executors,
-      const std::vector<Array<T>>& input_data, ReductionKind reduction_kind) {
-    const int64_t num_ranks = input_data.size();
-    const LaunchDimensions launch_dimensions = AllReduceLaunchDimensions(
-        input_data[0].num_elements(), num_ranks, params_.all_reduce_strategy,
-        executors[0]->GetDeviceDescription());
-
-    int64_t num_elements = input_data[0].num_elements();
-
-    ABSL_RETURN_IF_ERROR(executors[0]->EnablePeerAccessTo(executors[1]));
-    ABSL_RETURN_IF_ERROR(executors[1]->EnablePeerAccessTo(executors[0]));
-
-    std::unique_ptr<se::gpu::MulticastMemory> multicast_memory;
-    if (params_.all_reduce_strategy == AllReduceStrategy::kMultimem) {
-      ABSL_ASSIGN_OR_RETURN(
-          multicast_memory,
-          dynamic_cast<se::gpu::GpuExecutor*>(executors[0])
-              ->CreateMulticastMemory(num_elements * sizeof(T), num_ranks));
-
-      for (int i = 0; i < num_ranks; ++i) {
-        ABSL_RETURN_IF_ERROR(multicast_memory->SubscribeDevice(i));
-      }
-    }
-
-    std::vector<std::unique_ptr<se::Stream>> streams;
-    std::vector<se::DeviceAddressBase> allocated_buffers;
-    std::vector<se::DeviceAddressBase> input_buffers;
-    std::vector<se::DeviceAddressBase> output_buffers;
-    std::vector<se::DeviceAddressBase> symmetric_input_buffers;
-    std::vector<se::DeviceAddressBase> signal_flags_buffers;
-
-    uint64_t input_size = num_elements * sizeof(T);
-    uint64_t aligned_input_size =
-        xla::RoundUpTo<uint64_t>(input_size, kXlaAllocatedBufferAlignBytes);
-    uint64_t signal_size =
-        num_ranks * launch_dimensions.num_blocks() * sizeof(int32_t);
-    uint64_t aligned_signal_size =
-        xla::RoundUpTo<uint64_t>(signal_size, kXlaAllocatedBufferAlignBytes);
-    for (int i = 0; i < num_ranks; ++i) {
-      auto* executor = executors[i];
-      streams.push_back(executor->CreateStream().value());
-
-      uint64_t total_size =
-          /*input_buffer_size=*/aligned_input_size +
-          /*symmetric_input_buffer_size=*/aligned_input_size +
-          /*output_buffer_size=*/aligned_input_size +
-          /*signal_buffer_size=*/aligned_signal_size;
-      allocated_buffers.emplace_back(executor->AllocateArray<T>(
-          total_size,
-          static_cast<int64_t>(stream_executor::MemorySpace::kCollective)));
-      input_buffers.emplace_back(
-          allocated_buffers[i].GetByteSlice(0, aligned_input_size));
-      TF_RET_CHECK(!input_buffers[i].is_null());
-
-      symmetric_input_buffers.emplace_back(allocated_buffers[i].GetByteSlice(
-          aligned_input_size, aligned_input_size));
-      TF_RET_CHECK(!symmetric_input_buffers[i].is_null());
-
-      output_buffers.emplace_back(allocated_buffers[i].GetByteSlice(
-          2 * aligned_input_size, aligned_input_size));
-      TF_RET_CHECK(!output_buffers[i].is_null());
-      ABSL_RETURN_IF_ERROR(
-          streams[i]->MemZero(&output_buffers[i], aligned_input_size));
-
-      signal_flags_buffers.emplace_back(allocated_buffers[i].GetByteSlice(
-          3 * aligned_input_size, aligned_signal_size));
-      TF_RET_CHECK(!signal_flags_buffers[i].is_null());
-      ABSL_RETURN_IF_ERROR(
-          streams[i]->MemZero(&signal_flags_buffers[i], aligned_signal_size));
-      ABSL_RETURN_IF_ERROR(streams[i]->Memcpy(&input_buffers[i],
-                                         input_data[i].data(), input_size));
-      XLA_VLOG_DEVICE(1, i)
-          << "Allocated buffer: " << allocated_buffers[i].opaque()
-          << ", Input buffer: " << input_buffers[i].opaque()
-          << ", Symmetric input buffer: " << symmetric_input_buffers[i].opaque()
-          << ", Output buffer: " << output_buffers[i].opaque()
-          << ", Signal buffer: " << signal_flags_buffers[i].opaque();
-    }
-
-    std::vector<se::DeviceAddressBase> metadata_buffers;
-    for (int i = 0; i < num_ranks; ++i) {
-      CollectiveKernelMetadata metadata;
-      metadata.rank = i;
-      std::vector<void*> param_to_peers_ptrs;
-
-      if (params_.all_reduce_strategy == AllReduceStrategy::kMultimem) {
-        // Multimem also need to have an output buffer.
-        constexpr int kNumPeerParameters = 3;
-        // Multimem needs to address input and output buffers on the peer
-        // devices. Also an offset between the root exchanged pointer and
-        // the multimem address space should be the same.
-        for (int buffer_id = 0; buffer_id < kNumPeerParameters; ++buffer_id) {
-          for (int rank = 0; rank < num_ranks; ++rank) {
-            param_to_peers_ptrs.push_back(allocated_buffers[rank].opaque());
-          }
-        }
-
-        se::gpu::GpuExecutor* gpu_executor =
-            dynamic_cast<se::gpu::GpuExecutor*>(executors[i]);
-        TF_RET_CHECK(gpu_executor != nullptr);
-        ABSL_ASSIGN_OR_RETURN(
-            void* mapped_memory,
-            multicast_memory->MapMemory(allocated_buffers[i], gpu_executor));
-        std::vector<void*> param_to_multimem_addresses =
-            std::vector<void*>(kNumPeerParameters, mapped_memory);
-
-        const size_t param_to_peers_size =
-            sizeof(void*) * param_to_peers_ptrs.size();
-        const size_t param_to_multimem_addresses_byte_size =
-            sizeof(void*) * param_to_multimem_addresses.size();
-        // First map from parameter to peer ptrs and then metadata.
-        metadata_buffers.emplace_back(executors[i]->AllocateArray<uint64_t>(
-            sizeof(CollectiveKernelMetadata) + param_to_peers_size +
-            param_to_multimem_addresses_byte_size));
-
-        se::DeviceAddressBase param_to_multimem_addresses_buffer =
-            metadata_buffers[i].GetByteSlice(
-                sizeof(CollectiveKernelMetadata) + param_to_peers_size,
-                param_to_multimem_addresses_byte_size);
-        metadata.param_to_multimem_addresses = reinterpret_cast<void**>(
-            param_to_multimem_addresses_buffer.opaque());
-        ABSL_RETURN_IF_ERROR(
-            streams[i]->Memcpy(&param_to_multimem_addresses_buffer,
-                               param_to_multimem_addresses.data(),
-                               param_to_multimem_addresses_byte_size));
-        XLA_VLOG_DEVICE(1, i)
-            << "Constructed device state {"
-            << " metadata rank: " << metadata.rank << ", param_to_peers: ("
-            << absl::StrJoin(param_to_peers_ptrs, ", ", PtrFormatter{})
-            << "), multimem_addresses: ("
-            << absl::StrJoin(param_to_multimem_addresses, ", ", PtrFormatter{})
-            << ")}";
-      } else {
-        for (const se::DeviceAddressBase& input_buffer : input_buffers) {
-          param_to_peers_ptrs.push_back(input_buffer.opaque());
-        }
-        for (const se::DeviceAddressBase& signal_flags_buffer :
-             signal_flags_buffers) {
-          param_to_peers_ptrs.push_back(signal_flags_buffer.opaque());
-        }
-        metadata_buffers.emplace_back(executors[i]->AllocateArray<uint64_t>(
-            sizeof(CollectiveKernelMetadata) +
-            param_to_peers_ptrs.size() * sizeof(void*)));
-
-        XLA_VLOG_DEVICE(1, executors[i]->device_ordinal())
-            << "Constructed device state {"
-            << " metadata rank: " << metadata.rank << ", param_to_peers: ("
-            << absl::StrJoin(param_to_peers_ptrs, ", ", PtrFormatter{}) << ")}";
-      }
-
-      const size_t param_to_peers_size_bytes =
-          param_to_peers_ptrs.size() * sizeof(void*);
-      se::DeviceAddressBase param_to_peers_ptrs_buffer =
-          metadata_buffers[i].GetByteSlice(sizeof(CollectiveKernelMetadata),
-                                           param_to_peers_size_bytes);
-      metadata.param_to_peers =
-          reinterpret_cast<void**>(param_to_peers_ptrs_buffer.opaque());
-      ABSL_RETURN_IF_ERROR(streams[i]->Memcpy(&param_to_peers_ptrs_buffer,
-                                         param_to_peers_ptrs.data(),
-                                         param_to_peers_size_bytes));
-      ABSL_RETURN_IF_ERROR(streams[i]->Memcpy(&metadata_buffers[i], &metadata,
-                                         sizeof(CollectiveKernelMetadata)));
-    }
-
-    for (int i = 0; i < num_ranks; ++i) {
-      ABSL_RETURN_IF_ERROR(streams[i]->BlockHostUntilDone());
-    }
-
-    for (int i = 0; i < num_ranks; ++i) {
-      auto active_context = executors[i]->Activate();
-      ABSL_RETURN_IF_ERROR(RunAllReduceKernel(
-          streams[i].get(), launch_dimensions,
-          primitive_util::NativeToPrimitiveType<T>(),
-          /*reduction_kind=*/reduction_kind,
-          /*all_reduce_strategy=*/params_.all_reduce_strategy,
-          /*symmetric_input_buffer=*/symmetric_input_buffers[i],
-          // Memory is aliased for both input and output (similar to what nccl
-          // would do).
-          /*local_input_buffer=*/input_buffers[i],
-          /*output_buffer=*/output_buffers[i],
-          /*rank=*/RankId(i), /*num_ranks=*/num_ranks,
-          /*num_elements=*/num_elements,
-          /*symmetric_signal_buffer=*/signal_flags_buffers[i],
-          /*signal_value=*/1,
-          /*metadata=*/metadata_buffers[i]));
-    }
-
-    for (int i = 0; i < num_ranks; ++i) {
-      ABSL_RETURN_IF_ERROR(streams[i]->BlockHostUntilDone());
-    }
-
-    std::vector<Array<T>> results;
-    for (int i = 0; i < num_ranks; ++i) {
-      Array<T> output_results({num_elements});
-      ABSL_RETURN_IF_ERROR(streams[i]->Memcpy(
-          output_results.data(), output_buffers[i], num_elements * sizeof(T)));
-
-      results.push_back(std::move(output_results));
-    }
-
-    return results;
-  }
-
-  int64_t num_elements() const { return params_.num_elements; }
-
-  AllReduceStrategy strategy() const { return params_.all_reduce_strategy; }
-
- private:
-  TestParams params_;
-};
-
-TEST_P(AllReduceKernelTest, KernelTestAddF32) {
-  constexpr int64_t kNumRanks = 2;
-
-  std::vector<se::StreamExecutor*> executors = {GetGpuExecutor(0),
-                                                GetGpuExecutor(1)};
-  if (strategy() == AllReduceStrategy::kMultimem &&
-      !dynamic_cast<se::gpu::GpuExecutor*>(executors[0])
-           ->is_multicast_supported()) {
-    GTEST_SKIP() << "Multimem not supported on this device.";
-  }
-
-  if (!executors[0]->CanEnablePeerAccessTo(executors[1])) {
-    GTEST_SKIP() << "Test requires direct peer memory access between devices.";
-  }
-
-  Array<float> expected_output({num_elements()});
-  std::vector<Array<float>> inputs;
-
-  for (int i = 0; i < kNumRanks; ++i) {
-    Array<float> input_data({num_elements()});
-    input_data.FillRandom(0.0f, 10.0f, /*seed=*/i);
-
-    expected_output.Each([&](absl::Span<const int64_t> indices, float* val) {
-      *val += input_data(indices);
-    });
-
-    inputs.push_back(std::move(input_data));
-  }
-
-  ASSERT_OK_AND_ASSIGN(auto results,
-                       RunKernel<float>(executors, inputs, ReductionKind::SUM));
-
-  const Literal expected_output_literal =
-      LiteralUtil::CreateFromArray<float>(expected_output);
-  for (int i = 0; i < kNumRanks; ++i) {
-    const Literal actual_output_literal =
-        LiteralUtil::CreateFromArray<float>(results[i]);
-    EXPECT_TRUE(
-        LiteralTestUtil::Equal(expected_output_literal, actual_output_literal));
-  }
-}
-
-TEST_P(AllReduceKernelTest, KernelTestAddBF16) {
-  if (strategy() == AllReduceStrategy::kMultimem) {
-    GTEST_SKIP() << "Multimem does not support BF16.";
-  }
-  constexpr int64_t kNumRanks = 2;
-
-  std::vector<se::StreamExecutor*> executors = {GetGpuExecutor(0),
-                                                GetGpuExecutor(1)};
-
-  if (!executors[0]->CanEnablePeerAccessTo(executors[1])) {
-    GTEST_SKIP() << "Test requires direct peer memory access between devices.";
-  }
-
-  Array<bfloat16> expected_output({num_elements()});
-  std::vector<Array<bfloat16>> inputs;
-
-  for (int i = 0; i < kNumRanks; ++i) {
-    Array<bfloat16> input_data({num_elements()});
-    input_data.FillRandom(static_cast<bfloat16>(0.0f),
-                          static_cast<bfloat16>(10.0f), /*seed=*/i);
-
-    expected_output.Each([&](absl::Span<const int64_t> indices, bfloat16* val) {
-      *val += input_data(indices);
-    });
-
-    inputs.push_back(std::move(input_data));
-  }
-
-  ASSERT_OK_AND_ASSIGN(
-      auto results, RunKernel<bfloat16>(executors, inputs, ReductionKind::SUM));
-
-  for (int i = 0; i < kNumRanks; ++i) {
-    EXPECT_EQ(results[i], expected_output);
-  }
-}
-
-TEST_P(AllReduceKernelTest, KernelTestOrPred) {
-  if (strategy() == AllReduceStrategy::kMultimem) {
-    GTEST_SKIP() << "Multimem does not support predicates.";
-  }
-  constexpr int64_t kNumRanks = 2;
-
-  std::vector<se::StreamExecutor*> executors = {GetGpuExecutor(0),
-                                                GetGpuExecutor(1)};
-
-  if (!executors[0]->CanEnablePeerAccessTo(executors[1])) {
-    GTEST_SKIP() << "Test requires direct peer memory access between devices.";
-  }
-
-  Array<bool> expected_output({num_elements()});
-  std::vector<Array<bool>> inputs;
-
-  for (int i = 0; i < kNumRanks; ++i) {
-    Array<bool> input_data({num_elements()});
-    input_data.FillRandomBool(/*seed=*/i);
-
-    expected_output.Each([&](absl::Span<const int64_t> indices, bool* val) {
-      *val |= input_data(indices);
-    });
-
-    inputs.push_back(std::move(input_data));
-  }
-
-  // There are no logical operations in all-reduce reduction kind, so OR is
-  // simulated with MAX on uint8.
-  ASSERT_OK_AND_ASSIGN(auto results,
-                       RunKernel<bool>(executors, inputs, ReductionKind::MAX));
-
-  for (int i = 0; i < kNumRanks; ++i) {
-    EXPECT_EQ(results[i], expected_output);
-  }
-}
-
-TEST_P(AllReduceKernelTest, KernelTestAddPred_Unsupported) {
-  if (strategy() == AllReduceStrategy::kMultimem) {
-    GTEST_SKIP() << "Multimem does not support predicates.";
-  }
-  constexpr int64_t kNumRanks = 2;
-  std::vector<se::StreamExecutor*> executors = {GetGpuExecutor(0),
-                                                GetGpuExecutor(1)};
-
-  if (!executors[0]->CanEnablePeerAccessTo(executors[1])) {
-    GTEST_SKIP() << "Test requires direct peer memory access between devices.";
-  }
-
-  Array<bool> expected_output({num_elements()});
-  std::vector<Array<bool>> inputs(kNumRanks, Array<bool>({num_elements()}));
-
-  auto results = RunKernel<bool>(executors, inputs, ReductionKind::SUM);
-  EXPECT_THAT(results.status(), StatusIs(absl::StatusCode::kUnimplemented));
-  EXPECT_THAT(results.status().message(),
-              HasSubstr("combination is not supported"));
-}
-
-INSTANTIATE_TEST_SUITE_P(
-    AllReduceKernelTest, AllReduceKernelTest,
-    ::testing::ConvertGenerator(
-        ::testing::Combine(::testing::Values(AllReduceStrategy::kOneShot,
-                                             AllReduceStrategy::kTwoShot,
-                                             AllReduceStrategy::kMultimem),
-                           ::testing::Values(128000, 124000)),
-        [](const std::tuple<AllReduceStrategy, int64_t>& params) {
-          return TestParams{std::get<0>(params), std::get<1>(params)};
-        }),
-    [](const ::testing::TestParamInfo<TestParams>& info) {
-      return absl::StrFormat("%v_%d", info.param.all_reduce_strategy,
-                             info.param.num_elements);
-    });
-
 class AllReduceHloTest : public HloTestBase {};
 
 TEST_F(AllReduceHloTest, DefaultDeviceAssnWithHloRunner) {

From 14f66532dd1f93f8430e11685412d27f51c3685f Mon Sep 17 00:00:00 2001
From: Stanislav Bardyuk <sbardyuk@google.com>
Date: Wed, 14 Oct 2026 04:50:25 -0700
Subject: [PATCH 0761/1000] PR #48091: [XLA:CPU] Do not multi-output fuse
 scatter fusions
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/48091

## 📝 Summary of Changes

Fixes #47367.

With `--xla_backend_extra_options=xla_cpu_use_multi_output_fusion`, a
module containing two scatters that share operands crashes CPU
compilation with `symbolic_map.cc:196] Check failed: GetNumDims() ==
other.GetNumResults() (3 vs. 2)`.

`CpuMultiOutputFusion::IsFusible` accepts any loop fusion whose root is
not a dot, so two scatter-rooted wrapped fusions sharing indices/updates
get merged into a multi-output fusion rooted at `tuple(scatter,
scatter)`. The fusion kernel emitter dispatches to the dedicated scatter
emitter only when the fusion root itself is a scatter; the tuple root
falls through to the generic loop fusion emitter, whose indexing
analysis cannot express scatter (the operand indexing lives in the
update space, 3 dims here, while the subgraph indexing is over the root
space, 2 dims), and it dies on the CHECK.

This change excludes scatter-rooted fusions from multi-output fusion,
exactly like the existing dot exclusion in the same condition (both ops
have dedicated emitters that do not support multi-output roots). Nothing
regresses: the fusions this rejects currently crash the compiler.

## 🎯 Justification

The input is valid HLO (the same module compiles and runs fine without
the flag, and on GPU and interpreter); the flag is the CPU multi-output
fusion feature added in 52a47ae3ae. GPU's fusion machinery already treats
scatter specially (`IsInputFusibleScatter`) and never presents tuple-of-
scatter fusions to its emitters.

## 🚀 Kind of Contribution

🐛 Bug Fix

## 🧪 Unit Tests:

`MultiOutputFusionTest.DoesNotFuseScatterFusions` in
`cpu_multi_output_fusion_test.cc`: two wrapped scatter fusions sharing
indices and updates, results added; expects the pass to make no change
and asserts no tuple-rooted fusion is formed. Fails before this change
(the pass merges the scatters), passes after. All
`//xla/service/cpu/tests/...` targets pass.

## 🧪 Execution Tests:

With the flag set, the issue's original module and a reduced two-scatter
module both compile on CPU and produce results elementwise-identical to
the interpreter after this change; before it both abort on the CHECK.
Copybara import of the project:

--
9fb60eabf312a15fe96f387616b210d77859887f by Stanislav Bardyuk <sbardyuk@google.com>:

[XLA:CPU] Do not multi-output fuse scatter fusions

CpuMultiOutputFusion::IsFusible accepts any loop fusion whose root is
not a dot, so two scatter fusions that share operands get merged into a
multi-output fusion rooted at a tuple of scatters. The fusion kernel
emitter only dispatches to the dedicated scatter emitter when the
fusion root itself is a scatter, so the merged fusion falls through to
the generic loop fusion emitter, whose indexing analysis cannot express
scatter and crashes on the symbolic_map.cc:196 CHECK (the scatter
operand indexing lives in the update space while the subgraph indexing
is over the root space).

Exclude scatter-rooted fusions from multi-output fusion, like the
existing dot exclusion in the same condition: both ops are emitted by
dedicated emitters that do not support multi-output fusion roots. The
rejected fusions currently crash the compiler, so no working
configuration regresses.

Fixes #47367.

Merging this change closes #48091

PiperOrigin-RevId: 999287902
---
 .../service/cpu/cpu_multi_output_fusion.cc    |  5 +-
 .../cpu/cpu_multi_output_fusion_test.cc       | 69 +++++++++++++++++++
 2 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/cpu/cpu_multi_output_fusion.cc b/third_party/xla/xla/service/cpu/cpu_multi_output_fusion.cc
index 4068500d64d..ef2e0a45cb9 100644
--- a/third_party/xla/xla/service/cpu/cpu_multi_output_fusion.cc
+++ b/third_party/xla/xla/service/cpu/cpu_multi_output_fusion.cc
@@ -76,8 +76,11 @@ bool CpuMultiOutputFusion::IsFusible(HloInstruction* instr) {
     return false;
   }
 
+  // Dot and scatter fusions are emitted by dedicated emitters that do not
+  // support multi-output fusion roots.
   if (instr->IsLoopFusion() &&
-      instr->fused_expression_root()->opcode() != HloOpcode::kDot) {
+      instr->fused_expression_root()->opcode() != HloOpcode::kDot &&
+      instr->fused_expression_root()->opcode() != HloOpcode::kScatter) {
     return true;
   }
 
diff --git a/third_party/xla/xla/service/cpu/cpu_multi_output_fusion_test.cc b/third_party/xla/xla/service/cpu/cpu_multi_output_fusion_test.cc
index f5156706c08..5aac1f89e5e 100644
--- a/third_party/xla/xla/service/cpu/cpu_multi_output_fusion_test.cc
+++ b/third_party/xla/xla/service/cpu/cpu_multi_output_fusion_test.cc
@@ -158,5 +158,74 @@ TEST_F(MultiOutputFusionTest, DoesNotFuseInsideSortComparator) {
   }
 }
 
+// Regression test for https://github.com/openxla/xla/issues/47367: merging
+// scatter fusions creates a tuple-rooted fusion that bypasses the dedicated
+// scatter emitter and crashes the loop emitter's indexing analysis.
+TEST_F(MultiOutputFusionTest, DoesNotFuseScatterFusions) {
+  static constexpr absl::string_view kHloModule = R"(
+    HloModule module
+
+    %add_a {
+      %a = f32[] parameter(0)
+      %b = f32[] parameter(1)
+      ROOT %add = f32[] add(%a, %b)
+    }
+
+    %add_b {
+      %a = f32[] parameter(0)
+      %b = f32[] parameter(1)
+      ROOT %add = f32[] add(%a, %b)
+    }
+
+    %scatter_a {
+      %operand = f32[10,5] parameter(0)
+      %indices = s32[24,1] parameter(1)
+      %updates = f32[24,2,3] parameter(2)
+      ROOT %scatter = f32[10,5] scatter(%operand, %indices, %updates),
+        update_window_dims={1,2}, inserted_window_dims={},
+        scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%add_a
+    }
+
+    %scatter_b {
+      %operand = f32[10,5] parameter(0)
+      %indices = s32[24,1] parameter(1)
+      %updates = f32[24,2,3] parameter(2)
+      ROOT %scatter = f32[10,5] scatter(%operand, %indices, %updates),
+        update_window_dims={1,2}, inserted_window_dims={},
+        scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%add_b
+    }
+
+    ENTRY %main {
+      %arg0 = f32[10,5] parameter(0)
+      %arg1 = f32[10,5] parameter(1)
+      %indices = s32[24,1] parameter(2)
+      %updates = f32[24,2,3] parameter(3)
+      %scatter0 = f32[10,5] fusion(%arg0, %indices, %updates), kind=kLoop,
+        calls=%scatter_a
+      %scatter1 = f32[10,5] fusion(%arg1, %indices, %updates), kind=kLoop,
+        calls=%scatter_b
+      ROOT %sum = f32[10,5] add(%scatter0, %scatter1)
+    }
+  )";
+
+  TF_ASSERT_OK_AND_ASSIGN(auto hlo_module,
+                          ParseAndReturnVerifiedModule(kHloModule));
+  TF_ASSERT_OK_AND_ASSIGN(
+      bool changed, CpuMultiOutputFusion(&alias_info_).Run(hlo_module.get()));
+  EXPECT_FALSE(changed) << hlo_module->ToString();
+  int num_scatter_fusions = 0;
+  for (const HloInstruction* instr :
+       hlo_module->entry_computation()->instructions()) {
+    if (instr->opcode() == HloOpcode::kFusion) {
+      EXPECT_NE(instr->fused_expression_root()->opcode(), HloOpcode::kTuple)
+          << "Scatter fusions must not be merged into a multi-output fusion: "
+          << instr->ToString();
+      num_scatter_fusions +=
+          instr->fused_expression_root()->opcode() == HloOpcode::kScatter;
+    }
+  }
+  EXPECT_EQ(num_scatter_fusions, 2);
+}
+
 }  // namespace
 }  // namespace xla::cpu

From 99d6777b4edbb55fbf3203977e3ded496eca33d9 Mon Sep 17 00:00:00 2001
From: Levon Ter-Grigoryan <patrios@google.com>
Date: Wed, 14 Oct 2026 05:09:09 -0700
Subject: [PATCH 0762/1000] Add tile propagation support for
 `HloOpcode::kAllToAll` in `PropagateTileToInputForAllToAllOp`, mapping the
 output split-dimension offset `offset` to the intra-shard offset `offset %
 shard_size` and peer replica index `offset / shard_size` while validating
 that the tile size divides `shard_size`.

PiperOrigin-RevId: 999295758
---
 .../tiling/experimental/tile_propagation.cc   |  85 ++++++++++++++
 .../experimental/tile_propagation_test.cc     | 106 ++++++++++++++++++
 2 files changed, 191 insertions(+)

diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
index 6996fd03093..ac56ca5809d 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
@@ -1713,6 +1713,87 @@ absl::StatusOr<Tiles> PropagateTileToInputForAllGatherOp(
   return Tiles{input_tile};
 }
 
+absl::StatusOr<Tiles> PropagateTileToInputForAllToAllOp(
+    const TilingSpace& tiling_space, const HloInstruction& hlo,
+    const Tile& output_tile, int64_t output_index) {
+  const auto& all_to_all = *Cast<HloAllToAllInstruction>(&hlo);
+  if (!all_to_all.split_dimension().has_value()) {
+    return absl::UnimplementedError(absl::StrCat(
+        "Tuple all-to-all without split_dimension is not supported: ",
+        hlo.ToString()));
+  }
+  int64_t split_dim = *all_to_all.split_dimension();
+  if (output_index < 0 || output_index >= hlo.operand_count()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Output index ", output_index, " is out of range for ",
+                     hlo.ToString()));
+  }
+  const Shape& input_shape = hlo.operand(output_index)->shape();
+  int64_t num_replicas =
+      all_to_all.device_list()->replica_groups().empty()
+          ? 1
+          : all_to_all.device_list()->replica_groups()[0].replica_ids_size();
+  int64_t shard_size = input_shape.dimensions(split_dim) / num_replicas;
+
+  const DimTile& output_dim_tile = output_tile.dim_tiles()[split_dim];
+  if (!IsConstantValue(output_dim_tile.stride, 1)) {
+    return absl::UnimplementedError(absl::StrCat(
+        "Non-unit stride on the split dimension of ", all_to_all.ToString()));
+  }
+
+  if (!tiling_space.IsSymbolic()) {
+    DimTile simplified_dim_tile = output_dim_tile;
+    simplified_dim_tile.Simplify(tiling_space);
+    if (simplified_dim_tile.size.GetType() != SymbolicExprType::kConstant ||
+        simplified_dim_tile.size.GetValue() <= 0 ||
+        shard_size % simplified_dim_tile.size.GetValue() != 0 ||
+        !simplified_dim_tile.offset.Canonicalize().IsMultipleOf(
+            simplified_dim_tile.size.GetValue())) {
+      return absl::UnimplementedError(absl::StrCat(
+          "Output ", output_index, " of ", hlo.name(), " is tiled with size ",
+          simplified_dim_tile.size.ToString(), " at offset ",
+          simplified_dim_tile.offset.ToString(), " on split dimension ",
+          split_dim, ", expected tile size to divide per-rank shard size ",
+          shard_size,
+          " and offset to be a multiple of tile size (one peer per tile)"));
+    }
+  }
+
+  SymbolicExpr replica_id =
+      (output_dim_tile.offset / shard_size).Canonicalize();
+  SymbolicExpr input_offset =
+      (output_dim_tile.offset % shard_size).Canonicalize();
+
+  llvm::SmallVector<DimTile> input_dim_tiles;
+  input_dim_tiles.reserve(output_tile.num_dim_tiles());
+
+  for (int64_t i = 0; i < output_tile.num_dim_tiles(); ++i) {
+    if (i == split_dim) {
+      input_dim_tiles.push_back(DimTile{
+          input_offset, output_dim_tile.size, output_dim_tile.stride,
+          CreateSymbolicConstant(shard_size, output_tile.mlir_context())});
+    } else {
+      input_dim_tiles.push_back(output_tile.dim_tiles()[i]);
+    }
+  }
+
+  mlir::MLIRContext* ctx = output_tile.mlir_context();
+  llvm::SmallVector<DimTile> replica_id_dim_tiles =
+      llvm::to_vector(output_tile.replica_ids());
+  replica_id_dim_tiles.push_back(DimTile{
+      /*offset=*/replica_id,
+      /*size=*/CreateSymbolicConstant(1, ctx),
+      /*stride=*/CreateSymbolicConstant(1, ctx),
+      /*upper_bound=*/
+      CreateSymbolicConstant(num_replicas, ctx),
+  });
+
+  Tile input_tile(output_tile.tiling_space(), std::move(input_dim_tiles),
+                  std::move(replica_id_dim_tiles), output_tile.constraints());
+
+  return Tiles{input_tile};
+}
+
 absl::StatusOr<Tiles> PropagateTileToInput(TilingSpace& tiling_space,
                                            const HloInstruction& hlo,
                                            const Tile& output_tile,
@@ -1731,6 +1812,10 @@ absl::StatusOr<Tiles> PropagateTileToInput(TilingSpace& tiling_space,
     return PropagateTileToInputForAllGatherOp(tiling_space, hlo, output_tile,
                                               output_index);
   }
+  if (hlo.opcode() == HloOpcode::kAllToAll) {
+    return PropagateTileToInputForAllToAllOp(tiling_space, hlo, output_tile,
+                                             output_index);
+  }
   if (hlo.opcode() == HloOpcode::kReduceScatter) {
     return PropagateTileToInputForReduceScatterOp(
         tiling_space, *Cast<HloReduceScatterInstruction>(&hlo), output_tile);
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
index f4b6d815421..f061111e255 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
@@ -1281,6 +1281,112 @@ TEST_F(TilePropagationTest, CanPropagateReplicaIdThroughBroadcast) {
   )"));
 }
 
+TEST_F(TilePropagationTest, CanPropagateToInputOfAllToAllOp) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    ENTRY e {
+      p0 = f32[10, 64, 5] parameter(0)
+      ROOT all_to_all = f32[10, 64, 5] all-to-all(p0), replica_groups={{0,1}}, dimensions={1}
+    }
+  )");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TilingSpace> tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                          &mlir_context_));
+
+  ASSERT_OK_AND_ASSIGN(auto tiled_a2a_operands,
+                       PropagateTileToInput(*tiling_space, *root,
+                                            tiling_space->tiled_roots()[0], 0));
+
+  EXPECT_THAT(tiled_a2a_operands, MatchToString(R"(
+    0) (tid_0, tid_1, tid_2)
+      -> offsets [tid_0 * ts_0, (tid_1 * ts_1) mod 32, tid_2 * ts_2]
+         sizes [ts_0, ts_1, ts_2]
+         strides [1, 1, 1]
+         upper bounds [10, 32, 5]
+         replica ids {
+           offsets [(tid_1 * ts_1) / 32]
+           sizes [1]
+           strides [1]
+           upper bounds [2]
+         }
+  )"));
+
+  ASSERT_OK(tiling_space->AssignTileSizes({10, 16, 5}));
+  ASSERT_OK_AND_ASSIGN(auto concrete_tiled_a2a_operands,
+                       PropagateTileToInput(*tiling_space, *root,
+                                            tiling_space->tiled_roots()[0], 0));
+
+  EXPECT_THAT(concrete_tiled_a2a_operands, MatchToString(R"(
+    0) (tid_0, tid_1, tid_2)
+      -> offsets [0, (tid_1 mod 2) * 16, 0]
+         sizes [10, 16, 5]
+         strides [1, 1, 1]
+         upper bounds [10, 32, 5]
+         replica ids {
+           offsets [tid_1 / 2]
+           sizes [1]
+           strides [1]
+           upper bounds [2]
+         }
+  )"));
+}
+
+TEST_F(TilePropagationTest, AllToAllRejectsNonUnitSplitStride) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    ENTRY e {
+      p0 = f32[10, 64, 5] parameter(0)
+      ROOT all_to_all = f32[10, 64, 5] all-to-all(p0), replica_groups={{0,1}}, dimensions={1}
+    }
+  )");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TilingSpace> tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                          &mlir_context_));
+  // GetTestTile assigns stride = index + 1 = 2 on dimension 1.
+  EXPECT_THAT(PropagateTileToInput(
+                  *tiling_space, *root,
+                  GetTestTile(*tiling_space, root->shape().dimensions()), 0),
+              StatusIs(absl::StatusCode::kUnimplemented));
+}
+
+TEST_F(TilePropagationTest, AllToAllRejectsTileCrossingShards) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    ENTRY e {
+      p0 = f32[10, 64, 5] parameter(0)
+      ROOT all_to_all = f32[10, 64, 5] all-to-all(p0), replica_groups={{0,1}}, dimensions={1}
+    }
+  )");
+  {
+    ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<TilingSpace> tiling_space,
+        TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                            &mlir_context_));
+    // Shard size on dimension 1 is 32, which is not divisible by tile size 24.
+    ASSERT_OK(tiling_space->AssignTileSizes({10, 24, 5}));
+    EXPECT_THAT(PropagateTileToInput(*tiling_space, *root,
+                                     tiling_space->tiled_roots()[0], 0),
+                StatusIs(absl::StatusCode::kUnimplemented));
+  }
+  {
+    ASSERT_OK_AND_ASSIGN(
+        std::unique_ptr<TilingSpace> tiling_space,
+        TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                            &mlir_context_));
+    ASSERT_OK(tiling_space->AssignTileSizes({10, 16, 5}));
+    Tile unaligned_tile = tiling_space->tiled_roots()[0];
+    SmallVector<DimTile> dim_tiles(unaligned_tile.dim_tiles());
+    // Offset tid_1 * 16 + 8 is not a multiple of tile size 16.
+    dim_tiles[1].offset = dim_tiles[1].offset + 8;
+    EXPECT_THAT(PropagateTileToInput(
+                    *tiling_space, *root,
+                    unaligned_tile.CloneWithNewDims(std::move(dim_tiles)), 0),
+                StatusIs(absl::StatusCode::kUnimplemented));
+  }
+}
+
 TEST_F(TilePropagationTest, CanPropagateToInputsOfReduceOp) {
   HloInstruction* root = ParseAndGetRoot(R"(
     HloModule m

From 54324c7fb1da5224bfd76d9b6d2147605e737aa8 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2026 05:30:16 -0700
Subject: [PATCH 0763/1000] Optimize XLA buffer assignment by indexing values
 and live ranges by ID.

When `BufferAssigner` falls back to the default allocation manager on large modules with hundreds of thousands of buffers, `BufferAssigner::MaybeAssignBuffer` becomes a major compile-time bottleneck. Profiling shows most of the time is spent performing `HloLiveRange::buffer_live_ranges()` hash table lookups per assigned `HloValue*`, re-evaluating `opts_.must_not_live_out` and `BufferLivesOut` across candidate allocations, and running `CHECK` and `dynamic_cast` assertions inside the inner loop.

Replace pointer-keyed and ID-keyed hash tables in `BufferAssignment` with dense vectors indexed by `HloValue::Id` and `HloBuffer::Id` for allocation indices, buffer sizes, live ranges, `BufferLivesOut`, and `must_not_live_out` predicates. Store parallel contiguous vectors of assigned `HloValue*` and `{start, end}` live-range bounds on `BufferAllocation` so `MaybeAssignBuffer` scans live ranges sequentially without hash lookups or `HloValue*` dereferences unless live ranges touch at an endpoint. Replace hot-path `CHECK` and `dynamic_cast` assertions with `DCHECK`s, reorder early-exit checks before computing buffer colors, and move the `HloOpcode::kCopy` operand check after live-range interference.

| Benchmark | CPU Time (Before) | CPU Time (After) | Speedup |
| :--- | ---: | ---: | ---: |
| `BM_MaybeAssignBufferStress/1000` | 0.598 ms | 0.196 ms | **3.05x** |
| `BM_MaybeAssignBufferStress/4096` | 2.787 ms | 0.976 ms | **2.86x** |
| `BM_MaybeAssignBufferStress/32768` | 28.50 ms | 10.80 ms | **2.64x** |
| `BM_MaybeAssignBufferStress/262144` | 277.52 ms | 107.66 ms | **2.58x** |
| `BM_MaybeAssignBufferStress/1000000` | 1149.88 ms | 434.04 ms | **2.65x** |

PiperOrigin-RevId: 999303779
---
 third_party/xla/xla/service/BUILD             |   1 +
 .../xla/xla/service/buffer_assignment.cc      | 442 ++++++++++--------
 .../xla/xla/service/buffer_assignment.h       | 151 +++++-
 .../xla/xla/service/buffer_assignment_test.cc | 122 +++++
 third_party/xla/xla/service/hlo_buffer.cc     |  23 +
 third_party/xla/xla/service/hlo_buffer.h      |  27 +-
 6 files changed, 528 insertions(+), 238 deletions(-)

diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index 5f26fbc2818..bcf32652adc 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -3515,6 +3515,7 @@ cc_library(
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
diff --git a/third_party/xla/xla/service/buffer_assignment.cc b/third_party/xla/xla/service/buffer_assignment.cc
index 32f0bbd1e9f..8b8bfa1b891 100644
--- a/third_party/xla/xla/service/buffer_assignment.cc
+++ b/third_party/xla/xla/service/buffer_assignment.cc
@@ -92,16 +92,15 @@ struct BufferLiveRange {
   int64_t max_end = -1;
 };
 
-BufferLiveRange GetHloBufferLiveRange(
-    const HloBuffer* hlo_buffer,
-    const absl::flat_hash_map<const HloValue*, HloLiveRange::LiveRangeBounds>&
-        value_live_ranges) {
+BufferLiveRange GetHloBufferLiveRange(const HloBuffer* hlo_buffer,
+                                      const BufferAssignment& assignment) {
   BufferLiveRange result;
   for (const HloValue* value : hlo_buffer->values()) {
-    auto it = value_live_ranges.find(value);
-    if (it != value_live_ranges.end()) {
-      result.min_start = std::min(result.min_start, it->second.start);
-      result.max_end = std::max(result.max_end, it->second.end);
+    BufferAllocation::ValueLiveRange lr =
+        assignment.GetValueLiveRange(value->id());
+    if (lr.start >= 0) {
+      result.min_start = std::min(result.min_start, lr.start);
+      result.max_end = std::max(result.max_end, lr.end);
     }
   }
   return result;
@@ -154,7 +153,7 @@ class DefaultBufferAllocationsManagerForComputationsWithoutOrdering
     for (int allocation_index = allocation_indices_.size() - 1;
          allocation_index >= 0; --allocation_index) {
       BufferAllocation* allocation = assignment_->GetMutableAllocation(
-          allocation_indices_.at(allocation_index));
+          allocation_indices_[allocation_index]);
       ABSL_ASSIGN_OR_RETURN(bool success, assigner_->MaybeAssignBuffer(
                                          allocation, *hlo_buffer, assignment_));
       if (success) {
@@ -224,7 +223,7 @@ class FastMergeBufferAllocationsManagerForComputationsWithoutOrdering
     BufferLiveRange live_range = GetBufferLiveRange(hlo_buffer);
     UpdateLiveness(live_range.min_start);
 
-    ABSL_ASSIGN_OR_RETURN(auto buffer_color, hlo_buffer->color());
+    ABSL_ASSIGN_OR_RETURN(BufferValue::Color buffer_color, hlo_buffer->color());
     auto& pool = free_pool_[buffer_color];
     // Finds the smallest buffer in the free pool that is at least
     // `required_size`.
@@ -323,8 +322,7 @@ class FastMergeBufferAllocationsManagerForComputationsWithoutOrdering
   }
 
   BufferLiveRange GetBufferLiveRange(const HloBuffer* hlo_buffer) const {
-    return GetHloBufferLiveRange(
-        hlo_buffer, assignment_->hlo_live_range().buffer_live_ranges());
+    return GetHloBufferLiveRange(hlo_buffer, *assignment_);
   }
 
   // Advances the sweep-line state to the current start time, moving expired
@@ -409,10 +407,10 @@ std::optional<bool> CompareSize(
   return std::nullopt;
 };
 
-std::optional<bool> CompareLiveOut(const HloAliasAnalysis* alias_analysis,
+std::optional<bool> CompareLiveOut(const BufferAssignment* assignment,
                                    const HloBuffer* a, const HloBuffer* b) {
-  const bool a_live_out = alias_analysis->BufferLivesOut(*a);
-  const bool b_live_out = alias_analysis->BufferLivesOut(*b);
+  const bool a_live_out = assignment->BufferLivesOut(*a);
+  const bool b_live_out = assignment->BufferLivesOut(*b);
   if (a_live_out != b_live_out) {
     return a_live_out;
   }
@@ -438,19 +436,16 @@ std::optional<bool> ComparePosition(
 }
 
 // Compares two buffers by their earliest live range start time.
-std::optional<bool> CompareLiveRangeStart(
-    const absl::flat_hash_map<const HloValue*, HloLiveRange::LiveRangeBounds>*
-        buffer_live_ranges,
-    absl::flat_hash_map<const HloBuffer*, int64_t>* min_start_cache,
-    const HloBuffer* a, const HloBuffer* b) {
-  auto get_min_start = [buffer_live_ranges,
-                        min_start_cache](const HloBuffer* buffer) {
-    auto [it, inserted] =
-        min_start_cache->try_emplace(buffer, kUninitializedStartTime);
-    if (inserted) {
-      it->second = GetHloBufferLiveRange(buffer, *buffer_live_ranges).min_start;
+std::optional<bool> CompareLiveRangeStart(const BufferAssignment* assignment,
+                                          std::vector<int64_t>* min_start_cache,
+                                          const HloBuffer* a,
+                                          const HloBuffer* b) {
+  auto get_min_start = [assignment, min_start_cache](const HloBuffer* buffer) {
+    int64_t& cached = (*min_start_cache)[buffer->id()];
+    if (cached == kUninitializedStartTime) {
+      cached = GetHloBufferLiveRange(buffer, *assignment).min_start;
     }
-    return it->second;
+    return cached;
   };
   int64_t a_start = get_min_start(a);
   int64_t b_start = get_min_start(b);
@@ -684,11 +679,12 @@ BufferAllocation::Slice BufferAllocation::GetSlice(
 }
 
 absl::Status BufferAllocation::AddAssignment(const HloValue& buffer,
-                                             int64_t offset, int64_t size) {
+                                             int64_t offset, int64_t size,
+                                             ValueLiveRange live_range) {
   VLOG(4) << "Adding the following buffer to allocation #" << index()
           << absl::StrFormat(" (size=%d, offset=%d) %s", size, offset,
                              buffer.ToShortString());
-  CHECK(!assigned_buffers_.contains(&buffer))
+  DCHECK(!assigned_buffers_.contains(&buffer))
       << "LogicalBuffer " << buffer << " already assigned to allocation "
       << index_;
   // TF_RET_CHECK (rather than CHECK_LE) so that a malformed offset/size --
@@ -715,6 +711,8 @@ absl::Status BufferAllocation::AddAssignment(const HloValue& buffer,
   offset_size.offset = offset;
   offset_size.size = size;
   assigned_buffers_.emplace(&buffer, offset_size);
+  assigned_values_.push_back(&buffer);
+  assigned_live_ranges_.push_back(live_range);
   // For debugging purposes, store the assigned memory space in the
   // instruction's layout.
   for (const HloPosition& position : buffer.positions()) {
@@ -913,23 +911,64 @@ std::ostream& operator<<(std::ostream& out, const BufferAllocation::Slice& s) {
   return out;
 }
 
+void BufferAssignment::InitCaches() {
+  HloValue::Id max_value_id = -1;
+  for (const HloValue* value : dataflow_analysis().values()) {
+    max_value_id = std::max(max_value_id, value->id());
+  }
+  const size_t num_value_ids = static_cast<size_t>(max_value_id + 1);
+  allocation_index_for_value_id_.assign(num_value_ids, -1);
+  value_lives_out_.assign(num_value_ids, 0);
+  value_must_not_live_out_.assign(num_value_ids, kMustNotLiveOutUncomputed);
+
+  HloBuffer::Id max_buffer_id = -1;
+  for (const HloBuffer& buffer : alias_analysis_->buffers()) {
+    max_buffer_id = std::max(max_buffer_id, buffer.id());
+  }
+  const size_t num_buffer_ids = static_cast<size_t>(max_buffer_id + 1);
+  cached_buffer_sizes_.assign(num_buffer_ids, kUncomputedBufferSize);
+  buffer_lives_out_.assign(num_buffer_ids, 0);
+  buffer_must_not_live_out_.assign(num_buffer_ids, kMustNotLiveOutUncomputed);
+
+  for (const HloBuffer* live_out_buffer : alias_analysis_->LiveOutBuffers()) {
+    buffer_lives_out_[live_out_buffer->id()] = 1;
+    for (const HloValue* value : live_out_buffer->values()) {
+      value_lives_out_[value->id()] = 1;
+    }
+  }
+
+  if (hlo_live_range_ != nullptr) {
+    live_ranges_by_value_id_.resize(num_value_ids);
+    const auto& buffer_live_ranges = hlo_live_range_->buffer_live_ranges();
+    for (const HloValue* value : dataflow_analysis().values()) {
+      if (auto it = buffer_live_ranges.find(value);
+          it != buffer_live_ranges.end()) {
+        live_ranges_by_value_id_[value->id()] =
+            BufferAllocation::ValueLiveRange{it->second.start, it->second.end};
+      }
+    }
+  }
+}
+
 bool BufferAssignment::HasAllocation(const HloValue& value) const {
-  return allocation_index_for_value_.contains(&value);
+  return HasAllocation(value.id());
 }
 
 bool BufferAssignment::HasAllocation(HloValue::Id value_id) const {
-  return HasAllocation(dataflow_analysis().GetValue(value_id));
+  return value_id >= 0 &&
+         static_cast<size_t>(value_id) <
+             allocation_index_for_value_id_.size() &&
+         allocation_index_for_value_id_[value_id] >= 0;
 }
 
 bool BufferAssignment::HasAllocation(const HloBuffer& buffer) const {
-  return allocation_index_for_value_.contains(buffer.values()[0]);
+  return HasAllocation(buffer.values()[0]->id());
 }
 
 const BufferAllocation& BufferAssignment::GetAssignedAllocation(
     const HloValue& value) const {
-  auto it = allocation_index_for_value_.find(&value);
-  CHECK(it != allocation_index_for_value_.end());
-  return GetAllocation(it->second);
+  CHECK(HasAllocation(value.id()));
+  return GetAllocation(allocation_index_for_value_id_[value.id()]);
 }
 
 const BufferAllocation& BufferAssignment::GetAssignedAllocation(
@@ -956,8 +995,8 @@ std::set<BufferAllocation::Slice> BufferAssignment::GetAllSlices(
 
 const BufferAllocation& BufferAssignment::GetAllocation(
     BufferAllocation::Index index) const {
-  CHECK_GE(index, 0);
-  CHECK_LT(index, allocations_.size());
+  DCHECK_GE(index, 0);
+  DCHECK_LT(index, allocations_.size());
   return allocations_[index];
 }
 
@@ -984,7 +1023,7 @@ bool BufferAssignment::HasAllocationAt(const HloInstruction* instruction,
                                        const ShapeIndex& index) const {
   return absl::c_any_of(
       dataflow_analysis().GetValueSet(instruction, index).values(),
-      IsKeyIn(allocation_index_for_value_));
+      [this](const HloValue* value) { return HasAllocation(*value); });
 }
 
 bool BufferAssignment::HasTopLevelAllocation(
@@ -1128,18 +1167,20 @@ absl::StatusOr<BufferAllocation*> BufferAssignment::NewAllocation(
 absl::Status BufferAssignment::AddAssignment(BufferAllocation* allocation,
                                              const HloBuffer& buffer,
                                              int64_t offset, int64_t size) {
-  CHECK(allocation->is_reusable() || allocation->assigned_buffers().empty())
+  DCHECK(allocation->is_reusable() || allocation->assigned_buffers().empty())
       << "Non-reusable allocation already assigned a buffer: "
       << allocation->ToString();
 
   for (const HloValue* buffer_value : buffer.values()) {
-    CHECK(!allocation_index_for_value_.contains(buffer_value))
+    DCHECK(!HasAllocation(buffer_value->id()))
         << "BufferValue " << buffer_value << " already has an allocation.";
-    ABSL_RETURN_IF_ERROR(allocation->AddAssignment(*buffer_value, offset, size));
-    allocation_index_for_value_[buffer_value] = allocation->index();
+    ABSL_RETURN_IF_ERROR(allocation->AddAssignment(
+        *buffer_value, offset, size, GetValueLiveRange(buffer_value->id())));
+    allocation_index_for_value_id_[buffer_value->id()] = allocation->index();
+    ++num_assigned_values_;
   }
 
-  if (alias_analysis().BufferLivesOut(buffer)) {
+  if (BufferLivesOut(buffer)) {
     VLOG(3) << "HloBuffer lives out: " << buffer.ToString();
     VLOG(3) << "Set maybe live out: " << allocation->ToString();
     allocation->set_maybe_live_out(true);
@@ -1150,12 +1191,14 @@ absl::Status BufferAssignment::AddAssignment(BufferAllocation* allocation,
 absl::Status BufferAssignment::AddAssignment(BufferAllocation* allocation,
                                              const HloValue& value,
                                              int64_t offset, int64_t size) {
-  ABSL_RETURN_IF_ERROR(allocation->AddAssignment(value, offset, size));
-  allocation_index_for_value_[&value] = allocation->index();
-  const HloValue& hlo_value =
-      *CHECK_NOTNULL(dynamic_cast<const HloValue*>(&value));
-  if (alias_analysis().ValueLivesOut(hlo_value)) {
-    VLOG(3) << "HloValue lives out: " << hlo_value.ToString();
+  ABSL_RETURN_IF_ERROR(allocation->AddAssignment(value, offset, size,
+                                            GetValueLiveRange(value.id())));
+  if (allocation_index_for_value_id_[value.id()] < 0) {
+    ++num_assigned_values_;
+  }
+  allocation_index_for_value_id_[value.id()] = allocation->index();
+  if (ValueLivesOut(value)) {
+    VLOG(3) << "HloValue lives out: " << value.ToString();
     VLOG(3) << "Set maybe live out: " << allocation->ToString();
     allocation->set_maybe_live_out(true);
   }
@@ -1289,8 +1332,8 @@ absl::Status BufferAssignment::CombineTempAllocations(
       const HloValue* value = buffer_offset_size.first;
       const int64_t offset = buffer_offset_size.second.offset;
       const int64_t size = buffer_offset_size.second.size;
-      ABSL_RETURN_IF_ERROR(
-          combined_allocation->AddAssignment(*value, base + offset, size));
+      ABSL_RETURN_IF_ERROR(combined_allocation->AddAssignment(
+          *value, base + offset, size, GetValueLiveRange(value->id())));
     }
     if (!temp_allocation.HeapTraces().empty()) {
       CHECK_EQ(temp_allocation.HeapTraces().size(), 1);
@@ -1324,16 +1367,16 @@ absl::Status BufferAssignment::CombineTempAllocations(
   }
 
   // Update allocation indices to their new positions.
-  allocation_index_for_value_.erase(allocation_index_for_value_.begin(),
-                                    allocation_index_for_value_.end());
+  std::fill(allocation_index_for_value_id_.begin(),
+            allocation_index_for_value_id_.end(), -1);
+  num_assigned_values_ = 0;
   for (BufferAllocation::Index index = 0; index < allocations_.size();
        ++index) {
     BufferAllocation& allocation = allocations_[index];
     allocation.set_index(index);
-    // NOLINTNEXTLINE : the order of the loop is not important.
-    for (const auto& buffer_offset_size : allocation.assigned_buffers_) {
-      const HloValue* value = buffer_offset_size.first;
-      allocation_index_for_value_[value] = index;
+    for (const HloValue* value : allocation.assigned_values_) {
+      allocation_index_for_value_id_[value->id()] = index;
+      ++num_assigned_values_;
     }
   }
   return absl::OkStatus();
@@ -1348,7 +1391,9 @@ void BufferAssignment::ClearAllocationsForColors(
                      allocations_.end());
 
   temp_allocation_total_size_ = 0;
-  allocation_index_for_value_.clear();
+  std::fill(allocation_index_for_value_id_.begin(),
+            allocation_index_for_value_id_.end(), -1);
+  num_assigned_values_ = 0;
   absl::erase_if(fragmentation_bytes_by_color_, [&](const auto& entry) {
     return colors.contains(entry.first);
   });
@@ -1357,9 +1402,9 @@ void BufferAssignment::ClearAllocationsForColors(
        ++index) {
     BufferAllocation& allocation = allocations_[index];
     allocation.set_index(index);
-    // NOLINTNEXTLINE : the order of the loop is not important.
-    for (const auto& [value, offset_size] : allocation.assigned_buffers_) {
-      allocation_index_for_value_[value] = index;
+    for (const HloValue* value : allocation.assigned_values_) {
+      allocation_index_for_value_id_[value->id()] = index;
+      ++num_assigned_values_;
     }
   }
 }
@@ -1972,7 +2017,7 @@ absl::StatusOr<std::unique_ptr<BufferAssignment>> BufferAssignment::FromProto(
 
   // Ensure each buffer in the proto has an allocation assigned.
   TF_RET_CHECK(proto.logical_buffers_size() ==
-               buffer_assignment->allocation_index_for_value_.size());
+               buffer_assignment->num_assigned_values_);
   for (auto& logical_buffer_proto : proto.logical_buffers()) {
     TF_RET_CHECK(buffer_assignment->HasAllocation(
         *id_to_logical_buffer[logical_buffer_proto.id()]));
@@ -1984,10 +2029,16 @@ absl::StatusOr<std::unique_ptr<BufferAssignment>> BufferAssignment::FromProto(
 void BufferAssignment::Finalize() {
   hlo_ordering_.reset();
   hlo_live_range_.reset();
+  live_ranges_by_value_id_.clear();
+  live_ranges_by_value_id_.shrink_to_fit();
 
   for (BufferAllocation& allocation : allocations_) {
     allocation.heap_traces_.clear();
     allocation.heap_traces_.shrink_to_fit();
+    allocation.assigned_values_.clear();
+    allocation.assigned_values_.shrink_to_fit();
+    allocation.assigned_live_ranges_.clear();
+    allocation.assigned_live_ranges_.shrink_to_fit();
   }
 }
 
@@ -2015,10 +2066,10 @@ absl::StatusOr<std::unique_ptr<BufferAssignment>> BufferAssigner::Run(
 }
 
 bool BufferAssigner::LiveRangeInterferes(
-    const HloValue* buffer1, const HloLiveRange::LiveRangeBounds& live_range1,
-    const HloValue* buffer2, const HloLiveRange::LiveRangeBounds& live_range2,
+    const HloValue* buffer1, BufferAllocation::ValueLiveRange live_range1,
+    const HloValue* buffer2, BufferAllocation::ValueLiveRange live_range2,
     BufferAssignment* assignment) {
-  CHECK((assignment->hlo_live_range().total_order_scheduled()));
+  DCHECK((assignment->hlo_live_range().total_order_scheduled()));
 
   // An HloValue can hold multiple instruction positions during its lifetime
   // (e.g. when passed into a while loop tuple or bitcast view).
@@ -2028,24 +2079,22 @@ bool BufferAssigner::LiveRangeInterferes(
   // different position of the value (such as the definition instruction
   // before the loop). Therefore, we must check all positions of the
   // HloValue to see if any instruction can share its buffer with the user.
-  auto can_share_as_operand =
-      [&assignment, this](
-          const HloValue* user_value, const HloValue* operand_value,
-          const HloLiveRange::LiveRangeBounds& operand_live_range) {
-        if (user_value->instruction()->opcode() == HloOpcode::kCopy) {
-          return false;
-        }
-        for (const HloPosition& operand_pos : operand_value->positions()) {
-          if (user_value->instruction()->IsUserOf(operand_pos.instruction) &&
-              assignment->dataflow_analysis().CanShareOperandBufferWithUser(
-                  operand_pos.instruction, operand_pos.index,
-                  user_value->instruction(), user_value->index(),
-                  alias_info_)) {
-            return true;
-          }
-        }
-        return false;
-      };
+  auto can_share_as_operand = [&assignment, this](
+                                  const HloValue* user_value,
+                                  const HloValue* operand_value) {
+    if (user_value->instruction()->opcode() == HloOpcode::kCopy) {
+      return false;
+    }
+    for (const HloPosition& operand_pos : operand_value->positions()) {
+      if (user_value->instruction()->IsUserOf(operand_pos.instruction) &&
+          assignment->dataflow_analysis().CanShareOperandBufferWithUser(
+              operand_pos.instruction, operand_pos.index,
+              user_value->instruction(), user_value->index(), alias_info_)) {
+        return true;
+      }
+    }
+    return false;
+  };
 
   bool interferes = false;
   if (!(live_range1.start > live_range2.end ||
@@ -2053,7 +2102,7 @@ bool BufferAssigner::LiveRangeInterferes(
     if (live_range1.end == live_range2.start) {
       auto operand_value = buffer1;
       auto user_value = buffer2;
-      if (!can_share_as_operand(user_value, operand_value, live_range1)) {
+      if (!can_share_as_operand(user_value, operand_value)) {
         VLOG(4) << "End of live range of " << buffer1->ToShortString()
                 << " is equal to the start of live range of "
                 << buffer2->ToShortString() << ", buffer cannot be shared.";
@@ -2062,7 +2111,7 @@ bool BufferAssigner::LiveRangeInterferes(
     } else if (live_range2.end == live_range1.start) {
       auto operand_value = buffer2;
       auto user_value = buffer1;
-      if (!can_share_as_operand(user_value, operand_value, live_range2)) {
+      if (!can_share_as_operand(user_value, operand_value)) {
         VLOG(4) << "End of live range of " << buffer2->ToShortString()
                 << " is equal to the start of live range of "
                 << buffer1->ToShortString() << ", buffer cannot be shared.";
@@ -2096,17 +2145,60 @@ bool BufferAssigner::LiveRangeInterferes(
   return interferes;
 }
 
+bool BufferAssigner::HasLiveRangeInterference(
+    const BufferAllocation& allocation, const HloBuffer& hlo_buffer,
+    BufferAssignment* assignment) {
+  const auto& assigned_live_ranges = allocation.assigned_live_ranges_;
+  const auto& assigned_values = allocation.assigned_values_;
+  const size_t num_assigned = assigned_live_ranges.size();
+  for (const HloValue* new_value : hlo_buffer.values()) {
+    const BufferAllocation::ValueLiveRange new_lr =
+        assignment->GetValueLiveRange(new_value->id());
+    DCHECK_GE(new_lr.start, 0)
+        << "Buffer doesn't have a proper live range:" << new_value->ToString();
+    for (size_t j = 0; j < num_assigned; ++j) {
+      DCHECK_GE(assigned_live_ranges[j].start, 0)
+          << "Buffer doesn't have a proper live range:"
+          << assigned_values[j]->ToString();
+      if (LiveRangeInterferes(new_value, new_lr, assigned_values[j],
+                              assigned_live_ranges[j], assignment)) {
+        VLOG(4) << "Can't assign: assignee " << *assigned_values[j]
+                << " live range interferes with " << new_value->ToShortString();
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
 absl::StatusOr<bool> BufferAssigner::MaybeAssignBuffer(
     BufferAllocation* allocation, const HloBuffer& hlo_buffer,
     BufferAssignment* assignment) {
-  CHECK(!assignment->HasAllocation(hlo_buffer))
+  DCHECK(!assignment->HasAllocation(hlo_buffer))
       << "buffer " << hlo_buffer << " already has an allocation assigned.";
 
-  VLOG(4) << "Trying to assign " << hlo_buffer << " size "
-          << assignment->HloBufferSize(hlo_buffer)
+  const int64_t buffer_size = assignment->HloBufferSize(hlo_buffer);
+  VLOG(4) << "Trying to assign " << hlo_buffer << " size " << buffer_size
           << " to allocation: " << *allocation;
 
-  ABSL_ASSIGN_OR_RETURN(auto buffer_color, hlo_buffer.color());
+  if (buffer_size > allocation->size()) {
+    VLOG(4) << "Can't assign: buffer is larger than allocation (" << buffer_size
+            << " > " << allocation->size() << ")";
+    return false;
+  }
+
+  if (allocation->is_readonly()) {
+    VLOG(4) << "Can't assign: allocation is readonly";
+    return false;
+  }
+
+  if (!allocation->is_reusable()) {
+    VLOG(4) << "Can't assign: allocation is not reusable";
+    return false;
+  }
+
+  const auto& new_values = hlo_buffer.values();
+  ABSL_ASSIGN_OR_RETURN(BufferValue::Color buffer_color, hlo_buffer.color());
 
   // Input/output allocations are backed by caller-provided or live-out storage,
   // so buffer assignment cannot strengthen their runtime allocation
@@ -2121,21 +2213,21 @@ absl::StatusOr<bool> BufferAssigner::MaybeAssignBuffer(
 
   // This is a storage-compatibility check, not recoloring: the buffer keeps its
   // HLO color even when it can use an allocation with a different color.
-  if (!opts_.can_use_allocation(buffer_color, allocation->color())) {
+  if (buffer_color != allocation->color() &&
+      !opts_.can_use_allocation(buffer_color, allocation->color())) {
     VLOG(4) << "Can't assign: buffer has color " << buffer_color
             << " and allocation has color " << allocation->color() << ".";
     return false;
   }
 
-  if (assignment->HloBufferSize(hlo_buffer) > allocation->size()) {
-    VLOG(4) << "Can't assign: buffer is larger than allocation ("
-            << assignment->HloBufferSize(hlo_buffer) << " > "
-            << allocation->size() << ")";
-    return false;
-  }
+  const bool buffer_lives_out = assignment->BufferLivesOut(hlo_buffer);
 
-  if (allocation->is_readonly()) {
-    VLOG(4) << "Can't assign: allocation is readonly";
+  // If the buffer is live out of the computation then it should only be
+  // assigned a buffer which exactly fits the result to avoid wasting memory
+  // (result buffers can have arbitrary lifetimes).
+  if (buffer_lives_out && allocation->size() != buffer_size) {
+    VLOG(4) << "Can't assign: buffer " << hlo_buffer
+            << "is live out and size not the same as allocation";
     return false;
   }
 
@@ -2143,88 +2235,37 @@ absl::StatusOr<bool> BufferAssigner::MaybeAssignBuffer(
     if (allocation->maybe_live_out()) {
       // If a buffer maybe live out, the allocation cannot contain any node
       // where must_not_live_out returns true.
-      for (const HloValue* value : hlo_buffer.values()) {
-        if ((*opts_.must_not_live_out)(assignment->alias_analysis(),
-                                       value->instruction(), value->index())) {
-          VLOG(4) << "Can't assign: " << value->instruction()->ToString()
-                  << " cannot live out of the module";
-          return false;
-        }
+      if (assignment->BufferMustNotLiveOut(hlo_buffer,
+                                           *opts_.must_not_live_out)) {
+        VLOG(4) << "Can't assign: buffer " << hlo_buffer
+                << " cannot live out of the module";
+        return false;
       }
     }
     // The above check is not enough -- There could be the case where an
     // allocation can be not live out and contains an instruction where
     // must_not_live_out_ returns true, but assigning a live out buffer to
     // that allocation makes the allocation live out and also contain an
-    // instruction where ust_not_live_out_ returns true.
-    if (assignment->alias_analysis().BufferLivesOut(hlo_buffer)) {
-      for (const auto& buffer_offset_size : allocation->assigned_buffers()) {
-        const HloValue* value = buffer_offset_size.first;
-        if ((*opts_.must_not_live_out)(assignment->alias_analysis(),
-                                       value->instruction(), value->index())) {
-          VLOG(4) << "Can't assign: " << value->instruction()
-                  << " cannot live out of the module";
-          return false;
-        }
+    // instruction where must_not_live_out_ returns true.
+    if (buffer_lives_out) {
+      if (assignment->AllocationHasMustNotLiveOut(allocation,
+                                                  *opts_.must_not_live_out)) {
+        VLOG(4) << "Can't assign: allocation contains a buffer that cannot "
+                   "live out of the module";
+        return false;
       }
     }
   }
 
-  if (!allocation->is_reusable()) {
-    VLOG(4) << "Can't assign: allocation is not reusable";
-    return false;
-  }
-
-  // Pre-compute and cache the live ranges for `hlo_buffer.values()`.
-  // Although hash map lookups are O(1), performing them inside the nested loops
-  // below results in O(N*M) lookups. Caching them in a contiguous vector
-  // reduces the total lookups to O(N+M) and significantly improves CPU cache
-  // locality in the hot inner loop.
-  std::vector<const HloLiveRange::LiveRangeBounds*> cached_new_live_ranges;
   if (assignment->hlo_live_range().total_order_scheduled()) {
-    const auto& buffer_live_ranges =
-        assignment->hlo_live_range().buffer_live_ranges();
-    cached_new_live_ranges.reserve(hlo_buffer.values().size());
-    for (const HloValue* new_value : hlo_buffer.values()) {
-      auto it = buffer_live_ranges.find(new_value);
-      CHECK(it != buffer_live_ranges.end())
-          << "Buffer doesn't have a proper live range:"
-          << new_value->ToString();
-      cached_new_live_ranges.push_back(&it->second);
+    if (HasLiveRangeInterference(*allocation, hlo_buffer, assignment)) {
+      return false;
     }
-  }
-
-  for (const auto& buffer_offset_size : allocation->assigned_buffers()) {
-    // Pairwise compare.
-    const HloValue& assigned_buffer =
-        *CHECK_NOTNULL(dynamic_cast<const HloValue*>(buffer_offset_size.first));
-
-    const HloLiveRange::LiveRangeBounds* assigned_live_range = nullptr;
-    if (assignment->hlo_live_range().total_order_scheduled()) {
-      const auto& buffer_live_ranges =
-          assignment->hlo_live_range().buffer_live_ranges();
-      auto it2 = buffer_live_ranges.find(&assigned_buffer);
-      CHECK(it2 != buffer_live_ranges.end())
-          << "Buffer doesn't have a proper live range:"
-          << assigned_buffer.ToString();
-      assigned_live_range = &it2->second;
-    }
-
-    for (size_t i = 0; i < hlo_buffer.values().size(); ++i) {
-      const HloValue* new_value = hlo_buffer.values()[i];
-      if (assignment->hlo_live_range().total_order_scheduled()) {
-        CHECK_NOTNULL(assigned_live_range);
-        if (LiveRangeInterferes(new_value, *cached_new_live_ranges[i],
-                                &assigned_buffer, *assigned_live_range,
-                                assignment)) {
-          VLOG(4) << "Can't assign: assignee " << assigned_buffer
-                  << " live range interferes with "
-                  << new_value->ToShortString();
-          return false;
-        }
-      } else {
+  } else {
+    for (const HloValue* assigned_buffer : allocation->assigned_values_) {
+      for (const HloValue* new_value : new_values) {
         if (assignment->hlo_ordering().MayInterfere(
-                assigned_buffer, *new_value, assignment->dataflow_analysis(),
+                *assigned_buffer, *new_value, assignment->dataflow_analysis(),
                 alias_info_)) {
           std::optional<bool> custom_interference;
           if (opts_.live_range_interference_options.has_value()) {
@@ -2232,28 +2273,32 @@ absl::StatusOr<bool> BufferAssigner::MaybeAssignBuffer(
             if (options.has_custom_interference(assignment->alias_analysis(),
                                                 *new_value) &&
                 options.has_custom_interference(assignment->alias_analysis(),
-                                                assigned_buffer)) {
+                                                *assigned_buffer)) {
               custom_interference = options.interferes(
-                  assignment->alias_analysis(), *new_value, assigned_buffer);
+                  assignment->alias_analysis(), *new_value, *assigned_buffer);
             }
           }
           if (custom_interference.value_or(true)) {
             // Fallback to partial order based interference detection (slower)
             // when we don't have a total order scheduled module.
-            VLOG(4) << "Can't assign: assignee " << assigned_buffer
+            VLOG(4) << "Can't assign: assignee " << *assigned_buffer
                     << " may interfere with " << new_value->ToShortString();
             return false;
           }
         }
       }
+    }
+  }
 
-      // Copy instruction don't share a buffer with their input operand.
-      if (new_value->instruction()->opcode() == HloOpcode::kCopy) {
+  // Copy instruction don't share a buffer with their input operand.
+  for (const HloValue* new_value : new_values) {
+    if (new_value->instruction()->opcode() == HloOpcode::kCopy) {
+      const HloInstruction* copy_operand = new_value->instruction()->operand(0);
+      for (const HloValue* assigned_buffer : allocation->assigned_values_) {
         for (const HloPosition& assigned_buffer_position :
-             assigned_buffer.positions()) {
-          if (new_value->instruction()->IsUserOf(
-                  assigned_buffer_position.instruction)) {
-            VLOG(4) << "Can't assign: assignee " << assigned_buffer
+             assigned_buffer->positions()) {
+          if (assigned_buffer_position.instruction == copy_operand) {
+            VLOG(4) << "Can't assign: assignee " << *assigned_buffer
                     << " is used at copy instruction "
                     << new_value->ToShortString();
             return false;
@@ -2281,8 +2326,7 @@ absl::StatusOr<bool> BufferAssigner::MaybeAssignBuffer(
     if (req_offset.has_value()) {
       target_offset = *req_offset;
       if (target_offset < 0 ||
-          target_offset + assignment->HloBufferSize(hlo_buffer) >
-              allocation->size()) {
+          target_offset + buffer_size > allocation->size()) {
         return false;
       }
       // NOLINTNEXTLINE(*-custom-deterministic-iteration-order)
@@ -2312,9 +2356,7 @@ absl::StatusOr<bool> BufferAssigner::MaybeAssignBuffer(
         }
       }
     } else {
-      // NOLINTNEXTLINE(*-custom-deterministic-iteration-order)
-      for (const auto& buffer_offset_size : allocation->assigned_buffers()) {
-        const HloValue* val = buffer_offset_size.first;
+      for (const HloValue* val : allocation->assigned_values_) {
         if ((*opts_.fixed_offset)(assignment->alias_analysis(),
                                   val->instruction(), val->index())
                 .has_value()) {
@@ -2326,19 +2368,8 @@ absl::StatusOr<bool> BufferAssigner::MaybeAssignBuffer(
     }
   }
 
-  // If the buffer is live out of the computation then it should only be
-  // assigned a buffer which exactly fits the result to avoid wasting memory
-  // (result buffers can have arbitrary lifetimes).
-  if (assignment->alias_analysis().BufferLivesOut(hlo_buffer) &&
-      allocation->size() != assignment->HloBufferSize(hlo_buffer)) {
-    VLOG(4) << "Can't assign: buffer " << hlo_buffer
-            << "is live out and size not the same as allocation";
-    return false;
-  }
-
-  ABSL_RETURN_IF_ERROR(
-      assignment->AddAssignment(allocation, hlo_buffer, target_offset,
-                                assignment->HloBufferSize(hlo_buffer)));
+  ABSL_RETURN_IF_ERROR(assignment->AddAssignment(allocation, hlo_buffer,
+                                            target_offset, buffer_size));
   return true;
 }
 
@@ -2436,7 +2467,7 @@ bool BufferAssigner::DelayTemporaryBufferAssignment(
         buffers_to_assign_sequentially,
     BufferAssignment* assignment) {
   if (assignment->HasAllocation(*hlo_buffer) ||
-      assignment->alias_analysis().BufferLivesOut(*hlo_buffer)) {
+      assignment->BufferLivesOut(*hlo_buffer)) {
     return false;
   }
 
@@ -2690,14 +2721,13 @@ absl::Status BufferAssigner::AssignBuffersForGlobalComputations(
       const HloBuffer* a, const HloBuffer* b)>;
   auto size_of = absl::bind_front(&BufferAssignment::HloBufferSize, assignment);
   Comparator compare_size = absl::bind_front(CompareSize, size_of);
-  Comparator compare_live_out =
-      absl::bind_front(CompareLiveOut, &alias_analysis);
+  Comparator compare_live_out = absl::bind_front(CompareLiveOut, assignment);
   Comparator compare_position =
       absl::bind_front(ComparePosition, &post_order_position);
-  absl::flat_hash_map<const HloBuffer*, int64_t> min_start_cache;
-  Comparator compare_live_range_start = absl::bind_front(
-      CompareLiveRangeStart, &assignment->hlo_live_range().buffer_live_ranges(),
-      &min_start_cache);
+  std::vector<int64_t> min_start_cache(assignment->num_buffer_ids(),
+                                       kUninitializedStartTime);
+  Comparator compare_live_range_start =
+      absl::bind_front(CompareLiveRangeStart, assignment, &min_start_cache);
   BufferOrder buffer_order = opts_.buffer_order;
   if (algorithm ==
       buffer_assignment::
@@ -2783,7 +2813,7 @@ absl::Status BufferAssigner::AssignBuffersForGlobalComputations(
   for (const HloBuffer* buffer : sorted_buffers) {
     VLOG(3) << "=================================================";
     VLOG(3) << "Assigning buffer for " << *buffer;
-    ABSL_ASSIGN_OR_RETURN(auto color, buffer->color());
+    ABSL_ASSIGN_OR_RETURN(BufferValue::Color color, buffer->color());
     BufferAllocationsManagerForComputationsWithoutOrdering* allocation_manager =
         get_manager(buffer, color);
 
@@ -2995,10 +3025,11 @@ absl::StatusOr<int64_t> BufferAssigner::ReuseCompatibleTempHeaps(
   // `other_value`. If so, they cannot occupy overlapping storage.
   auto candidate_interferes_with = [&](const ReuseCandidate& candidate,
                                        const HloValue* other_value) {
-    const auto& other_range = live_ranges.at(other_value);
+    const auto other_range = assignment->GetValueLiveRange(other_value->id());
     return absl::c_any_of(candidate.hlo_buffer->values(),
                           [&](const HloValue* value) {
-                            const auto& candidate_range = live_ranges.at(value);
+                            const auto candidate_range =
+                                assignment->GetValueLiveRange(value->id());
                             // LiveRangeInterferes allows touching operand/user
                             // ranges for exact in-place sharing. This best-fit
                             // placer does not preserve the operand's exact
@@ -3106,7 +3137,8 @@ absl::Status BufferAssigner::AssignPresetBuffers(
           << "No preset value allocation for color " << value->color()
           << " for " << value->ToShortString() << " found.";
       ABSL_RETURN_IF_ERROR(preset_allocations_iter->second->AddAssignment(
-          *value, chunk.offset, chunk.size));
+          *value, chunk.offset, chunk.size,
+          assignment->GetValueLiveRange(value->id())));
     }
 
     assigned_buffers->insert(&buffer);
diff --git a/third_party/xla/xla/service/buffer_assignment.h b/third_party/xla/xla/service/buffer_assignment.h
index 72ac832d24f..e8d98e0de1a 100644
--- a/third_party/xla/xla/service/buffer_assignment.h
+++ b/third_party/xla/xla/service/buffer_assignment.h
@@ -394,6 +394,14 @@ class BufferAllocation {
 
   void set_constant(bool is_constant) { is_constant_ = is_constant; }
 
+  struct ValueLiveRange {
+    int64_t start;
+    int64_t end;
+    constexpr ValueLiveRange() : start(-1), end(-1) {}
+    constexpr ValueLiveRange(int64_t start, int64_t end)
+        : start(start), end(end) {}
+  };
+
  private:
   // Only BufferAssigner and BufferAssignment can modify BufferAllocation.
   friend class BufferAssigner;
@@ -401,7 +409,7 @@ class BufferAllocation {
 
   // Adds a LogicalBuffer to the set assigned to this buffer.
   absl::Status AddAssignment(const HloValue& buffer, int64_t offset,
-                             int64_t size);
+                             int64_t size, ValueLiveRange live_range = {});
 
   void set_index(Index index) {
     index_ = index;
@@ -453,6 +461,9 @@ class BufferAllocation {
   // See comment on the is_constant() accessor.
   bool is_constant_ : 1;
 
+  bool has_must_not_live_out_value_ = false;
+  size_t must_not_live_out_checked_count_ = 0;
+
   // If this buffer is for an entry computation parameter, which subshape of the
   // parameter is it for?
   ShapeIndex param_shape_index_;
@@ -461,6 +472,11 @@ class BufferAllocation {
   // logical offsets and sizes.
   absl::flat_hash_map<const HloValue*, OffsetSize> assigned_buffers_;
 
+  // Parallel contiguous arrays of assigned HloValue* and their live range
+  // bounds for fast sequential scanning in BufferAssigner::MaybeAssignBuffer.
+  std::vector<const HloValue*> assigned_values_;
+  std::vector<ValueLiveRange> assigned_live_ranges_;
+
   int64_t fragmentation_bytes_ = 0;
   std::vector<HeapSimulatorTrace> heap_traces_;
 
@@ -530,6 +546,8 @@ class BufferAssignment {
   std::vector<BufferAllocation> TakeAllocations() && {
     for (auto& allocation : allocations_) {
       allocation.assigned_buffers_.clear();
+      allocation.assigned_values_.clear();
+      allocation.assigned_live_ranges_.clear();
       allocation.peak_buffers_.clear();
       allocation.cross_color_buffers_.clear();
     }
@@ -564,6 +582,7 @@ class BufferAssignment {
   // Returns the allocation with the given index. CHECKs if no allocation exists
   // with the given index.
   const BufferAllocation& GetAllocation(BufferAllocation::Index index) const;
+  BufferAllocation* GetMutableAllocation(BufferAllocation::Index index);
 
   // Returns the allocation with the given instruction and shape index. nullptr
   // if no allocation exists.
@@ -723,7 +742,9 @@ class BufferAssignment {
   void ClearAllocations() {
     allocations_.clear();
     temp_allocation_total_size_ = 0;
-    allocation_index_for_value_.clear();
+    std::fill(allocation_index_for_value_id_.begin(),
+              allocation_index_for_value_id_.end(), -1);
+    num_assigned_values_ = 0;
     fragmentation_bytes_by_color_.clear();
     stats_ = Stats();
   }
@@ -738,14 +759,44 @@ class BufferAssignment {
                                        LogicalBuffer::Color color);
 
   int64_t HloBufferSize(const HloBuffer& buffer) {
-    auto [it, inserted] = cached_buffer_sizes_.try_emplace(buffer.id());
-    if (inserted) {
-      it->second = buffer.ComputeSize(buffer_size_);
+    DCHECK_GE(buffer.id(), 0);
+    DCHECK_LT(static_cast<size_t>(buffer.id()), cached_buffer_sizes_.size());
+    int64_t& cached_size = cached_buffer_sizes_[buffer.id()];
+    if (cached_size == kUncomputedBufferSize) {
+      cached_size = buffer.ComputeSize(buffer_size_);
     }
-    return it->second;
+    return cached_size;
   }
 
+  BufferAllocation::ValueLiveRange GetValueLiveRange(
+      HloValue::Id value_id) const {
+    if (value_id >= 0 &&
+        static_cast<size_t>(value_id) < live_ranges_by_value_id_.size()) {
+      return live_ranges_by_value_id_[value_id];
+    }
+    return {};
+  }
+
+  bool BufferLivesOut(const HloBuffer& buffer) const {
+    DCHECK_GE(buffer.id(), 0);
+    DCHECK_LT(static_cast<size_t>(buffer.id()), buffer_lives_out_.size());
+    return buffer_lives_out_[buffer.id()] != 0;
+  }
+
+  bool ValueLivesOut(const HloValue& value) const {
+    DCHECK_GE(value.id(), 0);
+    DCHECK_LT(static_cast<size_t>(value.id()), value_lives_out_.size());
+    return value_lives_out_[value.id()] != 0;
+  }
+
+  size_t num_buffer_ids() const { return cached_buffer_sizes_.size(); }
+
  private:
+  static constexpr int64_t kUncomputedBufferSize = -2;
+  static constexpr int8_t kMustNotLiveOutUncomputed = -1;
+  static constexpr int8_t kMustNotLiveOutFalse = 0;
+  static constexpr int8_t kMustNotLiveOutTrue = 1;
+
   // Only BufferAssigner can build or modify BufferAssignments.
   friend class BufferAssigner;
   friend class DefaultBufferAllocationsManagerForComputationsWithoutOrdering;
@@ -769,6 +820,62 @@ class BufferAssignment {
     // -1 means no constraint.
     multiheap_size_constraint_per_heap_ =
         (raw_value == -1) ? UINT64_MAX : raw_value;
+    InitCaches();
+  }
+
+  void InitCaches();
+
+  using MustNotLiveOutFn = std::function<bool(
+      const HloAliasAnalysis&, const HloInstruction*, const ShapeIndex&)>;
+
+  bool ValueMustNotLiveOut(const HloValue& value,
+                           const MustNotLiveOutFn& pred) {
+    DCHECK_GE(value.id(), 0);
+    DCHECK_LT(static_cast<size_t>(value.id()), value_must_not_live_out_.size());
+    int8_t& cached = value_must_not_live_out_[value.id()];
+    if (cached == kMustNotLiveOutUncomputed) {
+      cached = pred(*alias_analysis_, value.instruction(), value.index())
+                   ? kMustNotLiveOutTrue
+                   : kMustNotLiveOutFalse;
+    }
+    return cached == kMustNotLiveOutTrue;
+  }
+
+  bool BufferMustNotLiveOut(const HloBuffer& buffer,
+                            const MustNotLiveOutFn& pred) {
+    DCHECK_GE(buffer.id(), 0);
+    DCHECK_LT(static_cast<size_t>(buffer.id()),
+              buffer_must_not_live_out_.size());
+    int8_t& cached = buffer_must_not_live_out_[buffer.id()];
+    if (cached == kMustNotLiveOutUncomputed) {
+      bool result = false;
+      for (const HloValue* value : buffer.values()) {
+        if (ValueMustNotLiveOut(*value, pred)) {
+          result = true;
+          break;
+        }
+      }
+      cached = result ? kMustNotLiveOutTrue : kMustNotLiveOutFalse;
+    }
+    return cached == kMustNotLiveOutTrue;
+  }
+
+  bool AllocationHasMustNotLiveOut(BufferAllocation* allocation,
+                                   const MustNotLiveOutFn& pred) {
+    if (allocation->has_must_not_live_out_value_) {
+      return true;
+    }
+    while (allocation->must_not_live_out_checked_count_ <
+           allocation->assigned_values_.size()) {
+      const HloValue* value =
+          allocation->assigned_values_
+              [allocation->must_not_live_out_checked_count_++];
+      if (ValueMustNotLiveOut(*value, pred)) {
+        allocation->has_must_not_live_out_value_ = true;
+        return true;
+      }
+    }
+    return false;
   }
 
   // Helper that calls NewEmptyAllocation and AddAssignment in one call,
@@ -788,7 +895,6 @@ class BufferAssignment {
 
   // Mutable accessors for allocations.
   BufferAllocation* GetMutableAssignedAllocation(const HloBuffer& buffer);
-  BufferAllocation* GetMutableAllocation(BufferAllocation::Index index);
 
   // Combines allocations of temporary buffers into one big BufferAllocation.
   absl::Status CombineTempAllocations(
@@ -812,9 +918,19 @@ class BufferAssignment {
 
   uint64_t multiheap_size_constraint_per_heap_;
 
-  // Maps Buffers to the index of the BufferAllocation which holds the buffer.
-  absl::flat_hash_map<const HloValue*, BufferAllocation::Index>
-      allocation_index_for_value_;
+  // Maps HloValue::Id to the index of the BufferAllocation which holds the
+  // buffer (-1 if unassigned).
+  std::vector<BufferAllocation::Index> allocation_index_for_value_id_;
+  size_t num_assigned_values_ = 0;
+
+  // Maps HloValue::Id to its live range bounds.
+  std::vector<BufferAllocation::ValueLiveRange> live_ranges_by_value_id_;
+
+  // Precomputed live-out flags and lazy must_not_live_out caches indexed by ID.
+  std::vector<uint8_t> buffer_lives_out_;
+  std::vector<uint8_t> value_lives_out_;
+  std::vector<int8_t> value_must_not_live_out_;
+  std::vector<int8_t> buffer_must_not_live_out_;
 
   // Points to the associated HloModule. The module is not owned by this class
   // and must outlive this BufferAssignment.
@@ -836,7 +952,7 @@ class BufferAssignment {
 
   absl::btree_map<LogicalBuffer::Color, int64_t> fragmentation_bytes_by_color_;
 
-  absl::flat_hash_map<HloBuffer::Id, int64_t> cached_buffer_sizes_;
+  std::vector<int64_t> cached_buffer_sizes_;
 
   BufferAssignment(const BufferAssignment&) = delete;
   BufferAssignment& operator=(const BufferAssignment&) = delete;
@@ -1119,13 +1235,20 @@ class BufferAssigner {
       const std::vector<const HloComputation*>& computations,
       BufferAssignment* assignment);
 
-  // Returns true if buffer's live range interferences with buffer2's.
+  // Returns true if buffer1's live range interferes with buffer2's.
   bool LiveRangeInterferes(const HloValue* buffer1,
-                           const HloLiveRange::LiveRangeBounds& live_range1,
+                           BufferAllocation::ValueLiveRange live_range1,
                            const HloValue* buffer2,
-                           const HloLiveRange::LiveRangeBounds& live_range2,
+                           BufferAllocation::ValueLiveRange live_range2,
                            BufferAssignment* assignment);
 
+  // Returns true if any value in `hlo_buffer` has a live range that interferes
+  // with a value already assigned to `allocation`. Requires a total-order
+  // scheduled module.
+  bool HasLiveRangeInterference(const BufferAllocation& allocation,
+                                const HloBuffer& hlo_buffer,
+                                BufferAssignment* assignment);
+
   // Assigns pre-set assignments, if provided. These assignments will be added
   // to assigned_buffers and skip buffer allocation.
   absl::Status AssignPresetBuffers(
diff --git a/third_party/xla/xla/service/buffer_assignment_test.cc b/third_party/xla/xla/service/buffer_assignment_test.cc
index bba5c535cdb..3466c68b892 100644
--- a/third_party/xla/xla/service/buffer_assignment_test.cc
+++ b/third_party/xla/xla/service/buffer_assignment_test.cc
@@ -5865,6 +5865,128 @@ void BM_FastMergeManagerStress(::testing::benchmark::State& state) {
 
 BENCHMARK(BM_FastMergeManagerStress)->Range(1'000, 10'000'000);
 
+void BM_MaybeAssignBufferStress(::testing::benchmark::State& state) {
+  const int64_t num_buffers = state.range(0);
+  HloModuleConfig config;
+  config.set_debug_options(GetDebugOptionsFromFlags());
+  HloModule module("BM_MaybeAssignBufferStress", config);
+  auto builder = HloComputation::Builder("entry");
+
+  Shape shape = ShapeUtil::MakeShape(F32, {16});
+  auto param0 =
+      builder.AddInstruction(HloInstruction::CreateParameter(0, shape, "p"));
+  std::vector<HloInstruction*> insts;
+  insts.reserve(num_buffers);
+  insts.push_back(param0);
+
+  std::vector<HloInstruction*> root_elements;
+  for (int64_t i = 1; i < num_buffers; ++i) {
+    HloOpcode opcode = (i % 4 == 0) ? HloOpcode::kCopy : HloOpcode::kNegate;
+    HloInstruction* inst = builder.AddInstruction(
+        HloInstruction::CreateUnary(shape, opcode, insts.back()));
+    insts.push_back(inst);
+    if (i % 64 == 0) {
+      root_elements.push_back(inst);
+    }
+  }
+  if (root_elements.empty()) {
+    root_elements.push_back(insts.back());
+  }
+  HloInstruction* root =
+      builder.AddInstruction(HloInstruction::CreateTuple(root_elements));
+  insts.push_back(root);
+  HloComputation* entry = module.AddEntryComputation(builder.Build());
+
+  HloSchedule schedule(&module);
+  schedule.set_sequence(entry, insts);
+  CHECK_OK(module.set_schedule(schedule));
+
+  AliasInfo alias_info;
+  BufferAssigner::Options opts;
+  opts.assignment_algorithm_for_computations_without_ordering =
+      buffer_assignment::
+          AssignmentAlgorithmForComputationsWithoutOrderingProto::FAST_MERGE;
+  opts.buffer_assignment_algorithm =
+      buffer_assignment::BufferAssignmentAlgorithmProto::FAST_MERGE;
+  opts.color_memory_limit = [](LogicalBuffer::Color) -> int64_t {
+    return int64_t{1} << 40;
+  };
+  opts.must_not_live_out = [](const HloAliasAnalysis&,
+                              const HloInstruction* instr, const ShapeIndex&) {
+    return instr->opcode() == HloOpcode::kNegate &&
+           (instr->unique_id() % 5 == 0);
+  };
+  BufferAssigner assigner(&alias_info, std::move(opts));
+
+  ASSERT_OK_AND_ASSIGN(
+      auto base_assignment,
+      assigner.CreateAssignment(
+          &module, std::make_unique<SequentialHloOrdering>(schedule),
+          &BufferSizeBytes, [](LogicalBuffer::Color) { return 1; }));
+
+  std::vector<const HloBuffer*> buffers;
+  buffers.reserve(num_buffers);
+  for (int64_t i = 0; i < num_buffers; ++i) {
+    buffers.push_back(
+        &base_assignment->alias_analysis().GetBufferContainingValue(
+            base_assignment->dataflow_analysis().GetUniqueValueAt(insts[i],
+                                                                  {})));
+  }
+
+  constexpr int kWindowSize = 8;
+  constexpr int kMaxAssignedPerAlloc = 32;
+
+  for (auto s : state) {
+    base_assignment->ClearAllocations();
+    std::vector<BufferAllocation::Index> active_allocs;
+    active_allocs.reserve(kWindowSize);
+    for (int k = 0; k < kWindowSize; ++k) {
+      BufferAllocation* alloc = base_assignment->NewEmptyAllocation(
+          (k % 2 == 0) ? 64 : 128, LogicalBuffer::Color(0));
+      if (k < kWindowSize / 2) {
+        alloc->set_maybe_live_out(true);
+      }
+      active_allocs.push_back(alloc->index());
+    }
+
+    int replace_idx = 0;
+    for (int64_t i = 0; i < num_buffers; ++i) {
+      const HloBuffer* buffer = buffers[i];
+      bool assigned = false;
+      for (int k = 0; k < kWindowSize; ++k) {
+        BufferAllocation::Index idx = active_allocs[(i + k) % kWindowSize];
+        BufferAllocation* alloc = base_assignment->GetMutableAllocation(idx);
+        auto status_or =
+            assigner.MaybeAssignBuffer(alloc, *buffer, base_assignment.get());
+        if (status_or.ok() && status_or.value()) {
+          assigned = true;
+          if (alloc->assigned_buffers().size() >= kMaxAssignedPerAlloc) {
+            BufferAllocation* new_alloc = base_assignment->NewEmptyAllocation(
+                ((i + k) % 2 == 0) ? 64 : 128, LogicalBuffer::Color(0));
+            if (((i + k) % kWindowSize) < kWindowSize / 2) {
+              new_alloc->set_maybe_live_out(true);
+            }
+            active_allocs[(i + k) % kWindowSize] = new_alloc->index();
+          }
+          break;
+        }
+      }
+      if (!assigned) {
+        BufferAllocation* new_alloc =
+            base_assignment->NewEmptyAllocation(64, LogicalBuffer::Color(0));
+        active_allocs[replace_idx] = new_alloc->index();
+        replace_idx = (replace_idx + 1) % kWindowSize;
+        CHECK_OK(
+            assigner
+                .MaybeAssignBuffer(new_alloc, *buffer, base_assignment.get())
+                .status());
+      }
+    }
+  }
+}
+
+BENCHMARK(BM_MaybeAssignBufferStress)->Range(1'000, 1'000'000);
+
 // Tests that BufferAssignment::FromProto rejects a malformed proto containing
 // an assigned buffer with a negative offset or size.
 //
diff --git a/third_party/xla/xla/service/hlo_buffer.cc b/third_party/xla/xla/service/hlo_buffer.cc
index 7fe0f504182..fc5da34e88a 100644
--- a/third_party/xla/xla/service/hlo_buffer.cc
+++ b/third_party/xla/xla/service/hlo_buffer.cc
@@ -24,7 +24,10 @@ limitations under the License.
 #include "absl/algorithm/container.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
 #include "absl/strings/str_join.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_opcode.h"
@@ -33,6 +36,26 @@ limitations under the License.
 
 namespace xla {
 
+absl::StatusOr<BufferValue::Color> HloBuffer::ComputeColorSlowPath() const {
+  // Invariant: All values in the buffer should have the same color.
+  BufferValue::Color result = values()[0]->color();
+  for (const HloValue* value : values()) {
+    if (result != value->color()) {
+      std::string details = absl::StrFormat(
+          "Not all HloValues in the HloBuffer have the same color. "
+          "Buffer id=%d has %d values:",
+          id(), values().size());
+      for (const HloValue* v : values()) {
+        absl::StrAppendFormat(&details, "\n  value %d color=%d defined at %s",
+                              v->id(), v->color(),
+                              v->defining_position().ToString());
+      }
+      return absl::FailedPreconditionError(details);
+    }
+  }
+  return result;
+}
+
 bool HloBuffer::operator==(const HloBuffer& other) const {
   bool equal = id() == other.id();
   if (equal) {
diff --git a/third_party/xla/xla/service/hlo_buffer.h b/third_party/xla/xla/service/hlo_buffer.h
index 4be60c9097d..3cbafc883b2 100644
--- a/third_party/xla/xla/service/hlo_buffer.h
+++ b/third_party/xla/xla/service/hlo_buffer.h
@@ -22,11 +22,10 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/base/attributes.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/log/check.h"
-#include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "absl/strings/str_format.h"
 #include "xla/service/buffer_value.h"
 #include "xla/service/hlo_value.h"
 #include "xla/xla_data.pb.h"
@@ -119,24 +118,12 @@ class HloBuffer {
 
   // Memory space color. Used to indicate the memory space that the hlo buffer
   // needs to live in.
-  absl::StatusOr<BufferValue::Color> color() const {
-    // Invariant: All values in the buffer should have the same color.
-    BufferValue::Color result = values()[0]->color();
-    for (const HloValue* value : values()) {
-      if (result != value->color()) {
-        std::string details = absl::StrFormat(
-            "Not all HloValues in the HloBuffer have the same color. "
-            "Buffer id=%d has %d values:",
-            id(), values().size());
-        for (const HloValue* v : values()) {
-          absl::StrAppendFormat(&details, "\n  value %d color=%d defined at %s",
-                                v->id(), v->color(),
-                                v->defining_position().ToString());
-        }
-        return absl::FailedPreconditionError(details);
-      }
+  ABSL_ATTRIBUTE_ALWAYS_INLINE absl::StatusOr<BufferValue::Color> color()
+      const {
+    if (values_.size() == 1) {
+      return values_[0]->color();
     }
-    return result;
+    return ComputeColorSlowPath();
   }
 
   // Return the unique HLO value in the buffer. CHECK fails if the buffer does
@@ -160,6 +147,8 @@ class HloBuffer {
   bool operator!=(const HloBuffer& other) const { return !(*this == other); }
 
  private:
+  absl::StatusOr<BufferValue::Color> ComputeColorSlowPath() const;
+
   // Unique identifier for this HloBuffer.
   Id id_;
 

From 1b4aaea740148b37f45ebb7deb94aa1c651f070f Mon Sep 17 00:00:00 2001
From: Louis Lechevalier <louis@zml.ai>
Date: Wed, 14 Oct 2026 05:59:15 -0700
Subject: [PATCH 0764/1000] PR #50654: [XLA:GPU] Set fusion kind to kCustom
 when applying a cuDNN config
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50654

📝 Summary of Changes

`CudnnBackend::ApplyConfig` sets fusion kind to kCustom.

🎯 Justification

Follow-up to [#49686](https://github.com/openxla/xla/pull/49686), and based on the review, the cuDNN backend must still be a candidate for kLoop fusions. When cuDNN backend is candidate for non-kCustom fusion (e.g. tiny dot which is kLoop `--xla_gpu_cudnn_gemm_fusion_level=2`).
The identified issue was that the `CuDnnFusionCompiler` adds a workspace output and stores the graph under the computation's fingerprint. The hloDCE only skips kCustom fusions and removes the workspace. `CuDnnThunk::Initialize` threw std::out_of_range because the fingerprints didn't match anymore.

🚀 Kind of Contribution

🐛 Bug Fix, 🧪 Tests

🧪 Unit Tests:

Add `CudnnBackendTest.ApplyConfigToLoopFusionSetsCustomFusionKind` in `xla/backends/gpu/autotuner/cudnn_test.cc`: it applies a cuDNN config to a kLoop dot and checks that the fusion kind is kCustom.

Copybara import of the project:

--
2b311c5e80c24960e57c849ce10917516471b148 by Louis Lechevalier <louis@zml.ai>:

[XLA:GPU] Change fusion kind to kCustom when using the Cudnn backend

Merging this change closes #50654

PiperOrigin-RevId: 999314473
---
 .../xla/xla/backends/gpu/autotuner/cudnn.cc   |  1 +
 .../xla/backends/gpu/autotuner/cudnn_test.cc  | 33 +++++++++++++++++++
 2 files changed, 34 insertions(+)

diff --git a/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc b/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc
index 7689c859e09..baf5f977d07 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc
@@ -351,6 +351,7 @@ absl::Status ApplyConfigToCudnnFusion(HloInstruction& instr,
   backend_config->set_kind(kCuDnnFusionKind);
   backend_config->mutable_cudnn_fusion_config()->set_plan_id(config.algo_id());
   ABSL_RETURN_IF_ERROR(instr.set_backend_config(std::move(gpu_config)));
+  instr.set_fusion_kind(HloInstruction::FusionKind::kCustom);
   return absl::OkStatus();
 }
 
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 209154e9491..c47ee79b83d 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/cudnn_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/cudnn_test.cc
@@ -122,6 +122,21 @@ absl::string_view kTritonGemmFusionHlo = R"hlo(
       backend_config={"fusion_backend_config": {kind: "__triton_gemm"}}
   })hlo";
 
+absl::string_view kLoopDotFusionHlo = R"hlo(
+  wrapped_dot_computation {
+    p0 = f32[4,4] parameter(0)
+    p1 = f32[4,4] parameter(1)
+    ROOT d = f32[4,4] dot(p0, p1),
+      lhs_contracting_dims={1}, rhs_contracting_dims={0}
+  }
+
+  e {
+    p0 = f32[4,4] parameter(0)
+    p1 = f32[4,4] parameter(1)
+    ROOT wrapped_dot = f32[4,4] fusion(p0, p1), kind=kLoop,
+      calls=wrapped_dot_computation
+  })hlo";
+
 absl::string_view kF64GemmFusionHlo = R"hlo(
   fusion1 {
     p0 = f64[3,28,32] parameter(0)
@@ -409,6 +424,24 @@ TEST_F(CudnnBackendTest, ApplyConfigToTritonGemmFusionSetsCudnnKind) {
             1);
 }
 
+TEST_F(CudnnBackendTest, ApplyConfigToLoopFusionSetsCustomFusionKind) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> hlo_module,
+                       ParseAndReturnVerifiedModule(kLoopDotFusionHlo));
+  CudnnBackendConfig config;
+  config.set_algo_id(1);
+  HloInstruction* fusion_instr =
+      hlo_module->entry_computation()->root_instruction();
+  BackendConfig backend_config;
+  *backend_config.mutable_algorithm() = config;
+  ASSERT_OK(backend_->ApplyConfig(*fusion_instr, backend_config));
+  EXPECT_EQ(fusion_instr->fusion_kind(), HloInstruction::FusionKind::kCustom);
+  ASSERT_OK_AND_ASSIGN(GpuBackendConfig gpu_config,
+                       fusion_instr->backend_config<GpuBackendConfig>());
+  EXPECT_EQ(gpu_config.fusion_backend_config().kind(), kCuDnnFusionKind);
+  EXPECT_EQ(gpu_config.fusion_backend_config().cudnn_fusion_config().plan_id(),
+            1);
+}
+
 TEST_F(CudnnBackendTest, ApplyConfigToCudnnCustomCall) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> hlo_module,
                        ParseAndReturnVerifiedModule(kCudnnCustomCallHlo));

From 5d2cb438ead567048eacab328248fe14535a6456 Mon Sep 17 00:00:00 2001
From: Oleg Shyshkov <shyshkov@google.com>
Date: Wed, 14 Oct 2026 06:19:31 -0700
Subject: [PATCH 0765/1000] [XLA:GPU] Extract GetCollectiveInfo helper function
 into utils.

And use in CollectiveKernelStrategyAnnotator.

PiperOrigin-RevId: 999324234
---
 .../xla/xla/backends/gpu/codegen/triton/BUILD |   1 +
 .../gpu/codegen/triton/collective_emitter.cc  |  47 +----
 .../xla/xla/backends/gpu/runtime/BUILD        |  40 ++++
 .../gpu/runtime/collective_info_utils.cc      |  54 ++++++
 .../gpu/runtime/collective_info_utils.h       |  42 ++++
 .../gpu/runtime/collective_info_utils_test.cc | 180 ++++++++++++++++++
 .../backends/gpu/transforms/collectives/BUILD |   2 +
 .../collective_kernel_strategy_annotator.cc   | 159 ++++------------
 8 files changed, 364 insertions(+), 161 deletions(-)
 create mode 100644 third_party/xla/xla/backends/gpu/runtime/collective_info_utils.cc
 create mode 100644 third_party/xla/xla/backends/gpu/runtime/collective_info_utils.h
 create mode 100644 third_party/xla/xla/backends/gpu/runtime/collective_info_utils_test.cc

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/BUILD b/third_party/xla/xla/backends/gpu/codegen/triton/BUILD
index af53f9bc048..2af3c3b8b7c 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/BUILD
@@ -759,6 +759,7 @@ cc_library(
         "//xla/backends/gpu/codegen/triton/ir:triton_xla",
         "//xla/backends/gpu/runtime:all_gather",
         "//xla/backends/gpu/runtime:all_reduce",
+        "//xla/backends/gpu/runtime:collective_info_utils",
         "//xla/backends/gpu/runtime:collective_params",
         "//xla/backends/gpu/runtime:reduce_scatter",
         "//xla/backends/gpu/transforms/collectives:collective_ops_utils",
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
index 403eeb6e53d..0d911ef8d2e 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
@@ -64,6 +64,7 @@ limitations under the License.
 #include "xla/backends/gpu/codegen/triton/lowering_util.h"
 #include "xla/backends/gpu/runtime/all_gather.h"
 #include "xla/backends/gpu/runtime/all_reduce.h"
+#include "xla/backends/gpu/runtime/collective_info_utils.h"
 #include "xla/backends/gpu/runtime/collective_params.h"
 #include "xla/backends/gpu/runtime/reduce_scatter.h"
 #include "xla/backends/gpu/transforms/collectives/collective_ops_utils.h"
@@ -300,47 +301,9 @@ absl::StatusOr<ReduceScatterEmitterContext> CreateReduceScatterEmitterContext(
   return ctx;
 }
 
-using InfoStruct =
-    std::variant<AllReduceInfo, AllGatherInfo, ReduceScatterInfo>;
-
-// Builds helper struct for the given collective instruction.
-absl::StatusOr<InfoStruct> GetCollectiveInfo(
-    const HloInstruction* instr, const GpuTopology& gpu_topology,
-    const DeviceAssignment* device_assignment) {
-  switch (instr->opcode()) {
-    case HloOpcode::kAllReduce: {
-      const HloAllReduceInstruction* all_reduce =
-          Cast<HloAllReduceInstruction>(instr);
-      ABSL_ASSIGN_OR_RETURN(AllReduceInfo all_reduce_info,
-                       BuildAllReduceInfo(
-                           /*is_multimem_enabled=*/false, gpu_topology,
-                           all_reduce, device_assignment));
-      return all_reduce_info;
-    }
-    case HloOpcode::kAllGather: {
-      const HloAllGatherInstruction* all_gather =
-          Cast<HloAllGatherInstruction>(instr);
-      ABSL_ASSIGN_OR_RETURN(
-          AllGatherInfo all_gather_info,
-          BuildAllGatherInfo(gpu_topology, all_gather, device_assignment));
-      return all_gather_info;
-    }
-    case HloOpcode::kReduceScatter: {
-      ABSL_ASSIGN_OR_RETURN(
-          ReduceScatterInfo reduce_scatter_info,
-          BuildReduceScatterInfo(gpu_topology,
-                                 Cast<HloReduceScatterInstruction>(instr),
-                                 device_assignment));
-      return reduce_scatter_info;
-    }
-    default:
-      return absl::InvalidArgumentError(
-          absl::StrCat("Unsupported collective opcode: ", instr->opcode()));
-  }
-}
-
 absl::StatusOr<LaunchDimensions> GetLaunchDimensions(
-    const InfoStruct& collective_info, const GpuTopology& gpu_topology) {
+    const CollectiveKernelInfo& collective_info,
+    const GpuTopology& gpu_topology) {
   const se::DeviceDescription& device_info =
       gpu_topology.gpu_target_config().device_description;
   return std::visit(
@@ -365,7 +328,7 @@ absl::StatusOr<LaunchDimensions> GetLaunchDimensions(
 
 absl::Status ValidateBlockLevelFusionConfig(
     const BlockLevelFusionConfig& block_level_config,
-    const InfoStruct& collective_info) {
+    const CollectiveKernelInfo& collective_info) {
   TF_RET_CHECK(block_level_config.output_tiles_size() == 1)
       << "expected 1 output tile, but got "
       << block_level_config.output_tiles_size();
@@ -1358,7 +1321,7 @@ absl::StatusOr<BlockLevelFusionConfig> GetCollectiveBlockLevelFusionConfig(
         "Triton codegen. %s",
         instr->ToString()));
   }
-  ABSL_ASSIGN_OR_RETURN(InfoStruct collective_info,
+  ABSL_ASSIGN_OR_RETURN(CollectiveKernelInfo collective_info,
                    GetCollectiveInfo(instr, gpu_topology, device_assignment));
   ABSL_ASSIGN_OR_RETURN(const LaunchDimensions launch_dims,
                    GetLaunchDimensions(collective_info, gpu_topology));
diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index 3ccc68e99d3..700e94ea739 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -4245,6 +4245,46 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "collective_info_utils",
+    srcs = ["collective_info_utils.cc"],
+    hdrs = ["collective_info_utils.h"],
+    deps = [
+        ":all_gather",
+        ":all_reduce",
+        ":reduce_scatter",
+        "//xla/hlo/ir:hlo",
+        "//xla/service:device_assignment",
+        "//xla/service:gpu_topology",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+xla_cc_test(
+    name = "collective_info_utils_test",
+    srcs = ["collective_info_utils_test.cc"],
+    deps = [
+        ":all_gather",
+        ":all_reduce",
+        ":collective_info_utils",
+        ":reduce_scatter",
+        "//xla/backends/gpu/target_config",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+        "//xla/service:gpu_topology",
+        "//xla/service/gpu:gpu_device_info_for_tests",
+        "//xla/stream_executor:device_description",
+        "//xla/tsl/platform:test_main",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest",
+    ],
+)
+
 xla_cc_test(
     name = "all_reduce_build_info_test",
     srcs = ["all_reduce_build_info_test.cc"],
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_info_utils.cc b/third_party/xla/xla/backends/gpu/runtime/collective_info_utils.cc
new file mode 100644
index 00000000000..bf17b1e9a4d
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_info_utils.cc
@@ -0,0 +1,54 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/gpu/runtime/collective_info_utils.h"
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "xla/backends/gpu/runtime/all_gather.h"
+#include "xla/backends/gpu/runtime/all_reduce.h"
+#include "xla/backends/gpu/runtime/reduce_scatter.h"
+#include "xla/hlo/ir/hlo_casting_utils.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/service/device_assignment.h"
+#include "xla/service/gpu_topology.h"
+
+namespace xla::gpu {
+
+absl::StatusOr<CollectiveKernelInfo> GetCollectiveInfo(
+    const HloInstruction* instr, const GpuTopology& gpu_topology,
+    const DeviceAssignment* device_assignment, bool is_multimem_enabled) {
+  if (const auto* all_reduce = DynCast<HloAllReduceInstruction>(instr)) {
+    return BuildAllReduceInfo(is_multimem_enabled, gpu_topology, all_reduce,
+                              device_assignment);
+  }
+
+  if (const auto* all_gather = DynCast<HloAllGatherInstruction>(instr)) {
+    return BuildAllGatherInfo(gpu_topology, all_gather, device_assignment);
+  }
+
+  if (const auto* reduce_scatter =
+          DynCast<HloReduceScatterInstruction>(instr)) {
+    return BuildReduceScatterInfo(gpu_topology, reduce_scatter,
+                                  device_assignment);
+  }
+
+  return absl::InvalidArgumentError(
+      absl::StrCat("Unsupported collective opcode: ", instr->opcode()));
+}
+
+}  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_info_utils.h b/third_party/xla/xla/backends/gpu/runtime/collective_info_utils.h
new file mode 100644
index 00000000000..6146b9051d3
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_info_utils.h
@@ -0,0 +1,42 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_BACKENDS_GPU_RUNTIME_COLLECTIVE_INFO_UTILS_H_
+#define XLA_BACKENDS_GPU_RUNTIME_COLLECTIVE_INFO_UTILS_H_
+
+#include <variant>
+
+#include "absl/status/statusor.h"
+#include "xla/backends/gpu/runtime/all_gather.h"
+#include "xla/backends/gpu/runtime/all_reduce.h"
+#include "xla/backends/gpu/runtime/reduce_scatter.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/service/device_assignment.h"
+#include "xla/service/gpu_topology.h"
+
+namespace xla::gpu {
+
+using CollectiveKernelInfo =
+    std::variant<AllReduceInfo, AllGatherInfo, ReduceScatterInfo>;
+
+// Builds helper struct for the given collective instruction.
+absl::StatusOr<CollectiveKernelInfo> GetCollectiveInfo(
+    const HloInstruction* instr, const GpuTopology& gpu_topology,
+    const DeviceAssignment* device_assignment = nullptr,
+    bool is_multimem_enabled = false);
+
+}  // namespace xla::gpu
+
+#endif  // XLA_BACKENDS_GPU_RUNTIME_COLLECTIVE_INFO_UTILS_H_
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_info_utils_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_info_utils_test.cc
new file mode 100644
index 00000000000..f13d0e4140a
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_info_utils_test.cc
@@ -0,0 +1,180 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/gpu/runtime/collective_info_utils.h"
+
+#include <memory>
+#include <variant>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/log/check.h"
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "absl/strings/string_view.h"
+#include "xla/backends/gpu/runtime/all_gather.h"
+#include "xla/backends/gpu/runtime/all_reduce.h"
+#include "xla/backends/gpu/runtime/reduce_scatter.h"
+#include "xla/backends/gpu/target_config/target_config.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
+#include "xla/service/gpu/gpu_device_info_for_tests.h"
+#include "xla/service/gpu_topology.h"
+#include "xla/stream_executor/device_description.h"
+
+namespace xla::gpu {
+namespace {
+
+using ::absl_testing::StatusIs;
+
+class CollectiveInfoUtilsTest : public HloHardwareIndependentTestBase {
+ protected:
+  GpuTopology CreateGpuTopology() {
+    se::DeviceDescription device_info = TestGpuDeviceInfo::H100SXMDeviceInfo();
+    stream_executor::GpuTargetConfigProto target_config_proto;
+    *target_config_proto.mutable_gpu_device_info() = device_info.ToProto();
+    target_config_proto.mutable_gpu_device_info()
+        ->mutable_device_interconnect_info()
+        ->set_active_links(18);
+    target_config_proto.set_platform_name("CUDA");
+    auto target_config_or =
+        gpu::GpuTargetConfig::FromProto(target_config_proto);
+    CHECK_OK(target_config_or.status());
+    return GpuTopology("platform_version", /*num_partitions=*/1,
+                       /*num_hosts_per_partition=*/1,
+                       /*num_devices_per_host=*/8, *target_config_or);
+  }
+};
+
+TEST_F(CollectiveInfoUtilsTest, HandlesAllReduce) {
+  constexpr absl::string_view kModuleStr = R"(
+  HloModule test
+  apply_op {
+    x = s32[] parameter(0)
+    y = s32[] parameter(1)
+    ROOT apply_op = s32[] add(x, y)
+  }
+  ENTRY test_computation {
+    param_0 = s32[128] parameter(0)
+    ROOT all-reduce = s32[128] all-reduce(param_0), to_apply=apply_op,
+        replica_groups={{0,1}}
+  }
+  )";
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<HloModule> module,
+      ParseAndReturnVerifiedModule(kModuleStr, /*replica_count=*/2));
+  const HloInstruction* instr =
+      FindInstruction(module.get(), HloOpcode::kAllReduce);
+  GpuTopology gpu_topology = CreateGpuTopology();
+
+  auto info_or = GetCollectiveInfo(instr, gpu_topology);
+  ASSERT_OK(info_or);
+  EXPECT_TRUE(std::holds_alternative<AllReduceInfo>(*info_or));
+}
+
+TEST_F(CollectiveInfoUtilsTest, HandlesAllReduceStart) {
+  constexpr absl::string_view kModuleStr = R"(
+  HloModule test
+  apply_op {
+    x = s32[] parameter(0)
+    y = s32[] parameter(1)
+    ROOT apply_op = s32[] add(x, y)
+  }
+  ENTRY test_computation {
+    param_0 = s32[128] parameter(0)
+    start = s32[128] all-reduce-start(param_0), to_apply=apply_op,
+        replica_groups={{0,1}}
+    ROOT done = s32[128] all-reduce-done(start)
+  }
+  )";
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<HloModule> module,
+      ParseAndReturnVerifiedModule(kModuleStr, /*replica_count=*/2));
+  const HloInstruction* instr =
+      FindInstruction(module.get(), HloOpcode::kAllReduceStart);
+  GpuTopology gpu_topology = CreateGpuTopology();
+
+  auto info_or = GetCollectiveInfo(instr, gpu_topology);
+  ASSERT_OK(info_or);
+  EXPECT_TRUE(std::holds_alternative<AllReduceInfo>(*info_or));
+}
+
+TEST_F(CollectiveInfoUtilsTest, HandlesAllGather) {
+  constexpr absl::string_view kModuleStr = R"(
+  HloModule test
+  ENTRY test_computation {
+    param_0 = s32[128] parameter(0)
+    ROOT all-gather = s32[256] all-gather(param_0),
+        dimensions={0}, replica_groups={{0,1}}
+  }
+  )";
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<HloModule> module,
+      ParseAndReturnVerifiedModule(kModuleStr, /*replica_count=*/2));
+  const HloInstruction* instr =
+      FindInstruction(module.get(), HloOpcode::kAllGather);
+  GpuTopology gpu_topology = CreateGpuTopology();
+
+  auto info_or = GetCollectiveInfo(instr, gpu_topology);
+  ASSERT_OK(info_or);
+  EXPECT_TRUE(std::holds_alternative<AllGatherInfo>(*info_or));
+}
+
+TEST_F(CollectiveInfoUtilsTest, HandlesReduceScatter) {
+  constexpr absl::string_view kModuleStr = R"(
+  HloModule test
+  apply_op {
+    x = f32[] parameter(0)
+    y = f32[] parameter(1)
+    ROOT apply_op = f32[] add(x, y)
+  }
+  ENTRY test_computation {
+    param_0 = f32[256] parameter(0)
+    ROOT rs = f32[128] reduce-scatter(param_0),
+        dimensions={0}, replica_groups={{0,1}}, to_apply=apply_op
+  }
+  )";
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<HloModule> module,
+      ParseAndReturnVerifiedModule(kModuleStr, /*replica_count=*/2));
+  const HloInstruction* instr =
+      FindInstruction(module.get(), HloOpcode::kReduceScatter);
+  GpuTopology gpu_topology = CreateGpuTopology();
+
+  auto info_or = GetCollectiveInfo(instr, gpu_topology);
+  ASSERT_OK(info_or);
+  EXPECT_TRUE(std::holds_alternative<ReduceScatterInfo>(*info_or));
+}
+
+TEST_F(CollectiveInfoUtilsTest, RejectsUnsupportedInstruction) {
+  constexpr absl::string_view kModuleStr = R"(
+  HloModule test
+  ENTRY test_computation {
+    param_0 = s32[128] parameter(0)
+    ROOT add = s32[128] add(param_0, param_0)
+  }
+  )";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kModuleStr));
+  const HloInstruction* instr = FindInstruction(module.get(), HloOpcode::kAdd);
+  GpuTopology gpu_topology = CreateGpuTopology();
+
+  EXPECT_THAT(GetCollectiveInfo(instr, gpu_topology),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+}  // namespace
+}  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/BUILD b/third_party/xla/xla/backends/gpu/transforms/collectives/BUILD
index 55ea7b921f5..cf7a911f638 100644
--- a/third_party/xla/xla/backends/gpu/transforms/collectives/BUILD
+++ b/third_party/xla/xla/backends/gpu/transforms/collectives/BUILD
@@ -861,6 +861,7 @@ cc_library(
         "//xla:xla_proto_cc",
         "//xla/backends/gpu/runtime:all_gather",
         "//xla/backends/gpu/runtime:all_reduce",
+        "//xla/backends/gpu/runtime:collective_info_utils",
         "//xla/backends/gpu/runtime:reduce_scatter",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass",
@@ -869,6 +870,7 @@ cc_library(
         "//xla/stream_executor/gpu:all_reduce_kernel",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/functional:overload",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_kernel_strategy_annotator.cc b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_kernel_strategy_annotator.cc
index 1235f12f855..3e714679e3a 100644
--- a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_kernel_strategy_annotator.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_kernel_strategy_annotator.cc
@@ -15,10 +15,11 @@ limitations under the License.
 
 #include "xla/backends/gpu/transforms/collectives/collective_kernel_strategy_annotator.h"
 
-#include <utility>
+#include <variant>
 
 #include "absl/algorithm/container.h"
 #include "absl/container/flat_hash_set.h"
+#include "absl/functional/overload.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
@@ -27,12 +28,11 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "xla/backends/gpu/runtime/all_gather.h"
 #include "xla/backends/gpu/runtime/all_reduce.h"
+#include "xla/backends/gpu/runtime/collective_info_utils.h"
 #include "xla/backends/gpu/runtime/reduce_scatter.h"
 #include "xla/backends/gpu/transforms/collectives/collective_ops_utils.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/ir/hlo_opcode.h"
 #include "xla/service/gpu/backend_configs.pb.h"
@@ -53,18 +53,29 @@ using stream_executor::gpu::AllReduceStrategy;
 // Both the Triton codegen kernel and the built-in custom C++ kernel use the
 // same NVLink-based cost formula, so a single annotation covers both backends.
 CollectiveBackendConfig::CollectiveKernelStrategy ToProtoStrategy(
-    AllReduceStrategy strategy) {
-  switch (strategy) {
-    case AllReduceStrategy::kOneShot:
-      return CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT;
-    case AllReduceStrategy::kTwoShot:
-      return CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_TWO_SHOT;
-    case AllReduceStrategy::kMultimem:
-      // Not yet modelled in the cost model; fall back to default.
-      [[fallthrough]];
-    default:
-      return CollectiveBackendConfig::KERNEL_STRATEGY_DEFAULT;
-  }
+    const CollectiveKernelInfo& info) {
+  return std::visit(
+      absl::Overload{
+          [](const AllReduceInfo& all_reduce_info) {
+            switch (all_reduce_info.all_reduce_strategy) {
+              case AllReduceStrategy::kOneShot:
+                return CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT;
+              case AllReduceStrategy::kTwoShot:
+                return CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_TWO_SHOT;
+              case AllReduceStrategy::kMultimem:
+                // Not yet modelled in the cost model; fall back to default.
+                [[fallthrough]];
+              default:
+                return CollectiveBackendConfig::KERNEL_STRATEGY_DEFAULT;
+            }
+          },
+          [](const AllGatherInfo&) {
+            return CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT;
+          },
+          [](const ReduceScatterInfo&) {
+            return CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT;
+          }},
+      info);
 }
 
 bool HasCollectivesGroupAttribute(const HloInstruction* instruction) {
@@ -72,36 +83,30 @@ bool HasCollectivesGroupAttribute(const HloInstruction* instruction) {
       kCollectiveGroupMarkerAttr);
 }
 
-// Tries to determine the Triton kernel strategy for `instr` (which must be an
-// AllReduce or AllReduceStart) and writes the result into backend_config.
-// Returns true if the annotation was written.
-absl::StatusOr<bool> TryAnnotateAllReduce(HloInstruction* instr,
-                                          const GpuTopology& gpu_topology,
-                                          bool is_multimem_enabled) {
-  // Both kAllReduce and kAllReduceStart are HloAllReduceInstruction.
-  const auto* all_reduce = DynCast<HloAllReduceInstruction>(instr);
-  if (all_reduce == nullptr) {
-    return false;
-  }
-
+// Tries to determine the collective kernel strategy for `instr` and writes the
+// result into backend_config. Returns true if the annotation was written.
+absl::StatusOr<bool> TryAnnotateCollective(HloInstruction* instr,
+                                           const GpuTopology& gpu_topology,
+                                           bool is_multimem_enabled) {
   const DeviceAssignment* device_assignment = nullptr;
   if (instr->GetModule()->config().has_static_device_assignment()) {
     device_assignment =
         &instr->GetModule()->config().static_device_assignment();
   }
 
-  absl::StatusOr<AllReduceInfo> maybe_info = BuildAllReduceInfo(
-      is_multimem_enabled, gpu_topology, all_reduce, device_assignment);
+  absl::StatusOr<CollectiveKernelInfo> maybe_info = GetCollectiveInfo(
+      instr, gpu_topology, device_assignment, is_multimem_enabled);
   if (absl::IsUnimplemented(maybe_info.status())) {
     VLOG(3) << "[CollectiveKernelStrategyAnnotator] Collective kernel not "
                "supported for "
             << instr->name() << ": " << maybe_info.status();
     return false;
   }
-  ABSL_ASSIGN_OR_RETURN(AllReduceInfo info, std::move(maybe_info));
+
+  ABSL_ASSIGN_OR_RETURN(CollectiveKernelInfo info, maybe_info);
 
   CollectiveBackendConfig::CollectiveKernelStrategy proto_strategy =
-      ToProtoStrategy(info.all_reduce_strategy);
+      ToProtoStrategy(info);
 
   ABSL_ASSIGN_OR_RETURN(GpuBackendConfig gpu_config,
                    instr->backend_config<GpuBackendConfig>());
@@ -116,79 +121,6 @@ absl::StatusOr<bool> TryAnnotateAllReduce(HloInstruction* instr,
   return true;
 }
 
-// Tries to determine if the AllGather instruction should use the Triton
-// collective kernel and annotates it accordingly.
-// AllGather is always one-shot, so when eligible it is annotated with
-// KERNEL_STRATEGY_TRITON_ONE_SHOT.  Returns true if the annotation was written.
-absl::StatusOr<bool> TryAnnotateAllGather(HloInstruction* instr,
-                                          const GpuTopology& gpu_topology) {
-  const auto* all_gather = DynCast<HloAllGatherInstruction>(instr);
-  if (all_gather == nullptr) {
-    return false;
-  }
-
-  const DeviceAssignment* device_assignment = nullptr;
-  if (instr->GetModule()->config().has_static_device_assignment()) {
-    device_assignment =
-        &instr->GetModule()->config().static_device_assignment();
-  }
-
-  absl::StatusOr<AllGatherInfo> maybe_info =
-      BuildAllGatherInfo(gpu_topology, all_gather, device_assignment);
-  if (!maybe_info.ok()) {
-    VLOG(3) << "[CollectiveKernelStrategyAnnotator] Collective kernel not "
-               "supported for AllGather "
-            << instr->name() << ": " << maybe_info.status();
-    return false;
-  }
-
-  ABSL_ASSIGN_OR_RETURN(GpuBackendConfig gpu_config,
-                   instr->backend_config<GpuBackendConfig>());
-  gpu_config.mutable_collective_backend_config()->set_kernel_strategy(
-      CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT);
-  ABSL_RETURN_IF_ERROR(instr->set_backend_config(gpu_config));
-
-  VLOG(3) << "[CollectiveKernelStrategyAnnotator] Annotated AllGather "
-          << instr->name() << " with KERNEL_STRATEGY_TRITON_ONE_SHOT";
-  return true;
-}
-
-// Tries to determine if the ReduceScatter instruction should use the Triton
-// one-shot collective kernel and annotates it accordingly. Returns true if the
-// annotation was written.
-absl::StatusOr<bool> TryAnnotateReduceScatter(HloInstruction* instr,
-                                              const GpuTopology& gpu_topology) {
-  const auto* reduce_scatter = DynCast<HloReduceScatterInstruction>(instr);
-  if (reduce_scatter == nullptr) {
-    return false;
-  }
-
-  const DeviceAssignment* device_assignment = nullptr;
-  if (instr->GetModule()->config().has_static_device_assignment()) {
-    device_assignment =
-        &instr->GetModule()->config().static_device_assignment();
-  }
-
-  absl::StatusOr<ReduceScatterInfo> maybe_info =
-      BuildReduceScatterInfo(gpu_topology, reduce_scatter, device_assignment);
-  if (!maybe_info.ok()) {
-    VLOG(3) << "[CollectiveKernelStrategyAnnotator] Collective kernel not "
-               "supported for ReduceScatter "
-            << instr->name() << ": " << maybe_info.status();
-    return false;
-  }
-
-  ABSL_ASSIGN_OR_RETURN(GpuBackendConfig gpu_config,
-                   instr->backend_config<GpuBackendConfig>());
-  gpu_config.mutable_collective_backend_config()->set_kernel_strategy(
-      CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT);
-  ABSL_RETURN_IF_ERROR(instr->set_backend_config(gpu_config));
-
-  VLOG(3) << "[CollectiveKernelStrategyAnnotator] Annotated ReduceScatter "
-          << instr->name() << " with KERNEL_STRATEGY_TRITON_ONE_SHOT";
-  return true;
-}
-
 }  // namespace
 
 CollectiveKernelStrategyAnnotator::CollectiveKernelStrategyAnnotator(
@@ -220,21 +152,10 @@ absl::StatusOr<bool> CollectiveKernelStrategyAnnotator::RunImpl(
       if (!instructions_to_annotate.contains(instr->opcode())) {
         continue;
       }
-      if (instr->opcode() == HloOpcode::kAllReduce ||
-          instr->opcode() == HloOpcode::kAllReduceStart) {
-        ABSL_ASSIGN_OR_RETURN(
-            bool annotated,
-            TryAnnotateAllReduce(instr, gpu_topology_, is_multimem_enabled_));
-        changed |= annotated;
-      } else if (instr->opcode() == HloOpcode::kAllGather) {
-        ABSL_ASSIGN_OR_RETURN(bool annotated,
-                         TryAnnotateAllGather(instr, gpu_topology_));
-        changed |= annotated;
-      } else if (instr->opcode() == HloOpcode::kReduceScatter) {
-        ABSL_ASSIGN_OR_RETURN(bool annotated,
-                         TryAnnotateReduceScatter(instr, gpu_topology_));
-        changed |= annotated;
-      }
+      ABSL_ASSIGN_OR_RETURN(
+          bool annotated,
+          TryAnnotateCollective(instr, gpu_topology_, is_multimem_enabled_));
+      changed |= annotated;
     }
   }
   return changed;

From 154c9c3e1546b7b6615dbb6aff655fc865a64975 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 19:11:32 +0530
Subject: [PATCH 0766/1000] Revert rank-check test regression and guard scalar
 elements

- Restore the original assertion in
  testIterableFailsWithUnknownShapeDatasetSpec (unknown rank -> 'must
  have a valid rank'), and add a new
  testIterableFailsWithPartiallyDefinedShapeDatasetSpec covering the
  is_fully_defined() check with a partially-defined shape.
- Guard against scalar dataset elements: raise a clear ValueError in
  _validate_input when dataset_already_batched and elem_spec.shape.rank
  == 0, before indexing as_list()[0]; and raise before indexing
  layout.sharding_specs[0] when a batch_dim is specified but a layout
  has rank 0.
---
 tensorflow/dtensor/python/input_util.py         |  9 +++++++++
 .../dtensor/python/tests/input_util_test.py     | 17 +++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/tensorflow/dtensor/python/input_util.py b/tensorflow/dtensor/python/input_util.py
index da0c6e02e02..38762d3a178 100644
--- a/tensorflow/dtensor/python/input_util.py
+++ b/tensorflow/dtensor/python/input_util.py
@@ -245,6 +245,11 @@ def _validate_input(flattened_layouts: Sequence[layout_lib.Layout],
           (expected_rank, elem_spec, layout.sharding_specs))
 
     if dataset_already_batched:
+      if elem_spec.shape.rank == 0:
+        raise ValueError(
+            ('Dataset element shape must have rank >= 1 when the dataset is '
+             'already batched, got spec %s.') % elem_spec)
+
       # Check that the batch dimension size of all dataset elements match.
       batch_dim_size = first_elem_shape.as_list()[0]
       if batch_dim_size is None:
@@ -495,6 +500,10 @@ class DTensorDataset(dataset_ops.UnaryUnchangedStructureDataset):
               [loc[batch_dim] for loc in mesh.local_device_locations()]))
 
       for layout in flattened_layouts:
+        if layout.rank == 0:
+          raise ValueError(
+              ('batch_dim %s was specified but at least one layout has rank '
+               '0: %s') % (batch_dim, layout))
         if batch_dim != layout.sharding_specs[0]:
           raise ValueError(
               ('batch_dim %s was specified but at least one layout did not '
diff --git a/tensorflow/dtensor/python/tests/input_util_test.py b/tensorflow/dtensor/python/tests/input_util_test.py
index 72428872e9e..e26f0e9d6f2 100644
--- a/tensorflow/dtensor/python/tests/input_util_test.py
+++ b/tensorflow/dtensor/python/tests/input_util_test.py
@@ -82,6 +82,23 @@ class DTensorDatasetTest(test_util.DTensorBaseTest):
         output_signature=tensor_spec.TensorSpec(
             tensor_shape.TensorShape(None), dtype=dtypes.int32))
 
+    with self.assertRaisesRegex(
+        ValueError, 'Dataset element shape must have a valid rank'):
+      input_util.DTensorDataset(
+          dataset=dataset,
+          global_batch_size=8,
+          mesh=self.mesh,
+          layouts=Layout.replicated(self.mesh, rank=2))
+
+  def testIterableFailsWithPartiallyDefinedShapeDatasetSpec(self):
+    def gen():
+      yield constant_op.constant([1, 2], dtype=dtypes.int32)
+
+    dataset = dataset_ops.DatasetV2.from_generator(
+        gen,
+        output_signature=tensor_spec.TensorSpec(
+            tensor_shape.TensorShape([None]), dtype=dtypes.int32))
+
     with self.assertRaisesRegex(
         ValueError, 'Dataset element shape must be fully defined'):
       input_util.DTensorDataset(

From efd5c92f7ce9f1054d69d97c821f47104930d9bf Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 19:17:27 +0530
Subject: [PATCH 0767/1000] Fix
 test_checkpoint_deferred_restore_scalar_dvariable CI failure

trackable.base has no SimpleTrackable class, so
root = trackable.SimpleTrackable() raised AttributeError at test
collection time. Restructure the test to match the known-good pattern
in checkpointing_test.py's testInitializeFromCheckpoint: restore into
a bare checkpoint.Checkpoint() and call _preload_simple_restoration
directly on it, since the 'v' dependency saved via
Checkpoint(v=v_save) lives on the Checkpoint root itself, not on a
nested child object.

Also add the missing //tensorflow/python/trackable:base BUILD
dependency for the trackable.CheckpointInitialValueCallable import.
---
 tensorflow/dtensor/python/tests/BUILD                  |  1 +
 .../dtensor/python/tests/save_restore_v2_test.py       | 10 ++++++----
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/tensorflow/dtensor/python/tests/BUILD b/tensorflow/dtensor/python/tests/BUILD
index 0db3a7cf26e..0880818c1bc 100644
--- a/tensorflow/dtensor/python/tests/BUILD
+++ b/tensorflow/dtensor/python/tests/BUILD
@@ -932,6 +932,7 @@ dtensor_test(
         "//tensorflow/python/ops:math_ops",
         "//tensorflow/python/ops:stateless_random_ops",
         "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/trackable:base",
         "//third_party/py/numpy",
         "@absl_py//absl/testing:parameterized",
     ],
diff --git a/tensorflow/dtensor/python/tests/save_restore_v2_test.py b/tensorflow/dtensor/python/tests/save_restore_v2_test.py
index 8455cf3b6f0..6d00a246598 100644
--- a/tensorflow/dtensor/python/tests/save_restore_v2_test.py
+++ b/tensorflow/dtensor/python/tests/save_restore_v2_test.py
@@ -134,13 +134,15 @@ class DTensorSaveRestoreV2Test(test_util.DTensorBaseTest):
     ckpt = checkpoint.Checkpoint(v=v_save)
     saved_path = ckpt.save(self.get_temp_dir())
 
-    # Re-create checkpoint on an empty container to test deferred restoration
-    root = trackable.SimpleTrackable()
-    restore_ckpt = checkpoint.Checkpoint(root=root)
+    # Re-create an empty checkpoint object to test deferred restoration. The
+    # dependency named "v" lives directly on the Checkpoint root, matching
+    # the attribute name used when saving above (see checkpointing_test.py's
+    # testInitializeFromCheckpoint for the same pattern).
+    restore_ckpt = checkpoint.Checkpoint()
     restore_ckpt.restore(saved_path)
 
     # Preload simple restoration produces a CheckpointInitialValueCallable
-    init_callable = root._preload_simple_restoration('v')
+    init_callable = restore_ckpt._preload_simple_restoration('v')
     self.assertIsInstance(
         init_callable, trackable.CheckpointInitialValueCallable
     )

From a774de5f2a162f1bb0d2e4cae313c116c67269b7 Mon Sep 17 00:00:00 2001
From: Levon Ter-Grigoryan <patrios@google.com>
Date: Wed, 14 Oct 2026 06:20:23 -0700
Subject: [PATCH 0768/1000] Add `BarrierMode::kPeerSymmetric`
 (`peer_symmetric`) to `triton_xla.block_barrier` for 1-to-1 pairwise
 synchronization between corresponding blocks across peer GPUs.

This mode is needed to AllToAll collective where each block should communicate with a single peer block on a peer devices and these block ids are different.

PiperOrigin-RevId: 999324691
---
 .../gpu/codegen/triton/ir/tests/invalid.mlir  |  16 +
 .../gpu/codegen/triton/ir/triton_xla_attrs.td |   1 +
 .../gpu/codegen/triton/ir/triton_xla_ops.td   |  18 +-
 .../tests/triton_xla_block_barrier.mlir       |  75 ++++-
 .../triton_xla_lower_block_barrier_pass.cc    | 314 ++++++++++++------
 5 files changed, 308 insertions(+), 116 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/ir/tests/invalid.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/ir/tests/invalid.mlir
index 603a011b0af..1e04e942cbe 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/ir/tests/invalid.mlir
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/ir/tests/invalid.mlir
@@ -112,3 +112,19 @@ func.func @block_barrier_asymmetric_without_stride(%ptr: !tt.ptr<i64>, %rank: i3
   triton_xla.block_barrier %ptr, %rank, %sig, %slot <world_size = 2, barrier_mode = producer_symmetric> : (!tt.ptr<i64>, i32, i32, i32) -> ()
   return
 }
+
+// -----
+
+func.func @block_barrier_pairwise_symmetric_without_slot(%ptr: !tt.ptr<i64>, %rank: i32, %sig: i32) {
+  // expected-error @+1 {{'triton_xla.block_barrier' op asymmetric barrier_mode requires signal_slot to be specified}}
+  triton_xla.block_barrier %ptr, %rank, %sig <world_size = 2, signal_stride = 4, barrier_mode = pairwise_symmetric> : (!tt.ptr<i64>, i32, i32) -> ()
+  return
+}
+
+// -----
+
+func.func @block_barrier_pairwise_symmetric_without_stride(%ptr: !tt.ptr<i64>, %rank: i32, %sig: i32, %slot: i32) {
+  // expected-error @+1 {{'triton_xla.block_barrier' op asymmetric barrier_mode requires a positive signal_stride}}
+  triton_xla.block_barrier %ptr, %rank, %sig, %slot <world_size = 2, barrier_mode = pairwise_symmetric> : (!tt.ptr<i64>, i32, i32, i32) -> ()
+  return
+}
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_attrs.td b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_attrs.td
index de54d77026f..e128d8933d6 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_attrs.td
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_attrs.td
@@ -92,6 +92,7 @@ def TTXLA_BarrierModeEnum : I32EnumAttr<"BarrierMode",
       I32EnumAttrCase<"kSymmetric", 0, "symmetric">,
       I32EnumAttrCase<"kConsumerSymmetric", 1, "consumer_symmetric">,
       I32EnumAttrCase<"kProducerSymmetric", 2, "producer_symmetric">,
+      I32EnumAttrCase<"kPairwiseSymmetric", 3, "pairwise_symmetric">,
     ]> {
   let cppNamespace = "::mlir::triton::xla";
   let genSpecializedAttr = 0;
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.td b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.td
index b8ebaee318e..3bda38c1f66 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.td
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/ir/triton_xla_ops.td
@@ -337,7 +337,7 @@ def TTXLA_BlockBarrierOp : TTXLA_Op<"block_barrier", []> {
     of `i32` flags with shape `[num_blocks][world_size]`, flattened as
     `block_idx * world_size + slot_idx`.
 
-    This barrier supports three synchronization modes via `$barrier_mode`:
+    This barrier supports four synchronization modes via `$barrier_mode`:
 
     ### 1. `symmetric` (default; `$signal_stride == 0`, `$signal_slot` omitted)
     Used when block `block_idx` on every GPU produces and consumes data for the
@@ -379,6 +379,22 @@ def TTXLA_BlockBarrierOp : TTXLA_Op<"block_barrier", []> {
       buffer `$signal_buffers[r]` until all `world_size` slots at
       `target_block * world_size + [0..world_size)` reach `>= $signal_value`.
 
+    ### 4. `pairwise_symmetric` (`$signal_stride > 0`, `$signal_slot` provided)
+    Used for 1-to-1 pairwise synchronization where block `block_idx` on GPU `r`
+    communicates exclusively with block `target_block` on peer GPU
+    `$signal_slot` (`p = $signal_slot`):
+    - Let `base_block = block_idx - $signal_slot * $signal_stride` and
+      `target_block = base_block + r * $signal_stride`.
+    - **Signal (Write)**: Thread 0 of block `block_idx` on GPU `r` updates its
+      local invocation counter at `$signal_buffers[r][block_idx * world_size + r]`
+      and atomically writes `$signal_value` to index
+      `target_block * world_size + r` in peer GPU `$signal_slot`'s signal
+      buffer (`$signal_buffers[$signal_slot]`).
+    - **Wait (Read)**: Thread 0 of block `block_idx` on GPU `r` spin-waits on
+      its own buffer `$signal_buffers[r]` at index
+      `block_idx * world_size + $signal_slot` until it reaches
+      `>= $signal_value`.
+
     *NOTE* Block orchestration order is not guaranteed across GPUs,
     so there is a chance of a deadlock should the number of blocks be greater
     than the number of SMs on a GPU. That is to say that if BlockN is launched
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_block_barrier.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_block_barrier.mlir
index d587b6c3788..584e94d8ae9 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_block_barrier.mlir
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_block_barrier.mlir
@@ -25,14 +25,14 @@ tt.func @block_barrier_kernel(
   // CHECK-NEXT: %[[PROGRAM_ID:.+]] = tt.get_program_id x
   // CHECK-NEXT: %[[COND:.+]] = arith.cmpi ult, %[[TID]], %[[WORLD_SIZE]]
   // CHECK-NEXT: scf.if %[[COND]] {
-  // CHECK-NEXT:   %[[SPLAT_PTR:.+]] = tt.splat %[[PTR]]
+  // CHECK-NEXT:   %[[MULI:.+]] = arith.muli %[[PROGRAM_ID]], %[[WORLD_SIZE]]
+  // CHECK-NEXT:   %[[ADDI:.+]] = arith.addi %[[MULI]], %[[RANK]]
   // CHECK-NEXT:   %[[RANGE:.+]] = tt.make_range {end = 8 : i32, start = 0 : i32}
+  // CHECK-NEXT:   %[[SPLAT_PTR:.+]] = tt.splat %[[PTR]]
+  // CHECK-NEXT:   %[[SPLAT_ADDI:.+]] = tt.splat %[[ADDI]]
   // CHECK-NEXT:   %[[ADD_PTR:.+]] = tt.addptr %[[SPLAT_PTR]], %[[RANGE]]
   // CHECK-NEXT:   %[[LOAD:.+]] = tt.load %[[ADD_PTR]]
   // CHECK-NEXT:   %[[INT_TO_PTR:.+]] = tt.int_to_ptr %[[LOAD]]
-  // CHECK-NEXT:   %[[MULI:.+]] = arith.muli %[[PROGRAM_ID]], %[[WORLD_SIZE]]
-  // CHECK-NEXT:   %[[ADDI:.+]] = arith.addi %[[MULI]], %[[RANK]]
-  // CHECK-NEXT:   %[[SPLAT_ADDI:.+]] = tt.splat %[[ADDI]]
   // CHECK-NEXT:   %[[ADD_PTR_2:.+]] = tt.addptr %[[INT_TO_PTR]], %[[SPLAT_ADDI]]
   // CHECK-NEXT:   triton_xla.atomic_write sys, release, %[[ADD_PTR_2]], %[[SIGNAL_VALUE]]
   // CHECK-NEXT:   %[[ADD_PTR_3:.+]] = tt.addptr %[[PTR]], %[[RANK]]
@@ -63,11 +63,6 @@ tt.func @block_barrier_consumer_symmetric_kernel(
   // CHECK-NEXT: %[[PROGRAM_ID:.+]] = tt.get_program_id x
   // CHECK-NEXT: %[[COND:.+]] = arith.cmpi ult, %[[TID]], %[[WORLD_SIZE]]
   // CHECK-NEXT: scf.if %[[COND]] {
-  // CHECK-NEXT:   %[[SPLAT_PTR:.+]] = tt.splat %[[PTR]]
-  // CHECK-NEXT:   %[[RANGE:.+]] = tt.make_range {end = 2 : i32, start = 0 : i32}
-  // CHECK-NEXT:   %[[ADD_PTR:.+]] = tt.addptr %[[SPLAT_PTR]], %[[RANGE]]
-  // CHECK-NEXT:   %[[LOAD:.+]] = tt.load %[[ADD_PTR]]
-  // CHECK-NEXT:   %[[INT_TO_PTR:.+]] = tt.int_to_ptr %[[LOAD]]
   // CHECK-NEXT:   %[[MULI:.+]] = arith.muli %[[PROGRAM_ID]], %[[WORLD_SIZE]]
   // CHECK-NEXT:   %[[S_OFFSET:.+]] = arith.muli %[[SIGNAL_SLOT]], %[[STRIDE]]
   // CHECK-NEXT:   %[[BASE_T:.+]] = arith.subi %[[PROGRAM_ID]], %[[S_OFFSET]]
@@ -75,7 +70,12 @@ tt.func @block_barrier_consumer_symmetric_kernel(
   // CHECK-NEXT:   %[[WRITE_BLOCK:.+]] = arith.addi %[[BASE_T]], %[[RANK_OFFSET]]
   // CHECK-NEXT:   %[[WRITE_BLOCK_OFFSET:.+]] = arith.muli %[[WRITE_BLOCK]], %[[WORLD_SIZE]]
   // CHECK-NEXT:   %[[WRITE_OFFSET:.+]] = arith.addi %[[WRITE_BLOCK_OFFSET]], %[[SIGNAL_SLOT]]
+  // CHECK-NEXT:   %[[RANGE:.+]] = tt.make_range {end = 2 : i32, start = 0 : i32}
+  // CHECK-NEXT:   %[[SPLAT_PTR:.+]] = tt.splat %[[PTR]]
   // CHECK-NEXT:   %[[SPLAT_WRITE_OFFSET:.+]] = tt.splat %[[WRITE_OFFSET]]
+  // CHECK-NEXT:   %[[ADD_PTR:.+]] = tt.addptr %[[SPLAT_PTR]], %[[RANGE]]
+  // CHECK-NEXT:   %[[LOAD:.+]] = tt.load %[[ADD_PTR]]
+  // CHECK-NEXT:   %[[INT_TO_PTR:.+]] = tt.int_to_ptr %[[LOAD]]
   // CHECK-NEXT:   %[[ADD_PTR_2:.+]] = tt.addptr %[[INT_TO_PTR]], %[[SPLAT_WRITE_OFFSET]]
   // CHECK-NEXT:   triton_xla.atomic_write sys, release, %[[ADD_PTR_2]], %[[SIGNAL_VALUE]]
   // CHECK-NEXT:   %[[ADD_PTR_3:.+]] = tt.addptr %[[PTR]], %[[RANK]]
@@ -107,11 +107,6 @@ tt.func @block_barrier_producer_symmetric_kernel(
   // CHECK-NEXT: %[[PROGRAM_ID:.+]] = tt.get_program_id x
   // CHECK-NEXT: %[[COND:.+]] = arith.cmpi ult, %[[TID]], %[[WORLD_SIZE]]
   // CHECK-NEXT: scf.if %[[COND]] {
-  // CHECK-NEXT:   %[[SPLAT_PTR:.+]] = tt.splat %[[PTR]]
-  // CHECK-NEXT:   %[[RANGE:.+]] = tt.make_range {end = 2 : i32, start = 0 : i32}
-  // CHECK-NEXT:   %[[ADD_PTR:.+]] = tt.addptr %[[SPLAT_PTR]], %[[RANGE]]
-  // CHECK-NEXT:   %[[LOAD:.+]] = tt.load %[[ADD_PTR]]
-  // CHECK-NEXT:   %[[INT_TO_PTR:.+]] = tt.int_to_ptr %[[LOAD]]
   // CHECK-NEXT:   %[[MULI:.+]] = arith.muli %[[PROGRAM_ID]], %[[WORLD_SIZE]]
   // CHECK-NEXT:   %[[S_OFFSET:.+]] = arith.muli %[[SIGNAL_SLOT]], %[[STRIDE]]
   // CHECK-NEXT:   %[[BASE_T:.+]] = arith.subi %[[PROGRAM_ID]], %[[S_OFFSET]]
@@ -119,7 +114,12 @@ tt.func @block_barrier_producer_symmetric_kernel(
   // CHECK-NEXT:   %[[WAIT_BLOCK:.+]] = arith.addi %[[BASE_T]], %[[RANK_OFFSET]]
   // CHECK-NEXT:   %[[WAIT_BLOCK_OFFSET:.+]] = arith.muli %[[WAIT_BLOCK]], %[[WORLD_SIZE]]
   // CHECK-NEXT:   %[[WRITE_OFFSET:.+]] = arith.addi %[[MULI]], %[[RANK]]
+  // CHECK-NEXT:   %[[RANGE:.+]] = tt.make_range {end = 2 : i32, start = 0 : i32}
+  // CHECK-NEXT:   %[[SPLAT_PTR:.+]] = tt.splat %[[PTR]]
   // CHECK-NEXT:   %[[SPLAT_WRITE_OFFSET:.+]] = tt.splat %[[WRITE_OFFSET]]
+  // CHECK-NEXT:   %[[ADD_PTR:.+]] = tt.addptr %[[SPLAT_PTR]], %[[RANGE]]
+  // CHECK-NEXT:   %[[LOAD:.+]] = tt.load %[[ADD_PTR]]
+  // CHECK-NEXT:   %[[INT_TO_PTR:.+]] = tt.int_to_ptr %[[LOAD]]
   // CHECK-NEXT:   %[[ADD_PTR_2:.+]] = tt.addptr %[[INT_TO_PTR]], %[[SPLAT_WRITE_OFFSET]]
   // CHECK-NEXT:   triton_xla.atomic_write sys, release, %[[ADD_PTR_2]], %[[SIGNAL_VALUE]]
   // CHECK-NEXT:   %[[ADD_PTR_3:.+]] = tt.addptr %[[PTR]], %[[RANK]]
@@ -137,3 +137,50 @@ tt.func @block_barrier_producer_symmetric_kernel(
     (!tt.ptr<i64>, i32, i32, i32) -> ()
   tt.return
 }
+
+// CHECK-LABEL: block_barrier_pairwise_symmetric_kernel
+// CHECK-SAME: %[[PTR:[^:]+]]: !tt.ptr<i64>,
+// CHECK-SAME: %[[SIGNAL_VALUE:[^:]+]]: i32,
+// CHECK-SAME: %[[RANK:[^:]+]]: i32,
+// CHECK-SAME: %[[SIGNAL_SLOT:[^:]+]]: i32
+tt.func @block_barrier_pairwise_symmetric_kernel(
+  %ptr : !tt.ptr<i64>, %signal_value : i32, %rank: i32, %signal_slot: i32) {
+  // CHECK-DAG: %[[ZERO:.+]] = arith.constant 0 : i32
+  // CHECK-DAG: %[[STRIDE:.+]] = arith.constant 4 : i32
+  // CHECK-DAG: %[[WORLD_SIZE:.+]] = arith.constant 2 : i32
+  // CHECK: %[[TID:.+]] = triton_xla.get_tid
+  // CHECK-NEXT: %[[PROGRAM_ID:.+]] = tt.get_program_id x
+  // CHECK-NEXT: %[[COND:.+]] = arith.cmpi eq, %[[TID]], %[[ZERO]]
+  // CHECK-NEXT: scf.if %[[COND]] {
+  // CHECK-NEXT:   %[[BLOCK_OFFSET:.+]] = arith.muli %[[PROGRAM_ID]], %[[WORLD_SIZE]]
+  // CHECK-NEXT:   %[[LOCAL_CNT_OFFSET:.+]] = arith.addi %[[BLOCK_OFFSET]], %[[RANK]]
+  // CHECK-NEXT:   %[[LOCAL_BUF_PTR:.+]] = tt.addptr %[[PTR]], %[[RANK]]
+  // CHECK-NEXT:   %[[LOCAL_BUF_I64:.+]] = tt.load %[[LOCAL_BUF_PTR]]
+  // CHECK-NEXT:   %[[LOCAL_BUF:.+]] = tt.int_to_ptr %[[LOCAL_BUF_I64]]
+  // CHECK-NEXT:   %[[LOCAL_CNT_ADDR:.+]] = tt.addptr %[[LOCAL_BUF]], %[[LOCAL_CNT_OFFSET]]
+  // CHECK-NEXT:   triton_xla.atomic_write gpu, relaxed, %[[LOCAL_CNT_ADDR]], %[[SIGNAL_VALUE]]
+  // CHECK-NEXT:   %[[S_OFFSET:.+]] = arith.muli %[[SIGNAL_SLOT]], %[[STRIDE]]
+  // CHECK-NEXT:   %[[BASE_T:.+]] = arith.subi %[[PROGRAM_ID]], %[[S_OFFSET]]
+  // CHECK-NEXT:   %[[RANK_OFFSET:.+]] = arith.muli %[[RANK]], %[[STRIDE]]
+  // CHECK-NEXT:   %[[TARGET_BLOCK:.+]] = arith.addi %[[BASE_T]], %[[RANK_OFFSET]]
+  // CHECK-NEXT:   %[[TARGET_BLOCK_OFFSET:.+]] = arith.muli %[[TARGET_BLOCK]], %[[WORLD_SIZE]]
+  // CHECK-NEXT:   %[[PEER_WRITE_OFFSET:.+]] = arith.addi %[[TARGET_BLOCK_OFFSET]], %[[RANK]]
+  // CHECK-NEXT:   %[[PEER_BUF_PTR:.+]] = tt.addptr %[[PTR]], %[[SIGNAL_SLOT]]
+  // CHECK-NEXT:   %[[PEER_BUF_I64:.+]] = tt.load %[[PEER_BUF_PTR]]
+  // CHECK-NEXT:   %[[PEER_BUF:.+]] = tt.int_to_ptr %[[PEER_BUF_I64]]
+  // CHECK-NEXT:   %[[PEER_SIG_ADDR:.+]] = tt.addptr %[[PEER_BUF]], %[[PEER_WRITE_OFFSET]]
+  // CHECK-NEXT:   triton_xla.atomic_write sys, release, %[[PEER_SIG_ADDR]], %[[SIGNAL_VALUE]]
+  // CHECK-NEXT:   %[[WAIT_OFFSET:.+]] = arith.addi %[[BLOCK_OFFSET]], %[[SIGNAL_SLOT]]
+  // CHECK-NEXT:   %[[WAIT_BUF_PTR:.+]] = tt.addptr %[[PTR]], %[[RANK]]
+  // CHECK-NEXT:   %[[WAIT_BUF_I64:.+]] = tt.load %[[WAIT_BUF_PTR]]
+  // CHECK-NEXT:   %[[WAIT_BUF:.+]] = tt.int_to_ptr %[[WAIT_BUF_I64]]
+  // CHECK-NEXT:   %[[WAIT_ADDR:.+]] = tt.addptr %[[WAIT_BUF]], %[[WAIT_OFFSET]]
+  // CHECK-NEXT:   triton_xla.atomic_spin_wait sys, acquire, %[[WAIT_ADDR]], greater_than_or_equal_to, %[[SIGNAL_VALUE]]
+  // CHECK-NEXT: }
+  // CHECK-NEXT: ttg.barrier local
+  // CHECK-NEXT: tt.return
+  triton_xla.block_barrier %ptr, %rank, %signal_value, %signal_slot
+    <world_size = 2, signal_stride = 4, barrier_mode = pairwise_symmetric> :
+    (!tt.ptr<i64>, i32, i32, i32) -> ()
+  tt.return
+}
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_lower_block_barrier_pass.cc b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_lower_block_barrier_pass.cc
index 0cf6fabfe70..4019f9ce7a3 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_lower_block_barrier_pass.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_lower_block_barrier_pass.cc
@@ -17,13 +17,14 @@ limitations under the License.
 #include <utility>
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/IR/ArithAttributes.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeRange.h"
 #include "mlir/IR/Value.h"
-#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/Pass.h"  // IWYU pragma: keep
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -40,8 +41,199 @@ namespace mlir::triton::xla {
 
 namespace {
 
+// Suppose we have 4 blocks (0, 1, 2, 3), 2 ranks (world_size = 2),
+// and signal_stride = 2.
+//
+// Blocks 1 and 3 work on the same scratch tile and are `signal_stride` blocks
+// apart:
+//   - Block 1 has signal_slot = 0 (first block in the pair).
+//   - Block 3 has signal_slot = 1 (second block in the pair).
+//
+// 1. Find `base_block`, the first block in the pair (block 1):
+//    `slot_offset = signal_slot * signal_stride` is how many blocks
+//    this block is ahead of `base_block`.
+//    Subtracting it gives `base_block = block_id - slot_offset`:
+//      - Block 1: base_block = 1 - (0 * 2) = 1
+//      - Block 3: base_block = 3 - (1 * 2) = 1
+//
+// 2. Find `target_block` (`base_block + rank * signal_stride`).
+mlir::Value ComputeTargetBlock(mlir::ImplicitLocOpBuilder& builder,
+                               mlir::Value block_id, mlir::Value rank,
+                               mlir::Value signal_slot, int32_t signal_stride) {
+  const mlir::IntegerType i32_type = builder.getI32Type();
+  auto stride_op = mlir::arith::ConstantOp::create(
+      builder, builder.getI32IntegerAttr(signal_stride));
+  auto slot_offset =
+      mlir::arith::MulIOp::create(builder, i32_type, signal_slot, stride_op);
+  auto base_block =
+      mlir::arith::SubIOp::create(builder, i32_type, block_id, slot_offset);
+  auto rank_offset =
+      mlir::arith::MulIOp::create(builder, i32_type, rank, stride_op);
+  return mlir::arith::AddIOp::create(builder, i32_type, base_block,
+                                     rank_offset);
+}
+
+// Returns a pointer (or 1D tensor of pointers) to `&SignalBuffers[x][y]`.
+// `x` and `y` may be scalar `i32` values or 1D `i32` tensors. If `x` is a 1D
+// tensor and `y` is a scalar `i32`, `y` is splatted to match `x`'s shape.
+mlir::Value GetSignalBufferElementPtr(mlir::ImplicitLocOpBuilder& builder,
+                                      mlir::Value signal_buffers_arg,
+                                      mlir::Value x, mlir::Value y) {
+  constexpr auto kGlobalAddressSpace = mlir::triton::PtrAddrSpace::Global;
+  const mlir::IntegerType i32_type = builder.getI32Type();
+  const mlir::IntegerType i64_type = builder.getI64Type();
+  const auto ptr_to_i32_type =
+      mlir::triton::PointerType::get(i32_type, kGlobalAddressSpace);
+  const auto ptr_to_i64_type =
+      mlir::triton::PointerType::get(i64_type, kGlobalAddressSpace);
+
+  // Triton fails to do pointer arithmetic on pointer of pointers, so we cast
+  // the inner pointer to i64.
+  mlir::Value signal_buffers_i64 = signal_buffers_arg;
+  if (signal_buffers_arg.getType() != ptr_to_i64_type) {
+    signal_buffers_i64 = mlir::triton::BitcastOp::create(
+        builder, ptr_to_i64_type, signal_buffers_arg);
+  }
+
+  mlir::Type ptr_i64_type = ptr_to_i64_type;
+  mlir::Type ptr_i32_type = ptr_to_i32_type;
+  if (auto ranked_x = mlir::dyn_cast<mlir::RankedTensorType>(x.getType())) {
+    ptr_i64_type =
+        mlir::RankedTensorType::get(ranked_x.getShape(), ptr_to_i64_type);
+    ptr_i32_type =
+        mlir::RankedTensorType::get(ranked_x.getShape(), ptr_to_i32_type);
+    signal_buffers_i64 = mlir::triton::SplatOp::create(builder, ptr_i64_type,
+                                                       signal_buffers_i64);
+    if (!mlir::isa<mlir::RankedTensorType>(y.getType())) {
+      y = mlir::triton::SplatOp::create(builder, ranked_x, y);
+    }
+  }
+
+  auto buf_ptr_i64 = mlir::triton::AddPtrOp::create(builder, ptr_i64_type,
+                                                    signal_buffers_i64, x);
+  auto buf_i64 = mlir::triton::LoadOp::create(builder, /*ptr=*/buf_ptr_i64,
+                                              /*isVolatile=*/false);
+  auto buf = mlir::triton::IntToPtrOp::create(builder, ptr_i32_type, buf_i64);
+  return mlir::triton::AddPtrOp::create(builder, ptr_i32_type, buf, y);
+}
+
+// Atomically writes `signal_value` to `SignalBuffers[x][y]`.
+void WriteToSignalBuffer(mlir::ImplicitLocOpBuilder& builder,
+                         mlir::Value signal_buffers_arg, mlir::Value x,
+                         mlir::Value y, mlir::Value signal_value,
+                         mlir::triton::MemSyncScope mem_sync_scope,
+                         mlir::triton::MemSemantic mem_sync_semantic) {
+  mlir::Value signal_addr =
+      GetSignalBufferElementPtr(builder, signal_buffers_arg, x, y);
+  mlir::triton::xla::AtomicWriteOp::create(
+      builder,
+      /*resultTypes=*/mlir::TypeRange{},
+      /*ptr=*/signal_addr,
+      /*value=*/signal_value,
+      /*mask=*/mlir::Value{},
+      /*mem_sync_scope=*/mem_sync_scope,
+      /*mem_sync_semantic=*/mem_sync_semantic);
+}
+
+// PairwiseSymmetric barrier mode is lowered separately because it synchronizes
+// with a single peer rank using scalar ops on thread 0, whereas the other
+// modes synchronize across all world_size ranks using tensor ops on
+// world_size threads.
+//
+// Two atomic writes are needed because unlike other modes, which are
+// broadcasting to every rank's buffer and updating the local rank's own
+// invocation counter as part of the same write, PairwiseSymmetric
+// only signals one peer's buffer and needs a separate local write to advance
+// its own invocation counter.
+//
+// Because only the local and target peer slots are updated, the signal buffer
+// is left with different values across slots and must not be shared with
+// other operations or barrier modes.
+LogicalResult LowerPairwiseSymmetricBlockBarrierOp(BlockBarrierOp block_barrier,
+                                                   PatternRewriter& rewriter) {
+  mlir::ImplicitLocOpBuilder builder(block_barrier.getLoc(), rewriter);
+
+  mlir::TypedValue<triton::PointerType> signal_buffers_arg =
+      block_barrier.getSignalBuffers();
+  const mlir::TypedValue<mlir::IntegerType> rank =
+      block_barrier.getDeviceRank();
+  const mlir::TypedValue<mlir::IntegerType> signal_value =
+      block_barrier.getSignalValue();
+  const mlir::Value signal_slot = block_barrier.getSignalSlot();
+  const int32_t world_size = block_barrier.getWorldSize();
+  const int32_t signal_stride = block_barrier.getSignalStride();
+
+  const mlir::TypedValue<mlir::Type> world_size_op =
+      mlir::arith::ConstantOp::create(builder,
+                                      builder.getI32IntegerAttr(world_size));
+  const mlir::TypedValue<mlir::IntegerType> thread_id =
+      triton::xla::GetTidOp::create(builder);
+  const mlir::TypedValue<mlir::IntegerType> block_id =
+      triton::GetProgramIdOp::create(builder, 0);
+  auto zero_op =
+      mlir::arith::ConstantOp::create(builder, builder.getI32IntegerAttr(0));
+  auto tid_is_zero = mlir::arith::CmpIOp::create(
+      builder, mlir::arith::CmpIPredicate::eq, thread_id, zero_op);
+  mlir::scf::IfOp::create(
+      builder,
+      /*cond=*/tid_is_zero,
+      /*thenBuilder=*/
+      [&](mlir::OpBuilder& op_builder, mlir::Location location) {
+        mlir::ImplicitLocOpBuilder builder(location, op_builder);
+        const mlir::IntegerType i32_type = builder.getI32Type();
+
+        // Update local invocation counter slot
+        // SignalBuffers[rank][block_id * world_size + rank].
+        auto block_offset = mlir::arith::MulIOp::create(
+            builder, i32_type, block_id, world_size_op);
+        auto local_counter_offset =
+            mlir::arith::AddIOp::create(builder, i32_type, block_offset, rank);
+        WriteToSignalBuffer(builder, signal_buffers_arg, rank,
+                            local_counter_offset, signal_value,
+                            mlir::triton::MemSyncScope::GPU,
+                            mlir::triton::MemSemantic::RELAXED);
+
+        // Signal peer at
+        // SignalBuffers[signal_slot][target_block * world_size + rank].
+        mlir::Value target_block = ComputeTargetBlock(
+            builder, block_id, rank, signal_slot, signal_stride);
+        auto target_block_offset = mlir::arith::MulIOp::create(
+            builder, i32_type, target_block, world_size_op);
+        auto peer_write_offset = mlir::arith::AddIOp::create(
+            builder, i32_type, target_block_offset, rank);
+        WriteToSignalBuffer(builder, signal_buffers_arg, signal_slot,
+                            peer_write_offset, signal_value,
+                            mlir::triton::MemSyncScope::SYSTEM,
+                            mlir::triton::MemSemantic::RELEASE);
+
+        // Wait on SignalBuffers[rank][block_id * world_size + signal_slot].
+        auto wait_offset = mlir::arith::AddIOp::create(
+            builder, i32_type, block_offset, signal_slot);
+        mlir::Value wait_addr = GetSignalBufferElementPtr(
+            builder, signal_buffers_arg, rank, wait_offset);
+        mlir::triton::xla::AtomicSpinWaitOp::create(
+            builder,
+            /*resultTypes=*/mlir::TypeRange{},
+            /*ptr=*/wait_addr,
+            /*expected=*/signal_value,
+            /*mask=*/mlir::Value{},
+            /*mem_sync_scope=*/mlir::triton::MemSyncScope::SYSTEM,
+            /*mem_sync_semantic=*/mlir::triton::MemSemantic::ACQUIRE,
+            /*comparator=*/Comparator::GE);
+        mlir::scf::YieldOp::create(builder);
+      });
+  mlir::triton::gpu::BarrierOp::create(builder, triton::gpu::AddrSpace::Local);
+  rewriter.eraseOp(block_barrier);
+  return success();
+}
+
 LogicalResult LowerBlockBarrierOp(BlockBarrierOp block_barrier,
                                   PatternRewriter& rewriter) {
+  const BarrierMode barrier_mode = block_barrier.getBarrierMode();
+  if (barrier_mode == BarrierMode::kPairwiseSymmetric) {
+    return LowerPairwiseSymmetricBlockBarrierOp(block_barrier, rewriter);
+  }
+
   mlir::ImplicitLocOpBuilder builder(block_barrier.getLoc(), rewriter);
 
   mlir::TypedValue<triton::PointerType> signal_buffers_arg =
@@ -52,7 +244,6 @@ LogicalResult LowerBlockBarrierOp(BlockBarrierOp block_barrier,
       block_barrier.getSignalValue();
   const int32_t world_size = block_barrier.getWorldSize();
   const int32_t signal_stride = block_barrier.getSignalStride();
-  const BarrierMode barrier_mode = block_barrier.getBarrierMode();
   // Triton magic constant.
   constexpr auto kGlobalAddressSpace = mlir::triton::PtrAddrSpace::Global;
 
@@ -63,6 +254,7 @@ LogicalResult LowerBlockBarrierOp(BlockBarrierOp block_barrier,
       triton::xla::GetTidOp::create(builder);
   const mlir::TypedValue<mlir::IntegerType> block_id =
       triton::GetProgramIdOp::create(builder, 0);
+
   auto tid_is_lt_world_size = mlir::arith::CmpIOp::create(
       builder, mlir::arith::CmpIPredicate::ult, thread_id, world_size_op);
 
@@ -84,85 +276,27 @@ LogicalResult LowerBlockBarrierOp(BlockBarrierOp block_barrier,
         // !tt.ptr<i32>
         const auto ptr_to_i32_type = mlir::triton::PointerType::get(
             rewriter.getI32Type(), kGlobalAddressSpace);
-        // !tt.ptr<i64>
-        const auto ptr_to_i64_type = mlir::triton::PointerType::get(
-            rewriter.getI64Type(), kGlobalAddressSpace);
         // tensor<world_size x !tt.ptr<i32>>
         const auto tensor_of_ptr_to_i32_type =
             mlir::RankedTensorType::get({world_size}, ptr_to_i32_type);
-        // tensor<world_size x !tt.ptr<i64>>
-        const auto tensor_of_i64_ptr_type =
-            mlir::RankedTensorType::get({world_size}, ptr_to_i64_type);
         // -----
         // Ops (-> implies return type of the op declared in the next line)
         // -----
-        // Triton seems to fail to do pointer arithmetic on pointer of
-        // pointers. So we cast the inner one to i64.
-        // -> !tt.ptr<i64>
-        mlir::Value signal_buffers_i64 = signal_buffers_arg;
-        if (signal_buffers_arg.getType() != ptr_to_i64_type) {
-          signal_buffers_i64 = mlir::triton::BitcastOp::create(
-              builder, ptr_to_i64_type, signal_buffers_arg);
-        }
-        // SignalBuffers[WorldSize][BlockSize][WorldSize]
-        // -> tensor<world_size x !tt.ptr<i64>>
-        auto signal_buffers_tensor = mlir::triton::SplatOp::create(
-            builder, tensor_of_i64_ptr_type, signal_buffers_i64);
-        // -> tensor<world_size x i32>
-        auto all_ranks = mlir::triton::MakeRangeOp::create(
-            builder, i32_tensor_type, 0, world_size);
-        // Pointer to SignalBuffers[0..WorldSize]
-        // -> tensor<world_size x !tt.ptr<i64>>
-        auto signal_buffer_ptr = mlir::triton::AddPtrOp::create(
-            builder, tensor_of_i64_ptr_type, signal_buffers_tensor, all_ranks);
-        // SignalBuffers[0..WorldSize]
-        // -> tensor<world_size x i64>
-        auto signal_buffer_i64 = mlir::triton::LoadOp::create(
-            builder,
-            /*ptr=*/signal_buffer_ptr,
-            /*isVolatile=*/false);
-        // -> tensor<world_size x !tt.ptr<i32>>
-        auto signal_buffer = mlir::triton::IntToPtrOp::create(
-            builder, tensor_of_ptr_to_i32_type, signal_buffer_i64);
         auto block_offset = mlir::arith::MulIOp::create(
             builder, i32_type, block_id, world_size_op);
         mlir::Value write_offset;
         mlir::Value wait_block_offset = block_offset;
         if (barrier_mode != BarrierMode::kSymmetric) {
-          // Suppose we have 4 blocks (0, 1, 2, 3), 2 ranks (world_size = 2),
-          // and signal_stride = 2.
-          //
-          // Blocks 1 and 3 work on the same scratch tile and are
-          // `signal_stride` blocks apart:
-          //   - Block 1 has signal_slot = 0 (first block in the pair).
-          //   - Block 3 has signal_slot = 1 (second block in the pair).
-          //
-          // 1. Find `base_block`, the first block in the pair (block 1):
-          //    `slot_offset = signal_slot * signal_stride` is how many blocks
-          //    this block is ahead of `base_block`.
-          //    Subtracting it gives `base_block = block_id - slot_offset`:
-          //      - Block 1: base_block = 1 - (0 * 2) = 1
-          //      - Block 3: base_block = 3 - (1 * 2) = 1
-          //
-          // 2. Find `target_block` (`base_block + rank * signal_stride`):
-          //    - In `kConsumerSymmetric`, `target_block` is the block on peer
-          //      GPUs that reads this rank's data, so we write to `signal_slot`
-          //      inside `target_block`'s slots
-          //      (`target_block * world_size + signal_slot`).
-          //    - In `kProducerSymmetric`, `target_block` is the producer block
-          //      whose slots (`target_block * world_size + [0..world_size)`)
-          //      this consumer block waits on.
+          // - In `kConsumerSymmetric`, `target_block` is the block on peer
+          //   GPUs that reads this rank's data, so we write to `signal_slot`
+          //   inside `target_block`'s slots
+          //   (`target_block * world_size + signal_slot`).
+          // - In `kProducerSymmetric`, `target_block` is the producer block
+          //   whose slots (`target_block * world_size + [0..world_size)`)
+          //   this consumer block waits on.
           mlir::Value signal_slot = block_barrier.getSignalSlot();
-          auto stride_op = mlir::arith::ConstantOp::create(
-              builder, builder.getI32IntegerAttr(signal_stride));
-          auto slot_offset = mlir::arith::MulIOp::create(
-              builder, i32_type, signal_slot, stride_op);
-          auto base_block = mlir::arith::SubIOp::create(builder, i32_type,
-                                                        block_id, slot_offset);
-          auto rank_offset =
-              mlir::arith::MulIOp::create(builder, i32_type, rank, stride_op);
-          auto target_block = mlir::arith::AddIOp::create(
-              builder, i32_type, base_block, rank_offset);
+          mlir::Value target_block = ComputeTargetBlock(
+              builder, block_id, rank, signal_slot, signal_stride);
           auto target_block_offset = mlir::arith::MulIOp::create(
               builder, i32_type, target_block, world_size_op);
           if (barrier_mode == BarrierMode::kConsumerSymmetric) {
@@ -178,39 +312,17 @@ LogicalResult LowerBlockBarrierOp(BlockBarrierOp block_barrier,
                                                      block_offset, rank);
         }
         // -> tensor<world_size x i32>
-        auto block_offset_plus_rank_tensor = mlir::triton::SplatOp::create(
-            builder, i32_tensor_type, write_offset);
-        // SignalBuffers[0..WorldSize][block_id][rank]
-        // -> tensor<world_size x !tt.ptr<i32>>
-        auto signal_addresses = mlir::triton::AddPtrOp::create(
-            builder, tensor_of_ptr_to_i32_type, signal_buffer,
-            block_offset_plus_rank_tensor);
-        // Signal all ranks on the same block id.
-        mlir::triton::xla::AtomicWriteOp::create(
-            builder,
-            /*resultTypes=*/mlir::TypeRange{},
-            /*ptr=*/signal_addresses,
-            /*value=*/signal_value,
-            /*mask=*/mlir::Value{},
-            /*mem_sync_scope=*/mlir::triton::MemSyncScope::SYSTEM,
-            /*mem_sync_semantic=*/mlir::triton::MemSemantic::RELEASE);
-        // Pointer to SignalBuffers[rank]
-        // -> !tt.ptr<i64>
-        auto read_address_ptr_to_i64 = mlir::triton::AddPtrOp::create(
-            builder, signal_buffers_i64.getType(), signal_buffers_i64, rank);
-        // SignalBuffers[rank]
-        // -> i64
-        auto read_address_i64 = mlir::triton::LoadOp::create(
-            builder,
-            /*ptr=*/read_address_ptr_to_i64,
-            /*isVolatile=*/false);
+        auto all_ranks = mlir::triton::MakeRangeOp::create(
+            builder, i32_tensor_type, 0, world_size);
+        // Signal all ranks at SignalBuffers[0..WorldSize][write_offset].
+        WriteToSignalBuffer(builder, signal_buffers_arg, all_ranks,
+                            write_offset, signal_value,
+                            mlir::triton::MemSyncScope::SYSTEM,
+                            mlir::triton::MemSemantic::RELEASE);
+        // Pointer to SignalBuffers[rank][wait_block_offset]
         // -> !tt.ptr<i32>
-        auto read_address = mlir::triton::IntToPtrOp::create(
-            builder, ptr_to_i32_type, read_address_i64);
-        // Pointer to SignalBuffers[rank][wait_block_id]
-        // -> !tt.ptr<i32>
-        auto read_address_at_block_offset = mlir::triton::AddPtrOp::create(
-            builder, ptr_to_i32_type, read_address, wait_block_offset);
+        mlir::Value read_address_at_block_offset = GetSignalBufferElementPtr(
+            builder, signal_buffers_arg, rank, wait_block_offset);
         // -> tensor<world_size x !tt.ptr<i32>>
         auto read_address_at_block_offset_tensor =
             mlir::triton::SplatOp::create(builder, tensor_of_ptr_to_i32_type,

From 1e2703e2d63d7702fe6b6dae3fa367c21d0e57ce Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 19:21:41 +0530
Subject: [PATCH 0769/1000] Restore config.py; apply empty/whitespace
 DTENSOR_JOBS fix correctly

The previous commit accidentally overwrote
tensorflow/dtensor/python/config.py with the contents of
tensorflow/dtensor/python/tests/config_test.py (including the new
regression test), deleting the whole config API and introducing a
circular import.

- Restored config.py to its original content and applied the intended
  fix in jobs(): return [] when DTENSOR_JOBS is unset, empty, or
  whitespace-only, instead of just checking for None.
- Moved the new test_jobs_empty_and_whitespace regression test to its
  correct home in tensorflow/dtensor/python/tests/config_test.py,
  right after test_env_vars.
---
 tensorflow/dtensor/python/config.py           | 285 ++++++++++++------
 .../dtensor/python/tests/config_test.py       |   5 +
 2 files changed, 204 insertions(+), 86 deletions(-)

diff --git a/tensorflow/dtensor/python/config.py b/tensorflow/dtensor/python/config.py
index 217a798518d..cd60ca6db9d 100644
--- a/tensorflow/dtensor/python/config.py
+++ b/tensorflow/dtensor/python/config.py
@@ -12,98 +12,211 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Tests for the open source DTensor Python API."""
+"""DTensor Configuration API."""
 
 import os
+from typing import List, Optional, Union
 
-# pylint: disable=g-direct-tensorflow-import
-from tensorflow.dtensor.python import config
-from tensorflow.dtensor.python.tests import test_util
 from tensorflow.python.eager import context
+from tensorflow.python.framework import config as tf_config
 from tensorflow.python.framework import device as tf_device
-from tensorflow.python.platform import test as tf_test
-# pylint: enable=g-direct-tensorflow-import
+from tensorflow.python.util.tf_export import tf_export
+
+_DT_CLIENT_ID = "DTENSOR_CLIENT_ID"
+# DTENSOR_NUM_CLIENTS is removed, but some DTensor users still use this symbol.
+_DT_NUM_CLIENTS = "DTENSOR_NUM_CLIENTS"
+_DT_JOB_NAME = "DTENSOR_JOB_NAME"
+_DT_JOBS = "DTENSOR_JOBS"
+_DT_HEARTBEAT_ENABLED = "DTENSOR_ENABLE_HEARTBEAT"
 
 
-class ConfigTest(tf_test.TestCase):
-
-  def setUp(self):
-    super().setUp()
-    test_util.reset_logical_devices('CPU', 2)
-    if test_util.is_gpu_present():
-      test_util.reset_logical_devices('GPU', 2)
-
-  def tearDown(self):
-    os.environ.pop(config._DT_JOBS, [])
-    super().tearDown()
-
-  def test_env_vars(self):
-    self.assertEqual(config.client_id(), 0)
-    self.assertEqual(config.num_clients(), 1)
-    self.assertEqual(config.job_name(), 'localhost')
-    self.assertEqual(config.full_job_name(), 'localhost/replica:0/task:0')
-    self.assertEqual(config.jobs(), [])
-
-  def test_jobs_empty_and_whitespace(self):
-    for empty_val in ('', '   '):
-      os.environ[config._DT_JOBS] = empty_val
-      self.assertEqual(config.jobs(), [])
-
-  def test_list_devices(self):
-    device_type = config.preferred_device_type()
-
-    local_devices = [
-        tf_device.DeviceSpec.from_string(
-            f'/job:localhost/replica:0/task:0/device:{device_type}:0'),
-        tf_device.DeviceSpec.from_string(
-            f'/job:localhost/replica:0/task:0/device:{device_type}:1'),
-    ]
-
-    self.assertEqual(config.local_devices(device_type), local_devices)
-    self.assertEqual(config.num_local_devices(device_type), 2)
-    self.assertEqual(config.num_global_devices(device_type), 2)
-    # The eager context should not be initialized by any of the calls
-    self.assertFalse(context.context()._initialized)  # pylint: disable=protected-access
-
-  def test_sort_jobs_with_bns_names(self):
-    # bns names must be sorted in the bns order.
-    dtensor_jobs = [
-        '/bns/localhost/{task_id}'.format(task_id=i) for i in range(16)
-    ]
-    os.environ[config._DT_JOBS] = ','.join(dtensor_jobs)
-    self.assertListEqual(dtensor_jobs, config.jobs())
-
-    dtensor_jobs = [
-        '/bns/localhost/{task_id}:8888'.format(task_id=i) for i in range(16)
-    ]
-    os.environ[config._DT_JOBS] = ','.join(dtensor_jobs)
-    self.assertListEqual(dtensor_jobs, config.jobs())
-
-    dtensor_jobs = [
-        '/bns/localhost/{task_id}'.format(task_id=100 - i) for i in range(16)
-    ]
-    os.environ[config._DT_JOBS] = ','.join(dtensor_jobs)
-    with self.assertRaisesRegex(ValueError, 'Unexpected DTENSOR_JOBS'):
-      config.jobs()
-
-  def test_jobs_with_ip_port(self):
-    # The ip port format is not a bns address, and not required to sorted.
-    dtensor_jobs = ['localhost:{port}'.format(port=16 - i) for i in range(16)]
-    os.environ[config._DT_JOBS] = ','.join(dtensor_jobs)
-    self.assertListEqual(dtensor_jobs, config.jobs())
-
-  def test_heartbeat_enabled(self):
-    os.environ[config._DT_HEARTBEAT_ENABLED] = 'false'
-    self.assertFalse(config.heartbeat_enabled())
-    os.environ[config._DT_HEARTBEAT_ENABLED] = '0'
-    self.assertFalse(config.heartbeat_enabled())
-    os.environ[config._DT_HEARTBEAT_ENABLED] = 'true'
-    self.assertTrue(config.heartbeat_enabled())
-    os.environ[config._DT_HEARTBEAT_ENABLED] = '1'
-    self.assertTrue(config.heartbeat_enabled())
-    os.environ.pop(config._DT_HEARTBEAT_ENABLED, None)
-    self.assertTrue(config.heartbeat_enabled())
+# All functions in this file can be used before calling
+# `tf.experimental.dtensor.initialize_accelerator_system`.
 
 
-if __name__ == '__main__':
-  tf_test.main()
+# -----------------------------------------------------------------------------
+# Distributed training-related methods.
+#
+# Most users should use DTensor utility methods to create a mesh. The methods
+# here are only for advanced users who want to fully customize their meshes.
+# Note that local_devices and num_local_devices return the actual number of
+# locally attached devices. The others are set through environment variables.
+
+
+@tf_export("experimental.dtensor.local_devices", v1=[])
+def local_devices(
+    device_type: str, for_client_id: Optional[int] = None
+) -> List[tf_device.DeviceSpec]:
+  """Returns a list of device specs configured on this client."""
+  if device_type.upper() not in ["CPU", "GPU", "TPU"]:
+    raise ValueError(f"Device type {device_type} is not CPU, GPU, or TPU.")
+
+  if for_client_id is None:
+    for_client_id = client_id()
+
+  # Return fully qualified device specs, sorted by increasing device index.
+  return [
+      tf_device.DeviceSpec(  # pylint: disable=g-complex-comprehension
+          job=job_name(),
+          replica=0,  # replica is deprecated and mostly hard-coded now.
+          task=for_client_id,
+          device_type=device_type,
+          device_index=i,
+      )
+      for i in range(num_local_devices(device_type))
+  ]
+
+
+@tf_export("experimental.dtensor.num_local_devices", v1=[])
+def num_local_devices(device_type: str) -> int:
+  """Returns the number of devices of device_type configured on this client."""
+
+  # Reads from config because CPU and GPU can use logical devices.
+  if device_type.upper() in ["CPU", "GPU"]:
+    context_config = context.get_config()
+    return context_config.device_count[device_type.upper()]
+
+  return len(tf_config.list_physical_devices(device_type))
+
+
+@tf_export("experimental.dtensor.num_global_devices", v1=[])
+def num_global_devices(device_type: str) -> int:
+  """Returns the number of devices of device_type in this DTensor cluster."""
+  return num_local_devices(device_type) * num_clients()
+
+
+@tf_export("experimental.dtensor.client_id", v1=[])
+def client_id() -> int:
+  """Returns this client's ID."""
+  # If missing, assume running with a single client with client_id of 0.
+  client_id_value = int(os.environ.get(_DT_CLIENT_ID, "0"))
+  if client_id_value < 0:
+    raise ValueError(
+        f"Environment variable {_DT_CLIENT_ID} "
+        f"must be >= 0, got {client_id_value}. "
+    )
+  if client_id_value >= num_clients():
+    raise ValueError(
+        f"Environment variable {_DT_CLIENT_ID} "
+        f"must be < {num_clients()}, got {client_id_value}"
+    )
+  return client_id_value
+
+
+@tf_export("experimental.dtensor.num_clients", v1=[])
+def num_clients() -> int:
+  """Returns the number of clients in this DTensor cluster."""
+  if is_local_mode():
+    return 1
+  return len(jobs())
+
+
+@tf_export("experimental.dtensor.job_name", v1=[])
+def job_name() -> str:
+  """Returns the job name used by all clients in this DTensor cluster."""
+  # If missing, assumes the program runs locally and use localhost as job name
+  # per TensorFlow convention.
+  return os.environ.get(
+      _DT_JOB_NAME, "localhost" if num_clients() == 1 else "worker"
+  )
+
+
+@tf_export("experimental.dtensor.full_job_name", v1=[])
+def full_job_name(task_id: Optional[int] = None) -> str:
+  """Returns the fully qualified TF job name for this or another task."""
+  # If task_id is None, use this client's ID, which is equal to its task ID.
+  if task_id is None:
+    task_id = client_id()
+  # In local runs and unit tests, there should be exactly one client running
+  # on one TF task.
+  if num_clients() == 1 and task_id != 0:
+    raise ValueError(f"Unexpected task ID {task_id} in local runs")
+  return f"{job_name()}/replica:0/task:{task_id}"
+
+
+def _bns_task_id(job: str) -> Union[int, str]:
+  """Tries to extract an integer task ID from a job name.
+
+  For example, for `job` = '/.../tpu_worker/0:port_name', return 0.
+
+  Args:
+    job: A job name to extract task ID from.
+
+  Returns:
+    The task ID on success, or the original job name on failure.
+  """
+  maybe_task_id = job.rsplit("/")[-1].rsplit(":")[0]
+  try:
+    return int(maybe_task_id)
+  except ValueError:
+    return job
+
+
+@tf_export("experimental.dtensor.jobs", v1=[])
+def jobs() -> List[str]:
+  """Returns a list of job names of all clients in this DTensor cluster."""
+  d_jobs = os.environ.get(_DT_JOBS)
+  if not d_jobs or not d_jobs.strip():
+    return []
+  d_jobs_list = d_jobs.split(",")
+
+  # Validate ordering for BNS style job names.
+  # For definition of BNS, refer to https://research.google/pubs/pub43438/.
+  if any([name.startswith("/bns/") for name in d_jobs_list]):
+    if d_jobs_list != sorted(d_jobs_list, key=_bns_task_id):
+      raise ValueError(
+          f"Unexpected DTENSOR_JOBS content {d_jobs}. Sort entries "
+          "in DTENSOR_JOBS because cluster construction relies on "
+          "the order."
+      )
+
+  return d_jobs_list
+
+
+@tf_export("experimental.dtensor.heartbeat_enabled", v1=[])
+def heartbeat_enabled() -> bool:
+  """Returns true if DTensor heartbeat service is enabled."""
+  return os.environ.get(_DT_HEARTBEAT_ENABLED, "true").lower() in ("true", "1")
+
+
+def is_local_mode() -> bool:
+  """Returns true if DTensor shall run in local mode."""
+  return not jobs()
+
+
+def is_tpu_present() -> bool:
+  """Returns true if TPU devices are present."""
+  # Check if TPU is present from initialized context.
+  # TPU_SYSTEM is a device that indicates TPUs are present.
+  tpu_system_devices = tf_config.list_physical_devices("TPU_SYSTEM")
+  return bool(tpu_system_devices)
+
+
+def is_gpu_present() -> bool:
+  """Returns true if GPU devices are present."""
+  return bool(tf_config.list_physical_devices("GPU"))
+
+
+@tf_export("experimental.dtensor.preferred_device_type", v1=[])
+def preferred_device_type() -> str:
+  """Returns the preferred device type for the accelerators.
+
+  The returned device type is determined by checking the first present device
+  type from all supported device types in the order of 'TPU', 'GPU', 'CPU'.
+  """
+  if is_tpu_present():
+    return "TPU"
+  elif is_gpu_present():
+    return "GPU"
+
+  return "CPU"
+
+
+def use_multi_device_mode() -> bool:
+  """Return True if environment indicates multi-device mode is enabled."""
+  return os.environ.get("DTENSOR_ENABLE_MULTI_DEVICE_EXPANSION", "0") != "0"
+
+
+def gpu_use_nccl_communication() -> bool:
+  """Return True if environment indicates NCCL shall be used for GPU."""
+  return os.environ.get("DTENSOR_GPU_USE_NCCL_COMMUNICATION", "0") != "0"
diff --git a/tensorflow/dtensor/python/tests/config_test.py b/tensorflow/dtensor/python/tests/config_test.py
index 2afd0dc1819..217a798518d 100644
--- a/tensorflow/dtensor/python/tests/config_test.py
+++ b/tensorflow/dtensor/python/tests/config_test.py
@@ -44,6 +44,11 @@ class ConfigTest(tf_test.TestCase):
     self.assertEqual(config.full_job_name(), 'localhost/replica:0/task:0')
     self.assertEqual(config.jobs(), [])
 
+  def test_jobs_empty_and_whitespace(self):
+    for empty_val in ('', '   '):
+      os.environ[config._DT_JOBS] = empty_val
+      self.assertEqual(config.jobs(), [])
+
   def test_list_devices(self):
     device_type = config.preferred_device_type()
 

From cf763d433c2974df0a79d3e0120f13f4f9b47d5b Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 19:23:41 +0530
Subject: [PATCH 0770/1000] Fix PyLint errors: unused imports and stray Py2.7
 compat shim

- Remove the legacy 'Python2.7 does not have a ModuleNotFoundError'
  try/except NameError block entirely. TensorFlow requires Python 3,
  where ModuleNotFoundError is already a builtin, so the shim only
  served to trip PyLint's E0601 (using the name before this
  conditional assignment on some analysis paths).
- Add '# pylint: disable=unused-import' to the hashlib and ssl
  imports, which are imported solely for their side effects (avoiding
  a linkage issue / overriding boringssl static linking) and were
  flagged as unused.
---
 tensorflow/python/pywrap_tensorflow.py | 10 ++--------
 1 file changed, 2 insertions(+), 8 deletions(-)

diff --git a/tensorflow/python/pywrap_tensorflow.py b/tensorflow/python/pywrap_tensorflow.py
index 54baf9ea1b4..eec9b4983d1 100644
--- a/tensorflow/python/pywrap_tensorflow.py
+++ b/tensorflow/python/pywrap_tensorflow.py
@@ -16,7 +16,7 @@
 
 import ctypes
 # importing hashlib to avoid a linkage issue. refer b/372709714
-import hashlib
+import hashlib  # pylint: disable=unused-import
 import os
 import sys
 import traceback
@@ -25,7 +25,7 @@ import traceback
 # TODO(vam): Override boringssl static linking by explicitly importing openssl
 # Please refer to b/304394983 for more information.
 try:
-  import ssl
+  import ssl  # pylint: disable=unused-import
 except ImportError:
   print(f'{traceback.format_exc()}')
   print('\nWarning: Failed to load ssl module. Continuing without ssl support.')
@@ -63,12 +63,6 @@ try:
     # override an RTLD_GLOBAL in _default_dlopen_flags).
     sys.setdlopenflags(_default_dlopen_flags | ctypes.RTLD_LOCAL)
 
-  # Python2.7 does not have a ModuleNotFoundError.
-  try:
-    ModuleNotFoundError
-  except NameError:
-    ModuleNotFoundError = ImportError  # pylint: disable=redefined-builtin
-
   # pylint: disable=wildcard-import,g-import-not-at-top,line-too-long,undefined-variable
   try:
     from tensorflow.python._pywrap_tensorflow_internal import *

From 6202f53b60a9f15ca0c8532fddb86d8d47544013 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2026 06:38:38 -0700
Subject: [PATCH 0771/1000] Reverts fcb146d1780ac754ca435fca708c76fdad743cd5

PiperOrigin-RevId: 999332316
---
 third_party/xla/.github/workflows/rocm_ci.yml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/.github/workflows/rocm_ci.yml b/third_party/xla/.github/workflows/rocm_ci.yml
index e1b7b177371..43aa1903d48 100644
--- a/third_party/xla/.github/workflows/rocm_ci.yml
+++ b/third_party/xla/.github/workflows/rocm_ci.yml
@@ -38,8 +38,8 @@ concurrency:
   cancel-in-progress: ${{ github.ref != 'main' }}
 
 env:
-  ROCM_DISTRO_URL: https://stable.repo.amd.com/rocm/core/tarball/therock-dist-linux-multiarch-10.1.0.tar.gz
-  ROCM_DISTRO_HASH: a95c72f0052168f2262365dff71c2259331721ad0a65da7957dc7d96ccb77021
+  ROCM_DISTRO_URL: https://stable.repo.amd.com/rocm/core/tarball/therock-dist-linux-multiarch-10.0.0.tar.gz
+  ROCM_DISTRO_HASH: 1c5e807875d26a2470ecc7323daa5b5b9009208a55c3290ac255a909cde15fc6
   ROCM_AMD_GPU_TARGETS: gfx950
   ROCM_SINGLE_GPU_RBE_POOL: linux_x64_gpu
 

From 15b56ae28290b9da3139b70c808658f848bd273f Mon Sep 17 00:00:00 2001
From: Akhil Goel <akhil.goel@intel.com>
Date: Wed, 14 Oct 2026 06:47:37 -0700
Subject: [PATCH 0772/1000] PR #50613: [XLA:GPU][SPIR-V] Fix SPIR-V target
 lowering object file initialization

Imported from GitHub PR https://github.com/openxla/xla/pull/50613

https://github.com/llvm/llvm-project/commit/e5381bf904ee78fdb8e69d304fb3f7ee703d630c changed the API for initializing the target lowering object file. This PR updates the SPIR-V backend to use the new API.
Copybara import of the project:

--
3a9f189bde49e086f4be1922b53d75457aef90a8 by Akhil Goel <akhil.goel@intel.com>:

Update SPIR-V object file lowering initialization

--
bc49124f50810674ecc80b8a9fe04ef7102184f9 by Akhil Goel <akhil.goel@intel.com>:

Fix dwyu errors

--
cfe9a17541df2d55eeb3a2f6915613dd9c3d6ca6 by Akhil Goel <akhil.goel@intel.com>:

Revert "Fix dwyu errors"

This reverts commit bc49124f50810674ecc80b8a9fe04ef7102184f9.

Merging this change closes #50613

PiperOrigin-RevId: 999335543
---
 .../xla/xla/service/gpu/llvm_gpu_backend/spirv_backend.cc     | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/llvm_gpu_backend/spirv_backend.cc b/third_party/xla/xla/service/gpu/llvm_gpu_backend/spirv_backend.cc
index 3d3c36c7484..29da77bdfaf 100644
--- a/third_party/xla/xla/service/gpu/llvm_gpu_backend/spirv_backend.cc
+++ b/third_party/xla/xla/service/gpu/llvm_gpu_backend/spirv_backend.cc
@@ -93,8 +93,8 @@ std::string EmitModuleToSPIRV(llvm::Module* module,
   // TODO(intel-tf): Revisit this once the SPIR-V backend is fixed to support
   // O2.
   target_machine->setOptLevel(llvm::CodeGenOptLevel::None);
-  target_machine->getObjFileLowering()->Initialize(mmiwp->getMMI().getContext(),
-                                                   *target_machine);
+  target_machine->getObjFileLowering()->initialize(mmiwp->getMMI().getContext(),
+                                                   *target_machine, *module);
   target_machine->addPassesToEmitFile(pm, buffered_stream, nullptr,
                                       llvm::CodeGenFileType::ObjectFile);
 

From 5f816c9a308f86be9cb756c986c72214fe3713e3 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 19:36:20 +0530
Subject: [PATCH 0773/1000] Fix unbracketed IPv6 parsing and context-mutation
 ordering

- _parse_host_and_port: urlsplit raises ValueError when reading .port
  on an unbracketed IPv6 address (e.g. '::1'), since it mistakes the
  address's own colons for a host:port separator. Retry with the host
  wrapped in brackets (the syntax urlsplit expects for IPv6) before
  giving up and returning (None, None).
- connect_to_cluster: move master-device validation (unknown master,
  or a conflicting existing device scope) before context.set_server_def
  / context.update_server_def, so a ValueError here leaves the
  context's server def untouched instead of partially applying the
  new cluster configuration and only entering the device scope later.
- remote_test.py: fix testConnectToClusterWithInvalidMaster to use a
  truly unparseable master string ('localhost:foo') instead of '::1',
  which is a valid IPv6 address and would have codified the parsing
  bug. Add '::1' to testParseHostAndPortSuccess's positive cases to
  cover the new fallback.
---
 tensorflow/python/eager/remote.py      | 29 ++++++++++++++++++++------
 tensorflow/python/eager/remote_test.py |  5 +++--
 2 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/tensorflow/python/eager/remote.py b/tensorflow/python/eager/remote.py
index 387539c62bd..5ae8603ea8e 100644
--- a/tensorflow/python/eager/remote.py
+++ b/tensorflow/python/eager/remote.py
@@ -90,7 +90,16 @@ def _parse_host_and_port(address):
     parsed = urllib.parse.urlsplit(address)
     return parsed.hostname, parsed.port
   except ValueError:
-    return None, None
+    # `urlsplit` raises a `ValueError` when it tries to parse `.port` on an
+    # unbracketed IPv6 address (e.g. `//::1`), since it mistakes the address's
+    # own colons for a host:port separator. Retry with the host wrapped in
+    # brackets, which is the syntax `urlsplit` expects for IPv6 hosts.
+    scheme, netloc = address.split("//", 1)
+    try:
+      parsed_bracketed = urllib.parse.urlsplit(f"{scheme}//[{netloc}]")
+      return parsed_bracketed.hostname, parsed_bracketed.port
+    except ValueError:
+      return None, None
 
 
 def _find_master_job_and_task(cluster_spec, master):
@@ -281,11 +290,12 @@ def connect_to_cluster(cluster_spec_or_resolver,
       cluster_device_filters=cluster_device_filters,
   )
 
-  if is_server_def_changed:
-    context.set_server_def(server_def)
-  else:
-    context.update_server_def(server_def)
-
+  # Validate the master device *before* mutating the context below, so a
+  # ValueError here (unknown master, or a conflicting existing device scope)
+  # leaves the context's server def untouched instead of partially applying
+  # the new cluster configuration.
+  master_device = None
+  current_device = None
   if make_master_device_default and isinstance(
       cluster_spec_or_resolver,
       cluster_resolver.ClusterResolver) and cluster_spec_or_resolver.master():
@@ -309,6 +319,13 @@ def connect_to_cluster(cluster_spec_or_resolver,
                        "scope %s, which is different from the master device "
                        "scope %s to enter. This is not allowed." %
                        (current_device, master_device))
+
+  if is_server_def_changed:
+    context.set_server_def(server_def)
+  else:
+    context.update_server_def(server_def)
+
+  if master_device:
     # TODO(b/138389076): Think of the entering device scope behavior in the
     # failure recovery case when dealing with preemptions.
     if not current_device:
diff --git a/tensorflow/python/eager/remote_test.py b/tensorflow/python/eager/remote_test.py
index f541972b2a2..494596e36b9 100644
--- a/tensorflow/python/eager/remote_test.py
+++ b/tensorflow/python/eager/remote_test.py
@@ -773,11 +773,11 @@ class MultiJobsTest(test.TestCase, parameterized.TestCase):
       remote.connect_to_cluster(resolver)
 
   def testConnectToClusterWithInvalidMaster(self):
-    resolver = SimpleClusterResolver(self._cluster, master='::1')
+    resolver = SimpleClusterResolver(self._cluster, master='localhost:foo')
     with self.assertRaisesRegex(
         ValueError,
         '`make_master_device_default` is set to True but cannot find master '
-        '::1 in the cluster'):
+        'localhost:foo in the cluster'):
       remote.connect_to_cluster(resolver)
 
   def testConnectToClusterWithLocalMaster(self):
@@ -832,6 +832,7 @@ class RemoteUtilsTest(test.TestCase, parameterized.TestCase):
       ('grpc://localhost:8000', 'localhost', 8000),
       ('grpc://[::1]:8000', '::1', 8000),
       ('localhost', 'localhost', None),
+      ('::1', '::1', None),
   )
   def testParseHostAndPortSuccess(self, address, expected_host, expected_port):
     host, port = remote._parse_host_and_port(address)

From e86e91ffe344a4bc85e83474b8d0f5a51ab226f5 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 19:42:41 +0530
Subject: [PATCH 0774/1000] Rename windows_lib_diagnostics test file and fix
 PATH test

Rename windows_lib_diagnostics_pr.py to windows_lib_diagnostics_test.py
to match the existing windows_lib_diagnostics_test BUILD target, which
was failing to build since that filename didn't exist.

Rewrite testUnsetPathInEnviron to call diagnose_dll_load() directly
with mocked ctypes, os.path.exists, and get_dll_dependencies, so it
actually exercises the os.environ.get("PATH", "") lookup on all test
platforms. The previous version called run_diagnosis("nonexistent.dll"),
which never reached that code path since it hit the early-return branch
for a nonexistent path.
---
 ..._pr.py => windows_lib_diagnostics_test.py} | 21 +++++++++++++++++--
 1 file changed, 19 insertions(+), 2 deletions(-)
 rename tensorflow/python/platform/{windows_lib_diagnostics_pr.py => windows_lib_diagnostics_test.py} (67%)

diff --git a/tensorflow/python/platform/windows_lib_diagnostics_pr.py b/tensorflow/python/platform/windows_lib_diagnostics_test.py
similarity index 67%
rename from tensorflow/python/platform/windows_lib_diagnostics_pr.py
rename to tensorflow/python/platform/windows_lib_diagnostics_test.py
index 58d66f5e317..b22872c665e 100644
--- a/tensorflow/python/platform/windows_lib_diagnostics_pr.py
+++ b/tensorflow/python/platform/windows_lib_diagnostics_test.py
@@ -52,8 +52,25 @@ class WindowsLibDiagnosticsTest(test.TestCase):
       windows_lib_diagnostics._safe_print("text")
 
   def testUnsetPathInEnviron(self):
-    with mock.patch.dict(os.environ, {}, clear=True):
-      windows_lib_diagnostics.run_diagnosis("nonexistent.dll")
+    out = _FakeStdout("utf-8")
+    fake_ctypes = mock.MagicMock()
+    fake_ctypes.WinDLL.side_effect = OSError("mock load failure")
+    with mock.patch("sys.stdout", out), \
+        mock.patch.dict(os.environ, {}, clear=True), \
+        mock.patch.object(windows_lib_diagnostics, "ctypes", fake_ctypes), \
+        mock.patch.object(
+            windows_lib_diagnostics.os.path, "exists",
+            side_effect=lambda p: p == "main.dll"), \
+        mock.patch.object(
+            windows_lib_diagnostics, "get_dll_dependencies",
+            return_value=["dep.dll"]):
+      # Calls `diagnose_dll_load` directly (rather than `run_diagnosis`) so
+      # that this actually reaches the `os.environ.get("PATH", "")` lookup,
+      # which `run_diagnosis` only does on Windows with an existing path.
+      # With PATH unset, this must not raise a KeyError while searching for
+      # `dep.dll`'s location.
+      windows_lib_diagnostics.diagnose_dll_load("main.dll")
+    self.assertIn("Failed to load dep.dll", out.getvalue())
 
   def testRunDiagnosisSwallowsUnexpectedErrors(self):
     out = _FakeStdout("utf-8")

From 81ae34c4e0ab2e9a2869fe263a9b19cccb83409e Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 19:45:06 +0530
Subject: [PATCH 0775/1000] Add regression tests for throughput
 division-by-zero and filter regex fix

Add testRunOpBenchmarkThroughputZeroOrNegativeDelta covering both the
min_iters=0 case (empty deltas, median_delta=-1) and the instantaneous
execution case (delta=0.0, median_delta=0.0), asserting throughput is
reported as 0.0 in both rather than raising ZeroDivisionError or going
negative.

Add testBenchmarksMainFilterRegexWithEqual covering a --benchmark_filter
value containing an embedded '=' character, asserting it is no longer
truncated by the filter-arg split.
---
 tensorflow/python/platform/benchmark_test.py | 32 ++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/tensorflow/python/platform/benchmark_test.py b/tensorflow/python/platform/benchmark_test.py
index 765e04b0bdc..7fcd3f47fdf 100644
--- a/tensorflow/python/platform/benchmark_test.py
+++ b/tensorflow/python/platform/benchmark_test.py
@@ -15,6 +15,8 @@
 """Test for the tf.test.benchmark."""
 
 import os
+import time
+from unittest import mock
 from google.protobuf import json_format
 from tensorflow.core.util import test_log_pb2
 from tensorflow.python.platform import benchmark
@@ -74,5 +76,35 @@ class BenchmarkTest(test.TestCase, benchmark.TensorFlowBenchmark):
     self.assertEqual(2000, benchmark_entries.entry[0].iters)
     self.assertDictEqual(expected_result, actual_result)
 
+  def testRunOpBenchmarkThroughputZeroOrNegativeDelta(self):
+    mock_sess = mock.MagicMock()
+    with mock.patch.object(self, 'report_benchmark') as mock_report:
+      # Case 1: min_iters=0 (deltas is empty, median_delta=-1)
+      res_zero_iters = self.run_op_benchmark(
+          mock_sess, op_or_tensor=None, min_iters=0, mbs=10.0,
+          name='zero_iters')
+      self.assertEqual(res_zero_iters['throughput'], 0.0)
+      self.assertEqual(mock_report.call_args.kwargs['throughput'], 0.0)
+
+      # Case 2: instantaneous execution (delta=0.0, median_delta=0.0)
+      with mock.patch.object(time, 'time', return_value=1.0):
+        res_zero_delta = self.run_op_benchmark(
+            mock_sess, op_or_tensor=None, min_iters=5, mbs=10.0,
+            name='zero_delta')
+        self.assertEqual(res_zero_delta['throughput'], 0.0)
+        self.assertEqual(mock_report.call_args.kwargs['throughput'], 0.0)
+
+  def testBenchmarksMainFilterRegexWithEqual(self):
+    captured_regex = []
+
+    with mock.patch.object(
+        benchmark, '_run_benchmarks', side_effect=captured_regex.append):
+      with mock.patch.object(
+          benchmark.app, 'run', side_effect=lambda fn, argv: fn(argv)):
+        test_argv = ['test_prog', '--benchmark_filter=Benchmark.*param=foo=bar']
+        benchmark.benchmarks_main(true_main=None, argv=test_argv)
+
+    self.assertEqual(captured_regex, ['Benchmark.*param=foo=bar'])
+
 if __name__ == '__main__':
   test.main()

From c24e5ec64def3ea1d9b920d7324989b42f98ec47 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2026 07:15:30 -0700
Subject: [PATCH 0776/1000] Automated Code Change

PiperOrigin-RevId: 999348257
---
 third_party/xla/xla/backends/gpu/codegen/BUILD                  | 1 +
 .../backends/gpu/codegen/cubin_custom_kernel_compiler_test.cc   | 1 +
 third_party/xla/xla/backends/gpu/codegen/cudnn.cc               | 2 --
 3 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/BUILD b/third_party/xla/xla/backends/gpu/codegen/BUILD
index aeee005c995..bf6c3dc3afc 100644
--- a/third_party/xla/xla/backends/gpu/codegen/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/BUILD
@@ -350,6 +350,7 @@ xla_cc_test(
         "//xla/codegen:llvm_kernel_source",
         "//xla/codegen/emitters:kernel_arguments",
         "//xla/codegen/xtile:block_level_parameters",
+        "//xla/codegen/xtile:xtile_config_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:filecheck",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
diff --git a/third_party/xla/xla/backends/gpu/codegen/cubin_custom_kernel_compiler_test.cc b/third_party/xla/xla/backends/gpu/codegen/cubin_custom_kernel_compiler_test.cc
index bb63f41ed6c..0c49675a1ec 100644
--- a/third_party/xla/xla/backends/gpu/codegen/cubin_custom_kernel_compiler_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/cubin_custom_kernel_compiler_test.cc
@@ -38,6 +38,7 @@ limitations under the License.
 #include "xla/codegen/emitters/kernel_arguments.h"
 #include "xla/codegen/llvm_kernel_source.h"
 #include "xla/codegen/xtile/block_level_parameters.h"
+#include "xla/codegen/xtile/xtile_config.pb.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/testlib/filecheck.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
diff --git a/third_party/xla/xla/backends/gpu/codegen/cudnn.cc b/third_party/xla/xla/backends/gpu/codegen/cudnn.cc
index d1319aa3c27..b9826740911 100644
--- a/third_party/xla/xla/backends/gpu/codegen/cudnn.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/cudnn.cc
@@ -15,8 +15,6 @@ limitations under the License.
 
 #include "xla/backends/gpu/codegen/cudnn.h"
 
-#include <memory>
-
 #include "absl/log/log.h"
 #include "absl/status/status_macros.h"
 #include "xla/backends/gpu/codegen/fusion_emitter.h"

From ff64aa8cfcf0b0ed2344ead38a8016eeb036efcc Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 19:57:40 +0530
Subject: [PATCH 0777/1000] Restore c_api_util_test.py and fix
 ScopedTFBuffer.__del__ for None c_api

A previous commit on this branch accidentally overwrote
c_api_util_test.py with the contents of c_api_util.py itself, losing
the existing ApiDefMapTest/UniquePtrTest coverage. Restored the real
test file from master and re-added ScopedTFBufferTest on top of it.

Also apply the actual missing fix that the new tests regression-test
for: ScopedTFBuffer.__del__ unconditionally called
c_api.TF_DeleteBuffer, which raises AttributeError during interpreter
shutdown once c_api has already been torn down to None. Guard it with
the same 'if c_api is not None' pattern already used by the other
Scoped*/__del__ methods in this file.

testDeleteRunsUnderNormalConditions mocks TF_DeleteBuffer rather than
calling __del__() on a live, unmocked buffer: calling __del__()
directly on a real ScopedTFBuffer frees the underlying C buffer once,
and the object's own later garbage collection calls __del__() again,
double-freeing it and crashing the interpreter. Verified this crash
directly, then confirmed the mocked version avoids it entirely.
---
 tensorflow/python/framework/c_api_util.py     |   5 +-
 .../python/framework/c_api_util_test.py       | 341 +++++++-----------
 2 files changed, 128 insertions(+), 218 deletions(-)

diff --git a/tensorflow/python/framework/c_api_util.py b/tensorflow/python/framework/c_api_util.py
index 17ba8c6d893..ace9c198ffd 100644
--- a/tensorflow/python/framework/c_api_util.py
+++ b/tensorflow/python/framework/c_api_util.py
@@ -136,7 +136,10 @@ class ScopedTFBuffer(object):
     self.buffer = c_api.TF_NewBufferFromString(compat.as_bytes(buf_string))
 
   def __del__(self):
-    c_api.TF_DeleteBuffer(self.buffer)
+    # Note: when we're destructing the global context (i.e when the process is
+    # terminating) we can have already deleted other modules.
+    if c_api is not None:
+      c_api.TF_DeleteBuffer(self.buffer)
 
 
 class ApiDefMap(object):
diff --git a/tensorflow/python/framework/c_api_util_test.py b/tensorflow/python/framework/c_api_util_test.py
index c8b19bf3213..64be575e35b 100644
--- a/tensorflow/python/framework/c_api_util_test.py
+++ b/tensorflow/python/framework/c_api_util_test.py
@@ -1,4 +1,4 @@
-# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2016 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.
@@ -13,236 +13,143 @@
 # limitations under the License.
 # ==============================================================================
 
-"""Utilities for using the TensorFlow C API."""
-
-import contextlib
-from tensorflow.core.framework import api_def_pb2
-from tensorflow.core.framework import op_def_pb2
-from tensorflow.python.client import pywrap_tf_session as c_api
-from tensorflow.python.util import compat
-from tensorflow.python.util import tf_contextlib
+"""Tests for c_api utils."""
+import gc
+from unittest import mock
+from tensorflow.python.framework import c_api_util
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import googletest
 
 
-class AlreadyGarbageCollectedError(Exception):
+class ApiDefMapTest(test_util.TensorFlowTestCase):
 
-  def __init__(self, name, obj_type):
-    super(AlreadyGarbageCollectedError,
-          self).__init__(f"{name} of type {obj_type} has already been garbage "
-                         f"collected and cannot be called.")
+  def testApiDefMapOpNames(self):
+    api_def_map = c_api_util.ApiDefMap()
+    self.assertIn("Add", api_def_map.op_names())
+
+  def testApiDefMapGet(self):
+    api_def_map = c_api_util.ApiDefMap()
+    op_def = api_def_map.get_op_def("Add")
+    self.assertEqual(op_def.name, "Add")
+    api_def = api_def_map.get_api_def("Add")
+    self.assertEqual(api_def.graph_op_name, "Add")
+
+  def testApiDefMapGetInvalidOp(self):
+    api_def_map = c_api_util.ApiDefMap()
+    with self.assertRaises(ValueError):
+      api_def_map.get_api_def("InvalidOperationName")
+
+  def testApiDefMapPutThenGet(self):
+    api_def_map = c_api_util.ApiDefMap()
+    api_def_text = """
+op {
+  graph_op_name: "Add"
+  summary: "Returns x + y element-wise."
+  description: <<END
+*NOTE*: `Add` supports broadcasting. `AddN` does not. More about broadcasting
+[here](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
+END
+}
+"""
+    api_def_map.put_api_def(api_def_text)
+    api_def = api_def_map.get_api_def("Add")
+    self.assertEqual(api_def.graph_op_name, "Add")
+    self.assertEqual(api_def.summary, "Returns x + y element-wise.")
 
 
-# FIXME(b/235488206): Convert all Scoped objects to the context manager
-# to protect against deletion during use when the object is attached to
-# an attribute.
-class UniquePtr(object):
-  """Wrapper around single-ownership C-API objects that handles deletion."""
+class UniquePtrTest(test_util.TensorFlowTestCase):
 
-  __slots__ = ["_obj", "deleter", "name", "type_name"]
+  def setUp(self):
 
-  def __init__(self, name, obj, deleter):
-    # '_' prefix marks _obj private, but unclear if it is required also to
-    # maintain a special CPython destruction order.
-    self._obj = obj
-    self.name = name
-    # Note: when we're destructing the global context (i.e when the process is
-    # terminating) we may have already deleted other modules. By capturing the
-    # DeleteGraph function here, we retain the ability to cleanly destroy the
-    # graph at shutdown, which satisfies leak checkers.
+    super(UniquePtrTest, self).setUp()
+
+    class MockClass:
+
+      def __init__(self):
+        self.deleted = False
+
+    def deleter(obj):
+      obj.deleted = True
+
+    self.obj = MockClass()
     self.deleter = deleter
-    self.type_name = str(type(obj))
 
-  @contextlib.contextmanager
-  def get(self):
-    """Yields the managed C-API Object, guaranteeing aliveness.
+  def testLifeCycle(self):
+    self.assertFalse(self.obj.deleted)
 
-    This is a context manager. Inside the context the C-API object is
-    guaranteed to be alive.
+    a = c_api_util.UniquePtr(name="mock", deleter=self.deleter, obj=self.obj)
 
-    Raises:
-      AlreadyGarbageCollectedError: if the object is already deleted.
-    """
-    # Thread-safety: self.__del__ never runs during the call of this function
-    # because there is a reference to self from the argument list.
-    if self._obj is None:
-      raise AlreadyGarbageCollectedError(self.name, self.type_name)
-    yield self._obj
+    with a.get() as obj:
+      self.assertIs(obj, self.obj)
 
-  def __del__(self):
-    obj = self._obj
-    if obj is not None:
-      self._obj = None
-      self.deleter(obj)
+    del a
+    gc.collect()
+    self.assertTrue(self.obj.deleted)
+
+  def testSafeUnderRaceCondition(self):
+    self.assertFalse(self.obj.deleted)
+
+    a = c_api_util.UniquePtr(name="mock", deleter=self.deleter, obj=self.obj)
+
+    with a.get() as obj:
+      self.assertIs(obj, self.obj)
+      # The del below mimics a potential race condition.
+      # 'a' could be owned by a different thread, and this thread not
+      # necessarily hold a long-term reference to a.
+      del a
+      gc.collect()
+      self.assertFalse(obj.deleted)
+
+    gc.collect()
+    self.assertTrue(self.obj.deleted)
+
+  def testRaiseAfterDeleted(self):
+    self.assertFalse(self.obj.deleted)
+
+    a = c_api_util.UniquePtr(name="mock", deleter=self.deleter, obj=self.obj)
+
+    # The __del__ below mimics a partially started deletion, potentially
+    # started from another thread.
+    # 'a' could be owned by a different thread, and this thread not
+    # necessarily hold a long-term reference to a.
+    a.__del__()
+    self.assertTrue(self.obj.deleted)
+
+    with self.assertRaisesRegex(c_api_util.AlreadyGarbageCollectedError,
+                                "MockClass"):
+      with a.get():
+        pass
+
+    gc.collect()
+    self.assertTrue(self.obj.deleted)
 
 
-class ScopedTFStatus(object):
-  """Wrapper around TF_Status that handles deletion."""
+class ScopedTFBufferTest(test_util.TensorFlowTestCase):
 
-  __slots__ = ["status"]
+  def testDeleteRunsUnderNormalConditions(self):
+    # Mock TF_DeleteBuffer rather than letting __del__ actually free the
+    # buffer here: calling __del__() directly on a live ScopedTFBuffer frees
+    # the real C buffer immediately, and the object's own garbage collection
+    # later calls __del__() again for the same (already-freed) buffer,
+    # double-freeing it. Mocking out the delete call lets us call __del__()
+    # directly (CPython prints and swallows any exception raised inside a
+    # GC-triggered __del__ instead of propagating it, so a plain
+    # `del buf; gc.collect()` would let a broken destructor pass silently)
+    # while leaving the real deletion to happen exactly once, naturally,
+    # when the object is actually collected.
+    with mock.patch.object(c_api_util.c_api, "TF_DeleteBuffer") as mock_delete:
+      buf = c_api_util.ScopedTFBuffer(b"test")
+      buf.__del__()  # should not raise
+      mock_delete.assert_called_once()
 
-  def __init__(self):
-    self.status = c_api.TF_NewStatus()
+  def testDeleteGuardsAgainstNoneCApi(self):
+    buf = c_api_util.ScopedTFBuffer(b"test")
+    # Regression test for the __del__ fix: during interpreter shutdown,
+    # `c_api_util.c_api` can already have been torn down to None by the time
+    # this object is collected, so __del__ must guard against that rather
+    # than unconditionally calling into it.
+    with mock.patch.object(c_api_util, "c_api", None):
+      buf.__del__()  # should not raise
 
-  def __del__(self):
-    # Note: when we're destructing the global context (i.e when the process is
-    # terminating) we can have already deleted other modules.
-    if c_api is not None and c_api.TF_DeleteStatus is not None:
-      c_api.TF_DeleteStatus(self.status)
-
-
-class ScopedTFImportGraphDefOptions(object):
-  """Wrapper around TF_ImportGraphDefOptions that handles deletion."""
-
-  __slots__ = ["options"]
-
-  def __init__(self):
-    self.options = c_api.TF_NewImportGraphDefOptions()
-
-  def __del__(self):
-    # Note: when we're destructing the global context (i.e when the process is
-    # terminating) we can have already deleted other modules.
-    if c_api is not None and c_api.TF_DeleteImportGraphDefOptions is not None:
-      c_api.TF_DeleteImportGraphDefOptions(self.options)
-
-
-class ScopedTFImportGraphDefResults(object):
-  """Wrapper around TF_ImportGraphDefResults that handles deletion."""
-
-  __slots__ = ["results"]
-
-  def __init__(self, results):
-    self.results = results
-
-  def __del__(self):
-    # Note: when we're destructing the global context (i.e when the process is
-    # terminating) we can have already deleted other modules.
-    if c_api is not None and c_api.TF_DeleteImportGraphDefResults is not None:
-      c_api.TF_DeleteImportGraphDefResults(self.results)
-
-
-class ScopedTFFunction(UniquePtr):
-  """Wrapper around TF_Function that handles deletion."""
-
-  def __init__(self, func, name):
-    super(ScopedTFFunction, self).__init__(
-        name=name, obj=func, deleter=c_api.TF_DeleteFunction)
-
-
-class ScopedTFBuffer(object):
-  """An internal class to help manage the TF_Buffer lifetime."""
-
-  __slots__ = ["buffer"]
-
-  def __init__(self, buf_string):
-    self.buffer = c_api.TF_NewBufferFromString(compat.as_bytes(buf_string))
-
-  def __del__(self):
-    # Note: when we're destructing the global context (i.e when the process is
-    # terminating) we can have already deleted other modules.
-    # Also guard against __init__ having failed before `self.buffer` was
-    # ever assigned (e.g. compat.as_bytes raising on bad input): __del__
-    # still runs on a partially-constructed object, and accessing
-    # self.buffer directly would raise a second, masking AttributeError.
-    buffer = getattr(self, "buffer", None)
-    if (buffer is not None and c_api is not None and
-        c_api.TF_DeleteBuffer is not None):
-      c_api.TF_DeleteBuffer(buffer)
-
-
-class ApiDefMap(object):
-  """Wrapper around TF_ApiDefMap that handles querying and deletion.
-
-  The OpDef protos are also stored in this class so that they could
-  be queried by op name.
-  """
-
-  __slots__ = ["_api_def_map", "_op_per_name"]
-
-  def __init__(self):
-    op_def_proto = op_def_pb2.OpList()
-    buf = c_api.TF_GetAllOpList()
-    try:
-      op_def_proto.ParseFromString(c_api.TF_GetBuffer(buf))
-      self._api_def_map = c_api.TF_NewApiDefMap(buf)
-    finally:
-      c_api.TF_DeleteBuffer(buf)
-
-    self._op_per_name = {}
-    for op in op_def_proto.op:
-      self._op_per_name[op.name] = op
-
-  def __del__(self):
-    # Note: when we're destructing the global context (i.e when the process is
-    # terminating) we can have already deleted other modules.
-    if c_api is not None and c_api.TF_DeleteApiDefMap is not None:
-      c_api.TF_DeleteApiDefMap(self._api_def_map)
-
-  def put_api_def(self, text):
-    c_api.TF_ApiDefMapPut(self._api_def_map, text, len(text))
-
-  def get_api_def(self, op_name):
-    api_def_proto = api_def_pb2.ApiDef()
-    buf = c_api.TF_ApiDefMapGet(self._api_def_map, op_name, len(op_name))
-    if buf is None:
-      raise ValueError(f"No api_def found for op name {op_name}.")
-    try:
-      api_def_proto.ParseFromString(c_api.TF_GetBuffer(buf))
-    finally:
-      c_api.TF_DeleteBuffer(buf)
-    return api_def_proto
-
-  def get_op_def(self, op_name):
-    if op_name in self._op_per_name:
-      return self._op_per_name[op_name]
-    raise ValueError(f"No op_def found for op name {op_name}.")
-
-  def op_names(self):
-    return self._op_per_name.keys()
-
-
-@tf_contextlib.contextmanager
-def tf_buffer(data=None):
-  """Context manager that creates and deletes TF_Buffer.
-
-  Example usage:
-    with tf_buffer() as buf:
-      # get serialized graph def into buf
-      ...
-      proto_data = c_api.TF_GetBuffer(buf)
-      graph_def.ParseFromString(compat.as_bytes(proto_data))
-    # buf has been deleted
-
-    with tf_buffer(some_string) as buf:
-      c_api.TF_SomeFunction(buf)
-    # buf has been deleted
-
-  Args:
-    data: An optional `bytes`, `str`, or `unicode` object. If not None, the
-      yielded buffer will contain this data.
-
-  Yields:
-    Created TF_Buffer
-  """
-  if data:
-    buf = c_api.TF_NewBufferFromString(compat.as_bytes(data))
-  else:
-    buf = c_api.TF_NewBuffer()
-  try:
-    yield buf
-  finally:
-    c_api.TF_DeleteBuffer(buf)
-
-
-def tf_output(c_op, index):
-  """Returns a wrapped TF_Output with specified operation and index.
-
-  Args:
-    c_op: wrapped TF_Operation
-    index: integer
-
-  Returns:
-    Wrapped TF_Output
-  """
-  ret = c_api.TF_Output()
-  ret.oper = c_op
-  ret.index = index
-  return ret
+if __name__ == "__main__":
+  googletest.main()

From f13893046282b7a391d309390dea0886aa616980 Mon Sep 17 00:00:00 2001
From: Pablo Zimmermann <karupayun@google.com>
Date: Wed, 14 Oct 2026 07:19:23 -0700
Subject: [PATCH 0778/1000] [XLA:GPU] Pass k_loop_iterations directly to
 CalculateL2Time.

As part of migrating `gpu_dot_fusion_cost_model` into
`GpuIndexingPerformanceModel`, `detail::CalculateL2Time` needs to support
multi-dot fusions where K-loop iterations are accumulated across dots and no
single `(dot_k, tile_k)` pair exists.

Since `detail::CalculateL2Time` only used `dot_k` and `tile_k` to compute
`CeilOfRatio(dot_k, tile_k)`, accept `k_loop_iterations` directly to match
`detail::CalculatePipelinedLoopTime`.

PiperOrigin-RevId: 999350353
---
 .../gpu/model/gpu_dot_fusion_cost_model.cc     | 18 +++++++-----------
 .../gpu/model/gpu_dot_fusion_cost_model.h      |  6 +++---
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/model/gpu_dot_fusion_cost_model.cc b/third_party/xla/xla/service/gpu/model/gpu_dot_fusion_cost_model.cc
index 5f8770a27d7..17eab9b82bb 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_dot_fusion_cost_model.cc
+++ b/third_party/xla/xla/service/gpu/model/gpu_dot_fusion_cost_model.cc
@@ -303,12 +303,11 @@ absl::StatusOr<ComputeAndFlops> CalculateComputeTimeWithTileAndWaveQuantization(
   return result;
 }
 
-absl::StatusOr<absl::Duration> CalculateL2Time(
-    int64_t dot_k, int64_t tile_k, const se::DeviceDescription& device_info,
-    int64_t l2_bytes_read, bool is_tma_allowed) {
+absl::Duration CalculateL2Time(int64_t k_loop_iterations,
+                               const se::DeviceDescription& device_info,
+                               int64_t l2_bytes_read, bool is_tma_allowed) {
   double l2_cache_bandwidth_bytes_per_sec =
       GetMemoryBandwidthSpec(device_info).l2_cache_bandwidth_bytes_per_sec;
-  int64_t num_k_iters = CeilOfRatio<int64_t>(dot_k, tile_k);
 
   // Empirical overheads per K-dimension iteration.
   // The overhead is dictated by the memory instruction pathway rather than
@@ -323,7 +322,7 @@ absl::StatusOr<absl::Duration> CalculateL2Time(
 
   double base_time_seconds =
       1.0 * l2_bytes_read / l2_cache_bandwidth_bytes_per_sec;
-  return absl::Seconds(base_time_seconds + num_k_iters * k_loop_overhead);
+  return absl::Seconds(base_time_seconds + k_loop_iterations * k_loop_overhead);
 }
 
 // Returns the effective HBM bandwidth in bytes per second for a given dma_size.
@@ -766,12 +765,6 @@ absl::StatusOr<EstimateRunTimeData> EstimateRunTimeForDotOpWithBlockParameters(
   estimates.registers_per_thread = detail::CalculateRegistersPerThread(
       dot_info, dot_tile, block_params, device_info);
 
-  // Calculate L2 time.
-  ABSL_ASSIGN_OR_RETURN(absl::Duration l2_time,
-                   detail::CalculateL2Time(dot_info.k, dot_tile.k, device_info,
-                                           estimates.l2_bytes_read,
-                                           block_params.is_tma_allowed));
-
   TF_RET_CHECK(block_k_val > 0)
       << "block_k_val must be strictly positive, got " << block_k_val;
   TF_RET_CHECK(dot_info.k > 0)
@@ -792,6 +785,9 @@ absl::StatusOr<EstimateRunTimeData> EstimateRunTimeForDotOpWithBlockParameters(
       detail::CalculatePipelinedLoopTimeWithLaunchWaves(
           launch_config, compute_and_flops.compute_time, hbm_timing,
           device_info);
+  absl::Duration l2_time = detail::CalculateL2Time(
+      k_loop_iterations, device_info, estimates.l2_bytes_read,
+      block_params.is_tma_allowed);
 
   // Assuming perfect overlap between compute and memory for the rest,
   // but main loop is now modeled precisely.
diff --git a/third_party/xla/xla/service/gpu/model/gpu_dot_fusion_cost_model.h b/third_party/xla/xla/service/gpu/model/gpu_dot_fusion_cost_model.h
index 818aaad410f..cd630026d3d 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_dot_fusion_cost_model.h
+++ b/third_party/xla/xla/service/gpu/model/gpu_dot_fusion_cost_model.h
@@ -164,9 +164,9 @@ int CalculateRegistersPerThread(const DotProblemInfo& dot_info,
                                 const se::DeviceDescription& device_info);
 
 // Calculates the L2 time for a GPU DOT operation.
-absl::StatusOr<absl::Duration> CalculateL2Time(
-    int64_t dot_k, int64_t tile_k, const se::DeviceDescription& device_info,
-    int64_t l2_bytes, bool is_tma_allowed);
+absl::Duration CalculateL2Time(int64_t k_loop_iterations,
+                               const se::DeviceDescription& device_info,
+                               int64_t l2_bytes, bool is_tma_allowed);
 
 // Calculates the compute time for a GPU DOT operation with tile and wave
 // quantization effects taken into account.

From 9c4d5af0cbaef3287470241e3180a56ec5ee7327 Mon Sep 17 00:00:00 2001
From: Ilya Biryukov <ibiryukov@google.com>
Date: Wed, 14 Oct 2026 07:35:56 -0700
Subject: [PATCH 0779/1000] [XLA:GPU] Ignore `tileiras` older than CUDA 13.3 in
 the TensorIR emitter.

`//xla/backends/gpu/codegen/tensor_ir/tests:fusion_emitter_thunk_test` fails on
a CUDA 13.2 toolkit with

```
INTERNAL: TensorIrFusion: failed to compile fusion fusion: tileiras --list-versions failed with exit code 1
```

CUDA 13.3 is the first version that ships `tileiras` with the `--list-versions` probe TensorIR runs before assembling.

Have `SetUpTileIrAssembler` require `tileiras` 13.3 or newer, so an older one
is treated as absent and the bytecode fallback applies, as it already does for
toolkits that ship no assembler at all.

PiperOrigin-RevId: 999358912
---
 third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD | 2 ++
 .../gpu/codegen/tensor_ir/compilation_pipeline.cc        | 9 ++++++++-
 .../xla/xla/backends/gpu/codegen/tensor_ir/tests/BUILD   | 9 +++++----
 3 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD
index c19c9016151..e6840039565 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/BUILD
@@ -81,12 +81,14 @@ cc_library(
     data = if_google(if_cuda(["@local_config_cuda//cuda:runtime_tileiras"])),
     deps = [
         "//xla/stream_executor:device_description",
+        "//xla/stream_executor:semantic_version",
         "//xla/stream_executor/cuda:cuda_compute_capability",
         "//xla/stream_executor/cuda:subprocess_compilation",
         "@com_google_absl//absl/base",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/types:span",
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:Support",
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/compilation_pipeline.cc b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/compilation_pipeline.cc
index 842bc84801c..5a57f90da38 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/compilation_pipeline.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/compilation_pipeline.cc
@@ -27,11 +27,13 @@ limitations under the License.
 #include "absl/log/log.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
+#include "absl/types/span.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "mlir/Pass/PassManager.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/cuda/subprocess_compilation.h"
 #include "xla/stream_executor/device_description.h"
+#include "xla/stream_executor/semantic_version.h"
 #include "tsl/platform/path.h"
 
 namespace xla::gpu::tensor_ir {
@@ -64,8 +66,13 @@ void SetUpTileIrAssembler() {
     // The same search that finds `ptxas`: the CUDA toolkit roots first, `$PATH`
     // last. `tileiras` ships next to `ptxas` from CUDA 13.1 on; TensorIR itself
     // needs the 13.3 release or newer.
+    static constexpr stream_executor::SemanticVersion kMinimumTileirasVersion{
+        13, 3, 0};
+    static constexpr absl::Span<const stream_executor::SemanticVersion>
+        kNoExcludedVersions{};
     absl::StatusOr<std::string> path = stream_executor::FindCudaExecutable(
-        "tileiras", /*preferred_cuda_dir=*/"");
+        "tileiras", /*preferred_cuda_dir=*/"", kMinimumTileirasVersion,
+        kNoExcludedVersions);
     if (!path.ok()) {
       LOG(WARNING) << "TensorIR: no Tile IR assembler (CUDA 13.3 or newer "
                       "ships one), falling back to Tile IR bytecode for the "
diff --git a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/BUILD b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/BUILD
index 2964d0d1666..5857eea700d 100644
--- a/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/tensor_ir/tests/BUILD
@@ -73,10 +73,11 @@ xla_test(
 # `--xla_gpu_cuda_data_dir`, or on `PATH`). CUDA 13.0 ships no `tileiras`, and
 # the 13.1 and 13.2 releases lack the `--list-versions` probe TensorIR uses to
 # check bytecode compatibility, so 13.3 is the first release the compiler can
-# assemble with. Without an assembler the compiler hands back Tile IR bytecode
-# for the driver to JIT. TensorIR emits the CUDA 13.3 bytecode format by
-# default (`max(compatibility, 13.3)`, see the tensor_ir README), which needs
-# the matching r610 or newer driver; the r590 (CUDA 13.1) JIT only takes the
+# assemble with; `SetUpTileIrAssembler` ignores older ones. Without an
+# assembler the compiler hands back Tile IR bytecode for the driver to JIT.
+# TensorIR emits the CUDA 13.3 bytecode format by default
+# (`max(compatibility, 13.3)`, see the tensor_ir README), which needs the
+# matching r610 or newer driver; the r590 (CUDA 13.1) JIT only takes the
 # `compatibility` format, which the emitter does not request.
 #
 # Why Hopper. The compiler refuses to assemble a cubin below sm_90

From 64a66d9b28f97d9bc267758fbde8b6cf21c3db22 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Wed, 14 Oct 2026 07:40:57 -0700
Subject: [PATCH 0780/1000] [XLA:GPU] Propagate GPU errors to local collectives
 and coordination service.

PiperOrigin-RevId: 999361352
---
 .../xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc    | 63 +++++++++++++
 .../gpu/se_gpu_pjrt_client_multi_gpu_test.cc  | 50 ++++++++++
 third_party/xla/xla/service/gpu/BUILD         |  6 ++
 .../xla/xla/service/gpu/gpu_executable.cc     | 21 ++++-
 .../service/gpu/gpu_executable_run_options.cc | 17 ++++
 .../service/gpu/gpu_executable_run_options.h  | 28 ++++++
 .../xla/service/gpu/gpu_executable_test.cc    | 93 +++++++++++++++++++
 7 files changed, 275 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
index ec881f12410..f2f65fe0d3e 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/pjrt/gpu/se_gpu_pjrt_client.h"
 
+#include <atomic>
 #include <cstddef>
 #include <cstdint>
 #include <cstdlib>
@@ -1539,6 +1540,50 @@ void NameDeviceAndLauncherThread(const LocalTopologyProto& node,
   });
 }
 
+constexpr absl::string_view kFatalGpuErrorKey = "gpu_fatal_error";
+
+void RegisterDistributedDeviceErrorPropagation(
+    int node_id, const std::shared_ptr<KeyValueStoreInterface>& kv_store,
+    gpu::GpuExecutableRunOptions* gpu_run_options) {
+  auto already_fired = std::make_shared<std::atomic<bool>>(false);
+  gpu_run_options->set_device_error_callback(
+      [kv_store, already_fired, node_id](int device_ordinal,
+                                         const absl::Status& status) {
+        if (already_fired->exchange(true)) {
+          return;
+        }
+        LOG(ERROR) << absl::StreamFormat(
+            "Aborting local GPU collectives after fatal device error on "
+            "node %d device %d: %v",
+            node_id, device_ordinal, status);
+        if (absl::Status s = gpu::AbortAllCliques(); !s.ok()) {
+          LOG(WARNING) << "Failed to abort GPU cliques: " << s;
+        }
+        absl::Status set_status = kv_store->Set(
+            kFatalGpuErrorKey,
+            absl::StrFormat("Fatal GPU device error on node %d device %d: %v",
+                            node_id, device_ordinal, status));
+        if (!set_status.ok()) {
+          LOG(WARNING) << "Failed to broadcast GPU error via kv_store: "
+                       << set_status;
+        }
+      });
+
+  kv_store->AsyncGet(
+      kFatalGpuErrorKey,
+      [already_fired, node_id](const absl::StatusOr<std::string>& status) {
+        if (!status.ok() || already_fired->exchange(true)) {
+          return;
+        }
+        LOG(ERROR) << "Node " << node_id
+                   << " aborting local GPU collectives due to remote error: "
+                   << *status;
+        if (absl::Status s = gpu::AbortAllCliques(); !s.ok()) {
+          LOG(WARNING) << "Failed to abort GPU cliques: " << s;
+        }
+      });
+}
+
 }  // namespace
 
 // Creates allocator memory registration and adds the required suballocator
@@ -2016,6 +2061,19 @@ absl::StatusOr<std::unique_ptr<PjRtClient>> GetStreamExecutorGpuClient(
   if (kv_store == nullptr) {
     kv_store = std::make_shared<InMemoryKeyValueStore>();
   }
+
+  // Clear any stale fatal GPU error key left over from a prior incarnation
+  // before entering the topology exchange barrier. Because ExchangeTopologies
+  // inside BuildDistributedDevices is a full barrier across all nodes, no peer
+  // can finish client initialization and Set a fresh error until every node has
+  // completed this Delete. Note: AsyncGet in
+  // RegisterDistributedDeviceErrorPropagation is a single-shot subscription, so
+  // peers that do not restart alongside a restarting node will not re-arm their
+  // subscription after the first fatal error.
+  if (options.num_nodes > 1 && kv_store != nullptr) {
+    kv_store->Delete(kFatalGpuErrorKey).IgnoreError();
+  }
+
   ABSL_ASSIGN_OR_RETURN(
       PjRtDevicesAndTopology devices_and_topology,
       BuildDistributedDevices(
@@ -2024,6 +2082,11 @@ absl::StatusOr<std::unique_ptr<PjRtClient>> GetStreamExecutorGpuClient(
           options.enable_mock_nccl, options.mock_gpu_topology,
           options.partition_index, options.verify_topology_fingerprint));
 
+  if (options.num_nodes > 1 && kv_store != nullptr) {
+    RegisterDistributedDeviceErrorPropagation(options.node_id, kv_store,
+                                              gpu_run_options.get());
+  }
+
   se::StreamExecutor* first_executor =
       GetFirstExecutor(devices_and_topology.local_device_states);
 
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc
index ea17185ec79..7bbe9c14cdd 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_multi_gpu_test.cc
@@ -1949,6 +1949,56 @@ INSTANTIATE_TEST_SUITE_P(
         {2, 0}, {2, 1}, {2, 2}}),
     ShardedAutotuningTestInfo::Name);
 
+TEST(StreamExecutorGpuClientTest, DeviceErrorCallbackPropagatesToPeerNodes) {
+  constexpr int kNumNodes = 2;
+  auto kv_store = std::make_shared<InMemoryKeyValueStore>();
+  ASSERT_OK(kv_store->Set("gpu_fatal_error", "stale_error_from_previous_run"));
+
+  absl::StatusOr<std::unique_ptr<PjRtClient>> client_0_or;
+  absl::StatusOr<std::unique_ptr<PjRtClient>> client_1_or;
+  {
+    tsl::thread::ThreadPool pool(tsl::Env::Default(), "create_clients",
+                                 kNumNodes);
+    pool.Schedule([&]() {
+      GpuClientOptions options = GetTestGpuClientOptions(2);
+      options.node_id = 0;
+      options.num_nodes = kNumNodes;
+      options.allowed_devices = {0};
+      options.kv_store = kv_store;
+      client_0_or = GetStreamExecutorGpuClient(options);
+    });
+    pool.Schedule([&]() {
+      GpuClientOptions options = GetTestGpuClientOptions(2);
+      options.node_id = 1;
+      options.num_nodes = kNumNodes;
+      options.allowed_devices = {1};
+      options.kv_store = kv_store;
+      client_1_or = GetStreamExecutorGpuClient(options);
+    });
+  }
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtClient> client_0,
+                       std::move(client_0_or));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtClient> client_1,
+                       std::move(client_1_or));
+  EXPECT_THAT(kv_store->TryGet("gpu_fatal_error"),
+              StatusIs(absl::StatusCode::kNotFound));
+
+  auto* se_client_0 = absl::down_cast<CommonPjRtClientImpl*>(client_0.get());
+  auto* raw_client_0 =
+      absl::down_cast<PjRtStreamExecutorRawClient*>(se_client_0->raw_client());
+  ASSERT_NE(raw_client_0->gpu_run_options(), nullptr);
+  ASSERT_NE(raw_client_0->gpu_run_options()->device_error_callback(), nullptr);
+  ASSERT_TRUE(*raw_client_0->gpu_run_options()->device_error_callback());
+
+  (*raw_client_0->gpu_run_options()->device_error_callback())(
+      /*device_ordinal=*/0, absl::InternalError("CUDA_ERROR_ILLEGAL_ADDRESS"));
+
+  ASSERT_OK_AND_ASSIGN(std::string kv_error,
+                       kv_store->Get("gpu_fatal_error", absl::Seconds(5)));
+  EXPECT_THAT(kv_error, HasSubstr("Fatal GPU device error on node 0 device 0"));
+  EXPECT_THAT(kv_error, HasSubstr("CUDA_ERROR_ILLEGAL_ADDRESS"));
+}
+
 }  // namespace
 }  // namespace xla
 
diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 4784424d0c9..d74471d8f1a 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -94,6 +94,7 @@ cc_library(
         "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/functional:any_invocable",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/time",
@@ -952,8 +953,10 @@ xla_test(
         ":gpu_executable",
         ":gpu_executable_buffer_allocator",
         ":gpu_executable_proto_cc",
+        ":gpu_executable_run_options",
         ":launch_dimensions",
         "//xla:debug_options_flags",
+        "//xla:executable_run_options",
         "//xla:literal_util",
         "//xla:shape_layout",
         "//xla:shape_util",
@@ -987,10 +990,12 @@ xla_test(
         "//xla/service:shaped_slice",
         "//xla/service:xla_debug_info_manager",
         "//xla/stream_executor:device_address",
+        "//xla/stream_executor:device_address_allocator",
         "//xla/stream_executor:device_description",
         "//xla/stream_executor:mock_stream",
         "//xla/stream_executor:mock_stream_executor",
         "//xla/stream_executor:module_spec",
+        "//xla/stream_executor:platform_manager",
         "//xla/stream_executor:semantic_version",
         "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor/abi:executable_abi_version",
@@ -1003,6 +1008,7 @@ xla_test(
         "//xla/tsl/util/proto:parse_text_proto",
         "//xla/tsl/util/proto:proto_matchers",
         "//xla/util/split_proto:split_proto_reader",
+        "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:status_matchers",
diff --git a/third_party/xla/xla/service/gpu/gpu_executable.cc b/third_party/xla/xla/service/gpu/gpu_executable.cc
index a5f2eaf1efa..fc7682893ca 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable.cc
@@ -925,9 +925,24 @@ absl::Status GpuExecutable::ExecuteThunksImpl(
         HangWatchdog::Abort(device_watchdog_name, device_timeout)));
   }
 
-  if (!device_guards.empty()) {
-    ABSL_RETURN_IF_ERROR(
-        main_stream->DoHostCallback([guards = std::move(device_guards)] {}));
+  std::shared_ptr<const GpuExecutableRunOptions::DeviceErrorCallback> error_cb =
+      gpu_run_options ? gpu_run_options->device_error_callback() : nullptr;
+
+  if (!device_guards.empty() || error_cb != nullptr) {
+    ABSL_RETURN_IF_ERROR(main_stream->DoHostCallback(
+        [guards = std::move(device_guards)] {},
+        [device_ordinal = executor->device_ordinal(), module_name, run_id,
+         error_cb = std::move(error_cb)](absl::Status status) {
+          if (absl::IsCancelled(status)) {
+            return;
+          }
+          LOG(ERROR) << absl::StreamFormat(
+              "[%d] Async device error in `%s` (run_id=%v): %v", device_ordinal,
+              module_name, run_id, status);
+          if (error_cb != nullptr && *error_cb) {
+            (*error_cb)(device_ordinal, status);
+          }
+        }));
   }
 
   return MaybeSyncAndProfile(run_options, execution_timer.get(),
diff --git a/third_party/xla/xla/service/gpu/gpu_executable_run_options.cc b/third_party/xla/xla/service/gpu/gpu_executable_run_options.cc
index 66d5f4af591..3a7fd59bb36 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable_run_options.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable_run_options.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "xla/service/gpu/gpu_executable_run_options.h"
 
 #include <functional>
+#include <memory>
 #include <optional>
 #include <utility>
 #include <vector>
@@ -84,4 +85,20 @@ GpuExecutableRunOptions::execution_timeout_handlers() const {
   return execution_timeout_handlers_();
 }
 
+GpuExecutableRunOptions& GpuExecutableRunOptions::set_device_error_callback(
+    DeviceErrorCallback callback) {
+  if (callback) {
+    device_error_callback_ =
+        std::make_shared<const DeviceErrorCallback>(std::move(callback));
+  } else {
+    device_error_callback_.reset();
+  }
+  return *this;
+}
+
+std::shared_ptr<const GpuExecutableRunOptions::DeviceErrorCallback>
+GpuExecutableRunOptions::device_error_callback() const {
+  return device_error_callback_;
+}
+
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/service/gpu/gpu_executable_run_options.h b/third_party/xla/xla/service/gpu/gpu_executable_run_options.h
index 5157cf3a8ba..febdbf381f4 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable_run_options.h
+++ b/third_party/xla/xla/service/gpu/gpu_executable_run_options.h
@@ -17,12 +17,14 @@ limitations under the License.
 #define XLA_SERVICE_GPU_GPU_EXECUTABLE_RUN_OPTIONS_H_
 
 #include <functional>
+#include <memory>
 #include <optional>
 #include <vector>
 
 #include "absl/container/btree_map.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/functional/any_invocable.h"
+#include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/time/time.h"
@@ -124,6 +126,31 @@ class GpuExecutableRunOptions {
   // timeout handlers are not set.
   std::vector<ExecutionTimeoutHandler> execution_timeout_handlers() const;
 
+  // Callback invoked when an asynchronous device stream error (e.g.
+  // `CUDA_ERROR_ILLEGAL_ADDRESS` from a faulting kernel on `main_stream`) is
+  // detected at the end of an `ExecuteThunks` launch.
+  //
+  // Because GPU kernels execute asynchronously after host launch returns,
+  // device-side faults are detected asynchronously when XLA's background
+  // stream status monitor polls the poisoned stream (or synchronously on the
+  // host launch thread if host-callback registration fails).
+  //
+  // Typically used to abort local collective cliques (`AbortAllCliques`) and
+  // broadcast the fatal error via `kv_store` so peer devices/nodes abort their
+  // in-flight collectives immediately instead of hanging until timeout.
+  //
+  // Threading & lifetime: Invoked via `Stream::DoHostCallback`'s error path on
+  // XLA's background stream status monitor `tsl::Thread` (or synchronously on
+  // the host launch thread), not on a GPU driver callback thread. May be
+  // invoked after the run options are destroyed; capture only owning pointers
+  // or values.
+  using DeviceErrorCallback = absl::AnyInvocable<void(
+      int device_ordinal, const absl::Status& status) const>;
+
+  GpuExecutableRunOptions& set_device_error_callback(
+      DeviceErrorCallback callback);
+  std::shared_ptr<const DeviceErrorCallback> device_error_callback() const;
+
  private:
   bool requires_exclusive_lock_on_gpu_ = false;
   bool enable_mock_collectives_ = false;
@@ -134,6 +161,7 @@ class GpuExecutableRunOptions {
       incarnations_;
   std::function<std::vector<ExecutionTimeoutHandler>()>
       execution_timeout_handlers_;
+  std::shared_ptr<const DeviceErrorCallback> device_error_callback_;
 };
 
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/service/gpu/gpu_executable_test.cc b/third_party/xla/xla/service/gpu/gpu_executable_test.cc
index fdbc3615056..6b09b7e75cd 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable_test.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable_test.cc
@@ -25,6 +25,7 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/functional/any_invocable.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/status_matchers.h"
@@ -51,6 +52,7 @@ limitations under the License.
 #include "xla/client/executable_build_options.h"
 #include "xla/codegen/emitters/kernel_arguments.h"
 #include "xla/debug_options_flags.h"
+#include "xla/executable_run_options.h"
 #include "xla/hlo/analysis/alias_info.h"
 #include "xla/hlo/analysis/hlo_ordering.h"
 #include "xla/hlo/ir/hlo_input_output_alias_config.h"
@@ -65,6 +67,7 @@ limitations under the License.
 #include "xla/service/gpu/dense_data_intermediate.h"
 #include "xla/service/gpu/gpu_executable.pb.h"
 #include "xla/service/gpu/gpu_executable_buffer_allocator.h"
+#include "xla/service/gpu/gpu_executable_run_options.h"
 #include "xla/service/gpu/gpu_module_globals.h"
 #include "xla/service/gpu/launch_dimensions.h"
 #include "xla/service/gpu_topology.h"
@@ -80,11 +83,13 @@ limitations under the License.
 #include "xla/stream_executor/abi/executable_abi_version.pb.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/device_address.h"
+#include "xla/stream_executor/device_address_allocator.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/gpu/tma_metadata.h"
 #include "xla/stream_executor/mock_stream.h"
 #include "xla/stream_executor/mock_stream_executor.h"
 #include "xla/stream_executor/module_spec.h"
+#include "xla/stream_executor/platform_manager.h"
 #include "xla/stream_executor/semantic_version.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/lib/core/status_test_util.h"
@@ -1340,6 +1345,94 @@ TEST_F(GpuExecutableTest, ExecutableAbiVersion) {
               EqualsProto(executable_abi_version_proto));
 }
 
+class MockStreamWithErrorCallback : public se::MockStream {
+ public:
+  using se::MockStream::DoHostCallbackWithStatus;
+  MOCK_METHOD(absl::Status, DoHostCallbackWithStatus,
+              (absl::AnyInvocable<absl::Status() &&> callback,
+               absl::AnyInvocable<void(absl::Status) &&> error_cb),
+              (override));
+};
+
+class FakeDeviceAddressAllocator : public se::DeviceAddressAllocator {
+ public:
+  explicit FakeDeviceAddressAllocator(const se::Platform* platform)
+      : se::DeviceAddressAllocator(platform) {}
+  absl::StatusOr<se::ScopedDeviceAddress<uint8_t>> Allocate(
+      int device_ordinal, uint64_t size, bool retry_on_failure,
+      int64_t memory_space) override {
+    return se::ScopedDeviceAddress<uint8_t>();
+  }
+  absl::Status Deallocate(int device_ordinal,
+                          se::DeviceAddressBase mem) override {
+    return absl::OkStatus();
+  }
+  bool AllowsAsynchronousDeallocation() const override { return true; }
+  absl::StatusOr<se::Stream*> GetStream(int device_ordinal) override {
+    return nullptr;
+  }
+};
+
+TEST_F(GpuExecutableTest, InvokesDeviceErrorCallbackOnStreamError) {
+  ASSERT_OK_AND_ASSIGN(
+      se::Platform * platform,
+      se::PlatformManager::PlatformWithId(stream_executor_platform_id()));
+  se::MockStreamExecutor executor;
+  MockStreamWithErrorCallback stream;
+
+  EXPECT_CALL(executor, GetPlatform()).WillRepeatedly(Return(platform));
+  EXPECT_CALL(executor, device_ordinal()).WillRepeatedly(Return(0));
+  EXPECT_CALL(stream, parent()).WillRepeatedly(Return(&executor));
+  EXPECT_CALL(stream, GetCudaComputeCapability())
+      .WillRepeatedly(Return(se::CudaComputeCapability::Hopper()));
+
+  se::DeviceDescription device_description;
+  device_description.set_gpu_compute_capability(
+      se::GpuComputeCapability{se::CudaComputeCapability::Hopper()});
+
+  GpuExecutable::Params params;
+  params.module_name = "test_module";
+  params.executable = std::make_unique<ThunkExecutor>(ThunkSequence{});
+  params.device_description = device_description;
+  params.enable_debug_info_manager = false;
+  SetDummyBufferAssignment(params);
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<GpuExecutable> executable,
+                       GpuExecutable::Create(std::move(params)));
+
+  int reported_device = -1;
+  absl::Status reported_status = absl::OkStatus();
+  GpuExecutableRunOptions gpu_opts;
+  gpu_opts.set_device_error_callback(
+      [&](int device_ordinal, const absl::Status& status) {
+        reported_device = device_ordinal;
+        reported_status = status;
+      });
+
+  FakeDeviceAddressAllocator allocator(platform);
+  ExecutableRunOptions exec_run_options;
+  exec_run_options.set_stream(&stream);
+  exec_run_options.set_allocator(&allocator);
+  exec_run_options.set_gpu_executable_run_options(&gpu_opts);
+  ServiceExecutableRunOptions service_run_options(exec_run_options);
+
+  EXPECT_CALL(stream, DoHostCallbackWithStatus(::testing::_, ::testing::_))
+      .WillOnce([](absl::AnyInvocable<absl::Status() &&> /*ok_cb*/,
+                   absl::AnyInvocable<void(absl::Status) &&> error_cb) {
+        std::move(error_cb)(absl::InternalError("CUDA_ERROR_ILLEGAL_ADDRESS"));
+        return absl::OkStatus();
+      });
+
+  BufferAllocations buffer_allocations(/*buffers=*/{}, /*device_ordinal=*/0,
+                                       &allocator);
+  ASSERT_OK(
+      executable->ExecuteThunks(buffer_allocations, &service_run_options));
+  EXPECT_EQ(reported_device, 0);
+  EXPECT_THAT(reported_status,
+              absl_testing::StatusIs(
+                  absl::StatusCode::kInternal,
+                  ::testing::HasSubstr("CUDA_ERROR_ILLEGAL_ADDRESS")));
+}
+
 int64_t BufferSizeBytes(const BufferValue& buffer) {
   return ShapeUtil::ByteSizeOf(buffer.shape(), sizeof(void*));
 }

From 13beb4dac82b078ecf62563a43a0b81dea114cf7 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 20:26:56 +0530
Subject: [PATCH 0781/1000] Use _safe_print for all remaining console output in
 diagnose_dll_load

The except OSError block in diagnose_dll_load's per-dependency loop
still used plain print() for the 'Failed to load <dll_name>' message
and every error-code-specific message that follows it, including the
one printing dll_name itself, which is parsed from the target binary's
PE import table and is not guaranteed to be ASCII. On a console using
a legacy code page, this could raise UnicodeEncodeError and replace
the original ImportError being diagnosed, defeating the whole point of
_safe_print.

Verified: with a dll_name and exception message containing non-ASCII
characters and an ascii-only console encoding, the old plain print()
would raise; with _safe_print throughout, it prints the
backslash-escaped text instead and raises nothing.
---
 .../python/platform/windows_lib_diagnostics.py | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/tensorflow/python/platform/windows_lib_diagnostics.py b/tensorflow/python/platform/windows_lib_diagnostics.py
index f2b800668be..43fdf9f4cd8 100644
--- a/tensorflow/python/platform/windows_lib_diagnostics.py
+++ b/tensorflow/python/platform/windows_lib_diagnostics.py
@@ -172,29 +172,31 @@ def diagnose_dll_load(path, depth=0, visited=None):
       _ = ctypes.WinDLL(dll_name)
     except OSError as e:
       err = getattr(e, "winerror", None)
-      print(f"{indent}[Error] Failed to load {dll_name}: ", end="")
+      _safe_print(f"{indent}[Error] Failed to load {dll_name}: ", end="")
       if err == ERROR_MOD_NOT_FOUND:
-        print(
+        _safe_print(
             "MODULE NOT FOUND (0x7E) - Check if the file or its dependencies"
             " exist."
         )
       elif err == ERROR_BAD_EXE_FORMAT:
-        print("BAD EXE FORMAT (0xC1/193) - Likely a 32-bit vs 64-bit mismatch.")
+        _safe_print(
+            "BAD EXE FORMAT (0xC1/193) - Likely a 32-bit vs 64-bit mismatch."
+        )
       elif err == ERROR_PROC_NOT_FOUND:
-        print(
+        _safe_print(
             "PROC NOT FOUND (0x7F/127) - Version mismatch in a dependent DLL."
         )
       elif err == ERROR_ACCESS_DENIED:
-        print("ACCESS DENIED (0x05) - Check file permissions.")
+        _safe_print("ACCESS DENIED (0x05) - Check file permissions.")
       elif err == ERROR_DLL_INIT_FAILED:
-        print(
+        _safe_print(
             "INITIALIZATION FAILED (0x45A) - The DLL's DllMain returned false."
         )
-        print(
+        _safe_print(
             f"{indent}    Hint: This often happens if your CPU lacks required"
             " instructions (like AVX/AVX2)"
         )
-        print(
+        _safe_print(
             f"{indent}    or if the Microsoft Visual C++ Redistributable is"
             " outdated/missing."
         )

From 97612016b473973638dade7fb1fab2af5d5b9111 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 20:31:17 +0530
Subject: [PATCH 0782/1000] Parameterize testDatasetAlreadyBatchedOffsets for
 Eager and Graph modes

Follows the same pattern as testRangeIteration/testForInIteration: wrap
the iterator advance in a polymorphic_function.function when is_graph is
True, and barrier after it so any pending async error surfaces within
the test.

Verified by reloading this file's input_util.py and
input_util_test.py into a real installed TensorFlow build and running
both the Eager and Graph parameterizations directly (both pass), plus
the pre-existing validation-regression tests from earlier review
rounds on this PR (unknown-shape rejection, partially-defined-shape
rejection, incorrect-batch-size rejection, missing-drop_remainder
rejection) to confirm nothing regressed.
---
 tensorflow/dtensor/python/tests/input_util_test.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/tensorflow/dtensor/python/tests/input_util_test.py b/tensorflow/dtensor/python/tests/input_util_test.py
index e26f0e9d6f2..632697134b1 100644
--- a/tensorflow/dtensor/python/tests/input_util_test.py
+++ b/tensorflow/dtensor/python/tests/input_util_test.py
@@ -122,7 +122,8 @@ class DTensorDatasetTest(test_util.DTensorBaseTest):
           layouts=images_layout,
           batch_dim=MESH_DIM_BATCH)
 
-  def testDatasetAlreadyBatchedOffsets(self):
+  @parameterized.named_parameters(('Eager', False), ('Graph', True))
+  def testDatasetAlreadyBatchedOffsets(self, is_graph):
     layout = Layout(
         [MESH_DIM_BATCH, MESH_DIM_HEIGHT, MESH_DIM_WIDTH], self.mesh)
 
@@ -142,7 +143,15 @@ class DTensorDatasetTest(test_util.DTensorBaseTest):
         dataset_already_batched=True)
 
     iterator = iter(d_dataset)
-    d_tensor = next(iterator)
+
+    def get_next(it):
+      return next(it)
+
+    get_next_fn = (
+        polymorphic_function.function(get_next) if is_graph else get_next)
+    d_tensor = get_next_fn(iterator)
+    if is_graph:
+      mesh_util.barrier(self.mesh)
 
     self.assertDTensorEqual(local_images, layout, d_tensor)
 

From 7a2a442c28958940f412ae7bde073c640916e969 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 20:38:50 +0530
Subject: [PATCH 0783/1000] Fix deferred restoration of scalar DVariables

api.call_with_layout was not forwarding shape to the
CheckpointInitialValueCallable, so shape defaulted to None inside
CheckpointInitialValue and crashed with 'NoneType is not iterable'
when shard_info was truthy. Passing shape=shape without adjusting
shard_info also breaks: an empty-shape ShardInfo produces a single
space as the shape_and_slice spec, which RestoreV2 rejects since
scalar tensors can't be sliced/sharded.

Fix: pass shape=shape through, and use shard_info=None for scalar
(falsy shape) tensors, preserving the existing shard_info for
non-scalar shapes.

Verified by loading the patched module into a real installed
TensorFlow build and running test_checkpoint_deferred_restore_scalar_dvariable
plus two non-scalar test_checkpoint_simple cases: all pass. Confirmed
the prior code fails with exactly the reported TypeError.
---
 tensorflow/dtensor/python/d_variable.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/tensorflow/dtensor/python/d_variable.py b/tensorflow/dtensor/python/d_variable.py
index c501c597fdf..3fea2f8fe4d 100644
--- a/tensorflow/dtensor/python/d_variable.py
+++ b/tensorflow/dtensor/python/d_variable.py
@@ -182,12 +182,21 @@ class DVariable(resource_variable_ops.ResourceVariable):
         # layout of the tf.RestoreV2 that will be called
         #
         # Passing shard_info with the correct shape allows the tf.RestoreV2
-        # ShapeInference to extract the shape.
+        # ShapeInference to extract the shape. Scalar tensors (shape=()) can't
+        # be sliced/sharded, so they must use shard_info=None rather than a
+        # ShardInfo with an empty shape and offset (which would otherwise
+        # produce an invalid " " shape_and_slice spec for RestoreV2).
+        shard_info = (
+            trackable.ShardInfo(shape=shape, offset=[0] * len(shape))
+            if shape
+            else None
+        )
         initial_value = api.call_with_layout(
             initial_value,
             layout,
-            shard_info=trackable.ShardInfo(
-                shape=shape, offset=[0] * len(shape)))
+            shape=shape,
+            shard_info=shard_info,
+        )
       else:
         initial_value = initial_value()
 

From 021a44d8c16f3558399525fbfdafdb465f675fdd Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2026 08:06:53 -0700
Subject: [PATCH 0784/1000] Reverts 596eca059dd448a67178f25e31e7ab34222b1aaa

PiperOrigin-RevId: 999375638
---
 third_party/xla/xla/debug_options_flags.cc | 10 ++++++----
 third_party/xla/xla/xla.proto              |  4 +---
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index cab2e9bdae0..9f7a7708997 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -3901,9 +3901,12 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
       debug_options->xla_gpu_rocm_max_trace_events(),
       "Maximum number of ROCm trace events (applies to callback/activity/"
       "annotation). Set as high as memory allows; up to 1e9."));
-  flag_list->push_back(
-      tsl::Flag("xla_gpu_experimental_enable_tiling_propagation",
-                noop_flag_setter<bool>, true, "[Deprecated, do not use]."));
+  flag_list->push_back(tsl::Flag(
+      "xla_gpu_experimental_enable_tiling_propagation",
+      bool_setter_for(
+          &DebugOptions::set_xla_gpu_experimental_enable_tiling_propagation),
+      debug_options->xla_gpu_experimental_enable_tiling_propagation(),
+      "If true, enable experimental tiling propagation."));
 
   auto setter_for_xla_gpu_detect_nan =
       [debug_options, detection_mode](const std::string& value) {
@@ -4114,7 +4117,6 @@ DebugOptions GetDebugOptionsFromProtoAndFlags(
     DebugOptions* proto_debug_options) {
   if (proto_debug_options != nullptr) {
     DebugOptions debug_options = *proto_debug_options;
-    debug_options.set_xla_gpu_experimental_enable_tiling_propagation(true);
     // Apply overrides from XLA_FLAGS environment variable.
     std::vector<tsl::Flag> flag_list;
     MakeDebugOptionsFlags(&flag_list, &debug_options);
diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index 1d0936e2bb6..b0e67029819 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -1883,9 +1883,7 @@ message DebugOptions {
   optional int64 xla_gpu_async_copy_min_bytes = 522;
 
   // If true, enable experimental tiling propagation.
-  // TODO(b/514293537): drop the flag.
-  optional bool xla_gpu_experimental_enable_tiling_propagation = 456
-      [deprecated = true];
+  optional bool xla_gpu_experimental_enable_tiling_propagation = 456;
 
   // Enables and controls expetimental cost model-driven tiling choices for GEMM
   // fusions.

From c876640ff6542094a383494540462b513fb0676b Mon Sep 17 00:00:00 2001
From: Ce Zheng CA <zce@google.com>
Date: Wed, 14 Oct 2026 08:14:31 -0700
Subject: [PATCH 0785/1000] [XLA] Extract NearComparatorBase in
 xla/literal_comparison.cc to reduce template bloat

Previously, `NearComparator<NativeT>` in `xla/literal_comparison.cc` duplicated all of its member variables and type-independent reporting/bucket helpers (`Prepare`, `Finalize`, `ErrorMessage`, `UpdateErrorBucket`, `UpdateAbsValueBucket`, `InsertCandidate`, `UpdateCandidates`, and `CalculateFloatDistance`, which itself expands a 15-way `FloatingPointTypeSwitch` over `low_precision_fp_error_spec.type`) across all 17 floating-point and complex primitive types (`float4_e2m1fn` through `double`, `complex64`, and `complex128`). Extract a non-template, non-polymorphic `NearComparatorBase` class for the type-independent state and helpers while keeping `NearComparator<NativeT>` templated on `NativeT` for element comparison and `Mismatch` storage.

Runtime and memory impact: Zero runtime overhead or memory allocation change; element-wise comparison (`CompareLiterals`, `CompareLiteralsSlow`, `CompareValues`) and `Mismatch` remain typed on `NativeT` with no virtual dispatch and identical inlining and code generation, while deduplicating 16 redundant copies of the outer setup, error formatting, and float-distance dispatch reduces compiled binary size by ~128.7 KB.

PiperOrigin-RevId: 999379991
---
 third_party/xla/xla/literal_comparison.cc | 441 ++++++++++++----------
 1 file changed, 234 insertions(+), 207 deletions(-)

diff --git a/third_party/xla/xla/literal_comparison.cc b/third_party/xla/xla/literal_comparison.cc
index 9d53304f0a6..bfca975215b 100644
--- a/third_party/xla/xla/literal_comparison.cc
+++ b/third_party/xla/xla/literal_comparison.cc
@@ -38,6 +38,7 @@ limitations under the License.
 #include "absl/status/status_macros.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "Eigen/Core"
 #include "xla/error_spec.h"
@@ -256,62 +257,15 @@ inline double RoundUpTo1SigFig(double val) {
   return std::ceil(val / factor - 1e-9) * factor;
 }
 
-// Helper class for comparing floating-point literals within an error bound.
-template <typename NativeT>
-class NearComparator {
- public:
-  // Compares the two array literals elementwise and returns a comparison
-  // result. The comparison is ok() if all actual and expected elements are
-  // within the given error bound. In case of error, the status contains a
-  // detailed message about the discrepancy.
-  static absl::Status Compare(const LiteralSlice& expected,
-                              const LiteralSlice& actual,
-                              const ShapeIndex& shape_index, ErrorSpec error,
-                              bool detailed_message,
-                              const MiscompareCallback& miscompare_callback) {
-    NearComparator<NativeT> comparator(expected, actual, shape_index, error,
-                                       detailed_message, miscompare_callback);
-    return comparator.Run();
-  }
-
- private:
-  // Data structure encapsulating metadata about a single element mismatch.
-  struct Mismatch {
-    NativeT actual;
-    NativeT expected;
-    double rel_error;
-    double abs_error;
-
-    // The linear index of the failure within the shape. This linear index is
-    // from the 'actual' literal.
-    int64_t linear_index;
-
-    int float_distance = -1;
-
-    bool operator<(const Mismatch& other) const {
-      return rel_error < other.rel_error;
-    }
-
-    std::string ToString(const Shape& shape) const {
-      auto s = absl::StrFormat(
-          "actual %s, expected %s, index %s, rel error %8.3g, abs error "
-          "%8.3g",
-          FpValueToString(actual), FpValueToString(expected),
-          LiteralUtil::MultiIndexAsString(
-              IndexUtil::LinearIndexToMultidimensionalIndex(shape,
-                                                            linear_index)),
-          rel_error, abs_error);
-      if (float_distance >= 0) {
-        StrAppendFormat(&s, ", float distance %d", float_distance);
-      }
-      return s;
-    }
-  };
-
-  NearComparator(const LiteralSlice& expected, const LiteralSlice& actual,
-                 const ShapeIndex& shape_index, ErrorSpec error,
-                 bool detailed_message,
-                 const MiscompareCallback& miscompare_callback)
+// Base class for NearComparator<NativeT> factored out to avoid duplicating
+// type-independent state, setup, and error formatting across all 17
+// floating-point and complex primitive type instantiations.
+class NearComparatorBase {
+ protected:
+  NearComparatorBase(const LiteralSlice& expected, const LiteralSlice& actual,
+                     const ShapeIndex& shape_index, ErrorSpec error,
+                     bool detailed_message,
+                     const MiscompareCallback& miscompare_callback)
       : expected_(expected),
         actual_(actual),
         shape_index_(shape_index),
@@ -323,8 +277,8 @@ class NearComparator {
         rel_error_buckets_(kErrorBucketBounds.size(), 0),
         candidates_({error}) {}
 
-  // Runs the comparison between expected and actual literals.
-  absl::Status Run() {
+  // Validates the input shapes and initializes mismatches_ before comparing.
+  absl::Status Prepare() {
     // If the shapes mismatch, we simply fail the expectation instead of
     // printing out data, as it's a type error rather than a value error.
     ABSL_RETURN_IF_ERROR(EqualShapes(expected_.shape(), actual_.shape()));
@@ -335,22 +289,24 @@ class NearComparator {
 
     mismatches_ = Literal(ShapeUtil::ChangeElementType(actual_.shape(), PRED));
     mismatches_.PopulateWithValue(false);
+    return absl::OkStatus();
+  }
 
-    CompareLiterals();
-
-    if (num_mismatches_ == 0) {
-      return absl::OkStatus();
-    } else if (!VLOG_IS_ON(1) && miscompare_callback_ != nullptr) {
+  // Invokes the miscompare callback (if any) and returns the formatted error
+  // status when mismatches are present.
+  absl::Status Finalize(absl::string_view top_mismatches) {
+    if (!VLOG_IS_ON(1) && miscompare_callback_ != nullptr) {
       miscompare_callback_(
           expected_, actual_, mismatches_, shape_index_,
           ErrorBuckets(abs_error_buckets_, rel_error_buckets_));
     }
-    return InvalidArgument("%s", ErrorMessage());
+    return InvalidArgument("%s", ErrorMessage(top_mismatches));
   }
 
   // Insert the given absolute value into the absolute value bucket vector. The
   // bounds of the buckets are given by kAbsValueBucketBounds.
-  void UpdateAbsValueBucket(NativeT value, bool is_mismatch) {
+  template <typename T>
+  void UpdateAbsValueBucket(T value, bool is_mismatch) {
     // Adjust the bucket containing the absolute values of the 'actual'
     // elements.
     const double abs_value = FpAbsoluteValue(value);
@@ -453,6 +409,218 @@ class NearComparator {
         error_.low_precision_fp_error_spec.type);
   }
 
+  // Returns an error message string with a detailed breakdown of the
+  // mismatches. Called after calling Run().
+  std::string ErrorMessage(absl::string_view top_mismatches) {
+    std::string out;
+    int64_t element_count = ShapeUtil::ElementsIn(actual_.shape());
+
+    auto percent_string = [](double a, double b) {
+      double pct = b == 0.0 ? 0.0 : 100.0 * a / b;
+      return absl::StrFormat("%0.4f%%", pct);
+    };
+
+    StrAppendFormat(
+        &out,
+        "\nMismatch count %d (%s) in shape %s (%d elements), abs bound "
+        "%g, rel bound %g\n",
+        num_mismatches_, percent_string(num_mismatches_, element_count),
+        ShapeUtil::HumanString(actual_.shape()),
+        ShapeUtil::ElementsIn(actual_.shape()), error_.abs, error_.rel);
+    if (num_nan_mismatches_ > 0) {
+      StrAppend(&out, "nan mismatches ", num_nan_mismatches_, "\n");
+    }
+    StrAppendFormat(&out, "Top relative error mismatches:\n");
+    StrAppend(&out, top_mismatches);
+
+    if (!detailed_message_) {
+      return out;
+    }
+
+    StrAppend(&out, "Absolute magnitude breakdown of actual values:\n");
+    CHECK_EQ(abs_value_buckets_.size() + 1, kAbsValueBucketBounds.size());
+    for (int i = 0; i < abs_value_buckets_.size(); ++i) {
+      const int64_t bucket_size = abs_value_buckets_[i].first;
+      const int64_t bucket_mismatches = abs_value_buckets_[i].second;
+      std::string mismatch_str =
+          bucket_mismatches > 0
+              ? absl::StrFormat(", mismatches %d", bucket_mismatches)
+              : "";
+      StrAppendFormat(&out, "  %-6g <= x < %-6g : %7d (%9s)%s\n",
+                      kAbsValueBucketBounds[i], kAbsValueBucketBounds[i + 1],
+                      bucket_size, percent_string(bucket_size, element_count),
+                      mismatch_str);
+    }
+
+    auto print_accum_buckets = [&](const std::string& header, int64_t total,
+                                   absl::Span<const int64_t> buckets) {
+      StrAppend(&out, header, ":\n");
+      StrAppendFormat(&out, "  <  %-6g : %7d (%s)\n", kErrorBucketBounds[0],
+                      total - buckets[0],
+                      percent_string(total - buckets[0], total));
+      CHECK_EQ(buckets.size(), kErrorBucketBounds.size());
+      for (int i = 0; i < kErrorBucketBounds.size(); ++i) {
+        StrAppendFormat(&out, "  >= %-6g : %7d (%s)\n", kErrorBucketBounds[i],
+                        buckets[i], percent_string(buckets[i], total));
+      }
+    };
+    StrAppendFormat(&out, "Elements exceeding abs error bound %g: %d (%s)\n",
+                    error_.abs, num_abs_mismatches_,
+                    percent_string(num_abs_mismatches_, element_count));
+    print_accum_buckets(
+        "Relative error breakdown of elements exceeding abs error bound",
+        num_abs_mismatches_, rel_error_buckets_);
+    StrAppendFormat(&out, "Elements exceeding rel error bound %g: %d (%s)\n",
+                    error_.rel, num_rel_mismatches_,
+                    percent_string(num_rel_mismatches_, element_count));
+    print_accum_buckets(
+        "Absolute error breakdown of elements exceeding rel error bound",
+        num_rel_mismatches_, abs_error_buckets_);
+
+    if (num_mismatches_ > 0) {
+      StrAppendFormat(&out, "\nCurrent ErrorSpec: ErrorSpec{%g, %g}\n",
+                      error_.abs, error_.rel);
+      StrAppend(&out,
+                "Suggested ErrorSpec adjustments to make this test pass:\n");
+
+      std::vector<ErrorSpec> rounded_candidates;
+      for (const auto& c : candidates_) {
+        InsertCandidate(
+            ErrorSpec(RoundUpTo1SigFig(c.abs), RoundUpTo1SigFig(c.rel)),
+            rounded_candidates);
+      }
+
+      for (size_t i = 0; i < rounded_candidates.size(); ++i) {
+        StrAppendFormat(&out, "  Option %d: ErrorSpec{%g, %g}\n", i + 1,
+                        rounded_candidates[i].abs, rounded_candidates[i].rel);
+      }
+    }
+
+    return out;
+  }
+
+  // 'actual' and 'expected' literals being compared.
+  LiteralSlice expected_;
+  LiteralSlice actual_;
+
+  // The shape index of the LiteralSlice that is being compared.
+  ShapeIndex shape_index_;
+
+  // The error bounds of the comparison.
+  ErrorSpec error_;
+
+  // Whether to include detailed breakdown of mismatches in the error message.
+  bool detailed_message_;
+
+  // Callback to invoke on miscompare.
+  MiscompareCallback miscompare_callback_;
+
+  // Number of element mismatches encountered so far.
+  int64_t num_mismatches_ = 0;
+
+  // Number of elements with a nan mismatch.
+  int64_t num_nan_mismatches_ = 0;
+
+  // Number of elements which exceed the absolute/relative error bound.
+  int64_t num_abs_mismatches_ = 0;
+  int64_t num_rel_mismatches_ = 0;
+
+  // A Literal containing which elements did not match in the expected and
+  // actual literals. mismatches_ contains PREDs and is of the same sizes as
+  // the comparison literals.
+  Literal mismatches_;
+
+  // The number of mismatches to report in the output, sorted by relative error
+  // magnitude.
+  static constexpr int64_t kTopRelativeErrorCount = 5;
+
+  // Actual values are bucketed by absolute value. kAbsValueBucketBounds is the
+  // bounds of these buckets. abs_value_buckets_ contains a pair for each
+  // bucket: the element count and failure count.
+  static constexpr std::array<double, 7> kAbsValueBucketBounds = {
+      0.0, 0.0001, 0.001, 0.01, 0.1, 1, std::numeric_limits<double>::infinity(),
+  };
+  std::vector<std::pair<int64_t, int64_t>> abs_value_buckets_;
+
+  // Buckets for relative and absolute errors. The relative error buckets only
+  // contains those elements which exceed the *absolute* error bound, and vice
+  // versa. This makes it easy to see the effect of adjusting the relative (or
+  // absolute) error bound on the success of the comparison. kErrorBucketBounds
+  // are the lower bounds of the buckets in both vectors. The error buckets are
+  // a cumulative distribution so an error value may appear in more than one
+  // bucket. For example an error value of 0.003 may appear in the buckets
+  // bounded by 0.01, 0.1, and 1.0.
+  static constexpr std::array<double, 5> kErrorBucketBounds = {0.0001, 0.001,
+                                                               0.01, 0.1, 1};
+  std::vector<int64_t> abs_error_buckets_;
+  std::vector<int64_t> rel_error_buckets_;
+  std::vector<ErrorSpec> candidates_;
+};
+
+// Helper class for comparing floating-point literals within an error bound.
+template <typename NativeT>
+class NearComparator : private NearComparatorBase {
+ public:
+  // Compares the two array literals elementwise and returns a comparison
+  // result. The comparison is ok() if all actual and expected elements are
+  // within the given error bound. In case of error, the status contains a
+  // detailed message about the discrepancy.
+  static absl::Status Compare(const LiteralSlice& expected,
+                              const LiteralSlice& actual,
+                              const ShapeIndex& shape_index, ErrorSpec error,
+                              bool detailed_message,
+                              const MiscompareCallback& miscompare_callback) {
+    NearComparator<NativeT> comparator(expected, actual, shape_index, error,
+                                       detailed_message, miscompare_callback);
+    return comparator.Run();
+  }
+
+ private:
+  // Data structure encapsulating metadata about a single element mismatch.
+  struct Mismatch {
+    NativeT actual;
+    NativeT expected;
+    double rel_error;
+    double abs_error;
+
+    // The linear index of the failure within the shape. This linear index is
+    // from the 'actual' literal.
+    int64_t linear_index;
+
+    int float_distance = -1;
+
+    bool operator<(const Mismatch& other) const {
+      return rel_error < other.rel_error;
+    }
+
+    std::string ToString(const Shape& shape) const {
+      auto s = absl::StrFormat(
+          "actual %s, expected %s, index %s, rel error %8.3g, abs error "
+          "%8.3g",
+          FpValueToString(actual), FpValueToString(expected),
+          LiteralUtil::MultiIndexAsString(
+              IndexUtil::LinearIndexToMultidimensionalIndex(shape,
+                                                            linear_index)),
+          rel_error, abs_error);
+      if (float_distance >= 0) {
+        StrAppendFormat(&s, ", float distance %d", float_distance);
+      }
+      return s;
+    }
+  };
+
+  using NearComparatorBase::NearComparatorBase;
+
+  // Runs the comparison between expected and actual literals.
+  absl::Status Run() {
+    ABSL_RETURN_IF_ERROR(Prepare());
+    CompareLiterals();
+    if (num_mismatches_ == 0) {
+      return absl::OkStatus();
+    }
+    return Finalize(FormatTopMismatches());
+  }
+
   // Compares the two given elements from the expected and actual literals at
   // the given literal_index and keeps track of various mismatch statistics.
   template <typename T>
@@ -632,159 +800,18 @@ class NearComparator {
     }
   }
 
-  // Returns an error message string with a detailed breakdown of the
-  // mismatches. Called after calling Run().
-  std::string ErrorMessage() {
+  std::string FormatTopMismatches() const {
     std::string out;
-    int64_t element_count = ShapeUtil::ElementsIn(actual_.shape());
-
-    auto percent_string = [](double a, double b) {
-      double pct = b == 0.0 ? 0.0 : 100.0 * a / b;
-      return absl::StrFormat("%0.4f%%", pct);
-    };
-
-    StrAppendFormat(
-        &out,
-        "\nMismatch count %d (%s) in shape %s (%d elements), abs bound "
-        "%g, rel bound %g\n",
-        num_mismatches_, percent_string(num_mismatches_, element_count),
-        ShapeUtil::HumanString(actual_.shape()),
-        ShapeUtil::ElementsIn(actual_.shape()), error_.abs, error_.rel);
-    if (num_nan_mismatches_ > 0) {
-      StrAppend(&out, "nan mismatches ", num_nan_mismatches_, "\n");
-    }
-    StrAppendFormat(&out, "Top relative error mismatches:\n");
     for (auto it = top_rel_mismatches_.rbegin();
          it != top_rel_mismatches_.rend(); ++it) {
       StrAppend(&out, "  ", it->ToString(actual_.shape()), "\n");
     }
-
-    if (!detailed_message_) {
-      return out;
-    }
-
-    StrAppend(&out, "Absolute magnitude breakdown of actual values:\n");
-    CHECK_EQ(abs_value_buckets_.size() + 1, kAbsValueBucketBounds.size());
-    for (int i = 0; i < abs_value_buckets_.size(); ++i) {
-      const int64_t bucket_size = abs_value_buckets_[i].first;
-      const int64_t bucket_mismatches = abs_value_buckets_[i].second;
-      std::string mismatch_str =
-          bucket_mismatches > 0
-              ? absl::StrFormat(", mismatches %d", bucket_mismatches)
-              : "";
-      StrAppendFormat(&out, "  %-6g <= x < %-6g : %7d (%9s)%s\n",
-                      kAbsValueBucketBounds[i], kAbsValueBucketBounds[i + 1],
-                      bucket_size, percent_string(bucket_size, element_count),
-                      mismatch_str);
-    }
-
-    auto print_accum_buckets = [&](const std::string& header, int64_t total,
-                                   absl::Span<const int64_t> buckets) {
-      StrAppend(&out, header, ":\n");
-      StrAppendFormat(&out, "  <  %-6g : %7d (%s)\n", kErrorBucketBounds[0],
-                      total - buckets[0],
-                      percent_string(total - buckets[0], total));
-      CHECK_EQ(buckets.size(), kErrorBucketBounds.size());
-      for (int i = 0; i < kErrorBucketBounds.size(); ++i) {
-        StrAppendFormat(&out, "  >= %-6g : %7d (%s)\n", kErrorBucketBounds[i],
-                        buckets[i], percent_string(buckets[i], total));
-      }
-    };
-    StrAppendFormat(&out, "Elements exceeding abs error bound %g: %d (%s)\n",
-                    error_.abs, num_abs_mismatches_,
-                    percent_string(num_abs_mismatches_, element_count));
-    print_accum_buckets(
-        "Relative error breakdown of elements exceeding abs error bound",
-        num_abs_mismatches_, rel_error_buckets_);
-    StrAppendFormat(&out, "Elements exceeding rel error bound %g: %d (%s)\n",
-                    error_.rel, num_rel_mismatches_,
-                    percent_string(num_rel_mismatches_, element_count));
-    print_accum_buckets(
-        "Absolute error breakdown of elements exceeding rel error bound",
-        num_rel_mismatches_, abs_error_buckets_);
-
-    if (num_mismatches_ > 0) {
-      StrAppendFormat(&out, "\nCurrent ErrorSpec: ErrorSpec{%g, %g}\n",
-                      error_.abs, error_.rel);
-      StrAppend(&out,
-                "Suggested ErrorSpec adjustments to make this test pass:\n");
-
-      std::vector<ErrorSpec> rounded_candidates;
-      for (const auto& c : candidates_) {
-        InsertCandidate(
-            ErrorSpec(RoundUpTo1SigFig(c.abs), RoundUpTo1SigFig(c.rel)),
-            rounded_candidates);
-      }
-
-      for (size_t i = 0; i < rounded_candidates.size(); ++i) {
-        StrAppendFormat(&out, "  Option %d: ErrorSpec{%g, %g}\n", i + 1,
-                        rounded_candidates[i].abs, rounded_candidates[i].rel);
-      }
-    }
-
     return out;
   }
 
-  // 'actual' and 'expected' literals being compared.
-  LiteralSlice expected_;
-  LiteralSlice actual_;
-
-  // The shape index of the LiteralSlice that is being compared.
-  ShapeIndex shape_index_;
-
-  // The error bounds of the comparison.
-  ErrorSpec error_;
-
-  // Whether to include detailed breakdown of mismatches in the error message.
-  bool detailed_message_;
-
-  // Callback to invoke on miscompare.
-  MiscompareCallback miscompare_callback_;
-
-  // Number of element mismatches encountered so far.
-  int64_t num_mismatches_ = 0;
-
-  // Number of elements with a nan mismatch.
-  int64_t num_nan_mismatches_ = 0;
-
-  // Number of elements which exceed the absolute/relative error bound.
-  int64_t num_abs_mismatches_ = 0;
-  int64_t num_rel_mismatches_ = 0;
-
-  // A Literal containing which elements did not match in the expected and
-  // actual literals. mismatches_ contains PREDs and is of the same sizes as
-  // the comparison literals.
-  Literal mismatches_;
-
-  // The number of mismatches to report in the output, sorted by relative error
-  // magnitude.
-  static constexpr int64_t kTopRelativeErrorCount = 5;
-
   // The set of mismatches with the largest relative error. The size of this set
   // is bounded by kTopRelativeErrorCount.
   std::multiset<Mismatch> top_rel_mismatches_;
-
-  // Actual values are bucketed by absolute value. kAbsValueBucketBounds is the
-  // bounds of these buckets. abs_value_buckets_ contains a pair for each
-  // bucket: the element count and failure count.
-  static constexpr std::array<double, 7> kAbsValueBucketBounds = {
-      0.0, 0.0001, 0.001, 0.01, 0.1, 1, std::numeric_limits<double>::infinity(),
-  };
-  std::vector<std::pair<int64_t, int64_t>> abs_value_buckets_;
-
-  // Buckets for relative and absolute errors. The relative error buckets only
-  // contains those elements which exceed the *absolute* error bound, and vice
-  // versa. This makes it easy to see the effect of adjusting the relative (or
-  // absolute) error bound on the success of the comparison. kErrorBucketBounds
-  // are the lower bounds of the buckets in both vectors. The error buckets are
-  // a cumulative distribution so an error value may appear in more than one
-  // bucket. For example an error value of 0.003 may appear in the buckets
-  // bounded by 0.01, 0.1, and 1.0.
-  static constexpr std::array<double, 5> kErrorBucketBounds = {0.0001, 0.001,
-                                                               0.01, 0.1, 1};
-  std::vector<int64_t> abs_error_buckets_;
-  std::vector<int64_t> rel_error_buckets_;
-  std::vector<ErrorSpec> candidates_;
 };
 
 absl::Status EqualHelper(const LiteralSlice& expected,

From 900ad684790e6ed393edc0945791dff6339fcb1b Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Wed, 14 Oct 2026 08:15:13 -0700
Subject: [PATCH 0786/1000] [XLA:GPU]: Support for tiling variadic
 reduce-scatter.

PiperOrigin-RevId: 999380359
---
 .../xla/xla/codegen/tiling/experimental/BUILD |   2 +
 .../codegen/tiling/experimental/scheduling.cc |  62 +++--
 .../tiling/experimental/scheduling_test.cc    | 111 +++++++++
 .../tiling/experimental/tile_propagation.cc   |  32 ++-
 .../experimental/tile_propagation_test.cc     | 216 ++++++++++++++++++
 .../tiling/experimental/tiled_hlo_test.cc     | 170 ++++++++++++++
 .../tiling/experimental/tiling_space.cc       |  44 ++--
 .../tiling/experimental/tiling_space.h        |   9 +-
 .../tiling/experimental/tiling_space_test.cc  |  67 +++++-
 9 files changed, 655 insertions(+), 58 deletions(-)

diff --git a/third_party/xla/xla/codegen/tiling/experimental/BUILD b/third_party/xla/xla/codegen/tiling/experimental/BUILD
index df27fcac371..b7ad1878827 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/BUILD
+++ b/third_party/xla/xla/codegen/tiling/experimental/BUILD
@@ -196,6 +196,7 @@ xla_cc_test(
         "//xla/hlo/analysis:interval",
         "//xla/hlo/analysis:symbolic_map",
         "//xla/hlo/ir:hlo",
+        "//xla/hlo/testlib:filecheck",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
         "//xla/hlo/testlib:verified_hlo_module",
         "//xla/hlo/utils:hlo_traversal",
@@ -308,6 +309,7 @@ xla_cc_test(
         "//xla/hlo/analysis:symbolic_map",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/parser:hlo_parser",
+        "//xla/hlo/testlib:filecheck",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
         "//xla/hlo/testlib:verified_hlo_module",
         "//xla/hlo/utils:hlo_traversal",
diff --git a/third_party/xla/xla/codegen/tiling/experimental/scheduling.cc b/third_party/xla/xla/codegen/tiling/experimental/scheduling.cc
index 7c0bdc0debd..8cdb059c0af 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/scheduling.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/scheduling.cc
@@ -295,8 +295,7 @@ struct IndexWiseVariadicInfo {
   //   - Each GPU owns a 4-row slice, which is 4 / 2 = 2 tiles along `d0`.
   //   - So `tiles_per_peer_by_dim[d0] = 2`.
   llvm::DenseMap<int64_t, int64_t> tiles_per_peer_by_dim;
-  // Number of peers the outputs are read from. 1 if `tiles_per_peer_by_dim`
-  // is empty.
+  // Number of peers across `tiles_per_peer_by_dim` and `kReplica` dimensions.
   int64_t num_peers = 1;
 };
 
@@ -404,14 +403,26 @@ absl::StatusOr<IndexWiseVariadicInfo> BuildIndexWiseVariadicInfo(
       info.tiles_per_peer_by_dim[dim_id] = tiles_per_peer;
     }
   }
+  for (const TilingSpace::DimensionInfo& dim : dims) {
+    if (dim.type != TilingSpace::DimensionSemantics::kReplica) {
+      continue;
+    }
+    if (num_peers.has_value() && *num_peers != dim.dimension_size) {
+      return absl::InvalidArgumentError(absl::StrFormat(
+          "Replica dimension %d of %s has %d replicas, but other dimensions "
+          "have %d peers.",
+          dim.id.value(), dim.hlo->name(), dim.dimension_size, *num_peers));
+    }
+    num_peers = dim.dimension_size;
+  }
   info.num_peers = num_peers.value_or(1);
   return info;
 }
 
 // Internal struct for holding the parallel dimensions of a root.
 struct RootDims {
-  // Parallel output dimension whose tile ID encodes the peer index
-  // (all-gather).
+  // Parallel output dimension (all-gather) or `kReplica` dimension
+  // (reduce-scatter) whose tile ID encodes the peer index.
   std::optional<int64_t> peer_dim_id;
   int64_t tiles_per_peer = 1;
   llvm::SmallVector<int64_t, 4> non_peer_dim_ids;
@@ -423,7 +434,6 @@ absl::StatusOr<llvm::SmallVector<RootDims, 2>> BuildRootsForIndexWiseVariadic(
     const TiledHloComputation& tiled_computation,
     const IndexWiseVariadicInfo& info) {
   const TilingSpace& tiling_space = tiled_computation.tiling_space();
-  const bool has_peer_dim = !info.tiles_per_peer_by_dim.empty();
   // Group the parallel dimensions by the root (output) that defines them.
   llvm::SmallVector<RootDims, 2> roots;
   int64_t num_grouped_dims = 0;
@@ -440,22 +450,28 @@ absl::StatusOr<llvm::SmallVector<RootDims, 2>> BuildRootsForIndexWiseVariadic(
       ++num_grouped_dims;
       const int64_t dim_id = dim->id.value();
       auto it = info.tiles_per_peer_by_dim.find(dim_id);
-      if (it == info.tiles_per_peer_by_dim.end()) {
-        root.non_peer_dim_ids.push_back(dim_id);
-        root.non_peer_block_counts.push_back(
-            CeilOfRatio(dim->dimension_size, *dim->tile_size));
-        continue;
+      if (dim->replica_dim.has_value() ||
+          it != info.tiles_per_peer_by_dim.end()) {
+        if (root.peer_dim_id.has_value() ||
+            (dim->replica_dim.has_value() &&
+             it != info.tiles_per_peer_by_dim.end())) {
+          return absl::UnimplementedError(absl::StrFormat(
+              "Output %d of %s has more than one dimension carrying the peer "
+              "index.",
+              root_index, hlo.name()));
+        }
+        if (dim->replica_dim.has_value()) {
+          root.peer_dim_id = dim->replica_dim->value();
+        } else {
+          root.peer_dim_id = dim_id;
+          root.tiles_per_peer = it->second;
+          continue;
+        }
       }
-      if (root.peer_dim_id.has_value()) {
-        return absl::UnimplementedError(absl::StrFormat(
-            "Output %d of %s has more than one dimension carrying the peer "
-            "index.",
-            root_index, hlo.name()));
-      }
-      root.peer_dim_id = dim_id;
-      root.tiles_per_peer = it->second;
+      root.non_peer_dim_ids.push_back(dim_id);
+      root.non_peer_block_counts.push_back(tiling_space.GetNumTiles(*dim));
     }
-    if (has_peer_dim && !root.peer_dim_id.has_value()) {
+    if (info.num_peers > 1 && !root.peer_dim_id.has_value()) {
       return absl::UnimplementedError(
           absl::StrFormat("Output %d of %s does not depend on the peer index.",
                           root_index, hlo.name()));
@@ -472,11 +488,11 @@ absl::StatusOr<llvm::SmallVector<RootDims, 2>> BuildRootsForIndexWiseVariadic(
 absl::StatusOr<Schedule> BuildIndexWiseVariadicSchedule(
     mlir::MLIRContext* ctx, const IndexWiseVariadicInfo& info,
     llvm::ArrayRef<RootDims> roots, int64_t num_tiles_per_pid,
-    int64_t max_tiles, bool has_peer_dim, SymbolicExpr global_tile_id) {
+    int64_t max_tiles, SymbolicExpr global_tile_id) {
   Schedule schedule;
   SymbolicExpr peer = global_tile_id.floorDiv(max_tiles);
   SymbolicExpr intra_peer_tile_id =
-      has_peer_dim ? global_tile_id % max_tiles : global_tile_id;
+      info.num_peers > 1 ? global_tile_id % max_tiles : global_tile_id;
   for (const RootDims& root : roots) {
     llvm::SmallVector<int64_t, 4> block_counts = root.non_peer_block_counts;
     if (root.tiles_per_peer > 1) {
@@ -530,7 +546,6 @@ absl::StatusOr<Schedule> GetIndexWiseVariadicSchedule(
     const IndexWiseVariadicInfo& info, SymbolicExpr global_tile_id,
     int64_t num_tiles_per_pid) {
   mlir::MLIRContext* ctx = tiled_computation.GetMLIRContext();
-  const bool has_peer_dim = !info.tiles_per_peer_by_dim.empty();
   using Roots = llvm::SmallVector<RootDims, 2>;
   ABSL_ASSIGN_OR_RETURN(Roots roots,
                    BuildRootsForIndexWiseVariadic(tiled_computation, info));
@@ -542,8 +557,7 @@ absl::StatusOr<Schedule> GetIndexWiseVariadicSchedule(
   }
   // Now build the scheduled rooted on the output with the max number of tiles.
   return BuildIndexWiseVariadicSchedule(ctx, info, roots, num_tiles_per_pid,
-                                        max_tiles, has_peer_dim,
-                                        global_tile_id);
+                                        max_tiles, global_tile_id);
 }
 
 }  // namespace
diff --git a/third_party/xla/xla/codegen/tiling/experimental/scheduling_test.cc b/third_party/xla/xla/codegen/tiling/experimental/scheduling_test.cc
index bb4b43876b2..f22cf4a97ff 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/scheduling_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/scheduling_test.cc
@@ -556,6 +556,117 @@ TEST_F(SchedulingTest, VariadicAllGatherWithDotProducersSharesPeerIndex) {
           "num_pids=16, num_tiles=16")));
 }
 
+// Replica dims d4 and d5 share `pid / 2`; root 1 is padded on d5.
+TEST_F(SchedulingTest, VariadicReduceScatterSharesShardIndex) {
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"(
+    add {
+      lhs = f32[] parameter(0)
+      rhs = f32[] parameter(1)
+      ROOT add = f32[] add(lhs, rhs)
+    }
+
+    f {
+      p0 = f32[4,64] parameter(0)
+      p1 = f32[4,1] parameter(1)
+      rs = (f32[2,64], f32[2,1]) reduce-scatter(p0, p1),
+        replica_groups={{0,1}}, dimensions={0}, to_apply=add
+      gte0 = f32[2,64] get-tuple-element(rs), index=0
+      gte1 = f32[2,1] get-tuple-element(rs), index=1
+      ROOT t = (f32[2,64], f32[2,1]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      arg0 = f32[4,64] parameter(0)
+      arg1 = f32[4,1] parameter(1)
+      ROOT fusion = (f32[2,64], f32[2,1]) fusion(arg0, arg1),
+        kind=kCustom, calls=f
+    })",
+                                    {1, 32, 1, 1}));
+  EXPECT_THAT(
+      GetSchedule(tiled_computation),
+      IsOkAndHolds(MatchSchedule(
+          "d0 -> (pid mod 2) / 2, d1 -> pid mod 2 mod 2, d2 -> pid mod 2, "
+          "d3 -> 0, d4 -> pid / 2, d5 -> pid / 2 + (pid mod 2) * 2, "
+          "num_pids=4, num_tiles=4")));
+}
+
+// Output 0 has 2 scattered tiles/shard on d0; d4 and d5 carry the peer index.
+TEST_F(SchedulingTest, VariadicReduceScatterMultipleTilesPerShard) {
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"(
+    add {
+      lhs = f32[] parameter(0)
+      rhs = f32[] parameter(1)
+      ROOT add = f32[] add(lhs, rhs)
+    }
+
+    f {
+      p0 = f32[8,64] parameter(0)
+      p1 = f32[8,1] parameter(1)
+      rs = (f32[4,64], f32[4,1]) reduce-scatter(p0, p1),
+        replica_groups={{0,1}}, dimensions={0}, to_apply=add
+      gte0 = f32[4,64] get-tuple-element(rs), index=0
+      gte1 = f32[4,1] get-tuple-element(rs), index=1
+      ROOT t = (f32[4,64], f32[4,1]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      arg0 = f32[8,64] parameter(0)
+      arg1 = f32[8,1] parameter(1)
+      ROOT fusion = (f32[4,64], f32[4,1]) fusion(arg0, arg1),
+        kind=kCustom, calls=f
+    })",
+                                    {1, 64, 2, 1}));
+  EXPECT_THAT(GetSchedule(tiled_computation),
+              IsOkAndHolds(MatchSchedule(
+                  "d0 -> pid mod 2, d1 -> 0, d2 -> pid mod 2, d3 -> 0, "
+                  "d4 -> pid / 2, d5 -> pid / 2 + (pid mod 2) * 2, "
+                  "num_pids=4, num_tiles=4")));
+}
+
+TEST_F(SchedulingTest, VariadicReduceScatterWithDotProducers) {
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"(
+    add {
+      lhs = f32[] parameter(0)
+      rhs = f32[] parameter(1)
+      ROOT add = f32[] add(lhs, rhs)
+    }
+
+    f {
+      lhs0 = f32[16,8] parameter(0)
+      rhs0 = f32[8,16] parameter(1)
+      lhs1 = f32[32,8] parameter(2)
+      rhs1 = f32[8,32] parameter(3)
+      dot0 = f32[16,16] dot(lhs0, rhs0),
+        lhs_contracting_dims={1}, rhs_contracting_dims={0}
+      dot1 = f32[32,32] dot(lhs1, rhs1),
+        lhs_contracting_dims={1}, rhs_contracting_dims={0}
+      rs = (f32[4,16], f32[8,32]) reduce-scatter(dot0, dot1),
+        replica_groups={{0,1,2,3}}, dimensions={0}, to_apply=add
+      gte0 = f32[4,16] get-tuple-element(rs), index=0
+      gte1 = f32[8,32] get-tuple-element(rs), index=1
+      ROOT t = (f32[4,16], f32[8,32]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      lhs0 = f32[16,8] parameter(0)
+      rhs0 = f32[8,16] parameter(1)
+      lhs1 = f32[32,8] parameter(2)
+      rhs1 = f32[8,32] parameter(3)
+      ROOT fusion = (f32[4,16], f32[8,32]) fusion(lhs0, rhs0, lhs1, rhs1),
+        kind=kCustom, calls=f
+    })",
+                                    {1, 8, 1, 16, 8, 8}));
+  EXPECT_THAT(GetSchedule(tiled_computation),
+              IsOkAndHolds(MatchSchedule(
+                  "d0 -> (pid mod 4) / 2, d1 -> pid mod 4 mod 2, "
+                  "d2 -> (pid mod 4) / 2, d3 -> pid mod 4 mod 2, "
+                  "d4 -> pid / 4 + ((pid mod 4) / 2) * 4, d5 -> pid / 4, "
+                  "num_pids=16, num_tiles=16")));
+}
+
 TEST_F(SchedulingTest, SingleOperandReduceScatterReplicaDimension) {
   ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
                        ParseAndTile(R"(
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
index ac56ca5809d..cbf1da9cc44 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
@@ -1589,15 +1589,22 @@ absl::StatusOr<Tiles> PropagateTileToOutputForBitcastOp(
 //    (64 rows). This guarantees that the tile never crosses the boundary
 //    between shards (e.g., crossing row 64), meaning this thread block
 //    only ever has to send its data to exactly *one* target GPU.
+//
+// For a variadic reduce-scatter, output `output_index` depends only on operand
+// `output_index`; the scatter dimension and replica count are shared by all
+// operands, so the same mapping applies to each of them independently.
 absl::StatusOr<Tiles> PropagateTileToInputForReduceScatterOp(
     const TilingSpace& tiling_space,
-    const HloReduceScatterInstruction& reduce_scatter,
-    const Tile& output_tile) {
-  // Crash OK. Must be checked while forming the fusion.
-  CHECK_EQ(reduce_scatter.operand_count(), 1)
-      << "Multi-operand ReduceScatter is not yet supported.";
+    const HloReduceScatterInstruction& reduce_scatter, const Tile& output_tile,
+    int64_t output_index) {
+  if (output_index < 0 || output_index >= reduce_scatter.operand_count()) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Output index ", output_index, " is out of range for ",
+                     reduce_scatter.ToString()));
+  }
   const int64_t scatter_dim = reduce_scatter.scatter_dimension();
-  const int64_t shard_size = reduce_scatter.shape().dimensions(scatter_dim);
+  const int64_t shard_size =
+      GetFirstShape(&reduce_scatter, output_index).dimensions(scatter_dim);
   const int64_t num_replicas =
       reduce_scatter.device_list()->num_devices_per_group();
 
@@ -1611,9 +1618,11 @@ absl::StatusOr<Tiles> PropagateTileToInputForReduceScatterOp(
   if (num_replicas <= 1) {
     return Tiles{output_tile};
   }
+  const int64_t total_result_dims =
+      GetResultDimensionOffset(reduce_scatter, reduce_scatter.operand_count());
   const TilingSpace::DimensionInfo& replica_dim_info =
       tiling_space.GetDimensionInfo(reduce_scatter,
-                                    reduce_scatter.shape().dimensions().size());
+                                    total_result_dims + output_index);
   SymbolicExpr replica_id =
       CreateDimExpr(replica_dim_info.id.value(), tiling_space.mlir_context());
   scatter_tile.offset = (scatter_tile.offset + replica_id * shard_size -
@@ -1628,9 +1637,9 @@ absl::StatusOr<Tiles> PropagateTileToInputForReduceScatterOp(
     if (!size.has_value() || shard_size % *size != 0 ||
         !scatter_tile.offset.IsMultipleOf(*size)) {
       return absl::FailedPreconditionError(absl::StrCat(
-          "Tiling propagation rejected for ", reduce_scatter.ToString(),
-          ": operand tile ", scatter_tile.size.ToString(), " at ",
-          scatter_tile.offset.ToString(),
+          "Tiling propagation rejected for output ", output_index, " of ",
+          reduce_scatter.ToString(), ": operand tile ",
+          scatter_tile.size.ToString(), " at ", scatter_tile.offset.ToString(),
           " crosses shard boundary (shard_size=", shard_size, ")"));
     }
   }
@@ -1818,7 +1827,8 @@ absl::StatusOr<Tiles> PropagateTileToInput(TilingSpace& tiling_space,
   }
   if (hlo.opcode() == HloOpcode::kReduceScatter) {
     return PropagateTileToInputForReduceScatterOp(
-        tiling_space, *Cast<HloReduceScatterInstruction>(&hlo), output_tile);
+        tiling_space, *Cast<HloReduceScatterInstruction>(&hlo), output_tile,
+        output_index);
   }
   if (hlo.opcode() == HloOpcode::kBitcast) {
     return PropagateTileToInputForBitcastOp(hlo, output_tile);
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
index f061111e255..7737e6f806c 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
@@ -590,6 +590,222 @@ TEST_F(TilePropagationTest, ReduceScatterRejectsTileCrossingShards) {
               StatusIs(absl::StatusCode::kFailedPrecondition));
 }
 
+TEST_F(TilePropagationTest, CanPropagateToInputsOfVariadicReduceScatterOp) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    add {
+      lhs = f32[] parameter(0)
+      rhs = f32[] parameter(1)
+      ROOT add = f32[] add(lhs, rhs)
+    }
+    f {
+      p0 = f32[4,64] parameter(0)
+      p1 = f32[4,1] parameter(1)
+      rs = (f32[2,64], f32[2,1]) reduce-scatter(p0, p1),
+        replica_groups={{0,1}}, dimensions={0}, to_apply=add
+      gte0 = f32[2,64] get-tuple-element(rs), index=0
+      gte1 = f32[2,1] get-tuple-element(rs), index=1
+      ROOT t = (f32[2,64], f32[2,1]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      p0 = f32[4,64] parameter(0)
+      p1 = f32[4,1] parameter(1)
+      ROOT fusion = (f32[2,64], f32[2,1]) fusion(p0, p1), kind=kCustom,
+        calls=f
+    }
+  )");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TilingSpace> tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                          &mlir_context_));
+  const HloInstruction* rs =
+      root->fused_instructions_computation()->GetInstructionWithName("rs");
+  ASSERT_EQ(tiling_space->tiled_roots().size(), 2);
+  ASSERT_OK_AND_ASSIGN(
+      Tiles operand_0_tiles,
+      PropagateTileToInput(*tiling_space, *rs, tiling_space->tiled_roots()[0],
+                           /*output_index=*/0));
+  EXPECT_THAT(operand_0_tiles, MatchToString(R"(
+    0) (tid_0, tid_1, tid_2, tid_3, tid_4, tid_5)
+      -> offsets [tid_0 * ts_0 * 2 + tid_4 * 2, tid_1 * ts_1]
+         sizes [ts_0 * 2, ts_1]
+         strides [1, 1]
+         upper bounds [4, 64]
+  )"));
+
+  ASSERT_OK_AND_ASSIGN(
+      Tiles operand_1_tiles,
+      PropagateTileToInput(*tiling_space, *rs, tiling_space->tiled_roots()[1],
+                           /*output_index=*/1));
+  EXPECT_THAT(operand_1_tiles, MatchToString(R"(
+    0) (tid_0, tid_1, tid_2, tid_3, tid_4, tid_5)
+      -> offsets [tid_2 * ts_2 * 2 + tid_5 * 2, tid_3 * ts_3]
+         sizes [ts_2 * 2, ts_3]
+         strides [1, 1]
+         upper bounds [4, 1]
+  )"));
+}
+
+TEST_F(TilePropagationTest, CanPropagateFromVariadicReduceScatterToDotOps) {
+  HloInstruction* fusion = ParseAndGetRoot(R"(
+    HloModule m
+    add {
+      lhs = f32[] parameter(0)
+      rhs = f32[] parameter(1)
+      ROOT add = f32[] add(lhs, rhs)
+    }
+    fusion {
+      lhs0 = f32[128,256] parameter(0)
+      rhs0 = f32[256,512] parameter(1)
+      lhs1 = f32[64,128] parameter(2)
+      rhs1 = f32[128,256] parameter(3)
+      dot0 = f32[128,512] dot(lhs0, rhs0),
+        lhs_contracting_dims={1}, rhs_contracting_dims={0}
+      dot1 = f32[64,256] dot(lhs1, rhs1),
+        lhs_contracting_dims={1}, rhs_contracting_dims={0}
+      rs = (f32[64,512], f32[32,256]) reduce-scatter(dot0, dot1),
+        replica_groups={{0,1}}, dimensions={0}, to_apply=add
+      gte0 = f32[64,512] get-tuple-element(rs), index=0
+      gte1 = f32[32,256] get-tuple-element(rs), index=1
+      ROOT t = (f32[64,512], f32[32,256]) tuple(gte0, gte1)
+    }
+    ENTRY e {
+      lhs0 = f32[128,256] parameter(0)
+      rhs0 = f32[256,512] parameter(1)
+      lhs1 = f32[64,128] parameter(2)
+      rhs1 = f32[128,256] parameter(3)
+      ROOT fusion = (f32[64,512], f32[32,256])
+        fusion(lhs0, rhs0, lhs1, rhs1), kind=kCustom, calls=fusion
+    }
+  )");
+  const HloInstruction* rs =
+      fusion->fused_instructions_computation()->GetInstructionWithName("rs");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TilingSpace> tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(fusion),
+                          &mlir_context_));
+  ASSERT_OK_AND_ASSIGN(
+      Tiles rs_operand_0,
+      PropagateTileToInput(*tiling_space, *rs, tiling_space->tiled_roots()[0],
+                           /*output_index=*/0));
+  ASSERT_OK_AND_ASSIGN(
+      Tiles dot0_operands,
+      PropagateTileToInput(*tiling_space, *rs->operand(0), rs_operand_0[0], 0));
+  EXPECT_THAT(dot0_operands, MatchToString(R"(
+    0) (tid_0, tid_1, tid_2, tid_3, tid_4, tid_5, tid_6, tid_7)
+      -> offsets [tid_0 * ts_0 * 2 + tid_4 * 64, tid_7 * ts_7]
+         sizes [ts_0 * 2, ts_7]
+         strides [1, 1]
+         upper bounds [128, 256]
+    1) (tid_0, tid_1, tid_2, tid_3, tid_4, tid_5, tid_6, tid_7)
+      -> offsets [tid_7 * ts_7, tid_1 * ts_1]
+         sizes [ts_7, ts_1]
+         strides [1, 1]
+         upper bounds [256, 512]
+  )"));
+
+  ASSERT_OK_AND_ASSIGN(
+      Tiles rs_operand_1,
+      PropagateTileToInput(*tiling_space, *rs, tiling_space->tiled_roots()[1],
+                           /*output_index=*/1));
+  ASSERT_OK_AND_ASSIGN(
+      Tiles dot1_operands,
+      PropagateTileToInput(*tiling_space, *rs->operand(1), rs_operand_1[0], 0));
+  EXPECT_THAT(dot1_operands, MatchToString(R"(
+    0) (tid_0, tid_1, tid_2, tid_3, tid_4, tid_5, tid_6, tid_7)
+      -> offsets [tid_2 * ts_2 * 2 + tid_5 * 32, tid_6 * ts_6]
+         sizes [ts_2 * 2, ts_6]
+         strides [1, 1]
+         upper bounds [64, 128]
+    1) (tid_0, tid_1, tid_2, tid_3, tid_4, tid_5, tid_6, tid_7)
+      -> offsets [tid_6 * ts_6, tid_3 * ts_3]
+         sizes [ts_6, ts_3]
+         strides [1, 1]
+         upper bounds [128, 256]
+  )"));
+}
+
+TEST_F(TilePropagationTest, VariadicReduceScatterRejectsOutOfRangeOutputIndex) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    add {
+      lhs = f32[] parameter(0)
+      rhs = f32[] parameter(1)
+      ROOT add = f32[] add(lhs, rhs)
+    }
+    f {
+      p0 = f32[4,8] parameter(0)
+      p1 = f32[4,8] parameter(1)
+      rs = (f32[2,8], f32[2,8]) reduce-scatter(p0, p1),
+        replica_groups={{0,1}}, dimensions={0}, to_apply=add
+      gte0 = f32[2,8] get-tuple-element(rs), index=0
+      gte1 = f32[2,8] get-tuple-element(rs), index=1
+      ROOT t = (f32[2,8], f32[2,8]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      p0 = f32[4,8] parameter(0)
+      p1 = f32[4,8] parameter(1)
+      ROOT fusion = (f32[2,8], f32[2,8]) fusion(p0, p1), kind=kCustom, calls=f
+    }
+  )");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TilingSpace> tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                          &mlir_context_));
+  const HloInstruction* rs =
+      root->fused_instructions_computation()->GetInstructionWithName("rs");
+  EXPECT_THAT(
+      PropagateTileToInput(*tiling_space, *rs, tiling_space->tiled_roots()[0],
+                           /*output_index=*/2),
+      StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+TEST_F(TilePropagationTest, VariadicReduceScatterChecksShardBoundaryPerOutput) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    add {
+      lhs = f32[] parameter(0)
+      rhs = f32[] parameter(1)
+      ROOT add = f32[] add(lhs, rhs)
+    }
+    f {
+      p0 = f32[192,256] parameter(0)
+      p1 = f32[192,8] parameter(1)
+      rs = (f32[96,256], f32[96,8]) reduce-scatter(p0, p1),
+        replica_groups={{0,1}}, dimensions={0}, to_apply=add
+      gte0 = f32[96,256] get-tuple-element(rs), index=0
+      gte1 = f32[96,8] get-tuple-element(rs), index=1
+      ROOT t = (f32[96,256], f32[96,8]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      p0 = f32[192,256] parameter(0)
+      p1 = f32[192,8] parameter(1)
+      ROOT fusion = (f32[96,256], f32[96,8]) fusion(p0, p1), kind=kCustom,
+        calls=f
+    }
+  )");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TilingSpace> tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                          &mlir_context_));
+  const HloInstruction* rs =
+      root->fused_instructions_computation()->GetInstructionWithName("rs");
+  // Output 0: operand tile 32 * 2 = 64 rows; block 1 reads rows [64, 128),
+  // crossing the shard boundary at 96. Output 1: operand tile 48 * 2 = 96
+  // rows is exactly one shard and is fine.
+  ASSERT_OK(tiling_space->AssignTileSizes({32, 256, 48, 8}));
+  EXPECT_THAT(
+      PropagateTileToInput(*tiling_space, *rs, tiling_space->tiled_roots()[0],
+                           /*output_index=*/0),
+      StatusIs(absl::StatusCode::kFailedPrecondition));
+  EXPECT_OK(PropagateTileToInput(*tiling_space, *rs,
+                                 tiling_space->tiled_roots()[1],
+                                 /*output_index=*/1));
+}
+
 TEST_F(TilePropagationTest, CanPropagateToInputOfBroadcastOp) {
   HloInstruction* root = ParseAndGetRoot(R"(
     HloModule m
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc
index 372e0c9243a..b6dff55cfc0 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc
@@ -42,6 +42,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/parser/hlo_parser.h"
+#include "xla/hlo/testlib/filecheck.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/hlo/testlib/verified_hlo_module.h"
 #include "xla/hlo/utils/hlo_traversal.h"
@@ -51,6 +52,7 @@ limitations under the License.
 namespace xla::gpu::experimental {
 namespace {
 
+using ::absl_testing::IsOkAndHolds;
 using ::absl_testing::StatusIs;
 using ::mlir::MLIRContext;
 using ::testing::Contains;
@@ -1437,6 +1439,174 @@ TEST_F(IndexWiseVariadicTileAnalysisTest,
                HasSubstr("Output 1 of ag is not consumed in the fusion.")));
 }
 
+TEST_F(IndexWiseVariadicTileAnalysisTest,
+       VariadicReduceScatterIsTiledAsOneNodeWithPerOutputTiles) {
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"hlo(
+    add {
+      lhs = f32[] parameter(0)
+      rhs = f32[] parameter(1)
+      ROOT add = f32[] add(lhs, rhs)
+    }
+
+    f {
+      p0 = f32[4,64] parameter(0)
+      p1 = f32[4,1] parameter(1)
+      rs = (f32[2,64], f32[2,1]) reduce-scatter(p0, p1),
+        replica_groups={{0,1}}, dimensions={0}, to_apply=add
+      gte0 = f32[2,64] get-tuple-element(rs), index=0
+      gte1 = f32[2,1] get-tuple-element(rs), index=1
+      ROOT t = (f32[2,64], f32[2,1]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      arg0 = f32[4,64] parameter(0)
+      arg1 = f32[4,1] parameter(1)
+      ROOT fusion = (f32[2,64], f32[2,1]) fusion(arg0, arg1),
+        kind=kCustom, calls=f
+    })hlo",
+                                    {1, 64, 1, 1}));
+
+  EXPECT_THAT(RunFileCheck(tiled_computation.ToString(), R"(
+    CHECK: Dimensions:
+    CHECK:   4 type: replica size: 2 tile size: 1 dim ID:4
+    CHECK:   5 type: replica size: 2 tile size: 1 dim ID:5
+    CHECK: Root tiles:
+    CHECK:   0 root tile:  offsets [tid_4, 0] sizes [1, 64] strides [1, 1] upper bounds [2, 64]
+    CHECK:   1 root tile:  offsets [tid_5, 0] sizes [1, 1] strides [1, 1] upper bounds [2, 1]
+    CHECK: Tiled HLO:
+    CHECK:   arg0.tile_0 = parameter(0)  offsets [tid_4 * 2, 0] sizes [2, 64] strides [1, 1] upper bounds [4, 64]
+    CHECK:   arg1.tile_0 = parameter(1)  offsets [tid_5 * 2, 0] sizes [2, 1] strides [1, 1] upper bounds [4, 1]
+    CHECK:   rs.tile_0 = reduce-scatter(arg0.tile_0, arg1.tile_0)
+    CHECK:     #0  offsets [tid_4, 0] sizes [1, 64] strides [1, 1] upper bounds [2, 64]
+    CHECK:     #1  offsets [tid_5, 0] sizes [1, 1] strides [1, 1] upper bounds [2, 1]
+    CHECK:   gte0.tile_0 = get-tuple-element(rs.tile_0)  offsets [tid_4, 0] sizes [1, 64]
+    CHECK:   gte1.tile_0 = get-tuple-element(rs.tile_0)  offsets [tid_5, 0] sizes [1, 1]
+  )"),
+              IsOkAndHolds(true));
+
+  std::vector<const TiledHloInstruction*> reduce_scatters =
+      FindAll(tiled_computation, HloOpcode::kReduceScatter);
+  ASSERT_EQ(reduce_scatters.size(), 1);
+  const TiledHloInstruction* rs = reduce_scatters[0];
+  ASSERT_EQ(rs->operands().size(), 2);
+  ASSERT_EQ(rs->tiles().size(), 2);
+  EXPECT_EQ(rs->operand(0)->hlo()->name(), "arg0");
+  EXPECT_EQ(rs->operand(1)->hlo()->name(), "arg1");
+  EXPECT_EQ(rs->operand(0)->tile().num_replica_ids(), 0);
+  EXPECT_EQ(rs->operand(1)->tile().num_replica_ids(), 0);
+
+  // One GTE root per output, all backed by the same tiled reduce-scatter.
+  ASSERT_EQ(tiled_computation.roots().size(), 2);
+  EXPECT_EQ(tiled_computation.roots()[0]->operand(0), rs);
+  EXPECT_EQ(tiled_computation.roots()[1]->operand(0), rs);
+  EXPECT_EQ(rs->tile(0), tiled_computation.roots()[0]->tile());
+  EXPECT_EQ(rs->tile(1), tiled_computation.roots()[1]->tile());
+}
+
+TEST_F(IndexWiseVariadicTileAnalysisTest,
+       VariadicReduceScatterConsumedByGetTupleElementAndConvert) {
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"hlo(
+    add {
+      lhs = f32[] parameter(0)
+      rhs = f32[] parameter(1)
+      ROOT add = f32[] add(lhs, rhs)
+    }
+
+    f {
+      p0 = f32[16,16] parameter(0)
+      p1 = f32[16] parameter(1)
+      rs = (f32[4,16], f32[4]) reduce-scatter(p0, p1),
+        replica_groups={{0,1,2,3}}, dimensions={0}, to_apply=add
+      gte0 = f32[4,16] get-tuple-element(rs), index=0
+      gte1 = f32[4] get-tuple-element(rs), index=1
+      c0 = bf16[4,16] convert(gte0)
+      c1 = s32[4] convert(gte1)
+      ROOT t = (bf16[4,16], s32[4]) tuple(c0, c1)
+    }
+
+    ENTRY e {
+      arg0 = f32[16,16] parameter(0)
+      arg1 = f32[16] parameter(1)
+      ROOT fusion = (bf16[4,16], s32[4]) fusion(arg0, arg1),
+        kind=kCustom, calls=f
+    })hlo",
+                                    {1, 8, 1}));
+
+  std::vector<const TiledHloInstruction*> reduce_scatters =
+      FindAll(tiled_computation, HloOpcode::kReduceScatter);
+  ASSERT_EQ(reduce_scatters.size(), 1);
+  const TiledHloInstruction* rs = reduce_scatters[0];
+  ASSERT_EQ(rs->operands().size(), 2);
+  ASSERT_EQ(rs->tiles().size(), 2);
+
+  std::vector<const TiledHloInstruction*> gtes =
+      FindAll(tiled_computation, HloOpcode::kGetTupleElement);
+  ASSERT_EQ(gtes.size(), 2);
+  for (const TiledHloInstruction* gte : gtes) {
+    ASSERT_EQ(gte->operands().size(), 1);
+    EXPECT_EQ(gte->operand(0), rs);
+  }
+}
+
+TEST_F(IndexWiseVariadicTileAnalysisTest,
+       VariadicReduceScatterWithDotProducers) {
+  ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
+                       ParseAndTile(R"hlo(
+    add {
+      lhs = f32[] parameter(0)
+      rhs = f32[] parameter(1)
+      ROOT add = f32[] add(lhs, rhs)
+    }
+
+    f {
+      lhs0 = f32[16,8] parameter(0)
+      rhs0 = f32[8,16] parameter(1)
+      lhs1 = f32[32,8] parameter(2)
+      rhs1 = f32[8,32] parameter(3)
+      dot0 = f32[16,16] dot(lhs0, rhs0),
+        lhs_contracting_dims={1}, rhs_contracting_dims={0}
+      dot1 = f32[32,32] dot(lhs1, rhs1),
+        lhs_contracting_dims={1}, rhs_contracting_dims={0}
+      rs = (f32[4,16], f32[8,32]) reduce-scatter(dot0, dot1),
+        replica_groups={{0,1,2,3}}, dimensions={0}, to_apply=add
+      gte0 = f32[4,16] get-tuple-element(rs), index=0
+      gte1 = f32[8,32] get-tuple-element(rs), index=1
+      ROOT t = (f32[4,16], f32[8,32]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      lhs0 = f32[16,8] parameter(0)
+      rhs0 = f32[8,16] parameter(1)
+      lhs1 = f32[32,8] parameter(2)
+      rhs1 = f32[8,32] parameter(3)
+      ROOT fusion = (f32[4,16], f32[8,32]) fusion(lhs0, rhs0, lhs1, rhs1),
+        kind=kCustom, calls=f
+    })hlo",
+                                    {1, 8, 2, 16, 8, 8}));
+
+  EXPECT_THAT(RunFileCheck(tiled_computation.ToString(), R"(
+    CHECK: Tiled HLO:
+    CHECK:   dot0.tile_0 = dot(lhs0.1.tile_0, rhs0.1.tile_0)  offsets [tid_4 * 4, tid_1 * 8] sizes [4, 8]
+    CHECK:   region #0 {
+    CHECK:     lhs0.1.tile_0 = parameter(0)  offsets [tid_4 * 4, 0] sizes [4, 8]
+    CHECK:     rhs0.1.tile_0 = parameter(1)  offsets [0, tid_1 * 8] sizes [8, 8]
+    CHECK:   }
+    CHECK:   dot1.tile_0 = dot(lhs1.1.tile_0, rhs1.1.tile_0)  offsets [tid_5 * 8, tid_3 * 16] sizes [8, 16]
+    CHECK:   region #0 {
+    CHECK:     lhs1.1.tile_0 = parameter(2)  offsets [tid_5 * 8, 0] sizes [8, 8]
+    CHECK:     rhs1.1.tile_0 = parameter(3)  offsets [0, tid_3 * 16] sizes [8, 16]
+    CHECK:   }
+    CHECK:   rs.tile_0 = reduce-scatter(dot0.tile_0, dot1.tile_0)
+    CHECK:     #0  offsets [tid_4, tid_1 * 8] sizes [1, 8]
+    CHECK:     #1  offsets [tid_5 * 2, tid_3 * 16] sizes [2, 16]
+    CHECK:   gte0.tile_0 = get-tuple-element(rs.tile_0)  offsets [tid_4, tid_1 * 8] sizes [1, 8]
+    CHECK:   gte1.tile_0 = get-tuple-element(rs.tile_0)  offsets [tid_5 * 2, tid_3 * 16] sizes [2, 16]
+  )"),
+              IsOkAndHolds(true));
+}
+
 TEST_F(SameShapeMultiOutputFusionTileAnalysisTest,
        VariadicReduceIsTiledAsOneNode) {
   // A variadic reduce is not index-wise: every output depends on every input.
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc
index a82a13aa60e..d92f51eda38 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc
@@ -73,13 +73,13 @@ bool HasIndexWiseVariadic(const HloFusionAdaptor& fusion) {
                         });
 }
 
-// Returns the position of the first dimension of result `result_index` in the
-// ordered list of dimensions of `hlo`. For an index-wise variadic instruction
-// the dimensions of all results are laid out back to back: result k starts at
-// the sum of the ranks of results 0..k-1. For every other instruction the
-// results share the dimensions of the first one and the offset is 0.
+}  // namespace
+
 int64_t GetResultDimensionOffset(const HloInstruction& hlo,
                                  int64_t result_index) {
+  if (!hlo.shape().IsTuple()) {
+    return result_index * hlo.shape().dimensions().size();
+  }
   if (!IsIndexWiseVariadic(hlo)) {
     return 0;
   }
@@ -90,10 +90,10 @@ int64_t GetResultDimensionOffset(const HloInstruction& hlo,
   return offset;
 }
 
-}  // namespace
-
 bool IsIndexWiseVariadic(const HloInstruction& hlo) {
-  return hlo.opcode() == HloOpcode::kAllGather && hlo.operand_count() > 1;
+  return HloPredicateIsOp<HloOpcode::kAllGather, HloOpcode::kReduceScatter>(
+             &hlo) &&
+         hlo.operand_count() > 1;
 }
 
 llvm::DenseMap<SymbolicExpr, SymbolicExpr> GetTileSizeReplacementMap(
@@ -224,27 +224,29 @@ void TilingSpace::ProcessReduce(const HloInstruction& hlo) {
   }
 }
 
-// Adds a `kReplica` dimension for `reduce-scatter` and links it to the scatter
-// dimension.
+// Adds a `kReplica` dimension for each `reduce-scatter` output and links it to
+// the output's scatter dimension.
 void TilingSpace::ProcessReduceScatter(const HloInstruction& hlo) {
-  if (hlo.shape().IsTuple()) {
-    return;
-  }
   const auto* reduce_scatter = Cast<HloReduceScatterInstruction>(&hlo);
   const int64_t scatter_dim = reduce_scatter->scatter_dimension();
-  auto it = hlo_to_dimension_.find(std::make_pair(&hlo, scatter_dim));
-  if (it == hlo_to_dimension_.end()) {
-    return;
-  }
-  const TiledDimId scatter_dim_id = it->second->id;
   const int64_t num_replicas =
       reduce_scatter->device_list()->num_devices_per_group();
   if (num_replicas <= 1) {
     return;
   }
-  AppendDimension(&hlo, hlo.shape().dimensions().size(), num_replicas,
-                  DimensionSemantics::kReplica);
-  dimensions_[scatter_dim_id.value()].replica_dim = dimensions_.back().id;
+  const int64_t num_results = reduce_scatter->operand_count();
+  const int64_t total_result_dims = GetResultDimensionOffset(hlo, num_results);
+  for (int64_t result_index = 0; result_index < num_results; ++result_index) {
+    auto it = hlo_to_dimension_.find(std::make_pair(
+        &hlo, GetResultDimensionOffset(hlo, result_index) + scatter_dim));
+    if (it == hlo_to_dimension_.end()) {
+      continue;
+    }
+    const TiledDimId scatter_dim_id = it->second->id;
+    AppendDimension(&hlo, total_result_dims + result_index, num_replicas,
+                    DimensionSemantics::kReplica);
+    dimensions_[scatter_dim_id.value()].replica_dim = dimensions_.back().id;
+  }
 }
 
 // Register sequential dimensions and RTVars for a kRaggedDot instruction.
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h
index acd57f97058..b5c698edd9b 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h
@@ -317,9 +317,16 @@ const Shape& GetFirstShape(const HloInstruction* instr, int64_t index = 0);
 // Returns true if `hlo` is a tuple-producing instruction whose output `k`
 // depends only on operand `k`, with the same tile mapping for every `k`.
 //
-// Currently only multi-operand all-gather qualifies.
+// Currently multi-operand all-gather and reduce-scatter qualify.
 bool IsIndexWiseVariadic(const HloInstruction& hlo);
 
+// Returns the position of the first dimension of result `result_index` in the
+// ordered list of dimensions of `hlo`. For an index-wise variadic instruction
+// the dimensions of all results are laid out back to back: result k starts at
+// the sum of the ranks of results 0..k-1.
+int64_t GetResultDimensionOffset(const HloInstruction& hlo,
+                                 int64_t result_index);
+
 // Returns a symbol replacement map to set concrete tile sizes.
 llvm::DenseMap<SymbolicExpr, SymbolicExpr> GetTileSizeReplacementMap(
     const TilingSpace& tiling_space);
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc
index bfc37ddbaad..abc574bfe36 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc
@@ -36,6 +36,7 @@ limitations under the License.
 #include "xla/hlo/analysis/symbolic_map_serialization.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/testlib/filecheck.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/hlo/testlib/verified_hlo_module.h"
 #include "xla/hlo/utils/hlo_traversal.h"
@@ -44,6 +45,7 @@ limitations under the License.
 namespace xla::gpu::experimental {
 namespace {
 
+using ::absl_testing::IsOkAndHolds;
 using ::absl_testing::StatusIs;
 using ::mlir::MLIRContext;
 using ::testing::ElementsAre;
@@ -373,6 +375,11 @@ TEST_F(TilingSpaceTest, TwoOutputsEqualShapesParallelDims) {
 TEST_F(TilingSpaceTest, IsIndexWiseVariadic) {
   ParseAndGetRoot(R"(
     HloModule m
+    add {
+      x = f32[] parameter(0)
+      y = f32[] parameter(1)
+      ROOT add = f32[] add(x, y)
+    }
     add_pair {
       a = f32[] parameter(0)
       b = s32[] parameter(1)
@@ -387,16 +394,23 @@ TEST_F(TilingSpaceTest, IsIndexWiseVariadic) {
       p1 = s32[1] parameter(1)
       p2 = f32[4,8] parameter(2)
       p3 = s32[4,8] parameter(3)
+      p4 = f32[4,16] parameter(4)
       c0 = f32[] constant(0)
       c1 = s32[] constant(0)
       single_ag = f32[4,8] all-gather(p0), replica_groups={{0,1,2,3}},
         dimensions={0}
       variadic_ag = (f32[4,8], s32[4]) all-gather(p0, p1),
         replica_groups={{0,1,2,3}}, dimensions={0}
+      single_rs = f32[1,8] reduce-scatter(p2), replica_groups={{0,1,2,3}},
+        dimensions={0}, to_apply=add
+      variadic_rs = (f32[1,8], f32[1,16]) reduce-scatter(p2, p4),
+        replica_groups={{0,1,2,3}}, dimensions={0}, to_apply=add
       variadic_reduce = (f32[4], s32[4]) reduce(p2, p3, c0, c1),
         dimensions={1}, to_apply=add_pair
       gte = f32[4,8] get-tuple-element(variadic_ag), index=0
-      ROOT t = (f32[4,8], f32[4,8]) tuple(single_ag, gte)
+      gte_rs = f32[1,8] get-tuple-element(variadic_rs), index=0
+      ROOT t = (f32[4,8], f32[4,8], f32[1,8], f32[1,8])
+        tuple(single_ag, gte, single_rs, gte_rs)
     }
   )");
   HloComputation* entry = module_->entry_computation();
@@ -404,9 +418,14 @@ TEST_F(TilingSpaceTest, IsIndexWiseVariadic) {
       IsIndexWiseVariadic(*entry->GetInstructionWithName("variadic_ag")));
   EXPECT_FALSE(
       IsIndexWiseVariadic(*entry->GetInstructionWithName("single_ag")));
+  EXPECT_TRUE(
+      IsIndexWiseVariadic(*entry->GetInstructionWithName("variadic_rs")));
+  EXPECT_FALSE(
+      IsIndexWiseVariadic(*entry->GetInstructionWithName("single_rs")));
   EXPECT_FALSE(
       IsIndexWiseVariadic(*entry->GetInstructionWithName("variadic_reduce")));
   EXPECT_FALSE(IsIndexWiseVariadic(*entry->GetInstructionWithName("gte")));
+  EXPECT_FALSE(IsIndexWiseVariadic(*entry->GetInstructionWithName("gte_rs")));
   EXPECT_FALSE(IsIndexWiseVariadic(*entry->root_instruction()));
 }
 
@@ -569,6 +588,52 @@ TEST_F(TilingSpaceTest,
   EXPECT_EQ(tiling_space->num_parallel_dimensions(), 4);
 }
 
+TEST_F(TilingSpaceTest,
+       VariadicReduceScatterRootGetsReplicaDimensionsAndReplicaSplitRootTiles) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    add {
+      lhs = f32[] parameter(0)
+      rhs = f32[] parameter(1)
+      ROOT add = f32[] add(lhs, rhs)
+    }
+    f {
+      p0 = f32[32,16] parameter(0)
+      p1 = f32[16,2] parameter(1)
+      rs = (f32[8,16], f32[4,2]) reduce-scatter(p0, p1),
+        replica_groups={{0,1,2,3}}, dimensions={0}, to_apply=add
+      gte0 = f32[8,16] get-tuple-element(rs), index=0
+      gte1 = f32[4,2] get-tuple-element(rs), index=1
+      ROOT t = (f32[8,16], f32[4,2]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      p0 = f32[32,16] parameter(0)
+      p1 = f32[16,2] parameter(1)
+      ROOT fusion = (f32[8,16], f32[4,2]) fusion(p0, p1), kind=kCustom,
+        calls=f
+    }
+  )");
+  auto fusion_adaptor = HloFusionAdaptor::ForInstruction(root);
+  ASSERT_OK_AND_ASSIGN(auto tiling_space,
+                       TilingSpace::Create(*fusion_adaptor, &mlir_context_));
+  EXPECT_EQ(tiling_space->num_dimensions(), 6);
+  EXPECT_EQ(tiling_space->num_parallel_dimensions(), 4);
+  EXPECT_THAT(RunFileCheck(tiling_space->ToString(), R"(
+    CHECK: Dimensions:
+    CHECK:   0 type: parallel size: 8 dim ID:0
+    CHECK:   1 type: parallel size: 16 dim ID:1
+    CHECK:   2 type: parallel size: 4 dim ID:0
+    CHECK:   3 type: parallel size: 2 dim ID:1
+    CHECK:   4 type: replica size: 4 dim ID:4
+    CHECK:   5 type: replica size: 4 dim ID:5
+    CHECK: Root tiles:
+    CHECK:   0 root tile:  offsets [tid_0 * ts_0 * 4 + tid_4 * ts_0, tid_1 * ts_1] sizes [ts_0, ts_1] strides [1, 1] upper bounds [8, 16]
+    CHECK:   1 root tile:  offsets [tid_2 * ts_2 * 4 + tid_5 * ts_2, tid_3 * ts_3] sizes [ts_2, ts_3] strides [1, 1] upper bounds [4, 2]
+  )"),
+              IsOkAndHolds(true));
+}
+
 TEST_F(TilingSpaceTest,
        ReduceScatterRootGetsReplicaDimensionAndReplicaSplitRootTile) {
   HloInstruction* root = ParseAndGetRoot(R"(

From 23174295879992240ffe06f0086c52515a6a1c12 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2026 09:07:32 -0700
Subject: [PATCH 0787/1000] [XLA:GPU] Only pass
 CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY when graph nodes set
 non-default priority

`GraphInstantiate` in `cuda_command_buffer.cc` unconditionally passed `CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY` to `cuGraphInstantiate` (cl/766117515), which causes `cuGraphLaunch` to ignore the launch stream's priority and use each node's `CU_LAUNCH_ATTRIBUTE_PRIORITY` attribute instead. Because standard XLA `Command` nodes and `command_buffer_trace_stream` default to `StreamPriority::Default` (0), all CUDA Graph kernels were silently demoted to priority 0 even when launched on a high-priority stream, causing multi-stream GPU workloads to suffer cross-stream preemption at graph kernel boundaries.

Fix this by tracking `has_node_priority_` on `CudaCommandBuffer` (propagating from nested/child command buffers) and only passing `CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY` when at least one node in the graph hierarchy explicitly sets a non-default priority.

PiperOrigin-RevId: 999410431
---
 .../xla/xla/stream_executor/cuda/BUILD        |   1 +
 .../cuda/cuda_command_buffer.cc               |  36 ++++--
 .../cuda/cuda_command_buffer.h                |   5 +
 .../cuda/cuda_command_buffer_12_9.cc          |   3 +
 .../cuda/cuda_command_buffer_test.cc          | 114 ++++++++++++++++++
 5 files changed, 152 insertions(+), 7 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/cuda/BUILD b/third_party/xla/xla/stream_executor/cuda/BUILD
index cf50493f7e4..9131900c9e7 100644
--- a/third_party/xla/xla/stream_executor/cuda/BUILD
+++ b/third_party/xla/xla/stream_executor/cuda/BUILD
@@ -2124,6 +2124,7 @@ xla_test(
     backends = ["gpu"],
     tags = ["cuda-only"],
     deps = [
+        ":cuda_command_buffer",
         ":cuda_compute_capability",
         ":cudnn_plugin",
         "//xla/service:platform_util",
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc
index 1a260a57386..05ca850f395 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.cc
@@ -269,9 +269,12 @@ std::string ConditionalTypeToString(GpuCommandBuffer::ConditionType type) {
   }
 }
 
-absl::Status GraphInstantiate(CUgraphExec* exec, CUgraph graph) {
-  VLOG(2) << "Instantiate CUDA executable graph from graph " << graph;
-  uint64_t cu_flags = CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY;
+absl::Status GraphInstantiate(CUgraphExec* exec, CUgraph graph,
+                              bool use_node_priority) {
+  VLOG(2) << "Instantiate CUDA executable graph from graph " << graph
+          << "; use_node_priority=" << use_node_priority;
+  uint64_t cu_flags =
+      use_node_priority ? CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY : 0;
   return cuda::ToStatus(cuGraphInstantiate(exec, graph, cu_flags),
                         "Failed to instantiate CUDA graph");
 }
@@ -686,6 +689,9 @@ absl::StatusOr<GraphNodeHandle> CudaCommandBuffer::CreateClonedChildNode(
       absl::down_cast<const CudaCommandBuffer&>(nested);
   CHECK_EQ(child_command_buffer.parent_, nullptr)
       << "Nested command buffer's parent is not null";
+  if (child_command_buffer.has_node_priority_) {
+    SetHasNodePriority();
+  }
 
   CUgraph child_graph = child_command_buffer.graph_;
   std::vector<CUgraphNode> deps = ToCudaGraphHandles(dependencies);
@@ -716,8 +722,12 @@ absl::StatusOr<GraphNodeHandle> CudaCommandBuffer::CreateMovedChildNode(
 
 absl::Status CudaCommandBuffer::UpdateClonedChildNode(
     GraphNodeHandle node_handle, const CommandBuffer& nested) {
-  CUgraph child_graph =
-      absl::down_cast<const CudaCommandBuffer&>(nested).graph_;
+  const auto& child_command_buffer =
+      absl::down_cast<const CudaCommandBuffer&>(nested);
+  if (child_command_buffer.has_node_priority_) {
+    SetHasNodePriority();
+  }
+  CUgraph child_graph = child_command_buffer.graph_;
   VLOG(2) << "Set child node params " << node_handle << " in graph executable "
           << graph_exec() << " to params contained in " << child_graph;
 
@@ -850,6 +860,7 @@ absl::StatusOr<GraphNodeHandle> CudaCommandBuffer::CreateKernelNode(
       "Failed to add kernel node to a CUDA graph"));
 
   if (priority != StreamPriority::Default) {
+    SetHasNodePriority();
     CUlaunchAttributeValue value;
     value.priority = stream_exec_->GetGpuStreamPriority(priority);
     ABSL_RETURN_IF_ERROR(
@@ -1019,6 +1030,9 @@ absl::StatusOr<size_t> CudaCommandBuffer::GetNodeCount() const {
 }
 
 absl::Status CudaCommandBuffer::SetPriority(StreamPriority priority) {
+  if (priority != StreamPriority::Default) {
+    SetHasNodePriority();
+  }
   size_t num_nodes;
   ABSL_RETURN_IF_ERROR(
       cuda::ToStatus(cuGraphGetNodes(graph_, /*nodes=*/nullptr, &num_nodes)));
@@ -1045,6 +1059,13 @@ absl::Status CudaCommandBuffer::SetPriority(StreamPriority priority) {
   return absl::OkStatus();
 }
 
+void CudaCommandBuffer::SetHasNodePriority() {
+  for (CudaCommandBuffer* current = this; current != nullptr;
+       current = current->parent_) {
+    current->has_node_priority_ = true;
+  }
+}
+
 absl::Status CudaCommandBuffer::PrepareFinalization() {
   if (stream_exec_->GetDeviceDescription().driver_version() <
       SemanticVersion{12, 8, 0}) {
@@ -1092,7 +1113,8 @@ absl::Status CudaCommandBuffer::InstantiateGraph() {
   VLOG(2) << "CUDA Graph DOT:\n" << ToString();
   // If we get a "resource exhausted error" we retry instantiating Gpu graph
   // one more time after releasing unused device memory allocated for graphs.
-  auto instantiated = GraphInstantiate(&graph_exec_, graph_);
+  auto instantiated =
+      GraphInstantiate(&graph_exec_, graph_, has_node_priority_);
   if (instantiated.code() == absl::StatusCode::kResourceExhausted) {
     LOG(WARNING) << "Retry CUDA graph instantiation after OOM error";
     CUdevice device;
@@ -1101,7 +1123,7 @@ absl::Status CudaCommandBuffer::InstantiateGraph() {
                        "Failed call to cuDeviceGet"));
     ABSL_RETURN_IF_ERROR(cuda::ToStatus(cuDeviceGraphMemTrim(device),
                                    "Failed to trim device graph memory"));
-    ABSL_RETURN_IF_ERROR(GraphInstantiate(&graph_exec_, graph_));
+    ABSL_RETURN_IF_ERROR(GraphInstantiate(&graph_exec_, graph_, has_node_priority_));
   } else {
     ABSL_RETURN_IF_ERROR(instantiated);
   }
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h
index 48e4ced1630..8fb52da1394 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer.h
@@ -55,6 +55,8 @@ class CudaCommandBuffer final : public GpuCommandBuffer {
 
   std::string ToString() const override;
 
+  bool has_node_priority() const { return has_node_priority_; }
+
   ~CudaCommandBuffer() override;
 
  private:
@@ -220,6 +222,8 @@ class CudaCommandBuffer final : public GpuCommandBuffer {
 
   CUgraphExec graph_exec() const;
 
+  void SetHasNodePriority();
+
   absl::Status CheckCanBeUpdated() override;
 
   // Static helpers that are needed across different versions of CUDA.
@@ -266,6 +270,7 @@ class CudaCommandBuffer final : public GpuCommandBuffer {
 
   CUgraph graph_ = nullptr;
   bool is_owned_graph_ = true;
+  bool has_node_priority_ = false;
   CUgraphExec graph_exec_ = nullptr;
 };
 
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer_12_9.cc b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer_12_9.cc
index a1cd57a7610..b8dbff4caed 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer_12_9.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer_12_9.cc
@@ -74,6 +74,9 @@ absl::StatusOr<GraphNodeHandle> CudaCommandBuffer::CreateMovedChildNodeImpl(
 
   child_command_buffer->parent_ = this;
   child_command_buffer->is_owned_graph_ = false;
+  if (child_command_buffer->has_node_priority_) {
+    SetHasNodePriority();
+  }
 
   return FromCudaGraphHandle(node_handle);
 }
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer_test.cc b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer_test.cc
index 7696a0d2729..13254965997 100644
--- a/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer_test.cc
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_command_buffer_test.cc
@@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include "xla/stream_executor/cuda/cuda_command_buffer.h"
+
 #include <array>
 #include <cstddef>
 #include <cstdint>
@@ -394,5 +396,117 @@ TEST_F(CudaCommandBufferTest, MemcpyH2D2H) {
   EXPECT_EQ(dst2, 456);
   EXPECT_EQ(dst3, 789);
 }
+
+TEST_F(CudaCommandBufferTest,
+       GraphInstantiateUsesLaunchStreamPriorityByDefault) {
+  Platform* platform = CudaPlatform();
+  ASSERT_OK_AND_ASSIGN(StreamExecutor * executor,
+                       platform->ExecutorForDevice(0));
+  ASSERT_OK_AND_ASSIGN(
+      KernelLoaderSpec add_spec,
+      ::stream_executor::gpu::GetAddI32TestKernelSpec(platform->id()));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<Kernel> kernel,
+                       executor->LoadKernel(add_spec));
+
+  DeviceAddress<int32_t> a = executor->AllocateArray<int32_t>(4);
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<CommandBuffer> cmd_buffer,
+                       executor->CreateCommandBuffer(primary));
+  ASSERT_OK_AND_ASSIGN(
+      const CommandBuffer::Command* cmd,
+      cmd_buffer->CreateLaunch(ThreadDim(4, 1, 1), BlockDim(1, 1, 1),
+                               std::nullopt, *kernel,
+                               *stream_executor::PackKernelArgs(0, a, a, a), {},
+                               StreamPriority::Default));
+  ASSERT_NE(cmd, nullptr);
+  ASSERT_OK(cmd_buffer->Finalize());
+
+  auto* cuda_cmd_buffer =
+      dynamic_cast<gpu::CudaCommandBuffer*>(cmd_buffer.get());
+  ASSERT_NE(cuda_cmd_buffer, nullptr);
+  EXPECT_FALSE(cuda_cmd_buffer->has_node_priority());
+
+  auto* gpu_cmd = dynamic_cast<const gpu::GpuCommandBuffer::GpuCommand*>(cmd);
+  ASSERT_NE(gpu_cmd, nullptr);
+  CUlaunchAttributeValue value{};
+  ASSERT_EQ(cuGraphKernelNodeGetAttribute(
+                absl::bit_cast<CUgraphNode>(gpu_cmd->handle),
+                CU_LAUNCH_ATTRIBUTE_PRIORITY, &value),
+            CUDA_SUCCESS);
+  EXPECT_EQ(value.priority, 0);
+}
+
+TEST_F(CudaCommandBufferTest,
+       GraphInstantiateUsesNodePriorityWhenExplicitlySet) {
+  Platform* platform = CudaPlatform();
+  ASSERT_OK_AND_ASSIGN(StreamExecutor * executor,
+                       platform->ExecutorForDevice(0));
+  ASSERT_OK_AND_ASSIGN(
+      KernelLoaderSpec add_spec,
+      ::stream_executor::gpu::GetAddI32TestKernelSpec(platform->id()));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<Kernel> kernel,
+                       executor->LoadKernel(add_spec));
+
+  DeviceAddress<int32_t> a = executor->AllocateArray<int32_t>(4);
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<CommandBuffer> cmd_buffer,
+                       executor->CreateCommandBuffer(primary));
+  ASSERT_OK_AND_ASSIGN(
+      const CommandBuffer::Command* cmd,
+      cmd_buffer->CreateLaunch(ThreadDim(4, 1, 1), BlockDim(1, 1, 1),
+                               std::nullopt, *kernel,
+                               *stream_executor::PackKernelArgs(0, a, a, a), {},
+                               StreamPriority::Highest));
+  ASSERT_NE(cmd, nullptr);
+  ASSERT_OK(cmd_buffer->Finalize());
+
+  auto* cuda_cmd_buffer =
+      dynamic_cast<gpu::CudaCommandBuffer*>(cmd_buffer.get());
+  ASSERT_NE(cuda_cmd_buffer, nullptr);
+  EXPECT_TRUE(cuda_cmd_buffer->has_node_priority());
+
+  auto* gpu_cmd = dynamic_cast<const gpu::GpuCommandBuffer::GpuCommand*>(cmd);
+  ASSERT_NE(gpu_cmd, nullptr);
+  CUlaunchAttributeValue value{};
+  ASSERT_EQ(cuGraphKernelNodeGetAttribute(
+                absl::bit_cast<CUgraphNode>(gpu_cmd->handle),
+                CU_LAUNCH_ATTRIBUTE_PRIORITY, &value),
+            CUDA_SUCCESS);
+  EXPECT_EQ(value.priority,
+            executor->GetGpuStreamPriority(StreamPriority::Highest));
+}
+
+TEST_F(CudaCommandBufferTest,
+       NestedChildCommandPropagatesNodePriorityToParent) {
+  Platform* platform = CudaPlatform();
+  ASSERT_OK_AND_ASSIGN(StreamExecutor * executor,
+                       platform->ExecutorForDevice(0));
+  ASSERT_OK_AND_ASSIGN(
+      KernelLoaderSpec add_spec,
+      ::stream_executor::gpu::GetAddI32TestKernelSpec(platform->id()));
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<Kernel> kernel,
+                       executor->LoadKernel(add_spec));
+
+  DeviceAddress<int32_t> a = executor->AllocateArray<int32_t>(4);
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<CommandBuffer> nested_cmd,
+      executor->CreateCommandBuffer(CommandBuffer::Mode::kNested));
+  ASSERT_OK(nested_cmd
+                ->CreateLaunch(ThreadDim(4, 1, 1), BlockDim(1, 1, 1),
+                               std::nullopt, *kernel,
+                               *stream_executor::PackKernelArgs(0, a, a, a), {},
+                               StreamPriority::Default)
+                .status());
+  ASSERT_OK(nested_cmd->SetPriority(StreamPriority::Highest));
+  ASSERT_OK(nested_cmd->Finalize());
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<CommandBuffer> cmd_buffer,
+                       executor->CreateCommandBuffer(primary));
+  ASSERT_OK(cmd_buffer->CreateChildCommand(*nested_cmd, {}).status());
+  ASSERT_OK(cmd_buffer->Finalize());
+
+  auto* cuda_cmd_buffer =
+      dynamic_cast<gpu::CudaCommandBuffer*>(cmd_buffer.get());
+  ASSERT_NE(cuda_cmd_buffer, nullptr);
+  EXPECT_TRUE(cuda_cmd_buffer->has_node_priority());
+}
 }  // namespace
 }  // namespace stream_executor::cuda

From 5ed35c4e9b6b7d166ee36fb39887369c4c7bb432 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 21:49:04 +0530
Subject: [PATCH 0788/1000] Fix IPv6 fallback regression and missing break
 statements

_parse_host_and_port's bracketed-IPv6 fallback unconditionally
retried any address that failed initial parsing, including a
genuinely invalid one like 'localhost:foo', which the fallback then
misread as a literal IPv6 host. Restrict the fallback to netlocs
with more than one colon (the actual signature of an unbracketed
IPv6 address), so invalid addresses still correctly return
(None, None).

Also add the two missing break statements flagged as the same bug
class as the original master-lookup fix: the collective_group_leader
loop now stops at the first non-local job instead of letting later
jobs silently overwrite it, and the current_server_def cleanup loop
stops after deleting the matching job instead of continuing to
iterate (and shift indices) over the just-mutated repeated field.

Added a '[::1]:8000' case to testParseHostAndPortSuccess to cover
unbracketed-IPv6-with-port parsing.

Verified: all parse cases (valid host:port, grpc:// prefix, bracketed
and unbracketed IPv6, and both invalid-address cases) checked against
real execution; RemoteUtilsTest (16 tests) and MultiJobsTest's
testConnectToClusterWithMissingMaster/WithInvalidMaster (using real
local gRPC servers) all pass against the patched module.
---
 tensorflow/python/eager/remote.py      | 20 ++++++++++++++------
 tensorflow/python/eager/remote_test.py |  1 +
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/tensorflow/python/eager/remote.py b/tensorflow/python/eager/remote.py
index 5ae8603ea8e..19cafcc1da8 100644
--- a/tensorflow/python/eager/remote.py
+++ b/tensorflow/python/eager/remote.py
@@ -93,13 +93,19 @@ def _parse_host_and_port(address):
     # `urlsplit` raises a `ValueError` when it tries to parse `.port` on an
     # unbracketed IPv6 address (e.g. `//::1`), since it mistakes the address's
     # own colons for a host:port separator. Retry with the host wrapped in
-    # brackets, which is the syntax `urlsplit` expects for IPv6 hosts.
+    # brackets, which is the syntax `urlsplit` expects for IPv6 hosts. Only do
+    # this when the netloc has more than one colon (a true unbracketed IPv6
+    # address like `::1`); otherwise an address that's simply unparseable
+    # (e.g. `localhost:foo`) would get wrapped in brackets and misread as a
+    # literal IPv6 host instead of correctly failing to parse.
     scheme, netloc = address.split("//", 1)
-    try:
-      parsed_bracketed = urllib.parse.urlsplit(f"{scheme}//[{netloc}]")
-      return parsed_bracketed.hostname, parsed_bracketed.port
-    except ValueError:
-      return None, None
+    if not netloc.startswith("[") and netloc.count(":") > 1:
+      try:
+        parsed_bracketed = urllib.parse.urlsplit(f"{scheme}//[{netloc}]")
+        return parsed_bracketed.hostname, parsed_bracketed.port
+      except ValueError:
+        pass
+    return None, None
 
 
 def _find_master_job_and_task(cluster_spec, master):
@@ -221,6 +227,7 @@ def connect_to_cluster(cluster_spec_or_resolver,
     for i, job in enumerate(current_server_def.cluster.job):
       if job.name == job_name:
         del current_server_def.cluster.job[i]
+        break
   if (current_server_def is None or current_server_def.cluster != cluster_def or
       current_server_def.job_name != job_name or
       current_server_def.task_index != task_index):
@@ -278,6 +285,7 @@ def connect_to_cluster(cluster_spec_or_resolver,
     default_session_config.experimental.collective_group_leader = (
         f"/job:{name}/replica:0/task:0"
     )
+    break
 
   logging.info("default session config: %s", default_session_config)
 
diff --git a/tensorflow/python/eager/remote_test.py b/tensorflow/python/eager/remote_test.py
index 494596e36b9..e0d9d6106cd 100644
--- a/tensorflow/python/eager/remote_test.py
+++ b/tensorflow/python/eager/remote_test.py
@@ -831,6 +831,7 @@ class RemoteUtilsTest(test.TestCase, parameterized.TestCase):
       ('localhost:8000', 'localhost', 8000),
       ('grpc://localhost:8000', 'localhost', 8000),
       ('grpc://[::1]:8000', '::1', 8000),
+      ('[::1]:8000', '::1', 8000),
       ('localhost', 'localhost', None),
       ('::1', '::1', None),
   )

From 06c8af8fa91d428b41522eb1b19820a2fea58889 Mon Sep 17 00:00:00 2001
From: bodapatisaikrishna <saikrishnabodapati@gmail.com>
Date: Wed, 14 Oct 2026 21:52:40 +0530
Subject: [PATCH 0789/1000] Add initialization check in AssignUpdateVariableOp
 and expand tests

---
 .../core/kernels/resource_variable_ops.cc     |  3 ++
 .../python/training/training_ops_test.py      | 42 +++++++++++++++----
 2 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc
index 8a1c33a8f88..bd57da5bd44 100644
--- a/tensorflow/core/kernels/resource_variable_ops.cc
+++ b/tensorflow/core/kernels/resource_variable_ops.cc
@@ -648,6 +648,9 @@ class AssignUpdateVariableOp : public OpKernel {
     // ADD if value's refcount was 1.
     mutex_lock ml(*variable->mu());
     Tensor* var_tensor = variable->tensor();
+    OP_REQUIRES(context, var_tensor->IsInitialized(),
+                errors::FailedPrecondition(
+                    "Attempting to use uninitialized variables: ", handle.name()));
     OP_REQUIRES_OK(context, ValidateAssignUpdateVariableOpShapes(
                                 var_tensor->shape(), value.shape()));
     OP_REQUIRES(context, var_tensor->dtype() == value.dtype(),
diff --git a/tensorflow/python/training/training_ops_test.py b/tensorflow/python/training/training_ops_test.py
index d8d98753979..12daccab7fb 100644
--- a/tensorflow/python/training/training_ops_test.py
+++ b/tensorflow/python/training/training_ops_test.py
@@ -699,7 +699,7 @@ class TrainingOpsTest(TensorFlowTestCase):
             )
         )
 
-  @test_util.run_v2_only
+  @test_util.run_in_graph_and_eager_modes
   def testResourceVariableDtypeMismatch(self):
     # A resource variable whose dtype differs from optimizer dtype T
     # previously caused a fatal CHECK failure in CheckTypeAndIsAligned
@@ -750,19 +750,28 @@ class TrainingOpsTest(TensorFlowTestCase):
           )
       )
 
-  @test_util.run_v2_only
+  @test_util.run_in_graph_and_eager_modes
   def testResourceVariableUninitialized(self):
-    shared_name = (
-        context.anonymous_name() if context.executing_eagerly() else ""
-    )
     var_uninit_f32 = gen_resource_variable_ops.var_handle_op(
-        dtype=dtypes.float32, shape=[1, 2], shared_name=shared_name
+        dtype=dtypes.float32,
+        shape=[1, 2],
+        shared_name=(
+            context.anonymous_name() if context.executing_eagerly() else ""
+        ),
     )
     var_uninit_f16 = gen_resource_variable_ops.var_handle_op(
-        dtype=dtypes.float16, shape=[1, 2], shared_name=shared_name
+        dtype=dtypes.float16,
+        shape=[1, 2],
+        shared_name=(
+            context.anonymous_name() if context.executing_eagerly() else ""
+        ),
     )
     accum_uninit_f16 = gen_resource_variable_ops.var_handle_op(
-        dtype=dtypes.float16, shape=[1, 2], shared_name=shared_name
+        dtype=dtypes.float16,
+        shape=[1, 2],
+        shared_name=(
+            context.anonymous_name() if context.executing_eagerly() else ""
+        ),
     )
     lr = constant_op.constant(0.1, dtype=dtypes.float16)
     grad = constant_op.constant([[0.1, 0.2]], dtype=dtypes.float16)
@@ -818,6 +827,23 @@ class TrainingOpsTest(TensorFlowTestCase):
           )
       )
 
+    var_uninit_scalar_f16 = gen_resource_variable_ops.var_handle_op(
+        dtype=dtypes.float16,
+        shape=[],
+        shared_name=(
+            context.anonymous_name() if context.executing_eagerly() else ""
+        ),
+    )
+    with self.assertRaisesRegex(
+        errors.FailedPreconditionError, "uninitialized"
+    ):
+      self.evaluate(
+          gen_resource_variable_ops.assign_add_variable_op(
+              resource=var_uninit_scalar_f16,
+              value=constant_op.constant(0.1, dtype=dtypes.float16),
+          )
+      )
+
 
 if __name__ == '__main__':
   googletest.main()

From 8ebbb6d1a8809812ee9e200a3679107af6a03d0d Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 21:52:52 +0530
Subject: [PATCH 0790/1000] Fix scalar tf.TensorShape([]) edge case in deferred
 DVariable restore

bool(shape) is not a reliable scalar check: an empty tuple () is
falsy, but tf.TensorShape([]) is truthy even though it also has
length 0, so the prior 'if shape else None' guard let a scalar
TensorShape fall through to the ShardInfo branch, reproducing the
same 'Non numerical dimension in shape_and_slice' InvalidArgumentError
this fix was meant to eliminate.

Switch the check to len(shape) > 0, which is true for both
representations of a scalar shape.

Added a second assertion to test_checkpoint_deferred_restore_scalar_dvariable
covering shape=tf.TensorShape([]) (via a fresh restore_ckpt/init_callable,
since re-using the same CheckpointInitialValueCallable would double-consume
the deferred restore).

Verified by running the old (len(shape)-less) code against this new test
case: it fails with exactly the predicted InvalidArgumentError. The fixed
code passes both the tuple and TensorShape scalar cases, plus an existing
non-scalar test_checkpoint_simple case.
---
 tensorflow/dtensor/python/d_variable.py         |  7 +++++--
 .../python/tests/save_restore_v2_test.py        | 17 +++++++++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/tensorflow/dtensor/python/d_variable.py b/tensorflow/dtensor/python/d_variable.py
index 3fea2f8fe4d..875516df400 100644
--- a/tensorflow/dtensor/python/d_variable.py
+++ b/tensorflow/dtensor/python/d_variable.py
@@ -185,10 +185,13 @@ class DVariable(resource_variable_ops.ResourceVariable):
         # ShapeInference to extract the shape. Scalar tensors (shape=()) can't
         # be sliced/sharded, so they must use shard_info=None rather than a
         # ShardInfo with an empty shape and offset (which would otherwise
-        # produce an invalid " " shape_and_slice spec for RestoreV2).
+        # produce an invalid " " shape_and_slice spec for RestoreV2). Check
+        # `len(shape) > 0` rather than truthiness: an empty tuple () is
+        # falsy, but a scalar tf.TensorShape([]) is truthy even though it
+        # also has length 0, so a plain `if shape` would miss that case.
         shard_info = (
             trackable.ShardInfo(shape=shape, offset=[0] * len(shape))
-            if shape
+            if len(shape) > 0
             else None
         )
         initial_value = api.call_with_layout(
diff --git a/tensorflow/dtensor/python/tests/save_restore_v2_test.py b/tensorflow/dtensor/python/tests/save_restore_v2_test.py
index 6d00a246598..a1e2ab448e8 100644
--- a/tensorflow/dtensor/python/tests/save_restore_v2_test.py
+++ b/tensorflow/dtensor/python/tests/save_restore_v2_test.py
@@ -155,6 +155,23 @@ class DTensorSaveRestoreV2Test(test_util.DTensorBaseTest):
         constant_op.constant(42.0), layout, v_restored.read_value()
     )
 
+    # Regression coverage for the scalar tf.TensorShape([]) boundary case:
+    # unlike an empty tuple (), TensorShape([]) is truthy even though it also
+    # has length 0, so this must go through a fresh restore (re-using
+    # `init_callable` would double-consume the same deferred restore).
+    restore_ckpt_ts = checkpoint.Checkpoint()
+    restore_ckpt_ts.restore(saved_path)
+    init_callable_ts = restore_ckpt_ts._preload_simple_restoration('v')
+    v_restored_ts = d_variable.DVariable(
+        init_callable_ts,
+        shape=constant_op.constant(42.0).shape,
+        layout=layout,
+        dtype=dtypes.float32,
+    )
+    self.assertDTensorEqual(
+        constant_op.constant(42.0), layout, v_restored_ts.read_value()
+    )
+
   @parameterized.named_parameters(
       ('x_unsharded', [_MESH_DIM_X, UNSHARDED]),
       ('unsharded_x', [UNSHARDED, _MESH_DIM_X]),

From 9218c9add67bb91d85bb5be03d43c743bcd91db5 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 21:55:35 +0530
Subject: [PATCH 0791/1000] Revert double-free regression in
 ScopedTFBuffer.__del__

The previous round added 'if c_api is not None: c_api.TF_DeleteBuffer(
self.buffer)' ahead of the existing getattr-guarded block, but that
existing block already handled both interpreter shutdown (c_api is
None) and partial initialization (self.buffer never assigned) safely.
The addition caused every valid ScopedTFBuffer to have
TF_DeleteBuffer called on it twice.

Verified by direct execution: the file as it stood crashed the
interpreter with 'munmap_chunk(): invalid pointer' on a plain
ScopedTFBuffer.__del__() call, confirming the double free. Removed
the redundant first block, restoring the single guarded deletion
path.

Also added the suggested testDeletePartiallyInitializedObjectDoesNotRaise,
covering __del__ on an object created via __new__ (so self.buffer was
never assigned) to guard against this regressing again.

All 10 real tests in c_api_util_test.py (ApiDefMapTest, UniquePtrTest,
ScopedTFBufferTest x3) pass against the patched module.
---
 tensorflow/python/framework/c_api_util.py      | 2 --
 tensorflow/python/framework/c_api_util_test.py | 9 +++++++++
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/framework/c_api_util.py b/tensorflow/python/framework/c_api_util.py
index 2f223af2d89..99342037186 100644
--- a/tensorflow/python/framework/c_api_util.py
+++ b/tensorflow/python/framework/c_api_util.py
@@ -138,8 +138,6 @@ class ScopedTFBuffer(object):
   def __del__(self):
     # Note: when we're destructing the global context (i.e when the process is
     # terminating) we can have already deleted other modules.
-    if c_api is not None:
-      c_api.TF_DeleteBuffer(self.buffer)
     # Also guard against __init__ having failed before `self.buffer` was
     # ever assigned (e.g. compat.as_bytes raising on bad input): __del__
     # still runs on a partially-constructed object, and accessing
diff --git a/tensorflow/python/framework/c_api_util_test.py b/tensorflow/python/framework/c_api_util_test.py
index 64be575e35b..4d49fbeca17 100644
--- a/tensorflow/python/framework/c_api_util_test.py
+++ b/tensorflow/python/framework/c_api_util_test.py
@@ -151,5 +151,14 @@ class ScopedTFBufferTest(test_util.TensorFlowTestCase):
     with mock.patch.object(c_api_util, "c_api", None):
       buf.__del__()  # should not raise
 
+  def testDeletePartiallyInitializedObjectDoesNotRaise(self):
+    # Regression test: if __init__ raises before `self.buffer` is ever
+    # assigned (e.g. compat.as_bytes failing on bad input), __del__ still
+    # runs on the partially-constructed object. It must not raise an
+    # AttributeError from accessing the missing `self.buffer`, which would
+    # mask whatever real exception caused __init__ to fail.
+    buf = c_api_util.ScopedTFBuffer.__new__(c_api_util.ScopedTFBuffer)
+    buf.__del__()  # should not raise
+
 if __name__ == "__main__":
   googletest.main()

From f86884ba1b615b773d60f5d4a8f70d219c9c717e Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2026 09:15:12 -0700
Subject: [PATCH 0792/1000] Integrate LLVM at llvm/llvm-project@bb80dc2a4f6f

Updates LLVM usage to match
[bb80dc2a4f6f](https://github.com/llvm/llvm-project/commit/bb80dc2a4f6ffc8262701cb0417c6f08158a17df)

PiperOrigin-RevId: 999414872
---
 .../xla/third_party/llvm/workspace.bzl        |    4 +-
 .../xla/third_party/shardy/temporary.patch    | 2429 ++++++++++++++++-
 .../xla/third_party/stablehlo/temporary.patch | 2270 +++++++--------
 3 files changed, 3606 insertions(+), 1097 deletions(-)

diff --git a/third_party/xla/third_party/llvm/workspace.bzl b/third_party/xla/third_party/llvm/workspace.bzl
index 5906818eabd..b13ddf8949c 100644
--- a/third_party/xla/third_party/llvm/workspace.bzl
+++ b/third_party/xla/third_party/llvm/workspace.bzl
@@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "4a0a7e3ddc3bdc4fde7c5b9b9e6890fb4dbf9153"
-    LLVM_SHA256 = "1453c2abaa0fdb957a6108a1802d282bd01010dcde904288c698d44f7539c488"
+    LLVM_COMMIT = "bb80dc2a4f6ffc8262701cb0417c6f08158a17df"
+    LLVM_SHA256 = "27dd5e6295f58415653076a65b6b8d3e3210b09e778b102b086b47ff3b2bcb5e"
 
     tf_http_archive(
         name = name,
diff --git a/third_party/xla/third_party/shardy/temporary.patch b/third_party/xla/third_party/shardy/temporary.patch
index 4d53e31bb73..57bc9af6357 100644
--- a/third_party/xla/third_party/shardy/temporary.patch
+++ b/third_party/xla/third_party/shardy/temporary.patch
@@ -43,6 +43,36 @@ index 1ea9160e..8f149153 100644
  }
  
  TEST_F(DialectTest, DimensionShardingAttrGetSharded) {
+diff --git a/shardy/dialect/sdy/transforms/import/test/import_pipeline.mlir b/shardy/dialect/sdy/transforms/import/test/import_pipeline.mlir
+index d8200b88..60204b54 100644
+--- a/shardy/dialect/sdy/transforms/import/test/import_pipeline.mlir
++++ b/shardy/dialect/sdy/transforms/import/test/import_pipeline.mlir
+@@ -2,8 +2,8 @@
+ 
+ sdy.mesh @mesh = <["a"=2]>
+ 
+-// Verifies that `-apply-sharding-constraints` pass is applied after
+-// `-add-data_flow_edges` pass
++// Verifies that `-apply-sharding-constraints` pass is not applied
++// before -add-data_flow_edges pass.
+ // CHECK-LABEL: func @main
+ func.func @main(%arg0: tensor<32x96xf32>) -> tensor<32x96xf32> {
+   // CHECK-NEXT: %[[OPT_BARRIER:.*]] = stablehlo.optimization_barrier %arg0
+@@ -53,11 +53,9 @@ func.func @main(%arg0: tensor<8xf32>) -> tensor<8xf32> {
+ 
+ sdy.mesh @mesh = <["c"=2, "a"=2, "b"=2]>
+ 
+-// Due to the in_sharding being fully closed, the in_sharding is added to the
+-// func arg but with the manual axis added as replicated.
+-// CHECK-LABEL: func @main
+-// CHECK-SAME     %arg0: tensor<16x16xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"c"}, {}], replicated={"a"}>}
+-// CHECK-SAME     -> tensor<16x16xf32> {
++// CHECK-LABEL: func @main(
++// CHECK-SAME:      %arg0: tensor<8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"c"}], replicated={"a"}>})
++// CHECK-SAME:      -> tensor<8xf32> {
+ func.func @main(%arg0: tensor<8xf32>) -> tensor<8xf32> {
+   // CHECK-NEXT: %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{"c"}], replicated={"a"}>] out_shardings=[<@mesh, [{"c"}], replicated={"a"}>] manual_axes={"c", "a"} (%arg1: tensor<4xf32>) {
+   // CHECK-NEXT:   %2 = sdy.data_flow_edge %arg1 sharding=<@mesh, [{}]>
 diff --git a/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc b/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc
 index 8417874d..48217e35 100644
 --- a/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc
@@ -58,6 +88,76 @@ index 8417874d..48217e35 100644
                      return WalkResult::interrupt();
                    }
                  }
+diff --git a/shardy/dialect/sdy/transforms/propagation/debugging/source_sharding.cc b/shardy/dialect/sdy/transforms/propagation/debugging/source_sharding.cc
+index 183f25fd..b2933711 100644
+--- a/shardy/dialect/sdy/transforms/propagation/debugging/source_sharding.cc
++++ b/shardy/dialect/sdy/transforms/propagation/debugging/source_sharding.cc
+@@ -691,7 +691,7 @@ void prepareFuncResultToEdgesHandler(
+   }
+ }
+ 
+-OriginSharding lookUpValueOriginSharding(
++std::optional<OriginSharding> lookUpValueOriginSharding(
+     Value value, AxisRefAttr axisRef,
+     const ValueToOriginShardingMap& valueToOriginShardingMap) {
+   // NOTE: need to call `getShardableValue` in case the operand/result is
+@@ -700,7 +700,7 @@ OriginSharding lookUpValueOriginSharding(
+   auto valueIt = valueToOriginShardingMap.find(getShardableValue(value));
+   if (valueIt == valueToOriginShardingMap.end()) {
+     SDY_LOG(WARNING) << "Couldn't find sharding origin for value";
+-    return {};
++    return std::nullopt;
+   }
+   const AxisToOriginShardingMap& axisToOriginSharding = valueIt->second;
+   if (auto it = axisToOriginSharding.find(axisRef);
+@@ -717,7 +717,7 @@ OriginSharding lookUpValueOriginSharding(
+     }
+   }
+   SDY_LOG(WARNING) << "Couldn't find sharding origin";
+-  return {};
++  return std::nullopt;
+ }
+ 
+ }  // namespace
+@@ -752,8 +752,9 @@ void SourceShardingHandler::operator()(function_ref<void()> transform,
+   propagationStep++;
+   // If the new and old shardings are different, something was propagated to it.
+   // Find and save it.
+-  auto lookUpOriginSharding = [&](EdgeNode edgeNode,
+-                                  AxisRefAttr axisRef) -> OriginSharding {
++  auto lookUpOriginSharding =
++      [&](EdgeNode edgeNode,
++          AxisRefAttr axisRef) -> std::optional<OriginSharding> {
+     switch (edgeNode.type) {
+       case EdgeNodeType::OPERAND:
+         return lookUpValueOriginSharding(
+@@ -778,8 +779,11 @@ void SourceShardingHandler::operator()(function_ref<void()> transform,
+       }
+     }
+     if (mappings->debugShardingOrigins) {
+-      mappings->valueToOriginShardingMap[getShardableValue(value)].try_emplace(
+-          axisRef, lookUpOriginSharding(edge.source, axisRef));
++      if (std::optional<OriginSharding> originSharding =
++              lookUpOriginSharding(edge.source, axisRef)) {
++        mappings->valueToOriginShardingMap[getShardableValue(value)]
++            .try_emplace(axisRef, *originSharding);
++      }
+     }
+   };
+ 
+diff --git a/shardy/dialect/sdy/transforms/propagation/debugging/test/sharding_origins.mlir b/shardy/dialect/sdy/transforms/propagation/debugging/test/sharding_origins.mlir
+index 16d7d20b..f9c3e2d2 100644
+--- a/shardy/dialect/sdy/transforms/propagation/debugging/test/sharding_origins.mlir
++++ b/shardy/dialect/sdy/transforms/propagation/debugging/test/sharding_origins.mlir
+@@ -790,7 +790,7 @@ sdy.mesh @mesh = <["a"=2, "b"=2, "c"=8]>
+ 
+ // CHECK-LABEL: func.func @input_already_has_sharding(
+ // CHECK-SAME:    %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a", ?}, {"b", ?}]>,
+-// CHECK-SAME:                            sdy.sharding_origins = {a = "self", b = "constraint_0"}})
++// CHECK-SAME:                            sdy.sharding_origins = {b = "constraint_0"}})
+ // CHECK-SAME:    -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{?}, {"b", ?}]>,
+ // CHECK-SAME:                         sdy.sharding_origins = {b = "constraint_0"}}) {
+ func.func @input_already_has_sharding(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
 diff --git a/shardy/dialect/sdy/transforms/propagation/test/basic_propagation.mlir b/shardy/dialect/sdy/transforms/propagation/test/basic_propagation.mlir
 index c5cd8022..6d9f32c2 100644
 --- a/shardy/dialect/sdy/transforms/propagation/test/basic_propagation.mlir
@@ -78,7 +178,7 @@ index c5cd8022..6d9f32c2 100644
 +  return %1 : tensor<8x8xf32>
 +}
 diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
-index e42d7dbd..5906818e 100644
+index e42d7dbd..b13ddf89 100644
 --- a/third_party/llvm/workspace.bzl
 +++ b/third_party/llvm/workspace.bzl
 @@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
@@ -87,11 +187,2334 @@ index e42d7dbd..5906818e 100644
      """Imports LLVM."""
 -    LLVM_COMMIT = "2e0d9e12b36e853d9b108c4da09509cd54d1f3f9"
 -    LLVM_SHA256 = "8e4a851780af109d623fa6d6db9b1e62d7126272d17268fed1b149f0e737c1d7"
-+    LLVM_COMMIT = "4a0a7e3ddc3bdc4fde7c5b9b9e6890fb4dbf9153"
-+    LLVM_SHA256 = "1453c2abaa0fdb957a6108a1802d282bd01010dcde904288c698d44f7539c488"
++    LLVM_COMMIT = "bb80dc2a4f6ffc8262701cb0417c6f08158a17df"
++    LLVM_SHA256 = "27dd5e6295f58415653076a65b6b8d3e3210b09e778b102b086b47ff3b2bcb5e"
  
      tf_http_archive(
          name = name,
+diff --git a/third_party/stablehlo/temporary.patch b/third_party/stablehlo/temporary.patch
+index c0bd29a0..d2710427 100644
+--- a/third_party/stablehlo/temporary.patch
++++ b/third_party/stablehlo/temporary.patch
+@@ -15,755 +15,42 @@
+ #
+ # This file is automatically generated by the `generate_patch` tool.
+ # Do not edit directly.
+-diff --ruN a/stablehlo/.github/ISSUE_TEMPLATE/bug_report.yml b/stablehlo/.github/ISSUE_TEMPLATE/bug_report.yml
+---- stablehlo/.github/ISSUE_TEMPLATE/bug_report.yml
+-+++ stablehlo/.github/ISSUE_TEMPLATE/bug_report.yml
+-@@ -1,50 +0,0 @@
+--# Copyright 2024 The StableHLO Authors.
+--# 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.
+--
+--name: 🐞 Bug Report
+--description: Report an issue
+--body:
+--  - type: markdown
+--    attributes:
+--      value: |
+--        :star2: Thanks for taking the time to report this issue! :star2:
+--
+--        Please search through [other recent issues](https://github.com/openxla/stablehlo/issues) to see if your report overlaps with an existing issue.
+--  - type: textarea
+--    id: what-happened
+--    attributes:
+--      label: What happened?
+--      description: |
+--        Also tell us, what did you expect to happen?
+--
+--        Include any MLIR output leading up to whatever error occurred where possible.
+--      placeholder: Tell us what you see! Please also copy/paste or link to any relevant log output.
+--    validations:
+--      required: true
+--  - type: textarea
+--    id: reproducing
+--    attributes:
+--      label: Steps to reproduce your issue
+--      description: |
+--        Include any commands that you ran.
+--      value: |
+--        1. Go to '...'
+--        2. Click on '....'
+--        3. Scroll down to '....'
+--        4. See error
+--  - type: textarea
+--    id: version
+--    attributes:
+--      label: Version information
+--      description: What version of our software are you using?
+--      placeholder: This could be a git commit hash.
+-diff --ruN a/stablehlo/.github/ISSUE_TEMPLATE/config.yml b/stablehlo/.github/ISSUE_TEMPLATE/config.yml
+---- stablehlo/.github/ISSUE_TEMPLATE/config.yml
+-+++ stablehlo/.github/ISSUE_TEMPLATE/config.yml
+-@@ -1,24 +0,0 @@
+--# Copyright 2024 The StableHLO Authors.
+--# 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.
+--
+--blank_issues_enabled: true
+--contact_links:
+--  - name: 📄 Blank Issue
+--    url: https://github.com/openxla/stablehlo/issues/new
+--    about: If you know what you're doing (especially collaborators)
+--  - name: 🗣 Start a discussion on GitHub
+--    url: https://github.com/orgs/openxla/discussions/new
+--    about: For early feedback on a design or other proposal
+--  - name: 💬 Join us on Discord
+--    url: https://discord.gg/Wqr8vWZKKJ
+--    about: For realtime communication with the team and collaborators
+-diff --ruN a/stablehlo/.github/ISSUE_TEMPLATE/feature_request.yml b/stablehlo/.github/ISSUE_TEMPLATE/feature_request.yml
+---- stablehlo/.github/ISSUE_TEMPLATE/feature_request.yml
+-+++ stablehlo/.github/ISSUE_TEMPLATE/feature_request.yml
+-@@ -1,30 +0,0 @@
+--# Copyright 2024 The StableHLO Authors.
+--# 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.
+--
+--name: ➕ Feature Request
+--description: Ask for a new feature to be added
+--body:
+--  - type: markdown
+--    attributes:
+--      value: |
+--        :star2: Thanks for taking the time to file this request! :star2:
+--
+--        We also work closely with upstream communities like the MLIR project, which uses the [MLIR Discourse forum](https://discourse.llvm.org/c/mlir/31) for coordination.
+--  - type: textarea
+--    id: overview
+--    attributes:
+--      label: Request description
+--      description: |
+--        What would you like added or changed?
+--    validations:
+--      required: true
+-diff --ruN a/stablehlo/.github/actions/setup-build/action.yml b/stablehlo/.github/actions/setup-build/action.yml
+---- stablehlo/.github/actions/setup-build/action.yml
+-+++ stablehlo/.github/actions/setup-build/action.yml
+-@@ -1,58 +0,0 @@
+--# Copyright 2023 The StableHLO Authors.
+--#
+--# 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
+--#
+--#      https://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.
+--
+--# The setup-build action gets everything needed by buildAndTest into the workspace.
+--name: "Setup build environment (ninja, ccache, llvm, clang, lld)"
+--description: "Setup build environment (ninja, ccache, llvm, clang, lld)"
+--
+--inputs:
+--  llvm-version:
+--    description: |
+--      LLVM version for checkout and build. Used for ccache value and checkout.
+--    required: true
+--  clang-version:
+--    description: |
+--      Clang and LLD version to use for compilation. Independent of LLVM version.
+--    required: false
+--    default: '20'
+--  python-version:
+--    description: |
+--      Python version to install.
+--    required: false
+--    default: '3.11.13'
+--runs:
+--  # This is a composite action - has a list of steps to execute.
+--  using: "composite"
+--
+--  steps:
+--  # Checkout llvm at version specified in input argument.
+--  - uses: actions/checkout@v4
+--    with:
+--      repository: llvm/llvm-project
+--      ref: ${{ inputs.llvm-version }}
+--      path: llvm-project
+--      fetch-depth: 1
+--
+--  # Setup C++ caching using ccache.
+--  # Cache key is a combination of OS arch and LLVM version.
+--  - name: Ccache for C++ compilation
+--    uses: hendrikmuhs/ccache-action@v1.2
+--    with:
+--      key: ${{ runner.os }}-stablehlo_build_assets-${{ inputs.llvm-version }}
+--      max-size: 4G
+--
+--  - name: Install MLIR python requirements
+--    shell: bash
+--    run: |
+--      python3 -m pip install -r $GITHUB_WORKSPACE/llvm-project/mlir/python/requirements.txt
+ diff --ruN a/stablehlo/.github/workflows/buildAndTestBazel.yml b/stablehlo/.github/workflows/buildAndTestBazel.yml
+ --- stablehlo/.github/workflows/buildAndTestBazel.yml
+ +++ stablehlo/.github/workflows/buildAndTestBazel.yml
+-@@ -1,103 +0,0 @@
+--# Copyright 2023 The StableHLO Authors.
+--#
+--# 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
+--#
+--#      https://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.
+--
+--name: Bazel Build
+--
+--on:
+--  pull_request:
+--    branches: [ main ]
+--  push:
+--    branches: [ main ]
+--  schedule:
+--    # Run once a day
+--    - cron:  '0 12 * * *'
+--  workflow_dispatch:
+--    inputs:
+--      halt-for-connection:
+--        description: 'Should this workflow run wait for a remote connection?'
+--        type: choice
+--        required: true
+--        default: 'no'
+--        options:
+--          - 'yes'
+--          - 'no'
+--
+--# Ensure that only a single job or workflow using the same
+--# concurrency group will run at a time. This would cancel
+--# any in-progress jobs in the same github workflow and github
+--# ref (e.g. refs/heads/main or refs/pull/<pr_number>/merge).
+--concurrency:
+--  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+--  cancel-in-progress: ${{ github.ref != 'main' }}
+--
+--jobs:
+--  bazel-build:
+--    strategy:
+--      fail-fast: false
++@@ -45,7 +45,7 @@
++   bazel-build:
++     strategy:
++       fail-fast: false
+ -    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n2-16' || 'ubuntu-22.04' }}
+--    container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
+--
+--    steps:
+--    # Checkout StableHLO with all commits in the chain
+--    - name: Checkout StableHLO
+--      uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
+--      with:
+--        fetch-depth: $(( ${{ github.event_name == 'pull_request' && github.event.pull_request.commits || 0 }} + 1 ))
+--        set-safe-directory: true
+--
+--    - name: Mark project as safe for Git
+--      run: git config --global --add safe.directory '*'
+--
+--    - name: "Wait For Connection"
+--      uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
+--      with:
+--        halt-dispatch-input: ${{ inputs.halt-for-connection }}
+--
+--    - name: Validate LLVM Commit
+--      run: |
+--        ./build_tools/github_actions/lint_llvm_commit.sh .
+--
+--    - name: Get LLVM Version
+--      id: llvm-version
+--      shell: bash
+--      run: |
+--        echo "version=$(cat ./build_tools/llvm_version.txt)" >> $GITHUB_OUTPUT
+--
+--    - name: Build and check docs
+--      shell: bash
+--      run: |
+--        ./build_tools/github_actions/ci_build_docs.sh -c
+--
+--    - name: Build and Test StableHLO (Diff Only)
+--      if: ${{ github.event_name == 'pull_request' }}
+--      shell: bash
+--      run: |
+--        START_HASH=${{ github.event.pull_request.base.sha }}
+--        END_HASH=${{ github.event.pull_request.head.sha }}
+--        BAZEL_DIFF=/tmp/bazel-diff.jar
+--        curl -Lo "$BAZEL_DIFF" https://github.com/Tinder/bazel-diff/releases/download/17.0.2/bazel-diff_deploy.jar
+--        ./build_tools/github_actions/ci_build_bazel.sh "$BAZEL_DIFF" "$START_HASH" "$END_HASH"
+--
+--    - name: Build and Test StableHLO (All)
+--      if: ${{ github.event_name != 'pull_request' }}
+--      shell: bash
+--      run: |
+--        ./build_tools/github_actions/ci_build_bazel.sh
+--
+--    - name: "Wait For Connection - On Error"
+--      if: failure()
+--      uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
+--      with:
+--        halt-dispatch-input: ${{ inputs.halt-for-connection }}
+--
+++    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n4-16' || 'ubuntu-22.04' }}
++     container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
++ 
++     steps:
+ diff --ruN a/stablehlo/.github/workflows/buildAndTestCMake.yml b/stablehlo/.github/workflows/buildAndTestCMake.yml
+ --- stablehlo/.github/workflows/buildAndTestCMake.yml
+ +++ stablehlo/.github/workflows/buildAndTestCMake.yml
+-@@ -1,98 +0,0 @@
+--# Copyright 2023 The StableHLO Authors.
+--#
+--# 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
+--#
+--#      https://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.
+--
+--name: CMake Build
+--
+--on:
+--  pull_request:
+--    branches: [ main ]
+--  push:
+--    branches: [ main ]
+--  schedule:
+--    # Run once every 6hr building against llvm-project@HEAD
+--    - cron:  '0 */6 * * *'
+--  workflow_dispatch:
+--
+--# Ensure that only a single job or workflow using the same
+--# concurrency group will run at a time. This would cancel
+--# any in-progress jobs in the same github workflow and github
+--# ref (e.g. refs/heads/main or refs/pull/<pr_number>/merge).
+--concurrency:
+--  group: ${{ github.workflow }}-${{ github.ref }}
+--  cancel-in-progress: true
+--
+--jobs:
+--  cmake-build:
+--    # Only run scheduled CI on main repo
+--    if: (github.repository == 'openxla/stablehlo' || github.event_name != 'schedule')
+--    name: "cmake-build${{ github.event_name == 'schedule' && ' (llvm-project@HEAD)' || ''}}"
+--    env:
+--      LLVM_PROJECT_DIR: "llvm-project"
+--      LLVM_BUILD_DIR: "llvm-build"
+--      STABLEHLO_BUILD_DIR: "stablehlo-build"
+--      STABLEHLO_PYTHON_BUILD_DIR: "stablehlo-python-build"
+--    strategy:
+--      fail-fast: false
++@@ -44,7 +44,7 @@
++       STABLEHLO_PYTHON_BUILD_DIR: "stablehlo-python-build"
++     strategy:
++       fail-fast: false
+ -    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n2-16' || 'ubuntu-22.04' }}
+--    container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cmake:latest"
+--
+--    steps:
+--    - name: Checkout StableHLO
+--      uses: actions/checkout@v4
+--
+--    - name: Get LLVM Version
+--      id: llvm-version
+--      shell: bash
+--      run: |
+--        USE_LLVM_HEAD=${{ github.event_name == 'schedule' }}
+--        if [[ $USE_LLVM_HEAD = true ]]; then
+--          echo "version=main" >> $GITHUB_OUTPUT
+--        else
+--          echo "version=$(cat ./build_tools/llvm_version.txt)" >> $GITHUB_OUTPUT
+--        fi;
+--
+--    - name: Setup workspace
+--      uses: ./.github/actions/setup-build
+--      with:
+--        llvm-version: ${{ steps.llvm-version.outputs.version }}
+--
+--    - name: Configure and Build LLVM
+--      shell: bash
+--      run: |
+--        ./build_tools/github_actions/ci_build_cmake_llvm.sh "$LLVM_PROJECT_DIR" "$LLVM_BUILD_DIR"
+--      env:
+--          CMAKE_BUILD_TYPE: Release
+--          MLIR_ENABLE_BINDINGS_PYTHON: ON
+--
+--    - name: Build and Test StableHLO (with AddressSanitizer)
+--      shell: bash
+--      run: |
+--        ./build_tools/github_actions/ci_build_cmake.sh "$LLVM_BUILD_DIR" "$STABLEHLO_BUILD_DIR"
+--      env:
+--          CMAKE_BUILD_TYPE: Release
+--          STABLEHLO_ENABLE_BINDINGS_PYTHON: OFF
+--          STABLEHLO_ENABLE_SANITIZER: address
+--          # TODO: remove this once https://github.com/openxla/stablehlo/pull/2318 is fixed
+--          ASAN_OPTIONS: allow_user_poisoning=false
+--
+--    - name: Build and Test StableHLO (with Python bindings)
+--      shell: bash
+--      run: |
+--        python3 -m pip install tensorflow-cpu nanobind==2.9
+--        python3 -m pip install -r "$LLVM_PROJECT_DIR/mlir/python/requirements.txt"
+--        ./build_tools/github_actions/ci_build_cmake.sh "$LLVM_BUILD_DIR" "$STABLEHLO_BUILD_DIR"
+--      env:
+--          CMAKE_BUILD_TYPE: Release
+--          STABLEHLO_ENABLE_BINDINGS_PYTHON: ON
+--          STABLEHLO_ENABLE_PYTHON_TF_TESTS: ON
+-diff --ruN a/stablehlo/.github/workflows/execute_tutorials.yml b/stablehlo/.github/workflows/execute_tutorials.yml
+---- stablehlo/.github/workflows/execute_tutorials.yml
+-+++ stablehlo/.github/workflows/execute_tutorials.yml
+-@@ -1,46 +0,0 @@
+--# Copyright 2024 The StableHLO 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
+--
+--#      https://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.
+--
+--name: Execute Tutorial Notebooks
+--
+--on:
+--  push:
+--    paths:
+--      - 'docs/tutorials/**.ipynb'
+--      - .github/workflows/execute_tutorials.yml
+--  schedule:
+--    # Run once a day
+--    - cron:  '0 20 * * *'
+--  workflow_dispatch:
+--
+--jobs:
+--  test-notebooks:
+--    runs-on: ubuntu-latest
+--
+--    steps:
+--      - name: Checkout repository
+--        uses: actions/checkout@v4
+--
+--      - name: Set up Python and install dependencies
+--        uses: actions/setup-python@v5
+--        with:
+--          python-version: '3.11'
+--      - run:  |
+--          pip install --upgrade pip
+--          pip install ipython ipykernel nbclient
+--
+--      - name: Execute tutorial notebooks
+--        run: |
+--          jupyter execute docs/tutorials/*.ipynb
+--        shell: bash
+-diff --ruN a/stablehlo/.github/workflows/lint.yml b/stablehlo/.github/workflows/lint.yml
+---- stablehlo/.github/workflows/lint.yml
+-+++ stablehlo/.github/workflows/lint.yml
+-@@ -1,111 +0,0 @@
+--# Copyright 2023 The StableHLO Authors.
+--#
+--# 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
+--#
+--#      https://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.
+--
+--name: Lint
+--
+--on:
+--  pull_request:
+--    branches: [ main ]
+--  push:
+--    branches: [ main ]
+--  schedule:
+--    # Run once a day
+--    - cron:  '0 12 * * *'
+--  workflow_dispatch:
+--
+--env:
+--  # Have `go install` place binaries in $PATH
+--  GOBIN: "/usr/local/bin"
+--
+--jobs:
+--  clang-format:
+--    # This job can only be run on pull_request since GITHUB_BASE_REF is only set on PR.
+--    if: "github.event_name == 'pull_request'"
+--    runs-on: ubuntu-22.04
+--    steps:
+--      - name: Checking out repository
+--        uses: actions/checkout@v4
+--
+--      - name: Run clang-format
+--        run: |
+--          git fetch --no-tags --prune --depth=1 origin "${GITHUB_BASE_REF?}:${GITHUB_BASE_REF?}"
+--          ./build_tools/github_actions/lint_clang_format.sh -b "${GITHUB_BASE_REF}"
+--
+--  whitespace-checks:
+--    if: "github.event_name == 'pull_request'"
+--    runs-on: ubuntu-22.04
+--    steps:
+--      - name: Checking out repository
+--        uses: actions/checkout@v4
+--
+--      - name: Run Whitespace Checks
+--        run: |
+--          git fetch --no-tags --prune --depth=1 origin "${GITHUB_BASE_REF?}:${GITHUB_BASE_REF?}"
+--          ./build_tools/github_actions/lint_whitespace_checks.sh
+--
+--  version-checks:
+--    if: "github.event_name == 'pull_request'"
+--    runs-on: ubuntu-22.04
+--    steps:
+--      - name: Checking out repository
+--        uses: actions/checkout@v4
+--
+--      - name: Run Version Checks
+--        run: |
+--          ./build_tools/github_actions/lint_version.sh
+--
+--  buildifier-check:
+--    if: "github.event_name == 'pull_request'"
+--    runs-on: ubuntu-22.04
+--    defaults:
+--      run:
+--        shell: bash
+--    steps:
+--      - name: Checking out repository
+--        uses: actions/checkout@v4
+--
+--      - name: Install buildifier
+--        run: go install github.com/bazelbuild/buildtools/buildifier@433ea85 # 6.4.0
+--
+--      - name: Run buildifier checks
+--        run: ./build_tools/github_actions/lint_buildifier.sh
+--
+--  license-check:
+--    if: "github.event_name == 'pull_request'"
+--    runs-on: ubuntu-22.04
+--    defaults:
+--      run:
+--        shell: bash
+--    steps:
+--      - name: Checking out repository
+--        uses: actions/checkout@v4
+--
+--      - name: Run license checks
+--        run: |
+--          git fetch --no-tags --prune --depth=1 origin "${GITHUB_BASE_REF?}:${GITHUB_BASE_REF?}"
+--          ./build_tools/github_actions/lint_check_license.sh -b "${GITHUB_BASE_REF}"
+--
+--  shell-script-check:
+--    if: "github.event_name == 'pull_request'"
+--    runs-on: ubuntu-22.04
+--    defaults:
+--      run:
+--        shell: bash
+--    steps:
+--      - name: Checking out repository
+--        uses: actions/checkout@v4
+--
+--      # shellcheck is installed by default on Ubuntu runners: https://github.com/actions/runner-images/blob/37a171e44de1f134793d7f0cbff137ca481942ab/images/ubuntu/Ubuntu2204-Readme.md
+--      - name: Run shellcheck
+--        run: ./build_tools/github_actions/lint_shellcheck.sh
+-diff --ruN a/stablehlo/.github/workflows/markdown_lint.yml b/stablehlo/.github/workflows/markdown_lint.yml
+---- stablehlo/.github/workflows/markdown_lint.yml
+-+++ stablehlo/.github/workflows/markdown_lint.yml
+-@@ -1,35 +0,0 @@
+--# Copyright 2023 The StableHLO Authors.
+--#
+--# 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
+--#
+--#      https://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.
+--
+--name: Markdown Lint
+--
+--on:
+--  pull_request:
+--    branches: [ main ]
+--  push:
+--    branches: [ main ]
+--  schedule:
+--    # Run once a day
+--    - cron:  '0 12 * * *'
+--  workflow_dispatch:
+--
+--jobs:
+--  markdown-lint:
+--    runs-on: ubuntu-22.04
+--    steps:
+--    - name: Checking out repository
+--      uses: actions/checkout@v4
+--    - name: Running markdownlint
+--      run: |
+--        ./build_tools/github_actions/lint_markdown.sh --ignore "./docs/generated/**/*.md" "./**/*.md"
+++    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n4-16' || 'ubuntu-22.04' }}
++     container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cmake:latest"
++ 
++     steps:
+ diff --ruN a/stablehlo/.github/workflows/publishWheelRelease.yml b/stablehlo/.github/workflows/publishWheelRelease.yml
+ --- stablehlo/.github/workflows/publishWheelRelease.yml
+ +++ stablehlo/.github/workflows/publishWheelRelease.yml
+-@@ -1,136 +0,0 @@
+--# Copyright 2024 The StableHLO Authors.
+--#
+--# 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
+--#
+--#      https://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.
+--
+--name: Publish Python Wheel Release
+--
+--on:
+--  workflow_dispatch:
+--  schedule:
+--    # Runs at 6:00 AM UTC, which is 10:00 PM previous day PST (UTC-8)
+--    - cron: '0 6 * * *'
+--
+--# Ensure that only a single job or workflow using the same
+--# concurrency group will run at a time. This would cancel
+--# any in-progress jobs in the same github workflow and github
+--# ref (e.g. refs/heads/main or refs/pull/<pr_number>/merge).
+--concurrency:
+--  group: ${{ github.workflow }}-${{ github.ref }}
+--  cancel-in-progress: true
+--
+--jobs:
+--  # TODO(fzakaria): Consider running this job only if a new commit for the day has been done
+--  cmake-build-wheel:
+--    # Only run scheduled CI on main repo
+--    if: (github.repository == 'openxla/stablehlo' || github.event_name != 'schedule')
+--    name: "cmake-build ${{ github.event_name == 'schedule' && '(llvm-project@HEAD)' || ''}}"
+--    env:
+--      LLVM_PROJECT_DIR: "llvm-project"
+--      LLVM_BUILD_DIR: "llvm-build"
+--      # The setup.py script expects the StableHLO build to be in 'build'
+--      STABLEHLO_BUILD_DIR: "build"
+--    strategy:
+--      fail-fast: false
+--      matrix:
+--        python-version: [ "3.10", "3.11", "3.12" ]
++@@ -43,7 +43,7 @@
++       fail-fast: false
++       matrix:
++         python-version: [ "3.10", "3.11", "3.12" ]
+ -    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n2-16' ||  'ubuntu-22.04'  }}
+--    container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cmake:latest"
+--
+--    steps:
+--    - name: Checkout StableHLO
+--      uses: actions/checkout@v4
+--      # TODO(fzakaria): We need all the commits and the tags for the
+--      # git describe` command to work to find the correct version.
+--      # This can be removed if the version can be determmined in another way.
+--      with:
+--        fetch-depth: 0
+--    - name: Get LLVM Version
+--      id: llvm-version
+--      shell: bash
+--      run: |
+--        echo "version=$(cat ${{ github.workspace }}/build_tools/llvm_version.txt)" >> $GITHUB_OUTPUT
+--
+--    - name: Setup workspace
+--      uses: ./.github/actions/setup-build
+--      with:
+--        llvm-version: ${{ steps.llvm-version.outputs.version }}
+--        python-version: ${{ matrix.python-version }}
+--
+--    # The actual tool need to execute the python build
+--    # TODO(fzakaria): Consider using cibuildwheel to build multiple wheels at
+--    # once and mark them as manylinux compatible.
+--    - name: Install build dependencies
+--      shell: bash
+--      run: pip install build
+--
+--    - name: Configure and Build LLVM
+--      shell: bash
+--      run: |
+--        ./build_tools/github_actions/ci_build_cmake_llvm.sh "$LLVM_PROJECT_DIR" "$LLVM_BUILD_DIR"
+--      env:
+--          CMAKE_BUILD_TYPE: Release
+--          MLIR_ENABLE_BINDINGS_PYTHON: ON
+--
+--    - name: Build and Test StableHLO with Python
+--      shell: bash
+--      run: |
+--        pip install tensorflow-cpu nanobind==2.4
+--        pip install -r "$LLVM_PROJECT_DIR/mlir/python/requirements.txt"
+--        ./build_tools/github_actions/ci_build_cmake.sh "$LLVM_BUILD_DIR" "$STABLEHLO_BUILD_DIR"
+--      env:
+--          CMAKE_BUILD_TYPE: Release
+--          STABLEHLO_ENABLE_BINDINGS_PYTHON: ON
+--          STABLEHLO_ENABLE_PYTHON_TF_TESTS: ON
+--
+--    - name: Build and Test Python Wheel
+--      shell: bash
+--      run: |
+--        ./build_tools/github_actions/ci_build_python_wheel.sh "$GITHUB_WORKSPACE/${{ github.sha }}"
+--
+--    - name: Upload an artifact
+--      uses: actions/upload-artifact@v4
+--      with:
+--        if-no-files-found: error
+--        name: build_artifact-${{ matrix.python-version }}
+--        path: ${{ github.sha }}
+--
+--  upload-tarballs:
+--    runs-on: 'ubuntu-24.04'
+--    needs: cmake-build-wheel
+--    strategy:
+--      fail-fast: false
+--      matrix:
+--        python-version: [ "3.10", "3.11", "3.12" ]
+--
+--    steps:
+--      - name: Checkout
+--        uses: actions/checkout@v4
+--
+--      - name: Download artifacts
+--        uses: actions/download-artifact@v4
+--        with:
+--          name: build_artifact-${{ matrix.python-version }}
+--          path: ${{ github.sha }}
+--
+--      - name: Release current commit
+--        uses: ncipollo/release-action@v1.13.0
+--        with:
+--          artifacts: "${{ github.sha }}/*.whl"
+--          token: "${{ secrets.PUBLISH_ACCESS_TOKEN }}"
+--          # git describe in setup.py explicitly filters out this tag
+--          tag: "dev-wheels"
+--          name: "dev-wheels"
+--          removeArtifacts: false
+--          allowUpdates: true
+--          replacesArtifacts: true
+--          makeLatest: false
+-diff --ruN a/stablehlo/.gitignore b/stablehlo/.gitignore
+---- stablehlo/.gitignore
+-+++ stablehlo/.gitignore
+-@@ -1,14 +0,0 @@
+--bazel-bin
+--bazel-out
+--bazel-stablehlo
+--bazel-testlogs
+--build
+--cmake-build
+--llvm-project
+--llvm-build
+--.vscode
+--.cache
+--.idea/
+--**/__pycache__
+--*.whl
+--**/*.egg-info/
+++    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n4-16' ||  'ubuntu-22.04'  }}
++     container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cmake:latest"
++ 
++     steps:
+ diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel
+ --- stablehlo/BUILD.bazel
+ +++ stablehlo/BUILD.bazel
+@@ -3522,6 +2809,479 @@ diff --ruN a/stablehlo/stablehlo/integrations/cpp/builder/StablehloBuilderTest.c
+  TEST(MlirBuilderTest, GatherOp) {
+    std::string expected = R"mlir(module {
+    func.func @main(%arg0: tensor<3xi64>, %arg1: tensor<1x1xi64>) -> tensor<1xi64> {
++@@ -1021,7 +1097,7 @@
++ TEST(MlirBuilderTest, ConstantF4E2M1FNArray) {
++   std::string expected = R"mlir(module {
++   func.func @main() -> tensor<10xf4E2M1FN> {
++-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 0.000000e+00, 0.000000e+00, 3.000000e+00, 6.000000e+00, 6.000000e+00, 1.000000e+00, 6.000000e+00]> : tensor<10xf4E2M1FN>
+++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 0.000000e+00, 0.000000e+00, 3.000000e+00, 6.000000e+00, -6.000000e+00, 5.000000e-01, -5.000000e-01]> : tensor<10xf4E2M1FN>
++     return %cst : tensor<10xf4E2M1FN>
++   }
++ })mlir";
++@@ -1031,10 +1107,15 @@
++     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
++     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
++     auto type = makeTensorType(fb.getContext(), {10}, ElementType::F4E2M1FN);
+++    double posMax = 6.0;
+++    double negMax = -6.0;
+++    double posMinDenorm = 0.5;
+++    double negMinDenorm = -0.5;
++     auto cst = stablehlo::Constant(
++-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
++-                                           0x07, 0x0F, 0x01, 0x09}),
++-                         type));
+++        fb, makeConstant(
+++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
+++                                  negMax, posMinDenorm, negMinDenorm}),
+++                type));
++     func::Return(fb, {cst});
++   }
++ 
++@@ -1045,7 +1126,7 @@
++ TEST(MlirBuilderTest, ConstantF6E2M3FNArray) {
++   std::string expected = R"mlir(module {
++   func.func @main() -> tensor<10xf6E2M3FN> {
++-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.250000e-01, 3.250000e+00, 7.500000e+00, 7.500000e+00, 1.000000e+00, 7.500000e+00]> : tensor<10xf6E2M3FN>
+++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.250000e-01, 3.250000e+00, 7.500000e+00, -7.500000e+00, 1.250000e-01, -1.250000e-01]> : tensor<10xf6E2M3FN>
++     return %cst : tensor<10xf6E2M3FN>
++   }
++ })mlir";
++@@ -1055,10 +1136,15 @@
++     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
++     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
++     auto type = makeTensorType(fb.getContext(), {10}, ElementType::F6E2M3FN);
+++    double posMax = 7.5;
+++    double negMax = -7.5;
+++    double posMinDenorm = 0x1p-3;
+++    double negMinDenorm = -0x1p-3;
++     auto cst = stablehlo::Constant(
++-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
++-                                           0x1F, 0x3F, 0x01, 0x21}),
++-                         type));
+++        fb, makeConstant(
+++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
+++                                  negMax, posMinDenorm, negMinDenorm}),
+++                type));
++     func::Return(fb, {cst});
++   }
++ 
++@@ -1069,7 +1155,7 @@
++ TEST(MlirBuilderTest, ConstantF6E3M2FNArray) {
++   std::string expected = R"mlir(module {
++   func.func @main() -> tensor<10xf6E3M2FN> {
++-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.250000e-01, 3.000000e+00, 2.800000e+01, 2.800000e+01, 1.000000e+00, 2.800000e+01]> : tensor<10xf6E3M2FN>
+++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.250000e-01, 3.000000e+00, 2.800000e+01, -2.800000e+01, 6.250000e-02, -6.250000e-02]> : tensor<10xf6E3M2FN>
++     return %cst : tensor<10xf6E3M2FN>
++   }
++ })mlir";
++@@ -1079,10 +1165,15 @@
++     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
++     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
++     auto type = makeTensorType(fb.getContext(), {10}, ElementType::F6E3M2FN);
+++    double posMax = 28.0;
+++    double negMax = -28.0;
+++    double posMinDenorm = 0x1p-4;
+++    double negMinDenorm = -0x1p-4;
++     auto cst = stablehlo::Constant(
++-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
++-                                           0x1F, 0x3F, 0x01, 0x21}),
++-                         type));
+++        fb, makeConstant(
+++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
+++                                  negMax, posMinDenorm, negMinDenorm}),
+++                type));
++     func::Return(fb, {cst});
++   }
++ 
++@@ -1092,9 +1183,9 @@
++ 
++ TEST(MlirBuilderTest, ConstantF8E3M4Array) {
++   std::string expected = R"mlir(module {
++-  func.func @main() -> tensor<10xf8E3M4> {
++-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.125000e+00, 0x70, 0x70, 1.000000e+00, 0x70]> : tensor<10xf8E3M4>
++-    return %cst : tensor<10xf8E3M4>
+++  func.func @main() -> tensor<11xf8E3M4> {
+++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.125000e+00, 0x70, 0xF0, 0x78, 1.562500e-02, -1.562500e-02]> : tensor<11xf8E3M4>
+++    return %cst : tensor<11xf8E3M4>
++   }
++ })mlir";
++ 
++@@ -1102,11 +1193,17 @@
++   {
++     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
++     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
++-    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E3M4);
+++    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E3M4);
+++    double posInf = std::numeric_limits<double>::infinity();
+++    double negInf = -std::numeric_limits<double>::infinity();
+++    double nan = std::numeric_limits<double>::quiet_NaN();
+++    double posMinDenorm = 0x1p-6;
+++    double negMinDenorm = -0x1p-6;
++     auto cst = stablehlo::Constant(
++-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
++-                                           0x7F, 0xFF, 0x01, 0x81}),
++-                         type));
+++        fb, makeConstant(
+++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posInf,
+++                                  negInf, nan, posMinDenorm, negMinDenorm}),
+++                type));
++     func::Return(fb, {cst});
++   }
++ 
++@@ -1116,9 +1213,9 @@
++ 
++ TEST(MlirBuilderTest, ConstantF8E4M3B11FNUZArray) {
++   std::string expected = R"mlir(module {
++-  func.func @main() -> tensor<10xf8E4M3B11FNUZ> {
++-    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 0x80, 0x80, 1.000000e+00, 0x80]> : tensor<10xf8E4M3B11FNUZ>
++-    return %cst : tensor<10xf8E4M3B11FNUZ>
+++  func.func @main() -> tensor<11xf8E4M3B11FNUZ> {
+++    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 3.000000e+01, -3.000000e+01, 0x80, 1.220700e-04, -1.220700e-04]> : tensor<11xf8E4M3B11FNUZ>
+++    return %cst : tensor<11xf8E4M3B11FNUZ>
++   }
++ })mlir";
++ 
++@@ -1127,11 +1224,17 @@
++     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
++     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
++     auto type =
++-        makeTensorType(fb.getContext(), {10}, ElementType::F8E4M3B11FNUZ);
+++        makeTensorType(fb.getContext(), {11}, ElementType::F8E4M3B11FNUZ);
+++    double posMax = 30.0;
+++    double negMax = -30.0;
+++    double nan = std::numeric_limits<double>::quiet_NaN();
+++    double posMinDenorm = 0x1p-13;
+++    double negMinDenorm = -0x1p-13;
++     auto cst = stablehlo::Constant(
++-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
++-                                           0x7F, 0xFF, 0x01, 0x81}),
++-                         type));
+++        fb, makeConstant(
+++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
+++                                  negMax, nan, posMinDenorm, negMinDenorm}),
+++                type));
++     func::Return(fb, {cst});
++   }
++ 
++@@ -1141,9 +1244,9 @@
++ 
++ TEST(MlirBuilderTest, ConstantF8E4M3Array) {
++   std::string expected = R"mlir(module {
++-  func.func @main() -> tensor<10xf8E4M3> {
++-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 1.280000e+02, 0x78, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E4M3>
++-    return %cst : tensor<10xf8E4M3>
+++  func.func @main() -> tensor<11xf8E4M3> {
+++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 0x78, 0xF8, 0x7C, 1.953130e-03, -1.953130e-03]> : tensor<11xf8E4M3>
+++    return %cst : tensor<11xf8E4M3>
++   }
++ })mlir";
++ 
++@@ -1151,11 +1254,17 @@
++   {
++     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
++     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
++-    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E4M3);
+++    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E4M3);
+++    double posInf = std::numeric_limits<double>::infinity();
+++    double negInf = -std::numeric_limits<double>::infinity();
+++    double nan = std::numeric_limits<double>::quiet_NaN();
+++    double posMinDenorm = 0x1p-9;
+++    double negMinDenorm = -0x1p-9;
++     auto cst = stablehlo::Constant(
++-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
++-                                           0x7F, 0xFF, 0x01, 0x81}),
++-                         type));
+++        fb, makeConstant(
+++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posInf,
+++                                  negInf, nan, posMinDenorm, negMinDenorm}),
+++                type));
++     func::Return(fb, {cst});
++   }
++ 
++@@ -1165,9 +1274,9 @@
++ 
++ TEST(MlirBuilderTest, ConstantF8E4M3FNArray) {
++   std::string expected = R"mlir(module {
++-  func.func @main() -> tensor<10xf8E4M3FN> {
++-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 1.280000e+02, 2.560000e+02, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E4M3FN>
++-    return %cst : tensor<10xf8E4M3FN>
+++  func.func @main() -> tensor<11xf8E4M3FN> {
+++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 4.480000e+02, -4.480000e+02, 0x7F, 1.953130e-03, -1.953130e-03]> : tensor<11xf8E4M3FN>
+++    return %cst : tensor<11xf8E4M3FN>
++   }
++ })mlir";
++ 
++@@ -1175,11 +1284,17 @@
++   {
++     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
++     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
++-    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E4M3FN);
+++    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E4M3FN);
+++    double posMax = 448.0;
+++    double negMax = -448.0;
+++    double nan = std::numeric_limits<double>::quiet_NaN();
+++    double posMinDenorm = 0x1p-9;
+++    double negMinDenorm = -0x1p-9;
++     auto cst = stablehlo::Constant(
++-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
++-                                           0x7F, 0xFF, 0x01, 0x81}),
++-                         type));
+++        fb, makeConstant(
+++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
+++                                  negMax, nan, posMinDenorm, negMinDenorm}),
+++                type));
++     func::Return(fb, {cst});
++   }
++ 
++@@ -1189,9 +1304,9 @@
++ 
++ TEST(MlirBuilderTest, ConstantF8E4M3FNUZArray) {
++   std::string expected = R"mlir(module {
++-  func.func @main() -> tensor<10xf8E4M3FNUZ> {
++-    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 1.280000e+02, 0x80, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E4M3FNUZ>
++-    return %cst : tensor<10xf8E4M3FNUZ>
+++  func.func @main() -> tensor<11xf8E4M3FNUZ> {
+++    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 2.400000e+02, -2.400000e+02, 0x80, 9.765620e-04, -9.765620e-04]> : tensor<11xf8E4M3FNUZ>
+++    return %cst : tensor<11xf8E4M3FNUZ>
++   }
++ })mlir";
++ 
++@@ -1199,11 +1314,17 @@
++   {
++     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
++     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
++-    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E4M3FNUZ);
+++    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E4M3FNUZ);
+++    double posMax = 240.0;
+++    double negMax = -240.0;
+++    double nan = std::numeric_limits<double>::quiet_NaN();
+++    double posMinDenorm = 0x1p-10;
+++    double negMinDenorm = -0x1p-10;
++     auto cst = stablehlo::Constant(
++-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
++-                                           0x7F, 0xFF, 0x01, 0x81}),
++-                         type));
+++        fb, makeConstant(
+++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
+++                                  negMax, nan, posMinDenorm, negMinDenorm}),
+++                type));
++     func::Return(fb, {cst});
++   }
++ 
++@@ -1213,9 +1334,9 @@
++ 
++ TEST(MlirBuilderTest, ConstantF8E5M2Array) {
++   std::string expected = R"mlir(module {
++-  func.func @main() -> tensor<10xf8E5M2> {
++-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.000000e+00, 1.280000e+02, 2.560000e+02, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E5M2>
++-    return %cst : tensor<10xf8E5M2>
+++  func.func @main() -> tensor<11xf8E5M2> {
+++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.000000e+00, 0x7C, 0xFC, 0x7E, 1.525880e-05, -1.525880e-05]> : tensor<11xf8E5M2>
+++    return %cst : tensor<11xf8E5M2>
++   }
++ })mlir";
++ 
++@@ -1223,11 +1344,17 @@
++   {
++     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
++     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
++-    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E5M2);
+++    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E5M2);
+++    double posInf = std::numeric_limits<double>::infinity();
+++    double negInf = -std::numeric_limits<double>::infinity();
+++    double nan = std::numeric_limits<double>::quiet_NaN();
+++    double posMinDenorm = 0x1p-16;
+++    double negMinDenorm = -0x1p-16;
++     auto cst = stablehlo::Constant(
++-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
++-                                           0x7F, 0xFF, 0x01, 0x81}),
++-                         type));
+++        fb, makeConstant(
+++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posInf,
+++                                  negInf, nan, posMinDenorm, negMinDenorm}),
+++                type));
++     func::Return(fb, {cst});
++   }
++ 
++@@ -1237,9 +1364,9 @@
++ 
++ TEST(MlirBuilderTest, ConstantF8E5M2FNUZArray) {
++   std::string expected = R"mlir(module {
++-  func.func @main() -> tensor<10xf8E5M2FNUZ> {
++-    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.000000e+00, 1.280000e+02, 2.560000e+02, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E5M2FNUZ>
++-    return %cst : tensor<10xf8E5M2FNUZ>
+++  func.func @main() -> tensor<11xf8E5M2FNUZ> {
+++    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.000000e+00, 5.734400e+04, -5.734400e+04, 0x80, 7.629390e-06, -7.629390e-06]> : tensor<11xf8E5M2FNUZ>
+++    return %cst : tensor<11xf8E5M2FNUZ>
++   }
++ })mlir";
++ 
++@@ -1247,11 +1374,17 @@
++   {
++     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
++     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
++-    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E5M2FNUZ);
+++    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E5M2FNUZ);
+++    double posMax = 57344.0;
+++    double negMax = -57344.0;
+++    double nan = std::numeric_limits<double>::quiet_NaN();
+++    double posMinDenorm = 0x1p-17;
+++    double negMinDenorm = -0x1p-17;
++     auto cst = stablehlo::Constant(
++-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
++-                                           0x7F, 0xFF, 0x01, 0x81}),
++-                         type));
+++        fb, makeConstant(
+++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
+++                                  negMax, nan, posMinDenorm, negMinDenorm}),
+++                type));
++     func::Return(fb, {cst});
++   }
++ 
++@@ -1262,7 +1395,7 @@
++ TEST(MlirBuilderTest, ConstantF8E8M0FNUArray) {
++   std::string expected = R"mlir(module {
++   func.func @main() -> tensor<8xf8E8M0FNU> {
++-    %cst = stablehlo.constant dense<[5.877470e-39, 1.000000e+00, 1.250000e-01, 1.250000e-01, 4.000000e+00, 5.877470e-39, 1.280000e+02, 2.560000e+02]> : tensor<8xf8E8M0FNU>
+++    %cst = stablehlo.constant dense<[5.877470e-39, 1.000000e+00, 1.250000e-01, 1.250000e-01, 4.000000e+00, 5.877470e-39, 2.000000e+00, 0xFF]> : tensor<8xf8E8M0FNU>
++     return %cst : tensor<8xf8E8M0FNU>
++   }
++ })mlir";
++@@ -1272,10 +1405,12 @@
++     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
++     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
++     auto type = makeTensorType(fb.getContext(), {8}, ElementType::F8E8M0FNU);
+++    double min = 0x1p-127;
+++    double nan = std::numeric_limits<double>::quiet_NaN();
++     auto cst = stablehlo::Constant(
++-        fb, makeConstant(ArrayRef<double>(
++-                             {0.0, 1.0, 0.125, 0.1, 3.1415, 0x00, 0x80, 0xFF}),
++-                         type));
+++        fb, makeConstant(
+++                ArrayRef<double>({0.0, 1.0, 0.125, 0.1, 3.1415, min, 2.0, nan}),
+++                type));
++     func::Return(fb, {cst});
++   }
++ 
++@@ -1286,7 +1421,7 @@
++ TEST(MlirBuilderTest, ConstantBF16Array) {
++   std::string expected = R"mlir(module {
++   func.func @main() -> tensor<11xbf16> {
++-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000980e-01, 3.140630e+00, 3.264000e+04, 6.553600e+04, 3.276800e+04, 1.000000e+00, 3.276800e+04]> : tensor<11xbf16>
+++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000980e-01, 3.140630e+00, 0x7F80, 0xFF80, 0x7FC0, 9.183550e-41, -9.183550e-41]> : tensor<11xbf16>
++     return %cst : tensor<11xbf16>
++   }
++ })mlir";
++@@ -1296,11 +1431,16 @@
++     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
++     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
++     auto type = makeTensorType(fb.getContext(), {11}, ElementType::BF16);
+++    double posInf = std::numeric_limits<double>::infinity();
+++    double negInf = -std::numeric_limits<double>::infinity();
+++    double nan = std::numeric_limits<double>::quiet_NaN();
+++    double posMinDenorm = 0x1p-133;
+++    double negMinDenorm = -0x1p-133;
++     auto cst = stablehlo::Constant(
++-        fb,
++-        makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.140630,
++-                                       0x7F80, 0xFF80, 0x7FFF, 0x0001, 0x8001}),
++-                     type));
+++        fb, makeConstant(
+++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.140630, posInf,
+++                                  negInf, nan, posMinDenorm, negMinDenorm}),
+++                type));
++     func::Return(fb, {cst});
++   }
++ 
++@@ -1311,7 +1451,7 @@
++ TEST(MlirBuilderTest, ConstantF16Array) {
++   std::string expected = R"mlir(module {
++   func.func @main() -> tensor<11xf16> {
++-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.997550e-02, 3.140630e+00, 3.174400e+04, 6.451200e+04, 3.276800e+04, 1.000000e+00, 3.276800e+04]> : tensor<11xf16>
+++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.997550e-02, 3.140630e+00, 0x7C00, 0xFC00, 0x7E00, 5.960460e-08, -5.960460e-08]> : tensor<11xf16>
++     return %cst : tensor<11xf16>
++   }
++ })mlir";
++@@ -1321,11 +1461,16 @@
++     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
++     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
++     auto type = makeTensorType(fb.getContext(), {11}, ElementType::F16);
+++    double posInf = std::numeric_limits<double>::infinity();
+++    double negInf = -std::numeric_limits<double>::infinity();
+++    double nan = std::numeric_limits<double>::quiet_NaN();
+++    double posMinDenorm = 0x1p-24;
+++    double negMinDenorm = -0x1p-24;
++     auto cst = stablehlo::Constant(
++-        fb,
++-        makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.140630,
++-                                       0x7C00, 0xFC00, 0x7FFF, 0x0001, 0x8001}),
++-                     type));
+++        fb, makeConstant(
+++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.140630, posInf,
+++                                  negInf, nan, posMinDenorm, negMinDenorm}),
+++                type));
++     func::Return(fb, {cst});
++   }
++ 
++@@ -1336,7 +1481,7 @@
++ TEST(MlirBuilderTest, ConstantF32Array) {
++   std::string expected = R"mlir(module {
++   func.func @main() -> tensor<11xf32> {
++-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000000e-01, 3.14159274, 2.13909504E+9, 4.28657869E+9, 2.14748365E+9, 1.000000e+00, 2.14748365E+9]> : tensor<11xf32>
+++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000000e-01, 3.14159274, 0x7F800000, 0xFF800000, 0x7FC00000, 1.401300e-45, -1.401300e-45]> : tensor<11xf32>
++     return %cst : tensor<11xf32>
++   }
++ })mlir";
++@@ -1346,10 +1491,15 @@
++     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
++     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
++     auto type = makeTensorType(fb.getContext(), {11}, ElementType::F32);
+++    double posInf = std::numeric_limits<double>::infinity();
+++    double negInf = -std::numeric_limits<double>::infinity();
+++    double nan = std::numeric_limits<double>::quiet_NaN();
+++    double posMinDenorm = std::numeric_limits<float>::denorm_min();
+++    double negMinDenorm = -std::numeric_limits<float>::denorm_min();
++     auto cst = stablehlo::Constant(
++         fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1,
++-                                           3.14159274, 0x7F800000, 0xFF800000,
++-                                           0x7FFFFFFF, 0x00000001, 0x80000001}),
+++                                           3.14159274, posInf, negInf, nan,
+++                                           posMinDenorm, negMinDenorm}),
++                          type));
++     func::Return(fb, {cst});
++   }
++@@ -1361,7 +1511,7 @@
++ TEST(MlirBuilderTest, ConstantF64Array) {
++   std::string expected = R"mlir(module {
++   func.func @main() -> tensor<11xf64> {
++-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000000e-01, 3.1415926535897931, 9.2188684372274053E+18, 1.8442240474082181E+19, 9.2233720368547758E+18, 1.000000e+00, 9.2233720368547758E+18]> : tensor<11xf64>
+++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000000e-01, 3.1415926535897931, 0x7FF0000000000000, 0xFFF0000000000000, 0x7FF8000000000000, 4.940660e-324, -4.940660e-324]> : tensor<11xf64>
++     return %cst : tensor<11xf64>
++   }
++ })mlir";
++@@ -1371,13 +1521,16 @@
++     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
++     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
++     auto type = makeTensorType(fb.getContext(), {11}, ElementType::F64);
+++    double posInf = std::numeric_limits<double>::infinity();
+++    double negInf = -std::numeric_limits<double>::infinity();
+++    double nan = std::numeric_limits<double>::quiet_NaN();
+++    double posMinDenorm = std::numeric_limits<double>::denorm_min();
+++    double negMinDenorm = -std::numeric_limits<double>::denorm_min();
++     auto cst = stablehlo::Constant(
++-        fb, makeConstant(
++-                ArrayRef<long double>({0.0, -0.0, 1.0, 0.125, 0.1,
++-                                       3.1415926535897931, 0x7FF0000000000000,
++-                                       0xFFF0000000000000, 0x7FFFFFFFFFFFFFFF,
++-                                       0x0000000000000001, 0x8000000000000001}),
++-                type));
+++        fb, makeConstant(ArrayRef<long double>(
+++                             {0.0, -0.0, 1.0, 0.125, 0.1, 3.1415926535897931,
+++                              posInf, negInf, nan, posMinDenorm, negMinDenorm}),
+++                         type));
++     func::Return(fb, {cst});
++   }
++ 
+ diff --ruN a/stablehlo/stablehlo/integrations/python/StablehloApi.cpp b/stablehlo/stablehlo/integrations/python/StablehloApi.cpp
+ --- stablehlo/stablehlo/integrations/python/StablehloApi.cpp
+ +++ stablehlo/stablehlo/integrations/python/StablehloApi.cpp
+@@ -4182,6 +3942,77 @@ diff --ruN a/stablehlo/stablehlo/reference/Ops.cpp b/stablehlo/stablehlo/referen
+      Axes resultPermutation;
+      resultPermutation.push_back(outputBatchDimension);
+      resultPermutation.append(outputSpatialDimensions.begin(),
++@@ -1942,23 +1969,30 @@
++ Tensor __fft_1d(const Tensor& operand, ShapedType resultType,
++                 int64_t dimension) {
++   Tensor result(resultType);
++-
+++  auto complexTy = cast<ComplexType>(resultType.getElementType());
+++  const llvm::fltSemantics& sem =
+++      cast<FloatType>(complexTy.getElementType()).getFloatSemantics();
++   for (auto resultIt = result.index_begin(); resultIt != result.index_end();
++        ++resultIt) {
++     auto resultIndex = *resultIt;
++     auto N = operand.getShape()[dimension];
++ 
++-    Element X(resultType.getElementType(),
++-              mlir::Complex<APFloat>(APFloat(0.0), APFloat(0.0)));
+++    Element X(
+++        resultType.getElementType(),
+++        mlir::Complex<APFloat>(APFloat::getZero(sem), APFloat::getZero(sem)));
++ 
++     for (int n = 0; n < N; n++) {
++       auto operandIndex = Index(resultIndex);
++       operandIndex[dimension] = n;
++ 
++-      auto floatType =
++-          cast<ComplexType>(operand.getElementType()).getElementType();
++-      Element phase(floatType,
++-                    APFloat(-2 * M_PI * n * resultIndex[dimension] / N));
+++      auto floatType = cast<FloatType>(
+++          cast<ComplexType>(operand.getElementType()).getElementType());
+++
+++      auto phaseVal = APFloat(-2 * M_PI * n * resultIndex[dimension] / N);
+++      bool losesInfo;
+++      phaseVal.convert(floatType.getFloatSemantics(),
+++                       APFloat::rmNearestTiesToEven, &losesInfo);
+++      Element phase(floatType, phaseVal);
++       auto rotator = complex(cosine(phase), sine(phase));
++       X = X + operand.get(operandIndex) * rotator;
++     }
++@@ -1992,10 +2026,12 @@
++   } else if (fftType == FftType::IFFT) {
++     // compute IFFT as conjugation, FFT, conjugation and normalization
++     Tensor result(operand);
+++    auto complexTy = cast<ComplexType>(resultType.getElementType());
+++    const llvm::fltSemantics& sem =
+++        cast<FloatType>(complexTy.getElementType()).getFloatSemantics();
++     for (auto d = fftDims.begin(); d != fftDims.end(); ++d) {
++       auto divisorValue = mlir::Complex<APFloat>(
++-          APFloat(static_cast<double>(resultType.getDimSize(*d))),
++-          APFloat(0.0));
+++          APFloat(sem, resultType.getDimSize(*d)), APFloat::getZero(sem));
++       auto divisor =
++           constantOp(mlir::DenseElementsAttr::get(resultType, divisorValue));
++       result = divideOp(conjugate(__fft_1d(conjugate(result), resultType, *d)),
++@@ -2040,6 +2076,8 @@
++     replica = conjugate(reverseOp(replica, Axes{d_irfft}, replica.getType()));
++     auto complexType =
++         ComplexType::get(cast<FloatType>(resultType.getElementType()));
+++    const llvm::fltSemantics& sem =
+++        cast<FloatType>(complexType.getElementType()).getFloatSemantics();
++     result = concatenateOp({result, replica}, Axis{d_irfft},
++                            resultType.cloneWith(std::nullopt, complexType));
++ 
++@@ -2048,7 +2086,7 @@
++                     resultType);
++ 
++     // normalize result of inverse transformation
++-    auto divisorValue = APFloat(static_cast<double>(renormFactor));
+++    auto divisorValue = APFloat(sem, renormFactor);
++     auto divisor =
++         constantOp(mlir::DenseElementsAttr::get(resultType, divisorValue));
++     return divideOp(result, divisor, resultType);
+ diff --ruN a/stablehlo/stablehlo/reference/Ops.h b/stablehlo/stablehlo/reference/Ops.h
+ --- stablehlo/stablehlo/reference/Ops.h
+ +++ stablehlo/stablehlo/reference/Ops.h
+@@ -4828,6 +4659,249 @@ diff --ruN a/stablehlo/stablehlo/tests/interpret/convolution.mlir b/stablehlo/st
+ +    dense<[[[3.0], [4.0]], [[6.0], [8.0]]]> : tensor<2x2x1xf32>
+ +  func.return
+ +}
++diff --ruN a/stablehlo/stablehlo/tests/interpret/fft.mlir b/stablehlo/stablehlo/tests/interpret/fft.mlir
++--- stablehlo/stablehlo/tests/interpret/fft.mlir
+++++ stablehlo/stablehlo/tests/interpret/fft.mlir
++@@ -233,3 +233,239 @@
++   > : tensor<2x2x4xf64>
++   func.return
++ }
+++
+++// -----
+++
+++func.func @fft_op_test_fft_1d_c64() {
+++  %operand = stablehlo.constant dense<[(1.0, 2.0), (3.0, -4.0)]> : tensor<2xcomplex<f32>>
+++  %result = stablehlo.fft %operand, type = FFT, length = [2] : (tensor<2xcomplex<f32>>) -> tensor<2xcomplex<f32>>
+++  check.expect_almost_eq_const %result, dense<[(4.0, -2.0), (-2.0, 6.0)]> : tensor<2xcomplex<f32>>
+++  func.return
+++}
+++
+++// -----
+++
+++func.func @fft_op_test_fft_2d_c64() {
+++  %operand = stablehlo.constant dense<
+++    [
+++      [(1.0, 2.0), (3.0, -4.0)],
+++      [(-5.0, 6.0), (-7.0, -8.0)]
+++    ]
+++  > : tensor<2x2xcomplex<f32>>
+++  %result = stablehlo.fft %operand, type = FFT, length = [2, 2] : (tensor<2x2xcomplex<f32>>) -> tensor<2x2xcomplex<f32>>
+++  check.expect_almost_eq_const %result, dense<
+++    [
+++      [(-8.0, -4.0), (0.0, 20.0)],
+++      [(16.0, 0.0), (-4.0, -8.0)]
+++    ]
+++  > : tensor<2x2xcomplex<f32>>
+++  func.return
+++}
+++
+++// -----
+++
+++func.func @fft_op_test_fft_3d_c64() {
+++  %operand = stablehlo.constant dense<
+++    [
+++      [
+++        [(1.0, 2.0), (3.0, -4.0)],
+++        [(-5.0, 6.0), (-7.0, -8.0)]
+++      ],[
+++        [(-1.0, -2.0), (-3.0, 4.0)],
+++        [(5.0, -6.0), (7.0, 8.0)]
+++      ]
+++    ]
+++  > : tensor<2x2x2xcomplex<f32>>
+++  %result = stablehlo.fft %operand, type = FFT, length = [2, 2, 2] : (tensor<2x2x2xcomplex<f32>>) -> tensor<2x2x2xcomplex<f32>>
+++  check.expect_almost_eq_const %result, dense<
+++    [
+++      [
+++        [(0.0, 0.0), (0.0, 0.0)],
+++        [(0.0, 0.0), (0.0, 0.0)]
+++      ],[
+++        [(-16.0, -8.0), (0.0, 40.0)],
+++        [(32.0, 0.0), (-8.0, -16.0)]
+++      ]
+++    ]
+++  > : tensor<2x2x2xcomplex<f32>>
+++  func.return
+++}
+++
+++// -----
+++
+++func.func @fft_op_test_ifft_1d_c64() {
+++  %operand = stablehlo.constant dense<[(4.0, -2.0), (-2.0, 6.0)]> : tensor<2xcomplex<f32>>
+++  %result = stablehlo.fft %operand, type = IFFT, length = [2] : (tensor<2xcomplex<f32>>) -> tensor<2xcomplex<f32>>
+++  check.expect_almost_eq_const %result, dense<[(1.0, 2.0), (3.0, -4.0)]> : tensor<2xcomplex<f32>>
+++  func.return
+++}
+++
+++// -----
+++
+++func.func @fft_op_test_ifft_2d_c64() {
+++  %operand = stablehlo.constant dense<
+++    [
+++      [(-8.0, -4.0), (16.0, 0.0)],
+++      [(0.0, 20.0), (-4.0, -8.0)]
+++    ]
+++  > : tensor<2x2xcomplex<f32>>
+++  %result = stablehlo.fft %operand, type = IFFT, length = [2, 2] : (tensor<2x2xcomplex<f32>>) -> tensor<2x2xcomplex<f32>>
+++  check.expect_almost_eq_const %result, dense<
+++    [
+++      [(1.0, 2.0), (-5.0, 6.0)],
+++      [(3.0, -4.0), (-7.0, -8.0)]
+++    ]
+++  > : tensor<2x2xcomplex<f32>>
+++  func.return
+++}
+++
+++// -----
+++
+++func.func @fft_op_test_ifft_3d_c64() {
+++  %operand = stablehlo.constant dense<
+++    [
+++      [
+++        [(0.0, 0.0), (0.0, 0.0)],
+++        [(0.0, 0.0), (0.0, 0.0)]
+++      ],[
+++        [(-16.0, -8.0), (32.0, 0.0)],
+++        [(0.0, 40.0), (-8.0, -16.0)]
+++      ]
+++    ]
+++  > : tensor<2x2x2xcomplex<f32>>
+++  %result = stablehlo.fft %operand, type = IFFT, length = [2, 2, 2] : (tensor<2x2x2xcomplex<f32>>) -> tensor<2x2x2xcomplex<f32>>
+++  check.expect_almost_eq_const %result, dense<
+++    [
+++      [
+++        [(1.0, 2.0), (-5.0, 6.0)],
+++        [(3.0, -4.0), (-7.0, -8.0)]
+++      ],[
+++        [(-1.0, -2.0), (5.0, -6.0)],
+++        [(-3.0, 4.0), (7.0, 8.0)]
+++      ]
+++    ]
+++  > : tensor<2x2x2xcomplex<f32>>
+++  func.return
+++}
+++
+++// -----
+++
+++func.func @fft_op_test_rfft_1d_f32() {
+++  %operand = stablehlo.constant dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
+++  %result = stablehlo.fft %operand, type = RFFT, length = [4] : (tensor<4xf32>) -> tensor<3xcomplex<f32>>
+++  check.expect_almost_eq_const %result, dense<[(10.0, 0.0), (-2.0, 2.0), (-2.0, 0.0)]> : tensor<3xcomplex<f32>>
+++  func.return
+++}
+++
+++// -----
+++
+++func.func @fft_op_test_rfft_2d_f32() {
+++  %operand = stablehlo.constant dense<
+++    [
+++      [1.0, 2.0, 3.0, 4.0],
+++      [5.0, 6.0, 7.0, 8.0],
+++      [9.0, 10.0, 11.0, 12.0],
+++      [13.0, 14.0, 15.0, 16.0]
+++    ]
+++  > : tensor<4x4xf32>
+++  %result = stablehlo.fft %operand, type = RFFT, length = [4, 4] : (tensor<4x4xf32>) -> tensor<4x3xcomplex<f32>>
+++  check.expect_almost_eq_const %result, dense<
+++    [
+++      [(136.0, 0.0), (-8.0, 8.0), (-8.0, 0.0)],
+++      [(-32.0, 32.0), (0.0, 0.0), (0.0, 0.0)],
+++      [(-32.0, 0.0), (0.0, 0.0), (0.0, 0.0)],
+++      [(-32.0, -32.0), (0.0, 0.0), (0.0, 0.0)]
+++    ]
+++  > : tensor<4x3xcomplex<f32>>
+++  func.return
+++}
+++
+++// -----
+++
+++func.func @fft_op_test_rfft_3d_f32() {
+++  %operand = stablehlo.constant dense<
+++    [
+++      [
+++        [1.0, 2.0, 3.0, 4.0],
+++        [5.0, 6.0, 7.0, 8.0]
+++      ],[
+++        [9.0, 10.0, 11.0, 12.0],
+++        [13.0, 14.0, 15.0, 16.0]
+++      ]
+++    ]
+++  > : tensor<2x2x4xf32>
+++  %result = stablehlo.fft %operand, type = RFFT, length = [2, 2, 4] : (tensor<2x2x4xf32>) -> tensor<2x2x3xcomplex<f32>>
+++  check.expect_almost_eq_const %result, dense<
+++    [
+++      [
+++        [(136.0, 0.0), (-8.0, 8.0), (-8.0, 0.0)],
+++        [(-32.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
+++      ],[
+++        [(-64.0, 0.0), (0.0, 0.0), (0.0, 0.0)],
+++        [(0.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
+++      ]
+++    ]
+++  > : tensor<2x2x3xcomplex<f32>>
+++  func.return
+++}
+++
+++// -----
+++
+++func.func @fft_op_test_irfft_1d_c64() {
+++  %operand = stablehlo.constant dense<[(10.0, 0.0), (-2.0, 2.0), (-2.0, 0.0)]> : tensor<3xcomplex<f32>>
+++  %result = stablehlo.fft %operand, type = IRFFT, length = [4] : (tensor<3xcomplex<f32>>) -> tensor<4xf32>
+++  check.expect_almost_eq_const %result, dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
+++  func.return
+++}
+++
+++// -----
+++
+++func.func @fft_op_test_irfft_2d_c64() {
+++  %operand = stablehlo.constant dense<
+++    [
+++      [(136.0, 0.0), (-8.0, 8.0), (-8.0, 0.0)],
+++      [(-32.0, 32.0), (0.0, 0.0), (0.0, 0.0)],
+++      [(-32.0, 0.0), (0.0, 0.0), (0.0, 0.0)],
+++      [(-32.0, -32.0), (0.0, 0.0), (0.0, 0.0)]
+++    ]
+++  > : tensor<4x3xcomplex<f32>>
+++  %result = stablehlo.fft %operand, type = IRFFT, length = [4, 4] : (tensor<4x3xcomplex<f32>>) -> tensor<4x4xf32>
+++  check.expect_almost_eq_const %result, dense<
+++    [
+++      [1.0, 2.0, 3.0, 4.0],
+++      [5.0, 6.0, 7.0, 8.0],
+++      [9.0, 10.0, 11.0, 12.0],
+++      [13.0, 14.0, 15.0, 16.0]
+++    ]
+++  > : tensor<4x4xf32>
+++  func.return
+++}
+++
+++// -----
+++
+++func.func @fft_op_test_irfft_3d_c64() {
+++  %operand = stablehlo.constant dense<
+++    [
+++      [
+++        [(136.0, 0.0), (-8.0, 8.0), (-8.0, 0.0)],
+++        [(-32.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
+++      ],[
+++        [(-64.0, 0.0), (0.0, 0.0), (0.0, 0.0)],
+++        [(0.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
+++      ]
+++    ]
+++  > : tensor<2x2x3xcomplex<f32>>
+++  %result = stablehlo.fft %operand, type = IRFFT, length = [2, 2, 4] : (tensor<2x2x3xcomplex<f32>>) -> tensor<2x2x4xf32>
+++  check.expect_almost_eq_const %result, dense<
+++    [
+++      [
+++        [1.0, 2.0, 3.0, 4.0],
+++        [5.0, 6.0, 7.0, 8.0]
+++      ],[
+++        [9.0, 10.0, 11.0, 12.0],
+++        [13.0, 14.0, 15.0, 16.0]
+++      ]
+++    ]
+++  > : tensor<2x2x4xf32>
+++  func.return
+++}
+ diff --ruN a/stablehlo/stablehlo/tests/ops_chlo.mlir b/stablehlo/stablehlo/tests/ops_chlo.mlir
+ --- stablehlo/stablehlo/tests/ops_chlo.mlir
+ +++ stablehlo/stablehlo/tests/ops_chlo.mlir
+@@ -13148,363 +13222,375 @@ diff --ruN a/stablehlo/third_party/BUILD b/stablehlo/third_party/BUILD
+ @@ -1 +0,0 @@
+ -licenses(["notice"])  # Apache 2.0
+ \ No newline at end of file
+-diff --ruN a/stablehlo/third_party/repo.bzl b/stablehlo/third_party/repo.bzl
+---- stablehlo/third_party/repo.bzl
+-+++ stablehlo/third_party/repo.bzl
+-@@ -1,159 +0,0 @@
+--# Copyright 2017 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.
+--
+--"""Utilities for defining TensorFlow Bazel dependencies."""
+--
+--def tf_mirror_urls(url):
+--    """A helper for generating TF-mirror versions of URLs.
+--
+--    Given a URL, it returns a list of the TF-mirror cache version of that URL
+--    and the original URL, suitable for use in `urls` field of `tf_http_archive`.
+--    """
+--    if not url.startswith("https://"):
+--        return [url]
+--    return [
+--        "https://storage.googleapis.com/mirror.tensorflow.org/%s" % url[8:],
+--        url,
+--    ]
+--
+--def _get_env_var(ctx, name):
+--    if name in ctx.os.environ:
+--        return ctx.os.environ[name]
+--    else:
+--        return None
+--
+--# Checks if we should use the system lib instead of the bundled one
+--def _use_system_lib(ctx, name):
+--    syslibenv = _get_env_var(ctx, "TF_SYSTEM_LIBS")
+--    if not syslibenv:
+--        return False
+--    return name in [n.strip() for n in syslibenv.split(",")]
+--
+--def _get_link_dict(ctx, link_files, build_file):
+--    link_dict = {ctx.path(v): ctx.path(Label(k)) for k, v in link_files.items()}
+--    if build_file:
+--        # Use BUILD.bazel because it takes precedence over BUILD.
+--        link_dict[ctx.path("BUILD.bazel")] = ctx.path(Label(build_file))
+--    return link_dict
+--
+--def _tf_http_archive_impl(ctx):
+--    # Construct all paths early on to prevent rule restart. We want the
+--    # attributes to be strings instead of labels because they refer to files
+--    # in the TensorFlow repository, not files in repos depending on TensorFlow.
+--    # See also https://github.com/bazelbuild/bazel/issues/10515.
+--    link_dict = _get_link_dict(ctx, ctx.attr.link_files, ctx.attr.build_file)
+--
+--    # For some reason, we need to "resolve" labels once before the
+--    # download_and_extract otherwise it'll invalidate and re-download the
+--    # archive each time.
+--    # https://github.com/bazelbuild/bazel/issues/10515
+--    patch_files = ctx.attr.patch_file
+--    for patch_file in patch_files:
+--        if patch_file:
+--            ctx.path(Label(patch_file))
+--
+--    if _use_system_lib(ctx, ctx.attr.name):
+--        link_dict.update(_get_link_dict(
+--            ctx = ctx,
+--            link_files = ctx.attr.system_link_files,
+--            build_file = ctx.attr.system_build_file,
+--        ))
+--    else:
+--        ctx.download_and_extract(
+--            url = ctx.attr.urls,
+--            sha256 = ctx.attr.sha256,
+--            type = ctx.attr.type,
+--            stripPrefix = ctx.attr.strip_prefix,
+--        )
+--        if hasattr(ctx.attr, "patch_cmds") and ctx.attr.patch_cmds:
+--            for cmd in ctx.attr.patch_cmds:
+--                res = ctx.execute(["bash", "-c", cmd])
+--                if res.return_code != 0:
+--                    fail("patch_cmd failed: %s\n%s" % (cmd, res.stderr))
+--        if patch_files:
+--            for patch_file in patch_files:
+--                patch_file = ctx.path(Label(patch_file)) if patch_file else None
+--                if patch_file:
+--                    ctx.patch(patch_file, strip = 1)
+--
+--    for dst, src in link_dict.items():
+--        ctx.delete(dst)
+--        ctx.symlink(src, dst)
+--
+--_tf_http_archive = repository_rule(
+--    implementation = _tf_http_archive_impl,
+--    attrs = {
+--        "sha256": attr.string(mandatory = True),
+--        "urls": attr.string_list(mandatory = True),
+--        "strip_prefix": attr.string(),
+--        "type": attr.string(),
+--        "patch_file": attr.string_list(),
+--        "patch_cmds": attr.string_list(),
+--        "build_file": attr.string(),
+--        "system_build_file": attr.string(),
+--        "link_files": attr.string_dict(),
+--        "system_link_files": attr.string_dict(),
+--    },
+--    environ = ["TF_SYSTEM_LIBS"],
+--)
+--
+--def tf_http_archive(name, sha256, urls, **kwargs):
+--    """Downloads and creates Bazel repos for dependencies.
+--
+--    This is a swappable replacement for both http_archive() and
+--    new_http_archive() that offers some additional features. It also helps
+--    ensure best practices are followed.
+--
+--    File arguments are relative to the TensorFlow repository by default. Dependent
+--    repositories that use this rule should refer to files either with absolute
+--    labels (e.g. '@foo//:bar') or from a label created in their repository (e.g.
+--    'str(Label("//:bar"))').
+--    """
+--    if len(urls) < 2:
+--        fail("tf_http_archive(urls) must have redundant URLs.")
+--
+--    if not any([mirror in urls[0] for mirror in (
+--        "mirror.tensorflow.org",
+--        "mirror.bazel.build",
+--        "storage.googleapis.com",
+--    )]):
+--        fail("The first entry of tf_http_archive(urls) must be a mirror " +
+--             "URL, preferrably mirror.tensorflow.org. Even if you don't have " +
+--             "permission to mirror the file, please put the correctly " +
+--             "formatted mirror URL there anyway, because someone will come " +
+--             "along shortly thereafter and mirror the file.")
+--
+--    if native.existing_rule(name):
+--        return
+--
+--    _tf_http_archive(
+--        name = name,
+--        sha256 = sha256,
+--        urls = urls,
+--        **kwargs
+--    )
+--
+--def _tf_vendored_impl(ctx):
+--    ctx.symlink(ctx.path(ctx.attr._root).dirname.get_child(ctx.attr.path), ".")
+--
+--tf_vendored = repository_rule(
+--    implementation = _tf_vendored_impl,
+--    doc = "Similar to local_repository, but path is relative to the root of the " +
+--          "repository, not the root of the workspace.",
+--    attrs = {
+--        "_root": attr.label(default = "//:unused"),
+--        "path": attr.string(),
+--    },
+--    local = True,
+--)
++diff --ruN a/stablehlo/third_party/BUILD.bazel b/stablehlo/third_party/BUILD.bazel
++--- stablehlo/third_party/BUILD.bazel
+++++ stablehlo/third_party/BUILD.bazel
++@@ -0,0 +1,12 @@
+++package(
+++    # copybara:uncomment default_applicable_licenses = ["//:license"],
+++    licenses = ["notice"],
+++)
+++
+++# copybara:uncomment_begin
+++# filegroup(
+++#     name = "all_files",
+++#     srcs = glob(["**"]),
+++#     visibility = ["//:__subpackages__"],
+++# )
+++# copybara:uncomment_end
+ diff --ruN a/stablehlo/third_party/rules_ml_toolchain/BUILD b/stablehlo/third_party/rules_ml_toolchain/BUILD
+ --- stablehlo/third_party/rules_ml_toolchain/BUILD
+ +++ stablehlo/third_party/rules_ml_toolchain/BUILD
+ @@ -1 +0,0 @@
+ -exports_files(["fix_strip.patch"])
+-diff --ruN a/stablehlo/third_party/rules_ml_toolchain/fix_strip.patch b/stablehlo/third_party/rules_ml_toolchain/fix_strip.patch
+---- stablehlo/third_party/rules_ml_toolchain/fix_strip.patch
+-+++ stablehlo/third_party/rules_ml_toolchain/fix_strip.patch
+-@@ -1,30 +0,0 @@
+--diff --git a/cc/impls/linux_x86_64_linux_x86_64/wrappers/BUILD b/cc/impls/linux_x86_64_linux_x86_64/wrappers/BUILD
+--index a206127..fa119c10 100644
+----- a/cc/impls/linux_x86_64_linux_x86_64/wrappers/BUILD
+--+++ b/cc/impls/linux_x86_64_linux_x86_64/wrappers/BUILD
+--@@ -24,5 +24,6 @@ filegroup(
+--         "idler",
+--         "ld",
+--         "objcopy",
+--+        "strip",
+--     ],
+--     visibility = ["//visibility:public"],
+-- )
+--diff --git a/cc/impls/linux_x86_64_linux_x86_64/wrappers/strip b/cc/impls/linux_x86_64_linux_x86_64/wrappers/strip
+--new file mode 100755
+--index 0000000..e3b0923
+----- /dev/null
+--+++ b/cc/impls/linux_x86_64_linux_x86_64/wrappers/strip
+--@@ -0,0 +1,3 @@
+--+#!/bin/bash
+--+
+--+"$STRIP_PATH" "$@"
+--diff --git a/third_party/rules_cc_toolchain/toolchain_config.bzl b/third_party/rules_cc_toolchain/toolchain_config.bzl
+--index 7381bbb..23afe63 100644
+----- a/third_party/rules_cc_toolchain/toolchain_config.bzl
+--+++ b/third_party/rules_cc_toolchain/toolchain_config.bzl
+--@@ -58,3 +58,4 @@ ALL_ACTIONS = [
+--     ACTION_NAMES.cpp_link_static_library,
+--     ACTION_NAMES.clif_match,
+--+    ACTION_NAMES.strip,
+-- ]
+-diff --ruN a/stablehlo/third_party/zlib-ng.BUILD b/stablehlo/third_party/zlib-ng.BUILD
+---- stablehlo/third_party/zlib-ng.BUILD
+-+++ stablehlo/third_party/zlib-ng.BUILD
+-@@ -1,106 +0,0 @@
+--# This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+--# See https://llvm.org/LICENSE.txt for license information.
+--# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+--
+--load("@bazel_skylib//rules:copy_file.bzl", "copy_file")
+--load("@rules_cc//cc:cc_library.bzl", "cc_library")
+--
+--package(
+--    default_visibility = ["//visibility:public"],
+--    # BSD/MIT-like license (for zlib)
+--    licenses = ["notice"],
+--)
+--
+--copy_file(
+--    # The input template is identical to the CMake output.
+--    name = "zconf_gen",
+--    src = "zconf.h.in",
+--    out = "zconf.h",
+--    allow_symlink = True,
+--)
+--
+--cc_library(
+--    name = "zlib",
+--    srcs = select({
+--        "@llvm-project//third-party:llvm_zlib_enabled": [
+--            "adler32.c",
+--            "adler32_p.h",
+--            "chunkset.c",
+--            "chunkset_tpl.h",
+--            "compare258.c",
+--            "compress.c",
+--            "crc32.c",
+--            "crc32_comb.c",
+--            "crc32_comb_tbl.h",
+--            "crc32_p.h",
+--            "crc32_tbl.h",
+--            "deflate.c",
+--            "deflate.h",
+--            "deflate_fast.c",
+--            "deflate_medium.c",
+--            "deflate_p.h",
+--            "deflate_quick.c",
+--            "deflate_slow.c",
+--            "fallback_builtins.h",
+--            "functable.c",
+--            "functable.h",
+--            "infback.c",
+--            "inffast.c",
+--            "inffast.h",
+--            "inffixed_tbl.h",
+--            "inflate.c",
+--            "inflate.h",
+--            "inflate_p.h",
+--            "inftrees.c",
+--            "inftrees.h",
+--            "insert_string.c",
+--            "insert_string_tpl.h",
+--            "match_tpl.h",
+--            "trees.c",
+--            "trees.h",
+--            "trees_emit.h",
+--            "trees_tbl.h",
+--            "uncompr.c",
+--            "zbuild.h",
+--            "zendian.h",
+--            "zutil.c",
+--            "zutil.h",
+--            "zutil_p.h",
+--        ],
+--        "//conditions:default": [],
+--    }),
+--    hdrs = select({
+--        "@llvm-project//third-party:llvm_zlib_enabled": [
+--            "zlib.h",
+--            ":zconf_gen",
+--        ],
+--        "//conditions:default": [],
+--    }),
+--    copts = [
+--        "-std=c11",
+--        "-DZLIB_COMPAT",
+--        "-DWITH_GZFILEOP",
+--        "-DWITH_OPTIM",
+--        "-DWITH_NEW_STRATEGIES",
+--        # For local builds you might want to add "-DWITH_NATIVE_INSTRUCTIONS"
+--        # here to improve performance. Native instructions aren't enabled in
+--        # the default config for reproducibility.
+--    ],
+--    defines = select({
+--        "@llvm-project//third-party:llvm_zlib_enabled": [
+--            "LLVM_ENABLE_ZLIB=1",
+--        ],
+--        "//conditions:default": [
+--            "LLVM_ENABLE_ZLIB=0",
+--        ],
+--    }),
+--    # Clang includes zlib with angled instead of quoted includes, so we need
+--    # strip_include_prefix here.
+--    strip_include_prefix = ".",
+--    visibility = ["//visibility:public"],
+--)
+--
+--alias(
+--    name = "zlib-ng",
+--    actual = ":zlib",
+--)
+-diff --ruN a/stablehlo/third_party/zstd.BUILD b/stablehlo/third_party/zstd.BUILD
+---- stablehlo/third_party/zstd.BUILD
+-+++ stablehlo/third_party/zstd.BUILD
+-@@ -1,44 +0,0 @@
+--# This file is licensed under the Apache License v2.0 with LLVM Exceptions.
+--# See https://llvm.org/LICENSE.txt for license information.
+--# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+--load("@rules_cc//cc:cc_library.bzl", "cc_library")
+--
+--package(
+--    default_visibility = ["//visibility:public"],
+--    # BSD/MIT-like license (for zstd)
+--    licenses = ["notice"],
+--)
+--
+--cc_library(
+--    name = "zstd",
+--    srcs = select({
+--        "@llvm-project//third-party:llvm_zstd_enabled": glob([
+--            "lib/common/*.c",
+--            "lib/common/*.h",
+--            "lib/compress/*.c",
+--            "lib/compress/*.h",
+--            "lib/decompress/*.c",
+--            "lib/decompress/*.h",
+--            "lib/decompress/*.S",
+--            "lib/dictBuilder/*.c",
+--            "lib/dictBuilder/*.h",
+--        ]),
+--        "//conditions:default": [],
+--    }),
+--    hdrs = select({
+--        "@llvm-project//third-party:llvm_zstd_enabled": [
+--            "lib/zdict.h",
+--            "lib/zstd.h",
+--            "lib/zstd_errors.h",
+--        ],
+--        "//conditions:default": [],
+--    }),
+--    defines = select({
+--        "@llvm-project//third-party:llvm_zstd_enabled": [
+--            "LLVM_ENABLE_ZSTD=1",
+--            "ZSTD_MULTITHREAD",
+--        ],
+--        "//conditions:default": [],
+--    }),
+--    strip_include_prefix = "lib",
+--)
++diff --ruN a/stablehlo/third_party/rules_ml_toolchain/BUILD.bazel b/stablehlo/third_party/rules_ml_toolchain/BUILD.bazel
++--- stablehlo/third_party/rules_ml_toolchain/BUILD.bazel
+++++ stablehlo/third_party/rules_ml_toolchain/BUILD.bazel
++@@ -0,0 +1,14 @@
+++package(
+++    # copybara:uncomment default_applicable_licenses = ["//:license"],
+++    licenses = ["notice"],
+++)
+++
+++# copybara:uncomment_begin
+++# filegroup(
+++#     name = "all_files",
+++#     srcs = glob(["**"]),
+++#     visibility = ["//:__subpackages__"],
+++# )
+++# copybara:uncomment_end
+++
+++exports_files(["fix_strip.patch"])
++diff --ruN a/stablehlo/third_party/xz.BUILD b/stablehlo/third_party/xz.BUILD
++--- stablehlo/third_party/xz.BUILD
+++++ stablehlo/third_party/xz.BUILD
++@@ -0,0 +1,329 @@
+++# Copyright 2026 The StableHLO 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
+++#
+++#      https://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.
+++
+++load("@bazel_skylib//rules:write_file.bzl", "write_file")
+++load("@rules_cc//cc:cc_library.bzl", "cc_library")
+++
+++package(default_visibility = ["//visibility:private"])
+++
+++licenses(["unencumbered"])
+++
+++exports_files(["COPYING"])
+++
+++cc_library(
+++    name = "lzma",
+++    srcs = glob(
+++        include = [
+++            "src/**/*.h",
+++            "src/liblzma/**/*.c",
+++        ],
+++        exclude = [
+++            "src/**/*_tablegen.c",
+++            "src/liblzma/check/crc_clmul_consts_gen.c",
+++            "src/liblzma/check/sha256.c",
+++            "src/liblzma/check/crc*_small.c",
+++            "src/liblzma/common/stream_decoder_mt.c",
+++            "src/liblzma/common/stream_encoder_mt.c",
+++        ],
+++    ) + select({
+++        "@platforms//os:osx": [],
+++        "@platforms//os:ios": [],
+++        "//conditions:default": [
+++            "src/liblzma/check/sha256.c",
+++        ],
+++    }) + select({
+++        "@platforms//os:none": [],
+++        "//conditions:default": [
+++            "src/common/tuklib_cpucores.c",
+++            "src/common/tuklib_physmem.c",
+++            "src/liblzma/common/stream_decoder_mt.c",
+++            "src/liblzma/common/stream_encoder_mt.c",
+++        ],
+++    }),
+++    hdrs = glob(["src/liblzma/api/**/*.h"]),
+++    copts = select({
+++        "@platforms//os:windows": [],
+++        "//conditions:default": ["-std=c99"],
+++    }),
+++    defines = select({
+++        "@platforms//os:windows": ["LZMA_API_STATIC"],
+++        "//conditions:default": [],
+++    }),
+++    implementation_deps = [
+++        ":config",
+++        ":src_common",
+++        ":src_liblzma",
+++        ":src_liblzma_check",
+++        ":src_liblzma_common",
+++        ":src_liblzma_delta",
+++        ":src_liblzma_lz",
+++        ":src_liblzma_lzma",
+++        ":src_liblzma_rangecoder",
+++        ":src_liblzma_simple",
+++    ],
+++    linkopts = select({
+++        "@platforms//os:android": [],
+++        "@platforms//os:windows": [],
+++        "@platforms//os:none": [],
+++        "//conditions:default": ["-lpthread"],
+++    }),
+++    linkstatic = select({
+++        "@platforms//os:windows": True,
+++        "//conditions:default": False,
+++    }),
+++    local_defines = ["HAVE_CONFIG_H"],
+++    strip_include_prefix = "src/liblzma/api",
+++    visibility = ["//visibility:public"],
+++)
+++
+++write_file(
+++    name = "config_gen",
+++    out = "lzma_config/config.h",
+++    content = [
+++        "#define ASSUME_RAM 128",
+++        "#ifndef _WIN32",
+++        "#define HAVE_BSWAP_16 1",
+++        "#define HAVE_BSWAP_32 1",
+++        "#define HAVE_BSWAP_64 1",
+++        "#endif",
+++        "#if defined(__APPLE__) || defined(_WIN32)",
+++        "#undef HAVE_BYTESWAP_H",
+++        "#else",
+++        "#define HAVE_BYTESWAP_H 1",
+++        "#endif",
+++        "#undef HAVE_CAPSICUM",
+++        "#if defined(__APPLE__)",
+++        "#define HAVE_CC_SHA256_CTX 1",
+++        "#define HAVE_CC_SHA256_INIT 1",
+++        "#define HAVE_COMMONCRYPTO_COMMONDIGEST_H 1",
+++        "#endif",
+++        "#define HAVE_CHECK_CRC32 1",
+++        "#define HAVE_CHECK_CRC64 1",
+++        "#define HAVE_CHECK_SHA256 1",
+++        "#if !defined(__APPLE__)",
+++        "#define HAVE_CLOCK_GETTIME 1",
+++        "#define HAVE_DECL_CLOCK_MONOTONIC 1",
+++        "#endif",
+++        "#if defined(__ANDROID__) || defined(__APPLE__) || defined(_WIN32)",
+++        "#define HAVE_DECL_PROGRAM_INVOCATION_NAME 0",
+++        "#else",
+++        "#define HAVE_DECL_PROGRAM_INVOCATION_NAME 1",
+++        "#endif",
+++        "#define HAVE_DECODERS 1",
+++        "#define HAVE_DECODER_ARM 1",
+++        "#define HAVE_DECODER_ARM64 1",
+++        "#define HAVE_DECODER_ARMTHUMB 1",
+++        "#define HAVE_DECODER_DELTA 1",
+++        "#define HAVE_DECODER_IA64 1",
+++        "#define HAVE_DECODER_LZMA1 1",
+++        "#define HAVE_DECODER_LZMA2 1",
+++        "#define HAVE_DECODER_POWERPC 1",
+++        "#define HAVE_DECODER_SPARC 1",
+++        "#define HAVE_DECODER_X86 1",
+++        "#define HAVE_DLFCN_H 1",
+++        "#define HAVE_ENCODERS 1",
+++        "#define HAVE_ENCODER_ARM 1",
+++        "#define HAVE_ENCODER_ARM64 1",
+++        "#define HAVE_ENCODER_ARMTHUMB 1",
+++        "#define HAVE_ENCODER_DELTA 1",
+++        "#define HAVE_ENCODER_IA64 1",
+++        "#define HAVE_ENCODER_LZMA1 1",
+++        "#define HAVE_ENCODER_LZMA2 1",
+++        "#define HAVE_ENCODER_POWERPC 1",
+++        "#define HAVE_ENCODER_SPARC 1",
+++        "#define HAVE_ENCODER_X86 1",
+++        "#define HAVE_FCNTL_H 1",
+++        "#define HAVE_FUTIMENS 1",
+++        "#define HAVE_GETOPT_H 1",
+++        "#define HAVE_GETOPT_LONG 1",
+++        "#define HAVE_INTTYPES_H 1",
+++        "#define HAVE_LIMITS_H 1",
+++        "#define HAVE_MBRTOWC 1",
+++        "#define HAVE_MEMORY_H 1",
+++        "#define HAVE_MF_BT2 1",
+++        "#define HAVE_MF_BT3 1",
+++        "#define HAVE_MF_BT4 1",
+++        "#define HAVE_MF_HC3 1",
+++        "#define HAVE_MF_HC4 1",
+++        "#if !defined(__APPLE__)",
+++        "#define HAVE_POSIX_FADVISE 1",
+++        "#endif",
+++        "#if !defined(__APPLE__) && !defined(__ANDROID__)",
+++        "#define HAVE_PTHREAD_CONDATTR_SETCLOCK 1",
+++        "#endif",
+++        "#define HAVE_PTHREAD_PRIO_INHERIT 1",
+++        "#undef HAVE_SMALL",
+++        "#define HAVE_STDBOOL_H 1",
+++        "#define HAVE_STDINT_H 1",
+++        "#define HAVE_STDLIB_H 1",
+++        "#define HAVE_STRING_H 1",
+++        "#define HAVE_STRUCT_STAT_ST_ATIM_TV_NSEC 1",
+++        "#undef HAVE_SYS_CAPSICUM_H",
+++        "#ifndef _WIN32",
+++        "#define HAVE_SYS_PARAM_H 1",
+++        "#define HAVE_SYS_STAT_H 1",
+++        "#define HAVE_SYS_TIME_H 1",
+++        "#define HAVE_SYS_TYPES_H 1",
+++        "#endif",
+++        "#define HAVE_UINTPTR_T 1",
+++        "#define HAVE_UNISTD_H 1",
+++        "#define HAVE_VISIBILITY 0",
+++        "#ifndef _WIN32",
+++        "#define HAVE_WCWIDTH 1",
+++        "#endif",
+++        "#define HAVE__BOOL 1",
+++        "#undef HAVE__FUTIME",
+++        "#define LT_OBJDIR \".libs/\"",
+++        "#ifdef _WIN32",
+++        "#define MYTHREAD_VISTA 1",
+++        "#else",
+++        "#define MYTHREAD_POSIX 1",
+++        "#endif",
+++        "#define NDEBUG 1",
+++        "#define PACKAGE \"xz\"",
+++        "#define PACKAGE_BUGREPORT \"lasse.collin@tukaani.org\"",
+++        "#define PACKAGE_NAME \"XZ Utils\"",
+++        "#define PACKAGE_STRING \"XZ Utils 5.8.3\"",
+++        "#define PACKAGE_TARNAME \"xz\"",
+++        "#define PACKAGE_URL \"https://tukaani.org/xz/\"",
+++        "#define PACKAGE_VERSION \"5.8.3\"",
+++        "#define SIZEOF_SIZE_T 8",
+++        "#define STDC_HEADERS 1",
+++        "#define TUKLIB_CPUCORES_SYSCONF 1",
+++        "#define TUKLIB_FAST_UNALIGNED_ACCESS 1",
+++        "#if defined(__APPLE__)",
+++        "#undef TUKLIB_PHYSMEM_SYSCONF",
+++        "#else",
+++        "#define TUKLIB_PHYSMEM_SYSCONF 1",
+++        "#endif",
+++        "#ifndef _ALL_SOURCE",
+++        "#define _ALL_SOURCE 1",
+++        "#endif",
+++        "#ifndef _GNU_SOURCE",
+++        "#define _GNU_SOURCE 1",
+++        "#endif",
+++        "#ifndef _POSIX_PTHREAD_SEMANTICS",
+++        "#define _POSIX_PTHREAD_SEMANTICS 1",
+++        "#endif",
+++        "#ifndef _TANDEM_SOURCE",
+++        "#define _TANDEM_SOURCE 1",
+++        "#endif",
+++        "#ifndef __EXTENSIONS__",
+++        "#define __EXTENSIONS__ 1",
+++        "#endif",
+++        "#define VERSION \"5.8.3\"",
+++        "#if defined AC_APPLE_UNIVERSAL_BUILD",
+++        "#if defined __BIG_ENDIAN__",
+++        "#define WORDS_BIGENDIAN 1",
+++        "#endif",
+++        "#endif",
+++        "#ifndef _DARWIN_USE_64_BIT_INODE",
+++        "#define _DARWIN_USE_64_BIT_INODE 1",
+++        "#endif",
+++        "",
+++    ],
+++)
+++
+++cc_library(
+++    name = "config",
+++    hdrs = ["lzma_config/config.h"],
+++    strip_include_prefix = "lzma_config",
+++)
+++
+++cc_library(
+++    name = "src_common",
+++    srcs = select({
+++        "@platforms//os:none": [],
+++        "//conditions:default": [
+++            "src/common/tuklib_exit.c",
+++            "src/common/tuklib_progname.c",
+++        ],
+++    }),
+++    hdrs = glob(["src/common/*.h"]),
+++    defines = select({
+++        "@platforms//os:windows": ["LZMA_API_STATIC"],
+++        "//conditions:default": [],
+++    }),
+++    local_defines = ["HAVE_CONFIG_H"] + select({
+++        "@platforms//os:windows": ["TUKLIB_GETTEXT=0"],
+++        "//conditions:default": [],
+++    }),
+++    strip_include_prefix = "src/common",
+++    deps = [
+++        ":config",
+++        ":src_liblzma",
+++        ":src_liblzma_api",
+++        ":src_liblzma_common",
+++    ],
+++)
+++
+++cc_library(
+++    name = "src_liblzma",
+++    hdrs = glob([
+++        "src/liblzma/common/*.h",
+++        "src/liblzma/lzma/*.h",
+++    ]),
+++    strip_include_prefix = "src/liblzma",
+++)
+++
+++cc_library(
+++    name = "src_liblzma_api",
+++    hdrs = glob(
+++        include = ["src/liblzma/api/**/*.h"],
+++        exclude = ["src/liblzma/api/lzma.h"],
+++    ),
+++    strip_include_prefix = "src/liblzma/api",
+++)
+++
+++cc_library(
+++    name = "src_liblzma_check",
+++    hdrs = glob(["src/liblzma/check/*.h"]),
+++    strip_include_prefix = "src/liblzma/check",
+++)
+++
+++cc_library(
+++    name = "src_liblzma_common",
+++    hdrs = glob(["src/liblzma/common/*.h"]),
+++    strip_include_prefix = "src/liblzma/common",
+++)
+++
+++cc_library(
+++    name = "src_liblzma_delta",
+++    hdrs = glob(["src/liblzma/delta/*.h"]),
+++    strip_include_prefix = "src/liblzma/delta",
+++)
+++
+++cc_library(
+++    name = "src_liblzma_lz",
+++    hdrs = glob(["src/liblzma/lz/*.h"]),
+++    strip_include_prefix = "src/liblzma/lz",
+++)
+++
+++cc_library(
+++    name = "src_liblzma_lzma",
+++    hdrs = glob(["src/liblzma/lzma/*.h"]),
+++    strip_include_prefix = "src/liblzma/lzma",
+++)
+++
+++cc_library(
+++    name = "src_liblzma_rangecoder",
+++    hdrs = glob(["src/liblzma/rangecoder/*.h"]),
+++    strip_include_prefix = "src/liblzma/rangecoder",
+++)
+++
+++cc_library(
+++    name = "src_liblzma_simple",
+++    hdrs = glob(["src/liblzma/simple/*.h"]),
+++    strip_include_prefix = "src/liblzma/simple",
+++)
 diff --git a/third_party/stablehlo/workspace.bzl b/third_party/stablehlo/workspace.bzl
 index 5c603385..7c94cab1 100644
 --- a/third_party/stablehlo/workspace.bzl
diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch
index c0bd29a0066..d2710427181 100644
--- a/third_party/xla/third_party/stablehlo/temporary.patch
+++ b/third_party/xla/third_party/stablehlo/temporary.patch
@@ -15,755 +15,42 @@
 #
 # This file is automatically generated by the `generate_patch` tool.
 # Do not edit directly.
-diff --ruN a/stablehlo/.github/ISSUE_TEMPLATE/bug_report.yml b/stablehlo/.github/ISSUE_TEMPLATE/bug_report.yml
---- stablehlo/.github/ISSUE_TEMPLATE/bug_report.yml
-+++ stablehlo/.github/ISSUE_TEMPLATE/bug_report.yml
-@@ -1,50 +0,0 @@
--# Copyright 2024 The StableHLO Authors.
--# 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.
--
--name: 🐞 Bug Report
--description: Report an issue
--body:
--  - type: markdown
--    attributes:
--      value: |
--        :star2: Thanks for taking the time to report this issue! :star2:
--
--        Please search through [other recent issues](https://github.com/openxla/stablehlo/issues) to see if your report overlaps with an existing issue.
--  - type: textarea
--    id: what-happened
--    attributes:
--      label: What happened?
--      description: |
--        Also tell us, what did you expect to happen?
--
--        Include any MLIR output leading up to whatever error occurred where possible.
--      placeholder: Tell us what you see! Please also copy/paste or link to any relevant log output.
--    validations:
--      required: true
--  - type: textarea
--    id: reproducing
--    attributes:
--      label: Steps to reproduce your issue
--      description: |
--        Include any commands that you ran.
--      value: |
--        1. Go to '...'
--        2. Click on '....'
--        3. Scroll down to '....'
--        4. See error
--  - type: textarea
--    id: version
--    attributes:
--      label: Version information
--      description: What version of our software are you using?
--      placeholder: This could be a git commit hash.
-diff --ruN a/stablehlo/.github/ISSUE_TEMPLATE/config.yml b/stablehlo/.github/ISSUE_TEMPLATE/config.yml
---- stablehlo/.github/ISSUE_TEMPLATE/config.yml
-+++ stablehlo/.github/ISSUE_TEMPLATE/config.yml
-@@ -1,24 +0,0 @@
--# Copyright 2024 The StableHLO Authors.
--# 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.
--
--blank_issues_enabled: true
--contact_links:
--  - name: 📄 Blank Issue
--    url: https://github.com/openxla/stablehlo/issues/new
--    about: If you know what you're doing (especially collaborators)
--  - name: 🗣 Start a discussion on GitHub
--    url: https://github.com/orgs/openxla/discussions/new
--    about: For early feedback on a design or other proposal
--  - name: 💬 Join us on Discord
--    url: https://discord.gg/Wqr8vWZKKJ
--    about: For realtime communication with the team and collaborators
-diff --ruN a/stablehlo/.github/ISSUE_TEMPLATE/feature_request.yml b/stablehlo/.github/ISSUE_TEMPLATE/feature_request.yml
---- stablehlo/.github/ISSUE_TEMPLATE/feature_request.yml
-+++ stablehlo/.github/ISSUE_TEMPLATE/feature_request.yml
-@@ -1,30 +0,0 @@
--# Copyright 2024 The StableHLO Authors.
--# 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.
--
--name: ➕ Feature Request
--description: Ask for a new feature to be added
--body:
--  - type: markdown
--    attributes:
--      value: |
--        :star2: Thanks for taking the time to file this request! :star2:
--
--        We also work closely with upstream communities like the MLIR project, which uses the [MLIR Discourse forum](https://discourse.llvm.org/c/mlir/31) for coordination.
--  - type: textarea
--    id: overview
--    attributes:
--      label: Request description
--      description: |
--        What would you like added or changed?
--    validations:
--      required: true
-diff --ruN a/stablehlo/.github/actions/setup-build/action.yml b/stablehlo/.github/actions/setup-build/action.yml
---- stablehlo/.github/actions/setup-build/action.yml
-+++ stablehlo/.github/actions/setup-build/action.yml
-@@ -1,58 +0,0 @@
--# Copyright 2023 The StableHLO Authors.
--#
--# 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
--#
--#      https://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.
--
--# The setup-build action gets everything needed by buildAndTest into the workspace.
--name: "Setup build environment (ninja, ccache, llvm, clang, lld)"
--description: "Setup build environment (ninja, ccache, llvm, clang, lld)"
--
--inputs:
--  llvm-version:
--    description: |
--      LLVM version for checkout and build. Used for ccache value and checkout.
--    required: true
--  clang-version:
--    description: |
--      Clang and LLD version to use for compilation. Independent of LLVM version.
--    required: false
--    default: '20'
--  python-version:
--    description: |
--      Python version to install.
--    required: false
--    default: '3.11.13'
--runs:
--  # This is a composite action - has a list of steps to execute.
--  using: "composite"
--
--  steps:
--  # Checkout llvm at version specified in input argument.
--  - uses: actions/checkout@v4
--    with:
--      repository: llvm/llvm-project
--      ref: ${{ inputs.llvm-version }}
--      path: llvm-project
--      fetch-depth: 1
--
--  # Setup C++ caching using ccache.
--  # Cache key is a combination of OS arch and LLVM version.
--  - name: Ccache for C++ compilation
--    uses: hendrikmuhs/ccache-action@v1.2
--    with:
--      key: ${{ runner.os }}-stablehlo_build_assets-${{ inputs.llvm-version }}
--      max-size: 4G
--
--  - name: Install MLIR python requirements
--    shell: bash
--    run: |
--      python3 -m pip install -r $GITHUB_WORKSPACE/llvm-project/mlir/python/requirements.txt
 diff --ruN a/stablehlo/.github/workflows/buildAndTestBazel.yml b/stablehlo/.github/workflows/buildAndTestBazel.yml
 --- stablehlo/.github/workflows/buildAndTestBazel.yml
 +++ stablehlo/.github/workflows/buildAndTestBazel.yml
-@@ -1,103 +0,0 @@
--# Copyright 2023 The StableHLO Authors.
--#
--# 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
--#
--#      https://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.
--
--name: Bazel Build
--
--on:
--  pull_request:
--    branches: [ main ]
--  push:
--    branches: [ main ]
--  schedule:
--    # Run once a day
--    - cron:  '0 12 * * *'
--  workflow_dispatch:
--    inputs:
--      halt-for-connection:
--        description: 'Should this workflow run wait for a remote connection?'
--        type: choice
--        required: true
--        default: 'no'
--        options:
--          - 'yes'
--          - 'no'
--
--# Ensure that only a single job or workflow using the same
--# concurrency group will run at a time. This would cancel
--# any in-progress jobs in the same github workflow and github
--# ref (e.g. refs/heads/main or refs/pull/<pr_number>/merge).
--concurrency:
--  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
--  cancel-in-progress: ${{ github.ref != 'main' }}
--
--jobs:
--  bazel-build:
--    strategy:
--      fail-fast: false
+@@ -45,7 +45,7 @@
+   bazel-build:
+     strategy:
+       fail-fast: false
 -    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n2-16' || 'ubuntu-22.04' }}
--    container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
--
--    steps:
--    # Checkout StableHLO with all commits in the chain
--    - name: Checkout StableHLO
--      uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
--      with:
--        fetch-depth: $(( ${{ github.event_name == 'pull_request' && github.event.pull_request.commits || 0 }} + 1 ))
--        set-safe-directory: true
--
--    - name: Mark project as safe for Git
--      run: git config --global --add safe.directory '*'
--
--    - name: "Wait For Connection"
--      uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
--      with:
--        halt-dispatch-input: ${{ inputs.halt-for-connection }}
--
--    - name: Validate LLVM Commit
--      run: |
--        ./build_tools/github_actions/lint_llvm_commit.sh .
--
--    - name: Get LLVM Version
--      id: llvm-version
--      shell: bash
--      run: |
--        echo "version=$(cat ./build_tools/llvm_version.txt)" >> $GITHUB_OUTPUT
--
--    - name: Build and check docs
--      shell: bash
--      run: |
--        ./build_tools/github_actions/ci_build_docs.sh -c
--
--    - name: Build and Test StableHLO (Diff Only)
--      if: ${{ github.event_name == 'pull_request' }}
--      shell: bash
--      run: |
--        START_HASH=${{ github.event.pull_request.base.sha }}
--        END_HASH=${{ github.event.pull_request.head.sha }}
--        BAZEL_DIFF=/tmp/bazel-diff.jar
--        curl -Lo "$BAZEL_DIFF" https://github.com/Tinder/bazel-diff/releases/download/17.0.2/bazel-diff_deploy.jar
--        ./build_tools/github_actions/ci_build_bazel.sh "$BAZEL_DIFF" "$START_HASH" "$END_HASH"
--
--    - name: Build and Test StableHLO (All)
--      if: ${{ github.event_name != 'pull_request' }}
--      shell: bash
--      run: |
--        ./build_tools/github_actions/ci_build_bazel.sh
--
--    - name: "Wait For Connection - On Error"
--      if: failure()
--      uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
--      with:
--        halt-dispatch-input: ${{ inputs.halt-for-connection }}
--
++    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n4-16' || 'ubuntu-22.04' }}
+     container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
+ 
+     steps:
 diff --ruN a/stablehlo/.github/workflows/buildAndTestCMake.yml b/stablehlo/.github/workflows/buildAndTestCMake.yml
 --- stablehlo/.github/workflows/buildAndTestCMake.yml
 +++ stablehlo/.github/workflows/buildAndTestCMake.yml
-@@ -1,98 +0,0 @@
--# Copyright 2023 The StableHLO Authors.
--#
--# 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
--#
--#      https://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.
--
--name: CMake Build
--
--on:
--  pull_request:
--    branches: [ main ]
--  push:
--    branches: [ main ]
--  schedule:
--    # Run once every 6hr building against llvm-project@HEAD
--    - cron:  '0 */6 * * *'
--  workflow_dispatch:
--
--# Ensure that only a single job or workflow using the same
--# concurrency group will run at a time. This would cancel
--# any in-progress jobs in the same github workflow and github
--# ref (e.g. refs/heads/main or refs/pull/<pr_number>/merge).
--concurrency:
--  group: ${{ github.workflow }}-${{ github.ref }}
--  cancel-in-progress: true
--
--jobs:
--  cmake-build:
--    # Only run scheduled CI on main repo
--    if: (github.repository == 'openxla/stablehlo' || github.event_name != 'schedule')
--    name: "cmake-build${{ github.event_name == 'schedule' && ' (llvm-project@HEAD)' || ''}}"
--    env:
--      LLVM_PROJECT_DIR: "llvm-project"
--      LLVM_BUILD_DIR: "llvm-build"
--      STABLEHLO_BUILD_DIR: "stablehlo-build"
--      STABLEHLO_PYTHON_BUILD_DIR: "stablehlo-python-build"
--    strategy:
--      fail-fast: false
+@@ -44,7 +44,7 @@
+       STABLEHLO_PYTHON_BUILD_DIR: "stablehlo-python-build"
+     strategy:
+       fail-fast: false
 -    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n2-16' || 'ubuntu-22.04' }}
--    container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cmake:latest"
--
--    steps:
--    - name: Checkout StableHLO
--      uses: actions/checkout@v4
--
--    - name: Get LLVM Version
--      id: llvm-version
--      shell: bash
--      run: |
--        USE_LLVM_HEAD=${{ github.event_name == 'schedule' }}
--        if [[ $USE_LLVM_HEAD = true ]]; then
--          echo "version=main" >> $GITHUB_OUTPUT
--        else
--          echo "version=$(cat ./build_tools/llvm_version.txt)" >> $GITHUB_OUTPUT
--        fi;
--
--    - name: Setup workspace
--      uses: ./.github/actions/setup-build
--      with:
--        llvm-version: ${{ steps.llvm-version.outputs.version }}
--
--    - name: Configure and Build LLVM
--      shell: bash
--      run: |
--        ./build_tools/github_actions/ci_build_cmake_llvm.sh "$LLVM_PROJECT_DIR" "$LLVM_BUILD_DIR"
--      env:
--          CMAKE_BUILD_TYPE: Release
--          MLIR_ENABLE_BINDINGS_PYTHON: ON
--
--    - name: Build and Test StableHLO (with AddressSanitizer)
--      shell: bash
--      run: |
--        ./build_tools/github_actions/ci_build_cmake.sh "$LLVM_BUILD_DIR" "$STABLEHLO_BUILD_DIR"
--      env:
--          CMAKE_BUILD_TYPE: Release
--          STABLEHLO_ENABLE_BINDINGS_PYTHON: OFF
--          STABLEHLO_ENABLE_SANITIZER: address
--          # TODO: remove this once https://github.com/openxla/stablehlo/pull/2318 is fixed
--          ASAN_OPTIONS: allow_user_poisoning=false
--
--    - name: Build and Test StableHLO (with Python bindings)
--      shell: bash
--      run: |
--        python3 -m pip install tensorflow-cpu nanobind==2.9
--        python3 -m pip install -r "$LLVM_PROJECT_DIR/mlir/python/requirements.txt"
--        ./build_tools/github_actions/ci_build_cmake.sh "$LLVM_BUILD_DIR" "$STABLEHLO_BUILD_DIR"
--      env:
--          CMAKE_BUILD_TYPE: Release
--          STABLEHLO_ENABLE_BINDINGS_PYTHON: ON
--          STABLEHLO_ENABLE_PYTHON_TF_TESTS: ON
-diff --ruN a/stablehlo/.github/workflows/execute_tutorials.yml b/stablehlo/.github/workflows/execute_tutorials.yml
---- stablehlo/.github/workflows/execute_tutorials.yml
-+++ stablehlo/.github/workflows/execute_tutorials.yml
-@@ -1,46 +0,0 @@
--# Copyright 2024 The StableHLO 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
--
--#      https://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.
--
--name: Execute Tutorial Notebooks
--
--on:
--  push:
--    paths:
--      - 'docs/tutorials/**.ipynb'
--      - .github/workflows/execute_tutorials.yml
--  schedule:
--    # Run once a day
--    - cron:  '0 20 * * *'
--  workflow_dispatch:
--
--jobs:
--  test-notebooks:
--    runs-on: ubuntu-latest
--
--    steps:
--      - name: Checkout repository
--        uses: actions/checkout@v4
--
--      - name: Set up Python and install dependencies
--        uses: actions/setup-python@v5
--        with:
--          python-version: '3.11'
--      - run:  |
--          pip install --upgrade pip
--          pip install ipython ipykernel nbclient
--
--      - name: Execute tutorial notebooks
--        run: |
--          jupyter execute docs/tutorials/*.ipynb
--        shell: bash
-diff --ruN a/stablehlo/.github/workflows/lint.yml b/stablehlo/.github/workflows/lint.yml
---- stablehlo/.github/workflows/lint.yml
-+++ stablehlo/.github/workflows/lint.yml
-@@ -1,111 +0,0 @@
--# Copyright 2023 The StableHLO Authors.
--#
--# 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
--#
--#      https://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.
--
--name: Lint
--
--on:
--  pull_request:
--    branches: [ main ]
--  push:
--    branches: [ main ]
--  schedule:
--    # Run once a day
--    - cron:  '0 12 * * *'
--  workflow_dispatch:
--
--env:
--  # Have `go install` place binaries in $PATH
--  GOBIN: "/usr/local/bin"
--
--jobs:
--  clang-format:
--    # This job can only be run on pull_request since GITHUB_BASE_REF is only set on PR.
--    if: "github.event_name == 'pull_request'"
--    runs-on: ubuntu-22.04
--    steps:
--      - name: Checking out repository
--        uses: actions/checkout@v4
--
--      - name: Run clang-format
--        run: |
--          git fetch --no-tags --prune --depth=1 origin "${GITHUB_BASE_REF?}:${GITHUB_BASE_REF?}"
--          ./build_tools/github_actions/lint_clang_format.sh -b "${GITHUB_BASE_REF}"
--
--  whitespace-checks:
--    if: "github.event_name == 'pull_request'"
--    runs-on: ubuntu-22.04
--    steps:
--      - name: Checking out repository
--        uses: actions/checkout@v4
--
--      - name: Run Whitespace Checks
--        run: |
--          git fetch --no-tags --prune --depth=1 origin "${GITHUB_BASE_REF?}:${GITHUB_BASE_REF?}"
--          ./build_tools/github_actions/lint_whitespace_checks.sh
--
--  version-checks:
--    if: "github.event_name == 'pull_request'"
--    runs-on: ubuntu-22.04
--    steps:
--      - name: Checking out repository
--        uses: actions/checkout@v4
--
--      - name: Run Version Checks
--        run: |
--          ./build_tools/github_actions/lint_version.sh
--
--  buildifier-check:
--    if: "github.event_name == 'pull_request'"
--    runs-on: ubuntu-22.04
--    defaults:
--      run:
--        shell: bash
--    steps:
--      - name: Checking out repository
--        uses: actions/checkout@v4
--
--      - name: Install buildifier
--        run: go install github.com/bazelbuild/buildtools/buildifier@433ea85 # 6.4.0
--
--      - name: Run buildifier checks
--        run: ./build_tools/github_actions/lint_buildifier.sh
--
--  license-check:
--    if: "github.event_name == 'pull_request'"
--    runs-on: ubuntu-22.04
--    defaults:
--      run:
--        shell: bash
--    steps:
--      - name: Checking out repository
--        uses: actions/checkout@v4
--
--      - name: Run license checks
--        run: |
--          git fetch --no-tags --prune --depth=1 origin "${GITHUB_BASE_REF?}:${GITHUB_BASE_REF?}"
--          ./build_tools/github_actions/lint_check_license.sh -b "${GITHUB_BASE_REF}"
--
--  shell-script-check:
--    if: "github.event_name == 'pull_request'"
--    runs-on: ubuntu-22.04
--    defaults:
--      run:
--        shell: bash
--    steps:
--      - name: Checking out repository
--        uses: actions/checkout@v4
--
--      # shellcheck is installed by default on Ubuntu runners: https://github.com/actions/runner-images/blob/37a171e44de1f134793d7f0cbff137ca481942ab/images/ubuntu/Ubuntu2204-Readme.md
--      - name: Run shellcheck
--        run: ./build_tools/github_actions/lint_shellcheck.sh
-diff --ruN a/stablehlo/.github/workflows/markdown_lint.yml b/stablehlo/.github/workflows/markdown_lint.yml
---- stablehlo/.github/workflows/markdown_lint.yml
-+++ stablehlo/.github/workflows/markdown_lint.yml
-@@ -1,35 +0,0 @@
--# Copyright 2023 The StableHLO Authors.
--#
--# 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
--#
--#      https://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.
--
--name: Markdown Lint
--
--on:
--  pull_request:
--    branches: [ main ]
--  push:
--    branches: [ main ]
--  schedule:
--    # Run once a day
--    - cron:  '0 12 * * *'
--  workflow_dispatch:
--
--jobs:
--  markdown-lint:
--    runs-on: ubuntu-22.04
--    steps:
--    - name: Checking out repository
--      uses: actions/checkout@v4
--    - name: Running markdownlint
--      run: |
--        ./build_tools/github_actions/lint_markdown.sh --ignore "./docs/generated/**/*.md" "./**/*.md"
++    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n4-16' || 'ubuntu-22.04' }}
+     container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cmake:latest"
+ 
+     steps:
 diff --ruN a/stablehlo/.github/workflows/publishWheelRelease.yml b/stablehlo/.github/workflows/publishWheelRelease.yml
 --- stablehlo/.github/workflows/publishWheelRelease.yml
 +++ stablehlo/.github/workflows/publishWheelRelease.yml
-@@ -1,136 +0,0 @@
--# Copyright 2024 The StableHLO Authors.
--#
--# 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
--#
--#      https://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.
--
--name: Publish Python Wheel Release
--
--on:
--  workflow_dispatch:
--  schedule:
--    # Runs at 6:00 AM UTC, which is 10:00 PM previous day PST (UTC-8)
--    - cron: '0 6 * * *'
--
--# Ensure that only a single job or workflow using the same
--# concurrency group will run at a time. This would cancel
--# any in-progress jobs in the same github workflow and github
--# ref (e.g. refs/heads/main or refs/pull/<pr_number>/merge).
--concurrency:
--  group: ${{ github.workflow }}-${{ github.ref }}
--  cancel-in-progress: true
--
--jobs:
--  # TODO(fzakaria): Consider running this job only if a new commit for the day has been done
--  cmake-build-wheel:
--    # Only run scheduled CI on main repo
--    if: (github.repository == 'openxla/stablehlo' || github.event_name != 'schedule')
--    name: "cmake-build ${{ github.event_name == 'schedule' && '(llvm-project@HEAD)' || ''}}"
--    env:
--      LLVM_PROJECT_DIR: "llvm-project"
--      LLVM_BUILD_DIR: "llvm-build"
--      # The setup.py script expects the StableHLO build to be in 'build'
--      STABLEHLO_BUILD_DIR: "build"
--    strategy:
--      fail-fast: false
--      matrix:
--        python-version: [ "3.10", "3.11", "3.12" ]
+@@ -43,7 +43,7 @@
+       fail-fast: false
+       matrix:
+         python-version: [ "3.10", "3.11", "3.12" ]
 -    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n2-16' ||  'ubuntu-22.04'  }}
--    container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cmake:latest"
--
--    steps:
--    - name: Checkout StableHLO
--      uses: actions/checkout@v4
--      # TODO(fzakaria): We need all the commits and the tags for the
--      # git describe` command to work to find the correct version.
--      # This can be removed if the version can be determmined in another way.
--      with:
--        fetch-depth: 0
--    - name: Get LLVM Version
--      id: llvm-version
--      shell: bash
--      run: |
--        echo "version=$(cat ${{ github.workspace }}/build_tools/llvm_version.txt)" >> $GITHUB_OUTPUT
--
--    - name: Setup workspace
--      uses: ./.github/actions/setup-build
--      with:
--        llvm-version: ${{ steps.llvm-version.outputs.version }}
--        python-version: ${{ matrix.python-version }}
--
--    # The actual tool need to execute the python build
--    # TODO(fzakaria): Consider using cibuildwheel to build multiple wheels at
--    # once and mark them as manylinux compatible.
--    - name: Install build dependencies
--      shell: bash
--      run: pip install build
--
--    - name: Configure and Build LLVM
--      shell: bash
--      run: |
--        ./build_tools/github_actions/ci_build_cmake_llvm.sh "$LLVM_PROJECT_DIR" "$LLVM_BUILD_DIR"
--      env:
--          CMAKE_BUILD_TYPE: Release
--          MLIR_ENABLE_BINDINGS_PYTHON: ON
--
--    - name: Build and Test StableHLO with Python
--      shell: bash
--      run: |
--        pip install tensorflow-cpu nanobind==2.4
--        pip install -r "$LLVM_PROJECT_DIR/mlir/python/requirements.txt"
--        ./build_tools/github_actions/ci_build_cmake.sh "$LLVM_BUILD_DIR" "$STABLEHLO_BUILD_DIR"
--      env:
--          CMAKE_BUILD_TYPE: Release
--          STABLEHLO_ENABLE_BINDINGS_PYTHON: ON
--          STABLEHLO_ENABLE_PYTHON_TF_TESTS: ON
--
--    - name: Build and Test Python Wheel
--      shell: bash
--      run: |
--        ./build_tools/github_actions/ci_build_python_wheel.sh "$GITHUB_WORKSPACE/${{ github.sha }}"
--
--    - name: Upload an artifact
--      uses: actions/upload-artifact@v4
--      with:
--        if-no-files-found: error
--        name: build_artifact-${{ matrix.python-version }}
--        path: ${{ github.sha }}
--
--  upload-tarballs:
--    runs-on: 'ubuntu-24.04'
--    needs: cmake-build-wheel
--    strategy:
--      fail-fast: false
--      matrix:
--        python-version: [ "3.10", "3.11", "3.12" ]
--
--    steps:
--      - name: Checkout
--        uses: actions/checkout@v4
--
--      - name: Download artifacts
--        uses: actions/download-artifact@v4
--        with:
--          name: build_artifact-${{ matrix.python-version }}
--          path: ${{ github.sha }}
--
--      - name: Release current commit
--        uses: ncipollo/release-action@v1.13.0
--        with:
--          artifacts: "${{ github.sha }}/*.whl"
--          token: "${{ secrets.PUBLISH_ACCESS_TOKEN }}"
--          # git describe in setup.py explicitly filters out this tag
--          tag: "dev-wheels"
--          name: "dev-wheels"
--          removeArtifacts: false
--          allowUpdates: true
--          replacesArtifacts: true
--          makeLatest: false
-diff --ruN a/stablehlo/.gitignore b/stablehlo/.gitignore
---- stablehlo/.gitignore
-+++ stablehlo/.gitignore
-@@ -1,14 +0,0 @@
--bazel-bin
--bazel-out
--bazel-stablehlo
--bazel-testlogs
--build
--cmake-build
--llvm-project
--llvm-build
--.vscode
--.cache
--.idea/
--**/__pycache__
--*.whl
--**/*.egg-info/
++    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n4-16' ||  'ubuntu-22.04'  }}
+     container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cmake:latest"
+ 
+     steps:
 diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel
 --- stablehlo/BUILD.bazel
 +++ stablehlo/BUILD.bazel
@@ -3522,6 +2809,479 @@ diff --ruN a/stablehlo/stablehlo/integrations/cpp/builder/StablehloBuilderTest.c
  TEST(MlirBuilderTest, GatherOp) {
    std::string expected = R"mlir(module {
    func.func @main(%arg0: tensor<3xi64>, %arg1: tensor<1x1xi64>) -> tensor<1xi64> {
+@@ -1021,7 +1097,7 @@
+ TEST(MlirBuilderTest, ConstantF4E2M1FNArray) {
+   std::string expected = R"mlir(module {
+   func.func @main() -> tensor<10xf4E2M1FN> {
+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 0.000000e+00, 0.000000e+00, 3.000000e+00, 6.000000e+00, 6.000000e+00, 1.000000e+00, 6.000000e+00]> : tensor<10xf4E2M1FN>
++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 0.000000e+00, 0.000000e+00, 3.000000e+00, 6.000000e+00, -6.000000e+00, 5.000000e-01, -5.000000e-01]> : tensor<10xf4E2M1FN>
+     return %cst : tensor<10xf4E2M1FN>
+   }
+ })mlir";
+@@ -1031,10 +1107,15 @@
+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
+     auto type = makeTensorType(fb.getContext(), {10}, ElementType::F4E2M1FN);
++    double posMax = 6.0;
++    double negMax = -6.0;
++    double posMinDenorm = 0.5;
++    double negMinDenorm = -0.5;
+     auto cst = stablehlo::Constant(
+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
+-                                           0x07, 0x0F, 0x01, 0x09}),
+-                         type));
++        fb, makeConstant(
++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
++                                  negMax, posMinDenorm, negMinDenorm}),
++                type));
+     func::Return(fb, {cst});
+   }
+ 
+@@ -1045,7 +1126,7 @@
+ TEST(MlirBuilderTest, ConstantF6E2M3FNArray) {
+   std::string expected = R"mlir(module {
+   func.func @main() -> tensor<10xf6E2M3FN> {
+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.250000e-01, 3.250000e+00, 7.500000e+00, 7.500000e+00, 1.000000e+00, 7.500000e+00]> : tensor<10xf6E2M3FN>
++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.250000e-01, 3.250000e+00, 7.500000e+00, -7.500000e+00, 1.250000e-01, -1.250000e-01]> : tensor<10xf6E2M3FN>
+     return %cst : tensor<10xf6E2M3FN>
+   }
+ })mlir";
+@@ -1055,10 +1136,15 @@
+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
+     auto type = makeTensorType(fb.getContext(), {10}, ElementType::F6E2M3FN);
++    double posMax = 7.5;
++    double negMax = -7.5;
++    double posMinDenorm = 0x1p-3;
++    double negMinDenorm = -0x1p-3;
+     auto cst = stablehlo::Constant(
+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
+-                                           0x1F, 0x3F, 0x01, 0x21}),
+-                         type));
++        fb, makeConstant(
++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
++                                  negMax, posMinDenorm, negMinDenorm}),
++                type));
+     func::Return(fb, {cst});
+   }
+ 
+@@ -1069,7 +1155,7 @@
+ TEST(MlirBuilderTest, ConstantF6E3M2FNArray) {
+   std::string expected = R"mlir(module {
+   func.func @main() -> tensor<10xf6E3M2FN> {
+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.250000e-01, 3.000000e+00, 2.800000e+01, 2.800000e+01, 1.000000e+00, 2.800000e+01]> : tensor<10xf6E3M2FN>
++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.250000e-01, 3.000000e+00, 2.800000e+01, -2.800000e+01, 6.250000e-02, -6.250000e-02]> : tensor<10xf6E3M2FN>
+     return %cst : tensor<10xf6E3M2FN>
+   }
+ })mlir";
+@@ -1079,10 +1165,15 @@
+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
+     auto type = makeTensorType(fb.getContext(), {10}, ElementType::F6E3M2FN);
++    double posMax = 28.0;
++    double negMax = -28.0;
++    double posMinDenorm = 0x1p-4;
++    double negMinDenorm = -0x1p-4;
+     auto cst = stablehlo::Constant(
+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
+-                                           0x1F, 0x3F, 0x01, 0x21}),
+-                         type));
++        fb, makeConstant(
++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
++                                  negMax, posMinDenorm, negMinDenorm}),
++                type));
+     func::Return(fb, {cst});
+   }
+ 
+@@ -1092,9 +1183,9 @@
+ 
+ TEST(MlirBuilderTest, ConstantF8E3M4Array) {
+   std::string expected = R"mlir(module {
+-  func.func @main() -> tensor<10xf8E3M4> {
+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.125000e+00, 0x70, 0x70, 1.000000e+00, 0x70]> : tensor<10xf8E3M4>
+-    return %cst : tensor<10xf8E3M4>
++  func.func @main() -> tensor<11xf8E3M4> {
++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.125000e+00, 0x70, 0xF0, 0x78, 1.562500e-02, -1.562500e-02]> : tensor<11xf8E3M4>
++    return %cst : tensor<11xf8E3M4>
+   }
+ })mlir";
+ 
+@@ -1102,11 +1193,17 @@
+   {
+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
+-    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E3M4);
++    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E3M4);
++    double posInf = std::numeric_limits<double>::infinity();
++    double negInf = -std::numeric_limits<double>::infinity();
++    double nan = std::numeric_limits<double>::quiet_NaN();
++    double posMinDenorm = 0x1p-6;
++    double negMinDenorm = -0x1p-6;
+     auto cst = stablehlo::Constant(
+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
+-                                           0x7F, 0xFF, 0x01, 0x81}),
+-                         type));
++        fb, makeConstant(
++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posInf,
++                                  negInf, nan, posMinDenorm, negMinDenorm}),
++                type));
+     func::Return(fb, {cst});
+   }
+ 
+@@ -1116,9 +1213,9 @@
+ 
+ TEST(MlirBuilderTest, ConstantF8E4M3B11FNUZArray) {
+   std::string expected = R"mlir(module {
+-  func.func @main() -> tensor<10xf8E4M3B11FNUZ> {
+-    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 0x80, 0x80, 1.000000e+00, 0x80]> : tensor<10xf8E4M3B11FNUZ>
+-    return %cst : tensor<10xf8E4M3B11FNUZ>
++  func.func @main() -> tensor<11xf8E4M3B11FNUZ> {
++    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 3.000000e+01, -3.000000e+01, 0x80, 1.220700e-04, -1.220700e-04]> : tensor<11xf8E4M3B11FNUZ>
++    return %cst : tensor<11xf8E4M3B11FNUZ>
+   }
+ })mlir";
+ 
+@@ -1127,11 +1224,17 @@
+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
+     auto type =
+-        makeTensorType(fb.getContext(), {10}, ElementType::F8E4M3B11FNUZ);
++        makeTensorType(fb.getContext(), {11}, ElementType::F8E4M3B11FNUZ);
++    double posMax = 30.0;
++    double negMax = -30.0;
++    double nan = std::numeric_limits<double>::quiet_NaN();
++    double posMinDenorm = 0x1p-13;
++    double negMinDenorm = -0x1p-13;
+     auto cst = stablehlo::Constant(
+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
+-                                           0x7F, 0xFF, 0x01, 0x81}),
+-                         type));
++        fb, makeConstant(
++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
++                                  negMax, nan, posMinDenorm, negMinDenorm}),
++                type));
+     func::Return(fb, {cst});
+   }
+ 
+@@ -1141,9 +1244,9 @@
+ 
+ TEST(MlirBuilderTest, ConstantF8E4M3Array) {
+   std::string expected = R"mlir(module {
+-  func.func @main() -> tensor<10xf8E4M3> {
+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 1.280000e+02, 0x78, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E4M3>
+-    return %cst : tensor<10xf8E4M3>
++  func.func @main() -> tensor<11xf8E4M3> {
++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 0x78, 0xF8, 0x7C, 1.953130e-03, -1.953130e-03]> : tensor<11xf8E4M3>
++    return %cst : tensor<11xf8E4M3>
+   }
+ })mlir";
+ 
+@@ -1151,11 +1254,17 @@
+   {
+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
+-    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E4M3);
++    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E4M3);
++    double posInf = std::numeric_limits<double>::infinity();
++    double negInf = -std::numeric_limits<double>::infinity();
++    double nan = std::numeric_limits<double>::quiet_NaN();
++    double posMinDenorm = 0x1p-9;
++    double negMinDenorm = -0x1p-9;
+     auto cst = stablehlo::Constant(
+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
+-                                           0x7F, 0xFF, 0x01, 0x81}),
+-                         type));
++        fb, makeConstant(
++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posInf,
++                                  negInf, nan, posMinDenorm, negMinDenorm}),
++                type));
+     func::Return(fb, {cst});
+   }
+ 
+@@ -1165,9 +1274,9 @@
+ 
+ TEST(MlirBuilderTest, ConstantF8E4M3FNArray) {
+   std::string expected = R"mlir(module {
+-  func.func @main() -> tensor<10xf8E4M3FN> {
+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 1.280000e+02, 2.560000e+02, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E4M3FN>
+-    return %cst : tensor<10xf8E4M3FN>
++  func.func @main() -> tensor<11xf8E4M3FN> {
++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 4.480000e+02, -4.480000e+02, 0x7F, 1.953130e-03, -1.953130e-03]> : tensor<11xf8E4M3FN>
++    return %cst : tensor<11xf8E4M3FN>
+   }
+ })mlir";
+ 
+@@ -1175,11 +1284,17 @@
+   {
+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
+-    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E4M3FN);
++    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E4M3FN);
++    double posMax = 448.0;
++    double negMax = -448.0;
++    double nan = std::numeric_limits<double>::quiet_NaN();
++    double posMinDenorm = 0x1p-9;
++    double negMinDenorm = -0x1p-9;
+     auto cst = stablehlo::Constant(
+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
+-                                           0x7F, 0xFF, 0x01, 0x81}),
+-                         type));
++        fb, makeConstant(
++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
++                                  negMax, nan, posMinDenorm, negMinDenorm}),
++                type));
+     func::Return(fb, {cst});
+   }
+ 
+@@ -1189,9 +1304,9 @@
+ 
+ TEST(MlirBuilderTest, ConstantF8E4M3FNUZArray) {
+   std::string expected = R"mlir(module {
+-  func.func @main() -> tensor<10xf8E4M3FNUZ> {
+-    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 1.280000e+02, 0x80, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E4M3FNUZ>
+-    return %cst : tensor<10xf8E4M3FNUZ>
++  func.func @main() -> tensor<11xf8E4M3FNUZ> {
++    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 2.400000e+02, -2.400000e+02, 0x80, 9.765620e-04, -9.765620e-04]> : tensor<11xf8E4M3FNUZ>
++    return %cst : tensor<11xf8E4M3FNUZ>
+   }
+ })mlir";
+ 
+@@ -1199,11 +1314,17 @@
+   {
+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
+-    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E4M3FNUZ);
++    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E4M3FNUZ);
++    double posMax = 240.0;
++    double negMax = -240.0;
++    double nan = std::numeric_limits<double>::quiet_NaN();
++    double posMinDenorm = 0x1p-10;
++    double negMinDenorm = -0x1p-10;
+     auto cst = stablehlo::Constant(
+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
+-                                           0x7F, 0xFF, 0x01, 0x81}),
+-                         type));
++        fb, makeConstant(
++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
++                                  negMax, nan, posMinDenorm, negMinDenorm}),
++                type));
+     func::Return(fb, {cst});
+   }
+ 
+@@ -1213,9 +1334,9 @@
+ 
+ TEST(MlirBuilderTest, ConstantF8E5M2Array) {
+   std::string expected = R"mlir(module {
+-  func.func @main() -> tensor<10xf8E5M2> {
+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.000000e+00, 1.280000e+02, 2.560000e+02, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E5M2>
+-    return %cst : tensor<10xf8E5M2>
++  func.func @main() -> tensor<11xf8E5M2> {
++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.000000e+00, 0x7C, 0xFC, 0x7E, 1.525880e-05, -1.525880e-05]> : tensor<11xf8E5M2>
++    return %cst : tensor<11xf8E5M2>
+   }
+ })mlir";
+ 
+@@ -1223,11 +1344,17 @@
+   {
+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
+-    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E5M2);
++    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E5M2);
++    double posInf = std::numeric_limits<double>::infinity();
++    double negInf = -std::numeric_limits<double>::infinity();
++    double nan = std::numeric_limits<double>::quiet_NaN();
++    double posMinDenorm = 0x1p-16;
++    double negMinDenorm = -0x1p-16;
+     auto cst = stablehlo::Constant(
+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
+-                                           0x7F, 0xFF, 0x01, 0x81}),
+-                         type));
++        fb, makeConstant(
++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posInf,
++                                  negInf, nan, posMinDenorm, negMinDenorm}),
++                type));
+     func::Return(fb, {cst});
+   }
+ 
+@@ -1237,9 +1364,9 @@
+ 
+ TEST(MlirBuilderTest, ConstantF8E5M2FNUZArray) {
+   std::string expected = R"mlir(module {
+-  func.func @main() -> tensor<10xf8E5M2FNUZ> {
+-    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.000000e+00, 1.280000e+02, 2.560000e+02, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E5M2FNUZ>
+-    return %cst : tensor<10xf8E5M2FNUZ>
++  func.func @main() -> tensor<11xf8E5M2FNUZ> {
++    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.000000e+00, 5.734400e+04, -5.734400e+04, 0x80, 7.629390e-06, -7.629390e-06]> : tensor<11xf8E5M2FNUZ>
++    return %cst : tensor<11xf8E5M2FNUZ>
+   }
+ })mlir";
+ 
+@@ -1247,11 +1374,17 @@
+   {
+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
+-    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E5M2FNUZ);
++    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E5M2FNUZ);
++    double posMax = 57344.0;
++    double negMax = -57344.0;
++    double nan = std::numeric_limits<double>::quiet_NaN();
++    double posMinDenorm = 0x1p-17;
++    double negMinDenorm = -0x1p-17;
+     auto cst = stablehlo::Constant(
+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
+-                                           0x7F, 0xFF, 0x01, 0x81}),
+-                         type));
++        fb, makeConstant(
++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
++                                  negMax, nan, posMinDenorm, negMinDenorm}),
++                type));
+     func::Return(fb, {cst});
+   }
+ 
+@@ -1262,7 +1395,7 @@
+ TEST(MlirBuilderTest, ConstantF8E8M0FNUArray) {
+   std::string expected = R"mlir(module {
+   func.func @main() -> tensor<8xf8E8M0FNU> {
+-    %cst = stablehlo.constant dense<[5.877470e-39, 1.000000e+00, 1.250000e-01, 1.250000e-01, 4.000000e+00, 5.877470e-39, 1.280000e+02, 2.560000e+02]> : tensor<8xf8E8M0FNU>
++    %cst = stablehlo.constant dense<[5.877470e-39, 1.000000e+00, 1.250000e-01, 1.250000e-01, 4.000000e+00, 5.877470e-39, 2.000000e+00, 0xFF]> : tensor<8xf8E8M0FNU>
+     return %cst : tensor<8xf8E8M0FNU>
+   }
+ })mlir";
+@@ -1272,10 +1405,12 @@
+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
+     auto type = makeTensorType(fb.getContext(), {8}, ElementType::F8E8M0FNU);
++    double min = 0x1p-127;
++    double nan = std::numeric_limits<double>::quiet_NaN();
+     auto cst = stablehlo::Constant(
+-        fb, makeConstant(ArrayRef<double>(
+-                             {0.0, 1.0, 0.125, 0.1, 3.1415, 0x00, 0x80, 0xFF}),
+-                         type));
++        fb, makeConstant(
++                ArrayRef<double>({0.0, 1.0, 0.125, 0.1, 3.1415, min, 2.0, nan}),
++                type));
+     func::Return(fb, {cst});
+   }
+ 
+@@ -1286,7 +1421,7 @@
+ TEST(MlirBuilderTest, ConstantBF16Array) {
+   std::string expected = R"mlir(module {
+   func.func @main() -> tensor<11xbf16> {
+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000980e-01, 3.140630e+00, 3.264000e+04, 6.553600e+04, 3.276800e+04, 1.000000e+00, 3.276800e+04]> : tensor<11xbf16>
++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000980e-01, 3.140630e+00, 0x7F80, 0xFF80, 0x7FC0, 9.183550e-41, -9.183550e-41]> : tensor<11xbf16>
+     return %cst : tensor<11xbf16>
+   }
+ })mlir";
+@@ -1296,11 +1431,16 @@
+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
+     auto type = makeTensorType(fb.getContext(), {11}, ElementType::BF16);
++    double posInf = std::numeric_limits<double>::infinity();
++    double negInf = -std::numeric_limits<double>::infinity();
++    double nan = std::numeric_limits<double>::quiet_NaN();
++    double posMinDenorm = 0x1p-133;
++    double negMinDenorm = -0x1p-133;
+     auto cst = stablehlo::Constant(
+-        fb,
+-        makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.140630,
+-                                       0x7F80, 0xFF80, 0x7FFF, 0x0001, 0x8001}),
+-                     type));
++        fb, makeConstant(
++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.140630, posInf,
++                                  negInf, nan, posMinDenorm, negMinDenorm}),
++                type));
+     func::Return(fb, {cst});
+   }
+ 
+@@ -1311,7 +1451,7 @@
+ TEST(MlirBuilderTest, ConstantF16Array) {
+   std::string expected = R"mlir(module {
+   func.func @main() -> tensor<11xf16> {
+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.997550e-02, 3.140630e+00, 3.174400e+04, 6.451200e+04, 3.276800e+04, 1.000000e+00, 3.276800e+04]> : tensor<11xf16>
++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.997550e-02, 3.140630e+00, 0x7C00, 0xFC00, 0x7E00, 5.960460e-08, -5.960460e-08]> : tensor<11xf16>
+     return %cst : tensor<11xf16>
+   }
+ })mlir";
+@@ -1321,11 +1461,16 @@
+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
+     auto type = makeTensorType(fb.getContext(), {11}, ElementType::F16);
++    double posInf = std::numeric_limits<double>::infinity();
++    double negInf = -std::numeric_limits<double>::infinity();
++    double nan = std::numeric_limits<double>::quiet_NaN();
++    double posMinDenorm = 0x1p-24;
++    double negMinDenorm = -0x1p-24;
+     auto cst = stablehlo::Constant(
+-        fb,
+-        makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.140630,
+-                                       0x7C00, 0xFC00, 0x7FFF, 0x0001, 0x8001}),
+-                     type));
++        fb, makeConstant(
++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.140630, posInf,
++                                  negInf, nan, posMinDenorm, negMinDenorm}),
++                type));
+     func::Return(fb, {cst});
+   }
+ 
+@@ -1336,7 +1481,7 @@
+ TEST(MlirBuilderTest, ConstantF32Array) {
+   std::string expected = R"mlir(module {
+   func.func @main() -> tensor<11xf32> {
+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000000e-01, 3.14159274, 2.13909504E+9, 4.28657869E+9, 2.14748365E+9, 1.000000e+00, 2.14748365E+9]> : tensor<11xf32>
++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000000e-01, 3.14159274, 0x7F800000, 0xFF800000, 0x7FC00000, 1.401300e-45, -1.401300e-45]> : tensor<11xf32>
+     return %cst : tensor<11xf32>
+   }
+ })mlir";
+@@ -1346,10 +1491,15 @@
+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
+     auto type = makeTensorType(fb.getContext(), {11}, ElementType::F32);
++    double posInf = std::numeric_limits<double>::infinity();
++    double negInf = -std::numeric_limits<double>::infinity();
++    double nan = std::numeric_limits<double>::quiet_NaN();
++    double posMinDenorm = std::numeric_limits<float>::denorm_min();
++    double negMinDenorm = -std::numeric_limits<float>::denorm_min();
+     auto cst = stablehlo::Constant(
+         fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1,
+-                                           3.14159274, 0x7F800000, 0xFF800000,
+-                                           0x7FFFFFFF, 0x00000001, 0x80000001}),
++                                           3.14159274, posInf, negInf, nan,
++                                           posMinDenorm, negMinDenorm}),
+                          type));
+     func::Return(fb, {cst});
+   }
+@@ -1361,7 +1511,7 @@
+ TEST(MlirBuilderTest, ConstantF64Array) {
+   std::string expected = R"mlir(module {
+   func.func @main() -> tensor<11xf64> {
+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000000e-01, 3.1415926535897931, 9.2188684372274053E+18, 1.8442240474082181E+19, 9.2233720368547758E+18, 1.000000e+00, 9.2233720368547758E+18]> : tensor<11xf64>
++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000000e-01, 3.1415926535897931, 0x7FF0000000000000, 0xFFF0000000000000, 0x7FF8000000000000, 4.940660e-324, -4.940660e-324]> : tensor<11xf64>
+     return %cst : tensor<11xf64>
+   }
+ })mlir";
+@@ -1371,13 +1521,16 @@
+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
+     auto type = makeTensorType(fb.getContext(), {11}, ElementType::F64);
++    double posInf = std::numeric_limits<double>::infinity();
++    double negInf = -std::numeric_limits<double>::infinity();
++    double nan = std::numeric_limits<double>::quiet_NaN();
++    double posMinDenorm = std::numeric_limits<double>::denorm_min();
++    double negMinDenorm = -std::numeric_limits<double>::denorm_min();
+     auto cst = stablehlo::Constant(
+-        fb, makeConstant(
+-                ArrayRef<long double>({0.0, -0.0, 1.0, 0.125, 0.1,
+-                                       3.1415926535897931, 0x7FF0000000000000,
+-                                       0xFFF0000000000000, 0x7FFFFFFFFFFFFFFF,
+-                                       0x0000000000000001, 0x8000000000000001}),
+-                type));
++        fb, makeConstant(ArrayRef<long double>(
++                             {0.0, -0.0, 1.0, 0.125, 0.1, 3.1415926535897931,
++                              posInf, negInf, nan, posMinDenorm, negMinDenorm}),
++                         type));
+     func::Return(fb, {cst});
+   }
+ 
 diff --ruN a/stablehlo/stablehlo/integrations/python/StablehloApi.cpp b/stablehlo/stablehlo/integrations/python/StablehloApi.cpp
 --- stablehlo/stablehlo/integrations/python/StablehloApi.cpp
 +++ stablehlo/stablehlo/integrations/python/StablehloApi.cpp
@@ -4182,6 +3942,77 @@ diff --ruN a/stablehlo/stablehlo/reference/Ops.cpp b/stablehlo/stablehlo/referen
      Axes resultPermutation;
      resultPermutation.push_back(outputBatchDimension);
      resultPermutation.append(outputSpatialDimensions.begin(),
+@@ -1942,23 +1969,30 @@
+ Tensor __fft_1d(const Tensor& operand, ShapedType resultType,
+                 int64_t dimension) {
+   Tensor result(resultType);
+-
++  auto complexTy = cast<ComplexType>(resultType.getElementType());
++  const llvm::fltSemantics& sem =
++      cast<FloatType>(complexTy.getElementType()).getFloatSemantics();
+   for (auto resultIt = result.index_begin(); resultIt != result.index_end();
+        ++resultIt) {
+     auto resultIndex = *resultIt;
+     auto N = operand.getShape()[dimension];
+ 
+-    Element X(resultType.getElementType(),
+-              mlir::Complex<APFloat>(APFloat(0.0), APFloat(0.0)));
++    Element X(
++        resultType.getElementType(),
++        mlir::Complex<APFloat>(APFloat::getZero(sem), APFloat::getZero(sem)));
+ 
+     for (int n = 0; n < N; n++) {
+       auto operandIndex = Index(resultIndex);
+       operandIndex[dimension] = n;
+ 
+-      auto floatType =
+-          cast<ComplexType>(operand.getElementType()).getElementType();
+-      Element phase(floatType,
+-                    APFloat(-2 * M_PI * n * resultIndex[dimension] / N));
++      auto floatType = cast<FloatType>(
++          cast<ComplexType>(operand.getElementType()).getElementType());
++
++      auto phaseVal = APFloat(-2 * M_PI * n * resultIndex[dimension] / N);
++      bool losesInfo;
++      phaseVal.convert(floatType.getFloatSemantics(),
++                       APFloat::rmNearestTiesToEven, &losesInfo);
++      Element phase(floatType, phaseVal);
+       auto rotator = complex(cosine(phase), sine(phase));
+       X = X + operand.get(operandIndex) * rotator;
+     }
+@@ -1992,10 +2026,12 @@
+   } else if (fftType == FftType::IFFT) {
+     // compute IFFT as conjugation, FFT, conjugation and normalization
+     Tensor result(operand);
++    auto complexTy = cast<ComplexType>(resultType.getElementType());
++    const llvm::fltSemantics& sem =
++        cast<FloatType>(complexTy.getElementType()).getFloatSemantics();
+     for (auto d = fftDims.begin(); d != fftDims.end(); ++d) {
+       auto divisorValue = mlir::Complex<APFloat>(
+-          APFloat(static_cast<double>(resultType.getDimSize(*d))),
+-          APFloat(0.0));
++          APFloat(sem, resultType.getDimSize(*d)), APFloat::getZero(sem));
+       auto divisor =
+           constantOp(mlir::DenseElementsAttr::get(resultType, divisorValue));
+       result = divideOp(conjugate(__fft_1d(conjugate(result), resultType, *d)),
+@@ -2040,6 +2076,8 @@
+     replica = conjugate(reverseOp(replica, Axes{d_irfft}, replica.getType()));
+     auto complexType =
+         ComplexType::get(cast<FloatType>(resultType.getElementType()));
++    const llvm::fltSemantics& sem =
++        cast<FloatType>(complexType.getElementType()).getFloatSemantics();
+     result = concatenateOp({result, replica}, Axis{d_irfft},
+                            resultType.cloneWith(std::nullopt, complexType));
+ 
+@@ -2048,7 +2086,7 @@
+                     resultType);
+ 
+     // normalize result of inverse transformation
+-    auto divisorValue = APFloat(static_cast<double>(renormFactor));
++    auto divisorValue = APFloat(sem, renormFactor);
+     auto divisor =
+         constantOp(mlir::DenseElementsAttr::get(resultType, divisorValue));
+     return divideOp(result, divisor, resultType);
 diff --ruN a/stablehlo/stablehlo/reference/Ops.h b/stablehlo/stablehlo/reference/Ops.h
 --- stablehlo/stablehlo/reference/Ops.h
 +++ stablehlo/stablehlo/reference/Ops.h
@@ -4828,6 +4659,249 @@ diff --ruN a/stablehlo/stablehlo/tests/interpret/convolution.mlir b/stablehlo/st
 +    dense<[[[3.0], [4.0]], [[6.0], [8.0]]]> : tensor<2x2x1xf32>
 +  func.return
 +}
+diff --ruN a/stablehlo/stablehlo/tests/interpret/fft.mlir b/stablehlo/stablehlo/tests/interpret/fft.mlir
+--- stablehlo/stablehlo/tests/interpret/fft.mlir
++++ stablehlo/stablehlo/tests/interpret/fft.mlir
+@@ -233,3 +233,239 @@
+   > : tensor<2x2x4xf64>
+   func.return
+ }
++
++// -----
++
++func.func @fft_op_test_fft_1d_c64() {
++  %operand = stablehlo.constant dense<[(1.0, 2.0), (3.0, -4.0)]> : tensor<2xcomplex<f32>>
++  %result = stablehlo.fft %operand, type = FFT, length = [2] : (tensor<2xcomplex<f32>>) -> tensor<2xcomplex<f32>>
++  check.expect_almost_eq_const %result, dense<[(4.0, -2.0), (-2.0, 6.0)]> : tensor<2xcomplex<f32>>
++  func.return
++}
++
++// -----
++
++func.func @fft_op_test_fft_2d_c64() {
++  %operand = stablehlo.constant dense<
++    [
++      [(1.0, 2.0), (3.0, -4.0)],
++      [(-5.0, 6.0), (-7.0, -8.0)]
++    ]
++  > : tensor<2x2xcomplex<f32>>
++  %result = stablehlo.fft %operand, type = FFT, length = [2, 2] : (tensor<2x2xcomplex<f32>>) -> tensor<2x2xcomplex<f32>>
++  check.expect_almost_eq_const %result, dense<
++    [
++      [(-8.0, -4.0), (0.0, 20.0)],
++      [(16.0, 0.0), (-4.0, -8.0)]
++    ]
++  > : tensor<2x2xcomplex<f32>>
++  func.return
++}
++
++// -----
++
++func.func @fft_op_test_fft_3d_c64() {
++  %operand = stablehlo.constant dense<
++    [
++      [
++        [(1.0, 2.0), (3.0, -4.0)],
++        [(-5.0, 6.0), (-7.0, -8.0)]
++      ],[
++        [(-1.0, -2.0), (-3.0, 4.0)],
++        [(5.0, -6.0), (7.0, 8.0)]
++      ]
++    ]
++  > : tensor<2x2x2xcomplex<f32>>
++  %result = stablehlo.fft %operand, type = FFT, length = [2, 2, 2] : (tensor<2x2x2xcomplex<f32>>) -> tensor<2x2x2xcomplex<f32>>
++  check.expect_almost_eq_const %result, dense<
++    [
++      [
++        [(0.0, 0.0), (0.0, 0.0)],
++        [(0.0, 0.0), (0.0, 0.0)]
++      ],[
++        [(-16.0, -8.0), (0.0, 40.0)],
++        [(32.0, 0.0), (-8.0, -16.0)]
++      ]
++    ]
++  > : tensor<2x2x2xcomplex<f32>>
++  func.return
++}
++
++// -----
++
++func.func @fft_op_test_ifft_1d_c64() {
++  %operand = stablehlo.constant dense<[(4.0, -2.0), (-2.0, 6.0)]> : tensor<2xcomplex<f32>>
++  %result = stablehlo.fft %operand, type = IFFT, length = [2] : (tensor<2xcomplex<f32>>) -> tensor<2xcomplex<f32>>
++  check.expect_almost_eq_const %result, dense<[(1.0, 2.0), (3.0, -4.0)]> : tensor<2xcomplex<f32>>
++  func.return
++}
++
++// -----
++
++func.func @fft_op_test_ifft_2d_c64() {
++  %operand = stablehlo.constant dense<
++    [
++      [(-8.0, -4.0), (16.0, 0.0)],
++      [(0.0, 20.0), (-4.0, -8.0)]
++    ]
++  > : tensor<2x2xcomplex<f32>>
++  %result = stablehlo.fft %operand, type = IFFT, length = [2, 2] : (tensor<2x2xcomplex<f32>>) -> tensor<2x2xcomplex<f32>>
++  check.expect_almost_eq_const %result, dense<
++    [
++      [(1.0, 2.0), (-5.0, 6.0)],
++      [(3.0, -4.0), (-7.0, -8.0)]
++    ]
++  > : tensor<2x2xcomplex<f32>>
++  func.return
++}
++
++// -----
++
++func.func @fft_op_test_ifft_3d_c64() {
++  %operand = stablehlo.constant dense<
++    [
++      [
++        [(0.0, 0.0), (0.0, 0.0)],
++        [(0.0, 0.0), (0.0, 0.0)]
++      ],[
++        [(-16.0, -8.0), (32.0, 0.0)],
++        [(0.0, 40.0), (-8.0, -16.0)]
++      ]
++    ]
++  > : tensor<2x2x2xcomplex<f32>>
++  %result = stablehlo.fft %operand, type = IFFT, length = [2, 2, 2] : (tensor<2x2x2xcomplex<f32>>) -> tensor<2x2x2xcomplex<f32>>
++  check.expect_almost_eq_const %result, dense<
++    [
++      [
++        [(1.0, 2.0), (-5.0, 6.0)],
++        [(3.0, -4.0), (-7.0, -8.0)]
++      ],[
++        [(-1.0, -2.0), (5.0, -6.0)],
++        [(-3.0, 4.0), (7.0, 8.0)]
++      ]
++    ]
++  > : tensor<2x2x2xcomplex<f32>>
++  func.return
++}
++
++// -----
++
++func.func @fft_op_test_rfft_1d_f32() {
++  %operand = stablehlo.constant dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
++  %result = stablehlo.fft %operand, type = RFFT, length = [4] : (tensor<4xf32>) -> tensor<3xcomplex<f32>>
++  check.expect_almost_eq_const %result, dense<[(10.0, 0.0), (-2.0, 2.0), (-2.0, 0.0)]> : tensor<3xcomplex<f32>>
++  func.return
++}
++
++// -----
++
++func.func @fft_op_test_rfft_2d_f32() {
++  %operand = stablehlo.constant dense<
++    [
++      [1.0, 2.0, 3.0, 4.0],
++      [5.0, 6.0, 7.0, 8.0],
++      [9.0, 10.0, 11.0, 12.0],
++      [13.0, 14.0, 15.0, 16.0]
++    ]
++  > : tensor<4x4xf32>
++  %result = stablehlo.fft %operand, type = RFFT, length = [4, 4] : (tensor<4x4xf32>) -> tensor<4x3xcomplex<f32>>
++  check.expect_almost_eq_const %result, dense<
++    [
++      [(136.0, 0.0), (-8.0, 8.0), (-8.0, 0.0)],
++      [(-32.0, 32.0), (0.0, 0.0), (0.0, 0.0)],
++      [(-32.0, 0.0), (0.0, 0.0), (0.0, 0.0)],
++      [(-32.0, -32.0), (0.0, 0.0), (0.0, 0.0)]
++    ]
++  > : tensor<4x3xcomplex<f32>>
++  func.return
++}
++
++// -----
++
++func.func @fft_op_test_rfft_3d_f32() {
++  %operand = stablehlo.constant dense<
++    [
++      [
++        [1.0, 2.0, 3.0, 4.0],
++        [5.0, 6.0, 7.0, 8.0]
++      ],[
++        [9.0, 10.0, 11.0, 12.0],
++        [13.0, 14.0, 15.0, 16.0]
++      ]
++    ]
++  > : tensor<2x2x4xf32>
++  %result = stablehlo.fft %operand, type = RFFT, length = [2, 2, 4] : (tensor<2x2x4xf32>) -> tensor<2x2x3xcomplex<f32>>
++  check.expect_almost_eq_const %result, dense<
++    [
++      [
++        [(136.0, 0.0), (-8.0, 8.0), (-8.0, 0.0)],
++        [(-32.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
++      ],[
++        [(-64.0, 0.0), (0.0, 0.0), (0.0, 0.0)],
++        [(0.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
++      ]
++    ]
++  > : tensor<2x2x3xcomplex<f32>>
++  func.return
++}
++
++// -----
++
++func.func @fft_op_test_irfft_1d_c64() {
++  %operand = stablehlo.constant dense<[(10.0, 0.0), (-2.0, 2.0), (-2.0, 0.0)]> : tensor<3xcomplex<f32>>
++  %result = stablehlo.fft %operand, type = IRFFT, length = [4] : (tensor<3xcomplex<f32>>) -> tensor<4xf32>
++  check.expect_almost_eq_const %result, dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
++  func.return
++}
++
++// -----
++
++func.func @fft_op_test_irfft_2d_c64() {
++  %operand = stablehlo.constant dense<
++    [
++      [(136.0, 0.0), (-8.0, 8.0), (-8.0, 0.0)],
++      [(-32.0, 32.0), (0.0, 0.0), (0.0, 0.0)],
++      [(-32.0, 0.0), (0.0, 0.0), (0.0, 0.0)],
++      [(-32.0, -32.0), (0.0, 0.0), (0.0, 0.0)]
++    ]
++  > : tensor<4x3xcomplex<f32>>
++  %result = stablehlo.fft %operand, type = IRFFT, length = [4, 4] : (tensor<4x3xcomplex<f32>>) -> tensor<4x4xf32>
++  check.expect_almost_eq_const %result, dense<
++    [
++      [1.0, 2.0, 3.0, 4.0],
++      [5.0, 6.0, 7.0, 8.0],
++      [9.0, 10.0, 11.0, 12.0],
++      [13.0, 14.0, 15.0, 16.0]
++    ]
++  > : tensor<4x4xf32>
++  func.return
++}
++
++// -----
++
++func.func @fft_op_test_irfft_3d_c64() {
++  %operand = stablehlo.constant dense<
++    [
++      [
++        [(136.0, 0.0), (-8.0, 8.0), (-8.0, 0.0)],
++        [(-32.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
++      ],[
++        [(-64.0, 0.0), (0.0, 0.0), (0.0, 0.0)],
++        [(0.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
++      ]
++    ]
++  > : tensor<2x2x3xcomplex<f32>>
++  %result = stablehlo.fft %operand, type = IRFFT, length = [2, 2, 4] : (tensor<2x2x3xcomplex<f32>>) -> tensor<2x2x4xf32>
++  check.expect_almost_eq_const %result, dense<
++    [
++      [
++        [1.0, 2.0, 3.0, 4.0],
++        [5.0, 6.0, 7.0, 8.0]
++      ],[
++        [9.0, 10.0, 11.0, 12.0],
++        [13.0, 14.0, 15.0, 16.0]
++      ]
++    ]
++  > : tensor<2x2x4xf32>
++  func.return
++}
 diff --ruN a/stablehlo/stablehlo/tests/ops_chlo.mlir b/stablehlo/stablehlo/tests/ops_chlo.mlir
 --- stablehlo/stablehlo/tests/ops_chlo.mlir
 +++ stablehlo/stablehlo/tests/ops_chlo.mlir
@@ -13148,363 +13222,375 @@ diff --ruN a/stablehlo/third_party/BUILD b/stablehlo/third_party/BUILD
 @@ -1 +0,0 @@
 -licenses(["notice"])  # Apache 2.0
 \ No newline at end of file
-diff --ruN a/stablehlo/third_party/repo.bzl b/stablehlo/third_party/repo.bzl
---- stablehlo/third_party/repo.bzl
-+++ stablehlo/third_party/repo.bzl
-@@ -1,159 +0,0 @@
--# Copyright 2017 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.
--
--"""Utilities for defining TensorFlow Bazel dependencies."""
--
--def tf_mirror_urls(url):
--    """A helper for generating TF-mirror versions of URLs.
--
--    Given a URL, it returns a list of the TF-mirror cache version of that URL
--    and the original URL, suitable for use in `urls` field of `tf_http_archive`.
--    """
--    if not url.startswith("https://"):
--        return [url]
--    return [
--        "https://storage.googleapis.com/mirror.tensorflow.org/%s" % url[8:],
--        url,
--    ]
--
--def _get_env_var(ctx, name):
--    if name in ctx.os.environ:
--        return ctx.os.environ[name]
--    else:
--        return None
--
--# Checks if we should use the system lib instead of the bundled one
--def _use_system_lib(ctx, name):
--    syslibenv = _get_env_var(ctx, "TF_SYSTEM_LIBS")
--    if not syslibenv:
--        return False
--    return name in [n.strip() for n in syslibenv.split(",")]
--
--def _get_link_dict(ctx, link_files, build_file):
--    link_dict = {ctx.path(v): ctx.path(Label(k)) for k, v in link_files.items()}
--    if build_file:
--        # Use BUILD.bazel because it takes precedence over BUILD.
--        link_dict[ctx.path("BUILD.bazel")] = ctx.path(Label(build_file))
--    return link_dict
--
--def _tf_http_archive_impl(ctx):
--    # Construct all paths early on to prevent rule restart. We want the
--    # attributes to be strings instead of labels because they refer to files
--    # in the TensorFlow repository, not files in repos depending on TensorFlow.
--    # See also https://github.com/bazelbuild/bazel/issues/10515.
--    link_dict = _get_link_dict(ctx, ctx.attr.link_files, ctx.attr.build_file)
--
--    # For some reason, we need to "resolve" labels once before the
--    # download_and_extract otherwise it'll invalidate and re-download the
--    # archive each time.
--    # https://github.com/bazelbuild/bazel/issues/10515
--    patch_files = ctx.attr.patch_file
--    for patch_file in patch_files:
--        if patch_file:
--            ctx.path(Label(patch_file))
--
--    if _use_system_lib(ctx, ctx.attr.name):
--        link_dict.update(_get_link_dict(
--            ctx = ctx,
--            link_files = ctx.attr.system_link_files,
--            build_file = ctx.attr.system_build_file,
--        ))
--    else:
--        ctx.download_and_extract(
--            url = ctx.attr.urls,
--            sha256 = ctx.attr.sha256,
--            type = ctx.attr.type,
--            stripPrefix = ctx.attr.strip_prefix,
--        )
--        if hasattr(ctx.attr, "patch_cmds") and ctx.attr.patch_cmds:
--            for cmd in ctx.attr.patch_cmds:
--                res = ctx.execute(["bash", "-c", cmd])
--                if res.return_code != 0:
--                    fail("patch_cmd failed: %s\n%s" % (cmd, res.stderr))
--        if patch_files:
--            for patch_file in patch_files:
--                patch_file = ctx.path(Label(patch_file)) if patch_file else None
--                if patch_file:
--                    ctx.patch(patch_file, strip = 1)
--
--    for dst, src in link_dict.items():
--        ctx.delete(dst)
--        ctx.symlink(src, dst)
--
--_tf_http_archive = repository_rule(
--    implementation = _tf_http_archive_impl,
--    attrs = {
--        "sha256": attr.string(mandatory = True),
--        "urls": attr.string_list(mandatory = True),
--        "strip_prefix": attr.string(),
--        "type": attr.string(),
--        "patch_file": attr.string_list(),
--        "patch_cmds": attr.string_list(),
--        "build_file": attr.string(),
--        "system_build_file": attr.string(),
--        "link_files": attr.string_dict(),
--        "system_link_files": attr.string_dict(),
--    },
--    environ = ["TF_SYSTEM_LIBS"],
--)
--
--def tf_http_archive(name, sha256, urls, **kwargs):
--    """Downloads and creates Bazel repos for dependencies.
--
--    This is a swappable replacement for both http_archive() and
--    new_http_archive() that offers some additional features. It also helps
--    ensure best practices are followed.
--
--    File arguments are relative to the TensorFlow repository by default. Dependent
--    repositories that use this rule should refer to files either with absolute
--    labels (e.g. '@foo//:bar') or from a label created in their repository (e.g.
--    'str(Label("//:bar"))').
--    """
--    if len(urls) < 2:
--        fail("tf_http_archive(urls) must have redundant URLs.")
--
--    if not any([mirror in urls[0] for mirror in (
--        "mirror.tensorflow.org",
--        "mirror.bazel.build",
--        "storage.googleapis.com",
--    )]):
--        fail("The first entry of tf_http_archive(urls) must be a mirror " +
--             "URL, preferrably mirror.tensorflow.org. Even if you don't have " +
--             "permission to mirror the file, please put the correctly " +
--             "formatted mirror URL there anyway, because someone will come " +
--             "along shortly thereafter and mirror the file.")
--
--    if native.existing_rule(name):
--        return
--
--    _tf_http_archive(
--        name = name,
--        sha256 = sha256,
--        urls = urls,
--        **kwargs
--    )
--
--def _tf_vendored_impl(ctx):
--    ctx.symlink(ctx.path(ctx.attr._root).dirname.get_child(ctx.attr.path), ".")
--
--tf_vendored = repository_rule(
--    implementation = _tf_vendored_impl,
--    doc = "Similar to local_repository, but path is relative to the root of the " +
--          "repository, not the root of the workspace.",
--    attrs = {
--        "_root": attr.label(default = "//:unused"),
--        "path": attr.string(),
--    },
--    local = True,
--)
+diff --ruN a/stablehlo/third_party/BUILD.bazel b/stablehlo/third_party/BUILD.bazel
+--- stablehlo/third_party/BUILD.bazel
++++ stablehlo/third_party/BUILD.bazel
+@@ -0,0 +1,12 @@
++package(
++    # copybara:uncomment default_applicable_licenses = ["//:license"],
++    licenses = ["notice"],
++)
++
++# copybara:uncomment_begin
++# filegroup(
++#     name = "all_files",
++#     srcs = glob(["**"]),
++#     visibility = ["//:__subpackages__"],
++# )
++# copybara:uncomment_end
 diff --ruN a/stablehlo/third_party/rules_ml_toolchain/BUILD b/stablehlo/third_party/rules_ml_toolchain/BUILD
 --- stablehlo/third_party/rules_ml_toolchain/BUILD
 +++ stablehlo/third_party/rules_ml_toolchain/BUILD
 @@ -1 +0,0 @@
 -exports_files(["fix_strip.patch"])
-diff --ruN a/stablehlo/third_party/rules_ml_toolchain/fix_strip.patch b/stablehlo/third_party/rules_ml_toolchain/fix_strip.patch
---- stablehlo/third_party/rules_ml_toolchain/fix_strip.patch
-+++ stablehlo/third_party/rules_ml_toolchain/fix_strip.patch
-@@ -1,30 +0,0 @@
--diff --git a/cc/impls/linux_x86_64_linux_x86_64/wrappers/BUILD b/cc/impls/linux_x86_64_linux_x86_64/wrappers/BUILD
--index a206127..fa119c10 100644
----- a/cc/impls/linux_x86_64_linux_x86_64/wrappers/BUILD
--+++ b/cc/impls/linux_x86_64_linux_x86_64/wrappers/BUILD
--@@ -24,5 +24,6 @@ filegroup(
--         "idler",
--         "ld",
--         "objcopy",
--+        "strip",
--     ],
--     visibility = ["//visibility:public"],
-- )
--diff --git a/cc/impls/linux_x86_64_linux_x86_64/wrappers/strip b/cc/impls/linux_x86_64_linux_x86_64/wrappers/strip
--new file mode 100755
--index 0000000..e3b0923
----- /dev/null
--+++ b/cc/impls/linux_x86_64_linux_x86_64/wrappers/strip
--@@ -0,0 +1,3 @@
--+#!/bin/bash
--+
--+"$STRIP_PATH" "$@"
--diff --git a/third_party/rules_cc_toolchain/toolchain_config.bzl b/third_party/rules_cc_toolchain/toolchain_config.bzl
--index 7381bbb..23afe63 100644
----- a/third_party/rules_cc_toolchain/toolchain_config.bzl
--+++ b/third_party/rules_cc_toolchain/toolchain_config.bzl
--@@ -58,3 +58,4 @@ ALL_ACTIONS = [
--     ACTION_NAMES.cpp_link_static_library,
--     ACTION_NAMES.clif_match,
--+    ACTION_NAMES.strip,
-- ]
-diff --ruN a/stablehlo/third_party/zlib-ng.BUILD b/stablehlo/third_party/zlib-ng.BUILD
---- stablehlo/third_party/zlib-ng.BUILD
-+++ stablehlo/third_party/zlib-ng.BUILD
-@@ -1,106 +0,0 @@
--# This file is licensed under the Apache License v2.0 with LLVM Exceptions.
--# See https://llvm.org/LICENSE.txt for license information.
--# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
--
--load("@bazel_skylib//rules:copy_file.bzl", "copy_file")
--load("@rules_cc//cc:cc_library.bzl", "cc_library")
--
--package(
--    default_visibility = ["//visibility:public"],
--    # BSD/MIT-like license (for zlib)
--    licenses = ["notice"],
--)
--
--copy_file(
--    # The input template is identical to the CMake output.
--    name = "zconf_gen",
--    src = "zconf.h.in",
--    out = "zconf.h",
--    allow_symlink = True,
--)
--
--cc_library(
--    name = "zlib",
--    srcs = select({
--        "@llvm-project//third-party:llvm_zlib_enabled": [
--            "adler32.c",
--            "adler32_p.h",
--            "chunkset.c",
--            "chunkset_tpl.h",
--            "compare258.c",
--            "compress.c",
--            "crc32.c",
--            "crc32_comb.c",
--            "crc32_comb_tbl.h",
--            "crc32_p.h",
--            "crc32_tbl.h",
--            "deflate.c",
--            "deflate.h",
--            "deflate_fast.c",
--            "deflate_medium.c",
--            "deflate_p.h",
--            "deflate_quick.c",
--            "deflate_slow.c",
--            "fallback_builtins.h",
--            "functable.c",
--            "functable.h",
--            "infback.c",
--            "inffast.c",
--            "inffast.h",
--            "inffixed_tbl.h",
--            "inflate.c",
--            "inflate.h",
--            "inflate_p.h",
--            "inftrees.c",
--            "inftrees.h",
--            "insert_string.c",
--            "insert_string_tpl.h",
--            "match_tpl.h",
--            "trees.c",
--            "trees.h",
--            "trees_emit.h",
--            "trees_tbl.h",
--            "uncompr.c",
--            "zbuild.h",
--            "zendian.h",
--            "zutil.c",
--            "zutil.h",
--            "zutil_p.h",
--        ],
--        "//conditions:default": [],
--    }),
--    hdrs = select({
--        "@llvm-project//third-party:llvm_zlib_enabled": [
--            "zlib.h",
--            ":zconf_gen",
--        ],
--        "//conditions:default": [],
--    }),
--    copts = [
--        "-std=c11",
--        "-DZLIB_COMPAT",
--        "-DWITH_GZFILEOP",
--        "-DWITH_OPTIM",
--        "-DWITH_NEW_STRATEGIES",
--        # For local builds you might want to add "-DWITH_NATIVE_INSTRUCTIONS"
--        # here to improve performance. Native instructions aren't enabled in
--        # the default config for reproducibility.
--    ],
--    defines = select({
--        "@llvm-project//third-party:llvm_zlib_enabled": [
--            "LLVM_ENABLE_ZLIB=1",
--        ],
--        "//conditions:default": [
--            "LLVM_ENABLE_ZLIB=0",
--        ],
--    }),
--    # Clang includes zlib with angled instead of quoted includes, so we need
--    # strip_include_prefix here.
--    strip_include_prefix = ".",
--    visibility = ["//visibility:public"],
--)
--
--alias(
--    name = "zlib-ng",
--    actual = ":zlib",
--)
-diff --ruN a/stablehlo/third_party/zstd.BUILD b/stablehlo/third_party/zstd.BUILD
---- stablehlo/third_party/zstd.BUILD
-+++ stablehlo/third_party/zstd.BUILD
-@@ -1,44 +0,0 @@
--# This file is licensed under the Apache License v2.0 with LLVM Exceptions.
--# See https://llvm.org/LICENSE.txt for license information.
--# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
--load("@rules_cc//cc:cc_library.bzl", "cc_library")
--
--package(
--    default_visibility = ["//visibility:public"],
--    # BSD/MIT-like license (for zstd)
--    licenses = ["notice"],
--)
--
--cc_library(
--    name = "zstd",
--    srcs = select({
--        "@llvm-project//third-party:llvm_zstd_enabled": glob([
--            "lib/common/*.c",
--            "lib/common/*.h",
--            "lib/compress/*.c",
--            "lib/compress/*.h",
--            "lib/decompress/*.c",
--            "lib/decompress/*.h",
--            "lib/decompress/*.S",
--            "lib/dictBuilder/*.c",
--            "lib/dictBuilder/*.h",
--        ]),
--        "//conditions:default": [],
--    }),
--    hdrs = select({
--        "@llvm-project//third-party:llvm_zstd_enabled": [
--            "lib/zdict.h",
--            "lib/zstd.h",
--            "lib/zstd_errors.h",
--        ],
--        "//conditions:default": [],
--    }),
--    defines = select({
--        "@llvm-project//third-party:llvm_zstd_enabled": [
--            "LLVM_ENABLE_ZSTD=1",
--            "ZSTD_MULTITHREAD",
--        ],
--        "//conditions:default": [],
--    }),
--    strip_include_prefix = "lib",
--)
+diff --ruN a/stablehlo/third_party/rules_ml_toolchain/BUILD.bazel b/stablehlo/third_party/rules_ml_toolchain/BUILD.bazel
+--- stablehlo/third_party/rules_ml_toolchain/BUILD.bazel
++++ stablehlo/third_party/rules_ml_toolchain/BUILD.bazel
+@@ -0,0 +1,14 @@
++package(
++    # copybara:uncomment default_applicable_licenses = ["//:license"],
++    licenses = ["notice"],
++)
++
++# copybara:uncomment_begin
++# filegroup(
++#     name = "all_files",
++#     srcs = glob(["**"]),
++#     visibility = ["//:__subpackages__"],
++# )
++# copybara:uncomment_end
++
++exports_files(["fix_strip.patch"])
+diff --ruN a/stablehlo/third_party/xz.BUILD b/stablehlo/third_party/xz.BUILD
+--- stablehlo/third_party/xz.BUILD
++++ stablehlo/third_party/xz.BUILD
+@@ -0,0 +1,329 @@
++# Copyright 2026 The StableHLO 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
++#
++#      https://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.
++
++load("@bazel_skylib//rules:write_file.bzl", "write_file")
++load("@rules_cc//cc:cc_library.bzl", "cc_library")
++
++package(default_visibility = ["//visibility:private"])
++
++licenses(["unencumbered"])
++
++exports_files(["COPYING"])
++
++cc_library(
++    name = "lzma",
++    srcs = glob(
++        include = [
++            "src/**/*.h",
++            "src/liblzma/**/*.c",
++        ],
++        exclude = [
++            "src/**/*_tablegen.c",
++            "src/liblzma/check/crc_clmul_consts_gen.c",
++            "src/liblzma/check/sha256.c",
++            "src/liblzma/check/crc*_small.c",
++            "src/liblzma/common/stream_decoder_mt.c",
++            "src/liblzma/common/stream_encoder_mt.c",
++        ],
++    ) + select({
++        "@platforms//os:osx": [],
++        "@platforms//os:ios": [],
++        "//conditions:default": [
++            "src/liblzma/check/sha256.c",
++        ],
++    }) + select({
++        "@platforms//os:none": [],
++        "//conditions:default": [
++            "src/common/tuklib_cpucores.c",
++            "src/common/tuklib_physmem.c",
++            "src/liblzma/common/stream_decoder_mt.c",
++            "src/liblzma/common/stream_encoder_mt.c",
++        ],
++    }),
++    hdrs = glob(["src/liblzma/api/**/*.h"]),
++    copts = select({
++        "@platforms//os:windows": [],
++        "//conditions:default": ["-std=c99"],
++    }),
++    defines = select({
++        "@platforms//os:windows": ["LZMA_API_STATIC"],
++        "//conditions:default": [],
++    }),
++    implementation_deps = [
++        ":config",
++        ":src_common",
++        ":src_liblzma",
++        ":src_liblzma_check",
++        ":src_liblzma_common",
++        ":src_liblzma_delta",
++        ":src_liblzma_lz",
++        ":src_liblzma_lzma",
++        ":src_liblzma_rangecoder",
++        ":src_liblzma_simple",
++    ],
++    linkopts = select({
++        "@platforms//os:android": [],
++        "@platforms//os:windows": [],
++        "@platforms//os:none": [],
++        "//conditions:default": ["-lpthread"],
++    }),
++    linkstatic = select({
++        "@platforms//os:windows": True,
++        "//conditions:default": False,
++    }),
++    local_defines = ["HAVE_CONFIG_H"],
++    strip_include_prefix = "src/liblzma/api",
++    visibility = ["//visibility:public"],
++)
++
++write_file(
++    name = "config_gen",
++    out = "lzma_config/config.h",
++    content = [
++        "#define ASSUME_RAM 128",
++        "#ifndef _WIN32",
++        "#define HAVE_BSWAP_16 1",
++        "#define HAVE_BSWAP_32 1",
++        "#define HAVE_BSWAP_64 1",
++        "#endif",
++        "#if defined(__APPLE__) || defined(_WIN32)",
++        "#undef HAVE_BYTESWAP_H",
++        "#else",
++        "#define HAVE_BYTESWAP_H 1",
++        "#endif",
++        "#undef HAVE_CAPSICUM",
++        "#if defined(__APPLE__)",
++        "#define HAVE_CC_SHA256_CTX 1",
++        "#define HAVE_CC_SHA256_INIT 1",
++        "#define HAVE_COMMONCRYPTO_COMMONDIGEST_H 1",
++        "#endif",
++        "#define HAVE_CHECK_CRC32 1",
++        "#define HAVE_CHECK_CRC64 1",
++        "#define HAVE_CHECK_SHA256 1",
++        "#if !defined(__APPLE__)",
++        "#define HAVE_CLOCK_GETTIME 1",
++        "#define HAVE_DECL_CLOCK_MONOTONIC 1",
++        "#endif",
++        "#if defined(__ANDROID__) || defined(__APPLE__) || defined(_WIN32)",
++        "#define HAVE_DECL_PROGRAM_INVOCATION_NAME 0",
++        "#else",
++        "#define HAVE_DECL_PROGRAM_INVOCATION_NAME 1",
++        "#endif",
++        "#define HAVE_DECODERS 1",
++        "#define HAVE_DECODER_ARM 1",
++        "#define HAVE_DECODER_ARM64 1",
++        "#define HAVE_DECODER_ARMTHUMB 1",
++        "#define HAVE_DECODER_DELTA 1",
++        "#define HAVE_DECODER_IA64 1",
++        "#define HAVE_DECODER_LZMA1 1",
++        "#define HAVE_DECODER_LZMA2 1",
++        "#define HAVE_DECODER_POWERPC 1",
++        "#define HAVE_DECODER_SPARC 1",
++        "#define HAVE_DECODER_X86 1",
++        "#define HAVE_DLFCN_H 1",
++        "#define HAVE_ENCODERS 1",
++        "#define HAVE_ENCODER_ARM 1",
++        "#define HAVE_ENCODER_ARM64 1",
++        "#define HAVE_ENCODER_ARMTHUMB 1",
++        "#define HAVE_ENCODER_DELTA 1",
++        "#define HAVE_ENCODER_IA64 1",
++        "#define HAVE_ENCODER_LZMA1 1",
++        "#define HAVE_ENCODER_LZMA2 1",
++        "#define HAVE_ENCODER_POWERPC 1",
++        "#define HAVE_ENCODER_SPARC 1",
++        "#define HAVE_ENCODER_X86 1",
++        "#define HAVE_FCNTL_H 1",
++        "#define HAVE_FUTIMENS 1",
++        "#define HAVE_GETOPT_H 1",
++        "#define HAVE_GETOPT_LONG 1",
++        "#define HAVE_INTTYPES_H 1",
++        "#define HAVE_LIMITS_H 1",
++        "#define HAVE_MBRTOWC 1",
++        "#define HAVE_MEMORY_H 1",
++        "#define HAVE_MF_BT2 1",
++        "#define HAVE_MF_BT3 1",
++        "#define HAVE_MF_BT4 1",
++        "#define HAVE_MF_HC3 1",
++        "#define HAVE_MF_HC4 1",
++        "#if !defined(__APPLE__)",
++        "#define HAVE_POSIX_FADVISE 1",
++        "#endif",
++        "#if !defined(__APPLE__) && !defined(__ANDROID__)",
++        "#define HAVE_PTHREAD_CONDATTR_SETCLOCK 1",
++        "#endif",
++        "#define HAVE_PTHREAD_PRIO_INHERIT 1",
++        "#undef HAVE_SMALL",
++        "#define HAVE_STDBOOL_H 1",
++        "#define HAVE_STDINT_H 1",
++        "#define HAVE_STDLIB_H 1",
++        "#define HAVE_STRING_H 1",
++        "#define HAVE_STRUCT_STAT_ST_ATIM_TV_NSEC 1",
++        "#undef HAVE_SYS_CAPSICUM_H",
++        "#ifndef _WIN32",
++        "#define HAVE_SYS_PARAM_H 1",
++        "#define HAVE_SYS_STAT_H 1",
++        "#define HAVE_SYS_TIME_H 1",
++        "#define HAVE_SYS_TYPES_H 1",
++        "#endif",
++        "#define HAVE_UINTPTR_T 1",
++        "#define HAVE_UNISTD_H 1",
++        "#define HAVE_VISIBILITY 0",
++        "#ifndef _WIN32",
++        "#define HAVE_WCWIDTH 1",
++        "#endif",
++        "#define HAVE__BOOL 1",
++        "#undef HAVE__FUTIME",
++        "#define LT_OBJDIR \".libs/\"",
++        "#ifdef _WIN32",
++        "#define MYTHREAD_VISTA 1",
++        "#else",
++        "#define MYTHREAD_POSIX 1",
++        "#endif",
++        "#define NDEBUG 1",
++        "#define PACKAGE \"xz\"",
++        "#define PACKAGE_BUGREPORT \"lasse.collin@tukaani.org\"",
++        "#define PACKAGE_NAME \"XZ Utils\"",
++        "#define PACKAGE_STRING \"XZ Utils 5.8.3\"",
++        "#define PACKAGE_TARNAME \"xz\"",
++        "#define PACKAGE_URL \"https://tukaani.org/xz/\"",
++        "#define PACKAGE_VERSION \"5.8.3\"",
++        "#define SIZEOF_SIZE_T 8",
++        "#define STDC_HEADERS 1",
++        "#define TUKLIB_CPUCORES_SYSCONF 1",
++        "#define TUKLIB_FAST_UNALIGNED_ACCESS 1",
++        "#if defined(__APPLE__)",
++        "#undef TUKLIB_PHYSMEM_SYSCONF",
++        "#else",
++        "#define TUKLIB_PHYSMEM_SYSCONF 1",
++        "#endif",
++        "#ifndef _ALL_SOURCE",
++        "#define _ALL_SOURCE 1",
++        "#endif",
++        "#ifndef _GNU_SOURCE",
++        "#define _GNU_SOURCE 1",
++        "#endif",
++        "#ifndef _POSIX_PTHREAD_SEMANTICS",
++        "#define _POSIX_PTHREAD_SEMANTICS 1",
++        "#endif",
++        "#ifndef _TANDEM_SOURCE",
++        "#define _TANDEM_SOURCE 1",
++        "#endif",
++        "#ifndef __EXTENSIONS__",
++        "#define __EXTENSIONS__ 1",
++        "#endif",
++        "#define VERSION \"5.8.3\"",
++        "#if defined AC_APPLE_UNIVERSAL_BUILD",
++        "#if defined __BIG_ENDIAN__",
++        "#define WORDS_BIGENDIAN 1",
++        "#endif",
++        "#endif",
++        "#ifndef _DARWIN_USE_64_BIT_INODE",
++        "#define _DARWIN_USE_64_BIT_INODE 1",
++        "#endif",
++        "",
++    ],
++)
++
++cc_library(
++    name = "config",
++    hdrs = ["lzma_config/config.h"],
++    strip_include_prefix = "lzma_config",
++)
++
++cc_library(
++    name = "src_common",
++    srcs = select({
++        "@platforms//os:none": [],
++        "//conditions:default": [
++            "src/common/tuklib_exit.c",
++            "src/common/tuklib_progname.c",
++        ],
++    }),
++    hdrs = glob(["src/common/*.h"]),
++    defines = select({
++        "@platforms//os:windows": ["LZMA_API_STATIC"],
++        "//conditions:default": [],
++    }),
++    local_defines = ["HAVE_CONFIG_H"] + select({
++        "@platforms//os:windows": ["TUKLIB_GETTEXT=0"],
++        "//conditions:default": [],
++    }),
++    strip_include_prefix = "src/common",
++    deps = [
++        ":config",
++        ":src_liblzma",
++        ":src_liblzma_api",
++        ":src_liblzma_common",
++    ],
++)
++
++cc_library(
++    name = "src_liblzma",
++    hdrs = glob([
++        "src/liblzma/common/*.h",
++        "src/liblzma/lzma/*.h",
++    ]),
++    strip_include_prefix = "src/liblzma",
++)
++
++cc_library(
++    name = "src_liblzma_api",
++    hdrs = glob(
++        include = ["src/liblzma/api/**/*.h"],
++        exclude = ["src/liblzma/api/lzma.h"],
++    ),
++    strip_include_prefix = "src/liblzma/api",
++)
++
++cc_library(
++    name = "src_liblzma_check",
++    hdrs = glob(["src/liblzma/check/*.h"]),
++    strip_include_prefix = "src/liblzma/check",
++)
++
++cc_library(
++    name = "src_liblzma_common",
++    hdrs = glob(["src/liblzma/common/*.h"]),
++    strip_include_prefix = "src/liblzma/common",
++)
++
++cc_library(
++    name = "src_liblzma_delta",
++    hdrs = glob(["src/liblzma/delta/*.h"]),
++    strip_include_prefix = "src/liblzma/delta",
++)
++
++cc_library(
++    name = "src_liblzma_lz",
++    hdrs = glob(["src/liblzma/lz/*.h"]),
++    strip_include_prefix = "src/liblzma/lz",
++)
++
++cc_library(
++    name = "src_liblzma_lzma",
++    hdrs = glob(["src/liblzma/lzma/*.h"]),
++    strip_include_prefix = "src/liblzma/lzma",
++)
++
++cc_library(
++    name = "src_liblzma_rangecoder",
++    hdrs = glob(["src/liblzma/rangecoder/*.h"]),
++    strip_include_prefix = "src/liblzma/rangecoder",
++)
++
++cc_library(
++    name = "src_liblzma_simple",
++    hdrs = glob(["src/liblzma/simple/*.h"]),
++    strip_include_prefix = "src/liblzma/simple",
++)

From d07949689b666f50286920fd2e0ce9915e2102d2 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Wed, 14 Oct 2026 09:19:35 -0700
Subject: [PATCH 0793/1000] [XLA:GPU] Migrate MultiGpuBarrier kernel cache and
 CollectiveMemoryCache to PerDeviceState.

PiperOrigin-RevId: 999417480
---
 .../xla/xla/backends/gpu/runtime/BUILD        |  24 ++-
 .../runtime/collective_kernel_thunk_test.cc   |   4 +-
 .../backends/gpu/runtime/collective_memory.cc |  11 +-
 .../gpu/runtime/collective_memory_cache.cc    |  55 ++++--
 .../gpu/runtime/collective_memory_cache.h     |  35 ++--
 .../runtime/collective_memory_cache_test.cc   | 170 ++++++++++++++++++
 .../backends/gpu/runtime/multi_gpu_barrier.cc |  40 ++---
 .../xla/xla/service/gpu/gpu_executable.cc     |   2 +
 8 files changed, 279 insertions(+), 62 deletions(-)
 create mode 100644 third_party/xla/xla/backends/gpu/runtime/collective_memory_cache_test.cc

diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index 700e94ea739..8eabf313313 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -2921,18 +2921,37 @@ cc_library(
     hdrs = ["collective_memory_cache.h"],
     compatible_with = get_compatible_with_portable(),
     deps = [
+        ":per_device_state",
         "//xla/backends/gpu/collectives:gpu_clique_key",
         "//xla/core/collectives:symmetric_memory",
         "//xla/stream_executor:device_address",
-        "//xla/stream_executor:memory_allocation",
         "//xla/stream_executor/gpu:multicast_memory",
         "//xla/tsl/util:tied_ref",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/synchronization",
     ],
 )
 
+xla_cc_test(
+    name = "collective_memory_cache_test",
+    srcs = ["collective_memory_cache_test.cc"],
+    deps = [
+        ":collective_memory_cache",
+        "//xla/backends/gpu/collectives:gpu_clique_key",
+        "//xla/core/collectives:symmetric_memory",
+        "//xla/runtime:device_id",
+        "//xla/stream_executor:device_address",
+        "//xla/stream_executor/gpu:multicast_memory",
+        "//xla/tsl/platform:env",
+        "//xla/tsl/util:tied_ref",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "collective_cliques",
     srcs = ["collective_cliques.cc"],
@@ -3436,6 +3455,7 @@ cc_library(
     hdrs = ["multi_gpu_barrier.h"],
     compatible_with = get_compatible_with_portable(),
     deps = [
+        ":per_device_state",
         "//xla:status_macros",
         "//xla/core/collectives:rank_id",
         "//xla/core/collectives:symmetric_memory",
@@ -3450,11 +3470,9 @@ cc_library(
         "//xla/stream_executor/gpu:multi_gpu_barrier_kernel",
         "@com_google_absl//absl/base:no_destructor",
         "@com_google_absl//absl/base:nullability",
-        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/synchronization",
     ],
 )
 
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc
index 5cdbd51a9d7..eaabcb72a01 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_kernel_thunk_test.cc
@@ -376,7 +376,7 @@ absl::StatusOr<se::DeviceAddressBase> RunCollectiveKernelThunk(
                                       &buffer_allocations};
 
   ABSL_RETURN_IF_ERROR(metadata.thunk->Prepare(prepare_params));
-  CollectiveMemoryCache collective_memory_cache;
+  CollectiveMemoryCache collective_memory_cache(/*num_devices=*/2);
   CollectiveCliques collective_cliques;
   ABSL_ASSIGN_OR_RETURN(collective_cliques, AcquireCollectiveCliques(
                                            collective_params, clique_requests));
@@ -635,7 +635,7 @@ TEST(CollectiveKernelThunkTest, RecordCommandBufferCreateUpdate) {
                                       &allocations1};
   ASSERT_OK(collective_kernel_thunk->Prepare(prepare_params));
 
-  CollectiveMemoryCache collective_memory_cache;
+  CollectiveMemoryCache collective_memory_cache(/*num_devices=*/2);
   ASSERT_OK_AND_ASSIGN(
       CollectiveCliques collective_cliques,
       AcquireCollectiveCliques(collective_params, clique_requests));
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_memory.cc b/third_party/xla/xla/backends/gpu/runtime/collective_memory.cc
index dc7c4732082..f2dc5de53ae 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_memory.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_memory.cc
@@ -267,7 +267,8 @@ AcquireSymmetricMemory(
       se::DeviceAddressBase addr = buffers.GetDeviceAddress(i);
       CollectiveMemory::Key mem_key = std::make_pair(r.clique, i);
       // Check cache first to avoid redundant collective window registration.
-      if (auto cached = memory_cache.FindSymmetricMemory(r.clique, addr)) {
+      if (auto cached = memory_cache.FindSymmetricMemory(
+              params.executor->device_ordinal(), r.clique, addr)) {
         sym_memories[mem_key] = std::move(cached);
         continue;
       }
@@ -276,7 +277,8 @@ AcquireSymmetricMemory(
       ABSL_ASSIGN_OR_RETURN(tsl::TiedRef<SymmetricMemory> tied_symm,
                        cliques.Tie(r.clique, std::move(symm)));
       sym_memories[mem_key] =
-          memory_cache.AddSymmetricMemory(r.clique, addr, std::move(tied_symm));
+          memory_cache.AddSymmetricMemory(params.executor->device_ordinal(),
+                                          r.clique, addr, std::move(tied_symm));
     }
   }
 
@@ -416,8 +418,9 @@ absl::StatusOr<MulticastMemoryMap> AcquireMulticastMemory(
         se::DeviceAddressBase mcast_addr =
             params[0]->buffers.GetDeviceAddress(i);
         std::shared_ptr<se::gpu::MulticastMemory> cached =
-            memory_cache.AddMulticastMemory(r.clique, mcast_addr,
-                                            std::move(tied_multicast_memory));
+            memory_cache.AddMulticastMemory(
+                params[0]->executor->device_ordinal(), r.clique, mcast_addr,
+                std::move(tied_multicast_memory));
 
         clique_mcast_memories[std::make_pair(r.clique, i)] =
             MappedMulticastMemory{std::move(cached), std::move(mapped_ptrs)};
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_memory_cache.cc b/third_party/xla/xla/backends/gpu/runtime/collective_memory_cache.cc
index 5c9a1bc9e76..f62cecd6ef1 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_memory_cache.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_memory_cache.cc
@@ -18,56 +18,73 @@ limitations under the License.
 #include <memory>
 #include <utility>
 
+#include "absl/log/check.h"
+#include "absl/status/statusor.h"
 #include "absl/synchronization/mutex.h"
 #include "xla/backends/gpu/collectives/gpu_clique_key.h"
 #include "xla/core/collectives/symmetric_memory.h"
 #include "xla/stream_executor/device_address.h"
 #include "xla/stream_executor/gpu/multicast_memory.h"
-#include "xla/stream_executor/memory_allocation.h"
 #include "xla/tsl/util/tied_ref.h"
 
 namespace xla::gpu {
 
 std::shared_ptr<SymmetricMemory> CollectiveMemoryCache::FindSymmetricMemory(
-    const GpuCliqueKey& clique, stream_executor::DeviceAddressBase addr) {
-  absl::MutexLock lock(mutex_);
-  auto it = sym_memories_.find(Key{clique, addr});
-  if (it == sym_memories_.end()) {
+    int device_ordinal, const GpuCliqueKey& clique,
+    stream_executor::DeviceAddressBase addr) {
+  absl::StatusOr<DeviceCache*> device_cache =
+      device_caches_.GetOrCreate(device_ordinal);
+  CHECK_OK(device_cache.status());
+  absl::MutexLock lock((*device_cache)->mutex);
+  auto it = (*device_cache)->sym_memories.find(Key{clique, addr});
+  if (it == (*device_cache)->sym_memories.end()) {
     return nullptr;
   }
   return it->second.Lock();
 }
 
 std::shared_ptr<SymmetricMemory> CollectiveMemoryCache::AddSymmetricMemory(
-    const GpuCliqueKey& clique, stream_executor::DeviceAddressBase addr,
+    int device_ordinal, const GpuCliqueKey& clique,
+    stream_executor::DeviceAddressBase addr,
     tsl::TiedRef<SymmetricMemory> sym_memory) {
-  absl::MutexLock lock(mutex_);
-  auto [it, _] =
-      sym_memories_.insert_or_assign(Key{clique, addr}, std::move(sym_memory));
+  absl::StatusOr<DeviceCache*> device_cache =
+      device_caches_.GetOrCreate(device_ordinal);
+  CHECK_OK(device_cache.status());
+  absl::MutexLock lock((*device_cache)->mutex);
+  auto [it, _] = (*device_cache)
+                     ->sym_memories.insert_or_assign(Key{clique, addr},
+                                                     std::move(sym_memory));
   return it->second.Lock();
 }
 
 std::shared_ptr<stream_executor::gpu::MulticastMemory>
 CollectiveMemoryCache::AddMulticastMemory(
-    const GpuCliqueKey& clique, stream_executor::DeviceAddressBase addr,
+    int device_ordinal, const GpuCliqueKey& clique,
+    stream_executor::DeviceAddressBase addr,
     tsl::TiedRef<stream_executor::gpu::MulticastMemory> multicast_memory) {
-  absl::MutexLock lock(mutex_);
-  auto [it, _] = multicast_memories_.insert_or_assign(
-      Key{clique, addr}, std::move(multicast_memory));
+  absl::StatusOr<DeviceCache*> device_cache =
+      device_caches_.GetOrCreate(device_ordinal);
+  CHECK_OK(device_cache.status());
+  absl::MutexLock lock((*device_cache)->mutex);
+  auto [it, _] = (*device_cache)
+                     ->multicast_memories.insert_or_assign(
+                         Key{clique, addr}, std::move(multicast_memory));
   return it->second.Lock();
 }
 
 std::shared_ptr<stream_executor::gpu::MulticastMemory>
 CollectiveMemoryCache::FindMulticastMemory(
-    const GpuCliqueKey& clique, stream_executor::DeviceAddressBase addr) {
-  absl::MutexLock lock(mutex_);
-  auto it = multicast_memories_.find(Key{clique, addr});
-  if (it == multicast_memories_.end()) {
+    int device_ordinal, const GpuCliqueKey& clique,
+    stream_executor::DeviceAddressBase addr) {
+  absl::StatusOr<DeviceCache*> device_cache =
+      device_caches_.GetOrCreate(device_ordinal);
+  CHECK_OK(device_cache.status());
+  absl::MutexLock lock((*device_cache)->mutex);
+  auto it = (*device_cache)->multicast_memories.find(Key{clique, addr});
+  if (it == (*device_cache)->multicast_memories.end()) {
     return nullptr;
   }
   return it->second.Lock();
 }
 
-
-
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_memory_cache.h b/third_party/xla/xla/backends/gpu/runtime/collective_memory_cache.h
index 0c5e40f0dbc..272137876d5 100644
--- a/third_party/xla/xla/backends/gpu/runtime/collective_memory_cache.h
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_memory_cache.h
@@ -16,7 +16,6 @@ limitations under the License.
 #ifndef XLA_BACKENDS_GPU_RUNTIME_COLLECTIVE_MEMORY_CACHE_H_
 #define XLA_BACKENDS_GPU_RUNTIME_COLLECTIVE_MEMORY_CACHE_H_
 
-#include <cstdint>
 #include <memory>
 #include <utility>
 
@@ -24,10 +23,10 @@ limitations under the License.
 #include "absl/container/flat_hash_map.h"
 #include "absl/synchronization/mutex.h"
 #include "xla/backends/gpu/collectives/gpu_clique_key.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/core/collectives/symmetric_memory.h"
 #include "xla/stream_executor/device_address.h"
 #include "xla/stream_executor/gpu/multicast_memory.h"
-#include "xla/stream_executor/memory_allocation.h"
 #include "xla/tsl/util/tied_ref.h"
 
 namespace xla::gpu {
@@ -47,35 +46,45 @@ namespace xla::gpu {
 // valid as long as the clique exists.
 class CollectiveMemoryCache {
  public:
+  explicit CollectiveMemoryCache(int num_devices)
+      : device_caches_(num_devices) {}
+
   // Returns a locked shared_ptr to the cached memory, or nullptr if not cached
   // or if the underlying TiedRef has expired.
   std::shared_ptr<SymmetricMemory> FindSymmetricMemory(
-      const GpuCliqueKey& clique, stream_executor::DeviceAddressBase addr);
+      int device_ordinal, const GpuCliqueKey& clique,
+      stream_executor::DeviceAddressBase addr);
 
   std::shared_ptr<stream_executor::gpu::MulticastMemory> FindMulticastMemory(
-      const GpuCliqueKey& clique, stream_executor::DeviceAddressBase addr);
+      int device_ordinal, const GpuCliqueKey& clique,
+      stream_executor::DeviceAddressBase addr);
 
   // Inserts or replaces a cached entry. Returns a locked shared_ptr.
   std::shared_ptr<SymmetricMemory> AddSymmetricMemory(
-      const GpuCliqueKey& clique, stream_executor::DeviceAddressBase addr,
+      int device_ordinal, const GpuCliqueKey& clique,
+      stream_executor::DeviceAddressBase addr,
       tsl::TiedRef<SymmetricMemory> sym_memory);
 
   std::shared_ptr<stream_executor::gpu::MulticastMemory> AddMulticastMemory(
-      const GpuCliqueKey& clique, stream_executor::DeviceAddressBase addr,
+      int device_ordinal, const GpuCliqueKey& clique,
+      stream_executor::DeviceAddressBase addr,
       tsl::TiedRef<stream_executor::gpu::MulticastMemory> multicast_memory);
 
-
-
  private:
   using Key = std::pair<GpuCliqueKey, stream_executor::DeviceAddressBase>;
 
-  absl::Mutex mutex_;
+  struct DeviceCache {
+    absl::Mutex mutex;
 
-  absl::flat_hash_map<Key, tsl::TiedRef<SymmetricMemory>> sym_memories_
-      ABSL_GUARDED_BY(mutex_);
+    absl::flat_hash_map<Key, tsl::TiedRef<SymmetricMemory>> sym_memories
+        ABSL_GUARDED_BY(mutex);
 
-  absl::flat_hash_map<Key, tsl::TiedRef<stream_executor::gpu::MulticastMemory>>
-      multicast_memories_ ABSL_GUARDED_BY(mutex_);
+    absl::flat_hash_map<Key,
+                        tsl::TiedRef<stream_executor::gpu::MulticastMemory>>
+        multicast_memories ABSL_GUARDED_BY(mutex);
+  };
+
+  PerDeviceState<DeviceCache> device_caches_;
 };
 
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/runtime/collective_memory_cache_test.cc b/third_party/xla/xla/backends/gpu/runtime/collective_memory_cache_test.cc
new file mode 100644
index 00000000000..35481f0f615
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/runtime/collective_memory_cache_test.cc
@@ -0,0 +1,170 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/gpu/runtime/collective_memory_cache.h"
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include <gtest/gtest.h>
+#include "absl/synchronization/notification.h"
+#include "xla/backends/gpu/collectives/gpu_clique_key.h"
+#include "xla/core/collectives/symmetric_memory.h"
+#include "xla/runtime/device_id.h"
+#include "xla/stream_executor/device_address.h"
+#include "xla/stream_executor/gpu/multicast_memory.h"
+#include "xla/tsl/platform/env.h"
+#include "xla/tsl/platform/threadpool.h"
+#include "xla/tsl/util/tied_ref.h"
+
+namespace xla::gpu {
+namespace {
+
+class FakeSymmetricMemory : public SymmetricMemory {
+ public:
+  explicit FakeSymmetricMemory(stream_executor::DeviceAddressBase addr)
+      : addr_(addr) {}
+
+  stream_executor::DeviceAddressBase addr() const override { return addr_; }
+  std::string ToString() const override { return "FakeSymmetricMemory"; }
+  PackedKernelArg PackKernelArg() const override { return addr_.opaque(); }
+
+ private:
+  stream_executor::DeviceAddressBase addr_;
+};
+
+class FakeMulticastMemory : public stream_executor::gpu::MulticastMemory {};
+
+GpuCliqueKey MakeCliqueKey() {
+  return GpuCliqueKey({GlobalDeviceId(0), GlobalDeviceId(1)},
+                      /*num_local_participants=*/2);
+}
+
+stream_executor::DeviceAddressBase MakeAddress(uintptr_t ptr, uint64_t size) {
+  return stream_executor::DeviceAddressBase(reinterpret_cast<void*>(ptr), size);
+}
+
+TEST(CollectiveMemoryCacheTest, FindReturnsNullptrWhenNotCached) {
+  CollectiveMemoryCache cache(/*num_devices=*/2);
+  GpuCliqueKey clique = MakeCliqueKey();
+  stream_executor::DeviceAddressBase addr = MakeAddress(0x1000, 64);
+
+  EXPECT_EQ(cache.FindSymmetricMemory(0, clique, addr), nullptr);
+  EXPECT_EQ(cache.FindMulticastMemory(0, clique, addr), nullptr);
+  EXPECT_EQ(cache.FindSymmetricMemory(4, clique, addr), nullptr);
+  EXPECT_EQ(cache.FindMulticastMemory(4, clique, addr), nullptr);
+}
+
+TEST(CollectiveMemoryCacheTest, IsolatesEntriesByDeviceOrdinal) {
+  CollectiveMemoryCache cache(/*num_devices=*/2);
+  GpuCliqueKey clique = MakeCliqueKey();
+  stream_executor::DeviceAddressBase addr = MakeAddress(0x1000, 64);
+
+  tsl::TiedAny owner;
+  auto sym0 = std::make_unique<FakeSymmetricMemory>(addr);
+  auto* sym0_ptr = sym0.get();
+  auto sym1 = std::make_unique<FakeSymmetricMemory>(addr);
+  auto* sym1_ptr = sym1.get();
+  auto sym5 = std::make_unique<FakeSymmetricMemory>(addr);
+  auto* sym5_ptr = sym5.get();
+
+  EXPECT_EQ(cache
+                .AddSymmetricMemory(0, clique, addr,
+                                    owner.Tie<SymmetricMemory>(std::move(sym0)))
+                .get(),
+            sym0_ptr);
+  EXPECT_EQ(cache
+                .AddSymmetricMemory(1, clique, addr,
+                                    owner.Tie<SymmetricMemory>(std::move(sym1)))
+                .get(),
+            sym1_ptr);
+  EXPECT_EQ(cache
+                .AddSymmetricMemory(5, clique, addr,
+                                    owner.Tie<SymmetricMemory>(std::move(sym5)))
+                .get(),
+            sym5_ptr);
+
+  EXPECT_EQ(cache.FindSymmetricMemory(0, clique, addr).get(), sym0_ptr);
+  EXPECT_EQ(cache.FindSymmetricMemory(1, clique, addr).get(), sym1_ptr);
+  EXPECT_EQ(cache.FindSymmetricMemory(5, clique, addr).get(), sym5_ptr);
+  EXPECT_EQ(cache.FindSymmetricMemory(2, clique, addr), nullptr);
+}
+
+TEST(CollectiveMemoryCacheTest, ExpiredTiedRefReturnsNullptr) {
+  CollectiveMemoryCache cache(/*num_devices=*/2);
+  GpuCliqueKey clique = MakeCliqueKey();
+  stream_executor::DeviceAddressBase addr = MakeAddress(0x2000, 64);
+
+  {
+    tsl::TiedAny owner;
+    cache.AddSymmetricMemory(0, clique, addr,
+                             owner.Tie<SymmetricMemory>(
+                                 std::make_unique<FakeSymmetricMemory>(addr)));
+    cache.AddMulticastMemory(0, clique, addr,
+                             owner.Tie<stream_executor::gpu::MulticastMemory>(
+                                 std::make_unique<FakeMulticastMemory>()));
+    EXPECT_NE(cache.FindSymmetricMemory(0, clique, addr), nullptr);
+    EXPECT_NE(cache.FindMulticastMemory(0, clique, addr), nullptr);
+  }
+
+  EXPECT_EQ(cache.FindSymmetricMemory(0, clique, addr), nullptr);
+  EXPECT_EQ(cache.FindMulticastMemory(0, clique, addr), nullptr);
+}
+
+TEST(CollectiveMemoryCacheTest, ConcurrentAddAndFindAcrossDevices) {
+  constexpr int kNumDevices = 4;
+  constexpr int kNumOrdinals = 8;
+  constexpr int kThreadsPerOrdinal = 4;
+  constexpr int kNumThreads = kNumOrdinals * kThreadsPerOrdinal;
+
+  CollectiveMemoryCache cache(kNumDevices);
+  GpuCliqueKey clique = MakeCliqueKey();
+  tsl::TiedAny owner;
+  absl::Notification start;
+
+  {
+    tsl::thread::ThreadPool pool(tsl::Env::Default(), "cache_test",
+                                 kNumThreads);
+    for (int t = 0; t < kNumThreads; ++t) {
+      pool.Schedule([&, t] {
+        int ordinal = t % kNumOrdinals;
+        stream_executor::DeviceAddressBase addr =
+            MakeAddress(0x1000 * (t + 1), 64);
+        auto sym = owner.Tie<SymmetricMemory>(
+            std::make_unique<FakeSymmetricMemory>(addr));
+        auto mcast = owner.Tie<stream_executor::gpu::MulticastMemory>(
+            std::make_unique<FakeMulticastMemory>());
+
+        start.WaitForNotification();
+
+        auto added_sym =
+            cache.AddSymmetricMemory(ordinal, clique, addr, std::move(sym));
+        auto added_mcast =
+            cache.AddMulticastMemory(ordinal, clique, addr, std::move(mcast));
+        EXPECT_NE(added_sym, nullptr);
+        EXPECT_NE(added_mcast, nullptr);
+        EXPECT_EQ(cache.FindSymmetricMemory(ordinal, clique, addr), added_sym);
+        EXPECT_EQ(cache.FindMulticastMemory(ordinal, clique, addr),
+                  added_mcast);
+      });
+    }
+    start.Notify();
+  }
+}
+
+}  // namespace
+}  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.cc b/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.cc
index a478e23f367..753f740223a 100644
--- a/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/multi_gpu_barrier.cc
@@ -26,11 +26,10 @@ limitations under the License.
 
 #include "absl/base/no_destructor.h"
 #include "absl/base/nullability.h"
-#include "absl/container/flat_hash_map.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
-#include "absl/synchronization/mutex.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/core/collectives/rank_id.h"
 #include "xla/core/collectives/symmetric_memory.h"
 #include "xla/ffi/api/record_api.h"
@@ -54,25 +53,24 @@ namespace {
 template <typename Kernel>
 absl::StatusOr<typename Kernel::KernelType*> GetCachedKernel(
     se::StreamExecutor* executor) {
-  static absl::NoDestructor<absl::flat_hash_map<
-      se::StreamExecutor*, std::unique_ptr<typename Kernel::KernelType>>>
-      kernel_per_executor;
-  static absl::NoDestructor<absl::Mutex> kernel_mutex;
-
-  absl::MutexLock lock(*kernel_mutex);
-  auto it = kernel_per_executor->find(executor);
-  if (it == kernel_per_executor->end()) {
-    ABSL_ASSIGN_OR_RETURN(
-        auto new_kernel,
-        (stream_executor::gpu::GpuKernelRegistry::GetGlobalRegistry()
-             .LoadKernel<Kernel>(executor)));
-    it = kernel_per_executor
-             ->emplace(executor, std::make_unique<typename Kernel::KernelType>(
-                                     std::move(new_kernel)))
-             .first;
-  }
-
-  return it->second.get();
+  // Pre-size to 16 slots so all standard single-host GPU topologies
+  // (up to 16 GPUs, e.g. DGX H100/B200) hit the lock-free array fast path.
+  static absl::NoDestructor<
+      PerDeviceState<std::unique_ptr<typename Kernel::KernelType>>>
+      kernel_per_device(/*num_devices=*/16);
+  int device_ordinal = executor->device_ordinal();
+  ABSL_RETURN_IF_ERROR(kernel_per_device->GetOrCreateAndInitialize(
+      device_ordinal,
+      [&](std::unique_ptr<typename Kernel::KernelType>* slot) -> absl::Status {
+        ABSL_ASSIGN_OR_RETURN(
+            auto new_kernel,
+            stream_executor::gpu::GpuKernelRegistry::GetGlobalRegistry()
+                .LoadKernel<Kernel>(executor));
+        *slot = std::make_unique<typename Kernel::KernelType>(
+            std::move(new_kernel));
+        return absl::OkStatus();
+      }));
+  return kernel_per_device->Find(device_ordinal)->get();
 }
 
 }  // namespace
diff --git a/third_party/xla/xla/service/gpu/gpu_executable.cc b/third_party/xla/xla/service/gpu/gpu_executable.cc
index fc7682893ca..7182d9ff96b 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable.cc
@@ -518,6 +518,8 @@ GpuExecutable::GpuExecutable(
       thunk_sequence_proto_(std::move(thunk_sequence_proto)),
       executable_abi_version_(std::move(executable_abi_version)),
       cpu_target_machine_options_(std::move(cpu_target_machine_options)),
+      collective_memory_cache_(
+          gpu_topology.has_value() ? gpu_topology->num_devices_per_host() : 0),
       buffer_allocations_debug_summary_(
           std::move(buffer_allocations_debug_summary)),
       collective_use_minimal_resource_(collective_use_minimal_resource),

From 39f393f362b074b531193fce3eaf4d10bac2d61c Mon Sep 17 00:00:00 2001
From: Oleg Shyshkov <shyshkov@google.com>
Date: Wed, 14 Oct 2026 09:22:49 -0700
Subject: [PATCH 0794/1000] [XLA:GPU] Move FlattenCollectiveFusion to
 collective_fusion.

PiperOrigin-RevId: 999419508
---
 .../gpu/codegen/triton/collective_emitter.cc  | 39 ---------------
 .../gpu/codegen/triton/collective_emitter.h   | 10 ----
 .../codegen/triton/collective_emitter_test.cc | 37 --------------
 .../backends/gpu/transforms/collectives/BUILD |  3 ++
 .../collectives/collective_fusion.cc          | 48 +++++++++++++++++++
 .../collectives/collective_fusion_test.cc     | 29 +++++++++++
 6 files changed, 80 insertions(+), 86 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
index 0d911ef8d2e..ad27ee30ae5 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
@@ -1240,45 +1240,6 @@ class ReduceScatterEmitter : public ReductionEmitter {
 
 }  // namespace
 
-absl::Status FlattenCollectiveFusion(
-    HloFusionInstruction* absl_nonnull fusion_instr) {
-  HloComputation* entry_computation = fusion_instr->parent();
-  HloComputation* fused_computation =
-      fusion_instr->fused_instructions_computation();
-  Shape original_shape = fusion_instr->shape();
-  const int64_t total_elements = ShapeUtil::ElementsIn(original_shape);
-  if (fused_computation->parameter_instructions().size() != 1) {
-    return absl::InternalError(
-        "Flattening to 1D is not supported for variadic fused computations "
-        "with more than one parameter instruction");
-  }
-  DCHECK_EQ(fusion_instr->operand_count(), 1);
-  // Flatten all instructions in the fused computation to 1D.
-  for (HloInstruction* instr : fused_computation->instructions()) {
-    if (ShapeUtil::SameDimensions(instr->shape(), original_shape)) {
-      *instr->mutable_shape() = ShapeUtil::MakeShapeWithDenseLayout(
-          instr->shape().element_type(), {total_elements}, {0});
-    }
-  }
-  *fusion_instr->mutable_shape() = ShapeUtil::MakeShapeWithDenseLayout(
-      original_shape.element_type(), {total_elements}, {0});
-  HloInstruction* fusion_operand = fusion_instr->mutable_operand(0);
-  Shape flat_input_shape = ShapeUtil::MakeShapeWithDenseLayout(
-      fusion_operand->shape().element_type(), {total_elements}, {0});
-  HloInstruction* bitcast_to_1d = entry_computation->AddInstruction(
-      HloInstruction::CreateBitcast(flat_input_shape, fusion_operand));
-  ABSL_RETURN_IF_ERROR(
-      fusion_instr->ReplaceOperandWithDifferentShape(0, bitcast_to_1d));
-  HloInstruction* bitcast_to_original_shape = entry_computation->AddInstruction(
-      HloInstruction::CreateBitcast(original_shape, fusion_instr));
-  if (VLOG_IS_ON(3)) {
-    VLOG(3) << "Flattening original shape : " << original_shape.ToString()
-            << " to 1D shape: " << flat_input_shape.ToString();
-  }
-  return fusion_instr->ReplaceAllUsesWithDifferentShape(
-      bitcast_to_original_shape);
-}
-
 llvm::SmallVector<int64_t> GreedyPowerOfTwoTiles(const Shape& output_shape,
                                                  int32_t num_blocks) {
   CHECK_GT(num_blocks, 0) << "num_blocks must be positive. Was " << num_blocks;
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h
index 8e14ed631b3..de79607c935 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h
@@ -33,7 +33,6 @@ limitations under the License.
 #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_opcode.h"
 #include "xla/service/gpu/backend_configs.pb.h"
 #include "xla/service/gpu/launch_dimensions.h"
 #include "xla/shape.h"
@@ -46,15 +45,6 @@ class GpuTopology;
 
 namespace xla::gpu {
 
-// Flattens instructions inside a > 1D collective fusion to 1D instructions
-// inplace. For example, a 2D computation with shape [2, 3] will be flattened to
-// a 1D computation with shape [6] with surrounding bitcasts inserted in the
-// parent computation.
-// NB: This should only be used if the fused computation  only contains
-// elementwise operations (and the collective itself).
-absl::Status FlattenCollectiveFusion(
-    HloFusionInstruction* absl_nonnull fusion_instr);
-
 // Returns a tile such that each dimension is a power of two and that the total
 // number of blocks does not exceed num_blocks. Returns the tile sizes in
 // same order as the output shape. It picks tiles greedily from the most major
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter_test.cc
index 4cadf56733e..fef68f8dbec 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter_test.cc
@@ -484,43 +484,6 @@ INSTANTIATE_TEST_SUITE_P(
     [](const ::testing::TestParamInfo<GreedyPowerOfTwoTilesTest::ParamType>&
            info) { return info.param.test_name; });
 
-TEST_F(CollectiveEmitterTest, FlattenCollectiveComputation) {
-  Shape shape_3d = ShapeUtil::MakeShape(F32, {2, 4, 8});
-  Shape shape_1d = ShapeUtil::MakeShape(F32, {64});  // 2 * 4 * 8
-  ASSERT_OK_AND_ASSIGN(ModuleWithFusion module_with_fusion,
-                       BuildModuleWithFusion(GetModuleStr(shape_3d)));
-  HloFusionInstruction* fusion_instr = module_with_fusion.MutableFusionInstr();
-  EXPECT_THAT(fusion_instr,
-              AllOf(HasOpcode(HloOpcode::kFusion), HasShape(shape_3d)));
-
-  HloComputation* entry_computation =
-      module_with_fusion.module->entry_computation();
-  EXPECT_OK(FlattenCollectiveFusion(fusion_instr));
-
-  // Output bitcast should have the original 3D shape.
-  HloInstruction* root = entry_computation->root_instruction();
-  EXPECT_THAT(root, AllOf(HasOpcode(HloOpcode::kBitcast), HasShape(shape_3d),
-                          HasNumOperands(1)));
-
-  // All other ops should have the flattened 1D shape.
-  HloInstruction* new_fusion_instr = root->mutable_operand(0);
-  EXPECT_THAT(new_fusion_instr, AllOf(HasOpcode(HloOpcode::kFusion),
-                                      HasShape(shape_1d), HasNumOperands(1)));
-
-  HloInstruction* input_bitcast = new_fusion_instr->mutable_operand(0);
-  EXPECT_THAT(input_bitcast, AllOf(HasOpcode(HloOpcode::kBitcast),
-                                   HasShape(shape_1d), HasNumOperands(1)));
-  EXPECT_THAT(input_bitcast->operand(0), HasOpcode(HloOpcode::kParameter));
-
-  HloComputation* fused_comp =
-      new_fusion_instr->fused_instructions_computation();
-  EXPECT_THAT(
-      fused_comp->parameter_instructions(),
-      ElementsAre(AllOf(HasOpcode(HloOpcode::kParameter), HasShape(shape_1d))));
-  EXPECT_THAT(fused_comp->root_instruction(),
-              AllOf(HasOpcode(HloOpcode::kAllReduce), HasShape(shape_1d)));
-}
-
 TEST_F(CollectiveBlockLevelConfigTest, AllGatherBlockLevelConfig) {
   constexpr absl::string_view kAllGatherHloStr = R"(
     HloModule test
diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/BUILD b/third_party/xla/xla/backends/gpu/transforms/collectives/BUILD
index cf7a911f638..704ab31abd7 100644
--- a/third_party/xla/xla/backends/gpu/transforms/collectives/BUILD
+++ b/third_party/xla/xla/backends/gpu/transforms/collectives/BUILD
@@ -975,8 +975,11 @@ cc_library(
         "//xla/service:gpu_topology",
         "//xla/service/gpu:backend_configs_cc",
         "//xla/stream_executor/gpu:all_reduce_kernel",
+        "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/log:vlog_is_on",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion.cc b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion.cc
index 57952e3b3db..31577cf9150 100644
--- a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion.cc
@@ -18,8 +18,11 @@ limitations under the License.
 #include <cstdint>
 #include <vector>
 
+#include "absl/base/nullability.h"
 #include "absl/container/flat_hash_set.h"
+#include "absl/log/check.h"
 #include "absl/log/log.h"
+#include "absl/log/vlog_is_on.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
@@ -126,6 +129,51 @@ bool ShouldFlatten(const HloInstruction* instr) {
              se::gpu::AllReduceStrategy::kTwoShot;
 }
 
+// Flattens instructions inside a > 1D collective fusion to 1D instructions
+// inplace. For example, a 2D computation with shape [2, 3] will be flattened to
+// a 1D computation with shape [6] with surrounding bitcasts inserted in the
+// parent computation.
+// NB: This should only be used if the fused computation  only contains
+// elementwise operations (and the collective itself).
+absl::Status FlattenCollectiveFusion(
+    HloFusionInstruction* absl_nonnull fusion_instr) {
+  HloComputation* entry_computation = fusion_instr->parent();
+  HloComputation* fused_computation =
+      fusion_instr->fused_instructions_computation();
+  Shape original_shape = fusion_instr->shape();
+  const int64_t total_elements = ShapeUtil::ElementsIn(original_shape);
+  if (fused_computation->parameter_instructions().size() != 1) {
+    return absl::InternalError(
+        "Flattening to 1D is not supported for variadic fused computations "
+        "with more than one parameter instruction");
+  }
+  DCHECK_EQ(fusion_instr->operand_count(), 1);
+  // Flatten all instructions in the fused computation to 1D.
+  for (HloInstruction* instr : fused_computation->instructions()) {
+    if (ShapeUtil::SameDimensions(instr->shape(), original_shape)) {
+      *instr->mutable_shape() = ShapeUtil::MakeShapeWithDenseLayout(
+          instr->shape().element_type(), {total_elements}, {0});
+    }
+  }
+  *fusion_instr->mutable_shape() = ShapeUtil::MakeShapeWithDenseLayout(
+      original_shape.element_type(), {total_elements}, {0});
+  HloInstruction* fusion_operand = fusion_instr->mutable_operand(0);
+  Shape flat_input_shape = ShapeUtil::MakeShapeWithDenseLayout(
+      fusion_operand->shape().element_type(), {total_elements}, {0});
+  HloInstruction* bitcast_to_1d = entry_computation->AddInstruction(
+      HloInstruction::CreateBitcast(flat_input_shape, fusion_operand));
+  ABSL_RETURN_IF_ERROR(
+      fusion_instr->ReplaceOperandWithDifferentShape(0, bitcast_to_1d));
+  HloInstruction* bitcast_to_original_shape = entry_computation->AddInstruction(
+      HloInstruction::CreateBitcast(original_shape, fusion_instr));
+  if (VLOG_IS_ON(3)) {
+    VLOG(3) << "Flattening original shape : " << original_shape.ToString()
+            << " to 1D shape: " << flat_input_shape.ToString();
+  }
+  return fusion_instr->ReplaceAllUsesWithDifferentShape(
+      bitcast_to_original_shape);
+}
+
 // Normalizes legacy all-reduce-start/done into modern generic async boundaries.
 // Normally this does not exist but it can be created by explicitly calling
 // psum_start/end in jax.
diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion_test.cc b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion_test.cc
index 2c19a3459d4..8b0f22c873d 100644
--- a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion_test.cc
@@ -167,6 +167,35 @@ TEST_F(CollectiveFusionTest, FusesAndFlattensMediumAllReduceTwoShot) {
                             kExpected, /*after_pass_checks=*/nullptr, &config);
 }
 
+TEST_F(CollectiveFusionTest, FusesAndFlattens3DAllReduceTwoShot) {
+  // Use a 3D shape to test multi-dimensional flattening:
+  // [2, 4, 32768] = 262144 elements (1 MB -> TwoShot)
+  Shape shape = xla::ShapeUtil::MakeShape(F32, {2, 4, 32768});
+  std::string hlo = absl::StrFormat(kAllReduceHloTemplate, shape.ToString());
+  SCOPED_TRACE(::testing::Message() << "hlo: " << hlo);
+  static constexpr absl::string_view kExpected = R"(
+    // CHECK: %[[FUSION_COMPUTATION:.*]] (param_0: f32[262144]) -> f32[262144] {
+    // CHECK:   %[[P0:.*]] = f32[262144]{0} parameter(0)
+    // CHECK:   ROOT {{.*}} = f32[262144]{0} all-reduce(%[[P0]]),
+    // CHECK-SAME: to_apply=%add
+    // CHECK: }
+    //
+    // CHECK: ENTRY %e (p0.1: f32[2,4,32768]) -> f32[2,4,32768] {
+    // CHECK:   %[[P0:.*]] = f32[2,4,32768]{2,1,0} parameter(0)
+    // CHECK:   %[[BITCAST_TO_1D:.*]] = f32[262144]{0} bitcast(%[[P0]])
+    // CHECK:   %[[FUSION:.*]] = f32[262144]{0} fusion(%[[BITCAST_TO_1D]]),
+    // CHECK-SAME: kind=kCustom,
+    // CHECK-SAME: calls=%[[FUSION_COMPUTATION]],
+    // CHECK-SAME: backend_config={
+    // CHECK-SAME: "kind":"__triton_collective"
+    // CHECK-SAME: "block_level_fusion_config"
+    // CHECK:   ROOT {{.*}} = f32[2,4,32768]{2,1,0} bitcast(%[[FUSION]])
+  )";
+  HloModuleConfig config = GetModuleConfigForTest(/*replica_count=*/8);
+  RunAndFilecheckHloRewrite(hlo, AnnotateAndFusePipeline(*gpu_topology_),
+                            kExpected, /*after_pass_checks=*/nullptr, &config);
+}
+
 TEST_F(CollectiveFusionTest, DoesNotFuseUnannotated) {
   constexpr int64_t kNumElements = 32768;
   Shape shape = ShapeUtil::MakeShape(F32, {kNumElements});

From f2072e98f67671afc008d0be623ce2def4a1c200 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 22:27:15 +0530
Subject: [PATCH 0795/1000] Force GitHub PR sync


From fc058b6345f5d1b61c6da47a3bf06134df69f5f6 Mon Sep 17 00:00:00 2001
From: Vrishab Madduri <vrishab@google.com>
Date: Wed, 14 Oct 2026 09:31:51 -0700
Subject: [PATCH 0796/1000] Avoid driver-side waits on compute work in CC
 chunked staged transfers.

Under GPU Confidential Computing, a chunked staged host-to-device transfer made the shared host-to-device stream wait on the destination buffer's allocation event, which is recorded on the compute stream. Staged copies can then wait inside the CUDA driver for that compute work while holding driver locks, so under load concurrent transfers, event creation and CUDA graph capture could deadlock, and transfers into ready buffers queued behind compute.

When chunked staging is enabled (CC mode, or forced in tests), staged H2D and D2H transfers are now scheduled once the destination or source allocation's BufferSequencingEvent is ready (an ExecuteWhenReady dependency), so neither the transfer streams nor a host thread wait on compute work. Chunked H2D copies to a device run on a per-device single-threaded AsyncWorkRunner (LocalDeviceState::staged_host_to_device_runner(), created lazily), so one transfer at a time issues staged copies to the device; a transfer whose allocation is still pending does not hold up later transfers. D2H transfers get the same allocation dependency without serialization. Staging slot events come from the device's EventPool. Behavior with CC disabled is unchanged.
Per-device serialization of staged H2D copies is required under CC: without it, a B200 CC disaggregated serving workload livelocked in the CUDA driver on a cold start at 8-12 concurrent requests, with no throughput gain.

Testing: new se_gpu_pjrt_client_test cases with chunked staging forced: an H2D transfer into a ready buffer completes while the compute stream is stalled, and concurrent H2D transfers into buffers reusing memory of inputs read by queued executions complete with correct data. Existing pjrt/se and se_gpu_pjrt_client tests pass. Also validated end to end under GPU CC on B200 with a disaggregated serving workload that deadlocked without this change.
PiperOrigin-RevId: 999424667
---
 third_party/xla/xla/pjrt/gpu/BUILD            |   3 +
 .../xla/pjrt/gpu/se_gpu_pjrt_client_test.cc   | 158 ++++++++++++++++++
 third_party/xla/xla/pjrt/se/BUILD             |   1 +
 .../xla/xla/pjrt/se/local_device_state.cc     |  16 ++
 .../xla/xla/pjrt/se/local_device_state.h      |  10 ++
 third_party/xla/xla/pjrt/se/se_raw_buffer.cc  | 109 ++++++++++--
 6 files changed, 281 insertions(+), 16 deletions(-)

diff --git a/third_party/xla/xla/pjrt/gpu/BUILD b/third_party/xla/xla/pjrt/gpu/BUILD
index 44ffae3ae0e..4a2935e79be 100644
--- a/third_party/xla/xla/pjrt/gpu/BUILD
+++ b/third_party/xla/xla/pjrt/gpu/BUILD
@@ -451,7 +451,10 @@ xla_test(
         "//xla/service/gpu:gpu_executable_run_options",
         "//xla/service/gpu:gpu_memory_space_assignment",
         "//xla/stream_executor:device_address_allocator",
+        "//xla/stream_executor:event",
+        "//xla/stream_executor:stream",
         "//xla/stream_executor:stream_executor_address_allocator",
+        "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor/integrations:tf_allocator_adapter",
         "//xla/tests:literal_test_util",
         "//xla/tsl/concurrency:async_value",
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
index 50815923322..557fb946447 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
@@ -110,7 +110,10 @@ limitations under the License.
 #endif  // GOOGLE_CUDA
 #include "xla/pjrt/gpu/se_gpu_pjrt_client_test_helper.h"
 #include "xla/stream_executor/device_address_allocator.h"
+#include "xla/stream_executor/event.h"
 #include "xla/stream_executor/integrations/tf_allocator_adapter.h"
+#include "xla/stream_executor/stream.h"
+#include "xla/stream_executor/stream_executor.h"
 #include "xla/stream_executor/stream_executor_address_allocator.h"
 #include "xla/tests/literal_test_util.h"
 #include "xla/tsl/concurrency/async_value_ref.h"
@@ -1790,6 +1793,161 @@ TEST(StreamExecutorGpuClientTest, ChunkedStagingWithNonzeroOffset) {
                                 10 * kTestStagingChunkSize + 45);
 }
 
+// Forces chunked staging (the transfer path used under GPU Confidential
+// Computing) on `client`.
+PjRtStreamExecutorRawClient* ForceChunkedStaging(CommonPjRtClient* client) {
+  auto* raw_client =
+      absl::down_cast<PjRtStreamExecutorRawClient*>(client->raw_client());
+  raw_client->SetStagingChunkingForTesting(/*force=*/true,
+                                           kTestStagingChunkSize);
+  return raw_client;
+}
+
+void ExpectDeviceBytes(PjRtRawBufferInterface& raw_buffer,
+                       const std::vector<uint8_t>& expected) {
+  std::vector<uint8_t> actual(expected.size(), 0);
+  TF_ASSERT_OK(
+      raw_buffer.CopyRawDeviceToHost(actual.data(), 0, actual.size()).Await());
+  EXPECT_EQ(actual, expected);
+}
+
+// A chunked staged H2D transfer into a buffer whose allocation is not yet
+// ready on the compute stream waits on the host, not on the shared
+// host-to-device stream, so transfers into ready buffers do not queue behind
+// compute work.
+TEST(StreamExecutorGpuClientTest, ChunkedStagingH2DDoesNotQueueBehindCompute) {
+  GpuClientOptions options = GetTestGpuClientOptions();
+  options.should_stage_host_to_device_transfers = true;
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtClient> pjrt_client,
+                          GetStreamExecutorGpuClient(options));
+  auto* client = absl::down_cast<CommonPjRtClient*>(pjrt_client.get());
+  PjRtStreamExecutorRawClient* raw_client = ForceChunkedStaging(client);
+  PjRtDevice* device = client->addressable_devices()[0];
+  LocalDeviceState* local_device_state =
+      raw_client->device_state(device->local_device_id());
+  TF_ASSERT_OK_AND_ASSIGN(PjRtMemorySpace * memory_space,
+                          device->default_memory_space());
+  const int64_t size = 3 * kTestStagingChunkSize;
+  auto allocate = [&]() {
+    return client->AllocateRawBuffer(memory_space, size, /*retry_on_oom=*/true,
+                                     /*allocate_after=*/{});
+  };
+  const std::vector<uint8_t> src = MakeStagingTestPattern(size);
+  const std::vector<uint8_t> src2(src.rbegin(), src.rend());
+
+  // `ready`'s allocation is complete before the compute stream stalls.
+  TF_ASSERT_OK_AND_ASSIGN(PjRtRawBufferRef ready, allocate());
+  TF_ASSERT_OK(ready->CopyRawHostToDevice(src.data(), 0, size).Await());
+
+  auto release_compute = std::make_shared<absl::Notification>();
+  absl::Cleanup release = [&] {
+    if (!release_compute->HasBeenNotified()) {
+      release_compute->Notify();
+    }
+  };
+  TF_ASSERT_OK(local_device_state->compute_stream()->DoHostCallback(
+      [release_compute] { release_compute->WaitForNotification(); }));
+  // `pending`'s allocation event is recorded behind the stalled compute work.
+  TF_ASSERT_OK_AND_ASSIGN(PjRtRawBufferRef pending, allocate());
+  Future<> pending_h2d = pending->CopyRawHostToDevice(src.data(), 0, size);
+  Future<> ready_h2d = ready->CopyRawHostToDevice(src2.data(), 0, size);
+  // The transfers run asynchronously and expose no signal for "copies issued",
+  // so give both time to reach the host-to-device stream. Without this, the
+  // event below could be recorded before either transfer enqueues work on that
+  // stream and pass trivially. Their relative order does not matter.
+  absl::SleepFor(absl::Milliseconds(200));
+  // Host functions in a context run serially on one driver thread, so
+  // `ready_h2d`'s completion callback cannot run during the stall. Instead,
+  // poll an event recorded on the host-to-device stream.
+  se::Stream* h2d_stream = local_device_state->host_to_device_stream();
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<se::Event> h2d_done,
+                          h2d_stream->parent()->CreateEvent());
+  TF_ASSERT_OK(h2d_stream->RecordEvent(h2d_done.get()));
+  const absl::Time deadline = absl::Now() + absl::Seconds(10);
+  while (h2d_done->PollForStatus() == se::Event::Status::kPending &&
+         absl::Now() < deadline) {
+    absl::SleepFor(absl::Milliseconds(1));
+  }
+  EXPECT_EQ(h2d_done->PollForStatus(), se::Event::Status::kComplete);
+  EXPECT_FALSE(pending_h2d.IsReady());
+
+  std::move(release).Invoke();
+  TF_ASSERT_OK(pending_h2d.Await());
+  TF_ASSERT_OK(ready_h2d.Await());
+  ExpectDeviceBytes(*pending, src);
+  ExpectDeviceBytes(*ready, src2);
+}
+
+// Many concurrent chunked staged H2D transfers into buffers that may reuse the
+// memory of inputs still being read by queued executions complete without
+// deadlock, and neither the transfers nor the executions see corrupted data.
+TEST(StreamExecutorGpuClientTest, ChunkedStagingConcurrentH2DAfterExecute) {
+  GpuClientOptions options = GetTestGpuClientOptions();
+  options.should_stage_host_to_device_transfers = true;
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtClient> pjrt_client,
+                          GetStreamExecutorGpuClient(options));
+  auto* client = absl::down_cast<CommonPjRtClient*>(pjrt_client.get());
+  ForceChunkedStaging(client);
+  PjRtDevice* device = client->addressable_devices()[0];
+  TF_ASSERT_OK_AND_ASSIGN(PjRtMemorySpace * memory_space,
+                          device->default_memory_space());
+  static constexpr char const* kAddOneProgram = R"(
+HloModule add_one, entry_computation_layout={(f32[4096])->f32[4096]}
+
+ENTRY main {
+  x = f32[4096] parameter(0)
+  one = f32[] constant(1)
+  ones = f32[4096] broadcast(one), dimensions={}
+  ROOT result = f32[4096] add(x, ones)
+}
+)";
+  TF_ASSERT_OK_AND_ASSIGN(auto executable,
+                          CompileExecutable(kAddOneProgram, *client));
+  Shape shape = ShapeUtil::MakeShape(F32, {4096});
+  const int64_t size = ShapeUtil::ByteSizeOf(shape);
+  constexpr int kIterations = 4;
+  constexpr int kTransfers = 8;
+  for (int iter = 0; iter < kIterations; ++iter) {
+    std::vector<float> data(4096, static_cast<float>(iter));
+    TF_ASSERT_OK_AND_ASSIGN(
+        auto input,
+        client->BufferFromHostBuffer(
+            data.data(), shape.element_type(), shape.dimensions(),
+            /*byte_strides=*/std::nullopt,
+            PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall,
+            /*on_done_with_host_buffer=*/nullptr, memory_space,
+            /*device_layout=*/nullptr));
+    TF_ASSERT_OK_AND_ASSIGN(
+        auto result, executable->Execute({{input.get()}}, /*options=*/{}));
+    // Free the input while the execution may still be reading it, so the
+    // buffers below may reuse its memory.
+    input.reset();
+
+    std::vector<std::vector<uint8_t>> srcs;
+    std::vector<PjRtRawBufferRef> buffers;
+    std::vector<Future<>> transfers;
+    for (int i = 0; i < kTransfers; ++i) {
+      srcs.push_back(MakeStagingTestPattern(size));
+      srcs.back()[0] = static_cast<uint8_t>(iter * kTransfers + i);
+      TF_ASSERT_OK_AND_ASSIGN(
+          PjRtRawBufferRef buffer,
+          client->AllocateRawBuffer(memory_space, size, /*retry_on_oom=*/true,
+                                    /*allocate_after=*/{}));
+      transfers.push_back(
+          buffer->CopyRawHostToDevice(srcs.back().data(), 0, size));
+      buffers.push_back(std::move(buffer));
+    }
+    for (int i = 0; i < kTransfers; ++i) {
+      TF_ASSERT_OK(transfers[i].Await());
+      ExpectDeviceBytes(*buffers[i], srcs[i]);
+    }
+    TF_ASSERT_OK_AND_ASSIGN(auto literal, result[0][0]->ToLiteral().Await());
+    EXPECT_TRUE(LiteralTestUtil::Equal(
+        LiteralUtil::CreateR1<float>(std::vector<float>(4096, iter + 1)),
+        *literal));
+  }
+}
+
 TEST(StreamExecutorGpuClientTest, BufferFromHostBufferPinnedMemory) {
   TF_ASSERT_OK_AND_ASSIGN(
       auto client, GetStreamExecutorGpuClient(GetTestGpuClientOptions()));
diff --git a/third_party/xla/xla/pjrt/se/BUILD b/third_party/xla/xla/pjrt/se/BUILD
index 9356860796a..1b16f4ae26e 100644
--- a/third_party/xla/xla/pjrt/se/BUILD
+++ b/third_party/xla/xla/pjrt/se/BUILD
@@ -80,6 +80,7 @@ cc_library(
         "//xla/pjrt:device_event",
         "//xla/pjrt:pjrt_common",
         "//xla/pjrt:semaphore",
+        "//xla/pjrt:thread_pool_async_work_runner",
         "//xla/pjrt:worker_thread",
         "//xla/stream_executor:device_address",
         "//xla/stream_executor:event",
diff --git a/third_party/xla/xla/pjrt/se/local_device_state.cc b/third_party/xla/xla/pjrt/se/local_device_state.cc
index 13ab5fee8c0..e0685a7995f 100644
--- a/third_party/xla/xla/pjrt/se/local_device_state.cc
+++ b/third_party/xla/xla/pjrt/se/local_device_state.cc
@@ -23,6 +23,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/base/call_once.h"
 #include "absl/base/casts.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/functional/any_invocable.h"
@@ -36,6 +37,7 @@ limitations under the License.
 #include "xla/client/local_client.h"
 #include "xla/pjrt/async_work_runner.h"
 #include "xla/pjrt/se/buffer_sequencing_event.h"
+#include "xla/pjrt/thread_pool_async_work_runner.h"
 #include "xla/pjrt/worker_thread.h"
 #include "xla/stream_executor/device_address.h"
 #include "xla/stream_executor/platform.h"
@@ -158,6 +160,9 @@ LocalDeviceState::LocalDeviceState(se::StreamExecutor* executor,
 }
 
 LocalDeviceState::~LocalDeviceState() {
+  // Finish staged host-to-device copies before synchronizing the streams they
+  // use.
+  staged_host_to_device_runner_.reset();
   absl::Status status = SynchronizeAllActivity();
   if (!status.ok()) {
     LOG(ERROR) << "Error when closing device: " << status;
@@ -240,6 +245,17 @@ absl::Status LocalDeviceState::SynchronizeAllActivity() {
   return status;
 }
 
+AsyncWorkRunner* LocalDeviceState::staged_host_to_device_runner() {
+  absl::call_once(staged_host_to_device_runner_once_, [this] {
+    tsl::ThreadOptions thread_options;
+    thread_options.numa_node = executor_->numa_node();
+    staged_host_to_device_runner_ =
+        MakeThreadPoolAsyncWorkRunner(tsl::Env::Default(), "py_xla_staged_h2d",
+                                      /*num_threads=*/1, thread_options);
+  });
+  return staged_host_to_device_runner_.get();
+}
+
 absl::Status LocalDeviceState::ThenMemcpyDeviceToDevice(
     se::Stream* transfer_stream, se::Stream* dst_stream,
     se::DeviceAddressBase src_buffer, se::DeviceAddressBase dst_buffer) {
diff --git a/third_party/xla/xla/pjrt/se/local_device_state.h b/third_party/xla/xla/pjrt/se/local_device_state.h
index 1f8abb435e7..407eb593e3b 100644
--- a/third_party/xla/xla/pjrt/se/local_device_state.h
+++ b/third_party/xla/xla/pjrt/se/local_device_state.h
@@ -24,6 +24,7 @@ limitations under the License.
 #include <stack>
 #include <vector>
 
+#include "absl/base/call_once.h"
 #include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/functional/any_invocable.h"
@@ -213,6 +214,12 @@ class LocalDeviceState {
         nullptr, "ThenRelease");
   }
 
+  // Returns a single-threaded runner on which chunked staged host-to-device
+  // transfers (used under GPU Confidential Computing) issue their copies, so
+  // at most one staged transfer at a time issues copies to this device's
+  // `host_to_device_stream()`. Created on first use.
+  AsyncWorkRunner* staged_host_to_device_runner();
+
   std::optional<Semaphore>& compute_semaphore() { return compute_semaphore_; }
 
 
@@ -276,6 +283,9 @@ class LocalDeviceState {
   std::stack<std::unique_ptr<se::Stream>> usage_stream_pool_
       ABSL_GUARDED_BY(stream_pool_mu_);
 
+  absl::once_flag staged_host_to_device_runner_once_;
+  std::unique_ptr<AsyncWorkRunner> staged_host_to_device_runner_;
+
   // Callback map pairs callback stream with a device stream and is used for
   // running short host-side callbacks after device side events, without
   // preventing the device-side stream from doing useful work.
diff --git a/third_party/xla/xla/pjrt/se/se_raw_buffer.cc b/third_party/xla/xla/pjrt/se/se_raw_buffer.cc
index 23401ed95b3..757aae63525 100644
--- a/third_party/xla/xla/pjrt/se/se_raw_buffer.cc
+++ b/third_party/xla/xla/pjrt/se/se_raw_buffer.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include <cstdint>
 #include <cstring>
 #include <memory>
+#include <optional>
 #include <utility>
 #include <vector>
 
@@ -80,10 +81,12 @@ namespace {
 constexpr int64_t kMaxStagingChunksInFlight = 4;
 
 // A pinned host staging buffer for one chunk, and an event recorded after the
-// most recent DMA that used it.
+// most recent DMA that used it. The event comes from the device's EventPool and
+// returns to it when the slot is released, so staged transfers do not create or
+// destroy device events on the hot path.
 struct StagingSlot {
   std::shared_ptr<void> buffer;
-  std::unique_ptr<se::Event> dma_done;
+  std::optional<EventPool::Handle> dma_done;
 };
 
 // Allocates the staging slots for a chunked transfer of `transfer_size` bytes.
@@ -109,7 +112,9 @@ absl::StatusOr<std::vector<StagingSlot>> AllocateStagingSlots(
           "XLA_PJRT_GPU_CC_STAGING_CHUNK_MB).",
           slot_size);
     }
-    ABSL_ASSIGN_OR_RETURN(slot.dma_done, stream->parent()->CreateEvent());
+    ABSL_ASSIGN_OR_RETURN(slot.dma_done,
+                     local_device->event_pool().AllocateEvent(
+                         client->async_work_runner(), stream->parent()));
   }
   return slots;
 }
@@ -125,6 +130,12 @@ absl::StatusOr<std::vector<StagingSlot>> AllocateStagingSlots(
 // event is recorded after already-enqueued stream work, so it completes on GPU
 // progress alone and never waits on work from this runner. Clients that supply
 // a bounded runner would hold one thread per in-flight staged transfer.
+// Chunked staged transfers run only once the device has passed the
+// destination or source allocation event (see `ExecuteWhenAllocationReady`),
+// so the transfer streams never wait on compute work. Chunked H2D transfers to
+// a device run on its single-threaded
+// `LocalDeviceState::staged_host_to_device_runner()`, which only ever waits for
+// its own slots' previous DMAs on `host_to_device_stream()`.
 
 // Copies `transfer_size` bytes from host `src` to device `dst` in
 // `chunk_size`-byte chunks staged through `slots`. Host copies into the
@@ -143,7 +154,10 @@ absl::Status ChunkedStagedHostToDevice(se::Stream* stream, const void* src,
     const int64_t size = std::min(chunk_size, transfer_size - offset);
     StagingSlot& slot = slots[i % num_slots];
     if (i >= num_slots) {
-      ABSL_RETURN_IF_ERROR(slot.dma_done->Synchronize());
+      // Safe to wait in the driver: the slot event only follows this
+      // transfer's own copies on the host-to-device stream, which no longer
+      // waits on compute work.
+      ABSL_RETURN_IF_ERROR(slot.dma_done->event()->Synchronize());
     }
     {
       tsl::profiler::TraceMe trace("H2D Copy To Staging Chunk");
@@ -152,7 +166,7 @@ absl::Status ChunkedStagedHostToDevice(se::Stream* stream, const void* src,
     }
     se::DeviceAddressBase dst_chunk = dst.GetByteSlice(offset, size);
     ABSL_RETURN_IF_ERROR(stream->Memcpy(&dst_chunk, slot.buffer.get(), size));
-    ABSL_RETURN_IF_ERROR(stream->RecordEvent(slot.dma_done.get()));
+    ABSL_RETURN_IF_ERROR(stream->RecordEvent(slot.dma_done->event()));
   }
   return absl::OkStatus();
 }
@@ -174,7 +188,7 @@ absl::Status ChunkedStagedDeviceToHost(se::Stream* stream,
     StagingSlot& slot = slots[i % num_slots];
     ABSL_RETURN_IF_ERROR(stream->Memcpy(slot.buffer.get(),
                                    src.GetByteSlice(offset, size), size));
-    return stream->RecordEvent(slot.dma_done.get());
+    return stream->RecordEvent(slot.dma_done->event());
   };
   for (int64_t i = 0; i < num_slots; ++i) {
     ABSL_RETURN_IF_ERROR(enqueue_dma(i));
@@ -183,7 +197,7 @@ absl::Status ChunkedStagedDeviceToHost(se::Stream* stream,
     const int64_t offset = i * chunk_size;
     const int64_t size = std::min(chunk_size, transfer_size - offset);
     StagingSlot& slot = slots[i % num_slots];
-    ABSL_RETURN_IF_ERROR(slot.dma_done->Synchronize());
+    ABSL_RETURN_IF_ERROR(slot.dma_done->event()->Synchronize());
     {
       tsl::profiler::TraceMe trace("D2H Copy From Staging Chunk");
       std::memcpy(static_cast<char*>(dst) + offset, slot.buffer.get(), size);
@@ -204,6 +218,40 @@ absl::Status DrainStreamOnError(se::Stream* stream, absl::Status status) {
   return status;
 }
 
+// Executes `task` on `runner` once `dependencies` are ready and the device has
+// passed `buffer`'s allocation event, without making a stream or a thread wait
+// for compute work: the allocation event becomes ready from the host callback
+// that `LocalDeviceState::AllocateAndRecordEvent` enqueued when recording it.
+// As in `WaitForAllocation`, an allocation event that failed is not waited
+// for. Sets `device_event` to an error if the allocation event can't be
+// obtained.
+void ExecuteWhenAllocationReady(PjRtStreamExecutorRawClient* client,
+                                tsl::AsyncValueRef<RawSEDeviceMemory> buffer,
+                                BufferSequencingEventRef device_event,
+                                const PjRtDeviceEventRefVector& dependencies,
+                                AsyncWorkRunner* runner,
+                                AsyncWorkRunner::Task task) {
+  ExecuteWhenReady(dependencies, client->async_work_runner(),
+                   [client, buffer = std::move(buffer),
+                    device_event = std::move(device_event), runner,
+                    task = std::move(task)]() mutable {
+                     absl::StatusOr<BufferSequencingEventRef> allocation_event =
+                         buffer->GetDefinitionEvent(client->async_work_runner(),
+                                                    /*nullptr_if_past=*/true);
+                     if (!allocation_event.ok()) {
+                       client->SetEventAsError(device_event,
+                                               allocation_event.status());
+                       return;
+                     }
+                     if (!*allocation_event) {
+                       runner->Execute(std::move(task));
+                       return;
+                     }
+                     runner->ExecuteWhenReady({allocation_event->CopyRCRef()},
+                                              std::move(task));
+                   });
+}
+
 }  // namespace
 
 PjRtStreamExecutorDeviceEventPromise::PjRtStreamExecutorDeviceEventPromise(
@@ -274,7 +322,11 @@ PjRtStreamExecutorRawBuffer::CopyRawHostToDeviceAndReturnEvent(
     // the slots outlive the drain below if AllocateAndRecordEvent fails.
     std::shared_ptr<std::vector<StagingSlot>> staging_slots;
     auto status = [&]() -> absl::Status {
-      ABSL_RETURN_IF_ERROR(client->WaitForAllocation(stream, *buf));
+      // Chunked staged transfers run once their allocation is ready; see
+      // ExecuteWhenAllocationReady below.
+      if (!client->ShouldUseChunkedStaging()) {
+        ABSL_RETURN_IF_ERROR(client->WaitForAllocation(stream, *buf));
+      }
       if (transfer_size > 0) {
         if (client->ShouldStageHostToDeviceTransfers(src, transfer_size)) {
           if (client->GetHostMemoryAllocator() == nullptr) {
@@ -289,9 +341,16 @@ PjRtStreamExecutorRawBuffer::CopyRawHostToDeviceAndReturnEvent(
                                                   transfer_size, chunk_size));
             staging_slots =
                 std::make_shared<std::vector<StagingSlot>>(std::move(slots));
-            return ChunkedStagedHostToDevice(stream, src, sub_buffer,
-                                             transfer_size, chunk_size,
-                                             absl::MakeSpan(*staging_slots));
+            ABSL_RETURN_IF_ERROR(ChunkedStagedHostToDevice(
+                stream, src, sub_buffer, transfer_size, chunk_size,
+                absl::MakeSpan(*staging_slots)));
+            // Nothing waits on the slot events after the copies are issued, so
+            // return them to the pool for the next transfer to reuse. The slot
+            // buffers stay alive until the transfer's completion event.
+            for (StagingSlot& slot : *staging_slots) {
+              slot.dma_done.reset();
+            }
+            return absl::OkStatus();
           }
           HostMemoryAllocator::AllocateOptions alloc_opts;
           alloc_opts.numa_node = stream->parent()->numa_node();
@@ -345,8 +404,17 @@ PjRtStreamExecutorRawBuffer::CopyRawHostToDeviceAndReturnEvent(
   // `run_transfer` may be deferred until `dependencies` are ready, so record
   // the allocation event now.
   client_->MaterializeAllocationEvent(*this);
-  ExecuteWhenReady(dependencies, client_->async_work_runner(),
-                   std::move(run_transfer));
+  if (client_->ShouldUseChunkedStaging()) {
+    // Issue the copies on the device's staged host-to-device runner, one
+    // transfer at a time. A transfer whose allocation is not ready yet does
+    // not hold up later transfers to the same device.
+    ExecuteWhenAllocationReady(
+        client_, device_buffer_, device_event, dependencies,
+        local_device_->staged_host_to_device_runner(), std::move(run_transfer));
+  } else {
+    ExecuteWhenReady(dependencies, client_->async_work_runner(),
+                     std::move(run_transfer));
+  }
 
   return PjRtDeviceEventRef(std::move(device_event));
 }
@@ -376,7 +444,10 @@ PjRtStreamExecutorRawBuffer::CopyRawDeviceToHostAndReturnEvent(
     }
     std::shared_ptr<void> staging_buffer;
     auto status = [&]() -> absl::Status {
-      ABSL_RETURN_IF_ERROR(client->WaitForAllocation(stream, *buf));
+      // See CopyRawHostToDeviceAndReturnEvent.
+      if (!client->ShouldUseChunkedStaging()) {
+        ABSL_RETURN_IF_ERROR(client->WaitForAllocation(stream, *buf));
+      }
       if (transfer_size > 0) {
         if (client->ShouldStageHostToDeviceTransfers(dst, transfer_size)) {
           if (client->GetHostMemoryAllocator() == nullptr) {
@@ -442,8 +513,14 @@ PjRtStreamExecutorRawBuffer::CopyRawDeviceToHostAndReturnEvent(
   // `run_transfer` may be deferred until `dependencies` are ready, so record
   // the allocation event now.
   client_->MaterializeAllocationEvent(*this);
-  ExecuteWhenReady(dependencies, client_->async_work_runner(),
-                   std::move(run_transfer));
+  if (client_->ShouldUseChunkedStaging()) {
+    ExecuteWhenAllocationReady(client_, device_buffer_, device_event,
+                               dependencies, client_->async_work_runner(),
+                               std::move(run_transfer));
+  } else {
+    ExecuteWhenReady(dependencies, client_->async_work_runner(),
+                     std::move(run_transfer));
+  }
 
   return PjRtDeviceEventRef(std::move(device_event));
 }

From 0e1fd5f2c72f34625260b0cc45220381b3b4cb3f Mon Sep 17 00:00:00 2001
From: bodapatisaikrishna <saikrishnabodapati@gmail.com>
Date: Wed, 14 Oct 2026 22:34:39 +0530
Subject: [PATCH 0797/1000] Add initialization checks in resource scatter and
 gather ops

---
 .../core/kernels/resource_variable_ops.cc     | 16 ++++
 tensorflow/core/kernels/scatter_nd_op.cc      |  4 +
 .../python/training/training_ops_test.py      | 74 +++++++++++++++++++
 3 files changed, 94 insertions(+)

diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc
index bd57da5bd44..66b3f74e758 100644
--- a/tensorflow/core/kernels/resource_variable_ops.cc
+++ b/tensorflow/core/kernels/resource_variable_ops.cc
@@ -768,6 +768,10 @@ class ResourceGatherOp : public OpKernel {
     // reference count greater than one and make a copy of the
     // (potentially very large) tensor buffer.
     tf_shared_lock ml(*v->mu());
+    OP_REQUIRES(c, v->tensor()->IsInitialized(),
+                errors::FailedPrecondition(
+                    "Attempting to use uninitialized variables: ",
+                    handle.name()));
     OP_REQUIRES(
         c, v->tensor()->dtype() == DataTypeToEnum<T>::v(),
         absl::InvalidArgumentError(absl::StrCat(
@@ -961,6 +965,10 @@ class ResourceGatherNdOp : public OpKernel {
     // reference count greater than one and make a copy of the
     // (potentially very large) tensor buffer.
     tf_shared_lock ml(*v->mu());
+    OP_REQUIRES(c, v->tensor()->IsInitialized(),
+                errors::FailedPrecondition(
+                    "Attempting to use uninitialized variables: ",
+                    handle.name()));
     OP_REQUIRES(
         c, v->tensor()->dtype() == DataTypeToEnum<T>::v(),
         absl::InvalidArgumentError(absl::StrCat(
@@ -1192,6 +1200,10 @@ class ResourceScatterUpdateOp : public OpKernel {
         update_dtype == DT_STRING || update_dtype == DT_VARIANT;
     if (is_non_pod_dtype || use_exclusive_lock_) {
       mutex_lock ml(*v->mu());
+      OP_REQUIRES(c, v->tensor()->IsInitialized(),
+                  errors::FailedPrecondition(
+                      "Attempting to use uninitialized variables: ",
+                      handle.name()));
       OP_REQUIRES(c, v->tensor()->dtype() == update_dtype,
                   absl::InvalidArgumentError(
                       "DType of scatter resource and updates does not match."));
@@ -1199,6 +1211,10 @@ class ResourceScatterUpdateOp : public OpKernel {
     } else {
       // For POD dtypes, we can safely run the update without the mutex.
       tf_shared_lock ml(*v->mu());
+      OP_REQUIRES(c, v->tensor()->IsInitialized(),
+                  errors::FailedPrecondition(
+                      "Attempting to use uninitialized variables: ",
+                      handle.name()));
       OP_REQUIRES(c, v->tensor()->dtype() == update_dtype,
                   absl::InvalidArgumentError(
                       "DType of scatter resource and updates does not match."));
diff --git a/tensorflow/core/kernels/scatter_nd_op.cc b/tensorflow/core/kernels/scatter_nd_op.cc
index 49d8236dcd0..b2db94b89d3 100644
--- a/tensorflow/core/kernels/scatter_nd_op.cc
+++ b/tensorflow/core/kernels/scatter_nd_op.cc
@@ -323,6 +323,10 @@ class ScatterNdUpdateOp : public ScatterOpBase<Device> {
       OP_REQUIRES_OK(c, LookupResource(c, handle, &v));
       OP_REQUIRES_OK(c, EnsureSparseVariableAccess<Device, T>(c, v.get()));
       mutex_lock m(*v->mu());
+      OP_REQUIRES(c, v->tensor()->IsInitialized(),
+                  errors::FailedPrecondition(
+                      "Attempting to use uninitialized variables: ",
+                      handle.name()));
       DoCompute(c);
     } else if (use_exclusive_lock_) {
       // If we're here, it means the input type is a ref.
diff --git a/tensorflow/python/training/training_ops_test.py b/tensorflow/python/training/training_ops_test.py
index 12daccab7fb..c43e6adb9c4 100644
--- a/tensorflow/python/training/training_ops_test.py
+++ b/tensorflow/python/training/training_ops_test.py
@@ -29,6 +29,7 @@ from tensorflow.python.framework import test_util
 from tensorflow.python.framework.test_util import TensorFlowTestCase
 # Import resource_variable_ops for the variables-to-tensor implicit conversion.
 from tensorflow.python.ops import gen_resource_variable_ops
+from tensorflow.python.ops import gen_state_ops
 from tensorflow.python.ops import gen_training_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import resource_variable_ops  # pylint: disable=unused-import
@@ -845,5 +846,78 @@ class TrainingOpsTest(TensorFlowTestCase):
       )
 
 
+    scatter_var_uninit_f16 = gen_resource_variable_ops.var_handle_op(
+        dtype=dtypes.float16,
+        shape=[1, 2],
+        shared_name=(
+            context.anonymous_name() if context.executing_eagerly() else ""
+        ),
+    )
+    with self.assertRaisesRegex(
+        errors.FailedPreconditionError, "uninitialized"
+    ):
+      self.evaluate(
+          gen_resource_variable_ops.resource_scatter_update(
+              resource=scatter_var_uninit_f16,
+              indices=indices,
+              updates=grad,
+          )
+      )
+
+    gather_var_uninit_f16 = gen_resource_variable_ops.var_handle_op(
+        dtype=dtypes.float16,
+        shape=[1, 2],
+        shared_name=(
+            context.anonymous_name() if context.executing_eagerly() else ""
+        ),
+    )
+    with self.assertRaisesRegex(
+        errors.FailedPreconditionError, "uninitialized"
+    ):
+      self.evaluate(
+          gen_resource_variable_ops.resource_gather(
+              resource=gather_var_uninit_f16,
+              indices=indices,
+              dtype=dtypes.float16,
+          )
+      )
+
+    gather_nd_var_uninit_f16 = gen_resource_variable_ops.var_handle_op(
+        dtype=dtypes.float16,
+        shape=[1, 2],
+        shared_name=(
+            context.anonymous_name() if context.executing_eagerly() else ""
+        ),
+    )
+    with self.assertRaisesRegex(
+        errors.FailedPreconditionError, "uninitialized"
+    ):
+      self.evaluate(
+          gen_resource_variable_ops.resource_gather_nd(
+              resource=gather_nd_var_uninit_f16,
+              indices=constant_op.constant([[0]], dtype=dtypes.int32),
+              dtype=dtypes.float16,
+          )
+      )
+
+    scatter_nd_var_uninit_f16 = gen_resource_variable_ops.var_handle_op(
+        dtype=dtypes.float16,
+        shape=[1, 2],
+        shared_name=(
+            context.anonymous_name() if context.executing_eagerly() else ""
+        ),
+    )
+    with self.assertRaisesRegex(
+        errors.FailedPreconditionError, "uninitialized"
+    ):
+      self.evaluate(
+          gen_state_ops.resource_scatter_nd_update(
+              ref=scatter_nd_var_uninit_f16,
+              indices=constant_op.constant([[0]], dtype=dtypes.int32),
+              updates=grad,
+          )
+      )
+
+
 if __name__ == '__main__':
   googletest.main()

From 803243afb00b3a19d7ddd529fb189a07f26aa27c Mon Sep 17 00:00:00 2001
From: Dmitri Latushko <latushko@google.com>
Date: Wed, 14 Oct 2026 09:45:59 -0700
Subject: [PATCH 0798/1000] Restore --xla_cpu_experimental_ynn_fusion_type= in
 saved_model_aot_compile.py to fix tfcompile ynn-fusion failure

PiperOrigin-RevId: 999432777
---
 tensorflow/python/tools/saved_model_aot_compile.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tensorflow/python/tools/saved_model_aot_compile.py b/tensorflow/python/tools/saved_model_aot_compile.py
index 335979b295e..dc6dda79591 100644
--- a/tensorflow/python/tools/saved_model_aot_compile.py
+++ b/tensorflow/python/tools/saved_model_aot_compile.py
@@ -367,6 +367,9 @@ def aot_compile_cpu_meta_graph_def(checkpoint_path,
     else:
       xla_flags += ' --xla_cpu_multi_thread_eigen={}'.format(
           'true' if multithreading else 'false')
+    # Disable YNNPACK fusions for tfcompile, which does not support ynn-fusion
+    # thunks (b/452693819, b/573087829).
+    xla_flags += ' --xla_cpu_experimental_ynn_fusion_type= '
     os.environ['XLA_FLAGS'] = xla_flags
 
   temp_dir = test.get_temp_dir()

From f825dfd6c1c31133a7d6b1f00e072b28f3d3a681 Mon Sep 17 00:00:00 2001
From: Tori Baker <vwbaker@google.com>
Date: Wed, 14 Oct 2026 09:55:42 -0700
Subject: [PATCH 0799/1000] Look through bitcasts when checking if multiply
 operands are worth fusing as inputs.

This has some benchmark improvements.

PiperOrigin-RevId: 999438395
---
 third_party/xla/xla/service/gpu/BUILD         |  1 +
 .../service/gpu/triton_tiling_propagation.cc  | 16 +++++-
 .../gpu/triton_tiling_propagation_test.cc     | 52 +++++++++++++++++++
 3 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index d74471d8f1a..a3b3fcb2148 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -1222,6 +1222,7 @@ xla_cc_test(
     srcs = ["triton_tiling_propagation_test.cc"],
     deps = [
         ":triton_tiling_propagation",
+        "//xla:xla_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
         "//xla/tests:xla_internal_test_main",
diff --git a/third_party/xla/xla/service/gpu/triton_tiling_propagation.cc b/third_party/xla/xla/service/gpu/triton_tiling_propagation.cc
index 511afe95467..60dd7bfa249 100644
--- a/third_party/xla/xla/service/gpu/triton_tiling_propagation.cc
+++ b/third_party/xla/xla/service/gpu/triton_tiling_propagation.cc
@@ -1054,6 +1054,15 @@ bool IsCheapToDuplicate(const HloInstruction& hlo) {
          AllUsersAreSlices(hlo);
 }
 
+// Returns the first instruction in the operand chain of `hlo` (including `hlo`
+// itself) that is not a bitcast.
+const HloInstruction* SkipBitcasts(const HloInstruction* hlo) {
+  while (hlo->opcode() == HloOpcode::kBitcast) {
+    hlo = hlo->operand(0);
+  }
+  return hlo;
+}
+
 }  // namespace
 
 // Tells that fusing an instruction as an input is efficient.
@@ -1100,9 +1109,12 @@ bool IsInputWorthFusing(const HloInstruction& hlo) {
           .debug_options()
           .xla_gpu_experimental_enable_subchannel_dequantisation_fusion();
   if (hlo.opcode() == HloOpcode::kMultiply) {
+    // Bitcasts don't change the data volume, so on their own they always look
+    // worth fusing. Look through them to judge the actual producer instead,
+    // which may not be (e.g. another dot fusion).
     return enable_subchannel_dequantisation_fusion &&
-           IsInputWorthFusing(*hlo.operand(0)) &&
-           IsInputWorthFusing(*hlo.operand(1));
+           IsInputWorthFusing(*SkipBitcasts(hlo.operand(0))) &&
+           IsInputWorthFusing(*SkipBitcasts(hlo.operand(1)));
   }
   return hlo_query::AllOperandsAreParametersOrConstantsWithSingleUser(hlo);
 }
diff --git a/third_party/xla/xla/service/gpu/triton_tiling_propagation_test.cc b/third_party/xla/xla/service/gpu/triton_tiling_propagation_test.cc
index be646b81c9e..d6c3be6c7e6 100644
--- a/third_party/xla/xla/service/gpu/triton_tiling_propagation_test.cc
+++ b/third_party/xla/xla/service/gpu/triton_tiling_propagation_test.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include <gtest/gtest.h>
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
+#include "xla/xla.pb.h"
 
 namespace xla::gpu {
 namespace {
@@ -145,5 +146,56 @@ ENTRY entry {
   EXPECT_FALSE(triton_fusion::IsInputWorthFusing(*slice));
 }
 
+TEST_F(TritonTilingPropagationTest,
+       IsInputWorthFusingMultiplyLooksThroughBitcastsOfOperands) {
+  // `dot` reads more data than it produces, so the multiply is not worth fusing
+  // as an input. The bitcasts between `dot` and the multiply don't change that.
+  const char* hlo_string = R"(
+HloModule module
+
+ENTRY entry {
+  p0 = f32[32,64] parameter(0)
+  p1 = f32[64,128] parameter(1)
+  p2 = f32[32,2,64] parameter(2)
+  negate = f32[32,64] negate(p0)
+  dot = f32[32,128] dot(negate, p1),
+    lhs_contracting_dims={1}, rhs_contracting_dims={0}
+  bitcast0 = f32[4096] bitcast(dot)
+  bitcast1 = f32[32,2,64] bitcast(bitcast0)
+  ROOT multiply = f32[32,2,64] multiply(p2, bitcast1)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  module->mutable_config()
+      .mutable_debug_options()
+      .set_xla_gpu_experimental_enable_subchannel_dequantisation_fusion(true);
+  EXPECT_FALSE(triton_fusion::IsInputWorthFusing(
+      *module->entry_computation()->root_instruction()));
+}
+
+TEST_F(TritonTilingPropagationTest,
+       IsInputWorthFusingDequantizationWithBitcastBeforeMultiply) {
+  // Subchannel dequantization with a bitcast between the convert and the
+  // multiply is still worth fusing as an input.
+  const char* hlo_string = R"(
+HloModule module
+
+ENTRY entry {
+  w = s8[32,128] parameter(0)
+  s = f32[32,2] parameter(1)
+  w_f32 = f32[32,128] convert(w)
+  w_split = f32[32,2,64] bitcast(w_f32)
+  s_b = f32[32,2,64] broadcast(s), dimensions={0,1}
+  ROOT multiply = f32[32,2,64] multiply(w_split, s_b)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  module->mutable_config()
+      .mutable_debug_options()
+      .set_xla_gpu_experimental_enable_subchannel_dequantisation_fusion(true);
+  EXPECT_TRUE(triton_fusion::IsInputWorthFusing(
+      *module->entry_computation()->root_instruction()));
+}
+
 }  // namespace
 }  // namespace xla::gpu

From 2fe47e3a54cbcb2db0f92c01437b778e7896dd3a Mon Sep 17 00:00:00 2001
From: kaivalya-cyber <141600539+kaivalya-cyber@users.noreply.github.com>
Date: Wed, 14 Oct 2026 11:06:18 -0700
Subject: [PATCH 0800/1000] Catch ValueError alongside InvalidArgumentError in
 roll OOB tests

Shape inference can surface a statically-known out-of-bounds axis as
ValueError during @tf.function tracing, which would escape an
InvalidArgumentError-only assertRaisesRegex. Catch both, matching the
existing testFlip pattern, and parameterize the duplicated
out-of-bounds blocks into a single loop for maintainability.
---
 .../python/ops/numpy_ops/np_array_ops_test.py | 45 +++++++------------
 1 file changed, 15 insertions(+), 30 deletions(-)

diff --git a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
index a2ce3355c99..aa5406b5182 100644
--- a/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_array_ops_test.py
@@ -1691,34 +1691,19 @@ class ArrayMethodsTest(test.TestCase):
     _test(np.array(42.0), 1)
 
     # Out-of-bounds axes raise, checked by the backend kernel.
-    with self.assertRaisesRegex(
-        errors_impl.InvalidArgumentError, 'out of range'
-    ):
-      np_array_ops.roll(a, 1, axis=3)
-    with self.assertRaisesRegex(
-        errors_impl.InvalidArgumentError, 'out of range'
-    ):
-      np_array_ops.roll(a, 1, axis=-4)
-    with self.assertRaisesRegex(
-        errors_impl.InvalidArgumentError, 'out of range'
-    ):
-      np_array_ops.roll(a, (1, 2), axis=(0, 3))
-    with self.assertRaisesRegex(
-        errors_impl.InvalidArgumentError, 'out of range'
-    ):
-      np_array_ops.roll(a, (1, -1), axis=[0, -4])
-    with self.assertRaisesRegex(
-        errors_impl.InvalidArgumentError, 'out of range'
-    ):
-      np_array_ops.roll(a, range(2), axis=range(2, 4))
-    with self.assertRaisesRegex(
-        errors_impl.InvalidArgumentError, 'out of range'
-    ):
-      np_array_ops.roll(a, 1, axis=1 << 33)
-    with self.assertRaisesRegex(
-        errors_impl.InvalidArgumentError, 'out of range'
-    ):
-      np_array_ops.roll(a, 1, axis=-(1 << 33))
+    for shift, invalid_axis in [
+        (1, 3),
+        (1, -4),
+        ((1, 2), (0, 3)),
+        ((1, -1), [0, -4]),
+        (range(2), range(2, 4)),
+        (1, 1 << 33),
+        (1, -(1 << 33)),
+    ]:
+      with self.assertRaisesRegex(
+          (errors_impl.InvalidArgumentError, ValueError), 'out of range'
+      ):
+        np_array_ops.roll(a, shift, axis=invalid_axis)
 
     # Graph-mode tracing of roll with concrete axes.
     @def_function.function
@@ -1734,11 +1719,11 @@ class ArrayMethodsTest(test.TestCase):
         f(a, constant_op.constant(1), constant_op.constant(0)),
     )
     with self.assertRaisesRegex(
-        errors_impl.InvalidArgumentError, 'out of range'
+        (errors_impl.InvalidArgumentError, ValueError), 'out of range'
     ):
       f(a, 1, 3)
     with self.assertRaisesRegex(
-        errors_impl.InvalidArgumentError, 'out of range'
+        (errors_impl.InvalidArgumentError, ValueError), 'out of range'
     ):
       f(a, constant_op.constant(1), constant_op.constant(3))
 

From 63149a5674913650ebc970bf029ef8059b6955e2 Mon Sep 17 00:00:00 2001
From: Akhil Goel <akhil.goel@intel.com>
Date: Wed, 14 Oct 2026 10:48:57 -0700
Subject: [PATCH 0801/1000] PR #50663: [XLA:GPU][oneAPI] Align XPU Triton hook
 with upstream LLVM

Imported from GitHub PR https://github.com/openxla/xla/pull/50663

This PR aligns the XPU Triton hook with https://github.com/llvm/llvm-project/commit/4a0a7e3ddc3bdc4fde7c5b9b9e6890fb4dbf9153. It includes the fix for this breaking LLVM change https://github.com/llvm/llvm-project/commit/a0957bb89e3660ad034e42f05831c184e166adc5 that adds a noundef metadata to llvm.load op and is part of the most recent LLVM update.
Copybara import of the project:

--
4196397d78f245dbc0edb099f39463016d3ea9f6 by Akhil Goel <akhil.goel@intel.com>:

Update XPU Triton hook

Merging this change closes #50663

PiperOrigin-RevId: 999472727
---
 third_party/xla/third_party/triton/intel_xpu/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/third_party/triton/intel_xpu/workspace.bzl b/third_party/xla/third_party/triton/intel_xpu/workspace.bzl
index 594ae3e13ea..a2318a6dcef 100644
--- a/third_party/xla/third_party/triton/intel_xpu/workspace.bzl
+++ b/third_party/xla/third_party/triton/intel_xpu/workspace.bzl
@@ -1,7 +1,7 @@
 """Intel XPU Triton archive metadata."""
 
-XPU_TRITON_COMMIT = "616ef4c791a0feb79243689b985ca2a2afd9fa6b"
-XPU_TRITON_SHA256 = "e51b47dd8f6a59bd3574964337fc8b1deef04c34a2ff8fbe2a4ba2cba134f1cf"
+XPU_TRITON_COMMIT = "33f9cb578f45b8fb719c728896b2a75f29e66616"
+XPU_TRITON_SHA256 = "856f7c73cd5dd43a9157709f4c17ca70fbb1e2504803711e66d62a9e0dadadd1"
 
 def use_xpu_triton(repository_ctx):
     return repository_ctx.getenv("ENABLE_INTEL_XPU_TRITON", "").strip() == "1"

From b2653741b4dc0374765287d4af2d0db611e36ca4 Mon Sep 17 00:00:00 2001
From: Tori Baker <vwbaker@google.com>
Date: Wed, 14 Oct 2026 11:22:47 -0700
Subject: [PATCH 0802/1000] [XLA:GPU] Support moving f8e8m0fnu values around in
 Triton fusions.

Run TritonXLAConvertUnsupportedTypesPass (f8e8m0fnu -> i8) for all modules
instead of only those with tt.dot_scaled, and mark f8e8m0fnu as supported for
ops that don't change values. This lets GemmFusionV2 fuse scaled dot scale
producers. This allows our new support.cc rather than the legacy support to work for f8e8m0fnu.

Also reject reduce-precision on f8e8m0fnu, which would be miscompiled.

PiperOrigin-RevId: 999495751
---
 .../backends/gpu/codegen/triton/support.cc    | 46 ++++++++++++++++---
 .../triton_xla_convert_unsupported_types.mlir | 30 +++++++++++-
 .../triton_xla_convert_unsupported_types.cc   | 20 +++-----
 .../gpu/codegen/triton/xtile_compiler.cc      |  8 +---
 .../codegen/xtile/codegen/emitter_helpers.cc  | 10 +++-
 5 files changed, 84 insertions(+), 30 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/support.cc b/third_party/xla/xla/backends/gpu/codegen/triton/support.cc
index 883d6c00665..8d222408670 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/support.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/support.cc
@@ -101,6 +101,33 @@ CodegenDecision IsTritonSupportedDataType(
   }
 }
 
+// Like `IsTritonSupportedDataType`, but also allows types that are only
+// supported by some opcodes.
+CodegenDecision IsTritonSupportedDataTypeForOpcode(
+    PrimitiveType type, HloOpcode opcode,
+    const se::GpuComputeCapability& gpu_version) {
+  // F8E8M0FNU is lowered as i8, so it supports ops that don't change values.
+  if (type == F8E8M0FNU && gpu_version.IsCuda()) {
+    switch (opcode) {
+      case HloOpcode::kBitcast:
+      case HloOpcode::kBroadcast:
+      case HloOpcode::kConcatenate:
+      case HloOpcode::kConstant:
+      case HloOpcode::kCopy:
+      case HloOpcode::kPad:
+      case HloOpcode::kParameter:
+      case HloOpcode::kReshape:
+      case HloOpcode::kSelect:
+      case HloOpcode::kSlice:
+      case HloOpcode::kTranspose:
+        return CodegenDecision::Allow();
+      default:
+        break;
+    }
+  }
+  return IsTritonSupportedDataType(type, gpu_version);
+}
+
 // Set of unary elementwise ops that are genuinely supported by Triton.
 absl::flat_hash_set<HloOpcode> TritonSupportedUnaryElementwiseOps(
     PrimitiveType element_type, const se::GpuComputeCapability& gpu_version) {
@@ -223,6 +250,12 @@ CodegenDecision IsTritonSupportedConversion(
     return error_message("Unsupported PRED to U8 conversion.");
   }
 
+  // A no-op conversion only moves data around.
+  if (input == output) {
+    return IsTritonSupportedDataTypeForOpcode(input, HloOpcode::kCopy,
+                                              gpu_version);
+  }
+
   if (!IsTritonSupportedDataType(input, gpu_version)) {
     return CodegenDecision::Forbid(
         absl::StrCat("Unsupported input type for conversion: ",
@@ -825,12 +858,13 @@ CodegenDecision IsTritonSupportedInstructionImpl(
         return CodegenDecision::Forbid(
             "Nested tuple or token shape is not supported.");
       }
-      output_type_is_supported = output_type_is_supported.And(
-          IsTritonSupportedDataType(shape.element_type(), gpu_version));
+      output_type_is_supported =
+          output_type_is_supported.And(IsTritonSupportedDataTypeForOpcode(
+              shape.element_type(), instr.opcode(), gpu_version));
     }
   } else {
-    output_type_is_supported =
-        IsTritonSupportedDataType(instr.shape().element_type(), gpu_version);
+    output_type_is_supported = IsTritonSupportedDataTypeForOpcode(
+        instr.shape().element_type(), instr.opcode(), gpu_version);
   }
 
   if (!output_type_is_supported) {
@@ -840,8 +874,8 @@ CodegenDecision IsTritonSupportedInstructionImpl(
 
   CodegenDecision input_types_are_supported = CodegenDecision(
       absl::c_all_of(instr.operands(), [&](const HloInstruction* operand) {
-        return IsTritonSupportedDataType(operand->shape().element_type(),
-                                         gpu_version);
+        return IsTritonSupportedDataTypeForOpcode(
+            operand->shape().element_type(), instr.opcode(), gpu_version);
       }));
 
   if (!input_types_are_supported) {
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_convert_unsupported_types.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_convert_unsupported_types.mlir
index 0394a31a129..53c7facf3b1 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_convert_unsupported_types.mlir
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_convert_unsupported_types.mlir
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 // ==============================================================================
-// RUN: xla-opt --convert-unsupported-types %s | FileCheck %s
+// RUN: xla-opt --convert-unsupported-types --split-input-file %s | FileCheck %s
 
 module {
   // CHECK:   xtile.entry_func @triton_fn(
@@ -58,3 +58,31 @@ module {
     xtile.return
   }
 }
+
+
+// -----
+
+// f8E8M0FNU values are converted to i8 even if the module has no tt.dot_scaled.
+module {
+  // CHECK:   xtile.entry_func @move_e8m0(
+  // CHECK-SAME:   %arg0: memref<4x16xi8, #xtile.layout<[1, 0]>>,
+  // CHECK-SAME:   %arg1: memref<4x16xi8, #xtile.layout<[1, 0]>>,
+  xtile.entry_func @move_e8m0(
+      %arg0: memref<4x16xf8E8M0FNU, #xtile.layout<[1, 0]>>,
+      %arg1: memref<4x16xf8E8M0FNU, #xtile.layout<[1, 0]>>,
+      %tile_id: index) {
+    %c_0 = arith.constant 0 : index
+    // CHECK: %[[EXTRACTED:.*]] = xtile.extract %arg0{{.*}} -> tensor<4x16xi8>
+    %extracted = xtile.extract %arg0[%c_0, %c_0] [4, 16] [1, 1] : memref<4x16xf8E8M0FNU, #xtile.layout<[1, 0]>> -> tensor<4x16xf8E8M0FNU>
+    // CHECK: %[[CST:.*]] = arith.constant dense<0> : tensor<4x16xi8>
+    %cst = arith.constant dense<5.877470e-39> : tensor<4x16xf8E8M0FNU>
+    %as_int = arith.bitcast %extracted : tensor<4x16xf8E8M0FNU> to tensor<4x16xi8>
+    %zero = arith.constant dense<0> : tensor<4x16xi8>
+    %mask = arith.cmpi sge, %as_int, %zero : tensor<4x16xi8>
+    // CHECK: %[[SELECTED:.*]] = arith.select %{{.*}}, %[[EXTRACTED]], %[[CST]] : tensor<4x16xi1>, tensor<4x16xi8>
+    %selected = arith.select %mask, %extracted, %cst : tensor<4x16xi1>, tensor<4x16xf8E8M0FNU>
+    // CHECK: xtile.insert %[[SELECTED]] into %arg1{{.*}} : tensor<4x16xi8> -> memref<4x16xi8, #xtile.layout<[1, 0]>>
+    xtile.insert %selected into %arg1[%c_0, %c_0] [4, 16] [1, 1] : tensor<4x16xf8E8M0FNU> -> memref<4x16xf8E8M0FNU, #xtile.layout<[1, 0]>>
+    xtile.return
+  }
+}
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_convert_unsupported_types.cc b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_convert_unsupported_types.cc
index 2756c652d95..2ef59efeb88 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_convert_unsupported_types.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_convert_unsupported_types.cc
@@ -38,16 +38,6 @@ namespace mlir::triton::xla {
 
 namespace {
 
-bool FindScaledDotOp(const ModuleOp& module) {
-  auto walk_result = module->walk([&](Operation* op) {
-    if (auto extSI = dyn_cast<triton::DotScaledOp>(op)) {
-      return WalkResult::interrupt();
-    }
-    return WalkResult::advance();
-  });
-  return walk_result.wasInterrupted();
-}
-
 template <typename OpType>
 struct GenericOpConversionPattern final : public OpConversionPattern<OpType> {
   using OpConversionPattern<OpType>::OpConversionPattern;
@@ -97,9 +87,8 @@ class TritonXLAConvertUnsupportedTypesPass
 
  private:
   void runOnOperation() override {
-    if (!FindScaledDotOp(getOperation())) {
-      return;
-    }
+    // Triton has no support for f8E8M0FNU outside of tt.dot_scaled, so we
+    // store and move such values around as i8.
     TypeConverter converter;
     converter.addConversion([](Type type) { return type; });
     converter.addConversion([&](Float8E8M0FNUType type) {
@@ -138,15 +127,18 @@ class TritonXLAConvertUnsupportedTypesPass
         ConstantOpConversionPattern,
         GenericOpConversionPattern<::xla::xtile::ExtractTileOp>,
         GenericOpConversionPattern<::xla::xtile::InsertTileOp>,
+        GenericOpConversionPattern<BitcastOp>,
         GenericOpConversionPattern<BroadcastOp>,
         GenericOpConversionPattern<DotScaledOp>,
         GenericOpConversionPattern<ExpandDimsOp>,
         GenericOpConversionPattern<ReshapeOp>,
         GenericOpConversionPattern<SplatOp>,
         GenericOpConversionPattern<TransOp>,
+        GenericOpConversionPattern<UnsplatOp>,
         GenericOpConversionPattern<arith::BitcastOp>,
         GenericOpConversionPattern<arith::SelectOp>,
-        GenericOpConversionPattern<tensor::ExtractOp>
+        GenericOpConversionPattern<tensor::ExtractOp>,
+        GenericOpConversionPattern<tensor::FromElementsOp>
         // go/keep-sorted end
         >(converter, ctx);
     scf::populateSCFStructuralTypeConversions(converter, patterns);
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc b/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc
index 82af8119569..4c66f0ecacf 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc
@@ -734,13 +734,7 @@ absl::Status LowerXTileToTriton(
     pm.addPass(
         mlir::triton::xla::createTritonXLAMathToLibdevicePass(math_options));
 
-    if (fusion.GetModule()
-            ->config()
-            .debug_options()
-            .xla_gpu_experimental_scaled_dot_with_triton()) {
-      pm.addPass(
-          mlir::triton::xla::createTritonXLAConvertUnsupportedTypesPass());
-    }
+    pm.addPass(mlir::triton::xla::createTritonXLAConvertUnsupportedTypesPass());
     tsl::StatusScopedDiagnosticHandler diagnostic_handler(&mlir_context);
     if (absl::Status status =
             diagnostic_handler.consumeStatus(pm.run(xtile_dialect_module));
diff --git a/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.cc b/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.cc
index 067430edc09..b15a6ccdb4f 100644
--- a/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.cc
+++ b/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.cc
@@ -199,8 +199,14 @@ mlir::Value OnesLike(mlir::ImplicitLocOpBuilder& b, mlir::Type type) {
   return cst;
 }
 
-Value ReducePrecision(mlir::ImplicitLocOpBuilder& b, const HloInstruction& hlo,
-                      Value value) {
+absl::StatusOr<Value> ReducePrecision(mlir::ImplicitLocOpBuilder& b,
+                                      const HloInstruction& hlo, Value value) {
+  // The lowering assumes an IEEE-like layout with a sign bit and a mantissa,
+  // which F8E8M0FNU does not have.
+  if (hlo.operand(0)->shape().element_type() == F8E8M0FNU) {
+    return absl::UnimplementedError(
+        "ReducePrecision is not supported for F8E8M0FNU.");
+  }
   // Check if it's a no-op.
   auto elem_type =
       GetPrimitiveType(mlir::getElementTypeOrSelf(value.getType()));

From 67336f3ec399b2374baab728353bee31a881f35c Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Wed, 14 Oct 2026 12:51:00 -0700
Subject: [PATCH 0803/1000] Set memory_kind in IFRT Proxy CopyArrays when it
 differs from source.

PiperOrigin-RevId: 999556122
---
 .../xla/python/ifrt_proxy/client/client.cc    | 10 ++---
 .../python/ifrt_proxy/client/client_test.cc   | 44 +++++++++++++++++++
 2 files changed, 48 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/xla/python/ifrt_proxy/client/client.cc b/third_party/xla/xla/python/ifrt_proxy/client/client.cc
index 99146ac4f00..d5936ceb657 100644
--- a/third_party/xla/xla/python/ifrt_proxy/client/client.cc
+++ b/third_party/xla/xla/python/ifrt_proxy/client/client.cc
@@ -322,13 +322,11 @@ absl::StatusOr<std::vector<xla::ifrt::ArrayRef>> Client::CopyArrays(
       req->add_device_ids(device->Id().value());
     }
   }
-  if (memory_kind.has_value()) {
-    // Use an empty string to indicate the default memory kind.
+  if (memory_kind.has_value() &&
+      *memory_kind != arrays[0]->sharding().memory_kind()) {
     // OSS requires explicit string conversion
-    if (!memory_kind->is_default()) {
-      // NOLINTNEXTLINE(*-redundant-string-conversions)
-      req->set_memory_kind(std::string(memory_kind->value()));
-    }
+    // NOLINTNEXTLINE(*-redundant-string-conversions)
+    req->set_memory_kind(std::string(memory_kind->value()));
   }
   req->set_copy_semantics(ToArrayCopySemanticsProto(semantics));
 
diff --git a/third_party/xla/xla/python/ifrt_proxy/client/client_test.cc b/third_party/xla/xla/python/ifrt_proxy/client/client_test.cc
index 4d2494fe39a..6b79b692a54 100644
--- a/third_party/xla/xla/python/ifrt_proxy/client/client_test.cc
+++ b/third_party/xla/xla/python/ifrt_proxy/client/client_test.cc
@@ -15,6 +15,7 @@
 #include "xla/python/ifrt_proxy/client/client.h"
 
 #include <memory>
+#include <optional>
 #include <string>
 #include <utility>
 #include <vector>
@@ -299,6 +300,49 @@ TEST_P(ClientTest, CopyArraysDefaultLayoutSuccess) {
   EXPECT_EQ(layout_2, nullptr);
 }
 
+TEST_P(ClientTest, CopyArraysMemoryKindCombinations) {
+  IfrtResponse response;
+  response.mutable_copy_arrays_response()->add_array_handles(1);
+  EXPECT_CALL(*session_,
+              Enqueue(IfrtRequestOfType(IfrtRequest::kCopyArraysRequest)))
+      .WillRepeatedly(MockClientSessionReturnResponse(response));
+  EXPECT_CALL(*session_,
+              Enqueue(IfrtRequestOfType(IfrtRequest::kDestructArrayRequest)))
+      .WillRepeatedly(MockClientSessionReturnResponse(IfrtResponse()));
+
+  struct Case {
+    std::optional<MemoryKind> kind;
+    absl::string_view expected;
+  };
+  const std::vector<Case> kCases = {
+      {std::nullopt, "device"},
+      {MemoryKind(), "device"},
+      {MemoryKind("device"), "device"},
+      {MemoryKind("pinned_host"), "pinned_host"},
+  };
+
+  for (const Case& src : kCases) {
+    for (const Case& dst : kCases) {
+      auto sharding = SingleDeviceSharding::Create(
+          device_, src.kind.value_or(MemoryKind()));
+      std::vector<ArrayRef> arrays = {tsl::MakeRef<Array>(
+          client_.get(), rpc_helper_, DType(DType::kF64), Shape({1, 2, 3}),
+          std::move(sharding), ArrayHandle{1234}, /*layout=*/nullptr)};
+      ASSERT_OK_AND_ASSIGN(DeviceListRef devices,
+                           client_->MakeDeviceList({device_}));
+
+      ASSERT_OK_AND_ASSIGN(
+          std::vector<ArrayRef> copied,
+          client_->CopyArrays(absl::MakeSpan(arrays), std::move(devices),
+                              dst.kind, ArrayCopySemantics::kAlwaysCopy));
+      ASSERT_THAT(copied, SizeIs(1));
+
+      EXPECT_EQ(copied[0]->sharding().memory_kind().value(),
+                dst.kind.has_value() ? dst.expected : src.expected);
+    }
+  }
+}
+
 TEST_P(ClientTest, CopyArraysCustomLayoutSuccess) {
   std::shared_ptr<xla::ifrt::SingleDeviceSharding> sharding =
       xla::ifrt::SingleDeviceSharding::Create(device_, xla::ifrt::MemoryKind());

From c0987b4961e6bf9917d3ef683a8d2ffeca6e3182 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Wed, 14 Oct 2026 13:08:03 -0700
Subject: [PATCH 0804/1000] [XLA:GPU] Fix
 GpuExecutableTest.InvokesDeviceErrorCallbackOnStreamError on ROCm.

The test hardcoded a CUDA compute capability but ran against the real test
platform. On ROCm, ExecuteThunks reads the ROCm compute capability from the
executable's device description, which was null, and crashed. Use the test
device's real device description instead.

PiperOrigin-RevId: 999567327
---
 third_party/xla/xla/service/gpu/BUILD         |  1 +
 .../xla/service/gpu/gpu_executable_test.cc    | 20 +++++++++++++------
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index a3b3fcb2148..d10ad3500ed 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -1002,6 +1002,7 @@ xla_test(
         "//xla/stream_executor/abi:executable_abi_version_proto_cc",
         "//xla/stream_executor/cuda:cuda_compute_capability",
         "//xla/stream_executor/gpu:tma_metadata",
+        "//xla/stream_executor/rocm:rocm_compute_capability",
         "//xla/tsl/lib/core:status_test_util",
         "//xla/tsl/platform:env",
         "//xla/tsl/testing:temporary_directory",
diff --git a/third_party/xla/xla/service/gpu/gpu_executable_test.cc b/third_party/xla/xla/service/gpu/gpu_executable_test.cc
index 6b09b7e75cd..5c0eb4923db 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable_test.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable_test.cc
@@ -90,6 +90,7 @@ limitations under the License.
 #include "xla/stream_executor/mock_stream_executor.h"
 #include "xla/stream_executor/module_spec.h"
 #include "xla/stream_executor/platform_manager.h"
+#include "xla/stream_executor/rocm/rocm_compute_capability.h"
 #include "xla/stream_executor/semantic_version.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/lib/core/status_test_util.h"
@@ -1383,17 +1384,24 @@ TEST_F(GpuExecutableTest, InvokesDeviceErrorCallbackOnStreamError) {
   EXPECT_CALL(executor, GetPlatform()).WillRepeatedly(Return(platform));
   EXPECT_CALL(executor, device_ordinal()).WillRepeatedly(Return(0));
   EXPECT_CALL(stream, parent()).WillRepeatedly(Return(&executor));
-  EXPECT_CALL(stream, GetCudaComputeCapability())
-      .WillRepeatedly(Return(se::CudaComputeCapability::Hopper()));
 
-  se::DeviceDescription device_description;
-  device_description.set_gpu_compute_capability(
-      se::GpuComputeCapability{se::CudaComputeCapability::Hopper()});
+  const se::GpuComputeCapability& compute_capability =
+      device_description().gpu_compute_capability();
+  if (const se::CudaComputeCapability* cuda_cc =
+          compute_capability.cuda_compute_capability()) {
+    EXPECT_CALL(stream, GetCudaComputeCapability())
+        .WillRepeatedly(Return(*cuda_cc));
+  }
+  if (const se::RocmComputeCapability* rocm_cc =
+          compute_capability.rocm_compute_capability()) {
+    EXPECT_CALL(stream, GetRocmComputeCapability())
+        .WillRepeatedly(Return(*rocm_cc));
+  }
 
   GpuExecutable::Params params;
   params.module_name = "test_module";
   params.executable = std::make_unique<ThunkExecutor>(ThunkSequence{});
-  params.device_description = device_description;
+  params.device_description = device_description();
   params.enable_debug_info_manager = false;
   SetDummyBufferAssignment(params);
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<GpuExecutable> executable,

From 2fc67f76b15c2e04dec63b3ea827673a8bf62a96 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2026 13:20:26 -0700
Subject: [PATCH 0805/1000] [XLA] Size tuple results in the HLO isolator's
 default module size estimate

PiperOrigin-RevId: 999575403
---
 .../tools/hlo_isolation/hlo_isolation_api.cc  |  6 ++--
 .../hlo_isolation/hlo_isolation_api_test.cc   | 33 +++++++++++++++++++
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api.cc b/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api.cc
index d4ecd13f056..4bc08fb1b25 100644
--- a/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api.cc
+++ b/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api.cc
@@ -138,10 +138,12 @@ absl::Status InitIsolatorOptions(ModuleIsolationOptions& options) {
       int64_t total_size = 0;
       for (const auto* param :
            module.entry_computation()->parameter_instructions()) {
-        total_size += ShapeUtil::ByteSizeOf(param->shape());
+        total_size += ShapeUtil::ByteSizeOf(param->shape(),
+                                            /*pointer_size=*/sizeof(void*));
       }
       total_size += ShapeUtil::ByteSizeOf(
-          module.entry_computation()->root_instruction()->shape());
+          module.entry_computation()->root_instruction()->shape(),
+          /*pointer_size=*/sizeof(void*));
       return total_size;
     };
   }
diff --git a/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api_test.cc b/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api_test.cc
index b1c61eee679..feced371f14 100644
--- a/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api_test.cc
+++ b/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api_test.cc
@@ -402,6 +402,39 @@ ENTRY main {
   EXPECT_DOUBLE_EQ(details[0].rel_error, 0.25);
 }
 
+TEST(HloIsolationApiTest, DefaultModuleSizeEstimateHandlesTupleResults) {
+  const absl::string_view kHlo = R"hlo(
+HloModule m
+ENTRY main {
+  p0 = f32[4] parameter(0)
+  ROOT t = (f32[4], f32[4]) tuple(p0, p0)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       xla::ParseAndReturnUnverifiedModule(kHlo));
+
+  ModuleIsolationOptions options;
+  options.max_module_size_bytes = 1;
+  options.run_module_fn =
+      [](std::unique_ptr<HloModule> m, HloRunnerInterface* r,
+         absl::Span<const Literal> i,
+         const RunModuleOptions& run_opts) -> absl::StatusOr<Literal> {
+    if (absl::StrContains(m->name(), "defused")) {
+      return absl::ResourceExhaustedError("Ran out of memory on HBM");
+    }
+    Literal element = LiteralUtil::CreateR1<float>({0.0f, 0.0f, 0.0f, 0.0f});
+    return LiteralUtil::MakeTuple({&element, &element});
+  };
+
+  // Estimating the size of a tuple result used to abort; the module is now
+  // only too large for the reference run.
+  ASSERT_OK_AND_ASSIGN(HloIsolationTestResult result,
+                       RunIsolationTestOnModule(
+                           *module, nullptr, FakeReferenceRunner(), options));
+  EXPECT_EQ(result.state(), State::SKIPPED);
+  EXPECT_EQ(result.reason(), "DEFUSED_REFERENCE_UNAVAILABLE");
+}
+
 }  // namespace
 }  // namespace hlo_isolation
 }  // namespace xla

From 1c65ffaae13f1a1bce4d8b1a5da1c11b42b290fd Mon Sep 17 00:00:00 2001
From: Oleg Shyshkov <shyshkov@google.com>
Date: Wed, 14 Oct 2026 13:42:28 -0700
Subject: [PATCH 0806/1000] [XLA:GPU] Move collective launch dimension
 calculations to collective_emitter.

PiperOrigin-RevId: 999590884
---
 .../gpu/codegen/triton/collective_emitter.cc  | 118 +++++++++++++++---
 .../xla/backends/gpu/runtime/all_gather.cc    |  28 -----
 .../xla/xla/backends/gpu/runtime/all_gather.h |   9 --
 .../xla/backends/gpu/runtime/all_reduce.cc    |  28 -----
 .../xla/xla/backends/gpu/runtime/all_reduce.h |   7 --
 .../backends/gpu/runtime/reduce_scatter.cc    |  22 ----
 .../xla/backends/gpu/runtime/reduce_scatter.h |   5 -
 7 files changed, 102 insertions(+), 115 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
index ad27ee30ae5..0d03e7bfe79 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
@@ -301,28 +301,114 @@ absl::StatusOr<ReduceScatterEmitterContext> CreateReduceScatterEmitterContext(
   return ctx;
 }
 
+// Returns the launch dimensions for the all-reduce kernel.
+// The launch dimensions are determined by the number of elements, the
+// all-reduce strategy, and the warp size of the target device.
+LaunchDimensions AllReduceLaunchDimensions(
+    const AllReduceInfo& all_reduce_info,
+    const se::DeviceDescription& device_info) {
+  // Heuristic maxima after some benchmarking.
+  // It's nicer to have this as a power of 2 because we consume blocks in powers
+  // of 2 from a higher dimension to a lower dimension.
+  static constexpr int64_t kMaxBlocksPerGrid = 32;
+  static constexpr uint64_t kMaxThreadsPerBlock = 512;
+
+  int64_t threads_per_block;
+  int64_t blocks_per_grid;
+  const int64_t warp_size = WarpSize(device_info);
+  const int64_t elements_per_rank =
+      all_reduce_info.num_elements /
+      (all_reduce_info.all_reduce_strategy == AllReduceStrategy::kTwoShot
+           ? all_reduce_info.num_devices
+           : 1);
+  // Maximum number of threads such that each thread has elements to process.
+  const int64_t total_threads =
+      RoundUpTo(CeilOfRatio(elements_per_rank, se::gpu::kNumElementsPerThread),
+                warp_size);
+  // Triton expects power of 2 for threads_per_block / threads_per_warp.
+  // Since threads_per_warp is always a power of 2, this is guaranteed.
+  threads_per_block =
+      std::min(kMaxThreadsPerBlock,
+               llvm::bit_ceil(static_cast<uint64_t>(total_threads)));
+  blocks_per_grid = std::min(kMaxBlocksPerGrid,
+                             CeilOfRatio(total_threads, threads_per_block));
+  return LaunchDimensions(blocks_per_grid, threads_per_block);
+}
+
+// Returns the launch dimensions for the all-gather kernel.
+// All-gather uses a one-shot strategy: each rank contributes its local slice
+// to the symmetric buffer and then reads each peer's slice.
+// warp_size should be device_description.threads_per_warp() (32 on NVIDIA,
+// 64 on AMD) so that the thread count is a multiple of the hardware warp.
+LaunchDimensions AllGatherLaunchDimensions(
+    const AllGatherInfo& all_gather_info,
+    const se::DeviceDescription& device_info) {
+  // All-gather always uses a one-shot strategy: every rank reads its peers'
+  // slices directly, so the work is proportional to the full element count
+  // with no rank-based partitioning.
+  static constexpr uint64_t kMaxThreadsPerBlock = 512;
+
+  // Maximum number of threads such that each thread has elements to process.
+  // Round up to a multiple of warp_size so every warp is fully occupied.
+  const int64_t output_elements =
+      all_gather_info.num_elements * all_gather_info.num_devices;
+  // Assuming each thread writes se::gpu::kNumElementsPerThread output elements
+  // we need total threads corresponding to the output elements.
+  const int64_t total_threads =
+      RoundUpTo(CeilOfRatio(output_elements, se::gpu::kNumElementsPerThread),
+                device_info.threads_per_warp());
+  // Triton expects power of 2 for threads_per_block / threads_per_warp.
+  const int64_t threads_per_block =
+      std::min(kMaxThreadsPerBlock,
+               llvm::bit_ceil(static_cast<uint64_t>(total_threads)));
+  const int64_t blocks_per_grid =
+      std::min(kAllGatherMaxBlocksPerGrid,
+               CeilOfRatio(total_threads, threads_per_block));
+  return LaunchDimensions(blocks_per_grid, threads_per_block);
+}
+
+// Returns the launch dimensions for the reduce-scatter kernel.
+LaunchDimensions ReduceScatterLaunchDimensions(
+    const ReduceScatterInfo& reduce_scatter_info,
+    const se::DeviceDescription& device_info) {
+  constexpr int64_t kNumElementsPerThread = 4;
+  constexpr uint64_t kMaxThreadsPerBlock = 512;
+
+  const int64_t warp_size = device_info.threads_per_warp();
+  const int64_t elements_per_rank_tile = CeilOfRatio(
+      reduce_scatter_info.num_output_elements, reduce_scatter_info.num_devices);
+  const int64_t threads_per_rank_tile = RoundUpTo(
+      CeilOfRatio(elements_per_rank_tile, kNumElementsPerThread), warp_size);
+  const int64_t threads_per_block = static_cast<int64_t>(
+      std::min(kMaxThreadsPerBlock,
+               llvm::bit_ceil(static_cast<uint64_t>(threads_per_rank_tile))));
+  const int64_t blocks_per_rank_tile =
+      CeilOfRatio(threads_per_rank_tile, threads_per_block);
+  const int64_t blocks_per_grid = std::min(
+      kReduceScatterMaxBlocksPerGrid,
+      reduce_scatter_info.num_devices *
+          static_cast<int64_t>(
+              llvm::bit_floor(static_cast<uint64_t>(blocks_per_rank_tile))));
+  return LaunchDimensions(blocks_per_grid, threads_per_block);
+}
+
 absl::StatusOr<LaunchDimensions> GetLaunchDimensions(
     const CollectiveKernelInfo& collective_info,
     const GpuTopology& gpu_topology) {
   const se::DeviceDescription& device_info =
       gpu_topology.gpu_target_config().device_description;
   return std::visit(
-      absl::Overload{[&](const AllReduceInfo& all_reduce_info) {
-                       return AllReduceLaunchDimensions(
-                           all_reduce_info.num_elements,
-                           all_reduce_info.num_devices,
-                           all_reduce_info.all_reduce_strategy, device_info);
-                     },
-                     [&](const AllGatherInfo& all_gather_info) {
-                       return AllGatherLaunchDimensions(
-                           all_gather_info.num_elements,
-                           all_gather_info.num_devices, device_info);
-                     },
-                     [&](const ReduceScatterInfo& reduce_scatter_info) {
-                       return ReduceScatterLaunchDimensions(
-                           reduce_scatter_info.num_output_elements,
-                           reduce_scatter_info.num_devices, device_info);
-                     }},
+      absl::Overload{
+          [&](const AllReduceInfo& all_reduce_info) {
+            return AllReduceLaunchDimensions(all_reduce_info, device_info);
+          },
+          [&](const AllGatherInfo& all_gather_info) {
+            return AllGatherLaunchDimensions(all_gather_info, device_info);
+          },
+          [&](const ReduceScatterInfo& reduce_scatter_info) {
+            return ReduceScatterLaunchDimensions(reduce_scatter_info,
+                                                 device_info);
+          }},
       collective_info);
 }
 
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_gather.cc b/third_party/xla/xla/backends/gpu/runtime/all_gather.cc
index 992d11d4b4d..aa69b32fb5a 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_gather.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/all_gather.cc
@@ -165,34 +165,6 @@ absl::StatusOr<AllGatherInfo> BuildAllGatherInfo(
   };
 }
 
-namespace {
-// All-gather always uses a one-shot strategy: every rank reads its peers'
-// slices directly, so the work is proportional to the full element count
-// with no rank-based partitioning.
-static constexpr uint64_t kMaxThreadsPerBlock = 512;
-}  // namespace
-
-LaunchDimensions AllGatherLaunchDimensions(
-    int64_t elements, int64_t num_ranks,
-    const se::DeviceDescription& device_info) {
-  // Maximum number of threads such that each thread has elements to process.
-  // Round up to a multiple of warp_size so every warp is fully occupied.
-  const int64_t output_elements = elements * num_ranks;
-  // Assuming each thread writes se::gpu::kNumElementsPerThread output elements
-  // we need total threads corresponding to the output elements.
-  const int64_t total_threads =
-      RoundUpTo(CeilOfRatio(output_elements, se::gpu::kNumElementsPerThread),
-                device_info.threads_per_warp());
-  // Triton expects power of 2 for threads_per_block / threads_per_warp.
-  const int64_t threads_per_block =
-      std::min(kMaxThreadsPerBlock,
-               llvm::bit_ceil(static_cast<uint64_t>(total_threads)));
-  const int64_t blocks_per_grid =
-      std::min(kAllGatherMaxBlocksPerGrid,
-               CeilOfRatio(total_threads, threads_per_block));
-  return LaunchDimensions(blocks_per_grid, threads_per_block);
-}
-
 absl::StatusOr<CollectiveKernelSpec> CreateAllGatherKernelSpec(
     const HloInstruction* instr, const LaunchDimensions& launch_dimensions,
     SymmetricMemoryType scratch_memory_type) {
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_gather.h b/third_party/xla/xla/backends/gpu/runtime/all_gather.h
index 4a5be79053e..ba542216b80 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_gather.h
+++ b/third_party/xla/xla/backends/gpu/runtime/all_gather.h
@@ -82,15 +82,6 @@ absl::StatusOr<AllGatherInfo> BuildAllGatherInfo(
     const GpuTopology& gpu_topology, const HloAllGatherInstruction* all_gather,
     const DeviceAssignment* device_assignment);
 
-// Returns the launch dimensions for the all-gather kernel.
-// All-gather uses a one-shot strategy: each rank contributes its local slice
-// to the symmetric buffer and then reads each peer's slice.
-// warp_size should be device_description.threads_per_warp() (32 on NVIDIA,
-// 64 on AMD) so that the thread count is a multiple of the hardware warp.
-LaunchDimensions AllGatherLaunchDimensions(
-    int64_t elements, int64_t num_ranks,
-    const se::DeviceDescription& device_info);
-
 // Creates a CollectiveKernelSpec describing the resource requirements of a
 // Triton all-gather kernel.  The kernel argument layout matches all-reduce:
 //   [0] local input buffer (kInputBuffer, index 0)
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc b/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc
index 9d614471938..be2ec560a24 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc
@@ -62,12 +62,6 @@ static constexpr int64_t kMB = kKB * 1024;
 static constexpr int64_t kMaxOneShotAllReduceSizeBytes = 2 * kMB;
 static constexpr int64_t kMaxTwoShotAllReduceSizeBytes = 32 * kMB;
 
-// Heuristic maxima after some benchmarking.
-// It's nicer to have this as a power of 2 because we consume blocks in powers
-// of 2 from a higher dimension to a lower dimension.
-static constexpr int64_t kMaxBlocksPerGrid = 32;
-static constexpr uint64_t kMaxThreadsPerBlock = 512;
-
 bool IsElementReductionSupported(PrimitiveType element_type,
                                  ReductionKind reduction_kind) {
   absl::Span<const HloOpcode> supported_ops =
@@ -114,28 +108,6 @@ int64_t GetMaxSupportedAllReduceSizeBytes(AllReduceStrategy strategy) {
   }
 }
 
-LaunchDimensions AllReduceLaunchDimensions(
-    int64_t elements, int64_t num_ranks, AllReduceStrategy strategy,
-    const se::DeviceDescription& device_info) {
-  int64_t threads_per_block;
-  int64_t blocks_per_grid;
-  const int64_t warp_size = WarpSize(device_info);
-  const int64_t elements_per_rank =
-      elements / (strategy == AllReduceStrategy::kTwoShot ? num_ranks : 1);
-  // Maximum number of threads such that each thread has elements to process.
-  const int64_t total_threads =
-      RoundUpTo(CeilOfRatio(elements_per_rank, se::gpu::kNumElementsPerThread),
-                warp_size);
-  // Triton expects power of 2 for threads_per_block / threads_per_warp.
-  // Since threads_per_warp is always a power of 2, this is guaranteed.
-  threads_per_block =
-      std::min(kMaxThreadsPerBlock,
-               llvm::bit_ceil(static_cast<uint64_t>(total_threads)));
-  blocks_per_grid = std::min(kMaxBlocksPerGrid,
-                             CeilOfRatio(total_threads, threads_per_block));
-  return LaunchDimensions(blocks_per_grid, threads_per_block);
-}
-
 absl::Status IsAllReduceKernelSupported(int64_t num_ranks, int64_t num_elements,
                                         PrimitiveType element_type,
                                         ReductionKind reduction_kind,
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_reduce.h b/third_party/xla/xla/backends/gpu/runtime/all_reduce.h
index 77a99efe838..db16dd4cd03 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_reduce.h
+++ b/third_party/xla/xla/backends/gpu/runtime/all_reduce.h
@@ -83,13 +83,6 @@ se::gpu::AllReduceStrategy GetAllReduceStrategy(int64_t input_size_bytes,
 // strategy.
 int64_t GetMaxSupportedAllReduceSizeBytes(se::gpu::AllReduceStrategy strategy);
 
-// Returns the launch dimensions for the all-reduce kernel.
-// The launch dimensions are determined by the number of elements, the
-// all-reduce strategy, and the warp size of the target device.
-LaunchDimensions AllReduceLaunchDimensions(
-    int64_t elements, int64_t num_ranks, se::gpu::AllReduceStrategy strategy,
-    const se::DeviceDescription& device_info);
-
 // Creates a CollectiveKernelSpec for a given all-reduce instruction.
 // `scratch_memory_type` is used for both scratch buffers.
 absl::StatusOr<CollectiveKernelSpec> CreateAllReduceKernelSpec(
diff --git a/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.cc b/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.cc
index cc912ac42dd..7e555d76317 100644
--- a/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.cc
@@ -170,28 +170,6 @@ absl::StatusOr<ReduceScatterInfo> BuildReduceScatterInfo(
   };
 }
 
-LaunchDimensions ReduceScatterLaunchDimensions(
-    int64_t num_output_elements, int64_t num_ranks,
-    const se::DeviceDescription& device_info) {
-  constexpr int64_t kNumElementsPerThread = 4;
-  constexpr uint64_t kMaxThreadsPerBlock = 512;
-  const int64_t warp_size = device_info.threads_per_warp();
-  const int64_t elements_per_rank_tile =
-      CeilOfRatio(num_output_elements, num_ranks);
-  const int64_t threads_per_rank_tile = RoundUpTo(
-      CeilOfRatio(elements_per_rank_tile, kNumElementsPerThread), warp_size);
-  const int64_t threads_per_block = static_cast<int64_t>(
-      std::min(kMaxThreadsPerBlock,
-               llvm::bit_ceil(static_cast<uint64_t>(threads_per_rank_tile))));
-  const int64_t blocks_per_rank_tile =
-      CeilOfRatio(threads_per_rank_tile, threads_per_block);
-  const int64_t blocks_per_grid =
-      std::min(kReduceScatterMaxBlocksPerGrid,
-               num_ranks * static_cast<int64_t>(llvm::bit_floor(
-                               static_cast<uint64_t>(blocks_per_rank_tile))));
-  return LaunchDimensions(blocks_per_grid, threads_per_block);
-}
-
 absl::StatusOr<CollectiveKernelSpec> CreateReduceScatterKernelSpec(
     const HloInstruction* instr, const LaunchDimensions& launch_dimensions,
     SymmetricMemoryType scratch_memory_type) {
diff --git a/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.h b/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.h
index ac74dd006e8..002e93121a8 100644
--- a/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.h
+++ b/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.h
@@ -61,11 +61,6 @@ absl::StatusOr<ReduceScatterInfo> BuildReduceScatterInfo(
     const HloReduceScatterInstruction* reduce_scatter,
     const DeviceAssignment* device_assignment);
 
-// Returns the launch dimensions for the reduce-scatter kernel.
-LaunchDimensions ReduceScatterLaunchDimensions(
-    int64_t num_output_elements, int64_t num_ranks,
-    const se::DeviceDescription& device_info);
-
 // Creates a CollectiveKernelSpec describing the resource requirements of a
 // Triton one-shot reduce-scatter kernel. The kernel argument layout is:
 //   [0] input buffer  (kInputBuffer, index 0)

From 27f95aeefdbd2f1bf609cd098063a3f80c4b1690 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2026 13:58:17 -0700
Subject: [PATCH 0807/1000] Add direct lowering support for linear, gather, and
 quantization ops.

PiperOrigin-RevId: 999602043
---
 .../transforms/legalize_hlo_conversions/BUILD |   4 +
 .../legalize_hlo_conversions/custom_call.cc   | 331 ++++++++++++++++++
 2 files changed, 335 insertions(+)

diff --git a/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/BUILD b/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/BUILD
index 28f984492fe..a31ce430c84 100644
--- a/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/BUILD
+++ b/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/BUILD
@@ -111,8 +111,12 @@ cc_library(
     ],
     deps = [
         "//tensorflow/compiler/mlir/lite:tensorflow_lite",
+        "//tensorflow/compiler/mlir/lite:tensorflow_lite_quantize",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:AsmParser",
         "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:QuantOps",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TransformUtils",
         "@xla//xla/mlir_hlo",
diff --git a/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/custom_call.cc b/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/custom_call.cc
index 18d9b10d677..0d57945cc47 100644
--- a/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/custom_call.cc
+++ b/tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/custom_call.cc
@@ -15,16 +15,29 @@ limitations under the License.
 
 #include "tensorflow/compiler/mlir/lite/stablehlo/transforms/legalize_hlo_conversions/custom_call.h"
 
+#include <algorithm>
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
 #include <optional>
 
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
 #include "mlir/AsmParser/AsmParser.h"  // from @llvm-project
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/Quant.h"  // from @llvm-project
+#include "mlir/Dialect/Quant/IR/QuantTypes.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
 #include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
 #include "mlir/IR/OperationSupport.h"  // from @llvm-project
 #include "mlir/IR/PatternMatch.h"  // from @llvm-project
 #include "mlir/Support/LLVM.h"  // from @llvm-project
 #include "mlir/Support/LogicalResult.h"  // from @llvm-project
 #include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
 #include "tensorflow/compiler/mlir/lite/ir/tfl_ops.h"  // IWYU pragma: keep
+#include "tensorflow/compiler/mlir/lite/transforms/lower_quant_annotations_helper.h"
 #include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
 
 namespace mlir {
@@ -61,10 +74,242 @@ mlir::DictionaryAttr ParseSerializedTFLOpAttributes(
   return dict_attribute;
 }
 
+// Quantization parameters carried in the backend_config of a `tfl.quantize` /
+// `tfl.dequantize` custom call.
+struct QuantParams {
+  SmallVector<double, 4> scales;
+  SmallVector<int64_t, 4> zero_points;
+  int num_bits = 8;
+  bool is_signed = true;
+  bool narrow_range = false;
+  int32_t quantized_dimension = 0;
+};
+
+QuantParams ParseQuantParams(DictionaryAttr attributes) {
+  QuantParams params;
+
+  if (auto scale_attr = attributes.getAs<mlir::FloatAttr>("scale")) {
+    params.scales.push_back(scale_attr.getValueAsDouble());
+  } else if (auto scale_array = attributes.getAs<mlir::ArrayAttr>("scale")) {
+    for (auto elem : scale_array) {
+      if (auto f_attr = mlir::dyn_cast<mlir::FloatAttr>(elem)) {
+        params.scales.push_back(f_attr.getValueAsDouble());
+      } else if (auto i_attr = mlir::dyn_cast<mlir::IntegerAttr>(elem)) {
+        params.scales.push_back(static_cast<double>(i_attr.getInt()));
+      }
+    }
+  } else if (auto scale_dense =
+                 attributes.getAs<mlir::DenseFPElementsAttr>("scale")) {
+    for (auto val : scale_dense.getValues<APFloat>()) {
+      params.scales.push_back(val.convertToDouble());
+    }
+  } else if (auto scale_int = attributes.getAs<mlir::IntegerAttr>("scale")) {
+    params.scales.push_back(static_cast<double>(scale_int.getInt()));
+  }
+
+  if (auto zp_attr = attributes.getAs<mlir::IntegerAttr>("zero_point")) {
+    params.zero_points.push_back(zp_attr.getInt());
+  } else if (auto zp_array = attributes.getAs<mlir::ArrayAttr>("zero_point")) {
+    for (auto elem : zp_array) {
+      if (auto i_attr = mlir::dyn_cast<mlir::IntegerAttr>(elem)) {
+        params.zero_points.push_back(i_attr.getInt());
+      }
+    }
+  } else if (auto zp_dense =
+                 attributes.getAs<mlir::DenseIntElementsAttr>("zero_point")) {
+    for (auto val : zp_dense.getValues<APInt>()) {
+      params.zero_points.push_back(val.getSExtValue());
+    }
+  }
+
+  if (auto num_bits_attr = attributes.getAs<mlir::IntegerAttr>("num_bits")) {
+    params.num_bits = num_bits_attr.getInt();
+  }
+  if (auto is_signed_attr = attributes.getAs<mlir::BoolAttr>("is_signed")) {
+    params.is_signed = is_signed_attr.getValue();
+  }
+  if (auto narrow_range_attr =
+          attributes.getAs<mlir::BoolAttr>("narrow_range")) {
+    params.narrow_range = narrow_range_attr.getValue();
+  }
+  if (auto dim_attr =
+          attributes.getAs<mlir::IntegerAttr>("quantized_dimension")) {
+    params.quantized_dimension = dim_attr.getInt();
+  }
+
+  if (params.scales.empty()) {
+    params.scales.push_back(1.0);
+  }
+  if (params.zero_points.empty()) {
+    params.zero_points.push_back(0);
+  }
+  return params;
+}
+
+// Builds the quantized tensor type of `shape` described by `params`. Returns
+// null if `params.quantized_dimension` is out of range for `shape`.
+RankedTensorType BuildQuantizedTensorType(Builder& builder,
+                                          const QuantParams& params,
+                                          ArrayRef<int64_t> shape,
+                                          Type expressed_type, Location loc) {
+  Type quantized_element_type;
+  if (params.scales.size() == 1) {
+    quantized_element_type = TFL::GetPerTensorQuantizedTensorType(
+        builder, params.scales[0], params.zero_points[0], expressed_type,
+        params.num_bits, loc, params.narrow_range, params.is_signed);
+  } else {
+    const int64_t rank = static_cast<int64_t>(shape.size());
+    int64_t quant_dim = params.quantized_dimension;
+    if (quant_dim < 0) {
+      quant_dim += rank;
+    }
+    if (quant_dim < 0 || quant_dim >= rank) {
+      return nullptr;
+    }
+    quantized_element_type = TFL::GetPerAxisQuantizedTensorType(
+        builder, params.scales, params.zero_points,
+        static_cast<int32_t>(quant_dim), expressed_type, params.num_bits, loc,
+        params.narrow_range, params.is_signed);
+  }
+  if (!quantized_element_type) {
+    return nullptr;
+  }
+  return RankedTensorType::get(shape, quantized_element_type);
+}
+
+bool ApproxEqual(double a, double b) {
+  return std::abs(a - b) <= 1e-6 * std::max(std::abs(a), std::abs(b));
+}
+
+// Returns true if `qtype` has the scales and zero points in `params`.
+bool MatchesScalesAndZeroPoints(quant::QuantizedType qtype,
+                                const QuantParams& params) {
+  if (auto uniform = mlir::dyn_cast<quant::UniformQuantizedType>(qtype)) {
+    return params.scales.size() == 1 &&
+           ApproxEqual(uniform.getScale(), params.scales[0]) &&
+           uniform.getZeroPoint() == params.zero_points[0];
+  }
+  if (auto per_axis =
+          mlir::dyn_cast<quant::UniformQuantizedPerAxisType>(qtype)) {
+    ArrayRef<double> scales = per_axis.getScales();
+    ArrayRef<int64_t> zero_points = per_axis.getZeroPoints();
+    if (scales.size() != params.scales.size() ||
+        zero_points.size() != params.zero_points.size()) {
+      return false;
+    }
+    for (size_t i = 0; i < scales.size(); ++i) {
+      if (!ApproxEqual(scales[i], params.scales[i]) ||
+          zero_points[i] != params.zero_points[i]) {
+        return false;
+      }
+    }
+    return true;
+  }
+  return false;
+}
+
+// Rewrites a `tfl.quantize` custom call on `input` into `tfl.quantize`.
+LogicalResult RewriteQuantizeCustomCall(mhlo::CustomCallOp op, Value input,
+                                        PatternRewriter& rewriter) {
+  DictionaryAttr attributes =
+      ParseSerializedTFLOpAttributes(op.getBackendConfig(), op.getContext());
+  if (!attributes) {
+    return failure();
+  }
+  QuantParams params = ParseQuantParams(attributes);
+
+  ShapedType input_shaped_type = mlir::cast<ShapedType>(input.getType());
+  RankedTensorType output_type =
+      BuildQuantizedTensorType(rewriter, params, input_shaped_type.getShape(),
+                               input_shaped_type.getElementType(), op.getLoc());
+  if (!output_type) {
+    return op.emitError("invalid quantization parameters for tfl.quantize");
+  }
+  auto quant_op = rewriter.create<TFL::QuantizeOp>(
+      op.getLoc(), output_type, input, TypeAttr::get(output_type));
+  rewriter.replaceOp(op, quant_op.getOutput());
+  return success();
+}
+
+// Rewrites a `tfl.dequantize` custom call on `input` into `tfl.dequantize`.
+// - An already-quantized input is dequantized with its own type, which must
+//   match the scales and zero points in the backend_config.
+// - An integer constant input becomes a `tfl.pseudo_qconst` with the quantized
+//   type described by the backend_config.
+// Any other input (e.g. a runtime integer tensor) can't be given a quantized
+// type and is rejected.
+LogicalResult RewriteDequantizeCustomCall(mhlo::CustomCallOp op, Value input,
+                                          PatternRewriter& rewriter) {
+  DictionaryAttr attributes =
+      ParseSerializedTFLOpAttributes(op.getBackendConfig(), op.getContext());
+  if (!attributes) {
+    return failure();
+  }
+
+  if (auto call_op = input.getDefiningOp<mhlo::CustomCallOp>()) {
+    if (call_op.getCallTargetName() == "tfl.quantize") {
+      return failure();
+    }
+  }
+  ShapedType input_shaped_type = mlir::cast<ShapedType>(input.getType());
+  Type output_type = op.getResultTypes().front();
+  Type expressed_type = mlir::cast<ShapedType>(output_type).getElementType();
+  QuantParams params = ParseQuantParams(attributes);
+
+  Value tfl_dequant_input = input;
+  if (auto qtype = mlir::dyn_cast<quant::QuantizedType>(
+          input_shaped_type.getElementType())) {
+    if (!MatchesScalesAndZeroPoints(qtype, params)) {
+      return op.emitError(
+          "tfl.dequantize scale/zero_point don't match the quantized type of "
+          "its input");
+    }
+  } else {
+    ElementsAttr const_value;
+    Location const_loc = op.getLoc();
+    if (auto const_op = input.getDefiningOp<arith::ConstantOp>()) {
+      const_value = mlir::dyn_cast<ElementsAttr>(const_op.getValue());
+      const_loc = const_op.getLoc();
+    } else if (auto const_op = input.getDefiningOp<mhlo::ConstantOp>()) {
+      const_value = mlir::dyn_cast<ElementsAttr>(const_op.getValue());
+      const_loc = const_op.getLoc();
+    }
+    if (!const_value ||
+        !mlir::isa<mlir::IntegerType>(const_value.getElementType())) {
+      return rewriter.notifyMatchFailure(
+          op, "tfl.dequantize input must be quantized or an integer constant");
+    }
+    RankedTensorType qtensor_type =
+        BuildQuantizedTensorType(rewriter, params, input_shaped_type.getShape(),
+                                 expressed_type, op.getLoc());
+    if (!qtensor_type) {
+      return op.emitError("invalid quantization parameters for tfl.dequantize");
+    }
+    auto qconst_op = rewriter.create<TFL::QConstOp>(
+        const_loc, qtensor_type, TypeAttr::get(qtensor_type), const_value);
+    tfl_dequant_input = qconst_op.getResult();
+  }
+
+  auto dequant_op = rewriter.create<TFL::DequantizeOp>(op.getLoc(), output_type,
+                                                       tfl_dequant_input);
+  rewriter.replaceOp(op, dequant_op.getOutput());
+  return success();
+}
+
 LogicalResult ConvertCustomCallOp::matchAndRewrite(
     mhlo::CustomCallOp mhlo_custom_call, OpAdaptor adaptor,
     ConversionPatternRewriter& rewriter) const {
   auto call_target_name = mhlo_custom_call.getCallTargetName();
+
+  if (call_target_name == "tfl.quantize") {
+    return RewriteQuantizeCustomCall(mhlo_custom_call,
+                                     adaptor.getOperands().front(), rewriter);
+  }
+  if (call_target_name == "tfl.dequantize") {
+    return RewriteDequantizeCustomCall(mhlo_custom_call,
+                                       adaptor.getOperands().front(), rewriter);
+  }
+
   if (call_target_name.starts_with("tfl.")) {
     auto bc = mhlo_custom_call.getBackendConfig();
     if (mlir::DictionaryAttr attributes =
@@ -101,6 +346,34 @@ LogicalResult ConvertCustomCallOp::matchAndRewrite(
   return success();
 }
 
+class RewriteQuantizeCustomCallOp
+    : public OpRewritePattern<mhlo::CustomCallOp> {
+ public:
+  using OpRewritePattern<mhlo::CustomCallOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mhlo::CustomCallOp op,
+                                PatternRewriter& rewriter) const override {
+    if (op.getCallTargetName() != "tfl.quantize") {
+      return failure();
+    }
+    return RewriteQuantizeCustomCall(op, op.getOperand(0), rewriter);
+  }
+};
+
+class RewriteDequantizeCustomCallOp
+    : public OpRewritePattern<mhlo::CustomCallOp> {
+ public:
+  using OpRewritePattern<mhlo::CustomCallOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mhlo::CustomCallOp op,
+                                PatternRewriter& rewriter) const override {
+    if (op.getCallTargetName() != "tfl.dequantize") {
+      return failure();
+    }
+    return RewriteDequantizeCustomCall(op, op.getOperand(0), rewriter);
+  }
+};
+
 // Removes the `mhlo.custom_call @shape_assertion` custom call which represents
 // an assertion that the first operand (`assert_what`) evaluates to `true`.
 // This is a temporary workaround for unblocking dynamic model conversion
@@ -132,6 +405,61 @@ LogicalResult RemoveCustomCallWithShapeAssertion::matchAndRewrite(
   return success();
 }
 
+// TFL ops whose quantized output uses the same scale and zero point as their
+// first operand: they only move, select or replicate elements.
+constexpr llvm::StringLiteral kScalePreservingTflOps[] = {
+    "tfl.broadcast_to", "tfl.depth_to_space", "tfl.expand_dims",
+    "tfl.gather",       "tfl.gather_nd",      "tfl.reshape",
+    "tfl.reverse_v2",   "tfl.slice",          "tfl.space_to_depth",
+    "tfl.squeeze",      "tfl.strided_slice",  "tfl.tile",
+    "tfl.transpose",
+};
+
+// Gives a scale-preserving `tfl.*` custom call the quantized type of its first
+// operand, so that e.g. `reshape(quantize(x))` stays quantized. Ops that
+// compute new values (FC, BMM, add, ...) need their own output scale and are
+// deliberately not handled here.
+class PropagateQuantizedTypeToCustomCallOp
+    : public OpRewritePattern<mhlo::CustomCallOp> {
+ public:
+  using OpRewritePattern<mhlo::CustomCallOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mhlo::CustomCallOp op,
+                                PatternRewriter& rewriter) const override {
+    if (!llvm::is_contained(kScalePreservingTflOps, op.getCallTargetName())) {
+      return failure();
+    }
+    if (op.getNumOperands() == 0 || op.getNumResults() != 1) {
+      return failure();
+    }
+    auto operand_type = mlir::dyn_cast<ShapedType>(op.getOperand(0).getType());
+    if (!operand_type) {
+      return failure();
+    }
+    auto qtype =
+        mlir::dyn_cast<quant::QuantizedType>(operand_type.getElementType());
+    if (!qtype) {
+      return failure();
+    }
+    auto res_shaped_type =
+        mlir::dyn_cast<ShapedType>(op.getResult(0).getType());
+    if (!res_shaped_type) {
+      return failure();
+    }
+    // Only retype results that already hold the storage type, e.g. i8 -> !quant
+    // with i8 storage.
+    auto res_int_type =
+        mlir::dyn_cast<mlir::IntegerType>(res_shaped_type.getElementType());
+    if (!res_int_type ||
+        res_int_type.getWidth() != qtype.getStorageTypeIntegralWidth()) {
+      return failure();
+    }
+    rewriter.modifyOpInPlace(
+        op, [&] { op.getResult(0).setType(res_shaped_type.clone(qtype)); });
+    return success();
+  }
+};
+
 std::optional<bool> IsCustomCallLegal(mhlo::CustomCallOp op) {
   auto call_target_name = op.getCallTargetName();
   if (call_target_name.starts_with("custom_call.")) {
@@ -160,6 +488,9 @@ void PopulateCustomCallPatterns(MLIRContext* ctx, RewritePatternSet& patterns,
 void PopulateCustomCallPreparePatterns(MLIRContext* ctx,
                                        RewritePatternSet& patterns) {
   patterns.add<RemoveCustomCallWithShapeAssertion>(ctx);
+  patterns.add<RewriteQuantizeCustomCallOp>(ctx);
+  patterns.add<RewriteDequantizeCustomCallOp>(ctx);
+  patterns.add<PropagateQuantizedTypeToCustomCallOp>(ctx);
 }
 
 }  // namespace odml

From 91e6612c504f416fd013b9a9cb8d6c4ef3cc8d93 Mon Sep 17 00:00:00 2001
From: Dmitri Latushko <latushko@google.com>
Date: Wed, 14 Oct 2026 14:16:01 -0700
Subject: [PATCH 0808/1000] Subtract squared residual mean before reduction in
 tf.nn.moments

PiperOrigin-RevId: 999614913
---
 tensorflow/python/ops/nn_impl.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/tensorflow/python/ops/nn_impl.py b/tensorflow/python/ops/nn_impl.py
index ed578011628..c18000798a1 100644
--- a/tensorflow/python/ops/nn_impl.py
+++ b/tensorflow/python/ops/nn_impl.py
@@ -1459,9 +1459,17 @@ def moments(
     #       because that gradient is zero
     diff = math_ops.subtract(y, array_ops.stop_gradient(mean))
     variance = math_ops.maximum(
-        math_ops.subtract(
-            math_ops.reduce_mean(math_ops.square(diff), axes, keepdims=True),
-            math_ops.square(math_ops.reduce_mean(diff, axes, keepdims=True)),
+        math_ops.reduce_mean(
+            math_ops.subtract(
+                math_ops.square(diff),
+                math_ops.square(
+                    array_ops.stop_gradient(
+                        math_ops.reduce_mean(diff, axes, keepdims=True)
+                    )
+                ),
+            ),
+            axes,
+            keepdims=True,
         ),
         math_ops.cast(0, y.dtype),
         name="variance",

From 4f8f953971d929b81be90cda581e8f6bf36ba0c7 Mon Sep 17 00:00:00 2001
From: Dmitri Latushko <latushko@google.com>
Date: Wed, 14 Oct 2026 15:03:31 -0700
Subject: [PATCH 0809/1000] Shard bias_op_test across 4 shards to prevent TSAN
 timeouts.

PiperOrigin-RevId: 999646977
---
 tensorflow/python/kernel_tests/nn_ops/BUILD | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/python/kernel_tests/nn_ops/BUILD b/tensorflow/python/kernel_tests/nn_ops/BUILD
index 51d436d3531..c19ed445842 100644
--- a/tensorflow/python/kernel_tests/nn_ops/BUILD
+++ b/tensorflow/python/kernel_tests/nn_ops/BUILD
@@ -160,6 +160,7 @@ cuda_py_strict_test(
     name = "bias_op_test",
     size = "medium",
     srcs = ["bias_op_test.py"],
+    shard_count = 4,
     deps = [
         ":bias_op_base",
         "//tensorflow/python/platform:client_testlib",

From 4db83379a2bd03a977e09fdae9e42af23617201d Mon Sep 17 00:00:00 2001
From: Aliia Khasanova <aliia@google.com>
Date: Wed, 14 Oct 2026 15:06:08 -0700
Subject: [PATCH 0810/1000] Add v4 GPU AOT goldens for the collective ops tests

Dump after changing the comparison rules (and it's been broken at head)

PiperOrigin-RevId: 999648866
---
 .../v4/AllGatherMixedTypes.pbtxt              | 2467 +++++++++++++++++
 .../v4/ReduceScatter.pbtxt                    | 1901 +++++++++++++
 .../v4/AllGatherMixedTypes.pbtxt              | 2467 +++++++++++++++++
 .../v4/ReduceScatter.pbtxt                    | 1901 +++++++++++++
 .../v4/AllGatherMixedTypes.pbtxt              | 2459 ++++++++++++++++
 .../v4/ReduceScatter.pbtxt                    | 1893 +++++++++++++
 6 files changed, 13088 insertions(+)
 create mode 100644 third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_b200/v4/AllGatherMixedTypes.pbtxt
 create mode 100644 third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_b200/v4/ReduceScatter.pbtxt
 create mode 100644 third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_gb200/v4/AllGatherMixedTypes.pbtxt
 create mode 100644 third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_gb200/v4/ReduceScatter.pbtxt
 create mode 100644 third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_h100/v4/AllGatherMixedTypes.pbtxt
 create mode 100644 third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_h100/v4/ReduceScatter.pbtxt

diff --git a/third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_b200/v4/AllGatherMixedTypes.pbtxt b/third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_b200/v4/AllGatherMixedTypes.pbtxt
new file mode 100644
index 00000000000..b5875b34730
--- /dev/null
+++ b/third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_b200/v4/AllGatherMixedTypes.pbtxt
@@ -0,0 +1,2467 @@
+# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto
+# proto-message: xla.HumanReadableAotExecutable
+executable_and_options {
+  compile_options {
+    executable_build_options {
+      device_ordinal: -1
+      result_layout {
+        element_type: TUPLE
+        tuple_shapes {
+          element_type: U32
+          dimensions: 4
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+        tuple_shapes {
+          element_type: F32
+          dimensions: 4
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+      }
+      debug_options {
+        xla_disable_hlo_passes: "constant_folding"
+        xla_backend_optimization_level: 3
+        xla_eliminate_hlo_implicit_broadcast: true
+        xla_cpu_multi_thread_eigen: true
+        xla_gpu_cuda_data_dir: "./cuda_sdk_lib"
+        xla_llvm_enable_alias_scope_metadata: true
+        xla_llvm_enable_noalias_metadata: true
+        xla_llvm_enable_invariant_load_metadata: true
+        xla_llvm_disable_expensive_passes: false
+        xla_cpu_use_onednn: false
+        xla_cpu_enable_fast_math: false
+        xla_gpu_enable_fast_min_max: false
+        xla_force_host_platform_device_count: 1
+        xla_gpu_enable_dynamic_slice_fusion: true
+        xla_hlo_evaluator_use_fast_path: true
+        xla_dump_hlo_as_html: false
+        xla_cpu_fast_math_honor_nans: true
+        xla_cpu_fast_math_honor_infs: true
+        xla_allow_excess_precision: true
+        xla_gpu_autotune_level: 4
+        xla_cpu_fast_math_honor_division: true
+        xla_cpu_fast_math_honor_functions: true
+        xla_dump_include_timestamp: false
+        xla_dump_max_hlo_modules: -1
+        xla_cpu_enable_xprof_traceme: false
+        xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false
+        xla_cpu_enable_fast_min_max: true
+        xla_multiheap_size_constraint_per_heap: -1
+        xla_dump_module_metadata: false
+        xla_gpu_force_compilation_parallelism: 1
+        xla_dump_fusion_visualization: false
+        xla_gpu_all_reduce_combine_threshold_bytes: 31457287
+        xla_gpu_nccl_termination_timeout_seconds: -1
+        xla_dump_hlo_as_long_text: true
+        xla_gpu_enable_cublaslt: true
+        xla_gpu_shape_checks: RUNTIME
+        xla_gpu_use_runtime_fusion: false
+        xla_dump_latency_hiding_schedule: false
+        xla_dump_enable_mlir_pretty_form: true
+        xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP
+        xla_gpu_enable_triton_gemm: true
+        xla_gpu_enable_cudnn_int8x32_convolution_reordering: true
+        xla_gpu_triton_gemm_any: true
+        xla_gpu_enable_while_loop_reduce_scatter_code_motion: false
+        xla_gpu_collective_inflation_factor: 1
+        xla_gpu_graph_min_graph_size: 5
+        xla_gpu_enable_reassociation_for_converted_ar: true
+        xla_gpu_pgle_profile_file_or_directory_path: ""
+        xla_gpu_all_gather_combine_threshold_bytes: 31457287
+        xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287
+        xla_gpu_enable_highest_priority_async_stream: true
+        xla_gpu_exhaustive_tiling_search: false
+        xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0
+        xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1
+        xla_gpu_redzone_padding_bytes: 8388608
+        xla_gpu_copy_insertion_use_region_analysis: false
+        xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807
+        xla_gpu_collect_cost_model_stats: false
+        xla_gpu_enable_pipelined_p2p: false
+        xla_gpu_cublas_fallback: true
+        xla_gpu_enable_while_loop_double_buffering: false
+        xla_gpu_filter_kernels_spilling_registers_on_autotuning: true
+        xla_debug_buffer_assignment_show_max: 15
+        xla_detailed_logging: true
+        xla_enable_dumping: true
+        xla_gpu_enable_all_gather_combine_by_dim: false
+        xla_gpu_enable_analytical_latency_estimator: false
+        xla_gpu_llvm_verification_level: 0
+        xla_gpu_enable_reduce_scatter_combine_by_dim: false
+        xla_gpu_enable_command_buffer: COLLECTIVES_KERNEL
+        xla_gpu_enable_command_buffer: CONDITIONAL
+        xla_gpu_enable_command_buffer: CUBLAS
+        xla_gpu_enable_command_buffer: CUBLASLT
+        xla_gpu_enable_command_buffer: CUDNN
+        xla_gpu_enable_command_buffer: CUSTOM_CALL
+        xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION
+        xla_gpu_enable_command_buffer: FUSION
+        xla_gpu_enable_cub_radix_sort: true
+        xla_gpu_memory_limit_slop_factor: 95
+        xla_gpu_target_config_filename: ""
+        xla_gpu_enable_cudnn_layer_norm: false
+        xla_gpu_threshold_for_windowed_einsum_mib: 100000
+        xla_gpu_enable_nccl_user_buffers: false
+        xla_gpu_enable_libnvptxcompiler: true
+        xla_gpu_enable_nccl_comm_splitting: true
+        xla_gpu_nccl_collective_max_nchannels: 0
+        xla_gpu_nccl_p2p_max_nchannels: 0
+        xla_gpu_nccl_init_max_rank_per_root_ratio: 0
+        xla_gpu_multi_streamed_windowed_einsum: true
+        xla_gpu_gemm_rewrite_size_threshold: 100
+        xla_gpu_require_complete_aot_autotune_results: false
+        xla_gpu_cudnn_gemm_fusion_level: 0
+        xla_gpu_use_memcpy_local_p2p: false
+        xla_gpu_autotune_max_solutions: 0
+        xla_dump_large_constants: false
+        xla_gpu_verify_triton_fusion_numerics: false
+        xla_reduce_window_rewrite_base_length: 16
+        xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL
+        xla_gpu_enable_host_memory_offloading: false
+        xla_gpu_exclude_nondeterministic_ops: true
+        xla_gpu_nccl_terminate_on_error: false
+        xla_gpu_shard_autotuning: true
+        xla_gpu_enable_approx_costly_collectives: false
+        xla_cpu_prefer_vector_width: 256
+        xla_gpu_per_fusion_autotune_cache_dir: ""
+        xla_cmd_buffer_trace_cache_size: 16
+        xla_gpu_temp_buffer_use_separate_color: false
+        xla_syntax_sugar_async_ops: false
+        xla_gpu_autotune_gemm_rtol: 0.1
+        xla_enable_command_buffers_during_profiling: false
+        xla_gpu_cudnn_gemm_max_plans: 5
+        xla_cpu_parallel_codegen_split_count: 32
+        xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE
+        xla_gpu_executable_warn_stuck_timeout_seconds: 10
+        xla_gpu_executable_terminate_timeout_seconds: 30
+        xla_gpu_experimental_disable_binary_libraries: false
+        xla_gpu_dot_merger_threshold_mb: 64
+        xla_cpu_max_isa: ""
+        xla_gpu_experimental_enable_fusion_block_level_rewriter: false
+        xla_enable_fast_math: false
+        xla_gpu_experimental_parallel_collective_overlap_limit: 1
+        xla_cpu_copy_insertion_use_region_analysis: false
+        xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1
+        xla_gpu_experimental_enable_triton_heroless_priority_fusion: false
+        xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN
+        xla_gpu_experimental_stream_annotation: true
+        xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO
+        xla_pjrt_allow_auto_layout_in_hlo: false
+        xla_gpu_enable_scatter_determinism_expander: false
+        xla_gpu_require_exclusive_lock: false
+        xla_gpu_generate_debug_info: false
+        xla_gpu_generate_line_info: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false
+        xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE
+        xla_gpu_fail_ptx_compilation_on_register_spilling: false
+        xla_gpu_enable_analytical_sol_latency_estimator: true
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_prep_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_size_bytes"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "gpus_per_node"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nccl_op_launch_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nic_speed_gbps"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "rtt_us"
+          value: "-1"
+        }
+        xla_cpu_use_xnnpack: true
+        xla_gpu_experimental_pack_dot_operands_along_k_dimension: true
+        xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false
+        xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED
+        xla_gpu_experimental_enable_heuristic_collective_combining: false
+        xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false
+        xla_hlo_pass_fix_detect_cycles: false
+        xla_cpu_generate_unique_c_style_kernel_entry_points: false
+        xla_gpu_experimental_collective_cse_distance_threshold: 0
+        xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true
+        xla_gpu_experimental_collective_perf_table_path: ""
+        xla_gpu_collective_permute_combine_threshold_bytes: 31457287
+        xla_unsupported_crash_on_hlo_pass_noop_change: false
+        xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false
+        xla_dump_full_hlo_config: true
+        xla_gpu_unsupported_enable_triton_multi_output_fusion: false
+        xla_gpu_experimental_matmul_perf_table_path: ""
+        xla_gpu_unsupported_enable_all_reduce_decomposer: false
+        xla_gpu_nccl_blocking_communicators: true
+        xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20
+        xla_gpu_first_collective_call_terminate_timeout_seconds: 40
+        xla_gpu_nccl_async_execution: false
+        xla_cpu_emitter_verification_level: 1
+        xla_early_exit_with_layouts: false
+        xla_gpu_experimental_use_ragged_dot_fusion: false
+        xla_detect_unstable_reductions: DETECTION_MODE_NONE
+        xla_gpu_command_buffer_scheduling_mode: LHS
+        xla_gpu_experimental_enable_nccl_symmetric_buffers: false
+        xla_disable_automatic_host_compute_offload: false
+        xla_gpu_experimental_enable_fusion_autotuner: true
+        xla_gpu_experimental_scaled_dot_with_triton: true
+        xla_gpu_command_buffer_unroll_loops: false
+        xla_cpu_experimental_onednn_custom_call: false
+        xla_gpu_experimental_enable_checksum_tracing_on_thunks: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true
+        xla_cpu_collective_call_terminate_timeout_seconds: 40
+        xla_cpu_collective_call_warn_stuck_seconds: 20
+        xla_keep_shardings_after_spmd: false
+        xla_gpu_experimental_enable_triton_warp_specialization: false
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE
+        xla_cpu_enable_platform_dependent_math: true
+        xla_gpu_detect_nan: DETECTION_MODE_NONE
+        xla_gpu_detect_inf: DETECTION_MODE_NONE
+        xla_enable_enzyme_comms_opt: false
+        xla_gpu_experimental_enable_buffer_saver_on_thunks: false
+        xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE
+        xla_enable_scoped_logging_timers: true
+        xla_gpu_executable_embed_debug_info: true
+        xla_cpu_collective_timeout_seconds: 1800
+        xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false
+        xla_gpu_enable_pipelined_host_offloading: false
+        xla_gpu_default_to_alg_dot_bf16_bf16_f32: false
+        xla_gpu_experimental_autotune_backends: CUDNN
+        xla_gpu_experimental_autotune_backends: TRITON
+        xla_gpu_experimental_autotune_backends: CUBLASLT
+        xla_gpu_experimental_autotune_backends: HIPBLASLT
+        xla_gpu_experimental_autotune_backends: MIOPEN
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER
+        xla_gpu_experimental_autotune_backends: NATIVE_EMITTER
+        xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER
+        xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION
+        xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: YNNPACK
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU
+        xla_gpu_experimental_autotune_backends: TENSOR_IR
+        xla_gpu_experimental_all_fusions_with_triton: false
+        xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT
+        xla_gpu_rocm_max_trace_events: 4194304
+        xla_gpu_execution_terminate_timeout: "inf"
+        xla_enable_hlo_sharding_v3: true
+        xla_gpu_log_minmax: false
+        xla_gpu_print_compilation_stats: false
+        xla_gpu_experimental_enable_conv_fusion: false
+        xla_gpu_experimental_enable_tiling_propagation: true
+        xla_gpu_execution_progress_tracking: 0
+        xla_gpu_experimental_max_unroll_factor: 32
+        xla_gpu_enable_pdl: true
+        xla_gpu_collective_combine_threshold_count: 2048
+        xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true
+        xla_recognize_reduction_optimization_level: 0
+        xla_gpu_experimental_parallel_async_compute_limit: 2
+        xla_dump_buffer_assignment_analysis: true
+        xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE
+        xla_gpu_experimental_enable_collective_multi_streaming: false
+        xla_gpu_collective_permute_connected_components: false
+        xla_gpu_experimental_force_split_k: 0
+        xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_experimental_gemm_fusion_v2: true
+        xla_gpu_fusion_autotune_top_k_configs: 1
+        xla_gpu_native_emitter_tune_unroll_factor_for_loops: false
+        xla_enable_rgv3_materialization: true
+        xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK
+        xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false
+        xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_executable_num_compute_streams: 0
+        xla_gpu_executable_num_communication_streams: 0
+        xla_gpu_experimental_collective_start_as_early_as_possible: false
+        xla_cpu_experimental_enable_tiling_propagation: true
+        xla_gpu_enable_dus_accumulator_zero_init_elimination: false
+        xla_gpu_experimental_thunk_buffer_debug_module_outputs: false
+        xla_cpu_ftz: true
+        xla_gpu_enable_allocator_spatial_partitioning: true
+        xla_dump_compact_gte: true
+        xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES
+        xla_dump_hlo_as_riegeli: false
+        xla_gpu_experimental_use_ragged_dot_grouped_gemm: true
+        xla_dump_hlo_to_subfolder: false
+        xla_gpu_enable_pdl_launch: true
+        xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED
+        xla_gpu_enable_nccl_user_buffers_in_default_space: false
+        xla_sdy_export_all_reduce_scatter: false
+        xla_gpu_enable_gxl_ragged_all_to_all: false
+        xla_gpu_gxl_scratch_size_bytes: 67108864
+        xla_gpu_use_new_autotune_cache_format: true
+        xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE
+        xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON
+        xla_spmd_enable_dynamic_slice_collective_broadcast: false
+        xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false
+        xla_gpu_async_copy_min_bytes: -1
+        xla_gpu_blas_max_algorithms: 0
+        xla_gpu_experimental_enable_selective_memcpy_overlap: false
+        xla_gpu_experimental_enable_same_shape_multi_output_fusion: false
+        xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1
+        xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1
+        xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240
+        xla_cpu_use_new_xtile_lowering: true
+        xla_gpu_trace_annotation_level: 0
+        xla_gpu_experimental_cost_model_gemm_tiling_default: false
+        xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false
+        xla_gpu_all_reduce_splitter_ignore_profitability_check: false
+        xla_gpu_enable_cupti_multi_subscriber: true
+        xla_enable_hlo_modules_upload: true
+        xla_gpu_collective_domain_assignment: ""
+        xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1
+        xla_compile_all_supported_configs: false
+        xla_gpu_unsupported_use_cross_host_one_shot_kernel: ALLCOLLECTIVES
+        xla_deduplicate_backend_configs_min_size: 9223372036854775807
+        xla_force_config: ""
+        xla_candidate_configs_file: ""
+        xla_gpu_enable_persistent_symmetric_memory: false
+        xla_gpu_experimental_enable_raft_for_stable_topk: false
+        xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND
+        xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET
+        xla_gpu_experimental_emit_collective_reduce: false
+        xla_gpu_experimental_vmm_disabled: false
+        xla_gpu_experimental_enable_dynamic_slice_table_offsets: false
+        xla_gpu_device_execution_terminate_timeout: "inf"
+        xla_gpu_topology_filename: ""
+        xla_gpu_experimental_enable_dynamic_slice_extended_offsets: false
+        xla_gpu_experimental_deterministic_ops_triton: DETERMINISTIC_OPS_TRITON_MODE_OFF
+        xla_gpu_experimental_enable_tensor_ir: false
+      }
+      num_replicas: 2
+      num_partitions: 1
+      device_assignment {
+        replica_count: 2
+        computation_count: 1
+        computation_devices {
+          replica_device_ids: 0
+          replica_device_ids: 1
+        }
+      }
+      allow_spmd_sharding_propagation_to_output: false
+      comp_envs {
+      }
+      allow_spmd_sharding_propagation_to_parameters: false
+      process_count: 1
+      memory_fitting_level: EFFORT_O2
+    }
+  }
+  pjrt_client_name: "PjRtStreamExecutorClient"
+}
+gpu_executable {
+  hlo_module_with_config {
+    hlo_module {
+      name: "test"
+      entry_computation_name: "test_computation"
+      computations {
+        name: "fused_transpose"
+        instructions {
+          name: "param_0.1"
+          opcode: "parameter"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "param_0.1"
+          }
+          id: 4294967298
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "bitcast.6.1"
+          opcode: "bitcast"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "bitcast.6.1"
+          }
+          id: 4294967297
+          operand_ids: 4294967298
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "transpose.1.1"
+          opcode: "transpose"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "transpose.1.1"
+          }
+          dimensions: 1
+          dimensions: 0
+          id: 4294967296
+          operand_ids: 4294967297
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          parameters {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          result {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          parameter_names: "param_0.1"
+        }
+        id: 1
+        root_id: 4294967296
+        is_fusion_computation: true
+      }
+      computations {
+        name: "fused_transpose.1"
+        instructions {
+          name: "param_0.3"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "param_0.3"
+          }
+          id: 8589934594
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "bitcast.4.1"
+          opcode: "bitcast"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "bitcast.4.1"
+          }
+          id: 8589934593
+          operand_ids: 8589934594
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "transpose.3"
+          opcode: "transpose"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "transpose.3"
+          }
+          dimensions: 1
+          dimensions: 0
+          id: 8589934592
+          operand_ids: 8589934593
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          parameters {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          result {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          parameter_names: "param_0.3"
+        }
+        id: 2
+        root_id: 8589934592
+        is_fusion_computation: true
+      }
+      computations {
+        name: "fused_broadcast"
+        instructions {
+          name: "param_0.5"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "param_0.5"
+          }
+          id: 12884901890
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "p1.1.1"
+          opcode: "convert"
+          shape {
+            element_type: F32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "p1.1.1"
+          }
+          id: 12884901889
+          operand_ids: 12884901890
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "broadcast.13.1"
+          opcode: "broadcast"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "broadcast.13.1"
+          }
+          id: 12884901888
+          operand_ids: 12884901889
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          parameters {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          result {
+            element_type: F32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          parameter_names: "param_0.5"
+        }
+        id: 3
+        root_id: 12884901888
+        is_fusion_computation: true
+      }
+      computations {
+        name: "wrapped_broadcast_computation"
+        instructions {
+          name: "param_0.7"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "param_0.7"
+          }
+          id: 21474836481
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "broadcast.12.1"
+          opcode: "broadcast"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "broadcast.12.1"
+          }
+          id: 21474836480
+          operand_ids: 21474836481
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          parameters {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          result {
+            element_type: U32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          parameter_names: "param_0.7"
+        }
+        id: 5
+        root_id: 21474836480
+        is_fusion_computation: true
+      }
+      computations {
+        name: "test_computation"
+        instructions {
+          name: "id.1"
+          opcode: "replica-id"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "id.1"
+          }
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "wrapped_broadcast"
+          opcode: "fusion"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "wrapped_broadcast"
+          }
+          fusion_kind: "kLoop"
+          id: 9
+          operand_ids: 0
+          called_computation_ids: 5
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "loop_broadcast_fusion"
+          opcode: "fusion"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "loop_broadcast_fusion"
+          }
+          fusion_kind: "kLoop"
+          id: 8
+          operand_ids: 0
+          called_computation_ids: 3
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "allgather.3"
+          opcode: "all-gather"
+          shape {
+            element_type: TUPLE
+            tuple_shapes {
+              element_type: U32
+              dimensions: 2
+              dimensions: 2
+              layout {
+                minor_to_major: 0
+                minor_to_major: 1
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+              is_dynamic_dimension: false
+            }
+            tuple_shapes {
+              element_type: F32
+              dimensions: 2
+              dimensions: 2
+              layout {
+                minor_to_major: 0
+                minor_to_major: 1
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+              is_dynamic_dimension: false
+            }
+          }
+          metadata {
+            scheduling_name: "all-gather-start"
+          }
+          dimensions: 1
+          id: 10
+          operand_ids: 9
+          operand_ids: 8
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+          collective_device_list {
+          }
+          backend_config_payload {
+            id: 0
+          }
+        }
+        instructions {
+          name: "ag0.1"
+          opcode: "get-tuple-element"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "ag0.1"
+          }
+          id: 1
+          operand_ids: 10
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "loop_transpose_fusion.1"
+          opcode: "fusion"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "loop_transpose_fusion.1"
+          }
+          fusion_kind: "kLoop"
+          id: 7
+          operand_ids: 1
+          called_computation_ids: 2
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "bitcast.24"
+          opcode: "bitcast"
+          shape {
+            element_type: U32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "bitcast.24"
+          }
+          id: 4
+          operand_ids: 7
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "ag1.1"
+          opcode: "get-tuple-element"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "ag1.1"
+          }
+          tuple_index: 1
+          id: 2
+          operand_ids: 10
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "loop_transpose_fusion"
+          opcode: "fusion"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "loop_transpose_fusion"
+          }
+          fusion_kind: "kLoop"
+          id: 6
+          operand_ids: 2
+          called_computation_ids: 1
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "bitcast.1.0"
+          opcode: "bitcast"
+          shape {
+            element_type: F32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "bitcast.1.0"
+          }
+          id: 5
+          operand_ids: 6
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "out.1"
+          opcode: "tuple"
+          shape {
+            element_type: TUPLE
+            tuple_shapes {
+              element_type: U32
+              dimensions: 4
+              layout {
+                minor_to_major: 0
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+            }
+            tuple_shapes {
+              element_type: F32
+              dimensions: 4
+              layout {
+                minor_to_major: 0
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+            }
+          }
+          metadata {
+            scheduling_name: "out.1"
+          }
+          id: 3
+          operand_ids: 4
+          operand_ids: 5
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          result {
+            element_type: TUPLE
+            tuple_shapes {
+              element_type: U32
+              dimensions: 4
+              layout {
+                minor_to_major: 0
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+            }
+            tuple_shapes {
+              element_type: F32
+              dimensions: 4
+              layout {
+                minor_to_major: 0
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+            }
+          }
+        }
+        root_id: 3
+      }
+      host_program_shape {
+        result {
+          element_type: TUPLE
+          tuple_shapes {
+            element_type: U32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          tuple_shapes {
+            element_type: F32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+        }
+      }
+      schedule {
+        sequences {
+          key: 0
+          value {
+            instruction_ids: 0
+            instruction_ids: 8
+            instruction_ids: 9
+            instruction_ids: 10
+            instruction_ids: 1
+            instruction_ids: 7
+            instruction_ids: 2
+            instruction_ids: 6
+            instruction_ids: 4
+            instruction_ids: 5
+            instruction_ids: 3
+          }
+        }
+      }
+      input_output_alias {
+      }
+      device_assignment {
+        replica_count: 2
+        computation_count: 1
+        computation_devices {
+          replica_device_ids: 0
+          replica_device_ids: 1
+        }
+      }
+      buffer_donor {
+      }
+      frontend_attributes {
+        map {
+          key: "fingerprint_before_lhs"
+          value: "7b2e4f79c5524b5cbc566da178aa7fda"
+        }
+      }
+      payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}"
+    }
+    config {
+      entry_computation_layout {
+        result {
+          element_type: TUPLE
+          tuple_shapes {
+            element_type: U32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          tuple_shapes {
+            element_type: F32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+        }
+      }
+      replica_count: 2
+      num_partitions: 1
+      intra_op_parallelism_threads: -1
+      debug_options {
+        xla_disable_hlo_passes: "constant_folding"
+        xla_backend_optimization_level: 3
+        xla_eliminate_hlo_implicit_broadcast: true
+        xla_cpu_multi_thread_eigen: true
+        xla_gpu_cuda_data_dir: "./cuda_sdk_lib"
+        xla_llvm_enable_alias_scope_metadata: true
+        xla_llvm_enable_noalias_metadata: true
+        xla_llvm_enable_invariant_load_metadata: true
+        xla_llvm_disable_expensive_passes: false
+        xla_cpu_use_onednn: false
+        xla_cpu_enable_fast_math: false
+        xla_gpu_enable_fast_min_max: false
+        xla_force_host_platform_device_count: 1
+        xla_gpu_enable_dynamic_slice_fusion: true
+        xla_hlo_evaluator_use_fast_path: true
+        xla_dump_hlo_as_html: false
+        xla_cpu_fast_math_honor_nans: true
+        xla_cpu_fast_math_honor_infs: true
+        xla_allow_excess_precision: true
+        xla_gpu_autotune_level: 4
+        xla_cpu_fast_math_honor_division: true
+        xla_cpu_fast_math_honor_functions: true
+        xla_dump_include_timestamp: false
+        xla_dump_max_hlo_modules: -1
+        xla_cpu_enable_xprof_traceme: false
+        xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false
+        xla_cpu_enable_fast_min_max: true
+        xla_multiheap_size_constraint_per_heap: -1
+        xla_dump_module_metadata: false
+        xla_gpu_force_compilation_parallelism: 1
+        xla_dump_fusion_visualization: false
+        xla_gpu_all_reduce_combine_threshold_bytes: 31457287
+        xla_gpu_nccl_termination_timeout_seconds: -1
+        xla_dump_hlo_as_long_text: true
+        xla_gpu_enable_cublaslt: true
+        xla_gpu_shape_checks: RUNTIME
+        xla_gpu_use_runtime_fusion: false
+        xla_dump_latency_hiding_schedule: false
+        xla_dump_enable_mlir_pretty_form: true
+        xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP
+        xla_gpu_enable_triton_gemm: true
+        xla_gpu_enable_cudnn_int8x32_convolution_reordering: true
+        xla_gpu_triton_gemm_any: true
+        xla_gpu_enable_while_loop_reduce_scatter_code_motion: false
+        xla_gpu_collective_inflation_factor: 1
+        xla_gpu_graph_min_graph_size: 5
+        xla_gpu_enable_reassociation_for_converted_ar: true
+        xla_gpu_pgle_profile_file_or_directory_path: ""
+        xla_gpu_all_gather_combine_threshold_bytes: 31457287
+        xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287
+        xla_gpu_enable_highest_priority_async_stream: true
+        xla_gpu_exhaustive_tiling_search: false
+        xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0
+        xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1
+        xla_gpu_redzone_padding_bytes: 8388608
+        xla_gpu_copy_insertion_use_region_analysis: false
+        xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807
+        xla_gpu_collect_cost_model_stats: false
+        xla_gpu_enable_pipelined_p2p: false
+        xla_gpu_cublas_fallback: true
+        xla_gpu_enable_while_loop_double_buffering: false
+        xla_gpu_filter_kernels_spilling_registers_on_autotuning: true
+        xla_debug_buffer_assignment_show_max: 15
+        xla_detailed_logging: true
+        xla_enable_dumping: true
+        xla_gpu_enable_all_gather_combine_by_dim: false
+        xla_gpu_enable_analytical_latency_estimator: false
+        xla_gpu_llvm_verification_level: 0
+        xla_gpu_enable_reduce_scatter_combine_by_dim: false
+        xla_gpu_enable_command_buffer: COLLECTIVES_KERNEL
+        xla_gpu_enable_command_buffer: CONDITIONAL
+        xla_gpu_enable_command_buffer: CUBLAS
+        xla_gpu_enable_command_buffer: CUBLASLT
+        xla_gpu_enable_command_buffer: CUDNN
+        xla_gpu_enable_command_buffer: CUSTOM_CALL
+        xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION
+        xla_gpu_enable_command_buffer: FUSION
+        xla_gpu_enable_cub_radix_sort: true
+        xla_gpu_memory_limit_slop_factor: 95
+        xla_gpu_target_config_filename: ""
+        xla_gpu_enable_cudnn_layer_norm: false
+        xla_gpu_threshold_for_windowed_einsum_mib: 100000
+        xla_gpu_enable_nccl_user_buffers: false
+        xla_gpu_enable_libnvptxcompiler: true
+        xla_gpu_enable_nccl_comm_splitting: true
+        xla_gpu_nccl_collective_max_nchannels: 0
+        xla_gpu_nccl_p2p_max_nchannels: 0
+        xla_gpu_nccl_init_max_rank_per_root_ratio: 0
+        xla_gpu_multi_streamed_windowed_einsum: true
+        xla_gpu_gemm_rewrite_size_threshold: 100
+        xla_gpu_require_complete_aot_autotune_results: false
+        xla_gpu_cudnn_gemm_fusion_level: 0
+        xla_gpu_use_memcpy_local_p2p: false
+        xla_gpu_autotune_max_solutions: 0
+        xla_dump_large_constants: false
+        xla_gpu_verify_triton_fusion_numerics: false
+        xla_reduce_window_rewrite_base_length: 16
+        xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL
+        xla_gpu_enable_host_memory_offloading: false
+        xla_gpu_exclude_nondeterministic_ops: true
+        xla_gpu_nccl_terminate_on_error: false
+        xla_gpu_shard_autotuning: true
+        xla_gpu_enable_approx_costly_collectives: false
+        xla_cpu_prefer_vector_width: 256
+        xla_gpu_per_fusion_autotune_cache_dir: ""
+        xla_cmd_buffer_trace_cache_size: 16
+        xla_gpu_temp_buffer_use_separate_color: false
+        xla_syntax_sugar_async_ops: false
+        xla_gpu_autotune_gemm_rtol: 0.1
+        xla_enable_command_buffers_during_profiling: false
+        xla_gpu_cudnn_gemm_max_plans: 5
+        xla_cpu_parallel_codegen_split_count: 32
+        xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE
+        xla_gpu_executable_warn_stuck_timeout_seconds: 10
+        xla_gpu_executable_terminate_timeout_seconds: 30
+        xla_gpu_experimental_disable_binary_libraries: false
+        xla_gpu_dot_merger_threshold_mb: 64
+        xla_cpu_max_isa: ""
+        xla_gpu_experimental_enable_fusion_block_level_rewriter: false
+        xla_enable_fast_math: false
+        xla_gpu_experimental_parallel_collective_overlap_limit: 1
+        xla_cpu_copy_insertion_use_region_analysis: false
+        xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1
+        xla_gpu_experimental_enable_triton_heroless_priority_fusion: false
+        xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN
+        xla_gpu_experimental_stream_annotation: true
+        xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO
+        xla_pjrt_allow_auto_layout_in_hlo: false
+        xla_gpu_enable_scatter_determinism_expander: false
+        xla_gpu_require_exclusive_lock: false
+        xla_gpu_generate_debug_info: false
+        xla_gpu_generate_line_info: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false
+        xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE
+        xla_gpu_fail_ptx_compilation_on_register_spilling: false
+        xla_gpu_enable_analytical_sol_latency_estimator: true
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_prep_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_size_bytes"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "gpus_per_node"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nccl_op_launch_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nic_speed_gbps"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "rtt_us"
+          value: "-1"
+        }
+        xla_cpu_use_xnnpack: true
+        xla_gpu_experimental_pack_dot_operands_along_k_dimension: true
+        xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false
+        xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED
+        xla_gpu_experimental_enable_heuristic_collective_combining: false
+        xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false
+        xla_hlo_pass_fix_detect_cycles: false
+        xla_cpu_generate_unique_c_style_kernel_entry_points: false
+        xla_gpu_experimental_collective_cse_distance_threshold: 0
+        xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true
+        xla_gpu_experimental_collective_perf_table_path: ""
+        xla_gpu_collective_permute_combine_threshold_bytes: 31457287
+        xla_unsupported_crash_on_hlo_pass_noop_change: false
+        xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false
+        xla_dump_full_hlo_config: true
+        xla_gpu_unsupported_enable_triton_multi_output_fusion: false
+        xla_gpu_experimental_matmul_perf_table_path: ""
+        xla_gpu_unsupported_enable_all_reduce_decomposer: false
+        xla_gpu_nccl_blocking_communicators: true
+        xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20
+        xla_gpu_first_collective_call_terminate_timeout_seconds: 40
+        xla_gpu_nccl_async_execution: false
+        xla_cpu_emitter_verification_level: 1
+        xla_early_exit_with_layouts: false
+        xla_gpu_experimental_use_ragged_dot_fusion: false
+        xla_detect_unstable_reductions: DETECTION_MODE_NONE
+        xla_gpu_command_buffer_scheduling_mode: LHS
+        xla_gpu_experimental_enable_nccl_symmetric_buffers: false
+        xla_disable_automatic_host_compute_offload: false
+        xla_gpu_experimental_enable_fusion_autotuner: true
+        xla_gpu_experimental_scaled_dot_with_triton: true
+        xla_gpu_command_buffer_unroll_loops: false
+        xla_cpu_experimental_onednn_custom_call: false
+        xla_gpu_experimental_enable_checksum_tracing_on_thunks: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true
+        xla_cpu_collective_call_terminate_timeout_seconds: 40
+        xla_cpu_collective_call_warn_stuck_seconds: 20
+        xla_keep_shardings_after_spmd: false
+        xla_gpu_experimental_enable_triton_warp_specialization: false
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE
+        xla_cpu_enable_platform_dependent_math: true
+        xla_gpu_detect_nan: DETECTION_MODE_NONE
+        xla_gpu_detect_inf: DETECTION_MODE_NONE
+        xla_enable_enzyme_comms_opt: false
+        xla_gpu_experimental_enable_buffer_saver_on_thunks: false
+        xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE
+        xla_enable_scoped_logging_timers: true
+        xla_gpu_executable_embed_debug_info: true
+        xla_cpu_collective_timeout_seconds: 1800
+        xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false
+        xla_gpu_enable_pipelined_host_offloading: false
+        xla_gpu_default_to_alg_dot_bf16_bf16_f32: false
+        xla_gpu_experimental_autotune_backends: CUDNN
+        xla_gpu_experimental_autotune_backends: TRITON
+        xla_gpu_experimental_autotune_backends: CUBLASLT
+        xla_gpu_experimental_autotune_backends: HIPBLASLT
+        xla_gpu_experimental_autotune_backends: MIOPEN
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER
+        xla_gpu_experimental_autotune_backends: NATIVE_EMITTER
+        xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER
+        xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION
+        xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: YNNPACK
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU
+        xla_gpu_experimental_autotune_backends: TENSOR_IR
+        xla_gpu_experimental_all_fusions_with_triton: false
+        xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT
+        xla_gpu_rocm_max_trace_events: 4194304
+        xla_gpu_execution_terminate_timeout: "inf"
+        xla_enable_hlo_sharding_v3: true
+        xla_gpu_log_minmax: false
+        xla_gpu_print_compilation_stats: false
+        xla_gpu_experimental_enable_conv_fusion: false
+        xla_gpu_experimental_enable_tiling_propagation: true
+        xla_gpu_execution_progress_tracking: 0
+        xla_gpu_experimental_max_unroll_factor: 32
+        xla_gpu_enable_pdl: true
+        xla_gpu_collective_combine_threshold_count: 2048
+        xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true
+        xla_recognize_reduction_optimization_level: 0
+        xla_gpu_experimental_parallel_async_compute_limit: 2
+        xla_dump_buffer_assignment_analysis: true
+        xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE
+        xla_gpu_experimental_enable_collective_multi_streaming: false
+        xla_gpu_collective_permute_connected_components: false
+        xla_gpu_experimental_force_split_k: 0
+        xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_experimental_gemm_fusion_v2: true
+        xla_gpu_fusion_autotune_top_k_configs: 1
+        xla_gpu_native_emitter_tune_unroll_factor_for_loops: false
+        xla_enable_rgv3_materialization: true
+        xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK
+        xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false
+        xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_executable_num_compute_streams: 0
+        xla_gpu_executable_num_communication_streams: 0
+        xla_gpu_experimental_collective_start_as_early_as_possible: false
+        xla_cpu_experimental_enable_tiling_propagation: true
+        xla_gpu_enable_dus_accumulator_zero_init_elimination: false
+        xla_gpu_experimental_thunk_buffer_debug_module_outputs: false
+        xla_cpu_ftz: true
+        xla_gpu_enable_allocator_spatial_partitioning: true
+        xla_dump_compact_gte: true
+        xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES
+        xla_dump_hlo_as_riegeli: false
+        xla_gpu_experimental_use_ragged_dot_grouped_gemm: true
+        xla_dump_hlo_to_subfolder: false
+        xla_gpu_enable_pdl_launch: true
+        xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED
+        xla_gpu_enable_nccl_user_buffers_in_default_space: false
+        xla_sdy_export_all_reduce_scatter: false
+        xla_gpu_enable_gxl_ragged_all_to_all: false
+        xla_gpu_gxl_scratch_size_bytes: 67108864
+        xla_gpu_use_new_autotune_cache_format: true
+        xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE
+        xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON
+        xla_spmd_enable_dynamic_slice_collective_broadcast: false
+        xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false
+        xla_gpu_async_copy_min_bytes: -1
+        xla_gpu_blas_max_algorithms: 0
+        xla_gpu_experimental_enable_selective_memcpy_overlap: false
+        xla_gpu_experimental_enable_same_shape_multi_output_fusion: false
+        xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1
+        xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1
+        xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240
+        xla_cpu_use_new_xtile_lowering: true
+        xla_gpu_trace_annotation_level: 0
+        xla_gpu_experimental_cost_model_gemm_tiling_default: false
+        xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false
+        xla_gpu_all_reduce_splitter_ignore_profitability_check: false
+        xla_gpu_enable_cupti_multi_subscriber: true
+        xla_enable_hlo_modules_upload: true
+        xla_gpu_collective_domain_assignment: ""
+        xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1
+        xla_compile_all_supported_configs: false
+        xla_gpu_unsupported_use_cross_host_one_shot_kernel: ALLCOLLECTIVES
+        xla_deduplicate_backend_configs_min_size: 9223372036854775807
+        xla_force_config: ""
+        xla_candidate_configs_file: ""
+        xla_gpu_enable_persistent_symmetric_memory: false
+        xla_gpu_experimental_enable_raft_for_stable_topk: false
+        xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND
+        xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET
+        xla_gpu_experimental_emit_collective_reduce: false
+        xla_gpu_experimental_vmm_disabled: false
+        xla_gpu_experimental_enable_dynamic_slice_table_offsets: false
+        xla_gpu_device_execution_terminate_timeout: "inf"
+        xla_gpu_topology_filename: ""
+        xla_gpu_experimental_enable_dynamic_slice_extended_offsets: false
+        xla_gpu_experimental_deterministic_ops_triton: DETERMINISTIC_OPS_TRITON_MODE_OFF
+        xla_gpu_experimental_enable_tensor_ir: false
+      }
+      static_device_assignment {
+        replica_count: 2
+        computation_count: 1
+        computation_devices {
+          replica_device_ids: 0
+          replica_device_ids: 1
+        }
+      }
+      allow_spmd_sharding_propagation_to_output: false
+      allow_spmd_sharding_propagation_to_parameters: false
+      sharding_config {
+      }
+      memory_fitting_level: EFFORT_O2
+      schedule_config {
+      }
+      partition_size: 2
+    }
+  }
+  buffer_assignment {
+    logical_buffers {
+      id: 9
+      size: 4
+      defined_at {
+      }
+    }
+    logical_buffers {
+      id: 10
+      size: 8
+      defined_at {
+        instruction_id: 9
+      }
+    }
+    logical_buffers {
+      id: 11
+      size: 8
+      defined_at {
+        instruction_id: 8
+      }
+    }
+    logical_buffers {
+      id: 12
+      size: 16
+      defined_at {
+        instruction_id: 10
+      }
+    }
+    logical_buffers {
+      id: 13
+      size: 16
+      defined_at {
+        shape_index: 0
+        instruction_id: 10
+      }
+    }
+    logical_buffers {
+      id: 14
+      size: 16
+      defined_at {
+        shape_index: 1
+        instruction_id: 10
+      }
+    }
+    logical_buffers {
+      id: 15
+      size: 16
+      defined_at {
+        instruction_id: 7
+      }
+    }
+    logical_buffers {
+      id: 16
+      size: 16
+      defined_at {
+        instruction_id: 6
+      }
+    }
+    logical_buffers {
+      id: 17
+      size: 16
+      defined_at {
+        instruction_id: 3
+      }
+    }
+    buffer_allocations {
+      size: 16
+      maybe_live_out: true
+      assigned {
+        logical_buffer_id: 10
+        size: 8
+        element_type: U32
+      }
+      assigned {
+        logical_buffer_id: 15
+        size: 16
+        element_type: U32
+      }
+    }
+    buffer_allocations {
+      index: 1
+      size: 16
+      maybe_live_out: true
+      assigned {
+        logical_buffer_id: 9
+        size: 4
+        element_type: U32
+      }
+      assigned {
+        logical_buffer_id: 13
+        size: 16
+        element_type: U32
+      }
+      assigned {
+        logical_buffer_id: 16
+        size: 16
+        element_type: F32
+      }
+    }
+    buffer_allocations {
+      index: 2
+      size: 16
+      maybe_live_out: true
+      assigned {
+        logical_buffer_id: 17
+        size: 16
+        element_type: TUPLE
+      }
+      is_tuple: true
+    }
+    buffer_allocations {
+      index: 3
+      size: 520
+      assigned {
+        logical_buffer_id: 11
+        offset: 512
+        size: 8
+        element_type: F32
+      }
+      assigned {
+        logical_buffer_id: 12
+        size: 16
+        element_type: TUPLE
+      }
+      assigned {
+        logical_buffer_id: 14
+        offset: 256
+        size: 16
+        element_type: F32
+      }
+      is_tuple: true
+    }
+    heap_simulator_traces {
+      events {
+        buffer_id: 11
+        computation_name: "test_computation"
+        instruction_name: "loop_broadcast_fusion"
+      }
+      events {
+        buffer_id: 14
+        computation_name: "test_computation"
+        instruction_name: "allgather.3"
+      }
+      events {
+        kind: FREE
+        buffer_id: 11
+        computation_name: "test_computation"
+        instruction_name: "loop_broadcast_fusion"
+      }
+      events {
+        kind: FREE
+        buffer_id: 14
+        computation_name: "test_computation"
+        instruction_name: "allgather.3"
+      }
+      whole_module_simulation: true
+      buffer_allocation_index: 3
+    }
+    peak_buffers {
+      logical_buffer_ids: 15
+    }
+    peak_buffers {
+      index: 1
+      logical_buffer_ids: 16
+    }
+    peak_buffers {
+      index: 2
+      logical_buffer_ids: 17
+    }
+    peak_buffers {
+      index: 3
+      logical_buffer_ids: 12
+      logical_buffer_ids: 11
+      logical_buffer_ids: 14
+    }
+  }
+  gpu_compute_capability {
+    cuda_compute_capability {
+      major: 10
+      feature_extension: ACCELERATED_FEATURES
+    }
+  }
+  module_name: "test"
+  program_shape {
+    result {
+      element_type: TUPLE
+      tuple_shapes {
+        element_type: U32
+        dimensions: 4
+        layout {
+          minor_to_major: 0
+          tail_padding_alignment_in_elements: 1
+        }
+        is_dynamic_dimension: false
+      }
+      tuple_shapes {
+        element_type: F32
+        dimensions: 4
+        layout {
+          minor_to_major: 0
+          tail_padding_alignment_in_elements: 1
+        }
+        is_dynamic_dimension: false
+      }
+    }
+  }
+  buffer_allocations {
+    values {
+      size: 16
+      maybe_live_out: true
+    }
+    values {
+      index: 1
+      size: 16
+      maybe_live_out: true
+    }
+    values {
+      index: 2
+      size: 16
+      maybe_live_out: true
+      is_tuple: true
+    }
+    values {
+      index: 3
+      size: 520
+      is_tuple: true
+    }
+  }
+  output_info_map {
+    shape_index {
+    }
+    output_info {
+      allocation_index: 2
+    }
+  }
+  output_info_map {
+    shape_index {
+      indexes: 0
+    }
+    output_info {
+    }
+  }
+  output_info_map {
+    shape_index {
+      indexes: 1
+    }
+    output_info {
+      allocation_index: 1
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "id.1"
+      thunk_id: 1
+    }
+    replica_id_thunk {
+      dest_buffer {
+        size: 4
+        buffer_allocation_index: 1
+        element_type: U32
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "loop_broadcast_fusion"
+      thunk_id: 2
+    }
+    custom_kernel_thunk {
+      args {
+        slice {
+          size: 4
+          buffer_allocation_index: 1
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          layout {
+            tail_padding_alignment_in_elements: 1
+          }
+        }
+      }
+      args {
+        slice {
+          offset: 512
+          size: 8
+          buffer_allocation_index: 3
+          element_type: F32
+        }
+        shape {
+          element_type: F32
+          dimensions: 2
+          dimensions: 1
+          layout {
+            minor_to_major: 0
+            minor_to_major: 1
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+          is_dynamic_dimension: false
+        }
+      }
+      written: false
+      written: true
+      custom_kernel {
+        name: "loop_broadcast_fusion"
+        kernel_spec {
+          cubin {
+            data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\310\024\000\000\000\000\000\000H\017\000\000\000\000\000\000\002d\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_broadcast_fusion\000.nv.capmerc.text.loop_broadcast_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_broadcast_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\271\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\350\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000U\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000k\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004(\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\002\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 13.2, V13.2.51\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills  -arch sm_100a \000\000\014\000\000\000\010\000\000\000\350\003\000\000NVIDIA Corp\000\002\000d\000\204\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\002\t\001\000\002\002\001\000\002\005\005\000\003\007\001\001\002\003\000\000\004\013\010\000\t\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000\240\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\t\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\002\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\342/\000Er\007\000\002\000\000\000\000\020 \000\000\312O\000\206y\000\004\007\000\000\000\004\031\020\014\000\342\017\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\001\000\000\300\014\000\000\000\300\005\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\001\000\000\000\000\000\000\000\000\000\000\000\320\006\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\004\020\000\000\000\004\360\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004Z \000\212\235\"\244\261\235\024m\000\264*\363\367X\003\245\',!0\311\036\307\217\017\014Il\n/\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000D\000\000\000\000\000\000\000=\000\001\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\271\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\350\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000U\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\373\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000_\002\000\000\000\000\000\000\217\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\317\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000H\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\374\005\000\000\000\000\000\000 \000\000\000\000\000\000\000\005\000\000\000\010\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\006\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000y\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000`\006\000\000\000\000\000\000p\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\355\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\010\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\271\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\033\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000?\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\220\014\000\000\000\000\000\000\262\000\000\000\000\000\000\000\025\000\000\000\010\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000f\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000B\r\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000{\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\264\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\214\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\330\r\000\000\000\000\000\000|\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000X\016\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\315\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000p\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000p\016\000\000\000\000\000\000\330\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\310\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\310\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\200\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000"
+          }
+          arity: 2
+          kernel_name: "loop_broadcast_fusion"
+          kernel_args_packing_spec {
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 1
+              }
+            }
+          }
+        }
+        block_dims {
+          coordinates {
+            x: 1
+            y: 1
+            z: 1
+          }
+        }
+        thread_dims {
+          coordinates {
+            x: 2
+            y: 1
+            z: 1
+          }
+        }
+      }
+      use_pdl: true
+      tma_metadata {
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "wrapped_broadcast"
+      thunk_id: 3
+    }
+    custom_kernel_thunk {
+      args {
+        slice {
+          size: 4
+          buffer_allocation_index: 1
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          layout {
+            tail_padding_alignment_in_elements: 1
+          }
+        }
+      }
+      args {
+        slice {
+          size: 8
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          dimensions: 2
+          dimensions: 1
+          layout {
+            minor_to_major: 0
+            minor_to_major: 1
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+          is_dynamic_dimension: false
+        }
+      }
+      written: false
+      written: true
+      custom_kernel {
+        name: "wrapped_broadcast"
+        kernel_spec {
+          cubin {
+            data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000H\024\000\000\000\000\000\000\310\016\000\000\000\000\000\000\002d\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.wrapped_broadcast\000.nv.capmerc.text.wrapped_broadcast\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.wrapped_broadcast\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\375\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000-\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000[\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004$\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\002\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 13.2, V13.2.51\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills  -arch sm_100a \000\000\014\000\000\000\010\000\000\000\350\003\000\000NVIDIA Corp\000\002\000d\000\204\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\002\t\001\000\002\002\001\000\002\005\005\000\003\007\001\001\002\003\000\000\004\013\010\000\t\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000\220\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\t\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\312/\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\001\000\000\300\013\000\000\000\300\002\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\007\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\004\020\000\000\000\004\340\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004Z \000\212\235\"\244\261\235\024m\000\264*\363\367X\003\245\',!0\311\036\307\217\017\014Il\n/\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000\340\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000D\000\000\000\000\000\000\000=\000\001\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\375\000\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000-\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\343\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000G\002\000\000\000\000\000\000{\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\003\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\303\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\270\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\324\005\000\000\000\000\000\000 \000\000\000\000\000\000\000\005\000\000\000\010\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\364\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000u\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\006\000\000\000\000\000\000p\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\363\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\341\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\010\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\017\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000/\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\262\000\000\000\000\000\000\000\025\000\000\000\010\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\302\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000g\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0004\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000X\r\000\000\000\000\000\000|\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\233\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\330\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\265\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\323\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\r\000\000\000\000\000\000\330\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000H\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000H\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000"
+          }
+          arity: 2
+          kernel_name: "wrapped_broadcast"
+          kernel_args_packing_spec {
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 1
+              }
+            }
+          }
+        }
+        block_dims {
+          coordinates {
+            x: 1
+            y: 1
+            z: 1
+          }
+        }
+        thread_dims {
+          coordinates {
+            x: 2
+            y: 1
+            z: 1
+          }
+        }
+      }
+      use_pdl: true
+      tma_metadata {
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "allgather.3"
+      thunk_id: 4
+    }
+    all_gather_thunk {
+      collective_config {
+        operand_element_type: U32
+        operand_element_type: F32
+      }
+      buffers {
+        element_count: 2
+        source_buffer {
+          slice {
+            size: 8
+            element_type: U32
+          }
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+        }
+        destination_buffer {
+          slice {
+            size: 16
+            buffer_allocation_index: 1
+            element_type: U32
+          }
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+        }
+      }
+      buffers {
+        element_count: 2
+        source_buffer {
+          slice {
+            offset: 512
+            size: 8
+            buffer_allocation_index: 3
+            element_type: F32
+          }
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+        }
+        destination_buffer {
+          slice {
+            offset: 256
+            size: 16
+            buffer_allocation_index: 3
+            element_type: F32
+          }
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+        }
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "loop_transpose_fusion.1"
+      thunk_id: 5
+    }
+    custom_kernel_thunk {
+      args {
+        slice {
+          size: 16
+          buffer_allocation_index: 1
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          dimensions: 2
+          dimensions: 2
+          layout {
+            minor_to_major: 0
+            minor_to_major: 1
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+          is_dynamic_dimension: false
+        }
+      }
+      args {
+        slice {
+          size: 16
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          dimensions: 2
+          dimensions: 2
+          layout {
+            minor_to_major: 1
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+          is_dynamic_dimension: false
+        }
+      }
+      written: false
+      written: true
+      custom_kernel {
+        name: "loop_transpose_fusion_1"
+        kernel_spec {
+          cubin {
+            data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\310\024\000\000\000\000\000\000H\017\000\000\000\000\000\000\002d\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion_1\000.nv.capmerc.text.loop_transpose_fusion_1\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion_1\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\277\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\356\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\017\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000?\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000[\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000s\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\002\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 13.2, V13.2.51\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills  -arch sm_100a \000\000\014\000\000\000\010\000\000\000\350\003\000\000NVIDIA Corp\000\002\000d\000\204\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\002\t\001\000\002\002\001\000\002\005\005\000\003\007\001\001\002\003\000\000\004\013\010\000\t\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\t\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\001\000\000\300\017\000\000\000\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004Z \000\212\235\"\244\261\235\024m\000\264*\363\367X\003\245\',!0\311\036\307\217\017\014Il\n/\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000D\000\000\000\000\000\000\000=\000\001\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\277\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\356\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\017\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000?\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000[\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\007\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000k\002\000\000\000\000\000\000\231\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\004\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\325\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\370\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000`\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\024\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\005\000\000\000\010\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0004\006\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000X\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000{\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000p\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\001\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\350\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\363\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\007\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\010\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\277\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000!\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000G\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\220\014\000\000\000\000\000\000\262\000\000\000\000\000\000\000\025\000\000\000\010\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000B\r\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\205\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\264\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\226\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\330\r\000\000\000\000\000\000|\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\277\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000X\016\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\331\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000p\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\367\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000p\016\000\000\000\000\000\000\330\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\310\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\310\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\200\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000"
+          }
+          arity: 2
+          kernel_name: "loop_transpose_fusion_1"
+          kernel_args_packing_spec {
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 1
+              }
+            }
+          }
+        }
+        block_dims {
+          coordinates {
+            x: 1
+            y: 1
+            z: 1
+          }
+        }
+        thread_dims {
+          coordinates {
+            x: 4
+            y: 1
+            z: 1
+          }
+        }
+      }
+      use_pdl: true
+      tma_metadata {
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "loop_transpose_fusion"
+      thunk_id: 6
+    }
+    custom_kernel_thunk {
+      args {
+        slice {
+          offset: 256
+          size: 16
+          buffer_allocation_index: 3
+          element_type: F32
+        }
+        shape {
+          element_type: F32
+          dimensions: 2
+          dimensions: 2
+          layout {
+            minor_to_major: 0
+            minor_to_major: 1
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+          is_dynamic_dimension: false
+        }
+      }
+      args {
+        slice {
+          size: 16
+          buffer_allocation_index: 1
+          element_type: F32
+        }
+        shape {
+          element_type: F32
+          dimensions: 2
+          dimensions: 2
+          layout {
+            minor_to_major: 1
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+          is_dynamic_dimension: false
+        }
+      }
+      written: false
+      written: true
+      custom_kernel {
+        name: "loop_transpose_fusion"
+        kernel_spec {
+          cubin {
+            data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\310\024\000\000\000\000\000\000H\017\000\000\000\000\000\000\002d\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion\000.nv.capmerc.text.loop_transpose_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\271\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\350\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000U\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000k\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\002\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 13.2, V13.2.51\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills  -arch sm_100a \000\000\014\000\000\000\010\000\000\000\350\003\000\000NVIDIA Corp\000\002\000d\000\204\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\002\t\001\000\002\002\001\000\002\005\005\000\003\007\001\001\002\003\000\000\004\013\010\000\t\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\t\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\001\000\000\300\017\000\000\000\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004Z \000\212\235\"\244\261\235\024m\000\264*\363\367X\003\245\',!0\311\036\307\217\017\014Il\n/\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000D\000\000\000\000\000\000\000=\000\001\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\271\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\350\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000U\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\373\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000_\002\000\000\000\000\000\000\217\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\317\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000H\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\374\005\000\000\000\000\000\000 \000\000\000\000\000\000\000\005\000\000\000\010\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\006\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000y\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000`\006\000\000\000\000\000\000p\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\355\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\010\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\271\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\033\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000?\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\220\014\000\000\000\000\000\000\262\000\000\000\000\000\000\000\025\000\000\000\010\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000f\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000B\r\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000{\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\264\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\214\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\330\r\000\000\000\000\000\000|\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000X\016\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\315\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000p\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000p\016\000\000\000\000\000\000\330\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\310\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\310\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\200\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000"
+          }
+          arity: 2
+          kernel_name: "loop_transpose_fusion"
+          kernel_args_packing_spec {
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 1
+              }
+            }
+          }
+        }
+        block_dims {
+          coordinates {
+            x: 1
+            y: 1
+            z: 1
+          }
+        }
+        thread_dims {
+          coordinates {
+            x: 4
+            y: 1
+            z: 1
+          }
+        }
+      }
+      use_pdl: true
+      tma_metadata {
+      }
+    }
+  }
+  executable_abi_version {
+    platform_name: "CUDA"
+    cuda_platform_version {
+      cuda_toolkit_version: "13.2.0"
+      cudnn_version: "9.23.0"
+      cub_version: "3.1.2"
+    }
+  }
+  buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n             parameter allocation:         0B\n              constant allocation:         0B\n        maybe_live_out allocation:        48B\n     preallocated temp allocation:       520B\n  preallocated temp fragmentation:       240B (46.15%)\n                 total allocation:       568B\n              total fragmentation:       504B (88.73%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n"
+  gpu_topology {
+    platform_version: "Compute Capability 10.0a"
+    num_partitions: 1
+    num_hosts_per_partition: 1
+    num_devices_per_host: 2
+    gpu_target_config {
+      gpu_device_info {
+        threads_per_block_limit: 1024
+        threads_per_warp: 32
+        shared_memory_per_block: 49152
+        shared_memory_per_core: 233472
+        threads_per_core_limit: 2048
+        core_count: 148
+        fpus_per_core: 128
+        block_dim_limit_x: 2147483647
+        block_dim_limit_y: 65535
+        block_dim_limit_z: 65535
+        memory_bandwidth: 7672320000000
+        l2_cache_size: 132644864
+        clock_rate_ghz: 1.965
+        device_memory_size: 191514411008
+        shared_memory_per_block_optin: 232448
+        cuda_compute_capability {
+          major: 10
+          feature_extension: ACCELERATED_FEATURES
+        }
+        registers_per_core_limit: 65536
+        registers_per_block_limit: 65536
+        scalar_unit_description {
+          rate_infos {
+            key: 4
+            value {
+              clock_rate_ghz: 1.965
+              units_per_core: 64
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 8
+            value {
+              clock_rate_ghz: 1.965
+              units_per_core: 64
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 10
+            value {
+              clock_rate_ghz: 1.965
+              units_per_core: 128
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 11
+            value {
+              clock_rate_ghz: 1.965
+              units_per_core: 128
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 12
+            value {
+              clock_rate_ghz: 1.965
+              units_per_core: 64
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 16
+            value {
+              clock_rate_ghz: 1.965
+              units_per_core: 128
+              ops_per_clock: 1
+            }
+          }
+        }
+        matrix_unit_description {
+          rate_infos {
+            key: 2
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 6
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 10
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 11
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 512
+            }
+          }
+          rate_infos {
+            key: 12
+            value {
+              clock_rate_ghz: 1.965
+              units_per_core: 4
+              ops_per_clock: 16
+            }
+          }
+          rate_infos {
+            key: 16
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 19
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 20
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 23
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 24
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 25
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 28
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 29
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 32
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 4096
+            }
+          }
+          rate_infos {
+            key: 33
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 35
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 36
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+        }
+        driver_version: "13.2.0"
+        kernel_mode_driver_version: "570.133.20"
+        runtime_version: "13.2.0"
+        compile_time_toolkit_version: "13.2.0"
+        dnn_version: "9.23.0"
+        cub_version: "3.1.2"
+        device_interconnect_info {
+          active_links: 18
+        }
+        device_vendor: "NVIDIA Corporation"
+        platform_version: "Compute Capability 10.0a"
+        pci_bus_id: "<undefined>"
+        name: "NVIDIA B200"
+        model_str: "sm_10.0a with 191514411008B RAM, 148 cores, 1965000KHz clock, 3996000KHz mem clock, 132644864B L2$"
+        numa_node: -1
+        thread_dim_limit_x: 1024
+        thread_dim_limit_y: 1024
+        thread_dim_limit_z: 64
+        device_address_bits: 64
+        pcie_bandwidth: 68719476736
+        ecc_enabled: true
+        mem_clock_ghz: 3.996
+        reserved_shared_memory_per_block: 1024
+        max_blocks_per_multiprocessor: 32
+        collective_memory_granularity: 2097152
+      }
+      platform_name: "CUDA"
+      device_description_str: "NVIDIA B200"
+      runtime_version {
+        major: 13
+        minor: 2
+      }
+    }
+    host_target_machine_options {
+      triple: "x86_64-grtev4-linux-gnu"
+      cpu: "emeraldrapids"
+    }
+    num_devices_per_process: 2
+  }
+}
diff --git a/third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_b200/v4/ReduceScatter.pbtxt b/third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_b200/v4/ReduceScatter.pbtxt
new file mode 100644
index 00000000000..1168bf48fa8
--- /dev/null
+++ b/third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_b200/v4/ReduceScatter.pbtxt
@@ -0,0 +1,1901 @@
+# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto
+# proto-message: xla.HumanReadableAotExecutable
+executable_and_options {
+  compile_options {
+    executable_build_options {
+      device_ordinal: -1
+      result_layout {
+        element_type: U32
+        dimensions: 4
+        layout {
+          minor_to_major: 0
+          tail_padding_alignment_in_elements: 1
+        }
+        is_dynamic_dimension: false
+      }
+      debug_options {
+        xla_disable_hlo_passes: "constant_folding"
+        xla_backend_optimization_level: 3
+        xla_eliminate_hlo_implicit_broadcast: true
+        xla_cpu_multi_thread_eigen: true
+        xla_gpu_cuda_data_dir: "./cuda_sdk_lib"
+        xla_llvm_enable_alias_scope_metadata: true
+        xla_llvm_enable_noalias_metadata: true
+        xla_llvm_enable_invariant_load_metadata: true
+        xla_llvm_disable_expensive_passes: false
+        xla_cpu_use_onednn: false
+        xla_cpu_enable_fast_math: false
+        xla_gpu_enable_fast_min_max: false
+        xla_force_host_platform_device_count: 1
+        xla_gpu_enable_dynamic_slice_fusion: true
+        xla_hlo_evaluator_use_fast_path: true
+        xla_dump_hlo_as_html: false
+        xla_cpu_fast_math_honor_nans: true
+        xla_cpu_fast_math_honor_infs: true
+        xla_allow_excess_precision: true
+        xla_gpu_autotune_level: 4
+        xla_cpu_fast_math_honor_division: true
+        xla_cpu_fast_math_honor_functions: true
+        xla_dump_include_timestamp: false
+        xla_dump_max_hlo_modules: -1
+        xla_cpu_enable_xprof_traceme: false
+        xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false
+        xla_cpu_enable_fast_min_max: true
+        xla_multiheap_size_constraint_per_heap: -1
+        xla_dump_module_metadata: false
+        xla_gpu_force_compilation_parallelism: 1
+        xla_dump_fusion_visualization: false
+        xla_gpu_all_reduce_combine_threshold_bytes: 31457287
+        xla_gpu_nccl_termination_timeout_seconds: -1
+        xla_dump_hlo_as_long_text: true
+        xla_gpu_enable_cublaslt: true
+        xla_gpu_shape_checks: RUNTIME
+        xla_gpu_use_runtime_fusion: false
+        xla_dump_latency_hiding_schedule: false
+        xla_dump_enable_mlir_pretty_form: true
+        xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP
+        xla_gpu_enable_triton_gemm: true
+        xla_gpu_enable_cudnn_int8x32_convolution_reordering: true
+        xla_gpu_triton_gemm_any: true
+        xla_gpu_enable_while_loop_reduce_scatter_code_motion: false
+        xla_gpu_collective_inflation_factor: 1
+        xla_gpu_graph_min_graph_size: 5
+        xla_gpu_enable_reassociation_for_converted_ar: true
+        xla_gpu_pgle_profile_file_or_directory_path: ""
+        xla_gpu_all_gather_combine_threshold_bytes: 31457287
+        xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287
+        xla_gpu_enable_highest_priority_async_stream: true
+        xla_gpu_exhaustive_tiling_search: false
+        xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0
+        xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1
+        xla_gpu_redzone_padding_bytes: 8388608
+        xla_gpu_copy_insertion_use_region_analysis: false
+        xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807
+        xla_gpu_collect_cost_model_stats: false
+        xla_gpu_enable_pipelined_p2p: false
+        xla_gpu_cublas_fallback: true
+        xla_gpu_enable_while_loop_double_buffering: false
+        xla_gpu_filter_kernels_spilling_registers_on_autotuning: true
+        xla_debug_buffer_assignment_show_max: 15
+        xla_detailed_logging: true
+        xla_enable_dumping: true
+        xla_gpu_enable_all_gather_combine_by_dim: false
+        xla_gpu_enable_analytical_latency_estimator: false
+        xla_gpu_llvm_verification_level: 0
+        xla_gpu_enable_reduce_scatter_combine_by_dim: false
+        xla_gpu_enable_command_buffer: COLLECTIVES_KERNEL
+        xla_gpu_enable_command_buffer: CONDITIONAL
+        xla_gpu_enable_command_buffer: CUBLAS
+        xla_gpu_enable_command_buffer: CUBLASLT
+        xla_gpu_enable_command_buffer: CUDNN
+        xla_gpu_enable_command_buffer: CUSTOM_CALL
+        xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION
+        xla_gpu_enable_command_buffer: FUSION
+        xla_gpu_enable_cub_radix_sort: true
+        xla_gpu_memory_limit_slop_factor: 95
+        xla_gpu_target_config_filename: ""
+        xla_gpu_enable_cudnn_layer_norm: false
+        xla_gpu_threshold_for_windowed_einsum_mib: 100000
+        xla_gpu_enable_nccl_user_buffers: false
+        xla_gpu_enable_libnvptxcompiler: true
+        xla_gpu_enable_nccl_comm_splitting: true
+        xla_gpu_nccl_collective_max_nchannels: 0
+        xla_gpu_nccl_p2p_max_nchannels: 0
+        xla_gpu_nccl_init_max_rank_per_root_ratio: 0
+        xla_gpu_multi_streamed_windowed_einsum: true
+        xla_gpu_gemm_rewrite_size_threshold: 100
+        xla_gpu_require_complete_aot_autotune_results: false
+        xla_gpu_cudnn_gemm_fusion_level: 0
+        xla_gpu_use_memcpy_local_p2p: false
+        xla_gpu_autotune_max_solutions: 0
+        xla_dump_large_constants: false
+        xla_gpu_verify_triton_fusion_numerics: false
+        xla_reduce_window_rewrite_base_length: 16
+        xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL
+        xla_gpu_enable_host_memory_offloading: false
+        xla_gpu_exclude_nondeterministic_ops: true
+        xla_gpu_nccl_terminate_on_error: false
+        xla_gpu_shard_autotuning: true
+        xla_gpu_enable_approx_costly_collectives: false
+        xla_cpu_prefer_vector_width: 256
+        xla_gpu_per_fusion_autotune_cache_dir: ""
+        xla_cmd_buffer_trace_cache_size: 16
+        xla_gpu_temp_buffer_use_separate_color: false
+        xla_syntax_sugar_async_ops: false
+        xla_gpu_autotune_gemm_rtol: 0.1
+        xla_enable_command_buffers_during_profiling: false
+        xla_gpu_cudnn_gemm_max_plans: 5
+        xla_cpu_parallel_codegen_split_count: 32
+        xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE
+        xla_gpu_executable_warn_stuck_timeout_seconds: 10
+        xla_gpu_executable_terminate_timeout_seconds: 30
+        xla_gpu_experimental_disable_binary_libraries: false
+        xla_gpu_dot_merger_threshold_mb: 64
+        xla_cpu_max_isa: ""
+        xla_gpu_experimental_enable_fusion_block_level_rewriter: false
+        xla_enable_fast_math: false
+        xla_gpu_experimental_parallel_collective_overlap_limit: 1
+        xla_cpu_copy_insertion_use_region_analysis: false
+        xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1
+        xla_gpu_experimental_enable_triton_heroless_priority_fusion: false
+        xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN
+        xla_gpu_experimental_stream_annotation: true
+        xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO
+        xla_pjrt_allow_auto_layout_in_hlo: false
+        xla_gpu_enable_scatter_determinism_expander: false
+        xla_gpu_require_exclusive_lock: false
+        xla_gpu_generate_debug_info: false
+        xla_gpu_generate_line_info: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false
+        xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE
+        xla_gpu_fail_ptx_compilation_on_register_spilling: false
+        xla_gpu_enable_analytical_sol_latency_estimator: true
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_prep_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_size_bytes"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "gpus_per_node"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nccl_op_launch_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nic_speed_gbps"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "rtt_us"
+          value: "-1"
+        }
+        xla_cpu_use_xnnpack: true
+        xla_gpu_experimental_pack_dot_operands_along_k_dimension: true
+        xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false
+        xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED
+        xla_gpu_experimental_enable_heuristic_collective_combining: false
+        xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false
+        xla_hlo_pass_fix_detect_cycles: false
+        xla_cpu_generate_unique_c_style_kernel_entry_points: false
+        xla_gpu_experimental_collective_cse_distance_threshold: 0
+        xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true
+        xla_gpu_experimental_collective_perf_table_path: ""
+        xla_gpu_collective_permute_combine_threshold_bytes: 31457287
+        xla_unsupported_crash_on_hlo_pass_noop_change: false
+        xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false
+        xla_dump_full_hlo_config: true
+        xla_gpu_unsupported_enable_triton_multi_output_fusion: false
+        xla_gpu_experimental_matmul_perf_table_path: ""
+        xla_gpu_unsupported_enable_all_reduce_decomposer: false
+        xla_gpu_nccl_blocking_communicators: true
+        xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20
+        xla_gpu_first_collective_call_terminate_timeout_seconds: 40
+        xla_gpu_nccl_async_execution: false
+        xla_cpu_emitter_verification_level: 1
+        xla_early_exit_with_layouts: false
+        xla_gpu_experimental_use_ragged_dot_fusion: false
+        xla_detect_unstable_reductions: DETECTION_MODE_NONE
+        xla_gpu_command_buffer_scheduling_mode: LHS
+        xla_gpu_experimental_enable_nccl_symmetric_buffers: false
+        xla_disable_automatic_host_compute_offload: false
+        xla_gpu_experimental_enable_fusion_autotuner: true
+        xla_gpu_experimental_scaled_dot_with_triton: true
+        xla_gpu_command_buffer_unroll_loops: false
+        xla_cpu_experimental_onednn_custom_call: false
+        xla_gpu_experimental_enable_checksum_tracing_on_thunks: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true
+        xla_cpu_collective_call_terminate_timeout_seconds: 40
+        xla_cpu_collective_call_warn_stuck_seconds: 20
+        xla_keep_shardings_after_spmd: false
+        xla_gpu_experimental_enable_triton_warp_specialization: false
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE
+        xla_cpu_enable_platform_dependent_math: true
+        xla_gpu_detect_nan: DETECTION_MODE_NONE
+        xla_gpu_detect_inf: DETECTION_MODE_NONE
+        xla_enable_enzyme_comms_opt: false
+        xla_gpu_experimental_enable_buffer_saver_on_thunks: false
+        xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE
+        xla_enable_scoped_logging_timers: true
+        xla_gpu_executable_embed_debug_info: true
+        xla_cpu_collective_timeout_seconds: 1800
+        xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false
+        xla_gpu_enable_pipelined_host_offloading: false
+        xla_gpu_default_to_alg_dot_bf16_bf16_f32: false
+        xla_gpu_experimental_autotune_backends: CUDNN
+        xla_gpu_experimental_autotune_backends: TRITON
+        xla_gpu_experimental_autotune_backends: CUBLASLT
+        xla_gpu_experimental_autotune_backends: HIPBLASLT
+        xla_gpu_experimental_autotune_backends: MIOPEN
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER
+        xla_gpu_experimental_autotune_backends: NATIVE_EMITTER
+        xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER
+        xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION
+        xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: YNNPACK
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU
+        xla_gpu_experimental_autotune_backends: TENSOR_IR
+        xla_gpu_experimental_all_fusions_with_triton: false
+        xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT
+        xla_gpu_rocm_max_trace_events: 4194304
+        xla_gpu_execution_terminate_timeout: "inf"
+        xla_enable_hlo_sharding_v3: true
+        xla_gpu_log_minmax: false
+        xla_gpu_print_compilation_stats: false
+        xla_gpu_experimental_enable_conv_fusion: false
+        xla_gpu_experimental_enable_tiling_propagation: true
+        xla_gpu_execution_progress_tracking: 0
+        xla_gpu_experimental_max_unroll_factor: 32
+        xla_gpu_enable_pdl: true
+        xla_gpu_collective_combine_threshold_count: 2048
+        xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true
+        xla_recognize_reduction_optimization_level: 0
+        xla_gpu_experimental_parallel_async_compute_limit: 2
+        xla_dump_buffer_assignment_analysis: true
+        xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE
+        xla_gpu_experimental_enable_collective_multi_streaming: false
+        xla_gpu_collective_permute_connected_components: false
+        xla_gpu_experimental_force_split_k: 0
+        xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_experimental_gemm_fusion_v2: true
+        xla_gpu_fusion_autotune_top_k_configs: 1
+        xla_gpu_native_emitter_tune_unroll_factor_for_loops: false
+        xla_enable_rgv3_materialization: true
+        xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK
+        xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false
+        xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_executable_num_compute_streams: 0
+        xla_gpu_executable_num_communication_streams: 0
+        xla_gpu_experimental_collective_start_as_early_as_possible: false
+        xla_cpu_experimental_enable_tiling_propagation: true
+        xla_gpu_enable_dus_accumulator_zero_init_elimination: false
+        xla_gpu_experimental_thunk_buffer_debug_module_outputs: false
+        xla_cpu_ftz: true
+        xla_gpu_enable_allocator_spatial_partitioning: true
+        xla_dump_compact_gte: true
+        xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES
+        xla_dump_hlo_as_riegeli: false
+        xla_gpu_experimental_use_ragged_dot_grouped_gemm: true
+        xla_dump_hlo_to_subfolder: false
+        xla_gpu_enable_pdl_launch: true
+        xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED
+        xla_gpu_enable_nccl_user_buffers_in_default_space: false
+        xla_sdy_export_all_reduce_scatter: false
+        xla_gpu_enable_gxl_ragged_all_to_all: false
+        xla_gpu_gxl_scratch_size_bytes: 67108864
+        xla_gpu_use_new_autotune_cache_format: true
+        xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE
+        xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON
+        xla_spmd_enable_dynamic_slice_collective_broadcast: false
+        xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false
+        xla_gpu_async_copy_min_bytes: -1
+        xla_gpu_blas_max_algorithms: 0
+        xla_gpu_experimental_enable_selective_memcpy_overlap: false
+        xla_gpu_experimental_enable_same_shape_multi_output_fusion: false
+        xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1
+        xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1
+        xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240
+        xla_cpu_use_new_xtile_lowering: true
+        xla_gpu_trace_annotation_level: 0
+        xla_gpu_experimental_cost_model_gemm_tiling_default: false
+        xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false
+        xla_gpu_all_reduce_splitter_ignore_profitability_check: false
+        xla_gpu_enable_cupti_multi_subscriber: true
+        xla_enable_hlo_modules_upload: true
+        xla_gpu_collective_domain_assignment: ""
+        xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1
+        xla_compile_all_supported_configs: false
+        xla_gpu_unsupported_use_cross_host_one_shot_kernel: ALLCOLLECTIVES
+        xla_deduplicate_backend_configs_min_size: 9223372036854775807
+        xla_force_config: ""
+        xla_candidate_configs_file: ""
+        xla_gpu_enable_persistent_symmetric_memory: false
+        xla_gpu_experimental_enable_raft_for_stable_topk: false
+        xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND
+        xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET
+        xla_gpu_experimental_emit_collective_reduce: false
+        xla_gpu_experimental_vmm_disabled: false
+        xla_gpu_experimental_enable_dynamic_slice_table_offsets: false
+        xla_gpu_device_execution_terminate_timeout: "inf"
+        xla_gpu_topology_filename: ""
+        xla_gpu_experimental_enable_dynamic_slice_extended_offsets: false
+        xla_gpu_experimental_deterministic_ops_triton: DETERMINISTIC_OPS_TRITON_MODE_OFF
+        xla_gpu_experimental_enable_tensor_ir: false
+      }
+      num_replicas: 2
+      num_partitions: 1
+      device_assignment {
+        replica_count: 2
+        computation_count: 1
+        computation_devices {
+          replica_device_ids: 0
+          replica_device_ids: 1
+        }
+      }
+      allow_spmd_sharding_propagation_to_output: false
+      comp_envs {
+      }
+      allow_spmd_sharding_propagation_to_parameters: false
+      process_count: 1
+      memory_fitting_level: EFFORT_O2
+    }
+  }
+  pjrt_client_name: "PjRtStreamExecutorClient"
+}
+gpu_executable {
+  hlo_module_with_config {
+    hlo_module {
+      name: "test"
+      entry_computation_name: "main"
+      computations {
+        name: "add"
+        instructions {
+          name: "lhs"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "lhs"
+          }
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "rhs"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "rhs"
+          }
+          parameter_number: 1
+          id: 1
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "add.1"
+          opcode: "add"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "add.1"
+          }
+          id: 2
+          operand_ids: 0
+          operand_ids: 1
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          parameters {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          parameters {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          result {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          parameter_names: "lhs"
+          parameter_names: "rhs"
+        }
+        root_id: 2
+      }
+      computations {
+        name: "fused_select"
+        instructions {
+          name: "param_2.2"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "param_2.2"
+          }
+          parameter_number: 2
+          id: 8589934597
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "zero_2"
+          opcode: "constant"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+          }
+          literal {
+            shape {
+              element_type: U32
+              layout {
+                tail_padding_alignment_in_elements: 1
+              }
+            }
+            u32s: 0
+          }
+          id: 8589934598
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "p.2"
+          opcode: "compare"
+          shape {
+            element_type: PRED
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "p.2"
+          }
+          id: 8589934596
+          operand_ids: 8589934597
+          operand_ids: 8589934598
+          comparison_direction: "EQ"
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+          comparison_order: "TOTAL"
+        }
+        instructions {
+          name: "pb.2"
+          opcode: "broadcast"
+          shape {
+            element_type: PRED
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "pb.2"
+          }
+          id: 8589934593
+          operand_ids: 8589934596
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "param_0.1"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "param_0.1"
+          }
+          id: 8589934594
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "param_1.1"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "param_1.1"
+          }
+          parameter_number: 1
+          id: 8589934595
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "data.2"
+          opcode: "select"
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "data.2"
+          }
+          id: 8589934592
+          operand_ids: 8589934593
+          operand_ids: 8589934594
+          operand_ids: 8589934595
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          parameters {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          parameters {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          parameters {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          result {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          parameter_names: "param_0.1"
+          parameter_names: "param_1.1"
+          parameter_names: "param_2.2"
+        }
+        id: 2
+        root_id: 8589934592
+        is_fusion_computation: true
+      }
+      computations {
+        name: "main"
+        instructions {
+          name: "c0_1"
+          opcode: "constant"
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+          }
+          literal {
+            shape {
+              element_type: U32
+              dimensions: 8
+              layout {
+                minor_to_major: 0
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+            }
+            u32s: 1
+            u32s: 2
+            u32s: 3
+            u32s: 4
+            u32s: 5
+            u32s: 6
+            u32s: 7
+            u32s: 8
+          }
+          id: 4294967297
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "c1_1"
+          opcode: "constant"
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+          }
+          literal {
+            shape {
+              element_type: U32
+              dimensions: 8
+              layout {
+                minor_to_major: 0
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+            }
+            u32s: 10
+            u32s: 11
+            u32s: 12
+            u32s: 13
+            u32s: 14
+            u32s: 15
+            u32s: 16
+            u32s: 17
+          }
+          id: 4294967298
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "id.1"
+          opcode: "replica-id"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "id.1"
+          }
+          id: 4294967296
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "loop_select_fusion"
+          opcode: "fusion"
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "loop_select_fusion"
+          }
+          fusion_kind: "kLoop"
+          id: 4294967299
+          operand_ids: 4294967297
+          operand_ids: 4294967298
+          operand_ids: 4294967296
+          called_computation_ids: 2
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "ars.3"
+          opcode: "reduce-scatter"
+          shape {
+            element_type: U32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "reduce-scatter-start"
+          }
+          dimensions: 0
+          id: 4294967300
+          operand_ids: 4294967299
+          called_computation_ids: 0
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+          collective_device_list {
+          }
+          backend_config_payload {
+            id: 0
+          }
+        }
+        program_shape {
+          result {
+            element_type: U32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+        }
+        id: 1
+        root_id: 4294967300
+      }
+      host_program_shape {
+        result {
+          element_type: U32
+          dimensions: 4
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+      }
+      entry_computation_id: 1
+      schedule {
+        sequences {
+          key: 0
+          value {
+            instruction_ids: 1
+            instruction_ids: 0
+            instruction_ids: 2
+          }
+        }
+        sequences {
+          key: 1
+          value {
+            instruction_ids: 4294967296
+            instruction_ids: 4294967298
+            instruction_ids: 4294967297
+            instruction_ids: 4294967299
+            instruction_ids: 4294967300
+          }
+        }
+      }
+      input_output_alias {
+      }
+      device_assignment {
+        replica_count: 2
+        computation_count: 1
+        computation_devices {
+          replica_device_ids: 0
+          replica_device_ids: 1
+        }
+      }
+      buffer_donor {
+      }
+      frontend_attributes {
+        map {
+          key: "fingerprint_before_lhs"
+          value: "c36d8870be58d238935bebefbbc5b827"
+        }
+      }
+      payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}"
+    }
+    config {
+      entry_computation_layout {
+        result {
+          element_type: U32
+          dimensions: 4
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+      }
+      replica_count: 2
+      num_partitions: 1
+      intra_op_parallelism_threads: -1
+      debug_options {
+        xla_disable_hlo_passes: "constant_folding"
+        xla_backend_optimization_level: 3
+        xla_eliminate_hlo_implicit_broadcast: true
+        xla_cpu_multi_thread_eigen: true
+        xla_gpu_cuda_data_dir: "./cuda_sdk_lib"
+        xla_llvm_enable_alias_scope_metadata: true
+        xla_llvm_enable_noalias_metadata: true
+        xla_llvm_enable_invariant_load_metadata: true
+        xla_llvm_disable_expensive_passes: false
+        xla_cpu_use_onednn: false
+        xla_cpu_enable_fast_math: false
+        xla_gpu_enable_fast_min_max: false
+        xla_force_host_platform_device_count: 1
+        xla_gpu_enable_dynamic_slice_fusion: true
+        xla_hlo_evaluator_use_fast_path: true
+        xla_dump_hlo_as_html: false
+        xla_cpu_fast_math_honor_nans: true
+        xla_cpu_fast_math_honor_infs: true
+        xla_allow_excess_precision: true
+        xla_gpu_autotune_level: 4
+        xla_cpu_fast_math_honor_division: true
+        xla_cpu_fast_math_honor_functions: true
+        xla_dump_include_timestamp: false
+        xla_dump_max_hlo_modules: -1
+        xla_cpu_enable_xprof_traceme: false
+        xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false
+        xla_cpu_enable_fast_min_max: true
+        xla_multiheap_size_constraint_per_heap: -1
+        xla_dump_module_metadata: false
+        xla_gpu_force_compilation_parallelism: 1
+        xla_dump_fusion_visualization: false
+        xla_gpu_all_reduce_combine_threshold_bytes: 31457287
+        xla_gpu_nccl_termination_timeout_seconds: -1
+        xla_dump_hlo_as_long_text: true
+        xla_gpu_enable_cublaslt: true
+        xla_gpu_shape_checks: RUNTIME
+        xla_gpu_use_runtime_fusion: false
+        xla_dump_latency_hiding_schedule: false
+        xla_dump_enable_mlir_pretty_form: true
+        xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP
+        xla_gpu_enable_triton_gemm: true
+        xla_gpu_enable_cudnn_int8x32_convolution_reordering: true
+        xla_gpu_triton_gemm_any: true
+        xla_gpu_enable_while_loop_reduce_scatter_code_motion: false
+        xla_gpu_collective_inflation_factor: 1
+        xla_gpu_graph_min_graph_size: 5
+        xla_gpu_enable_reassociation_for_converted_ar: true
+        xla_gpu_pgle_profile_file_or_directory_path: ""
+        xla_gpu_all_gather_combine_threshold_bytes: 31457287
+        xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287
+        xla_gpu_enable_highest_priority_async_stream: true
+        xla_gpu_exhaustive_tiling_search: false
+        xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0
+        xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1
+        xla_gpu_redzone_padding_bytes: 8388608
+        xla_gpu_copy_insertion_use_region_analysis: false
+        xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807
+        xla_gpu_collect_cost_model_stats: false
+        xla_gpu_enable_pipelined_p2p: false
+        xla_gpu_cublas_fallback: true
+        xla_gpu_enable_while_loop_double_buffering: false
+        xla_gpu_filter_kernels_spilling_registers_on_autotuning: true
+        xla_debug_buffer_assignment_show_max: 15
+        xla_detailed_logging: true
+        xla_enable_dumping: true
+        xla_gpu_enable_all_gather_combine_by_dim: false
+        xla_gpu_enable_analytical_latency_estimator: false
+        xla_gpu_llvm_verification_level: 0
+        xla_gpu_enable_reduce_scatter_combine_by_dim: false
+        xla_gpu_enable_command_buffer: COLLECTIVES_KERNEL
+        xla_gpu_enable_command_buffer: CONDITIONAL
+        xla_gpu_enable_command_buffer: CUBLAS
+        xla_gpu_enable_command_buffer: CUBLASLT
+        xla_gpu_enable_command_buffer: CUDNN
+        xla_gpu_enable_command_buffer: CUSTOM_CALL
+        xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION
+        xla_gpu_enable_command_buffer: FUSION
+        xla_gpu_enable_cub_radix_sort: true
+        xla_gpu_memory_limit_slop_factor: 95
+        xla_gpu_target_config_filename: ""
+        xla_gpu_enable_cudnn_layer_norm: false
+        xla_gpu_threshold_for_windowed_einsum_mib: 100000
+        xla_gpu_enable_nccl_user_buffers: false
+        xla_gpu_enable_libnvptxcompiler: true
+        xla_gpu_enable_nccl_comm_splitting: true
+        xla_gpu_nccl_collective_max_nchannels: 0
+        xla_gpu_nccl_p2p_max_nchannels: 0
+        xla_gpu_nccl_init_max_rank_per_root_ratio: 0
+        xla_gpu_multi_streamed_windowed_einsum: true
+        xla_gpu_gemm_rewrite_size_threshold: 100
+        xla_gpu_require_complete_aot_autotune_results: false
+        xla_gpu_cudnn_gemm_fusion_level: 0
+        xla_gpu_use_memcpy_local_p2p: false
+        xla_gpu_autotune_max_solutions: 0
+        xla_dump_large_constants: false
+        xla_gpu_verify_triton_fusion_numerics: false
+        xla_reduce_window_rewrite_base_length: 16
+        xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL
+        xla_gpu_enable_host_memory_offloading: false
+        xla_gpu_exclude_nondeterministic_ops: true
+        xla_gpu_nccl_terminate_on_error: false
+        xla_gpu_shard_autotuning: true
+        xla_gpu_enable_approx_costly_collectives: false
+        xla_cpu_prefer_vector_width: 256
+        xla_gpu_per_fusion_autotune_cache_dir: ""
+        xla_cmd_buffer_trace_cache_size: 16
+        xla_gpu_temp_buffer_use_separate_color: false
+        xla_syntax_sugar_async_ops: false
+        xla_gpu_autotune_gemm_rtol: 0.1
+        xla_enable_command_buffers_during_profiling: false
+        xla_gpu_cudnn_gemm_max_plans: 5
+        xla_cpu_parallel_codegen_split_count: 32
+        xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE
+        xla_gpu_executable_warn_stuck_timeout_seconds: 10
+        xla_gpu_executable_terminate_timeout_seconds: 30
+        xla_gpu_experimental_disable_binary_libraries: false
+        xla_gpu_dot_merger_threshold_mb: 64
+        xla_cpu_max_isa: ""
+        xla_gpu_experimental_enable_fusion_block_level_rewriter: false
+        xla_enable_fast_math: false
+        xla_gpu_experimental_parallel_collective_overlap_limit: 1
+        xla_cpu_copy_insertion_use_region_analysis: false
+        xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1
+        xla_gpu_experimental_enable_triton_heroless_priority_fusion: false
+        xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN
+        xla_gpu_experimental_stream_annotation: true
+        xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO
+        xla_pjrt_allow_auto_layout_in_hlo: false
+        xla_gpu_enable_scatter_determinism_expander: false
+        xla_gpu_require_exclusive_lock: false
+        xla_gpu_generate_debug_info: false
+        xla_gpu_generate_line_info: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false
+        xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE
+        xla_gpu_fail_ptx_compilation_on_register_spilling: false
+        xla_gpu_enable_analytical_sol_latency_estimator: true
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_prep_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_size_bytes"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "gpus_per_node"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nccl_op_launch_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nic_speed_gbps"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "rtt_us"
+          value: "-1"
+        }
+        xla_cpu_use_xnnpack: true
+        xla_gpu_experimental_pack_dot_operands_along_k_dimension: true
+        xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false
+        xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED
+        xla_gpu_experimental_enable_heuristic_collective_combining: false
+        xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false
+        xla_hlo_pass_fix_detect_cycles: false
+        xla_cpu_generate_unique_c_style_kernel_entry_points: false
+        xla_gpu_experimental_collective_cse_distance_threshold: 0
+        xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true
+        xla_gpu_experimental_collective_perf_table_path: ""
+        xla_gpu_collective_permute_combine_threshold_bytes: 31457287
+        xla_unsupported_crash_on_hlo_pass_noop_change: false
+        xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false
+        xla_dump_full_hlo_config: true
+        xla_gpu_unsupported_enable_triton_multi_output_fusion: false
+        xla_gpu_experimental_matmul_perf_table_path: ""
+        xla_gpu_unsupported_enable_all_reduce_decomposer: false
+        xla_gpu_nccl_blocking_communicators: true
+        xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20
+        xla_gpu_first_collective_call_terminate_timeout_seconds: 40
+        xla_gpu_nccl_async_execution: false
+        xla_cpu_emitter_verification_level: 1
+        xla_early_exit_with_layouts: false
+        xla_gpu_experimental_use_ragged_dot_fusion: false
+        xla_detect_unstable_reductions: DETECTION_MODE_NONE
+        xla_gpu_command_buffer_scheduling_mode: LHS
+        xla_gpu_experimental_enable_nccl_symmetric_buffers: false
+        xla_disable_automatic_host_compute_offload: false
+        xla_gpu_experimental_enable_fusion_autotuner: true
+        xla_gpu_experimental_scaled_dot_with_triton: true
+        xla_gpu_command_buffer_unroll_loops: false
+        xla_cpu_experimental_onednn_custom_call: false
+        xla_gpu_experimental_enable_checksum_tracing_on_thunks: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true
+        xla_cpu_collective_call_terminate_timeout_seconds: 40
+        xla_cpu_collective_call_warn_stuck_seconds: 20
+        xla_keep_shardings_after_spmd: false
+        xla_gpu_experimental_enable_triton_warp_specialization: false
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE
+        xla_cpu_enable_platform_dependent_math: true
+        xla_gpu_detect_nan: DETECTION_MODE_NONE
+        xla_gpu_detect_inf: DETECTION_MODE_NONE
+        xla_enable_enzyme_comms_opt: false
+        xla_gpu_experimental_enable_buffer_saver_on_thunks: false
+        xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE
+        xla_enable_scoped_logging_timers: true
+        xla_gpu_executable_embed_debug_info: true
+        xla_cpu_collective_timeout_seconds: 1800
+        xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false
+        xla_gpu_enable_pipelined_host_offloading: false
+        xla_gpu_default_to_alg_dot_bf16_bf16_f32: false
+        xla_gpu_experimental_autotune_backends: CUDNN
+        xla_gpu_experimental_autotune_backends: TRITON
+        xla_gpu_experimental_autotune_backends: CUBLASLT
+        xla_gpu_experimental_autotune_backends: HIPBLASLT
+        xla_gpu_experimental_autotune_backends: MIOPEN
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER
+        xla_gpu_experimental_autotune_backends: NATIVE_EMITTER
+        xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER
+        xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION
+        xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: YNNPACK
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU
+        xla_gpu_experimental_autotune_backends: TENSOR_IR
+        xla_gpu_experimental_all_fusions_with_triton: false
+        xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT
+        xla_gpu_rocm_max_trace_events: 4194304
+        xla_gpu_execution_terminate_timeout: "inf"
+        xla_enable_hlo_sharding_v3: true
+        xla_gpu_log_minmax: false
+        xla_gpu_print_compilation_stats: false
+        xla_gpu_experimental_enable_conv_fusion: false
+        xla_gpu_experimental_enable_tiling_propagation: true
+        xla_gpu_execution_progress_tracking: 0
+        xla_gpu_experimental_max_unroll_factor: 32
+        xla_gpu_enable_pdl: true
+        xla_gpu_collective_combine_threshold_count: 2048
+        xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true
+        xla_recognize_reduction_optimization_level: 0
+        xla_gpu_experimental_parallel_async_compute_limit: 2
+        xla_dump_buffer_assignment_analysis: true
+        xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE
+        xla_gpu_experimental_enable_collective_multi_streaming: false
+        xla_gpu_collective_permute_connected_components: false
+        xla_gpu_experimental_force_split_k: 0
+        xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_experimental_gemm_fusion_v2: true
+        xla_gpu_fusion_autotune_top_k_configs: 1
+        xla_gpu_native_emitter_tune_unroll_factor_for_loops: false
+        xla_enable_rgv3_materialization: true
+        xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK
+        xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false
+        xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_executable_num_compute_streams: 0
+        xla_gpu_executable_num_communication_streams: 0
+        xla_gpu_experimental_collective_start_as_early_as_possible: false
+        xla_cpu_experimental_enable_tiling_propagation: true
+        xla_gpu_enable_dus_accumulator_zero_init_elimination: false
+        xla_gpu_experimental_thunk_buffer_debug_module_outputs: false
+        xla_cpu_ftz: true
+        xla_gpu_enable_allocator_spatial_partitioning: true
+        xla_dump_compact_gte: true
+        xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES
+        xla_dump_hlo_as_riegeli: false
+        xla_gpu_experimental_use_ragged_dot_grouped_gemm: true
+        xla_dump_hlo_to_subfolder: false
+        xla_gpu_enable_pdl_launch: true
+        xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED
+        xla_gpu_enable_nccl_user_buffers_in_default_space: false
+        xla_sdy_export_all_reduce_scatter: false
+        xla_gpu_enable_gxl_ragged_all_to_all: false
+        xla_gpu_gxl_scratch_size_bytes: 67108864
+        xla_gpu_use_new_autotune_cache_format: true
+        xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE
+        xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON
+        xla_spmd_enable_dynamic_slice_collective_broadcast: false
+        xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false
+        xla_gpu_async_copy_min_bytes: -1
+        xla_gpu_blas_max_algorithms: 0
+        xla_gpu_experimental_enable_selective_memcpy_overlap: false
+        xla_gpu_experimental_enable_same_shape_multi_output_fusion: false
+        xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1
+        xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1
+        xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240
+        xla_cpu_use_new_xtile_lowering: true
+        xla_gpu_trace_annotation_level: 0
+        xla_gpu_experimental_cost_model_gemm_tiling_default: false
+        xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false
+        xla_gpu_all_reduce_splitter_ignore_profitability_check: false
+        xla_gpu_enable_cupti_multi_subscriber: true
+        xla_enable_hlo_modules_upload: true
+        xla_gpu_collective_domain_assignment: ""
+        xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1
+        xla_compile_all_supported_configs: false
+        xla_gpu_unsupported_use_cross_host_one_shot_kernel: ALLCOLLECTIVES
+        xla_deduplicate_backend_configs_min_size: 9223372036854775807
+        xla_force_config: ""
+        xla_candidate_configs_file: ""
+        xla_gpu_enable_persistent_symmetric_memory: false
+        xla_gpu_experimental_enable_raft_for_stable_topk: false
+        xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND
+        xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET
+        xla_gpu_experimental_emit_collective_reduce: false
+        xla_gpu_experimental_vmm_disabled: false
+        xla_gpu_experimental_enable_dynamic_slice_table_offsets: false
+        xla_gpu_device_execution_terminate_timeout: "inf"
+        xla_gpu_topology_filename: ""
+        xla_gpu_experimental_enable_dynamic_slice_extended_offsets: false
+        xla_gpu_experimental_deterministic_ops_triton: DETERMINISTIC_OPS_TRITON_MODE_OFF
+        xla_gpu_experimental_enable_tensor_ir: false
+      }
+      static_device_assignment {
+        replica_count: 2
+        computation_count: 1
+        computation_devices {
+          replica_device_ids: 0
+          replica_device_ids: 1
+        }
+      }
+      allow_spmd_sharding_propagation_to_output: false
+      allow_spmd_sharding_propagation_to_parameters: false
+      sharding_config {
+      }
+      memory_fitting_level: EFFORT_O2
+      schedule_config {
+      }
+      partition_size: 2
+    }
+  }
+  buffer_assignment {
+    logical_buffers {
+      size: 4
+      defined_at {
+      }
+    }
+    logical_buffers {
+      id: 1
+      size: 4
+      defined_at {
+        instruction_id: 1
+      }
+    }
+    logical_buffers {
+      id: 2
+      size: 4
+      defined_at {
+        instruction_id: 2
+      }
+    }
+    logical_buffers {
+      id: 10
+      size: 32
+      defined_at {
+        instruction_id: 4294967297
+      }
+    }
+    logical_buffers {
+      id: 11
+      size: 32
+      defined_at {
+        instruction_id: 4294967298
+      }
+    }
+    logical_buffers {
+      id: 12
+      size: 4
+      defined_at {
+        instruction_id: 4294967296
+      }
+    }
+    logical_buffers {
+      id: 13
+      size: 32
+      defined_at {
+        instruction_id: 4294967299
+      }
+    }
+    logical_buffers {
+      id: 14
+      size: 16
+      defined_at {
+        instruction_id: 4294967300
+      }
+    }
+    buffer_allocations {
+      size: 32
+      assigned {
+        logical_buffer_id: 10
+        size: 32
+        element_type: U32
+      }
+      is_constant: true
+    }
+    buffer_allocations {
+      index: 1
+      size: 32
+      assigned {
+        logical_buffer_id: 11
+        size: 32
+        element_type: U32
+      }
+      is_constant: true
+    }
+    buffer_allocations {
+      index: 2
+      size: 16
+      maybe_live_out: true
+      assigned {
+        logical_buffer_id: 12
+        size: 4
+        element_type: U32
+      }
+      assigned {
+        logical_buffer_id: 14
+        size: 16
+        element_type: U32
+      }
+    }
+    buffer_allocations {
+      index: 3
+      size: 4
+      is_thread_local: true
+      assigned {
+        logical_buffer_id: 2
+        size: 4
+        element_type: U32
+      }
+    }
+    buffer_allocations {
+      index: 4
+      size: 4
+      is_thread_local: true
+      assigned {
+        size: 4
+        element_type: U32
+      }
+    }
+    buffer_allocations {
+      index: 5
+      size: 4
+      is_thread_local: true
+      assigned {
+        logical_buffer_id: 1
+        size: 4
+        element_type: U32
+      }
+    }
+    buffer_allocations {
+      index: 6
+      size: 32
+      assigned {
+        logical_buffer_id: 13
+        size: 32
+        element_type: U32
+      }
+    }
+    heap_simulator_traces {
+      events {
+        buffer_id: 13
+        computation_name: "main"
+        instruction_name: "loop_select_fusion"
+      }
+      events {
+        kind: FREE
+        buffer_id: 13
+        computation_name: "main"
+        instruction_name: "loop_select_fusion"
+      }
+      whole_module_simulation: true
+      buffer_allocation_index: 6
+    }
+    peak_buffers {
+      logical_buffer_ids: 10
+    }
+    peak_buffers {
+      index: 1
+      logical_buffer_ids: 11
+    }
+    peak_buffers {
+      index: 2
+      logical_buffer_ids: 14
+    }
+    peak_buffers {
+      index: 3
+      logical_buffer_ids: 2
+    }
+    peak_buffers {
+      index: 4
+      logical_buffer_ids: 0
+    }
+    peak_buffers {
+      index: 5
+      logical_buffer_ids: 1
+    }
+    peak_buffers {
+      index: 6
+      logical_buffer_ids: 13
+    }
+  }
+  binary: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000(\014\000\000\000\000\000\000(\007\000\000\000\000\000\000\002d\000\006@\0008\000\004\000@\000\024\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.nv.shared.reserved.0\000.nv.global\000.nv.constant4\000\000\000\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant4\000\000\000\000.nv.callgraph\000.nv.prototype\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.rela.nv.constant.pic\000.nv.merc.nv.constant.pic\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.global\000.nv.constant.pic\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant.pic\000.nv.callgraph\000.nv.prototype\000buffer_for_c1_1\000buffer_for_c0_1\000.nv.constant4\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\214\000\000\000 \240\014\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\003\000\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000=\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\001\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\035\001\000\000\021\000\r\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000-\001\000\000\021\000\r\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\002\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 13.2, V13.2.51\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills  -arch sm_100a \000\000\014\000\000\000\010\000\000\000\350\003\000\000NVIDIA Corp\000\002\000d\000\204\000\000\000\003_\001\001\002\t\001\000\002\002\001\000\002\005\005\000\003\007\001\001\002\003\000\000\004\013\010\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\010\000\000\000\000\000\000\000\002\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\003_\001\001\000\000\000\000\010\000\000\000\000\000\000\000\002\000\001\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\001\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\214\000\000\000 \240\022\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\003\000\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\270\000\000\000\003\000\021\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\001\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\035\001\000\000\035 \r\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000-\001\000\000\035 \r\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000n\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\322\001\000\000\000\000\000\000K\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\003\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\217\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000(\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\334\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\005\000\000\000\010\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\004\000\000\000\000\000\000\004\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\307\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\261\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\005\000\000\000\000\000\0000\000\000\000\000\000\000\000\003\000\000\000\013\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000~\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\005\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000s\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\343\000\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000(\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\370\000\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\000\004\000\000\000\000\000\000\000\023\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\010\006\000\000\000\000\000\0000\000\000\000\000\000\000\000\023\000\000\000\021\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\'\001\000\000}\000\000p\002\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000p\005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0008\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000^\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0008\006\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000(\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000(\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000p\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\000\000\000\000\000\000\000\300\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000"
+  gpu_compute_capability {
+    cuda_compute_capability {
+      major: 10
+      feature_extension: ACCELERATED_FEATURES
+    }
+  }
+  module_name: "test"
+  program_shape {
+    result {
+      element_type: U32
+      dimensions: 4
+      layout {
+        minor_to_major: 0
+        tail_padding_alignment_in_elements: 1
+      }
+      is_dynamic_dimension: false
+    }
+  }
+  buffer_allocations {
+    values {
+      size: 32
+      is_constant: true
+    }
+    values {
+      index: 1
+      size: 32
+      is_constant: true
+    }
+    values {
+      index: 2
+      size: 16
+      maybe_live_out: true
+    }
+    values {
+      index: 3
+      size: 4
+      is_thread_local: true
+    }
+    values {
+      index: 4
+      size: 4
+      is_thread_local: true
+    }
+    values {
+      index: 5
+      size: 4
+      is_thread_local: true
+    }
+    values {
+      index: 6
+      size: 32
+    }
+  }
+  output_info_map {
+    shape_index {
+    }
+    output_info {
+      allocation_index: 2
+    }
+  }
+  constants {
+    symbol_name: "buffer_for_c1_1"
+    allocation_index: 1
+  }
+  constants {
+    symbol_name: "buffer_for_c0_1"
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "id.1"
+      thunk_id: 1
+    }
+    replica_id_thunk {
+      dest_buffer {
+        size: 4
+        buffer_allocation_index: 2
+        element_type: U32
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "loop_select_fusion"
+      thunk_id: 2
+    }
+    custom_kernel_thunk {
+      args {
+        slice {
+          size: 32
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          dimensions: 8
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+      }
+      args {
+        slice {
+          size: 32
+          buffer_allocation_index: 1
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          dimensions: 8
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+      }
+      args {
+        slice {
+          size: 4
+          buffer_allocation_index: 2
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          layout {
+            tail_padding_alignment_in_elements: 1
+          }
+        }
+      }
+      args {
+        slice {
+          size: 32
+          buffer_allocation_index: 6
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          dimensions: 8
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+      }
+      written: false
+      written: false
+      written: false
+      written: true
+      custom_kernel {
+        name: "loop_select_fusion"
+        kernel_spec {
+          cubin {
+            data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\330\025\000\000\000\000\000\000X\020\000\000\000\000\000\000\002d\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_select_fusion\000.nv.capmerc.text.loop_select_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_select_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_select_fusion\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\337\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000_\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\004<\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\002\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 13.2, V13.2.51\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills  -arch sm_100a \000\000\014\000\000\000\010\000\000\000\350\003\000\000NVIDIA Corp\000\002\000d\000\204\000\000\000\004/\010\000\010\000\000\000\014\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\002\t\001\000\002\002\001\000\002\005\005\000\003\007\001\001\002\003\000\000\004\013\010\000\t\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\003\031 \000\004\n\010\000\t\000\000\000\200\003 \000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\344\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000.\016\000\202{\004\377\000\340\000\000\000\n\000\000\000b\016\000\031y\t\000\000\000\000\000\000!\000\000\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\314\017\000\202{\006\377\000\346\000\000\000\n\000\000\000\342\016\000\201y\002\002\004\000\000\000\000\231\036\014\000$\037\000\014r\000\002\377\000\000\000pP\360\003\000\332\017\001\202\013\004\377\000\342\000\000\000\010\000\000\000\260.\000\202\013\005\377\000\343\000\000\000\010\000\000\000\244\016\000%x\004\t\004\000\000\000\004\000\216\007\000\314O\000\201y\005\004\004\000\000\000\000\231\036\014\000\242\016\000%x\006\t\004\000\000\000\006\000\216\007\000\312\217\000\206y\000\006\005\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\001\000\000\300\021\000\000\000\200\271\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\004\n\370\000\004\000\000\000A\002\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\203\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\030\000\000\000\002\"\016\006\000\000B\000\000\000\001\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\002\"\016\006\000\000B\000\000\000A\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\0028\0162\370\000@\021\000\000\000\000\202\001\n\000\000\002\001@\001\000\000\000\000\000\000\000\000\000\000\000P\010\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\004\020\000\000\000\004\220\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\010\000\000\000\014\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004Z \000\212\235\"\244\261\235\024m\000\264*\363\367X\003\245\',!0\311\036\307\217\017\014Il\n/\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000\220\001\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000D\000\000\000\000\000\000\000=\000\001\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\337\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\351\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000M\002\000\000\000\000\000\000\200\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\003\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\306\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000(\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\334\005\000\000\000\000\000\000 \000\000\000\000\000\000\000\005\000\000\000\010\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \006\000\000\000\000\000\000 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000v\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\006\000\000\000\000\000\000\220\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\344\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\007\000\000\000\000\000\000\000\002\000\000\000\000\000\000\003\000\000\000\010\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\t\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\022\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\t\000\000\000\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0003\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000 \r\000\000\000\000\000\000\022\001\000\000\000\000\000\000\025\000\000\000\010\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000W\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0002\016\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000l\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\244\016\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000}\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\016\000\000\000\000\000\000\234\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000h\017\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\273\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\200\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\331\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\200\017\000\000\000\000\000\000\330\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\330\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\330\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\200\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000\240\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000"
+          }
+          arity: 4
+          kernel_name: "loop_select_fusion"
+          kernel_args_packing_spec {
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 1
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 2
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 3
+              }
+            }
+          }
+        }
+        block_dims {
+          coordinates {
+            x: 1
+            y: 1
+            z: 1
+          }
+        }
+        thread_dims {
+          coordinates {
+            x: 8
+            y: 1
+            z: 1
+          }
+        }
+      }
+      use_pdl: true
+      tma_metadata {
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "ars.3"
+      thunk_id: 3
+    }
+    reduce_scatter_thunk {
+      collective_config {
+        operand_element_type: U32
+      }
+      buffers {
+        element_count: 8
+        source_buffer {
+          slice {
+            size: 32
+            buffer_allocation_index: 6
+            element_type: U32
+          }
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+        }
+        destination_buffer {
+          slice {
+            size: 16
+            buffer_allocation_index: 2
+            element_type: U32
+          }
+          shape {
+            element_type: U32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+        }
+      }
+      reduction_kind: REDUCTION_KIND_SUM
+    }
+  }
+  executable_abi_version {
+    platform_name: "CUDA"
+    cuda_platform_version {
+      cuda_toolkit_version: "13.2.0"
+      cudnn_version: "9.23.0"
+      cub_version: "3.1.2"
+    }
+  }
+  buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n             parameter allocation:         0B\n              constant allocation:        64B\n        maybe_live_out allocation:        16B\n     preallocated temp allocation:        32B\n  preallocated temp fragmentation:         0B (0.00%)\n                 total allocation:       124B\n              total fragmentation:        76B (61.29%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n"
+  gpu_topology {
+    platform_version: "Compute Capability 10.0a"
+    num_partitions: 1
+    num_hosts_per_partition: 1
+    num_devices_per_host: 2
+    gpu_target_config {
+      gpu_device_info {
+        threads_per_block_limit: 1024
+        threads_per_warp: 32
+        shared_memory_per_block: 49152
+        shared_memory_per_core: 233472
+        threads_per_core_limit: 2048
+        core_count: 148
+        fpus_per_core: 128
+        block_dim_limit_x: 2147483647
+        block_dim_limit_y: 65535
+        block_dim_limit_z: 65535
+        memory_bandwidth: 7672320000000
+        l2_cache_size: 132644864
+        clock_rate_ghz: 1.965
+        device_memory_size: 191514411008
+        shared_memory_per_block_optin: 232448
+        cuda_compute_capability {
+          major: 10
+          feature_extension: ACCELERATED_FEATURES
+        }
+        registers_per_core_limit: 65536
+        registers_per_block_limit: 65536
+        scalar_unit_description {
+          rate_infos {
+            key: 4
+            value {
+              clock_rate_ghz: 1.965
+              units_per_core: 64
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 8
+            value {
+              clock_rate_ghz: 1.965
+              units_per_core: 64
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 10
+            value {
+              clock_rate_ghz: 1.965
+              units_per_core: 128
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 11
+            value {
+              clock_rate_ghz: 1.965
+              units_per_core: 128
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 12
+            value {
+              clock_rate_ghz: 1.965
+              units_per_core: 64
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 16
+            value {
+              clock_rate_ghz: 1.965
+              units_per_core: 128
+              ops_per_clock: 1
+            }
+          }
+        }
+        matrix_unit_description {
+          rate_infos {
+            key: 2
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 6
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 10
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 11
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 512
+            }
+          }
+          rate_infos {
+            key: 12
+            value {
+              clock_rate_ghz: 1.965
+              units_per_core: 4
+              ops_per_clock: 16
+            }
+          }
+          rate_infos {
+            key: 16
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 19
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 20
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 23
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 24
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 25
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 28
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 29
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 32
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 4096
+            }
+          }
+          rate_infos {
+            key: 33
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 35
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 36
+            value {
+              clock_rate_ghz: 1.8353101
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+        }
+        driver_version: "13.2.0"
+        kernel_mode_driver_version: "570.133.20"
+        runtime_version: "13.2.0"
+        compile_time_toolkit_version: "13.2.0"
+        dnn_version: "9.23.0"
+        cub_version: "3.1.2"
+        device_interconnect_info {
+          active_links: 18
+        }
+        device_vendor: "NVIDIA Corporation"
+        platform_version: "Compute Capability 10.0a"
+        pci_bus_id: "<undefined>"
+        name: "NVIDIA B200"
+        model_str: "sm_10.0a with 191514411008B RAM, 148 cores, 1965000KHz clock, 3996000KHz mem clock, 132644864B L2$"
+        numa_node: -1
+        thread_dim_limit_x: 1024
+        thread_dim_limit_y: 1024
+        thread_dim_limit_z: 64
+        device_address_bits: 64
+        pcie_bandwidth: 68719476736
+        ecc_enabled: true
+        mem_clock_ghz: 3.996
+        reserved_shared_memory_per_block: 1024
+        max_blocks_per_multiprocessor: 32
+        collective_memory_granularity: 2097152
+      }
+      platform_name: "CUDA"
+      device_description_str: "NVIDIA B200"
+      runtime_version {
+        major: 13
+        minor: 2
+      }
+    }
+    host_target_machine_options {
+      triple: "x86_64-grtev4-linux-gnu"
+      cpu: "emeraldrapids"
+    }
+    num_devices_per_process: 2
+  }
+}
diff --git a/third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_gb200/v4/AllGatherMixedTypes.pbtxt b/third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_gb200/v4/AllGatherMixedTypes.pbtxt
new file mode 100644
index 00000000000..b49543a0ff0
--- /dev/null
+++ b/third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_gb200/v4/AllGatherMixedTypes.pbtxt
@@ -0,0 +1,2467 @@
+# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto
+# proto-message: xla.HumanReadableAotExecutable
+executable_and_options {
+  compile_options {
+    executable_build_options {
+      device_ordinal: -1
+      result_layout {
+        element_type: TUPLE
+        tuple_shapes {
+          element_type: U32
+          dimensions: 4
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+        tuple_shapes {
+          element_type: F32
+          dimensions: 4
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+      }
+      debug_options {
+        xla_disable_hlo_passes: "constant_folding"
+        xla_backend_optimization_level: 3
+        xla_eliminate_hlo_implicit_broadcast: true
+        xla_cpu_multi_thread_eigen: true
+        xla_gpu_cuda_data_dir: "./cuda_sdk_lib"
+        xla_llvm_enable_alias_scope_metadata: true
+        xla_llvm_enable_noalias_metadata: true
+        xla_llvm_enable_invariant_load_metadata: true
+        xla_llvm_disable_expensive_passes: false
+        xla_cpu_use_onednn: false
+        xla_cpu_enable_fast_math: false
+        xla_gpu_enable_fast_min_max: false
+        xla_force_host_platform_device_count: 1
+        xla_gpu_enable_dynamic_slice_fusion: true
+        xla_hlo_evaluator_use_fast_path: true
+        xla_dump_hlo_as_html: false
+        xla_cpu_fast_math_honor_nans: true
+        xla_cpu_fast_math_honor_infs: true
+        xla_allow_excess_precision: true
+        xla_gpu_autotune_level: 4
+        xla_cpu_fast_math_honor_division: true
+        xla_cpu_fast_math_honor_functions: true
+        xla_dump_include_timestamp: false
+        xla_dump_max_hlo_modules: -1
+        xla_cpu_enable_xprof_traceme: false
+        xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false
+        xla_cpu_enable_fast_min_max: true
+        xla_multiheap_size_constraint_per_heap: -1
+        xla_dump_module_metadata: false
+        xla_gpu_force_compilation_parallelism: 1
+        xla_dump_fusion_visualization: false
+        xla_gpu_all_reduce_combine_threshold_bytes: 31457287
+        xla_gpu_nccl_termination_timeout_seconds: -1
+        xla_dump_hlo_as_long_text: true
+        xla_gpu_enable_cublaslt: true
+        xla_gpu_shape_checks: RUNTIME
+        xla_gpu_use_runtime_fusion: false
+        xla_dump_latency_hiding_schedule: false
+        xla_dump_enable_mlir_pretty_form: true
+        xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP
+        xla_gpu_enable_triton_gemm: true
+        xla_gpu_enable_cudnn_int8x32_convolution_reordering: true
+        xla_gpu_triton_gemm_any: true
+        xla_gpu_enable_while_loop_reduce_scatter_code_motion: false
+        xla_gpu_collective_inflation_factor: 1
+        xla_gpu_graph_min_graph_size: 5
+        xla_gpu_enable_reassociation_for_converted_ar: true
+        xla_gpu_pgle_profile_file_or_directory_path: ""
+        xla_gpu_all_gather_combine_threshold_bytes: 31457287
+        xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287
+        xla_gpu_enable_highest_priority_async_stream: true
+        xla_gpu_exhaustive_tiling_search: false
+        xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0
+        xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1
+        xla_gpu_redzone_padding_bytes: 8388608
+        xla_gpu_copy_insertion_use_region_analysis: false
+        xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807
+        xla_gpu_collect_cost_model_stats: false
+        xla_gpu_enable_pipelined_p2p: false
+        xla_gpu_cublas_fallback: true
+        xla_gpu_enable_while_loop_double_buffering: false
+        xla_gpu_filter_kernels_spilling_registers_on_autotuning: true
+        xla_debug_buffer_assignment_show_max: 15
+        xla_detailed_logging: true
+        xla_enable_dumping: true
+        xla_gpu_enable_all_gather_combine_by_dim: false
+        xla_gpu_enable_analytical_latency_estimator: false
+        xla_gpu_llvm_verification_level: 0
+        xla_gpu_enable_reduce_scatter_combine_by_dim: false
+        xla_gpu_enable_command_buffer: COLLECTIVES_KERNEL
+        xla_gpu_enable_command_buffer: CONDITIONAL
+        xla_gpu_enable_command_buffer: CUBLAS
+        xla_gpu_enable_command_buffer: CUBLASLT
+        xla_gpu_enable_command_buffer: CUDNN
+        xla_gpu_enable_command_buffer: CUSTOM_CALL
+        xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION
+        xla_gpu_enable_command_buffer: FUSION
+        xla_gpu_enable_cub_radix_sort: true
+        xla_gpu_memory_limit_slop_factor: 95
+        xla_gpu_target_config_filename: ""
+        xla_gpu_enable_cudnn_layer_norm: false
+        xla_gpu_threshold_for_windowed_einsum_mib: 100000
+        xla_gpu_enable_nccl_user_buffers: false
+        xla_gpu_enable_libnvptxcompiler: true
+        xla_gpu_enable_nccl_comm_splitting: true
+        xla_gpu_nccl_collective_max_nchannels: 0
+        xla_gpu_nccl_p2p_max_nchannels: 0
+        xla_gpu_nccl_init_max_rank_per_root_ratio: 0
+        xla_gpu_multi_streamed_windowed_einsum: true
+        xla_gpu_gemm_rewrite_size_threshold: 100
+        xla_gpu_require_complete_aot_autotune_results: false
+        xla_gpu_cudnn_gemm_fusion_level: 0
+        xla_gpu_use_memcpy_local_p2p: false
+        xla_gpu_autotune_max_solutions: 0
+        xla_dump_large_constants: false
+        xla_gpu_verify_triton_fusion_numerics: false
+        xla_reduce_window_rewrite_base_length: 16
+        xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL
+        xla_gpu_enable_host_memory_offloading: false
+        xla_gpu_exclude_nondeterministic_ops: true
+        xla_gpu_nccl_terminate_on_error: false
+        xla_gpu_shard_autotuning: true
+        xla_gpu_enable_approx_costly_collectives: false
+        xla_cpu_prefer_vector_width: 256
+        xla_gpu_per_fusion_autotune_cache_dir: ""
+        xla_cmd_buffer_trace_cache_size: 16
+        xla_gpu_temp_buffer_use_separate_color: false
+        xla_syntax_sugar_async_ops: false
+        xla_gpu_autotune_gemm_rtol: 0.1
+        xla_enable_command_buffers_during_profiling: false
+        xla_gpu_cudnn_gemm_max_plans: 5
+        xla_cpu_parallel_codegen_split_count: 32
+        xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE
+        xla_gpu_executable_warn_stuck_timeout_seconds: 10
+        xla_gpu_executable_terminate_timeout_seconds: 30
+        xla_gpu_experimental_disable_binary_libraries: false
+        xla_gpu_dot_merger_threshold_mb: 64
+        xla_cpu_max_isa: "NEON"
+        xla_gpu_experimental_enable_fusion_block_level_rewriter: false
+        xla_enable_fast_math: false
+        xla_gpu_experimental_parallel_collective_overlap_limit: 1
+        xla_cpu_copy_insertion_use_region_analysis: false
+        xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1
+        xla_gpu_experimental_enable_triton_heroless_priority_fusion: false
+        xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN
+        xla_gpu_experimental_stream_annotation: true
+        xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO
+        xla_pjrt_allow_auto_layout_in_hlo: false
+        xla_gpu_enable_scatter_determinism_expander: false
+        xla_gpu_require_exclusive_lock: false
+        xla_gpu_generate_debug_info: false
+        xla_gpu_generate_line_info: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false
+        xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE
+        xla_gpu_fail_ptx_compilation_on_register_spilling: false
+        xla_gpu_enable_analytical_sol_latency_estimator: true
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_prep_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_size_bytes"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "gpus_per_node"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nccl_op_launch_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nic_speed_gbps"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "rtt_us"
+          value: "-1"
+        }
+        xla_cpu_use_xnnpack: true
+        xla_gpu_experimental_pack_dot_operands_along_k_dimension: true
+        xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false
+        xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED
+        xla_gpu_experimental_enable_heuristic_collective_combining: false
+        xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false
+        xla_hlo_pass_fix_detect_cycles: false
+        xla_cpu_generate_unique_c_style_kernel_entry_points: false
+        xla_gpu_experimental_collective_cse_distance_threshold: 0
+        xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true
+        xla_gpu_experimental_collective_perf_table_path: ""
+        xla_gpu_collective_permute_combine_threshold_bytes: 31457287
+        xla_unsupported_crash_on_hlo_pass_noop_change: false
+        xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false
+        xla_dump_full_hlo_config: true
+        xla_gpu_unsupported_enable_triton_multi_output_fusion: false
+        xla_gpu_experimental_matmul_perf_table_path: ""
+        xla_gpu_unsupported_enable_all_reduce_decomposer: false
+        xla_gpu_nccl_blocking_communicators: true
+        xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20
+        xla_gpu_first_collective_call_terminate_timeout_seconds: 40
+        xla_gpu_nccl_async_execution: false
+        xla_cpu_emitter_verification_level: 1
+        xla_early_exit_with_layouts: false
+        xla_gpu_experimental_use_ragged_dot_fusion: false
+        xla_detect_unstable_reductions: DETECTION_MODE_NONE
+        xla_gpu_command_buffer_scheduling_mode: LHS
+        xla_gpu_experimental_enable_nccl_symmetric_buffers: false
+        xla_disable_automatic_host_compute_offload: false
+        xla_gpu_experimental_enable_fusion_autotuner: true
+        xla_gpu_experimental_scaled_dot_with_triton: true
+        xla_gpu_command_buffer_unroll_loops: false
+        xla_cpu_experimental_onednn_custom_call: false
+        xla_gpu_experimental_enable_checksum_tracing_on_thunks: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true
+        xla_cpu_collective_call_terminate_timeout_seconds: 40
+        xla_cpu_collective_call_warn_stuck_seconds: 20
+        xla_keep_shardings_after_spmd: false
+        xla_gpu_experimental_enable_triton_warp_specialization: false
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE
+        xla_cpu_enable_platform_dependent_math: true
+        xla_gpu_detect_nan: DETECTION_MODE_NONE
+        xla_gpu_detect_inf: DETECTION_MODE_NONE
+        xla_enable_enzyme_comms_opt: false
+        xla_gpu_experimental_enable_buffer_saver_on_thunks: false
+        xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE
+        xla_enable_scoped_logging_timers: true
+        xla_gpu_executable_embed_debug_info: true
+        xla_cpu_collective_timeout_seconds: 1800
+        xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false
+        xla_gpu_enable_pipelined_host_offloading: false
+        xla_gpu_default_to_alg_dot_bf16_bf16_f32: false
+        xla_gpu_experimental_autotune_backends: CUDNN
+        xla_gpu_experimental_autotune_backends: TRITON
+        xla_gpu_experimental_autotune_backends: CUBLASLT
+        xla_gpu_experimental_autotune_backends: HIPBLASLT
+        xla_gpu_experimental_autotune_backends: MIOPEN
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER
+        xla_gpu_experimental_autotune_backends: NATIVE_EMITTER
+        xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER
+        xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION
+        xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: YNNPACK
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU
+        xla_gpu_experimental_autotune_backends: TENSOR_IR
+        xla_gpu_experimental_all_fusions_with_triton: false
+        xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT
+        xla_gpu_rocm_max_trace_events: 4194304
+        xla_gpu_execution_terminate_timeout: "inf"
+        xla_enable_hlo_sharding_v3: true
+        xla_gpu_log_minmax: false
+        xla_gpu_print_compilation_stats: false
+        xla_gpu_experimental_enable_conv_fusion: false
+        xla_gpu_experimental_enable_tiling_propagation: true
+        xla_gpu_execution_progress_tracking: 0
+        xla_gpu_experimental_max_unroll_factor: 32
+        xla_gpu_enable_pdl: true
+        xla_gpu_collective_combine_threshold_count: 2048
+        xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true
+        xla_recognize_reduction_optimization_level: 0
+        xla_gpu_experimental_parallel_async_compute_limit: 2
+        xla_dump_buffer_assignment_analysis: true
+        xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE
+        xla_gpu_experimental_enable_collective_multi_streaming: false
+        xla_gpu_collective_permute_connected_components: false
+        xla_gpu_experimental_force_split_k: 0
+        xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_experimental_gemm_fusion_v2: true
+        xla_gpu_fusion_autotune_top_k_configs: 1
+        xla_gpu_native_emitter_tune_unroll_factor_for_loops: false
+        xla_enable_rgv3_materialization: true
+        xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK
+        xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false
+        xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_executable_num_compute_streams: 0
+        xla_gpu_executable_num_communication_streams: 0
+        xla_gpu_experimental_collective_start_as_early_as_possible: false
+        xla_cpu_experimental_enable_tiling_propagation: true
+        xla_gpu_enable_dus_accumulator_zero_init_elimination: false
+        xla_gpu_experimental_thunk_buffer_debug_module_outputs: false
+        xla_cpu_ftz: true
+        xla_gpu_enable_allocator_spatial_partitioning: true
+        xla_dump_compact_gte: true
+        xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES
+        xla_dump_hlo_as_riegeli: false
+        xla_gpu_experimental_use_ragged_dot_grouped_gemm: true
+        xla_dump_hlo_to_subfolder: false
+        xla_gpu_enable_pdl_launch: true
+        xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED
+        xla_gpu_enable_nccl_user_buffers_in_default_space: false
+        xla_sdy_export_all_reduce_scatter: false
+        xla_gpu_enable_gxl_ragged_all_to_all: false
+        xla_gpu_gxl_scratch_size_bytes: 67108864
+        xla_gpu_use_new_autotune_cache_format: true
+        xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE
+        xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON
+        xla_spmd_enable_dynamic_slice_collective_broadcast: false
+        xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false
+        xla_gpu_async_copy_min_bytes: -1
+        xla_gpu_blas_max_algorithms: 0
+        xla_gpu_experimental_enable_selective_memcpy_overlap: false
+        xla_gpu_experimental_enable_same_shape_multi_output_fusion: false
+        xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1
+        xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1
+        xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240
+        xla_cpu_use_new_xtile_lowering: true
+        xla_gpu_trace_annotation_level: 0
+        xla_gpu_experimental_cost_model_gemm_tiling_default: false
+        xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false
+        xla_gpu_all_reduce_splitter_ignore_profitability_check: false
+        xla_gpu_enable_cupti_multi_subscriber: true
+        xla_enable_hlo_modules_upload: true
+        xla_gpu_collective_domain_assignment: ""
+        xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1
+        xla_compile_all_supported_configs: false
+        xla_gpu_unsupported_use_cross_host_one_shot_kernel: ALLCOLLECTIVES
+        xla_deduplicate_backend_configs_min_size: 9223372036854775807
+        xla_force_config: ""
+        xla_candidate_configs_file: ""
+        xla_gpu_enable_persistent_symmetric_memory: false
+        xla_gpu_experimental_enable_raft_for_stable_topk: false
+        xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND
+        xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET
+        xla_gpu_experimental_emit_collective_reduce: false
+        xla_gpu_experimental_vmm_disabled: false
+        xla_gpu_experimental_enable_dynamic_slice_table_offsets: false
+        xla_gpu_device_execution_terminate_timeout: "inf"
+        xla_gpu_topology_filename: ""
+        xla_gpu_experimental_enable_dynamic_slice_extended_offsets: false
+        xla_gpu_experimental_deterministic_ops_triton: DETERMINISTIC_OPS_TRITON_MODE_OFF
+        xla_gpu_experimental_enable_tensor_ir: false
+      }
+      num_replicas: 2
+      num_partitions: 1
+      device_assignment {
+        replica_count: 2
+        computation_count: 1
+        computation_devices {
+          replica_device_ids: 0
+          replica_device_ids: 1
+        }
+      }
+      allow_spmd_sharding_propagation_to_output: false
+      comp_envs {
+      }
+      allow_spmd_sharding_propagation_to_parameters: false
+      process_count: 1
+      memory_fitting_level: EFFORT_O2
+    }
+  }
+  pjrt_client_name: "PjRtStreamExecutorClient"
+}
+gpu_executable {
+  hlo_module_with_config {
+    hlo_module {
+      name: "test"
+      entry_computation_name: "test_computation"
+      computations {
+        name: "fused_transpose"
+        instructions {
+          name: "param_0.1"
+          opcode: "parameter"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "param_0.1"
+          }
+          id: 4294967298
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "bitcast.6.1"
+          opcode: "bitcast"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "bitcast.6.1"
+          }
+          id: 4294967297
+          operand_ids: 4294967298
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "transpose.1.1"
+          opcode: "transpose"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "transpose.1.1"
+          }
+          dimensions: 1
+          dimensions: 0
+          id: 4294967296
+          operand_ids: 4294967297
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          parameters {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          result {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          parameter_names: "param_0.1"
+        }
+        id: 1
+        root_id: 4294967296
+        is_fusion_computation: true
+      }
+      computations {
+        name: "fused_transpose.1"
+        instructions {
+          name: "param_0.3"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "param_0.3"
+          }
+          id: 8589934594
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "bitcast.4.1"
+          opcode: "bitcast"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "bitcast.4.1"
+          }
+          id: 8589934593
+          operand_ids: 8589934594
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "transpose.3"
+          opcode: "transpose"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "transpose.3"
+          }
+          dimensions: 1
+          dimensions: 0
+          id: 8589934592
+          operand_ids: 8589934593
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          parameters {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          result {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          parameter_names: "param_0.3"
+        }
+        id: 2
+        root_id: 8589934592
+        is_fusion_computation: true
+      }
+      computations {
+        name: "fused_broadcast"
+        instructions {
+          name: "param_0.5"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "param_0.5"
+          }
+          id: 12884901890
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "p1.1.1"
+          opcode: "convert"
+          shape {
+            element_type: F32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "p1.1.1"
+          }
+          id: 12884901889
+          operand_ids: 12884901890
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "broadcast.13.1"
+          opcode: "broadcast"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "broadcast.13.1"
+          }
+          id: 12884901888
+          operand_ids: 12884901889
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          parameters {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          result {
+            element_type: F32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          parameter_names: "param_0.5"
+        }
+        id: 3
+        root_id: 12884901888
+        is_fusion_computation: true
+      }
+      computations {
+        name: "wrapped_broadcast_computation"
+        instructions {
+          name: "param_0.7"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "param_0.7"
+          }
+          id: 21474836481
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "broadcast.12.1"
+          opcode: "broadcast"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "broadcast.12.1"
+          }
+          id: 21474836480
+          operand_ids: 21474836481
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          parameters {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          result {
+            element_type: U32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          parameter_names: "param_0.7"
+        }
+        id: 5
+        root_id: 21474836480
+        is_fusion_computation: true
+      }
+      computations {
+        name: "test_computation"
+        instructions {
+          name: "id.1"
+          opcode: "replica-id"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "id.1"
+          }
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "wrapped_broadcast"
+          opcode: "fusion"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "wrapped_broadcast"
+          }
+          fusion_kind: "kLoop"
+          id: 9
+          operand_ids: 0
+          called_computation_ids: 5
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "loop_broadcast_fusion"
+          opcode: "fusion"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "loop_broadcast_fusion"
+          }
+          fusion_kind: "kLoop"
+          id: 8
+          operand_ids: 0
+          called_computation_ids: 3
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "allgather.3"
+          opcode: "all-gather"
+          shape {
+            element_type: TUPLE
+            tuple_shapes {
+              element_type: U32
+              dimensions: 2
+              dimensions: 2
+              layout {
+                minor_to_major: 0
+                minor_to_major: 1
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+              is_dynamic_dimension: false
+            }
+            tuple_shapes {
+              element_type: F32
+              dimensions: 2
+              dimensions: 2
+              layout {
+                minor_to_major: 0
+                minor_to_major: 1
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+              is_dynamic_dimension: false
+            }
+          }
+          metadata {
+            scheduling_name: "all-gather-start"
+          }
+          dimensions: 1
+          id: 10
+          operand_ids: 9
+          operand_ids: 8
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+          collective_device_list {
+          }
+          backend_config_payload {
+            id: 0
+          }
+        }
+        instructions {
+          name: "ag0.1"
+          opcode: "get-tuple-element"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "ag0.1"
+          }
+          id: 1
+          operand_ids: 10
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "loop_transpose_fusion.1"
+          opcode: "fusion"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "loop_transpose_fusion.1"
+          }
+          fusion_kind: "kLoop"
+          id: 7
+          operand_ids: 1
+          called_computation_ids: 2
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "bitcast.24"
+          opcode: "bitcast"
+          shape {
+            element_type: U32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "bitcast.24"
+          }
+          id: 4
+          operand_ids: 7
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "ag1.1"
+          opcode: "get-tuple-element"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "ag1.1"
+          }
+          tuple_index: 1
+          id: 2
+          operand_ids: 10
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "loop_transpose_fusion"
+          opcode: "fusion"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "loop_transpose_fusion"
+          }
+          fusion_kind: "kLoop"
+          id: 6
+          operand_ids: 2
+          called_computation_ids: 1
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "bitcast.1.0"
+          opcode: "bitcast"
+          shape {
+            element_type: F32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "bitcast.1.0"
+          }
+          id: 5
+          operand_ids: 6
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "out.1"
+          opcode: "tuple"
+          shape {
+            element_type: TUPLE
+            tuple_shapes {
+              element_type: U32
+              dimensions: 4
+              layout {
+                minor_to_major: 0
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+            }
+            tuple_shapes {
+              element_type: F32
+              dimensions: 4
+              layout {
+                minor_to_major: 0
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+            }
+          }
+          metadata {
+            scheduling_name: "out.1"
+          }
+          id: 3
+          operand_ids: 4
+          operand_ids: 5
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          result {
+            element_type: TUPLE
+            tuple_shapes {
+              element_type: U32
+              dimensions: 4
+              layout {
+                minor_to_major: 0
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+            }
+            tuple_shapes {
+              element_type: F32
+              dimensions: 4
+              layout {
+                minor_to_major: 0
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+            }
+          }
+        }
+        root_id: 3
+      }
+      host_program_shape {
+        result {
+          element_type: TUPLE
+          tuple_shapes {
+            element_type: U32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          tuple_shapes {
+            element_type: F32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+        }
+      }
+      schedule {
+        sequences {
+          key: 0
+          value {
+            instruction_ids: 0
+            instruction_ids: 8
+            instruction_ids: 9
+            instruction_ids: 10
+            instruction_ids: 1
+            instruction_ids: 7
+            instruction_ids: 2
+            instruction_ids: 6
+            instruction_ids: 4
+            instruction_ids: 5
+            instruction_ids: 3
+          }
+        }
+      }
+      input_output_alias {
+      }
+      device_assignment {
+        replica_count: 2
+        computation_count: 1
+        computation_devices {
+          replica_device_ids: 0
+          replica_device_ids: 1
+        }
+      }
+      buffer_donor {
+      }
+      frontend_attributes {
+        map {
+          key: "fingerprint_before_lhs"
+          value: "7b2e4f79c5524b5cbc566da178aa7fda"
+        }
+      }
+      payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}"
+    }
+    config {
+      entry_computation_layout {
+        result {
+          element_type: TUPLE
+          tuple_shapes {
+            element_type: U32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          tuple_shapes {
+            element_type: F32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+        }
+      }
+      replica_count: 2
+      num_partitions: 1
+      intra_op_parallelism_threads: -1
+      debug_options {
+        xla_disable_hlo_passes: "constant_folding"
+        xla_backend_optimization_level: 3
+        xla_eliminate_hlo_implicit_broadcast: true
+        xla_cpu_multi_thread_eigen: true
+        xla_gpu_cuda_data_dir: "./cuda_sdk_lib"
+        xla_llvm_enable_alias_scope_metadata: true
+        xla_llvm_enable_noalias_metadata: true
+        xla_llvm_enable_invariant_load_metadata: true
+        xla_llvm_disable_expensive_passes: false
+        xla_cpu_use_onednn: false
+        xla_cpu_enable_fast_math: false
+        xla_gpu_enable_fast_min_max: false
+        xla_force_host_platform_device_count: 1
+        xla_gpu_enable_dynamic_slice_fusion: true
+        xla_hlo_evaluator_use_fast_path: true
+        xla_dump_hlo_as_html: false
+        xla_cpu_fast_math_honor_nans: true
+        xla_cpu_fast_math_honor_infs: true
+        xla_allow_excess_precision: true
+        xla_gpu_autotune_level: 4
+        xla_cpu_fast_math_honor_division: true
+        xla_cpu_fast_math_honor_functions: true
+        xla_dump_include_timestamp: false
+        xla_dump_max_hlo_modules: -1
+        xla_cpu_enable_xprof_traceme: false
+        xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false
+        xla_cpu_enable_fast_min_max: true
+        xla_multiheap_size_constraint_per_heap: -1
+        xla_dump_module_metadata: false
+        xla_gpu_force_compilation_parallelism: 1
+        xla_dump_fusion_visualization: false
+        xla_gpu_all_reduce_combine_threshold_bytes: 31457287
+        xla_gpu_nccl_termination_timeout_seconds: -1
+        xla_dump_hlo_as_long_text: true
+        xla_gpu_enable_cublaslt: true
+        xla_gpu_shape_checks: RUNTIME
+        xla_gpu_use_runtime_fusion: false
+        xla_dump_latency_hiding_schedule: false
+        xla_dump_enable_mlir_pretty_form: true
+        xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP
+        xla_gpu_enable_triton_gemm: true
+        xla_gpu_enable_cudnn_int8x32_convolution_reordering: true
+        xla_gpu_triton_gemm_any: true
+        xla_gpu_enable_while_loop_reduce_scatter_code_motion: false
+        xla_gpu_collective_inflation_factor: 1
+        xla_gpu_graph_min_graph_size: 5
+        xla_gpu_enable_reassociation_for_converted_ar: true
+        xla_gpu_pgle_profile_file_or_directory_path: ""
+        xla_gpu_all_gather_combine_threshold_bytes: 31457287
+        xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287
+        xla_gpu_enable_highest_priority_async_stream: true
+        xla_gpu_exhaustive_tiling_search: false
+        xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0
+        xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1
+        xla_gpu_redzone_padding_bytes: 8388608
+        xla_gpu_copy_insertion_use_region_analysis: false
+        xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807
+        xla_gpu_collect_cost_model_stats: false
+        xla_gpu_enable_pipelined_p2p: false
+        xla_gpu_cublas_fallback: true
+        xla_gpu_enable_while_loop_double_buffering: false
+        xla_gpu_filter_kernels_spilling_registers_on_autotuning: true
+        xla_debug_buffer_assignment_show_max: 15
+        xla_detailed_logging: true
+        xla_enable_dumping: true
+        xla_gpu_enable_all_gather_combine_by_dim: false
+        xla_gpu_enable_analytical_latency_estimator: false
+        xla_gpu_llvm_verification_level: 0
+        xla_gpu_enable_reduce_scatter_combine_by_dim: false
+        xla_gpu_enable_command_buffer: COLLECTIVES_KERNEL
+        xla_gpu_enable_command_buffer: CONDITIONAL
+        xla_gpu_enable_command_buffer: CUBLAS
+        xla_gpu_enable_command_buffer: CUBLASLT
+        xla_gpu_enable_command_buffer: CUDNN
+        xla_gpu_enable_command_buffer: CUSTOM_CALL
+        xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION
+        xla_gpu_enable_command_buffer: FUSION
+        xla_gpu_enable_cub_radix_sort: true
+        xla_gpu_memory_limit_slop_factor: 95
+        xla_gpu_target_config_filename: ""
+        xla_gpu_enable_cudnn_layer_norm: false
+        xla_gpu_threshold_for_windowed_einsum_mib: 100000
+        xla_gpu_enable_nccl_user_buffers: false
+        xla_gpu_enable_libnvptxcompiler: true
+        xla_gpu_enable_nccl_comm_splitting: true
+        xla_gpu_nccl_collective_max_nchannels: 0
+        xla_gpu_nccl_p2p_max_nchannels: 0
+        xla_gpu_nccl_init_max_rank_per_root_ratio: 0
+        xla_gpu_multi_streamed_windowed_einsum: true
+        xla_gpu_gemm_rewrite_size_threshold: 100
+        xla_gpu_require_complete_aot_autotune_results: false
+        xla_gpu_cudnn_gemm_fusion_level: 0
+        xla_gpu_use_memcpy_local_p2p: false
+        xla_gpu_autotune_max_solutions: 0
+        xla_dump_large_constants: false
+        xla_gpu_verify_triton_fusion_numerics: false
+        xla_reduce_window_rewrite_base_length: 16
+        xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL
+        xla_gpu_enable_host_memory_offloading: false
+        xla_gpu_exclude_nondeterministic_ops: true
+        xla_gpu_nccl_terminate_on_error: false
+        xla_gpu_shard_autotuning: true
+        xla_gpu_enable_approx_costly_collectives: false
+        xla_cpu_prefer_vector_width: 256
+        xla_gpu_per_fusion_autotune_cache_dir: ""
+        xla_cmd_buffer_trace_cache_size: 16
+        xla_gpu_temp_buffer_use_separate_color: false
+        xla_syntax_sugar_async_ops: false
+        xla_gpu_autotune_gemm_rtol: 0.1
+        xla_enable_command_buffers_during_profiling: false
+        xla_gpu_cudnn_gemm_max_plans: 5
+        xla_cpu_parallel_codegen_split_count: 32
+        xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE
+        xla_gpu_executable_warn_stuck_timeout_seconds: 10
+        xla_gpu_executable_terminate_timeout_seconds: 30
+        xla_gpu_experimental_disable_binary_libraries: false
+        xla_gpu_dot_merger_threshold_mb: 64
+        xla_cpu_max_isa: "NEON"
+        xla_gpu_experimental_enable_fusion_block_level_rewriter: false
+        xla_enable_fast_math: false
+        xla_gpu_experimental_parallel_collective_overlap_limit: 1
+        xla_cpu_copy_insertion_use_region_analysis: false
+        xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1
+        xla_gpu_experimental_enable_triton_heroless_priority_fusion: false
+        xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN
+        xla_gpu_experimental_stream_annotation: true
+        xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO
+        xla_pjrt_allow_auto_layout_in_hlo: false
+        xla_gpu_enable_scatter_determinism_expander: false
+        xla_gpu_require_exclusive_lock: false
+        xla_gpu_generate_debug_info: false
+        xla_gpu_generate_line_info: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false
+        xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE
+        xla_gpu_fail_ptx_compilation_on_register_spilling: false
+        xla_gpu_enable_analytical_sol_latency_estimator: true
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_prep_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_size_bytes"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "gpus_per_node"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nccl_op_launch_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nic_speed_gbps"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "rtt_us"
+          value: "-1"
+        }
+        xla_cpu_use_xnnpack: true
+        xla_gpu_experimental_pack_dot_operands_along_k_dimension: true
+        xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false
+        xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED
+        xla_gpu_experimental_enable_heuristic_collective_combining: false
+        xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false
+        xla_hlo_pass_fix_detect_cycles: false
+        xla_cpu_generate_unique_c_style_kernel_entry_points: false
+        xla_gpu_experimental_collective_cse_distance_threshold: 0
+        xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true
+        xla_gpu_experimental_collective_perf_table_path: ""
+        xla_gpu_collective_permute_combine_threshold_bytes: 31457287
+        xla_unsupported_crash_on_hlo_pass_noop_change: false
+        xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false
+        xla_dump_full_hlo_config: true
+        xla_gpu_unsupported_enable_triton_multi_output_fusion: false
+        xla_gpu_experimental_matmul_perf_table_path: ""
+        xla_gpu_unsupported_enable_all_reduce_decomposer: false
+        xla_gpu_nccl_blocking_communicators: true
+        xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20
+        xla_gpu_first_collective_call_terminate_timeout_seconds: 40
+        xla_gpu_nccl_async_execution: false
+        xla_cpu_emitter_verification_level: 1
+        xla_early_exit_with_layouts: false
+        xla_gpu_experimental_use_ragged_dot_fusion: false
+        xla_detect_unstable_reductions: DETECTION_MODE_NONE
+        xla_gpu_command_buffer_scheduling_mode: LHS
+        xla_gpu_experimental_enable_nccl_symmetric_buffers: false
+        xla_disable_automatic_host_compute_offload: false
+        xla_gpu_experimental_enable_fusion_autotuner: true
+        xla_gpu_experimental_scaled_dot_with_triton: true
+        xla_gpu_command_buffer_unroll_loops: false
+        xla_cpu_experimental_onednn_custom_call: false
+        xla_gpu_experimental_enable_checksum_tracing_on_thunks: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true
+        xla_cpu_collective_call_terminate_timeout_seconds: 40
+        xla_cpu_collective_call_warn_stuck_seconds: 20
+        xla_keep_shardings_after_spmd: false
+        xla_gpu_experimental_enable_triton_warp_specialization: false
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE
+        xla_cpu_enable_platform_dependent_math: true
+        xla_gpu_detect_nan: DETECTION_MODE_NONE
+        xla_gpu_detect_inf: DETECTION_MODE_NONE
+        xla_enable_enzyme_comms_opt: false
+        xla_gpu_experimental_enable_buffer_saver_on_thunks: false
+        xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE
+        xla_enable_scoped_logging_timers: true
+        xla_gpu_executable_embed_debug_info: true
+        xla_cpu_collective_timeout_seconds: 1800
+        xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false
+        xla_gpu_enable_pipelined_host_offloading: false
+        xla_gpu_default_to_alg_dot_bf16_bf16_f32: false
+        xla_gpu_experimental_autotune_backends: CUDNN
+        xla_gpu_experimental_autotune_backends: TRITON
+        xla_gpu_experimental_autotune_backends: CUBLASLT
+        xla_gpu_experimental_autotune_backends: HIPBLASLT
+        xla_gpu_experimental_autotune_backends: MIOPEN
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER
+        xla_gpu_experimental_autotune_backends: NATIVE_EMITTER
+        xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER
+        xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION
+        xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: YNNPACK
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU
+        xla_gpu_experimental_autotune_backends: TENSOR_IR
+        xla_gpu_experimental_all_fusions_with_triton: false
+        xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT
+        xla_gpu_rocm_max_trace_events: 4194304
+        xla_gpu_execution_terminate_timeout: "inf"
+        xla_enable_hlo_sharding_v3: true
+        xla_gpu_log_minmax: false
+        xla_gpu_print_compilation_stats: false
+        xla_gpu_experimental_enable_conv_fusion: false
+        xla_gpu_experimental_enable_tiling_propagation: true
+        xla_gpu_execution_progress_tracking: 0
+        xla_gpu_experimental_max_unroll_factor: 32
+        xla_gpu_enable_pdl: true
+        xla_gpu_collective_combine_threshold_count: 2048
+        xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true
+        xla_recognize_reduction_optimization_level: 0
+        xla_gpu_experimental_parallel_async_compute_limit: 2
+        xla_dump_buffer_assignment_analysis: true
+        xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE
+        xla_gpu_experimental_enable_collective_multi_streaming: false
+        xla_gpu_collective_permute_connected_components: false
+        xla_gpu_experimental_force_split_k: 0
+        xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_experimental_gemm_fusion_v2: true
+        xla_gpu_fusion_autotune_top_k_configs: 1
+        xla_gpu_native_emitter_tune_unroll_factor_for_loops: false
+        xla_enable_rgv3_materialization: true
+        xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK
+        xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false
+        xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_executable_num_compute_streams: 0
+        xla_gpu_executable_num_communication_streams: 0
+        xla_gpu_experimental_collective_start_as_early_as_possible: false
+        xla_cpu_experimental_enable_tiling_propagation: true
+        xla_gpu_enable_dus_accumulator_zero_init_elimination: false
+        xla_gpu_experimental_thunk_buffer_debug_module_outputs: false
+        xla_cpu_ftz: true
+        xla_gpu_enable_allocator_spatial_partitioning: true
+        xla_dump_compact_gte: true
+        xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES
+        xla_dump_hlo_as_riegeli: false
+        xla_gpu_experimental_use_ragged_dot_grouped_gemm: true
+        xla_dump_hlo_to_subfolder: false
+        xla_gpu_enable_pdl_launch: true
+        xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED
+        xla_gpu_enable_nccl_user_buffers_in_default_space: false
+        xla_sdy_export_all_reduce_scatter: false
+        xla_gpu_enable_gxl_ragged_all_to_all: false
+        xla_gpu_gxl_scratch_size_bytes: 67108864
+        xla_gpu_use_new_autotune_cache_format: true
+        xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE
+        xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON
+        xla_spmd_enable_dynamic_slice_collective_broadcast: false
+        xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false
+        xla_gpu_async_copy_min_bytes: -1
+        xla_gpu_blas_max_algorithms: 0
+        xla_gpu_experimental_enable_selective_memcpy_overlap: false
+        xla_gpu_experimental_enable_same_shape_multi_output_fusion: false
+        xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1
+        xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1
+        xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240
+        xla_cpu_use_new_xtile_lowering: true
+        xla_gpu_trace_annotation_level: 0
+        xla_gpu_experimental_cost_model_gemm_tiling_default: false
+        xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false
+        xla_gpu_all_reduce_splitter_ignore_profitability_check: false
+        xla_gpu_enable_cupti_multi_subscriber: true
+        xla_enable_hlo_modules_upload: true
+        xla_gpu_collective_domain_assignment: ""
+        xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1
+        xla_compile_all_supported_configs: false
+        xla_gpu_unsupported_use_cross_host_one_shot_kernel: ALLCOLLECTIVES
+        xla_deduplicate_backend_configs_min_size: 9223372036854775807
+        xla_force_config: ""
+        xla_candidate_configs_file: ""
+        xla_gpu_enable_persistent_symmetric_memory: false
+        xla_gpu_experimental_enable_raft_for_stable_topk: false
+        xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND
+        xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET
+        xla_gpu_experimental_emit_collective_reduce: false
+        xla_gpu_experimental_vmm_disabled: false
+        xla_gpu_experimental_enable_dynamic_slice_table_offsets: false
+        xla_gpu_device_execution_terminate_timeout: "inf"
+        xla_gpu_topology_filename: ""
+        xla_gpu_experimental_enable_dynamic_slice_extended_offsets: false
+        xla_gpu_experimental_deterministic_ops_triton: DETERMINISTIC_OPS_TRITON_MODE_OFF
+        xla_gpu_experimental_enable_tensor_ir: false
+      }
+      static_device_assignment {
+        replica_count: 2
+        computation_count: 1
+        computation_devices {
+          replica_device_ids: 0
+          replica_device_ids: 1
+        }
+      }
+      allow_spmd_sharding_propagation_to_output: false
+      allow_spmd_sharding_propagation_to_parameters: false
+      sharding_config {
+      }
+      memory_fitting_level: EFFORT_O2
+      schedule_config {
+      }
+      partition_size: 2
+    }
+  }
+  buffer_assignment {
+    logical_buffers {
+      id: 9
+      size: 4
+      defined_at {
+      }
+    }
+    logical_buffers {
+      id: 10
+      size: 8
+      defined_at {
+        instruction_id: 9
+      }
+    }
+    logical_buffers {
+      id: 11
+      size: 8
+      defined_at {
+        instruction_id: 8
+      }
+    }
+    logical_buffers {
+      id: 12
+      size: 16
+      defined_at {
+        instruction_id: 10
+      }
+    }
+    logical_buffers {
+      id: 13
+      size: 16
+      defined_at {
+        shape_index: 0
+        instruction_id: 10
+      }
+    }
+    logical_buffers {
+      id: 14
+      size: 16
+      defined_at {
+        shape_index: 1
+        instruction_id: 10
+      }
+    }
+    logical_buffers {
+      id: 15
+      size: 16
+      defined_at {
+        instruction_id: 7
+      }
+    }
+    logical_buffers {
+      id: 16
+      size: 16
+      defined_at {
+        instruction_id: 6
+      }
+    }
+    logical_buffers {
+      id: 17
+      size: 16
+      defined_at {
+        instruction_id: 3
+      }
+    }
+    buffer_allocations {
+      size: 16
+      maybe_live_out: true
+      assigned {
+        logical_buffer_id: 10
+        size: 8
+        element_type: U32
+      }
+      assigned {
+        logical_buffer_id: 15
+        size: 16
+        element_type: U32
+      }
+    }
+    buffer_allocations {
+      index: 1
+      size: 16
+      maybe_live_out: true
+      assigned {
+        logical_buffer_id: 9
+        size: 4
+        element_type: U32
+      }
+      assigned {
+        logical_buffer_id: 13
+        size: 16
+        element_type: U32
+      }
+      assigned {
+        logical_buffer_id: 16
+        size: 16
+        element_type: F32
+      }
+    }
+    buffer_allocations {
+      index: 2
+      size: 16
+      maybe_live_out: true
+      assigned {
+        logical_buffer_id: 17
+        size: 16
+        element_type: TUPLE
+      }
+      is_tuple: true
+    }
+    buffer_allocations {
+      index: 3
+      size: 520
+      assigned {
+        logical_buffer_id: 11
+        offset: 512
+        size: 8
+        element_type: F32
+      }
+      assigned {
+        logical_buffer_id: 12
+        size: 16
+        element_type: TUPLE
+      }
+      assigned {
+        logical_buffer_id: 14
+        offset: 256
+        size: 16
+        element_type: F32
+      }
+      is_tuple: true
+    }
+    heap_simulator_traces {
+      events {
+        buffer_id: 11
+        computation_name: "test_computation"
+        instruction_name: "loop_broadcast_fusion"
+      }
+      events {
+        buffer_id: 14
+        computation_name: "test_computation"
+        instruction_name: "allgather.3"
+      }
+      events {
+        kind: FREE
+        buffer_id: 11
+        computation_name: "test_computation"
+        instruction_name: "loop_broadcast_fusion"
+      }
+      events {
+        kind: FREE
+        buffer_id: 14
+        computation_name: "test_computation"
+        instruction_name: "allgather.3"
+      }
+      whole_module_simulation: true
+      buffer_allocation_index: 3
+    }
+    peak_buffers {
+      logical_buffer_ids: 15
+    }
+    peak_buffers {
+      index: 1
+      logical_buffer_ids: 16
+    }
+    peak_buffers {
+      index: 2
+      logical_buffer_ids: 17
+    }
+    peak_buffers {
+      index: 3
+      logical_buffer_ids: 12
+      logical_buffer_ids: 11
+      logical_buffer_ids: 14
+    }
+  }
+  gpu_compute_capability {
+    cuda_compute_capability {
+      major: 10
+      feature_extension: ACCELERATED_FEATURES
+    }
+  }
+  module_name: "test"
+  program_shape {
+    result {
+      element_type: TUPLE
+      tuple_shapes {
+        element_type: U32
+        dimensions: 4
+        layout {
+          minor_to_major: 0
+          tail_padding_alignment_in_elements: 1
+        }
+        is_dynamic_dimension: false
+      }
+      tuple_shapes {
+        element_type: F32
+        dimensions: 4
+        layout {
+          minor_to_major: 0
+          tail_padding_alignment_in_elements: 1
+        }
+        is_dynamic_dimension: false
+      }
+    }
+  }
+  buffer_allocations {
+    values {
+      size: 16
+      maybe_live_out: true
+    }
+    values {
+      index: 1
+      size: 16
+      maybe_live_out: true
+    }
+    values {
+      index: 2
+      size: 16
+      maybe_live_out: true
+      is_tuple: true
+    }
+    values {
+      index: 3
+      size: 520
+      is_tuple: true
+    }
+  }
+  output_info_map {
+    shape_index {
+    }
+    output_info {
+      allocation_index: 2
+    }
+  }
+  output_info_map {
+    shape_index {
+      indexes: 0
+    }
+    output_info {
+    }
+  }
+  output_info_map {
+    shape_index {
+      indexes: 1
+    }
+    output_info {
+      allocation_index: 1
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "id.1"
+      thunk_id: 1
+    }
+    replica_id_thunk {
+      dest_buffer {
+        size: 4
+        buffer_allocation_index: 1
+        element_type: U32
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "loop_broadcast_fusion"
+      thunk_id: 2
+    }
+    custom_kernel_thunk {
+      args {
+        slice {
+          size: 4
+          buffer_allocation_index: 1
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          layout {
+            tail_padding_alignment_in_elements: 1
+          }
+        }
+      }
+      args {
+        slice {
+          offset: 512
+          size: 8
+          buffer_allocation_index: 3
+          element_type: F32
+        }
+        shape {
+          element_type: F32
+          dimensions: 2
+          dimensions: 1
+          layout {
+            minor_to_major: 0
+            minor_to_major: 1
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+          is_dynamic_dimension: false
+        }
+      }
+      written: false
+      written: true
+      custom_kernel {
+        name: "loop_broadcast_fusion"
+        kernel_spec {
+          cubin {
+            data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\310\024\000\000\000\000\000\000H\017\000\000\000\000\000\000\002d\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_broadcast_fusion\000.nv.capmerc.text.loop_broadcast_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_broadcast_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\271\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\350\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000U\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000k\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004(\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\002\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 13.2, V13.2.51\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills  -arch sm_100a \000\000\014\000\000\000\010\000\000\000\350\003\000\000NVIDIA Corp\000\002\000d\000\204\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\002\t\001\000\002\002\001\000\002\005\005\000\003\007\001\001\002\003\000\000\004\013\010\000\t\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000\240\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\t\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\002\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\342/\000Er\007\000\002\000\000\000\000\020 \000\000\312O\000\206y\000\004\007\000\000\000\004\031\020\014\000\342\017\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\001\000\000\300\014\000\000\000\300\005\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\001\000\000\000\000\000\000\000\000\000\000\000\320\006\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\004\020\000\000\000\004\360\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004Z \000\212\235\"\244\261\235\024m\000\264*\363\367X\003\245\',!0\311\036\307\217\017\014Il\n/\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000D\000\000\000\000\000\000\000=\000\001\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\271\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\350\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000U\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\373\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000_\002\000\000\000\000\000\000\217\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\317\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000H\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\374\005\000\000\000\000\000\000 \000\000\000\000\000\000\000\005\000\000\000\010\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\006\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000y\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000`\006\000\000\000\000\000\000p\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\355\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\010\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\271\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\033\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000?\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\220\014\000\000\000\000\000\000\262\000\000\000\000\000\000\000\025\000\000\000\010\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000f\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000B\r\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000{\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\264\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\214\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\330\r\000\000\000\000\000\000|\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000X\016\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\315\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000p\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000p\016\000\000\000\000\000\000\330\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\310\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\310\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\200\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000"
+          }
+          arity: 2
+          kernel_name: "loop_broadcast_fusion"
+          kernel_args_packing_spec {
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 1
+              }
+            }
+          }
+        }
+        block_dims {
+          coordinates {
+            x: 1
+            y: 1
+            z: 1
+          }
+        }
+        thread_dims {
+          coordinates {
+            x: 2
+            y: 1
+            z: 1
+          }
+        }
+      }
+      use_pdl: true
+      tma_metadata {
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "wrapped_broadcast"
+      thunk_id: 3
+    }
+    custom_kernel_thunk {
+      args {
+        slice {
+          size: 4
+          buffer_allocation_index: 1
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          layout {
+            tail_padding_alignment_in_elements: 1
+          }
+        }
+      }
+      args {
+        slice {
+          size: 8
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          dimensions: 2
+          dimensions: 1
+          layout {
+            minor_to_major: 0
+            minor_to_major: 1
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+          is_dynamic_dimension: false
+        }
+      }
+      written: false
+      written: true
+      custom_kernel {
+        name: "wrapped_broadcast"
+        kernel_spec {
+          cubin {
+            data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000H\024\000\000\000\000\000\000\310\016\000\000\000\000\000\000\002d\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.wrapped_broadcast\000.nv.capmerc.text.wrapped_broadcast\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.wrapped_broadcast\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\375\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000-\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000[\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004$\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\002\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 13.2, V13.2.51\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills  -arch sm_100a \000\000\014\000\000\000\010\000\000\000\350\003\000\000NVIDIA Corp\000\002\000d\000\204\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\002\t\001\000\002\002\001\000\002\005\005\000\003\007\001\001\002\003\000\000\004\013\010\000\t\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000\220\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\t\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\340\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000\"\016\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\342\000\000\000\n\000\000\000b\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\312/\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\001\000\000\300\013\000\000\000\300\002\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\007\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\004\020\000\000\000\004\340\000\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004Z \000\212\235\"\244\261\235\024m\000\264*\363\367X\003\245\',!0\311\036\307\217\017\014Il\n/\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000\340\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000D\000\000\000\000\000\000\000=\000\001\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\375\000\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000-\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\360\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\343\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000G\002\000\000\000\000\000\000{\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\003\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\303\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\270\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\324\005\000\000\000\000\000\000 \000\000\000\000\000\000\000\005\000\000\000\010\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\364\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000u\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\006\000\000\000\000\000\000p\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\363\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\341\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\310\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\010\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\017\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000/\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\020\014\000\000\000\000\000\000\262\000\000\000\000\000\000\000\025\000\000\000\010\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\302\014\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000g\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0004\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000X\r\000\000\000\000\000\000|\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\233\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\330\r\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\265\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\323\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\360\r\000\000\000\000\000\000\330\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000H\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000H\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\000\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000"
+          }
+          arity: 2
+          kernel_name: "wrapped_broadcast"
+          kernel_args_packing_spec {
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 1
+              }
+            }
+          }
+        }
+        block_dims {
+          coordinates {
+            x: 1
+            y: 1
+            z: 1
+          }
+        }
+        thread_dims {
+          coordinates {
+            x: 2
+            y: 1
+            z: 1
+          }
+        }
+      }
+      use_pdl: true
+      tma_metadata {
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "allgather.3"
+      thunk_id: 4
+    }
+    all_gather_thunk {
+      collective_config {
+        operand_element_type: U32
+        operand_element_type: F32
+      }
+      buffers {
+        element_count: 2
+        source_buffer {
+          slice {
+            size: 8
+            element_type: U32
+          }
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+        }
+        destination_buffer {
+          slice {
+            size: 16
+            buffer_allocation_index: 1
+            element_type: U32
+          }
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+        }
+      }
+      buffers {
+        element_count: 2
+        source_buffer {
+          slice {
+            offset: 512
+            size: 8
+            buffer_allocation_index: 3
+            element_type: F32
+          }
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+        }
+        destination_buffer {
+          slice {
+            offset: 256
+            size: 16
+            buffer_allocation_index: 3
+            element_type: F32
+          }
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+        }
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "loop_transpose_fusion.1"
+      thunk_id: 5
+    }
+    custom_kernel_thunk {
+      args {
+        slice {
+          size: 16
+          buffer_allocation_index: 1
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          dimensions: 2
+          dimensions: 2
+          layout {
+            minor_to_major: 0
+            minor_to_major: 1
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+          is_dynamic_dimension: false
+        }
+      }
+      args {
+        slice {
+          size: 16
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          dimensions: 2
+          dimensions: 2
+          layout {
+            minor_to_major: 1
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+          is_dynamic_dimension: false
+        }
+      }
+      written: false
+      written: true
+      custom_kernel {
+        name: "loop_transpose_fusion_1"
+        kernel_spec {
+          cubin {
+            data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\310\024\000\000\000\000\000\000H\017\000\000\000\000\000\000\002d\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion_1\000.nv.capmerc.text.loop_transpose_fusion_1\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion_1\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\277\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\356\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\017\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000?\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000[\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000s\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\002\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 13.2, V13.2.51\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills  -arch sm_100a \000\000\014\000\000\000\010\000\000\000\350\003\000\000NVIDIA Corp\000\002\000d\000\204\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\002\t\001\000\002\002\001\000\002\005\005\000\003\007\001\001\002\003\000\000\004\013\010\000\t\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\t\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\001\000\000\300\017\000\000\000\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004Z \000\212\235\"\244\261\235\024m\000\264*\363\367X\003\245\',!0\311\036\307\217\017\014Il\n/\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000D\000\000\000\000\000\000\000=\000\001\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\277\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\356\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\017\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000?\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000[\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\007\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000k\002\000\000\000\000\000\000\231\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\004\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\325\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\370\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000`\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\024\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\005\000\000\000\010\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0004\006\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000X\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000{\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\006\000\000\000\000\000\000p\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\001\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\350\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\363\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\007\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\010\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\277\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000!\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000G\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\220\014\000\000\000\000\000\000\262\000\000\000\000\000\000\000\025\000\000\000\010\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000B\r\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\205\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\264\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\226\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\330\r\000\000\000\000\000\000|\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\277\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000X\016\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\331\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000p\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\367\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000p\016\000\000\000\000\000\000\330\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\310\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\310\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\200\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000"
+          }
+          arity: 2
+          kernel_name: "loop_transpose_fusion_1"
+          kernel_args_packing_spec {
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 1
+              }
+            }
+          }
+        }
+        block_dims {
+          coordinates {
+            x: 1
+            y: 1
+            z: 1
+          }
+        }
+        thread_dims {
+          coordinates {
+            x: 4
+            y: 1
+            z: 1
+          }
+        }
+      }
+      use_pdl: true
+      tma_metadata {
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "loop_transpose_fusion"
+      thunk_id: 6
+    }
+    custom_kernel_thunk {
+      args {
+        slice {
+          offset: 256
+          size: 16
+          buffer_allocation_index: 3
+          element_type: F32
+        }
+        shape {
+          element_type: F32
+          dimensions: 2
+          dimensions: 2
+          layout {
+            minor_to_major: 0
+            minor_to_major: 1
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+          is_dynamic_dimension: false
+        }
+      }
+      args {
+        slice {
+          size: 16
+          buffer_allocation_index: 1
+          element_type: F32
+        }
+        shape {
+          element_type: F32
+          dimensions: 2
+          dimensions: 2
+          layout {
+            minor_to_major: 1
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+          is_dynamic_dimension: false
+        }
+      }
+      written: false
+      written: true
+      custom_kernel {
+        name: "loop_transpose_fusion"
+        kernel_spec {
+          cubin {
+            data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\310\024\000\000\000\000\000\000H\017\000\000\000\000\000\000\002d\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion\000.nv.capmerc.text.loop_transpose_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_transpose_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\271\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\350\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000U\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000k\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\002\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 13.2, V13.2.51\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills  -arch sm_100a \000\000\014\000\000\000\010\000\000\000\350\003\000\000NVIDIA Corp\000\002\000d\000\204\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\002\t\001\000\002\002\001\000\002\005\005\000\003\007\001\001\002\003\000\000\004\013\010\000\t\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\t\000\000\000\200\003\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\340\000\000\000\n\000\000\000b\016\000\254w\004\377\000k\000\000\000\n\000\010\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\342\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242N\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\001\000\000\300\017\000\000\000\340-\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\001\013\004\n\370\000\004\000\000\000\301\001\001\002\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\0028\0162\370\000@\021\000\000\000\000\002\001\n\000\000\002\001\300\000\000\000\000\000\000\000\000\000\000\000\000P\005\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\004\020\000\000\000\004`\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004Z \000\212\235\"\244\261\235\024m\000\264*\363\367X\003\245\',!0\311\036\307\217\017\014Il\n/\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000`\001\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000D\000\000\000\000\000\000\000=\000\001\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\271\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\350\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000U\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000p\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\373\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000_\002\000\000\000\000\000\000\217\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\003\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\317\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000H\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\374\005\000\000\000\000\000\000 \000\000\000\000\000\000\000\005\000\000\000\010\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\034\006\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000y\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000`\006\000\000\000\000\000\000p\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\355\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\007\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\010\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\271\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\033\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\000\000\000\000\000\000\220\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000?\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\220\014\000\000\000\000\000\000\262\000\000\000\000\000\000\000\025\000\000\000\010\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000f\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000B\r\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000{\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\264\r\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\214\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\330\r\000\000\000\000\000\000|\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\263\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000X\016\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\315\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000p\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\353\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000p\016\000\000\000\000\000\000\330\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\310\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\310\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\200\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\220\003\000\000\000\000\000\000\220\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000"
+          }
+          arity: 2
+          kernel_name: "loop_transpose_fusion"
+          kernel_args_packing_spec {
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 1
+              }
+            }
+          }
+        }
+        block_dims {
+          coordinates {
+            x: 1
+            y: 1
+            z: 1
+          }
+        }
+        thread_dims {
+          coordinates {
+            x: 4
+            y: 1
+            z: 1
+          }
+        }
+      }
+      use_pdl: true
+      tma_metadata {
+      }
+    }
+  }
+  executable_abi_version {
+    platform_name: "CUDA"
+    cuda_platform_version {
+      cuda_toolkit_version: "13.2.0"
+      cudnn_version: "9.23.0"
+      cub_version: "3.1.2"
+    }
+  }
+  buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n             parameter allocation:         0B\n              constant allocation:         0B\n        maybe_live_out allocation:        48B\n     preallocated temp allocation:       520B\n  preallocated temp fragmentation:       240B (46.15%)\n                 total allocation:       568B\n              total fragmentation:       504B (88.73%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n"
+  gpu_topology {
+    platform_version: "Compute Capability 10.0a"
+    num_partitions: 1
+    num_hosts_per_partition: 1
+    num_devices_per_host: 2
+    gpu_target_config {
+      gpu_device_info {
+        threads_per_block_limit: 1024
+        threads_per_warp: 32
+        shared_memory_per_block: 49152
+        shared_memory_per_core: 233472
+        threads_per_core_limit: 2048
+        core_count: 152
+        fpus_per_core: 128
+        block_dim_limit_x: 2147483647
+        block_dim_limit_y: 65535
+        block_dim_limit_z: 65535
+        memory_bandwidth: 7928064000000
+        l2_cache_size: 135528448
+        clock_rate_ghz: 2.062
+        device_memory_size: 197897617408
+        shared_memory_per_block_optin: 232448
+        cuda_compute_capability {
+          major: 10
+          feature_extension: ACCELERATED_FEATURES
+        }
+        registers_per_core_limit: 65536
+        registers_per_block_limit: 65536
+        scalar_unit_description {
+          rate_infos {
+            key: 4
+            value {
+              clock_rate_ghz: 2.062
+              units_per_core: 64
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 8
+            value {
+              clock_rate_ghz: 2.062
+              units_per_core: 64
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 10
+            value {
+              clock_rate_ghz: 2.062
+              units_per_core: 128
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 11
+            value {
+              clock_rate_ghz: 2.062
+              units_per_core: 128
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 12
+            value {
+              clock_rate_ghz: 2.062
+              units_per_core: 64
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 16
+            value {
+              clock_rate_ghz: 2.062
+              units_per_core: 128
+              ops_per_clock: 1
+            }
+          }
+        }
+        matrix_unit_description {
+          rate_infos {
+            key: 2
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 6
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 10
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 11
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 512
+            }
+          }
+          rate_infos {
+            key: 12
+            value {
+              clock_rate_ghz: 2.062
+              units_per_core: 4
+              ops_per_clock: 16
+            }
+          }
+          rate_infos {
+            key: 16
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 19
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 20
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 23
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 24
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 25
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 28
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 29
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 32
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 4096
+            }
+          }
+          rate_infos {
+            key: 33
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 35
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 36
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+        }
+        driver_version: "13.2.0"
+        kernel_mode_driver_version: "580.126.20"
+        runtime_version: "13.2.0"
+        compile_time_toolkit_version: "13.2.0"
+        dnn_version: "9.23.0"
+        cub_version: "3.1.2"
+        device_interconnect_info {
+          active_links: 18
+        }
+        device_vendor: "NVIDIA Corporation"
+        platform_version: "Compute Capability 10.0a"
+        pci_bus_id: "<undefined>"
+        name: "NVIDIA GB200"
+        model_str: "sm_10.0a with 197897617408B RAM, 152 cores, 2062000KHz clock, 3996000KHz mem clock, 135528448B L2$"
+        numa_node: -1
+        thread_dim_limit_x: 1024
+        thread_dim_limit_y: 1024
+        thread_dim_limit_z: 64
+        device_address_bits: 64
+        pcie_bandwidth: 2147483648
+        ecc_enabled: true
+        mem_clock_ghz: 3.996
+        reserved_shared_memory_per_block: 1024
+        max_blocks_per_multiprocessor: 32
+        collective_memory_granularity: 2097152
+      }
+      platform_name: "CUDA"
+      device_description_str: "NVIDIA GB200"
+      runtime_version {
+        major: 13
+        minor: 2
+      }
+    }
+    host_target_machine_options {
+      triple: "aarch64-unknown-linux-gnu"
+      cpu: "neoverse-v2"
+    }
+    num_devices_per_process: 2
+  }
+}
diff --git a/third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_gb200/v4/ReduceScatter.pbtxt b/third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_gb200/v4/ReduceScatter.pbtxt
new file mode 100644
index 00000000000..92f90e61e36
--- /dev/null
+++ b/third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_gb200/v4/ReduceScatter.pbtxt
@@ -0,0 +1,1901 @@
+# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto
+# proto-message: xla.HumanReadableAotExecutable
+executable_and_options {
+  compile_options {
+    executable_build_options {
+      device_ordinal: -1
+      result_layout {
+        element_type: U32
+        dimensions: 4
+        layout {
+          minor_to_major: 0
+          tail_padding_alignment_in_elements: 1
+        }
+        is_dynamic_dimension: false
+      }
+      debug_options {
+        xla_disable_hlo_passes: "constant_folding"
+        xla_backend_optimization_level: 3
+        xla_eliminate_hlo_implicit_broadcast: true
+        xla_cpu_multi_thread_eigen: true
+        xla_gpu_cuda_data_dir: "./cuda_sdk_lib"
+        xla_llvm_enable_alias_scope_metadata: true
+        xla_llvm_enable_noalias_metadata: true
+        xla_llvm_enable_invariant_load_metadata: true
+        xla_llvm_disable_expensive_passes: false
+        xla_cpu_use_onednn: false
+        xla_cpu_enable_fast_math: false
+        xla_gpu_enable_fast_min_max: false
+        xla_force_host_platform_device_count: 1
+        xla_gpu_enable_dynamic_slice_fusion: true
+        xla_hlo_evaluator_use_fast_path: true
+        xla_dump_hlo_as_html: false
+        xla_cpu_fast_math_honor_nans: true
+        xla_cpu_fast_math_honor_infs: true
+        xla_allow_excess_precision: true
+        xla_gpu_autotune_level: 4
+        xla_cpu_fast_math_honor_division: true
+        xla_cpu_fast_math_honor_functions: true
+        xla_dump_include_timestamp: false
+        xla_dump_max_hlo_modules: -1
+        xla_cpu_enable_xprof_traceme: false
+        xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false
+        xla_cpu_enable_fast_min_max: true
+        xla_multiheap_size_constraint_per_heap: -1
+        xla_dump_module_metadata: false
+        xla_gpu_force_compilation_parallelism: 1
+        xla_dump_fusion_visualization: false
+        xla_gpu_all_reduce_combine_threshold_bytes: 31457287
+        xla_gpu_nccl_termination_timeout_seconds: -1
+        xla_dump_hlo_as_long_text: true
+        xla_gpu_enable_cublaslt: true
+        xla_gpu_shape_checks: RUNTIME
+        xla_gpu_use_runtime_fusion: false
+        xla_dump_latency_hiding_schedule: false
+        xla_dump_enable_mlir_pretty_form: true
+        xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP
+        xla_gpu_enable_triton_gemm: true
+        xla_gpu_enable_cudnn_int8x32_convolution_reordering: true
+        xla_gpu_triton_gemm_any: true
+        xla_gpu_enable_while_loop_reduce_scatter_code_motion: false
+        xla_gpu_collective_inflation_factor: 1
+        xla_gpu_graph_min_graph_size: 5
+        xla_gpu_enable_reassociation_for_converted_ar: true
+        xla_gpu_pgle_profile_file_or_directory_path: ""
+        xla_gpu_all_gather_combine_threshold_bytes: 31457287
+        xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287
+        xla_gpu_enable_highest_priority_async_stream: true
+        xla_gpu_exhaustive_tiling_search: false
+        xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0
+        xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1
+        xla_gpu_redzone_padding_bytes: 8388608
+        xla_gpu_copy_insertion_use_region_analysis: false
+        xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807
+        xla_gpu_collect_cost_model_stats: false
+        xla_gpu_enable_pipelined_p2p: false
+        xla_gpu_cublas_fallback: true
+        xla_gpu_enable_while_loop_double_buffering: false
+        xla_gpu_filter_kernels_spilling_registers_on_autotuning: true
+        xla_debug_buffer_assignment_show_max: 15
+        xla_detailed_logging: true
+        xla_enable_dumping: true
+        xla_gpu_enable_all_gather_combine_by_dim: false
+        xla_gpu_enable_analytical_latency_estimator: false
+        xla_gpu_llvm_verification_level: 0
+        xla_gpu_enable_reduce_scatter_combine_by_dim: false
+        xla_gpu_enable_command_buffer: COLLECTIVES_KERNEL
+        xla_gpu_enable_command_buffer: CONDITIONAL
+        xla_gpu_enable_command_buffer: CUBLAS
+        xla_gpu_enable_command_buffer: CUBLASLT
+        xla_gpu_enable_command_buffer: CUDNN
+        xla_gpu_enable_command_buffer: CUSTOM_CALL
+        xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION
+        xla_gpu_enable_command_buffer: FUSION
+        xla_gpu_enable_cub_radix_sort: true
+        xla_gpu_memory_limit_slop_factor: 95
+        xla_gpu_target_config_filename: ""
+        xla_gpu_enable_cudnn_layer_norm: false
+        xla_gpu_threshold_for_windowed_einsum_mib: 100000
+        xla_gpu_enable_nccl_user_buffers: false
+        xla_gpu_enable_libnvptxcompiler: true
+        xla_gpu_enable_nccl_comm_splitting: true
+        xla_gpu_nccl_collective_max_nchannels: 0
+        xla_gpu_nccl_p2p_max_nchannels: 0
+        xla_gpu_nccl_init_max_rank_per_root_ratio: 0
+        xla_gpu_multi_streamed_windowed_einsum: true
+        xla_gpu_gemm_rewrite_size_threshold: 100
+        xla_gpu_require_complete_aot_autotune_results: false
+        xla_gpu_cudnn_gemm_fusion_level: 0
+        xla_gpu_use_memcpy_local_p2p: false
+        xla_gpu_autotune_max_solutions: 0
+        xla_dump_large_constants: false
+        xla_gpu_verify_triton_fusion_numerics: false
+        xla_reduce_window_rewrite_base_length: 16
+        xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL
+        xla_gpu_enable_host_memory_offloading: false
+        xla_gpu_exclude_nondeterministic_ops: true
+        xla_gpu_nccl_terminate_on_error: false
+        xla_gpu_shard_autotuning: true
+        xla_gpu_enable_approx_costly_collectives: false
+        xla_cpu_prefer_vector_width: 256
+        xla_gpu_per_fusion_autotune_cache_dir: ""
+        xla_cmd_buffer_trace_cache_size: 16
+        xla_gpu_temp_buffer_use_separate_color: false
+        xla_syntax_sugar_async_ops: false
+        xla_gpu_autotune_gemm_rtol: 0.1
+        xla_enable_command_buffers_during_profiling: false
+        xla_gpu_cudnn_gemm_max_plans: 5
+        xla_cpu_parallel_codegen_split_count: 32
+        xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE
+        xla_gpu_executable_warn_stuck_timeout_seconds: 10
+        xla_gpu_executable_terminate_timeout_seconds: 30
+        xla_gpu_experimental_disable_binary_libraries: false
+        xla_gpu_dot_merger_threshold_mb: 64
+        xla_cpu_max_isa: "NEON"
+        xla_gpu_experimental_enable_fusion_block_level_rewriter: false
+        xla_enable_fast_math: false
+        xla_gpu_experimental_parallel_collective_overlap_limit: 1
+        xla_cpu_copy_insertion_use_region_analysis: false
+        xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1
+        xla_gpu_experimental_enable_triton_heroless_priority_fusion: false
+        xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN
+        xla_gpu_experimental_stream_annotation: true
+        xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO
+        xla_pjrt_allow_auto_layout_in_hlo: false
+        xla_gpu_enable_scatter_determinism_expander: false
+        xla_gpu_require_exclusive_lock: false
+        xla_gpu_generate_debug_info: false
+        xla_gpu_generate_line_info: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false
+        xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE
+        xla_gpu_fail_ptx_compilation_on_register_spilling: false
+        xla_gpu_enable_analytical_sol_latency_estimator: true
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_prep_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_size_bytes"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "gpus_per_node"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nccl_op_launch_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nic_speed_gbps"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "rtt_us"
+          value: "-1"
+        }
+        xla_cpu_use_xnnpack: true
+        xla_gpu_experimental_pack_dot_operands_along_k_dimension: true
+        xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false
+        xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED
+        xla_gpu_experimental_enable_heuristic_collective_combining: false
+        xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false
+        xla_hlo_pass_fix_detect_cycles: false
+        xla_cpu_generate_unique_c_style_kernel_entry_points: false
+        xla_gpu_experimental_collective_cse_distance_threshold: 0
+        xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true
+        xla_gpu_experimental_collective_perf_table_path: ""
+        xla_gpu_collective_permute_combine_threshold_bytes: 31457287
+        xla_unsupported_crash_on_hlo_pass_noop_change: false
+        xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false
+        xla_dump_full_hlo_config: true
+        xla_gpu_unsupported_enable_triton_multi_output_fusion: false
+        xla_gpu_experimental_matmul_perf_table_path: ""
+        xla_gpu_unsupported_enable_all_reduce_decomposer: false
+        xla_gpu_nccl_blocking_communicators: true
+        xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20
+        xla_gpu_first_collective_call_terminate_timeout_seconds: 40
+        xla_gpu_nccl_async_execution: false
+        xla_cpu_emitter_verification_level: 1
+        xla_early_exit_with_layouts: false
+        xla_gpu_experimental_use_ragged_dot_fusion: false
+        xla_detect_unstable_reductions: DETECTION_MODE_NONE
+        xla_gpu_command_buffer_scheduling_mode: LHS
+        xla_gpu_experimental_enable_nccl_symmetric_buffers: false
+        xla_disable_automatic_host_compute_offload: false
+        xla_gpu_experimental_enable_fusion_autotuner: true
+        xla_gpu_experimental_scaled_dot_with_triton: true
+        xla_gpu_command_buffer_unroll_loops: false
+        xla_cpu_experimental_onednn_custom_call: false
+        xla_gpu_experimental_enable_checksum_tracing_on_thunks: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true
+        xla_cpu_collective_call_terminate_timeout_seconds: 40
+        xla_cpu_collective_call_warn_stuck_seconds: 20
+        xla_keep_shardings_after_spmd: false
+        xla_gpu_experimental_enable_triton_warp_specialization: false
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE
+        xla_cpu_enable_platform_dependent_math: true
+        xla_gpu_detect_nan: DETECTION_MODE_NONE
+        xla_gpu_detect_inf: DETECTION_MODE_NONE
+        xla_enable_enzyme_comms_opt: false
+        xla_gpu_experimental_enable_buffer_saver_on_thunks: false
+        xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE
+        xla_enable_scoped_logging_timers: true
+        xla_gpu_executable_embed_debug_info: true
+        xla_cpu_collective_timeout_seconds: 1800
+        xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false
+        xla_gpu_enable_pipelined_host_offloading: false
+        xla_gpu_default_to_alg_dot_bf16_bf16_f32: false
+        xla_gpu_experimental_autotune_backends: CUDNN
+        xla_gpu_experimental_autotune_backends: TRITON
+        xla_gpu_experimental_autotune_backends: CUBLASLT
+        xla_gpu_experimental_autotune_backends: HIPBLASLT
+        xla_gpu_experimental_autotune_backends: MIOPEN
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER
+        xla_gpu_experimental_autotune_backends: NATIVE_EMITTER
+        xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER
+        xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION
+        xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: YNNPACK
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU
+        xla_gpu_experimental_autotune_backends: TENSOR_IR
+        xla_gpu_experimental_all_fusions_with_triton: false
+        xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT
+        xla_gpu_rocm_max_trace_events: 4194304
+        xla_gpu_execution_terminate_timeout: "inf"
+        xla_enable_hlo_sharding_v3: true
+        xla_gpu_log_minmax: false
+        xla_gpu_print_compilation_stats: false
+        xla_gpu_experimental_enable_conv_fusion: false
+        xla_gpu_experimental_enable_tiling_propagation: true
+        xla_gpu_execution_progress_tracking: 0
+        xla_gpu_experimental_max_unroll_factor: 32
+        xla_gpu_enable_pdl: true
+        xla_gpu_collective_combine_threshold_count: 2048
+        xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true
+        xla_recognize_reduction_optimization_level: 0
+        xla_gpu_experimental_parallel_async_compute_limit: 2
+        xla_dump_buffer_assignment_analysis: true
+        xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE
+        xla_gpu_experimental_enable_collective_multi_streaming: false
+        xla_gpu_collective_permute_connected_components: false
+        xla_gpu_experimental_force_split_k: 0
+        xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_experimental_gemm_fusion_v2: true
+        xla_gpu_fusion_autotune_top_k_configs: 1
+        xla_gpu_native_emitter_tune_unroll_factor_for_loops: false
+        xla_enable_rgv3_materialization: true
+        xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK
+        xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false
+        xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_executable_num_compute_streams: 0
+        xla_gpu_executable_num_communication_streams: 0
+        xla_gpu_experimental_collective_start_as_early_as_possible: false
+        xla_cpu_experimental_enable_tiling_propagation: true
+        xla_gpu_enable_dus_accumulator_zero_init_elimination: false
+        xla_gpu_experimental_thunk_buffer_debug_module_outputs: false
+        xla_cpu_ftz: true
+        xla_gpu_enable_allocator_spatial_partitioning: true
+        xla_dump_compact_gte: true
+        xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES
+        xla_dump_hlo_as_riegeli: false
+        xla_gpu_experimental_use_ragged_dot_grouped_gemm: true
+        xla_dump_hlo_to_subfolder: false
+        xla_gpu_enable_pdl_launch: true
+        xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED
+        xla_gpu_enable_nccl_user_buffers_in_default_space: false
+        xla_sdy_export_all_reduce_scatter: false
+        xla_gpu_enable_gxl_ragged_all_to_all: false
+        xla_gpu_gxl_scratch_size_bytes: 67108864
+        xla_gpu_use_new_autotune_cache_format: true
+        xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE
+        xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON
+        xla_spmd_enable_dynamic_slice_collective_broadcast: false
+        xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false
+        xla_gpu_async_copy_min_bytes: -1
+        xla_gpu_blas_max_algorithms: 0
+        xla_gpu_experimental_enable_selective_memcpy_overlap: false
+        xla_gpu_experimental_enable_same_shape_multi_output_fusion: false
+        xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1
+        xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1
+        xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240
+        xla_cpu_use_new_xtile_lowering: true
+        xla_gpu_trace_annotation_level: 0
+        xla_gpu_experimental_cost_model_gemm_tiling_default: false
+        xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false
+        xla_gpu_all_reduce_splitter_ignore_profitability_check: false
+        xla_gpu_enable_cupti_multi_subscriber: true
+        xla_enable_hlo_modules_upload: true
+        xla_gpu_collective_domain_assignment: ""
+        xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1
+        xla_compile_all_supported_configs: false
+        xla_gpu_unsupported_use_cross_host_one_shot_kernel: ALLCOLLECTIVES
+        xla_deduplicate_backend_configs_min_size: 9223372036854775807
+        xla_force_config: ""
+        xla_candidate_configs_file: ""
+        xla_gpu_enable_persistent_symmetric_memory: false
+        xla_gpu_experimental_enable_raft_for_stable_topk: false
+        xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND
+        xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET
+        xla_gpu_experimental_emit_collective_reduce: false
+        xla_gpu_experimental_vmm_disabled: false
+        xla_gpu_experimental_enable_dynamic_slice_table_offsets: false
+        xla_gpu_device_execution_terminate_timeout: "inf"
+        xla_gpu_topology_filename: ""
+        xla_gpu_experimental_enable_dynamic_slice_extended_offsets: false
+        xla_gpu_experimental_deterministic_ops_triton: DETERMINISTIC_OPS_TRITON_MODE_OFF
+        xla_gpu_experimental_enable_tensor_ir: false
+      }
+      num_replicas: 2
+      num_partitions: 1
+      device_assignment {
+        replica_count: 2
+        computation_count: 1
+        computation_devices {
+          replica_device_ids: 0
+          replica_device_ids: 1
+        }
+      }
+      allow_spmd_sharding_propagation_to_output: false
+      comp_envs {
+      }
+      allow_spmd_sharding_propagation_to_parameters: false
+      process_count: 1
+      memory_fitting_level: EFFORT_O2
+    }
+  }
+  pjrt_client_name: "PjRtStreamExecutorClient"
+}
+gpu_executable {
+  hlo_module_with_config {
+    hlo_module {
+      name: "test"
+      entry_computation_name: "main"
+      computations {
+        name: "add"
+        instructions {
+          name: "lhs"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "lhs"
+          }
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "rhs"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "rhs"
+          }
+          parameter_number: 1
+          id: 1
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "add.1"
+          opcode: "add"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "add.1"
+          }
+          id: 2
+          operand_ids: 0
+          operand_ids: 1
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          parameters {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          parameters {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          result {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          parameter_names: "lhs"
+          parameter_names: "rhs"
+        }
+        root_id: 2
+      }
+      computations {
+        name: "fused_select"
+        instructions {
+          name: "param_2.2"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "param_2.2"
+          }
+          parameter_number: 2
+          id: 8589934597
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "zero_2"
+          opcode: "constant"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+          }
+          literal {
+            shape {
+              element_type: U32
+              layout {
+                tail_padding_alignment_in_elements: 1
+              }
+            }
+            u32s: 0
+          }
+          id: 8589934598
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "p.2"
+          opcode: "compare"
+          shape {
+            element_type: PRED
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "p.2"
+          }
+          id: 8589934596
+          operand_ids: 8589934597
+          operand_ids: 8589934598
+          comparison_direction: "EQ"
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+          comparison_order: "TOTAL"
+        }
+        instructions {
+          name: "pb.2"
+          opcode: "broadcast"
+          shape {
+            element_type: PRED
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "pb.2"
+          }
+          id: 8589934593
+          operand_ids: 8589934596
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "param_0.1"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "param_0.1"
+          }
+          id: 8589934594
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "param_1.1"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "param_1.1"
+          }
+          parameter_number: 1
+          id: 8589934595
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "data.2"
+          opcode: "select"
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "data.2"
+          }
+          id: 8589934592
+          operand_ids: 8589934593
+          operand_ids: 8589934594
+          operand_ids: 8589934595
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          parameters {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          parameters {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          parameters {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          result {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          parameter_names: "param_0.1"
+          parameter_names: "param_1.1"
+          parameter_names: "param_2.2"
+        }
+        id: 2
+        root_id: 8589934592
+        is_fusion_computation: true
+      }
+      computations {
+        name: "main"
+        instructions {
+          name: "c0_1"
+          opcode: "constant"
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+          }
+          literal {
+            shape {
+              element_type: U32
+              dimensions: 8
+              layout {
+                minor_to_major: 0
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+            }
+            u32s: 1
+            u32s: 2
+            u32s: 3
+            u32s: 4
+            u32s: 5
+            u32s: 6
+            u32s: 7
+            u32s: 8
+          }
+          id: 4294967297
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "c1_1"
+          opcode: "constant"
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+          }
+          literal {
+            shape {
+              element_type: U32
+              dimensions: 8
+              layout {
+                minor_to_major: 0
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+            }
+            u32s: 10
+            u32s: 11
+            u32s: 12
+            u32s: 13
+            u32s: 14
+            u32s: 15
+            u32s: 16
+            u32s: 17
+          }
+          id: 4294967298
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "id.1"
+          opcode: "replica-id"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "id.1"
+          }
+          id: 4294967296
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "loop_select_fusion"
+          opcode: "fusion"
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "loop_select_fusion"
+          }
+          fusion_kind: "kLoop"
+          id: 4294967299
+          operand_ids: 4294967297
+          operand_ids: 4294967298
+          operand_ids: 4294967296
+          called_computation_ids: 2
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "ars.3"
+          opcode: "reduce-scatter"
+          shape {
+            element_type: U32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "reduce-scatter-start"
+          }
+          dimensions: 0
+          id: 4294967300
+          operand_ids: 4294967299
+          called_computation_ids: 0
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+          collective_device_list {
+          }
+          backend_config_payload {
+            id: 0
+          }
+        }
+        program_shape {
+          result {
+            element_type: U32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+        }
+        id: 1
+        root_id: 4294967300
+      }
+      host_program_shape {
+        result {
+          element_type: U32
+          dimensions: 4
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+      }
+      entry_computation_id: 1
+      schedule {
+        sequences {
+          key: 0
+          value {
+            instruction_ids: 1
+            instruction_ids: 0
+            instruction_ids: 2
+          }
+        }
+        sequences {
+          key: 1
+          value {
+            instruction_ids: 4294967296
+            instruction_ids: 4294967298
+            instruction_ids: 4294967297
+            instruction_ids: 4294967299
+            instruction_ids: 4294967300
+          }
+        }
+      }
+      input_output_alias {
+      }
+      device_assignment {
+        replica_count: 2
+        computation_count: 1
+        computation_devices {
+          replica_device_ids: 0
+          replica_device_ids: 1
+        }
+      }
+      buffer_donor {
+      }
+      frontend_attributes {
+        map {
+          key: "fingerprint_before_lhs"
+          value: "c36d8870be58d238935bebefbbc5b827"
+        }
+      }
+      payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}"
+    }
+    config {
+      entry_computation_layout {
+        result {
+          element_type: U32
+          dimensions: 4
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+      }
+      replica_count: 2
+      num_partitions: 1
+      intra_op_parallelism_threads: -1
+      debug_options {
+        xla_disable_hlo_passes: "constant_folding"
+        xla_backend_optimization_level: 3
+        xla_eliminate_hlo_implicit_broadcast: true
+        xla_cpu_multi_thread_eigen: true
+        xla_gpu_cuda_data_dir: "./cuda_sdk_lib"
+        xla_llvm_enable_alias_scope_metadata: true
+        xla_llvm_enable_noalias_metadata: true
+        xla_llvm_enable_invariant_load_metadata: true
+        xla_llvm_disable_expensive_passes: false
+        xla_cpu_use_onednn: false
+        xla_cpu_enable_fast_math: false
+        xla_gpu_enable_fast_min_max: false
+        xla_force_host_platform_device_count: 1
+        xla_gpu_enable_dynamic_slice_fusion: true
+        xla_hlo_evaluator_use_fast_path: true
+        xla_dump_hlo_as_html: false
+        xla_cpu_fast_math_honor_nans: true
+        xla_cpu_fast_math_honor_infs: true
+        xla_allow_excess_precision: true
+        xla_gpu_autotune_level: 4
+        xla_cpu_fast_math_honor_division: true
+        xla_cpu_fast_math_honor_functions: true
+        xla_dump_include_timestamp: false
+        xla_dump_max_hlo_modules: -1
+        xla_cpu_enable_xprof_traceme: false
+        xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false
+        xla_cpu_enable_fast_min_max: true
+        xla_multiheap_size_constraint_per_heap: -1
+        xla_dump_module_metadata: false
+        xla_gpu_force_compilation_parallelism: 1
+        xla_dump_fusion_visualization: false
+        xla_gpu_all_reduce_combine_threshold_bytes: 31457287
+        xla_gpu_nccl_termination_timeout_seconds: -1
+        xla_dump_hlo_as_long_text: true
+        xla_gpu_enable_cublaslt: true
+        xla_gpu_shape_checks: RUNTIME
+        xla_gpu_use_runtime_fusion: false
+        xla_dump_latency_hiding_schedule: false
+        xla_dump_enable_mlir_pretty_form: true
+        xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP
+        xla_gpu_enable_triton_gemm: true
+        xla_gpu_enable_cudnn_int8x32_convolution_reordering: true
+        xla_gpu_triton_gemm_any: true
+        xla_gpu_enable_while_loop_reduce_scatter_code_motion: false
+        xla_gpu_collective_inflation_factor: 1
+        xla_gpu_graph_min_graph_size: 5
+        xla_gpu_enable_reassociation_for_converted_ar: true
+        xla_gpu_pgle_profile_file_or_directory_path: ""
+        xla_gpu_all_gather_combine_threshold_bytes: 31457287
+        xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287
+        xla_gpu_enable_highest_priority_async_stream: true
+        xla_gpu_exhaustive_tiling_search: false
+        xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0
+        xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1
+        xla_gpu_redzone_padding_bytes: 8388608
+        xla_gpu_copy_insertion_use_region_analysis: false
+        xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807
+        xla_gpu_collect_cost_model_stats: false
+        xla_gpu_enable_pipelined_p2p: false
+        xla_gpu_cublas_fallback: true
+        xla_gpu_enable_while_loop_double_buffering: false
+        xla_gpu_filter_kernels_spilling_registers_on_autotuning: true
+        xla_debug_buffer_assignment_show_max: 15
+        xla_detailed_logging: true
+        xla_enable_dumping: true
+        xla_gpu_enable_all_gather_combine_by_dim: false
+        xla_gpu_enable_analytical_latency_estimator: false
+        xla_gpu_llvm_verification_level: 0
+        xla_gpu_enable_reduce_scatter_combine_by_dim: false
+        xla_gpu_enable_command_buffer: COLLECTIVES_KERNEL
+        xla_gpu_enable_command_buffer: CONDITIONAL
+        xla_gpu_enable_command_buffer: CUBLAS
+        xla_gpu_enable_command_buffer: CUBLASLT
+        xla_gpu_enable_command_buffer: CUDNN
+        xla_gpu_enable_command_buffer: CUSTOM_CALL
+        xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION
+        xla_gpu_enable_command_buffer: FUSION
+        xla_gpu_enable_cub_radix_sort: true
+        xla_gpu_memory_limit_slop_factor: 95
+        xla_gpu_target_config_filename: ""
+        xla_gpu_enable_cudnn_layer_norm: false
+        xla_gpu_threshold_for_windowed_einsum_mib: 100000
+        xla_gpu_enable_nccl_user_buffers: false
+        xla_gpu_enable_libnvptxcompiler: true
+        xla_gpu_enable_nccl_comm_splitting: true
+        xla_gpu_nccl_collective_max_nchannels: 0
+        xla_gpu_nccl_p2p_max_nchannels: 0
+        xla_gpu_nccl_init_max_rank_per_root_ratio: 0
+        xla_gpu_multi_streamed_windowed_einsum: true
+        xla_gpu_gemm_rewrite_size_threshold: 100
+        xla_gpu_require_complete_aot_autotune_results: false
+        xla_gpu_cudnn_gemm_fusion_level: 0
+        xla_gpu_use_memcpy_local_p2p: false
+        xla_gpu_autotune_max_solutions: 0
+        xla_dump_large_constants: false
+        xla_gpu_verify_triton_fusion_numerics: false
+        xla_reduce_window_rewrite_base_length: 16
+        xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL
+        xla_gpu_enable_host_memory_offloading: false
+        xla_gpu_exclude_nondeterministic_ops: true
+        xla_gpu_nccl_terminate_on_error: false
+        xla_gpu_shard_autotuning: true
+        xla_gpu_enable_approx_costly_collectives: false
+        xla_cpu_prefer_vector_width: 256
+        xla_gpu_per_fusion_autotune_cache_dir: ""
+        xla_cmd_buffer_trace_cache_size: 16
+        xla_gpu_temp_buffer_use_separate_color: false
+        xla_syntax_sugar_async_ops: false
+        xla_gpu_autotune_gemm_rtol: 0.1
+        xla_enable_command_buffers_during_profiling: false
+        xla_gpu_cudnn_gemm_max_plans: 5
+        xla_cpu_parallel_codegen_split_count: 32
+        xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE
+        xla_gpu_executable_warn_stuck_timeout_seconds: 10
+        xla_gpu_executable_terminate_timeout_seconds: 30
+        xla_gpu_experimental_disable_binary_libraries: false
+        xla_gpu_dot_merger_threshold_mb: 64
+        xla_cpu_max_isa: "NEON"
+        xla_gpu_experimental_enable_fusion_block_level_rewriter: false
+        xla_enable_fast_math: false
+        xla_gpu_experimental_parallel_collective_overlap_limit: 1
+        xla_cpu_copy_insertion_use_region_analysis: false
+        xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1
+        xla_gpu_experimental_enable_triton_heroless_priority_fusion: false
+        xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN
+        xla_gpu_experimental_stream_annotation: true
+        xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO
+        xla_pjrt_allow_auto_layout_in_hlo: false
+        xla_gpu_enable_scatter_determinism_expander: false
+        xla_gpu_require_exclusive_lock: false
+        xla_gpu_generate_debug_info: false
+        xla_gpu_generate_line_info: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false
+        xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE
+        xla_gpu_fail_ptx_compilation_on_register_spilling: false
+        xla_gpu_enable_analytical_sol_latency_estimator: true
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_prep_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_size_bytes"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "gpus_per_node"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nccl_op_launch_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nic_speed_gbps"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "rtt_us"
+          value: "-1"
+        }
+        xla_cpu_use_xnnpack: true
+        xla_gpu_experimental_pack_dot_operands_along_k_dimension: true
+        xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false
+        xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED
+        xla_gpu_experimental_enable_heuristic_collective_combining: false
+        xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false
+        xla_hlo_pass_fix_detect_cycles: false
+        xla_cpu_generate_unique_c_style_kernel_entry_points: false
+        xla_gpu_experimental_collective_cse_distance_threshold: 0
+        xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true
+        xla_gpu_experimental_collective_perf_table_path: ""
+        xla_gpu_collective_permute_combine_threshold_bytes: 31457287
+        xla_unsupported_crash_on_hlo_pass_noop_change: false
+        xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false
+        xla_dump_full_hlo_config: true
+        xla_gpu_unsupported_enable_triton_multi_output_fusion: false
+        xla_gpu_experimental_matmul_perf_table_path: ""
+        xla_gpu_unsupported_enable_all_reduce_decomposer: false
+        xla_gpu_nccl_blocking_communicators: true
+        xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20
+        xla_gpu_first_collective_call_terminate_timeout_seconds: 40
+        xla_gpu_nccl_async_execution: false
+        xla_cpu_emitter_verification_level: 1
+        xla_early_exit_with_layouts: false
+        xla_gpu_experimental_use_ragged_dot_fusion: false
+        xla_detect_unstable_reductions: DETECTION_MODE_NONE
+        xla_gpu_command_buffer_scheduling_mode: LHS
+        xla_gpu_experimental_enable_nccl_symmetric_buffers: false
+        xla_disable_automatic_host_compute_offload: false
+        xla_gpu_experimental_enable_fusion_autotuner: true
+        xla_gpu_experimental_scaled_dot_with_triton: true
+        xla_gpu_command_buffer_unroll_loops: false
+        xla_cpu_experimental_onednn_custom_call: false
+        xla_gpu_experimental_enable_checksum_tracing_on_thunks: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true
+        xla_cpu_collective_call_terminate_timeout_seconds: 40
+        xla_cpu_collective_call_warn_stuck_seconds: 20
+        xla_keep_shardings_after_spmd: false
+        xla_gpu_experimental_enable_triton_warp_specialization: false
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE
+        xla_cpu_enable_platform_dependent_math: true
+        xla_gpu_detect_nan: DETECTION_MODE_NONE
+        xla_gpu_detect_inf: DETECTION_MODE_NONE
+        xla_enable_enzyme_comms_opt: false
+        xla_gpu_experimental_enable_buffer_saver_on_thunks: false
+        xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE
+        xla_enable_scoped_logging_timers: true
+        xla_gpu_executable_embed_debug_info: true
+        xla_cpu_collective_timeout_seconds: 1800
+        xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false
+        xla_gpu_enable_pipelined_host_offloading: false
+        xla_gpu_default_to_alg_dot_bf16_bf16_f32: false
+        xla_gpu_experimental_autotune_backends: CUDNN
+        xla_gpu_experimental_autotune_backends: TRITON
+        xla_gpu_experimental_autotune_backends: CUBLASLT
+        xla_gpu_experimental_autotune_backends: HIPBLASLT
+        xla_gpu_experimental_autotune_backends: MIOPEN
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER
+        xla_gpu_experimental_autotune_backends: NATIVE_EMITTER
+        xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER
+        xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION
+        xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: YNNPACK
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU
+        xla_gpu_experimental_autotune_backends: TENSOR_IR
+        xla_gpu_experimental_all_fusions_with_triton: false
+        xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT
+        xla_gpu_rocm_max_trace_events: 4194304
+        xla_gpu_execution_terminate_timeout: "inf"
+        xla_enable_hlo_sharding_v3: true
+        xla_gpu_log_minmax: false
+        xla_gpu_print_compilation_stats: false
+        xla_gpu_experimental_enable_conv_fusion: false
+        xla_gpu_experimental_enable_tiling_propagation: true
+        xla_gpu_execution_progress_tracking: 0
+        xla_gpu_experimental_max_unroll_factor: 32
+        xla_gpu_enable_pdl: true
+        xla_gpu_collective_combine_threshold_count: 2048
+        xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true
+        xla_recognize_reduction_optimization_level: 0
+        xla_gpu_experimental_parallel_async_compute_limit: 2
+        xla_dump_buffer_assignment_analysis: true
+        xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE
+        xla_gpu_experimental_enable_collective_multi_streaming: false
+        xla_gpu_collective_permute_connected_components: false
+        xla_gpu_experimental_force_split_k: 0
+        xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_experimental_gemm_fusion_v2: true
+        xla_gpu_fusion_autotune_top_k_configs: 1
+        xla_gpu_native_emitter_tune_unroll_factor_for_loops: false
+        xla_enable_rgv3_materialization: true
+        xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK
+        xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false
+        xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_executable_num_compute_streams: 0
+        xla_gpu_executable_num_communication_streams: 0
+        xla_gpu_experimental_collective_start_as_early_as_possible: false
+        xla_cpu_experimental_enable_tiling_propagation: true
+        xla_gpu_enable_dus_accumulator_zero_init_elimination: false
+        xla_gpu_experimental_thunk_buffer_debug_module_outputs: false
+        xla_cpu_ftz: true
+        xla_gpu_enable_allocator_spatial_partitioning: true
+        xla_dump_compact_gte: true
+        xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES
+        xla_dump_hlo_as_riegeli: false
+        xla_gpu_experimental_use_ragged_dot_grouped_gemm: true
+        xla_dump_hlo_to_subfolder: false
+        xla_gpu_enable_pdl_launch: true
+        xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED
+        xla_gpu_enable_nccl_user_buffers_in_default_space: false
+        xla_sdy_export_all_reduce_scatter: false
+        xla_gpu_enable_gxl_ragged_all_to_all: false
+        xla_gpu_gxl_scratch_size_bytes: 67108864
+        xla_gpu_use_new_autotune_cache_format: true
+        xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE
+        xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON
+        xla_spmd_enable_dynamic_slice_collective_broadcast: false
+        xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false
+        xla_gpu_async_copy_min_bytes: -1
+        xla_gpu_blas_max_algorithms: 0
+        xla_gpu_experimental_enable_selective_memcpy_overlap: false
+        xla_gpu_experimental_enable_same_shape_multi_output_fusion: false
+        xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1
+        xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1
+        xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240
+        xla_cpu_use_new_xtile_lowering: true
+        xla_gpu_trace_annotation_level: 0
+        xla_gpu_experimental_cost_model_gemm_tiling_default: false
+        xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false
+        xla_gpu_all_reduce_splitter_ignore_profitability_check: false
+        xla_gpu_enable_cupti_multi_subscriber: true
+        xla_enable_hlo_modules_upload: true
+        xla_gpu_collective_domain_assignment: ""
+        xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1
+        xla_compile_all_supported_configs: false
+        xla_gpu_unsupported_use_cross_host_one_shot_kernel: ALLCOLLECTIVES
+        xla_deduplicate_backend_configs_min_size: 9223372036854775807
+        xla_force_config: ""
+        xla_candidate_configs_file: ""
+        xla_gpu_enable_persistent_symmetric_memory: false
+        xla_gpu_experimental_enable_raft_for_stable_topk: false
+        xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND
+        xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET
+        xla_gpu_experimental_emit_collective_reduce: false
+        xla_gpu_experimental_vmm_disabled: false
+        xla_gpu_experimental_enable_dynamic_slice_table_offsets: false
+        xla_gpu_device_execution_terminate_timeout: "inf"
+        xla_gpu_topology_filename: ""
+        xla_gpu_experimental_enable_dynamic_slice_extended_offsets: false
+        xla_gpu_experimental_deterministic_ops_triton: DETERMINISTIC_OPS_TRITON_MODE_OFF
+        xla_gpu_experimental_enable_tensor_ir: false
+      }
+      static_device_assignment {
+        replica_count: 2
+        computation_count: 1
+        computation_devices {
+          replica_device_ids: 0
+          replica_device_ids: 1
+        }
+      }
+      allow_spmd_sharding_propagation_to_output: false
+      allow_spmd_sharding_propagation_to_parameters: false
+      sharding_config {
+      }
+      memory_fitting_level: EFFORT_O2
+      schedule_config {
+      }
+      partition_size: 2
+    }
+  }
+  buffer_assignment {
+    logical_buffers {
+      size: 4
+      defined_at {
+      }
+    }
+    logical_buffers {
+      id: 1
+      size: 4
+      defined_at {
+        instruction_id: 1
+      }
+    }
+    logical_buffers {
+      id: 2
+      size: 4
+      defined_at {
+        instruction_id: 2
+      }
+    }
+    logical_buffers {
+      id: 10
+      size: 32
+      defined_at {
+        instruction_id: 4294967297
+      }
+    }
+    logical_buffers {
+      id: 11
+      size: 32
+      defined_at {
+        instruction_id: 4294967298
+      }
+    }
+    logical_buffers {
+      id: 12
+      size: 4
+      defined_at {
+        instruction_id: 4294967296
+      }
+    }
+    logical_buffers {
+      id: 13
+      size: 32
+      defined_at {
+        instruction_id: 4294967299
+      }
+    }
+    logical_buffers {
+      id: 14
+      size: 16
+      defined_at {
+        instruction_id: 4294967300
+      }
+    }
+    buffer_allocations {
+      size: 32
+      assigned {
+        logical_buffer_id: 10
+        size: 32
+        element_type: U32
+      }
+      is_constant: true
+    }
+    buffer_allocations {
+      index: 1
+      size: 32
+      assigned {
+        logical_buffer_id: 11
+        size: 32
+        element_type: U32
+      }
+      is_constant: true
+    }
+    buffer_allocations {
+      index: 2
+      size: 16
+      maybe_live_out: true
+      assigned {
+        logical_buffer_id: 12
+        size: 4
+        element_type: U32
+      }
+      assigned {
+        logical_buffer_id: 14
+        size: 16
+        element_type: U32
+      }
+    }
+    buffer_allocations {
+      index: 3
+      size: 4
+      is_thread_local: true
+      assigned {
+        logical_buffer_id: 2
+        size: 4
+        element_type: U32
+      }
+    }
+    buffer_allocations {
+      index: 4
+      size: 4
+      is_thread_local: true
+      assigned {
+        size: 4
+        element_type: U32
+      }
+    }
+    buffer_allocations {
+      index: 5
+      size: 4
+      is_thread_local: true
+      assigned {
+        logical_buffer_id: 1
+        size: 4
+        element_type: U32
+      }
+    }
+    buffer_allocations {
+      index: 6
+      size: 32
+      assigned {
+        logical_buffer_id: 13
+        size: 32
+        element_type: U32
+      }
+    }
+    heap_simulator_traces {
+      events {
+        buffer_id: 13
+        computation_name: "main"
+        instruction_name: "loop_select_fusion"
+      }
+      events {
+        kind: FREE
+        buffer_id: 13
+        computation_name: "main"
+        instruction_name: "loop_select_fusion"
+      }
+      whole_module_simulation: true
+      buffer_allocation_index: 6
+    }
+    peak_buffers {
+      logical_buffer_ids: 10
+    }
+    peak_buffers {
+      index: 1
+      logical_buffer_ids: 11
+    }
+    peak_buffers {
+      index: 2
+      logical_buffer_ids: 14
+    }
+    peak_buffers {
+      index: 3
+      logical_buffer_ids: 2
+    }
+    peak_buffers {
+      index: 4
+      logical_buffer_ids: 0
+    }
+    peak_buffers {
+      index: 5
+      logical_buffer_ids: 1
+    }
+    peak_buffers {
+      index: 6
+      logical_buffer_ids: 13
+    }
+  }
+  binary: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000(\014\000\000\000\000\000\000(\007\000\000\000\000\000\000\002d\000\006@\0008\000\004\000@\000\024\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.nv.shared.reserved.0\000.nv.global\000.nv.constant4\000\000\000\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant4\000\000\000\000.nv.callgraph\000.nv.prototype\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.rela.nv.constant.pic\000.nv.merc.nv.constant.pic\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.global\000.nv.constant.pic\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant.pic\000.nv.callgraph\000.nv.prototype\000buffer_for_c1_1\000buffer_for_c0_1\000.nv.constant4\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\214\000\000\000 \240\014\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\003\000\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000=\001\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\001\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\035\001\000\000\021\000\r\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000-\001\000\000\021\000\r\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\002\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 13.2, V13.2.51\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills  -arch sm_100a \000\000\014\000\000\000\010\000\000\000\350\003\000\000NVIDIA Corp\000\002\000d\000\204\000\000\000\003_\001\001\002\t\001\000\002\002\001\000\002\005\005\000\003\007\001\001\002\003\000\000\004\013\010\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\010\000\000\000\000\000\000\000\002\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\003_\001\001\000\000\000\000\010\000\000\000\000\000\000\000\002\000\001\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\001\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\214\000\000\000 \240\022\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\003\000\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\270\000\000\000\003\000\021\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\001\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\035\001\000\000\035 \r\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000-\001\000\000\035 \r\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000n\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\322\001\000\000\000\000\000\000K\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\003\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\217\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000(\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\334\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\005\000\000\000\010\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\004\000\000\000\000\000\000\004\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\307\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \005\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\261\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\005\000\000\000\000\000\0000\000\000\000\000\000\000\000\003\000\000\000\013\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000~\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\005\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000s\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\343\000\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000(\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\370\000\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\000\000\000\004\000\000\000\000\000\000\000\023\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\010\006\000\000\000\000\000\0000\000\000\000\000\000\000\000\023\000\000\000\021\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\'\001\000\000}\000\000p\002\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000p\005\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0008\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000^\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0008\006\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000(\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000(\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000p\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\000\000\000\000\000\000\000\300\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000"
+  gpu_compute_capability {
+    cuda_compute_capability {
+      major: 10
+      feature_extension: ACCELERATED_FEATURES
+    }
+  }
+  module_name: "test"
+  program_shape {
+    result {
+      element_type: U32
+      dimensions: 4
+      layout {
+        minor_to_major: 0
+        tail_padding_alignment_in_elements: 1
+      }
+      is_dynamic_dimension: false
+    }
+  }
+  buffer_allocations {
+    values {
+      size: 32
+      is_constant: true
+    }
+    values {
+      index: 1
+      size: 32
+      is_constant: true
+    }
+    values {
+      index: 2
+      size: 16
+      maybe_live_out: true
+    }
+    values {
+      index: 3
+      size: 4
+      is_thread_local: true
+    }
+    values {
+      index: 4
+      size: 4
+      is_thread_local: true
+    }
+    values {
+      index: 5
+      size: 4
+      is_thread_local: true
+    }
+    values {
+      index: 6
+      size: 32
+    }
+  }
+  output_info_map {
+    shape_index {
+    }
+    output_info {
+      allocation_index: 2
+    }
+  }
+  constants {
+    symbol_name: "buffer_for_c1_1"
+    allocation_index: 1
+  }
+  constants {
+    symbol_name: "buffer_for_c0_1"
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "id.1"
+      thunk_id: 1
+    }
+    replica_id_thunk {
+      dest_buffer {
+        size: 4
+        buffer_allocation_index: 2
+        element_type: U32
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "loop_select_fusion"
+      thunk_id: 2
+    }
+    custom_kernel_thunk {
+      args {
+        slice {
+          size: 32
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          dimensions: 8
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+      }
+      args {
+        slice {
+          size: 32
+          buffer_allocation_index: 1
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          dimensions: 8
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+      }
+      args {
+        slice {
+          size: 4
+          buffer_allocation_index: 2
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          layout {
+            tail_padding_alignment_in_elements: 1
+          }
+        }
+      }
+      args {
+        slice {
+          size: 32
+          buffer_allocation_index: 6
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          dimensions: 8
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+      }
+      written: false
+      written: false
+      written: false
+      written: true
+      custom_kernel {
+        name: "loop_select_fusion"
+        kernel_spec {
+          cubin {
+            data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\330\025\000\000\000\000\000\000X\020\000\000\000\000\000\000\002d\000\006@\0008\000\005\000@\000\026\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_select_fusion\000.nv.capmerc.text.loop_select_fusion\000.nv.merc.debug_frame\000.nv.merc.nv.info\000.nv.merc.nv.info.loop_select_fusion\000.nv.merc.rela.debug_frame\000.nv.merc.nv.shared.reserved.0\000.nv.merc.symtab\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_select_fusion\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\000\000\000!\000\000\000@\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\337\000\000\000 \240\r\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000_\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\004<\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\002\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 13.2, V13.2.51\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills  -arch sm_100a \000\000\014\000\000\000\010\000\000\000\350\003\000\000NVIDIA Corp\000\002\000d\000\204\000\000\000\004/\010\000\010\000\000\000\014\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\002\t\001\000\002\002\001\000\002\005\005\000\003\007\001\001\002\003\000\000\004\013\010\000\t\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000\360\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\003\031 \000\004\n\010\000\t\000\000\000\200\003 \000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\337\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\344\000\000\000\n\000\000\000\"\016\000\254w\004\377\000k\000\000\000\n\000\010\000.\016\000\202{\004\377\000\340\000\000\000\n\000\000\000b\016\000\031y\t\000\000\000\000\000\000!\000\000\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\314\017\000\202{\006\377\000\346\000\000\000\n\000\000\000\342\016\000\201y\002\002\004\000\000\000\000\231\036\014\000$\037\000\014r\000\002\377\000\000\000pP\360\003\000\332\017\001\202\013\004\377\000\342\000\000\000\010\000\000\000\260.\000\202\013\005\377\000\343\000\000\000\010\000\000\000\244\016\000%x\004\t\004\000\000\000\004\000\216\007\000\314O\000\201y\005\004\004\000\000\000\000\231\036\014\000\242\016\000%x\006\t\004\000\000\000\006\000\216\007\000\312\217\000\206y\000\006\005\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\001\000\000\300\021\000\000\000\200\271\000\000\001\013\004\n\370\000\004\000\000\000A\000\000\004\000\000\002\"\016\006\370\000R\000\000\000\203\000@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\001\013\016\n\372\000\005\000\000\000\003\0019\004\000\000\002\"\016\006\370\000R\000\000\000\003\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\013\004\n\370\000\004\000\000\000A\002\001\002\000\000\001b\000\n\370\000\000\000\000\000\000\000\000\000\000\000\002\"\016\006\370\000R\000\000\000\203\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\030\000\000\000\002\"\016\006\000\000B\000\000\000\001\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\002\"\016\006\000\000B\000\000\000A\001@\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\014\000\000\000\0028\0162\370\000@\021\000\000\000\000\202\001\n\000\000\002\001@\001\000\000\000\000\000\000\000\000\000\000\000P\010\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\001|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\3774\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\004\020\000\000\000\004\220\001\000\000\014\201\200\200(\000\004\360\377\377\377\000\000\000\000\000\000\000\000\000\004/\010\000\010\000\000\000\014\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004Z \000\212\235\"\244\261\235\024m\000\264*\363\367X\003\245\',!0\311\036\307\217\017\014Il\n/\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\002J\000\000\004\034\004\000\220\001\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\000\000\000\000D\000\000\000\000\000\000\000=\000\001\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\337\000\000\000 \240\024\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\001\000\000\022\020\017\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\351\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000M\002\000\000\000\000\000\000\200\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\003\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\306\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\300\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000(\005\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\334\005\000\000\000\000\000\000 \000\000\000\000\000\000\000\005\000\000\000\010\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\374\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \006\000\000\000\000\000\000 \000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000v\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\006\000\000\000\000\000\000\220\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\320\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\344\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\360\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\007\000\000\000\000\000\000\000\002\000\000\000\000\000\000\003\000\000\000\010\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\t\000\000\000\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\022\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\t\000\000\000\000\000\000\240\003\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0003\001\000\000\026\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000 \r\000\000\000\000\000\000\022\001\000\000\000\000\000\000\025\000\000\000\010\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000W\001\000\000\001\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\0002\016\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000l\001\000\000\203\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\244\016\000\000\000\000\000\000$\000\000\000\000\000\000\000\025\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000}\001\000\000\203\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\310\016\000\000\000\000\000\000\234\000\000\000\000\000\000\000\025\000\000\000\017\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\241\001\000\000\202\000\000p@\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000h\017\000\000\000\000\000\000\030\000\000\000\000\000\000\000\025\000\000\000\020\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\273\001\000\000\025\000\000p\003\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\200\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\331\001\000\000\205\000\000p\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\200\017\000\000\000\000\000\000\330\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\330\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\330\025\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\200\007\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\003\000\000\000\000\000\000\240\003\000\000\000\000\000\000\010\000\000\000\000\000\000\000"
+          }
+          arity: 4
+          kernel_name: "loop_select_fusion"
+          kernel_args_packing_spec {
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 1
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 2
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 3
+              }
+            }
+          }
+        }
+        block_dims {
+          coordinates {
+            x: 1
+            y: 1
+            z: 1
+          }
+        }
+        thread_dims {
+          coordinates {
+            x: 8
+            y: 1
+            z: 1
+          }
+        }
+      }
+      use_pdl: true
+      tma_metadata {
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "ars.3"
+      thunk_id: 3
+    }
+    reduce_scatter_thunk {
+      collective_config {
+        operand_element_type: U32
+      }
+      buffers {
+        element_count: 8
+        source_buffer {
+          slice {
+            size: 32
+            buffer_allocation_index: 6
+            element_type: U32
+          }
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+        }
+        destination_buffer {
+          slice {
+            size: 16
+            buffer_allocation_index: 2
+            element_type: U32
+          }
+          shape {
+            element_type: U32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+        }
+      }
+      reduction_kind: REDUCTION_KIND_SUM
+    }
+  }
+  executable_abi_version {
+    platform_name: "CUDA"
+    cuda_platform_version {
+      cuda_toolkit_version: "13.2.0"
+      cudnn_version: "9.23.0"
+      cub_version: "3.1.2"
+    }
+  }
+  buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n             parameter allocation:         0B\n              constant allocation:        64B\n        maybe_live_out allocation:        16B\n     preallocated temp allocation:        32B\n  preallocated temp fragmentation:         0B (0.00%)\n                 total allocation:       124B\n              total fragmentation:        76B (61.29%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n"
+  gpu_topology {
+    platform_version: "Compute Capability 10.0a"
+    num_partitions: 1
+    num_hosts_per_partition: 1
+    num_devices_per_host: 2
+    gpu_target_config {
+      gpu_device_info {
+        threads_per_block_limit: 1024
+        threads_per_warp: 32
+        shared_memory_per_block: 49152
+        shared_memory_per_core: 233472
+        threads_per_core_limit: 2048
+        core_count: 152
+        fpus_per_core: 128
+        block_dim_limit_x: 2147483647
+        block_dim_limit_y: 65535
+        block_dim_limit_z: 65535
+        memory_bandwidth: 7928064000000
+        l2_cache_size: 135528448
+        clock_rate_ghz: 2.062
+        device_memory_size: 197897617408
+        shared_memory_per_block_optin: 232448
+        cuda_compute_capability {
+          major: 10
+          feature_extension: ACCELERATED_FEATURES
+        }
+        registers_per_core_limit: 65536
+        registers_per_block_limit: 65536
+        scalar_unit_description {
+          rate_infos {
+            key: 4
+            value {
+              clock_rate_ghz: 2.062
+              units_per_core: 64
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 8
+            value {
+              clock_rate_ghz: 2.062
+              units_per_core: 64
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 10
+            value {
+              clock_rate_ghz: 2.062
+              units_per_core: 128
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 11
+            value {
+              clock_rate_ghz: 2.062
+              units_per_core: 128
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 12
+            value {
+              clock_rate_ghz: 2.062
+              units_per_core: 64
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 16
+            value {
+              clock_rate_ghz: 2.062
+              units_per_core: 128
+              ops_per_clock: 1
+            }
+          }
+        }
+        matrix_unit_description {
+          rate_infos {
+            key: 2
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 6
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 10
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 11
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 512
+            }
+          }
+          rate_infos {
+            key: 12
+            value {
+              clock_rate_ghz: 2.062
+              units_per_core: 4
+              ops_per_clock: 16
+            }
+          }
+          rate_infos {
+            key: 16
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 19
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 20
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 23
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 24
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 25
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 28
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 29
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 32
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 4096
+            }
+          }
+          rate_infos {
+            key: 33
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 35
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+          rate_infos {
+            key: 36
+            value {
+              clock_rate_ghz: 1.9259081
+              units_per_core: 4
+              ops_per_clock: 2048
+            }
+          }
+        }
+        driver_version: "13.2.0"
+        kernel_mode_driver_version: "580.126.20"
+        runtime_version: "13.2.0"
+        compile_time_toolkit_version: "13.2.0"
+        dnn_version: "9.23.0"
+        cub_version: "3.1.2"
+        device_interconnect_info {
+          active_links: 18
+        }
+        device_vendor: "NVIDIA Corporation"
+        platform_version: "Compute Capability 10.0a"
+        pci_bus_id: "<undefined>"
+        name: "NVIDIA GB200"
+        model_str: "sm_10.0a with 197897617408B RAM, 152 cores, 2062000KHz clock, 3996000KHz mem clock, 135528448B L2$"
+        numa_node: -1
+        thread_dim_limit_x: 1024
+        thread_dim_limit_y: 1024
+        thread_dim_limit_z: 64
+        device_address_bits: 64
+        pcie_bandwidth: 2147483648
+        ecc_enabled: true
+        mem_clock_ghz: 3.996
+        reserved_shared_memory_per_block: 1024
+        max_blocks_per_multiprocessor: 32
+        collective_memory_granularity: 2097152
+      }
+      platform_name: "CUDA"
+      device_description_str: "NVIDIA GB200"
+      runtime_version {
+        major: 13
+        minor: 2
+      }
+    }
+    host_target_machine_options {
+      triple: "aarch64-unknown-linux-gnu"
+      cpu: "neoverse-v2"
+    }
+    num_devices_per_process: 2
+  }
+}
diff --git a/third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_h100/v4/AllGatherMixedTypes.pbtxt b/third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_h100/v4/AllGatherMixedTypes.pbtxt
new file mode 100644
index 00000000000..d94951cc120
--- /dev/null
+++ b/third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_h100/v4/AllGatherMixedTypes.pbtxt
@@ -0,0 +1,2459 @@
+# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto
+# proto-message: xla.HumanReadableAotExecutable
+executable_and_options {
+  compile_options {
+    executable_build_options {
+      device_ordinal: -1
+      result_layout {
+        element_type: TUPLE
+        tuple_shapes {
+          element_type: U32
+          dimensions: 4
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+        tuple_shapes {
+          element_type: F32
+          dimensions: 4
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+      }
+      debug_options {
+        xla_disable_hlo_passes: "constant_folding"
+        xla_backend_optimization_level: 3
+        xla_eliminate_hlo_implicit_broadcast: true
+        xla_cpu_multi_thread_eigen: true
+        xla_gpu_cuda_data_dir: "./cuda_sdk_lib"
+        xla_llvm_enable_alias_scope_metadata: true
+        xla_llvm_enable_noalias_metadata: true
+        xla_llvm_enable_invariant_load_metadata: true
+        xla_llvm_disable_expensive_passes: false
+        xla_cpu_use_onednn: false
+        xla_cpu_enable_fast_math: false
+        xla_gpu_enable_fast_min_max: false
+        xla_force_host_platform_device_count: 1
+        xla_gpu_enable_dynamic_slice_fusion: true
+        xla_hlo_evaluator_use_fast_path: true
+        xla_dump_hlo_as_html: false
+        xla_cpu_fast_math_honor_nans: true
+        xla_cpu_fast_math_honor_infs: true
+        xla_allow_excess_precision: true
+        xla_gpu_autotune_level: 4
+        xla_cpu_fast_math_honor_division: true
+        xla_cpu_fast_math_honor_functions: true
+        xla_dump_include_timestamp: false
+        xla_dump_max_hlo_modules: -1
+        xla_cpu_enable_xprof_traceme: false
+        xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false
+        xla_cpu_enable_fast_min_max: true
+        xla_multiheap_size_constraint_per_heap: -1
+        xla_dump_module_metadata: false
+        xla_gpu_force_compilation_parallelism: 1
+        xla_dump_fusion_visualization: false
+        xla_gpu_all_reduce_combine_threshold_bytes: 31457287
+        xla_gpu_nccl_termination_timeout_seconds: -1
+        xla_dump_hlo_as_long_text: true
+        xla_gpu_enable_cublaslt: true
+        xla_gpu_shape_checks: RUNTIME
+        xla_gpu_use_runtime_fusion: false
+        xla_dump_latency_hiding_schedule: false
+        xla_dump_enable_mlir_pretty_form: true
+        xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP
+        xla_gpu_enable_triton_gemm: true
+        xla_gpu_enable_cudnn_int8x32_convolution_reordering: true
+        xla_gpu_triton_gemm_any: true
+        xla_gpu_enable_while_loop_reduce_scatter_code_motion: false
+        xla_gpu_collective_inflation_factor: 1
+        xla_gpu_graph_min_graph_size: 5
+        xla_gpu_enable_reassociation_for_converted_ar: true
+        xla_gpu_pgle_profile_file_or_directory_path: ""
+        xla_gpu_all_gather_combine_threshold_bytes: 31457287
+        xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287
+        xla_gpu_enable_highest_priority_async_stream: true
+        xla_gpu_exhaustive_tiling_search: false
+        xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0
+        xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1
+        xla_gpu_redzone_padding_bytes: 8388608
+        xla_gpu_copy_insertion_use_region_analysis: false
+        xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807
+        xla_gpu_collect_cost_model_stats: false
+        xla_gpu_enable_pipelined_p2p: false
+        xla_gpu_cublas_fallback: true
+        xla_gpu_enable_while_loop_double_buffering: false
+        xla_gpu_filter_kernels_spilling_registers_on_autotuning: true
+        xla_debug_buffer_assignment_show_max: 15
+        xla_detailed_logging: true
+        xla_enable_dumping: true
+        xla_gpu_enable_all_gather_combine_by_dim: false
+        xla_gpu_enable_analytical_latency_estimator: false
+        xla_gpu_llvm_verification_level: 0
+        xla_gpu_enable_reduce_scatter_combine_by_dim: false
+        xla_gpu_enable_command_buffer: COLLECTIVES_KERNEL
+        xla_gpu_enable_command_buffer: CONDITIONAL
+        xla_gpu_enable_command_buffer: CUBLAS
+        xla_gpu_enable_command_buffer: CUBLASLT
+        xla_gpu_enable_command_buffer: CUDNN
+        xla_gpu_enable_command_buffer: CUSTOM_CALL
+        xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION
+        xla_gpu_enable_command_buffer: FUSION
+        xla_gpu_enable_cub_radix_sort: true
+        xla_gpu_memory_limit_slop_factor: 95
+        xla_gpu_target_config_filename: ""
+        xla_gpu_enable_cudnn_layer_norm: false
+        xla_gpu_threshold_for_windowed_einsum_mib: 100000
+        xla_gpu_enable_nccl_user_buffers: false
+        xla_gpu_enable_libnvptxcompiler: true
+        xla_gpu_enable_nccl_comm_splitting: true
+        xla_gpu_nccl_collective_max_nchannels: 0
+        xla_gpu_nccl_p2p_max_nchannels: 0
+        xla_gpu_nccl_init_max_rank_per_root_ratio: 0
+        xla_gpu_multi_streamed_windowed_einsum: true
+        xla_gpu_gemm_rewrite_size_threshold: 100
+        xla_gpu_require_complete_aot_autotune_results: false
+        xla_gpu_cudnn_gemm_fusion_level: 0
+        xla_gpu_use_memcpy_local_p2p: false
+        xla_gpu_autotune_max_solutions: 0
+        xla_dump_large_constants: false
+        xla_gpu_verify_triton_fusion_numerics: false
+        xla_reduce_window_rewrite_base_length: 16
+        xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL
+        xla_gpu_enable_host_memory_offloading: false
+        xla_gpu_exclude_nondeterministic_ops: true
+        xla_gpu_nccl_terminate_on_error: false
+        xla_gpu_shard_autotuning: true
+        xla_gpu_enable_approx_costly_collectives: false
+        xla_cpu_prefer_vector_width: 256
+        xla_gpu_per_fusion_autotune_cache_dir: ""
+        xla_cmd_buffer_trace_cache_size: 16
+        xla_gpu_temp_buffer_use_separate_color: false
+        xla_syntax_sugar_async_ops: false
+        xla_gpu_autotune_gemm_rtol: 0.1
+        xla_enable_command_buffers_during_profiling: false
+        xla_gpu_cudnn_gemm_max_plans: 5
+        xla_cpu_parallel_codegen_split_count: 32
+        xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE
+        xla_gpu_executable_warn_stuck_timeout_seconds: 10
+        xla_gpu_executable_terminate_timeout_seconds: 30
+        xla_gpu_experimental_disable_binary_libraries: false
+        xla_gpu_dot_merger_threshold_mb: 64
+        xla_cpu_max_isa: ""
+        xla_gpu_experimental_enable_fusion_block_level_rewriter: false
+        xla_enable_fast_math: false
+        xla_gpu_experimental_parallel_collective_overlap_limit: 1
+        xla_cpu_copy_insertion_use_region_analysis: false
+        xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1
+        xla_gpu_experimental_enable_triton_heroless_priority_fusion: false
+        xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN
+        xla_gpu_experimental_stream_annotation: true
+        xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO
+        xla_pjrt_allow_auto_layout_in_hlo: false
+        xla_gpu_enable_scatter_determinism_expander: false
+        xla_gpu_require_exclusive_lock: false
+        xla_gpu_generate_debug_info: false
+        xla_gpu_generate_line_info: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false
+        xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE
+        xla_gpu_fail_ptx_compilation_on_register_spilling: false
+        xla_gpu_enable_analytical_sol_latency_estimator: true
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_prep_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_size_bytes"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "gpus_per_node"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nccl_op_launch_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nic_speed_gbps"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "rtt_us"
+          value: "-1"
+        }
+        xla_cpu_use_xnnpack: true
+        xla_gpu_experimental_pack_dot_operands_along_k_dimension: true
+        xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false
+        xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED
+        xla_gpu_experimental_enable_heuristic_collective_combining: false
+        xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false
+        xla_hlo_pass_fix_detect_cycles: false
+        xla_cpu_generate_unique_c_style_kernel_entry_points: false
+        xla_gpu_experimental_collective_cse_distance_threshold: 0
+        xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true
+        xla_gpu_experimental_collective_perf_table_path: ""
+        xla_gpu_collective_permute_combine_threshold_bytes: 31457287
+        xla_unsupported_crash_on_hlo_pass_noop_change: false
+        xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false
+        xla_dump_full_hlo_config: true
+        xla_gpu_unsupported_enable_triton_multi_output_fusion: false
+        xla_gpu_experimental_matmul_perf_table_path: ""
+        xla_gpu_unsupported_enable_all_reduce_decomposer: false
+        xla_gpu_nccl_blocking_communicators: true
+        xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20
+        xla_gpu_first_collective_call_terminate_timeout_seconds: 40
+        xla_gpu_nccl_async_execution: false
+        xla_cpu_emitter_verification_level: 1
+        xla_early_exit_with_layouts: false
+        xla_gpu_experimental_use_ragged_dot_fusion: false
+        xla_detect_unstable_reductions: DETECTION_MODE_NONE
+        xla_gpu_command_buffer_scheduling_mode: LHS
+        xla_gpu_experimental_enable_nccl_symmetric_buffers: false
+        xla_disable_automatic_host_compute_offload: false
+        xla_gpu_experimental_enable_fusion_autotuner: true
+        xla_gpu_experimental_scaled_dot_with_triton: true
+        xla_gpu_command_buffer_unroll_loops: false
+        xla_cpu_experimental_onednn_custom_call: false
+        xla_gpu_experimental_enable_checksum_tracing_on_thunks: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true
+        xla_cpu_collective_call_terminate_timeout_seconds: 40
+        xla_cpu_collective_call_warn_stuck_seconds: 20
+        xla_keep_shardings_after_spmd: false
+        xla_gpu_experimental_enable_triton_warp_specialization: false
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE
+        xla_cpu_enable_platform_dependent_math: true
+        xla_gpu_detect_nan: DETECTION_MODE_NONE
+        xla_gpu_detect_inf: DETECTION_MODE_NONE
+        xla_enable_enzyme_comms_opt: false
+        xla_gpu_experimental_enable_buffer_saver_on_thunks: false
+        xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE
+        xla_enable_scoped_logging_timers: true
+        xla_gpu_executable_embed_debug_info: true
+        xla_cpu_collective_timeout_seconds: 1800
+        xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false
+        xla_gpu_enable_pipelined_host_offloading: false
+        xla_gpu_default_to_alg_dot_bf16_bf16_f32: false
+        xla_gpu_experimental_autotune_backends: CUDNN
+        xla_gpu_experimental_autotune_backends: TRITON
+        xla_gpu_experimental_autotune_backends: CUBLASLT
+        xla_gpu_experimental_autotune_backends: HIPBLASLT
+        xla_gpu_experimental_autotune_backends: MIOPEN
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER
+        xla_gpu_experimental_autotune_backends: NATIVE_EMITTER
+        xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER
+        xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION
+        xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: YNNPACK
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU
+        xla_gpu_experimental_autotune_backends: TENSOR_IR
+        xla_gpu_experimental_all_fusions_with_triton: false
+        xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT
+        xla_gpu_rocm_max_trace_events: 4194304
+        xla_gpu_execution_terminate_timeout: "inf"
+        xla_enable_hlo_sharding_v3: true
+        xla_gpu_log_minmax: false
+        xla_gpu_print_compilation_stats: false
+        xla_gpu_experimental_enable_conv_fusion: false
+        xla_gpu_experimental_enable_tiling_propagation: true
+        xla_gpu_execution_progress_tracking: 0
+        xla_gpu_experimental_max_unroll_factor: 32
+        xla_gpu_enable_pdl: true
+        xla_gpu_collective_combine_threshold_count: 2048
+        xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true
+        xla_recognize_reduction_optimization_level: 0
+        xla_gpu_experimental_parallel_async_compute_limit: 2
+        xla_dump_buffer_assignment_analysis: true
+        xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE
+        xla_gpu_experimental_enable_collective_multi_streaming: false
+        xla_gpu_collective_permute_connected_components: false
+        xla_gpu_experimental_force_split_k: 0
+        xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_experimental_gemm_fusion_v2: true
+        xla_gpu_fusion_autotune_top_k_configs: 1
+        xla_gpu_native_emitter_tune_unroll_factor_for_loops: false
+        xla_enable_rgv3_materialization: true
+        xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK
+        xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false
+        xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_executable_num_compute_streams: 0
+        xla_gpu_executable_num_communication_streams: 0
+        xla_gpu_experimental_collective_start_as_early_as_possible: false
+        xla_cpu_experimental_enable_tiling_propagation: true
+        xla_gpu_enable_dus_accumulator_zero_init_elimination: false
+        xla_gpu_experimental_thunk_buffer_debug_module_outputs: false
+        xla_cpu_ftz: true
+        xla_gpu_enable_allocator_spatial_partitioning: true
+        xla_dump_compact_gte: true
+        xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES
+        xla_dump_hlo_as_riegeli: false
+        xla_gpu_experimental_use_ragged_dot_grouped_gemm: true
+        xla_dump_hlo_to_subfolder: false
+        xla_gpu_enable_pdl_launch: true
+        xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED
+        xla_gpu_enable_nccl_user_buffers_in_default_space: false
+        xla_sdy_export_all_reduce_scatter: false
+        xla_gpu_enable_gxl_ragged_all_to_all: false
+        xla_gpu_gxl_scratch_size_bytes: 67108864
+        xla_gpu_use_new_autotune_cache_format: true
+        xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE
+        xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON
+        xla_spmd_enable_dynamic_slice_collective_broadcast: false
+        xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false
+        xla_gpu_async_copy_min_bytes: -1
+        xla_gpu_blas_max_algorithms: 0
+        xla_gpu_experimental_enable_selective_memcpy_overlap: false
+        xla_gpu_experimental_enable_same_shape_multi_output_fusion: false
+        xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1
+        xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1
+        xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240
+        xla_cpu_use_new_xtile_lowering: true
+        xla_gpu_trace_annotation_level: 0
+        xla_gpu_experimental_cost_model_gemm_tiling_default: false
+        xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false
+        xla_gpu_all_reduce_splitter_ignore_profitability_check: false
+        xla_gpu_enable_cupti_multi_subscriber: true
+        xla_enable_hlo_modules_upload: true
+        xla_gpu_collective_domain_assignment: ""
+        xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1
+        xla_compile_all_supported_configs: false
+        xla_gpu_unsupported_use_cross_host_one_shot_kernel: ALLCOLLECTIVES
+        xla_deduplicate_backend_configs_min_size: 9223372036854775807
+        xla_force_config: ""
+        xla_candidate_configs_file: ""
+        xla_gpu_enable_persistent_symmetric_memory: false
+        xla_gpu_experimental_enable_raft_for_stable_topk: false
+        xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND
+        xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET
+        xla_gpu_experimental_emit_collective_reduce: false
+        xla_gpu_experimental_vmm_disabled: false
+        xla_gpu_experimental_enable_dynamic_slice_table_offsets: false
+        xla_gpu_device_execution_terminate_timeout: "inf"
+        xla_gpu_topology_filename: ""
+        xla_gpu_experimental_enable_dynamic_slice_extended_offsets: false
+        xla_gpu_experimental_deterministic_ops_triton: DETERMINISTIC_OPS_TRITON_MODE_OFF
+        xla_gpu_experimental_enable_tensor_ir: false
+      }
+      num_replicas: 2
+      num_partitions: 1
+      device_assignment {
+        replica_count: 2
+        computation_count: 1
+        computation_devices {
+          replica_device_ids: 0
+          replica_device_ids: 1
+        }
+      }
+      allow_spmd_sharding_propagation_to_output: false
+      comp_envs {
+      }
+      allow_spmd_sharding_propagation_to_parameters: false
+      process_count: 1
+      memory_fitting_level: EFFORT_O2
+    }
+  }
+  pjrt_client_name: "PjRtStreamExecutorClient"
+}
+gpu_executable {
+  hlo_module_with_config {
+    hlo_module {
+      name: "test"
+      entry_computation_name: "test_computation"
+      computations {
+        name: "fused_transpose"
+        instructions {
+          name: "param_0.1"
+          opcode: "parameter"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "param_0.1"
+          }
+          id: 4294967298
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "bitcast.6.1"
+          opcode: "bitcast"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "bitcast.6.1"
+          }
+          id: 4294967297
+          operand_ids: 4294967298
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "transpose.1.1"
+          opcode: "transpose"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "transpose.1.1"
+          }
+          dimensions: 1
+          dimensions: 0
+          id: 4294967296
+          operand_ids: 4294967297
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          parameters {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          result {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          parameter_names: "param_0.1"
+        }
+        id: 1
+        root_id: 4294967296
+        is_fusion_computation: true
+      }
+      computations {
+        name: "fused_transpose.1"
+        instructions {
+          name: "param_0.3"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "param_0.3"
+          }
+          id: 8589934594
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "bitcast.4.1"
+          opcode: "bitcast"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "bitcast.4.1"
+          }
+          id: 8589934593
+          operand_ids: 8589934594
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "transpose.3"
+          opcode: "transpose"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "transpose.3"
+          }
+          dimensions: 1
+          dimensions: 0
+          id: 8589934592
+          operand_ids: 8589934593
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          parameters {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          result {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          parameter_names: "param_0.3"
+        }
+        id: 2
+        root_id: 8589934592
+        is_fusion_computation: true
+      }
+      computations {
+        name: "fused_broadcast"
+        instructions {
+          name: "param_0.5"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "param_0.5"
+          }
+          id: 12884901890
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "p1.1.1"
+          opcode: "convert"
+          shape {
+            element_type: F32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "p1.1.1"
+          }
+          id: 12884901889
+          operand_ids: 12884901890
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "broadcast.13.1"
+          opcode: "broadcast"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "broadcast.13.1"
+          }
+          id: 12884901888
+          operand_ids: 12884901889
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          parameters {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          result {
+            element_type: F32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          parameter_names: "param_0.5"
+        }
+        id: 3
+        root_id: 12884901888
+        is_fusion_computation: true
+      }
+      computations {
+        name: "wrapped_broadcast_computation"
+        instructions {
+          name: "param_0.7"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "param_0.7"
+          }
+          id: 21474836481
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "broadcast.12.1"
+          opcode: "broadcast"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "broadcast.12.1"
+          }
+          id: 21474836480
+          operand_ids: 21474836481
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          parameters {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          result {
+            element_type: U32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          parameter_names: "param_0.7"
+        }
+        id: 5
+        root_id: 21474836480
+        is_fusion_computation: true
+      }
+      computations {
+        name: "test_computation"
+        instructions {
+          name: "id.1"
+          opcode: "replica-id"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "id.1"
+          }
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "wrapped_broadcast"
+          opcode: "fusion"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "wrapped_broadcast"
+          }
+          fusion_kind: "kLoop"
+          id: 9
+          operand_ids: 0
+          called_computation_ids: 5
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "loop_broadcast_fusion"
+          opcode: "fusion"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "loop_broadcast_fusion"
+          }
+          fusion_kind: "kLoop"
+          id: 8
+          operand_ids: 0
+          called_computation_ids: 3
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "allgather.3"
+          opcode: "all-gather"
+          shape {
+            element_type: TUPLE
+            tuple_shapes {
+              element_type: U32
+              dimensions: 2
+              dimensions: 2
+              layout {
+                minor_to_major: 0
+                minor_to_major: 1
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+              is_dynamic_dimension: false
+            }
+            tuple_shapes {
+              element_type: F32
+              dimensions: 2
+              dimensions: 2
+              layout {
+                minor_to_major: 0
+                minor_to_major: 1
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+              is_dynamic_dimension: false
+            }
+          }
+          metadata {
+            scheduling_name: "all-gather-start"
+          }
+          dimensions: 1
+          id: 10
+          operand_ids: 9
+          operand_ids: 8
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+          collective_device_list {
+          }
+          backend_config_payload {
+            id: 0
+          }
+        }
+        instructions {
+          name: "ag0.1"
+          opcode: "get-tuple-element"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "ag0.1"
+          }
+          id: 1
+          operand_ids: 10
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "loop_transpose_fusion.1"
+          opcode: "fusion"
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "loop_transpose_fusion.1"
+          }
+          fusion_kind: "kLoop"
+          id: 7
+          operand_ids: 1
+          called_computation_ids: 2
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "bitcast.24"
+          opcode: "bitcast"
+          shape {
+            element_type: U32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "bitcast.24"
+          }
+          id: 4
+          operand_ids: 7
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "ag1.1"
+          opcode: "get-tuple-element"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "ag1.1"
+          }
+          tuple_index: 1
+          id: 2
+          operand_ids: 10
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "loop_transpose_fusion"
+          opcode: "fusion"
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 1
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "loop_transpose_fusion"
+          }
+          fusion_kind: "kLoop"
+          id: 6
+          operand_ids: 2
+          called_computation_ids: 1
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "bitcast.1.0"
+          opcode: "bitcast"
+          shape {
+            element_type: F32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "bitcast.1.0"
+          }
+          id: 5
+          operand_ids: 6
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "out.1"
+          opcode: "tuple"
+          shape {
+            element_type: TUPLE
+            tuple_shapes {
+              element_type: U32
+              dimensions: 4
+              layout {
+                minor_to_major: 0
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+            }
+            tuple_shapes {
+              element_type: F32
+              dimensions: 4
+              layout {
+                minor_to_major: 0
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+            }
+          }
+          metadata {
+            scheduling_name: "out.1"
+          }
+          id: 3
+          operand_ids: 4
+          operand_ids: 5
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          result {
+            element_type: TUPLE
+            tuple_shapes {
+              element_type: U32
+              dimensions: 4
+              layout {
+                minor_to_major: 0
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+            }
+            tuple_shapes {
+              element_type: F32
+              dimensions: 4
+              layout {
+                minor_to_major: 0
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+            }
+          }
+        }
+        root_id: 3
+      }
+      host_program_shape {
+        result {
+          element_type: TUPLE
+          tuple_shapes {
+            element_type: U32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          tuple_shapes {
+            element_type: F32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+        }
+      }
+      schedule {
+        sequences {
+          key: 0
+          value {
+            instruction_ids: 0
+            instruction_ids: 8
+            instruction_ids: 9
+            instruction_ids: 10
+            instruction_ids: 1
+            instruction_ids: 7
+            instruction_ids: 2
+            instruction_ids: 6
+            instruction_ids: 4
+            instruction_ids: 5
+            instruction_ids: 3
+          }
+        }
+      }
+      input_output_alias {
+      }
+      device_assignment {
+        replica_count: 2
+        computation_count: 1
+        computation_devices {
+          replica_device_ids: 0
+          replica_device_ids: 1
+        }
+      }
+      buffer_donor {
+      }
+      frontend_attributes {
+        map {
+          key: "fingerprint_before_lhs"
+          value: "7b2e4f79c5524b5cbc566da178aa7fda"
+        }
+      }
+      payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}"
+    }
+    config {
+      entry_computation_layout {
+        result {
+          element_type: TUPLE
+          tuple_shapes {
+            element_type: U32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          tuple_shapes {
+            element_type: F32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+        }
+      }
+      replica_count: 2
+      num_partitions: 1
+      intra_op_parallelism_threads: -1
+      debug_options {
+        xla_disable_hlo_passes: "constant_folding"
+        xla_backend_optimization_level: 3
+        xla_eliminate_hlo_implicit_broadcast: true
+        xla_cpu_multi_thread_eigen: true
+        xla_gpu_cuda_data_dir: "./cuda_sdk_lib"
+        xla_llvm_enable_alias_scope_metadata: true
+        xla_llvm_enable_noalias_metadata: true
+        xla_llvm_enable_invariant_load_metadata: true
+        xla_llvm_disable_expensive_passes: false
+        xla_cpu_use_onednn: false
+        xla_cpu_enable_fast_math: false
+        xla_gpu_enable_fast_min_max: false
+        xla_force_host_platform_device_count: 1
+        xla_gpu_enable_dynamic_slice_fusion: true
+        xla_hlo_evaluator_use_fast_path: true
+        xla_dump_hlo_as_html: false
+        xla_cpu_fast_math_honor_nans: true
+        xla_cpu_fast_math_honor_infs: true
+        xla_allow_excess_precision: true
+        xla_gpu_autotune_level: 4
+        xla_cpu_fast_math_honor_division: true
+        xla_cpu_fast_math_honor_functions: true
+        xla_dump_include_timestamp: false
+        xla_dump_max_hlo_modules: -1
+        xla_cpu_enable_xprof_traceme: false
+        xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false
+        xla_cpu_enable_fast_min_max: true
+        xla_multiheap_size_constraint_per_heap: -1
+        xla_dump_module_metadata: false
+        xla_gpu_force_compilation_parallelism: 1
+        xla_dump_fusion_visualization: false
+        xla_gpu_all_reduce_combine_threshold_bytes: 31457287
+        xla_gpu_nccl_termination_timeout_seconds: -1
+        xla_dump_hlo_as_long_text: true
+        xla_gpu_enable_cublaslt: true
+        xla_gpu_shape_checks: RUNTIME
+        xla_gpu_use_runtime_fusion: false
+        xla_dump_latency_hiding_schedule: false
+        xla_dump_enable_mlir_pretty_form: true
+        xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP
+        xla_gpu_enable_triton_gemm: true
+        xla_gpu_enable_cudnn_int8x32_convolution_reordering: true
+        xla_gpu_triton_gemm_any: true
+        xla_gpu_enable_while_loop_reduce_scatter_code_motion: false
+        xla_gpu_collective_inflation_factor: 1
+        xla_gpu_graph_min_graph_size: 5
+        xla_gpu_enable_reassociation_for_converted_ar: true
+        xla_gpu_pgle_profile_file_or_directory_path: ""
+        xla_gpu_all_gather_combine_threshold_bytes: 31457287
+        xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287
+        xla_gpu_enable_highest_priority_async_stream: true
+        xla_gpu_exhaustive_tiling_search: false
+        xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0
+        xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1
+        xla_gpu_redzone_padding_bytes: 8388608
+        xla_gpu_copy_insertion_use_region_analysis: false
+        xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807
+        xla_gpu_collect_cost_model_stats: false
+        xla_gpu_enable_pipelined_p2p: false
+        xla_gpu_cublas_fallback: true
+        xla_gpu_enable_while_loop_double_buffering: false
+        xla_gpu_filter_kernels_spilling_registers_on_autotuning: true
+        xla_debug_buffer_assignment_show_max: 15
+        xla_detailed_logging: true
+        xla_enable_dumping: true
+        xla_gpu_enable_all_gather_combine_by_dim: false
+        xla_gpu_enable_analytical_latency_estimator: false
+        xla_gpu_llvm_verification_level: 0
+        xla_gpu_enable_reduce_scatter_combine_by_dim: false
+        xla_gpu_enable_command_buffer: COLLECTIVES_KERNEL
+        xla_gpu_enable_command_buffer: CONDITIONAL
+        xla_gpu_enable_command_buffer: CUBLAS
+        xla_gpu_enable_command_buffer: CUBLASLT
+        xla_gpu_enable_command_buffer: CUDNN
+        xla_gpu_enable_command_buffer: CUSTOM_CALL
+        xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION
+        xla_gpu_enable_command_buffer: FUSION
+        xla_gpu_enable_cub_radix_sort: true
+        xla_gpu_memory_limit_slop_factor: 95
+        xla_gpu_target_config_filename: ""
+        xla_gpu_enable_cudnn_layer_norm: false
+        xla_gpu_threshold_for_windowed_einsum_mib: 100000
+        xla_gpu_enable_nccl_user_buffers: false
+        xla_gpu_enable_libnvptxcompiler: true
+        xla_gpu_enable_nccl_comm_splitting: true
+        xla_gpu_nccl_collective_max_nchannels: 0
+        xla_gpu_nccl_p2p_max_nchannels: 0
+        xla_gpu_nccl_init_max_rank_per_root_ratio: 0
+        xla_gpu_multi_streamed_windowed_einsum: true
+        xla_gpu_gemm_rewrite_size_threshold: 100
+        xla_gpu_require_complete_aot_autotune_results: false
+        xla_gpu_cudnn_gemm_fusion_level: 0
+        xla_gpu_use_memcpy_local_p2p: false
+        xla_gpu_autotune_max_solutions: 0
+        xla_dump_large_constants: false
+        xla_gpu_verify_triton_fusion_numerics: false
+        xla_reduce_window_rewrite_base_length: 16
+        xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL
+        xla_gpu_enable_host_memory_offloading: false
+        xla_gpu_exclude_nondeterministic_ops: true
+        xla_gpu_nccl_terminate_on_error: false
+        xla_gpu_shard_autotuning: true
+        xla_gpu_enable_approx_costly_collectives: false
+        xla_cpu_prefer_vector_width: 256
+        xla_gpu_per_fusion_autotune_cache_dir: ""
+        xla_cmd_buffer_trace_cache_size: 16
+        xla_gpu_temp_buffer_use_separate_color: false
+        xla_syntax_sugar_async_ops: false
+        xla_gpu_autotune_gemm_rtol: 0.1
+        xla_enable_command_buffers_during_profiling: false
+        xla_gpu_cudnn_gemm_max_plans: 5
+        xla_cpu_parallel_codegen_split_count: 32
+        xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE
+        xla_gpu_executable_warn_stuck_timeout_seconds: 10
+        xla_gpu_executable_terminate_timeout_seconds: 30
+        xla_gpu_experimental_disable_binary_libraries: false
+        xla_gpu_dot_merger_threshold_mb: 64
+        xla_cpu_max_isa: ""
+        xla_gpu_experimental_enable_fusion_block_level_rewriter: false
+        xla_enable_fast_math: false
+        xla_gpu_experimental_parallel_collective_overlap_limit: 1
+        xla_cpu_copy_insertion_use_region_analysis: false
+        xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1
+        xla_gpu_experimental_enable_triton_heroless_priority_fusion: false
+        xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN
+        xla_gpu_experimental_stream_annotation: true
+        xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO
+        xla_pjrt_allow_auto_layout_in_hlo: false
+        xla_gpu_enable_scatter_determinism_expander: false
+        xla_gpu_require_exclusive_lock: false
+        xla_gpu_generate_debug_info: false
+        xla_gpu_generate_line_info: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false
+        xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE
+        xla_gpu_fail_ptx_compilation_on_register_spilling: false
+        xla_gpu_enable_analytical_sol_latency_estimator: true
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_prep_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_size_bytes"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "gpus_per_node"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nccl_op_launch_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nic_speed_gbps"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "rtt_us"
+          value: "-1"
+        }
+        xla_cpu_use_xnnpack: true
+        xla_gpu_experimental_pack_dot_operands_along_k_dimension: true
+        xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false
+        xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED
+        xla_gpu_experimental_enable_heuristic_collective_combining: false
+        xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false
+        xla_hlo_pass_fix_detect_cycles: false
+        xla_cpu_generate_unique_c_style_kernel_entry_points: false
+        xla_gpu_experimental_collective_cse_distance_threshold: 0
+        xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true
+        xla_gpu_experimental_collective_perf_table_path: ""
+        xla_gpu_collective_permute_combine_threshold_bytes: 31457287
+        xla_unsupported_crash_on_hlo_pass_noop_change: false
+        xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false
+        xla_dump_full_hlo_config: true
+        xla_gpu_unsupported_enable_triton_multi_output_fusion: false
+        xla_gpu_experimental_matmul_perf_table_path: ""
+        xla_gpu_unsupported_enable_all_reduce_decomposer: false
+        xla_gpu_nccl_blocking_communicators: true
+        xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20
+        xla_gpu_first_collective_call_terminate_timeout_seconds: 40
+        xla_gpu_nccl_async_execution: false
+        xla_cpu_emitter_verification_level: 1
+        xla_early_exit_with_layouts: false
+        xla_gpu_experimental_use_ragged_dot_fusion: false
+        xla_detect_unstable_reductions: DETECTION_MODE_NONE
+        xla_gpu_command_buffer_scheduling_mode: LHS
+        xla_gpu_experimental_enable_nccl_symmetric_buffers: false
+        xla_disable_automatic_host_compute_offload: false
+        xla_gpu_experimental_enable_fusion_autotuner: true
+        xla_gpu_experimental_scaled_dot_with_triton: true
+        xla_gpu_command_buffer_unroll_loops: false
+        xla_cpu_experimental_onednn_custom_call: false
+        xla_gpu_experimental_enable_checksum_tracing_on_thunks: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true
+        xla_cpu_collective_call_terminate_timeout_seconds: 40
+        xla_cpu_collective_call_warn_stuck_seconds: 20
+        xla_keep_shardings_after_spmd: false
+        xla_gpu_experimental_enable_triton_warp_specialization: false
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE
+        xla_cpu_enable_platform_dependent_math: true
+        xla_gpu_detect_nan: DETECTION_MODE_NONE
+        xla_gpu_detect_inf: DETECTION_MODE_NONE
+        xla_enable_enzyme_comms_opt: false
+        xla_gpu_experimental_enable_buffer_saver_on_thunks: false
+        xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE
+        xla_enable_scoped_logging_timers: true
+        xla_gpu_executable_embed_debug_info: true
+        xla_cpu_collective_timeout_seconds: 1800
+        xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false
+        xla_gpu_enable_pipelined_host_offloading: false
+        xla_gpu_default_to_alg_dot_bf16_bf16_f32: false
+        xla_gpu_experimental_autotune_backends: CUDNN
+        xla_gpu_experimental_autotune_backends: TRITON
+        xla_gpu_experimental_autotune_backends: CUBLASLT
+        xla_gpu_experimental_autotune_backends: HIPBLASLT
+        xla_gpu_experimental_autotune_backends: MIOPEN
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER
+        xla_gpu_experimental_autotune_backends: NATIVE_EMITTER
+        xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER
+        xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION
+        xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: YNNPACK
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU
+        xla_gpu_experimental_autotune_backends: TENSOR_IR
+        xla_gpu_experimental_all_fusions_with_triton: false
+        xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT
+        xla_gpu_rocm_max_trace_events: 4194304
+        xla_gpu_execution_terminate_timeout: "inf"
+        xla_enable_hlo_sharding_v3: true
+        xla_gpu_log_minmax: false
+        xla_gpu_print_compilation_stats: false
+        xla_gpu_experimental_enable_conv_fusion: false
+        xla_gpu_experimental_enable_tiling_propagation: true
+        xla_gpu_execution_progress_tracking: 0
+        xla_gpu_experimental_max_unroll_factor: 32
+        xla_gpu_enable_pdl: true
+        xla_gpu_collective_combine_threshold_count: 2048
+        xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true
+        xla_recognize_reduction_optimization_level: 0
+        xla_gpu_experimental_parallel_async_compute_limit: 2
+        xla_dump_buffer_assignment_analysis: true
+        xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE
+        xla_gpu_experimental_enable_collective_multi_streaming: false
+        xla_gpu_collective_permute_connected_components: false
+        xla_gpu_experimental_force_split_k: 0
+        xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_experimental_gemm_fusion_v2: true
+        xla_gpu_fusion_autotune_top_k_configs: 1
+        xla_gpu_native_emitter_tune_unroll_factor_for_loops: false
+        xla_enable_rgv3_materialization: true
+        xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK
+        xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false
+        xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_executable_num_compute_streams: 0
+        xla_gpu_executable_num_communication_streams: 0
+        xla_gpu_experimental_collective_start_as_early_as_possible: false
+        xla_cpu_experimental_enable_tiling_propagation: true
+        xla_gpu_enable_dus_accumulator_zero_init_elimination: false
+        xla_gpu_experimental_thunk_buffer_debug_module_outputs: false
+        xla_cpu_ftz: true
+        xla_gpu_enable_allocator_spatial_partitioning: true
+        xla_dump_compact_gte: true
+        xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES
+        xla_dump_hlo_as_riegeli: false
+        xla_gpu_experimental_use_ragged_dot_grouped_gemm: true
+        xla_dump_hlo_to_subfolder: false
+        xla_gpu_enable_pdl_launch: true
+        xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED
+        xla_gpu_enable_nccl_user_buffers_in_default_space: false
+        xla_sdy_export_all_reduce_scatter: false
+        xla_gpu_enable_gxl_ragged_all_to_all: false
+        xla_gpu_gxl_scratch_size_bytes: 67108864
+        xla_gpu_use_new_autotune_cache_format: true
+        xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE
+        xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON
+        xla_spmd_enable_dynamic_slice_collective_broadcast: false
+        xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false
+        xla_gpu_async_copy_min_bytes: -1
+        xla_gpu_blas_max_algorithms: 0
+        xla_gpu_experimental_enable_selective_memcpy_overlap: false
+        xla_gpu_experimental_enable_same_shape_multi_output_fusion: false
+        xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1
+        xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1
+        xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240
+        xla_cpu_use_new_xtile_lowering: true
+        xla_gpu_trace_annotation_level: 0
+        xla_gpu_experimental_cost_model_gemm_tiling_default: false
+        xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false
+        xla_gpu_all_reduce_splitter_ignore_profitability_check: false
+        xla_gpu_enable_cupti_multi_subscriber: true
+        xla_enable_hlo_modules_upload: true
+        xla_gpu_collective_domain_assignment: ""
+        xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1
+        xla_compile_all_supported_configs: false
+        xla_gpu_unsupported_use_cross_host_one_shot_kernel: ALLCOLLECTIVES
+        xla_deduplicate_backend_configs_min_size: 9223372036854775807
+        xla_force_config: ""
+        xla_candidate_configs_file: ""
+        xla_gpu_enable_persistent_symmetric_memory: false
+        xla_gpu_experimental_enable_raft_for_stable_topk: false
+        xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND
+        xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET
+        xla_gpu_experimental_emit_collective_reduce: false
+        xla_gpu_experimental_vmm_disabled: false
+        xla_gpu_experimental_enable_dynamic_slice_table_offsets: false
+        xla_gpu_device_execution_terminate_timeout: "inf"
+        xla_gpu_topology_filename: ""
+        xla_gpu_experimental_enable_dynamic_slice_extended_offsets: false
+        xla_gpu_experimental_deterministic_ops_triton: DETERMINISTIC_OPS_TRITON_MODE_OFF
+        xla_gpu_experimental_enable_tensor_ir: false
+      }
+      static_device_assignment {
+        replica_count: 2
+        computation_count: 1
+        computation_devices {
+          replica_device_ids: 0
+          replica_device_ids: 1
+        }
+      }
+      allow_spmd_sharding_propagation_to_output: false
+      allow_spmd_sharding_propagation_to_parameters: false
+      sharding_config {
+      }
+      memory_fitting_level: EFFORT_O2
+      schedule_config {
+      }
+      partition_size: 2
+    }
+  }
+  buffer_assignment {
+    logical_buffers {
+      id: 9
+      size: 4
+      defined_at {
+      }
+    }
+    logical_buffers {
+      id: 10
+      size: 8
+      defined_at {
+        instruction_id: 9
+      }
+    }
+    logical_buffers {
+      id: 11
+      size: 8
+      defined_at {
+        instruction_id: 8
+      }
+    }
+    logical_buffers {
+      id: 12
+      size: 16
+      defined_at {
+        instruction_id: 10
+      }
+    }
+    logical_buffers {
+      id: 13
+      size: 16
+      defined_at {
+        shape_index: 0
+        instruction_id: 10
+      }
+    }
+    logical_buffers {
+      id: 14
+      size: 16
+      defined_at {
+        shape_index: 1
+        instruction_id: 10
+      }
+    }
+    logical_buffers {
+      id: 15
+      size: 16
+      defined_at {
+        instruction_id: 7
+      }
+    }
+    logical_buffers {
+      id: 16
+      size: 16
+      defined_at {
+        instruction_id: 6
+      }
+    }
+    logical_buffers {
+      id: 17
+      size: 16
+      defined_at {
+        instruction_id: 3
+      }
+    }
+    buffer_allocations {
+      size: 16
+      maybe_live_out: true
+      assigned {
+        logical_buffer_id: 10
+        size: 8
+        element_type: U32
+      }
+      assigned {
+        logical_buffer_id: 15
+        size: 16
+        element_type: U32
+      }
+    }
+    buffer_allocations {
+      index: 1
+      size: 16
+      maybe_live_out: true
+      assigned {
+        logical_buffer_id: 9
+        size: 4
+        element_type: U32
+      }
+      assigned {
+        logical_buffer_id: 13
+        size: 16
+        element_type: U32
+      }
+      assigned {
+        logical_buffer_id: 16
+        size: 16
+        element_type: F32
+      }
+    }
+    buffer_allocations {
+      index: 2
+      size: 16
+      maybe_live_out: true
+      assigned {
+        logical_buffer_id: 17
+        size: 16
+        element_type: TUPLE
+      }
+      is_tuple: true
+    }
+    buffer_allocations {
+      index: 3
+      size: 520
+      assigned {
+        logical_buffer_id: 11
+        offset: 512
+        size: 8
+        element_type: F32
+      }
+      assigned {
+        logical_buffer_id: 12
+        size: 16
+        element_type: TUPLE
+      }
+      assigned {
+        logical_buffer_id: 14
+        offset: 256
+        size: 16
+        element_type: F32
+      }
+      is_tuple: true
+    }
+    heap_simulator_traces {
+      events {
+        buffer_id: 11
+        computation_name: "test_computation"
+        instruction_name: "loop_broadcast_fusion"
+      }
+      events {
+        buffer_id: 14
+        computation_name: "test_computation"
+        instruction_name: "allgather.3"
+      }
+      events {
+        kind: FREE
+        buffer_id: 11
+        computation_name: "test_computation"
+        instruction_name: "loop_broadcast_fusion"
+      }
+      events {
+        kind: FREE
+        buffer_id: 14
+        computation_name: "test_computation"
+        instruction_name: "allgather.3"
+      }
+      whole_module_simulation: true
+      buffer_allocation_index: 3
+    }
+    peak_buffers {
+      logical_buffer_ids: 15
+    }
+    peak_buffers {
+      index: 1
+      logical_buffer_ids: 16
+    }
+    peak_buffers {
+      index: 2
+      logical_buffer_ids: 17
+    }
+    peak_buffers {
+      index: 3
+      logical_buffer_ids: 12
+      logical_buffer_ids: 11
+      logical_buffer_ids: 14
+    }
+  }
+  gpu_compute_capability {
+    cuda_compute_capability {
+      major: 9
+      feature_extension: ACCELERATED_FEATURES
+    }
+  }
+  module_name: "test"
+  program_shape {
+    result {
+      element_type: TUPLE
+      tuple_shapes {
+        element_type: U32
+        dimensions: 4
+        layout {
+          minor_to_major: 0
+          tail_padding_alignment_in_elements: 1
+        }
+        is_dynamic_dimension: false
+      }
+      tuple_shapes {
+        element_type: F32
+        dimensions: 4
+        layout {
+          minor_to_major: 0
+          tail_padding_alignment_in_elements: 1
+        }
+        is_dynamic_dimension: false
+      }
+    }
+  }
+  buffer_allocations {
+    values {
+      size: 16
+      maybe_live_out: true
+    }
+    values {
+      index: 1
+      size: 16
+      maybe_live_out: true
+    }
+    values {
+      index: 2
+      size: 16
+      maybe_live_out: true
+      is_tuple: true
+    }
+    values {
+      index: 3
+      size: 520
+      is_tuple: true
+    }
+  }
+  output_info_map {
+    shape_index {
+    }
+    output_info {
+      allocation_index: 2
+    }
+  }
+  output_info_map {
+    shape_index {
+      indexes: 0
+    }
+    output_info {
+    }
+  }
+  output_info_map {
+    shape_index {
+      indexes: 1
+    }
+    output_info {
+      allocation_index: 1
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "id.1"
+      thunk_id: 1
+    }
+    replica_id_thunk {
+      dest_buffer {
+        size: 4
+        buffer_allocation_index: 1
+        element_type: U32
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "loop_broadcast_fusion"
+      thunk_id: 2
+    }
+    custom_kernel_thunk {
+      args {
+        slice {
+          size: 4
+          buffer_allocation_index: 1
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          layout {
+            tail_padding_alignment_in_elements: 1
+          }
+        }
+      }
+      args {
+        slice {
+          offset: 512
+          size: 8
+          buffer_allocation_index: 3
+          element_type: F32
+        }
+        shape {
+          element_type: F32
+          dimensions: 2
+          dimensions: 1
+          layout {
+            minor_to_major: 0
+            minor_to_major: 1
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+          is_dynamic_dimension: false
+        }
+      }
+      written: false
+      written: true
+      custom_kernel {
+        name: "loop_broadcast_fusion"
+        kernel_spec {
+          cubin {
+            data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000 \n\000\000\000\000\000\000\004Z\000\006@\0008\000\005\000@\000\017\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_broadcast_fusion\000.nv.info.loop_broadcast_fusion\000.nv.shared.loop_broadcast_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_broadcast_fusion\000.nv.constant0.loop_broadcast_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\271\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\350\000\000\000 \240\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000U\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000k\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004(\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\002\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 13.2, V13.2.51\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills  -arch sm_90a \000\000\000\014\000\000\000\010\000\000\000\350\003\000\000NVIDIA Corp\000\002\000Z\000\204\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\002\t\001\000\002\002\001\000\003\007\001\001\002\003\000\000\004\013\010\000\000\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\004\034\004\000\240\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\t\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\204\000\000\000\n\000\000\000\"\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\206\000\000\000\n\000\000\000b\016\000\201y\002\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\342/\000Er\007\000\002\000\000\000\000\020 \000\000\312O\000\206y\000\004\007\000\000\000\004\031\020\014\000\342\017\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000?\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\264\001\000\000\000\000\000\000\217\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\003\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\317\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\240\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000T\005\000\000\000\000\000\000 \000\000\000\000\000\000\000\005\000\000\000\010\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000t\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\230\005\000\000\000\000\000\000\034\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000y\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\264\005\000\000\000\000\000\000l\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\355\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\010\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\271\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\033\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000"
+          }
+          arity: 2
+          kernel_name: "loop_broadcast_fusion"
+          kernel_args_packing_spec {
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 1
+              }
+            }
+          }
+        }
+        block_dims {
+          coordinates {
+            x: 1
+            y: 1
+            z: 1
+          }
+        }
+        thread_dims {
+          coordinates {
+            x: 2
+            y: 1
+            z: 1
+          }
+        }
+      }
+      use_pdl: true
+      tma_metadata {
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "wrapped_broadcast"
+      thunk_id: 3
+    }
+    custom_kernel_thunk {
+      args {
+        slice {
+          size: 4
+          buffer_allocation_index: 1
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          layout {
+            tail_padding_alignment_in_elements: 1
+          }
+        }
+      }
+      args {
+        slice {
+          size: 8
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          dimensions: 2
+          dimensions: 1
+          layout {
+            minor_to_major: 0
+            minor_to_major: 1
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+          is_dynamic_dimension: false
+        }
+      }
+      written: false
+      written: true
+      custom_kernel {
+        name: "wrapped_broadcast"
+        kernel_spec {
+          cubin {
+            data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000 \n\000\000\000\000\000\000\004Z\000\006@\0008\000\005\000@\000\017\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.wrapped_broadcast\000.nv.info.wrapped_broadcast\000.nv.shared.wrapped_broadcast\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000wrapped_broadcast\000.nv.constant0.wrapped_broadcast\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\334\000\000\000 \240\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\375\000\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000-\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000[\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\004$\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\002\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 13.2, V13.2.51\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills  -arch sm_90a \000\000\000\014\000\000\000\010\000\000\000\350\003\000\000NVIDIA Corp\000\002\000Z\000\204\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\002\t\001\000\002\002\001\000\003\007\001\001\002\003\000\000\004\013\010\000\000\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\004\034\004\000\220\000\000\000\004\020\014\000\002\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\t\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\204\000\000\000\n\000\000\000\"\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000b\016\000.x\000\000\000\000\000\000\000\000\000\000\000\312\017\000\202{\004\377\000\206\000\000\000\n\000\000\000b\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\036\000%x\004\007\004\000\000\000\004\000\216\007\000\312/\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000/\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\240\001\000\000\000\000\000\000{\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \003\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\303\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000x\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000,\005\000\000\000\000\000\000 \000\000\000\000\000\000\000\005\000\000\000\010\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\005\000\000\000\000\000\000\034\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000u\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\214\005\000\000\000\000\000\000l\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\363\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\370\005\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\341\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\010\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\017\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000"
+          }
+          arity: 2
+          kernel_name: "wrapped_broadcast"
+          kernel_args_packing_spec {
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 1
+              }
+            }
+          }
+        }
+        block_dims {
+          coordinates {
+            x: 1
+            y: 1
+            z: 1
+          }
+        }
+        thread_dims {
+          coordinates {
+            x: 2
+            y: 1
+            z: 1
+          }
+        }
+      }
+      use_pdl: true
+      tma_metadata {
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "allgather.3"
+      thunk_id: 4
+    }
+    all_gather_thunk {
+      collective_config {
+        operand_element_type: U32
+        operand_element_type: F32
+      }
+      buffers {
+        element_count: 2
+        source_buffer {
+          slice {
+            size: 8
+            element_type: U32
+          }
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+        }
+        destination_buffer {
+          slice {
+            size: 16
+            buffer_allocation_index: 1
+            element_type: U32
+          }
+          shape {
+            element_type: U32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+        }
+      }
+      buffers {
+        element_count: 2
+        source_buffer {
+          slice {
+            offset: 512
+            size: 8
+            buffer_allocation_index: 3
+            element_type: F32
+          }
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 1
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+        }
+        destination_buffer {
+          slice {
+            offset: 256
+            size: 16
+            buffer_allocation_index: 3
+            element_type: F32
+          }
+          shape {
+            element_type: F32
+            dimensions: 2
+            dimensions: 2
+            layout {
+              minor_to_major: 0
+              minor_to_major: 1
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+            is_dynamic_dimension: false
+          }
+        }
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "loop_transpose_fusion.1"
+      thunk_id: 5
+    }
+    custom_kernel_thunk {
+      args {
+        slice {
+          size: 16
+          buffer_allocation_index: 1
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          dimensions: 2
+          dimensions: 2
+          layout {
+            minor_to_major: 0
+            minor_to_major: 1
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+          is_dynamic_dimension: false
+        }
+      }
+      args {
+        slice {
+          size: 16
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          dimensions: 2
+          dimensions: 2
+          layout {
+            minor_to_major: 1
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+          is_dynamic_dimension: false
+        }
+      }
+      written: false
+      written: true
+      custom_kernel {
+        name: "loop_transpose_fusion_1"
+        kernel_spec {
+          cubin {
+            data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000 \n\000\000\000\000\000\000\004Z\000\006@\0008\000\005\000@\000\017\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_transpose_fusion_1\000.nv.info.loop_transpose_fusion_1\000.nv.shared.loop_transpose_fusion_1\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion_1\000.nv.constant0.loop_transpose_fusion_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\277\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\356\000\000\000 \240\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\017\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000?\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000[\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000s\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\002\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 13.2, V13.2.51\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills  -arch sm_90a \000\000\000\014\000\000\000\010\000\000\000\350\003\000\000NVIDIA Corp\000\002\000Z\000\204\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\002\t\001\000\002\002\001\000\003\007\001\001\002\003\000\000\004\013\010\000\000\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\t\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\204\000\000\000\n\000\000\000b\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\206\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\016\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000G\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\276\001\000\000\000\000\000\000\231\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000X\003\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\325\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\260\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000d\005\000\000\000\000\000\000 \000\000\000\000\000\000\000\005\000\000\000\010\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\204\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\005\000\000\000\000\000\000\034\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000{\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\304\005\000\000\000\000\000\000l\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\001\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\363\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000P\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\010\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\277\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000!\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000"
+          }
+          arity: 2
+          kernel_name: "loop_transpose_fusion_1"
+          kernel_args_packing_spec {
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 1
+              }
+            }
+          }
+        }
+        block_dims {
+          coordinates {
+            x: 1
+            y: 1
+            z: 1
+          }
+        }
+        thread_dims {
+          coordinates {
+            x: 4
+            y: 1
+            z: 1
+          }
+        }
+      }
+      use_pdl: true
+      tma_metadata {
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "loop_transpose_fusion"
+      thunk_id: 6
+    }
+    custom_kernel_thunk {
+      args {
+        slice {
+          offset: 256
+          size: 16
+          buffer_allocation_index: 3
+          element_type: F32
+        }
+        shape {
+          element_type: F32
+          dimensions: 2
+          dimensions: 2
+          layout {
+            minor_to_major: 0
+            minor_to_major: 1
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+          is_dynamic_dimension: false
+        }
+      }
+      args {
+        slice {
+          size: 16
+          buffer_allocation_index: 1
+          element_type: F32
+        }
+        shape {
+          element_type: F32
+          dimensions: 2
+          dimensions: 2
+          layout {
+            minor_to_major: 1
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+          is_dynamic_dimension: false
+        }
+      }
+      written: false
+      written: true
+      custom_kernel {
+        name: "loop_transpose_fusion"
+        kernel_spec {
+          cubin {
+            data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\340\r\000\000\000\000\000\000 \n\000\000\000\000\000\000\004Z\000\006@\0008\000\005\000@\000\017\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_transpose_fusion\000.nv.info.loop_transpose_fusion\000.nv.shared.loop_transpose_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_transpose_fusion\000.nv.constant0.loop_transpose_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\271\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\350\000\000\000 \240\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\t\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000U\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000k\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\0044\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\002\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 13.2, V13.2.51\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills  -arch sm_90a \000\000\000\014\000\000\000\010\000\000\000\350\003\000\000NVIDIA Corp\000\002\000Z\000\204\000\000\000\004/\010\000\010\000\000\000\n\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\002\t\001\000\002\002\001\000\003\007\001\001\002\003\000\000\004\013\010\000\000\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\004\034\004\000\320\000\000\000\004\020\014\000\004\000\000\000\001\000\000\000\001\000\000\000\003\031\020\000\004\n\010\000\t\000\000\000\020\002\020\000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\342\017\000\031y\007\000\000\000\000\000\000!\000\000\000.\016\000\202{\002\377\000\204\000\000\000\n\000\000\000b\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\342\017\000.x\000\000\000\000\000\000\000\000\000\000\000\342\017\000\031x\000\007\001\000\000\000\377\006\000\000\000\310\037\000\022x\000\000\002\000\000\000\377\300\216\007\000\310\017\000\021r\005\007\000\000\000\000\377\370\217\007\000\312\017\000%x\002\005\004\000\000\000\002\000\216\007\000\344/\000\202{\004\377\000\206\000\000\000\n\000\000\000(\016\000\201y\003\002\004\000\000\000\000\231\036\014\000\242\016\000%x\004\007\004\000\000\000\004\000\216\007\000\312\037\000\206y\000\004\003\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000?\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\264\001\000\000\000\000\000\000\217\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\003\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\317\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\240\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000T\005\000\000\000\000\000\000 \000\000\000\000\000\000\000\005\000\000\000\010\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000t\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\230\005\000\000\000\000\000\000\034\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000y\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\264\005\000\000\000\000\000\000l\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\355\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\200\001\000\000\000\000\000\000\003\000\000\000\010\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\271\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\033\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000 \002\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\340\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \002\000\000\000\000\000\000 \002\000\000\000\000\000\000\010\000\000\000\000\000\000\000"
+          }
+          arity: 2
+          kernel_name: "loop_transpose_fusion"
+          kernel_args_packing_spec {
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 1
+              }
+            }
+          }
+        }
+        block_dims {
+          coordinates {
+            x: 1
+            y: 1
+            z: 1
+          }
+        }
+        thread_dims {
+          coordinates {
+            x: 4
+            y: 1
+            z: 1
+          }
+        }
+      }
+      use_pdl: true
+      tma_metadata {
+      }
+    }
+  }
+  executable_abi_version {
+    platform_name: "CUDA"
+    cuda_platform_version {
+      cuda_toolkit_version: "13.2.0"
+      cudnn_version: "9.23.0"
+      cub_version: "3.1.2"
+    }
+  }
+  buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n             parameter allocation:         0B\n              constant allocation:         0B\n        maybe_live_out allocation:        48B\n     preallocated temp allocation:       520B\n  preallocated temp fragmentation:       240B (46.15%)\n                 total allocation:       568B\n              total fragmentation:       504B (88.73%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 16B\n\t\tXLA Label: all-gather\n\t\tShape: (u32[2,2], f32[2,2])\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 16B\n\t\tXLA Label: tuple\n\t\tShape: (u32[4], f32[4])\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,2]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 16B\n\t\tXLA Label: fusion\n\t\tShape: u32[2,2]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 8B\n\t\tXLA Label: fusion\n\t\tShape: f32[2,1]\n\t\t==========================\n"
+  gpu_topology {
+    platform_version: "Compute Capability 9.0a"
+    num_partitions: 1
+    num_hosts_per_partition: 1
+    num_devices_per_host: 2
+    gpu_target_config {
+      gpu_device_info {
+        threads_per_block_limit: 1024
+        threads_per_warp: 32
+        shared_memory_per_block: 49152
+        shared_memory_per_core: 233472
+        threads_per_core_limit: 2048
+        core_count: 132
+        fpus_per_core: 128
+        block_dim_limit_x: 2147483647
+        block_dim_limit_y: 65535
+        block_dim_limit_z: 65535
+        memory_bandwidth: 3352320000000
+        l2_cache_size: 52428800
+        clock_rate_ghz: 1.98
+        device_memory_size: 84942979072
+        shared_memory_per_block_optin: 232448
+        cuda_compute_capability {
+          major: 9
+          feature_extension: ACCELERATED_FEATURES
+        }
+        registers_per_core_limit: 65536
+        registers_per_block_limit: 65536
+        scalar_unit_description {
+          rate_infos {
+            key: 4
+            value {
+              clock_rate_ghz: 1.98
+              units_per_core: 64
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 8
+            value {
+              clock_rate_ghz: 1.98
+              units_per_core: 64
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 10
+            value {
+              clock_rate_ghz: 1.98
+              units_per_core: 128
+              ops_per_clock: 2
+            }
+          }
+          rate_infos {
+            key: 11
+            value {
+              clock_rate_ghz: 1.98
+              units_per_core: 128
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 12
+            value {
+              clock_rate_ghz: 1.98
+              units_per_core: 64
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 16
+            value {
+              clock_rate_ghz: 1.98
+              units_per_core: 128
+              ops_per_clock: 2
+            }
+          }
+        }
+        matrix_unit_description {
+          rate_infos {
+            key: 2
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 6
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 10
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 512
+            }
+          }
+          rate_infos {
+            key: 11
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 256
+            }
+          }
+          rate_infos {
+            key: 12
+            value {
+              clock_rate_ghz: 1.98
+              units_per_core: 4
+              ops_per_clock: 32
+            }
+          }
+          rate_infos {
+            key: 16
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 512
+            }
+          }
+          rate_infos {
+            key: 19
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 20
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 23
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 24
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 25
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 28
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 29
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 33
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 35
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 36
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+        }
+        driver_version: "13.2.0"
+        kernel_mode_driver_version: "535.183.6"
+        runtime_version: "13.2.0"
+        compile_time_toolkit_version: "13.2.0"
+        dnn_version: "9.23.0"
+        cub_version: "3.1.2"
+        device_interconnect_info {
+          active_links: 18
+        }
+        device_vendor: "NVIDIA Corporation"
+        platform_version: "Compute Capability 9.0a"
+        pci_bus_id: "<undefined>"
+        name: "NVIDIA H100 80GB HBM3"
+        model_str: "sm_9.0a with 84942979072B RAM, 132 cores, 1980000KHz clock, 2619000KHz mem clock, 52428800B L2$"
+        numa_node: -1
+        thread_dim_limit_x: 1024
+        thread_dim_limit_y: 1024
+        thread_dim_limit_z: 64
+        device_address_bits: 64
+        pcie_bandwidth: 34359738368
+        ecc_enabled: true
+        mem_clock_ghz: 2.619
+        reserved_shared_memory_per_block: 1024
+        max_blocks_per_multiprocessor: 32
+        collective_memory_granularity: 2097152
+      }
+      platform_name: "CUDA"
+      device_description_str: "NVIDIA H100 80GB HBM3"
+      runtime_version {
+        major: 13
+        minor: 2
+      }
+    }
+    host_target_machine_options {
+      triple: "x86_64-grtev4-linux-gnu"
+      cpu: "sapphirerapids"
+    }
+    num_devices_per_process: 2
+  }
+}
diff --git a/third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_h100/v4/ReduceScatter.pbtxt b/third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_h100/v4/ReduceScatter.pbtxt
new file mode 100644
index 00000000000..6b3bb48bdf4
--- /dev/null
+++ b/third_party/xla/xla/tests/aot_compatibility/gpu/executables/collective_ops_aot_test_2gpu_h100/v4/ReduceScatter.pbtxt
@@ -0,0 +1,1893 @@
+# proto-file: third_party/tensorflow/compiler/xla/util/split_proto/human_readable_aot_executable.proto
+# proto-message: xla.HumanReadableAotExecutable
+executable_and_options {
+  compile_options {
+    executable_build_options {
+      device_ordinal: -1
+      result_layout {
+        element_type: U32
+        dimensions: 4
+        layout {
+          minor_to_major: 0
+          tail_padding_alignment_in_elements: 1
+        }
+        is_dynamic_dimension: false
+      }
+      debug_options {
+        xla_disable_hlo_passes: "constant_folding"
+        xla_backend_optimization_level: 3
+        xla_eliminate_hlo_implicit_broadcast: true
+        xla_cpu_multi_thread_eigen: true
+        xla_gpu_cuda_data_dir: "./cuda_sdk_lib"
+        xla_llvm_enable_alias_scope_metadata: true
+        xla_llvm_enable_noalias_metadata: true
+        xla_llvm_enable_invariant_load_metadata: true
+        xla_llvm_disable_expensive_passes: false
+        xla_cpu_use_onednn: false
+        xla_cpu_enable_fast_math: false
+        xla_gpu_enable_fast_min_max: false
+        xla_force_host_platform_device_count: 1
+        xla_gpu_enable_dynamic_slice_fusion: true
+        xla_hlo_evaluator_use_fast_path: true
+        xla_dump_hlo_as_html: false
+        xla_cpu_fast_math_honor_nans: true
+        xla_cpu_fast_math_honor_infs: true
+        xla_allow_excess_precision: true
+        xla_gpu_autotune_level: 4
+        xla_cpu_fast_math_honor_division: true
+        xla_cpu_fast_math_honor_functions: true
+        xla_dump_include_timestamp: false
+        xla_dump_max_hlo_modules: -1
+        xla_cpu_enable_xprof_traceme: false
+        xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false
+        xla_cpu_enable_fast_min_max: true
+        xla_multiheap_size_constraint_per_heap: -1
+        xla_dump_module_metadata: false
+        xla_gpu_force_compilation_parallelism: 1
+        xla_dump_fusion_visualization: false
+        xla_gpu_all_reduce_combine_threshold_bytes: 31457287
+        xla_gpu_nccl_termination_timeout_seconds: -1
+        xla_dump_hlo_as_long_text: true
+        xla_gpu_enable_cublaslt: true
+        xla_gpu_shape_checks: RUNTIME
+        xla_gpu_use_runtime_fusion: false
+        xla_dump_latency_hiding_schedule: false
+        xla_dump_enable_mlir_pretty_form: true
+        xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP
+        xla_gpu_enable_triton_gemm: true
+        xla_gpu_enable_cudnn_int8x32_convolution_reordering: true
+        xla_gpu_triton_gemm_any: true
+        xla_gpu_enable_while_loop_reduce_scatter_code_motion: false
+        xla_gpu_collective_inflation_factor: 1
+        xla_gpu_graph_min_graph_size: 5
+        xla_gpu_enable_reassociation_for_converted_ar: true
+        xla_gpu_pgle_profile_file_or_directory_path: ""
+        xla_gpu_all_gather_combine_threshold_bytes: 31457287
+        xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287
+        xla_gpu_enable_highest_priority_async_stream: true
+        xla_gpu_exhaustive_tiling_search: false
+        xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0
+        xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1
+        xla_gpu_redzone_padding_bytes: 8388608
+        xla_gpu_copy_insertion_use_region_analysis: false
+        xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807
+        xla_gpu_collect_cost_model_stats: false
+        xla_gpu_enable_pipelined_p2p: false
+        xla_gpu_cublas_fallback: true
+        xla_gpu_enable_while_loop_double_buffering: false
+        xla_gpu_filter_kernels_spilling_registers_on_autotuning: true
+        xla_debug_buffer_assignment_show_max: 15
+        xla_detailed_logging: true
+        xla_enable_dumping: true
+        xla_gpu_enable_all_gather_combine_by_dim: false
+        xla_gpu_enable_analytical_latency_estimator: false
+        xla_gpu_llvm_verification_level: 0
+        xla_gpu_enable_reduce_scatter_combine_by_dim: false
+        xla_gpu_enable_command_buffer: COLLECTIVES_KERNEL
+        xla_gpu_enable_command_buffer: CONDITIONAL
+        xla_gpu_enable_command_buffer: CUBLAS
+        xla_gpu_enable_command_buffer: CUBLASLT
+        xla_gpu_enable_command_buffer: CUDNN
+        xla_gpu_enable_command_buffer: CUSTOM_CALL
+        xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION
+        xla_gpu_enable_command_buffer: FUSION
+        xla_gpu_enable_cub_radix_sort: true
+        xla_gpu_memory_limit_slop_factor: 95
+        xla_gpu_target_config_filename: ""
+        xla_gpu_enable_cudnn_layer_norm: false
+        xla_gpu_threshold_for_windowed_einsum_mib: 100000
+        xla_gpu_enable_nccl_user_buffers: false
+        xla_gpu_enable_libnvptxcompiler: true
+        xla_gpu_enable_nccl_comm_splitting: true
+        xla_gpu_nccl_collective_max_nchannels: 0
+        xla_gpu_nccl_p2p_max_nchannels: 0
+        xla_gpu_nccl_init_max_rank_per_root_ratio: 0
+        xla_gpu_multi_streamed_windowed_einsum: true
+        xla_gpu_gemm_rewrite_size_threshold: 100
+        xla_gpu_require_complete_aot_autotune_results: false
+        xla_gpu_cudnn_gemm_fusion_level: 0
+        xla_gpu_use_memcpy_local_p2p: false
+        xla_gpu_autotune_max_solutions: 0
+        xla_dump_large_constants: false
+        xla_gpu_verify_triton_fusion_numerics: false
+        xla_reduce_window_rewrite_base_length: 16
+        xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL
+        xla_gpu_enable_host_memory_offloading: false
+        xla_gpu_exclude_nondeterministic_ops: true
+        xla_gpu_nccl_terminate_on_error: false
+        xla_gpu_shard_autotuning: true
+        xla_gpu_enable_approx_costly_collectives: false
+        xla_cpu_prefer_vector_width: 256
+        xla_gpu_per_fusion_autotune_cache_dir: ""
+        xla_cmd_buffer_trace_cache_size: 16
+        xla_gpu_temp_buffer_use_separate_color: false
+        xla_syntax_sugar_async_ops: false
+        xla_gpu_autotune_gemm_rtol: 0.1
+        xla_enable_command_buffers_during_profiling: false
+        xla_gpu_cudnn_gemm_max_plans: 5
+        xla_cpu_parallel_codegen_split_count: 32
+        xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE
+        xla_gpu_executable_warn_stuck_timeout_seconds: 10
+        xla_gpu_executable_terminate_timeout_seconds: 30
+        xla_gpu_experimental_disable_binary_libraries: false
+        xla_gpu_dot_merger_threshold_mb: 64
+        xla_cpu_max_isa: ""
+        xla_gpu_experimental_enable_fusion_block_level_rewriter: false
+        xla_enable_fast_math: false
+        xla_gpu_experimental_parallel_collective_overlap_limit: 1
+        xla_cpu_copy_insertion_use_region_analysis: false
+        xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1
+        xla_gpu_experimental_enable_triton_heroless_priority_fusion: false
+        xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN
+        xla_gpu_experimental_stream_annotation: true
+        xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO
+        xla_pjrt_allow_auto_layout_in_hlo: false
+        xla_gpu_enable_scatter_determinism_expander: false
+        xla_gpu_require_exclusive_lock: false
+        xla_gpu_generate_debug_info: false
+        xla_gpu_generate_line_info: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false
+        xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE
+        xla_gpu_fail_ptx_compilation_on_register_spilling: false
+        xla_gpu_enable_analytical_sol_latency_estimator: true
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_prep_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_size_bytes"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "gpus_per_node"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nccl_op_launch_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nic_speed_gbps"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "rtt_us"
+          value: "-1"
+        }
+        xla_cpu_use_xnnpack: true
+        xla_gpu_experimental_pack_dot_operands_along_k_dimension: true
+        xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false
+        xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED
+        xla_gpu_experimental_enable_heuristic_collective_combining: false
+        xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false
+        xla_hlo_pass_fix_detect_cycles: false
+        xla_cpu_generate_unique_c_style_kernel_entry_points: false
+        xla_gpu_experimental_collective_cse_distance_threshold: 0
+        xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true
+        xla_gpu_experimental_collective_perf_table_path: ""
+        xla_gpu_collective_permute_combine_threshold_bytes: 31457287
+        xla_unsupported_crash_on_hlo_pass_noop_change: false
+        xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false
+        xla_dump_full_hlo_config: true
+        xla_gpu_unsupported_enable_triton_multi_output_fusion: false
+        xla_gpu_experimental_matmul_perf_table_path: ""
+        xla_gpu_unsupported_enable_all_reduce_decomposer: false
+        xla_gpu_nccl_blocking_communicators: true
+        xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20
+        xla_gpu_first_collective_call_terminate_timeout_seconds: 40
+        xla_gpu_nccl_async_execution: false
+        xla_cpu_emitter_verification_level: 1
+        xla_early_exit_with_layouts: false
+        xla_gpu_experimental_use_ragged_dot_fusion: false
+        xla_detect_unstable_reductions: DETECTION_MODE_NONE
+        xla_gpu_command_buffer_scheduling_mode: LHS
+        xla_gpu_experimental_enable_nccl_symmetric_buffers: false
+        xla_disable_automatic_host_compute_offload: false
+        xla_gpu_experimental_enable_fusion_autotuner: true
+        xla_gpu_experimental_scaled_dot_with_triton: true
+        xla_gpu_command_buffer_unroll_loops: false
+        xla_cpu_experimental_onednn_custom_call: false
+        xla_gpu_experimental_enable_checksum_tracing_on_thunks: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true
+        xla_cpu_collective_call_terminate_timeout_seconds: 40
+        xla_cpu_collective_call_warn_stuck_seconds: 20
+        xla_keep_shardings_after_spmd: false
+        xla_gpu_experimental_enable_triton_warp_specialization: false
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE
+        xla_cpu_enable_platform_dependent_math: true
+        xla_gpu_detect_nan: DETECTION_MODE_NONE
+        xla_gpu_detect_inf: DETECTION_MODE_NONE
+        xla_enable_enzyme_comms_opt: false
+        xla_gpu_experimental_enable_buffer_saver_on_thunks: false
+        xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE
+        xla_enable_scoped_logging_timers: true
+        xla_gpu_executable_embed_debug_info: true
+        xla_cpu_collective_timeout_seconds: 1800
+        xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false
+        xla_gpu_enable_pipelined_host_offloading: false
+        xla_gpu_default_to_alg_dot_bf16_bf16_f32: false
+        xla_gpu_experimental_autotune_backends: CUDNN
+        xla_gpu_experimental_autotune_backends: TRITON
+        xla_gpu_experimental_autotune_backends: CUBLASLT
+        xla_gpu_experimental_autotune_backends: HIPBLASLT
+        xla_gpu_experimental_autotune_backends: MIOPEN
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER
+        xla_gpu_experimental_autotune_backends: NATIVE_EMITTER
+        xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER
+        xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION
+        xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: YNNPACK
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU
+        xla_gpu_experimental_autotune_backends: TENSOR_IR
+        xla_gpu_experimental_all_fusions_with_triton: false
+        xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT
+        xla_gpu_rocm_max_trace_events: 4194304
+        xla_gpu_execution_terminate_timeout: "inf"
+        xla_enable_hlo_sharding_v3: true
+        xla_gpu_log_minmax: false
+        xla_gpu_print_compilation_stats: false
+        xla_gpu_experimental_enable_conv_fusion: false
+        xla_gpu_experimental_enable_tiling_propagation: true
+        xla_gpu_execution_progress_tracking: 0
+        xla_gpu_experimental_max_unroll_factor: 32
+        xla_gpu_enable_pdl: true
+        xla_gpu_collective_combine_threshold_count: 2048
+        xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true
+        xla_recognize_reduction_optimization_level: 0
+        xla_gpu_experimental_parallel_async_compute_limit: 2
+        xla_dump_buffer_assignment_analysis: true
+        xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE
+        xla_gpu_experimental_enable_collective_multi_streaming: false
+        xla_gpu_collective_permute_connected_components: false
+        xla_gpu_experimental_force_split_k: 0
+        xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_experimental_gemm_fusion_v2: true
+        xla_gpu_fusion_autotune_top_k_configs: 1
+        xla_gpu_native_emitter_tune_unroll_factor_for_loops: false
+        xla_enable_rgv3_materialization: true
+        xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK
+        xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false
+        xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_executable_num_compute_streams: 0
+        xla_gpu_executable_num_communication_streams: 0
+        xla_gpu_experimental_collective_start_as_early_as_possible: false
+        xla_cpu_experimental_enable_tiling_propagation: true
+        xla_gpu_enable_dus_accumulator_zero_init_elimination: false
+        xla_gpu_experimental_thunk_buffer_debug_module_outputs: false
+        xla_cpu_ftz: true
+        xla_gpu_enable_allocator_spatial_partitioning: true
+        xla_dump_compact_gte: true
+        xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES
+        xla_dump_hlo_as_riegeli: false
+        xla_gpu_experimental_use_ragged_dot_grouped_gemm: true
+        xla_dump_hlo_to_subfolder: false
+        xla_gpu_enable_pdl_launch: true
+        xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED
+        xla_gpu_enable_nccl_user_buffers_in_default_space: false
+        xla_sdy_export_all_reduce_scatter: false
+        xla_gpu_enable_gxl_ragged_all_to_all: false
+        xla_gpu_gxl_scratch_size_bytes: 67108864
+        xla_gpu_use_new_autotune_cache_format: true
+        xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE
+        xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON
+        xla_spmd_enable_dynamic_slice_collective_broadcast: false
+        xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false
+        xla_gpu_async_copy_min_bytes: -1
+        xla_gpu_blas_max_algorithms: 0
+        xla_gpu_experimental_enable_selective_memcpy_overlap: false
+        xla_gpu_experimental_enable_same_shape_multi_output_fusion: false
+        xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1
+        xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1
+        xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240
+        xla_cpu_use_new_xtile_lowering: true
+        xla_gpu_trace_annotation_level: 0
+        xla_gpu_experimental_cost_model_gemm_tiling_default: false
+        xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false
+        xla_gpu_all_reduce_splitter_ignore_profitability_check: false
+        xla_gpu_enable_cupti_multi_subscriber: true
+        xla_enable_hlo_modules_upload: true
+        xla_gpu_collective_domain_assignment: ""
+        xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1
+        xla_compile_all_supported_configs: false
+        xla_gpu_unsupported_use_cross_host_one_shot_kernel: ALLCOLLECTIVES
+        xla_deduplicate_backend_configs_min_size: 9223372036854775807
+        xla_force_config: ""
+        xla_candidate_configs_file: ""
+        xla_gpu_enable_persistent_symmetric_memory: false
+        xla_gpu_experimental_enable_raft_for_stable_topk: false
+        xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND
+        xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET
+        xla_gpu_experimental_emit_collective_reduce: false
+        xla_gpu_experimental_vmm_disabled: false
+        xla_gpu_experimental_enable_dynamic_slice_table_offsets: false
+        xla_gpu_device_execution_terminate_timeout: "inf"
+        xla_gpu_topology_filename: ""
+        xla_gpu_experimental_enable_dynamic_slice_extended_offsets: false
+        xla_gpu_experimental_deterministic_ops_triton: DETERMINISTIC_OPS_TRITON_MODE_OFF
+        xla_gpu_experimental_enable_tensor_ir: false
+      }
+      num_replicas: 2
+      num_partitions: 1
+      device_assignment {
+        replica_count: 2
+        computation_count: 1
+        computation_devices {
+          replica_device_ids: 0
+          replica_device_ids: 1
+        }
+      }
+      allow_spmd_sharding_propagation_to_output: false
+      comp_envs {
+      }
+      allow_spmd_sharding_propagation_to_parameters: false
+      process_count: 1
+      memory_fitting_level: EFFORT_O2
+    }
+  }
+  pjrt_client_name: "PjRtStreamExecutorClient"
+}
+gpu_executable {
+  hlo_module_with_config {
+    hlo_module {
+      name: "test"
+      entry_computation_name: "main"
+      computations {
+        name: "add"
+        instructions {
+          name: "lhs"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "lhs"
+          }
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "rhs"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "rhs"
+          }
+          parameter_number: 1
+          id: 1
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "add.1"
+          opcode: "add"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "add.1"
+          }
+          id: 2
+          operand_ids: 0
+          operand_ids: 1
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          parameters {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          parameters {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          result {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          parameter_names: "lhs"
+          parameter_names: "rhs"
+        }
+        root_id: 2
+      }
+      computations {
+        name: "fused_select"
+        instructions {
+          name: "param_2.2"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "param_2.2"
+          }
+          parameter_number: 2
+          id: 8589934597
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "zero_2"
+          opcode: "constant"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+          }
+          literal {
+            shape {
+              element_type: U32
+              layout {
+                tail_padding_alignment_in_elements: 1
+              }
+            }
+            u32s: 0
+          }
+          id: 8589934598
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "p.2"
+          opcode: "compare"
+          shape {
+            element_type: PRED
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "p.2"
+          }
+          id: 8589934596
+          operand_ids: 8589934597
+          operand_ids: 8589934598
+          comparison_direction: "EQ"
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+          comparison_order: "TOTAL"
+        }
+        instructions {
+          name: "pb.2"
+          opcode: "broadcast"
+          shape {
+            element_type: PRED
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "pb.2"
+          }
+          id: 8589934593
+          operand_ids: 8589934596
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "param_0.1"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "param_0.1"
+          }
+          id: 8589934594
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "param_1.1"
+          opcode: "parameter"
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "param_1.1"
+          }
+          parameter_number: 1
+          id: 8589934595
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "data.2"
+          opcode: "select"
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "data.2"
+          }
+          id: 8589934592
+          operand_ids: 8589934593
+          operand_ids: 8589934594
+          operand_ids: 8589934595
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        program_shape {
+          parameters {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          parameters {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          parameters {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          result {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          parameter_names: "param_0.1"
+          parameter_names: "param_1.1"
+          parameter_names: "param_2.2"
+        }
+        id: 2
+        root_id: 8589934592
+        is_fusion_computation: true
+      }
+      computations {
+        name: "main"
+        instructions {
+          name: "c0_1"
+          opcode: "constant"
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+          }
+          literal {
+            shape {
+              element_type: U32
+              dimensions: 8
+              layout {
+                minor_to_major: 0
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+            }
+            u32s: 1
+            u32s: 2
+            u32s: 3
+            u32s: 4
+            u32s: 5
+            u32s: 6
+            u32s: 7
+            u32s: 8
+          }
+          id: 4294967297
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "c1_1"
+          opcode: "constant"
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+          }
+          literal {
+            shape {
+              element_type: U32
+              dimensions: 8
+              layout {
+                minor_to_major: 0
+                tail_padding_alignment_in_elements: 1
+              }
+              is_dynamic_dimension: false
+            }
+            u32s: 10
+            u32s: 11
+            u32s: 12
+            u32s: 13
+            u32s: 14
+            u32s: 15
+            u32s: 16
+            u32s: 17
+          }
+          id: 4294967298
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "id.1"
+          opcode: "replica-id"
+          shape {
+            element_type: U32
+            layout {
+              tail_padding_alignment_in_elements: 1
+            }
+          }
+          metadata {
+            scheduling_name: "id.1"
+          }
+          id: 4294967296
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "loop_select_fusion"
+          opcode: "fusion"
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "loop_select_fusion"
+          }
+          fusion_kind: "kLoop"
+          id: 4294967299
+          operand_ids: 4294967297
+          operand_ids: 4294967298
+          operand_ids: 4294967296
+          called_computation_ids: 2
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+        }
+        instructions {
+          name: "ars.3"
+          opcode: "reduce-scatter"
+          shape {
+            element_type: U32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+          metadata {
+            scheduling_name: "reduce-scatter-start"
+          }
+          dimensions: 0
+          id: 4294967300
+          operand_ids: 4294967299
+          called_computation_ids: 0
+          frontend_attributes {
+          }
+          statistics_viz {
+          }
+          collective_device_list {
+          }
+          backend_config_payload {
+            id: 0
+          }
+        }
+        program_shape {
+          result {
+            element_type: U32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+        }
+        id: 1
+        root_id: 4294967300
+      }
+      host_program_shape {
+        result {
+          element_type: U32
+          dimensions: 4
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+      }
+      entry_computation_id: 1
+      schedule {
+        sequences {
+          key: 0
+          value {
+            instruction_ids: 1
+            instruction_ids: 0
+            instruction_ids: 2
+          }
+        }
+        sequences {
+          key: 1
+          value {
+            instruction_ids: 4294967296
+            instruction_ids: 4294967298
+            instruction_ids: 4294967297
+            instruction_ids: 4294967299
+            instruction_ids: 4294967300
+          }
+        }
+      }
+      input_output_alias {
+      }
+      device_assignment {
+        replica_count: 2
+        computation_count: 1
+        computation_devices {
+          replica_device_ids: 0
+          replica_device_ids: 1
+        }
+      }
+      buffer_donor {
+      }
+      frontend_attributes {
+        map {
+          key: "fingerprint_before_lhs"
+          value: "c36d8870be58d238935bebefbbc5b827"
+        }
+      }
+      payloads: "{\"collective_backend_config\":{\"backend\":\"DEFAULT\",\"collectives_mode\":\"COLLECTIVES_MODE_INVALID\",\"communication_domain\":\"COLLECTIVE_COMMUNICATION_DOMAIN_UNSPECIFIED\",\"is_pipelined\":false,\"is_spmd_generated\":false,\"is_sync\":true,\"kernel_strategy\":\"KERNEL_STRATEGY_DEFAULT\"},\"device_type\":\"DEVICE_TYPE_INVALID\",\"force_earliest_schedule\":false,\"operation_queue_id\":\"0\",\"reification_cost\":[]}"
+    }
+    config {
+      entry_computation_layout {
+        result {
+          element_type: U32
+          dimensions: 4
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+      }
+      replica_count: 2
+      num_partitions: 1
+      intra_op_parallelism_threads: -1
+      debug_options {
+        xla_disable_hlo_passes: "constant_folding"
+        xla_backend_optimization_level: 3
+        xla_eliminate_hlo_implicit_broadcast: true
+        xla_cpu_multi_thread_eigen: true
+        xla_gpu_cuda_data_dir: "./cuda_sdk_lib"
+        xla_llvm_enable_alias_scope_metadata: true
+        xla_llvm_enable_noalias_metadata: true
+        xla_llvm_enable_invariant_load_metadata: true
+        xla_llvm_disable_expensive_passes: false
+        xla_cpu_use_onednn: false
+        xla_cpu_enable_fast_math: false
+        xla_gpu_enable_fast_min_max: false
+        xla_force_host_platform_device_count: 1
+        xla_gpu_enable_dynamic_slice_fusion: true
+        xla_hlo_evaluator_use_fast_path: true
+        xla_dump_hlo_as_html: false
+        xla_cpu_fast_math_honor_nans: true
+        xla_cpu_fast_math_honor_infs: true
+        xla_allow_excess_precision: true
+        xla_gpu_autotune_level: 4
+        xla_cpu_fast_math_honor_division: true
+        xla_cpu_fast_math_honor_functions: true
+        xla_dump_include_timestamp: false
+        xla_dump_max_hlo_modules: -1
+        xla_cpu_enable_xprof_traceme: false
+        xla_gpu_unsafe_fallback_to_driver_on_ptxas_not_found: false
+        xla_cpu_enable_fast_min_max: true
+        xla_multiheap_size_constraint_per_heap: -1
+        xla_dump_module_metadata: false
+        xla_gpu_force_compilation_parallelism: 1
+        xla_dump_fusion_visualization: false
+        xla_gpu_all_reduce_combine_threshold_bytes: 31457287
+        xla_gpu_nccl_termination_timeout_seconds: -1
+        xla_dump_hlo_as_long_text: true
+        xla_gpu_enable_cublaslt: true
+        xla_gpu_shape_checks: RUNTIME
+        xla_gpu_use_runtime_fusion: false
+        xla_dump_latency_hiding_schedule: false
+        xla_dump_enable_mlir_pretty_form: true
+        xla_partitioning_algorithm: PARTITIONING_ALGORITHM_NOOP
+        xla_gpu_enable_triton_gemm: true
+        xla_gpu_enable_cudnn_int8x32_convolution_reordering: true
+        xla_gpu_triton_gemm_any: true
+        xla_gpu_enable_while_loop_reduce_scatter_code_motion: false
+        xla_gpu_collective_inflation_factor: 1
+        xla_gpu_graph_min_graph_size: 5
+        xla_gpu_enable_reassociation_for_converted_ar: true
+        xla_gpu_pgle_profile_file_or_directory_path: ""
+        xla_gpu_all_gather_combine_threshold_bytes: 31457287
+        xla_gpu_reduce_scatter_combine_threshold_bytes: 31457287
+        xla_gpu_enable_highest_priority_async_stream: true
+        xla_gpu_exhaustive_tiling_search: false
+        xla_gpu_auto_spmd_partitioning_memory_budget_gb: 0
+        xla_gpu_auto_spmd_partitioning_memory_budget_ratio: 1.1
+        xla_gpu_redzone_padding_bytes: 8388608
+        xla_gpu_copy_insertion_use_region_analysis: false
+        xla_gpu_collective_permute_decomposer_threshold: 9223372036854775807
+        xla_gpu_collect_cost_model_stats: false
+        xla_gpu_enable_pipelined_p2p: false
+        xla_gpu_cublas_fallback: true
+        xla_gpu_enable_while_loop_double_buffering: false
+        xla_gpu_filter_kernels_spilling_registers_on_autotuning: true
+        xla_debug_buffer_assignment_show_max: 15
+        xla_detailed_logging: true
+        xla_enable_dumping: true
+        xla_gpu_enable_all_gather_combine_by_dim: false
+        xla_gpu_enable_analytical_latency_estimator: false
+        xla_gpu_llvm_verification_level: 0
+        xla_gpu_enable_reduce_scatter_combine_by_dim: false
+        xla_gpu_enable_command_buffer: COLLECTIVES_KERNEL
+        xla_gpu_enable_command_buffer: CONDITIONAL
+        xla_gpu_enable_command_buffer: CUBLAS
+        xla_gpu_enable_command_buffer: CUBLASLT
+        xla_gpu_enable_command_buffer: CUDNN
+        xla_gpu_enable_command_buffer: CUSTOM_CALL
+        xla_gpu_enable_command_buffer: DYNAMIC_SLICE_FUSION
+        xla_gpu_enable_command_buffer: FUSION
+        xla_gpu_enable_cub_radix_sort: true
+        xla_gpu_memory_limit_slop_factor: 95
+        xla_gpu_target_config_filename: ""
+        xla_gpu_enable_cudnn_layer_norm: false
+        xla_gpu_threshold_for_windowed_einsum_mib: 100000
+        xla_gpu_enable_nccl_user_buffers: false
+        xla_gpu_enable_libnvptxcompiler: true
+        xla_gpu_enable_nccl_comm_splitting: true
+        xla_gpu_nccl_collective_max_nchannels: 0
+        xla_gpu_nccl_p2p_max_nchannels: 0
+        xla_gpu_nccl_init_max_rank_per_root_ratio: 0
+        xla_gpu_multi_streamed_windowed_einsum: true
+        xla_gpu_gemm_rewrite_size_threshold: 100
+        xla_gpu_require_complete_aot_autotune_results: false
+        xla_gpu_cudnn_gemm_fusion_level: 0
+        xla_gpu_use_memcpy_local_p2p: false
+        xla_gpu_autotune_max_solutions: 0
+        xla_dump_large_constants: false
+        xla_gpu_verify_triton_fusion_numerics: false
+        xla_reduce_window_rewrite_base_length: 16
+        xla_gpu_enable_while_loop_unrolling: WHILE_LOOP_UNROLLING_AUTO_UNROLL
+        xla_gpu_enable_host_memory_offloading: false
+        xla_gpu_exclude_nondeterministic_ops: true
+        xla_gpu_nccl_terminate_on_error: false
+        xla_gpu_shard_autotuning: true
+        xla_gpu_enable_approx_costly_collectives: false
+        xla_cpu_prefer_vector_width: 256
+        xla_gpu_per_fusion_autotune_cache_dir: ""
+        xla_cmd_buffer_trace_cache_size: 16
+        xla_gpu_temp_buffer_use_separate_color: false
+        xla_syntax_sugar_async_ops: false
+        xla_gpu_autotune_gemm_rtol: 0.1
+        xla_enable_command_buffers_during_profiling: false
+        xla_gpu_cudnn_gemm_max_plans: 5
+        xla_cpu_parallel_codegen_split_count: 32
+        xla_gpu_experimental_autotune_cache_mode: AUTOTUNE_CACHE_MODE_UPDATE
+        xla_gpu_executable_warn_stuck_timeout_seconds: 10
+        xla_gpu_executable_terminate_timeout_seconds: 30
+        xla_gpu_experimental_disable_binary_libraries: false
+        xla_gpu_dot_merger_threshold_mb: 64
+        xla_cpu_max_isa: ""
+        xla_gpu_experimental_enable_fusion_block_level_rewriter: false
+        xla_enable_fast_math: false
+        xla_gpu_experimental_parallel_collective_overlap_limit: 1
+        xla_cpu_copy_insertion_use_region_analysis: false
+        xla_gpu_operand_bytes_threshold_for_windowed_einsum: -1
+        xla_gpu_experimental_enable_triton_heroless_priority_fusion: false
+        xla_gpu_pgle_accuracy_checker: PGLE_STRICTNESS_LEVEL_WARN
+        xla_gpu_experimental_stream_annotation: true
+        xla_gpu_libnvjitlink_mode: LIB_NV_JIT_LINK_MODE_AUTO
+        xla_pjrt_allow_auto_layout_in_hlo: false
+        xla_gpu_enable_scatter_determinism_expander: false
+        xla_gpu_require_exclusive_lock: false
+        xla_gpu_generate_debug_info: false
+        xla_gpu_generate_line_info: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_decomposer: false
+        xla_gpu_experimental_pipeline_parallelism_opt_level: PIPELINE_PARALLELISM_OPT_LEVEL_DISABLE
+        xla_gpu_fail_ptx_compilation_on_register_spilling: false
+        xla_gpu_enable_analytical_sol_latency_estimator: true
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_prep_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "chunk_size_bytes"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "gpus_per_node"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nccl_op_launch_us"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "nic_speed_gbps"
+          value: "-1"
+        }
+        xla_gpu_analytical_latency_estimator_options {
+          key: "rtt_us"
+          value: "-1"
+        }
+        xla_cpu_use_xnnpack: true
+        xla_gpu_experimental_pack_dot_operands_along_k_dimension: true
+        xla_unsupported_crash_on_hlo_pass_fix_max_iterations: false
+        xla_cpu_experimental_xnn_graph_fusion_mode: XNN_GRAPH_FUSION_MODE_DISABLED
+        xla_gpu_experimental_enable_heuristic_collective_combining: false
+        xla_gpu_experimental_enable_subchannel_dequantisation_fusion: false
+        xla_hlo_pass_fix_detect_cycles: false
+        xla_cpu_generate_unique_c_style_kernel_entry_points: false
+        xla_gpu_experimental_collective_cse_distance_threshold: 0
+        xla_gpu_unsupported_use_ragged_all_to_all_one_shot_kernel: true
+        xla_gpu_experimental_collective_perf_table_path: ""
+        xla_gpu_collective_permute_combine_threshold_bytes: 31457287
+        xla_unsupported_crash_on_hlo_pass_noop_change: false
+        xla_unsupported_crash_on_hlo_pass_silent_hlo_change: false
+        xla_dump_full_hlo_config: true
+        xla_gpu_unsupported_enable_triton_multi_output_fusion: false
+        xla_gpu_experimental_matmul_perf_table_path: ""
+        xla_gpu_unsupported_enable_all_reduce_decomposer: false
+        xla_gpu_nccl_blocking_communicators: true
+        xla_gpu_first_collective_call_warn_stuck_timeout_seconds: 20
+        xla_gpu_first_collective_call_terminate_timeout_seconds: 40
+        xla_gpu_nccl_async_execution: false
+        xla_cpu_emitter_verification_level: 1
+        xla_early_exit_with_layouts: false
+        xla_gpu_experimental_use_ragged_dot_fusion: false
+        xla_detect_unstable_reductions: DETECTION_MODE_NONE
+        xla_gpu_command_buffer_scheduling_mode: LHS
+        xla_gpu_experimental_enable_nccl_symmetric_buffers: false
+        xla_disable_automatic_host_compute_offload: false
+        xla_gpu_experimental_enable_fusion_autotuner: true
+        xla_gpu_experimental_scaled_dot_with_triton: true
+        xla_gpu_command_buffer_unroll_loops: false
+        xla_cpu_experimental_onednn_custom_call: false
+        xla_gpu_experimental_enable_checksum_tracing_on_thunks: false
+        xla_gpu_unsupported_enable_ragged_all_to_all_multi_host_decomposer: true
+        xla_cpu_collective_call_terminate_timeout_seconds: 40
+        xla_cpu_collective_call_warn_stuck_seconds: 20
+        xla_keep_shardings_after_spmd: false
+        xla_gpu_experimental_enable_triton_warp_specialization: false
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_DOT
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_INDIVIDUAL_CONVOLUTION
+        xla_cpu_experimental_ynn_fusion_type: LIBRARY_FUSION_TYPE_REDUCE
+        xla_cpu_enable_platform_dependent_math: true
+        xla_gpu_detect_nan: DETECTION_MODE_NONE
+        xla_gpu_detect_inf: DETECTION_MODE_NONE
+        xla_enable_enzyme_comms_opt: false
+        xla_gpu_experimental_enable_buffer_saver_on_thunks: false
+        xla_detect_unstable_reductions_post_optimizations: DETECTION_MODE_NONE
+        xla_enable_scoped_logging_timers: true
+        xla_gpu_executable_embed_debug_info: true
+        xla_cpu_collective_timeout_seconds: 1800
+        xla_allow_h2h_copy_when_automatic_host_compute_offload_disabled: false
+        xla_gpu_enable_pipelined_host_offloading: false
+        xla_gpu_default_to_alg_dot_bf16_bf16_f32: false
+        xla_gpu_experimental_autotune_backends: CUDNN
+        xla_gpu_experimental_autotune_backends: TRITON
+        xla_gpu_experimental_autotune_backends: CUBLASLT
+        xla_gpu_experimental_autotune_backends: HIPBLASLT
+        xla_gpu_experimental_autotune_backends: MIOPEN
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER
+        xla_gpu_experimental_autotune_backends: NATIVE_EMITTER
+        xla_gpu_experimental_autotune_backends: LLVM_KERNEL_EMITTER
+        xla_gpu_experimental_autotune_backends: CUBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: CUSTOM_KERNEL_FISSION
+        xla_gpu_experimental_autotune_backends: HIPBLASLT_FISSION
+        xla_gpu_experimental_autotune_backends: YNNPACK
+        xla_gpu_experimental_autotune_backends: BLOCK_LEVEL_EMITTER_CPU
+        xla_gpu_experimental_autotune_backends: TENSOR_IR
+        xla_gpu_experimental_all_fusions_with_triton: false
+        xla_cpu_scheduler_type: CPU_SCHEDULER_TYPE_DEFAULT
+        xla_gpu_rocm_max_trace_events: 4194304
+        xla_gpu_execution_terminate_timeout: "inf"
+        xla_enable_hlo_sharding_v3: true
+        xla_gpu_log_minmax: false
+        xla_gpu_print_compilation_stats: false
+        xla_gpu_experimental_enable_conv_fusion: false
+        xla_gpu_experimental_enable_tiling_propagation: true
+        xla_gpu_execution_progress_tracking: 0
+        xla_gpu_experimental_max_unroll_factor: 32
+        xla_gpu_enable_pdl: true
+        xla_gpu_collective_combine_threshold_count: 2048
+        xla_gpu_experimental_ragged_all_to_all_use_barrier_with_nccl: true
+        xla_recognize_reduction_optimization_level: 0
+        xla_gpu_experimental_parallel_async_compute_limit: 2
+        xla_dump_buffer_assignment_analysis: true
+        xla_gpu_command_buffer_update_mode: ALWAYS_UPDATE
+        xla_gpu_experimental_enable_collective_multi_streaming: false
+        xla_gpu_collective_permute_connected_components: false
+        xla_gpu_experimental_force_split_k: 0
+        xla_gpu_collective_permute_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_experimental_gemm_fusion_v2: true
+        xla_gpu_fusion_autotune_top_k_configs: 1
+        xla_gpu_native_emitter_tune_unroll_factor_for_loops: false
+        xla_enable_rgv3_materialization: true
+        xla_gpu_ragged_all_to_all_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_deviceless_cub_mode: DEVICELESS_CUB_WITH_FALLBACK
+        xla_gpu_experimental_dynamic_slice_fusion_verify_offsets: false
+        xla_gpu_all_gather_mode: COLLECTIVES_PRIVATE_MEMORY
+        xla_gpu_executable_num_compute_streams: 0
+        xla_gpu_executable_num_communication_streams: 0
+        xla_gpu_experimental_collective_start_as_early_as_possible: false
+        xla_cpu_experimental_enable_tiling_propagation: true
+        xla_gpu_enable_dus_accumulator_zero_init_elimination: false
+        xla_gpu_experimental_thunk_buffer_debug_module_outputs: false
+        xla_cpu_ftz: true
+        xla_gpu_enable_allocator_spatial_partitioning: true
+        xla_dump_compact_gte: true
+        xla_gpu_enable_collectives_command_buffer_filter: ALLCOLLECTIVES
+        xla_dump_hlo_as_riegeli: false
+        xla_gpu_experimental_use_ragged_dot_grouped_gemm: true
+        xla_dump_hlo_to_subfolder: false
+        xla_gpu_enable_pdl_launch: true
+        xla_gpu_cudnn_deviceless_compilation_mode: CUDNN_DEVICELESS_COMPILATION_DISABLED
+        xla_gpu_enable_nccl_user_buffers_in_default_space: false
+        xla_sdy_export_all_reduce_scatter: false
+        xla_gpu_enable_gxl_ragged_all_to_all: false
+        xla_gpu_gxl_scratch_size_bytes: 67108864
+        xla_gpu_use_new_autotune_cache_format: true
+        xla_gpu_experimental_use_collective_kernels: COLLECTIVE_KERNEL_ALL_REDUCE
+        xla_gpu_pipeline_all_gather: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_all_reduce: COLLECTIVE_PIPELINING_MODE_DEFAULT
+        xla_gpu_pipeline_reduce_scatter: COLLECTIVE_PIPELINING_MODE_ON
+        xla_spmd_enable_dynamic_slice_collective_broadcast: false
+        xla_gpu_experimental_ragged_all_to_all_use_device_kernel: false
+        xla_gpu_async_copy_min_bytes: -1
+        xla_gpu_blas_max_algorithms: 0
+        xla_gpu_experimental_enable_selective_memcpy_overlap: false
+        xla_gpu_experimental_enable_same_shape_multi_output_fusion: false
+        xla_gpu_experimental_scheduler_memory_fencing_threshold_bytes: -1
+        xla_gpu_experimental_scheduler_memory_fencing_slack_windows: 1
+        xla_while_loop_all_reduce_dus_code_motion_max_size_bytes: 10240
+        xla_cpu_use_new_xtile_lowering: true
+        xla_gpu_trace_annotation_level: 0
+        xla_gpu_experimental_cost_model_gemm_tiling_default: false
+        xla_gpu_allow_ragged_all_to_all_nccl_send_recv_fallback: false
+        xla_gpu_all_reduce_splitter_ignore_profitability_check: false
+        xla_gpu_enable_cupti_multi_subscriber: true
+        xla_enable_hlo_modules_upload: true
+        xla_gpu_collective_domain_assignment: ""
+        xla_gpu_experimental_parallel_scale_up_collective_overlap_limit: 1
+        xla_compile_all_supported_configs: false
+        xla_gpu_unsupported_use_cross_host_one_shot_kernel: ALLCOLLECTIVES
+        xla_deduplicate_backend_configs_min_size: 9223372036854775807
+        xla_force_config: ""
+        xla_candidate_configs_file: ""
+        xla_gpu_enable_persistent_symmetric_memory: false
+        xla_gpu_experimental_enable_raft_for_stable_topk: false
+        xla_autotuner_preferred_backend: UNSPECIFIED_BACKEND
+        xla_gpu_experimental_early_exit: EARLY_EXIT_POINT_UNSET
+        xla_gpu_experimental_emit_collective_reduce: false
+        xla_gpu_experimental_vmm_disabled: false
+        xla_gpu_experimental_enable_dynamic_slice_table_offsets: false
+        xla_gpu_device_execution_terminate_timeout: "inf"
+        xla_gpu_topology_filename: ""
+        xla_gpu_experimental_enable_dynamic_slice_extended_offsets: false
+        xla_gpu_experimental_deterministic_ops_triton: DETERMINISTIC_OPS_TRITON_MODE_OFF
+        xla_gpu_experimental_enable_tensor_ir: false
+      }
+      static_device_assignment {
+        replica_count: 2
+        computation_count: 1
+        computation_devices {
+          replica_device_ids: 0
+          replica_device_ids: 1
+        }
+      }
+      allow_spmd_sharding_propagation_to_output: false
+      allow_spmd_sharding_propagation_to_parameters: false
+      sharding_config {
+      }
+      memory_fitting_level: EFFORT_O2
+      schedule_config {
+      }
+      partition_size: 2
+    }
+  }
+  buffer_assignment {
+    logical_buffers {
+      size: 4
+      defined_at {
+      }
+    }
+    logical_buffers {
+      id: 1
+      size: 4
+      defined_at {
+        instruction_id: 1
+      }
+    }
+    logical_buffers {
+      id: 2
+      size: 4
+      defined_at {
+        instruction_id: 2
+      }
+    }
+    logical_buffers {
+      id: 10
+      size: 32
+      defined_at {
+        instruction_id: 4294967297
+      }
+    }
+    logical_buffers {
+      id: 11
+      size: 32
+      defined_at {
+        instruction_id: 4294967298
+      }
+    }
+    logical_buffers {
+      id: 12
+      size: 4
+      defined_at {
+        instruction_id: 4294967296
+      }
+    }
+    logical_buffers {
+      id: 13
+      size: 32
+      defined_at {
+        instruction_id: 4294967299
+      }
+    }
+    logical_buffers {
+      id: 14
+      size: 16
+      defined_at {
+        instruction_id: 4294967300
+      }
+    }
+    buffer_allocations {
+      size: 32
+      assigned {
+        logical_buffer_id: 10
+        size: 32
+        element_type: U32
+      }
+      is_constant: true
+    }
+    buffer_allocations {
+      index: 1
+      size: 32
+      assigned {
+        logical_buffer_id: 11
+        size: 32
+        element_type: U32
+      }
+      is_constant: true
+    }
+    buffer_allocations {
+      index: 2
+      size: 16
+      maybe_live_out: true
+      assigned {
+        logical_buffer_id: 12
+        size: 4
+        element_type: U32
+      }
+      assigned {
+        logical_buffer_id: 14
+        size: 16
+        element_type: U32
+      }
+    }
+    buffer_allocations {
+      index: 3
+      size: 4
+      is_thread_local: true
+      assigned {
+        logical_buffer_id: 2
+        size: 4
+        element_type: U32
+      }
+    }
+    buffer_allocations {
+      index: 4
+      size: 4
+      is_thread_local: true
+      assigned {
+        size: 4
+        element_type: U32
+      }
+    }
+    buffer_allocations {
+      index: 5
+      size: 4
+      is_thread_local: true
+      assigned {
+        logical_buffer_id: 1
+        size: 4
+        element_type: U32
+      }
+    }
+    buffer_allocations {
+      index: 6
+      size: 32
+      assigned {
+        logical_buffer_id: 13
+        size: 32
+        element_type: U32
+      }
+    }
+    heap_simulator_traces {
+      events {
+        buffer_id: 13
+        computation_name: "main"
+        instruction_name: "loop_select_fusion"
+      }
+      events {
+        kind: FREE
+        buffer_id: 13
+        computation_name: "main"
+        instruction_name: "loop_select_fusion"
+      }
+      whole_module_simulation: true
+      buffer_allocation_index: 6
+    }
+    peak_buffers {
+      logical_buffer_ids: 10
+    }
+    peak_buffers {
+      index: 1
+      logical_buffer_ids: 11
+    }
+    peak_buffers {
+      index: 2
+      logical_buffer_ids: 14
+    }
+    peak_buffers {
+      index: 3
+      logical_buffer_ids: 2
+    }
+    peak_buffers {
+      index: 4
+      logical_buffer_ids: 0
+    }
+    peak_buffers {
+      index: 5
+      logical_buffer_ids: 1
+    }
+    peak_buffers {
+      index: 6
+      logical_buffer_ids: 13
+    }
+  }
+  binary: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\000\005\000\000\000\000\000\000\004Z\000\006@\0008\000\004\000@\000\016\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.nv.shared.reserved.0\000.nv.global\000.nv.constant4\000\000\000\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant4\000\000\000\000.nv.callgraph\000.nv.prototype\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.nv.global\000.nv.constant4\000\000\000\000.debug_frame\000.rel.nv.constant.pic\000.rela.nv.constant.pic\000.nv.callgraph\000.nv.prototype\000buffer_for_c1_1\000buffer_for_c0_1\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\214\000\000\000 \240\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\255\000\000\000\003\000\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\270\000\000\000\003\000\013\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\001\000\000\003\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\035\001\000\000\021\000\r\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000-\001\000\000\021\000\r\000\000\001\000\000\000\000\000\000@\000\000\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\002\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 13.2, V13.2.51\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills  -arch sm_90a \000\000\000\014\000\000\000\010\000\000\000\350\003\000\000NVIDIA Corp\000\002\000Z\000\204\000\000\000\003_\001\001\002\t\001\000\002\002\001\000\003\007\001\001\002\003\000\000\004\013\010\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377\000\000\000\000\010\000\000\000\000\000\000\000\002\000\000\000\t\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\000\343\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0001\001\000\000\000\000\000\000=\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000p\002\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\217\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000`\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000`\003\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\024\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\005\000\000\000\010\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0004\004\000\000\000\000\000\000\004\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0008\004\000\000\000\000\000\000\034\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\307\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\004\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\261\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\004\000\000\000\000\000\0000\000\000\000\000\000\000\000\003\000\000\000\013\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000~\000\000\000\001\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\250\004\000\000\000\000\000\000\020\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\270\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000s\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\005\000\000\000\000\000\000@\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\340\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\250\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\020\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\270\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000H\000\000\000\000\000\000\000\210\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000"
+  gpu_compute_capability {
+    cuda_compute_capability {
+      major: 9
+      feature_extension: ACCELERATED_FEATURES
+    }
+  }
+  module_name: "test"
+  program_shape {
+    result {
+      element_type: U32
+      dimensions: 4
+      layout {
+        minor_to_major: 0
+        tail_padding_alignment_in_elements: 1
+      }
+      is_dynamic_dimension: false
+    }
+  }
+  buffer_allocations {
+    values {
+      size: 32
+      is_constant: true
+    }
+    values {
+      index: 1
+      size: 32
+      is_constant: true
+    }
+    values {
+      index: 2
+      size: 16
+      maybe_live_out: true
+    }
+    values {
+      index: 3
+      size: 4
+      is_thread_local: true
+    }
+    values {
+      index: 4
+      size: 4
+      is_thread_local: true
+    }
+    values {
+      index: 5
+      size: 4
+      is_thread_local: true
+    }
+    values {
+      index: 6
+      size: 32
+    }
+  }
+  output_info_map {
+    shape_index {
+    }
+    output_info {
+      allocation_index: 2
+    }
+  }
+  constants {
+    symbol_name: "buffer_for_c1_1"
+    allocation_index: 1
+  }
+  constants {
+    symbol_name: "buffer_for_c0_1"
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "id.1"
+      thunk_id: 1
+    }
+    replica_id_thunk {
+      dest_buffer {
+        size: 4
+        buffer_allocation_index: 2
+        element_type: U32
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "loop_select_fusion"
+      thunk_id: 2
+    }
+    custom_kernel_thunk {
+      args {
+        slice {
+          size: 32
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          dimensions: 8
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+      }
+      args {
+        slice {
+          size: 32
+          buffer_allocation_index: 1
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          dimensions: 8
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+      }
+      args {
+        slice {
+          size: 4
+          buffer_allocation_index: 2
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          layout {
+            tail_padding_alignment_in_elements: 1
+          }
+        }
+      }
+      args {
+        slice {
+          size: 32
+          buffer_allocation_index: 6
+          element_type: U32
+        }
+        shape {
+          element_type: U32
+          dimensions: 8
+          layout {
+            minor_to_major: 0
+            tail_padding_alignment_in_elements: 1
+          }
+          is_dynamic_dimension: false
+        }
+      }
+      written: false
+      written: false
+      written: false
+      written: true
+      custom_kernel {
+        name: "loop_select_fusion"
+        kernel_spec {
+          cubin {
+            data: "\177ELF\002\001\001A\010\000\000\000\000\000\000\000\002\000\276\000\001\000\000\000\000\000\000\000\000\000\000\000p\016\000\000\000\000\000\000\260\n\000\000\000\000\000\000\004Z\000\006@\0008\000\005\000@\000\017\000\001\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.shared.reserved.0\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000.shstrtab\000.strtab\000.symtab\000.symtab_shndx\000.note.nv.tkinfo\000.note.nv.cuinfo\000.nv.info\000.nv.compat\000.text.loop_select_fusion\000.nv.info.loop_select_fusion\000.nv.shared.loop_select_fusion\000.nv.reservedSmem.offset0\000.nv.shared.reserved.0\000__nv_reservedSMEM_offset_0_alias\000.debug_frame\000.rel.debug_frame\000.rela.debug_frame\000.nv.callgraph\000.nv.prototype\000loop_select_fusion\000.nv.constant0.loop_select_fusion\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\003\000\005\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\003\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000]\000\000\000\003\000\014\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\000\000\000!\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\337\000\000\000 \240\r\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\003\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\001\000\000\003\000\n\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000L\001\000\000\022\020\014\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000_\001\000\000\003\000\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\377\377\377\377$\000\000\000\000\000\000\000\377\377\377\377\377\377\377\377\003\000\004|\377\377\377\377\017\014\201\200\200(\000\010\377\201\200(\010\201\200\200(\000\000\000\377\377\377\377,\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\004<\000\000\000\004\004\000\000\000\014\201\200\200(\000\000\000\000\000\014\000\000\000\234\000\000\000\320\007\000\000NVIDIA Corp\000\002\000\000\000\000\000\000\000\001\000\000\000\007\000\000\0006\000\000\000f\000\000\000\000ptxas\000Cuda compilation tools, release 13.2, V13.2.51\000Build system must define TOOLS_VERSION_EXTENDED\000-warn-spills  -arch sm_90a \000\000\000\014\000\000\000\010\000\000\000\350\003\000\000NVIDIA Corp\000\002\000Z\000\204\000\000\000\004/\010\000\010\000\000\000\014\000\000\000\004\021\010\000\010\000\000\000\000\000\000\000\004\022\010\000\010\000\000\000\000\000\000\000\002\t\001\000\002\002\001\000\003\007\001\001\002\003\000\000\004\013\010\000\000\000\000\000\000\000\000\000\004f\004\000\003\000\000\000\0047\004\000\204\000\000\000\004\027\014\000\000\000\000\000\003\000\030\000\010\365!\000\004\027\014\000\000\000\000\000\002\000\020\000\010\365!\000\004\027\014\000\000\000\000\000\001\000\010\000\010\365!\000\004\027\014\000\000\000\000\000\000\000\000\000\010\365!\000\003P\000\000\003\033\377\000\003_\001\001\004\034\004\000\360\000\000\000\004\020\014\000\010\000\000\000\001\000\000\000\001\000\000\000\003\031 \000\004\n\010\000\t\000\000\000\020\002 \000\0046\004\000\010\000\000\000\000\000\000\000\377\377\377\377\000\000\000\000\376\377\377\377\000\000\000\000\375\377\377\377\000\000\000\000\374\377\377\377D\000\000\000\000\000\000\000\002\000\000\000\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\202{\001\377\000\n\000\000\000\010\000\000\000\360\017\000\202{\002\377\000\210\000\000\000\n\000\000\000\"\016\000\271z\004\000\000\202\000\000\000\n\000\000\000\316\017\000\202{\004\377\000\204\000\000\000\n\000\000\000b\016\000\031y\t\000\000\000\000\000\000!\000\000\000\242\016\000.x\000\000\000\000\000\000\000\000\000\000\000\314\017\000\202{\006\377\000\212\000\000\000\n\000\000\000\342\016\000\201y\002\002\004\000\000\000\000\231\036\014\000$\037\000\014r\000\002\377\000\000\000pP\360\003\000\332\017\001\202\013\004\377\000\206\000\000\000\010\000\000\000\260.\000\202\013\005\377\000\207\000\000\000\010\000\000\000\244\016\000%x\004\t\004\000\000\000\004\000\216\007\000\314O\000\201y\005\004\004\000\000\000\000\231\036\014\000\242\016\000%x\006\t\004\000\000\000\006\000\216\007\000\312\217\000\206y\000\006\005\000\000\000\004\031\020\014\000\342O\000My\000\000\000\000\000\000\000\000\200\003\000\352\017\000Gy\374\000\374\377\377\377\377\377\203\003\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\030y\000\000\000\000\000\000\000\000\000\000\000\300\017\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\000\000\000\000\000\000\0003\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\013\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\245\001\000\000\000\000\000\000\200\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\023\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000(\003\000\000\000\000\000\000\360\000\000\000\000\000\000\000\002\000\000\000\n\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000\306\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\004\000\000\000\000\000\000h\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000)\000\000\000\007\000\000\000\000\000\000\002\000\000\000\000\000\000\000\000\000\000\000\000\200\004\000\000\000\000\000\000\264\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0009\000\000\000\007\000\000\000@\000\000\001\000\000\000\000\000\000\000\000\000\000\000\0004\005\000\000\000\000\000\000 \000\000\000\000\000\000\000\005\000\000\000\010\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000I\000\000\000\000\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000T\005\000\000\000\000\000\000$\000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000R\000\000\000\206\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000x\005\000\000\000\000\000\000\034\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000v\000\000\000\000\000\000p@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\224\005\000\000\000\000\000\000\214\000\000\000\000\000\000\000\003\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\366\000\000\000\001\000\000p\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000 \006\000\000\000\000\000\000 \000\000\000\000\000\000\000\003\000\000\000\000\000\000\000\004\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\344\000\000\000\004\000\000\000@\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000@\006\000\000\000\000\000\000\030\000\000\000\000\000\000\000\003\000\000\000\004\000\000\000\010\000\000\000\000\000\000\000\030\000\000\000\000\000\000\000]\000\000\000\001\000\000\000\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\006\000\000\000\000\000\000\000\002\000\000\000\000\000\000\003\000\000\000\010\000\000\000\200\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\260\000\000\000\010\000\000\000\003\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\001\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\022\001\000\000\001\000\000\000B\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\200\010\000\000\000\000\000\0000\002\000\000\000\000\000\000\000\000\000\000\014\000\000\000\004\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\006\000\000\000\004\000\000\000p\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000p\016\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\030\001\000\000\000\000\000\000\030\001\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\005\000\000\000\200\006\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\002\000\000\000\000\000\000\000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\006\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\010\000\000\000\000\000\000\000\001\000\000\000\004\000\000\000\200\010\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\000\0000\002\000\000\000\000\000\0000\002\000\000\000\000\000\000\010\000\000\000\000\000\000\000"
+          }
+          arity: 4
+          kernel_name: "loop_select_fusion"
+          kernel_args_packing_spec {
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 1
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 2
+              }
+            }
+            kernel_arguments {
+              relocations {
+                kind: KIND_BITS64_ABSOLUTE
+                argument_index: 3
+              }
+            }
+          }
+        }
+        block_dims {
+          coordinates {
+            x: 1
+            y: 1
+            z: 1
+          }
+        }
+        thread_dims {
+          coordinates {
+            x: 8
+            y: 1
+            z: 1
+          }
+        }
+      }
+      use_pdl: true
+      tma_metadata {
+      }
+    }
+  }
+  thunks {
+    thunk_info {
+      profile_annotation: "ars.3"
+      thunk_id: 3
+    }
+    reduce_scatter_thunk {
+      collective_config {
+        operand_element_type: U32
+      }
+      buffers {
+        element_count: 8
+        source_buffer {
+          slice {
+            size: 32
+            buffer_allocation_index: 6
+            element_type: U32
+          }
+          shape {
+            element_type: U32
+            dimensions: 8
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+        }
+        destination_buffer {
+          slice {
+            size: 16
+            buffer_allocation_index: 2
+            element_type: U32
+          }
+          shape {
+            element_type: U32
+            dimensions: 4
+            layout {
+              minor_to_major: 0
+              tail_padding_alignment_in_elements: 1
+            }
+            is_dynamic_dimension: false
+          }
+        }
+      }
+      reduction_kind: REDUCTION_KIND_SUM
+    }
+  }
+  executable_abi_version {
+    platform_name: "CUDA"
+    cuda_platform_version {
+      cuda_toolkit_version: "13.2.0"
+      cudnn_version: "9.23.0"
+      cub_version: "3.1.2"
+    }
+  }
+  buffer_allocations_debug_summary: "BufferAssignment OOM Debugging.\nBufferAssignment stats:\n             parameter allocation:         0B\n              constant allocation:        64B\n        maybe_live_out allocation:        16B\n     preallocated temp allocation:        32B\n  preallocated temp fragmentation:         0B (0.00%)\n                 total allocation:       124B\n              total fragmentation:        76B (61.29%)\nPeak buffers:\n\tBuffer 1:\n\t\tSize: 32B\n\t\tXLA Label: fusion\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 2:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 3:\n\t\tSize: 32B\n\t\tXLA Label: constant\n\t\tShape: u32[8]\n\t\t==========================\n\n\tBuffer 4:\n\t\tSize: 16B\n\t\tXLA Label: reduce-scatter\n\t\tShape: u32[4]\n\t\t==========================\n\n\tBuffer 5:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 6:\n\t\tSize: 4B\n\t\tXLA Label: parameter\n\t\tShape: u32[]\n\t\t==========================\n\n\tBuffer 7:\n\t\tSize: 4B\n\t\tXLA Label: add\n\t\tShape: u32[]\n\t\t==========================\n"
+  gpu_topology {
+    platform_version: "Compute Capability 9.0a"
+    num_partitions: 1
+    num_hosts_per_partition: 1
+    num_devices_per_host: 2
+    gpu_target_config {
+      gpu_device_info {
+        threads_per_block_limit: 1024
+        threads_per_warp: 32
+        shared_memory_per_block: 49152
+        shared_memory_per_core: 233472
+        threads_per_core_limit: 2048
+        core_count: 132
+        fpus_per_core: 128
+        block_dim_limit_x: 2147483647
+        block_dim_limit_y: 65535
+        block_dim_limit_z: 65535
+        memory_bandwidth: 3352320000000
+        l2_cache_size: 52428800
+        clock_rate_ghz: 1.98
+        device_memory_size: 84942979072
+        shared_memory_per_block_optin: 232448
+        cuda_compute_capability {
+          major: 9
+          feature_extension: ACCELERATED_FEATURES
+        }
+        registers_per_core_limit: 65536
+        registers_per_block_limit: 65536
+        scalar_unit_description {
+          rate_infos {
+            key: 4
+            value {
+              clock_rate_ghz: 1.98
+              units_per_core: 64
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 8
+            value {
+              clock_rate_ghz: 1.98
+              units_per_core: 64
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 10
+            value {
+              clock_rate_ghz: 1.98
+              units_per_core: 128
+              ops_per_clock: 2
+            }
+          }
+          rate_infos {
+            key: 11
+            value {
+              clock_rate_ghz: 1.98
+              units_per_core: 128
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 12
+            value {
+              clock_rate_ghz: 1.98
+              units_per_core: 64
+              ops_per_clock: 1
+            }
+          }
+          rate_infos {
+            key: 16
+            value {
+              clock_rate_ghz: 1.98
+              units_per_core: 128
+              ops_per_clock: 2
+            }
+          }
+        }
+        matrix_unit_description {
+          rate_infos {
+            key: 2
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 6
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 10
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 512
+            }
+          }
+          rate_infos {
+            key: 11
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 256
+            }
+          }
+          rate_infos {
+            key: 12
+            value {
+              clock_rate_ghz: 1.98
+              units_per_core: 4
+              ops_per_clock: 32
+            }
+          }
+          rate_infos {
+            key: 16
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 512
+            }
+          }
+          rate_infos {
+            key: 19
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 20
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 23
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 24
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 25
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 28
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 29
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 33
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 35
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+          rate_infos {
+            key: 36
+            value {
+              clock_rate_ghz: 1.8295201
+              units_per_core: 4
+              ops_per_clock: 1024
+            }
+          }
+        }
+        driver_version: "13.2.0"
+        kernel_mode_driver_version: "535.183.6"
+        runtime_version: "13.2.0"
+        compile_time_toolkit_version: "13.2.0"
+        dnn_version: "9.23.0"
+        cub_version: "3.1.2"
+        device_interconnect_info {
+          active_links: 18
+        }
+        device_vendor: "NVIDIA Corporation"
+        platform_version: "Compute Capability 9.0a"
+        pci_bus_id: "<undefined>"
+        name: "NVIDIA H100 80GB HBM3"
+        model_str: "sm_9.0a with 84942979072B RAM, 132 cores, 1980000KHz clock, 2619000KHz mem clock, 52428800B L2$"
+        numa_node: -1
+        thread_dim_limit_x: 1024
+        thread_dim_limit_y: 1024
+        thread_dim_limit_z: 64
+        device_address_bits: 64
+        pcie_bandwidth: 34359738368
+        ecc_enabled: true
+        mem_clock_ghz: 2.619
+        reserved_shared_memory_per_block: 1024
+        max_blocks_per_multiprocessor: 32
+        collective_memory_granularity: 2097152
+      }
+      platform_name: "CUDA"
+      device_description_str: "NVIDIA H100 80GB HBM3"
+      runtime_version {
+        major: 13
+        minor: 2
+      }
+    }
+    host_target_machine_options {
+      triple: "x86_64-grtev4-linux-gnu"
+      cpu: "sapphirerapids"
+    }
+    num_devices_per_process: 2
+  }
+}

From b5d3f043120eca3f48770d08f3afc27271ae2d62 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2026 15:13:27 -0700
Subject: [PATCH 0811/1000] Add Python 3.14 PyPI classifier, Bzlmod toolchain,
 and 2.22.0 release notes.

Adds `'Programming Language :: Python :: 3.14'` to `CLASSIFIERS` in `setup.oss.py.tpl`, registers the Python 3.14 toolchain and `requirements_lock_3_14.txt` in `MODULE.bazel`, documents Python 3.14 support under Release 2.22.0 in `RELEASE.md`, and removes a stale comment in `ml_build/Dockerfile`.

Fixes #102890

PiperOrigin-RevId: 999653577
---
 MODULE.bazel                               | 2 ++
 RELEASE.md                                 | 1 +
 ci/official/containers/ml_build/Dockerfile | 1 -
 tensorflow/tools/pip_package/setup.py.tpl  | 1 +
 4 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/MODULE.bazel b/MODULE.bazel
index c83bc7b5055..6a0d7850d68 100644
--- a/MODULE.bazel
+++ b/MODULE.bazel
@@ -219,6 +219,7 @@ python.toolchain(python_version = "3.10")
 python.toolchain(python_version = "3.11")
 python.toolchain(python_version = "3.12")
 python.toolchain(python_version = "3.13")
+python.toolchain(python_version = "3.14")
 
 pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip")
 
@@ -237,6 +238,7 @@ pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip")
     "3.11",
     "3.12",
     "3.13",
+    "3.14",
 ]]
 
 use_repo(pip, "pypi_mods", pypi = "tf_pypi")
diff --git a/RELEASE.md b/RELEASE.md
index 65d8fb8b31e..bc6ed23056f 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -52,6 +52,7 @@ In `tensorflow/c/experimental/filesystem/filesystem_interface.h`, removed `TF_Tr
 
 ### Major Features and Improvements
 
+*   Support for Python 3.14 has been added.
 * `tf.lite`
     * Adds support for QUI4 (Quantized Unsigned 4-bit) in Dequantize operator.
     * Adds support for FP16 and BF16 in Unpack operator.
diff --git a/ci/official/containers/ml_build/Dockerfile b/ci/official/containers/ml_build/Dockerfile
index 56f02d226cf..4fa9d259688 100644
--- a/ci/official/containers/ml_build/Dockerfile
+++ b/ci/official/containers/ml_build/Dockerfile
@@ -71,7 +71,6 @@ RUN if [ "${INSTALL_CMAKE_TOOLCHAIN}" = "true" ]; then \
 # Setup links for TensorFlow to compile.
 # Referenced in devel.usertools/*.bazelrc.
 # Set python3.12 as the default python version.
-# TF does not support python3.13.
 RUN ln -sf /usr/bin/python3.12 /usr/bin/python3
 RUN ln -sf /usr/bin/python3.12 /usr/bin/python
 RUN ln -sf /usr/lib/python3.12 /usr/lib/tf_python
diff --git a/tensorflow/tools/pip_package/setup.py.tpl b/tensorflow/tools/pip_package/setup.py.tpl
index 5bfb14eabe4..df84b257242 100644
--- a/tensorflow/tools/pip_package/setup.py.tpl
+++ b/tensorflow/tools/pip_package/setup.py.tpl
@@ -466,6 +466,7 @@ CLASSIFIERS = [
     'Programming Language :: Python :: 3.11',
     'Programming Language :: Python :: 3.12',
     'Programming Language :: Python :: 3.13',
+    'Programming Language :: Python :: 3.14',
     'Programming Language :: Python :: 3 :: Only',
     'Topic :: Scientific/Engineering',
     'Topic :: Scientific/Engineering :: Mathematics',

From 97d83d86acadecd5bdb80bb03459adbdc7b80b28 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2026 15:25:23 -0700
Subject: [PATCH 0812/1000] [XLA] Copy instead of move async results in the
 evaluator

PiperOrigin-RevId: 999661283
---
 .../xla/xla/hlo/evaluator/hlo_evaluator.cc    | 10 +++-----
 .../xla/hlo/evaluator/hlo_evaluator_test.cc   | 25 +++++++++++++++++++
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
index ed3179972b7..ac014cb48ac 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
@@ -3448,9 +3448,8 @@ absl::Status HloEvaluator::HandleAsyncStart(const HloInstruction* async_start) {
         Literal result,
         embedded_evaluator->Evaluate(*async_start->async_wrapped_computation(),
                                      arg_literals));
-    // Move the output value to the index {1} of the output.
-    ABSL_RETURN_IF_ERROR(
-        literal.MoveFrom(std::move(result), /*dest_shape_index=*/{1}));
+    // Copy the output value to the index {1} of the output.
+    ABSL_RETURN_IF_ERROR(literal.CopyFrom(result, /*dest_shape_index=*/{1}));
   }
 
   SetEvaluatedLiteralFor(async_start, std::move(literal));
@@ -3510,9 +3509,8 @@ absl::Status HloEvaluator::HandleAsyncUpdate(
         Literal result,
         embedded_evaluator->Evaluate(*async_update->async_wrapped_computation(),
                                      sub_literals));
-    // Move the output value to the index {1} of the output.
-    ABSL_RETURN_IF_ERROR(
-        literal.MoveFrom(std::move(result), /*dest_shape_index=*/{1}));
+    // Copy the output value to the index {1} of the output.
+    ABSL_RETURN_IF_ERROR(literal.CopyFrom(result, /*dest_shape_index=*/{1}));
 
     // propagate the output subshapes back to previous async ops
     ABSL_RETURN_IF_ERROR(
diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc
index bacb11ead17..fed4b1e043c 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_test.cc
@@ -6888,6 +6888,31 @@ TEST_F(HloEvaluatorTest, AsyncOpsWithLayout) {
   EXPECT_TRUE(LiteralTestUtil::Equal(expected, result));
 }
 
+TEST_F(HloEvaluatorTest, AsyncStartWithTiledLayout) {
+  // The reshape drops the tiling, so the result can't be moved.
+  const absl::string_view hlo_text = R"(
+  HloModule test
+
+  async_computation {
+    p0 = f32[8,128]{1,0:T(8,128)} parameter(0)
+    ROOT reshape = f32[8,128]{1,0:T(8,128)} reshape(p0)
+  }, execution_thread="thread"
+
+  ENTRY entry {
+    x = f32[8,128]{1,0:T(8,128)} parameter(0)
+    async-start = ((f32[8,128]{1,0:T(8,128)}), f32[8,128]{1,0:T(8,128)}) async-start(x), calls=async_computation, async_execution_thread="thread"
+    ROOT async-done = f32[8,128]{1,0:T(8,128)} async-done(async-start)
+  }
+  )";
+  ASSERT_OK_AND_ASSIGN(m_, ParseAndReturnVerifiedModule(hlo_text));
+  Array2D<float> input(8, 128);
+  input.FillIota(1.0f);
+  Literal arg = LiteralUtil::CreateR2FromArray2D<float>(input);
+  ASSERT_OK_AND_ASSIGN(Literal result, HloEvaluator().Evaluate(
+                                           *m_->entry_computation(), {&arg}));
+  EXPECT_TRUE(LiteralTestUtil::Equal(arg, result));
+}
+
 TEST_F(HloEvaluatorTest, LateBindingAsyncStartAndUpdate) {
   const absl::string_view hlo_text = R"(
   HloModule test

From ba721bca5b7ddbc26e5e4b5c25143ee3659e1774 Mon Sep 17 00:00:00 2001
From: Tori Baker <vwbaker@google.com>
Date: Wed, 14 Oct 2026 15:51:18 -0700
Subject: [PATCH 0813/1000] [XLA:GPU] Fuse scaled dots with GemmFusionV2.

Scaled dots now always use GemmFusionV2: the V2 helpers accept kScaledDot in
addition to kDot (scale operands are fused without shape tracking). Update tests that depended on
the legacy fusion names. Other scaled dot tests already exist in gemm_fusion_test.cc

PiperOrigin-RevId: 999678120
---
 .../scaled_dot_online_quantization.hlo        |   4 +-
 ...scaled_dot_online_quantization_bitcast.hlo |   4 +-
 .../triton/tests/scaled_dot_device_test.cc    |   4 +-
 .../backends/gpu/transforms/gemm_fusion.cc    | 160 ++++++++----------
 4 files changed, 78 insertions(+), 94 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization.hlo
index 5619eee8d23..f0f2196d3ed 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization.hlo
@@ -46,12 +46,12 @@ ENTRY %main.4 (a.1: f8e4m3fn[2,256], b.1: f8e4m3fn[256,256], a_scale.1: f8e8m0fn
   ROOT %tuple.3 = (f8e3m4[2,256]{1,0}, f8e8m0fnu[2,8]{1,0}) tuple(%jit_online_quantize_.4, %jit_online_quantize_.5)
 }
 
-// CHECK: %fusion_scaled-dot (parameter
+// CHECK: %[[FUSION_DOT:[^ ]*scaled-dot[^ ]*]] (param
 // CHECK: %fused_reduce (param
 // CHECK: %fused_convert (param
 // CHECK: %fused_convert.1 (param
 // CHECK: ENTRY
-// CHECK: %[[fusion_dot:.*]] = f32[2,256]{1,0} fusion({{.*}}), kind=kCustom, calls=%fusion_scaled-dot
+// CHECK: %[[fusion_dot:.*]] = f32[2,256]{1,0} fusion({{.*}}), kind=kCustom, calls=%[[FUSION_DOT]]
 // CHECK: %[[fusion_reduce:.*]] = f32[2,8]{1,0} fusion(%[[fusion_dot]]), kind=kInput, calls=%fused_reduce
 // CHECK: %[[fusion_convert1:.*]] = f8e3m4[2,256]{1,0} fusion(%[[fusion_dot]], %[[fusion_reduce]]), kind=kLoop, calls=%fused_convert
 // CHECK: %[[fusion_convert2:.*]] = f8e8m0fnu[2,8]{1,0} fusion(%[[fusion_reduce]]), kind=kLoop, calls=%fused_convert.1
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization_bitcast.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization_bitcast.hlo
index 20fd0b6c248..ff21e7c54c4 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization_bitcast.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization_bitcast.hlo
@@ -59,12 +59,12 @@ ENTRY %main.5 (a.1: f8e4m3fn[2,256], b.1: f8e4m3fn[256,256], a_scale.1: f8e8m0fn
   ROOT %tuple.0 = (f8e3m4[2,256]{1,0}, f8e8m0fnu[2,8]{1,0}) tuple(%convert_element_type.1, %convert_element_type.2)
 }
 
-// CHECK: %fusion_scaled-dot (parameter
+// CHECK: %[[FUSION_DOT:[^ ]*scaled-dot[^ ]*]] (param
 // CHECK: %fused_reduce (param
 // CHECK: %fused_convert (param
 // CHECK: %fused_convert.1 (param
 // CHECK: ENTRY
-// CHECK: %[[fusion_dot:.*]] = f32[2,256]{1,0} fusion({{.*}}), kind=kCustom, calls=%fusion_scaled-dot
+// CHECK: %[[fusion_dot:.*]] = f32[2,256]{1,0} fusion({{.*}}), kind=kCustom, calls=%[[FUSION_DOT]]
 // CHECK: %[[fusion_reduce:.*]] = f32[2,8]{1,0} fusion(%[[fusion_dot]]), kind=kInput, calls=%fused_reduce
 // CHECK: %[[fusion_convert1:.*]] = f8e3m4[2,256]{1,0} fusion(%[[fusion_dot]], %[[fusion_reduce]]), kind=kLoop, calls=%fused_convert
 // CHECK: %[[fusion_convert2:.*]] = f8e8m0fnu[2,8]{1,0} fusion(%[[fusion_reduce]]), kind=kLoop, calls=%fused_convert.1
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc
index abf8ca66f8f..ee848028d19 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc
@@ -706,7 +706,7 @@ ENTRY e {
   ASSERT_OK_AND_ASSIGN(auto optimized_module,
                        GetOptimizedModule(kHloTextTemplate));
   constexpr absl::string_view kExpectedOptimizedHLO = R"(
-    CHECK: %fusion
+    CHECK: %{{.*}} (param
     CHECK: %{{.*}} = f8e8m0fnu[3,128,4]{2,1,0} broadcast(%{{.*}}), dimensions={0,1}
     CHECK: %{{.*}} = f8e8m0fnu[3,128,4]{2,1,0} broadcast(%{{.*}}), dimensions={0,1}
     CHECK: ROOT {{.*}} scaled-dot
@@ -894,7 +894,7 @@ ENTRY e {
   ASSERT_OK_AND_ASSIGN(auto optimized_module,
                        GetOptimizedModule(kHloTextTemplate));
   constexpr absl::string_view kExpectedOptimizedHLO = R"(
-    CHECK: %[[fusion_name:.*]] (parameter
+    CHECK: %[[fusion_name:.*]] (param
     CHECK: %[[scaled_dot:.*]] = bf16[3,128,128]{2,1,0} scaled-dot
     CHECK: %[[global_scaler:.*]] = bf16[3,128,128]{2,1,0} broadcast
     CHECK: ROOT %{{.*}} = bf16[3,128,128]{2,1,0} multiply(%[[scaled_dot]], %[[global_scaler]])
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 d28df71b899..5ed108f5149 100644
--- a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc
@@ -718,8 +718,9 @@ bool AllowedInGemmFusion(const HloInstruction& instr) {
     // the moment. Includes kFusion, kReduce, kAllReduce, etc.
     return false;
   }
-  return HloPredicateIsNotOp<HloOpcode::kDot, HloOpcode::kParameter,
-                             HloOpcode::kAllGather, HloOpcode::kAllGatherStart,
+  return HloPredicateIsNotOp<HloOpcode::kDot, HloOpcode::kScaledDot,
+                             HloOpcode::kParameter, HloOpcode::kAllGather,
+                             HloOpcode::kAllGatherStart,
                              HloOpcode::kAllGatherDone>(&instr);
 }
 
@@ -1265,7 +1266,8 @@ FusionDecision CanUnpackS4ParametersInFusion(
   }
 
   if (tiled_computation->roots().size() != 1 ||
-      tiled_computation->roots()[0]->hlo()->opcode() != HloOpcode::kDot) {
+      !HloPredicateIsOp<HloOpcode::kDot, HloOpcode::kScaledDot>(
+          tiled_computation->roots()[0]->hlo())) {
     return FusionDecision::Forbid(
         "Expected tiled computation root to be a single dot instruction.");
   }
@@ -1315,7 +1317,8 @@ FusionDecision CanTile(mlir::MLIRContext& mlir_context,
     // a tile size to be divisible by 3, but all Triton tiles are powers of 2.
     std::optional<HloInstructionAdaptor> dot =
         HloBfsFindIf(fusion.GetRoots(), fusion, [](HloInstructionAdaptor node) {
-          return node.opcode() == HloOpcode::kDot;
+          return HloPredicateIsOp<HloOpcode::kDot, HloOpcode::kScaledDot>(
+              &node.instruction());
         });
     if (!dot.has_value()) {
       return FusionDecision::Forbid("No dot instruction found in the fusion.");
@@ -1668,22 +1671,26 @@ FusionDecision ShouldFuseUserTranspose(const HloInstruction& transpose,
                                        const HloInstruction& original_transpose,
                                        const HloInstruction& fusion,
                                        const ShapeTracker& tracker) {
-  const HloInstruction* dot = hlo_query::FindInstruction(
-      fusion.fused_instructions_computation(), HloOpcode::kDot);
+  const HloInstruction* dot = hlo_query::GetFirstInstructionWithOpcode(
+      *fusion.fused_instructions_computation(),
+      {HloOpcode::kDot, HloOpcode::kScaledDot});
   if (dot == nullptr) {
     return FusionDecision::Forbid("Dot not found in fusion.");
   }
 
-  absl::StatusOr<std::array<DotOperandDims, 2>> dot_dims =
-      DotOperandDims::FromDot(dot);
-  if (!dot_dims.ok()) {
+  // The output of a scaled dot has the same layout as the output of a dot of
+  // its LHS and RHS, so the dims of operands 0 and 1 are sufficient.
+  absl::StatusOr<DotOperandDims> lhs_dims =
+      DotOperandDims::FromDotOperand(dot, 0);
+  absl::StatusOr<DotOperandDims> rhs_dims =
+      DotOperandDims::FromDotOperand(dot, 1);
+  if (!lhs_dims.ok() || !rhs_dims.ok()) {
     return FusionDecision::Forbid("Failed to get dot operand dims.");
   }
-  const auto& [lhs_dims, rhs_dims] = *dot_dims;
 
-  int64_t num_batch = lhs_dims.Rank(DotOperandDims::kBatch);
-  int64_t num_lhs_nc = lhs_dims.Rank(DotOperandDims::kNonContracting);
-  int64_t num_rhs_nc = rhs_dims.Rank(DotOperandDims::kNonContracting);
+  int64_t num_batch = lhs_dims->Rank(DotOperandDims::kBatch);
+  int64_t num_lhs_nc = lhs_dims->Rank(DotOperandDims::kNonContracting);
+  int64_t num_rhs_nc = rhs_dims->Rank(DotOperandDims::kNonContracting);
 
   // Dot result always maps to [batch..., lhs_nc..., rhs_nc...].
   absl::InlinedVector<int64_t, 4> batch_dims(num_batch);
@@ -1910,18 +1917,22 @@ absl::StatusOr<std::variant<Fusion, FusionDecision>> CreateTileableFusion(
     return FusionDecision::Forbid("Cannot tile the dot instruction.");
   }
 
-  // Initialize queue with trackers for dot operands.
-  absl::StatusOr<std::array<DotOperandDims, 2>> dot_dims =
-      DotOperandDims::FromDot(dot);
-  if (!dot_dims.ok()) {
-    return FusionDecision::Forbid("Failed to get dot operand dims.");
+  // Initialize queue with trackers for the LHS and RHS. Scaled dot scales get
+  // no tracker, so transposes and concatenates are not fused into them.
+  std::queue<std::pair<HloInstruction*, std::optional<TrackerInfo>>> queue;
+  for (int64_t i = 0; i < dot->operand_count(); ++i) {
+    HloInstruction* operand = dot->mutable_operand(i);
+    std::optional<TrackerInfo> tracker;
+    if (i < 2) {
+      absl::StatusOr<DotOperandDims> dims =
+          DotOperandDims::FromDotOperand(dot, i);
+      if (!dims.ok()) {
+        return FusionDecision::Forbid("Failed to get dot operand dims.");
+      }
+      tracker = TrackerInfo{ShapeTracker(operand->shape()), i, *dims};
+    }
+    queue.push({operand, std::move(tracker)});
   }
-  std::queue<std::pair<HloInstruction*, std::optional<TrackerInfo>>> queue(
-      {{dot->mutable_operand(0),
-        TrackerInfo{ShapeTracker(dot->operand(0)->shape()), 0, (*dot_dims)[0]}},
-       {dot->mutable_operand(1),
-        TrackerInfo{ShapeTracker(dot->operand(1)->shape()), 1,
-                    (*dot_dims)[1]}}});
 
   // Start with a fusion containing only the dot instruction.
   auto entry = fusion_search_space.entry();
@@ -2013,9 +2024,13 @@ absl::StatusOr<std::variant<Fusion, FusionDecision>> CreateTileableFusion(
                 fusion_computation->Clone(std::string(name)), original_output};
 }
 
+// Creates a GEMM fusion around a dot-like instruction (kDot or kScaledDot)
+// using the GemmFusionV2 algorithm.
 absl::StatusOr<std::variant<Fusion, FusionDecision>> CreateDotFusionV2(
-    HloDotInstruction& dot, const se::DeviceDescription& device_description,
+    HloInstruction& dot, const se::DeviceDescription& device_description,
     absl::string_view name) {
+  CHECK((HloPredicateIsOp<HloOpcode::kDot, HloOpcode::kScaledDot>(&dot)))
+      << dot.ToString();
   VLOG(3) << "Creating dot fusion v2 around dot: " << dot.ToString();
   FusionSearchSpace fusion_search_space(&dot, device_description);
   VLOG(3) << "Found fusion search space: \n"
@@ -2129,8 +2144,7 @@ absl::StatusOr<std::variant<Fusion, FusionDecision>> CreateDotFusion(
 class GemmFusionVisitor : public DfsHloRewriteVisitor {
  public:
   explicit GemmFusionVisitor(const se::DeviceDescription& device_description)
-      : device_description_(device_description),
-        gpu_version_(device_description.gpu_compute_capability()) {}
+      : device_description_(device_description) {}
   // Checks that a dot() should be targeting the triton GEMM emitter;
   // if so - fuses all its compatible inputs and outputs as a new computation
   // and replaces the original dot() with a call to the computation.
@@ -2162,18 +2176,43 @@ class GemmFusionVisitor : public DfsHloRewriteVisitor {
       return absl::OkStatus();
     }
 
-    Fusion fusion = std::get<Fusion>(std::move(fusion_or_decision));
+    return InsertGemmFusion(
+        dot, std::get<Fusion>(std::move(fusion_or_decision)), fusion_name);
+  }
 
+  absl::Status HandleScaledDot(HloInstruction* scaled_dot) override {
+    CHECK_EQ(scaled_dot->opcode(), HloOpcode::kScaledDot);
+    std::string fusion_name = absl::StrCat("gemm_fusion_", scaled_dot->name());
+    ABSL_ASSIGN_OR_RETURN(
+        auto fusion_or_decision,
+        CreateDotFusionV2(*scaled_dot, device_description_,
+                          absl::StrCat(fusion_name, "_computation")));
+    if (std::holds_alternative<FusionDecision>(fusion_or_decision)) {
+      VLOG(3) << "Not fusing: "
+              << std::get<FusionDecision>(fusion_or_decision).Explain();
+      return absl::OkStatus();
+    }
+    return InsertGemmFusion(scaled_dot,
+                            std::get<Fusion>(std::move(fusion_or_decision)),
+                            fusion_name);
+  }
+
+ private:
+  // Wraps `fusion` into a fusion instruction and replaces `fusion.output` in
+  // the original computation with it. `dot_like` is the dot-like instruction
+  // the fusion was created around.
+  absl::Status InsertGemmFusion(HloInstruction* dot_like, Fusion fusion,
+                                absl::string_view fusion_name) {
     HloComputation* computation =
-        dot->GetModule()->AddComputationAndUnifyNamesAndIds(
+        dot_like->GetModule()->AddComputationAndUnifyNamesAndIds(
             std::move(fusion.computation), /*is_entry=*/false);
     HloInstruction* dot_fusion =
-        dot->parent()->AddInstruction(HloInstruction::CreateFusion(
+        dot_like->parent()->AddInstruction(HloInstruction::CreateFusion(
             computation->root_instruction()->shape(),
             HloInstruction::FusionKind::kCustom, fusion.inputs, computation));
-    // Copy the metadata of the `dot` to the newly created `fusion` op. This
+    // Copy the metadata of `dot_like` to the newly created `fusion` op. This
     // is convenient for handling metadata in split-k rewriting subsequently.
-    dot_fusion->set_metadata(dot->metadata());
+    dot_fusion->set_metadata(dot_like->metadata());
     dot_fusion->GetModule()->SetAndUniquifyInstrName(dot_fusion, fusion_name);
 
     ABSL_ASSIGN_OR_RETURN(auto gpu_config,
@@ -2185,7 +2224,7 @@ class GemmFusionVisitor : public DfsHloRewriteVisitor {
 
     HloInstruction* replacement = dot_fusion;
     if (fusion.output->shape() != dot_fusion->shape()) {
-      replacement = dot->parent()->AddInstruction(
+      replacement = dot_like->parent()->AddInstruction(
           HloInstruction::CreateBitcast(fusion.output->shape(), dot_fusion));
     }
 
@@ -2202,62 +2241,7 @@ class GemmFusionVisitor : public DfsHloRewriteVisitor {
     return absl::OkStatus();
   }
 
-  absl::Status HandleScaledDot(HloInstruction* scaled_dot) override {
-    CHECK_EQ(scaled_dot->opcode(), HloOpcode::kScaledDot);
-    HloComputation::Builder builder(
-        absl::StrCat("fusion_", scaled_dot->name()));
-
-    std::vector<HloInstruction*> fusion_inputs;
-
-    std::vector<HlosAndRequirements> hlos_and_reqs;
-    hlos_and_reqs.reserve(scaled_dot->operand_count());
-    ABSL_ASSIGN_OR_RETURN(HlosAndRequirements lhs_hlos_and_reqs,
-                     FuseDotOperand(*scaled_dot, /*operand_index=*/0,
-                                    gpu_version_, builder, fusion_inputs));
-    hlos_and_reqs.push_back(lhs_hlos_and_reqs);
-    ABSL_ASSIGN_OR_RETURN(HlosAndRequirements rhs_hlos_and_reqs,
-                     FuseDotOperand(*scaled_dot, /*operand_index=*/1,
-                                    gpu_version_, builder, fusion_inputs));
-    hlos_and_reqs.push_back(rhs_hlos_and_reqs);
-    ABSL_ASSIGN_OR_RETURN(HlosAndRequirements lhs_scale_hlos_and_reqs,
-                     FuseDotOperand(*scaled_dot, /*operand_index=*/2,
-                                    gpu_version_, builder, fusion_inputs));
-    hlos_and_reqs.push_back(lhs_scale_hlos_and_reqs);
-    ABSL_ASSIGN_OR_RETURN(HlosAndRequirements rhs_scale_hlos_and_reqs,
-                     FuseDotOperand(*scaled_dot, /*operand_index=*/3,
-                                    gpu_version_, builder, fusion_inputs));
-    hlos_and_reqs.push_back(rhs_scale_hlos_and_reqs);
-
-    HloInstruction& fused_dot = FuseDot(*scaled_dot, hlos_and_reqs, builder);
-
-    HlosAndRequirements fused_output_and_reqs =
-        FuseDotOutput(*scaled_dot, fused_dot, gpu_version_,
-                      lhs_hlos_and_reqs.requirements, builder, fusion_inputs);
-    HloComputation* computation =
-        scaled_dot->GetModule()->AddComputationAndUnifyNamesAndIds(
-            builder.Build(),
-            /*is_entry=*/false);
-
-    HloInstruction* fusion =
-        scaled_dot->parent()->AddInstruction(HloInstruction::CreateFusion(
-            computation->root_instruction()->shape(),
-            HloInstruction::FusionKind::kCustom, fusion_inputs, computation));
-
-    ABSL_ASSIGN_OR_RETURN(auto gpu_config,
-                     fusion->backend_config<GpuBackendConfig>());
-    FusionBackendConfig& backend_config =
-        *gpu_config.mutable_fusion_backend_config();
-    backend_config.set_kind(kTritonGemmFusionKind);
-    ABSL_RETURN_IF_ERROR(fusion->set_backend_config(gpu_config));
-    HloInstruction* fusion_output = fused_output_and_reqs.original_hlo;
-    ABSL_RETURN_IF_ERROR(ReplaceInstruction(fusion_output, fusion));
-    MarkAsChanged();
-    return absl::OkStatus();
-  }
-
- private:
   const se::DeviceDescription& device_description_;
-  se::GpuComputeCapability gpu_version_;
 };
 
 absl::StatusOr<bool> RunOnComputation(

From 3fa7838c6bdcf27af978e491a4ada9007e1221ea Mon Sep 17 00:00:00 2001
From: Bill Varcho <varcho@google.com>
Date: Wed, 14 Oct 2026 16:29:03 -0700
Subject: [PATCH 0814/1000] [SDY][perf] Defer SymbolTable construction in
 SdyRoundTripDedupMeshesPass.

PiperOrigin-RevId: 999701751
---
 .../xla/xla/service/spmd/shardy/sdy_round_trip/dedup_meshes.cc  | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/dedup_meshes.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/dedup_meshes.cc
index a2571510a9a..3a6c6fce006 100644
--- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/dedup_meshes.cc
+++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/dedup_meshes.cc
@@ -509,7 +509,6 @@ class SdyRoundTripDedupMeshesPass
 
   void runOnOperation() final {
     ModuleOp moduleOp = getOperation();
-    SymbolTable symbolTable(moduleOp);
 
     // Exit early if there are no meshes or just one mesh.
     auto meshIter = moduleOp.getOps<MeshOp>();
@@ -528,6 +527,7 @@ class SdyRoundTripDedupMeshesPass
       // No meshes to dedup.
       return;
     }
+    SymbolTable symbolTable(moduleOp);
     dedupMeshes(moduleOp, symbolTable, duplicateMeshesToAxisMap);
     eraseMeshes(symbolTable, duplicateMeshesToAxisMap);
   }

From c66b5eed56fc6ef591e1a8b8d302e140f3347bdc Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2026 17:08:33 -0700
Subject: [PATCH 0815/1000] Integrate LLVM at llvm/llvm-project@2b35fb9283a4

Updates LLVM usage to match
[2b35fb9283a4](https://github.com/llvm/llvm-project/commit/2b35fb9283a44aae92cb72a41d077ee8de2be7ad)

PiperOrigin-RevId: 999724350
---
 .../xla/third_party/llvm/workspace.bzl        |    4 +-
 .../xla/third_party/shardy/temporary.patch    | 1282 ++++++++++++++++-
 2 files changed, 1280 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/third_party/llvm/workspace.bzl b/third_party/xla/third_party/llvm/workspace.bzl
index b13ddf8949c..ba37bcd14fb 100644
--- a/third_party/xla/third_party/llvm/workspace.bzl
+++ b/third_party/xla/third_party/llvm/workspace.bzl
@@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "bb80dc2a4f6ffc8262701cb0417c6f08158a17df"
-    LLVM_SHA256 = "27dd5e6295f58415653076a65b6b8d3e3210b09e778b102b086b47ff3b2bcb5e"
+    LLVM_COMMIT = "2b35fb9283a44aae92cb72a41d077ee8de2be7ad"
+    LLVM_SHA256 = "b61c85c97cde1c3e96ea7cfd09737c24578d6d0656c392d79ce402da6904b828"
 
     tf_http_archive(
         name = name,
diff --git a/third_party/xla/third_party/shardy/temporary.patch b/third_party/xla/third_party/shardy/temporary.patch
index 57bc9af6357..95b473b0325 100644
--- a/third_party/xla/third_party/shardy/temporary.patch
+++ b/third_party/xla/third_party/shardy/temporary.patch
@@ -1,5 +1,5 @@
 diff --git a/shardy/dialect/sdy/ir/dialect.cc b/shardy/dialect/sdy/ir/dialect.cc
-index 719686a5..53082307 100644
+index 719686a5..fc0561f4 100644
 --- a/shardy/dialect/sdy/ir/dialect.cc
 +++ b/shardy/dialect/sdy/ir/dialect.cc
 @@ -969,9 +969,7 @@ Type TensorShardingAttr::getLocalType(Type globalType, MeshAttr mesh,
@@ -13,6 +13,34 @@ index 719686a5..53082307 100644
        getLocalType(globalTensorType, mesh, allowNonDivisible));
  }
  
+@@ -1232,16 +1230,17 @@ TensorShardingAttr addFreeAxesToManualComputationSharding(
+   TensorShardingAttr returnedSharding =
+       eraseFreeAxes(outerManualSharding, manualAxes);
+ 
+-  SmallVector<DimensionShardingAttr> resultDimShardings(
+-      returnedSharding.getDimShardings());
+-  for (auto [resultDimSharding, newDimSharding] :
+-       llvm::zip(resultDimShardings, newSharding.getDimShardings())) {
+-    resultDimSharding = DimensionShardingAttr::get(
+-        resultDimSharding.getContext(),
+-        llvm::to_vector(llvm::concat<const AxisRefAttr>(
+-            resultDimSharding.getAxes(), newDimSharding.getAxes())),
+-        newDimSharding.getIsClosed(), newDimSharding.getPriority());
+-  }
++  SmallVector<DimensionShardingAttr> resultDimShardings = llvm::map_to_vector(
++      llvm::zip_equal(returnedSharding.getDimShardings(),
++                      newSharding.getDimShardings()),
++      [](auto pair) {
++        auto [returnedDimSharding, newDimSharding] = pair;
++        return DimensionShardingAttr::get(
++            returnedDimSharding.getContext(),
++            llvm::to_vector(llvm::concat<const AxisRefAttr>(
++                returnedDimSharding.getAxes(), newDimSharding.getAxes())),
++            newDimSharding.getIsClosed(), newDimSharding.getPriority());
++      });
+   return TensorShardingAttr::get(newSharding.getContext(),
+                                  newSharding.getMeshOrRef(), resultDimShardings,
+                                  outerManualSharding.getReplicatedAxes(),
 diff --git a/shardy/dialect/sdy/ir/dialect_test.cc b/shardy/dialect/sdy/ir/dialect_test.cc
 index 1ea9160e..8f149153 100644
 --- a/shardy/dialect/sdy/ir/dialect_test.cc
@@ -43,6 +71,1252 @@ index 1ea9160e..8f149153 100644
  }
  
  TEST_F(DialectTest, DimensionShardingAttrGetSharded) {
+diff --git a/shardy/dialect/sdy/transforms/export/convert_global_to_local.cc b/shardy/dialect/sdy/transforms/export/convert_global_to_local.cc
+index c6b1a677..1b545c49 100644
+--- a/shardy/dialect/sdy/transforms/export/convert_global_to_local.cc
++++ b/shardy/dialect/sdy/transforms/export/convert_global_to_local.cc
+@@ -280,19 +280,6 @@ LogicalResult localizeGenericOp(Operation* op, ValueRange operands,
+   return success();
+ }
+ 
+-void checkUnreducedResultSharding(TensorShardingAttr resSharding,
+-                                  ArrayRef<AxisRefAttr> expectedAxes,
+-                                  StringRef opName) {
+-  SDY_CHECK(resSharding && !resSharding.getUnreducedAxes().empty())
+-      << "Expected sharded reduction " << opName.str()
+-      << " result to have unreduced axes sharding.";
+-  SDY_CHECK(resSharding.getUnreducedAxes() == expectedAxes)
+-      << "The unreduced axes of the " << opName.str()
+-      << " result sharding should match the sharded reduction dimensions of "
+-         "the "
+-      << opName.str() << ".";
+-}
+-
+ // Pattern for generic ops that do not require special handling beyond type
+ // conversion, such as StableHLO ops, sdy.all_reduce, etc.
+ class GenericOpPattern : public ConversionPattern {
+@@ -305,16 +292,15 @@ class GenericOpPattern : public ConversionPattern {
+   LogicalResult matchAndRewrite(
+       Operation* op, ArrayRef<Value> operands,
+       ConversionPatternRewriter& rewriter) const override {
+-    Dialect* dialect = op->getDialect();
+-    if ((dialect && dialect->getNamespace() != "stablehlo" &&
++    if ((isa_and_nonnull<SdyDialect, func::FuncDialect>(op->getDialect()) &&
+          !isa<sdy::ReturnOp, func::CallOp, func::ReturnOp>(op)) ||
+         isa<stablehlo::ConcatenateOp, stablehlo::ConvolutionOp,
+             stablehlo::DotGeneralOp, stablehlo::DotOp,
+-            stablehlo::DynamicUpdateSliceOp, stablehlo::GatherOp,
+-            stablehlo::IotaOp, stablehlo::PadOp, stablehlo::ReduceOp,
+-            stablehlo::ReduceWindowOp, stablehlo::RngBitGeneratorOp,
+-            stablehlo::ScatterOp, stablehlo::SelectAndScatterOp,
+-            stablehlo::SliceOp>(op)) {
++            stablehlo::DynamicSliceOp, stablehlo::DynamicUpdateSliceOp,
++            stablehlo::GatherOp, stablehlo::IotaOp, stablehlo::PadOp,
++            stablehlo::ReduceOp, stablehlo::ReduceWindowOp,
++            stablehlo::RngBitGeneratorOp, stablehlo::ScatterOp,
++            stablehlo::SelectAndScatterOp, stablehlo::SliceOp>(op)) {
+       return failure();
+     }
+     return localizeGenericOp(
+@@ -432,23 +418,70 @@ std::pair<Attribute, int64_t> getReplicaGroupsAndSize(
+   return {replicaGroups, groupSize};
+ }
+ 
+-bool hasUnreducedResultSharding(Operation* op,
+-                                ArrayRef<AxisRefAttr> reductionAxes,
+-                                StringRef opName,
+-                                const GlobalToLocalTypeConverter* converter) {
+-  bool hasUnreducedAxes = false;
+-  for (Value res : op->getResults()) {
+-    TensorShardingAttr resSharding = converter->getSharding(res);
+-    if (resSharding && !resSharding.getUnreducedAxes().empty()) {
+-      checkUnreducedResultSharding(resSharding, reductionAxes, opName);
+-      hasUnreducedAxes = true;
+-    } else if (llvm::any_of(res.getUsers(), [](Operation* user) {
+-                 return isa<sdy::AllReduceOp, sdy::ReduceScatterOp>(user);
+-               })) {
+-      hasUnreducedAxes = true;
+-    }
+-  }
+-  return hasUnreducedAxes;
++// Returns whether the op introduces partial results related to the given
++// reductionAxes for its operands with kReduction factors.
++//
++// The op introduces partial results if either:
++//   - The result's sharding has unreduced axes containing reductionAxes, or
++//   - the result is fed directly into an AllReduce or ReduceScatter on
++//     reductionAxes.
++//
++// We assume if an op that can introduce partial result and produce multiple
++// results, all of its results should be sharded the same way. This is the case
++// with the current StableHLO ops.
++bool doesOpIntroducePartialResult(Operation* op,
++                                  ArrayRef<AxisRefAttr> reductionAxes,
++                                  StringRef opName,
++                                  const GlobalToLocalTypeConverter* converter) {
++  auto containsReductionAxes = [&](ArrayRef<AxisRefAttr> axes, MeshAttr mesh) {
++    if (axes.empty()) {
++      return false;
++    }
++    SmallVector<AxisRefAttr> missingReductionAxes =
++        getAxisSetDiff(reductionAxes, axes, mesh);
++    if (missingReductionAxes.empty()) {
++      return true;
++    }
++    // Otherwise, none of the sharded reduction axes of the op should be in
++    // `axes` (the axes only come from operands' unreduced axes).
++    SDY_CHECK(ArrayRef<AxisRefAttr>(missingReductionAxes) == reductionAxes)
++        << "The unreduced or reduction axes for the " << opName.str()
++        << " result must either contain all or none of the sharded "
++           "reduction dimensions of the "
++        << opName.str() << ".";
++    return false;
++  };
++
++  // All results of a reduction op share the same sharding and unreduced axes,
++  // so we only need to inspect the first result.
++  Value res = op->getResult(0);
++  TensorShardingAttr resSharding = converter->getSharding(res);
++  const SymbolTable& symbolTable = converter->getSymbolTable();
++  if (resSharding && containsReductionAxes(resSharding.getUnreducedAxes(),
++                                           resSharding.getMesh(symbolTable))) {
++    return true;
++  }
++
++  // Otherwise, if the op is fed directly into an AllReduce or ReduceScatter on
++  // reductionAxes, we consider it introduce a partial result and the upstream
++  // passes forgot to properly annotate the result with unreduced axes.
++  return llvm::any_of(res.getUsers(), [&](Operation* user) {
++    if (auto allReduce = dyn_cast<sdy::AllReduceOp>(user)) {
++      return containsReductionAxes(
++          allReduce.getReductionAxes(),
++          allReduce.getOutSharding().getMesh(symbolTable));
++    }
++    if (auto reduceScatter = dyn_cast<sdy::ReduceScatterOp>(user)) {
++      SmallVector<AxisRefAttr> reduceScatterAxes;
++      for (AxisRefListAttr axesList : reduceScatter.getReduceScatterAxes()) {
++        llvm::append_range(reduceScatterAxes, axesList.getValue());
++      }
++      return containsReductionAxes(
++          reduceScatterAxes,
++          reduceScatter.getOutSharding().getMesh(symbolTable));
++    }
++    return false;
++  });
+ }
+ 
+ ValueRange createAllReduceForUnreducedAxes(
+@@ -708,6 +741,24 @@ class AllReduceOpPattern : public OpConversionPattern<sdy::AllReduceOp> {
+   bool enableRGV3;
+ };
+ 
++// Looks up a per-device scalar value from a 1D constant `tableAttr` indexed by
++// `deviceId` (or `getDeviceId` if `deviceId` is null).
++Value lookupPerDeviceValue(Location loc, DenseElementsAttr tableAttr,
++                           const ConversionState& conversionState,
++                           ConversionPatternRewriter& rewriter,
++                           Value deviceId = nullptr) {
++  if (!deviceId) {
++    deviceId = getDeviceId(loc, conversionState, rewriter);
++  }
++  Type elemTy = tableAttr.getElementType();
++  auto tableConst = stablehlo::ConstantOp::create(rewriter, loc, tableAttr);
++  auto slice = stablehlo::DynamicSliceOp::create(
++      rewriter, loc, RankedTensorType::get({1}, elemTy), tableConst,
++      ValueRange{deviceId}, rewriter.getDenseI64ArrayAttr({1}));
++  return stablehlo::ReshapeOp::create(rewriter, loc,
++                                      RankedTensorType::get({}, elemTy), slice);
++}
++
+ // Returns a 0-rank i32 tensor containing the global offset for the given shard
+ // axes and local shard size.
+ Value getDimensionOffset(Location loc, MeshAttr mesh,
+@@ -716,9 +767,6 @@ Value getDimensionOffset(Location loc, MeshAttr mesh,
+                          ConversionPatternRewriter& rewriter) {
+   int64_t numDevices = mesh.getTotalSize();
+   Type i32Ty = rewriter.getI32Type();
+-  auto indexTy = RankedTensorType::get({}, i32Ty);
+-
+-  Value deviceId = getDeviceId(loc, conversionState, rewriter);
+ 
+   // Calculate a compile-time offset table for this dimension.
+   SmallVector<int32_t> offsetsTable = llvm::map_to_vector(
+@@ -728,16 +776,11 @@ Value getDimensionOffset(Location loc, MeshAttr mesh,
+         return static_cast<int32_t>(offset);
+       });
+ 
+-  // Create the offset table and look up the value for the current device.
+-  auto tableConst = stablehlo::ConstantOp::create(
+-      rewriter, loc,
++  return lookupPerDeviceValue(
++      loc,
+       DenseIntElementsAttr::get(RankedTensorType::get({numDevices}, i32Ty),
+-                                offsetsTable));
+-  auto offsetSlice = stablehlo::DynamicSliceOp::create(
+-      rewriter, loc, RankedTensorType::get({1}, i32Ty), tableConst,
+-      ValueRange{deviceId}, rewriter.getDenseI64ArrayAttr({1}));
+-
+-  return stablehlo::ReshapeOp::create(rewriter, loc, indexTy, offsetSlice);
++                                offsetsTable),
++      conversionState, rewriter);
+ }
+ 
+ Value emitDynamicSliceForAxes(Location loc, Value globalTensor, MeshAttr mesh,
+@@ -767,6 +810,43 @@ Value emitDynamicSliceForAxes(Location loc, Value globalTensor, MeshAttr mesh,
+       .getResult();
+ }
+ 
++// Masks non-leader shards (where shard index across `axes` != 0) to
++// `nonLeaderVal` (or zero if `nonLeaderVal` is null).
++SmallVector<Value> maskNonLeaderShards(
++    Location loc, ValueRange inputs, MeshAttr mesh, ArrayRef<AxisRefAttr> axes,
++    const ConversionState& conversionState, ConversionPatternRewriter& rewriter,
++    Value nonLeaderVal = nullptr, Value deviceId = nullptr) {
++  int64_t numDevices = mesh.getTotalSize();
++  SmallVector<bool> leaderTable = llvm::map_to_vector(
++      llvm::seq<int64_t>(0, numDevices),
++      [&](int64_t devId) { return getShardIndex(devId, mesh, axes) == 0; });
++  Value isLeader = lookupPerDeviceValue(
++      loc,
++      DenseElementsAttr::get(
++          RankedTensorType::get({numDevices}, rewriter.getI1Type()),
++          leaderTable),
++      conversionState, rewriter, deviceId);
++  if (!nonLeaderVal) {
++    auto localType = cast<RankedTensorType>(inputs.front().getType());
++    nonLeaderVal = createZeroConstant(rewriter, loc, localType);
++  }
++  return llvm::map_to_vector(inputs, [&](Value input) -> Value {
++    return stablehlo::SelectOp::create(rewriter, loc, isLeader, input,
++                                       nonLeaderVal);
++  });
++}
++
++// Masks non-leader shards (where shard index across `axes` != 0) to zero.
++Value maskNonLeaderShardsToZero(Location loc, Value input, MeshAttr mesh,
++                                ArrayRef<AxisRefAttr> axes,
++                                const ConversionState& conversionState,
++                                ConversionPatternRewriter& rewriter,
++                                Value deviceId = nullptr) {
++  return maskNonLeaderShards(loc, input, mesh, axes, conversionState, rewriter,
++                             /*nonLeaderVal=*/nullptr, deviceId)
++      .front();
++}
++
+ class AllSliceOpPattern : public OpConversionPattern<AllSliceOp> {
+  public:
+   AllSliceOpPattern(TypeConverter& converter, MLIRContext* ctx,
+@@ -1119,7 +1199,14 @@ class ConstantOpPattern : public OpConversionPattern<sdy::ConstantOp> {
+     auto localType = cast<RankedTensorType>(converter->convertType(op));
+     ElementsAttr elementsAttr = op.getValue();
+     Location loc = op.getLoc();
++    const SymbolTable& symbolTable = converter->getSymbolTable();
++    TensorShardingAttr sharding = getSharding(op.getResult());
++    MeshAttr mesh = sharding ? sharding.getMesh(symbolTable) : nullptr;
++    if (sharding && !mesh) {
++      return op.emitOpError("failed to resolve mesh for constant");
++    }
+ 
++    Value result;
+     // Unsharded or splat constants.
+     if (localType == globalType || elementsAttr.isSplat()) {
+       Attribute newAttr =
+@@ -1128,33 +1215,35 @@ class ConstantOpPattern : public OpConversionPattern<sdy::ConstantOp> {
+               : SplatElementsAttr::get(localType,
+                                        elementsAttr.getSplatValue<Attribute>());
+ 
+-      rewriter.replaceOp(
+-          op, stablehlo::ConstantOp::create(rewriter, loc, localType,
+-                                            cast<ElementsAttr>(newAttr)));
+-      conversionState.removeToConvertOp(op);
+-      return success();
+-    }
++      result = stablehlo::ConstantOp::create(rewriter, loc, localType,
++                                             cast<ElementsAttr>(newAttr));
++    } else {
++      // Sharded dense constants.
++      auto globalConst = stablehlo::ConstantOp::create(
++          rewriter, loc, globalType, elementsAttr);
+ 
+-    // Sharded dense constants.
+-    const SymbolTable& symbolTable = converter->getSymbolTable();
+-    TensorShardingAttr sharding = getSharding(op.getResult());
+-    MeshAttr mesh = sharding.getMesh(symbolTable);
+-    if (!mesh) {
+-      return op.emitOpError("failed to resolve mesh for constant");
++      SmallVector<AxisRefListAttr> slicingAxesPerDim;
++      slicingAxesPerDim.reserve(localType.getRank());
++      for (DimensionShardingAttr dimSharding : sharding.getDimShardings()) {
++        slicingAxesPerDim.push_back(
++            AxisRefListAttr::get(op.getContext(), dimSharding.getAxes()));
++      }
++
++      result =
++          emitDynamicSliceForAxes(loc, globalConst, mesh, slicingAxesPerDim,
++                                  localType, conversionState, rewriter);
+     }
+-    auto globalConst =
+-        stablehlo::ConstantOp::create(rewriter, loc, globalType, elementsAttr);
+ 
+-    SmallVector<AxisRefListAttr> slicingAxesPerDim;
+-    slicingAxesPerDim.reserve(localType.getRank());
+-    for (DimensionShardingAttr dimSharding : sharding.getDimShardings()) {
+-      slicingAxesPerDim.push_back(
+-          AxisRefListAttr::get(op.getContext(), dimSharding.getAxes()));
++    if (sharding && !sharding.getUnreducedAxes().empty()) {
++      SDY_CHECK(sharding.getReductionOp() == ReductionOp::SUM)
++          << "Only SUM reduction is supported for constants with unreduced "
++             "axes.";
++      result = maskNonLeaderShardsToZero(loc, result, mesh,
++                                         sharding.getUnreducedAxes(),
++                                         conversionState, rewriter);
+     }
+ 
+-    rewriter.replaceOp(
+-        op, emitDynamicSliceForAxes(loc, globalConst, mesh, slicingAxesPerDim,
+-                                    localType, conversionState, rewriter));
++    rewriter.replaceOp(op, result);
+     conversionState.removeToConvertOp(op);
+ 
+     return success();
+@@ -1425,6 +1514,97 @@ class ReduceScatterOpPattern : public OpConversionPattern<ReduceScatterOp> {
+   bool enableRGV3;
+ };
+ 
++class ReplicatedToUnreducedOpPattern
++    : public OpConversionPattern<ReplicatedToUnreducedOp> {
++ public:
++  ReplicatedToUnreducedOpPattern(TypeConverter& converter, MLIRContext* ctx,
++                                 ConversionState& state)
++      : OpConversionPattern<ReplicatedToUnreducedOp>(converter, ctx),
++        conversionState(state) {}
++
++  LogicalResult matchAndRewrite(
++      ReplicatedToUnreducedOp op, OpAdaptor adaptor,
++      ConversionPatternRewriter& rewriter) const override {
++    TensorShardingAttr outSharding = op.getOutSharding();
++    SDY_CHECK(outSharding.getReductionOp() == ReductionOp::SUM)
++        << "Only SUM reduction is supported for ReplicatedToUnreducedOp.";
++    auto* converter =
++        static_cast<const GlobalToLocalTypeConverter*>(getTypeConverter());
++    MeshAttr mesh = outSharding.getMesh(converter->getSymbolTable());
++    if (!mesh) {
++      return op.emitOpError("failed to resolve mesh");
++    }
++
++    Value result =
++        maskNonLeaderShardsToZero(op.getLoc(), adaptor.getTensor(), mesh,
++                                  op.getAxes(), conversionState, rewriter);
++    rewriter.replaceOp(op, result);
++    conversionState.removeToConvertOp(op);
++    return success();
++  }
++
++ private:
++  ConversionState& conversionState;
++};
++
++class ShardedToUnreducedOpPattern
++    : public OpConversionPattern<ShardedToUnreducedOp> {
++ public:
++  ShardedToUnreducedOpPattern(TypeConverter& converter, MLIRContext* ctx,
++                              ConversionState& state)
++      : OpConversionPattern<ShardedToUnreducedOp>(converter, ctx),
++        conversionState(state) {}
++
++  LogicalResult matchAndRewrite(
++      ShardedToUnreducedOp op, OpAdaptor adaptor,
++      ConversionPatternRewriter& rewriter) const override {
++    Location loc = op.getLoc();
++    TensorShardingAttr outSharding = op.getOutSharding();
++    SDY_CHECK(outSharding.getReductionOp() == ReductionOp::SUM)
++        << "Only SUM reduction is supported for ShardedToUnreducedOp.";
++    auto* converter =
++        static_cast<const GlobalToLocalTypeConverter*>(getTypeConverter());
++    MeshAttr mesh = outSharding.getMesh(converter->getSymbolTable());
++    if (!mesh) {
++      return op.emitOpError("failed to resolve mesh");
++    }
++
++    Value localInput = adaptor.getTensor();
++    auto localInputType = cast<RankedTensorType>(localInput.getType());
++    auto localResultType = cast<RankedTensorType>(converter->convertType(op));
++    if (!localResultType.hasStaticShape()) {
++      return op.emitOpError("dynamic shape not supported");
++    }
++    Value zeroTensor = createZeroConstant(rewriter, loc, localResultType);
++
++    auto indexTy = RankedTensorType::get({}, rewriter.getI32Type());
++    Value zeroOffset;
++    SmallVector<Value> startIndices;
++    startIndices.reserve(localResultType.getRank());
++    for (int64_t i = 0; i < localResultType.getRank(); ++i) {
++      ArrayRef<AxisRefAttr> axes = op.getAxes()[i].getValue();
++      if (axes.empty()) {
++        if (!zeroOffset) {
++          zeroOffset = createZeroConstant(rewriter, loc, indexTy);
++        }
++        startIndices.push_back(zeroOffset);
++        continue;
++      }
++      startIndices.push_back(getDimensionOffset(loc, mesh, axes,
++                                                localInputType.getDimSize(i),
++                                                conversionState, rewriter));
++    }
++
++    rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
++        op, localResultType, zeroTensor, localInput, startIndices);
++    conversionState.removeToConvertOp(op);
++    return success();
++  }
++
++ private:
++  ConversionState& conversionState;
++};
++
+ class NamedComputationOpPattern
+     : public OpConversionPattern<NamedComputationOp> {
+  public:
+@@ -1608,11 +1788,11 @@ class StablehloConvolutionOpPattern
+       addRhsReductionAxes(spatialDim);
+     }
+ 
+-    bool hasUnreducedAxes = true;
++    bool introducesPartialResult = true;
+     if (!reductionAxes.empty()) {
+       sortAndMergeAxes(reductionAxes, mesh);
+-      hasUnreducedAxes = hasUnreducedResultSharding(op, reductionAxes,
+-                                                    "convolution", converter);
++      introducesPartialResult = doesOpIntroducePartialResult(
++          op, reductionAxes, "convolution", converter);
+     }
+ 
+     auto getAxes = [&](TensorShardingAttr sharding, int64_t dim) {
+@@ -1643,7 +1823,7 @@ class StablehloConvolutionOpPattern
+         localBatchGroupCount, op.getPrecisionConfigAttr());
+ 
+     ValueRange results = localConv->getResults();
+-    if (!reductionAxes.empty() && !hasUnreducedAxes) {
++    if (!reductionAxes.empty() && !introducesPartialResult) {
+       Attribute meshOrRef =
+           rhsSharding ? rhsSharding.getMeshOrRef() : lhsSharding.getMeshOrRef();
+       results = createAllReduceForUnreducedAxes(
+@@ -1727,8 +1907,8 @@ class StablehloDotOpPattern : public OpConversionPattern<stablehlo::DotOp> {
+ 
+         // Verify that an AllReduce has been inserted by a previous pass unless
+         // the result is explicitly unreduced.
+-        if (!hasUnreducedResultSharding(op, contractingAxes, "dot",
+-                                        converter)) {
++        if (!doesOpIntroducePartialResult(op, contractingAxes, "dot",
++                                          converter)) {
+           Operation* localDot = createLocalGenericOp(op, adaptor.getOperands(),
+                                                      rewriter, converter);
+           ValueRange results = createAllReduceForUnreducedAxes(
+@@ -1804,8 +1984,8 @@ class StablehloDotGeneralOpPattern
+ 
+         // Verify that an AllReduce has been inserted by a previous pass unless
+         // the result is explicitly unreduced.
+-        if (!hasUnreducedResultSharding(op, reductionAxes, "dot_general",
+-                                        converter)) {
++        if (!doesOpIntroducePartialResult(op, reductionAxes, "dot_general",
++                                          converter)) {
+           Operation* localDotGeneral = createLocalGenericOp(
+               op, adaptor.getOperands(), rewriter, converter);
+           ValueRange results = createAllReduceForUnreducedAxes(
+@@ -1828,6 +2008,71 @@ class StablehloDotGeneralOpPattern
+   bool enableRGV3;
+ };
+ 
++class StablehloDynamicSliceOpPattern
++    : public OpConversionPattern<stablehlo::DynamicSliceOp> {
++ public:
++  StablehloDynamicSliceOpPattern(TypeConverter& converter, MLIRContext* ctx,
++                                 ConversionState& state)
++      : OpConversionPattern<stablehlo::DynamicSliceOp>(converter, ctx),
++        conversionState(state) {}
++
++  LogicalResult matchAndRewrite(
++      stablehlo::DynamicSliceOp op, OpAdaptor adaptor,
++      ConversionPatternRewriter& rewriter) const override {
++    auto* converter =
++        static_cast<const GlobalToLocalTypeConverter*>(getTypeConverter());
++    TensorShardingAttr inSharding = converter->getSharding(op.getOperand());
++    TensorShardingAttr outSharding = converter->getSharding(op.getResult());
++
++    if (isFullyReplicated(inSharding) && isFullyReplicated(outSharding)) {
++      return localizeGenericOp(op, adaptor.getOperands(), rewriter, converter,
++                               conversionState);
++    }
++
++    auto globalInputType = op.getOperand().getType();
++    auto localResultType =
++        cast<RankedTensorType>(converter->convertType(op.getResult()));
++    MeshAttr mesh =
++        inSharding ? inSharding.getMesh(converter->getSymbolTable()) : nullptr;
++
++    for (int64_t i = 0; i < globalInputType.getRank(); ++i) {
++      ArrayRef<AxisRefAttr> inAxes =
++          inSharding ? inSharding.getDimShardings()[i].getAxes()
++                     : ArrayRef<AxisRefAttr>();
++      ArrayRef<AxisRefAttr> outAxes =
++          outSharding ? outSharding.getDimShardings()[i].getAxes()
++                      : ArrayRef<AxisRefAttr>();
++
++      if (inAxes != outAxes) {
++        return op.emitOpError()
++               << "dimension " << i << " has mismatched sharding axes ("
++               << strippedAttrsString(inAxes) << " vs "
++               << strippedAttrsString(outAxes)
++               << "). This requires a reshard which should have been handled "
++                  "by previous passes.";
++      }
++
++      int64_t globalSliceSize = op.getSliceSizes()[i];
++      if (!inAxes.empty() && globalSliceSize != globalInputType.getDimSize(i)) {
++        if (!mesh || inSharding.getDimShardings()[i].getShardedSize(mesh) > 1) {
++          return op.emitOpError() << "dimension " << i
++                                  << " is sharded but the slice is not a full "
++                                     "slice and requires device communication.";
++        }
++      }
++    }
++
++    rewriter.replaceOpWithNewOp<stablehlo::DynamicSliceOp>(
++        op, localResultType, adaptor.getOperand(), adaptor.getStartIndices(),
++        localResultType.getShape());
++    conversionState.removeToConvertOp(op);
++    return success();
++  }
++
++ private:
++  ConversionState& conversionState;
++};
++
+ Value createScalarInt(OpBuilder& b, Location loc, RankedTensorType type,
+                       int64_t val) {
+   return stablehlo::ConstantOp::create(
+@@ -2095,7 +2340,7 @@ class StablehloReduceOpPattern
+     }
+ 
+     if (!reductionAxes.empty() &&
+-        !hasUnreducedResultSharding(op, reductionAxes, "reduce", converter)) {
++        !doesOpIntroducePartialResult(op, reductionAxes, "reduce", converter)) {
+       auto localReduceOp = cast<stablehlo::ReduceOp>(
+           createLocalGenericOp(op, adaptor.getOperands(), rewriter, converter));
+       if (op.getNumResults() == 1) {
+@@ -2555,8 +2800,8 @@ class StablehloGatherOpPattern
+         mask, bcastDims);
+ 
+     TensorShardingAttr resSharding = getSharding(op.getResult());
+-    bool hasUnreducedAxes =
+-        hasUnreducedResultSharding(op, reductionAxes, "gather", converter);
++    bool introducesPartialResult =
++        doesOpIntroducePartialResult(op, reductionAxes, "gather", converter);
+     ReductionOp reductionOp =
+         resSharding ? resSharding.getReductionOp() : ReductionOp::SUM;
+ 
+@@ -2573,7 +2818,7 @@ class StablehloGatherOpPattern
+     result = stablehlo::SelectOp::create(rewriter, loc, bcastMask, result,
+                                          identityVal);
+ 
+-    if (!hasUnreducedAxes) {
++    if (!introducesPartialResult) {
+       result =
+           createAllReduceForUnreducedAxes(loc, result, reductionAxes, mesh,
+                                           operandSharding.getMeshOrRef(),
+@@ -2747,30 +2992,8 @@ class StablehloRngBitGeneratorOpPattern
+     // Mask non-leader shards (where shard index across dataAxes != 0) to zero
+     // and all-reduce across dataAxes so all partitions receive device 0's
+     // updated output_state.
+-    SmallVector<bool> leaderTable = llvm::map_to_vector(
+-        llvm::seq<int64_t>(0, numDevices), [&](int64_t devId) {
+-          return getShardIndex(devId, mesh, dataAxes) == 0;
+-        });
+-    auto leaderTableConst = stablehlo::ConstantOp::create(
+-        rewriter, loc,
+-        DenseElementsAttr::get(
+-            RankedTensorType::get({numDevices}, rewriter.getI1Type()),
+-            leaderTable));
+-    Value isLeader = stablehlo::DynamicSliceOp::create(
+-        rewriter, loc, RankedTensorType::get({1}, rewriter.getI1Type()),
+-        leaderTableConst, ValueRange{deviceId},
+-        rewriter.getDenseI64ArrayAttr({1}));
+-    isLeader = stablehlo::ReshapeOp::create(
+-        rewriter, loc, RankedTensorType::get({}, rewriter.getI1Type()),
+-        isLeader);
+-    isLeader = stablehlo::BroadcastInDimOp::create(
+-        rewriter, loc,
+-        RankedTensorType::get(localStateType.getShape(), rewriter.getI1Type()),
+-        isLeader, rewriter.getDenseI64ArrayAttr({}));
+-
+-    Value zeroState = createZeroConstant(rewriter, loc, localStateType);
+-    Value maskedState = stablehlo::SelectOp::create(rewriter, loc, isLeader,
+-                                                    finalState, zeroState);
++    Value maskedState = maskNonLeaderShardsToZero(
++        loc, finalState, mesh, dataAxes, conversionState, rewriter, deviceId);
+ 
+     finalState =
+         createAllReduceForUnreducedAxes(loc, ValueRange{maskedState}, dataAxes,
+@@ -3121,35 +3344,13 @@ class StablehloScatterOpPattern
+       if (!identity) {
+         return op.emitError("failed to get scatter reduction identity.");
+       }
+-      // Generate leader mask: (shardIndex(reductionAxes) == 0).
+-      int64_t numDevices = mesh.getTotalSize();
+-      SmallVector<bool> leaderTable = llvm::map_to_vector(
+-          llvm::seq<int64_t>(0, numDevices), [&](int64_t devId) {
+-            return getShardIndex(devId, mesh, reductionAxes) == 0;
+-          });
+-      Value partitionId = getDeviceId(loc, conversionState, rewriter);
+-      auto tableConst = stablehlo::ConstantOp::create(
+-          rewriter, loc,
+-          DenseElementsAttr::get(
+-              RankedTensorType::get({numDevices}, rewriter.getI1Type()),
+-              leaderTable));
+-      Value isLeader =
+-          stablehlo::DynamicSliceOp::create(
+-              rewriter, loc, RankedTensorType::get({1}, rewriter.getI1Type()),
+-              tableConst, {partitionId}, rewriter.getDenseI64ArrayAttr({1}))
+-              .getResult();
+-
+-      isLeader = stablehlo::ReshapeOp::create(
+-          rewriter, loc, RankedTensorType::get({}, rewriter.getI1Type()),
+-          isLeader);
+       // Broadcast the scalar identity to the shape of the current local input.
+       Value broadcastIdentity = stablehlo::BroadcastInDimOp::create(
+           rewriter, loc, localInputs.front().getType(), identity,
+           rewriter.getDenseI64ArrayAttr({}));
+-      llvm::transform(localInputs, localInputs.begin(), [&](Value input) {
+-        return stablehlo::SelectOp::create(rewriter, loc, isLeader, input,
+-                                           broadcastIdentity);
+-      });
++      localInputs = maskNonLeaderShards(loc, localInputs, mesh, reductionAxes,
++                                        conversionState, rewriter,
++                                        broadcastIdentity);
+     }
+ 
+     auto globalInputType = cast<RankedTensorType>(input.getType());
+@@ -3207,7 +3408,8 @@ class StablehloScatterOpPattern
+ 
+     ValueRange results = localScatter.getResults();
+     if (!reductionAxes.empty() &&
+-        !hasUnreducedResultSharding(op, reductionAxes, "scatter", converter)) {
++        !doesOpIntroducePartialResult(op, reductionAxes, "scatter",
++                                      converter)) {
+       Attribute meshOrRef = inputSharding ? inputSharding.getMeshOrRef()
+                                           : updateSharding.getMeshOrRef();
+       results = createAllReduceForUnreducedAxes(
+@@ -3503,7 +3705,8 @@ struct ConvertGlobalToLocalPass
+     patterns.add<AllSliceOpPattern, CollectivePermuteOpPattern,
+                  ConstantOpPattern, FuncOpSignaturePattern, GenericOpPattern,
+                  ManualComputationOpPattern, NamedComputationOpPattern,
+-                 StablehloConcatenateOpPattern,
++                 ReplicatedToUnreducedOpPattern, ShardedToUnreducedOpPattern,
++                 StablehloConcatenateOpPattern, StablehloDynamicSliceOpPattern,
+                  StablehloDynamicUpdateSliceOpPattern, StablehloIotaOpPattern,
+                  StablehloPadOpPattern,
+                  StablehloWindowedOpPattern<stablehlo::ReduceWindowOp>,
+diff --git a/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc b/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc
+index 4cef1c10..7831b6ca 100644
+--- a/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc
++++ b/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc
+@@ -78,6 +78,20 @@ Value createConstant(OpBuilder& b, Location loc, Type elementType,
+   llvm_unreachable("invalid PaddingValueKind");
+ }
+ 
++// Returns the PaddingValueKind of the given scalar constant, or std::nullopt if
++// not a recognized padding value (e.g. 0 or 1).
++std::optional<PaddingValueKind> getPaddingValueKind(Value paddingValue) {
++  if (matchPattern(paddingValue, m_Zero()) ||
++      matchPattern(paddingValue, m_AnyZeroFloat())) {
++    return PaddingValueKind::kZero;
++  }
++  if (matchPattern(paddingValue, m_One()) ||
++      matchPattern(paddingValue, m_OneFloat())) {
++    return PaddingValueKind::kOne;
++  }
++  return std::nullopt;
++}
++
+ // Helper methods for commonly created ops (stablehlo.slice and stablehlo.pad).
+ 
+ // Slices a tensor to `targetType` along dimensions that need trimming.
+@@ -946,8 +960,9 @@ class StablehloSliceOpPattern : public OpConversionPattern<stablehlo::SliceOp> {
+ 
+ class StablehloPadOpPattern : public OpConversionPattern<stablehlo::PadOp> {
+  public:
+-  explicit StablehloPadOpPattern(TypeConverter& converter, MLIRContext* ctx)
+-      : OpConversionPattern(converter, ctx) {}
++  StablehloPadOpPattern(TypeConverter& converter, MLIRContext* ctx,
++                        PaddingCache& cache)
++      : OpConversionPattern(converter, ctx), cache(cache) {}
+ 
+   LogicalResult matchAndRewrite(
+       stablehlo::PadOp op, OpAdaptor adaptor,
+@@ -1017,11 +1032,20 @@ class StablehloPadOpPattern : public OpConversionPattern<stablehlo::PadOp> {
+     // Copy sharding attribute to the new result.
+     setSharding(newPadOp.getResult(), getSharding(result));
+ 
+-    // TODO(b/537380378): We could parse the padding value from the original pad
+-    // op to cache the padding kind and avoid downstream selects.
++    if (std::optional<PaddingValueKind> paddingKind =
++            getPaddingValueKind(adaptor.getOperands()[1])) {
++      if (input.getType() == op.getOperand().getType() || highPaddingReduced ||
++          cache.getPadding(input) == paddingKind) {
++        cache.setPadding(newPadOp.getResult(), *paddingKind);
++      }
++    }
++
+     rewriter.replaceOp(op, newPadOp.getResult());
+     return success();
+   }
++
++ private:
++  PaddingCache& cache;
+ };
+ 
+ class StablehloDotGeneralOpPattern
+@@ -1406,9 +1430,10 @@ struct PadForDivisibilityPass
+     patterns.add<FuncOpSignaturePattern>(typeConverter, &getContext(),
+                                          paddingCache);
+     patterns.add<GenericOpPattern, StablehloSliceOpPattern,
+-                 StablehloGatherOpPattern>(typeConverter, &getContext());
+-    patterns.add<StablehloPadOpPattern, StablehloReshapeOpPattern>(
++                 StablehloGatherOpPattern, StablehloReshapeOpPattern>(
+         typeConverter, &getContext());
++    patterns.add<StablehloPadOpPattern>(typeConverter, &getContext(),
++                                        paddingCache);
+     // Sharing the padding cache reference across pattern instances is safe from
+     // data races because pattern application within a function is sequential.
+     // TODO(b/553579414): remove the pattern for stablehlo.ConstantOp.
+diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/generic_ops.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/generic_ops.mlir
+index c840e9f1..f11dcdbe 100644
+--- a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/generic_ops.mlir
++++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/generic_ops.mlir
+@@ -40,6 +40,17 @@ func.func @unknown_dialect_op(
+   // CHECK-NEXT:  return %[[ARG0]] : tensor<8xf32>
+   return %arg0 : tensor<16xf32>
+ }
++
++// CHECK-LABEL: func.func @registered_non_stablehlo_dialect_op
++// CHECK-SAME: (%[[ARG0:.*]]: tensor<2x8xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}, {}]>}) -> (tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}]>})
++func.func @registered_non_stablehlo_dialect_op(
++  %arg0: tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}, {}]>})
++  -> (tensor<32xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}]>}) {
++  // CHECK-NEXT: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[ARG0]] {{\[\[}}0, 1]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_2, [{"x"}]>]>} : tensor<2x8xf32> into tensor<16xf32>
++  %0 = tensor.collapse_shape %arg0 [[0, 1]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_2, [{"x"}]>]>} : tensor<4x8xf32> into tensor<32xf32>
++  // CHECK-NEXT: return %[[COLLAPSE]] : tensor<16xf32>
++  return %0 : tensor<32xf32>
++}
+ // CHECK-LABEL: func.func @callee_sharded
+ // CHECK-SAME: (%arg0: tensor<8xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}]>}) -> (tensor<8xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}]>})
+ func.func @callee_sharded(%arg0: tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}]>}) -> (tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}]>}) {
+diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_constant.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_constant.mlir
+index 2e01e015..15a45782 100644
+--- a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_constant.mlir
++++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_constant.mlir
+@@ -45,3 +45,18 @@ func.func @unsharded_dense() -> (tensor<4x2xi32> {sdy.sharding = #sdy.sharding<@
+   // CHECK-NEXT: return %[[CST]] : tensor<4x2xi32>
+   return %0 : tensor<4x2xi32>
+ }
++
++// CHECK-LABEL: func.func @unreduced_splat
++// CHECK-SAME:    -> (tensor<4x4xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {}], unreduced={"x"}>})
++func.func @unreduced_splat() -> (tensor<4x4xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {}], unreduced={"x"}>}) {
++  // CHECK-NEXT: %[[CST:.*]] = stablehlo.constant dense<1.000000e+00> : tensor<4x4xf32>
++  // CHECK-NEXT: %[[PID:.*]] = stablehlo.partition_id : tensor<ui32>
++  // CHECK-NEXT: %[[TABLE:.*]] = stablehlo.constant dense<[true, true, true, true, false, false, false, false]> : tensor<8xi1>
++  // CHECK-NEXT: %[[SLICE:.*]] = stablehlo.dynamic_slice %[[TABLE]], %[[PID]], sizes = [1] : (tensor<8xi1>, tensor<ui32>) -> tensor<1xi1>
++  // CHECK-NEXT: %[[SCALAR:.*]] = stablehlo.reshape %[[SLICE]] : (tensor<1xi1>) -> tensor<i1>
++  // CHECK-NEXT: %[[ZERO:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<4x4xf32>
++  // CHECK-NEXT: %[[RES:.*]] = stablehlo.select %[[SCALAR]], %[[CST]], %[[ZERO]] : tensor<i1>, tensor<4x4xf32>
++  %0 = sdy.constant {sdy.sharding = #sdy.sharding_per_value<[<@mesh_2_4, [{}, {}], unreduced={"x"}>]>} dense<1.0> : tensor<4x4xf32>
++  // CHECK-NEXT: return %[[RES]] : tensor<4x4xf32>
++  return %0 : tensor<4x4xf32>
++}
+diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_replicated_to_unreduced.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_replicated_to_unreduced.mlir
+new file mode 100644
+index 00000000..045684b3
+--- /dev/null
++++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_replicated_to_unreduced.mlir
+@@ -0,0 +1,55 @@
++// RUN: sdy_opt %s -sdy-convert-global-to-local | FileCheck %s
++
++// CHECK: sdy.mesh @mesh_2_4 = <["x"=2, "y"=4]>
++sdy.mesh @mesh_2_4 = <["x"=2, "y"=4]>
++
++// CHECK-LABEL: func.func @partial_mesh_major_axis
++// CHECK-SAME: (%[[ARG0:.*]]: tensor<4x4xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {"y"}]>})
++// CHECK-SAME: -> (tensor<4x4xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {"y"}], unreduced={"x"}>})
++func.func @partial_mesh_major_axis(
++    %arg0: tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {"y"}]>})
++    -> (tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {"y"}], unreduced={"x"}>}) {
++  // CHECK-NEXT: %[[PID:.*]] = stablehlo.partition_id : tensor<ui32>
++  // CHECK-NEXT: %[[TABLE:.*]] = stablehlo.constant dense<[true, true, true, true, false, false, false, false]> : tensor<8xi1>
++  // CHECK-NEXT: %[[SLICE:.*]] = stablehlo.dynamic_slice %[[TABLE]], %[[PID]], sizes = [1] : (tensor<8xi1>, tensor<ui32>) -> tensor<1xi1>
++  // CHECK-NEXT: %[[SCALAR:.*]] = stablehlo.reshape %[[SLICE]] : (tensor<1xi1>) -> tensor<i1>
++  // CHECK-NEXT: %[[ZERO:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<4x4xf32>
++  // CHECK-NEXT: %[[RES:.*]] = stablehlo.select %[[SCALAR]], %[[ARG0]], %[[ZERO]] : tensor<i1>, tensor<4x4xf32>
++  %0 = sdy.replicated_to_unreduced {"x"} %arg0 out_sharding=<@mesh_2_4, [{}, {"y"}], unreduced={"x"}> : tensor<4x16xf32>
++  // CHECK-NEXT: return %[[RES]] : tensor<4x4xf32>
++  return %0 : tensor<4x16xf32>
++}
++
++// CHECK-LABEL: func.func @partial_mesh_minor_axis
++// CHECK-SAME: (%[[ARG0:.*]]: tensor<2x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}]>})
++// CHECK-SAME: -> (tensor<2x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>})
++func.func @partial_mesh_minor_axis(
++    %arg0: tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}]>})
++    -> (tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>}) {
++  // CHECK-NEXT: %[[PID:.*]] = stablehlo.partition_id : tensor<ui32>
++  // CHECK-NEXT: %[[TABLE:.*]] = stablehlo.constant dense<[true, false, false, false, true, false, false, false]> : tensor<8xi1>
++  // CHECK-NEXT: %[[SLICE:.*]] = stablehlo.dynamic_slice %[[TABLE]], %[[PID]], sizes = [1] : (tensor<8xi1>, tensor<ui32>) -> tensor<1xi1>
++  // CHECK-NEXT: %[[SCALAR:.*]] = stablehlo.reshape %[[SLICE]] : (tensor<1xi1>) -> tensor<i1>
++  // CHECK-NEXT: %[[ZERO:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<2x16xf32>
++  // CHECK-NEXT: %[[RES:.*]] = stablehlo.select %[[SCALAR]], %[[ARG0]], %[[ZERO]] : tensor<i1>, tensor<2x16xf32>
++  %0 = sdy.replicated_to_unreduced {"y"} %arg0 out_sharding=<@mesh_2_4, [{"x"}, {}], unreduced={"y"}> : tensor<4x16xf32>
++  // CHECK-NEXT: return %[[RES]] : tensor<2x16xf32>
++  return %0 : tensor<4x16xf32>
++}
++
++// CHECK-LABEL: func.func @full_mesh_axes
++// CHECK-SAME: (%[[ARG0:.*]]: tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {}]>})
++// CHECK-SAME: -> (tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {}], unreduced={"x", "y"}>})
++func.func @full_mesh_axes(
++    %arg0: tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {}]>})
++    -> (tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {}], unreduced={"x", "y"}>}) {
++  // CHECK-NEXT: %[[PID:.*]] = stablehlo.partition_id : tensor<ui32>
++  // CHECK-NEXT: %[[TABLE:.*]] = stablehlo.constant dense<[true, false, false, false, false, false, false, false]> : tensor<8xi1>
++  // CHECK-NEXT: %[[SLICE:.*]] = stablehlo.dynamic_slice %[[TABLE]], %[[PID]], sizes = [1] : (tensor<8xi1>, tensor<ui32>) -> tensor<1xi1>
++  // CHECK-NEXT: %[[SCALAR:.*]] = stablehlo.reshape %[[SLICE]] : (tensor<1xi1>) -> tensor<i1>
++  // CHECK-NEXT: %[[ZERO:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<4x16xf32>
++  // CHECK-NEXT: %[[RES:.*]] = stablehlo.select %[[SCALAR]], %[[ARG0]], %[[ZERO]] : tensor<i1>, tensor<4x16xf32>
++  %0 = sdy.replicated_to_unreduced {"x", "y"} %arg0 out_sharding=<@mesh_2_4, [{}, {}], unreduced={"x", "y"}> : tensor<4x16xf32>
++  // CHECK-NEXT: return %[[RES]] : tensor<4x16xf32>
++  return %0 : tensor<4x16xf32>
++}
+diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_sharded_to_unreduced.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_sharded_to_unreduced.mlir
+new file mode 100644
+index 00000000..eaf6ba99
+--- /dev/null
++++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_sharded_to_unreduced.mlir
+@@ -0,0 +1,22 @@
++// RUN: sdy_opt %s -sdy-convert-global-to-local | FileCheck %s
++
++// CHECK: sdy.mesh @mesh_2_4 = <["x"=2, "y"=4]>
++sdy.mesh @mesh_2_4 = <["x"=2, "y"=4]>
++
++// CHECK-LABEL: func.func @sharded_to_unreduced_single_dim
++// CHECK-SAME: (%[[ARG0:.*]]: tensor<4x4xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {"y"}]>})
++// CHECK-SAME: -> (tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>})
++func.func @sharded_to_unreduced_single_dim(
++    %arg0: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {"y"}]>})
++    -> (tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>}) {
++  // CHECK-NEXT: %[[ZERO:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<4x16xf32>
++  // CHECK-NEXT: %[[OFF0:.*]] = stablehlo.constant dense<0> : tensor<i32>
++  // CHECK-NEXT: %[[PID:.*]] = stablehlo.partition_id : tensor<ui32>
++  // CHECK-NEXT: %[[TABLE:.*]] = stablehlo.constant dense<[0, 4, 8, 12, 0, 4, 8, 12]> : tensor<8xi32>
++  // CHECK-NEXT: %[[DS:.*]] = stablehlo.dynamic_slice %[[TABLE]], %[[PID]], sizes = [1] : (tensor<8xi32>, tensor<ui32>) -> tensor<1xi32>
++  // CHECK-NEXT: %[[OFF1:.*]] = stablehlo.reshape %[[DS]] : (tensor<1xi32>) -> tensor<i32>
++  // CHECK-NEXT: %[[RES:.*]] = stablehlo.dynamic_update_slice %[[ZERO]], %[[ARG0]], %[[OFF0]], %[[OFF1]] : (tensor<4x16xf32>, tensor<4x4xf32>, tensor<i32>, tensor<i32>) -> tensor<4x16xf32>
++  %0 = sdy.sharded_to_unreduced [{}, {"y"}] %arg0 out_sharding=<@mesh_2_4, [{"x"}, {}], unreduced={"y"}> : tensor<8x16xf32>
++  // CHECK-NEXT: return %[[RES]] : tensor<4x16xf32>
++  return %0 : tensor<8x16xf32>
++}
+diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_dynamic_slice.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_dynamic_slice.mlir
+new file mode 100644
+index 00000000..a5aa58eb
+--- /dev/null
++++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_dynamic_slice.mlir
+@@ -0,0 +1,69 @@
++// RUN: sdy_opt %s -sdy-convert-global-to-local -split-input-file -verify-diagnostics | FileCheck %s
++
++sdy.mesh @mesh_4_2 = <["x"=4, "y"=2]>
++sdy.mesh @mesh_1_2 = <["x"=1, "y"=2]>
++
++// CHECK-LABEL: func @replicated
++// CHECK-SAME:    (%[[ARG0:.*]]: tensor<16x8xf32>, %[[IDX0:.*]]: tensor<i32>, %[[IDX1:.*]]: tensor<i32>) -> tensor<4x8xf32>
++func.func @replicated(%arg0: tensor<16x8xf32>, %arg1: tensor<i32>, %arg2: tensor<i32>) -> tensor<4x8xf32> {
++  // CHECK-NEXT: %[[RES:.*]] = stablehlo.dynamic_slice %[[ARG0]], %[[IDX0]], %[[IDX1]], sizes = [4, 8] : (tensor<16x8xf32>, tensor<i32>, tensor<i32>) -> tensor<4x8xf32>
++  %0 = stablehlo.dynamic_slice %arg0, %arg1, %arg2, sizes = [4, 8] : (tensor<16x8xf32>, tensor<i32>, tensor<i32>) -> tensor<4x8xf32>
++  // CHECK-NEXT: return %[[RES]] : tensor<4x8xf32>
++  return %0 : tensor<4x8xf32>
++}
++
++// CHECK-LABEL: func @slicing_dim_not_sharded
++// CHECK-SAME:    (%[[ARG0:.*]]: tensor<16x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{}, {"x"}]>},
++// CHECK-SAME:     %[[IDX0:.*]]: tensor<i32>, %[[IDX1:.*]]: tensor<i32>)
++// CHECK-SAME:    -> (tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{}, {"x"}]>})
++func.func @slicing_dim_not_sharded(
++    %arg0: tensor<16x32xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{}, {"x"}]>},
++    %arg1: tensor<i32>, %arg2: tensor<i32>)
++    -> (tensor<4x32xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{}, {"x"}]>}) {
++  // CHECK-NEXT: %[[RES:.*]] = stablehlo.dynamic_slice %[[ARG0]], %[[IDX0]], %[[IDX1]], sizes = [4, 8] : (tensor<16x8xf32>, tensor<i32>, tensor<i32>) -> tensor<4x8xf32>
++  %0 = stablehlo.dynamic_slice %arg0, %arg1, %arg2, sizes = [4, 32] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4_2, [{}, {"x"}]>]>} : (tensor<16x32xf32>, tensor<i32>, tensor<i32>) -> tensor<4x32xf32>
++  // CHECK-NEXT: return %[[RES]] : tensor<4x8xf32>
++  return %0 : tensor<4x32xf32>
++}
++
++// CHECK-LABEL: func @slicing_dim_sharded_on_size_one_axis
++// CHECK-SAME:    (%[[ARG0:.*]]: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1_2, [{"x"}, {"y"}]>},
++// CHECK-SAME:     %[[IDX0:.*]]: tensor<i32>, %[[IDX1:.*]]: tensor<i32>)
++// CHECK-SAME:    -> (tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1_2, [{"x"}, {"y"}]>})
++func.func @slicing_dim_sharded_on_size_one_axis(
++    %arg0: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_1_2, [{"x"}, {"y"}]>},
++    %arg1: tensor<i32>, %arg2: tensor<i32>)
++    -> (tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_1_2, [{"x"}, {"y"}]>}) {
++  // CHECK-NEXT: %[[RES:.*]] = stablehlo.dynamic_slice %[[ARG0]], %[[IDX0]], %[[IDX1]], sizes = [4, 8] : (tensor<8x8xf32>, tensor<i32>, tensor<i32>) -> tensor<4x8xf32>
++  %0 = stablehlo.dynamic_slice %arg0, %arg1, %arg2, sizes = [4, 16] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_1_2, [{"x"}, {"y"}]>]>} : (tensor<8x16xf32>, tensor<i32>, tensor<i32>) -> tensor<4x16xf32>
++  // CHECK-NEXT: return %[[RES]] : tensor<4x8xf32>
++  return %0 : tensor<4x16xf32>
++}
++
++// -----
++
++sdy.mesh @mesh_4_2 = <["x"=4, "y"=2]>
++
++func.func @mismatched_sharding_axes(
++    %arg0: tensor<16x32xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{}, {"x"}]>},
++    %arg1: tensor<i32>, %arg2: tensor<i32>)
++    -> (tensor<4x32xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{}, {"y"}]>}) {
++  // expected-error @+2 {{'stablehlo.dynamic_slice' op dimension 1 has mismatched sharding axes ([axis_ref"x"] vs [axis_ref"y"]). This requires a reshard which should have been handled by previous passes.}}
++  // expected-error @+1 {{failed to legalize operation 'stablehlo.dynamic_slice' that was explicitly marked illegal}}
++  %0 = stablehlo.dynamic_slice %arg0, %arg1, %arg2, sizes = [4, 32] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4_2, [{}, {"y"}]>]>} : (tensor<16x32xf32>, tensor<i32>, tensor<i32>) -> tensor<4x32xf32>
++  return %0 : tensor<4x32xf32>
++}
++
++// -----
++
++sdy.mesh @mesh_4_2 = <["x"=4, "y"=2]>
++
++func.func @slicing_dim_sharded_non_full_slice(
++    %arg0: tensor<16x32xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{"x"}, {}]>},
++    %arg1: tensor<i32>, %arg2: tensor<i32>)
++    -> (tensor<4x32xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{"x"}, {}]>}) {
++  // expected-error @+2 {{'stablehlo.dynamic_slice' op dimension 0 is sharded but the slice is not a full slice and requires device communication.}}
++  // expected-error @+1 {{failed to legalize operation 'stablehlo.dynamic_slice' that was explicitly marked illegal}}
++  %0 = stablehlo.dynamic_slice %arg0, %arg1, %arg2, sizes = [4, 32] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4_2, [{"x"}, {}]>]>} : (tensor<16x32xf32>, tensor<i32>, tensor<i32>) -> tensor<4x32xf32>
++  return %0 : tensor<4x32xf32>
++}
+diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_reduce.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_reduce.mlir
+index 527e6631..4c162c29 100644
+--- a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_reduce.mlir
++++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_reduce.mlir
+@@ -86,3 +86,77 @@ func.func @stablehlo_reduce_multi_result_sharded(
+     }
+   return %0#0, %0#1 : tensor<8xf32>, tensor<8xi32>
+ }
++
++// CHECK-LABEL: func.func @stablehlo_reduce_unreduced_operand_and_sharded_reduction_dim
++// CHECK-SAME: (%[[ARG0:.*]]: tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>})
++// CHECK-SAME: -> (tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}], unreduced={"x", "y"}>})
++func.func @stablehlo_reduce_unreduced_operand_and_sharded_reduction_dim(
++    %arg0: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>})
++    -> (tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}], unreduced={"x", "y"}>}) {
++  // CHECK-NEXT: %[[INIT:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
++  // CHECK-NEXT: %[[LOCAL_RED:.*]] = stablehlo.reduce(%[[ARG0]] init: %[[INIT]]) applies stablehlo.add across dimensions = [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_2_4, [{}], unreduced={"x", "y"}>]>} : (tensor<4x16xf32>, tensor<f32>) -> tensor<16xf32>
++  // CHECK-NEXT: return %[[LOCAL_RED]] : tensor<16xf32>
++  %init = stablehlo.constant dense<0.0> : tensor<f32>
++  %0 = stablehlo.reduce(%arg0 init: %init) across dimensions = [0]
++      {sdy.sharding = #sdy.sharding_per_value<[<@mesh_2_4, [{}], unreduced={"x", "y"}>]>}
++      : (tensor<8x16xf32>, tensor<f32>) -> tensor<16xf32>
++    reducer(%lhs: tensor<f32>, %rhs: tensor<f32>) {
++      %add = stablehlo.add %lhs, %rhs : tensor<f32>
++      stablehlo.return %add : tensor<f32>
++    }
++  return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: func.func @stablehlo_reduce_multiple_unreduced_operands_and_sharded_reduction_dim
++// CHECK-SAME: (%[[ARG0:.*]]: tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>},
++// CHECK-SAME:  %[[ARG1:.*]]: tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>})
++// CHECK-SAME: -> (tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}], unreduced={"x", "y"}>},
++// CHECK-SAME:     tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}], unreduced={"x", "y"}>})
++func.func @stablehlo_reduce_multiple_unreduced_operands_and_sharded_reduction_dim(
++    %arg0: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>},
++    %arg1: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>})
++    -> (tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}], unreduced={"x", "y"}>},
++        tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}], unreduced={"x", "y"}>}) {
++  // CHECK-NEXT: %[[INIT:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
++  // CHECK-NEXT: %[[LOCAL_RED:.*]]:2 = stablehlo.reduce(%[[ARG0]] init: %[[INIT]]), (%[[ARG1]] init: %[[INIT]]) across dimensions = [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_2_4, [{}], unreduced={"x", "y"}>, <@mesh_2_4, [{}], unreduced={"x", "y"}>]>} : (tensor<4x16xf32>, tensor<4x16xf32>, tensor<f32>, tensor<f32>) -> (tensor<16xf32>, tensor<16xf32>)
++  // CHECK:      return %[[LOCAL_RED]]#0, %[[LOCAL_RED]]#1 : tensor<16xf32>, tensor<16xf32>
++  %init = stablehlo.constant dense<0.0> : tensor<f32>
++  %0:2 = stablehlo.reduce(%arg0 init: %init), (%arg1 init: %init) across dimensions = [0]
++      {sdy.sharding = #sdy.sharding_per_value<[<@mesh_2_4, [{}], unreduced={"x", "y"}>, <@mesh_2_4, [{}], unreduced={"x", "y"}>]>}
++      : (tensor<8x16xf32>, tensor<8x16xf32>, tensor<f32>, tensor<f32>) -> (tensor<16xf32>, tensor<16xf32>)
++    reducer(%lhs0: tensor<f32>, %rhs0: tensor<f32>) (%lhs1: tensor<f32>, %rhs1: tensor<f32>) {
++      %add0 = stablehlo.add %lhs0, %rhs0 : tensor<f32>
++      %add1 = stablehlo.add %lhs1, %rhs1 : tensor<f32>
++      stablehlo.return %add0, %add1 : tensor<f32>, tensor<f32>
++    }
++  return %0#0, %0#1 : tensor<16xf32>, tensor<16xf32>
++}
++
++// CHECK-LABEL: func.func @stablehlo_reduce_unreduced_operand_reduces_newly_introduced_axes
++// CHECK-SAME: (%[[ARG0:.*]]: tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>})
++// CHECK-SAME: -> tensor<16xf32>
++func.func @stablehlo_reduce_unreduced_operand_reduces_newly_introduced_axes(
++    %arg0: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>})
++    -> tensor<16xf32> {
++  // CHECK-NEXT: %[[INIT:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
++  // CHECK-NEXT: %[[LOCAL_RED:.*]] = stablehlo.reduce(%[[ARG0]] init: %[[INIT]]) applies stablehlo.add across dimensions = [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_2_4, [{}], unreduced={"y"}>]>} : (tensor<4x16xf32>, tensor<f32>) -> tensor<16xf32>
++  // CHECK:      %[[ALL_RED_X:.*]] = "stablehlo.all_reduce"(%[[LOCAL_RED]])
++  // CHECK-SAME:   channel_handle = #stablehlo.channel_handle<handle = 4, type = 1>
++  // CHECK-SAME:   replica_groups = #stablehlo.replica_group_mesh_axes<mesh = @mesh_2_4, axes = [#stablehlo.axis_ref<name = "x">]>
++  // CHECK-SAME:   use_global_device_ids
++  // CHECK:      %[[ALL_RED_Y:.*]] = "stablehlo.all_reduce"(%[[ALL_RED_X]])
++  // CHECK-SAME:   channel_handle = #stablehlo.channel_handle<handle = 5, type = 1>
++  // CHECK-SAME:   replica_groups = #stablehlo.replica_group_mesh_axes<mesh = @mesh_2_4, axes = [#stablehlo.axis_ref<name = "y">]>
++  // CHECK-SAME:   use_global_device_ids
++  // CHECK:      return %[[ALL_RED_Y]] : tensor<16xf32>
++  %init = stablehlo.constant dense<0.0> : tensor<f32>
++  %0 = stablehlo.reduce(%arg0 init: %init) across dimensions = [0]
++      {sdy.sharding = #sdy.sharding_per_value<[<@mesh_2_4, [{}], unreduced={"y"}>]>}
++      : (tensor<8x16xf32>, tensor<f32>) -> tensor<16xf32>
++    reducer(%lhs: tensor<f32>, %rhs: tensor<f32>) {
++      %add = stablehlo.add %lhs, %rhs : tensor<f32>
++      stablehlo.return %add : tensor<f32>
++    }
++  %1 = sdy.all_reduce {"y"} %0 out_sharding=<@mesh_2_4, [{}]> : tensor<16xf32>
++  return %1 : tensor<16xf32>
++}
+diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_rng_bit_generator.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_rng_bit_generator.mlir
+index 67ffc38f..54d2a403 100644
+--- a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_rng_bit_generator.mlir
++++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_rng_bit_generator.mlir
+@@ -32,10 +32,9 @@ func.func @sharded_1d_ui64_state(
+   // UNSAFE-NEXT: %[[LOCAL_STATE:.*]], %[[LOCAL_DATA:.*]] = stablehlo.rng_bit_generator %[[ADJ_STATE]], algorithm = DEFAULT : (tensor<2xui64>) -> (tensor<2xui64>, tensor<2x16xf32>)
+   // UNSAFE-NEXT: %[[LEADER_TABLE:.*]] = stablehlo.constant dense<[true, false, false, false]> : tensor<4xi1>
+   // UNSAFE-NEXT: %[[IS_LEADER_SLICE:.*]] = stablehlo.dynamic_slice %[[LEADER_TABLE]], %[[PID]], sizes = [1] : (tensor<4xi1>, tensor<ui32>) -> tensor<1xi1>
+-  // UNSAFE-NEXT: %[[IS_LEADER_SCALAR:.*]] = stablehlo.reshape %[[IS_LEADER_SLICE]] : (tensor<1xi1>) -> tensor<i1>
+-  // UNSAFE-NEXT: %[[IS_LEADER:.*]] = stablehlo.broadcast_in_dim %[[IS_LEADER_SCALAR]], dims = [] : (tensor<i1>) -> tensor<2xi1>
++  // UNSAFE-NEXT: %[[IS_LEADER:.*]] = stablehlo.reshape %[[IS_LEADER_SLICE]] : (tensor<1xi1>) -> tensor<i1>
+   // UNSAFE-NEXT: %[[ZERO_STATE:.*]] = stablehlo.constant dense<0> : tensor<2xui64>
+-  // UNSAFE-NEXT: %[[MASKED_STATE:.*]] = stablehlo.select %[[IS_LEADER]], %[[LOCAL_STATE]], %[[ZERO_STATE]] : tensor<2xi1>, tensor<2xui64>
++  // UNSAFE-NEXT: %[[MASKED_STATE:.*]] = stablehlo.select %[[IS_LEADER]], %[[LOCAL_STATE]], %[[ZERO_STATE]] : tensor<i1>, tensor<2xui64>
+   // UNSAFE-NEXT: %[[FINAL_STATE:.*]] = "stablehlo.all_reduce"(%[[MASKED_STATE]])
+   // UNSAFE-SAME{LITERAL}: <{channel_handle = #stablehlo.channel_handle<handle = 1, type = 1>, replica_groups = #stablehlo.replica_group_mesh_axes<mesh = @mesh_4, axes = [#stablehlo.axis_ref<name = "x">]>, use_global_device_ids}> ({
+   // UNSAFE-NEXT: ^bb0(%[[ARG0:.*]]: tensor<ui64>, %[[ARG1:.*]]: tensor<ui64>):
+@@ -76,10 +75,9 @@ func.func @sharded_partial_mesh_ui32_state(
+   // UNSAFE-NEXT: %[[LOCAL_STATE:.*]], %[[LOCAL_DATA:.*]] = stablehlo.rng_bit_generator %[[ADJ_STATE]], algorithm = THREE_FRY : (tensor<4xui32>) -> (tensor<4xui32>, tensor<8x8xf32>)
+   // UNSAFE-NEXT: %[[LEADER_TABLE:.*]] = stablehlo.constant dense<[true, false, true, false]> : tensor<4xi1>
+   // UNSAFE-NEXT: %[[IS_LEADER_SLICE:.*]] = stablehlo.dynamic_slice %[[LEADER_TABLE]], %[[PID]], sizes = [1] : (tensor<4xi1>, tensor<ui32>) -> tensor<1xi1>
+-  // UNSAFE-NEXT: %[[IS_LEADER_SCALAR:.*]] = stablehlo.reshape %[[IS_LEADER_SLICE]] : (tensor<1xi1>) -> tensor<i1>
+-  // UNSAFE-NEXT: %[[IS_LEADER:.*]] = stablehlo.broadcast_in_dim %[[IS_LEADER_SCALAR]], dims = [] : (tensor<i1>) -> tensor<4xi1>
++  // UNSAFE-NEXT: %[[IS_LEADER:.*]] = stablehlo.reshape %[[IS_LEADER_SLICE]] : (tensor<1xi1>) -> tensor<i1>
+   // UNSAFE-NEXT: %[[ZERO_STATE:.*]] = stablehlo.constant dense<0> : tensor<4xui32>
+-  // UNSAFE-NEXT: %[[MASKED_STATE:.*]] = stablehlo.select %[[IS_LEADER]], %[[LOCAL_STATE]], %[[ZERO_STATE]] : tensor<4xi1>, tensor<4xui32>
++  // UNSAFE-NEXT: %[[MASKED_STATE:.*]] = stablehlo.select %[[IS_LEADER]], %[[LOCAL_STATE]], %[[ZERO_STATE]] : tensor<i1>, tensor<4xui32>
+   // UNSAFE-NEXT: %[[FINAL_STATE:.*]] = "stablehlo.all_reduce"(%[[MASKED_STATE]])
+   // UNSAFE-SAME{LITERAL}: <{channel_handle = #stablehlo.channel_handle<handle = 2, type = 1>, replica_groups = #stablehlo.replica_group_mesh_axes<mesh = @mesh_2_2, axes = [#stablehlo.axis_ref<name = "y">]>, use_global_device_ids}> ({
+   // UNSAFE-NEXT: ^bb0(%[[ARG0:.*]]: tensor<ui32>, %[[ARG1:.*]]: tensor<ui32>):
+diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_scatter.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_scatter.mlir
+index e45919c6..3661799c 100644
+--- a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_scatter.mlir
++++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_scatter.mlir
+@@ -18,11 +18,11 @@ func.func @input_not_sharded_scatter_indices_update_sharded_on_implicit_batch_di
+     %arg2: tensor<2x3x1xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}, {}]>})
+  -> tensor<3x4x2xf32> {
+   // CHECK-DAG:  %[[CST:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
++  // CHECK-DAG:  %[[ID_BCAST:.*]] = stablehlo.broadcast_in_dim %[[CST]], dims = []
+   // CHECK-DAG:  %[[PID:.*]] = stablehlo.partition_id : tensor<ui32>
+   // CHECK:      %[[TABLE:.*]] = stablehlo.constant dense<[true, true, true, true, false, false, false, false]> : tensor<8xi1>
+   // CHECK:      %[[IS_LEADER_SLICE:.*]] = stablehlo.dynamic_slice %[[TABLE]], %[[PID]], sizes = [1]
+   // CHECK:      %[[IS_LEADER:.*]] = stablehlo.reshape %[[IS_LEADER_SLICE]] : (tensor<1xi1>) -> tensor<i1>
+-  // CHECK:      %[[ID_BCAST:.*]] = stablehlo.broadcast_in_dim %[[CST]], dims = []
+   // CHECK:      %[[INPUT_SEL:.*]] = stablehlo.select %[[IS_LEADER]], %[[ARG0]], %[[ID_BCAST]]
+   // CHECK:      %[[SCATTER:.*]] = "stablehlo.scatter"(%[[INPUT_SEL]], %[[ARG1]], %[[ARG2]])
+   // CHECK-SAME: scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [2], inserted_window_dims = [0, 1], scatter_dims_to_operand_dims = [0, 1], index_vector_dim = 2>
+@@ -63,11 +63,11 @@ func.func @input_scalar_scatter_indices_update_sharded_on_implicit_batch_dim(
+     %arg2: tensor<2xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}]>})
+  -> tensor<f32> {
+   // CHECK-DAG:  %[[CST:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
++  // CHECK-DAG:  %[[ID_BCAST:.*]] = stablehlo.broadcast_in_dim %[[CST]], dims = []
+   // CHECK-DAG:  %[[PID:.*]] = stablehlo.partition_id : tensor<ui32>
+   // CHECK:      %[[TABLE:.*]] = stablehlo.constant dense<[true, true, true, true, false, false, false, false]> : tensor<8xi1>
+   // CHECK:      %[[IS_LEADER_SLICE:.*]] = stablehlo.dynamic_slice %[[TABLE]], %[[PID]], sizes = [1]
+   // CHECK:      %[[IS_LEADER:.*]] = stablehlo.reshape %[[IS_LEADER_SLICE]] : (tensor<1xi1>) -> tensor<i1>
+-  // CHECK:      %[[ID_BCAST:.*]] = stablehlo.broadcast_in_dim %[[CST]], dims = []
+   // CHECK:      %[[INPUT_SEL:.*]] = stablehlo.select %[[IS_LEADER]], %[[ARG0]], %[[ID_BCAST]]
+   // CHECK:      %[[SCATTER:.*]] = "stablehlo.scatter"(%[[INPUT_SEL]], %[[ARG1]], %[[ARG2]])
+   // CHECK-SAME: scatter_dimension_numbers = #stablehlo.scatter<index_vector_dim = 1>
+@@ -113,11 +113,11 @@ func.func @input_sharded_not_on_indexed_dim(
+     %arg2: tensor<2x3x1xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}, {}]>})
+  -> (tensor<3x4x2xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {}, {"y":(2)2}]>}) {
+   // CHECK-DAG:  %[[CST:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
++  // CHECK-DAG:  %[[ID_BCAST:.*]] = stablehlo.broadcast_in_dim %[[CST]], dims = []
+   // CHECK-DAG:  %[[PID:.*]] = stablehlo.partition_id : tensor<ui32>
+   // CHECK:      %[[TABLE:.*]] = stablehlo.constant dense<[true, true, true, true, false, false, false, false]> : tensor<8xi1>
+   // CHECK:      %[[IS_LEADER_SLICE:.*]] = stablehlo.dynamic_slice %[[TABLE]], %[[PID]], sizes = [1]
+   // CHECK:      %[[IS_LEADER:.*]] = stablehlo.reshape %[[IS_LEADER_SLICE]] : (tensor<1xi1>) -> tensor<i1>
+-  // CHECK:      %[[ID_BCAST:.*]] = stablehlo.broadcast_in_dim %[[CST]], dims = []
+   // CHECK:      %[[INPUT_SEL:.*]] = stablehlo.select %[[IS_LEADER]], %[[ARG0]], %[[ID_BCAST]]
+   // CHECK:      %[[SCATTER:.*]] = "stablehlo.scatter"(%[[INPUT_SEL]], %[[ARG1]], %[[ARG2]])
+   // CHECK-SAME: scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [2], inserted_window_dims = [0, 1], scatter_dims_to_operand_dims = [0, 1], index_vector_dim = 2>
+diff --git a/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/stablehlo_pad.mlir b/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/stablehlo_pad.mlir
+index 59cc6b86..208aa2c2 100644
+--- a/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/stablehlo_pad.mlir
++++ b/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/stablehlo_pad.mlir
+@@ -56,10 +56,22 @@ func.func @no_op(%arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4, [
+   return %1 : tensor<12x8xf32>
+ }
+ 
++// The pad below is padded for divisibility along dim 0 (9 -> 12) with its own
++// zero padding value, so the trailing rows 9-11 are already zero. The pad
++// pattern registers this zero padding in the padding cache, so the dot_general
++// pattern can contract over dim 0 without inserting an iota/compare/select mask
++// on the lhs. Without caching, the lhs would be masked with a stablehlo.select.
+ // CHECK-LABEL: func @pad_then_dot_general
+ // CHECK-SAME:    %[[ARG0:.*]]: tensor<7x7xf32>, %[[ARG1:.*]]: tensor<9x5xf32>
+-// CHECK:         stablehlo.select
+-// CHECK:         stablehlo.dot_general
++// CHECK-DAG:     %[[CST0:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
++// CHECK-DAG:     %[[PAD0:.*]] = stablehlo.pad %[[ARG0]], %[[CST0]], low = [1, 0], high = [4, 0], interior = [0, 0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4, [{"x"}, {}]>]>} : (tensor<7x7xf32>, tensor<f32>) -> tensor<12x7xf32>
++// CHECK-DAG:     %[[CST1:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
++// CHECK-DAG:     %[[PAD1:.*]] = stablehlo.pad %[[ARG1]], %[[CST1]], low = [0, 0], high = [3, 0], interior = [0, 0] : (tensor<9x5xf32>, tensor<f32>) -> tensor<12x5xf32>
++// CHECK-DAG:     %[[SLICE:.*]] = sdy.all_slice [{"x"}, {}] %[[PAD1]] out_sharding=<@mesh_4, [{"x"}, {}]> : tensor<12x5xf32>
++// CHECK-NOT:     stablehlo.select
++// The dot_general consumes %[[PAD0]] directly, i.e. no mask is applied.
++// CHECK:         %[[DOT:.*]] = stablehlo.dot_general %[[PAD0]], %[[SLICE]], contracting_dims = [0] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4, [{}, {}]>]>} : (tensor<12x7xf32>, tensor<12x5xf32>) -> tensor<7x5xf32>
++// CHECK-NEXT:    return %[[DOT]] : tensor<7x5xf32>
+ func.func @pad_then_dot_general(%arg0: tensor<7x7xf32>, %arg1: tensor<9x5xf32>) -> tensor<7x5xf32> {
+   %cst = stablehlo.constant dense<0.000000e+00> : tensor<f32>
+   %0 = stablehlo.pad %arg0, %cst, low = [1, 0], high = [1, 0], interior = [0, 0] {sdy.sharding = #sdy.sharding_per_value<[#sdy.sharding<@mesh_4, [{"x"}, {}]>]>} : (tensor<7x7xf32>, tensor<f32>) -> tensor<9x7xf32>
+@@ -68,6 +80,48 @@ func.func @pad_then_dot_general(%arg0: tensor<7x7xf32>, %arg1: tensor<9x5xf32>)
+   return %1 : tensor<7x5xf32>
+ }
+ 
++// Same idea for stablehlo.reduce: the pad result is sharded along the
++// reduction dim and padded for divisibility (7 -> 8) with zero. Zero is the
++// identity of add, so the reduce consumes the padded tensor directly instead
++// of masking column 7 with a stablehlo.select.
++// CHECK-LABEL: func @zero_pad_then_reduce_add
++// CHECK-SAME:    %[[ARG0:.*]]: tensor<4x5xf32>
++// CHECK-NEXT:    %[[CST:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
++// CHECK-NEXT:    %[[PAD:.*]] = stablehlo.pad %[[ARG0]], %[[CST]], low = [0, 0], high = [0, 3], interior = [0, 0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4, [{}, {"x"}]>]>} : (tensor<4x5xf32>, tensor<f32>) -> tensor<4x8xf32>
++// CHECK-NEXT:    %[[REDUCE:.*]] = stablehlo.reduce(%[[PAD]] init: %[[CST]]) applies stablehlo.add across dimensions = [1]
++// CHECK-SAME:    : (tensor<4x8xf32>, tensor<f32>) -> tensor<4xf32>
++// CHECK-NEXT:    return %[[REDUCE]] : tensor<4xf32>
++func.func @zero_pad_then_reduce_add(%arg0: tensor<4x5xf32>) -> tensor<4xf32> {
++  %cst = stablehlo.constant dense<0.000000e+00> : tensor<f32>
++  %0 = stablehlo.pad %arg0, %cst, low = [0, 0], high = [0, 2], interior = [0, 0] {sdy.sharding = #sdy.sharding_per_value<[#sdy.sharding<@mesh_4, [{}, {"x"}]>]>} : (tensor<4x5xf32>, tensor<f32>) -> tensor<4x7xf32>
++  %1 = stablehlo.reduce(%0 init: %cst) applies stablehlo.add across dimensions = [1] {sdy.sharding = #sdy.sharding_per_value<[#sdy.sharding<@mesh_4, [{}]>]>} : (tensor<4x7xf32>, tensor<f32>) -> tensor<4xf32>
++  return %1 : tensor<4xf32>
++}
++
++// Negative case: the pad value is 1, so the padded column 7 holds 1, which is
++// not the identity of add. The mask must stay: column 7 (iota >= 7) is replaced
++// with the init value 0 before the reduce. The real padding in columns 5-6
++// (iota < 7) is kept, since those 1s are part of the original tensor.
++// CHECK-LABEL: func @one_pad_then_reduce_add
++// CHECK-SAME:    %[[ARG0:.*]]: tensor<4x5xf32>
++// CHECK-DAG:     %[[ONE:.*]] = stablehlo.constant dense<1.000000e+00> : tensor<f32>
++// CHECK-DAG:     %[[ZERO:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
++// CHECK:         %[[PAD:.*]] = stablehlo.pad %[[ARG0]], %[[ONE]], low = [0, 0], high = [0, 3], interior = [0, 0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4, [{}, {"x"}]>]>} : (tensor<4x5xf32>, tensor<f32>) -> tensor<4x8xf32>
++// CHECK:         %[[IOTA:.*]] = stablehlo.iota dim = 1
++// CHECK:         %[[LIMIT:.*]] = stablehlo.constant dense<7> : tensor<i32>
++// CHECK:         %[[MASK:.*]] = stablehlo.compare  LT, %[[IOTA]], %{{.*}}
++// CHECK:         %[[SELECT:.*]] = stablehlo.select %[[MASK]], %[[PAD]], %{{.*}}
++// CHECK:         %[[REDUCE:.*]] = stablehlo.reduce(%[[SELECT]] init: %[[ZERO]]) applies stablehlo.add across dimensions = [1]
++// CHECK-SAME:    : (tensor<4x8xf32>, tensor<f32>) -> tensor<4xf32>
++// CHECK-NEXT:    return %[[REDUCE]] : tensor<4xf32>
++func.func @one_pad_then_reduce_add(%arg0: tensor<4x5xf32>) -> tensor<4xf32> {
++  %one = stablehlo.constant dense<1.000000e+00> : tensor<f32>
++  %zero = stablehlo.constant dense<0.000000e+00> : tensor<f32>
++  %0 = stablehlo.pad %arg0, %one, low = [0, 0], high = [0, 2], interior = [0, 0] {sdy.sharding = #sdy.sharding_per_value<[#sdy.sharding<@mesh_4, [{}, {"x"}]>]>} : (tensor<4x5xf32>, tensor<f32>) -> tensor<4x7xf32>
++  %1 = stablehlo.reduce(%0 init: %zero) applies stablehlo.add across dimensions = [1] {sdy.sharding = #sdy.sharding_per_value<[#sdy.sharding<@mesh_4, [{}]>]>} : (tensor<4x7xf32>, tensor<f32>) -> tensor<4xf32>
++  return %1 : tensor<4xf32>
++}
++
+ // CHECK-LABEL: func @negative_high_padding(
+ // CHECK-SAME:    %[[ARG0:.*]]: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4, [{"x"}, {}]>}) -> tensor<7x8xf32> {
+ // CHECK-NEXT:    %[[SLICE:.*]] = stablehlo.slice %[[ARG0]] [0:8, 0:8] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4, [{"x"}, {}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32>
+diff --git a/shardy/dialect/sdy/transforms/export/test/resolve_permutation_factors.mlir b/shardy/dialect/sdy/transforms/export/test/resolve_permutation_factors.mlir
+index 162bb3db..b1e2bd3d 100644
+--- a/shardy/dialect/sdy/transforms/export/test/resolve_permutation_factors.mlir
++++ b/shardy/dialect/sdy/transforms/export/test/resolve_permutation_factors.mlir
+@@ -1760,6 +1760,22 @@ func.func @dynamic_update_slice_dynamic_operand_shape(
+   return %0 : tensor<?x8xf32>
+ }
+ 
++// CHECK-LABEL: func @dynamic_update_slice_missing_operand_sharding
++// CHECK-SAME: (%[[ARG0:.*]]: tensor<16x8xf32>, %[[ARG1:.*]]: tensor<2x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>}, %[[ARG2:.*]]: tensor<i32>, %[[ARG3:.*]]: tensor<i32>)
++// CHECK-NEXT: %[[DUS:.*]] = stablehlo.dynamic_update_slice %[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]
++// CHECK-NEXT: return %[[DUS]]
++func.func @dynamic_update_slice_missing_operand_sharding(
++    %arg0: tensor<16x8xf32>,
++    %arg1: tensor<2x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>},
++    %arg2: tensor<i32>,
++    %arg3: tensor<i32>)
++    -> tensor<16x8xf32> {
++  %0 = stablehlo.dynamic_update_slice %arg0, %arg1, %arg2, %arg3
++      : (tensor<16x8xf32>, tensor<2x8xf32>, tensor<i32>, tensor<i32>)
++      -> tensor<16x8xf32>
++  return %0 : tensor<16x8xf32>
++}
++
+ //===----------------------------------------------------------------------===//
+ // stablehlo.rng_bit_generator tests
+ //===----------------------------------------------------------------------===//
+diff --git a/shardy/dialect/sdy/transforms/export/utils.cc b/shardy/dialect/sdy/transforms/export/utils.cc
+index 92bb4bc3..54e5d6f3 100644
+--- a/shardy/dialect/sdy/transforms/export/utils.cc
++++ b/shardy/dialect/sdy/transforms/export/utils.cc
+@@ -77,6 +77,13 @@ SmallVector<AxisRefAttr> getOrderedAxes(ArrayRef<AxisList> axesPerDim) {
+   return result;
+ }
+ 
++// Returns the sharded size of dimension `dim` in `sharding`, or 1 if `sharding`
++// is null.
++int64_t getDimShardedSize(TensorShardingAttr sharding, int64_t dim,
++                          MeshAttr mesh) {
++  return sharding ? sharding.getDimSharding(dim).getShardedSize(mesh) : 1;
++}
++
+ }  // namespace
+ 
+ void alignSubAxesByDecomposition(AxisList& axes,
+@@ -156,7 +163,7 @@ TensorShardingAttr updateSharding(TensorShardingAttr sharding,
+ 
+ bool isCommunicationFreeSliceDim(int64_t dimIdx, stablehlo::SliceOp sliceOp,
+                                  TensorShardingAttr sharding, MeshAttr mesh) {
+-  int64_t shardCount = sharding.getDimShardings()[dimIdx].getShardedSize(mesh);
++  int64_t shardCount = getDimShardedSize(sharding, dimIdx, mesh);
+ 
+   if (shardCount <= 1) {
+     return true;
+@@ -183,7 +190,7 @@ bool isCommunicationFreeSliceDim(int64_t dimIdx, stablehlo::SliceOp sliceOp,
+ 
+ bool isCommunicationFreePadDim(int64_t dimIdx, stablehlo::PadOp padOp,
+                                TensorShardingAttr sharding, MeshAttr mesh) {
+-  int64_t shardCount = sharding.getDimShardings()[dimIdx].getShardedSize(mesh);
++  int64_t shardCount = getDimShardedSize(sharding, dimIdx, mesh);
+   if (shardCount <= 1) {
+     return true;
+   }
+@@ -229,8 +236,7 @@ bool isCommunicationFreeDynamicUpdateSliceDim(
+     return false;
+   }
+ 
+-  int64_t shardCount =
+-      operandSharding.getDimShardings()[dimIdx].getShardedSize(mesh);
++  int64_t shardCount = getDimShardedSize(operandSharding, dimIdx, mesh);
+   // Return true if the dimension is unsharded, not sliced, or a single-element
+   // slice.
+   if (shardCount <= 1 || operandDimSize == updateDimSize ||
+@@ -408,10 +414,8 @@ RankedTensorType getDivisiblePaddedType(RankedTensorType globalType,
+       paddedShape.push_back(dimSize);
+       continue;
+     }
+-    int64_t inShardSize =
+-        inSharding ? inSharding.getDimSharding(dim).getShardedSize(mesh) : 1;
+-    int64_t outShardSize =
+-        outSharding ? outSharding.getDimSharding(dim).getShardedSize(mesh) : 1;
++    int64_t inShardSize = getDimShardedSize(inSharding, dim, mesh);
++    int64_t outShardSize = getDimShardedSize(outSharding, dim, mesh);
+     int64_t divisor = std::lcm(inShardSize, outShardSize);
+     paddedShape.push_back(llvm::alignTo(dimSize, divisor));
+   }
 diff --git a/shardy/dialect/sdy/transforms/import/test/import_pipeline.mlir b/shardy/dialect/sdy/transforms/import/test/import_pipeline.mlir
 index d8200b88..60204b54 100644
 --- a/shardy/dialect/sdy/transforms/import/test/import_pipeline.mlir
@@ -178,7 +1452,7 @@ index c5cd8022..6d9f32c2 100644
 +  return %1 : tensor<8x8xf32>
 +}
 diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
-index e42d7dbd..b13ddf89 100644
+index e42d7dbd..ba37bcd1 100644
 --- a/third_party/llvm/workspace.bzl
 +++ b/third_party/llvm/workspace.bzl
 @@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
@@ -187,8 +1461,8 @@ index e42d7dbd..b13ddf89 100644
      """Imports LLVM."""
 -    LLVM_COMMIT = "2e0d9e12b36e853d9b108c4da09509cd54d1f3f9"
 -    LLVM_SHA256 = "8e4a851780af109d623fa6d6db9b1e62d7126272d17268fed1b149f0e737c1d7"
-+    LLVM_COMMIT = "bb80dc2a4f6ffc8262701cb0417c6f08158a17df"
-+    LLVM_SHA256 = "27dd5e6295f58415653076a65b6b8d3e3210b09e778b102b086b47ff3b2bcb5e"
++    LLVM_COMMIT = "2b35fb9283a44aae92cb72a41d077ee8de2be7ad"
++    LLVM_SHA256 = "b61c85c97cde1c3e96ea7cfd09737c24578d6d0656c392d79ce402da6904b828"
  
      tf_http_archive(
          name = name,

From 7d96cfad22b42caf5e12d340968cc781c5799d72 Mon Sep 17 00:00:00 2001
From: Bixia Zheng <bixia@google.com>
Date: Wed, 14 Oct 2026 17:24:30 -0700
Subject: [PATCH 0816/1000] Register Shardy pointwise sharding rules for
 elementwise MHLO ops.

Generalize CopyShardingRuleOpInterface into PointwiseShardingRuleOpInterface<OpTy> and register pointwise sharding rules for elementwise MHLO operations that remain in the mhlo dialect after hlo-legalize-to-stablehlo (acos, acosh, asin, asinh, atanh, copy, cosh, erf, mulhi, sinh, and stochastic_convert).

PiperOrigin-RevId: 999732267
---
 .../spmd/shardy/extensions/mhlo_extensions.cc | 17 +++++++++--
 .../shardy/test/mhlo_extensions_test.mlir     | 30 +++++++++++++++++++
 2 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/shardy/extensions/mhlo_extensions.cc b/third_party/xla/xla/service/spmd/shardy/extensions/mhlo_extensions.cc
index d54c65b3c73..d3549836751 100644
--- a/third_party/xla/xla/service/spmd/shardy/extensions/mhlo_extensions.cc
+++ b/third_party/xla/xla/service/spmd/shardy/extensions/mhlo_extensions.cc
@@ -49,14 +49,22 @@ using ::mlir::sdy::getTensorShape;
 using ::mlir::sdy::kNullDim;
 using ::mlir::sdy::OpShardingRuleBuilder;
 
-struct CopyShardingRuleOpInterface
+template <typename OpTy>
+struct PointwiseShardingRuleOpInterface
     : public mlir::sdy::ShardingRuleOpInterface::ExternalModel<
-          CopyShardingRuleOpInterface, mhlo::CopyOp> {
+          PointwiseShardingRuleOpInterface<OpTy>, OpTy> {
   mlir::sdy::OpShardingRuleAttr getShardingRule(mlir::Operation* op) const {
     return OpShardingRuleBuilder::buildPointwise(op);
   }
 };
 
+template <typename... OpTys>
+void attachPointwiseShardingRule(mlir::MLIRContext* ctx) {
+  (OpTys::template attachInterface<PointwiseShardingRuleOpInterface<OpTys>>(
+       *ctx),
+   ...);
+}
+
 enum RaggedDotMode {
   // Ragged non-contracting (m): [b,m,k], [g,b,k,n], [b,g] -> [b,m,n].
   kNonContracting,
@@ -443,7 +451,10 @@ struct ScanShardingRuleOpInterface
 
 void registerMhloExtensions(mlir::DialectRegistry& registry) {
   registry.addExtension(+[](mlir::MLIRContext* ctx, mhlo::MhloDialect*) {
-    mhlo::CopyOp::attachInterface<CopyShardingRuleOpInterface>(*ctx);
+    attachPointwiseShardingRule<mhlo::AcosOp, mhlo::AcoshOp, mhlo::AsinOp,
+                                mhlo::AsinhOp, mhlo::AtanhOp, mhlo::CopyOp,
+                                mhlo::CoshOp, mhlo::ErfOp, mhlo::MulhiOp,
+                                mhlo::SinhOp, mhlo::StochasticConvertOp>(ctx);
     mhlo::RaggedDotOp::attachInterface<RaggedDotShardingRuleOpInterface>(*ctx);
     mhlo::ScanOp::attachInterface<ScanShardingRuleOpInterface>(*ctx);
     mhlo::TopKOp::attachInterface<TopKShardingRuleOpInterface>(*ctx);
diff --git a/third_party/xla/xla/service/spmd/shardy/test/mhlo_extensions_test.mlir b/third_party/xla/xla/service/spmd/shardy/test/mhlo_extensions_test.mlir
index ec92448d4b2..012bc49b157 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/mhlo_extensions_test.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/mhlo_extensions_test.mlir
@@ -24,6 +24,36 @@ func.func @copy(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
   return %0 : tensor<16x8xf32>
 }
 
+// CHECK-LABEL: func @pointwise_unary_ops
+func.func @pointwise_unary_ops(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
+  // CHECK-NEXT: mhlo.acos %arg0 {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=16, j=8}>}
+  %0 = mhlo.acos %arg0 : tensor<16x8xf32>
+  // CHECK-NEXT: mhlo.acosh %0 {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=16, j=8}>}
+  %1 = mhlo.acosh %0 : tensor<16x8xf32>
+  // CHECK-NEXT: mhlo.asin %1 {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=16, j=8}>}
+  %2 = mhlo.asin %1 : tensor<16x8xf32>
+  // CHECK-NEXT: mhlo.asinh %2 {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=16, j=8}>}
+  %3 = mhlo.asinh %2 : tensor<16x8xf32>
+  // CHECK-NEXT: mhlo.atanh %3 {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=16, j=8}>}
+  %4 = mhlo.atanh %3 : tensor<16x8xf32>
+  // CHECK-NEXT: mhlo.cosh %4 {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=16, j=8}>}
+  %5 = mhlo.cosh %4 : tensor<16x8xf32>
+  // CHECK-NEXT: mhlo.erf %5 {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=16, j=8}>}
+  %6 = mhlo.erf %5 : tensor<16x8xf32>
+  // CHECK-NEXT: mhlo.sinh %6 {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j])->([i, j]) {i=16, j=8}>}
+  %7 = mhlo.sinh %6 : tensor<16x8xf32>
+  return %7 : tensor<16x8xf32>
+}
+
+// CHECK-LABEL: func @pointwise_binary_ops
+func.func @pointwise_binary_ops(%arg0: tensor<16x8xi32>, %arg1: tensor<16x8xf32>, %arg2: tensor<16x8xui32>) -> (tensor<16x8xi32>, tensor<16x8xf16>) {
+  // CHECK-NEXT: mhlo.mulhi %arg0, %arg0 {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j], [i, j])->([i, j]) {i=16, j=8}>}
+  %0 = mhlo.mulhi %arg0, %arg0 : tensor<16x8xi32>
+  // CHECK-NEXT: "mhlo.stochastic_convert"(%arg1, %arg2) {sdy.sharding_rule = #sdy.op_sharding_rule<([i, j], [i, j])->([i, j]) {i=16, j=8}>}
+  %1 = "mhlo.stochastic_convert"(%arg1, %arg2) : (tensor<16x8xf32>, tensor<16x8xui32>) -> tensor<16x8xf16>
+  return %0, %1 : tensor<16x8xi32>, tensor<16x8xf16>
+}
+
 // CHECK-LABEL: func @ragged_dot_mode_non_contracting
 func.func @ragged_dot_mode_non_contracting(
     %arg0: tensor<16x32x64xf32> {sdy.sharding=#sdy.sharding<@mesh_abcd, [{"a"}, {"b"}, {"c"}]>},

From 3ce16812d55fecc1de183aa3fa555fe6a822a46b Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2026 17:36:33 -0700
Subject: [PATCH 0817/1000] Ensure LaunchAsyncOpKernel and
 FallbackBatchResource invoke DoneCallback on all error paths.

In kernel_runner_utils.cc and fallback_batch_kernel.h, use absl::MakeCleanup so the done callback and deferred-op counters are always invoked even if an op kernel returns early on error.

PiperOrigin-RevId: 999738289
---
 .../runtime/fallback_batch_kernel.h           |  13 +-
 tensorflow/core/tfrt/mlrt/kernel/BUILD        |   3 +
 .../tfrt/mlrt/kernel/kernel_runner_utils.cc   |  62 ++++--
 .../core/tfrt/mlrt/kernel/kernel_test.cc      | 194 ++++++++++++++++++
 4 files changed, 254 insertions(+), 18 deletions(-)

diff --git a/tensorflow/core/runtime_fallback/runtime/fallback_batch_kernel.h b/tensorflow/core/runtime_fallback/runtime/fallback_batch_kernel.h
index ef0cf92e71b..04abc051a03 100644
--- a/tensorflow/core/runtime_fallback/runtime/fallback_batch_kernel.h
+++ b/tensorflow/core/runtime_fallback/runtime/fallback_batch_kernel.h
@@ -145,10 +145,11 @@ void BatchFunctionFallbackKernel<BatchResourceType>::ComputeAsync(
                             : std::nullopt,
                         GetModelName(c));
   RecordBatchParamNumBatchThreads(num_batch_threads_, GetModelName(c));
-  OP_REQUIRES_VALUE(tfrt::ResourceContext * client_graph_resource_context, c,
-                    BatchResourceType::GetClientGraphResourceContext(c));
+  absl::StatusOr<tfrt::ResourceContext*> client_graph_resource_context =
+      BatchResourceType::GetClientGraphResourceContext(c);
+  OP_REQUIRES_OK_ASYNC(c, client_graph_resource_context.status(), done);
   OP_REQUIRES_ASYNC(
-      c, client_graph_resource_context != nullptr,
+      c, *client_graph_resource_context != nullptr,
       absl::FailedPreconditionError("client graph resource context not found"),
       done);
   std::function<
@@ -269,8 +270,10 @@ void BatchFunctionFallbackKernel<BatchResourceType>::ComputeAsync(
     };
   }
 
-  auto br = client_graph_resource_context->GetOrCreateResource<
-      tensorflow::core::RefCountPtr<BatchResourceType>>(shared_name_, creator);
+  auto br = (*client_graph_resource_context)
+                ->GetOrCreateResource<
+                    tensorflow::core::RefCountPtr<BatchResourceType>>(
+                    shared_name_, creator);
   if (!br.ok()) OP_REQUIRES_OK_ASYNC(c, br.status(), done);
   auto expected_name = BatchResourceType::GetBatchFunctionName(batch_function_);
   auto received_name =
diff --git a/tensorflow/core/tfrt/mlrt/kernel/BUILD b/tensorflow/core/tfrt/mlrt/kernel/BUILD
index 6457ebfce14..49d58f176b6 100644
--- a/tensorflow/core/tfrt/mlrt/kernel/BUILD
+++ b/tensorflow/core/tfrt/mlrt/kernel/BUILD
@@ -163,7 +163,10 @@ cc_library(
         "//tensorflow/core/tfrt/utils:fallback_tensor",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/cleanup",
+        "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
         "@xla//xla/tsl/platform:errors",
     ],
 )
diff --git a/tensorflow/core/tfrt/mlrt/kernel/kernel_runner_utils.cc b/tensorflow/core/tfrt/mlrt/kernel/kernel_runner_utils.cc
index 168344320bb..d306df4e8f7 100644
--- a/tensorflow/core/tfrt/mlrt/kernel/kernel_runner_utils.cc
+++ b/tensorflow/core/tfrt/mlrt/kernel/kernel_runner_utils.cc
@@ -14,12 +14,15 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/core/tfrt/mlrt/kernel/kernel_runner_utils.h"
 
+#include <functional>
 #include <memory>
 #include <utility>
 #include <vector>
 
-#include "absl/cleanup/cleanup.h"
 #include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "tensorflow/core/framework/device_base.h"
 #include "tensorflow/core/framework/op_kernel.h"
 #include "tensorflow/core/tfrt/fallback/op_kernel_runner.h"
@@ -35,19 +38,58 @@ void LaunchAsyncOpKernel(
     const tfrt_stub::OpKernelRunState& run_state,
     const OpKernelContext::Params& params, mlrt::RegisterSpan results,
     std::shared_ptr<tensorflow::DeviceBase> custom_device) {
+  struct DeferredOpsGuard {
+    explicit DeferredOpsGuard(const OpKernelContext::Params& params)
+        : dec_fn(params.dec_num_deferred_ops_function) {
+      if (params.inc_num_deferred_ops_function) {
+        params.inc_num_deferred_ops_function();
+      }
+    }
+    DeferredOpsGuard(const DeferredOpsGuard&) = delete;
+    DeferredOpsGuard& operator=(const DeferredOpsGuard&) = delete;
+    ~DeferredOpsGuard() {
+      if (dec_fn) {
+        dec_fn();
+      }
+    }
+    std::function<void()> dec_fn;
+  };
+
   struct AsyncState {
     explicit AsyncState(const tfrt_stub::OpKernelRunState& rs,
                         const OpKernelContext::Params& params, int num_outputs,
                         std::shared_ptr<tensorflow::DeviceBase> device)
-        : run_state(rs.input_tf_tensor_values, params, device.get()),
+        : deferred_ops_guard(params),
+          run_state(rs.input_tf_tensor_values, params, device.get()),
           context(&run_state.params, num_outputs),
           custom_device(std::move(device)) {}
 
+    ~AsyncState() {
+      if (!done_called) {
+        LOG(ERROR) << "AsyncOpKernel " << context.op_kernel().name()
+                   << " did not invoke DoneCallback; context status: "
+                   << context.status();
+        absl::Status status =
+            !context.status().ok()
+                ? context.status()
+                : absl::FailedPreconditionError(
+                      absl::StrCat("AsyncOpKernel ", context.op_kernel().name(),
+                                   " did not invoke DoneCallback"));
+        for (auto& result : results) {
+          std::move(result).SetError(status);
+        }
+      }
+    }
+
+    // Declared first so its destructor runs last, after `context` and
+    // `run_state` have been destroyed.
+    DeferredOpsGuard deferred_ops_guard;
     tfrt_stub::OpKernelRunState run_state;
     OpKernelContext context;
 
     std::vector<mlrt::Promise> results;
     std::shared_ptr<tensorflow::DeviceBase> custom_device;
+    bool done_called = false;
   };
 
   DCHECK_EQ(results.size(), kernel_runner.op_kernel()->num_outputs());
@@ -65,17 +107,11 @@ void LaunchAsyncOpKernel(
 
   auto* op_kernel_context_ptr = &async_state->context;
 
-  if (params.inc_num_deferred_ops_function) {
-    params.inc_num_deferred_ops_function();
-  }
-
-  auto done_callback = [async_state = std::move(async_state),
-                        dec_fn = params.dec_num_deferred_ops_function]() {
-    auto cleanup = absl::MakeCleanup([dec_fn] {
-      if (dec_fn) {
-        dec_fn();
-      }
-    });
+  auto done_callback = [async_state = std::move(async_state)]() {
+    if (async_state->done_called) {
+      return;
+    }
+    async_state->done_called = true;
 
     auto& op_kernel_context = async_state->context;
 
diff --git a/tensorflow/core/tfrt/mlrt/kernel/kernel_test.cc b/tensorflow/core/tfrt/mlrt/kernel/kernel_test.cc
index 6ccbc08cd50..88ac3afa217 100644
--- a/tensorflow/core/tfrt/mlrt/kernel/kernel_test.cc
+++ b/tensorflow/core/tfrt/mlrt/kernel/kernel_test.cc
@@ -2716,6 +2716,200 @@ TEST(KernelTest, AsyncWhileOpError) {
       absl_testing::StatusIs(absl::StatusCode::kInternal, "Test error"));
 }
 
+REGISTER_OP("TestAsyncMissingDoneError")
+    .Input("x: T")
+    .Output("y: T")
+    .Attr("T: {int32}")
+    .SetShapeFn(::tensorflow::shape_inference::UnchangedShape);
+
+class TestAsyncMissingDoneErrorKernel : public AsyncOpKernel {
+ public:
+  explicit TestAsyncMissingDoneErrorKernel(OpKernelConstruction* context)
+      : AsyncOpKernel(context) {}
+
+  void ComputeAsync(OpKernelContext* ctx, DoneCallback done) override {
+    // Intentionally return early without invoking `done` to verify
+    // LaunchAsyncOpKernel still pairs inc/dec_num_deferred_ops_function and
+    // fulfills result promises when DoneCallback is destroyed uninvoked.
+    if (ctx->input(0).scalar<int32_t>()() >= 0) {
+      ctx->SetStatus(absl::InvalidArgumentError("missing done error"));
+    }
+  }
+};
+
+REGISTER_KERNEL_BUILDER(Name("TestAsyncMissingDoneError").Device(DEVICE_CPU),
+                        TestAsyncMissingDoneErrorKernel);
+
+TEST(KernelTest, AsyncExecuteOpMissingDoneCallbackStillDecrementsDeferredOps) {
+  mlrt::bc::Buffer buffer;
+  mlrt::bc::Allocator allocator(&buffer);
+
+  auto executable_ctor = mlrt::bc::New<mlrt::bc::Executable>(&allocator);
+  mlrt::testing::SymbolTable kernels;
+  std::vector<std::string> kernel_names = {
+      "tf_mlrt.createop", "tf_mlrt.async_executeop", "tf_mlrt.await", "return"};
+  executable_ctor.construct_kernel_names(kernel_names.size())
+      .Assign(kernel_names);
+  kernels.Def(kernel_names);
+
+  mlrt::testing::AttributeTable attributes(
+      executable_ctor.construct_attributes(2));
+  attributes.Add("node_def_str",
+                 R"pb(name: "TestAsyncMissingDoneError"
+                      op: "TestAsyncMissingDoneError"
+                      input: "dummy_arg"
+                      device: "/job:localhost/replica:0/task:0/device:CPU:0"
+                      attr {
+                        key: "T"
+                        value { type: DT_INT32 }
+                      })pb");
+  attributes.Add("op_key", 0);
+
+  auto functions_ctor = executable_ctor.construct_functions(1);
+  auto function_ctor = functions_ctor.ConstructAt(0);
+  function_ctor.construct_name("main");
+
+  mlrt::testing::SymbolTable regs;
+  function_ctor.construct_input_regs(1).Assign({regs.Def("input")});
+
+  auto kernels_ctor = function_ctor.construct_kernels(4);
+  {
+    auto kernel_ctor = kernels_ctor.ConstructAt(0);
+    kernel_ctor.set_code(kernels.Use("tf_mlrt.createop"));
+    kernel_ctor.construct_attributes(2).Assign(
+        {attributes.GetHandle("node_def_str"), attributes.GetHandle("op_key")});
+  }
+  {
+    auto kernel_ctor = kernels_ctor.ConstructAt(1);
+    kernel_ctor.set_code(kernels.Use("tf_mlrt.async_executeop"));
+    kernel_ctor.construct_arguments(1).Assign({regs.Use("input")});
+    kernel_ctor.construct_results(1).Assign({regs.Def("result_future")});
+    kernel_ctor.construct_attributes(2).Assign(
+        {attributes.GetHandle("node_def_str"), attributes.GetHandle("op_key")});
+    kernel_ctor.construct_last_uses(1).Assign({true});
+  }
+  {
+    auto kernel_ctor = kernels_ctor.ConstructAt(2);
+    kernel_ctor.set_code(kernels.Use("tf_mlrt.await"));
+    kernel_ctor.construct_arguments(1).Assign({regs.Use("result_future")});
+    kernel_ctor.construct_last_uses(1).Assign({true});
+    kernel_ctor.construct_results(1).Assign({regs.Def("result")});
+  }
+  {
+    auto kernel_ctor = kernels_ctor.ConstructAt(3);
+    kernel_ctor.set_code(kernels.Use("return"));
+    kernel_ctor.construct_arguments(1).Assign({regs.Use("result")});
+  }
+  function_ctor.set_num_regs(regs.size());
+  function_ctor.construct_output_regs(1).Assign({regs.Use("result")});
+
+  mlrt::bc::Executable executable(buffer.data());
+  mlrt::KernelRegistry registry;
+  RegisterTfMlrtKernels(registry);
+  mlrt::LoadedExecutable loaded_executable(executable, registry);
+
+  auto work_queue = tfrt::CreateMultiThreadedWorkQueue(
+      /*num_threads=*/4, /*num_blocking_threads=*/4);
+  mlrt::ExecutionContext execution_context(&loaded_executable);
+  execution_context.set_work_queue(work_queue.get());
+
+  tensorflow::SessionOptions session_options;
+  tensorflow::FunctionDefLibrary fdef_lib;
+  TF_ASSERT_OK_AND_ASSIGN(auto fallback_state, tfrt_stub::FallbackState::Create(
+                                                   session_options, fdef_lib));
+
+  std::function<void(std::function<void()>)> runner =
+      [](const std::function<void()>& f) { f(); };
+  tfrt_stub::OpKernelRunnerTable runner_table;
+  tfd::FallbackResourceArray resource_array;
+  tfd::KernelFallbackCompatRequestState fallback_request_state(
+      &runner, &fallback_state->device_manager(), /*step_id=*/0, &runner_table,
+      &resource_array, /*user_intra_op_threadpool=*/nullptr,
+      /*model_metadata=*/std::nullopt,
+      &fallback_state->process_function_library_runtime());
+
+  tfrt::ResourceContext resource_context;
+  auto tf_context =
+      std::make_unique<Context>(&fallback_request_state, &resource_context);
+  int num_deferred_ops = 0;
+  int inc_calls = 0;
+  int dec_calls = 0;
+  tf_context->params().inc_num_deferred_ops_function = [&]() {
+    ++inc_calls;
+    ++num_deferred_ops;
+  };
+  tf_context->params().dec_num_deferred_ops_function = [&]() {
+    ++dec_calls;
+    --num_deferred_ops;
+  };
+  execution_context.AddUserContext(std::move(tf_context));
+
+  tensorflow::Tensor input_tensor(100);
+  mlrt::Value arg(tfrt_stub::FallbackTensor(std::move(input_tensor)));
+  mlrt::Value result;
+
+  absl::Notification notification;
+  execution_context.set_exit_handler(
+      [&notification]() { notification.Notify(); });
+  std::vector<uint8_t> last_uses = {true};
+  execution_context.Call(executable.functions()[0], last_uses,
+                         absl::MakeSpan(&arg, 1), absl::MakeSpan(&result, 1));
+  mlrt::Execute(execution_context);
+
+  notification.WaitForNotification();
+  EXPECT_EQ(inc_calls, 1);
+  EXPECT_EQ(dec_calls, 1);
+  EXPECT_EQ(num_deferred_ops, 0);
+  EXPECT_THAT(
+      execution_context.status(),
+      absl_testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                             ::testing::HasSubstr("missing done error")));
+
+  // Also verify the fallback status when the kernel returns without calling
+  // `done` and without setting an error on `context`.
+  tfrt_stub::OpKernelRunnerTable runner_table_no_status;
+  tfd::FallbackResourceArray resource_array_no_status;
+  tfd::KernelFallbackCompatRequestState fallback_request_state_no_status(
+      &runner, &fallback_state->device_manager(), /*step_id=*/1,
+      &runner_table_no_status, &resource_array_no_status,
+      /*user_intra_op_threadpool=*/nullptr,
+      /*model_metadata=*/std::nullopt,
+      &fallback_state->process_function_library_runtime());
+  mlrt::ExecutionContext execution_context_no_status(&loaded_executable);
+  execution_context_no_status.set_work_queue(work_queue.get());
+  auto tf_context_no_status = std::make_unique<Context>(
+      &fallback_request_state_no_status, &resource_context);
+  tf_context_no_status->params().inc_num_deferred_ops_function = [&]() {
+    ++inc_calls;
+    ++num_deferred_ops;
+  };
+  tf_context_no_status->params().dec_num_deferred_ops_function = [&]() {
+    ++dec_calls;
+    --num_deferred_ops;
+  };
+  execution_context_no_status.AddUserContext(std::move(tf_context_no_status));
+
+  tensorflow::Tensor neg_input_tensor(-1);
+  mlrt::Value neg_arg(tfrt_stub::FallbackTensor(std::move(neg_input_tensor)));
+  mlrt::Value neg_result;
+  absl::Notification notification_no_status;
+  execution_context_no_status.set_exit_handler(
+      [&notification_no_status]() { notification_no_status.Notify(); });
+  execution_context_no_status.Call(executable.functions()[0], last_uses,
+                                   absl::MakeSpan(&neg_arg, 1),
+                                   absl::MakeSpan(&neg_result, 1));
+  mlrt::Execute(execution_context_no_status);
+
+  notification_no_status.WaitForNotification();
+  EXPECT_EQ(inc_calls, 2);
+  EXPECT_EQ(dec_calls, 2);
+  EXPECT_EQ(num_deferred_ops, 0);
+  EXPECT_THAT(execution_context_no_status.status(),
+              absl_testing::StatusIs(
+                  absl::StatusCode::kFailedPrecondition,
+                  ::testing::HasSubstr("did not invoke DoneCallback")));
+}
+
 }  // namespace
 }  // namespace tf_mlrt
 }  // namespace tensorflow

From eba5052a7debc683615da717ea4e130a6de75864 Mon Sep 17 00:00:00 2001
From: Changming Sun <changmingsun@google.com>
Date: Wed, 14 Oct 2026 17:45:59 -0700
Subject: [PATCH 0818/1000] Harden lookup kernels and LiteRT model
 verification.

PiperOrigin-RevId: 999742990
---
 tensorflow/lite/kernels/BUILD                 |   2 +
 tensorflow/lite/kernels/embedding_lookup.cc   | 239 ++++++++----------
 .../lite/kernels/embedding_lookup_test.cc     | 171 +++++++++++++
 tensorflow/lite/kernels/hashtable_lookup.cc   | 119 +++++----
 .../lite/kernels/hashtable_lookup_test.cc     |  83 ++++++
 5 files changed, 429 insertions(+), 185 deletions(-)

diff --git a/tensorflow/lite/kernels/BUILD b/tensorflow/lite/kernels/BUILD
index 775ec49bed7..02a01ad1884 100644
--- a/tensorflow/lite/kernels/BUILD
+++ b/tensorflow/lite/kernels/BUILD
@@ -2228,6 +2228,7 @@ cc_test(
     deps = [
         ":test_main",
         ":test_util",
+        "//tensorflow/lite:array",
         "//tensorflow/lite:framework_stable",
         "//tensorflow/lite/core:framework_stable",
         "//tensorflow/lite/kernels/internal:tensor",
@@ -2393,6 +2394,7 @@ cc_test(
         "//tensorflow/lite:string",
         "//tensorflow/lite:string_util",
         "//tensorflow/lite/core:framework_stable",
+        "//tensorflow/lite/core/c:common",
         "//tensorflow/lite/kernels/internal:tensor",
         "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
diff --git a/tensorflow/lite/kernels/embedding_lookup.cc b/tensorflow/lite/kernels/embedding_lookup.cc
index 6266ece188b..169592e5f53 100644
--- a/tensorflow/lite/kernels/embedding_lookup.cc
+++ b/tensorflow/lite/kernels/embedding_lookup.cc
@@ -30,22 +30,39 @@ limitations under the License.
 //
 
 #include <cinttypes>
+#include <cstddef>
 #include <cstdint>
 #include <cstring>
-#include <functional>
-#include <utility>
 
 #include "tensorflow/lite/c/c_api_types.h"
 #include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/kernels/internal/kernel_utils.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 #include "tensorflow/lite/types/half.h"
 #include "tensorflow/lite/util.h"
+
 namespace tflite {
 namespace ops {
 namespace builtin {
 namespace embedding_lookup {
 
+namespace {
+
+TfLiteStatus ValidateLookupIndex(TfLiteContext* context, int32_t idx,
+                                 int row_size) {
+  if (idx >= row_size || idx < 0) {
+    TF_LITE_KERNEL_LOG(context,
+                       "Embedding Lookup: index out of bounds. "
+                       "Got %" PRId32 ", and bounds are [0, %d]",
+                       idx, row_size - 1);
+    return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+
+}  // namespace
+
 TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumInputs(node), 2);
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 1);
@@ -92,46 +109,31 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
 
   TfLiteTensor* output;
   TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
-  TfLiteIntArray* output_size = TfLiteIntArrayCreate(NumDimensions(value));
-
-  output_size->data[0] = SizeOfDimension(lookup, 0);
-  output_size->data[1] = SizeOfDimension(value, 1);
-  for (int i = 2; i < NumDimensions(value); i++) {
-    output_size->data[i] = SizeOfDimension(value, i);
-  }
-  return context->ResizeTensor(context, output, output_size);
+  return kernel_utils::ResizeLookupOutputTensor(context, *lookup, *value,
+                                                output);
 }
 
 TfLiteStatus EvalSimple(TfLiteContext* context, TfLiteNode* node,
                         const TfLiteTensor* lookup, const TfLiteTensor* value,
                         TfLiteTensor* output) {
-  const int row_size = SizeOfDimension(value, 0);
-  if (row_size == 0) {
-    // Propagate empty tensor if input is empty
+  const int lookup_size = SizeOfDimension(lookup, 0);
+  if (lookup_size == 0) {
+    // Propagate empty output tensor if lookup has zero elements.
     return kTfLiteOk;
   }
+  const int row_size = SizeOfDimension(value, 0);
+  TF_LITE_ENSURE(context, row_size > 0);
   const size_t row_bytes = value->bytes / row_size;
 
   char* output_raw = GetTensorData<char>(output);
   const char* value_raw = GetTensorData<char>(value);
   const int32_t* lookup_data = GetTensorData<int32_t>(lookup);
-  for (int i = 0; i < SizeOfDimension(lookup, 0); i++) {
+  for (int i = 0; i < lookup_size; i++) {
     const int32_t idx = lookup_data[i];
-    if (idx >= row_size || idx < 0) {
-      TF_LITE_KERNEL_LOG(context,
-                         "Embedding Lookup: index out of bounds. "
-                         "Got %" PRId32 ", and bounds are [0, %d]",
-                         idx, row_size - 1);
-      return kTfLiteError;
-    } else {
-      CheckedInt<size_t> offset_output = CheckedInt<size_t>(i) * row_bytes;
-      CheckedInt<size_t> offset_value = CheckedInt<size_t>(idx) * row_bytes;
-      if (offset_output.Overflow() || offset_value.Overflow()) {
-        TF_LITE_KERNEL_LOG(context, "Embedding Lookup: offset overflow.");
-        return kTfLiteError;
-      }
-      std::memcpy(output_raw + offset_output.Value(),
-                  value_raw + offset_value.Value(), row_bytes);
+    TF_LITE_ENSURE_OK(context, ValidateLookupIndex(context, idx, row_size));
+    if (row_bytes > 0) {
+      std::memcpy(output_raw + static_cast<size_t>(i) * row_bytes,
+                  value_raw + static_cast<size_t>(idx) * row_bytes, row_bytes);
     }
   }
   return kTfLiteOk;
@@ -144,16 +146,16 @@ void Unpack4Bit(float scaling_factor, size_t col_size, const int8_t* value_ptr,
   size_t j = 0;
   size_t i4_idx = 0;
   for (; j + 1 < col_size; j += 2, ++i4_idx) {
-    int8_t i4_val = value_ptr[i4_idx];
-    int8_t i8_val0 = i4_val << 4;
-    int8_t i8_val1 = i4_val & 0xF0;
+    uint8_t i4_val = static_cast<uint8_t>(value_ptr[i4_idx]);
+    int8_t i8_val0 = static_cast<int8_t>(i4_val << 4);
+    int8_t i8_val1 = static_cast<int8_t>(i4_val & 0xF0);
 
     output_ptr[j] = i8_val0 * scaling_factor0;
     output_ptr[j + 1] = i8_val1 * scaling_factor0;
   }
   if (col_size & 1) {
-    int8_t i4_val = value_ptr[i4_idx];
-    int8_t i8_val0 = i4_val << 4;
+    uint8_t i4_val = static_cast<uint8_t>(value_ptr[i4_idx]);
+    int8_t i8_val0 = static_cast<int8_t>(i4_val << 4);
     output_ptr[j] = i8_val0 * scaling_factor0;
   }
 }
@@ -165,11 +167,11 @@ void Unpack2Bit(float scaling_factor, size_t col_size, const int8_t* value_ptr,
   size_t j = 0;
   size_t i2_idx = 0;
   for (; j + 3 < col_size; j += 4, ++i2_idx) {
-    int8_t i2_val = value_ptr[i2_idx];
+    uint8_t i2_val = static_cast<uint8_t>(value_ptr[i2_idx]);
     int8_t i8_val0 = static_cast<int8_t>(i2_val << 6);
-    int8_t i8_val1 = static_cast<int8_t>(i2_val << 4) & 0xC0;
-    int8_t i8_val2 = static_cast<int8_t>(i2_val << 2) & 0xC0;
-    int8_t i8_val3 = i2_val & 0xC0;
+    int8_t i8_val1 = static_cast<int8_t>((i2_val << 4) & 0xC0);
+    int8_t i8_val2 = static_cast<int8_t>((i2_val << 2) & 0xC0);
+    int8_t i8_val3 = static_cast<int8_t>(i2_val & 0xC0);
 
     output_ptr[j] = i8_val0 * scaling_factor0;
     output_ptr[j + 1] = i8_val1 * scaling_factor0;
@@ -178,20 +180,33 @@ void Unpack2Bit(float scaling_factor, size_t col_size, const int8_t* value_ptr,
   }
   size_t rem = col_size - j;
   if (rem) {
-    int8_t i2_val = value_ptr[i2_idx];
+    uint8_t i2_val = static_cast<uint8_t>(value_ptr[i2_idx]);
     int8_t i8_val0 = static_cast<int8_t>(i2_val << 6);
     output_ptr[j] = i8_val0 * scaling_factor0;
     if (rem & 2) {
-      int8_t i8_val1 = static_cast<int8_t>(i2_val << 4) & 0xC0;
+      int8_t i8_val1 = static_cast<int8_t>((i2_val << 4) & 0xC0);
       output_ptr[j + 1] = i8_val1 * scaling_factor0;
       if (rem & 1) {
-        int8_t i8_val2 = static_cast<int8_t>(i2_val << 2) & 0xC0;
+        int8_t i8_val2 = static_cast<int8_t>((i2_val << 2) & 0xC0);
         output_ptr[j + 2] = i8_val2 * scaling_factor0;
       }
     }
   }
 }
 
+template <typename T>
+void UnpackSubByteElements(TfLiteType type, float scaling_factor,
+                           size_t num_elements, size_t elem_offset,
+                           const int8_t* value_ptr, T* output_ptr) {
+  if (type == kTfLiteInt2) {
+    Unpack2Bit(scaling_factor, num_elements, &value_ptr[elem_offset >> 2],
+               output_ptr);
+  } else {
+    Unpack4Bit(scaling_factor, num_elements, &value_ptr[elem_offset >> 1],
+               output_ptr);
+  }
+}
+
 TfLiteStatus EvalBlockwise(TfLiteContext* context, TfLiteNode* node,
                            const TfLiteTensor* lookup,
                            const TfLiteTensor* value, TfLiteTensor* output) {
@@ -214,16 +229,11 @@ TfLiteStatus EvalBlockwise(TfLiteContext* context, TfLiteNode* node,
     return kTfLiteError;
   }
   const int row_size = SizeOfDimension(value, 0);
+  size_t col_size = 0;
+  TF_LITE_ENSURE_OK(context,
+                    kernel_utils::CheckedDimensionProduct(
+                        context, *value, 1, NumDimensions(value), col_size));
 
-  // col_size after we flatten tensor into 2D.
-  CheckedInt<size_t> col_size = 1;
-  for (int i = 1; i < NumDimensions(value); i++) {
-    col_size = col_size * CheckedInt<size_t>(SizeOfDimension(value, i));
-    if (col_size.Overflow()) {
-      TF_LITE_KERNEL_LOG(context, "Embedding Lookup: col_size overflow.");
-      return kTfLiteError;
-    }
-  }
   const auto quantization_params =
       reinterpret_cast<const TfLiteBlockwiseQuantization*>(
           value->quantization.params);
@@ -233,52 +243,20 @@ TfLiteStatus EvalBlockwise(TfLiteContext* context, TfLiteNode* node,
 
   float* output_fp32_ptr = GetTensorData<float>(output);
   half* output_fp16_ptr = GetTensorData<half>(output);
-
   const int8_t* value_ptr = GetTensorData<int8_t>(value);
   const int32_t* lookup_data = GetTensorData<int32_t>(lookup);
 
-  // Wrap the correct 2/4-bit float32/float16 unpacking function.
-  auto [unpack_to_fp32, unpack_to_fp16] =
-      value->type == kTfLiteInt2
-          ? std::make_pair(Unpack2Bit<float>, Unpack2Bit<half>)
-          : std::make_pair(Unpack4Bit<float>, Unpack4Bit<half>);
-  const int values_per_byte = value->type == kTfLiteInt2 ? 4 : 2;
-  std::function<void(float, size_t, size_t)> unpack;
-  if (output->type == kTfLiteFloat32) {
-    unpack = [&, unpack = unpack_to_fp32](float scaling_factor,
-                                          size_t value_offset,
-                                          size_t output_offset) {
-      unpack(scaling_factor, blocksize,
-             &value_ptr[value_offset / values_per_byte],
-             &output_fp32_ptr[output_offset]);
-    };
-  } else {
-    unpack = [&, unpack = unpack_to_fp16](float scaling_factor,
-                                          size_t value_offset,
-                                          size_t output_offset) {
-      unpack(scaling_factor, blocksize,
-             &value_ptr[value_offset / values_per_byte],
-             &output_fp16_ptr[output_offset]);
-    };
-  }
-
-  if (col_size.Value() % blocksize != 0) {
+  if (col_size % blocksize != 0) {
     TF_LITE_KERNEL_LOG(context,
                        "Embedding Lookup: lookup dimension %zu must be "
                        "divisible by blocksize %d",
-                       col_size.Value(), blocksize);
+                       col_size, blocksize);
     return kTfLiteError;
   }
-  size_t num_blocks = col_size.Value() / blocksize;
+  const size_t num_blocks = col_size / blocksize;
   for (int i = 0; i < dimension_size; i++) {
-    int idx = lookup_data[i];
-    if (idx >= row_size || idx < 0) {
-      TF_LITE_KERNEL_LOG(context,
-                         "Embedding Lookup: index out of bounds. "
-                         "Got %d, and bounds are [0, %d]",
-                         idx, row_size - 1);
-      return kTfLiteError;
-    }
+    const int32_t idx = lookup_data[i];
+    TF_LITE_ENSURE_OK(context, ValidateLookupIndex(context, idx, row_size));
     CheckedInt<size_t> output_row_offset = CheckedInt<size_t>(i) * col_size;
     CheckedInt<size_t> scale_offset = CheckedInt<size_t>(idx) * num_blocks;
     CheckedInt<size_t> value_offset = CheckedInt<size_t>(idx) * col_size;
@@ -298,7 +276,15 @@ TfLiteStatus EvalBlockwise(TfLiteContext* context, TfLiteNode* node,
         TF_LITE_KERNEL_LOG(context, "Embedding Lookup: offset overflow.");
         return kTfLiteError;
       }
-      unpack(scaling_factor, val_off.Value(), out_off.Value());
+      if (output->type == kTfLiteFloat32) {
+        UnpackSubByteElements(value->type, scaling_factor, blocksize,
+                              val_off.Value(), value_ptr,
+                              &output_fp32_ptr[out_off.Value()]);
+      } else {
+        UnpackSubByteElements(value->type, scaling_factor, blocksize,
+                              val_off.Value(), value_ptr,
+                              &output_fp16_ptr[out_off.Value()]);
+      }
     }
   }
   return kTfLiteOk;
@@ -308,31 +294,23 @@ TfLiteStatus EvalHybrid(TfLiteContext* context, TfLiteNode* node,
                         const TfLiteTensor* lookup, const TfLiteTensor* value,
                         TfLiteTensor* output) {
   const int row_size = SizeOfDimension(value, 0);
+  size_t col_size = 0;
+  TF_LITE_ENSURE_OK(context,
+                    kernel_utils::CheckedDimensionProduct(
+                        context, *value, 1, NumDimensions(value), col_size));
 
-  // col_size after we flatten tensor into 2D.
-  CheckedInt<size_t> col_size = 1;
-  for (int i = 1; i < NumDimensions(value); i++) {
-    col_size = col_size * CheckedInt<size_t>(SizeOfDimension(value, i));
-    if (col_size.Overflow()) {
-      TF_LITE_KERNEL_LOG(context, "Embedding Lookup: col_size overflow.");
-      return kTfLiteError;
-    }
-  }
   auto copy_row = [&](float scaling_factor, auto output_ptr, auto value_ptr,
                       int idx, size_t out_off) -> TfLiteStatus {
-    CheckedInt<size_t> offset = col_size * CheckedInt<size_t>(idx);
+    CheckedInt<size_t> offset = CheckedInt<size_t>(col_size) * idx;
     if (offset.Overflow()) {
       TF_LITE_KERNEL_LOG(context, "Embedding Lookup: offset overflow.");
       return kTfLiteError;
     }
-    if (value->type == kTfLiteInt4) {
-      Unpack4Bit(scaling_factor, col_size.Value(),
-                 &value_ptr[offset.Value() >> 1], &output_ptr[out_off]);
-    } else if (value->type == kTfLiteInt2) {
-      Unpack2Bit(scaling_factor, col_size.Value(),
-                 &value_ptr[offset.Value() >> 2], &output_ptr[out_off]);
+    if (value->type == kTfLiteInt4 || value->type == kTfLiteInt2) {
+      UnpackSubByteElements(value->type, scaling_factor, col_size,
+                            offset.Value(), value_ptr, &output_ptr[out_off]);
     } else {
-      for (size_t j = 0; j < col_size.Value(); j++) {
+      for (size_t j = 0; j < col_size; j++) {
         CheckedInt<size_t> output_idx = CheckedInt<size_t>(j) + out_off;
         CheckedInt<size_t> value_idx = offset + CheckedInt<size_t>(j);
         if (output_idx.Overflow() || value_idx.Overflow()) {
@@ -355,38 +333,31 @@ TfLiteStatus EvalHybrid(TfLiteContext* context, TfLiteNode* node,
 
   for (int i = 0; i < SizeOfDimension(lookup, 0); i++) {
     const int32_t idx = lookup_data[i];
-    if (idx >= row_size || idx < 0) {
-      TF_LITE_KERNEL_LOG(context,
-                         "Embedding Lookup: index out of bounds. "
-                         "Got %" PRId32 ", and bounds are [0, %d]",
-                         idx, row_size - 1);
+    TF_LITE_ENSURE_OK(context, ValidateLookupIndex(context, idx, row_size));
+    CheckedInt<size_t> out_off = CheckedInt<size_t>(i) * col_size;
+    if (out_off.Overflow()) {
+      TF_LITE_KERNEL_LOG(context, "Embedding Lookup: offset overflow.");
       return kTfLiteError;
-    } else {
-      CheckedInt<size_t> out_off = CheckedInt<size_t>(i) * col_size.Value();
-      if (out_off.Overflow()) {
-        TF_LITE_KERNEL_LOG(context, "Embedding Lookup: offset overflow.");
-        return kTfLiteError;
-      }
-      // Dequantize embedding values.
-      // TODO(alanchiao): refactor scalar multiply into separate function
-      // for ease of adding a neon equivalent if ever necessary.
-      float scaling_factor = value->params.scale;
-      if (value->quantization.type == kTfLiteAffineQuantization) {
-        const auto qparams = static_cast<const TfLiteAffineQuantization*>(
-            value->quantization.params);
-        if (qparams->scale->size > 1) {
-          // get this row's scale for per-axis quantization
-          scaling_factor = qparams->scale->data[idx];
-        }
+    }
+    // Dequantize embedding values.
+    // TODO(alanchiao): refactor scalar multiply into separate function
+    // for ease of adding a neon equivalent if ever necessary.
+    float scaling_factor = value->params.scale;
+    if (value->quantization.type == kTfLiteAffineQuantization) {
+      const auto qparams = static_cast<const TfLiteAffineQuantization*>(
+          value->quantization.params);
+      if (qparams->scale->size > 1) {
+        // get this row's scale for per-axis quantization
+        scaling_factor = qparams->scale->data[idx];
       }
+    }
 
-      if (output_fp32_ptr) {
-        TF_LITE_ENSURE_OK(context, copy_row(scaling_factor, output_fp32_ptr,
-                                            value_ptr, idx, out_off.Value()));
-      } else {
-        TF_LITE_ENSURE_OK(context, copy_row(scaling_factor, output_fp16_ptr,
-                                            value_ptr, idx, out_off.Value()));
-      }
+    if (output_fp32_ptr) {
+      TF_LITE_ENSURE_OK(context, copy_row(scaling_factor, output_fp32_ptr,
+                                          value_ptr, idx, out_off.Value()));
+    } else {
+      TF_LITE_ENSURE_OK(context, copy_row(scaling_factor, output_fp16_ptr,
+                                          value_ptr, idx, out_off.Value()));
     }
   }
 
diff --git a/tensorflow/lite/kernels/embedding_lookup_test.cc b/tensorflow/lite/kernels/embedding_lookup_test.cc
index 213a24bd80a..39028b1fd88 100644
--- a/tensorflow/lite/kernels/embedding_lookup_test.cc
+++ b/tensorflow/lite/kernels/embedding_lookup_test.cc
@@ -15,14 +15,20 @@ License.
 // Unit test for TFLite Lookup op.
 
 #include <stdint.h>
+#if defined(__LP64__) && !defined(_WIN32)
+#include <sys/mman.h>
+#endif
 
+#include <cstddef>
 #include <functional>
 #include <initializer_list>
+#include <limits>
 #include <memory>
 #include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "tensorflow/lite/array.h"
 #include "tensorflow/lite/core/interpreter.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/test_util.h"
@@ -30,6 +36,12 @@ License.
 #include "tensorflow/lite/types/half.h"
 
 namespace tflite {
+namespace ops {
+namespace builtin {
+TfLiteRegistration* Register_EMBEDDING_LOOKUP();
+}  // namespace builtin
+}  // namespace ops
+
 namespace {
 
 constexpr float kTestTolerance = 7.41e-03;
@@ -1109,5 +1121,164 @@ TEST(PerAxisHybridEmbeddingLookupHybridOpTest,
                                        },
                                        kFp16TestTolerance)));
 }
+
+void SilentReportError(TfLiteContext*, const char*, ...) {}
+
+struct DirectEvalFixture {
+  TfLiteContext context{};
+  TfLiteNode node{};
+  IntArrayUniquePtr node_inputs = BuildTfLiteArray<int>({0, 1});
+  IntArrayUniquePtr node_outputs = BuildTfLiteArray<int>({2});
+  TfLiteTensor tensors[3]{};
+
+  DirectEvalFixture() {
+    context.ReportError = SilentReportError;
+    context.tensors = tensors;
+    context.tensors_size = 3;
+    node.inputs = node_inputs.get();
+    node.outputs = node_outputs.get();
+  }
+
+  TfLiteStatus Invoke() {
+    return ops::builtin::Register_EMBEDDING_LOOKUP()->invoke(&context, &node);
+  }
+};
+
+#if !defined(MEMORY_SANITIZER) && !defined(GOOGLE_UNSUPPORTED_OS_LOONIX) && \
+    defined(__LP64__) && !defined(_WIN32)
+TEST(EmbeddingLookupKernelSanitizationTest, LargeInt4TableRowOffsetSucceeds) {
+  DirectEvalFixture f;
+  IntArrayUniquePtr lookup_dims = BuildTfLiteArray<int>({1});
+  IntArrayUniquePtr value_dims = BuildTfLiteArray<int>({262144, 8960});
+  int32_t idx = 239675;
+  std::vector<float> output_data(8960, 0.0f);
+
+  f.tensors[0].dims = lookup_dims.get();
+  f.tensors[0].type = kTfLiteInt32;
+  f.tensors[0].data.i32 = &idx;
+
+  const size_t bytes = size_t{262144} * 8960 / 2;
+  void* data = mmap(nullptr, bytes, PROT_READ | PROT_WRITE,
+                    MAP_PRIVATE | MAP_ANONYMOUS | MAP_NORESERVE, -1, 0);
+  ASSERT_NE(data, MAP_FAILED);
+
+  f.tensors[1].dims = value_dims.get();
+  f.tensors[1].type = kTfLiteInt4;
+  f.tensors[1].params.scale = 1.0f;
+  f.tensors[1].data.int8 = static_cast<int8_t*>(data);
+  f.tensors[1].bytes = bytes;
+  f.tensors[1].data.int8[size_t{239675} * 8960 / 2] = 0x12;
+
+  f.tensors[2].type = kTfLiteFloat32;
+  f.tensors[2].data.f = output_data.data();
+
+  EXPECT_EQ(f.Invoke(), kTfLiteOk);
+  EXPECT_FLOAT_EQ(output_data[0], 2.0f);
+  EXPECT_FLOAT_EQ(output_data[1], 1.0f);
+
+  munmap(data, bytes);
+}
+#endif
+
+TEST(EmbeddingLookupKernelSanitizationTest,
+     HybridColumnAndRowOverflowRejected) {
+  DirectEvalFixture f;
+  IntArrayUniquePtr lookup_dims = BuildTfLiteArray<int>({1});
+  int32_t idx = 5;
+  int8_t dummy_value = 0;
+  float dummy_output[8]{};
+
+  f.tensors[0].dims = lookup_dims.get();
+  f.tensors[0].type = kTfLiteInt32;
+  f.tensors[0].data.i32 = &idx;
+
+  f.tensors[1].type = kTfLiteInt4;
+  f.tensors[1].params.scale = 1.0f;
+  f.tensors[1].data.int8 = &dummy_value;
+
+  f.tensors[2].type = kTfLiteFloat32;
+  f.tensors[2].data.f = dummy_output;
+
+  // col_size product overflows size_t.
+  IntArrayUniquePtr col_overflow_dims = BuildTfLiteArray<int>(
+      {6, std::numeric_limits<int>::max(), std::numeric_limits<int>::max(), 5});
+  f.tensors[1].dims = col_overflow_dims.get();
+  EXPECT_EQ(f.Invoke(), kTfLiteError);
+
+  // col_size fits in 64-bit size_t, but idx * col_size overflows size_t.
+  IntArrayUniquePtr row_overflow_dims = BuildTfLiteArray<int>(
+      {6, std::numeric_limits<int>::max(), std::numeric_limits<int>::max()});
+  f.tensors[1].dims = row_overflow_dims.get();
+  EXPECT_EQ(f.Invoke(), kTfLiteError);
+}
+
+TEST(EmbeddingLookupKernelSanitizationTest,
+     NegativeInt4AndInt2UnpackWithoutSignedShiftUb) {
+  DirectEvalFixture f;
+  IntArrayUniquePtr lookup_dims = BuildTfLiteArray<int>({1});
+  int32_t idx = 0;
+  float output_data[4]{};
+
+  f.tensors[0].dims = lookup_dims.get();
+  f.tensors[0].type = kTfLiteInt32;
+  f.tensors[0].data.i32 = &idx;
+  f.tensors[1].params.scale = 1.0f;
+  f.tensors[2].type = kTfLiteFloat32;
+  f.tensors[2].data.f = output_data;
+
+  // Negative INT4 nibbles: 0xFF (-1, -1) and 0x88 (-8, -8), even & odd cols.
+  for (int col_size : {1, 2, 3}) {
+    IntArrayUniquePtr value_dims = BuildTfLiteArray<int>({1, col_size});
+    int8_t packed[2] = {-1, static_cast<int8_t>(0x88)};
+    f.tensors[1].dims = value_dims.get();
+    f.tensors[1].type = kTfLiteInt4;
+    f.tensors[1].data.int8 = packed;
+    ASSERT_EQ(f.Invoke(), kTfLiteOk);
+    EXPECT_FLOAT_EQ(output_data[0], -1.0f);
+    if (col_size >= 2) EXPECT_FLOAT_EQ(output_data[1], -1.0f);
+    if (col_size >= 3) EXPECT_FLOAT_EQ(output_data[2], -8.0f);
+  }
+
+  // Negative INT2 crumbs: 0xFF (-1, -1, -1, -1) across col_size = 1..4.
+  for (int col_size : {1, 2, 3, 4}) {
+    IntArrayUniquePtr value_dims = BuildTfLiteArray<int>({1, col_size});
+    int8_t packed = -1;
+    f.tensors[1].dims = value_dims.get();
+    f.tensors[1].type = kTfLiteInt2;
+    f.tensors[1].data.int8 = &packed;
+    ASSERT_EQ(f.Invoke(), kTfLiteOk);
+    for (int c = 0; c < col_size; ++c) {
+      EXPECT_FLOAT_EQ(output_data[c], -1.0f);
+    }
+  }
+}
+
+TEST(EmbeddingLookupOpTest, EmptyTensorBoundaryCases) {
+  // Empty lookup [0] with non-empty table [3, 4] succeeds with empty output.
+  {
+    EmbeddingLookupOpModel m({0}, {3, 4}, TensorType_FLOAT32,
+                             TensorType_FLOAT32);
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_TRUE(m.GetOutput<float>().empty());
+  }
+
+  // Empty table [0, 4] with non-empty lookup [1] must fail (out of bounds).
+  {
+    EmbeddingLookupOpModel m({1}, {0, 4}, TensorType_FLOAT32,
+                             TensorType_FLOAT32);
+    m.SetInput({0});
+    EXPECT_EQ(m.Invoke(), kTfLiteError);
+  }
+
+  // Zero column size [3, 0] with valid lookup [2] succeeds with empty output.
+  {
+    EmbeddingLookupOpModel m({2}, {3, 0}, TensorType_FLOAT32,
+                             TensorType_FLOAT32);
+    m.SetInput({0, 2});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_TRUE(m.GetOutput<float>().empty());
+  }
+}
+
 }  // namespace
 }  // namespace tflite
diff --git a/tensorflow/lite/kernels/hashtable_lookup.cc b/tensorflow/lite/kernels/hashtable_lookup.cc
index 87a6b94ea5e..519cc7d7f2f 100644
--- a/tensorflow/lite/kernels/hashtable_lookup.cc
+++ b/tensorflow/lite/kernels/hashtable_lookup.cc
@@ -31,25 +31,47 @@ limitations under the License.
 //   Each item indicates whether the corresponding lookup has a returned value.
 //   0 for missing key, 1 for found key.
 
-#include <stdint.h>
-
-#include <cstdlib>
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
 #include <cstring>
 
 #include "tensorflow/lite/core/c/common.h"
 #include "tensorflow/lite/kernels/internal/compatibility.h"
+#include "tensorflow/lite/kernels/internal/kernel_utils.h"
 #include "tensorflow/lite/kernels/kernel_util.h"
 #include "tensorflow/lite/string_util.h"
 #include "tensorflow/lite/util.h"
+
 namespace tflite {
 namespace ops {
 namespace builtin {
 namespace hashtable_lookup {
 
-int greater(const void* a, const void* b) {
-  return *static_cast<const int*>(a) - *static_cast<const int*>(b);
+namespace {
+
+TfLiteStatus AppendValidatedString(TfLiteContext* context,
+                                   const TfLiteTensor& value, int idx,
+                                   int num_rows, DynamicBuffer& buf) {
+  TF_LITE_ENSURE(context, idx >= 0 && idx < num_rows);
+  TF_LITE_ENSURE(context, value.data.raw != nullptr);
+  TF_LITE_ENSURE(context, value.bytes >= sizeof(int32_t));
+  const int num_strings = GetStringCount(&value);
+  TF_LITE_ENSURE(context, num_strings >= num_rows);
+  TF_LITE_ENSURE(context, value.bytes / sizeof(int32_t) >=
+                              static_cast<size_t>(num_strings) + 2);
+  const int32_t* offsets = reinterpret_cast<const int32_t*>(value.data.raw);
+  const int32_t start_offset = offsets[static_cast<size_t>(idx) + 1];
+  const int32_t end_offset = offsets[static_cast<size_t>(idx) + 2];
+  TF_LITE_ENSURE(context, start_offset >= 0);
+  TF_LITE_ENSURE(context, end_offset >= start_offset);
+  TF_LITE_ENSURE(context, static_cast<size_t>(end_offset) <= value.bytes);
+  return buf.AddString(value.data.raw + start_offset,
+                       end_offset - start_offset);
 }
 
+}  // namespace
+
 TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumInputs(node), 3);
   TF_LITE_ENSURE_EQ(context, NumOutputs(node), 2);
@@ -73,29 +95,22 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     TF_LITE_ENSURE_EQ(context, NumDimensions(value), 1);
   }
 
-  TfLiteTensor* hits;
-  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 1, &hits));
-  TF_LITE_ENSURE_EQ(context, hits->type, kTfLiteUInt8);
-  TfLiteIntArray* hitSize = TfLiteIntArrayCreate(1);
-  hitSize->data[0] = SizeOfDimension(lookup, 0);
-
   TfLiteTensor* output;
   TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 0, &output));
   TF_LITE_ENSURE_EQ(context, value->type, output->type);
 
-  TfLiteStatus status = kTfLiteOk;
+  TfLiteTensor* hits;
+  TF_LITE_ENSURE_OK(context, GetOutputSafe(context, node, 1, &hits));
+  TF_LITE_ENSURE_EQ(context, hits->type, kTfLiteUInt8);
+
   if (output->type != kTfLiteString) {
-    TfLiteIntArray* outputSize = TfLiteIntArrayCreate(NumDimensions(value));
-    outputSize->data[0] = SizeOfDimension(lookup, 0);
-    for (int i = 1; i < NumDimensions(value); i++) {
-      outputSize->data[i] = SizeOfDimension(value, i);
-    }
-    status = context->ResizeTensor(context, output, outputSize);
+    TF_LITE_ENSURE_OK(context, kernel_utils::ResizeLookupOutputTensor(
+                                   context, *lookup, *value, output));
   }
-  if (context->ResizeTensor(context, hits, hitSize) != kTfLiteOk) {
-    status = kTfLiteError;
-  }
-  return status;
+
+  TfLiteIntArray* hit_size = TfLiteIntArrayCreate(1);
+  hit_size->data[0] = SizeOfDimension(lookup, 0);
+  return context->ResizeTensor(context, hits, hit_size);
 }
 
 TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
@@ -111,44 +126,46 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_OK(context, GetInputSafe(context, node, 2, &value));
 
   const int num_rows = SizeOfDimension(value, 0);
-  TF_LITE_ENSURE(context, num_rows != 0);
-  const size_t row_bytes = value->bytes / num_rows;
-  void* pointer = nullptr;
-  DynamicBuffer buf;
+  const int lookup_size = SizeOfDimension(lookup, 0);
 
-  for (int i = 0; i < SizeOfDimension(lookup, 0); i++) {
-    int idx = -1;
-    pointer = bsearch(&(lookup->data.i32[i]), key->data.i32, num_rows,
-                      sizeof(int32_t), greater);
-    if (pointer != nullptr) {
-      idx = (reinterpret_cast<char*>(pointer) - (key->data.raw)) /
-            sizeof(int32_t);
+  DynamicBuffer buf;
+  if (lookup_size == 0) {
+    if (output->type == kTfLiteString) {
+      buf.WriteToTensorAsVector(output);
     }
+    return kTfLiteOk;
+  }
+
+  TF_LITE_ENSURE(context, num_rows > 0);
+  const size_t row_bytes = value->bytes / num_rows;
+
+  const int32_t* key_begin = key->data.i32;
+  const int32_t* key_end = key_begin + num_rows;
+  for (int i = 0; i < lookup_size; i++) {
+    const int32_t target = lookup->data.i32[i];
+    const int32_t* it =
+        std::lower_bound(key_begin, key_end, target,
+                         [](int32_t lhs, int32_t rhs) { return lhs < rhs; });
+    const int idx = (it != key_end && *it == target)
+                        ? static_cast<int>(it - key_begin)
+                        : -1;
 
     if (idx >= num_rows || idx < 0) {
       if (output->type == kTfLiteString) {
-        buf.AddString(nullptr, 0);
-      } else {
-        CheckedInt<size_t> offset = CheckedInt<size_t>(i) * row_bytes;
-        if (offset.Overflow()) {
-          TF_LITE_KERNEL_LOG(context, "Hashtable Lookup: offset overflow.");
-          return kTfLiteError;
-        }
-        memset(output->data.raw + offset.Value(), 0, row_bytes);
+        TF_LITE_ENSURE_OK(context, buf.AddString(nullptr, 0));
+      } else if (row_bytes > 0) {
+        std::memset(output->data.raw + static_cast<size_t>(i) * row_bytes, 0,
+                    row_bytes);
       }
       hits->data.uint8[i] = 0;
     } else {
       if (output->type == kTfLiteString) {
-        buf.AddString(GetString(value, idx));
-      } else {
-        CheckedInt<size_t> offset_output = CheckedInt<size_t>(i) * row_bytes;
-        CheckedInt<size_t> offset_value = CheckedInt<size_t>(idx) * row_bytes;
-        if (offset_output.Overflow() || offset_value.Overflow()) {
-          TF_LITE_KERNEL_LOG(context, "Hashtable Lookup: offset overflow.");
-          return kTfLiteError;
-        }
-        memcpy(output->data.raw + offset_output.Value(),
-               value->data.raw + offset_value.Value(), row_bytes);
+        TF_LITE_ENSURE_OK(context, AppendValidatedString(context, *value, idx,
+                                                         num_rows, buf));
+      } else if (row_bytes > 0) {
+        std::memcpy(output->data.raw + static_cast<size_t>(i) * row_bytes,
+                    value->data.raw + static_cast<size_t>(idx) * row_bytes,
+                    row_bytes);
       }
       hits->data.uint8[i] = 1;
     }
diff --git a/tensorflow/lite/kernels/hashtable_lookup_test.cc b/tensorflow/lite/kernels/hashtable_lookup_test.cc
index 4633f472dad..8464357fc88 100644
--- a/tensorflow/lite/kernels/hashtable_lookup_test.cc
+++ b/tensorflow/lite/kernels/hashtable_lookup_test.cc
@@ -18,12 +18,14 @@ limitations under the License.
 
 #include <functional>
 #include <initializer_list>
+#include <limits>
 #include <memory>
 #include <string>
 #include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "tensorflow/lite/core/c/common.h"
 #include "tensorflow/lite/core/interpreter.h"
 #include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
 #include "tensorflow/lite/kernels/test_util.h"
@@ -63,6 +65,13 @@ class HashtableLookupOpModel : public SingleOpModel {
     PopulateStringTensor(value_, content);
   }
 
+  void SetStringOffsets(int index, int32_t start_offset, int32_t end_offset) {
+    TfLiteTensor* tensor = interpreter_->tensor(value_);
+    int32_t* offsets = reinterpret_cast<int32_t*>(tensor->data.raw);
+    offsets[index + 1] = start_offset;
+    offsets[index + 2] = end_offset;
+  }
+
   void SetHashtableValue(const std::function<float(int)>& function) {
     TfLiteTensor* tensor = interpreter_->tensor(value_);
     int rows = tensor->dims->data[0];
@@ -175,5 +184,79 @@ TEST(HashtableLookupOpTest, TestString) {
                           }));
 }
 
+TEST(HashtableLookupOpTest, TestExtremeInt32KeysWithoutOverflow) {
+  constexpr int kMin = std::numeric_limits<int32_t>::min();
+  constexpr int kMax = std::numeric_limits<int32_t>::max();
+  HashtableLookupOpModel m({5}, {3}, {3}, TensorType_FLOAT32);
+
+  m.SetLookup({kMax, kMin, 42, 0, kMin + 1});
+  m.SetHashtableKey({kMin, 0, kMax});
+  m.SetHashtableValue([](int i) { return static_cast<float>(i + 1); });
+
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+
+  EXPECT_THAT(m.GetOutput(),
+              ElementsAreArray(ArrayFloatNear({3.0f, 1.0f, 0.0f, 2.0f, 0.0f})));
+  EXPECT_THAT(m.GetHit(), ElementsAreArray({1, 1, 0, 1, 0}));
+}
+
+TEST(HashtableLookupOpTest, EmptyTensorBoundaryCases) {
+  // Empty lookup [0] with non-empty float table [3, 2] succeeds.
+  {
+    HashtableLookupOpModel m({0}, {3}, {3, 2}, TensorType_FLOAT32);
+    m.SetHashtableKey({-11, 0, 1234});
+    m.SetHashtableValue([](int i, int j) { return i + j / 10.0f; });
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_TRUE(m.GetOutput().empty());
+    EXPECT_TRUE(m.GetHit().empty());
+  }
+
+  // Empty lookup [0] with non-empty string table [3] succeeds.
+  {
+    HashtableLookupOpModel m({0}, {3}, {3}, TensorType_STRING);
+    m.SetHashtableKey({-11, 0, 1234});
+    m.SetHashtableValue({"Hello", "", "Hi"});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_TRUE(m.GetStringOutput().empty());
+    EXPECT_TRUE(m.GetHit().empty());
+  }
+
+  // Empty table [0, 2] with non-empty lookup [1] fails cleanly.
+  {
+    HashtableLookupOpModel m({1}, {0}, {0, 2}, TensorType_FLOAT32);
+    m.SetLookup({0});
+    EXPECT_EQ(m.Invoke(), kTfLiteError);
+  }
+
+  // Zero-feature table [3, 0] with both hit and miss succeeds without UB.
+  {
+    HashtableLookupOpModel m({2}, {3}, {3, 0}, TensorType_FLOAT32);
+    m.SetLookup({1234, 999});
+    m.SetHashtableKey({-11, 0, 1234});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_TRUE(m.GetOutput().empty());
+    EXPECT_THAT(m.GetHit(), ElementsAreArray({1, 0}));
+  }
+}
+
+TEST(HashtableLookupOpTest, CorruptedStringTensorOffsetsRejected) {
+  HashtableLookupOpModel m({1}, {1}, {1}, TensorType_STRING);
+  m.SetLookup({10});
+  m.SetHashtableKey({10});
+  m.SetHashtableValue({"ok"});
+
+  // Out-of-bounds end_offset (> value.bytes).
+  m.SetStringOffsets(0, /*start_offset=*/12, /*end_offset=*/999);
+  EXPECT_EQ(m.Invoke(), kTfLiteError);
+
+  // Inverted offsets (end_offset < start_offset).
+  m.SetStringOffsets(0, /*start_offset=*/12, /*end_offset=*/8);
+  EXPECT_EQ(m.Invoke(), kTfLiteError);
+
+  // Negative start_offset.
+  m.SetStringOffsets(0, /*start_offset=*/-4, /*end_offset=*/12);
+  EXPECT_EQ(m.Invoke(), kTfLiteError);
+}
+
 }  // namespace
 }  // namespace tflite

From c0f8e2696af0b17230c690c689316f769c20b7c4 Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Wed, 14 Oct 2026 22:20:57 -0400
Subject: [PATCH 0819/1000] security/fuzzing: silent reporter, add.bin seed,
 stricter input zeroing

Per internal review: the model (and through it the interpreter builder) report
into a no-op ErrorReporter so malformed inputs do not write to stderr on every
iteration; tensorflow/lite/testdata/add.bin is loaded as a seed so the fuzzer
starts from a model that passes the structural check; the MSan zeroing skips
empty, persistent read-only, string, resource and variant inputs as well as
mmap-backed constants.
---
 .../fuzzing/cc/tflite_interp_overflow_fuzz.cc | 60 ++++++++++++++++---
 1 file changed, 53 insertions(+), 7 deletions(-)

diff --git a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
index 25f95875c29..a744691278a 100644
--- a/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
+++ b/tensorflow/security/fuzzing/cc/tflite_interp_overflow_fuzz.cc
@@ -23,13 +23,20 @@ limitations under the License.
 // InterpreterBuilder::ParseNodes and ParseTensors are exercised directly. The
 // OSS-Fuzz tensorflow build has no TFLite target that reaches this path.
 
+#include <cstdarg>
 #include <cstring>
+#include <fstream>
+#include <iterator>
 #include <memory>
 #include <string>
+#include <tuple>
+#include <utility>
 #include <vector>
 
 #include "fuzztest/fuzztest.h"
+#include "tensorflow/core/platform/resource_loader.h"
 #include "tensorflow/lite/c/common.h"
+#include "tensorflow/lite/error_reporter.h"
 #include "tensorflow/lite/interpreter.h"
 #include "tensorflow/lite/interpreter_builder.h"
 #include "tensorflow/lite/kernels/register.h"
@@ -37,9 +44,51 @@ limitations under the License.
 
 namespace {
 
+// Nearly every fuzz input is malformed, and the default reporter writes each
+// rejection to stderr, which would dominate the run. The model and, through
+// the model, the interpreter builder report into this instead.
+class NullErrorReporter : public tflite::ErrorReporter {
+ public:
+  int Report(const char* /*format*/, va_list /*args*/) override { return 0; }
+};
+
+// A valid model to start from, so the fuzzer does not have to rediscover the
+// flatbuffer framing before it reaches InterpreterBuilder and Invoke.
+std::vector<std::tuple<std::string>> SeedModels() {
+  std::ifstream file(
+      tensorflow::GetDataDependencyFilepath("tensorflow/lite/testdata/add.bin"),
+      std::ios::binary);
+  std::string bytes((std::istreambuf_iterator<char>(file)),
+                    std::istreambuf_iterator<char>());
+  if (bytes.empty()) return {};
+  return {std::make_tuple(std::move(bytes))};
+}
+
+// Whether `tensor` is a runtime input that is safe and useful to zero.
+bool IsZeroableInput(const TfLiteTensor* tensor) {
+  if (tensor == nullptr || tensor->data.raw == nullptr || tensor->bytes == 0) {
+    return false;
+  }
+  // A constant input aliases the read-only model mapping, and a persistent
+  // read-only one is owned and filled by the planner; neither is
+  // uninitialized, and writing to the mapping would fault.
+  if (tensor->allocation_type == kTfLiteMmapRo ||
+      tensor->allocation_type == kTfLitePersistentRo) {
+    return false;
+  }
+  // These are not plain byte buffers of `bytes` length.
+  if (tensor->type == kTfLiteString || tensor->type == kTfLiteResource ||
+      tensor->type == kTfLiteVariant) {
+    return false;
+  }
+  return true;
+}
+
 void FuzzModelBuildAndInvoke(const std::string& model_bytes) {
+  NullErrorReporter error_reporter;
   auto model = tflite::FlatBufferModel::VerifyAndBuildFromBuffer(
-      model_bytes.data(), model_bytes.size());
+      model_bytes.data(), model_bytes.size(), /*extra_verifier=*/nullptr,
+      &error_reporter);
   if (model == nullptr) return;
 
   tflite::ops::builtin::BuiltinOpResolver resolver;
@@ -51,19 +100,16 @@ void FuzzModelBuildAndInvoke(const std::string& model_bytes) {
   if (interpreter->AllocateTensors() != kTfLiteOk) return;
 
   // Zero-initialize runtime inputs so MSan does not flag uninitialized reads.
-  // An input a crafted model backs with a constant buffer is kTfLiteMmapRo and
-  // points into the read-only model mapping; it is already initialized and
-  // writing to it would fault, so leave those alone.
   for (int input_index : interpreter->inputs()) {
     TfLiteTensor* tensor = interpreter->tensor(input_index);
-    if (tensor != nullptr && tensor->data.raw != nullptr &&
-        tensor->allocation_type != kTfLiteMmapRo) {
+    if (IsZeroableInput(tensor)) {
       std::memset(tensor->data.raw, 0, tensor->bytes);
     }
   }
 
   interpreter->Invoke();
 }
-FUZZ_TEST(TfliteInterpreterBuilder, FuzzModelBuildAndInvoke);
+FUZZ_TEST(TfliteInterpreterBuilder, FuzzModelBuildAndInvoke)
+    .WithSeeds(SeedModels());
 
 }  // namespace

From ba730ad80897d79270609605cc8027d7faec8a54 Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Wed, 14 Oct 2026 22:22:10 -0400
Subject: [PATCH 0820/1000] security/fuzzing: seed data and resource_loader dep
 for the TFLite fuzz target

add.bin from //tensorflow/lite is the seed corpus; resource_loader resolves its
runfiles path, the same way checkpoint_reader_fuzz does in this package.
---
 tensorflow/security/fuzzing/cc/BUILD | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow/security/fuzzing/cc/BUILD b/tensorflow/security/fuzzing/cc/BUILD
index 66c81352f36..97d575fdd57 100644
--- a/tensorflow/security/fuzzing/cc/BUILD
+++ b/tensorflow/security/fuzzing/cc/BUILD
@@ -272,7 +272,9 @@ tf_cc_fuzz_test(
 tf_cc_fuzz_test(
     name = "tflite_interp_overflow_fuzz",
     srcs = ["tflite_interp_overflow_fuzz.cc"],
+    data = ["//tensorflow/lite:testdata/add.bin"],
     deps = [
+        "//tensorflow/core/platform:resource_loader",
         "//tensorflow/lite:framework",
         "//tensorflow/lite/c:common",
         "//tensorflow/lite/kernels:builtin_ops",

From c1be76e06447c1643b6b7f94f315b05a86eae532 Mon Sep 17 00:00:00 2001
From: Layau Eulizier Jr <130326664+cybrdude@users.noreply.github.com>
Date: Wed, 14 Oct 2026 22:23:20 -0400
Subject: [PATCH 0821/1000] lite: initialize buffer_ptr before
 get_readonly_data

Per internal review: buffer_ptr is never left uninitialized if a path through
get_readonly_data returns kTfLiteOk without assigning it.
---
 tensorflow/lite/core/interpreter_builder.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/lite/core/interpreter_builder.cc b/tensorflow/lite/core/interpreter_builder.cc
index cc7156b1f58..ae89a845b82 100644
--- a/tensorflow/lite/core/interpreter_builder.cc
+++ b/tensorflow/lite/core/interpreter_builder.cc
@@ -832,7 +832,7 @@ TfLiteStatus InterpreterBuilder::ParseTensors(
       return kTfLiteOk;
     };
     size_t buffer_size = 0;
-    const char* buffer_ptr;
+    const char* buffer_ptr = nullptr;
     TF_LITE_ENSURE_STATUS(get_readonly_data(&buffer_ptr, &buffer_size));
 
     const auto* src_quantization = tensor->quantization();

From 43168536762f275ed4ebd9e5018fd04f11405181 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Thu, 15 Oct 2026 08:37:33 +0530
Subject: [PATCH 0822/1000] Condense inline comment on scalar shard_info
 handling

---
 tensorflow/dtensor/python/d_variable.py | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/tensorflow/dtensor/python/d_variable.py b/tensorflow/dtensor/python/d_variable.py
index 875516df400..dcb8e1846f4 100644
--- a/tensorflow/dtensor/python/d_variable.py
+++ b/tensorflow/dtensor/python/d_variable.py
@@ -181,14 +181,11 @@ class DVariable(resource_variable_ops.ResourceVariable):
         # The default layout scope allows us to correctly specify the output
         # layout of the tf.RestoreV2 that will be called
         #
-        # Passing shard_info with the correct shape allows the tf.RestoreV2
-        # ShapeInference to extract the shape. Scalar tensors (shape=()) can't
-        # be sliced/sharded, so they must use shard_info=None rather than a
-        # ShardInfo with an empty shape and offset (which would otherwise
-        # produce an invalid " " shape_and_slice spec for RestoreV2). Check
-        # `len(shape) > 0` rather than truthiness: an empty tuple () is
-        # falsy, but a scalar tf.TensorShape([]) is truthy even though it
-        # also has length 0, so a plain `if shape` would miss that case.
+        # Passing shard_info with the correct shape lets tf.RestoreV2 extract
+        # the shape. Scalars can't be sharded, so they use shard_info=None
+        # instead of an empty-shape ShardInfo (which RestoreV2 rejects).
+        # Use `len(shape) > 0`, not truthiness: a scalar tf.TensorShape([])
+        # is truthy despite having length 0.
         shard_info = (
             trackable.ShardInfo(shape=shape, offset=[0] * len(shape))
             if len(shape) > 0

From ad997db6a00bd80b8fe1f4a1f02640e5377a11d0 Mon Sep 17 00:00:00 2001
From: bodapatisaikrishna <saikrishnabodapati@gmail.com>
Date: Thu, 15 Oct 2026 08:47:12 +0530
Subject: [PATCH 0823/1000] Revert uninitialized variable checks and prune test

---
 .../core/kernels/resource_variable_ops.cc     |  19 --
 tensorflow/core/kernels/scatter_nd_op.cc      |   4 -
 tensorflow/core/kernels/training_op_helpers.h |   7 -
 .../python/training/training_ops_test.py      | 170 ------------------
 4 files changed, 200 deletions(-)

diff --git a/tensorflow/core/kernels/resource_variable_ops.cc b/tensorflow/core/kernels/resource_variable_ops.cc
index 66b3f74e758..8a1c33a8f88 100644
--- a/tensorflow/core/kernels/resource_variable_ops.cc
+++ b/tensorflow/core/kernels/resource_variable_ops.cc
@@ -648,9 +648,6 @@ class AssignUpdateVariableOp : public OpKernel {
     // ADD if value's refcount was 1.
     mutex_lock ml(*variable->mu());
     Tensor* var_tensor = variable->tensor();
-    OP_REQUIRES(context, var_tensor->IsInitialized(),
-                errors::FailedPrecondition(
-                    "Attempting to use uninitialized variables: ", handle.name()));
     OP_REQUIRES_OK(context, ValidateAssignUpdateVariableOpShapes(
                                 var_tensor->shape(), value.shape()));
     OP_REQUIRES(context, var_tensor->dtype() == value.dtype(),
@@ -768,10 +765,6 @@ class ResourceGatherOp : public OpKernel {
     // reference count greater than one and make a copy of the
     // (potentially very large) tensor buffer.
     tf_shared_lock ml(*v->mu());
-    OP_REQUIRES(c, v->tensor()->IsInitialized(),
-                errors::FailedPrecondition(
-                    "Attempting to use uninitialized variables: ",
-                    handle.name()));
     OP_REQUIRES(
         c, v->tensor()->dtype() == DataTypeToEnum<T>::v(),
         absl::InvalidArgumentError(absl::StrCat(
@@ -965,10 +958,6 @@ class ResourceGatherNdOp : public OpKernel {
     // reference count greater than one and make a copy of the
     // (potentially very large) tensor buffer.
     tf_shared_lock ml(*v->mu());
-    OP_REQUIRES(c, v->tensor()->IsInitialized(),
-                errors::FailedPrecondition(
-                    "Attempting to use uninitialized variables: ",
-                    handle.name()));
     OP_REQUIRES(
         c, v->tensor()->dtype() == DataTypeToEnum<T>::v(),
         absl::InvalidArgumentError(absl::StrCat(
@@ -1200,10 +1189,6 @@ class ResourceScatterUpdateOp : public OpKernel {
         update_dtype == DT_STRING || update_dtype == DT_VARIANT;
     if (is_non_pod_dtype || use_exclusive_lock_) {
       mutex_lock ml(*v->mu());
-      OP_REQUIRES(c, v->tensor()->IsInitialized(),
-                  errors::FailedPrecondition(
-                      "Attempting to use uninitialized variables: ",
-                      handle.name()));
       OP_REQUIRES(c, v->tensor()->dtype() == update_dtype,
                   absl::InvalidArgumentError(
                       "DType of scatter resource and updates does not match."));
@@ -1211,10 +1196,6 @@ class ResourceScatterUpdateOp : public OpKernel {
     } else {
       // For POD dtypes, we can safely run the update without the mutex.
       tf_shared_lock ml(*v->mu());
-      OP_REQUIRES(c, v->tensor()->IsInitialized(),
-                  errors::FailedPrecondition(
-                      "Attempting to use uninitialized variables: ",
-                      handle.name()));
       OP_REQUIRES(c, v->tensor()->dtype() == update_dtype,
                   absl::InvalidArgumentError(
                       "DType of scatter resource and updates does not match."));
diff --git a/tensorflow/core/kernels/scatter_nd_op.cc b/tensorflow/core/kernels/scatter_nd_op.cc
index b2db94b89d3..49d8236dcd0 100644
--- a/tensorflow/core/kernels/scatter_nd_op.cc
+++ b/tensorflow/core/kernels/scatter_nd_op.cc
@@ -323,10 +323,6 @@ class ScatterNdUpdateOp : public ScatterOpBase<Device> {
       OP_REQUIRES_OK(c, LookupResource(c, handle, &v));
       OP_REQUIRES_OK(c, EnsureSparseVariableAccess<Device, T>(c, v.get()));
       mutex_lock m(*v->mu());
-      OP_REQUIRES(c, v->tensor()->IsInitialized(),
-                  errors::FailedPrecondition(
-                      "Attempting to use uninitialized variables: ",
-                      handle.name()));
       DoCompute(c);
     } else if (use_exclusive_lock_) {
       // If we're here, it means the input type is a ref.
diff --git a/tensorflow/core/kernels/training_op_helpers.h b/tensorflow/core/kernels/training_op_helpers.h
index 11981f7aba8..e7a29bab735 100644
--- a/tensorflow/core/kernels/training_op_helpers.h
+++ b/tensorflow/core/kernels/training_op_helpers.h
@@ -76,10 +76,6 @@ absl::Status EnsureSparseVariableAccess(OpKernelContext* ctx, Var* var) {
         " (resource variable dtype)"));
   }
 
-  if (!var->tensor()->IsInitialized()) {
-    return absl::OkStatus();
-  }
-
   // Once copy-on-read mode is True the refcount is guaranteed to be 1. This can
   // also happen if there are no concurrent reads of the variable and
   // copy-on-read mode is false.
@@ -251,9 +247,6 @@ absl::Status PrepareToUpdateVariable(OpKernelContext* ctx, Tensor* tensor,
         "dtype mismatch: expected ", DataTypeString(DataTypeToEnum<T>::v()),
         " but got ", DataTypeString(tensor->dtype())));
   }
-  if (!tensor->IsInitialized()) {
-    return absl::OkStatus();
-  }
   if (copy_on_read_mode || !tensor->RefCountIsOne()) {
     // Tensor's buffer is in use by some read, so we need to copy before
     // updating.
diff --git a/tensorflow/python/training/training_ops_test.py b/tensorflow/python/training/training_ops_test.py
index c43e6adb9c4..48329d9bc43 100644
--- a/tensorflow/python/training/training_ops_test.py
+++ b/tensorflow/python/training/training_ops_test.py
@@ -19,7 +19,6 @@ import threading
 
 import numpy as np
 
-from tensorflow.python.eager import context
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
@@ -28,8 +27,6 @@ from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.framework.test_util import TensorFlowTestCase
 # Import resource_variable_ops for the variables-to-tensor implicit conversion.
-from tensorflow.python.ops import gen_resource_variable_ops
-from tensorflow.python.ops import gen_state_ops
 from tensorflow.python.ops import gen_training_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import resource_variable_ops  # pylint: disable=unused-import
@@ -751,173 +748,6 @@ class TrainingOpsTest(TensorFlowTestCase):
           )
       )
 
-  @test_util.run_in_graph_and_eager_modes
-  def testResourceVariableUninitialized(self):
-    var_uninit_f32 = gen_resource_variable_ops.var_handle_op(
-        dtype=dtypes.float32,
-        shape=[1, 2],
-        shared_name=(
-            context.anonymous_name() if context.executing_eagerly() else ""
-        ),
-    )
-    var_uninit_f16 = gen_resource_variable_ops.var_handle_op(
-        dtype=dtypes.float16,
-        shape=[1, 2],
-        shared_name=(
-            context.anonymous_name() if context.executing_eagerly() else ""
-        ),
-    )
-    accum_uninit_f16 = gen_resource_variable_ops.var_handle_op(
-        dtype=dtypes.float16,
-        shape=[1, 2],
-        shared_name=(
-            context.anonymous_name() if context.executing_eagerly() else ""
-        ),
-    )
-    lr = constant_op.constant(0.1, dtype=dtypes.float16)
-    grad = constant_op.constant([[0.1, 0.2]], dtype=dtypes.float16)
-    indices = constant_op.constant([0], dtype=dtypes.int32)
-
-    # Mismatched dtype should raise InvalidArgumentError
-    with self.assertRaisesRegex(
-        errors.InvalidArgumentError, "dtype mismatch"
-    ):
-      self.evaluate(
-          gen_training_ops.resource_apply_gradient_descent(
-              var=var_uninit_f32,
-              alpha=lr,
-              delta=grad,
-          )
-      )
-
-    with self.assertRaisesRegex(
-        errors.InvalidArgumentError, "dtype mismatch"
-    ):
-      self.evaluate(
-          gen_training_ops.resource_sparse_apply_adagrad(
-              var=var_uninit_f32,
-              accum=accum_uninit_f16,
-              lr=lr,
-              grad=grad,
-              indices=indices,
-          )
-      )
-
-    # Matched dtype should raise FailedPreconditionError
-    with self.assertRaisesRegex(
-        errors.FailedPreconditionError, "uninitialized"
-    ):
-      self.evaluate(
-          gen_training_ops.resource_apply_gradient_descent(
-              var=var_uninit_f16,
-              alpha=lr,
-              delta=grad,
-          )
-      )
-
-    with self.assertRaisesRegex(
-        errors.FailedPreconditionError, "uninitialized"
-    ):
-      self.evaluate(
-          gen_training_ops.resource_sparse_apply_adagrad(
-              var=var_uninit_f16,
-              accum=accum_uninit_f16,
-              lr=lr,
-              grad=grad,
-              indices=indices,
-          )
-      )
-
-    var_uninit_scalar_f16 = gen_resource_variable_ops.var_handle_op(
-        dtype=dtypes.float16,
-        shape=[],
-        shared_name=(
-            context.anonymous_name() if context.executing_eagerly() else ""
-        ),
-    )
-    with self.assertRaisesRegex(
-        errors.FailedPreconditionError, "uninitialized"
-    ):
-      self.evaluate(
-          gen_resource_variable_ops.assign_add_variable_op(
-              resource=var_uninit_scalar_f16,
-              value=constant_op.constant(0.1, dtype=dtypes.float16),
-          )
-      )
-
-
-    scatter_var_uninit_f16 = gen_resource_variable_ops.var_handle_op(
-        dtype=dtypes.float16,
-        shape=[1, 2],
-        shared_name=(
-            context.anonymous_name() if context.executing_eagerly() else ""
-        ),
-    )
-    with self.assertRaisesRegex(
-        errors.FailedPreconditionError, "uninitialized"
-    ):
-      self.evaluate(
-          gen_resource_variable_ops.resource_scatter_update(
-              resource=scatter_var_uninit_f16,
-              indices=indices,
-              updates=grad,
-          )
-      )
-
-    gather_var_uninit_f16 = gen_resource_variable_ops.var_handle_op(
-        dtype=dtypes.float16,
-        shape=[1, 2],
-        shared_name=(
-            context.anonymous_name() if context.executing_eagerly() else ""
-        ),
-    )
-    with self.assertRaisesRegex(
-        errors.FailedPreconditionError, "uninitialized"
-    ):
-      self.evaluate(
-          gen_resource_variable_ops.resource_gather(
-              resource=gather_var_uninit_f16,
-              indices=indices,
-              dtype=dtypes.float16,
-          )
-      )
-
-    gather_nd_var_uninit_f16 = gen_resource_variable_ops.var_handle_op(
-        dtype=dtypes.float16,
-        shape=[1, 2],
-        shared_name=(
-            context.anonymous_name() if context.executing_eagerly() else ""
-        ),
-    )
-    with self.assertRaisesRegex(
-        errors.FailedPreconditionError, "uninitialized"
-    ):
-      self.evaluate(
-          gen_resource_variable_ops.resource_gather_nd(
-              resource=gather_nd_var_uninit_f16,
-              indices=constant_op.constant([[0]], dtype=dtypes.int32),
-              dtype=dtypes.float16,
-          )
-      )
-
-    scatter_nd_var_uninit_f16 = gen_resource_variable_ops.var_handle_op(
-        dtype=dtypes.float16,
-        shape=[1, 2],
-        shared_name=(
-            context.anonymous_name() if context.executing_eagerly() else ""
-        ),
-    )
-    with self.assertRaisesRegex(
-        errors.FailedPreconditionError, "uninitialized"
-    ):
-      self.evaluate(
-          gen_state_ops.resource_scatter_nd_update(
-              ref=scatter_nd_var_uninit_f16,
-              indices=constant_op.constant([[0]], dtype=dtypes.int32),
-              updates=grad,
-          )
-      )
-
 
 if __name__ == '__main__':
   googletest.main()

From 447c167eac1987936e489abc5313c12ef8aa906b Mon Sep 17 00:00:00 2001
From: Esun Kim <veblush@google.com>
Date: Wed, 14 Oct 2026 20:29:36 -0700
Subject: [PATCH 0824/1000] Fix float-to-double narrowing in TFLite quantized
 Mul kernel.

Cast input1->params.scale, input2->params.scale, and output->params.scale to double before computing real_multiplier in Prepare so that QuantizeMultiplier receives full double precision rather than a float-narrowed intermediate value.

PiperOrigin-RevId: 999814764
---
 tensorflow/lite/kernels/mul.cc | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/tensorflow/lite/kernels/mul.cc b/tensorflow/lite/kernels/mul.cc
index 73ffd3488f9..e2d516ac733 100644
--- a/tensorflow/lite/kernels/mul.cc
+++ b/tensorflow/lite/kernels/mul.cc
@@ -122,8 +122,9 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     TF_LITE_ENSURE_STATUS(CalculateActivationRangeQuantized(
         context, params->activation, output, &data->output_activation_min,
         &data->output_activation_max));
-    double real_multiplier =
-        input1->params.scale * input2->params.scale / output->params.scale;
+    double real_multiplier = static_cast<double>(input1->params.scale) *
+                             static_cast<double>(input2->params.scale) /
+                             static_cast<double>(output->params.scale);
     QuantizeMultiplier(real_multiplier, &data->output_multiplier,
                        &data->output_shift);
   }

From 5703b56be47de728fcad51fb2fb29f41c6799fe2 Mon Sep 17 00:00:00 2001
From: Esun Kim <veblush@google.com>
Date: Wed, 14 Oct 2026 20:30:01 -0700
Subject: [PATCH 0825/1000] Fix float-to-double narrowing in TFLite reference
 kernels.

Corrects a numerical precision issue in several TFLite reference kernels where intermediate scale variables intended for double precision calculations were evaluated in float precision before promotion to double.

PiperOrigin-RevId: 999814910
---
 tensorflow/lite/core/c/common.h               | 13 ++++++----
 tensorflow/lite/kernels/activations.cc        | 24 ++++++++++++-------
 tensorflow/lite/kernels/add.cc                | 11 +++++----
 tensorflow/lite/kernels/div.cc                |  6 ++---
 tensorflow/lite/kernels/elementwise.cc        |  7 ++++--
 tensorflow/lite/kernels/lstm.cc               |  3 ++-
 tensorflow/lite/kernels/squared_difference.cc | 12 ++++++----
 tensorflow/lite/kernels/sub.cc                | 13 ++++++----
 tensorflow/lite/kernels/svdf.cc               | 14 ++++++-----
 9 files changed, 62 insertions(+), 41 deletions(-)

diff --git a/tensorflow/lite/core/c/common.h b/tensorflow/lite/core/c/common.h
index 66c8bf6fad6..ffc910d986c 100644
--- a/tensorflow/lite/core/c/common.h
+++ b/tensorflow/lite/core/c/common.h
@@ -269,11 +269,14 @@ void TfLiteFloatArrayFree(TfLiteFloatArray* a);
 
 #define TF_LITE_ENSURE_NEAR(context, a, b, epsilon)                          \
   do {                                                                       \
-    auto delta = ((a) > (b)) ? ((a) - (b)) : ((b) - (a));                    \
-    if (delta > epsilon) {                                                   \
-      TF_LITE_KERNEL_LOG((context), "%s:%d %s not near %s (%f != %f)",       \
-                         __FILE__, __LINE__, #a, #b, static_cast<double>(a), \
-                         static_cast<double>(b));                            \
+    const auto _val_a = (a);                                                 \
+    const auto _val_b = (b);                                                 \
+    const auto _eps = (epsilon);                                             \
+    const auto _delta = _val_a - _val_b;                                     \
+    if (_delta > _eps || _delta < -_eps) {                                   \
+      TF_LITE_KERNEL_LOG(                                                    \
+          (context), "%s:%d %s not near %s (%f != %f)", __FILE__, __LINE__,  \
+          #a, #b, static_cast<double>(_val_a), static_cast<double>(_val_b)); \
       return kTfLiteError;                                                   \
     }                                                                        \
   } while (0)
diff --git a/tensorflow/lite/kernels/activations.cc b/tensorflow/lite/kernels/activations.cc
index 20ffc244a31..b6d7c4cdbf8 100644
--- a/tensorflow/lite/kernels/activations.cc
+++ b/tensorflow/lite/kernels/activations.cc
@@ -231,7 +231,8 @@ TfLiteStatus ReluPrepare(TfLiteContext* context, TfLiteNode* node) {
 
   if (input->type == kTfLiteInt8 || input->type == kTfLiteUInt8 ||
       input->type == kTfLiteInt16) {
-    double real_multiplier = input->params.scale / output->params.scale;
+    double real_multiplier = static_cast<double>(input->params.scale) /
+                             static_cast<double>(output->params.scale);
     QuantizeMultiplier(real_multiplier, &data->output_multiplier,
                        &data->output_shift);
   }
@@ -315,11 +316,13 @@ TfLiteStatus LeakyReluPrepare(TfLiteContext* context, TfLiteNode* node) {
     const auto* params =
         reinterpret_cast<TfLiteLeakyReluParams*>(node->builtin_data);
 
-    double alpha_multiplier =
-        input->params.scale * params->alpha / output->params.scale;
+    double alpha_multiplier = static_cast<double>(input->params.scale) *
+                              static_cast<double>(params->alpha) /
+                              static_cast<double>(output->params.scale);
     QuantizeMultiplier(alpha_multiplier, &data->output_multiplier_alpha,
                        &data->output_shift_alpha);
-    double identity_multiplier = input->params.scale / output->params.scale;
+    double identity_multiplier = static_cast<double>(input->params.scale) /
+                                 static_cast<double>(output->params.scale);
     QuantizeMultiplier(identity_multiplier, &data->output_multiplier_identity,
                        &data->output_shift_identity);
   }
@@ -641,8 +644,9 @@ TfLiteStatus SoftmaxPrepare(TfLiteContext* context, TfLiteNode* node) {
     data->params.zero_point = output->params.zero_point;
     data->params.scale = output->params.scale;
 
-    double input_scale_beta_rescale =
-        input->params.scale * params->beta /
+    const double input_scale_beta_rescale =
+        static_cast<double>(input->params.scale) *
+        static_cast<double>(params->beta) /
         (10.0 / 65535.0);  // scale the input_diff such that [-65535, 0]
                            // correspond to [-10.0, 0.0]
     QuantizeMultiplier(input_scale_beta_rescale, &data->params.input_multiplier,
@@ -736,9 +740,11 @@ TfLiteStatus PreluPrepare(TfLiteContext* context, TfLiteNode* node) {
     // output real multiplier 1 is input_scale / output_scale;
     // for input_q - input_zp < 0:
     // output real multiplier 2 is input_scale  * alpha_scale/ output_scale.
-    double real_multiplier_1 = input->params.scale / output->params.scale;
-    double real_multiplier_2 =
-        input->params.scale * alpha->params.scale / output->params.scale;
+    double real_multiplier_1 = static_cast<double>(input->params.scale) /
+                               static_cast<double>(output->params.scale);
+    double real_multiplier_2 = static_cast<double>(input->params.scale) *
+                               static_cast<double>(alpha->params.scale) /
+                               static_cast<double>(output->params.scale);
     QuantizeMultiplier(real_multiplier_1, &data->output_multiplier_1,
                        &data->output_shift_1);
     QuantizeMultiplier(real_multiplier_2, &data->output_multiplier_2,
diff --git a/tensorflow/lite/kernels/add.cc b/tensorflow/lite/kernels/add.cc
index 9e5dfcd228c..ee61fd93de3 100644
--- a/tensorflow/lite/kernels/add.cc
+++ b/tensorflow/lite/kernels/add.cc
@@ -181,14 +181,15 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     // therefore the addition will still fit in a 32 bit accumulator.
     data->left_shift = general_scale_int16 ? 15 : 20;
     const double twice_max_input_scale =
-        2 * std::max(input1->params.scale, input2->params.scale);
+        2 * static_cast<double>(
+                std::max(input1->params.scale, input2->params.scale));
     const double real_input1_multiplier =
-        input1->params.scale / twice_max_input_scale;
+        static_cast<double>(input1->params.scale) / twice_max_input_scale;
     const double real_input2_multiplier =
-        input2->params.scale / twice_max_input_scale;
+        static_cast<double>(input2->params.scale) / twice_max_input_scale;
     const double real_output_multiplier =
-        twice_max_input_scale /
-        ((1 << data->left_shift) * output->params.scale);
+        twice_max_input_scale / (static_cast<double>(1 << data->left_shift) *
+                                 static_cast<double>(output->params.scale));
 
     QuantizeMultiplierSmallerThanOneExp(
         real_input1_multiplier, &data->input1_multiplier, &data->input1_shift);
diff --git a/tensorflow/lite/kernels/div.cc b/tensorflow/lite/kernels/div.cc
index d7cebed68f9..c5a59cd2494 100644
--- a/tensorflow/lite/kernels/div.cc
+++ b/tensorflow/lite/kernels/div.cc
@@ -104,8 +104,9 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     TF_LITE_ENSURE_STATUS(CalculateActivationRangeQuantized(
         context, params->activation, output, &data->output_activation_min,
         &data->output_activation_max));
-    const double real_multiplier =
-        input1->params.scale / (input2->params.scale * output->params.scale);
+    const double real_multiplier = static_cast<double>(input1->params.scale) /
+                                   (static_cast<double>(input2->params.scale) *
+                                    static_cast<double>(output->params.scale));
     QuantizeMultiplier(real_multiplier, &data->output_multiplier,
                        &data->output_shift);
   }
@@ -279,7 +280,6 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_OK(context,
                     GetOutputSafe(context, node, kOutputTensor, &output));
 
-
   if (output->type == kTfLiteFloat32) {
     // Div by zero seems ok in this case, we don't do a check at this point.
     // However, unlike in TF where infinities are returned, here we return an
diff --git a/tensorflow/lite/kernels/elementwise.cc b/tensorflow/lite/kernels/elementwise.cc
index 0bff391ffb3..cead70011b2 100644
--- a/tensorflow/lite/kernels/elementwise.cc
+++ b/tensorflow/lite/kernels/elementwise.cc
@@ -86,13 +86,16 @@ bool IsLogSupportedType(const TfLiteType type) {
 inline void SetAbsOutputMultiplier(const float input_scale,
                                    const float output_scale,
                                    int32_t* multiplier, int32_t* shift) {
-  QuantizeMultiplier(input_scale / output_scale, multiplier, shift);
+  QuantizeMultiplier(
+      static_cast<double>(input_scale) / static_cast<double>(output_scale),
+      multiplier, shift);
 }
 
 inline void SetRsqrtOutputMultiplier(const float input_scale,
                                      const float output_scale,
                                      int32_t* multiplier, int32_t* shift) {
-  const double scale = 1. / (std::sqrt(input_scale) * output_scale);
+  const double scale = 1.0 / (std::sqrt(static_cast<double>(input_scale)) *
+                              static_cast<double>(output_scale));
   QuantizeMultiplier(scale, multiplier, shift);
 }
 
diff --git a/tensorflow/lite/kernels/lstm.cc b/tensorflow/lite/kernels/lstm.cc
index 10dde7f4987..2c1114d48ce 100644
--- a/tensorflow/lite/kernels/lstm.cc
+++ b/tensorflow/lite/kernels/lstm.cc
@@ -2325,7 +2325,8 @@ TfLiteStatus Eval(TfLiteContext* context, TfLiteNode* node) {
       return kTfLiteError;
     }
 
-    double real_accum_multiplier = 4096 * bias->params.scale;
+    double real_accum_multiplier =
+        4096.0 * static_cast<double>(bias->params.scale);
     int32 accum_multiplier;
     int accum_shift;
     tflite::QuantizeMultiplier(real_accum_multiplier, &accum_multiplier,
diff --git a/tensorflow/lite/kernels/squared_difference.cc b/tensorflow/lite/kernels/squared_difference.cc
index 57bf92cc37c..21597d0f424 100644
--- a/tensorflow/lite/kernels/squared_difference.cc
+++ b/tensorflow/lite/kernels/squared_difference.cc
@@ -106,16 +106,18 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     // shift to make integer for scales.
     data->arithmetic_params.left_shift = 7;
     const double twice_max_input_scale =
-        2 * std::max(input1_quantization_params.scale,
-                     input2_quantization_params.scale);
+        2 * static_cast<double>(std::max(input1_quantization_params.scale,
+                                         input2_quantization_params.scale));
     const double real_input1_multiplier =
-        input1_quantization_params.scale / twice_max_input_scale;
+        static_cast<double>(input1_quantization_params.scale) /
+        twice_max_input_scale;
     double real_input2_multiplier =
-        input2_quantization_params.scale / twice_max_input_scale;
+        static_cast<double>(input2_quantization_params.scale) /
+        twice_max_input_scale;
     const double real_output_multiplier =
         (twice_max_input_scale * twice_max_input_scale) /
         ((1 << data->arithmetic_params.left_shift * 2) *
-         output_quantization_params.scale);
+         static_cast<double>(output_quantization_params.scale));
     tflite::QuantizeMultiplierSmallerThanOneExp(
         real_input1_multiplier, &data->arithmetic_params.input1_multiplier,
         &data->arithmetic_params.input1_shift);
diff --git a/tensorflow/lite/kernels/sub.cc b/tensorflow/lite/kernels/sub.cc
index 0136e94ab21..7cb5508227a 100644
--- a/tensorflow/lite/kernels/sub.cc
+++ b/tensorflow/lite/kernels/sub.cc
@@ -136,15 +136,18 @@ TfLiteStatus PrepareGeneralSubOp(TfLiteContext* context,
   // 31, therefore the addition will still fit in a 32 bit accumulator.
   op_params->left_shift = output->type == kTfLiteInt16 ? 15 : 20;
   const double twice_max_input_scale =
-      2 * std::max(input1_quantization_params.scale,
-                   input2_quantization_params.scale);
+      2 * static_cast<double>(std::max(input1_quantization_params.scale,
+                                       input2_quantization_params.scale));
   const double real_input1_multiplier =
-      input1_quantization_params.scale / twice_max_input_scale;
+      static_cast<double>(input1_quantization_params.scale) /
+      twice_max_input_scale;
   const double real_input2_multiplier =
-      input2_quantization_params.scale / twice_max_input_scale;
+      static_cast<double>(input2_quantization_params.scale) /
+      twice_max_input_scale;
   const double real_output_multiplier =
       twice_max_input_scale /
-      ((1 << op_params->left_shift) * output_quantization_params.scale);
+      ((1 << op_params->left_shift) *
+       static_cast<double>(output_quantization_params.scale));
 
   tflite::QuantizeMultiplierSmallerThanOneExp(real_input1_multiplier,
                                               &op_params->input1_multiplier,
diff --git a/tensorflow/lite/kernels/svdf.cc b/tensorflow/lite/kernels/svdf.cc
index 74a38153631..5173322f260 100644
--- a/tensorflow/lite/kernels/svdf.cc
+++ b/tensorflow/lite/kernels/svdf.cc
@@ -275,12 +275,14 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
     TF_LITE_ENSURE(context, output->quantization.type != kTfLiteNoQuantization);
     auto* output_params = reinterpret_cast<TfLiteAffineQuantization*>(
         output->quantization.params);
-    const double effective_scale_1 = input_params->scale->data[0] *
-                                     weights_feature_params->scale->data[0] /
-                                     state_params->scale->data[0];
-    const double effective_scale_2 = state_params->scale->data[0] *
-                                     weight_time_params->scale->data[0] /
-                                     output_params->scale->data[0];
+    const double effective_scale_1 =
+        static_cast<double>(input_params->scale->data[0]) *
+        static_cast<double>(weights_feature_params->scale->data[0]) /
+        static_cast<double>(state_params->scale->data[0]);
+    const double effective_scale_2 =
+        static_cast<double>(state_params->scale->data[0]) *
+        static_cast<double>(weight_time_params->scale->data[0]) /
+        static_cast<double>(output_params->scale->data[0]);
     QuantizeMultiplier(effective_scale_1, &op_data->effective_scale_1_a,
                        &op_data->effective_scale_1_b);
     QuantizeMultiplier(effective_scale_2, &op_data->effective_scale_2_a,

From 2a4f4d85539ad49581157655878de046c0059eb1 Mon Sep 17 00:00:00 2001
From: Esun Kim <veblush@google.com>
Date: Wed, 14 Oct 2026 20:30:16 -0700
Subject: [PATCH 0826/1000] Fix float-to-double narrowing in
 GetQuantizedConvolutionMultipler for int8/int16.

Evaluate input->params.scale * filter->params.scale in double precision for non-uint8 quantized tensors in the 5-argument GetQuantizedConvolutionMultipler overload (matching the 6-argument overload), while preserving legacy float scale multiplication for deprecated kTfLiteUInt8 quantization to maintain backward compatibility with existing golden tests.

PiperOrigin-RevId: 999815030
---
 tensorflow/lite/kernels/kernel_util.cc | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/tensorflow/lite/kernels/kernel_util.cc b/tensorflow/lite/kernels/kernel_util.cc
index 36ad5178158..a4757abb0aa 100644
--- a/tensorflow/lite/kernels/kernel_util.cc
+++ b/tensorflow/lite/kernels/kernel_util.cc
@@ -330,8 +330,13 @@ TfLiteStatus GetQuantizedConvolutionMultipler(TfLiteContext* context,
                                               const TfLiteTensor* filter,
                                               TfLiteTensor* output,
                                               double* multiplier) {
+  // Keep float scale multiplication for legacy uint8 quantization for backward
+  // compatibility with existing golden tests.
   const double input_product_scale =
-      static_cast<double>(input->params.scale * filter->params.scale);
+      (input->type == kTfLiteUInt8)
+          ? static_cast<double>(input->params.scale * filter->params.scale)
+          : static_cast<double>(input->params.scale) *
+                static_cast<double>(filter->params.scale);
   TF_LITE_ENSURE(context, input_product_scale >= 0);
   *multiplier = input_product_scale / static_cast<double>(output->params.scale);
 

From b2cddb6e283350085d5bb17560b8e834622bfbaf Mon Sep 17 00:00:00 2001
From: "Patrick C. Toulme" <ptoulme@google.com>
Date: Wed, 14 Oct 2026 20:43:22 -0700
Subject: [PATCH 0827/1000] [XLA] Fill elided HLO constants per parse, not per
 process. The HLO parser filled elided constants (`constant({...})`) from
 function local static counters, so the same text parsed to different literals
 depending on earlier parses in the process, and concurrent parses raced on
 the counters. The counters are now members of the parser, so every parse of
 the same text gives the same literals while constants within one module still
 differ. This also fixes the fill of the bytes after the last whole 32 bit
 word, which overwrote bytes of the first words and left the last bytes zero.

PiperOrigin-RevId: 999820137
---
 third_party/xla/xla/hlo/parser/BUILD          |   3 +
 third_party/xla/xla/hlo/parser/hlo_parser.cc  |  28 +++--
 .../xla/xla/hlo/parser/hlo_parser_test.cc     | 108 ++++++++++++++++++
 3 files changed, 129 insertions(+), 10 deletions(-)

diff --git a/third_party/xla/xla/hlo/parser/BUILD b/third_party/xla/xla/hlo/parser/BUILD
index 1d0bbd1a6dd..c836adbb5c2 100644
--- a/third_party/xla/xla/hlo/parser/BUILD
+++ b/third_party/xla/xla/hlo/parser/BUILD
@@ -80,6 +80,7 @@ xla_cc_test(
         ":hlo_parser",
         "//xla:array",
         "//xla:comparison_util",
+        "//xla:literal",
         "//xla:shape_util",
         "//xla:window_util",
         "//xla:xla_data_proto_cc",
@@ -92,6 +93,7 @@ xla_cc_test(
         "//xla/service:hlo_proto_cc",
         "//xla/service:pattern_matcher",
         "//xla/tsl/lib/core:status_test_util",
+        "//xla/tsl/platform:env",
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:test",
         "//xla/tsl/platform:test_main",
@@ -102,6 +104,7 @@ xla_cc_test(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest",
     ],
diff --git a/third_party/xla/xla/hlo/parser/hlo_parser.cc b/third_party/xla/xla/hlo/parser/hlo_parser.cc
index 757db3152d6..63ae942c349 100644
--- a/third_party/xla/xla/hlo/parser/hlo_parser.cc
+++ b/third_party/xla/xla/hlo/parser/hlo_parser.cc
@@ -764,6 +764,11 @@ class HloParserImpl : public HloParser {
   // Tracks recursion depth to prevent stack overflow from deeply nested input.
   int current_recursion_depth_ = 0;
   StackFrameIndexProto stack_frame_index_proto_;
+
+  // Next values used to fill elided constants (`{...}`) when
+  // fill_shortform_constants_with_random_values() is set.
+  uint32_t shortform_fill_word_ = 0;
+  uint8_t shortform_fill_byte_ = 0;
 };
 
 bool SplitToInt64s(absl::string_view s, char delim, std::vector<int64_t>* out) {
@@ -5754,11 +5759,13 @@ bool HloParserImpl::ParseDenseLiteral(Literal* literal, const Shape& shape) {
         if (!options_.fill_shortform_constants_with_random_values()) {
           break;
         }
-        // Fill data with deterministic (garbage) values. Use static to avoid
-        // creating identical constants which could potentially got CSE'ed
-        // away. This is a best-effort approach to make sure replaying a HLO
-        // gives us same optimized HLO graph.
-        static uint32_t data = 0;
+        // Fill data with deterministic (garbage) values. The counters keep
+        // running across the elided constants of one parse, so constants of
+        // the same shape get different values and are not CSE'ed away. They
+        // are parser members rather than statics, so every parse of the same
+        // text produces the same values, regardless of what was parsed before
+        // or on which thread. This is a best effort approach to make sure
+        // replaying an HLO gives us the same optimized HLO graph.
 
         // According to the System V ABI not all 8 bit values are valid booleans
         // - only the values 0 and 1 are allowed. So to avoid undefined
@@ -5772,13 +5779,14 @@ bool HloParserImpl::ParseDenseLiteral(Literal* literal, const Shape& shape) {
             (shape.element_type() == PRED) ? kBooleanMask : kNoMask;
 
         uint32_t* raw_data = static_cast<uint32_t*>(literal->untyped_data());
-        for (int64_t i = 0; i < literal->size_bytes() / 4; ++i) {
-          raw_data[i] = data++ & mask;
+        const int64_t num_words = literal->size_bytes() / 4;
+        for (int64_t i = 0; i < num_words; ++i) {
+          raw_data[i] = shortform_fill_word_++ & mask;
         }
+        // The bytes after the last whole word.
         uint8_t* raw_data_int8 = static_cast<uint8_t*>(literal->untyped_data());
-        static uint8_t data_int8 = 0;
-        for (int64_t i = 0; i < literal->size_bytes() % 4; ++i) {
-          raw_data_int8[literal->size_bytes() / 4 + i] = data_int8++ & mask;
+        for (int64_t i = num_words * 4; i < literal->size_bytes(); ++i) {
+          raw_data_int8[i] = shortform_fill_byte_++ & mask;
         }
         break;
       }
diff --git a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc
index 9fc0944aba3..f357f37292a 100644
--- a/third_party/xla/xla/hlo/parser/hlo_parser_test.cc
+++ b/third_party/xla/xla/hlo/parser/hlo_parser_test.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "xla/hlo/parser/hlo_parser.h"
 
 #include <cstdint>
+#include <cstring>
 #include <memory>
 #include <string>
 #include <utility>
@@ -32,6 +33,7 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
+#include "absl/synchronization/notification.h"
 #include "absl/types/span.h"
 #include "xla/array.h"
 #include "xla/comparison_util.h"
@@ -50,14 +52,17 @@ limitations under the License.
 #include "xla/hlo/testlib/verified_hlo_module.h"
 #include "xla/layout.h"
 #include "xla/layout_util.h"
+#include "xla/literal.h"
 #include "xla/service/hlo.pb.h"
 #include "xla/service/hlo_module_config.h"
 #include "xla/service/pattern_matcher.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/tsl/lib/core/status_test_util.h"
+#include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/test.h"
+#include "xla/tsl/platform/threadpool.h"
 #include "xla/tsl/util/proto/proto_matchers.h"
 #include "xla/window_util.h"
 #include "xla/xla_data.pb.h"
@@ -4055,6 +4060,109 @@ ENTRY %ShortConstant.v4 () -> f32[67,89] {
   EXPECT_EQ(result.value()->ToString(HloPrintOptions()), original);
 }
 
+// Returns the literals of the entry computation's constants, in parse order.
+std::vector<Literal> EntryConstantLiterals(const HloModule& module) {
+  std::vector<Literal> literals;
+  for (const HloInstruction* instruction :
+       module.entry_computation()->instructions()) {
+    if (instruction->opcode() == HloOpcode::kConstant) {
+      literals.push_back(instruction->literal().Clone());
+    }
+  }
+  return literals;
+}
+
+std::vector<uint8_t> RawBytes(const Literal& literal) {
+  const auto* data = static_cast<const uint8_t*>(literal.untyped_data());
+  return std::vector<uint8_t>(data, data + literal.size_bytes());
+}
+
+constexpr absl::string_view kElidedConstantsModule = R"(
+HloModule elided_constants
+
+ENTRY main {
+  a = s32[4] constant({...})
+  b = s32[4] constant({...})
+  c = s8[3] constant({...})
+  d = s8[3] constant({...})
+  ROOT t = (s32[4], s32[4], s8[3], s8[3]) tuple(a, b, c, d)
+}
+)";
+
+TEST_F(HloParserTest, ElidedConstantsAreTheSameInEveryParse) {
+  ASSERT_OK_AND_ASSIGN(auto first,
+                       ParseAndReturnUnverifiedModule(kElidedConstantsModule));
+  ASSERT_OK_AND_ASSIGN(auto second,
+                       ParseAndReturnUnverifiedModule(kElidedConstantsModule));
+  const std::vector<Literal> literals = EntryConstantLiterals(*first);
+  ASSERT_EQ(literals.size(), 4);
+  // Elided constants of one module still differ, so CSE cannot merge them.
+  EXPECT_NE(literals[0], literals[1]);
+  EXPECT_NE(literals[2], literals[3]);
+  EXPECT_EQ(EntryConstantLiterals(*second), literals);
+}
+
+TEST_F(HloParserTest, ElidedConstantFillsTheBytesAfterItsLastWord) {
+  constexpr absl::string_view kHlo = R"(
+HloModule elided_constant_tails
+
+ENTRY main {
+  a = s32[2] constant({...})
+  b = s8[7] constant({...})
+  c = bf16[3] constant({...})
+  d = pred[5] constant({...})
+  ROOT t = (s32[2], s8[7], bf16[3], pred[5]) tuple(a, b, c, d)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(kHlo));
+  const std::vector<Literal> literals = EntryConstantLiterals(*module);
+  ASSERT_EQ(literals.size(), 4);
+  // Whole 32 bit words count 0, 1, 2, ... in host byte order. The bytes after
+  // the last whole word count 0, 1, 2, ... separately.
+  auto words_then_bytes = [](std::vector<uint32_t> words,
+                             std::vector<uint8_t> tail) {
+    std::vector<uint8_t> bytes(words.size() * sizeof(uint32_t));
+    std::memcpy(bytes.data(), words.data(), bytes.size());
+    bytes.insert(bytes.end(), tail.begin(), tail.end());
+    return bytes;
+  };
+  EXPECT_EQ(RawBytes(literals[0]), words_then_bytes({0, 1}, {}));
+  EXPECT_EQ(RawBytes(literals[1]), words_then_bytes({2}, {0, 1, 2}));
+  EXPECT_EQ(RawBytes(literals[2]), words_then_bytes({3}, {3, 4}));
+  // PRED keeps only the lowest bit of each byte: word 4 and byte 5 become 0
+  // and 1.
+  EXPECT_EQ(RawBytes(literals[3]), words_then_bytes({0}, {1}));
+}
+
+TEST_F(HloParserTest, ConcurrentParsesFillElidedConstantsAlike) {
+  ASSERT_OK_AND_ASSIGN(auto reference,
+                       ParseAndReturnUnverifiedModule(kElidedConstantsModule));
+  const std::vector<Literal> expected = EntryConstantLiterals(*reference);
+
+  constexpr int kNumThreads = 8;
+  constexpr int kParsesPerThread = 4;
+  std::vector<absl::StatusOr<std::unique_ptr<HloModule>>> modules(
+      kNumThreads * kParsesPerThread);
+  {
+    absl::Notification start;
+    tsl::thread::ThreadPool pool(tsl::Env::Default(), "parse", kNumThreads);
+    for (int t = 0; t < kNumThreads; ++t) {
+      pool.Schedule([&, t] {
+        start.WaitForNotification();
+        for (int i = 0; i < kParsesPerThread; ++i) {
+          modules[t * kParsesPerThread + i] =
+              ParseAndReturnUnverifiedModule(kElidedConstantsModule);
+        }
+      });
+    }
+    start.Notify();
+  }  // The pool's destructor waits for every parse.
+  for (const absl::StatusOr<std::unique_ptr<HloModule>>& module : modules) {
+    ASSERT_OK(module);
+    EXPECT_EQ(EntryConstantLiterals(**module), expected);
+  }
+}
+
 TEST_F(HloParserTest, NegativeNan) {
   const std::string original =
       R"(HloModule NegativeNan_module, entry_computation_layout={()->bf16[2]{0}}

From f79e44d86ccd930acb10512d84d642edd18c9614 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2026 21:01:18 -0700
Subject: [PATCH 0828/1000] Automated Code Change

PiperOrigin-RevId: 999827465
---
 tensorflow/lite/tools/BUILD                       | 3 ++-
 tensorflow/lite/tools/gen_op_registration.cc      | 4 ++--
 tensorflow/lite/tools/gen_op_registration_main.cc | 2 --
 tensorflow/lite/tools/gen_op_registration_test.cc | 5 +++++
 tensorflow/lite/tools/list_flex_ops.cc            | 6 +++---
 tensorflow/lite/tools/list_flex_ops_main.cc       | 1 +
 tensorflow/lite/tools/list_flex_ops_no_kernel.cc  | 1 +
 tensorflow/lite/tools/list_flex_ops_test.cc       | 4 +++-
 tensorflow/lite/tools/model_loader_test.cc        | 2 +-
 9 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/tensorflow/lite/tools/BUILD b/tensorflow/lite/tools/BUILD
index c329ceccd37..9e4bfa6b725 100644
--- a/tensorflow/lite/tools/BUILD
+++ b/tensorflow/lite/tools/BUILD
@@ -329,6 +329,7 @@ cc_library(
         "//tensorflow/lite:util",
         "//tensorflow/lite/core:framework",
         "//tensorflow/lite/schema:schema_utils",
+        "@com_google_absl//absl/log",
         "@flatbuffers",
         "@jsoncpp_git//:jsoncpp",
     ],
@@ -401,6 +402,7 @@ tf_cc_test(
         "//tensorflow/core/platform:protobuf",
         "//tensorflow/core/platform:resource_loader",
         "//tensorflow/lite/kernels:test_util",
+        "//third_party/protobuf",
         "@com_google_googletest//:gtest_main",
         "@flatbuffers",
     ],
@@ -465,7 +467,6 @@ cc_test(
         ":model_loader",
         "//tensorflow/lite:model_builder",
         "//tensorflow/lite/schema:schema_fbs",
-        "@com_google_absl//absl/strings:str_format",
         "@com_google_googletest//:gtest_main",
         "@flatbuffers//:runtime_cc",
     ],
diff --git a/tensorflow/lite/tools/gen_op_registration.cc b/tensorflow/lite/tools/gen_op_registration.cc
index 832469fb526..00a28305093 100644
--- a/tensorflow/lite/tools/gen_op_registration.cc
+++ b/tensorflow/lite/tools/gen_op_registration.cc
@@ -15,8 +15,8 @@ limitations under the License.
 #include "tensorflow/lite/tools/gen_op_registration.h"
 
 #include <algorithm>
-#include <string>
-#include <vector>
+#include <cctype>
+#include <utility>
 
 #include "re2/re2.h"
 #include "tensorflow/lite/core/model.h"
diff --git a/tensorflow/lite/tools/gen_op_registration_main.cc b/tensorflow/lite/tools/gen_op_registration_main.cc
index c959945b334..42213053961 100644
--- a/tensorflow/lite/tools/gen_op_registration_main.cc
+++ b/tensorflow/lite/tools/gen_op_registration_main.cc
@@ -14,13 +14,11 @@ limitations under the License.
 ==============================================================================*/
 
 #include <fstream>
-#include <map>
 #include <sstream>
 #include <string>
 #include <vector>
 
 #include "absl/strings/str_split.h"
-#include "absl/strings/strip.h"
 #include "tensorflow/lite/tools/command_line_flags.h"
 #include "tensorflow/lite/tools/gen_op_registration.h"
 #include "tensorflow/lite/util.h"
diff --git a/tensorflow/lite/tools/gen_op_registration_test.cc b/tensorflow/lite/tools/gen_op_registration_test.cc
index f4ca9f11609..f7885bd2f4b 100644
--- a/tensorflow/lite/tools/gen_op_registration_test.cc
+++ b/tensorflow/lite/tools/gen_op_registration_test.cc
@@ -14,6 +14,11 @@ limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/lite/tools/gen_op_registration.h"
+
+#include <map>
+#include <utility>
+#include <vector>
+
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
diff --git a/tensorflow/lite/tools/list_flex_ops.cc b/tensorflow/lite/tools/list_flex_ops.cc
index 188efc57f90..17e514faf1a 100644
--- a/tensorflow/lite/tools/list_flex_ops.cc
+++ b/tensorflow/lite/tools/list_flex_ops.cc
@@ -14,11 +14,11 @@ limitations under the License.
 ==============================================================================*/
 #include "tensorflow/lite/tools/list_flex_ops.h"
 
-#include <fstream>
-#include <sstream>
+#include <cstdint>
+#include <cstring>
 #include <string>
-#include <vector>
 
+#include "absl/log/log.h"
 #include "flatbuffers/flexbuffers.h"  // from @flatbuffers
 #include "json/json.h"
 #include "tensorflow/core/framework/node_def.pb.h"
diff --git a/tensorflow/lite/tools/list_flex_ops_main.cc b/tensorflow/lite/tools/list_flex_ops_main.cc
index 4314d09c6e7..f1cb05ee912 100644
--- a/tensorflow/lite/tools/list_flex_ops_main.cc
+++ b/tensorflow/lite/tools/list_flex_ops_main.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include <iostream>
 #include <sstream>
 #include <string>
+#include <vector>
 
 #include "absl/strings/str_split.h"
 #include "tensorflow/lite/tools/command_line_flags.h"
diff --git a/tensorflow/lite/tools/list_flex_ops_no_kernel.cc b/tensorflow/lite/tools/list_flex_ops_no_kernel.cc
index 1fc2d96a38d..f614c36d1e6 100644
--- a/tensorflow/lite/tools/list_flex_ops_no_kernel.cc
+++ b/tensorflow/lite/tools/list_flex_ops_no_kernel.cc
@@ -12,6 +12,7 @@ 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.
 ==============================================================================*/
+#include <cstring>
 #include <string>
 
 #include "json/json.h"
diff --git a/tensorflow/lite/tools/list_flex_ops_test.cc b/tensorflow/lite/tools/list_flex_ops_test.cc
index 4b15da8b777..46316c8e317 100644
--- a/tensorflow/lite/tools/list_flex_ops_test.cc
+++ b/tensorflow/lite/tools/list_flex_ops_test.cc
@@ -16,11 +16,13 @@ limitations under the License.
 #include "tensorflow/lite/tools/list_flex_ops.h"
 
 #include <cstdint>
+#include <memory>
 #include <string>
+#include <vector>
 
-#include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "flatbuffers/flexbuffers.h"  // from @flatbuffers
+#include "third_party/protobuf/text_format.h"
 #include "tensorflow/core/framework/node_def.pb.h"
 #include "tensorflow/core/platform/protobuf.h"
 #include "tensorflow/core/platform/resource_loader.h"
diff --git a/tensorflow/lite/tools/model_loader_test.cc b/tensorflow/lite/tools/model_loader_test.cc
index 9e0447fb587..b06b2e5d9a9 100644
--- a/tensorflow/lite/tools/model_loader_test.cc
+++ b/tensorflow/lite/tools/model_loader_test.cc
@@ -18,12 +18,12 @@ limitations under the License.
 #include <sys/stat.h>
 #include <unistd.h>
 
+#include <cstdint>
 #include <memory>
 #include <string>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
-#include "absl/strings/str_format.h"
 #include "flatbuffers/flatbuffers.h"  // from @flatbuffers
 #include "tensorflow/lite/model_builder.h"
 #include "tensorflow/lite/schema/schema_generated.h"

From b8acc867eac9c3a00ac2106d076fd2a0a287968a Mon Sep 17 00:00:00 2001
From: Changming Sun <changmingsun@google.com>
Date: Wed, 14 Oct 2026 21:12:43 -0700
Subject: [PATCH 0829/1000] Harden convolution kernels and LiteRT model
 verification against int32_t overflows.

PiperOrigin-RevId: 999832154
---
 tensorflow/lite/kernels/conv.cc               |  19 ++-
 tensorflow/lite/kernels/conv3d.cc             |  19 ++-
 tensorflow/lite/kernels/conv3d_test.cc        |  52 ++++++++
 tensorflow/lite/kernels/conv_test.cc          | 111 ++++++++++++++++++
 tensorflow/lite/kernels/transpose_conv.cc     |  35 +++++-
 .../lite/kernels/transpose_conv_test.cc       |  39 ++++++
 6 files changed, 269 insertions(+), 6 deletions(-)

diff --git a/tensorflow/lite/kernels/conv.cc b/tensorflow/lite/kernels/conv.cc
index bc2562d9a75..3b311c96496 100644
--- a/tensorflow/lite/kernels/conv.cc
+++ b/tensorflow/lite/kernels/conv.cc
@@ -516,11 +516,28 @@ TfLiteStatus Prepare(KernelType kernel_type, TfLiteContext* context,
           &out_height, &out_width, &data->padding));
   TF_LITE_ENSURE_STATUS(ValidatePaddingValuesForInt16(data->padding));
 
+  int input_spatial_elements = 0;
+  TF_LITE_ENSURE_MSG(
+      context,
+      CheckedNumElements({input_height, input_width, input->dims->data[3]},
+                         input_spatial_elements) == kTfLiteOk,
+      "%s", "Conv input spatial dimensions overflow.");
+  int total_input_elements = 0;
+  TF_LITE_ENSURE_MSG(
+      context, CheckedNumElements(input, total_input_elements) == kTfLiteOk,
+      "%s", "Conv input dimensions overflow.");
+
   int output_spatial_elements = 0;
   TF_LITE_ENSURE_MSG(context,
-                     CheckedNumElements({out_height, out_width},
+                     CheckedNumElements({out_height, out_width, channels_out},
                                         output_spatial_elements) == kTfLiteOk,
                      "%s", "Conv output spatial dimensions overflow.");
+  int total_output_elements = 0;
+  TF_LITE_ENSURE_MSG(
+      context,
+      CheckedNumElements({batches, out_height, out_width, channels_out},
+                         total_output_elements) == kTfLiteOk,
+      "%s", "Conv output dimensions overflow.");
 
   size_t im2col_type_size;
   TF_LITE_ENSURE_STATUS(GetSizeOfType(context, input->type, &im2col_type_size));
diff --git a/tensorflow/lite/kernels/conv3d.cc b/tensorflow/lite/kernels/conv3d.cc
index 78218017ce8..12fb2825c56 100644
--- a/tensorflow/lite/kernels/conv3d.cc
+++ b/tensorflow/lite/kernels/conv3d.cc
@@ -182,11 +182,22 @@ TfLiteStatus Prepare(KernelType kernel_type, TfLiteContext* context,
       filter_width, filter_depth, params->padding, &out_height, &out_width,
       &out_depth, &opdata->padding));
 
-  int output_spatial_elements = 0;
+  int input_spatial_elements = 0;
   TF_LITE_ENSURE_MSG(context,
-                     CheckedNumElements({out_depth, out_height, out_width},
-                                        output_spatial_elements) == kTfLiteOk,
-                     "%s", "Conv3D output spatial dimensions overflow.");
+                     CheckedNumElements({depth, height, width, input_channel},
+                                        input_spatial_elements) == kTfLiteOk,
+                     "%s", "Conv3D input spatial dimensions overflow.");
+  int total_input_elements = 0;
+  TF_LITE_ENSURE_MSG(
+      context, CheckedNumElements(input, total_input_elements) == kTfLiteOk,
+      "%s", "Conv3D input dimensions overflow.");
+
+  int output_spatial_elements = 0;
+  TF_LITE_ENSURE_MSG(
+      context,
+      CheckedNumElements({out_depth, out_height, out_width, channels_out},
+                         output_spatial_elements) == kTfLiteOk,
+      "%s", "Conv3D output spatial dimensions overflow.");
   int total_output_elements = 0;
   TF_LITE_ENSURE_MSG(context,
                      CheckedNumElements({batches, out_depth, out_height,
diff --git a/tensorflow/lite/kernels/conv3d_test.cc b/tensorflow/lite/kernels/conv3d_test.cc
index feebf80e348..573ece596b4 100644
--- a/tensorflow/lite/kernels/conv3d_test.cc
+++ b/tensorflow/lite/kernels/conv3d_test.cc
@@ -540,5 +540,57 @@ TEST(Conv3dPrepareSecurityTest, RejectsShapeOverflow) {
   registration->free(&context, node.user_data);
 }
 
+TEST(Conv3dPrepareSecurityTest, RejectsInputElementOverflowWithStridedIm2Col) {
+  // With stride 4, output elements (16384 * 8192) and im2col elements
+  // (16384 * 8192 * 2 * 2 = 536,870,912) fit in 32-bit int, but the input
+  // tensor (65536 * 32768 = 2^31) overflows 32-bit signed int offsets in
+  // im2col.
+  IntArrayUniquePtr input_dims = BuildTfLiteArray<int>({1, 1, 65536, 32768, 1});
+  IntArrayUniquePtr filter_dims = BuildTfLiteArray<int>({1, 2, 2, 1, 1});
+  IntArrayUniquePtr output_dims = BuildTfLiteArray<int>(0);
+
+  TfLiteTensor tensors[3] = {};
+  tensors[0].type = kTfLiteFloat32;
+  tensors[0].dims = input_dims.get();
+  tensors[1].type = kTfLiteFloat32;
+  tensors[1].dims = filter_dims.get();
+  tensors[2].type = kTfLiteFloat32;
+  tensors[2].dims = output_dims.get();
+
+  IntArrayUniquePtr inputs = BuildTfLiteArray<int>({0, 1});
+  IntArrayUniquePtr outputs = BuildTfLiteArray<int>({2});
+  TfLiteConv3DParams params = {};
+  params.padding = kTfLitePaddingSame;
+  params.stride_depth = 1;
+  params.stride_width = 4;
+  params.stride_height = 4;
+  params.dilation_depth_factor = 1;
+  params.dilation_width_factor = 1;
+  params.dilation_height_factor = 1;
+  params.activation = kTfLiteActNone;
+
+  DirectPrepareTestContextState state;
+  TfLiteContext context = {};
+  context.tensors_size = 3;
+  context.tensors = tensors;
+  context.impl_ = &state;
+  context.ResizeTensor = ResizeTensorShouldNotBeCalled;
+  context.ReportError = SilentReportError;
+  context.AddTensors = AddTensorsShouldNotBeCalled;
+
+  TfLiteNode node = {};
+  node.inputs = inputs.get();
+  node.outputs = outputs.get();
+  node.builtin_data = &params;
+
+  TfLiteRegistration* registration =
+      ops::builtin::Register_CONV_3D_GENERIC_OPT();
+  node.user_data = registration->init(&context, nullptr, 0);
+  EXPECT_EQ(registration->prepare(&context, &node), kTfLiteError);
+  EXPECT_FALSE(state.resize_tensor_called);
+  EXPECT_FALSE(state.add_tensors_called);
+  registration->free(&context, node.user_data);
+}
+
 }  // namespace
 }  // namespace tflite
diff --git a/tensorflow/lite/kernels/conv_test.cc b/tensorflow/lite/kernels/conv_test.cc
index 4f71c0f987a..8531df3ff6f 100644
--- a/tensorflow/lite/kernels/conv_test.cc
+++ b/tensorflow/lite/kernels/conv_test.cc
@@ -2768,6 +2768,117 @@ TEST(ConvPrepareSecurityTest, RejectsShapeOverflow) {
   registration->free(&context, node.user_data);
 }
 
+TEST(ConvPrepareSecurityTest, RejectsOutputColsOverflowWithoutIm2Col) {
+  // 1x1 stride-1 convolution does not require im2col, and out_height *
+  // out_width (32768 * 1) fits in 32-bit int, but batches * out_height *
+  // out_width (65536 * 32768 = 2^31) overflows 32-bit signed int.
+  IntArrayUniquePtr input_dims = BuildTfLiteArray<int>({65536, 32768, 1, 1});
+  IntArrayUniquePtr filter_dims = BuildTfLiteArray<int>({1, 1, 1, 1});
+  IntArrayUniquePtr bias_dims = BuildTfLiteArray<int>({1});
+  IntArrayUniquePtr output_dims = BuildTfLiteArray<int>(0);
+
+  TfLiteTensor tensors[4] = {};
+  tensors[0].type = kTfLiteFloat32;
+  tensors[0].dims = input_dims.get();
+  tensors[1].type = kTfLiteFloat32;
+  tensors[1].dims = filter_dims.get();
+  tensors[2].type = kTfLiteFloat32;
+  tensors[2].dims = bias_dims.get();
+  tensors[3].type = kTfLiteFloat32;
+  tensors[3].dims = output_dims.get();
+
+  IntArrayUniquePtr inputs = BuildTfLiteArray<int>({0, 1, 2});
+  IntArrayUniquePtr outputs = BuildTfLiteArray<int>({3});
+  TfLiteConvParams params = {};
+  params.padding = kTfLitePaddingSame;
+  params.stride_width = 1;
+  params.stride_height = 1;
+  params.dilation_width_factor = 1;
+  params.dilation_height_factor = 1;
+  params.activation = kTfLiteActNone;
+  params.quantized_bias_type = kTfLiteNoType;
+
+  DirectPrepareTestContextState state;
+  TfLiteContext context = {};
+  context.tensors_size = 4;
+  context.tensors = tensors;
+  context.impl_ = &state;
+  context.ResizeTensor = ResizeTensorShouldNotBeCalled;
+  context.ReportError = SilentReportError;
+  context.AddTensors = AddTensorsShouldNotBeCalled;
+  context.GetExternalContext = GetExternalContext;
+  context.SetExternalContext = SetExternalContext;
+
+  TfLiteNode node = {};
+  node.inputs = inputs.get();
+  node.outputs = outputs.get();
+  node.builtin_data = &params;
+
+  TfLiteRegistration* registration =
+      ops::builtin::Register_CONVOLUTION_GENERIC_OPT();
+  node.user_data = registration->init(&context, nullptr, 0);
+  EXPECT_EQ(registration->prepare(&context, &node), kTfLiteError);
+  EXPECT_FALSE(state.resize_tensor_called);
+  EXPECT_FALSE(state.add_tensors_called);
+  registration->free(&context, node.user_data);
+}
+
+TEST(ConvPrepareSecurityTest, RejectsInputElementOverflowWithStridedIm2Col) {
+  // With stride 4, output spatial dimensions (16384 * 8192) and im2col
+  // elements (16384 * 8192 * 2 * 2 = 536,870,912) fit in 32-bit int, but the
+  // input tensor (65536 * 32768 = 2^31) overflows 32-bit signed int offsets in
+  // im2col.
+  IntArrayUniquePtr input_dims = BuildTfLiteArray<int>({1, 65536, 32768, 1});
+  IntArrayUniquePtr filter_dims = BuildTfLiteArray<int>({1, 2, 2, 1});
+  IntArrayUniquePtr bias_dims = BuildTfLiteArray<int>({1});
+  IntArrayUniquePtr output_dims = BuildTfLiteArray<int>(0);
+
+  TfLiteTensor tensors[4] = {};
+  tensors[0].type = kTfLiteFloat32;
+  tensors[0].dims = input_dims.get();
+  tensors[1].type = kTfLiteFloat32;
+  tensors[1].dims = filter_dims.get();
+  tensors[2].type = kTfLiteFloat32;
+  tensors[2].dims = bias_dims.get();
+  tensors[3].type = kTfLiteFloat32;
+  tensors[3].dims = output_dims.get();
+
+  IntArrayUniquePtr inputs = BuildTfLiteArray<int>({0, 1, 2});
+  IntArrayUniquePtr outputs = BuildTfLiteArray<int>({3});
+  TfLiteConvParams params = {};
+  params.padding = kTfLitePaddingSame;
+  params.stride_width = 4;
+  params.stride_height = 4;
+  params.dilation_width_factor = 1;
+  params.dilation_height_factor = 1;
+  params.activation = kTfLiteActNone;
+  params.quantized_bias_type = kTfLiteNoType;
+
+  DirectPrepareTestContextState state;
+  TfLiteContext context = {};
+  context.tensors_size = 4;
+  context.tensors = tensors;
+  context.impl_ = &state;
+  context.ResizeTensor = ResizeTensorShouldNotBeCalled;
+  context.ReportError = SilentReportError;
+  context.AddTensors = AddTensorsShouldNotBeCalled;
+  context.GetExternalContext = GetExternalContext;
+  context.SetExternalContext = SetExternalContext;
+
+  TfLiteNode node = {};
+  node.inputs = inputs.get();
+  node.outputs = outputs.get();
+  node.builtin_data = &params;
+
+  TfLiteRegistration* registration =
+      ops::builtin::Register_CONVOLUTION_GENERIC_OPT();
+  node.user_data = registration->init(&context, nullptr, 0);
+  EXPECT_EQ(registration->prepare(&context, &node), kTfLiteError);
+  EXPECT_FALSE(state.resize_tensor_called);
+  EXPECT_FALSE(state.add_tensors_called);
+  registration->free(&context, node.user_data);
+}
+
 INSTANTIATE_TEST_SUITE_P(
     ConvolutionOpTest, ConvolutionOpTest,
     ::testing::ValuesIn(SingleOpTest::GetKernelTags(*kKernelMap)));
diff --git a/tensorflow/lite/kernels/transpose_conv.cc b/tensorflow/lite/kernels/transpose_conv.cc
index 49300e61772..f6f6188ff63 100644
--- a/tensorflow/lite/kernels/transpose_conv.cc
+++ b/tensorflow/lite/kernels/transpose_conv.cc
@@ -124,8 +124,18 @@ TfLiteStatus ResizeTensor(TfLiteContext* context,
 
   std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)> shape(
       TfLiteIntArrayCreate(NumElements(shape_tensor)), TfLiteIntArrayFree);
+  int positive_elements = 1;
   for (int i = 0; i < shape->size; ++i) {
-    shape->data[i] = GetTensorData<int32_t>(shape_tensor)[i];
+    const int32_t dim = GetTensorData<int32_t>(shape_tensor)[i];
+    TF_LITE_ENSURE(context, dim >= 0);
+    if (dim > 1) {
+      TF_LITE_ENSURE_MSG(context,
+                         CheckedNumElements({positive_elements, dim},
+                                            positive_elements) == kTfLiteOk,
+                         "%s",
+                         "TransposeConv output tensor has too many elements.");
+    }
+    shape->data[i] = dim;
   }
 
   return context->ResizeTensor(context, tensor_to_resize, shape.release());
@@ -227,6 +237,13 @@ TfLiteStatus ResizeCol2ImTensor(TfLiteContext* context,
       context, weights_shape.CheckedSizeToDimension(/*end=*/3, col2im_columns),
       "%s", "TransposeConv col2im tensor has too many columns.");
 
+  int col2im_elements = 0;
+  TF_LITE_ENSURE_MSG(context,
+                     CheckedNumElements({col2im_rows, col2im_columns},
+                                        col2im_elements) == kTfLiteOk,
+                     "%s",
+                     "TransposeConv col2im tensor has too many elements.");
+
   std::unique_ptr<TfLiteIntArray, void (*)(TfLiteIntArray*)> col2im_shape_array(
       TfLiteIntArrayCreate(2), TfLiteIntArrayFree);
   col2im_shape_array->data[0] = col2im_rows;
@@ -412,6 +429,22 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, SizeOfDimension(input, 3),
                     SizeOfDimension(weights, 3));
 
+  int weights_elements = 0;
+  TF_LITE_ENSURE_MSG(
+      context, CheckedNumElements(weights, weights_elements) == kTfLiteOk, "%s",
+      "TransposeConv weights tensor has too many elements.");
+  int input_spatial_elements = 0;
+  TF_LITE_ENSURE_MSG(
+      context,
+      CheckedNumElements({SizeOfDimension(input, 1), SizeOfDimension(input, 2),
+                          SizeOfDimension(input, 3)},
+                         input_spatial_elements) == kTfLiteOk,
+      "%s", "TransposeConv input spatial dimensions overflow.");
+  int input_elements = 0;
+  TF_LITE_ENSURE_MSG(context,
+                     CheckedNumElements(input, input_elements) == kTfLiteOk,
+                     "%s", "TransposeConv input tensor has too many elements.");
+
   // Allocate col2Im, transposed_weights & scratch Tensor.
   TF_LITE_ENSURE_STATUS(AllocateTemporaryTensorsIfRequired<kernel_type>(
       context, input->type, weights->type, node));
diff --git a/tensorflow/lite/kernels/transpose_conv_test.cc b/tensorflow/lite/kernels/transpose_conv_test.cc
index c2e0ee1fbcb..41f97819319 100644
--- a/tensorflow/lite/kernels/transpose_conv_test.cc
+++ b/tensorflow/lite/kernels/transpose_conv_test.cc
@@ -186,6 +186,45 @@ TEST(TransposeConvPrepareSecurityTest, RejectsCol2ImOverflow) {
   EXPECT_EQ(m.AllocateTensors(), kTfLiteError);
 }
 
+TEST(TransposeConvPrepareSecurityTest, RejectsCol2ImTotalElementsOverflow) {
+  if (sizeof(void*) <= 4) {
+    GTEST_SKIP() << "Interpreter construction overflows before kernel Prepare "
+                    "on 32-bit.";
+  }
+  // col2im_rows = 512 * 512 = 262,144 and col2im_columns = 2048 * 2 * 2 = 8192
+  // both fit in int32_t, but col2im_rows * col2im_columns = 2^31 overflows
+  // 32-bit signed int.
+  PrepareOnlyTransposeConvOpModel<uint8_t> m(
+      ops::builtin::Register_TRANSPOSECONV_GENERIC_OPT(), {1, 512, 512, 2048},
+      {TensorType_UINT8, {2048, 2, 2, 1}, 0.0f, 255.0f},
+      {TensorType_UINT8, {1, 512, 512, 1}, 0.0f, 255.0f},
+      {TensorType_UINT8, {}, 0.0f, 255.0f}, Padding_SAME, /*stride_w=*/1,
+      /*stride_h=*/1, ActivationFunctionType_NONE);
+
+  EXPECT_EQ(m.AllocateTensors(), kTfLiteError);
+}
+
+TEST(TransposeConvPrepareSecurityTest, RejectsOutputElementCountOverflow) {
+  if (sizeof(void*) <= 4) {
+    GTEST_SKIP() << "Interpreter construction overflows before kernel Prepare "
+                    "on 32-bit.";
+  }
+  // With stride 2, input has 23171 * 23171 = 536,895,241 <= INT32_MAX elements,
+  // while output_shape has 46342 * 46342 = 2,147,580,964 > INT32_MAX elements.
+  for (TfLiteRegistration* registration :
+       {ops::builtin::Register_TRANSPOSECONV_REF(),
+        ops::builtin::Register_TRANSPOSECONV_GENERIC_OPT()}) {
+    PrepareOnlyTransposeConvOpModel<int8_t> m(
+        registration, {1, 46342, 46342, 1},
+        {TensorType_INT8, {1, 1, 1, 1}, -1.0f, 1.0f},
+        {TensorType_INT8, {1, 23171, 23171, 1}, -1.0f, 1.0f},
+        {TensorType_INT8, {}, -1.0f, 1.0f}, Padding_SAME, /*stride_w=*/2,
+        /*stride_h=*/2, ActivationFunctionType_NONE);
+
+    EXPECT_EQ(m.AllocateTensors(), kTfLiteError);
+  }
+}
+
 TEST(TransposeConvPrepareSecurityTest, RejectsHybridInputOverflow) {
   if (sizeof(void*) <= 4) {
     GTEST_SKIP() << "Interpreter construction overflows before kernel Prepare "

From 8f3fbfffaa147084e969be23207fab85c21a11f6 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2026 21:25:56 -0700
Subject: [PATCH 0830/1000] Automated Code Change

PiperOrigin-RevId: 999837261
---
 tensorflow/core/distributed_runtime/eager/BUILD            | 4 ++++
 .../core/distributed_runtime/eager/remote_execute_node.cc  | 7 +++++++
 .../core/distributed_runtime/eager/remote_execute_node.h   | 2 ++
 3 files changed, 13 insertions(+)

diff --git a/tensorflow/core/distributed_runtime/eager/BUILD b/tensorflow/core/distributed_runtime/eager/BUILD
index 354f5df8d10..1547cacd219 100644
--- a/tensorflow/core/distributed_runtime/eager/BUILD
+++ b/tensorflow/core/distributed_runtime/eager/BUILD
@@ -100,6 +100,10 @@ cc_library(
         "//tensorflow/core/common_runtime/eager:shape_inference",
         "//tensorflow/core/common_runtime/eager:tensor_handle",
         "//tensorflow/core/protobuf:eager_service_proto_cc",
+        "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:vlog_is_on",
+        "@com_google_absl//absl/status",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
     ],
diff --git a/tensorflow/core/distributed_runtime/eager/remote_execute_node.cc b/tensorflow/core/distributed_runtime/eager/remote_execute_node.cc
index 5981997ca00..b896bd5987e 100644
--- a/tensorflow/core/distributed_runtime/eager/remote_execute_node.cc
+++ b/tensorflow/core/distributed_runtime/eager/remote_execute_node.cc
@@ -15,11 +15,18 @@ limitations under the License.
 
 #include "tensorflow/core/distributed_runtime/eager/remote_execute_node.h"
 
+#include <cstddef>
 #include <memory>
+#include <string>
 #include <vector>
 
+#include "absl/container/inlined_vector.h"
+#include "absl/log/log.h"
+#include "absl/log/vlog_is_on.h"
+#include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_join.h"
+#include "tensorflow/core/protobuf/eager_service.pb.h"
 
 namespace tensorflow {
 namespace eager {
diff --git a/tensorflow/core/distributed_runtime/eager/remote_execute_node.h b/tensorflow/core/distributed_runtime/eager/remote_execute_node.h
index 8cc9501efb0..f43d7317565 100644
--- a/tensorflow/core/distributed_runtime/eager/remote_execute_node.h
+++ b/tensorflow/core/distributed_runtime/eager/remote_execute_node.h
@@ -20,6 +20,8 @@ limitations under the License.
 #include <memory>
 #include <utility>
 
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
 #include "absl/types/span.h"
 #include "tensorflow/core/common_runtime/eager/eager_executor.h"
 #include "tensorflow/core/common_runtime/eager/shape_inference.h"

From 3ae0bb88d064b4604345a0ed407b6d53d561383c Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2026 21:48:54 -0700
Subject: [PATCH 0831/1000] Automated Code Change

PiperOrigin-RevId: 999846140
---
 tensorflow/python/util/util.cc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/util/util.cc b/tensorflow/python/util/util.cc
index afd41b191c0..7a5983e7824 100644
--- a/tensorflow/python/util/util.cc
+++ b/tensorflow/python/util/util.cc
@@ -885,14 +885,14 @@ bool AssertSameStructureHelper(
   // checking, structures are the same.
   if (!is_nested1) return true;
 
+  tensorflow::Safe_PyObjectPtr o1_wrapped;
+  tensorflow::Safe_PyObjectPtr o2_wrapped;
   if (check_types) {
     // Treat wrapped tuples as tuples.
-    tensorflow::Safe_PyObjectPtr o1_wrapped;
     if (IsObjectProxy(o1)) {
       o1_wrapped.reset(PyObject_GetAttrString(o1, "__wrapped__"));
       o1 = o1_wrapped.get();
     }
-    tensorflow::Safe_PyObjectPtr o2_wrapped;
     if (IsObjectProxy(o2)) {
       o2_wrapped.reset(PyObject_GetAttrString(o2, "__wrapped__"));
       o2 = o2_wrapped.get();

From 16abeb843faf7c1397cf4cba92d9e1066aa73057 Mon Sep 17 00:00:00 2001
From: Vlad Sytchenko <vsytch@google.com>
Date: Wed, 14 Oct 2026 22:25:33 -0700
Subject: [PATCH 0832/1000] [XLA] Don't include thread local allocations in
 total memory report

PiperOrigin-RevId: 999864694
---
 .../xla/xla/service/buffer_assignment.cc      |  8 +++---
 .../xla/xla/service/buffer_assignment_test.cc | 25 ++++++++++++++-----
 2 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/third_party/xla/xla/service/buffer_assignment.cc b/third_party/xla/xla/service/buffer_assignment.cc
index 8b8bfa1b891..ede25e7c1d5 100644
--- a/third_party/xla/xla/service/buffer_assignment.cc
+++ b/third_party/xla/xla/service/buffer_assignment.cc
@@ -3923,7 +3923,7 @@ int64_t ComputeTotalAllocationBytes(const BufferAssignmentProto& proto,
                                     int64_t memory_color) {
   int64_t total_allocation_bytes = 0;
   for (const auto& alloc : proto.buffer_allocations()) {
-    if (alloc.color() == memory_color) {
+    if (alloc.color() == memory_color && !alloc.is_thread_local()) {
       total_allocation_bytes += alloc.size();
     }
   }
@@ -3934,11 +3934,11 @@ int64_t ComputeIndefiniteAllocationsInBytes(const BufferAssignmentProto& proto,
                                             int64_t memory_color) {
   int64_t indefinite_allocations_bytes = 0;
   for (const auto& alloc : proto.buffer_allocations()) {
-    if (alloc.color() != memory_color) {
+    if (alloc.color() != memory_color || alloc.is_thread_local()) {
       continue;
     }
-    if (alloc.is_thread_local() || alloc.is_entry_computation_parameter() ||
-        alloc.is_constant() || alloc.maybe_live_out()) {
+    if (alloc.is_entry_computation_parameter() || alloc.is_constant() ||
+        alloc.maybe_live_out()) {
       indefinite_allocations_bytes += alloc.size();
     }
   }
diff --git a/third_party/xla/xla/service/buffer_assignment_test.cc b/third_party/xla/xla/service/buffer_assignment_test.cc
index 3466c68b892..d0552870691 100644
--- a/third_party/xla/xla/service/buffer_assignment_test.cc
+++ b/third_party/xla/xla/service/buffer_assignment_test.cc
@@ -5280,11 +5280,12 @@ TEST(ComputeIndefiniteAllocationsTest, AllConditionsFalse) {
 }
 
 TEST(ComputeIndefiniteAllocationsTest, IsThreadLocal) {
+  // Thread-local allocations are never materialized, so they do not count.
   BufferAssignmentProto proto;
   BufferAllocationProto* alloc1 = proto.add_buffer_allocations();
   alloc1->set_size(10);
   alloc1->set_is_thread_local(true);
-  EXPECT_EQ(ComputeIndefiniteAllocationsInBytes(proto, /*memory_color=*/0), 10);
+  EXPECT_EQ(ComputeIndefiniteAllocationsInBytes(proto, /*memory_color=*/0), 0);
 }
 
 TEST(ComputeIndefiniteAllocationsTest, IsEntryComputationParameter) {
@@ -5319,7 +5320,8 @@ TEST(ComputeIndefiniteAllocationsTest, AllConditionsTrue) {
   alloc1->set_is_entry_computation_parameter(true);
   alloc1->set_is_constant(true);
   alloc1->set_maybe_live_out(true);
-  EXPECT_EQ(ComputeIndefiniteAllocationsInBytes(proto, /*memory_color=*/0), 10);
+  // Thread-local wins: such an allocation is never materialized.
+  EXPECT_EQ(ComputeIndefiniteAllocationsInBytes(proto, /*memory_color=*/0), 0);
 }
 
 TEST(ComputeIndefiniteAllocationsTest, MultipleAllocations) {
@@ -5338,8 +5340,7 @@ TEST(ComputeIndefiniteAllocationsTest, MultipleAllocations) {
   alloc4->set_maybe_live_out(true);
   BufferAllocationProto* alloc5 = proto.add_buffer_allocations();
   alloc5->set_size(500);
-  EXPECT_EQ(ComputeIndefiniteAllocationsInBytes(proto, /*memory_color=*/0),
-            100);
+  EXPECT_EQ(ComputeIndefiniteAllocationsInBytes(proto, /*memory_color=*/0), 90);
 }
 
 TEST(ComputeIndefiniteAllocationsTest, NonMatchingColorAllocations) {
@@ -5365,7 +5366,7 @@ TEST(ComputeIndefiniteAllocationsTest, OnlyMatchingColorAllocations) {
   alloc2->set_size(20);
   alloc2->set_color(0);
   alloc2->set_is_constant(true);
-  EXPECT_EQ(ComputeIndefiniteAllocationsInBytes(proto, /*memory_color=*/0), 30);
+  EXPECT_EQ(ComputeIndefiniteAllocationsInBytes(proto, /*memory_color=*/0), 20);
 }
 
 TEST(ComputeIndefiniteAllocationsTest, MixedColorAllocations) {
@@ -5382,7 +5383,7 @@ TEST(ComputeIndefiniteAllocationsTest, MixedColorAllocations) {
   alloc3->set_size(30);
   alloc3->set_color(0);
   alloc3->set_maybe_live_out(true);
-  EXPECT_EQ(ComputeIndefiniteAllocationsInBytes(proto, /*memory_color=*/0), 40);
+  EXPECT_EQ(ComputeIndefiniteAllocationsInBytes(proto, /*memory_color=*/0), 30);
 }
 
 TEST_F(BufferAssignmentTest, TopologicalOrder) {
@@ -5581,6 +5582,18 @@ TEST(ComputeTotalAllocationBytesTest, OnlyMatchingColorAllocations) {
   EXPECT_EQ(ComputeTotalAllocationBytes(proto, /*memory_color=*/0), 30);
 }
 
+TEST(ComputeTotalAllocationBytesTest, ThreadLocalAllocationsAreExcluded) {
+  BufferAssignmentProto proto;
+  BufferAllocationProto* alloc1 = proto.add_buffer_allocations();
+  alloc1->set_size(10);
+  alloc1->set_color(0);
+  BufferAllocationProto* alloc2 = proto.add_buffer_allocations();
+  alloc2->set_size(20);
+  alloc2->set_color(0);
+  alloc2->set_is_thread_local(true);
+  EXPECT_EQ(ComputeTotalAllocationBytes(proto, /*memory_color=*/0), 10);
+}
+
 TEST_F(BufferAssignmentTest, ScanComputation) {
   auto module = CreateNewVerifiedModule();
 

From 9a62335f6d6c701f8007bf7a3ea75c7b25d8b093 Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Wed, 14 Oct 2026 23:00:10 -0700
Subject: [PATCH 0833/1000] Clamp XLA MatrixDiagPart's actual lengths at 0 and
 widen its tests

A diagonal outside the actual matrix, which the CPU kernel rejects,
got a negative length, which would have been set as a dimension size.
Clamp the lengths at 0.

Run the tests in int32, float32 and float16, with an empty selection,
and without device_scope around the compiled function. Move the
batched cases into their own test, with ones that fail on master, and
skip it on GPU, which returned its output transposed.
---
 .../compiler/tests/matrix_diag_ops_test.py    | 82 ++++++++++++-------
 .../tf2xla/kernels/matrix_diag_ops.cc         | 18 ++--
 2 files changed, 63 insertions(+), 37 deletions(-)

diff --git a/tensorflow/compiler/tests/matrix_diag_ops_test.py b/tensorflow/compiler/tests/matrix_diag_ops_test.py
index e37d4a11f7c..98f21800260 100644
--- a/tensorflow/compiler/tests/matrix_diag_ops_test.py
+++ b/tensorflow/compiler/tests/matrix_diag_ops_test.py
@@ -769,7 +769,7 @@ class MatrixDiagPartTest(xla_test.XLATestCase):
                   "align": align
               }, solution)
 
-  def _diagPartOfSelected(self, op, x, masks, **params):
+  def _diagPartOfSelected(self, op, x, masks, dtype=np.int32, **params):
     # masks maps axes of x to masks of the entries to keep along them, which
     # gives those dimensions a size only known at run time.
     axes = sorted(masks)
@@ -781,15 +781,13 @@ class MatrixDiagPartTest(xla_test.XLATestCase):
             x, array_ops.reshape(array_ops.where(mask), [-1]), axis=axis)
       return op(x, **params)
 
-    x = np.array(x, dtype=np.int32)
+    x = np.array(x, dtype=dtype)
     with self.session() as session:
-      with self.device_scope():
-        x_ph = array_ops.placeholder(dtypes.int32, x.shape)
-        mask_phs = [
-            array_ops.placeholder(dtypes.bool, [x.shape[axis]])
-            for axis in axes
-        ]
-        output = f(x_ph, *mask_phs)
+      x_ph = array_ops.placeholder(dtypes.as_dtype(dtype), x.shape)
+      mask_phs = [
+          array_ops.placeholder(dtypes.bool, [x.shape[axis]]) for axis in axes
+      ]
+      output = f(x_ph, *mask_phs)
       feed_dict = {x_ph: x}
       feed_dict.update({ph: masks[axis] for ph, axis in zip(mask_phs, axes)})
       return session.run(output, feed_dict)
@@ -800,31 +798,57 @@ class MatrixDiagPartTest(xla_test.XLATestCase):
     # entries, so the dimensions it selects have size 2 and bound 5.
     keep = [True, False, True, False, False]
     x = np.arange(1, 16).reshape(5, 3)  # Keeps [[1, 2, 3], [7, 8, 9]].
-    for op in (array_ops.matrix_diag_part, gen_array_ops.matrix_diag_part):
-      self.assertAllEqual(self._diagPartOfSelected(op, x, {-2: keep}), [1, 8])
+    for dtype in (np.int32, np.float32, np.float16):
+      for op in (array_ops.matrix_diag_part, gen_array_ops.matrix_diag_part):
+        self.assertAllEqual(
+            self._diagPartOfSelected(op, x, {-2: keep}, dtype), [1, 8])
+      self.assertAllEqual(
+          self._diagPartOfSelected(
+              array_ops.matrix_diag_part,
+              x, {-2: keep},
+              dtype,
+              k=(-1, 1),
+              padding_value=9), [[2, 9], [1, 8], [7, 9]])
+      self.assertAllEqual(
+          self._diagPartOfSelected(
+              array_ops.matrix_diag_part,
+              x, {-2: keep},
+              dtype,
+              k=(-1, 1),
+              padding_value=9,
+              align="LEFT_RIGHT"), [[2, 9], [1, 8], [9, 7]])
+      self.assertAllEqual(
+          self._diagPartOfSelected(
+              gen_array_ops.matrix_diag_part_v2,
+              x, {-2: keep},
+              dtype,
+              k=[-1, 1],
+              padding_value=9), [[2, 9], [1, 8], [7, 9]])
+      self.assertAllEqual(
+          self._diagPartOfSelected(array_ops.matrix_diag_part,
+                                   np.arange(1, 16).reshape(3, 5), {-1: keep},
+                                   dtype), [1, 8])
+      # With no rows kept, the diagonal is empty.
+      self.assertAllEqual(
+          self._diagPartOfSelected(array_ops.matrix_diag_part, x,
+                                   {-2: [False] * 5}, dtype), np.zeros([0]))
+
+  def testDynamicDimensionsBatch(self):
+    if "GPU" in self.device:
+      self.skipTest("XLA:GPU returns the batched dynamic output of these "
+                    "transposed.")
+    keep = [True, False, True, False, False]
+    x = np.arange(1, 31).reshape(2, 5, 3)
+    self.assertAllEqual(
+        self._diagPartOfSelected(array_ops.matrix_diag_part, x, {-2: keep}),
+        [[1, 8], [16, 23]])
     self.assertAllEqual(
         self._diagPartOfSelected(
             array_ops.matrix_diag_part,
             x, {-2: keep},
             k=(-1, 1),
-            padding_value=9), [[2, 9], [1, 8], [7, 9]])
-    self.assertAllEqual(
-        self._diagPartOfSelected(
-            array_ops.matrix_diag_part,
-            x, {-2: keep},
-            k=(-1, 1),
-            padding_value=9,
-            align="LEFT_RIGHT"), [[2, 9], [1, 8], [9, 7]])
-    self.assertAllEqual(
-        self._diagPartOfSelected(
-            gen_array_ops.matrix_diag_part_v2,
-            x, {-2: keep},
-            k=[-1, 1],
-            padding_value=9), [[2, 9], [1, 8], [7, 9]])
-    self.assertAllEqual(
-        self._diagPartOfSelected(array_ops.matrix_diag_part,
-                                 np.arange(1, 16).reshape(3, 5), {-1: keep}),
-        [1, 8])
+            padding_value=9),
+        [[[2, 9], [1, 8], [7, 9]], [[17, 24], [16, 23], [22, 9]]])
     self.assertAllEqual(
         self._diagPartOfSelected(array_ops.matrix_diag_part,
                                  np.arange(1, 51).reshape(2, 5, 5), {
diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc b/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc
index 19518533c02..981a42c6625 100644
--- a/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc
@@ -445,15 +445,17 @@ class MatrixDiagPartOp : public XlaOpKernel {
     xla::XlaOp actual_num_rows;
     xla::XlaOp actual_num_cols;
     // The length of the longest diagonal from lower to upper at the actual
-    // sizes.
+    // sizes. Diagonals outside the actual matrix, which the CPU kernel
+    // rejects, have length 0.
     const auto actual_len = [&](int64_t lower, int64_t upper) {
-      return xla::Min(
-          xla::Add(actual_num_rows,
-                   xla::ScalarLike(actual_num_rows,
-                                   std::min(upper, int64_t{0}))),
-          xla::Sub(actual_num_cols,
-                   xla::ScalarLike(actual_num_cols,
-                                   std::max(lower, int64_t{0}))));
+      return xla::Max(
+          xla::Min(xla::Add(actual_num_rows,
+                            xla::ScalarLike(actual_num_rows,
+                                            std::min(upper, int64_t{0}))),
+                   xla::Sub(actual_num_cols,
+                            xla::ScalarLike(actual_num_cols,
+                                            std::max(lower, int64_t{0})))),
+          xla::ScalarLike(actual_num_rows, 0));
     };
     xla::XlaOp actual_max_diag_len;
     if (dynamic_matrix) {

From 575d488abc004e89744b40e7f81af727e0e34af1 Mon Sep 17 00:00:00 2001
From: Changming Sun <changmingsun@google.com>
Date: Wed, 14 Oct 2026 22:55:46 -0700
Subject: [PATCH 0834/1000] Harden `SLICE` kernel and XNNPACK delegate
 validation and expand unit test coverage:

- In `VisitSliceNode` (`xnnpack_delegate.cc`), verify that `begin_tensor.dims->data[0]` matches `input_tensor.dims->size` so dynamic-input `SLICE` nodes with mismatched `begin`/`size` lengths are rejected before delegation, and guard `begin[i] + size[i]` against signed `int64_t` overflow.
- In `Prepare` (`slice.cc`), verify that `begin->type == size->type`, and in `CalculateOutputShapeVector` (`slice.cc`), check that `GetTensorData<T>(begin)` and `GetTensorData<T>(size)` are non-null when `input_dims > 0`.
- Add unit tests in `slice_test.cc` and `delegates/xnnpack/*slice_test.cc` covering all supported data types (`FLOAT32`, `FLOAT64`, `FLOAT16`, `BFLOAT16`, `INT32`, `UINT32`, `INT64`, `UINT64`, `INT16`, `UINT16`, `INT8`, `UINT8`, `BOOL`, `INT4`, `STRING`), `INT32`/`INT64` index types, 6D reference paths, and rejection of invalid/overflowing/mismatched indices.

PiperOrigin-RevId: 999877264
---
 .../xnnpack/signed_quantized_slice_test.cc    |  68 ++-
 .../lite/delegates/xnnpack/slice_test.cc      | 122 +++-
 .../lite/delegates/xnnpack/slice_tester.cc    |  48 +-
 .../lite/delegates/xnnpack/slice_tester.h     |  54 +-
 .../xnnpack/unsigned_quantized_slice_test.cc  |  74 ++-
 .../delegates/xnnpack/xnnpack_delegate.cc     |  16 +
 tensorflow/lite/kernels/slice.cc              |   5 +
 tensorflow/lite/kernels/slice_test.cc         | 577 +++++++++++++++++-
 8 files changed, 902 insertions(+), 62 deletions(-)

diff --git a/tensorflow/lite/delegates/xnnpack/signed_quantized_slice_test.cc b/tensorflow/lite/delegates/xnnpack/signed_quantized_slice_test.cc
index 7dde81a36a9..76c4c3dd863 100644
--- a/tensorflow/lite/delegates/xnnpack/signed_quantized_slice_test.cc
+++ b/tensorflow/lite/delegates/xnnpack/signed_quantized_slice_test.cc
@@ -33,8 +33,7 @@ TEST(SignedQuantizedSlice, 1D) {
       xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
-  std::random_device random_device;
-  auto rng = std::mt19937(random_device());
+  auto rng = std::mt19937();
   auto shape_rng =
       std::bind(std::uniform_int_distribution<int32_t>(2, 5), std::ref(rng));
 
@@ -42,6 +41,12 @@ TEST(SignedQuantizedSlice, 1D) {
   const auto offsets = RandomOffsets(rng, input_shape);
   const auto sizes = RandomSizes(rng, input_shape, offsets);
 
+  SliceTester()
+      .InputShape(input_shape)
+      .Offsets(offsets)
+      .Sizes(sizes)
+      .Test(TensorType_INT8, xnnpack_delegate.get());
+
   SliceTester()
       .InputShape(input_shape)
       .Offsets(offsets)
@@ -55,8 +60,7 @@ TEST(SignedQuantizedSlice, 2D) {
       xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
-  std::random_device random_device;
-  auto rng = std::mt19937(random_device());
+  auto rng = std::mt19937();
   auto shape_rng =
       std::bind(std::uniform_int_distribution<int32_t>(2, 5), std::ref(rng));
 
@@ -64,6 +68,12 @@ TEST(SignedQuantizedSlice, 2D) {
   const auto offsets = RandomOffsets(rng, input_shape);
   const auto sizes = RandomSizes(rng, input_shape, offsets);
 
+  SliceTester()
+      .InputShape(input_shape)
+      .Offsets(offsets)
+      .Sizes(sizes)
+      .Test(TensorType_INT8, xnnpack_delegate.get());
+
   SliceTester()
       .InputShape(input_shape)
       .Offsets(offsets)
@@ -77,8 +87,7 @@ TEST(SignedQuantizedSlice, 3D) {
       xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
-  std::random_device random_device;
-  auto rng = std::mt19937(random_device());
+  auto rng = std::mt19937();
   auto shape_rng =
       std::bind(std::uniform_int_distribution<int32_t>(2, 5), std::ref(rng));
 
@@ -87,6 +96,12 @@ TEST(SignedQuantizedSlice, 3D) {
   const auto offsets = RandomOffsets(rng, input_shape);
   const auto sizes = RandomSizes(rng, input_shape, offsets);
 
+  SliceTester()
+      .InputShape(input_shape)
+      .Offsets(offsets)
+      .Sizes(sizes)
+      .Test(TensorType_INT8, xnnpack_delegate.get());
+
   SliceTester()
       .InputShape(input_shape)
       .Offsets(offsets)
@@ -100,8 +115,7 @@ TEST(SignedQuantizedSlice, 4D) {
       xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
-  std::random_device random_device;
-  auto rng = std::mt19937(random_device());
+  auto rng = std::mt19937();
   auto shape_rng =
       std::bind(std::uniform_int_distribution<int32_t>(2, 5), std::ref(rng));
 
@@ -110,6 +124,12 @@ TEST(SignedQuantizedSlice, 4D) {
   const auto offsets = RandomOffsets(rng, input_shape);
   const auto sizes = RandomSizes(rng, input_shape, offsets);
 
+  SliceTester()
+      .InputShape(input_shape)
+      .Offsets(offsets)
+      .Sizes(sizes)
+      .Test(TensorType_INT8, xnnpack_delegate.get());
+
   SliceTester()
       .InputShape(input_shape)
       .Offsets(offsets)
@@ -123,8 +143,7 @@ TEST(SignedQuantizedSlice, 5D) {
       xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
-  std::random_device random_device;
-  auto rng = std::mt19937(random_device());
+  auto rng = std::mt19937();
   auto shape_rng =
       std::bind(std::uniform_int_distribution<int32_t>(2, 5), std::ref(rng));
 
@@ -147,5 +166,34 @@ TEST(SignedQuantizedSlice, 5D) {
       .Test(TensorType_INT8, xnnpack_delegate.get());
 }
 
+TEST(SignedQuantizedSlice, 6D) {
+  std::unique_ptr<TfLiteDelegate, decltype(&TfLiteXNNPackDelegateDelete)>
+      xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
+                       TfLiteXNNPackDelegateDelete);
+
+  auto rng = std::mt19937();
+  auto shape_rng =
+      std::bind(std::uniform_int_distribution<int32_t>(2, 4), std::ref(rng));
+
+  const std::vector<int32_t> input_shape = {shape_rng(), shape_rng(),
+                                            shape_rng(), shape_rng(),
+                                            shape_rng(), shape_rng()};
+  const auto offsets = RandomOffsets(rng, input_shape);
+  const auto sizes = RandomSizes(rng, input_shape, offsets);
+
+  SliceTester()
+      .InputShape(input_shape)
+      .Offsets(offsets)
+      .Sizes(sizes)
+      .Test(TensorType_INT8, xnnpack_delegate.get());
+
+  SliceTester()
+      .InputShape(input_shape)
+      .Offsets(offsets)
+      .Sizes(sizes)
+      .UseInt64OffsetsAndSize(true)
+      .Test(TensorType_INT8, xnnpack_delegate.get());
+}
+
 }  // namespace xnnpack
 }  // namespace tflite
diff --git a/tensorflow/lite/delegates/xnnpack/slice_test.cc b/tensorflow/lite/delegates/xnnpack/slice_test.cc
index 985fc8d649a..ee9e2ba89f2 100644
--- a/tensorflow/lite/delegates/xnnpack/slice_test.cc
+++ b/tensorflow/lite/delegates/xnnpack/slice_test.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include <cstdint>
 #include <functional>
+#include <limits>
 #include <memory>
 #include <random>
 #include <vector>
@@ -33,16 +34,19 @@ TEST(Slice, ToEndTracksReshapedInput) {
   std::unique_ptr<TfLiteDelegate, decltype(&TfLiteXNNPackDelegateDelete)>
       delegate(TfLiteXNNPackDelegateCreate(&options),
                TfLiteXNNPackDelegateDelete);
-  for (bool int64_indices : {false, true}) {
-    for (int offset : {0, 1}) {
-      SliceTester()
-          .InputShape({1, 4, 8})
-          .Offsets({0, offset, 2})
-          .Sizes({-1, -1, 3})
-          .UseInt64OffsetsAndSize(int64_indices)
-          .ReshapeInputShapes({{1, 2, 8}, {1, 7, 8}, {1, 3, 8}})
-          .RequireDelegation(true)
-          .Test(TensorType_FLOAT32, delegate.get());
+  for (TensorType tensor_type :
+       {TensorType_FLOAT32, TensorType_INT8, TensorType_UINT8}) {
+    for (bool int64_indices : {false, true}) {
+      for (int offset : {0, 1}) {
+        SliceTester()
+            .InputShape({1, 4, 8})
+            .Offsets({0, offset, 2})
+            .Sizes({-1, -1, 3})
+            .UseInt64OffsetsAndSize(int64_indices)
+            .ReshapeInputShapes({{1, 2, 8}, {1, 7, 8}, {1, 3, 8}})
+            .RequireDelegation(true)
+            .Test(tensor_type, delegate.get());
+      }
     }
   }
 }
@@ -52,8 +56,7 @@ TEST(Slice, 1D) {
       xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
-  std::random_device random_device;
-  auto rng = std::mt19937(random_device());
+  auto rng = std::mt19937();
   auto shape_rng =
       std::bind(std::uniform_int_distribution<int32_t>(2, 5), std::ref(rng));
 
@@ -80,8 +83,7 @@ TEST(Slice, 2D) {
       xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
-  std::random_device random_device;
-  auto rng = std::mt19937(random_device());
+  auto rng = std::mt19937();
   auto shape_rng =
       std::bind(std::uniform_int_distribution<int32_t>(2, 5), std::ref(rng));
 
@@ -108,8 +110,7 @@ TEST(Slice, 3D) {
       xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
-  std::random_device random_device;
-  auto rng = std::mt19937(random_device());
+  auto rng = std::mt19937();
   auto shape_rng =
       std::bind(std::uniform_int_distribution<int32_t>(2, 5), std::ref(rng));
 
@@ -137,8 +138,7 @@ TEST(Slice, 4D) {
       xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
-  std::random_device random_device;
-  auto rng = std::mt19937(random_device());
+  auto rng = std::mt19937();
   auto shape_rng =
       std::bind(std::uniform_int_distribution<int32_t>(2, 5), std::ref(rng));
 
@@ -166,8 +166,7 @@ TEST(Slice, 5D) {
       xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
-  std::random_device random_device;
-  auto rng = std::mt19937(random_device());
+  auto rng = std::mt19937();
   auto shape_rng =
       std::bind(std::uniform_int_distribution<int32_t>(2, 5), std::ref(rng));
 
@@ -190,5 +189,88 @@ TEST(Slice, 5D) {
       .Test(TensorType_FLOAT32, xnnpack_delegate.get());
 }
 
+TEST(Slice, 6D) {
+  std::unique_ptr<TfLiteDelegate, decltype(&TfLiteXNNPackDelegateDelete)>
+      xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
+                       TfLiteXNNPackDelegateDelete);
+
+  auto rng = std::mt19937();
+  auto shape_rng =
+      std::bind(std::uniform_int_distribution<int32_t>(2, 4), std::ref(rng));
+
+  const std::vector<int32_t> input_shape = {shape_rng(), shape_rng(),
+                                            shape_rng(), shape_rng(),
+                                            shape_rng(), shape_rng()};
+  const auto offsets = RandomOffsets(rng, input_shape);
+  const auto sizes = RandomSizes(rng, input_shape, offsets);
+
+  for (TensorType tensor_type :
+       {TensorType_FLOAT32, TensorType_INT8, TensorType_UINT8}) {
+    for (bool int64_indices : {false, true}) {
+      SliceTester()
+          .InputShape(input_shape)
+          .Offsets(offsets)
+          .Sizes(sizes)
+          .UseInt64OffsetsAndSize(int64_indices)
+          .RequireDelegation(true)
+          .Test(tensor_type, xnnpack_delegate.get());
+    }
+  }
+}
+
+TEST(Slice, RejectsMismatchedIndexLengthWithDynamicInput) {
+  std::unique_ptr<TfLiteDelegate, decltype(&TfLiteXNNPackDelegateDelete)>
+      xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
+                       TfLiteXNNPackDelegateDelete);
+
+  for (bool int64_indices : {false, true}) {
+    SliceTester()
+        .InputShape({2, 3})
+        .InputShapeSignature({-1, 3})
+        .Offsets({0})
+        .Sizes({1})
+        .OutputShape({1, 3})
+        .UseInt64OffsetsAndSize(int64_indices)
+        .TestUnsupported(TensorType_FLOAT32, xnnpack_delegate.get());
+  }
+}
+
+TEST(Slice, RejectsBeginPlusSizeInt64Overflow) {
+  std::unique_ptr<TfLiteDelegate, decltype(&TfLiteXNNPackDelegateDelete)>
+      xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
+                       TfLiteXNNPackDelegateDelete);
+
+  SliceTester()
+      .InputShape({2})
+      .InputShapeSignature({-1})
+      .OffsetsInt64({std::numeric_limits<int64_t>::max()})
+      .SizesInt64({2})
+      .OutputShape({2})
+      .UseInt64OffsetsAndSize(true)
+      .TestUnsupported(TensorType_FLOAT32, xnnpack_delegate.get());
+}
+
+TEST(Slice, RejectsNegativeBeginAndInvalidSize) {
+  std::unique_ptr<TfLiteDelegate, decltype(&TfLiteXNNPackDelegateDelete)>
+      xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
+                       TfLiteXNNPackDelegateDelete);
+
+  SliceTester()
+      .InputShape({4})
+      .InputShapeSignature({-1})
+      .Offsets({-1})
+      .Sizes({2})
+      .OutputShape({2})
+      .TestUnsupported(TensorType_FLOAT32, xnnpack_delegate.get());
+
+  SliceTester()
+      .InputShape({4})
+      .InputShapeSignature({-1})
+      .Offsets({0})
+      .Sizes({-2})
+      .OutputShape({2})
+      .TestUnsupported(TensorType_FLOAT32, xnnpack_delegate.get());
+}
+
 }  // namespace xnnpack
 }  // namespace tflite
diff --git a/tensorflow/lite/delegates/xnnpack/slice_tester.cc b/tensorflow/lite/delegates/xnnpack/slice_tester.cc
index 810e4470b74..1cb0fb6db5f 100644
--- a/tensorflow/lite/delegates/xnnpack/slice_tester.cc
+++ b/tensorflow/lite/delegates/xnnpack/slice_tester.cc
@@ -55,8 +55,7 @@ std::function<float(std::mt19937&)> GetDist() {
 template <typename T>
 void SliceTester::Test(Interpreter* default_interpreter,
                        Interpreter* delegate_interpreter) const {
-  std::random_device random_device;
-  auto rng = std::mt19937(random_device());
+  auto rng = std::mt19937();
   auto input_distribution = GetDist<T>();
   auto input_rng = std::bind(input_distribution, std::ref(rng));
   const auto* input = default_interpreter->input_tensor(0);
@@ -168,6 +167,35 @@ void SliceTester::Test(TensorType tensor_type, TfLiteDelegate* delegate) const {
   }
 }
 
+void SliceTester::TestUnsupported(TensorType tensor_type,
+                                  TfLiteDelegate* delegate) const {
+  const std::vector<char> buffer = CreateTfLiteModel(tensor_type);
+  const Model* model = GetModel(buffer.data());
+
+  std::unique_ptr<Interpreter> delegate_interpreter;
+  ASSERT_EQ(
+      InterpreterBuilder(
+          model,
+          ::tflite::ops::builtin::BuiltinOpResolverWithoutDefaultDelegates())(
+          &delegate_interpreter),
+      kTfLiteOk);
+  ASSERT_TRUE(delegate_interpreter);
+
+  ASSERT_EQ(delegate_interpreter->AllocateTensors(), kTfLiteOk);
+  ASSERT_EQ(delegate_interpreter->ModifyGraphWithDelegate(delegate), kTfLiteOk);
+  for (const auto& shape : reshape_input_shapes_) {
+    ASSERT_EQ(delegate_interpreter->ResizeInputTensor(0, shape), kTfLiteOk);
+    ASSERT_EQ(delegate_interpreter->AllocateTensors(), kTfLiteOk);
+  }
+  ASSERT_EQ(delegate_interpreter->execution_plan().size(), 1);
+  EXPECT_EQ(
+      delegate_interpreter
+          ->node_and_registration(delegate_interpreter->execution_plan()[0])
+          ->first.delegate,
+      nullptr);
+  EXPECT_NE(delegate_interpreter->Invoke(), kTfLiteOk);
+}
+
 std::vector<char> SliceTester::CreateTfLiteModel(TensorType tensor_type) const {
   flatbuffers::FlatBufferBuilder builder;
   flatbuffers::Offset<OperatorCode> operator_code =
@@ -189,22 +217,28 @@ std::vector<char> SliceTester::CreateTfLiteModel(TensorType tensor_type) const {
           builder.CreateVector<float>({/*scale=*/1.0f}),
           builder.CreateVector<int64_t>({/*zero_point=*/0}));
 
-  const int32_t num_dims = Offsets().size();
+  const int32_t begin_num_dims = OffsetsCount();
+  const int32_t size_num_dims = SizesCount();
   TensorType offsets_and_sizes_tensor_type =
       UseInt64OffsetsAndSize() ? TensorType_INT64 : TensorType_INT32;
+  const auto input_signature =
+      input_shape_signature_.empty()
+          ? flatbuffers::Offset<flatbuffers::Vector<int32_t>>()
+          : builder.CreateVector(input_shape_signature_);
   const auto output_signature =
       reshape_input_shapes_.empty()
           ? flatbuffers::Offset<flatbuffers::Vector<int32_t>>()
-          : builder.CreateVector(std::vector<int32_t>(num_dims, -1));
+          : builder.CreateVector(std::vector<int32_t>(begin_num_dims, -1));
 
   const std::array<flatbuffers::Offset<Tensor>, 4> tensors{{
       CreateTensor(builder,
                    builder.CreateVector<int32_t>(InputShape().data(),
                                                  InputShape().size()),
-                   tensor_type, /*buffer=*/0, /*name=*/0, quantization_params),
-      CreateTensor(builder, builder.CreateVector<int32_t>({num_dims}),
+                   tensor_type, /*buffer=*/0, /*name=*/0, quantization_params,
+                   /*is_variable=*/false, /*sparsity=*/0, input_signature),
+      CreateTensor(builder, builder.CreateVector<int32_t>({begin_num_dims}),
                    offsets_and_sizes_tensor_type, /*buffer=*/1),
-      CreateTensor(builder, builder.CreateVector<int32_t>({num_dims}),
+      CreateTensor(builder, builder.CreateVector<int32_t>({size_num_dims}),
                    offsets_and_sizes_tensor_type, /*buffer=*/2),
       CreateTensor(builder,
                    builder.CreateVector<int32_t>(OutputShape().data(),
diff --git a/tensorflow/lite/delegates/xnnpack/slice_tester.h b/tensorflow/lite/delegates/xnnpack/slice_tester.h
index 8af6a6d4be8..b28ad3da911 100644
--- a/tensorflow/lite/delegates/xnnpack/slice_tester.h
+++ b/tensorflow/lite/delegates/xnnpack/slice_tester.h
@@ -45,6 +45,16 @@ class SliceTester {
 
   inline const std::vector<int32_t>& InputShape() const { return input_shape_; }
 
+  inline SliceTester& InputShapeSignature(
+      const std::vector<int32_t>& shape_signature) {
+    input_shape_signature_ = shape_signature;
+    return *this;
+  }
+
+  inline const std::vector<int32_t>& InputShapeSignature() const {
+    return input_shape_signature_;
+  }
+
   inline SliceTester& Offsets(const std::vector<int32_t>& offsets) {
     offsets_ = std::vector<int32_t>(offsets);
     offsets_int64_.resize(offsets.size());
@@ -52,8 +62,18 @@ class SliceTester {
     return *this;
   }
 
+  inline SliceTester& OffsetsInt64(const std::vector<int64_t>& offsets) {
+    offsets_int64_ = offsets;
+    use_int64_offsets_and_sizes_ = true;
+    return *this;
+  }
+
   inline const std::vector<int32_t> Offsets() const { return offsets_; }
 
+  inline size_t OffsetsCount() const {
+    return UseInt64OffsetsAndSize() ? offsets_int64_.size() : offsets_.size();
+  }
+
   inline const void* OffsetsData() const {
     return UseInt64OffsetsAndSize() ? (void*)offsets_int64_.data()
                                     : (void*)offsets_.data();
@@ -61,7 +81,7 @@ class SliceTester {
 
   inline const size_t OffsetsSizeInBytes() const {
     if (use_int64_offsets_and_sizes_) {
-      return Offsets().size() * sizeof(int64_t);
+      return offsets_int64_.size() * sizeof(int64_t);
     } else {
       return Offsets().size() * sizeof(int32_t);
     }
@@ -71,17 +91,30 @@ class SliceTester {
     sizes_ = std::vector<int32_t>(sizes);
     sizes_int64_.resize(sizes.size());
     std::copy(sizes_.begin(), sizes_.end(), sizes_int64_.begin());
-    output_shape_ = std::vector<int32_t>(sizes);
-    for (size_t i = 0; i < sizes.size(); i++) {
-      if (output_shape_[i] < 0) {
-        output_shape_[i] = input_shape_[i] - offsets_[i];
+    if (!has_explicit_output_shape_) {
+      output_shape_ = std::vector<int32_t>(sizes);
+      for (size_t i = 0; i < sizes.size(); i++) {
+        if (output_shape_[i] < 0 && i < input_shape_.size() &&
+            i < offsets_.size()) {
+          output_shape_[i] = input_shape_[i] - offsets_[i];
+        }
       }
     }
     return *this;
   }
 
+  inline SliceTester& SizesInt64(const std::vector<int64_t>& sizes) {
+    sizes_int64_ = sizes;
+    use_int64_offsets_and_sizes_ = true;
+    return *this;
+  }
+
   inline const std::vector<int32_t> Sizes() const { return sizes_; }
 
+  inline size_t SizesCount() const {
+    return UseInt64OffsetsAndSize() ? sizes_int64_.size() : sizes_.size();
+  }
+
   inline const void* SizesData() const {
     return UseInt64OffsetsAndSize() ? (void*)sizes_int64_.data()
                                     : (void*)sizes_.data();
@@ -89,12 +122,18 @@ class SliceTester {
 
   inline const size_t SizesSizeInBytes() const {
     if (use_int64_offsets_and_sizes_) {
-      return Sizes().size() * sizeof(int64_t);
+      return sizes_int64_.size() * sizeof(int64_t);
     } else {
       return Sizes().size() * sizeof(int32_t);
     }
   }
 
+  inline SliceTester& OutputShape(const std::vector<int32_t>& shape) {
+    output_shape_ = shape;
+    has_explicit_output_shape_ = true;
+    return *this;
+  }
+
   std::vector<int32_t> OutputShape() const { return output_shape_; }
 
   inline SliceTester& UseInt64OffsetsAndSize(bool use_int64) {
@@ -118,6 +157,7 @@ class SliceTester {
   }
 
   void Test(TensorType tensor_type, TfLiteDelegate* delegate) const;
+  void TestUnsupported(TensorType tensor_type, TfLiteDelegate* delegate) const;
   template <typename T>
   void Test(Interpreter* default_interpreter,
             Interpreter* delegate_interpreter) const;
@@ -126,11 +166,13 @@ class SliceTester {
   std::vector<char> CreateTfLiteModel(TensorType tensor_type) const;
 
   std::vector<int32_t> input_shape_;
+  std::vector<int32_t> input_shape_signature_;
   std::vector<int32_t> offsets_;
   std::vector<int64_t> offsets_int64_;
   std::vector<int32_t> sizes_;
   std::vector<int64_t> sizes_int64_;
   std::vector<int32_t> output_shape_;
+  bool has_explicit_output_shape_ = false;
   bool use_int64_offsets_and_sizes_ = false;
   std::vector<std::vector<int32_t>> reshape_input_shapes_;
   bool require_delegation_ = false;
diff --git a/tensorflow/lite/delegates/xnnpack/unsigned_quantized_slice_test.cc b/tensorflow/lite/delegates/xnnpack/unsigned_quantized_slice_test.cc
index f39612f8f35..565d379767c 100644
--- a/tensorflow/lite/delegates/xnnpack/unsigned_quantized_slice_test.cc
+++ b/tensorflow/lite/delegates/xnnpack/unsigned_quantized_slice_test.cc
@@ -31,8 +31,7 @@ TEST(UnsignedQuantizedSlice, 1D) {
       xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
-  std::random_device random_device;
-  auto rng = std::mt19937(random_device());
+  auto rng = std::mt19937();
   auto shape_rng =
       std::bind(std::uniform_int_distribution<int32_t>(2, 5), std::ref(rng));
 
@@ -40,6 +39,12 @@ TEST(UnsignedQuantizedSlice, 1D) {
   const auto offsets = RandomOffsets(rng, input_shape);
   const auto sizes = RandomSizes(rng, input_shape, offsets);
 
+  SliceTester()
+      .InputShape(input_shape)
+      .Offsets(offsets)
+      .Sizes(sizes)
+      .Test(TensorType_UINT8, xnnpack_delegate.get());
+
   SliceTester()
       .InputShape(input_shape)
       .Offsets(offsets)
@@ -53,8 +58,7 @@ TEST(UnsignedQuantizedSlice, 2D) {
       xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
-  std::random_device random_device;
-  auto rng = std::mt19937(random_device());
+  auto rng = std::mt19937();
   auto shape_rng =
       std::bind(std::uniform_int_distribution<int32_t>(2, 5), std::ref(rng));
 
@@ -62,6 +66,12 @@ TEST(UnsignedQuantizedSlice, 2D) {
   const auto offsets = RandomOffsets(rng, input_shape);
   const auto sizes = RandomSizes(rng, input_shape, offsets);
 
+  SliceTester()
+      .InputShape(input_shape)
+      .Offsets(offsets)
+      .Sizes(sizes)
+      .Test(TensorType_UINT8, xnnpack_delegate.get());
+
   SliceTester()
       .InputShape(input_shape)
       .Offsets(offsets)
@@ -75,8 +85,7 @@ TEST(UnsignedQuantizedSlice, 3D) {
       xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
-  std::random_device random_device;
-  auto rng = std::mt19937(random_device());
+  auto rng = std::mt19937();
   auto shape_rng =
       std::bind(std::uniform_int_distribution<int32_t>(2, 5), std::ref(rng));
 
@@ -85,6 +94,12 @@ TEST(UnsignedQuantizedSlice, 3D) {
   const auto offsets = RandomOffsets(rng, input_shape);
   const auto sizes = RandomSizes(rng, input_shape, offsets);
 
+  SliceTester()
+      .InputShape(input_shape)
+      .Offsets(offsets)
+      .Sizes(sizes)
+      .Test(TensorType_UINT8, xnnpack_delegate.get());
+
   SliceTester()
       .InputShape(input_shape)
       .Offsets(offsets)
@@ -98,8 +113,7 @@ TEST(UnsignedQuantizedSlice, 4D) {
       xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
-  std::random_device random_device;
-  auto rng = std::mt19937(random_device());
+  auto rng = std::mt19937();
   auto shape_rng =
       std::bind(std::uniform_int_distribution<int32_t>(2, 5), std::ref(rng));
 
@@ -108,6 +122,12 @@ TEST(UnsignedQuantizedSlice, 4D) {
   const auto offsets = RandomOffsets(rng, input_shape);
   const auto sizes = RandomSizes(rng, input_shape, offsets);
 
+  SliceTester()
+      .InputShape(input_shape)
+      .Offsets(offsets)
+      .Sizes(sizes)
+      .Test(TensorType_UINT8, xnnpack_delegate.get());
+
   SliceTester()
       .InputShape(input_shape)
       .Offsets(offsets)
@@ -121,8 +141,7 @@ TEST(UnsignedQuantizedSlice, 5D) {
       xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
                        TfLiteXNNPackDelegateDelete);
 
-  std::random_device random_device;
-  auto rng = std::mt19937(random_device());
+  auto rng = std::mt19937();
   auto shape_rng =
       std::bind(std::uniform_int_distribution<int32_t>(2, 5), std::ref(rng));
 
@@ -131,6 +150,41 @@ TEST(UnsignedQuantizedSlice, 5D) {
   const auto offsets = RandomOffsets(rng, input_shape);
   const auto sizes = RandomSizes(rng, input_shape, offsets);
 
+  SliceTester()
+      .InputShape(input_shape)
+      .Offsets(offsets)
+      .Sizes(sizes)
+      .Test(TensorType_UINT8, xnnpack_delegate.get());
+
+  SliceTester()
+      .InputShape(input_shape)
+      .Offsets(offsets)
+      .Sizes(sizes)
+      .UseInt64OffsetsAndSize(true)
+      .Test(TensorType_UINT8, xnnpack_delegate.get());
+}
+
+TEST(UnsignedQuantizedSlice, 6D) {
+  std::unique_ptr<TfLiteDelegate, decltype(&TfLiteXNNPackDelegateDelete)>
+      xnnpack_delegate(TfLiteXNNPackDelegateCreate(nullptr),
+                       TfLiteXNNPackDelegateDelete);
+
+  auto rng = std::mt19937();
+  auto shape_rng =
+      std::bind(std::uniform_int_distribution<int32_t>(2, 4), std::ref(rng));
+
+  const std::vector<int32_t> input_shape = {shape_rng(), shape_rng(),
+                                            shape_rng(), shape_rng(),
+                                            shape_rng(), shape_rng()};
+  const auto offsets = RandomOffsets(rng, input_shape);
+  const auto sizes = RandomSizes(rng, input_shape, offsets);
+
+  SliceTester()
+      .InputShape(input_shape)
+      .Offsets(offsets)
+      .Sizes(sizes)
+      .Test(TensorType_UINT8, xnnpack_delegate.get());
+
   SliceTester()
       .InputShape(input_shape)
       .Offsets(offsets)
diff --git a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
index f96b28b4c3b..8a6e1958130 100644
--- a/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
+++ b/tensorflow/lite/delegates/xnnpack/xnnpack_delegate.cc
@@ -5990,6 +5990,14 @@ class Subgraph {
         logging_context, begin_tensor, size_tensor, 0, node_index, "SLICE"));
 
     const int num_dims = begin_tensor.dims->data[0];
+    if (num_dims != input_tensor.dims->size) {
+      TF_LITE_MAYBE_KERNEL_LOG(
+          logging_context,
+          "begin_tensor size (%d) must be equal to input rank (%d) in SLICE "
+          "node #%d",
+          num_dims, input_tensor.dims->size, node_index);
+      return kTfLiteError;
+    }
     if (num_dims > XNN_MAX_TENSOR_DIMS) {
       TF_LITE_MAYBE_KERNEL_LOG(
           logging_context,
@@ -6026,6 +6034,14 @@ class Subgraph {
                                  size[i], node_index);
         return kTfLiteError;
       }
+      if (size[i] > 0 &&
+          begin[i] > std::numeric_limits<int64_t>::max() - size[i]) {
+        TF_LITE_MAYBE_KERNEL_LOG(logging_context,
+                                 "begin %" PRId64 " + size %" PRId64
+                                 " overflows in SLICE node #%d",
+                                 begin[i], size[i], node_index);
+        return kTfLiteError;
+      }
     }
 
     if (subgraph != nullptr) {
diff --git a/tensorflow/lite/kernels/slice.cc b/tensorflow/lite/kernels/slice.cc
index 8d9bc3af5f2..e308eca0c9c 100644
--- a/tensorflow/lite/kernels/slice.cc
+++ b/tensorflow/lite/kernels/slice.cc
@@ -64,6 +64,10 @@ TfLiteStatus CalculateOutputShapeVector(TfLiteContext* context,
   // the input rank, so establish that here before either is dereferenced.
   TF_LITE_ENSURE_EQ(context, NumElements(begin), input_dims);
   TF_LITE_ENSURE_EQ(context, NumElements(size), input_dims);
+  if (input_dims > 0) {
+    TF_LITE_ENSURE(context, GetTensorData<T>(begin) != nullptr);
+    TF_LITE_ENSURE(context, GetTensorData<T>(size) != nullptr);
+  }
   for (int idx = 0; idx < input_dims; ++idx) {
     const T dim_value = static_cast<T>(SizeOfDimension(input, idx));
     const T begin_value = GetTensorData<T>(begin)[idx];
@@ -186,6 +190,7 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
                  begin->type == kTfLiteInt32 || begin->type == kTfLiteInt64);
   TF_LITE_ENSURE(context,
                  size->type == kTfLiteInt32 || size->type == kTfLiteInt64);
+  TF_LITE_ENSURE_TYPES_EQ(context, begin->type, size->type);
   TF_LITE_ENSURE_EQ(context, NumDimensions(begin), 1);
   TF_LITE_ENSURE_EQ(context, NumDimensions(size), 1);
   TF_LITE_ENSURE_EQ(context, NumElements(begin), NumElements(size));
diff --git a/tensorflow/lite/kernels/slice_test.cc b/tensorflow/lite/kernels/slice_test.cc
index 55ef81e9bb8..ecf43985fe5 100644
--- a/tensorflow/lite/kernels/slice_test.cc
+++ b/tensorflow/lite/kernels/slice_test.cc
@@ -15,6 +15,8 @@ limitations under the License.
 #include <stdint.h>
 
 #include <initializer_list>
+#include <limits>
+#include <memory>
 #include <string>
 #include <type_traits>
 #include <vector>
@@ -31,9 +33,16 @@ limitations under the License.
 #include "tensorflow/lite/types/half.h"
 
 namespace tflite {
+namespace ops {
+namespace builtin {
+TfLiteRegistration* Register_SLICE_REF();
+}  // namespace builtin
+}  // namespace ops
+
 namespace {
 
 using ::testing::ElementsAreArray;
+using ::testing::IsEmpty;
 
 enum class TestType {
   kConst = 0,
@@ -50,7 +59,9 @@ class SliceOpModel : public SingleOpModel {
                std::initializer_list<index_type> size_data,
                TensorType tensor_index_type, TensorType tensor_input_type,
                TestType input_tensor_types,
-               std::initializer_list<int> output_shape = {}) {
+               std::initializer_list<int> output_shape = {},
+               bool allocate_and_delegate = true,
+               bool use_reference_kernel = false) {
     input_ = AddInput(tensor_input_type);
     if (input_tensor_types == TestType::kDynamic) {
       begin_ = AddInput(tensor_index_type);
@@ -63,9 +74,15 @@ class SliceOpModel : public SingleOpModel {
     output_ = AddOutput(TensorData(tensor_input_type, output_shape));
     SetBuiltinOp(BuiltinOperator_SLICE, BuiltinOptions_SliceOptions,
                  CreateSliceOptions(builder_).Union());
-    BuildInterpreter({input_shape, begin_shape, size_shape});
+    if (use_reference_kernel) {
+      SetResolver(std::make_unique<SingleOpResolver>(
+          BuiltinOperator_SLICE, ops::builtin::Register_SLICE_REF()));
+    }
+    BuildInterpreter({input_shape, begin_shape, size_shape}, /*num_threads=*/-1,
+                     /*allow_fp32_relax_to_fp16=*/false,
+                     /*apply_delegate=*/true, allocate_and_delegate);
 
-    if (input_tensor_types == TestType::kDynamic) {
+    if (allocate_and_delegate && input_tensor_types == TestType::kDynamic) {
       PopulateTensor<index_type>(begin_, begin_data);
       PopulateTensor<index_type>(size_, size_data);
     }
@@ -100,6 +117,40 @@ class SliceOpModel : public SingleOpModel {
   int output_;
 };
 
+template <typename begin_type, typename size_type>
+class MismatchedIndexSliceOpModel : public SingleOpModel {
+ public:
+  MismatchedIndexSliceOpModel(std::initializer_list<int> input_shape,
+                              std::initializer_list<int> begin_shape,
+                              std::initializer_list<begin_type> begin_data,
+                              std::initializer_list<int> size_shape,
+                              std::initializer_list<size_type> size_data,
+                              TestType input_tensor_types) {
+    input_ = AddInput(TensorType_FLOAT32);
+    if (input_tensor_types == TestType::kDynamic) {
+      begin_ = AddInput(GetTensorType<begin_type>());
+      size_ = AddInput(GetTensorType<size_type>());
+    } else {
+      begin_ =
+          AddConstInput(GetTensorType<begin_type>(), begin_data, begin_shape);
+      size_ = AddConstInput(GetTensorType<size_type>(), size_data, size_shape);
+    }
+    output_ = AddOutput(TensorType_FLOAT32);
+    SetBuiltinOp(BuiltinOperator_SLICE, BuiltinOptions_SliceOptions,
+                 CreateSliceOptions(builder_).Union());
+    BuildInterpreter({input_shape, begin_shape, size_shape}, /*num_threads=*/-1,
+                     /*allow_fp32_relax_to_fp16=*/false,
+                     /*apply_delegate=*/true,
+                     /*allocate_and_delegate=*/false);
+  }
+
+ private:
+  int input_;
+  int begin_;
+  int size_;
+  int output_;
+};
+
 class SliceOpTest : public ::testing::TestWithParam<TestType> {};
 
 TEST_P(SliceOpTest, In1D) {
@@ -408,6 +459,388 @@ TEST_P(SliceOpTest, BeginNonZeroSizeMinus1Axis1BFloat16) {
                                 Eigen::bfloat16(8), Eigen::bfloat16(9)}));
 }
 
+TEST_P(SliceOpTest, SliceUint16) {
+  SliceOpModel<uint16_t, int32_t> m({3, 2, 3, 1}, {4}, {1, 0, 0, 0}, {4},
+                                    {2, 1, -1, 1}, TensorType_INT32,
+                                    TensorType_UINT16, GetParam());
+  m.SetInput({1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5, 6, 6, 6});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 1, 3, 1}));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({3, 3, 3, 5, 5, 5}));
+}
+
+TEST_P(SliceOpTest, SliceUint64) {
+  SliceOpModel<uint64_t, int32_t> m({3, 2, 3, 1}, {4}, {1, 0, 0, 0}, {4},
+                                    {2, 1, -1, 1}, TensorType_INT32,
+                                    TensorType_UINT64, GetParam());
+  m.SetInput({1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5, 6, 6, 6});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 1, 3, 1}));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({3, 3, 3, 5, 5, 5}));
+}
+
+TEST_P(SliceOpTest, SliceFloat64) {
+  SliceOpModel<double, int32_t> m({3, 2, 3, 1}, {4}, {1, 0, 0, 0}, {4},
+                                  {2, 1, -1, 1}, TensorType_INT32,
+                                  TensorType_FLOAT64, GetParam());
+  m.SetInput({1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4, 5, 5, 5, 6, 6, 6});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 1, 3, 1}));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({3, 3, 3, 5, 5, 5}));
+}
+
+TEST_P(SliceOpTest, AllDataTypesWithInt64Indices) {
+  {
+    SliceOpModel<uint8_t, int64_t> m({2, 3}, {2}, {1, 1}, {2}, {1, -1},
+                                     TensorType_INT64, TensorType_UINT8,
+                                     GetParam());
+    m.SetInput({1, 2, 3, 4, 5, 6});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
+    EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 6}));
+  }
+  {
+    SliceOpModel<int8_t, int64_t> m({2, 3}, {2}, {1, 1}, {2}, {1, -1},
+                                    TensorType_INT64, TensorType_INT8,
+                                    GetParam());
+    m.SetInput({-1, -2, -3, -4, 5, 6});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
+    EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 6}));
+  }
+  {
+    SliceOpModel<bool, int64_t> m({2, 3}, {2}, {1, 1}, {2}, {1, -1},
+                                  TensorType_INT64, TensorType_BOOL,
+                                  GetParam());
+    m.SetInput({true, false, true, false, true, false});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
+    EXPECT_THAT(m.GetOutput(), ElementsAreArray({true, false}));
+  }
+  {
+    SliceOpModel<int16_t, int64_t> m({2, 3}, {2}, {1, 1}, {2}, {1, -1},
+                                     TensorType_INT64, TensorType_INT16,
+                                     GetParam());
+    m.SetInput({10, 20, 30, 40, -50, 60});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
+    EXPECT_THAT(m.GetOutput(), ElementsAreArray({-50, 60}));
+  }
+  {
+    SliceOpModel<uint16_t, int64_t> m({2, 3}, {2}, {1, 1}, {2}, {1, -1},
+                                      TensorType_INT64, TensorType_UINT16,
+                                      GetParam());
+    m.SetInput({10, 20, 30, 40, 50, 60});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
+    EXPECT_THAT(m.GetOutput(), ElementsAreArray({50, 60}));
+  }
+  {
+    SliceOpModel<half, int64_t> m({2, 3}, {2}, {1, 1}, {2}, {1, -1},
+                                  TensorType_INT64, TensorType_FLOAT16,
+                                  GetParam());
+    m.SetInput({half(1), half(2), half(3), half(4), half(5), half(6)});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
+    EXPECT_THAT(m.GetOutput(), ElementsAreArray({half(5), half(6)}));
+  }
+  {
+    SliceOpModel<Eigen::bfloat16, int64_t> m({2, 3}, {2}, {1, 1}, {2}, {1, -1},
+                                             TensorType_INT64,
+                                             TensorType_BFLOAT16, GetParam());
+    m.SetInput({Eigen::bfloat16(1), Eigen::bfloat16(2), Eigen::bfloat16(3),
+                Eigen::bfloat16(4), Eigen::bfloat16(5), Eigen::bfloat16(6)});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
+    EXPECT_THAT(m.GetOutput(),
+                ElementsAreArray({Eigen::bfloat16(5), Eigen::bfloat16(6)}));
+  }
+  {
+    SliceOpModel<int32_t, int64_t> m({2, 3}, {2}, {1, 1}, {2}, {1, -1},
+                                     TensorType_INT64, TensorType_INT32,
+                                     GetParam());
+    m.SetInput({1, 2, 3, 4, 5, 6});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
+    EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 6}));
+  }
+  {
+    SliceOpModel<uint32_t, int64_t> m({2, 3}, {2}, {1, 1}, {2}, {1, -1},
+                                      TensorType_INT64, TensorType_UINT32,
+                                      GetParam());
+    m.SetInput({1, 2, 3, 4, 5, 6});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
+    EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 6}));
+  }
+  {
+    SliceOpModel<int64_t, int64_t> m({2, 3}, {2}, {1, 1}, {2}, {1, -1},
+                                     TensorType_INT64, TensorType_INT64,
+                                     GetParam());
+    m.SetInput({1, 2, 3, 4, 5, 6});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
+    EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 6}));
+  }
+  {
+    SliceOpModel<uint64_t, int64_t> m({2, 3}, {2}, {1, 1}, {2}, {1, -1},
+                                      TensorType_INT64, TensorType_UINT64,
+                                      GetParam());
+    m.SetInput({1, 2, 3, 4, 5, 6});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
+    EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 6}));
+  }
+  {
+    SliceOpModel<double, int64_t> m({2, 3}, {2}, {1, 1}, {2}, {1, -1},
+                                    TensorType_INT64, TensorType_FLOAT64,
+                                    GetParam());
+    m.SetInput({1.0, 2.0, 3.0, 4.0, 5.5, 6.5});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
+    EXPECT_THAT(m.GetOutput(), ElementsAreArray({5.5, 6.5}));
+  }
+  {
+    SliceOpModel<int8_t, int64_t> m({2, 3}, {2}, {1, 1}, {2}, {1, -1},
+                                    TensorType_INT64, TensorType_INT4,
+                                    GetParam());
+    m.SetInput({1, 2, 3, 4, -5, 6});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
+    const TfLiteTensor* output_tensor = m.GetOutputTensor();
+    const int num_elements = NumElements(output_tensor);
+    std::vector<int8_t> unpacked_output(num_elements);
+    tensor_utils::UnpackPackedIntToInt8(GetTensorData<int8_t>(output_tensor),
+                                        num_elements, /*bit_width=*/4,
+                                        unpacked_output.data());
+    EXPECT_THAT(unpacked_output, ElementsAreArray({-5, 6}));
+  }
+  {
+    SliceOpModel<string, int64_t> m({2, 3}, {2}, {1, 1}, {2}, {1, -1},
+                                    TensorType_INT64, TensorType_STRING,
+                                    GetParam());
+    m.SetStringInput({"a", "b", "c", "d", "e", "f"});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
+    EXPECT_THAT(m.GetOutput(), ElementsAreArray({"e", "f"}));
+  }
+}
+
+// Covers the `input_dims > kMaxDim` (6D) dynamic-rank reference path in Eval()
+// across all supported bit widths (8, 16, 32, 64), INT4, and STRING, for both
+// optimized and reference registrations.
+TEST_P(SliceOpTest, In6DAllDataTypesAndKernelVariants) {
+  for (bool use_reference : {false, true}) {
+    {
+      SliceOpModel<float, int64_t> m(
+          {2, 1, 1, 1, 1, 3}, {6}, {1, 0, 0, 0, 0, 1}, {6}, {1, 1, 1, 1, 1, -1},
+          TensorType_INT64, TensorType_FLOAT32, GetParam(), {},
+          /*allocate_and_delegate=*/true, use_reference);
+      m.SetInput({1, 2, 3, 4, 5, 6});
+      ASSERT_EQ(m.Invoke(), kTfLiteOk);
+      EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 1, 1, 1, 2}));
+      EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 6}));
+    }
+    {
+      SliceOpModel<int32_t, int64_t> m(
+          {2, 2, 1, 1, 2, 2}, {6}, {0, 1, 0, 0, 0, 1}, {6}, {2, 1, 1, 1, -1, 1},
+          TensorType_INT64, TensorType_INT32, GetParam(), {},
+          /*allocate_and_delegate=*/true, use_reference);
+      m.SetInput({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16});
+      ASSERT_EQ(m.Invoke(), kTfLiteOk);
+      EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({2, 1, 1, 1, 2, 1}));
+      EXPECT_THAT(m.GetOutput(), ElementsAreArray({6, 8, 14, 16}));
+    }
+    {
+      SliceOpModel<uint32_t, int32_t> m(
+          {2, 1, 1, 1, 1, 3}, {6}, {1, 0, 0, 0, 0, 1}, {6}, {1, 1, 1, 1, 1, 2},
+          TensorType_INT32, TensorType_UINT32, GetParam(), {},
+          /*allocate_and_delegate=*/true, use_reference);
+      m.SetInput({1, 2, 3, 4, 5, 6});
+      ASSERT_EQ(m.Invoke(), kTfLiteOk);
+      EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 1, 1, 1, 2}));
+      EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 6}));
+    }
+    {
+      SliceOpModel<int64_t, int64_t> m(
+          {2, 1, 1, 1, 1, 3}, {6}, {1, 0, 0, 0, 0, 1}, {6}, {1, 1, 1, 1, 1, 2},
+          TensorType_INT64, TensorType_INT64, GetParam(), {},
+          /*allocate_and_delegate=*/true, use_reference);
+      m.SetInput({1, 2, 3, 4, 5, 6});
+      ASSERT_EQ(m.Invoke(), kTfLiteOk);
+      EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 1, 1, 1, 2}));
+      EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 6}));
+    }
+    {
+      SliceOpModel<uint64_t, int32_t> m(
+          {2, 1, 1, 1, 1, 3}, {6}, {1, 0, 0, 0, 0, 1}, {6}, {1, 1, 1, 1, 1, 2},
+          TensorType_INT32, TensorType_UINT64, GetParam(), {},
+          /*allocate_and_delegate=*/true, use_reference);
+      m.SetInput({1, 2, 3, 4, 5, 6});
+      ASSERT_EQ(m.Invoke(), kTfLiteOk);
+      EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 1, 1, 1, 2}));
+      EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 6}));
+    }
+    {
+      SliceOpModel<double, int32_t> m(
+          {2, 1, 1, 1, 1, 3}, {6}, {1, 0, 0, 0, 0, 1}, {6}, {1, 1, 1, 1, 1, 2},
+          TensorType_INT32, TensorType_FLOAT64, GetParam(), {},
+          /*allocate_and_delegate=*/true, use_reference);
+      m.SetInput({1.0, 2.0, 3.0, 4.0, 5.0, 6.0});
+      ASSERT_EQ(m.Invoke(), kTfLiteOk);
+      EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 1, 1, 1, 2}));
+      EXPECT_THAT(m.GetOutput(), ElementsAreArray({5.0, 6.0}));
+    }
+    {
+      SliceOpModel<int16_t, int32_t> m(
+          {2, 1, 1, 1, 1, 3}, {6}, {1, 0, 0, 0, 0, 1}, {6}, {1, 1, 1, 1, 1, 2},
+          TensorType_INT32, TensorType_INT16, GetParam(), {},
+          /*allocate_and_delegate=*/true, use_reference);
+      m.SetInput({1, 2, 3, 4, 5, 6});
+      ASSERT_EQ(m.Invoke(), kTfLiteOk);
+      EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 1, 1, 1, 2}));
+      EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 6}));
+    }
+    {
+      SliceOpModel<uint16_t, int32_t> m(
+          {2, 1, 1, 1, 1, 3}, {6}, {1, 0, 0, 0, 0, 1}, {6}, {1, 1, 1, 1, 1, 2},
+          TensorType_INT32, TensorType_UINT16, GetParam(), {},
+          /*allocate_and_delegate=*/true, use_reference);
+      m.SetInput({1, 2, 3, 4, 5, 6});
+      ASSERT_EQ(m.Invoke(), kTfLiteOk);
+      EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 1, 1, 1, 2}));
+      EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 6}));
+    }
+    {
+      SliceOpModel<half, int32_t> m(
+          {2, 1, 1, 1, 1, 3}, {6}, {1, 0, 0, 0, 0, 1}, {6}, {1, 1, 1, 1, 1, 2},
+          TensorType_INT32, TensorType_FLOAT16, GetParam(), {},
+          /*allocate_and_delegate=*/true, use_reference);
+      m.SetInput({half(1), half(2), half(3), half(4), half(5), half(6)});
+      ASSERT_EQ(m.Invoke(), kTfLiteOk);
+      EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 1, 1, 1, 2}));
+      EXPECT_THAT(m.GetOutput(), ElementsAreArray({half(5), half(6)}));
+    }
+    {
+      SliceOpModel<Eigen::bfloat16, int32_t> m(
+          {2, 1, 1, 1, 1, 3}, {6}, {1, 0, 0, 0, 0, 1}, {6}, {1, 1, 1, 1, 1, 2},
+          TensorType_INT32, TensorType_BFLOAT16, GetParam(), {},
+          /*allocate_and_delegate=*/true, use_reference);
+      m.SetInput({Eigen::bfloat16(1), Eigen::bfloat16(2), Eigen::bfloat16(3),
+                  Eigen::bfloat16(4), Eigen::bfloat16(5), Eigen::bfloat16(6)});
+      ASSERT_EQ(m.Invoke(), kTfLiteOk);
+      EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 1, 1, 1, 2}));
+      EXPECT_THAT(m.GetOutput(),
+                  ElementsAreArray({Eigen::bfloat16(5), Eigen::bfloat16(6)}));
+    }
+    {
+      SliceOpModel<int8_t, int32_t> m(
+          {2, 1, 1, 1, 1, 3}, {6}, {1, 0, 0, 0, 0, 1}, {6}, {1, 1, 1, 1, 1, 2},
+          TensorType_INT32, TensorType_INT8, GetParam(), {},
+          /*allocate_and_delegate=*/true, use_reference);
+      m.SetInput({1, 2, 3, 4, 5, 6});
+      ASSERT_EQ(m.Invoke(), kTfLiteOk);
+      EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 1, 1, 1, 2}));
+      EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 6}));
+    }
+    {
+      SliceOpModel<uint8_t, int32_t> m(
+          {2, 1, 1, 1, 1, 3}, {6}, {1, 0, 0, 0, 0, 1}, {6}, {1, 1, 1, 1, 1, 2},
+          TensorType_INT32, TensorType_UINT8, GetParam(), {},
+          /*allocate_and_delegate=*/true, use_reference);
+      m.SetInput({1, 2, 3, 4, 5, 6});
+      ASSERT_EQ(m.Invoke(), kTfLiteOk);
+      EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 1, 1, 1, 2}));
+      EXPECT_THAT(m.GetOutput(), ElementsAreArray({5, 6}));
+    }
+    {
+      SliceOpModel<bool, int32_t> m(
+          {2, 1, 1, 1, 1, 3}, {6}, {1, 0, 0, 0, 0, 1}, {6}, {1, 1, 1, 1, 1, 2},
+          TensorType_INT32, TensorType_BOOL, GetParam(), {},
+          /*allocate_and_delegate=*/true, use_reference);
+      m.SetInput({true, false, true, false, true, false});
+      ASSERT_EQ(m.Invoke(), kTfLiteOk);
+      EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 1, 1, 1, 2}));
+      EXPECT_THAT(m.GetOutput(), ElementsAreArray({true, false}));
+    }
+    {
+      SliceOpModel<int8_t, int32_t> m(
+          {2, 1, 1, 1, 1, 3}, {6}, {1, 0, 0, 0, 0, 1}, {6}, {1, 1, 1, 1, 1, 2},
+          TensorType_INT32, TensorType_INT4, GetParam(), {},
+          /*allocate_and_delegate=*/true, use_reference);
+      m.SetInput({1, 2, 3, 4, -5, 6});
+      ASSERT_EQ(m.Invoke(), kTfLiteOk);
+      EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 1, 1, 1, 2}));
+      const TfLiteTensor* output_tensor = m.GetOutputTensor();
+      const int num_elements = NumElements(output_tensor);
+      std::vector<int8_t> unpacked_output(num_elements);
+      tensor_utils::UnpackPackedIntToInt8(GetTensorData<int8_t>(output_tensor),
+                                          num_elements, /*bit_width=*/4,
+                                          unpacked_output.data());
+      EXPECT_THAT(unpacked_output, ElementsAreArray({-5, 6}));
+    }
+    {
+      SliceOpModel<string, int32_t> m(
+          {2, 1, 1, 1, 1, 3}, {6}, {1, 0, 0, 0, 0, 1}, {6}, {1, 1, 1, 1, 1, 2},
+          TensorType_INT32, TensorType_STRING, GetParam(), {},
+          /*allocate_and_delegate=*/true, use_reference);
+      m.SetStringInput({"a", "b", "c", "d", "e", "f"});
+      ASSERT_EQ(m.Invoke(), kTfLiteOk);
+      EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 1, 1, 1, 1, 2}));
+      EXPECT_THAT(m.GetOutput(), ElementsAreArray({"e", "f"}));
+    }
+  }
+}
+
+TEST_P(SliceOpTest, ReferenceKernelUpTo5D) {
+  {
+    SliceOpModel<float, int32_t> m(
+        {2, 3}, {2}, {1, 0}, {2}, {1, 2}, TensorType_INT32, TensorType_FLOAT32,
+        GetParam(), {}, /*allocate_and_delegate=*/true,
+        /*use_reference_kernel=*/true);
+    m.SetInput({1, 2, 3, 4, 5, 6});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
+    EXPECT_THAT(m.GetOutput(), ElementsAreArray({4, 5}));
+  }
+  {
+    SliceOpModel<int8_t, int32_t> m(
+        {2, 3}, {2}, {1, 1}, {2}, {1, -1}, TensorType_INT32, TensorType_INT4,
+        GetParam(), {}, /*allocate_and_delegate=*/true,
+        /*use_reference_kernel=*/true);
+    m.SetInput({1, 2, 3, 4, -5, 6});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
+    const TfLiteTensor* output_tensor = m.GetOutputTensor();
+    const int num_elements = NumElements(output_tensor);
+    std::vector<int8_t> unpacked_output(num_elements);
+    tensor_utils::UnpackPackedIntToInt8(GetTensorData<int8_t>(output_tensor),
+                                        num_elements, /*bit_width=*/4,
+                                        unpacked_output.data());
+    EXPECT_THAT(unpacked_output, ElementsAreArray({-5, 6}));
+  }
+  {
+    SliceOpModel<string, int32_t> m(
+        {2, 3}, {2}, {1, 1}, {2}, {1, -1}, TensorType_INT32, TensorType_STRING,
+        GetParam(), {}, /*allocate_and_delegate=*/true,
+        /*use_reference_kernel=*/true);
+    m.SetStringInput({"a", "b", "c", "d", "e", "f"});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 2}));
+    EXPECT_THAT(m.GetOutput(), ElementsAreArray({"e", "f"}));
+  }
+}
+
+TEST_P(SliceOpTest, ZeroSizeSlice) {
+  SliceOpModel<float, int32_t> m({2, 3}, {2}, {1, 0}, {2}, {0, 2},
+                                 TensorType_INT32, TensorType_FLOAT32,
+                                 GetParam());
+  m.SetInput({1, 2, 3, 4, 5, 6});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({0, 2}));
+  EXPECT_THAT(m.GetOutput(), IsEmpty());
+}
+
 // A statically shaped output must keep its static allocation. Prepare() returns
 // early for such an output; this asserts that validating it does not convert it
 // to a dynamic tensor, which would change arena planning for valid models.
@@ -420,6 +853,28 @@ TEST(SliceOpTest, StaticOutputKeepsStaticAllocation) {
   EXPECT_NE(m.GetOutputTensor()->allocation_type, kTfLiteDynamic);
 }
 
+// `begin` and `size` must have matching index types (`INT32` or `INT64`).
+TEST_P(SliceOpTest, RejectsMismatchedBeginAndSizeIndexTypes) {
+  {
+    MismatchedIndexSliceOpModel<int32_t, int64_t> m({2, 3}, {2}, {0, 0}, {2},
+                                                    {1, 2}, GetParam());
+    EXPECT_NE(m.AllocateTensors(), kTfLiteOk);
+  }
+  {
+    MismatchedIndexSliceOpModel<int64_t, int32_t> m({2, 3}, {2}, {0, 0}, {2},
+                                                    {1, 2}, GetParam());
+    EXPECT_NE(m.AllocateTensors(), kTfLiteOk);
+  }
+}
+
+TEST_P(SliceOpTest, RejectsMismatchedBeginAndSizeElementCounts) {
+  SliceOpModel<float, int32_t> m({2, 3}, {1}, {0}, {2}, {1, 2},
+                                 TensorType_INT32, TensorType_FLOAT32,
+                                 GetParam(), {},
+                                 /*allocate_and_delegate=*/false);
+  EXPECT_NE(m.AllocateTensors(), kTfLiteOk);
+}
+
 // `begin` and `size` supplied at run time are not visible to Prepare(), which
 // returns early because the output shape is fully specified. Without validation
 // in Eval() the kernel reads past the end of the input: begin 3 + size 2
@@ -447,15 +902,119 @@ TEST(SliceOpTest, NegativeBeginWithStaticOutput) {
   EXPECT_NE(m.Invoke(), kTfLiteOk);
 }
 
+TEST(SliceOpTest, NegativeBeginRejectedForConstAndDynamic) {
+  {
+    SliceOpModel<float, int32_t> m({4}, {1}, {-1}, {1}, {2}, TensorType_INT32,
+                                   TensorType_FLOAT32, TestType::kConst, {},
+                                   /*allocate_and_delegate=*/false);
+    EXPECT_NE(m.AllocateTensors(), kTfLiteOk);
+  }
+  {
+    SliceOpModel<float, int64_t> m({4}, {1}, {-1}, {1}, {2}, TensorType_INT64,
+                                   TensorType_FLOAT32, TestType::kConst, {2},
+                                   /*allocate_and_delegate=*/false);
+    EXPECT_NE(m.AllocateTensors(), kTfLiteOk);
+  }
+  {
+    SliceOpModel<float, int64_t> m({4}, {1}, {-1}, {1}, {2}, TensorType_INT64,
+                                   TensorType_FLOAT32, TestType::kDynamic);
+    m.SetInput({1, 2, 3, 4});
+    EXPECT_NE(m.Invoke(), kTfLiteOk);
+  }
+}
+
+TEST(SliceOpTest, BeginGreaterThanInputDimensionRejected) {
+  {
+    SliceOpModel<float, int32_t> m({4}, {1}, {5}, {1}, {0}, TensorType_INT32,
+                                   TensorType_FLOAT32, TestType::kConst, {},
+                                   /*allocate_and_delegate=*/false);
+    EXPECT_NE(m.AllocateTensors(), kTfLiteOk);
+  }
+  {
+    SliceOpModel<float, int64_t> m({4}, {1}, {5}, {1}, {0}, TensorType_INT64,
+                                   TensorType_FLOAT32, TestType::kDynamic);
+    m.SetInput({1, 2, 3, 4});
+    EXPECT_NE(m.Invoke(), kTfLiteOk);
+  }
+}
+
+TEST(SliceOpTest, SizeLessThanMinusOneRejected) {
+  {
+    SliceOpModel<float, int32_t> m({4}, {1}, {0}, {1}, {-2}, TensorType_INT32,
+                                   TensorType_FLOAT32, TestType::kConst, {},
+                                   /*allocate_and_delegate=*/false);
+    EXPECT_NE(m.AllocateTensors(), kTfLiteOk);
+  }
+  {
+    SliceOpModel<float, int64_t> m({4}, {1}, {0}, {1}, {-2}, TensorType_INT64,
+                                   TensorType_FLOAT32, TestType::kDynamic, {2});
+    m.SetInput({1, 2, 3, 4});
+    EXPECT_NE(m.Invoke(), kTfLiteOk);
+  }
+}
+
+TEST(SliceOpTest, BeginPlusSizeOverflowRejected) {
+  {
+    SliceOpModel<float, int32_t> m(
+        {4}, {1}, {2}, {1}, {std::numeric_limits<int32_t>::max()},
+        TensorType_INT32, TensorType_FLOAT32, TestType::kDynamic, {2});
+    m.SetInput({1, 2, 3, 4});
+    EXPECT_NE(m.Invoke(), kTfLiteOk);
+  }
+  {
+    SliceOpModel<float, int64_t> m(
+        {4}, {1}, {2}, {1}, {std::numeric_limits<int64_t>::max()},
+        TensorType_INT64, TensorType_FLOAT32, TestType::kDynamic, {2});
+    m.SetInput({1, 2, 3, 4});
+    EXPECT_NE(m.Invoke(), kTfLiteOk);
+  }
+}
+
 // CalculateOutputShapeVector() indexes `begin` and `size` by input dimension.
 // Prepare() only checks the two against each other, so an index vector shorter
-// than the input rank is rejected here rather than read out of bounds.
+// or longer than the input rank is rejected rather than read out of bounds.
 TEST(SliceOpTest, IndexVectorShorterThanInputRank) {
-  SliceOpModel<float, int32_t> m({2, 3}, {1}, {0}, {1}, {1}, TensorType_INT32,
-                                 TensorType_FLOAT32, TestType::kDynamic,
-                                 {1, 1});
-  m.SetInput({1, 2, 3, 4, 5, 6});
-  EXPECT_NE(m.Invoke(), kTfLiteOk);
+  {
+    SliceOpModel<float, int32_t> m({2, 3}, {1}, {0}, {1}, {1}, TensorType_INT32,
+                                   TensorType_FLOAT32, TestType::kDynamic,
+                                   {1, 1});
+    m.SetInput({1, 2, 3, 4, 5, 6});
+    EXPECT_NE(m.Invoke(), kTfLiteOk);
+  }
+  {
+    SliceOpModel<float, int64_t> m({2, 3}, {1}, {0}, {1}, {1}, TensorType_INT64,
+                                   TensorType_FLOAT32, TestType::kDynamic);
+    m.SetInput({1, 2, 3, 4, 5, 6});
+    EXPECT_NE(m.Invoke(), kTfLiteOk);
+  }
+  {
+    SliceOpModel<float, int32_t> m({2, 3}, {1}, {0}, {1}, {1}, TensorType_INT32,
+                                   TensorType_FLOAT32, TestType::kConst, {},
+                                   /*allocate_and_delegate=*/false);
+    EXPECT_NE(m.AllocateTensors(), kTfLiteOk);
+  }
+  {
+    SliceOpModel<float, int64_t> m({2, 3}, {1}, {0}, {1}, {1}, TensorType_INT64,
+                                   TensorType_FLOAT32, TestType::kConst, {1, 1},
+                                   /*allocate_and_delegate=*/false);
+    EXPECT_NE(m.AllocateTensors(), kTfLiteOk);
+  }
+}
+
+TEST(SliceOpTest, IndexVectorLongerThanInputRank) {
+  {
+    SliceOpModel<float, int32_t> m(
+        {2}, {2}, {0, 0}, {2}, {1, 1}, TensorType_INT32, TensorType_FLOAT32,
+        TestType::kConst, {}, /*allocate_and_delegate=*/false);
+    EXPECT_NE(m.AllocateTensors(), kTfLiteOk);
+  }
+  {
+    SliceOpModel<float, int64_t> m({2}, {2}, {0, 0}, {2}, {1, 1},
+                                   TensorType_INT64, TensorType_FLOAT32,
+                                   TestType::kDynamic, {1});
+    m.SetInput({1, 2});
+    EXPECT_NE(m.Invoke(), kTfLiteOk);
+  }
 }
 
 // A valid window whose computed shape disagrees with the declared static output

From b5a93a5c161f43d749570f54dbfa275c47a258ed Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Wed, 14 Oct 2026 23:25:44 -0700
Subject: [PATCH 0835/1000] Test rank 5 in testHighRank and name the last
 dimension

Rank 5 is the highest that the oneDNN softmax kernel takes without
collapsing the leading dimensions, so test it next to ranks 6 and 7.
Read the last dimension's size into a variable before reassigning
src_shape.
---
 tensorflow/core/kernels/mkl/mkl_softmax_op.cc            | 4 ++--
 tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py | 5 +++--
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/kernels/mkl/mkl_softmax_op.cc b/tensorflow/core/kernels/mkl/mkl_softmax_op.cc
index c58684ecdcf..6dc6252a36e 100644
--- a/tensorflow/core/kernels/mkl/mkl_softmax_op.cc
+++ b/tensorflow/core/kernels/mkl/mkl_softmax_op.cc
@@ -235,8 +235,8 @@ class MklSoftmaxOp : public OpKernel {
         for (int i = 0; i < src_shape.dims() - 1; ++i) {
           num_rows *= src_shape.dim_size(i);
         }
-        src_shape =
-            TensorShape({num_rows, src_shape.dim_size(src_shape.dims() - 1)});
+        const int64_t last_dim = src_shape.dim_size(src_shape.dims() - 1);
+        src_shape = TensorShape({num_rows, last_dim});
       }
       const int input_dims = src_shape.dims();
       memory::format_tag src_fmt;
diff --git a/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py b/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py
index 91fb22161c1..8bd45a6b74d 100644
--- a/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py
+++ b/tensorflow/python/kernel_tests/nn_ops/softmax_op_test.py
@@ -279,9 +279,10 @@ class SoftmaxTest(test.TestCase):
   @test_util.run_in_graph_and_eager_modes(use_gpu=False)
   def testHighRank(self):
     # Regression test for GitHub issue 128814: the oneDNN kernel rejected
-    # inputs of more than 5 dimensions.
+    # inputs of more than 5 dimensions. It now collapses the leading
+    # dimensions of those, and rank 5 is the highest it takes as is.
     rng = np.random.RandomState(0)
-    for rank in (6, 7):
+    for rank in (5, 6, 7):
       for empty_dim in (None, 2, -1):
         shape = [2] * (rank - 1) + [3]
         if empty_dim is not None:

From 488d68dde37163ab247bb88271680727a5d56449 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Thu, 15 Oct 2026 13:04:39 +0530
Subject: [PATCH 0836/1000] Address non-blocking review suggestions: ordering
 invariant test, defensive None check, collective_group_leader test

---
 tensorflow/python/eager/remote.py      |  4 +++-
 tensorflow/python/eager/remote_test.py | 23 +++++++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/tensorflow/python/eager/remote.py b/tensorflow/python/eager/remote.py
index 19cafcc1da8..5224c2c56d1 100644
--- a/tensorflow/python/eager/remote.py
+++ b/tensorflow/python/eager/remote.py
@@ -84,7 +84,9 @@ def _parse_host_and_port(address):
     A tuple `(host, port)`. `port` is None if the address has no port, and
     both are None if the address cannot be parsed.
   """
-  if "://" not in address:
+  if not address or not isinstance(address, str):
+    return None, None
+  if "://" not in address and not address.startswith("//"):
     address = "//" + address
   try:
     parsed = urllib.parse.urlsplit(address)
diff --git a/tensorflow/python/eager/remote_test.py b/tensorflow/python/eager/remote_test.py
index e0d9d6106cd..dd15d98bf0b 100644
--- a/tensorflow/python/eager/remote_test.py
+++ b/tensorflow/python/eager/remote_test.py
@@ -766,19 +766,37 @@ class MultiJobsTest(test.TestCase, parameterized.TestCase):
 
   def testConnectToClusterWithMissingMaster(self):
     resolver = SimpleClusterResolver(self._cluster, master='localhost:80')
+    # The master lookup must be validated before the eager context's
+    # ServerDef is mutated, so a validation failure here must leave the
+    # context's ServerDef untouched (see testConnectToClusterOrderingInvariant
+    # for the dedicated regression test of that ordering).
+    initial_server_def = context.get_server_def()
     with self.assertRaisesRegex(
         ValueError,
         '`make_master_device_default` is set to True but cannot find master '
         'localhost:80 in the cluster'):
       remote.connect_to_cluster(resolver)
+    self.assertEqual(context.get_server_def(), initial_server_def)
 
   def testConnectToClusterWithInvalidMaster(self):
     resolver = SimpleClusterResolver(self._cluster, master='localhost:foo')
+    initial_server_def = context.get_server_def()
     with self.assertRaisesRegex(
         ValueError,
         '`make_master_device_default` is set to True but cannot find master '
         'localhost:foo in the cluster'):
       remote.connect_to_cluster(resolver)
+    self.assertEqual(context.get_server_def(), initial_server_def)
+
+  def testCollectiveGroupLeaderIsFirstNonLocalJob(self):
+    # Regression test for the `break` added to the `collective_group_leader`
+    # search loop: with multiple non-local jobs in the cluster spec, the
+    # leader must be the first one (`my_worker`), not a later one
+    # (`my_ps`) that the loop would otherwise overwrite it with.
+    remote.connect_to_cluster(self._cluster, job_name='localhost')
+    self.assertEqual(
+        context.get_server_def().default_session_config.experimental
+        .collective_group_leader, '/job:my_worker/replica:0/task:0')
 
   def testConnectToClusterWithLocalMaster(self):
     local_resolver = SimpleClusterResolver(ClusterSpec({}), master='local')
@@ -821,8 +839,13 @@ class RemoteUtilsTest(test.TestCase, parameterized.TestCase):
   @parameterized.parameters(
       ('localhost:foo',),
       ('localhost:6553600',),
+      (None,),
+      (123,),
+      ('',),
   )
   def testParseHostAndPortError(self, address):
+    # None, a non-string, and an empty string must be rejected up front
+    # rather than raising TypeError when `"://" in address` is evaluated.
     host, port = remote._parse_host_and_port(address)
     self.assertIsNone(host)
     self.assertIsNone(port)

From 7eefc14b9e46707a214609771dbb07c45d85a9a2 Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Thu, 15 Oct 2026 00:35:04 -0700
Subject: [PATCH 0837/1000] Gather diagonals on GPU only for static XLA
 MatrixDiagPart inputs

XLA:GPU can lay out the gather's output with the batch dimensions
minor, and a batched dynamic output laid out that way came back
transposed, so take the diagonals of dynamic inputs with
GetMatrixDiagonal on GPU as well, and run testDynamicDimensionsBatch
there too.

Test float64 instead of float16 in testDynamicDimensions. The CPU
runtime reads the dynamic sizes of an output from right after its data
as int32_t, and a float16 output of 3 entries puts them at an offset
of 6 bytes.
---
 tensorflow/compiler/tests/matrix_diag_ops_test.py     | 5 +----
 tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc | 9 ++++++---
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/tensorflow/compiler/tests/matrix_diag_ops_test.py b/tensorflow/compiler/tests/matrix_diag_ops_test.py
index 98f21800260..fdfae1cd557 100644
--- a/tensorflow/compiler/tests/matrix_diag_ops_test.py
+++ b/tensorflow/compiler/tests/matrix_diag_ops_test.py
@@ -798,7 +798,7 @@ class MatrixDiagPartTest(xla_test.XLATestCase):
     # entries, so the dimensions it selects have size 2 and bound 5.
     keep = [True, False, True, False, False]
     x = np.arange(1, 16).reshape(5, 3)  # Keeps [[1, 2, 3], [7, 8, 9]].
-    for dtype in (np.int32, np.float32, np.float16):
+    for dtype in (np.int32, np.float32, np.float64):
       for op in (array_ops.matrix_diag_part, gen_array_ops.matrix_diag_part):
         self.assertAllEqual(
             self._diagPartOfSelected(op, x, {-2: keep}, dtype), [1, 8])
@@ -834,9 +834,6 @@ class MatrixDiagPartTest(xla_test.XLATestCase):
                                    {-2: [False] * 5}, dtype), np.zeros([0]))
 
   def testDynamicDimensionsBatch(self):
-    if "GPU" in self.device:
-      self.skipTest("XLA:GPU returns the batched dynamic output of these "
-                    "transposed.")
     keep = [True, False, True, False, False]
     x = np.arange(1, 31).reshape(2, 5, 3)
     self.assertAllEqual(
diff --git a/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc b/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc
index 981a42c6625..6753fe955ba 100644
--- a/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/matrix_diag_ops.cc
@@ -468,8 +468,11 @@ class MatrixDiagPartOp : public XlaOpKernel {
     std::vector<xla::XlaOp> diag_list;
     xla::PaddingConfig padding_config =
         xla::MakeNoPaddingConfig(input_rank - 1);
+    // XLA:GPU can lay out the gather's output with the batch dimensions minor,
+    // and a batched dynamic output laid out that way comes back transposed.
+    const bool use_gather = is_gpu_ && !input_xla_shape.is_dynamic();
     if (num_diags == 1) {
-      output = is_gpu_
+      output = use_gather
                    ? xla::GetMatrixDiagonalViaGather(input, upper_diag_index)
                    : xla::GetMatrixDiagonal(input, upper_diag_index);
     } else {
@@ -479,8 +482,8 @@ class MatrixDiagPartOp : public XlaOpKernel {
       for (int diag_index = upper_diag_index; diag_index >= lower_diag_index;
            --diag_index) {
         xla::XlaOp single_diag =
-            is_gpu_ ? xla::GetMatrixDiagonalViaGather(input, diag_index)
-                    : xla::GetMatrixDiagonal(input, diag_index);
+            use_gather ? xla::GetMatrixDiagonalViaGather(input, diag_index)
+                       : xla::GetMatrixDiagonal(input, diag_index);
         const int64_t diag_len =
             ComputeDiagLen(diag_index, num_rows, num_cols);
         const int64_t padding_len = max_diag_len - diag_len;

From 8a6d5054a951546f7db1d4f7a1db28b953abb85a Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Thu, 15 Oct 2026 13:07:14 +0530
Subject: [PATCH 0838/1000] Use mock.patch.object instead of mock.patch for
 static-analysis friendliness

---
 .../python/platform/windows_lib_diagnostics_test.py  | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/tensorflow/python/platform/windows_lib_diagnostics_test.py b/tensorflow/python/platform/windows_lib_diagnostics_test.py
index b22872c665e..c73a48b7d75 100644
--- a/tensorflow/python/platform/windows_lib_diagnostics_test.py
+++ b/tensorflow/python/platform/windows_lib_diagnostics_test.py
@@ -14,8 +14,10 @@
 # ==============================================================================
 """Tests for windows_lib_diagnostics."""
 
+import builtins
 import io
 import os
+import sys
 from unittest import mock
 
 from tensorflow.python.platform import test
@@ -37,25 +39,25 @@ class WindowsLibDiagnosticsTest(test.TestCase):
 
   def testSafePrintWithAsciiConsole(self):
     out = _FakeStdout("ascii")
-    with mock.patch("sys.stdout", out):
+    with mock.patch.object(sys, "stdout", out):
       windows_lib_diagnostics._safe_print("C:\\Users\\Jos\u00e9\\x.dll")
     self.assertIn("Jos\\xe9", out.getvalue())
 
   def testSafePrintWithNoneEncoding(self):
     out = _FakeStdout(None)
-    with mock.patch("sys.stdout", out):
+    with mock.patch.object(sys, "stdout", out):
       windows_lib_diagnostics._safe_print("caf\u00e9")
     self.assertIn("caf\\xe9", out.getvalue())
 
   def testSafePrintNeverRaises(self):
-    with mock.patch("builtins.print", side_effect=ValueError("boom")):
+    with mock.patch.object(builtins, "print", side_effect=ValueError("boom")):
       windows_lib_diagnostics._safe_print("text")
 
   def testUnsetPathInEnviron(self):
     out = _FakeStdout("utf-8")
     fake_ctypes = mock.MagicMock()
     fake_ctypes.WinDLL.side_effect = OSError("mock load failure")
-    with mock.patch("sys.stdout", out), \
+    with mock.patch.object(sys, "stdout", out), \
         mock.patch.dict(os.environ, {}, clear=True), \
         mock.patch.object(windows_lib_diagnostics, "ctypes", fake_ctypes), \
         mock.patch.object(
@@ -74,7 +76,7 @@ class WindowsLibDiagnosticsTest(test.TestCase):
 
   def testRunDiagnosisSwallowsUnexpectedErrors(self):
     out = _FakeStdout("utf-8")
-    with mock.patch("sys.stdout", out), mock.patch.object(
+    with mock.patch.object(sys, "stdout", out), mock.patch.object(
         windows_lib_diagnostics.os.path, "abspath",
         side_effect=RuntimeError("unexpected")):
       windows_lib_diagnostics.run_diagnosis(None)

From 1a1b41be7b7740fe4995454806b8360b73747953 Mon Sep 17 00:00:00 2001
From: Alejandro Pacheco <alepacheco99@gmail.com>
Date: Wed, 14 Oct 2026 08:56:27 +0200
Subject: [PATCH 0839/1000] Check every batch when initializing a lookup table

InitializableLookupTable::Initialize checked the key and value shapes of
the first batch from the init iterator only, then passed every batch to
DoInsert. The dataset iterator yields one batch per element, so a later
element with more keys than values made HashTable::DoInsert read past
the end of the values tensor.

Check each batch with CheckKeyAndValueTensorsForInsert before inserting
it. Initialize is the only caller of DoInsert, so this covers every
initializable table and iterator.
---
 tensorflow/core/kernels/BUILD                 |  16 +++
 .../kernels/initializable_lookup_table.cc     |   5 +
 .../core/kernels/initializable_lookup_table.h |   3 +-
 .../initializable_lookup_table_test.cc        | 112 ++++++++++++++++++
 4 files changed, 135 insertions(+), 1 deletion(-)
 create mode 100644 tensorflow/core/kernels/initializable_lookup_table_test.cc

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index c3adbc128a3..0c4d567d2df 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -3682,6 +3682,22 @@ tf_cc_test(
     ],
 )
 
+tf_cc_test(
+    name = "initializable_lookup_table_test",
+    size = "small",
+    srcs = ["initializable_lookup_table_test.cc"],
+    deps = [
+        ":initializable_lookup_table",
+        ":lookup_table_op",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+        "@com_google_absl//absl/status",
+    ],
+)
+
 MATH_DEPS = [
     ":fill_functor",
     "//tensorflow/core:core_cpu",
diff --git a/tensorflow/core/kernels/initializable_lookup_table.cc b/tensorflow/core/kernels/initializable_lookup_table.cc
index a546fe8ca78..4df69d0e3d5 100644
--- a/tensorflow/core/kernels/initializable_lookup_table.cc
+++ b/tensorflow/core/kernels/initializable_lookup_table.cc
@@ -89,6 +89,11 @@ absl::Status InitializableLookupTable::Initialize(
   }
   TF_RETURN_IF_ERROR(DoLazyPrepare([&iter]() { return iter.total_size(); }));
   while (iter.Valid()) {
+    // An iterator can yield several batches (e.g. one per dataset element).
+    // DoInsert assumes matching key and value shapes, so check every batch,
+    // not just the first one checked above.
+    TF_RETURN_IF_ERROR(
+        CheckKeyAndValueTensorsForInsert(iter.keys(), iter.values()));
     TF_RETURN_IF_ERROR(DoInsert(iter.keys(), iter.values()));
     iter.Next();
   }
diff --git a/tensorflow/core/kernels/initializable_lookup_table.h b/tensorflow/core/kernels/initializable_lookup_table.h
index 5b3bb1e68f3..d3f73053aa6 100644
--- a/tensorflow/core/kernels/initializable_lookup_table.h
+++ b/tensorflow/core/kernels/initializable_lookup_table.h
@@ -193,7 +193,8 @@ class InitializableLookupTable : public LookupInterface {
   }
 
   // Populates the table in batches given keys and values as tensors into the
-  // underlying data structure.
+  // underlying data structure. Initialize() checks every batch with
+  // CheckKeyAndValueTensorsForInsert() before passing it here.
   virtual absl::Status DoInsert(const Tensor& keys, const Tensor& values) = 0;
 
   // Performs the batch find operation on the underlying data structure.
diff --git a/tensorflow/core/kernels/initializable_lookup_table_test.cc b/tensorflow/core/kernels/initializable_lookup_table_test.cc
new file mode 100644
index 00000000000..58608e4c4b6
--- /dev/null
+++ b/tensorflow/core/kernels/initializable_lookup_table_test.cc
@@ -0,0 +1,112 @@
+/* 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.
+==============================================================================*/
+
+#include "tensorflow/core/kernels/initializable_lookup_table.h"
+
+#include <cstdint>
+#include <utility>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/kernels/lookup_table_op.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/platform/refcount.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace lookup {
+namespace {
+
+// Yields the given (keys, values) pairs one batch at a time, like the dataset
+// iterator used by InitializeTableFromDataset.
+class BatchIterator : public InitializableLookupTable::InitTableIterator {
+ public:
+  explicit BatchIterator(std::vector<std::pair<Tensor, Tensor>> batches)
+      : batches_(std::move(batches)) {}
+
+  void Next() override { ++index_; }
+  bool Valid() const override { return index_ < batches_.size(); }
+  const Tensor& keys() const override { return batches_[index_].first; }
+  const Tensor& values() const override { return batches_[index_].second; }
+  absl::Status status() const override {
+    return Valid() ? absl::OkStatus() : absl::OutOfRangeError("No more data.");
+  }
+  int64_t total_size() const override { return 0; }
+
+ private:
+  std::vector<std::pair<Tensor, Tensor>> batches_;
+  size_t index_ = 0;
+};
+
+core::RefCountPtr<HashTable<int64_t, int64_t>> MakeTable() {
+  return core::RefCountPtr<HashTable<int64_t, int64_t>>(
+      new HashTable<int64_t, int64_t>(nullptr, nullptr));
+}
+
+TEST(InitializableLookupTableTest, InsertsEveryBatch) {
+  auto table = MakeTable();
+  BatchIterator iter({
+      {test::AsTensor<int64_t>({1}), test::AsTensor<int64_t>({10})},
+      {test::AsTensor<int64_t>({2, 3}), test::AsTensor<int64_t>({20, 30})},
+  });
+  TF_ASSERT_OK(table->Initialize(iter));
+  EXPECT_EQ(table->size(), 3u);
+
+  Tensor keys = test::AsTensor<int64_t>({1, 2, 3, 4});
+  Tensor values(DT_INT64, keys.shape());
+  TF_ASSERT_OK(
+      table->Find(nullptr, keys, &values, test::AsScalar<int64_t>(-1)));
+  test::ExpectTensorEqual<int64_t>(values,
+                                   test::AsTensor<int64_t>({10, 20, 30, -1}));
+}
+
+TEST(InitializableLookupTableTest, RejectsFirstBatchWithMismatchedShapes) {
+  auto table = MakeTable();
+  BatchIterator iter({
+      {test::AsTensor<int64_t>({1, 2}), test::AsTensor<int64_t>({10})},
+  });
+  EXPECT_TRUE(absl::IsInvalidArgument(table->Initialize(iter)));
+  EXPECT_FALSE(table->is_initialized());
+}
+
+TEST(InitializableLookupTableTest, RejectsLaterBatchWithMoreKeysThanValues) {
+  auto table = MakeTable();
+  // Only the second batch is malformed: inserting it would read one value per
+  // key, past the end of its single-element values tensor.
+  BatchIterator iter({
+      {test::AsTensor<int64_t>({1}), test::AsTensor<int64_t>({10})},
+      {test::AsTensor<int64_t>({2, 3, 4, 5}), test::AsTensor<int64_t>({20})},
+  });
+  EXPECT_TRUE(absl::IsInvalidArgument(table->Initialize(iter)));
+  EXPECT_FALSE(table->is_initialized());
+}
+
+TEST(InitializableLookupTableTest, RejectsLaterBatchWithExtraValueDims) {
+  auto table = MakeTable();
+  BatchIterator iter({
+      {test::AsTensor<int64_t>({1}), test::AsTensor<int64_t>({10})},
+      {test::AsTensor<int64_t>({2}),
+       test::AsTensor<int64_t>({20, 21}, TensorShape({1, 2}))},
+  });
+  EXPECT_TRUE(absl::IsInvalidArgument(table->Initialize(iter)));
+  EXPECT_FALSE(table->is_initialized());
+}
+
+}  // namespace
+}  // namespace lookup
+}  // namespace tensorflow

From 648a70ad4cf4a6ced5b46d95a8c118196d68cf78 Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Thu, 15 Oct 2026 00:48:44 -0700
Subject: [PATCH 0840/1000] Free the Go SessionOptions config copy when
 TF_SetConfig fails

(*SessionOptions).c() copied Config into C memory and freed the copy
only in the function it returns on success, so each call with a Config
that TF_SetConfig rejects leaked len(Config) bytes. TF_SetConfig only
parses the buffer, so free it right after the call on both paths, as
ContextOptions.c() already does.

Fixes #129133
---
 tensorflow/go/session.go | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/tensorflow/go/session.go b/tensorflow/go/session.go
index 309dc21432a..7ab21930b74 100644
--- a/tensorflow/go/session.go
+++ b/tensorflow/go/session.go
@@ -330,24 +330,21 @@ func (o *SessionOptions) c() (ret *C.TF_SessionOptions, done func(), err error)
 	C.TF_SetTarget(opt, t)
 	C.free(unsafe.Pointer(t))
 
-	var cConfig unsafe.Pointer
 	if sz := len(o.Config); sz > 0 {
 		status := newStatus()
 		// Copying into C-memory is the simplest thing to do in terms
 		// of memory safety and cgo rules ("C code may not keep a copy
 		// of a Go pointer after the call returns" from
 		// https://golang.org/cmd/cgo/#hdr-Passing_pointers).
-		cConfig = C.CBytes(o.Config)
+		cConfig := C.CBytes(o.Config)
 		C.TF_SetConfig(opt, cConfig, C.size_t(sz), status.c)
+		C.free(cConfig)
 		if err := status.Err(); err != nil {
 			C.TF_DeleteSessionOptions(opt)
 			return nil, func() {}, fmt.Errorf("invalid SessionOptions.Config: %v", err)
 		}
 	}
-	return opt, func() {
-		C.TF_DeleteSessionOptions(opt)
-		C.free(cConfig)
-	}, nil
+	return opt, func() { C.TF_DeleteSessionOptions(opt) }, nil
 }
 
 // cRunArgs translates the arguments to Session.Run and PartialRun.Run into

From 971afd296288e1dd539a4675143d174cfbe6283b Mon Sep 17 00:00:00 2001
From: George Necula <necula@google.com>
Date: Wed, 14 Oct 2026 23:48:32 -0700
Subject: [PATCH 0841/1000] [tflite] Set hoist_constants_as_args=False in
 TFLiteJaxConverterV2

With the upcoming jax_use_simplified_jaxpr_constants=True, JAX hoists closed-over array constants larger than 32 bytes into leading const_args parameters of the lowered HLO @main function by default. Because TFLiteJaxConverterV2 exports a standalone HLO proto and binds only explicit ordered_inputs, set hoist_constants_as_args=False in mlir.LoweringParameters.

PiperOrigin-RevId: 999901656
---
 tensorflow/lite/python/lite.py | 8 +++++++-
 tensorflow/lite/python/util.py | 2 ++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/tensorflow/lite/python/lite.py b/tensorflow/lite/python/lite.py
index 65a6416d132..6f0fe369724 100644
--- a/tensorflow/lite/python/lite.py
+++ b/tensorflow/lite/python/lite.py
@@ -65,6 +65,7 @@ from tensorflow.lite.python.op_hint import is_ophint_converted as _is_ophint_con
 from tensorflow.lite.python.op_hint import OpHint  # pylint: disable=unused-import
 from tensorflow.lite.python.optimize import calibrator as _calibrator
 from tensorflow.lite.python.util import _jit
+from tensorflow.lite.python.util import _mlir
 from tensorflow.lite.python.util import build_debug_info_func as _build_debug_info_func
 from tensorflow.lite.python.util import convert_debug_info_func as _convert_debug_info_func
 from tensorflow.lite.python.util import freeze_graph as _freeze_graph
@@ -2113,7 +2114,12 @@ class TFLiteJaxConverterV2(TFLiteConverterBaseV2):
       hlo_proto = (
           _jit(self._serving_funcs[0])
           .trace(*ordered_inputs)
-          .lower(lowering_platforms=("cpu",))
+          .lower(
+              lowering_platforms=("cpu",),
+              _private_parameters=_mlir.LoweringParameters(
+                  hoist_constants_as_args=False
+              ),
+          )
           .compiler_ir("hlo")
           .as_serialized_hlo_module_proto()
       )
diff --git a/tensorflow/lite/python/util.py b/tensorflow/lite/python/util.py
index 5f05881764a..b308710f5bb 100644
--- a/tensorflow/lite/python/util.py
+++ b/tensorflow/lite/python/util.py
@@ -51,8 +51,10 @@ get_save_spec = _tflite_keras_util.get_save_spec
 # pylint: disable=unused-import
 try:
   from jax import jit as _jit
+  from jax.interpreters import mlir as _mlir
 except ImportError:
   _jit = None
+  _mlir = None
 # pylint: enable=g-import-not-at-top
 # pylint: enable=unused-import
 

From fd4c4a7932460ada8d1e7bf1e4c815551df1668d Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2026 23:48:49 -0700
Subject: [PATCH 0842/1000] Add XNNPACK unit test for variable assignment with
 no outputs.

Adds VariableOpsTester::TestAssignOnly and ReadAssignVariable.SimpleAssignOnly to exercise delegating a subgraph containing only VAR_HANDLE and ASSIGN_VARIABLE with zero subgraph outputs. Under UndefinedBehaviorSanitizer with nonnull-attribute checks enabled, this reproduces the null source pointer passed to memcpy in CopyVectorToTfLiteIntArray.

PiperOrigin-RevId: 999901794
---
 tensorflow/lite/delegates/xnnpack/BUILD       |  4 +-
 .../delegates/xnnpack/variable_ops_test.cc    |  6 ++
 .../delegates/xnnpack/variable_ops_tester.cc  | 73 ++++++++++++++++++-
 .../delegates/xnnpack/variable_ops_tester.h   | 11 +++
 4 files changed, 91 insertions(+), 3 deletions(-)

diff --git a/tensorflow/lite/delegates/xnnpack/BUILD b/tensorflow/lite/delegates/xnnpack/BUILD
index 3792804d145..baff3db7ede 100644
--- a/tensorflow/lite/delegates/xnnpack/BUILD
+++ b/tensorflow/lite/delegates/xnnpack/BUILD
@@ -1276,11 +1276,11 @@ cc_library(
     hdrs = ["variable_ops_tester.h"],
     deps = [
         ":xnnpack_delegate_test_mode",
+        "//tensorflow/compiler/mlir/lite/schema:schema_conversion_utils",
         "//tensorflow/lite:framework",
         "//tensorflow/lite:schema_fbs_version",
         "//tensorflow/lite/core/c:common",
         "//tensorflow/lite/core/kernels:builtin_ops",
-        "//tensorflow/lite/schema:schema_conversion_utils",
         "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
         "@flatbuffers",
@@ -1298,11 +1298,11 @@ cc_library(
     hdrs = ["variable_ops_tester.h"],
     deps = [
         ":xnnpack_delegate",
+        "//tensorflow/compiler/mlir/lite/schema:schema_conversion_utils",
         "//tensorflow/lite:framework",
         "//tensorflow/lite:schema_fbs_version",
         "//tensorflow/lite/core/c:common",
         "//tensorflow/lite/core/kernels:builtin_ops",
-        "//tensorflow/lite/schema:schema_conversion_utils",
         "//tensorflow/lite/schema:schema_fbs",
         "@com_google_googletest//:gtest",
         "@flatbuffers",
diff --git a/tensorflow/lite/delegates/xnnpack/variable_ops_test.cc b/tensorflow/lite/delegates/xnnpack/variable_ops_test.cc
index f751dc72b13..b9169967831 100644
--- a/tensorflow/lite/delegates/xnnpack/variable_ops_test.cc
+++ b/tensorflow/lite/delegates/xnnpack/variable_ops_test.cc
@@ -22,6 +22,12 @@ limitations under the License.
 namespace tflite {
 namespace xnnpack {
 
+TEST(ReadAssignVariable, SimpleAssignOnly) {
+  auto xnnpack_delegate = NewXnnPackDelegateSupportingVariableOps();
+  VariableOpsTester().NumInputs(1).NumOutputs(0).TestAssignOnly(
+      xnnpack_delegate.get());
+}
+
 TEST(ReadAssignVariable, SimpleAssignThenRead) {
   auto xnnpack_delegate = NewXnnPackDelegateSupportingVariableOps();
   TfLiteDelegate* delegate = xnnpack_delegate.get();
diff --git a/tensorflow/lite/delegates/xnnpack/variable_ops_tester.cc b/tensorflow/lite/delegates/xnnpack/variable_ops_tester.cc
index 989ce948673..df0969c831e 100644
--- a/tensorflow/lite/delegates/xnnpack/variable_ops_tester.cc
+++ b/tensorflow/lite/delegates/xnnpack/variable_ops_tester.cc
@@ -25,11 +25,13 @@ limitations under the License.
 #include <vector>
 
 #include <gtest/gtest.h>
+#include "flatbuffers/buffer.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "tensorflow/compiler/mlir/lite/schema/schema_conversion_utils.h"
 #include "tensorflow/lite/core/c/common.h"
 #include "tensorflow/lite/core/kernels/register.h"
 #include "tensorflow/lite/delegates/xnnpack/xnnpack_delegate.h"
 #include "tensorflow/lite/interpreter.h"
-#include "tensorflow/lite/schema/schema_conversion_utils.h"
 #include "tensorflow/lite/schema/schema_generated.h"
 #include "tensorflow/lite/version.h"
 
@@ -54,6 +56,75 @@ NewXnnPackDelegateSupportingVariableOps() {
   return xnnpack_delegate;
 }
 
+std::vector<char> VariableOpsTester::CreateModelAssignOnly() const {
+  constexpr int32_t kInputTensorIndex = 0;
+  constexpr int32_t kVarHandleTensorIndex = 1;
+  constexpr size_t kNumSubgraphs = 1;
+  constexpr char kContainerName[] = "container";
+  constexpr char kSharedName[] = "shared_name";
+  constexpr char kModelDescription[] = "AssignVariable model";
+
+  flatbuffers::FlatBufferBuilder builder;
+
+  // Keep the same operator code indices (VAR_HANDLE, READ_VARIABLE,
+  // ASSIGN_VARIABLE) used across all VariableOpsTester models.
+  const std::vector<flatbuffers::Offset<OperatorCode>> operator_codes = {
+      CreateOperatorCode(builder, BuiltinOperator_VAR_HANDLE),
+      CreateOperatorCode(builder, BuiltinOperator_READ_VARIABLE),
+      CreateOperatorCode(builder, BuiltinOperator_ASSIGN_VARIABLE),
+  };
+
+  const std::vector<flatbuffers::Offset<Buffer>> buffers{{
+      CreateBuffer(builder, builder.CreateVector({})),
+  }};
+
+  // Tensor 0 is the graph input; tensor 1 is the VAR_HANDLE resource output.
+  const std::vector<flatbuffers::Offset<Tensor>> tensors{{
+      CreateTensor(
+          builder,
+          builder.CreateVector<int32_t>(Shape().data(), Shape().size()),
+          TensorType_FLOAT32),
+      CreateTensor(builder,
+                   builder.CreateVector<int32_t>(ResourceShape().data(),
+                                                 ResourceShape().size()),
+                   TensorType_RESOURCE),
+  }};
+
+  // Create VAR_HANDLE and ASSIGN_VARIABLE operators with no subgraph outputs.
+  const flatbuffers::Offset<Operator> var_handle_op = CreateOperator(
+      builder, /*opcode_index=*/VAR_HANDLE,
+      /*inputs=*/builder.CreateVector<int32_t>({}),
+      /*outputs=*/builder.CreateVector<int32_t>({kVarHandleTensorIndex}),
+      tflite::BuiltinOptions_VarHandleOptions,
+      CreateVarHandleOptions(builder, builder.CreateString(kContainerName),
+                             builder.CreateString(kSharedName))
+          .Union());
+
+  const flatbuffers::Offset<Operator> assign_op = CreateOperator(
+      builder, /*opcode_index=*/ASSIGN_VARIABLE,
+      /*inputs=*/
+      builder.CreateVector<int32_t>({kVarHandleTensorIndex, kInputTensorIndex}),
+      /*outputs=*/builder.CreateVector<int32_t>({}));
+
+  const flatbuffers::Offset<SubGraph> subgraph = CreateSubGraph(
+      builder, builder.CreateVector(tensors.data(), tensors.size()),
+      /*inputs=*/builder.CreateVector<int32_t>({kInputTensorIndex}),
+      /*outputs=*/builder.CreateVector<int32_t>({}),
+      builder.CreateVector({var_handle_op, assign_op}));
+
+  const flatbuffers::Offset<Model> model_buffer = CreateModel(
+      builder, /*version=*/TFLITE_SCHEMA_VERSION,
+      builder.CreateVector(operator_codes.data(), operator_codes.size()),
+      builder.CreateVector(&subgraph, /*len=*/kNumSubgraphs),
+      builder.CreateString(kModelDescription),
+      builder.CreateVector(buffers.data(), buffers.size()));
+
+  builder.Finish(model_buffer);
+
+  return std::vector<char>(builder.GetBufferPointer(),
+                           builder.GetBufferPointer() + builder.GetSize());
+}
+
 std::vector<char> VariableOpsTester::CreateModelAssignThenRead() const {
   flatbuffers::FlatBufferBuilder builder;
 
diff --git a/tensorflow/lite/delegates/xnnpack/variable_ops_tester.h b/tensorflow/lite/delegates/xnnpack/variable_ops_tester.h
index 7da1b62fe03..4a5383c34e3 100644
--- a/tensorflow/lite/delegates/xnnpack/variable_ops_tester.h
+++ b/tensorflow/lite/delegates/xnnpack/variable_ops_tester.h
@@ -32,6 +32,11 @@ NewXnnPackDelegateSupportingVariableOps();
 
 class VariableOpsTester {
  public:
+  void TestAssignOnly(TfLiteDelegate* delegate) const {
+    const std::vector<char> model = CreateModelAssignOnly();
+    Test(delegate, model);
+  }
+
   void TestAssignThenRead(TfLiteDelegate *delegate) const {
     const std::vector<char> model = CreateModelAssignThenRead();
     Test(delegate, model);
@@ -70,6 +75,12 @@ class VariableOpsTester {
     Test(delegate, model);
   }
 
+  // Creates a model with this subgraph:
+  //   (initial_input) ----\
+  //                        \
+  //   VAR_HANDLE ----------> AV
+  std::vector<char> CreateModelAssignOnly() const;
+
   // Creates a model with this subgraph:
   //   (initial_input) ----\
   //                        \

From c5c65157fa2c8eed901569d0df8cd0ab52f74e71 Mon Sep 17 00:00:00 2001
From: Ibrahim Umit Akgun <umitakgun@google.com>
Date: Wed, 14 Oct 2026 23:49:37 -0700
Subject: [PATCH 0843/1000] [XLA:GPU] Avoid grouping GPU reductions that only
 share broadcasted inputs.

When multiple reductions in a multi-output `kInput` fusion share a lower-rank input through a `kBroadcast` instruction (such as the shared bit-shift constant in `jnp.packbits`), `GroupDisjointReductions` previously merged all of those reductions into a single reduction group. Placing up to 8 independent reductions into the same inner per-thread loop increases loop body size beyond LLVM's unroll threshold, preventing full unrolling of the inner vector loop and causing dynamic vector indexing to spill into `.local` stack memory with high register pressure and low occupancy.

This change:
1. Adds `IsBroadcastCheapToRecompute`, which checks whether a `kBroadcast` inside the fusion only depends on cheap-to-recompute operations (`kParameter`, `kConstant`, `kIota`, `kConvert`, `kBitcast`, `kReshape`, and `kBroadcast`).
2. Runs a separate second `HloBfsConsumersFirstTraversal` per root to compute `unbroadcasted_reachable_outputs` (returning `TraversalResult::kSkip` when visiting a cheap-to-recompute `kBroadcast`), so that BFS shortest-hop ordering cannot shadow longer unbroadcasted paths in diamond graphs.
3. Requires `unbroadcasted_reachable_outputs[i].contains(output)` before merging more than one reduction root (`has_unbroadcasted_reduce || added_to_reduce`), placing reductions that only share cheap-to-recompute broadcasted inputs into separate reduction groups executed in parallel across `blockIdx.z` within the same kernel launch while still allowing non-reduction side-output roots to group with a single reduction root.

PiperOrigin-RevId: 999902124
---
 .../xla/backends/gpu/codegen/emitters/BUILD   |   3 -
 .../gpu/codegen/emitters/reduction_base.cc    |  94 ++++++++-
 .../codegen/emitters/reduction_base_test.cc   | 192 ++++++++++++++++++
 3 files changed, 279 insertions(+), 10 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/BUILD b/third_party/xla/xla/backends/gpu/codegen/emitters/BUILD
index 5bddbf11fbd..079c82495cd 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/BUILD
@@ -344,11 +344,8 @@ cc_library(
         "//xla/hlo/ir:hlo",
         "//xla/hlo/utils:hlo_query",
         "//xla/hlo/utils:hlo_traversal",
-        "//xla/service/gpu:gpu_fusible",
         "//xla/service/gpu:hlo_fusion_analysis",
-        "//xla/service/gpu:ir_emission_utils",
         "//xla/service/gpu:reduction_utils",
-        "//xla/stream_executor:device_description",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/reduction_base.cc b/third_party/xla/xla/backends/gpu/codegen/emitters/reduction_base.cc
index 3ba36f95c55..1c3139fdfec 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/reduction_base.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/reduction_base.cc
@@ -33,21 +33,53 @@ limitations under the License.
 #include "xla/hlo/analysis/indexing_map.h"
 #include "xla/hlo/analysis/symbolic_map.h"
 #include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/utils/hlo_query.h"
 #include "xla/hlo/utils/hlo_traversal.h"
 #include "xla/primitive_util.h"
-#include "xla/service/gpu/gpu_fusible.h"
 #include "xla/service/gpu/hlo_fusion_analysis.h"
-#include "xla/service/gpu/ir_emission_utils.h"
 #include "xla/service/gpu/reduction_utils.h"
 #include "xla/shape.h"
-#include "xla/stream_executor/device_description.h"
 #include "xla/union_find.h"
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
 
 namespace xla {
 namespace gpu {
+namespace {
+
+// Returns true if `instr` is a broadcast inside the fusion whose operand chain
+// within the fusion consists only of cheap-to-recompute instructions
+// (`parameter`, `constant`, `iota`, `convert`, `bitcast`, `reshape`, and
+// `broadcast`).
+bool IsBroadcastCheapToRecompute(HloInstructionAdaptor instr) {
+  const HloFusionAdaptor& fusion = instr.parent();
+  if (!fusion.ContainsInstruction(instr) ||
+      instr.opcode() != HloOpcode::kBroadcast) {
+    return false;
+  }
+  HloInstructionAdaptor curr = instr.GetOperand(0);
+  while (fusion.ContainsInstruction(curr)) {
+    switch (curr.opcode()) {
+      case HloOpcode::kParameter:
+      case HloOpcode::kConstant:
+      case HloOpcode::kIota:
+        return true;
+      case HloOpcode::kConvert:
+      case HloOpcode::kBitcast:
+      case HloOpcode::kReshape:
+      case HloOpcode::kBroadcast:
+        curr = curr.GetOperand(0);
+        break;
+      default:
+        return false;
+    }
+  }
+  // `curr` is outside the fusion, which means it is a fusion parameter.
+  return true;
+}
+
+}  // namespace
 
 int GetVectorSizeForMlir(const HloFusionAnalysis& analysis, int64_t minor_dim,
                          int num_threads) {
@@ -103,6 +135,22 @@ ReductionGroups GroupDisjointReductions(const HloFusionAnalysis& analysis) {
   absl::node_hash_map<HloInstructionAdaptor,
                       absl::flat_hash_set<HloInstructionAdaptor>>
       reachable_outputs;
+  // Tracks fusion roots that can be reached from an instruction without
+  // traversing any cheap-to-recompute `kBroadcast` instruction. In multi-output
+  // `kInput` reduction fusions, all roots in the same reduction group are
+  // emitted into a single per-thread inner reduction loop, whereas separate
+  // reduction groups execute in parallel across `blockIdx.z` within the same
+  // kernel. Merging reductions that only share a cheap broadcasted input (such
+  // as a small lookup table or shift constant) bloats the inner reduction loop
+  // body and can exceed LLVM's loop unroll threshold, preventing full
+  // unrolling of the inner vector loop and causing dynamic vector indexing to
+  // spill into `.local` stack memory with high register pressure and low
+  // occupancy. Keeping `unbroadcasted_reachable_outputs` allows us to only
+  // merge reduction roots that share an unbroadcasted or non-trivial
+  // predecessor.
+  absl::node_hash_map<HloInstructionAdaptor,
+                      absl::flat_hash_set<HloInstructionAdaptor>>
+      unbroadcasted_reachable_outputs;
   absl::flat_hash_set<HloInstructionAdaptor> roots_with_reduction;
   absl::flat_hash_map<const HloInstruction*, int> root_indices;
   const auto& roots = analysis.fusion().GetRoots();
@@ -147,16 +195,46 @@ ReductionGroups GroupDisjointReductions(const HloFusionAnalysis& analysis) {
 
   for (auto root : roots) {
     visit({root});
+    HloBfsConsumersFirstTraversal(
+        {root}, analysis.fusion(), [&](HloInstructionAdaptor consumer) {
+          if (IsBroadcastCheapToRecompute(consumer)) {
+            return TraversalResult::kSkip;
+          }
+          unbroadcasted_reachable_outputs[consumer].insert(root);
+          for (auto producer : consumer.GetOperands()) {
+            if (!analysis.fusion().ContainsInstruction(producer)) {
+              unbroadcasted_reachable_outputs[producer].insert(root);
+            }
+          }
+          return TraversalResult::kAdvance;
+        });
   }
 
   for (auto instr : instructions) {
     const auto& reachable = reachable_outputs[instr];
+    const auto& unbroadcasted_reachable =
+        unbroadcasted_reachable_outputs[instr];
+    bool has_unbroadcasted_reduce =
+        absl::c_any_of(roots_with_reduction, [&](HloInstructionAdaptor root) {
+          return unbroadcasted_reachable.contains(root);
+        });
     std::vector<HloInstructionAdaptor> reached_output_ids;
     bool added_to_reduce = false;
     for (auto output : roots) {
       bool has_real_hero = roots_with_reduction.contains(output);
+      if (has_real_hero && !unbroadcasted_reachable.contains(output)) {
+        if (has_unbroadcasted_reduce || added_to_reduce) {
+          // Do not group more than one output reduce instructions through an
+          // instruction that only reaches reductions via a cheap-to-recompute
+          // broadcast.
+          VLOG(3) << "Skip instruction " << instr.ToString()
+                  << " reaching reduction output " << output.ToString()
+                  << " only through a cheap broadcast";
+          continue;
+        }
+      }
       if (has_real_hero &&
-          (hlo_query::IsBroadcastedConstantOrScalar(instr.instruction()))) {
+          hlo_query::IsBroadcastedConstantOrScalar(instr.instruction())) {
         if (added_to_reduce) {
           // Do not group more than one output reduce instructions through
           // broadcasted constants or scalars, as the recomputation should be
@@ -175,9 +253,11 @@ ReductionGroups GroupDisjointReductions(const HloFusionAnalysis& analysis) {
         }
       }
     }
-    auto& first_reached_output = disjoint_sets.at(reached_output_ids.front());
-    for (size_t j = 1; j < reached_output_ids.size(); ++j) {
-      first_reached_output.Merge(&disjoint_sets.at(reached_output_ids[j]));
+    if (!reached_output_ids.empty()) {
+      auto& first_reached_output = disjoint_sets.at(reached_output_ids.front());
+      for (size_t j = 1; j < reached_output_ids.size(); ++j) {
+        first_reached_output.Merge(&disjoint_sets.at(reached_output_ids[j]));
+      }
     }
   }
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/reduction_base_test.cc b/third_party/xla/xla/backends/gpu/codegen/emitters/reduction_base_test.cc
index bed62b2d2e3..6311c9a508c 100644
--- a/third_party/xla/xla/backends/gpu/codegen/emitters/reduction_base_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/emitters/reduction_base_test.cc
@@ -90,6 +90,198 @@ TEST_F(MlirReductionBaseTest, OneGroup) {
   EXPECT_THAT(reduction_groups.grouped_roots, SizeIs(1));
 }
 
+TEST_F(MlirReductionBaseTest, TwoGroupsSharedBroadcast) {
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
+    add {
+      p0 = f32[] parameter(0)
+      p1 = f32[] parameter(1)
+      ROOT add = f32[] add(p0, p1)
+    }
+    fusion {
+      %p0 = s16[4,8] parameter(0)
+      %conv = f32[4,8] convert(%p0)
+      %reshape = f32[32] reshape(%conv)
+      %bitcast = f32[32] bitcast(%reshape)
+      %b0 = f32[32,64] broadcast(%bitcast), dimensions={0}
+      %iota = f32[64] iota(), iota_dimension=0
+      %b1 = f32[32,64] broadcast(%iota), dimensions={1}
+      %p1 = f32[32,64] parameter(1)
+      %p2 = f32[32,64] parameter(2)
+      %a0 = f32[32,64] add(%p1, %b0)
+      %a0_iota = f32[32,64] add(%a0, %b1)
+      %a1 = f32[32,64] add(%p2, %b0)
+      %a1_iota = f32[32,64] add(%a1, %b1)
+      %c0 = f32[] constant(0)
+      %r0 = f32[32] reduce(%a0_iota, %c0), dimensions={1}, to_apply=add
+      %r1 = f32[32] reduce(%a1_iota, %c0), dimensions={1}, to_apply=add
+      ROOT %tuple = (f32[32], f32[32]) tuple(%r0, %r1)
+    }
+    ENTRY entry {
+      %p0 = s16[4,8] parameter(0)
+      %p1 = f32[32,64] parameter(1)
+      %p2 = f32[32,64] parameter(2)
+      ROOT %fusion = (f32[32], f32[32]) fusion(%p0, %p1, %p2), kind=kInput, calls=fusion
+    })"));
+
+  auto device_info = TestGpuDeviceInfo::CudaOrRocmDeviceInfo();
+  auto* root = module->entry_computation()->root_instruction();
+  auto analysis = HloFusionAnalysis::Create(*root, device_info);
+  auto reduction_groups = GroupDisjointReductions(analysis);
+  EXPECT_THAT(reduction_groups.grouped_roots,
+              ElementsAre(ElementsAre(&analysis.fusion_root(0).instruction()),
+                          ElementsAre(&analysis.fusion_root(1).instruction())));
+}
+
+TEST_F(MlirReductionBaseTest, OneGroupSharedExpensiveBroadcast) {
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
+    add {
+      p0 = f32[] parameter(0)
+      p1 = f32[] parameter(1)
+      ROOT add = f32[] add(p0, p1)
+    }
+    fusion {
+      %p0 = f32[32] parameter(0)
+      %p1 = f32[32] parameter(1)
+      %sum = f32[32] add(%p0, %p1)
+      %b0 = f32[32,64] broadcast(%sum), dimensions={0}
+      %p2 = f32[32,64] parameter(2)
+      %p3 = f32[32,64] parameter(3)
+      %a0 = f32[32,64] add(%p2, %b0)
+      %a1 = f32[32,64] add(%p3, %b0)
+      %c0 = f32[] constant(0)
+      %r0 = f32[32] reduce(%a0, %c0), dimensions={1}, to_apply=add
+      %r1 = f32[32] reduce(%a1, %c0), dimensions={1}, to_apply=add
+      ROOT %tuple = (f32[32], f32[32]) tuple(%r0, %r1)
+    }
+    ENTRY entry {
+      %p0 = f32[32] parameter(0)
+      %p1 = f32[32] parameter(1)
+      %p2 = f32[32,64] parameter(2)
+      %p3 = f32[32,64] parameter(3)
+      ROOT %fusion = (f32[32], f32[32]) fusion(%p0, %p1, %p2, %p3), kind=kInput, calls=fusion
+    })"));
+
+  auto device_info = TestGpuDeviceInfo::CudaOrRocmDeviceInfo();
+  auto* root = module->entry_computation()->root_instruction();
+  auto analysis = HloFusionAnalysis::Create(*root, device_info);
+  auto reduction_groups = GroupDisjointReductions(analysis);
+  EXPECT_THAT(reduction_groups.grouped_roots, SizeIs(1));
+}
+
+TEST_F(MlirReductionBaseTest, TwoGroupsBroadcastToFirstRootDirectToSecondRoot) {
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
+    add {
+      p0 = f32[] parameter(0)
+      p1 = f32[] parameter(1)
+      ROOT add = f32[] add(p0, p1)
+    }
+    fusion {
+      %p0 = f32[32] parameter(0)
+      %b0 = f32[32,64] broadcast(%p0), dimensions={0}
+      %p1 = f32[32,64] parameter(1)
+      %a0 = f32[32,64] add(%p1, %b0)
+      %c0 = f32[] constant(0)
+      %r0 = f32[32] reduce(%a0, %c0), dimensions={1}, to_apply=add
+      %r1 = f32[] reduce(%p0, %c0), dimensions={0}, to_apply=add
+      ROOT %tuple = (f32[32], f32[]) tuple(%r0, %r1)
+    }
+    ENTRY entry {
+      %p0 = f32[32] parameter(0)
+      %p1 = f32[32,64] parameter(1)
+      ROOT %fusion = (f32[32], f32[]) fusion(%p0, %p1), kind=kInput, calls=fusion
+    })"));
+
+  auto device_info = TestGpuDeviceInfo::CudaOrRocmDeviceInfo();
+  auto* root = module->entry_computation()->root_instruction();
+  auto analysis = HloFusionAnalysis::Create(*root, device_info);
+  auto reduction_groups = GroupDisjointReductions(analysis);
+  EXPECT_THAT(reduction_groups.grouped_roots,
+              ElementsAre(ElementsAre(&analysis.fusion_root(0).instruction()),
+                          ElementsAre(&analysis.fusion_root(1).instruction())));
+}
+
+TEST_F(MlirReductionBaseTest, OneGroupDiamondWithShortBroadcastPath) {
+  // %conv (and %p0) reaches %r0 via BOTH a shorter cheap-broadcast path
+  // (%r0 -> %r0_bcast -> %b0 -> %conv, 3 hops) and a longer unbroadcasted path
+  // (%r0 -> %r0_epilogue -> %neg2 -> %neg1 -> %conv, 4 hops), while also
+  // reaching %r1 directly (%r1 -> %conv). Even though BFS visits %b0 before
+  // %neg1, skipping %b0 ensures the unbroadcasted path still reaches %conv and
+  // %p0, grouping %r0 and %r1 into a single reduction group.
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
+    add {
+      p0 = f32[] parameter(0)
+      p1 = f32[] parameter(1)
+      ROOT add = f32[] add(p0, p1)
+    }
+    fusion {
+      %p0 = s16[32] parameter(0)
+      %conv = f32[32] convert(%p0)
+      %b0 = f32[32,64] broadcast(%conv), dimensions={0}
+      %p1 = f32[32,64] parameter(1)
+      %p2 = f32[32,64] parameter(2)
+      %c0 = f32[] constant(0)
+      %r0_base = f32[32] reduce(%p1, %c0), dimensions={1}, to_apply=add
+      %neg1 = f32[32] negate(%conv)
+      %neg2 = f32[32] negate(%neg1)
+      %r0_epilogue = f32[32] add(%r0_base, %neg2)
+      %r0_bcast = f32[32] reduce(%b0, %c0), dimensions={1}, to_apply=add
+      %r0 = f32[32] add(%r0_epilogue, %r0_bcast)
+      %r1_base = f32[32] reduce(%p2, %c0), dimensions={1}, to_apply=add
+      %r1 = f32[32] add(%r1_base, %conv)
+      ROOT %tuple = (f32[32], f32[32]) tuple(%r0, %r1)
+    }
+    ENTRY entry {
+      %p0 = s16[32] parameter(0)
+      %p1 = f32[32,64] parameter(1)
+      %p2 = f32[32,64] parameter(2)
+      ROOT %fusion = (f32[32], f32[32]) fusion(%p0, %p1, %p2), kind=kInput, calls=fusion
+    })"));
+
+  auto device_info = TestGpuDeviceInfo::CudaOrRocmDeviceInfo();
+  auto* root = module->entry_computation()->root_instruction();
+  auto analysis = HloFusionAnalysis::Create(*root, device_info);
+  auto reduction_groups = GroupDisjointReductions(analysis);
+  EXPECT_THAT(reduction_groups.grouped_roots, SizeIs(1));
+}
+
+TEST_F(MlirReductionBaseTest, TwoGroupsWithBroadcastSideOutput) {
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
+    add {
+      p0 = f32[] parameter(0)
+      p1 = f32[] parameter(1)
+      ROOT add = f32[] add(p0, p1)
+    }
+    fusion {
+      %p0 = s16[32] parameter(0)
+      %conv = f32[32] convert(%p0)
+      %b0 = f32[32,64] broadcast(%conv), dimensions={0}
+      %p1 = f32[32,64] parameter(1)
+      %p2 = f32[32,64] parameter(2)
+      %a0 = f32[32,64] add(%p1, %b0)
+      %a1 = f32[32,64] add(%p2, %b0)
+      %c0 = f32[] constant(0)
+      %r0 = f32[32] reduce(%a0, %c0), dimensions={1}, to_apply=add
+      %r1 = f32[32] reduce(%a1, %c0), dimensions={1}, to_apply=add
+      ROOT %tuple = (f32[32], f32[32], f32[32,64]) tuple(%r0, %r1, %b0)
+    }
+    ENTRY entry {
+      %p0 = s16[32] parameter(0)
+      %p1 = f32[32,64] parameter(1)
+      %p2 = f32[32,64] parameter(2)
+      ROOT %fusion = (f32[32], f32[32], f32[32,64])
+          fusion(%p0, %p1, %p2), kind=kInput, calls=fusion
+    })"));
+
+  auto device_info = TestGpuDeviceInfo::CudaOrRocmDeviceInfo();
+  auto* root = module->entry_computation()->root_instruction();
+  auto analysis = HloFusionAnalysis::Create(*root, device_info);
+  auto reduction_groups = GroupDisjointReductions(analysis);
+  EXPECT_THAT(reduction_groups.grouped_roots,
+              ElementsAre(ElementsAre(&analysis.fusion_root(0).instruction(),
+                                      &analysis.fusion_root(2).instruction()),
+                          ElementsAre(&analysis.fusion_root(1).instruction())));
+}
+
 }  // namespace
 }  // namespace gpu
 }  // namespace xla

From 9bd2dd02afc6fc214b12a108f6239c94e66741b1 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Wed, 14 Oct 2026 23:49:42 -0700
Subject: [PATCH 0844/1000] Move zizmor dangerous-triggers ignore comment onto
 pull_request_target line.

Upstream `zizmor` v1.31.0 narrowed the `dangerous-triggers` diagnostic span from the outer `on:` block to the `pull_request_target:` key (`zizmorcore/zizmor#2402`), causing the preceding `# zizmor: ignore[dangerous-triggers]` comment on line 19 of `.github/workflows/pr_size_check.yml` to fall outside the finding span and fail `GitHub Actions Security Analysis` across `main` and open pull requests.

Move `# zizmor: ignore[dangerous-triggers]` onto the `pull_request_target:` line in `.github/workflows/pr_size_check.yml` so the suppression lies inside the `pull_request_target:` span.

PiperOrigin-RevId: 999902149
---
 third_party/xla/.github/workflows/pr_size_check.yml | 3 +--
 third_party/xla/.github/workflows/zizmor.yml        | 3 ++-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/.github/workflows/pr_size_check.yml b/third_party/xla/.github/workflows/pr_size_check.yml
index 6b98851c6be..cd60a316f74 100644
--- a/third_party/xla/.github/workflows/pr_size_check.yml
+++ b/third_party/xla/.github/workflows/pr_size_check.yml
@@ -16,8 +16,7 @@ name: PR Size Check
 permissions:
   pull-requests: write
 on:
-  # zizmor: ignore[dangerous-triggers]
-  pull_request_target:
+  pull_request_target: # zizmor: ignore[dangerous-triggers]
     branches:
       - main
 
diff --git a/third_party/xla/.github/workflows/zizmor.yml b/third_party/xla/.github/workflows/zizmor.yml
index 76e48445905..5767aec9996 100644
--- a/third_party/xla/.github/workflows/zizmor.yml
+++ b/third_party/xla/.github/workflows/zizmor.yml
@@ -36,6 +36,7 @@ jobs:
         uses: astral-sh/setup-uv@c18668ad3cf93ea998bef934396af7bb5c839dc7 # v10.2.0
 
       - name: Run zizmor
-        run:  uvx zizmor --format=github .github xla
+        # Pinned version; bumped on a regular schedule by the XLA zizmor release monitor.
+        run:  uvx zizmor@1.31.0 --format=github .github xla
         env:
           GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}

From 8adf67a9ecb01338ea6f5f301b792f4c0770fb1e Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 15 Oct 2026 00:10:16 -0700
Subject: [PATCH 0845/1000] Automated Code Change

PiperOrigin-RevId: 999913039
---
 tensorflow/core/kernels/conv_grad_filter_ops_launcher.cc | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/tensorflow/core/kernels/conv_grad_filter_ops_launcher.cc b/tensorflow/core/kernels/conv_grad_filter_ops_launcher.cc
index 4f0e24e6fe2..739a168b9ed 100644
--- a/tensorflow/core/kernels/conv_grad_filter_ops_launcher.cc
+++ b/tensorflow/core/kernels/conv_grad_filter_ops_launcher.cc
@@ -403,10 +403,11 @@ void LaunchConv2DBackpropFilterOpImpl(
   // Validate filter element count before allocation to prevent OOM on invalid
   // inputs. GPU transformation uses 32-bit indexing via To32Bit().
   OP_REQUIRES(
-      ctx, filter_backprop->NumElements() <= std::numeric_limits<int32>::max(),
-      errors::InvalidArgument("Filter tensor num elements (",
-                              filter_backprop->NumElements(),
-                              ") exceeds 32-bit limit for GPU transformation"));
+      ctx,
+      filter_backprop->NumElements() <= std::numeric_limits<int32_t>::max(),
+      absl::InvalidArgumentError(absl::StrCat(
+          "Filter tensor num elements (", filter_backprop->NumElements(),
+          ") exceeds 32-bit limit for GPU transformation")));
 
   Tensor pre_transformed_filter_backprop;
   OP_REQUIRES_OK(

From 4cda1bb6103d48107da783f79c2880edc57df5f3 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 15 Oct 2026 00:43:41 -0700
Subject: [PATCH 0846/1000] Reverts ba721bca5b7ddbc26e5e4b5c25143ee3659e1774

PiperOrigin-RevId: 999929341
---
 .../scaled_dot_online_quantization.hlo        |   4 +-
 ...scaled_dot_online_quantization_bitcast.hlo |   4 +-
 .../triton/tests/scaled_dot_device_test.cc    |   4 +-
 .../backends/gpu/transforms/gemm_fusion.cc    | 160 ++++++++++--------
 4 files changed, 94 insertions(+), 78 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization.hlo
index f0f2196d3ed..5619eee8d23 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization.hlo
@@ -46,12 +46,12 @@ ENTRY %main.4 (a.1: f8e4m3fn[2,256], b.1: f8e4m3fn[256,256], a_scale.1: f8e8m0fn
   ROOT %tuple.3 = (f8e3m4[2,256]{1,0}, f8e8m0fnu[2,8]{1,0}) tuple(%jit_online_quantize_.4, %jit_online_quantize_.5)
 }
 
-// CHECK: %[[FUSION_DOT:[^ ]*scaled-dot[^ ]*]] (param
+// CHECK: %fusion_scaled-dot (parameter
 // CHECK: %fused_reduce (param
 // CHECK: %fused_convert (param
 // CHECK: %fused_convert.1 (param
 // CHECK: ENTRY
-// CHECK: %[[fusion_dot:.*]] = f32[2,256]{1,0} fusion({{.*}}), kind=kCustom, calls=%[[FUSION_DOT]]
+// CHECK: %[[fusion_dot:.*]] = f32[2,256]{1,0} fusion({{.*}}), kind=kCustom, calls=%fusion_scaled-dot
 // CHECK: %[[fusion_reduce:.*]] = f32[2,8]{1,0} fusion(%[[fusion_dot]]), kind=kInput, calls=%fused_reduce
 // CHECK: %[[fusion_convert1:.*]] = f8e3m4[2,256]{1,0} fusion(%[[fusion_dot]], %[[fusion_reduce]]), kind=kLoop, calls=%fused_convert
 // CHECK: %[[fusion_convert2:.*]] = f8e8m0fnu[2,8]{1,0} fusion(%[[fusion_reduce]]), kind=kLoop, calls=%fused_convert.1
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization_bitcast.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization_bitcast.hlo
index ff21e7c54c4..20fd0b6c248 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization_bitcast.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization_bitcast.hlo
@@ -59,12 +59,12 @@ ENTRY %main.5 (a.1: f8e4m3fn[2,256], b.1: f8e4m3fn[256,256], a_scale.1: f8e8m0fn
   ROOT %tuple.0 = (f8e3m4[2,256]{1,0}, f8e8m0fnu[2,8]{1,0}) tuple(%convert_element_type.1, %convert_element_type.2)
 }
 
-// CHECK: %[[FUSION_DOT:[^ ]*scaled-dot[^ ]*]] (param
+// CHECK: %fusion_scaled-dot (parameter
 // CHECK: %fused_reduce (param
 // CHECK: %fused_convert (param
 // CHECK: %fused_convert.1 (param
 // CHECK: ENTRY
-// CHECK: %[[fusion_dot:.*]] = f32[2,256]{1,0} fusion({{.*}}), kind=kCustom, calls=%[[FUSION_DOT]]
+// CHECK: %[[fusion_dot:.*]] = f32[2,256]{1,0} fusion({{.*}}), kind=kCustom, calls=%fusion_scaled-dot
 // CHECK: %[[fusion_reduce:.*]] = f32[2,8]{1,0} fusion(%[[fusion_dot]]), kind=kInput, calls=%fused_reduce
 // CHECK: %[[fusion_convert1:.*]] = f8e3m4[2,256]{1,0} fusion(%[[fusion_dot]], %[[fusion_reduce]]), kind=kLoop, calls=%fused_convert
 // CHECK: %[[fusion_convert2:.*]] = f8e8m0fnu[2,8]{1,0} fusion(%[[fusion_reduce]]), kind=kLoop, calls=%fused_convert.1
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc
index ee848028d19..abf8ca66f8f 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc
@@ -706,7 +706,7 @@ ENTRY e {
   ASSERT_OK_AND_ASSIGN(auto optimized_module,
                        GetOptimizedModule(kHloTextTemplate));
   constexpr absl::string_view kExpectedOptimizedHLO = R"(
-    CHECK: %{{.*}} (param
+    CHECK: %fusion
     CHECK: %{{.*}} = f8e8m0fnu[3,128,4]{2,1,0} broadcast(%{{.*}}), dimensions={0,1}
     CHECK: %{{.*}} = f8e8m0fnu[3,128,4]{2,1,0} broadcast(%{{.*}}), dimensions={0,1}
     CHECK: ROOT {{.*}} scaled-dot
@@ -894,7 +894,7 @@ ENTRY e {
   ASSERT_OK_AND_ASSIGN(auto optimized_module,
                        GetOptimizedModule(kHloTextTemplate));
   constexpr absl::string_view kExpectedOptimizedHLO = R"(
-    CHECK: %[[fusion_name:.*]] (param
+    CHECK: %[[fusion_name:.*]] (parameter
     CHECK: %[[scaled_dot:.*]] = bf16[3,128,128]{2,1,0} scaled-dot
     CHECK: %[[global_scaler:.*]] = bf16[3,128,128]{2,1,0} broadcast
     CHECK: ROOT %{{.*}} = bf16[3,128,128]{2,1,0} multiply(%[[scaled_dot]], %[[global_scaler]])
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 5ed108f5149..d28df71b899 100644
--- a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc
@@ -718,9 +718,8 @@ bool AllowedInGemmFusion(const HloInstruction& instr) {
     // the moment. Includes kFusion, kReduce, kAllReduce, etc.
     return false;
   }
-  return HloPredicateIsNotOp<HloOpcode::kDot, HloOpcode::kScaledDot,
-                             HloOpcode::kParameter, HloOpcode::kAllGather,
-                             HloOpcode::kAllGatherStart,
+  return HloPredicateIsNotOp<HloOpcode::kDot, HloOpcode::kParameter,
+                             HloOpcode::kAllGather, HloOpcode::kAllGatherStart,
                              HloOpcode::kAllGatherDone>(&instr);
 }
 
@@ -1266,8 +1265,7 @@ FusionDecision CanUnpackS4ParametersInFusion(
   }
 
   if (tiled_computation->roots().size() != 1 ||
-      !HloPredicateIsOp<HloOpcode::kDot, HloOpcode::kScaledDot>(
-          tiled_computation->roots()[0]->hlo())) {
+      tiled_computation->roots()[0]->hlo()->opcode() != HloOpcode::kDot) {
     return FusionDecision::Forbid(
         "Expected tiled computation root to be a single dot instruction.");
   }
@@ -1317,8 +1315,7 @@ FusionDecision CanTile(mlir::MLIRContext& mlir_context,
     // a tile size to be divisible by 3, but all Triton tiles are powers of 2.
     std::optional<HloInstructionAdaptor> dot =
         HloBfsFindIf(fusion.GetRoots(), fusion, [](HloInstructionAdaptor node) {
-          return HloPredicateIsOp<HloOpcode::kDot, HloOpcode::kScaledDot>(
-              &node.instruction());
+          return node.opcode() == HloOpcode::kDot;
         });
     if (!dot.has_value()) {
       return FusionDecision::Forbid("No dot instruction found in the fusion.");
@@ -1671,26 +1668,22 @@ FusionDecision ShouldFuseUserTranspose(const HloInstruction& transpose,
                                        const HloInstruction& original_transpose,
                                        const HloInstruction& fusion,
                                        const ShapeTracker& tracker) {
-  const HloInstruction* dot = hlo_query::GetFirstInstructionWithOpcode(
-      *fusion.fused_instructions_computation(),
-      {HloOpcode::kDot, HloOpcode::kScaledDot});
+  const HloInstruction* dot = hlo_query::FindInstruction(
+      fusion.fused_instructions_computation(), HloOpcode::kDot);
   if (dot == nullptr) {
     return FusionDecision::Forbid("Dot not found in fusion.");
   }
 
-  // The output of a scaled dot has the same layout as the output of a dot of
-  // its LHS and RHS, so the dims of operands 0 and 1 are sufficient.
-  absl::StatusOr<DotOperandDims> lhs_dims =
-      DotOperandDims::FromDotOperand(dot, 0);
-  absl::StatusOr<DotOperandDims> rhs_dims =
-      DotOperandDims::FromDotOperand(dot, 1);
-  if (!lhs_dims.ok() || !rhs_dims.ok()) {
+  absl::StatusOr<std::array<DotOperandDims, 2>> dot_dims =
+      DotOperandDims::FromDot(dot);
+  if (!dot_dims.ok()) {
     return FusionDecision::Forbid("Failed to get dot operand dims.");
   }
+  const auto& [lhs_dims, rhs_dims] = *dot_dims;
 
-  int64_t num_batch = lhs_dims->Rank(DotOperandDims::kBatch);
-  int64_t num_lhs_nc = lhs_dims->Rank(DotOperandDims::kNonContracting);
-  int64_t num_rhs_nc = rhs_dims->Rank(DotOperandDims::kNonContracting);
+  int64_t num_batch = lhs_dims.Rank(DotOperandDims::kBatch);
+  int64_t num_lhs_nc = lhs_dims.Rank(DotOperandDims::kNonContracting);
+  int64_t num_rhs_nc = rhs_dims.Rank(DotOperandDims::kNonContracting);
 
   // Dot result always maps to [batch..., lhs_nc..., rhs_nc...].
   absl::InlinedVector<int64_t, 4> batch_dims(num_batch);
@@ -1917,22 +1910,18 @@ absl::StatusOr<std::variant<Fusion, FusionDecision>> CreateTileableFusion(
     return FusionDecision::Forbid("Cannot tile the dot instruction.");
   }
 
-  // Initialize queue with trackers for the LHS and RHS. Scaled dot scales get
-  // no tracker, so transposes and concatenates are not fused into them.
-  std::queue<std::pair<HloInstruction*, std::optional<TrackerInfo>>> queue;
-  for (int64_t i = 0; i < dot->operand_count(); ++i) {
-    HloInstruction* operand = dot->mutable_operand(i);
-    std::optional<TrackerInfo> tracker;
-    if (i < 2) {
-      absl::StatusOr<DotOperandDims> dims =
-          DotOperandDims::FromDotOperand(dot, i);
-      if (!dims.ok()) {
-        return FusionDecision::Forbid("Failed to get dot operand dims.");
-      }
-      tracker = TrackerInfo{ShapeTracker(operand->shape()), i, *dims};
-    }
-    queue.push({operand, std::move(tracker)});
+  // Initialize queue with trackers for dot operands.
+  absl::StatusOr<std::array<DotOperandDims, 2>> dot_dims =
+      DotOperandDims::FromDot(dot);
+  if (!dot_dims.ok()) {
+    return FusionDecision::Forbid("Failed to get dot operand dims.");
   }
+  std::queue<std::pair<HloInstruction*, std::optional<TrackerInfo>>> queue(
+      {{dot->mutable_operand(0),
+        TrackerInfo{ShapeTracker(dot->operand(0)->shape()), 0, (*dot_dims)[0]}},
+       {dot->mutable_operand(1),
+        TrackerInfo{ShapeTracker(dot->operand(1)->shape()), 1,
+                    (*dot_dims)[1]}}});
 
   // Start with a fusion containing only the dot instruction.
   auto entry = fusion_search_space.entry();
@@ -2024,13 +2013,9 @@ absl::StatusOr<std::variant<Fusion, FusionDecision>> CreateTileableFusion(
                 fusion_computation->Clone(std::string(name)), original_output};
 }
 
-// Creates a GEMM fusion around a dot-like instruction (kDot or kScaledDot)
-// using the GemmFusionV2 algorithm.
 absl::StatusOr<std::variant<Fusion, FusionDecision>> CreateDotFusionV2(
-    HloInstruction& dot, const se::DeviceDescription& device_description,
+    HloDotInstruction& dot, const se::DeviceDescription& device_description,
     absl::string_view name) {
-  CHECK((HloPredicateIsOp<HloOpcode::kDot, HloOpcode::kScaledDot>(&dot)))
-      << dot.ToString();
   VLOG(3) << "Creating dot fusion v2 around dot: " << dot.ToString();
   FusionSearchSpace fusion_search_space(&dot, device_description);
   VLOG(3) << "Found fusion search space: \n"
@@ -2144,7 +2129,8 @@ absl::StatusOr<std::variant<Fusion, FusionDecision>> CreateDotFusion(
 class GemmFusionVisitor : public DfsHloRewriteVisitor {
  public:
   explicit GemmFusionVisitor(const se::DeviceDescription& device_description)
-      : device_description_(device_description) {}
+      : device_description_(device_description),
+        gpu_version_(device_description.gpu_compute_capability()) {}
   // Checks that a dot() should be targeting the triton GEMM emitter;
   // if so - fuses all its compatible inputs and outputs as a new computation
   // and replaces the original dot() with a call to the computation.
@@ -2176,43 +2162,18 @@ class GemmFusionVisitor : public DfsHloRewriteVisitor {
       return absl::OkStatus();
     }
 
-    return InsertGemmFusion(
-        dot, std::get<Fusion>(std::move(fusion_or_decision)), fusion_name);
-  }
+    Fusion fusion = std::get<Fusion>(std::move(fusion_or_decision));
 
-  absl::Status HandleScaledDot(HloInstruction* scaled_dot) override {
-    CHECK_EQ(scaled_dot->opcode(), HloOpcode::kScaledDot);
-    std::string fusion_name = absl::StrCat("gemm_fusion_", scaled_dot->name());
-    ABSL_ASSIGN_OR_RETURN(
-        auto fusion_or_decision,
-        CreateDotFusionV2(*scaled_dot, device_description_,
-                          absl::StrCat(fusion_name, "_computation")));
-    if (std::holds_alternative<FusionDecision>(fusion_or_decision)) {
-      VLOG(3) << "Not fusing: "
-              << std::get<FusionDecision>(fusion_or_decision).Explain();
-      return absl::OkStatus();
-    }
-    return InsertGemmFusion(scaled_dot,
-                            std::get<Fusion>(std::move(fusion_or_decision)),
-                            fusion_name);
-  }
-
- private:
-  // Wraps `fusion` into a fusion instruction and replaces `fusion.output` in
-  // the original computation with it. `dot_like` is the dot-like instruction
-  // the fusion was created around.
-  absl::Status InsertGemmFusion(HloInstruction* dot_like, Fusion fusion,
-                                absl::string_view fusion_name) {
     HloComputation* computation =
-        dot_like->GetModule()->AddComputationAndUnifyNamesAndIds(
+        dot->GetModule()->AddComputationAndUnifyNamesAndIds(
             std::move(fusion.computation), /*is_entry=*/false);
     HloInstruction* dot_fusion =
-        dot_like->parent()->AddInstruction(HloInstruction::CreateFusion(
+        dot->parent()->AddInstruction(HloInstruction::CreateFusion(
             computation->root_instruction()->shape(),
             HloInstruction::FusionKind::kCustom, fusion.inputs, computation));
-    // Copy the metadata of `dot_like` to the newly created `fusion` op. This
+    // Copy the metadata of the `dot` to the newly created `fusion` op. This
     // is convenient for handling metadata in split-k rewriting subsequently.
-    dot_fusion->set_metadata(dot_like->metadata());
+    dot_fusion->set_metadata(dot->metadata());
     dot_fusion->GetModule()->SetAndUniquifyInstrName(dot_fusion, fusion_name);
 
     ABSL_ASSIGN_OR_RETURN(auto gpu_config,
@@ -2224,7 +2185,7 @@ class GemmFusionVisitor : public DfsHloRewriteVisitor {
 
     HloInstruction* replacement = dot_fusion;
     if (fusion.output->shape() != dot_fusion->shape()) {
-      replacement = dot_like->parent()->AddInstruction(
+      replacement = dot->parent()->AddInstruction(
           HloInstruction::CreateBitcast(fusion.output->shape(), dot_fusion));
     }
 
@@ -2241,7 +2202,62 @@ class GemmFusionVisitor : public DfsHloRewriteVisitor {
     return absl::OkStatus();
   }
 
+  absl::Status HandleScaledDot(HloInstruction* scaled_dot) override {
+    CHECK_EQ(scaled_dot->opcode(), HloOpcode::kScaledDot);
+    HloComputation::Builder builder(
+        absl::StrCat("fusion_", scaled_dot->name()));
+
+    std::vector<HloInstruction*> fusion_inputs;
+
+    std::vector<HlosAndRequirements> hlos_and_reqs;
+    hlos_and_reqs.reserve(scaled_dot->operand_count());
+    ABSL_ASSIGN_OR_RETURN(HlosAndRequirements lhs_hlos_and_reqs,
+                     FuseDotOperand(*scaled_dot, /*operand_index=*/0,
+                                    gpu_version_, builder, fusion_inputs));
+    hlos_and_reqs.push_back(lhs_hlos_and_reqs);
+    ABSL_ASSIGN_OR_RETURN(HlosAndRequirements rhs_hlos_and_reqs,
+                     FuseDotOperand(*scaled_dot, /*operand_index=*/1,
+                                    gpu_version_, builder, fusion_inputs));
+    hlos_and_reqs.push_back(rhs_hlos_and_reqs);
+    ABSL_ASSIGN_OR_RETURN(HlosAndRequirements lhs_scale_hlos_and_reqs,
+                     FuseDotOperand(*scaled_dot, /*operand_index=*/2,
+                                    gpu_version_, builder, fusion_inputs));
+    hlos_and_reqs.push_back(lhs_scale_hlos_and_reqs);
+    ABSL_ASSIGN_OR_RETURN(HlosAndRequirements rhs_scale_hlos_and_reqs,
+                     FuseDotOperand(*scaled_dot, /*operand_index=*/3,
+                                    gpu_version_, builder, fusion_inputs));
+    hlos_and_reqs.push_back(rhs_scale_hlos_and_reqs);
+
+    HloInstruction& fused_dot = FuseDot(*scaled_dot, hlos_and_reqs, builder);
+
+    HlosAndRequirements fused_output_and_reqs =
+        FuseDotOutput(*scaled_dot, fused_dot, gpu_version_,
+                      lhs_hlos_and_reqs.requirements, builder, fusion_inputs);
+    HloComputation* computation =
+        scaled_dot->GetModule()->AddComputationAndUnifyNamesAndIds(
+            builder.Build(),
+            /*is_entry=*/false);
+
+    HloInstruction* fusion =
+        scaled_dot->parent()->AddInstruction(HloInstruction::CreateFusion(
+            computation->root_instruction()->shape(),
+            HloInstruction::FusionKind::kCustom, fusion_inputs, computation));
+
+    ABSL_ASSIGN_OR_RETURN(auto gpu_config,
+                     fusion->backend_config<GpuBackendConfig>());
+    FusionBackendConfig& backend_config =
+        *gpu_config.mutable_fusion_backend_config();
+    backend_config.set_kind(kTritonGemmFusionKind);
+    ABSL_RETURN_IF_ERROR(fusion->set_backend_config(gpu_config));
+    HloInstruction* fusion_output = fused_output_and_reqs.original_hlo;
+    ABSL_RETURN_IF_ERROR(ReplaceInstruction(fusion_output, fusion));
+    MarkAsChanged();
+    return absl::OkStatus();
+  }
+
+ private:
   const se::DeviceDescription& device_description_;
+  se::GpuComputeCapability gpu_version_;
 };
 
 absl::StatusOr<bool> RunOnComputation(

From 767dedfd9057225519566837a45fb2468b760f19 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Thu, 15 Oct 2026 15:51:34 +0530
Subject: [PATCH 0847/1000] Fix doubled-brace f-string bug in
 embedding_lookup_sparse error message

embedding_lookup_sparse's type-mismatch error used
{{type(sp_ids).__name__!r}} with doubled braces, which an f-string renders
as literal text instead of substituting the value. The raised TypeError
showed the literal string '{type(sp_ids).__name__!r}' instead of the
actual type name, making the message useless for debugging.

Fixed by using single braces so the type names are actually interpolated.
Verified against the installed TensorFlow build: reproduces the broken
message before the fix, correct message after.

Added a regression test (ragged_embedding_ops_test.py) that fails on the
old code and passes on the fix, and registered it as a py_test target in
tensorflow/python/ops/ragged/BUILD so the test suite actually runs it.

This commit is based directly on tensorflow:master and intentionally
excludes the unrelated ArgMax/ArgMin CPU kernel and DTensor d_random
changes that were previously mixed into this branch, per review feedback.
---
 tensorflow/python/ops/ragged/BUILD            | 14 +++++
 .../python/ops/ragged/ragged_embedding_ops.py |  4 +-
 .../ops/ragged/ragged_embedding_ops_test.py   | 53 +++++++++++++++++++
 3 files changed, 69 insertions(+), 2 deletions(-)
 create mode 100644 tensorflow/python/ops/ragged/ragged_embedding_ops_test.py

diff --git a/tensorflow/python/ops/ragged/BUILD b/tensorflow/python/ops/ragged/BUILD
index 908ff550c0e..075585a15ae 100644
--- a/tensorflow/python/ops/ragged/BUILD
+++ b/tensorflow/python/ops/ragged/BUILD
@@ -905,6 +905,20 @@ py_test(
     ],
 )
 
+py_test(
+    name = "ragged_embedding_ops_test",
+    srcs = ["ragged_embedding_ops_test.py"],
+    strict_deps = True,
+    deps = [
+        ":ragged_embedding_ops",
+        ":ragged_factory_ops",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:test",
+    ],
+)
+
 py_test(
     name = "ragged_range_op_test",
     srcs = ["ragged_range_op_test.py"],
diff --git a/tensorflow/python/ops/ragged/ragged_embedding_ops.py b/tensorflow/python/ops/ragged/ragged_embedding_ops.py
index a9ea994bfa2..c6cb0939a64 100644
--- a/tensorflow/python/ops/ragged/ragged_embedding_ops.py
+++ b/tensorflow/python/ops/ragged/ragged_embedding_ops.py
@@ -188,8 +188,8 @@ def embedding_lookup_sparse(
     if not isinstance(rt_weights, ragged_tensor.RaggedTensor):
       raise TypeError(
           f"sp_ids must be of the same type as sp_weights, "
-          f"received {{type(sp_ids).__name__!r}} for sp_ids and "
-          f"{{type(sp_weights).__name__!r}} for sp_weights."
+          f"received {type(sp_ids).__name__!r} for sp_ids and "
+          f"{type(sp_weights).__name__!r} for sp_weights."
       )
     rt_ids.values.get_shape().assert_is_compatible_with(
         rt_weights.values.get_shape()
diff --git a/tensorflow/python/ops/ragged/ragged_embedding_ops_test.py b/tensorflow/python/ops/ragged/ragged_embedding_ops_test.py
new file mode 100644
index 00000000000..aec4b1fb6ec
--- /dev/null
+++ b/tensorflow/python/ops/ragged/ragged_embedding_ops_test.py
@@ -0,0 +1,53 @@
+# 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 ragged_embedding_ops."""
+
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import variables
+from tensorflow.python.ops.ragged import ragged_embedding_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import googletest
+
+
+class RaggedEmbeddingOpsTest(test_util.TensorFlowTestCase):
+
+  def testEmbeddingLookupSparseMismatchedTypesErrorMessage(self):
+    # Regression test: the error message's `received {{type(sp_ids)...}}`
+    # f-string used doubled braces, which f-strings treat as literal `{`
+    # and `}` characters rather than substitution syntax. As a result the
+    # raised TypeError displayed the literal text
+    # "{type(sp_ids).__name__!r}" instead of the actual type name, which
+    # is useless for debugging a real type mismatch.
+    params = variables.Variable([[1.0, 2.0], [3.0, 4.0]])
+    sp_ids = ragged_factory_ops.constant([[1], [0]])
+    sp_weights = sparse_tensor.SparseTensor(
+        indices=[[0, 0]], values=[1.0], dense_shape=[1, 1])
+
+    with self.assertRaises(TypeError) as cm:
+      ragged_embedding_ops.embedding_lookup_sparse(params, sp_ids,
+                                                    sp_weights)
+
+    message = str(cm.exception)
+    # The substituted type names must appear in the message...
+    self.assertIn('RaggedTensor', message)
+    self.assertIn('SparseTensor', message)
+    # ...and the unsubstituted placeholder text must not.
+    self.assertNotIn('{type(sp_ids)', message)
+    self.assertNotIn('{type(sp_weights)', message)
+
+
+if __name__ == '__main__':
+  googletest.main()

From a0bb7c3c82157dc9cd801f09c256d5f8c895974f Mon Sep 17 00:00:00 2001
From: Alejandro Pacheco <alepacheco99@gmail.com>
Date: Thu, 15 Oct 2026 12:41:59 +0200
Subject: [PATCH 0848/1000] Match the shape error in the lookup table init
 tests

Check that the rejected batches fail with the shape error from
CheckKeyAndValueTensorsForInsert, not just any InvalidArgument.
---
 tensorflow/core/kernels/BUILD                       |  1 +
 .../core/kernels/initializable_lookup_table_test.cc | 13 ++++++++++---
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 0c4d567d2df..8fe0111cced 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -3695,6 +3695,7 @@ tf_cc_test(
         "//tensorflow/core:test_main",
         "//tensorflow/core:testlib",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
     ],
 )
 
diff --git a/tensorflow/core/kernels/initializable_lookup_table_test.cc b/tensorflow/core/kernels/initializable_lookup_table_test.cc
index 58608e4c4b6..da1f4fd8598 100644
--- a/tensorflow/core/kernels/initializable_lookup_table_test.cc
+++ b/tensorflow/core/kernels/initializable_lookup_table_test.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
 #include "tensorflow/core/framework/tensor.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/tensor_testutil.h"
@@ -80,7 +81,9 @@ TEST(InitializableLookupTableTest, RejectsFirstBatchWithMismatchedShapes) {
   BatchIterator iter({
       {test::AsTensor<int64_t>({1, 2}), test::AsTensor<int64_t>({10})},
   });
-  EXPECT_TRUE(absl::IsInvalidArgument(table->Initialize(iter)));
+  EXPECT_THAT(table->Initialize(iter),
+              absl_testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                                     ::testing::HasSubstr("Expected shape")));
   EXPECT_FALSE(table->is_initialized());
 }
 
@@ -92,7 +95,9 @@ TEST(InitializableLookupTableTest, RejectsLaterBatchWithMoreKeysThanValues) {
       {test::AsTensor<int64_t>({1}), test::AsTensor<int64_t>({10})},
       {test::AsTensor<int64_t>({2, 3, 4, 5}), test::AsTensor<int64_t>({20})},
   });
-  EXPECT_TRUE(absl::IsInvalidArgument(table->Initialize(iter)));
+  EXPECT_THAT(table->Initialize(iter),
+              absl_testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                                     ::testing::HasSubstr("Expected shape")));
   EXPECT_FALSE(table->is_initialized());
 }
 
@@ -103,7 +108,9 @@ TEST(InitializableLookupTableTest, RejectsLaterBatchWithExtraValueDims) {
       {test::AsTensor<int64_t>({2}),
        test::AsTensor<int64_t>({20, 21}, TensorShape({1, 2}))},
   });
-  EXPECT_TRUE(absl::IsInvalidArgument(table->Initialize(iter)));
+  EXPECT_THAT(table->Initialize(iter),
+              absl_testing::StatusIs(absl::StatusCode::kInvalidArgument,
+                                     ::testing::HasSubstr("Expected shape")));
   EXPECT_FALSE(table->is_initialized());
 }
 

From fc2a52a8d2719489746784b245d2b8d25f475816 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Thu, 15 Oct 2026 00:45:26 -0700
Subject: [PATCH 0849/1000] [XLA:GPU] Partially rollback device_error_callback
 in GpuExecutable.

Enqueuing a host callback on every thunk execution when `device_error_callback` is set causes a step time regression in multi-node runs. Disable using it in `GpuExecutable` until a lower-overhead mechanism is implemented.

PiperOrigin-RevId: 999930161
---
 third_party/xla/xla/service/gpu/gpu_executable.cc      | 4 +++-
 third_party/xla/xla/service/gpu/gpu_executable_test.cc | 4 +++-
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/gpu_executable.cc b/third_party/xla/xla/service/gpu/gpu_executable.cc
index 7182d9ff96b..9667ea18f52 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable.cc
@@ -927,8 +927,10 @@ absl::Status GpuExecutable::ExecuteThunksImpl(
         HangWatchdog::Abort(device_watchdog_name, device_timeout)));
   }
 
+  // TODO(b/573818657): Re-enable device_error_callback once the host callback
+  // performance regression is fixed.
   std::shared_ptr<const GpuExecutableRunOptions::DeviceErrorCallback> error_cb =
-      gpu_run_options ? gpu_run_options->device_error_callback() : nullptr;
+      nullptr;
 
   if (!device_guards.empty() || error_cb != nullptr) {
     ABSL_RETURN_IF_ERROR(main_stream->DoHostCallback(
diff --git a/third_party/xla/xla/service/gpu/gpu_executable_test.cc b/third_party/xla/xla/service/gpu/gpu_executable_test.cc
index 5c0eb4923db..4db82daba54 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable_test.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable_test.cc
@@ -1374,7 +1374,9 @@ class FakeDeviceAddressAllocator : public se::DeviceAddressAllocator {
   }
 };
 
-TEST_F(GpuExecutableTest, InvokesDeviceErrorCallbackOnStreamError) {
+// TODO(b/573818657): Re-enable once device_error_callback is wired without
+// per-step DoHostCallback overhead.
+TEST_F(GpuExecutableTest, DISABLED_InvokesDeviceErrorCallbackOnStreamError) {
   ASSERT_OK_AND_ASSIGN(
       se::Platform * platform,
       se::PlatformManager::PlatformWithId(stream_executor_platform_id()));

From 40014889fb9313abd8f5b7c0d29d58376756a162 Mon Sep 17 00:00:00 2001
From: Alexandros Theodoridis <atheodor@amd.com>
Date: Thu, 15 Oct 2026 01:30:29 -0700
Subject: [PATCH 0850/1000] PR #50746: [ROCm] Remove flaky attempts to not
 pollute rocm rbe cluster on test failures
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50746

📝 Summary of Changes
Remove flaky attemps so there are no flaky tests in rocm ci

🎯 Justification
When test start to fail we execute it 3 times.
At some point if it is flaky it get get merged
due to one of these attempts succeed.
Later on it pollutes our CI runners with
flaky executions not allowing other jobs to run.

🚀 Kind of Contribution
Please remove what does not apply: 🐛 Bug Fix

📊 Benchmark (for Performance Improvements)
Not relevant

🧪 Unit Tests:
CI

🧪 Execution Tests:
CI

Copybara import of the project:

--
28e88ef9fd0c6d6dda2da13f31756ab808bc9f17 by Alexandros Theodoridis <atheodor@amd.com>:

Remove flaky attempts to not pollute rocm rbe cluster on test failures

Merging this change closes #50746

PiperOrigin-RevId: 999953031
---
 third_party/xla/build_tools/rocm/rocm_xla.bazelrc | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/build_tools/rocm/rocm_xla.bazelrc b/third_party/xla/build_tools/rocm/rocm_xla.bazelrc
index f69f7bbd880..2da9532b930 100644
--- a/third_party/xla/build_tools/rocm/rocm_xla.bazelrc
+++ b/third_party/xla/build_tools/rocm/rocm_xla.bazelrc
@@ -54,12 +54,12 @@ build:tsan --run_under=//build_tools/rocm:sanitizer_wrapper
 
 build:asan --run_under=//build_tools/rocm:sanitizer_wrapper
 
-build:ci_single_gpu --flaky_test_attempts=3
+build:ci_single_gpu --config=rocm_base
 
+build:ci_multi_gpu --config=rocm_base
 build:ci_multi_gpu --action_env=XLA_FLAGS="--xla_gpu_force_compilation_parallelism=16"
 build:ci_multi_gpu --action_env=NCCL_MAX_NCHANNELS=1
 build:ci_multi_gpu --test_sharding_strategy=disabled
-build:ci_multi_gpu --flaky_test_attempts=3
 build:ci_multi_gpu --experimental_guard_against_concurrent_changes
 build:ci_multi_gpu --test_env=HIP_VISIBLE_DEVICES=0,1,2,3
 

From 9d3dc4d1550eba6473665bf34427aa701ef296be Mon Sep 17 00:00:00 2001
From: Alejandro Pacheco <alepacheco99@gmail.com>
Date: Thu, 15 Oct 2026 13:15:48 +0200
Subject: [PATCH 0851/1000] Depend on platform:refcount in the lookup table
 init test

---
 tensorflow/core/kernels/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 8fe0111cced..9229669e6a4 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -3690,10 +3690,10 @@ tf_cc_test(
         ":initializable_lookup_table",
         ":lookup_table_op",
         "//tensorflow/core:framework",
-        "//tensorflow/core:lib",
         "//tensorflow/core:test",
         "//tensorflow/core:test_main",
         "//tensorflow/core:testlib",
+        "//tensorflow/core/platform:refcount",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
     ],

From 14c10e5b74c4a8faccbed64fbf0084c088822b3a Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Thu, 15 Oct 2026 01:32:47 -0700
Subject: [PATCH 0852/1000] [XLA:GPU]: Add variadic xtile.reduce_scatter op and
 migrate reduce-scatter lowering.

StableHLO does not have a variadic reduce-scatter. So we add a shadow xtile instruction to lower through that instead of stablehlo.

PiperOrigin-RevId: 999954157
---
 .../gpu/codegen/triton/collective_emitter.cc  |  6 +-
 .../gpu/codegen/triton/collective_emitter.h   |  5 +-
 .../fusion_emitter_shared_dialect_test.cc     | 50 ++++++++++++-
 .../transforms/stablehlo_lower_to_triton.cc   |  5 +-
 .../tests/stable_hlo_to_triton_lowering.mlir  | 24 +++----
 .../codegen/xtile/codegen/emitter_helpers.cc  |  6 +-
 .../codegen/xtile/codegen/emitter_helpers.h   |  5 +-
 .../codegen/experimental_fusion_emitter.cc    | 69 +++++++++++++-----
 .../codegen/xtile/codegen/fusion_emitter.cc   |  8 +--
 .../xla/xla/codegen/xtile/ir/tests/ops.mlir   | 71 +++++++++++++++++++
 .../ir/transforms/verify_legal_xtile_ops.cc   |  7 +-
 .../xla/xla/codegen/xtile/ir/xtile_ops.cc     | 66 +++++++++++++++++
 .../xla/xla/codegen/xtile/ir/xtile_ops.td     | 42 +++++++++++
 13 files changed, 311 insertions(+), 53 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
index 0d03e7bfe79..232d20bd061 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
@@ -249,7 +249,7 @@ absl::StatusOr<AllReduceEmitterContext> CreateAllReduceEmitterContext(
 // happens when the collective fusion is formed) and creates the
 // ReduceScatterEmitterContext.
 absl::StatusOr<ReduceScatterEmitterContext> CreateReduceScatterEmitterContext(
-    mlir::stablehlo::ReduceScatterOp op) {
+    xtile::ReduceScatterOp op) {
   ABSL_ASSIGN_OR_RETURN(ReduceScatterEmitterContext ctx,
                    CreateReductionEmitterContext<ReduceScatterEmitterContext>(
                        op, /*has_invocation_count=*/false));
@@ -295,7 +295,7 @@ absl::StatusOr<ReduceScatterEmitterContext> CreateReduceScatterEmitterContext(
   ctx.shard_size = shard_size;
   ctx.signal_stride = signal_stride;
   const llvm::ArrayRef<int64_t> output_tile_shape =
-      mlir::cast<mlir::RankedTensorType>(op.getType()).getShape();
+      mlir::cast<mlir::RankedTensorType>(op.getType(0)).getShape();
   ctx.output_tile_shape.assign(output_tile_shape.begin(),
                                output_tile_shape.end());
   return ctx;
@@ -1499,7 +1499,7 @@ mlir::LogicalResult RewriteAllReduce(mlir::stablehlo::AllReduceOp op,
   return AllReduceEmitter::Emit(*std::move(maybe_context), rewriter);
 }
 
-mlir::LogicalResult RewriteReduceScatter(mlir::stablehlo::ReduceScatterOp op,
+mlir::LogicalResult RewriteReduceScatter(xtile::ReduceScatterOp op,
                                          mlir::PatternRewriter& rewriter) {
   const mlir::Location loc = op->getLoc();
   absl::StatusOr<ReduceScatterEmitterContext> maybe_context =
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h
index de79607c935..86c93904ec8 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h
@@ -29,6 +29,7 @@ limitations under the License.
 #include "mlir/Support/LLVM.h"
 #include "stablehlo/dialect/StablehloOps.h"
 #include "xla/backends/gpu/runtime/collective_params.h"
+#include "xla/codegen/xtile/ir/xtile_ops.h"
 #include "xla/codegen/xtile/xtile_config.pb.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
@@ -95,8 +96,8 @@ mlir::LogicalResult RewriteAllReduce(mlir::stablehlo::AllReduceOp op,
 mlir::LogicalResult RewriteAllGather(mlir::stablehlo::AllGatherOp op,
                                      mlir::PatternRewriter& rewriter);
 
-// Rewrites stablehlo reduce-scatter op to a triton implementation.
-mlir::LogicalResult RewriteReduceScatter(mlir::stablehlo::ReduceScatterOp op,
+// Rewrites xtile reduce-scatter op to a triton implementation.
+mlir::LogicalResult RewriteReduceScatter(xla::xtile::ReduceScatterOp op,
                                          mlir::PatternRewriter& rewriter);
 
 // Returns the symmetric memory type for the scratch buffers of the collective
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc
index 0361c93ba3a..ffc129a9651 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_shared_dialect_test.cc
@@ -460,7 +460,55 @@ TEST_F(XTileDialectTest, HloReduceScatterIsLoweredToStableHloReduceScatter) {
       block_level_parameters,
       R"(
 CHECK: %[[INPUT_TILE:.*]] = xtile.extract %arg0[%{{.*}}, %{{.*}}] [2, 1024] [1, 1] : memref<4x16384xf32> -> tensor<2x1024xf32>
-CHECK: stablehlo.reduce_scatter
+CHECK: xtile.reduce_scatter
+CHECK: stablehlo.add
+)"));
+}
+
+TEST_F(XTileDialectTest,
+       VariadicHloReduceScatterIsLoweredToXTileReduceScatter) {
+  constexpr absl::string_view kHloText =
+      R"(
+      HloModule wrapped_module_variadic_reduce-scatter
+
+      %apply_op {
+        %x = f32[] parameter(0)
+        %y = f32[] parameter(1)
+        ROOT %apply_op = f32[] add(%x, %y)
+      }
+
+      %wrapped_reduce-scatter {
+        %p0 = f32[4,16384]{1,0} parameter(0)
+        %p1 = f32[4,8]{1,0} parameter(1)
+        %rs = (f32[2,16384]{1,0}, f32[2,8]{1,0}) reduce-scatter(%p0, %p1),
+          replica_groups={{0,1}}, dimensions={0}, to_apply=%apply_op
+        %gte0 = f32[2,16384]{1,0} get-tuple-element(%rs), index=0
+        %gte1 = f32[2,8]{1,0} get-tuple-element(%rs), index=1
+        ROOT %tuple = (f32[2,16384]{1,0}, f32[2,8]{1,0}) tuple(%gte0, %gte1)
+      }
+
+      ENTRY %entry {
+        %p0 = f32[4,16384]{1,0} parameter(0)
+        %p1 = f32[4,8]{1,0} parameter(1)
+        ROOT %fusion = (f32[2,16384]{1,0}, f32[2,8]{1,0}) fusion(%p0, %p1),
+          kind=kLoop, calls=%wrapped_reduce-scatter,
+          backend_config={"fusion_backend_config":{"kind":"__triton_collective","block_level_fusion_config":{"num_warps":"16","output_tiles":[{"sizes":["1","1024"]},{"sizes":["1","8"]}],"num_ctas":1,"num_stages":1,"is_tma_allowed":false,"is_warp_specialization_allowed":false}}}
+      }
+    )";
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> hlo_module,
+                       ParseAndReturnUnverifiedModule(kHloText));
+
+  BlockLevelParameters block_level_parameters;
+  block_level_parameters.output_tile_sizes = {{1, 1024}, {1, 8}};
+
+  EXPECT_OK(CreateXTileIrAndFileCheck(
+      *hlo_module->GetComputationWithName("wrapped_reduce-scatter"),
+      block_level_parameters,
+      R"(
+CHECK-DAG: xtile.extract %arg0[%{{.*}}, %{{.*}}] [2, 1024] [1, 1] : memref<4x16384xf32> -> tensor<2x1024xf32>
+CHECK-DAG: xtile.extract %arg1[%{{.*}}, %{{.*}}] [2, 8] [1, 1] : memref<4x8xf32> -> tensor<2x8xf32>
+CHECK: xtile.reduce_scatter(%{{.*}}, %{{.*}}) scatter_dimension = 0 {{.*}} : (tensor<2x1024xf32>, tensor<2x8xf32>) -> (tensor<1x1024xf32>, tensor<1x8xf32>)
 CHECK: stablehlo.add
 )"));
 }
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/stablehlo_lower_to_triton.cc b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/stablehlo_lower_to_triton.cc
index 8e0164a994d..6531fb2aa49 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/stablehlo_lower_to_triton.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/stablehlo_lower_to_triton.cc
@@ -64,6 +64,7 @@ limitations under the License.
 #include "xla/backends/gpu/codegen/triton/transforms/passes.h"
 #include "xla/codegen/xtile/codegen/dot_algorithms.h"
 #include "xla/codegen/xtile/codegen/emitter_helpers.h"
+#include "xla/codegen/xtile/ir/xtile_ops.h"
 #include "xla/hlo/translate/mhlo_to_hlo/attribute_exporter.h"
 #include "xla/service/algorithm_util.h"
 #include "xla/tsl/platform/errors.h"
@@ -1055,13 +1056,13 @@ class LowerAllReduce : public mlir::OpRewritePattern<stablehlo::AllReduceOp> {
 };
 
 class LowerReduceScatter
-    : public mlir::OpRewritePattern<stablehlo::ReduceScatterOp> {
+    : public mlir::OpRewritePattern<::xla::xtile::ReduceScatterOp> {
  public:
   using OpRewritePattern::OpRewritePattern;
 
  private:
   mlir::LogicalResult matchAndRewrite(
-      stablehlo::ReduceScatterOp op,
+      ::xla::xtile::ReduceScatterOp op,
       mlir::PatternRewriter& rewriter) const override {
     return ::xla::gpu::RewriteReduceScatter(op, rewriter);
   }
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/stable_hlo_to_triton_lowering.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/stable_hlo_to_triton_lowering.mlir
index e17bcd81500..d14ae9fa827 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/stable_hlo_to_triton_lowering.mlir
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/stable_hlo_to_triton_lowering.mlir
@@ -468,12 +468,12 @@ xtile.entry_func @reduce_scatter_one_shot(%input: memref<4x1024xf32>, %output: m
   %c0 = arith.constant 0 : index
   %tile = xtile.extract %input[%c0, %tile_id][2, 64][1, 1] : memref<4x1024xf32> -> tensor<2x64xf32>
   // CHECK: triton_xla.block_barrier {{.*}} <world_size = 2, signal_stride = 16, barrier_mode = producer_symmetric>
-  // CHECK-NOT: stablehlo.reduce_scatter
-  %reduce_scatter = "stablehlo.reduce_scatter"(%tile) <{replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>, scatter_dimension = 0 : i64}> ({
-    ^bb0(%arg7: tensor<f32>, %arg8: tensor<f32>):
-      %4 = arith.addf %arg7, %arg8 : tensor<f32>
-      stablehlo.return %4 : tensor<f32>
-    }) : (tensor<2x64xf32>) -> tensor<1x64xf32>
+  // CHECK-NOT: xtile.reduce_scatter
+  %reduce_scatter = xtile.reduce_scatter(%tile) scatter_dimension = 0 <replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>> : (tensor<2x64xf32>) -> (tensor<1x64xf32>) {
+  ^bb0(%arg7: tensor<f32>, %arg8: tensor<f32>):
+    %4 = arith.addf %arg7, %arg8 : tensor<f32>
+    stablehlo.return %4 : tensor<f32>
+  }
   xtile.insert %reduce_scatter into %output[%c0, %tile_id][1, 64][1, 1] : tensor<1x64xf32> -> memref<2x1024xf32>
   xtile.return
 }
@@ -482,12 +482,12 @@ xtile.entry_func @reduce_scatter_one_shot(%input: memref<4x1024xf32>, %output: m
 xtile.entry_func @reduce_scatter_invalid_tile_size_doesnt_lower(%input: memref<4x1024xf32>, %output: memref<2x1024xf32>, %device_rank: i32, %signal_buffer: !tt.ptr<i64>, %remote_input_buffer: !tt.ptr<i64>, %tile_id: index) attributes {num_opaque_args = 3 : i32} {
   %c0 = arith.constant 0 : index
   %tile = xtile.extract %input[%c0, %tile_id][4, 64][1, 1] : memref<4x1024xf32> -> tensor<4x64xf32>
-  // CHECK: stablehlo.reduce_scatter
-  %reduce_scatter = "stablehlo.reduce_scatter"(%tile) <{replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>, scatter_dimension = 0 : i64}> ({
-    ^bb0(%arg7: tensor<f32>, %arg8: tensor<f32>):
-      %4 = arith.addf %arg7, %arg8 : tensor<f32>
-      stablehlo.return %4 : tensor<f32>
-    }) : (tensor<4x64xf32>) -> tensor<2x64xf32>
+  // CHECK: xtile.reduce_scatter
+  %reduce_scatter = xtile.reduce_scatter(%tile) scatter_dimension = 0 <replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>> : (tensor<4x64xf32>) -> (tensor<2x64xf32>) {
+  ^bb0(%arg7: tensor<f32>, %arg8: tensor<f32>):
+    %4 = arith.addf %arg7, %arg8 : tensor<f32>
+    stablehlo.return %4 : tensor<f32>
+  }
   xtile.insert %reduce_scatter into %output[%c0, %tile_id][2, 64][1, 1] : tensor<2x64xf32> -> memref<2x1024xf32>
   xtile.return
 }
diff --git a/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.cc b/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.cc
index b15a6ccdb4f..39078ba1a47 100644
--- a/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.cc
+++ b/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.cc
@@ -1260,12 +1260,10 @@ TensorValue EmitTiledTranspose(mlir::ImplicitLocOpBuilder& b,
 }
 
 absl::Status EmitReduceComputation(mlir::ImplicitLocOpBuilder& b,
-                                   const HloInstruction* hlo_reduction,
+                                   PrimitiveType element_type,
                                    const HloComputation* reduction_computation,
                                    mlir::Operation* reduction) {
-  ABSL_ASSIGN_OR_RETURN(
-      Type result_ty,
-      PrimitiveTypeToMlirType(b, hlo_reduction->shape().element_type()));
+  ABSL_ASSIGN_OR_RETURN(Type result_ty, PrimitiveTypeToMlirType(b, element_type));
   result_ty = mlir::RankedTensorType::get({}, result_ty);
 
   mlir::Location loc = b.getLoc();
diff --git a/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.h b/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.h
index f8333680506..1a606216fa2 100644
--- a/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.h
+++ b/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.h
@@ -555,9 +555,10 @@ TensorValue EmitTiledTranspose(mlir::ImplicitLocOpBuilder& b,
                                llvm::SmallVector<int64_t> dimensions,
                                TensorValue input);
 
-// Emits a reduction computation.
+// Emits a reduction computation (two scalar arguments of `element_type`) as
+// the region of `reduction`.
 absl::Status EmitReduceComputation(mlir::ImplicitLocOpBuilder& b,
-                                   const HloInstruction* hlo_reduction,
+                                   PrimitiveType element_type,
                                    const HloComputation* reduction_computation,
                                    mlir::Operation* reduction);
 }  // namespace xla::xtile
diff --git a/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc b/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc
index 4facf30eec6..53b141ffc85 100644
--- a/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc
+++ b/third_party/xla/xla/codegen/xtile/codegen/experimental_fusion_emitter.cc
@@ -255,12 +255,13 @@ absl::StatusOr<TensorValue> EmitAllReduce(
       b, output_type, operands, replica_groups_attr, channel_handle_attr,
       use_global_device_ids);
 
-  ABSL_RETURN_IF_ERROR(EmitReduceComputation(b, all_reduce, all_reduce->to_apply(),
-                                        all_reduce_op));
+  ABSL_RETURN_IF_ERROR(EmitReduceComputation(b, all_reduce->shape().element_type(),
+                                        all_reduce->to_apply(), all_reduce_op));
   return mlir::cast<TensorValue>(all_reduce_op.getResult(0));
 }
 
-absl::StatusOr<TensorValue> EmitReduceScatter(
+// Emits a (possibly variadic) reduce-scatter and returns one tensor per output.
+absl::StatusOr<EmittedResults> EmitReduceScatter(
     EmitterContext& emitter_ctx,
     const HloReduceScatterInstruction* reduce_scatter,
     const ge::TiledHloInstruction& tiled_reduce_scatter, ValueRange operands) {
@@ -279,13 +280,32 @@ absl::StatusOr<TensorValue> EmitReduceScatter(
   }
 
   ImplicitLocOpBuilder& b = emitter_ctx.b();
-  ABSL_ASSIGN_OR_RETURN(auto output_element_type,
-                   xtile::PrimitiveTypeToMlirType(
-                       b, reduce_scatter->shape().element_type()));
-  ABSL_ASSIGN_OR_RETURN(SmallVector<int64_t> tile_sizes,
-                   tiled_reduce_scatter.tile().GetStaticTileSizes());
-  auto output_type = mlir::RankedTensorType::get(GetPaddedTileSizes(tile_sizes),
-                                                 output_element_type);
+  // One result per operand. For a variadic reduce-scatter the HLO shape is a
+  // tuple and the tiled instruction holds one tile per output.
+  const bool is_variadic = reduce_scatter->shape().IsTuple();
+  const int64_t num_outputs =
+      is_variadic ? reduce_scatter->shape().tuple_shapes().size() : 1;
+  TF_RET_CHECK(operands.size() == num_outputs)
+      << "Expected one operand per reduce-scatter output, got "
+      << operands.size() << " operands and " << num_outputs << " outputs.";
+  SmallVector<Type> output_types;
+  output_types.reserve(num_outputs);
+  for (int64_t k = 0; k < num_outputs; ++k) {
+    const Shape& output_shape = is_variadic
+                                    ? reduce_scatter->shape().tuple_shapes(k)
+                                    : reduce_scatter->shape();
+    ABSL_ASSIGN_OR_RETURN(
+        auto output_element_type,
+        xtile::PrimitiveTypeToMlirType(b, output_shape.element_type()));
+    ABSL_ASSIGN_OR_RETURN(SmallVector<int64_t> tile_sizes,
+                     tiled_reduce_scatter.tile(k).GetStaticTileSizes());
+    output_types.push_back(mlir::RankedTensorType::get(
+        GetPaddedTileSizes(tile_sizes), output_element_type));
+  }
+  // All operands share the reducer, so they share the element type.
+  const PrimitiveType element_type =
+      is_variadic ? reduce_scatter->shape().tuple_shapes(0).element_type()
+                  : reduce_scatter->shape().element_type();
 
   auto replica_groups_type = mlir::RankedTensorType::get(
       {static_cast<int64_t>(reduce_scatter->replica_groups().size()),
@@ -294,13 +314,23 @@ absl::StatusOr<TensorValue> EmitReduceScatter(
       b.getI64Type());
   auto replica_groups_attr = mlir::DenseIntElementsAttr::get(
       replica_groups_type, flattened_replica_group_ids);
-  auto reduce_scatter_op = mlir::stablehlo::ReduceScatterOp::create(
-      b, output_type, operands[0], reduce_scatter->scatter_dimension(),
-      replica_groups_attr, /*channel_handle=*/nullptr);
+
+  // `stablehlo.reduce_scatter` is single-operand/single-result (unlike
+  // `stablehlo.all_gather`/`all_reduce`), so reduce-scatter is emitted as
+  // `xtile.reduce_scatter` instead of changing the StableHLO dialect. See the
+  // op description in xtile_ops.td.
+  auto reduce_scatter_op = xtile::ReduceScatterOp::create(
+      b, mlir::TypeRange{output_types}, operands,
+      reduce_scatter->scatter_dimension(), replica_groups_attr);
 
   ABSL_RETURN_IF_ERROR(EmitReduceComputation(
-      b, reduce_scatter, reduce_scatter->to_apply(), reduce_scatter_op));
-  return mlir::cast<TensorValue>(reduce_scatter_op.getResult());
+      b, element_type, reduce_scatter->to_apply(), reduce_scatter_op));
+  EmittedResults results;
+  results.reserve(num_outputs);
+  for (Value result : reduce_scatter_op->getResults()) {
+    results.push_back(mlir::cast<TensorValue>(result));
+  }
+  return results;
 }
 
 absl::StatusOr<TensorValue> EmitBroadcast(
@@ -2128,8 +2158,9 @@ absl::StatusOr<TensorValue> EmitReduceWithNoRegion(
       input_value, mask_dim_bounds, neutral_value));
   stablehlo::ReduceOp reduction = stablehlo::ReduceOp::create(
       b, input_value, init_value, reduce_hlo.dimensions());
-  ABSL_RETURN_IF_ERROR(EmitReduceComputation(
-      b, &reduce_hlo, tiled_hlo.hlo()->to_apply(), reduction));
+  ABSL_RETURN_IF_ERROR(EmitReduceComputation(b, reduce_hlo.shape().element_type(),
+                                        tiled_hlo.hlo()->to_apply(),
+                                        reduction));
   return mlir::cast<TensorValue>(reduction.getResult(0));
 }
 
@@ -2275,8 +2306,8 @@ absl::StatusOr<TensorValue> EmitReduceWithRegion(
 
   stablehlo::ReduceOp reduction = stablehlo::ReduceOp::create(
       b, loop_result, init_value, reduce_hlo.dimensions());
-  ABSL_RETURN_IF_ERROR(
-      EmitReduceComputation(b, &reduce_hlo, reduce_hlo.to_apply(), reduction));
+  ABSL_RETURN_IF_ERROR(EmitReduceComputation(b, reduce_hlo.shape().element_type(),
+                                        reduce_hlo.to_apply(), reduction));
   return mlir::cast<TensorValue>(reduction.getResult(0));
 }
 
diff --git a/third_party/xla/xla/codegen/xtile/codegen/fusion_emitter.cc b/third_party/xla/xla/codegen/xtile/codegen/fusion_emitter.cc
index 83b44deb93b..77f11982fbf 100644
--- a/third_party/xla/xla/codegen/xtile/codegen/fusion_emitter.cc
+++ b/third_party/xla/xla/codegen/xtile/codegen/fusion_emitter.cc
@@ -169,8 +169,8 @@ absl::StatusOr<TensorValue> EmitReduce(
 
   stablehlo::ReduceOp reduction = stablehlo::ReduceOp::create(
       b, input, init_value, hlo_reduce.dimensions());
-  ABSL_RETURN_IF_ERROR(
-      EmitReduceComputation(b, &hlo_reduce, hlo_reduce.to_apply(), reduction));
+  ABSL_RETURN_IF_ERROR(EmitReduceComputation(b, hlo_reduce.shape().element_type(),
+                                        hlo_reduce.to_apply(), reduction));
 
   return mlir::cast<TensorValue>(reduction.getResult(0));
 }
@@ -945,8 +945,8 @@ absl::StatusOr<TensorValue> EmitAllReduce(
       b, b.getLoc(), output_type, mlir::ValueRange(operands),
       replica_groups_attr, channel_handle_attr, use_global_device_ids);
 
-  ABSL_RETURN_IF_ERROR(EmitReduceComputation(b, &all_reduce, all_reduce.to_apply(),
-                                        all_reduce_op));
+  ABSL_RETURN_IF_ERROR(EmitReduceComputation(b, all_reduce.shape().element_type(),
+                                        all_reduce.to_apply(), all_reduce_op));
 
   return mlir::cast<TensorValue>(all_reduce_op.getResult(0));
 }
diff --git a/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir b/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir
index 2f9e60aacc1..6bb45d84ef7 100644
--- a/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir
+++ b/third_party/xla/xla/codegen/xtile/ir/tests/ops.mlir
@@ -271,6 +271,77 @@ func.func @scan_op_carry_type_mismatch(%input: tensor<10x20x30xf32>, %init: tens
 
 // -----
 
+func.func @reduce_scatter_op(%input0: tensor<2x64xf32>, %input1: tensor<2x8xf32>) -> (tensor<1x64xf32>, tensor<1x8xf32>) {
+  %out0, %out1 = xtile.reduce_scatter(%input0, %input1) scatter_dimension = 0 <replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>> : (tensor<2x64xf32>, tensor<2x8xf32>) -> (tensor<1x64xf32>, tensor<1x8xf32>) {
+  ^bb0(%arg0: tensor<f32>, %arg1: tensor<f32>):
+    %add = stablehlo.add %arg0, %arg1 : tensor<f32>
+    stablehlo.return %add : tensor<f32>
+  }
+  return %out0, %out1 : tensor<1x64xf32>, tensor<1x8xf32>
+}
+
+// -----
+
+func.func @reduce_scatter_op_output_count_mismatch(%input0: tensor<2x64xf32>, %input1: tensor<2x8xf32>) -> tensor<1x64xf32> {
+  // expected-error @+1 {{expects one output per input, but got 2 inputs and 1 outputs}}
+  %out0 = xtile.reduce_scatter(%input0, %input1) scatter_dimension = 0 <replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>> : (tensor<2x64xf32>, tensor<2x8xf32>) -> (tensor<1x64xf32>) {
+  ^bb0(%arg0: tensor<f32>, %arg1: tensor<f32>):
+    %add = stablehlo.add %arg0, %arg1 : tensor<f32>
+    stablehlo.return %add : tensor<f32>
+  }
+  return %out0 : tensor<1x64xf32>
+}
+
+// -----
+
+func.func @reduce_scatter_op_element_type_mismatch(%input0: tensor<2x64xf32>, %input1: tensor<2x8xbf16>) -> (tensor<1x64xf32>, tensor<1x8xbf16>) {
+  // expected-error @+1 {{input #1 element type 'bf16' does not match the element type of input #0 'f32'}}
+  %out0, %out1 = xtile.reduce_scatter(%input0, %input1) scatter_dimension = 0 <replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>> : (tensor<2x64xf32>, tensor<2x8xbf16>) -> (tensor<1x64xf32>, tensor<1x8xbf16>) {
+  ^bb0(%arg0: tensor<f32>, %arg1: tensor<f32>):
+    %add = stablehlo.add %arg0, %arg1 : tensor<f32>
+    stablehlo.return %add : tensor<f32>
+  }
+  return %out0, %out1 : tensor<1x64xf32>, tensor<1x8xbf16>
+}
+
+// -----
+
+func.func @reduce_scatter_op_scatter_dimension_out_of_range(%input: tensor<2x64xf32>) -> tensor<1x64xf32> {
+  // expected-error @+1 {{scatter dimension 2 is out of range for input #0 of rank 2}}
+  %out = xtile.reduce_scatter(%input) scatter_dimension = 2 <replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>> : (tensor<2x64xf32>) -> (tensor<1x64xf32>) {
+  ^bb0(%arg0: tensor<f32>, %arg1: tensor<f32>):
+    %add = stablehlo.add %arg0, %arg1 : tensor<f32>
+    stablehlo.return %add : tensor<f32>
+  }
+  return %out : tensor<1x64xf32>
+}
+
+// -----
+
+func.func @reduce_scatter_op_non_scatter_dimension_mismatch(%input: tensor<2x64xf32>) -> tensor<1x32xf32> {
+  // expected-error @+1 {{output #0 type 'tensor<1x32xf32>' differs from input type 'tensor<2x64xf32>' in non-scatter dimension 1}}
+  %out = xtile.reduce_scatter(%input) scatter_dimension = 0 <replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>> : (tensor<2x64xf32>) -> (tensor<1x32xf32>) {
+  ^bb0(%arg0: tensor<f32>, %arg1: tensor<f32>):
+    %add = stablehlo.add %arg0, %arg1 : tensor<f32>
+    stablehlo.return %add : tensor<f32>
+  }
+  return %out : tensor<1x32xf32>
+}
+
+// -----
+
+func.func @reduce_scatter_op_non_scalar_reducer(%input: tensor<2x64xf32>) -> tensor<1x64xf32> {
+  // expected-error @+1 {{expects reduction computation arguments of type 'tensor<f32>', but got 'tensor<1xf32>'}}
+  %out = xtile.reduce_scatter(%input) scatter_dimension = 0 <replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>> : (tensor<2x64xf32>) -> (tensor<1x64xf32>) {
+  ^bb0(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>):
+    %add = stablehlo.add %arg0, %arg1 : tensor<1xf32>
+    stablehlo.return %add : tensor<1xf32>
+  }
+  return %out : tensor<1x64xf32>
+}
+
+// -----
+
 // expected-error @+1 {{layout has 0 dimensions, but shape has 1}}
 func.func @memref_layout_shape_size_mismatch(%arg0: memref<1024xf32, #xtile.layout<[]>>) {
   return
diff --git a/third_party/xla/xla/codegen/xtile/ir/transforms/verify_legal_xtile_ops.cc b/third_party/xla/xla/codegen/xtile/ir/transforms/verify_legal_xtile_ops.cc
index bdd7470fbd9..91a22486650 100644
--- a/third_party/xla/xla/codegen/xtile/ir/transforms/verify_legal_xtile_ops.cc
+++ b/third_party/xla/xla/codegen/xtile/ir/transforms/verify_legal_xtile_ops.cc
@@ -93,10 +93,9 @@ std::optional<absl::string_view> IsLegalStablehloOp(mlir::Operation* op) {
           mlir::stablehlo::MinOp, mlir::stablehlo::MulOp,
           mlir::stablehlo::NegOp, mlir::stablehlo::OrOp, mlir::stablehlo::PowOp,
           mlir::stablehlo::RealOp, mlir::stablehlo::ReduceOp,
-          mlir::stablehlo::ReduceScatterOp, mlir::stablehlo::RemOp,
-          mlir::stablehlo::ReshapeOp, mlir::stablehlo::ReturnOp,
-          mlir::stablehlo::SubtractOp, mlir::stablehlo::TransposeOp,
-          mlir::stablehlo::XorOp
+          mlir::stablehlo::RemOp, mlir::stablehlo::ReshapeOp,
+          mlir::stablehlo::ReturnOp, mlir::stablehlo::SubtractOp,
+          mlir::stablehlo::TransposeOp, mlir::stablehlo::XorOp
           // go/keep-sorted end
           >(op)) {
     return std::nullopt;
diff --git a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.cc b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.cc
index 3e6a2c34a28..a44488d1ca4 100644
--- a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.cc
+++ b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.cc
@@ -334,6 +334,72 @@ mlir::LogicalResult ScanOp::verify() {
   return mlir::success();
 }
 
+mlir::LogicalResult ReduceScatterOp::verify() {
+  mlir::ValueRange inputs = getInputs();
+  if (inputs.empty()) {
+    return emitOpError() << "expects at least one input";
+  }
+  if (getOutputs().size() != inputs.size()) {
+    return emitOpError() << "expects one output per input, but got "
+                         << inputs.size() << " inputs and "
+                         << getOutputs().size() << " outputs";
+  }
+
+  const int64_t scatter_dimension = getScatterDimension();
+  mlir::Type element_type =
+      mlir::cast<mlir::RankedTensorType>(inputs.front().getType())
+          .getElementType();
+  for (auto [index, input] : llvm::enumerate(inputs)) {
+    mlir::RankedTensorType input_type =
+        mlir::cast<mlir::RankedTensorType>(input.getType());
+    mlir::RankedTensorType output_type =
+        mlir::cast<mlir::RankedTensorType>(getOutputs()[index].getType());
+    if (input_type.getElementType() != element_type) {
+      return emitOpError() << "input #" << index << " element type "
+                           << input_type.getElementType()
+                           << " does not match the element type of input #0 "
+                           << element_type;
+    }
+    if (scatter_dimension < 0 || scatter_dimension >= input_type.getRank()) {
+      return emitOpError() << "scatter dimension " << scatter_dimension
+                           << " is out of range for input #" << index
+                           << " of rank " << input_type.getRank();
+    }
+    if (output_type.getRank() != input_type.getRank() ||
+        output_type.getElementType() != input_type.getElementType()) {
+      return emitOpError() << "output #" << index << " type " << output_type
+                           << " does not match the rank and element type of "
+                              "input type "
+                           << input_type;
+    }
+    for (int64_t dim = 0; dim < input_type.getRank(); ++dim) {
+      if (dim != scatter_dimension &&
+          output_type.getDimSize(dim) != input_type.getDimSize(dim)) {
+        return emitOpError() << "output #" << index << " type " << output_type
+                             << " differs from input type " << input_type
+                             << " in non-scatter dimension " << dim;
+      }
+    }
+  }
+
+  mlir::Block& body = getComputation().front();
+  if (body.getNumArguments() != 2) {
+    return emitOpError() << "expects a binary reduction computation, but the "
+                            "region has "
+                         << body.getNumArguments() << " arguments";
+  }
+  mlir::RankedTensorType scalar_type =
+      mlir::RankedTensorType::get({}, element_type);
+  for (mlir::BlockArgument argument : body.getArguments()) {
+    if (argument.getType() != scalar_type) {
+      return emitOpError() << "expects reduction computation arguments of "
+                              "type "
+                           << scalar_type << ", but got " << argument.getType();
+    }
+  }
+  return mlir::success();
+}
+
 static bool IsSupportedScaleElementType(mlir::Type type) {
   return mlir::isa<mlir::Float8E8M0FNUType, mlir::Float8E4M3FNType,
                    mlir::Float8E5M2Type>(type) ||
diff --git a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td
index 468a9b05823..e08bb1b7019 100644
--- a/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td
+++ b/third_party/xla/xla/codegen/xtile/ir/xtile_ops.td
@@ -358,6 +358,48 @@ def ScanOp : XTile_Op<"scan", [Pure, AttrSizedOperandSegments, AttrSizedResultSe
   let hasVerifier = 1;
 }
 
+def ReduceScatterOp : XTile_Op<"reduce_scatter", [Pure]> {
+  let summary = "variadic reduce-scatter operation";
+  let description = [{
+    Mirrors HLO's kReduceScatter operation with one or more operands.
+
+    Output `k` on every replica is the element-wise reduction, over all
+    replicas of its replica group, of that replica's shard of input `k` along
+    `scatter_dimension`. The same binary reduction `computation` (a region
+    with two 0-d tensor arguments, terminated by `stablehlo.return`) is applied
+    to all operands.
+
+    This op exists because `stablehlo.reduce_scatter` takes exactly one operand
+    and produces one result, whereas HLO's reduce-scatter (like
+    `stablehlo.all_gather` and `stablehlo.all_reduce`) is variadic. Extending
+    StableHLO is out of scope for XTile, so reduce-scatters are emitted as this
+    op.
+
+    Example: for inputs of type `tensor<2x64xf32>`, `scatter_dimension = 0` and
+    a replica group of size 2, the outputs have type `tensor<1x64xf32>`.
+  }];
+
+  let arguments = (
+    ins
+    Variadic<AnyRankedTensor>:$inputs,
+    I64Attr:$scatter_dimension,
+    I64ElementsAttr:$replica_groups
+  );
+
+  let results = (outs Variadic<AnyRankedTensor>:$outputs);
+
+  let regions = (region SizedRegion<1>:$computation);
+
+  let assemblyFormat = [{
+    `(` $inputs `)` `scatter_dimension` `=` $scatter_dimension
+    prop-dict attr-dict
+    `:` `(` type($inputs) `)` `->` `(` type($outputs) `)`
+    $computation
+  }];
+
+  let hasVerifier = 1;
+}
+
 //
 // DotScaled Op
 //

From ca4bce011f168600d8a855b8fd50f693c731159b Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Thu, 15 Oct 2026 01:45:52 -0700
Subject: [PATCH 0853/1000] PR #50775: Bump lit from 23.1.2 to 23.1.3

Imported from GitHub PR https://github.com/openxla/xla/pull/50775

Bumps [lit](https://github.com/llvm/llvm-project) from 23.1.2 to 23.1.3.
<details>
<summary>Release notes</summary>
<p><em>Sourced from <a href="https://github.com/llvm/llvm-project/releases">lit's releases</a>.</em></p>
<blockquote>
<h2>LLVM 23.1.3 Release</h2>
<ul>
<li>Linux x86_64: <a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/LLVM-23.1.3-Linux-X64.tar.xz">xz archive</a> (<a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/LLVM-23.1.3-Linux-X64.tar.xz.jsonl">signature</a>), <a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/LLVM-23.1.3-Linux-X64.tar.zst">zstd archive</a> (<a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/LLVM-23.1.3-Linux-X64.tar.zst.jsonl">signature</a>)</li>
<li>Linux Arm64: <a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/LLVM-23.1.3-Linux-ARM64.tar.xz">xz archive</a> (<a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/LLVM-23.1.3-Linux-ARM64.tar.xz.jsonl">signature</a>), <a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/LLVM-23.1.3-Linux-ARM64.tar.zst">zstd archive</a> (<a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/LLVM-23.1.3-Linux-ARM64.tar.zst.jsonl">signature</a>)</li>
</ul>
<!-- raw HTML omitted -->
<ul>
<li>macOS Apple Silicon (ARM64): <a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/LLVM-23.1.3-macOS-ARM64.tar.xz">xz archive</a> (<a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/LLVM-23.1.3-macOS-ARM64.tar.xz.jsonl">signature</a>), <a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/LLVM-23.1.3-macOS-ARM64.tar.zst">zstd archive</a> (<a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/LLVM-23.1.3-macOS-ARM64.tar.zst.jsonl">signature</a>)</li>
</ul>
<!-- raw HTML omitted -->
<ul>
<li>Windows x64 (64-bit): <a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/LLVM-23.1.3-win64.msi">installer</a> (<a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/LLVM-23.1.3-win64.msi.jsonl">signature</a>), <a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/clang+llvm-23.1.3-x86_64-pc-windows-msvc.tar.xz">xz archive</a> (<a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/clang+llvm-23.1.3-x86_64-pc-windows-msvc.tar.xz.jsonl">signature</a>), <a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/clang+llvm-23.1.3-x86_64-pc-windows-msvc.tar.zst">zstd archive</a> (<a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/clang+llvm-23.1.3-x86_64-pc-windows-msvc.tar.zst.jsonl">signature</a>)</li>
</ul>
<!-- raw HTML omitted -->
<ul>
<li>Windows on Arm (ARM64): <a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/LLVM-23.1.3-woa64.msi">installer</a> (<a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/LLVM-23.1.3-woa64.msi.jsonl">signature</a>), <a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/clang+llvm-23.1.3-aarch64-pc-windows-msvc.tar.xz">xz archive</a> (<a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/clang+llvm-23.1.3-aarch64-pc-windows-msvc.tar.xz.jsonl">signature</a>), <a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/clang+llvm-23.1.3-aarch64-pc-windows-msvc.tar.zst">zstd archive</a> (<a href="https://github.com/llvm/llvm-project/releases/download/llvmorg-23.1.3/clang+llvm-23.1.3-aarch64-pc-windows-msvc.tar.zst.jsonl">signature</a>)</li>
</ul>
<p>Download links for common platforms will appear above once builds have completed, if they are available. Check the full list of release packages at the bottom of this release page if you do not find a link above.</p>
<p>If you do not find a release package for your platform, you may be able to find a community built package on the LLVM Discourse forum thread for this release. Remember that these are built by volunteers and may not always be available. If you rely on a platform or configuration that is not one of the defaults, we suggest you use the binaries that your platform provides, or build your own release packages.</p>
<h2>Package Types</h2>
<p>Each platform has binary release packages. The file name starts with either <code>LLVM-</code> or <code>clang+llvm-</code> and ends with the platform's name. For example, <code>LLVM-23.1.3-Linux-ARM64.tar.xz</code> contains LLVM binaries for Arm64 Linux. Binary archive packages may be available as <code>.tar.xz</code> or <code>.tar.zst</code> files. The <code>.tar.zst</code> files contain the same package contents, but use zstd compression.</p>
<p>Except for Windows. Where <code>LLVM-*.msi</code> is an installer intended for using LLVM as a toolchain and the archive <code>clang+llvm-</code> contains the contents of the installer, plus libraries and tools not normally used in a toolchain. You most likely want the <code>LLVM-</code> installer, unless you are developing software which itself uses LLVM, in which case choose <code>clang+llvm-</code>.</p>
<p>In addition, source archives are available:</p>
<ul>
<li>To get all the <code>llvm-project</code> source code for this release, choose <code>llvm-project-23.1.3.src.tar.xz</code>.</li>
<li><code>test-suite-23.1.3.src.tar.xz</code> is an archive of the <a href="https://github.com/llvm/llvm-test-suite">LLVM Test Suite</a>) for this release.</li>
</ul>
<h2>Verifying Packages</h2>
<p>All packages come with a matching <code>.sig</code> and/or <code>.jsonl</code> file. You should use these to verify the integrity of the packages.</p>
<p>If it has a <code>.sig</code> file, it should have been signed by the release managers using GPG. Download the keys from the <a href="https://releases.llvm.org/release-keys.asc">LLVM website</a>, import them into your keyring and use them to verify the file:</p>
<pre><code>$ gpg --import release-keys.asc
$ gpg --verify &lt;package file name&gt;.sig &lt;package file name&gt;
</code></pre>
<p>If it has a <code>.jsonl</code> file, use <a href="https://cli.github.com/manual/gh_attestation_verify">gh</a> to verify the package:</p>
<pre><code>$ gh attestation verify --repo llvm/llvm-project &lt;package file name&gt;
(if you are able to connect to GitHub)
$ gh attestation verify --repo llvm/llvm-project &lt;package file name&gt; --bundle &lt;package file name&gt;.jsonl
(using attestation file on disk)
</code></pre>
</blockquote>
</details>
<details>
<summary>Commits</summary>
<ul>
<li><a href="https://github.com/llvm/llvm-project/commit/9efc50ac8a20e6f8b127a4204943e25eaa6f3313"><code>9efc50a</code></a> [Windows] Don't export clang::interp symbols for plugins (<a href="https://redirect.github.com/llvm/llvm-project/issues/221295">#221295</a>)</li>
<li><a href="https://github.com/llvm/llvm-project/commit/3273cdaa5ea1188a6dc630039f0be90d8d8084b6"><code>3273cda</code></a> [GlobalISel][AArch64] Do not remove dst type predicates from pattern children...</li>
<li><a href="https://github.com/llvm/llvm-project/commit/615ef6c76842c0046e94523b2042d3cfb2f8e803"><code>615ef6c</code></a> github-upload-release: Create draft releases (<a href="https://redirect.github.com/llvm/llvm-project/issues/209876">#209876</a>)</li>
<li><a href="https://github.com/llvm/llvm-project/commit/67f4a076a0978eb34564241c3eceb886bf6246f8"><code>67f4a07</code></a> Bump version to 23.1.3</li>
<li>See full diff in <a href="https://github.com/llvm/llvm-project/compare/llvmorg-23.1.2...llvmorg-23.1.3">compare view</a></li>
</ul>
</details>
<br />

[![Dependabot compatibility score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=lit&package-manager=pip&previous-version=23.1.2&new-version=23.1.3)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores)

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`.

[//]: # (dependabot-automerge-start)
[//]: # (dependabot-automerge-end)

---

<details>
<summary>Dependabot commands and options</summary>
<br />

You can trigger Dependabot actions by commenting on this PR:
- `@dependabot rebase` will rebase this PR
- `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it
- `@dependabot show <dependency name> ignore conditions` will show all of the ignore conditions of the specified dependency
- `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
- `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
- `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)

</details>
Copybara import of the project:

--
5f643f19411b7950f9f907bb2e39a302c85ebb17 by dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>:

Bump lit from 23.1.2 to 23.1.3

Bumps [lit](https://github.com/llvm/llvm-project) from 23.1.2 to 23.1.3.
- [Release notes](https://github.com/llvm/llvm-project/releases)
- [Commits](https://github.com/llvm/llvm-project/compare/llvmorg-23.1.2...llvmorg-23.1.3)

---
updated-dependencies:
- dependency-name: lit
  dependency-version: 23.1.3
  dependency-type: direct:production
  update-type: version-update:semver-patch
...

Signed-off-by: dependabot[bot] <support@github.com>

Merging this change closes #50775

PiperOrigin-RevId: 999961000
---
 third_party/xla/requirements_lock_3_11.txt | 6 +++---
 third_party/xla/requirements_lock_3_12.txt | 6 +++---
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/requirements_lock_3_11.txt b/third_party/xla/requirements_lock_3_11.txt
index e9e1f83393a..04256775b25 100644
--- a/third_party/xla/requirements_lock_3_11.txt
+++ b/third_party/xla/requirements_lock_3_11.txt
@@ -71,9 +71,9 @@ numpy==2.4.6 \
     --hash=sha256:f3a3570c4a2a16746ac2c31a7c7c7b0c186b95ce902e33db6f28094ed7387dda \
     --hash=sha256:f407cb6b8e9d6d8c626bc73c945db1706035af8fd632295547bf1c9e46d092d6 \
     --hash=sha256:f74a575920ab21fe304421a3fc28793d82e299cae9eccb37084e9fc7f3617c20
-lit==23.1.2 \
-    --hash=sha256:1a839a3f187ae74ce96d9b29a6b5b69671359e436b674e706f2ea5a343ca3cd6 \
-    --hash=sha256:d3b90c32edf5d4f2686d387ac58f35e2f5eb9c87e62ca00d360fb6b7bde080ce
+lit==23.1.3 \
+    --hash=sha256:559c274005e81bd1cf55c2c828fb6b9674fb108a9ed45fdc1feed2f54e148a5a \
+    --hash=sha256:a3d02402ddbc5ecd4df89fac6e64008e0fe6e4c7bafbb8fdb5dab5c50e1396ec
 ml-dtypes==0.6.0 \
     --hash=sha256:008382aeab529df5d3f00501ad9a7dcd64494d4b5b1971fc4c79019e6c1f5010 \
     --hash=sha256:03ce583adfce34ad33aa9e1fc7a8344dcf90ea776cc4ef0e5a48d4eae84e5d20 \
diff --git a/third_party/xla/requirements_lock_3_12.txt b/third_party/xla/requirements_lock_3_12.txt
index 1bd14c14498..de449717dc4 100644
--- a/third_party/xla/requirements_lock_3_12.txt
+++ b/third_party/xla/requirements_lock_3_12.txt
@@ -43,9 +43,9 @@ numpy==2.5.1 \
     --hash=sha256:f7119ebff1a9829e9f431a4f9d28e703023bb6b9fe7c8f724467dbfc27c94ab3 \
     --hash=sha256:f7d60026c0bdb1380e83bfa7a0419c4577ee4b9a08880afcb6dadeb74c649fa2 \
     --hash=sha256:f7feb014281029e628ba2d5a007407443b06e418b6fe451d1e2adcbc8eba0107
-lit==23.1.2 \
-    --hash=sha256:1a839a3f187ae74ce96d9b29a6b5b69671359e436b674e706f2ea5a343ca3cd6 \
-    --hash=sha256:d3b90c32edf5d4f2686d387ac58f35e2f5eb9c87e62ca00d360fb6b7bde080ce
+lit==23.1.3 \
+    --hash=sha256:559c274005e81bd1cf55c2c828fb6b9674fb108a9ed45fdc1feed2f54e148a5a \
+    --hash=sha256:a3d02402ddbc5ecd4df89fac6e64008e0fe6e4c7bafbb8fdb5dab5c50e1396ec
 ml-dtypes==0.6.0 \
     --hash=sha256:008382aeab529df5d3f00501ad9a7dcd64494d4b5b1971fc4c79019e6c1f5010 \
     --hash=sha256:03ce583adfce34ad33aa9e1fc7a8344dcf90ea776cc4ef0e5a48d4eae84e5d20 \

From 9f748cd5b7670f64af107a736dbc6b7e81912b13 Mon Sep 17 00:00:00 2001
From: Ilya Tikhonovskiy <loislo@google.com>
Date: Thu, 15 Oct 2026 01:51:32 -0700
Subject: [PATCH 0854/1000] [XLA:GPU] Restore scale swizzling and support
 checks for cuDNN kScaledDot fusions.

cuDNN block-scaled GEMM kernels require `F8_128x4` scale tensor reordering and 128x4 tile padding. When `kScaledDot` was wired into `CudnnFusionCompiler` and `CudnnBackend`, `CudnnBackend::ApplyConfigToCudnnFusion` did not invoke `CudnnScaledDotHelper::AddScaleSwizzle`, and `IsSupportedCudnnFusion` did not check Blackwell compute capability, the minimum cuDNN version, or `CudnnScaledDotHelper::IsSupported`. As a result, `kScaledDot` fusions lowered to `__cudnn$fusion` read unswizzled scales and produced incorrect outputs.

Restore the `CudnnScaledDotHelper::IsSupported` check in `IsSupportedCudnnFusion` and call `CudnnScaledDotHelper::AddScaleSwizzle` in `ApplyConfigToCudnnFusion` when the fusion contains a `kScaledDot` instruction. Also pass `MakeScaledDotArguments` to `RunAndCompareTwoModules` in the cuDNN scaled dot coverage test so that MXFP8 scales are non-trivial and NVFP4 scales are non-negative, and update the NVFP4 coverage expectation from `kFailed` to `kSupported`.

PiperOrigin-RevId: 999963919
---
 .../xla/xla/backends/gpu/autotuner/BUILD      |  1 +
 .../xla/xla/backends/gpu/autotuner/cudnn.cc   | 28 +++++++++++++++----
 .../xla/backends/gpu/autotuner/cudnn_test.cc  | 25 +++++++++++++++++
 .../gpu/codegen/cudnn_scaled_dot_test.cc      | 11 ++++++--
 4 files changed, 58 insertions(+), 7 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/autotuner/BUILD b/third_party/xla/xla/backends/gpu/autotuner/BUILD
index 6f74a53ae44..8003f53fc7d 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/BUILD
+++ b/third_party/xla/xla/backends/gpu/autotuner/BUILD
@@ -320,6 +320,7 @@ cc_library(
         "//xla/backends/autotuner:backends_proto_cc",
         "//xla/backends/autotuner:codegen_backend",
         "//xla/backends/gpu/target_config",
+        "//xla/backends/gpu/transforms:block_scaling_rewriter",
         "//xla/backends/gpu/transforms:cudnn_fusion_compiler",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/utils:hlo_query",
diff --git a/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc b/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc
index baf5f977d07..b23f39490c6 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc
@@ -28,6 +28,7 @@ limitations under the License.
 #include "xla/autotuning.pb.h"
 #include "xla/backends/autotuner/codegen_backend.h"
 #include "xla/backends/gpu/target_config/target_config.h"
+#include "xla/backends/gpu/transforms/block_scaling_rewriter.h"
 #include "xla/backends/gpu/transforms/cudnn_fusion_compiler.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_instruction.h"
@@ -148,11 +149,20 @@ bool IsSupportedCudnnFusion(const HloInstruction& instr,
   }
 
   if (hero->opcode() == HloOpcode::kConvolution ||
-      hero->opcode() == HloOpcode::kRaggedDot ||
-      hero->opcode() == HloOpcode::kScaledDot) {
+      hero->opcode() == HloOpcode::kRaggedDot) {
     return true;
   }
 
+  if (hero->opcode() == HloOpcode::kScaledDot) {
+    return target_config.device_description.cuda_compute_capability()
+               .IsAtLeastBlackwell() &&
+           se::dnn::VersionInfo(
+               target_config.device_description.dnn_version()) >=
+               kCudnnSupportsBlockScaledDot &&
+           CudnnScaledDotHelper::IsSupported(
+               Cast<HloScaledDotInstruction>(hero));
+  }
+
   if (hero->shape().element_type() == PrimitiveType::F64) {
     VLOG(1) << "cuDNN GEMM fusion does not support F64.";
     return false;
@@ -344,14 +354,22 @@ GetConvolutionCustomCallConfigs(const HloCustomCallInstruction* instr,
 
 absl::Status ApplyConfigToCudnnFusion(HloInstruction& instr,
                                       const CudnnBackendConfig& config) {
+  HloInstruction* fusion_instr = &instr;
+  auto* fusion = Cast<HloFusionInstruction>(&instr);
+  if (hlo_query::GetFirstInstructionWithOpcode(
+          *fusion->fused_instructions_computation(), HloOpcode::kScaledDot) !=
+      nullptr) {
+    ABSL_ASSIGN_OR_RETURN(fusion_instr,
+                     CudnnScaledDotHelper::AddScaleSwizzle(fusion));
+  }
   ABSL_ASSIGN_OR_RETURN(GpuBackendConfig gpu_config,
-                   instr.backend_config<GpuBackendConfig>());
+                   fusion_instr->backend_config<GpuBackendConfig>());
   FusionBackendConfig* backend_config =
       gpu_config.mutable_fusion_backend_config();
   backend_config->set_kind(kCuDnnFusionKind);
   backend_config->mutable_cudnn_fusion_config()->set_plan_id(config.algo_id());
-  ABSL_RETURN_IF_ERROR(instr.set_backend_config(std::move(gpu_config)));
-  instr.set_fusion_kind(HloInstruction::FusionKind::kCustom);
+  ABSL_RETURN_IF_ERROR(fusion_instr->set_backend_config(std::move(gpu_config)));
+  fusion_instr->set_fusion_kind(HloInstruction::FusionKind::kCustom);
   return absl::OkStatus();
 }
 
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 c47ee79b83d..665842650c4 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/cudnn_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/cudnn_test.cc
@@ -30,6 +30,7 @@ limitations under the License.
 #include "xla/backends/gpu/target_config/target_config.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/service/compiler.h"
 #include "xla/service/gpu/backend_configs.pb.h"
@@ -442,6 +443,30 @@ TEST_F(CudnnBackendTest, ApplyConfigToLoopFusionSetsCustomFusionKind) {
             1);
 }
 
+TEST_F(CudnnBackendTest, ApplyConfigToScaledDotGemmFusionAddsScaleSwizzle) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> hlo_module,
+                       ParseAndReturnVerifiedModule(kScaledDotGemmFusionHlo));
+  CudnnBackendConfig config;
+  config.set_algo_id(1);
+  HloInstruction* fusion_instr =
+      hlo_module->entry_computation()->root_instruction();
+  BackendConfig backend_config;
+  *backend_config.mutable_algorithm() = config;
+  ASSERT_OK(backend_->ApplyConfig(*fusion_instr, backend_config));
+  HloInstruction* root_instr =
+      hlo_module->entry_computation()->root_instruction();
+  ASSERT_EQ(root_instr->opcode(), HloOpcode::kFusion);
+  ASSERT_OK_AND_ASSIGN(GpuBackendConfig gpu_config,
+                       root_instr->backend_config<GpuBackendConfig>());
+  EXPECT_EQ(gpu_config.fusion_backend_config().kind(), kCuDnnFusionKind);
+  EXPECT_EQ(gpu_config.fusion_backend_config().cudnn_fusion_config().plan_id(),
+            1);
+  EXPECT_EQ(root_instr->operand(2)->opcode(), HloOpcode::kGetTupleElement);
+  EXPECT_EQ(root_instr->operand(2)->operand(0)->opcode(), HloOpcode::kCall);
+  EXPECT_EQ(root_instr->operand(3)->opcode(), HloOpcode::kGetTupleElement);
+  EXPECT_EQ(root_instr->operand(3)->operand(0)->opcode(), HloOpcode::kCall);
+}
+
 TEST_F(CudnnBackendTest, ApplyConfigToCudnnCustomCall) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> hlo_module,
                        ParseAndReturnVerifiedModule(kCudnnCustomCallHlo));
diff --git a/third_party/xla/xla/backends/gpu/codegen/cudnn_scaled_dot_test.cc b/third_party/xla/xla/backends/gpu/codegen/cudnn_scaled_dot_test.cc
index 8a9a98e060a..198f26d8f45 100644
--- a/third_party/xla/xla/backends/gpu/codegen/cudnn_scaled_dot_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/cudnn_scaled_dot_test.cc
@@ -290,7 +290,7 @@ std::vector<ScaledDotCoverageTestCase> GetCoverageTestCases() {
       {F4E2M1FN, F4E2M1FN, F8E4M3FN, 16, true, false,
        CuDnnSupportResult::kNotSupported},
       {F4E2M1FN, F4E2M1FN, F8E4M3FN, 16, true, true,
-       CuDnnSupportResult::kFailed},
+       CuDnnSupportResult::kSupported},
   };
 }
 
@@ -510,9 +510,16 @@ ENTRY %main {
         auto ref_optimized_or =
             GetOptimizedModule((*ref_module_or)->ToString());
         if (ref_optimized_or.ok()) {
+          auto args_or = MakeScaledDotArguments(optimized_module_or->get());
+          CHECK_OK(args_or.status());
+          std::vector<const Literal*> arg_ptrs;
+          arg_ptrs.reserve(args_or->size());
+          for (const Literal& arg : *args_or) {
+            arg_ptrs.push_back(&arg);
+          }
           auto run_result = RunAndCompareTwoModules(
               std::move(*optimized_module_or), std::move(*ref_optimized_or),
-              ErrorSpec{/*aabs=*/0.02, /*arel=*/0.05},
+              arg_ptrs, ErrorSpec{/*aabs=*/0.02, /*arel=*/0.05},
               /*run_hlo_passes=*/false);
           status = run_result ? "SUCCESS" : "FAILURE";
           if (!run_result) {

From f683e946bc1183029cf04dbe5c363f52696d3270 Mon Sep 17 00:00:00 2001
From: Anutosh Bhat <andersonbhat491@gmail.com>
Date: Thu, 15 Oct 2026 01:58:02 -0700
Subject: [PATCH 0855/1000] PR #50733: [XLA] Fix scalar index iteration on
 32-bit platforms

Imported from GitHub PR https://github.com/openxla/xla/pull/50733

I found this while porting JAX/XLA to WebAssembly for CPU execution in the browser, while investigating a compiler trap during scalar constant hashing.

For rank-zero shapes, `indices.size() - 1` underflows in unsigned `size_t` before being assigned to `int64_t`. On wasm32 this gives `4294967295` instead of `-1`, so `BumpIndices` enters the loop with an invalid dimension.

This casts the size before subtracting and adds a scalar case to the existing test.
Copybara import of the project:

--
d5839264569438b4769a371267073845ac3e966a by anutosh491 <andersonbhat491@gmail.com>:

[XLA] Fix scalar index iteration on 32-bit platforms

Merging this change closes #50733

PiperOrigin-RevId: 999967392
---
 third_party/xla/xla/index_util.h       | 3 ++-
 third_party/xla/xla/index_util_test.cc | 6 ++++++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/index_util.h b/third_party/xla/xla/index_util.h
index 9cf51389414..8a651271a57 100644
--- a/third_party/xla/xla/index_util.h
+++ b/third_party/xla/xla/index_util.h
@@ -160,7 +160,8 @@ class IndexUtil {
 
 inline bool ABSL_ATTRIBUTE_ALWAYS_INLINE
 IndexUtil::BumpIndices(const Shape& shape, absl::Span<int64_t> indices) {
-  for (int64_t dimno = indices.size() - 1; dimno >= 0; --dimno) {
+  for (int64_t dimno = static_cast<int64_t>(indices.size()) - 1; dimno >= 0;
+       --dimno) {
     int64_t limit = shape.dimensions(dimno);
     if (indices[dimno] + 1 < limit) {
       indices[dimno]++;
diff --git a/third_party/xla/xla/index_util_test.cc b/third_party/xla/xla/index_util_test.cc
index a828e7c7674..a8b94d29f27 100644
--- a/third_party/xla/xla/index_util_test.cc
+++ b/third_party/xla/xla/index_util_test.cc
@@ -144,6 +144,12 @@ TEST(IndexUtilTest, LinearToMultiToLinear) {
   }
 }
 
+TEST(IndexUtilTest, BumpIndicesScalar) {
+  const Shape shape = ShapeUtil::MakeShape(S32, {});
+  std::vector<int64_t> indices;
+  EXPECT_FALSE(IndexUtil::BumpIndices(shape, absl::MakeSpan(indices)));
+}
+
 TEST(IndexUtilTest, BumpIndices2x2) {
   auto shape = ShapeUtil::MakeShape(S32, {2, 2});
   std::vector<int64_t> indices = {0, 0};

From a4dacdbb318ed6c149a4e44a776a0b13b36ba2cf Mon Sep 17 00:00:00 2001
From: Vladimir Suvorov <suvorovv@google.com>
Date: Thu, 15 Oct 2026 02:02:59 -0700
Subject: [PATCH 0856/1000] Integrate Triton up to
 b575c0c38d428a609e25487db64051a5fa7bfab8

[List of integrated commits](https://github.com/triton-lang/triton/compare/0ccad974ee86942e704ac06dafdc0c3f214eb229...b575c0c38d428a609e25487db64051a5fa7bfab8)

PiperOrigin-RevId: 999970383
---
 .../third_party/triton/common/allocator.patch |  12 +
 .../triton/common/amd_range_analysis.patch    |   2 +-
 .../triton/common/cxx17_support.patch         | 252 ++++--
 .../triton/common/disable_amd_test.patch      |   2 +-
 .../triton/common/enable_peer_access.patch    |   4 +-
 .../third_party/triton/common/launcher.patch  |   2 +-
 .../triton/common/llvm_cl969980821.patch      |   2 +-
 .../triton/common/llvm_cl974500093.patch      |  32 +-
 .../triton/common/llvm_cl979179688.patch      |  11 +-
 .../triton/common/llvm_cl979526570.patch      |  19 +-
 .../triton/common/llvm_cl983115707.patch      | 804 +++++++++++++++++-
 .../triton/common/llvm_cl991819890.patch      |   6 +-
 .../triton/common/llvm_cl996124567.patch      |  10 +-
 .../triton/common/mig_limits_pytests.patch    |  10 +-
 .../third_party/triton/common/old_ptxas.patch |  61 ++
 .../third_party/triton/common/tmemfix.patch   |  37 +-
 .../triton/oss_only/build_files.patch         |   5 +-
 .../xla/third_party/triton/workspace.bzl      |   4 +-
 .../triton/compilation_pipeline_rocm.cc       |   4 +-
 .../transforms/tests/fp8_to_llvm_hopper.mlir  |   2 +-
 .../gpu/autotuning/autotune_cache_key.h       |   2 +-
 21 files changed, 1102 insertions(+), 181 deletions(-)

diff --git a/third_party/xla/third_party/triton/common/allocator.patch b/third_party/xla/third_party/triton/common/allocator.patch
index e1799c029cf..8bc9bde47b6 100644
--- a/third_party/xla/third_party/triton/common/allocator.patch
+++ b/third_party/xla/third_party/triton/common/allocator.patch
@@ -18,3 +18,15 @@
    res = cuDevicePrimaryCtxRetain(&ctx, device);
    if (res != CUDA_SUCCESS)
 
+--- a/python/triton/experimental/gsan/src/GSan.h
++++ b/python/triton/experimental/gsan/src/GSan.h
+@@ -35,7 +35,7 @@ struct LaunchState {
+ // A normal/write-once pair for each power-of-two granularity from 1 to 16.
+ // The low pool-index bit selects the mode; the remaining bits encode log2
+ // granularity. Six unused slots keep the reservation a power of two.
+-static constexpr size_t kPoolReserveSize = 1ull << 40;
++static constexpr size_t kPoolReserveSize = 1ull << 36;
+ static constexpr int kNumPools = 10;
+ static constexpr size_t kReserveSize = 16 * kPoolReserveSize;
+ // Default normal granularity. Write-once allocations default to one byte.
+
diff --git a/third_party/xla/third_party/triton/common/amd_range_analysis.patch b/third_party/xla/third_party/triton/common/amd_range_analysis.patch
index 819eadaf74f..13fd6ef4835 100644
--- a/third_party/xla/third_party/triton/common/amd_range_analysis.patch
+++ b/third_party/xla/third_party/triton/common/amd_range_analysis.patch
@@ -32,7 +32,7 @@ Lifecycle: Upstream this fix to triton-lang/triton and remove the patch in the n
  }
  
  bool isEmptyInitializedRange(ConstantIntRanges rv) {
-@@ -384,11 +387,16 @@ int64_t
+@@ -379,11 +382,16 @@ int64_t
  TritonIntegerRangeAnalysis::getTotalLoopTripCount(LoopLikeOpInterface loop) {
    SmallVector<LoopLikeOpInterface> loops{loop};
    getEnclosingLoops(*loop, loops);
diff --git a/third_party/xla/third_party/triton/common/cxx17_support.patch b/third_party/xla/third_party/triton/common/cxx17_support.patch
index 917968da5c9..93d3547d687 100644
--- a/third_party/xla/third_party/triton/common/cxx17_support.patch
+++ b/third_party/xla/third_party/triton/common/cxx17_support.patch
@@ -6,7 +6,7 @@ compiled with C++17.
 
 --- a/CMakeLists.txt
 +++ b/CMakeLists.txt
-@@ -9,7 +9,7 @@
+@@ -9,7 +9,7 @@ endif()
  project(triton CXX C)
  include(CTest)
  
@@ -18,7 +18,7 @@ compiled with C++17.
 
 --- a/include/triton/Analysis/Membar.h
 +++ b/include/triton/Analysis/Membar.h
-@@ -137,7 +137,10 @@
+@@ -137,7 +137,10 @@ struct ThreadIssuer {
    Kind kind = Unknown;
    Region *region = nullptr;
  
@@ -30,7 +30,7 @@ compiled with C++17.
  };
  
  // Empty summaries contain no operations. Nonempty summaries preserve an issuer
-@@ -151,7 +154,12 @@
+@@ -151,7 +154,12 @@ public:
    bool hasSameKnownIssuer(const IssuerSummary &other) const;
    void join(const IssuerSummary &other);
  
@@ -44,7 +44,7 @@ compiled with C++17.
  
  private:
    std::optional<ThreadIssuer> issuer;
-@@ -176,7 +184,14 @@
+@@ -176,7 +184,14 @@ struct BlockInfo {
      CompletionSync completion = CompletionSync::None;
      bool peerWaitNeedsSync = false;
  
@@ -60,7 +60,7 @@ compiled with C++17.
    } threadSync;
    struct ThreadDemandState {
      bool hasDemand = false;
-@@ -186,7 +201,15 @@
+@@ -186,7 +201,15 @@ struct BlockInfo {
      IssuerSummary publicationIssuers;
      bool mayNotifyPeer = false;
  
@@ -77,7 +77,7 @@ compiled with C++17.
    } threadDemands;
  
    BlockInfo() = default;
-@@ -275,7 +298,13 @@
+@@ -275,7 +298,13 @@ struct BlockInfo {
      threadDemands = {};
    }
  
@@ -92,7 +92,7 @@ compiled with C++17.
  
    /// Checks one pair of access maps with the same alias and operation filters.
    static bool isIntersected(const SliceMapT &lhsSlices,
-@@ -359,7 +388,13 @@
+@@ -359,7 +388,13 @@ struct MembarInfo {
      entryBlockInfo.transformSlices(transform);
    }
  
@@ -108,9 +108,25 @@ compiled with C++17.
  
  /// Classify the barriers that synchronize local memory accesses in `op`
 
+--- a/include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h
++++ b/include/triton/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVMBase.h
+@@ -150,8 +150,11 @@ protected:
+   ModuleAxisInfoAnalysis &axisAnalysisPass;
+ };
+ 
++template <typename DestOp> void noopCustomize(DestOp) {}
++
+ // Trivial case where we map elementwise to an existing LLVM operator
+-template <typename SourceOp, typename DestOp, auto Customize = [](DestOp) {}>
++template <typename SourceOp, typename DestOp,
++          void (*Customize)(DestOp) = noopCustomize<DestOp>>
+ struct ElementwiseOpConversion
+     : public ElementwiseOpConversionBase<
+           SourceOp, ElementwiseOpConversion<SourceOp, DestOp, Customize>> {
+
 --- a/include/triton/Tools/Sys/GetEnv.h
 +++ b/include/triton/Tools/Sys/GetEnv.h
-@@ -67,8 +67,8 @@
+@@ -67,8 +67,8 @@ inline const std::set<std::string> CACHE_NEUTRAL_ENV_VARS = {
  namespace tools {
  
  inline void assertIsRecognized(const std::string &env) {
@@ -122,9 +138,41 @@ compiled with C++17.
                               "Please add it to triton/Tools/Sys/GetEnv.h";
    assert((is_invalidating || is_neutral) && errmsg.c_str());
 
+--- a/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp
++++ b/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp
+@@ -689,22 +689,24 @@ void mlir::triton::populateElementwiseOpToLLVMPatterns(
+                                    benefit);
+ }
+ 
++static void setIntMinNotPoison(LLVM::AbsOp op) {
++  op.setIsIntMinPoison(false);
++}
++
+ void mlir::triton::populateElementwiseOpToLLVMPatterns(
+     LLVMTypeConverter &typeConverter, RewritePatternSet &patterns,
+     ModuleAxisInfoAnalysis &axisInfoAnalysis, PatternBenefit benefit) {
+ #define POPULATE_UNARY_OP(SRC_OP, DST_OP, ...)                                 \
+-  patterns.add<                                                                \
+-      ElementwiseOpConversion<SRC_OP, DST_OP __VA_OPT__(, ) __VA_ARGS__>>(     \
++  patterns.add<ElementwiseOpConversion<SRC_OP, DST_OP, ##__VA_ARGS__>>(        \
+       typeConverter, axisInfoAnalysis, benefit);
+ 
+   POPULATE_UNARY_OP(arith::TruncIOp, LLVM::TruncOp)
+   POPULATE_UNARY_OP(arith::ExtSIOp, LLVM::SExtOp)
+   POPULATE_UNARY_OP(arith::ExtUIOp, LLVM::ZExtOp)
+   POPULATE_UNARY_OP(arith::FPToUIOp, LLVM::FPToUIOp)
+   POPULATE_UNARY_OP(arith::UIToFPOp, LLVM::UIToFPOp)
+   POPULATE_UNARY_OP(arith::NegFOp, LLVM::FNegOp)
+-  POPULATE_UNARY_OP(math::AbsIOp, LLVM::AbsOp,
+-                    [](LLVM::AbsOp op) { op.setIsIntMinPoison(false); })
++  POPULATE_UNARY_OP(math::AbsIOp, LLVM::AbsOp, setIntMinNotPoison)
+   POPULATE_UNARY_OP(math::FloorOp, math::FloorOp)
+   POPULATE_UNARY_OP(math::CeilOp, math::CeilOp)
+   POPULATE_UNARY_OP(math::LogOp, math::LogOp)
+
 --- a/lib/Dialect/TritonGPU/Transforms/Utility.cpp
 +++ b/lib/Dialect/TritonGPU/Transforms/Utility.cpp
-@@ -318,7 +318,7 @@
+@@ -318,7 +318,7 @@ std::string GraphDumper::emitEdge(const std::string &srcId,
  
  std::string GraphDumper::emitValueNode(Value value) const {
    NodeInfo info = onValue(value);
@@ -133,7 +181,7 @@ compiled with C++17.
      std::string shapeStr = getShapeStr(value.getType());
      if (auto arg = mlir::dyn_cast<BlockArgument>(value))
        info["label"] =
-@@ -331,7 +331,7 @@
+@@ -331,7 +331,7 @@ std::string GraphDumper::emitValueNode(Value value) const {
  
  std::string GraphDumper::emitOperationNode(Operation *op) const {
    NodeInfo info = onOperation(op);
@@ -145,7 +193,7 @@ compiled with C++17.
 
 --- a/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionScheduling.cpp
 +++ b/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionScheduling.cpp
-@@ -340,7 +340,7 @@
+@@ -340,7 +340,7 @@ void deserializeManualPartitions(Operation *region, Graph *graph) {
              cast<DenseI32ArrayAttr>(op->getAttr(kPartitionAttrName))
                  .asArrayRef();
          for (auto id : partitionIds) {
@@ -154,7 +202,7 @@ compiled with C++17.
              auto partition = graph->addPartition();
              partition->addFlag(Flags::MANUAL);
              manual_partitions[id] = partition;
-@@ -1213,7 +1213,7 @@
+@@ -1213,7 +1213,7 @@ void duplicateCheapOps(Graph *graph, std::string funcName,
                  } else if (child->getPartition() == startPartition) {
                    // found a path, set all nodes on the path to the partition
                    node->addPartition(startPartition);
@@ -166,7 +214,7 @@ compiled with C++17.
 
 --- a/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionSchedulingUtility.cpp
 +++ b/lib/Dialect/TritonGPU/Transforms/WarpSpecialization/PartitionSchedulingUtility.cpp
-@@ -200,7 +200,7 @@
+@@ -203,7 +203,7 @@ void visualize(std::string key, std::string filename, std::string title,
        tools::getBoolEnv("TRITON_PARTITION_SCHEDULING_DUMP_LOOP_ONLY");
  
    static std::map<std::string, int> keys;
@@ -178,7 +226,7 @@ compiled with C++17.
 
 --- a/lib/Tools/LinearLayout.cpp
 +++ b/lib/Tools/LinearLayout.cpp
-@@ -1189,7 +1189,7 @@
+@@ -1189,7 +1189,7 @@ LinearLayout::getFreeVariableMasks() const {
    for (StringAttr dim : getInDimNames()) {
      int32_t mask = 0;
      for (int i = 0; i < getInDimSizeLog2(dim); i++, c++) {
@@ -196,7 +244,7 @@ compiled with C++17.
  #include <cmath>
  #include <cstddef>
  #include <cstdint>
-@@ -260,12 +259,23 @@
+@@ -260,12 +259,23 @@ constexpr uint32_t ToFloatBits(uint16_t h) {
    }
  }
  
@@ -225,7 +273,7 @@ compiled with C++17.
 
 --- a/python/src/llvm.cc
 +++ b/python/src/llvm.cc
-@@ -945,7 +945,7 @@
+@@ -980,7 +980,7 @@ void init_triton_llvm(py::module_ &m) {
        // Mark linked-in functions as internal because backends use external
        // linkage as a signifier of kernel functions.
        for (llvm::Function &fn : dstMod->functions()) {
@@ -237,7 +285,7 @@ compiled with C++17.
 
 --- a/third_party/amd/backend/codegen/CMakeLists.txt
 +++ b/third_party/amd/backend/codegen/CMakeLists.txt
-@@ -12,7 +12,7 @@
+@@ -12,7 +12,7 @@ set(TRITON_AMD_LLVM_REVISION "unknown" CACHE STRING
  set(TRITON_SOURCE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../../..")
  add_library(triton_amd_codegen SHARED amdgpu_codegen.cc
              "${TRITON_SOURCE_DIR}/lib/Tools/LLVMOptions.cpp")
@@ -249,7 +297,7 @@ compiled with C++17.
 
 --- a/third_party/amd/lib/TritonAMDGPUToLLVM/ConvertFpCastOpToLLVM.cpp
 +++ b/third_party/amd/lib/TritonAMDGPUToLLVM/ConvertFpCastOpToLLVM.cpp
-@@ -433,8 +433,9 @@
+@@ -472,8 +472,9 @@ SmallVector<Value> UnscaledF8x4ToFp32(ISAFamily family, Location loc,
  }
  
  // OCP Bf8/Fp8 -> Bf16
@@ -261,7 +309,7 @@ compiled with C++17.
  SmallVector<Value> OcpF8ToBf16SW(Location loc,
                                   ConversionPatternRewriter &rewriter,
                                   const SmallVector<Value> &v) {
-@@ -622,11 +623,13 @@
+@@ -661,11 +662,13 @@ SmallVector<Value> scalePk8DowncastToFp8(Location loc,
  // In saturation mode, inf and out-of-range numbers are converted to the largest
  // normal number, i.e. ±448. NaNs are converted to NaNs.
  // For UZ formats please check: https://onnx.ai/onnx/technical/float8.html
@@ -283,7 +331,7 @@ compiled with C++17.
 
 --- a/third_party/amd/lib/TritonAMDGPUToLLVM/ConvertWarpPipeline.cpp
 +++ b/third_party/amd/lib/TritonAMDGPUToLLVM/ConvertWarpPipeline.cpp
-@@ -129,7 +129,7 @@
+@@ -129,7 +129,7 @@ static LogicalResult validatePipelinedForBody(scf::ForOp forOp) {
              "non-warp-pipeline scf.execute_region inside pipelined_for body");
        ++numClusters;
      } else if (isWarpPipelineIgnorableBarrier(&op)) {
@@ -292,7 +340,7 @@ compiled with C++17.
          return op.emitError("multiple pre-existing barriers between pipeline "
                              "stages; insert a dummy stage instead");
        existingBarrierMap[numClusters] = &op;
-@@ -144,8 +144,8 @@
+@@ -144,8 +144,8 @@ static LogicalResult validatePipelinedForBody(scf::ForOp forOp) {
    if (numClusters < 2)
      return forOp.emitError(
          "pipelined_for body must contain at least two pipeline stages");
@@ -303,7 +351,7 @@ compiled with C++17.
      return forOp.emitError("pipelined_for body has both top-of-loop and "
                             "bottom-of-loop pre-existing barriers");
    return success();
-@@ -438,7 +438,7 @@
+@@ -438,7 +438,7 @@ private:
      // Normally, we don't expect a pipelined loop begins with a barrier
      // but sometimes required by memory prefetching pattern.
      auto bottomBar = existingBarrierMap.find(numClusters);
@@ -315,7 +363,7 @@ compiled with C++17.
 
 --- a/third_party/amd/lib/TritonAMDGPUTransforms/BlockPingpong.cpp
 +++ b/third_party/amd/lib/TritonAMDGPUTransforms/BlockPingpong.cpp
-@@ -1160,18 +1160,23 @@
+@@ -1160,18 +1160,23 @@ void Pingponger::getDotPingponged() {
    // Prune Memory operations that may be moved to only those involved in dot
    // computation. To understand the "cluster assumptions" we also estimate
    // the impact of any additional loads/stores.
@@ -354,7 +402,7 @@ compiled with C++17.
 
 --- a/third_party/amd/lib/TritonAMDGPUTransforms/ConvertToBufferOps.cpp
 +++ b/third_party/amd/lib/TritonAMDGPUTransforms/ConvertToBufferOps.cpp
-@@ -499,9 +499,10 @@
+@@ -499,9 +499,10 @@ private:
    bool analyzeSmallTensorOfst;
  };
  
@@ -371,7 +419,7 @@ compiled with C++17.
 
 --- a/third_party/amd/lib/TritonAMDGPUTransforms/LowerLoops.cpp
 +++ b/third_party/amd/lib/TritonAMDGPUTransforms/LowerLoops.cpp
-@@ -679,8 +679,8 @@
+@@ -679,8 +679,8 @@ LogicalResult initSchedule(int maxDist, Stages &stages, int numStages,
    // Make assignments
    Clusters clusterVec;
    schedule.clusters.clear();
@@ -385,7 +433,7 @@ compiled with C++17.
 
 --- a/third_party/amd/lib/TritonAMDGPUTransforms/ScheduleLoops.cpp
 +++ b/third_party/amd/lib/TritonAMDGPUTransforms/ScheduleLoops.cpp
-@@ -352,8 +352,8 @@
+@@ -352,8 +352,8 @@ void initSymbolicSchedule(int maxDist, Stages &stages, Clusters &clusters,
    stages[SCHED_ASYNC_WAIT] = stages[SCHED_LOCAL_LOAD];
  
    Clusters clusterVec;
@@ -396,7 +444,7 @@ compiled with C++17.
  
    // This is a symbolic cluster assignment. In this stage, we only focus on
    // global load and compute ops.
-@@ -491,8 +491,8 @@
+@@ -491,8 +491,8 @@ buildSchedule(scf::ForOp &forOp, int numStages, const LoadToInfoMap &loadToInfo,
    LDBG("Build ChainedDotSchedule");
    tt::CoarseSchedule schedule(numStages);
    ChainedDotClusters clusters;
@@ -419,7 +467,7 @@ compiled with C++17.
  
  namespace mlir {
  namespace triton {
-@@ -175,7 +175,7 @@
+@@ -175,7 +175,7 @@ public:
        currDesc.leadDimensionBaseOffset = 0;
      }
      int32_t smemByteOffsetb128 = smemByteOffsetb8 >> 4;
@@ -428,7 +476,7 @@ compiled with C++17.
      // Add the base address to the descriptor
      Value low = tb.add(tb.i32_val(uint32_t(descBits)), baseb128);
      Value high = tb.i32_val(descBits >> 32);
-@@ -288,7 +288,7 @@
+@@ -288,7 +288,7 @@ private:
            loc, "MMA instruction exceeds the shared-memory layout");
  
      auto desc =
@@ -440,21 +488,12 @@ compiled with C++17.
 
 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv2.cpp
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/DotOpToLLVM/MMAv2.cpp
-@@ -782,7 +782,7 @@
+@@ -680,7 +680,7 @@ LogicalResult convertMMA(triton::DotOp op, triton::DotOp::Adaptor adaptor,
  
    TensorCoreType mmaType = getMmaTypeDot(op, aTensorTy, bTensorTy, dTensorTy);
-   const auto &instrMap = isTuring ? mmaInstrPtxTuring : mmaInstrPtxAmpere;
--  if (!instrMap.contains(mmaType))
-+  if (!instrMap.count(mmaType))
-     return op.emitError(
-         "unsupported MMA instruction for the given operand/result types");
- 
-@@ -816,7 +816,7 @@
- 
-   TensorCoreType mmaType =
-       getMmaTypeDotScaled(op, aTensorTy, bTensorTy, dTensorTy);
--  if (!mmaInstrPtxScaled.contains(mmaType))
-+  if (!mmaInstrPtxScaled.count(mmaType))
+   if (mmaType == TensorCoreType::NOT_APPLICABLE ||
+-      (isTuring && !mmaInstrPtxTuring.contains(mmaType)))
++      (isTuring && !mmaInstrPtxTuring.count(mmaType)))
      return op.emitError(
          "unsupported MMA instruction for the given operand/result types");
  
@@ -471,7 +510,7 @@ compiled with C++17.
  
  using namespace mlir;
  using namespace mlir::triton;
-@@ -201,7 +201,7 @@
+@@ -201,7 +201,7 @@ static Value createInstDescriptor(ConversionPatternRewriter &rewriter,
    if (kSize == 64)
      desc.kSize = 1;
  
@@ -480,7 +519,7 @@ compiled with C++17.
  }
  
  static Value createScaleInstDescriptorFp8(ConversionPatternRewriter &rewriter,
-@@ -267,7 +267,7 @@
+@@ -267,7 +267,7 @@ static Value createScaleInstDescriptorFp8(ConversionPatternRewriter &rewriter,
      desc.BScaleFactor *= 2;
    }
  
@@ -489,7 +528,7 @@ compiled with C++17.
  }
  
  static Value createScaleInstDescriptorFp4(
-@@ -313,7 +313,7 @@
+@@ -313,7 +313,7 @@ static Value createScaleInstDescriptorFp4(
      // K96 selectors are byte offsets into three (block32) or six (block16)
      // consecutive scales, including across scale words.
      desc.kSizeLower = 1;
@@ -498,7 +537,7 @@ compiled with C++17.
    }
    if (kSize == 128) {
      desc.kSizeUpper = 1;
-@@ -345,7 +345,7 @@
+@@ -345,7 +345,7 @@ static Value createScaleInstDescriptorFp4(
             "MMAv5 with kind=mxf4nvf4 and .block16 only supports SFB_ID 0");
    }
  
@@ -508,9 +547,50 @@ compiled with C++17.
  
  //===----------------------------------------------------------------------===//
 
+--- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp
++++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp
+@@ -884,31 +884,32 @@ private:
+ 
+ } // namespace mlir::triton
+ 
++static void setApproxDivFastmathFlags(LLVM::FDivOp op) {
++  op.setFastmathFlags(LLVM::FastmathFlags::afn);
++}
++
+ void mlir::triton::NVIDIA::populateElementwiseOpToLLVMPatterns(
+     LLVMTypeConverter &typeConverter, RewritePatternSet &patterns,
+     ModuleAxisInfoAnalysis &axisInfoAnalysis, int computeCapability,
+     const TargetInfo &targetInfo, PatternBenefit benefit) {
+   using namespace mlir::triton::gpu;
+ 
+   mlir::triton::populateElementwiseOpToLLVMPatterns(typeConverter, patterns,
+                                                     axisInfoAnalysis, benefit);
+ 
+   patterns.add<MulhiUIOpConversion>(typeConverter, axisInfoAnalysis,
+                                     benefit.getBenefit() + 1);
+ 
+ #define POPULATE_OP(SRC_OP, DST_OP, ...)                                       \
+-  patterns.add<                                                                \
+-      ElementwiseOpConversion<SRC_OP, DST_OP __VA_OPT__(, ) __VA_ARGS__>>(     \
++  patterns.add<ElementwiseOpConversion<SRC_OP, DST_OP, ##__VA_ARGS__>>(        \
+       typeConverter, axisInfoAnalysis, benefit)
+ 
+   POPULATE_OP(arith::SubFOp, LLVM::FSubOp);
+   POPULATE_OP(arith::AddFOp, LLVM::FAddOp);
+   POPULATE_OP(arith::MulFOp, LLVM::FMulOp);
+   POPULATE_OP(triton::PreciseDivFOp, LLVM::FDivOp);
+   POPULATE_OP(triton::PreciseSqrtOp, LLVM::SqrtOp);
+-  POPULATE_OP(triton::ApproxDivFOp, LLVM::FDivOp, [](LLVM::FDivOp op) {
+-    op.setFastmathFlags(LLVM::FastmathFlags::afn);
+-  });
++  POPULATE_OP(triton::ApproxDivFOp, LLVM::FDivOp, setApproxDivFastmathFlags);
+ 
+   POPULATE_OP(arith::ExtFOp, LLVM::FPExtOp);
+   POPULATE_OP(arith::TruncFOp, LLVM::FPTruncOp);
+
 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
-@@ -841,7 +841,7 @@
+@@ -841,7 +841,7 @@ public:
               triton::nvgpu::MemSyncScope::SYSTEM}};
      const bool doPTXLDPromotion = !useRed && isPromotableToNVPTXLD(op) &&
                                    vec == 1 && packed == 1 &&
@@ -522,7 +602,7 @@ compiled with C++17.
 
 --- a/third_party/proton/common/lib/TraceDataIO/CircularLayoutParser.cpp
 +++ b/third_party/proton/common/lib/TraceDataIO/CircularLayoutParser.cpp
-@@ -110,14 +110,14 @@
+@@ -110,14 +110,14 @@ void CircularLayoutParser::parseSegment(
          trace.asyncEvents.push_back(entry);
          continue;
        }
@@ -539,7 +619,7 @@ compiled with C++17.
          prevState = scopeState[entry->scopeId];
  
        if (entry->isStart) {
-@@ -174,7 +174,7 @@
+@@ -174,7 +174,7 @@ void CircularLayoutParser::pairAsyncEvents(
      if (endpoint.entry->isStart) {
        // A warp cannot have two outstanding instances of the same static event
        // because a later end would not identify which instance it completes.
@@ -551,7 +631,7 @@ compiled with C++17.
 
 --- a/third_party/proton/common/lib/TraceDataIO/TraceWriter.cpp
 +++ b/third_party/proton/common/lib/TraceDataIO/TraceWriter.cpp
-@@ -91,7 +91,7 @@
+@@ -91,7 +91,7 @@ void populateTraceInfo(std::shared_ptr<CircularLayoutParserResult> result,
  
      // Group block traces by proc id
      int procId = bt.procId;
@@ -560,7 +640,7 @@ compiled with C++17.
        procToBlockTraces[procId] = {};
      }
      procToBlockTraces[procId].push_back(&bt);
-@@ -203,12 +203,12 @@
+@@ -203,12 +203,12 @@ void writeProfileEvents(json &object, const BlockTrace &blockTrace,
      for (const auto &event : trace.profileEvents) {
        int lineId = eventLineIds.at(event.first.get());
        int scopeId = event.first->scopeId;
@@ -575,7 +655,7 @@ compiled with C++17.
                               ? "scope_" + std::to_string(scopeId)
                               : metadata.scopeName.at(scopeId);
        std::string tid = "warp " + std::to_string(warpId) + " (line " +
-@@ -250,12 +250,12 @@
+@@ -250,12 +250,12 @@ void writeAsyncEvents(json &object, const BlockTrace &blockTrace,
    constexpr double freq = 1000.0;
    for (const auto &link : blockTrace.asyncLinks) {
      int scopeId = link.first.entry->scopeId;
@@ -593,17 +673,15 @@ compiled with C++17.
 
 --- a/third_party/proton/csrc/include/Session/Session.h
 +++ b/third_party/proton/csrc/include/Session/Session.h
-@@ -151,10 +151,10 @@
+@@ -151,10 +151,10 @@ private:
    size_t getSessionId(const std::string &path) { return sessionPaths[path]; }
  
    bool hasSession(const std::string &path) {
 -    return sessionPaths.contains(path);
--  }
--
--  bool hasSession(size_t sessionId) { return sessions.contains(sessionId); }
 +    return sessionPaths.count(path);
-+  }
-+
+   }
+ 
+-  bool hasSession(size_t sessionId) { return sessions.contains(sessionId); }
 +  bool hasSession(size_t sessionId) { return sessions.count(sessionId); }
  
    void removeSession(size_t sessionId);
@@ -611,7 +689,7 @@ compiled with C++17.
 
 --- a/third_party/proton/csrc/include/Utility/Map.h
 +++ b/third_party/proton/csrc/include/Utility/Map.h
-@@ -68,7 +68,7 @@
+@@ -68,7 +68,7 @@ public:
  
    bool contains(const Key &key) const {
      std::shared_lock<std::shared_mutex> lock(mutex);
@@ -623,7 +701,7 @@ compiled with C++17.
 
 --- a/third_party/proton/csrc/include/Utility/Set.h
 +++ b/third_party/proton/csrc/include/Utility/Set.h
-@@ -19,7 +19,7 @@
+@@ -19,7 +19,7 @@ public:
  
    bool contains(const Key &key) const {
      std::shared_lock<std::shared_mutex> lock(mutex);
@@ -635,7 +713,7 @@ compiled with C++17.
 
 --- a/third_party/proton/csrc/lib/Data/Dump/HatchetMsgPack.cpp
 +++ b/third_party/proton/csrc/lib/Data/Dump/HatchetMsgPack.cpp
-@@ -102,14 +102,14 @@
+@@ -103,14 +103,14 @@ countMetricEntries(const std::map<MetricKind, std::unique_ptr<Metric>> &metrics,
    }
    if (isRoot) {
      if (metricSummary.hasKernelMetric &&
@@ -653,7 +731,7 @@ compiled with C++17.
        metricEntries += kCycleInclusiveCount;
      }
    }
-@@ -202,14 +202,14 @@
+@@ -203,14 +203,14 @@ void packMetrics(MsgPackWriter &writer,
    }
    if (isRoot) {
      if (metricSummary.hasKernelMetric &&
@@ -670,7 +748,7 @@ compiled with C++17.
        PCSamplingMetric pcSamplingMetric;
        for (size_t i = 0; i < PCSamplingMetric::Count; i++) {
          const auto valueName = pcSamplingMetric.getValueName(i);
-@@ -217,7 +217,7 @@
+@@ -218,7 +218,7 @@ void packMetrics(MsgPackWriter &writer,
          writer.packUInt(0);
        }
      }
@@ -679,7 +757,7 @@ compiled with C++17.
        writer.packStr(CycleMetric::getValueName(CycleMetric::Duration));
        writer.packUInt(0);
        writer.packStr(
-@@ -348,7 +348,7 @@
+@@ -348,7 +348,7 @@ TreeData::buildHatchetMsgPack(TreeData::Tree *tree,
      linkedChildren.reserve(virtualNode.children.size());
      for (const auto &child : virtualNode.children) {
        const auto &childNode = virtualTree->getNode(child.id);
@@ -691,7 +769,7 @@ compiled with C++17.
 
 --- a/third_party/proton/csrc/lib/Data/Dump/TraceDataDump.cpp
 +++ b/third_party/proton/csrc/lib/Data/Dump/TraceDataDump.cpp
-@@ -215,7 +215,7 @@
+@@ -215,7 +215,7 @@ convertToTimelineTrace(std::vector<CycleEvent> &cycleEvents) {
            }
  
            auto scopeName = event.contexts.back().name;
@@ -730,7 +808,7 @@ compiled with C++17.
  
  std::map<uint64_t, MetricBuffer::MetricDescriptor>
      MetricBuffer::metricDescriptors;
-@@ -101,22 +115,22 @@
+@@ -101,22 +115,22 @@ collectTensorMetrics(Runtime *runtime,
                                     stream);
      runtime->synchronizeStream(stream);
      if (tensorMetric.typeIndex == variant_index_v<double, MetricValueType>) {
@@ -757,7 +835,7 @@ compiled with C++17.
        }
        tensorMetricsHost[name] = std::move(values);
      } else {
-@@ -166,7 +180,7 @@
+@@ -166,7 +180,7 @@ void MetricBuffer::queue(uint64_t seqId, MetricValueType scalarMetric,
          } else {
            static_assert(sizeof(T) == sizeof(uint64_t),
                          "MetricValueType alternative must be 8 bytes");
@@ -766,7 +844,7 @@ compiled with C++17.
          }
        },
        scalarMetric);
-@@ -204,7 +218,7 @@
+@@ -204,7 +218,7 @@ void MetricBuffer::synchronize(DeviceBuffer &buffer) {
  MetricBuffer::DeviceBuffer &MetricBuffer::getOrCreateBuffer() {
    std::lock_guard<std::mutex> lock(bufferMutex);
    auto device = runtime->getDevice();
@@ -778,7 +856,7 @@ compiled with C++17.
 
 --- a/third_party/proton/csrc/lib/Data/TraceData.cpp
 +++ b/third_party/proton/csrc/lib/Data/TraceData.cpp
-@@ -53,7 +53,7 @@
+@@ -53,7 +53,7 @@ public:
      void addChild(const Context &context, size_t id) { children[context] = id; }
  
      bool hasChild(const Context &context) const {
@@ -787,7 +865,7 @@ compiled with C++17.
      }
  
      size_t getChild(const Context &context) const {
-@@ -152,7 +152,7 @@
+@@ -152,7 +152,7 @@ public:
      return nextEventId++;
    }
  
@@ -796,7 +874,7 @@ compiled with C++17.
  
    Event &getEvent(size_t eventId) {
      auto it = traceEvents.find(eventId);
-@@ -253,7 +253,7 @@
+@@ -253,7 +253,7 @@ void TraceData::addMetrics(
    auto &event = currentTrace->getEvent(eventId);
    auto &flexibleMetrics = event.metricSet.flexibleMetrics;
    for (auto [metricName, metricValue] : metrics) {
@@ -808,7 +886,7 @@ compiled with C++17.
 
 --- a/third_party/proton/csrc/lib/Profiler/Cupti/CuptiPCSampling.cpp
 +++ b/third_party/proton/csrc/lib/Profiler/Cupti/CuptiPCSampling.cpp
-@@ -365,7 +365,7 @@
+@@ -365,7 +365,7 @@ void CuptiPCSampling::processPCSamplingData(ConfigureData *configureData,
        auto *cubinData = getCubinData(pcData->cubinCrc);
        auto key =
            CubinData::LineInfoKey{pcData->functionIndex, pcData->pcOffset};
@@ -817,7 +895,7 @@ compiled with C++17.
          auto [lineNumber, fileName, dirName] =
              getSassToSourceCorrelation(pcData->functionName, pcData->pcOffset,
                                         cubinData->cubin, cubinData->cubinSize);
-@@ -376,7 +376,7 @@
+@@ -376,7 +376,7 @@ void CuptiPCSampling::processPCSamplingData(ConfigureData *configureData,
        auto &lineInfo = cubinData->lineInfo[key];
        for (size_t j = 0; j < pcData->stallReasonCount; ++j) {
          auto *stallReason = &pcData->stallReason[j];
@@ -826,7 +904,7 @@ compiled with C++17.
                  stallReason->pcSamplingStallReasonIndex))
            throw makeOutOfRange("Invalid stall reason index");
          for (const auto &[data, baseEntry] : dataToEntry) {
-@@ -392,7 +392,7 @@
+@@ -392,7 +392,7 @@ void CuptiPCSampling::processPCSamplingData(ConfigureData *configureData,
                    [stallReason->pcSamplingStallReasonIndex]);
            auto samples = stallReason->samples;
            auto stalledSamples =
@@ -867,7 +945,7 @@ compiled with C++17.
  
  void GraphState::recordNode(uint64_t nodeId, const std::string &name,
                              std::optional<MetricNodeState> metricNodeState,
-@@ -69,7 +84,7 @@
+@@ -69,7 +84,7 @@ void GraphState::recordNode(uint64_t nodeId, const std::string &name,
  void GraphState::buildLaunchEntries(const DataToEntryMap &dataToEntry,
                                      DataToEntryMap &dataToGraphEntry) const {
    for (const auto &[data, entry] : dataToEntry) {
@@ -876,7 +954,7 @@ compiled with C++17.
        // This data object was not enabled during graph capture.
        continue;
      dataToGraphEntry.insert({data, entry});
-@@ -153,11 +168,11 @@
+@@ -153,11 +168,11 @@ void emitMetricRecords(MetricBuffer &metricBuffer, uint64_t *hostBasePtr,
        break;
      }
      case variant_index_v<int64_t, MetricValueType>: {
@@ -890,7 +968,7 @@ compiled with C++17.
        break;
      }
      case variant_index_v<std::vector<uint64_t>, MetricValueType>: {
-@@ -171,7 +186,7 @@
+@@ -171,7 +186,7 @@ void emitMetricRecords(MetricBuffer &metricBuffer, uint64_t *hostBasePtr,
      case variant_index_v<std::vector<int64_t>, MetricValueType>: {
        std::vector<int64_t> values(metricDesc.size);
        for (size_t j = 0; j < metricDesc.size; ++j) {
@@ -899,7 +977,7 @@ compiled with C++17.
        }
        metricValueVariant = std::move(values);
        break;
-@@ -179,7 +194,7 @@
+@@ -179,7 +194,7 @@ void emitMetricRecords(MetricBuffer &metricBuffer, uint64_t *hostBasePtr,
      case variant_index_v<std::vector<double>, MetricValueType>: {
        std::vector<double> values(metricDesc.size);
        for (size_t j = 0; j < metricDesc.size; ++j) {
@@ -911,7 +989,7 @@ compiled with C++17.
 
 --- a/third_party/proton/csrc/lib/Profiler/Instrumentation/InstrumentationProfiler.cpp
 +++ b/third_party/proton/csrc/lib/Profiler/Instrumentation/InstrumentationProfiler.cpp
-@@ -90,7 +90,7 @@
+@@ -90,7 +90,7 @@ std::vector<uint32_t>
  getUnitIdVector(const std::map<std::string, std::string> &modeOptions,
                  size_t totalUnits) {
    std::vector<uint32_t> unitIdVector;
@@ -920,7 +998,7 @@ compiled with C++17.
      auto &samplingOption = modeOptions.at("sampling_options");
      auto unitIds = proton::split(samplingOption, ",");
      for (auto uintId : unitIds) {
-@@ -118,7 +118,7 @@
+@@ -118,7 +118,7 @@ InstrumentationProfiler::getParserConfig(uint64_t functionId,
    auto config = std::make_shared<CircularLayoutParserConfig>();
    config->scratchMemSize =
        functionMetadata.at(functionId).getScratchMemorySize();
@@ -929,7 +1007,7 @@ compiled with C++17.
        modeOptions.at("granularity") != "GRANULARITY.WARP") {
      throw makeInvalidArgument("Only warp granularity is supported for now");
    }
-@@ -146,7 +146,7 @@
+@@ -146,7 +146,7 @@ void InstrumentationProfiler::initFunctionMetadata(
      const std::vector<std::pair<size_t, std::string>> &scopeIdPairs,
      const std::vector<std::pair<size_t, size_t>> &scopeIdParentPairs,
      const std::string &metadataPath) {
@@ -938,7 +1016,7 @@ compiled with C++17.
      throw makeInvalidArgument(
          "Duplicate function id: " + std::to_string(functionId) +
          " for function " + functionName);
-@@ -155,7 +155,7 @@
+@@ -155,7 +155,7 @@ void InstrumentationProfiler::initFunctionMetadata(
    for (auto &pair : scopeIdPairs) {
      auto scopeId = pair.first;
      auto scopeName = pair.second;
@@ -947,7 +1025,7 @@ compiled with C++17.
        throw makeInvalidArgument(
            "Duplicate scope id: " + std::to_string(scopeId) + " for function " +
            functionName);
-@@ -173,7 +173,7 @@
+@@ -173,7 +173,7 @@ void InstrumentationProfiler::initFunctionMetadata(
      std::vector<Context> reversedContexts;
      reversedContexts.emplace_back(name);
      auto currentId = scopeId;
@@ -956,7 +1034,7 @@ compiled with C++17.
        auto parentId = scopeIdParentMap[currentId];
        auto parentName = functionScopeIdNames[functionId].at(parentId);
        reversedContexts.emplace_back(parentName);
-@@ -239,7 +239,7 @@
+@@ -239,7 +239,7 @@ void InstrumentationProfiler::exitInstrumentedOp(uint64_t streamId,
    }
  
    int64_t timeShiftCost = 0;
@@ -968,7 +1046,7 @@ compiled with C++17.
 
 --- a/third_party/proton/csrc/lib/Profiler/RocprofSDK/RocprofSDKPCSamplingSourceLocations.cpp
 +++ b/third_party/proton/csrc/lib/Profiler/RocprofSDK/RocprofSDKPCSamplingSourceLocations.cpp
-@@ -229,8 +229,8 @@
+@@ -229,8 +229,8 @@ void RocprofSDKPCSampling::clearSourceLocationCacheLocked(
  void RocprofSDKPCSampling::tryReleaseCodeObject(uint64_t codeObjectId) {
    metadataState.withLock([&](MetadataState &metadata) {
      samplingState.withLock([&](SamplingState &sampling) {
@@ -982,7 +1060,7 @@ compiled with C++17.
 
 --- a/third_party/proton/csrc/lib/Session/Session.cpp
 +++ b/third_party/proton/csrc/lib/Session/Session.cpp
-@@ -66,7 +66,7 @@
+@@ -66,7 +66,7 @@ makeContextSource(const std::string &contextSourceName) {
  void throwIfSessionNotInitialized(
      const std::map<size_t, std::unique_ptr<Session>> &sessions,
      size_t sessionId) {
@@ -1001,7 +1079,7 @@ compiled with C++17.
  #include <cstring>
  #include <limits>
  #include <type_traits>
-@@ -78,7 +77,10 @@
+@@ -78,7 +77,10 @@ void MsgPackWriter::packInt(int64_t value) {
  
  void MsgPackWriter::packDouble(double value) {
    out.push_back(0xcb);
diff --git a/third_party/xla/third_party/triton/common/disable_amd_test.patch b/third_party/xla/third_party/triton/common/disable_amd_test.patch
index f42f4eb0b1c..07f61831969 100644
--- a/third_party/xla/third_party/triton/common/disable_amd_test.patch
+++ b/third_party/xla/third_party/triton/common/disable_amd_test.patch
@@ -4,7 +4,7 @@ would not be able to get run on presubmit anyway.
 
 --- a/python/test/unit/language/test_compile_only.py
 +++ b/python/test/unit/language/test_compile_only.py
-@@ -443,6 +443,8 @@ def test_signature_ordering():
+@@ -472,6 +472,8 @@ def test_signature_ordering():
      triton.compile(src=src, target=target)
  
  
diff --git a/third_party/xla/third_party/triton/common/enable_peer_access.patch b/third_party/xla/third_party/triton/common/enable_peer_access.patch
index f624c016e33..c7188f0b5e9 100644
--- a/third_party/xla/third_party/triton/common/enable_peer_access.patch
+++ b/third_party/xla/third_party/triton/common/enable_peer_access.patch
@@ -33,7 +33,7 @@ new file mode 100644
 
 --- a/third_party/nvidia/backend/driver.c
 +++ b/third_party/nvidia/backend/driver.c
-@@ -984,6 +984,37 @@ cleanup:
+@@ -991,6 +991,37 @@ cleanup:
    return NULL;
  }
  
@@ -71,7 +71,7 @@ new file mode 100644
  static void _launch(int gridX, int gridY, int gridZ, int num_warps,
                      int num_ctas, int launch_cooperative_grid, int launch_pdl,
                      int shared_memory, CUstream stream, CUfunction function,
-@@ -1095,6 +1126,9 @@ bool extractPointer(void *ptr, PyObject *obj) {
+@@ -1102,6 +1133,9 @@ bool extractPointer(void *ptr, PyObject *obj) {
    if (*dev_ptr == 0) {
      return true; // valid nullptr
    }
diff --git a/third_party/xla/third_party/triton/common/launcher.patch b/third_party/xla/third_party/triton/common/launcher.patch
index faf100557c7..97f2230a0a6 100644
--- a/third_party/xla/third_party/triton/common/launcher.patch
+++ b/third_party/xla/third_party/triton/common/launcher.patch
@@ -4,7 +4,7 @@
 
 --- a/third_party/nvidia/backend/driver.py
 +++ b/third_party/nvidia/backend/driver.py
-@@ -99,13 +99,8 @@ class CudaUtils(object):
+@@ -106,13 +106,8 @@ class CudaUtils(object):
          return cls.instance
  
      def __init__(self):
diff --git a/third_party/xla/third_party/triton/common/llvm_cl969980821.patch b/third_party/xla/third_party/triton/common/llvm_cl969980821.patch
index 23c2ed63a72..21543e81e9a 100644
--- a/third_party/xla/third_party/triton/common/llvm_cl969980821.patch
+++ b/third_party/xla/third_party/triton/common/llvm_cl969980821.patch
@@ -1,6 +1,6 @@
 --- a/python/src/llvm.cc
 +++ b/python/src/llvm.cc
-@@ -704,7 +704,7 @@ void init_triton_llvm(py::module_ &m) {
+@@ -739,7 +739,7 @@ void init_triton_llvm(py::module_ &m) {
            // registry name "vector-combine".
            const StringRef kVectorCombinePassName = "VectorCombinePass";
            passInstrCb.registerShouldRunOptionalPassCallback(
diff --git a/third_party/xla/third_party/triton/common/llvm_cl974500093.patch b/third_party/xla/third_party/triton/common/llvm_cl974500093.patch
index b9963ad02c2..6839f1466c0 100644
--- a/third_party/xla/third_party/triton/common/llvm_cl974500093.patch
+++ b/third_party/xla/third_party/triton/common/llvm_cl974500093.patch
@@ -58,7 +58,7 @@
      // CHECK-NOT: nvvm.ldmatrix
      %AA_DOT = ttg.local_load %AA : !ttg.memdesc<16x16xf16, #shared0, #smem> -> tensor<16x16xf16, #dot_operand_a>
      %BB_DOT = ttg.local_load %BB : !ttg.memdesc<16x16xf16, #shared0, #smem> -> tensor<16x16xf16, #dot_operand_b>
-@@ -1424,8 +1424,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+@@ -1421,8 +1421,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
    tt.func @convert_dot_ldmatrix_swizzle(%A: tensor<16x16xf16, #blocked0>, %B: tensor<16x16xf16, #blocked0>) {
      %AA = ttg.local_alloc %A : (tensor<16x16xf16, #blocked0>) -> !ttg.memdesc<16x16xf16, #shared0, #smem>
      %BB = ttg.local_alloc %B : (tensor<16x16xf16, #blocked0>) -> !ttg.memdesc<16x16xf16, #shared0, #smem>
@@ -69,7 +69,7 @@
      // CHECK-NOT: nvvm.ldmatrix
      %AA_DOT = ttg.local_load %AA : !ttg.memdesc<16x16xf16, #shared0, #smem> -> tensor<16x16xf16, #dot_operand_a>
      %BB_DOT = ttg.local_load %BB : !ttg.memdesc<16x16xf16, #shared0, #smem> -> tensor<16x16xf16, #dot_operand_b>
-@@ -1506,7 +1506,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+@@ -1497,7 +1497,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
    tt.func @convert_dot_fp8(%A: tensor<16x16xf8E5M2, #blocked0>, %B: tensor<16x16xf8E5M2, #blocked0>) {
      %AA = ttg.local_alloc %A : (tensor<16x16xf8E5M2, #blocked0>) -> !ttg.memdesc<16x16xf8E5M2, #shared0, #smem>
      %BB = ttg.local_alloc %B : (tensor<16x16xf8E5M2, #blocked0>) -> !ttg.memdesc<16x16xf8E5M2, #shared0, #smem>
@@ -81,7 +81,7 @@
 
 --- a/test/Conversion/tritongpu_to_llvm_blackwell.mlir
 +++ b/test/Conversion/tritongpu_to_llvm_blackwell.mlir
-@@ -1003,10 +1003,10 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.thr
+@@ -1012,10 +1012,10 @@ tt.func @load_store_x1_unpacked(%arg0: !ttg.memdesc<128x2xf16, #tmem_x1_unpacked
  
  // CHECK-LABEL: max_reduction
  //       CHECK-DAG:  %[[M:.+]] = llvm.mlir.constant(-1 : i32) : i32
@@ -94,7 +94,7 @@
  //  CHECK-NEXT:   llvm.return
  #blocked = #ttg.blocked<{sizePerThread = [1, 4], threadsPerWarp = [1, 32], warpsPerCTA = [1, 4], order = [1, 0]}>
  module attributes {"ttg.target" = "cuda:100", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
-@@ -1050,7 +1050,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 2 : i32} {
+@@ -1059,7 +1059,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 2 : i32} {
    // CHECK-LABEL: lower_ldmatrix_trans_b8
    tt.func @lower_ldmatrix_trans_b8(%A: !ttg.memdesc<128x64xf8E4M3FN, #shared, #smem, mutable, 1x128x64>) {
      %0 = ttg.local_load %A : !ttg.memdesc<128x64xf8E4M3FN, #shared, #smem, mutable, 1x128x64> -> tensor<128x64xf8E4M3FN, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 4}>>
@@ -103,7 +103,7 @@
      tt.return
    }
  }
-@@ -1067,7 +1067,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+@@ -1076,7 +1076,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
    // CHECK-LABEL: @ldmatrix_b8_minimise_conflicts_preserve_words
    // CHECK-DAG: %[[ONE:.*]] = llvm.mlir.constant(1 : i32)
    // CHECK-DAG: %[[TWO:.*]] = llvm.mlir.constant(2 : i32)
@@ -112,7 +112,7 @@
    // CHECK: %[[WORD:.*]] = llvm.extractvalue %[[LOAD]][0]
    // CHECK: %[[BYTES:.*]] = llvm.bitcast %[[WORD]] : i32 to vector<4xi8>
    // CHECK: %[[BYTE1:.*]] = llvm.extractelement %[[BYTES]][%[[ONE]] : i32]
-@@ -1081,7 +1081,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+@@ -1090,7 +1090,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
    // Selecting register [8] avoids the bank conflict caused by register [64].
    // CHECK-LABEL: @ldmatrix_b16_minimise_conflicts
    // CHECK-DAG: %[[ONE:.*]] = llvm.mlir.constant(1 : i32)
@@ -121,7 +121,7 @@
    // CHECK: %[[WORD:.*]] = llvm.extractvalue %[[LOAD]][0]
    // CHECK: %[[HALVES:.*]] = llvm.bitcast %[[WORD]] : i32 to vector<2xf16>
    // CHECK: %[[HALF1:.*]] = llvm.extractelement %[[HALVES]][%[[ONE]] : i32]
-@@ -1098,7 +1098,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+@@ -1107,7 +1107,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
    // CHECK: %[[PACK2:.*]] = llvm.insertelement %{{.*}}, %[[PACK1]][
    // CHECK: %[[PACK3:.*]] = llvm.insertelement %{{.*}}, %[[PACK2]][
    // CHECK: %[[WORD:.*]] = llvm.bitcast %[[PACK3]] : vector<4xi8> to i32
@@ -130,7 +130,7 @@
    tt.func private @stmatrix_b8_minimise_conflicts(%A: !ttg.memdesc<512xi8, #shared, #smem, mutable>, %data: tensor<512xi8, #stb8>) {
      ttg.local_store %data, %A : tensor<512xi8, #stb8> -> !ttg.memdesc<512xi8, #shared, #smem, mutable>
      tt.return
-@@ -1108,7 +1108,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+@@ -1117,7 +1117,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
    // CHECK: %[[INPUT2:.*]] = llvm.extractvalue %{{.*}}[2]
    // CHECK: %[[PACK:.*]] = llvm.insertelement %[[INPUT2]], %{{.*}}[%[[ONE]] : i32]
    // CHECK: %[[WORD:.*]] = llvm.bitcast %[[PACK]] : vector<2xf16> to i32
@@ -139,7 +139,7 @@
    tt.func private @stmatrix_b16_minimise_conflicts(%A: !ttg.memdesc<256xf16, #shared, #smem, mutable>, %data: tensor<256xf16, #b16>) {
      ttg.local_store %data, %A : tensor<256xf16, #b16> -> !ttg.memdesc<256xf16, #shared, #smem, mutable>
      tt.return
-@@ -1123,7 +1123,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+@@ -1132,7 +1132,7 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
  module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
    // CHECK-LABEL: @stmatrix_b8_trans_linear
    tt.func public @stmatrix_b8_trans_linear(%data: tensor<1x1x1x16x256xf8E4M3FN, #linear3>) {
@@ -151,7 +151,7 @@
 
 --- a/test/Conversion/tritongpu_to_llvm_hopper.mlir
 +++ b/test/Conversion/tritongpu_to_llvm_hopper.mlir
-@@ -470,7 +470,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
+@@ -473,7 +473,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
  // CHECK-LABEL: distribute_to_shared_st_matrix_local_store
  module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
    tt.func @distribute_to_shared_st_matrix_local_store(%a: tensor<64x128xf16, #linear>) {
@@ -160,7 +160,7 @@
      //          CHECK: llvm.return
      %b = ttg.local_alloc {allocation.offset = 0 : i32} : () -> !ttg.memdesc<64x128xf16, #shared, #smem, mutable>
      ttg.local_store %a, %b : tensor<64x128xf16, #linear> -> !ttg.memdesc<64x128xf16, #shared, #smem, mutable>
-@@ -502,7 +502,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
+@@ -505,7 +505,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
  // CHECK-LABEL: linear_to_swizzled_st_matrix_local_store
  module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
    tt.func @linear_to_swizzled_st_matrix_local_store(%a: tensor<64x32xf16, #linear>) {
@@ -169,7 +169,7 @@
      //          CHECK: llvm.return
      %b = ttg.local_alloc {allocation.offset = 0 : i32} : () -> !ttg.memdesc<64x32xf16, #shared, #smem, mutable>
      ttg.local_store %a, %b : tensor<64x32xf16, #linear> -> !ttg.memdesc<64x32xf16, #shared, #smem, mutable>
-@@ -524,7 +524,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
+@@ -527,7 +527,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
  // CHECK-LABEL: linear_to_swizzled_st_matrix_local_store
  module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
    tt.func @linear_to_swizzled_st_matrix_local_store(%a: tensor<32x32xf16, #linear>) {
@@ -178,7 +178,7 @@
      //          CHECK: llvm.return
      %b = ttg.local_alloc {allocation.offset = 0 : i32} : () -> !ttg.memdesc<32x32xf16, #shared, #smem, mutable>
      ttg.local_store %a, %b : tensor<32x32xf16, #linear> -> !ttg.memdesc<32x32xf16, #shared, #smem, mutable>
-@@ -540,7 +540,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
+@@ -543,7 +543,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
  // CHECK-LABEL: linear_to_swizzled_st_matrix_x2_local_store_fp8
  module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
    tt.func @linear_to_swizzled_st_matrix_x2_local_store_fp8(%a: tensor<64x16xf8E4M3FNUZ, #linear>) {
@@ -187,7 +187,7 @@
      //          CHECK: llvm.return
      %b = ttg.local_alloc {allocation.offset = 0 : i32} : () -> !ttg.memdesc<64x16xf8E4M3FNUZ, #shared, #smem, mutable>
      ttg.local_store %a, %b : tensor<64x16xf8E4M3FNUZ, #linear> -> !ttg.memdesc<64x16xf8E4M3FNUZ, #shared, #smem, mutable>
-@@ -556,7 +556,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
+@@ -559,7 +559,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
  // CHECK-LABEL: linear_to_swizzled_st_matrix_local_store_fp32
  module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
    tt.func @linear_to_swizzled_st_matrix_local_store_fp32(%a: tensor<64x16xf32, #linear>) {
@@ -196,7 +196,7 @@
      //          CHECK: llvm.return
      %b = ttg.local_alloc {allocation.offset = 0 : i32} : () -> !ttg.memdesc<64x16xf32, #shared, #smem, mutable>
      ttg.local_store %a, %b : tensor<64x16xf32, #linear> -> !ttg.memdesc<64x16xf32, #shared, #smem, mutable>
-@@ -573,7 +573,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
+@@ -576,7 +576,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
  // CHECK-LABEL: linear_to_swizzled_st_matrix_trans_local_store
  module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
    tt.func @linear_to_swizzled_st_matrix_trans_local_store(%a: tensor<64x32xf16, #linear>) {
@@ -205,7 +205,7 @@
      //          CHECK: llvm.return
      %b = ttg.local_alloc {allocation.offset = 0 : i32} : () -> !ttg.memdesc<64x32xf16, #shared, #smem, mutable>
      ttg.local_store %a, %b : tensor<64x32xf16, #linear> -> !ttg.memdesc<64x32xf16, #shared, #smem, mutable>
-@@ -595,7 +595,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
+@@ -598,7 +598,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
  // CHECK-LABEL: linear_to_swizzled_st_matrix_trans_local_store
  module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
    tt.func @linear_to_swizzled_st_matrix_trans_local_store(%a: tensor<16x32xf16, #linear>) {
diff --git a/third_party/xla/third_party/triton/common/llvm_cl979179688.patch b/third_party/xla/third_party/triton/common/llvm_cl979179688.patch
index 15258224806..b10b7baf8a8 100644
--- a/third_party/xla/third_party/triton/common/llvm_cl979179688.patch
+++ b/third_party/xla/third_party/triton/common/llvm_cl979179688.patch
@@ -21,13 +21,14 @@
 
 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp
-@@ -389,7 +389,9 @@ struct FpToFpOpConversion
-           "unsupported rounding mode for f32->bf16 conversion: " +
-           stringifyRoundingMode(rounding) + "\n");
+@@ -334,8 +334,10 @@ struct FpToFpOpConversion
+           "unsupported rounding mode for f32->" + dstName +
+           " conversion: " + stringifyRoundingMode(rounding) + "\n");
      }
--    return LLVM::createLLVMIntrinsicCallOp(rewriter, loc, name, bf16_ty, {v})
 +    Value pzo = LLVM::createConstantI1(loc, rewriter, false);
-+    return LLVM::createLLVMIntrinsicCallOp(rewriter, loc, name, bf16_ty,
+     auto name = ("llvm.nvvm.f2" + dstName + "." + suffix).str();
+-    return LLVM::createLLVMIntrinsicCallOp(rewriter, loc, name, dstTy, {v})
++    return LLVM::createLLVMIntrinsicCallOp(rewriter, loc, name, dstTy,
 +                                           {v, pzo})
          .getResult(0);
    }
diff --git a/third_party/xla/third_party/triton/common/llvm_cl979526570.patch b/third_party/xla/third_party/triton/common/llvm_cl979526570.patch
index 162119c597a..6afafc4457c 100644
--- a/third_party/xla/third_party/triton/common/llvm_cl979526570.patch
+++ b/third_party/xla/third_party/triton/common/llvm_cl979526570.patch
@@ -1,6 +1,6 @@
 --- a/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp
 +++ b/lib/Conversion/TritonGPUToLLVM/ElementwiseOpToLLVM.cpp
-@@ -268,7 +268,8 @@ struct ElementwiseInlineAsmOpConversion
+@@ -277,7 +277,8 @@ struct ElementwiseInlineAsmOpConversion
                             /*asm_string=*/op.getAsmString(),
                             /*constraints=*/op.getConstraints(),
                             /*has_side_effects=*/!op.getPure(),
@@ -10,7 +10,7 @@
                             /*asm_dialect=*/
                             LLVM::AsmDialectAttr::get(rewriter.getContext(),
                                                       LLVM::AsmDialect::AD_ATT),
-@@ -417,6 +418,7 @@ struct InlineAsmOpConversion
+@@ -426,6 +427,7 @@ struct InlineAsmOpConversion
      auto call = LLVM::InlineAsmOp::create(
          rewriter, loc, returnType, operands, op.getAsmString(),
          op.getConstraints(), !op.getPure(), false, LLVM::TailCallKind::None,
@@ -30,20 +30,9 @@
                 /*operand_attrs=*/ArrayAttr())
          ->getResult(0);
 
---- a/third_party/amd/lib/TritonAMDGPUToLLVM/GCNAsmFormat.cpp
-+++ b/third_party/amd/lib/TritonAMDGPUToLLVM/GCNAsmFormat.cpp
-@@ -83,6 +83,7 @@ mlir::Value GCNBuilder::launch(RewriterBase &rewriter, Location loc, Type resTy,
-       hasSideEffect,                          // has_side_effects
-       isAlignStack,                           // is_align_stack
-       LLVM::TailCallKind::None,
-+      /*convergent=*/false,
-       LLVM::AsmDialectAttr::get(ctx,
-                                 LLVM::AsmDialect::AD_ATT), // asm_dialect
-       ArrayAttr::get(ctx, attrs)                           // operand_attrs
-
 --- a/third_party/amd/lib/TritonAMDGPUToLLVM/TDMUtility.cpp
 +++ b/third_party/amd/lib/TritonAMDGPUToLLVM/TDMUtility.cpp
-@@ -37,6 +37,7 @@ Value sequenceGroup0(RewriterBase &rewriter, Location loc, Value group0) {
+@@ -39,6 +39,7 @@ Value sequenceGroup0(RewriterBase &rewriter, Location loc, Value group0) {
               /*asm_string=*/"", /*constraints=*/"=s,0",
               /*has_side_effects=*/true, /*is_align_stack=*/false,
               LLVM::TailCallKind::None,
@@ -54,7 +43,7 @@
 
 --- a/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp
 +++ b/third_party/amd/lib/TritonAMDGPUToLLVM/TargetInfo.cpp
-@@ -749,6 +749,7 @@ void TargetInfo::assertFail(RewriterBase &rewriter, Location loc,
+@@ -726,6 +726,7 @@ void TargetInfo::assertTrap(RewriterBase &rewriter, Location loc) const {
        "s_trap 2", /*constraints=*/"",
        /*has_side_effects=*/true, /*is_align_stack=*/false,
        LLVM::TailCallKind::None,
diff --git a/third_party/xla/third_party/triton/common/llvm_cl983115707.patch b/third_party/xla/third_party/triton/common/llvm_cl983115707.patch
index 3ba6c6920c9..14750c1bd06 100644
--- a/third_party/xla/third_party/triton/common/llvm_cl983115707.patch
+++ b/third_party/xla/third_party/triton/common/llvm_cl983115707.patch
@@ -96,7 +96,112 @@
 
 --- a/test/Conversion/tritongpu_to_llvm.mlir
 +++ b/test/Conversion/tritongpu_to_llvm.mlir
-@@ -2844,7 +2844,7 @@ module attributes {"ttg.target" = "cuda:75", "ttg.num-ctas" = 1 : i32, "ttg.num-
+@@ -1401,7 +1401,7 @@
+     %BB_DOT = ttg.local_load %BB : !ttg.memdesc<16x16xf16, #shared0, #smem> -> tensor<16x16xf16, #dot_operand_b>
+     %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf32, #mma0>
+ 
+-    // CHECK-COUNT-2: nvvm.mma.sync {{.*}}layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>{{.*}}shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (vector<2xf16>, vector<2xf16>, f32)
++    // CHECK-COUNT-2: nvvm.mma.sync {{.*}}shape = <m = 16, n = 8, k = 16>, layout_a = row, layout_b = col : (vector<2xf16>, vector<2xf16>, f32)
+     %D = tt.dot %AA_DOT, %BB_DOT, %cst0 : tensor<16x16xf16, #dot_operand_a> * tensor<16x16xf16, #dot_operand_b> -> tensor<16x16xf32, #mma0>
+ 
+     tt.return
+@@ -1428,7 +1428,7 @@
+     %BB_DOT = ttg.local_load %BB : !ttg.memdesc<16x16xf16, #shared0, #smem> -> tensor<16x16xf16, #dot_operand_b>
+     %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf32, #mma0>
+ 
+-    // CHECK-COUNT-2: nvvm.mma.sync {{.*}}layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>{{.*}}shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (vector<2xf16>, vector<2xf16>, f32)
++    // CHECK-COUNT-2: nvvm.mma.sync {{.*}}shape = <m = 16, n = 8, k = 16>, layout_a = row, layout_b = col : (vector<2xf16>, vector<2xf16>, f32)
+     %D = tt.dot %AA_DOT, %BB_DOT, %cst0 : tensor<16x16xf16, #dot_operand_a> * tensor<16x16xf16, #dot_operand_b> -> tensor<16x16xf32, #mma0>
+ 
+     tt.return
+@@ -1453,7 +1453,7 @@
+     %BB_DOT = ttg.local_load %BB : !ttg.memdesc<16x16xf16, #shared0, #smem> -> tensor<16x16xf16, #dot_operand_b>
+     %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf32, #mma0>
+ 
+-    // CHECK-COUNT-2: nvvm.mma.sync {{.*}}layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>{{.*}}shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (vector<2xf16>, vector<2xf16>, f32)
++    // CHECK-COUNT-2: nvvm.mma.sync {{.*}}shape = <m = 16, n = 8, k = 16>, layout_a = row, layout_b = col : (vector<2xf16>, vector<2xf16>, f32)
+     %D = tt.dot %AA_DOT, %BB_DOT, %cst0 : tensor<16x16xf16, #dot_operand_a> * tensor<16x16xf16, #dot_operand_b> -> tensor<16x16xf32, #mma0>
+ 
+     tt.return
+@@ -1503,7 +1503,7 @@
+     %BB_DOT = ttg.local_load %BB : !ttg.memdesc<16x16xf8E5M2, #shared0, #smem> -> tensor<16x16xf8E5M2, #dot_operand_b>
+     %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf32, #mma0>
+ 
+-    // CHECK-COUNT-2: nvvm.mma.sync {{.*}}layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>{{.*}}multiplicandAPtxType = #nvvm.mma_type<e5m2>, multiplicandBPtxType = #nvvm.mma_type<e5m2>{{.*}}shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, f32)
++    // CHECK-COUNT-2: nvvm.mma.sync {{.*}}shape = <m = 16, n = 8, k = 32>, layout_a = row, layout_b = col, multiplicand_a_ptx_type = e5m2, multiplicand_b_ptx_type = e5m2 : (i32, i32, f32)
+     %D = tt.dot %AA_DOT, %BB_DOT, %cst0 : tensor<16x16xf8E5M2, #dot_operand_a> * tensor<16x16xf8E5M2, #dot_operand_b> -> tensor<16x16xf32, #mma0>
+ 
+     tt.return
+@@ -1946,7 +1946,7 @@
+     %a_mat = ttg.local_load %a : !ttg.memdesc<32x16xf32, #shared, #smem> -> tensor<32x16xf32, #dot_operand_a>
+     %b_mat = ttg.local_load %b : !ttg.memdesc<16x32xf32, #shared, #smem> -> tensor<16x32xf32, #dot_operand_b>
+ 
+-    // CHECK-COUNT-4: nvvm.mma.sync {{.*}}layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>{{.*}}multiplicandAPtxType = #nvvm.mma_type<tf32>, multiplicandBPtxType = #nvvm.mma_type<tf32>{{.*}}shape = #nvvm.shape<m = 16, n = 8, k = 8>} : (i32, i32, f32)
++    // CHECK-COUNT-4: nvvm.mma.sync {{.*}}shape = <m = 16, n = 8, k = 8>, layout_a = row, layout_b = col, multiplicand_a_ptx_type = tf32, multiplicand_b_ptx_type = tf32 : (i32, i32, f32)
+     %28 = tt.dot %a_mat, %b_mat, %cst, inputPrecision = tf32 : tensor<32x16xf32, #dot_operand_a> * tensor<16x32xf32, #dot_operand_b> -> tensor<32x32xf32, #mma>
+     %38 = ttg.convert_layout %28 : tensor<32x32xf32, #mma> -> tensor<32x32xf32, #blocked>
+ 
+@@ -1967,7 +1967,7 @@
+   tt.func @matmul_signed_i8dot(%a: tensor<16x32xi8, #dot_operand_a>,
+                                %b: tensor<32x8xi8, #dot_operand_b>,
+                                %c: tensor<16x8xi32, #mma>) {
+-    // CHECK: nvvm.mma.sync {{.*}}intOverflowBehavior = #nvvm.mma_int_overflow<satfinite>{{.*}}layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>{{.*}}multiplicandAPtxType = #nvvm.mma_type<s8>, multiplicandBPtxType = #nvvm.mma_type<s8>{{.*}}shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, i32)
++    // CHECK: nvvm.mma.sync {{.*}}shape = <m = 16, n = 8, k = 32>, int_overflow = satfinite, layout_a = row, layout_b = col, multiplicand_a_ptx_type = s8, multiplicand_b_ptx_type = s8 : (i32, i32, i32)
+     %d = tt.dot %a, %b, %c : tensor<16x32xi8, #dot_operand_a> * tensor<32x8xi8, #dot_operand_b> -> tensor<16x8xi32, #mma>
+     tt.return
+   }
+@@ -1984,10 +1984,10 @@
+                                      %b: tensor<32x8xi8, #dot_operand_b>,
+                                      %c: tensor<16x8xi32, #mma>) {
+     // CHECK-NOT: satfinite
+-    // CHECK: nvvm.mma.sync {{.*}}intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>{{.*}}layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>{{.*}}multiplicandAPtxType = #nvvm.mma_type<s8>, multiplicandBPtxType = #nvvm.mma_type<s8>{{.*}}shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, i32)
+-    // CHECK: nvvm.mma.sync {{.*}}intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>{{.*}}layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>{{.*}}multiplicandAPtxType = #nvvm.mma_type<s8>, multiplicandBPtxType = #nvvm.mma_type<u8>{{.*}}shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, i32)
+-    // CHECK: nvvm.mma.sync {{.*}}intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>{{.*}}layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>{{.*}}multiplicandAPtxType = #nvvm.mma_type<u8>, multiplicandBPtxType = #nvvm.mma_type<s8>{{.*}}shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, i32)
+-    // CHECK: nvvm.mma.sync {{.*}}intOverflowBehavior = #nvvm.mma_int_overflow<wrapped>{{.*}}layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>{{.*}}multiplicandAPtxType = #nvvm.mma_type<u8>, multiplicandBPtxType = #nvvm.mma_type<u8>{{.*}}shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, i32)
++    // CHECK: nvvm.mma.sync {{.*}}shape = <m = 16, n = 8, k = 32>, int_overflow = wrapped, layout_a = row, layout_b = col, multiplicand_a_ptx_type = s8, multiplicand_b_ptx_type = s8 : (i32, i32, i32)
++    // CHECK: nvvm.mma.sync {{.*}}shape = <m = 16, n = 8, k = 32>, int_overflow = wrapped, layout_a = row, layout_b = col, multiplicand_a_ptx_type = s8, multiplicand_b_ptx_type = u8 : (i32, i32, i32)
++    // CHECK: nvvm.mma.sync {{.*}}shape = <m = 16, n = 8, k = 32>, int_overflow = wrapped, layout_a = row, layout_b = col, multiplicand_a_ptx_type = u8, multiplicand_b_ptx_type = s8 : (i32, i32, i32)
++    // CHECK: nvvm.mma.sync {{.*}}shape = <m = 16, n = 8, k = 32>, int_overflow = wrapped, layout_a = row, layout_b = col, multiplicand_a_ptx_type = u8, multiplicand_b_ptx_type = u8 : (i32, i32, i32)
+     // CHECK-NOT: satfinite
+     %d0 = tti.dot_i8 %a, %b, %c, aSigned = true, bSigned = true : tensor<16x32xi8, #dot_operand_a> * tensor<32x8xi8, #dot_operand_b> -> tensor<16x8xi32, #mma>
+     %d1 = tti.dot_i8 %a, %b, %d0, aSigned = true, bSigned = false : tensor<16x32xi8, #dot_operand_a> * tensor<32x8xi8, #dot_operand_b> -> tensor<16x8xi32, #mma>
+@@ -2469,7 +2469,7 @@
+   // CHECK: %[[CST:.+]] = llvm.mlir.constant(1.000000e+00 : f32) : f32
+   // CHECK: %[[V:.+]] = llvm.insertelement %[[CST]], {{.*}} : vector<1xf32>
+   // CHECK: llvm.bitcast %[[V]] : vector<1xf32> to i32
+-  // CHECK: nvvm.mma.sync {{.*}}layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>{{.*}}multiplicandAPtxType = #nvvm.mma_type<tf32>, multiplicandBPtxType = #nvvm.mma_type<tf32>{{.*}}shape = #nvvm.shape<m = 16, n = 8, k = 8>} : (i32, i32, f32)
++  // CHECK: nvvm.mma.sync {{.*}}shape = <m = 16, n = 8, k = 8>, layout_a = row, layout_b = col, multiplicand_a_ptx_type = tf32, multiplicand_b_ptx_type = tf32 : (i32, i32, f32)
+     %b_mat = arith.constant dense<1.000000e+00> : tensor<16x32xf32, #dot_operand_b>
+     %28 = tt.dot %a, %b_mat, %c, inputPrecision = tf32 : tensor<32x16xf32, #dot_operand_a> * tensor<16x32xf32, #dot_operand_b> -> tensor<32x32xf32, #mma>
+     %38 = ttg.convert_layout %28 : tensor<32x32xf32, #mma> -> tensor<32x32xf32, #blocked>
+@@ -2494,7 +2494,7 @@
+   // CHECK: %[[V0:.+]] = llvm.insertelement %{{.*}}, %[[U]][%[[C0]] : i32] : vector<2xf16>
+   // CHECK: %[[V1:.+]] = llvm.insertelement %{{.*}}, %[[V0]][%[[C1]] : i32] : vector<2xf16>
+   // CHECK: nvvm.mma.sync A[%[[V1]],
+-  // CHECK-SAME: shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (vector<2xf16>,
++  // CHECK-SAME: shape = <m = 16, n = 8, k = 16>, layout_a = row, layout_b = col : (vector<2xf16>,
+     %cst_0 = arith.constant dense<1.000000e+00> : tensor<32x32xf16, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 2}>>
+     %cst_1 = arith.constant dense<1.000000e+00> : tensor<32x32xf16, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 2}>>
+     %cst_2 = arith.constant dense<32> : tensor<32x1xi32, #blocked>
+@@ -2738,7 +2738,7 @@
+     %converted_i16 = arith.sitofp %i16_dot : tensor<16x16xi16, #dot_operand_b> to tensor<16x16xf16, #dot_operand_b>
+     %cst0 = arith.constant dense<0.000000e+00> : tensor<16x16xf32, #mma>
+ 
+-    // CHECK-COUNT-2: nvvm.mma.sync {{.*}}layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>{{.*}}shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (vector<2xf16>, vector<2xf16>, f32)
++    // CHECK-COUNT-2: nvvm.mma.sync {{.*}}shape = <m = 16, n = 8, k = 16>, layout_a = row, layout_b = col : (vector<2xf16>, vector<2xf16>, f32)
+ 
+     %out = tt.dot %f16_dot, %converted_i16, %cst0 : tensor<16x16xf16, #dot_operand_a> * tensor<16x16xf16, #dot_operand_b> -> tensor<16x16xf32, #mma>
+ 
+@@ -2763,7 +2763,7 @@
+ 
+     %3 = ttg.local_load %1 : !ttg.memdesc<16x16xf64, #shared1, #smem, mutable> -> tensor<16x16xf64, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 1}>>
+ 
+-    // CHECK-COUNT-2: nvvm.mma.sync {{.*}}layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>{{.*}}shape = #nvvm.shape<m = 8, n = 8, k = 4>} : (f64, f64, f64)
++    // CHECK-COUNT-2: nvvm.mma.sync {{.*}}shape = <m = 8, n = 8, k = 4>, layout_a = row, layout_b = col : (f64, f64, f64)
+ 
+     %out = tt.dot %2, %3, %cst, inputPrecision = tf32 : tensor<16x16xf64, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 1}>> * tensor<16x16xf64, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 1}>> -> tensor<16x16xf64, #mma>
+ 
+@@ -2815,7 +2815,7 @@ module attributes {"ttg.target" = "cuda:75", "ttg.num-ctas" = 1 : i32, "ttg.num-
  module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, "ttg.threads-per-warp" = 32 : i32} {
    // CHECK-LABEL: @vectorize_shmem_load
    // CHECK: llvm.load
@@ -105,7 +210,7 @@
    // CHECK-NOT: llvm.load
    tt.func private @vectorize_shmem_load(%shmem : !ttg.memdesc<16x16xi8, #shared, #smem>) -> tensor<16x16xi8, #blocked> {
      %0 = ttg.local_load %shmem : !ttg.memdesc<16x16xi8, #shared, #smem> -> tensor<16x16xi8, #blocked>
-@@ -2859,7 +2859,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
+@@ -2830,7 +2830,7 @@ module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-
  #smem = #ttg.shared_memory
  module attributes {"ttg.target" = "cuda:90", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 8 : i32, "ttg.threads-per-warp" = 32 : i32} {
    // CHECK-LABEL: @vectorize_shmem_store
@@ -114,18 +219,469 @@
    tt.func public @vectorize_shmem_store(%block : tensor<64x64xi32, #blocked>) {
      %0 = ttg.local_alloc %block : (tensor<64x64xi32, #blocked>) -> !ttg.memdesc<64x64xi32, #shared, #smem>
      tt.return
+@@ -3284,28 +3284,28 @@
+ module attributes {"ttg.num-warps" = 8 : i32, ttg.target = "cuda:120"} {
+   // CHECK-LABEL: mmav2_e5m2_e5m2_fp16
+   tt.func public @mmav2_e5m2_e5m2_fp16(%arg0: tensor<32x32xf8E5M2, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 4}>>, %arg1: tensor<32x32xf8E5M2, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 4}>>, %arg2: tensor<32x32xf16, #mma>) {
+-    // CHECK: nvvm.mma.sync {{.*}}multiplicandAPtxType = #nvvm.mma_type<e5m2>, multiplicandBPtxType = #nvvm.mma_type<e5m2>{{.*}}shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, vector<2xf16>)
++    // CHECK: nvvm.mma.sync {{.*}}shape = <m = 16, n = 8, k = 32>, layout_a = row, layout_b = col, multiplicand_a_ptx_type = e5m2, multiplicand_b_ptx_type = e5m2 : (i32, i32, vector<2xf16>)
+     %0 = tt.dot %arg0, %arg1, %arg2 {maxNumImpreciseAcc = 1073741824 : i32} : tensor<32x32xf8E5M2, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 4}>> * tensor<32x32xf8E5M2, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 4}>> -> tensor<32x32xf16, #mma>
+     tt.return
+   }
+ 
+   // CHECK-LABEL: mmav2_e5m2_e4m3_fp16
+   tt.func public @mmav2_e5m2_e4m3_fp16(%arg0: tensor<32x32xf8E5M2, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 4}>>, %arg1: tensor<32x32xf8E4M3FN, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 4}>>, %arg2: tensor<32x32xf16, #mma>) {
+-    // CHECK: nvvm.mma.sync {{.*}}multiplicandAPtxType = #nvvm.mma_type<e5m2>, multiplicandBPtxType = #nvvm.mma_type<e4m3>{{.*}}shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, vector<2xf16>)
++    // CHECK: nvvm.mma.sync {{.*}}shape = <m = 16, n = 8, k = 32>, layout_a = row, layout_b = col, multiplicand_a_ptx_type = e5m2, multiplicand_b_ptx_type = e4m3 : (i32, i32, vector<2xf16>)
+     %0 = tt.dot %arg0, %arg1, %arg2 {maxNumImpreciseAcc = 1073741824 : i32} : tensor<32x32xf8E5M2, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 4}>> * tensor<32x32xf8E4M3FN, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 4}>> -> tensor<32x32xf16, #mma>
+     tt.return
+   }
+ 
+   // CHECK-LABEL: mmav2_e4m3_e5m2_fp16
+   tt.func public @mmav2_e4m3_e5m2_fp16(%arg0: tensor<32x32xf8E4M3FN, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 4}>>, %arg1: tensor<32x32xf8E5M2, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 4}>>, %arg2: tensor<32x32xf16, #mma>) {
+-    // CHECK: nvvm.mma.sync {{.*}}multiplicandAPtxType = #nvvm.mma_type<e4m3>, multiplicandBPtxType = #nvvm.mma_type<e5m2>{{.*}}shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, vector<2xf16>)
++    // CHECK: nvvm.mma.sync {{.*}}shape = <m = 16, n = 8, k = 32>, layout_a = row, layout_b = col, multiplicand_a_ptx_type = e4m3, multiplicand_b_ptx_type = e5m2 : (i32, i32, vector<2xf16>)
+     %0 = tt.dot %arg0, %arg1, %arg2 {maxNumImpreciseAcc = 1073741824 : i32} : tensor<32x32xf8E4M3FN, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 4}>> * tensor<32x32xf8E5M2, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 4}>> -> tensor<32x32xf16, #mma>
+     tt.return
+   }
+ 
+   // CHECK-LABEL: mmav2_e4m3_e4m3_fp16
+   tt.func public @mmav2_e4m3_e4m3_fp16(%arg0: tensor<32x32xf8E4M3FN, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 4}>>, %arg1: tensor<32x32xf8E4M3FN, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 4}>>, %arg2: tensor<32x32xf16, #mma>) {
+-    // CHECK: nvvm.mma.sync {{.*}}multiplicandAPtxType = #nvvm.mma_type<e4m3>, multiplicandBPtxType = #nvvm.mma_type<e4m3>{{.*}}shape = #nvvm.shape<m = 16, n = 8, k = 32>} : (i32, i32, vector<2xf16>)
++    // CHECK: nvvm.mma.sync {{.*}}shape = <m = 16, n = 8, k = 32>, layout_a = row, layout_b = col, multiplicand_a_ptx_type = e4m3, multiplicand_b_ptx_type = e4m3 : (i32, i32, vector<2xf16>)
+     %0 = tt.dot %arg0, %arg1, %arg2 {maxNumImpreciseAcc = 1073741824 : i32} : tensor<32x32xf8E4M3FN, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 4}>> * tensor<32x32xf8E4M3FN, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 4}>> -> tensor<32x32xf16, #mma>
+     tt.return
+   }
+@@ -3579,7 +3579,7 @@
+ 
+   // CHECK-LABEL: approx_div
+   tt.func private @approx_div(%arg0 : tensor<256xf32, #blocked>, %arg1 : tensor<256xf32, #blocked>) -> tensor<256xf32, #blocked> {
+-    // CHECK: llvm.fdiv %{{.*}}, %{{.*}} {fastmathFlags = #llvm.fastmath<afn>} : f32
++    // CHECK: llvm.fdiv %{{.*}}, %{{.*}} fastmath<afn> : f32
+     %0 = tt.approx_divf %arg0, %arg1 : tensor<256xf32, #blocked>
+     tt.return %0 : tensor<256xf32, #blocked>
+   }
+@@ -3602,7 +3602,7 @@
+     %arg1: tensor<32x32x32xf16, #dot_operand_b>
+   ) {
+     %cst = arith.constant dense<0.000000e+00> : tensor<32x1x32xf32, #mma>
+-    // CHECK: nvvm.mma.sync {{.*}}layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>{{.*}}shape = #nvvm.shape<m = 16, n = 8, k = 16>} : (vector<2xf16>, vector<2xf16>, f32)
++    // CHECK: nvvm.mma.sync {{.*}}shape = <m = 16, n = 8, k = 16>, layout_a = row, layout_b = col : (vector<2xf16>, vector<2xf16>, f32)
+     %result = tt.dot %arg0, %arg1, %cst, inputPrecision = tf32 :
+       tensor<32x1x32xf16, #dot_operand_a> * tensor<32x32x32xf16, #dot_operand_b> -> tensor<32x1x32xf32, #mma>
+     tt.return
+@@ -3929,7 +3929,7 @@
+ module attributes {"ttg.target" = "cuda:80", "ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+   // SM89-LABEL: @fp32_to_fp8e5_rtne
+   // SM89-NOT: llvm.nvvm.f2f16.rz
+-  // SM89: nvvm.convert.f32x2.to.f8x2 {{.*}} {rnd = #nvvm.fp_rnd_mode<rn>, sat = #nvvm.sat_mode<satfinite>} : vector<2xi8>(f8E5M2)
++  // SM89: nvvm.convert.f32x2.to.f8x2 {{.*}} rnd = <rn> sat = <satfinite> : vector<2xi8>(f8E5M2)
+   // CHECK-LABEL: @fp32_to_fp8e5_rtne
+   // CHECK-NOT: llvm.fpext
+   // CHECK-NOT: llvm.fcmp
+
+--- a/test/Conversion/tritongpu_to_llvm_blackwell.mlir
++++ b/test/Conversion/tritongpu_to_llvm_blackwell.mlir
+@@ -253,7 +253,7 @@
+   // CHECK: nvg.tensor_memory_base
+   // CHECK: tcgen05.st.sync.aligned.32x32b.x128.b32
+   // CHECK: nvvm.tcgen05.wait <store>
+-  // CHECK: nvvm.tcgen05.ld {{.*}}shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<128xi32>
++  // CHECK: nvvm.tcgen05.ld {{.*}}shape = shape_32x32b : vector<128xi32>
+   // CHECK: nvvm.tcgen05.wait <load>
+   tt.func public @tensor_memory_ld(%arg0: !tt.ptr<f16>, %arg1: !tt.ptr<f16>, %arg2: !tt.ptr<f16>) {
+     %cst_0 = arith.constant dense<0.000000e+00> : tensor<128x128xf32, #blocked1>
+@@ -271,7 +271,7 @@
+   // CHECK-LABEL: @tensor_memory_ld_16x256
+   // CHECK: tcgen05.st.sync.aligned.16x256b.x16.b32
+   // CHECK: tcgen05.st.sync.aligned.16x256b.x16.b32
+-  // CHECK-COUNT-2: nvvm.tcgen05.ld {{.*}}shape = #nvvm.tcgen05_ldst_shape<shape_16x256b>} : vector<64xi32>
++  // CHECK-COUNT-2: nvvm.tcgen05.ld {{.*}}shape = shape_16x256b : vector<64xi32>
+   tt.func public @tensor_memory_ld_16x256(%arg0: !tt.ptr<f16>, %arg1: !tt.ptr<f16>, %arg2: !tt.ptr<f16>) {
+     %cst_0 = arith.constant dense<0.000000e+00> : tensor<128x128xf32, #linear>
+     %0 = ttng.tmem_alloc %cst_0 {tensor_memory_col_offset = 0 : i32, tensor_memory_row_offset = 0 : i32} : (tensor<128x128xf32, #linear>) -> !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>
+@@ -301,7 +301,7 @@
+   // CHECK: nvg.tensor_memory_base
+   // CHECK: tcgen05.st.sync.aligned.32x32b.x128.b32
+   // CHECK: nvvm.tcgen05.wait <store>
+-  // CHECK: nvvm.tcgen05.ld {{.*}}shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<128xi32>
++  // CHECK: nvvm.tcgen05.ld {{.*}}shape = shape_32x32b : vector<128xi32>
+   // CHECK: nvvm.tcgen05.wait <load>
+   tt.func public @tensor_memory_ld_m64(%arg0: !tt.ptr<f16>, %arg1: !tt.ptr<f16>, %arg2: !tt.ptr<f16>) {
+     %cst_0 = arith.constant dense<0.000000e+00> : tensor<128x128xf32, #linear>
+@@ -321,7 +321,7 @@
+   // CHECK-DAG: nvg.tensor_memory_base
+   // CHECK: tcgen05.st.sync.aligned.32x32b.x64.unpack::16b.b32
+   // CHECK: nvvm.tcgen05.wait <store>
+-  // CHECK: nvvm.tcgen05.ld {{.*}} pack {shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<64xi32>
++  // CHECK: nvvm.tcgen05.ld {{.*}} pack shape = shape_32x32b : vector<64xi32>
+   // CHECK: nvvm.tcgen05.wait <load>
+   tt.func public @tensor_memory_unpack_f16() {
+     %cst_0 = arith.constant dense<0.000000e+00> : tensor<128x128xf16, #blocked1>
+@@ -647,7 +647,7 @@
+   // CHECK-COUNT-4: tcgen05.st.sync.aligned.32x32b.x64.b32
+   // CHECK-NOT: tcgen05.st
+   // CHECK: nvvm.tcgen05.wait <store>
+-  // CHECK-COUNT-4: nvvm.tcgen05.ld {{.*}}shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<64xi32>
++  // CHECK-COUNT-4: nvvm.tcgen05.ld {{.*}}shape = shape_32x32b : vector<64xi32>
+   // CHECK-NOT: tcgen05.ld
+   // CHECK: nvvm.tcgen05.wait <load>
+   tt.func public @tensor_memory_ld_128x256(%arg0: !tt.ptr<f16>, %arg1: !tt.ptr<f16>, %arg2: !tt.ptr<f16>) {
+@@ -667,7 +667,7 @@
+   // CHECK-LABEL: @tensor_memory_ld_128x256_8_warps
+   // CHECK: tcgen05.st.sync.aligned.32x32b.x128.b32
+   // CHECK: nvvm.tcgen05.wait <store>
+-  // CHECK: nvvm.tcgen05.ld {{.*}}shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<128xi32>
++  // CHECK: nvvm.tcgen05.ld {{.*}}shape = shape_32x32b : vector<128xi32>
+   // CHECK: nvvm.tcgen05.wait <load>
+   tt.func public @tensor_memory_ld_128x256_8_warps(%arg0: !tt.ptr<f16>, %arg1: !tt.ptr<f16>, %arg2: !tt.ptr<f16>) {
+     %cst_0 = arith.constant dense<0.000000e+00> : tensor<128x256xf32, #blocked>
+@@ -685,7 +685,7 @@
+ module attributes {"ttg.num-warps" = 8 : i32} {
+   // CHECK-LABEL: @tensor_memory_ld_256x64_8_warps_blocked
+   tt.func public @tensor_memory_ld_256x64_8_warps_blocked(%tmem: !ttg.memdesc<256x64xf32, #tmem, #ttng.tensor_memory, mutable>) {
+-    // CHECK-COUNT-1: nvvm.tcgen05.ld {{.*}}shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<64xi32>
++    // CHECK-COUNT-1: nvvm.tcgen05.ld {{.*}}shape = shape_32x32b : vector<64xi32>
+     // CHECK-NOT: tcgen05.ld
+     %result = ttng.tmem_load %tmem : !ttg.memdesc<256x64xf32, #tmem, #ttng.tensor_memory, mutable> -> tensor<256x64xf32, #blocked>
+     tt.return
+@@ -700,7 +700,7 @@
+ module attributes {"ttg.num-warps" = 8 : i32} {
+   // CHECK-LABEL: @tensor_memory_ld_256x64_8_warps_splitM
+   tt.func public @tensor_memory_ld_256x64_8_warps_splitM(%tmem: !ttg.memdesc<256x64xf32, #tmem, #ttng.tensor_memory, mutable>) {
+-    // CHECK: nvvm.tcgen05.ld {{.*}}shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<64xi32>
++    // CHECK: nvvm.tcgen05.ld {{.*}}shape = shape_32x32b : vector<64xi32>
+     // CHECK-NOT: tcgen05.ld
+     %result = ttng.tmem_load %tmem : !ttg.memdesc<256x64xf32, #tmem, #ttng.tensor_memory, mutable> -> tensor<256x64xf32, #linear>
+     tt.return
+@@ -715,7 +715,7 @@
+ module attributes {"ttg.num-warps" = 8 : i32} {
+   // CHECK-LABEL: @tensor_memory_ld_128x128_8_warps_splitM
+   tt.func public @tensor_memory_ld_128x128_8_warps_splitM(%tmem: !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>) {
+-    // CHECK-COUNT-1: nvvm.tcgen05.ld {{.*}}shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<64xi32>
++    // CHECK-COUNT-1: nvvm.tcgen05.ld {{.*}}shape = shape_32x32b : vector<64xi32>
+     // CHECK-NOT: tcgen05.ld
+     %result = ttng.tmem_load %tmem : !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable> -> tensor<128x128xf32, #linear>
+     tt.return
+@@ -730,7 +730,7 @@
+ module attributes {"ttg.num-warps" = 8 : i32} {
+   // CHECK-LABEL: @tensor_memory_ld_128x64_8_warps_splitM
+   tt.func public @tensor_memory_ld_128x64_8_warps_splitM(%tmem: !ttg.memdesc<128x64xf32, #tmem, #ttng.tensor_memory, mutable>) {
+-    // CHECK-COUNT-1: nvvm.tcgen05.ld {{.*}}shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<32xi32>
++    // CHECK-COUNT-1: nvvm.tcgen05.ld {{.*}}shape = shape_32x32b : vector<32xi32>
+     // CHECK-NOT: tcgen05.ld
+     %result = ttng.tmem_load %tmem : !ttg.memdesc<128x64xf32, #tmem, #ttng.tensor_memory, mutable> -> tensor<128x64xf32, #linear>
+     tt.return
+@@ -747,9 +747,9 @@
+ // CHECK-LABEL: @tmem_message_maxnreg_80
+ tt.func public @tmem_message_maxnreg_80(%desc: !ttg.memdesc<128x64xf32, #tmem, #ttng.tensor_memory>) {
+   // CHECK: [[BASE:%.*]] = llvm.inttoptr {{.*}} : i32 to !llvm.ptr<6>
+-  // CHECK: nvvm.tcgen05.ld [[BASE]] {{.*}}{shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<32xi32>
++  // CHECK: nvvm.tcgen05.ld [[BASE]] {{.*}}shape = shape_32x32b : vector<32xi32>
+   // CHECK: [[PTR:%.*]] = llvm.getelementptr [[BASE]][32]
+-  // CHECK: nvvm.tcgen05.ld [[PTR]] {{.*}}{shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<32xi32>
++  // CHECK: nvvm.tcgen05.ld [[PTR]] {{.*}}shape = shape_32x32b : vector<32xi32>
+   // CHECK-NOT: tcgen05.ld
+   ttng.tmem_load %desc : !ttg.memdesc<128x64xf32, #tmem, #ttng.tensor_memory> -> tensor<128x64xf32, #blocked>
+   tt.return
+@@ -759,14 +759,14 @@
+ tt.func public @module_constraint_supercedes_local(%desc: !ttg.memdesc<128x64xf32, #tmem, #ttng.tensor_memory>) {
+   ttg.warp_specialize(%desc) attributes {actualRegisters = array<i32: 256, 256>}
+   default {
+-    // CHECK-COUNT-2: nvvm.tcgen05.ld {{.*}}shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<32xi32>
++    // CHECK-COUNT-2: nvvm.tcgen05.ld {{.*}}shape = shape_32x32b : vector<32xi32>
+     // CHECK-NOT: tcgen05.ld
+     // CHECK: ttg.warp_yield
+     ttng.tmem_load %desc : !ttg.memdesc<128x64xf32, #tmem, #ttng.tensor_memory> -> tensor<128x64xf32, #blocked>
+     ttg.warp_yield
+   }
+   partition0(%arg0: !ttg.memdesc<128x64xf32, #tmem, #ttng.tensor_memory>) num_warps(4) {
+-    // CHECK-COUNT-2: nvvm.tcgen05.ld {{.*}}shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<32xi32>
++    // CHECK-COUNT-2: nvvm.tcgen05.ld {{.*}}shape = shape_32x32b : vector<32xi32>
+     // CHECK-NOT: tcgen05.ld
+     // CHECK: ttg.warp_return
+     ttng.tmem_load %arg0 : !ttg.memdesc<128x64xf32, #tmem, #ttng.tensor_memory> -> tensor<128x64xf32, #blocked>
+@@ -784,9 +784,9 @@
+   ttg.warp_specialize(%desc) attributes {actualRegisters = array<i32: 80, 48>}
+   default {
+     // CHECK: [[BASE:%.*]] = llvm.inttoptr {{.*}} : i32 to !llvm.ptr<6>
+-    // CHECK: nvvm.tcgen05.ld [[BASE]] {{.*}}{shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<32xi32>
++    // CHECK: nvvm.tcgen05.ld [[BASE]] {{.*}}shape = shape_32x32b : vector<32xi32>
+     // CHECK: [[PTR:%.*]] = llvm.getelementptr [[BASE]][32]
+-    // CHECK: nvvm.tcgen05.ld [[PTR]] {{.*}}{shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<32xi32>
++    // CHECK: nvvm.tcgen05.ld [[PTR]] {{.*}}shape = shape_32x32b : vector<32xi32>
+     // CHECK-NOT: tcgen05.ld
+     // CHECK: ttg.warp_yield
+     ttng.tmem_load %desc : !ttg.memdesc<128x64xf32, #tmem, #ttng.tensor_memory> -> tensor<128x64xf32, #blocked>
+@@ -794,13 +794,13 @@
+   }
+   partition0(%arg0: !ttg.memdesc<128x64xf32, #tmem, #ttng.tensor_memory>) num_warps(4) {
+     // CHECK: [[BASE:%.*]] = llvm.inttoptr {{.*}} : i32 to !llvm.ptr<6>
+-    // CHECK: nvvm.tcgen05.ld [[BASE]] {{.*}}{shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<16xi32>
++    // CHECK: nvvm.tcgen05.ld [[BASE]] {{.*}}shape = shape_32x32b : vector<16xi32>
+     // CHECK: [[PTR:%.*]] = llvm.getelementptr [[BASE]][16]
+-    // CHECK: nvvm.tcgen05.ld [[PTR]] {{.*}}{shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<16xi32>
++    // CHECK: nvvm.tcgen05.ld [[PTR]] {{.*}}shape = shape_32x32b : vector<16xi32>
+     // CHECK: [[PTR:%.*]] = llvm.getelementptr [[BASE]][32]
+-    // CHECK: nvvm.tcgen05.ld [[PTR]] {{.*}}{shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<16xi32>
++    // CHECK: nvvm.tcgen05.ld [[PTR]] {{.*}}shape = shape_32x32b : vector<16xi32>
+     // CHECK: [[PTR:%.*]] = llvm.getelementptr [[BASE]][48]
+-    // CHECK: nvvm.tcgen05.ld [[PTR]] {{.*}}{shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<16xi32>
++    // CHECK: nvvm.tcgen05.ld [[PTR]] {{.*}}shape = shape_32x32b : vector<16xi32>
+     // CHECK-NOT: tcgen05.ld
+     // CHECK: ttg.warp_return
+     ttng.tmem_load %arg0 : !ttg.memdesc<128x64xf32, #tmem, #ttng.tensor_memory> -> tensor<128x64xf32, #blocked>
+@@ -977,7 +977,7 @@
+ // CHECK-LABEL: @load_store_x1
+ tt.func @load_store_x1(%arg0: !ttg.memdesc<128x2xf16, #tmem_x1, #ttng.tensor_memory, mutable>) {
+   %true = arith.constant true
+-  // CHECK: [[LOAD:%.*]] = nvvm.tcgen05.ld {{.*}}shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<1xi32>
++  // CHECK: [[LOAD:%.*]] = nvvm.tcgen05.ld {{.*}}shape = shape_32x32b : vector<1xi32>
+   // CHECK: [[V:%.*]] = llvm.extractelement [[LOAD]]{{.*}} : vector<1xi32>
+   // CHECK: [[F:%.*]] = llvm.bitcast [[V]] : i32 to vector<2xf16>
+   // CHECK: [[E0:%.*]] = llvm.extractelement [[F]]{{.*}} : vector<2xf16>
+@@ -996,7 +996,7 @@
+ // CHECK-DAG: [[C1:%.*]] = llvm.mlir.constant(1 : i32)
+ tt.func @load_store_x1_unpacked(%arg0: !ttg.memdesc<128x2xf16, #tmem_x1_unpacked, #ttng.tensor_memory, mutable>) {
+   %true = arith.constant true
+-  // CHECK: [[LOAD:%.*]] = nvvm.tcgen05.ld {{.*}}shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<1xi32>
++  // CHECK: [[LOAD:%.*]] = nvvm.tcgen05.ld {{.*}}shape = shape_32x32b : vector<1xi32>
+   // CHECK: [[V:%.*]] = llvm.extractelement [[LOAD]]{{.*}} : vector<1xi32>
+   // CHECK: [[F:%.*]] = llvm.bitcast [[V]] : i32 to vector<2xf16>
+   // CHECK: extractelement [[F]][[[C0]] : i32]
+@@ -1294,7 +1294,7 @@
+ #tmem = #ttng.tensor_memory_encoding<blockM = 128, blockN = 128, colStride = 1>
+ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.shared = 65544 : i32, ttg.target = "cuda:103", ttg.tensor_memory_size = 128 : i32, "ttg.threads-per-warp" = 32 : i32} {
+   // CHECK-LABEL: @tensor_memory_ld_red_min
+-  // CHECK: nvvm.tcgen05.ld.red min {{.*}} {shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<128xf32>, f32
++  // CHECK: nvvm.tcgen05.ld.red min {{.*}} shape = shape_32x32b : vector<128xf32>, f32
+   // CHECK: tcgen05.wait <load>
+   tt.func public @tensor_memory_ld_red_min() {
+     %cst_0 = arith.constant dense<0.000000e+00> : tensor<128x128xf32, #blocked1>
+@@ -1430,7 +1430,7 @@
+ #tmem = #ttng.tensor_memory_encoding<blockM = 128, blockN = 128, colStride = 1>
+ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.shared = 65544 : i32, ttg.target = "cuda:103", ttg.tensor_memory_size = 128 : i32, "ttg.threads-per-warp" = 32 : i32} {
+   // CHECK-LABEL: @tensor_memory_ld_red_max
+-  // CHECK: nvvm.tcgen05.ld.red max {{.*}} {shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<128xf32>, f32
++  // CHECK: nvvm.tcgen05.ld.red max {{.*}} shape = shape_32x32b : vector<128xf32>, f32
+   // CHECK: tcgen05.wait <load>
+   tt.func public @tensor_memory_ld_red_max() {
+     %cst_0 = arith.constant dense<0.000000e+00> : tensor<128x128xf32, #blocked1>
+@@ -1448,7 +1448,7 @@
+ #tmem = #ttng.tensor_memory_encoding<blockM = 128, blockN = 128, colStride = 1>
+ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.shared = 65544 : i32, ttg.target = "cuda:103", ttg.tensor_memory_size = 128 : i32, "ttg.threads-per-warp" = 32 : i32} {
+   // CHECK-LABEL: @tensor_memory_ld_red_min_abs
+-  // CHECK: nvvm.tcgen05.ld.red min {{.*}} {abs, shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<128xf32>, f32
++  // CHECK: nvvm.tcgen05.ld.red min {{.*}} shape = shape_32x32b abs : vector<128xf32>, f32
+   // CHECK: tcgen05.wait <load>
+   tt.func public @tensor_memory_ld_red_min_abs() {
+     %cst_0 = arith.constant dense<0.000000e+00> : tensor<128x128xf32, #blocked1>
+@@ -1466,7 +1466,7 @@
+ #tmem = #ttng.tensor_memory_encoding<blockM = 128, blockN = 128, colStride = 1>
+ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.shared = 65544 : i32, ttg.target = "cuda:103", ttg.tensor_memory_size = 128 : i32, "ttg.threads-per-warp" = 32 : i32} {
+   // CHECK-LABEL: @tensor_memory_ld_red_max_nan
+-  // CHECK: nvvm.tcgen05.ld.red max {{.*}} {nan, shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<128xf32>, f32
++  // CHECK: nvvm.tcgen05.ld.red max {{.*}} shape = shape_32x32b nan : vector<128xf32>, f32
+   // CHECK: tcgen05.wait <load>
+   tt.func public @tensor_memory_ld_red_max_nan() {
+     %cst_0 = arith.constant dense<0.000000e+00> : tensor<128x128xf32, #blocked1>
+@@ -1484,7 +1484,7 @@
+ #tmem = #ttng.tensor_memory_encoding<blockM = 128, blockN = 128, colStride = 1>
+ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.shared = 65544 : i32, ttg.target = "cuda:103", ttg.tensor_memory_size = 128 : i32, "ttg.threads-per-warp" = 32 : i32} {
+   // CHECK-LABEL: @tensor_memory_ld_red_max_abs_nan
+-  // CHECK: nvvm.tcgen05.ld.red max {{.*}} {abs, nan, shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<128xf32>, f32
++  // CHECK: nvvm.tcgen05.ld.red max {{.*}} shape = shape_32x32b abs nan : vector<128xf32>, f32
+   // CHECK: tcgen05.wait <load>
+   tt.func public @tensor_memory_ld_red_max_abs_nan() {
+     %cst_0 = arith.constant dense<0.000000e+00> : tensor<128x128xf32, #blocked1>
+@@ -1504,7 +1504,7 @@
+ #tmem_8w = #ttng.tensor_memory_encoding<blockM = 128, blockN = 64, colStride = 1>
+ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 8 : i32, ttg.shared = 65544 : i32, ttg.target = "cuda:103", ttg.tensor_memory_size = 128 : i32, "ttg.threads-per-warp" = 32 : i32} {
+   // CHECK-LABEL: @tensor_memory_ld_red_min_8_warps
+-  // CHECK: nvvm.tcgen05.ld.red min {{.*}} {shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<64xf32>, f32
++  // CHECK: nvvm.tcgen05.ld.red min {{.*}} shape = shape_32x32b : vector<64xf32>, f32
+   // CHECK: tcgen05.wait <load>
+   tt.func public @tensor_memory_ld_red_min_8_warps() {
+     %cst_0 = arith.constant dense<0.000000e+00> : tensor<256x64xf32, #blocked_8w>
+@@ -1525,7 +1525,7 @@
+ #tmem_256N = #ttng.tensor_memory_encoding<blockM = 128, blockN = 256, colStride = 1>
+ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.shared = 65544 : i32, ttg.target = "cuda:103", ttg.tensor_memory_size = 128 : i32, "ttg.threads-per-warp" = 32 : i32} {
+   // CHECK-LABEL: @tensor_memory_ld_red_min_128x256_4_warps
+-  // CHECK-COUNT-4: nvvm.tcgen05.ld.red min {{.*}} {shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<64xf32>, f32
++  // CHECK-COUNT-4: nvvm.tcgen05.ld.red min {{.*}} shape = shape_32x32b : vector<64xf32>, f32
+   // CHECK-NOT: tcgen05.wait <load>
+   // CHECK-COUNT-3: llvm.intr.minnum
+   // CHECK: tcgen05.wait <load>
+@@ -1537,7 +1537,7 @@
+   }
+ 
+   // CHECK-LABEL: @tensor_memory_ld_red_max_128x256_4_warps
+-  // CHECK-COUNT-4: nvvm.tcgen05.ld.red max {{.*}} {shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<64xf32>, f32
++  // CHECK-COUNT-4: nvvm.tcgen05.ld.red max {{.*}} shape = shape_32x32b : vector<64xf32>, f32
+   // CHECK-COUNT-3: llvm.intr.maxnum
+   // CHECK: tcgen05.wait <load>
+   tt.func private @tensor_memory_ld_red_max_128x256_4_warps() -> tensor<128xf32, #blocked_red_256N_4w> {
+@@ -1557,7 +1557,7 @@
+ #tmem_256N_nan = #ttng.tensor_memory_encoding<blockM = 128, blockN = 256, colStride = 1>
+ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.shared = 65544 : i32, ttg.target = "cuda:103", ttg.tensor_memory_size = 128 : i32, "ttg.threads-per-warp" = 32 : i32} {
+   // CHECK-LABEL: @tensor_memory_ld_red_min_128x256_4_warps_nan
+-  // CHECK-COUNT-4: nvvm.tcgen05.ld.red min {{.*}} {nan, shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<64xf32>, f32
++  // CHECK-COUNT-4: nvvm.tcgen05.ld.red min {{.*}} shape = shape_32x32b nan : vector<64xf32>, f32
+   // CHECK-COUNT-3: llvm.intr.minimum
+   // CHECK: tcgen05.wait <load>
+   tt.func private @tensor_memory_ld_red_min_128x256_4_warps_nan() -> tensor<128xf32, #blocked_red_256N_4w_nan> {
+@@ -1568,7 +1568,7 @@
+   }
+ 
+   // CHECK-LABEL: @tensor_memory_ld_red_max_128x256_4_warps_nan
+-  // CHECK-COUNT-4: nvvm.tcgen05.ld.red max {{.*}} {nan, shape = #nvvm.tcgen05_ldst_shape<shape_32x32b>} : vector<64xf32>, f32
++  // CHECK-COUNT-4: nvvm.tcgen05.ld.red max {{.*}} shape = shape_32x32b nan : vector<64xf32>, f32
+   // CHECK-COUNT-3: llvm.intr.maximum
+   // CHECK: tcgen05.wait <load>
+   tt.func private @tensor_memory_ld_red_max_128x256_4_warps_nan() -> tensor<128xf32, #blocked_red_256N_4w_nan> {
+
+--- a/test/Conversion/tritongpu_to_llvm_blackwell_256bit.mlir
++++ b/test/Conversion/tritongpu_to_llvm_blackwell_256bit.mlir
+@@ -100,16 +100,16 @@
+ #blocked_fp8 = #ttg.blocked<{sizePerThread = [2], threadsPerWarp = [32], warpsPerCTA = [1], order = [0]}>
+ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+   // FP8-LABEL: @bf16_to_fp8e4
+-  // PACKED: nvvm.convert.bf16x2.to.f8x2 {{.*}} {rnd = #nvvm.fp_rnd_mode<rn>, sat = #nvvm.sat_mode<satfinite>} : vector<2xbf16> -> vector<2xi8>(f8E4M3FN)
+-  // LEGACY: nvvm.convert.f32x2.to.f8x2 {{.*}} {rnd = #nvvm.fp_rnd_mode<rn>, sat = #nvvm.sat_mode<satfinite>} : vector<2xi8>(f8E4M3FN)
++  // PACKED: nvvm.convert.bf16x2.to.f8x2 {{.*}} rnd = <rn> sat = <satfinite> : vector<2xbf16> -> vector<2xi8>(f8E4M3FN)
++  // LEGACY: nvvm.convert.f32x2.to.f8x2 {{.*}} rnd = <rn> sat = <satfinite> : vector<2xi8>(f8E4M3FN)
+   tt.func private @bf16_to_fp8e4(%in: tensor<64xbf16, #blocked_fp8>) -> tensor<64xf8E4M3FN, #blocked_fp8> {
+     %out = tt.fp_to_fp %in, rounding = rtne : tensor<64xbf16, #blocked_fp8> -> tensor<64xf8E4M3FN, #blocked_fp8>
+     tt.return %out : tensor<64xf8E4M3FN, #blocked_fp8>
+   }
+ 
+   // FP8-LABEL: @bf16_to_fp8e5
+-  // PACKED: nvvm.convert.bf16x2.to.f8x2 {{.*}} {rnd = #nvvm.fp_rnd_mode<rn>, sat = #nvvm.sat_mode<satfinite>} : vector<2xbf16> -> vector<2xi8>(f8E5M2)
+-  // LEGACY: nvvm.convert.f32x2.to.f8x2 {{.*}} {rnd = #nvvm.fp_rnd_mode<rn>, sat = #nvvm.sat_mode<satfinite>} : vector<2xi8>(f8E5M2)
++  // PACKED: nvvm.convert.bf16x2.to.f8x2 {{.*}} rnd = <rn> sat = <satfinite> : vector<2xbf16> -> vector<2xi8>(f8E5M2)
++  // LEGACY: nvvm.convert.f32x2.to.f8x2 {{.*}} rnd = <rn> sat = <satfinite> : vector<2xi8>(f8E5M2)
+   tt.func private @bf16_to_fp8e5(%in: tensor<64xbf16, #blocked_fp8>) -> tensor<64xf8E5M2, #blocked_fp8> {
+     %out = tt.fp_to_fp %in, rounding = rtne : tensor<64xbf16, #blocked_fp8> -> tensor<64xf8E5M2, #blocked_fp8>
+     tt.return %out : tensor<64xf8E5M2, #blocked_fp8>
+
+--- a/test/Conversion/tritongpu_to_llvm_hopper.mlir
++++ b/test/Conversion/tritongpu_to_llvm_hopper.mlir
+@@ -224,9 +224,9 @@
+     // CHECK-COUNT-2: nvvm.convert.f16x2.to.f8x2 {{.*}} : vector<2xf16> -> vector<2xi8>(f8E4M3FN)
+     %out4 = tt.fp_to_fp %in2, rounding = rtne : tensor<128xf16, #blocked> -> tensor<128xf8E4M3FN, #blocked>
+ 
+-    // CHECK-COUNT-2: nvvm.convert.f32x2.to.f8x2 {{.*}} {rnd = #nvvm.fp_rnd_mode<rn>, sat = #nvvm.sat_mode<satfinite>} : vector<2xi8>(f8E5M2)
++    // CHECK-COUNT-2: nvvm.convert.f32x2.to.f8x2 {{.*}} rnd = <rn> sat = <satfinite> : vector<2xi8>(f8E5M2)
+     %out5 = tt.fp_to_fp %in3, rounding = rtne : tensor<128xf32, #blocked> -> tensor<128xf8E5M2, #blocked>
+-    // CHECK-COUNT-2: nvvm.convert.f32x2.to.f8x2 {{.*}} {rnd = #nvvm.fp_rnd_mode<rn>, sat = #nvvm.sat_mode<satfinite>} : vector<2xi8>(f8E4M3FN)
++    // CHECK-COUNT-2: nvvm.convert.f32x2.to.f8x2 {{.*}} rnd = <rn> sat = <satfinite> : vector<2xi8>(f8E4M3FN)
+     %out6 = tt.fp_to_fp %in3, rounding = rtne : tensor<128xf32, #blocked> -> tensor<128xf8E4M3FN, #blocked>
+     tt.return %out0, %out1, %out2, %out3, %out4, %out5, %out6 :
+         tensor<128xf16, #blocked>,
+@@ -686,7 +686,7 @@
+ 
+     %3 = ttg.local_load %1 : !ttg.memdesc<16x16xf64, #shared1, #smem, mutable> -> tensor<16x16xf64, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 1}>>
+ 
+-    // CHECK-COUNT-2: nvvm.mma.sync {{.*}}layoutA = #nvvm.mma_layout<row>, layoutB = #nvvm.mma_layout<col>{{.*}}shape = #nvvm.shape<m = 8, n = 8, k = 4>} : (f64, f64, f64)
++    // CHECK-COUNT-2: nvvm.mma.sync {{.*}}shape = <m = 8, n = 8, k = 4>, layout_a = row, layout_b = col : (f64, f64, f64)
+ 
+     %out = tt.dot %2, %3, %cst, inputPrecision = tf32 : tensor<16x16xf64, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 1}>> * tensor<16x16xf64, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 1}>> -> tensor<16x16xf64, #mma>
+ 
+
+--- a/test/Conversion/tritongpu_to_llvm_rubin.mlir
++++ b/test/Conversion/tritongpu_to_llvm_rubin.mlir
+@@ -489,13 +489,13 @@
+   // CHECK: nvvm.addf {{.*}} : vector<2xf16>
+   // CHECK: nvvm.subf {{.*}} : vector<2xf16>
+   // CHECK: llvm.inline_asm {{.*}} "mul.f16x2 $0, $1, $2;", "=r,r,r"
+-  // CHECK: nvvm.fma {{.*}} {rnd = #nvvm.fp_rnd_mode<rn>} : vector<2xf16>
++  // CHECK: nvvm.fma {{.*}} rnd = <rn> : vector<2xf16>
+   // CHECK: llvm.call_intrinsic "llvm.nvvm.fmin.f16x2"
+   // CHECK: llvm.call_intrinsic "llvm.nvvm.fmax.f16x2"
+   // CHECK: nvvm.addf {{.*}} : vector<2xbf16>
+   // CHECK: nvvm.subf {{.*}} : vector<2xbf16>
+   // CHECK: llvm.inline_asm {{.*}} "mul.bf16x2 $0, $1, $2;", "=r,r,r"
+-  // CHECK: nvvm.fma {{.*}} {rnd = #nvvm.fp_rnd_mode<rn>} : vector<2xbf16>
++  // CHECK: nvvm.fma {{.*}} rnd = <rn> : vector<2xbf16>
+   // CHECK: llvm.call_intrinsic "llvm.nvvm.fmin.bf16x2"
+   // CHECK: llvm.call_intrinsic "llvm.nvvm.fmax.bf16x2"
+   tt.func private @packed_arith_homogeneous_half(
+@@ -573,7 +573,7 @@
+ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:107", "ttg.threads-per-warp" = 32 : i32} {
+   // CHECK-LABEL: @packed_arith_mixed
+   // CHECK: llvm.inline_asm {{.*}} "mul.f16x2 $0, $1, $2;", "=r,r,r" {{.*}} : (i32, i32) -> i32
+-  // CHECK: nvvm.fma {{.*}} {rnd = #nvvm.fp_rnd_mode<rn>} : vector<2xbf16>
++  // CHECK: nvvm.fma {{.*}} rnd = <rn> : vector<2xbf16>
+   // CHECK: llvm.inline_asm {{.*}} "add.f32x2.f16x2.f32x2 $0, $1, $2;", "=l,r,l" {{.*}} : (i32, i64) -> i64
+   // CHECK: llvm.inline_asm {{.*}} "add.f32x2.bf16x2.f32x2 $0, $1, $2;", "=l,r,l" {{.*}} : (i32, i64) -> i64
+   // CHECK: llvm.inline_asm {{.*}} "sub.f32x2.f16x2.f32x2 $0, $1, $2;", "=l,r,l" {{.*}} : (i32, i64) -> i64
 
 --- a/test/Conversion/warp_specialize_to_llvm.mlir
 +++ b/test/Conversion/warp_specialize_to_llvm.mlir
-@@ -212,7 +212,7 @@ llvm.func @generate_switch_loop() attributes {allocation.offset = 32 : i32} {
+@@ -197,7 +197,7 @@
+   // CHECK-NEXT: llvm.cond_br [[IS_DEFAULT]], [[BODY:\^.*]], [[SWITCH_LOOP:\^.*]]
+ 
+   // CHECK: [[SWITCH_LOOP]]:
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: [[SMEM_BASE:%.*]] = llvm.getelementptr [[SMEM_ADDR]][32] : (!llvm.ptr<3>) -> !llvm.ptr<3>, i8
+   // CHECK-NEXT: [[REL_WID:%.*]] = llvm.sub [[WARP_ID]], [[C4]]
+ 
+@@ -210,29 +210,29 @@
+   // CHECK-NEXT: 3: [[EXIT:\^.*]]
+ 
    // CHECK: [[DEFAULT]]:
-   // CHECK-NEXT: "llvm.nvvm.barrier.cta.sync.all"([[C1]])
-   // CHECK-NEXT: "llvm.nvvm.barrier.cta.sync.all"([[C1]])
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
 -  // CHECK-NEXT: llvm.br [[SWITCH_LOOP]] {loop_annotation = #llvm.loop_annotation<licm = <disable = true>>}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
 +  // CHECK-NEXT: llvm.br [[SWITCH_LOOP]] loop_annotation = <licm = <disable = true>>
  
    // CHECK: [[EXIT]]:
    // CHECK-NEXT: llvm.return
+ 
+   // CHECK: [[PARTITION0]]:
+-  // CHECK: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: "partition0"
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: llvm.br [[SWITCH_LOOP]]
+ 
+   // CHECK: [[PARTITION1]]:
+-  // CHECK: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: "partition1"
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: llvm.br [[SWITCH_LOOP]]
+ 
+   // CHECK: [[PARTITION2]]:
+-  // CHECK: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: "partition2"
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: llvm.br [[SWITCH_LOOP]]
+ 
+   // CHECK: [[BODY]]:
+@@ -255,12 +255,12 @@
+   // CHECK-NEXT: [[PTR:%.*]] = llvm.getelementptr [[SMEM_BASE]][6]
+   // CHECK-NEXT: llvm.store [[C2_i8]], [[PTR]]
+ 
+-  // CHECK: nvvm.barrier id = [[C1]] {aligned = false}
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK: nvvm.barrier id = [[C1]] aligned = false
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: llvm.br [[DEFAULT_PARTITION:\^.*]]
+   // CHECK: [[DEFAULT_PARTITION]]:
+   // CHECK-NEXT: "default"
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: llvm.br [[AFTER:\^.*]]
+ 
+   // AMD: [[WID:%.*]] = rocdl.wave.id : i32
 @@ -283,7 +283,7 @@ llvm.func @generate_switch_loop() attributes {allocation.offset = 32 : i32} {
    // AMD: [[DEFAULT]]:
    // AMD-NEXT: rocdl.barrier
@@ -135,7 +691,16 @@
  
    // AMD: [[EXIT]]:
    // AMD-NEXT: llvm.return
-@@ -414,9 +414,9 @@ llvm.func @pass_captures() attributes {allocation.offset = 32 : i32} {
+@@ -370,7 +370,7 @@
+   // CHECK-NEXT: llvm.store [[C3_i8]], [[PTR]]
+   // CHECK-NEXT: [[PTR:%.*]] = llvm.getelementptr [[SMEM_BASE]][6]
+   // CHECK-NEXT: llvm.store [[C3_i8]], [[PTR]]
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: llvm.return
+ 
+   // AMD: [[AFTER:\^bb[0-9]+]]:
+@@ -414,36 +414,36 @@ llvm.func @pass_captures() attributes {allocation.offset = 32 : i32} {
  
    // CHECK: ^bb4:
    // CHECK-NEXT: [[ARG0_PTR:%.*]] = llvm.getelementptr [[SMEM_ADDR]][0, 0] : (!llvm.ptr<3>) -> !llvm.ptr<3>, !llvm.struct<packed (i32, i64)>
@@ -143,11 +708,13 @@
 +  // CHECK-NEXT: [[ARG0:%.*]] = llvm.load [[ARG0_PTR]] <alignment = 1> : !llvm.ptr<3> -> i32
    // CHECK-NEXT: [[ARG1_PTR:%.*]] = llvm.getelementptr [[SMEM_ADDR]][0, 1] : (!llvm.ptr<3>) -> !llvm.ptr<3>, !llvm.struct<packed (i32, i64)>
 -  // CHECK-NEXT: [[ARG1:%.*]] = llvm.load [[ARG1_PTR]] {alignment = 1 : i64}
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
 +  // CHECK-NEXT: [[ARG1:%.*]] = llvm.load [[ARG1_PTR]] <alignment = 1> : !llvm.ptr<3> -> i64
-   // CHECK-NEXT: "llvm.nvvm.barrier.cta.sync.all"([[C1]])
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
    // CHECK-NEXT: "use"([[ARG0]], [[ARG1]])
-   // CHECK-NEXT: "llvm.nvvm.barrier.cta.sync.all"([[C1]])
-@@ -424,17 +424,17 @@ llvm.func @pass_captures() attributes {allocation.offset = 32 : i32} {
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+ 
    // CHECK: ^bb5:
    // CHECK: [[INS:%.*]]:2 = "produce"()
    // CHECK: [[ARG0_PTR:%.*]] = llvm.getelementptr [[SMEM_ADDR]][0, 0] : (!llvm.ptr<3>) -> !llvm.ptr<3>, !llvm.struct<packed (i32, i64)>
@@ -155,9 +722,11 @@
 +  // CHECK-NEXT: llvm.store [[INS]]#0, [[ARG0_PTR]] <alignment = 1> : i32, !llvm.ptr<3>
    // CHECK-NEXT: [[ARG1_PTR:%.*]] = llvm.getelementptr [[SMEM_ADDR]][0, 1] : (!llvm.ptr<3>) -> !llvm.ptr<3>, !llvm.struct<packed (i32, i64)>
 -  // CHECK-NEXT: llvm.store [[INS]]#1, [[ARG1_PTR]] {alignment = 1 : i64}
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
 +  // CHECK-NEXT: llvm.store [[INS]]#1, [[ARG1_PTR]] <alignment = 1> : i64, !llvm.ptr<3>
-   // CHECK-NEXT: "llvm.nvvm.barrier.cta.sync.all"([[C1]])
-   // CHECK-NEXT: "llvm.nvvm.barrier.cta.sync.all"([[C1]])
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
  
    // AMD: ^bb4:
    // AMD-NEXT: [[ARG0_PTR:%.*]] = llvm.getelementptr [[SMEM_ADDR]][0, 0] : (!llvm.ptr<3>) -> !llvm.ptr<3>, !llvm.struct<packed (i32, i64)>
@@ -169,7 +738,7 @@
    // AMD-NEXT: rocdl.barrier
    // AMD-NEXT: "use"([[ARG0]], [[ARG1]])
    // AMD-NEXT: rocdl.barrier
-@@ -442,9 +442,9 @@ llvm.func @pass_captures() attributes {allocation.offset = 32 : i32} {
+ 
    // AMD: ^bb5:
    // AMD: [[INS:%.*]]:2 = "produce"()
    // AMD: [[ARG0_PTR:%.*]] = llvm.getelementptr [[SMEM_ADDR]][0, 0] : (!llvm.ptr<3>) -> !llvm.ptr<3>, !llvm.struct<packed (i32, i64)>
@@ -180,7 +749,218 @@
 +  // AMD-NEXT: llvm.store [[INS]]#1, [[ARG1_PTR]] <alignment = 1> : i64, !llvm.ptr<3>
    // AMD-NEXT: rocdl.barrier
    // AMD-NEXT: rocdl.barrier
+@@ -642,7 +642,7 @@
+   // CHECK-NEXT: llvm.store [[C2_i8]], [[PTR]]
+   // CHECK-NEXT: [[PTR:%.*]] = llvm.getelementptr %{{[0-9]+}}[7]
+   // CHECK-NEXT: llvm.store [[Cn1_i8]], [[PTR]]
+-  // CHECK: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK: "ws0_default"
  
+   // AMD: llvm.switch
+@@ -720,7 +720,7 @@
+   // CHECK-NEXT: llvm.store [[C3_i8]], [[PTR]]
+   // CHECK-NEXT: [[PTR:%.*]] = llvm.getelementptr %{{[0-9]+}}[7]
+   // CHECK-NEXT: llvm.store [[C3_i8]], [[PTR]]
+-  // CHECK: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK: "ws1_default"
+ 
+   // AMD: getelementptr
+@@ -772,7 +772,7 @@
+   // CHECK-NEXT: llvm.store [[Cn1_i8]], [[PTR]]
+   // CHECK-NEXT: [[PTR:%.*]] = llvm.getelementptr %{{[0-9]+}}[7]
+   // CHECK-NEXT: llvm.store [[Cn1_i8]], [[PTR]]
+-  // CHECK: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK: "ws2_default"
+ 
+   // AMD: getelementptr
+@@ -816,7 +816,7 @@
+   // CHECK-NEXT: llvm.store [[C5_i8]], [[PTR]]
+   // CHECK-NEXT: [[PTR:%.*]] = llvm.getelementptr %{{[0-9]+}}[7]
+   // CHECK-NEXT: llvm.store [[C5_i8]], [[PTR]]
+-  // CHECK: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK: "ws3_default"
+ 
+   // AMD: getelementptr
+@@ -868,29 +868,29 @@
+   // COMMON-NEXT: 1: [[EXIT:\^.*]]
+ 
+   // CHECK: [[PARTITION]]:
+-  // CHECK: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: "something"()[[[A:\^.*]], [[B:\^.*]]]
+   // CHECK: [[A]]:
+   // CHECK-NEXT: "A"
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: llvm.br [[SWITCH_LOOP]]
+   // CHECK: [[B]]:
+   // CHECK-NEXT: "B"
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: llvm.br [[SWITCH_LOOP]]
+ 
+-  // CHECK: nvvm.barrier id = [[C1]] {aligned = false}
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK: nvvm.barrier id = [[C1]] aligned = false
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK: llvm.br [[DEFAULT:\^.*]]
+   // CHECK: [[DEFAULT]]:
+   // CHECK-NEXT: "something"()[[[A:\^.*]], [[B:\^.*]]]
+   // CHECK: [[A]]:
+   // CHECK-NEXT: "A"
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: llvm.br [[AFTER:\^.*]]
+   // CHECK: [[B]]:
+   // CHECK-NEXT: "B"
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: llvm.br [[AFTER]]
+ 
+   // AMD: [[PARTITION]]:
+@@ -1045,13 +1045,13 @@
+   }
+   partition0(%arg0: i32, %arg1: !llvm.ptr<3>) num_warps(1) {
+   // CHECK: ^bb4:
+-    // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++    // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+     // CHECK-NEXT: "use"([[CAP0]], [[CAP1]])
+   // AMD: ^bb4:
+     // AMD-NEXT: rocdl.barrier
+     // AMD-NEXT: "use"([[CAP0]], [[CAP1]])
+     "use"(%arg0, %arg1) : (i32, !llvm.ptr<3>) -> ()
+-    // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++    // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+     // AMD-NEXT: rocdl.barrier
+     ttg.warp_return
+   } : (i32, !llvm.ptr<3>) -> ()
+@@ -1074,7 +1074,7 @@
+   }
+   partition0(%arg2: !llvm.ptr<3>, %arg3: i32) num_warps(1) {
+   // CHECK: ^bb4:
+-    // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++    // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+     // CHECK-NEXT: [[ADD:%.*]] = llvm.add %arg0, %arg1 : i32
+     // CHECK-NEXT: [[MUL:%.*]] = llvm.mul [[ADD]], %arg1 : i32
+     // CHECK-NEXT: [[UREM:%.*]] = llvm.urem [[ADD]], [[MUL]] : i32
+@@ -1088,7 +1088,7 @@
+     // AMD-NEXT: [[PTR:%.*]] = llvm.getelementptr [[ADDR]][%arg0]
+     // AMD-NEXT: "use"([[PTR]], [[UREM]])
+     "use"(%arg2, %arg3) : (!llvm.ptr<3>, i32) -> ()
+-    // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++    // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+     // AMD-NEXT: rocdl.barrier
+     ttg.warp_return
+   } : (!llvm.ptr<3>, i32) -> ()
+@@ -1111,7 +1111,7 @@
+ 
+   // CHECK: [[SWITCH_LOOP]]:
+   // CHECK-NEXT: nvvm.setmaxregister decrease 24
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK: llvm.switch
+   // CHECK-NEXT: 0: [[PARTITION0:\^.*]],
+   // CHECK-NEXT: 1: [[PARTITION1:\^.*]],
+@@ -1120,33 +1120,33 @@
+ 
+   // CHECK: [[PARTITION0]]:
+   // CHECK-NEXT: nvvm.setmaxregister increase 80
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: "partition0"()
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: nvvm.setmaxregister decrease 24
+ 
+   // CHECK: [[PARTITION1]]:
+   // CHECK-NEXT: nvvm.setmaxregister increase 48
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: "partition1"()
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: nvvm.setmaxregister decrease 24
+ 
+   // CHECK: [[PARTITION2]]:
+   // CHECK-NEXT: nvvm.setmaxregister increase 128
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: "partition2"()
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: nvvm.setmaxregister decrease 24
+ 
+   // CHECK: [[ENTRY]]:
+   // CHECK-NEXT: nvvm.setmaxregister increase 248
+ 
+-  // CHECK: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: setmaxregister decrease 152
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK: "default"
+-  // CHECK: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: setmaxregister increase 248
+ 
+   ttg.warp_specialize() attributes {allocation.offset = 0 : i32, warpGroupStartIds = array<i32: 4, 8, 12>, actualRegisters = array<i32: 152, 80, 48, 128>}
+@@ -1185,7 +1185,7 @@
+ 
+   // CHECK: [[SWITCH_LOOP]]:
+   // CHECK-NEXT: nvvm.setmaxregister decrease 80
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK: llvm.switch
+   // CHECK-NEXT: 0: [[PARTITION0:\^.*]],
+   // CHECK-NEXT: 1: [[PARTITION1:\^.*]],
+@@ -1194,33 +1194,33 @@
+ 
+   // CHECK: [[PARTITION0]]:
+   // CHECK-NEXT: nvvm.setmaxregister decrease 24
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: "partition0"()
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: nvvm.setmaxregister increase 80
+ 
+   // CHECK: [[PARTITION1]]:
+   // CHECK-NEXT: nvvm.setmaxregister increase 192
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: "partition1"()
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: nvvm.setmaxregister decrease 80
+ 
+   // CHECK: [[PARTITION2]]:
+   // CHECK-NEXT: nvvm.setmaxregister increase 192
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: "partition2"()
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: nvvm.setmaxregister decrease 80
+ 
+   // CHECK: [[ENTRY]]:
+   // CHECK-NEXT: nvvm.setmaxregister increase 256
+ 
+-  // CHECK: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: setmaxregister decrease 104
+-  // CHECK-NEXT: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK-NEXT: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK: "default"
+-  // CHECK: nvvm.barrier id = [[C1]] {aligned = false}
++  // CHECK: nvvm.barrier id = [[C1]] aligned = false
+   // CHECK-NEXT: setmaxregister increase 256
+ 
+   ttg.warp_specialize() attributes {allocation.offset = 0 : i32, warpGroupStartIds = array<i32: 4, 8, 12>, actualRegisters = array<i32: 104, 24, 192, 192>}
+
 
 --- a/third_party/amd/lib/TritonAMDGPUToLLVM/ConvertWarpSpecializeToLLVM.cpp
 +++ b/third_party/amd/lib/TritonAMDGPUToLLVM/ConvertWarpSpecializeToLLVM.cpp
diff --git a/third_party/xla/third_party/triton/common/llvm_cl991819890.patch b/third_party/xla/third_party/triton/common/llvm_cl991819890.patch
index 791e53a811b..5d0644dc289 100644
--- a/third_party/xla/third_party/triton/common/llvm_cl991819890.patch
+++ b/third_party/xla/third_party/triton/common/llvm_cl991819890.patch
@@ -1,6 +1,6 @@
 --- a/python/src/llvm.cc
 +++ b/python/src/llvm.cc
-@@ -753,7 +753,7 @@ void init_triton_llvm(py::module_ &m) {
+@@ -788,7 +788,7 @@ void init_triton_llvm(py::module_ &m) {
          if (!pluginFile.empty()) {
            // TODO: Add some logging here that we inserted a pass into the LLVM
            // pass pipeline
@@ -12,7 +12,7 @@
 
 --- a/python/test/unit/language/test_compile_only.py
 +++ b/python/test/unit/language/test_compile_only.py
-@@ -99,8 +99,7 @@ def test_compile_only_mulhi_survives_llvm_opt(bits) -> None:
+@@ -99,8 +99,7 @@ def test_umulhi_truncated_input(seed_type, dtype):
      for kind, stage in stages.items():
          module = stage(module, metadata)
          if kind == "llir":
@@ -46,7 +46,7 @@
 
 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/ElementwiseOpToLLVM.cpp
-@@ -949,8 +949,7 @@ struct MulhiUIOpConversion
+@@ -871,8 +871,7 @@ struct MulhiUIOpConversion
                                     Location loc) const {
      unsigned bitWidth = elemTy.getIntOrFloatBitWidth();
      assert(bitWidth == 32 || bitWidth == 64);
diff --git a/third_party/xla/third_party/triton/common/llvm_cl996124567.patch b/third_party/xla/third_party/triton/common/llvm_cl996124567.patch
index b19f3c07ff9..d06dd10426b 100644
--- a/third_party/xla/third_party/triton/common/llvm_cl996124567.patch
+++ b/third_party/xla/third_party/triton/common/llvm_cl996124567.patch
@@ -1,6 +1,6 @@
 --- a/test/Conversion/tritongpu_to_ptx_mmav3.mlir
 +++ b/test/Conversion/tritongpu_to_ptx_mmav3.mlir
-@@ -63,7 +63,8 @@
+@@ -63,7 +63,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
      // CHECK-DAG: and.b32 [[LANEID_FIXED_BITS:%.*]], [[TID]], 28;
      // CHECK-DAG: shl.b32 [[L0_TEMP:%.*]], [[L0_VAL]], 1;
      // CHECK-DAG: or.b32  [[LANEID_PART_PERM:%.*]], [[L0_TEMP]], [[LANEID_FIXED_BITS]];
@@ -10,7 +10,7 @@
      // CHECK-DAG: or.b32  [[LANEID_PERM:%.*]], [[LANEID_PART_PERM]], [[L1_TEMP]];
  
      // The index shuffles have source lane dependent on the value of the r2 bit.
-@@ -97,7 +98,6 @@
+@@ -97,7 +98,6 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
      // on, we have selectors 0x7632 and 0x3276. These are 21520, 4180, 30258,
      // and 12918 in decimal, respectively.
  
@@ -31,9 +31,9 @@
  #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 +#include "mlir/Dialect/LLVMIR/ROCDLTargetInfo.h"
  #include "mlir/Pass/Pass.h"
- #include "third_party/amd/include/Analysis/AMDGPUAllocation.h"
  #include "third_party/amd/include/Analysis/AxisInfoExt.h"
-@@ -213,16 +213,16 @@
+ #include "third_party/amd/include/Dialect/TritonAMDGPU/IR/Dialect.h"
+@@ -203,16 +203,16 @@ struct ConvertTritonAMDGPUToLLVM
      mlir::triton::AMD::populateWarpIdOpToLLVMPattern(typeConverter, targetInfo,
                                                       patterns, commonBenefit);
  
@@ -68,7 +68,7 @@
  #include "mlir/Pass/Pass.h"
  #include "third_party/amd/lib/TritonAMDGPUToLLVM/PatternTritonGPUOpToLLVM.h"
  #include "triton/Conversion/TritonGPUToLLVM/TypeConverter.h"
-@@ -65,15 +65,15 @@
+@@ -65,15 +65,15 @@ struct ConvertProtonAMDGPUToLLVM
                                                         tritonTargetInfo);
      mlir::arith::populateArithToLLVMConversionPatterns(typeConverter, patterns);
  
diff --git a/third_party/xla/third_party/triton/common/mig_limits_pytests.patch b/third_party/xla/third_party/triton/common/mig_limits_pytests.patch
index 89f926edc5b..dc27248e243 100644
--- a/third_party/xla/third_party/triton/common/mig_limits_pytests.patch
+++ b/third_party/xla/third_party/triton/common/mig_limits_pytests.patch
@@ -1,10 +1,10 @@
 --- a/python/test/conftest.py
 +++ b/python/test/conftest.py
-@@ -2,7 +2,18 @@ import pytest
+@@ -2,6 +2,18 @@ import pytest
  
  
  def pytest_configure(config):
--    config.addinivalue_line("markers", "gsan_fine_granularity(reason): requires sub-16-byte GSan tracking")
+     config.addinivalue_line("markers", "gsan_fine_granularity(reason): requires sub-16-byte GSan tracking")
 +    # In internal hermetic GPU test environments (e.g. a100, h100, b200 test targets), Forge enforces go/forge-accel-or-fail
 +    # to verify that GPU test binaries actively utilize the GPU. For pure AST/compiler tests (e.g. test_cache_determinism),
 +    # lazy CUDA initialization causes 0 GPU usage and fails the runner. We actively guarantee GPU utilization during pytest setup.
@@ -23,7 +23,7 @@
 
 --- a/python/test/unit/language/test_core.py
 +++ b/python/test/unit/language/test_core.py
-@@ -2351,7 +2351,13 @@ def test_atomic_cas(sem, num_ctas, dtype_str, device):
+@@ -2377,7 +2377,13 @@ def test_atomic_cas(sem, num_ctas, dtype_str, device):
  
  @pytest.mark.skipif(not is_cuda() or torch.cuda.get_device_capability()[0] < 9,
                      reason="num_ctas > 1 requires NVIDIA SM90+ (Hopper)")
@@ -38,7 +38,7 @@
  def test_scalar_atomic_cas_multicta_result(device):
  
      @triton.jit
-@@ -2367,8 +2373,8 @@ def test_scalar_atomic_cas_multicta_result(device):
+@@ -2393,8 +2399,8 @@ def test_scalar_atomic_cas_multicta_result(device):
      torch.testing.assert_close(output, torch.full_like(output, 7))
  
  
@@ -49,7 +49,7 @@
  @pytest.mark.parametrize("size", [1, 4, 16, 128, 512])
  def test_tensor_atomic_cas_multicta_result(size, device):
  
-@@ -2419,8 +2425,8 @@ def test_tensor_atomic_cas(sem, size, dtype_str, num_ctas, device):
+@@ -2445,8 +2451,8 @@ def test_tensor_atomic_cas(sem, size, dtype_str, num_ctas, device):
  
  
  @pytest.mark.interpreter
diff --git a/third_party/xla/third_party/triton/common/old_ptxas.patch b/third_party/xla/third_party/triton/common/old_ptxas.patch
index 40bb19ea4da..1081ddb4382 100644
--- a/third_party/xla/third_party/triton/common/old_ptxas.patch
+++ b/third_party/xla/third_party/triton/common/old_ptxas.patch
@@ -31,6 +31,67 @@ On Blackwell architectures with PTX 8.8+, upstream Triton generates 256-bit glob
    // PRE_BW-LABEL: global_load_v4_b64
    // PRE_BW-NOT: ld.global.v4.b64
    // PRE_BW: ld.global.v2.b64
+@@ -156,10 +156,10 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 1 : i32} {
+ 
+   // FP4-LABEL: @fp4_to_fp16
+   // FP4-NOT: cvt.rn.bf16x2.e2m1x2
+-  // FP16-NATIVE-COUNT-4: nvvm.convert.f4x2.to.f16x2 {{.*}} : i8(f4E2M1FN) -> vector<2xf16>
+-  // FP16-LEGACY-NOT: nvvm.convert.f4x2.to.f16x2
++  // FP16-NATIVE-COUNT-4: cvt.rn.f16x2.e2m1x2
++  // FP16-LEGACY-NOT: cvt.rn.f16x2.e2m1x2
+   // FP16-LEGACY: cvt.rn.f16x2.e4m3x2
+-  // FP16-LEGACY-NOT: nvvm.convert.f4x2.to.f16x2
++  // FP16-LEGACY-NOT: cvt.rn.f16x2.e2m1x2
+   // FP4-NOT: cvt.rn.bf16x2.e2m1x2
+   tt.func private @fp4_to_fp16(%in: tensor<128xi8, #packed_fp4>) -> tensor<256xf16, #unpacked_fp4> {
+     %out = ttg.fp4_to_fp %in {axis = 0 : i32} : tensor<128xi8, #packed_fp4> -> tensor<256xf16, #unpacked_fp4>
+
+--- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/Fp4ToFpOpToLLVM.cpp
++++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/Fp4ToFpOpToLLVM.cpp
+@@ -44,6 +44,16 @@ static constexpr const char *FP4ToBf16NativePtx =
+     "cvt.rn.bf16x2.e2m1x2 $3, b3;\n"
+     "}";
+ 
++static constexpr const char *FP4ToFp16NativePtx =
++    "{\n"
++    ".reg .b8 b<4>;\n"
++    "mov.b32 {b0, b1, b2, b3}, $4;\n"
++    "cvt.rn.f16x2.e2m1x2 $0, b0;\n"
++    "cvt.rn.f16x2.e2m1x2 $1, b1;\n"
++    "cvt.rn.f16x2.e2m1x2 $2, b2;\n"
++    "cvt.rn.f16x2.e2m1x2 $3, b3;\n"
++    "}";
++
+ static constexpr const char *FP4ToFP16Ptx =
+     "{\n"
+     ".reg .b32           a<11>;\n"
+@@ -100,23 +110,10 @@ class Fp4ToFpOpPattern : public Fp4ToFpOpConversionBase {
+     bool toFp16 = elemType == f16_ty;
+     auto b = TritonLLVMOpBuilder(loc, rewriter);
+     std::array<Value, 8> results;
+-    if (toFp16 && targetInfo.getComputeCapability() >= 100 &&
+-        targetInfo.getPtxVersion() >= 86) {
+-      Value bytes = b.bitcast(packedVec, vec_ty(i8_ty, 4));
+-      for (int i = 0; i < 4; ++i) {
+-        Value byte = b.extract_element(bytes, b.i32_val(i));
+-        Value elements = NVVM::ConvertF4x2ToF16x2Op::create(
+-            rewriter, loc, vec_ty(f16_ty, 2), byte, Float4E2M1FNType::get(ctx),
+-            false);
+-        results[2 * i] = b.extract_element(elements, b.i32_val(0));
+-        results[2 * i + 1] = b.extract_element(elements, b.i32_val(1));
+-      }
+-      return results;
+-    }
+     std::string ptx = toFp16 ? FP4ToFP16Ptx : FP4ToBP16Ptx;
+-    if (!toFp16 && targetInfo.getComputeCapability() >= 100 &&
+-        targetInfo.getPtxVersion() >= 92)
+-      ptx = FP4ToBf16NativePtx;
++    if (targetInfo.getComputeCapability() >= 100 &&
++        targetInfo.getPtxVersion() >= (toFp16 ? 86 : 92))
++      ptx = toFp16 ? FP4ToFp16NativePtx : FP4ToBf16NativePtx;
+     SmallVector<Type> rets(4, i32_ty);
+     Type retType = struct_ty(rets);
+     Value ret = createInlineAsmUpcast(loc, rewriter, ptx, retType, packedVec);
 
 --- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
 +++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/LoadStoreOpToLLVM.cpp
diff --git a/third_party/xla/third_party/triton/common/tmemfix.patch b/third_party/xla/third_party/triton/common/tmemfix.patch
index 75d0f89a0d4..f5d1787db5b 100644
--- a/third_party/xla/third_party/triton/common/tmemfix.patch
+++ b/third_party/xla/third_party/triton/common/tmemfix.patch
@@ -157,7 +157,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
 
 --- a/python/test/unit/language/test_compile_only.py
 +++ b/python/test/unit/language/test_compile_only.py
-@@ -413,7 +413,7 @@ def test_compile_only_dot_mxfp() -> None:
+@@ -442,7 +442,7 @@ def test_compile_only_dot_mxfp() -> None:
      assert re.search(pattern, str(ttgir)), "The TTGIR does not match the expected pattern."
  
      ptx = k.asm["ptx"]
@@ -169,7 +169,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
 
 --- a/python/test/unit/language/test_core.py
 +++ b/python/test/unit/language/test_core.py
-@@ -4388,10 +4388,7 @@ def get_test_small_dots_cases():
+@@ -4450,10 +4450,7 @@ def get_test_small_dots_cases():
      if not is_cuda():
          return []
      return [(2, 4, 32, 1, False, False, 'None', 'ieee', 'float16', 'float32', 1, None),
@@ -181,7 +181,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
  
  
  @pytest.mark.interpreter
-@@ -4630,21 +4627,19 @@ def test_dot(M, N, K, num_warps, col_a, col_b, epilogue, input_precision, in_dty
+@@ -4692,21 +4689,19 @@ def test_dot(M, N, K, num_warps, col_a, col_b, epilogue, input_precision, in_dty
      # make sure ld/st are vectorized
      ptx = pgm.asm['ptx']
  
@@ -252,7 +252,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
 
 --- a/test/Conversion/tritongpu_to_llvm_blackwell.mlir
 +++ b/test/Conversion/tritongpu_to_llvm_blackwell.mlir
-@@ -341,10 +341,10 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+@@ -340,10 +340,10 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
    // CHECK-LABEL: @tc_gen5_mma_scaled_a_tmem
    // CHECK: %[[TMEM_BASE:.+]] = llvm.ptrtoint %arg2 : !llvm.ptr<3> to i32
    // CHECK: %[[A_BASE:.+]] = llvm.ptrtoint %arg0 : !llvm.ptr<3> to i32
@@ -267,7 +267,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    tt.func @tc_gen5_mma_scaled_a_tmem(
        %a: !ttg.memdesc<128x256xf8E5M2, #tmem, #ttng.tensor_memory>,
        %b: !ttg.memdesc<256x128xf8E5M2, #shared, #ttg.shared_memory>,
-@@ -381,8 +381,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+@@ -380,8 +380,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
    // CHECK: %[[P0:.+]] = llvm.icmp "eq" %[[WID]], %[[C0]] : i32
    // CHECK: %[[P1:.+]] = llvm.and %{{.*}}, %[[P0]]  : i1
    // CHECK: llvm.cond_br %[[P1]]
@@ -278,7 +278,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    tt.func @tc_gen5_mma_block_scale(%a: !ttg.memdesc<128x64xf8E4M3FN, #shared, #ttg.shared_memory>,
                         %b: !ttg.memdesc<32x128xi8, #shared1, #ttg.shared_memory>,
                         %c: !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>,
-@@ -416,10 +416,10 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+@@ -415,10 +415,10 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
    // CHECK-DAG: %[[DESC1:.+]] = llvm.mlir.constant(681640592 : i32) : i32
    // CHECK-DAG: %[[DESC2:.+]] = llvm.mlir.constant(1218511520 : i32) : i32
    // CHECK-DAG: %[[DESC3:.+]] = llvm.mlir.constant(1755382448 : i32) : i32
@@ -293,7 +293,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    tt.func @tc_gen5_mma_block_scale_fp4_a(%a: !ttg.memdesc<128x64xi8, #shared1, #ttg.shared_memory>,
                         %b: !ttg.memdesc<128x128xi8, #shared, #ttg.shared_memory>,
                         %c: !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>,
-@@ -582,8 +582,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+@@ -581,8 +581,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
    // CHECK-LABEL: @tc_gen5_mma_block_scale_nvfp4
    // CHECK-DAG: %[[DESC0:.+]] = llvm.mlir.constant(138413184 : i32) : i32
    // CHECK: %[[TMEM_BASE:.+]] = llvm.ptrtoint %arg2 : !llvm.ptr<3> to i32
@@ -304,7 +304,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    tt.func @tc_gen5_mma_block_scale_nvfp4(%a: !ttg.memdesc<128x64xi8, #shared, #ttg.shared_memory>,
                         %b: !ttg.memdesc<64x256xi8, #shared1, #ttg.shared_memory>,
                         %c: !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable>,
-@@ -616,8 +616,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+@@ -615,8 +615,8 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
    // CHECK-DAG: %[[DESC0:.+]] = llvm.mlir.constant(146801792 : i32) : i32
    // CHECK-DAG: %[[DESC1:.+]] = llvm.mlir.constant(1220543648 : i32) : i32
    // CHECK-DAG: %[[TMEM_BASE:.+]] = llvm.ptrtoint %arg2 : !llvm.ptr<3> to i32
@@ -315,7 +315,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    tt.func @tc_gen5_mma_block_scale_mxfp4(%a: !ttg.memdesc<128x64xi8, #shared, #ttg.shared_memory>,
                         %b: !ttg.memdesc<64x256xi8, #shared1, #ttg.shared_memory>,
                         %c: !ttg.memdesc<128x256xf32, #tmem, #ttng.tensor_memory, mutable>,
-@@ -1237,7 +1237,7 @@ module attributes {"ttg.num-warps" = 8 : i32} {
+@@ -1246,7 +1246,7 @@ module attributes {"ttg.num-warps" = 8 : i32} {
  // CHECK-LABEL: @load_store_16x32bx1_broadcast
  tt.func private @load_store_16x32bx1_broadcast(%arg0: !ttg.memdesc<16x8xi8, #tmem_scales, #ttng.tensor_memory, mutable>, %arg1: tensor<16x8xi8, #linear>) {
    %true = arith.constant true
@@ -324,7 +324,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    ttng.tmem_store %arg1, %arg0, %true : tensor<16x8xi8, #linear> -> !ttg.memdesc<16x8xi8, #tmem_scales, #ttng.tensor_memory, mutable>
    tt.return
  }
-@@ -1252,7 +1252,7 @@ module attributes {"ttg.num-warps" = 4 : i32} {
+@@ -1261,7 +1261,7 @@ module attributes {"ttg.num-warps" = 4 : i32} {
  // CHECK-LABEL: @load_store_16x2_scales_uses_zero_second_half
  tt.func private @load_store_16x2_scales_uses_zero_second_half(%arg0: !ttg.memdesc<16x2xi8, #tmem_scales, #ttng.tensor_memory, mutable>, %arg1: tensor<16x2xi8, #linear>) {
    %true = arith.constant true
@@ -333,7 +333,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    // CHECK: llvm.return
    ttng.tmem_store %arg1, %arg0, %true : tensor<16x2xi8, #linear> -> !ttg.memdesc<16x2xi8, #tmem_scales, #ttng.tensor_memory, mutable>
    tt.return
-@@ -1316,22 +1316,22 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+@@ -1325,22 +1325,22 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
        %pred: i1) {
      // CHECK: %[[ACC_BASE:.+]] = llvm.ptrtoint %arg4 : !llvm.ptr<3> to i32
      // CHECK: %[[A_E4M3_BASE:.+]] = llvm.ptrtoint %arg0 : !llvm.ptr<3> to i32
@@ -364,7 +364,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
      // CHECK-SAME: %[[ACC_BASE]], %[[A_E5M2_BASE]]
      ttng.tc_gen5_mma_scaled %a_e4m3, %b_e5m2, %c, %scale_a, %scale_b, %useAcc, %pred lhs = e4m3 rhs = e5m2 :
         !ttg.memdesc<128x128xf8E4M3FN, #tmem_fp8_lhs, #ttng.tensor_memory>,
-@@ -1359,13 +1359,13 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+@@ -1368,13 +1368,13 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
        %pred: i1) {
      // CHECK-DAG: %[[ACC_BASE:.+]] = llvm.ptrtoint %arg2 : !llvm.ptr<3> to i32
      // CHECK-DAG: %[[A_BASE:.+]] = llvm.ptrtoint %arg0 : !llvm.ptr<3> to i32
@@ -382,7 +382,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
      // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]]
      ttng.tc_gen5_mma_scaled %a, %b, %c, %scale_a, %scale_b, %useAcc, %pred lhs = e2m1 rhs = e5m2 :
         !ttg.memdesc<128x64xi8, #tmem_fp4_padded_lhs, #ttng.tensor_memory>,
-@@ -1395,13 +1395,13 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
+@@ -1404,13 +1404,13 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
        %pred: i1) {
      // CHECK-DAG: %[[ACC_BASE:.+]] = llvm.ptrtoint %arg2 : !llvm.ptr<3> to i32
      // CHECK-DAG: %[[A_BASE:.+]] = llvm.ptrtoint %arg0 : !llvm.ptr<3> to i32
@@ -400,7 +400,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
      // CHECK-SAME: %[[ACC_BASE]], %[[A_BASE]]
      ttng.tc_gen5_mma_scaled %a, %b, %c, %scale_a, %scale_b, %useAcc, %pred lhs = e2m1 rhs = e2m1 :
         !ttg.memdesc<128x128xi8, #tmem_fp4_dense_lhs, #ttng.tensor_memory>,
-@@ -1651,13 +1651,13 @@ module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttng.tw
+@@ -1660,13 +1660,13 @@ module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttng.tw
    // SM103-DAG: %[[K96_0:.+]] = llvm.mlir.constant(-1868561280 : i32)
    // SM103-DAG: %[[K96_96:.+]] = llvm.mlir.constant(-257948496 : i32)
    // SM103-DAG: %[[K64_192:.+]] = llvm.mlir.constant(1352664224 : i32)
@@ -418,7 +418,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    // SM103-NOT: tcgen05.mma
    // CHECK: llvm.return
    tt.func @tc_gen5_mma_mxfp4_k256(
-@@ -1681,20 +1681,20 @@ module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttng.tw
+@@ -1690,20 +1690,20 @@ module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttng.tw
    // SM103-DAG: %[[K64_192:.+]] = llvm.mlir.constant(272630912 : i32)
    // SM103-DAG: %[[C16:.+]] = llvm.mlir.constant(16 : i32)
    // SM103-DAG: %[[C32:.+]] = llvm.mlir.constant(32 : i32)
@@ -446,7 +446,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    // SM103-NOT: tcgen05.mma
    // CHECK: llvm.return
    tt.func @tc_gen5_mma_nvfp4_k512(
-@@ -1716,11 +1716,11 @@ module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttng.tw
+@@ -1725,11 +1725,11 @@ module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttng.tw
    // SM103-DAG: %[[ABS_HIGH:.+]] = llvm.mlir.constant(1074806848 : i32)
    // SM103-DAG: %[[ADDR_MASK:.+]] = llvm.mlir.constant(16383 : i64)
    // SM103-DAG: %[[SHIFT:.+]] = llvm.mlir.constant(16 : i64)
@@ -460,7 +460,7 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
    // SM103: %[[ADESC:.+]] = llvm.bitcast
    // SM103: %[[ANEXT:.+]] = llvm.bitcast
    // SM103: %[[AADDR:.+]] = llvm.and %[[ANEXT]], %[[ADDR_MASK]]
-@@ -1732,8 +1732,8 @@ module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttng.tw
+@@ -1741,8 +1741,8 @@ module attributes {"ttg.num-ctas" = 2 : i32, "ttg.num-warps" = 4 : i32, "ttng.tw
    // SM103: %[[BADDR:.+]] = llvm.and %[[BNEXT]], %[[ADDR_MASK]]
    // SM103: %[[BLEAD:.+]] = llvm.shl %[[BADDR]], %[[SHIFT]]
    // SM103: %[[BABS:.+]] = llvm.or %[[BDESC]], %[[BLEAD]]
@@ -616,7 +616,4 @@ This patch allows to use CUDA 12.8 with PTXAS 8.7 instead of CUDA 13.2
      int baseK = dot.getKOffset(k) / dot.mmaSizeK;
      int subWordIdx = baseK % mmasPerScaleWord;
      int wordIdx = baseK / mmasPerScaleWord;
-     int scaleIdxA = linearizeScaleBlockIdx(op.getAScale(), m, wordIdx, numRepM,
-                                            numRepKWords);
-     int scaleIdxB = linearizeScaleBlockIdx(op.getBScale(), n, wordIdx, numRepN,
 
diff --git a/third_party/xla/third_party/triton/oss_only/build_files.patch b/third_party/xla/third_party/triton/oss_only/build_files.patch
index c0f09c75e0f..5519e5cace5 100644
--- a/third_party/xla/third_party/triton/oss_only/build_files.patch
+++ b/third_party/xla/third_party/triton/oss_only/build_files.patch
@@ -4049,7 +4049,7 @@ diff --git b/unittest/BUILD b/unittest/BUILD
 new file mode 100644
 --- /dev/null
 +++ b/unittest/BUILD
-@@ -0,0 +1,159 @@
+@@ -0,0 +1,162 @@
 +load("@rules_cc//cc:cc_test.bzl", "cc_test")
 +load("//tools/build_defs/build_test:build_test.bzl", "build_test")
 +
@@ -4065,10 +4065,13 @@ new file mode 100644
 +    deps = [
 +        "@com_google_googletest//:gtest_main",
 +        "@llvm-project//llvm:Support",
++        "@llvm-project//mlir:IR",
 +        "@llvm-project//mlir:Parser",
 +        "@triton//:TritonDialects",
++        "@triton//:TritonGPUTransforms",
 +        "@triton//:TritonNvidiaGPUTransforms",
 +        "@triton//:TritonTransforms",
++        "@triton//:WarpSpecialization",
 +    ],
 +)
 +
diff --git a/third_party/xla/third_party/triton/workspace.bzl b/third_party/xla/third_party/triton/workspace.bzl
index 4ba5608b5fb..3b843055089 100644
--- a/third_party/xla/third_party/triton/workspace.bzl
+++ b/third_party/xla/third_party/triton/workspace.bzl
@@ -71,8 +71,8 @@ triton_archive = repository_rule(
 def repo():
     """Imports Triton."""
 
-    TRITON_COMMIT = "0ccad974ee86942e704ac06dafdc0c3f214eb229"
-    TRITON_SHA256 = "05f3bf327718825595400fc0d797f4f78ec470cef731e41eef07b9bb3ef9b3d6"
+    TRITON_COMMIT = "b575c0c38d428a609e25487db64051a5fa7bfab8"
+    TRITON_SHA256 = "8af8fff829127873870396b6fc0a6e402b1fb5ca96692126c9138ba2710b0346"
     triton_archive(
         name = "triton",
         sha256 = TRITON_SHA256,
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/compilation_pipeline_rocm.cc b/third_party/xla/xla/backends/gpu/codegen/triton/compilation_pipeline_rocm.cc
index 2468daa31cb..39184beaba7 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/compilation_pipeline_rocm.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/compilation_pipeline_rocm.cc
@@ -85,8 +85,8 @@ static void MakeTTGIR(mlir::OpPassManager* pm,
   pm->addPass(mt::gpu::createTritonGPURemoveLayoutConversions());
   // TODO ROCm Check if we want to compare MI100 and greater
   pm->addPass(mlir::createTritonAMDGPUOptimizeEpilogue());
-  pm->addPass(mt::amdgpu::createTritonAMDGPUOptimizeDotOperands(
-      {rocm_cc.gfx_version()}));
+  pm->addPass(
+      mlir::createTritonAMDGPUOptimizeDotOperands({rocm_cc.gfx_version()}));
   pm->addNestedPass<mlir::triton::FuncOp>(
       mlir::createTritonAMDGPUHoistLayoutConversions());
   pm->addNestedPass<mlir::triton::FuncOp>(
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/fp8_to_llvm_hopper.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/fp8_to_llvm_hopper.mlir
index c2974b82e3b..23555478aa7 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/fp8_to_llvm_hopper.mlir
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/fp8_to_llvm_hopper.mlir
@@ -38,7 +38,7 @@ tt.func @e4m3_mapping(
     %arg1: tensor<256x16xf8E4M3FN, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 4}>>
   ) {
   %cst = arith.constant dense<0.000000e+00> : tensor<16x16xf32, #mma>
-  // CHECK: mma.{{.*}}.e4m3.e4m3.f32
+  // CHECK: nvvm.mma.sync {{.*}}shape = <m = 16, n = 8, k = 32>, layout_a = row, layout_b = col, multiplicand_a_ptx_type = e4m3, multiplicand_b_ptx_type = e4m3 : (i32, i32, f32)
   %res = tt.dot %arg0, %arg1, %cst {allowTF32 = true, maxNumImpreciseAcc = 0 : i32}
       : tensor<16x256xf8E4M3FN, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 4}>> *
         tensor<256x16xf8E4M3FN, #ttg.dot_op<{opIdx = 1, parent = #mma, kWidth = 4}>>
diff --git a/third_party/xla/xla/service/gpu/autotuning/autotune_cache_key.h b/third_party/xla/xla/service/gpu/autotuning/autotune_cache_key.h
index b9370cb6647..5e3153b598d 100644
--- a/third_party/xla/xla/service/gpu/autotuning/autotune_cache_key.h
+++ b/third_party/xla/xla/service/gpu/autotuning/autotune_cache_key.h
@@ -34,7 +34,7 @@ class AutotuneCacheKey {
   // changes that may affect the autotuning results.
   // To prevent accidental merges of concurrent increments, update the comment
   // to explain why the version is bumped.
-  static constexpr int kCurrentVersion = 55;  // Triton integration 1.40
+  static constexpr int kCurrentVersion = 56;  // Triton integration 1.41
 
   AutotuneCacheKey(const se::DeviceDescription& device_description,
                    const HloInstruction& instruction,

From 8928d3ca8aac20ef73fd504b21214f6384c3aa5b Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Thu, 15 Oct 2026 02:05:03 -0700
Subject: [PATCH 0857/1000] PR #50776: Bump actions/upload-artifact from 4.4.3
 to 7.0.2
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50776

Bumps [actions/upload-artifact](https://github.com/actions/upload-artifact) from 4.4.3 to 7.0.2.
<details>
<summary>Release notes</summary>
<p><em>Sourced from <a href="https://github.com/actions/upload-artifact/releases">actions/upload-artifact's releases</a>.</em></p>
<blockquote>
<h2>v7.0.2</h2>
<h2>What's Changed</h2>
<ul>
<li>Improves artifact download retries when the service returns HTTP 429 (rate limiting), including honoring valid Retry-After headers.</li>
<li>Updates <code>@​actions/artifact</code> to v6.3.1.</li>
</ul>
<h2>New Contributors</h2>
<ul>
<li><a href="https://github.com/tunc-d"><code>@​tunc-d</code></a> made their first contribution in <a href="https://redirect.github.com/actions/upload-artifact/pull/821">actions/upload-artifact#821</a></li>
</ul>
<p><strong>Full Changelog</strong>: <a href="https://github.com/actions/upload-artifact/compare/v7.0.1...v7.0.2">https://github.com/actions/upload-artifact/compare/v7.0.1...v7.0.2</a></p>
<h2>v7.0.1</h2>
<h2>What's Changed</h2>
<ul>
<li>Update the readme with direct upload details by <a href="https://github.com/danwkennedy"><code>@​danwkennedy</code></a> in <a href="https://redirect.github.com/actions/upload-artifact/pull/795">actions/upload-artifact#795</a></li>
<li>Readme: bump all the example versions to v7 by <a href="https://github.com/danwkennedy"><code>@​danwkennedy</code></a> in <a href="https://redirect.github.com/actions/upload-artifact/pull/796">actions/upload-artifact#796</a></li>
<li>Include changes in typespec/ts-http-runtime 0.3.5 by <a href="https://github.com/yacaovsnc"><code>@​yacaovsnc</code></a> in <a href="https://redirect.github.com/actions/upload-artifact/pull/797">actions/upload-artifact#797</a></li>
</ul>
<p><strong>Full Changelog</strong>: <a href="https://github.com/actions/upload-artifact/compare/v7...v7.0.1">https://github.com/actions/upload-artifact/compare/v7...v7.0.1</a></p>
<h2>v7.0.0</h2>
<h2>v7 What's new</h2>
<h3>Direct Uploads</h3>
<p>Adds support for uploading single files directly (unzipped). Callers can set the new <code>archive</code> parameter to <code>false</code> to skip zipping the file during upload. Right now, we only support single files. The action will fail if the glob passed resolves to multiple files. The <code>name</code> parameter is also ignored with this setting. Instead, the name of the artifact will be the name of the uploaded file.</p>
<h3>ESM</h3>
<p>To support new versions of the <code>@actions/*</code> packages, we've upgraded the package to ESM.</p>
<h2>What's Changed</h2>
<ul>
<li>Add proxy integration test by <a href="https://github.com/Link"><code>@​Link</code></a>- in <a href="https://redirect.github.com/actions/upload-artifact/pull/754">actions/upload-artifact#754</a></li>
<li>Upgrade the module to ESM and bump dependencies by <a href="https://github.com/danwkennedy"><code>@​danwkennedy</code></a> in <a href="https://redirect.github.com/actions/upload-artifact/pull/762">actions/upload-artifact#762</a></li>
<li>Support direct file uploads by <a href="https://github.com/danwkennedy"><code>@​danwkennedy</code></a> in <a href="https://redirect.github.com/actions/upload-artifact/pull/764">actions/upload-artifact#764</a></li>
</ul>
<h2>New Contributors</h2>
<ul>
<li><a href="https://github.com/Link"><code>@​Link</code></a>- made their first contribution in <a href="https://redirect.github.com/actions/upload-artifact/pull/754">actions/upload-artifact#754</a></li>
</ul>
<p><strong>Full Changelog</strong>: <a href="https://github.com/actions/upload-artifact/compare/v6...v7.0.0">https://github.com/actions/upload-artifact/compare/v6...v7.0.0</a></p>
<h2>v6.0.0</h2>
<h2>v6 - What's new</h2>
<blockquote>
<p>[!IMPORTANT]
actions/upload-artifact@v6 now runs on Node.js 24 (<code>runs.using: node24</code>) and requires a minimum Actions Runner version of 2.327.1. If you are using self-hosted runners, ensure they are updated before upgrading.</p>
</blockquote>
<h3>Node.js 24</h3>
<p>This release updates the runtime to Node.js 24. v5 had preliminary support for Node.js 24, however this action was by default still running on Node.js 20. Now this action by default will run on Node.js 24.</p>
<!-- raw HTML omitted -->
</blockquote>
<p>... (truncated)</p>
</details>
<details>
<summary>Commits</summary>
<ul>
<li><a href="https://github.com/actions/upload-artifact/commit/cf430e030ddbb5b0abf93d22962f4752f3646cd9"><code>cf430e0</code></a> Merge pull request <a href="https://redirect.github.com/actions/upload-artifact/issues/822">#822</a> from actions/tunc-d-prepare-upload-v7-0-2</li>
<li><a href="https://github.com/actions/upload-artifact/commit/c2acdedd35a8824a1552c7137f8c18288c3ea469"><code>c2acded</code></a> chore: prepare v7.0.2 release</li>
<li><a href="https://github.com/actions/upload-artifact/commit/4687c1bbaa0a58abb517d90b4bbade6754c80c1a"><code>4687c1b</code></a> Merge pull request <a href="https://redirect.github.com/actions/upload-artifact/issues/821">#821</a> from actions/tunc-d-upload-artifact-package-bump</li>
<li><a href="https://github.com/actions/upload-artifact/commit/8cbe1bebdb858105d8fdd88f81404d4de4595719"><code>8cbe1be</code></a> Bump <code>@​actions/artifact</code> to 6.3.1</li>
<li><a href="https://github.com/actions/upload-artifact/commit/043fb46d1a93c77aae656e7c1c64a875d1fc6a0a"><code>043fb46</code></a> Merge pull request <a href="https://redirect.github.com/actions/upload-artifact/issues/797">#797</a> from actions/yacaovsnc/update-dependency</li>
<li><a href="https://github.com/actions/upload-artifact/commit/634250c1388765ea7ed0f053e636f1f399000b94"><code>634250c</code></a> Include changes in typespec/ts-http-runtime 0.3.5</li>
<li><a href="https://github.com/actions/upload-artifact/commit/e454baaac2be505c9450e11b8f3215c6fc023ce8"><code>e454baa</code></a> Readme: bump all the example versions to v7 (<a href="https://redirect.github.com/actions/upload-artifact/issues/796">#796</a>)</li>
<li><a href="https://github.com/actions/upload-artifact/commit/74fad66b98a6d799dc004d3353ccd0e6f6b2530e"><code>74fad66</code></a> Update the readme with direct upload details (<a href="https://redirect.github.com/actions/upload-artifact/issues/795">#795</a>)</li>
<li><a href="https://github.com/actions/upload-artifact/commit/bbbca2ddaa5d8feaa63e36b76fdaad77386f024f"><code>bbbca2d</code></a> Support direct file uploads (<a href="https://redirect.github.com/actions/upload-artifact/issues/764">#764</a>)</li>
<li><a href="https://github.com/actions/upload-artifact/commit/589182c5a4cec8920b8c1bce3e2fab1c97a02296"><code>589182c</code></a> Upgrade the module to ESM and bump dependencies (<a href="https://redirect.github.com/actions/upload-artifact/issues/762">#762</a>)</li>
<li>Additional commits viewable in <a href="https://github.com/actions/upload-artifact/compare/v4.4.3...cf430e030ddbb5b0abf93d22962f4752f3646cd9">compare view</a></li>
</ul>
</details>
<br />

[![Dependabot compatibility score](https://dependabot-badges.githubapp.com/badges/compatibility_score?dependency-name=actions/upload-artifact&package-manager=github_actions&previous-version=4.4.3&new-version=7.0.2)](https://docs.github.com/en/github/managing-security-vulnerabilities/about-dependabot-security-updates#about-compatibility-scores)

Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting `@dependabot rebase`.

[//]: # (dependabot-automerge-start)
[//]: # (dependabot-automerge-end)

---

<details>
<summary>Dependabot commands and options</summary>
<br />

You can trigger Dependabot actions by commenting on this PR:
- `@dependabot rebase` will rebase this PR
- `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it
- `@dependabot show <dependency name> ignore conditions` will show all of the ignore conditions of the specified dependency
- `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
- `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
- `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)

</details>
Copybara import of the project:

--
46942a38cfc17c9cf8373605ce44eda8b810ef23 by dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>:

Bump actions/upload-artifact from 4.4.3 to 7.0.2

Bumps [actions/upload-artifact](https://github.com/actions/upload-artifact) from 4.4.3 to 7.0.2.
- [Release notes](https://github.com/actions/upload-artifact/releases)
- [Commits](https://github.com/actions/upload-artifact/compare/v4.4.3...cf430e030ddbb5b0abf93d22962f4752f3646cd9)

---
updated-dependencies:
- dependency-name: actions/upload-artifact
  dependency-version: 7.0.2
  dependency-type: direct:production
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>

Merging this change closes #50776

PiperOrigin-RevId: 999971744
---
 third_party/xla/.github/workflows/rocm_ci.yml             | 4 ++--
 third_party/xla/.github/workflows/scorecards-analysis.yml | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/.github/workflows/rocm_ci.yml b/third_party/xla/.github/workflows/rocm_ci.yml
index 43aa1903d48..8afc130fc60 100644
--- a/third_party/xla/.github/workflows/rocm_ci.yml
+++ b/third_party/xla/.github/workflows/rocm_ci.yml
@@ -117,7 +117,7 @@ jobs:
 
       - name: Upload JAX BEP File
         if: always()
-        uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882 # v4.4.3
+        uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2
         with:
           name: jax-bep
           path: jax/jax_build_events.json
@@ -190,7 +190,7 @@ jobs:
 
       - name: Upload XLA BEP Files
         if: always()
-        uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882 # v4.4.3
+        uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2
         with:
           name: xla-bep
           path: |
diff --git a/third_party/xla/.github/workflows/scorecards-analysis.yml b/third_party/xla/.github/workflows/scorecards-analysis.yml
index dac2a4a620f..65af0d95b3d 100644
--- a/third_party/xla/.github/workflows/scorecards-analysis.yml
+++ b/third_party/xla/.github/workflows/scorecards-analysis.yml
@@ -58,7 +58,7 @@ jobs:
       # Upload the results as artifacts (optional). Commenting out will disable uploads of run results in SARIF
       # format to the repository Actions tab.
       - name: "Upload artifact"
-        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
+        uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2
         with:
           name: SARIF file
           path: results.sarif

From 21df4b83e49beacd465bcaa1a426546277706a5a Mon Sep 17 00:00:00 2001
From: Manjunath Gaonkar <magaonka@amd.com>
Date: Thu, 15 Oct 2026 02:07:00 -0700
Subject: [PATCH 0858/1000] PR #50677: [ROCm] Drop per-node hipFuncSetAttribute
 from command buffer kernel nodes
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50677

## 📝 Summary of Changes

`RocmCommandBuffer` no longer calls `hipFuncSetAttribute(hipFuncAttributeMaxDynamicSharedMemorySize, ...)` on every kernel node it creates or updates with dynamic shared memory.

## 🎯 Justification

HIP has no CUDA-style opt-in for dynamic shared memory: a kernel's limit already defaults to the maximum, and HIP never reads this attribute when it builds a dispatch. The call also fails with `hipErrorInvalidDeviceFunction` for in-process kernels loaded with `hipGetFuncBySymbol` (e.g. the TopK custom kernel), so they could not be recorded into a ROCm command buffer, and it takes a process-wide lock in the HIP runtime on every call.

## 🚀 Kind of Contribution

🐛 Bug Fix

## 🧪 Unit Tests

`GpuCommandBufferTest.LaunchKernelWithDynamicSharedMemory` records an in-process kernel with dynamic shared memory, updates it to a larger request, and checks the result. It fails on ROCm without this change.

Copybara import of the project:

--
2df6b8d9b0ef0a6ea228d5e5817579b8d9e9fc4f by magaonka <magaonka@amd.com>:

[ROCm] Drop per-node hipFuncSetAttribute from command buffer kernel nodes

HIP never reads this attribute when it builds a dispatch, and the call
fails for in-process kernels from hipGetFuncBySymbol (e.g. TopK).

Merging this change closes #50677

PiperOrigin-RevId: 999972921
---
 third_party/xla/xla/stream_executor/gpu/BUILD |  2 +
 .../gpu/gpu_command_buffer_test.cc            | 61 +++++++++++++++++++
 .../rocm/rocm_command_buffer.cc               | 16 -----
 3 files changed, 63 insertions(+), 16 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/gpu/BUILD b/third_party/xla/xla/stream_executor/gpu/BUILD
index 3f051c01c5b..4c759fc0f0d 100644
--- a/third_party/xla/xla/stream_executor/gpu/BUILD
+++ b/third_party/xla/xla/stream_executor/gpu/BUILD
@@ -664,6 +664,7 @@ xla_test(
         "//xla/stream_executor:command_buffer",
         "//xla/stream_executor:device_address",
         "//xla/stream_executor:kernel",
+        "//xla/stream_executor:kernel_args",
         "//xla/stream_executor:kernel_spec",
         "//xla/stream_executor:launch_dim",
         "//xla/stream_executor:platform",
@@ -682,6 +683,7 @@ xla_test(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest",
diff --git a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc
index dbd8a113e95..217290dfae0 100644
--- a/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc
+++ b/third_party/xla/xla/stream_executor/gpu/gpu_command_buffer_test.cc
@@ -24,6 +24,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/status_matchers.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/ascii.h"
 #include "absl/types/span.h"
 #include "xla/service/platform_util.h"
@@ -32,6 +33,7 @@ limitations under the License.
 #include "xla/stream_executor/device_address.h"
 #include "xla/stream_executor/gpu/gpu_test_kernels.h"
 #include "xla/stream_executor/kernel.h"
+#include "xla/stream_executor/kernel_args.h"
 #include "xla/stream_executor/kernel_spec.h"
 #include "xla/stream_executor/launch_dim.h"
 #include "xla/stream_executor/platform.h"
@@ -142,6 +144,65 @@ TEST_F(GpuCommandBufferTest, LaunchSingleKernel) {
   ASSERT_EQ(dst, expected);
 }
 
+// DynShmemKernel is an in-process kernel (loaded from its host symbol), so this
+// also covers recording such kernels with dynamic shared memory.
+TEST_F(GpuCommandBufferTest, LaunchKernelWithDynamicSharedMemory) {
+  Platform* platform = GpuPlatform();
+  StreamExecutor* executor = platform->ExecutorForDevice(0).value();
+
+  ASSERT_OK_AND_ASSIGN(auto stream, executor->CreateStream());
+  ASSERT_OK_AND_ASSIGN(auto kernel, LoadDynShmemTestKernel(executor));
+
+  // DynShmemKernel reverses an n_cols x n_rows byte buffer through dynamic
+  // shared memory of n_cols * n_rows bytes, one thread per column.
+  constexpr uint32_t kCols = 64;
+  constexpr uint32_t kMaxRows = 512;
+  DeviceAddress<uint8_t> buf =
+      executor->AllocateArray<uint8_t>(kCols * kMaxRows);
+  DeviceAddress<uint32_t> cols = executor->AllocateScalar<uint32_t>();
+  DeviceAddress<uint32_t> rows = executor->AllocateScalar<uint32_t>();
+  ASSERT_OK(stream->Memcpy(&cols, &kCols, sizeof(kCols)));
+
+  auto upload = [&](uint32_t n_rows) -> absl::StatusOr<std::vector<uint8_t>> {
+    uint32_t n = kCols * n_rows;
+    std::vector<uint8_t> host(n);
+    for (uint32_t i = 0; i < n; ++i) {
+      host[i] = static_cast<uint8_t>(i * 7 + 3);
+    }
+    ABSL_RETURN_IF_ERROR(stream->Memcpy(&buf, host.data(), n));
+    ABSL_RETURN_IF_ERROR(stream->Memcpy(&rows, &n_rows, sizeof(n_rows)));
+    return host;
+  };
+  auto expect_reversed = [&](const std::vector<uint8_t>& before) {
+    std::vector<uint8_t> after(before.size());
+    ASSERT_OK(stream->Memcpy(after.data(), buf, after.size()));
+    ASSERT_OK(stream->BlockHostUntilDone());
+    std::vector<uint8_t> expected(before.rbegin(), before.rend());
+    ASSERT_EQ(after, expected);
+  };
+
+  // Create with a small dynamic shared memory request.
+  ASSERT_OK_AND_ASSIGN(std::vector<uint8_t> before, upload(16));
+  ASSERT_OK_AND_ASSIGN(auto cmd_buffer, executor->CreateCommandBuffer(primary));
+  ASSERT_OK_AND_ASSIGN(auto* launch,
+                       cmd_buffer->CreateLaunch(
+                           ThreadDim(kCols), BlockDim(1), {}, *kernel,
+                           *PackKernelArgs(kCols * 16, buf, cols, rows), {}));
+  ASSERT_OK(cmd_buffer->Finalize());
+  ASSERT_OK(cmd_buffer->Submit(stream.get()));
+  ASSERT_NO_FATAL_FAILURE(expect_reversed(before));
+
+  // Update to a larger request.
+  ASSERT_OK_AND_ASSIGN(before, upload(kMaxRows));
+  ASSERT_OK(cmd_buffer->Update());
+  ASSERT_OK(cmd_buffer->UpdateLaunch(
+      launch, ThreadDim(kCols), BlockDim(1), {}, *kernel,
+      *PackKernelArgs(kCols * kMaxRows, buf, cols, rows)));
+  ASSERT_OK(cmd_buffer->Finalize());
+  ASSERT_OK(cmd_buffer->Submit(stream.get()));
+  ASSERT_NO_FATAL_FAILURE(expect_reversed(before));
+}
+
 TEST_F(GpuCommandBufferTest, TraceSingleKernel) {
   Platform* platform = GpuPlatform();
   StreamExecutor* executor = platform->ExecutorForDevice(0).value();
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.cc b/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.cc
index 3e8d3aef633..f85e6098520 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.cc
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_command_buffer.cc
@@ -372,14 +372,6 @@ absl::StatusOr<GraphNodeHandle> RocmCommandBuffer::CreateKernelNode(
   params.kernelParams = const_cast<void**>(args.argument_addresses().data());
   params.extra = nullptr;
 
-  if (shared_mem_bytes != 0) {
-    ABSL_RETURN_IF_ERROR(
-        ToStatus(hipFuncSetAttribute(function,
-                                     hipFuncAttributeMaxDynamicSharedMemorySize,
-                                     shared_mem_bytes),
-                 "Failed to set shared memory size"));
-  }
-
   std::vector<hipGraphNode_t> deps = ToHipGraphHandles(dependencies);
 
   hipGraphNode_t node_handle = nullptr;
@@ -437,14 +429,6 @@ absl::Status RocmCommandBuffer::UpdateKernelNode(
   params.kernelParams = const_cast<void**>(args.argument_addresses().data());
   params.extra = nullptr;
 
-  if (shared_mem_bytes != 0) {
-    ABSL_RETURN_IF_ERROR(
-        ToStatus(hipFuncSetAttribute(function,
-                                     hipFuncAttributeMaxDynamicSharedMemorySize,
-                                     shared_mem_bytes),
-                 "Failed to set shared memory size"));
-  }
-
   return ToStatus(hipGraphExecKernelNodeSetParams(
                       exec_, ToHipGraphHandle(node_handle), &params),
                   "Failed to set HIP graph kernel node params");

From 84a72f22cebc207a839230427d09039575fdfaca Mon Sep 17 00:00:00 2001
From: Anutosh Bhat <andersonbhat491@gmail.com>
Date: Thu, 15 Oct 2026 02:07:32 -0700
Subject: [PATCH 0859/1000] PR #50772: Fix dynamic kernel counts on 32-bit
 targets

Imported from GitHub PR https://github.com/openxla/xla/pull/50772

Found this while building JAX for WebAssembly. On wasm32, converting the `-1` dynamic-count sentinel to `size_t` makes `KernelThunk` select enormous fixed arrays instead of dynamic storage.

Keep the comparison signed and clamp sizes before converting to `size_t`. Checked header instantiation before/after on wasm32 and wasm64.
Copybara import of the project:

--
ce41ec4348f20afd4fe5daf9a4c83b3ff202ac94 by anutosh491 <andersonbhat491@gmail.com>:

Fix dynamic kernel counts on 32-bit targets

Merging this change closes #50772

PiperOrigin-RevId: 999973260
---
 third_party/xla/xla/backends/cpu/runtime/kernel_thunk.h | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/backends/cpu/runtime/kernel_thunk.h b/third_party/xla/xla/backends/cpu/runtime/kernel_thunk.h
index e8949a8f028..f36bc96594d 100644
--- a/third_party/xla/xla/backends/cpu/runtime/kernel_thunk.h
+++ b/third_party/xla/xla/backends/cpu/runtime/kernel_thunk.h
@@ -107,12 +107,12 @@ class KernelThunk : public KernelThunkBase {
  private:
   friend class ::xla::cpu::KernelThunk;
 
-  static constexpr bool IsDynamic(size_t n) {
+  static constexpr bool IsDynamic(int64_t n) {
     return n == kDynamicKernelParameter;
   }
 
   static constexpr size_t Size(int64_t size) {
-    return std::max<size_t>(size, 0);
+    return static_cast<size_t>(std::max<int64_t>(size, 0));
   }
 
   // If we know the number of arguments and results at compile time, we use

From fb4a2e96b15eec2e1bc92903572aed8665f0774a Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 15 Oct 2026 02:12:02 -0700
Subject: [PATCH 0860/1000] - Replace the previous primary/fallback approach by
 testing a list of algorithm strategies and choosing the first one that fits
 the estimated memory limits (or the most memory-efficient one if all exceed
 limits).

The list of the strategies to try:

1. FAST_MERGE for AssignBuffersForComputations, FAST_MERGE for AssignBuffersWithSequentialOrdering (was used before this change).

2. FAST_MERGE for AssignBuffersForComputations, FAST_SPLIT for AssignBuffersWithSequentialOrdering (this one is new).

3. DEFAULT for AssignBuffersForComputations, DEFAULT for AssignBuffersWithSequentialOrdering (was used before this change).

- Cache and reuse the non-sequential buffer assignment results across different strategies.
- Introduce `BufferAssignment::State` to snapshot and restore allocations while evaluating different strategies.

PiperOrigin-RevId: 999975757
---
 .../xla/xla/service/buffer_assignment.cc      | 240 ++++++++++++++----
 .../xla/xla/service/buffer_assignment.h       |  62 ++++-
 .../xla/xla/service/buffer_assignment_test.cc | 216 ++++++++++++++++
 3 files changed, 463 insertions(+), 55 deletions(-)

diff --git a/third_party/xla/xla/service/buffer_assignment.cc b/third_party/xla/xla/service/buffer_assignment.cc
index ede25e7c1d5..4b2cf591855 100644
--- a/third_party/xla/xla/service/buffer_assignment.cc
+++ b/third_party/xla/xla/service/buffer_assignment.cc
@@ -1213,6 +1213,9 @@ absl::Status BufferAssignment::CombineTempAllocations(
         const HloAliasAnalysis&, const HloInstruction*, const ShapeIndex&)>>
         fixed_offset) {
   VLOG(1) << "CombineTempAllocations()";
+  // All temp allocations, including the ones combined by an earlier call, are
+  // processed below, so the total is recomputed from scratch.
+  temp_allocation_total_size_ = 0;
 
   // Move all temp allocations into a single run at the end of the allocations
   // vector.
@@ -1382,26 +1385,41 @@ absl::Status BufferAssignment::CombineTempAllocations(
   return absl::OkStatus();
 }
 
-void BufferAssignment::ClearAllocationsForColors(
+void BufferAssignment::RestoreStateForColors(
+    const State& state,
     const absl::flat_hash_set<LogicalBuffer::Color>& colors) {
   allocations_.erase(std::remove_if(allocations_.begin(), allocations_.end(),
                                     [&](const BufferAllocation& alloc) {
                                       return colors.contains(alloc.color());
                                     }),
                      allocations_.end());
+  for (const BufferAllocation& alloc : state.allocations) {
+    if (colors.contains(alloc.color())) {
+      allocations_.push_back(alloc);
+    }
+  }
 
   temp_allocation_total_size_ = 0;
   std::fill(allocation_index_for_value_id_.begin(),
             allocation_index_for_value_id_.end(), -1);
   num_assigned_values_ = 0;
-  absl::erase_if(fragmentation_bytes_by_color_, [&](const auto& entry) {
-    return colors.contains(entry.first);
-  });
+  // NOLINTNEXTLINE : the order of the loop is not important.
+  for (LogicalBuffer::Color color : colors) {
+    auto it = state.fragmentation_bytes_by_color.find(color);
+    if (it != state.fragmentation_bytes_by_color.end()) {
+      fragmentation_bytes_by_color_[color] = it->second;
+    } else {
+      fragmentation_bytes_by_color_.erase(color);
+    }
+  }
   stats_ = Stats();
   for (BufferAllocation::Index index = 0; index < allocations_.size();
        ++index) {
     BufferAllocation& allocation = allocations_[index];
     allocation.set_index(index);
+    if (allocation.IsPreallocatedTempBuffer()) {
+      temp_allocation_total_size_ += allocation.size();
+    }
     for (const HloValue* value : allocation.assigned_values_) {
       allocation_index_for_value_id_[value->id()] = index;
       ++num_assigned_values_;
@@ -1409,6 +1427,11 @@ void BufferAssignment::ClearAllocationsForColors(
   }
 }
 
+void BufferAssignment::ClearAllocationsForColors(
+    const absl::flat_hash_set<LogicalBuffer::Color>& colors) {
+  RestoreStateForColors({}, colors);
+}
+
 void BufferAssignment::ComputeSummaryStats() {
   for (auto& allocation : Allocations()) {
     if (allocation.is_entry_computation_parameter()) {
@@ -3666,30 +3689,93 @@ absl::Status BufferAssigner::RunAssignBuffersWithFallback(
     backup_presets = opts_.preset_assignments->ClonePresetAssignments();
   }
 
-  // 2. Primary Run: Execute buffer assignment using the primary configured
-  // algorithms across all colors.
-  absl::Status primary_status = RunAssignBuffers(
-      module, global_computations, thread_local_computations, assignment,
-      opts_.buffer_assignment_algorithm,
-      opts_.assignment_algorithm_for_computations_without_ordering);
+  // 2. Define the assignment strategies to try.
+  const std::vector<BufferAssigner::AssignmentStrategy> assignment_strategies =
+      {
+          {opts_.assignment_algorithm_for_computations_without_ordering,
+           opts_.buffer_assignment_algorithm},
+          {opts_.assignment_algorithm_for_computations_without_ordering,
+           buffer_assignment::BufferAssignmentAlgorithmProto::FAST_SPLIT},
+          {opts_.fallback_algorithm_for_computations_without_ordering,
+           opts_.fallback_algorithm},
+      };
 
-  // 3. Check if the primary run succeeded and met memory limits per color.
-  // Collect only the colors that exceeded their adjusted memory limit
-  // (subtracting a safety margin of 2.5 GiB).
-  absl::flat_hash_set<LogicalBuffer::Color> fallback_colors;
-  if (absl::IsResourceExhausted(primary_status)) {
-    VLOG(1) << "Primary BufferAssignment failed with ResourceExhausted. "
-            << "Triggering in-place fallback for all colors.";
-    for (const HloBuffer& buffer : assignment->alias_analysis().buffers()) {
-      ABSL_ASSIGN_OR_RETURN(LogicalBuffer::Color color, buffer.color());
-      fallback_colors.insert(color);
+  // The first step of RunAssignBuffers (AssignBuffersForGlobalComputations)
+  // only depends on the non-sequential algorithm. Its result, i.e. the
+  // allocations it creates and the temp buffers it leaves for heap simulation,
+  // is cached and reused by the following strategies with the same
+  // non-sequential algorithm.
+  std::optional<
+      buffer_assignment::
+          AssignmentAlgorithmForComputationsWithoutOrderingProto::Value>
+      cached_non_sequential_algorithm;
+  BufferAssignment::State cached_non_sequential_assignment_state;
+  absl::flat_hash_map<const HloComputation*,
+                      absl::flat_hash_set<const HloValue*>>
+      cached_buffers_to_assign_sequentially;
+
+  // The lowest memory usage and the assignment with it for each color that
+  // exceeded its memory limit with every strategy so far. Colors are assigned
+  // independently, so each one keeps its own best assignment.
+  struct BestAssignment {
+    int64_t memory_usage;
+    std::shared_ptr<const BufferAssignment::State> snapshot;
+  };
+  absl::flat_hash_map<LogicalBuffer::Color, BestAssignment> best_assignments;
+
+  std::optional<absl::flat_hash_set<LogicalBuffer::Color>> active_colors;
+
+  // 3. Iterate through the assignment strategies. Return the first result that
+  // fits the limit for all colors, or the best per color if none fits.
+  for (const BufferAssigner::AssignmentStrategy& strategy :
+       assignment_strategies) {
+    LOG(INFO) << "Trying strategy: " << strategy.non_sequential_algorithm
+              << ", " << strategy.sequential_algorithm;
+    const absl::flat_hash_set<LogicalBuffer::Color>* colors_to_assign =
+        active_colors.has_value() ? &*active_colors : nullptr;
+
+    if (cached_non_sequential_algorithm == strategy.non_sequential_algorithm) {
+      // Reuse the cached AssignBuffersForGlobalComputations result for the
+      // colors that are reassigned.
+      if (colors_to_assign != nullptr) {
+        assignment->RestoreStateForColors(
+            cached_non_sequential_assignment_state, *colors_to_assign);
+        // NOLINTNEXTLINE : the order of the loop is not important.
+        for (auto& [comp, buffers] : cached_buffers_to_assign_sequentially) {
+          absl::erase_if(buffers, [&](const HloValue* value) {
+            return !colors_to_assign->contains(value->color());
+          });
+        }
+      } else {
+        assignment->RestoreState(cached_non_sequential_assignment_state);
+      }
+    } else {
+      if (backup_presets) {
+        opts_.preset_assignments = backup_presets->ClonePresetAssignments();
+      }
+      if (colors_to_assign != nullptr) {
+        assignment->ClearAllocationsForColors(*colors_to_assign);
+      } else {
+        assignment->ClearAllocations();
+      }
+      // Clear the buffers collected for the previous non-sequential algorithm.
+      cached_buffers_to_assign_sequentially.clear();
+      ABSL_RETURN_IF_ERROR(AssignBuffersForGlobalComputations(
+          global_computations, &cached_buffers_to_assign_sequentially,
+          assignment, strategy.non_sequential_algorithm, colors_to_assign));
+
+      cached_non_sequential_algorithm = strategy.non_sequential_algorithm;
+      cached_non_sequential_assignment_state = assignment->GetState();
     }
-  } else if (primary_status.ok()) {
-    // Aggregate total allocated bytes across all allocations by color.
-    // Ensure we account for alignment fragmentation exactly the way
-    // CombineTempAllocations will.
-    absl::btree_map<LogicalBuffer::Color, int64_t> allocated_bytes_by_color;
 
+    // Proceed identically to what RunAssignBuffers does, but skip step 1
+    // (AssignBuffersForGlobalComputations) since we just did it.
+    ABSL_RETURN_IF_ERROR(RunAssignBuffers(
+        module, global_computations, thread_local_computations, assignment,
+        strategy.sequential_algorithm, strategy.non_sequential_algorithm,
+        colors_to_assign, &cached_buffers_to_assign_sequentially));
+
+    absl::btree_map<LogicalBuffer::Color, int64_t> allocated_bytes_by_color;
     for (const BufferAllocation& alloc : assignment->Allocations()) {
       LogicalBuffer::Color color = alloc.color();
       if (alloc.IsPreallocatedTempBuffer()) {
@@ -3702,6 +3788,7 @@ absl::Status BufferAssigner::RunAssignBuffersWithFallback(
       }
     }
 
+    absl::flat_hash_set<LogicalBuffer::Color> fallback_colors;
     for (const auto& [color, total_allocated_bytes] :
          allocated_bytes_by_color) {
       int64_t memory_limit = GetMemoryLimit(*assignment, color);
@@ -3713,37 +3800,72 @@ absl::Status BufferAssigner::RunAssignBuffersWithFallback(
 
         if (total_allocated_bytes > adjusted_limit) {
           fallback_colors.insert(color);
-          VLOG(1) << "Primary BufferAssignment exceeded memory limit for color "
-                  << color << " (" << total_allocated_bytes
-                  << " bytes vs limit of " << adjusted_limit
-                  << " bytes). Triggering in-place fallback to DEFAULT.";
+          VLOG(1) << "BufferAssignment strategy (non_sequential: "
+                  << strategy.non_sequential_algorithm
+                  << ", sequential: " << strategy.sequential_algorithm
+                  << ") exceeded memory limit for color " << color << " ("
+                  << total_allocated_bytes << " bytes vs limit of "
+                  << adjusted_limit << " bytes).";
         } else {
-          VLOG(1) << "Primary BufferAssignment did not exceed memory limit for "
-                     "color "
-                  << color << " (" << total_allocated_bytes
-                  << " bytes vs limit of " << adjusted_limit << " bytes).";
+          VLOG(1) << "BufferAssignment strategy (non_sequential: "
+                  << strategy.non_sequential_algorithm
+                  << ", sequential: " << strategy.sequential_algorithm
+                  << ") did not exceed memory limit for color " << color << " ("
+                  << total_allocated_bytes << " bytes vs limit of "
+                  << adjusted_limit << " bytes).";
         }
       } else {
-        VLOG(1)
-            << "Primary BufferAssignment did not set memory limit for color "
-            << color << ". Fallback not triggered.";
+        VLOG(1) << "BufferAssignment strategy (non_sequential: "
+                << strategy.non_sequential_algorithm
+                << ", sequential: " << strategy.sequential_algorithm
+                << ") did not set memory limit for color " << color
+                << ". Fallback not triggered.";
       }
     }
-  } else {
-    return primary_status;
+
+    if (fallback_colors.empty()) {
+      return absl::OkStatus();
+    }
+    // The colors that fit now keep this assignment.
+    absl::erase_if(best_assignments, [&](const auto& entry) {
+      return !fallback_colors.contains(entry.first);
+    });
+    // Record the colors whose memory usage improved, with a snapshot since the
+    // next strategy reassigns them. After the last strategy, the current
+    // assignment is kept, so no snapshot is taken.
+    const bool is_last_strategy = &strategy == &assignment_strategies.back();
+    std::shared_ptr<const BufferAssignment::State> snapshot;
+    for (const auto& [color, memory_usage] : allocated_bytes_by_color) {
+      if (!fallback_colors.contains(color)) {
+        continue;
+      }
+      auto it = best_assignments.find(color);
+      if (it != best_assignments.end() &&
+          it->second.memory_usage <= memory_usage) {
+        continue;
+      }
+      if (snapshot == nullptr && !is_last_strategy) {
+        snapshot = std::make_shared<const BufferAssignment::State>(
+            assignment->GetState());
+      }
+      best_assignments[color] = {memory_usage, snapshot};
+    }
+    active_colors = std::move(fallback_colors);
   }
 
-  // 4. Execute In-Place Fallback only for the colors that exceeded their limit.
-  if (!fallback_colors.empty()) {
-    assignment->ClearAllocationsForColors(fallback_colors);
-    if (backup_presets) {
-      opts_.preset_assignments = backup_presets->ClonePresetAssignments();
+  // No strategy fits all colors. Each color that still exceeds its limit gets
+  // its best assignment back, unless it is the current one.
+  std::vector<LogicalBuffer::Color> colors_to_restore;
+  // NOLINTNEXTLINE(*-custom-deterministic-iteration-order)
+  for (const auto& [color, best_assignment] : best_assignments) {
+    if (best_assignment.snapshot != nullptr) {
+      colors_to_restore.push_back(color);
     }
-    ABSL_RETURN_IF_ERROR(RunAssignBuffers(
-        module, global_computations, thread_local_computations, assignment,
-        opts_.fallback_algorithm,
-        opts_.fallback_algorithm_for_computations_without_ordering,
-        &fallback_colors));
+  }
+  absl::c_sort(colors_to_restore);
+  for (LogicalBuffer::Color color : colors_to_restore) {
+    assignment->RestoreStateForColors(*best_assignments.at(color).snapshot,
+                                      {color});
   }
   return absl::OkStatus();
 }
@@ -3757,15 +3879,25 @@ absl::Status BufferAssigner::RunAssignBuffers(
         sequential_algorithm,
     buffer_assignment::AssignmentAlgorithmForComputationsWithoutOrderingProto::
         Value non_sequential_algorithm,
-    const absl::flat_hash_set<LogicalBuffer::Color>* colors_to_assign) {
+    const absl::flat_hash_set<LogicalBuffer::Color>* colors_to_assign,
+    const flat_hash_map<const HloComputation*, flat_hash_set<const HloValue*>>*
+        precomputed_sequential_buffers) {
   // First assign buffers for global computations. Temporary buffers for
   // sequential computations are collected in
-  // 'buffers_to_assign_sequentially'.
+  // 'buffers_to_assign_sequentially'. Skipped if the caller has already done
+  // this and passed the collected buffers in `precomputed_sequential_buffers`.
   flat_hash_map<const HloComputation*, flat_hash_set<const HloValue*>>
-      buffers_to_assign_sequentially;
-  ABSL_RETURN_IF_ERROR(AssignBuffersForGlobalComputations(
-      global_computations, &buffers_to_assign_sequentially, assignment,
-      non_sequential_algorithm, colors_to_assign));
+      local_buffers_to_assign_sequentially;
+  if (precomputed_sequential_buffers == nullptr) {
+    ABSL_RETURN_IF_ERROR(AssignBuffersForGlobalComputations(
+        global_computations, &local_buffers_to_assign_sequentially, assignment,
+        non_sequential_algorithm, colors_to_assign));
+  }
+  const flat_hash_map<const HloComputation*, flat_hash_set<const HloValue*>>&
+      buffers_to_assign_sequentially =
+          precomputed_sequential_buffers != nullptr
+              ? *precomputed_sequential_buffers
+              : local_buffers_to_assign_sequentially;
   // Assign buffers with sequential ordering, if any. If all global
   // computations are sequential, we can run heap simulation on the whole
   // module, which reduces memory usage.
diff --git a/third_party/xla/xla/service/buffer_assignment.h b/third_party/xla/xla/service/buffer_assignment.h
index e8d98e0de1a..70ca294e8df 100644
--- a/third_party/xla/xla/service/buffer_assignment.h
+++ b/third_party/xla/xla/service/buffer_assignment.h
@@ -749,6 +749,44 @@ class BufferAssignment {
     stats_ = Stats();
   }
 
+  // Encapsulates the state of the buffer assignment.
+  struct State {
+    std::vector<BufferAllocation> allocations;
+    int64_t temp_allocation_total_size = 0;
+    std::vector<BufferAllocation::Index> allocation_index_for_value_id;
+    size_t num_assigned_values = 0;
+    absl::btree_map<LogicalBuffer::Color, int64_t> fragmentation_bytes_by_color;
+    Stats stats;
+  };
+
+  // Takes a snapshot of the current state.
+  State GetState() const {
+    return {allocations_,
+            temp_allocation_total_size_,
+            allocation_index_for_value_id_,
+            num_assigned_values_,
+            fragmentation_bytes_by_color_,
+            stats_};
+  }
+
+  // Restores the state from a snapshot.
+  void RestoreState(State state) {
+    allocations_ = std::move(state.allocations);
+    temp_allocation_total_size_ = state.temp_allocation_total_size;
+    allocation_index_for_value_id_ =
+        std::move(state.allocation_index_for_value_id);
+    num_assigned_values_ = state.num_assigned_values;
+    fragmentation_bytes_by_color_ =
+        std::move(state.fragmentation_bytes_by_color);
+    stats_ = state.stats;
+  }
+
+  // Replaces allocations and fragmentation stats for `colors` with those from
+  // `state`, leaving all other colors untouched.
+  void RestoreStateForColors(
+      const State& state,
+      const absl::flat_hash_set<LogicalBuffer::Color>& colors);
+
   // Clears only the buffer allocations belonging to `colors`.
   void ClearAllocationsForColors(
       const absl::flat_hash_set<LogicalBuffer::Color>& colors);
@@ -1012,6 +1050,14 @@ class BufferAllocationsManagerForComputationsWithoutOrdering {
 // A class which constructs a buffer assignment.
 class BufferAssigner {
  public:
+  // Specifies the buffer assignment algorithms for
+  // AssignBuffersForGlobalComputations and AssignBuffersWithSequentialOrdering.
+  struct AssignmentStrategy {
+    buffer_assignment::AssignmentAlgorithmForComputationsWithoutOrderingProto::
+        Value non_sequential_algorithm;
+    buffer_assignment::BufferAssignmentAlgorithmProto::Value
+        sequential_algorithm;
+  };
   friend class BufferAssignmentTest;
   using Colorer =
       std::function<absl::Status(HloAliasAnalysis*, const HloOrdering&)>;
@@ -1193,6 +1239,17 @@ class BufferAssigner {
                          LogicalBuffer::Color color) const;
 
  private:
+  // Assigns buffers of `colors_to_assign` (all colors if null) in three steps:
+  //  1. AssignBuffersForGlobalComputations with `non_sequential_algorithm`,
+  //     which also collects the temp buffers of sequentially ordered
+  //     computations,
+  //  2. heap simulation of the collected buffers with `sequential_algorithm`
+  //     (AssignBuffersWithSequentialOrdering),
+  //  3. thread-local computations, live-out marking and
+  //     CombineTempAllocations.
+  // If `precomputed_sequential_buffers` is set, step 1 is skipped and step 2
+  // uses the given buffers instead. The caller must have already run step 1 on
+  // `assignment`.
   absl::Status RunAssignBuffers(
       const HloModule* module,
       const std::vector<const HloComputation*>& global_computations,
@@ -1204,7 +1261,10 @@ class BufferAssigner {
           AssignmentAlgorithmForComputationsWithoutOrderingProto::Value
               non_sequential_algorithm,
       const absl::flat_hash_set<LogicalBuffer::Color>* colors_to_assign =
-          nullptr);
+          nullptr,
+      const absl::flat_hash_map<const HloComputation*,
+                                absl::flat_hash_set<const HloValue*>>*
+          precomputed_sequential_buffers = nullptr);
 
   absl::Status RunAssignBuffersWithFallback(
       const HloModule* module,
diff --git a/third_party/xla/xla/service/buffer_assignment_test.cc b/third_party/xla/xla/service/buffer_assignment_test.cc
index d0552870691..60bb3cced0c 100644
--- a/third_party/xla/xla/service/buffer_assignment_test.cc
+++ b/third_party/xla/xla/service/buffer_assignment_test.cc
@@ -949,6 +949,222 @@ TEST_F(BufferAssignmentTest, OOMFallbackOnlyForExceededColor) {
   EXPECT_EQ(color2_bytes, 400);
 }
 
+// Verifies which assignment the fallback keeps, using the chain from
+// OOMFallbackToDefault. The strategies are evaluated in this order:
+//  * FAST_MERGE (408 bytes)
+//  * FAST_SPLIT (288 bytes)
+//  * the fallback algorithm
+// The fallback algorithm is set to FAST_MERGE (408 bytes), so the FAST_SPLIT
+// result can be told apart from the result of the last strategy.
+TEST_F(BufferAssignmentTest, OOMFallbackStrategySelection) {
+  auto builder = HloComputation::Builder(TestName());
+  auto param = builder.AddInstruction(
+      HloInstruction::CreateParameter(0, ShapeUtil::MakeShape(F32, {}), "p1"));
+  auto a = builder.AddInstruction(HloInstruction::CreateCustomCall(
+      ShapeUtil::MakeShape(F32, {10}), {param}, "dummy"));
+  auto b = builder.AddInstruction(HloInstruction::CreateCustomCall(
+      ShapeUtil::MakeShape(F32, {20}), {a}, "dummy"));
+  auto c = builder.AddInstruction(HloInstruction::CreateCustomCall(
+      ShapeUtil::MakeShape(F32, {30}), {b}, "dummy"));
+  auto d = builder.AddInstruction(HloInstruction::CreateCustomCall(
+      ShapeUtil::MakeShape(F32, {40}), {c}, "dummy"));
+  auto root = builder.AddInstruction(HloInstruction::CreateCustomCall(
+      ShapeUtil::MakeShape(F32, {1}), {d}, "dummy"));
+  auto module = CreateNewVerifiedModule();
+  module->AddEntryComputation(builder.Build());
+
+  HloSchedule schedule(module.get());
+  schedule.set_sequence(module->entry_computation(), {param, a, b, c, d, root});
+  CHECK_OK(module->set_schedule(schedule));
+
+  constexpr int64_t kSafetyMargin = int64_t{5} << 29;
+  auto run = [&](int64_t limit) {
+    BufferAssigner::Options opts;
+    opts.enable_fallback = true;
+    opts.assignment_algorithm_for_computations_without_ordering =
+        buffer_assignment::
+            AssignmentAlgorithmForComputationsWithoutOrderingProto::FAST_MERGE;
+    opts.buffer_assignment_algorithm =
+        buffer_assignment::BufferAssignmentAlgorithmProto::FAST_MERGE;
+    opts.fallback_algorithm =
+        buffer_assignment::BufferAssignmentAlgorithmProto::FAST_MERGE;
+    opts.color_memory_limit = [limit](LogicalBuffer::Color) {
+      return kSafetyMargin + limit;
+    };
+    return BufferAssigner::Run(
+        module.get(), std::make_unique<SequentialHloOrdering>(schedule),
+        &BufferSizeBytes, &alias_info_, [](LogicalBuffer::Color) { return 1; },
+        std::move(opts));
+  };
+
+  // FAST_MERGE fits: it is kept, although FAST_SPLIT would use less memory.
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<BufferAssignment> fast_merge_fits,
+                       run(/*limit=*/1000));
+  EXPECT_EQ(fast_merge_fits->GetStats().total_allocation_bytes, 408);
+
+  // FAST_MERGE exceeds the limit and FAST_SPLIT fits: FAST_SPLIT is kept.
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<BufferAssignment> fast_split_fits,
+                       run(/*limit=*/300));
+  EXPECT_EQ(fast_split_fits->GetStats().total_allocation_bytes, 288);
+
+  // No strategy fits: the assignment with the lowest memory usage (FAST_SPLIT)
+  // is restored, although the fallback algorithm was evaluated last.
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<BufferAssignment> none_fits,
+                       run(/*limit=*/10));
+  EXPECT_EQ(none_fits->GetStats().total_allocation_bytes, 288);
+}
+
+// Verifies that clearing and restoring the allocations of one color from a
+// snapshot leaves the other colors untouched and keeps the indices consistent.
+TEST_F(BufferAssignmentTest, RestoreStateForColorsOnlyReplacesGivenColors) {
+  auto with_space = [](Shape s, int64_t space) {
+    s.mutable_layout()->set_memory_space(space);
+    return s;
+  };
+  auto builder = HloComputation::Builder(TestName());
+  auto param = builder.AddInstruction(
+      HloInstruction::CreateParameter(0, ShapeUtil::MakeShape(F32, {}), "p1"));
+  // A chain of two buffers in each of memory spaces 0 and 2.
+  auto a0 = builder.AddInstruction(HloInstruction::CreateCustomCall(
+      with_space(ShapeUtil::MakeShape(F32, {10}), 0), {param}, "dummy"));
+  auto b0 = builder.AddInstruction(HloInstruction::CreateCustomCall(
+      with_space(ShapeUtil::MakeShape(F32, {20}), 0), {a0}, "dummy"));
+  auto a2 = builder.AddInstruction(HloInstruction::CreateCustomCall(
+      with_space(ShapeUtil::MakeShape(F32, {10}), 2), {param}, "dummy"));
+  auto b2 = builder.AddInstruction(HloInstruction::CreateCustomCall(
+      with_space(ShapeUtil::MakeShape(F32, {20}), 2), {a2}, "dummy"));
+  auto root = builder.AddInstruction(HloInstruction::CreateCustomCall(
+      with_space(ShapeUtil::MakeShape(F32, {1}), 0), {b0, b2}, "dummy"));
+  auto module = CreateNewVerifiedModule();
+  module->AddEntryComputation(builder.Build());
+
+  HloSchedule schedule(module.get());
+  schedule.set_sequence(module->entry_computation(),
+                        {param, a0, b0, a2, b2, root});
+  CHECK_OK(module->set_schedule(schedule));
+
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<BufferAssignment> assignment,
+      BufferAssigner::Run(
+          module.get(), std::make_unique<SequentialHloOrdering>(schedule),
+          &BufferSizeBytes, &alias_info_,
+          [](LogicalBuffer::Color) { return 1; }, BufferAssigner::Options()));
+
+  auto bytes_of_color = [&](LogicalBuffer::Color color) {
+    int64_t bytes = 0;
+    for (const BufferAllocation& allocation : assignment->Allocations()) {
+      if (allocation.color() == color) {
+        bytes += allocation.size();
+      }
+    }
+    return bytes;
+  };
+  const int64_t color0_bytes = bytes_of_color(0);
+  const int64_t color2_bytes = bytes_of_color(2);
+  ASSERT_GT(color0_bytes, 0);
+  ASSERT_GT(color2_bytes, 0);
+  const int64_t temp_allocation_total_size =
+      assignment->temp_allocation_total_size();
+  const BufferAssignment::State state = assignment->GetState();
+
+  assignment->ClearAllocationsForColors({0});
+  EXPECT_EQ(bytes_of_color(0), 0);
+  EXPECT_EQ(bytes_of_color(2), color2_bytes);
+  // All allocations of color 2 are temp allocations.
+  EXPECT_EQ(assignment->temp_allocation_total_size(), color2_bytes);
+
+  assignment->RestoreStateForColors(state, {0});
+  EXPECT_EQ(bytes_of_color(0), color0_bytes);
+  EXPECT_EQ(bytes_of_color(2), color2_bytes);
+  EXPECT_EQ(assignment->temp_allocation_total_size(),
+            temp_allocation_total_size);
+  // The allocation indices and the value to allocation map are rebuilt.
+  for (BufferAllocation::Index index = 0;
+       index < assignment->Allocations().size(); ++index) {
+    const BufferAllocation& allocation = assignment->GetAllocation(index);
+    EXPECT_EQ(allocation.index(), index);
+    // NOLINTNEXTLINE(*-custom-deterministic-iteration-order)
+    for (const auto& [value, offset_size] : allocation.assigned_buffers()) {
+      EXPECT_EQ(assignment->GetAssignedAllocation(*value).index(), index);
+    }
+  }
+}
+
+// When no strategy fits all colors, each color keeps its own best assignment.
+//  * Color 0 is the chain from OOMFallbackToDefault: FAST_MERGE (408 bytes)
+//    exceeds its limit, and FAST_SPLIT (288 bytes) fits.
+//  * Color 2 is a chain of more than 500 buffers of the same size, for which
+//    FAST_SPLIT uses more memory than FAST_MERGE (800 bytes). Its expensive
+//    pass handles 500 buffers, and the others are placed on top.
+// Color 2 never fits, but the total memory usage over all colors is lowest
+// with FAST_MERGE, where color 0 doesn't fit.
+TEST_F(BufferAssignmentTest, OOMFallbackKeepsBestAssignmentPerColor) {
+  auto with_space = [](Shape s, int64_t space) {
+    s.mutable_layout()->set_memory_space(space);
+    return s;
+  };
+  auto builder = HloComputation::Builder(TestName());
+  std::vector<HloInstruction*> sequence;
+  auto add = [&](std::unique_ptr<HloInstruction> instruction) {
+    sequence.push_back(builder.AddInstruction(std::move(instruction)));
+    return sequence.back();
+  };
+  HloInstruction* param = add(
+      HloInstruction::CreateParameter(0, ShapeUtil::MakeShape(F32, {}), "p1"));
+  HloInstruction* chain0 = param;
+  for (int64_t size : {10, 20, 30, 40}) {
+    chain0 = add(HloInstruction::CreateCustomCall(
+        with_space(ShapeUtil::MakeShape(F32, {size}), 0), {chain0}, "dummy"));
+  }
+  HloInstruction* chain2 = param;
+  for (int i = 0; i < 502; ++i) {
+    chain2 = add(HloInstruction::CreateCustomCall(
+        with_space(ShapeUtil::MakeShape(F32, {100}), 2), {chain2}, "dummy"));
+  }
+  add(HloInstruction::CreateCustomCall(
+      with_space(ShapeUtil::MakeShape(F32, {1}), 0), {chain0, chain2},
+      "dummy"));
+  auto module = CreateNewVerifiedModule();
+  module->AddEntryComputation(builder.Build());
+
+  HloSchedule schedule(module.get());
+  schedule.set_sequence(module->entry_computation(), sequence);
+  CHECK_OK(module->set_schedule(schedule));
+
+  // The strategies are FAST_MERGE, FAST_SPLIT, and the fallback, FAST_SPLIT.
+  BufferAssigner::Options opts;
+  opts.enable_fallback = true;
+  opts.assignment_algorithm_for_computations_without_ordering =
+      buffer_assignment::
+          AssignmentAlgorithmForComputationsWithoutOrderingProto::FAST_MERGE;
+  opts.buffer_assignment_algorithm =
+      buffer_assignment::BufferAssignmentAlgorithmProto::FAST_MERGE;
+  opts.fallback_algorithm =
+      buffer_assignment::BufferAssignmentAlgorithmProto::FAST_SPLIT;
+  constexpr int64_t kSafetyMargin = int64_t{5} << 29;
+  opts.color_memory_limit = [](LogicalBuffer::Color color) -> int64_t {
+    return kSafetyMargin + (color == 0 ? 300 : 100);
+  };
+
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<BufferAssignment> assignment,
+      BufferAssigner::Run(
+          module.get(), std::make_unique<SequentialHloOrdering>(schedule),
+          &BufferSizeBytes, &alias_info_,
+          [](LogicalBuffer::Color) { return 1; }, std::move(opts)));
+
+  absl::flat_hash_map<LogicalBuffer::Color, int64_t> bytes_by_color;
+  for (const BufferAllocation& allocation : assignment->Allocations()) {
+    bytes_by_color[allocation.color()] += allocation.size();
+  }
+  // Color 0 keeps the FAST_SPLIT assignment, which fits, and color 2 gets its
+  // FAST_MERGE assignment back. The parameter and the root (4 bytes each) are
+  // the only allocations that are not temp allocations.
+  EXPECT_EQ(bytes_by_color[0], 288);
+  EXPECT_EQ(bytes_by_color[2], 800);
+  EXPECT_EQ(assignment->temp_allocation_total_size(), 288 - 8 + 800);
+}
+
 MATCHER(IdEq, "") {
   auto* actual = std::get<0>(arg);
   auto* expected = std::get<1>(arg);

From cc84f2e024e077854417a04e3a41e444a79ddc6a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Eusebio=20Dur=C3=A1n=20Monta=C3=B1a?= <eusebiodm@google.com>
Date: Thu, 15 Oct 2026 02:15:45 -0700
Subject: [PATCH 0861/1000] Remove unnecessary LLVM CodeGen deps from
 `cpu_aot_loader`

`cpu_aot_loader` never runs any LLVM optimization or CodeGen passes.

PiperOrigin-RevId: 999978131
---
 third_party/xla/xla/service/cpu/BUILD | 17 +----------------
 1 file changed, 1 insertion(+), 16 deletions(-)

diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD
index cbc6da2e898..3a47e7be030 100644
--- a/third_party/xla/xla/service/cpu/BUILD
+++ b/third_party/xla/xla/service/cpu/BUILD
@@ -1837,7 +1837,6 @@ cc_library(
         "//xla:util",
         "//xla/backends/cpu:target_machine_options",
         "//xla/backends/cpu/codegen:builtin_definition_generator",
-        "//xla/backends/cpu/codegen:cpu_features",
         "//xla/backends/cpu/codegen:execution_engine",
         "//xla/backends/cpu/codegen:object_loader",
         "//xla/backends/cpu/runtime:function_library",
@@ -1846,30 +1845,16 @@ cc_library(
         "//xla/service:compiler",
         "//xla/service:executable",
         "//xla/service:hlo_module_config",
-        "//xla/tsl/platform:errors",
-        "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/log",
-        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/types:span",
         "@llvm-project//llvm:Support",
         "@llvm-project//llvm:Target",
         "@llvm-project//llvm:TargetParser",
         "@llvm-project//llvm:ir_headers",
-    ] + if_llvm_aarch64_available([
-        "@llvm-project//llvm:AArch64CodeGen",  # fixdeps: keep
-    ]) + if_llvm_powerpc_available([
-        "@llvm-project//llvm:PowerPCCodeGen",  # fixdeps: keep
-    ]) + if_llvm_riscv_available([
-        "@llvm-project//llvm:RISCVCodeGen",  # fixdeps: keep
-    ]) + if_llvm_system_z_available([
-        "@llvm-project//llvm:SystemZCodeGen",  # fixdeps: keep
-    ]) + if_llvm_x86_available([
-        "@llvm-project//llvm:X86CodeGen",  # fixdeps: keep
-    ]),
+    ],
 )
 
 xla_cc_test(

From 496630d769310b05fdca984cb8cdf80fdd7b5e37 Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Thu, 15 Oct 2026 02:26:14 -0700
Subject: [PATCH 0862/1000] Remove --end-of-options from git checkout in
 bazel_diff

When `change_detector.py` was extracted, `--end-of-options` was added alongside `change_detector.validate_git_ref` to the `git checkout -f` calls in `bazel_diff.py` to satisfy a static-analysis check for flag injection. However, prior to Git 2.43.1, `git checkout` did not consume `--end-of-options` during option parsing and treated both `--end-of-options` and the following commit SHA as pathspecs. Since the OpenXLA CI container runs Ubuntu 22.04 with Git 2.34.1, `git checkout -f --end-of-options <sha>` failed on every PR, and because `compute_impacted_targets` is designed to fail open to a full test run on any error, the regression went unnoticed while CI jobs kept passing.

Since all revision arguments are already validated by `change_detector.validate_git_ref`, `--end-of-options` is not needed here. Therefore this change removes `--end-of-options` from the `git checkout -f` calls in `compute_impacted_targets`.

PiperOrigin-RevId: 999983661
---
 third_party/xla/build_tools/ci/bazel_diff.py  |   6 +-
 .../xla/build_tools/ci/bazel_diff_test.py     | 109 ++++++++++++++++++
 2 files changed, 112 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/build_tools/ci/bazel_diff.py b/third_party/xla/build_tools/ci/bazel_diff.py
index d7cb221508d..c071fbf19e8 100644
--- a/third_party/xla/build_tools/ci/bazel_diff.py
+++ b/third_party/xla/build_tools/ci/bazel_diff.py
@@ -547,7 +547,7 @@ def compute_impacted_targets(
 
     # Generate base hashes
     subprocess.run(
-        ["git", "checkout", "-f", "--end-of-options", diff_base],
+        ["git", "checkout", "-f", diff_base],
         cwd=workspace_dir,
         check=True,
         capture_output=True,
@@ -563,7 +563,7 @@ def compute_impacted_targets(
 
     # Generate head hashes
     subprocess.run(
-        ["git", "checkout", "-f", "--end-of-options", safe_head],
+        ["git", "checkout", "-f", safe_head],
         cwd=workspace_dir,
         check=True,
         capture_output=True,
@@ -601,7 +601,7 @@ def compute_impacted_targets(
     try:
       safe_original_head = change_detector.validate_git_ref(original_head)
       change_detector.run_git_best_effort(
-          ["checkout", "-f", "--end-of-options", safe_original_head],
+          ["checkout", "-f", safe_original_head],
           cwd=workspace_dir,
       )
     except ValueError:
diff --git a/third_party/xla/build_tools/ci/bazel_diff_test.py b/third_party/xla/build_tools/ci/bazel_diff_test.py
index 0bb6384ac60..2c9d3755c58 100644
--- a/third_party/xla/build_tools/ci/bazel_diff_test.py
+++ b/third_party/xla/build_tools/ci/bazel_diff_test.py
@@ -14,6 +14,7 @@
 # ============================================================================
 """Tests for bazel_diff module."""
 
+import dataclasses
 import hashlib
 import os
 import subprocess
@@ -27,6 +28,28 @@ from build_tools.ci import bazel_diff
 from build_tools.ci import change_detector
 
 
+@dataclasses.dataclass(frozen=True)
+class _FakeBuildConfig:
+  target_patterns: tuple[str, ...] = ("//xla/...",)
+  configs: tuple[str, ...] = ()
+  repo_env: dict[str, str] = dataclasses.field(default_factory=dict)
+  options: dict[str, object] = dataclasses.field(default_factory=dict)
+  startup_options: dict[str, object] = dataclasses.field(default_factory=dict)
+  bazel_diff_use_cquery: bool = True
+
+
+def _write_impacted_targets(
+    jar_path: str,
+    workspace_dir: str,
+    base_hashes: str,
+    head_hashes: str,
+    output_path: str,
+) -> None:
+  del jar_path, workspace_dir, base_hashes, head_hashes
+  with open(output_path, "w", encoding="utf-8") as f:
+    f.write("//xla/service/cpu:cpu_compiler_test\n")
+
+
 class BazelDiffTest(parameterized.TestCase):
 
   @parameterized.parameters(
@@ -336,6 +359,92 @@ class BazelDiffTest(parameterized.TestCase):
     self.assertEqual(decision.decision, bazel_diff.BazelDiffDecisionType.FULL)
     self.assertIn("Failed to get git changed files", decision.reason)
 
+  @mock.patch.object(
+      bazel_diff, "run_get_impacted_targets", autospec=True, spec_set=True
+  )
+  @mock.patch.object(
+      bazel_diff, "run_generate_hashes", autospec=True, spec_set=True
+  )
+  @mock.patch.object(
+      bazel_diff, "download_and_verify_jar", autospec=True, spec_set=True
+  )
+  @mock.patch.object(bazel_diff.subprocess, "run", autospec=True, spec_set=True)
+  @mock.patch.object(
+      change_detector, "get_diff_base", autospec=True, spec_set=True
+  )
+  @mock.patch.object(
+      change_detector, "get_changed_files", autospec=True, spec_set=True
+  )
+  @mock.patch.object(
+      change_detector, "ensure_base_fetched", autospec=True, spec_set=True
+  )
+  @mock.patch.object(
+      change_detector, "run_git_best_effort", autospec=True, spec_set=True
+  )
+  def test_compute_impacted_targets_checkout_and_restore(
+      self,
+      mock_git_best_effort,
+      mock_ensure_fetched,
+      mock_changed_files,
+      mock_diff_base,
+      mock_subprocess_run,
+      mock_download_jar,
+      mock_generate_hashes,
+      mock_get_impacted,
+  ):
+    del mock_ensure_fetched, mock_generate_hashes
+    mock_git_best_effort.side_effect = [
+        subprocess.CompletedProcess(
+            args=["git", "rev-parse", "HEAD"],
+            returncode=0,
+            stdout="orig_head_sha\n",
+        ),
+        subprocess.CompletedProcess(
+            args=["git", "checkout", "-f", "orig_head_sha"],
+            returncode=0,
+            stdout="",
+        ),
+    ]
+    mock_changed_files.return_value = ["xla/service/cpu/cpu_compiler.cc"]
+    mock_diff_base.return_value = "merge_base_sha"
+    mock_download_jar.return_value = "/tmp/bazel-diff.jar"
+    mock_get_impacted.side_effect = _write_impacted_targets
+
+    decision = bazel_diff.compute_impacted_targets(
+        _FakeBuildConfig(), "base_sha", "HEAD", workspace_dir="/workspace"
+    )
+
+    self.assertEqual(
+        decision.decision, bazel_diff.BazelDiffDecisionType.IMPACTED
+    )
+    self.assertEqual(
+        decision.impacted_targets, ("//xla/service/cpu:cpu_compiler_test",)
+    )
+    self.assertEqual(
+        mock_subprocess_run.call_args_list,
+        [
+            mock.call(
+                ["git", "checkout", "-f", "merge_base_sha"],
+                cwd="/workspace",
+                check=True,
+                capture_output=True,
+            ),
+            mock.call(
+                ["git", "checkout", "-f", "orig_head_sha"],
+                cwd="/workspace",
+                check=True,
+                capture_output=True,
+            ),
+        ],
+    )
+    self.assertEqual(
+        mock_git_best_effort.call_args_list,
+        [
+            mock.call(["rev-parse", "HEAD"], cwd="/workspace"),
+            mock.call(["checkout", "-f", "orig_head_sha"], cwd="/workspace"),
+        ],
+    )
+
 
 if __name__ == "__main__":
   absltest.main()

From 075d0a4f8a2ccc31d39fc1a6c259a31a80fa4dd2 Mon Sep 17 00:00:00 2001
From: Sevin Varoglu <svaroglu@nvidia.com>
Date: Thu, 15 Oct 2026 03:54:26 -0700
Subject: [PATCH 0863/1000] PR #50088: [XLA:GPU] ragged_all_to_all dev_kernel:
 fire remote GIN puts before LSA copy
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50088

📝 Summary of Changes
Reorders `RaggedAllToAllCopy` so remote-peer GIN puts are submitted first, then the LSA copies run — fabric ingestion overlaps with NVLink stores instead of serializing after them. `RaggedAllToAllDeviceKernelImpl` still waits on the GIN signal after `RaggedAllToAllCopy` returns, so exit-side ordering is unchanged.

🎯 Justification
Today, LSA finishes before GIN starts. Submitting GIN first lets fabric latency overlap with LSA.

🚀 Kind of Contribution
⚡️ Performance Improvement

📊 Benchmark (for Performance Improvements)
N/A

🧪 Unit Tests:
N/A

🧪 Execution Tests:
N/A
Copybara import of the project:

--
2da5d0fbf32e140d7e2456e42acd2d08fbd8d62e by Sevin Varoglu <svaroglu@nvidia.com>:

[XLA:GPU] ragged_all_to_all dev_kernel: fire remote GIN puts before LSA copy

Submits the remote-peer GIN puts first in RaggedAllToAllCopy, then does
the LSA copy loop. Fabric ingestion overlaps with NVLink stores instead
of running back-to-back.

The submission is async and returns as soon as each descriptor is
accepted, so Phase 1 completes quickly and Phase 2 (LSA) begins while
puts are still in flight to remote peers. The caller in
RaggedAllToAllDeviceKernelImpl still waits on the GIN signal after
RaggedAllToAllCopy returns, so the exit-side ordering is unchanged; only
the intra-kernel schedule moves.

No signature change, no thunk plumbing, no interaction with the acquire
or release barriers. LSA-only shapes are unaffected (gin == nullptr
branch skips Phase 1).

Expected gain caps at min(LSA_time, fabric_latency); most useful on
shapes where the two are comparable (moderate remote fanout with medium
LSA payload). Small-remote-fanout or LSA-dominated shapes see little
change.

Merging this change closes #50088

PiperOrigin-RevId: 1000028566
---
 .../ragged_all_to_all_device_kernel_lib.cu.h  | 53 +++++++++----------
 1 file changed, 26 insertions(+), 27 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/gpu/ragged_all_to_all_device_kernel_lib.cu.h b/third_party/xla/xla/stream_executor/gpu/ragged_all_to_all_device_kernel_lib.cu.h
index 2fbc76f4e6f..90dcfe7e87d 100644
--- a/third_party/xla/xla/stream_executor/gpu/ragged_all_to_all_device_kernel_lib.cu.h
+++ b/third_party/xla/xla/stream_executor/gpu/ragged_all_to_all_device_kernel_lib.cu.h
@@ -80,6 +80,32 @@ __device__ void RaggedAllToAllCopy(
     unsigned int signal_index) {
   using T = DeviceVec<kVectorSize>;
 
+  if (gin != nullptr) {
+    const int tid = threadIdx.x + blockIdx.x * blockDim.x;
+    const int nthreads = blockDim.x * gridDim.x;
+    const int64_t total_updates = num_updates_per_replica * num_ranks;
+
+    for (int64_t flat_idx = tid; flat_idx < total_updates;
+         flat_idx += nthreads) {
+      const int peer = flat_idx / num_updates_per_replica;
+      if (peer >= start_lsa && peer < start_lsa + lsa_size) {
+        continue;
+      }
+
+      RaggedAllToAllUpdateMetadata<kVectorSize> meta;
+      if (!LoadRaggedAllToAllUpdateMetadata<kVectorSize>(
+              flat_idx, num_updates_per_replica, num_row_elements,
+              input_buffer_offset_bytes, output_buffer_offset_bytes,
+              input_offsets_ptr, send_sizes_ptr, output_offsets_ptr, &meta)) {
+        continue;
+      }
+
+      gin->put(world, meta.peer, recv_win, meta.dst_byte_offset, send_win,
+               meta.src_byte_offset, meta.byte_count,
+               ncclGin_SignalInc{signal_index});
+    }
+  }
+
   if (lsa_size > 0) {
     // Work-proportional CTA assignment: CTA k copies the global element range
     // [total*k/grid, total*(k+1)/grid), walking update boundaries as needed.
@@ -140,33 +166,6 @@ __device__ void RaggedAllToAllCopy(
       update_begin = update_end;
     }
   }
-
-  if (gin == nullptr) {
-    return;
-  }
-
-  const int tid = threadIdx.x + blockIdx.x * blockDim.x;
-  const int nthreads = blockDim.x * gridDim.x;
-  const int64_t total_updates = num_updates_per_replica * num_ranks;
-
-  for (int64_t flat_idx = tid; flat_idx < total_updates; flat_idx += nthreads) {
-    const int peer = flat_idx / num_updates_per_replica;
-    if (peer >= start_lsa && peer < start_lsa + lsa_size) {
-      continue;
-    }
-
-    RaggedAllToAllUpdateMetadata<kVectorSize> meta;
-    if (!LoadRaggedAllToAllUpdateMetadata<kVectorSize>(
-            flat_idx, num_updates_per_replica, num_row_elements,
-            input_buffer_offset_bytes, output_buffer_offset_bytes,
-            input_offsets_ptr, send_sizes_ptr, output_offsets_ptr, &meta)) {
-      continue;
-    }
-
-    gin->put(world, meta.peer, recv_win, meta.dst_byte_offset, send_win,
-             meta.src_byte_offset, meta.byte_count,
-             ncclGin_SignalInc{signal_index});
-  }
 }
 
 template <int64_t kVectorSize>

From 7372c807887111dffdafc0088ec431bc6abdbe1a Mon Sep 17 00:00:00 2001
From: Oleg Shyshkov <shyshkov@google.com>
Date: Thu, 15 Oct 2026 04:58:00 -0700
Subject: [PATCH 0864/1000] [XLA:GPU] Inline collective launch dimension
 calculations into Triton collective emitter.

PiperOrigin-RevId: 1000055201
---
 .../xla/xla/backends/gpu/codegen/triton/BUILD |   2 -
 .../gpu/codegen/triton/collective_emitter.cc  | 357 +++++++++---------
 2 files changed, 168 insertions(+), 191 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/BUILD b/third_party/xla/xla/backends/gpu/codegen/triton/BUILD
index 2af3c3b8b7c..bd13299ce68 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/BUILD
@@ -759,7 +759,6 @@ cc_library(
         "//xla/backends/gpu/codegen/triton/ir:triton_xla",
         "//xla/backends/gpu/runtime:all_gather",
         "//xla/backends/gpu/runtime:all_reduce",
-        "//xla/backends/gpu/runtime:collective_info_utils",
         "//xla/backends/gpu/runtime:collective_params",
         "//xla/backends/gpu/runtime:reduce_scatter",
         "//xla/backends/gpu/transforms/collectives:collective_ops_utils",
@@ -781,7 +780,6 @@ cc_library(
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/functional:any_invocable",
-        "@com_google_absl//absl/functional:overload",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/log:vlog_is_on",
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
index 232d20bd061..46164b66cc9 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
@@ -21,13 +21,11 @@ limitations under the License.
 #include <numeric>
 #include <string>
 #include <utility>
-#include <variant>
 #include <vector>
 
 #include "absl/algorithm/container.h"
 #include "absl/base/nullability.h"
 #include "absl/functional/any_invocable.h"
-#include "absl/functional/overload.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/log/vlog_is_on.h"
@@ -64,7 +62,6 @@ limitations under the License.
 #include "xla/backends/gpu/codegen/triton/lowering_util.h"
 #include "xla/backends/gpu/runtime/all_gather.h"
 #include "xla/backends/gpu/runtime/all_reduce.h"
-#include "xla/backends/gpu/runtime/collective_info_utils.h"
 #include "xla/backends/gpu/runtime/collective_params.h"
 #include "xla/backends/gpu/runtime/reduce_scatter.h"
 #include "xla/backends/gpu/transforms/collectives/collective_ops_utils.h"
@@ -301,139 +298,6 @@ absl::StatusOr<ReduceScatterEmitterContext> CreateReduceScatterEmitterContext(
   return ctx;
 }
 
-// Returns the launch dimensions for the all-reduce kernel.
-// The launch dimensions are determined by the number of elements, the
-// all-reduce strategy, and the warp size of the target device.
-LaunchDimensions AllReduceLaunchDimensions(
-    const AllReduceInfo& all_reduce_info,
-    const se::DeviceDescription& device_info) {
-  // Heuristic maxima after some benchmarking.
-  // It's nicer to have this as a power of 2 because we consume blocks in powers
-  // of 2 from a higher dimension to a lower dimension.
-  static constexpr int64_t kMaxBlocksPerGrid = 32;
-  static constexpr uint64_t kMaxThreadsPerBlock = 512;
-
-  int64_t threads_per_block;
-  int64_t blocks_per_grid;
-  const int64_t warp_size = WarpSize(device_info);
-  const int64_t elements_per_rank =
-      all_reduce_info.num_elements /
-      (all_reduce_info.all_reduce_strategy == AllReduceStrategy::kTwoShot
-           ? all_reduce_info.num_devices
-           : 1);
-  // Maximum number of threads such that each thread has elements to process.
-  const int64_t total_threads =
-      RoundUpTo(CeilOfRatio(elements_per_rank, se::gpu::kNumElementsPerThread),
-                warp_size);
-  // Triton expects power of 2 for threads_per_block / threads_per_warp.
-  // Since threads_per_warp is always a power of 2, this is guaranteed.
-  threads_per_block =
-      std::min(kMaxThreadsPerBlock,
-               llvm::bit_ceil(static_cast<uint64_t>(total_threads)));
-  blocks_per_grid = std::min(kMaxBlocksPerGrid,
-                             CeilOfRatio(total_threads, threads_per_block));
-  return LaunchDimensions(blocks_per_grid, threads_per_block);
-}
-
-// Returns the launch dimensions for the all-gather kernel.
-// All-gather uses a one-shot strategy: each rank contributes its local slice
-// to the symmetric buffer and then reads each peer's slice.
-// warp_size should be device_description.threads_per_warp() (32 on NVIDIA,
-// 64 on AMD) so that the thread count is a multiple of the hardware warp.
-LaunchDimensions AllGatherLaunchDimensions(
-    const AllGatherInfo& all_gather_info,
-    const se::DeviceDescription& device_info) {
-  // All-gather always uses a one-shot strategy: every rank reads its peers'
-  // slices directly, so the work is proportional to the full element count
-  // with no rank-based partitioning.
-  static constexpr uint64_t kMaxThreadsPerBlock = 512;
-
-  // Maximum number of threads such that each thread has elements to process.
-  // Round up to a multiple of warp_size so every warp is fully occupied.
-  const int64_t output_elements =
-      all_gather_info.num_elements * all_gather_info.num_devices;
-  // Assuming each thread writes se::gpu::kNumElementsPerThread output elements
-  // we need total threads corresponding to the output elements.
-  const int64_t total_threads =
-      RoundUpTo(CeilOfRatio(output_elements, se::gpu::kNumElementsPerThread),
-                device_info.threads_per_warp());
-  // Triton expects power of 2 for threads_per_block / threads_per_warp.
-  const int64_t threads_per_block =
-      std::min(kMaxThreadsPerBlock,
-               llvm::bit_ceil(static_cast<uint64_t>(total_threads)));
-  const int64_t blocks_per_grid =
-      std::min(kAllGatherMaxBlocksPerGrid,
-               CeilOfRatio(total_threads, threads_per_block));
-  return LaunchDimensions(blocks_per_grid, threads_per_block);
-}
-
-// Returns the launch dimensions for the reduce-scatter kernel.
-LaunchDimensions ReduceScatterLaunchDimensions(
-    const ReduceScatterInfo& reduce_scatter_info,
-    const se::DeviceDescription& device_info) {
-  constexpr int64_t kNumElementsPerThread = 4;
-  constexpr uint64_t kMaxThreadsPerBlock = 512;
-
-  const int64_t warp_size = device_info.threads_per_warp();
-  const int64_t elements_per_rank_tile = CeilOfRatio(
-      reduce_scatter_info.num_output_elements, reduce_scatter_info.num_devices);
-  const int64_t threads_per_rank_tile = RoundUpTo(
-      CeilOfRatio(elements_per_rank_tile, kNumElementsPerThread), warp_size);
-  const int64_t threads_per_block = static_cast<int64_t>(
-      std::min(kMaxThreadsPerBlock,
-               llvm::bit_ceil(static_cast<uint64_t>(threads_per_rank_tile))));
-  const int64_t blocks_per_rank_tile =
-      CeilOfRatio(threads_per_rank_tile, threads_per_block);
-  const int64_t blocks_per_grid = std::min(
-      kReduceScatterMaxBlocksPerGrid,
-      reduce_scatter_info.num_devices *
-          static_cast<int64_t>(
-              llvm::bit_floor(static_cast<uint64_t>(blocks_per_rank_tile))));
-  return LaunchDimensions(blocks_per_grid, threads_per_block);
-}
-
-absl::StatusOr<LaunchDimensions> GetLaunchDimensions(
-    const CollectiveKernelInfo& collective_info,
-    const GpuTopology& gpu_topology) {
-  const se::DeviceDescription& device_info =
-      gpu_topology.gpu_target_config().device_description;
-  return std::visit(
-      absl::Overload{
-          [&](const AllReduceInfo& all_reduce_info) {
-            return AllReduceLaunchDimensions(all_reduce_info, device_info);
-          },
-          [&](const AllGatherInfo& all_gather_info) {
-            return AllGatherLaunchDimensions(all_gather_info, device_info);
-          },
-          [&](const ReduceScatterInfo& reduce_scatter_info) {
-            return ReduceScatterLaunchDimensions(reduce_scatter_info,
-                                                 device_info);
-          }},
-      collective_info);
-}
-
-absl::Status ValidateBlockLevelFusionConfig(
-    const BlockLevelFusionConfig& block_level_config,
-    const CollectiveKernelInfo& collective_info) {
-  TF_RET_CHECK(block_level_config.output_tiles_size() == 1)
-      << "expected 1 output tile, but got "
-      << block_level_config.output_tiles_size();
-  const int64_t linear_tile_size =
-      Product(block_level_config.output_tiles(0).sizes());
-  if (std::holds_alternative<AllReduceInfo>(collective_info)) {
-    const AllReduceInfo& all_reduce_info =
-        std::get<AllReduceInfo>(collective_info);
-    if (all_reduce_info.all_reduce_strategy == AllReduceStrategy::kTwoShot &&
-        linear_tile_size % all_reduce_info.num_devices != 0) {
-      return absl::InvalidArgumentError(absl::StrFormat(
-          "Two-shot all-reduce linear_tile_size(%d) %% num_devices(%d) != 0. "
-          "Codegen will not be supported.",
-          linear_tile_size, all_reduce_info.num_devices));
-    }
-  }
-  return absl::OkStatus();
-}
-
 // Unmanaged kernel arguments of collectives that use the AllReduce argument
 // layout: rank, signal value (except for reduce-scatter), signal buffers and
 // one remote buffer per parameter.
@@ -1354,6 +1218,145 @@ llvm::SmallVector<int64_t> GreedyPowerOfTwoTiles(const Shape& output_shape,
   return tile_sizes;
 }
 
+namespace {
+
+// Heuristic maximum warps per block across collective kernels.
+// Triton expects a power of 2 for num_warps.
+constexpr int64_t kMaxWarpsPerBlock = 16;
+
+BlockLevelFusionConfig CreateBlockLevelFusionConfig(
+    int64_t num_warps, absl::Span<const int64_t> tile_sizes) {
+  BlockLevelFusionConfig config;
+  config.set_num_warps(num_warps);
+  config.set_num_ctas(1);    // No block-level clustering.
+  config.set_num_stages(1);  // No pipelining of loops.
+  config.add_output_tiles()->mutable_sizes()->Assign(tile_sizes.begin(),
+                                                     tile_sizes.end());
+  return config;
+}
+
+absl::StatusOr<BlockLevelFusionConfig> GetBlockLevelConfigForAllReduce(
+    const GpuTopology& gpu_topology, const HloAllReduceInstruction* instr,
+    const DeviceAssignment* device_assignment) {
+  // Heuristic maximum blocks per grid after benchmarking. Power of 2 so blocks
+  // are consumed cleanly in powers of 2 from major to minor dimensions.
+  static constexpr int64_t kAllReduceMaxBlocksPerGrid = 32;
+
+  ABSL_ASSIGN_OR_RETURN(const AllReduceInfo all_reduce_info,
+                   BuildAllReduceInfo(/*is_multimem_enabled=*/false,
+                                      gpu_topology, instr, device_assignment));
+
+  const int64_t elements_per_warp =
+      WarpSize(gpu_topology.gpu_target_config().device_description) *
+      se::gpu::kNumElementsPerThread;
+  const int64_t elements_per_rank =
+      all_reduce_info.num_elements /
+      (all_reduce_info.all_reduce_strategy == AllReduceStrategy::kTwoShot
+           ? all_reduce_info.num_devices
+           : 1);
+  const int64_t total_warps = CeilOfRatio(elements_per_rank, elements_per_warp);
+  const int64_t num_warps = std::min<int64_t>(
+      kMaxWarpsPerBlock, llvm::bit_ceil(static_cast<uint64_t>(total_warps)));
+  const int64_t num_blocks =
+      std::min(kAllReduceMaxBlocksPerGrid, CeilOfRatio(total_warps, num_warps));
+
+  llvm::SmallVector<int64_t> tile_sizes =
+      GreedyPowerOfTwoTiles(instr->shape(), num_blocks);
+  const int64_t linear_tile_size = Product(tile_sizes);
+  if (all_reduce_info.all_reduce_strategy == AllReduceStrategy::kTwoShot &&
+      linear_tile_size % all_reduce_info.num_devices != 0) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        "Two-shot all-reduce linear_tile_size(%d) %% num_devices(%d) != 0. "
+        "Codegen will not be supported.",
+        linear_tile_size, all_reduce_info.num_devices));
+  }
+
+  return CreateBlockLevelFusionConfig(num_warps, tile_sizes);
+}
+
+absl::StatusOr<BlockLevelFusionConfig> GetBlockLevelConfigForAllGather(
+    const GpuTopology& gpu_topology, const HloAllGatherInstruction* instr,
+    const DeviceAssignment* device_assignment) {
+  ABSL_ASSIGN_OR_RETURN(const AllGatherInfo all_gather_info,
+                   BuildAllGatherInfo(gpu_topology, instr, device_assignment));
+
+  const Shape& output_shape = instr->shape();
+
+  const int64_t elements_per_warp =
+      WarpSize(gpu_topology.gpu_target_config().device_description) *
+      se::gpu::kNumElementsPerThread;
+  // All-gather always uses a one-shot strategy: every rank reads its peers'
+  // slices directly, so the work is proportional to the full output element
+  // count with no rank-based partitioning.
+  const int64_t output_elements =
+      all_gather_info.num_elements * all_gather_info.num_devices;
+  const int64_t total_warps = CeilOfRatio(output_elements, elements_per_warp);
+  const int64_t num_warps = std::min<int64_t>(
+      kMaxWarpsPerBlock, llvm::bit_ceil(static_cast<uint64_t>(total_warps)));
+  const int64_t num_blocks =
+      std::min(kAllGatherMaxBlocksPerGrid, CeilOfRatio(total_warps, num_warps));
+
+  llvm::SmallVector<int64_t> tile_sizes =
+      GreedyPowerOfTwoTiles(output_shape, num_blocks);
+
+  // For AllGather, the tile on the gather dimension must not exceed the
+  // per-rank size. Otherwise a single tile would span multiple replicas,
+  // but the tiling framework assigns a single replica_id per tile.
+  const int64_t gather_dim = instr->all_gather_dimension();
+  const int64_t per_rank_size =
+      output_shape.dimensions(gather_dim) / all_gather_info.num_devices;
+  tile_sizes[gather_dim] = std::min(tile_sizes[gather_dim],
+                                    static_cast<int64_t>(llvm::bit_floor(
+                                        static_cast<uint64_t>(per_rank_size))));
+
+  return CreateBlockLevelFusionConfig(num_warps, tile_sizes);
+}
+
+absl::StatusOr<BlockLevelFusionConfig> GetBlockLevelConfigForReduceScatter(
+    const GpuTopology& gpu_topology, const HloReduceScatterInstruction* instr,
+    const DeviceAssignment* device_assignment) {
+  ABSL_ASSIGN_OR_RETURN(
+      const ReduceScatterInfo reduce_scatter_info,
+      BuildReduceScatterInfo(gpu_topology, instr, device_assignment));
+
+  const Shape& output_shape = instr->shape();
+  const int64_t num_devices = reduce_scatter_info.num_devices;
+  const int64_t scatter_dim = instr->scatter_dimension();
+  TF_RET_CHECK(output_shape.dimensions(scatter_dim) % num_devices == 0)
+      << "Reduce-scatter output scatter dimension "
+      << output_shape.dimensions(scatter_dim)
+      << " is not divisible by the number of devices " << num_devices;
+
+  const int64_t elements_per_warp =
+      WarpSize(gpu_topology.gpu_target_config().device_description) *
+      se::gpu::kNumElementsPerThread;
+  const int64_t elements_per_rank_tile =
+      CeilOfRatio(reduce_scatter_info.num_output_elements, num_devices);
+  const int64_t warps_per_rank_tile =
+      CeilOfRatio(elements_per_rank_tile, elements_per_warp);
+  const int64_t num_warps = std::min<int64_t>(
+      kMaxWarpsPerBlock,
+      llvm::bit_ceil(static_cast<uint64_t>(warps_per_rank_tile)));
+  const int64_t blocks_per_rank_tile =
+      CeilOfRatio(warps_per_rank_tile, num_warps);
+  const int64_t num_blocks =
+      std::min(kReduceScatterMaxBlocksPerGrid,
+               num_devices * static_cast<int64_t>(llvm::bit_floor(
+                                 static_cast<uint64_t>(blocks_per_rank_tile))));
+
+  llvm::SmallVector<int64_t> tile_sizes =
+      GreedyPowerOfTwoTiles(output_shape, num_blocks);
+  // The pull-based reduce-scatter kernel requires the input tile of every
+  // program to cover exactly one shard along the scatter dimension, i.e. an
+  // output tile of `shard_size / num_devices` along the scatter dimension.
+  // FlattenReduceScatterFusion reshapes the fusion so that this is 1.
+  tile_sizes[scatter_dim] = output_shape.dimensions(scatter_dim) / num_devices;
+
+  return CreateBlockLevelFusionConfig(num_warps, tile_sizes);
+}
+
+}  // namespace
+
 absl::StatusOr<BlockLevelFusionConfig> GetCollectiveBlockLevelFusionConfig(
     const GpuTopology& gpu_topology, const HloFusionInstruction* fusion_instr,
     const DeviceAssignment* device_assignment) {
@@ -1364,61 +1367,36 @@ absl::StatusOr<BlockLevelFusionConfig> GetCollectiveBlockLevelFusionConfig(
           gpu_config.collective_backend_config().kernel_strategy())) {
     return absl::InvalidArgumentError(absl::StrFormat(
         "Instruction was not annotated with Triton strategy but still called "
-        "for "
-        "Triton codegen. %s",
+        "for Triton codegen. %s",
         instr->ToString()));
   }
-  ABSL_ASSIGN_OR_RETURN(CollectiveKernelInfo collective_info,
-                   GetCollectiveInfo(instr, gpu_topology, device_assignment));
-  ABSL_ASSIGN_OR_RETURN(const LaunchDimensions launch_dims,
-                   GetLaunchDimensions(collective_info, gpu_topology));
-  const Shape& output_shape = instr->shape();
-  const se::DeviceDescription& device_info =
-      gpu_topology.gpu_target_config().device_description;
-  BlockLevelFusionConfig block_level_config;
-  block_level_config.set_num_warps(xla::CeilOfRatio(
-      static_cast<int64_t>(launch_dims.num_threads_per_block()),
-      WarpSize(device_info)));
-  block_level_config.set_num_ctas(1);    // No block-level clustering.
-  block_level_config.set_num_stages(1);  // No pipelining of loops.
-  xtile::Tile* output_tile = block_level_config.add_output_tiles();
-  llvm::SmallVector<int64_t> tile_sizes =
-      GreedyPowerOfTwoTiles(output_shape, launch_dims.num_blocks());
-  // For AllGather, the tile on the gather dimension must not exceed the
-  // per-rank size. Otherwise a single tile would span multiple replicas,
-  // but the tiling framework assigns a single replica_id per tile.
-  if (instr->opcode() == HloOpcode::kAllGather) {
-    const auto* all_gather = Cast<HloAllGatherInstruction>(instr);
-    const int64_t gather_dim = all_gather->all_gather_dimension();
-    const int64_t num_devices =
-        std::get<AllGatherInfo>(collective_info).num_devices;
-    const int64_t per_rank_size =
-        output_shape.dimensions(gather_dim) / num_devices;
-    tile_sizes[gather_dim] = std::min(
-        tile_sizes[gather_dim], static_cast<int64_t>(llvm::bit_floor(
-                                    static_cast<uint64_t>(per_rank_size))));
-  } else if (instr->opcode() == HloOpcode::kReduceScatter) {
-    // The pull-based reduce-scatter kernel requires the input tile of every
-    // program to cover exactly one shard along the scatter dimension, i.e. an
-    // output tile of `shard_size / num_devices` along the scatter dimension.
-    // FlattenReduceScatterFusion reshapes the fusion so that this is 1.
-    const auto* rs = Cast<HloReduceScatterInstruction>(instr);
-    const int64_t scatter_dim = rs->scatter_dimension();
-    const int64_t num_devices =
-        std::get<ReduceScatterInfo>(collective_info).num_devices;
-    TF_RET_CHECK(output_shape.dimensions(scatter_dim) % num_devices == 0)
-        << "Reduce-scatter output scatter dimension "
-        << output_shape.dimensions(scatter_dim)
-        << " is not divisible by the number of devices " << num_devices;
-    tile_sizes[scatter_dim] =
-        output_shape.dimensions(scatter_dim) / num_devices;
+  if (!gpu_topology.has_gpu_target_config()) {
+    return absl::InvalidArgumentError(
+        "GpuTopology must have a target config to build "
+        "BlockLevelFusionConfig.");
   }
-  output_tile->mutable_sizes()->Assign(tile_sizes.begin(), tile_sizes.end());
-  ABSL_RETURN_IF_ERROR(
-      ValidateBlockLevelFusionConfig(block_level_config, collective_info));
-  VLOG(3) << "Block level fusion config for " << instr->name() << ": "
-          << block_level_config;
-  return block_level_config;
+
+  if (const HloAllReduceInstruction* all_reduce =
+          DynCast<HloAllReduceInstruction>(instr)) {
+    return GetBlockLevelConfigForAllReduce(gpu_topology, all_reduce,
+                                           device_assignment);
+  }
+
+  if (const HloAllGatherInstruction* all_gather =
+          DynCast<HloAllGatherInstruction>(instr)) {
+    return GetBlockLevelConfigForAllGather(gpu_topology, all_gather,
+                                           device_assignment);
+  }
+
+  if (const HloReduceScatterInstruction* reduce_scatter =
+          DynCast<HloReduceScatterInstruction>(instr)) {
+    return GetBlockLevelConfigForReduceScatter(gpu_topology, reduce_scatter,
+                                               device_assignment);
+  }
+
+  return absl::UnimplementedError(
+      absl::StrFormat("Unsupported collective opcode for Triton codegen: %s",
+                      HloOpcodeString(instr->opcode())));
 }
 
 absl::Status TrySetGpuBackendConfigForCollective(
@@ -1849,7 +1827,8 @@ absl::Status FlattenReduceScatterFusion(
   // When `inner_elements` is not a power of two (e.g. 12288), a 2D
   // [R, OutputSize / R] shape rounds the inner tile size up to a power of two
   // in GreedyPowerOfTwoTiles, which can drop the grid below the target block
-  // count from ReduceScatterLaunchDimensions (e.g. 48 blocks instead of 64).
+  // count from GetBlockLevelConfigForReduceScatter (e.g. 48 blocks instead of
+  // 64).
   // Factor out up to `kReduceScatterMaxBlocksPerGrid / R` blocks into a middle
   // dimension [R, B, OutputSize / (R * B)] so GreedyPowerOfTwoTiles can assign
   // blocks along that dimension without rounding `inner_elements` as a single

From 936742253d64ba83cbe61bc950ab88da2e3643b0 Mon Sep 17 00:00:00 2001
From: Quentin Khan <qkhan@google.com>
Date: Thu, 15 Oct 2026 05:26:04 -0700
Subject: [PATCH 0865/1000] Use a `size_t` instead of an `int` in
 `LoadLastBuildStep`.

This allows us to build cache files that are over 2 GiB.

PiperOrigin-RevId: 1000067200
---
 tensorflow/lite/delegates/xnnpack/weight_cache.cc | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/lite/delegates/xnnpack/weight_cache.cc b/tensorflow/lite/delegates/xnnpack/weight_cache.cc
index 1806f3ccc73..9daecb0abd9 100644
--- a/tensorflow/lite/delegates/xnnpack/weight_cache.cc
+++ b/tensorflow/lite/delegates/xnnpack/weight_cache.cc
@@ -610,7 +610,7 @@ bool MMapWeightCacheProvider::LoadLastBuildStep() {
   // - or add a new mapping handle.
   {
     MMapHandle& last_mmap_handle = mmap_handles_.back();
-    const int last_mmap_size = last_mmap_handle.size();
+    const size_t last_mmap_size = last_mmap_handle.size();
     if (!last_mmap_handle.Resize(last_mmap_size +
                                  builder_.LastBuildStepSize())) {
       mmap_handles_.emplace_back();

From 5c8408ccd09632ecf57bab8cd7dcac9c4dcdd12e Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 15 Oct 2026 05:53:47 -0700
Subject: [PATCH 0866/1000] [XLA] Fix while-loop buffer double counting in
 NoFragmentationStatsHeap.

Previously, NoFragmentationStatsHeap::ShareWith only delegated to Alloc, which caused shared buffers (such as while-loop init/parameter/result/root tuples) and reused operand buffers to be double counted in current_heap_size_.

Track shared buffers via a disjoint-set union with path compression (FindRoot) and reference counting in NoFragmentationStatsHeap so that memory is counted once per shared buffer group and deallocated only when the last reference is freed.

PiperOrigin-RevId: 1000078634
---
 .../service/heap_simulator/heap_simulator.cc  | 50 ++++++++++++--
 .../service/heap_simulator/heap_simulator.h   | 11 +++
 .../heap_simulator/heap_simulator_test.cc     | 68 +++++++++++++++++++
 3 files changed, 125 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
index 89fe04a9dde..94b73067687 100644
--- a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
+++ b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
@@ -640,9 +640,13 @@ void HeapSimulator::FillDebugTrace(HeapSimulatorTrace::Event::Kind kind,
 template <typename BufferType>
 void NoFragmentationStatsHeap<BufferType>::Alloc(const BufferType* buffer,
                                                  int64_t size) {
-  current_heap_size_ += size;
-  if (current_heap_size_ > max_heap_size_) {
-    max_heap_size_ = current_heap_size_;
+  BufferInfo& info = buffer_info_[buffer];
+  info.ref_count++;
+  if (info.ref_count == 1) {
+    current_heap_size_ += size;
+    if (current_heap_size_ > max_heap_size_) {
+      max_heap_size_ = current_heap_size_;
+    }
   }
 }
 
@@ -664,10 +668,48 @@ void NoFragmentationStatsHeap<BufferType>::AccountForSubcomputationMemory(
       std::max(max_heap_size_, current_heap_size_ + max_subcomputation_bytes);
 }
 
+template <typename BufferType>
+const BufferType* NoFragmentationStatsHeap<BufferType>::FindRoot(
+    const BufferType* buffer) {
+  const BufferType* root = buffer;
+  while (buffer_info_[root].underlying != nullptr) {
+    root = buffer_info_[root].underlying;
+  }
+  const BufferType* curr = buffer;
+  while (curr != root) {
+    const BufferType* next = buffer_info_[curr].underlying;
+    buffer_info_[curr].underlying = root;
+    curr = next;
+  }
+  return root;
+}
+
 template <typename BufferType>
 void NoFragmentationStatsHeap<BufferType>::Free(const BufferType* buffer,
                                                 int64_t size) {
-  current_heap_size_ -= size;
+  const BufferType* root = FindRoot(buffer);
+  BufferInfo& root_info = buffer_info_[root];
+  root_info.ref_count--;
+  if (root_info.ref_count == 0) {
+    current_heap_size_ -= size;
+  }
+}
+
+template <typename BufferType>
+void NoFragmentationStatsHeap<BufferType>::ShareWith(
+    const BufferType* buffer, const BufferType* share_with, int64_t size) {
+  const BufferType* root = FindRoot(share_with);
+  if (buffer != root) {
+    buffer_info_[buffer].underlying = root;
+  }
+  BufferInfo& root_info = buffer_info_[root];
+  root_info.ref_count++;
+  if (root_info.ref_count == 1) {
+    current_heap_size_ += size;
+    if (current_heap_size_ > max_heap_size_) {
+      max_heap_size_ = current_heap_size_;
+    }
+  }
 }
 
 template <typename BufferType>
diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator.h b/third_party/xla/xla/service/heap_simulator/heap_simulator.h
index edacc1ca4e7..949e4338f2b 100644
--- a/third_party/xla/xla/service/heap_simulator/heap_simulator.h
+++ b/third_party/xla/xla/service/heap_simulator/heap_simulator.h
@@ -341,9 +341,20 @@ class NoFragmentationStatsHeap : public HeapAlgorithm<BufferType> {
 
   void Free(const BufferType* buffer, int64_t size) override;
 
+  void ShareWith(const BufferType* buffer, const BufferType* share_with,
+                 int64_t size) override;
+
   absl::StatusOr<Result> Finish() override;
 
  private:
+  struct BufferInfo {
+    const BufferType* underlying = nullptr;
+    int64_t ref_count = 0;
+  };
+
+  const BufferType* FindRoot(const BufferType* buffer);
+
+  absl::flat_hash_map<const BufferType*, BufferInfo> buffer_info_;
   int64_t current_heap_size_ = 0;
   int64_t max_heap_size_ = 0;
 };
diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator_test.cc b/third_party/xla/xla/service/heap_simulator/heap_simulator_test.cc
index 1fc11b0d65e..cf3d0514428 100644
--- a/third_party/xla/xla/service/heap_simulator/heap_simulator_test.cc
+++ b/third_party/xla/xla/service/heap_simulator/heap_simulator_test.cc
@@ -245,6 +245,53 @@ TEST_F(MinimumMemoryForSequenceTest, SubcomputationAccounting) {
                     .value());
 }
 
+TEST_F(MinimumMemoryForSequenceTest, WhileLoopDoubleCounting) {
+  // Demonstrates that MinimumMemoryForModule (via NoFragmentationStatsHeap)
+  // does not double count shared buffers between the while loop condition,
+  // body, and init.
+  constexpr absl::string_view kHloString = R"hlo(
+    HloModule WhileLoopDoubleCounting, is_scheduled=true
+
+    WhileCond {
+      cond_param = f32[1000] parameter(0)
+      slice = f32[1] slice(cond_param), slice={[0:1]}
+      reshape = f32[] reshape(slice)
+      zero = f32[] constant(0)
+      ROOT cond_comparison = pred[] compare(reshape, zero), direction=NE
+    }
+
+    WhileBody {
+      body_param = f32[1000] parameter(0)
+      ROOT negate = f32[1000] negate(body_param)
+    }
+
+    ENTRY Entry {
+      init = f32[1000] parameter(0)
+      ROOT while_loop = f32[1000] while(init), condition=WhileCond, body=WhileBody
+    }
+  )hlo";
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kHloString));
+  ASSERT_OK(module->schedule().Verify());
+
+  BufferValue::SizeFunction size_fn = [](const BufferValue& buffer) {
+    return ShapeUtil::ByteSizeOf(buffer.shape());
+  };
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloAliasAnalysis> alias_analysis,
+                       HloAliasAnalysis::Run(module.get(), &alias_info_));
+
+  // With the NoFragmentationStatsHeap fix in place, the peak memory correctly
+  // avoids double counting the 4000-byte shared buffer across the condition and
+  // body, resulting in 4008 bytes (4000 + 4 slice + 4 reshape).
+  ASSERT_OK_AND_ASSIGN(
+      int64_t min_memory,
+      HeapSimulator::MinimumMemoryForModule(module->schedule(), *alias_analysis,
+                                            &alias_info_, &size_fn));
+  EXPECT_EQ(min_memory, 4008);
+}
+
 const char kAlloc[] = "Alloc";
 const char kFree[] = "Free";
 const char kShare[] = "Share";
@@ -1107,6 +1154,27 @@ TEST_F(NoFragmentationStatsHeapTest, Mixed) {
   EXPECT_EQ(40, result.heap_size);
 }
 
+TEST_F(NoFragmentationStatsHeapTest, ShareWith) {
+  NoFragmentationStatsHeap<HloValue> heap;
+  heap.Alloc(buffer_a_, 100);                 // current: 100, max: 100
+  heap.ShareWith(buffer_b_, buffer_a_, 100);  // current: 100, max: 100
+  heap.ShareWith(buffer_c_, buffer_b_, 100);  // current: 100, max: 100
+  heap.Free(buffer_b_, 100);                  // current: 100, max: 100
+  heap.ShareWith(buffer_d_, buffer_b_, 100);  // current: 100, max: 100
+  heap.Free(buffer_a_, 100);                  // current: 100, max: 100
+  heap.Free(buffer_c_, 100);                  // current: 100, max: 100
+  heap.Free(buffer_d_, 100);                  // current: 0, max: 100
+
+  heap.Alloc(buffer_e_, 50);  // current: 50, max: 100
+  heap.Alloc(buffer_f_, 60);  // current: 110, max: 110
+  heap.Free(buffer_e_, 50);
+  heap.Free(buffer_f_, 60);
+
+  ASSERT_OK_AND_ASSIGN(const HeapSimulator::Result<HloValue> result,
+                       heap.Finish());
+  EXPECT_EQ(110, result.heap_size);
+}
+
 class GlobalDecreasingSizeBestFitHeapTest : public HeapAlgorithmTestBase {};
 
 TEST_F(GlobalDecreasingSizeBestFitHeapTest, Empty) {

From 1075230b766be1cbc41afcf858a7fc599d0343f1 Mon Sep 17 00:00:00 2001
From: Benjamin Chetioui <bchetioui@google.com>
Date: Thu, 15 Oct 2026 05:54:50 -0700
Subject: [PATCH 0867/1000] [XLA:CPU] Restore piecewise math.expm1 expansion
 when C++ bitcode intrinsics are unavailable.

When `AreCppGenIntrinsicsAvailable()` returns `false` (such as on Windows, where `empty_bitcode` embeds an empty buffer), `xla::codegen::intrinsics::Expm1::SupportedVectorTypes()` is empty. Without a fallback rewrite in `ExpandFloatOpsPass`, `math.expm1` falls through to Microsoft UCRT `expm1f`, which underflows inputs with `|x| < 2^-24` to `0.0f` instead of `x` (causing `log(-expm1(-x))` to evaluate to `-inf` and zeroing small `complex64` components in `sin`, `tan`, `tanh`, and `expm1`).

Restore `RewriteExpm1Pattern` in `ExpandFloatOpsPass` when `!AreCppGenIntrinsicsAvailable()`:
- `|x| > 0.5`: `exp(x) - 1`
- `|x| <= 0.5`: `tanh(x / 2) * (exp(x) + 1)`

Platforms with C++ bitcode intrinsics continue to use `xla::codegen::intrinsics::Expm1`.

PiperOrigin-RevId: 1000079164
---
 .../cpu/codegen/emitters/transforms/BUILD     |  1 +
 .../emitters/transforms/expand_float_ops.cc   | 56 +++++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/BUILD b/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/BUILD
index 6026d1178f1..46a4ab9c6d2 100644
--- a/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/BUILD
+++ b/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/BUILD
@@ -59,6 +59,7 @@ cc_library(
         ":xla_cpu_rewrite_patterns",
         "//xla/backends/cpu/codegen/emitters/ir:xla_cpu",
         "//xla/codegen/emitters/ir:xla",
+        "//xla/codegen/intrinsic/cpp:_cpp_gen_intrinsics",
         "//xla/hlo/analysis:indexing_analysis",
         "//xla/hlo/analysis:symbolic_map",
         "@com_google_absl//absl/container:inlined_vector",
diff --git a/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/expand_float_ops.cc b/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/expand_float_ops.cc
index 0f135558882..eec2f97b493 100644
--- a/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/expand_float_ops.cc
+++ b/third_party/xla/xla/backends/cpu/codegen/emitters/transforms/expand_float_ops.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include <utility>
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/IR/ArithAttributes.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Math/IR/Math.h"  // IWYU pragma: keep
 #include "mlir/IR/Builders.h"
@@ -35,6 +36,7 @@ limitations under the License.
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "xla/codegen/intrinsic/cpp/cpp_gen_intrinsics.h"
 
 namespace xla::cpu {
 
@@ -100,6 +102,57 @@ struct RewriteExtFPattern : public mlir::OpRewritePattern<ma::ExtFOp> {
   }
 };
 
+// Fallback lowering for expm1(x) on platforms (such as Windows) where CppGen
+// bitcode intrinsics are unavailable:
+// |x| > 0.5: exp(x) - 1
+// |x| <= 0.5: tanh(x / 2) * (exp(x) + 1)
+class RewriteExpm1Pattern : public mlir::OpRewritePattern<mlir::math::ExpM1Op> {
+ public:
+  using OpRewritePattern::OpRewritePattern;
+
+  mlir::LogicalResult matchAndRewrite(
+      mlir::math::ExpM1Op op, mlir::PatternRewriter& rewriter) const override {
+    mlir::ImplicitLocOpBuilder b(op.getLoc(), rewriter);
+
+    mlir::Type type = op.getType();
+    mlir::Type element_type = mlir::getElementTypeOrSelf(type);
+    mlir::Value one = GetConst(b, type, b.getFloatAttr(element_type, 1.0));
+    mlir::Value half = GetConst(b, type, b.getFloatAttr(element_type, 0.5));
+    mlir::Value zero = GetConst(b, type, b.getFloatAttr(element_type, 0.0));
+    mlir::Value x = op.getOperand();
+
+    mlir::arith::FastMathFlagsAttr fastmath = op.getFastmathAttr();
+
+    mlir::Value exp_x = mlir::math::ExpOp::create(b, x, fastmath);
+
+    mlir::Value exp_x_minus_1 =
+        mlir::arith::SubFOp::create(b, exp_x, one, fastmath);
+
+    mlir::Value half_x = mlir::arith::MulFOp::create(b, x, half, fastmath);
+    mlir::Value tanh_half_x = mlir::math::TanhOp::create(b, half_x, fastmath);
+    mlir::Value exp_x_plus_1 =
+        mlir::arith::AddFOp::create(b, exp_x, one, fastmath);
+    mlir::Value small_result =
+        mlir::arith::MulFOp::create(b, tanh_half_x, exp_x_plus_1, fastmath);
+
+    mlir::Value abs_x = mlir::math::AbsFOp::create(b, x, fastmath);
+    mlir::Value x_is_large = mlir::arith::CmpFOp::create(
+        b, mlir::arith::CmpFPredicate::OGT, abs_x, half);
+    mlir::Value normal_result = mlir::arith::SelectOp::create(
+        b, x_is_large, exp_x_minus_1, small_result);
+
+    // half_x can underflow resulting in zero.
+    // TODO(willfroom): Do we actually need this check? tanh(0) == 0.
+    mlir::Value half_x_is_zero = mlir::arith::CmpFOp::create(
+        b, mlir::arith::CmpFPredicate::OEQ, half_x, zero);
+    mlir::Value result =
+        mlir::arith::SelectOp::create(b, half_x_is_zero, x, normal_result);
+
+    rewriter.replaceOp(op, result);
+    return mlir::success();
+  }
+};
+
 class ExpandFloatOpsPass
     : public impl::ExpandFloatOpsPassBase<ExpandFloatOpsPass> {
  public:
@@ -108,6 +161,9 @@ class ExpandFloatOpsPass
   void runOnOperation() override {
     mlir::RewritePatternSet patterns(&getContext());
     patterns.add<RewriteExtFPattern>(&getContext());
+    if (!xla::codegen::AreCppGenIntrinsicsAvailable()) {
+      patterns.add<RewriteExpm1Pattern>(&getContext());
+    }
 
     if (mlir::failed(
             mlir::applyPatternsGreedily(getOperation(), std::move(patterns)))) {

From 5e450b6e6653e891f359c40a318f6602fa9c4bfe Mon Sep 17 00:00:00 2001
From: Mohammed Anany <manany@google.com>
Date: Thu, 15 Oct 2026 05:55:09 -0700
Subject: [PATCH 0868/1000] [XLA:GPU] Remove
 xla_gpu_unsupported_enable_triton_multi_output_fusion check from TilingSpace.
 Also disable corresponding tests.

This should only be added back when it is properly supported.

PiperOrigin-RevId: 1000079303
---
 .../tests/bitcast/bitcast_reduce_stride1.hlo  |  2 +-
 .../tests/fusion_emitter_device_test.cc       |  6 ++-
 .../tests/reduce/reduce_minor_extra_out.hlo   |  2 +-
 .../reduce/reduce_to_scalar_extra_out.hlo     |  2 +-
 .../tests/slice/slice_extra_out_padding.hlo   |  2 +-
 .../triton/tests/slice/slice_padding.hlo      |  2 +-
 .../xla/xla/codegen/tiling/experimental/BUILD |  1 -
 .../tiling/experimental/tiled_hlo_test.cc     | 44 ++++++-------------
 .../tiling/experimental/tiling_space.cc       |  5 ---
 .../tiling/experimental/tiling_space_test.cc  |  2 +
 .../model/triton_emitter_constraints_test.cc  |  2 +
 11 files changed, 26 insertions(+), 44 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reduce_stride1.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reduce_stride1.hlo
index cb96399ce98..9d696d29fb8 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reduce_stride1.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/bitcast/bitcast_reduce_stride1.hlo
@@ -1,5 +1,5 @@
 // RUN: fusion_to_triton --xla_gpu_unsupported_enable_triton_multi_output_fusion %s 2>&1 | FileCheck %s --check-prefixes=CHECK-ERROR
-// CHECK-ERROR: Only single-result fusions are supported for now. Received 2 roots.
+// CHECK-ERROR: TilingSpace does not support fusions with multiple roots.
 // TODO: b/502910372 - support multi-output fusion.
 // RUN-NOT: fusion_to_triton --xla_gpu_unsupported_enable_triton_multi_output_fusion %s | FileCheck %s
 // RUN-NOT: triton_test_correctness --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc
index 29a1327e315..0b6e036e54b 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/fusion_emitter_device_test.cc
@@ -218,8 +218,10 @@ ENTRY entry_computation {
       CreateTritonIrFromHloTextAndFileCheck(kHloText, "fused_computation", "");
   EXPECT_THAT(
       status,
-      StatusIs(tsl::error::UNIMPLEMENTED,
-               HasSubstr("Only single-result fusions are supported for now")));
+      StatusIs(
+          tsl::error::UNIMPLEMENTED,
+          HasSubstr(
+              "TilingSpace does not support fusions with multiple roots.")));
 }
 
 class TritonEmitterTestWithOffsetParam
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_minor_extra_out.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_minor_extra_out.hlo
index 1f419226222..756faef45e1 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_minor_extra_out.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_minor_extra_out.hlo
@@ -1,5 +1,5 @@
 // RUN: fusion_to_triton --xla_gpu_unsupported_enable_triton_multi_output_fusion %s 2>&1 | FileCheck %s --check-prefixes=CHECK-ERROR
-// CHECK-ERROR: Only single-result fusions are supported for now. Received 2 roots.
+// CHECK-ERROR: TilingSpace does not support fusions with multiple roots.
 // TODO: b/502910372 - support multi-output fusion.
 // RUN-NOT: fusion_to_triton --xla_gpu_unsupported_enable_triton_multi_output_fusion %s | FileCheck %s
 // RUN-NOT: triton_test_correctness --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_to_scalar_extra_out.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_to_scalar_extra_out.hlo
index 45bacf8c748..19a4b246946 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_to_scalar_extra_out.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_to_scalar_extra_out.hlo
@@ -1,5 +1,5 @@
 // RUN: fusion_to_triton --xla_gpu_unsupported_enable_triton_multi_output_fusion %s 2>&1 | FileCheck %s --check-prefixes=CHECK-ERROR
-// CHECK-ERROR: Only single-result fusions are supported for now. Received 2 roots.
+// CHECK-ERROR: TilingSpace does not support fusions with multiple roots.
 // TODO: b/502910372 - enable when multi-output is supported.
 // RUN-NOT: fusion_to_triton --xla_gpu_unsupported_enable_triton_multi_output_fusion %s | FileCheck %s
 // RUN-NOT: triton_test_correctness --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/slice_extra_out_padding.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/slice_extra_out_padding.hlo
index 2309066fff4..fd1de9fc3a7 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/slice_extra_out_padding.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/slice_extra_out_padding.hlo
@@ -1,5 +1,5 @@
 // RUN: fusion_to_triton --xla_gpu_unsupported_enable_triton_multi_output_fusion %s 2>&1 | FileCheck %s --check-prefixes=CHECK-ERROR
-// CHECK-ERROR: Only single-result fusions are supported for now. Received 2 roots.
+// CHECK-ERROR: TilingSpace does not support fusions with multiple roots.
 // TODO: b/502910372 - support multi-output fusion.
 // RUN-NOT: triton_test_correctness --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/slice_padding.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/slice_padding.hlo
index dc239f63d5f..8ac0244c018 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/slice_padding.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/slice/slice_padding.hlo
@@ -1,5 +1,5 @@
 // RUN: fusion_to_triton --xla_gpu_unsupported_enable_triton_multi_output_fusion %s 2>&1 | FileCheck %s --check-prefixes=CHECK-ERROR
-// CHECK-ERROR: Only single-result fusions are supported for now. Received 2 roots.
+// CHECK-ERROR: TilingSpace does not support fusions with multiple roots.
 // TODO: b/502910372 - support multi-output fusion.
 // RUN-NOT: triton_test_correctness --xla_gpu_unsupported_enable_triton_multi_output_fusion %s
 
diff --git a/third_party/xla/xla/codegen/tiling/experimental/BUILD b/third_party/xla/xla/codegen/tiling/experimental/BUILD
index b7ad1878827..8d46ed00c18 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/BUILD
+++ b/third_party/xla/xla/codegen/tiling/experimental/BUILD
@@ -318,7 +318,6 @@ xla_cc_test(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:status_matchers",
-        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest_main",
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc
index b6dff55cfc0..047648b9830 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiled_hlo_test.cc
@@ -18,7 +18,6 @@ limitations under the License.
 #include <cstdint>
 #include <memory>
 #include <string>
-#include <tuple>
 #include <utility>
 #include <vector>
 
@@ -29,7 +28,6 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/status_matchers.h"
-#include "absl/strings/str_join.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "llvm/ADT/STLExtras.h"
@@ -245,41 +243,27 @@ class TileAnalysisTestBase : public HloHardwareIndependentTestBase {
 };
 
 // For tests that don't depend on the flags in TileAnalysisTestParams.
-class TileAnalysisTest
-    : public TileAnalysisTestBase,
-      public ::testing::WithParamInterface<std::tuple<bool, bool>> {
+class TileAnalysisTest : public TileAnalysisTestBase,
+                         public ::testing::WithParamInterface<bool> {
   DebugOptions GetDebugOptionsForTest() const override {
     DebugOptions debug_options = TileAnalysisTestBase::GetDebugOptionsForTest();
-    debug_options.set_xla_gpu_unsupported_enable_triton_multi_output_fusion(
-        EnableTritonMultiOutputFusion());
     debug_options
         .set_xla_gpu_experimental_enable_same_shape_multi_output_fusion(
             EnableSameShapeMultiOutputFusion());
     return debug_options;
   }
 
-  bool EnableTritonMultiOutputFusion() const { return std::get<0>(GetParam()); }
-  bool EnableSameShapeMultiOutputFusion() const {
-    return std::get<1>(GetParam());
-  }
+  bool EnableSameShapeMultiOutputFusion() const { return GetParam(); }
 };
 
-INSTANTIATE_TEST_SUITE_P(
-    TileAnalysisTestInstantiation, TileAnalysisTest,
-    testing::Combine(testing::Bool(), testing::Bool()),
-    [](const ::testing::TestParamInfo<std::tuple<bool, bool>>& info) {
-      std::vector<std::string> parts;
-      if (std::get<0>(info.param)) {
-        parts.push_back("WithTritonMOF");
-      }
-      if (std::get<1>(info.param)) {
-        parts.push_back("WithSameShapeMOF");
-      }
-      if (parts.empty()) {
-        return std::string("Default");
-      }
-      return absl::StrJoin(parts, "_");
-    });
+INSTANTIATE_TEST_SUITE_P(TileAnalysisTestInstantiation, TileAnalysisTest,
+                         testing::Bool(),
+                         [](const ::testing::TestParamInfo<bool>& info) {
+                           if (info.param) {
+                             return std::string("WithSameShapeMOF");
+                           }
+                           return std::string("Default");
+                         });
 
 // For tests that assume no multi-output fusion support.
 class SingleOutputFusionTileAnalysisTest : public TileAnalysisTestBase {};
@@ -302,10 +286,6 @@ class SameShapeMultiOutputFusionTileAnalysisTest : public TileAnalysisTestBase {
     DebugOptions debug_options = TileAnalysisTestBase::GetDebugOptionsForTest();
     debug_options
         .set_xla_gpu_experimental_enable_same_shape_multi_output_fusion(true);
-    // Triton multi-output fusion is preferred over same-shape MOF so we must
-    // disable it to test same-shape MOF.
-    debug_options.set_xla_gpu_unsupported_enable_triton_multi_output_fusion(
-        false);
     return debug_options;
   }
 };
@@ -961,6 +941,8 @@ TEST_F(SingleOutputFusionTileAnalysisTest, SimpleFusion) {
 }
 
 TEST_F(TritonMultiOutputFusionTileAnalysisTest, DuplicateFusionRoots) {
+  GTEST_SKIP() << "Arbitrary multi-output fusions are currently not supported "
+                  "in TilingSpace (b/502910372).";
   ASSERT_OK_AND_ASSIGN(const TiledHloComputation tiled_computation,
                        ParseAndTile(R"hlo(
     fusion {
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc
index d92f51eda38..a50825475fa 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc
@@ -612,11 +612,6 @@ absl::StatusOr<std::unique_ptr<TilingSpace>> TilingSpace::Create(
           .xla_gpu_experimental_enable_same_shape_multi_output_fusion()) {
     ABSL_RETURN_IF_ERROR(
         tiling_space->InitializeDimensionsForSameShapeMultiOutputFusion(roots));
-  } else if (debug_options
-                 .xla_gpu_unsupported_enable_triton_multi_output_fusion()) {
-    // xla_gpu_unsupported_enable_triton_multi_output_fusion flag name is
-    // misleading, it is not GPU specific.
-    ABSL_RETURN_IF_ERROR(tiling_space->InitializeDimensions(roots));
   } else {
     // TODO(b/502910372): Support arbitrary multi-output fusions.
     return absl::UnimplementedError(
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc
index abc574bfe36..e416022793a 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc
@@ -296,6 +296,8 @@ TEST_F(TilingSpaceTest, DynamicSlice) {
 }
 
 TEST_F(TilingSpaceTest, TwoOutputsParallelDims) {
+  GTEST_SKIP() << "Arbitrary multi-output fusions are currently not supported "
+                  "in TilingSpace (b/502910372).";
   HloInstruction* root = ParseAndGetRoot(R"(
     HloModule m
     f {
diff --git a/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc b/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc
index 96d32c61608..82d3aa460d4 100644
--- a/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc
+++ b/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc
@@ -613,6 +613,8 @@ ENTRY entry_computation {
 }
 
 TEST_F(VerifyTritonConstraintsTest, MultiOutputFusionHasPowerOfTwoTileSizes) {
+  GTEST_SKIP() << "Arbitrary multi-output fusions are currently not supported "
+                  "in TilingSpace (b/502910372).";
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"hlo(
 HloModule m

From e9b22e2fa9115d3fd3d27c837127daa8c388d369 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Thu, 15 Oct 2026 06:57:35 -0700
Subject: [PATCH 0869/1000] [XLA:CPU] Improve float32 rsqrt precision to 1 ULP
 on CPU

- Replace unfused `CreateFAdd(CreateFMul(x, y), z)` in `PMAdd` with `llvm::Intrinsic::fmuladd` in `xla::codegen::intrinsics::Rsqrt` so Newton-Raphson iterations use FMA when available.
- Always perform 2 Newton-Raphson iterations across both Intel and AMD CPUs rather than 1 iteration on AMD non-AVX512 targets.
- Tighten the CPU `rsqrt` `float32` precision bound to 1 ULP in JAX's `power_test.py` and XLA's `exhaustive_unary_test_functions.cc`.

PiperOrigin-RevId: 1000107347
---
 third_party/xla/xla/codegen/intrinsic/rsqrt.cc     | 13 ++++---------
 .../exhaustive/exhaustive_unary_test_functions.cc  | 14 ++++++++++++++
 .../tests/exhaustive/exhaustive_unary_test_ops.inc |  7 ++++++-
 3 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/third_party/xla/xla/codegen/intrinsic/rsqrt.cc b/third_party/xla/xla/codegen/intrinsic/rsqrt.cc
index 25d1975dd26..72492bc64a2 100644
--- a/third_party/xla/xla/codegen/intrinsic/rsqrt.cc
+++ b/third_party/xla/xla/codegen/intrinsic/rsqrt.cc
@@ -46,7 +46,8 @@ namespace xla::codegen::intrinsics {
 namespace {
 llvm::Value* PMAdd(llvm::IRBuilder<>& builder, llvm::Value* x, llvm::Value* y,
                    llvm::Value* z) {
-  return builder.CreateFAdd(builder.CreateFMul(x, y), z);
+  return builder.CreateIntrinsic(llvm::Intrinsic::fmuladd, {x->getType()},
+                                 {x, y, z});
 }
 }  // namespace
 
@@ -194,14 +195,8 @@ absl::StatusOr<llvm::Function*> Rsqrt::CreateDefinition(
   llvm::Value* y_approx = rsqrt_intrinsic.CreateCall(builder, x);
 
   // Eigen only does 1 step for F32, but that only gives within 2 ULPs and we
-  // are targeting 1. AMD's SSE/AVX rsqrt intrinsics are more accurate; their
-  // avx512f intrinsics have the same accuracy as Intel's avx512f intrinsics.
-  const bool using_avx512 =
-      options.Contains("+avx512f") &&
-      (type.element_type() == F64 ||
-       (type.element_type() == F32 && type.vector_width().value_or(1) > 8));
-  const bool is_amd = options.device_type == DeviceType::kAmdCpu;
-  const size_t num_steps = (is_amd && !using_avx512) ? 1 : 2;
+  // are targeting 1.
+  const size_t num_steps = 2;
   llvm::Value* refined_result =
       NewtonRaphsonRsqrtIteration(builder, x, y_approx, input_type, num_steps);
 
diff --git a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_functions.cc b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_functions.cc
index 9f8eef54404..ae6fb272665 100644
--- a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_functions.cc
+++ b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_functions.cc
@@ -130,6 +130,20 @@ UNARY_TEST(Rsqrt, {
             .strict_signed_zeros()
             .build();
       })
+      .CpuError(+[](NativeT x) {
+        if constexpr (std::is_same_v<NativeT, double>) {
+          // The reference 1.0 / std::sqrt(x) computed in double can itself be
+          // 1 ULP off from the exact value in the opposite direction.
+          return ErrorSpec::Builder()
+              .distance_err(2)
+              .strict_signed_zeros()
+              .build();
+        }
+        return ErrorSpec::Builder()
+            .distance_err(1)
+            .strict_signed_zeros()
+            .build();
+      })
       .Run();
 })
 UNARY_TEST(Sqrt, {
diff --git a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_ops.inc b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_ops.inc
index 89591d79ca2..050f69537ed 100644
--- a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_ops.inc
+++ b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_ops.inc
@@ -208,7 +208,12 @@ DEFINE_UNARY_TEST_OP(
     { return +[](Traits::NativeRefT x) { return std::exp2(x); }; });
 DEFINE_UNARY_TEST_OP(
     RsqrtOp, { return [](XlaOp x) { return Rsqrt(x); }; },
-    { return +[](Traits::NativeRefT x) { return 1 / std::sqrt(x); }; });
+    {
+      return +[](typename Traits::NativeRefT x) {
+        return static_cast<typename Traits::NativeRefT>(
+            1.0 / std::sqrt(static_cast<double>(x)));
+      };
+    });
 DEFINE_UNARY_TEST_OP(
     SqrtOp, { return [](XlaOp x) { return Sqrt(x); }; }, { return std::sqrt; });
 DEFINE_UNARY_TEST_OP(

From eb03ef04f8ee75887efe2c6c881549ba7725c82c Mon Sep 17 00:00:00 2001
From: Jake VanderPlas <vanderplas@google.com>
Date: Thu, 15 Oct 2026 06:59:20 -0700
Subject: [PATCH 0870/1000] [CHLO] Avoid underflow in chlo.zeta Euler-Maclaurin
 correction term

The Euler-Maclaurin integral term of the chlo.zeta decomposition was
computed as (acc^-x * acc) / (x - 1), where acc = q + 10. For large acc,
acc^-x underflows to zero (or is flushed to zero) even though acc^(1-x),
the dominant term of the result, is representable, so zeta(x, q) silently
returned 0. For q = +inf the same product is 0 * inf = NaN, although
zeta(x, inf) = 0 for x > 1.

Compute acc^(1-x) directly instead (same identity as the unmerged
https://github.com/openxla/stablehlo/pull/3030). This is algebraically
identical, never underflows earlier than the old product, and evaluates to 0
at acc = inf, so no special-casing of q = inf is needed. Since
polygamma(n, x) is lowered via zeta(n + 1, x), this also fixes
polygamma(n >= 1, inf) returning NaN.

This fixes the following JAX issues (jax.scipy.special.zeta and polygamma
lower to chlo.zeta / chlo.polygamma):
- https://github.com/jax-ml/jax/issues/41366
- https://github.com/jax-ml/jax/issues/38637
- https://github.com/jax-ml/jax/issues/38608

PiperOrigin-RevId: 1000108122
---
 .../xla/third_party/stablehlo/temporary.patch | 82 +++++++++++++++++--
 third_party/xla/xla/hlo/builder/lib/math.cc   |  5 +-
 .../xla/xla/hlo/builder/lib/math_test.cc      | 24 ++++++
 .../xla/hlo/builder/tests/math_polygamma.hlo  | 10 ++-
 .../Dialect/chlo/chlo_legalize_to_mhlo.mlir   | 15 ++--
 5 files changed, 118 insertions(+), 18 deletions(-)

diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch
index d2710427181..ccd5f019ae8 100644
--- a/third_party/xla/third_party/stablehlo/temporary.patch
+++ b/third_party/xla/third_party/stablehlo/temporary.patch
@@ -4509,7 +4509,59 @@ diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/st
  // CHECK:         }
  func.func @asin_f64(%arg : tensor<f64>) -> tensor<f64> {
    %result = "chlo.asin"(%arg) : (tensor<f64>) -> tensor<f64>
-@@ -5424,6 +5416,55 @@
+@@ -2432,9 +2424,10 @@
+ // CHECK:           %[[ADD_18:.*]] = stablehlo.add %[[ADD_16]], %[[CONSTANT_1]] : tensor<f32>
+ // CHECK:           %[[POWER_10:.*]] = stablehlo.power %[[ADD_18]], %[[NEGATE_0]] : tensor<f32>
+ // CHECK:           %[[CONSTANT_2:.*]] = stablehlo.constant dense<1.000000e+00> : tensor<f32>
+-// CHECK:           %[[MULTIPLY_0:.*]] = stablehlo.multiply %[[POWER_10]], %[[ADD_18]] : tensor<f32>
+ // CHECK:           %[[SUBTRACT_0:.*]] = stablehlo.subtract %[[CONVERT_0]], %[[CONSTANT_2]] : tensor<f32>
+-// CHECK:           %[[DIVIDE_0:.*]] = stablehlo.divide %[[MULTIPLY_0]], %[[SUBTRACT_0]] : tensor<f32>
++// CHECK:           %[[NEGATE_1:.*]] = stablehlo.negate %[[SUBTRACT_0]] : tensor<f32>
++// CHECK:           %[[POWER_11:.*]] = stablehlo.power %[[ADD_18]], %[[NEGATE_1]] : tensor<f32>
++// CHECK:           %[[DIVIDE_0:.*]] = stablehlo.divide %[[POWER_11]], %[[SUBTRACT_0]] : tensor<f32>
+ // CHECK:           %[[MULTIPLY_1:.*]] = stablehlo.multiply %[[ADD_18]], %[[ADD_18]] : tensor<f32>
+ // CHECK:           %[[DIVIDE_1:.*]] = stablehlo.divide %[[CONSTANT_1]], %[[MULTIPLY_1]] : tensor<f32>
+ // CHECK:           %[[CONSTANT_3:.*]] = stablehlo.constant dense<2.200000e+01> : tensor<f32>
+@@ -2710,9 +2703,10 @@
+ // CHECK:           %[[ADD_40:.*]] = stablehlo.add %[[ADD_38]], %[[CONSTANT_29]] : tensor<f32>
+ // CHECK:           %[[POWER_10:.*]] = stablehlo.power %[[ADD_40]], %[[NEGATE_2]] : tensor<f32>
+ // CHECK:           %[[CONSTANT_30:.*]] = stablehlo.constant dense<1.000000e+00> : tensor<f32>
+-// CHECK:           %[[MULTIPLY_3:.*]] = stablehlo.multiply %[[POWER_10]], %[[ADD_40]] : tensor<f32>
+ // CHECK:           %[[SUBTRACT_7:.*]] = stablehlo.subtract %[[ADD_0]], %[[CONSTANT_30]] : tensor<f32>
+-// CHECK:           %[[DIVIDE_10:.*]] = stablehlo.divide %[[MULTIPLY_3]], %[[SUBTRACT_7]] : tensor<f32>
++// CHECK:           %[[NEGATE_4:.*]] = stablehlo.negate %[[SUBTRACT_7]] : tensor<f32>
++// CHECK:           %[[POWER_11:.*]] = stablehlo.power %[[ADD_40]], %[[NEGATE_4]] : tensor<f32>
++// CHECK:           %[[DIVIDE_10:.*]] = stablehlo.divide %[[POWER_11]], %[[SUBTRACT_7]] : tensor<f32>
+ // CHECK:           %[[MULTIPLY_4:.*]] = stablehlo.multiply %[[ADD_40]], %[[ADD_40]] : tensor<f32>
+ // CHECK:           %[[DIVIDE_11:.*]] = stablehlo.divide %[[CONSTANT_29]], %[[MULTIPLY_4]] : tensor<f32>
+ // CHECK:           %[[CONSTANT_31:.*]] = stablehlo.constant dense<2.200000e+01> : tensor<f32>
+@@ -3103,9 +3097,10 @@
+ // CHECK:           %[[ADD_40:.*]] = stablehlo.add %[[ADD_38]], %[[CONSTANT_29]] : tensor<f64>
+ // CHECK:           %[[POWER_10:.*]] = stablehlo.power %[[ADD_40]], %[[NEGATE_2]] : tensor<f64>
+ // CHECK:           %[[CONSTANT_30:.*]] = stablehlo.constant dense<1.000000e+00> : tensor<f64>
+-// CHECK:           %[[MULTIPLY_3:.*]] = stablehlo.multiply %[[POWER_10]], %[[ADD_40]] : tensor<f64>
+ // CHECK:           %[[SUBTRACT_7:.*]] = stablehlo.subtract %[[ADD_0]], %[[CONSTANT_30]] : tensor<f64>
+-// CHECK:           %[[DIVIDE_10:.*]] = stablehlo.divide %[[MULTIPLY_3]], %[[SUBTRACT_7]] : tensor<f64>
++// CHECK:           %[[NEGATE_4:.*]] = stablehlo.negate %[[SUBTRACT_7]] : tensor<f64>
++// CHECK:           %[[POWER_11:.*]] = stablehlo.power %[[ADD_40]], %[[NEGATE_4]] : tensor<f64>
++// CHECK:           %[[DIVIDE_10:.*]] = stablehlo.divide %[[POWER_11]], %[[SUBTRACT_7]] : tensor<f64>
+ // CHECK:           %[[MULTIPLY_4:.*]] = stablehlo.multiply %[[ADD_40]], %[[ADD_40]] : tensor<f64>
+ // CHECK:           %[[DIVIDE_11:.*]] = stablehlo.divide %[[CONSTANT_29]], %[[MULTIPLY_4]] : tensor<f64>
+ // CHECK:           %[[CONSTANT_31:.*]] = stablehlo.constant dense<2.200000e+01> : tensor<f64>
+@@ -3498,9 +3493,10 @@
+ // CHECK:           %[[ADD_40:.*]] = stablehlo.add %[[ADD_38]], %[[CONSTANT_29]] : tensor<f32>
+ // CHECK:           %[[POWER_10:.*]] = stablehlo.power %[[ADD_40]], %[[NEGATE_2]] : tensor<f32>
+ // CHECK:           %[[CONSTANT_30:.*]] = stablehlo.constant dense<1.000000e+00> : tensor<f32>
+-// CHECK:           %[[MULTIPLY_3:.*]] = stablehlo.multiply %[[POWER_10]], %[[ADD_40]] : tensor<f32>
+ // CHECK:           %[[SUBTRACT_7:.*]] = stablehlo.subtract %[[ADD_0]], %[[CONSTANT_30]] : tensor<f32>
+-// CHECK:           %[[DIVIDE_10:.*]] = stablehlo.divide %[[MULTIPLY_3]], %[[SUBTRACT_7]] : tensor<f32>
++// CHECK:           %[[NEGATE_4:.*]] = stablehlo.negate %[[SUBTRACT_7]] : tensor<f32>
++// CHECK:           %[[POWER_11:.*]] = stablehlo.power %[[ADD_40]], %[[NEGATE_4]] : tensor<f32>
++// CHECK:           %[[DIVIDE_10:.*]] = stablehlo.divide %[[POWER_11]], %[[SUBTRACT_7]] : tensor<f32>
+ // CHECK:           %[[MULTIPLY_4:.*]] = stablehlo.multiply %[[ADD_40]], %[[ADD_40]] : tensor<f32>
+ // CHECK:           %[[DIVIDE_11:.*]] = stablehlo.divide %[[CONSTANT_29]], %[[MULTIPLY_4]] : tensor<f32>
+ // CHECK:           %[[CONSTANT_31:.*]] = stablehlo.constant dense<2.200000e+01> : tensor<f32>
+@@ -5424,6 +5420,55 @@
  
  // -----
  
@@ -12601,7 +12653,23 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stable
    }
    return {};
  }
-@@ -2520,7 +2522,7 @@
+@@ -1747,9 +1749,13 @@
+   acc = AddOp::create(rewriter, loc, acc, one);
+   qNegPower = PowOp::create(rewriter, loc, acc, negX);
+   Value oneLikeX = getConstantLike(rewriter, loc, 1.0, x);
++  // Compute acc^(1-x) directly rather than as acc^-x * acc: the latter
++  // underflows to 0 for large acc, and is inf * 0 = NaN for acc = inf.
++  Value xMinusOne = SubtractOp::create(rewriter, loc, x, oneLikeX);
++  Value accPowOneMinusX = PowOp::create(
++      rewriter, loc, acc, NegOp::create(rewriter, loc, xMinusOne));
+   Value correctionEulerMaclaurin =
+-      DivOp::create(rewriter, loc, MulOp::create(rewriter, loc, qNegPower, acc),
+-                    SubtractOp::create(rewriter, loc, x, oneLikeX));
++      DivOp::create(rewriter, loc, accPowOneMinusX, xMinusOne);
+ 
+   // Manual reciprocal of the square root as RsqrtOp produces different results
+   Value rsqrtAcc =
+@@ -2520,7 +2526,7 @@
      }
      Value dotGeneral = mlir::stablehlo::DotGeneralOp::create(
          rewriter, op.getLoc(), resultType, ValueRange{lhsMasked, rhsReshaped},
@@ -12610,7 +12678,7 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stable
  
      if (i == 0) {
        out = dotGeneral;
-@@ -2672,7 +2674,7 @@
+@@ -2672,7 +2678,7 @@
      }
      Value dotGeneral = mlir::stablehlo::DotGeneralOp::create(
          rewriter, op.getLoc(), RankedTensorType::get(resultShape, outDType),
@@ -12619,7 +12687,7 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stable
  
      results.push_back(dotGeneral);
      start = limit;
-@@ -2747,7 +2749,7 @@
+@@ -2747,7 +2753,7 @@
    }
    rewriter.replaceOpWithNewOp<mlir::stablehlo::DotGeneralOp>(
        op, op.getResult().getType(), ValueRange{op.getLhs(), op.getRhs()},
@@ -12628,7 +12696,7 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stable
    return success();
  }
  
-@@ -2930,8 +2932,18 @@
+@@ -2930,8 +2936,18 @@
        SmallVector<int64_t> sliceShape(inputType.getShape().begin(),
                                        inputType.getShape().end());
        sliceShape[dim] = 1;
@@ -12649,7 +12717,7 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stable
  
        Value slice;
        if (inputType.hasStaticShape()) {
-@@ -2980,9 +2992,18 @@
+@@ -2980,9 +2996,18 @@
        for (int64_t d = 0; d < rank; ++d) {
          if (d != dim) resultShape.push_back(sliceShape[d]);
        }
@@ -12669,7 +12737,7 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stable
            slice);
        slicedInputs.push_back(reshaped);
      }
-@@ -3004,9 +3025,20 @@
+@@ -3004,9 +3029,20 @@
        SmallVector<int64_t> elementShape(outputType.getShape().begin(),
                                          outputType.getShape().end());
        elementShape[dim] = 1;
diff --git a/third_party/xla/xla/hlo/builder/lib/math.cc b/third_party/xla/xla/hlo/builder/lib/math.cc
index fe7b94d54f7..85e87c861dd 100644
--- a/third_party/xla/xla/hlo/builder/lib/math.cc
+++ b/third_party/xla/xla/hlo/builder/lib/math.cc
@@ -2128,7 +2128,10 @@ XlaOp Zeta(XlaOp x, XlaOp q) {
     neg_power = Pow(acc, Neg(x));
 
     // Calculating `I`.
-    XlaOp I = neg_power * acc / (x - ScalarLike(acc, 1.));
+    // Compute acc^(1-x) directly rather than as acc^-x * acc: the latter
+    // underflows to 0 for large acc, and is 0 * inf = NaN for acc = inf.
+    XlaOp x_minus_one = x - ScalarLike(acc, 1.);
+    XlaOp I = Pow(acc, Neg(x_minus_one)) / x_minus_one;
 
     // Calculating the rest of the formula – `T`.
     // Use Horner's rule for this.
diff --git a/third_party/xla/xla/hlo/builder/lib/math_test.cc b/third_party/xla/xla/hlo/builder/lib/math_test.cc
index 0b126329390..826c1f2eff8 100644
--- a/third_party/xla/xla/hlo/builder/lib/math_test.cc
+++ b/third_party/xla/xla/hlo/builder/lib/math_test.cc
@@ -884,6 +884,30 @@ TEST_F(MathTest, ZetaF64) {
                               ErrorSpec{0.00000000000001});
 }
 
+TEST_F(MathTest, ZetaLargeQ) {
+  XlaBuilder builder(TestName());
+  const float inf = std::numeric_limits<float>::infinity();
+  auto x = ConstantR1<float>(&builder, {2.0f, 4.0f, 4.0f});
+  auto q = ConstantR1<float>(&builder, {inf, inf, 1e12f});
+  Zeta(x, q);
+  std::vector<float> expected = {0.0f, 0.0f, 1e-36f / 3.0f};
+
+  ComputeAndCompareR1<float>(&builder, expected, {},
+                             ErrorSpec{/*aabs=*/0.0, /*arel=*/1e-4});
+
+  if (test::BackendSupportsFloat64()) {
+    XlaBuilder builder_f64(TestName());
+    const double inf_f64 = std::numeric_limits<double>::infinity();
+    auto x_f64 = ConstantR1<double>(&builder_f64, {2.0, 4.0, 4.0});
+    auto q_f64 = ConstantR1<double>(&builder_f64, {inf_f64, inf_f64, 1e100});
+    Zeta(x_f64, q_f64);
+    std::vector<double> expected_f64 = {0.0, 0.0, 1e-300 / 3.0};
+
+    ComputeAndCompareR1<double>(&builder_f64, expected_f64, {},
+                                ErrorSpec{/*aabs=*/0.0, /*arel=*/1e-12});
+  }
+}
+
 TEST_F(MathTest, IgammaDynamicDimensionMismatch) {
   XlaBuilder builder(TestName());
   // Logical shape is identical, but one operand carries a dynamic-dimension
diff --git a/third_party/xla/xla/hlo/builder/tests/math_polygamma.hlo b/third_party/xla/xla/hlo/builder/tests/math_polygamma.hlo
index 52156a188b4..23255fa1980 100644
--- a/third_party/xla/xla/hlo/builder/tests/math_polygamma.hlo
+++ b/third_party/xla/xla/hlo/builder/tests/math_polygamma.hlo
@@ -360,10 +360,11 @@
 // CHECK-NEXT:  %[[constant_322:[^ ]+]] = f32[] constant(1.1920929e-07)
 // CHECK-NEXT:  %[[multiply_171:[^ ]+]] = f32[] multiply(%[[abs_14]], %[[constant_322]])
 // CHECK-NEXT:  %[[compare_43:[^ ]+]] = pred[] compare(%[[abs_13]], %[[multiply_171]]), direction=LT
-// CHECK-NEXT:  %[[multiply_134:[^ ]+]] = f32[] multiply(%[[power_32]], %[[add_288]])
 // CHECK-NEXT:  %[[constant_283:[^ ]+]] = f32[] constant(1)
 // CHECK-NEXT:  %[[subtract_36:[^ ]+]] = f32[] subtract(%[[add_236]], %[[constant_283]])
-// CHECK-NEXT:  %[[divide_64:[^ ]+]] = f32[] divide(%[[multiply_134]], %[[subtract_36]])
+// CHECK-NEXT:  %[[negate_36_1:[^ ]+]] = f32[] negate(%[[subtract_36]])
+// CHECK-NEXT:  %[[power_32_1:[^ ]+]] = f32[] power(%[[add_288]], %[[negate_36_1]])
+// CHECK-NEXT:  %[[divide_64:[^ ]+]] = f32[] divide(%[[power_32_1]], %[[subtract_36]])
 // CHECK-NEXT:  %[[add_324:[^ ]+]] = f32[] add(%[[add_287]], %[[divide_64]])
 // CHECK-NEXT:  %[[constant_320:[^ ]+]] = f32[] constant(0.5)
 // CHECK-NEXT:  %[[divide_66:[^ ]+]] = f32[] divide(%[[add_236]], %[[add_288]])
@@ -717,10 +718,11 @@
 // CHECK-NEXT:  %[[constant_442:[^ ]+]] = f32[] constant(1.1920929e-07)
 // CHECK-NEXT:  %[[multiply_224:[^ ]+]] = f32[] multiply(%[[abs_19]], %[[constant_442]])
 // CHECK-NEXT:  %[[compare_63:[^ ]+]] = pred[] compare(%[[abs_18]], %[[multiply_224]]), direction=LT
-// CHECK-NEXT:  %[[multiply_187:[^ ]+]] = f32[] multiply(%[[power_43]], %[[add_434]])
 // CHECK-NEXT:  %[[constant_403:[^ ]+]] = f32[] constant(1)
 // CHECK-NEXT:  %[[subtract_55:[^ ]+]] = f32[] subtract(%[[add_236]], %[[constant_403]])
-// CHECK-NEXT:  %[[divide_97:[^ ]+]] = f32[] divide(%[[multiply_187]], %[[subtract_55]])
+// CHECK-NEXT:  %[[negate_55_1:[^ ]+]] = f32[] negate(%[[subtract_55]])
+// CHECK-NEXT:  %[[power_43_1:[^ ]+]] = f32[] power(%[[add_434]], %[[negate_55_1]])
+// CHECK-NEXT:  %[[divide_97:[^ ]+]] = f32[] divide(%[[power_43_1]], %[[subtract_55]])
 // CHECK-NEXT:  %[[add_470:[^ ]+]] = f32[] add(%[[add_433]], %[[divide_97]])
 // CHECK-NEXT:  %[[constant_440:[^ ]+]] = f32[] constant(0.5)
 // CHECK-NEXT:  %[[divide_99:[^ ]+]] = f32[] divide(%[[add_236]], %[[add_434]])
diff --git a/third_party/xla/xla/mlir_hlo/tests/Dialect/chlo/chlo_legalize_to_mhlo.mlir b/third_party/xla/xla/mlir_hlo/tests/Dialect/chlo/chlo_legalize_to_mhlo.mlir
index 09752e3d03b..0a3eaaf1c88 100644
--- a/third_party/xla/xla/mlir_hlo/tests/Dialect/chlo/chlo_legalize_to_mhlo.mlir
+++ b/third_party/xla/xla/mlir_hlo/tests/Dialect/chlo/chlo_legalize_to_mhlo.mlir
@@ -1565,8 +1565,9 @@ func.func @zeta_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
   // CHECK: %[[TMP_33:.*]] = mhlo.add %[[TMP_30]], %[[TMP_3]]
   // CHECK: %[[TMP_34:.*]] = mhlo.power %[[TMP_33]], %[[TMP_4]]
   // CHECK: %[[TMP_35:.*]] = mhlo.constant dense<1.000000e+00>
-  // CHECK-DAG: %[[TMP_36:.*]] = mhlo.multiply %[[TMP_34]], %[[TMP_33]]
-  // CHECK-DAG: %[[TMP_37:.*]] = mhlo.subtract %[[TMP_0]], %[[TMP_35]]
+  // CHECK: %[[TMP_37:.*]] = mhlo.subtract %[[TMP_0]], %[[TMP_35]]
+  // CHECK: %[[TMP_36_NEG:.*]] = mhlo.negate %[[TMP_37]]
+  // CHECK: %[[TMP_36:.*]] = mhlo.power %[[TMP_33]], %[[TMP_36_NEG]]
   // CHECK: %[[TMP_38:.*]] = mhlo.divide %[[TMP_36]], %[[TMP_37]]
   // CHECK: %[[TMP_39:.*]] = mhlo.multiply %[[TMP_33]], %[[TMP_33]]
   // CHECK: %[[TMP_40:.*]] = mhlo.divide %[[TMP_3]], %[[TMP_39]]
@@ -1840,8 +1841,9 @@ func.func @polygamma_f32(%lhs : tensor<f32>, %rhs : tensor<f32>) -> tensor<f32>
   // CHECK: %[[TMP_121:.*]] = mhlo.add %[[TMP_118]], %[[TMP_91]]
   // CHECK: %[[TMP_122:.*]] = mhlo.power %[[TMP_121]], %[[TMP_92]]
   // CHECK: %[[TMP_123:.*]] = mhlo.constant dense<1.000000e+00>
-  // CHECK-DAG: %[[TMP_124:.*]] = mhlo.multiply %[[TMP_122]], %[[TMP_121]]
-  // CHECK-DAG: %[[TMP_125:.*]] = mhlo.subtract %[[TMP_5]], %[[TMP_123]]
+  // CHECK: %[[TMP_125:.*]] = mhlo.subtract %[[TMP_5]], %[[TMP_123]]
+  // CHECK: %[[TMP_124_NEG:.*]] = mhlo.negate %[[TMP_125]]
+  // CHECK: %[[TMP_124:.*]] = mhlo.power %[[TMP_121]], %[[TMP_124_NEG]]
   // CHECK: %[[TMP_126:.*]] = mhlo.divide %[[TMP_124]], %[[TMP_125]]
   // CHECK: %[[TMP_127:.*]] = mhlo.multiply %[[TMP_121]], %[[TMP_121]]
   // CHECK: %[[TMP_128:.*]] = mhlo.divide %[[TMP_91]], %[[TMP_127]]
@@ -2230,8 +2232,9 @@ func.func @polygamma_f64(%lhs : tensor<f64>, %rhs : tensor<f64>) -> tensor<f64>
   // CHECK: %[[TMP_121:.*]] = mhlo.add %[[TMP_118]], %[[TMP_91]]
   // CHECK: %[[TMP_122:.*]] = mhlo.power %[[TMP_121]], %[[TMP_92]]
   // CHECK: %[[TMP_123:.*]] = mhlo.constant dense<1.000000e+00>
-  // CHECK-DAG: %[[TMP_124:.*]] = mhlo.multiply %[[TMP_122]], %[[TMP_121]]
-  // CHECK-DAG: %[[TMP_125:.*]] = mhlo.subtract %[[TMP_5]], %[[TMP_123]]
+  // CHECK: %[[TMP_125:.*]] = mhlo.subtract %[[TMP_5]], %[[TMP_123]]
+  // CHECK: %[[TMP_124_NEG:.*]] = mhlo.negate %[[TMP_125]]
+  // CHECK: %[[TMP_124:.*]] = mhlo.power %[[TMP_121]], %[[TMP_124_NEG]]
   // CHECK: %[[TMP_126:.*]] = mhlo.divide %[[TMP_124]], %[[TMP_125]]
   // CHECK: %[[TMP_127:.*]] = mhlo.multiply %[[TMP_121]], %[[TMP_121]]
   // CHECK: %[[TMP_128:.*]] = mhlo.divide %[[TMP_91]], %[[TMP_127]]

From 642635c63e5a413a4a1c3b69a976a814caae7742 Mon Sep 17 00:00:00 2001
From: Jake VanderPlas <vanderplas@google.com>
Date: Thu, 15 Oct 2026 07:02:49 -0700
Subject: [PATCH 0871/1000] [XLA] Use CloneWithNewOperands in HandleReverse for
 binary ops.

When simplifying a full-slice scatter followed by a reverse (such as
`jnp.zeros(4).at[:].set(x)[::-1]`), `HandleScatter` rewrites the
0-element index scatter into a 2-operand `kMap` instruction. In the same
pass, `HandleReverse` visits `reverse(map)` where `map` satisfies
`IsElementwiseBinary()`. When one of the operands is a scalar broadcast,
`HandleReverse` previously called `HloInstruction::CreateBinary()`,
which crashes with a fatal error on `kMap` because `CreateBinary` does
not support `kMap`.

Replace `HloInstruction::CreateBinary()` with `CloneWithNewOperands()`,
guard against side-effecting instructions, and unwrap `MakeReverseHlo`
with `ASSIGN_OR_RETURN`.

Fixes https://github.com/jax-ml/jax/issues/41497

PiperOrigin-RevId: 1000109848
---
 .../simplifiers/algebraic_simplifier.cc       | 33 ++++----
 .../simplifiers/algebraic_simplifier_test.cc  | 82 +++++++++++++++++++
 2 files changed, 99 insertions(+), 16 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
index 98c101a5671..e66f267ad7e 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
@@ -7118,28 +7118,29 @@ absl::Status AlgebraicSimplifierVisitor::HandleReverse(
     };
     if (inner->IsElementwiseBinary()) {
       // produces incorrect result for rng.
-      if (inner->opcode() == HloOpcode::kRng ||
-          inner_opcode == HloOpcode::kCompare) {
+      if (inner_opcode == HloOpcode::kCompare ||
+          inner_opcode == HloOpcode::kFusion || inner->HasSideEffect()) {
         return absl::OkStatus();
       }
-      HloInstruction* hlo;
       if (match_with_scalar(inner->mutable_operand(0))) {
-        hlo = inner->mutable_operand(1);
+        HloInstruction* hlo = inner->mutable_operand(1);
+        ABSL_ASSIGN_OR_RETURN(HloInstruction * reverse_hlo,
+                         MakeReverseHlo(hlo, reverse_dims));
         return ReplaceWithNewInstruction(
             reverse,
-            HloInstruction::CreateBinary(inner->shape(), inner_opcode,
-                                         inner->mutable_operand(0),
-                                         *MakeReverseHlo(hlo, reverse_dims)));
-      } else if (match_with_scalar(inner->mutable_operand(1))) {
-        hlo = inner->mutable_operand(0);
-        return ReplaceWithNewInstruction(
-            reverse,
-            HloInstruction::CreateBinary(inner->shape(), inner_opcode,
-                                         *MakeReverseHlo(hlo, reverse_dims),
-                                         inner->mutable_operand(1)));
-      } else {
-        return absl::OkStatus();
+            inner->CloneWithNewOperands(
+                inner->shape(), {inner->mutable_operand(0), reverse_hlo}));
       }
+      if (match_with_scalar(inner->mutable_operand(1))) {
+        HloInstruction* hlo = inner->mutable_operand(0);
+        ABSL_ASSIGN_OR_RETURN(HloInstruction * reverse_hlo,
+                         MakeReverseHlo(hlo, reverse_dims));
+        return ReplaceWithNewInstruction(
+            reverse,
+            inner->CloneWithNewOperands(
+                inner->shape(), {reverse_hlo, inner->mutable_operand(1)}));
+      }
+      return absl::OkStatus();
     }
     // reverse(DegenerateDimensionAddingReshape(x)) ==>
     // DegenerateDimensionAddingReshape(reverse(x))
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
index 3b7af4365b8..80b9590d0d8 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
@@ -13755,6 +13755,88 @@ TEST_F(AlgebraicSimplifierTest, SwapConstantEwboWithReverse2) {
                                       m::Reverse(m::Parameter(0)))));
 }
 
+TEST_F(AlgebraicSimplifierTest, SwapConstantMapWithReverse) {
+  constexpr absl::string_view kModuleStr = R"(
+    HloModule m
+    mapped_computation {
+      lhs = bf16[] parameter(0)
+      rhs = bf16[] parameter(1)
+      mul = bf16[] multiply(lhs, rhs)
+      ROOT add = bf16[] add(mul, rhs)
+    }
+    ENTRY test {
+      p0 = bf16[100,1,100,512] parameter(0)
+      constant0 = bf16[] constant(1)
+      const = bf16[100,1,100,512] broadcast(constant0), dimensions={}
+      map0 = bf16[100,1,100,512] map(const, p0), dimensions={0,1,2,3},
+          to_apply=mapped_computation
+      ROOT r0 = bf16[100,1,100,512] reverse(map0), dimensions={0}
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  SCOPED_TRACE("Before rewrite\n" + m->ToString());
+  AlgebraicSimplifier simplifier(default_options_);
+  ASSERT_OK_AND_ASSIGN(bool result, simplifier.Run(m.get()));
+  SCOPED_TRACE("After rewrite\n" + m->ToString());
+  ASSERT_TRUE(result);
+  auto* root = m->entry_computation()->root_instruction();
+  EXPECT_THAT(root, GmockMatch(m::Map(m::Broadcast(m::Constant()),
+                                      m::Reverse(m::Parameter(0)))));
+}
+
+TEST_F(AlgebraicSimplifierTest, DoNotSwapConstantFusionWithReverse) {
+  constexpr absl::string_view kModuleStr = R"(
+    HloModule m
+    fused_computation {
+      lhs = f32[4] parameter(0)
+      rhs = f32[4] parameter(1)
+      c = f32[4] constant({1, 2, 3, 4})
+      add0 = f32[4] add(lhs, rhs)
+      ROOT add1 = f32[4] add(add0, c)
+    }
+    ENTRY test {
+      p0 = f32[4] parameter(0)
+      c0 = f32[] constant(1)
+      b0 = f32[4] broadcast(c0), dimensions={}
+      fusion = f32[4] fusion(b0, p0), kind=kLoop, calls=fused_computation
+      ROOT r0 = f32[4] reverse(fusion), dimensions={0}
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  AlgebraicSimplifier simplifier(default_options_);
+  ASSERT_OK_AND_ASSIGN(bool result, simplifier.Run(m.get()));
+  EXPECT_FALSE(result);
+}
+
+TEST_F(AlgebraicSimplifierTest, FullSliceScatterFollowedByReverse) {
+  constexpr absl::string_view kModuleStr = R"(
+    HloModule m
+    update_computation {
+      lhs = f32[] parameter(0)
+      ROOT rhs = f32[] parameter(1)
+    }
+    ENTRY main {
+      p0 = f32[4] parameter(0)
+      zero = f32[] constant(0)
+      operand = f32[4] broadcast(zero), dimensions={}
+      indices = s32[0] constant({})
+      scatter = f32[4] scatter(operand, indices, p0),
+          update_window_dims={0},
+          inserted_window_dims={},
+          scatter_dims_to_operand_dims={},
+          index_vector_dim=0,
+          to_apply=update_computation
+      ROOT reverse = f32[4] reverse(scatter), dimensions={0}
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  HloPassFix<AlgebraicSimplifier> simplifier(default_options_);
+  ASSERT_OK_AND_ASSIGN(bool result, simplifier.Run(m.get()));
+  ASSERT_TRUE(result);
+  auto* root = m->entry_computation()->root_instruction();
+  EXPECT_THAT(root, GmockMatch(m::Reverse(m::Parameter(0))));
+}
+
 TEST_F(AlgebraicSimplifierTest, SquaredComplexSqrtIsFloat) {
   const absl::string_view kModuleStr = R"(
   HloModule module

From 42cb08536cee30233d4120bf7229d63fedeeff37 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 15 Oct 2026 07:14:37 -0700
Subject: [PATCH 0872/1000] Automated Code Change

PiperOrigin-RevId: 1000115884
---
 tensorflow/core/kernels/BUILD | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index c3adbc128a3..34b9e54ac80 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -8376,9 +8376,7 @@ cc_library(
     name = "loose_headers",
     tags = ["avoid_dep"],
     textual_hdrs = glob(["*.h"]),
-    visibility = [
-        "//tensorflow:__subpackages__",
-    ] + LOOSE_HEADER_VISIBILITY,
+    visibility = LOOSE_HEADER_VISIBILITY,
     deps = [
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/log:check",

From ba00b30766eef9c3f4bb42907887b6ffa36b7fec Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 15 Oct 2026 08:01:56 -0700
Subject: [PATCH 0873/1000] Reverts 5c8408ccd09632ecf57bab8cd7dcac9c4dcdd12e

PiperOrigin-RevId: 1000140438
---
 .../service/heap_simulator/heap_simulator.cc  | 50 ++------------
 .../service/heap_simulator/heap_simulator.h   | 11 ---
 .../heap_simulator/heap_simulator_test.cc     | 68 -------------------
 3 files changed, 4 insertions(+), 125 deletions(-)

diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
index 94b73067687..89fe04a9dde 100644
--- a/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
+++ b/third_party/xla/xla/service/heap_simulator/heap_simulator.cc
@@ -640,13 +640,9 @@ void HeapSimulator::FillDebugTrace(HeapSimulatorTrace::Event::Kind kind,
 template <typename BufferType>
 void NoFragmentationStatsHeap<BufferType>::Alloc(const BufferType* buffer,
                                                  int64_t size) {
-  BufferInfo& info = buffer_info_[buffer];
-  info.ref_count++;
-  if (info.ref_count == 1) {
-    current_heap_size_ += size;
-    if (current_heap_size_ > max_heap_size_) {
-      max_heap_size_ = current_heap_size_;
-    }
+  current_heap_size_ += size;
+  if (current_heap_size_ > max_heap_size_) {
+    max_heap_size_ = current_heap_size_;
   }
 }
 
@@ -668,48 +664,10 @@ void NoFragmentationStatsHeap<BufferType>::AccountForSubcomputationMemory(
       std::max(max_heap_size_, current_heap_size_ + max_subcomputation_bytes);
 }
 
-template <typename BufferType>
-const BufferType* NoFragmentationStatsHeap<BufferType>::FindRoot(
-    const BufferType* buffer) {
-  const BufferType* root = buffer;
-  while (buffer_info_[root].underlying != nullptr) {
-    root = buffer_info_[root].underlying;
-  }
-  const BufferType* curr = buffer;
-  while (curr != root) {
-    const BufferType* next = buffer_info_[curr].underlying;
-    buffer_info_[curr].underlying = root;
-    curr = next;
-  }
-  return root;
-}
-
 template <typename BufferType>
 void NoFragmentationStatsHeap<BufferType>::Free(const BufferType* buffer,
                                                 int64_t size) {
-  const BufferType* root = FindRoot(buffer);
-  BufferInfo& root_info = buffer_info_[root];
-  root_info.ref_count--;
-  if (root_info.ref_count == 0) {
-    current_heap_size_ -= size;
-  }
-}
-
-template <typename BufferType>
-void NoFragmentationStatsHeap<BufferType>::ShareWith(
-    const BufferType* buffer, const BufferType* share_with, int64_t size) {
-  const BufferType* root = FindRoot(share_with);
-  if (buffer != root) {
-    buffer_info_[buffer].underlying = root;
-  }
-  BufferInfo& root_info = buffer_info_[root];
-  root_info.ref_count++;
-  if (root_info.ref_count == 1) {
-    current_heap_size_ += size;
-    if (current_heap_size_ > max_heap_size_) {
-      max_heap_size_ = current_heap_size_;
-    }
-  }
+  current_heap_size_ -= size;
 }
 
 template <typename BufferType>
diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator.h b/third_party/xla/xla/service/heap_simulator/heap_simulator.h
index 949e4338f2b..edacc1ca4e7 100644
--- a/third_party/xla/xla/service/heap_simulator/heap_simulator.h
+++ b/third_party/xla/xla/service/heap_simulator/heap_simulator.h
@@ -341,20 +341,9 @@ class NoFragmentationStatsHeap : public HeapAlgorithm<BufferType> {
 
   void Free(const BufferType* buffer, int64_t size) override;
 
-  void ShareWith(const BufferType* buffer, const BufferType* share_with,
-                 int64_t size) override;
-
   absl::StatusOr<Result> Finish() override;
 
  private:
-  struct BufferInfo {
-    const BufferType* underlying = nullptr;
-    int64_t ref_count = 0;
-  };
-
-  const BufferType* FindRoot(const BufferType* buffer);
-
-  absl::flat_hash_map<const BufferType*, BufferInfo> buffer_info_;
   int64_t current_heap_size_ = 0;
   int64_t max_heap_size_ = 0;
 };
diff --git a/third_party/xla/xla/service/heap_simulator/heap_simulator_test.cc b/third_party/xla/xla/service/heap_simulator/heap_simulator_test.cc
index cf3d0514428..1fc11b0d65e 100644
--- a/third_party/xla/xla/service/heap_simulator/heap_simulator_test.cc
+++ b/third_party/xla/xla/service/heap_simulator/heap_simulator_test.cc
@@ -245,53 +245,6 @@ TEST_F(MinimumMemoryForSequenceTest, SubcomputationAccounting) {
                     .value());
 }
 
-TEST_F(MinimumMemoryForSequenceTest, WhileLoopDoubleCounting) {
-  // Demonstrates that MinimumMemoryForModule (via NoFragmentationStatsHeap)
-  // does not double count shared buffers between the while loop condition,
-  // body, and init.
-  constexpr absl::string_view kHloString = R"hlo(
-    HloModule WhileLoopDoubleCounting, is_scheduled=true
-
-    WhileCond {
-      cond_param = f32[1000] parameter(0)
-      slice = f32[1] slice(cond_param), slice={[0:1]}
-      reshape = f32[] reshape(slice)
-      zero = f32[] constant(0)
-      ROOT cond_comparison = pred[] compare(reshape, zero), direction=NE
-    }
-
-    WhileBody {
-      body_param = f32[1000] parameter(0)
-      ROOT negate = f32[1000] negate(body_param)
-    }
-
-    ENTRY Entry {
-      init = f32[1000] parameter(0)
-      ROOT while_loop = f32[1000] while(init), condition=WhileCond, body=WhileBody
-    }
-  )hlo";
-
-  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
-                       ParseAndReturnVerifiedModule(kHloString));
-  ASSERT_OK(module->schedule().Verify());
-
-  BufferValue::SizeFunction size_fn = [](const BufferValue& buffer) {
-    return ShapeUtil::ByteSizeOf(buffer.shape());
-  };
-
-  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloAliasAnalysis> alias_analysis,
-                       HloAliasAnalysis::Run(module.get(), &alias_info_));
-
-  // With the NoFragmentationStatsHeap fix in place, the peak memory correctly
-  // avoids double counting the 4000-byte shared buffer across the condition and
-  // body, resulting in 4008 bytes (4000 + 4 slice + 4 reshape).
-  ASSERT_OK_AND_ASSIGN(
-      int64_t min_memory,
-      HeapSimulator::MinimumMemoryForModule(module->schedule(), *alias_analysis,
-                                            &alias_info_, &size_fn));
-  EXPECT_EQ(min_memory, 4008);
-}
-
 const char kAlloc[] = "Alloc";
 const char kFree[] = "Free";
 const char kShare[] = "Share";
@@ -1154,27 +1107,6 @@ TEST_F(NoFragmentationStatsHeapTest, Mixed) {
   EXPECT_EQ(40, result.heap_size);
 }
 
-TEST_F(NoFragmentationStatsHeapTest, ShareWith) {
-  NoFragmentationStatsHeap<HloValue> heap;
-  heap.Alloc(buffer_a_, 100);                 // current: 100, max: 100
-  heap.ShareWith(buffer_b_, buffer_a_, 100);  // current: 100, max: 100
-  heap.ShareWith(buffer_c_, buffer_b_, 100);  // current: 100, max: 100
-  heap.Free(buffer_b_, 100);                  // current: 100, max: 100
-  heap.ShareWith(buffer_d_, buffer_b_, 100);  // current: 100, max: 100
-  heap.Free(buffer_a_, 100);                  // current: 100, max: 100
-  heap.Free(buffer_c_, 100);                  // current: 100, max: 100
-  heap.Free(buffer_d_, 100);                  // current: 0, max: 100
-
-  heap.Alloc(buffer_e_, 50);  // current: 50, max: 100
-  heap.Alloc(buffer_f_, 60);  // current: 110, max: 110
-  heap.Free(buffer_e_, 50);
-  heap.Free(buffer_f_, 60);
-
-  ASSERT_OK_AND_ASSIGN(const HeapSimulator::Result<HloValue> result,
-                       heap.Finish());
-  EXPECT_EQ(110, result.heap_size);
-}
-
 class GlobalDecreasingSizeBestFitHeapTest : public HeapAlgorithmTestBase {};
 
 TEST_F(GlobalDecreasingSizeBestFitHeapTest, Empty) {

From ccc81fe6c020af2139ba336eccc5ba21d6a60c74 Mon Sep 17 00:00:00 2001
From: Mohammed Anany <manany@google.com>
Date: Thu, 15 Oct 2026 08:02:16 -0700
Subject: [PATCH 0874/1000] [XLA:GPU] Support sequential dimension tiling in
 BlockLevelFusionConfig. This should allow representation for reduce and scan
 operations.

Enable candidate tile generation for `kReduce` in `TilingSpace::GetValidTilings`
and populate `sequential_tile_sizes` only when the fusion contains a single
reduce that reduces a single dimension.

`xla_gpu_experimental_enable_reduction_dimension_tiling` is introduced
(currently disabled) to control the feature.

Enablement is pending changes to the cost model and ensuring performance gains
are realized.

PiperOrigin-RevId: 1000140675
---
 .../reduce/reduce_tiled_block_config.hlo      | 52 ++++++++++
 .../tiling/experimental/tiling_space.cc       | 40 +++++---
 .../tiling/experimental/tiling_space.h        |  9 +-
 .../tiling/experimental/tiling_space_test.cc  | 94 ++++++++++++++++++
 .../codegen/xtile/block_level_parameters.h    |  8 ++
 .../xtile/block_level_parameters_test.cc      | 23 +++++
 .../xtile/tiling_from_block_parameters.cc     | 42 ++++++--
 .../tiling_from_block_parameters_test.cc      | 98 +++++++++++++++++++
 .../xla/xla/codegen/xtile/xtile_config.proto  |  6 ++
 third_party/xla/xla/debug_options_flags.cc    |  8 ++
 third_party/xla/xla/service/gpu/model/BUILD   |  1 +
 .../model/gpu_indexing_performance_model.cc   | 58 ++++++++++-
 third_party/xla/xla/xla.proto                 |  9 +-
 13 files changed, 423 insertions(+), 25 deletions(-)
 create mode 100644 third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_tiled_block_config.hlo

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_tiled_block_config.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_tiled_block_config.hlo
new file mode 100644
index 00000000000..f8e023030cf
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/reduce/reduce_tiled_block_config.hlo
@@ -0,0 +1,52 @@
+// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s | FileCheck %s
+// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation --abs_error_bound=1e-7 --rel_error_bound=1e-7 %s
+
+// Same as reduce_tiled.hlo, but the reduction tile size comes from
+// `sequential_tile_sizes` in the block-level fusion config instead of the
+// backend config of the reduce instruction.
+
+region {
+  p0 = f32[] parameter(0)
+  p1 = f32[] parameter(1)
+  ROOT add = f32[] add(p0, p1)
+}
+
+fusion {
+  p = f32[8,16] parameter(0)
+  c = f32[] constant(0)
+  ROOT reduce = f32[8] reduce(p, c), dimensions={1}, to_apply=region
+}
+
+ENTRY main {
+  p = f32[8,16] parameter(0)
+  ROOT fusion = f32[8] fusion(p), kind=kCustom, calls=fusion,
+    backend_config={
+      "fusion_backend_config":{
+        "kind":"__triton",
+        "block_level_fusion_config":{
+          "output_tiles":[{"sizes":["4"]}],
+          "sequential_tile_sizes":["4"],
+          "num_warps":"1",
+          "num_ctas":"1",
+          "num_stages":"1"}}}
+}
+
+// CHECK-DAG:  %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG:  %[[C4:.*]] = arith.constant 4 : index
+// CHECK-DAG:  %[[C1:.*]] = arith.constant 1 : index
+// CHECK-DAG:  %[[INIT:.*]] = arith.constant dense<0.000000e+00> : tensor<4x4xf32>
+
+// CHECK: %[[FOR:.*]] = scf.for %[[IV:.*]] = %[[C0]] to %[[C4]] step %[[C1]]
+// CHECK-SAME: iter_args(%[[ACC:.*]] = %[[INIT]]) -> (tensor<4x4xf32>) {
+
+// CHECK:   %[[OFFSET_P:.*]] = xla.apply_indexing
+// CHECK:   %[[OFFSET_R:.*]] = xla.apply_indexing
+// CHECK:   %[[EXTRACT:.*]] = xtile.extract %{{.*}}[%[[OFFSET_P]], %[[OFFSET_R]]] [4, 4]
+// CHECK:   %[[COMBINE:.*]] = arith.addf %[[ACC]], %[[EXTRACT]] : tensor<4x4xf32>
+// CHECK:   scf.yield %[[COMBINE]] : tensor<4x4xf32>
+// CHECK: }
+// CHECK: %[[REDUCE:.*]] = "tt.reduce"(%[[FOR]]) <{axis = 1 : i32}> ({
+// CHECK: ^bb0(%[[ARG0:.*]]: f32, %[[ARG1:.*]]: f32):
+// CHECK:   %[[SUM:.*]] = arith.addf %[[ARG0]], %[[ARG1]] : f32
+// CHECK:   tt.reduce.return %[[SUM]] : f32
+// CHECK: }) : (tensor<4x4xf32>) -> tensor<4xf32>
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc
index a50825475fa..5c63f41c3ae 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.cc
@@ -751,6 +751,12 @@ int64_t TilingSpace::num_parallel_dimensions() const {
   });
 }
 
+int64_t TilingSpace::num_sequential_dimensions() const {
+  return absl::c_count_if(dimensions_, [](const DimensionInfo& dim) {
+    return dim.type == DimensionSemantics::kSequential;
+  });
+}
+
 void TilingSpace::InitSimplificationIndexing() {
   CHECK(!is_symbolic_) << "Tile sizes must be assigned before initializing "
                           "cached indexing map variables.";
@@ -821,24 +827,35 @@ TilingSpace::SimplificationResult TilingSpace::SimplifyExpressions(
 }
 
 absl::StatusOr<std::vector<llvm::SmallVector<int64_t, 4>>>
-TilingSpace::GetValidTilings() {
+TilingSpace::GetValidTilings(bool tile_reduction_dimensions) {
   // TODO: b/511080616 - returned tilings should be valid. Right now we return
   // all possible tilings and rely on the downstream to check the validity.
   llvm::SmallVector<int64_t, 4> input_space;
 
-  // Sequential reduce dimensions are not tiled yet. To work around the
-  // limitation of `GetFlatTilingsForInputSpace`, we set the tile size to 1 here
-  // and later replace with the actual dimension size.
+  // TODO(b/534676634): Lift this condition once multi-reduction and
+  // multi-dimension reduction fusions are supported.
+  // Reduction dimensions are only tiled when requested and the fusion has a
+  // single reduce that reduces a single dimension. Otherwise, sequential reduce
+  // dimensions are not tiled. To work around the limitation of
+  // `GetFlatTilingsForInputSpace`, we set the tile size to 1 here and later
+  // replace with the actual dimension size.
+  const int64_t num_reduction_dims =
+      absl::c_count_if(dimensions_, [](const DimensionInfo& dim) {
+        return dim.type == DimensionSemantics::kSequential &&
+               dim.hlo->opcode() == HloOpcode::kReduce;
+      });
+  const bool tile_reduction_dims =
+      tile_reduction_dimensions && num_reduction_dims == 1;
+  auto is_untiled_dim = [&](const DimensionInfo& dim) {
+    return !tile_reduction_dims &&
+           dim.type == DimensionSemantics::kSequential &&
+           dim.hlo->opcode() == HloOpcode::kReduce;
+  };
   for (const auto& dim : dimensions_) {
     if (dim.type == DimensionSemantics::kReplica) {
       continue;
     }
-    if (dim.type == DimensionSemantics::kSequential &&
-        dim.hlo->opcode() == HloOpcode::kReduce) {
-      input_space.push_back(1);
-    } else {
-      input_space.push_back(dim.dimension_size);
-    }
+    input_space.push_back(is_untiled_dim(dim) ? 1 : dim.dimension_size);
   }
 
   ABSL_ASSIGN_OR_RETURN(auto flat_tilings, GetFlatTilingsForInputSpace(input_space));
@@ -849,8 +866,7 @@ TilingSpace::GetValidTilings() {
       if (dim.type == DimensionSemantics::kReplica) {
         continue;
       }
-      if (dim.type == DimensionSemantics::kSequential &&
-          dim.hlo->opcode() == HloOpcode::kReduce) {
+      if (is_untiled_dim(dim)) {
         flat_tiling[idx] = dim.dimension_size;
       }
       ++idx;
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h
index b5c698edd9b..c69f7db9998 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space.h
@@ -201,6 +201,7 @@ class TilingSpace {
 
   int64_t num_dimensions() const { return dimensions_.size(); }
   int64_t num_parallel_dimensions() const;
+  int64_t num_sequential_dimensions() const;
   int64_t num_rt_vars() const { return rt_vars_.size(); }
 
   void AppendDimension(const HloInstruction* hlo, int64_t dim_position,
@@ -243,8 +244,12 @@ class TilingSpace {
       const llvm::SmallVector<SymbolicExpr>& expressions,
       llvm::ArrayRef<std::pair<SymbolicExpr, Interval>> constraints = {}) const;
 
-  // Returns the list of valid tilings for the tiling space.
-  absl::StatusOr<std::vector<llvm::SmallVector<int64_t, 4>>> GetValidTilings();
+  // Returns the list of valid tilings for the tiling space. If
+  // `tile_reduction_dimensions` is set, tile sizes for the sequential
+  // dimensions of reduces are enumerated too; otherwise each of them is covered
+  // by a single tile of the dimension size.
+  absl::StatusOr<std::vector<llvm::SmallVector<int64_t, 4>>> GetValidTilings(
+      bool tile_reduction_dimensions = false);
 
  private:
   absl::Status InitializeDimensions(
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc
index e416022793a..2f4c5ee8121 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tiling_space_test.cc
@@ -1113,5 +1113,99 @@ TEST_F(TilingSpaceTest,
   }
 }
 
+TEST_F(TilingSpaceTest, GetValidTilingsTilesReductionDimOfSingleDimReduce) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+      HloModule m
+      add {
+        p0 = f32[] parameter(0)
+        p1 = f32[] parameter(1)
+        ROOT res = f32[] add(p0, p1)
+      }
+      ENTRY e {
+        p0 = f32[2,4] parameter(0)
+        c0 = f32[] constant(0)
+        ROOT r = f32[2] reduce(p0, c0), dimensions={1}, to_apply=add
+      }
+  )");
+  auto fusion_adaptor = HloFusionAdaptor::ForInstruction(root);
+  ASSERT_OK_AND_ASSIGN(auto tiling_space,
+                       TilingSpace::Create(*fusion_adaptor, &mlir_context_));
+  ASSERT_OK_AND_ASSIGN(
+      auto valid_tilings,
+      tiling_space->GetValidTilings(/*tile_reduction_dimensions=*/true));
+  EXPECT_THAT(valid_tilings, ElementsAre(ElementsAre(1, 1), ElementsAre(1, 2),
+                                         ElementsAre(1, 4), ElementsAre(2, 1),
+                                         ElementsAre(2, 2), ElementsAre(2, 4)));
+
+  // By default, the reduction dimension is covered by a single tile.
+  ASSERT_OK_AND_ASSIGN(auto untiled_valid_tilings,
+                       tiling_space->GetValidTilings());
+  EXPECT_THAT(untiled_valid_tilings,
+              ElementsAre(ElementsAre(1, 4), ElementsAre(2, 4)));
+}
+
+// TODO(b/534676634): Remove this test once tiling reduction dimensions of
+// multi-reduction fusions is supported.
+TEST_F(TilingSpaceTest,
+       GetValidTilingsDoesNotTileReductionDimsOfMultipleReduces) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+      HloModule m
+      add {
+        p0 = f32[] parameter(0)
+        p1 = f32[] parameter(1)
+        ROOT res = f32[] add(p0, p1)
+      }
+      fused_computation {
+        p0 = f32[2,4,8] parameter(0)
+        c0 = f32[] constant(0)
+        r0 = f32[2,4] reduce(p0, c0), dimensions={2}, to_apply=add
+        ROOT r1 = f32[2] reduce(r0, c0), dimensions={1}, to_apply=add
+      }
+      ENTRY e {
+        p0 = f32[2,4,8] parameter(0)
+        ROOT fusion = f32[2] fusion(p0), kind=kLoop, calls=fused_computation
+      }
+  )");
+  auto fusion_adaptor = HloFusionAdaptor::ForInstruction(root);
+  ASSERT_OK_AND_ASSIGN(auto tiling_space,
+                       TilingSpace::Create(*fusion_adaptor, &mlir_context_));
+  ASSERT_OK_AND_ASSIGN(
+      auto valid_tilings,
+      tiling_space->GetValidTilings(/*tile_reduction_dimensions=*/true));
+  // Dimensions are ordered use-before-def: r1's reduction dimension (4) comes
+  // before r0's (8). Both are kept at their full size.
+  EXPECT_THAT(valid_tilings,
+              ElementsAre(ElementsAre(1, 4, 8), ElementsAre(2, 4, 8)));
+}
+
+// TODO(b/534676634): Remove this test once tiling reduction dimensions of
+// multi-dimension reductions is supported.
+TEST_F(TilingSpaceTest,
+       GetValidTilingsDoesNotTileReductionDimsOfMultiDimReduce) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+      HloModule m
+      add {
+        p0 = f32[] parameter(0)
+        p1 = f32[] parameter(1)
+        ROOT res = f32[] add(p0, p1)
+      }
+      ENTRY e {
+        p0 = f32[2,4,8] parameter(0)
+        c0 = f32[] constant(0)
+        ROOT r = f32[2] reduce(p0, c0), dimensions={1,2}, to_apply=add
+      }
+  )");
+  auto fusion_adaptor = HloFusionAdaptor::ForInstruction(root);
+  ASSERT_OK_AND_ASSIGN(auto tiling_space,
+                       TilingSpace::Create(*fusion_adaptor, &mlir_context_));
+  ASSERT_OK_AND_ASSIGN(
+      auto valid_tilings,
+      tiling_space->GetValidTilings(/*tile_reduction_dimensions=*/true));
+  // Multi-dimension reductions keep both reduction dimensions at their full
+  // size.
+  EXPECT_THAT(valid_tilings,
+              ElementsAre(ElementsAre(1, 4, 8), ElementsAre(2, 4, 8)));
+}
+
 }  // namespace
 }  // namespace xla::gpu::experimental
diff --git a/third_party/xla/xla/codegen/xtile/block_level_parameters.h b/third_party/xla/xla/codegen/xtile/block_level_parameters.h
index a50cad4f6eb..c24ccc01b86 100644
--- a/third_party/xla/xla/codegen/xtile/block_level_parameters.h
+++ b/third_party/xla/xla/codegen/xtile/block_level_parameters.h
@@ -40,6 +40,10 @@ struct BlockLevelParameters {
   int num_tiles_per_pid = 1;
   int waves_per_eu = 0;
 
+  // Tile sizes of the sequential dimensions, in tiling space order. Empty when
+  // sequential tile sizes are not set in the fusion config.
+  std::vector<int64_t> sequential_tile_sizes;
+
   // Returns a BlockLevelParameters struct from a BlockLevelFusionConfig proto.
   static BlockLevelParameters FromBlockLevelFusionConfig(
       const xla::xtile::BlockLevelFusionConfig& config) {
@@ -57,6 +61,8 @@ struct BlockLevelParameters {
       result.output_tile_sizes.push_back(
           std::vector<int64_t>(tile.sizes().begin(), tile.sizes().end()));
     }
+    result.sequential_tile_sizes.assign(config.sequential_tile_sizes().begin(),
+                                        config.sequential_tile_sizes().end());
     return result;
   }
 
@@ -68,6 +74,8 @@ struct BlockLevelParameters {
       tile.mutable_sizes()->Add(tile_sizes.begin(), tile_sizes.end());
       *config.add_output_tiles() = tile;
     }
+    config.mutable_sequential_tile_sizes()->Add(sequential_tile_sizes.begin(),
+                                                sequential_tile_sizes.end());
     config.set_num_warps(num_warps);
     config.set_num_ctas(num_ctas);
     config.set_num_stages(num_stages);
diff --git a/third_party/xla/xla/codegen/xtile/block_level_parameters_test.cc b/third_party/xla/xla/codegen/xtile/block_level_parameters_test.cc
index ffff13b1b16..3e924483bd2 100644
--- a/third_party/xla/xla/codegen/xtile/block_level_parameters_test.cc
+++ b/third_party/xla/xla/codegen/xtile/block_level_parameters_test.cc
@@ -23,6 +23,7 @@ namespace xla::xtile {
 namespace {
 
 using ::testing::ElementsAre;
+using ::testing::IsEmpty;
 using ::xla::xtile::BlockLevelFusionConfig;
 
 TEST(BlockLevelParametersTest,
@@ -32,6 +33,8 @@ TEST(BlockLevelParametersTest,
   tile.mutable_sizes()->Add(18);
   tile.mutable_sizes()->Add(19);
   *block_level_fusion_config.add_output_tiles() = tile;
+  block_level_fusion_config.add_sequential_tile_sizes(32);
+  block_level_fusion_config.add_sequential_tile_sizes(64);
   block_level_fusion_config.set_num_warps(12);
   block_level_fusion_config.set_num_ctas(13);
   block_level_fusion_config.set_num_stages(14);
@@ -44,6 +47,8 @@ TEST(BlockLevelParametersTest,
           block_level_fusion_config);
   EXPECT_EQ(block_level_parameters.output_tile_sizes.size(), 1);
   EXPECT_THAT(block_level_parameters.output_tile_sizes[0], ElementsAre(18, 19));
+  EXPECT_THAT(block_level_parameters.sequential_tile_sizes,
+              ElementsAre(32, 64));
   EXPECT_THAT(block_level_parameters.num_warps, 12);
   EXPECT_THAT(block_level_parameters.num_ctas, 13);
   EXPECT_THAT(block_level_parameters.num_stages, 14);
@@ -56,6 +61,7 @@ TEST(BlockLevelParametersTest,
      BlockLevelParametersCanBeConvertedToBlockLevelFusionConfig) {
   BlockLevelParameters block_level_parameters;
   block_level_parameters.output_tile_sizes = {{18, 19}};
+  block_level_parameters.sequential_tile_sizes = {32, 64};
   block_level_parameters.num_warps = 12;
   block_level_parameters.num_ctas = 13;
   block_level_parameters.num_stages = 14;
@@ -69,6 +75,8 @@ TEST(BlockLevelParametersTest,
   EXPECT_EQ(block_level_fusion_config.output_tiles_size(), 1);
   EXPECT_THAT(block_level_fusion_config.output_tiles(0).sizes(),
               ElementsAre(18, 19));
+  EXPECT_THAT(block_level_fusion_config.sequential_tile_sizes(),
+              ElementsAre(32, 64));
   EXPECT_THAT(block_level_fusion_config.num_warps(), 12);
   EXPECT_THAT(block_level_fusion_config.num_ctas(), 13);
   EXPECT_THAT(block_level_fusion_config.num_stages(), 14);
@@ -77,5 +85,20 @@ TEST(BlockLevelParametersTest,
   EXPECT_THAT(block_level_fusion_config.waves_per_eu(), 4);
 }
 
+TEST(BlockLevelParametersTest,
+     EmptySequentialTileSizesAreNotSetInBlockLevelFusionConfig) {
+  BlockLevelParameters block_level_parameters;
+  block_level_parameters.output_tile_sizes = {{18, 19}};
+
+  BlockLevelFusionConfig block_level_fusion_config =
+      block_level_parameters.ToBlockLevelFusionConfig();
+
+  EXPECT_THAT(block_level_fusion_config.sequential_tile_sizes(), IsEmpty());
+  EXPECT_THAT(BlockLevelParameters::FromBlockLevelFusionConfig(
+                  block_level_fusion_config)
+                  .sequential_tile_sizes,
+              IsEmpty());
+}
+
 }  // namespace
 }  // namespace xla::xtile
diff --git a/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.cc b/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.cc
index bdddb6a2001..273214417fb 100644
--- a/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.cc
+++ b/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters.cc
@@ -176,9 +176,21 @@ absl::StatusOr<llvm::SmallVector<int64_t>> GetTilingSpaceConcreteSizes(
         "(%d).",
         num_parallel_dims, parallel_tile_sizes.size());
   }
+  const int64_t num_sequential_dims = tiling_space.num_sequential_dimensions();
+  const auto& sequential_tile_sizes =
+      block_level_parameters.sequential_tile_sizes;
+  if (!sequential_tile_sizes.empty() &&
+      sequential_tile_sizes.size() != num_sequential_dims) {
+    return Internal(
+        "Number of sequential dimensions in the tiling space (%d) does not "
+        "match the number of sequential tile sizes in the block level fusion "
+        "config (%d).",
+        num_sequential_dims, sequential_tile_sizes.size());
+  }
   llvm::SmallVector<int64_t> tile_sizes;
   tile_sizes.reserve(tiling_space.dimensions().size());
   int parallel_dim_count = 0;
+  int sequential_dim_count = 0;
   for (const xla::gpu::experimental::TilingSpace::DimensionInfo& dim :
        tiling_space.dimensions()) {
     switch (dim.type) {
@@ -189,6 +201,18 @@ absl::StatusOr<llvm::SmallVector<int64_t>> GetTilingSpaceConcreteSizes(
       case DimensionSemantics::kReplica:
         break;
       case DimensionSemantics::kSequential: {
+        if (!sequential_tile_sizes.empty()) {
+          if (dim.hlo->has_backend_config()) {
+            return Internal(
+                "Instruction %s has a backend config while "
+                "sequential_tile_sizes is set in the block level fusion "
+                "config.",
+                dim.hlo->name());
+          }
+          tile_sizes.push_back(sequential_tile_sizes[sequential_dim_count]);
+          sequential_dim_count++;
+          break;
+        }
         if (dim.hlo->has_backend_config()) {
           ABSL_ASSIGN_OR_RETURN(xla::xtile::Tile config,
                            dim.hlo->backend_config<xla::xtile::Tile>());
@@ -197,23 +221,23 @@ absl::StatusOr<llvm::SmallVector<int64_t>> GetTilingSpaceConcreteSizes(
           // in config.sizes is offset by output_rank. For scans, the sequential
           // dimension is in the output and config.sizes contains only the scan
           // dimension at index 0.
-          int64_t seq_dim_idx = 0;
+          int64_t config_dim_idx = 0;
           if (dim.hlo->opcode() != HloOpcode::kScan) {
             int64_t output_rank = dim.hlo->shape().dimensions().size();
-            seq_dim_idx = dim.dim_position - output_rank;
+            config_dim_idx = dim.dim_position - output_rank;
           }
-          if (seq_dim_idx < 0 || seq_dim_idx >= config.sizes_size()) {
+          if (config_dim_idx < 0 || config_dim_idx >= config.sizes_size()) {
             return Internal(
                 "Sequential dimension index %d is out of bounds for backend "
                 "config sizes of size %d.",
-                seq_dim_idx, config.sizes_size());
+                config_dim_idx, config.sizes_size());
           }
-          tile_sizes.push_back(config.sizes(seq_dim_idx));
-        } else {
-          VLOG(1) << "No backend_config set for HLO instruction of dimension "
-                  << dim.ToString() << ". Using dimension size as tile size.";
-          tile_sizes.push_back(dim.dimension_size);
+          tile_sizes.push_back(config.sizes(config_dim_idx));
+          break;
         }
+        VLOG(1) << "No backend_config set for HLO instruction of dimension "
+                << dim.ToString() << ". Using dimension size as tile size.";
+        tile_sizes.push_back(dim.dimension_size);
         break;
       }
       default:
diff --git a/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters_test.cc b/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters_test.cc
index ff6fab791a5..9c6f6f43831 100644
--- a/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters_test.cc
+++ b/third_party/xla/xla/codegen/xtile/tiling_from_block_parameters_test.cc
@@ -163,6 +163,104 @@ ENTRY entry {
   EXPECT_THAT(tile_sizes, ElementsAre(16, 128, 256));
 }
 
+TEST_F(GetTileTilingSpaceConcreteSizesTest,
+       MultipleReductionsWithSequentialTileSizes) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(R"hlo(
+add {
+  lhs = f32[] parameter(0)
+  rhs = f32[] parameter(1)
+  ROOT add = f32[] add(lhs, rhs)
+}
+
+f {
+  p0 = f32[64,128,256,512] parameter(0)
+  c0 = f32[] constant(0)
+  r0 = f32[64,128] reduce(p0, c0), dimensions={2,3}, to_apply=add
+  ROOT r1 = f32[64] reduce(r0, c0), dimensions={1}, to_apply=add
+}
+
+ENTRY entry {
+  param_0 = f32[64,128,256,512] parameter(0)
+  ROOT fusion = f32[64] fusion(param_0),
+    kind=kInput, calls=f,
+    backend_config={fusion_backend_config:{block_level_fusion_config:{
+      output_tiles:[{sizes:[16]}],
+      sequential_tile_sizes:[32, 64, 128]
+    }}}
+}
+)hlo"));
+  const HloInstruction* root = module->entry_computation()->root_instruction();
+  ASSERT_OK_AND_ASSIGN(llvm::SmallVector<int64_t> tile_sizes,
+                       ComputeConcreteTileSizesOfFusion(root));
+  // TilingSpace visits instructions in reverse post-order (use-before-def),
+  // so r1's sequential dimension comes before r0's two sequential dimensions.
+  EXPECT_THAT(tile_sizes, ElementsAre(16, 32, 64, 128));
+}
+
+TEST_F(GetTileTilingSpaceConcreteSizesTest,
+       FailsWhenSequentialTileSizesDoNotMatchTilingSpace) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(R"hlo(
+add {
+  lhs = f32[] parameter(0)
+  rhs = f32[] parameter(1)
+  ROOT add = f32[] add(lhs, rhs)
+}
+
+f {
+  p0 = f32[64,128,256] parameter(0)
+  c0 = f32[] constant(0)
+  ROOT reduce = f32[64] reduce(p0, c0), dimensions={1,2}, to_apply=add
+}
+
+ENTRY entry {
+  param_0 = f32[64,128,256] parameter(0)
+  ROOT fusion = f32[64] fusion(param_0),
+    kind=kInput, calls=f,
+    backend_config={fusion_backend_config:{block_level_fusion_config:{
+      output_tiles:[{sizes:[16]}],
+      sequential_tile_sizes:[32]
+    }}}
+}
+)hlo"));
+  const HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(ComputeConcreteTileSizesOfFusion(root),
+              StatusIs(absl::StatusCode::kInternal));
+}
+
+TEST_F(GetTileTilingSpaceConcreteSizesTest,
+       FailsWhenBothSequentialTileSizesAndInstructionBackendConfigAreSet) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(R"hlo(
+add {
+  lhs = f32[] parameter(0)
+  rhs = f32[] parameter(1)
+  ROOT add = f32[] add(lhs, rhs)
+}
+
+f {
+  p0 = f32[64,128] parameter(0)
+  c0 = f32[] constant(0)
+  ROOT reduce = f32[64] reduce(p0, c0), dimensions={1}, to_apply=add,
+    backend_config={sizes:[32]}
+}
+
+ENTRY entry {
+  param_0 = f32[64,128] parameter(0)
+  ROOT fusion = f32[64] fusion(param_0),
+    kind=kInput, calls=f,
+    backend_config={fusion_backend_config:{block_level_fusion_config:{
+      output_tiles:[{sizes:[16]}],
+      sequential_tile_sizes:[32]
+    }}}
+}
+)hlo"));
+  const HloInstruction* root = module->entry_computation()->root_instruction();
+  EXPECT_THAT(ComputeConcreteTileSizesOfFusion(root),
+              StatusIs(absl::StatusCode::kInternal));
+}
+
 TEST_F(GetTileTilingSpaceConcreteSizesTest, FailsOnMultiOutputFusions) {
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"hlo(
diff --git a/third_party/xla/xla/codegen/xtile/xtile_config.proto b/third_party/xla/xla/codegen/xtile/xtile_config.proto
index 33a18703861..8242c70bbde 100644
--- a/third_party/xla/xla/codegen/xtile/xtile_config.proto
+++ b/third_party/xla/xla/codegen/xtile/xtile_config.proto
@@ -67,4 +67,10 @@ message BlockLevelFusionConfig {
   // reordering (grouping) in ragged-dot fusions. A value <= 1 disables the
   // reordering.
   int32 group_size = 10;
+
+  // Tile sizes of the sequential dimensions in the fusion, in TilingSpace
+  // order. When empty, sequential tile sizes are not set in the fusion config,
+  // and each sequential dimension uses the `Tile` backend config of its
+  // instruction or its full size.
+  repeated int64 sequential_tile_sizes = 11;
 }
diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index 9f7a7708997..bc01ecf1dfc 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -3907,6 +3907,14 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
           &DebugOptions::set_xla_gpu_experimental_enable_tiling_propagation),
       debug_options->xla_gpu_experimental_enable_tiling_propagation(),
       "If true, enable experimental tiling propagation."));
+  flag_list->push_back(tsl::Flag(
+      "xla_gpu_experimental_enable_reduction_dimension_tiling",
+      bool_setter_for(
+          &DebugOptions::
+              set_xla_gpu_experimental_enable_reduction_dimension_tiling),
+      debug_options->xla_gpu_experimental_enable_reduction_dimension_tiling(),
+      "If true, the block-level (Triton) tiling search also enumerates tile "
+      "sizes for reduction dimensions."));
 
   auto setter_for_xla_gpu_detect_nan =
       [debug_options, detection_mode](const std::string& value) {
diff --git a/third_party/xla/xla/service/gpu/model/BUILD b/third_party/xla/xla/service/gpu/model/BUILD
index 4897fb4064c..8f2a3a9104d 100644
--- a/third_party/xla/xla/service/gpu/model/BUILD
+++ b/third_party/xla/xla/service/gpu/model/BUILD
@@ -579,6 +579,7 @@ cc_library(
         ":triton_emitter_constraints",
         "//xla:shape_util",
         "//xla:util",
+        "//xla:xla_proto_cc",
         "//xla/backends/gpu/codegen/triton:fusion",
         "//xla/codegen/tiling:symbolic_tile_analysis",
         "//xla/codegen/tiling:tiled_hlo_computation",
diff --git a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc
index a15bd65296f..03b3bcda8d4 100644
--- a/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc
+++ b/third_party/xla/xla/service/gpu/model/gpu_indexing_performance_model.cc
@@ -79,6 +79,7 @@ limitations under the License.
 #include "xla/tsl/concurrency/future.h"
 #include "xla/tsl/platform/env.h"
 #include "xla/util.h"
+#include "xla/xla.pb.h"
 #include "tsl/platform/path.h"
 
 namespace xla {
@@ -669,6 +670,45 @@ int64_t EstimateNumWarpsImpl(
   return GetNumWarps(largest_live_tile_size);
 }
 
+// SymbolicTileAnalysis does not tile sequential reduction/scan dimensions.
+std::vector<int64_t> GetSequentialTileSizes(
+    const TiledHloComputation& tiled_hlo_computation) {
+  return {};
+}
+
+// TODO(b/534676634): Support multi-reduction and multi-dimension reduction
+// fusions. When supporting multi-reduction fusions, `sequential_tile_sizes`
+// must be ordered in the reverse post-order that the merged fusion computation
+// will have after `MergeFusionInstruction` (e.g., via a post-order DFS from
+// `fusion_adaptor.GetRoots()` following
+// `HloInstructionAdaptor::GetOperands()`), because
+// `HloFusionAdaptor::MakeInstructionPostOrder()` on a producer-consumer adaptor
+// concatenates the producer's post-order before the consumer's post-order,
+// which can differ from the post-order of the merged computation.
+std::vector<int64_t> GetSequentialTileSizes(
+    const experimental::TiledHloComputation& tiled_hlo_computation) {
+  const experimental::TilingSpace& tiling_space =
+      tiled_hlo_computation.tiling_space();
+  std::vector<int64_t> sequential_tile_sizes;
+  bool all_dimensions_untiled = true;
+  for (const auto& dim : tiling_space.dimensions()) {
+    if (dim.type ==
+            experimental::TilingSpace::DimensionSemantics::kSequential &&
+        dim.hlo->opcode() == HloOpcode::kReduce) {
+      CHECK(dim.tile_size.has_value()) << dim.ToString();
+      sequential_tile_sizes.push_back(*dim.tile_size);
+      all_dimensions_untiled &= *dim.tile_size >= dim.dimension_size;
+    }
+  }
+  // An empty result means that the sequential dimensions are not tiled. Keep it
+  // empty when every sequential dimension is covered by a single tile, so that
+  // the config of such tilings is the same as before reduction tiling.
+  if (all_dimensions_untiled || sequential_tile_sizes.size() != 1) {
+    return {};
+  }
+  return sequential_tile_sizes;
+}
+
 template <typename TiledHloComputationType>
 absl::StatusOr<std::optional<TiledRunTimeData>> EstimateTiledRunTimeDataImpl(
     const HloFusionAdaptor& fusion_adaptor,
@@ -685,6 +725,8 @@ absl::StatusOr<std::optional<TiledRunTimeData>> EstimateTiledRunTimeDataImpl(
     block_level_parameters.output_tile_sizes.emplace_back(tile_sizes.begin(),
                                                           tile_sizes.end());
   }
+  block_level_parameters.sequential_tile_sizes =
+      GetSequentialTileSizes(tiled_hlo_computation);
   block_level_parameters.num_warps =
       EstimateNumWarpsImpl(tiled_hlo_computation);
 
@@ -702,6 +744,17 @@ absl::StatusOr<std::optional<TiledRunTimeData>> EstimateTiledRunTimeDataImpl(
   return TiledRunTimeData{estimate_run_time_data, block_level_parameters};
 }
 
+// Whether the tiling search for `fusion_adaptor` should enumerate tile sizes
+// for the reduction dimensions, i.e. consider emitting reductions as loops.
+bool ShouldTileReductionDimensions(const HloFusionAdaptor& fusion_adaptor) {
+  const HloModule* module =
+      fusion_adaptor.GetRoots().front().instruction().GetModule();
+  return module != nullptr &&
+         module->config()
+             .debug_options()
+             .xla_gpu_experimental_enable_reduction_dimension_tiling();
+}
+
 absl::Status DumpAndLogTopKCandidates(
     const HloFusionAdaptor& fusion_adaptor,
     absl::Span<const TiledRunTimeData> candidates) {
@@ -857,7 +910,10 @@ tsl::Future<TopKTiledRunTimeDataOrError> TryFindTopKBestTilingsWithTilingSpace(
       base->tiling_space,
       experimental::TilingSpace::Create(fusion_adaptor, base->context.get()));
 
-  ABSL_ASSIGN_OR_RETURN(auto tilings, base->tiling_space->GetValidTilings());
+  ABSL_ASSIGN_OR_RETURN(auto tilings,
+                   base->tiling_space->GetValidTilings(
+                       /*tile_reduction_dimensions=*/
+                       ShouldTileReductionDimensions(fusion_adaptor)));
   VLOG(1) << absl::StrCat(
       "TryFindTopKBestTilingsForFusionAsync tiling_space evaluating ",
       tilings.size(), " tilings.");
diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index b0e67029819..8b6e735625f 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -1952,7 +1952,14 @@ message DebugOptions {
   // Requires sm_90+ and CUDA 13.3+ (tileiras) at runtime. Off by default.
   optional bool xla_gpu_experimental_enable_tensor_ir = 557;
 
-  // Next id: 558
+  // If true, the block-level (Triton) tiling search also enumerates tile sizes
+  // for reduction dimensions, i.e. reductions may be emitted as a loop over
+  // tiles.
+  // TODO(b/534676634): Only fusions with a single reduce over a single
+  // dimension are tiled for now; extend to the other cases.
+  optional bool xla_gpu_experimental_enable_reduction_dimension_tiling = 558;
+
+  // Next id: 559
 
   // Extra options to pass to the compilation backend (e.g. LLVM); specific
   // interpretation of these values is left to the backend.

From a64db63c0e41315e76ee8fb42345f19f06df2ff9 Mon Sep 17 00:00:00 2001
From: bodapatisaikrishna <saikrishnabodapati@gmail.com>
Date: Thu, 15 Oct 2026 22:26:47 +0530
Subject: [PATCH 0875/1000] Disable XLA auto-clustering in
 testResourceVariableDtypeMismatch

---
 tensorflow/python/training/training_ops_test.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow/python/training/training_ops_test.py b/tensorflow/python/training/training_ops_test.py
index 48329d9bc43..f22396fbaf2 100644
--- a/tensorflow/python/training/training_ops_test.py
+++ b/tensorflow/python/training/training_ops_test.py
@@ -698,6 +698,7 @@ class TrainingOpsTest(TensorFlowTestCase):
         )
 
   @test_util.run_in_graph_and_eager_modes
+  @test_util.disable_xla("Tests C++ kernel validation in training_op_helpers.h")
   def testResourceVariableDtypeMismatch(self):
     # A resource variable whose dtype differs from optimizer dtype T
     # previously caused a fatal CHECK failure in CheckTypeAndIsAligned

From 11a47256df68a4912946280d7aa2f63afc89a890 Mon Sep 17 00:00:00 2001
From: Vladimir Belitskiy <belitskiy@google.com>
Date: Thu, 15 Oct 2026 08:49:55 -0700
Subject: [PATCH 0876/1000] Clear moved-from `flat_hash_map`s in
 `PullTable::Reset()` to fix ASAN crash in `streaming_test`.

In `PullTable::Reset()` (`xla/python/transfer/streaming.cc`), `entries_` and `paused_fetches_` are moved into local variables while holding `mu_`, leaving the member `absl::flat_hash_map` instances in a moved-from state when `mu_` is unlocked.

When `PullTable::Handle()` calls `entry->Handle(...)` without holding `mu_` and the entry callback triggers `table_->Reset()` (as in `PullTable.PullTableRace` in `//xla/python/transfer:streaming_test`), `PullTable::Handle()` subsequently re-locks `mu_` and calls `entries_.find(req.uuid())` on the moved-from `entries_` table, triggering an AddressSanitizer / Abseil raw_hash_set abort (`[raw_hash_set.h] RAW: Use of moved-from hash table.`).

Explicitly call `entries_.clear()` and `paused_fetches_.clear()` before unlocking `mu_` in `PullTable::Reset()` so both member maps are reinitialized to valid empty hash tables.

PiperOrigin-RevId: 1000167680
---
 third_party/xla/xla/python/transfer/streaming.cc | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/third_party/xla/xla/python/transfer/streaming.cc b/third_party/xla/xla/python/transfer/streaming.cc
index f73e1c55fdf..e9ead1cd784 100644
--- a/third_party/xla/xla/python/transfer/streaming.cc
+++ b/third_party/xla/xla/python/transfer/streaming.cc
@@ -506,6 +506,8 @@ void PullTable::Reset() {
   mu_.lock();
   auto entries = std::move(entries_);
   auto paused_fetches_by_uuid = std::move(paused_fetches_);
+  entries_.clear();
+  paused_fetches_.clear();
   entry_heap_.clear();
   fetch_heap_.clear();
   mu_.unlock();

From a4a8e1c2e9ab3e5085bd1296362f00e101f682ef Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 15 Oct 2026 09:16:35 -0700
Subject: [PATCH 0877/1000] [tflite] Constant-fold PAD with a constant input.

`PadOp::fold` only handled the no-op case (input and output shapes equal), so a
`PAD` whose input and paddings are both constant survived into the flatbuffer.
Delegates that require a runtime data input reject such a node; ML Drift
reports `PAD: inputs[0] should not be a constant tensor.`

This shows up in Swin: torchvision's `shifted_window_attention` builds the
shifted-window attention mask from static shapes, and its decomposition starts
from a constant `PAD` in each shifted stage. Fold a constant input padded by
constant amounts (with zeros, as `PAD` does), subject to the usual
`ShouldFoldOperation` size limit. Quantized element types are left unfolded.

PiperOrigin-RevId: 1000183464
---
 tensorflow/compiler/mlir/lite/ir/tfl_ops.cc   | 48 ++++++++++++++++++-
 .../compiler/mlir/lite/tests/const-fold.mlir  | 31 ++++++++++++
 2 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc b/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc
index b9bc4fcd47e..4cd59932c74 100644
--- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc
+++ b/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc
@@ -6018,11 +6018,55 @@ int64_t L2NormalizationOp::GetArithmeticCount(Operation* op) {
 // PadOp
 //===----------------------------------------------------------------------===//
 
-OpFoldResult PadOp::fold(FoldAdaptor) {
+OpFoldResult PadOp::fold(FoldAdaptor adaptor) {
   if (InputOutputHasSameShape(getInput().getType(), getOutput().getType()))
     return getInput();
 
-  return {};
+  // Fold a constant input padded by constant amounts. Left unfolded, a
+  // constant-input PAD survives into the flatbuffer, and delegates that
+  // require a runtime input (e.g. ML Drift) reject it.
+  if (!ShouldFoldOperation(this->getOperation())) return {};
+
+  auto input = dyn_cast_or_null<DenseElementsAttr>(adaptor.getInput());
+  if (!input) return {};
+  auto output_type = mlir::dyn_cast<RankedTensorType>(getOutput().getType());
+  if (!output_type || !output_type.hasStaticShape()) return {};
+  // DenseElementsAttr does not support quantized element types.
+  Type element_type = output_type.getElementType();
+  if (!element_type.isSignlessIntOrFloat()) return {};
+
+  auto input_type = mlir::cast<ShapedType>(input.getType());
+  const int64_t rank = input_type.getRank();
+  const llvm::SmallVector<int64_t> paddings =
+      UnpackIndexOperand(adaptor.getPadding());
+  if (paddings.size() != 2 * rank) return {};
+  for (int64_t d = 0; d < rank; ++d) {
+    const int64_t before = paddings[2 * d], after = paddings[2 * d + 1];
+    if (before < 0 || after < 0 ||
+        input_type.getDimSize(d) + before + after !=
+            output_type.getDimSize(d)) {
+      return {};
+    }
+  }
+
+  // TFL PAD fills with zero (PADV2 takes an explicit value).
+  Attribute zero = mlir::isa<FloatType>(element_type)
+                       ? Attribute(FloatAttr::get(element_type, 0.0))
+                       : Attribute(IntegerAttr::get(element_type, 0));
+  std::vector<Attribute> result(output_type.getNumElements(), zero);
+
+  const FlatIndHelper read_inds(input_type);
+  const FlatIndHelper write_inds(output_type);
+  llvm::SmallVector<int64_t> offset(rank);
+  for (int64_t d = 0; d < rank; ++d) offset[d] = paddings[2 * d];
+
+  int64_t read_flat_ind = 0;
+  for (Attribute value : input.getValues<Attribute>()) {
+    llvm::SmallVector<int64_t> ind = read_inds.GetShapedInd(read_flat_ind++);
+    FlatIndHelper::AddOffset(ind, offset);
+    result[write_inds.GetFlatInd(ind)] = value;
+  }
+  return DenseElementsAttr::get(output_type, result);
 }
 
 // When padding amounts are constants, cast them to i32. XNN can only
diff --git a/tensorflow/compiler/mlir/lite/tests/const-fold.mlir b/tensorflow/compiler/mlir/lite/tests/const-fold.mlir
index e4b93ca18db..dcdda38f84b 100644
--- a/tensorflow/compiler/mlir/lite/tests/const-fold.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/const-fold.mlir
@@ -537,6 +537,37 @@ func.func @ConstFoldPad(%arg0: tensor<15600xf32>) -> tensor<15600xf32> {
   // CHECK:  return %arg0
 }
 
+// CHECK-LABEL: @ConstFoldPadConstInput
+func.func @ConstFoldPadConstInput() -> tensor<3x4xf32> {
+  %0 = "tfl.pseudo_const"() {value = dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xf32>} : () -> tensor<2x2xf32>
+  %1 = "tfl.pseudo_const"() {value = dense<[[1, 0], [0, 2]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
+  %2 = "tfl.pad"(%0, %1) : (tensor<2x2xf32>, tensor<2x2xi32>) -> tensor<3x4xf32>
+  func.return %2 : tensor<3x4xf32>
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[}}0.000000e+00, 0.000000e+00, 0.000000e+00, 0.000000e+00], [1.000000e+00, 2.000000e+00, 0.000000e+00, 0.000000e+00], [3.000000e+00, 4.000000e+00, 0.000000e+00, 0.000000e+00]]> : tensor<3x4xf32>
+  // CHECK-NOT: tfl.pad
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstFoldPadConstInputI64Paddings
+func.func @ConstFoldPadConstInputI64Paddings() -> tensor<4xi32> {
+  %0 = "tfl.pseudo_const"() {value = dense<[7, 8]> : tensor<2xi32>} : () -> tensor<2xi32>
+  %1 = "tfl.pseudo_const"() {value = dense<[[1, 1]]> : tensor<1x2xi64>} : () -> tensor<1x2xi64>
+  %2 = "tfl.pad"(%0, %1) : (tensor<2xi32>, tensor<1x2xi64>) -> tensor<4xi32>
+  func.return %2 : tensor<4xi32>
+  // CHECK: %[[CST:.*]] = arith.constant dense<[0, 7, 8, 0]> : tensor<4xi32>
+  // CHECK-NOT: tfl.pad
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @NotFoldPadRuntimeInput
+func.func @NotFoldPadRuntimeInput(%arg0: tensor<2x2xf32>) -> tensor<3x4xf32> {
+  %0 = "tfl.pseudo_const"() {value = dense<[[1, 0], [0, 2]]> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
+  %1 = "tfl.pad"(%arg0, %0) : (tensor<2x2xf32>, tensor<2x2xi32>) -> tensor<3x4xf32>
+  func.return %1 : tensor<3x4xf32>
+  // CHECK: %[[PAD:.*]] = "tfl.pad"(%arg0
+  // CHECK: return %[[PAD]]
+}
+
 func.func @ConstFoldPadV2(%arg0: tensor<15600xf32>) -> tensor<15600xf32> {
   %0 = "tfl.pseudo_const"() {value = dense<0> : tensor<1x2xi32>} : () -> tensor<1x2xi32>
   %1 = "tfl.pseudo_const"() {value = dense<0.0> : tensor<f32>} : () -> tensor<f32>

From e961b4d62b936d23367589fe3dd990d747a63b1d Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Thu, 15 Oct 2026 09:45:40 -0700
Subject: [PATCH 0878/1000] [XLA] Do not simplify integer pow(A, -1) to 1 / A
 in AlgebraicSimplifier

`AlgebraicSimplifierVisitor::HandlePower` rewrites `power(A, -1)` into `divide(1, A)` whenever `IsAll(rhs, -1)` holds. `LiteralBase::IsAll(int8_t{-1})` returns true for signed integer constants `S8`, `S16`, `S32`, and `S64`, so this rewrite fires on signed integer power instructions.

That rewrite is unsound for integers at `A = 0`. In XLA HLO semantics, signed integer `pow(0, -1)` evaluates to `0`, whereas signed integer `1 / 0` evaluates to `-1`. In `divide(X, power(A, -1))` at `A = 0`, `X / pow(0, -1)` evaluates to `X / 0 = -1`. Rewriting the denominator to `1 / A` instead produces `X / (1 / 0) = X / -1 = -X`.

Guard the `pow(A, -1)` to `1 / A` rewrite in `HandlePower` with `!ShapeUtil::ElementIsIntegral(power->shape())` so that integer `pow(A, -1)` is preserved. Floating point and complex `pow(A, -1)` continue to simplify to `1 / A`.

Tests:
- `//third_party/tensorflow/compiler/xla/hlo/transforms/simplifiers:algebraic_simplifier_test`: `IntegerPowNegative1NotSimplified` verifies that `s32[] power(a, -1)` is not rewritten to division.
PiperOrigin-RevId: 1000200286
---
 .../simplifiers/algebraic_simplifier.cc       |  2 +-
 .../simplifiers/algebraic_simplifier_test.cc  | 19 +++++++++++++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
index e66f267ad7e..37d4350d67e 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier.cc
@@ -6302,7 +6302,7 @@ absl::Status AlgebraicSimplifierVisitor::HandlePower(HloInstruction* power) {
   }
 
   VLOG(10) << "trying transform [pow(A, -1) => 1/A]: " << power->ToString();
-  if (IsAll(rhs, -1)) {
+  if (!ShapeUtil::ElementIsIntegral(power->shape()) && IsAll(rhs, -1)) {
     return ReplaceWithNewInstruction(
         power, HloInstruction::CreateBinary(power->shape(), HloOpcode::kDivide,
                                             MakeScalarLike(lhs, 1), lhs));
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
index 80b9590d0d8..6e494140a32 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/algebraic_simplifier_test.cc
@@ -3529,6 +3529,25 @@ TEST_F(AlgebraicSimplifierTest, PowNegative1) {
   EXPECT_EQ(root->operand(0)->literal().GetFirstElement<float>(), 1);
 }
 
+// Test that an integer pow(A, -1) is not simplified to 1/A because integer
+// pow(0, -1) is 0 while integer 1/0 is -1.
+TEST_F(AlgebraicSimplifierTest, IntegerPowNegative1NotSimplified) {
+  constexpr absl::string_view kModuleStr = R"(
+    HloModule m
+    test {
+      a = s32[] parameter(0)
+      minus_one = s32[] constant(-1)
+      ROOT p = s32[] power(a, minus_one)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(kModuleStr));
+  ASSERT_OK_AND_ASSIGN(bool changed,
+                       AlgebraicSimplifier(default_options_).Run(m.get()));
+  EXPECT_FALSE(changed);
+  EXPECT_THAT(m->entry_computation()->root_instruction(),
+              GmockMatch(m::Power(m::Parameter(0), m::ConstantScalar(-1))));
+}
+
 // pow(A, 0.5) => sqrt(A), for A >= 0
 TEST_F(AlgebraicSimplifierTest, PowHalf) {
   constexpr absl::string_view kModuleStr = R"(

From ed7d056a96de22e2692aab7cad0dac071fe3aa17 Mon Sep 17 00:00:00 2001
From: Vladimir Belitskiy <belitskiy@google.com>
Date: Thu, 15 Oct 2026 09:51:46 -0700
Subject: [PATCH 0879/1000] [CI] Add BCR mirror fallbacks and Windows JVM DNS
 caching to XLA CI.

Bazel wraps `java.net.UnknownHostException` in `UnrecoverableHttpException` (`HttpConnector.java`), which skips per-URL retries in `HttpConnectorMultiplexer` and only falls back if alternative URLs are provided via `--experimental_downloader_config`.

Add `bcr.bazel.build` rewrite rules (`storage.googleapis.com/bcr.bazel.build` and `raw.githubusercontent.com/bazelbuild/bazel-central-registry`) to `bazel_downloader.cfg` and switch `--experimental_downloader_config=bazel_downloader.cfg` to `common` in `tensorflow.bazelrc` so `cquery` also uses the downloader config.

Also add `build_tools/ci/bazel_dns_windows.security` (`networkaddress.cache.ttl=60`, `networkaddress.cache.negative.ttl=1`) and pass it via `--host_jvm_args` on `XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS` and `JAX_WINDOWS_X86_CPU_GITHUB_ACTIONS` in `build_tools/ci/build.py`.

PiperOrigin-RevId: 1000204208
---
 third_party/xla/bazel_downloader.cfg          |  4 +++
 .../build_tools/ci/bazel_dns_windows.security | 33 +++++++++++++++++++
 third_party/xla/build_tools/ci/build.py       |  6 ++++
 third_party/xla/build_tools/ci/build_test.py  | 11 +++++++
 .../xla/build_tools/ci/golden_commands.txt    |  4 +--
 third_party/xla/tensorflow.bazelrc            |  8 ++---
 6 files changed, 58 insertions(+), 8 deletions(-)
 create mode 100644 third_party/xla/build_tools/ci/bazel_dns_windows.security

diff --git a/third_party/xla/bazel_downloader.cfg b/third_party/xla/bazel_downloader.cfg
index 9ce7de4bfef..6fa4703284c 100644
--- a/third_party/xla/bazel_downloader.cfg
+++ b/third_party/xla/bazel_downloader.cfg
@@ -20,3 +20,7 @@ rewrite (github.com)/(.*) https://$1/$2
 rewrite (files\.pythonhosted\.org)/(.*) https://mirror.bazel.build/$1/$2
 rewrite (files\.pythonhosted\.org)/(.*) https://storage.googleapis.com/mirror.tensorflow.org/$1/$2
 rewrite (files\.pythonhosted\.org)/(.*) https://$1/$2
+
+rewrite (bcr\.bazel\.build)/(.*) https://$1/$2
+rewrite (bcr\.bazel\.build)/(.*) https://storage.googleapis.com/$1/$2
+rewrite (bcr\.bazel\.build)/(.*) https://raw.githubusercontent.com/bazelbuild/bazel-central-registry/main/$2
diff --git a/third_party/xla/build_tools/ci/bazel_dns_windows.security b/third_party/xla/build_tools/ci/bazel_dns_windows.security
new file mode 100644
index 00000000000..51c503b2cbe
--- /dev/null
+++ b/third_party/xla/build_tools/ci/bazel_dns_windows.security
@@ -0,0 +1,33 @@
+# Copyright 2026 The OpenXLA Authors.
+#
+# 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.
+# ==============================================================================
+
+# --- Rationale ---
+#
+# Bazel, by default, retries transient remote-execution failures five times
+# using exponential backoff, at roughly 100ms, 200ms, 400ms, 800ms, and 1.6s.
+#
+# Java normally caches a failed DNS lookup for 10 seconds, causing all of those
+# retries to receive the same cached failure.
+#
+# To work around this:
+# - Keep successful DNS lookups longer to bridge occasional, brief resolver
+#   outages.
+# - Allow Bazel's remote retries to retry failed lookups much quicker
+#   instead of using the JVM's default ten-second negative cache.
+
+# Up from 30 seconds
+networkaddress.cache.ttl=60
+# Down from 10 seconds
+networkaddress.cache.negative.ttl=1
diff --git a/third_party/xla/build_tools/ci/build.py b/third_party/xla/build_tools/ci/build.py
index dc4310e2b1e..7af54c9130c 100755
--- a/third_party/xla/build_tools/ci/build.py
+++ b/third_party/xla/build_tools/ci/build.py
@@ -73,6 +73,10 @@ _XLA_GPU_PRESUBMIT_BENCHMARKS_DEFAULT_TARGET_PATTERNS = (
     "//xla/tools:compute_xspace_stats_main_gpu",
 )
 _GITHUB_WORKSPACE = os.environ.get("GITHUB_WORKSPACE", "$GITHUB_WORKSPACE")
+_WINDOWS_BAZEL_DNS_SECURITY_PROPERTIES = (
+    f"-Djava.security.properties={_GITHUB_WORKSPACE}"
+    r"\openxla\xla\build_tools\ci\bazel_dns_windows.security"
+)
 
 
 def retry(
@@ -486,6 +490,7 @@ Build(
     subcommand="build",
     startup_options={
         "output_user_root": "C:/x",
+        "host_jvm_args": _WINDOWS_BAZEL_DNS_SECURITY_PROPERTIES,
     },
     # Retry failed Bazel builds against the warm local cache to mitigate
     # transient Windows lld-link wcifs.sys file-open races (b/571438464) until
@@ -832,6 +837,7 @@ Build(
     subcommand="build",
     startup_options={
         "output_base": f"{_GITHUB_WORKSPACE}\\bazel_output_base",
+        "host_jvm_args": _WINDOWS_BAZEL_DNS_SECURITY_PROPERTIES,
     },
     # Retry failed Bazel builds against the warm local cache to mitigate
     # transient Windows lld-link wcifs.sys file-open races (b/571438464) until
diff --git a/third_party/xla/build_tools/ci/build_test.py b/third_party/xla/build_tools/ci/build_test.py
index ad4377b8fda..6974afc24e5 100644
--- a/third_party/xla/build_tools/ci/build_test.py
+++ b/third_party/xla/build_tools/ci/build_test.py
@@ -47,6 +47,17 @@ class BuildTest(parameterized.TestCase):
         0,
     )
 
+  def test_windows_builds_configure_dns_security_properties(self) -> None:
+    builds = build.Build.all_builds()
+    for build_type in (
+        build.BuildType.XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS,
+        build.BuildType.JAX_WINDOWS_X86_CPU_GITHUB_ACTIONS,
+    ):
+      self.assertIn(
+          r"\openxla\xla\build_tools\ci\bazel_dns_windows.security",
+          builds[build_type].startup_options.get("host_jvm_args", ""),
+      )
+
   @parameterized.named_parameters(
       dict(
           testcase_name="xla_windows_succeeds_on_first_attempt",
diff --git a/third_party/xla/build_tools/ci/golden_commands.txt b/third_party/xla/build_tools/ci/golden_commands.txt
index ec31f8d099c..65b208da349 100644
--- a/third_party/xla/build_tools/ci/golden_commands.txt
+++ b/third_party/xla/build_tools/ci/golden_commands.txt
@@ -7,7 +7,7 @@ parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_fi
 bazel test --build_tag_filters=-multiaccelerator --test_tag_filters=-multiaccelerator --config=rbe_linux_x86_64_cuda13 --test_env=JAX_SKIP_SLOW_TESTS=1 --test_env=TF_CPP_MIN_LOG_LEVEL=0 --test_env=JAX_EXCLUDE_TEST_TARGETS=PmapTest.testSizeOverflow --repo_env=HERMETIC_PYTHON_VERSION=3.12 --override_repository=xla=$GITHUB_WORKSPACE/openxla/xla --override_module=xla=$GITHUB_WORKSPACE/openxla/xla --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --@local_config_cuda//cuda:override_include_cuda_libs -- //tests:gpu_tests //tests:backend_independent_tests //tests/pallas:gpu_tests //tests/pallas:backend_independent_tests //jaxlib/tools:check_gpu_wheel_sources_test //jaxlib/tools:jax_cuda_plugin_wheel_size_test //jaxlib/tools:jax_cuda_pjrt_wheel_size_test //jaxlib/tools:jaxlib_wheel_size_test //:jax_wheel_size_test
 # END BuildType.JAX_LINUX_X86_GPU_T4_GITHUB_ACTIONS
 # BEGIN BuildType.JAX_WINDOWS_X86_CPU_GITHUB_ACTIONS
-bazel --output_base=$GITHUB_WORKSPACE\bazel_output_base build --build_tag_filters= --test_tag_filters= --config=rbe_windows_amd64 --test_env=JAX_NUM_GENERATED_CASES=25 --test_env=JAX_SKIP_SLOW_TESTS=1 --repo_env=HERMETIC_PYTHON_VERSION=3.12 --override_repository=xla=$GITHUB_WORKSPACE\openxla\xla --override_module=xla=$GITHUB_WORKSPACE\openxla\xla --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_runfile_links=False -- //tests:cpu_tests //tests:backend_independent_tests //jax/experimental/jax2tf/tests:jax2tf_test_cpu //tests/multiprocess:cpu_tests //jax/experimental/jax2tf/tests/multiprocess:cpu_tests //jaxlib/tools:check_cpu_wheel_sources_test //jaxlib/tools:jaxlib_wheel_size_test //:jax_wheel_size_test
+bazel --output_base=$GITHUB_WORKSPACE\bazel_output_base --host_jvm_args=-Djava.security.properties=$GITHUB_WORKSPACE\openxla\xla\build_tools\ci\bazel_dns_windows.security build --build_tag_filters= --test_tag_filters= --config=rbe_windows_amd64 --test_env=JAX_NUM_GENERATED_CASES=25 --test_env=JAX_SKIP_SLOW_TESTS=1 --repo_env=HERMETIC_PYTHON_VERSION=3.12 --override_repository=xla=$GITHUB_WORKSPACE\openxla\xla --override_module=xla=$GITHUB_WORKSPACE\openxla\xla --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_runfile_links=False -- //tests:cpu_tests //tests:backend_independent_tests //jax/experimental/jax2tf/tests:jax2tf_test_cpu //tests/multiprocess:cpu_tests //jax/experimental/jax2tf/tests/multiprocess:cpu_tests //jaxlib/tools:check_cpu_wheel_sources_test //jaxlib/tools:jaxlib_wheel_size_test //:jax_wheel_size_test
 # END BuildType.JAX_WINDOWS_X86_CPU_GITHUB_ACTIONS
 # BEGIN BuildType.XLA_LINUX_ARM64_CPU_48_VCPU_PRESUBMIT_GITHUB_ACTIONS
 parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --test_tag_filters=-no_oss,-gpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel,-not_run:arm --config=warnings --config=rbe_cross_compile_linux_arm64 --config=nonccl --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --build_tests_only=False --//xla/tsl:ci_build --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main //xla/tools:compute_xspace_stats_main
@@ -91,5 +91,5 @@ mkdir -p /Volumes/BuildData/bazel_output
 bazel test --build_tag_filters=-no_oss,-gpu,-no_mac,-mac_excluded,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_oss,-gpu,-no_mac,-mac_excluded,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=nonccl --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --macos_minimum_os=10.15 --test_tmpdir=/Volumes/BuildData/bazel_output --//xla/tsl:ci_build -- //xla/... -//xla/python_api/... -//xla/python/... -//xla/service/gpu/...
 # END BuildType.XLA_MACOS_X86_CPU_KOKORO
 # BEGIN BuildType.XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS
-bazel --output_user_root=C:/x build --build_tag_filters=-no_windows,-no_oss,-gpu,-tpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_windows,-no_oss,-gpu,-tpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_xla_windows_x86_cpu_2022 --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/... -//xla/tpu/... -//xla/backends/cpu/collectives:gloo_collectives_test -//xla/backends/cpu/collectives:mpi_collectives -//xla/backends/cpu/collectives:mpi_communicator -//xla/python/transfer/... -//xla/backends/profiler/subprocess:subprocess_profiling_session -//xla/backends/profiler/subprocess:subprocess_profiling_session_test -//xla/backends/profiler/subprocess:subprocess_registry -//xla/backends/profiler/subprocess:subprocess_registry_test -//xla/backends/cpu/runtime/ynnpack:ynn_fusion_thunk -//xla/backends/cpu/runtime/ynnpack:ynn_interop -//xla/backends/cpu/runtime/ynnpack:ynn_threadpool -//xla/backends/gpu/... -//xla/codegen/emitters/tests/... -//xla/service/gpu/... -//xla/codegen/intrinsic/cpp:eigen_unary_test
+bazel --output_user_root=C:/x --host_jvm_args=-Djava.security.properties=$GITHUB_WORKSPACE\openxla\xla\build_tools\ci\bazel_dns_windows.security build --build_tag_filters=-no_windows,-no_oss,-gpu,-tpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --test_tag_filters=-no_windows,-no_oss,-gpu,-tpu,-requires-gpu-nvidia,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=nonccl --config=rbe_xla_windows_x86_cpu_2022 --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build -- //xla/... //build_tools/... @tsl//tsl/... -//xla/tpu/... -//xla/backends/cpu/collectives:gloo_collectives_test -//xla/backends/cpu/collectives:mpi_collectives -//xla/backends/cpu/collectives:mpi_communicator -//xla/python/transfer/... -//xla/backends/profiler/subprocess:subprocess_profiling_session -//xla/backends/profiler/subprocess:subprocess_profiling_session_test -//xla/backends/profiler/subprocess:subprocess_registry -//xla/backends/profiler/subprocess:subprocess_registry_test -//xla/backends/cpu/runtime/ynnpack:ynn_fusion_thunk -//xla/backends/cpu/runtime/ynnpack:ynn_interop -//xla/backends/cpu/runtime/ynnpack:ynn_threadpool -//xla/backends/gpu/... -//xla/codegen/emitters/tests/... -//xla/service/gpu/... -//xla/codegen/intrinsic/cpp:eigen_unary_test
 # END BuildType.XLA_WINDOWS_X86_CPU_GITHUB_ACTIONS
diff --git a/third_party/xla/tensorflow.bazelrc b/third_party/xla/tensorflow.bazelrc
index 844a51f0096..e6f2b794e9e 100644
--- a/third_party/xla/tensorflow.bazelrc
+++ b/third_party/xla/tensorflow.bazelrc
@@ -81,13 +81,9 @@ common --@rules_python//python/config_settings:bootstrap_impl=script --repo_env=
 
 common --incompatible_default_to_explicit_init_py
 
-# Use a downloader config to mirror GitHub and Python package downloads.
+# Use a downloader config to mirror GitHub, Python package, and BCR downloads.
 # We use the experimental_ name for better compatibility with different Bazel versions.
-build --experimental_downloader_config=bazel_downloader.cfg
-test --experimental_downloader_config=bazel_downloader.cfg
-fetch --experimental_downloader_config=bazel_downloader.cfg
-sync --experimental_downloader_config=bazel_downloader.cfg
-query --experimental_downloader_config=bazel_downloader.cfg
+common --experimental_downloader_config=bazel_downloader.cfg
 
 # For local build, without dependency on rbe
 common:linux_cpu --config=avx_linux

From 7b7e9ae848611e0d4d777c566c6e4e83a6cb09fb Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 15 Oct 2026 09:58:40 -0700
Subject: [PATCH 0880/1000] Integrate LLVM at llvm/llvm-project@1ee8805f01a6

Updates LLVM usage to match
[1ee8805f01a6](https://github.com/llvm/llvm-project/commit/1ee8805f01a663b7f2b09215a8eab5ab2d15ce35)

PiperOrigin-RevId: 1000208619
---
 .../xla/third_party/llvm/workspace.bzl        |   4 +-
 .../xla/third_party/shardy/temporary.patch    | 198 ++++++++++++++++--
 2 files changed, 177 insertions(+), 25 deletions(-)

diff --git a/third_party/xla/third_party/llvm/workspace.bzl b/third_party/xla/third_party/llvm/workspace.bzl
index ba37bcd14fb..caf6ccb58d8 100644
--- a/third_party/xla/third_party/llvm/workspace.bzl
+++ b/third_party/xla/third_party/llvm/workspace.bzl
@@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "2b35fb9283a44aae92cb72a41d077ee8de2be7ad"
-    LLVM_SHA256 = "b61c85c97cde1c3e96ea7cfd09737c24578d6d0656c392d79ce402da6904b828"
+    LLVM_COMMIT = "1ee8805f01a663b7f2b09215a8eab5ab2d15ce35"
+    LLVM_SHA256 = "273fa5e9618fa8085aea51d6eb40129677fea8cf4c8825a2c4ff6ea582706a55"
 
     tf_http_archive(
         name = name,
diff --git a/third_party/xla/third_party/shardy/temporary.patch b/third_party/xla/third_party/shardy/temporary.patch
index 95b473b0325..8a259e4c00e 100644
--- a/third_party/xla/third_party/shardy/temporary.patch
+++ b/third_party/xla/third_party/shardy/temporary.patch
@@ -72,7 +72,7 @@ index 1ea9160e..8f149153 100644
  
  TEST_F(DialectTest, DimensionShardingAttrGetSharded) {
 diff --git a/shardy/dialect/sdy/transforms/export/convert_global_to_local.cc b/shardy/dialect/sdy/transforms/export/convert_global_to_local.cc
-index c6b1a677..1b545c49 100644
+index c6b1a677..c77a0ad4 100644
 --- a/shardy/dialect/sdy/transforms/export/convert_global_to_local.cc
 +++ b/shardy/dialect/sdy/transforms/export/convert_global_to_local.cc
 @@ -280,19 +280,6 @@ LogicalResult localizeGenericOp(Operation* op, ValueRange operands,
@@ -206,7 +206,81 @@ index c6b1a677..1b545c49 100644
  }
  
  ValueRange createAllReduceForUnreducedAxes(
-@@ -708,6 +741,24 @@ class AllReduceOpPattern : public OpConversionPattern<sdy::AllReduceOp> {
+@@ -528,18 +561,23 @@ class AllGatherOpPattern : public OpConversionPattern<AllGatherOp> {
+     }
+ 
+     SDY_CHECK(gatheringDims.size() > 1);
+-
+-    // Add a leading dimension of size 1 to gather all partitions at once.
+-    SmallVector<int64_t> tmpShape = {1};
+-    llvm::append_range(tmpShape, localShape);
++    int64_t firstGatheringDim = gatheringDims.front();
++
++    // Add a dimension of size 1 at `firstGatheringDim` to gather all
++    // partitions at once, keeping dimensions 0..firstGatheringDim-1 untouched.
++    SmallVector<int64_t> tmpShape =
++        llvm::to_vector(localShape.take_front(firstGatheringDim));
++    tmpShape.push_back(1);
++    llvm::append_range(tmpShape, localShape.drop_front(firstGatheringDim));
+     Value curInput = stablehlo::ReshapeOp::create(
+         rewriter, op.getLoc(),
+         RankedTensorType::get(tmpShape, localType.getElementType()), input);
+ 
+-    // Perform All-Gather on dim 0.
++    // Perform All-Gather on `firstGatheringDim`.
+     Attribute replicaGroups;
+-    std::tie(replicaGroups, tmpShape[0]) = getReplicaGroupsAndSize(
+-        allGatheringAxes, mesh, meshOrRef, enableRGV3, rewriter);
++    std::tie(replicaGroups, tmpShape[firstGatheringDim]) =
++        getReplicaGroupsAndSize(allGatheringAxes, mesh, meshOrRef, enableRGV3,
++                                rewriter);
+     auto channelHandle = stablehlo::ChannelHandleAttr::get(
+         ctx, /*handle=*/conversionState.getNextChannelId(),
+         kCrossPartitionChannelHandleType);
+@@ -548,13 +586,15 @@ class AllGatherOpPattern : public OpConversionPattern<AllGatherOp> {
+     auto allGather = stablehlo::AllGatherOp::create(
+         rewriter, op.getLoc(),
+         TypeRange{RankedTensorType::get(tmpShape, localType.getElementType())},
+-        curInput, /*all_gather_dim=*/0, replicaGroups, channelHandle,
+-        /*use_global_device_ids=*/true);
++        curInput, /*all_gather_dim=*/firstGatheringDim, replicaGroups,
++        channelHandle, /*use_global_device_ids=*/true);
+     curInput = allGather.getResult(0);
+ 
+-    // Split the leading dimension into individual gathering factors.
+-    SmallVector<int64_t> splitShape = gatheringFactors;
+-    llvm::append_range(splitShape, localShape);
++    // Split `firstGatheringDim` into individual gathering factors.
++    SmallVector<int64_t> splitShape =
++        llvm::to_vector(localShape.take_front(firstGatheringDim));
++    llvm::append_range(splitShape, gatheringFactors);
++    llvm::append_range(splitShape, localShape.drop_front(firstGatheringDim));
+     curInput = stablehlo::ReshapeOp::create(
+         rewriter, op.getLoc(),
+         RankedTensorType::get(splitShape, localType.getElementType()),
+@@ -562,14 +602,16 @@ class AllGatherOpPattern : public OpConversionPattern<AllGatherOp> {
+ 
+     // Transpose factors next to their corresponding dimensions.
+     SmallVector<int64_t> xposePerm;
+-    int64_t nextFactorIdx = 0;
+-    for (int64_t i = 0; i < rank; ++i) {
++    xposePerm.reserve(rank + gatheringDims.size());
++    for (int64_t i = 0; i < firstGatheringDim; ++i) {
++      xposePerm.push_back(i);
++    }
++    int64_t nextFactorIdx = firstGatheringDim;
++    for (int64_t i = firstGatheringDim; i < rank; ++i) {
+       if (llvm::is_contained(gatheringDims, i)) {
+         xposePerm.push_back(nextFactorIdx++);
+-        xposePerm.push_back(i + gatheringDims.size());
+-      } else {
+-        xposePerm.push_back(i + gatheringDims.size());
+       }
++      xposePerm.push_back(i + gatheringDims.size());
+     }
+     curInput = stablehlo::TransposeOp::create(
+         rewriter, op.getLoc(), curInput,
+@@ -708,6 +750,24 @@ class AllReduceOpPattern : public OpConversionPattern<sdy::AllReduceOp> {
    bool enableRGV3;
  };
  
@@ -231,7 +305,7 @@ index c6b1a677..1b545c49 100644
  // Returns a 0-rank i32 tensor containing the global offset for the given shard
  // axes and local shard size.
  Value getDimensionOffset(Location loc, MeshAttr mesh,
-@@ -716,9 +767,6 @@ Value getDimensionOffset(Location loc, MeshAttr mesh,
+@@ -716,9 +776,6 @@ Value getDimensionOffset(Location loc, MeshAttr mesh,
                           ConversionPatternRewriter& rewriter) {
    int64_t numDevices = mesh.getTotalSize();
    Type i32Ty = rewriter.getI32Type();
@@ -241,7 +315,7 @@ index c6b1a677..1b545c49 100644
  
    // Calculate a compile-time offset table for this dimension.
    SmallVector<int32_t> offsetsTable = llvm::map_to_vector(
-@@ -728,16 +776,11 @@ Value getDimensionOffset(Location loc, MeshAttr mesh,
+@@ -728,16 +785,11 @@ Value getDimensionOffset(Location loc, MeshAttr mesh,
          return static_cast<int32_t>(offset);
        });
  
@@ -262,7 +336,7 @@ index c6b1a677..1b545c49 100644
  }
  
  Value emitDynamicSliceForAxes(Location loc, Value globalTensor, MeshAttr mesh,
-@@ -767,6 +810,43 @@ Value emitDynamicSliceForAxes(Location loc, Value globalTensor, MeshAttr mesh,
+@@ -767,6 +819,43 @@ Value emitDynamicSliceForAxes(Location loc, Value globalTensor, MeshAttr mesh,
        .getResult();
  }
  
@@ -306,7 +380,7 @@ index c6b1a677..1b545c49 100644
  class AllSliceOpPattern : public OpConversionPattern<AllSliceOp> {
   public:
    AllSliceOpPattern(TypeConverter& converter, MLIRContext* ctx,
-@@ -1119,7 +1199,14 @@ class ConstantOpPattern : public OpConversionPattern<sdy::ConstantOp> {
+@@ -1119,7 +1208,14 @@ class ConstantOpPattern : public OpConversionPattern<sdy::ConstantOp> {
      auto localType = cast<RankedTensorType>(converter->convertType(op));
      ElementsAttr elementsAttr = op.getValue();
      Location loc = op.getLoc();
@@ -321,7 +395,7 @@ index c6b1a677..1b545c49 100644
      // Unsharded or splat constants.
      if (localType == globalType || elementsAttr.isSplat()) {
        Attribute newAttr =
-@@ -1128,33 +1215,35 @@ class ConstantOpPattern : public OpConversionPattern<sdy::ConstantOp> {
+@@ -1128,33 +1224,35 @@ class ConstantOpPattern : public OpConversionPattern<sdy::ConstantOp> {
                : SplatElementsAttr::get(localType,
                                         elementsAttr.getSplatValue<Attribute>());
  
@@ -379,7 +453,7 @@ index c6b1a677..1b545c49 100644
      conversionState.removeToConvertOp(op);
  
      return success();
-@@ -1425,6 +1514,97 @@ class ReduceScatterOpPattern : public OpConversionPattern<ReduceScatterOp> {
+@@ -1425,6 +1523,97 @@ class ReduceScatterOpPattern : public OpConversionPattern<ReduceScatterOp> {
    bool enableRGV3;
  };
  
@@ -477,7 +551,7 @@ index c6b1a677..1b545c49 100644
  class NamedComputationOpPattern
      : public OpConversionPattern<NamedComputationOp> {
   public:
-@@ -1608,11 +1788,11 @@ class StablehloConvolutionOpPattern
+@@ -1608,11 +1797,11 @@ class StablehloConvolutionOpPattern
        addRhsReductionAxes(spatialDim);
      }
  
@@ -492,7 +566,7 @@ index c6b1a677..1b545c49 100644
      }
  
      auto getAxes = [&](TensorShardingAttr sharding, int64_t dim) {
-@@ -1643,7 +1823,7 @@ class StablehloConvolutionOpPattern
+@@ -1643,7 +1832,7 @@ class StablehloConvolutionOpPattern
          localBatchGroupCount, op.getPrecisionConfigAttr());
  
      ValueRange results = localConv->getResults();
@@ -501,7 +575,7 @@ index c6b1a677..1b545c49 100644
        Attribute meshOrRef =
            rhsSharding ? rhsSharding.getMeshOrRef() : lhsSharding.getMeshOrRef();
        results = createAllReduceForUnreducedAxes(
-@@ -1727,8 +1907,8 @@ class StablehloDotOpPattern : public OpConversionPattern<stablehlo::DotOp> {
+@@ -1727,8 +1916,8 @@ class StablehloDotOpPattern : public OpConversionPattern<stablehlo::DotOp> {
  
          // Verify that an AllReduce has been inserted by a previous pass unless
          // the result is explicitly unreduced.
@@ -512,7 +586,7 @@ index c6b1a677..1b545c49 100644
            Operation* localDot = createLocalGenericOp(op, adaptor.getOperands(),
                                                       rewriter, converter);
            ValueRange results = createAllReduceForUnreducedAxes(
-@@ -1804,8 +1984,8 @@ class StablehloDotGeneralOpPattern
+@@ -1804,8 +1993,8 @@ class StablehloDotGeneralOpPattern
  
          // Verify that an AllReduce has been inserted by a previous pass unless
          // the result is explicitly unreduced.
@@ -523,7 +597,7 @@ index c6b1a677..1b545c49 100644
            Operation* localDotGeneral = createLocalGenericOp(
                op, adaptor.getOperands(), rewriter, converter);
            ValueRange results = createAllReduceForUnreducedAxes(
-@@ -1828,6 +2008,71 @@ class StablehloDotGeneralOpPattern
+@@ -1828,6 +2017,71 @@ class StablehloDotGeneralOpPattern
    bool enableRGV3;
  };
  
@@ -595,7 +669,7 @@ index c6b1a677..1b545c49 100644
  Value createScalarInt(OpBuilder& b, Location loc, RankedTensorType type,
                        int64_t val) {
    return stablehlo::ConstantOp::create(
-@@ -2095,7 +2340,7 @@ class StablehloReduceOpPattern
+@@ -2095,7 +2349,7 @@ class StablehloReduceOpPattern
      }
  
      if (!reductionAxes.empty() &&
@@ -604,7 +678,7 @@ index c6b1a677..1b545c49 100644
        auto localReduceOp = cast<stablehlo::ReduceOp>(
            createLocalGenericOp(op, adaptor.getOperands(), rewriter, converter));
        if (op.getNumResults() == 1) {
-@@ -2555,8 +2800,8 @@ class StablehloGatherOpPattern
+@@ -2555,8 +2809,8 @@ class StablehloGatherOpPattern
          mask, bcastDims);
  
      TensorShardingAttr resSharding = getSharding(op.getResult());
@@ -615,7 +689,7 @@ index c6b1a677..1b545c49 100644
      ReductionOp reductionOp =
          resSharding ? resSharding.getReductionOp() : ReductionOp::SUM;
  
-@@ -2573,7 +2818,7 @@ class StablehloGatherOpPattern
+@@ -2573,7 +2827,7 @@ class StablehloGatherOpPattern
      result = stablehlo::SelectOp::create(rewriter, loc, bcastMask, result,
                                           identityVal);
  
@@ -624,7 +698,7 @@ index c6b1a677..1b545c49 100644
        result =
            createAllReduceForUnreducedAxes(loc, result, reductionAxes, mesh,
                                            operandSharding.getMeshOrRef(),
-@@ -2747,30 +2992,8 @@ class StablehloRngBitGeneratorOpPattern
+@@ -2747,30 +3001,8 @@ class StablehloRngBitGeneratorOpPattern
      // Mask non-leader shards (where shard index across dataAxes != 0) to zero
      // and all-reduce across dataAxes so all partitions receive device 0's
      // updated output_state.
@@ -657,7 +731,7 @@ index c6b1a677..1b545c49 100644
  
      finalState =
          createAllReduceForUnreducedAxes(loc, ValueRange{maskedState}, dataAxes,
-@@ -3121,35 +3344,13 @@ class StablehloScatterOpPattern
+@@ -3121,35 +3353,13 @@ class StablehloScatterOpPattern
        if (!identity) {
          return op.emitError("failed to get scatter reduction identity.");
        }
@@ -696,7 +770,7 @@ index c6b1a677..1b545c49 100644
      }
  
      auto globalInputType = cast<RankedTensorType>(input.getType());
-@@ -3207,7 +3408,8 @@ class StablehloScatterOpPattern
+@@ -3207,7 +3417,8 @@ class StablehloScatterOpPattern
  
      ValueRange results = localScatter.getResults();
      if (!reductionAxes.empty() &&
@@ -706,7 +780,7 @@ index c6b1a677..1b545c49 100644
        Attribute meshOrRef = inputSharding ? inputSharding.getMeshOrRef()
                                            : updateSharding.getMeshOrRef();
        results = createAllReduceForUnreducedAxes(
-@@ -3503,7 +3705,8 @@ struct ConvertGlobalToLocalPass
+@@ -3503,7 +3714,8 @@ struct ConvertGlobalToLocalPass
      patterns.add<AllSliceOpPattern, CollectivePermuteOpPattern,
                   ConstantOpPattern, FuncOpSignaturePattern, GenericOpPattern,
                   ManualComputationOpPattern, NamedComputationOpPattern,
@@ -811,6 +885,38 @@ index c840e9f1..f11dcdbe 100644
  // CHECK-LABEL: func.func @callee_sharded
  // CHECK-SAME: (%arg0: tensor<8xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}]>}) -> (tensor<8xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}]>})
  func.func @callee_sharded(%arg0: tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}]>}) -> (tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}]>}) {
+diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_all_gather.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_all_gather.mlir
+index d79f9dae..7e6ae755 100644
+--- a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_all_gather.mlir
++++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_all_gather.mlir
+@@ -174,3 +174,27 @@ func.func @two_dims_merge_adjacent_sub_axes(%arg0 : tensor<8x16xf32> {sdy.shardi
+   return %0 : tensor<8x16xf32>
+ }
+ 
++// CHECK-LABEL: func @two_dims_non_zero_first_dim
++// CHECK-SAME: (%[[ARG0:.*]]: tensor<3x5x7x11xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4_2, [{}, {"y"}, {}, {"z"}]>})
++// CHECK-SAME: -> (tensor<3x20x7x22xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4_2, [{}, {}, {}, {}]>})
++func.func @two_dims_non_zero_first_dim(%arg0 : tensor<3x20x7x22xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4_2, [{}, {"y"}, {}, {"z"}]>})
++  -> (tensor<3x20x7x22xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4_2, [{}, {}, {}, {}]>}){
++
++  // --- Combined Dimensions All-Gather Strategy (per-dim-all-gather=false) ---
++  // COMBINED: %[[RESHAPE1:.*]] = stablehlo.reshape %[[ARG0]] : (tensor<3x5x7x11xf32>) -> tensor<3x1x5x7x11xf32>
++  // COMBINED: %[[GATHER:.*]] = "stablehlo.all_gather"(%[[RESHAPE1]]) <{
++  // COMBINED-SAME: all_gather_dim = 1 : i64,
++  // COMBINED-SAME: channel_handle = #stablehlo.channel_handle<handle = 8, type = 1>,
++  // COMBINED-V1-SAME{LITERAL}: replica_groups = dense<[[0, 1, 2, 3, 4, 5, 6, 7], [8, 9, 10, 11, 12, 13, 14, 15]]>
++  // COMBINED-V3-SAME: replica_groups = #stablehlo.replica_group_mesh_axes<mesh = @mesh_2_4_2, axes = [#stablehlo.axis_ref<name = "y">, #stablehlo.axis_ref<name = "z">]>
++  // COMBINED-SAME: }> : (tensor<3x1x5x7x11xf32>) -> tensor<3x8x5x7x11xf32>
++  // COMBINED: %[[RESHAPE2:.*]] = stablehlo.reshape %[[GATHER]] : (tensor<3x8x5x7x11xf32>) -> tensor<3x4x2x5x7x11xf32>
++  // COMBINED: %[[TRANSPOSE:.*]] = stablehlo.transpose %[[RESHAPE2]], dims = [0, 1, 3, 4, 2, 5] : (tensor<3x4x2x5x7x11xf32>) -> tensor<3x4x5x7x2x11xf32>
++  // COMBINED: %[[RESULT:.*]] = stablehlo.reshape %[[TRANSPOSE]] : (tensor<3x4x5x7x2x11xf32>) -> tensor<3x20x7x22xf32>
++
++  %0 = sdy.all_gather[{}, {"y"}, {}, {"z"}] %arg0 out_sharding = <@mesh_2_4_2, [{}, {}, {}, {}]> : tensor<3x20x7x22xf32>
++
++  // COMBINED: return %[[RESULT]] : tensor<3x20x7x22xf32>
++  return %0 : tensor<3x20x7x22xf32>
++}
++
 diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_constant.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_constant.mlir
 index 2e01e015..15a45782 100644
 --- a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_constant.mlir
@@ -1153,6 +1259,52 @@ index e45919c6..3661799c 100644
    // CHECK:      %[[INPUT_SEL:.*]] = stablehlo.select %[[IS_LEADER]], %[[ARG0]], %[[ID_BCAST]]
    // CHECK:      %[[SCATTER:.*]] = "stablehlo.scatter"(%[[INPUT_SEL]], %[[ARG1]], %[[ARG2]])
    // CHECK-SAME: scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [2], inserted_window_dims = [0, 1], scatter_dims_to_operand_dims = [0, 1], index_vector_dim = 2>
+diff --git a/shardy/dialect/sdy/transforms/export/test/executable_convert_global_to_local/sdy_all_gather.mlir b/shardy/dialect/sdy/transforms/export/test/executable_convert_global_to_local/sdy_all_gather.mlir
+index 2377e4a7..47c5801a 100644
+--- a/shardy/dialect/sdy/transforms/export/test/executable_convert_global_to_local/sdy_all_gather.mlir
++++ b/shardy/dialect/sdy/transforms/export/test/executable_convert_global_to_local/sdy_all_gather.mlir
+@@ -30,6 +30,14 @@ func.func @all_gather_test(
+   return %0 : tensor<4x4xi32>
+ }
+ 
++// This function gathers non-zero dims 1 and 3 of a 2x4x2x2 tensor.
++func.func @all_gather_test_non_zero_first_dim(
++  %arg0: tensor<2x4x2x2xi32> {sdy.sharding = #sdy.sharding<@mesh_2_2, [{}, {"x"}, {}, {"y"}]>}
++  ) -> tensor<2x4x2x2xi32> {
++  %0 = sdy.all_gather [{}, {"x"}, {}, {"y"}] %arg0 out_sharding=<@mesh_2_2, [{}, {}, {}, {}]> : tensor<2x4x2x2xi32>
++  return %0 : tensor<2x4x2x2xi32>
++}
++
+ //--- part2.mlir
+ 
+ func.func @main() {
+@@ -63,5 +71,26 @@ func.func @main() {
+   "check.expect_eq"(%res_comb#0, %res_pdim#0) : (tensor<4x4xi32>, tensor<4x4xi32>) -> ()
+   "check.expect_eq"(%res_comb#0, %cst) : (tensor<4x4xi32>, tensor<4x4xi32>) -> ()
+ 
++  // Test gathering on non-zero dims 1 and 3 with non-gathering dims 0 and 2.
++  %iota = stablehlo.iota dim = 0 : tensor<32xi32>
++  %cst_4d = stablehlo.reshape %iota : (tensor<32xi32>) -> tensor<2x4x2x2xi32>
++  %s0_4d = stablehlo.slice %cst_4d [0:2, 0:2, 0:2, 0:1] : (tensor<2x4x2x2xi32>) -> tensor<2x2x2x1xi32>
++  %s1_4d = stablehlo.slice %cst_4d [0:2, 0:2, 0:2, 1:2] : (tensor<2x4x2x2xi32>) -> tensor<2x2x2x1xi32>
++  %s2_4d = stablehlo.slice %cst_4d [0:2, 2:4, 0:2, 0:1] : (tensor<2x4x2x2xi32>) -> tensor<2x2x2x1xi32>
++  %s3_4d = stablehlo.slice %cst_4d [0:2, 2:4, 0:2, 1:2] : (tensor<2x4x2x2xi32>) -> tensor<2x2x2x1xi32>
++
++  %res_comb_4d:4 = "interpreter.run_parallel"(%s0_4d, %s1_4d, %s2_4d, %s3_4d) {
++    programs = [[@combined_gather_non_zero_first_dim, @combined_gather_non_zero_first_dim, @combined_gather_non_zero_first_dim, @combined_gather_non_zero_first_dim]]
++  } : (tensor<2x2x2x1xi32>, tensor<2x2x2x1xi32>, tensor<2x2x2x1xi32>, tensor<2x2x2x1xi32>) ->
++      (tensor<2x4x2x2xi32>, tensor<2x4x2x2xi32>, tensor<2x4x2x2xi32>, tensor<2x4x2x2xi32>)
++
++  %res_pdim_4d:4 = "interpreter.run_parallel"(%s0_4d, %s1_4d, %s2_4d, %s3_4d) {
++    programs = [[@per_dim_gather_non_zero_first_dim, @per_dim_gather_non_zero_first_dim, @per_dim_gather_non_zero_first_dim, @per_dim_gather_non_zero_first_dim]]
++  } : (tensor<2x2x2x1xi32>, tensor<2x2x2x1xi32>, tensor<2x2x2x1xi32>, tensor<2x2x2x1xi32>) ->
++      (tensor<2x4x2x2xi32>, tensor<2x4x2x2xi32>, tensor<2x4x2x2xi32>, tensor<2x4x2x2xi32>)
++
++  "check.expect_eq"(%res_comb_4d#0, %res_pdim_4d#0) : (tensor<2x4x2x2xi32>, tensor<2x4x2x2xi32>) -> ()
++  "check.expect_eq"(%res_comb_4d#0, %cst_4d) : (tensor<2x4x2x2xi32>, tensor<2x4x2x2xi32>) -> ()
++
+   return
+ }
 diff --git a/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/stablehlo_pad.mlir b/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/stablehlo_pad.mlir
 index 59cc6b86..208aa2c2 100644
 --- a/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/stablehlo_pad.mlir
@@ -1452,7 +1604,7 @@ index c5cd8022..6d9f32c2 100644
 +  return %1 : tensor<8x8xf32>
 +}
 diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
-index e42d7dbd..ba37bcd1 100644
+index e42d7dbd..caf6ccb5 100644
 --- a/third_party/llvm/workspace.bzl
 +++ b/third_party/llvm/workspace.bzl
 @@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
@@ -1461,8 +1613,8 @@ index e42d7dbd..ba37bcd1 100644
      """Imports LLVM."""
 -    LLVM_COMMIT = "2e0d9e12b36e853d9b108c4da09509cd54d1f3f9"
 -    LLVM_SHA256 = "8e4a851780af109d623fa6d6db9b1e62d7126272d17268fed1b149f0e737c1d7"
-+    LLVM_COMMIT = "2b35fb9283a44aae92cb72a41d077ee8de2be7ad"
-+    LLVM_SHA256 = "b61c85c97cde1c3e96ea7cfd09737c24578d6d0656c392d79ce402da6904b828"
++    LLVM_COMMIT = "1ee8805f01a663b7f2b09215a8eab5ab2d15ce35"
++    LLVM_SHA256 = "273fa5e9618fa8085aea51d6eb40129677fea8cf4c8825a2c4ff6ea582706a55"
  
      tf_http_archive(
          name = name,

From d27685655244809e48def3318ca1d79dc7b88cfe Mon Sep 17 00:00:00 2001
From: Artsiom Ablavatski <artsiom@google.com>
Date: Thu, 15 Oct 2026 10:41:17 -0700
Subject: [PATCH 0881/1000] Fold f32, DRQ, and Weight-Only Conv2D,
 DepthwiseConv2D, TransposeConv, and FC ops.

PiperOrigin-RevId: 1000238610
---
 tensorflow/compiler/mlir/lite/ir/tfl_ops.cc   | 864 +++++++++++++++++-
 tensorflow/compiler/mlir/lite/ir/tfl_ops.td   |   4 +
 .../compiler/mlir/lite/tests/const-fold.mlir  | 443 ++++++++-
 .../compiler/mlir/lite/tests/optimize.mlir    |  41 +
 .../mlir/lite/transforms/optimize_pass.cc     |  10 +-
 5 files changed, 1326 insertions(+), 36 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc b/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc
index 4cd59932c74..4b50a8af5cf 100644
--- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc
+++ b/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc
@@ -2180,6 +2180,262 @@ LogicalResult FullyConnectedOp::verify() {
   return mlir::success();
 }
 
+namespace {
+
+bool ExtractDequantizedWeightsFromQConst(QConstOp qconst_op,
+                                         ShapedType weights_type,
+                                         std::vector<float>& weights_values) {
+  if (!qconst_op || !weights_type.hasStaticShape() ||
+      weights_type.getRank() < 1) {
+    return false;
+  }
+  auto quant_type =
+      mlir::dyn_cast<quant::QuantizedType>(weights_type.getElementType());
+  if (!quant_type || !quant_type.getExpressedType().isF32()) {
+    return false;
+  }
+
+  const unsigned bit_width = quant_type.getStorageTypeIntegralWidth();
+  if (bit_width != 8 && bit_width != 4) {
+    return false;
+  }
+  const bool is_signed = quant_type.isSigned();
+
+  const int64_t num_elements = weights_type.getNumElements();
+
+  int32_t quant_dim = -1;
+  ArrayRef<double> scales;
+  ArrayRef<int64_t> zero_points;
+  double single_scale = 0.0;
+  int64_t single_zp = 0;
+
+  if (auto uq_type = mlir::dyn_cast<quant::UniformQuantizedType>(quant_type)) {
+    single_scale = uq_type.getScale();
+    single_zp = uq_type.getZeroPoint();
+    scales = ArrayRef<double>(&single_scale, 1);
+    zero_points = ArrayRef<int64_t>(&single_zp, 1);
+  } else if (auto per_axis_type =
+                 mlir::dyn_cast<quant::UniformQuantizedPerAxisType>(
+                     quant_type)) {
+    quant_dim = per_axis_type.getQuantizedDimension();
+    if (quant_dim < 0 || quant_dim >= weights_type.getRank()) {
+      return false;
+    }
+    scales = per_axis_type.getScales();
+    zero_points = per_axis_type.getZeroPoints();
+    const int64_t expected_channels = weights_type.getDimSize(quant_dim);
+    if (static_cast<int64_t>(scales.size()) != expected_channels ||
+        static_cast<int64_t>(zero_points.size()) != expected_channels) {
+      return false;
+    }
+  } else {
+    return false;
+  }
+
+  int64_t dim_size = 1;
+  int64_t stride = 1;
+  if (quant_dim >= 0) {
+    dim_size = weights_type.getDimSize(quant_dim);
+    if (dim_size <= 0) return false;
+    for (int i = quant_dim + 1; i < weights_type.getRank(); ++i) {
+      const int64_t d = weights_type.getDimSize(i);
+      if (d <= 0) return false;
+      stride *= d;
+    }
+  }
+
+  auto dequantize_at = [&](int64_t flat_idx, int64_t q) -> float {
+    const int64_t channel =
+        (quant_dim >= 0) ? ((flat_idx / stride) % dim_size) : 0;
+    return static_cast<float>(q - zero_points[channel]) *
+           static_cast<float>(scales[channel]);
+  };
+
+  ElementsAttr value_attr = qconst_op.getValue();
+  if (!value_attr) return false;
+  auto value_shaped_type = mlir::dyn_cast<ShapedType>(value_attr.getType());
+  if (!value_shaped_type || !value_shaped_type.hasStaticShape() ||
+      value_shaped_type.getNumElements() != num_elements ||
+      !value_shaped_type.getElementType().isInteger(bit_width)) {
+    return false;
+  }
+
+  weights_values.resize(num_elements);
+  if (auto dense_attr = mlir::dyn_cast<DenseElementsAttr>(value_attr)) {
+    int64_t idx = 0;
+    for (const APInt& val : dense_attr.getValues<APInt>()) {
+      const int64_t q = is_signed ? val.getSExtValue() : val.getZExtValue();
+      weights_values[idx] = dequantize_at(idx, q);
+      ++idx;
+    }
+    return idx == num_elements;
+  }
+
+  if (auto res_attr = mlir::dyn_cast<DenseResourceElementsAttr>(value_attr)) {
+    AsmResourceBlob* blob = res_attr.getRawHandle().getBlob();
+    if (!blob) return false;
+    ArrayRef<char> raw_data = blob->getData();
+    if (bit_width == 8) {
+      if (static_cast<int64_t>(raw_data.size()) != num_elements) {
+        return false;
+      }
+      for (int64_t idx = 0; idx < num_elements; ++idx) {
+        const int64_t q =
+            is_signed
+                ? static_cast<int64_t>(static_cast<int8_t>(raw_data[idx]))
+                : static_cast<int64_t>(static_cast<uint8_t>(raw_data[idx]));
+        weights_values[idx] = dequantize_at(idx, q);
+      }
+      return true;
+    }
+    if (bit_width == 4) {
+      const int64_t packed_size = (num_elements + 1) / 2;
+      if (static_cast<int64_t>(raw_data.size()) == num_elements) {
+        for (int64_t idx = 0; idx < num_elements; ++idx) {
+          const uint8_t byte = static_cast<uint8_t>(raw_data[idx]) & 0x0F;
+          const int64_t q =
+              is_signed
+                  ? static_cast<int64_t>(static_cast<int8_t>(byte << 4) >> 4)
+                  : static_cast<int64_t>(byte);
+          weights_values[idx] = dequantize_at(idx, q);
+        }
+        return true;
+      }
+      if (static_cast<int64_t>(raw_data.size()) == packed_size) {
+        for (int64_t idx = 0; idx < num_elements; ++idx) {
+          const uint8_t byte = static_cast<uint8_t>(raw_data[idx / 2]);
+          const uint8_t nibble = (byte >> ((idx % 2) * 4)) & 0x0F;
+          const int64_t q =
+              is_signed
+                  ? static_cast<int64_t>(static_cast<int8_t>(nibble << 4) >> 4)
+                  : static_cast<int64_t>(nibble);
+          weights_values[idx] = dequantize_at(idx, q);
+        }
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
+class FloatElementsAccessor {
+ public:
+  static std::optional<FloatElementsAccessor> Create(ElementsAttr attr) {
+    if (!attr) return std::nullopt;
+    auto shaped_type = mlir::dyn_cast<ShapedType>(attr.getType());
+    if (!shaped_type || !shaped_type.hasStaticShape() ||
+        !shaped_type.getElementType().isF32()) {
+      return std::nullopt;
+    }
+    FloatElementsAccessor accessor;
+    if (auto dense_attr = mlir::dyn_cast<DenseElementsAttr>(attr)) {
+      accessor.dense_values_ = dense_attr.getValues<float>();
+      return accessor;
+    }
+    if (auto res_attr = mlir::dyn_cast<DenseResourceElementsAttr>(attr)) {
+      ArrayRef<float> data = GetValues<float>(res_attr);
+      if (static_cast<int64_t>(data.size()) != shaped_type.getNumElements()) {
+        return std::nullopt;
+      }
+      accessor.resource_values_ = data;
+      return accessor;
+    }
+    return std::nullopt;
+  }
+
+  static std::optional<FloatElementsAccessor> CreateFromQConst(
+      QConstOp qconst_op, ShapedType weights_type) {
+    FloatElementsAccessor accessor;
+    if (!ExtractDequantizedWeightsFromQConst(qconst_op, weights_type,
+                                             accessor.owned_values_)) {
+      return std::nullopt;
+    }
+    return accessor;
+  }
+
+  float operator[](int64_t index) const {
+    if (dense_values_.has_value()) {
+      return (*dense_values_)[index];
+    }
+    if (!owned_values_.empty()) {
+      return owned_values_[index];
+    }
+    return resource_values_[index];
+  }
+
+ private:
+  std::optional<mlir::detail::ElementsAttrRange<
+      DenseElementsAttr::ElementIterator<float>>>
+      dense_values_;
+  ArrayRef<float> resource_values_;
+  std::vector<float> owned_values_;
+};
+
+QConstOp GetWeightsQConstOp(Value weights) {
+  if (auto qconst_op = weights.getDefiningOp<QConstOp>()) {
+    return qconst_op;
+  }
+  if (auto dequant_op = weights.getDefiningOp<DequantizeOp>()) {
+    return dequant_op.getInput().getDefiningOp<QConstOp>();
+  }
+  return nullptr;
+}
+
+std::optional<FloatElementsAccessor> CreateWeightsAccessor(
+    ElementsAttr weights_tensor, QConstOp weights_qconst_op,
+    ShapedType weights_type) {
+  if (weights_tensor && weights_type.getElementType().isF32()) {
+    return FloatElementsAccessor::Create(weights_tensor);
+  }
+  if (weights_qconst_op &&
+      mlir::isa<quant::QuantizedType>(weights_type.getElementType())) {
+    return FloatElementsAccessor::CreateFromQConst(weights_qconst_op,
+                                                   weights_type);
+  }
+  return std::nullopt;
+}
+
+void NeumaierAdd(float addend, float& sum, float& compensation) {
+  const float new_sum = sum + addend;
+  // DO NOT enable -funsafe-math-optimizations here.
+  // There is a test detecting unsafe optimizations.
+  // Unsafe math optimizations can reorder float formulas, and set the
+  // compensation to constant 0. The formula must be evaluated as written
+  // for the algorithm to work.
+  // (Note: -ffast-math is a superset of -funsafe-math-optimizations.)
+  if (std::abs(sum) >= std::abs(addend)) {
+    compensation += (sum - new_sum) + addend;
+  } else {
+    compensation += (addend - new_sum) + sum;
+  }
+  sum = new_sum;
+}
+
+// Maximum number of multiply-accumulate operations allowed during constant
+// folding of FullyConnected and Conv ops, to avoid excessive compile times.
+constexpr int64_t kMaxFoldMultiplyAdds = 100'000'000;
+
+bool IsWithinFoldLimits(llvm::ArrayRef<int64_t> loop_bounds,
+                        ShapedType input_type, ShapedType weights_type,
+                        ShapedType output_type) {
+  int64_t total_macs = 1;
+  for (int64_t bound : loop_bounds) {
+    if (bound <= 0 || total_macs > kMaxFoldMultiplyAdds / bound) {
+      return false;
+    }
+    total_macs *= bound;
+  }
+  // Avoid folding when the output constant has more elements than the input
+  // and weight constants combined, which would inflate model size.
+  if (output_type.getNumElements() >
+      input_type.getNumElements() + weights_type.getNumElements()) {
+    return false;
+  }
+  return true;
+}
+
+}  // namespace
+
 LogicalResult FullyConnectedOp::fold(FoldAdaptor adaptor,
                                      SmallVectorImpl<OpFoldResult>& results) {
   assert(adaptor.getOperands().size() == 3);
@@ -2188,6 +2444,7 @@ LogicalResult FullyConnectedOp::fold(FoldAdaptor adaptor,
   // besides the default.
   if (getFusedActivationFunction() != "NONE") return failure();
   if (getWeightsFormat() != "DEFAULT") return failure();
+  if (getNumResults() != 1) return failure();
 
   // Bias tensor is optional.
   const bool has_bias =
@@ -2195,32 +2452,32 @@ LogicalResult FullyConnectedOp::fold(FoldAdaptor adaptor,
 
   // Get the tensors.
   auto operands = adaptor.getOperands();
-  DenseElementsAttr input_tensor =
-      dyn_cast_or_null<DenseElementsAttr>(operands[0]);
-  DenseElementsAttr weights_tensor =
-      dyn_cast_or_null<DenseElementsAttr>(operands[1]);
-  DenseElementsAttr bias_tensor;
+  ElementsAttr input_tensor = dyn_cast_or_null<ElementsAttr>(operands[0]);
+  ElementsAttr weights_tensor = dyn_cast_or_null<ElementsAttr>(operands[1]);
+  QConstOp weights_qconst_op = GetWeightsQConstOp(getFilter());
+  ElementsAttr bias_tensor;
 
-  if (!input_tensor || !weights_tensor) {
+  if (!input_tensor || (!weights_tensor && !weights_qconst_op)) {
     return failure();
   }
 
   if (has_bias) {
-    bias_tensor = dyn_cast_or_null<DenseElementsAttr>(operands[2]);
+    bias_tensor = dyn_cast_or_null<ElementsAttr>(operands[2]);
     if (!bias_tensor) return failure();
   }
 
   // Get the tensor types.
   const auto input_type = mlir::cast<ShapedType>(input_tensor.getType());
-  const auto weights_type = mlir::cast<ShapedType>(weights_tensor.getType());
+  const auto weights_type = mlir::cast<ShapedType>(
+      weights_qconst_op ? weights_qconst_op.getOutput().getType()
+                        : getFilter().getType());
   const auto bias_type =
       has_bias ? mlir::cast<ShapedType>(bias_tensor.getType()) : ShapedType{};
 
   const auto output_type = mlir::cast<ShapedType>(getType(0));
 
-  // Folding only implemented for float tensors.
+  // Folding only implemented for float input/bias/output tensors.
   if (!input_type.getElementType().isF32() ||
-      !weights_type.getElementType().isF32() ||
       !output_type.getElementType().isF32() ||
       (has_bias && !bias_type.getElementType().isF32())) {
     return failure();
@@ -2228,6 +2485,7 @@ LogicalResult FullyConnectedOp::fold(FoldAdaptor adaptor,
 
   // Folding only implemented for static shapes
   if (!input_type.hasStaticShape() || !weights_type.hasStaticShape() ||
+      !output_type.hasStaticShape() ||
       (has_bias && !bias_type.hasStaticShape())) {
     return failure();
   }
@@ -2238,6 +2496,10 @@ LogicalResult FullyConnectedOp::fold(FoldAdaptor adaptor,
 
   const int64_t in_dim = weights_type.getDimSize(1);
   const int64_t out_dim = weights_type.getDimSize(0);
+  if (in_dim <= 0 || out_dim <= 0 ||
+      input_type.getNumElements() % in_dim != 0) {
+    return failure();
+  }
 
   if (has_bias) {
     if (bias_type.getRank() > 2 ||
@@ -2253,36 +2515,37 @@ LogicalResult FullyConnectedOp::fold(FoldAdaptor adaptor,
     return failure();
   }
 
-  auto input_values_range = input_tensor.getValues<float>();
-  auto weights_values_range = weights_tensor.getValues<float>();
-  std::optional<decltype(input_values_range)> bias_values_range;
-  if (has_bias) bias_values_range = bias_tensor.getValues<float>();
+  if (!IsWithinFoldLimits({batch_size, out_dim, in_dim}, input_type,
+                          weights_type, output_type)) {
+    return failure();
+  }
+
+  auto input_values_range = FloatElementsAccessor::Create(input_tensor);
+  if (!input_values_range.has_value()) return failure();
+
+  auto weights_values_range =
+      CreateWeightsAccessor(weights_tensor, weights_qconst_op, weights_type);
+  if (!weights_values_range.has_value()) return failure();
+
+  std::optional<FloatElementsAccessor> bias_values_range;
+  if (has_bias) {
+    bias_values_range = FloatElementsAccessor::Create(bias_tensor);
+    if (!bias_values_range.has_value()) return failure();
+  }
 
   // Do the actual folding, one output at a time.
   std::vector<float> result_values;
   result_values.reserve(batch_size * out_dim);
 
-  for (int b = 0; b < batch_size; ++b) {
-    for (int o = 0; o < out_dim; ++o) {
+  for (int64_t b = 0; b < batch_size; ++b) {
+    for (int64_t o = 0; o < out_dim; ++o) {
       // Dot product with Kahan/Neumaier summation to minimize numeric errors.
       float sum = has_bias ? (*bias_values_range)[o] : 0.0f;
       float compensation = 0.0f;
-      for (int i = 0; i < in_dim; ++i) {
-        const float addend = input_values_range[b * in_dim + i] *
-                             weights_values_range[o * in_dim + i];
-        const float new_sum = sum + addend;
-        // DO NOT enable -funsafe-math-optimizations here.
-        // There is a test detecting unsafe optimizations.
-        // Unsafe math optimizations can reorder float formulas, and set the
-        // compensation to constant 0. The formula must be evaluated as written
-        // for the algorithm to work.
-        // (Note: -ffast-math is a superset of -funsafe-math-optimizations.)
-        if (std::abs(sum) >= std::abs(addend)) {
-          compensation += (sum - new_sum) + addend;
-        } else {
-          compensation += (addend - new_sum) + sum;
-        }
-        sum = new_sum;
+      for (int64_t i = 0; i < in_dim; ++i) {
+        const float addend = (*input_values_range)[b * in_dim + i] *
+                             (*weights_values_range)[o * in_dim + i];
+        NeumaierAdd(addend, sum, compensation);
       }
       result_values.push_back(sum + compensation);
     }
@@ -2432,6 +2695,183 @@ int64_t Conv2DOp::GetArithmeticCount(Operation* op) {
   return -1;
 }
 
+OpFoldResult Conv2DOp::fold(FoldAdaptor adaptor) {
+  assert(adaptor.getOperands().size() >= 3);
+
+  if (getFusedActivationFunction() != "NONE") return {};
+
+  const bool has_bias =
+      !(!getBias() || mlir::isa<NoneType>(getBias().getType()));
+
+  auto operands = adaptor.getOperands();
+  ElementsAttr input_tensor = dyn_cast_or_null<ElementsAttr>(operands[0]);
+  ElementsAttr weights_tensor = dyn_cast_or_null<ElementsAttr>(operands[1]);
+  QConstOp weights_qconst_op = GetWeightsQConstOp(getFilter());
+
+  if (!input_tensor || (!weights_tensor && !weights_qconst_op)) {
+    return {};
+  }
+
+  ElementsAttr bias_tensor;
+  if (has_bias) {
+    bias_tensor = dyn_cast_or_null<ElementsAttr>(operands[2]);
+    if (!bias_tensor) return {};
+  }
+
+  const auto input_type = mlir::dyn_cast<ShapedType>(input_tensor.getType());
+  const auto weights_type = mlir::dyn_cast<ShapedType>(
+      weights_qconst_op ? weights_qconst_op.getOutput().getType()
+                        : getFilter().getType());
+  const auto bias_type = has_bias
+                             ? mlir::dyn_cast<ShapedType>(bias_tensor.getType())
+                             : ShapedType{};
+  const auto output_type = mlir::dyn_cast<ShapedType>(getType());
+
+  if (!input_type || !weights_type || !output_type ||
+      (has_bias && !bias_type)) {
+    return {};
+  }
+
+  if (!input_type.getElementType().isF32() ||
+      !output_type.getElementType().isF32() ||
+      (has_bias && !bias_type.getElementType().isF32())) {
+    return {};
+  }
+
+  if (!input_type.hasStaticShape() || !weights_type.hasStaticShape() ||
+      !output_type.hasStaticShape() ||
+      (has_bias && !bias_type.hasStaticShape())) {
+    return {};
+  }
+
+  if (input_type.getRank() != 4 || weights_type.getRank() != 4 ||
+      output_type.getRank() != 4) {
+    return {};
+  }
+
+  const int64_t batches = input_type.getDimSize(0);
+  const int64_t input_h = input_type.getDimSize(1);
+  const int64_t input_w = input_type.getDimSize(2);
+  const int64_t input_c = input_type.getDimSize(3);
+
+  const int64_t output_depth = weights_type.getDimSize(0);
+  const int64_t filter_h = weights_type.getDimSize(1);
+  const int64_t filter_w = weights_type.getDimSize(2);
+  const int64_t filter_input_c = weights_type.getDimSize(3);
+
+  if (batches <= 0 || input_h <= 0 || input_w <= 0 || input_c <= 0 ||
+      output_depth <= 0 || filter_h <= 0 || filter_w <= 0 ||
+      filter_input_c <= 0) {
+    return {};
+  }
+
+  if (input_c % filter_input_c != 0) return {};
+  const int64_t groups = input_c / filter_input_c;
+  if (output_depth % groups != 0) return {};
+  const int64_t filters_per_group = output_depth / groups;
+
+  const int64_t stride_h = getStrideH();
+  const int64_t stride_w = getStrideW();
+  const int64_t dilation_h = getDilationHFactor();
+  const int64_t dilation_w = getDilationWFactor();
+  if (stride_h <= 0 || stride_w <= 0 || dilation_h <= 0 || dilation_w <= 0) {
+    return {};
+  }
+
+  tensorflow::Padding padding_enum;
+  if (!GetPaddingFromString(getPadding().str(), &padding_enum).ok()) {
+    return {};
+  }
+
+  int64_t out_h = 0, pad_top = 0, pad_bottom = 0;
+  if (!tensorflow::GetWindowedOutputSizeVerbose(input_h, filter_h, dilation_h,
+                                                stride_h, padding_enum, &out_h,
+                                                &pad_top, &pad_bottom)
+           .ok()) {
+    return {};
+  }
+
+  int64_t out_w = 0, pad_left = 0, pad_right = 0;
+  if (!tensorflow::GetWindowedOutputSizeVerbose(input_w, filter_w, dilation_w,
+                                                stride_w, padding_enum, &out_w,
+                                                &pad_left, &pad_right)
+           .ok()) {
+    return {};
+  }
+
+  if (output_type.getDimSize(0) != batches ||
+      output_type.getDimSize(1) != out_h ||
+      output_type.getDimSize(2) != out_w ||
+      output_type.getDimSize(3) != output_depth) {
+    return {};
+  }
+
+  if (has_bias) {
+    if (bias_type.getNumElements() != output_depth) {
+      return {};
+    }
+  }
+
+  if (!IsWithinFoldLimits({batches, out_h, out_w, output_depth, filter_h,
+                           filter_w, filter_input_c},
+                          input_type, weights_type, output_type)) {
+    return {};
+  }
+
+  auto input_values_range = FloatElementsAccessor::Create(input_tensor);
+  if (!input_values_range.has_value()) return {};
+
+  auto weights_values_range =
+      CreateWeightsAccessor(weights_tensor, weights_qconst_op, weights_type);
+  if (!weights_values_range.has_value()) return {};
+
+  std::optional<FloatElementsAccessor> bias_values_range;
+  if (has_bias) {
+    bias_values_range = FloatElementsAccessor::Create(bias_tensor);
+    if (!bias_values_range.has_value()) return {};
+  }
+
+  const int64_t total_outputs = batches * out_h * out_w * output_depth;
+  std::vector<float> result_values;
+  result_values.reserve(total_outputs);
+
+  for (int64_t b = 0; b < batches; ++b) {
+    for (int64_t out_y = 0; out_y < out_h; ++out_y) {
+      const int64_t in_y_origin = out_y * stride_h - pad_top;
+      for (int64_t out_x = 0; out_x < out_w; ++out_x) {
+        const int64_t in_x_origin = out_x * stride_w - pad_left;
+        for (int64_t oc = 0; oc < output_depth; ++oc) {
+          const int64_t group = oc / filters_per_group;
+          float sum = has_bias ? (*bias_values_range)[oc] : 0.0f;
+          float compensation = 0.0f;
+          for (int64_t fy = 0; fy < filter_h; ++fy) {
+            const int64_t in_y = in_y_origin + dilation_h * fy;
+            if (in_y < 0 || in_y >= input_h) continue;
+            for (int64_t fx = 0; fx < filter_w; ++fx) {
+              const int64_t in_x = in_x_origin + dilation_w * fx;
+              if (in_x < 0 || in_x >= input_w) continue;
+              for (int64_t in_c = 0; in_c < filter_input_c; ++in_c) {
+                const int64_t ic = in_c + group * filter_input_c;
+                const int64_t in_idx =
+                    ((b * input_h + in_y) * input_w + in_x) * input_c + ic;
+                const int64_t filter_idx =
+                    ((oc * filter_h + fy) * filter_w + fx) * filter_input_c +
+                    in_c;
+                const float addend = (*input_values_range)[in_idx] *
+                                     (*weights_values_range)[filter_idx];
+                NeumaierAdd(addend, sum, compensation);
+              }
+            }
+          }
+          result_values.push_back(sum + compensation);
+        }
+      }
+    }
+  }
+
+  return DenseElementsAttr::get(output_type, ArrayRef<float>(result_values));
+}
+
 //===----------------------------------------------------------------------===//
 // DepthwiseConv2DOp
 //===----------------------------------------------------------------------===//
@@ -2452,6 +2892,182 @@ int64_t DepthwiseConv2DOp::GetArithmeticCount(Operation* op) {
   return -1;
 }
 
+OpFoldResult DepthwiseConv2DOp::fold(FoldAdaptor adaptor) {
+  assert(adaptor.getOperands().size() >= 3);
+
+  if (getFusedActivationFunction() != "NONE") return {};
+
+  const bool has_bias =
+      !(!getBias() || mlir::isa<NoneType>(getBias().getType()));
+
+  auto operands = adaptor.getOperands();
+  ElementsAttr input_tensor = dyn_cast_or_null<ElementsAttr>(operands[0]);
+  ElementsAttr weights_tensor = dyn_cast_or_null<ElementsAttr>(operands[1]);
+  QConstOp weights_qconst_op = GetWeightsQConstOp(getFilter());
+
+  if (!input_tensor || (!weights_tensor && !weights_qconst_op)) {
+    return {};
+  }
+
+  ElementsAttr bias_tensor;
+  if (has_bias) {
+    bias_tensor = dyn_cast_or_null<ElementsAttr>(operands[2]);
+    if (!bias_tensor) return {};
+  }
+
+  const auto input_type = mlir::dyn_cast<ShapedType>(input_tensor.getType());
+  const auto weights_type = mlir::dyn_cast<ShapedType>(
+      weights_qconst_op ? weights_qconst_op.getOutput().getType()
+                        : getFilter().getType());
+  const auto bias_type = has_bias
+                             ? mlir::dyn_cast<ShapedType>(bias_tensor.getType())
+                             : ShapedType{};
+  const auto output_type = mlir::dyn_cast<ShapedType>(getType());
+
+  if (!input_type || !weights_type || !output_type ||
+      (has_bias && !bias_type)) {
+    return {};
+  }
+
+  if (!input_type.getElementType().isF32() ||
+      !output_type.getElementType().isF32() ||
+      (has_bias && !bias_type.getElementType().isF32())) {
+    return {};
+  }
+
+  if (!input_type.hasStaticShape() || !weights_type.hasStaticShape() ||
+      !output_type.hasStaticShape() ||
+      (has_bias && !bias_type.hasStaticShape())) {
+    return {};
+  }
+
+  if (input_type.getRank() != 4 || weights_type.getRank() != 4 ||
+      output_type.getRank() != 4) {
+    return {};
+  }
+
+  if (weights_type.getDimSize(0) != 1) {
+    return {};
+  }
+
+  const int64_t batches = input_type.getDimSize(0);
+  const int64_t input_h = input_type.getDimSize(1);
+  const int64_t input_w = input_type.getDimSize(2);
+  const int64_t input_c = input_type.getDimSize(3);
+
+  const int64_t filter_h = weights_type.getDimSize(1);
+  const int64_t filter_w = weights_type.getDimSize(2);
+  const int64_t output_depth = weights_type.getDimSize(3);
+  const int64_t depth_multiplier = getDepthMultiplier();
+
+  if (batches <= 0 || input_h <= 0 || input_w <= 0 || input_c <= 0 ||
+      filter_h <= 0 || filter_w <= 0 || output_depth <= 0 ||
+      depth_multiplier <= 0) {
+    return {};
+  }
+
+  if (output_depth != input_c * depth_multiplier) {
+    return {};
+  }
+
+  const int64_t stride_h = getStrideH();
+  const int64_t stride_w = getStrideW();
+  const int64_t dilation_h = getDilationHFactor();
+  const int64_t dilation_w = getDilationWFactor();
+  if (stride_h <= 0 || stride_w <= 0 || dilation_h <= 0 || dilation_w <= 0) {
+    return {};
+  }
+
+  tensorflow::Padding padding_enum;
+  if (!GetPaddingFromString(getPadding().str(), &padding_enum).ok()) {
+    return {};
+  }
+
+  int64_t out_h = 0, pad_top = 0, pad_bottom = 0;
+  if (!tensorflow::GetWindowedOutputSizeVerbose(input_h, filter_h, dilation_h,
+                                                stride_h, padding_enum, &out_h,
+                                                &pad_top, &pad_bottom)
+           .ok()) {
+    return {};
+  }
+
+  int64_t out_w = 0, pad_left = 0, pad_right = 0;
+  if (!tensorflow::GetWindowedOutputSizeVerbose(input_w, filter_w, dilation_w,
+                                                stride_w, padding_enum, &out_w,
+                                                &pad_left, &pad_right)
+           .ok()) {
+    return {};
+  }
+
+  if (output_type.getDimSize(0) != batches ||
+      output_type.getDimSize(1) != out_h ||
+      output_type.getDimSize(2) != out_w ||
+      output_type.getDimSize(3) != output_depth) {
+    return {};
+  }
+
+  if (has_bias) {
+    if (bias_type.getNumElements() != output_depth) {
+      return {};
+    }
+  }
+
+  if (!IsWithinFoldLimits(
+          {batches, out_h, out_w, output_depth, filter_h, filter_w}, input_type,
+          weights_type, output_type)) {
+    return {};
+  }
+
+  auto input_values_range = FloatElementsAccessor::Create(input_tensor);
+  if (!input_values_range.has_value()) return {};
+
+  auto weights_values_range =
+      CreateWeightsAccessor(weights_tensor, weights_qconst_op, weights_type);
+  if (!weights_values_range.has_value()) return {};
+
+  std::optional<FloatElementsAccessor> bias_values_range;
+  if (has_bias) {
+    bias_values_range = FloatElementsAccessor::Create(bias_tensor);
+    if (!bias_values_range.has_value()) return {};
+  }
+
+  const int64_t total_outputs = batches * out_h * out_w * output_depth;
+  std::vector<float> result_values;
+  result_values.reserve(total_outputs);
+
+  for (int64_t b = 0; b < batches; ++b) {
+    for (int64_t out_y = 0; out_y < out_h; ++out_y) {
+      const int64_t in_y_origin = out_y * stride_h - pad_top;
+      for (int64_t out_x = 0; out_x < out_w; ++out_x) {
+        const int64_t in_x_origin = out_x * stride_w - pad_left;
+        for (int64_t oc = 0; oc < output_depth; ++oc) {
+          const int64_t ic = oc / depth_multiplier;
+          float sum = has_bias ? (*bias_values_range)[oc] : 0.0f;
+          float compensation = 0.0f;
+          for (int64_t fy = 0; fy < filter_h; ++fy) {
+            const int64_t in_y = in_y_origin + dilation_h * fy;
+            if (in_y < 0 || in_y >= input_h) continue;
+            for (int64_t fx = 0; fx < filter_w; ++fx) {
+              const int64_t in_x = in_x_origin + dilation_w * fx;
+              if (in_x < 0 || in_x >= input_w) continue;
+              const int64_t in_idx =
+                  ((b * input_h + in_y) * input_w + in_x) * input_c + ic;
+              const int64_t filter_idx =
+                  (fy * filter_w + fx) * output_depth + oc;
+              const float addend = (*input_values_range)[in_idx] *
+                                   (*weights_values_range)[filter_idx];
+              NeumaierAdd(addend, sum, compensation);
+            }
+          }
+          result_values.push_back(sum + compensation);
+        }
+      }
+    }
+  }
+
+  return DenseElementsAttr::get(output_type, ArrayRef<float>(result_values));
+}
+
 //===----------------------------------------------------------------------===//
 // GatherOp
 //===----------------------------------------------------------------------===//
@@ -5171,6 +5787,190 @@ int64_t TransposeConvOp::GetArithmeticCount(Operation* op) {
   return count;
 }
 
+OpFoldResult TransposeConvOp::fold(FoldAdaptor adaptor) {
+  assert(adaptor.getOperands().size() >= 3);
+
+  if (getFusedActivationFunction() != "NONE") return {};
+
+  const bool has_bias =
+      !(!getBias() || mlir::isa<NoneType>(getBias().getType()));
+
+  auto operands = adaptor.getOperands();
+  ElementsAttr weights_tensor = dyn_cast_or_null<ElementsAttr>(operands[1]);
+  ElementsAttr input_tensor = dyn_cast_or_null<ElementsAttr>(operands[2]);
+  QConstOp weights_qconst_op = GetWeightsQConstOp(getWeights());
+
+  if (!input_tensor || (!weights_tensor && !weights_qconst_op)) {
+    return {};
+  }
+
+  ElementsAttr bias_tensor;
+  if (has_bias) {
+    if (operands.size() < 4) return {};
+    bias_tensor = dyn_cast_or_null<ElementsAttr>(operands[3]);
+    if (!bias_tensor) return {};
+  }
+
+  const auto input_type = mlir::dyn_cast<ShapedType>(input_tensor.getType());
+  const auto weights_type = mlir::dyn_cast<ShapedType>(
+      weights_qconst_op ? weights_qconst_op.getOutput().getType()
+                        : getWeights().getType());
+  const auto bias_type = has_bias
+                             ? mlir::dyn_cast<ShapedType>(bias_tensor.getType())
+                             : ShapedType{};
+  const auto output_type = mlir::dyn_cast<ShapedType>(getType());
+
+  if (!input_type || !weights_type || !output_type ||
+      (has_bias && !bias_type)) {
+    return {};
+  }
+
+  if (!input_type.getElementType().isF32() ||
+      !output_type.getElementType().isF32() ||
+      (has_bias && !bias_type.getElementType().isF32())) {
+    return {};
+  }
+
+  if (!input_type.hasStaticShape() || !weights_type.hasStaticShape() ||
+      !output_type.hasStaticShape() ||
+      (has_bias && !bias_type.hasStaticShape())) {
+    return {};
+  }
+
+  if (input_type.getRank() != 4 || weights_type.getRank() != 4 ||
+      output_type.getRank() != 4) {
+    return {};
+  }
+
+  const int64_t batches = input_type.getDimSize(0);
+  const int64_t input_h = input_type.getDimSize(1);
+  const int64_t input_w = input_type.getDimSize(2);
+  const int64_t input_c = input_type.getDimSize(3);
+
+  const int64_t output_depth = weights_type.getDimSize(0);
+  const int64_t filter_h = weights_type.getDimSize(1);
+  const int64_t filter_w = weights_type.getDimSize(2);
+  const int64_t filter_input_c = weights_type.getDimSize(3);
+
+  const int64_t out_batches = output_type.getDimSize(0);
+  const int64_t out_h = output_type.getDimSize(1);
+  const int64_t out_w = output_type.getDimSize(2);
+  const int64_t out_c = output_type.getDimSize(3);
+
+  if (batches <= 0 || input_h <= 0 || input_w <= 0 || input_c <= 0 ||
+      output_depth <= 0 || filter_h <= 0 || filter_w <= 0 ||
+      filter_input_c <= 0 || out_batches <= 0 || out_h <= 0 || out_w <= 0 ||
+      out_c <= 0) {
+    return {};
+  }
+
+  if (batches != out_batches || output_depth != out_c ||
+      input_c != filter_input_c) {
+    return {};
+  }
+
+  const int64_t stride_h = getStrideH();
+  const int64_t stride_w = getStrideW();
+  if (stride_h <= 0 || stride_w <= 0) {
+    return {};
+  }
+
+  tensorflow::Padding padding_enum;
+  if (!GetPaddingFromString(getPadding().str(), &padding_enum).ok()) {
+    return {};
+  }
+
+  int64_t computed_in_h = 0, pad_top = 0, pad_bottom = 0;
+  if (!tensorflow::GetWindowedOutputSizeVerbose(
+           out_h, filter_h, /*dilation_rate=*/1, stride_h, padding_enum,
+           &computed_in_h, &pad_top, &pad_bottom)
+           .ok()) {
+    return {};
+  }
+  if (computed_in_h != input_h) return {};
+
+  int64_t computed_in_w = 0, pad_left = 0, pad_right = 0;
+  if (!tensorflow::GetWindowedOutputSizeVerbose(
+           out_w, filter_w, /*dilation_rate=*/1, stride_w, padding_enum,
+           &computed_in_w, &pad_left, &pad_right)
+           .ok()) {
+    return {};
+  }
+  if (computed_in_w != input_w) return {};
+
+  if (has_bias) {
+    if (bias_type.getNumElements() != output_depth) {
+      return {};
+    }
+  }
+
+  if (!IsWithinFoldLimits({batches, input_h, input_w, input_c, filter_h,
+                           filter_w, output_depth},
+                          input_type, weights_type, output_type)) {
+    return {};
+  }
+
+  auto input_values_range = FloatElementsAccessor::Create(input_tensor);
+  if (!input_values_range.has_value()) return {};
+
+  auto weights_values_range =
+      CreateWeightsAccessor(weights_tensor, weights_qconst_op, weights_type);
+  if (!weights_values_range.has_value()) return {};
+
+  std::optional<FloatElementsAccessor> bias_values_range;
+  if (has_bias) {
+    bias_values_range = FloatElementsAccessor::Create(bias_tensor);
+    if (!bias_values_range.has_value()) return {};
+  }
+
+  const int64_t total_outputs = batches * out_h * out_w * output_depth;
+  std::vector<float> result_values(total_outputs, 0.0f);
+  std::vector<float> compensations(total_outputs, 0.0f);
+  if (has_bias) {
+    for (int64_t i = 0; i < total_outputs; ++i) {
+      result_values[i] = (*bias_values_range)[i % output_depth];
+    }
+  }
+
+  for (int64_t b = 0; b < batches; ++b) {
+    for (int64_t in_y = 0; in_y < input_h; ++in_y) {
+      const int64_t out_y_origin = in_y * stride_h - pad_top;
+      for (int64_t in_x = 0; in_x < input_w; ++in_x) {
+        const int64_t out_x_origin = in_x * stride_w - pad_left;
+        for (int64_t in_c = 0; in_c < input_c; ++in_c) {
+          const int64_t in_idx =
+              ((b * input_h + in_y) * input_w + in_x) * input_c + in_c;
+          const float in_val = (*input_values_range)[in_idx];
+          for (int64_t fy = 0; fy < filter_h; ++fy) {
+            const int64_t out_y = out_y_origin + fy;
+            if (out_y < 0 || out_y >= out_h) continue;
+            for (int64_t fx = 0; fx < filter_w; ++fx) {
+              const int64_t out_x = out_x_origin + fx;
+              if (out_x < 0 || out_x >= out_w) continue;
+              for (int64_t oc = 0; oc < output_depth; ++oc) {
+                const int64_t filter_idx =
+                    ((oc * filter_h + fy) * filter_w + fx) * input_c + in_c;
+                const int64_t out_idx =
+                    ((b * out_h + out_y) * out_w + out_x) * output_depth + oc;
+                const float addend =
+                    in_val * (*weights_values_range)[filter_idx];
+                NeumaierAdd(addend, result_values[out_idx],
+                            compensations[out_idx]);
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+
+  for (int64_t i = 0; i < total_outputs; ++i) {
+    result_values[i] += compensations[i];
+  }
+
+  return DenseElementsAttr::get(output_type, ArrayRef<float>(result_values));
+}
+
 //===----------------------------------------------------------------------===//
 // StridedSliceOp
 //===----------------------------------------------------------------------===//
diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td
index 7a2778dc978..2d08b9af4bc 100644
--- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td
+++ b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td
@@ -744,6 +744,8 @@ def TFL_TransposeConvOp: TFL_Op<"transpose_conv", [
 
   let hasVerifier = 1;
 
+  let hasFolder = 1;
+
   let extraClassDeclaration = [{
     // AffineQuantizedOpInterface:
     int GetChannelDimIndex() { return 0; }
@@ -1018,6 +1020,7 @@ def TFL_Conv2DOp : TFL_ConvOp<"conv_2d", "Convolution", 0,
   );
 
   let hasCanonicalizer = 1;
+  let hasFolder = 1;
 
   let extraClassDeclaration = [{
     // AffineQuantizedOpInterface:
@@ -1108,6 +1111,7 @@ def TFL_DepthwiseConv2DOp :
   );
 
   let hasCanonicalizer = 1;
+  let hasFolder = 1;
 
   let extraClassDeclaration = [{
     // AffineQuantizedOpInterface:
diff --git a/tensorflow/compiler/mlir/lite/tests/const-fold.mlir b/tensorflow/compiler/mlir/lite/tests/const-fold.mlir
index dcdda38f84b..4765f783c32 100644
--- a/tensorflow/compiler/mlir/lite/tests/const-fold.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/const-fold.mlir
@@ -485,6 +485,439 @@ func.func @ConstantFoldFullyConnectedBatched() -> tensor<13x1536xf32> {
   // CHECK:  return %[[CST]]
 }
 
+// CHECK-LABEL: @ConstantFoldFullyConnectedDRQW8A32PerTensor
+func.func @ConstantFoldFullyConnectedDRQW8A32PerTensor() -> tensor<1x3xf32> {
+  %cst_input = arith.constant dense<[[2.0, 3.0]]> : tensor<1x2xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<3x2x!quant.uniform<i8:f32, 0.5:0>>, value = dense<[[2, 4], [-6, 8], [10, -12]]> : tensor<3x2xi8>}> : () -> tensor<3x2x!quant.uniform<i8:f32, 0.5:0>>
+  %cst_bias = arith.constant dense<[1.0, 2.0, 3.0]> : tensor<3xf32>
+
+  %0 = "tfl.fully_connected"(%cst_input, %cst_weights, %cst_bias) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<1x2xf32>, tensor<3x2x!quant.uniform<i8:f32, 0.5:0>>, tensor<3xf32>) -> tensor<1x3xf32>
+  func.return %0 : tensor<1x3xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[}}9.000000e+00, 8.000000e+00, -5.000000e+00{{\]\]}}> : tensor<1x3xf32>
+  // CHECK-NOT: fully_connected
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldFullyConnectedDRQW8A32PerAxis
+func.func @ConstantFoldFullyConnectedDRQW8A32PerAxis() -> tensor<1x3xf32> {
+  %cst_input = arith.constant dense<[[2.0, 3.0]]> : tensor<1x2xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<3x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0, 0.25}>>, value = dense<[[2, 4], [-1, 3], [8, -4]]> : tensor<3x2xi8>}> : () -> tensor<3x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0, 0.25}>>
+  %cst_bias = arith.constant dense<[1.0, 2.0, 3.0]> : tensor<3xf32>
+
+  %0 = "tfl.fully_connected"(%cst_input, %cst_weights, %cst_bias) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<1x2xf32>, tensor<3x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0, 0.25}>>, tensor<3xf32>) -> tensor<1x3xf32>
+  func.return %0 : tensor<1x3xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[}}9.000000e+00, 1.600000e+01, 4.000000e+00{{\]\]}}> : tensor<1x3xf32>
+  // CHECK-NOT: fully_connected
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldFullyConnectedDRQW4A32PerTensor
+func.func @ConstantFoldFullyConnectedDRQW4A32PerTensor() -> tensor<1x2xf32> {
+  %cst_input = arith.constant dense<[[2.0, 3.0]]> : tensor<1x2xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x!quant.uniform<i4:f32, 0.5:0>>, value = dense<[[2, -4], [6, -2]]> : tensor<2x2xi4>}> : () -> tensor<2x2x!quant.uniform<i4:f32, 0.5:0>>
+  %cst_bias = "tfl.no_value"() {value = unit} : () -> none
+
+  %0 = "tfl.fully_connected"(%cst_input, %cst_weights, %cst_bias) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<1x2xf32>, tensor<2x2x!quant.uniform<i4:f32, 0.5:0>>, none) -> tensor<1x2xf32>
+  func.return %0 : tensor<1x2xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[}}-4.000000e+00, 3.000000e+00{{\]\]}}> : tensor<1x2xf32>
+  // CHECK-NOT: fully_connected
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldFullyConnectedDRQW4A32PerAxis
+func.func @ConstantFoldFullyConnectedDRQW4A32PerAxis() -> tensor<1x2xf32> {
+  %cst_input = arith.constant dense<[[2.0, 3.0]]> : tensor<1x2xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>, value = dense<[[2, 4], [-3, 1]]> : tensor<2x2xi4>}> : () -> tensor<2x2x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.fully_connected"(%cst_input, %cst_weights, %cst_bias) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<1x2xf32>, tensor<2x2x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x2xf32>
+  func.return %0 : tensor<1x2xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[}}9.000000e+00, -4.000000e+00{{\]\]}}> : tensor<1x2xf32>
+  // CHECK-NOT: fully_connected
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldFullyConnectedWeightOnlyW8A32PerAxis
+func.func @ConstantFoldFullyConnectedWeightOnlyW8A32PerAxis() -> tensor<1x3xf32> {
+  %cst_input = arith.constant dense<[[2.0, 3.0]]> : tensor<1x2xf32>
+  %cst_qweights = "tfl.pseudo_qconst"() <{qtype = tensor<3x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0, 0.25}>>, value = dense<[[2, 4], [-1, 3], [8, -4]]> : tensor<3x2xi8>}> : () -> tensor<3x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0, 0.25}>>
+  %cst_weights = "tfl.dequantize"(%cst_qweights) : (tensor<3x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0, 0.25}>>) -> tensor<3x2xf32>
+  %cst_bias = arith.constant dense<[1.0, 2.0, 3.0]> : tensor<3xf32>
+
+  %0 = "tfl.fully_connected"(%cst_input, %cst_weights, %cst_bias) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<1x2xf32>, tensor<3x2xf32>, tensor<3xf32>) -> tensor<1x3xf32>
+  func.return %0 : tensor<1x3xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[}}9.000000e+00, 1.600000e+01, 4.000000e+00{{\]\]}}> : tensor<1x3xf32>
+  // CHECK-NOT: fully_connected
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldConv2DDRQW8A32PerAxis
+func.func @ConstantFoldConv2DDRQW8A32PerAxis() -> tensor<1x1x1x2xf32> {
+  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2], [4]], [[6], [8]]], [[[1], [-1]], [[2], [-2]]]]> : tensor<2x2x2x1xi8>}> : () -> tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.conv_2d"(%cst_input, %cst_weights, %cst_bias) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
+  func.return %0 : tensor<1x1x1x2xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}3.100000e+01, -4.000000e+00{{\]\]\]\]}}> : tensor<1x1x1x2xf32>
+  // CHECK-NOT: conv_2d
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldConv2DDRQW4A32PerAxis
+func.func @ConstantFoldConv2DDRQW4A32PerAxis() -> tensor<1x1x1x2xf32> {
+  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x2x1x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2], [4]], [[6], [-2]]], [[[1], [-1]], [[3], [-2]]]]> : tensor<2x2x2x1xi4>}> : () -> tensor<2x2x2x1x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>
+  %cst_bias = arith.constant dense<[5.0, -3.0]> : tensor<2xf32>
+
+  %0 = "tfl.conv_2d"(%cst_input, %cst_weights, %cst_bias) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<2x2x2x1x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
+  func.return %0 : tensor<1x1x1x2xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}1.500000e+01, -3.000000e+00{{\]\]\]\]}}> : tensor<1x1x1x2xf32>
+  // CHECK-NOT: conv_2d
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldConv2DWeightOnlyW8A32PerAxis
+func.func @ConstantFoldConv2DWeightOnlyW8A32PerAxis() -> tensor<1x1x1x2xf32> {
+  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
+  %cst_qweights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2], [4]], [[6], [8]]], [[[1], [-1]], [[2], [-2]]]]> : tensor<2x2x2x1xi8>}> : () -> tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>
+  %cst_weights = "tfl.dequantize"(%cst_qweights) : (tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>) -> tensor<2x2x2x1xf32>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.conv_2d"(%cst_input, %cst_weights, %cst_bias) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<2x2x2x1xf32>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
+  func.return %0 : tensor<1x1x1x2xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}3.100000e+01, -4.000000e+00{{\]\]\]\]}}> : tensor<1x1x1x2xf32>
+  // CHECK-NOT: conv_2d
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @NoConstantFoldConv2DFusedRelu
+func.func @NoConstantFoldConv2DFusedRelu() -> tensor<1x1x1x2xf32> {
+  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2], [4]], [[6], [8]]], [[[1], [-1]], [[2], [-2]]]]> : tensor<2x2x2x1xi8>}> : () -> tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.conv_2d"(%cst_input, %cst_weights, %cst_bias) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "RELU", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
+  func.return %0 : tensor<1x1x1x2xf32>
+
+  // CHECK: tfl.conv_2d
+}
+
+// CHECK-LABEL: @ConstantFoldDepthwiseConv2DDRQW8A32PerAxis
+func.func @ConstantFoldDepthwiseConv2DDRQW8A32PerAxis() -> tensor<1x1x1x2xf32> {
+  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>, value = dense<[[[[2, 1], [4, -1]], [[6, 2], [8, -2]]]]> : tensor<1x2x2x2xi8>}> : () -> tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.depthwise_conv_2d"(%cst_input, %cst_weights, %cst_bias) <{depth_multiplier = 2 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
+  func.return %0 : tensor<1x1x1x2xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}3.100000e+01, -4.000000e+00{{\]\]\]\]}}> : tensor<1x1x1x2xf32>
+  // CHECK-NOT: depthwise_conv_2d
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldDepthwiseConv2DDRQW4A32PerAxis
+func.func @ConstantFoldDepthwiseConv2DDRQW4A32PerAxis() -> tensor<1x1x1x2xf32> {
+  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<1x2x2x2x!quant.uniform<i4<-7:7>:f32:3, {0.5, 2.0}>>, value = dense<[[[[2, 1], [4, -1]], [[6, 3], [-2, -2]]]]> : tensor<1x2x2x2xi4>}> : () -> tensor<1x2x2x2x!quant.uniform<i4<-7:7>:f32:3, {0.5, 2.0}>>
+  %cst_bias = arith.constant dense<[5.0, -3.0]> : tensor<2xf32>
+
+  %0 = "tfl.depthwise_conv_2d"(%cst_input, %cst_weights, %cst_bias) <{depth_multiplier = 2 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<1x2x2x2x!quant.uniform<i4<-7:7>:f32:3, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
+  func.return %0 : tensor<1x1x1x2xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}1.500000e+01, -3.000000e+00{{\]\]\]\]}}> : tensor<1x1x1x2xf32>
+  // CHECK-NOT: depthwise_conv_2d
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldDepthwiseConv2DWeightOnlyW8A32PerAxis
+func.func @ConstantFoldDepthwiseConv2DWeightOnlyW8A32PerAxis() -> tensor<1x1x1x2xf32> {
+  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
+  %cst_qweights = "tfl.pseudo_qconst"() <{qtype = tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>, value = dense<[[[[2, 1], [4, -1]], [[6, 2], [8, -2]]]]> : tensor<1x2x2x2xi8>}> : () -> tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>
+  %cst_weights = "tfl.dequantize"(%cst_qweights) : (tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>) -> tensor<1x2x2x2xf32>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.depthwise_conv_2d"(%cst_input, %cst_weights, %cst_bias) <{depth_multiplier = 2 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<1x2x2x2xf32>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
+  func.return %0 : tensor<1x1x1x2xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}3.100000e+01, -4.000000e+00{{\]\]\]\]}}> : tensor<1x1x1x2xf32>
+  // CHECK-NOT: depthwise_conv_2d
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldTransposeConvDRQW8A32PerAxis
+func.func @ConstantFoldTransposeConvDRQW8A32PerAxis() -> tensor<1x2x2x2xf32> {
+  %cst_shape = arith.constant dense<[1, 2, 2, 2]> : tensor<4xi32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2], [4]], [[6], [8]]], [[[1], [-1]], [[2], [-2]]]]> : tensor<2x2x2x1xi8>}> : () -> tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>
+  %cst_input = arith.constant dense<[[[[2.0]]]]> : tensor<1x1x1x1xf32>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.transpose_conv"(%cst_shape, %cst_weights, %cst_input, %cst_bias) <{fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<4xi32>, tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, tensor<1x1x1x1xf32>, tensor<2xf32>) -> tensor<1x2x2x2xf32>
+  func.return %0 : tensor<1x2x2x2xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}3.000000e+00, 6.000000e+00], [5.000000e+00, -2.000000e+00]], {{\[\[}}7.000000e+00, 1.000000e+01], [9.000000e+00, -6.000000e+00{{\]\]\]\]}}> : tensor<1x2x2x2xf32>
+  // CHECK-NOT: transpose_conv
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldTransposeConvDRQW4A32PerAxis
+func.func @ConstantFoldTransposeConvDRQW4A32PerAxis() -> tensor<1x2x2x2xf32> {
+  %cst_shape = arith.constant dense<[1, 2, 2, 2]> : tensor<4xi32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x2x1x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2], [4]], [[6], [-2]]], [[[1], [-1]], [[3], [-2]]]]> : tensor<2x2x2x1xi4>}> : () -> tensor<2x2x2x1x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>
+  %cst_input = arith.constant dense<[[[[2.0]]]]> : tensor<1x1x1x1xf32>
+  %cst_bias = arith.constant dense<[5.0, -3.0]> : tensor<2xf32>
+
+  %0 = "tfl.transpose_conv"(%cst_shape, %cst_weights, %cst_input, %cst_bias) <{fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<4xi32>, tensor<2x2x2x1x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>, tensor<1x1x1x1xf32>, tensor<2xf32>) -> tensor<1x2x2x2xf32>
+  func.return %0 : tensor<1x2x2x2xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}7.000000e+00, 1.000000e+00], [9.000000e+00, -7.000000e+00]], {{\[\[}}1.100000e+01, 9.000000e+00], [3.000000e+00, -1.100000e+01{{\]\]\]\]}}> : tensor<1x2x2x2xf32>
+  // CHECK-NOT: transpose_conv
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldTransposeConvWeightOnlyW8A32PerAxis
+func.func @ConstantFoldTransposeConvWeightOnlyW8A32PerAxis() -> tensor<1x2x2x2xf32> {
+  %cst_shape = arith.constant dense<[1, 2, 2, 2]> : tensor<4xi32>
+  %cst_qweights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2], [4]], [[6], [8]]], [[[1], [-1]], [[2], [-2]]]]> : tensor<2x2x2x1xi8>}> : () -> tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>
+  %cst_weights = "tfl.dequantize"(%cst_qweights) : (tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>) -> tensor<2x2x2x1xf32>
+  %cst_input = arith.constant dense<[[[[2.0]]]]> : tensor<1x1x1x1xf32>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.transpose_conv"(%cst_shape, %cst_weights, %cst_input, %cst_bias) <{fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<4xi32>, tensor<2x2x2x1xf32>, tensor<1x1x1x1xf32>, tensor<2xf32>) -> tensor<1x2x2x2xf32>
+  func.return %0 : tensor<1x2x2x2xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}3.000000e+00, 6.000000e+00], [5.000000e+00, -2.000000e+00]], {{\[\[}}7.000000e+00, 1.000000e+01], [9.000000e+00, -6.000000e+00{{\]\]\]\]}}> : tensor<1x2x2x2xf32>
+  // CHECK-NOT: transpose_conv
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldFullyConnectedDRQU8A32NonZeroPoint
+func.func @ConstantFoldFullyConnectedDRQU8A32NonZeroPoint() -> tensor<1x2xf32> {
+  %cst_input = arith.constant dense<[[2.0, 3.0]]> : tensor<1x2xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x!quant.uniform<u8:f32:0, {0.5:128, 2.0:130}>>, value = dense<[[-126, -124], [-127, -125]]> : tensor<2x2xi8>}> : () -> tensor<2x2x!quant.uniform<u8:f32:0, {0.5:128, 2.0:130}>>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.fully_connected"(%cst_input, %cst_weights, %cst_bias) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<1x2xf32>, tensor<2x2x!quant.uniform<u8:f32:0, {0.5:128, 2.0:130}>>, tensor<2xf32>) -> tensor<1x2xf32>
+  func.return %0 : tensor<1x2xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[}}9.000000e+00, 4.000000e+00{{\]\]}}> : tensor<1x2xf32>
+  // CHECK-NOT: fully_connected
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldFullyConnectedDRQW8A32DenseResource
+func.func @ConstantFoldFullyConnectedDRQW8A32DenseResource() -> tensor<1x2xf32> {
+  %cst_input = arith.constant dense<[[2.0, 3.0]]> : tensor<1x2xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, value = dense_resource<fc_w8_2x2> : tensor<2x2xi8>}> : () -> tensor<2x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.fully_connected"(%cst_input, %cst_weights, %cst_bias) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<1x2xf32>, tensor<2x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x2xf32>
+  func.return %0 : tensor<1x2xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[}}9.000000e+00, 1.600000e+01{{\]\]}}> : tensor<1x2xf32>
+  // CHECK-NOT: fully_connected
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldFullyConnectedDRQW4A32DenseResource
+func.func @ConstantFoldFullyConnectedDRQW4A32DenseResource() -> tensor<1x2xf32> {
+  %cst_input = arith.constant dense<[[2.0, 3.0]]> : tensor<1x2xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>, value = dense_resource<fc_w4_2x2> : tensor<2x2xi4>}> : () -> tensor<2x2x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.fully_connected"(%cst_input, %cst_weights, %cst_bias) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<1x2xf32>, tensor<2x2x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x2xf32>
+  func.return %0 : tensor<1x2xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[}}9.000000e+00, -4.000000e+00{{\]\]}}> : tensor<1x2xf32>
+  // CHECK-NOT: fully_connected
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldConv2DGrouped
+func.func @ConstantFoldConv2DGrouped() -> tensor<1x1x1x2xf32> {
+  %cst_input = arith.constant dense<[[[[1.0, 2.0, 3.0, 4.0]]]]> : tensor<1x1x1x4xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x1x1x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2, 4]]], [[[1, -1]]]]> : tensor<2x1x1x2xi8>}> : () -> tensor<2x1x1x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.conv_2d"(%cst_input, %cst_weights, %cst_bias) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x1x1x4xf32>, tensor<2x1x1x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
+  func.return %0 : tensor<1x1x1x2xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}6.000000e+00, 0.000000e+00{{\]\]\]\]}}> : tensor<1x1x1x2xf32>
+  // CHECK-NOT: conv_2d
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldConv2DDilated
+func.func @ConstantFoldConv2DDilated() -> tensor<1x1x1x1xf32> {
+  %cst_input = arith.constant dense<[[[[1.0], [0.0], [2.0]], [[0.0], [0.0], [0.0]], [[3.0], [0.0], [4.0]]]]> : tensor<1x3x3x1xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<1x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5}>>, value = dense<[[[[2], [4]], [[6], [8]]]]> : tensor<1x2x2x1xi8>}> : () -> tensor<1x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5}>>
+  %cst_bias = arith.constant dense<[1.0]> : tensor<1xf32>
+
+  %0 = "tfl.conv_2d"(%cst_input, %cst_weights, %cst_bias) <{dilation_h_factor = 2 : i32, dilation_w_factor = 2 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x3x3x1xf32>, tensor<1x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5}>>, tensor<1xf32>) -> tensor<1x1x1x1xf32>
+  func.return %0 : tensor<1x1x1x1xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<3.100000e+01> : tensor<1x1x1x1xf32>
+  // CHECK-NOT: conv_2d
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldConv2DF32
+func.func @ConstantFoldConv2DF32() -> tensor<1x1x1x2xf32> {
+  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
+  %cst_weights = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]], [[[2.0], [-2.0]], [[4.0], [-4.0]]]]> : tensor<2x2x2x1xf32>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.conv_2d"(%cst_input, %cst_weights, %cst_bias) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<2x2x2x1xf32>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
+  func.return %0 : tensor<1x1x1x2xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}3.100000e+01, -4.000000e+00{{\]\]\]\]}}> : tensor<1x1x1x2xf32>
+  // CHECK-NOT: conv_2d
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldDepthwiseConv2DF32
+func.func @ConstantFoldDepthwiseConv2DF32() -> tensor<1x1x1x2xf32> {
+  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
+  %cst_weights = arith.constant dense<[[[[1.0, 2.0], [2.0, -2.0]], [[3.0, 4.0], [4.0, -4.0]]]]> : tensor<1x2x2x2xf32>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.depthwise_conv_2d"(%cst_input, %cst_weights, %cst_bias) <{depth_multiplier = 2 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<1x2x2x2xf32>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
+  func.return %0 : tensor<1x1x1x2xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}3.100000e+01, -4.000000e+00{{\]\]\]\]}}> : tensor<1x1x1x2xf32>
+  // CHECK-NOT: depthwise_conv_2d
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldTransposeConvF32
+func.func @ConstantFoldTransposeConvF32() -> tensor<1x2x2x2xf32> {
+  %cst_shape = arith.constant dense<[1, 2, 2, 2]> : tensor<4xi32>
+  %cst_weights = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]], [[[2.0], [-2.0]], [[4.0], [-4.0]]]]> : tensor<2x2x2x1xf32>
+  %cst_input = arith.constant dense<[[[[2.0]]]]> : tensor<1x1x1x1xf32>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.transpose_conv"(%cst_shape, %cst_weights, %cst_input, %cst_bias) <{fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<4xi32>, tensor<2x2x2x1xf32>, tensor<1x1x1x1xf32>, tensor<2xf32>) -> tensor<1x2x2x2xf32>
+  func.return %0 : tensor<1x2x2x2xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}3.000000e+00, 6.000000e+00], [5.000000e+00, -2.000000e+00]], {{\[\[}}7.000000e+00, 1.000000e+01], [9.000000e+00, -6.000000e+00{{\]\]\]\]}}> : tensor<1x2x2x2xf32>
+  // CHECK-NOT: transpose_conv
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldConv2DSameAsymmetricPadStride2
+func.func @ConstantFoldConv2DSameAsymmetricPadStride2() -> tensor<1x2x2x1xf32> {
+  %cst_input = arith.constant dense<1.0> : tensor<1x3x3x1xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<1x4x4x1x!quant.uniform<i8<-127:127>:f32:0, {0.5}>>, value = dense<[[[[1], [1], [1], [1]], [[2], [2], [2], [2]], [[3], [3], [3], [3]], [[4], [4], [4], [4]]]]> : tensor<1x4x4x1xi8>}> : () -> tensor<1x4x4x1x!quant.uniform<i8<-127:127>:f32:0, {0.5}>>
+  %cst_bias = arith.constant dense<[1.0]> : tensor<1xf32>
+
+  %0 = "tfl.conv_2d"(%cst_input, %cst_weights, %cst_bias) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<1x3x3x1xf32>, tensor<1x4x4x1x!quant.uniform<i8<-127:127>:f32:0, {0.5}>>, tensor<1xf32>) -> tensor<1x2x2x1xf32>
+  func.return %0 : tensor<1x2x2x1xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}1.450000e+01], [1.000000e+01]], {{\[\[}}5.500000e+00], [4.000000e+00{{\]\]\]\]}}> : tensor<1x2x2x1xf32>
+  // CHECK-NOT: conv_2d
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldDepthwiseConv2DSameAsymmetricPadStride2
+func.func @ConstantFoldDepthwiseConv2DSameAsymmetricPadStride2() -> tensor<1x2x2x1xf32> {
+  %cst_input = arith.constant dense<1.0> : tensor<1x3x3x1xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<1x4x4x1x!quant.uniform<i8<-127:127>:f32:3, {0.5}>>, value = dense<[[[[1], [1], [1], [1]], [[2], [2], [2], [2]], [[3], [3], [3], [3]], [[4], [4], [4], [4]]]]> : tensor<1x4x4x1xi8>}> : () -> tensor<1x4x4x1x!quant.uniform<i8<-127:127>:f32:3, {0.5}>>
+  %cst_bias = arith.constant dense<[1.0]> : tensor<1xf32>
+
+  %0 = "tfl.depthwise_conv_2d"(%cst_input, %cst_weights, %cst_bias) <{depth_multiplier = 1 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<1x3x3x1xf32>, tensor<1x4x4x1x!quant.uniform<i8<-127:127>:f32:3, {0.5}>>, tensor<1xf32>) -> tensor<1x2x2x1xf32>
+  func.return %0 : tensor<1x2x2x1xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}1.450000e+01], [1.000000e+01]], {{\[\[}}5.500000e+00], [4.000000e+00{{\]\]\]\]}}> : tensor<1x2x2x1xf32>
+  // CHECK-NOT: depthwise_conv_2d
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @ConstantFoldTransposeConvSameNonZeroPad
+func.func @ConstantFoldTransposeConvSameNonZeroPad() -> tensor<1x3x3x1xf32> {
+  %cst_shape = arith.constant dense<[1, 3, 3, 1]> : tensor<4xi32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<1x3x3x1x!quant.uniform<i8<-127:127>:f32:0, {0.5}>>, value = dense<2> : tensor<1x3x3x1xi8>}> : () -> tensor<1x3x3x1x!quant.uniform<i8<-127:127>:f32:0, {0.5}>>
+  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
+  %cst_bias = arith.constant dense<[0.0]> : tensor<1xf32>
+
+  %0 = "tfl.transpose_conv"(%cst_shape, %cst_weights, %cst_input, %cst_bias) <{fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<4xi32>, tensor<1x3x3x1x!quant.uniform<i8<-127:127>:f32:0, {0.5}>>, tensor<1x2x2x1xf32>, tensor<1xf32>) -> tensor<1x3x3x1xf32>
+  func.return %0 : tensor<1x3x3x1xf32>
+
+  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}1.000000e+00], [3.000000e+00], [2.000000e+00]], {{\[\[}}4.000000e+00], [1.000000e+01], [6.000000e+00]], {{\[\[}}3.000000e+00], [7.000000e+00], [4.000000e+00{{\]\]\]\]}}> : tensor<1x3x3x1xf32>
+  // CHECK-NOT: transpose_conv
+  // CHECK: return %[[CST]]
+}
+
+// CHECK-LABEL: @NoConstantFoldDepthwiseConv2DFusedRelu
+func.func @NoConstantFoldDepthwiseConv2DFusedRelu() -> tensor<1x1x1x2xf32> {
+  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>, value = dense<[[[[2, 1], [4, -1]], [[6, 2], [8, -2]]]]> : tensor<1x2x2x2xi8>}> : () -> tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.depthwise_conv_2d"(%cst_input, %cst_weights, %cst_bias) <{depth_multiplier = 2 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "RELU", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
+  func.return %0 : tensor<1x1x1x2xf32>
+
+  // CHECK: tfl.depthwise_conv_2d
+}
+
+// CHECK-LABEL: @NoConstantFoldDepthwiseConv2DDynamicShape
+func.func @NoConstantFoldDepthwiseConv2DDynamicShape() -> tensor<1x?x?x2xf32> {
+  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>, value = dense<[[[[2, 1], [4, -1]], [[6, 2], [8, -2]]]]> : tensor<1x2x2x2xi8>}> : () -> tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.depthwise_conv_2d"(%cst_input, %cst_weights, %cst_bias) <{depth_multiplier = 2 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x?x?x2xf32>
+  func.return %0 : tensor<1x?x?x2xf32>
+
+  // CHECK: tfl.depthwise_conv_2d
+}
+
+// CHECK-LABEL: @NoConstantFoldTransposeConvFusedRelu
+func.func @NoConstantFoldTransposeConvFusedRelu() -> tensor<1x2x2x2xf32> {
+  %cst_shape = arith.constant dense<[1, 2, 2, 2]> : tensor<4xi32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2], [4]], [[6], [8]]], [[[1], [-1]], [[2], [-2]]]]> : tensor<2x2x2x1xi8>}> : () -> tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>
+  %cst_input = arith.constant dense<[[[[2.0]]]]> : tensor<1x1x1x1xf32>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.transpose_conv"(%cst_shape, %cst_weights, %cst_input, %cst_bias) <{fused_activation_function = "RELU", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<4xi32>, tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, tensor<1x1x1x1xf32>, tensor<2xf32>) -> tensor<1x2x2x2xf32>
+  func.return %0 : tensor<1x2x2x2xf32>
+
+  // CHECK: tfl.transpose_conv
+}
+
+// CHECK-LABEL: @NoConstantFoldTransposeConvDynamicShape
+func.func @NoConstantFoldTransposeConvDynamicShape() -> tensor<1x?x?x2xf32> {
+  %cst_shape = arith.constant dense<[1, 2, 2, 2]> : tensor<4xi32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2], [4]], [[6], [8]]], [[[1], [-1]], [[2], [-2]]]]> : tensor<2x2x2x1xi8>}> : () -> tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>
+  %cst_input = arith.constant dense<[[[[2.0]]]]> : tensor<1x1x1x1xf32>
+  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
+
+  %0 = "tfl.transpose_conv"(%cst_shape, %cst_weights, %cst_input, %cst_bias) <{fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<4xi32>, tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, tensor<1x1x1x1xf32>, tensor<2xf32>) -> tensor<1x?x?x2xf32>
+  func.return %0 : tensor<1x?x?x2xf32>
+
+  // CHECK: tfl.transpose_conv
+}
+
+// CHECK-LABEL: @NoConstantFoldFullyConnectedExceedsSizeLimit
+func.func @NoConstantFoldFullyConnectedExceedsSizeLimit() -> tensor<5x5xf32> {
+  %cst_input = arith.constant dense<1.0> : tensor<5x2xf32>
+  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<5x2x!quant.uniform<i8:f32, 0.5:0>>, value = dense<2> : tensor<5x2xi8>}> : () -> tensor<5x2x!quant.uniform<i8:f32, 0.5:0>>
+  %cst_bias = arith.constant dense<0.0> : tensor<5xf32>
+
+  %0 = "tfl.fully_connected"(%cst_input, %cst_weights, %cst_bias) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<5x2xf32>, tensor<5x2x!quant.uniform<i8:f32, 0.5:0>>, tensor<5xf32>) -> tensor<5x5xf32>
+  func.return %0 : tensor<5x5xf32>
+
+  // CHECK: tfl.fully_connected
+}
+
 // CHECK-LABEL: @ShapeOpI32
 func.func @ShapeOpI32(%arg0 : tensor<576x72xf32>) -> tensor<2xi32> {
   %0 = "tfl.shape"(%arg0) : (tensor<576x72xf32>) -> tensor<2xi32>
@@ -1313,5 +1746,11 @@ func.func @slice_big_float() -> tensor<1x1x1792x256xf32> {
 // CHECK-LITERAL: %cst = arith.constant dense<9.000000e+00> : tensor<1x1x1792x256xf32>
 // CHECK-NOT:     slice
 
-
-
+{-#
+  dialect_resources: {
+    builtin: {
+      fc_w8_2x2: "0x040000000204FF03",
+      fc_w4_2x2: "0x04000000421D"
+    }
+  }
+#-}
diff --git a/tensorflow/compiler/mlir/lite/tests/optimize.mlir b/tensorflow/compiler/mlir/lite/tests/optimize.mlir
index 2861c01f2d5..c18da5a22ed 100644
--- a/tensorflow/compiler/mlir/lite/tests/optimize.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/optimize.mlir
@@ -5048,6 +5048,47 @@ func.func @FullyConnectedSwapOperandsWhenLHSIsConstLHSRank3(%arg0: tensor<512x51
   // CHECK:  %0 = "tfl.fully_connected"(%cst, %arg0, %arg1)
 }
 
+// CHECK-LABEL: @FullyConnectedSwapOperandsWhenLHSIsConstQConstFilter
+func.func @FullyConnectedSwapOperandsWhenLHSIsConstQConstFilter(%arg0: none) -> tensor<2x2xf32> {
+  %cst = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xf32>
+  %qconst = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x!quant.uniform<i16:f32, 0.5:0>>, value = dense<2> : tensor<2x2xi16>}> : () -> tensor<2x2x!quant.uniform<i16:f32, 0.5:0>>
+  %0 = "tfl.fully_connected"(%cst, %qconst, %arg0) {asymmetric_quantize_inputs = true, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<2x2xf32>, tensor<2x2x!quant.uniform<i16:f32, 0.5:0>>, none) -> tensor<2x2xf32>
+  func.return %0 : tensor<2x2xf32>
+
+  // CHECK: %[[QCONST:.*]] = "tfl.pseudo_qconst"
+  // CHECK: %[[FC:.*]] = "tfl.fully_connected"(%cst, %[[QCONST]], %arg0)
+  // CHECK-NOT: tfl.transpose
+  // CHECK: return %[[FC]] : tensor<2x2xf32>
+}
+
+// CHECK-LABEL: @FullyConnectedSwapOperandsWhenLHSIsConstDequantQConstFilter
+func.func @FullyConnectedSwapOperandsWhenLHSIsConstDequantQConstFilter(%arg0: none) -> tensor<2x2xf32> {
+  %cst = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xf32>
+  %qconst = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x!quant.uniform<i16:f32, 0.5:0>>, value = dense<2> : tensor<2x2xi16>}> : () -> tensor<2x2x!quant.uniform<i16:f32, 0.5:0>>
+  %dequant = "tfl.dequantize"(%qconst) : (tensor<2x2x!quant.uniform<i16:f32, 0.5:0>>) -> tensor<2x2xf32>
+  %0 = "tfl.fully_connected"(%cst, %dequant, %arg0) {asymmetric_quantize_inputs = true, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<2x2xf32>, tensor<2x2xf32>, none) -> tensor<2x2xf32>
+  func.return %0 : tensor<2x2xf32>
+
+  // CHECK: %[[DEQUANT:.*]] = "tfl.dequantize"
+  // CHECK: %[[FC:.*]] = "tfl.fully_connected"(%cst, %[[DEQUANT]], %arg0)
+  // CHECK-NOT: tfl.transpose
+  // CHECK: return %[[FC]] : tensor<2x2xf32>
+}
+
+// CHECK-LABEL: @FullyConnectedSwapOperandsWhenLHSIsConstDequantF16Filter
+func.func @FullyConnectedSwapOperandsWhenLHSIsConstDequantF16Filter(%arg0: none) -> tensor<2x2xf32> {
+  %cst = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xf32>
+  %cst_f16 = "tfl.pseudo_const"() <{value = dense<1.0> : tensor<2x2xf16>}> : () -> tensor<2x2xf16>
+  %dequant = "tfl.dequantize"(%cst_f16) : (tensor<2x2xf16>) -> tensor<2x2xf32>
+  %0 = "tfl.fully_connected"(%cst, %dequant, %arg0) {asymmetric_quantize_inputs = true, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<2x2xf32>, tensor<2x2xf32>, none) -> tensor<2x2xf32>
+  func.return %0 : tensor<2x2xf32>
+
+  // CHECK: %[[DEQUANT:.*]] = "tfl.dequantize"
+  // CHECK: %[[FC:.*]] = "tfl.fully_connected"(%cst, %[[DEQUANT]], %arg0)
+  // CHECK-NOT: tfl.transpose
+  // CHECK: return %[[FC]] : tensor<2x2xf32>
+}
+
 // CHECK-LABEL: @AddComputedZero
 func.func @AddComputedZero(%arg0: tensor<512x512xf32>, %arg1: tensor<1x512xf32>) -> tensor<512x512xf32> {
   %0 = "tfl.sub"(%arg1, %arg1) {fused_activation_function = "NONE"} : (tensor<1x512xf32>, tensor<1x512xf32>) -> tensor<1x512xf32>
diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize_pass.cc b/tensorflow/compiler/mlir/lite/transforms/optimize_pass.cc
index 0da83d8c342..16d9983252e 100644
--- a/tensorflow/compiler/mlir/lite/transforms/optimize_pass.cc
+++ b/tensorflow/compiler/mlir/lite/transforms/optimize_pass.cc
@@ -3331,8 +3331,14 @@ struct FullyConnectedSwapOperandsWhenLHSIsConst
     auto input = fc.getInput();
     auto filter = fc.getFilter();
 
-    if (!matchPattern(input, m_Constant()) ||
-        matchPattern(filter, m_Constant()))
+    auto is_const_or_qconst = [](Value v) {
+      return matchPattern(v, m_Constant()) ||
+             llvm::isa_and_nonnull<TFL::QConstOp, TFL::SparseConstOp,
+                                   TFL::SparseQConstOp>(v.getDefiningOp());
+    };
+    auto filter_dequant = filter.getDefiningOp<TFL::DequantizeOp>();
+    if (!matchPattern(input, m_Constant()) || is_const_or_qconst(filter) ||
+        (filter_dequant && is_const_or_qconst(filter_dequant.getInput())))
       return failure();
 
     auto input_type = mlir::dyn_cast<RankedTensorType>(input.getType());

From 50c5820fa425ed6006c95bb0c3ddbe7152729c7f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 15 Oct 2026 10:50:30 -0700
Subject: [PATCH 0882/1000] [XLA:GPU] Don't read the autotune cache when
 determinism is required.

Note that deterministic models already do not write to the cache.

The autotune cache is shared by all modules in a process. Suppose we have models A and B, where model A is the same model as B except B is annotated to require determinism. Further suppose the XLA flags are such that we are using autotuning by default. If model A runs inference first it will save its results to the autotune cache, and then when B runs inference it will check the autotune cache before deciding what to do with the remaining ops (e.g. autotune them, use cost model, etc). However this violates the requirements of the deterministic model. This patch prevents deterministic models from reading from the autotune cache so it does not get polluted.

PiperOrigin-RevId: 1000245268
---
 .../xla/backends/autotuner/config_assigner.cc | 18 +++-
 .../xla/backends/autotuner/config_assigner.h  | 12 ++-
 .../autotuner/config_assigner_test.cc         | 48 ++++++++++
 .../xla/xla/service/gpu/autotuning/BUILD      |  3 +
 .../gpu/autotuning/config_assigner_pass.cc    | 10 ++
 .../autotuning/config_assigner_pass_test.cc   | 92 +++++++++++++++++++
 6 files changed, 175 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/xla/backends/autotuner/config_assigner.cc b/third_party/xla/xla/backends/autotuner/config_assigner.cc
index 4d21fa148cc..07847a6b4b4 100644
--- a/third_party/xla/xla/backends/autotuner/config_assigner.cc
+++ b/third_party/xla/xla/backends/autotuner/config_assigner.cc
@@ -144,7 +144,7 @@ absl::StatusOr<std::optional<ConfigAssigner::Config>> ParseForcedConfig(
 
 absl::StatusOr<std::unique_ptr<ConfigAssigner>> ConfigAssigner::Create(
     Options options,
-    std::unique_ptr<AutotunerCacheInterface> absl_nonnull cache,
+    std::unique_ptr<AutotunerCacheInterface> absl_nullable cache,
     std::unique_ptr<CodegenOrchestrator> absl_nonnull orchestrator,
     std::unique_ptr<Autotuner> absl_nullable autotuner,
     tsl::thread::ThreadPool* thread_pool) {
@@ -189,10 +189,11 @@ absl::Status ConfigAssigner::AssignConfigs(
     HloModule* module, const InstructionFilterFn& should_assign_config,
     MultiProcessKeyValueStore& sharding_kv_store) {
   // Sharding the instructions only makes sense if we can have different
-  // configs for different shards, which only happens due to online tuning.
-  if (!options_.allow_autotuning) {
+  // configs for different shards, which only happens due to online tuning, and
+  // requires the cache to exchange results across shards.
+  if (!options_.allow_autotuning || optimal_config_cache_ == nullptr) {
     VLOG(1) << "Falling back to non-sharded config assignment as online "
-               "tuning is disabled.";
+               "tuning or caching is disabled.";
     return AssignConfigs(module, should_assign_config);
   }
 
@@ -488,6 +489,9 @@ ConfigAssigner::GetFirstCompilableConfigOrDefault(const HloInstruction* instr) {
 
 std::optional<ConfigAssigner::Config> ConfigAssigner::LookUp(
     const HloInstruction* instr) const {
+  if (optimal_config_cache_ == nullptr) {
+    return std::nullopt;
+  }
   auto cached_config = optimal_config_cache_->Lookup(instr);
   if (!cached_config.has_value()) {
     return std::nullopt;
@@ -510,6 +514,9 @@ std::optional<ConfigAssigner::Config> ConfigAssigner::LookUp(
 
 absl::Status ConfigAssigner::Insert(const HloInstruction* instr,
                                     const ConfigAssigner::Config& config) {
+  if (optimal_config_cache_ == nullptr) {
+    return absl::OkStatus();
+  }
   AutotunerCacheInterface::Config cached_config;
   cached_config.codegen_backend = config.codegen_backend->backend();
   cached_config.backend_config = *config.backend_config;
@@ -597,6 +604,9 @@ std::string ConfigAssigner::Options::ToString() const {
 }
 
 AutotunerCacheInterface::CacheStats ConfigAssigner::GetCacheStats() const {
+  if (optimal_config_cache_ == nullptr) {
+    return {};
+  }
   return optimal_config_cache_->GetCacheStats();
 }
 
diff --git a/third_party/xla/xla/backends/autotuner/config_assigner.h b/third_party/xla/xla/backends/autotuner/config_assigner.h
index 27ec7b25ff5..68cfe4cc8e8 100644
--- a/third_party/xla/xla/backends/autotuner/config_assigner.h
+++ b/third_party/xla/xla/backends/autotuner/config_assigner.h
@@ -70,9 +70,11 @@ class ConfigAssigner {
 
   using Config = CodegenOrchestrator::Config;
 
+  // If `optimal_config_cache` is null, the config-assigner neither reads from
+  // nor writes to any autotune cache.
   static absl::StatusOr<std::unique_ptr<ConfigAssigner>> Create(
       Options options,
-      absl_nonnull std::unique_ptr<AutotunerCacheInterface>
+      absl_nullable std::unique_ptr<AutotunerCacheInterface>
           optimal_config_cache,
       absl_nonnull std::unique_ptr<CodegenOrchestrator> orchestrator,
       absl_nullable std::unique_ptr<Autotuner> autotuner,
@@ -94,7 +96,7 @@ class ConfigAssigner {
 
  private:
   ConfigAssigner(Options options,
-                 absl_nonnull std::unique_ptr<AutotunerCacheInterface> cache,
+                 absl_nullable std::unique_ptr<AutotunerCacheInterface> cache,
                  absl_nonnull std::unique_ptr<CodegenOrchestrator> orchestrator,
                  absl_nullable std::unique_ptr<Autotuner> autotuner,
                  tsl::thread::ThreadPool* thread_pool,
@@ -142,16 +144,18 @@ class ConfigAssigner {
       const HloInstruction* instr);
 
   // Returns the cached config for the given HLO instruction, if any.
-  // Otherwise, returns std::nullopt.
+  // Otherwise (including when there is no cache), returns std::nullopt.
   std::optional<Config> LookUp(const HloInstruction* instr) const;
   // Inserts the given config into the cache for the given HLO instruction.
+  // No-op if there is no cache.
   absl::Status Insert(const HloInstruction* instr, const Config& config);
 
   // Dumps HLO before and after applying the config.
   absl::Status DumpHlo(const HloInstruction& instr, const Config& config);
 
   Options options_;
-  absl_nonnull std::unique_ptr<AutotunerCacheInterface> optimal_config_cache_;
+  // Null if the config-assigner must not use any autotune cache.
+  absl_nullable std::unique_ptr<AutotunerCacheInterface> optimal_config_cache_;
   absl_nonnull std::unique_ptr<CodegenOrchestrator> orchestrator_;
   absl_nullable std::unique_ptr<Autotuner> autotuner_;
   tsl::thread::ThreadPool* thread_pool_ = nullptr;
diff --git a/third_party/xla/xla/backends/autotuner/config_assigner_test.cc b/third_party/xla/xla/backends/autotuner/config_assigner_test.cc
index 6096857bf1d..8d470ff40c7 100644
--- a/third_party/xla/xla/backends/autotuner/config_assigner_test.cc
+++ b/third_party/xla/xla/backends/autotuner/config_assigner_test.cc
@@ -245,6 +245,54 @@ TEST_F(ConfigAssignerTest, CacheHit) {
   EXPECT_THAT(config_assigner->AssignConfig(dummy_instr.get()), IsOk());
 }
 
+TEST_F(ConfigAssignerTest, NullCacheAutotunesWithoutCaching) {
+  config_.allow_autotuning = true;
+
+  auto backend = std::make_unique<MockCodegenBackend>();
+  EXPECT_CALL(*backend, name()).WillRepeatedly(Return("mock_backend"));
+  EXPECT_CALL(*backend, GetSupportedConfigs(_))
+      .WillOnce([](const HloInstruction&) {
+        std::vector<std::unique_ptr<BackendConfig>> configs;
+        configs.push_back(GetTestConfig("test_config_1"));
+        configs.push_back(GetTestConfig("test_config_2"));
+        return configs;
+      });
+  EXPECT_CALL(*backend, Compile(_, _)).Times(2).WillRepeatedly([] {
+    return std::unique_ptr<Executable>();
+  });
+  EXPECT_CALL(*backend, ApplyConfig(_, ConfigMatcher("test_config_1")))
+      .WillOnce(Return(absl::OkStatus()));
+
+  auto profiler = std::make_unique<MockProfiler>();
+  EXPECT_CALL(*profiler, CreateInputBuffers(_, _))
+      .WillOnce(Return(std::make_unique<InputBuffers>()));
+  EXPECT_CALL(*profiler, Profile(_, _))
+      .WillOnce(Return(ProfileResult({absl::Seconds(1)})))
+      .WillOnce(Return(ProfileResult({absl::Seconds(2)})));
+
+  std::vector<std::unique_ptr<CodegenBackend>> backends;
+  backends.push_back(std::move(backend));
+  ASSERT_OK_AND_ASSIGN(auto orchestrator,
+                       CodegenOrchestrator::Create(std::move(backends), {}));
+  Autotuner::Options autotuner_options;
+  autotuner_options.correctness_check_options.enable_correctness_check = false;
+  std::vector<std::unique_ptr<Profiler>> profilers;
+  profilers.push_back(std::move(profiler));
+  ASSERT_OK_AND_ASSIGN(
+      auto autotuner,
+      Autotuner::Create(*orchestrator, std::move(profilers), autotuner_options,
+                        /*thread_pool=*/nullptr));
+  ASSERT_OK_AND_ASSIGN(
+      auto config_assigner,
+      ConfigAssigner::Create(config_, /*optimal_config_cache=*/nullptr,
+                             std::move(orchestrator), std::move(autotuner)));
+
+  auto dummy_instr = HloInstruction::CreateConstant(LiteralUtil::CreateR0(1));
+  EXPECT_THAT(config_assigner->AssignConfig(dummy_instr.get()), IsOk());
+  EXPECT_EQ(config_assigner->GetCacheStats(),
+            AutotunerCacheInterface::CacheStats());
+}
+
 TEST_F(ConfigAssignerTest, ForceConfigBypassesCacheAndAutotuning) {
   auto cache_manager = std::make_unique<MockAutotunerCache>();
   EXPECT_CALL(*cache_manager, Lookup(_)).Times(0);
diff --git a/third_party/xla/xla/service/gpu/autotuning/BUILD b/third_party/xla/xla/service/gpu/autotuning/BUILD
index c2f420473f5..fd030fd2183 100644
--- a/third_party/xla/xla/service/gpu/autotuning/BUILD
+++ b/third_party/xla/xla/service/gpu/autotuning/BUILD
@@ -83,7 +83,10 @@ xla_test(
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/log:globals",
         "@com_google_absl//absl/log:scoped_mock_log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest_main",
         "@llvm-project//mlir:IR",
diff --git a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc
index 89822e5e713..146cc60d759 100644
--- a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc
+++ b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc
@@ -247,6 +247,15 @@ std::unique_ptr<AutotunerCacheInterface> CreateAutotunerCache(
     const DebugOptions& debug_options,
     const Compiler::GpuTargetConfig& target_config,
     const std::vector<std::unique_ptr<CodegenBackend>>& backends) {
+  // The autotune cache is process-wide and keyed without debug options, so it
+  // may contain timing-based picks inserted by other (non-deterministic)
+  // modules. Using them would make a deterministic compilation depend on what
+  // else was compiled in the process first, so don't use any cache.
+  if (debug_options.xla_gpu_deterministic_ops() ||
+      debug_options.xla_gpu_exclude_nondeterministic_ops()) {
+    return nullptr;
+  }
+
   std::string cache_dir = debug_options.xla_gpu_per_fusion_autotune_cache_dir();
 
   if (!debug_options.xla_gpu_use_new_autotune_cache_format()) {
@@ -471,6 +480,7 @@ absl::StatusOr<std::unique_ptr<ConfigAssignerPass>> ConfigAssignerPass::Create(
       GetConfigAssignerOptions(debug_options, is_deviceless);
   CodegenOrchestrator::Options orchestrator_options =
       GetCodegenOrchestratorOptions(debug_options);
+  // Null in deterministic mode.
   std::unique_ptr<AutotunerCacheInterface> cache =
       CreateAutotunerCache(debug_options, *target_config, backends);
 
diff --git a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc
index e63c8f02985..9a766507004 100644
--- a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc
+++ b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass_test.cc
@@ -27,7 +27,10 @@ limitations under the License.
 #include "absl/log/globals.h"
 #include "absl/log/log.h"
 #include "absl/log/scoped_mock_log.h"
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
 #include "absl/status/status_matchers.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/ascii.h"
 #include "absl/strings/str_cat.h"
 #include "mlir/IR/MLIRContext.h"
@@ -1197,6 +1200,95 @@ TEST_F(ConfigAssignerPassTest, OtherOpcodeAllowsSpills) {
       *instr, autotuner::Backend::BLOCK_LEVEL_EMITTER));
 }
 
+// Regression test: the autotune caches are process-wide, so a deterministic
+// compilation must not reuse configs that a non-deterministic compilation of a
+// different module selected via (timing-based) autotuning. Otherwise the
+// deterministic result depends on what else was compiled in the process first.
+class DeterministicCacheIsolationTest
+    : public ConfigAssignerPassTest,
+      public ::testing::WithParamInterface<bool /*use_new_cache_format*/> {
+ protected:
+  absl::StatusOr<bool> RunPass(HloModule* module,
+                               tsl::thread::ThreadPool* thread_pool) {
+    GpuCompiler::GpuTargetConfig target_config(stream_executor_);
+    std::vector<std::unique_ptr<CodegenBackend>> backends;
+    backends.push_back(std::make_unique<CublasLtBackend>(
+        stream_executor_, &module->config().debug_options(), &compiler_,
+        &target_config));
+    auto get_backends_fn =
+        [backends =
+             std::make_shared<std::vector<std::unique_ptr<CodegenBackend>>>(
+                 std::move(backends))]() mutable {
+          return std::move(*backends);
+        };
+    ABSL_ASSIGN_OR_RETURN(
+        std::unique_ptr<ConfigAssignerPass> pass,
+        ConfigAssignerPass::Create(
+            std::move(get_backends_fn), module->config().debug_options(),
+            target_config.device_description.gpu_compute_capability(),
+            stream_executor_, thread_pool, &target_config,
+            /*alias_info=*/nullptr, /*mlir_context=*/nullptr,
+            /*shape_size_fn=*/[](const Shape& shape) { return 0; },
+            allocator_.get()));
+    return pass->Run(module, /*execution_threads=*/{});
+  }
+};
+
+TEST_P(DeterministicCacheIsolationTest,
+       DeterministicCompileIgnoresNonDeterministicCacheEntries) {
+  const bool use_new_cache_format = GetParam();
+  tsl::thread::ThreadPool thread_pool(tsl::Env::Default(), "autotuning",
+                                      /*num_threads=*/4);
+
+  // 1. A non-deterministic compilation autotunes the GEMM and populates the
+  //    (in-memory, process-wide) autotune cache.
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> autotuned_module,
+                       ParseAndReturnVerifiedModule(kCublasCustomCallHlo));
+  autotuned_module->mutable_config()
+      .mutable_debug_options()
+      .set_xla_gpu_use_new_autotune_cache_format(use_new_cache_format);
+  EXPECT_THAT(RunPass(autotuned_module.get(), &thread_pool),
+              absl_testing::IsOkAndHolds(true));
+
+  // 2. A deterministic compilation of the same GEMM, requiring every config to
+  //    come from the cache, must not see the entry inserted above.
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> deterministic_module,
+                       ParseAndReturnVerifiedModule(kCublasCustomCallHlo));
+  DebugOptions& debug_options =
+      deterministic_module->mutable_config().mutable_debug_options();
+  debug_options.set_xla_gpu_use_new_autotune_cache_format(use_new_cache_format);
+  debug_options.set_xla_gpu_exclude_nondeterministic_ops(true);
+  debug_options.set_xla_gpu_require_complete_aot_autotune_results(true);
+  EXPECT_THAT(RunPass(deterministic_module.get(), &thread_pool),
+              absl_testing::StatusIs(absl::StatusCode::kNotFound));
+
+  // 3. Without `require_complete_aot_autotune_results`, the deterministic
+  //    compilation falls back to the first valid config, exactly as if the
+  //    non-deterministic compilation had never happened.
+  ASSERT_OK_AND_ASSIGN(deterministic_module,
+                       ParseAndReturnVerifiedModule(kCublasCustomCallHlo));
+  DebugOptions& debug_options_2 =
+      deterministic_module->mutable_config().mutable_debug_options();
+  debug_options_2.set_xla_gpu_use_new_autotune_cache_format(
+      use_new_cache_format);
+  debug_options_2.set_xla_gpu_exclude_nondeterministic_ops(true);
+  EXPECT_THAT(RunPass(deterministic_module.get(), &thread_pool),
+              absl_testing::IsOkAndHolds(true));
+  const HloInstruction* gemm =
+      deterministic_module->entry_computation()->GetInstructionWithName(
+          "custom-call.1");
+  ASSERT_OK_AND_ASSIGN(auto gpu_backend_config,
+                       gemm->backend_config<GpuBackendConfig>());
+  EXPECT_EQ(gpu_backend_config.gemm_backend_config().selected_algorithm(), 0);
+}
+
+INSTANTIATE_TEST_SUITE_P(CacheFormats, DeterministicCacheIsolationTest,
+                         ::testing::Bool(),
+                         [](const ::testing::TestParamInfo<bool>& info) {
+                           return info.param ? "NewCacheFormat"
+                                             : "LegacyCacheFormat";
+                         });
+
 }  // namespace
 }  // namespace gpu
 }  // namespace xla

From e7c099deeea4f21c01539c1a2ab8b564e3a3eeed Mon Sep 17 00:00:00 2001
From: Akhil Goel <akhil.goel@intel.com>
Date: Thu, 15 Oct 2026 10:58:49 -0700
Subject: [PATCH 0883/1000] PR #50820: [XLA:GPU][oneAPI] Align XPU Triton hook
 with upstream Triton

Imported from GitHub PR https://github.com/openxla/xla/pull/50820

This PR aligns the XPU Triton hook with https://github.com/triton-lang/triton/commit/b575c0c38d428a609e25487db64051a5fa7bfab8
Copybara import of the project:

--
a95502b1e179ff1b8b638b7c04615e6c55cbbc75 by Akhil Goel <akhil.goel@intel.com>:

Update XPU Triton hook

Merging this change closes #50820

PiperOrigin-RevId: 1000251170
---
 third_party/xla/third_party/triton/intel_xpu/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/third_party/triton/intel_xpu/workspace.bzl b/third_party/xla/third_party/triton/intel_xpu/workspace.bzl
index a2318a6dcef..bf183ad297a 100644
--- a/third_party/xla/third_party/triton/intel_xpu/workspace.bzl
+++ b/third_party/xla/third_party/triton/intel_xpu/workspace.bzl
@@ -1,7 +1,7 @@
 """Intel XPU Triton archive metadata."""
 
-XPU_TRITON_COMMIT = "33f9cb578f45b8fb719c728896b2a75f29e66616"
-XPU_TRITON_SHA256 = "856f7c73cd5dd43a9157709f4c17ca70fbb1e2504803711e66d62a9e0dadadd1"
+XPU_TRITON_COMMIT = "e45abef5f100e02e2f9e98f5f5a2fd8dfe3c6c34"
+XPU_TRITON_SHA256 = "a0d191feab8e91bc710f96fb0aefde8a22b0c99daac4f2c4153adc8f6c98b26c"
 
 def use_xpu_triton(repository_ctx):
     return repository_ctx.getenv("ENABLE_INTEL_XPU_TRITON", "").strip() == "1"

From 09e4e18214bc830115facb92a54cda0754c2a938 Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Thu, 15 Oct 2026 12:01:37 -0700
Subject: [PATCH 0884/1000] Test NewSession with an unparseable
 SessionOptions.Config

Cover the error path of SessionOptions.c() that now frees its copy of
Config as well.
---
 tensorflow/go/session_test.go | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/tensorflow/go/session_test.go b/tensorflow/go/session_test.go
index ff7cbb7c89d..b41c8561c78 100644
--- a/tensorflow/go/session_test.go
+++ b/tensorflow/go/session_test.go
@@ -19,6 +19,7 @@ package tensorflow
 import (
 	"fmt"
 	"reflect"
+	"strings"
 	"testing"
 )
 
@@ -398,6 +399,20 @@ func TestSessionConfig(t *testing.T) {
 	}
 }
 
+func TestSessionConfigInvalid(t *testing.T) {
+	// SessionOptions.c() returns early when TF_SetConfig rejects Config, and
+	// has to free its C copy of Config there too (GitHub issue 129133).
+	graph := NewGraph()
+	opts := SessionOptions{Config: []byte("invalid config")}
+	_, err := NewSession(graph, &opts)
+	if err == nil {
+		t.Fatal("Expected error from NewSession with invalid config, got nil")
+	}
+	if got, want := err.Error(), "invalid SessionOptions.Config"; !strings.Contains(got, want) {
+		t.Errorf("Got error %q, want it to contain %q", got, want)
+	}
+}
+
 func TestListDevices(t *testing.T) {
 	s, err := NewSession(NewGraph(), nil)
 	if err != nil {

From b6b866c2e8c815235791ddb4178dbaf7009ff40d Mon Sep 17 00:00:00 2001
From: Ilya Tikhonovskiy <loislo@google.com>
Date: Thu, 15 Oct 2026 11:43:40 -0700
Subject: [PATCH 0885/1000] [XLA:GPU] Enable MXFP8 E5M2xE5M2 and MXFP4xMXFP4 in
 cuDNN block-scaled dot.

cuDNN 9.10+ on Blackwell supports E5M2xE5M2 MXFP8 and E2M1FNxE2M1FN MXFP4 (with E8M0FNU scale and block size 32) block-scaled matrix multiplication. Add `CudnnMxType::MXFP4` to `GetCudnnMxType` and allow `MXFP8_E5M2 x MXFP8_E5M2` and `MXFP4 x MXFP4` in `IsSupportedByCudnn`.

PiperOrigin-RevId: 1000284389
---
 .../gpu/codegen/cudnn_scaled_dot_test.cc      |  8 +--
 .../gpu/transforms/block_scaling_rewriter.cc  | 14 +++-
 .../gpu/transforms/block_scaling_rewriter.h   |  1 +
 .../block_scaling_rewriter_cudnn_test.cc      | 72 +++++++++++++++++++
 .../transforms/block_scaling_rewriter_test.cc | 34 ++++++++-
 5 files changed, 122 insertions(+), 7 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/cudnn_scaled_dot_test.cc b/third_party/xla/xla/backends/gpu/codegen/cudnn_scaled_dot_test.cc
index 198f26d8f45..40a7f2ce220 100644
--- a/third_party/xla/xla/backends/gpu/codegen/cudnn_scaled_dot_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/cudnn_scaled_dot_test.cc
@@ -222,7 +222,7 @@ std::vector<ScaledDotCoverageTestCase> GetCoverageTestCases() {
       {F8E5M2, F8E4M3FN, F8E8M0FNU, 32, true, true,
        CuDnnSupportResult::kSupported},
 
-      // 5. MXFP8 E5M2 x E5M2 (scale E8M0FNU, block 32) - unsupported by cuDNN
+      // 5. MXFP8 E5M2 x E5M2 (scale E8M0FNU, block 32)
       {F8E5M2, F8E5M2, F8E8M0FNU, 32, false, false,
        CuDnnSupportResult::kNotSupported},
       {F8E5M2, F8E5M2, F8E8M0FNU, 32, false, true,
@@ -230,7 +230,7 @@ std::vector<ScaledDotCoverageTestCase> GetCoverageTestCases() {
       {F8E5M2, F8E5M2, F8E8M0FNU, 32, true, false,
        CuDnnSupportResult::kNotSupported},
       {F8E5M2, F8E5M2, F8E8M0FNU, 32, true, true,
-       CuDnnSupportResult::kNotSupported},
+       CuDnnSupportResult::kSupported},
 
       // 6. E5M2 x F4E2M1FN (scale E8M0FNU, block 32)
       {F8E5M2, F4E2M1FN, F8E8M0FNU, 32, false, false,
@@ -270,7 +270,7 @@ std::vector<ScaledDotCoverageTestCase> GetCoverageTestCases() {
       {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 32, true, false,
        CuDnnSupportResult::kNotSupported},
       {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 32, true, true,
-       CuDnnSupportResult::kNotSupported},
+       CuDnnSupportResult::kSupported},
 
       // 10. F4E2M1FN x F4E2M1FN (scale E8M0FNU, block 16)
       {F4E2M1FN, F4E2M1FN, F8E8M0FNU, 16, false, false,
@@ -519,7 +519,7 @@ ENTRY %main {
           }
           auto run_result = RunAndCompareTwoModules(
               std::move(*optimized_module_or), std::move(*ref_optimized_or),
-              arg_ptrs, ErrorSpec{/*aabs=*/0.02, /*arel=*/0.05},
+              arg_ptrs, ErrorSpec{/*aabs=*/1e-4, /*arel=*/1e-5},
               /*run_hlo_passes=*/false);
           status = run_result ? "SUCCESS" : "FAILURE";
           if (!run_result) {
diff --git a/third_party/xla/xla/backends/gpu/transforms/block_scaling_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/block_scaling_rewriter.cc
index ab2d67e89a4..a9bea21a5c5 100644
--- a/third_party/xla/xla/backends/gpu/transforms/block_scaling_rewriter.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/block_scaling_rewriter.cc
@@ -257,6 +257,8 @@ enum class CudnnMxType {
   MXFP8_E4M3FN,
   // Input: E5M2, scale: E8M0FNU, block size: 32.
   MXFP8_E5M2,
+  // Input: E2M1FN, scale: E8M0FNU, block size: 32.
+  MXFP4,
   // Input: E2M1FN, scale: E4M3FN, block size: 16.
   NVFP4,
 };
@@ -292,6 +294,14 @@ CudnnMxType GetCudnnMxType(const Shape& input_shape, const Shape& scale_shape,
     return CudnnMxType::MXFP8_E5M2;
   }
 
+  // MXFP4: the input is E2M1FN and the scale is E8M0FNU.
+  if (input_shape.element_type() == PrimitiveType::F4E2M1FN &&
+      scale_shape.element_type() == PrimitiveType::F8E8M0FNU &&
+      actual_block_size == BlockScalingRewriter::kBlockSizeMXFP4 &&
+      contracting_size % actual_block_size == 0) {
+    return CudnnMxType::MXFP4;
+  }
+
   // NVFP4: the input is E2M1FN and the scale is E4M3FN.
   if (input_shape.element_type() == PrimitiveType::F4E2M1FN &&
       scale_shape.element_type() == PrimitiveType::F8E4M3FN &&
@@ -310,12 +320,12 @@ CudnnMxType GetCudnnMxType(const Shape& input_shape, const Shape& scale_shape,
 }
 
 bool IsSupportedByCudnn(CudnnMxType lhs, CudnnMxType rhs) {
-  // cuDNN supports mixing input types for MXFP8, but the E5M2/E5M2 combination
-  // is not supported.
   return (lhs == CudnnMxType::MXFP8_E4M3FN &&
           rhs == CudnnMxType::MXFP8_E4M3FN) ||
          (lhs == CudnnMxType::MXFP8_E4M3FN && rhs == CudnnMxType::MXFP8_E5M2) ||
          (lhs == CudnnMxType::MXFP8_E5M2 && rhs == CudnnMxType::MXFP8_E4M3FN) ||
+         (lhs == CudnnMxType::MXFP8_E5M2 && rhs == CudnnMxType::MXFP8_E5M2) ||
+         (lhs == CudnnMxType::MXFP4 && rhs == CudnnMxType::MXFP4) ||
          (lhs == CudnnMxType::NVFP4 && rhs == CudnnMxType::NVFP4);
 }
 
diff --git a/third_party/xla/xla/backends/gpu/transforms/block_scaling_rewriter.h b/third_party/xla/xla/backends/gpu/transforms/block_scaling_rewriter.h
index 76b3870fca2..30d9c52780f 100644
--- a/third_party/xla/xla/backends/gpu/transforms/block_scaling_rewriter.h
+++ b/third_party/xla/xla/backends/gpu/transforms/block_scaling_rewriter.h
@@ -93,6 +93,7 @@ class BlockScalingRewriter : public OpExpanderPass {
 
   // Common block size constants.
   static constexpr int kBlockSizeMXFP8 = 32;
+  static constexpr int kBlockSizeMXFP4 = 32;
   static constexpr int kBlockSizeNVFP4 = 16;
 
  private:
diff --git a/third_party/xla/xla/backends/gpu/transforms/block_scaling_rewriter_cudnn_test.cc b/third_party/xla/xla/backends/gpu/transforms/block_scaling_rewriter_cudnn_test.cc
index 22b28f544f5..5dc3dce7808 100644
--- a/third_party/xla/xla/backends/gpu/transforms/block_scaling_rewriter_cudnn_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/block_scaling_rewriter_cudnn_test.cc
@@ -225,5 +225,77 @@ ENTRY main {
                             "CHECK: __cudnn$blockScaledDot");
 }
 
+TEST_F(BlockScalingRewriterCudnnTest, Mxfp8_E5M2) {
+  constexpr absl::string_view hlo_string = R"(
+HloModule test
+
+ENTRY main {
+  %lhs = f8e5m2[256,256] parameter(0)
+  %rhs = f8e5m2[256,256] parameter(1)
+  %lhs_scale = f8e8m0fnu[256,8] parameter(2)
+  %rhs_scale = f8e8m0fnu[256,8] parameter(3)
+  ROOT %result = f32[256,256] custom-call(%lhs, %rhs, %lhs_scale, %rhs_scale),
+      custom_call_target="__op$block_scaled_dot"
+})";
+
+  EXPECT_TRUE(RunAndCompare(
+      hlo_string, ErrorSpec(/*aabs=*/1e-4, /*arel=*/1e-5),
+      /*reference_preprocessor=*/
+      [](HloModule* reference_module) {
+        BlockScalingRewriter pass(kCudnnDisabled);
+        EXPECT_THAT(RunHloPass(&pass, reference_module),
+                    absl_testing::IsOkAndHolds(true));
+      },
+      /*test_preprocessor=*/
+      [](HloModule* test_module) {
+        BlockScalingRewriter pass(kCudnnVersion);
+        EXPECT_THAT(RunHloPass(&pass, test_module),
+                    absl_testing::IsOkAndHolds(true));
+      }));
+
+  RunAndFilecheckHloRewrite(hlo_string, BlockScalingRewriter(kCudnnDisabled),
+                            "CHECK-NOT: __cudnn$blockScaledDot");
+  RunAndFilecheckHloRewrite(hlo_string, BlockScalingRewriter(kCudnnVersion),
+                            "CHECK: __cudnn$blockScaledDot");
+}
+
+TEST_F(BlockScalingRewriterCudnnTest, Mxfp4) {
+  constexpr absl::string_view hlo_string = R"(
+HloModule test
+
+ENTRY main {
+  %mult_scalar = f16[] constant(6)
+  %mult = f16[256,256] broadcast(%mult_scalar), dimensions={}
+  %p0 = f16[256,256] parameter(0)
+  %p1 = f16[256,256] parameter(1)
+  %lhs = f4e2m1fn[256,256] convert(f16[256,256] multiply(%p0, %mult))
+  %rhs = f4e2m1fn[256,256] convert(f16[256,256] multiply(%p1, %mult))
+  %lhs_scale = f8e8m0fnu[256,8] parameter(2)
+  %rhs_scale = f8e8m0fnu[256,8] parameter(3)
+  ROOT %result = f32[256,256] custom-call(%lhs, %rhs, %lhs_scale, %rhs_scale),
+      custom_call_target="__op$block_scaled_dot"
+})";
+
+  EXPECT_TRUE(RunAndCompare(
+      hlo_string, ErrorSpec(/*aabs=*/1e-4, /*arel=*/1e-5),
+      /*reference_preprocessor=*/
+      [](HloModule* reference_module) {
+        BlockScalingRewriter pass(kCudnnDisabled);
+        EXPECT_THAT(RunHloPass(&pass, reference_module),
+                    absl_testing::IsOkAndHolds(true));
+      },
+      /*test_preprocessor=*/
+      [](HloModule* test_module) {
+        BlockScalingRewriter pass(kCudnnVersion);
+        EXPECT_THAT(RunHloPass(&pass, test_module),
+                    absl_testing::IsOkAndHolds(true));
+      }));
+
+  RunAndFilecheckHloRewrite(hlo_string, BlockScalingRewriter(kCudnnDisabled),
+                            "CHECK-NOT: __cudnn$blockScaledDot");
+  RunAndFilecheckHloRewrite(hlo_string, BlockScalingRewriter(kCudnnVersion),
+                            "CHECK: __cudnn$blockScaledDot");
+}
+
 }  // namespace
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/transforms/block_scaling_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/block_scaling_rewriter_test.cc
index 9ea3f006426..d409c98f99b 100644
--- a/third_party/xla/xla/backends/gpu/transforms/block_scaling_rewriter_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/block_scaling_rewriter_test.cc
@@ -376,7 +376,7 @@ ENTRY main {
       test_module->entry_computation()->root_instruction())));
 }
 
-TEST_F(BlockScalingRewriterTest, CudnnFusionUnsupportedE5M2) {
+TEST_F(BlockScalingRewriterTest, CudnnFusionSupportedE5M2) {
   constexpr absl::string_view hlo_string = R"(
 ENTRY main {
   %lhs = f8e5m2[128,256] parameter(0)
@@ -385,6 +385,38 @@ ENTRY main {
   %rhs_scale = f8e8m0fnu[128,8] parameter(3)
   ROOT %result = f16[128,128] scaled-dot(%lhs, %rhs, %lhs_scale, %rhs_scale),
       lhs_contracting_dims={1}, rhs_contracting_dims={1}
+})";
+  ASSERT_OK_AND_ASSIGN(auto test_module,
+                       ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_TRUE(CudnnScaledDotHelper::IsSupported(Cast<HloScaledDotInstruction>(
+      test_module->entry_computation()->root_instruction())));
+}
+
+TEST_F(BlockScalingRewriterTest, CudnnFusionSupportedMxfp4) {
+  constexpr absl::string_view hlo_string = R"(
+ENTRY main {
+  %lhs = f4e2m1fn[128,256] parameter(0)
+  %rhs = f4e2m1fn[128,256] parameter(1)
+  %lhs_scale = f8e8m0fnu[128,8] parameter(2)
+  %rhs_scale = f8e8m0fnu[128,8] parameter(3)
+  ROOT %result = f16[128,128] scaled-dot(%lhs, %rhs, %lhs_scale, %rhs_scale),
+      lhs_contracting_dims={1}, rhs_contracting_dims={1}
+})";
+  ASSERT_OK_AND_ASSIGN(auto test_module,
+                       ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_TRUE(CudnnScaledDotHelper::IsSupported(Cast<HloScaledDotInstruction>(
+      test_module->entry_computation()->root_instruction())));
+}
+
+TEST_F(BlockScalingRewriterTest, CudnnFusionUnsupportedMixedFp8Fp4) {
+  constexpr absl::string_view hlo_string = R"(
+ENTRY main {
+  %lhs = f8e4m3fn[128,256] parameter(0)
+  %rhs = f4e2m1fn[128,256] parameter(1)
+  %lhs_scale = f8e8m0fnu[128,8] parameter(2)
+  %rhs_scale = f8e8m0fnu[128,8] parameter(3)
+  ROOT %result = f16[128,128] scaled-dot(%lhs, %rhs, %lhs_scale, %rhs_scale),
+      lhs_contracting_dims={1}, rhs_contracting_dims={1}
 })";
   ASSERT_OK_AND_ASSIGN(auto test_module,
                        ParseAndReturnVerifiedModule(hlo_string));

From 72972b72de227ac0cec95de44398e2c1845c7cf9 Mon Sep 17 00:00:00 2001
From: Alexandre Boulgakov <sashabu@google.com>
Date: Thu, 15 Oct 2026 12:08:52 -0700
Subject: [PATCH 0886/1000] Add HloInstructionHasName matcher to
 hlo_matchers.h.

PiperOrigin-RevId: 1000301893
---
 third_party/xla/xla/hlo/utils/hlo_matchers.h       | 12 ++++++++++++
 third_party/xla/xla/hlo/utils/hlo_matchers_test.cc |  7 +++++++
 2 files changed, 19 insertions(+)

diff --git a/third_party/xla/xla/hlo/utils/hlo_matchers.h b/third_party/xla/xla/hlo/utils/hlo_matchers.h
index 14b18c53d9b..cac43b220da 100644
--- a/third_party/xla/xla/hlo/utils/hlo_matchers.h
+++ b/third_party/xla/xla/hlo/utils/hlo_matchers.h
@@ -18,6 +18,7 @@ limitations under the License.
 
 #include <optional>
 #include <string>
+#include <tuple>
 #include <utility>
 
 #include "absl/strings/string_view.h"
@@ -645,6 +646,17 @@ std::vector<const HloInstruction*> Pointers(const Container& container) {
   return result;
 }
 
+MATCHER_P(HloInstructionHasName, name, "") {
+  return ::testing::ExplainMatchResult(
+      ::testing::Property("name", &HloInstruction::name, name), *arg,
+      result_listener);
+}
+
+MATCHER(HloInstructionHasName, "") {
+  return ::testing::ExplainMatchResult(HloInstructionHasName(std::get<1>(arg)),
+                                       std::get<0>(arg), result_listener);
+}
+
 }  // namespace testing
 
 // Tell GMock to print HloInstruction* by value, so error messages are nice.
diff --git a/third_party/xla/xla/hlo/utils/hlo_matchers_test.cc b/third_party/xla/xla/hlo/utils/hlo_matchers_test.cc
index 916ceac9da8..959b7e8d7c3 100644
--- a/third_party/xla/xla/hlo/utils/hlo_matchers_test.cc
+++ b/third_party/xla/xla/hlo/utils/hlo_matchers_test.cc
@@ -34,6 +34,7 @@ namespace op = xla::testing::opcode_matchers;
 using ::testing::_;
 using ::testing::Eq;
 using ::testing::HasSubstr;
+using ::xla::testing::HloInstructionHasName;
 
 namespace xla {
 namespace {
@@ -432,5 +433,11 @@ TEST_F(HloMatchersTest, MetadataMatcher) {
   EXPECT_THAT(DescribeHloMatcher(op::Metadata(p0->metadata())),
               R"( (metadata: op_type1 op_name1  0 0 0 0))");
 }
+
+TEST_F(HloMatchersTest, HloInstructionHasName) {
+  std::unique_ptr<HloInstruction> p0 =
+      HloInstruction::CreateParameter(0, ShapeUtil::MakeShape(F32, {1}), "p0");
+  EXPECT_THAT(p0, HloInstructionHasName("p0"));
+}
 }  // namespace
 }  // namespace xla

From 18be7466a79569a25432e853c9138ef7e908318a Mon Sep 17 00:00:00 2001
From: Amir Saadat <amirsaadat@google.com>
Date: Thu, 15 Oct 2026 12:12:09 -0700
Subject: [PATCH 0887/1000] [XLA] Trap subnormal flush to zero in HLO isolator

Traps cases where expected floating-point subnormals are flushed to signed zero
under a nonzero ErrorSpec tolerance in HLO isolation (by default on
bit-preserving ops, or on all ops when trap_subnormal_flush is enabled),
coerces f32 fusion debug callback literals to the reference shape, and extracts
integer mismatch coordinates for HTML mismatch reports.

PiperOrigin-RevId: 1000304064
---
 third_party/xla/xla/tools/hlo_isolation/BUILD |   4 +
 .../tools/hlo_isolation/hlo_isolation_api.cc  | 438 +++++++++++++++---
 .../hlo_isolation/hlo_isolation_api_test.cc   | 211 +++++++++
 3 files changed, 590 insertions(+), 63 deletions(-)

diff --git a/third_party/xla/xla/tools/hlo_isolation/BUILD b/third_party/xla/xla/tools/hlo_isolation/BUILD
index f8e0cdcc4c2..a105a394a78 100644
--- a/third_party/xla/xla/tools/hlo_isolation/BUILD
+++ b/third_party/xla/xla/tools/hlo_isolation/BUILD
@@ -56,6 +56,8 @@ cc_library(
         "//xla:literal",
         "//xla:literal_comparison",
         "//xla:shape_util",
+        "//xla:types",
+        "//xla:util",
         "//xla:xla_data_proto_cc",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass_pipeline",
@@ -168,8 +170,10 @@ xla_cc_test(
     deps = [
         ":hlo_inf_nan_intent_analyzer",
         ":hlo_isolation_api",
+        ":hlo_isolation_proto_cc",
         "//xla:literal",
         "//xla:literal_util",
+        "//xla:shape_util",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/parser:hlo_parser",
         "//xla/service:hlo_module_config",
diff --git a/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api.cc b/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api.cc
index 4bc08fb1b25..ac9eeb9ce3f 100644
--- a/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api.cc
+++ b/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api.cc
@@ -16,15 +16,18 @@ limitations under the License.
 #include "xla/tools/hlo_isolation/hlo_isolation_api.h"
 
 #include <algorithm>
+#include <cmath>
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
 #include <functional>
+#include <limits>
 #include <memory>
 #include <optional>
 #include <queue>
 #include <sstream>
 #include <string>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
@@ -38,6 +41,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
+#include "absl/strings/ascii.h"
 #include "absl/strings/match.h"
 #include "absl/strings/numbers.h"
 #include "absl/strings/str_cat.h"
@@ -60,9 +64,12 @@ limitations under the License.
 #include "xla/hlo/transforms/defuser.h"
 #include "xla/hlo/transforms/despecializer.h"
 #include "xla/hlo/transforms/simplifiers/hlo_memory_scheduler.h"
+#include "xla/index_util.h"
+#include "xla/layout_util.h"
 #include "xla/literal.h"
 #include "xla/literal_comparison.h"
 #include "xla/pjrt/pjrt_executable.h"
+#include "xla/primitive_util.h"
 #include "xla/service/hlo_module_config.h"
 #include "xla/service/hlo_runner_interface.h"
 #include "xla/shape.h"
@@ -75,6 +82,8 @@ limitations under the License.
 #include "xla/tools/hlo_module_loader.h"
 #include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/test.h"
+#include "xla/types.h"
+#include "xla/util.h"
 #include "xla/xla_data.pb.h"
 #include "tsl/platform/path.h"
 
@@ -230,6 +239,22 @@ void ApplyBoundingBoxToMismatch(
     }
     mismatch.set_mismatch_count(bbox.mismatch_count);
     mismatch.set_total_elements(bbox.total_elements);
+    // A mismatch parsed from an integer or subnormal flush message carries no
+    // percentages. Derive them from the mask so the report still shows the
+    // fraction of flushed elements.
+    if (bbox.total_elements > 0 && bbox.mismatch_count > 0) {
+      const double pct = 100.0 * static_cast<double>(bbox.mismatch_count) /
+                         static_cast<double>(bbox.total_elements);
+      if (!mismatch.has_percentage_of_elems_exceeding_both_errors()) {
+        mismatch.set_percentage_of_elems_exceeding_both_errors(pct);
+      }
+      if (!mismatch.has_percentage_of_elems_exceeding_abs_error()) {
+        mismatch.set_percentage_of_elems_exceeding_abs_error(pct);
+      }
+      if (!mismatch.has_percentage_of_elems_exceeding_rel_error()) {
+        mismatch.set_percentage_of_elems_exceeding_rel_error(pct);
+      }
+    }
   }
 }
 
@@ -261,6 +286,272 @@ auto MakeMiscompareCallback(
   };
 }
 
+template <typename T>
+bool IsSubnormal(T v) {
+  if constexpr (!has_positive_or_negative_zero_v<T> ||
+                std::numeric_limits<T>::has_denorm == std::denorm_absent) {
+    return false;
+  } else {
+    return v != T{0} && Eigen::numext::isfinite(v) &&
+           Eigen::numext::abs(v) < (std::numeric_limits<T>::min)();
+  }
+}
+
+template <typename T>
+bool IsSignedZero(T v) {
+  if constexpr (!has_positive_or_negative_zero_v<T>) {
+    return false;
+  } else {
+    return v == T{0};
+  }
+}
+
+// Result of scanning an expected/actual pair for subnormals that the device
+// flushed to signed zero. `mask` marks the flushed elements (PRED, same
+// dimensions as the array) and `first_index` is the linear index of the first
+// one, which `message` describes.
+struct SubnormalFlush {
+  int64_t first_index = 0;
+  std::string message;
+  Literal mask;
+};
+
+std::optional<SubnormalFlush> DetectSubnormalFlushInArray(
+    const Shape& shape, const uint8_t* expected_ptr,
+    const uint8_t* actual_ptr) {
+  const PrimitiveType element_type = shape.element_type();
+  if (!primitive_util::IsFloatingPointType(element_type)) {
+    return std::nullopt;
+  }
+  return primitive_util::FloatingPointTypeSwitch(
+      [&](auto kType) -> std::optional<SubnormalFlush> {
+        using NativeT = primitive_util::NativeTypeOf<kType>;
+        using UintT = std::conditional_t<
+            sizeof(NativeT) == 1, uint8_t,
+            std::conditional_t<
+                sizeof(NativeT) == 2, uint16_t,
+                std::conditional_t<sizeof(NativeT) == 4, uint32_t, uint64_t>>>;
+        constexpr absl::PadSpec kHexPad =
+            sizeof(NativeT) == 1   ? absl::kZeroPad2
+            : sizeof(NativeT) == 2 ? absl::kZeroPad4
+            : sizeof(NativeT) == 4 ? absl::kZeroPad8
+                                   : absl::kZeroPad16;
+        const int64_t num_elements = ShapeUtil::ElementsIn(shape);
+        std::optional<SubnormalFlush> flush;
+        for (int64_t i = 0; i < num_elements; ++i) {
+          NativeT exp_v, act_v;
+          std::memcpy(&exp_v, expected_ptr + i * sizeof(NativeT),
+                      sizeof(NativeT));
+          std::memcpy(&act_v, actual_ptr + i * sizeof(NativeT),
+                      sizeof(NativeT));
+          if (!IsSubnormal(exp_v) || !IsSignedZero(act_v)) {
+            continue;
+          }
+          if (!flush.has_value()) {
+            UintT exp_bits = 0, act_bits = 0;
+            std::memcpy(&exp_bits, &exp_v, sizeof(UintT));
+            std::memcpy(&act_bits, &act_v, sizeof(UintT));
+            std::string type_name = absl::AsciiStrToUpper(
+                primitive_util::LowercasePrimitiveTypeName(element_type));
+            flush.emplace();
+            flush->first_index = i;
+            flush->message = absl::StrCat(
+                "SUBNORMAL FLUSH-TO-ZERO DETECTED at element ", i,
+                ": Expected subnormal ", type_name, " 0x",
+                absl::Hex(exp_bits, kHexPad), ", but actual is signed zero 0x",
+                absl::Hex(act_bits, kHexPad));
+            flush->mask = Literal(ShapeUtil::ChangeElementType(shape, PRED));
+            flush->mask.PopulateWithValue(false);
+          }
+          flush->mask.data<bool>()[i] = true;
+        }
+        return flush;
+      },
+      element_type);
+}
+
+// Subnormal flushes found across all array leaves of a literal pair.
+struct SubnormalFlushReport {
+  // One line per flushed leaf, for the error status message.
+  std::string message;
+  // One entry per flushed leaf describing its first flushed element, in the
+  // same form that ExtractTopMismatches produces for a tolerance failure.
+  std::vector<NumericMismatch> mismatches;
+};
+
+// Scans every array leaf of `expected` against `actual`. For each leaf with a
+// flushed subnormal, invokes `callback` with the flush mask so the caller can
+// record bounding boxes for exactly those elements.
+void DetectSubnormalFlushInLiteral(
+    const LiteralSlice& expected, const LiteralSlice& actual,
+    const literal_comparison::MiscompareCallback& callback,
+    const ShapeIndex& shape_index, SubnormalFlushReport& report) {
+  if (expected.shape().IsTuple() && actual.shape().IsTuple()) {
+    const int64_t count = std::min(expected.shape().tuple_shapes_size(),
+                                   actual.shape().tuple_shapes_size());
+    for (int64_t i = 0; i < count; ++i) {
+      ShapeIndex next_index = shape_index;
+      next_index.push_back(i);
+      DetectSubnormalFlushInLiteral(LiteralSlice(expected, {i}),
+                                    LiteralSlice(actual, {i}), callback,
+                                    next_index, report);
+    }
+    return;
+  }
+  if (!expected.shape().IsArray() || !actual.shape().IsArray() ||
+      expected.shape().element_type() != actual.shape().element_type() ||
+      !ShapeUtil::SameDimensions(expected.shape(), actual.shape())) {
+    return;
+  }
+  std::optional<Literal> relayout_actual;
+  LiteralSlice actual_aligned = actual;
+  if (!LayoutUtil::Equal(expected.shape().layout(), actual.shape().layout())) {
+    relayout_actual = actual.Relayout(expected.shape().layout());
+    actual_aligned = LiteralSlice(*relayout_actual);
+  }
+  if (expected.size_bytes() != actual_aligned.size_bytes()) {
+    return;
+  }
+  std::optional<SubnormalFlush> flush = DetectSubnormalFlushInArray(
+      expected.shape(), static_cast<const uint8_t*>(expected.untyped_data()),
+      static_cast<const uint8_t*>(actual_aligned.untyped_data()));
+  if (!flush.has_value()) {
+    return;
+  }
+  if (callback) {
+    callback(expected, actual_aligned, LiteralSlice(flush->mask), shape_index,
+             literal_comparison::ErrorBuckets());
+  }
+  absl::StrAppend(
+      &report.message, report.message.empty() ? "" : "\n",
+      shape_index.empty()
+          ? ""
+          : absl::StrCat("Array at shape index ", shape_index.ToString(), ": "),
+      flush->message);
+  NumericMismatch& mismatch = report.mismatches.emplace_back();
+  mismatch.set_output_shape_index(shape_index.empty() ? 0
+                                                      : shape_index.front());
+  DimensionVector multi_index = IndexUtil::LinearIndexToMultidimensionalIndex(
+      expected.shape(), flush->first_index);
+  mismatch.set_expected(expected.GetAsDouble(multi_index).value_or(0.0));
+  mismatch.set_actual(0.0);
+  mismatch.set_rel_error(1.0);
+  for (int64_t coord : multi_index) {
+    mismatch.add_top_mismatch_index(coord);
+  }
+}
+
+// Returns a report of flushed subnormals across `expected` and `actual`, or
+// nullopt when nothing was flushed.
+std::optional<SubnormalFlushReport> DetectSubnormalFlushInLiteral(
+    const LiteralSlice& expected, const LiteralSlice& actual,
+    const literal_comparison::MiscompareCallback& callback = nullptr) {
+  SubnormalFlushReport report;
+  DetectSubnormalFlushInLiteral(expected, actual, callback, ShapeIndex(),
+                                report);
+  if (report.mismatches.empty()) {
+    return std::nullopt;
+  }
+  return report;
+}
+
+// Combines a tolerance comparison with the subnormal flush scan. A flushed
+// subnormal fails the comparison even when it is within `error_spec`; the
+// flush mask, not a zero tolerance rerun, drives the miscompare callback so
+// bounding boxes cover only the flushed elements. When the tolerance check
+// passed, `flush_mismatches` receives one entry per flushed leaf for the
+// numeric check, since the status message then has no `Near` mismatch lines.
+absl::Status NearWithSubnormalFlushCheck(
+    const LiteralSlice& expected, const LiteralSlice& actual,
+    const ErrorSpec& error_spec,
+    const literal_comparison::MiscompareCallback& callback,
+    std::vector<NumericMismatch>& flush_mismatches) {
+  absl::Status status = literal_comparison::Near(
+      expected, actual, error_spec, /*detailed_message=*/true, callback);
+  std::optional<SubnormalFlushReport> report = DetectSubnormalFlushInLiteral(
+      expected, actual, status.ok() ? callback : nullptr);
+  if (!report.has_value()) {
+    return status;
+  }
+  if (status.ok()) {
+    flush_mismatches = std::move(report->mismatches);
+  }
+  return absl::InternalError(
+      absl::StrCat(status.message(), "\n", report->message));
+}
+
+bool LiteralContainsSubnormal(const LiteralSlice& literal) {
+  if (literal.shape().IsTuple()) {
+    for (int64_t i = 0; i < literal.shape().tuple_shapes_size(); ++i) {
+      if (LiteralContainsSubnormal(LiteralSlice(literal, {i}))) {
+        return true;
+      }
+    }
+    return false;
+  }
+  if (!literal.shape().IsArray() ||
+      !primitive_util::IsFloatingPointType(literal.shape().element_type())) {
+    return false;
+  }
+  const auto* ptr = static_cast<const uint8_t*>(literal.untyped_data());
+  const int64_t num_elements = ShapeUtil::ElementsIn(literal.shape());
+  return primitive_util::FloatingPointTypeSwitch(
+      [&](auto kType) -> bool {
+        using NativeT = primitive_util::NativeTypeOf<kType>;
+        for (int64_t i = 0; i < num_elements; ++i) {
+          NativeT val;
+          std::memcpy(&val, ptr + i * sizeof(NativeT), sizeof(NativeT));
+          if (IsSubnormal(val)) {
+            return true;
+          }
+        }
+        return false;
+      },
+      literal.shape().element_type());
+}
+
+bool InputsContainSubnormal(absl::Span<const Literal> inputs) {
+  for (const Literal& input : inputs) {
+    if (LiteralContainsSubnormal(LiteralSlice(input))) {
+      return true;
+    }
+  }
+  return false;
+}
+
+absl::StatusOr<Literal> CoerceActualLiteralToExpectedShape(
+    const Literal& actual, const Shape& expected_shape) {
+  Shape host_expected = ShapeUtil::DeviceShapeToHostShape(expected_shape);
+  if (ShapeUtil::Compatible(host_expected, actual.shape())) {
+    return actual.Clone();
+  }
+  if (actual.shape().IsArray() && host_expected.IsArray() &&
+      ShapeUtil::SameDimensions(host_expected, actual.shape()) &&
+      actual.shape().element_type() == F32) {
+    const PrimitiveType expected_type = host_expected.element_type();
+    if (expected_type == BF16 || primitive_util::IsF8Type(expected_type) ||
+        expected_type == F4E2M1FN) {
+      return actual.Convert(expected_type);
+    }
+    if (primitive_util::IsIntegralType(expected_type)) {
+      const int bit_width = primitive_util::BitWidth(expected_type);
+      if (bit_width == 32) {
+        return actual.BitcastConvert(host_expected);
+      }
+      if (bit_width < 32) {
+        ABSL_ASSIGN_OR_RETURN(Literal as_u32,
+                         actual.BitcastConvert(ShapeUtil::ChangeElementType(
+                             actual.shape(), U32)));
+        return as_u32.Convert(expected_type);
+      }
+    }
+  }
+  return absl::InvalidArgumentError(
+      absl::StrCat("Reference literal shape ", expected_shape.ToString(),
+                   " is not compatible with actual literal shape ",
+                   actual.shape().ToString()));
+}
+
 // Compares `test_output` against `reference_output` and records a NumericCheck
 // named `check_name` on `result`.
 //
@@ -283,9 +574,11 @@ absl::Status CompareOutputs(const HloModule& module, const Literal& test_output,
   ErrorSpec error_spec(options.abs_error_bound, options.rel_error_bound);
   absl::flat_hash_map<int64_t, numerics::debug_info::MismatchBoundingBox>
       computed_bboxes;
-  absl::Status status = literal_comparison::Near(
-      reference_output, test_output, error_spec, true,
-      MakeMiscompareCallback(std::string(literal_prefix), &computed_bboxes));
+  std::vector<NumericMismatch> flush_mismatches;
+  absl::Status status = NearWithSubnormalFlushCheck(
+      reference_output, test_output, error_spec,
+      MakeMiscompareCallback(std::string(literal_prefix), &computed_bboxes),
+      flush_mismatches);
   NumericCheck* numeric_check = result.add_numeric_checks();
   numeric_check->set_name(check_name);
   numeric_check->set_expected_contains_inf_or_nan(
@@ -298,7 +591,9 @@ absl::Status CompareOutputs(const HloModule& module, const Literal& test_output,
                         check_name, module.name(), status.message()));
     mismatch_messages.push_back(std::string(status.message()));
     absl::StatusOr<std::vector<NumericMismatch>> top_mismatches =
-        ExtractAndEnrichTopMismatches(std::string(status.message()), &module);
+        flush_mismatches.empty() ? ExtractAndEnrichTopMismatches(
+                                       std::string(status.message()), &module)
+                                 : std::move(flush_mismatches);
     if (top_mismatches.ok()) {
       for (NumericMismatch& mismatch : *top_mismatches) {
         ApplyBoundingBoxToMismatch(computed_bboxes, mismatch);
@@ -474,32 +769,28 @@ std::vector<HloOutputCallback> CreateComparisonHloOutputCallbacks(
               return;
             }
 
-            Literal expected_literal;
-            if (ShapeUtil::Compatible(expected_literal_ptr->shape(),
-                                      literals[0]->shape())) {
-              expected_literal = expected_literal_ptr->Clone();
-            } else {
-              LOG(WARNING) << "Reference literal shape "
-                           << expected_literal_ptr->shape().ToString()
-                           << " is not compatible with actual literal shape "
-                           << literals[0]->shape().ToString() << " for op "
+            Literal expected_literal = expected_literal_ptr->Clone();
+            absl::StatusOr<Literal> actual_literal_or =
+                CoerceActualLiteralToExpectedShape(*literals[0],
+                                                   expected_literal.shape());
+            if (!actual_literal_or.ok()) {
+              LOG(WARNING) << actual_literal_or.status().message() << " for op "
                            << op_name << " within fusion " << module_name;
               return;
             }
+            const Literal& actual_literal = *actual_literal_or;
 
             absl::flat_hash_map<int64_t,
                                 numerics::debug_info::MismatchBoundingBox>
                 computed_bboxes;
             xla::ErrorSpec error_spec(static_cast<float>(abs_error),
                                       static_cast<float>(rel_error));
-            absl::Status matched = xla::literal_comparison::Near(
-                /*expected=*/expected_literal,
-                /*actual=*/*literals[0],
-                /*error=*/error_spec,
-                /*detailed_message=*/true,
-                /*miscompare_callback=*/
+            std::vector<NumericMismatch> flush_mismatches;
+            absl::Status matched = NearWithSubnormalFlushCheck(
+                expected_literal, actual_literal, error_spec,
                 MakeMiscompareCallback(absl::StrCat(module_name, "-", op_name),
-                                       &computed_bboxes));
+                                       &computed_bboxes),
+                flush_mismatches);
 
             if (!matched.ok()) {
               std::string error_message = absl::StrFormat(
@@ -515,11 +806,13 @@ std::vector<HloOutputCallback> CreateComparisonHloOutputCallbacks(
               numeric_check->set_expected_contains_inf_or_nan(
                   LiteralContainsInfOrNan(expected_literal));
               numeric_check->set_actual_contains_inf_or_nan(
-                  LiteralContainsInfOrNan(*literals[0]));
+                  LiteralContainsInfOrNan(actual_literal));
 
               absl::StatusOr<std::vector<NumericMismatch>> top_mismatches =
-                  ExtractTopMismatches(std::string(matched.message()),
-                                       literals[0]->shape().IsTuple());
+                  flush_mismatches.empty()
+                      ? ExtractTopMismatches(std::string(matched.message()),
+                                             actual_literal.shape().IsTuple())
+                      : std::move(flush_mismatches);
               if (top_mismatches.ok()) {
                 for (NumericMismatch& mismatch : *top_mismatches) {
                   ApplyBoundingBoxToMismatch(computed_bboxes, mismatch);
@@ -719,47 +1012,8 @@ absl::StatusOr<HloIsolationTestResult> RunIsolationTestOnModule(
                  << defused_output.status();
   }
 
-  if (defused_output.ok()) {
-    // Compare Test vs Defused Test.
-    absl::Status compare_status =
-        CompareOutputs(module, test_literal, *defused_output, result, options,
-                       "TPU_VS_DEFUSED_TPU", module.name(), mismatch_messages);
-    if (compare_status.ok()) {
-      result.set_state(State::SUCCESS);
-      result.set_reason("STAGE_1_DEFUSED_TPU_SUCCESS");
-      return result;
-    }
-
-    // Compare Test vs Defused Test again, this time with excess precision
-    // disabled on both sides.
-    //
-    // XLA lets a backend evaluate an instruction in a wider type than the HLO
-    // declares. A fusion can therefore hold an intermediate in f32 where the
-    // defused module, which materializes that intermediate in its declared
-    // (say bf16) type, is forced to round. Both results are correct, but they
-    // differ, and the difference is unbounded relative to the reference
-    // wherever the rounded reference cancels to exactly zero. Pinning both
-    // sides to the declared types removes that degree of freedom.
-    //
-    // Note that this recompiles the module under test, not just the reference,
-    // so it cannot distinguish "the two runs merely rounded differently" from
-    // "the backend miscompiles this fusion only on the excess-precision path".
-    // Set retry_without_excess_precision to false to investigate the latter.
-    //
-    // This is best effort: any failure to produce the comparison leaves the
-    // original mismatch standing.
-    if (options.retry_without_excess_precision &&
-        module.config().debug_options().xla_allow_excess_precision() &&
-        RetryWithoutExcessPrecision(module, test_runner, input_data, result,
-                                    options, mismatch_messages)) {
-      result.set_state(State::SUCCESS);
-      result.set_reason("STAGE_1B_NO_EXCESS_PRECISION_SUCCESS");
-      return result;
-    }
-  }
-
   // Potentially skip reference run.
-  if (options.max_module_size_bytes > 0) {
+  if (reference_runner != nullptr && options.max_module_size_bytes > 0) {
     int64_t size = options.estimate_module_size_fn(module);
     if (size > options.max_module_size_bytes) {
       LOG(INFO) << "Skipping reference run for module: " << module.name()
@@ -773,6 +1027,34 @@ absl::StatusOr<HloIsolationTestResult> RunIsolationTestOnModule(
     reference_runner = nullptr;
   }
 
+  // The backend may flush subnormals to zero in both the fused and the
+  // defused run, so a matching stage 1 proves nothing about subnormal inputs.
+  // Only skip the stage 1 short circuit when a reference will actually run.
+  const bool inputs_have_subnormals =
+      reference_runner != nullptr && InputsContainSubnormal(input_data);
+  if (defused_output.ok()) {
+    // Compare Test vs Defused Test.
+    absl::Status compare_status =
+        CompareOutputs(module, test_literal, *defused_output, result, options,
+                       "TPU_VS_DEFUSED_TPU", module.name(), mismatch_messages);
+    if (compare_status.ok()) {
+      if (!inputs_have_subnormals) {
+        result.set_state(State::SUCCESS);
+        result.set_reason("STAGE_1_DEFUSED_TPU_SUCCESS");
+        return result;
+      }
+    } else if (!inputs_have_subnormals &&
+               options.retry_without_excess_precision &&
+               module.config().debug_options().xla_allow_excess_precision() &&
+               RetryWithoutExcessPrecision(module, test_runner, input_data,
+                                           result, options,
+                                           mismatch_messages)) {
+      result.set_state(State::SUCCESS);
+      result.set_reason("STAGE_1B_NO_EXCESS_PRECISION_SUCCESS");
+      return result;
+    }
+  }
+
   // The defused run exhausted device memory (or timed out), which means the
   // module's materialized intermediates do not fit. The interpreter reference
   // materializes those same intermediates on the host and is orders of
@@ -1193,6 +1475,36 @@ absl::StatusOr<std::vector<NumericMismatch>> ExtractTopMismatches(
   if (current_mismatch.has_value()) {
     mismatches.push_back(std::move(*current_mismatch));
   }
+  if (mismatches.empty()) {
+    std::string index_str, expected_str, actual_str;
+    if (RE2::PartialMatch(
+            error_message,
+            R"(first mismatch at array index\s*\{([^}]*)\}:\s*expected value:\s*([^\s]+)\s*actual value:\s*([^\s]+))",
+            &index_str, &expected_str, &actual_str)) {
+      double expected_double = 0.0;
+      double actual_double = 0.0;
+      if (absl::SimpleAtod(expected_str, &expected_double) &&
+          absl::SimpleAtod(actual_str, &actual_double)) {
+        NumericMismatch data;
+        data.set_output_shape_index(shape_index.value_or(0));
+        data.set_expected(expected_double);
+        data.set_actual(actual_double);
+        const double abs_diff = std::abs(actual_double - expected_double);
+        const double rel_err = expected_double != 0.0
+                                   ? abs_diff / std::abs(expected_double)
+                                   : (actual_double != 0.0 ? 1.0 : 0.0);
+        data.set_rel_error(rel_err);
+        for (absl::string_view idx_part :
+             absl::StrSplit(index_str, ',', absl::SkipEmpty())) {
+          int64_t coord;
+          if (absl::SimpleAtoi(absl::StripAsciiWhitespace(idx_part), &coord)) {
+            data.add_top_mismatch_index(coord);
+          }
+        }
+        mismatches.push_back(std::move(data));
+      }
+    }
+  }
   return mismatches;
 }
 
diff --git a/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api_test.cc b/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api_test.cc
index feced371f14..8404b99c4c2 100644
--- a/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api_test.cc
+++ b/third_party/xla/xla/tools/hlo_isolation/hlo_isolation_api_test.cc
@@ -18,8 +18,11 @@ limitations under the License.
 #include <cstdint>
 #include <limits>
 #include <memory>
+#include <string>
+#include <utility>
 #include <vector>
 
+#include <gtest/gtest-spi.h>
 #include <gtest/gtest.h>
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
@@ -28,12 +31,14 @@ limitations under the License.
 #include "absl/types/span.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/parser/hlo_parser.h"
+#include "xla/layout_util.h"
 #include "xla/literal.h"
 #include "xla/literal_util.h"
 #include "xla/service/hlo_module_config.h"
 #include "xla/service/hlo_runner_interface.h"
 #include "xla/tools/hlo_dump/hlo_dump_utils.h"
 #include "xla/tools/hlo_isolation/hlo_inf_nan_intent_analyzer.h"
+#include "xla/tools/hlo_isolation/hlo_isolation.pb.h"
 #include "xla/tsl/platform/test.h"
 
 namespace xla {
@@ -435,6 +440,212 @@ ENTRY main {
   EXPECT_EQ(result.reason(), "DEFUSED_REFERENCE_UNAVAILABLE");
 }
 
+TEST(HloIsolationApiTest, SubnormalFlushToZeroTrappedUnderDefaultTolerance) {
+  const absl::string_view kHlo = R"hlo(
+HloModule subnormal_flush
+ENTRY main {
+  p0 = f32[4] parameter(0)
+  ROOT add = f32[4] add(p0, p0)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       xla::ParseAndReturnUnverifiedModule(kHlo));
+
+  const float subnormal_f32 = std::numeric_limits<float>::denorm_min();
+  std::vector<Literal> inputs;
+  inputs.push_back(
+      LiteralUtil::CreateR1<float>({subnormal_f32, 1.0f, subnormal_f32, 1.0f}));
+
+  ModuleIsolationOptions options;
+  // Both fused and defused device runs flush the subnormal to 0.0f, while the
+  // interpreter reference preserves 2 * denorm_min(). Even under the default
+  // tolerance (abs=0.01, rel=0.1), the subnormal flush to zero must fail. The
+  // normal elements differ by rounding noise that the tolerance accepts and
+  // must not be counted as mismatches.
+  options.run_module_fn =
+      [subnormal_f32](
+          std::unique_ptr<HloModule> m, HloRunnerInterface* r,
+          absl::Span<const Literal> i,
+          const RunModuleOptions& run_opts) -> absl::StatusOr<Literal> {
+    if (r == FakeReferenceRunner()) {
+      const float expected = subnormal_f32 * 2.0f;
+      Literal ref =
+          LiteralUtil::CreateR1<float>({expected, 2.0f, expected, 2.0f});
+      if (run_opts.expected_literals != nullptr) {
+        (*run_opts.expected_literals)["add"] =
+            std::make_shared<Literal>(ref.Clone());
+      }
+      return ref;
+    }
+    Literal flushed = LiteralUtil::CreateR1<float>({0.0f, 2.001f, 0.0f, 2.0f});
+    for (const auto& cb : run_opts.hlo_output_callbacks) {
+      auto ptr = std::make_shared<Literal>(flushed.Clone());
+      cb.callback(0, 0, {ptr});
+    }
+    return flushed;
+  };
+  int mismatch_reports = 0;
+  std::string reported_message;
+  options.on_mismatch_fn = [&](const HloModule& m, const Literal& test_output,
+                               const Literal& ref_output,
+                               const absl::Status& status) {
+    ++mismatch_reports;
+    reported_message = std::string(status.message());
+  };
+
+  ::testing::TestPartResultArray failures;
+  HloIsolationTestResult result;
+  {
+    ::testing::ScopedFakeTestPartResultReporter reporter(
+        ::testing::ScopedFakeTestPartResultReporter::INTERCEPT_ALL_THREADS,
+        &failures);
+    ASSERT_OK_AND_ASSIGN(
+        result, RunIsolationTestOnModule(
+                    *module, nullptr, FakeReferenceRunner(), options, inputs));
+  }
+  EXPECT_EQ(result.state(), State::FAILURE);
+  EXPECT_EQ(result.reason(), "NUMERIC_MISMATCH");
+  EXPECT_EQ(mismatch_reports, 1);
+  EXPECT_TRUE(
+      absl::StrContains(reported_message, "SUBNORMAL FLUSH-TO-ZERO DETECTED"));
+  EXPECT_EQ(failures.size(), 1);
+
+  // The interpreter check and the fusion debugger check both record only the
+  // two flushed elements, not the rounding noise on the normal elements. The
+  // defused device run flushed the same elements, so that check passes.
+  int flush_checks = 0;
+  for (const NumericCheck& check : result.numeric_checks()) {
+    if (check.name() == "TPU_VS_DEFUSED_TPU") {
+      EXPECT_EQ(check.top_mismatches_size(), 0);
+      continue;
+    }
+    ++flush_checks;
+    ASSERT_EQ(check.top_mismatches_size(), 1) << check.name();
+    const NumericMismatch& mismatch = check.top_mismatches(0);
+    EXPECT_DOUBLE_EQ(mismatch.expected(), subnormal_f32 * 2.0f) << check.name();
+    EXPECT_DOUBLE_EQ(mismatch.actual(), 0.0) << check.name();
+    EXPECT_EQ(mismatch.mismatch_count(), 2) << check.name();
+    EXPECT_EQ(mismatch.total_elements(), 4) << check.name();
+    EXPECT_DOUBLE_EQ(mismatch.percentage_of_elems_exceeding_both_errors(), 50.0)
+        << check.name();
+    ASSERT_EQ(mismatch.top_mismatch_index_size(), 1) << check.name();
+    EXPECT_EQ(mismatch.top_mismatch_index(0), 0) << check.name();
+    ASSERT_EQ(mismatch.mismatch_box_min_size(), 1) << check.name();
+    EXPECT_EQ(mismatch.mismatch_box_min(0), 0) << check.name();
+    EXPECT_EQ(mismatch.mismatch_box_max(0), 2) << check.name();
+  }
+  EXPECT_EQ(flush_checks, 2);
+}
+
+TEST(HloIsolationApiTest, SubnormalInputsStillSucceedWhenReferenceIsSkipped) {
+  const absl::string_view kHlo = R"hlo(
+HloModule subnormal_inputs_no_reference
+ENTRY main {
+  p0 = f32[4] parameter(0)
+  ROOT add = f32[4] add(p0, p0)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       xla::ParseAndReturnUnverifiedModule(kHlo));
+
+  const float subnormal_f32 = std::numeric_limits<float>::denorm_min();
+  std::vector<Literal> inputs;
+  inputs.push_back(LiteralUtil::CreateR1<float>(
+      {subnormal_f32, subnormal_f32, subnormal_f32, subnormal_f32}));
+
+  ModuleIsolationOptions options;
+  // The size gate removes the reference runner, so a matching defused run is
+  // the only evidence available and the module must be accepted as before.
+  options.max_module_size_bytes = 1;
+  options.estimate_module_size_fn = [](const HloModule&) { return 2; };
+  int reference_runs = 0;
+  options.run_module_fn =
+      [&reference_runs](
+          std::unique_ptr<HloModule> m, HloRunnerInterface* r,
+          absl::Span<const Literal> i,
+          const RunModuleOptions& run_opts) -> absl::StatusOr<Literal> {
+    if (r == FakeReferenceRunner()) {
+      ++reference_runs;
+    }
+    return LiteralUtil::CreateR1<float>({0.0f, 0.0f, 0.0f, 0.0f});
+  };
+  int mismatch_reports = 0;
+  options.on_mismatch_fn = [&](const HloModule&, const Literal&, const Literal&,
+                               const absl::Status&) { ++mismatch_reports; };
+
+  ASSERT_OK_AND_ASSIGN(
+      HloIsolationTestResult result,
+      RunIsolationTestOnModule(*module, nullptr, FakeReferenceRunner(), options,
+                               inputs));
+  EXPECT_EQ(result.state(), State::SUCCESS);
+  EXPECT_EQ(result.reason(), "STAGE_1_DEFUSED_TPU_SUCCESS");
+  EXPECT_EQ(reference_runs, 0);
+  EXPECT_EQ(mismatch_reports, 0);
+}
+
+TEST(HloIsolationApiTest, SubnormalFlushRespectsNonDefaultLayouts) {
+  const absl::string_view kHlo = R"hlo(
+HloModule subnormal_layout_test
+ENTRY main {
+  p0 = f32[2,2]{0,1} parameter(0)
+  ROOT add = f32[2,2]{0,1} add(p0, p0)
+}
+)hlo";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       xla::ParseAndReturnUnverifiedModule(kHlo));
+
+  const float subnormal_f32 = std::numeric_limits<float>::denorm_min();
+  Literal input = LiteralUtil::CreateR2WithLayout<float>(
+      {{1.0f, subnormal_f32}, {1.0f, 1.0f}}, LayoutUtil::MakeLayout({0, 1}));
+  std::vector<Literal> inputs;
+  inputs.push_back(std::move(input));
+
+  ModuleIsolationOptions options;
+  options.run_module_fn =
+      [subnormal_f32](
+          std::unique_ptr<HloModule> m, HloRunnerInterface* r,
+          absl::Span<const Literal> i,
+          const RunModuleOptions& run_opts) -> absl::StatusOr<Literal> {
+    if (r == FakeReferenceRunner()) {
+      return LiteralUtil::CreateR2WithLayout<float>(
+          {{2.0f, subnormal_f32 * 2.0f}, {2.0f, 2.0f}},
+          LayoutUtil::MakeLayout({0, 1}));
+    }
+    return LiteralUtil::CreateR2WithLayout<float>(
+        {{2.0f, 0.0f}, {2.0f, 2.0f}}, LayoutUtil::MakeLayout({1, 0}));
+  };
+  options.on_mismatch_fn = [](const HloModule&, const Literal&, const Literal&,
+                              const absl::Status&) {};
+
+  ASSERT_OK_AND_ASSIGN(
+      HloIsolationTestResult result,
+      RunIsolationTestOnModule(*module, nullptr, FakeReferenceRunner(), options,
+                               inputs));
+  EXPECT_EQ(result.state(), State::FAILURE);
+  EXPECT_EQ(result.reason(), "NUMERIC_MISMATCH");
+
+  bool found_interp_check = false;
+  for (const NumericCheck& check : result.numeric_checks()) {
+    if (check.name() != "TPU_VS_INTERPRETER") {
+      continue;
+    }
+    found_interp_check = true;
+    ASSERT_EQ(check.top_mismatches_size(), 1);
+    const NumericMismatch& mismatch = check.top_mismatches(0);
+    EXPECT_EQ(mismatch.mismatch_count(), 1);
+    EXPECT_EQ(mismatch.total_elements(), 4);
+    ASSERT_EQ(mismatch.top_mismatch_index_size(), 2);
+    EXPECT_EQ(mismatch.top_mismatch_index(0), 0);
+    EXPECT_EQ(mismatch.top_mismatch_index(1), 1);
+    ASSERT_EQ(mismatch.mismatch_box_min_size(), 2);
+    EXPECT_EQ(mismatch.mismatch_box_min(0), 0);
+    EXPECT_EQ(mismatch.mismatch_box_min(1), 1);
+    EXPECT_EQ(mismatch.mismatch_box_max(0), 0);
+    EXPECT_EQ(mismatch.mismatch_box_max(1), 1);
+  }
+  EXPECT_TRUE(found_interp_check);
+}
+
 }  // namespace
 }  // namespace hlo_isolation
 }  // namespace xla

From f7e8dd9ca459c54aef52e972f127749e95d68afb Mon Sep 17 00:00:00 2001
From: Guoren Li <guoren@google.com>
Date: Thu, 15 Oct 2026 12:14:22 -0700
Subject: [PATCH 0888/1000] Expose PJRT C-API Megascale and multi-slice
 extension headers to Megascale users.

Add `//xla/megascale:internal` to the visibility of
`pjrt_c_api_megascale_extension_hdrs` and
`pjrt_c_api_multi_slice_extension_hdrs`, next to the `//xla:internal` they
already inherit. Visibility-only; OSS builds are unaffected because
`internal_visibility` resolves to public there.

PiperOrigin-RevId: 1000305466
---
 third_party/xla/xla/pjrt/c/BUILD | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/third_party/xla/xla/pjrt/c/BUILD b/third_party/xla/xla/pjrt/c/BUILD
index 514cb384a9a..14b3917f259 100644
--- a/third_party/xla/xla/pjrt/c/BUILD
+++ b/third_party/xla/xla/pjrt/c/BUILD
@@ -138,6 +138,10 @@ xla_cc_test(
 cc_library(
     name = "pjrt_c_api_megascale_extension_hdrs",
     hdrs = ["pjrt_c_api_megascale_extension.h"],
+    visibility = internal_visibility([
+        "//xla:internal",
+        "//xla/megascale:internal",
+    ]),
     deps = [
         ":pjrt_c_api_hdrs",
         ":pjrt_c_api_multi_slice_extension_hdrs",
@@ -147,6 +151,10 @@ cc_library(
 cc_library(
     name = "pjrt_c_api_multi_slice_extension_hdrs",
     hdrs = ["pjrt_c_api_multi_slice_extension.h"],
+    visibility = internal_visibility([
+        "//xla:internal",
+        "//xla/megascale:internal",
+    ]),
     deps = [":pjrt_c_api_hdrs"],
 )
 

From 9079ebed7a9913218f4d8e957094ba78383aaff9 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Sat, 10 Oct 2026 08:44:50 +0530
Subject: [PATCH 0889/1000] Refactor fatal logging to use _fatal function

---
 tensorflow/dtensor/python/heartbeat.py | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/tensorflow/dtensor/python/heartbeat.py b/tensorflow/dtensor/python/heartbeat.py
index d1f74c933f8..ad33d0540b5 100644
--- a/tensorflow/dtensor/python/heartbeat.py
+++ b/tensorflow/dtensor/python/heartbeat.py
@@ -23,6 +23,7 @@ the functions below are called, in a thread or not.
 """
 
 import atexit
+import os
 import threading
 import time
 
@@ -41,6 +42,19 @@ _failure_count = 0
 _heartbeat_timer = None
 
 
+def _fatal(msg, *args):
+  """Logs `msg` at FATAL level and unconditionally crashes the process.
+
+  `logging.fatal` (tf_logging -> the stdlib `logging` module) only logs at
+  CRITICAL level and returns normally; it does not terminate the process.
+  Since this can be called from the background heartbeat thread, raising
+  or calling `sys.exit` would only end that thread, not the process. Use
+  `os._exit` to guarantee a hard process crash regardless of caller.
+  """
+  logging.fatal(msg, *args)
+  os._exit(1)  # pylint: disable=protected-access
+
+
 def _heartbeat(
     period: int,  # in seconds
     timer: threading.Event,
@@ -87,13 +101,13 @@ def _heartbeat(
                         _CONSECUTIVE_FAILURES_LIMIT - _failure_count, e)
         continue
       else:
-        logging.fatal('Heartbeat failure %d, limit of %d reached: %s',
+        _fatal('Heartbeat failure %d, limit of %d reached: %s',
                       _failure_count, _CONSECUTIVE_FAILURES_LIMIT, e)
     logging.vlog(2, 'Received heartbeat signal %s', signal)
 
     # Out of sync workers will cause this, crash immediately.
     if not np.all(signal == token):
-      logging.fatal('Unexpected heartbeat signal received: %s', signal)
+      _fatal('Unexpected heartbeat signal received: %s', signal)
 
     # Any success resets the failure counter.
     _failure_count = 0
@@ -152,10 +166,10 @@ def start(period: int) -> threading.Event:
   # out of sync, and we should terminate all workers.
   if task_id == 0:
     if not np.all(signal == token):  # pyrefly: ignore[unbound-name]
-      logging.fatal('Merged heartbeat signal has value != %d', token)
+      _fatal('Merged heartbeat signal has value != %d', token)
   else:
     if len(set(signal)) != 1:
-      logging.fatal('Merged heartbeat signal has unequal elements')
+      _fatal('Merged heartbeat signal has unequal elements')
     token = signal[0]
 
   # On normal main process exit, set the timer to stop the heartbeat thread.

From 7e9c5d179172be0a4c2f56bf813a50a69a1c7675 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Sat, 10 Oct 2026 08:50:18 +0530
Subject: [PATCH 0890/1000] Fix formatting issues in heartbeat.py

---
 tensorflow/dtensor/python/heartbeat.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/dtensor/python/heartbeat.py b/tensorflow/dtensor/python/heartbeat.py
index ad33d0540b5..fbcf5260eb5 100644
--- a/tensorflow/dtensor/python/heartbeat.py
+++ b/tensorflow/dtensor/python/heartbeat.py
@@ -102,7 +102,7 @@ def _heartbeat(
         continue
       else:
         _fatal('Heartbeat failure %d, limit of %d reached: %s',
-                      _failure_count, _CONSECUTIVE_FAILURES_LIMIT, e)
+               _failure_count, _CONSECUTIVE_FAILURES_LIMIT, e)
     logging.vlog(2, 'Received heartbeat signal %s', signal)
 
     # Out of sync workers will cause this, crash immediately.

From 98e0670b21d30d38cfac37a7ef4528b92626d576 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 18:01:53 +0530
Subject: [PATCH 0891/1000] Update _fatal function to accept keyword arguments

---
 tensorflow/dtensor/python/heartbeat.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorflow/dtensor/python/heartbeat.py b/tensorflow/dtensor/python/heartbeat.py
index fbcf5260eb5..d49cc64f2c9 100644
--- a/tensorflow/dtensor/python/heartbeat.py
+++ b/tensorflow/dtensor/python/heartbeat.py
@@ -18,7 +18,7 @@ In normal cases, all workers will exchange the same randomly generated number
 until normal program termination. If any worker stops or restarts, other workers
 will detect that and crash themselves.
 
-In this module, logging.fatal is used to guarantee a worker crash no matter how
+In this module, _fatal is used to guarantee a worker crash no matter how
 the functions below are called, in a thread or not.
 """
 
@@ -42,7 +42,7 @@ _failure_count = 0
 _heartbeat_timer = None
 
 
-def _fatal(msg, *args):
+def _fatal(msg, *args, **kwargs):
   """Logs `msg` at FATAL level and unconditionally crashes the process.
 
   `logging.fatal` (tf_logging -> the stdlib `logging` module) only logs at
@@ -51,7 +51,7 @@ def _fatal(msg, *args):
   or calling `sys.exit` would only end that thread, not the process. Use
   `os._exit` to guarantee a hard process crash regardless of caller.
   """
-  logging.fatal(msg, *args)
+  logging.fatal(msg, *args, **kwargs)
   os._exit(1)  # pylint: disable=protected-access
 
 

From c97dfba9355ef762bf90809fa938b34310a27e20 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 18:03:25 +0530
Subject: [PATCH 0892/1000] Add unit tests for heartbeat module functionality

This test suite verifies the behavior of the heartbeat module, ensuring that fatal errors lead to process crashes and that logging functions as expected.
---
 .../dtensor/python/tests/heartbeat_test.py    | 57 +++++++++++++++++++
 1 file changed, 57 insertions(+)
 create mode 100644 tensorflow/dtensor/python/tests/heartbeat_test.py

diff --git a/tensorflow/dtensor/python/tests/heartbeat_test.py b/tensorflow/dtensor/python/tests/heartbeat_test.py
new file mode 100644
index 00000000000..7212a91f99b
--- /dev/null
+++ b/tensorflow/dtensor/python/tests/heartbeat_test.py
@@ -0,0 +1,57 @@
+# 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 tensorflow.dtensor.python.heartbeat."""
+
+from unittest import mock
+
+# pylint: disable=g-direct-tensorflow-import
+from tensorflow.dtensor.python import heartbeat
+from tensorflow.python.platform import test as tf_test
+# pylint: enable=g-direct-tensorflow-import
+
+
+class HeartbeatFatalTest(tf_test.TestCase):
+
+  def test_fatal_crashes_the_process(self):
+    with mock.patch.object(heartbeat.os, '_exit') as mock_exit:
+      with mock.patch.object(heartbeat.logging, 'fatal') as mock_log_fatal:
+        heartbeat._fatal('Heartbeat failure %d, limit of %d reached: %s',
+                         1, 3, 'boom')
+
+    mock_log_fatal.assert_called_once_with(
+        'Heartbeat failure %d, limit of %d reached: %s', 1, 3, 'boom')
+    mock_exit.assert_called_once_with(1)
+
+  def test_fatal_still_crashes_if_logging_raises(self):
+    # _fatal must guarantee a process crash even if the logging call itself
+    # fails, since this is the last-resort path for a desynced worker.
+    with mock.patch.object(heartbeat.os, '_exit') as mock_exit:
+      with mock.patch.object(
+          heartbeat.logging, 'fatal', side_effect=ValueError('bad format')):
+        with self.assertRaises(ValueError):
+          heartbeat._fatal('Unexpected heartbeat signal received: %s', [1, 0])
+
+    mock_exit.assert_not_called()
+
+  def test_fatal_forwards_kwargs(self):
+    with mock.patch.object(heartbeat.os, '_exit'):
+      with mock.patch.object(heartbeat.logging, 'fatal') as mock_log_fatal:
+        heartbeat._fatal('msg', exc_info=True)
+
+    mock_log_fatal.assert_called_once_with('msg', exc_info=True)
+
+
+if __name__ == '__main__':
+  tf_test.main()

From d9b8aff2faead67a423f46583ffdbcd55324e547 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 19:19:43 +0530
Subject: [PATCH 0893/1000] Guarantee _fatal crashes even if logging.fatal
 raises; add BUILD rule

- Wrap the logging.fatal() call in _fatal() with try/finally so
  os._exit(1) always runs, even if logging.fatal itself raises (e.g.
  an invalid format string). Previously such an exception would only
  kill the background heartbeat thread, leaving the worker process
  hanging instead of crashing.
- Update test_fatal_still_crashes_if_logging_raises to assert the new
  (correct) behavior: the ValueError still propagates via
  assertRaisesRegex, but os._exit(1) is now called once beforehand.
- Add the missing py_test BUILD rule for heartbeat_test.py so these
  tests actually run in CI.
---
 tensorflow/dtensor/python/heartbeat.py            |  6 ++++--
 tensorflow/dtensor/python/tests/BUILD             | 12 ++++++++++++
 tensorflow/dtensor/python/tests/heartbeat_test.py |  4 ++--
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/tensorflow/dtensor/python/heartbeat.py b/tensorflow/dtensor/python/heartbeat.py
index d49cc64f2c9..c9736cef352 100644
--- a/tensorflow/dtensor/python/heartbeat.py
+++ b/tensorflow/dtensor/python/heartbeat.py
@@ -51,8 +51,10 @@ def _fatal(msg, *args, **kwargs):
   or calling `sys.exit` would only end that thread, not the process. Use
   `os._exit` to guarantee a hard process crash regardless of caller.
   """
-  logging.fatal(msg, *args, **kwargs)
-  os._exit(1)  # pylint: disable=protected-access
+  try:
+    logging.fatal(msg, *args, **kwargs)
+  finally:
+    os._exit(1)  # pylint: disable=protected-access
 
 
 def _heartbeat(
diff --git a/tensorflow/dtensor/python/tests/BUILD b/tensorflow/dtensor/python/tests/BUILD
index cdab09849ba..410424bf9aa 100644
--- a/tensorflow/dtensor/python/tests/BUILD
+++ b/tensorflow/dtensor/python/tests/BUILD
@@ -109,6 +109,18 @@ py_test(
     ],
 )
 
+py_test(
+    name = "heartbeat_test",
+    srcs = [
+        "heartbeat_test.py",
+    ],
+    strict_deps = True,
+    deps = [
+        "//tensorflow/dtensor/python:heartbeat",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
 # TODO(b/301286466): Investigate why python annotation type mismatch is not catptured by the type
 # strict BUILD rules.
 
diff --git a/tensorflow/dtensor/python/tests/heartbeat_test.py b/tensorflow/dtensor/python/tests/heartbeat_test.py
index 7212a91f99b..6c231f9b5a9 100644
--- a/tensorflow/dtensor/python/tests/heartbeat_test.py
+++ b/tensorflow/dtensor/python/tests/heartbeat_test.py
@@ -40,10 +40,10 @@ class HeartbeatFatalTest(tf_test.TestCase):
     with mock.patch.object(heartbeat.os, '_exit') as mock_exit:
       with mock.patch.object(
           heartbeat.logging, 'fatal', side_effect=ValueError('bad format')):
-        with self.assertRaises(ValueError):
+        with self.assertRaisesRegex(ValueError, r'bad format'):
           heartbeat._fatal('Unexpected heartbeat signal received: %s', [1, 0])
 
-    mock_exit.assert_not_called()
+    mock_exit.assert_called_once_with(1)
 
   def test_fatal_forwards_kwargs(self):
     with mock.patch.object(heartbeat.os, '_exit'):

From 77abf4ce6761058fad91c81c83819345d2c0b7e7 Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 18:44:08 +0530
Subject: [PATCH 0894/1000] Improve control flow context management in TPU
 utility

Refactor context management to ensure proper restoration of control flow context using try-finally.
---
 tensorflow/python/distribute/tpu_util.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tensorflow/python/distribute/tpu_util.py b/tensorflow/python/distribute/tpu_util.py
index 3167e069df3..fccc346ed07 100644
--- a/tensorflow/python/distribute/tpu_util.py
+++ b/tensorflow/python/distribute/tpu_util.py
@@ -54,8 +54,10 @@ def outside_or_skip_tpu_context():
   else:
     saved_context = graph._get_control_flow_context()  # pylint: disable=protected-access
     graph._set_control_flow_context(ctx.outer_context)  # pylint: disable=protected-access
-    yield
-    graph._set_control_flow_context(saved_context)  # pylint: disable=protected-access
+    try:
+      yield
+    finally:
+      graph._set_control_flow_context(saved_context)  # pylint: disable=protected-access
 
 
 @contextlib.contextmanager

From 87536405a35e3c1ee512aada690639b0c2aa5d6a Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Tue, 13 Oct 2026 21:03:16 +0530
Subject: [PATCH 0895/1000] Address review: add regression test for
 control-flow context restore

Add tpu_util_test.py with
test_outside_or_skip_tpu_context_restores_on_exception, verifying the
graph's control-flow context is restored to the TPUReplicateContext
even when an exception is raised inside outside_or_skip_tpu_context(),
plus test_outside_or_skip_tpu_context_no_context for the no-enclosing-
context case. Registers the new test target in
python/distribute/BUILD.
---
 tensorflow/python/distribute/BUILD            | 12 +++++
 tensorflow/python/distribute/tpu_util_test.py | 54 +++++++++++++++++++
 2 files changed, 66 insertions(+)
 create mode 100644 tensorflow/python/distribute/tpu_util_test.py

diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD
index b5e63d595ad..69f9c6e0017 100644
--- a/tensorflow/python/distribute/BUILD
+++ b/tensorflow/python/distribute/BUILD
@@ -350,6 +350,18 @@ py_library(
     ],
 )
 
+distribute_py_strict_test(
+    name = "tpu_util_test",
+    srcs = ["tpu_util_test.py"],
+    deps = [
+        ":tpu_util",
+        "//tensorflow/python/eager:test",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/ops:control_flow_ops",
+        "//tensorflow/python/tpu:tpu_replication",
+    ],
+)
+
 py_library(
     name = "mirrored_strategy",
     srcs = ["mirrored_strategy.py"],
diff --git a/tensorflow/python/distribute/tpu_util_test.py b/tensorflow/python/distribute/tpu_util_test.py
new file mode 100644
index 00000000000..5e42747eb3e
--- /dev/null
+++ b/tensorflow/python/distribute/tpu_util_test.py
@@ -0,0 +1,54 @@
+# Copyright 2023 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 tpu_util."""
+
+from tensorflow.python.distribute import tpu_util
+from tensorflow.python.eager import test
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.tpu import tpu_replication
+
+
+class TPUUtilTest(test.TestCase):
+
+  def test_outside_or_skip_tpu_context_restores_on_exception(self):
+    with ops.Graph().as_default() as graph:
+      pivot = control_flow_ops.no_op()
+      context = tpu_replication.TPUReplicateContext(
+          b"tpu_context", 1, pivot=pivot
+      )
+      context.Enter()
+      try:
+        self.assertEqual(graph._get_control_flow_context(), context)
+        with self.assertRaisesRegex(RuntimeError, "simulated failure"):
+          with tpu_util.outside_or_skip_tpu_context():
+            self.assertEqual(
+                graph._get_control_flow_context(), context.outer_context
+            )
+            raise RuntimeError("simulated failure")
+        self.assertEqual(graph._get_control_flow_context(), context)
+      finally:
+        context.Exit()
+
+  def test_outside_or_skip_tpu_context_no_context(self):
+    with ops.Graph().as_default() as graph:
+      self.assertIsNone(graph._get_control_flow_context())
+      with tpu_util.outside_or_skip_tpu_context():
+        self.assertIsNone(graph._get_control_flow_context())
+      self.assertIsNone(graph._get_control_flow_context())
+
+
+if __name__ == "__main__":
+  test.main()

From a05e6f7964b82baf345dee8a3c90c89542be93e6 Mon Sep 17 00:00:00 2001
From: Gerardo Carranza <gcarranza@google.com>
Date: Thu, 15 Oct 2026 12:19:59 -0700
Subject: [PATCH 0896/1000] Mark decomposition subgraphs as
 delegation-skippable in YNNPackDelegate.

PiperOrigin-RevId: 1000308956
---
 .../lite/delegates/ynnpack/ynnpack_delegate.cc     | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/tensorflow/lite/delegates/ynnpack/ynnpack_delegate.cc b/tensorflow/lite/delegates/ynnpack/ynnpack_delegate.cc
index 15c2bf3dceb..b134b1d60d1 100644
--- a/tensorflow/lite/delegates/ynnpack/ynnpack_delegate.cc
+++ b/tensorflow/lite/delegates/ynnpack/ynnpack_delegate.cc
@@ -312,6 +312,20 @@ class YNNPackDelegateKernel : public SimpleDelegateKernelInterface {
       } else if (IsMoe(reg, node)) {
         node_info.composite_op_type = CompositeOpType::kMoe;
       }
+      if (node_info.composite_op_type != CompositeOpType::kNone) {
+        auto* subgraph = reinterpret_cast<tflite::Subgraph*>(context->impl_);
+        const auto* composite_params =
+            static_cast<const TfLiteStablehloCompositeParams*>(
+                node->builtin_data);
+        if (subgraph != nullptr && subgraph->GetSubgraphs() != nullptr &&
+            composite_params != nullptr &&
+            composite_params->subgraph_index > 0 &&
+            static_cast<size_t>(composite_params->subgraph_index) <
+                subgraph->GetSubgraphs()->size()) {
+          TF_LITE_ENSURE_STATUS(subgraph->MarkSubgraphAsDelegationSkippable(
+              composite_params->subgraph_index));
+        }
+      }
       nodes_info_.push_back(node_info);
     }
 

From a14be9af9a239ddb38c8b427ba2c77ae6481b086 Mon Sep 17 00:00:00 2001
From: David Dunleavy <ddunleavy@google.com>
Date: Thu, 15 Oct 2026 12:20:11 -0700
Subject: [PATCH 0897/1000] Move `core/tpu:tpu_compile_interface` and TPU
 embedding proto utils to XLA

PiperOrigin-RevId: 1000309060
---
 tensorflow/core/tpu/BUILD                     |  24 +-
 tensorflow/core/tpu/tpu_compile_interface.h   |  29 +-
 .../tpu/tpu_embedding_configuration_utils.h   |  18 +-
 ...embedding_optimization_parameters_utils.cc | 447 -----------------
 ..._embedding_optimization_parameters_utils.h |  95 +---
 third_party/xla/xla/tpu/BUILD                 |  13 +
 third_party/xla/xla/tpu/protobuf/BUILD        |  37 ++
 .../tpu_embedding_configuration_utils.cc      |  10 +-
 .../tpu_embedding_configuration_utils.h       |  37 ++
 ...embedding_optimization_parameters_utils.cc | 472 ++++++++++++++++++
 ..._embedding_optimization_parameters_utils.h | 121 +++++
 .../xla/xla}/tpu/tpu_compile_interface.cc     |  10 +-
 .../xla/xla/tpu/tpu_compile_interface.h       |  47 ++
 13 files changed, 744 insertions(+), 616 deletions(-)
 rename {tensorflow/core/tpu => third_party/xla/xla/tpu/protobuf}/tpu_embedding_configuration_utils.cc (84%)
 create mode 100644 third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration_utils.h
 create mode 100644 third_party/xla/xla/tpu/protobuf/tpu_embedding_optimization_parameters_utils.cc
 create mode 100644 third_party/xla/xla/tpu/protobuf/tpu_embedding_optimization_parameters_utils.h
 rename {tensorflow/core => third_party/xla/xla}/tpu/tpu_compile_interface.cc (81%)
 create mode 100644 third_party/xla/xla/tpu/tpu_compile_interface.h

diff --git a/tensorflow/core/tpu/BUILD b/tensorflow/core/tpu/BUILD
index e3d661185ed..7caf30cd621 100644
--- a/tensorflow/core/tpu/BUILD
+++ b/tensorflow/core/tpu/BUILD
@@ -40,16 +40,10 @@ package(
 
 cc_library(
     name = "tpu_embedding_configuration_utils",
-    srcs = ["tpu_embedding_configuration_utils.cc"],
     hdrs = ["tpu_embedding_configuration_utils.h"],
     visibility = ["//visibility:public"],
     deps = [
-        ":tpu_embedding_optimization_parameters_utils",
-        "//tensorflow/core/protobuf/tpu:optimization_parameters_proto_cc",
-        "//tensorflow/core/protobuf/tpu:tpu_embedding_configuration_proto_cc",
-        "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings:str_format",
+        "@xla//xla/tpu/protobuf:tpu_embedding_configuration_utils",
     ],
 )
 
@@ -90,16 +84,9 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         "//tensorflow/core:framework",
-        "//tensorflow/core:lib",
-        "//tensorflow/core:lib_proto_parsing",
-        "//tensorflow/core/protobuf/tpu:optimization_parameters_proto_cc",
-        "@com_google_absl//absl/base",
-        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/status",
-        "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/strings:str_format",
-        "@xla//xla:xla_data_proto_cc",
-        "@xla//xla/service:hlo_proto_cc",
+        "@xla//xla/tpu/protobuf:tpu_embedding_optimization_parameters_utils",
+        "@xla//xla/tsl/platform:errors",
     ],
 )
 
@@ -185,12 +172,9 @@ cc_library(
 
 cc_library(
     name = "tpu_compile_interface",
-    srcs = ["tpu_compile_interface.cc"],
     hdrs = ["tpu_compile_interface.h"],
     deps = [
-        "//tensorflow/core:lib",
-        "@com_google_absl//absl/log",
-        "@com_google_absl//absl/strings",
+        "@xla//xla/tpu:tpu_compile_interface",
     ],
 )
 
diff --git a/tensorflow/core/tpu/tpu_compile_interface.h b/tensorflow/core/tpu/tpu_compile_interface.h
index a97e721b016..85979f83add 100644
--- a/tensorflow/core/tpu/tpu_compile_interface.h
+++ b/tensorflow/core/tpu/tpu_compile_interface.h
@@ -15,33 +15,6 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_TPU_TPU_COMPILE_INTERFACE_H_
 #define TENSORFLOW_CORE_TPU_TPU_COMPILE_INTERFACE_H_
 
-#include <cstdint>
-
-#include "absl/strings/string_view.h"
-
-// Some legacy code requires different implementations for operations like
-// fingerprint/hashing during compilation and/or graph rewriting. These
-// alternate implementations can be registered (via a module initializer) to
-// change the default behavior.
-class TpuCompileInterface {
- public:
-  virtual ~TpuCompileInterface() {}
-  static TpuCompileInterface* Get();
-  static bool RegisterImplementation(TpuCompileInterface* impl);
-
-  virtual uint64_t FingerprintString(absl::string_view str) = 0;
-
-  // Proto: tensorflow::tpu::CompilationResultProto
-  // Location: tensorflow/core/protobuf/tpu/compilation_result.proto
-  static inline constexpr char kTpuCompileErrorPayloadKey[] =
-      "type.googleapis.com/tensorflow.tpu.CompilationResultProto";
-
-  // Unique string added to the error message for permanent errors during
-  // XLA:TPU compilation. This can be used by TensorFlow models to distinguish
-  // compilation errors from transient errors created by TPU worker preemptions
-  // and restarts.
-  static inline constexpr char kTpuCompileErrorMessage[] =
-      "XLA:TPU compile permanent error";
-};
+#include "xla/tpu/tpu_compile_interface.h"  // IWYU pragma: export
 
 #endif  // TENSORFLOW_CORE_TPU_TPU_COMPILE_INTERFACE_H_
diff --git a/tensorflow/core/tpu/tpu_embedding_configuration_utils.h b/tensorflow/core/tpu/tpu_embedding_configuration_utils.h
index 3ac55d17ff6..70abe3087d3 100644
--- a/tensorflow/core/tpu/tpu_embedding_configuration_utils.h
+++ b/tensorflow/core/tpu/tpu_embedding_configuration_utils.h
@@ -16,22 +16,6 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_TPU_TPU_EMBEDDING_CONFIGURATION_UTILS_H_
 #define TENSORFLOW_CORE_TPU_TPU_EMBEDDING_CONFIGURATION_UTILS_H_
 
-#include <cstdint>
-
-#include "absl/status/statusor.h"
-#include "tensorflow/core/protobuf/tpu/tpu_embedding_configuration.pb.h"
-
-namespace tensorflow {
-namespace tpu {
-
-// Returns the total number of unique dynamic input tags used in optimizers. If
-// the tag specific is erroneous, returns an invalid argument error. For correct
-// tag specification, see the comment next to the OptimizerDynamicInput proto in
-// //third_party/tensorflow/core/protobuf/tpu/optimization_parameters.proto.
-absl::StatusOr<int32_t> ComputeTotalTagCountForOptimizerDynamicInputs(
-    const tensorflow::tpu::TPUEmbeddingConfiguration& tpu_embedding_config);
-
-}  // namespace tpu
-}  // namespace tensorflow
+#include "xla/tpu/protobuf/tpu_embedding_configuration_utils.h"  // IWYU pragma: export
 
 #endif  // TENSORFLOW_CORE_TPU_TPU_EMBEDDING_CONFIGURATION_UTILS_H_
diff --git a/tensorflow/core/tpu/tpu_embedding_optimization_parameters_utils.cc b/tensorflow/core/tpu/tpu_embedding_optimization_parameters_utils.cc
index 8a60fb3a259..0343fcab0c4 100644
--- a/tensorflow/core/tpu/tpu_embedding_optimization_parameters_utils.cc
+++ b/tensorflow/core/tpu/tpu_embedding_optimization_parameters_utils.cc
@@ -16,462 +16,15 @@ limitations under the License.
 #include "tensorflow/core/tpu/tpu_embedding_optimization_parameters_utils.h"
 
 #include <string>
-#include <utility>
 #include <vector>
 
-#include "absl/container/flat_hash_set.h"
 #include "absl/status/status.h"
-#include "absl/strings/str_cat.h"
-#include "absl/strings/str_format.h"
-#include "xla/service/hlo.pb.h"
 #include "xla/tsl/platform/errors.h"
-#include "xla/xla_data.pb.h"
 #include "tensorflow/core/framework/shape_inference.h"
-#include "tensorflow/core/platform/errors.h"
-#include "tensorflow/core/platform/stringprintf.h"
-#include "tensorflow/core/platform/types.h"
-#include "tensorflow/core/protobuf/tpu/optimization_parameters.pb.h"
 
 namespace tensorflow {
 namespace tpu {
 
-std::string GetOptimizationAlgorithmName(OptimizationAlgorithm alg) {
-  switch (alg) {
-    case OptimizationAlgorithm::kAdagrad:
-      return "Adagrad";
-    case OptimizationAlgorithm::kAdagradMomentum:
-      return "AdagradMomentum";
-    case OptimizationAlgorithm::kBoundedAdagrad:
-      return "BoundedAdagrad";
-    case OptimizationAlgorithm::kFrequencyAwareAdagrad:
-      return "FrequencyAwareAdagrad";
-    case OptimizationAlgorithm::kStochasticGradientDescent:
-      return "StochasticGradientDescent";
-    case OptimizationAlgorithm::kFtrl:
-      return "FTRL";
-    case OptimizationAlgorithm::kAdam:
-      return "ADAM";
-    case OptimizationAlgorithm::kMomentum:
-      return "Momentum";
-    case OptimizationAlgorithm::kLion:
-      return "Lion";
-    case OptimizationAlgorithm::kRmsProp:
-      return "RMSProp";
-    case OptimizationAlgorithm::kCenteredRmsProp:
-      return "CenteredRMSProp";
-    case OptimizationAlgorithm::kMdlAdagradLight:
-      return "MDLAdagradLight";
-    case OptimizationAlgorithm::kAdadelta:
-      return "Adadelta";
-    case OptimizationAlgorithm::kProximalAdagrad:
-      return "ProximalAdagrad";
-    case OptimizationAlgorithm::kOnlineYogi:
-      return "OnlineYogi";
-    case OptimizationAlgorithm::kProximalYogi:
-      return "ProximalYogi";
-    case OptimizationAlgorithm::kFrequencyEstimator:
-      return "FrequencyEstimator";
-    case OptimizationAlgorithm::kUserDefinedProgram:
-      return "UserDefinedProgram";
-    case OptimizationAlgorithm::kAssign:
-      return "Assign";
-    case OptimizationAlgorithm::PARAMETERS_NOT_SET:
-      return "*** Not set ***";
-  }
-  return "*** Not set ***";
-}
-
-std::string GetOptimizationAlgorithmFriendlyName(OptimizationAlgorithm alg) {
-  switch (alg) {
-    case OptimizationAlgorithm::kAdagrad:
-      return "Adagrad";
-    case OptimizationAlgorithm::kAdagradMomentum:
-      return "Adagrad with Momentum";
-    case OptimizationAlgorithm::kBoundedAdagrad:
-      return "Bounded Adagrad";
-    case OptimizationAlgorithm::kFrequencyAwareAdagrad:
-      return "Frequency Aware Adagrad";
-    case OptimizationAlgorithm::kStochasticGradientDescent:
-      return "stochastic gradient descent";
-    case OptimizationAlgorithm::kFtrl:
-      return "FTRL";
-    case OptimizationAlgorithm::kAdam:
-      return "ADAM";
-    case OptimizationAlgorithm::kMomentum:
-      return "Momentum";
-    case OptimizationAlgorithm::kLion:
-      return "Lion";
-    case OptimizationAlgorithm::kRmsProp:
-      return "RMSProp";
-    case OptimizationAlgorithm::kCenteredRmsProp:
-      return "centered RMSProp";
-    case OptimizationAlgorithm::kMdlAdagradLight:
-      return "MDL Adagrad Light";
-    case OptimizationAlgorithm::kAdadelta:
-      return "Adadelta";
-    case OptimizationAlgorithm::kProximalAdagrad:
-      return "proximal Adagrad";
-    case OptimizationAlgorithm::kOnlineYogi:
-      return "online Yogi";
-    case OptimizationAlgorithm::kProximalYogi:
-      return "proximal Yogi";
-    case OptimizationAlgorithm::kFrequencyEstimator:
-      return "frequency estimator";
-    case OptimizationAlgorithm::kUserDefinedProgram:
-      return "UserDefinedProgram";
-    case OptimizationAlgorithm::kAssign:
-      return "Assign";
-    case OptimizationAlgorithm::PARAMETERS_NOT_SET:
-      return "unknown (not specified)";
-  }
-  return "unknown (not specified)";
-}
-
-// Returns the number of optimization parameter vectors used by the optimization
-// algorithm, excluding the weights themselves and assuming no gradient
-// accumulation.
-absl::Status GetBaseAuxiliaryParameterCount(
-    const OptimizationParameters& params, int* count) {
-  switch (params.parameters_case()) {
-    case OptimizationAlgorithm::kAdagrad:
-      *count = 1;
-      return absl::OkStatus();
-    case OptimizationAlgorithm::kAdagradMomentum:
-      *count = 2;
-      return absl::OkStatus();
-    case OptimizationAlgorithm::kBoundedAdagrad:
-      *count = 1;
-      return absl::OkStatus();
-    case OptimizationAlgorithm::kFrequencyAwareAdagrad:
-      *count = 2;
-      return absl::OkStatus();
-    case OptimizationAlgorithm::kStochasticGradientDescent:
-      *count = 0;
-      return absl::OkStatus();
-    case OptimizationAlgorithm::kFtrl:
-      *count = 2;
-      return absl::OkStatus();
-    case OptimizationAlgorithm::kAdam:
-      *count = 2;
-      return absl::OkStatus();
-    case OptimizationAlgorithm::kMomentum:
-      *count = 1;
-      return absl::OkStatus();
-    case OptimizationAlgorithm::kLion:
-      *count = 1;
-      return absl::OkStatus();
-    case OptimizationAlgorithm::kRmsProp:
-      *count = 2;
-      return absl::OkStatus();
-    case OptimizationAlgorithm::kCenteredRmsProp:
-      *count = 3;
-      return absl::OkStatus();
-    case OptimizationAlgorithm::kMdlAdagradLight:
-      *count = 3;
-      return absl::OkStatus();
-    case OptimizationAlgorithm::kAdadelta:
-      *count = 2;
-      return absl::OkStatus();
-    case OptimizationAlgorithm::kProximalAdagrad:
-      *count = 1;
-      return absl::OkStatus();
-    case OptimizationAlgorithm::kOnlineYogi:
-      *count = 2;
-      return absl::OkStatus();
-    case OptimizationAlgorithm::kProximalYogi:
-      *count = 2;
-      return absl::OkStatus();
-    case OptimizationAlgorithm::kFrequencyEstimator:
-      *count = 1;
-      return absl::OkStatus();
-    case OptimizationAlgorithm::kUserDefinedProgram: {
-      const xla::ProgramShapeProto& program_shape =
-          params.user_defined_program().program().host_program_shape();
-
-      const int num_inputs = program_shape.parameters_size();
-      const int num_outputs = program_shape.result().tuple_shapes_size();
-
-      if ((num_inputs < 2) || ((num_inputs != num_outputs + 1) &&
-                               (num_inputs != num_outputs + 2))) {
-        return absl::InvalidArgumentError(absl::StrCat(
-            "User-defined TPU embedding optimizer program must have at least "
-            "two inputs and the number of outputs must be 1 or 2 less than the "
-            "number of inputs. Received ",
-            num_inputs, " input(s) and ", num_outputs, "output(s)."));
-      }
-
-      *count = num_outputs - 1;
-
-      return absl::OkStatus();
-    }
-    case OptimizationAlgorithm::kAssign:
-      *count = 0;
-      return absl::OkStatus();
-    case OptimizationAlgorithm::PARAMETERS_NOT_SET:
-      return absl::InvalidArgumentError("No optimization algorithm specified");
-  }
-  return absl::InvalidArgumentError("No optimization algorithm specified");
-}
-
-absl::Status GetGradientAccumulationSupport(
-    const OptimizationParameters& params,
-    GradientAccumulationSupport* support) {
-  int auxiliary_parameter_count;
-  TF_RETURN_IF_ERROR(
-      GetBaseAuxiliaryParameterCount(params, &auxiliary_parameter_count));
-  *support = auxiliary_parameter_count + 1 <= kMaxAuxiliaryParameterCount
-                 ? GradientAccumulationSupport::kSupported
-                 : GradientAccumulationSupport::kNotSupported;
-  return absl::OkStatus();
-}
-
-absl::Status UseGradientAccumulation(const OptimizationParameters& params,
-                                     bool* use_gradient_accumulation) {
-  GradientAccumulationSupport support;
-  TF_RETURN_IF_ERROR(GetGradientAccumulationSupport(params, &support));
-  bool raw_gradient_accumulation_status = false;
-  switch (params.gradient_accumulation_status()) {
-    case GradientAccumulationStatus::UNSPECIFIED: {
-      // Default is now to turn gradient accumulation on by default.
-      raw_gradient_accumulation_status = true;
-      break;
-    }
-    case GradientAccumulationStatus::DISABLED: {
-      raw_gradient_accumulation_status = false;
-      break;
-    }
-    case GradientAccumulationStatus::ENABLED: {
-      raw_gradient_accumulation_status = true;
-      break;
-    }
-    default:
-      return absl::InternalError(
-          absl::StrCat("Unsupported gradient accumulation status ",
-                       GradientAccumulationStatus_Status_Name(
-                           params.gradient_accumulation_status())));
-  }
-  switch (support) {
-    case GradientAccumulationSupport::kSupported: {
-      *use_gradient_accumulation = raw_gradient_accumulation_status;
-      break;
-    }
-    case GradientAccumulationSupport::kNotSupported: {
-      if (raw_gradient_accumulation_status) {
-        return absl::InvalidArgumentError(absl::StrFormat(
-            "Optimization algorithm %s does not support gradient accumulation "
-            "but parameters specify it.",
-            GetOptimizationAlgorithmName(params.parameters_case()).c_str()));
-      }
-      *use_gradient_accumulation = false;
-      break;
-    }
-  }
-  return absl::OkStatus();
-}
-
-absl::Status GetOptimizationAlgorithmStateVariables(
-    const OptimizationParameters& params,
-    std::vector<StateVariableSpecification>* state_variables) {
-  // The parameter set for the weights themselves is required to be named
-  // "parameters". The rest should stay stable for compatibility. There is an
-  // internal function, GetOptimizationAlgorithmStateVariableInternalIndices,
-  // that needs to be updated along with this one.
-  bool use_gradient_accumulation;
-  TF_RETURN_IF_ERROR(
-      UseGradientAccumulation(params, &use_gradient_accumulation));
-
-  auto add_state_variable = [&](const std::string& name) {
-    StateVariableSpecification spec;
-    spec.set_name(name);
-    (void)spec.mutable_user_defined();
-    state_variables->push_back(spec);
-  };
-
-  switch (params.parameters_case()) {
-    case OptimizationAlgorithm::kAdagrad: {
-      add_state_variable("parameters");
-      add_state_variable("accumulators");
-      break;
-    }
-    case OptimizationAlgorithm::kAdagradMomentum: {
-      add_state_variable("parameters");
-      add_state_variable("accumulators");
-      add_state_variable("momenta");
-      break;
-    }
-    case OptimizationAlgorithm::kBoundedAdagrad: {
-      add_state_variable("parameters");
-      add_state_variable("accumulators");
-      break;
-    }
-    case OptimizationAlgorithm::kFrequencyAwareAdagrad: {
-      add_state_variable("parameters");
-      add_state_variable("accumulators");
-      add_state_variable("counters");
-      break;
-    }
-    case OptimizationAlgorithm::kStochasticGradientDescent: {
-      add_state_variable("parameters");
-      break;
-    }
-    case OptimizationAlgorithm::kFtrl: {
-      add_state_variable("parameters");
-      add_state_variable("accumulators");
-      add_state_variable("linears");
-      break;
-    }
-    case OptimizationAlgorithm::kAdam: {
-      add_state_variable("parameters");
-      add_state_variable("momenta");
-      add_state_variable("velocities");
-      break;
-    }
-    case OptimizationAlgorithm::kMomentum: {
-      add_state_variable("parameters");
-      add_state_variable("momenta");
-      break;
-    }
-    case OptimizationAlgorithm::kLion: {
-      add_state_variable("parameters");
-      add_state_variable("momenta");
-      break;
-    }
-    case OptimizationAlgorithm::kRmsProp: {
-      add_state_variable("parameters");
-      add_state_variable("ms");
-      add_state_variable("mom");
-      break;
-    }
-    case OptimizationAlgorithm::kCenteredRmsProp: {
-      add_state_variable("parameters");
-      add_state_variable("ms");
-      add_state_variable("mom");
-      add_state_variable("mg");
-      break;
-    }
-    case OptimizationAlgorithm::kMdlAdagradLight: {
-      add_state_variable("parameters");
-      add_state_variable("accumulators");
-      add_state_variable("weights");
-      add_state_variable("benefits");
-      break;
-    }
-    case OptimizationAlgorithm::kAdadelta: {
-      add_state_variable("parameters");
-      add_state_variable("accumulators");
-      add_state_variable("updates");
-      break;
-    }
-    case OptimizationAlgorithm::kProximalAdagrad: {
-      add_state_variable("parameters");
-      add_state_variable("accumulators");
-      break;
-    }
-    case OptimizationAlgorithm::kOnlineYogi: {
-      add_state_variable("parameters");
-      add_state_variable("vs");
-      add_state_variable("linears");
-      break;
-    }
-    case OptimizationAlgorithm::kProximalYogi: {
-      add_state_variable("parameters");
-      add_state_variable("v");
-      add_state_variable("m");
-      break;
-    }
-    case OptimizationAlgorithm::kFrequencyEstimator: {
-      add_state_variable("parameters");
-      add_state_variable("last_hit_step");
-      break;
-    }
-    case OptimizationAlgorithm::kUserDefinedProgram: {
-      add_state_variable("parameters");
-      int num_slots = -1;
-      TF_RETURN_IF_ERROR(GetBaseAuxiliaryParameterCount(params, &num_slots));
-      for (int i = 0; i < num_slots; ++i) {
-        add_state_variable(absl::StrCat("Slot_", i));
-      }
-      break;
-    }
-    case OptimizationAlgorithm::kAssign: {
-      add_state_variable("parameters");
-      break;
-    }
-    case OptimizationAlgorithm::PARAMETERS_NOT_SET: {
-      return absl::InvalidArgumentError("No optimization algorithm specified");
-    }
-  }
-
-  // This needs to be last for compatibility.
-  if (use_gradient_accumulation) {
-    StateVariableSpecification gradient_acc;
-    gradient_acc.set_name("gradient_accumulators");
-    gradient_acc.mutable_fill_with_constant()->set_initial_value(
-        GradientAccumulatorInitialValue());
-    state_variables->push_back(std::move(gradient_acc));
-  }
-
-  if (state_variables->size() > kMaxAuxiliaryParameterCount + 1) {
-    return absl::InvalidArgumentError(
-        absl::StrCat("Optimization algorithm",
-                     GetOptimizationAlgorithmName(params.parameters_case()),
-                     "does not support gradient accumulation because it "
-                     "already has too many other accumulators"));
-  }
-  return absl::OkStatus();
-}
-
-absl::flat_hash_set<int> GetOptimizerDynamicInputTags(
-    const OptimizationParameters& params) {
-  absl::flat_hash_set<int> tags;
-  if (params.learning_rate().has_dynamic()) {
-    tags.insert(params.learning_rate().dynamic().tag());
-  }
-  tags.merge(GetOptimizerHyperParameterTags(params));
-  return tags;
-}
-
-absl::flat_hash_set<int> GetOptimizerHyperParameterTags(
-    const OptimizationParameters& params) {
-  absl::flat_hash_set<int> tags;
-  switch (params.parameters_case()) {
-    case OptimizationAlgorithm::kFrequencyAwareAdagrad:
-      tags.insert(params.frequency_aware_adagrad().step_counter().tag());
-      break;
-    default:
-      break;
-  }
-  return tags;
-}
-
-bool UsesDynamicInputsInOptimizer(const OptimizationParameters& params) {
-  return !GetOptimizerDynamicInputTags(params).empty();
-}
-
-std::vector<OptimizationAlgorithm> GetOptimizationAlgorithms() {
-  return {
-      OptimizationAlgorithm::kAdagrad,
-      OptimizationAlgorithm::kAdagradMomentum,
-      OptimizationAlgorithm::kBoundedAdagrad,
-      OptimizationAlgorithm::kFrequencyAwareAdagrad,
-      OptimizationAlgorithm::kStochasticGradientDescent,
-      OptimizationAlgorithm::kFtrl,
-      OptimizationAlgorithm::kAdam,
-      OptimizationAlgorithm::kMomentum,
-      OptimizationAlgorithm::kLion,
-      OptimizationAlgorithm::kRmsProp,
-      OptimizationAlgorithm::kCenteredRmsProp,
-      OptimizationAlgorithm::kMdlAdagradLight,
-      OptimizationAlgorithm::kAdadelta,
-      OptimizationAlgorithm::kProximalAdagrad,
-      OptimizationAlgorithm::kOnlineYogi,
-      OptimizationAlgorithm::kProximalYogi,
-      OptimizationAlgorithm::kFrequencyEstimator,
-      OptimizationAlgorithm::kUserDefinedProgram,
-      OptimizationAlgorithm::kAssign,
-  };
-}
-
 absl::Status LoadOpShapeFunction::operator()(
     shape_inference::InferenceContext* c) const {
   int table_id;
diff --git a/tensorflow/core/tpu/tpu_embedding_optimization_parameters_utils.h b/tensorflow/core/tpu/tpu_embedding_optimization_parameters_utils.h
index 43643fbdb90..00c64d93b69 100644
--- a/tensorflow/core/tpu/tpu_embedding_optimization_parameters_utils.h
+++ b/tensorflow/core/tpu/tpu_embedding_optimization_parameters_utils.h
@@ -16,106 +16,13 @@ limitations under the License.
 #ifndef TENSORFLOW_CORE_TPU_TPU_EMBEDDING_OPTIMIZATION_PARAMETERS_UTILS_H_
 #define TENSORFLOW_CORE_TPU_TPU_EMBEDDING_OPTIMIZATION_PARAMETERS_UTILS_H_
 
-#include <cstdint>
-#include <string>
-#include <vector>
-
-#include "absl/base/casts.h"
-#include "absl/container/flat_hash_set.h"
 #include "absl/status/status.h"
+#include "xla/tpu/protobuf/tpu_embedding_optimization_parameters_utils.h"  // IWYU pragma: export
 #include "tensorflow/core/framework/op_def_builder.h"
-#include "tensorflow/core/protobuf/tpu/optimization_parameters.pb.h"
 
 namespace tensorflow {
 namespace tpu {
 
-using OptimizationAlgorithm = OptimizationParameters::ParametersCase;
-
-// Returns the name of the optimization algorithm.
-std::string GetOptimizationAlgorithmName(OptimizationAlgorithm alg);
-
-// Returns a user-friendly name for the optimization algorithm.
-std::string GetOptimizationAlgorithmFriendlyName(OptimizationAlgorithm alg);
-
-// Returns all supported optimization algorithms.
-std::vector<OptimizationAlgorithm> GetOptimizationAlgorithms();
-
-enum class GradientAccumulationSupport {
-  // Accumulation cannot be used with this optimizer.
-  kNotSupported,
-
-  // Accumulation is allowed and changes optimizer behavior.
-  kSupported,
-};
-
-// Returns the number of optimization parameter vectors used by the optimization
-// algorithm, excluding the weights themselves and assuming no gradient
-// accumulation.
-absl::Status GetBaseAuxiliaryParameterCount(
-    const OptimizationParameters &params, int *count);
-
-// Returns whether (and how) an optimization algorithm supports gradient
-// accumulation.
-absl::Status GetGradientAccumulationSupport(
-    const OptimizationParameters &params, GradientAccumulationSupport *support);
-
-// Returns whether both the given set of optimization parameters has gradient
-// accumulation turned on and that the algorithm used supports it or should
-// ignore that setting. Returns an error if gradient accumulation is enabled and
-// the algorithm does not support it.
-absl::Status UseGradientAccumulation(const OptimizationParameters &params,
-                                     bool *use_gradient_accumulation);
-
-// Returns the parameter specifications for the optimization algorithm (the main
-// parameters first, followed by any auxiliary parameters such as Adagrad
-// accumulators).
-absl::Status GetOptimizationAlgorithmStateVariables(
-    const OptimizationParameters &params,
-    std::vector<StateVariableSpecification> *state_variables);
-
-// Returns the set of dynamic input tags used by the optimization algorithm.
-// This includes both dynamic learning rates and other hyperparameters (e.g.,
-// step counters for the frequency aware Adagrad optimizer).
-absl::flat_hash_set<int> GetOptimizerDynamicInputTags(
-    const OptimizationParameters &params);
-
-// Returns the set of dynamic hyperparameter tags used by the optimization
-// algorithm. This includes other hyperparameters used by the optimization
-// algorithm (e.g., step counters for the frequency aware Adagrad optimizer). It
-// excludes the dynamic learning rate tag.
-absl::flat_hash_set<int> GetOptimizerHyperParameterTags(
-    const OptimizationParameters &params);
-
-// Returns true if the optimization algorithm uses dynamic inputs in its
-// computation.
-bool UsesDynamicInputsInOptimizer(const OptimizationParameters &params);
-
-// Maximum value of auxiliary_parametery_count for any optimization algorithm.
-// This count is used by TPU embedding load/retrieve and needs to be independent
-// of any particular TPU version and hence, we take the maximum across all TPU
-// versions.
-static constexpr int kMaxAuxiliaryParameterCount = 7;
-
-// Fill value for gradient accumulators. This is a denormal so that it will be
-// flushed to zero on the current TPU platforms and needs to continue to have
-// the following properties in the future:
-//
-// 1. Does not have the same bit pattern as a zero and can be distinguished from
-// it using integer operations.
-// 2. Treated as zero by floating-point arithmetic operations (at least addition
-// and subtraction).
-// 3. Cannot be produced by any floating-point arithmetic operation, including
-// those involving itself.
-//
-// It does not need to compare equal or not equal to zero in floating point. We
-// need to use a non-zero value here because some optimization algorithms are
-// not no-ops on zero gradients, so we need to distinguish an accumulated
-// gradient of zero from one that has been cleared after its gradients have
-// already been applied to the parameters and accumulators.
-inline float GradientAccumulatorInitialValue() {
-  return absl::bit_cast<float, uint32_t>(1);
-}
-
 // Generic shape function for per-optimization-algorithm load ops.
 class LoadOpShapeFunction {
  public:
diff --git a/third_party/xla/xla/tpu/BUILD b/third_party/xla/xla/tpu/BUILD
index 9fce2731bcb..931de025e3e 100644
--- a/third_party/xla/xla/tpu/BUILD
+++ b/third_party/xla/xla/tpu/BUILD
@@ -283,3 +283,16 @@ cc_library(
         "@tsl//tsl/platform:types",
     ],
 )
+
+cc_library(
+    name = "tpu_compile_interface",
+    srcs = ["tpu_compile_interface.cc"],
+    hdrs = ["tpu_compile_interface.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/strings",
+        "@tsl//tsl/platform:fingerprint",
+        "@tsl//tsl/platform:logging",
+    ],
+)
diff --git a/third_party/xla/xla/tpu/protobuf/BUILD b/third_party/xla/xla/tpu/protobuf/BUILD
index f732b3e19e7..1a62300e115 100644
--- a/third_party/xla/xla/tpu/protobuf/BUILD
+++ b/third_party/xla/xla/tpu/protobuf/BUILD
@@ -18,6 +18,10 @@ load(
     "//xla/tsl/platform:build_config.bzl",
     "tf_proto_library",
 )
+load(
+    "//xla/tsl/platform:rules_cc.bzl",
+    "cc_library",
+)
 
 package(
     # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
@@ -45,6 +49,39 @@ tf_proto_library(
     deps = ["@com_google_protobuf//:wrappers_proto"],
 )
 
+cc_library(
+    name = "tpu_embedding_optimization_parameters_utils",
+    srcs = ["tpu_embedding_optimization_parameters_utils.cc"],
+    hdrs = ["tpu_embedding_optimization_parameters_utils.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":optimization_parameters_proto_cc",
+        "//xla:xla_data_proto_cc",
+        "//xla/service:hlo_proto_cc",
+        "//xla/tsl/platform:errors",
+        "@com_google_absl//absl/base",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+    ],
+)
+
+cc_library(
+    name = "tpu_embedding_configuration_utils",
+    srcs = ["tpu_embedding_configuration_utils.cc"],
+    hdrs = ["tpu_embedding_configuration_utils.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":optimization_parameters_proto_cc",
+        ":tpu_embedding_configuration_proto_cc",
+        ":tpu_embedding_optimization_parameters_utils",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:str_format",
+    ],
+)
+
 # copybara:uncomment_begin(google-only)
 # py_proto_library(
 #     name = "tpu_embedding_configuration_py_pb2",
diff --git a/tensorflow/core/tpu/tpu_embedding_configuration_utils.cc b/third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration_utils.cc
similarity index 84%
rename from tensorflow/core/tpu/tpu_embedding_configuration_utils.cc
rename to third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration_utils.cc
index 3da9697b2e8..71a2cb8723f 100644
--- a/tensorflow/core/tpu/tpu_embedding_configuration_utils.cc
+++ b/third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration_utils.cc
@@ -1,4 +1,4 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2026 The OpenXLA Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -13,7 +13,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "tensorflow/core/tpu/tpu_embedding_configuration_utils.h"
+#include "xla/tpu/protobuf/tpu_embedding_configuration_utils.h"
 
 #include <cstdint>
 #include <set>
@@ -21,9 +21,9 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_format.h"
-#include "tensorflow/core/protobuf/tpu/optimization_parameters.pb.h"
-#include "tensorflow/core/protobuf/tpu/tpu_embedding_configuration.pb.h"
-#include "tensorflow/core/tpu/tpu_embedding_optimization_parameters_utils.h"
+#include "xla/tpu/protobuf/optimization_parameters.pb.h"
+#include "xla/tpu/protobuf/tpu_embedding_configuration.pb.h"
+#include "xla/tpu/protobuf/tpu_embedding_optimization_parameters_utils.h"
 
 namespace tensorflow {
 namespace tpu {
diff --git a/third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration_utils.h b/third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration_utils.h
new file mode 100644
index 00000000000..e6ce461fb99
--- /dev/null
+++ b/third_party/xla/xla/tpu/protobuf/tpu_embedding_configuration_utils.h
@@ -0,0 +1,37 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_TPU_PROTOBUF_TPU_EMBEDDING_CONFIGURATION_UTILS_H_
+#define XLA_TPU_PROTOBUF_TPU_EMBEDDING_CONFIGURATION_UTILS_H_
+
+#include <cstdint>
+
+#include "absl/status/statusor.h"
+#include "xla/tpu/protobuf/tpu_embedding_configuration.pb.h"
+
+namespace tensorflow {
+namespace tpu {
+
+// Returns the total number of unique dynamic input tags used in optimizers. If
+// the tag specific is erroneous, returns an invalid argument error. For correct
+// tag specification, see the comment next to the OptimizerDynamicInput proto in
+// //third_party/tensorflow/compiler/xla/tpu/protobuf/optimization_parameters.proto.
+absl::StatusOr<int32_t> ComputeTotalTagCountForOptimizerDynamicInputs(
+    const tensorflow::tpu::TPUEmbeddingConfiguration& tpu_embedding_config);
+
+}  // namespace tpu
+}  // namespace tensorflow
+
+#endif  // XLA_TPU_PROTOBUF_TPU_EMBEDDING_CONFIGURATION_UTILS_H_
diff --git a/third_party/xla/xla/tpu/protobuf/tpu_embedding_optimization_parameters_utils.cc b/third_party/xla/xla/tpu/protobuf/tpu_embedding_optimization_parameters_utils.cc
new file mode 100644
index 00000000000..c0d66153b90
--- /dev/null
+++ b/third_party/xla/xla/tpu/protobuf/tpu_embedding_optimization_parameters_utils.cc
@@ -0,0 +1,472 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/tpu/protobuf/tpu_embedding_optimization_parameters_utils.h"
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
+#include "xla/service/hlo.pb.h"
+#include "xla/tpu/protobuf/optimization_parameters.pb.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/xla_data.pb.h"
+
+namespace tensorflow {
+namespace tpu {
+
+std::string GetOptimizationAlgorithmName(OptimizationAlgorithm alg) {
+  switch (alg) {
+    case OptimizationAlgorithm::kAdagrad:
+      return "Adagrad";
+    case OptimizationAlgorithm::kAdagradMomentum:
+      return "AdagradMomentum";
+    case OptimizationAlgorithm::kBoundedAdagrad:
+      return "BoundedAdagrad";
+    case OptimizationAlgorithm::kFrequencyAwareAdagrad:
+      return "FrequencyAwareAdagrad";
+    case OptimizationAlgorithm::kStochasticGradientDescent:
+      return "StochasticGradientDescent";
+    case OptimizationAlgorithm::kFtrl:
+      return "FTRL";
+    case OptimizationAlgorithm::kAdam:
+      return "ADAM";
+    case OptimizationAlgorithm::kMomentum:
+      return "Momentum";
+    case OptimizationAlgorithm::kLion:
+      return "Lion";
+    case OptimizationAlgorithm::kRmsProp:
+      return "RMSProp";
+    case OptimizationAlgorithm::kCenteredRmsProp:
+      return "CenteredRMSProp";
+    case OptimizationAlgorithm::kMdlAdagradLight:
+      return "MDLAdagradLight";
+    case OptimizationAlgorithm::kAdadelta:
+      return "Adadelta";
+    case OptimizationAlgorithm::kProximalAdagrad:
+      return "ProximalAdagrad";
+    case OptimizationAlgorithm::kOnlineYogi:
+      return "OnlineYogi";
+    case OptimizationAlgorithm::kProximalYogi:
+      return "ProximalYogi";
+    case OptimizationAlgorithm::kFrequencyEstimator:
+      return "FrequencyEstimator";
+    case OptimizationAlgorithm::kUserDefinedProgram:
+      return "UserDefinedProgram";
+    case OptimizationAlgorithm::kAssign:
+      return "Assign";
+    case OptimizationAlgorithm::PARAMETERS_NOT_SET:
+      return "*** Not set ***";
+  }
+  return "*** Not set ***";
+}
+
+std::string GetOptimizationAlgorithmFriendlyName(OptimizationAlgorithm alg) {
+  switch (alg) {
+    case OptimizationAlgorithm::kAdagrad:
+      return "Adagrad";
+    case OptimizationAlgorithm::kAdagradMomentum:
+      return "Adagrad with Momentum";
+    case OptimizationAlgorithm::kBoundedAdagrad:
+      return "Bounded Adagrad";
+    case OptimizationAlgorithm::kFrequencyAwareAdagrad:
+      return "Frequency Aware Adagrad";
+    case OptimizationAlgorithm::kStochasticGradientDescent:
+      return "stochastic gradient descent";
+    case OptimizationAlgorithm::kFtrl:
+      return "FTRL";
+    case OptimizationAlgorithm::kAdam:
+      return "ADAM";
+    case OptimizationAlgorithm::kMomentum:
+      return "Momentum";
+    case OptimizationAlgorithm::kLion:
+      return "Lion";
+    case OptimizationAlgorithm::kRmsProp:
+      return "RMSProp";
+    case OptimizationAlgorithm::kCenteredRmsProp:
+      return "centered RMSProp";
+    case OptimizationAlgorithm::kMdlAdagradLight:
+      return "MDL Adagrad Light";
+    case OptimizationAlgorithm::kAdadelta:
+      return "Adadelta";
+    case OptimizationAlgorithm::kProximalAdagrad:
+      return "proximal Adagrad";
+    case OptimizationAlgorithm::kOnlineYogi:
+      return "online Yogi";
+    case OptimizationAlgorithm::kProximalYogi:
+      return "proximal Yogi";
+    case OptimizationAlgorithm::kFrequencyEstimator:
+      return "frequency estimator";
+    case OptimizationAlgorithm::kUserDefinedProgram:
+      return "UserDefinedProgram";
+    case OptimizationAlgorithm::kAssign:
+      return "Assign";
+    case OptimizationAlgorithm::PARAMETERS_NOT_SET:
+      return "unknown (not specified)";
+  }
+  return "unknown (not specified)";
+}
+
+// Returns the number of optimization parameter vectors used by the optimization
+// algorithm, excluding the weights themselves and assuming no gradient
+// accumulation.
+absl::Status GetBaseAuxiliaryParameterCount(
+    const OptimizationParameters& params, int* count) {
+  switch (params.parameters_case()) {
+    case OptimizationAlgorithm::kAdagrad:
+      *count = 1;
+      return absl::OkStatus();
+    case OptimizationAlgorithm::kAdagradMomentum:
+      *count = 2;
+      return absl::OkStatus();
+    case OptimizationAlgorithm::kBoundedAdagrad:
+      *count = 1;
+      return absl::OkStatus();
+    case OptimizationAlgorithm::kFrequencyAwareAdagrad:
+      *count = 2;
+      return absl::OkStatus();
+    case OptimizationAlgorithm::kStochasticGradientDescent:
+      *count = 0;
+      return absl::OkStatus();
+    case OptimizationAlgorithm::kFtrl:
+      *count = 2;
+      return absl::OkStatus();
+    case OptimizationAlgorithm::kAdam:
+      *count = 2;
+      return absl::OkStatus();
+    case OptimizationAlgorithm::kMomentum:
+      *count = 1;
+      return absl::OkStatus();
+    case OptimizationAlgorithm::kLion:
+      *count = 1;
+      return absl::OkStatus();
+    case OptimizationAlgorithm::kRmsProp:
+      *count = 2;
+      return absl::OkStatus();
+    case OptimizationAlgorithm::kCenteredRmsProp:
+      *count = 3;
+      return absl::OkStatus();
+    case OptimizationAlgorithm::kMdlAdagradLight:
+      *count = 3;
+      return absl::OkStatus();
+    case OptimizationAlgorithm::kAdadelta:
+      *count = 2;
+      return absl::OkStatus();
+    case OptimizationAlgorithm::kProximalAdagrad:
+      *count = 1;
+      return absl::OkStatus();
+    case OptimizationAlgorithm::kOnlineYogi:
+      *count = 2;
+      return absl::OkStatus();
+    case OptimizationAlgorithm::kProximalYogi:
+      *count = 2;
+      return absl::OkStatus();
+    case OptimizationAlgorithm::kFrequencyEstimator:
+      *count = 1;
+      return absl::OkStatus();
+    case OptimizationAlgorithm::kUserDefinedProgram: {
+      const xla::ProgramShapeProto& program_shape =
+          params.user_defined_program().program().host_program_shape();
+
+      const int num_inputs = program_shape.parameters_size();
+      const int num_outputs = program_shape.result().tuple_shapes_size();
+
+      if ((num_inputs < 2) || ((num_inputs != num_outputs + 1) &&
+                               (num_inputs != num_outputs + 2))) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "User-defined TPU embedding optimizer program must have at least "
+            "two inputs and the number of outputs must be 1 or 2 less than the "
+            "number of inputs. Received ",
+            num_inputs, " input(s) and ", num_outputs, "output(s)."));
+      }
+
+      *count = num_outputs - 1;
+
+      return absl::OkStatus();
+    }
+    case OptimizationAlgorithm::kAssign:
+      *count = 0;
+      return absl::OkStatus();
+    case OptimizationAlgorithm::PARAMETERS_NOT_SET:
+      return absl::InvalidArgumentError("No optimization algorithm specified");
+  }
+  return absl::InvalidArgumentError("No optimization algorithm specified");
+}
+
+absl::Status GetGradientAccumulationSupport(
+    const OptimizationParameters& params,
+    GradientAccumulationSupport* support) {
+  int auxiliary_parameter_count;
+  TF_RETURN_IF_ERROR(
+      GetBaseAuxiliaryParameterCount(params, &auxiliary_parameter_count));
+  *support = auxiliary_parameter_count + 1 <= kMaxAuxiliaryParameterCount
+                 ? GradientAccumulationSupport::kSupported
+                 : GradientAccumulationSupport::kNotSupported;
+  return absl::OkStatus();
+}
+
+absl::Status UseGradientAccumulation(const OptimizationParameters& params,
+                                     bool* use_gradient_accumulation) {
+  GradientAccumulationSupport support;
+  TF_RETURN_IF_ERROR(GetGradientAccumulationSupport(params, &support));
+  bool raw_gradient_accumulation_status = false;
+  switch (params.gradient_accumulation_status()) {
+    case GradientAccumulationStatus::UNSPECIFIED: {
+      // Default is now to turn gradient accumulation on by default.
+      raw_gradient_accumulation_status = true;
+      break;
+    }
+    case GradientAccumulationStatus::DISABLED: {
+      raw_gradient_accumulation_status = false;
+      break;
+    }
+    case GradientAccumulationStatus::ENABLED: {
+      raw_gradient_accumulation_status = true;
+      break;
+    }
+    default:
+      return absl::InternalError(
+          absl::StrCat("Unsupported gradient accumulation status ",
+                       GradientAccumulationStatus_Status_Name(
+                           params.gradient_accumulation_status())));
+  }
+  switch (support) {
+    case GradientAccumulationSupport::kSupported: {
+      *use_gradient_accumulation = raw_gradient_accumulation_status;
+      break;
+    }
+    case GradientAccumulationSupport::kNotSupported: {
+      if (raw_gradient_accumulation_status) {
+        return absl::InvalidArgumentError(absl::StrFormat(
+            "Optimization algorithm %s does not support gradient accumulation "
+            "but parameters specify it.",
+            GetOptimizationAlgorithmName(params.parameters_case()).c_str()));
+      }
+      *use_gradient_accumulation = false;
+      break;
+    }
+  }
+  return absl::OkStatus();
+}
+
+absl::Status GetOptimizationAlgorithmStateVariables(
+    const OptimizationParameters& params,
+    std::vector<StateVariableSpecification>* state_variables) {
+  // The parameter set for the weights themselves is required to be named
+  // "parameters". The rest should stay stable for compatibility. There is an
+  // internal function, GetOptimizationAlgorithmStateVariableInternalIndices,
+  // that needs to be updated along with this one.
+  bool use_gradient_accumulation;
+  TF_RETURN_IF_ERROR(
+      UseGradientAccumulation(params, &use_gradient_accumulation));
+
+  auto add_state_variable = [&](const std::string& name) {
+    StateVariableSpecification spec;
+    spec.set_name(name);
+    (void)spec.mutable_user_defined();
+    state_variables->push_back(spec);
+  };
+
+  switch (params.parameters_case()) {
+    case OptimizationAlgorithm::kAdagrad: {
+      add_state_variable("parameters");
+      add_state_variable("accumulators");
+      break;
+    }
+    case OptimizationAlgorithm::kAdagradMomentum: {
+      add_state_variable("parameters");
+      add_state_variable("accumulators");
+      add_state_variable("momenta");
+      break;
+    }
+    case OptimizationAlgorithm::kBoundedAdagrad: {
+      add_state_variable("parameters");
+      add_state_variable("accumulators");
+      break;
+    }
+    case OptimizationAlgorithm::kFrequencyAwareAdagrad: {
+      add_state_variable("parameters");
+      add_state_variable("accumulators");
+      add_state_variable("counters");
+      break;
+    }
+    case OptimizationAlgorithm::kStochasticGradientDescent: {
+      add_state_variable("parameters");
+      break;
+    }
+    case OptimizationAlgorithm::kFtrl: {
+      add_state_variable("parameters");
+      add_state_variable("accumulators");
+      add_state_variable("linears");
+      break;
+    }
+    case OptimizationAlgorithm::kAdam: {
+      add_state_variable("parameters");
+      add_state_variable("momenta");
+      add_state_variable("velocities");
+      break;
+    }
+    case OptimizationAlgorithm::kMomentum: {
+      add_state_variable("parameters");
+      add_state_variable("momenta");
+      break;
+    }
+    case OptimizationAlgorithm::kLion: {
+      add_state_variable("parameters");
+      add_state_variable("momenta");
+      break;
+    }
+    case OptimizationAlgorithm::kRmsProp: {
+      add_state_variable("parameters");
+      add_state_variable("ms");
+      add_state_variable("mom");
+      break;
+    }
+    case OptimizationAlgorithm::kCenteredRmsProp: {
+      add_state_variable("parameters");
+      add_state_variable("ms");
+      add_state_variable("mom");
+      add_state_variable("mg");
+      break;
+    }
+    case OptimizationAlgorithm::kMdlAdagradLight: {
+      add_state_variable("parameters");
+      add_state_variable("accumulators");
+      add_state_variable("weights");
+      add_state_variable("benefits");
+      break;
+    }
+    case OptimizationAlgorithm::kAdadelta: {
+      add_state_variable("parameters");
+      add_state_variable("accumulators");
+      add_state_variable("updates");
+      break;
+    }
+    case OptimizationAlgorithm::kProximalAdagrad: {
+      add_state_variable("parameters");
+      add_state_variable("accumulators");
+      break;
+    }
+    case OptimizationAlgorithm::kOnlineYogi: {
+      add_state_variable("parameters");
+      add_state_variable("vs");
+      add_state_variable("linears");
+      break;
+    }
+    case OptimizationAlgorithm::kProximalYogi: {
+      add_state_variable("parameters");
+      add_state_variable("v");
+      add_state_variable("m");
+      break;
+    }
+    case OptimizationAlgorithm::kFrequencyEstimator: {
+      add_state_variable("parameters");
+      add_state_variable("last_hit_step");
+      break;
+    }
+    case OptimizationAlgorithm::kUserDefinedProgram: {
+      add_state_variable("parameters");
+      int num_slots = -1;
+      TF_RETURN_IF_ERROR(GetBaseAuxiliaryParameterCount(params, &num_slots));
+      for (int i = 0; i < num_slots; ++i) {
+        add_state_variable(absl::StrCat("Slot_", i));
+      }
+      break;
+    }
+    case OptimizationAlgorithm::kAssign: {
+      add_state_variable("parameters");
+      break;
+    }
+    case OptimizationAlgorithm::PARAMETERS_NOT_SET: {
+      return absl::InvalidArgumentError("No optimization algorithm specified");
+    }
+  }
+
+  // This needs to be last for compatibility.
+  if (use_gradient_accumulation) {
+    StateVariableSpecification gradient_acc;
+    gradient_acc.set_name("gradient_accumulators");
+    gradient_acc.mutable_fill_with_constant()->set_initial_value(
+        GradientAccumulatorInitialValue());
+    state_variables->push_back(std::move(gradient_acc));
+  }
+
+  if (state_variables->size() > kMaxAuxiliaryParameterCount + 1) {
+    return absl::InvalidArgumentError(
+        absl::StrCat("Optimization algorithm",
+                     GetOptimizationAlgorithmName(params.parameters_case()),
+                     "does not support gradient accumulation because it "
+                     "already has too many other accumulators"));
+  }
+  return absl::OkStatus();
+}
+
+absl::flat_hash_set<int> GetOptimizerDynamicInputTags(
+    const OptimizationParameters& params) {
+  absl::flat_hash_set<int> tags;
+  if (params.learning_rate().has_dynamic()) {
+    tags.insert(params.learning_rate().dynamic().tag());
+  }
+  tags.merge(GetOptimizerHyperParameterTags(params));
+  return tags;
+}
+
+absl::flat_hash_set<int> GetOptimizerHyperParameterTags(
+    const OptimizationParameters& params) {
+  absl::flat_hash_set<int> tags;
+  switch (params.parameters_case()) {
+    case OptimizationAlgorithm::kFrequencyAwareAdagrad:
+      tags.insert(params.frequency_aware_adagrad().step_counter().tag());
+      break;
+    default:
+      break;
+  }
+  return tags;
+}
+
+bool UsesDynamicInputsInOptimizer(const OptimizationParameters& params) {
+  return !GetOptimizerDynamicInputTags(params).empty();
+}
+
+std::vector<OptimizationAlgorithm> GetOptimizationAlgorithms() {
+  return {
+      OptimizationAlgorithm::kAdagrad,
+      OptimizationAlgorithm::kAdagradMomentum,
+      OptimizationAlgorithm::kBoundedAdagrad,
+      OptimizationAlgorithm::kFrequencyAwareAdagrad,
+      OptimizationAlgorithm::kStochasticGradientDescent,
+      OptimizationAlgorithm::kFtrl,
+      OptimizationAlgorithm::kAdam,
+      OptimizationAlgorithm::kMomentum,
+      OptimizationAlgorithm::kLion,
+      OptimizationAlgorithm::kRmsProp,
+      OptimizationAlgorithm::kCenteredRmsProp,
+      OptimizationAlgorithm::kMdlAdagradLight,
+      OptimizationAlgorithm::kAdadelta,
+      OptimizationAlgorithm::kProximalAdagrad,
+      OptimizationAlgorithm::kOnlineYogi,
+      OptimizationAlgorithm::kProximalYogi,
+      OptimizationAlgorithm::kFrequencyEstimator,
+      OptimizationAlgorithm::kUserDefinedProgram,
+      OptimizationAlgorithm::kAssign,
+  };
+}
+
+}  // namespace tpu
+}  // namespace tensorflow
diff --git a/third_party/xla/xla/tpu/protobuf/tpu_embedding_optimization_parameters_utils.h b/third_party/xla/xla/tpu/protobuf/tpu_embedding_optimization_parameters_utils.h
new file mode 100644
index 00000000000..938b843b9ee
--- /dev/null
+++ b/third_party/xla/xla/tpu/protobuf/tpu_embedding_optimization_parameters_utils.h
@@ -0,0 +1,121 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_TPU_PROTOBUF_TPU_EMBEDDING_OPTIMIZATION_PARAMETERS_UTILS_H_
+#define XLA_TPU_PROTOBUF_TPU_EMBEDDING_OPTIMIZATION_PARAMETERS_UTILS_H_
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include "absl/base/casts.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
+#include "xla/tpu/protobuf/optimization_parameters.pb.h"
+
+namespace tensorflow {
+namespace tpu {
+
+using OptimizationAlgorithm = OptimizationParameters::ParametersCase;
+
+// Returns the name of the optimization algorithm.
+std::string GetOptimizationAlgorithmName(OptimizationAlgorithm alg);
+
+// Returns a user-friendly name for the optimization algorithm.
+std::string GetOptimizationAlgorithmFriendlyName(OptimizationAlgorithm alg);
+
+// Returns all supported optimization algorithms.
+std::vector<OptimizationAlgorithm> GetOptimizationAlgorithms();
+
+enum class GradientAccumulationSupport {
+  // Accumulation cannot be used with this optimizer.
+  kNotSupported,
+
+  // Accumulation is allowed and changes optimizer behavior.
+  kSupported,
+};
+
+// Returns the number of optimization parameter vectors used by the optimization
+// algorithm, excluding the weights themselves and assuming no gradient
+// accumulation.
+absl::Status GetBaseAuxiliaryParameterCount(
+    const OptimizationParameters& params, int* count);
+
+// Returns whether (and how) an optimization algorithm supports gradient
+// accumulation.
+absl::Status GetGradientAccumulationSupport(
+    const OptimizationParameters& params, GradientAccumulationSupport* support);
+
+// Returns whether both the given set of optimization parameters has gradient
+// accumulation turned on and that the algorithm used supports it or should
+// ignore that setting. Returns an error if gradient accumulation is enabled and
+// the algorithm does not support it.
+absl::Status UseGradientAccumulation(const OptimizationParameters& params,
+                                     bool* use_gradient_accumulation);
+
+// Returns the parameter specifications for the optimization algorithm (the main
+// parameters first, followed by any auxiliary parameters such as Adagrad
+// accumulators).
+absl::Status GetOptimizationAlgorithmStateVariables(
+    const OptimizationParameters& params,
+    std::vector<StateVariableSpecification>* state_variables);
+
+// Returns the set of dynamic input tags used by the optimization algorithm.
+// This includes both dynamic learning rates and other hyperparameters (e.g.,
+// step counters for the frequency aware Adagrad optimizer).
+absl::flat_hash_set<int> GetOptimizerDynamicInputTags(
+    const OptimizationParameters& params);
+
+// Returns the set of dynamic hyperparameter tags used by the optimization
+// algorithm. This includes other hyperparameters used by the optimization
+// algorithm (e.g., step counters for the frequency aware Adagrad optimizer). It
+// excludes the dynamic learning rate tag.
+absl::flat_hash_set<int> GetOptimizerHyperParameterTags(
+    const OptimizationParameters& params);
+
+// Returns true if the optimization algorithm uses dynamic inputs in its
+// computation.
+bool UsesDynamicInputsInOptimizer(const OptimizationParameters& params);
+
+// Maximum value of auxiliary_parametery_count for any optimization algorithm.
+// This count is used by TPU embedding load/retrieve and needs to be independent
+// of any particular TPU version and hence, we take the maximum across all TPU
+// versions.
+static constexpr int kMaxAuxiliaryParameterCount = 7;
+
+// Fill value for gradient accumulators. This is a denormal so that it will be
+// flushed to zero on the current TPU platforms and needs to continue to have
+// the following properties in the future:
+//
+// 1. Does not have the same bit pattern as a zero and can be distinguished from
+// it using integer operations.
+// 2. Treated as zero by floating-point arithmetic operations (at least addition
+// and subtraction).
+// 3. Cannot be produced by any floating-point arithmetic operation, including
+// those involving itself.
+//
+// It does not need to compare equal or not equal to zero in floating point. We
+// need to use a non-zero value here because some optimization algorithms are
+// not no-ops on zero gradients, so we need to distinguish an accumulated
+// gradient of zero from one that has been cleared after its gradients have
+// already been applied to the parameters and accumulators.
+inline float GradientAccumulatorInitialValue() {
+  return absl::bit_cast<float, uint32_t>(1);
+}
+
+}  // namespace tpu
+}  // namespace tensorflow
+
+#endif  // XLA_TPU_PROTOBUF_TPU_EMBEDDING_OPTIMIZATION_PARAMETERS_UTILS_H_
diff --git a/tensorflow/core/tpu/tpu_compile_interface.cc b/third_party/xla/xla/tpu/tpu_compile_interface.cc
similarity index 81%
rename from tensorflow/core/tpu/tpu_compile_interface.cc
rename to third_party/xla/xla/tpu/tpu_compile_interface.cc
index 4bbc230a4b0..3bc55fcdcfa 100644
--- a/tensorflow/core/tpu/tpu_compile_interface.cc
+++ b/third_party/xla/xla/tpu/tpu_compile_interface.cc
@@ -1,4 +1,4 @@
-/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+/* Copyright 2026 The OpenXLA Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -13,19 +13,19 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "tensorflow/core/tpu/tpu_compile_interface.h"
+#include "xla/tpu/tpu_compile_interface.h"
 
 #include <cstdint>
 
 #include "absl/log/log.h"
 #include "absl/strings/string_view.h"
-#include "tensorflow/core/platform/fingerprint.h"
-#include "tensorflow/core/platform/logging.h"
+#include "tsl/platform/fingerprint.h"
+#include "tsl/platform/logging.h"
 
 class TpuCompileInterfaceExternal : public TpuCompileInterface {
  public:
   uint64_t FingerprintString(absl::string_view str) override {
-    return ::tensorflow::Fingerprint64(str);
+    return tsl::Fingerprint64(str);
   }
 };
 
diff --git a/third_party/xla/xla/tpu/tpu_compile_interface.h b/third_party/xla/xla/tpu/tpu_compile_interface.h
new file mode 100644
index 00000000000..892ff49017e
--- /dev/null
+++ b/third_party/xla/xla/tpu/tpu_compile_interface.h
@@ -0,0 +1,47 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+#ifndef XLA_TPU_TPU_COMPILE_INTERFACE_H_
+#define XLA_TPU_TPU_COMPILE_INTERFACE_H_
+
+#include <cstdint>
+
+#include "absl/strings/string_view.h"
+
+// Some legacy code requires different implementations for operations like
+// fingerprint/hashing during compilation and/or graph rewriting. These
+// alternate implementations can be registered (via a module initializer) to
+// change the default behavior.
+class TpuCompileInterface {
+ public:
+  virtual ~TpuCompileInterface() {}
+  static TpuCompileInterface* Get();
+  static bool RegisterImplementation(TpuCompileInterface* impl);
+
+  virtual uint64_t FingerprintString(absl::string_view str) = 0;
+
+  // Proto: tensorflow::tpu::CompilationResultProto
+  // Location: tensorflow/core/protobuf/tpu/compilation_result.proto
+  static inline constexpr char kTpuCompileErrorPayloadKey[] =
+      "type.googleapis.com/tensorflow.tpu.CompilationResultProto";
+
+  // Unique string added to the error message for permanent errors during
+  // XLA:TPU compilation. This can be used by TensorFlow models to distinguish
+  // compilation errors from transient errors created by TPU worker preemptions
+  // and restarts.
+  static inline constexpr char kTpuCompileErrorMessage[] =
+      "XLA:TPU compile permanent error";
+};
+
+#endif  // XLA_TPU_TPU_COMPILE_INTERFACE_H_

From 745eaa30d4ba62fb6633f406973c5c022987eb8f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 15 Oct 2026 12:29:56 -0700
Subject: [PATCH 0898/1000] Introduce DtypeDict and normalize dtype keys in
 PlatformInfo

Fixes dtype key hash mismatches in PlatformInfo cadence and latency tables by canonicalizing keys to dtype instances on construction and lookup.

PiperOrigin-RevId: 1000315152
---
 .../xla/xla/benchmarks/core/benchmark.py      |   5 +-
 .../xla/xla/benchmarks/core/benchmark_test.py |  73 ++++++++++++
 .../xla/xla/benchmarks/core/platform_info.py  | 110 ++++++++++++------
 3 files changed, 152 insertions(+), 36 deletions(-)

diff --git a/third_party/xla/xla/benchmarks/core/benchmark.py b/third_party/xla/xla/benchmarks/core/benchmark.py
index 7d6230384a2..80f282094fb 100644
--- a/third_party/xla/xla/benchmarks/core/benchmark.py
+++ b/third_party/xla/xla/benchmarks/core/benchmark.py
@@ -26,6 +26,7 @@ import jax
 import jax.numpy as jnp
 import numpy as np
 
+from xla.benchmarks.core import platform_info
 from xla.benchmarks.jax_microbenchmarks import jax_profiler_utils
 
 
@@ -60,10 +61,10 @@ def str_to_dtype(dtype_str: str) -> jnp.dtype:
   """Converts a string to a JAX/NumPy dtype."""
   if dtype_str not in STR_TO_DTYPE_MAPPING:
     raise ValueError(f"Unsupported dtype: {dtype_str}")
-  return STR_TO_DTYPE_MAPPING[dtype_str]  # pyrefly: ignore[bad-return]
+  return jnp.dtype(STR_TO_DTYPE_MAPPING[dtype_str])
 
 
-DTYPE_TO_STR_MAPPING = immutabledict.immutabledict(
+DTYPE_TO_STR_MAPPING = platform_info.DtypeDict(
     {v: k for k, v in STR_TO_DTYPE_MAPPING.items()}
 )
 
diff --git a/third_party/xla/xla/benchmarks/core/benchmark_test.py b/third_party/xla/xla/benchmarks/core/benchmark_test.py
index 441fc76d4f6..7b3596165cc 100644
--- a/third_party/xla/xla/benchmarks/core/benchmark_test.py
+++ b/third_party/xla/xla/benchmarks/core/benchmark_test.py
@@ -20,10 +20,12 @@ from unittest import mock
 
 from absl.testing import absltest
 import jax
+from jax.experimental.pallas import tpu as pltpu
 import jax.numpy as jnp
 import numpy as np
 
 from xla.benchmarks.core import benchmark
+from xla.benchmarks.core import platform_info
 from xla.benchmarks.jax_microbenchmarks import jax_profiler_utils
 
 InputSpec = benchmark.InputSpec
@@ -200,6 +202,77 @@ class BenchmarkTest(absltest.TestCase):
     with self.assertRaises(ValueError):
       bm.run(runs=0)
 
+  def test_dtype_dict(self):
+    d = platform_info.DtypeDict({jnp.bfloat16: 16, jnp.float32: 32})
+    self.assertEqual(d[jnp.bfloat16], 16)
+    self.assertEqual(d[jnp.dtype("bfloat16")], 16)
+    self.assertEqual(d["bfloat16"], 16)
+    self.assertIn(jnp.bfloat16, d)
+    self.assertIn(jnp.dtype("bfloat16"), d)
+    self.assertIn("bfloat16", d)
+    self.assertNotIn(jnp.int32, d)
+    self.assertNotIn("invalid_dtype_str", d)
+    self.assertNotIn(12345, d)
+    self.assertEqual(d.get(jnp.dtype("bfloat16")), 16)
+    self.assertEqual(d.get(jnp.int32, 0), 0)
+
+  def test_dtype_to_str_accepts_dtype_objects(self):
+    for name, scalar_type in benchmark.STR_TO_DTYPE_MAPPING.items():
+      with self.subTest(name=name):
+        self.assertEqual(benchmark.dtype_to_str(scalar_type), name)
+        self.assertEqual(benchmark.dtype_to_str(jnp.dtype(scalar_type)), name)
+    self.assertEqual(benchmark.dtype_to_str("foo"), "foo")
+
+  def test_dtype_objects_match_scalar_types(self):
+    for chip in (
+        pltpu.ChipVersion.TPU_V5E,
+        pltpu.ChipVersion.TPU_V5P,
+        pltpu.ChipVersion.TPU_V6E,
+        pltpu.ChipVersion.TPU_7X,
+        pltpu.ChipVersion.TPU_8I,
+        pltpu.ChipVersion.TPU_8T,
+    ):
+      pinfo = platform_info.get_platform_info(chip)
+      for scalar_type in (
+          jnp.float32,
+          jnp.bfloat16,
+          jnp.float8_e4m3fn,
+          jnp.float8_e5m2,
+          jnp.float8_e4m3b11fnuz,
+          jnp.int8,
+          jnp.uint8,
+          jnp.int4,
+          jnp.uint4,
+          jnp.float4_e2m1fn,
+      ):
+        dtype_object = jnp.dtype(scalar_type)
+        with self.subTest(chip=chip.value, dtype=dtype_object.name):
+          for mapping in (
+              pinfo.matmul_cadence_cycles_by_dtype,
+              pinfo.matmul_latency_cycles_by_dtype,
+              pinfo.latch_cadence_cycles_by_dtype,
+          ):
+            expected = pinfo.get_emulated_dtype(scalar_type, mapping)
+            actual = pinfo.get_emulated_dtype(dtype_object, mapping)
+            self.assertEqual(actual, expected)
+            self.assertIn(actual, mapping)
+          self.assertEqual(
+              pinfo.get_matmul_cadence_cycles(dtype_object),
+              pinfo.get_matmul_cadence_cycles(scalar_type),
+          )
+          self.assertEqual(
+              pinfo.get_matmul_latency_cycles(dtype_object),
+              pinfo.get_matmul_latency_cycles(scalar_type),
+          )
+          self.assertEqual(
+              pinfo.get_latch_cadence_cycles(dtype_object),
+              pinfo.get_latch_cadence_cycles(scalar_type),
+          )
+          self.assertEqual(
+              pinfo.mxu_size_by_dtype(dtype_object),
+              pinfo.mxu_size_by_dtype(scalar_type),
+          )
+
 
 if __name__ == "__main__":
   absltest.main()
diff --git a/third_party/xla/xla/benchmarks/core/platform_info.py b/third_party/xla/xla/benchmarks/core/platform_info.py
index af6f0c562ec..820788248d8 100644
--- a/third_party/xla/xla/benchmarks/core/platform_info.py
+++ b/third_party/xla/xla/benchmarks/core/platform_info.py
@@ -16,6 +16,7 @@
 
 from collections.abc import Mapping
 import dataclasses
+from typing import Any, Generic, TypeVar, cast
 
 import immutabledict
 import jax
@@ -23,6 +24,40 @@ from jax.experimental.pallas import tpu as pltpu
 import jax.numpy as jnp
 
 immutabledict = immutabledict.immutabledict
+_V = TypeVar("_V")
+
+
+class DtypeDict(immutabledict[jnp.dtype, _V], Generic[_V]):
+  """An immutable dictionary that normalizes keys to jnp.dtype.
+
+  Keys are converted to `jnp.dtype(key)` on dictionary construction and
+  lookup, ensuring that scalar types (e.g. `jnp.bfloat16`), string names
+  (e.g. `'bfloat16'`), and `jnp.dtype` instances can be used interchangeably.
+  """
+
+  def __new__(
+      cls,
+      *args: Any,
+      **kwargs: Any,
+  ) -> "DtypeDict[_V]":
+    raw_dict = dict(*args, **kwargs)
+    normalized = {jnp.dtype(k): v for k, v in raw_dict.items()}
+    return cast("DtypeDict[_V]", super().__new__(cls, normalized))
+
+  def __getitem__(self, key: Any) -> _V:
+    return super().__getitem__(jnp.dtype(key))
+
+  def __contains__(self, key: Any) -> bool:
+    try:
+      return super().__contains__(jnp.dtype(key))
+    except (TypeError, ValueError):
+      return False
+
+  def get(self, key: Any, default: Any = None) -> Any:
+    try:
+      return super().get(jnp.dtype(key), default)
+    except (TypeError, ValueError):
+      return default
 
 
 @dataclasses.dataclass(frozen=True)
@@ -52,13 +87,13 @@ class PlatformInfo:
 
   chip_version: pltpu.ChipVersion
   default_internal_scratch_bytes: int
-  matmul_cadence_cycles_by_dtype: immutabledict[jnp.dtype, int]
-  latch_cadence_cycles_by_dtype: immutabledict[jnp.dtype, int]
+  matmul_cadence_cycles_by_dtype: DtypeDict[int | float]
+  latch_cadence_cycles_by_dtype: DtypeDict[int]
   clock_speed_ghz_by_p_state: immutabledict[int | None, float]
   hbm_to_vmem_bandwidth_gb_per_sec: int
   hbm_to_vmem_latency_ns: int
   vmem_to_hbm_latency_ns: int
-  matmul_latency_cycles_by_dtype: immutabledict[jnp.dtype, int]
+  matmul_latency_cycles_by_dtype: DtypeDict[int]
   default_p_state: int | None
   default_vmem_limit_kib: int
   _tpu_info: pltpu.TpuInfo = dataclasses.field(init=False, repr=False)
@@ -103,25 +138,30 @@ class PlatformInfo:
     return self._tpu_info.generation
 
   def get_emulated_dtype(
-      self, dtype: jnp.dtype, supported_dtypes: Mapping[jnp.dtype, int]
+      self,
+      dtype: jax.typing.DTypeLike,
+      supported_dtypes: Mapping[jnp.dtype, int | float],
   ) -> jnp.dtype:
     """Returns the emulated dtype for the given dtype."""
-    if dtype in supported_dtypes:
-      return dtype
+    target_dtype = jnp.dtype(dtype)
+    if target_dtype in supported_dtypes:
+      return target_dtype
     if (
-        dtype == jnp.float8_e4m3fn
+        target_dtype == jnp.dtype(jnp.float8_e4m3fn)
         and jnp.float8_e4m3b11fnuz in supported_dtypes
     ):
-      return jnp.float8_e4m3b11fnuz
-    dtype_bits = jax.dtypes.itemsize_bits(dtype)
+      return jnp.dtype(jnp.float8_e4m3b11fnuz)
+    dtype_bits = jax.dtypes.itemsize_bits(target_dtype)
     if dtype_bits < 8 and jnp.float8_e4m3fn in supported_dtypes:
-      return jnp.float8_e4m3fn
+      return jnp.dtype(jnp.float8_e4m3fn)
     elif dtype_bits == 8 and jnp.bfloat16 in supported_dtypes:
-      return jnp.bfloat16
+      return jnp.dtype(jnp.bfloat16)
     else:
-      return jnp.float32  # pyrefly: ignore[bad-return]
+      return jnp.dtype(jnp.float32)
 
-  def get_matmul_cadence_cycles(self, lhs_dtype: jnp.dtype) -> int:
+  def get_matmul_cadence_cycles(
+      self, lhs_dtype: jax.typing.DTypeLike
+  ) -> int | float:
     """Returns the number of cycles between each matmul on a single MXU.
 
     Args:
@@ -135,7 +175,7 @@ class PlatformInfo:
     )
     return self.matmul_cadence_cycles_by_dtype[emulated_dtype]
 
-  def get_matmul_latency_cycles(self, lhs_dtype: jnp.dtype) -> int:
+  def get_matmul_latency_cycles(self, lhs_dtype: jax.typing.DTypeLike) -> int:
     """Returns the number of cycles of matmul latency for the given LHS dtype.
 
     Args:
@@ -149,7 +189,7 @@ class PlatformInfo:
     )
     return self.matmul_latency_cycles_by_dtype[emulated_dtype]
 
-  def get_latch_cadence_cycles(self, rhs_dtype: jnp.dtype) -> int:
+  def get_latch_cadence_cycles(self, rhs_dtype: jax.typing.DTypeLike) -> int:
     """Returns the number of cycles between each latch operation.
 
     Args:
@@ -163,7 +203,9 @@ class PlatformInfo:
     )
     return self.latch_cadence_cycles_by_dtype[emulated_dtype]
 
-  def mxu_size_by_dtype(self, lhs_dtype: jnp.dtype) -> tuple[int, int]:
+  def mxu_size_by_dtype(
+      self, lhs_dtype: jax.typing.DTypeLike
+  ) -> tuple[int, int]:
     """Returns the MXU size for the given LHS dtype."""
     if (
         jax.dtypes.itemsize_bits(lhs_dtype) < 8
@@ -180,7 +222,7 @@ _PLATFORM_INFOS = (
     PlatformInfo(
         chip_version=pltpu.ChipVersion.TPU_V5E,
         default_internal_scratch_bytes=73728,
-        matmul_cadence_cycles_by_dtype=immutabledict({  # pyrefly: ignore[bad-argument-type]
+        matmul_cadence_cycles_by_dtype=DtypeDict({
             jnp.float8_e4m3b11fnuz: 32,
             jnp.float8_e5m2: 32,
             jnp.bfloat16: 16,
@@ -190,7 +232,7 @@ _PLATFORM_INFOS = (
             jnp.int4: 16,
             jnp.uint4: 16,
         }),
-        latch_cadence_cycles_by_dtype=immutabledict({  # pyrefly: ignore[bad-argument-type]
+        latch_cadence_cycles_by_dtype=DtypeDict({
             jnp.float8_e4m3b11fnuz: 4,
             jnp.float8_e5m2: 4,
             jnp.bfloat16: 4,
@@ -207,7 +249,7 @@ _PLATFORM_INFOS = (
         hbm_to_vmem_bandwidth_gb_per_sec=1459,
         hbm_to_vmem_latency_ns=1133,
         vmem_to_hbm_latency_ns=439,
-        matmul_latency_cycles_by_dtype=immutabledict({  # pyrefly: ignore[bad-argument-type]
+        matmul_latency_cycles_by_dtype=DtypeDict({
             jnp.float8_e4m3b11fnuz: 131,
             jnp.float8_e5m2: 131,
             jnp.bfloat16: 131,
@@ -223,7 +265,7 @@ _PLATFORM_INFOS = (
     PlatformInfo(
         chip_version=pltpu.ChipVersion.TPU_V5P,
         default_internal_scratch_bytes=73728,
-        matmul_cadence_cycles_by_dtype=immutabledict({  # pyrefly: ignore[bad-argument-type]
+        matmul_cadence_cycles_by_dtype=DtypeDict({
             jnp.float8_e4m3b11fnuz: 32,
             jnp.float8_e5m2: 32,
             jnp.bfloat16: 16,
@@ -233,7 +275,7 @@ _PLATFORM_INFOS = (
             jnp.int4: 16,
             jnp.uint4: 16,
         }),
-        latch_cadence_cycles_by_dtype=immutabledict({  # pyrefly: ignore[bad-argument-type]
+        latch_cadence_cycles_by_dtype=DtypeDict({
             jnp.float8_e4m3b11fnuz: 4,
             jnp.float8_e5m2: 4,
             jnp.bfloat16: 4,
@@ -250,7 +292,7 @@ _PLATFORM_INFOS = (
         hbm_to_vmem_bandwidth_gb_per_sec=1459,
         hbm_to_vmem_latency_ns=667,
         vmem_to_hbm_latency_ns=624,
-        matmul_latency_cycles_by_dtype=immutabledict({  # pyrefly: ignore[bad-argument-type]
+        matmul_latency_cycles_by_dtype=DtypeDict({
             jnp.float8_e4m3b11fnuz: 131,
             jnp.float8_e5m2: 131,
             jnp.bfloat16: 131,
@@ -266,7 +308,7 @@ _PLATFORM_INFOS = (
     PlatformInfo(
         chip_version=pltpu.ChipVersion.TPU_V6E,
         default_internal_scratch_bytes=73728,
-        matmul_cadence_cycles_by_dtype=immutabledict({  # pyrefly: ignore[bad-argument-type]
+        matmul_cadence_cycles_by_dtype=DtypeDict({
             jnp.float8_e4m3b11fnuz: 16,
             jnp.float8_e5m2: 16,
             jnp.bfloat16: 8,
@@ -276,7 +318,7 @@ _PLATFORM_INFOS = (
             jnp.int4: 8,
             jnp.uint4: 8,
         }),
-        latch_cadence_cycles_by_dtype=immutabledict({  # pyrefly: ignore[bad-argument-type]
+        latch_cadence_cycles_by_dtype=DtypeDict({
             jnp.float8_e4m3b11fnuz: 4,
             jnp.float8_e5m2: 4,
             jnp.bfloat16: 4,
@@ -293,7 +335,7 @@ _PLATFORM_INFOS = (
         hbm_to_vmem_bandwidth_gb_per_sec=1373,
         hbm_to_vmem_latency_ns=628,
         vmem_to_hbm_latency_ns=523,
-        matmul_latency_cycles_by_dtype=immutabledict({  # pyrefly: ignore[bad-argument-type]
+        matmul_latency_cycles_by_dtype=DtypeDict({
             jnp.float8_e4m3b11fnuz: 192,
             jnp.float8_e5m2: 192,
             jnp.bfloat16: 192,
@@ -309,13 +351,13 @@ _PLATFORM_INFOS = (
     PlatformInfo(
         chip_version=pltpu.ChipVersion.TPU_7X,
         default_internal_scratch_bytes=73728,
-        matmul_cadence_cycles_by_dtype=immutabledict({  # pyrefly: ignore[bad-argument-type]
+        matmul_cadence_cycles_by_dtype=DtypeDict({
             jnp.float8_e4m3fn: 8,
             jnp.float8_e5m2: 8,
             jnp.bfloat16: 8,
             jnp.float32: 4,
         }),
-        latch_cadence_cycles_by_dtype=immutabledict({  # pyrefly: ignore[bad-argument-type]
+        latch_cadence_cycles_by_dtype=DtypeDict({
             jnp.float8_e4m3fn: 4,
             jnp.float8_e5m2: 4,
             jnp.bfloat16: 4,
@@ -334,7 +376,7 @@ _PLATFORM_INFOS = (
         hbm_to_vmem_bandwidth_gb_per_sec=3207,
         hbm_to_vmem_latency_ns=692,
         vmem_to_hbm_latency_ns=594,
-        matmul_latency_cycles_by_dtype=immutabledict({  # pyrefly: ignore[bad-argument-type]
+        matmul_latency_cycles_by_dtype=DtypeDict({
             jnp.float8_e4m3fn: 204,
             jnp.float8_e5m2: 204,
             jnp.bfloat16: 211,
@@ -346,13 +388,13 @@ _PLATFORM_INFOS = (
     PlatformInfo(
         chip_version=pltpu.ChipVersion.TPU_8I,
         default_internal_scratch_bytes=73728,
-        matmul_cadence_cycles_by_dtype=immutabledict({  # pyrefly: ignore[bad-argument-type]
+        matmul_cadence_cycles_by_dtype=DtypeDict({
             jnp.float8_e4m3fn: 4,
             jnp.float8_e5m2: 4,
             jnp.bfloat16: 16,
             jnp.float32: 8,
         }),
-        latch_cadence_cycles_by_dtype=immutabledict({  # pyrefly: ignore[bad-argument-type]
+        latch_cadence_cycles_by_dtype=DtypeDict({
             jnp.int4: 2,
             jnp.float8_e4m3fn: 2,
             jnp.float8_e5m2: 2,
@@ -372,7 +414,7 @@ _PLATFORM_INFOS = (
         hbm_to_vmem_bandwidth_gb_per_sec=3741,
         hbm_to_vmem_latency_ns=645,
         vmem_to_hbm_latency_ns=512,
-        matmul_latency_cycles_by_dtype=immutabledict({  # pyrefly: ignore[bad-argument-type]
+        matmul_latency_cycles_by_dtype=DtypeDict({
             jnp.float8_e4m3fn: 211,
             jnp.float8_e5m2: 211,
             jnp.bfloat16: 235,
@@ -384,7 +426,7 @@ _PLATFORM_INFOS = (
     PlatformInfo(
         chip_version=pltpu.ChipVersion.TPU_8T,
         default_internal_scratch_bytes=278528,
-        matmul_cadence_cycles_by_dtype=immutabledict({  # pyrefly: ignore[bad-argument-type]
+        matmul_cadence_cycles_by_dtype=DtypeDict({
             jnp.int4: 16 / 3,
             jnp.float4_e2m1fn: 16 / 3,
             jnp.float8_e4m3fn: 16 / 3,
@@ -392,7 +434,7 @@ _PLATFORM_INFOS = (
             jnp.bfloat16: 16,
             jnp.float32: 8,
         }),
-        latch_cadence_cycles_by_dtype=immutabledict({  # pyrefly: ignore[bad-argument-type]
+        latch_cadence_cycles_by_dtype=DtypeDict({
             jnp.int4: 4,
             jnp.float4_e2m1fn: 4,
             jnp.float8_e4m3fn: 4,
@@ -410,7 +452,7 @@ _PLATFORM_INFOS = (
         hbm_to_vmem_bandwidth_gb_per_sec=5679,
         hbm_to_vmem_latency_ns=821,
         vmem_to_hbm_latency_ns=821,
-        matmul_latency_cycles_by_dtype=immutabledict({  # pyrefly: ignore[bad-argument-type]
+        matmul_latency_cycles_by_dtype=DtypeDict({
             jnp.int4: 225,
             jnp.float4_e2m1fn: 225,
             jnp.float8_e4m3fn: 225,

From a665106fd5f22f9b91114ec62ae11d4aeefa6c81 Mon Sep 17 00:00:00 2001
From: Sean Talts <talts@google.com>
Date: Thu, 15 Oct 2026 12:39:51 -0700
Subject: [PATCH 0899/1000] [XLA:CPU] Add vectorized sin and cos intrinsics for
 f32 and f64

Add custom target-independent C++ intrinsics (`xla/codegen/intrinsic/cpp/custom_cpp_intrinsics.{h,cc}`) for `sin` and `cos` across scalar and SIMD vector widths (`f32`, `v4f32`, `v8f32`, `v16f32`, `f64`, `v2f64`, `v4f64`, `v8f64`), compiled to LLVM bitcode (`custom_cpp_intrinsics_ll`), linked through `GetCppGenLibraries`, and registered via `intrinsics::Sin` and `intrinsics::Cos`.

Key details:
- f32 is evaluated in double precision with one final rounding: 3-part Cody-Waite reduction for `|x| < 2^28` (2-part when every lane is below 2^12), Payne-Hanek above, and short minimax polynomials. Max error is 0.5000x real ULP (702 sin / 812 cos misrounded inputs out of all 2^32 floats, each below 0.5 + 2^-18 ULP).
- f64 uses 4-part Cody-Waite reduction for `|x| < 2^20` and Payne-Hanek above, with double-word accumulation of the r^3, r^4 and r^5 terms. Max error is about 0.504 real ULP over 1M random inputs.
- The Cody-Waite splits of pi/2 have enough trailing zero bits that the reduction is exact with or without hardware FMA. Error terms use `@llvm.fmuladd`, so CPUs without FMA get no `fma` libcalls; there the error terms vanish and f64 max error rises to about 1.46 ULP over 1M inputs (f32 stays at 0.5000 over 2^24 inputs).
- Products whose rounding error is captured (`TwoProd`) use `fmuladd(a, b, +0.0)` so the `contract` flag that `IrCompiler` adds to every `fmul` cannot fuse them into a later add.
- 16-lane f32 and 8-lane f64 variants are only advertised with `+avx512f`.
- Lowers `math.sin` and `math.cos` to `xla.sin.*` and `xla.cos.*` in `xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc`.

<details>

<summary>Numerics Design</summary>

# Numerics Design: Vectorized Sine and Cosine Intrinsics

## 1. Overview
This implementation provides target-independent C++ intrinsics for `sin` and `cos` across scalar and SIMD vector widths for both `f32` and `f64`. The design prioritizes achieving an error bound tightly constrained to $\le 0.5$ ULP (Unit in the Last Place) for most inputs, while maintaining high performance by avoiding expensive fallback library calls where possible.

## 2. Argument Reduction Patterns
Trigonometric functions require mapping the input argument $x$ into a primary range (typically $[-\pi/4, \pi/4]$) by calculating $n = \text{round}(x \times 2/\pi)$ and $r = x - n \times \pi/2$.

### 2.1 Fast-Path Reduction: Cody-Waite
For inputs that are relatively small, the code uses a **Cody-Waite reduction**, splitting $\pi/2$ into multiple parts to preserve exactness in the subtraction.
*   **f32 Implementation:** The reduction relies on a 3-part split for inputs where $|x| < 2^{28}$. For highly localized SIMD vectors where every lane is strictly below $2^{12}$, it takes an even faster 2-part split path.
*   **f64 Implementation:** The reduction requires a 4-part split for inputs where $|x| < 2^{20}$.

### 2.2 Slow-Path Reduction: Payne-Hanek
For inputs exceeding the Cody-Waite thresholds, catastrophic cancellation would occur. The code falls back to a **Payne-Hanek range reduction**, which evaluates $x \times 2/\pi$ using an extended-precision table lookup.
*   The code maintains a 21-word table representing $1280$ bits of $2/\pi$.
*   To minimize stack spilling in hot vector code, the Payne-Hanek loop is marked `noinline, cold` and passes variables by value via a `ReducedArgument` struct.

## 3. Core Approximation Patterns

### 3.1 Magic Number Rounding (Quadrant Extraction)
To compute the integer $n$ (the quadrant multiplier), the code uses the "magic number" addition pattern.
*   By adding and subtracting the magic constant $1.5 \times 2^{52}$, the hardware forces a round-to-nearest-even operation into an integer.
*   This specific magic number cleverly places the value $n \pmod 4$ directly into the lowest 2 mantissa bits of the result, which is then extracted to determine the quadrant sign and function flip (sin vs. cos).

### 3.2 Error-Free Transformations (EFT) and Double-Word Arithmetic
For double precision (`f64`), standard polynomial evaluation doesn't have enough precision to stay within 0.5 ULP. The implementation uses Error-Free Transformations (EFT) to track and accumulate rounding errors.
*   **FastTwoSum (Dekker's):** Used to compute the exact sum and rounding error of two numbers when their magnitudes are ordered.
*   **TwoProd:** Used to capture the exact product and its rounding error via fused multiply-add operations.
*   **Double-word Accumulation:** The largest terms of the Taylor series (like $r^3$ and $r^5$ for sine, and $r^4$ for cosine) are accumulated as double-words (a high and low component) to preserve the rounding errors of the coefficients and the powers of $r$.

### 3.3 Minimax Polynomial Evaluation (Horner's Method)
After extracting the lower-order error components using double-word arithmetic, the high-order polynomial tails are evaluated using Minimax approximations.
*   **Tail Polynomials:** Terms of order $r^6$ and above are evaluated as standard polynomials and accumulated into the "low" word component.
*   They are computed optimally using Horner's method (`HornerPoly`), evaluating the sequence from highest degree down to lowest.

### 3.4 Type Promotion / Upcasting (`f32`)
Because `f32` vectors require less precision, the complexity of double-word tracking is avoided completely.
*   The `f32` implementation simply upcasts to `double` internally and evaluates the entire calculation in 64-bit precision, falling back to 5-coefficient Minimax polynomials before applying a single rounding operation at the end.
*   This makes the `f32` implementation mathematically robust (maximum error of 0.5000x real ULP) and extremely fast on hardware that easily interleaves double precision.

## 4. Hardware FMA Safety & Graceful Degradation
A major design constraint of this library is avoiding catastrophic performance drops on hardware lacking Fused Multiply-Add (FMA) instructions, which would otherwise fallback to slow `fma` library calls.
*   **`TwoProd` Safety:** The exact product is emitted as `@llvm.fmuladd(a, b, +0.0)`, preventing downstream compiler passes from aggressively fusing the result into subsequent addition operations, which would break the error term calculation.
*   **EFT Graceful Degradation:** The error term in `TwoProd` is calculated as `fmuladd(a, b, -p)`. On non-FMA hardware, the `fmuladd` expands to separate instructions, and the error term cleanly evaluates to `0`. This allows the exact reductions to stay exact, but causes the double-word tracking to quietly vanish. Consequently, accuracy degrades slightly on non-FMA hardware (to $\sim 1.46$ ULP for `f64`) while maintaining high execution speed without triggering costly software fallbacks.
</details>

PiperOrigin-RevId: 1000321344
---
 .../transforms/lower_xla_intrinsic_lib.cc     |   2 +
 .../tests/lower_xla_intrinsic_lib.mlir        |  14 +
 .../xla/xla/codegen/intrinsic/cpp/BUILD       |  42 ++
 .../intrinsic/cpp/cpp_gen_intrinsics.cc       |   8 +-
 .../intrinsic/cpp/custom_cpp_intrinsics.cc    | 422 ++++++++++++++++++
 .../intrinsic/cpp/custom_cpp_intrinsics.h     |  45 ++
 .../cpp/custom_cpp_intrinsics_test.cc         | 175 ++++++++
 .../intrinsic/cpp/intrinsic_declarations.h    |  55 ++-
 .../xla/codegen/intrinsic/cpp/vector_ops.h    | 203 ++++++---
 third_party/xla/xla/codegen/intrinsic_lib.cc  |   4 +
 .../xla/xla/codegen/intrinsic_lib_test.cc     |  47 ++
 11 files changed, 947 insertions(+), 70 deletions(-)
 create mode 100644 third_party/xla/xla/codegen/intrinsic/cpp/custom_cpp_intrinsics.cc
 create mode 100644 third_party/xla/xla/codegen/intrinsic/cpp/custom_cpp_intrinsics.h
 create mode 100644 third_party/xla/xla/codegen/intrinsic/cpp/custom_cpp_intrinsics_test.cc

diff --git a/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc b/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc
index a88e4357252..8d8f5026931 100644
--- a/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc
+++ b/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc
@@ -341,6 +341,8 @@ class LowerXlaIntrinsicLibPass
                  LowerIntrinsicPattern<ci::Rsqrt, mm::RsqrtOp>,
                  LowerIntrinsicPattern<ci::Tanh, mm::TanhOp>,
                  LowerIntrinsicPattern<ci::EigenAtan, mm::AtanOp>,
+                 LowerIntrinsicPattern<ci::Sin, mm::SinOp>,
+                 LowerIntrinsicPattern<ci::Cos, mm::CosOp>,
                  LowerIntrinsicPattern<ci::FpTrunc, ma::TruncFOp>,
                  LowerIntrinsicPattern<ci::Erf, mm::ErfOp>>(context,
                                                             cpu_features);
diff --git a/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir b/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir
index d0df27f8185..5d258683b66 100644
--- a/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir
+++ b/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir
@@ -438,3 +438,17 @@ func.func @expm1_f32_vector_16(%arg0: vector<16xf32>) -> vector<16xf32> {
 // AVX512-LABEL: @expm1_f32_vector_16
 // AVX512: %[[CALL:.*]] = call @xla.expm1.v16f32(%arg0) : (vector<16xf32>) -> vector<16xf32>
 // AVX512: return %[[CALL]]
+
+// -----
+
+func.func @sin_cos_vector(%arg0: vector<8xf32>, %arg1: vector<4xf64>) -> (vector<8xf32>, vector<4xf64>) {
+  %s = math.sin %arg0 : vector<8xf32>
+  %c = math.cos %arg1 : vector<4xf64>
+  return %s, %c : vector<8xf32>, vector<4xf64>
+}
+// CHECK-LABEL: @sin_cos_vector
+// CHECK-NOT: math.sin
+// CHECK-NOT: math.cos
+// CHECK: %[[SIN:.*]] = call @xla.sin.v8f32(%arg0) : (vector<8xf32>) -> vector<8xf32>
+// CHECK: %[[COS:.*]] = call @xla.cos.v4f64(%arg1) : (vector<4xf64>) -> vector<4xf64>
+// CHECK: return %[[SIN]], %[[COS]]
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/BUILD b/third_party/xla/xla/codegen/intrinsic/cpp/BUILD
index f0ef230b757..2c280ebd283 100644
--- a/third_party/xla/xla/codegen/intrinsic/cpp/BUILD
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/BUILD
@@ -312,11 +312,53 @@ cc_binary(
     ],
 )
 
+cc_library(
+    name = "custom_cpp_intrinsics_lib",
+    srcs = ["custom_cpp_intrinsics.cc"],
+    hdrs = ["custom_cpp_intrinsics.h"],
+    copts = [
+        "-Wno-psabi",
+    ],
+    deps = [
+        ":vector_ops",
+    ],
+)
+
+cc_ir_header(
+    name = "custom_cpp_intrinsics_ll",
+    src = "custom_cpp_intrinsics.cc",
+    deps = [
+        ":cpp_context_provider",
+        ":custom_cpp_intrinsics_lib",
+        ":vector_ops",
+    ],
+)
+
+xla_cc_test(
+    name = "custom_cpp_intrinsics_test",
+    srcs = ["custom_cpp_intrinsics_test.cc"],
+    copts = [
+        "-Wno-psabi",
+    ],
+    deps = [
+        ":_cpp_gen_intrinsics",
+        ":custom_cpp_intrinsics_lib",
+        ":custom_cpp_intrinsics_ll",
+        ":vector_ops",
+        "//xla/codegen/intrinsic:test_matchers",
+        "@com_google_absl//absl/base",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//llvm:ir_headers",
+    ],
+)
+
 cc_library(
     name = "_cpp_gen_intrinsics",
     srcs = ["cpp_gen_intrinsics.cc"],
     hdrs = ["cpp_gen_intrinsics.h"],
     deps = [
+        ":custom_cpp_intrinsics_ll",
         ":eigen_unary_16_ll",
         ":eigen_unary_32_ll",
         ":eigen_unary_64_ll",
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/cpp_gen_intrinsics.cc b/third_party/xla/xla/codegen/intrinsic/cpp/cpp_gen_intrinsics.cc
index af0108bbc1f..0ad8be14e79 100644
--- a/third_party/xla/xla/codegen/intrinsic/cpp/cpp_gen_intrinsics.cc
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/cpp_gen_intrinsics.cc
@@ -44,6 +44,7 @@ limitations under the License.
 #include "llvm/Support/Alignment.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/raw_ostream.h"
+#include "xla/codegen/intrinsic/cpp/custom_cpp_intrinsics_ll.h"
 #include "xla/codegen/intrinsic/cpp/eigen_unary_16_ll.h"
 #include "xla/codegen/intrinsic/cpp/eigen_unary_32_ll.h"
 #include "xla/codegen/intrinsic/cpp/eigen_unary_64_ll.h"
@@ -97,7 +98,8 @@ std::vector<CppGenIntrinsicLibrary> GetCppGenLibraries(
   std::vector<CppGenIntrinsicLibrary> libraries;
   if (options.device_type != intrinsics::DeviceType::kIntelCpu &&
       options.device_type != intrinsics::DeviceType::kAmdCpu &&
-      options.device_type != intrinsics::DeviceType::kArmCpu) {
+      options.device_type != intrinsics::DeviceType::kArmCpu &&
+      options.device_type != intrinsics::DeviceType::kSystemZCpu) {
     return libraries;
   }
   const std::string& eigen_ir = GetCppGenIrString(options);
@@ -109,6 +111,10 @@ std::vector<CppGenIntrinsicLibrary> GetCppGenLibraries(
   if (!expm1_ir.empty()) {
     libraries.emplace_back(expm1_ir, "expm1");
   }
+  if (!::llvm_ir::kCustomCppIntrinsicsLlIr.empty()) {
+    libraries.emplace_back(::llvm_ir::kCustomCppIntrinsicsLlIr,
+                           "custom_cpp_intrinsics");
+  }
   return libraries;
 }
 
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/custom_cpp_intrinsics.cc b/third_party/xla/xla/codegen/intrinsic/cpp/custom_cpp_intrinsics.cc
new file mode 100644
index 00000000000..010157eb007
--- /dev/null
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/custom_cpp_intrinsics.cc
@@ -0,0 +1,422 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#if defined(__has_attribute) && __has_attribute(ext_vector_type) && \
+    defined(__has_builtin) && __has_builtin(__builtin_vectorelements)
+
+#include "xla/codegen/intrinsic/cpp/custom_cpp_intrinsics.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <type_traits>
+
+#include "xla/codegen/intrinsic/cpp/vector_ops.h"
+
+namespace xla::codegen {
+namespace {
+
+constexpr double kInf = std::numeric_limits<double>::infinity();
+
+// Every multiply-add below is `Fmuladd`, so targets without FMA get a separate
+// multiply and add instead of an `fma` libcall. The Cody-Waite reductions are
+// built from products that are exact without FMA, so they stay accurate on
+// such targets; only the error terms of `TwoProd` in the polynomials are lost.
+
+// Lanes with `|x|` at or above this use Payne-Hanek instead of Cody-Waite.
+// Below them `n = round(|x| * 2/pi)` has at most 20 (f64) or 28 (f32) bits,
+// which bounds the width of the pi/2 parts that `n` can multiply exactly.
+constexpr double kF64FastThreshold = 0x1.0p+20;
+constexpr double kF32FastThreshold = 0x1.0p+28;
+
+// Adding 1.5 * 2^52 sets the ULP of `j` to 1.0 so round-to-nearest-even rounds
+// `ax * (2/pi)` to an integer `n` and places `n mod 4` in the low 2 mantissa
+// bits of `j`. Without FMA `n` can be off by one near a half-integer, which
+// only lets |r| exceed pi/4 by a few ULPs.
+constexpr double kMagic = 0x1.8p+52;
+constexpr double kTwoOverPi = 0x1.45f306dc9c883p-1;
+// pi/2 rounded to 53 bits and the next 53 bits, used by Payne-Hanek.
+constexpr double kPiOver2Hi = 0x1.921fb54442d18p+0;
+constexpr double kPiOver2Mid = 0x1.1a62633145c07p-54;
+
+// pi/2 split for the f64 fast path. Each part is the remainder of pi/2 after
+// the previous parts, rounded to nearest with at most 33, 32, 32 and 53 bits,
+// which leaves 31, 32, 28 and 53 significant bits; the sum is within 2^-159 of
+// pi/2. With |n| < 2^20 the products of n with the first three parts fit in 53
+// bits and are exact.
+constexpr double kF64PiOver2A = 0x1.921fb544p+0;
+constexpr double kF64PiOver2B = 0x1.0b4611a6p-34;
+constexpr double kF64PiOver2C = 0x1.3198a2ep-69;
+constexpr double kF64PiOver2D = 0x1.b839a252049c1p-104;
+
+// pi/2 split for the f32 fast path: `kPiOver2Hi` cut after 25 bits into
+// 0x1.921fb5p+0 (25 bits) and 0x1.110b46p-26 (24 bits) so that their products
+// with the 28-bit n are exact, followed by `kPiOver2Mid`.
+constexpr double kF32PiOver2A = 0x1.921fb5p+0;
+constexpr double kF32PiOver2B = 0x1.110b46p-26;
+
+// Vectors with every |x| below 2^12 (so |n| < 2^12) use a cheaper 2-part
+// split: pi/2 rounded to 41 bits, so that its product with n is exact, and the
+// rounded remainder.
+constexpr double kF32SmallThreshold = 0x1.0p+12;
+constexpr double kF32SmallPiOver2A = 0x1.921fb54443p+0;
+constexpr double kF32SmallPiOver2B = -0x1.73dcb3b399d74p-43;
+
+// -1/6 and 1/24 as hi + lo pairs. The r^3 term of sin is up to 11% and the r^4
+// term of cos up to 2% of the result, so a single rounded coefficient would
+// cost a visible fraction of an ULP.
+constexpr double kSinS1 = -0x1.5555555555555p-3;
+constexpr double kSinS1Lo = -0x1.5555555555555p-57;
+constexpr double kSinS2 = 0x1.1111111111111p-7;
+constexpr double kSinS2Lo = 0x1.1111111111111p-63;
+constexpr double kCosC1 = 0x1.5555555555555p-5;
+constexpr double kCosC1Lo = 0x1.5555555555555p-59;
+
+// (sin(r) - r + r^3/6 - r^5/120) / r^7 as a polynomial in r^2 on
+// [-pi/4, pi/4].
+template <typename T>
+T EvalF64SinTailPoly(T r2) {
+  return HornerPoly(r2, -0x1.2f49b46814157p-57, 0x1.952c77030ad4ap-49,
+                    -0x1.ae7f3e733b81fp-41, 0x1.6124613a86d09p-33,
+                    -0x1.ae64567f544e4p-26, 0x1.71de3a556c734p-19,
+                    -0x1.a01a01a01a01ap-13);
+}
+
+// (cos(r) - 1 + r^2/2 - r^4/24) / r^6 as a polynomial in r^2 on [-pi/4, pi/4].
+template <typename T>
+T EvalF64CosTailPoly(T r2) {
+  return HornerPoly(r2, 0x1.e542ba4020225p-62, -0x1.6827863b97d97p-53,
+                    0x1.ae7f3e733b81fp-45, -0x1.93974a8c07c9dp-37,
+                    0x1.1eed8eff8d898p-29, -0x1.27e4fb7789f5cp-22,
+                    0x1.a01a01a01a01ap-16, -0x1.6c16c16c16c17p-10);
+}
+
+// Minimax fit of (sin(r) - r) / r^3 as a polynomial in r^2 on [-pi/4, pi/4].
+// r + r^3 * P(r^2) has relative error below 2^-47.
+template <typename T>
+T EvalF32SinPoly(T r2) {
+  return HornerPoly(r2, -0x1.a961a17edc631p-26, 0x1.71d76d1676743p-19,
+                    -0x1.a019f93968c5bp-13, 0x1.1111110c87334p-7,
+                    -0x1.5555555552239p-3);
+}
+
+// Minimax fit of (cos(r) - 1) / r^2 as a polynomial in r^2 on [-pi/4, pi/4].
+// 1 + r^2 * P(r^2) has relative error below 2^-43.
+template <typename T>
+T EvalF32CosPoly(T r2) {
+  return HornerPoly(r2, -0x1.23c97e6419bd1p-22, 0x1.a00eb9af46668p-16,
+                    -0x1.6c16b348bc373p-10, 0x1.55555545c514ep-5,
+                    -0x1.ffffffffe98afp-2);
+}
+
+// 21 words of 2/pi: word 0 is the integer part (0), and words 1..20 hold bits
+// [1..1280] after the binary point of 2/pi (`sum_{k=1}^{20} W[k] * 2^{-64*k}`).
+static constexpr uint64_t kTwoOverPiTable[21] = {
+    0x0000000000000000ULL, 0xa2f9836e4e441529ULL, 0xfc2757d1f534ddc0ULL,
+    0xdb6295993c439041ULL, 0xfe5163abdebbc561ULL, 0xb7246e3a424dd2e0ULL,
+    0x06492eea09d1921cULL, 0xfe1deb1cb129a73eULL, 0xe88235f52ebb4484ULL,
+    0xe99c7026b45f7e41ULL, 0x3991d639835339f4ULL, 0x9c845f8bbdf9283bULL,
+    0x1ff897ffde05980fULL, 0xef2f118b5a0a6d1fULL, 0x6d367ecf27cb09b7ULL,
+    0x4f463f669e5fea2dULL, 0x7527bac7ebe5f17bULL, 0x3d0739f78a5292eaULL,
+    0x6bfb5fb11f8d5d08ULL, 0x56033046fc7b6babULL, 0xf0cfbc209af4361dULL,
+};
+
+struct PayneHanekResult {
+  uint64_t q_raw;
+  double r_hi;
+  double r_lo;
+};
+
+// Cold-path Payne-Hanek range reduction for a finite double lane at or above
+// the fast-path threshold (`ax >= 2^20`).
+inline PayneHanekResult PayneHanekLaneF64(double ax) {
+  using uint128_t = unsigned __int128;
+
+  uint64_t bits = __builtin_bit_cast(uint64_t, ax);
+  int exp = static_cast<int>((bits >> 52) & 0x7ff) - 1023;
+  // Shift 53-bit significand left by 2 so that with e = exp + 10, the binary
+  // point of the product lies at bit 64 of p2.
+  uint64_t m = ((bits & ((1ULL << 52) - 1)) | (1ULL << 52)) << 2;
+  int e = exp + 10;
+  int idx = e >> 6;
+  int shift = e & 63;
+
+  uint64_t w0 = kTwoOverPiTable[idx];
+  uint64_t w1 = kTwoOverPiTable[idx + 1];
+  uint64_t w2 = kTwoOverPiTable[idx + 2];
+
+  uint64_t hi = (shift == 0) ? w0 : ((w0 << shift) | (w1 >> (64 - shift)));
+  uint64_t mid = (shift == 0) ? w1 : ((w1 << shift) | (w2 >> (64 - shift)));
+
+  uint128_t p1 = static_cast<uint128_t>(m) * mid;
+  uint128_t p2 = static_cast<uint128_t>(m) * hi + (p1 >> 64);
+  uint128_t u128 = (p2 << 64) | static_cast<uint64_t>(p1);
+  __int128 s128 = static_cast<__int128>(u128);
+  int64_t i_hi = static_cast<int64_t>(s128 >> 76);
+
+  // When fewer than 14 fractional bits cancel, the 2-word product (hi, mid)
+  // already provides >= 90 bits of relative precision; otherwise fold in the
+  // third word (lo) and low 24 bits (i_lo).
+  int64_t top14 = i_hi >> 38;
+  double f_hi, f_lo;
+  if (__builtin_expect(top14 != 0 && top14 != -1 && top14 != -(1LL << 13), 1)) {
+    int64_t i_mid = static_cast<int64_t>((u128 >> 24) & ((1ULL << 52) - 1));
+    double d_hi = static_cast<double>(i_hi) * 0x1.0p-52;
+    double d_mid = static_cast<double>(i_mid) * 0x1.0p-104;
+    auto [fh, fl] = FastTwoSum(d_hi, d_mid);
+    f_hi = fh;
+    f_lo = fl;
+  } else {
+    uint64_t w3 = kTwoOverPiTable[idx + 3];
+    uint64_t lo = (shift == 0) ? w2 : ((w2 << shift) | (w3 >> (64 - shift)));
+    p1 += (static_cast<uint128_t>(m) * lo) >> 64;
+    p2 = static_cast<uint128_t>(m) * hi + (p1 >> 64);
+    u128 = (p2 << 64) | static_cast<uint64_t>(p1);
+    s128 = static_cast<__int128>(u128);
+    i_hi = static_cast<int64_t>(s128 >> 76);
+    int64_t i_mid = static_cast<int64_t>((u128 >> 24) & ((1ULL << 52) - 1));
+    int64_t i_lo = static_cast<int64_t>(u128 & ((1ULL << 24) - 1));
+    double d_hi = static_cast<double>(i_hi) * 0x1.0p-52;
+    double d_mid = static_cast<double>(i_mid) * 0x1.0p-104;
+    double d_lo = static_cast<double>(i_lo) * 0x1.0p-128;
+    auto [s1, e1] = FastTwoSum(d_mid, d_lo);
+    auto [f0_hi, f0_lo] = FastTwoSum(d_hi, s1);
+    auto [fh, fl] = FastTwoSum(f0_hi, f0_lo + e1);
+    f_hi = fh;
+    f_lo = fl;
+  }
+
+  uint64_t q_raw =
+      static_cast<uint64_t>((p2 + (uint128_t{1} << 63)) >> 64) & 3ULL;
+  auto [prod_hi, prod_err] = TwoProd(f_hi, kPiOver2Hi);
+  double prod_lo =
+      Fmuladd(f_lo, kPiOver2Hi, Fmuladd(f_hi, kPiOver2Mid, prod_err));
+  auto [r_hi, r_lo] = FastTwoSum(prod_hi, prod_lo);
+  return {q_raw, r_hi, r_lo};
+}
+
+template <typename T, typename IntT>
+struct ReducedArgument {
+  IntT q_raw;
+  T r_hi;
+  T r_lo;
+};
+
+// Returns `fast` with the Payne-Hanek reduction in every finite lane with
+// `ax >= threshold`. Marked `noinline, cold` so `LinkIntoModule` /
+// `PrepareForInlining` keeps the table lookup out of hot vector code while
+// avoiding per-lane `noinline` calls and dynamic-index stack spills. The
+// reduction is passed by value: passing the caller's variables by reference
+// would make their addresses escape and force stack stores on the hot path.
+template <typename T, typename IntT>
+__attribute__((noinline, cold)) ReducedArgument<T, IntT> PayneHanekColdLoop(
+    T ax, double threshold, ReducedArgument<T, IntT> fast) {
+  for (size_t lane = 0; lane < NumLanes<T>; ++lane) {
+    double ax_lane = GetLane(ax, lane);
+    if (ax_lane >= threshold && ax_lane < kInf) {
+      auto [q_lane, rh_lane, rl_lane] = PayneHanekLaneF64(ax_lane);
+      SetLane(fast.q_raw, lane, q_lane);
+      SetLane(fast.r_hi, lane, rh_lane);
+      SetLane(fast.r_lo, lane, rl_lane);
+    }
+  }
+  return fast;
+}
+
+// Picks sin or cos of the reduced argument by quadrant, applies the quadrant
+// sign (and for sin the sign of `x`), and maps +-Inf/NaN inputs to NaN.
+template <bool kIsCos, typename T, typename IntT>
+inline T ApplyQuadrant(T x, T ax, IntT q_raw, T sin_val, T cos_val) {
+  IntT q = (q_raw + Splat<IntT>(kIsCos ? 1 : 0)) & Splat<IntT>(3);
+  T y = ((q & Splat<IntT>(1)) != Splat<IntT>(0)) ? cos_val : sin_val;
+
+  // Bit 1 of `q` is the quadrant sign (`q in {2, 3}`); shift it into the
+  // IEEE-754 sign bit 63.
+  IntT sign_bit = (q & Splat<IntT>(2)) << 62;
+  if constexpr (!kIsCos) {
+    sign_bit ^= (__builtin_bit_cast(IntT, x) & Splat<IntT>(1ULL << 63));
+  }
+  y = __builtin_bit_cast(T, __builtin_bit_cast(IntT, y) ^ sign_bit);
+  return (ax < Splat<T>(kInf)) ? y : (x - x);
+}
+
+template <typename To, typename From>
+inline To ConvertLanes(From v) {
+  if constexpr (std::is_arithmetic_v<From>) {
+    return static_cast<To>(v);
+  } else {
+    return __builtin_convertvector(v, To);
+  }
+}
+
+// f64 sin/cos with a double-word reduced argument. With |r| <= pi/4 the result
+// is at least 0.7 * |r| (sin) or 0.7 (cos), so dropping the rounding error of a
+// term that is a fraction f of the result costs up to about f / 0.7 ULP, and
+// even a few hundredths of an ULP visibly increase the misrounding rate. The
+// r^3 and r^5 terms of sin (f <= 0.11 and 0.0035) and the r^4 term of cos
+// (f <= 0.022) are therefore accumulated as double-words, including the
+// rounding of the powers of r and of the split coefficients. Terms of order r^6
+// and above (f <= 0.0005) are plain polynomials added into the low word.
+template <typename T, bool kIsCos>
+inline T SinCosF64(T x) {
+  using IntT = typename internal::CorrespondingIntVector<T>::type;
+
+  T ax = BitwiseAbs(x);
+  T j = Fmuladd(ax, Splat<T>(kTwoOverPi), Splat<T>(kMagic));
+  T n = j - Splat<T>(kMagic);
+  IntT q_raw = __builtin_bit_cast(IntT, j) & Splat<IntT>(3);
+  // r = ax - n * pi/2 as a double-word with every rounding error kept. The
+  // products n * A, n * B and n * C are exact. `ax - n * A` is exact because
+  // both are multiples of 2^-53 (ax > 0.5 unless n = 0) and the difference is
+  // below 1. `FastTwoSum` needs |a| >= |b| unless `a + b` is exact, and it is:
+  // if |t1| < |n * B| then t1 - n * B is a multiple of 2^-65 below
+  // 2 * 2^20 * 2^-33.9 = 2^-12.9, and if |s1| < |n * C| then s1 - n * C is a
+  // multiple of 2^-96 below 2^-47.7, so both fit in 53 bits. What remains is
+  // the rounding of e1 + e2 and of n * D (< 2^-83), far below the smallest
+  // |r| of 2^-60.5 (at x = 0x1.6c6cbc45dc8dep+5) times 2^-53.
+  T t1 = Fmuladd(-n, Splat<T>(kF64PiOver2A), ax);
+  auto [s1, e1] = FastTwoSum(t1, -n * Splat<T>(kF64PiOver2B));
+  auto [s2, e2] = FastTwoSum(s1, -n * Splat<T>(kF64PiOver2C));
+  T lo = Fmuladd(-n, Splat<T>(kF64PiOver2D), e1 + e2);
+  auto [r_hi, r_lo] = FastTwoSum(s2, lo);
+
+  auto need_cold = (ax >= Splat<T>(kF64FastThreshold)) & (ax < Splat<T>(kInf));
+  if (__builtin_expect(AnyLane(need_cold), 0)) {
+    auto cold = PayneHanekColdLoop(ax, kF64FastThreshold,
+                                   ReducedArgument<T, IntT>{q_raw, r_hi, r_lo});
+    q_raw = cold.q_raw;
+    r_hi = cold.r_hi;
+    r_lo = cold.r_lo;
+  }
+
+  auto [r2, e_r2] = TwoProd(r_hi, r_hi);
+  auto [r3, r3_err] = TwoProd(r2, r_hi);
+  T e_r3 = Fmuladd(e_r2, r_hi, r3_err);
+  auto [r4, r4_err] = TwoProd(r2, r2);
+  T e_r4 = Fmuladd(Splat<T>(2.0) * r2, e_r2, r4_err);
+
+  T w = Fmuladd(Splat<T>(-0.5), r2, Splat<T>(1.0));
+  T w_lo = Fmuladd(Splat<T>(-0.5), e_r2,
+                   Fmuladd(Splat<T>(-0.5), r2, Splat<T>(1.0) - w));
+  auto [c4_hi, c4_err] = TwoProd(Splat<T>(kCosC1), r4);
+  T c4_lo =
+      Fmuladd(Splat<T>(kCosC1), e_r4, Fmuladd(Splat<T>(kCosC1Lo), r4, c4_err));
+  auto [c_hi, c_lo] = FastTwoSum(w, c4_hi);
+
+  auto [t3_hi, t3_err] = TwoProd(Splat<T>(kSinS1), r3);
+  T t3_lo =
+      Fmuladd(Splat<T>(kSinS1), e_r3, Fmuladd(Splat<T>(kSinS1Lo), r3, t3_err));
+  auto [s_hi, s_lo] = FastTwoSum(r_hi, t3_hi);
+  auto [r5, r5_err] = TwoProd(r3, r2);
+  T e_r5 = Fmuladd(r3, e_r2, Fmuladd(e_r3, r2, r5_err));
+  auto [t5_hi, t5_err] = TwoProd(Splat<T>(kSinS2), r5);
+  T t5_lo =
+      Fmuladd(Splat<T>(kSinS2), e_r5, Fmuladd(Splat<T>(kSinS2Lo), r5, t5_err));
+  auto [s2_hi, s2_lo] = FastTwoSum(s_hi, t5_hi);
+  // r_lo is scaled by d/dr sin(r) = cos(r) ~= c_hi and d/dr cos(r) = -sin(r)
+  // ~= -s2_hi; the dropped higher-order terms are below 2^-60 relative.
+  T sin_lo = Fmuladd(r5 * r2, EvalF64SinTailPoly(r2),
+                     Fmuladd(r_lo, c_hi, (s_lo + s2_lo) + (t3_lo + t5_lo)));
+  T sin_val = s2_hi + sin_lo;
+
+  T cos_lo = Fmuladd(r4 * r2, EvalF64CosTailPoly(r2),
+                     Fmuladd(-s2_hi, r_lo, (c_lo + w_lo) + c4_lo));
+  T cos_val = c_hi + cos_lo;
+
+  return ApplyQuadrant<kIsCos>(x, ax, q_raw, sin_val, cos_val);
+}
+
+// f32 sin/cos evaluated entirely in double with a single rounding to float at
+// the end. An exhaustive scan shows |x mod pi/2| >= 2^-27.8 for every float in
+// [pi/4, 2^28), and both reductions below give r with relative error <= 2^-53
+// with or without FMA. With the polynomial errors above, the double result is
+// within 2^-19 float ULP of the exact value.
+//
+// When every lane is below 2^12, n * kF32SmallPiOver2A is exact and so is
+// `ax - n * A`: both are multiples of 2^-40 and the difference is below 1.
+// n * B is below 2^-31.1 and rounds by < 2^-84, and the split is within 2^-96
+// of pi/2.
+//
+// Otherwise the products of n with kF32PiOver2A and kF32PiOver2B are exact, and
+// so are t and u: ax is a multiple of 2^-24 (ax > 0.5 unless n = 0), n * A of
+// 2^-24 and n * B of 2^-49, and |u| < 1. Only n * kPiOver2Mid (< 2^-26.5,
+// rounding error < 2^-80) and the last subtraction round, and the split is
+// within 2^-109 of pi/2 (< 2^-81 after scaling by n).
+template <typename T, bool kIsCos>
+inline T SinCosF32(T x) {
+  constexpr size_t kLanes = NumLanes<T>;
+  using D =
+      std::conditional_t<kLanes == 1, double,
+                         double
+                         __attribute__((vector_size(kLanes * sizeof(double))))>;
+  using Q = std::conditional_t<kLanes == 1, uint64_t,
+                               uint64_t __attribute__((
+                                   vector_size(kLanes * sizeof(uint64_t))))>;
+
+  D xd = ConvertLanes<D>(x);
+  D ax = ConvertLanes<D>(BitwiseAbs(x));
+  D j = Fmuladd(ax, Splat<D>(kTwoOverPi), Splat<D>(kMagic));
+  D n = j - Splat<D>(kMagic);
+  Q q_raw = __builtin_bit_cast(Q, j) & Splat<Q>(3);
+  D r = Fmuladd(-n, Splat<D>(kF32SmallPiOver2B),
+                Fmuladd(-n, Splat<D>(kF32SmallPiOver2A), ax));
+  // NaN lanes compare false and stay on the short path; their result is NaN
+  // either way.
+  if (__builtin_expect(AnyLane(ax >= Splat<D>(kF32SmallThreshold)), 0)) {
+    D t = Fmuladd(-n, Splat<D>(kF32PiOver2A), ax);
+    D u = Fmuladd(-n, Splat<D>(kF32PiOver2B), t);
+    r = Fmuladd(-n, Splat<D>(kPiOver2Mid), u);
+    auto need_cold =
+        (ax >= Splat<D>(kF32FastThreshold)) & (ax < Splat<D>(kInf));
+    if (__builtin_expect(AnyLane(need_cold), 0)) {
+      auto cold =
+          PayneHanekColdLoop(ax, kF32FastThreshold,
+                             ReducedArgument<D, Q>{q_raw, r, Splat<D>(0.0)});
+      q_raw = cold.q_raw;
+      r = cold.r_hi + cold.r_lo;
+    }
+  }
+
+  D r2 = r * r;
+  D sin_val = Fmuladd(r2 * r, EvalF32SinPoly(r2), r);
+  D cos_val = Fmuladd(r2, EvalF32CosPoly(r2), Splat<D>(1.0));
+  return ConvertLanes<T>(
+      ApplyQuadrant<kIsCos>(xd, ax, q_raw, sin_val, cos_val));
+}
+
+}  // namespace
+
+float sin_f32(float x) { return SinCosF32<float, false>(x); }
+Vec4f sin_v4f32(Vec4f x) { return SinCosF32<Vec4f, false>(x); }
+Vec8f sin_v8f32(Vec8f x) { return SinCosF32<Vec8f, false>(x); }
+Vec16f sin_v16f32(Vec16f x) { return SinCosF32<Vec16f, false>(x); }
+double sin_f64(double x) { return SinCosF64<double, false>(x); }
+Vec2d sin_v2f64(Vec2d x) { return SinCosF64<Vec2d, false>(x); }
+Vec4d sin_v4f64(Vec4d x) { return SinCosF64<Vec4d, false>(x); }
+Vec8d sin_v8f64(Vec8d x) { return SinCosF64<Vec8d, false>(x); }
+
+float cos_f32(float x) { return SinCosF32<float, true>(x); }
+Vec4f cos_v4f32(Vec4f x) { return SinCosF32<Vec4f, true>(x); }
+Vec8f cos_v8f32(Vec8f x) { return SinCosF32<Vec8f, true>(x); }
+Vec16f cos_v16f32(Vec16f x) { return SinCosF32<Vec16f, true>(x); }
+double cos_f64(double x) { return SinCosF64<double, true>(x); }
+Vec2d cos_v2f64(Vec2d x) { return SinCosF64<Vec2d, true>(x); }
+Vec4d cos_v4f64(Vec4d x) { return SinCosF64<Vec4d, true>(x); }
+Vec8d cos_v8f64(Vec8d x) { return SinCosF64<Vec8d, true>(x); }
+
+}  // namespace xla::codegen
+
+#endif  // defined(__has_attribute) && __has_attribute(ext_vector_type) &&
+        // defined(__has_builtin) && __has_builtin(__builtin_vectorelements)
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/custom_cpp_intrinsics.h b/third_party/xla/xla/codegen/intrinsic/cpp/custom_cpp_intrinsics.h
new file mode 100644
index 00000000000..d0fb2eef3c9
--- /dev/null
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/custom_cpp_intrinsics.h
@@ -0,0 +1,45 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_CODEGEN_INTRINSIC_CPP_CUSTOM_CPP_INTRINSICS_H_
+#define XLA_CODEGEN_INTRINSIC_CPP_CUSTOM_CPP_INTRINSICS_H_
+
+#include "xla/codegen/intrinsic/cpp/vector_ops.h"
+
+namespace xla::codegen {
+
+float sin_f32(float x) asm("xla.sin.f32");
+Vec4f sin_v4f32(Vec4f x) asm("xla.sin.v4f32");
+Vec8f sin_v8f32(Vec8f x) asm("xla.sin.v8f32");
+Vec16f sin_v16f32(Vec16f x) asm("xla.sin.v16f32");
+
+double sin_f64(double x) asm("xla.sin.f64");
+Vec2d sin_v2f64(Vec2d x) asm("xla.sin.v2f64");
+Vec4d sin_v4f64(Vec4d x) asm("xla.sin.v4f64");
+Vec8d sin_v8f64(Vec8d x) asm("xla.sin.v8f64");
+
+float cos_f32(float x) asm("xla.cos.f32");
+Vec4f cos_v4f32(Vec4f x) asm("xla.cos.v4f32");
+Vec8f cos_v8f32(Vec8f x) asm("xla.cos.v8f32");
+Vec16f cos_v16f32(Vec16f x) asm("xla.cos.v16f32");
+
+double cos_f64(double x) asm("xla.cos.f64");
+Vec2d cos_v2f64(Vec2d x) asm("xla.cos.v2f64");
+Vec4d cos_v4f64(Vec4d x) asm("xla.cos.v4f64");
+Vec8d cos_v8f64(Vec8d x) asm("xla.cos.v8f64");
+
+}  // namespace xla::codegen
+
+#endif  // XLA_CODEGEN_INTRINSIC_CPP_CUSTOM_CPP_INTRINSICS_H_
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/custom_cpp_intrinsics_test.cc b/third_party/xla/xla/codegen/intrinsic/cpp/custom_cpp_intrinsics_test.cc
new file mode 100644
index 00000000000..88edb100d32
--- /dev/null
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/custom_cpp_intrinsics_test.cc
@@ -0,0 +1,175 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/codegen/intrinsic/cpp/custom_cpp_intrinsics.h"
+
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <memory>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/base/casts.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Verifier.h"
+#include "llvm/Support/raw_ostream.h"
+#include "xla/codegen/intrinsic/cpp/cpp_gen_intrinsics.h"
+#include "xla/codegen/intrinsic/cpp/custom_cpp_intrinsics_ll.h"
+#include "xla/codegen/intrinsic/cpp/vector_ops.h"
+#include "xla/codegen/intrinsic/test_matchers.h"
+
+namespace xla::codegen {
+namespace {
+
+using ::testing::ContainsRegex;
+using ::testing::Not;
+using ::xla::codegen::intrinsic::NearUlps;
+
+constexpr int kTrigUlps = 1;
+
+template <typename VecT, typename Fn>
+void ExpectVectorNearTrig(Fn fn, const ElementType<VecT>* xs, size_t n,
+                          bool is_cos) {
+  using Scalar = ElementType<VecT>;
+  constexpr size_t kLanes = NumLanes<VecT>;
+  for (size_t i = 0; i < n; i += kLanes) {
+    VecT v = Splat<VecT>(0);
+    for (size_t lane = 0; lane < kLanes; ++lane) {
+      SetLane(v, lane, xs[i + lane]);
+    }
+    VecT out = fn(v);
+    for (size_t lane = 0; lane < kLanes; ++lane) {
+      Scalar x = xs[i + lane];
+      Scalar actual = GetLane(out, lane);
+      double ref = is_cos ? std::cos(static_cast<double>(x))
+                          : std::sin(static_cast<double>(x));
+      Scalar expected = static_cast<Scalar>(ref);
+      EXPECT_THAT(actual, NearUlps(expected, kTrigUlps)) << "x=" << x;
+      if (x == Scalar{0} && !is_cos) {
+        EXPECT_EQ(std::signbit(actual), std::signbit(x));
+      }
+    }
+  }
+}
+
+TEST(CustomCppIntrinsicsTest, SinAndCosF32WithinOneUlp) {
+  using Limits = std::numeric_limits<float>;
+  alignas(64) constexpr float kInputs[16] = {
+      0.0f,
+      -0.0f,
+      Limits::denorm_min(),
+      -Limits::denorm_min(),
+      0.5f,
+      -1.0f,
+      1.5707963f,
+      3.1415927f,
+      -10.0f,
+      0x1.0p+17f,
+      0x1.5p+25f,
+      -0x1.3p+40f,
+      0x1.f37c8ap+95f,
+      Limits::infinity(),
+      -Limits::infinity(),
+      Limits::quiet_NaN(),
+  };
+  ExpectVectorNearTrig<float>(sin_f32, kInputs, 16, /*is_cos=*/false);
+  ExpectVectorNearTrig<Vec4f>(sin_v4f32, kInputs, 16, /*is_cos=*/false);
+  ExpectVectorNearTrig<Vec8f>(sin_v8f32, kInputs, 16, /*is_cos=*/false);
+  ExpectVectorNearTrig<Vec16f>(sin_v16f32, kInputs, 16, /*is_cos=*/false);
+
+  ExpectVectorNearTrig<float>(cos_f32, kInputs, 16, /*is_cos=*/true);
+  ExpectVectorNearTrig<Vec4f>(cos_v4f32, kInputs, 16, /*is_cos=*/true);
+  ExpectVectorNearTrig<Vec8f>(cos_v8f32, kInputs, 16, /*is_cos=*/true);
+  ExpectVectorNearTrig<Vec16f>(cos_v16f32, kInputs, 16, /*is_cos=*/true);
+}
+
+// Inputs found by exhaustive search where a float-float implementation
+// returned a result 1 ULP away from the correctly rounded value.
+TEST(CustomCppIntrinsicsTest, SinAndCosF32CorrectlyRoundedAtKnownHardCases) {
+  auto f = [](uint32_t bits) { return absl::bit_cast<float>(bits); };
+  EXPECT_EQ(absl::bit_cast<uint32_t>(cos_f32(f(0x3f4cc5a7))), 0x3f326080);
+  EXPECT_EQ(absl::bit_cast<uint32_t>(cos_f32(f(0x4731b173))), 0x3f3010ac);
+  EXPECT_EQ(absl::bit_cast<uint32_t>(cos_f32(f(0x40140780))), 0xbf2d06a2);
+  EXPECT_EQ(absl::bit_cast<uint32_t>(sin_f32(f(0x3f524c80))), 0x3f3b6e7c);
+  Vec8f v = {f(0x3f4cc5a7), f(0x4731b173), f(0x40140780), f(0x3f524c80)};
+  Vec8f c = cos_v8f32(v);
+  Vec8f s = sin_v8f32(v);
+  EXPECT_EQ(absl::bit_cast<uint32_t>(c[0]), 0x3f326080);
+  EXPECT_EQ(absl::bit_cast<uint32_t>(c[1]), 0x3f3010ac);
+  EXPECT_EQ(absl::bit_cast<uint32_t>(c[2]), 0xbf2d06a2);
+  EXPECT_EQ(absl::bit_cast<uint32_t>(s[3]), 0x3f3b6e7c);
+}
+
+TEST(CustomCppIntrinsicsTest, SinAndCosF64WithinOneUlp) {
+  using Limits = std::numeric_limits<double>;
+  alignas(64) constexpr double kInputs[8] = {
+      0.0,  -0.0,      Limits::denorm_min(),  0.5,
+      -2.5, 0x1.0p+30, 0x1.45f306dc9c883p+50, Limits::infinity(),
+  };
+  ExpectVectorNearTrig<double>(sin_f64, kInputs, 8, /*is_cos=*/false);
+  ExpectVectorNearTrig<Vec2d>(sin_v2f64, kInputs, 8, /*is_cos=*/false);
+  ExpectVectorNearTrig<Vec4d>(sin_v4f64, kInputs, 8, /*is_cos=*/false);
+  ExpectVectorNearTrig<Vec8d>(sin_v8f64, kInputs, 8, /*is_cos=*/false);
+
+  ExpectVectorNearTrig<double>(cos_f64, kInputs, 8, /*is_cos=*/true);
+  ExpectVectorNearTrig<Vec2d>(cos_v2f64, kInputs, 8, /*is_cos=*/true);
+  ExpectVectorNearTrig<Vec4d>(cos_v4f64, kInputs, 8, /*is_cos=*/true);
+  ExpectVectorNearTrig<Vec8d>(cos_v8f64, kInputs, 8, /*is_cos=*/true);
+
+  // Kahan/Muller worst-case near a multiple of pi/2 (~61 bits of cancellation);
+  // compared against the exact rounded reference rather than std::cos because
+  // older glibc (< 2.28) std::cos is 8 ULPs off at this input.
+  constexpr double kWorstCase = 0x1.6ac5b262ca1ffp+849;
+  EXPECT_THAT(sin_f64(kWorstCase), NearUlps(1.0, kTrigUlps));
+  EXPECT_THAT(cos_f64(kWorstCase), NearUlps(-0x1.14ae72e6ba22fp-61, kTrigUlps));
+
+  // Below the 2^20 Payne-Hanek threshold: the double closest to a multiple of
+  // pi/2 (|r| = 2^-60.5), and one with |r| = 2^-52 where rounding n * pi/2 at
+  // 2^-104 already makes r off by 2^-52 relative.
+  EXPECT_THAT(cos_f64(0x1.6c6cbc45dc8dep+5),
+              NearUlps(-0x1.6d61b58c99c43p-61, kTrigUlps));
+  EXPECT_THAT(sin_f64(0x1.6c6cbc45dc8dep+5), NearUlps(1.0, kTrigUlps));
+  EXPECT_THAT(cos_f64(0x1.93c05c9ed3cbcp+18),
+              NearUlps(-0x1.065d73720c4f9p-52, kTrigUlps));
+}
+
+TEST(CustomCppIntrinsicsTest, EmbeddedBitcodeHasNoTargetIntrinsics) {
+  ASSERT_TRUE(AreCppGenIntrinsicsAvailable());
+  llvm::LLVMContext context;
+  std::unique_ptr<llvm::Module> module =
+      ParseEmbeddedBitcode(context, llvm_ir::kCustomCppIntrinsicsLlIr);
+  ASSERT_NE(module, nullptr);
+  EXPECT_FALSE(llvm::verifyModule(*module, &llvm::errs()));
+
+  std::string ir;
+  llvm::raw_string_ostream stream(ir);
+  module->print(stream, nullptr);
+  // `llvm.fma` would become a libcall per lane on targets without FMA.
+  EXPECT_THAT(ir, Not(ContainsRegex("llvm\\.fma\\.")));
+  EXPECT_THAT(ir, ContainsRegex("llvm\\.fmuladd\\.v2f64"));
+  EXPECT_THAT(ir, ContainsRegex("llvm\\.fmuladd\\.v4f64"));
+  EXPECT_THAT(ir, ContainsRegex("llvm\\.fmuladd\\.v8f64"));
+  EXPECT_THAT(ir, ContainsRegex("llvm\\.fmuladd\\.v16f64"));
+  EXPECT_THAT(ir, Not(ContainsRegex("llvm\\.x86")));
+  EXPECT_THAT(ir, Not(ContainsRegex("llvm\\.aarch64")));
+  EXPECT_THAT(ir, Not(ContainsRegex("__udivti3")));
+}
+
+}  // namespace
+}  // namespace xla::codegen
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/intrinsic_declarations.h b/third_party/xla/xla/codegen/intrinsic/cpp/intrinsic_declarations.h
index ac6abb08950..4ca66284dca 100644
--- a/third_party/xla/xla/codegen/intrinsic/cpp/intrinsic_declarations.h
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/intrinsic_declarations.h
@@ -84,7 +84,6 @@ class EigenAtan : public Intrinsic<EigenAtan> {
                              UnaryFunctionType(module, type));
   }
 };
-
 class Expm1 : public Intrinsic<Expm1> {
  public:
   static constexpr absl::string_view kName = "expm1";
@@ -111,6 +110,60 @@ class Expm1 : public Intrinsic<Expm1> {
                              UnaryFunctionType(module, type));
   }
 };
+
+class Sin : public Intrinsic<Sin> {
+ public:
+  static constexpr absl::string_view kName = "sin";
+
+  static std::vector<std::vector<Type>> SupportedVectorTypes(
+      absl::string_view features) {
+    if (!AreCppGenIntrinsicsAvailable()) {
+      return {};
+    }
+    std::vector<std::vector<Type>> types = {
+        {Type::S(xla::F32)}, {Type::V(xla::F32, 4)}, {Type::V(xla::F32, 8)},
+        {Type::S(xla::F64)}, {Type::V(xla::F64, 2)}, {Type::V(xla::F64, 4)},
+    };
+    if (absl::StrContains(features, "+avx512f")) {
+      types.push_back({Type::V(xla::F32, 16)});
+      types.push_back({Type::V(xla::F64, 8)});
+    }
+    return types;
+  }
+
+  static absl::StatusOr<llvm::Function*> CreateDefinition(
+      llvm::Module* module, const IntrinsicOptions& options, Type type) {
+    return GetCppGenFunction(module, Name(type),
+                             UnaryFunctionType(module, type));
+  }
+};
+
+class Cos : public Intrinsic<Cos> {
+ public:
+  static constexpr absl::string_view kName = "cos";
+
+  static std::vector<std::vector<Type>> SupportedVectorTypes(
+      absl::string_view features) {
+    if (!AreCppGenIntrinsicsAvailable()) {
+      return {};
+    }
+    std::vector<std::vector<Type>> types = {
+        {Type::S(xla::F32)}, {Type::V(xla::F32, 4)}, {Type::V(xla::F32, 8)},
+        {Type::S(xla::F64)}, {Type::V(xla::F64, 2)}, {Type::V(xla::F64, 4)},
+    };
+    if (absl::StrContains(features, "+avx512f")) {
+      types.push_back({Type::V(xla::F32, 16)});
+      types.push_back({Type::V(xla::F64, 8)});
+    }
+    return types;
+  }
+
+  static absl::StatusOr<llvm::Function*> CreateDefinition(
+      llvm::Module* module, const IntrinsicOptions& options, Type type) {
+    return GetCppGenFunction(module, Name(type),
+                             UnaryFunctionType(module, type));
+  }
+};
 }  // namespace xla::codegen::intrinsics
 
 #endif  // XLA_CODEGEN_INTRINSIC_CPP_INTRINSIC_DECLARATIONS_H_
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/vector_ops.h b/third_party/xla/xla/codegen/intrinsic/cpp/vector_ops.h
index 62cb244487c..b708bf8884b 100644
--- a/third_party/xla/xla/codegen/intrinsic/cpp/vector_ops.h
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/vector_ops.h
@@ -61,6 +61,20 @@ typedef int64_t Vec4l __attribute__((vector_size(32)));
 typedef int64_t Vec8l __attribute__((vector_size(64)));
 
 namespace internal {
+template <typename T, bool kIsArithmetic = std::is_arithmetic_v<T>>
+struct ElementTypeImpl {
+  using Vec = std::remove_cv_t<T>;
+  using type = std::remove_reference_t<decltype(Vec{}[0])>;
+};
+
+template <typename T>
+struct ElementTypeImpl<T, true> {
+  using type = std::remove_cv_t<T>;
+};
+
+template <typename T>
+using ElementType = typename ElementTypeImpl<T>::type;
+
 // Helper type to select the corresponding unsigned integer vector type.
 template <typename ScalarInt, size_t Width>
 struct MakeIntVec;
@@ -119,18 +133,31 @@ struct MakeSignedIntVec<int64_t, 8> {
   using type = Vec8l;
 };
 
-// This trait takes a float vector and provides its unsigned integer vector
+// This trait takes a float scalar or vector and provides its unsigned integer
 // partner.
 template <typename FloatVec>
 struct CorrespondingIntVector {
  private:
-  using ScalarFloat = decltype(FloatVec{}[0]);
+  using ScalarFloat = ElementType<FloatVec>;
   static constexpr size_t kWidth = sizeof(FloatVec) / sizeof(ScalarFloat);
   using ScalarInt =
       std::conditional_t<sizeof(ScalarFloat) == 4, uint32_t, uint64_t>;
 
  public:
   using type = typename MakeIntVec<ScalarInt, kWidth>::type;
+  using scalar_type = ScalarInt;
+};
+
+template <>
+struct CorrespondingIntVector<float> {
+  using type = uint32_t;
+  using scalar_type = uint32_t;
+};
+
+template <>
+struct CorrespondingIntVector<double> {
+  using type = uint64_t;
+  using scalar_type = uint64_t;
 };
 
 // This trait takes a float vector and provides its signed integer vector
@@ -146,8 +173,105 @@ struct CorrespondingSignedIntVector {
  public:
   using type = typename MakeSignedIntVec<ScalarInt, kWidth>::type;
 };
+
+template <typename T, typename S, size_t... Is>
+constexpr T SplatImpl(S val, std::index_sequence<Is...>) {
+  return T{(static_cast<void>(Is), static_cast<ElementType<T>>(val))...};
+}
 }  // namespace internal
 
+template <typename T>
+using ElementType = internal::ElementType<T>;
+
+template <typename T>
+inline constexpr size_t NumLanes = sizeof(T) / sizeof(ElementType<T>);
+
+template <typename T, typename S>
+constexpr T Splat(S val) {
+  if constexpr (std::is_arithmetic_v<T>) {
+    return static_cast<T>(val);
+  } else {
+    return internal::SplatImpl<T>(val, std::make_index_sequence<NumLanes<T>>{});
+  }
+}
+
+template <typename T>
+constexpr ElementType<T> GetLane(T v, size_t idx) {
+  if constexpr (std::is_arithmetic_v<T>) {
+    (void)idx;
+    return v;
+  } else {
+    return v[idx];
+  }
+}
+
+template <typename T>
+constexpr void SetLane(T& v, size_t idx, ElementType<T> val) {
+  if constexpr (std::is_arithmetic_v<T>) {
+    (void)idx;
+    v = val;
+  } else {
+    v[idx] = val;
+  }
+}
+
+template <typename MaskT>
+inline bool AnyLane(MaskT mask) {
+  if constexpr (std::is_arithmetic_v<MaskT>) {
+    return mask != 0;
+  } else {
+    return __builtin_reduce_or(mask) != 0;
+  }
+}
+
+// Emits `@llvm.fmuladd` (`a * b + c`), which targets without FMA lower to a
+// separate multiply and add instead of an `fma` libcall.
+template <typename T>
+inline T Fmuladd(T a, T b, T c) {
+#pragma clang fp contract(on)
+  return a * b + c;
+}
+
+// Error-free transformation `p + e = a * b` on targets with FMA. Without FMA
+// the error term evaluates to 0 and only `p`, the rounded product, is kept.
+// `p` is emitted as `@llvm.fmuladd(a, b, +0.0)` instead of a bare `fmul`
+// because `e` is only the rounding error of `p` if `p` is rounded on its own:
+// downstream `SetAllowContractOnFpArithmetic` or aggressive-FMA targets could
+// otherwise fuse it into an addition that consumes it.
+template <typename T>
+inline std::pair<T, T> TwoProd(T a, T b) {
+  T p = Fmuladd(a, b, Splat<T>(0.0));
+  T e = Fmuladd(a, b, -p);
+  return {p, e};
+}
+
+// Evaluates a polynomial c_0 * x^n + c_1 * x^(n-1) + ... + c_n using Horner's
+// method with `Fmuladd`, where `coeffs` are given in descending degree order.
+template <typename T, typename S, typename... Coeffs>
+constexpr T HornerPoly(T x, S c0, Coeffs... coeffs) {
+  T acc = Splat<T>(c0);
+  ((acc = Fmuladd(acc, x, Splat<T>(coeffs))), ...);
+  return acc;
+}
+
+// Dekker's error-free transformation `s + e = a + b` assuming `|a| >= |b|`.
+template <typename T>
+inline std::pair<T, T> FastTwoSum(T a, T b) {
+  T s = a + b;
+  T e = (a - s) + b;
+  return {s, e};
+}
+
+// Knuth's error-free transformation `s + e = a + b` without any magnitude
+// ordering precondition on `a` and `b`.
+template <typename T>
+inline std::pair<T, T> TwoSum(T a, T b) {
+  T s = a + b;
+  T v = s - a;
+  T e = (a - (s - v)) + (b - v);
+  return {s, e};
+}
+
 // ===--------------------------------------------------------------------===//
 // Eigen Array type mapping
 // ===--------------------------------------------------------------------===//
@@ -191,54 +315,25 @@ struct ArrayMap<Vec8d> {
   using type = Eigen::Array<double, 8, 1>;
 };
 
-namespace internal {
-template <typename T, typename S, size_t... Is>
-constexpr T SplatImpl(S val, std::index_sequence<Is...>) {
-  return T{(static_cast<void>(Is), static_cast<decltype(T{}[0])>(val))...};
-}
-}  // namespace internal
-
-template <typename T, typename S>
-constexpr T Splat(S val) {
-  if constexpr (std::is_fundamental_v<T>) {
-    return static_cast<T>(val);
-  } else {
-    constexpr size_t kWidth = sizeof(T) / sizeof(decltype(T{}[0]));
-    return internal::SplatImpl<T>(val, std::make_index_sequence<kWidth>{});
-  }
-}
-
-// Computes the absolute value of a vector using bitwise operations.
-// FloatVec: The floating-point vector type (e.g., Vec4f).
-// x: The input vector.
-// Returns a new vector containing the absolute value of each element in x.
+// Computes the absolute value of a scalar or vector using bitwise operations.
 template <typename FloatVec>
 FloatVec BitwiseAbs(FloatVec x) {
   using IntVec = typename internal::CorrespondingIntVector<FloatVec>::type;
-  // Get the underlying scalar integer type (e.g., uint32_t from Vec4u).
-  using ScalarInt = decltype(IntVec{}[0]);
-
-  // Create a mask to clear the sign bit (e.g., 0x7FFFFFFF for int).
-  // This is a vector where every element is the mask.
-  const IntVec abs_mask =
-      IntVec{~(static_cast<ScalarInt>(1) << (sizeof(ScalarInt) * 8 - 1))};
-
-  // Reinterpret float as int, apply the mask, and reinterpret back.
+  using ScalarInt =
+      typename internal::CorrespondingIntVector<FloatVec>::scalar_type;
+  const IntVec abs_mask = Splat<IntVec>(
+      ~(static_cast<ScalarInt>(1) << (sizeof(ScalarInt) * 8 - 1)));
   return __builtin_bit_cast(FloatVec, __builtin_bit_cast(IntVec, x) & abs_mask);
 }
 
-// Copies the sign of one vector to the value of another.
-// FloatVec: The floating-point vector type (e.g., Vec4f).
-// value: The vector providing the magnitude.
-// sign_source: The vector providing the sign.
-// Returns a new vector with the magnitude of `value` and the sign of
-// `sign_source`.
+// Copies the sign of one scalar or vector to the value of another.
 template <typename FloatVec>
 FloatVec BitwiseCopysign(FloatVec value, FloatVec sign_source) {
   using IntVec = typename internal::CorrespondingIntVector<FloatVec>::type;
-  using ScalarInt = decltype(IntVec{}[0]);
+  using ScalarInt =
+      typename internal::CorrespondingIntVector<FloatVec>::scalar_type;
   const IntVec sign_mask =
-      IntVec{static_cast<ScalarInt>(1) << (sizeof(ScalarInt) * 8 - 1)};
+      Splat<IntVec>(static_cast<ScalarInt>(1) << (sizeof(ScalarInt) * 8 - 1));
   FloatVec value_abs = BitwiseAbs<FloatVec>(value);
   IntVec sign_bits = __builtin_bit_cast(IntVec, sign_source) & sign_mask;
   return __builtin_bit_cast(FloatVec,
@@ -250,34 +345,6 @@ Vec Clamp(Vec x, Scalar min, Scalar max) {
   return x < min ? min : x > max ? max : x;
 }
 
-// Evaluates a polynomial c_0 * x^n + c_1 * x^(n-1) + ... + c_n using Horner's
-// method, where `coeffs` are given in descending degree order.
-template <typename T, typename S, typename... Coeffs>
-constexpr T HornerPoly(T x, S c0, Coeffs... coeffs) {
-  T acc = Splat<T>(c0);
-  ((acc = acc * x + Splat<T>(coeffs)), ...);
-  return acc;
-}
-
-// Emits `@llvm.fmuladd` to compute `a * b + c`. When the target CPU has
-// hardware FMA, this lowers to a single fused multiply-add instruction;
-// without hardware FMA, `@llvm.fmuladd` falls back to unfused multiply/add.
-template <typename VecT>
-inline VecT Fmuladd(VecT a, VecT b, VecT c) {
-#pragma clang fp contract(on)
-  return a * b + c;
-}
-
-// Computes the error-free transformation `s + e = a + b` assuming `|a| >= |b|`
-// (Dekker's FastTwoSum), returning `{s, e}` where `s = RN(a + b)` and `e` is
-// the rounding error.
-template <typename VecT>
-inline std::pair<VecT, VecT> FastTwoSum(VecT a, VecT b) {
-  VecT s = a + b;
-  VecT e = (a - s) + b;
-  return {s, e};
-}
-
 template <typename VecT>
 inline auto ToSignedInt(VecT x) {
   // Guard NaN lanes to 0 before float-to-int conversion (`fptosi` on NaN is UB
diff --git a/third_party/xla/xla/codegen/intrinsic_lib.cc b/third_party/xla/xla/codegen/intrinsic_lib.cc
index 9a61ed2f7c3..3d278f223ae 100644
--- a/third_party/xla/xla/codegen/intrinsic_lib.cc
+++ b/third_party/xla/xla/codegen/intrinsic_lib.cc
@@ -143,6 +143,10 @@ IntrinsicFunctionLib::IntrinsicFunctionLib(const IntrinsicOptions& options)
       std::make_unique<IntrinsicAdapter<intrinsics::Tanh>>());
   intrinsic_functions_.push_back(
       std::make_unique<IntrinsicAdapter<intrinsics::EigenAtan>>());
+  intrinsic_functions_.push_back(
+      std::make_unique<IntrinsicAdapter<intrinsics::Sin>>());
+  intrinsic_functions_.push_back(
+      std::make_unique<IntrinsicAdapter<intrinsics::Cos>>());
 }
 
 namespace {
diff --git a/third_party/xla/xla/codegen/intrinsic_lib_test.cc b/third_party/xla/xla/codegen/intrinsic_lib_test.cc
index b2242c4ec49..41ae9b5ce97 100644
--- a/third_party/xla/xla/codegen/intrinsic_lib_test.cc
+++ b/third_party/xla/xla/codegen/intrinsic_lib_test.cc
@@ -170,6 +170,53 @@ TEST(IntrinsicLibTest, Expm1Vectorizations) {
                            "xla.expm1.f64:xla.expm1.v4f64:4:_ZGV_LLVM_N4v",
                            "xla.expm1.f64:xla.expm1.v8f64:8:_ZGV_LLVM_N8v"));
 }
+
+TEST(IntrinsicLibTest, SinAndCosVectorizations) {
+  auto descs_with_prefix = [](absl::string_view features,
+                              absl::string_view prefix) {
+    IntrinsicOptions options;
+    options.features = features;
+    auto lib = IntrinsicFunctionLib(options);
+    std::vector<std::string> descs;
+    for (const auto& vec_desc : lib.Vectorizations()) {
+      if (vec_desc.getScalarFnName().starts_with(prefix)) {
+        descs.push_back(ToString(vec_desc));
+      }
+    }
+    return descs;
+  };
+
+  EXPECT_THAT(
+      descs_with_prefix("", "xla.sin"),
+      UnorderedElementsAre("xla.sin.f32:xla.sin.v4f32:4:_ZGV_LLVM_N4v",
+                           "xla.sin.f32:xla.sin.v8f32:8:_ZGV_LLVM_N8v",
+                           "xla.sin.f64:xla.sin.v2f64:2:_ZGV_LLVM_N2v",
+                           "xla.sin.f64:xla.sin.v4f64:4:_ZGV_LLVM_N4v"));
+  EXPECT_THAT(
+      descs_with_prefix("", "xla.cos"),
+      UnorderedElementsAre("xla.cos.f32:xla.cos.v4f32:4:_ZGV_LLVM_N4v",
+                           "xla.cos.f32:xla.cos.v8f32:8:_ZGV_LLVM_N8v",
+                           "xla.cos.f64:xla.cos.v2f64:2:_ZGV_LLVM_N2v",
+                           "xla.cos.f64:xla.cos.v4f64:4:_ZGV_LLVM_N4v"));
+
+  EXPECT_THAT(
+      descs_with_prefix("+avx512f", "xla.sin"),
+      UnorderedElementsAre("xla.sin.f32:xla.sin.v4f32:4:_ZGV_LLVM_N4v",
+                           "xla.sin.f32:xla.sin.v8f32:8:_ZGV_LLVM_N8v",
+                           "xla.sin.f32:xla.sin.v16f32:16:_ZGV_LLVM_N16v",
+                           "xla.sin.f64:xla.sin.v2f64:2:_ZGV_LLVM_N2v",
+                           "xla.sin.f64:xla.sin.v4f64:4:_ZGV_LLVM_N4v",
+                           "xla.sin.f64:xla.sin.v8f64:8:_ZGV_LLVM_N8v"));
+  EXPECT_THAT(
+      descs_with_prefix("+avx512f", "xla.cos"),
+      UnorderedElementsAre("xla.cos.f32:xla.cos.v4f32:4:_ZGV_LLVM_N4v",
+                           "xla.cos.f32:xla.cos.v8f32:8:_ZGV_LLVM_N8v",
+                           "xla.cos.f32:xla.cos.v16f32:16:_ZGV_LLVM_N16v",
+                           "xla.cos.f64:xla.cos.v2f64:2:_ZGV_LLVM_N2v",
+                           "xla.cos.f64:xla.cos.v4f64:4:_ZGV_LLVM_N4v",
+                           "xla.cos.f64:xla.cos.v8f64:8:_ZGV_LLVM_N8v"));
+}
+
 TEST(IntrinsicLibTest, CppGenIntrinsicLibraryPreservesNoInline) {
   llvm::LLVMContext context;
   llvm::Module dst_module("dst_module", context);

From 2c9330186cbaf4053bb60c9aac67a8bb00e508e4 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 15 Oct 2026 12:42:59 -0700
Subject: [PATCH 0900/1000] Fold f32, DRQ, and Weight-Only Conv2D,
 DepthwiseConv2D, TransposeConv, and FC ops.

Reverts d27685655244809e48def3318ca1d79dc7b88cfe

PiperOrigin-RevId: 1000323289
---
 tensorflow/compiler/mlir/lite/ir/tfl_ops.cc   | 864 +-----------------
 tensorflow/compiler/mlir/lite/ir/tfl_ops.td   |   4 -
 .../compiler/mlir/lite/tests/const-fold.mlir  | 443 +--------
 .../compiler/mlir/lite/tests/optimize.mlir    |  41 -
 .../mlir/lite/transforms/optimize_pass.cc     |  10 +-
 5 files changed, 36 insertions(+), 1326 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc b/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc
index 4b50a8af5cf..4cd59932c74 100644
--- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc
+++ b/tensorflow/compiler/mlir/lite/ir/tfl_ops.cc
@@ -2180,262 +2180,6 @@ LogicalResult FullyConnectedOp::verify() {
   return mlir::success();
 }
 
-namespace {
-
-bool ExtractDequantizedWeightsFromQConst(QConstOp qconst_op,
-                                         ShapedType weights_type,
-                                         std::vector<float>& weights_values) {
-  if (!qconst_op || !weights_type.hasStaticShape() ||
-      weights_type.getRank() < 1) {
-    return false;
-  }
-  auto quant_type =
-      mlir::dyn_cast<quant::QuantizedType>(weights_type.getElementType());
-  if (!quant_type || !quant_type.getExpressedType().isF32()) {
-    return false;
-  }
-
-  const unsigned bit_width = quant_type.getStorageTypeIntegralWidth();
-  if (bit_width != 8 && bit_width != 4) {
-    return false;
-  }
-  const bool is_signed = quant_type.isSigned();
-
-  const int64_t num_elements = weights_type.getNumElements();
-
-  int32_t quant_dim = -1;
-  ArrayRef<double> scales;
-  ArrayRef<int64_t> zero_points;
-  double single_scale = 0.0;
-  int64_t single_zp = 0;
-
-  if (auto uq_type = mlir::dyn_cast<quant::UniformQuantizedType>(quant_type)) {
-    single_scale = uq_type.getScale();
-    single_zp = uq_type.getZeroPoint();
-    scales = ArrayRef<double>(&single_scale, 1);
-    zero_points = ArrayRef<int64_t>(&single_zp, 1);
-  } else if (auto per_axis_type =
-                 mlir::dyn_cast<quant::UniformQuantizedPerAxisType>(
-                     quant_type)) {
-    quant_dim = per_axis_type.getQuantizedDimension();
-    if (quant_dim < 0 || quant_dim >= weights_type.getRank()) {
-      return false;
-    }
-    scales = per_axis_type.getScales();
-    zero_points = per_axis_type.getZeroPoints();
-    const int64_t expected_channels = weights_type.getDimSize(quant_dim);
-    if (static_cast<int64_t>(scales.size()) != expected_channels ||
-        static_cast<int64_t>(zero_points.size()) != expected_channels) {
-      return false;
-    }
-  } else {
-    return false;
-  }
-
-  int64_t dim_size = 1;
-  int64_t stride = 1;
-  if (quant_dim >= 0) {
-    dim_size = weights_type.getDimSize(quant_dim);
-    if (dim_size <= 0) return false;
-    for (int i = quant_dim + 1; i < weights_type.getRank(); ++i) {
-      const int64_t d = weights_type.getDimSize(i);
-      if (d <= 0) return false;
-      stride *= d;
-    }
-  }
-
-  auto dequantize_at = [&](int64_t flat_idx, int64_t q) -> float {
-    const int64_t channel =
-        (quant_dim >= 0) ? ((flat_idx / stride) % dim_size) : 0;
-    return static_cast<float>(q - zero_points[channel]) *
-           static_cast<float>(scales[channel]);
-  };
-
-  ElementsAttr value_attr = qconst_op.getValue();
-  if (!value_attr) return false;
-  auto value_shaped_type = mlir::dyn_cast<ShapedType>(value_attr.getType());
-  if (!value_shaped_type || !value_shaped_type.hasStaticShape() ||
-      value_shaped_type.getNumElements() != num_elements ||
-      !value_shaped_type.getElementType().isInteger(bit_width)) {
-    return false;
-  }
-
-  weights_values.resize(num_elements);
-  if (auto dense_attr = mlir::dyn_cast<DenseElementsAttr>(value_attr)) {
-    int64_t idx = 0;
-    for (const APInt& val : dense_attr.getValues<APInt>()) {
-      const int64_t q = is_signed ? val.getSExtValue() : val.getZExtValue();
-      weights_values[idx] = dequantize_at(idx, q);
-      ++idx;
-    }
-    return idx == num_elements;
-  }
-
-  if (auto res_attr = mlir::dyn_cast<DenseResourceElementsAttr>(value_attr)) {
-    AsmResourceBlob* blob = res_attr.getRawHandle().getBlob();
-    if (!blob) return false;
-    ArrayRef<char> raw_data = blob->getData();
-    if (bit_width == 8) {
-      if (static_cast<int64_t>(raw_data.size()) != num_elements) {
-        return false;
-      }
-      for (int64_t idx = 0; idx < num_elements; ++idx) {
-        const int64_t q =
-            is_signed
-                ? static_cast<int64_t>(static_cast<int8_t>(raw_data[idx]))
-                : static_cast<int64_t>(static_cast<uint8_t>(raw_data[idx]));
-        weights_values[idx] = dequantize_at(idx, q);
-      }
-      return true;
-    }
-    if (bit_width == 4) {
-      const int64_t packed_size = (num_elements + 1) / 2;
-      if (static_cast<int64_t>(raw_data.size()) == num_elements) {
-        for (int64_t idx = 0; idx < num_elements; ++idx) {
-          const uint8_t byte = static_cast<uint8_t>(raw_data[idx]) & 0x0F;
-          const int64_t q =
-              is_signed
-                  ? static_cast<int64_t>(static_cast<int8_t>(byte << 4) >> 4)
-                  : static_cast<int64_t>(byte);
-          weights_values[idx] = dequantize_at(idx, q);
-        }
-        return true;
-      }
-      if (static_cast<int64_t>(raw_data.size()) == packed_size) {
-        for (int64_t idx = 0; idx < num_elements; ++idx) {
-          const uint8_t byte = static_cast<uint8_t>(raw_data[idx / 2]);
-          const uint8_t nibble = (byte >> ((idx % 2) * 4)) & 0x0F;
-          const int64_t q =
-              is_signed
-                  ? static_cast<int64_t>(static_cast<int8_t>(nibble << 4) >> 4)
-                  : static_cast<int64_t>(nibble);
-          weights_values[idx] = dequantize_at(idx, q);
-        }
-        return true;
-      }
-    }
-  }
-  return false;
-}
-
-class FloatElementsAccessor {
- public:
-  static std::optional<FloatElementsAccessor> Create(ElementsAttr attr) {
-    if (!attr) return std::nullopt;
-    auto shaped_type = mlir::dyn_cast<ShapedType>(attr.getType());
-    if (!shaped_type || !shaped_type.hasStaticShape() ||
-        !shaped_type.getElementType().isF32()) {
-      return std::nullopt;
-    }
-    FloatElementsAccessor accessor;
-    if (auto dense_attr = mlir::dyn_cast<DenseElementsAttr>(attr)) {
-      accessor.dense_values_ = dense_attr.getValues<float>();
-      return accessor;
-    }
-    if (auto res_attr = mlir::dyn_cast<DenseResourceElementsAttr>(attr)) {
-      ArrayRef<float> data = GetValues<float>(res_attr);
-      if (static_cast<int64_t>(data.size()) != shaped_type.getNumElements()) {
-        return std::nullopt;
-      }
-      accessor.resource_values_ = data;
-      return accessor;
-    }
-    return std::nullopt;
-  }
-
-  static std::optional<FloatElementsAccessor> CreateFromQConst(
-      QConstOp qconst_op, ShapedType weights_type) {
-    FloatElementsAccessor accessor;
-    if (!ExtractDequantizedWeightsFromQConst(qconst_op, weights_type,
-                                             accessor.owned_values_)) {
-      return std::nullopt;
-    }
-    return accessor;
-  }
-
-  float operator[](int64_t index) const {
-    if (dense_values_.has_value()) {
-      return (*dense_values_)[index];
-    }
-    if (!owned_values_.empty()) {
-      return owned_values_[index];
-    }
-    return resource_values_[index];
-  }
-
- private:
-  std::optional<mlir::detail::ElementsAttrRange<
-      DenseElementsAttr::ElementIterator<float>>>
-      dense_values_;
-  ArrayRef<float> resource_values_;
-  std::vector<float> owned_values_;
-};
-
-QConstOp GetWeightsQConstOp(Value weights) {
-  if (auto qconst_op = weights.getDefiningOp<QConstOp>()) {
-    return qconst_op;
-  }
-  if (auto dequant_op = weights.getDefiningOp<DequantizeOp>()) {
-    return dequant_op.getInput().getDefiningOp<QConstOp>();
-  }
-  return nullptr;
-}
-
-std::optional<FloatElementsAccessor> CreateWeightsAccessor(
-    ElementsAttr weights_tensor, QConstOp weights_qconst_op,
-    ShapedType weights_type) {
-  if (weights_tensor && weights_type.getElementType().isF32()) {
-    return FloatElementsAccessor::Create(weights_tensor);
-  }
-  if (weights_qconst_op &&
-      mlir::isa<quant::QuantizedType>(weights_type.getElementType())) {
-    return FloatElementsAccessor::CreateFromQConst(weights_qconst_op,
-                                                   weights_type);
-  }
-  return std::nullopt;
-}
-
-void NeumaierAdd(float addend, float& sum, float& compensation) {
-  const float new_sum = sum + addend;
-  // DO NOT enable -funsafe-math-optimizations here.
-  // There is a test detecting unsafe optimizations.
-  // Unsafe math optimizations can reorder float formulas, and set the
-  // compensation to constant 0. The formula must be evaluated as written
-  // for the algorithm to work.
-  // (Note: -ffast-math is a superset of -funsafe-math-optimizations.)
-  if (std::abs(sum) >= std::abs(addend)) {
-    compensation += (sum - new_sum) + addend;
-  } else {
-    compensation += (addend - new_sum) + sum;
-  }
-  sum = new_sum;
-}
-
-// Maximum number of multiply-accumulate operations allowed during constant
-// folding of FullyConnected and Conv ops, to avoid excessive compile times.
-constexpr int64_t kMaxFoldMultiplyAdds = 100'000'000;
-
-bool IsWithinFoldLimits(llvm::ArrayRef<int64_t> loop_bounds,
-                        ShapedType input_type, ShapedType weights_type,
-                        ShapedType output_type) {
-  int64_t total_macs = 1;
-  for (int64_t bound : loop_bounds) {
-    if (bound <= 0 || total_macs > kMaxFoldMultiplyAdds / bound) {
-      return false;
-    }
-    total_macs *= bound;
-  }
-  // Avoid folding when the output constant has more elements than the input
-  // and weight constants combined, which would inflate model size.
-  if (output_type.getNumElements() >
-      input_type.getNumElements() + weights_type.getNumElements()) {
-    return false;
-  }
-  return true;
-}
-
-}  // namespace
-
 LogicalResult FullyConnectedOp::fold(FoldAdaptor adaptor,
                                      SmallVectorImpl<OpFoldResult>& results) {
   assert(adaptor.getOperands().size() == 3);
@@ -2444,7 +2188,6 @@ LogicalResult FullyConnectedOp::fold(FoldAdaptor adaptor,
   // besides the default.
   if (getFusedActivationFunction() != "NONE") return failure();
   if (getWeightsFormat() != "DEFAULT") return failure();
-  if (getNumResults() != 1) return failure();
 
   // Bias tensor is optional.
   const bool has_bias =
@@ -2452,32 +2195,32 @@ LogicalResult FullyConnectedOp::fold(FoldAdaptor adaptor,
 
   // Get the tensors.
   auto operands = adaptor.getOperands();
-  ElementsAttr input_tensor = dyn_cast_or_null<ElementsAttr>(operands[0]);
-  ElementsAttr weights_tensor = dyn_cast_or_null<ElementsAttr>(operands[1]);
-  QConstOp weights_qconst_op = GetWeightsQConstOp(getFilter());
-  ElementsAttr bias_tensor;
+  DenseElementsAttr input_tensor =
+      dyn_cast_or_null<DenseElementsAttr>(operands[0]);
+  DenseElementsAttr weights_tensor =
+      dyn_cast_or_null<DenseElementsAttr>(operands[1]);
+  DenseElementsAttr bias_tensor;
 
-  if (!input_tensor || (!weights_tensor && !weights_qconst_op)) {
+  if (!input_tensor || !weights_tensor) {
     return failure();
   }
 
   if (has_bias) {
-    bias_tensor = dyn_cast_or_null<ElementsAttr>(operands[2]);
+    bias_tensor = dyn_cast_or_null<DenseElementsAttr>(operands[2]);
     if (!bias_tensor) return failure();
   }
 
   // Get the tensor types.
   const auto input_type = mlir::cast<ShapedType>(input_tensor.getType());
-  const auto weights_type = mlir::cast<ShapedType>(
-      weights_qconst_op ? weights_qconst_op.getOutput().getType()
-                        : getFilter().getType());
+  const auto weights_type = mlir::cast<ShapedType>(weights_tensor.getType());
   const auto bias_type =
       has_bias ? mlir::cast<ShapedType>(bias_tensor.getType()) : ShapedType{};
 
   const auto output_type = mlir::cast<ShapedType>(getType(0));
 
-  // Folding only implemented for float input/bias/output tensors.
+  // Folding only implemented for float tensors.
   if (!input_type.getElementType().isF32() ||
+      !weights_type.getElementType().isF32() ||
       !output_type.getElementType().isF32() ||
       (has_bias && !bias_type.getElementType().isF32())) {
     return failure();
@@ -2485,7 +2228,6 @@ LogicalResult FullyConnectedOp::fold(FoldAdaptor adaptor,
 
   // Folding only implemented for static shapes
   if (!input_type.hasStaticShape() || !weights_type.hasStaticShape() ||
-      !output_type.hasStaticShape() ||
       (has_bias && !bias_type.hasStaticShape())) {
     return failure();
   }
@@ -2496,10 +2238,6 @@ LogicalResult FullyConnectedOp::fold(FoldAdaptor adaptor,
 
   const int64_t in_dim = weights_type.getDimSize(1);
   const int64_t out_dim = weights_type.getDimSize(0);
-  if (in_dim <= 0 || out_dim <= 0 ||
-      input_type.getNumElements() % in_dim != 0) {
-    return failure();
-  }
 
   if (has_bias) {
     if (bias_type.getRank() > 2 ||
@@ -2515,37 +2253,36 @@ LogicalResult FullyConnectedOp::fold(FoldAdaptor adaptor,
     return failure();
   }
 
-  if (!IsWithinFoldLimits({batch_size, out_dim, in_dim}, input_type,
-                          weights_type, output_type)) {
-    return failure();
-  }
-
-  auto input_values_range = FloatElementsAccessor::Create(input_tensor);
-  if (!input_values_range.has_value()) return failure();
-
-  auto weights_values_range =
-      CreateWeightsAccessor(weights_tensor, weights_qconst_op, weights_type);
-  if (!weights_values_range.has_value()) return failure();
-
-  std::optional<FloatElementsAccessor> bias_values_range;
-  if (has_bias) {
-    bias_values_range = FloatElementsAccessor::Create(bias_tensor);
-    if (!bias_values_range.has_value()) return failure();
-  }
+  auto input_values_range = input_tensor.getValues<float>();
+  auto weights_values_range = weights_tensor.getValues<float>();
+  std::optional<decltype(input_values_range)> bias_values_range;
+  if (has_bias) bias_values_range = bias_tensor.getValues<float>();
 
   // Do the actual folding, one output at a time.
   std::vector<float> result_values;
   result_values.reserve(batch_size * out_dim);
 
-  for (int64_t b = 0; b < batch_size; ++b) {
-    for (int64_t o = 0; o < out_dim; ++o) {
+  for (int b = 0; b < batch_size; ++b) {
+    for (int o = 0; o < out_dim; ++o) {
       // Dot product with Kahan/Neumaier summation to minimize numeric errors.
       float sum = has_bias ? (*bias_values_range)[o] : 0.0f;
       float compensation = 0.0f;
-      for (int64_t i = 0; i < in_dim; ++i) {
-        const float addend = (*input_values_range)[b * in_dim + i] *
-                             (*weights_values_range)[o * in_dim + i];
-        NeumaierAdd(addend, sum, compensation);
+      for (int i = 0; i < in_dim; ++i) {
+        const float addend = input_values_range[b * in_dim + i] *
+                             weights_values_range[o * in_dim + i];
+        const float new_sum = sum + addend;
+        // DO NOT enable -funsafe-math-optimizations here.
+        // There is a test detecting unsafe optimizations.
+        // Unsafe math optimizations can reorder float formulas, and set the
+        // compensation to constant 0. The formula must be evaluated as written
+        // for the algorithm to work.
+        // (Note: -ffast-math is a superset of -funsafe-math-optimizations.)
+        if (std::abs(sum) >= std::abs(addend)) {
+          compensation += (sum - new_sum) + addend;
+        } else {
+          compensation += (addend - new_sum) + sum;
+        }
+        sum = new_sum;
       }
       result_values.push_back(sum + compensation);
     }
@@ -2695,183 +2432,6 @@ int64_t Conv2DOp::GetArithmeticCount(Operation* op) {
   return -1;
 }
 
-OpFoldResult Conv2DOp::fold(FoldAdaptor adaptor) {
-  assert(adaptor.getOperands().size() >= 3);
-
-  if (getFusedActivationFunction() != "NONE") return {};
-
-  const bool has_bias =
-      !(!getBias() || mlir::isa<NoneType>(getBias().getType()));
-
-  auto operands = adaptor.getOperands();
-  ElementsAttr input_tensor = dyn_cast_or_null<ElementsAttr>(operands[0]);
-  ElementsAttr weights_tensor = dyn_cast_or_null<ElementsAttr>(operands[1]);
-  QConstOp weights_qconst_op = GetWeightsQConstOp(getFilter());
-
-  if (!input_tensor || (!weights_tensor && !weights_qconst_op)) {
-    return {};
-  }
-
-  ElementsAttr bias_tensor;
-  if (has_bias) {
-    bias_tensor = dyn_cast_or_null<ElementsAttr>(operands[2]);
-    if (!bias_tensor) return {};
-  }
-
-  const auto input_type = mlir::dyn_cast<ShapedType>(input_tensor.getType());
-  const auto weights_type = mlir::dyn_cast<ShapedType>(
-      weights_qconst_op ? weights_qconst_op.getOutput().getType()
-                        : getFilter().getType());
-  const auto bias_type = has_bias
-                             ? mlir::dyn_cast<ShapedType>(bias_tensor.getType())
-                             : ShapedType{};
-  const auto output_type = mlir::dyn_cast<ShapedType>(getType());
-
-  if (!input_type || !weights_type || !output_type ||
-      (has_bias && !bias_type)) {
-    return {};
-  }
-
-  if (!input_type.getElementType().isF32() ||
-      !output_type.getElementType().isF32() ||
-      (has_bias && !bias_type.getElementType().isF32())) {
-    return {};
-  }
-
-  if (!input_type.hasStaticShape() || !weights_type.hasStaticShape() ||
-      !output_type.hasStaticShape() ||
-      (has_bias && !bias_type.hasStaticShape())) {
-    return {};
-  }
-
-  if (input_type.getRank() != 4 || weights_type.getRank() != 4 ||
-      output_type.getRank() != 4) {
-    return {};
-  }
-
-  const int64_t batches = input_type.getDimSize(0);
-  const int64_t input_h = input_type.getDimSize(1);
-  const int64_t input_w = input_type.getDimSize(2);
-  const int64_t input_c = input_type.getDimSize(3);
-
-  const int64_t output_depth = weights_type.getDimSize(0);
-  const int64_t filter_h = weights_type.getDimSize(1);
-  const int64_t filter_w = weights_type.getDimSize(2);
-  const int64_t filter_input_c = weights_type.getDimSize(3);
-
-  if (batches <= 0 || input_h <= 0 || input_w <= 0 || input_c <= 0 ||
-      output_depth <= 0 || filter_h <= 0 || filter_w <= 0 ||
-      filter_input_c <= 0) {
-    return {};
-  }
-
-  if (input_c % filter_input_c != 0) return {};
-  const int64_t groups = input_c / filter_input_c;
-  if (output_depth % groups != 0) return {};
-  const int64_t filters_per_group = output_depth / groups;
-
-  const int64_t stride_h = getStrideH();
-  const int64_t stride_w = getStrideW();
-  const int64_t dilation_h = getDilationHFactor();
-  const int64_t dilation_w = getDilationWFactor();
-  if (stride_h <= 0 || stride_w <= 0 || dilation_h <= 0 || dilation_w <= 0) {
-    return {};
-  }
-
-  tensorflow::Padding padding_enum;
-  if (!GetPaddingFromString(getPadding().str(), &padding_enum).ok()) {
-    return {};
-  }
-
-  int64_t out_h = 0, pad_top = 0, pad_bottom = 0;
-  if (!tensorflow::GetWindowedOutputSizeVerbose(input_h, filter_h, dilation_h,
-                                                stride_h, padding_enum, &out_h,
-                                                &pad_top, &pad_bottom)
-           .ok()) {
-    return {};
-  }
-
-  int64_t out_w = 0, pad_left = 0, pad_right = 0;
-  if (!tensorflow::GetWindowedOutputSizeVerbose(input_w, filter_w, dilation_w,
-                                                stride_w, padding_enum, &out_w,
-                                                &pad_left, &pad_right)
-           .ok()) {
-    return {};
-  }
-
-  if (output_type.getDimSize(0) != batches ||
-      output_type.getDimSize(1) != out_h ||
-      output_type.getDimSize(2) != out_w ||
-      output_type.getDimSize(3) != output_depth) {
-    return {};
-  }
-
-  if (has_bias) {
-    if (bias_type.getNumElements() != output_depth) {
-      return {};
-    }
-  }
-
-  if (!IsWithinFoldLimits({batches, out_h, out_w, output_depth, filter_h,
-                           filter_w, filter_input_c},
-                          input_type, weights_type, output_type)) {
-    return {};
-  }
-
-  auto input_values_range = FloatElementsAccessor::Create(input_tensor);
-  if (!input_values_range.has_value()) return {};
-
-  auto weights_values_range =
-      CreateWeightsAccessor(weights_tensor, weights_qconst_op, weights_type);
-  if (!weights_values_range.has_value()) return {};
-
-  std::optional<FloatElementsAccessor> bias_values_range;
-  if (has_bias) {
-    bias_values_range = FloatElementsAccessor::Create(bias_tensor);
-    if (!bias_values_range.has_value()) return {};
-  }
-
-  const int64_t total_outputs = batches * out_h * out_w * output_depth;
-  std::vector<float> result_values;
-  result_values.reserve(total_outputs);
-
-  for (int64_t b = 0; b < batches; ++b) {
-    for (int64_t out_y = 0; out_y < out_h; ++out_y) {
-      const int64_t in_y_origin = out_y * stride_h - pad_top;
-      for (int64_t out_x = 0; out_x < out_w; ++out_x) {
-        const int64_t in_x_origin = out_x * stride_w - pad_left;
-        for (int64_t oc = 0; oc < output_depth; ++oc) {
-          const int64_t group = oc / filters_per_group;
-          float sum = has_bias ? (*bias_values_range)[oc] : 0.0f;
-          float compensation = 0.0f;
-          for (int64_t fy = 0; fy < filter_h; ++fy) {
-            const int64_t in_y = in_y_origin + dilation_h * fy;
-            if (in_y < 0 || in_y >= input_h) continue;
-            for (int64_t fx = 0; fx < filter_w; ++fx) {
-              const int64_t in_x = in_x_origin + dilation_w * fx;
-              if (in_x < 0 || in_x >= input_w) continue;
-              for (int64_t in_c = 0; in_c < filter_input_c; ++in_c) {
-                const int64_t ic = in_c + group * filter_input_c;
-                const int64_t in_idx =
-                    ((b * input_h + in_y) * input_w + in_x) * input_c + ic;
-                const int64_t filter_idx =
-                    ((oc * filter_h + fy) * filter_w + fx) * filter_input_c +
-                    in_c;
-                const float addend = (*input_values_range)[in_idx] *
-                                     (*weights_values_range)[filter_idx];
-                NeumaierAdd(addend, sum, compensation);
-              }
-            }
-          }
-          result_values.push_back(sum + compensation);
-        }
-      }
-    }
-  }
-
-  return DenseElementsAttr::get(output_type, ArrayRef<float>(result_values));
-}
-
 //===----------------------------------------------------------------------===//
 // DepthwiseConv2DOp
 //===----------------------------------------------------------------------===//
@@ -2892,182 +2452,6 @@ int64_t DepthwiseConv2DOp::GetArithmeticCount(Operation* op) {
   return -1;
 }
 
-OpFoldResult DepthwiseConv2DOp::fold(FoldAdaptor adaptor) {
-  assert(adaptor.getOperands().size() >= 3);
-
-  if (getFusedActivationFunction() != "NONE") return {};
-
-  const bool has_bias =
-      !(!getBias() || mlir::isa<NoneType>(getBias().getType()));
-
-  auto operands = adaptor.getOperands();
-  ElementsAttr input_tensor = dyn_cast_or_null<ElementsAttr>(operands[0]);
-  ElementsAttr weights_tensor = dyn_cast_or_null<ElementsAttr>(operands[1]);
-  QConstOp weights_qconst_op = GetWeightsQConstOp(getFilter());
-
-  if (!input_tensor || (!weights_tensor && !weights_qconst_op)) {
-    return {};
-  }
-
-  ElementsAttr bias_tensor;
-  if (has_bias) {
-    bias_tensor = dyn_cast_or_null<ElementsAttr>(operands[2]);
-    if (!bias_tensor) return {};
-  }
-
-  const auto input_type = mlir::dyn_cast<ShapedType>(input_tensor.getType());
-  const auto weights_type = mlir::dyn_cast<ShapedType>(
-      weights_qconst_op ? weights_qconst_op.getOutput().getType()
-                        : getFilter().getType());
-  const auto bias_type = has_bias
-                             ? mlir::dyn_cast<ShapedType>(bias_tensor.getType())
-                             : ShapedType{};
-  const auto output_type = mlir::dyn_cast<ShapedType>(getType());
-
-  if (!input_type || !weights_type || !output_type ||
-      (has_bias && !bias_type)) {
-    return {};
-  }
-
-  if (!input_type.getElementType().isF32() ||
-      !output_type.getElementType().isF32() ||
-      (has_bias && !bias_type.getElementType().isF32())) {
-    return {};
-  }
-
-  if (!input_type.hasStaticShape() || !weights_type.hasStaticShape() ||
-      !output_type.hasStaticShape() ||
-      (has_bias && !bias_type.hasStaticShape())) {
-    return {};
-  }
-
-  if (input_type.getRank() != 4 || weights_type.getRank() != 4 ||
-      output_type.getRank() != 4) {
-    return {};
-  }
-
-  if (weights_type.getDimSize(0) != 1) {
-    return {};
-  }
-
-  const int64_t batches = input_type.getDimSize(0);
-  const int64_t input_h = input_type.getDimSize(1);
-  const int64_t input_w = input_type.getDimSize(2);
-  const int64_t input_c = input_type.getDimSize(3);
-
-  const int64_t filter_h = weights_type.getDimSize(1);
-  const int64_t filter_w = weights_type.getDimSize(2);
-  const int64_t output_depth = weights_type.getDimSize(3);
-  const int64_t depth_multiplier = getDepthMultiplier();
-
-  if (batches <= 0 || input_h <= 0 || input_w <= 0 || input_c <= 0 ||
-      filter_h <= 0 || filter_w <= 0 || output_depth <= 0 ||
-      depth_multiplier <= 0) {
-    return {};
-  }
-
-  if (output_depth != input_c * depth_multiplier) {
-    return {};
-  }
-
-  const int64_t stride_h = getStrideH();
-  const int64_t stride_w = getStrideW();
-  const int64_t dilation_h = getDilationHFactor();
-  const int64_t dilation_w = getDilationWFactor();
-  if (stride_h <= 0 || stride_w <= 0 || dilation_h <= 0 || dilation_w <= 0) {
-    return {};
-  }
-
-  tensorflow::Padding padding_enum;
-  if (!GetPaddingFromString(getPadding().str(), &padding_enum).ok()) {
-    return {};
-  }
-
-  int64_t out_h = 0, pad_top = 0, pad_bottom = 0;
-  if (!tensorflow::GetWindowedOutputSizeVerbose(input_h, filter_h, dilation_h,
-                                                stride_h, padding_enum, &out_h,
-                                                &pad_top, &pad_bottom)
-           .ok()) {
-    return {};
-  }
-
-  int64_t out_w = 0, pad_left = 0, pad_right = 0;
-  if (!tensorflow::GetWindowedOutputSizeVerbose(input_w, filter_w, dilation_w,
-                                                stride_w, padding_enum, &out_w,
-                                                &pad_left, &pad_right)
-           .ok()) {
-    return {};
-  }
-
-  if (output_type.getDimSize(0) != batches ||
-      output_type.getDimSize(1) != out_h ||
-      output_type.getDimSize(2) != out_w ||
-      output_type.getDimSize(3) != output_depth) {
-    return {};
-  }
-
-  if (has_bias) {
-    if (bias_type.getNumElements() != output_depth) {
-      return {};
-    }
-  }
-
-  if (!IsWithinFoldLimits(
-          {batches, out_h, out_w, output_depth, filter_h, filter_w}, input_type,
-          weights_type, output_type)) {
-    return {};
-  }
-
-  auto input_values_range = FloatElementsAccessor::Create(input_tensor);
-  if (!input_values_range.has_value()) return {};
-
-  auto weights_values_range =
-      CreateWeightsAccessor(weights_tensor, weights_qconst_op, weights_type);
-  if (!weights_values_range.has_value()) return {};
-
-  std::optional<FloatElementsAccessor> bias_values_range;
-  if (has_bias) {
-    bias_values_range = FloatElementsAccessor::Create(bias_tensor);
-    if (!bias_values_range.has_value()) return {};
-  }
-
-  const int64_t total_outputs = batches * out_h * out_w * output_depth;
-  std::vector<float> result_values;
-  result_values.reserve(total_outputs);
-
-  for (int64_t b = 0; b < batches; ++b) {
-    for (int64_t out_y = 0; out_y < out_h; ++out_y) {
-      const int64_t in_y_origin = out_y * stride_h - pad_top;
-      for (int64_t out_x = 0; out_x < out_w; ++out_x) {
-        const int64_t in_x_origin = out_x * stride_w - pad_left;
-        for (int64_t oc = 0; oc < output_depth; ++oc) {
-          const int64_t ic = oc / depth_multiplier;
-          float sum = has_bias ? (*bias_values_range)[oc] : 0.0f;
-          float compensation = 0.0f;
-          for (int64_t fy = 0; fy < filter_h; ++fy) {
-            const int64_t in_y = in_y_origin + dilation_h * fy;
-            if (in_y < 0 || in_y >= input_h) continue;
-            for (int64_t fx = 0; fx < filter_w; ++fx) {
-              const int64_t in_x = in_x_origin + dilation_w * fx;
-              if (in_x < 0 || in_x >= input_w) continue;
-              const int64_t in_idx =
-                  ((b * input_h + in_y) * input_w + in_x) * input_c + ic;
-              const int64_t filter_idx =
-                  (fy * filter_w + fx) * output_depth + oc;
-              const float addend = (*input_values_range)[in_idx] *
-                                   (*weights_values_range)[filter_idx];
-              NeumaierAdd(addend, sum, compensation);
-            }
-          }
-          result_values.push_back(sum + compensation);
-        }
-      }
-    }
-  }
-
-  return DenseElementsAttr::get(output_type, ArrayRef<float>(result_values));
-}
-
 //===----------------------------------------------------------------------===//
 // GatherOp
 //===----------------------------------------------------------------------===//
@@ -5787,190 +5171,6 @@ int64_t TransposeConvOp::GetArithmeticCount(Operation* op) {
   return count;
 }
 
-OpFoldResult TransposeConvOp::fold(FoldAdaptor adaptor) {
-  assert(adaptor.getOperands().size() >= 3);
-
-  if (getFusedActivationFunction() != "NONE") return {};
-
-  const bool has_bias =
-      !(!getBias() || mlir::isa<NoneType>(getBias().getType()));
-
-  auto operands = adaptor.getOperands();
-  ElementsAttr weights_tensor = dyn_cast_or_null<ElementsAttr>(operands[1]);
-  ElementsAttr input_tensor = dyn_cast_or_null<ElementsAttr>(operands[2]);
-  QConstOp weights_qconst_op = GetWeightsQConstOp(getWeights());
-
-  if (!input_tensor || (!weights_tensor && !weights_qconst_op)) {
-    return {};
-  }
-
-  ElementsAttr bias_tensor;
-  if (has_bias) {
-    if (operands.size() < 4) return {};
-    bias_tensor = dyn_cast_or_null<ElementsAttr>(operands[3]);
-    if (!bias_tensor) return {};
-  }
-
-  const auto input_type = mlir::dyn_cast<ShapedType>(input_tensor.getType());
-  const auto weights_type = mlir::dyn_cast<ShapedType>(
-      weights_qconst_op ? weights_qconst_op.getOutput().getType()
-                        : getWeights().getType());
-  const auto bias_type = has_bias
-                             ? mlir::dyn_cast<ShapedType>(bias_tensor.getType())
-                             : ShapedType{};
-  const auto output_type = mlir::dyn_cast<ShapedType>(getType());
-
-  if (!input_type || !weights_type || !output_type ||
-      (has_bias && !bias_type)) {
-    return {};
-  }
-
-  if (!input_type.getElementType().isF32() ||
-      !output_type.getElementType().isF32() ||
-      (has_bias && !bias_type.getElementType().isF32())) {
-    return {};
-  }
-
-  if (!input_type.hasStaticShape() || !weights_type.hasStaticShape() ||
-      !output_type.hasStaticShape() ||
-      (has_bias && !bias_type.hasStaticShape())) {
-    return {};
-  }
-
-  if (input_type.getRank() != 4 || weights_type.getRank() != 4 ||
-      output_type.getRank() != 4) {
-    return {};
-  }
-
-  const int64_t batches = input_type.getDimSize(0);
-  const int64_t input_h = input_type.getDimSize(1);
-  const int64_t input_w = input_type.getDimSize(2);
-  const int64_t input_c = input_type.getDimSize(3);
-
-  const int64_t output_depth = weights_type.getDimSize(0);
-  const int64_t filter_h = weights_type.getDimSize(1);
-  const int64_t filter_w = weights_type.getDimSize(2);
-  const int64_t filter_input_c = weights_type.getDimSize(3);
-
-  const int64_t out_batches = output_type.getDimSize(0);
-  const int64_t out_h = output_type.getDimSize(1);
-  const int64_t out_w = output_type.getDimSize(2);
-  const int64_t out_c = output_type.getDimSize(3);
-
-  if (batches <= 0 || input_h <= 0 || input_w <= 0 || input_c <= 0 ||
-      output_depth <= 0 || filter_h <= 0 || filter_w <= 0 ||
-      filter_input_c <= 0 || out_batches <= 0 || out_h <= 0 || out_w <= 0 ||
-      out_c <= 0) {
-    return {};
-  }
-
-  if (batches != out_batches || output_depth != out_c ||
-      input_c != filter_input_c) {
-    return {};
-  }
-
-  const int64_t stride_h = getStrideH();
-  const int64_t stride_w = getStrideW();
-  if (stride_h <= 0 || stride_w <= 0) {
-    return {};
-  }
-
-  tensorflow::Padding padding_enum;
-  if (!GetPaddingFromString(getPadding().str(), &padding_enum).ok()) {
-    return {};
-  }
-
-  int64_t computed_in_h = 0, pad_top = 0, pad_bottom = 0;
-  if (!tensorflow::GetWindowedOutputSizeVerbose(
-           out_h, filter_h, /*dilation_rate=*/1, stride_h, padding_enum,
-           &computed_in_h, &pad_top, &pad_bottom)
-           .ok()) {
-    return {};
-  }
-  if (computed_in_h != input_h) return {};
-
-  int64_t computed_in_w = 0, pad_left = 0, pad_right = 0;
-  if (!tensorflow::GetWindowedOutputSizeVerbose(
-           out_w, filter_w, /*dilation_rate=*/1, stride_w, padding_enum,
-           &computed_in_w, &pad_left, &pad_right)
-           .ok()) {
-    return {};
-  }
-  if (computed_in_w != input_w) return {};
-
-  if (has_bias) {
-    if (bias_type.getNumElements() != output_depth) {
-      return {};
-    }
-  }
-
-  if (!IsWithinFoldLimits({batches, input_h, input_w, input_c, filter_h,
-                           filter_w, output_depth},
-                          input_type, weights_type, output_type)) {
-    return {};
-  }
-
-  auto input_values_range = FloatElementsAccessor::Create(input_tensor);
-  if (!input_values_range.has_value()) return {};
-
-  auto weights_values_range =
-      CreateWeightsAccessor(weights_tensor, weights_qconst_op, weights_type);
-  if (!weights_values_range.has_value()) return {};
-
-  std::optional<FloatElementsAccessor> bias_values_range;
-  if (has_bias) {
-    bias_values_range = FloatElementsAccessor::Create(bias_tensor);
-    if (!bias_values_range.has_value()) return {};
-  }
-
-  const int64_t total_outputs = batches * out_h * out_w * output_depth;
-  std::vector<float> result_values(total_outputs, 0.0f);
-  std::vector<float> compensations(total_outputs, 0.0f);
-  if (has_bias) {
-    for (int64_t i = 0; i < total_outputs; ++i) {
-      result_values[i] = (*bias_values_range)[i % output_depth];
-    }
-  }
-
-  for (int64_t b = 0; b < batches; ++b) {
-    for (int64_t in_y = 0; in_y < input_h; ++in_y) {
-      const int64_t out_y_origin = in_y * stride_h - pad_top;
-      for (int64_t in_x = 0; in_x < input_w; ++in_x) {
-        const int64_t out_x_origin = in_x * stride_w - pad_left;
-        for (int64_t in_c = 0; in_c < input_c; ++in_c) {
-          const int64_t in_idx =
-              ((b * input_h + in_y) * input_w + in_x) * input_c + in_c;
-          const float in_val = (*input_values_range)[in_idx];
-          for (int64_t fy = 0; fy < filter_h; ++fy) {
-            const int64_t out_y = out_y_origin + fy;
-            if (out_y < 0 || out_y >= out_h) continue;
-            for (int64_t fx = 0; fx < filter_w; ++fx) {
-              const int64_t out_x = out_x_origin + fx;
-              if (out_x < 0 || out_x >= out_w) continue;
-              for (int64_t oc = 0; oc < output_depth; ++oc) {
-                const int64_t filter_idx =
-                    ((oc * filter_h + fy) * filter_w + fx) * input_c + in_c;
-                const int64_t out_idx =
-                    ((b * out_h + out_y) * out_w + out_x) * output_depth + oc;
-                const float addend =
-                    in_val * (*weights_values_range)[filter_idx];
-                NeumaierAdd(addend, result_values[out_idx],
-                            compensations[out_idx]);
-              }
-            }
-          }
-        }
-      }
-    }
-  }
-
-  for (int64_t i = 0; i < total_outputs; ++i) {
-    result_values[i] += compensations[i];
-  }
-
-  return DenseElementsAttr::get(output_type, ArrayRef<float>(result_values));
-}
-
 //===----------------------------------------------------------------------===//
 // StridedSliceOp
 //===----------------------------------------------------------------------===//
diff --git a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td
index 2d08b9af4bc..7a2778dc978 100644
--- a/tensorflow/compiler/mlir/lite/ir/tfl_ops.td
+++ b/tensorflow/compiler/mlir/lite/ir/tfl_ops.td
@@ -744,8 +744,6 @@ def TFL_TransposeConvOp: TFL_Op<"transpose_conv", [
 
   let hasVerifier = 1;
 
-  let hasFolder = 1;
-
   let extraClassDeclaration = [{
     // AffineQuantizedOpInterface:
     int GetChannelDimIndex() { return 0; }
@@ -1020,7 +1018,6 @@ def TFL_Conv2DOp : TFL_ConvOp<"conv_2d", "Convolution", 0,
   );
 
   let hasCanonicalizer = 1;
-  let hasFolder = 1;
 
   let extraClassDeclaration = [{
     // AffineQuantizedOpInterface:
@@ -1111,7 +1108,6 @@ def TFL_DepthwiseConv2DOp :
   );
 
   let hasCanonicalizer = 1;
-  let hasFolder = 1;
 
   let extraClassDeclaration = [{
     // AffineQuantizedOpInterface:
diff --git a/tensorflow/compiler/mlir/lite/tests/const-fold.mlir b/tensorflow/compiler/mlir/lite/tests/const-fold.mlir
index 4765f783c32..dcdda38f84b 100644
--- a/tensorflow/compiler/mlir/lite/tests/const-fold.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/const-fold.mlir
@@ -485,439 +485,6 @@ func.func @ConstantFoldFullyConnectedBatched() -> tensor<13x1536xf32> {
   // CHECK:  return %[[CST]]
 }
 
-// CHECK-LABEL: @ConstantFoldFullyConnectedDRQW8A32PerTensor
-func.func @ConstantFoldFullyConnectedDRQW8A32PerTensor() -> tensor<1x3xf32> {
-  %cst_input = arith.constant dense<[[2.0, 3.0]]> : tensor<1x2xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<3x2x!quant.uniform<i8:f32, 0.5:0>>, value = dense<[[2, 4], [-6, 8], [10, -12]]> : tensor<3x2xi8>}> : () -> tensor<3x2x!quant.uniform<i8:f32, 0.5:0>>
-  %cst_bias = arith.constant dense<[1.0, 2.0, 3.0]> : tensor<3xf32>
-
-  %0 = "tfl.fully_connected"(%cst_input, %cst_weights, %cst_bias) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<1x2xf32>, tensor<3x2x!quant.uniform<i8:f32, 0.5:0>>, tensor<3xf32>) -> tensor<1x3xf32>
-  func.return %0 : tensor<1x3xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[}}9.000000e+00, 8.000000e+00, -5.000000e+00{{\]\]}}> : tensor<1x3xf32>
-  // CHECK-NOT: fully_connected
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldFullyConnectedDRQW8A32PerAxis
-func.func @ConstantFoldFullyConnectedDRQW8A32PerAxis() -> tensor<1x3xf32> {
-  %cst_input = arith.constant dense<[[2.0, 3.0]]> : tensor<1x2xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<3x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0, 0.25}>>, value = dense<[[2, 4], [-1, 3], [8, -4]]> : tensor<3x2xi8>}> : () -> tensor<3x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0, 0.25}>>
-  %cst_bias = arith.constant dense<[1.0, 2.0, 3.0]> : tensor<3xf32>
-
-  %0 = "tfl.fully_connected"(%cst_input, %cst_weights, %cst_bias) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<1x2xf32>, tensor<3x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0, 0.25}>>, tensor<3xf32>) -> tensor<1x3xf32>
-  func.return %0 : tensor<1x3xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[}}9.000000e+00, 1.600000e+01, 4.000000e+00{{\]\]}}> : tensor<1x3xf32>
-  // CHECK-NOT: fully_connected
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldFullyConnectedDRQW4A32PerTensor
-func.func @ConstantFoldFullyConnectedDRQW4A32PerTensor() -> tensor<1x2xf32> {
-  %cst_input = arith.constant dense<[[2.0, 3.0]]> : tensor<1x2xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x!quant.uniform<i4:f32, 0.5:0>>, value = dense<[[2, -4], [6, -2]]> : tensor<2x2xi4>}> : () -> tensor<2x2x!quant.uniform<i4:f32, 0.5:0>>
-  %cst_bias = "tfl.no_value"() {value = unit} : () -> none
-
-  %0 = "tfl.fully_connected"(%cst_input, %cst_weights, %cst_bias) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<1x2xf32>, tensor<2x2x!quant.uniform<i4:f32, 0.5:0>>, none) -> tensor<1x2xf32>
-  func.return %0 : tensor<1x2xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[}}-4.000000e+00, 3.000000e+00{{\]\]}}> : tensor<1x2xf32>
-  // CHECK-NOT: fully_connected
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldFullyConnectedDRQW4A32PerAxis
-func.func @ConstantFoldFullyConnectedDRQW4A32PerAxis() -> tensor<1x2xf32> {
-  %cst_input = arith.constant dense<[[2.0, 3.0]]> : tensor<1x2xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>, value = dense<[[2, 4], [-3, 1]]> : tensor<2x2xi4>}> : () -> tensor<2x2x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.fully_connected"(%cst_input, %cst_weights, %cst_bias) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<1x2xf32>, tensor<2x2x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x2xf32>
-  func.return %0 : tensor<1x2xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[}}9.000000e+00, -4.000000e+00{{\]\]}}> : tensor<1x2xf32>
-  // CHECK-NOT: fully_connected
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldFullyConnectedWeightOnlyW8A32PerAxis
-func.func @ConstantFoldFullyConnectedWeightOnlyW8A32PerAxis() -> tensor<1x3xf32> {
-  %cst_input = arith.constant dense<[[2.0, 3.0]]> : tensor<1x2xf32>
-  %cst_qweights = "tfl.pseudo_qconst"() <{qtype = tensor<3x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0, 0.25}>>, value = dense<[[2, 4], [-1, 3], [8, -4]]> : tensor<3x2xi8>}> : () -> tensor<3x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0, 0.25}>>
-  %cst_weights = "tfl.dequantize"(%cst_qweights) : (tensor<3x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0, 0.25}>>) -> tensor<3x2xf32>
-  %cst_bias = arith.constant dense<[1.0, 2.0, 3.0]> : tensor<3xf32>
-
-  %0 = "tfl.fully_connected"(%cst_input, %cst_weights, %cst_bias) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<1x2xf32>, tensor<3x2xf32>, tensor<3xf32>) -> tensor<1x3xf32>
-  func.return %0 : tensor<1x3xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[}}9.000000e+00, 1.600000e+01, 4.000000e+00{{\]\]}}> : tensor<1x3xf32>
-  // CHECK-NOT: fully_connected
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldConv2DDRQW8A32PerAxis
-func.func @ConstantFoldConv2DDRQW8A32PerAxis() -> tensor<1x1x1x2xf32> {
-  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2], [4]], [[6], [8]]], [[[1], [-1]], [[2], [-2]]]]> : tensor<2x2x2x1xi8>}> : () -> tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.conv_2d"(%cst_input, %cst_weights, %cst_bias) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
-  func.return %0 : tensor<1x1x1x2xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}3.100000e+01, -4.000000e+00{{\]\]\]\]}}> : tensor<1x1x1x2xf32>
-  // CHECK-NOT: conv_2d
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldConv2DDRQW4A32PerAxis
-func.func @ConstantFoldConv2DDRQW4A32PerAxis() -> tensor<1x1x1x2xf32> {
-  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x2x1x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2], [4]], [[6], [-2]]], [[[1], [-1]], [[3], [-2]]]]> : tensor<2x2x2x1xi4>}> : () -> tensor<2x2x2x1x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>
-  %cst_bias = arith.constant dense<[5.0, -3.0]> : tensor<2xf32>
-
-  %0 = "tfl.conv_2d"(%cst_input, %cst_weights, %cst_bias) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<2x2x2x1x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
-  func.return %0 : tensor<1x1x1x2xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}1.500000e+01, -3.000000e+00{{\]\]\]\]}}> : tensor<1x1x1x2xf32>
-  // CHECK-NOT: conv_2d
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldConv2DWeightOnlyW8A32PerAxis
-func.func @ConstantFoldConv2DWeightOnlyW8A32PerAxis() -> tensor<1x1x1x2xf32> {
-  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
-  %cst_qweights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2], [4]], [[6], [8]]], [[[1], [-1]], [[2], [-2]]]]> : tensor<2x2x2x1xi8>}> : () -> tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>
-  %cst_weights = "tfl.dequantize"(%cst_qweights) : (tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>) -> tensor<2x2x2x1xf32>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.conv_2d"(%cst_input, %cst_weights, %cst_bias) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<2x2x2x1xf32>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
-  func.return %0 : tensor<1x1x1x2xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}3.100000e+01, -4.000000e+00{{\]\]\]\]}}> : tensor<1x1x1x2xf32>
-  // CHECK-NOT: conv_2d
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @NoConstantFoldConv2DFusedRelu
-func.func @NoConstantFoldConv2DFusedRelu() -> tensor<1x1x1x2xf32> {
-  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2], [4]], [[6], [8]]], [[[1], [-1]], [[2], [-2]]]]> : tensor<2x2x2x1xi8>}> : () -> tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.conv_2d"(%cst_input, %cst_weights, %cst_bias) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "RELU", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
-  func.return %0 : tensor<1x1x1x2xf32>
-
-  // CHECK: tfl.conv_2d
-}
-
-// CHECK-LABEL: @ConstantFoldDepthwiseConv2DDRQW8A32PerAxis
-func.func @ConstantFoldDepthwiseConv2DDRQW8A32PerAxis() -> tensor<1x1x1x2xf32> {
-  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>, value = dense<[[[[2, 1], [4, -1]], [[6, 2], [8, -2]]]]> : tensor<1x2x2x2xi8>}> : () -> tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.depthwise_conv_2d"(%cst_input, %cst_weights, %cst_bias) <{depth_multiplier = 2 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
-  func.return %0 : tensor<1x1x1x2xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}3.100000e+01, -4.000000e+00{{\]\]\]\]}}> : tensor<1x1x1x2xf32>
-  // CHECK-NOT: depthwise_conv_2d
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldDepthwiseConv2DDRQW4A32PerAxis
-func.func @ConstantFoldDepthwiseConv2DDRQW4A32PerAxis() -> tensor<1x1x1x2xf32> {
-  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<1x2x2x2x!quant.uniform<i4<-7:7>:f32:3, {0.5, 2.0}>>, value = dense<[[[[2, 1], [4, -1]], [[6, 3], [-2, -2]]]]> : tensor<1x2x2x2xi4>}> : () -> tensor<1x2x2x2x!quant.uniform<i4<-7:7>:f32:3, {0.5, 2.0}>>
-  %cst_bias = arith.constant dense<[5.0, -3.0]> : tensor<2xf32>
-
-  %0 = "tfl.depthwise_conv_2d"(%cst_input, %cst_weights, %cst_bias) <{depth_multiplier = 2 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<1x2x2x2x!quant.uniform<i4<-7:7>:f32:3, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
-  func.return %0 : tensor<1x1x1x2xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}1.500000e+01, -3.000000e+00{{\]\]\]\]}}> : tensor<1x1x1x2xf32>
-  // CHECK-NOT: depthwise_conv_2d
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldDepthwiseConv2DWeightOnlyW8A32PerAxis
-func.func @ConstantFoldDepthwiseConv2DWeightOnlyW8A32PerAxis() -> tensor<1x1x1x2xf32> {
-  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
-  %cst_qweights = "tfl.pseudo_qconst"() <{qtype = tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>, value = dense<[[[[2, 1], [4, -1]], [[6, 2], [8, -2]]]]> : tensor<1x2x2x2xi8>}> : () -> tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>
-  %cst_weights = "tfl.dequantize"(%cst_qweights) : (tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>) -> tensor<1x2x2x2xf32>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.depthwise_conv_2d"(%cst_input, %cst_weights, %cst_bias) <{depth_multiplier = 2 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<1x2x2x2xf32>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
-  func.return %0 : tensor<1x1x1x2xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}3.100000e+01, -4.000000e+00{{\]\]\]\]}}> : tensor<1x1x1x2xf32>
-  // CHECK-NOT: depthwise_conv_2d
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldTransposeConvDRQW8A32PerAxis
-func.func @ConstantFoldTransposeConvDRQW8A32PerAxis() -> tensor<1x2x2x2xf32> {
-  %cst_shape = arith.constant dense<[1, 2, 2, 2]> : tensor<4xi32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2], [4]], [[6], [8]]], [[[1], [-1]], [[2], [-2]]]]> : tensor<2x2x2x1xi8>}> : () -> tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>
-  %cst_input = arith.constant dense<[[[[2.0]]]]> : tensor<1x1x1x1xf32>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.transpose_conv"(%cst_shape, %cst_weights, %cst_input, %cst_bias) <{fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<4xi32>, tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, tensor<1x1x1x1xf32>, tensor<2xf32>) -> tensor<1x2x2x2xf32>
-  func.return %0 : tensor<1x2x2x2xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}3.000000e+00, 6.000000e+00], [5.000000e+00, -2.000000e+00]], {{\[\[}}7.000000e+00, 1.000000e+01], [9.000000e+00, -6.000000e+00{{\]\]\]\]}}> : tensor<1x2x2x2xf32>
-  // CHECK-NOT: transpose_conv
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldTransposeConvDRQW4A32PerAxis
-func.func @ConstantFoldTransposeConvDRQW4A32PerAxis() -> tensor<1x2x2x2xf32> {
-  %cst_shape = arith.constant dense<[1, 2, 2, 2]> : tensor<4xi32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x2x1x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2], [4]], [[6], [-2]]], [[[1], [-1]], [[3], [-2]]]]> : tensor<2x2x2x1xi4>}> : () -> tensor<2x2x2x1x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>
-  %cst_input = arith.constant dense<[[[[2.0]]]]> : tensor<1x1x1x1xf32>
-  %cst_bias = arith.constant dense<[5.0, -3.0]> : tensor<2xf32>
-
-  %0 = "tfl.transpose_conv"(%cst_shape, %cst_weights, %cst_input, %cst_bias) <{fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<4xi32>, tensor<2x2x2x1x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>, tensor<1x1x1x1xf32>, tensor<2xf32>) -> tensor<1x2x2x2xf32>
-  func.return %0 : tensor<1x2x2x2xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}7.000000e+00, 1.000000e+00], [9.000000e+00, -7.000000e+00]], {{\[\[}}1.100000e+01, 9.000000e+00], [3.000000e+00, -1.100000e+01{{\]\]\]\]}}> : tensor<1x2x2x2xf32>
-  // CHECK-NOT: transpose_conv
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldTransposeConvWeightOnlyW8A32PerAxis
-func.func @ConstantFoldTransposeConvWeightOnlyW8A32PerAxis() -> tensor<1x2x2x2xf32> {
-  %cst_shape = arith.constant dense<[1, 2, 2, 2]> : tensor<4xi32>
-  %cst_qweights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2], [4]], [[6], [8]]], [[[1], [-1]], [[2], [-2]]]]> : tensor<2x2x2x1xi8>}> : () -> tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>
-  %cst_weights = "tfl.dequantize"(%cst_qweights) : (tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>) -> tensor<2x2x2x1xf32>
-  %cst_input = arith.constant dense<[[[[2.0]]]]> : tensor<1x1x1x1xf32>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.transpose_conv"(%cst_shape, %cst_weights, %cst_input, %cst_bias) <{fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<4xi32>, tensor<2x2x2x1xf32>, tensor<1x1x1x1xf32>, tensor<2xf32>) -> tensor<1x2x2x2xf32>
-  func.return %0 : tensor<1x2x2x2xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}3.000000e+00, 6.000000e+00], [5.000000e+00, -2.000000e+00]], {{\[\[}}7.000000e+00, 1.000000e+01], [9.000000e+00, -6.000000e+00{{\]\]\]\]}}> : tensor<1x2x2x2xf32>
-  // CHECK-NOT: transpose_conv
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldFullyConnectedDRQU8A32NonZeroPoint
-func.func @ConstantFoldFullyConnectedDRQU8A32NonZeroPoint() -> tensor<1x2xf32> {
-  %cst_input = arith.constant dense<[[2.0, 3.0]]> : tensor<1x2xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x!quant.uniform<u8:f32:0, {0.5:128, 2.0:130}>>, value = dense<[[-126, -124], [-127, -125]]> : tensor<2x2xi8>}> : () -> tensor<2x2x!quant.uniform<u8:f32:0, {0.5:128, 2.0:130}>>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.fully_connected"(%cst_input, %cst_weights, %cst_bias) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<1x2xf32>, tensor<2x2x!quant.uniform<u8:f32:0, {0.5:128, 2.0:130}>>, tensor<2xf32>) -> tensor<1x2xf32>
-  func.return %0 : tensor<1x2xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[}}9.000000e+00, 4.000000e+00{{\]\]}}> : tensor<1x2xf32>
-  // CHECK-NOT: fully_connected
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldFullyConnectedDRQW8A32DenseResource
-func.func @ConstantFoldFullyConnectedDRQW8A32DenseResource() -> tensor<1x2xf32> {
-  %cst_input = arith.constant dense<[[2.0, 3.0]]> : tensor<1x2xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, value = dense_resource<fc_w8_2x2> : tensor<2x2xi8>}> : () -> tensor<2x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.fully_connected"(%cst_input, %cst_weights, %cst_bias) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<1x2xf32>, tensor<2x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x2xf32>
-  func.return %0 : tensor<1x2xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[}}9.000000e+00, 1.600000e+01{{\]\]}}> : tensor<1x2xf32>
-  // CHECK-NOT: fully_connected
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldFullyConnectedDRQW4A32DenseResource
-func.func @ConstantFoldFullyConnectedDRQW4A32DenseResource() -> tensor<1x2xf32> {
-  %cst_input = arith.constant dense<[[2.0, 3.0]]> : tensor<1x2xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>, value = dense_resource<fc_w4_2x2> : tensor<2x2xi4>}> : () -> tensor<2x2x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.fully_connected"(%cst_input, %cst_weights, %cst_bias) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<1x2xf32>, tensor<2x2x!quant.uniform<i4<-7:7>:f32:0, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x2xf32>
-  func.return %0 : tensor<1x2xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[}}9.000000e+00, -4.000000e+00{{\]\]}}> : tensor<1x2xf32>
-  // CHECK-NOT: fully_connected
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldConv2DGrouped
-func.func @ConstantFoldConv2DGrouped() -> tensor<1x1x1x2xf32> {
-  %cst_input = arith.constant dense<[[[[1.0, 2.0, 3.0, 4.0]]]]> : tensor<1x1x1x4xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x1x1x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2, 4]]], [[[1, -1]]]]> : tensor<2x1x1x2xi8>}> : () -> tensor<2x1x1x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.conv_2d"(%cst_input, %cst_weights, %cst_bias) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x1x1x4xf32>, tensor<2x1x1x2x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
-  func.return %0 : tensor<1x1x1x2xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}6.000000e+00, 0.000000e+00{{\]\]\]\]}}> : tensor<1x1x1x2xf32>
-  // CHECK-NOT: conv_2d
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldConv2DDilated
-func.func @ConstantFoldConv2DDilated() -> tensor<1x1x1x1xf32> {
-  %cst_input = arith.constant dense<[[[[1.0], [0.0], [2.0]], [[0.0], [0.0], [0.0]], [[3.0], [0.0], [4.0]]]]> : tensor<1x3x3x1xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<1x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5}>>, value = dense<[[[[2], [4]], [[6], [8]]]]> : tensor<1x2x2x1xi8>}> : () -> tensor<1x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5}>>
-  %cst_bias = arith.constant dense<[1.0]> : tensor<1xf32>
-
-  %0 = "tfl.conv_2d"(%cst_input, %cst_weights, %cst_bias) <{dilation_h_factor = 2 : i32, dilation_w_factor = 2 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x3x3x1xf32>, tensor<1x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5}>>, tensor<1xf32>) -> tensor<1x1x1x1xf32>
-  func.return %0 : tensor<1x1x1x1xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<3.100000e+01> : tensor<1x1x1x1xf32>
-  // CHECK-NOT: conv_2d
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldConv2DF32
-func.func @ConstantFoldConv2DF32() -> tensor<1x1x1x2xf32> {
-  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
-  %cst_weights = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]], [[[2.0], [-2.0]], [[4.0], [-4.0]]]]> : tensor<2x2x2x1xf32>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.conv_2d"(%cst_input, %cst_weights, %cst_bias) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<2x2x2x1xf32>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
-  func.return %0 : tensor<1x1x1x2xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}3.100000e+01, -4.000000e+00{{\]\]\]\]}}> : tensor<1x1x1x2xf32>
-  // CHECK-NOT: conv_2d
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldDepthwiseConv2DF32
-func.func @ConstantFoldDepthwiseConv2DF32() -> tensor<1x1x1x2xf32> {
-  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
-  %cst_weights = arith.constant dense<[[[[1.0, 2.0], [2.0, -2.0]], [[3.0, 4.0], [4.0, -4.0]]]]> : tensor<1x2x2x2xf32>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.depthwise_conv_2d"(%cst_input, %cst_weights, %cst_bias) <{depth_multiplier = 2 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<1x2x2x2xf32>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
-  func.return %0 : tensor<1x1x1x2xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}3.100000e+01, -4.000000e+00{{\]\]\]\]}}> : tensor<1x1x1x2xf32>
-  // CHECK-NOT: depthwise_conv_2d
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldTransposeConvF32
-func.func @ConstantFoldTransposeConvF32() -> tensor<1x2x2x2xf32> {
-  %cst_shape = arith.constant dense<[1, 2, 2, 2]> : tensor<4xi32>
-  %cst_weights = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]], [[[2.0], [-2.0]], [[4.0], [-4.0]]]]> : tensor<2x2x2x1xf32>
-  %cst_input = arith.constant dense<[[[[2.0]]]]> : tensor<1x1x1x1xf32>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.transpose_conv"(%cst_shape, %cst_weights, %cst_input, %cst_bias) <{fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<4xi32>, tensor<2x2x2x1xf32>, tensor<1x1x1x1xf32>, tensor<2xf32>) -> tensor<1x2x2x2xf32>
-  func.return %0 : tensor<1x2x2x2xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}3.000000e+00, 6.000000e+00], [5.000000e+00, -2.000000e+00]], {{\[\[}}7.000000e+00, 1.000000e+01], [9.000000e+00, -6.000000e+00{{\]\]\]\]}}> : tensor<1x2x2x2xf32>
-  // CHECK-NOT: transpose_conv
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldConv2DSameAsymmetricPadStride2
-func.func @ConstantFoldConv2DSameAsymmetricPadStride2() -> tensor<1x2x2x1xf32> {
-  %cst_input = arith.constant dense<1.0> : tensor<1x3x3x1xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<1x4x4x1x!quant.uniform<i8<-127:127>:f32:0, {0.5}>>, value = dense<[[[[1], [1], [1], [1]], [[2], [2], [2], [2]], [[3], [3], [3], [3]], [[4], [4], [4], [4]]]]> : tensor<1x4x4x1xi8>}> : () -> tensor<1x4x4x1x!quant.uniform<i8<-127:127>:f32:0, {0.5}>>
-  %cst_bias = arith.constant dense<[1.0]> : tensor<1xf32>
-
-  %0 = "tfl.conv_2d"(%cst_input, %cst_weights, %cst_bias) <{dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<1x3x3x1xf32>, tensor<1x4x4x1x!quant.uniform<i8<-127:127>:f32:0, {0.5}>>, tensor<1xf32>) -> tensor<1x2x2x1xf32>
-  func.return %0 : tensor<1x2x2x1xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}1.450000e+01], [1.000000e+01]], {{\[\[}}5.500000e+00], [4.000000e+00{{\]\]\]\]}}> : tensor<1x2x2x1xf32>
-  // CHECK-NOT: conv_2d
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldDepthwiseConv2DSameAsymmetricPadStride2
-func.func @ConstantFoldDepthwiseConv2DSameAsymmetricPadStride2() -> tensor<1x2x2x1xf32> {
-  %cst_input = arith.constant dense<1.0> : tensor<1x3x3x1xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<1x4x4x1x!quant.uniform<i8<-127:127>:f32:3, {0.5}>>, value = dense<[[[[1], [1], [1], [1]], [[2], [2], [2], [2]], [[3], [3], [3], [3]], [[4], [4], [4], [4]]]]> : tensor<1x4x4x1xi8>}> : () -> tensor<1x4x4x1x!quant.uniform<i8<-127:127>:f32:3, {0.5}>>
-  %cst_bias = arith.constant dense<[1.0]> : tensor<1xf32>
-
-  %0 = "tfl.depthwise_conv_2d"(%cst_input, %cst_weights, %cst_bias) <{depth_multiplier = 1 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<1x3x3x1xf32>, tensor<1x4x4x1x!quant.uniform<i8<-127:127>:f32:3, {0.5}>>, tensor<1xf32>) -> tensor<1x2x2x1xf32>
-  func.return %0 : tensor<1x2x2x1xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}1.450000e+01], [1.000000e+01]], {{\[\[}}5.500000e+00], [4.000000e+00{{\]\]\]\]}}> : tensor<1x2x2x1xf32>
-  // CHECK-NOT: depthwise_conv_2d
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @ConstantFoldTransposeConvSameNonZeroPad
-func.func @ConstantFoldTransposeConvSameNonZeroPad() -> tensor<1x3x3x1xf32> {
-  %cst_shape = arith.constant dense<[1, 3, 3, 1]> : tensor<4xi32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<1x3x3x1x!quant.uniform<i8<-127:127>:f32:0, {0.5}>>, value = dense<2> : tensor<1x3x3x1xi8>}> : () -> tensor<1x3x3x1x!quant.uniform<i8<-127:127>:f32:0, {0.5}>>
-  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
-  %cst_bias = arith.constant dense<[0.0]> : tensor<1xf32>
-
-  %0 = "tfl.transpose_conv"(%cst_shape, %cst_weights, %cst_input, %cst_bias) <{fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<4xi32>, tensor<1x3x3x1x!quant.uniform<i8<-127:127>:f32:0, {0.5}>>, tensor<1x2x2x1xf32>, tensor<1xf32>) -> tensor<1x3x3x1xf32>
-  func.return %0 : tensor<1x3x3x1xf32>
-
-  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[\[\[\[}}1.000000e+00], [3.000000e+00], [2.000000e+00]], {{\[\[}}4.000000e+00], [1.000000e+01], [6.000000e+00]], {{\[\[}}3.000000e+00], [7.000000e+00], [4.000000e+00{{\]\]\]\]}}> : tensor<1x3x3x1xf32>
-  // CHECK-NOT: transpose_conv
-  // CHECK: return %[[CST]]
-}
-
-// CHECK-LABEL: @NoConstantFoldDepthwiseConv2DFusedRelu
-func.func @NoConstantFoldDepthwiseConv2DFusedRelu() -> tensor<1x1x1x2xf32> {
-  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>, value = dense<[[[[2, 1], [4, -1]], [[6, 2], [8, -2]]]]> : tensor<1x2x2x2xi8>}> : () -> tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.depthwise_conv_2d"(%cst_input, %cst_weights, %cst_bias) <{depth_multiplier = 2 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "RELU", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x1x1x2xf32>
-  func.return %0 : tensor<1x1x1x2xf32>
-
-  // CHECK: tfl.depthwise_conv_2d
-}
-
-// CHECK-LABEL: @NoConstantFoldDepthwiseConv2DDynamicShape
-func.func @NoConstantFoldDepthwiseConv2DDynamicShape() -> tensor<1x?x?x2xf32> {
-  %cst_input = arith.constant dense<[[[[1.0], [2.0]], [[3.0], [4.0]]]]> : tensor<1x2x2x1xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>, value = dense<[[[[2, 1], [4, -1]], [[6, 2], [8, -2]]]]> : tensor<1x2x2x2xi8>}> : () -> tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.depthwise_conv_2d"(%cst_input, %cst_weights, %cst_bias) <{depth_multiplier = 2 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x2x2x1xf32>, tensor<1x2x2x2x!quant.uniform<i8<-127:127>:f32:3, {0.5, 2.0}>>, tensor<2xf32>) -> tensor<1x?x?x2xf32>
-  func.return %0 : tensor<1x?x?x2xf32>
-
-  // CHECK: tfl.depthwise_conv_2d
-}
-
-// CHECK-LABEL: @NoConstantFoldTransposeConvFusedRelu
-func.func @NoConstantFoldTransposeConvFusedRelu() -> tensor<1x2x2x2xf32> {
-  %cst_shape = arith.constant dense<[1, 2, 2, 2]> : tensor<4xi32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2], [4]], [[6], [8]]], [[[1], [-1]], [[2], [-2]]]]> : tensor<2x2x2x1xi8>}> : () -> tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>
-  %cst_input = arith.constant dense<[[[[2.0]]]]> : tensor<1x1x1x1xf32>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.transpose_conv"(%cst_shape, %cst_weights, %cst_input, %cst_bias) <{fused_activation_function = "RELU", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<4xi32>, tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, tensor<1x1x1x1xf32>, tensor<2xf32>) -> tensor<1x2x2x2xf32>
-  func.return %0 : tensor<1x2x2x2xf32>
-
-  // CHECK: tfl.transpose_conv
-}
-
-// CHECK-LABEL: @NoConstantFoldTransposeConvDynamicShape
-func.func @NoConstantFoldTransposeConvDynamicShape() -> tensor<1x?x?x2xf32> {
-  %cst_shape = arith.constant dense<[1, 2, 2, 2]> : tensor<4xi32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, value = dense<[[[[2], [4]], [[6], [8]]], [[[1], [-1]], [[2], [-2]]]]> : tensor<2x2x2x1xi8>}> : () -> tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>
-  %cst_input = arith.constant dense<[[[[2.0]]]]> : tensor<1x1x1x1xf32>
-  %cst_bias = arith.constant dense<[1.0, 2.0]> : tensor<2xf32>
-
-  %0 = "tfl.transpose_conv"(%cst_shape, %cst_weights, %cst_input, %cst_bias) <{fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32}> : (tensor<4xi32>, tensor<2x2x2x1x!quant.uniform<i8<-127:127>:f32:0, {0.5, 2.0}>>, tensor<1x1x1x1xf32>, tensor<2xf32>) -> tensor<1x?x?x2xf32>
-  func.return %0 : tensor<1x?x?x2xf32>
-
-  // CHECK: tfl.transpose_conv
-}
-
-// CHECK-LABEL: @NoConstantFoldFullyConnectedExceedsSizeLimit
-func.func @NoConstantFoldFullyConnectedExceedsSizeLimit() -> tensor<5x5xf32> {
-  %cst_input = arith.constant dense<1.0> : tensor<5x2xf32>
-  %cst_weights = "tfl.pseudo_qconst"() <{qtype = tensor<5x2x!quant.uniform<i8:f32, 0.5:0>>, value = dense<2> : tensor<5x2xi8>}> : () -> tensor<5x2x!quant.uniform<i8:f32, 0.5:0>>
-  %cst_bias = arith.constant dense<0.0> : tensor<5xf32>
-
-  %0 = "tfl.fully_connected"(%cst_input, %cst_weights, %cst_bias) <{asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"}> : (tensor<5x2xf32>, tensor<5x2x!quant.uniform<i8:f32, 0.5:0>>, tensor<5xf32>) -> tensor<5x5xf32>
-  func.return %0 : tensor<5x5xf32>
-
-  // CHECK: tfl.fully_connected
-}
-
 // CHECK-LABEL: @ShapeOpI32
 func.func @ShapeOpI32(%arg0 : tensor<576x72xf32>) -> tensor<2xi32> {
   %0 = "tfl.shape"(%arg0) : (tensor<576x72xf32>) -> tensor<2xi32>
@@ -1746,11 +1313,5 @@ func.func @slice_big_float() -> tensor<1x1x1792x256xf32> {
 // CHECK-LITERAL: %cst = arith.constant dense<9.000000e+00> : tensor<1x1x1792x256xf32>
 // CHECK-NOT:     slice
 
-{-#
-  dialect_resources: {
-    builtin: {
-      fc_w8_2x2: "0x040000000204FF03",
-      fc_w4_2x2: "0x04000000421D"
-    }
-  }
-#-}
+
+
diff --git a/tensorflow/compiler/mlir/lite/tests/optimize.mlir b/tensorflow/compiler/mlir/lite/tests/optimize.mlir
index c18da5a22ed..2861c01f2d5 100644
--- a/tensorflow/compiler/mlir/lite/tests/optimize.mlir
+++ b/tensorflow/compiler/mlir/lite/tests/optimize.mlir
@@ -5048,47 +5048,6 @@ func.func @FullyConnectedSwapOperandsWhenLHSIsConstLHSRank3(%arg0: tensor<512x51
   // CHECK:  %0 = "tfl.fully_connected"(%cst, %arg0, %arg1)
 }
 
-// CHECK-LABEL: @FullyConnectedSwapOperandsWhenLHSIsConstQConstFilter
-func.func @FullyConnectedSwapOperandsWhenLHSIsConstQConstFilter(%arg0: none) -> tensor<2x2xf32> {
-  %cst = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xf32>
-  %qconst = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x!quant.uniform<i16:f32, 0.5:0>>, value = dense<2> : tensor<2x2xi16>}> : () -> tensor<2x2x!quant.uniform<i16:f32, 0.5:0>>
-  %0 = "tfl.fully_connected"(%cst, %qconst, %arg0) {asymmetric_quantize_inputs = true, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<2x2xf32>, tensor<2x2x!quant.uniform<i16:f32, 0.5:0>>, none) -> tensor<2x2xf32>
-  func.return %0 : tensor<2x2xf32>
-
-  // CHECK: %[[QCONST:.*]] = "tfl.pseudo_qconst"
-  // CHECK: %[[FC:.*]] = "tfl.fully_connected"(%cst, %[[QCONST]], %arg0)
-  // CHECK-NOT: tfl.transpose
-  // CHECK: return %[[FC]] : tensor<2x2xf32>
-}
-
-// CHECK-LABEL: @FullyConnectedSwapOperandsWhenLHSIsConstDequantQConstFilter
-func.func @FullyConnectedSwapOperandsWhenLHSIsConstDequantQConstFilter(%arg0: none) -> tensor<2x2xf32> {
-  %cst = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xf32>
-  %qconst = "tfl.pseudo_qconst"() <{qtype = tensor<2x2x!quant.uniform<i16:f32, 0.5:0>>, value = dense<2> : tensor<2x2xi16>}> : () -> tensor<2x2x!quant.uniform<i16:f32, 0.5:0>>
-  %dequant = "tfl.dequantize"(%qconst) : (tensor<2x2x!quant.uniform<i16:f32, 0.5:0>>) -> tensor<2x2xf32>
-  %0 = "tfl.fully_connected"(%cst, %dequant, %arg0) {asymmetric_quantize_inputs = true, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<2x2xf32>, tensor<2x2xf32>, none) -> tensor<2x2xf32>
-  func.return %0 : tensor<2x2xf32>
-
-  // CHECK: %[[DEQUANT:.*]] = "tfl.dequantize"
-  // CHECK: %[[FC:.*]] = "tfl.fully_connected"(%cst, %[[DEQUANT]], %arg0)
-  // CHECK-NOT: tfl.transpose
-  // CHECK: return %[[FC]] : tensor<2x2xf32>
-}
-
-// CHECK-LABEL: @FullyConnectedSwapOperandsWhenLHSIsConstDequantF16Filter
-func.func @FullyConnectedSwapOperandsWhenLHSIsConstDequantF16Filter(%arg0: none) -> tensor<2x2xf32> {
-  %cst = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xf32>
-  %cst_f16 = "tfl.pseudo_const"() <{value = dense<1.0> : tensor<2x2xf16>}> : () -> tensor<2x2xf16>
-  %dequant = "tfl.dequantize"(%cst_f16) : (tensor<2x2xf16>) -> tensor<2x2xf32>
-  %0 = "tfl.fully_connected"(%cst, %dequant, %arg0) {asymmetric_quantize_inputs = true, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<2x2xf32>, tensor<2x2xf32>, none) -> tensor<2x2xf32>
-  func.return %0 : tensor<2x2xf32>
-
-  // CHECK: %[[DEQUANT:.*]] = "tfl.dequantize"
-  // CHECK: %[[FC:.*]] = "tfl.fully_connected"(%cst, %[[DEQUANT]], %arg0)
-  // CHECK-NOT: tfl.transpose
-  // CHECK: return %[[FC]] : tensor<2x2xf32>
-}
-
 // CHECK-LABEL: @AddComputedZero
 func.func @AddComputedZero(%arg0: tensor<512x512xf32>, %arg1: tensor<1x512xf32>) -> tensor<512x512xf32> {
   %0 = "tfl.sub"(%arg1, %arg1) {fused_activation_function = "NONE"} : (tensor<1x512xf32>, tensor<1x512xf32>) -> tensor<1x512xf32>
diff --git a/tensorflow/compiler/mlir/lite/transforms/optimize_pass.cc b/tensorflow/compiler/mlir/lite/transforms/optimize_pass.cc
index 16d9983252e..0da83d8c342 100644
--- a/tensorflow/compiler/mlir/lite/transforms/optimize_pass.cc
+++ b/tensorflow/compiler/mlir/lite/transforms/optimize_pass.cc
@@ -3331,14 +3331,8 @@ struct FullyConnectedSwapOperandsWhenLHSIsConst
     auto input = fc.getInput();
     auto filter = fc.getFilter();
 
-    auto is_const_or_qconst = [](Value v) {
-      return matchPattern(v, m_Constant()) ||
-             llvm::isa_and_nonnull<TFL::QConstOp, TFL::SparseConstOp,
-                                   TFL::SparseQConstOp>(v.getDefiningOp());
-    };
-    auto filter_dequant = filter.getDefiningOp<TFL::DequantizeOp>();
-    if (!matchPattern(input, m_Constant()) || is_const_or_qconst(filter) ||
-        (filter_dequant && is_const_or_qconst(filter_dequant.getInput())))
+    if (!matchPattern(input, m_Constant()) ||
+        matchPattern(filter, m_Constant()))
       return failure();
 
     auto input_type = mlir::dyn_cast<RankedTensorType>(input.getType());

From 20f071377e8c43fac2b343a65fc8848bdbebffab Mon Sep 17 00:00:00 2001
From: Maxim Ermilov CA <maximermilov@google.com>
Date: Thu, 15 Oct 2026 12:47:35 -0700
Subject: [PATCH 0901/1000] Fix complex GetDistanceErr and restore skipping
 extreme inputs in C128 Log test.

`GetDistanceErr` in `xla/tests/exhaustive/exhaustive_op_test_base.cc` had a typo where `CalculateDistanceInFloats(expected.imag(), expected.imag())` was computed instead of `CalculateDistanceInFloats(expected.imag(), actual.imag())`, causing `distance_err` checks on complex outputs to ignore errors in the imaginary component.

Fixing `GetDistanceErr` exposes the imaginary-component precision issues in `ExhaustiveC128UnaryTest.Log` for inputs with component magnitudes above `1e+300` or below `1e-300`, as well as subnormal output components flushed to zero on CPU. Restore skipping those extreme inputs and allow `DBL_MIN` absolute error when `std::log(x)` has a subnormal component.

PiperOrigin-RevId: 1000325915
---
 .../exhaustive/exhaustive_op_test_base.cc     |  5 ++--
 .../exhaustive_unary_complex_test.cc          | 23 +++++++++++++++++--
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/third_party/xla/xla/tests/exhaustive/exhaustive_op_test_base.cc b/third_party/xla/xla/tests/exhaustive/exhaustive_op_test_base.cc
index f6c16cfe4de..9f449d13f69 100644
--- a/third_party/xla/xla/tests/exhaustive/exhaustive_op_test_base.cc
+++ b/third_party/xla/xla/tests/exhaustive/exhaustive_op_test_base.cc
@@ -438,9 +438,8 @@ template <typename NativeT>
 int64_t GetDistanceErr(NativeT expected, NativeT actual) {
   if constexpr (std::is_same_v<NativeT, xla::complex64> ||
                 std::is_same_v<NativeT, xla::complex128>) {
-    return std::max(
-        CalculateDistanceInFloats(expected.real(), actual.real()),
-        CalculateDistanceInFloats(expected.imag(), expected.imag()));
+    return std::max(CalculateDistanceInFloats(expected.real(), actual.real()),
+                    CalculateDistanceInFloats(expected.imag(), actual.imag()));
   } else {
     return CalculateDistanceInFloats(expected, actual);
   }
diff --git a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_complex_test.cc b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_complex_test.cc
index c7571950fe9..89846df8176 100644
--- a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_complex_test.cc
+++ b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_complex_test.cc
@@ -256,11 +256,30 @@ UNARY_TEST_COMPLEX_128(Log, {
   };
 
   if (IsCpu() || IsGpu()) {
-    error_spec_gen = +[](complex128 x) {
+    error_spec_gen = [&](complex128 x) {
       // TODO(rmlarsen): see b/162664705 and b/138578594
       bool should_skip = std::isnan(x.real()) || std::isnan(x.imag());
+
+      double abs_err = 0;
+      if (IsCpu() && std::abs(x.imag()) < 1.0e-6) {
+        // Subnormal imaginary parts are flashed to zero.
+        abs_err = std::numeric_limits<double>::min();
+      }
+      if (IsCpu() && std::abs(x.imag()) < 1.0e-280) {
+        // TODO(b/138578313) produces wrong result.
+        should_skip = true;
+      }
+
+      int distance_err = 1;
+      if (IsGpu() &&
+          (std::abs(x.imag()) < 1.0e-300 || std::abs(x.imag()) > 1.0e+300)) {
+        // __nv_atan2 is with in 2 ULP.
+        // https://docs.nvidia.com/cuda/cuda-programming-guide/05-appendices/mathematical-functions.html#trigonometric-functions
+        distance_err = 2;
+      }
       return ErrorSpec::Builder()
-          .distance_err(1)
+          .abs_err(abs_err)
+          .distance_err(distance_err)
           .skip_comparison(should_skip)
           .strict_signed_zeros()
           .build();

From d36d9ee22b674850ccfee310e289f981b5b3773d Mon Sep 17 00:00:00 2001
From: Vladimir Belitskiy <belitskiy@google.com>
Date: Thu, 15 Oct 2026 12:50:16 -0700
Subject: [PATCH 0902/1000] Fix `NanoRtExecutable::ManagedTemp` inline buffer
 alignment under AddressSanitizer.

`NanoRtExecutable::ManagedTemp<n>` previously applied `alignas(Align())` (`alignas(64)`) to the outer `absl::FixedArray<std::byte, n, Allocator> data_` member, assuming the inline storage buffer `buff_` is at offset `0` of `FixedArray`.

However, when `ABSL_HAVE_ADDRESS_SANITIZER` is enabled, `absl::FixedArray::NonEmptyInlinedStorage` places an 8-byte `ABSL_ADDRESS_SANITIZER_REDZONE(redzone_begin_)` member before `buff_`, and aligns `buff_` only to `alignof(StorageElement)` (`alignof(std::byte) == 1`). This places `buff_` at offset `8` inside `FixedArray`, causing `temp.data()` to be misaligned (`addr % 64 == 8`) for all inline sizes and failing `NanoRtClientStandaloneTest.ManagedTempAlignment` under ASAN.

Store `alignas(Align())` chunks in `absl::FixedArray` so that `FixedArray`'s inline buffer `buff_` is guaranteed to be aligned to `Align()` (`64`) even when preceded by `redzone_begin_` under ASAN, and rely on C++17 over-aligned `std::allocator<AlignedChunk>` for heap allocations (which also avoids a `free-size-mismatch` bug in `tsl::port::AlignedAllocator::deallocate` when `sizeof(T) > 1`).

In addition, use `ABSL_ANNOTATE_CONTIGUOUS_CONTAINER` in `ManagedTemp`'s constructor and destructor to poison the rounded chunk tail beyond the requested byte size `size_` and restore the full chunk extent before `FixedArray` destroys its storage, and use overflow-safe ceiling division in `CeilOfRatio`.

PiperOrigin-RevId: 1000327474
---
 third_party/xla/xla/backends/cpu/nanort/BUILD |  1 -
 .../backends/cpu/nanort/nanort_client_test.cc | 39 ++++++++++++++++++
 .../backends/cpu/nanort/nanort_executable.h   | 40 +++++++++++++++----
 3 files changed, 72 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/xla/backends/cpu/nanort/BUILD b/third_party/xla/xla/backends/cpu/nanort/BUILD
index 10f96d816e3..1b8aa55cc85 100644
--- a/third_party/xla/xla/backends/cpu/nanort/BUILD
+++ b/third_party/xla/xla/backends/cpu/nanort/BUILD
@@ -153,7 +153,6 @@ cc_library(
         "@com_google_absl//absl/types:span",
         "@eigen_archive//:eigen3",
         "@tsl//tsl/platform:casts",
-        "@tsl//tsl/platform:platform_port",
         "@tsl//tsl/profiler/lib:traceme",
         "@tsl//tsl/profiler/lib:traceme_encode",
     ],
diff --git a/third_party/xla/xla/backends/cpu/nanort/nanort_client_test.cc b/third_party/xla/xla/backends/cpu/nanort/nanort_client_test.cc
index 3369fd167cc..37348044d79 100644
--- a/third_party/xla/xla/backends/cpu/nanort/nanort_client_test.cc
+++ b/third_party/xla/xla/backends/cpu/nanort/nanort_client_test.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include <cstdint>
 #include <functional>
 #include <memory>
+#include <new>
 #include <optional>
 #include <string>
 #include <utility>
@@ -109,6 +110,44 @@ TEST(NanoRtClientStandaloneTest, ManagedTempAlignment) {
   EXPECT_EQ(reinterpret_cast<uintptr_t>(&temp3.data()[0]) % Align(), 0);
 }
 
+template <size_t n>
+void ExpectManagedTempSizeAndContents() {
+  constexpr size_t sizes[] = {0, 1, 3, 7, 8, 63, 64, 65, 127, 128, 129};
+  for (size_t size : sizes) {
+    SCOPED_TRACE(size);
+    NanoRtExecutable::ManagedTemp<n> temp(size);
+    absl::Span<std::byte> bytes = temp.data();
+    EXPECT_EQ(bytes.size(), size);
+    EXPECT_EQ(reinterpret_cast<uintptr_t>(bytes.data()) % Align(), 0);
+    for (std::byte& byte : bytes) {
+      byte = std::byte{0x42};
+    }
+    for (std::byte byte : bytes) {
+      EXPECT_EQ(byte, std::byte{0x42});
+    }
+  }
+}
+
+TEST(NanoRtClientStandaloneTest, ManagedTempSizeAndContents) {
+  ExpectManagedTempSizeAndContents<0>();
+  ExpectManagedTempSizeAndContents<3>();
+  ExpectManagedTempSizeAndContents<2 * Align()>();
+}
+
+TEST(NanoRtClientStandaloneTest, ManagedTempUnpoisonsStorageOnDestruction) {
+  using Temp = NanoRtExecutable::ManagedTemp<3>;
+  alignas(Temp) std::byte storage[sizeof(Temp)];
+  Temp* temp = new (storage) Temp(3);
+  temp->~Temp();
+  // Writes to the former inline buffer must not encounter stale ASAN poison.
+  volatile std::byte* bytes = storage;
+  for (size_t i = 0; i < sizeof(storage); ++i) {
+    bytes[i] = std::byte{0x42};
+    std::byte byte = bytes[i];
+    EXPECT_EQ(byte, std::byte{0x42});
+  }
+}
+
 TEST_P(NanoRtClientTest, CompileAndRunScalarComputation) {
   constexpr absl::string_view hlo = R"(
     HloModule add
diff --git a/third_party/xla/xla/backends/cpu/nanort/nanort_executable.h b/third_party/xla/xla/backends/cpu/nanort/nanort_executable.h
index dfc8a4f2bf7..0b36807bd93 100644
--- a/third_party/xla/xla/backends/cpu/nanort/nanort_executable.h
+++ b/third_party/xla/xla/backends/cpu/nanort/nanort_executable.h
@@ -39,7 +39,6 @@ limitations under the License.
 #include "xla/shape.h"
 #include "xla/tsl/concurrency/async_value_ref.h"
 #include "xla/tsl/concurrency/chain.h"
-#include "tsl/platform/mem.h"
 
 #define EIGEN_USE_THREADS
 
@@ -164,21 +163,48 @@ class NanoRtExecutable {
   template <size_t n>
   class ManagedTemp {
    public:
-    explicit ManagedTemp(size_t size) : data_(size) {
+    explicit ManagedTemp(size_t size)
+        : data_(CeilOfRatio(size, Align())), size_(size) {
       ABSL_ANNOTATE_MEMORY_IS_INITIALIZED(data_.data(), data_.memsize());
+      // Keep the rounded chunk tail inaccessible beyond the requested size.
+      AnnotateMemory(data_.memsize(), size_);
+    }
+
+    ~ManagedTemp() {
+      // Restore the full chunk extent before FixedArray destroys its storage.
+      AnnotateMemory(size_, data_.memsize());
     }
 
     ManagedTemp(const ManagedTemp&) = delete;
     ManagedTemp& operator=(const ManagedTemp&) = delete;
 
-    PreallocatedTemp data() { return absl::MakeSpan(data_); }
+    PreallocatedTemp data() {
+      return absl::MakeSpan(reinterpret_cast<std::byte*>(data_.data()), size_);
+    }
 
    private:
     friend class NanoRtExecutable;
-    using Allocator =
-        tsl::port::AlignedAllocator<std::byte,
-                                    static_cast<std::align_val_t>(Align())>;
-    alignas(Align()) absl::FixedArray<std::byte, n, Allocator> data_;
+
+    struct alignas(Align()) AlignedChunk {
+      std::byte data[Align()];
+    };
+
+    static constexpr size_t CeilOfRatio(size_t num, size_t div) {
+      return num / div + (num % div != 0);
+    }
+
+    void AnnotateMemory(size_t old_size, size_t new_size) {
+      if (size_ != 0) {
+        [[maybe_unused]] std::byte* begin =
+            reinterpret_cast<std::byte*>(data_.data());
+        [[maybe_unused]] std::byte* end = begin + data_.memsize();
+        ABSL_ANNOTATE_CONTIGUOUS_CONTAINER(begin, end, begin + old_size,
+                                           begin + new_size);
+      }
+    }
+
+    absl::FixedArray<AlignedChunk, CeilOfRatio(n, Align())> data_;
+    size_t size_;
   };
 
   tsl::AsyncValueRef<ExecuteEvent> Execute(absl::Span<const Argument> arguments,

From 21f3e33634715b643924d17f571309984f87cbef Mon Sep 17 00:00:00 2001
From: Matthias Kramm <kramm@google.com>
Date: Thu, 15 Oct 2026 13:03:19 -0700
Subject: [PATCH 0903/1000] Weave 'multithreaded' through all tsl file system
 implementations.

PiperOrigin-RevId: 1000335941
---
 .../tsl/platform/default/posix_file_system.cc |  9 +++++-
 .../tsl/platform/default/posix_file_system.h  |  4 +++
 third_party/xla/xla/tsl/platform/env.cc       |  8 ++++-
 third_party/xla/xla/tsl/platform/env.h        | 14 ++++++++
 third_party/xla/xla/tsl/platform/env_test.cc  | 32 +++++++++++++++++++
 .../xla/xla/tsl/platform/file_system.h        |  6 ++++
 .../xla/xla/tsl/platform/file_system_helper.h |  1 +
 .../platform/windows/windows_file_system.cc   | 10 ++++--
 .../platform/windows/windows_file_system.h    |  4 +++
 9 files changed, 84 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/tsl/platform/default/posix_file_system.cc b/third_party/xla/xla/tsl/platform/default/posix_file_system.cc
index 5537c886cd9..e6e14b6bfab 100644
--- a/third_party/xla/xla/tsl/platform/default/posix_file_system.cc
+++ b/third_party/xla/xla/tsl/platform/default/posix_file_system.cc
@@ -326,7 +326,14 @@ absl::Status PosixFileSystem::GetChildren(const std::string& dir,
 
 absl::Status PosixFileSystem::GetMatchingPaths(
     const std::string& pattern, std::vector<std::string>* results) {
-  return internal::GetMatchingPaths(this, Env::Default(), pattern, results);
+  return GetMatchingPaths(pattern, results, /*multithreaded=*/true);
+}
+
+absl::Status PosixFileSystem::GetMatchingPaths(
+    const std::string& pattern, std::vector<std::string>* results,
+    bool multithreaded) {
+  return internal::GetMatchingPaths(this, Env::Default(), pattern, results,
+                                    multithreaded);
 }
 
 absl::Status PosixFileSystem::DeleteFile(const std::string& fname) {
diff --git a/third_party/xla/xla/tsl/platform/default/posix_file_system.h b/third_party/xla/xla/tsl/platform/default/posix_file_system.h
index a0a12cff800..c3cc6f78c98 100644
--- a/third_party/xla/xla/tsl/platform/default/posix_file_system.h
+++ b/third_party/xla/xla/tsl/platform/default/posix_file_system.h
@@ -52,6 +52,10 @@ class PosixFileSystem : public FileSystem {
   absl::Status GetMatchingPaths(const std::string& pattern,
                                 std::vector<std::string>* results) override;
 
+  absl::Status GetMatchingPaths(const std::string& pattern,
+                                std::vector<std::string>* results,
+                                bool multithreaded) override;
+
   absl::Status DeleteFile(const std::string& fname) override;
 
   absl::Status CreateDir(const std::string& name) override;
diff --git a/third_party/xla/xla/tsl/platform/env.cc b/third_party/xla/xla/tsl/platform/env.cc
index 5a03cf87a17..41231b6063c 100644
--- a/third_party/xla/xla/tsl/platform/env.cc
+++ b/third_party/xla/xla/tsl/platform/env.cc
@@ -297,9 +297,15 @@ absl::Status Env::GetChildren(const std::string& dir,
 
 absl::Status Env::GetMatchingPaths(const std::string& pattern,
                                    std::vector<std::string>* results) {
+  return GetMatchingPaths(pattern, results, /*multithreaded=*/true);
+}
+
+absl::Status Env::GetMatchingPaths(const std::string& pattern,
+                                   std::vector<std::string>* results,
+                                   bool multithreaded) {
   FileSystem* fs;
   ABSL_RETURN_IF_ERROR(GetFileSystemForFile(pattern, &fs));
-  return fs->GetMatchingPaths(pattern, results);
+  return fs->GetMatchingPaths(pattern, results, multithreaded);
 }
 
 absl::Status Env::DeleteFile(const std::string& fname) {
diff --git a/third_party/xla/xla/tsl/platform/env.h b/third_party/xla/xla/tsl/platform/env.h
index 38c087e216a..c057b484d45 100644
--- a/third_party/xla/xla/tsl/platform/env.h
+++ b/third_party/xla/xla/tsl/platform/env.h
@@ -205,6 +205,10 @@ class Env {
   virtual absl::Status GetMatchingPaths(const std::string& pattern,
                                         std::vector<std::string>* results);
 
+  virtual absl::Status GetMatchingPaths(const std::string& pattern,
+                                        std::vector<std::string>* results,
+                                        bool multithreaded);
+
   // TODO(b/485502789): Remove the const std::string& versions of these
   // functions and move the actual implementation here, avoiding the string
   // copy.
@@ -220,6 +224,16 @@ class Env {
     return GetMatchingPaths(std::string(pattern), results);
   }
 
+  template <typename T,
+            typename = std::enable_if_t<
+                std::is_convertible_v<const T&, absl::string_view> &&
+                !std::is_same_v<std::decay_t<T>, std::string>>>
+  absl::Status GetMatchingPaths(const T& pattern,
+                                std::vector<std::string>* results,
+                                bool multithreaded) {
+    return GetMatchingPaths(std::string(pattern), results, multithreaded);
+  }
+
   /// Deletes the named file.
   absl::Status DeleteFile(const std::string& fname);
 
diff --git a/third_party/xla/xla/tsl/platform/env_test.cc b/third_party/xla/xla/tsl/platform/env_test.cc
index f491e5f25c2..473e56b0d0a 100644
--- a/third_party/xla/xla/tsl/platform/env_test.cc
+++ b/third_party/xla/xla/tsl/platform/env_test.cc
@@ -34,6 +34,7 @@ namespace tsl {
 namespace {
 
 using absl_testing::IsOk;
+using ::testing::UnorderedElementsAre;
 
 TEST(EnvTest, StartDetachedThread) {
   Env* env = Env::Default();
@@ -119,6 +120,37 @@ TEST(EnvTest, RenameFile) {
   EXPECT_EQ(target_content, "source content");
 }
 
+TEST(EnvTest, GetMatchingPaths) {
+  Env* env = Env::Default();
+  ASSERT_OK_AND_ASSIGN(
+      tsl::testing::TemporaryDirectory temp_dir,
+      tsl::testing::TemporaryDirectory::CreateForCurrentTestcase());
+  std::string dir_a = tsl::io::JoinPath(temp_dir.path(), "dir_a");
+  std::string dir_b = tsl::io::JoinPath(temp_dir.path(), "dir_b");
+  ASSERT_THAT(env->RecursivelyCreateDir(dir_a), IsOk());
+  ASSERT_THAT(env->RecursivelyCreateDir(dir_b), IsOk());
+
+  std::string file_1 = tsl::io::JoinPath(dir_a, "file_1.txt");
+  std::string file_2 = tsl::io::JoinPath(dir_b, "file_2.txt");
+  ASSERT_THAT(WriteStringToFile(env, file_1, "1"), IsOk());
+  ASSERT_THAT(WriteStringToFile(env, file_2, "2"), IsOk());
+
+  std::string pattern =
+      tsl::io::JoinPath(temp_dir.path(), "dir_*", "file_*.txt");
+  absl::string_view pattern_sv = pattern;
+  for (bool multithreaded : {false, true}) {
+    std::vector<std::string> results;
+    EXPECT_THAT(env->GetMatchingPaths(pattern, &results, multithreaded),
+                IsOk());
+    EXPECT_THAT(results, UnorderedElementsAre(file_1, file_2));
+
+    results.clear();
+    EXPECT_THAT(env->GetMatchingPaths(pattern_sv, &results, multithreaded),
+                IsOk());
+    EXPECT_THAT(results, UnorderedElementsAre(file_1, file_2));
+  }
+}
+
 #if defined(__linux__)
 TEST(EnvTest, ReadFileToStringAllowsShortReadsForSystemFiles) {
   std::string content;
diff --git a/third_party/xla/xla/tsl/platform/file_system.h b/third_party/xla/xla/tsl/platform/file_system.h
index d6262756e70..306122a783c 100644
--- a/third_party/xla/xla/tsl/platform/file_system.h
+++ b/third_party/xla/xla/tsl/platform/file_system.h
@@ -159,6 +159,12 @@ class FileSystem {
   virtual absl::Status GetMatchingPaths(const std::string& pattern,
                                         std::vector<std::string>* results) = 0;
 
+  virtual absl::Status GetMatchingPaths(const std::string& pattern,
+                                        std::vector<std::string>* results,
+                                        bool multithreaded) {
+    return GetMatchingPaths(pattern, results);
+  }
+
   /// \brief Checks if the given filename matches the pattern.
   ///
   /// This function provides the equivalent of posix fnmatch, however it is
diff --git a/third_party/xla/xla/tsl/platform/file_system_helper.h b/third_party/xla/xla/tsl/platform/file_system_helper.h
index 3bcfb64a77f..4a3473816ed 100644
--- a/third_party/xla/xla/tsl/platform/file_system_helper.h
+++ b/third_party/xla/xla/tsl/platform/file_system_helper.h
@@ -46,6 +46,7 @@ namespace internal {
 //   env: may not be null and will be used to check if a match has been found.
 //   pattern: see FileSystem::GetMatchingPaths() for details.
 //   results: will be cleared and may not be null.
+//   multithreaded: if true, scans subdirectories in parallel where supported.
 //
 // Returns an error status if any call to 'fs' failed.
 absl::Status GetMatchingPaths(FileSystem* fs, Env* env,
diff --git a/third_party/xla/xla/tsl/platform/windows/windows_file_system.cc b/third_party/xla/xla/tsl/platform/windows/windows_file_system.cc
index 00783938da0..efaeeda3ac6 100644
--- a/third_party/xla/xla/tsl/platform/windows/windows_file_system.cc
+++ b/third_party/xla/xla/tsl/platform/windows/windows_file_system.cc
@@ -701,6 +701,12 @@ absl::Status WindowsFileSystem::RenameFile(const std::string& src,
 
 absl::Status WindowsFileSystem::GetMatchingPaths(
     const std::string& pattern, std::vector<std::string>* results) {
+  return GetMatchingPaths(pattern, results, /*multithreaded=*/true);
+}
+
+absl::Status WindowsFileSystem::GetMatchingPaths(
+    const std::string& pattern, std::vector<std::string>* results,
+    bool multithreaded) {
   // NOTE(mrry): The existing implementation of FileSystem::GetMatchingPaths()
   // does not handle Windows paths containing backslashes correctly. Since
   // Windows APIs will accept forward and backslashes equivalently, we
@@ -709,8 +715,8 @@ absl::Status WindowsFileSystem::GetMatchingPaths(
   // but no code appears to rely on this behavior.
   std::string converted_pattern(pattern);
   std::replace(converted_pattern.begin(), converted_pattern.end(), '\\', '/');
-  ABSL_RETURN_IF_ERROR(internal::GetMatchingPaths(this, Env::Default(),
-                                             converted_pattern, results));
+  ABSL_RETURN_IF_ERROR(internal::GetMatchingPaths(
+      this, Env::Default(), converted_pattern, results, multithreaded));
   for (std::string& result : *results) {
     std::replace(result.begin(), result.end(), '/', '\\');
   }
diff --git a/third_party/xla/xla/tsl/platform/windows/windows_file_system.h b/third_party/xla/xla/tsl/platform/windows/windows_file_system.h
index 9ded0f7b595..54e9aed892c 100644
--- a/third_party/xla/xla/tsl/platform/windows/windows_file_system.h
+++ b/third_party/xla/xla/tsl/platform/windows/windows_file_system.h
@@ -61,6 +61,10 @@ class WindowsFileSystem : public FileSystem {
   absl::Status GetMatchingPaths(const std::string& pattern,
                                 std::vector<std::string>* result) override;
 
+  absl::Status GetMatchingPaths(const std::string& pattern,
+                                std::vector<std::string>* result,
+                                bool multithreaded) override;
+
   bool Match(absl::string_view filename, absl::string_view pattern) override;
 
   absl::Status Stat(const std::string& fname, FileStatistics* stat) override;

From 58dd296f800071c8caf7f820096a028da94b2888 Mon Sep 17 00:00:00 2001
From: Michael Kuperstein <mkuper@google.com>
Date: Thu, 15 Oct 2026 13:15:17 -0700
Subject: [PATCH 0904/1000] Prevent copy removal from eliding copies that
 overwrite live values.

In `ComputeRelativeLocation::InstructionCanIntercept`, a `kCopy` instruction in `def` was treated as non-intercepting whenever `region.contains(operand)` was true after peeling at most one `kGetTupleElement`. When a value in `region` flows into a `kConditional` or `kWhile` via a tuple parameter and `kGetTupleElement`, `operand` is the branch `kParameter`, which is excluded from `LiveRangeRegions`. However, if another branch of the `kConditional` reads that same `kParameter` via `kGetTupleElement`, the `kParameter` is still recorded as a use in `region`, causing `region.contains(operand)` to return true for a `kCopy` of an unrelated value and allowing the copy to be unsoundly elided. Conversely, if a `kCopy` genuinely copies a value from `region` through more than one `kGetTupleElement` or through a tuple parameter that has no other direct uses in `region`, `region.contains(operand)` returns false and unnecessarily retains the copy.

Because `LiveRangeRegions` records every direct use of a value in `region` (with a non-null `value_definition` pointing to the source `HloInstruction`), a `kCopy` copies a value from `region` if and only if `region.contains(instr)` is true with a non-null `value_definition`. Check `region.contains(instr)` and `value_definition != nullptr` instead of inspecting `operand`.

PiperOrigin-RevId: 1000343848
---
 .../xla/xla/service/copy_insertion_test.cc    | 105 ++++++++++++++++++
 third_party/xla/xla/service/copy_removal.cc   |  13 +--
 2 files changed, 110 insertions(+), 8 deletions(-)

diff --git a/third_party/xla/xla/service/copy_insertion_test.cc b/third_party/xla/xla/service/copy_insertion_test.cc
index 0c20d1bfce6..469505e99c7 100644
--- a/third_party/xla/xla/service/copy_insertion_test.cc
+++ b/third_party/xla/xla/service/copy_insertion_test.cc
@@ -6136,6 +6136,111 @@ ENTRY main {
   }
 }
 
+TEST_F(CopyInsertionTest,
+       ConditionalBranchMustCopyWhenOperandReadAfterConditional) {
+  absl::string_view hlo_string = R"(
+HloModule ConditionalOperandReadAfter
+
+branch_0 {
+  p0 = (f32[4]) parameter(0)
+  v0 = f32[4] get-tuple-element(p0), index=0
+  bar = f32[4] opt-barrier(v0)
+  ROOT r0 = (f32[4]) tuple(bar)
+}
+
+branch_1 {
+  p1 = (f32[4]) parameter(0)
+  v1 = f32[4] get-tuple-element(p1), index=0
+  neg = f32[4] negate(v1)
+  ROOT r1 = (f32[4]) tuple(neg)
+}
+
+ENTRY main {
+  cond_pred = pred[] parameter(0)
+  x = f32[4] parameter(1)
+  e = f32[4] negate(x)
+  t = (f32[4]) tuple(e)
+  cond = (f32[4]) conditional(cond_pred, t, t),
+    true_computation=branch_0, false_computation=branch_1
+  k = f32[4] get-tuple-element(cond), index=0
+  ROOT out = f32[4] add(e, k)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_string));
+  CopyInsertion copy_insertion(&alias_info_,
+                               /*use_region_based_live_range_analysis=*/-1);
+  ASSERT_THAT(copy_insertion.Run(module.get()), ::absl_testing::IsOk());
+
+  HloComputation* branch_0 = module->GetComputationWithName("branch_0");
+  HloComputation* branch_1 = module->GetComputationWithName("branch_1");
+  ASSERT_NE(branch_0, nullptr);
+  ASSERT_NE(branch_1, nullptr);
+  EXPECT_EQ(CountCopies(*branch_0), 1);
+  EXPECT_EQ(CountCopies(*branch_1), 0);
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloAliasAnalysis> alias_analysis,
+                       HloAliasAnalysis::Run(module.get(), &alias_info_));
+  const HloInstruction* e = FindInstruction(module.get(), "e");
+  const HloInstruction* cond = FindInstruction(module.get(), "cond");
+  ASSERT_NE(e, nullptr);
+  ASSERT_NE(cond, nullptr);
+  EXPECT_NE(&alias_analysis->GetUniqueBufferAt(e),
+            &alias_analysis->GetUniqueBufferAt(cond, {0}));
+}
+
+TEST_F(CopyInsertionTest,
+       ConditionalAllPassthroughBranchesElideCopiesWhenOperandReadAfter) {
+  // Both branches pass `e` through unchanged via a tuple parameter and GTE, and
+  // `e` is also read after the conditional. Because the GTE operand (`p0`/`p1`)
+  // is a kParameter, it is not in `src_live_range`, whereas the inserted branch
+  // root copies themselves are recorded as users in `src_live_range`. Checking
+  // `region.contains(instr)` (without requiring `region.contains(operand)`)
+  // allows both branch copies to be elided.
+  absl::string_view hlo_string = R"(
+HloModule ConditionalAllPassthroughOperandReadAfter
+
+branch_0 {
+  p0 = (f32[4]) parameter(0)
+  v0 = f32[4] get-tuple-element(p0), index=0
+  ROOT r0 = (f32[4]) tuple(v0)
+}
+
+branch_1 {
+  p1 = (f32[4]) parameter(0)
+  v1 = f32[4] get-tuple-element(p1), index=0
+  ROOT r1 = (f32[4]) tuple(v1)
+}
+
+ENTRY main {
+  cond_pred = pred[] parameter(0)
+  x = f32[4] parameter(1)
+  e = f32[4] negate(x)
+  t = (f32[4]) tuple(e)
+  cond = (f32[4]) conditional(cond_pred, t, t),
+    true_computation=branch_0, false_computation=branch_1
+  k = f32[4] get-tuple-element(cond), index=0
+  ROOT out = f32[4] add(e, k)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_string));
+  CopyInsertion copy_insertion(&alias_info_,
+                               /*use_region_based_live_range_analysis=*/-1);
+  ASSERT_THAT(copy_insertion.Run(module.get()), ::absl_testing::IsOk());
+
+  EXPECT_EQ(CountCopies(*module), 0);
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloAliasAnalysis> alias_analysis,
+                       HloAliasAnalysis::Run(module.get(), &alias_info_));
+  const HloInstruction* e = FindInstruction(module.get(), "e");
+  const HloInstruction* cond = FindInstruction(module.get(), "cond");
+  ASSERT_NE(e, nullptr);
+  ASSERT_NE(cond, nullptr);
+  EXPECT_EQ(&alias_analysis->GetUniqueBufferAt(e),
+            &alias_analysis->GetUniqueBufferAt(cond, {0}));
+}
+
 INSTANTIATE_TEST_SUITE_P(CondOrder, CopyInsertionCondOrderTest,
                          ::testing::Combine(::testing::Bool(),
                                             ::testing::Values(int64_t{0},
diff --git a/third_party/xla/xla/service/copy_removal.cc b/third_party/xla/xla/service/copy_removal.cc
index 123df196121..33dcb85e895 100644
--- a/third_party/xla/xla/service/copy_removal.cc
+++ b/third_party/xla/xla/service/copy_removal.cc
@@ -464,14 +464,11 @@ bool ComputeRelativeLocation::InstructionCanIntercept(
     // it does not overwrite the combined buffer with new values.
     case HloOpcode::kCopy: {
       // Checking the copy simply copies from the other live range with no
-      // layout conflicts.
-      HloInstruction* operand = instr->mutable_operand(0);
-      if (operand->opcode() == HloOpcode::kGetTupleElement) {
-        // kGetTupleElement only creates an alias among HloValues and is not
-        // included in the live range region. We check its operand instead.
-        operand = operand->mutable_operand(0);
-      }
-      if (region.contains(operand) &&
+      // layout conflicts. If `instr` copies a value from `region`, `instr`
+      // itself is recorded as a use (with non-null `value_definition`) in
+      // `region`.
+      if (region.contains(instr) &&
+          region[instr->parent()].at(instr).value_definition != nullptr &&
           ShapeUtil::Equal(instr->shape(), instr->operand(0)->shape())) {
         return false;  // Cannot intercept.
       }

From b5f33859cb1b389d046f57ebc1d0ae184f259547 Mon Sep 17 00:00:00 2001
From: Bill Varcho <varcho@google.com>
Date: Thu, 15 Oct 2026 13:17:45 -0700
Subject: [PATCH 0905/1000] [SDY][cleanup] Replace sdy maximal mesh terminology
 with single_device under `tensorflow/` .

PiperOrigin-RevId: 1000345482
---
 third_party/xla/xla/hlo/ir/mesh_and_axis.h    |  3 --
 .../xla/xla/hlo/ir/mesh_and_axis_test.cc      |  7 +--
 third_party/xla/xla/hlo/ir/named_sharding.h   |  2 +-
 .../xla/xla/hlo/ir/named_sharding_test.cc     |  6 +--
 third_party/xla/xla/hlo/parser/hlo_parser.cc  |  4 +-
 .../spmd/shardy/shardy_xla_pass_test.cc       | 12 ++---
 .../export_callback_custom_calls.h            |  2 +-
 .../stablehlo_round_trip/export_shardings.cc  | 15 +++---
 .../stablehlo_round_trip/shard_map_export.cc  |  2 +-
 .../stablehlo_round_trip/stablehlo_import.cc  | 46 +++++++++----------
 .../stablehlo_round_trip/stablehlo_import.h   |  2 +-
 .../stablehlo_import_test.cc                  | 45 +++++++++---------
 .../spmd/shardy/test/dedup_meshes.mlir        | 20 ++++----
 .../spmd/shardy/test/import_shardings.mlir    |  6 +--
 ..._trip_import_pipeline_hlo_sharding_v3.mlir | 22 ++++-----
 ...ound_trip_import_pipeline_no_sharding.mlir |  8 ++--
 ...rt_pipeline_sharding_in_frontend_attr.mlir | 20 ++++----
 .../test/stablehlo_export_pipeline.mlir       | 20 ++++----
 ...o_export_pipeline_simplify_replicated.mlir |  4 +-
 ...stablehlo_round_trip_shard_map_export.mlir |  4 +-
 .../xla/xla/service/spmd/shardy/utils.cc      |  4 +-
 21 files changed, 125 insertions(+), 129 deletions(-)

diff --git a/third_party/xla/xla/hlo/ir/mesh_and_axis.h b/third_party/xla/xla/hlo/ir/mesh_and_axis.h
index 8d18c6e85da..70dcc4c0ec1 100644
--- a/third_party/xla/xla/hlo/ir/mesh_and_axis.h
+++ b/third_party/xla/xla/hlo/ir/mesh_and_axis.h
@@ -84,9 +84,6 @@ class Mesh {
     return axes_names_.empty() && device_assignment_.num_elements() == 1;
   }
 
-  // Backward compatibility alias for `IsSingleDevice()`.
-  bool IsMaximal() const { return IsSingleDevice(); }
-
   bool operator==(const Mesh& other) const {
     return device_assignment_ == other.device_assignment_ &&
            axes_names_ == other.axes_names_;
diff --git a/third_party/xla/xla/hlo/ir/mesh_and_axis_test.cc b/third_party/xla/xla/hlo/ir/mesh_and_axis_test.cc
index 0ab1571d1b4..ccaf054f0b6 100644
--- a/third_party/xla/xla/hlo/ir/mesh_and_axis_test.cc
+++ b/third_party/xla/xla/hlo/ir/mesh_and_axis_test.cc
@@ -191,7 +191,7 @@ TEST(MeshAndAxisTest, ValidatesAxisRef) {
 
 TEST(MeshAndAxisTest, ValidatesMeshEmptyMesh) { Mesh(); }
 
-TEST(MeshAndAxisTest, ValidatesMeshMaximalMesh) { Mesh(5); }
+TEST(MeshAndAxisTest, ValidatesMeshSingleDeviceMesh) { Mesh(5); }
 
 TEST(MeshAndAxisTest, ValidatesMesh) {
   EXPECT_DEATH(
@@ -358,7 +358,7 @@ TEST(MeshAndAxisTest, EmptyMesh) {
   EXPECT_EQ(empty_mesh, Mesh());
   EXPECT_NE(empty_mesh, Mesh(5));
   EXPECT_NE(empty_mesh, Mesh({1}, {"a"}));
-  EXPECT_FALSE(empty_mesh.IsMaximal());
+  EXPECT_FALSE(empty_mesh.IsSingleDevice());
   EXPECT_THAT(empty_mesh.ToProto(), EqualsProto(MeshProto()));
   EXPECT_EQ(empty_mesh, Mesh::FromProto(MeshProto()));
   EXPECT_EQ(empty_mesh, Mesh::FromProto(empty_mesh.ToProto()));
@@ -367,13 +367,10 @@ TEST(MeshAndAxisTest, EmptyMesh) {
 TEST(MeshAndAxisTest, SingleDeviceMesh) {
   Mesh single_device_mesh(5);
   EXPECT_TRUE(single_device_mesh.IsSingleDevice());
-  EXPECT_TRUE(single_device_mesh.IsMaximal());
   Mesh non_single_device_mesh({2, 3}, {"a", "b"});
   EXPECT_FALSE(non_single_device_mesh.IsSingleDevice());
-  EXPECT_FALSE(non_single_device_mesh.IsMaximal());
   Mesh mesh_single_axis({1}, {"a"});
   EXPECT_FALSE(mesh_single_axis.IsSingleDevice());
-  EXPECT_FALSE(mesh_single_axis.IsMaximal());
 
   EXPECT_EQ(single_device_mesh, Mesh(5));
   EXPECT_NE(single_device_mesh, Mesh(6));
diff --git a/third_party/xla/xla/hlo/ir/named_sharding.h b/third_party/xla/xla/hlo/ir/named_sharding.h
index aeaeef6ab5e..112a87d44e8 100644
--- a/third_party/xla/xla/hlo/ir/named_sharding.h
+++ b/third_party/xla/xla/hlo/ir/named_sharding.h
@@ -172,7 +172,7 @@ class NamedSharding {
            unreduced_axes_.empty() && manual_axes_.empty();
   }
 
-  bool IsSingleDevice() const { return mesh_.IsMaximal(); }
+  bool IsSingleDevice() const { return mesh_.IsSingleDevice(); }
 
   // This checks for replicated or single-device sharding.
   bool IsReplicatedOrSingleDevice() const {
diff --git a/third_party/xla/xla/hlo/ir/named_sharding_test.cc b/third_party/xla/xla/hlo/ir/named_sharding_test.cc
index 464935986a6..72b18ff69e4 100644
--- a/third_party/xla/xla/hlo/ir/named_sharding_test.cc
+++ b/third_party/xla/xla/hlo/ir/named_sharding_test.cc
@@ -574,8 +574,8 @@ TEST(NamedShardingTest, NumDevices) {
   NamedSharding sharding(mesh, {});
   EXPECT_EQ(sharding.num_devices(), 2 * 4 * 3 * 8);
 
-  Mesh maximal_mesh(5);
-  NamedSharding single_device_sharding(maximal_mesh);
+  Mesh single_device_mesh(5);
+  NamedSharding single_device_sharding(single_device_mesh);
   EXPECT_EQ(single_device_sharding.num_devices(), 1);
 
   Mesh empty_mesh;
@@ -863,7 +863,7 @@ TEST(NamedShardingTest, NamedShardingProtoConversion) {
   EXPECT_EQ(sharding, from_proto);
 }
 
-TEST(NamedShardingPredicatesTest, HasPartialReplication_Maximal) {
+TEST(NamedShardingPredicatesTest, HasPartialReplication_SingleDevice) {
   EXPECT_FALSE(NamedSharding(Mesh(5)).HasPartialReplication());
 }
 
diff --git a/third_party/xla/xla/hlo/parser/hlo_parser.cc b/third_party/xla/xla/hlo/parser/hlo_parser.cc
index 63ae942c349..eddac0539d9 100644
--- a/third_party/xla/xla/hlo/parser/hlo_parser.cc
+++ b/third_party/xla/xla/hlo/parser/hlo_parser.cc
@@ -4634,8 +4634,8 @@ bool HloParserImpl::ParseSingleSharding(
     return true;
   };
 
-  // Maximal sharding
-  if (mesh->IsMaximal()) {
+  // Single-device sharding
+  if (mesh->IsSingleDevice()) {
     std::vector<OpMetadata> metadata;
     if (!parse_metadata_if_present(metadata)) {
       return false;
diff --git a/third_party/xla/xla/service/spmd/shardy/shardy_xla_pass_test.cc b/third_party/xla/xla/service/spmd/shardy/shardy_xla_pass_test.cc
index d72aec0a28f..c08d8cd499f 100644
--- a/third_party/xla/xla/service/spmd/shardy/shardy_xla_pass_test.cc
+++ b/third_party/xla/xla/service/spmd/shardy/shardy_xla_pass_test.cc
@@ -1168,16 +1168,16 @@ TEST_P(ShardyXLATest, ShardyGenerateDeviceLocalCodeCanonicalizeLayout) {
             "f32[8,16]{0,1:T(8,128)})");
 }
 
-TEST_P(ShardyXLATest, TestMaximalShardingNoResults) {
+TEST_P(ShardyXLATest, TestSingleDeviceShardingNoResults) {
   const char* const hloString = R"(
-HloModule maximal_sharding_module, entry_computation_layout={(s64[2]{0})->s64[2]{0}},
-    frontend_attributes={xla.sdy.meshes={maximal_mesh_0 = #sdy.mesh<[], device_ids=[0]>}}
+HloModule single_device_sharding_module, entry_computation_layout={(s64[2]{0})->s64[2]{0}},
+    frontend_attributes={xla.sdy.meshes={single_device_0 = #sdy.mesh<[], device_ids=[0]>}}
 
 ENTRY %main.0 (Arg_0.0: s64[2]) -> s64[2] {
   ROOT %Arg_0.0 = s64[2] parameter(0)
   %custom-call.0 = () custom-call(s64[2] %Arg_0.0), custom_call_target="xla_ffi_python_cpu_callback",
       operand_layout_constraints={s64[2]{0}}, custom_call_has_side_effect=true, api_version=API_VERSION_TYPED_FFI,
-      frontend_attributes={xla.sdy.sharding="#sdy.sharding_per_value<[<@maximal_mesh_0, []>]>"},
+      frontend_attributes={xla.sdy.sharding="#sdy.sharding_per_value<[<@single_device_0, []>]>"},
       sharding={{maximal device=0}}
 }
 )";
@@ -1249,7 +1249,7 @@ TEST_P(ShardyXLATest, WhileShardingOnlyOnFreeVariables) {
 
 TEST_P(ShardyXLATest, EmptyResultLayout) {
   const char* const hloString = R"(
-    HloModule pjit_f_, entry_computation_layout={(s64[2,2,2]{2,1,0})->()}, num_partitions=2, frontend_attributes={xla.sdy.meshes={maximal_mesh_0 = #sdy.mesh<[], device_ids=[0]>, mesh = #sdy.mesh<["x"=2]>}}
+    HloModule pjit_f_, entry_computation_layout={(s64[2,2,2]{2,1,0})->()}, num_partitions=2, frontend_attributes={xla.sdy.meshes={single_device_0 = #sdy.mesh<[], device_ids=[0]>, mesh = #sdy.mesh<["x"=2]>}}
 
     ENTRY %main.5 (Arg_0.1: s64[2,2,2]) -> () {
       %Arg_0.0 = s64[2,2,2]{2,1,0} parameter(0), sharding={devices=[2,1,1]<=[2]}, metadata={op_name="x"}
@@ -1265,7 +1265,7 @@ TEST_P(ShardyXLATest, EmptyResultLayout) {
 
 TEST_P(ShardyXLATest, EmptyOperandLayout) {
   const char* const hloString = R"(
-    HloModule pjit_f_, entry_computation_layout={()->s64[2,2]{1,0}}, num_partitions=2, frontend_attributes={xla.sdy.meshes={maximal_mesh_0 = #sdy.mesh<[], device_ids=[0]>, mesh = #sdy.mesh<["x"=2]>}}
+    HloModule pjit_f_, entry_computation_layout={()->s64[2,2]{1,0}}, num_partitions=2, frontend_attributes={xla.sdy.meshes={single_device_0 = #sdy.mesh<[], device_ids=[0]>, mesh = #sdy.mesh<["x"=2]>}}
 
     ENTRY %main.5 () -> s64[2,2] {
       ROOT %constant = s64[2,2]{1,0} constant({{1,1},{1,1}})
diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_callback_custom_calls.h b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_callback_custom_calls.h
index f50c2983851..da0c58f3752 100644
--- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_callback_custom_calls.h
+++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_callback_custom_calls.h
@@ -30,7 +30,7 @@ namespace sdy {
 // The rest of the XLA pipeline expects host callback custom calls to either be
 // a tuple with a get_tuple_element or no results (which we changed due to
 // shardy shardings expecting at least one result, and needing to attach a
-// maximal sharding to the callbacks).
+// single-device sharding to the callbacks).
 std::unique_ptr<mlir::Pass>
 createStablehloRoundTripExportCallbackCustomCallsPass();
 
diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc
index 4a4925565d5..a2d25143927 100644
--- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc
+++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_shardings.cc
@@ -388,11 +388,12 @@ HloSharding getHloShardingForOp(
     return manualAxes;
   };
 
-  bool isNoResultMaximal = op->getNumResults() == 0 && shardings.size() == 1 &&
-                           (getMeshAttr(shardings.front()).isMaximal() ||
-                            shardings.front().isFullyReplicated());
-  CHECK(shardings.size() == op->getNumResults() || isNoResultMaximal);
-  if (op->getNumResults() == 1 || isNoResultMaximal) {
+  bool isNoResultSingleDevice =
+      op->getNumResults() == 0 && shardings.size() == 1 &&
+      (getMeshAttr(shardings.front()).isSingleDevice() ||
+       shardings.front().isFullyReplicated());
+  CHECK(shardings.size() == op->getNumResults() || isNoResultSingleDevice);
+  if (op->getNumResults() == 1 || isNoResultSingleDevice) {
     ArrayRef<StringAttr> resultManualAxes =
         op->getNumResults() == 1
             ? getManualAxesForType(op->getResultTypes().front(),
@@ -453,7 +454,7 @@ NamedSharding convertToNamedSharding(
     return NamedSharding::Replicate();
   }
   // If there are no axes, convert to:
-  // - maximal sharding if the mesh has a device id
+  // - single-device sharding if the mesh has a device id
   // - else replicated sharding
   if (sdyMesh.getAxes().empty()) {
     if (!sdyMesh.getDeviceIds().empty()) {
@@ -559,7 +560,7 @@ HloSharding convertToHloSharding(
   MeshAttr mesh = getMeshAttr(sdySharding);
 
   // If there are no axes, convert to:
-  // - maximal sharding if the mesh has a device id
+  // - single-device sharding if the mesh has a device id
   // - else replicated sharding
   if (mesh.getAxes().empty()) {
     return mesh.getDeviceIds().empty()
diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc
index 71b6c94c5a3..6e01d92bf69 100644
--- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc
+++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/shard_map_export.cc
@@ -173,7 +173,7 @@ void setOpManualAxes(Operation* op, ManualAxesAttr manualAxes,
     if (opSharding.getMeshOrRef() != meshOrRef) {
       hasOtherMesh = true;
       MeshAttr otherMesh = opSharding.getMesh(op);
-      CHECK(otherMesh.getAxes().empty() || otherMesh.isMaximal());
+      CHECK(otherMesh.getAxes().empty() || otherMesh.isSingleDevice());
     }
   }
   if (hasOtherMesh) {
diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.cc b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.cc
index f321b8d56a1..a6bab4afbd9 100644
--- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.cc
+++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.cc
@@ -155,9 +155,9 @@ absl::flat_hash_set<xla::HloSharding> collectXlaHloShardings(
 
 struct MeshAxesAndIds {
   SmallVector<MeshAxisAttr> namedAxes;
-  // The device ids for maximal shardings sorted in ascending order
+  // The device ids for single-device shardings sorted in ascending order
   // without duplicates.
-  SmallVector<int64_t> maximalDeviceIds;
+  SmallVector<int64_t> singleDeviceIds;
   std::vector<int64_t> meshDeviceIds;
 };
 
@@ -165,7 +165,7 @@ MeshAxesAndIds findMeshAxesAndIds(
     const absl::flat_hash_set<xla::HloSharding>& oldShardings,
     mlir::MLIRContext* context) {
   MeshAxesAndIds result;
-  auto& [namedAxes, maximalDeviceIds, meshDeviceIds] = result;
+  auto& [namedAxes, singleDeviceIds, meshDeviceIds] = result;
 
   const xla::NamedSharding* firstNamedSharding = nullptr;
   for (const xla::HloSharding& hloSharding : oldShardings) {
@@ -192,12 +192,12 @@ MeshAxesAndIds findMeshAxesAndIds(
 
   // Find common axes of old shardings.
   SmallVector<int64_t> axes;
-  llvm::SmallDenseSet<int64_t> maximalDeviceIdSet;
+  llvm::SmallDenseSet<int64_t> singleDeviceIdSet;
   for (const xla::HloSharding& hloSharding : oldShardings) {
-    // If the sharding is a maximal sharding, we do not need to find common
-    // axes but add the device id to `deviceIdsForMaximalMesh`.
+    // If the sharding is a single-device sharding, we do not need to find
+    // common axes but add the device id to `deviceIdsForSingleDeviceMesh`.
     if (hloSharding.IsSingleDevice()) {
-      maximalDeviceIdSet.insert(hloSharding.GetUniqueDevice());
+      singleDeviceIdSet.insert(hloSharding.GetUniqueDevice());
       continue;
     }
 
@@ -243,13 +243,13 @@ MeshAxesAndIds findMeshAxesAndIds(
     }
   }
 
-  maximalDeviceIds = llvm::to_vector(maximalDeviceIdSet);
-  llvm::sort(maximalDeviceIds);
+  singleDeviceIds = llvm::to_vector(singleDeviceIdSet);
+  llvm::sort(singleDeviceIds);
   return result;
 }
 
 // Collect shardings with the attr name kXlaShardingAttr. Find common axes for
-// these shardings and device ids for maximal shardings.
+// these shardings and device ids for single-device shardings.
 MeshAxesAndIds findMeshAxesAndIds(ModuleOp moduleOp) {
   // Collect old shardings in the format of xla::HloSharding.
   const absl::flat_hash_set<xla::HloSharding> oldShardings =
@@ -323,7 +323,7 @@ std::string convertToSdySharding(const xla::OpSharding& opSharding,
 // `globalMesh`.
 TensorShardingAttr convertToSdySharding(
     const xla::HloSharding& hloSharding, MeshAttr globalMesh,
-    const SmallDenseMap<int64_t, StringRef>& deviceIdToMaximalMeshName,
+    const SmallDenseMap<int64_t, StringRef>& deviceIdToSingleDeviceMeshName,
     int64_t rank, bool openDims, bool inlineMesh) {
   mlir::MLIRContext* ctx = globalMesh.getContext();
 
@@ -406,9 +406,9 @@ TensorShardingAttr convertToSdySharding(
                                    reductionOp);
   }
 
-  // If the sharding is a maximal sharding, return a fully closed sharding.
-  // The exact sharding does not matter since the tensor can only exist on one
-  // device.
+  // If the sharding is a single-device sharding, return a fully closed
+  // sharding. The exact sharding does not matter since the tensor can only
+  // exist on one device.
   if (hloSharding.IsSingleDevice()) {
     if (inlineMesh) {
       return TensorShardingAttr::getFullyClosed(
@@ -417,7 +417,7 @@ TensorShardingAttr convertToSdySharding(
     }
     return TensorShardingAttr::getFullyClosed(
         ctx, /*rank=*/0,
-        deviceIdToMaximalMeshName.lookup(hloSharding.GetUniqueDevice()));
+        deviceIdToSingleDeviceMeshName.lookup(hloSharding.GetUniqueDevice()));
   }
   CHECK(!hloSharding.IsTuple());
 
@@ -592,7 +592,7 @@ bool shouldOpenDims(ArrayRef<bool> allowPropagationToTensors, int64_t index) {
 // Convert the shardings in `funcOp` from kXlaShardingAttr into kShardingAttr.
 LogicalResult importShardings(
     FuncOp funcOp, MeshAttr globalMesh,
-    const SmallDenseMap<int64_t, StringRef>& deviceIdToMaximalMeshName,
+    const SmallDenseMap<int64_t, StringRef>& deviceIdToSingleDeviceMeshName,
     ArrayRef<bool> allowPropagationToArgs,
     ArrayRef<bool> allowPropagationToResults, bool inlineMesh) {
   for (auto [argNum, argType] : llvm::enumerate(funcOp.getArgumentTypes())) {
@@ -602,7 +602,7 @@ LogicalResult importShardings(
           argNum, kShardingAttr,
           convertToSdySharding(
               parseShardingFromString(oldSharding), globalMesh,
-              deviceIdToMaximalMeshName, mlir::sdy::getTensorRank(argType),
+              deviceIdToSingleDeviceMeshName, mlir::sdy::getTensorRank(argType),
               shouldOpenDims(allowPropagationToArgs, argNum), inlineMesh));
       funcOp.removeArgAttr(argNum, kXlaShardingAttr);
     }
@@ -615,7 +615,7 @@ LogicalResult importShardings(
           resNum, kShardingAttr,
           convertToSdySharding(
               parseShardingFromString(oldSharding), globalMesh,
-              deviceIdToMaximalMeshName, mlir::sdy::getTensorRank(resType),
+              deviceIdToSingleDeviceMeshName, mlir::sdy::getTensorRank(resType),
               shouldOpenDims(allowPropagationToResults, resNum), inlineMesh));
       funcOp.removeResultAttr(
           resNum, StringAttr::get(funcOp.getContext(), kXlaShardingAttr));
@@ -633,14 +633,14 @@ LogicalResult importShardings(
       if (op->getNumResults() == 0) {
         const xla::HloSharding& leafSharding = flatHloSharding.front();
         newShardings.push_back(convertToSdySharding(
-            leafSharding, globalMesh, deviceIdToMaximalMeshName,
+            leafSharding, globalMesh, deviceIdToSingleDeviceMeshName,
             /*rank=*/0, /*openDims=*/false, inlineMesh));
       } else {
         newShardings.reserve(op->getNumResults());
         for (const auto& [resHloSharding, resType] :
              llvm::zip_equal(flatHloSharding, op->getResultTypes())) {
           newShardings.push_back(convertToSdySharding(
-              resHloSharding, globalMesh, deviceIdToMaximalMeshName,
+              resHloSharding, globalMesh, deviceIdToSingleDeviceMeshName,
               mlir::sdy::getTensorRank(resType),
               /*openDims=*/false, inlineMesh));
         }
@@ -667,9 +667,9 @@ class ImportShardingsPass
   void runOnOperation() final {
     ModuleOp moduleOp = getOperation();
 
-    auto [namedAxes, deviceIdsForMaximalMesh, meshDeviceIds] =
+    auto [namedAxes, deviceIdsForSingleDeviceMesh, meshDeviceIds] =
         findMeshAxesAndIds(moduleOp);
-    if (namedAxes.empty() && deviceIdsForMaximalMesh.empty()) {
+    if (namedAxes.empty() && deviceIdsForSingleDeviceMesh.empty()) {
       // There are no shardings in the `moduleOp`.
       return;
     }
@@ -688,7 +688,7 @@ class ImportShardingsPass
                                       kGlobalMeshName, globalMesh));
 
     SmallDenseMap<int64_t, StringRef> deviceIdToSingleDeviceMeshName;
-    for (int64_t deviceId : deviceIdsForMaximalMesh) {
+    for (int64_t deviceId : deviceIdsForSingleDeviceMesh) {
       // Create a mesh name with its deviceId as a suffix for each single-device
       // mesh.
       std::string meshName = absl::StrCat("single_device_", deviceId);
diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.h b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.h
index 9f0c2c3a3c5..44fd0c5b5dc 100644
--- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.h
+++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import.h
@@ -44,7 +44,7 @@ xla::HloSharding parseShardingFromString(const mlir::StringAttr& sharding);
 mlir::sdy::TensorShardingAttr convertToSdySharding(
     const xla::HloSharding& hloSharding, mlir::sdy::MeshAttr globalMesh,
     const llvm::SmallDenseMap<int64_t, mlir::StringRef>&
-        deviceIdToMaximalMeshName,
+        deviceIdToSingleDeviceMeshName,
     int64_t rank, bool openDims = false, bool inlineMesh = false);
 // Same as above, but takes an `xla::OpSharding` proto and uses fake mesh from
 // given `opSharding`. Returns a string representation of `TensorShardingAttr`
diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import_test.cc b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import_test.cc
index 697646a6769..7b92a62362e 100644
--- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import_test.cc
+++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/stablehlo_import_test.cc
@@ -43,7 +43,7 @@ TEST(StablehloImportTest, FullyReplicatedEmptyMesh) {
   TensorShardingAttr sharding = xla::sdy::convertToSdySharding(
       /*hloSharding=*/xla::HloSharding::Replicate(),
       /*globalMesh=*/mesh,
-      /*deviceIdToMaximalMeshName=*/
+      /*deviceIdToSingleDeviceMeshName=*/
       llvm::SmallDenseMap<int64_t, mlir::StringRef>(), /*rank=*/2,
       /*openDims=*/true);
   EXPECT_EQ(attributeToString(sharding), "#sdy.sharding<@mesh, [{?}, {?}]>");
@@ -60,7 +60,7 @@ TEST(StablehloImportTest, FullyReplicatedNonEmptyMesh) {
   TensorShardingAttr sharding = xla::sdy::convertToSdySharding(
       /*hloSharding=*/xla::HloSharding::Replicate(),
       /*globalMesh=*/mesh,
-      /*deviceIdToMaximalMeshName=*/
+      /*deviceIdToSingleDeviceMeshName=*/
       llvm::SmallDenseMap<int64_t, mlir::StringRef>(), /*rank=*/2,
       /*openDims=*/true);
   EXPECT_EQ(attributeToString(sharding), "#sdy.sharding<@mesh, [{?}, {?}]>");
@@ -78,7 +78,7 @@ TEST(StablehloImportTest, SkipFirstAxisOfSize1) {
   TensorShardingAttr sharding = xla::sdy::convertToSdySharding(
       /*hloSharding=*/xla::HloSharding::IotaTile({4, 2}),
       /*globalMesh=*/mesh,
-      /*deviceIdToMaximalMeshName=*/
+      /*deviceIdToSingleDeviceMeshName=*/
       llvm::SmallDenseMap<int64_t, mlir::StringRef>(), /*rank=*/2,
       /*openDims=*/true);
   EXPECT_EQ(attributeToString(sharding),
@@ -98,7 +98,7 @@ TEST(StablehloImportTest, SkipSecondAxisOfSize1) {
   TensorShardingAttr sharding = xla::sdy::convertToSdySharding(
       /*hloSharding=*/xla::HloSharding::IotaTile({4, 2}),
       /*globalMesh=*/mesh,
-      /*deviceIdToMaximalMeshName=*/
+      /*deviceIdToSingleDeviceMeshName=*/
       llvm::SmallDenseMap<int64_t, mlir::StringRef>(), /*rank=*/2,
       /*openDims=*/true);
   EXPECT_EQ(attributeToString(sharding),
@@ -118,7 +118,7 @@ TEST(StablehloImportTest, TransposedWithReplicatedAxis) {
       /*hloSharding=*/xla::HloSharding::PartialTile(
           xla::TileAssignment({2, 2, 4}, {2, 4, 2}, {2, 0, 1})),
       /*globalMesh=*/mesh,
-      /*deviceIdToMaximalMeshName=*/
+      /*deviceIdToSingleDeviceMeshName=*/
       llvm::SmallDenseMap<int64_t, mlir::StringRef>(), /*rank=*/2,
       /*openDims=*/false);
   EXPECT_EQ(attributeToString(sharding),
@@ -136,7 +136,7 @@ TEST(StablehloImportTest, IotaTileRequiresSubAxes) {
   TensorShardingAttr sharding = xla::sdy::convertToSdySharding(
       /*hloSharding=*/xla::HloSharding::IotaTile({2, 4}),
       /*globalMesh=*/mesh,
-      /*deviceIdToMaximalMeshName=*/
+      /*deviceIdToSingleDeviceMeshName=*/
       llvm::SmallDenseMap<int64_t, mlir::StringRef>(), /*rank=*/2,
       /*openDims=*/false);
   EXPECT_EQ(attributeToString(sharding),
@@ -153,7 +153,7 @@ TEST(StablehloImportTest, IotaTileRequiresSubAxes2) {
   TensorShardingAttr sharding = xla::sdy::convertToSdySharding(
       /*hloSharding=*/xla::HloSharding::IotaTile({4, 4}),
       /*globalMesh=*/mesh,
-      /*deviceIdToMaximalMeshName=*/
+      /*deviceIdToSingleDeviceMeshName=*/
       llvm::SmallDenseMap<int64_t, mlir::StringRef>(), /*rank=*/2,
       /*openDims=*/false);
   EXPECT_EQ(attributeToString(sharding),
@@ -172,7 +172,7 @@ TEST(StablehloImportTest, TransposedWithReplicatedRequiresSubAxes) {
       /*hloSharding=*/xla::HloSharding::PartialTile(
           xla::TileAssignment({2, 2, 2}, {2, 2, 2}, {2, 0, 1})),
       /*globalMesh=*/mesh,
-      /*deviceIdToMaximalMeshName=*/
+      /*deviceIdToSingleDeviceMeshName=*/
       llvm::SmallDenseMap<int64_t, mlir::StringRef>(), /*rank=*/2,
       /*openDims=*/false);
   EXPECT_EQ(attributeToString(sharding),
@@ -299,31 +299,32 @@ TEST(StablehloImportTest, NamedShardingEmptyDimShardings) {
   axes.emplace_back(mlir::sdy::MeshAxisAttr::get(&context, "x", 2));
   axes.emplace_back(mlir::sdy::MeshAxisAttr::get(&context, "y", 2));
   auto mesh = sdy::MeshAttr::get(&context, axes);
-  llvm::SmallDenseMap<int64_t, mlir::StringRef> emptyMaximalMap;
+  llvm::SmallDenseMap<int64_t, mlir::StringRef> emptySingleDeviceMap;
 
   xla::HloSharding replicated =
       xla::ParseSharding("{mesh[], replicated}").value();
-  EXPECT_EQ(
-      attributeToString(xla::sdy::convertToSdySharding(
-          replicated, mesh, emptyMaximalMap, /*rank=*/2, /*openDims=*/false)),
-      "#sdy.sharding<@mesh, [{}, {}]>");
-  EXPECT_EQ(
-      attributeToString(xla::sdy::convertToSdySharding(
-          replicated, mesh, emptyMaximalMap, /*rank=*/2, /*openDims=*/true)),
-      "#sdy.sharding<@mesh, [{?}, {?}]>");
+  EXPECT_EQ(attributeToString(xla::sdy::convertToSdySharding(
+                replicated, mesh, emptySingleDeviceMap, /*rank=*/2,
+                /*openDims=*/false)),
+            "#sdy.sharding<@mesh, [{}, {}]>");
+  EXPECT_EQ(attributeToString(xla::sdy::convertToSdySharding(
+                replicated, mesh, emptySingleDeviceMap, /*rank=*/2,
+                /*openDims=*/true)),
+            "#sdy.sharding<@mesh, [{?}, {?}]>");
 
   xla::HloSharding manual =
       xla::ParseSharding("{mesh['x'=2, 'y'=2], manual}").value();
   EXPECT_EQ(attributeToString(xla::sdy::convertToSdySharding(
-                manual, mesh, emptyMaximalMap, /*rank=*/2, /*openDims=*/false)),
+                manual, mesh, emptySingleDeviceMap, /*rank=*/2,
+                /*openDims=*/false)),
             "#sdy.sharding<@mesh, [{}, {}]>");
 
   xla::HloSharding unreduced =
       xla::ParseSharding("{mesh['x'=2, 'y'=2], unreduced}").value();
-  EXPECT_EQ(
-      attributeToString(xla::sdy::convertToSdySharding(
-          unreduced, mesh, emptyMaximalMap, /*rank=*/2, /*openDims=*/false)),
-      "#sdy.sharding<@mesh, [{}, {}], unreduced={\"x\", \"y\"}>");
+  EXPECT_EQ(attributeToString(xla::sdy::convertToSdySharding(
+                unreduced, mesh, emptySingleDeviceMap, /*rank=*/2,
+                /*openDims=*/false)),
+            "#sdy.sharding<@mesh, [{}, {}], unreduced={\"x\", \"y\"}>");
 }
 
 }  // namespace
diff --git a/third_party/xla/xla/service/spmd/shardy/test/dedup_meshes.mlir b/third_party/xla/xla/service/spmd/shardy/test/dedup_meshes.mlir
index 2b416b28821..d16715b4f20 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/dedup_meshes.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/dedup_meshes.mlir
@@ -36,10 +36,10 @@ sdy.mesh @meshA = <["a"=4, "b"=4]>
 sdy.mesh @mesh_one_device_before_empty = <["x"=1, "y"=1]>
 sdy.mesh @empty_mesh1 = <[]>
 sdy.mesh @empty_mesh2 = <[]>
-// CHECK:     sdy.mesh @maximal_mesh1 = <[], device_ids=[0]>
-// CHECK-NOT: sdy.mesh @maximal_mesh2 = <[], device_ids=[0]>
-sdy.mesh @maximal_mesh1 = <[], device_ids=[0]>
-sdy.mesh @maximal_mesh2 = <[], device_ids=[0]>
+// CHECK:     sdy.mesh @single_device_mesh1 = <[], device_ids=[0]>
+// CHECK-NOT: sdy.mesh @single_device_mesh2 = <[], device_ids=[0]>
+sdy.mesh @single_device_mesh1 = <[], device_ids=[0]>
+sdy.mesh @single_device_mesh2 = <[], device_ids=[0]>
 
 // CHECK-NOT: sdy.mesh @meshC = <["x"=4]>
 // CHECK:     sdy.mesh @meshD = <["a"=2, "b"=2]>
@@ -136,12 +136,12 @@ func.func @empty_mesh(
   return %arg0 : tensor<8x8xf32>
 }
 
-// CHECK-LABEL: @maximal_mesh
-// CHECK-SAME:  (%arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@maximal_mesh1, []>})
-// CHECK-SAME:  -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@maximal_mesh1, []>}) {
-func.func @maximal_mesh(
-  %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@maximal_mesh1, []>}
-) -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@maximal_mesh2, []>}) {
+// CHECK-LABEL: @single_device_mesh
+// CHECK-SAME:  (%arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@single_device_mesh1, []>})
+// CHECK-SAME:  -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@single_device_mesh1, []>}) {
+func.func @single_device_mesh(
+  %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@single_device_mesh1, []>}
+) -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@single_device_mesh2, []>}) {
   return %arg0 : tensor<8x8xf32>
 }
 
diff --git a/third_party/xla/xla/service/spmd/shardy/test/import_shardings.mlir b/third_party/xla/xla/service/spmd/shardy/test/import_shardings.mlir
index c579b8e38f0..9a4305d5f67 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/import_shardings.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/import_shardings.mlir
@@ -143,10 +143,10 @@ func.func @unknown_sharding(%arg0: tensor<8x8xf32> {mhlo.sharding = "{devices=[4
 // CHECK-LABEL: sdy.mesh @mesh = <[]>
 // CHECK-LABEL: sdy.mesh @single_device_0 = <[], device_ids=[0]>
 
-// CHECK-LABEL: func @one_maximal_mesh(
+// CHECK-LABEL: func @one_single_device_mesh(
 // CHECK-SAME:      %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@single_device_0, []>}
-func.func @one_maximal_mesh(%arg0: tensor<8x8xf32> {mhlo.sharding = "{maximal device=0}"},
-                            %arg1: tensor<8x8xf32>) -> tensor<8x8xf32> {
+func.func @one_single_device_mesh(%arg0: tensor<8x8xf32> {mhlo.sharding = "{maximal device=0}"},
+                                  %arg1: tensor<8x8xf32>) -> tensor<8x8xf32> {
   %0 = stablehlo.add %arg0, %arg1 : tensor<8x8xf32>
   return %0 : tensor<8x8xf32>
 }
diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir
index eeadf165681..dd79f764211 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_hlo_sharding_v3.mlir
@@ -244,13 +244,13 @@ module @module_1 {
       channel_handle = #stablehlo.channel_handle<handle = 1, type = 2>,
       is_host_transfer = true,
       mhlo.frontend_attributes = {baz = 1 : i32},
-      mhlo.sharding = "{maximal_mesh[device_id=0]}"
+      mhlo.sharding = "{single_device[device_id=0]}"
     } : (tensor<16x8xf32>, !stablehlo.token) -> !stablehlo.token
     %1:2 = "stablehlo.recv"(%0) {
       channel_handle = #stablehlo.channel_handle<handle = 1, type = 3>,
       is_host_transfer = true,
       mhlo.frontend_attributes = {baz = 1 : i32},
-      mhlo.sharding = "{{maximal_mesh[device_id=0]}, {maximal_mesh[device_id=0]}}"
+      mhlo.sharding = "{{single_device[device_id=0]}, {single_device[device_id=0]}}"
     } : (!stablehlo.token) -> (tensor<16x8xf32>, !stablehlo.token)
     %2 = stablehlo.add %1#0, %arg1 : tensor<16x8xf32>
     return %2 : tensor<16x8xf32>
@@ -259,9 +259,9 @@ module @module_1 {
 
 // -----
 
-module @maximal_sharding_module {
-  // CHECK-LABEL: @maximal_sharding_empty_tuple
-  func.func @maximal_sharding_empty_tuple(%arg0: tensor<2xi64>) -> tensor<2xi64> {
+module @single_device_sharding_module {
+  // CHECK-LABEL: @single_device_sharding_empty_tuple
+  func.func @single_device_sharding_empty_tuple(%arg0: tensor<2xi64>) -> tensor<2xi64> {
     // CHECK-NEXT: stablehlo.custom_call @xla_ffi_python_cpu_callback(%arg0) {
     // CHECK-SAME:   api_version = 4 : i32, backend_config = {descriptor = 126001424235520 : ui64},
     // CHECK-SAME:   has_side_effect = true, operand_layouts = [dense<0> : tensor<1xindex>], result_layouts = [],
@@ -271,14 +271,14 @@ module @maximal_sharding_module {
     %2 = stablehlo.custom_call @xla_ffi_python_cpu_callback(%arg0) {
       api_version = 4 : i32, backend_config = {descriptor = 126001424235520 : ui64},
       has_side_effect = true,
-      mhlo.sharding = "{maximal_mesh[device_id=0]}",
+      mhlo.sharding = "{single_device[device_id=0]}",
       operand_layouts = [dense<0> : tensor<1xindex>], result_layouts = [], xla_shape = "()"
     } : (tensor<2xi64>) -> tuple<>
     return %arg0 : tensor<2xi64>
   }
 
-  // CHECK-LABEL: @maximal_tuple_sharding_empty_tuple
-  func.func @maximal_tuple_sharding_empty_tuple(%arg0: tensor<2xi64>) -> tensor<2xi64> {
+  // CHECK-LABEL: @single_device_tuple_sharding_empty_tuple
+  func.func @single_device_tuple_sharding_empty_tuple(%arg0: tensor<2xi64>) -> tensor<2xi64> {
     // CHECK-NEXT: stablehlo.custom_call @xla_ffi_python_cpu_callback(%arg0) {
     // CHECK-SAME:   api_version = 4 : i32, backend_config = {descriptor = 126001424235520 : ui64},
     // CHECK-SAME:   has_side_effect = true, operand_layouts = [dense<0> : tensor<1xindex>], result_layouts = [],
@@ -288,15 +288,15 @@ module @maximal_sharding_module {
     %2 = stablehlo.custom_call @xla_ffi_python_cpu_callback(%arg0) {
       api_version = 4 : i32, backend_config = {descriptor = 126001424235520 : ui64},
       has_side_effect = true,
-      mhlo.sharding = "{{maximal_mesh[device_id=0]}}",
+      mhlo.sharding = "{{single_device[device_id=0]}}",
       operand_layouts = [dense<0> : tensor<1xindex>], result_layouts = [], xla_shape = "()"
     } : (tensor<2xi64>) -> tuple<>
     return %arg0 : tensor<2xi64>
   }
 
-  // CHECK-LABEL: func @result_with_maximal_sharding
+  // CHECK-LABEL: func @result_with_single_device_sharding
   // CHECK-SAME:    (%arg0: tensor<32xi32>) -> tensor<32xi32> {
-  func.func @result_with_maximal_sharding(%arg0: tensor<32xi32>) -> (tensor<32xi32> {mhlo.sharding = "{maximal_mesh[device_id=0]}"}) {
+  func.func @result_with_single_device_sharding(%arg0: tensor<32xi32>) -> (tensor<32xi32> {mhlo.sharding = "{single_device[device_id=0]}"}) {
     // CHECK-NEXT: return %arg0 : tensor<32xi32>
     return %arg0 : tensor<32xi32>
   }
diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_no_sharding.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_no_sharding.mlir
index 0e98bfff3ab..0b91539bc35 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_no_sharding.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_no_sharding.mlir
@@ -96,15 +96,15 @@ func.func @callback_no_result(%arg0: tensor<f64>) {
 
 // -----
 module @send_with_sdy_sharding_module {
-  // CHECK: sdy.mesh @maximal_mesh_0
-  sdy.mesh @maximal_mesh_0 = <[], device_ids=[0]>
+  // CHECK: sdy.mesh @single_device_0
+  sdy.mesh @single_device_0 = <[], device_ids=[0]>
   // CHECK-LABEL: func.func @send_with_sdy_sharding
   func.func @send_with_sdy_sharding(%arg0: tensor<i32>,
                                      %arg1: !stablehlo.token) -> !stablehlo.token {
-  // CHECK-NEXT: %0 = "stablehlo.send"(%arg0, %arg1) <{channel_handle = #stablehlo.channel_handle<handle = 1, type = 2>, is_host_transfer = true}> {mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "_host_callback_dtoh_0"}, sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh_0, []>]>, xla_shape = "token[]"}
+  // CHECK-NEXT: %0 = "stablehlo.send"(%arg0, %arg1) <{channel_handle = #stablehlo.channel_handle<handle = 1, type = 2>, is_host_transfer = true}> {mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "_host_callback_dtoh_0"}, sdy.sharding = #sdy.sharding_per_value<[<@single_device_0, []>]>, xla_shape = "token[]"}
     %1 = "stablehlo.send"(%arg0, %arg1) <{channel_handle = #stablehlo.channel_handle<handle = 1, type = 2>, is_host_transfer = true}>
       {mhlo.frontend_attributes = {_xla_host_transfer_handler_name = "tf_rendezvous", _xla_host_transfer_rendezvous = "_host_callback_dtoh_0"},
-      sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh_0, []>]>, xla_shape = "token[]"} : (tensor<i32>, !stablehlo.token) -> !stablehlo.token
+      sdy.sharding = #sdy.sharding_per_value<[<@single_device_0, []>]>, xla_shape = "token[]"} : (tensor<i32>, !stablehlo.token) -> !stablehlo.token
     return %1 : !stablehlo.token
   }
 }
diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_sharding_in_frontend_attr.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_sharding_in_frontend_attr.mlir
index 01d36fc862a..c797487cf83 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_sharding_in_frontend_attr.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_sharding_in_frontend_attr.mlir
@@ -16,7 +16,7 @@
 
 // CHECK-LABEL: module @multiple_func_result_shardings
 module @multiple_func_result_shardings attributes {mhlo.frontend_attributes = {xla.sdy.meshes =
-    "{mesh = #sdy.mesh<[\"a\"=8, \"b\"=8, \"c\"=8]>, mesh2 = #sdy.mesh<[\"a\"=1, \"b\"=4, \"c\"=1]>, maximal_mesh = #sdy.mesh<[], device_ids=[0]>}"}} {
+    "{mesh = #sdy.mesh<[\"a\"=8, \"b\"=8, \"c\"=8]>, mesh2 = #sdy.mesh<[\"a\"=1, \"b\"=4, \"c\"=1]>, single_device = #sdy.mesh<[], device_ids=[0]>}"}} {
   // CHECK: sdy.mesh @mesh = <["a"=8, "b"=8, "c"=8]>
   // CHECK: sdy.mesh @mesh2 = <["a"=1, "b"=4, "c"=1]>
 
@@ -259,19 +259,19 @@ module @multiple_func_result_shardings attributes {mhlo.frontend_attributes = {x
     %arg0: tensor<16x8xf32> {mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding<@mesh2, [{\"b\"}, {?}]>"}},
     %arg1: tensor<16x8xf32> {mhlo.frontend_attributes = {baz = 1 : i32, foo = "bar"}},
     %arg2: !stablehlo.token) -> tensor<16x8xf32> {
-    // CHECK-NEXT: %[[SEND:.*]] = "stablehlo.send"(%arg0, %arg2) <{channel_handle = #stablehlo.channel_handle<handle = 1, type = 2>, is_host_transfer = true}> {mhlo.frontend_attributes = {baz = 1 : i32}, sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh, []>]>} : (tensor<16x8xf32>, !stablehlo.token) -> !stablehlo.token
-    // CHECK-NEXT: %[[RECV:.*]]:2 = "stablehlo.recv"(%[[SEND]]) <{channel_handle = #stablehlo.channel_handle<handle = 1, type = 3>, is_host_transfer = true}> {mhlo.frontend_attributes = {baz = 1 : i32}, sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh, []>, <@maximal_mesh, []>]>} : (!stablehlo.token) -> (tensor<16x8xf32>, !stablehlo.token)
+    // CHECK-NEXT: %[[SEND:.*]] = "stablehlo.send"(%arg0, %arg2) <{channel_handle = #stablehlo.channel_handle<handle = 1, type = 2>, is_host_transfer = true}> {mhlo.frontend_attributes = {baz = 1 : i32}, sdy.sharding = #sdy.sharding_per_value<[<@single_device, []>]>} : (tensor<16x8xf32>, !stablehlo.token) -> !stablehlo.token
+    // CHECK-NEXT: %[[RECV:.*]]:2 = "stablehlo.recv"(%[[SEND]]) <{channel_handle = #stablehlo.channel_handle<handle = 1, type = 3>, is_host_transfer = true}> {mhlo.frontend_attributes = {baz = 1 : i32}, sdy.sharding = #sdy.sharding_per_value<[<@single_device, []>, <@single_device, []>]>} : (!stablehlo.token) -> (tensor<16x8xf32>, !stablehlo.token)
     // CHECK-NEXT: %[[ADD:.*]] = stablehlo.add %[[RECV]]#0, %arg1 : tensor<16x8xf32>
     // CHECK-NEXT: return %[[ADD]] : tensor<16x8xf32>
     %0 = "stablehlo.send"(%arg0, %arg2) {
       channel_handle = #stablehlo.channel_handle<handle = 1, type = 2>,
       is_host_transfer = true,
-      mhlo.frontend_attributes = {baz = 1 : i32, xla.sdy.sharding = "#sdy.sharding_per_value<[<@maximal_mesh, []>]>"}
+      mhlo.frontend_attributes = {baz = 1 : i32, xla.sdy.sharding = "#sdy.sharding_per_value<[<@single_device, []>]>"}
     } : (tensor<16x8xf32>, !stablehlo.token) -> !stablehlo.token
     %1:2 = "stablehlo.recv"(%0) {
       channel_handle = #stablehlo.channel_handle<handle = 1, type = 3>,
       is_host_transfer = true,
-      mhlo.frontend_attributes = {baz = 1 : i32, xla.sdy.sharding = "#sdy.sharding_per_value<[<@maximal_mesh, []>, <@maximal_mesh, []>]>"}
+      mhlo.frontend_attributes = {baz = 1 : i32, xla.sdy.sharding = "#sdy.sharding_per_value<[<@single_device, []>, <@single_device, []>]>"}
     } : (!stablehlo.token) -> (tensor<16x8xf32>, !stablehlo.token)
     %2 = stablehlo.add %1#0, %arg1 : tensor<16x8xf32>
     return %2 : tensor<16x8xf32>
@@ -280,18 +280,18 @@ module @multiple_func_result_shardings attributes {mhlo.frontend_attributes = {x
 
 // -----
 
-module @maximal_sharding_module attributes {mhlo.frontend_attributes = {xla.sdy.meshes = "{maximal_mesh_0 = #sdy.mesh<[], device_ids=[0]>}"}} {
-  // CHECK-LABEL: @maximal_sharding_empty_tuple
-  func.func @maximal_sharding_empty_tuple(%arg0: tensor<2xi64>) -> tensor<2xi64> {
+module @single_device_sharding_module attributes {mhlo.frontend_attributes = {xla.sdy.meshes = "{single_device_0 = #sdy.mesh<[], device_ids=[0]>}"}} {
+  // CHECK-LABEL: @single_device_sharding_empty_tuple
+  func.func @single_device_sharding_empty_tuple(%arg0: tensor<2xi64>) -> tensor<2xi64> {
     // CHECK-NEXT: stablehlo.custom_call @xla_ffi_python_cpu_callback(%arg0) {
     // CHECK-SAME:   api_version = 4 : i32, backend_config = {descriptor = 126001424235520 : ui64},
     // CHECK-SAME:   has_side_effect = true, operand_layouts = [dense<0> : tensor<1xindex>], result_layouts = [],
-    // CHECK-SAME:   sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh_0, []>]>, xla_shape = "()"
+    // CHECK-SAME:   sdy.sharding = #sdy.sharding_per_value<[<@single_device_0, []>]>, xla_shape = "()"
     // CHECK-SAME: } : (tensor<2xi64>) -> ()
     // CHECK-NEXT: return %arg0 : tensor<2xi64>
     %2 = stablehlo.custom_call @xla_ffi_python_cpu_callback(%arg0) {
       api_version = 4 : i32, backend_config = {descriptor = 126001424235520 : ui64},
-      has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@maximal_mesh_0, []>]>"},
+      has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@single_device_0, []>]>"},
       mhlo.sharding = "{{maximal device=0}}",
       operand_layouts = [dense<0> : tensor<1xindex>], result_layouts = [], xla_shape = "()"
     } : (tensor<2xi64>) -> tuple<>
diff --git a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_pipeline.mlir b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_pipeline.mlir
index 63d2f308741..3bedfc3380c 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_pipeline.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_pipeline.mlir
@@ -21,8 +21,8 @@ sdy.mesh @mesh_2 = <["x"=8, "y"=4]>
 sdy.mesh @mesh_3 = <["a"=2, "b"=2, "c"=2, "d"=2]>
 sdy.mesh @mesh_4 = <["axis_0"=2, "axis_1"=2, "axis_2"=2], device_ids=[0,2,4,6,1,3,5,7]>
 sdy.mesh @mesh_5 = <["i"=2, "j"=2]>
-sdy.mesh @maximal_mesh_0 = <[], device_ids=[0]>
-sdy.mesh @maximal_mesh_1 = <[], device_ids=[1]>
+sdy.mesh @single_device_0 = <[], device_ids=[0]>
+sdy.mesh @single_device_1 = <[], device_ids=[1]>
 sdy.mesh @empty_mesh_0 = <[]>
 sdy.mesh @empty_mesh_1 = <[]>
 
@@ -262,12 +262,12 @@ func.func @sharding_in_manual_computation_body(%arg0: tensor<8x16xf32> {sdy.shar
 
 // CHECK-V2-LABEL: func @mesh_with_device_id_should_be_converted_to_maximal_sharding(%arg0: tensor<8x8xf32> {mhlo.sharding = "{maximal device=0}"}, %arg1: tensor<8x8xf32>)
 // CHECK-V3-LABEL: func @mesh_with_device_id_should_be_converted_to_maximal_sharding(%arg0: tensor<8x8xf32> {mhlo.sharding = "{single_device[device_id=0]}"}, %arg1: tensor<8x8xf32>)
-func.func @mesh_with_device_id_should_be_converted_to_maximal_sharding(%arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@maximal_mesh_0, []>}, %arg1: tensor<8x8xf32>) -> tensor<8x8xf32> {
+func.func @mesh_with_device_id_should_be_converted_to_maximal_sharding(%arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@single_device_0, []>}, %arg1: tensor<8x8xf32>) -> tensor<8x8xf32> {
     // CHECK: %[[ADD:.*]] = stablehlo.add %arg0, %arg1
     %0 = stablehlo.add %arg0, %arg1 : tensor<8x8xf32>
     // CHECK-V2: %[[ADD_WITH_SHARDING:.*]] = stablehlo.add %[[ADD]], %[[ADD]] {mhlo.sharding = "{maximal device=1}"}
     // CHECK-V3: %[[ADD_WITH_SHARDING:.*]] = stablehlo.add %[[ADD]], %[[ADD]] {mhlo.sharding = "{single_device[device_id=1]}"}
-    %1 = stablehlo.add %0, %0 {sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh_1, []>]>} : tensor<8x8xf32>
+    %1 = stablehlo.add %0, %0 {sdy.sharding = #sdy.sharding_per_value<[<@single_device_1, []>]>} : tensor<8x8xf32>
     return %1 : tensor<8x8xf32>
 }
 
@@ -426,7 +426,7 @@ func.func public @callback_no_result(%arg0: tensor<f64>) {
   // CHECK-SAME:   operand_layouts = [dense<> : tensor<0xindex>, dense<> : tensor<0xindex>],
   // CHECK-SAME: } : (tensor<i64>, tensor<f64>) -> ()
   %c = stablehlo.constant dense<56238273106176> : tensor<i64>
-  stablehlo.custom_call @xla_python_cpu_callback(%c, %arg0) {api_version = 2 : i32, backend_config = "56238273106176", has_side_effect = true, operand_layouts = [dense<> : tensor<0xindex>, dense<> : tensor<0xindex>], result_layouts = [], sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh_0, []>]>} : (tensor<i64>, tensor<f64>) -> ()
+  stablehlo.custom_call @xla_python_cpu_callback(%c, %arg0) {api_version = 2 : i32, backend_config = "56238273106176", has_side_effect = true, operand_layouts = [dense<> : tensor<0xindex>, dense<> : tensor<0xindex>], result_layouts = [], sdy.sharding = #sdy.sharding_per_value<[<@single_device_0, []>]>} : (tensor<i64>, tensor<f64>) -> ()
   return
 }
 
@@ -443,7 +443,7 @@ func.func public @callback_tuple_result_token_used(%arg0: !stablehlo.token, %arg
   // CHECK-SAME: } : (tensor<i64>, !stablehlo.token, tensor<2xi64>) -> tuple<!stablehlo.token>
   // CHECK-NEXT: %[[TOKEN:.*]] = stablehlo.get_tuple_element %[[CALLBACK]][0] : (tuple<!stablehlo.token>) -> !stablehlo.token
   // CHECK-NEXT: return %[[TOKEN]] : !stablehlo.token
-  %0 = stablehlo.custom_call @xla_python_cpu_callback(%c, %arg0, %arg1) {api_version = 2 : i32, backend_config = "56238119409280", has_side_effect = true, operand_layouts = [dense<> : tensor<0xindex>, dense<> : tensor<0xindex>, dense<0> : tensor<1xindex>], result_layouts = [dense<> : tensor<0xindex>], sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh_0, []>]>} : (tensor<i64>, !stablehlo.token, tensor<2xi64>) -> tuple<!stablehlo.token>
+  %0 = stablehlo.custom_call @xla_python_cpu_callback(%c, %arg0, %arg1) {api_version = 2 : i32, backend_config = "56238119409280", has_side_effect = true, operand_layouts = [dense<> : tensor<0xindex>, dense<> : tensor<0xindex>, dense<0> : tensor<1xindex>], result_layouts = [dense<> : tensor<0xindex>], sdy.sharding = #sdy.sharding_per_value<[<@single_device_0, []>]>} : (tensor<i64>, !stablehlo.token, tensor<2xi64>) -> tuple<!stablehlo.token>
   %1 = stablehlo.get_tuple_element %0[0] : (tuple<!stablehlo.token>) -> !stablehlo.token
   return %1 : !stablehlo.token
 }
@@ -466,7 +466,7 @@ func.func @maximal_sharding_no_results(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32
   // CHECK-V2-NEXT: stablehlo.custom_call @foo(%arg0) {has_side_effect = true, mhlo.sharding = "{maximal device=0}"} : (tensor<8x8xf32>) -> ()
   // CHECK-V3-NEXT: stablehlo.custom_call @foo(%arg0) {has_side_effect = true, mhlo.sharding = "{single_device[device_id=0]}"} : (tensor<8x8xf32>) -> ()
   // CHECK-NEXT: return %arg0 : tensor<8x8xf32>
-  stablehlo.custom_call @foo(%arg0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh_0, []>]>} : (tensor<8x8xf32>) -> ()
+  stablehlo.custom_call @foo(%arg0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@single_device_0, []>]>} : (tensor<8x8xf32>) -> ()
   return %arg0 : tensor<8x8xf32>
 }
 
@@ -482,12 +482,12 @@ func.func @reshard_maximal_sharding(%arg0: tensor<8x8xf32>) -> (tensor<8x8xf32>,
   // CHECK-V3-NEXT: %[[COPY_1:.*]] = mhlo.copy %[[CALL]] {mhlo.sharding = "{mesh['x'=8,'y'=4], [{'y'}, {}]}"} : tensor<8x8xf32>
   // CHECK-V3-NEXT: %[[COPY_2:.*]] = mhlo.copy %[[CALL]] {mhlo.sharding = "{single_device[device_id=1]}"} : tensor<8x8xf32>
   // CHECK-NEXT: return %[[CALL]], %[[COPY_0]], %[[COPY_1]], %[[COPY_2]] : tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>
-  %0 = sdy.reshard %arg0 <@maximal_mesh_0, []> : tensor<8x8xf32>
-  %1 = stablehlo.custom_call @foo(%0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh_0, []>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32>
+  %0 = sdy.reshard %arg0 <@single_device_0, []> : tensor<8x8xf32>
+  %1 = stablehlo.custom_call @foo(%0) {has_side_effect = true, sdy.sharding = #sdy.sharding_per_value<[<@single_device_0, []>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32>
   %2 = sdy.reshard %1 <@mesh_2, [{}, {}]> : tensor<8x8xf32>
   %3 = sdy.reshard %2 <@mesh_2, [{"x"}, {}]> : tensor<8x8xf32>
   %4 = sdy.reshard %1 <@mesh_2, [{"y"}, {}]> : tensor<8x8xf32>
-  %5 = sdy.reshard %1 <@maximal_mesh_1, []> : tensor<8x8xf32>
+  %5 = sdy.reshard %1 <@single_device_1, []> : tensor<8x8xf32>
   return %2, %3, %4, %5 : tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>, tensor<8x8xf32>
 }
 
diff --git a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_pipeline_simplify_replicated.mlir b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_pipeline_simplify_replicated.mlir
index 1f716b800cd..e3adbb07203 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_pipeline_simplify_replicated.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_pipeline_simplify_replicated.mlir
@@ -50,7 +50,7 @@ func.func @do_not_simplify_unreduced(%arg0: tensor<8x8xf32> {sdy.sharding = #sdy
 // -----
 
 sdy.mesh @empty_mesh = <[]>
-sdy.mesh @maximal_mesh = <[], device_ids=[0]>
+sdy.mesh @single_device = <[], device_ids=[0]>
 
 // CHECK-LABEL: func @simplify_replicated_sharding_empty_mesh(
 // CHECK-SAME:      %arg0: tensor<8x8xf32> {mhlo.sharding = "{mesh[], replicated}"},
@@ -58,7 +58,7 @@ sdy.mesh @maximal_mesh = <[], device_ids=[0]>
 // CHECK-SAME:      %arg2: tensor<8x8xf32> {mhlo.sharding = "{single_device[device_id=0]}"})
 func.func @simplify_replicated_sharding_empty_mesh(%arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@empty_mesh, [{}, {}]>},
                                                    %arg1: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@empty_mesh, [{?}, {}]>},
-                                                   %arg2: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@maximal_mesh, []>}) -> tensor<8x8xf32> {
+                                                   %arg2: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@single_device, []>}) -> tensor<8x8xf32> {
   %0 = stablehlo.add %arg0, %arg1 : tensor<8x8xf32>
   return %0 : tensor<8x8xf32>
 }
diff --git a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_round_trip_shard_map_export.mlir b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_round_trip_shard_map_export.mlir
index 5ca34525922..de087b10204 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_round_trip_shard_map_export.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_round_trip_shard_map_export.mlir
@@ -17,7 +17,7 @@
 sdy.mesh @mesh_0 = <["a"=4, "b"=2]>
 sdy.mesh @mesh_1 = <["a"=2, "b"=2, "c"=2, "d"=2]>
 sdy.mesh @mesh_2 = <["a"=4, "b"=4, "c"=4]>
-sdy.mesh @maximal_mesh_0 = <[], device_ids=[0]>
+sdy.mesh @single_device_0 = <[], device_ids=[0]>
 
 // CHECK-LABEL: func @single_manual_comp
 func.func @single_manual_comp(%arg0: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a", ?}, {"b", ?}]>}, %arg1: tensor<16x32xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"b", ?}, {?}]>}) -> (tensor<8x32xf32> {sdy.sharding = #sdy.sharding<@mesh_0, [{"a"}, {}]>}) {
@@ -193,7 +193,7 @@ func.func @manual_computation_with_tokens(
       out_shardings=[<@mesh_0, []>, <@mesh_0, [{"b"}]>]
       manual_axes={"a", "b"} (%arg2: !stablehlo.token, %arg3: tensor<1xi64>) {
     %1 = stablehlo.custom_call @sdy_testonly(%arg2) : (!stablehlo.token) -> (!stablehlo.token)
-    stablehlo.custom_call @sdy_testonly(%1) {sdy.sharding = #sdy.sharding_per_value<[<@maximal_mesh_0, []>]>} : (!stablehlo.token) -> ()
+    stablehlo.custom_call @sdy_testonly(%1) {sdy.sharding = #sdy.sharding_per_value<[<@single_device_0, []>]>} : (!stablehlo.token) -> ()
     sdy.return %1, %arg3 : !stablehlo.token, tensor<1xi64>
   } : (!stablehlo.token, tensor<2xi64>) -> (!stablehlo.token, tensor<2xi64>)
   return %0#0, %0#1 : !stablehlo.token, tensor<2xi64>
diff --git a/third_party/xla/xla/service/spmd/shardy/utils.cc b/third_party/xla/xla/service/spmd/shardy/utils.cc
index 032bdc9f615..b56ea82462d 100644
--- a/third_party/xla/xla/service/spmd/shardy/utils.cc
+++ b/third_party/xla/xla/service/spmd/shardy/utils.cc
@@ -586,8 +586,8 @@ mlir::sdy::TensorShardingPerValueAttr convertToSdySharding(
   // the whole value instead, and let the caller skip setting it.
 
   // An op can carry a sharding without having any values, e.g. a custom call
-  // whose only result is an empty tuple. Such a sharding is maximal, so its
-  // rank is irrelevant, but we still need to keep it.
+  // whose only result is an empty tuple. Such a sharding is single-device, so
+  // its rank is irrelevant, but we still need to keep it.
   if (types.empty()) {
     TensorShardingAttr sdySharding =
         convertToSdyShardingAttr(leafShardings.front(), /*rank=*/0, context);

From 9f80dcb030817e4cad56ed6a11429ce119033e46 Mon Sep 17 00:00:00 2001
From: Vladimir Belitskiy <belitskiy@google.com>
Date: Thu, 15 Oct 2026 13:20:37 -0700
Subject: [PATCH 0906/1000] Preserve `LD_LIBRARY_PATH` in
 `SubProcessTest.SetEnviron` for shared-sanitizer builds.

`SubProcessTest.SetEnviron` replaces the child process's entire environment via `proc.SetEnviron(env)` with only `{"MY_VAR": "my_value"}` before spawning the `test_getenv` C++ helper binary.

When running under `--config=asan_shared` (`-shared-libasan`) with a hermetic LLVM toolchain, the shared ASAN runtime (`libclang_rt.asan.so`) is provided via `LD_LIBRARY_PATH` in the test's runfiles tree. Stripping `LD_LIBRARY_PATH` from the child environment causes the dynamic linker to fail to load `libclang_rt.asan.so` when launching `test_getenv` on RBE workers.

Forward `LD_LIBRARY_PATH` (if present in the parent test environment) into `env` in `SubProcessTest.SetEnviron` so `test_getenv` can load shared runtime libraries while still testing custom environment replacement.

PiperOrigin-RevId: 1000347331
---
 third_party/xla/xla/tsl/platform/subprocess_test.cc | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/third_party/xla/xla/tsl/platform/subprocess_test.cc b/third_party/xla/xla/tsl/platform/subprocess_test.cc
index da4d7a94733..a74416ae0c5 100644
--- a/third_party/xla/xla/tsl/platform/subprocess_test.cc
+++ b/third_party/xla/xla/tsl/platform/subprocess_test.cc
@@ -518,6 +518,9 @@ TEST_F(SubProcessTest, SetEnviron) {
   proc.SetProgram(GetEnvProgram(), {GetEnvProgram(), "MY_VAR"});
   SubProcess::EnvMap env;
   env["MY_VAR"] = "my_value";
+  if (const char* ld_library_path = getenv("LD_LIBRARY_PATH")) {
+    env["LD_LIBRARY_PATH"] = ld_library_path;
+  }
   proc.SetEnviron(env);
   proc.SetChannelAction(CHAN_STDOUT, ACTION_PIPE);
   proc.SetChannelAction(CHAN_STDERR, ACTION_PIPE);

From c5ecd9a24a9a198344994237038bcbce2768b5eb Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 15 Oct 2026 13:31:26 -0700
Subject: [PATCH 0907/1000] PickRematerializationCandidates previously called
 GetCostOfRecompute from scratch on the entire block vector each time a
 candidate block grew by one instruction. For block size B, this performed
 O(B^2) buffer scans, allocated a hash set per step, and deducted shared dead
 external operand buffers multiple times across instructions in the same
 block.

Replace full block reevaluation with an incremental BlockRematState and ExtendRecomputeBlock helper that updates memory_reduced, zero_cost_move, and deduplicated extended_dead_buffers in O(1) per appended instruction, terminating expansion early when an instruction cannot be part of a valid block. Synchronize nonfusion computation schedules across all execution threads after dead code removal.

PiperOrigin-RevId: 1000354259
---
 .../simplifiers/hlo_rematerialization.cc      | 232 ++++++++----------
 .../simplifiers/hlo_rematerialization_test.cc |  42 ++++
 2 files changed, 144 insertions(+), 130 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc
index 34fe885e638..0ef32e97c34 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization.cc
@@ -337,31 +337,6 @@ class MemoryUsageTracker {
   // EndInstruction memory for dead operand(s) is freed.
   absl::Status BeginInstruction(HloRematItem* item);
 
-  int64_t RematerializationCost(const std::vector<HloRematItem*>& items,
-                                int64_t memory_reduced,
-                                int64_t memory_limit_bytes) const {
-    // If none of the users of any 'item' have been placed in the
-    // sequence (as tracked by memory_tracker), then rematerialization of
-    // 'item' is a zero-cost move of 'item->instruction' in the sequence.
-    bool zero_cost_move = true;
-    for (auto* item : items) {
-      auto* instruction = item->instruction;
-      if (absl::c_any_of(
-              instruction->users(),
-              [this](const HloInstruction* inst) { return IsPlaced(inst); })) {
-        zero_cost_move = false;
-        break;
-      }
-    }
-    if (zero_cost_move) {
-      return 0;
-    }
-
-    CHECK_GT(memory_reduced, 0);
-    // Return the inverse of the benefit of rematerialization.
-    return memory_limit_bytes / memory_reduced;
-  }
-
   // Finishes the placement of the current instruction. This frees any dead
   // operands or dead result of the instruction. This must be called after
   // each call to BeginInstruction.
@@ -372,11 +347,6 @@ class MemoryUsageTracker {
   int64_t MemoryReducedIfCompressed(const HloRematItem* item,
                                     const Shape& compact_shape) const;
 
-  // Returns the number of bytes that the current memory usage will be reduced
-  // by if the given sequence of instructions is rematerialized.
-  int64_t MemoryReducedIfRematerialized(
-      absl::Span<const HloRematItem* const> items) const;
-
   absl::Status AddCompressInstructions(HloRematItem* original_item,
                                        HloRematItem* compressed_item,
                                        HloRematItem* uncompressed_item);
@@ -423,11 +393,6 @@ class MemoryUsageTracker {
   std::optional<int64_t> GetCostOfHostOffload(
       const HloRematItem* candidate_item, int64_t memory_limit_bytes) const;
 
-  // Calculates the cost of recomputing the candidate_item's output.
-  std::optional<int64_t> GetCostOfRecompute(
-      const std::vector<HloRematItem*>& candidate_items,
-      int64_t memory_limit_bytes) const;
-
   // Selects and returns the best candidate instructions for rematerialization.
   // A sequence of candidate instructions of length between min_block_size and
   // max_block_size (both inclusive) with the lowest rematerialization cost is
@@ -531,6 +496,22 @@ class MemoryUsageTracker {
   // scope.
   void CountAllocatedMemory(HloRematItem* item);
 
+  // Tracks incremental memory reduction, placed user dependencies, and operand
+  // live range extensions during candidate block evaluation in
+  // ExtendRecomputeBlock.
+  struct BlockRematState {
+    int64_t memory_reduced = 0;
+    bool zero_cost_move = true;
+    absl::InlinedVector<const HloRematItem*, 16> items;
+    absl::InlinedVector<BufferId, 16> extended_dead_buffers;
+  };
+
+  // Incrementally evaluates whether appending item to the candidate block is
+  // legal and updates state with net memory changes and required live range
+  // extensions. Returns false if the block cannot be extended with item.
+  bool ExtendRecomputeBlock(const HloRematItem* item,
+                            BlockRematState* state) const;
+
   // Adjusts our tracked memory usage as a result of this item going out of
   // scope.
   absl::Status CountFreedMemory(HloRematItem* item);
@@ -852,62 +833,60 @@ int64_t MemoryUsageTracker::MemoryReducedIfCompressed(
   return memory_reduced;
 }
 
-int64_t MemoryUsageTracker::MemoryReducedIfRematerialized(
-    absl::Span<const HloRematItem* const> items) const {
+bool MemoryUsageTracker::ExtendRecomputeBlock(const HloRematItem* item,
+                                              BlockRematState* state) const {
   CHECK_NE(in_progress_item_, nullptr);
-  int64_t memory_reduced = 0;
-  absl::flat_hash_set<const HloRematItem*> remat_candidates;
-
-  for (const HloRematItem* item : items) {
-    if (!item->placed || item == in_progress_item_) {
-      LOG(WARNING) << "Unplaced item or in progress item being checked for "
-                      "rematerialization.";
-      return 0;
-    }
-
-    // Compute the amount of memory reduced (if any) by rematerializing
-    // 'item->instruction'. The HloValues defined by 'item->instruction'
-    // will no longer be live at this program point, so initially set
-    // memory_reduced to the size of its defined values.
-    for (BufferId buffer_id : item->buffers_defined) {
-      const Buffer& buffer = buffers_.at(buffer_id);
-      // Avoid rematerializing instructions with indirect uses as it is
-      // difficult to reason about liveness after rematerializing the
-      // instruction.
-      // Avoid rematerializing instructions with live out buffers.
-      // Avoid rematerializing buffers that are in nested tuples.
-      // TODO(mpurohit): Check why live_out buffers are an issue here.
-      if (buffer.has_indirect_uses || buffer.live_out ||
-          buffer.index.size() > 1) {
-        return 0;
-      }
-      if (IsInUse(buffer_id)) {
-        return 0;
-      }
-      if (IsCurrentlyLive(buffer_id)) {
-        memory_reduced += AllocatedSize(buffer_id);
-      }
-    }
-
-    // Account for any logical buffers whose live range must be extended across
-    // this program point.
-    for (BufferId buffer_id : item->buffers_used) {
-      if (!IsCurrentlyLive(buffer_id)) {
-        // This logical buffer is used by 'item->instruction' but is not live at
-        // this program point. Rematerializing 'item->instruction' will extend
-        // the buffer's live range across this program point unless it is
-        // defined by an instruction that is also being rematerialized.
-        HloRematItem* defining_instruction =
-            buffers_.at(buffer_id).defining_instruction;
-        if (!remat_candidates.contains(defining_instruction)) {
-          memory_reduced -= AllocatedSize(buffer_id);
-        }
-      }
-    }
-    remat_candidates.insert(item);
+  if (!item->placed || item == in_progress_item_) {
+    LOG(WARNING) << "Unplaced item or in progress item being checked for "
+                    "rematerialization.";
+    return false;
   }
 
-  return memory_reduced;
+  const HloInstruction* candidate = item->instruction;
+  if (absl::c_any_of(
+          candidate->control_successors(),
+          [this](const HloInstruction* inst) { return IsPlaced(inst); })) {
+    return false;
+  }
+
+  for (BufferId buffer_id : item->buffers_defined) {
+    const Buffer& buffer = buffers_.at(buffer_id);
+    // Avoid rematerializing instructions with indirect uses as it is
+    // difficult to reason about liveness after rematerializing the
+    // instruction.
+    // Avoid rematerializing instructions with live out buffers.
+    // Avoid rematerializing buffers that are in nested tuples.
+    // TODO(mpurohit): Check why live_out buffers are an issue here.
+    if (buffer.has_indirect_uses || buffer.live_out ||
+        buffer.index.size() > 1 || IsInUse(buffer_id)) {
+      return false;
+    }
+    if (IsCurrentlyLive(buffer_id)) {
+      state->memory_reduced += AllocatedSize(buffer_id);
+    }
+  }
+
+  for (BufferId buffer_id : item->buffers_used) {
+    if (!IsCurrentlyLive(buffer_id)) {
+      HloRematItem* defining_instruction =
+          buffers_.at(buffer_id).defining_instruction;
+      if (!absl::c_linear_search(state->items, defining_instruction) &&
+          !absl::c_linear_search(state->extended_dead_buffers, buffer_id)) {
+        state->extended_dead_buffers.push_back(buffer_id);
+        state->memory_reduced -= AllocatedSize(buffer_id);
+      }
+    }
+  }
+
+  if (state->zero_cost_move &&
+      absl::c_any_of(candidate->users(), [this](const HloInstruction* inst) {
+        return IsPlaced(inst);
+      })) {
+    state->zero_cost_move = false;
+  }
+
+  state->items.push_back(item);
+  return true;
 }
 
 std::tuple<UsesList, UsesList> MemoryUsageTracker::GetPlacedAndUnplacedUsers(
@@ -1655,35 +1634,6 @@ std::optional<int64_t> MemoryUsageTracker::GetCostOfHostOffload(
   return memory_limit_bytes / bytes_used_by_buffers;
 }
 
-std::optional<int64_t> MemoryUsageTracker::GetCostOfRecompute(
-    const std::vector<HloRematItem*>& candidate_items,
-    int64_t memory_limit_bytes) const {
-  // If any of the candidate's control successor has been placed, we need
-  // to skip this candidate. Otherwise we will violate control dependency.
-  for (auto* item : candidate_items) {
-    HloInstruction* candidate = item->instruction;
-    if (std::any_of(
-            candidate->control_successors().begin(),
-            candidate->control_successors().end(),
-            [this](const HloInstruction* inst) { return IsPlaced(inst); })) {
-      return {};
-    }
-  }
-
-  // Evaluate this block as a candidate for recompute rematerialization.
-  VLOG(5) << "Block contains:";
-  for (auto* hlo : candidate_items) {
-    VLOG(5) << hlo->instruction->name();
-  }
-  const int64_t memory_reduced = MemoryReducedIfRematerialized(candidate_items);
-  if (memory_reduced <= 0) {
-    return {};
-  }
-
-  return RematerializationCost(candidate_items, memory_reduced,
-                               memory_limit_bytes);
-}
-
 std::tuple<std::vector<HloRematItem*>, RematStrategy, int>
 MemoryUsageTracker::PickRematerializationCandidates(
     const HloRematInstructionList& instruction_list, int64_t memory_limit_bytes,
@@ -1777,25 +1727,47 @@ MemoryUsageTracker::PickRematerializationCandidates(
       continue;
     }
 
-    while (block.size() <= max_block_size) {
-      auto cost = GetCostOfRecompute(block, memory_limit_bytes);
+    BlockRematState block_state;
+    bool valid_initial_block = true;
+    for (const HloRematItem* item : block) {
+      if (!ExtendRecomputeBlock(item, &block_state)) {
+        valid_initial_block = false;
+        break;
+      }
+    }
+    if (!valid_initial_block) {
       ++effort;
-      if (cost && *cost < best_cost) {
-        VLOG(1) << "Found new best cost; from " << best_cost << " to " << *cost
-                << " with strategy kRecompute on block of size "
-                << block.size();
-        best_strategy.kind = RematStrategy::kRecompute;
-        best_items = block;
-        best_cost = *cost;
+      continue;
+    }
+
+    while (block.size() <= max_block_size) {
+      ++effort;
+      if (block_state.memory_reduced > 0) {
+        int64_t cost = block_state.zero_cost_move
+                           ? 0
+                           : memory_limit_bytes / block_state.memory_reduced;
+        if (cost < best_cost) {
+          VLOG(1) << "Found new best cost; from " << best_cost << " to " << cost
+                  << " with strategy kRecompute on block of size "
+                  << block.size();
+          best_strategy.kind = RematStrategy::kRecompute;
+          best_items = block;
+          best_cost = cost;
+        }
+      }
+
+      if (block.size() == max_block_size) {
+        break;
       }
 
       // Try to add the next instruction to this block to evaluate as a possibly
       // better candidate for rematerialization.
-      auto* last_item = block[block.size() - 1];
+      auto* last_item = block.back();
       auto* next_item = instruction_list.next(last_item);
       if (next_item == nullptr || next_item->denylisted || !next_item->placed ||
           next_item == in_progress_item_ ||
-          !CanBeRematerialized(next_item->instruction, rematerializable_map)) {
+          !CanBeRematerialized(next_item->instruction, rematerializable_map) ||
+          !ExtendRecomputeBlock(next_item, &block_state)) {
         break;
       }
       block.push_back(next_item);
@@ -3365,8 +3337,8 @@ absl::StatusOr<bool> HloRematerialization::RunImpl(
   changed |= dead_code_removed;
 
   // After DCE, the module sequence may include instructions which no longer
-  // exist. Update the schedule and restore it.
-  ABSL_RETURN_IF_ERROR(saved_schedule.Update(execution_threads));
+  // exist. Update the schedule across all execution threads and restore it.
+  ABSL_RETURN_IF_ERROR(saved_schedule.Update());
   ABSL_RETURN_IF_ERROR(module->set_schedule(std::move(saved_schedule)));
   VLOG(1) << "Rematerialized " << instructions_rematerialized_
           << " instructions in module " << module->name() << "; "
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization_test.cc
index f44551ae4ca..c8aebc680fe 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/hlo_rematerialization_test.cc
@@ -2625,5 +2625,47 @@ ENTRY %entry {
       ::testing::UnitTest::GetInstance()->current_test_info()->name());
 }
 
+TEST_F(RecomputeAndCompressHloRematerializationTest,
+       BlockRematSharedDeadOperand) {
+  constexpr absl::string_view hlo_string = R"hlo(
+HloModule BlockRematSharedDeadOperand, is_scheduled=true
+
+ENTRY %entry {
+  %param = f32[] parameter(0)
+  %shared_dead = f32[1024]{0} custom-call(%param), custom_call_target="Producer"
+  %cand0 = f32[1024]{0} negate(%shared_dead)
+  %cand1 = f32[1024]{0} abs(%shared_dead)
+  %early_use = f32[1024]{0} custom-call(%cand0, %cand1), custom_call_target="EarlyUse"
+  %peak_op = f32[1024]{0} custom-call(%early_use), custom_call_target="Peak"
+  %final_op = f32[1024]{0} custom-call(%peak_op), custom_call_target="Final"
+  %late_use = f32[1024]{0} add(%cand0, %cand1)
+  ROOT %out = (f32[1024]{0}, f32[1024]{0}) tuple(%final_op, %late_use)
+}
+)hlo";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  HloRematerialization::RematerializationSizes sizes;
+  constexpr int64_t kMemoryLimitBytes = 20 * 1024 + 16;
+  ASSERT_OK_AND_ASSIGN(bool changed,
+                       RunHloRematerialization(
+                           kMemoryLimitBytes, module.get(),
+                           /*min_remat_size=*/0,
+                           HloRematerialization::RematAlgorithm::kAlwaysRemat,
+                           /*block_size_limit=*/2,
+                           /*on_rematerialized=*/nullptr, &sizes));
+  EXPECT_TRUE(changed);
+  EXPECT_EQ(sizes.before_bytes, 24 * 1024 + 16);
+  EXPECT_EQ(sizes.after_bytes, 20 * 1024 + 16);
+
+  const HloInstruction* late_use =
+      module->entry_computation()->GetInstructionWithName("late_use");
+  ASSERT_NE(late_use, nullptr);
+  EXPECT_EQ(late_use->operand(0)->name(), "cand0.remat");
+  EXPECT_EQ(late_use->operand(1)->name(), "cand1.remat");
+
+  CheckForRematInInstructionNames(
+      ::testing::UnitTest::GetInstance()->current_test_info()->name());
+}
+
 }  // namespace
 }  // namespace xla

From 899832f2c6afa3ad816ce30979157be633e8601e Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 15 Oct 2026 14:14:09 -0700
Subject: [PATCH 0908/1000] Fix asan bug introduced in the test

PiperOrigin-RevId: 1000384688
---
 tensorflow/core/tfrt/mlrt/kernel/ifrt_ops_kernel_test.cc | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/tensorflow/core/tfrt/mlrt/kernel/ifrt_ops_kernel_test.cc b/tensorflow/core/tfrt/mlrt/kernel/ifrt_ops_kernel_test.cc
index 6bb62454411..641bc548b27 100644
--- a/tensorflow/core/tfrt/mlrt/kernel/ifrt_ops_kernel_test.cc
+++ b/tensorflow/core/tfrt/mlrt/kernel/ifrt_ops_kernel_test.cc
@@ -1166,10 +1166,6 @@ TEST_P(KernelTest, IfrtRestoreVariableOpWaitsForAsyncShardsBeforeDestruction) {
   resource_context_.DeleteResource(ifrt_serving::kIfrtModelRestoreContextName);
   EXPECT_TRUE(restored_future.IsReady());
   fallback_state_.reset();
-
-  if (!unblock_restore.HasBeenNotified()) {
-    unblock_restore.Notify();
-  }
   restore_work_queue_->Quiesce();
 
   absl::StatusOr<tensorflow::Tensor> restored_tensor = restored_future.Await();

From 54a5f77c732fedf6532a5ecf36624e24e60804da Mon Sep 17 00:00:00 2001
From: Bill Varcho <varcho@google.com>
Date: Thu, 15 Oct 2026 14:24:05 -0700
Subject: [PATCH 0909/1000] [SDY][perf][2/2] Disable folding and constant CSE
 in SdyRoundTripExportPipeline canonicalizer.

PiperOrigin-RevId: 1000392609
---
 .../spmd/shardy/sdy_round_trip/pipelines.cc   | 22 +++++++++++++------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/pipelines.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/pipelines.cc
index 7db3eefc7bd..8004f842315 100644
--- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/pipelines.cc
+++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/pipelines.cc
@@ -49,10 +49,23 @@ using ::mlir::PassPipelineOptions;
 using ::mlir::PassPipelineRegistration;
 using ::mlir::func::FuncOp;
 
+namespace {
+
+mlir::GreedyRewriteConfig getCanonicalizerConfig() {
+  return mlir::GreedyRewriteConfig()
+      .setUseTopDownTraversal(true)
+      .setRegionSimplificationLevel(mlir::GreedySimplifyRegionLevel::Normal)
+      .enableFolding(false)
+      .enableConstantCSE(false);
+}
+
+}  // namespace
+
 void addSdyRoundTripExportPipeline(mlir::OpPassManager& pm,
                                    bool keepMeshesInlined,
                                    bool enableHloShardingV3) {
-  pm.addPass(mlir::createCanonicalizerPass());
+  // Constants are still sdy.constant here, so disable folding and constant CSE.
+  pm.addPass(mlir::createCanonicalizerPass(getCanonicalizerConfig()));
   // Lift meshes before deduping, since the dedup meshes pass ignores inlined
   // meshes.
   if (!keepMeshesInlined) {
@@ -117,12 +130,7 @@ void addSdyRoundTripImportPipeline(mlir::OpPassManager& pm,
             /*useSingleDeviceMeshNames=*/true}));
     pm.addPass(createSdyRoundTripDedupMeshesPass());
   }
-  pm.addPass(createCanonicalizerPass(
-      mlir::GreedyRewriteConfig()
-          .setUseTopDownTraversal(true)
-          .setRegionSimplificationLevel(mlir::GreedySimplifyRegionLevel::Normal)
-          .enableFolding(false)
-          .enableConstantCSE(false)));
+  pm.addPass(mlir::createCanonicalizerPass(getCanonicalizerConfig()));
 }
 
 namespace {

From 09a8a9d6c2bfd8e366a63cd9fa0f766b7f4d49d2 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Thu, 15 Oct 2026 14:35:07 -0700
Subject: [PATCH 0910/1000] PR #3022: Add exp2 and log2 operations.

Imported from GitHub PR https://github.com/openxla/stablehlo/pull/3022

Add both to the StableHLO and VHLO dialects along with reference interpreter and Linalg/Arith conversion implementations.
Copybara import of the project:

--
11c2db2c73b638b3d15bb273780037c68d77c4b5 by Peter Hawkins <phawkins@google.com>:

Add exp2 and log2 operations.

Add both to the StableHLO and VHLO dialects along with reference interpreter and Linalg/Arith conversion implementations.

Merging this change closes #3022

PiperOrigin-RevId: 1000400074
---
 .../xla/third_party/stablehlo/temporary.patch | 4060 ++++++++++++++++-
 .../mhlo/transforms/map_stablehlo_to_hlo_op.h |    2 +
 .../mhlo/hlo-legalize-to-stablehlo.mlir       |   18 +
 .../mhlo/stablehlo-legalize-to-hlo.mlir       |   18 +
 4 files changed, 4077 insertions(+), 21 deletions(-)

diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch
index ccd5f019ae8..ecdcdce23c6 100644
--- a/third_party/xla/third_party/stablehlo/temporary.patch
+++ b/third_party/xla/third_party/stablehlo/temporary.patch
@@ -135,6 +135,41 @@ diff --ruN a/stablehlo/build_tools/lit_test_suite.bzl b/stablehlo/build_tools/li
              tags = tags,
              size = size_per_test,
          )
+diff --ruN a/stablehlo/docs/generated/StablehloBuilder.md b/stablehlo/docs/generated/StablehloBuilder.md
+--- stablehlo/docs/generated/StablehloBuilder.md
++++ stablehlo/docs/generated/StablehloBuilder.md
+@@ -424,6 +424,15 @@
+ MlirOp Einsum(Type resultType, MlirOp &lhs, MlirOp &rhs, ::llvm::StringRef einsum_config);
+ ```
+ 
++### `stablehlo::Exp2Op`
++
++Creates a new [`stablehlo.exp2`](https://openxla.org/stablehlo/spec#exp2)
++operation.
++
++```c++
++MlirOp Exp2(MlirOp &operand, /*optional*/::mlir::stablehlo::ResultAccuracyAttr result_accuracy = {});
++```
++
+ ### `stablehlo::ExpOp`
+ 
+ Creates a new [`stablehlo.exponential`](https://openxla.org/stablehlo/spec#exponential)
+@@ -543,6 +552,15 @@
+ MlirOp Log1p(MlirOp &operand, /*optional*/::mlir::stablehlo::ResultAccuracyAttr result_accuracy = {});
+ ```
+ 
++### `stablehlo::Log2Op`
++
++Creates a new [`stablehlo.log2`](https://openxla.org/stablehlo/spec#log2)
++operation.
++
++```c++
++MlirOp Log2(MlirOp &operand, /*optional*/::mlir::stablehlo::ResultAccuracyAttr result_accuracy = {});
++```
++
+ ### `stablehlo::LogOp`
+ 
+ Creates a new [`stablehlo.log`](https://openxla.org/stablehlo/spec#log)
 diff --ruN a/stablehlo/docs/generated/chlo.md b/stablehlo/docs/generated/chlo.md
 --- stablehlo/docs/generated/chlo.md
 +++ stablehlo/docs/generated/chlo.md
@@ -1115,7 +1150,37 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgU
 diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToScalarOp.h b/stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToScalarOp.h
 --- stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToScalarOp.h
 +++ stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToScalarOp.h
-@@ -204,7 +204,8 @@
+@@ -17,6 +17,7 @@
+ #ifndef STABLEHLO_CONVERSIONS_LINALG_TRANSFORMS_MAP_STABLEHLO_TO_SCALAR_OP_H
+ #define STABLEHLO_CONVERSIONS_LINALG_TRANSFORMS_MAP_STABLEHLO_TO_SCALAR_OP_H
+ 
++#include <cmath>
+ #include <optional>
+ 
+ #include "llvm/ADT/ArrayRef.h"
+@@ -93,6 +94,10 @@
+   using COp = ::mlir::complex::Expm1Op;
+ };
+ template <>
++struct StablehloToScalarOp<stablehlo::Exp2Op> {
++  using FOp = ::mlir::math::Exp2Op;
++};
++template <>
+ struct StablehloToScalarOp<stablehlo::FloorOp> {
+   using FOp = ::mlir::math::FloorOp;
+ };
+@@ -107,6 +112,10 @@
+   using COp = ::mlir::complex::Log1pOp;
+ };
+ template <>
++struct StablehloToScalarOp<stablehlo::Log2Op> {
++  using FOp = ::mlir::math::Log2Op;
++};
++template <>
+ struct StablehloToScalarOp<stablehlo::MulOp> {
+   using FOp = ::mlir::arith::MulFOp;
+   using IOp = ::mlir::arith::MulIOp;
+@@ -204,7 +213,8 @@
    Value operator()(Location loc, ArrayRef<Type> resultTypes,
                     ArrayRef<Type> /*argTypes*/, ValueRange args, OpBuilder* b) {
      return StdScalarOp::create(*b, loc, resultTypes, args,
@@ -1125,7 +1190,7 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToSca
    }
  };
  
-@@ -215,7 +216,8 @@
+@@ -215,7 +225,8 @@
      Type elementType = getElementTypeOrSelf(argTypes.front());
      if (SupportedType{}(elementType)) {
        return StdScalarOp::create(*b, loc, resultTypes, args,
@@ -1135,7 +1200,7 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToSca
      }
      return MapStablehloOpToScalarOpImpl<Args...>{}(loc, resultTypes, argTypes,
                                                     args, b);
-@@ -456,6 +458,60 @@
+@@ -456,6 +467,60 @@
        assert(predicate.has_value() && "expected valid comparison direction");
        return arith::CmpIOp::create(*b, loc, *predicate, lhsInt, rhsInt);
      }
@@ -1196,7 +1261,54 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToSca
      std::optional<arith::CmpFPredicate> predicate =
          getCmpPredicate<arith::CmpFPredicate>(comparisonDirection,
                                                /*is_signed=*/true);
-@@ -688,32 +744,39 @@
+@@ -670,6 +735,46 @@
+       loc, resultTypes, argTypes, adaptor.getOperands(), b);
+ }
+ 
++template <>
++inline Value mapStablehloOpToStdScalarOp<stablehlo::Exp2Op>(
++    Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
++    stablehlo::Exp2Op::Adaptor adaptor, OpBuilder* b) {
++  Type type = adaptor.getOperand().getType();
++  if (!isa<ComplexType>(type))
++    return MapStablehloOpToScalarOpImpl<IsFloatType, math::Exp2Op>{}(
++        loc, resultTypes, argTypes, adaptor.getOperands(), b);
++  auto complexTy = cast<ComplexType>(type);
++  auto floatTy = complexTy.getElementType();
++  Value ln2 =
++      arith::ConstantOp::create(*b, loc, b->getFloatAttr(floatTy, M_LN2));
++  Value re = complex::ReOp::create(*b, loc, floatTy, adaptor.getOperand());
++  Value im = complex::ImOp::create(*b, loc, floatTy, adaptor.getOperand());
++  Value scaled = complex::CreateOp::create(
++      *b, loc, complexTy, arith::MulFOp::create(*b, loc, re, ln2),
++      arith::MulFOp::create(*b, loc, im, ln2));
++  return complex::ExpOp::create(*b, loc, scaled);
++}
++
++template <>
++inline Value mapStablehloOpToStdScalarOp<stablehlo::Log2Op>(
++    Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
++    stablehlo::Log2Op::Adaptor adaptor, OpBuilder* b) {
++  Type type = adaptor.getOperand().getType();
++  if (!isa<ComplexType>(type))
++    return MapStablehloOpToScalarOpImpl<IsFloatType, math::Log2Op>{}(
++        loc, resultTypes, argTypes, adaptor.getOperands(), b);
++  auto complexTy = cast<ComplexType>(type);
++  auto floatTy = complexTy.getElementType();
++  Value logZ = complex::LogOp::create(*b, loc, adaptor.getOperand());
++  Value oneOverLn2 =
++      arith::ConstantOp::create(*b, loc, b->getFloatAttr(floatTy, M_LOG2E));
++  Value re = complex::ReOp::create(*b, loc, floatTy, logZ);
++  Value im = complex::ImOp::create(*b, loc, floatTy, logZ);
++  return complex::CreateOp::create(
++      *b, loc, complexTy, arith::MulFOp::create(*b, loc, re, oneOverLn2),
++      arith::MulFOp::create(*b, loc, im, oneOverLn2));
++}
++
+ // 'target_types' is the unconverted type (signed or unsigned if integer),
+ // 'ResultTypes' is the converted type (signless if integer).
+ inline Value mapConvertOpToStdScalarOp(Location loc, ArrayRef<Type> targetTypes,
+@@ -688,32 +793,39 @@
    // A boolean value is considered to be unsigned when converting to
    // floating-point. Otherwise, it will become `-1`.
    if (IsUnsignedIntegerType{}(sourceType) &&
@@ -1249,7 +1361,7 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToSca
      }
      // No conversion is needed for identical float types.
      return args.front();
-@@ -738,17 +801,20 @@
+@@ -738,17 +850,20 @@
      auto src = cast<IntegerType>(sourceType);
      auto res = cast<IntegerType>(targetType);
      if (src.getWidth() > res.getWidth()) {
@@ -1276,6 +1388,51 @@ diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToSca
      }
      // No conversion is needed for the same width integers
      return args.front();
+diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/StablehloToArith.cpp b/stablehlo/stablehlo/conversions/linalg/transforms/StablehloToArith.cpp
+--- stablehlo/stablehlo/conversions/linalg/transforms/StablehloToArith.cpp
++++ stablehlo/stablehlo/conversions/linalg/transforms/StablehloToArith.cpp
+@@ -119,10 +119,12 @@
+       ScalarHloToArithmeticPattern<mlir::stablehlo::DivOp>,
+       ScalarHloToArithmeticPattern<mlir::stablehlo::ExpOp>,
+       ScalarHloToArithmeticPattern<mlir::stablehlo::Expm1Op>,
++      ScalarHloToArithmeticPattern<mlir::stablehlo::Exp2Op>,
+       ScalarHloToArithmeticPattern<mlir::stablehlo::FloorOp>,
+       ScalarHloToArithmeticPattern<mlir::stablehlo::ImagOp>,
+       ScalarHloToArithmeticPattern<mlir::stablehlo::IsFiniteOp>,
+       ScalarHloToArithmeticPattern<mlir::stablehlo::Log1pOp>,
++      ScalarHloToArithmeticPattern<mlir::stablehlo::Log2Op>,
+       ScalarHloToArithmeticPattern<mlir::stablehlo::LogOp>,
+       ScalarHloToArithmeticPattern<mlir::stablehlo::LogisticOp>,
+       ScalarHloToArithmeticPattern<mlir::stablehlo::MaxOp>,
+diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/StablehloToLinalgPointwise.cpp b/stablehlo/stablehlo/conversions/linalg/transforms/StablehloToLinalgPointwise.cpp
+--- stablehlo/stablehlo/conversions/linalg/transforms/StablehloToLinalgPointwise.cpp
++++ stablehlo/stablehlo/conversions/linalg/transforms/StablehloToLinalgPointwise.cpp
+@@ -285,10 +285,12 @@
+         PointwiseToLinalgMapConverter<mlir::stablehlo::DivOp>,
+         PointwiseToLinalgMapConverter<mlir::stablehlo::ExpOp>,
+         PointwiseToLinalgMapConverter<mlir::stablehlo::Expm1Op>,
++        PointwiseToLinalgMapConverter<mlir::stablehlo::Exp2Op>,
+         PointwiseToLinalgMapConverter<mlir::stablehlo::FloorOp>,
+         PointwiseToLinalgMapConverter<mlir::stablehlo::ImagOp>,
+         PointwiseToLinalgMapConverter<mlir::stablehlo::IsFiniteOp>,
+         PointwiseToLinalgMapConverter<mlir::stablehlo::Log1pOp>,
++        PointwiseToLinalgMapConverter<mlir::stablehlo::Log2Op>,
+         PointwiseToLinalgMapConverter<mlir::stablehlo::LogOp>,
+         PointwiseToLinalgMapConverter<mlir::stablehlo::LogisticOp>,
+         PointwiseToLinalgMapConverter<mlir::stablehlo::MaxOp>,
+@@ -336,10 +338,12 @@
+             PointwiseToLinalgConverter<mlir::stablehlo::DivOp>,
+             PointwiseToLinalgConverter<mlir::stablehlo::ExpOp>,
+             PointwiseToLinalgConverter<mlir::stablehlo::Expm1Op>,
++            PointwiseToLinalgConverter<mlir::stablehlo::Exp2Op>,
+             PointwiseToLinalgConverter<mlir::stablehlo::FloorOp>,
+             PointwiseToLinalgConverter<mlir::stablehlo::ImagOp>,
+             PointwiseToLinalgConverter<mlir::stablehlo::IsFiniteOp>,
+             PointwiseToLinalgConverter<mlir::stablehlo::Log1pOp>,
++            PointwiseToLinalgConverter<mlir::stablehlo::Log2Op>,
+             PointwiseToLinalgConverter<mlir::stablehlo::LogOp>,
+             PointwiseToLinalgConverter<mlir::stablehlo::LogisticOp>,
+             PointwiseToLinalgConverter<mlir::stablehlo::MaxOp>,
 diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir
 --- stablehlo/stablehlo/conversions/tosa/tests/binary.mlir
 +++ stablehlo/stablehlo/conversions/tosa/tests/binary.mlir
@@ -2261,7 +2418,59 @@ diff --ruN a/stablehlo/stablehlo/dialect/StablehloEnums.td b/stablehlo/stablehlo
 diff --ruN a/stablehlo/stablehlo/dialect/StablehloOps.cpp b/stablehlo/stablehlo/dialect/StablehloOps.cpp
 --- stablehlo/stablehlo/dialect/StablehloOps.cpp
 +++ stablehlo/stablehlo/dialect/StablehloOps.cpp
-@@ -3283,7 +3283,10 @@
+@@ -249,9 +249,11 @@
+ INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(DivOp)
+ INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(ExpOp)
+ INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(Expm1Op)
++INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(Exp2Op)
+ INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(FloorOp)
+ INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(LogOp)
+ INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(Log1pOp)
++INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(Log2Op)
+ INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(LogisticOp)
+ INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(MaxOp)
+ INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(MinOp)
+@@ -1129,6 +1131,19 @@
+   return success();
+ }
+ 
++// ===---------------------------------------------------------------------===//
++// Exp2Op
++//===----------------------------------------------------------------------===//
++
++LogicalResult Exp2Op::verify() {
++  if (auto attr = getResultAccuracyAttr()) {
++    return ResultAccuracyAttr::verify([&] { return emitError(); },
++                                      attr.getAtol(), attr.getRtol(),
++                                      attr.getUlps(), attr.getMode());
++  }
++  return success();
++}
++
+ // ===---------------------------------------------------------------------===//
+ // LogOp
+ //===----------------------------------------------------------------------===//
+@@ -1155,6 +1170,19 @@
+   return success();
+ }
+ 
++// ===---------------------------------------------------------------------===//
++// Log2Op
++//===----------------------------------------------------------------------===//
++
++LogicalResult Log2Op::verify() {
++  if (auto attr = getResultAccuracyAttr()) {
++    return ResultAccuracyAttr::verify([&] { return emitError(); },
++                                      attr.getAtol(), attr.getRtol(),
++                                      attr.getUlps(), attr.getMode());
++  }
++  return success();
++}
++
+ // ===---------------------------------------------------------------------===//
+ // LogisticOp
+ //===----------------------------------------------------------------------===//
+@@ -3283,7 +3311,10 @@
      RegionRange regions,
      SmallVectorImpl<ShapedTypeComponents>& inferredReturnShapes) {
    CompareOp::Adaptor adaptor(operands, attributes, properties, regions);
@@ -2273,7 +2482,7 @@ diff --ruN a/stablehlo/stablehlo/dialect/StablehloOps.cpp b/stablehlo/stablehlo/
                               inferredReturnShapes);
  }
  
-@@ -4332,6 +4335,23 @@
+@@ -4332,6 +4363,23 @@
    ReturnOp::create(*builder, loc, compare);
  }
  
@@ -2297,7 +2506,7 @@ diff --ruN a/stablehlo/stablehlo/dialect/StablehloOps.cpp b/stablehlo/stablehlo/
  void buildMaxAndArgmaxBody(Type elementType, Type indices_type, Region& body,
                             OpBuilder& builder) {
    OpBuilder::InsertionGuard guard(builder);
-@@ -4366,15 +4386,18 @@
+@@ -4366,15 +4414,18 @@
    // Final lhs Selection Condition: (gt_pred) OR (tie_breaker_condition)
    auto final_lhs_condition =
        OrOp::create(builder, loc, gt_pred, tie_breaker_condition).getResult();
@@ -2334,7 +2543,75 @@ diff --ruN a/stablehlo/stablehlo/dialect/StablehloOps.td b/stablehlo/stablehlo/d
  }
  
  class StableHLO_Op<string mnemonic, list<Trait> traits = []> :
-@@ -2525,8 +2526,9 @@
+@@ -393,6 +394,33 @@
+   }];
+ }
+ 
++def StableHLO_Exp2Op: StableHLO_UnaryElementwiseOp<"exp2",
++    [HLO_CompatibleOperandsAndResultType],
++    HLO_FpComplexOrQuantizedIntTensor> {
++  let summary = "Exp2 operation";
++  let description = [{
++    Performs element-wise base-2 exponential operation on `operand` tensor and
++    produces a `result` tensor.
++
++    See:
++    https://github.com/openxla/stablehlo/blob/main/docs/spec.md#exp2
++
++    Example:
++    ```mlir
++    %result = stablehlo.exp2 %operand : tensor<2x2xf64>
++    ```
++  }];
++  let arguments = (ins HLO_FpComplexOrQuantizedIntTensor:$operand,
++                       DefaultValuedOptionalAttr<StableHLO_ResultAccuracyAttr,
++                       "::mlir::stablehlo::ResultAccuracyMode::DEFAULT">:$result_accuracy);
++  let results = (outs HLO_FpComplexOrQuantizedIntTensor:$result);
++  let hasVerifier = 1;
++
++  let assemblyFormat = [{
++    $operand attr-dict `:` custom<SameOperandsAndResultType>(type($operand), type($result))
++  }];
++}
++
+ def StableHLO_FloorOp: StableHLO_UnaryElementwiseOp<"floor",
+     [HLO_CompatibleOperandsAndResultType], HLO_FpOrQuantizedIntTensor> {
+   let summary = "Floor operation";
+@@ -506,6 +534,33 @@
+   }];
+ }
+ 
++def StableHLO_Log2Op: StableHLO_UnaryElementwiseOp<"log2",
++    [HLO_CompatibleOperandsAndResultType],
++    HLO_FpComplexOrQuantizedIntTensor> {
++  let summary = "Log2 operation";
++  let description = [{
++    Performs element-wise base-2 logarithm operation on `operand` tensor and
++    produces a `result` tensor.
++
++    See:
++    https://github.com/openxla/stablehlo/blob/main/docs/spec.md#log2
++
++    Example:
++    ```mlir
++    %result = stablehlo.log2 %operand : tensor<2x2xf64>
++    ```
++  }];
++  let arguments = (ins HLO_FpComplexOrQuantizedIntTensor:$operand,
++                       DefaultValuedOptionalAttr<StableHLO_ResultAccuracyAttr,
++                       "::mlir::stablehlo::ResultAccuracyMode::DEFAULT">:$result_accuracy);
++  let results = (outs HLO_FpComplexOrQuantizedIntTensor:$result);
++  let hasVerifier = 1;
++
++  let assemblyFormat = [{
++    $operand attr-dict `:` custom<SameOperandsAndResultType>(type($operand), type($result))
++  }];
++}
++
+ def StableHLO_LogisticOp: StableHLO_UnaryElementwiseOp<"logistic",
+     [HLO_CompatibleOperandsAndResultType /*logistic_c1*/],
+     HLO_FpComplexOrQuantizedIntTensor /*logistic_i1*/> { /*logistic_c1*/
+@@ -2525,8 +2580,9 @@
        "::mlir::TypeRange":$resultTypes, "::mlir::ValueRange":$inputs,
        "::mlir::ArrayRef<::mlir::NamedAttribute>":$attributes
      ), [{
@@ -2395,7 +2672,7 @@ diff --ruN a/stablehlo/stablehlo/dialect/Version.h b/stablehlo/stablehlo/dialect
  
    /// Return a Version representing the current VHLO dialect version.
 -  static Version getCurrentVersion() { return Version(1, 20, 1); }
-+  static Version getCurrentVersion() { return Version(1, 22, 0); }
++  static Version getCurrentVersion() { return Version(1, 23, 0); }
  
    /// Return a Version representing the minimum supported VHLO dialect version.
    static Version getMinimumVersion() { return Version(0, 9, 0); }
@@ -2414,12 +2691,13 @@ diff --ruN a/stablehlo/stablehlo/dialect/VhloBytecode.h b/stablehlo/stablehlo/di
 diff --ruN a/stablehlo/stablehlo/dialect/VhloDialect.td b/stablehlo/stablehlo/dialect/VhloDialect.td
 --- stablehlo/stablehlo/dialect/VhloDialect.td
 +++ stablehlo/stablehlo/dialect/VhloDialect.td
-@@ -59,10 +59,13 @@
+@@ -59,10 +59,14 @@
        1.18.0: Add `result_tilings` attribute to `custom_call` op.
        1.19.0: Add CollectiveReduceOp.
        1.20.0: Add has_dynamic_root to collective_broadcast op. Allow mixed fp8 operands in `convolution` and `dynamic_conv` ops.
 +      1.21.0: Add F8E4M3FN_F8E4M3FN_F32_X3 and F8E4M3FN_F8E4M3FN_F32_X4 dot algorithms.
 +      1.22.0: Add WEAKORDER comparison type to compare op.
++      1.23.0: Add `exp2` and `log2` ops.
    }];
  
    let useDefaultAttributePrinterParser = 0;
@@ -2562,6 +2840,34 @@ diff --ruN a/stablehlo/stablehlo/dialect/VhloOps.td b/stablehlo/stablehlo/dialec
    let arguments = (ins
      VHLO_AnyType:$lhs,
      VHLO_AnyType:$rhs,
+@@ -702,6 +703,13 @@
+   let results = (outs VHLO_AnyType:$result);
+ }
+ 
++def VHLO_Exp2OpV1 : VHLO_Op<"exp2_v1", "1.23.0", "current"> {
++  let arguments = (ins
++    VHLO_AnyType:$operand,
++    VHLO_AnyAttr:$result_accuracy);
++  let results = (outs VHLO_AnyType:$result);
++}
++
+ def VHLO_ExpOpV1 : VHLO_Op<"exponential_v1", "0.9.0", "1.8.0"> {
+   let arguments = (ins VHLO_AnyType:$operand);
+   let results = (outs VHLO_AnyType:$result);
+@@ -860,6 +868,13 @@
+   let results = (outs VHLO_AnyType:$result);
+ }
+ 
++def VHLO_Log2OpV1 : VHLO_Op<"log2_v1", "1.23.0", "current"> {
++  let arguments = (ins
++    VHLO_AnyType:$operand,
++    VHLO_AnyAttr:$result_accuracy);
++  let results = (outs VHLO_AnyType:$result);
++}
++
+ def VHLO_MapOpV1 : VHLO_Op<"map_v1", "0.9.0", "current"> {
+   let arguments = (ins
+     Variadic<VHLO_AnyType>:$inputs,
 diff --ruN a/stablehlo/stablehlo/integrations/c/ChloAttributes.cpp b/stablehlo/stablehlo/integrations/c/ChloAttributes.cpp
 --- stablehlo/stablehlo/integrations/c/ChloAttributes.cpp
 +++ stablehlo/stablehlo/integrations/c/ChloAttributes.cpp
@@ -3470,7 +3776,35 @@ diff --ruN a/stablehlo/stablehlo/reference/Configuration.h b/stablehlo/stablehlo
 diff --ruN a/stablehlo/stablehlo/reference/Element.cpp b/stablehlo/stablehlo/reference/Element.cpp
 --- stablehlo/stablehlo/reference/Element.cpp
 +++ stablehlo/stablehlo/reference/Element.cpp
-@@ -1159,6 +1159,40 @@
+@@ -844,6 +844,14 @@
+       });
+ }
+ 
++Element exp2(const Element& el) {
++  return mapWithUpcastToDouble(
++      el, [](double e) { return std::exp2(e); },
++      [](mlir::Complex<double> e) {
++        return std::pow(mlir::Complex<double>(2.0, 0.0), e);
++      });
++}
++
+ Element floor(const Element& el) {
+   APFloat val = el.getFloatValue();
+   val.roundToIntegral(APFloat::rmTowardNegative);
+@@ -890,6 +898,12 @@
+       });
+ }
+ 
++Element log2(const Element& el) {
++  return mapWithUpcastToDouble(
++      el, [](double e) { return std::log2(e); },
++      [](mlir::Complex<double> e) { return std::log(e) / std::log(2.0); });
++}
++
+ Element logistic(const Element& el) {
+   auto one = convert(el.getType(), 1.0);
+   return one / (one + exponential(-el));
+@@ -1159,6 +1173,40 @@
        [](mlir::Complex<double> e) { return std::tanh(e); });
  }
  
@@ -3514,7 +3848,27 @@ diff --ruN a/stablehlo/stablehlo/reference/Element.cpp b/stablehlo/stablehlo/ref
 diff --ruN a/stablehlo/stablehlo/reference/Element.h b/stablehlo/stablehlo/reference/Element.h
 --- stablehlo/stablehlo/reference/Element.h
 +++ stablehlo/stablehlo/reference/Element.h
-@@ -314,6 +314,12 @@
+@@ -237,6 +237,9 @@
+ /// Returns exponential_minus_one of Element object.
+ Element exponentialMinusOne(const Element& el);
+ 
++/// Returns exp2 of Element object.
++Element exp2(const Element& el);
++
+ /// Returns floor of Element object.
+ Element floor(const Element& e);
+ 
+@@ -253,6 +256,9 @@
+ /// Returns log1p of Element object.
+ Element logPlusOne(const Element& el);
+ 
++/// Returns log2 of Element object.
++Element log2(const Element& el);
++
+ /// Returns logistic of Element object.
+ Element logistic(const Element& el);
+ 
+@@ -314,6 +320,12 @@
  /// Returns tanh of Element object.
  Element tanh(const Element& e);
  
@@ -3864,7 +4218,29 @@ diff --ruN a/stablehlo/stablehlo/reference/Ops.cpp b/stablehlo/stablehlo/referen
        scope.add(op.getResult(), result);
      } else if (auto op = dyn_cast<ComplexOp>(operation)) {
        auto lhs = scope.findTensor(op.getLhs());
-@@ -1597,27 +1599,56 @@
+@@ -771,6 +773,10 @@
+       auto operand = scope.findTensor(op.getOperand());
+       auto result = expm1Op(operand, op.getType());
+       scope.add(op.getResult(), result);
++    } else if (auto op = dyn_cast<Exp2Op>(operation)) {
++      auto operand = scope.findTensor(op.getOperand());
++      auto result = exp2Op(operand, op.getType());
++      scope.add(op.getResult(), result);
+     } else if (auto op = dyn_cast<FftOp>(operation)) {
+       const auto operand = scope.findTensor(op.getOperand());
+       const auto fftType = op.getFftType();
+@@ -828,6 +834,10 @@
+       auto operand = scope.findTensor(op.getOperand());
+       auto result = log1pOp(operand, op.getType());
+       scope.add(op.getResult(), result);
++    } else if (auto op = dyn_cast<Log2Op>(operation)) {
++      auto operand = scope.findTensor(op.getOperand());
++      auto result = log2Op(operand, op.getType());
++      scope.add(op.getResult(), result);
+     } else if (auto op = dyn_cast<LogOp>(operation)) {
+       auto operand = scope.findTensor(op.getOperand());
+       auto result = logOp(operand, op.getType());
+@@ -1597,27 +1607,56 @@
  
  Tensor compareOp(const Tensor& lhs, const Tensor& rhs,
                   ComparisonDirection comparisonDirection,
@@ -3927,7 +4303,7 @@ diff --ruN a/stablehlo/stablehlo/reference/Ops.cpp b/stablehlo/stablehlo/referen
          break;
      }
    }
-@@ -1777,12 +1808,8 @@
+@@ -1777,12 +1816,8 @@
          dotGeneralOp(reversedLhsWindow, rhs, lhsContractingDimensions,
                       rhsContractingDimensions, resultType.getElementType());
  
@@ -3942,7 +4318,21 @@ diff --ruN a/stablehlo/stablehlo/reference/Ops.cpp b/stablehlo/stablehlo/referen
      Axes resultPermutation;
      resultPermutation.push_back(outputBatchDimension);
      resultPermutation.append(outputSpatialDimensions.begin(),
-@@ -1942,23 +1969,30 @@
+@@ -1932,6 +1967,13 @@
+   return result;
+ }
+ 
++Tensor exp2Op(const Tensor& operand, ShapedType resultType) {
++  Tensor result(resultType);
++  for (auto it = result.index_begin(); it != result.index_end(); ++it)
++    result.set(*it, exp2(operand.get(*it)));
++  return result;
++}
++
+ Tensor exponentialOp(const Tensor& operand, ShapedType resultType) {
+   Tensor result(resultType);
+   for (auto it = result.index_begin(); it != result.index_end(); ++it)
+@@ -1942,23 +1984,30 @@
  Tensor __fft_1d(const Tensor& operand, ShapedType resultType,
                  int64_t dimension) {
    Tensor result(resultType);
@@ -3980,7 +4370,7 @@ diff --ruN a/stablehlo/stablehlo/reference/Ops.cpp b/stablehlo/stablehlo/referen
        auto rotator = complex(cosine(phase), sine(phase));
        X = X + operand.get(operandIndex) * rotator;
      }
-@@ -1992,10 +2026,12 @@
+@@ -1992,10 +2041,12 @@
    } else if (fftType == FftType::IFFT) {
      // compute IFFT as conjugation, FFT, conjugation and normalization
      Tensor result(operand);
@@ -3995,7 +4385,7 @@ diff --ruN a/stablehlo/stablehlo/reference/Ops.cpp b/stablehlo/stablehlo/referen
        auto divisor =
            constantOp(mlir::DenseElementsAttr::get(resultType, divisorValue));
        result = divideOp(conjugate(__fft_1d(conjugate(result), resultType, *d)),
-@@ -2040,6 +2076,8 @@
+@@ -2040,6 +2091,8 @@
      replica = conjugate(reverseOp(replica, Axes{d_irfft}, replica.getType()));
      auto complexType =
          ComplexType::get(cast<FloatType>(resultType.getElementType()));
@@ -4004,7 +4394,7 @@ diff --ruN a/stablehlo/stablehlo/reference/Ops.cpp b/stablehlo/stablehlo/referen
      result = concatenateOp({result, replica}, Axis{d_irfft},
                             resultType.cloneWith(std::nullopt, complexType));
  
-@@ -2048,7 +2086,7 @@
+@@ -2048,7 +2101,7 @@
                      resultType);
  
      // normalize result of inverse transformation
@@ -4013,6 +4403,20 @@ diff --ruN a/stablehlo/stablehlo/reference/Ops.cpp b/stablehlo/stablehlo/referen
      auto divisor =
          constantOp(mlir::DenseElementsAttr::get(resultType, divisorValue));
      return divideOp(result, divisor, resultType);
+@@ -2189,6 +2242,13 @@
+   return result;
+ }
+ 
++Tensor log2Op(const Tensor& operand, ShapedType resultType) {
++  Tensor result(resultType);
++  for (auto it = result.index_begin(); it != result.index_end(); ++it)
++    result.set(*it, log2(operand.get(*it)));
++  return result;
++}
++
+ Tensor logOp(const Tensor& operand, ShapedType resultType) {
+   Tensor result(resultType);
+   for (auto it = result.index_begin(); it != result.index_end(); ++it)
 diff --ruN a/stablehlo/stablehlo/reference/Ops.h b/stablehlo/stablehlo/reference/Ops.h
 --- stablehlo/stablehlo/reference/Ops.h
 +++ stablehlo/stablehlo/reference/Ops.h
@@ -4024,6 +4428,22 @@ diff --ruN a/stablehlo/stablehlo/reference/Ops.h b/stablehlo/stablehlo/reference
                   ShapedType resultType);
  Tensor complexOp(const Tensor& lhs, const Tensor& rhs, ShapedType resultType);
  Tensor concatenateOp(ArrayRef<Tensor> inputs, Axis dimension,
+@@ -102,6 +103,7 @@
+                             ArrayRef<Tensor> startIndices,
+                             ShapedType resultType);
+ Tensor expm1Op(const Tensor& operand, ShapedType resultType);
++Tensor exp2Op(const Tensor& operand, ShapedType resultType);
+ Tensor exponentialOp(const Tensor& operand, ShapedType resultType);
+ Tensor fftOp(const Tensor& operand, const FftType fftType,
+              const ArrayRef<int64_t>& fftLength, ShapedType resultType);
+@@ -124,6 +126,7 @@
+ Tensor iotaOp(Axis iotaDimension, ShapedType resultType);
+ Tensor isFiniteOp(const Tensor& operand, ShapedType resultType);
+ Tensor log1pOp(const Tensor& operand, ShapedType resultType);
++Tensor log2Op(const Tensor& operand, ShapedType resultType);
+ Tensor logOp(const Tensor& operand, ShapedType resultType);
+ Tensor logisticOp(const Tensor& operand, ShapedType resultType);
+ Tensor mapOp(ArrayRef<Tensor> inputs, Region& computation, Process* process,
 diff --ruN a/stablehlo/stablehlo/reference/Process.cpp b/stablehlo/stablehlo/reference/Process.cpp
 --- stablehlo/stablehlo/reference/Process.cpp
 +++ stablehlo/stablehlo/reference/Process.cpp
@@ -12259,6 +12679,3464 @@ diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_22_0.ml
 +}
 diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_22_0.mlir.bc b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_22_0.mlir.bc
 Binary files /dev/null and stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_22_0.mlir.bc differ
+diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_23_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_23_0.mlir
+--- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_23_0.mlir
++++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_23_0.mlir
+@@ -0,0 +1,3452 @@
++// RUN: stablehlo-opt --mlir-print-op-generic %s.bc | FileCheck %s
++// RUN: stablehlo-translate --deserialize %s.bc | stablehlo-translate --serialize --target=1.23.0 | stablehlo-opt --mlir-print-op-generic | FileCheck %s
++// RUN: stablehlo-translate --deserialize %s.bc | stablehlo-opt > %t.0
++// RUN: stablehlo-opt --strip-debuginfo %s > %t.1
++// RUN: diff %t.0 %t.1
++// RUN: stablehlo-translate --serialize --target=1.23.0 --strip-debuginfo %s > %t.2
++// RUN: diff %s.bc %t.2
++// RUN: %if asserts %{ stablehlo-opt --stablehlo-legalize-to-vhlo -emit-bytecode -debug-only=vhlo-bytecode %s 2>&1 | FileCheck --check-prefix=CHECK-WARN %s %}
++// RUN: %if asserts %{ stablehlo-opt --stablehlo-legalize-to-vhlo -emit-bytecode %s | stablehlo-opt -debug-only=vhlo-bytecode 2>&1 | FileCheck --check-prefix=CHECK-WARN %s %}
++
++// CHECK-WARN-NOT: Not Implemented
++
++func.func private @mesh()
++
++// ============ ATTRIBUTES ============
++
++// CHECK-LABEL: "attr_comparison_direction_eq"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
++    comparison_direction = #stablehlo<comparison_direction EQ>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_direction_ne"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
++    comparison_direction = #stablehlo<comparison_direction NE>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_direction_ge"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
++    comparison_direction = #stablehlo<comparison_direction GE>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_direction_gt"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
++    comparison_direction = #stablehlo<comparison_direction GT>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_direction_le"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
++    comparison_direction = #stablehlo<comparison_direction LE>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_direction_lt"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
++    comparison_direction = #stablehlo<comparison_direction LT>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_type_notype"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction EQ>
++    // CHECK: compare_type = #vhlo<comparison_type_v1 NOTYPE>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_type_float"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction EQ>,
++    // CHECK: compare_type = #vhlo<comparison_type_v1 FLOAT>,
++    compare_type = #stablehlo<comparison_type FLOAT>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_type_totalorder"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction EQ>,
++    // CHECK: compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
++    compare_type = #stablehlo<comparison_type TOTALORDER>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_type_signed"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction EQ>,
++    // CHECK: compare_type = #vhlo<comparison_type_v1 SIGNED>,
++    compare_type = #stablehlo<comparison_type SIGNED>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_type_unsigned"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction EQ>,
++    // CHECK: compare_type = #vhlo<comparison_type_v1 UNSIGNED>,
++    compare_type = #stablehlo<comparison_type UNSIGNED>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "attr_comparison_type_weakorder"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_comparison_type_weakorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction EQ>,
++    // CHECK: compare_type = #vhlo<comparison_type_v1 WEAKORDER>,
++    compare_type = #stablehlo<comparison_type WEAKORDER>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// ConvDimensionNumbers aka #stablehlo.conv is covered below.
++
++// CHECK-LABEL: "attr_custom_call_api_version_unspecified"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "foo",
++    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_UNSPECIFIED>
++    api_version = 0 : i32
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "attr_custom_call_api_version_original"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "foo",
++    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>
++    api_version = 1 : i32
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "attr_custom_call_api_version_status_returning"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "foo",
++    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>
++    api_version = 2 : i32
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "foo",
++    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING_UNIFIED>
++    api_version = 3 : i32
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "attr_dict"
++// CHECK: #vhlo.dict_v1<{#vhlo.string_v1<"attr1"> = #vhlo.integer_v1<1 : i32>, #vhlo.string_v1<"attr2"> = #vhlo.integer_v1<2 : i32>}
++func.func @attr_dict() attributes {stablehlo.attr = {attr1 = 1 : i32, attr2 = 2 : i32}} {
++  return
++}
++
++// CHECK-LABEL: "attr_custom_call_api_version_typed_ffi"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++// CHECK: api_version = #vhlo<api_version_v1 API_VERSION_TYPED_FFI>
++// CHECK-SAME: backend_config = #vhlo.dict_v1<{#vhlo.string_v1<"bar"> = #vhlo.integer_v1<42 : i32>}>
++func.func @attr_custom_call_api_version_typed_ffi(%arg0: tensor<f32>) -> tensor<f32> {
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "foo",
++    backend_config= {bar = 42 : i32},
++    api_version = 4 : i32
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++
++// CHECK-LABEL: "attr_custom_call_api_version_typed_ffi_no_backend_config"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++// CHECK: api_version = #vhlo<api_version_v1 API_VERSION_TYPED_FFI>
++// CHECK-SAME: backend_config = #vhlo.dict_v1<{}>
++func.func @attr_custom_call_api_version_typed_ffi_no_backend_config(%arg0: tensor<f32>) -> tensor<f32> {
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "foo",
++    api_version = 4 : i32
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// DotDimensionNumbers aka #stablehlo.dot is covered below.
++
++// CHECK-LABEL: "attr_fft_type_fft"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
++  %0 = "stablehlo.fft"(%arg0) {
++    // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
++    fft_type = #stablehlo<fft_type FFT>,
++    fft_length = array<i64: 16>
++  } : (tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>>
++  func.return %0 : tensor<16xcomplex<f32>>
++}
++
++// CHECK-LABEL: "attr_fft_type_ifft"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
++  %0 = "stablehlo.fft"(%arg0) {
++    // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
++    fft_type = #stablehlo<fft_type IFFT>,
++    fft_length = array<i64: 16>
++  } : (tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>>
++  func.return %0 : tensor<16xcomplex<f32>>
++}
++
++// CHECK-LABEL: "attr_fft_type_rfft"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
++  %0 = "stablehlo.fft"(%arg0) {
++    // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
++    fft_type = #stablehlo<fft_type RFFT>,
++    fft_length = array<i64: 16>
++  } : (tensor<16xf32>) -> tensor<9xcomplex<f32>>
++  func.return %0 : tensor<9xcomplex<f32>>
++}
++
++// CHECK-LABEL: "attr_fft_type_irfft"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
++  %0 = "stablehlo.fft"(%arg0) {
++    // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
++    fft_type = #stablehlo<fft_type IRFFT>,
++    fft_length = array<i64: 16>
++  } : (tensor<9xcomplex<f32>>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "exponential_HIGHEST"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}
++func.func @exponential_HIGHEST(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
++  %0 = "stablehlo.exponential"(%arg0) {
++    // CHECK: result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 HIGHEST>>
++    result_accuracy = #stablehlo.result_accuracy<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #stablehlo.result_accuracy_mode<HIGHEST>>
++  } : (tensor<8x16xf32>) -> tensor<8x16xf32>
++  func.return %0 : tensor<8x16xf32>
++}
++
++// CHECK-LABEL: "exponential_TOLERANCE"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}
++func.func @exponential_TOLERANCE(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
++  %0 = "stablehlo.exponential"(%arg0) {
++    // CHECK: result_accuracy = #vhlo.result_accuracy_v1<atol = 1.000000e-05, rtol = 0.000000e+00, ulps = 1, mode = #vhlo<result_accuracy_mode_v1 TOLERANCE>>
++    result_accuracy = #stablehlo.result_accuracy<atol = 1.000000e-5, rtol = 0.000000e+00, ulps = 1, mode = #stablehlo.result_accuracy_mode<TOLERANCE>>
++  } : (tensor<8x16xf32>) -> tensor<8x16xf32>
++  func.return %0 : tensor<8x16xf32>
++}
++
++// CHECK-LABEL: "exponential_DEFAULT"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}
++func.func @exponential_DEFAULT(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
++  %0 = "stablehlo.exponential"(%arg0) {
++    // CHECK: result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>
++    result_accuracy = #stablehlo.result_accuracy<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #stablehlo.result_accuracy_mode<DEFAULT>>
++  } : (tensor<8x16xf32>) -> tensor<8x16xf32>
++  func.return %0 : tensor<8x16xf32>
++}
++
++// GatherDimensionNumbers aka #stablehlo.gather is covered below.
++
++// CHECK-LABEL: "attr_precision_config_default"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
++  %0 = "stablehlo.dot"(%arg0, %arg1) {
++    // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
++  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
++  func.return %0 : tensor<8x8xf32>
++}
++
++// CHECK-LABEL: "attr_precision_config_high"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
++  %0 = "stablehlo.dot"(%arg0, %arg1) {
++    // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
++    precision_config = [#stablehlo<precision HIGH>, #stablehlo<precision HIGH>]
++  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
++  func.return %0 : tensor<8x8xf32>
++}
++
++// CHECK-LABEL: "attr_precision_config_highest"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
++  %0 = "stablehlo.dot"(%arg0, %arg1) {
++    // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
++    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
++  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
++  func.return %0 : tensor<8x8xf32>
++}
++
++// CHECK-LABEL: "attr_rng_algorithm_default"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
++  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
++    // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
++    rng_algorithm = #stablehlo<rng_algorithm DEFAULT>
++  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
++  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
++}
++
++// CHECK-LABEL: "attr_rng_algorithm_three_fry"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
++  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
++    // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
++    rng_algorithm = #stablehlo<rng_algorithm THREE_FRY>
++  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
++  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
++}
++
++// CHECK-LABEL: "attr_rng_algorithm_philox"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
++  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
++    // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
++    rng_algorithm = #stablehlo<rng_algorithm PHILOX>
++  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
++  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
++}
++
++// CHECK-LABEL: "attr_rng_distribution_uniform"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
++  %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
++    // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
++    rng_distribution = #stablehlo<rng_distribution UNIFORM>
++  } : (tensor<f32>, tensor<f32>, tensor<0xindex>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "attr_rng_distribution_normal"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
++  %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
++    // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
++    rng_distribution = #stablehlo<rng_distribution NORMAL>
++  } : (tensor<f32>, tensor<f32>, tensor<0xindex>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
++
++// CHECK-LABEL: "attr_transpose_no_transpose"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
++  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
++    left_side = true,
++    lower = true,
++    unit_diagonal = true,
++    // transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
++    transpose_a = #stablehlo<transpose NO_TRANSPOSE>
++  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
++  func.return %0 : tensor<16x16xf32>
++}
++
++// CHECK-LABEL: "attr_transpose_transpose"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
++  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
++    left_side = true,
++    lower = true,
++    unit_diagonal = true,
++    // transpose_a = #vhlo<transpose_v1 TRANSPOSE>,
++    transpose_a = #stablehlo<transpose TRANSPOSE>
++  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
++  func.return %0 : tensor<16x16xf32>
++}
++
++// CHECK-LABEL: "attr_transpose_adjoint"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
++  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
++    left_side = true,
++    lower = true,
++    unit_diagonal = true,
++    // transpose_a = #vhlo<transpose_v1 ADJOINT>,
++    transpose_a = #stablehlo<transpose ADJOINT>
++  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
++  func.return %0 : tensor<16x16xf32>
++}
++
++// TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
++
++// CHECK-LABEL: "attr_type_extensions_bounds"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
++  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
++  func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
++}
++
++// CHECK-LABEL: "attr_frontend_attributes"
++func.func @attr_frontend_attributes(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: some.unregistered_attr
++  %1 = stablehlo.cosine %arg0 {some.unregistered_attr = 1 : i32} : tensor<f32>
++  return %1 : tensor<f32>
++}
++
++// Builtin attriubute tests
++
++// CHECK-LABEL: "byte_packed_boolean"
++func.func @byte_packed_boolean() -> (tensor<8xi1>, tensor<8xi1>, tensor<4xi1>, tensor<4xi1>, tensor<16xi1>) {
++  // CHECK: #vhlo.tensor_v1<dense<[true, false, false, false, false, false, false, false]
++  // CHECK-NEXT: #vhlo.tensor_v1<dense<true> : tensor<8xi1>>
++  // CHECK-NEXT: #vhlo.tensor_v1<dense<[true, false, false, false]> : tensor<4xi1>>
++  // CHECK-NEXT: #vhlo.tensor_v1<dense<true> : tensor<4xi1>>
++  // CHECK-NEXT: #vhlo.tensor_v1<dense<[true, false, false, false, false, false, false, false, true, false, false, false, false, false, false, false]> : tensor<16xi1>>
++  %c = stablehlo.constant dense<[true, false, false, false, false, false, false, false]> : tensor<8xi1>
++  %c_0 = stablehlo.constant dense<true> : tensor<8xi1>
++  %c_1 = stablehlo.constant dense<[true, false, false, false]> : tensor<4xi1>
++  %c_2 = stablehlo.constant dense<true> : tensor<4xi1>
++  %c_3 = stablehlo.constant dense<[true, false, false, false, false, false, false, false, true, false, false, false, false, false, false, false]> : tensor<16xi1>
++  return %c, %c_0, %c_1, %c_2, %c_3 : tensor<8xi1>, tensor<8xi1>, tensor<4xi1>, tensor<4xi1>, tensor<16xi1>
++}
++
++// ============ DEFAULTS ============
++
++// CHECK-LABEL: "default_all_gather"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
++  //               CHECK: "vhlo.all_gather_v2"(%[[ARG0]]) <{
++  //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
++  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
++  %0 = "stablehlo.all_gather"(%arg0) {
++    all_gather_dim = 1 : i64,
++    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
++  } : (tensor<16x8xf32>) -> tensor<16x16xf32>
++  func.return %0 : tensor<16x16xf32>
++}
++
++// CHECK-LABEL: "default_all_gather_variadic"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_all_gather_variadic(%arg0: tensor<16x8xf32>, %arg1: tensor<16x8xf32>) -> (tensor<16x16xf32>, tensor<16x16xf32>) {
++  %0:2 = "stablehlo.all_gather"(%arg0, %arg1) {
++    all_gather_dim = 1 : i64,
++    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
++  } : (tensor<16x8xf32>, tensor<16x8xf32>) -> (tensor<16x16xf32>, tensor<16x16xf32>)
++  func.return %0#0, %0#1 : tensor<16x16xf32>, tensor<16x16xf32>
++}
++
++// CHECK-LABEL: "default_all_reduce"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
++  //               CHECK: "vhlo.all_reduce_v2"(%[[ARG0]])
++  //          CHECK-SAME: <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
++  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
++  //          CHECK-SAME: }> ({
++  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++
++  %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 = dense<[[0], [1]]> : tensor<2x1xi64>
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "attr_replica_group_mesh_axes"
++func.func @attr_replica_group_mesh_axes(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.all_reduce_v2"
++  // CHECK-SAME: replica_groups = #vhlo.replica_group_mesh_axes_v1<
++  // CHECK-SAME:   mesh = #vhlo.string_v1<"mesh">,
++  // CHECK-SAME:   axes = #vhlo.array_v1<[#vhlo.axis_ref_v1<name = #vhlo.string_v1<"x">, sub_axis_info = #vhlo.sub_axis_info_v1<pre_size = 1, size = 2>>, #vhlo.axis_ref_v1<name = #vhlo.string_v1<"y">>]>
++  // CHECK-SAME: >
++  %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 = @mesh,
++      axes = [
++        #stablehlo.axis_ref<name = "x", sub_axis_info = (1)2>,
++        #stablehlo.axis_ref<name = "y">
++      ]
++    >
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++
++
++// CHECK-LABEL: "default_all_to_all"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
++  //               CHECK: "vhlo.all_to_all_v2"(%[[ARG0]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
++  //          CHECK-SAME:   split_count = #vhlo.integer_v1<4 : i64>
++  //          CHECK-SAME:   split_dimension = #vhlo.integer_v1<1 : i64>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x4x!vhlo.f32_v1>
++  %0 = "stablehlo.all_to_all"(%arg0) {
++    split_dimension = 1 : i64,
++    concat_dimension = 0 : i64,
++    split_count = 4 : i64,
++    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>
++  } : (tensor<4x16xf32>) -> tensor<16x4xf32>
++  func.return %0 : tensor<16x4xf32>
++}
++
++// CHECK-LABEL: "default_all_to_all_variadic"
++func.func @default_all_to_all_variadic(%arg0: tensor<4x16xf32>, %arg1: tensor<5x16xf32>) -> (tensor<16x4xf32>, tensor<20x4xf32>) {
++  %0:2 = "stablehlo.all_to_all"(%arg0, %arg1) {
++    split_dimension = 1 : i64,
++    concat_dimension = 0 : i64,
++    split_count = 4 : i64,
++    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
++    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>
++  } : (tensor<4x16xf32>, tensor<5x16xf32>) -> (tensor<16x4xf32>, tensor<20x4xf32>)
++  func.return %0#0, %0#1 : tensor<16x4xf32>, tensor<20x4xf32>
++}
++
++// CHECK-LABEL: "default_cholesky"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
++  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   lower = #vhlo.bool_v1<false>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
++  %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
++  func.return %0 : tensor<1x16x16xf32>
++}
++
++// CHECK-LABEL: "default_collective_permute"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
++  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
++  %0 = "stablehlo.collective_permute"(%arg0) {
++    source_target_pairs = dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>
++  } : (tensor<16x8xf32>) -> tensor<16x8xf32>
++  func.return %0 : tensor<16x8xf32>
++}
++
++// CHECK-LABEL: "default_collective_broadcast"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_collective_broadcast(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
++  //               CHECK: "vhlo.collective_broadcast_v2"(%[[ARG0]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<false>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
++  %0 = "stablehlo.collective_broadcast"(%arg0) {
++    replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
++  } : (tensor<16x8xf32>) -> tensor<16x8xf32>
++  func.return %0 : tensor<16x8xf32>
++}
++
++// CHECK-LABEL: "default_collective_reduce"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_collective_reduce(%arg0: tensor<f32>) -> tensor<f32> {
++  //               CHECK: "vhlo.collective_reduce_v1"(%[[ARG0]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<false>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>,
++  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
++  //          CHECK-SAME: }> ({
++  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.collective_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 = dense<[[0, 1]]> : tensor<1x2xi64>
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "default_compare"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
++  // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction EQ>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "default_composite"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_composite(%arg0: tensor<f32>) -> tensor<f32> {
++  //               CHECK: "vhlo.composite_v2"(%[[ARG0]]) <{
++  //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{}>
++  //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
++  //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
++  //          CHECK-SAME:   version = #vhlo.integer_v1<0 : i64>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.composite"(%arg0) {
++    name = "stablehlo.composite_target",
++    decomposition = @composite_target
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "default_convolution"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
++  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
++  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<2x2xi64>>,
++  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
++  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
++  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<false> : tensor<2xi1>>,
++  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x6x6x16x!vhlo.f32_v1>
++  %0 = "stablehlo.convolution"(%arg0, %arg1) {
++    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
++    feature_group_count = 1 : i64,
++    batch_group_count = 1 : i64
++  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32>
++  func.return %0 : tensor<1x6x6x16xf32>
++}
++
++// CHECK-LABEL: "default_custom_call"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
++  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
++  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
++  // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
++  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
++  // CHECK-SAME:   called_computations = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<false>,
++  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[]>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "foo"
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "default_dot_general"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
++  //      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   allow_imprecise_accumulation = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   lhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++  // CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   num_primitive_operations = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
++  // CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   rhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
++  // CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
++  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
++    dot_dimension_numbers = #stablehlo.dot<
++      lhs_batching_dimensions = [0],
++      lhs_contracting_dimensions = [2],
++      rhs_batching_dimensions = [0],
++      rhs_contracting_dimensions = [1]
++    >
++  } : (tensor<8x8x16xf32>, tensor<8x16x8xf32>) -> tensor<8x8x8xf32>
++  func.return %0 : tensor<8x8x8xf32>
++}
++
++// CHECK-LABEL: "dot_general_algorithm"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @dot_general_algorithm(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
++//      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
++// CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.f32_v1>,
++// CHECK-SAME:   allow_imprecise_accumulation = #vhlo.bool_v1<false>,
++// CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++// CHECK-SAME:   lhs_component_count = #vhlo.integer_v1<1 : i64>,
++// CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++// CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.tf31_v1>,
++// CHECK-SAME:   num_primitive_operations = #vhlo.integer_v1<1 : i64>,
++// CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
++// CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++// CHECK-SAME:   rhs_component_count = #vhlo.integer_v1<1 : i64>,
++// CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
++// CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.tf31_v1>
++// CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
++  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
++    dot_dimension_numbers = #stablehlo.dot<
++      lhs_batching_dimensions = [0],
++      lhs_contracting_dimensions = [2],
++      rhs_batching_dimensions = [0],
++      rhs_contracting_dimensions = [1]
++    >,
++    algorithm = #stablehlo.dot_algorithm<
++      lhs_precision_type = tf32,
++      rhs_precision_type = tf32,
++      accumulation_type = f32,
++      lhs_component_count = 1,
++      rhs_component_count = 1,
++      num_primitive_operations = 1,
++      allow_imprecise_accumulation = false
++    >
++  } : (tensor<8x8x16xf32>, tensor<8x16x8xf32>) -> tensor<8x8x8xf32>
++  func.return %0 : tensor<8x8x8xf32>
++}
++
++// CHECK-LABEL: "dot_general_algorithm_f8e4m3fn_x3"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @dot_general_algorithm_f8e4m3fn_x3(%arg0: tensor<8x8x16xbf16>, %arg1: tensor<8x16x8xbf16>) -> tensor<8x8x8xf32> {
++//      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
++// CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.f32_v1>,
++// CHECK-SAME:   allow_imprecise_accumulation = #vhlo.bool_v1<false>,
++// CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++// CHECK-SAME:   lhs_component_count = #vhlo.integer_v1<1 : i64>,
++// CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++// CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>,
++// CHECK-SAME:   num_primitive_operations = #vhlo.integer_v1<3 : i64>,
++// CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
++// CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++// CHECK-SAME:   rhs_component_count = #vhlo.integer_v1<1 : i64>,
++// CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
++// CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>
++// CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.bf16_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.bf16_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
++  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
++    dot_dimension_numbers = #stablehlo.dot<
++      lhs_batching_dimensions = [0],
++      lhs_contracting_dimensions = [2],
++      rhs_batching_dimensions = [0],
++      rhs_contracting_dimensions = [1]
++    >,
++    algorithm = #stablehlo.dot_algorithm<
++      lhs_precision_type = f8E4M3FN,
++      rhs_precision_type = f8E4M3FN,
++      accumulation_type = f32,
++      lhs_component_count = 1,
++      rhs_component_count = 1,
++      num_primitive_operations = 3,
++      allow_imprecise_accumulation = false
++    >
++  } : (tensor<8x8x16xbf16>, tensor<8x16x8xbf16>) -> tensor<8x8x8xf32>
++  func.return %0 : tensor<8x8x8xf32>
++}
++
++// CHECK-LABEL: "dot_general_algorithm_f8e4m3fn_x4"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @dot_general_algorithm_f8e4m3fn_x4(%arg0: tensor<8x8x16xbf16>, %arg1: tensor<8x16x8xbf16>) -> tensor<8x8x8xf32> {
++//      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
++// CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.f32_v1>,
++// CHECK-SAME:   allow_imprecise_accumulation = #vhlo.bool_v1<false>,
++// CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++// CHECK-SAME:   lhs_component_count = #vhlo.integer_v1<1 : i64>,
++// CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++// CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>,
++// CHECK-SAME:   num_primitive_operations = #vhlo.integer_v1<4 : i64>,
++// CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
++// CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++// CHECK-SAME:   rhs_component_count = #vhlo.integer_v1<1 : i64>,
++// CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
++// CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>
++// CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.bf16_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.bf16_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
++  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
++    dot_dimension_numbers = #stablehlo.dot<
++      lhs_batching_dimensions = [0],
++      lhs_contracting_dimensions = [2],
++      rhs_batching_dimensions = [0],
++      rhs_contracting_dimensions = [1]
++    >,
++    algorithm = #stablehlo.dot_algorithm<
++      lhs_precision_type = f8E4M3FN,
++      rhs_precision_type = f8E4M3FN,
++      accumulation_type = f32,
++      lhs_component_count = 1,
++      rhs_component_count = 1,
++      num_primitive_operations = 4,
++      allow_imprecise_accumulation = false
++    >
++  } : (tensor<8x8x16xbf16>, tensor<8x16x8xbf16>) -> tensor<8x8x8xf32>
++  func.return %0 : tensor<8x8x8xf32>
++}
++
++// CHECK-LABEL: "default_dynamic_broadcast_in_dim"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
++  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
++  // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_broadcast_in_dim"(%arg0, %arg1) {
++    broadcast_dimensions = array<i64: 0, 1>
++  } : (tensor<?x?xf32>, tensor<2xindex>) -> tensor<?x?xf32>
++  func.return %0 : tensor<?x?xf32>
++}
++
++// CHECK-LABEL: "default_dynamic_conv"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<2x2xi64>) -> tensor<1x?x?x16xf32> {
++  //      CHECK: "vhlo.dynamic_conv_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
++  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
++  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
++  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<false> : tensor<2xi1>>,
++  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x2x!vhlo.i64_v1>) -> !vhlo.tensor_v1<1x?x?x16x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_conv"(%arg0, %arg1, %arg2) {
++    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
++    feature_group_count = 1 : i64,
++    batch_group_count = 1 : i64
++  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>, tensor<2x2xi64>) -> tensor<1x?x?x16xf32>
++  func.return %0 : tensor<1x?x?x16xf32>
++}
++
++// CHECK-LABEL: "default_dynamic_gather"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
++  //      CHECK: "vhlo.dynamic_gather_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
++  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
++  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>, !vhlo.tensor_v1<3x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x8x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_gather"(%arg0, %arg1, %arg2) {
++    dimension_numbers = #stablehlo.gather<
++      offset_dims = [2],
++      collapsed_slice_dims = [0, 1],
++      start_index_map = [0, 1],
++      index_vector_dim = 2
++    >
++  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<1x5x8xf32>
++  func.return %0 : tensor<1x5x8xf32>
++}
++
++func.func @default_func(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK:      "vhlo.func_v1"() <{
++  // CHECK-SAME:   arg_attrs = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>>>,
++  // CHECK-SAME:   res_attrs = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
++  // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : () -> ()
++  func.return %arg0 : tensor<f32>
++}
++
++// CHECK-LABEL: "default_gather"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @default_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
++  //      CHECK: "vhlo.gather_v2"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
++  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
++  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<1> : tensor<3xi64>>,
++  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x1x!vhlo.f32_v1>
++  %0 = "stablehlo.gather"(%arg0, %arg1) {
++    dimension_numbers = #stablehlo.gather<
++      offset_dims = [2],
++      collapsed_slice_dims = [0, 1],
++      start_index_map = [0, 1],
++      index_vector_dim = 2
++    >,
++    slice_sizes = array<i64: 1, 1, 1>
++  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>) -> tensor<1x5x1xf32>
++  func.return %0 : tensor<1x5x1xf32>
++}
++
++// CHECK-LABEL: "default_infeed"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
++  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
++  //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
++  // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
++  //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
++  %0:2 = "stablehlo.infeed"(%arg0) : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
++  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
++}
++
++// CHECK-LABEL: "default_outfeed"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
++  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
++  %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
++  func.return %0 : !stablehlo.token
++}
++
++// CHECK-LABEL: "op_recv_with_source_target_pairs"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_recv_with_source_target_pairs(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
++  //      CHECK: "vhlo.recv_v2"(%[[ARG0]]) <{
++  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>,
++  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>>
++  // CHECK-SAME{LITERAL}: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
++  %0:2 = "stablehlo.recv"(%arg0) {
++    channel_handle = #stablehlo.channel_handle<handle = 0, type = 1>,
++    source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>
++  } : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
++  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
++}
++
++// CHECK-LABEL: "default_send"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
++  //      CHECK: "vhlo.send_v2"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>,
++  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>>
++  // CHECK-SAME{LITERAL}: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
++  %0 = "stablehlo.send"(%arg0, %arg1) {
++    channel_handle = #stablehlo.channel_handle<handle = 0, type = 1>,
++    source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>
++  } : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
++  func.return %0 : !stablehlo.token
++}
++
++// CHECK-LABEL: "default_reduce_scatter"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
++  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
++  //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
++  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
++  //          CHECK-SAME: }> ({
++  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.reduce_scatter"(%arg0) ({
++    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
++      %1 = "stablehlo.add"(%arg1, %arg2) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    scatter_dimension = 0 : i64,
++    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
++  } : (tensor<16xf32>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "default_reduce_window"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
++  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
++  //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
++  // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
++  //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
++  //          CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
++  //          CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
++  //          CHECK-SAME: }> ({
++  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.maximum_v1"(%[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<2x17x31x7x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<2x16x30x7x!vhlo.f32_v1>
++  %0 = "stablehlo.reduce_window"(%arg0, %arg1) ({
++    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
++      %1 = "stablehlo.maximum"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    window_dimensions = array<i64: 1, 2, 2, 1>
++  } : (tensor<2x17x31x7xf32>, tensor<f32>) -> tensor<2x16x30x7xf32>
++  func.return %0 : tensor<2x16x30x7xf32>
++}
++
++// CHECK-LABEL: "default_scatter"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
++  //      CHECK: "vhlo.scatter_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
++  // CHECK-SAME:   input_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
++  // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   scatter_dims_to_operand_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   scatter_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
++  // CHECK-SAME:   unique_indices = #vhlo.bool_v1<false>,
++  // CHECK-SAME:   update_window_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>
++  %0 = "stablehlo.scatter"(%arg0, %arg1, %arg2) ({
++    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
++      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    scatter_dimension_numbers = #stablehlo.scatter<
++      update_window_dims = [1],
++      inserted_window_dims = [0, 1],
++      scatter_dims_to_operand_dims = [0, 1],
++      index_vector_dim = 1
++    >
++  } : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) -> tensor<200x100x300xf32>
++  func.return %0 : tensor<200x100x300xf32>
++}
++
++// CHECK-LABEL: "default_select_and_scatter"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
++  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
++  // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
++  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG31:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG41:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL11:.*]] = "vhlo.compare_v1"(%[[ARG31]], %[[ARG41]]) <{compare_type = #vhlo<comparison_type_v1 TOTALORDER>, comparison_direction = #vhlo<comparison_direction_v1 GE>}>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL11]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
++  // CHECK-NEXT: }, {
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG32:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG42:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL12:.*]] = "vhlo.add_v1"(%[[ARG32]], %[[ARG42]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL12]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>, !vhlo.tensor_v1<10x23x23x64x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>
++  %0 = "stablehlo.select_and_scatter"(%arg0, %arg1, %arg2) ({
++    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
++      %1 = "stablehlo.compare"(%arg3, %arg4) {compare_type = #stablehlo<comparison_type TOTALORDER>, comparison_direction = #stablehlo<comparison_direction GE>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
++      "stablehlo.return"(%1) : (tensor<i1>) -> ()
++  }, {
++    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
++      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    window_dimensions = array<i64: 1, 2, 2, 1>
++  } : (tensor<10x24x24x64xf32>, tensor<10x23x23x64xf32>, tensor<f32>) -> tensor<10x24x24x64xf32>
++  func.return %0 : tensor<10x24x24x64xf32>
++}
++
++// CHECK-LABEL: "default_sort"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
++  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
++  // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.compare_v1"(%[[ARG1]], %[[ARG2]]) <{compare_type = #vhlo<comparison_type_v1 FLOAT>, comparison_direction = #vhlo<comparison_direction_v1 GT>}>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.sort"(%arg0) ({
++    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
++      %1 = "stablehlo.compare"(%arg1, %arg2) {compare_type = #stablehlo<comparison_type FLOAT>, comparison_direction = #stablehlo<comparison_direction GT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
++      "stablehlo.return"(%1) : (tensor<i1>) -> ()
++  }) : (tensor<16xf32>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// ============ OPS ============
++
++// CHECK-LABEL: "op_abs"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_add"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_after_all"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
++  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
++  %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
++  func.return %0 : !stablehlo.token
++}
++
++// CHECK-LABEL: "op_all_gather"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
++  //               CHECK: "vhlo.all_gather_v2"(%[[ARG0]]) <{
++  //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
++  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
++  %0 = "stablehlo.all_gather"(%arg0) {
++    all_gather_dim = 1 : i64,
++    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>,
++    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
++    use_global_device_ids
++  } : (tensor<16x8xf32>) -> tensor<16x16xf32>
++  func.return %0 : tensor<16x16xf32>
++}
++
++// CHECK-LABEL: "op_all_reduce"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
++  //               CHECK: "vhlo.all_reduce_v2"(%[[ARG0]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
++  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
++  //          CHECK-SAME: }> ({
++  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %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 = dense<[[0], [1]]> : tensor<2x1xi64>,
++    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
++    use_global_device_ids
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_all_reduce_with_promotable_types"
++func.func @op_all_reduce_with_promotable_types(%operand: tensor<f32>) -> tensor<f64> {
++  //  CHECK: "vhlo.all_reduce_v2"(%[[ARG0:.*]])
++  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
++  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
++  //  CHECK: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
++  %result = "stablehlo.all_reduce"(%operand) ({
++    ^bb0(%arg0: tensor<f64>, %arg1: tensor<f64>):
++      %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
++      "stablehlo.return"(%0) : (tensor<f64>) -> ()
++  }) {
++    replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
++    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
++    use_global_device_ids
++  } : (tensor<f32>) -> tensor<f64>
++
++  func.return %result : tensor<f64>
++}
++
++// CHECK-LABEL: "default_all_reduce_variadic"
++func.func @default_all_reduce_variadic(%arg0: tensor<f32>, %arg1: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
++  %0:2 = "stablehlo.all_reduce"(%arg0, %arg1) ({
++    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
++      %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> (tensor<f32>)
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
++  } : (tensor<f32>, tensor<f32>) -> (tensor<f32>, tensor<f32>)
++  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
++}
++
++// CHECK-LABEL: "op_all_to_all"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
++  //               CHECK: "vhlo.all_to_all_v2"(%[[ARG0]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
++  //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
++  //          CHECK-SAME:   split_count = #vhlo.integer_v1<4 : i64>
++  //          CHECK-SAME:   split_dimension = #vhlo.integer_v1<1 : i64>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x4x!vhlo.f32_v1>
++  %0 = "stablehlo.all_to_all"(%arg0) {
++    split_dimension = 1 : i64,
++    concat_dimension = 0 : i64,
++    split_count = 4 : i64,
++    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
++    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>
++  } : (tensor<4x16xf32>) -> tensor<16x4xf32>
++  func.return %0 : tensor<16x4xf32>
++}
++
++// CHECK-LABEL: "op_and"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
++  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "op_atan2"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_batch_norm_grad"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
++func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
++  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
++  // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
++  // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
++  %0:3 = "stablehlo.batch_norm_grad"(%arg0, %arg1, %arg2, %arg3, %arg4) {
++    epsilon = 0.001 : f32,
++    feature_index = 0 : i64
++  } : (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>)
++  func.return %0#0, %0#1, %0#2 : tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_batch_norm_inference"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
++func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
++  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
++  // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
++  // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
++  %0 = "stablehlo.batch_norm_inference"(%arg0, %arg1, %arg2, %arg3, %arg4) {
++    epsilon = 0.001 : f32,
++    feature_index = 0 : i64
++  } : (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>) -> tensor<16x16x16x16xf32>
++  func.return %0 : tensor<16x16x16x16xf32>
++}
++
++// CHECK-LABEL: "op_batch_norm_training"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
++  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
++  // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
++  %0:3 = "stablehlo.batch_norm_training"(%arg0, %arg1, %arg2) {
++    epsilon = 0.001 : f32,
++    feature_index = 0 : i64
++  } : (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>)
++  func.return %0#0, %0#1, %0#2 : tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_bitcast_convert"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
++  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_broadcast_in_dim"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
++  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
++  %0 = "stablehlo.broadcast_in_dim"(%arg0) {
++    broadcast_dimensions = array<i64: 1>
++  } : (tensor<16xf32>) -> tensor<16x16xf32>
++  func.return %0 : tensor<16x16xf32>
++}
++
++// CHECK-LABEL: "op_broadcast"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
++  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
++  %0 = "stablehlo.broadcast"(%arg0) {
++    broadcast_sizes = array<i64: 16>
++  } : (tensor<16xf32>) -> tensor<16x16xf32>
++  func.return %0 : tensor<16x16xf32>
++}
++
++// CHECK-LABEL: "op_case"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
++  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.case"(%arg0) ({
++    "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
++  }) : (tensor<i32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_cbrt"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.cbrt_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_ceil"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_cholesky"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
++  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   lower = #vhlo.bool_v1<true>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
++  %0 = "stablehlo.cholesky"(%arg0) {
++    lower = true
++  } : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
++  func.return %0 : tensor<1x16x16xf32>
++}
++
++// CHECK-LABEL: "op_clamp"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_count_leading_zeros"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
++  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
++  %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
++  func.return %0 : tensor<i32>
++}
++
++// CHECK-LABEL: "op_collective_broadcast"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_collective_broadcast(%arg0: tensor<16x8xf32>, %arg1: tensor<1xi32>) -> tensor<16x8xf32> {
++  //               CHECK: "vhlo.collective_broadcast_v2"(%[[ARG0]], %[[ARG1]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
++  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<true>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
++  %0 = "stablehlo.collective_broadcast"(%arg0, %arg1) {
++    replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
++    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
++    has_dynamic_root
++  } : (tensor<16x8xf32>, tensor<1xi32>) -> tensor<16x8xf32>
++  func.return %0 : tensor<16x8xf32>
++}
++
++// CHECK-LABEL: "op_collective_reduce"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_collective_reduce(%arg0: tensor<f32>, %arg1: tensor<1xi32>) -> tensor<f32> {
++  //               CHECK: "vhlo.collective_reduce_v1"(%[[ARG0]], %[[ARG1]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
++  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<true>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
++  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
++  //          CHECK-SAME: }> ({
++  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.collective_reduce"(%arg0, %arg1) ({
++    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
++      %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>,
++    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
++    use_global_device_ids,
++    has_dynamic_root
++  } : (tensor<f32>, tensor<1xi32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_collective_permute"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
++  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
++  %0 = "stablehlo.collective_permute"(%arg0) {
++    source_target_pairs = dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>,
++    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>
++  } : (tensor<16x8xf32>) -> tensor<16x8xf32>
++  func.return %0 : tensor<16x8xf32>
++}
++
++// CHECK-LABEL: "op_compare"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
++  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
++  // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  %0 = "stablehlo.compare"(%arg0, %arg1) {
++    comparison_direction = #stablehlo<comparison_direction EQ>,
++    compare_type = #stablehlo<comparison_type TOTALORDER>
++  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "op_complex"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
++  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
++  %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
++  func.return %0 : tensor<complex<f32>>
++}
++
++// CHECK-LABEL: "op_composite"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_composite(%arg0: tensor<f32>) -> tensor<f32> {
++  //               CHECK: "vhlo.composite_v2"(%[[ARG0]]) <{
++  //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{#vhlo.string_v1<"my_int"> = #vhlo.integer_v1<1 : i64>, #vhlo.string_v1<"my_string"> = #vhlo.string_v1<"foo">}>
++  //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
++  //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
++  //          CHECK-SAME:   version = #vhlo.integer_v1<1 : i32>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.composite"(%arg0) {
++    name = "stablehlo.composite_target",
++    decomposition = @composite_target,
++    version = 1 : i32,
++    composite_attributes = {
++      my_int = 1 : i64,
++      my_string = "foo"
++    }
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "composite_regions"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @composite_regions(%arg0: tensor<f32>) -> tensor<f32> {
++  //               CHECK: "vhlo.composite_v2"(%[[ARG0]]) <{
++  //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{}>
++  //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
++  //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
++  //          CHECK-SAME:   version = #vhlo.integer_v1<1 : i32>
++  //          CHECK-SAME: }> ({
++  //          CHECK-NEXT: ^bb0(%[[ARG1:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //          CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.composite"(%arg0) ({
++    ^bb0(%arg1: tensor<f32>):
++      "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
++  }) {
++    name = "stablehlo.composite_target",
++    decomposition = @composite_target,
++    version = 1 : i32
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_concatenate"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
++  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.concatenate"(%arg0, %arg1) {
++    dimension = 0 : i64
++  } : (tensor<8xf32>, tensor<8xf32>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_constant"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
++  //      CHECK: "vhlo.constant_v1"() <{
++  // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
++  // CHECK-SAME: }> : () -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.constant"() {
++    value = dense<0.0> : tensor<f32>
++  } : () -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_convert"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
++  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_convolution"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
++  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
++  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<2x2xi64>>,
++  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
++  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
++  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<true> : tensor<2xi1>>,
++  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x7x7x16x!vhlo.f32_v1>
++  %0 = "stablehlo.convolution"(%arg0, %arg1) {
++    window_strides = array<i64: 2, 2>,
++    padding = dense<1> : tensor<2x2xi64>,
++    lhs_dilation = array<i64: 2, 2>,
++    rhs_dilation = array<i64: 2, 2>,
++    window_reversal = array<i1: true, true>,
++    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
++    feature_group_count = 1 : i64,
++    batch_group_count = 1 : i64,
++    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
++  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32>
++  func.return %0 : tensor<1x7x7x16xf32>
++}
++
++// CHECK-LABEL: "op_cosine"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.cosine_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_create_token"
++func.func @op_create_token() -> !stablehlo.token {
++  // CHECK: "vhlo.create_token_v1"() : () -> !vhlo.token_v1
++  %0 = "stablehlo.create_token"() : () -> !stablehlo.token
++  func.return %0 : !stablehlo.token
++}
++
++// CHECK-LABEL: "op_cross_replica_sum"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
++  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
++  //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.cross-replica-sum"(%arg0) {
++    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_custom_call"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
++  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
++  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
++  // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
++  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
++  // CHECK-SAME:   called_computations = #vhlo.array_v1<[#vhlo.string_v1<"foo">]>,
++  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<> : tensor<0xindex>>]>,
++  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[
++  // CHECK-SAME:     #vhlo.output_operand_alias_v1<
++  // CHECK-SAME:       outputTupleIndices = [],
++  // CHECK-SAME:       operandIndex = 0,
++  // CHECK-SAME:       operandTupleIndices = []>]>,
++  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<> : tensor<0xindex>>]>,
++  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[]>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "foo",
++    has_side_effect = true,
++    backend_config = "\08\03\1A\02",
++    api_version = 2 : i32,
++    called_computations = [@foo],
++    operand_layouts = [dense<> : tensor<0xindex>],
++    output_operand_aliases = [
++      #stablehlo.output_operand_alias<output_tuple_indices = [],
++                                 operand_index = 0,
++                                 operand_tuple_indices = []>],
++    result_layouts = [dense<> : tensor<0xindex>]
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_custom_call_empty_result_layout"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func public @op_custom_call_empty_result_layout(%arg0: tensor<i64>) -> tensor<i64> {
++  // %0 = "vhlo.custom_call_v2"(%arg0) <{>}> : (!vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tuple_v1<>
++  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
++  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
++  // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
++  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"empty_output">,
++  // CHECK-SAME:   called_computations = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<> : tensor<0xindex>>]>,
++  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[]>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tuple_v1<>
++  %0 = "stablehlo.custom_call"(%arg0) <{
++    api_version = 2 : i32,
++    call_target_name = "empty_output",
++    has_side_effect = true,
++    operand_layouts = [dense<> : tensor<0xindex>],
++    result_layouts = []
++  }> : (tensor<i64>) -> tuple<>
++  return %arg0 : tensor<i64>
++}
++
++// CHECK-LABEL: "op_custom_call_with_result_tilings"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_custom_call_with_result_tilings(%arg0: tensor<64x256xf32>) -> tensor<64x256xf32> {
++  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
++  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
++  // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
++  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
++  // CHECK-SAME:   called_computations = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<false>,
++  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<[1, 0]> : tensor<2xindex>>]>,
++  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[]>,
++  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<[1, 0]> : tensor<2xindex>>]>,
++  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[#vhlo.array_v1<[#vhlo.tensor_v1<dense<[8, 128]> : tensor<2xindex>>]>]>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<64x256x!vhlo.f32_v1>) -> !vhlo.tensor_v1<64x256x!vhlo.f32_v1>
++  %0 = "stablehlo.custom_call"(%arg0) {
++    api_version = 2 : i32,
++    call_target_name = "foo",
++    operand_layouts = [dense<[1, 0]> : tensor<2xindex>],
++    result_layouts = [dense<[1, 0]> : tensor<2xindex>],
++    result_tilings = [[dense<[8, 128]> : tensor<2xindex>]]
++  } : (tensor<64x256xf32>) -> tensor<64x256xf32>
++  func.return %0 : tensor<64x256xf32>
++}
++
++// CHECK-LABEL: "op_divide"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_dot_general"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
++  //      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   allow_imprecise_accumulation = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   lhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++  // CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   num_primitive_operations = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
++  // CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   rhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
++  // CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
++  // CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
++  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
++    dot_dimension_numbers = #stablehlo.dot<
++      lhs_batching_dimensions = [0],
++      lhs_contracting_dimensions = [2],
++      rhs_batching_dimensions = [0],
++      rhs_contracting_dimensions = [1]
++    >,
++    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
++  } : (tensor<8x8x16xf32>, tensor<8x16x8xf32>) -> tensor<8x8x8xf32>
++  func.return %0 : tensor<8x8x8xf32>
++}
++
++// CHECK-LABEL: "op_dot"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
++  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
++  %0 = "stablehlo.dot"(%arg0, %arg1) {
++    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
++  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
++  func.return %0 : tensor<8x8xf32>
++}
++
++// CHECK-LABEL: "op_dynamic_broadcast_in_dim"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
++  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_broadcast_in_dim"(%arg0, %arg1) {
++    broadcast_dimensions = array<i64: 0, 1>,
++    known_expanding_dimensions = array<i64: 0>,
++    known_nonexpanding_dimensions = array<i64: 1>
++  } : (tensor<?x?xf32>, tensor<2xindex>) -> tensor<?x?xf32>
++  func.return %0 : tensor<?x?xf32>
++}
++
++// CHECK-LABEL: "op_dynamic_conv"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<2x2xi64>) -> tensor<1x?x?x16xf32> {
++  //      CHECK: "vhlo.dynamic_conv_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
++  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
++  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
++  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<true> : tensor<2xi1>>,
++  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x2x!vhlo.i64_v1>) -> !vhlo.tensor_v1<1x?x?x16x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_conv"(%arg0, %arg1, %arg2) {
++    window_strides = array<i64: 2, 2>,
++    lhs_dilation = array<i64: 2, 2>,
++    rhs_dilation = array<i64: 2, 2>,
++    window_reversal = array<i1: true, true>,
++    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
++    feature_group_count = 1 : i64,
++    batch_group_count = 1 : i64,
++    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
++  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>, tensor<2x2xi64>) -> tensor<1x?x?x16xf32>
++  func.return %0 : tensor<1x?x?x16xf32>
++}
++
++// CHECK-LABEL: "op_dynamic_gather"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
++  //      CHECK: "vhlo.dynamic_gather_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
++  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>, !vhlo.tensor_v1<3x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x8x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_gather"(%arg0, %arg1, %arg2) {
++    dimension_numbers = #stablehlo.gather<
++      offset_dims = [2],
++      collapsed_slice_dims = [0, 1],
++      start_index_map = [0, 1],
++      index_vector_dim = 2
++    >,
++    indices_are_sorted = true
++  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<1x5x8xf32>
++  func.return %0 : tensor<1x5x8xf32>
++}
++
++// CHECK-LABEL: "op_dynamic_gather_with_batching_dims"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_dynamic_gather_with_batching_dims(%arg0 : tensor<5x2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<4xi32>) -> tensor<1x5x8xf32> {
++  //      CHECK: "vhlo.dynamic_gather_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<5x2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>, !vhlo.tensor_v1<4x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x8x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_gather"(%arg0, %arg1, %arg2) {
++    dimension_numbers = #stablehlo.gather<
++      offset_dims = [2],
++      collapsed_slice_dims = [1, 2],
++      operand_batching_dims = [0],
++      start_indices_batching_dims = [1],
++      start_index_map = [1, 2],
++      index_vector_dim = 2
++    >,
++    indices_are_sorted = true
++  } : (tensor<5x2x4x9xf32>, tensor<1x5x2xi32>, tensor<4xi32>) -> tensor<1x5x8xf32>
++  func.return %0 : tensor<1x5x8xf32>
++}
++
++// CHECK-LABEL: "op_dynamic_iota"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
++  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_iota"(%arg0) {
++    iota_dimension = 0 : i64
++  } : (tensor<1xindex>) -> tensor<?xf32>
++  func.return %0 : tensor<?xf32>
++}
++
++// CHECK-LABEL: "op_dynamic_pad"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
++func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
++  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
++  func.return %0 : tensor<?xf32>
++}
++
++// CHECK-LABEL: "op_dynamic_reshape"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
++  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
++  func.return %0 : tensor<?x?xf32>
++}
++
++// CHECK-LABEL: "op_dynamic_slice"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
++  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
++    slice_sizes = array<i64: 4>
++  } : (tensor<16xf32>, tensor<i64>) -> tensor<4xf32>
++  func.return %0 : tensor<4xf32>
++}
++
++// CHECK-LABEL: "op_dynamic_update_slice"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
++  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_einsum"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
++  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
++  %0 = "stablehlo.einsum"(%arg0, %arg1) {
++    einsum_config = "ab,bc->ac"
++  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
++  func.return %0 : tensor<8x8xf32>
++}
++
++// CHECK-LABEL: "op_exponential_minus_one"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.exponential_minus_one_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_exponential"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.exponential_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_exp2"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_exp2(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.exp2_v1"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.exp2"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_fft"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
++  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
++  // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
++  %0 = "stablehlo.fft"(%arg0) {
++    fft_type = #stablehlo<fft_type FFT>,
++    fft_length = array<i64: 16>
++  } : (tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>>
++  func.return %0 : tensor<16xcomplex<f32>>
++}
++
++// CHECK-LABEL: "op_floor"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++func.func private @op_func(%arg0: tensor<f32> {stablehlo.arg = "0"}) -> (tensor<f32> {stablehlo.result = "0"}) {
++  // CHECK:      "vhlo.func_v1"() <{
++  // CHECK-SAME:   arg_attrs = #vhlo.array_v1<[#vhlo.dict_v1<{#vhlo.string_v1<"stablehlo.arg"> = #vhlo.string_v1<"0">}>]>,
++  // CHECK-SAME:   function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>>>,
++  // CHECK-SAME:   res_attrs = #vhlo.array_v1<[#vhlo.dict_v1<{#vhlo.string_v1<"stablehlo.result"> = #vhlo.string_v1<"0">}>]>,
++  // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
++  // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : () -> ()
++
++  func.return %arg0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_gather"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
++  //      CHECK: "vhlo.gather_v2"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
++  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<1> : tensor<3xi64>>,
++  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x1x!vhlo.f32_v1>
++  %0 = "stablehlo.gather"(%arg0, %arg1) {
++    dimension_numbers = #stablehlo.gather<
++      offset_dims = [2],
++      collapsed_slice_dims = [0, 1],
++      start_index_map = [0, 1],
++      index_vector_dim = 2
++    >,
++    slice_sizes = array<i64: 1, 1, 1>,
++    indices_are_sorted = true
++  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>) -> tensor<1x5x1xf32>
++  func.return %0 : tensor<1x5x1xf32>
++}
++
++// CHECK-LABEL: "op_gather_with_batching_dims"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_gather_with_batching_dims(%arg0 : tensor<5x2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
++  //      CHECK: "vhlo.gather_v2"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
++  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
++  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<5x2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x1x!vhlo.f32_v1>
++  %0 = "stablehlo.gather"(%arg0, %arg1) {
++    dimension_numbers = #stablehlo.gather<
++      offset_dims = [2],
++      collapsed_slice_dims = [1, 2],
++      operand_batching_dims = [0],
++      start_indices_batching_dims = [1],
++      start_index_map = [1, 2],
++      index_vector_dim = 2
++    >,
++    slice_sizes = array<i64: 1, 1, 1, 1>,
++    indices_are_sorted = true
++  } : (tensor<5x2x4x9xf32>, tensor<1x5x2xi32>) -> tensor<1x5x1xf32>
++  func.return %0 : tensor<1x5x1xf32>
++}
++
++// CHECK-LABEL: "op_get_dimension_size"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
++  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
++  %0 = "stablehlo.get_dimension_size"(%arg0) {
++    dimension = 0 : i64
++  } : (tensor<?xf32>) -> tensor<i32>
++  func.return %0 : tensor<i32>
++}
++
++// CHECK-LABEL: "op_get_tuple_element"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
++  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
++  // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.get_tuple_element"(%arg0) {
++    index = 0 : i32
++  } : (tuple<tensor<f32>, tensor<i32>>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_if"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
++  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }, {
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.if"(%arg0) ({
++    "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
++  }, {
++    "stablehlo.return"(%arg2) : (tensor<f32>) -> ()
++  }) : (tensor<i1>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_imag"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
++  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_infeed"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
++  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
++  //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
++  // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
++  //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
++  %0:2 = "stablehlo.infeed"(%arg0) {
++    infeed_config = "foo",
++    layout = [[]]
++  } : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
++  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
++}
++
++// CHECK-LABEL: "op_iota"
++func.func @op_iota() -> tensor<16xf32> {
++  //      CHECK: "vhlo.iota_v1"() <{
++  // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME: }> : () -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.iota"() {
++    iota_dimension = 0 : i64
++  } : () -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_is_finite"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
++  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "op_log"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.log_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_log2"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_log2(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.log2_v1"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.log2"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_log_plus_one"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.log_plus_one_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_logistic"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.logistic_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_map"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
++  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.abs_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.map"(%arg0) ({
++    ^bb0(%arg1: tensor<f32>):
++      %1 = "stablehlo.abs"(%arg1) : (tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    dimensions = array<i64: 0>
++  } : (tensor<16xf32>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_maximum"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_minimum"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_multiply"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_negate"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_not"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
++  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "op_optimization_barrier"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_or"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
++  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "op_outfeed"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
++  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
++  %0 = "stablehlo.outfeed"(%arg0, %arg1) {
++    outfeed_config = "foo"
++  } : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
++  func.return %0 : !stablehlo.token
++}
++
++// CHECK-LABEL: "op_pad"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
++  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
++  // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
++  // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.pad"(%arg0, %arg1) {
++    edge_padding_high = array<i64: 4>,
++    edge_padding_low = array<i64: 4>,
++    interior_padding = array<i64: 0>
++  } : (tensor<8xf32>, tensor<f32>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_popcnt"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
++  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
++  %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
++  func.return %0 : tensor<i32>
++}
++
++// CHECK-LABEL: "op_power"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_real_dynamic_slice"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
++func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
++  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
++  %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
++  func.return %0 : tensor<?xf32>
++}
++
++// CHECK-LABEL: "op_real"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
++  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_recv_no_source_target_pairs"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_recv_no_source_target_pairs(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
++  //      CHECK: "vhlo.recv_v2"(%[[ARG0]]) <{
++  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
++  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>,
++  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<> : tensor<0xi64>
++  // CHECK-SAME{LITERAL}: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
++  %0:2 = "stablehlo.recv"(%arg0) {
++    channel_handle = #stablehlo.channel_handle<handle = 0, type = 3>,
++    is_host_transfer = true
++  } : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
++  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
++}
++
++// CHECK-LABEL: "op_reduce"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
++  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
++  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.reduce"(%arg0, %arg1) ({
++    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
++      %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    dimensions = array<i64: 0>
++  } : (tensor<16xf32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_reduce_precision"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
++  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
++  // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.reduce_precision"(%arg0) {
++    exponent_bits = 8 : i32,
++    mantissa_bits = 10 : i32
++  } : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK_lABEL: "op_reduce_with_promotable_types"
++func.func @op_reduce_with_promotable_types(%arg0: tensor<4x4xf32>, %arg1 : tensor<f32>)
++    -> (tensor<4xf64>) {
++  //  CHECK: "vhlo.reduce_v1"(%[[ARG0:.*]], %[[ARG1:.*]])
++  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
++  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
++  //  CHECK: }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x!vhlo.f64_v1>
++  %0 = "stablehlo.reduce"(%arg0, %arg1) ({
++  ^bb0(%arg2: tensor<f64>, %arg3: tensor<f64> ):
++    %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f64>, tensor<f64>) -> tensor<f64>
++    "stablehlo.return"(%1) : (tensor<f64>) -> ()
++
++  }) {dimensions = array<i64: 0>} : (tensor<4x4xf32>, tensor<f32>) -> tensor<4xf64>
++
++  func.return %0: tensor<4xf64>
++}
++
++// CHECK-LABEL: "op_reduce_scatter"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
++  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
++  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
++  //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
++  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
++  //          CHECK-SAME: }> ({
++  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.reduce_scatter"(%arg0) ({
++    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
++      %1 = "stablehlo.add"(%arg1, %arg2) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    scatter_dimension = 0 : i64,
++    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>,
++    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
++    use_global_device_ids
++  } : (tensor<16xf32>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK_lABEL: "op_reduce_scatter_with_promotable_types"
++func.func @op_reduce_scatter_with_promotable_types(%data: tensor<4x16xf32>) -> tensor<4x4xf64> {
++  //  CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0:.*]])
++  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
++  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
++  //  CHECK: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f64_v1>
++  %0 = "stablehlo.reduce_scatter"(%data) ({
++    ^bb0(%arg2: tensor<f64>, %arg3: tensor<f64>):
++    %1 = stablehlo.add %arg2, %arg3 : tensor<f64>
++    "stablehlo.return"(%1) : (tensor<f64>) -> ()
++  }) {replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
++      scatter_dimension = 1 : i64,
++      channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
++      use_global_device_ids} : (tensor<4x16xf32>) -> tensor<4x4xf64>
++  func.return %0 : tensor<4x4xf64>
++}
++
++
++// CHECK-LABEL: "op_reduce_window"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
++  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
++  //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
++  // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
++  //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
++  //          CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
++  //          CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 4, 4, 1]> : tensor<4xi64>>
++  //          CHECK-SAME: }> ({
++  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.maximum_v1"(%[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<2x17x31x7x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<2x9x16x7x!vhlo.f32_v1>
++  %0 = "stablehlo.reduce_window"(%arg0, %arg1) ({
++    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
++      %1 = "stablehlo.maximum"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    window_dimensions = array<i64: 1, 2, 2, 1>,
++    window_strides = array<i64: 1, 4, 4, 1>,
++    base_dilations = array<i64: 1, 2, 2, 1>,
++    window_dilations = array<i64: 1, 2, 2, 1>,
++    padding = dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>
++  } : (tensor<2x17x31x7xf32>, tensor<f32>) -> tensor<2x9x16x7xf32>
++  func.return %0 : tensor<2x9x16x7xf32>
++}
++
++// CHECK-LABEL: "op_reduce_window_with_promotable_types"
++func.func @op_reduce_window_with_promotable_types(%arg0: tensor<4x2xf32>,
++    %arg1: tensor<4x2xf32>, %init0: tensor<f32>, %init1: tensor<f32>) ->
++    (tensor<2x2xf64>, tensor<2x2xf32>) {
++  //  CHECK: "vhlo.reduce_window_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]], %[[ARG3:.*]])
++  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]], %[[VAL2:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  //  CHECK: }) : (!vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<2x2x!vhlo.f64_v1>, !vhlo.tensor_v1<2x2x!vhlo.f32_v1>)
++  %0:2 = "stablehlo.reduce_window"(%arg0, %arg1, %init0, %init1) ({
++         ^bb0(%a0: tensor<f64>, %a1: tensor<f32>, %b0: tensor<f64>,
++                %b1: tensor<f32>):
++              %2 = stablehlo.add %a0, %b0 : tensor<f64>
++              %3 = stablehlo.add %a1, %b1 : tensor<f32>
++              "stablehlo.return"(%2,%3) : (tensor<f64>, tensor<f32>) -> ()
++            })
++         { padding = dense<[[2, 2], [0, 0]]> : tensor<2x2xi64>,
++         window_dimensions = array<i64: 5, 1>,
++         window_strides = array<i64: 3, 1> }
++         : (tensor<4x2xf32>, tensor<4x2xf32>, tensor<f32>, tensor<f32>) ->
++              (tensor<2x2xf64>, tensor<2x2xf32>)
++  func.return %0#0, %0#1 : tensor<2x2xf64>, tensor<2x2xf32>
++}
++
++// CHECK-LABEL: "op_remainder"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_replica_id"
++func.func @op_replica_id() -> tensor<ui32> {
++  // CHECK: "vhlo.replica_id_v1"() : () -> !vhlo.tensor_v1<!vhlo.ui32_v1>
++  %0 = "stablehlo.replica_id"() : () -> tensor<ui32>
++  func.return %0 : tensor<ui32>
++}
++
++// CHECK-LABEL: "op_partition_id"
++func.func @op_partition_id() -> tensor<ui32> {
++  // CHECK: "vhlo.partition_id_v1"() : () -> !vhlo.tensor_v1<!vhlo.ui32_v1>
++  %0 = "stablehlo.partition_id"() : () -> tensor<ui32>
++  func.return %0 : tensor<ui32>
++}
++
++// CHECK-LABEL: "op_reshape"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
++  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
++  %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
++  func.return %0 : tensor<4x4xf32>
++}
++
++// CHECK-LABEL: "op_return"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
++  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.case"(%arg0) ({
++    "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
++  }) : (tensor<i32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_reverse"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
++  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.reverse"(%arg0) {
++    dimensions = array<i64: 0>
++  } : (tensor<16xf32>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_rng_bit_generator"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
++  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
++  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
++    rng_algorithm = #stablehlo<rng_algorithm PHILOX>
++  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
++  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
++}
++
++// CHECK-LABEL: "op_rng"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
++  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
++    rng_distribution = #stablehlo<rng_distribution NORMAL>
++  } : (tensor<f32>, tensor<f32>, tensor<0xindex>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_round_nearest_afz"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_round_nearest_even"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_rsqrt"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.rsqrt_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_scatter"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
++  //      CHECK: "vhlo.scatter_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   input_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
++  // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   scatter_dims_to_operand_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
++  // CHECK-SAME:   scatter_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
++  // CHECK-SAME:   unique_indices = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   update_window_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>
++  %0 = "stablehlo.scatter"(%arg0, %arg1, %arg2) ({
++    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
++      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    scatter_dimension_numbers = #stablehlo.scatter<
++      update_window_dims = [1],
++      inserted_window_dims = [0, 1],
++      scatter_dims_to_operand_dims = [0, 1],
++      index_vector_dim = 1
++    >,
++    indices_are_sorted = true,
++    unique_indices = true
++  } : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) -> tensor<200x100x300xf32>
++  func.return %0 : tensor<200x100x300xf32>
++}
++
++// CHECK-LABEL: "op_scatter_with_batching_dims"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_scatter_with_batching_dims(%arg0: tensor<10x200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<10x200x100x300xf32> {
++  //      CHECK: "vhlo.scatter_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   input_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   scatter_dims_to_operand_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
++  // CHECK-SAME:   scatter_indices_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   unique_indices = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   update_window_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<10x200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x200x100x300x!vhlo.f32_v1>
++  %0 = "stablehlo.scatter"(%arg0, %arg1, %arg2) ({
++    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
++      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    scatter_dimension_numbers = #stablehlo.scatter<
++      update_window_dims = [1],
++      inserted_window_dims = [1, 2],
++      input_batching_dims = [0],
++      scatter_dims_to_operand_dims = [1, 2],
++      scatter_indices_batching_dims = [0],
++      index_vector_dim = 1
++    >,
++    indices_are_sorted = true,
++    unique_indices = true
++  } : (tensor<10x200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) -> tensor<10x200x100x300xf32>
++  func.return %0 : tensor<10x200x100x300xf32>
++}
++
++// CHECK_lABEL: "op_scatter_with_promotable_types"
++func.func @op_scatter_with_promotable_types(%input_tensor: tensor<200x100x300xf32>,
++    %scatter_indices: tensor<10x2xi32>, %updates: tensor<10x300xf32>) ->
++      tensor<200x100x300xf64> {
++  //  CHECK: "vhlo.scatter_v2"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]])
++  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
++  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
++  //  CHECK: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f64_v1>
++  %0 = "stablehlo.scatter" (%input_tensor, %scatter_indices, %updates) ({
++  ^bb0(%lhs: tensor<f64>, %rhs: tensor<f64>):
++    %add = stablehlo.add %lhs, %rhs : tensor<f64>
++    "stablehlo.return"(%add) : (tensor<f64>) -> ()
++  }) {
++    scatter_dimension_numbers = #stablehlo.scatter<
++      update_window_dims = [1],
++      inserted_window_dims = [0, 1],
++      scatter_dims_to_operand_dims = [0, 1],
++      index_vector_dim = 1
++    >,
++    indices_are_sorted = true,
++    unique_indices = true
++  } : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) ->
++      tensor<200x100x300xf64>
++  func.return %0 : tensor<200x100x300xf64>
++}
++
++// CHECK-LABEL: "op_select_and_scatter"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
++  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
++  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
++  // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
++  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG31:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG41:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL11:.*]] = "vhlo.compare_v1"(%[[ARG31]], %[[ARG41]]) <{compare_type = #vhlo<comparison_type_v1 TOTALORDER>, comparison_direction = #vhlo<comparison_direction_v1 GE>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL11]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
++  // CHECK-NEXT: }, {
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG32:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG42:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL12:.*]] = "vhlo.add_v1"(%[[ARG32]], %[[ARG42]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL12]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>, !vhlo.tensor_v1<12x13x13x66x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>
++  %0 = "stablehlo.select_and_scatter"(%arg0, %arg1, %arg2) ({
++    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
++      %1 = "stablehlo.compare"(%arg3, %arg4) {compare_type = #stablehlo<comparison_type TOTALORDER>, comparison_direction = #stablehlo<comparison_direction GE>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
++      "stablehlo.return"(%1) : (tensor<i1>) -> ()
++  }, {
++    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
++      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%1) : (tensor<f32>) -> ()
++  }) {
++    window_dimensions = array<i64: 1, 2, 2, 1>,
++    window_strides = array<i64: 1, 2, 2, 1>,
++    padding = dense<1> : tensor<4x2xi64>
++  } : (tensor<10x24x24x64xf32>, tensor<12x13x13x66xf32>, tensor<f32>) -> tensor<10x24x24x64xf32>
++  func.return %0 : tensor<10x24x24x64xf32>
++}
++
++// CHECK-LABEL: "op_select_and_scatter_with_promotable_types"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_select_and_scatter_with_promotable_types(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf64> {
++  // CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
++  // CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
++  // CHECK:     %[[VAL:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
++  // CHECK:     "vhlo.return_v1"(%[[VAL]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
++  // CHECK: }) : (!vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>, !vhlo.tensor_v1<12x13x13x66x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x24x24x64x!vhlo.f64_v1>
++  %0 = "stablehlo.select_and_scatter"(%arg0, %arg1, %arg2) ({
++    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
++      %1 = "stablehlo.compare"(%arg3, %arg4) {compare_type = #stablehlo<comparison_type TOTALORDER>, comparison_direction = #stablehlo<comparison_direction GE>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
++      "stablehlo.return"(%1) : (tensor<i1>) -> ()
++  }, {
++    ^bb0(%arg3: tensor<f64>, %arg4: tensor<f64>):
++      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f64>, tensor<f64>) -> tensor<f64>
++      "stablehlo.return"(%1) : (tensor<f64>) -> ()
++  }) {
++    window_dimensions = array<i64: 1, 2, 2, 1>,
++    window_strides = array<i64: 1, 2, 2, 1>,
++    padding = dense<1> : tensor<4x2xi64>
++  } : (tensor<10x24x24x64xf32>, tensor<12x13x13x66xf32>, tensor<f32>) -> tensor<10x24x24x64xf64>
++  func.return %0 : tensor<10x24x24x64xf64>
++}
++
++// CHECK-LABEL: "op_select"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
++func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_send_no_source_target_pairs"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_send_no_source_target_pairs(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
++  //      CHECK: "vhlo.send_v2"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
++  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>,
++  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
++  // CHECK-SAME{LITERAL}: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
++  %0 = "stablehlo.send"(%arg0, %arg1) {
++    channel_handle = #stablehlo.channel_handle<handle = 0, type = 2>,
++    is_host_transfer = true
++  } : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
++  func.return %0 : !stablehlo.token
++}
++
++// CHECK-LABEL: "op_set_dimension_size"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
++  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
++    dimension = 0 : i64
++  } : (tensor<?xf32>, tensor<i32>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_shift_left"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
++  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
++  %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
++  func.return %0 : tensor<i32>
++}
++
++// CHECK-LABEL: "op_shift_right_arithmetic"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
++  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
++  %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
++  func.return %0 : tensor<i32>
++}
++
++// CHECK-LABEL: "op_shift_right_logical"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
++  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
++  %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
++  func.return %0 : tensor<i32>
++}
++
++// CHECK-LABEL: "op_sign"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_sine"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.sine_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_slice"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
++  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
++  // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
++  // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
++  %0 = "stablehlo.slice"(%arg0) {
++    start_indices = array<i64: 0>,
++    limit_indices = array<i64: 4>,
++    strides = array<i64: 1>
++  } : (tensor<16xf32>) -> tensor<4xf32>
++  func.return %0 : tensor<4xf32>
++}
++
++// CHECK-LABEL: "op_sort"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
++  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
++  // CHECK-SAME: }> ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.compare_v1"(%[[ARG1]], %[[ARG2]]) <{compare_type = #vhlo<comparison_type_v1 FLOAT>, comparison_direction = #vhlo<comparison_direction_v1 GT>}>
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
++  %0 = "stablehlo.sort"(%arg0) ({
++    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
++      %1 = "stablehlo.compare"(%arg1, %arg2) {compare_type = #stablehlo<comparison_type FLOAT>, comparison_direction = #stablehlo<comparison_direction GT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
++      "stablehlo.return"(%1) : (tensor<i1>) -> ()
++  }) {
++    dimension = 0 : i64,
++    is_stable = true
++  } : (tensor<16xf32>) -> tensor<16xf32>
++  func.return %0 : tensor<16xf32>
++}
++
++// CHECK-LABEL: "op_sqrt"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.sqrt_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_subtract"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_tan"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_tan(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.tan_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.tan"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_tanh"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.tanh_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_torch_index_select"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
++  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
++  %0 = "stablehlo.torch_index_select"(%arg0, %arg1) {
++    dim = 0 : i64,
++    batch_dims = 0 : i64
++  } : (tensor<5x1x5xf32>, tensor<2xi32>) -> tensor<2x1x5xf32>
++  func.return %0 : tensor<2x1x5xf32>
++}
++
++// CHECK-LABEL: "op_transpose"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
++  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
++  %0 = "stablehlo.transpose"(%arg0) {
++    permutation = array<i64: 1, 0>
++  } : (tensor<16x8xf32>) -> tensor<8x16xf32>
++  func.return %0 : tensor<8x16xf32>
++}
++
++// CHECK-LABEL: "op_triangular_solve"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
++  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
++  // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
++  // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
++  // CHECK-SAME:   unit_diagonal = #vhlo.bool_v1<true>
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
++  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
++    left_side = true,
++    lower = true,
++    unit_diagonal = true,
++    transpose_a = #stablehlo<transpose NO_TRANSPOSE>
++  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
++  func.return %0 : tensor<16x16xf32>
++}
++
++// CHECK-LABEL: "op_tuple"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
++  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
++  %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
++  func.return %0 : tuple<tensor<f32>>
++}
++
++// CHECK-LABEL: "op_unary_einsum"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
++  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
++  // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
++  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
++  %0 = "stablehlo.unary_einsum"(%arg0) {
++    einsum_config = "ab->a"
++  } : (tensor<8x16xf32>) -> tensor<8xf32>
++  func.return %0 : tensor<8xf32>
++}
++
++// CHECK-LABEL: "op_uniform_dequantize"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
++  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "op_uniform_quantize"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
++  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
++  %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
++  func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
++}
++
++// CHECK-LABEL: "op_while"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
++  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
++  // CHECK-NEXT:   }, {
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>)
++  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  %0 = "stablehlo.while"(%arg0) ({
++    ^bb0(%arg1: tensor<i1>):
++      "stablehlo.return"(%arg1) : (tensor<i1>) -> ()
++    }, {
++    ^bb0(%arg1: tensor<i1>):
++      "stablehlo.return"(%arg1) : (tensor<i1>) -> ()
++  }) : (tensor<i1>) -> tensor<i1>
++  func.return %0: tensor<i1>
++}
++
++// CHECK-LABEL: "op_xor"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
++  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// ============ TYPES ============
++
++// CHECK-LABEL: "type_i1"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
++  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
++  %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
++  func.return %0 : tensor<i1>
++}
++
++// CHECK-LABEL: "type_i2"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_i2(%arg0: tensor<i2>, %arg1: tensor<i2>) -> tensor<i2> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i2_v1>, !vhlo.tensor_v1<!vhlo.i2_v1>) -> !vhlo.tensor_v1<!vhlo.i2_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i2>, tensor<i2>) -> tensor<i2>
++  func.return %0 : tensor<i2>
++}
++
++// CHECK-LABEL: "type_i4"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
++  func.return %0 : tensor<i4>
++}
++
++// CHECK-LABEL: "type_i8"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
++  func.return %0 : tensor<i8>
++}
++
++// CHECK-LABEL: "type_i16"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
++  func.return %0 : tensor<i16>
++}
++
++// CHECK-LABEL: "type_i32"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
++  func.return %0 : tensor<i32>
++}
++
++// CHECK-LABEL: "type_i64"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
++  func.return %0 : tensor<i64>
++}
++
++// CHECK-LABEL: "type_ui2"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_ui2(%arg0: tensor<ui2>, %arg1: tensor<ui2>) -> tensor<ui2> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui2_v1>, !vhlo.tensor_v1<!vhlo.ui2_v1>) -> !vhlo.tensor_v1<!vhlo.ui2_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui2>, tensor<ui2>) -> tensor<ui2>
++  func.return %0 : tensor<ui2>
++}
++
++// CHECK-LABEL: "type_ui4"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
++  func.return %0 : tensor<ui4>
++}
++
++// CHECK-LABEL: "type_ui8"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
++  func.return %0 : tensor<ui8>
++}
++
++// CHECK-LABEL: "type_ui16"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
++  func.return %0 : tensor<ui16>
++}
++
++// CHECK-LABEL: "type_ui32"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
++  func.return %0 : tensor<ui32>
++}
++
++// CHECK-LABEL: "type_ui64"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
++  func.return %0 : tensor<ui64>
++}
++
++// CHECK-LABEL: "type_f4E2M1FN"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f4E2M1FN(%arg0: tensor<f4E2M1FN>, %arg1: tensor<f4E2M1FN>) -> tensor<f4E2M1FN> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f4E2M1FN_v1>, !vhlo.tensor_v1<!vhlo.f4E2M1FN_v1>) -> !vhlo.tensor_v1<!vhlo.f4E2M1FN_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f4E2M1FN>, tensor<f4E2M1FN>) -> tensor<f4E2M1FN>
++  func.return %0 : tensor<f4E2M1FN>
++}
++
++// CHECK-LABEL: "type_f6E2M3FN"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f6E2M3FN(%arg0: tensor<f6E2M3FN>, %arg1: tensor<f6E2M3FN>) -> tensor<f6E2M3FN> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f6E2M3FN_v1>, !vhlo.tensor_v1<!vhlo.f6E2M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f6E2M3FN_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f6E2M3FN>, tensor<f6E2M3FN>) -> tensor<f6E2M3FN>
++  func.return %0 : tensor<f6E2M3FN>
++}
++
++// CHECK-LABEL: "type_f6E3M2FN"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f6E3M2FN(%arg0: tensor<f6E3M2FN>, %arg1: tensor<f6E3M2FN>) -> tensor<f6E3M2FN> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f6E3M2FN_v1>, !vhlo.tensor_v1<!vhlo.f6E3M2FN_v1>) -> !vhlo.tensor_v1<!vhlo.f6E3M2FN_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f6E3M2FN>, tensor<f6E3M2FN>) -> tensor<f6E3M2FN>
++  func.return %0 : tensor<f6E3M2FN>
++}
++
++// CHECK-LABEL: "type_f8E3M4"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f8E3M4(%arg0: tensor<f8E3M4>, %arg1: tensor<f8E3M4>) -> tensor<f8E3M4> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E3M4_v1>, !vhlo.tensor_v1<!vhlo.f8E3M4_v1>) -> !vhlo.tensor_v1<!vhlo.f8E3M4_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E3M4>, tensor<f8E3M4>) -> tensor<f8E3M4>
++  func.return %0 : tensor<f8E3M4>
++}
++
++// CHECK-LABEL: "type_f8E4M3"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f8E4M3(%arg0: tensor<f8E4M3>, %arg1: tensor<f8E4M3>) -> tensor<f8E4M3> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3>, tensor<f8E4M3>) -> tensor<f8E4M3>
++  func.return %0 : tensor<f8E4M3>
++}
++
++// CHECK-LABEL: "type_f8E4M3FN"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
++  func.return %0 : tensor<f8E4M3FN>
++}
++
++// CHECK-LABEL: "type_f8E5M2"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
++  func.return %0 : tensor<f8E5M2>
++}
++
++// CHECK-LABEL: "type_f8E4M3FNUZ"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
++  func.return %0 : tensor<f8E4M3FNUZ>
++}
++
++// CHECK-LABEL: "type_f8E4M3B11FNUZ"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
++  func.return %0 : tensor<f8E4M3B11FNUZ>
++}
++
++// CHECK-LABEL: "type_f8E5M2FNUZ"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
++  func.return %0 : tensor<f8E5M2FNUZ>
++}
++
++// CHECK-LABEL: "type_f8E8M0FNU"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f8E8M0FNU(%arg0: tensor<f8E8M0FNU>, %arg1: tensor<f8E8M0FNU>) -> tensor<f8E8M0FNU> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E8M0FNU_v1>, !vhlo.tensor_v1<!vhlo.f8E8M0FNU_v1>) -> !vhlo.tensor_v1<!vhlo.f8E8M0FNU_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E8M0FNU>, tensor<f8E8M0FNU>) -> tensor<f8E8M0FNU>
++  func.return %0 : tensor<f8E8M0FNU>
++}
++
++// CHECK-LABEL: "type_bf16"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
++  func.return %0 : tensor<bf16>
++}
++
++// CHECK-LABEL: "type_f16"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
++  func.return %0 : tensor<f16>
++}
++
++// CHECK-LABEL: "type_f32"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
++// CHECK-LABEL: "type_f64"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
++  func.return %0 : tensor<f64>
++}
++
++// CHECK-LABEL: "type_complex_f32"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
++  func.return %0 : tensor<complex<f32>>
++}
++
++// CHECK-LABEL: "type_complex_f64"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
++  func.return %0 : tensor<complex<f64>>
++}
++
++// CHECK-LABEL: "type_tf32"
++// CHECK: #vhlo.type_v1<!vhlo.tf31_v1>
++func.func @type_tf32() attributes {stablehlo.attr = tf32 } {
++  return
++}
++
++// CHECK-LABEL: "type_none"
++// CHECK: #vhlo.type_v1<!vhlo.none_v1>
++func.func @type_none() attributes {stablehlo.attr = none } {
++  return
++}
++
++// CHECK-LABEL: "type_dynamism_ranked"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
++  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
++  %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
++  func.return %0 : tensor<?xf32>
++}
++
++// CHECK-LABEL: "type_per_tensor_quantization"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
++func.func @type_per_tensor_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
++  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
++  func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
++}
++
++// CHECK-LABEL: "type_per_axis_quantization"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @type_per_axis_quantization(%arg0: tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>) -> tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>> {
++  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG0]]) : (!vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, [3.400000e+01, 3.400000e+01], [16, 16], -128:127, 1>>, !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, [3.400000e+01, 3.400000e+01], [16, 16], -128:127, 1>>) -> !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, [3.400000e+01, 3.400000e+01], [16, 16], -128:127, 1>>
++  %0 = stablehlo.add %arg0, %arg0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
++  func.return %0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
++}
++
++//       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
++// CHECK-LABEL: "type_token_callee"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
++  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
++  return %arg0 : !stablehlo.token
++}
++
++//       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
++// CHECK-LABEL: "type_token_caller"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
++  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
++  // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
++  %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
++  return %0 : !stablehlo.token
++}
++
++// CHECK-LABEL: "type_tuple"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "foo"
++  // CHECK: (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>) -> !vhlo.tuple_v1<!vhlo.token_v1>
++  } : (tuple<tensor<f32>>) -> tuple<!stablehlo.token>
++  return %0 : tuple<!stablehlo.token>
++}
++
++// CHECK-LABEL: type_buffer_function_input_output
++// CHECK-NEXT: (%[[ARG0:.*]]: !vhlo.buffer_v1<2x!vhlo.f32_v1>)
++func.func @type_buffer_function_input_output(%arg0: memref<2xf32>) -> memref<2xf32> {
++  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.buffer_v1<2x!vhlo.f32_v1>) -> ()
++  func.return %arg0 : memref<2xf32>
++}
++
++// CHECK-LABEL: type_buffer_special_custom_calls
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @type_buffer_special_custom_calls(%arg0: tensor<2xf32>) -> tensor<2xf32> {
++  //               CHECK: %[[CALL0:.*]] = "vhlo.custom_call_v2"(%[[ARG0]])
++  //          CHECK-SAME: call_target_name = #vhlo.string_v1<"Pin">
++  //          CHECK-SAME: : (!vhlo.tensor_v1<2x!vhlo.f32_v1>) -> !vhlo.buffer_v1<2x!vhlo.f32_v1>
++  %0 = "stablehlo.custom_call"(%arg0) {
++    call_target_name = "Pin",
++    api_version = 4 : i32
++  } : (tensor<2xf32>) -> memref<2xf32>
++  //               CHECK: %{{.*}} = "vhlo.custom_call_v2"(%[[CALL0]])
++  //          CHECK-SAME: call_target_name = #vhlo.string_v1<"Unpin">
++  //          CHECK-SAME: : (!vhlo.buffer_v1<2x!vhlo.f32_v1>) -> !vhlo.tensor_v1<2x!vhlo.f32_v1>
++  %1 = "stablehlo.custom_call"(%0) {
++    call_target_name = "Unpin",
++    api_version = 4 : i32
++  } : (memref<2xf32>) -> tensor<2xf32>
++  func.return %1 : tensor<2xf32>
++}
++
++// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>>>
++// CHECK-LABEL: "type_future"
++func.func private @type_future() -> !stablehlo.future<tensor<f32>>
++
++// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<2x3x!vhlo.f32_v1>>>>
++// CHECK-LABEL: "type_future_ranked"
++func.func private @type_future_ranked() -> !stablehlo.future<tensor<2x3xf32>>
++
++// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<?x!vhlo.f32_v1>>>>
++// CHECK-LABEL: "type_future_dynamic"
++func.func private @type_future_dynamic() -> !stablehlo.future<tensor<?xf32>>
++
++// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>>>
++// CHECK-LABEL: "type_future_multiple"
++func.func private @type_future_multiple() -> !stablehlo.future<tensor<f32>, tensor<i32>>
++
++// CHECK-LABEL: "op_async_start_all_reduce"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_async_start_all_reduce(%arg0: tensor<4x4xf32>) -> !stablehlo.future<tensor<4x4xf32>> {
++  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
++  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.all_reduce_v2"(%[[BARG0]]) <{
++  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>,
++  // CHECK-SAME:     use_global_device_ids = #vhlo.bool_v1<false>
++  // CHECK-SAME:   }> ({
++  // CHECK-NEXT:     ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:       %[[VAL2:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  // CHECK-NEXT:       "vhlo.return_v1"(%[[VAL2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT:   }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<4x4x!vhlo.f32_v1>>
++  %0 = "stablehlo.async_start"(%arg0) ({
++    ^bb0(%barg0: tensor<4x4xf32>):
++    %1 = "stablehlo.all_reduce"(%barg0) ({
++    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
++      %2 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%2) : (tensor<f32>) -> ()
++    }) {replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
++    "stablehlo.return"(%1) : (tensor<4x4xf32>) -> ()
++  }) : (tensor<4x4xf32>) -> !stablehlo.future<tensor<4x4xf32>>
++  func.return %0: !stablehlo.future<tensor<4x4xf32>>
++}
++
++// CHECK-LABEL: "op_async_start_all_gather"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_async_start_all_gather(%arg0: tensor<8x2xf32>) -> !stablehlo.future<tensor<8x8xf32>> {
++  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
++  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.all_gather_v2"(%[[BARG0]]) <{
++  // CHECK-SAME:     all_gather_dim = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 2, 4, 6], [1, 3, 5, 7]]> : tensor<2x4xi64>>,
++  // CHECK-SAME:     use_global_device_ids = #vhlo.bool_v1<false>
++  // CHECK-SAME:   }> : (!vhlo.tensor_v1<8x2x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<8x8x!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<8x2x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<8x8x!vhlo.f32_v1>>
++  %0 = "stablehlo.async_start"(%arg0) ({
++    ^bb0(%barg0: tensor<8x2xf32>):
++    %1 = "stablehlo.all_gather"(%barg0) {
++      all_gather_dim = 1 : i64,
++      replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7]]> : tensor<2x4xi64>
++    } : (tensor<8x2xf32>) -> tensor<8x8xf32>
++    "stablehlo.return"(%1) : (tensor<8x8xf32>) -> ()
++  }) : (tensor<8x2xf32>) -> !stablehlo.future<tensor<8x8xf32>>
++  func.return %0: !stablehlo.future<tensor<8x8xf32>>
++}
++
++// CHECK-LABEL: "op_async_start_all_to_all"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_async_start_all_to_all(%arg0: tensor<4x16xf32>) -> !stablehlo.future<tensor<16x4xf32>> {
++  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
++  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.all_to_all_v2"(%[[BARG0]]) <{
++  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:     concat_dimension = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
++  // CHECK-SAME:     split_count = #vhlo.integer_v1<4 : i64>,
++  // CHECK-SAME:     split_dimension = #vhlo.integer_v1<1 : i64>
++  // CHECK-SAME:   }> : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x4x!vhlo.f32_v1>
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<16x4x!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<16x4x!vhlo.f32_v1>>
++  %0 = "stablehlo.async_start"(%arg0) ({
++    ^bb0(%barg0: tensor<4x16xf32>):
++    %1 = "stablehlo.all_to_all"(%barg0) {
++      split_dimension = 1 : i64,
++      concat_dimension = 0 : i64,
++      split_count = 4 : i64,
++      replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>
++    } : (tensor<4x16xf32>) -> tensor<16x4xf32>
++    "stablehlo.return"(%1) : (tensor<16x4xf32>) -> ()
++  }) : (tensor<4x16xf32>) -> !stablehlo.future<tensor<16x4xf32>>
++  func.return %0: !stablehlo.future<tensor<16x4xf32>>
++}
++
++// CHECK-LABEL: "op_async_start_collective_permute"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_async_start_collective_permute(%arg0: tensor<128x32xf32>) -> !stablehlo.future<tensor<128x32xf32>> {
++  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
++  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.collective_permute_v1"(%[[BARG0]]) <{
++  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:     source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
++  // CHECK-SAME:   }> : (!vhlo.tensor_v1<128x32x!vhlo.f32_v1>) -> !vhlo.tensor_v1<128x32x!vhlo.f32_v1>
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<128x32x!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<128x32x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<128x32x!vhlo.f32_v1>>
++  %0 = "stablehlo.async_start"(%arg0) ({
++    ^bb0(%barg0: tensor<128x32xf32>):
++    %1 = "stablehlo.collective_permute"(%barg0) {
++      source_target_pairs = dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>
++    } : (tensor<128x32xf32>) -> tensor<128x32xf32>
++    "stablehlo.return"(%1) : (tensor<128x32xf32>) -> ()
++  }) : (tensor<128x32xf32>) -> !stablehlo.future<tensor<128x32xf32>>
++  func.return %0: !stablehlo.future<tensor<128x32xf32>>
++}
++
++// CHECK-LABEL: "op_async_start_collective_broadcast"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_async_start_collective_broadcast(%arg0: tensor<16x8xf32>) -> !stablehlo.future<tensor<16x8xf32>> {
++  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
++  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.collective_broadcast_v2"(%[[BARG0]]) <{
++  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME:     has_dynamic_root = #vhlo.bool_v1<false>,
++  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
++  // CHECK-SAME:   }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<16x8x!vhlo.f32_v1>>
++  %0 = "stablehlo.async_start"(%arg0) ({
++    ^bb0(%barg0: tensor<16x8xf32>):
++    %1 = "stablehlo.collective_broadcast"(%barg0) {
++      replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
++    } : (tensor<16x8xf32>) -> tensor<16x8xf32>
++    "stablehlo.return"(%1) : (tensor<16x8xf32>) -> ()
++  }) : (tensor<16x8xf32>) -> !stablehlo.future<tensor<16x8xf32>>
++  func.return %0: !stablehlo.future<tensor<16x8xf32>>
++}
++
++// CHECK-LABEL: "op_async_start_reduce_scatter"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_async_start_reduce_scatter(%arg0: tensor<4x16xf32>) -> !stablehlo.future<tensor<4x4xf32>> {
++  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
++  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
++  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.reduce_scatter_v1"(%[[BARG0]]) <{
++  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
++  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
++  // CHECK-SAME:     scatter_dimension = #vhlo.integer_v1<1 : i64>,
++  // CHECK-SAME:     use_global_device_ids = #vhlo.bool_v1<false>
++  // CHECK-SAME:   }> ({
++  // CHECK-NEXT:     ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
++  // CHECK-NEXT:       %[[VAL2:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  // CHECK-NEXT:       "vhlo.return_v1"(%[[VAL2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT:   }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
++  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> ()
++  // CHECK-NEXT: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<4x4x!vhlo.f32_v1>>
++  %0 = "stablehlo.async_start"(%arg0) ({
++    ^bb0(%barg0: tensor<4x16xf32>):
++    %1 = "stablehlo.reduce_scatter"(%barg0) ({
++      ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
++      %2 = stablehlo.add %arg2, %arg3 : tensor<f32>
++      "stablehlo.return"(%2) : (tensor<f32>) -> ()
++    }) {replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
++        scatter_dimension = 1 : i64} : (tensor<4x16xf32>) -> tensor<4x4xf32>
++    "stablehlo.return"(%1) : (tensor<4x4xf32>) -> ()
++  }) : (tensor<4x16xf32>) -> !stablehlo.future<tensor<4x4xf32>>
++  func.return %0: !stablehlo.future<tensor<4x4xf32>>
++}
++
++// CHECK-LABEL: "op_async_done"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_async_done(%arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
++  // CHECK: %[[VAL1:.*]] = "vhlo.async_start_v1"(%[[ARG0]])
++  // CHECK: %[[VAL2:.*]] = "vhlo.async_done_v1"(%[[VAL1]]) : (!vhlo.future_v1<!vhlo.tensor_v1<4x4x!vhlo.f32_v1>>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
++  %0 = "stablehlo.async_start"(%arg0) ({
++    ^bb0(%barg0: tensor<4x4xf32>):
++    %1 = "stablehlo.all_reduce"(%barg0) ({
++    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
++      %2 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
++      "stablehlo.return"(%2) : (tensor<f32>) -> ()
++    }) {replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
++    "stablehlo.return"(%1) : (tensor<4x4xf32>) -> ()
++  }) : (tensor<4x4xf32>) -> !stablehlo.future<tensor<4x4xf32>>
++  %1 = "stablehlo.async_done"(%0) : (!stablehlo.future<tensor<4x4xf32>>) -> tensor<4x4xf32>
++  func.return %1: tensor<4x4xf32>
++}
++
++// ============ DEPENDENCIES  ============
++
++func.func @composite_target(%arg0: tensor<f32>) -> tensor<f32> {
++  return %arg0: tensor<f32>
++}
+diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_23_0.mlir.bc b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_23_0.mlir.bc
+Binary files /dev/null and stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_23_0.mlir.bc differ
 diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir
 --- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir
 +++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir
@@ -12361,6 +16239,36 @@ diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir b/st
  // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
  // CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
  func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
+@@ -1826,6 +1911,14 @@
+   func.return %0 : tensor<f32>
+ }
+ 
++// CHECK-LABEL: "op_exp2"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_exp2(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.exp2_v1"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.exp2"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
+ // CHECK-LABEL: "op_fft"
+ // CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
+ func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
+@@ -2006,6 +2099,14 @@
+   func.return %0 : tensor<f32>
+ }
+ 
++// CHECK-LABEL: "op_log2"
++// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
++func.func @op_log2(%arg0: tensor<f32>) -> tensor<f32> {
++  // CHECK: "vhlo.log2_v1"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
++  %0 = "stablehlo.log2"(%arg0) : (tensor<f32>) -> tensor<f32>
++  func.return %0 : tensor<f32>
++}
++
+ // CHECK-LABEL: "op_log_plus_one"
+ // CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
+ func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
 diff --ruN a/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_20_0.mlir b/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_20_0.mlir
 --- stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_20_0.mlir
 +++ stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_20_0.mlir
@@ -12441,6 +16349,31 @@ diff --ruN a/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_
 +  }
 +}
 +
+diff --ruN a/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_22_0.mlir b/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_22_0.mlir
+--- stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_22_0.mlir
++++ stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_22_0.mlir
+@@ -0,0 +1,21 @@
++// RUN: stablehlo-opt --stablehlo-legalize-to-vhlo --vhlo-to-version='target=1.22.0' --verify-diagnostics --split-input-file %s
++
++// expected-error @+1 {{failed to convert VHLO to v1.22.0}}
++module {
++  func.func @exp2(%arg0: tensor<f32>) -> tensor<f32> {
++    // expected-error @+1 {{failed to legalize operation 'vhlo.exp2_v1' that was explicitly marked illegal}}
++    %0 = "stablehlo.exp2"(%arg0) : (tensor<f32>) -> tensor<f32>
++    func.return %0 : tensor<f32>
++  }
++}
++
++// -----
++
++// expected-error @+1 {{failed to convert VHLO to v1.22.0}}
++module {
++  func.func @log2(%arg0: tensor<f32>) -> tensor<f32> {
++    // expected-error @+1 {{failed to legalize operation 'vhlo.log2_v1' that was explicitly marked illegal}}
++    %0 = "stablehlo.log2"(%arg0) : (tensor<f32>) -> tensor<f32>
++    func.return %0 : tensor<f32>
++  }
++}
 diff --ruN a/stablehlo/stablehlo/tools/StablehloLspServerMain.cpp b/stablehlo/stablehlo/tools/StablehloLspServerMain.cpp
 --- stablehlo/stablehlo/tools/StablehloLspServerMain.cpp
 +++ stablehlo/stablehlo/tools/StablehloLspServerMain.cpp
@@ -12759,6 +16692,25 @@ diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stable
            newElement);
  
        int64_t outRank = outputType.getRank();
+diff --ruN a/stablehlo/stablehlo/transforms/MapStablehloToVhlo.h b/stablehlo/stablehlo/transforms/MapStablehloToVhlo.h
+--- stablehlo/stablehlo/transforms/MapStablehloToVhlo.h
++++ stablehlo/stablehlo/transforms/MapStablehloToVhlo.h
+@@ -97,6 +97,7 @@
+ MAP_STABLEHLO_TO_VHLO(DynamicUpdateSliceOp, V1)
+ MAP_STABLEHLO_TO_VHLO(EinsumOp, V1)
+ MAP_STABLEHLO_TO_VHLO(Expm1Op, V2)
++MAP_STABLEHLO_TO_VHLO(Exp2Op, V1)
+ MAP_STABLEHLO_TO_VHLO(ExpOp, V2)
+ MAP_STABLEHLO_TO_VHLO(FftOp, V1)
+ MAP_STABLEHLO_TO_VHLO(FloorOp, V1)
+@@ -109,6 +110,7 @@
+ MAP_STABLEHLO_TO_VHLO(IotaOp, V1)
+ MAP_STABLEHLO_TO_VHLO(IsFiniteOp, V1)
+ MAP_STABLEHLO_TO_VHLO(Log1pOp, V2)
++MAP_STABLEHLO_TO_VHLO(Log2Op, V1)
+ MAP_STABLEHLO_TO_VHLO(LogisticOp, V2)
+ MAP_STABLEHLO_TO_VHLO(LogOp, V2)
+ MAP_STABLEHLO_TO_VHLO(MapOp, V1)
 diff --ruN a/stablehlo/stablehlo/transforms/PassPipelines.cpp b/stablehlo/stablehlo/transforms/PassPipelines.cpp
 --- stablehlo/stablehlo/transforms/PassPipelines.cpp
 +++ stablehlo/stablehlo/transforms/PassPipelines.cpp
@@ -13082,7 +17034,18 @@ diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeQuantizedOpToQDQ.cp
 diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeToVhlo.cpp b/stablehlo/stablehlo/transforms/StablehloLegalizeToVhlo.cpp
 --- stablehlo/stablehlo/transforms/StablehloLegalizeToVhlo.cpp
 +++ stablehlo/stablehlo/transforms/StablehloLegalizeToVhlo.cpp
-@@ -1131,15 +1131,17 @@
+@@ -969,8 +969,10 @@
+                 std::is_same<StablehloOpTy, stablehlo::CosineOp>::value ||
+                 std::is_same<StablehloOpTy, stablehlo::ExpOp>::value ||
+                 std::is_same<StablehloOpTy, stablehlo::Expm1Op>::value ||
++                std::is_same<StablehloOpTy, stablehlo::Exp2Op>::value ||
+                 std::is_same<StablehloOpTy, stablehlo::LogOp>::value ||
+                 std::is_same<StablehloOpTy, stablehlo::Log1pOp>::value ||
++                std::is_same<StablehloOpTy, stablehlo::Log2Op>::value ||
+                 std::is_same<StablehloOpTy, stablehlo::LogisticOp>::value ||
+                 std::is_same<StablehloOpTy, stablehlo::RsqrtOp>::value ||
+                 std::is_same<StablehloOpTy, stablehlo::SineOp>::value ||
+@@ -1131,15 +1133,17 @@
      // This can almost be done in a generic fashion, except for
      // stablehlo.case that uses a variadic number of regions which means an
      // additional argument for the generic builder.
@@ -13121,7 +17084,18 @@ diff --ruN a/stablehlo/stablehlo/transforms/StablehloRefineArguments.cpp b/stabl
 diff --ruN a/stablehlo/stablehlo/transforms/VhloLegalizeToStablehlo.cpp b/stablehlo/stablehlo/transforms/VhloLegalizeToStablehlo.cpp
 --- stablehlo/stablehlo/transforms/VhloLegalizeToStablehlo.cpp
 +++ stablehlo/stablehlo/transforms/VhloLegalizeToStablehlo.cpp
-@@ -1121,15 +1121,17 @@
+@@ -1045,8 +1045,10 @@
+                 std::is_same<VhloOpTy, vhlo::CosineOpV2>::value ||
+                 std::is_same<VhloOpTy, vhlo::ExpOpV2>::value ||
+                 std::is_same<VhloOpTy, vhlo::Expm1OpV2>::value ||
++                std::is_same<VhloOpTy, vhlo::Exp2OpV1>::value ||
+                 std::is_same<VhloOpTy, vhlo::LogOpV2>::value ||
+                 std::is_same<VhloOpTy, vhlo::Log1pOpV2>::value ||
++                std::is_same<VhloOpTy, vhlo::Log2OpV1>::value ||
+                 std::is_same<VhloOpTy, vhlo::LogisticOpV2>::value ||
+                 std::is_same<VhloOpTy, vhlo::RsqrtOpV2>::value ||
+                 std::is_same<VhloOpTy, vhlo::SineOpV2>::value ||
+@@ -1121,15 +1123,17 @@
      // Convert the VHLO operation to a StableHLO equivalent. This can almost be
      // done in a generic fashion, except for ops with a variadic number of
      // regions which means an additional argument for the generic builder.
@@ -13273,6 +17247,50 @@ diff --ruN a/stablehlo/stablehlo/transforms/optimization/StablehloAggressiveFold
        if (kind != ComparisonType::FLOAT && kind != ComparisonType::TOTALORDER)
          llvm::report_fatal_error("invalid float comparison");
  
+@@ -1661,6 +1691,15 @@
+   }
+ };
+ 
++struct FoldExp2OpPattern
++    : public FoldUnaryOpPattern<FoldExp2OpPattern, Exp2Op> {
++  using FoldUnaryOpPattern::FoldUnaryOpPattern;
++
++  static std::optional<double> EvaluateOp(double operand) {
++    return std::exp2(operand);
++  }
++};
++
+ struct FoldLogOpPattern : public FoldUnaryOpPattern<FoldLogOpPattern, LogOp> {
+   using FoldUnaryOpPattern::FoldUnaryOpPattern;
+ 
+@@ -1670,6 +1709,16 @@
+   }
+ };
+ 
++struct FoldLog2OpPattern
++    : public FoldUnaryOpPattern<FoldLog2OpPattern, Log2Op> {
++  using FoldUnaryOpPattern::FoldUnaryOpPattern;
++
++  static std::optional<double> EvaluateOp(double operand) {
++    if (operand <= 0.0) return std::nullopt;
++    return std::log2(operand);
++  }
++};
++
+ struct FoldLogisticOpPattern
+     : public FoldUnaryOpPattern<FoldLogisticOpPattern, LogisticOp> {
+   using FoldUnaryOpPattern::FoldUnaryOpPattern;
+@@ -2222,8 +2271,10 @@
+                 FoldCosineOpPattern,                  //
+                 FoldErfOpPattern,                     //
+                 FoldExpOpPattern,                     //
++                FoldExp2OpPattern,                    //
+                 FoldIotaOpPattern,                    //
+                 FoldLogOpPattern,                     //
++                FoldLog2OpPattern,                    //
+                 FoldLogisticOpPattern,                //
+                 FoldNegOpPattern,                     //
+                 FoldNotOpPattern,                     //
 diff --ruN a/stablehlo/stablehlo/transforms/optimization/StablehloAggressiveSimplification.cpp b/stablehlo/stablehlo/transforms/optimization/StablehloAggressiveSimplification.cpp
 --- stablehlo/stablehlo/transforms/optimization/StablehloAggressiveSimplification.cpp
 +++ stablehlo/stablehlo/transforms/optimization/StablehloAggressiveSimplification.cpp
diff --git a/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_stablehlo_to_hlo_op.h b/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_stablehlo_to_hlo_op.h
index 1361ab1a232..cf42694cc58 100644
--- a/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_stablehlo_to_hlo_op.h
+++ b/third_party/xla/xla/mlir_hlo/mhlo/transforms/map_stablehlo_to_hlo_op.h
@@ -92,6 +92,7 @@ MAP_STABLEHLO_TO_HLO(DynamicUpdateSliceOp)
 MAP_STABLEHLO_TO_HLO(EinsumOp)
 MAP_STABLEHLO_TO_HLO(Expm1Op)
 MAP_STABLEHLO_TO_HLO(ExpOp)
+MAP_STABLEHLO_TO_HLO(Exp2Op)
 MAP_STABLEHLO_TO_HLO(FftOp)
 MAP_STABLEHLO_TO_HLO(FloorOp)
 MAP_STABLEHLO_TO_HLO(GatherOp)
@@ -105,6 +106,7 @@ MAP_STABLEHLO_TO_HLO(IsFiniteOp)
 MAP_STABLEHLO_TO_HLO(Log1pOp)
 MAP_STABLEHLO_TO_HLO(LogisticOp)
 MAP_STABLEHLO_TO_HLO(LogOp)
+MAP_STABLEHLO_TO_HLO(Log2Op)
 MAP_STABLEHLO_TO_HLO(MapOp)
 MAP_STABLEHLO_TO_HLO(MaxOp)
 MAP_STABLEHLO_TO_HLO(MinOp)
diff --git a/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/hlo-legalize-to-stablehlo.mlir b/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/hlo-legalize-to-stablehlo.mlir
index 71a5466039b..0d6c9dbc9d9 100644
--- a/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/hlo-legalize-to-stablehlo.mlir
+++ b/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/hlo-legalize-to-stablehlo.mlir
@@ -2617,3 +2617,21 @@ func.func @op_scan(%arg0: tensor<10xf32>, %arg1: tensor<f32>) -> (tensor<10xf32>
   } : (tensor<10xf32>, tensor<f32>) -> (tensor<10xf32>, tensor<f32>)
   func.return %0#0, %0#1 : tensor<10xf32>, tensor<f32>
 }
+
+// -----
+
+// CHECK-LABEL: "op_exp2"
+func.func @op_exp2(%arg0: tensor<f32>) -> tensor<f32> {
+  // CHECK: "stablehlo.exp2"([[ARG0:%arg[0-9]+]]) : (tensor<f32>) -> tensor<f32>
+  %0 = "mhlo.exp2"(%arg0) : (tensor<f32>) -> tensor<f32>
+  func.return %0 : tensor<f32>
+}
+
+// -----
+
+// CHECK-LABEL: "op_log2"
+func.func @op_log2(%arg0: tensor<f32>) -> tensor<f32> {
+  // CHECK: "stablehlo.log2"([[ARG0:%arg[0-9]+]]) : (tensor<f32>) -> tensor<f32>
+  %0 = "mhlo.log2"(%arg0) : (tensor<f32>) -> tensor<f32>
+  func.return %0 : tensor<f32>
+}
diff --git a/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/stablehlo-legalize-to-hlo.mlir b/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/stablehlo-legalize-to-hlo.mlir
index 5a5ca25d714..ce31f517e30 100644
--- a/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/stablehlo-legalize-to-hlo.mlir
+++ b/third_party/xla/xla/mlir_hlo/tests/Dialect/mhlo/stablehlo-legalize-to-hlo.mlir
@@ -2802,6 +2802,24 @@ func.func @op_unary_einsum_deprecated(%arg0: tensor<8x16xf32>) -> tensor<8xf32>
 
 // -----
 
+// CHECK-LABEL: "op_exp2"
+func.func @op_exp2(%arg0: tensor<f32>) -> tensor<f32> {
+  // CHECK: "mhlo.exp2"([[ARG0:%arg[0-9]+]]) : (tensor<f32>) -> tensor<f32>
+  %0 = "stablehlo.exp2"(%arg0) : (tensor<f32>) -> tensor<f32>
+  func.return %0 : tensor<f32>
+}
+
+// -----
+
+// CHECK-LABEL: "op_log2"
+func.func @op_log2(%arg0: tensor<f32>) -> tensor<f32> {
+  // CHECK: "mhlo.log2"([[ARG0:%arg[0-9]+]]) : (tensor<f32>) -> tensor<f32>
+  %0 = "stablehlo.log2"(%arg0) : (tensor<f32>) -> tensor<f32>
+  func.return %0 : tensor<f32>
+}
+
+// -----
+
 // CHECK-LABEL: "custom_call_op_create_buffer"
 func.func @custom_call_op_create_buffer() -> memref<2xf32> {
   //      CHECK: "mhlo.custom_call"

From 48d24df72eacbc64ffbb54fe01cc9582fc200b83 Mon Sep 17 00:00:00 2001
From: Viktor Moros <vgm@google.com>
Date: Thu, 15 Oct 2026 14:41:20 -0700
Subject: [PATCH 0911/1000] [XLA Highway] Dynamic dispatch for transpose
 microkernels (not called yet)

PiperOrigin-RevId: 1000403827
---
 third_party/xla/xla/pjrt/BUILD                |   2 +
 third_party/xla/xla/pjrt/transpose_kernels.cc | 521 ++++++++++++++++++
 third_party/xla/xla/pjrt/transpose_kernels.h  | 389 +------------
 third_party/xla/xla/pjrt/transpose_test.cc    | 196 +++++--
 4 files changed, 685 insertions(+), 423 deletions(-)
 create mode 100644 third_party/xla/xla/pjrt/transpose_kernels.cc

diff --git a/third_party/xla/xla/pjrt/BUILD b/third_party/xla/xla/pjrt/BUILD
index 2c8cb2366ed..81e4f44e1e6 100644
--- a/third_party/xla/xla/pjrt/BUILD
+++ b/third_party/xla/xla/pjrt/BUILD
@@ -947,6 +947,7 @@ xla_cc_test(
 
 cc_library(
     name = "transpose_kernels",
+    srcs = ["transpose_kernels.cc"],
     hdrs = ["transpose_kernels.h"],
     deps = [
         "//xla:compiler_macros",
@@ -1022,6 +1023,7 @@ xla_cc_test(
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/types:span",
         "@eigen_archive//:eigen3",
+        "@highway//:hwy",
     ],
 )
 
diff --git a/third_party/xla/xla/pjrt/transpose_kernels.cc b/third_party/xla/xla/pjrt/transpose_kernels.cc
new file mode 100644
index 00000000000..224ef3b6c52
--- /dev/null
+++ b/third_party/xla/xla/pjrt/transpose_kernels.cc
@@ -0,0 +1,521 @@
+/* Copyright 2021 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/pjrt/transpose_kernels.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <type_traits>
+#include <utility>
+
+// clang-format off
+#undef HWY_TARGET_INCLUDE
+#define HWY_TARGET_INCLUDE "xla/pjrt/transpose_kernels.cc"  // NOLINT(whitespace/line_length)
+// clang-format on
+#include "hwy//foreach_target.h"  // IWYU pragma: keep
+#include "hwy//highway.h"
+
+HWY_BEFORE_NAMESPACE();
+namespace xla {
+namespace HWY_NAMESPACE {
+namespace hn = hwy::HWY_NAMESPACE;
+
+// Highway's model: Vectors are composed of 128-bit blocks, blocks hold lanes.
+// Interleave operations work within blocks. Cross-block movement requires
+// block-level shuffles.
+constexpr size_t kBlockBytes = 16;   // 128-bit blocks
+constexpr size_t kMaxLaneBytes = 8;  // Lanes are at most 64 bits
+
+enum class Extract { kLo, kHi };
+
+// Unpack: Interleave lanes within each 128-bit block
+template <size_t element_size, Extract extract, class V>
+HWY_INLINE V Unpack(V a, V b) {
+  static_assert(element_size <= kMaxLaneBytes);
+  using D = hn::DFromV<V>;
+  using ElementT = hwy::UnsignedFromSize<element_size>;
+  const hn::Repartition<ElementT, D> d_elem;
+  auto a_elem = hn::BitCast(d_elem, a);
+  auto b_elem = hn::BitCast(d_elem, b);
+  if constexpr (extract == Extract::kLo) {
+    return hn::BitCast(D(), hn::InterleaveLower(d_elem, a_elem, b_elem));
+  } else {
+    return hn::BitCast(D(), hn::InterleaveUpper(d_elem, a_elem, b_elem));
+  }
+}
+
+// Load/Store
+template <class D, size_t bytes>
+HWY_INLINE hn::Vec<D> LoadBytes(D d, const void* p) {
+  const hn::FixedTag<uint8_t, bytes> d_fixed;
+  return hn::ResizeBitCast(
+      d, hn::LoadU(d_fixed, reinterpret_cast<const uint8_t*>(p)));
+}
+
+template <class D, size_t bytes, size_t lane>
+HWY_INLINE void StoreLane(D d, void* p, hn::Vec<D> v) {
+  if constexpr (bytes <= kMaxLaneBytes && lane > 0) {
+    using ElementT = hwy::UnsignedFromSize<bytes>;
+    const hn::Repartition<ElementT, D> d_elem;
+    const ElementT scalar = hn::ExtractLane(hn::BitCast(d_elem, v), lane);
+    hwy::CopyBytes(&scalar, reinterpret_cast<ElementT*>(p), bytes);
+  } else {
+    const hn::Repartition<uint8_t, D> d_u8;
+    auto v_u8 = hn::BitCast(d_u8, v);
+    if constexpr (lane > 0) {
+      v_u8 = hn::SlideDownLanes(d_u8, v_u8, lane * bytes);
+    }
+    const hn::FixedTag<uint8_t, bytes> d_fixed;
+    hn::StoreU(hn::ResizeBitCast(d_fixed, v_u8), d_fixed,
+               reinterpret_cast<uint8_t*>(p));
+  }
+}
+
+template <class D, size_t bytes, size_t... lane>
+HWY_INLINE void StoreLanes(D d, char* b, int64_t ldb, hn::Vec<D> v, size_t i,
+                           std::index_sequence<lane...>) {
+  (StoreLane<D, bytes, lane>(d, b + ldb * (i + lane), v), ...);
+}
+
+// Extracts the I-th element from a parameter pack (0-indexed).
+template <size_t I, class T, class... Ts>
+HWY_INLINE decltype(auto) GetFromPack(T&& t, Ts&&... ts) {
+  static_assert(I < 1 + sizeof...(Ts), "Index out of bounds");
+  if constexpr (I == 0) {
+    return std::forward<T>(t);
+  } else {
+    return GetFromPack<I - 1>(std::forward<Ts>(ts)...);
+  }
+}
+
+// Pack iteration helper (pack -> stream): ForEachIndexed
+template <class F, size_t... I, class... Ts>
+HWY_INLINE void ForEachIndexedImpl(F&& f, std::index_sequence<I...>,
+                                   Ts&&... ts) {
+  (std::forward<F>(f)(std::integral_constant<size_t, I>{},
+                      std::forward<Ts>(ts)),
+   ...);
+}
+
+template <class F, class... Ts>
+HWY_INLINE void ForEachIndexed(F&& f, Ts&&... ts) {
+  ForEachIndexedImpl(std::forward<F>(f),
+                     std::make_index_sequence<sizeof...(Ts)>{},
+                     std::forward<Ts>(ts)...);
+}
+
+// CPS: UnpackStep (pack -> pack), lvalue-only inputs
+template <size_t out_i, size_t element_size, size_t step_size, class... Vs>
+HWY_INLINE auto UnpackOutFromPack(Vs&... vs) {
+  constexpr size_t group = 2 * step_size;
+  constexpr size_t base = (out_i / group) * group;
+  constexpr size_t p = out_i - base;
+  constexpr size_t j = p / 2;
+  constexpr bool is_lo = (p % 2) == 0;
+  constexpr size_t ia = base + j;
+  constexpr size_t ib = base + j + step_size;
+  auto& a = GetFromPack<ia>(vs...);
+  auto& b = GetFromPack<ib>(vs...);
+  if constexpr (is_lo) {
+    return Unpack<element_size * step_size, Extract::kLo>(a, b);
+  } else {
+    return Unpack<element_size * step_size, Extract::kHi>(a, b);
+  }
+}
+
+template <size_t element_size, size_t step_size, class Cont, size_t... I,
+          class... Vs>
+HWY_INLINE void UnpackStepCPSImpl(Cont&& cont, std::index_sequence<I...>,
+                                  Vs&... vs) {
+  std::forward<Cont>(cont)(
+      UnpackOutFromPack<I, element_size, step_size>(vs...)...);
+}
+
+template <size_t element_size, size_t step_size, class Cont, class... Vs>
+HWY_INLINE void UnpackStepCPS(Cont&& cont, Vs&... in) {
+  constexpr size_t N = sizeof...(Vs);
+  static_assert(N % (step_size * 2) == 0);
+  static_assert(element_size * step_size <= kMaxLaneBytes);
+  UnpackStepCPSImpl<element_size, step_size>(
+      std::forward<Cont>(cont), std::make_index_sequence<N>{}, in...);
+}
+
+// CPS: InterleaveWithinBlocks (pack -> pack).
+// Lambdas need explicit HWY_ATTR:
+// https://github.com/google/highway/blob/85cb4a5a266f7cc998ee4a270bdcd8f09926f63e/hwy/base.h#L206-L208
+template <size_t element_size, size_t step_size, size_t unpack_limit,
+          class Cont, class... Vs>
+HWY_INLINE void InterleaveWithinBlocksCPS(Cont&& cont, Vs... in) {
+  if constexpr (element_size * step_size < unpack_limit &&
+                element_size * step_size <= kMaxLaneBytes) {
+    UnpackStepCPS<element_size, step_size>(
+        [&](auto... out) HWY_ATTR {
+          // out... are by-value parameters; do not forward as rvalues.
+          InterleaveWithinBlocksCPS<element_size, step_size * 2, unpack_limit>(
+              std::forward<Cont>(cont), out...);
+        },
+        in...);  // lvalues inside this function
+  } else {
+    std::forward<Cont>(cont)(in...);
+  }
+}
+
+// Helper template that applies the pairwise interleaving reduction.
+template <size_t element_size, class Cont, size_t... I, class... Vs>
+HWY_INLINE void CombinePairsAccImpl(Cont&& cont, std::index_sequence<I...>,
+                                    Vs&&... vs) {
+  std::forward<Cont>(cont)(Unpack<element_size, Extract::kLo>(
+      GetFromPack<2 * I>(std::forward<Vs>(vs)...),
+      GetFromPack<2 * I + 1>(std::forward<Vs>(vs)...))...);
+}
+
+// Applies a pairwise interleaving reduction across an arbitrary sequence of
+// vectors. Adjacent vectors are combined using their lower halves, and the
+// resulting halved-length sequence is forwarded to the provided continuation.
+template <size_t element_size, class Cont, class... Vs>
+HWY_INLINE void CombinePairsAcc(Cont&& cont, Vs&&... in) {
+  CombinePairsAccImpl<element_size>(
+      std::forward<Cont>(cont), std::make_index_sequence<sizeof...(Vs) / 2>{},
+      std::forward<Vs>(in)...);
+}
+
+template <size_t element_size, size_t bs, size_t vector_bytes, class Cont,
+          class... Vs>
+HWY_INLINE void CombineRowsCPS(Cont&& cont, Vs... in) {
+  constexpr size_t N = sizeof...(Vs);
+  if constexpr (N > 1 && element_size * bs < vector_bytes) {
+    static_assert(N % 2 == 0);
+    CombinePairsAcc<element_size>(
+        [&](auto... paired) HWY_ATTR {
+          CombineRowsCPS<element_size * 2, bs, vector_bytes>(
+              std::forward<Cont>(cont), paired...);
+        },
+        in...);
+  } else {
+    std::forward<Cont>(cont)(in...);
+  }
+}
+
+// Stream -> Pack adapter for loads
+template <int I, typename T, int bs, class D, size_t row_bytes, class Cont,
+          class... Vs>
+HWY_INLINE void LoadRowsCPS(D d, const char* a, int64_t lda, Cont&& cont,
+                            Vs... vs) {
+  if constexpr (I == bs) {
+    std::forward<Cont>(cont)(vs...);
+  } else {
+    auto v = LoadBytes<D, row_bytes>(d, a + lda * I);
+    LoadRowsCPS<I + 1, T, bs, D, row_bytes>(d, a, lda, std::forward<Cont>(cont),
+                                            vs..., v);
+  }
+}
+
+// Pack -> Stream adapter for 128-bit stores
+template <class D, size_t row_bytes, size_t stores_per_row, class... Vs>
+HWY_INLINE void StoreVecs128(D d, char* b, int64_t ldb, Vs... v) {
+  ForEachIndexed(
+      [&](auto idx_c, auto&& vec) HWY_ATTR {
+        constexpr size_t i = decltype(idx_c)::value;
+        StoreLanes<D, row_bytes>(d, b, ldb, vec, i * stores_per_row,
+                                 std::make_index_sequence<stores_per_row>{});
+      },
+      v...);
+}
+
+// 128-bit (single block) transpose
+template <typename T, int bs>
+HWY_FLATTEN void TransposeMicroKernel128(const char* HWY_RESTRICT a,
+                                         int64_t lda, char* HWY_RESTRICT b,
+                                         int64_t ldb) {
+  using D = hn::Full128<uint8_t>;
+  static constexpr size_t element_size = sizeof(T);
+  static constexpr size_t row_bytes = element_size * bs;
+  static_assert(row_bytes <= kBlockBytes);
+  const D d;
+  LoadRowsCPS<0, T, bs, D, row_bytes>(d, a, lda, [&](auto... loads) HWY_ATTR {
+    CombineRowsCPS<element_size, bs, kBlockBytes>(
+        [&](auto... vecs) HWY_ATTR {
+          constexpr size_t kNumVecs = sizeof...(vecs);
+          constexpr size_t kBytesInMatrix = element_size * bs * bs;
+          constexpr size_t kStoresPerRow = bs / kNumVecs;
+          constexpr size_t kElementsPerVec = kBytesInMatrix / (bs * kNumVecs);
+          InterleaveWithinBlocksCPS<kElementsPerVec, 1, row_bytes>(
+              [&](auto... transposed) HWY_ATTR {
+                StoreVecs128<D, row_bytes, kStoresPerRow>(d, b, ldb,
+                                                          transposed...);
+              },
+              vecs...);
+        },
+        loads...);
+  });
+}
+
+// 256-bit (two block) transposes
+#if HWY_MIN_BYTES >= 32
+template <class D, class... Vs>
+HWY_INLINE void StoreRows256(D d, char* b, int64_t ldb, Vs... v) {
+  ForEachIndexed(
+      [&](auto idx_c, auto&& vec) HWY_ATTR {
+        constexpr size_t i = decltype(idx_c)::value;
+        hn::StoreU(vec, d, reinterpret_cast<hn::TFromD<D>*>(b + ldb * i));
+      },
+      v...);
+}
+
+// Invokes a continuation with an explicit parameter pack, rearranging the
+// arguments such that all even-indexed elements precede all odd-indexed
+// elements.
+template <size_t... I, class Cont, class... Vs>
+HWY_INLINE void CallContWithSeparatedImpl(Cont&& cont,
+                                          std::index_sequence<I...>,
+                                          Vs&&... vs) {
+  std::forward<Cont>(cont)(GetFromPack<2 * I>(std::forward<Vs>(vs)...)...,
+                           GetFromPack<2 * I + 1>(std::forward<Vs>(vs)...)...);
+}
+
+template <class Cont, class... Vs>
+HWY_INLINE void CallContWithSeparated(Cont&& cont, Vs&&... vs) {
+  CallContWithSeparatedImpl(std::forward<Cont>(cont),
+                            std::make_index_sequence<sizeof...(Vs) / 2>{},
+                            std::forward<Vs>(vs)...);
+}
+
+// Loads a square block of memory into 256-bit vectors, partitioning the
+// data spatially. Invokes the continuation with the vectors corresponding
+// to the left half of the block followed by the right half.
+template <int I, typename T, int bs, class D, class DHalf, class Cont,
+          class... Vs>
+HWY_INLINE void BuildVecs256SquareCPS(D d, DHalf d_half, const char* a,
+                                      int64_t lda, Cont& cont, Vs... vs) {
+  if constexpr (I == bs / 2) {
+    CallContWithSeparated(cont, vs...);
+  } else {
+    auto* row0 = reinterpret_cast<const hn::TFromD<DHalf>*>(a + lda * I);
+    auto row1 =
+        reinterpret_cast<const hn::TFromD<DHalf>*>(a + lda * (I + bs / 2));
+    auto row0_lo = hn::LoadU(d_half, row0);
+    auto row0_hi = hn::LoadU(d_half, row0 + kBlockBytes / sizeof(*row0));
+    auto row1_lo = hn::LoadU(d_half, row1);
+    auto row1_hi = hn::LoadU(d_half, row1 + kBlockBytes / sizeof(*row0));
+    auto v_lo = hn::Combine(d, row1_lo, row0_lo);
+    auto v_hi = hn::Combine(d, row1_hi, row0_hi);
+    BuildVecs256SquareCPS<I + 1, T, bs>(d, d_half, a, lda, cont, vs..., v_lo,
+                                        v_hi);
+  }
+}
+
+template <typename T, int bs, class D, class DHalf, class Cont>
+HWY_INLINE void BuildVecs256SquareCPS(D d, DHalf d_half, const char* a,
+                                      int64_t lda, Cont& cont) {
+  BuildVecs256SquareCPS<0, T, bs>(d, d_half, a, lda, cont);
+}
+
+template <typename T, int bs>
+HWY_FLATTEN void TransposeMicroKernel256Square(const char* HWY_RESTRICT a,
+                                               int64_t lda,
+                                               char* HWY_RESTRICT b,
+                                               int64_t ldb) {
+  using D = hn::FixedTag<uint8_t, 32>;
+  using DHalf = hn::Half<D>;
+  static constexpr size_t element_size = sizeof(T);
+  static constexpr size_t row_bytes = element_size * bs;
+  static_assert(row_bytes == 2 * kBlockBytes);
+  static_assert(bs % 2 == 0);
+  const D d;
+  const DHalf d_half;
+  auto cont = [&](auto... vecs) HWY_ATTR {
+    InterleaveWithinBlocksCPS<element_size, 1, kBlockBytes>(
+        [&](auto... transposed)
+            HWY_ATTR { StoreRows256(d, b, ldb, transposed...); },
+        vecs...);
+  };
+  BuildVecs256SquareCPS<T, bs>(d, d_half, a, lda, cont);
+}
+
+// Permutes 64-bit lanes across 128-bit blocks in a 256-bit vector.
+// Specifically, reorders the 4 64-bit lanes from [0, 1, 2, 3] to [0, 2, 1, 3].
+template <class V>
+HWY_INLINE V ShuffleAcrossBlocks_3120(V v) {
+  const hn::DFromV<V> d;
+  const hn::Repartition<uint64_t, decltype(d)> d64;
+  return hn::BitCast(d,
+                     hn::Per4LaneBlockShuffle<3, 1, 2, 0>(hn::BitCast(d64, v)));
+}
+
+template <class DHalf, class... Vs>
+HWY_INLINE void StoreRectPairs(DHalf d_half, char* b, int64_t ldb, Vs... v) {
+  ForEachIndexed(
+      [&](auto idx_c, auto&& vec) HWY_ATTR {
+        constexpr size_t i = decltype(idx_c)::value;
+        const auto lo = hn::LowerHalf(d_half, vec);
+        const auto hi = hn::UpperHalf(d_half, vec);
+        hn::StoreU(lo, d_half,
+                   reinterpret_cast<hn::TFromD<DHalf>*>(b + ldb * (i * 2)));
+        hn::StoreU(hi, d_half,
+                   reinterpret_cast<hn::TFromD<DHalf>*>(b + ldb * (i * 2 + 1)));
+      },
+      v...);
+}
+
+// Loads a rectangular block of memory into 256-bit vectors. Invokes the
+// continuation with the constructed sequence of vectors.
+template <int I, typename T, int bs, class D, class DHalf, class Cont,
+          class... Vs>
+HWY_INLINE void BuildVecs256RectCPS(D d, DHalf d_half, const char* a,
+                                    int64_t lda, Cont& cont, Vs... vs) {
+  if constexpr (I == bs / 2) {
+    cont(vs...);
+  } else {
+    auto lo = hn::LoadU(
+        d_half, reinterpret_cast<const hn::TFromD<DHalf>*>(a + lda * I));
+    auto hi = hn::LoadU(d_half, reinterpret_cast<const hn::TFromD<DHalf>*>(
+                                    a + lda * (I + bs / 2)));
+    auto v = hn::Combine(d, hi, lo);
+    BuildVecs256RectCPS<I + 1, T, bs>(d, d_half, a, lda, cont, vs..., v);
+  }
+}
+
+template <typename T, int bs, class D, class DHalf, class Cont>
+HWY_INLINE void BuildVecs256RectCPS(D d, DHalf d_half, const char* a,
+                                    int64_t lda, Cont& cont) {
+  BuildVecs256RectCPS<0, T, bs>(d, d_half, a, lda, cont);
+}
+
+template <typename T, int bs>
+HWY_FLATTEN void TransposeMicroKernel256Rect(const char* HWY_RESTRICT a,
+                                             int64_t lda, char* HWY_RESTRICT b,
+                                             int64_t ldb) {
+  using D = hn::FixedTag<uint8_t, 32>;
+  using DHalf = hn::Half<D>;
+  static constexpr size_t element_size = sizeof(T);
+  static constexpr size_t row_bytes = element_size * bs;
+  static_assert(row_bytes == kBlockBytes);
+  static_assert(bs % 2 == 0);
+  const D d;
+  const DHalf d_half;
+  auto cont = [&](auto... vecs) HWY_ATTR {
+    InterleaveWithinBlocksCPS<element_size, 1, kBlockBytes / 2>(
+        [&](auto... interleaved) HWY_ATTR {
+          StoreRectPairs(d_half, b, ldb,
+                         ShuffleAcrossBlocks_3120(interleaved)...);
+        },
+        vecs...);
+  };
+  BuildVecs256RectCPS<T, bs>(d, d_half, a, lda, cont);
+}
+#endif  // HWY_MIN_BYTES >= 32
+
+// Scalar fallback
+template <typename T, int bs>
+HWY_FLATTEN void TransposeMicroKernelScalar(const char* HWY_RESTRICT a,
+                                            int64_t lda, char* HWY_RESTRICT b,
+                                            int64_t ldb) {
+  for (int i = 0; i < bs; ++i) {
+    for (int j = 0; j < bs; ++j) {
+      std::memcpy(b + i * ldb + j * sizeof(T), a + j * lda + i * sizeof(T),
+                  sizeof(T));
+    }
+  }
+}
+
+// Entry point
+template <typename T, int bs>
+HWY_FLATTEN void TransposeMicroKernel(const char* HWY_RESTRICT a, int64_t lda,
+                                      char* HWY_RESTRICT b, int64_t ldb) {
+  constexpr size_t row_bytes = sizeof(T) * bs;
+  if constexpr (bs % 2 != 0) {
+    TransposeMicroKernelScalar<T, bs>(a, lda, b, ldb);
+  } else {
+#if HWY_MIN_BYTES >= 32
+    if constexpr (row_bytes == 2 * kBlockBytes) {
+      TransposeMicroKernel256Square<T, bs>(a, lda, b, ldb);
+      return;
+    } else if constexpr (row_bytes == kBlockBytes) {
+      if (lda >= ldb && lda <= kMaxSquare128StrideBytes) {
+        TransposeMicroKernel128<T, bs>(a, lda, b, ldb);
+      } else {
+        TransposeMicroKernel256Rect<T, bs>(a, lda, b, ldb);
+      }
+      return;
+    }
+#endif
+    if constexpr (row_bytes <= kBlockBytes) {
+      TransposeMicroKernel128<T, bs>(a, lda, b, ldb);
+    } else {
+      TransposeMicroKernelScalar<T, bs>(a, lda, b, ldb);
+    }
+  }
+}
+
+template <typename T, int bs>
+HWY_FLATTEN void TransposeMacroKernel(const char* HWY_RESTRICT a, int64_t lda,
+                                      int outer_bs_a, char* HWY_RESTRICT b,
+                                      int64_t ldb, int outer_bs_b) {
+  for (int i = 0; i < outer_bs_a; ++i) {
+    for (int j = 0; j < outer_bs_b; ++j) {
+      TransposeMicroKernel<T, bs>(a + bs * j * lda + i * bs * sizeof(T), lda,
+                                  b + bs * i * ldb + j * bs * sizeof(T), ldb);
+    }
+  }
+}
+
+template <int packed>
+HWY_FLATTEN void TransposeMacroKernelPacked(const char* HWY_RESTRICT a,
+                                            int64_t lda, int outer_bs_a,
+                                            char* HWY_RESTRICT b, int64_t ldb,
+                                            int outer_bs_b) {
+  constexpr size_t element_size = packed & 0xFFFF;
+  constexpr int bs = packed >> 16;
+  using T = hwy::UnsignedFromSize<element_size>;
+  TransposeMacroKernel<T, bs>(a, lda, outer_bs_a, b, ldb, outer_bs_b);
+}
+
+}  // namespace HWY_NAMESPACE
+}  // namespace xla
+HWY_AFTER_NAMESPACE();
+
+#if HWY_ONCE
+namespace xla {
+
+template <int packed>
+TransposeMacroKernelFn GetTransposeMacroKernelPacked() {
+  HWY_EXPORT_T(Table, TransposeMacroKernelPacked<packed>);
+  return HWY_DYNAMIC_POINTER_T(Table);
+}
+
+#define EXPLICIT_INSTANTIATION_PACKED(packed) \
+  template TransposeMacroKernelFn GetTransposeMacroKernelPacked<packed>()
+
+#define EXPLICIT_INSTANTIATION_FOR_SIZE(bs)      \
+  EXPLICIT_INSTANTIATION_PACKED(1 | (bs << 16)); \
+  EXPLICIT_INSTANTIATION_PACKED(2 | (bs << 16)); \
+  EXPLICIT_INSTANTIATION_PACKED(4 | (bs << 16)); \
+  EXPLICIT_INSTANTIATION_PACKED(8 | (bs << 16)); \
+  EXPLICIT_INSTANTIATION_PACKED(16 | (bs << 16))
+
+EXPLICIT_INSTANTIATION_FOR_SIZE(1);
+EXPLICIT_INSTANTIATION_FOR_SIZE(2);
+EXPLICIT_INSTANTIATION_FOR_SIZE(4);
+EXPLICIT_INSTANTIATION_FOR_SIZE(8);
+EXPLICIT_INSTANTIATION_FOR_SIZE(16);
+EXPLICIT_INSTANTIATION_FOR_SIZE(32);
+EXPLICIT_INSTANTIATION_FOR_SIZE(64);
+
+#undef EXPLICIT_INSTANTIATION_FOR_SIZE
+#undef EXPLICIT_INSTANTIATION_PACKED
+
+}  // namespace xla
+#endif
diff --git a/third_party/xla/xla/pjrt/transpose_kernels.h b/third_party/xla/xla/pjrt/transpose_kernels.h
index b4016543d26..9b47abd0e5c 100644
--- a/third_party/xla/xla/pjrt/transpose_kernels.h
+++ b/third_party/xla/xla/pjrt/transpose_kernels.h
@@ -23,392 +23,35 @@ limitations under the License.
 #include <type_traits>
 #include <utility>
 
-#include "hwy//highway.h"
 #include "xla/compiler_macros.h"
 
-HWY_BEFORE_NAMESPACE();
 namespace xla {
-namespace HWY_NAMESPACE {
-namespace hn = hwy::HWY_NAMESPACE;
 
-// Highway's model: Vectors are composed of 128-bit blocks, blocks hold lanes.
-// Interleave operations work within blocks. Cross-block movement requires
-// block-level shuffles.
-constexpr size_t kBlockBytes = 16;   // 128-bit blocks
-constexpr size_t kMaxLaneBytes = 8;  // Lanes are at most 64 bits
+using TransposeMacroKernelFn = void (*)(const char* a, int64_t lda,
+                                        int outer_bs_a, char* b, int64_t ldb,
+                                        int outer_bs_b);
 
-enum class Extract { kLo, kHi };
-// Unpack: Interleave lanes within each 128-bit block
-template <size_t element_size, Extract extract, class V>
-HWY_INLINE V Unpack(V a, V b) {
-  static_assert(element_size <= kMaxLaneBytes);
-  using D = hn::DFromV<V>;
-  using ElementT = hwy::UnsignedFromSize<element_size>;
-  const hn::Repartition<ElementT, D> d_elem;
-  auto a_elem = hn::BitCast(d_elem, a);
-  auto b_elem = hn::BitCast(d_elem, b);
-  if constexpr (extract == Extract::kLo) {
-    return hn::BitCast(D(), hn::InterleaveLower(d_elem, a_elem, b_elem));
-  } else {
-    return hn::BitCast(D(), hn::InterleaveUpper(d_elem, a_elem, b_elem));
-  }
-}
-// Load/Store
-template <class D, size_t bytes>
-HWY_INLINE hn::Vec<D> LoadBytes(D d, const void* p) {
-  const hn::FixedTag<uint8_t, bytes> d_fixed;
-  return hn::ResizeBitCast(
-      d, hn::LoadU(d_fixed, reinterpret_cast<const uint8_t*>(p)));
-}
-template <class D, size_t bytes, size_t lane>
-HWY_INLINE void StoreLane(D d, void* p, hn::Vec<D> v) {
-  if constexpr (bytes <= kMaxLaneBytes && lane > 0) {
-    using ElementT = hwy::UnsignedFromSize<bytes>;
-    const hn::Repartition<ElementT, D> d_elem;
-    const ElementT scalar = hn::ExtractLane(hn::BitCast(d_elem, v), lane);
-    hwy::CopyBytes(&scalar, reinterpret_cast<ElementT*>(p), bytes);
-  } else {
-    const hn::Repartition<uint8_t, D> d_u8;
-    auto v_u8 = hn::BitCast(d_u8, v);
-    if constexpr (lane > 0) {
-      v_u8 = hn::SlideDownLanes(d_u8, v_u8, lane * bytes);
-    }
-    const hn::FixedTag<uint8_t, bytes> d_fixed;
-    hn::StoreU(hn::ResizeBitCast(d_fixed, v_u8), d_fixed,
-               reinterpret_cast<uint8_t*>(p));
-  }
-}
-template <class D, size_t bytes, size_t... lane>
-HWY_INLINE void StoreLanes(D d, char* b, int64_t ldb, hn::Vec<D> v, size_t i,
-                           std::index_sequence<lane...>) {
-  (StoreLane<D, bytes, lane>(d, b + ldb * (i + lane), v), ...);
-}
+// Defined in transpose_kernels.cc (where foreach_target.h generates per-target
+// kernels); 0-arg getter avoids spilling argument registers in GetChosenTarget.
+template <int packed>
+TransposeMacroKernelFn GetTransposeMacroKernelPacked();
 
-// Extracts the I-th element from a parameter pack (0-indexed).
-template <size_t I, class T, class... Ts>
-HWY_INLINE decltype(auto) GetFromPack(T&& t, Ts&&... ts) {
-  static_assert(I < 1 + sizeof...(Ts), "Index out of bounds");
-  if constexpr (I == 0) {
-    return std::forward<T>(t);
-  } else {
-    return GetFromPack<I - 1>(std::forward<Ts>(ts)...);
-  }
-}
-
-// Pack iteration helper (pack -> stream): ForEachIndexed
-template <class F, size_t... I, class... Ts>
-HWY_INLINE void ForEachIndexedImpl(F&& f, std::index_sequence<I...>,
-                                   Ts&&... ts) {
-  (std::forward<F>(f)(std::integral_constant<size_t, I>{},
-                      std::forward<Ts>(ts)),
-   ...);
-}
-template <class F, class... Ts>
-HWY_INLINE void ForEachIndexed(F&& f, Ts&&... ts) {
-  ForEachIndexedImpl(std::forward<F>(f),
-                     std::make_index_sequence<sizeof...(Ts)>{},
-                     std::forward<Ts>(ts)...);
-}
-
-// CPS: UnpackStep (pack -> pack), lvalue-only inputs
-template <size_t out_i, size_t element_size, size_t step_size, class... Vs>
-HWY_INLINE auto UnpackOutFromPack(Vs&... vs) {
-  constexpr size_t group = 2 * step_size;
-  constexpr size_t base = (out_i / group) * group;
-  constexpr size_t p = out_i - base;
-  constexpr size_t j = p / 2;
-  constexpr bool is_lo = (p % 2) == 0;
-  constexpr size_t ia = base + j;
-  constexpr size_t ib = base + j + step_size;
-  auto& a = GetFromPack<ia>(vs...);
-  auto& b = GetFromPack<ib>(vs...);
-  if constexpr (is_lo) {
-    return Unpack<element_size * step_size, Extract::kLo>(a, b);
-  } else {
-    return Unpack<element_size * step_size, Extract::kHi>(a, b);
-  }
-}
-
-template <size_t element_size, size_t step_size, class Cont, size_t... I,
-          class... Vs>
-HWY_INLINE void UnpackStepCPSImpl(Cont&& cont, std::index_sequence<I...>,
-                                  Vs&... vs) {
-  std::forward<Cont>(cont)(
-      UnpackOutFromPack<I, element_size, step_size>(vs...)...);
-}
-
-template <size_t element_size, size_t step_size, class Cont, class... Vs>
-HWY_INLINE void UnpackStepCPS(Cont&& cont, Vs&... in) {
-  constexpr size_t N = sizeof...(Vs);
-  static_assert(N % (step_size * 2) == 0);
-  static_assert(element_size * step_size <= kMaxLaneBytes);
-  UnpackStepCPSImpl<element_size, step_size>(
-      std::forward<Cont>(cont), std::make_index_sequence<N>{}, in...);
-}
-
-// CPS: InterleaveWithinBlocks (pack -> pack)
-template <size_t element_size, size_t step_size, size_t unpack_limit,
-          class Cont, class... Vs>
-HWY_INLINE void InterleaveWithinBlocksCPS(Cont&& cont, Vs... in) {
-  if constexpr (element_size * step_size < unpack_limit &&
-                element_size * step_size <= kMaxLaneBytes) {
-    UnpackStepCPS<element_size, step_size>(
-        [&](auto... out) {
-          // out... are by-value parameters; do not forward as rvalues.
-          InterleaveWithinBlocksCPS<element_size, step_size * 2, unpack_limit>(
-              std::forward<Cont>(cont), out...);
-        },
-        in...);  // lvalues inside this function
-  } else {
-    std::forward<Cont>(cont)(in...);
-  }
-}
-
-// Helper template that applies the pairwise interleaving reduction.
-template <size_t element_size, class Cont, size_t... I, class... Vs>
-HWY_INLINE void CombinePairsAccImpl(Cont&& cont, std::index_sequence<I...>,
-                                    Vs&&... vs) {
-  std::forward<Cont>(cont)(Unpack<element_size, Extract::kLo>(
-      GetFromPack<2 * I>(std::forward<Vs>(vs)...),
-      GetFromPack<2 * I + 1>(std::forward<Vs>(vs)...))...);
-}
-
-// Applies a pairwise interleaving reduction across an arbitrary sequence of
-// vectors. Adjacent vectors are combined using their lower halves, and the
-// resulting halved-length sequence is forwarded to the provided continuation.
-template <size_t element_size, class Cont, class... Vs>
-HWY_INLINE void CombinePairsAcc(Cont&& cont, Vs&&... in) {
-  CombinePairsAccImpl<element_size>(
-      std::forward<Cont>(cont), std::make_index_sequence<sizeof...(Vs) / 2>{},
-      std::forward<Vs>(in)...);
-}
-
-template <size_t element_size, size_t bs, size_t vector_bytes, class Cont,
-          class... Vs>
-HWY_INLINE void CombineRowsCPS(Cont&& cont, Vs... in) {
-  constexpr size_t N = sizeof...(Vs);
-  if constexpr (N > 1 && element_size * bs < vector_bytes) {
-    static_assert(N % 2 == 0);
-    CombinePairsAcc<element_size>(
-        [&](auto... paired) {
-          CombineRowsCPS<element_size * 2, bs, vector_bytes>(
-              std::forward<Cont>(cont), paired...);
-        },
-        in...);
-  } else {
-    std::forward<Cont>(cont)(in...);
-  }
-}
-
-// Stream -> Pack adapter for loads
-template <int I, typename T, int bs, class D, size_t row_bytes, class Cont,
-          class... Vs>
-HWY_INLINE void LoadRowsCPS(D d, const char* a, int64_t lda, Cont&& cont,
-                            Vs... vs) {
-  if constexpr (I == bs) {
-    std::forward<Cont>(cont)(vs...);
-  } else {
-    auto v = LoadBytes<D, row_bytes>(d, a + lda * I);
-    LoadRowsCPS<I + 1, T, bs, D, row_bytes>(d, a, lda, std::forward<Cont>(cont),
-                                            vs..., v);
-  }
-}
-// Pack -> Stream adapter for 128-bit stores
-template <class D, size_t row_bytes, size_t stores_per_row, class... Vs>
-HWY_INLINE void StoreVecs128(D d, char* b, int64_t ldb, Vs... v) {
-  ForEachIndexed(
-      [&](auto idx_c, auto&& vec) {
-        constexpr size_t i = decltype(idx_c)::value;
-        StoreLanes<D, row_bytes>(d, b, ldb, vec, i * stores_per_row,
-                                 std::make_index_sequence<stores_per_row>{});
-      },
-      v...);
-}
-
-// 128-bit (single block) transpose
 template <typename T, int bs>
-HWY_FLATTEN void TransposeMicroKernel128(const char* HWY_RESTRICT a,
-                                         int64_t lda, char* HWY_RESTRICT b,
-                                         int64_t ldb) {
-  using D = hn::Full128<uint8_t>;
-  static constexpr size_t element_size = sizeof(T);
-  static constexpr size_t row_bytes = element_size * bs;
-  static_assert(row_bytes <= kBlockBytes);
-  const D d;
-  LoadRowsCPS<0, T, bs, D, row_bytes>(d, a, lda, [&](auto... loads) {
-    CombineRowsCPS<element_size, bs, kBlockBytes>(
-        [&](auto... vecs) {
-          constexpr size_t kNumVecs = sizeof...(vecs);
-          constexpr size_t kBytesInMatrix = element_size * bs * bs;
-          constexpr size_t kStoresPerRow = bs / kNumVecs;
-          constexpr size_t kElementsPerVec = kBytesInMatrix / (bs * kNumVecs);
-          InterleaveWithinBlocksCPS<kElementsPerVec, 1, row_bytes>(
-              [&](auto... transposed) {
-                StoreVecs128<D, row_bytes, kStoresPerRow>(d, b, ldb,
-                                                          transposed...);
-              },
-              vecs...);
-        },
-        loads...);
-  });
-}
-
-// 256-bit (two block) transposes
-#if HWY_MIN_BYTES >= 32
-template <class D, class... Vs>
-HWY_INLINE void StoreRows256(D d, char* b, int64_t ldb, Vs... v) {
-  ForEachIndexed(
-      [&](auto idx_c, auto&& vec) {
-        constexpr size_t i = decltype(idx_c)::value;
-        hn::StoreU(vec, d, reinterpret_cast<hn::TFromD<D>*>(b + ldb * i));
-      },
-      v...);
-}
-
-// Invokes a continuation with an explicit parameter pack, rearranging the
-// arguments such that all even-indexed elements precede all odd-indexed
-// elements.
-template <size_t... I, class Cont, class... Vs>
-HWY_INLINE void CallContWithSeparatedImpl(Cont&& cont,
-                                          std::index_sequence<I...>,
-                                          Vs&&... vs) {
-  std::forward<Cont>(cont)(GetFromPack<2 * I>(std::forward<Vs>(vs)...)...,
-                           GetFromPack<2 * I + 1>(std::forward<Vs>(vs)...)...);
-}
-
-template <class Cont, class... Vs>
-HWY_INLINE void CallContWithSeparated(Cont&& cont, Vs&&... vs) {
-  CallContWithSeparatedImpl(std::forward<Cont>(cont),
-                            std::make_index_sequence<sizeof...(Vs) / 2>{},
-                            std::forward<Vs>(vs)...);
-}
-
-// Loads a square block of memory into 256-bit vectors, partitioning the
-// data spatially. Invokes the continuation with the vectors corresponding
-// to the left half of the block followed by the right half.
-template <int I, typename T, int bs, class D, class DHalf, class Cont,
-          class... Vs>
-HWY_INLINE void BuildVecs256SquareCPS(D d, DHalf d_half, const char* a,
-                                      int64_t lda, Cont& cont, Vs... vs) {
-  if constexpr (I == bs / 2) {
-    CallContWithSeparated(cont, vs...);
-  } else {
-    auto* row0 = reinterpret_cast<const hn::TFromD<DHalf>*>(a + lda * I);
-    auto row1 =
-        reinterpret_cast<const hn::TFromD<DHalf>*>(a + lda * (I + bs / 2));
-    auto row0_lo = hn::LoadU(d_half, row0);
-    auto row0_hi = hn::LoadU(d_half, row0 + kBlockBytes / sizeof(*row0));
-    auto row1_lo = hn::LoadU(d_half, row1);
-    auto row1_hi = hn::LoadU(d_half, row1 + kBlockBytes / sizeof(*row0));
-    auto v_lo = hn::Combine(d, row1_lo, row0_lo);
-    auto v_hi = hn::Combine(d, row1_hi, row0_hi);
-    BuildVecs256SquareCPS<I + 1, T, bs>(d, d_half, a, lda, cont, vs..., v_lo,
-                                        v_hi);
-  }
-}
-
-template <typename T, int bs, class D, class DHalf, class Cont>
-HWY_INLINE void BuildVecs256SquareCPS(D d, DHalf d_half, const char* a,
-                                      int64_t lda, Cont& cont) {
-  BuildVecs256SquareCPS<0, T, bs>(d, d_half, a, lda, cont);
+void TransposeMacroKernelDispatch(const char* a, int64_t lda, int outer_bs_a,
+                                  char* b, int64_t ldb, int outer_bs_b) {
+  constexpr int packed = sizeof(T) | (bs << 16);
+  GetTransposeMacroKernelPacked<packed>()(a, lda, outer_bs_a, b, ldb,
+                                          outer_bs_b);
 }
 
 template <typename T, int bs>
-HWY_FLATTEN void TransposeMicroKernel256Square(const char* HWY_RESTRICT a,
-                                               int64_t lda,
-                                               char* HWY_RESTRICT b,
-                                               int64_t ldb) {
-  using D = hn::FixedTag<uint8_t, 32>;
-  using DHalf = hn::Half<D>;
-  static constexpr size_t element_size = sizeof(T);
-  static constexpr size_t row_bytes = element_size * bs;
-  static_assert(row_bytes == 2 * kBlockBytes);
-  static_assert(bs % 2 == 0);
-  const D d;
-  const DHalf d_half;
-  auto cont = [&](auto... vecs) {
-    InterleaveWithinBlocksCPS<element_size, 1, kBlockBytes>(
-        [&](auto... transposed) { StoreRows256(d, b, ldb, transposed...); },
-        vecs...);
-  };
-  BuildVecs256SquareCPS<T, bs>(d, d_half, a, lda, cont);
-}
-// Permutes 64-bit lanes across 128-bit blocks in a 256-bit vector.
-// Specifically, reorders the 4 64-bit lanes from [0, 1, 2, 3] to [0, 2, 1, 3].
-template <class V>
-HWY_INLINE V ShuffleAcrossBlocks_3120(V v) {
-  const hn::DFromV<V> d;
-  const hn::Repartition<uint64_t, decltype(d)> d64;
-  return hn::BitCast(d,
-                     hn::Per4LaneBlockShuffle<3, 1, 2, 0>(hn::BitCast(d64, v)));
+void TransposeMicroKernelDispatch(const char* a, int64_t lda, char* b,
+                                  int64_t ldb) {
+  TransposeMacroKernelDispatch<T, bs>(a, lda, /*outer_bs_a=*/1, b, ldb,
+                                      /*outer_bs_b=*/1);
 }
 
-template <class DHalf, class... Vs>
-HWY_INLINE void StoreRectPairs(DHalf d_half, char* b, int64_t ldb, Vs... v) {
-  ForEachIndexed(
-      [&](auto idx_c, auto&& vec) {
-        constexpr size_t i = decltype(idx_c)::value;
-        const auto lo = hn::LowerHalf(d_half, vec);
-        const auto hi = hn::UpperHalf(d_half, vec);
-        hn::StoreU(lo, d_half,
-                   reinterpret_cast<hn::TFromD<DHalf>*>(b + ldb * (i * 2)));
-        hn::StoreU(hi, d_half,
-                   reinterpret_cast<hn::TFromD<DHalf>*>(b + ldb * (i * 2 + 1)));
-      },
-      v...);
-}
-
-// Loads a rectangular block of memory into 256-bit vectors. Invokes the
-// continuation with the constructed sequence of vectors.
-template <int I, typename T, int bs, class D, class DHalf, class Cont,
-          class... Vs>
-HWY_INLINE void BuildVecs256RectCPS(D d, DHalf d_half, const char* a,
-                                    int64_t lda, Cont& cont, Vs... vs) {
-  if constexpr (I == bs / 2) {
-    cont(vs...);
-  } else {
-    auto lo = hn::LoadU(
-        d_half, reinterpret_cast<const hn::TFromD<DHalf>*>(a + lda * I));
-    auto hi = hn::LoadU(d_half, reinterpret_cast<const hn::TFromD<DHalf>*>(
-                                    a + lda * (I + bs / 2)));
-    auto v = hn::Combine(d, hi, lo);
-    BuildVecs256RectCPS<I + 1, T, bs>(d, d_half, a, lda, cont, vs..., v);
-  }
-}
-
-template <typename T, int bs, class D, class DHalf, class Cont>
-HWY_INLINE void BuildVecs256RectCPS(D d, DHalf d_half, const char* a,
-                                    int64_t lda, Cont& cont) {
-  BuildVecs256RectCPS<0, T, bs>(d, d_half, a, lda, cont);
-}
-
-template <typename T, int bs>
-HWY_FLATTEN void TransposeMicroKernel256Rect(const char* HWY_RESTRICT a,
-                                             int64_t lda, char* HWY_RESTRICT b,
-                                             int64_t ldb) {
-  using D = hn::FixedTag<uint8_t, 32>;
-  using DHalf = hn::Half<D>;
-  static constexpr size_t element_size = sizeof(T);
-  static constexpr size_t row_bytes = element_size * bs;
-  static_assert(row_bytes == kBlockBytes);
-  static_assert(bs % 2 == 0);
-  const D d;
-  const DHalf d_half;
-  auto cont = [&](auto... vecs) {
-    InterleaveWithinBlocksCPS<element_size, 1, kBlockBytes / 2>(
-        [&](auto... interleaved) {
-          StoreRectPairs(d_half, b, ldb,
-                         ShuffleAcrossBlocks_3120(interleaved)...);
-        },
-        vecs...);
-  };
-  BuildVecs256RectCPS<T, bs>(d, d_half, a, lda, cont);
-}
-#endif  // HWY_MIN_BYTES >= 32
-
-}  // namespace HWY_NAMESPACE
 }  // namespace xla
-HWY_AFTER_NAMESPACE();
 
 #ifdef XLA_HAS_SSE2
 #include <immintrin.h>  // IWYU pragma: keep
diff --git a/third_party/xla/xla/pjrt/transpose_test.cc b/third_party/xla/xla/pjrt/transpose_test.cc
index 520b450ffb8..b804505c96b 100644
--- a/third_party/xla/xla/pjrt/transpose_test.cc
+++ b/third_party/xla/xla/pjrt/transpose_test.cc
@@ -46,6 +46,7 @@ limitations under the License.
 #include "absl/strings/str_join.h"
 #include "absl/types/span.h"
 #include "unsupported/Eigen/CXX11/Tensor"
+#include "hwy//targets.h"
 #include "xla/array.h"
 #include "xla/ef57.h"
 #include "xla/hlo/testlib/test.h"
@@ -66,26 +67,38 @@ ABSL_FLAG(int, transpose_num_random_testcases, 10,
 
 namespace xla {
 
+// Number of bs x bs tiles along each dimension of the macroblocks used by the
+// macro-kernel test and benchmark.
 template <typename T, int bs>
-void TestMicroKernelEquivalence() {
-  // Test both tightly-packed tiles (0 padding) and tiles with 64 bytes of
-  // padding between rows (which exercises wide-stride / rectangular kernels).
-  static_assert(64 % sizeof(T) == 0);
+constexpr int BenchmarkOuterBlockSize() {
+  int max_by_elems = kMaxOuterBlockElems / bs;
+  int max_by_bytes = kMaxSquare128StrideBytes / (bs * sizeof(T));
+  return std::max(1, std::min(max_by_elems, max_by_bytes));
+}
 
-  for (int input_stride : {bs, bs + static_cast<int>(64 / sizeof(T))}) {
-    for (int output_stride : {bs, bs + static_cast<int>(64 / sizeof(T))}) {
-      std::vector<T> input(bs * input_stride);
-      std::vector<T> expected_output(bs * output_stride);
-      std::vector<T> actual_output(bs * output_stride);
+// Checks that `kernel(a, lda, b, ldb)` transposes an n x n matrix of T, where
+// n = outer_bs * bs.
+template <typename T, int bs, typename Kernel>
+void CheckTransposeKernel(int outer_bs, Kernel kernel) {
+  // Test both tightly-packed rows (0 padding) and rows with 64 bytes of
+  // padding (which exercises wide-stride / rectangular kernels).
+  static_assert(64 % sizeof(T) == 0);
+  const int n = outer_bs * bs;
+
+  for (int input_stride : {n, n + static_cast<int>(64 / sizeof(T))}) {
+    for (int output_stride : {n, n + static_cast<int>(64 / sizeof(T))}) {
+      std::vector<T> input(n * input_stride);
+      std::vector<T> expected_output(n * output_stride);
+      std::vector<T> actual_output(n * output_stride);
 
       // Initialize input
-      for (int i = 0; i < bs * input_stride; ++i) {
+      for (int i = 0; i < n * input_stride; ++i) {
         input[i] = static_cast<T>(static_cast<float>(i % 100));
       }
 
       // Compute reference
-      for (int row = 0; row < bs; ++row) {
-        for (int col = 0; col < bs; ++col) {
+      for (int row = 0; row < n; ++row) {
+        for (int col = 0; col < n; ++col) {
           expected_output[col * output_stride + row] =
               input[row * input_stride + col];
         }
@@ -93,47 +106,143 @@ void TestMicroKernelEquivalence() {
 
       const int64_t lda = input_stride * sizeof(T);
       const int64_t ldb = output_stride * sizeof(T);
-      TransposeMicroKernel<T, bs>::Apply(
-          reinterpret_cast<const char*>(input.data()), lda,
-          reinterpret_cast<char*>(actual_output.data()), ldb);
+      kernel(reinterpret_cast<const char*>(input.data()), lda,
+             reinterpret_cast<char*>(actual_output.data()), ldb);
 
-      EXPECT_EQ(0, std::memcmp(expected_output.data(), actual_output.data(),
-                               bs * ldb))
+      EXPECT_EQ(
+          0, std::memcmp(expected_output.data(), actual_output.data(), n * ldb))
           << "Mismatch for sizeof(T)=" << sizeof(T) << " bs=" << bs
-          << " lda=" << lda << " ldb=" << ldb;
+          << " outer_bs=" << outer_bs << " lda=" << lda << " ldb=" << ldb;
     }
   }
 }
 
+template <typename T, int bs>
+void TestMicroKernelEquivalence() {
+  CheckTransposeKernel<T, bs>(
+      /*outer_bs=*/1, [](const char* a, int64_t lda, char* b, int64_t ldb) {
+        TransposeMicroKernelDispatch<T, bs>(a, lda, b, ldb);
+      });
+}
+
+template <typename T, int bs>
+void TestMacroKernelEquivalence() {
+  constexpr int kOuterBs = BenchmarkOuterBlockSize<T, bs>();
+  CheckTransposeKernel<T, bs>(
+      kOuterBs, [](const char* a, int64_t lda, char* b, int64_t ldb) {
+        TransposeMacroKernelDispatch<T, bs>(a, lda, kOuterBs, b, ldb, kOuterBs);
+      });
+}
+
+// TODO(b/543906735): Remove along with the legacy SIMD intrinsic kernels once
+// transpose.cc uses the Highway kernels.
+template <typename T, int bs>
+void TestLegacyMicroKernelEquivalence() {
+  CheckTransposeKernel<T, bs>(
+      /*outer_bs=*/1, [](const char* a, int64_t lda, char* b, int64_t ldb) {
+        TransposeMicroKernel<T, bs>::Apply(a, lda, b, ldb);
+      });
+}
+
 TEST(TransposeMicroKernelTest, ExactEquivalence) {
+  for (int64_t target : hwy::SupportedAndGeneratedTargets()) {
+    hwy::SetSupportedTargetsForTest(target);
+
+    // TransposeMicroKernel256Square is triggered when a logical row of the tile
+    // (bs * sizeof(T)) is 256 bits (32 bytes).
+    TestMicroKernelEquivalence<int8_t, 32>();
+    TestMicroKernelEquivalence<int16_t, 16>();
+    TestMicroKernelEquivalence<float, 8>();
+    TestMicroKernelEquivalence<int64_t, 4>();
+    TestMicroKernelEquivalence<absl::uint128, 2>();
+
+    // TransposeMicroKernel128 or TransposeMicroKernel256Rect is triggered when
+    // a logical row of the tile (bs * sizeof(T)) is 128 bits (16 bytes).
+    // TransposeMicroKernel128 operates on 128-bit blocks when gathering from
+    // memory (lda >= ldb && lda <= 64B) while TransposeMicroKernel256Rect packs
+    // two rows into 256-bit vectors when scattering to memory (ldb > lda) or
+    // for large strides.
+    TestMicroKernelEquivalence<int8_t, 16>();
+    TestMicroKernelEquivalence<int16_t, 8>();
+    TestMicroKernelEquivalence<bfloat16, 8>();
+    TestMicroKernelEquivalence<float, 4>();
+    TestMicroKernelEquivalence<int64_t, 2>();
+
+    // Smaller or larger cases fall back to TransposeMicroKernel128 or scalar
+    // loop.
+    TestMicroKernelEquivalence<int8_t, 8>();
+    TestMicroKernelEquivalence<int16_t, 4>();
+    TestMicroKernelEquivalence<bfloat16, 4>();
+    TestMicroKernelEquivalence<float, 2>();
+    TestMicroKernelEquivalence<int8_t, 4>();
+    TestMicroKernelEquivalence<int16_t, 2>();
+    TestMicroKernelEquivalence<int8_t, 2>();
+    TestMicroKernelEquivalence<int8_t, 64>();
+  }
+  hwy::SetSupportedTargetsForTest(0);
+}
+
+// Same tile shapes as above, but transposing a macroblock of
+// BenchmarkOuterBlockSize() x BenchmarkOuterBlockSize() tiles.
+TEST(TransposeMacroKernelTest, ExactEquivalence) {
+  for (int64_t target : hwy::SupportedAndGeneratedTargets()) {
+    hwy::SetSupportedTargetsForTest(target);
+
+    TestMacroKernelEquivalence<int8_t, 32>();
+    TestMacroKernelEquivalence<int16_t, 16>();
+    TestMacroKernelEquivalence<float, 8>();
+    TestMacroKernelEquivalence<int64_t, 4>();
+    TestMacroKernelEquivalence<absl::uint128, 2>();
+
+    TestMacroKernelEquivalence<int8_t, 16>();
+    TestMacroKernelEquivalence<int16_t, 8>();
+    TestMacroKernelEquivalence<bfloat16, 8>();
+    TestMacroKernelEquivalence<float, 4>();
+    TestMacroKernelEquivalence<int64_t, 2>();
+
+    TestMacroKernelEquivalence<int8_t, 8>();
+    TestMacroKernelEquivalence<int16_t, 4>();
+    TestMacroKernelEquivalence<bfloat16, 4>();
+    TestMacroKernelEquivalence<float, 2>();
+    TestMacroKernelEquivalence<int8_t, 4>();
+    TestMacroKernelEquivalence<int16_t, 2>();
+    TestMacroKernelEquivalence<int8_t, 2>();
+    TestMacroKernelEquivalence<int8_t, 64>();
+  }
+  hwy::SetSupportedTargetsForTest(0);
+}
+
+// TODO(b/543906735): Remove along with the legacy SIMD intrinsic kernels once
+// transpose.cc uses the Highway kernels.
+TEST(TransposeMicroKernelTest, LegacyExactEquivalence) {
   // AvxSquareTransposeMicroKernelImpl is triggered when a logical row of the
   // tile (bs * sizeof(T)) is exactly 256 bits to fit in __m256i.
-  TestMicroKernelEquivalence<int8_t, 32>();
-  TestMicroKernelEquivalence<int16_t, 16>();
-  TestMicroKernelEquivalence<float, 8>();
-  TestMicroKernelEquivalence<int64_t, 4>();
-  TestMicroKernelEquivalence<absl::uint128, 2>();
+  TestLegacyMicroKernelEquivalence<int8_t, 32>();
+  TestLegacyMicroKernelEquivalence<int16_t, 16>();
+  TestLegacyMicroKernelEquivalence<float, 8>();
+  TestLegacyMicroKernelEquivalence<int64_t, 4>();
+  TestLegacyMicroKernelEquivalence<absl::uint128, 2>();
 
   // SseSquareTransposeMicroKernelImpl or AvxRectangularTransposeMicroKernelImpl
   // is triggered when a logical row of the tile (bs * sizeof(T)) is exactly
   // 128 bits. SseSquare operates directly on __m128i when gathering from memory
   // (lda >= ldb && lda <= 64B) while AvxRectangular packs two rows into __m256i
   // when scattering to memory (ldb > lda) or for large strides.
-  TestMicroKernelEquivalence<int8_t, 16>();
-  TestMicroKernelEquivalence<int16_t, 8>();
-  TestMicroKernelEquivalence<bfloat16, 8>();
-  TestMicroKernelEquivalence<float, 4>();
-  TestMicroKernelEquivalence<int64_t, 2>();
+  TestLegacyMicroKernelEquivalence<int8_t, 16>();
+  TestLegacyMicroKernelEquivalence<int16_t, 8>();
+  TestLegacyMicroKernelEquivalence<bfloat16, 8>();
+  TestLegacyMicroKernelEquivalence<float, 4>();
+  TestLegacyMicroKernelEquivalence<int64_t, 2>();
 
   // Smaller or larger cases fall back to either Vec128 or a for loop.
-  TestMicroKernelEquivalence<int8_t, 8>();
-  TestMicroKernelEquivalence<int16_t, 4>();
-  TestMicroKernelEquivalence<bfloat16, 4>();
-  TestMicroKernelEquivalence<float, 2>();
-  TestMicroKernelEquivalence<int8_t, 4>();
-  TestMicroKernelEquivalence<int16_t, 2>();
-  TestMicroKernelEquivalence<int8_t, 2>();
-  TestMicroKernelEquivalence<int8_t, 64>();
+  TestLegacyMicroKernelEquivalence<int8_t, 8>();
+  TestLegacyMicroKernelEquivalence<int16_t, 4>();
+  TestLegacyMicroKernelEquivalence<bfloat16, 4>();
+  TestLegacyMicroKernelEquivalence<float, 2>();
+  TestLegacyMicroKernelEquivalence<int8_t, 4>();
+  TestLegacyMicroKernelEquivalence<int16_t, 2>();
+  TestLegacyMicroKernelEquivalence<int8_t, 2>();
+  TestLegacyMicroKernelEquivalence<int8_t, 64>();
 }
 
 class TestTransposePlan : public TransposePlan {
@@ -1386,13 +1495,6 @@ enum class StrideMode {
   kWideScatter,    // lda < ldb (exercises AvxRectangular scatter)
 };
 
-template <typename T, int bs>
-constexpr int BenchmarkOuterBlockSize() {
-  int max_by_elems = kMaxOuterBlockElems / bs;
-  int max_by_bytes = kMaxSquare128StrideBytes / (bs * sizeof(T));
-  return std::max(1, std::min(max_by_elems, max_by_bytes));
-}
-
 template <typename T, int bs, StrideMode kStrideMode = StrideMode::kTight>
 void BM_TransposeMicroKernel(::testing::benchmark::State& state) {
   constexpr int kOuterBs = BenchmarkOuterBlockSize<T, bs>();
@@ -1416,13 +1518,7 @@ void BM_TransposeMicroKernel(::testing::benchmark::State& state) {
     benchmark::DoNotOptimize(lda);
     benchmark::DoNotOptimize(ldb);
     benchmark::DoNotOptimize(outer_bs);
-    for (int i = 0; i < outer_bs; ++i) {
-      for (int j = 0; j < outer_bs; ++j) {
-        TransposeMicroKernel<T, bs>::Apply(
-            a + bs * j * lda + i * bs * sizeof(T), lda,
-            b + bs * i * ldb + j * bs * sizeof(T), ldb);
-      }
-    }
+    TransposeMacroKernelDispatch<T, bs>(a, lda, outer_bs, b, ldb, outer_bs);
     benchmark::ClobberMemory();
   }
   state.SetBytesProcessed(state.iterations() * kOuterBs * kOuterBs * bs * bs *

From b69a193c0069eef980e2192eda59e2d0c1b49272 Mon Sep 17 00:00:00 2001
From: Mikhail Goncharov <goncharov@google.com>
Date: Thu, 15 Oct 2026 14:45:27 -0700
Subject: [PATCH 0912/1000] [XLA:GPU] Deprecate
 --xla_gpu_experimental_enable_tiling_propagation flag and proto field

This path is being depricated and check allows us and users to see "last warning" before we drop the code completely.
updated concat_prune_inactive_small_tile that was added recently.

PiperOrigin-RevId: 1000406333
---
 .../gpu/autotuner/block_level_emitter.h       |   2 +-
 .../backends/gpu/autotuner/fission_backend.cc |   3 +-
 .../xla/xla/backends/gpu/autotuner/triton.cc  |  43 ++-----
 .../triton/cost_model_config_optimization.cc  |   5 +-
 .../backends/gpu/codegen/triton/support.cc    |  10 --
 .../concat_prune_inactive_small_tile.hlo      |   9 +-
 .../gpu/codegen/triton/xtile_compiler.cc      |  95 +++++---------
 .../gpu/codegen/triton/xtile_test_base.cc     |  69 +++++-----
 .../transforms/convert_triton_gemm_config.cc  | 118 +-----------------
 .../transforms/fusion_block_level_rewriter.cc |  22 +---
 .../backends/gpu/transforms/gemm_fusion.cc    |  77 +++++-------
 .../gpu/transforms/priority_fusion.cc         |   5 +-
 .../gpu/transforms/softmax_rewriter_triton.cc |  48 +++----
 .../codegen/tiling/symbolic_tile_analysis.cc  |   9 +-
 .../tiling/symbolic_tile_analysis_test.cc     |  39 ++++++
 .../tiling/tiled_hlo_computation_test.cc      |   4 +
 .../codegen/tiling/tiled_hlo_schedule_test.cc |   3 +
 .../tiling/tiling_specification_test.cc       |   4 +
 .../codegen/xtile/codegen/fusion_emitter.cc   |   3 +-
 .../codegen/tiled_emitter_constraints_test.cc |   4 +
 third_party/xla/xla/debug_options_flags.cc    |   9 +-
 third_party/xla/xla/service/gpu/BUILD         |  59 ---------
 .../xla/xla/service/gpu/gpu_compiler.cc       |   8 +-
 .../xla/xla/service/gpu/ir_emission_utils.cc  |   9 --
 .../model/triton_emitter_constraints_test.cc  |   4 +
 third_party/xla/xla/xla.proto                 |   4 +-
 26 files changed, 207 insertions(+), 458 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter.h b/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter.h
index eb6cb7416d3..364948add68 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter.h
+++ b/third_party/xla/xla/backends/gpu/autotuner/block_level_emitter.h
@@ -66,7 +66,7 @@ class BlockLevelEmitterBackend : public GpuCodegenBackend {
         indexing_performance_model_(
             &target_config->device_description, &fusion_analysis_cache_,
             shape_size_fn_, &mlir_context_,
-            debug_options->xla_gpu_experimental_enable_tiling_propagation(),
+            /*use_experimental_tiling=*/true,
             debug_options
                 ->xla_gpu_experimental_enable_same_shape_multi_output_fusion(),
             mlir_context_pool),
diff --git a/third_party/xla/xla/backends/gpu/autotuner/fission_backend.cc b/third_party/xla/xla/backends/gpu/autotuner/fission_backend.cc
index 451386976bb..1a0f8e90a73 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/fission_backend.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/fission_backend.cc
@@ -288,8 +288,7 @@ absl::StatusOr<absl::Duration> FissionBackend::EstimateFissionPrologueEpilogue(
   GpuPerformanceModelWithIndexingAnalysis indexing_cost_model{
       &device_info, &fusion_analysis_cache, HloCostAnalysis::DefaultShapeSize,
       mlir_context->get(),
-      /*use_experimental_tiling=*/
-      debug_options().xla_gpu_experimental_enable_tiling_propagation(),
+      /*use_experimental_tiling=*/true,
       /*enable_same_shape_multi_output_fusion=*/
       debug_options()
           .xla_gpu_experimental_enable_same_shape_multi_output_fusion(),
diff --git a/third_party/xla/xla/backends/gpu/autotuner/triton.cc b/third_party/xla/xla/backends/gpu/autotuner/triton.cc
index c16e42306e9..390ecc56087 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/triton.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/triton.cc
@@ -371,38 +371,21 @@ bool TritonBackend::IsSupported(const HloInstruction& instr) {
       VLOG(1) << "Failed to borrow MLIRContext: " << mlir_context.status();
       return false;
     }
-    if (instr.GetModule()
-            ->config()
-            .debug_options()
-            .xla_gpu_experimental_enable_tiling_propagation()) {
-      auto ts = experimental::TilingSpace::Create(*fusion_adaptor,
-                                                  (*mlir_context)->get());
-      if (!ts.ok()) {
-        VLOG(1) << "Failed to create tiling space: " << ts.status().message();
-        return false;
-      }
-      auto tiled_computation_or = experimental::TiledHloComputation::Tile(
-          *fusion_adaptor, std::move(ts.value()));
-      if (!tiled_computation_or.ok()) {
-        VLOG(1) << "Fusion is not tileable with experimental tiling: "
-                << tiled_computation_or.status().message();
-        return false;
-      }
-      // We don't have concrete tile sizes here and don't validate Triton
-      // constraints here.
-      return true;
-    }
-
-    auto device_info = target_config().device_description;
-    SymbolicTileAnalysisOrError analysis_or_error =
-        SymbolicTileAnalysis::AnalyzeFusion(
-            *fusion_adaptor, (*mlir_context)->get(),
-            TritonEmitterConstraints::GetBuilder(device_info));
-    if (const auto* fusion_decision =
-            std::get_if<FusionDecision>(&analysis_or_error)) {
-      VLOG(1) << "Fusion not tileable: " << fusion_decision->Explain();
+    auto ts = experimental::TilingSpace::Create(*fusion_adaptor,
+                                                (*mlir_context)->get());
+    if (!ts.ok()) {
+      VLOG(1) << "Failed to create tiling space: " << ts.status().message();
       return false;
     }
+    auto tiled_computation_or = experimental::TiledHloComputation::Tile(
+        *fusion_adaptor, std::move(ts.value()));
+    if (!tiled_computation_or.ok()) {
+      VLOG(1) << "Fusion is not tileable with experimental tiling: "
+              << tiled_computation_or.status().message();
+      return false;
+    }
+    // We don't have concrete tile sizes here and don't validate Triton
+    // constraints here.
     return true;
   }
   return backend_config.kind() == kCuDnnFusionKind ||
diff --git a/third_party/xla/xla/backends/gpu/autotuner/triton/cost_model_config_optimization.cc b/third_party/xla/xla/backends/gpu/autotuner/triton/cost_model_config_optimization.cc
index 5db981d0c9b..75697d9ba93 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/triton/cost_model_config_optimization.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/triton/cost_model_config_optimization.cc
@@ -347,8 +347,7 @@ absl::StatusOr<std::vector<TritonGemmConfig>> OptimizeConfigsWithCostModel(
     const DebugOptions& debug_options, mlir::MLIRContext* mlir_context) {
   namespace detail = cost_model_config_optimization_detail;
 
-  const bool use_experimental_tiling =
-      debug_options.xla_gpu_experimental_enable_tiling_propagation();
+  const bool use_experimental_tiling = true;
   const bool enable_same_shape_multi_output_fusion =
       debug_options
           .xla_gpu_experimental_enable_same_shape_multi_output_fusion();
@@ -468,7 +467,7 @@ EstimateConfigsWithCostModel(const HloDotInstruction* dot,
       detail::OrderedEstimatesAndConfigs estimated_configs,
       detail::EstimateConfigs(
           extracted.context, configs, mlir_context,
-          debug_options.xla_gpu_experimental_enable_tiling_propagation(),
+          /*use_experimental_tiling=*/true,
           debug_options
               .xla_gpu_experimental_enable_same_shape_multi_output_fusion()));
 
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/support.cc b/third_party/xla/xla/backends/gpu/codegen/triton/support.cc
index 8d222408670..cf5a4976583 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/support.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/support.cc
@@ -420,11 +420,6 @@ CodegenDecision CanTritonHandleReduce(
         "Unsupported reduction computation by Triton.");
   }
 
-  bool tiling_propagation_enabled =
-      reduce.GetModule()
-          ->config()
-          .debug_options()
-          .xla_gpu_experimental_enable_tiling_propagation();
 
   if (reduce.operand_count() != 2) {
     return CodegenDecision::Forbid(
@@ -436,11 +431,6 @@ CodegenDecision CanTritonHandleReduce(
         "Reduction must have at least one dimension.");
   }
 
-  if (reduce.dimensions().size() > 1 && !tiling_propagation_enabled) {
-    return CodegenDecision::Forbid(
-        "Multi-dimension reductions are only supported when experimental "
-        "tiling propagation is enabled.");
-  }
 
   return CodegenDecision::Allow();
 }
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive_small_tile.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive_small_tile.hlo
index dba7de98b75..a85a1167b3a 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive_small_tile.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/concat/concat_prune_inactive_small_tile.hlo
@@ -1,8 +1,5 @@
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation=false %s 2>&1 | FileCheck %s
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation=false %s
-
-// RUN: fusion_to_triton --xla_gpu_experimental_enable_tiling_propagation %s 2>&1 | FileCheck %s --check-prefixes=CHECK-TILING
-// RUN: triton_test_correctness --xla_gpu_experimental_enable_tiling_propagation %s
+// RUN: fusion_to_triton %s 2>&1 | FileCheck %s
+// RUN: triton_test_correctness %s
 
 // This is a fork of concat_prune_inactive.hlo that uses a smaller concat-dim
 // tile size (16 instead of 64). With the sign-agnostic boundary-alignment
@@ -17,8 +14,6 @@
 
 // CHECK:      xtile.entry_func @triton_fn
 
-// CHECK-TILING:      xtile.entry_func @triton_fn
-
 HloModule min_repro, is_scheduled=true, entry_computation_layout={(f32[64,10]{1,0}, f32[1,16,10]{2,1,0}, f32[1,16,10]{2,1,0})->f32[64,10]{1,0}}
 
 %fused_computation (param_0: f32[64,10], param_1: f32[1,16,10], param_2: f32[1,16,10]) -> f32[64,10] {
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc b/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc
index 4c66f0ecacf..c47f69cb67e 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/xtile_compiler.cc
@@ -310,72 +310,45 @@ absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> TileAndEmitXTileModule(
     const BlockLevelParameters& block_level_parameters,
     absl::Span<mlir::Type> opaque_args_types, mlir::MLIRContext& mlir_context,
     bool use_experimental_tiling, bool enable_same_shape_multi_output_fusion) {
-  const HloComputation* computation = fusion.fused_instructions_computation();
+  using experimental::TiledHloComputation;
+  using experimental::TilingSpace;
 
-  if (use_experimental_tiling || IsCollectiveFusion(fusion)) {
-    using experimental::TiledHloComputation;
-    using experimental::TilingSpace;
+  auto fusion_adaptor = HloFusionAdaptor::ForInstruction(&fusion);
+  ABSL_ASSIGN_OR_RETURN(std::unique_ptr<TilingSpace> tiling_space,
+                   TilingSpace::Create(*fusion_adaptor, &mlir_context));
 
-    auto fusion_adaptor = HloFusionAdaptor::ForInstruction(&fusion);
-    ABSL_ASSIGN_OR_RETURN(std::unique_ptr<TilingSpace> tiling_space,
-                     TilingSpace::Create(*fusion_adaptor, &mlir_context));
-
-    VLOG(3) << "fusion instruction: " << fusion.ToString() << "\n";
-    VLOG(3) << "tiling space: " << tiling_space->ToString();
-    if (VLOG_IS_ON(4)) {
-      XLA_VLOG_LINES(
-          4, absl::StrCat("HLO module to reproduce:\n",
-                          ExtractInstructionIntoNewModule(fusion)->ToString(
-                              HloPrintOptions::ShortParsable())));
-    }
-    ABSL_ASSIGN_OR_RETURN(
-        llvm::SmallVector<int64_t> tile_sizes,
-        GetTilingSpaceConcreteSizes(*tiling_space, block_level_parameters,
-                                    enable_same_shape_multi_output_fusion));
-    ABSL_RETURN_IF_ERROR(
-        tiling_space->AssignTileSizes(xtile::GetPaddedTileSizes(tile_sizes)));
-
-    ABSL_ASSIGN_OR_RETURN(
-        TiledHloComputation tiled_computation,
-        TiledHloComputation::Tile(*fusion_adaptor, std::move(tiling_space)));
-    tiled_computation.Simplify();
-    tiled_computation.SortInstructionsPostOrder();
-    if (Decision constraints = experimental::VerifyTritonConstraints(
-            tiled_computation, device_info);
-        !constraints) {
-      return absl::InvalidArgumentError(
-          absl::StrCat("Triton constraints violated during codegen: ",
-                       constraints.Explain()));
-    }
-    VLOG(4) << "tiled computation: " << tiled_computation.ToString();
-    return xtile::EmitXTileModule(
-        fn_name, fusion, tiled_computation, mlir_context,
-        absl::MakeSpan(opaque_args_types),
-        std::make_optional(device_info.gpu_compute_capability()),
-        block_level_parameters.num_tiles_per_pid);
+  VLOG(3) << "fusion instruction: " << fusion.ToString() << "\n";
+  VLOG(3) << "tiling space: " << tiling_space->ToString();
+  if (VLOG_IS_ON(4)) {
+    XLA_VLOG_LINES(
+        4, absl::StrCat("HLO module to reproduce:\n",
+                        ExtractInstructionIntoNewModule(fusion)->ToString(
+                            HloPrintOptions::ShortParsable())));
   }
-  SymbolicTileAnalysisOrError symbolic_tile_analysis_or =
-      SymbolicTileAnalysis::AnalyzeComputation(
-          *computation, &mlir_context,
-          TritonEmitterConstraints::GetBuilder(device_info));
+  ABSL_ASSIGN_OR_RETURN(
+      llvm::SmallVector<int64_t> tile_sizes,
+      GetTilingSpaceConcreteSizes(*tiling_space, block_level_parameters,
+                                  enable_same_shape_multi_output_fusion));
+  ABSL_RETURN_IF_ERROR(
+      tiling_space->AssignTileSizes(xtile::GetPaddedTileSizes(tile_sizes)));
 
-  if (std::holds_alternative<FusionDecision>(symbolic_tile_analysis_or)) {
-    return Internal(
-        "Unsupported fusion in CreateTritonModule: %s",
-        std::get<FusionDecision>(symbolic_tile_analysis_or).Explain());
+  ABSL_ASSIGN_OR_RETURN(
+      TiledHloComputation tiled_computation,
+      TiledHloComputation::Tile(*fusion_adaptor, std::move(tiling_space)));
+  tiled_computation.Simplify();
+  tiled_computation.SortInstructionsPostOrder();
+  if (Decision constraints =
+          experimental::VerifyTritonConstraints(tiled_computation, device_info);
+      !constraints) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Triton constraints violated during codegen: ", constraints.Explain()));
   }
-
-  const auto& symbolic_tile_analysis =
-      std::get<SymbolicTileAnalysis>(symbolic_tile_analysis_or);
-
-  ABSL_ASSIGN_OR_RETURN(Tiling tiling,
-                   TilingFromAnnotatedFusion(symbolic_tile_analysis,
-                                             block_level_parameters));
-
+  VLOG(4) << "tiled computation: " << tiled_computation.ToString();
   return xtile::EmitXTileModule(
-      fn_name, fusion, symbolic_tile_analysis, tiling, mlir_context,
+      fn_name, fusion, tiled_computation, mlir_context,
       absl::MakeSpan(opaque_args_types),
-      std::make_optional(device_info.gpu_compute_capability()));
+      std::make_optional(device_info.gpu_compute_capability()),
+      block_level_parameters.num_tiles_per_pid);
 }
 
 absl::StatusOr<TritonKernelSource> CreateTritonModule(
@@ -387,9 +360,7 @@ absl::StatusOr<TritonKernelSource> CreateTritonModule(
 
   const DebugOptions& debug_options =
       fusion.GetModule()->config().debug_options();
-  bool use_experimental_tiling =
-      debug_options.xla_gpu_experimental_enable_tiling_propagation() ||
-      IsCollectiveFusion(fusion);
+  bool use_experimental_tiling = true;
   bool enable_same_shape_multi_output_fusion =
       debug_options
           .xla_gpu_experimental_enable_same_shape_multi_output_fusion();
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/xtile_test_base.cc b/third_party/xla/xla/backends/gpu/codegen/triton/xtile_test_base.cc
index eac48fb3d8b..6507ba2761a 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/xtile_test_base.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/xtile_test_base.cc
@@ -120,47 +120,36 @@ XTileTestBase::CreateXTileIrAndFileCheck(
     absl::string_view filecheck_pattern) {
   mlir::OwningOpRef<mlir::ModuleOp> xtile_dialect_module;
   LoadMlirDialectsForTriton(*mlir_context());
-  if (computation.parent()
-          ->config()
-          .debug_options()
-          .xla_gpu_experimental_enable_tiling_propagation()) {
-    namespace ge = ::xla::gpu::experimental;
-    auto* fusion = Cast<HloFusionInstruction>(computation.FusionInstruction());
-    auto fusion_adaptor = HloFusionAdaptor::ForInstruction(fusion);
-    ABSL_ASSIGN_OR_RETURN(std::unique_ptr<ge::TilingSpace> tiling_space,
-                     ge::TilingSpace::Create(*fusion_adaptor, mlir_context()));
-    ABSL_ASSIGN_OR_RETURN(
-        llvm::SmallVector<int64_t> concrete_sizes,
-        GetTilingSpaceConcreteSizes(
-            *tiling_space, block_level_parameters,
-            computation.parent()
-                ->config()
-                .debug_options()
-                .xla_gpu_experimental_enable_same_shape_multi_output_fusion()));
-    ABSL_RETURN_IF_ERROR(tiling_space->AssignTileSizes(
-        xtile::GetPaddedTileSizes(concrete_sizes)));
-    ABSL_ASSIGN_OR_RETURN(ge::TiledHloComputation tiled_computation,
-                     ge::TiledHloComputation::Tile(*fusion_adaptor,
-                                                   std::move(tiling_space)));
-    tiled_computation.Simplify();
-    tiled_computation.SortInstructionsPostOrder();
-    if (Decision constraints = ge::VerifyTritonConstraints(
-            tiled_computation, TestGpuDeviceInfo::RTXA6000DeviceInfo());
-        !constraints) {
-      return absl::InternalError(
-          absl::StrCat("Triton constraints violated during test codegen: ",
-                       constraints.Explain()));
-    }
-    ABSL_ASSIGN_OR_RETURN(
-        xtile_dialect_module,
-        xtile::EmitXTileModule("xtile_dialect_fn", *fusion, tiled_computation,
-                               *mlir_context()));
-  } else {
-    ABSL_ASSIGN_OR_RETURN(xtile_dialect_module,
-                     CreateXTileIrAndFileCheckLegacy(
-                         mlir_context(), computation, block_level_parameters,
-                         filecheck_pattern));
+  namespace ge = ::xla::gpu::experimental;
+  auto* fusion = Cast<HloFusionInstruction>(computation.FusionInstruction());
+  auto fusion_adaptor = HloFusionAdaptor::ForInstruction(fusion);
+  ABSL_ASSIGN_OR_RETURN(std::unique_ptr<ge::TilingSpace> tiling_space,
+                   ge::TilingSpace::Create(*fusion_adaptor, mlir_context()));
+  ABSL_ASSIGN_OR_RETURN(
+      llvm::SmallVector<int64_t> concrete_sizes,
+      GetTilingSpaceConcreteSizes(
+          *tiling_space, block_level_parameters,
+          computation.parent()
+              ->config()
+              .debug_options()
+              .xla_gpu_experimental_enable_same_shape_multi_output_fusion()));
+  ABSL_RETURN_IF_ERROR(
+      tiling_space->AssignTileSizes(xtile::GetPaddedTileSizes(concrete_sizes)));
+  ABSL_ASSIGN_OR_RETURN(
+      ge::TiledHloComputation tiled_computation,
+      ge::TiledHloComputation::Tile(*fusion_adaptor, std::move(tiling_space)));
+  tiled_computation.Simplify();
+  tiled_computation.SortInstructionsPostOrder();
+  if (Decision constraints = ge::VerifyTritonConstraints(
+          tiled_computation, TestGpuDeviceInfo::RTXA6000DeviceInfo());
+      !constraints) {
+    return absl::InternalError(
+        absl::StrCat("Triton constraints violated during test codegen: ",
+                     constraints.Explain()));
   }
+  ABSL_ASSIGN_OR_RETURN(xtile_dialect_module,
+                   xtile::EmitXTileModule("xtile_dialect_fn", *fusion,
+                                          tiled_computation, *mlir_context()));
   std::string out;
   llvm::raw_string_ostream os(out);
   xtile_dialect_module->print(os);
diff --git a/third_party/xla/xla/backends/gpu/transforms/convert_triton_gemm_config.cc b/third_party/xla/xla/backends/gpu/transforms/convert_triton_gemm_config.cc
index cea8b5eb781..f7401515039 100644
--- a/third_party/xla/xla/backends/gpu/transforms/convert_triton_gemm_config.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/convert_triton_gemm_config.cc
@@ -324,121 +324,9 @@ absl::StatusOr<BlockLevelParameters> FindBlockLevelParameters(
   // instruction that have a contracting dimension and all tile sizes except
   // m, n, and k are 1.
   const HloComputation* computation = dot->parent();
-  if (dot->GetModule()
-          ->config()
-          .debug_options()
-          .xla_gpu_experimental_enable_tiling_propagation()) {
-    return FindBlockLevelParameters(
-        *HloFusionAdaptor::ForComputation(computation), dot, config,
-        mlir_context, device_description);
-  }
-
-  VLOG(3) << "FindOutputTileSizesForEpilogue of computation: "
-          << computation->ToString();
-  SymbolicTileAnalysisOrError analysis_or =
-      SymbolicTileAnalysis::AnalyzeComputation(
-          *computation, mlir_context,
-          TritonEmitterConstraints::GetBuilder(device_description));
-
-  if (const auto* fusion_decision = std::get_if<FusionDecision>(&analysis_or)) {
-    std::unique_ptr<HloModule> extracted_computation_module =
-        ExtractInstructionIntoNewModule(*computation->FusionInstruction());
-    return absl::InternalError(absl::StrCat(
-        "Failed to analyze the computation (", fusion_decision->Explain(),
-        "):\n", extracted_computation_module->ToString()));
-  }
-
-  const auto& analysis = std::get<SymbolicTileAnalysis>(analysis_or);
-  const auto& tiled_instructions = analysis.GetSymbolicTiledHloComputation();
-  auto is_dot = [&](const auto& instr) { return instr->hlo() == dot; };
-  auto tiled_dot_it = absl::c_find_if(tiled_instructions, is_dot);
-  if (tiled_dot_it == tiled_instructions.end()) {
-    return absl::InternalError(absl::StrCat(
-        "Couldn't find a symbolic tiled instruction for ", dot->ToString()));
-  }
-  const SymbolicTiledHloInstruction& tiled_dot = **tiled_dot_it;
-
-  auto get_tile_sizes = [&](int64_t rank) {
-    QCHECK_GE(rank, 2) << "Expected at least rank 2 for the dot, got " << rank
-                       << " in computation " << computation->ToString();
-    // We always expect the shape to be [1, ..., block_m, block_n], by
-    // construction of GemmFusions.
-    llvm::SmallVector<int64_t> tile_sizes(rank - 2, 1);
-    tile_sizes.append({config.block_m, config.block_n});
-    return tile_sizes;
-  };
-
-  VLOG(3) << "FindOutputTileSizesForEpilogue: dot shape: "
-          << dot->shape().ToString();
-  auto expected_dot_tile_sizes =
-      get_tile_sizes(dot->shape().dimensions().size());
-  VLOG(2) << "FindOutputTileSizesForEpilogue: " << tiled_dot.ToString()
-          << "\nConstraints: "
-          << analysis.GetTilingSpecification().constraints().ToString(
-                 analysis.GetTilingSpecification().num_parameters())
-          << "Expected dot tile sizes: "
-          << absl::StrJoin(expected_dot_tile_sizes, " ");
-
-  // Try all permutations of the dot tile sizes to see if any of them satisfy
-  // the constraints of the analysis and map to the given config of the dot.
-  int64_t out_rank =
-      computation->root_instruction()->shape().dimensions().size();
-  VLOG(3) << "FindOutputTileSizesForEpilogue: computation root shape: "
-          << computation->root_instruction()->shape().ToString();
-  llvm::SmallVector<int64_t> output_tile_sizes = get_tile_sizes(out_rank);
-
-  absl::c_sort(output_tile_sizes);
-
-  const TilingSpecification& tiling_specification =
-      analysis.GetTilingSpecification();
-
-  do {
-    VLOG(4) << "trying output_tile_sizes = ("
-            << absl::StrJoin(output_tile_sizes, ",") << ")";
-    Tiling::TileMapping tile_mapping;
-    tile_mapping[dot] = {config.block_k};
-    // If the `dot` is a root, we need to assign both the hidden parameter and
-    // the output parameters to it.
-    if (dot->IsRoot()) {
-      tile_mapping[dot].insert(tile_mapping[dot].end(),
-                               output_tile_sizes.begin(),
-                               output_tile_sizes.end());
-    } else {
-      tile_mapping[dot->parent()->root_instruction()] = {
-          output_tile_sizes.begin(), output_tile_sizes.end()};
-    }
-
-    Tiling tiling(std::move(tile_mapping));
-    ABSL_ASSIGN_OR_RETURN(bool parameters_satisfy_constraints,
-                     analysis.ParametersSatisfyConstraints(tiling));
-    if (!parameters_satisfy_constraints) {
-      VLOG(4) << "Parameters don't satisfy constraints";
-      continue;
-    }
-    ABSL_ASSIGN_OR_RETURN(FlatTiling flat_tiling_parameters,
-                     tiling.Flatten(tiling_specification));
-    llvm::SmallVector<int64_t> mapped_dot_tile_sizes =
-        EvaluateTileSizes(tiled_dot.symbolic_tile(), flat_tiling_parameters);
-    if (mapped_dot_tile_sizes == expected_dot_tile_sizes) {
-      BlockLevelParameters params;
-      params.output_tile_sizes = {std::vector<int64_t>(
-          output_tile_sizes.begin(), output_tile_sizes.end())};
-      params.num_warps = config.num_warps;
-      params.num_ctas = config.num_ctas;
-      params.num_stages = config.num_stages;
-      params.is_tma_allowed = config.is_tma_allowed;
-      params.is_warp_specialization_allowed =
-          config.is_warp_specialization_allowed;
-      params.waves_per_eu = config.waves_per_eu;
-      return params;
-    }
-    VLOG(4) << "mapped_dot_tile_sizes: "
-            << absl::StrJoin(mapped_dot_tile_sizes, ",")
-            << " != " << absl::StrJoin(expected_dot_tile_sizes, ",");
-  } while (absl::c_next_permutation(output_tile_sizes));
-
-  return absl::InternalError(absl::StrCat(
-      "Couldn't find output tile sizes that satisfy ", tiled_dot.ToString()));
+  return FindBlockLevelParameters(
+      *HloFusionAdaptor::ForComputation(computation), dot, config, mlir_context,
+      device_description);
 }
 
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc
index 86f941bb4b6..93668556d76 100644
--- a/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter.cc
@@ -185,8 +185,7 @@ absl::StatusOr<bool> ShouldTryRewriteFusion(
   if (IsSameShapeMultiOutputFusion(*fusion,
                                    Shape::Equal().IgnoreElementType()) &&
       debug_options
-          .xla_gpu_experimental_enable_same_shape_multi_output_fusion() &&
-      debug_options.xla_gpu_experimental_enable_tiling_propagation()) {
+          .xla_gpu_experimental_enable_same_shape_multi_output_fusion()) {
     return true;
   }
 
@@ -195,12 +194,6 @@ absl::StatusOr<bool> ShouldTryRewriteFusion(
     return false;
   }
 
-  // Scan fusions require the new tiling propagation infrastructure.
-  if (IsScanFusion(fusion) &&
-      !debug_options.xla_gpu_experimental_enable_tiling_propagation()) {
-    return false;
-  }
-
   if (debug_options.xla_gpu_experimental_enable_fusion_block_level_rewriter()) {
     return true;
   }
@@ -351,14 +344,11 @@ absl::StatusOr<bool> FusionBlockLevelRewriter::RunImpl(
     }
     HloFusionInstruction* fusion_instruction =
         ::xla::Cast<HloFusionInstruction>(computation->FusionInstruction());
-    ABSL_ASSIGN_OR_RETURN(
-        bool changed,
-        ProcessFusionInstruction(
-            fusion_instruction, device_info_, shape_size_, mlir_context_,
-            module->config()
-                .debug_options()
-                .xla_gpu_experimental_enable_tiling_propagation(),
-            thread_pool_, mlir_context_pool_));
+    ABSL_ASSIGN_OR_RETURN(bool changed, ProcessFusionInstruction(
+                                       fusion_instruction, device_info_,
+                                       shape_size_, mlir_context_,
+                                       /*use_experimental_tiling=*/true,
+                                       thread_pool_, mlir_context_pool_));
 
     has_changed |= changed;
   }
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 d28df71b899..7e9797180ec 100644
--- a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc
@@ -1290,54 +1290,41 @@ FusionDecision CanUnpackS4ParametersInFusion(
 FusionDecision CanTile(mlir::MLIRContext& mlir_context,
                        const HloFusionAdaptor& fusion,
                        const se::DeviceDescription& device_description) {
-  if (fusion.GetRoots()[0]
-          .instruction()
-          .GetModule()
-          ->config()
-          .debug_options()
-          .xla_gpu_experimental_enable_tiling_propagation()) {
-    namespace ge = ::xla::gpu::experimental;
-    auto ts = ge::TilingSpace::Create(fusion, &mlir_context);
-    if (!ts.ok()) {
-      return FusionDecision::Forbid(absl::StrCat(
-          "Failed to create tiling space: ", ts.status().message()));
-    }
-    auto tiled_computation =
-        ge::TiledHloComputation::Tile(fusion, std::move(ts.value()));
-    if (!tiled_computation.ok()) {
-      return FusionDecision::Forbid(
-          absl::StrCat("Fusion is not tileable with experimental tiling: ",
-                       tiled_computation.status().message()));
-    }
-
-    // The symbolic tiling may be satisfiable in theory, but we don't have any
-    // tiles that we will actually try. For example, the constraints may require
-    // a tile size to be divisible by 3, but all Triton tiles are powers of 2.
-    std::optional<HloInstructionAdaptor> dot =
-        HloBfsFindIf(fusion.GetRoots(), fusion, [](HloInstructionAdaptor node) {
-          return node.opcode() == HloOpcode::kDot;
-        });
-    if (!dot.has_value()) {
-      return FusionDecision::Forbid("No dot instruction found in the fusion.");
-    }
-    for (const TritonGemmConfig& config :
-         GetDefaultTritonConfigs(device_description.gpu_compute_capability())) {
-      absl::StatusOr<xtile::BlockLevelParameters> params =
-          FindBlockLevelParameters(fusion, &dot->instruction(), config,
-                                   &mlir_context, device_description);
-      if (params.ok()) {
-        return FusionDecision::Allow();
-      }
-    }
+  namespace ge = ::xla::gpu::experimental;
+  auto ts = ge::TilingSpace::Create(fusion, &mlir_context);
+  if (!ts.ok()) {
     return FusionDecision::Forbid(
-        "None of the default Triton configs can be used to tile the fusion.");
+        absl::StrCat("Failed to create tiling space: ", ts.status().message()));
   }
-  auto fusion_analysis =
-      SymbolicTileAnalysis::AnalyzeFusion(fusion, &mlir_context);
-  if (std::holds_alternative<FusionDecision>(fusion_analysis)) {
-    return std::get<FusionDecision>(fusion_analysis);
+  auto tiled_computation =
+      ge::TiledHloComputation::Tile(fusion, std::move(ts.value()));
+  if (!tiled_computation.ok()) {
+    return FusionDecision::Forbid(
+        absl::StrCat("Fusion is not tileable with experimental tiling: ",
+                     tiled_computation.status().message()));
   }
-  return FusionDecision::Allow();
+
+  // The symbolic tiling may be satisfiable in theory, but we don't have any
+  // tiles that we will actually try. For example, the constraints may require
+  // a tile size to be divisible by 3, but all Triton tiles are powers of 2.
+  std::optional<HloInstructionAdaptor> dot =
+      HloBfsFindIf(fusion.GetRoots(), fusion, [](HloInstructionAdaptor node) {
+        return node.opcode() == HloOpcode::kDot;
+      });
+  if (!dot.has_value()) {
+    return FusionDecision::Forbid("No dot instruction found in the fusion.");
+  }
+  for (const TritonGemmConfig& config :
+       GetDefaultTritonConfigs(device_description.gpu_compute_capability())) {
+    absl::StatusOr<xtile::BlockLevelParameters> params =
+        FindBlockLevelParameters(fusion, &dot->instruction(), config,
+                                 &mlir_context, device_description);
+    if (params.ok()) {
+      return FusionDecision::Allow();
+    }
+  }
+  return FusionDecision::Forbid(
+      "None of the default Triton configs can be used to tile the fusion.");
 }
 
 // Returns true if fusing `producer` into `consumer` is possible and supported.
diff --git a/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc b/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc
index 8a6f2334e64..0d0f4157be0 100644
--- a/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/priority_fusion.cc
@@ -1390,10 +1390,7 @@ absl::StatusOr<bool> PriorityFusion::RunImpl(
           .debug_options()
           .xla_gpu_experimental_enable_triton_heroless_priority_fusion();
 
-  bool use_experimental_tiling =
-      module->config()
-          .debug_options()
-          .xla_gpu_experimental_enable_tiling_propagation();
+  bool use_experimental_tiling = true;
 
   FusionDeduplicationCache fusion_deduplication_cache =
       FusionDeduplicationCache::Create(*module, IsFusible);
diff --git a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.cc b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.cc
index d16ae5e238a..520d13ef24d 100644
--- a/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/softmax_rewriter_triton.cc
@@ -451,41 +451,21 @@ absl::StatusOr<bool> CanSymbolicTileAnalysisTileDiamond(
   ABSL_ASSIGN_OR_RETURN(HloFusionInstruction * normalization_fusion,
                    MakeFusionForDiamond(diamond));
   mlir::MLIRContext mlir_context;
-  bool use_experimental_tiling =
-      normalization_fusion->GetModule()
-          ->config()
-          .debug_options()
-          .xla_gpu_experimental_enable_tiling_propagation();
-
-  bool can_tile = false;
-  if (use_experimental_tiling) {
-    using experimental::TiledHloComputation;
-    using experimental::TilingSpace;
-    std::unique_ptr<HloFusionAdaptor> fusion_adaptor =
-        HloFusionAdaptor::ForInstruction(normalization_fusion);
-    ABSL_ASSIGN_OR_RETURN(std::unique_ptr<TilingSpace> tiling_space,
-                     TilingSpace::Create(*fusion_adaptor, &mlir_context));
-    absl::StatusOr<TiledHloComputation> tiled_computation_or =
-        TiledHloComputation::Tile(*fusion_adaptor, std::move(tiling_space));
-    // We don't have concrete tile sizes here and don't validate Triton
-    // constraints here.
-    if (!tiled_computation_or.ok()) {
-      VLOG(2) << absl::StrCat("TiledHloComputation::Tile failed: ",
-                              tiled_computation_or.status());
-    }
-    can_tile = tiled_computation_or.ok();
-  } else {
-    SymbolicTileAnalysisOrError symbolic_tile_analysis_or_error =
-        SymbolicTileAnalysis::AnalyzeComputation(
-            *normalization_fusion->called_computation(), &mlir_context,
-            TritonEmitterConstraints::GetBuilder(device_info));
-    if (std::holds_alternative<Decision>(symbolic_tile_analysis_or_error)) {
-      VLOG(2) << "SymbolicTileAnalysis failed: "
-              << std::get<Decision>(symbolic_tile_analysis_or_error).Explain();
-    }
-    can_tile = std::holds_alternative<SymbolicTileAnalysis>(
-        symbolic_tile_analysis_or_error);
+  using experimental::TiledHloComputation;
+  using experimental::TilingSpace;
+  std::unique_ptr<HloFusionAdaptor> fusion_adaptor =
+      HloFusionAdaptor::ForInstruction(normalization_fusion);
+  ABSL_ASSIGN_OR_RETURN(std::unique_ptr<TilingSpace> tiling_space,
+                   TilingSpace::Create(*fusion_adaptor, &mlir_context));
+  absl::StatusOr<TiledHloComputation> tiled_computation_or =
+      TiledHloComputation::Tile(*fusion_adaptor, std::move(tiling_space));
+  // We don't have concrete tile sizes here and don't validate Triton
+  // constraints here.
+  if (!tiled_computation_or.ok()) {
+    VLOG(2) << absl::StrCat("TiledHloComputation::Tile failed: ",
+                            tiled_computation_or.status());
   }
+  bool can_tile = tiled_computation_or.ok();
 
   ABSL_RETURN_IF_ERROR(diamond.root->GetModule()->RemoveEmbeddedComputation(
       normalization_fusion->called_computation()));
diff --git a/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis.cc b/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis.cc
index 38e915a268b..ddd77bc6e4d 100644
--- a/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis.cc
+++ b/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis.cc
@@ -1398,13 +1398,8 @@ SymbolicTileAnalysis::AnalyzeFromInstruction(
 /*static*/ SymbolicTileAnalysisOrError SymbolicTileAnalysis::AnalyzeFusion(
     const HloFusionAdaptor& fusion, MLIRContext* mlir_context,
     EmitterSpecificConstraintsBuilder emitter_specific_constraints_builder) {
-  CHECK(!fusion.GetRoots()
-             .front()
-             .instruction()
-             .GetModule()
-             ->config()
-             .debug_options()
-             .xla_gpu_experimental_enable_tiling_propagation());
+  LOG(FATAL) << "Calling into symbolic tile analysis is disallowed. "
+             << "Please use experimental tiling propagation instead.";
   RegisterSymbolicExprStorage(mlir_context);
   auto real_root_index_or = GetRealRootIndex(fusion.GetRoots());
   if (!real_root_index_or.ok()) {
diff --git a/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis_test.cc b/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis_test.cc
index d7dc6c536ac..ae4e5bbc456 100644
--- a/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis_test.cc
+++ b/third_party/xla/xla/codegen/tiling/symbolic_tile_analysis_test.cc
@@ -177,10 +177,49 @@ class FakeEmitterSpecificConstraints : public EmitterSpecificConstraints {
   int64_t dim0_tile_size_;
 };
 
+using SymbolicTileAnalysisDeathTest = HloHardwareIndependentTestBase;
+
+TEST_F(SymbolicTileAnalysisDeathTest, CallingSymbolicTileAnalysisFailsCheck) {
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
+HloModule m
+
+fused_computation {
+  p0 = f32[10,10] parameter(0)
+  ROOT r = f32[10,10] abs(p0)
+}
+
+ENTRY main {
+  p0 = f32[10,10] parameter(0)
+  ROOT fusion = f32[10,10] fusion(p0), kind=kCustom, calls=fused_computation
+}
+)"));
+  mlir::MLIRContext mlir_context;
+  EXPECT_DEATH(SymbolicTileAnalysis::AnalyzeComputation(
+                   *module->entry_computation()
+                        ->root_instruction()
+                        ->fused_instructions_computation(),
+                   &mlir_context),
+               "Calling into symbolic tile analysis is disallowed");
+
+  module->mutable_config()
+      .mutable_debug_options()
+      .set_xla_gpu_experimental_enable_tiling_propagation(false);
+  EXPECT_DEATH(SymbolicTileAnalysis::AnalyzeComputation(
+                   *module->entry_computation()
+                        ->root_instruction()
+                        ->fused_instructions_computation(),
+                   &mlir_context),
+               "Calling into symbolic tile analysis is disallowed");
+}
+
 class SymbolicTileAnalysisTest : public HloHardwareIndependentTestBase {
  public:
   SymbolicTileAnalysisTest() = default;
 
+  void SetUp() override {
+    GTEST_SKIP() << "Legacy symbolic tile analysis is disallowed.";
+  }
+
   std::optional<SymbolicTileAnalysis> TryAnalyzeModule(
       HloModule* module,
       EmitterSpecificConstraintsBuilder emitter_specific_constraints_builder =
diff --git a/third_party/xla/xla/codegen/tiling/tiled_hlo_computation_test.cc b/third_party/xla/xla/codegen/tiling/tiled_hlo_computation_test.cc
index 853b69b13c5..d2e3c2e9736 100644
--- a/third_party/xla/xla/codegen/tiling/tiled_hlo_computation_test.cc
+++ b/third_party/xla/xla/codegen/tiling/tiled_hlo_computation_test.cc
@@ -45,6 +45,10 @@ class TiledHloComputationTest : public HloHardwareIndependentTestBase {
  public:
   TiledHloComputationTest() { RegisterSymbolicExprStorage(&mlir_context_); }
 
+  void SetUp() override {
+    GTEST_SKIP() << "Legacy symbolic tile analysis is disallowed.";
+  }
+
  protected:
   DebugOptions GetDebugOptionsForTest() const override {
     DebugOptions debug_options =
diff --git a/third_party/xla/xla/codegen/tiling/tiled_hlo_schedule_test.cc b/third_party/xla/xla/codegen/tiling/tiled_hlo_schedule_test.cc
index a58d9ef83f4..22a324ab1ce 100644
--- a/third_party/xla/xla/codegen/tiling/tiled_hlo_schedule_test.cc
+++ b/third_party/xla/xla/codegen/tiling/tiled_hlo_schedule_test.cc
@@ -52,6 +52,9 @@ using ::testing::HasSubstr;
 class TiledHloScheduleTest : public HloHardwareIndependentTestBase {
  protected:
   TiledHloScheduleTest() { RegisterSymbolicExprStorage(&mlir_context_); }
+  void SetUp() override {
+    GTEST_SKIP() << "Legacy symbolic tile analysis is disallowed.";
+  }
   DebugOptions GetDebugOptionsForTest() const override {
     DebugOptions debug_options =
         HloHardwareIndependentTestBase::GetDebugOptionsForTest();
diff --git a/third_party/xla/xla/codegen/tiling/tiling_specification_test.cc b/third_party/xla/xla/codegen/tiling/tiling_specification_test.cc
index 6a302a41de5..1d72d1c88ce 100644
--- a/third_party/xla/xla/codegen/tiling/tiling_specification_test.cc
+++ b/third_party/xla/xla/codegen/tiling/tiling_specification_test.cc
@@ -51,6 +51,10 @@ class TilingSpecificationTest : public HloHardwareIndependentTestBase {
  public:
   TilingSpecificationTest() { RegisterSymbolicExprStorage(&mlir_context_); }
 
+  void SetUp() override {
+    GTEST_SKIP() << "Legacy symbolic tile analysis is disallowed.";
+  }
+
  protected:
   DebugOptions GetDebugOptionsForTest() const override {
     DebugOptions debug_options =
diff --git a/third_party/xla/xla/codegen/xtile/codegen/fusion_emitter.cc b/third_party/xla/xla/codegen/xtile/codegen/fusion_emitter.cc
index 77f11982fbf..950f2f42454 100644
--- a/third_party/xla/xla/codegen/xtile/codegen/fusion_emitter.cc
+++ b/third_party/xla/xla/codegen/xtile/codegen/fusion_emitter.cc
@@ -1256,7 +1256,8 @@ absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> EmitXTileModule(
     return absl::InvalidArgumentError("Multi-output fusion is disabled.");
   }
 
-  CHECK(!debug_options.xla_gpu_experimental_enable_tiling_propagation());
+  LOG(FATAL) << "Calling into legacy fusion emitter is disallowed. "
+             << "Please use experimental tiling propagation instead.";
 
   const HloComputation* hlo_computation =
       fusion.fused_instructions_computation();
diff --git a/third_party/xla/xla/codegen/xtile/codegen/tiled_emitter_constraints_test.cc b/third_party/xla/xla/codegen/xtile/codegen/tiled_emitter_constraints_test.cc
index 863b83b4af1..ef96001e82e 100644
--- a/third_party/xla/xla/codegen/xtile/codegen/tiled_emitter_constraints_test.cc
+++ b/third_party/xla/xla/codegen/xtile/codegen/tiled_emitter_constraints_test.cc
@@ -53,6 +53,10 @@ class TiledEmitterConstraintsTest : public HloHardwareIndependentTestBase {
 
  public:
   TiledEmitterConstraintsTest() { RegisterSymbolicExprStorage(&mlir_context_); }
+
+  void SetUp() override {
+    GTEST_SKIP() << "Legacy symbolic tile analysis is disallowed.";
+  }
   std::optional<SymbolicTileAnalysis> TryAnalyzeModule(
       HloModule* module, bool with_tiled_emitter_specific_constraints = true) {
     EmitterSpecificConstraintsBuilder constraints_builder = nullptr;
diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index bc01ecf1dfc..77c317693c7 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -3901,12 +3901,9 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
       debug_options->xla_gpu_rocm_max_trace_events(),
       "Maximum number of ROCm trace events (applies to callback/activity/"
       "annotation). Set as high as memory allows; up to 1e9."));
-  flag_list->push_back(tsl::Flag(
-      "xla_gpu_experimental_enable_tiling_propagation",
-      bool_setter_for(
-          &DebugOptions::set_xla_gpu_experimental_enable_tiling_propagation),
-      debug_options->xla_gpu_experimental_enable_tiling_propagation(),
-      "If true, enable experimental tiling propagation."));
+  flag_list->push_back(
+      tsl::Flag("xla_gpu_experimental_enable_tiling_propagation",
+                noop_flag_setter<bool>, true, "Deprecated, do not use."));
   flag_list->push_back(tsl::Flag(
       "xla_gpu_experimental_enable_reduction_dimension_tiling",
       bool_setter_for(
diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index d10ad3500ed..c51ad03cf3f 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -1784,13 +1784,11 @@ cc_library(
         ":conv_layout_normalization",
         ":cublas_cudnn",
         ":early_exit_compilation_result",
-        ":execution_stream_assignment",
         ":export_hlo",
         ":flag_utils",
         ":fusion_dispatch_pipeline",
         ":fusion_pipeline",
         ":gpu_aot_compilation_result",
-        ":gpu_constants",
         ":gpu_executable",
         ":gpu_executable_proto_cc",
         ":gpu_float_support",
@@ -1800,21 +1798,16 @@ cc_library(
         ":gpu_spmd_pipeline",
         ":hlo_fusion_stats",
         ":ir_emission_utils",
-        ":ir_emitter_context",
         ":kernel_reuse_cache",
-        ":launch_dimensions",
         ":matmul_utils",
         ":metrics",
         ":mlir_context_pool",
         ":pre_scheduling_copy_insertion_pipeline",
         ":reduction_utils",
         ":stream_executor_util",
-        ":thunk_emitter",
         "//xla:autotune_results_proto_cc",
-        "//xla:debug_options_flags",
         "//xla:shape_util",
         "//xla:status_macros",
-        "//xla:types",
         "//xla:util",
         "//xla:xla_data_proto_cc",
         "//xla:xla_proto_cc",
@@ -1824,17 +1817,10 @@ cc_library(
         "//xla/backends/cpu:target_machine_options",
         "//xla/backends/cpu/nanort:nanort_client",
         "//xla/backends/cpu/nanort:nanort_executable",
-        "//xla/backends/gpu/autotuner:block_level_emitter",
-        "//xla/backends/gpu/autotuner:factory",
-        "//xla/backends/gpu/autotuner:native_emitter",
         "//xla/backends/gpu/codegen:cubin_custom_kernel_compiler",
         "//xla/backends/gpu/codegen:kernel_compiler",
         "//xla/backends/gpu/codegen/emitters:mlir_kernel_emitter",
-        "//xla/backends/gpu/codegen/kernels:custom_kernel",
-        "//xla/backends/gpu/codegen/kernels:ptx_custom_kernel",
         "//xla/backends/gpu/codegen/triton:support",
-        "//xla/backends/gpu/runtime:custom_kernel_thunk",
-        "//xla/backends/gpu/runtime:execution_stream_id",
         "//xla/backends/gpu/runtime:host_execute_thunk",
         "//xla/backends/gpu/runtime:runtime_intrinsics",
         "//xla/backends/gpu/runtime:sequential_thunk",
@@ -1847,7 +1833,6 @@ cc_library(
         "//xla/backends/gpu/transforms:algorithm_checker",
         "//xla/backends/gpu/transforms:async_wrapper",
         "//xla/backends/gpu/transforms:composite_rewriter",
-        "//xla/backends/gpu/transforms:conv_kind_assignment",
         "//xla/backends/gpu/transforms:conv_rewriter",
         "//xla/backends/gpu/transforms:convert_triton_gemm_config",
         "//xla/backends/gpu/transforms:cudnn_custom_call_converter",
@@ -1867,7 +1852,6 @@ cc_library(
         "//xla/backends/gpu/transforms:estimate_cub_sort_scratch_size",
         "//xla/backends/gpu/transforms:explicit_collectives_group_async_wrapper",
         "//xla/backends/gpu/transforms:explicit_stream_annotation_async_wrapper",
-        "//xla/backends/gpu/transforms:fusion_block_level_rewriter",
         "//xla/backends/gpu/transforms:fusion_wrapper",
         "//xla/backends/gpu/transforms:gemm_broadcast_folding_rewriter",
         "//xla/backends/gpu/transforms:gemm_fusion",
@@ -1910,33 +1894,25 @@ cc_library(
         "//xla/backends/gpu/transforms/collectives:all_reduce_blueconnect",
         "//xla/backends/gpu/transforms/collectives:all_reduce_decomposer",
         "//xla/backends/gpu/transforms/collectives:all_reduce_splitter",
-        "//xla/backends/gpu/transforms/collectives:async_collective_annotator",
         "//xla/backends/gpu/transforms/collectives:collective_backend_assigner",
         "//xla/backends/gpu/transforms/collectives:collective_combiner_annotator",
         "//xla/backends/gpu/transforms/collectives:collective_domain_assigner",
         "//xla/backends/gpu/transforms/collectives:collective_fusion",
         "//xla/backends/gpu/transforms/collectives:collective_kernel_strategy_annotator",
-        "//xla/backends/gpu/transforms/collectives:collective_ops_utils",
         "//xla/backends/gpu/transforms/collectives:collective_permute_cycle_decomposer",
         "//xla/backends/gpu/transforms/collectives:collective_pipelining_analyzer",
-        "//xla/backends/gpu/transforms/collectives:collective_select_folder",
         "//xla/backends/gpu/transforms/collectives:convert_async_collectives_to_sync",
         "//xla/backends/gpu/transforms/collectives:gpu_all_gather_combiner",
         "//xla/backends/gpu/transforms/collectives:gpu_all_reduce_combiner",
         "//xla/backends/gpu/transforms/collectives:gpu_collective_combiner_utils",
         "//xla/backends/gpu/transforms/collectives:gpu_reduce_scatter_combiner",
         "//xla/backends/gpu/transforms/collectives:legalize_collective_domain",
-        "//xla/codegen:llvm_kernel_source",
-        "//xla/codegen/emitters:kernel_arguments",
         "//xla/core/host_offloading:hlo_host_device_type_call_wrapper",
         "//xla/core/host_offloading:host_compute_asyncifier",
         "//xla/core/host_offloading:host_offloading_executable",
         "//xla/ffi:ffi_registry",
         "//xla/hlo/analysis:alias_info",
         "//xla/hlo/analysis:hlo_alias_analysis",
-        "//xla/hlo/analysis:hlo_dataflow_analysis",
-        "//xla/hlo/analysis:hlo_ordering",
-        "//xla/hlo/analysis:symbolic_map",
         "//xla/hlo/builder:xla_computation",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass",
@@ -1950,20 +1926,15 @@ cc_library(
         "//xla/hlo/transforms:propagate_call_metadata",
         "//xla/hlo/transforms:while_loop_trip_count_annotator",
         "//xla/hlo/transforms/collectives:all_gather_broadcast_reorder",
-        "//xla/hlo/transforms/collectives:all_gather_combiner",
         "//xla/hlo/transforms/collectives:all_gather_remove_degenerate_dims",
-        "//xla/hlo/transforms/collectives:all_reduce_combiner",
         "//xla/hlo/transforms/collectives:all_reduce_contiguous",
         "//xla/hlo/transforms/collectives:all_reduce_promotion",
         "//xla/hlo/transforms/collectives:all_reduce_reassociate",
         "//xla/hlo/transforms/collectives:all_reduce_simplifier",
-        "//xla/hlo/transforms/collectives:async_collective_creator",
         "//xla/hlo/transforms/collectives:collective_permute_combiner",
         "//xla/hlo/transforms/collectives:collective_permute_decomposer",
         "//xla/hlo/transforms/collectives:collective_quantizer",
         "//xla/hlo/transforms/collectives:collectives_schedule_linearizer",
-        "//xla/hlo/transforms/collectives:convert_async_collectives_to_sync",
-        "//xla/hlo/transforms/collectives:reduce_scatter_combiner",
         "//xla/hlo/transforms/collectives:reduce_scatter_reassociate",
         "//xla/hlo/transforms/collectives:while_loop_all_reduce_code_motion",
         "//xla/hlo/transforms/expanders:bitcast_dtypes_expander",
@@ -2016,18 +1987,14 @@ cc_library(
         "//xla/hlo/transforms/simplifiers:sub_byte_normalization",
         "//xla/hlo/transforms/simplifiers:tuple_simplifier",
         "//xla/hlo/transforms/simplifiers:zero_sized_hlo_elimination",
-        "//xla/hlo/utils:hlo_query",
-        "//xla/hlo/utils:hlo_traversal",
         "//xla/pjrt/distributed:key_value_store_interface",
         "//xla/pjrt/proto:compile_options_proto_cc",
-        "//xla/runtime:object_pool",
         "//xla/service:async_collective_custom_call_rewriter",
         "//xla/service:batched_gather_scatter_normalizer",
         "//xla/service:batchnorm_expander",
         "//xla/service:buffer_assignment",
         "//xla/service:buffer_value",
         "//xla/service:call_inliner",
-        "//xla/service:collective_ops_utils",
         "//xla/service:collective_pipeliner",
         "//xla/service:collective_pipeliner_utils",
         "//xla/service:collective_utils",
@@ -2056,8 +2023,6 @@ cc_library(
         "//xla/service:layout_assignment",
         "//xla/service:layout_normalization",
         "//xla/service:llvm_compiler",
-        "//xla/service:logical_buffer",
-        "//xla/service:loop_schedule_linearizer",
         "//xla/service:memory_annotations_hdr",
         "//xla/service:nullary_function_wrap_inliner",
         "//xla/service:scan_expander",
@@ -2083,10 +2048,8 @@ cc_library(
         "//xla/service/llvm_ir:error_handler",
         "//xla/service/llvm_ir:llvm_command_line_options",
         "//xla/service/llvm_ir:llvm_util",
-        "//xla/service/spmd:collective_permute_motion",
         "//xla/service/spmd:schedule_aware_collective_ops_cse",
         "//xla/service/spmd/shardy:shardy_xla_pass",
-        "//xla/stream_executor:device_address_allocator",
         "//xla/stream_executor:device_description",
         "//xla/stream_executor:device_description_proto_cc",
         "//xla/stream_executor:dnn",
@@ -2099,17 +2062,9 @@ cc_library(
         "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor/abi:executable_abi_version",
         "//xla/stream_executor/cuda:cuda_compute_capability",
-        "//xla/stream_executor/integrations:device_mem_allocator",
-        "//xla/stream_executor/platform:platform_object_registry",
-        "//xla/tsl/concurrency:future",
-        "//xla/tsl/lib/monitoring:counter",
         "//xla/tsl/platform:env",
-        "//xla/tsl/platform:errors",
         "//xla/tsl/platform:logging",
-        "//xla/tsl/platform:status",
-        "//xla/tsl/platform:statusor",
         "//xla/tsl/util:maybe_owning",
-        "//xla/util/split_proto:split_gpu_executable_writer",
         "//xla/util/split_proto:split_proto_reader",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base",
@@ -2128,34 +2083,20 @@ cc_library(
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/synchronization",
-        "@com_google_absl//absl/types:span",
-        "@com_google_absl//absl/types:variant",
         "@com_google_protobuf//:protobuf",
         "@llvm-project//llvm:AsmParser",
-        "@llvm-project//llvm:BitReader",
-        "@llvm-project//llvm:BitWriter",
         "@llvm-project//llvm:Core",
-        "@llvm-project//llvm:JITLink",
         "@llvm-project//llvm:Support",
-        "@llvm-project//llvm:TargetParser",
         "@llvm-project//llvm:TransformUtils",
-        "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
-        "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:Support",
         "@riegeli//riegeli/bytes:string_reader",
         "@riegeli//riegeli/bytes:string_writer",
         "@tsl//tsl/platform",
-        "@tsl//tsl/platform:casts",
-        "@tsl//tsl/platform:env",
-        "@tsl//tsl/platform:errors",
-        "@tsl//tsl/platform:logging",
         "@tsl//tsl/platform:numbers",
         "@tsl//tsl/platform:path",
         "@tsl//tsl/platform:platform_port",
         "@tsl//tsl/platform:protobuf",
-        "@tsl//tsl/platform:stacktrace",
-        "@tsl//tsl/platform:statusor",
         "@tsl//tsl/profiler/lib:scoped_annotation",
         "@tsl//tsl/profiler/lib:traceme",
     ] + xla_internal(["//xla/service/gpu/google:export_hlo"]),
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 62fdb262f45..82826a68593 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -2213,9 +2213,7 @@ absl::Status GpuCompiler::OptimizeHloPostLayoutAssignment(
           gpu_target_config.device_description, ShapeSizeBytesFunction(),
           alias_info, mlir_context,
           /*only_fuse_if_profitable=*/true,
-          /*use_experimental_tiling=*/
-          debug_options.xla_gpu_experimental_enable_tiling_propagation(),
-          thread_pool, &mlir_context_pool_);
+          /*use_experimental_tiling=*/true, thread_pool, &mlir_context_pool_);
     }
 
     pipeline.AddPass<ReductionDimensionGrouper>();
@@ -3126,9 +3124,7 @@ absl::Status GpuCompiler::RunPreSchedulingPasses(
     // Cost model analysis for compute.
     pipeline.AddPass<GpuCostModelStatsCollection>(
         gpu_device_info, cost_analysis_options, mlir_context,
-        module->config()
-            .debug_options()
-            .xla_gpu_experimental_enable_tiling_propagation(),
+        /*use_experimental_tiling=*/true,
         module->config()
             .debug_options()
             .xla_gpu_experimental_enable_same_shape_multi_output_fusion());
diff --git a/third_party/xla/xla/service/gpu/ir_emission_utils.cc b/third_party/xla/xla/service/gpu/ir_emission_utils.cc
index cc2418abf67..40762d1fc82 100644
--- a/third_party/xla/xla/service/gpu/ir_emission_utils.cc
+++ b/third_party/xla/xla/service/gpu/ir_emission_utils.cc
@@ -109,15 +109,6 @@ bool IsTritonSupportedRaggedDot(
     return false;
   }
 
-  // Tiling propagation must be enabled.  TritonBackend::IsSupported gates
-  // kRaggedDot fusions on this flag — the xtile pipeline needs
-  // TiledHloComputation::Tile to drive code generation.
-  if (!instr.GetModule()
-           ->config()
-           .debug_options()
-           .xla_gpu_experimental_enable_tiling_propagation()) {
-    return false;
-  }
 
   // The xtile emitter's PrimitiveTypeToMlirType conversion supports the
   // following element types for ragged-dot operands and output.
diff --git a/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc b/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc
index 82d3aa460d4..e9c2c769853 100644
--- a/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc
+++ b/third_party/xla/xla/service/gpu/model/triton_emitter_constraints_test.cc
@@ -59,6 +59,10 @@ using ::testing::HasSubstr;
 class TritonEmitterConstraintsTest : public HloHardwareIndependentTestBase {
  public:
   TritonEmitterConstraintsTest() = default;
+
+  void SetUp() override {
+    GTEST_SKIP() << "Legacy symbolic tile analysis is disallowed.";
+  }
   std::optional<SymbolicTileAnalysis> TryAnalyzeModule(
       HloModule* module, bool with_triton_emitter_specific_constraints = true) {
     EmitterSpecificConstraintsBuilder constraints_builder = nullptr;
diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index 8b6e735625f..6aa03f34c95 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -1883,7 +1883,9 @@ message DebugOptions {
   optional int64 xla_gpu_async_copy_min_bytes = 522;
 
   // If true, enable experimental tiling propagation.
-  optional bool xla_gpu_experimental_enable_tiling_propagation = 456;
+  // TODO(b/514293537): drop the flag.
+  optional bool xla_gpu_experimental_enable_tiling_propagation = 456
+      [deprecated = true];
 
   // Enables and controls expetimental cost model-driven tiling choices for GEMM
   // fusions.

From a9c5f6a5d3aa1c16cdfb9de5a4cf5af6350fcf17 Mon Sep 17 00:00:00 2001
From: David Dunleavy <ddunleavy@google.com>
Date: Thu, 15 Oct 2026 17:08:19 -0700
Subject: [PATCH 0913/1000] Move `tpu_compilation_result.proto` to `xla/tpu`

PiperOrigin-RevId: 1000490301
---
 tensorflow/core/protobuf/tpu/BUILD            |  9 ++--
 .../protobuf/tpu/compilation_result.proto     | 25 +----------
 third_party/xla/xla/tpu/protobuf/BUILD        | 19 +++++++++
 .../xla/tpu/protobuf/compilation_result.proto | 41 +++++++++++++++++++
 .../xla/xla/tpu/tpu_compile_interface.h       |  2 +-
 5 files changed, 68 insertions(+), 28 deletions(-)
 create mode 100644 third_party/xla/xla/tpu/protobuf/compilation_result.proto

diff --git a/tensorflow/core/protobuf/tpu/BUILD b/tensorflow/core/protobuf/tpu/BUILD
index 08bf50b04d7..8d3f7036d75 100644
--- a/tensorflow/core/protobuf/tpu/BUILD
+++ b/tensorflow/core/protobuf/tpu/BUILD
@@ -16,7 +16,6 @@
 load("@rules_python//python:proto.bzl", "py_proto_library")
 load(
     "//tensorflow/core/platform:build_config.bzl",
-    "tf_additional_all_protos",
     "tf_proto_library",
     "tf_pyclif_proto_library",
 )
@@ -75,11 +74,13 @@ tf_proto_library(
     srcs = [
         "compilation_result.proto",
     ],
-    protodeps = tf_additional_all_protos() + [
-        "@xla//xla:xla_data_proto",
-        "@xla//xla/service:hlo_proto",
+    protodeps = [
+        "@xla//xla/tpu/protobuf:compilation_result_proto",
     ],
     visibility = ["//visibility:public"],
+    exports = [
+        "@xla//xla/tpu/protobuf:compilation_result_proto",
+    ],
 )
 
 tf_proto_library(
diff --git a/tensorflow/core/protobuf/tpu/compilation_result.proto b/tensorflow/core/protobuf/tpu/compilation_result.proto
index 1de6b5a0399..f2a3e600e00 100644
--- a/tensorflow/core/protobuf/tpu/compilation_result.proto
+++ b/tensorflow/core/protobuf/tpu/compilation_result.proto
@@ -15,27 +15,6 @@
 
 syntax = "proto3";
 
-package tensorflow.tpu;
+package tensorflow.tpu.dummy;
 
-import "xla/service/hlo.proto";
-import "tensorflow/core/protobuf/error_codes.proto";
-
-option cc_enable_arenas = true;
-
-// Describes the result of a TPU compilation. This is also used as TPU
-// compilation result status payload.
-// URI: "type.googleapis.com/tensorflow.tpu.CompilationResultProto"
-message CompilationResultProto {
-  // The error message, if any, returned during compilation.
-  error.Code status_code = 1;
-  string status_error_message = 2;
-
-  // HLO proto.
-  repeated xla.HloProto hlo_protos = 3;
-
-  enum ErrorCode {
-    UNKNOWN = 0;
-    OUT_OF_MEMORY = 1;
-  }
-  ErrorCode error_code = 4;
-}
+import public "xla/tpu/protobuf/compilation_result.proto";
diff --git a/third_party/xla/xla/tpu/protobuf/BUILD b/third_party/xla/xla/tpu/protobuf/BUILD
index 1a62300e115..97c445d2d80 100644
--- a/third_party/xla/xla/tpu/protobuf/BUILD
+++ b/third_party/xla/xla/tpu/protobuf/BUILD
@@ -49,6 +49,19 @@ tf_proto_library(
     deps = ["@com_google_protobuf//:wrappers_proto"],
 )
 
+tf_proto_library(
+    name = "compilation_result_proto",
+    srcs = [
+        "compilation_result.proto",
+    ],
+    protodeps = [
+        "//xla:xla_data_proto",
+        "//xla/service:hlo_proto",
+        "//xla/tsl/protobuf:error_codes_proto_impl",
+    ],
+    visibility = ["//visibility:public"],
+)
+
 cc_library(
     name = "tpu_embedding_optimization_parameters_utils",
     srcs = ["tpu_embedding_optimization_parameters_utils.cc"],
@@ -94,4 +107,10 @@ cc_library(
 #     visibility = ["//visibility:public"],
 #     deps = [":optimization_parameters_proto"],
 # )
+#
+# py_proto_library(
+#     name = "compilation_result_py_pb2",
+#     visibility = ["//visibility:public"],
+#     deps = [":compilation_result_proto"],
+# )
 # copybara:uncomment_end
diff --git a/third_party/xla/xla/tpu/protobuf/compilation_result.proto b/third_party/xla/xla/tpu/protobuf/compilation_result.proto
new file mode 100644
index 00000000000..dad626ef376
--- /dev/null
+++ b/third_party/xla/xla/tpu/protobuf/compilation_result.proto
@@ -0,0 +1,41 @@
+// Copyright 2026 The OpenXLA 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.
+// ==============================================================================
+
+syntax = "proto3";
+
+package tensorflow.tpu;
+
+import "xla/service/hlo.proto";
+import "xla/tsl/protobuf/error_codes.proto";
+
+option cc_enable_arenas = true;
+
+// Describes the result of a TPU compilation. This is also used as TPU
+// compilation result status payload.
+// URI: "type.googleapis.com/tensorflow.tpu.CompilationResultProto"
+message CompilationResultProto {
+  // The error message, if any, returned during compilation.
+  error.Code status_code = 1;
+  string status_error_message = 2;
+
+  // HLO proto.
+  repeated xla.HloProto hlo_protos = 3;
+
+  enum ErrorCode {
+    UNKNOWN = 0;
+    OUT_OF_MEMORY = 1;
+  }
+  ErrorCode error_code = 4;
+}
diff --git a/third_party/xla/xla/tpu/tpu_compile_interface.h b/third_party/xla/xla/tpu/tpu_compile_interface.h
index 892ff49017e..ade1f07ec82 100644
--- a/third_party/xla/xla/tpu/tpu_compile_interface.h
+++ b/third_party/xla/xla/tpu/tpu_compile_interface.h
@@ -32,7 +32,7 @@ class TpuCompileInterface {
   virtual uint64_t FingerprintString(absl::string_view str) = 0;
 
   // Proto: tensorflow::tpu::CompilationResultProto
-  // Location: tensorflow/core/protobuf/tpu/compilation_result.proto
+  // Location: xla/tpu/protobuf/compilation_result.proto
   static inline constexpr char kTpuCompileErrorPayloadKey[] =
       "type.googleapis.com/tensorflow.tpu.CompilationResultProto";
 

From 06d718aa259522e3af78651fc6b0ba9474c80ce5 Mon Sep 17 00:00:00 2001
From: Michael Kuperstein <mkuper@google.com>
Date: Thu, 15 Oct 2026 17:08:32 -0700
Subject: [PATCH 0914/1000] Flatten tuple-shaped outputs when extracting HLO
 modules in hlo_decomposer.

`ExtractCollectiveOperationsIntoNewModule` and `ExtractProducerConsumersIntoNewModule` combine multiple extracted output instructions into a single `kTuple` root instruction. When any of the extracted instructions already produces a tuple (such as a multi-output collective or a multi-output consumer fusion), wrapping it directly in `HloInstruction::CreateTuple` creates a nested tuple as the root of the extracted module.

Simplify the root output shape of extracted modules by decomposing tuple-shaped instructions via `GetTupleElement` instructions before assembling the root `kTuple` so that extracted modules always return a flat tuple of leaf elements.

PiperOrigin-RevId: 1000490416
---
 third_party/xla/xla/tools/hlo_decomposer.cc   | 33 ++++++-
 .../xla/xla/tools/hlo_decomposer_test.cc      | 92 +++++++++++++++++++
 2 files changed, 121 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/tools/hlo_decomposer.cc b/third_party/xla/xla/tools/hlo_decomposer.cc
index 1412140480a..c5c7caddbbd 100644
--- a/third_party/xla/xla/tools/hlo_decomposer.cc
+++ b/third_party/xla/xla/tools/hlo_decomposer.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/tools/hlo_decomposer.h"
 
+#include <cstdint>
 #include <memory>
 #include <string>
 #include <utility>
@@ -37,6 +38,32 @@ limitations under the License.
 namespace xla {
 namespace {
 
+// Appends the leaf elements of `instruction` to `elements`, decomposing any
+// tuple-shaped instruction via GetTupleElement instructions so that a flat
+// tuple can be constructed without nesting tuples.
+void AppendFlattenedTupleElements(HloInstruction* instruction,
+                                  HloComputation::Builder& builder,
+                                  std::vector<HloInstruction*>& elements) {
+  if (!instruction->shape().IsTuple()) {
+    elements.push_back(instruction);
+    return;
+  }
+  for (int64_t i = 0; i < instruction->shape().tuple_shapes().size(); ++i) {
+    HloInstruction* gte = builder.AddInstruction(
+        HloInstruction::CreateGetTupleElement(instruction, i));
+    AppendFlattenedTupleElements(gte, builder, elements);
+  }
+}
+
+void AddFlatTupleRoot(const std::vector<HloInstruction*>& instructions,
+                      HloComputation::Builder& builder) {
+  std::vector<HloInstruction*> tuple_elements;
+  for (HloInstruction* instruction : instructions) {
+    AppendFlattenedTupleElements(instruction, builder, tuple_elements);
+  }
+  builder.AddInstruction(HloInstruction::CreateTuple(tuple_elements));
+}
+
 // Returns whether it makes sense to run the given instruction in isolation
 // (e.g. whether it can run without dependent instructions).
 bool ShouldIsolateOpcode(HloOpcode opcode) {
@@ -167,9 +194,7 @@ std::unique_ptr<HloModule> ExtractCollectiveOperationsIntoNewModule(
   }
 
   if (return_tuple) {
-    std::unique_ptr<HloInstruction> tuple_instruction =
-        HloInstruction::CreateTuple(result_instructions);
-    builder.AddInstruction(std::move(tuple_instruction));
+    AddFlatTupleRoot(result_instructions, builder);
   }
   HloComputation* entry_computation =
       new_hlo_module->AddEntryComputationWithLayouts(builder.Build());
@@ -255,7 +280,7 @@ std::unique_ptr<HloModule> ExtractProducerConsumersIntoNewModule(
   }
 
   if (consumer_outputs.size() > 1) {
-    builder.AddInstruction(HloInstruction::CreateTuple(consumer_outputs));
+    AddFlatTupleRoot(consumer_outputs, builder);
   }
 
   HloComputation* entry_computation =
diff --git a/third_party/xla/xla/tools/hlo_decomposer_test.cc b/third_party/xla/xla/tools/hlo_decomposer_test.cc
index 1ed819c11bf..67f325ce2e8 100644
--- a/third_party/xla/xla/tools/hlo_decomposer_test.cc
+++ b/third_party/xla/xla/tools/hlo_decomposer_test.cc
@@ -23,6 +23,7 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/testlib/filecheck.h"
 #include "xla/tests/hlo_pjrt_test_base.h"
 #include "xla/tsl/platform/statusor.h"
@@ -300,5 +301,96 @@ CHECK-THEN: ROOT %add.1 = f32[10,10]{1,0} add(%dot.1, %parameter.2)
   EXPECT_TRUE(filecheck_result);
 }
 
+TEST_F(HloDecomposerTest,
+       ExtractCollectiveOperationsIntoNewModuleFlattensTupleResults) {
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          ParseAndReturnVerifiedModule(R"(
+HloModule module
+
+add {
+  lhs = f32[] parameter(0)
+  rhs = f32[] parameter(1)
+  ROOT add = f32[] add(lhs, rhs)
+}
+
+ENTRY main {
+  p0 = f32[10] parameter(0)
+  p1 = f32[10] parameter(1)
+  ar = (f32[10], f32[10]) all-reduce(p0, p1), replica_groups={{0,1}}, to_apply=add
+  ag = f32[20] all-gather(p0), dimensions={0}, replica_groups={{0,1}}
+  gte0 = f32[10] get-tuple-element(ar), index=0
+  gte1 = f32[10] get-tuple-element(ar), index=1
+  ROOT result = (f32[10], f32[10], f32[20]) tuple(gte0, gte1, ag)
+})"));
+  HloInstruction* ar =
+      module->entry_computation()->GetInstructionWithName("ar");
+  HloInstruction* ag =
+      module->entry_computation()->GetInstructionWithName("ag");
+  std::unique_ptr<HloModule> new_module =
+      ExtractCollectiveOperationsIntoNewModule(
+          {ar, ag}, /*done_ops=*/{},
+          /*non_optimized_ops=*/{HloOpcode::kAllReduce, HloOpcode::kAllGather},
+          /*return_tuple=*/true);
+  TF_ASSERT_OK_AND_ASSIGN(bool filecheck_result,
+                          RunFileCheck(new_module->ToString(), R"(
+CHECK: ENTRY
+CHECK-THEN: %[[P0:.*]] = f32[10]{0} parameter(0)
+CHECK-THEN: %[[P1:.*]] = f32[10]{0} parameter(1)
+CHECK-THEN: %[[AR:.*]] = (f32[10]{0}, f32[10]{0}) all-reduce(%[[P0]], %[[P1]])
+CHECK-THEN: %[[P2:.*]] = f32[10]{0} parameter(2)
+CHECK-THEN: %[[AG:.*]] = f32[20]{0} all-gather(%[[P2]])
+CHECK-THEN: %[[GTE0:.*]] = f32[10]{0} get-tuple-element(%[[AR]]), index=0
+CHECK-THEN: %[[GTE1:.*]] = f32[10]{0} get-tuple-element(%[[AR]]), index=1
+CHECK-THEN: ROOT %{{.*}} = (f32[10]{0}, f32[10]{0}, f32[20]{0}) tuple(%[[GTE0]], %[[GTE1]], %[[AG]])
+  )"));
+  EXPECT_TRUE(filecheck_result);
+}
+
+TEST_F(HloDecomposerTest,
+       ExtractProducerConsumersIntoNewModuleFlattensTupleConsumers) {
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          ParseAndReturnVerifiedModule(R"(
+HloModule module
+
+fused_comp {
+  p0 = f32[10,10] parameter(0)
+  p1 = f32[10,10] parameter(1)
+  add = f32[10,10] add(p0, p1)
+  mul = f32[10,10] multiply(p0, p1)
+  ROOT tuple = (f32[10,10], f32[10,10]) tuple(add, mul)
+}
+
+ENTRY main {
+  p0 = f32[10,10] parameter(0)
+  p1 = f32[10,10] parameter(1)
+  dot = f32[10,10] dot(p0, p1), lhs_contracting_dims={1}, rhs_contracting_dims={0}
+  p2 = f32[10,10] parameter(2)
+  fusion = (f32[10,10], f32[10,10]) fusion(dot, p2), kind=kLoop, calls=fused_comp
+  subtract = f32[10,10] subtract(dot, p2)
+  gte0 = f32[10,10] get-tuple-element(fusion), index=0
+  gte1 = f32[10,10] get-tuple-element(fusion), index=1
+  ROOT r = (f32[10,10], f32[10,10], f32[10,10]) tuple(gte0, gte1, subtract)
+})"));
+  HloInstruction* dot =
+      module->entry_computation()->GetInstructionWithName("dot");
+  std::unique_ptr<HloModule> new_module =
+      ExtractProducerConsumersIntoNewModule(*dot);
+  TF_ASSERT_OK_AND_ASSIGN(bool filecheck_result,
+                          RunFileCheck(new_module->ToString(), R"(
+CHECK: ENTRY
+CHECK-THEN: %[[P0:.*]] = f32[10,10]{1,0} parameter(0)
+CHECK-THEN: %[[P1:.*]] = f32[10,10]{1,0} parameter(1)
+CHECK-THEN: %[[DOT:.*]] = f32[10,10]{1,0} dot(%[[P0]], %[[P1]])
+CHECK-THEN: %[[P2:.*]] = f32[10,10]{1,0} parameter(2)
+CHECK-THEN: %[[FUSION:.*]] = (f32[10,10]{1,0}, f32[10,10]{1,0}) fusion(%[[DOT]], %[[P2]])
+CHECK-THEN: %[[P3:.*]] = f32[10,10]{1,0} parameter(3)
+CHECK-THEN: %[[SUB:.*]] = f32[10,10]{1,0} subtract(%[[DOT]], %[[P3]])
+CHECK-THEN: %[[GTE0:.*]] = f32[10,10]{1,0} get-tuple-element(%[[FUSION]]), index=0
+CHECK-THEN: %[[GTE1:.*]] = f32[10,10]{1,0} get-tuple-element(%[[FUSION]]), index=1
+CHECK-THEN: ROOT %{{.*}} = (f32[10,10]{1,0}, f32[10,10]{1,0}, f32[10,10]{1,0}) tuple(%[[GTE0]], %[[GTE1]], %[[SUB]])
+  )"));
+  EXPECT_TRUE(filecheck_result);
+}
+
 }  // namespace
 }  // namespace xla

From 4e83e5a34d8bf41ce25f30e668f2a7465c342233 Mon Sep 17 00:00:00 2001
From: Jian Cai <jiancai@google.com>
Date: Thu, 15 Oct 2026 17:18:03 -0700
Subject: [PATCH 0915/1000] Preserve HLO original values across while-loop
 concat code motion.

When `WhileLoopConcatCodeMotion` groups and concatenates loop-carried tuple elements, shapes of the loop tuple instructions and the first instruction of each `ConcatGroup` in the loop body are mutated in place, while slices of the concatenated result are sunk after the loop. Previously, the pass did not update `OriginalValue` metadata, causing shape-incompatible or stale original values on concatenated instructions and losing the original values of the post-loop outputs.

Update `RewriteLoopWithConcatGroups` to preserve original values for the post-loop output tuple and sunk slices, retain ungrouped tuple element origins on the loop tuple instructions while clearing concatenated elements, propagate original values to slices created for non-grouped consumers inside the loop body, and clear original values on in-place expanded body instructions so they do not overwrite their operands during replacement.

PiperOrigin-RevId: 1000495406
---
 .../service/while_loop_concat_code_motion.cc  |  52 ++++-
 .../while_loop_concat_code_motion_test.cc     | 207 ++++++++++++++++++
 2 files changed, 254 insertions(+), 5 deletions(-)

diff --git a/third_party/xla/xla/service/while_loop_concat_code_motion.cc b/third_party/xla/xla/service/while_loop_concat_code_motion.cc
index 8d655dbe7ce..9efe4bf58a4 100644
--- a/third_party/xla/xla/service/while_loop_concat_code_motion.cc
+++ b/third_party/xla/xla/service/while_loop_concat_code_motion.cc
@@ -17,6 +17,7 @@ limitations under the License.
 
 #include <cstdint>
 #include <map>
+#include <memory>
 #include <optional>
 #include <set>
 #include <utility>
@@ -38,6 +39,8 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/hlo/ir/hlo_original_value.h"
+#include "xla/hlo/ir/hlo_original_value_util.h"
 #include "xla/hlo/pass/hlo_pass_pipeline.h"
 #include "xla/hlo/transforms/simplifiers/hlo_dce.h"
 #include "xla/hlo/transforms/simplifiers/tuple_simplifier.h"
@@ -731,17 +734,35 @@ absl::Status RewriteLoopWithConcatGroups(
   auto param = body->parameter_instruction(0);
   auto cond_param = loop->while_condition()->parameter_instruction(0);
 
+  std::shared_ptr<OriginalValue> old_loop_ov = loop->original_value();
+  absl::flat_hash_map<const HloInstruction*, std::shared_ptr<OriginalValue>>
+      old_element_original_values;
+  for (const ConcatGroup& group : groups.Groups()) {
+    for (HloInstruction* element : group.elements) {
+      if (std::shared_ptr<OriginalValue> ov = element->original_value()) {
+        old_element_original_values[element] = std::move(ov);
+      }
+    }
+    group.elements[0]->set_original_value(nullptr);
+  }
+
   // First, modify loop signature and operands/users.
+  HloInstruction* old_while_init = loop->mutable_operand(0);
   std::vector<HloInstruction*> init_elements(
       loop->shape().tuple_shapes().size());
   for (int64_t i = 0; i < param_gtes.size(); ++i) {
     init_elements[i] =
         loop->parent()->AddInstruction(HloInstruction::CreateGetTupleElement(
-            loop->shape().tuple_shapes(i), loop->mutable_operand(0), i));
+            loop->shape().tuple_shapes(i), old_while_init, i));
   }
+  std::vector<int64_t> ungrouped_tuple_indices(param_gtes.size(), -1);
   for (int64_t i = 0; i < param_gtes.size(); ++i) {
     const auto& group_and_index = groups.GetGroupIndex(param_gtes[i]);
-    if (!group_and_index.has_value() || group_and_index->second != 0) {
+    if (!group_and_index.has_value()) {
+      ungrouped_tuple_indices[i] = i;
+      continue;
+    }
+    if (group_and_index->second != 0) {
       continue;
     }
     const auto& group = groups.GetGroup(group_and_index->first);
@@ -762,9 +783,15 @@ absl::Status RewriteLoopWithConcatGroups(
     init_elements[i] =
         group.CreateConcat(std::move(input_concat_elements), loop->parent());
   }
-  ABSL_RETURN_IF_ERROR(loop->ReplaceOperandWithDifferentShape(
-      0, loop->parent()->AddInstruction(
-             HloInstruction::CreateTuple(init_elements))));
+  HloInstruction* new_init_tuple = loop->parent()->AddInstruction(
+      HloInstruction::CreateTuple(init_elements));
+  CopyOriginalValue(old_while_init, new_init_tuple, ungrouped_tuple_indices);
+  CopyOriginalValue(loop, loop, ungrouped_tuple_indices);
+  CopyOriginalValue(param, param, ungrouped_tuple_indices);
+  CopyOriginalValue(cond_param, cond_param, ungrouped_tuple_indices);
+  CopyOriginalValue(body->root_instruction(), body->root_instruction(),
+                    ungrouped_tuple_indices);
+  ABSL_RETURN_IF_ERROR(loop->ReplaceOperandWithDifferentShape(0, new_init_tuple));
   // Adjust loop users.
   auto original_loop_users = loop->users();
   const bool loop_is_root = loop == loop->parent()->root_instruction();
@@ -790,6 +817,17 @@ absl::Status RewriteLoopWithConcatGroups(
   }
   auto new_output_tuple = loop->parent()->AddInstruction(
       HloInstruction::CreateTuple(output_elements));
+  if (old_loop_ov != nullptr && !old_loop_ov->is_synthetic_call()) {
+    for (int64_t i = 0; i < output_elements.size(); ++i) {
+      ABSL_ASSIGN_OR_RETURN(auto subtree, old_loop_ov->tree().Subtree({i}));
+      auto elem_ov = std::make_shared<OriginalValue>(std::move(subtree));
+      if (!elem_ov->IsEmpty()) {
+        output_elements[i]->set_original_value(std::move(elem_ov));
+      }
+    }
+    new_output_tuple->set_original_value(
+        std::make_shared<OriginalValue>(old_loop_ov->tree()));
+  }
   for (auto user : original_loop_users) {
     ABSL_RETURN_IF_ERROR(loop->ReplaceUseWithDifferentShape(user, new_output_tuple));
   }
@@ -945,6 +983,10 @@ absl::Status RewriteLoopWithConcatGroups(
         const auto& operand_group = groups.GetGroup(operand_group_index->first);
         auto slice = operand_group.CreateSlice(
             operand_group.elements[0], operand_group_index->second, body);
+        if (auto it = old_element_original_values.find(operand);
+            it != old_element_original_values.end()) {
+          slice->set_original_value(it->second);
+        }
         ABSL_RETURN_IF_ERROR(hlo->ReplaceOperandWithDifferentShape(i, slice));
       }
     }
diff --git a/third_party/xla/xla/service/while_loop_concat_code_motion_test.cc b/third_party/xla/xla/service/while_loop_concat_code_motion_test.cc
index 67f1a3c1cf8..22c514d7b28 100644
--- a/third_party/xla/xla/service/while_loop_concat_code_motion_test.cc
+++ b/third_party/xla/xla/service/while_loop_concat_code_motion_test.cc
@@ -20,6 +20,7 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/log/log.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_instruction.h"
@@ -538,5 +539,211 @@ TEST_F(WhileLoopConcatCodeMotionTest, NonElementwiseOps) {
                         op::Slice(op::GetTupleElement(loop))));
 }
 
+TEST_F(WhileLoopConcatCodeMotionTest, SimpleMotionWithOriginalValue) {
+  constexpr absl::string_view kHloModule = R"(
+    HloModule test
+
+    %cond {
+      %param = (s32[], f32[1024,1024], f32[1024,1024]) parameter(0), origin={({"cp0"}, {"cp1"}, {"cp2"})}
+      %gte.0 = s32[] get-tuple-element(%param), index=0, origin={{"cond_gte.0"}}
+      %constant = s32[] constant(5)
+      ROOT result = pred[] compare(%gte.0, %constant), direction=LT
+    }
+
+    %body {
+      %param = (s32[], f32[1024,1024], f32[1024,1024]) parameter(0), origin={({"bp0"}, {"bp1"}, {"bp2"})}
+      %gte.0 = s32[] get-tuple-element(%param), index=0, origin={{"gte.0"}}
+      %gte.1 = f32[1024,1024] get-tuple-element(%param), index=1, origin={{"gte.1"}}
+      %gte.2 = f32[1024,1024] get-tuple-element(%param), index=2, origin={{"gte.2"}}
+      %concat = f32[2048,1024] concatenate(%gte.1, %gte.2), dimensions={0}, origin={{"concat"}}
+      %ccall = f32[2048,1024] custom-call(%concat), custom_call_target="test", origin={{"ccall"}}
+      %slice.0 = f32[1024,1024] slice(%ccall), slice={[0:1024], [0:1024]}, origin={{"slice.0"}}
+      %slice.1 = f32[1024,1024] slice(%ccall), slice={[1024:2048], [0:1024]}, origin={{"slice.1"}}
+      %ccall2 = f32[1024,1024] custom-call(), custom_call_target="test2", origin={{"ccall2"}}
+      %add.0 = f32[1024,1024] add(%slice.0, %ccall2), origin={{"add.0"}}
+      %add.1 = f32[1024,1024] add(%slice.1, %ccall2), origin={{"add.1"}}
+      %t0 = token[] after-all()
+      %outfeed = token[] outfeed(%slice.1, %t0)
+      %constant = s32[] constant(1)
+      %increment_iteration = s32[] add(s32[] %gte.0, s32[] %constant), origin={{"inc"}}
+      ROOT %loop_result = (s32[], f32[1024,1024], f32[1024,1024])
+        tuple(%increment_iteration, %add.0, %add.1), origin={({"br0"}, {"br1"}, {"br2"})}
+    }
+
+    ENTRY test_main {
+      %param.0 = f32[1024,1024] parameter(0), origin={{"p0"}}
+      %param.1 = f32[1024,1024] parameter(1), origin={{"p1"}}
+      %constant.0 = s32[] constant(0), origin={{"c0"}}
+      %while_init = (s32[], f32[1024,1024], f32[1024,1024]) tuple(%constant.0, %param.0, %param.1), origin={({"c0"}, {"p0"}, {"p1"})}
+      ROOT %while = (s32[], f32[1024,1024], f32[1024,1024]) while(%while_init), condition=%cond, body=%body, origin={({"w0"}, {"w1"}, {"w2"}),["while#$"]}
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kHloModule));
+  ASSERT_OK_AND_ASSIGN(bool changed,
+                       WhileLoopConcatCodeMotion(2).Run(module.get()));
+  ASSERT_TRUE(changed);
+  ASSERT_OK(HloVerifier(/*layout_sensitive=*/false,
+                        /*allow_mixed_precision=*/false)
+                .Run(module.get())
+                .status());
+
+  const HloInstruction* entry_root =
+      module->entry_computation()->root_instruction();
+  ASSERT_NE(entry_root->original_value(), nullptr);
+  EXPECT_EQ(entry_root->original_value()->ToString(),
+            R"(({"w0"}, {"w1"}, {"w2"}))");
+
+  const HloInstruction* out_gte0 = entry_root->operand(0);
+  ASSERT_NE(out_gte0->original_value(), nullptr);
+  EXPECT_EQ(out_gte0->original_value()->ToString(), R"({"w0"})");
+
+  const HloInstruction* out_slice0 = entry_root->operand(1);
+  ASSERT_NE(out_slice0->original_value(), nullptr);
+  EXPECT_EQ(out_slice0->original_value()->ToString(), R"({"w1"})");
+  EXPECT_EQ(out_slice0->operand(0)->original_value(), nullptr);
+
+  const HloInstruction* out_slice1 = entry_root->operand(2);
+  ASSERT_NE(out_slice1->original_value(), nullptr);
+  EXPECT_EQ(out_slice1->original_value()->ToString(), R"({"w2"})");
+
+  const HloInstruction* while_op = out_gte0->operand(0);
+  ASSERT_NE(while_op->original_value(), nullptr);
+  EXPECT_EQ(while_op->original_value()->ToString(),
+            R"(({"w0"}, {}),["while#$"])");
+
+  const HloInstruction* while_init = while_op->while_init();
+  ASSERT_NE(while_init->original_value(), nullptr);
+  EXPECT_EQ(while_init->original_value()->ToString(), R"(({"c0"}, {}))");
+
+  const HloComputation* body = while_op->while_body();
+  ASSERT_NE(body->parameter_instruction(0)->original_value(), nullptr);
+  EXPECT_EQ(body->parameter_instruction(0)->original_value()->ToString(),
+            R"(({"bp0"}, {}))");
+
+  const HloComputation* cond = while_op->while_condition();
+  ASSERT_NE(cond->parameter_instruction(0)->original_value(), nullptr);
+  EXPECT_EQ(cond->parameter_instruction(0)->original_value()->ToString(),
+            R"(({"cp0"}, {}))");
+
+  const HloInstruction* body_root = body->root_instruction();
+  ASSERT_NE(body_root->original_value(), nullptr);
+  EXPECT_EQ(body_root->original_value()->ToString(), R"(({"br0"}, {}))");
+
+  const HloInstruction* inc = body_root->operand(0);
+  ASSERT_NE(inc->original_value(), nullptr);
+  EXPECT_EQ(inc->original_value()->ToString(), R"({"inc"})");
+
+  const HloInstruction* combined_add = body_root->operand(1);
+  EXPECT_EQ(combined_add->original_value(), nullptr);
+
+  const HloInstruction* ccall = combined_add->operand(0);
+  ASSERT_NE(ccall->original_value(), nullptr);
+  EXPECT_EQ(ccall->original_value()->ToString(), R"({"ccall"})");
+
+  const HloInstruction* outfeed =
+      FindInstruction(module.get(), HloOpcode::kOutfeed);
+  ASSERT_NE(outfeed, nullptr);
+  ASSERT_NE(outfeed->operand(0)->original_value(), nullptr);
+  EXPECT_EQ(outfeed->operand(0)->original_value()->ToString(),
+            R"({"slice.1"})");
+}
+
+TEST_F(WhileLoopConcatCodeMotionTest, NonElementwiseOpsWithOriginalValue) {
+  constexpr absl::string_view kHloModule = R"(
+    HloModule test
+
+    %cond {
+      %param = (s32[], f32[1024,1024], f32[1024,1024], f32[1024], f32[1024], f32[1], f32[1]) parameter(0)
+      %gte.0 = s32[] get-tuple-element(%param), index=0
+      %constant = s32[] constant(5)
+      ROOT result = pred[] compare(%gte.0, %constant), direction=LT
+    }
+
+    %sum {
+      %a = f32[] parameter(0)
+      %b = f32[] parameter(1)
+      ROOT %add = f32[] add(%a, %b)
+    }
+
+    %body {
+      %param = (s32[], f32[1024,1024], f32[1024,1024], f32[1024], f32[1024], f32[1], f32[1]) parameter(0)
+      %gte.0 = s32[] get-tuple-element(%param), index=0
+      %gte.1 = f32[1024,1024] get-tuple-element(%param), index=1
+      %gte.2 = f32[1024,1024] get-tuple-element(%param), index=2
+      %reshape.0 = f32[1,1024,1024] reshape(%gte.1)
+      %reshape.1 = f32[1,1024,1024] reshape(%gte.2)
+      %concat = f32[2,1024,1024] concatenate(%reshape.0, %reshape.1), dimensions={0}
+      %ccall = f32[2,1024,1024] custom-call(%concat), custom_call_target="test"
+      %slice.0 = f32[1,1024,1024] slice(%ccall), slice={[0:1], [0:1024], [0:1024]}
+      %slice.1 = f32[1,1024,1024] slice(%ccall), slice={[1:2], [0:1024], [0:1024]}
+      %reshape.2 = f32[1024,1024] reshape(%slice.0)
+      %reshape.3 = f32[1024,1024] reshape(%slice.1)
+      %gte.3 = f32[1024] get-tuple-element(%param), index=3
+      %gte.4 = f32[1024] get-tuple-element(%param), index=4
+      %constant.0 = f32[] constant(0)
+      %reduce.0 = f32[1024] reduce(%reshape.0, %constant.0), to_apply=%sum, dimensions={0,1}
+      %reduce.1 = f32[1024] reduce(%reshape.1, %constant.0), to_apply=%sum, dimensions={0,1}
+      %add.0 = f32[1024] add(%reduce.0, %gte.3)
+      %add.1 = f32[1024] add(%reduce.1, %gte.4)
+      %br0 = f32[1024,1024] broadcast(%add.0), dimensions={1}
+      %br1 = f32[1024,1024] broadcast(%add.1), dimensions={1}
+      %sub.0 = f32[1024,1024] subtract(%reshape.2, %br0)
+      %sub.1 = f32[1024,1024] subtract(%reshape.3, %br1)
+      %gte.5 = f32[1] get-tuple-element(%param), index=5
+      %gte.6 = f32[1] get-tuple-element(%param), index=6
+      %reshape.4 = f32[] reshape(%gte.5)
+      %reshape.5 = f32[] reshape(%gte.6)
+      %br2 = f32[1024] broadcast(%reshape.4), dimensions={}
+      %br3 = f32[1024] broadcast(%reshape.5), dimensions={}
+      %add.2 = f32[1024] add(%add.0, %br2)
+      %add.3 = f32[1024] add(%add.1, %br3)
+      %inc0 = f32[] add(%constant.0, %reshape.4)
+      %inc1 = f32[] add(%constant.0, %reshape.5)
+      %reshape.6 = f32[1] reshape(%inc0)
+      %reshape.7 = f32[1] reshape(%inc1)
+      %constant = s32[] constant(1)
+      %increment_iteration = s32[] add(s32[] %gte.0, s32[] %constant)
+      ROOT %loop_result = (s32[], f32[1024,1024], f32[1024,1024], f32[1024], f32[1024], f32[1], f32[1])
+        tuple(%increment_iteration, %sub.0, %sub.1, %add.2, %add.3, %reshape.6, %reshape.7)
+    }
+
+    ENTRY test_main {
+      %param.0 = f32[1024,1024] parameter(0)
+      %param.1 = f32[1024,1024] parameter(1)
+      %param.2 = f32[1024] parameter(2)
+      %param.3 = f32[1024] parameter(3)
+      %param.4 = f32[1] parameter(4)
+      %param.5 = f32[1] parameter(5)
+      %constant.0 = s32[] constant(0)
+      %while_init = (s32[], f32[1024,1024], f32[1024,1024], f32[1024], f32[1024], f32[1], f32[1])
+        tuple(%constant.0, %param.0, %param.1, %param.2, %param.3, %param.4, %param.5)
+      ROOT %while = (s32[], f32[1024,1024], f32[1024,1024], f32[1024], f32[1024], f32[1], f32[1])
+        while(%while_init), condition=%cond, body=%body,
+        origin={({"w0"}, {"w1"}, {"w2"}, {"w3"}, {"w4"}, {"w5"}, {"w6"})}
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kHloModule));
+  ASSERT_OK_AND_ASSIGN(bool changed,
+                       WhileLoopConcatCodeMotion(2).Run(module.get()));
+  ASSERT_TRUE(changed);
+  ASSERT_OK(HloVerifier(/*layout_sensitive=*/false,
+                        /*allow_mixed_precision=*/false)
+                .Run(module.get())
+                .status());
+
+  const HloInstruction* entry_root =
+      module->entry_computation()->root_instruction();
+  ASSERT_NE(entry_root->original_value(), nullptr);
+  EXPECT_EQ(entry_root->original_value()->ToString(),
+            R"(({"w0"}, {"w1"}, {"w2"}, {"w3"}, {"w4"}, {"w5"}, {"w6"}))");
+  for (int i = 0; i < 7; ++i) {
+    ASSERT_NE(entry_root->operand(i)->original_value(), nullptr);
+    EXPECT_EQ(entry_root->operand(i)->original_value()->ToString(),
+              absl::StrCat("{\"w", i, "\"}"));
+  }
+}
+
 }  // namespace
 }  // namespace xla

From 42d6aae315957bdb67839f5715377f49ff6a7b10 Mon Sep 17 00:00:00 2001
From: Matthias Guenther <mrguenther@google.com>
Date: Thu, 15 Oct 2026 17:27:08 -0700
Subject: [PATCH 0916/1000] Integrate StableHLO at openxla/stablehlo@e7876d86

PiperOrigin-RevId: 1000499806
---
 .../xla/third_party/stablehlo/temporary.patch | 17665 ----------------
 .../xla/third_party/stablehlo/workspace.bzl   |     4 +-
 2 files changed, 2 insertions(+), 17667 deletions(-)

diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch
index ecdcdce23c6..cd541efb2de 100644
--- a/third_party/xla/third_party/stablehlo/temporary.patch
+++ b/third_party/xla/third_party/stablehlo/temporary.patch
@@ -15,17668 +15,3 @@
 #
 # This file is automatically generated by the `generate_patch` tool.
 # Do not edit directly.
-diff --ruN a/stablehlo/.github/workflows/buildAndTestBazel.yml b/stablehlo/.github/workflows/buildAndTestBazel.yml
---- stablehlo/.github/workflows/buildAndTestBazel.yml
-+++ stablehlo/.github/workflows/buildAndTestBazel.yml
-@@ -45,7 +45,7 @@
-   bazel-build:
-     strategy:
-       fail-fast: false
--    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n2-16' || 'ubuntu-22.04' }}
-+    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n4-16' || 'ubuntu-22.04' }}
-     container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
- 
-     steps:
-diff --ruN a/stablehlo/.github/workflows/buildAndTestCMake.yml b/stablehlo/.github/workflows/buildAndTestCMake.yml
---- stablehlo/.github/workflows/buildAndTestCMake.yml
-+++ stablehlo/.github/workflows/buildAndTestCMake.yml
-@@ -44,7 +44,7 @@
-       STABLEHLO_PYTHON_BUILD_DIR: "stablehlo-python-build"
-     strategy:
-       fail-fast: false
--    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n2-16' || 'ubuntu-22.04' }}
-+    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n4-16' || 'ubuntu-22.04' }}
-     container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cmake:latest"
- 
-     steps:
-diff --ruN a/stablehlo/.github/workflows/publishWheelRelease.yml b/stablehlo/.github/workflows/publishWheelRelease.yml
---- stablehlo/.github/workflows/publishWheelRelease.yml
-+++ stablehlo/.github/workflows/publishWheelRelease.yml
-@@ -43,7 +43,7 @@
-       fail-fast: false
-       matrix:
-         python-version: [ "3.10", "3.11", "3.12" ]
--    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n2-16' ||  'ubuntu-22.04'  }}
-+    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n4-16' ||  'ubuntu-22.04'  }}
-     container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cmake:latest"
- 
-     steps:
-diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel
---- stablehlo/BUILD.bazel
-+++ stablehlo/BUILD.bazel
-@@ -1960,3 +1960,17 @@
-     name = "all_tests",
-     tests = [":stablehlo_ci_tests"],
- )
-+
-+# copybara:uncomment_begin
-+# alias(
-+#     name = "check_ops",
-+#     actual = "//stablehlo/tests:check_ops",
-+#     deprecation = "Use //stablehlo/tests:check_ops instead.",
-+# )
-+#
-+# alias(
-+#     name = "test_utils",
-+#     actual = "//stablehlo/tests:test_utils",
-+#     deprecation = "Use //stablehlo/tests:test_utils instead.",
-+# )
-+# copybara:uncomment_end
-diff --ruN a/stablehlo/WORKSPACE.bazel b/stablehlo/WORKSPACE.bazel
---- stablehlo/WORKSPACE.bazel
-+++ stablehlo/WORKSPACE.bazel
-@@ -56,9 +56,9 @@
- 
- http_archive(
-     name = "rules_python",
--    sha256 = "8964aa1e7525fea5244ba737458694a057ada1be96a92998a41caa1983562d00",
--    strip_prefix = "rules_python-1.8.5",
--    url = "https://github.com/bazel-contrib/rules_python/releases/download/1.8.5/rules_python-1.8.5.tar.gz",
-+    sha256 = "e11d2e1efce1589e5bdfa93986712c74fc7467a0f093143d489d2ef5ebb1ed2a",
-+    strip_prefix = "rules_python-2.2.0",
-+    url = "https://github.com/bazel-contrib/rules_python/releases/download/2.2.0/rules_python-2.2.0.tar.gz",
- )
- 
- load("@rules_python//python:repositories.bzl", "py_repositories", "python_register_toolchains")
-@@ -119,6 +119,18 @@
-     ],
- )
- 
-+# Required by the LLVM Bazel overlay (`@llvm-project//third-party:lzma`).
-+http_archive(
-+    name = "llvm_xz",
-+    build_file = "//third_party:xz.BUILD",
-+    sha256 = "3d3a1b973af218114f4f889bbaa2f4c037deaae0c8e815eec381c3d546b974a0",
-+    strip_prefix = "xz-5.8.3",
-+    urls = [
-+        "https://storage.googleapis.com/mirror.tensorflow.org/github.com/tukaani-project/xz/releases/download/v5.8.3/xz-5.8.3.tar.gz",
-+        "https://github.com/tukaani-project/xz/releases/download/v5.8.3/xz-5.8.3.tar.gz",
-+    ],
-+)
-+
- # These need to come in this specific order or else Bazel will complain about
- # missing/circular dependencies.
- load("//third_party/llvm:workspace.bzl", llvm = "repo")
-diff --ruN a/stablehlo/build_tools/lit_test_suite.bzl b/stablehlo/build_tools/lit_test_suite.bzl
---- stablehlo/build_tools/lit_test_suite.bzl
-+++ stablehlo/build_tools/lit_test_suite.bzl
-@@ -44,11 +44,13 @@
-     tests_in_suite = []
- 
-     for src in srcs:
--        test_name = ("%s.test" % src).replace("/", "_")
-+        test_name = ("%s.test" % src).replace("/", "_")  # copybara:comment
-+
-+        # copybara:uncomment test_name = "%s.test" % src
-         lit_test(
-             name = test_name,  # copybara:comment
-             # copybara:uncomment name = src,
--            srcs = [src],
-+            srcs = [src],  # copybara:comment
-             # copybara:uncomment src = src,
-             args = args,
-             data = data + native.glob(
-@@ -57,7 +59,7 @@
-                 allow_empty = True,
-             ),
-             deps = ["@rules_python//python/runfiles"],  # copybara:comment
--            # copybara:uncomment driver = "//third_party/llvm/llvm-project/mlir:run_lit.sh",
-+            # copybara:uncomment driver = "@llvm-project//mlir:run_lit.sh",
-             tags = tags,
-             size = size_per_test,
-         )
-diff --ruN a/stablehlo/docs/generated/StablehloBuilder.md b/stablehlo/docs/generated/StablehloBuilder.md
---- stablehlo/docs/generated/StablehloBuilder.md
-+++ stablehlo/docs/generated/StablehloBuilder.md
-@@ -424,6 +424,15 @@
- MlirOp Einsum(Type resultType, MlirOp &lhs, MlirOp &rhs, ::llvm::StringRef einsum_config);
- ```
- 
-+### `stablehlo::Exp2Op`
-+
-+Creates a new [`stablehlo.exp2`](https://openxla.org/stablehlo/spec#exp2)
-+operation.
-+
-+```c++
-+MlirOp Exp2(MlirOp &operand, /*optional*/::mlir::stablehlo::ResultAccuracyAttr result_accuracy = {});
-+```
-+
- ### `stablehlo::ExpOp`
- 
- Creates a new [`stablehlo.exponential`](https://openxla.org/stablehlo/spec#exponential)
-@@ -543,6 +552,15 @@
- MlirOp Log1p(MlirOp &operand, /*optional*/::mlir::stablehlo::ResultAccuracyAttr result_accuracy = {});
- ```
- 
-+### `stablehlo::Log2Op`
-+
-+Creates a new [`stablehlo.log2`](https://openxla.org/stablehlo/spec#log2)
-+operation.
-+
-+```c++
-+MlirOp Log2(MlirOp &operand, /*optional*/::mlir::stablehlo::ResultAccuracyAttr result_accuracy = {});
-+```
-+
- ### `stablehlo::LogOp`
- 
- Creates a new [`stablehlo.log`](https://openxla.org/stablehlo/spec#log)
-diff --ruN a/stablehlo/docs/generated/chlo.md b/stablehlo/docs/generated/chlo.md
---- stablehlo/docs/generated/chlo.md
-+++ stablehlo/docs/generated/chlo.md
-@@ -2068,7 +2068,7 @@
- 
- ```
- #chlo.comparison_type<
--  `NOTYPE` | `FLOAT` | `TOTALORDER` | `SIGNED` | `UNSIGNED`   # value
-+  `NOTYPE` | `FLOAT` | `TOTALORDER` | `SIGNED` | `UNSIGNED` | `WEAKORDER`   # value
- >
- ```
- 
-@@ -2141,6 +2141,7 @@
- | TOTALORDER | `2` | TOTALORDER |
- | SIGNED | `3` | SIGNED |
- | UNSIGNED | `4` | UNSIGNED |
-+| WEAKORDER | `5` | WEAKORDER |
- 
- ### Precision
- 
-diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
---- stablehlo/docs/spec.md
-+++ stablehlo/docs/spec.md
-@@ -2380,6 +2380,7 @@
- in order to be able to express windows of `lhs`:
- 
- <!-- markdownlint-disable line-length -->
-+
- * `lhs_window_dimensions = lhs_shape(dim(lhs, input_batch_dimension), dim(rhs, kernel_spatial_dimensions), dim(lhs, input_feature_dimension))`.
- * `lhs_window_strides = lhs_shape(1, window_strides, 1)`.
- * `lhs_padding = lhs_shape([0, 0], padding, [0, 0])`.
-@@ -2476,6 +2477,7 @@
- #### Constraints
- 
- <!-- markdownlint-disable line-length -->
-+
- * (C1) `N = rank(lhs) = rank(rhs)`.
- * (C2) `size(window_strides) = N - 2`.
- * (C3) `0 < window_strides`.
-@@ -2808,6 +2810,7 @@
- More formally, `result[result_index] = dot_product`, where:
- 
- <!-- markdownlint-disable line-length -->
-+
- * `lhs_result_dimensions = [d for d in axes(lhs) and d not in lhs_batching_dimensions and d not in lhs_contracting_dimensions]`.
- * `rhs_result_dimensions = [d for d in axes(rhs) and d not in rhs_batching_dimensions and d not in rhs_contracting_dimensions]`.
- * `result_batching_index + result_lhs_index + result_rhs_index = result_index`
-@@ -3148,6 +3151,7 @@
- #### Constraints
- 
- <!-- markdownlint-disable line-length -->
-+
- * (C1) `N = rank(lhs) = rank(rhs)`.
- * (C2) `size(window_strides) = N - 2`.
- * (C3) `0 < window_strides`.
-@@ -3626,6 +3630,49 @@
- 
- &nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/dynamic_update_slice.mlir)
- 
-+### exp2
-+
-+#### Semantics
-+
-+Performs element-wise base-2 exponential operation on `operand` tensor and
-+produces a `result` tensor. Depending on the element type, does the following:
-+
-+* For floats: `exp2` from IEEE-754.
-+* For complex numbers: complex base-2 exponential.
-+* For quantized types: `dequantize_op_quantize(exp2, operand, type(result))`.
-+
-+#### Inputs
-+
-+| Label | Name              | Type                 | Constraints |
-+| ----- | ----------------- | -------------------- | ----------- |
-+| (I1)  | `operand`         | tensor of            | (C1)        |
-+:       :                   : floating-point or    :             :
-+:       :                   : complex type or      :             :
-+:       :                   : per-tensor quantized :             :
-+:       :                   : tensor               :             :
-+| (I2)  | `result_accuracy` | optional             |             |
-+:       :                   : `ResultAccuracyAttr` :             :
-+:       :                   : (default `DEFAULT`)  :             :
-+
-+#### Outputs
-+
-+| Name     | Type                                        | Constraints |
-+| -------- | ------------------------------------------- | ----------- |
-+| `result` | tensor of floating-point or complex type or | (C1)        |
-+:          : per-tensor quantized tensor                 :             :
-+
-+#### Constraints
-+
-+* (C1) `baseline_type(operand) = baseline_type(result)`.
-+
-+#### Examples
-+
-+```mlir
-+// %operand: [[0.0, 1.0], [2.0, 3.0]]
-+%result = "stablehlo.exp2"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
-+// %result: [[1.0, 2.0], [4.0, 8.0]]
-+```
-+
- ### exponential
- 
- #### Semantics
-@@ -3798,6 +3845,7 @@
- * (C3) `1 <= size(fft_length) <= 3`.
- * (C4) If among `operand` and `result`, there is a tensor `real` of a
- floating-point type, then `shape(real)[-size(fft_length):] = fft_length`.
-+
- * (C5) `shape(result) = shape(operand)` except for:
-   * If `fft_type = RFFT`,
-     `dim(result, -1) = dim(operand, -1) = 0 ? 0 : dim(operand, -1) / 2 + 1`.
-@@ -3866,6 +3914,7 @@
- More formally, `result[result_index] = operand[operand_index]` where:
- 
- <!-- markdownlint-disable line-length -->
-+
- * `batch_dims = [d for d in axes(result) and d not in offset_dims]`.
- * `batch_index = result_index[batch_dims...]`.
- * `start_index` is defined as:
-@@ -4355,6 +4404,49 @@
- 
- &nbsp;[More Examples](https://github.com/openxla/stablehlo/tree/main/stablehlo/tests/interpret/log.mlir)
- 
-+### log2
-+
-+#### Semantics
-+
-+Performs element-wise base-2 logarithm operation on `operand` tensor and
-+produces a `result` tensor. Depending on the element type, does the following:
-+
-+* For floats: `log2` from IEEE-754.
-+* For complex numbers: complex base-2 logarithm.
-+* For quantized types: `dequantize_op_quantize(log2, operand, type(result))`.
-+
-+#### Inputs
-+
-+| Label | Name              | Type                 | Constraints |
-+| ----- | ----------------- | -------------------- | ----------- |
-+| (I1)  | `operand`         | tensor of            | (C1)        |
-+:       :                   : floating-point or    :             :
-+:       :                   : complex type or      :             :
-+:       :                   : per-tensor quantized :             :
-+:       :                   : tensor               :             :
-+| (I2)  | `result_accuracy` | optional             |             |
-+:       :                   : `ResultAccuracyAttr` :             :
-+:       :                   : (default `DEFAULT`)  :             :
-+
-+#### Outputs
-+
-+| Name     | Type                                        | Constraints |
-+| -------- | ------------------------------------------- | ----------- |
-+| `result` | tensor of floating-point or complex type or | (C1)        |
-+:          : per-tensor quantized tensor                 :             :
-+
-+#### Constraints
-+
-+* (C1) `baseline_type(operand) = baseline_type(result)`.
-+
-+#### Examples
-+
-+```mlir
-+// %operand: [[1.0, 2.0], [4.0, 8.0]]
-+%result = "stablehlo.log2"(%operand) : (tensor<2x2xf64>) -> tensor<2x2xf64>
-+// %result: [[0.0, 1.0], [2.0, 3.0]]
-+```
-+
- ### log_plus_one
- 
- #### Semantics
-@@ -5358,6 +5450,7 @@
- #### Constraints
- 
- <!-- markdownlint-disable line-length -->
-+
- * (C1) `0 < size(inputs) = size(init_values) = size(results) = N`.
- * (C2) `same(shape(inputs...))`.
- * (C3) `element_type(inputs...) = element_type(init_values...)`.
-@@ -5897,6 +5990,7 @@
- #### Constraints
- 
- <!-- markdownlint-disable line-length -->
-+
- * (C1) `same(shape(inputs...))`.
- * (C2) `rank(inputs[0]) = size(update_window_dims) + size(inserted_window_dims) + size(input_batching_dims)`.
- * (C3) `same(shape(updates...))`.
-@@ -6081,6 +6175,7 @@
-    ([#731](https://github.com/openxla/stablehlo/issues/731)).
- * `result[result_index] = reduce([source_values], [init_value], [0], scatter)`
-  where:
-+
-   * `source_values = [source[source_index] for source_index in
-    source_indices]`.
-   * `selected_index(source_index) = operand_index` if
-@@ -6111,6 +6206,7 @@
- #### Constraints
- 
- <!-- markdownlint-disable line-length -->
-+
- * (C1) `element_type(operand) = element_type(source)`.
- * (C2) `shape(source) = num_windows` where:
-   * `padded_operand_shape = padding[:, 0] + shape(operand) + padding[:, 1]`.
-@@ -7501,7 +7597,7 @@
- syntax ([ISO/IEC 14977:1996](https://www.iso.org/standard/26153.html),
- [Wikipedia](https://en.wikipedia.org/wiki/Extended_Backus%E2%80%93Naur_form)),
- with two modifications: 1) rules are defined using `::=` rather than `=`,
--2) concatenation is expressed using juxtaposition rather than `,`.
-+2\) concatenation is expressed using juxtaposition rather than `,`.
- 
- For describing semantics (i.e. within "Types", "Constants" and "Ops" sections),
- we are using formulas which are based on Python syntax extended with support
-@@ -7571,7 +7667,7 @@
- ### Names
- 
- In formulas, lexical scope includes: 1) global functions, 2) member definitions,
--3) local definitions. The list of global functions is provided below. The list
-+3\) local definitions. The list of global functions is provided below. The list
- of element definitions depends on the program element that the notation is
- applied to:
- 
-diff --ruN a/stablehlo/rfcs/20260821-low-precision-complex.md b/stablehlo/rfcs/20260821-low-precision-complex.md
---- stablehlo/rfcs/20260821-low-precision-complex.md
-+++ stablehlo/rfcs/20260821-low-precision-complex.md
-@@ -0,0 +1,599 @@
-+# [RFC] Support f16- and bf16-based complex types in StableHLO
-+
-+Status: In Review<br/>
-+Initial version: 08/21/2026<br/>
-+Last updated: 09/28/2026<br/>
-+Discussion thread: [openxla/stablehlo#2993](https://github.com/openxla/stablehlo/pull/2993)
-+
-+## Overview
-+
-+This RFC proposes extending the set of complex types supported by StableHLO
-+from:
-+
-+```text
-+complex<f32>
-+complex<f64>
-+```
-+
-+to:
-+
-+```text
-+complex<f16>
-+complex<bf16>
-+complex<f32>
-+complex<f64>
-+```
-+
-+The proposal uses the existing MLIR builtin `ComplexType`; it does not add a
-+new StableHLO type. It extends the set of floating-point component types that
-+StableHLO accepts inside `complex<T>`.
-+
-+This RFC defines the StableHLO and VHLO type and serialization contract for
-+low-precision complex types. It does not implement XLA/HLO, the
-+StableHLO-to-XLA bridge, runtime ABI exposure, or backend kernels. The feature
-+is intended to land as part of a coordinated implementation sequence with the
-+XLA proposal in [openxla/xla#6136](https://github.com/openxla/xla/issues/6136).
-+Approval of this RFC does not imply that every StableHLO consumer can execute
-+`complex<f16>` or `complex<bf16>`; consumers without support must reject these
-+types with a clear diagnostic before code generation.
-+
-+### Scope and current implementation status
-+
-+This RFC specifies the StableHLO/VHLO side of a cross-repository feature. The
-+current status of the related layers is:
-+
-+| Layer | Current status | Boundary of this RFC |
-+| --- | --- | --- |
-+| StableHLO type specification | Only `complex<f32>` and `complex<f64>` are currently specified | Adds `complex<f16>` and `complex<bf16>` |
-+| VHLO compatibility | `ComplexV1` is the only complex representation, and its current shallow verifier does not constrain the component type | Adds a target-version-aware component-type constraint on `ComplexV1` |
-+| XLA/HLO core types | `C32`/`BC32` are proposed in [openxla/xla#6136](https://github.com/openxla/xla/issues/6136), but are not yet implemented | Coordinated prerequisite; not implemented here |
-+| StableHLO-to-XLA bridge | No `C32`/`BC32` type or constant mapping is available | Follow-up implementation |
-+| Runtime and backend support | [cuFFT](https://docs.nvidia.com/cuda/cufft/index.html) exposes 16-bit complex formats, but XLA does not yet expose a corresponding path | Separate runtime and backend work |
-+
-+The StableHLO specification may define a valid type domain before every
-+consumer implements it. This RFC therefore defines the observable rejection
-+behavior for consumers that have not yet added execution support; it does not
-+claim end-to-end execution support from StableHLO serialization alone.
-+
-+## Motivation
-+
-+[openxla/stablehlo#1794](https://github.com/openxla/stablehlo/issues/1794)
-+requests support for 16-bit FFT input and output types. Today, the StableHLO
-+verifier prevents such programs even though MLIR can represent both
-+`complex<f16>` and `complex<bf16>`.
-+
-+There are two independent concerns:
-+
-+1. Whether StableHLO can represent and specify a low-precision complex value.
-+2. Whether a particular consumer or backend can execute that value efficiently.
-+
-+StableHLO should decide the first concern based on the coherence of its type
-+system and semantics rather than on the current implementation status of one
-+consumer. The related XLA proposal currently remains at the design stage; it
-+proposes adding `C32 = complex<f16>` and `BC32 = complex<bf16>` as a shared
-+core-type change before bridge and backend work. Consumers that do not support
-+the new types may reject them with a diagnostic until they add a lowering,
-+legalization, or native implementation. Such rejection is part of the staged
-+rollout, not an implicit promotion to `complex<f32>`.
-+
-+The body of #1794 suggests allowing any floating-point component type. This
-+RFC intentionally narrows that suggestion to f16 and bf16, the two widely used
-+16-bit floating-point formats, in addition to the existing f32 and f64 types.
-+Supporting only `bf16` would not address the issue's original `f16` request,
-+while allowing every member of StableHLO's broader floating-point type set
-+would also introduce complex types with FP4, FP6, FP8, and other specialized
-+component formats. Those formats have distinct representability and
-+compatibility questions and should not be included implicitly in this
-+proposal.
-+
-+## Current behavior
-+
-+The current StableHLO specification lists only `complex<f32>` and
-+`complex<f64>` as supported complex types.
-+
-+The implementation reflects this restriction in multiple places:
-+
-+- `HLO_Complex` is defined as `Complex<HLO_Float32Or64>`.
-+- The operands of `stablehlo.complex` are restricted to f32/f64 tensors.
-+- RFFT type inference has a hand-written check that requires f32 or f64 input.
-+- The StableHLO test suite contains a negative test requiring f16 RFFT input to
-+  be rejected.
-+
-+Changing only the RFFT check would be incomplete: RFFT could infer a type that
-+the common StableHLO complex constraint and `stablehlo.complex` still reject.
-+
-+VHLO currently represents MLIR builtin complex types with `ComplexV1`, whose
-+verifier only checks that the component is a VHLO type. This is a shallow
-+structural check: the representation can therefore spell component
-+combinations that were never in the historical StableHLO complex type domain.
-+Historically valid `ComplexV1` values are limited to f32 and f64 components.
-+
-+The current VHLO version-conversion pass has no type-to-type version
-+conversions because no VHLO type has previously required one. It recursively
-+checks whether nested types are legal for the target version, and that check
-+is already target-version-aware, but complex-type legality today does not
-+depend on the component type. Version-dependent constraint checks already
-+exist at the operation level: `validateConstraint` gates the reduce-family
-+operations and `custom_call_v1` on the target version, for example allowing
-+mismatched operand and result element types in reduce operations only from
-+0.17.0. This RFC extends that existing constraint pattern to the complex
-+component type rather than introducing new type-conversion infrastructure.
-+
-+## Proposal
-+
-+### Supported complex component types
-+
-+StableHLO will define an explicit complex component type set:
-+
-+```text
-+ComplexComponent ::= f16 | bf16 | f32 | f64
-+ComplexType      ::= complex<ComplexComponent>
-+```
-+
-+The implementation should use a dedicated ODS constraint for this set rather
-+than define `HLO_Complex` in terms of the full `HLO_Float` constraint. This
-+keeps future additions to `HLO_Float` from automatically changing the
-+StableHLO complex type system without review.
-+
-+No mixed-component complex type is introduced. A complex value continues to
-+have real and imaginary components with the same element type.
-+
-+### Effect on existing operations
-+
-+This proposal does not add an operation and does not add complex semantics to
-+an operation that currently excludes complex types.
-+
-+For every existing operation whose specified input or result type domain
-+already includes a complex type, `complex<f16>` and `complex<bf16>` become
-+additional valid instantiations of that existing type domain. This RFC
-+proposes no operation-specific exception. If review identifies an operation
-+whose semantics cannot support either new instantiation, that exception must
-+be specified in a revision of this RFC rather than chosen during
-+implementation.
-+
-+This consequence is intentional. Restricting the new types to FFT alone would
-+allow a low-precision complex value to be produced but prevent it from flowing
-+through the same basic construction, extraction, arithmetic, conversion, and
-+structural operations as other StableHLO complex values.
-+
-+The observable result type of an operation remains the type written in the
-+StableHLO program. This RFC does not insert implicit promotion to f32 and does
-+not change `complex<f16>` or `complex<bf16>` into another IR type. A consumer
-+may use wider intermediate precision where permitted by the existing operation
-+semantics and accuracy rules, but the declared input and result types remain
-+unchanged.
-+
-+The normative affected surface is grouped as follows:
-+
-+- Direct construction and extraction operations, including `complex`, `real`,
-+  and `imag`, accept the two new instantiations.
-+- Existing parametric complex arithmetic and transcendental operations accept
-+  them under their existing semantics and accuracy contracts. This includes
-+  the direct complex-domain uses in `abs`, `atan2`, `cbrt`, `cosine`, `divide`,
-+  `exponential`, `exponential_minus_one`, `log`, `log_plus_one`, `logistic`,
-+  `negate`, `power`, `remainder`, `rsqrt`, `sign`, `sine`, `sqrt`, `subtract`,
-+  `tan`, and `tanh`, as well as operations such as `add` and `multiply` whose
-+  generic tensor domains already include complex element types.
-+- Existing complex linear-algebra operations, including `cholesky` and
-+  `triangular_solve`, accept the two new instantiations.
-+- FFT, IFFT, RFFT, and IRFFT accept the relationships specified below.
-+- Constants, `iota`, conversion, data-movement, shape-preserving, control-flow,
-+  and structural operations continue to accept complex values wherever their
-+  current type domains already do so.
-+- A supported complex type may occur wherever an existing StableHLO container
-+  admits a complex element or nested value, including tensors, buffers,
-+  tuples, futures, function signatures, region arguments, and custom calls.
-+
-+Because `HLO_Complex` is used transitively, the implementation audit must cover
-+the complete ODS constraint graph rather than only literal `HLO_Complex`
-+occurrences. In particular, derived constraints such as `HLO_Tensor`,
-+`HLO_NonQuantizedTensor`, `HLO_Buffer`, `HLO_Tuple`, `HLO_CustomCallValue`, and
-+their static-shape or unranked variants expand with the common complex
-+component set. Every custom verifier and inference path reached through these
-+constraints must be checked against the specification above.
-+
-+This shared-constraint change also reaches the CHLO dialect: CHLO operations
-+such as `broadcast_complex` and the inverse-trigonometric and hyperbolic
-+complex math operations use the same complex tensor predicates. The
-+implementation must either include those operations and their
-+legalization/decomposition paths in the affected-surface audit and tests, or
-+deliberately separate their constraints so that this RFC does not expand CHLO
-+acceptance accidentally.
-+
-+This expansion does not create separate kinds of complex type. It also does
-+not define the physical layout or ABI of a buffer containing a low-precision
-+complex element; StableHLO buffer acceptance remains an abstract type-system
-+property under the existing buffer contract.
-+
-+### `complex`, `real`, and `imag`
-+
-+For every supported complex component type `T`:
-+
-+```text
-+stablehlo.complex : (tensor<...xT>, tensor<...xT>)
-+                 -> tensor<...xcomplex<T>>
-+
-+stablehlo.real    : tensor<...xcomplex<T>> -> tensor<...xT>
-+stablehlo.imag    : tensor<...xcomplex<T>> -> tensor<...xT>
-+```
-+
-+The existing shape and same-element-type constraints continue to apply.
-+
-+In particular, the following remain invalid:
-+
-+- constructing a complex value from one f16 tensor and one bf16 tensor;
-+- declaring a `complex<f32>` result for f16 components;
-+- declaring a real or imaginary result whose component type differs from the
-+  input complex component type.
-+
-+### FFT type relationships
-+
-+For `T` in `{f16, bf16, f32, f64}`, the FFT type relationships are:
-+
-+```text
-+FFT:   complex<T> -> complex<T>
-+IFFT:  complex<T> -> complex<T>
-+RFFT:          T  -> complex<T>
-+IRFFT: complex<T> -> T
-+```
-+
-+The existing FFT rank, shape, and `fft_length` constraints are unchanged.
-+
-+The hand-written RFFT type check should use the same supported complex
-+component predicate as the common complex type constraint, avoiding a second
-+component-type list that can drift from the specification.
-+
-+### Constants and conversion
-+
-+Complex constants with f16 or bf16 components follow the existing StableHLO
-+complex constant definition: the real and imaginary literals are interpreted
-+using the floating-point semantics of the component type.
-+
-+`stablehlo.convert` continues to use its existing specified conversion
-+semantics. This RFC makes the new complex types valid source or result types in
-+cases where the existing `convert` type domain includes complex types; it does
-+not introduce an implicit conversion or a new fallback rule.
-+
-+## Compatibility and VHLO
-+
-+Low-precision complex support is a new StableHLO feature and therefore requires
-+a StableHLO minor version bump. The implementation must use the next available
-+minor version at the time it lands. If implemented immediately after StableHLO
-+1.20.0, the feature version would be 1.21.0.
-+
-+The required observable compatibility behavior is:
-+
-+- serialization targeting the feature version or a later version succeeds;
-+- serialization of a program using `complex<f16>` or `complex<bf16>` to an
-+  earlier target version fails with a diagnostic;
-+- no compatibility path may silently promote the value to `complex<f32>` or
-+  decompose it into unrelated tensor types;
-+- programs using only `complex<f32>` and `complex<f64>` retain their existing
-+  ability to target older supported StableHLO versions;
-+- current consumers remain able to deserialize portable artifacts containing
-+  historical f32/f64 complex programs.
-+
-+StableHLO version compatibility and consumer execution capability are separate
-+contracts. Serialization to the feature version or a later version only
-+establishes that the artifact is representable by that StableHLO version. It
-+does not prove that an XLA, PJRT, or other consumer can lower every operation
-+containing the new type. A consumer that understands the StableHLO version but
-+lacks `C32` or `BC32` support must reject the module before code generation
-+with a diagnostic that identifies the unsupported type and operation.
-+
-+### Proposed VHLO representation
-+
-+VHLO retains `ComplexV1` as the only VHLO representation of complex types.
-+`ComplexV1` is structurally unchanged: one element-type parameter with the
-+existing shallow verifier, and an unchanged version range. The
-+version-dependent validity of component types is expressed as a
-+target-version-aware type constraint, following the pattern VHLO already uses
-+for operation constraints (`validateConstraint` on the reduce-family
-+operations and `custom_call_v1`).
-+
-+The constraint is normative:
-+
-+- For a target version at or after the feature version, `ComplexV1` accepts
-+  exactly f16-, bf16-, f32-, and f64-based components. FP8, integer, and all
-+  other component types are rejected.
-+- For a target version before the feature version, `ComplexV1` accepts only
-+  f32- and f64-based components. A `ComplexV1<f16>` or `ComplexV1<bf16>` value
-+  makes serialization to that target fail with a diagnostic, in the same way
-+  that a reduce operation with mismatched operand and result element types
-+  fails when the target version is older than 0.17.0.
-+- No new VHLO type or attribute is introduced, and no type-to-type conversion
-+  or type rewriting occurs during version conversion. Serialization of
-+  programs that use only f32/f64 complex types is therefore unaffected.
-+
-+The implementation must add a `VHLO_VersionedTypeConstraintInterface` as the
-+type-level analogue of `VHLO_VersionedOpConstraintInterface`. It provides a
-+`validateConstraint(Type, Version)` hook. `ComplexV1` must implement the hook,
-+and the recursive `isLegalType` check must invoke it. A local helper may
-+provide the hook's implementation, but the target-aware check must have a
-+type-level entry point; operation-only validation is insufficient because
-+complex types can be nested in signatures, containers, and type-bearing
-+attributes. This interface is an implementation hook and does not change the
-+serialized VHLO type identity or its version range.
-+
-+This models a component-domain relaxation as a versioned constraint rather
-+than as a new type version. VHLO is a serialization dialect: it is produced by
-+conversion from StableHLO and consumed by version conversion, not authored by
-+hand, and the dialect does not guarantee rejection of hand-authored payloads
-+that could never be produced through those flows. Every artifact is tagged
-+with the version it targets, and a consumer must be at least that version, so
-+a payload cannot be presented to an older consumer than it was serialized for.
-+
-+Officially produced bytecode is validated against the requested target version
-+before emission. Consequently, no official pre-feature artifact can contain a
-+low-precision complex value. Hand-authored VHLO payloads that violate this
-+producer contract are outside this RFC's portable-artifact guarantees; a
-+consumer may reject them during parsing or conversion, and accepting one must
-+not establish that the pre-feature target historically supported the type.
-+
-+### Recursive legality requirements
-+
-+The component-type constraint must be evaluated by the recursive type
-+legality check with access to the requested target version. The current check
-+already recurses through several nested VHLO types and attributes, but the
-+implementation for this feature must also add explicit coverage for every
-+type-bearing location listed below. Checking only an operation's immediate
-+result types would leave invalid values hidden in containers or signatures.
-+The constraint must therefore be enforced at every location, including:
-+
-+- `RankedTensorV1`, `UnrankedTensorV1`, and `RankedBufferV1` element types;
-+- `TupleV1` elements, nested tuples, and `FutureV1` element types;
-+- `FunctionV1` inputs and results;
-+- operation operands and results, function signatures, and region block
-+  arguments;
-+- `TypeV1Attr`, `TensorV1Attr`, `FloatV1Attr`, `IntegerV1Attr`, and every other
-+  attribute that carries a type;
-+- ranked-tensor encodings and nested array or dictionary attributes; and
-+- the storage and expressed types of `UniformQuantizedV1` and
-+  `UniformQuantizedPerAxisV1`, which must be traversed even when the
-+  quantized-type semantics reject a complex inhabitant.
-+
-+If any nested occurrence fails the component-type constraint for the requested
-+target, the entire version conversion must fail before bytecode is written.
-+The diagnostic must identify the target version and the unsupported complex
-+component type. No path may partially convert the module or emit an artifact
-+containing a complex value whose component type is invalid for the target
-+version.
-+
-+### Conversion behavior and invariants
-+
-+Because no VHLO type changes, version conversion requires no type rewriting.
-+The observable behavior is:
-+
-+1. Validate the target version as today.
-+2. Preserve the existing operation-version rewrites; complex types are never
-+   rewritten.
-+3. Extend the recursive type legality check to evaluate the component-type
-+   constraint against the requested target version over every type-bearing
-+   location listed above.
-+4. If any check fails, fail the pass before bytecode emission with a
-+   diagnostic naming the target version and the unsupported component type.
-+
-+The transformation remains all-or-nothing from the serializer's perspective,
-+and serialization must not emit bytecode from a module whose complex component
-+types are invalid for the requested target version.
-+
-+## Consumer behavior
-+
-+StableHLO verifier acceptance is not a claim that every consumer can execute
-+the new types.
-+
-+StableHLO serialization validates only StableHLO/VHLO representability. A
-+StableHLO-to-XLA bridge must reject `complex<f16>` or `complex<bf16>` before
-+constructing the XLA/HLO module when the corresponding XLA primitive type is
-+unavailable. The diagnostic must identify the unsupported type, operation, and
-+target. Other consumers must provide an equivalent pre-code-generation
-+capability check or a separately specified legalization path.
-+
-+A consumer that lacks support may:
-+
-+- reject a module containing the new types with a clear unsupported-type
-+  diagnostic;
-+- apply a separately specified legalization;
-+- implement the types natively.
-+
-+This RFC does not prescribe which option a consumer must choose. It does
-+require that unsupported consumers fail explicitly before code generation. In
-+particular, a consumer must not reinterpret a four-byte value as `C64`, fall
-+through an existing `C64`/`C128` path, or silently promote the declared value to
-+`complex<f32>`. XLA support is not a prerequisite for StableHLO type validity,
-+but the StableHLO implementation and the XLA/HLO work must be coordinated
-+before the feature is presented as end-to-end executable.
-+
-+## Verification and testing
-+
-+The implementation must include positive tests for:
-+
-+- StableHLO verifier acceptance and exact type inference for the complete
-+  specified domain, independently of whether a particular consumer executes
-+  every operation;
-+- parsing and printing `complex<f16>` and `complex<bf16>` in StableHLO
-+  programs;
-+- `stablehlo.complex`, `stablehlo.real`, and `stablehlo.imag` type relations;
-+- an acceptance case for every existing operation with a direct complex-domain
-+  numeric, generator, or linear-algebra constraint, including `iota`, plus
-+  representative cases for every transitive structural and container
-+  constraint family listed above;
-+- FFT, IFFT, RFFT, and IRFFT for both f16 and bf16 component types;
-+- exact return-type inference, including dynamic-shape cases;
-+- StableHLO-to-current-VHLO-to-StableHLO round trips for both `complex<f16>`
-+  and `complex<bf16>`;
-+- serialization and deserialization at the new version for both
-+  `complex<f16>` and `complex<bf16>`;
-+- deserialization of historical f32/f64 complex artifacts, which requires no
-+  type rewriting under this proposal;
-+- serialization of `complex<f16>` and `complex<bf16>` values, including values
-+  nested in every semantically valid type-bearing container and attribute
-+  listed above, to the feature version and later targets; and
-+- continued old-version serialization of f32/f64 complex programs.
-+
-+The implementation must include negative tests for:
-+
-+- consumer-side rejection before code generation when `C32` or `BC32` support
-+  is absent, including a diagnostic that names the unsupported type, operation,
-+  and target;
-+- complex values with integer component types;
-+- at least one floating-point component type outside this RFC, such as an FP8
-+  type;
-+- rejection of serialization to any target version before the feature version
-+  when a `ComplexV1<f16>` or `ComplexV1<bf16>` value is present, both directly
-+  and when nested in every semantically valid supported type-bearing container
-+  or attribute listed above, enforced by the version-conversion constraint
-+  rather than by tightening the shallow V1 parser or verifier;
-+- rejection of complex types placed in quantized-type storage or expressed
-+  type parameters by the semantic constraint owned by the corresponding
-+  StableHLO or consumer layer, while those locations are still traversed for
-+  VHLO version legality;
-+- a complex value with an FP8, integer, or other unsupported component, for
-+  any target version;
-+- mismatched f16/bf16 inputs to `stablehlo.complex`;
-+- mismatched real/complex component types in FFT input and result types;
-+- attempting to serialize a program containing `complex<f16>` or
-+  `complex<bf16>` to the version immediately before the feature version; and
-+- failure before bytecode emission, with a diagnostic that names both the
-+  requested target version and the unsupported nested type.
-+
-+The VHLO compatibility suite must cover the feature version following the
-+existing VHLO checklist, including fixtures that exercise the component-type
-+constraint against old and new target versions.
-+
-+Reference-interpreter numerical tests may be delivered in a separate
-+StableHLO-only follow-up if the required tensor storage and constant
-+materialization changes would obscure review of the type-system and
-+compatibility change. If maintainers require reference execution as an opset
-+acceptance criterion, that follow-up can be stacked or folded into the
-+implementation change without adding XLA or backend work.
-+
-+## Non-goals
-+
-+This RFC does not propose:
-+
-+- the implementation of an XLA primitive type or StableHLO-to-XLA bridge;
-+  those pieces are coordinated with [openxla/xla#6136](https://github.com/openxla/xla/issues/6136)
-+  and are required before claiming end-to-end execution;
-+- JAX dtype, promotion, tracing, or lowering changes;
-+- portable decomposition into real and imaginary tensor planes;
-+- CPU, GPU, TPU, or accelerator kernels;
-+- FFT library integration;
-+- a memory layout, storage ABI, or claim that a value occupies four bytes in a
-+  particular runtime;
-+- mandated accumulation precision or a new numerical-accuracy contract;
-+- automatic fallback or promotion to `complex<f32>`;
-+- support for complex values whose components use FP4, FP6, FP8, E8M0, or any
-+  future floating-point type not listed by this RFC;
-+- performance claims.
-+
-+These concerns are separate implementation stages, but they are not an
-+indefinite dependency: the enabling StableHLO implementation must have an
-+agreed XLA/HLO core-type plan and an explicit unsupported-consumer path before
-+it lands. The bridge, runtime, and backend changes may remain separate pull
-+requests.
-+
-+## Alternatives considered
-+
-+### Support only `complex<bf16>`
-+
-+This is the smallest change for bf16-specific use cases, but it does not address
-+the f16 request in #1794 and introduces an asymmetric exception between the two
-+widely used 16-bit floating-point formats. The additional StableHLO and VHLO
-+work required to support f16 and bf16 together is substantially shared, so this
-+RFC does not recommend a bf16-only type domain.
-+
-+### Support `Complex<HLO_Float>`
-+
-+This is mechanically concise but would include every floating-point type in
-+`HLO_Float`, including specialized FP4, FP6, FP8, and E8M0 formats. It would
-+also cause future additions to `HLO_Float` to expand the StableHLO complex type
-+domain without a separate compatibility decision. This RFC instead proposes a
-+closed component set.
-+
-+### Support low-precision complex types only on FFT
-+
-+This narrows the immediate verifier change but produces a fragmented type
-+system in which FFT can create a complex value that basic complex construction,
-+extraction, arithmetic, conversion, or structural operations may reject. If a
-+strictly FFT-only feature is desired, it should be proposed as an explicit
-+operation-specific type extension rather than as general StableHLO support for
-+low-precision complex types.
-+
-+### Reuse `ComplexV1` with target-aware validation (chosen)
-+
-+VHLO keeps `ComplexV1` as the only complex representation and expresses the
-+version-dependent validity of component types as a target-version-aware type
-+constraint, the same pattern the dialect already uses for operation
-+constraints. Review of this RFC directed this approach over a new type
-+version: VHLO is a serialization dialect produced by conversion and consumed
-+by version conversion rather than authored by hand, every artifact is tagged
-+with its target version, and a downgrade containing values invalid for the
-+target already fails today. Modeling a component-domain relaxation with the
-+constraint mechanism keeps the category of change consistent with existing
-+practice and requires far less code.
-+
-+### Add `ComplexV2` at the feature version (declined)
-+
-+A new `ComplexV2` with a closed component set keeps the meaning of a complex
-+type intrinsic to the type version and makes the compatibility boundary
-+inspectable without target context. It would, however, require the dialect's
-+first type version split together with new type-to-type conversion
-+infrastructure: recursive rewriting of enclosing types, function signatures,
-+and type-bearing attributes for a change that does not alter the structure of
-+the type. Given VHLO's serialization-dialect stance, that cost is not
-+justified for this category of change.
-+
-+The alternatives are:
-+
-+| Design | Benefit | Cost |
-+| --- | --- | --- |
-+| Target-aware validation on `ComplexV1` | Consistent with the existing versioned-constraint pattern; no type rewriting; far less code | Complex-type validity depends on the target version rather than the type alone |
-+| `ComplexV2` with V1/V2 conversion | Type-intrinsic meaning and an inspectable boundary without target context | First VHLO type version split plus recursive type-conversion infrastructure |
-+
-+Reusing `ComplexV1` with no version check at all would allow new programs to
-+appear serializable to StableHLO versions that never specified these
-+component combinations, and remains rejected.
-+
-+## Rollout and pull request boundaries
-+
-+The proposed upstream sequence is deliberately split by capability:
-+
-+1. Keep this RFC as a design-only PR and confirm the cross-repository contract
-+   with [openxla/xla#6136](https://github.com/openxla/xla/issues/6136). RFC
-+   approval does not itself enable serialization or execution.
-+2. Obtain agreement on the XLA/HLO core representation (`C32` and `BC32`),
-+   host storage, `LiteralProto`, parser/printer behavior, and explicit
-+   unsupported-execution diagnostics. The current XLA proposal places this
-+   core-type change before bridge and backend work.
-+3. After the RFC is approved, submit the compatibility-atomic StableHLO
-+   implementation, including:
-+
-+   - specification, ODS, verifier, and type-inference changes;
-+   - the target-version-aware `ComplexV1` component-type constraint and its
-+     recursive legality coverage;
-+   - positive, negative, round-trip, and serialization tests;
-+   - tests that reject unsupported target versions before bytecode emission.
-+
-+   The affected-surface audit must include CHLO operations that inherit the
-+   shared complex predicates, or record an explicit constraint boundary that
-+   excludes them.
-+
-+   The implementation may be developed in parallel with the XLA core change,
-+   but it must not be presented as end-to-end execution support. No released
-+   state may emit the new StableHLO types without the corresponding VHLO
-+   boundary.
-+4. Implement the StableHLO-to-XLA bridge, Shape mapping, constant materialization,
-+   and bit-exact round trips in a separate follow-up. A bridge that has not yet
-+   added `C32`/`BC32` must reject those types before code generation.
-+5. Add explicit legalization policies, keeping pair decomposition and
-+   promotion to `C64` as separate choices. The core type change must not choose
-+   either policy implicitly.
-+6. Add runtime ABI exposure and backend implementations, beginning with a
-+   concrete backend path such as the cuFFT 16-bit complex formats and then
-+   extending support as maintainers approve.
-+
-+The RFC PR must not close #1794 by itself because design approval does not
-+implement the requested functionality. The implementation PR can close the
-+issue only after the agreed StableHLO support and the required consumer path
-+are present.
-diff --ruN a/stablehlo/rfcs/20260907-exp2-log2.md b/stablehlo/rfcs/20260907-exp2-log2.md
---- stablehlo/rfcs/20260907-exp2-log2.md
-+++ stablehlo/rfcs/20260907-exp2-log2.md
-@@ -0,0 +1,51 @@
-+# [RFC] Add Exp2Op and Log2Op to the StableHLO specification
-+
-+Status: Review<br/>
-+Initial version: 09/07/2026<br/>
-+Last updated: 10/02/2026<br/>
-+Discussion thread: [GitHub](https://github.com/openxla/stablehlo/pulls)
-+
-+## Motivation
-+
-+IEEE-754 (`exp2`, `log2`), C99 `<math.h>`, and LLVM intrinsics (`llvm.exp2`,
-+`llvm.log2`) include base-2 exponential (`exp2`) and base-2 logarithm (`log2`)
-+as standard elementary math functions, and hardware accelerators (such as GPU
-+special function units and TPU EUP instructions) support them directly.
-+
-+Currently, StableHLO lacks native `exp2` and `log2` operations:
-+
-+* Frameworks expose `exp2` and `log2` directly: JAX provides [`jnp.exp2`][jax1]
-+  / `lax.exp2` and [`jnp.log2`][jax2] / `lax.log2`, and PyTorch provides
-+  [`torch.exp2`][torch1] and [`torch.log2`][torch2].
-+* Without first-class ops in StableHLO, frameworks must decompose `exp2(x)` into
-+  `exponential(x * ln(2))` and `log2(x)` into `log(x) / ln(2)`.
-+* Decomposing `exp2` and `log2` into natural `exponential` and `log` scaled by
-+  `ln(2)` degrades numerical accuracy (introducing rounding error on integer
-+  powers of two where `exp2(n)` and `log2(2^n)` should evaluate exactly) and
-+  prevents backends from emitting single-instruction base-2 hardware operations
-+  without fragile algebraic pattern-matching.
-+
-+Following the precedent of `TanOp` ([20240312-tanop.md](20240312-tanop.md)) and
-+the transcendental accuracy specification
-+([20241015-result-accuracy.md](20241015-result-accuracy.md)), this RFC adds
-+`exp2` (`Exp2Op`) and `log2` (`Log2Op`) to StableHLO with optional
-+`result_accuracy` attributes.
-+
-+## Proposed Specification Changes
-+
-+Please refer to [`docs/spec.md`](../docs/spec.md) for the specification of
-+[`exp2`](../docs/spec.md#exp2) and [`log2`](../docs/spec.md#log2).
-+
-+## Exp2Op, Log2Op, and Compatibility
-+
-+* **VHLO Serialization**: StableHLO version `1.21.0` introduces `VHLO_Exp2OpV1`
-+  (`exp2_v1`) and `VHLO_Log2OpV1` (`log2_v1`).
-+* **Forward Compatibility & Decomposition**: For runtimes targeting StableHLO
-+  versions prior to `1.21.0`, compatibility expansion lowers `exp2(x)` to
-+  `exponential(multiply(x, constant(ln(2))))` and `log2(x)` to
-+  `divide(log(x), constant(ln(2)))`.
-+
-+[jax1]: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.exp2.html
-+[jax2]: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.log2.html
-+[torch1]: https://pytorch.org/docs/stable/generated/torch.exp2.html
-+[torch2]: https://pytorch.org/docs/stable/generated/torch.log2.html
-diff --ruN a/stablehlo/stablehlo/conversions/linalg/tests/pointwise.mlir b/stablehlo/stablehlo/conversions/linalg/tests/pointwise.mlir
---- stablehlo/stablehlo/conversions/linalg/tests/pointwise.mlir
-+++ stablehlo/stablehlo/conversions/linalg/tests/pointwise.mlir
-@@ -728,6 +728,35 @@
- 
- // -----
- 
-+// CHECK-LABEL: func @float_cmp_weakorder
-+// CHECK-PRIMITIVE-LABEL: func @float_cmp_weakorder
-+func.func @float_cmp_weakorder(%lhs: tensor<2x2xbf16>,
-+                %rhs: tensor<2x2xbf16>) -> (tensor<2x2xi1>) {
-+  %0 = "stablehlo.compare"(%lhs, %rhs) {
-+    comparison_direction = #stablehlo<comparison_direction LT>,
-+    compare_type = #stablehlo<comparison_type WEAKORDER>
-+  } : (tensor<2x2xbf16>, tensor<2x2xbf16>) -> tensor<2x2xi1>
-+  func.return %0 : tensor<2x2xi1>
-+}
-+// CHECK: tensor.empty() : tensor<2x2xi1>
-+// CHECK: linalg.generic
-+// CHECK-NEXT: ^bb0(%[[LHS_IN:.*]]: bf16, %[[RHS_IN:.*]]: bf16, %{{.*}}: i1):
-+// CHECK-NEXT:   %[[CMP_ULT:.*]] = arith.cmpf ult, %[[LHS_IN]], %[[RHS_IN]] : bf16
-+// CHECK-NEXT:   %[[CMP_ORD:.*]] = arith.cmpf ord, %[[LHS_IN]], %[[LHS_IN]] : bf16
-+// CHECK-NEXT:   %[[RESULT:.*]] = arith.andi %[[CMP_ULT]], %[[CMP_ORD]] : i1
-+// CHECK-NEXT:   linalg.yield %[[RESULT]] : i1
-+
-+// CHECK-PRIMITIVE: linalg.map
-+// CHECK-PRIMITIVE-SAME: ins(
-+// CHECK-PRIMITIVE-SAME: outs(
-+// CHECK-PRIMITIVE-NEXT: (%[[LHS_IN:[a-zA-Z0-9]*]]: bf16, %[[RHS_IN:.*]]: bf16, %[[RESULT_OUT:.*]]: i1) {
-+// CHECK-PRIMITIVE-NEXT:   %[[CMP_ULT:.*]] = arith.cmpf ult, %[[LHS_IN]], %[[RHS_IN]] : bf16
-+// CHECK-PRIMITIVE-NEXT:   %[[CMP_ORD:.*]] = arith.cmpf ord, %[[LHS_IN]], %[[LHS_IN]] : bf16
-+// CHECK-PRIMITIVE-NEXT:   %[[RESULT:.*]] = arith.andi %[[CMP_ULT]], %[[CMP_ORD]] : i1
-+// CHECK-PRIMITIVE-NEXT:   linalg.yield %[[RESULT]] : i1
-+
-+// -----
-+
- // CHECK-LABEL: func @int_cmp
- // CHECK-PRIMITIVE-LABEL: func @int_cmp
- func.func @int_cmp(%lhs: tensor<2x2xi32>,
-diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h b/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h
---- stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h
-+++ stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.h
-@@ -60,26 +60,26 @@
-     unsigned nParallelLoops);
- 
- /// Generates an init sparse tensor.
--Value getEmptySparseTensor(OpBuilder &b, Location loc, ShapedType type,
-+Value getEmptySparseTensor(OpBuilder& b, Location loc, ShapedType type,
-                            ArrayRef<Value> dynSizes);
- 
- /// Generates a tensor.empty op.
--Value getEmptyTensor(OpBuilder &b, Location loc, ShapedType type,
-+Value getEmptyTensor(OpBuilder& b, Location loc, ShapedType type,
-                      ArrayRef<Value> dynSizes);
- 
- /// Generates an empty tensor for the result of the operation, which could be a
- /// dense tensor or a sparse tensor.
--Value getEmptyTensorFor(OpBuilder &b, Location loc, ShapedType resultType,
--                        Operation *op, ValueRange operands);
-+Value getEmptyTensorFor(OpBuilder& b, Location loc, ShapedType resultType,
-+                        Operation* op, ValueRange operands);
- 
- /// Ensures a tensor has the same shape (not including the element type) as
- /// another.
--Value coerceTensorShape(OpBuilder &builder, Location loc,
-+Value coerceTensorShape(OpBuilder& builder, Location loc,
-                         TypedValue<ShapedType> value, ShapedType targetType);
- 
- /// Fills |tensor| with a zero constant of the matching type. Returns the new
- /// value.
--Value fillTensorWithZeros(OpBuilder &builder, Location loc, Value tensor);
-+Value fillTensorWithZeros(OpBuilder& builder, Location loc, Value tensor);
- 
- /// Sparsifies a (block of) operation(s) that cannot be handled directly
- /// by the sparse compiler but has well-known semi-ring semantics.
-@@ -94,17 +94,17 @@
- ///     }
- ///     absent={}
- ///   linalg.yield %result
--Value preSparsify(Operation *op, llvm::SmallVector<Value, 2> &values, Type rtp,
--                  OpBuilder *b);
-+Value preSparsify(Operation* op, llvm::SmallVector<Value, 2>& values, Type rtp,
-+                  OpBuilder* b);
- 
- /// Finalizes sparse semi-ring construction.
--Value postSparsify(Operation *op, Value semiring, Value result, OpBuilder *b);
-+Value postSparsify(Operation* op, Value semiring, Value result, OpBuilder* b);
- 
- /// Returns true if all operands are tensors with rank 0.
--bool allOperandsAreScalarTensors(Operation *op);
-+bool allOperandsAreScalarTensors(Operation* op);
- 
- /// Returns true if parent op is linalg.
--bool isInBodyOfLinalgOps(Operation *op);
-+bool isInBodyOfLinalgOps(Operation* op);
- 
- /// Extracts integer values from the attribute |elements|.
- SmallVector<int64_t> extract1DVector(DenseIntElementsAttr elements);
-diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToScalarOp.h b/stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToScalarOp.h
---- stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToScalarOp.h
-+++ stablehlo/stablehlo/conversions/linalg/transforms/MapStablehloToScalarOp.h
-@@ -17,6 +17,7 @@
- #ifndef STABLEHLO_CONVERSIONS_LINALG_TRANSFORMS_MAP_STABLEHLO_TO_SCALAR_OP_H
- #define STABLEHLO_CONVERSIONS_LINALG_TRANSFORMS_MAP_STABLEHLO_TO_SCALAR_OP_H
- 
-+#include <cmath>
- #include <optional>
- 
- #include "llvm/ADT/ArrayRef.h"
-@@ -93,6 +94,10 @@
-   using COp = ::mlir::complex::Expm1Op;
- };
- template <>
-+struct StablehloToScalarOp<stablehlo::Exp2Op> {
-+  using FOp = ::mlir::math::Exp2Op;
-+};
-+template <>
- struct StablehloToScalarOp<stablehlo::FloorOp> {
-   using FOp = ::mlir::math::FloorOp;
- };
-@@ -107,6 +112,10 @@
-   using COp = ::mlir::complex::Log1pOp;
- };
- template <>
-+struct StablehloToScalarOp<stablehlo::Log2Op> {
-+  using FOp = ::mlir::math::Log2Op;
-+};
-+template <>
- struct StablehloToScalarOp<stablehlo::MulOp> {
-   using FOp = ::mlir::arith::MulFOp;
-   using IOp = ::mlir::arith::MulIOp;
-@@ -204,7 +213,8 @@
-   Value operator()(Location loc, ArrayRef<Type> resultTypes,
-                    ArrayRef<Type> /*argTypes*/, ValueRange args, OpBuilder* b) {
-     return StdScalarOp::create(*b, loc, resultTypes, args,
--                               ArrayRef<NamedAttribute>());
-+                               typename StdScalarOp::Properties{},
-+                               /*discardableAttributes=*/{});
-   }
- };
- 
-@@ -215,7 +225,8 @@
-     Type elementType = getElementTypeOrSelf(argTypes.front());
-     if (SupportedType{}(elementType)) {
-       return StdScalarOp::create(*b, loc, resultTypes, args,
--                                 ArrayRef<NamedAttribute>());
-+                                 typename StdScalarOp::Properties{},
-+                                 /*discardableAttributes=*/{});
-     }
-     return MapStablehloOpToScalarOpImpl<Args...>{}(loc, resultTypes, argTypes,
-                                                    args, b);
-@@ -456,6 +467,60 @@
-       assert(predicate.has_value() && "expected valid comparison direction");
-       return arith::CmpIOp::create(*b, loc, *predicate, lhsInt, rhsInt);
-     }
-+    if (adaptor.getCompareType() &&
-+        *adaptor.getCompareType() == stablehlo::ComparisonType::WEAKORDER) {
-+      // Weak order treats -0.0 == +0.0 and orders all NaNs equal and > +Inf.
-+      switch (comparisonDirection) {
-+        case stablehlo::ComparisonDirection::LE: {
-+          Value ole = arith::CmpFOp::create(*b, loc, arith::CmpFPredicate::OLE,
-+                                            lhs, rhs);
-+          Value rhsNan = arith::CmpFOp::create(
-+              *b, loc, arith::CmpFPredicate::UNO, rhs, rhs);
-+          return arith::OrIOp::create(*b, loc, ole, rhsNan);
-+        }
-+        case stablehlo::ComparisonDirection::GE: {
-+          Value oge = arith::CmpFOp::create(*b, loc, arith::CmpFPredicate::OGE,
-+                                            lhs, rhs);
-+          Value lhsNan = arith::CmpFOp::create(
-+              *b, loc, arith::CmpFPredicate::UNO, lhs, lhs);
-+          return arith::OrIOp::create(*b, loc, oge, lhsNan);
-+        }
-+        case stablehlo::ComparisonDirection::LT: {
-+          Value ult = arith::CmpFOp::create(*b, loc, arith::CmpFPredicate::ULT,
-+                                            lhs, rhs);
-+          Value lhsOrd = arith::CmpFOp::create(
-+              *b, loc, arith::CmpFPredicate::ORD, lhs, lhs);
-+          return arith::AndIOp::create(*b, loc, ult, lhsOrd);
-+        }
-+        case stablehlo::ComparisonDirection::GT: {
-+          Value ugt = arith::CmpFOp::create(*b, loc, arith::CmpFPredicate::UGT,
-+                                            lhs, rhs);
-+          Value rhsOrd = arith::CmpFOp::create(
-+              *b, loc, arith::CmpFPredicate::ORD, rhs, rhs);
-+          return arith::AndIOp::create(*b, loc, ugt, rhsOrd);
-+        }
-+        case stablehlo::ComparisonDirection::EQ: {
-+          Value oeq = arith::CmpFOp::create(*b, loc, arith::CmpFPredicate::OEQ,
-+                                            lhs, rhs);
-+          Value lhsNan = arith::CmpFOp::create(
-+              *b, loc, arith::CmpFPredicate::UNO, lhs, lhs);
-+          Value rhsNan = arith::CmpFOp::create(
-+              *b, loc, arith::CmpFPredicate::UNO, rhs, rhs);
-+          Value bothNan = arith::AndIOp::create(*b, loc, lhsNan, rhsNan);
-+          return arith::OrIOp::create(*b, loc, oeq, bothNan);
-+        }
-+        case stablehlo::ComparisonDirection::NE: {
-+          Value une = arith::CmpFOp::create(*b, loc, arith::CmpFPredicate::UNE,
-+                                            lhs, rhs);
-+          Value lhsOrd = arith::CmpFOp::create(
-+              *b, loc, arith::CmpFPredicate::ORD, lhs, lhs);
-+          Value rhsOrd = arith::CmpFOp::create(
-+              *b, loc, arith::CmpFPredicate::ORD, rhs, rhs);
-+          Value anyOrd = arith::OrIOp::create(*b, loc, lhsOrd, rhsOrd);
-+          return arith::AndIOp::create(*b, loc, une, anyOrd);
-+        }
-+      }
-+    }
-     std::optional<arith::CmpFPredicate> predicate =
-         getCmpPredicate<arith::CmpFPredicate>(comparisonDirection,
-                                               /*is_signed=*/true);
-@@ -670,6 +735,46 @@
-       loc, resultTypes, argTypes, adaptor.getOperands(), b);
- }
- 
-+template <>
-+inline Value mapStablehloOpToStdScalarOp<stablehlo::Exp2Op>(
-+    Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
-+    stablehlo::Exp2Op::Adaptor adaptor, OpBuilder* b) {
-+  Type type = adaptor.getOperand().getType();
-+  if (!isa<ComplexType>(type))
-+    return MapStablehloOpToScalarOpImpl<IsFloatType, math::Exp2Op>{}(
-+        loc, resultTypes, argTypes, adaptor.getOperands(), b);
-+  auto complexTy = cast<ComplexType>(type);
-+  auto floatTy = complexTy.getElementType();
-+  Value ln2 =
-+      arith::ConstantOp::create(*b, loc, b->getFloatAttr(floatTy, M_LN2));
-+  Value re = complex::ReOp::create(*b, loc, floatTy, adaptor.getOperand());
-+  Value im = complex::ImOp::create(*b, loc, floatTy, adaptor.getOperand());
-+  Value scaled = complex::CreateOp::create(
-+      *b, loc, complexTy, arith::MulFOp::create(*b, loc, re, ln2),
-+      arith::MulFOp::create(*b, loc, im, ln2));
-+  return complex::ExpOp::create(*b, loc, scaled);
-+}
-+
-+template <>
-+inline Value mapStablehloOpToStdScalarOp<stablehlo::Log2Op>(
-+    Location loc, ArrayRef<Type> resultTypes, ArrayRef<Type> argTypes,
-+    stablehlo::Log2Op::Adaptor adaptor, OpBuilder* b) {
-+  Type type = adaptor.getOperand().getType();
-+  if (!isa<ComplexType>(type))
-+    return MapStablehloOpToScalarOpImpl<IsFloatType, math::Log2Op>{}(
-+        loc, resultTypes, argTypes, adaptor.getOperands(), b);
-+  auto complexTy = cast<ComplexType>(type);
-+  auto floatTy = complexTy.getElementType();
-+  Value logZ = complex::LogOp::create(*b, loc, adaptor.getOperand());
-+  Value oneOverLn2 =
-+      arith::ConstantOp::create(*b, loc, b->getFloatAttr(floatTy, M_LOG2E));
-+  Value re = complex::ReOp::create(*b, loc, floatTy, logZ);
-+  Value im = complex::ImOp::create(*b, loc, floatTy, logZ);
-+  return complex::CreateOp::create(
-+      *b, loc, complexTy, arith::MulFOp::create(*b, loc, re, oneOverLn2),
-+      arith::MulFOp::create(*b, loc, im, oneOverLn2));
-+}
-+
- // 'target_types' is the unconverted type (signed or unsigned if integer),
- // 'ResultTypes' is the converted type (signless if integer).
- inline Value mapConvertOpToStdScalarOp(Location loc, ArrayRef<Type> targetTypes,
-@@ -688,32 +793,39 @@
-   // A boolean value is considered to be unsigned when converting to
-   // floating-point. Otherwise, it will become `-1`.
-   if (IsUnsignedIntegerType{}(sourceType) &&
--      mlir::arith::UIToFPOp::areCastCompatible(convertedSourceType,
--                                               targetType)) {
--    return mlir::arith::UIToFPOp::create(*b, loc, resultTypes, args,
--                                         ArrayRef<NamedAttribute>());
-+      arith::UIToFPOp::areCastCompatible(convertedSourceType, targetType)) {
-+    return arith::UIToFPOp::create(*b, loc, resultTypes, args,
-+                                   typename arith::UIToFPOp::Properties{},
-+                                   /*discardableAttributes=*/{});
-   }
--  if (mlir::arith::SIToFPOp::areCastCompatible(sourceType, targetType)) {
--    return mlir::arith::SIToFPOp::create(*b, loc, resultTypes, args,
--                                         ArrayRef<NamedAttribute>());
-+  if (arith::SIToFPOp::areCastCompatible(sourceType, targetType)) {
-+    return arith::SIToFPOp::create(*b, loc, resultTypes, args,
-+                                   typename arith::SIToFPOp::Properties{},
-+                                   /*discardableAttributes=*/{});
-   }
-   if (isa<FloatType>(sourceType) && isa<FloatType>(targetType)) {
-     auto src = cast<FloatType>(sourceType);
-     auto res = cast<FloatType>(targetType);
-     if (src.getWidth() > res.getWidth()) {
--      return mlir::arith::TruncFOp::create(*b, loc, resultTypes, args,
--                                           ArrayRef<NamedAttribute>());
-+      return arith::TruncFOp::create(*b, loc, resultTypes, args,
-+                                     typename arith::TruncFOp::Properties{},
-+                                     /*discardableAttributes=*/{});
-     }
-     if (src.getWidth() < res.getWidth()) {
--      return mlir::arith::ExtFOp::create(*b, loc, resultTypes, args,
--                                         ArrayRef<NamedAttribute>());
-+      return arith::ExtFOp::create(*b, loc, resultTypes, args,
-+                                   typename arith::ExtFOp::Properties{},
-+                                   /*discardableAttributes=*/{});
-     }
-     // There's no direct conversion between different 16 bit floating point
-     // types, so go through 32 bit float.
-     if (sourceType != targetType) {
-       assert(sourceType.isBF16() || targetType.isBF16());
--      Value ext = arith::ExtFOp::create(*b, loc, b->getF32Type(), args);
--      return arith::TruncFOp::create(*b, loc, resultTypes, ext);
-+      Value ext = arith::ExtFOp::create(*b, loc, b->getF32Type(), args,
-+                                        typename arith::ExtFOp::Properties{},
-+                                        /*discardableAttributes=*/{});
-+      return arith::TruncFOp::create(*b, loc, resultTypes, ext,
-+                                     typename arith::TruncFOp::Properties{},
-+                                     /*discardableAttributes=*/{});
-     }
-     // No conversion is needed for identical float types.
-     return args.front();
-@@ -738,17 +850,20 @@
-     auto src = cast<IntegerType>(sourceType);
-     auto res = cast<IntegerType>(targetType);
-     if (src.getWidth() > res.getWidth()) {
--      return mlir::arith::TruncIOp::create(*b, loc, resultTypes, args,
--                                           ArrayRef<NamedAttribute>());
-+      return arith::TruncIOp::create(*b, loc, resultTypes, args,
-+                                     typename arith::TruncIOp::Properties{},
-+                                     /*discardableAttributes=*/{});
-     }
-     if (src.getWidth() < res.getWidth()) {
-       // Special case boolean values, so they get casted to `1` instead of `-1`.
-       if (IsUnsignedIntegerType{}(src)) {
--        return mlir::arith::ExtUIOp::create(*b, loc, resultTypes, args,
--                                            ArrayRef<NamedAttribute>());
-+        return arith::ExtUIOp::create(*b, loc, resultTypes, args,
-+                                      typename arith::ExtUIOp::Properties{},
-+                                      /*discardableAttributes=*/{});
-       }
--      return mlir::arith::ExtSIOp::create(*b, loc, resultTypes, args,
--                                          ArrayRef<NamedAttribute>());
-+      return arith::ExtSIOp::create(*b, loc, resultTypes, args,
-+                                    typename arith::ExtSIOp::Properties{},
-+                                    /*discardableAttributes=*/{});
-     }
-     // No conversion is needed for the same width integers
-     return args.front();
-diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/StablehloToArith.cpp b/stablehlo/stablehlo/conversions/linalg/transforms/StablehloToArith.cpp
---- stablehlo/stablehlo/conversions/linalg/transforms/StablehloToArith.cpp
-+++ stablehlo/stablehlo/conversions/linalg/transforms/StablehloToArith.cpp
-@@ -119,10 +119,12 @@
-       ScalarHloToArithmeticPattern<mlir::stablehlo::DivOp>,
-       ScalarHloToArithmeticPattern<mlir::stablehlo::ExpOp>,
-       ScalarHloToArithmeticPattern<mlir::stablehlo::Expm1Op>,
-+      ScalarHloToArithmeticPattern<mlir::stablehlo::Exp2Op>,
-       ScalarHloToArithmeticPattern<mlir::stablehlo::FloorOp>,
-       ScalarHloToArithmeticPattern<mlir::stablehlo::ImagOp>,
-       ScalarHloToArithmeticPattern<mlir::stablehlo::IsFiniteOp>,
-       ScalarHloToArithmeticPattern<mlir::stablehlo::Log1pOp>,
-+      ScalarHloToArithmeticPattern<mlir::stablehlo::Log2Op>,
-       ScalarHloToArithmeticPattern<mlir::stablehlo::LogOp>,
-       ScalarHloToArithmeticPattern<mlir::stablehlo::LogisticOp>,
-       ScalarHloToArithmeticPattern<mlir::stablehlo::MaxOp>,
-diff --ruN a/stablehlo/stablehlo/conversions/linalg/transforms/StablehloToLinalgPointwise.cpp b/stablehlo/stablehlo/conversions/linalg/transforms/StablehloToLinalgPointwise.cpp
---- stablehlo/stablehlo/conversions/linalg/transforms/StablehloToLinalgPointwise.cpp
-+++ stablehlo/stablehlo/conversions/linalg/transforms/StablehloToLinalgPointwise.cpp
-@@ -285,10 +285,12 @@
-         PointwiseToLinalgMapConverter<mlir::stablehlo::DivOp>,
-         PointwiseToLinalgMapConverter<mlir::stablehlo::ExpOp>,
-         PointwiseToLinalgMapConverter<mlir::stablehlo::Expm1Op>,
-+        PointwiseToLinalgMapConverter<mlir::stablehlo::Exp2Op>,
-         PointwiseToLinalgMapConverter<mlir::stablehlo::FloorOp>,
-         PointwiseToLinalgMapConverter<mlir::stablehlo::ImagOp>,
-         PointwiseToLinalgMapConverter<mlir::stablehlo::IsFiniteOp>,
-         PointwiseToLinalgMapConverter<mlir::stablehlo::Log1pOp>,
-+        PointwiseToLinalgMapConverter<mlir::stablehlo::Log2Op>,
-         PointwiseToLinalgMapConverter<mlir::stablehlo::LogOp>,
-         PointwiseToLinalgMapConverter<mlir::stablehlo::LogisticOp>,
-         PointwiseToLinalgMapConverter<mlir::stablehlo::MaxOp>,
-@@ -336,10 +338,12 @@
-             PointwiseToLinalgConverter<mlir::stablehlo::DivOp>,
-             PointwiseToLinalgConverter<mlir::stablehlo::ExpOp>,
-             PointwiseToLinalgConverter<mlir::stablehlo::Expm1Op>,
-+            PointwiseToLinalgConverter<mlir::stablehlo::Exp2Op>,
-             PointwiseToLinalgConverter<mlir::stablehlo::FloorOp>,
-             PointwiseToLinalgConverter<mlir::stablehlo::ImagOp>,
-             PointwiseToLinalgConverter<mlir::stablehlo::IsFiniteOp>,
-             PointwiseToLinalgConverter<mlir::stablehlo::Log1pOp>,
-+            PointwiseToLinalgConverter<mlir::stablehlo::Log2Op>,
-             PointwiseToLinalgConverter<mlir::stablehlo::LogOp>,
-             PointwiseToLinalgConverter<mlir::stablehlo::LogisticOp>,
-             PointwiseToLinalgConverter<mlir::stablehlo::MaxOp>,
-diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/binary.mlir
---- stablehlo/stablehlo/conversions/tosa/tests/binary.mlir
-+++ stablehlo/stablehlo/conversions/tosa/tests/binary.mlir
-@@ -38,7 +38,7 @@
- 
- // CHECK-LABEL: @concatenate
- func.func @concatenate(%arg0 : tensor<3x3xf32>, %arg1 : tensor<3x3xf32>) -> tensor<6x3xf32> {
--  // CHECK: tosa.concat %arg0, %arg1 {axis = 0 : i32} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<6x3xf32>
-+  // CHECK: tosa.concat %arg0, %arg1 axis(0) : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<6x3xf32>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1) {dimension = 0 : i64} : (tensor<3x3xf32>, tensor<3x3xf32>) -> tensor<6x3xf32>
-   return %0 : tensor<6x3xf32>
- }
-@@ -52,7 +52,7 @@
- 
- // CHECK-LABEL: @divide_f32
- func.func @divide_f32(%arg0 : tensor<10xf32>, %arg1 : tensor<10xf32>) -> tensor<10xf32> {
--  // CHECK-DAG: %[[VAR0:.*]] = "tosa.const"() <{values = dense<0> : tensor<1xi8>}
-+  // CHECK-DAG: %[[VAR0:.*]] = tosa.const values(dense<0> : tensor<1xi8>)
-   // CHECK-DAG: %[[VAR1:.*]] = tosa.reciprocal %arg1
-   // CHECK: tosa.mul %arg0, %[[VAR1]], %[[VAR0]]
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
-@@ -61,10 +61,10 @@
- 
- // CHECK-LABEL: @dot_vector_vector
- func.func @dot_vector_vector(%arg0 : tensor<3xf32>, %arg1 : tensor<3xf32>) -> tensor<f32> {
--  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<> : tensor<0xindex>} : () -> !tosa.shape<0>
--  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 1, 3]> : tensor<3xindex>} : () -> !tosa.shape<3>
-+  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<> : tensor<0xindex>) : () -> !tosa.shape<0>
-+  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 1, 3]> : tensor<3xindex>) : () -> !tosa.shape<3>
-   // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]]
--  // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 1]> : tensor<3xindex>} : () -> !tosa.shape<3>
-+  // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 1]> : tensor<3xindex>) : () -> !tosa.shape<3>
-   // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]]
-   // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]]
-   // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]]
-@@ -74,10 +74,10 @@
- 
- // CHECK-LABEL: @dot_vector_matrix
- func.func @dot_vector_matrix(%arg0 : tensor<2xf32>, %arg1 : tensor<2x3xf32>) -> tensor<3xf32> {
--  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<3> : tensor<1xindex>} : () -> !tosa.shape<1>
--  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 1, 2]> : tensor<3xindex>} : () -> !tosa.shape<3>
-+  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<3> : tensor<1xindex>) : () -> !tosa.shape<1>
-+  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 1, 2]> : tensor<3xindex>) : () -> !tosa.shape<3>
-   // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]]
--  // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3>
-+  // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3>
-   // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]]
-   // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]]
-   // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]]
-@@ -87,10 +87,10 @@
- 
- // CHECK-LABEL: @dot_matrix_vector
- func.func @dot_matrix_vector(%arg0 : tensor<2x3xf32>, %arg1 : tensor<3xf32>) -> tensor<2xf32> {
--  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<2> : tensor<1xindex>} : () -> !tosa.shape<1>
--  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3>
-+  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<2> : tensor<1xindex>) : () -> !tosa.shape<1>
-+  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3>
-   // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]]
--  // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 1]> : tensor<3xindex>} : () -> !tosa.shape<3>
-+  // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 1]> : tensor<3xindex>) : () -> !tosa.shape<3>
-   // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]]
-   // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]]
-   // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]]
-@@ -100,10 +100,10 @@
- 
- // CHECK-LABEL: @dot_matrix_matrix
- func.func @dot_matrix_matrix(%arg0 : tensor<2x3xf32>, %arg1 : tensor<3x4xf32>) -> tensor<2x4xf32> {
--  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<[2, 4]> : tensor<2xindex>} : () -> !tosa.shape<2>
--  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3>
-+  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<[2, 4]> : tensor<2xindex>) : () -> !tosa.shape<2>
-+  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3>
-   // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]]
--  // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 4]> : tensor<3xindex>} : () -> !tosa.shape<3>
-+  // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 4]> : tensor<3xindex>) : () -> !tosa.shape<3>
-   // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]]
-   // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]]
-   // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]]
-@@ -113,10 +113,10 @@
- 
- // CHECK-LABEL: @dot_general_vector_vector
- func.func @dot_general_vector_vector(%arg0: tensor<3xf32>, %arg1: tensor<3xf32>) -> tensor<f32> {
--  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<> : tensor<0xindex>} : () -> !tosa.shape<0>
--  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 1, 3]> : tensor<3xindex>} : () -> !tosa.shape<3>
-+  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<> : tensor<0xindex>) : () -> !tosa.shape<0>
-+  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 1, 3]> : tensor<3xindex>) : () -> !tosa.shape<3>
-   // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]]
--  // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 1]> : tensor<3xindex>} : () -> !tosa.shape<3>
-+  // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 1]> : tensor<3xindex>) : () -> !tosa.shape<3>
-   // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]]
-   // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]]
-   // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]]
-@@ -126,10 +126,10 @@
- 
- // CHECK-LABEL: @dot_general_vector_matrix
- func.func @dot_general_vector_matrix(%arg0: tensor<2xf32>, %arg1: tensor<2x3xf32>) -> tensor<3xf32> {
--  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<3> : tensor<1xindex>} : () -> !tosa.shape<1>
--  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 1, 2]> : tensor<3xindex>} : () -> !tosa.shape<3>
-+  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<3> : tensor<1xindex>) : () -> !tosa.shape<1>
-+  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 1, 2]> : tensor<3xindex>) : () -> !tosa.shape<3>
-   // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]]
--  // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3>
-+  // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3>
-   // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]]
-   // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]]
-   // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]]
-@@ -139,10 +139,10 @@
- 
- // CHECK-LABEL: @dot_general_matrix_vector
- func.func @dot_general_matrix_vector(%arg0: tensor<2x3xf32>, %arg1: tensor<3xf32>) -> tensor<2xf32> {
--  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<2> : tensor<1xindex>} : () -> !tosa.shape<1>
--  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3>
-+  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<2> : tensor<1xindex>) : () -> !tosa.shape<1>
-+  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3>
-   // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]]
--  // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 1]> : tensor<3xindex>} : () -> !tosa.shape<3>
-+  // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 1]> : tensor<3xindex>) : () -> !tosa.shape<3>
-   // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]]
-   // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]]
-   // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]]
-@@ -152,10 +152,10 @@
- 
- // CHECK-LABEL: @dot_general_matrix_matrix
- func.func @dot_general_matrix_matrix(%arg0: tensor<2x3xf32>, %arg1: tensor<3x4xf32>) -> tensor<2x4xf32> {
--  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<[2, 4]> : tensor<2xindex>} : () -> !tosa.shape<2>
--  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape {values = dense<[1, 2, 3]> : tensor<3xindex>} : () -> !tosa.shape<3>
-+  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<[2, 4]> : tensor<2xindex>) : () -> !tosa.shape<2>
-+  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 2, 3]> : tensor<3xindex>) : () -> !tosa.shape<3>
-   // CHECK-DAG: %[[VAR2:.*]] = tosa.reshape %arg0, %[[VAR1]]
--  // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape {values = dense<[1, 3, 4]> : tensor<3xindex>} : () -> !tosa.shape<3>
-+  // CHECK-DAG: %[[VAR3:.*]] = tosa.const_shape values(dense<[1, 3, 4]> : tensor<3xindex>) : () -> !tosa.shape<3>
-   // CHECK-DAG: %[[VAR4:.*]] = tosa.reshape %arg1, %[[VAR3]]
-   // CHECK-DAG: %[[VAR5:.*]] = tosa.matmul %[[VAR2]], %[[VAR4]]
-   // CHECK-DAG: %[[VAR6:.*]] = tosa.reshape %[[VAR5]], %[[VAR0]]
-diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir b/stablehlo/stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir
---- stablehlo/stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir
-+++ stablehlo/stablehlo/conversions/tosa/tests/legalize_quant_ops_to_tosa_rescale.mlir
-@@ -4,17 +4,17 @@
- // CHECK-LABEL: @add
- func.func @add(%arg0 : tensor<2x2x!quant.uniform<i8:f32, 0.025:-1>>,
-                %arg1 : tensor<2x2x!quant.uniform<i8:f32, 0.075:-1>>) -> tensor<2x2x!quant.uniform<i8:f32, 1.5e-01:-1>> {
--  // CHECK-DAG: %[[SHIFT50:.+]] = "tosa.const"() <{values = dense<50> : tensor<1xi8>}>
--  // CHECK-DAG: %[[SHIFT11:.+]] = "tosa.const"() <{values = dense<11> : tensor<1xi8>}>
--  // CHECK-DAG: %[[SHIFT13:.+]] = "tosa.const"() <{values = dense<13> : tensor<1xi8>}>
--  // CHECK-DAG: %[[MULTIPLIER_1:.+]] = "tosa.const"() <{values = dense<1073741824> : tensor<1xi32>}>
--  // CHECK-DAG: %[[MULTIPLIER_2:.+]] = "tosa.const"() <{values = dense<1431655765> : tensor<1xi32>}>
--  // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}>
--  // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}>
--  // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT13]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
--  // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT11]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
-+  // CHECK-DAG: %[[SHIFT50:.+]] = tosa.const values(dense<50> : tensor<1xi8>)
-+  // CHECK-DAG: %[[SHIFT11:.+]] = tosa.const values(dense<11> : tensor<1xi8>)
-+  // CHECK-DAG: %[[SHIFT13:.+]] = tosa.const values(dense<13> : tensor<1xi8>)
-+  // CHECK-DAG: %[[MULTIPLIER_1:.+]] = tosa.const values(dense<1073741824> : tensor<1xi32>)
-+  // CHECK-DAG: %[[MULTIPLIER_2:.+]] = tosa.const values(dense<1431655765> : tensor<1xi32>)
-+  // CHECK-DAG: %[[ZP_MINUS_1:.+]] = tosa.const values(dense<-1> : tensor<1xi8>)
-+  // CHECK-DAG: %[[ZP_0:.+]] = tosa.const values(dense<0> : tensor<1xi32>)
-+  // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT13]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-+  // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT11]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-   // CHECK: %[[V2:.+]] = stablehlo.add %[[V0]], %[[V1]] : tensor<2x2xi32>
--  // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT50]], %[[ZP_0]], %[[ZP_MINUS_1]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
-+  // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT50]], %[[ZP_0]], %[[ZP_MINUS_1]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-   // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform<i8:f32, 1.500000e-01:-1>>
-   %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<2x2x!quant.uniform<i8:f32, 0.025:-1>>, tensor<2x2x!quant.uniform<i8:f32, 0.075:-1>>)
-             -> tensor<2x2x!quant.uniform<i8:f32, 1.5e-01:-1>>
-@@ -25,17 +25,17 @@
- // CHECK-LABEL: @sub
- func.func @sub(%arg0 : tensor<2x2x!quant.uniform<i8:f32, 0.025:-1>>,
-                %arg1 : tensor<2x2x!quant.uniform<i8:f32, 0.075:-1>>) -> tensor<2x2x!quant.uniform<i8:f32, 1.5e-01:-1>> {
--  // CHECK-DAG: %[[SHIFT50:.+]] = "tosa.const"() <{values = dense<50> : tensor<1xi8>}>
--  // CHECK-DAG: %[[SHIFT11:.+]] = "tosa.const"() <{values = dense<11> : tensor<1xi8>}>
--  // CHECK-DAG: %[[SHIFT13:.+]] = "tosa.const"() <{values = dense<13> : tensor<1xi8>}>
--  // CHECK-DAG: %[[MULTIPLIER_1:.+]] = "tosa.const"() <{values = dense<1073741824> : tensor<1xi32>}>
--  // CHECK-DAG: %[[MULTIPLIER_2:.+]] = "tosa.const"() <{values = dense<1431655765> : tensor<1xi32>}>
--  // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}>
--  // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}>
--  // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT13]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
--  // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT11]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
-+  // CHECK-DAG: %[[SHIFT50:.+]] = tosa.const values(dense<50> : tensor<1xi8>)
-+  // CHECK-DAG: %[[SHIFT11:.+]] = tosa.const values(dense<11> : tensor<1xi8>)
-+  // CHECK-DAG: %[[SHIFT13:.+]] = tosa.const values(dense<13> : tensor<1xi8>)
-+  // CHECK-DAG: %[[MULTIPLIER_1:.+]] = tosa.const values(dense<1073741824> : tensor<1xi32>)
-+  // CHECK-DAG: %[[MULTIPLIER_2:.+]] = tosa.const values(dense<1431655765> : tensor<1xi32>)
-+  // CHECK-DAG: %[[ZP_MINUS_1:.+]] = tosa.const values(dense<-1> : tensor<1xi8>)
-+  // CHECK-DAG: %[[ZP_0:.+]] = tosa.const values(dense<0> : tensor<1xi32>)
-+  // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT13]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-+  // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT11]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-   // CHECK: %[[V2:.+]] = stablehlo.subtract %[[V0]], %[[V1]] : tensor<2x2xi32>
--  // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT50]], %[[ZP_0]], %[[ZP_MINUS_1]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
-+  // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT50]], %[[ZP_0]], %[[ZP_MINUS_1]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-   // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform<i8:f32, 1.500000e-01:-1>>
-   %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<2x2x!quant.uniform<i8:f32, 0.025:-1>>, tensor<2x2x!quant.uniform<i8:f32, 0.075:-1>>)
-             -> tensor<2x2x!quant.uniform<i8:f32, 1.5e-01:-1>>
-@@ -46,16 +46,16 @@
- // CHECK-LABEL: @mul
- func.func @mul(%arg0 : tensor<2x2x!quant.uniform<i8:f32, 0.025:-1>>,
-                %arg1 : tensor<2x2x!quant.uniform<i8:f32, 0.075:-1>>) -> tensor<2x2x!quant.uniform<i8:f32, 1.5e-01:-1>> {
--  // CHECK-DAG: %[[SHIFT37:.+]] = "tosa.const"() <{values = dense<37> : tensor<1xi8>}>
--  // CHECK-DAG: %[[SHIFT30:.+]] = "tosa.const"() <{values = dense<30> : tensor<1xi8>}>
--  // CHECK-DAG: %[[MULTIPLIER_1:.+]] = "tosa.const"() <{values = dense<1073741824> : tensor<1xi32>}>
--  // CHECK-DAG: %[[MULTIPLIER_2:.+]] = "tosa.const"() <{values = dense<1717986918> : tensor<1xi32>}>
--  // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}>
--  // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}>
--  // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
--  // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
-+  // CHECK-DAG: %[[SHIFT37:.+]] = tosa.const values(dense<37> : tensor<1xi8>)
-+  // CHECK-DAG: %[[SHIFT30:.+]] = tosa.const values(dense<30> : tensor<1xi8>)
-+  // CHECK-DAG: %[[MULTIPLIER_1:.+]] = tosa.const values(dense<1073741824> : tensor<1xi32>)
-+  // CHECK-DAG: %[[MULTIPLIER_2:.+]] = tosa.const values(dense<1717986918> : tensor<1xi32>)
-+  // CHECK-DAG: %[[ZP_MINUS_1:.+]] = tosa.const values(dense<-1> : tensor<1xi8>)
-+  // CHECK-DAG: %[[ZP_0:.+]] = tosa.const values(dense<0> : tensor<1xi32>)
-+  // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-+  // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-   // CHECK: %[[V2:.+]] = stablehlo.multiply %[[V0]], %[[V1]] : tensor<2x2xi32>
--  // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_2]], %[[SHIFT37]], %[[ZP_0]], %[[ZP_MINUS_1]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
-+  // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_2]], %[[SHIFT37]], %[[ZP_0]], %[[ZP_MINUS_1]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-   // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform<i8:f32, 1.500000e-01:-1>>
-   %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<2x2x!quant.uniform<i8:f32, 0.025:-1>>, tensor<2x2x!quant.uniform<i8:f32, 0.075:-1>>)
-             -> tensor<2x2x!quant.uniform<i8:f32, 1.5e-01:-1>>
-@@ -66,18 +66,18 @@
- // CHECK-LABEL: @div
- func.func @div(%arg0 : tensor<2x2x!quant.uniform<i8:f32, 0.025:-1>>,
-                %arg1 : tensor<2x2x!quant.uniform<i8:f32, 0.075:-2>>) -> tensor<2x2x!quant.uniform<i8:f32, 1.5e-01:-3>> {
--  // CHECK-DAG: %[[SHIFT37:.+]] = "tosa.const"() <{values = dense<37> : tensor<1xi8>}>
--  // CHECK-DAG: %[[SHIFT30:.+]] = "tosa.const"() <{values = dense<30> : tensor<1xi8>}>
--  // CHECK-DAG: %[[MULTIPLIER_1:.+]] = "tosa.const"() <{values = dense<1073741824> : tensor<1xi32>}>
--  // CHECK-DAG: %[[MULTIPLIER_2:.+]] = "tosa.const"() <{values = dense<1717986918> : tensor<1xi32>}>
--  // CHECK-DAG: %[[ZP_MINUS_3:.+]] = "tosa.const"() <{values = dense<-3> : tensor<1xi8>}>
--  // CHECK-DAG: %[[ZP_MINUS_2:.+]] = "tosa.const"() <{values = dense<-2> : tensor<1xi8>}>
--  // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}>
--  // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}>
--  // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
--  // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_2]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
-+  // CHECK-DAG: %[[SHIFT37:.+]] = tosa.const values(dense<37> : tensor<1xi8>)
-+  // CHECK-DAG: %[[SHIFT30:.+]] = tosa.const values(dense<30> : tensor<1xi8>)
-+  // CHECK-DAG: %[[MULTIPLIER_1:.+]] = tosa.const values(dense<1073741824> : tensor<1xi32>)
-+  // CHECK-DAG: %[[MULTIPLIER_2:.+]] = tosa.const values(dense<1717986918> : tensor<1xi32>)
-+  // CHECK-DAG: %[[ZP_MINUS_3:.+]] = tosa.const values(dense<-3> : tensor<1xi8>)
-+  // CHECK-DAG: %[[ZP_MINUS_2:.+]] = tosa.const values(dense<-2> : tensor<1xi8>)
-+  // CHECK-DAG: %[[ZP_MINUS_1:.+]] = tosa.const values(dense<-1> : tensor<1xi8>)
-+  // CHECK-DAG: %[[ZP_0:.+]] = tosa.const values(dense<0> : tensor<1xi32>)
-+  // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-+  // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT30]], %[[ZP_MINUS_2]], %[[ZP_0]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-   // CHECK: %[[V2:.+]] = stablehlo.divide %[[V0]], %[[V1]] : tensor<2x2xi32>
--  // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_2]], %[[SHIFT37]], %[[ZP_0]], %[[ZP_MINUS_3]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
-+  // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_2]], %[[SHIFT37]], %[[ZP_0]], %[[ZP_MINUS_3]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-   // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform<i8:f32, 1.500000e-01:-3>>
-   %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<2x2x!quant.uniform<i8:f32, 0.025:-1>>, tensor<2x2x!quant.uniform<i8:f32, 0.075:-2>>)
-             -> tensor<2x2x!quant.uniform<i8:f32, 1.5e-01:-3>>
-@@ -88,19 +88,19 @@
- // CHECK-LABEL: @max
- func.func @max(%arg0 : tensor<2x2x!quant.uniform<i8:f32, 0.025:-1>>,
-                %arg1 : tensor<2x2x!quant.uniform<i8:f32, 0.075:-2>>) -> tensor<2x2x!quant.uniform<i8:f32, 1.5e-01:-3>> {
--  // CHECK-DAG: %[[ZP_MINUS_3:.+]] = "tosa.const"() <{values = dense<-3> : tensor<1xi8>}>
--  // CHECK-DAG: %[[SHIFT51:.+]] = "tosa.const"() <{values = dense<51> : tensor<1xi8>}>
--  // CHECK-DAG: %[[ZP_MINUS_2:.+]] = "tosa.const"() <{values = dense<-2> : tensor<1xi8>}>
--  // CHECK-DAG: %[[SHIFT10:.+]] = "tosa.const"() <{values = dense<10> : tensor<1xi8>}>
--  // CHECK-DAG: %[[MULTIPLIER_1:.+]] = "tosa.const"() <{values = dense<1073741824> : tensor<1xi32>}>
--  // CHECK-DAG: %[[MULTIPLIER_2:.+]] = "tosa.const"() <{values = dense<1431655765> : tensor<1xi32>}>
--  // CHECK-DAG: %[[SHIFT12:.+]] = "tosa.const"() <{values = dense<12> : tensor<1xi8>}>
--  // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}>
--  // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}>
--  // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
--  // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
-+  // CHECK-DAG: %[[ZP_MINUS_3:.+]] = tosa.const values(dense<-3> : tensor<1xi8>)
-+  // CHECK-DAG: %[[SHIFT51:.+]] = tosa.const values(dense<51> : tensor<1xi8>)
-+  // CHECK-DAG: %[[ZP_MINUS_2:.+]] = tosa.const values(dense<-2> : tensor<1xi8>)
-+  // CHECK-DAG: %[[SHIFT10:.+]] = tosa.const values(dense<10> : tensor<1xi8>)
-+  // CHECK-DAG: %[[MULTIPLIER_1:.+]] = tosa.const values(dense<1073741824> : tensor<1xi32>)
-+  // CHECK-DAG: %[[MULTIPLIER_2:.+]] = tosa.const values(dense<1431655765> : tensor<1xi32>)
-+  // CHECK-DAG: %[[SHIFT12:.+]] = tosa.const values(dense<12> : tensor<1xi8>)
-+  // CHECK-DAG: %[[ZP_MINUS_1:.+]] = tosa.const values(dense<-1> : tensor<1xi8>)
-+  // CHECK-DAG: %[[ZP_0:.+]] = tosa.const values(dense<0> : tensor<1xi32>)
-+  // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-+  // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-   // CHECK: %[[V2:.+]] = stablehlo.maximum %[[V0]], %[[V1]] : tensor<2x2xi32>
--  // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT51]], %[[ZP_0]], %[[ZP_MINUS_3]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
-+  // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT51]], %[[ZP_0]], %[[ZP_MINUS_3]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-   // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform<i8:f32, 1.500000e-01:-3>>
-   %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<2x2x!quant.uniform<i8:f32, 0.025:-1>>, tensor<2x2x!quant.uniform<i8:f32, 0.075:-2>>)
-             -> tensor<2x2x!quant.uniform<i8:f32, 1.5e-01:-3>>
-@@ -111,19 +111,19 @@
- // CHECK-LABEL: @min
- func.func @min(%arg0 : tensor<2x2x!quant.uniform<i8:f32, 0.025:-1>>,
-                %arg1 : tensor<2x2x!quant.uniform<i8:f32, 0.075:-2>>) -> tensor<2x2x!quant.uniform<i8:f32, 1.5e-01:-3>> {
--  // CHECK-DAG: %[[ZP_MINUS_3:.+]] = "tosa.const"() <{values = dense<-3> : tensor<1xi8>}>
--  // CHECK-DAG: %[[SHIFT51:.+]] = "tosa.const"() <{values = dense<51> : tensor<1xi8>}>
--  // CHECK-DAG: %[[ZP_MINUS_2:.+]] = "tosa.const"() <{values = dense<-2> : tensor<1xi8>}>
--  // CHECK-DAG: %[[SHIFT10:.+]] = "tosa.const"() <{values = dense<10> : tensor<1xi8>}>
--  // CHECK-DAG: %[[MULTIPLIER_1:.+]] = "tosa.const"() <{values = dense<1073741824> : tensor<1xi32>}>
--  // CHECK-DAG: %[[MULTIPLIER_2:.+]] = "tosa.const"() <{values = dense<1431655765> : tensor<1xi32>}>
--  // CHECK-DAG: %[[SHIFT12:.+]] = "tosa.const"() <{values = dense<12> : tensor<1xi8>}>
--  // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}>
--  // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}>
--  // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
--  // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
-+  // CHECK-DAG: %[[ZP_MINUS_3:.+]] = tosa.const values(dense<-3> : tensor<1xi8>)
-+  // CHECK-DAG: %[[SHIFT51:.+]] = tosa.const values(dense<51> : tensor<1xi8>)
-+  // CHECK-DAG: %[[ZP_MINUS_2:.+]] = tosa.const values(dense<-2> : tensor<1xi8>)
-+  // CHECK-DAG: %[[SHIFT10:.+]] = tosa.const values(dense<10> : tensor<1xi8>)
-+  // CHECK-DAG: %[[MULTIPLIER_1:.+]] = tosa.const values(dense<1073741824> : tensor<1xi32>)
-+  // CHECK-DAG: %[[MULTIPLIER_2:.+]] = tosa.const values(dense<1431655765> : tensor<1xi32>)
-+  // CHECK-DAG: %[[SHIFT12:.+]] = tosa.const values(dense<12> : tensor<1xi8>)
-+  // CHECK-DAG: %[[ZP_MINUS_1:.+]] = tosa.const values(dense<-1> : tensor<1xi8>)
-+  // CHECK-DAG: %[[ZP_0:.+]] = tosa.const values(dense<0> : tensor<1xi32>)
-+  // CHECK-DAG: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-+  // CHECK-DAG: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-   // CHECK: %[[V2:.+]] = stablehlo.minimum %[[V0]], %[[V1]] : tensor<2x2xi32>
--  // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT51]], %[[ZP_0]], %[[ZP_MINUS_3]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
-+  // CHECK: %[[V3:.+]] = tosa.rescale %[[V2]], %[[MULTIPLIER_1]], %[[SHIFT51]], %[[ZP_0]], %[[ZP_MINUS_3]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-   // CHECK: return %[[V3]] : tensor<2x2x!quant.uniform<i8:f32, 1.500000e-01:-3>>
-   %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<2x2x!quant.uniform<i8:f32, 0.025:-1>>, tensor<2x2x!quant.uniform<i8:f32, 0.075:-2>>)
-             -> tensor<2x2x!quant.uniform<i8:f32, 1.5e-01:-3>>
-@@ -133,16 +133,16 @@
- // -----
- // CHECK-LABEL: @abs
- func.func @abs(%arg0 : tensor<20x20x!quant.uniform<i8:f32, 0.025:-1>>) -> tensor<20x20x!quant.uniform<i8:f32, 1.5e-01:-128>> {
--  // CHECK-DAG: %[[ZP_MINUS_128:.+]] = "tosa.const"() <{values = dense<-128> : tensor<1xi8>}>
--  // CHECK-DAG: %[[SHIFT33:.+]] = "tosa.const"() <{values = dense<33> : tensor<1xi8>}>
--  // CHECK-DAG: %[[MULTIPLIER_1:.+]] = "tosa.const"() <{values = dense<1431655765> : tensor<1xi32>}>
--  // CHECK-DAG: %[[MULTIPLIER_2:.+]] = "tosa.const"() <{values = dense<1073741824> : tensor<1xi32>}>
--  // CHECK-DAG: %[[SHIFT30:.+]] = "tosa.const"() <{values = dense<30> : tensor<1xi8>}>
--  // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}>
--  // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}>
--  // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
-+  // CHECK-DAG: %[[ZP_MINUS_128:.+]] = tosa.const values(dense<-128> : tensor<1xi8>)
-+  // CHECK-DAG: %[[SHIFT33:.+]] = tosa.const values(dense<33> : tensor<1xi8>)
-+  // CHECK-DAG: %[[MULTIPLIER_1:.+]] = tosa.const values(dense<1431655765> : tensor<1xi32>)
-+  // CHECK-DAG: %[[MULTIPLIER_2:.+]] = tosa.const values(dense<1073741824> : tensor<1xi32>)
-+  // CHECK-DAG: %[[SHIFT30:.+]] = tosa.const values(dense<30> : tensor<1xi8>)
-+  // CHECK-DAG: %[[ZP_MINUS_1:.+]] = tosa.const values(dense<-1> : tensor<1xi8>)
-+  // CHECK-DAG: %[[ZP_0:.+]] = tosa.const values(dense<0> : tensor<1xi32>)
-+  // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT30]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-   // CHECK: %[[V1:.+]] = stablehlo.abs %[[V0]] : tensor<20x20xi32>
--  // CHECK: %[[V3:.+]] = tosa.rescale %[[V1]], %[[MULTIPLIER_1]], %[[SHIFT33]], %[[ZP_0]], %[[ZP_MINUS_128]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
-+  // CHECK: %[[V3:.+]] = tosa.rescale %[[V1]], %[[MULTIPLIER_1]], %[[SHIFT33]], %[[ZP_0]], %[[ZP_MINUS_128]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-   // CHECK: return %[[V3]] : tensor<20x20x!quant.uniform<i8:f32, 1.500000e-01:-128>>
-   %0 = "stablehlo.abs"(%arg0) : (tensor<20x20x!quant.uniform<i8:f32, 0.025:-1>>) -> tensor<20x20x!quant.uniform<i8:f32, 1.5e-01:-128>>
-   return %0 : tensor<20x20x!quant.uniform<i8:f32, 1.5e-01:-128>>
-@@ -152,15 +152,15 @@
- // CHECK-LABEL: @compareGE
- func.func @compareGE(%arg0 : tensor<20x20x!quant.uniform<i8:f32, 0.025:-1>>,
-                    %arg1 : tensor<20x20x!quant.uniform<i8:f32, 0.075:-2>>) -> tensor<20x20xi1> {
--  // CHECK-DAG: %[[ZP_MINUS_2:.+]] = "tosa.const"() <{values = dense<-2> : tensor<1xi8>}>
--  // CHECK-DAG: %[[SHIFT10:.+]] = "tosa.const"() <{values = dense<10> : tensor<1xi8>}>
--  // CHECK-DAG: %[[MULTIPLIER_1:.+]] = "tosa.const"() <{values = dense<1073741824> : tensor<1xi32>}>
--  // CHECK-DAG: %[[MULTIPLIER_2:.+]] = "tosa.const"() <{values = dense<1431655765> : tensor<1xi32>}>
--  // CHECK-DAG: %[[SHIFT12:.+]] = "tosa.const"() <{values = dense<12> : tensor<1xi8>}>
--  // CHECK-DAG: %[[ZP_MINUS_1:.+]] = "tosa.const"() <{values = dense<-1> : tensor<1xi8>}>
--  // CHECK-DAG: %[[ZP_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}>
--  // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
--  // CHECK: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
-+  // CHECK-DAG: %[[ZP_MINUS_2:.+]] = tosa.const values(dense<-2> : tensor<1xi8>)
-+  // CHECK-DAG: %[[SHIFT10:.+]] = tosa.const values(dense<10> : tensor<1xi8>)
-+  // CHECK-DAG: %[[MULTIPLIER_1:.+]] = tosa.const values(dense<1073741824> : tensor<1xi32>)
-+  // CHECK-DAG: %[[MULTIPLIER_2:.+]] = tosa.const values(dense<1431655765> : tensor<1xi32>)
-+  // CHECK-DAG: %[[SHIFT12:.+]] = tosa.const values(dense<12> : tensor<1xi8>)
-+  // CHECK-DAG: %[[ZP_MINUS_1:.+]] = tosa.const values(dense<-1> : tensor<1xi8>)
-+  // CHECK-DAG: %[[ZP_0:.+]] = tosa.const values(dense<0> : tensor<1xi32>)
-+  // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT12]], %[[ZP_MINUS_1]], %[[ZP_0]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-+  // CHECK: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT10]], %[[ZP_MINUS_2]], %[[ZP_0]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-   // CHECK: %[[V2:.+]] = stablehlo.compare GE, %[[V0]], %[[V1]], TOTALORDER :
-   // CHECK: return %[[V2]]
-   %0 = stablehlo.compare GE, %arg0, %arg1, TOTALORDER : (tensor<20x20x!quant.uniform<i8:f32, 0.025:-1>>, tensor<20x20x!quant.uniform<i8:f32, 0.075:-2>>) -> tensor<20x20xi1>
-@@ -171,14 +171,14 @@
- // CHECK-LABEL: @compareLT
- func.func @compareLT(%arg0 : tensor<20x20x!quant.uniform<i16:f32, 0.025:0>>,
-                    %arg1 : tensor<20x20x!quant.uniform<i16:f32, 0.075:0>>) -> tensor<20x20xi1> {
--  // CHECK-DAG: %[[SHIFT17:.+]] = "tosa.const"() <{values = dense<17> : tensor<1xi8>}>
--  // CHECK-DAG: %[[MULTIPLIER_1:.+]] = "tosa.const"() <{values = dense<1073741824> : tensor<1xi32>}>
--  // CHECK-DAG: %[[MULTIPLIER_2:.+]] = "tosa.const"() <{values = dense<1431655765> : tensor<1xi32>}>
--  // CHECK-DAG: %[[SHIFT15:.+]] = "tosa.const"() <{values = dense<15> : tensor<1xi8>}>
--  // CHECK-DAG: %[[ZP16_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi16>}>
--  // CHECK-DAG: %[[ZP32_0:.+]] = "tosa.const"() <{values = dense<0> : tensor<1xi32>}>
--  // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT17]], %[[ZP16_0]], %[[ZP32_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
--  // CHECK: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT15]], %[[ZP16_0]], %[[ZP32_0]] {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true}
-+  // CHECK-DAG: %[[SHIFT17:.+]] = tosa.const values(dense<17> : tensor<1xi8>)
-+  // CHECK-DAG: %[[MULTIPLIER_1:.+]] = tosa.const values(dense<1073741824> : tensor<1xi32>)
-+  // CHECK-DAG: %[[MULTIPLIER_2:.+]] = tosa.const values(dense<1431655765> : tensor<1xi32>)
-+  // CHECK-DAG: %[[SHIFT15:.+]] = tosa.const values(dense<15> : tensor<1xi8>)
-+  // CHECK-DAG: %[[ZP16_0:.+]] = tosa.const values(dense<0> : tensor<1xi16>)
-+  // CHECK-DAG: %[[ZP32_0:.+]] = tosa.const values(dense<0> : tensor<1xi32>)
-+  // CHECK: %[[V0:.+]] = tosa.rescale %arg0, %[[MULTIPLIER_2]], %[[SHIFT17]], %[[ZP16_0]], %[[ZP32_0]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-+  // CHECK: %[[V1:.+]] = tosa.rescale %arg1, %[[MULTIPLIER_1]], %[[SHIFT15]], %[[ZP16_0]], %[[ZP32_0]] scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false)
-   // CHECK: %[[V2:.+]] = stablehlo.compare LT, %[[V0]], %[[V1]], TOTALORDER :
-   // CHECK: return %[[V2]]
-   %0 = stablehlo.compare LT, %arg0, %arg1, TOTALORDER : (tensor<20x20x!quant.uniform<i16:f32, 0.025:0>>, tensor<20x20x!quant.uniform<i16:f32, 0.075:0>>) -> tensor<20x20xi1>
-diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir b/stablehlo/stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir
---- stablehlo/stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir
-+++ stablehlo/stablehlo/conversions/tosa/tests/legalize_tosa_rescale_to_stablehlo.mlir
-@@ -3,11 +3,11 @@
- // -----
- // CHECK-LABEL: @rescale
- func.func @rescale(%arg0 : tensor<2x2x!quant.uniform<i8:f32, 0.025:-1>>) -> tensor<2x2xi32> {
--  %multiplier = "tosa.const"() {values = dense<1431655765> : tensor<1xi32>} : () -> tensor<1xi32>
--  %shift = "tosa.const"() {values = dense<13> : tensor<1xi8>} : () -> tensor<1xi8>
--  %input_zp = "tosa.const"() {values = dense<-1> : tensor<1xi8>} : () -> tensor<1xi8>
--  %output_zp = "tosa.const"() {values = dense<0> : tensor<1xi32>} : () -> tensor<1xi32>
--  %0 = tosa.rescale %arg0, %multiplier, %shift, %input_zp, %output_zp {input_unsigned = false, output_unsigned = false, per_channel = false, rounding_mode = SINGLE_ROUND, scale32 = true} :
-+  %multiplier = tosa.const values(dense<1431655765> : tensor<1xi32>) : () -> tensor<1xi32>
-+  %shift = tosa.const values(dense<13> : tensor<1xi8>) : () -> tensor<1xi8>
-+  %input_zp = tosa.const values(dense<-1> : tensor<1xi8>) : () -> tensor<1xi8>
-+  %output_zp = tosa.const values(dense<0> : tensor<1xi32>) : () -> tensor<1xi32>
-+  %0 = tosa.rescale %arg0, %multiplier, %shift, %input_zp, %output_zp scale32(true) rounding_mode<SINGLE_ROUND> per_channel(false) input_unsigned(false) output_unsigned(false) :
-             (tensor<2x2x!quant.uniform<i8:f32, 0.025:-1>>, tensor<1xi32>, tensor<1xi8>, tensor<1xi8>, tensor<1xi32>) -> tensor<2x2xi32>
- 
-   // convert input quantized type to storage type
-diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir
---- stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir
-+++ stablehlo/stablehlo/conversions/tosa/tests/nullary.mlir
-@@ -16,9 +16,9 @@
- 
- // CHECK-LABEL: @iota_dimension_0
- func.func @iota_dimension_0() -> tensor<4x8xf32> {
--  // CHECK-DAG: %[[VAR0:.*]] = "tosa.const"()
--  // CHECK-SAME{LITERAL}: <{values = dense<[[0.000000e+00], [1.000000e+00], [2.000000e+00], [3.000000e+00]]> : tensor<4x1xf32>}> : () -> tensor<4x1xf32>
--  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape  {values = dense<[1, 8]> : vector<2xindex>} : () -> !tosa.shape<2>
-+  // CHECK-DAG: %[[VAR0:.*]] = tosa.const
-+  // CHECK-SAME{LITERAL}: values(dense<[[0.000000e+00], [1.000000e+00], [2.000000e+00], [3.000000e+00]]> : tensor<4x1xf32>) : () -> tensor<4x1xf32>
-+  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[1, 8]> : vector<2xindex>) : () -> !tosa.shape<2>
-   // CHECK-DAG: %[[VAR2:.*]] = tosa.tile %[[VAR0]], %[[VAR1]]
-   %0 = "stablehlo.iota"() {iota_dimension = 0 : i64} : () -> (tensor<4x8xf32>)
-   return %0 : tensor<4x8xf32>
-@@ -26,9 +26,9 @@
- 
- // CHECK-LABEL: @iota_dimension_1
- func.func @iota_dimension_1() -> tensor<4x8xi32> {
--  // CHECK-DAG: %[[VAR0:.*]] = "tosa.const"()
--  // CHECK-SAME{LITERAL}: <{values = dense<[[0, 1, 2, 3, 4, 5, 6, 7]]> : tensor<1x8xi32>}> : () -> tensor<1x8xi32>
--  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape  {values = dense<[4, 1]> : vector<2xindex>} : () -> !tosa.shape<2>
-+  // CHECK-DAG: %[[VAR0:.*]] = tosa.const
-+  // CHECK-SAME{LITERAL}: values(dense<[[0, 1, 2, 3, 4, 5, 6, 7]]> : tensor<1x8xi32>) : () -> tensor<1x8xi32>
-+  // CHECK-DAG: %[[VAR1:.*]] = tosa.const_shape values(dense<[4, 1]> : vector<2xindex>) : () -> !tosa.shape<2>
-   // CHECK-DAG: %[[VAR2:.*]] = tosa.tile %[[VAR0]], %[[VAR1]]
-   %0 = "stablehlo.iota"() {iota_dimension = 1 : i64} : () -> (tensor<4x8xi32>)
-   return %0 : tensor<4x8xi32>
-diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/ternary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/ternary.mlir
---- stablehlo/stablehlo/conversions/tosa/tests/ternary.mlir
-+++ stablehlo/stablehlo/conversions/tosa/tests/ternary.mlir
-@@ -2,7 +2,7 @@
- 
- // CHECK-LABEL: @concatenate
- func.func @concatenate(%arg0 : tensor<5x2xf32>, %arg1 : tensor<5x5xf32>, %arg2 : tensor<5x7xf32>) -> tensor<5x14xf32> {
--  // CHECK: tosa.concat %arg0, %arg1, %arg2 {axis = 1 : i32} : (tensor<5x2xf32>, tensor<5x5xf32>, tensor<5x7xf32>) -> tensor<5x14xf32>
-+  // CHECK: tosa.concat %arg0, %arg1, %arg2 axis(1) : (tensor<5x2xf32>, tensor<5x5xf32>, tensor<5x7xf32>) -> tensor<5x14xf32>
-   %0 = "stablehlo.concatenate"(%arg0, %arg1, %arg2) {dimension = 1 : i64} : (tensor<5x2xf32>, tensor<5x5xf32>, tensor<5x7xf32>) -> tensor<5x14xf32>
-   return %0 : tensor<5x14xf32>
- }
-diff --ruN a/stablehlo/stablehlo/conversions/tosa/tests/unary.mlir b/stablehlo/stablehlo/conversions/tosa/tests/unary.mlir
---- stablehlo/stablehlo/conversions/tosa/tests/unary.mlir
-+++ stablehlo/stablehlo/conversions/tosa/tests/unary.mlir
-@@ -30,7 +30,7 @@
- 
- // CHECK-LABEL: @exponential_minus_one
- func.func @exponential_minus_one(%arg : tensor<10xf32>) -> tensor<10xf32> {
--  // CHECK-DAG: %[[VAR0:.*]] = "tosa.const"() <{values = dense<1.000000e+00>
-+  // CHECK-DAG: %[[VAR0:.*]] = tosa.const values(dense<1.000000e+00>
-   // CHECK-DAG: %[[VAR1:.*]] = tosa.exp %arg0
-   // CHECK-DAG: %[[VAR2:.*]] = tosa.sub %[[VAR1]], %[[VAR0]]
-   %0 = "stablehlo.exponential_minus_one"(%arg) : (tensor<10xf32>) -> tensor<10xf32>
-@@ -46,7 +46,7 @@
- 
- // CHECK-LABEL: @is_finite
- func.func @is_finite(%arg : tensor<10xf32>) -> tensor<10xi1> {
--  // CHECK-DAG: %[[VAR0:.*]] = "tosa.const"() <{values = dense<0x7F800000>
-+  // CHECK-DAG: %[[VAR0:.*]] = tosa.const values(dense<0x7F800000>
-   // CHECK-DAG: %[[VAR1:.*]] = tosa.abs %arg0
-   // CHECK-DAG: %[[VAR2:.*]] = tosa.equal %[[VAR1]], %[[VAR0]]
-   // CHECK-DAG: %[[VAR3:.*]] = tosa.logical_not %[[VAR2]]
-@@ -63,7 +63,7 @@
- 
- // CHECK-LABEL: @log_plus_one
- func.func @log_plus_one(%arg : tensor<10xf16>) -> tensor<10xf16> {
--  // CHECK-DAG: %[[VAR0:.*]] = "tosa.const"() <{values = dense<1.000000e+00>
-+  // CHECK-DAG: %[[VAR0:.*]] = tosa.const values(dense<1.000000e+00>
-   // CHECK-DAG: %[[VAR1:.*]] = tosa.add %arg0, %[[VAR0]]
-   // CHECK-DAG: %[[VAR2:.*]] = tosa.log %[[VAR1]]
-   %0 = "stablehlo.log_plus_one"(%arg) : (tensor<10xf16>) -> tensor<10xf16>
-@@ -79,8 +79,8 @@
- 
- // CHECK-LABEL: @slice
- func.func @slice(%arg : tensor<4x3xf32>) -> tensor<2x2xf32> {
--  // CHECK-DAG: %[[SIZE:.*]] = tosa.const_shape {values = dense<[2, 1]> : tensor<2xindex>} : () -> !tosa.shape<2>
--  // CHECK-DAG: %[[START:.*]] = tosa.const_shape {values = dense<2> : tensor<2xindex>} : () -> !tosa.shape<2>
-+  // CHECK-DAG: %[[SIZE:.*]] = tosa.const_shape values(dense<[2, 1]> : tensor<2xindex>) : () -> !tosa.shape<2>
-+  // CHECK-DAG: %[[START:.*]] = tosa.const_shape values(dense<2> : tensor<2xindex>) : () -> !tosa.shape<2>
-   // CHECK: tosa.slice %arg0, %[[SIZE]], %[[START]]
-   %0 = "stablehlo.slice"(%arg) {
-     start_indices = array<i64: 2, 1>,
-@@ -123,15 +123,15 @@
- 
- // CHECK-LABEL: @transpose
- func.func @transpose(%arg0: tensor<1x2x3xf32>) -> tensor<3x2x1xf32> {
--  // CHECK: %[[VAR0:.*]] = tosa.transpose %arg0 {perms = array<i32: 2, 1, 0>}
-+  // CHECK: %[[VAR0:.*]] = tosa.transpose %arg0 perms([2, 1, 0])
-   %0 = "stablehlo.transpose"(%arg0) {permutation = array<i64: 2, 1, 0>} : (tensor<1x2x3xf32>) -> tensor<3x2x1xf32>
-   return %0 : tensor<3x2x1xf32>
- }
- 
- // CHECK-LABEL: @while
- func.func @while(%arg0: tensor<i32>) -> tensor<i32> {
--  // CHECK-DAG: %[[VAR0:.*]] = "tosa.const"() <{values = dense<3> : tensor<i32>}
--  // CHECK-DAG: %[[VAR1:.*]] = "tosa.const"() <{values = dense<1> : tensor<i32>}
-+  // CHECK-DAG: %[[VAR0:.*]] = tosa.const values(dense<3> : tensor<i32>)
-+  // CHECK-DAG: %[[VAR1:.*]] = tosa.const values(dense<1> : tensor<i32>)
-   // CHECK:     %[[VAR2:.*]] = tosa.while_loop (%[[ARG1:.+]] = %arg0) : (tensor<i32>) -> tensor<i32> {
-   // CHECK:       %[[VAR3:.*]] = tosa.equal %[[ARG1]], %[[VAR0]]
-   // CHECK:       tosa.yield %[[VAR3]]
-@@ -158,7 +158,7 @@
- 
- // CHECK-LABEL: @reshape
- func.func @reshape(%arg0 : tensor<2x3xf32>) -> tensor<6xf32> {
--  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape {values = dense<6> : tensor<1xindex>} : () -> !tosa.shape<1>
-+  // CHECK-DAG: %[[VAR0:.*]] = tosa.const_shape values(dense<6> : tensor<1xindex>) : () -> !tosa.shape<1>
-   // CHECK: tosa.reshape %arg0, %[[VAR0]]
-   %0 = "stablehlo.reshape"(%arg0) : (tensor<2x3xf32>) -> tensor<6xf32>
-   return %0 : tensor<6xf32>
-diff --ruN a/stablehlo/stablehlo/dialect/Base.cpp b/stablehlo/stablehlo/dialect/Base.cpp
---- stablehlo/stablehlo/dialect/Base.cpp
-+++ stablehlo/stablehlo/dialect/Base.cpp
-@@ -671,6 +671,7 @@
-   StringRef f32 = Float32Type::name;
-   StringRef f64 = Float64Type::name;
-   StringRef tf32 = FloatTF32Type::name;
-+  StringRef f8e4m3fn = Float8E4M3FNType::name;
-   std::map<std::tuple<StringRef, StringRef, StringRef, int64_t>,
-            KnownDotAlgorithm>
-       knownDotAlgorithms{
-@@ -679,6 +680,10 @@
-           {{bf16, bf16, bf16, 1}, KnownDotAlgorithm::BF16_BF16_BF16},
-           {{bf16, bf16, f32, 1}, KnownDotAlgorithm::BF16_BF16_F32},
-           {{bf16, bf16, f32, 3}, KnownDotAlgorithm::BF16_BF16_F32_X3},
-+          {{f8e4m3fn, f8e4m3fn, f32, 3},
-+           KnownDotAlgorithm::F8E4M3FN_F8E4M3FN_F32_X3},
-+          {{f8e4m3fn, f8e4m3fn, f32, 4},
-+           KnownDotAlgorithm::F8E4M3FN_F8E4M3FN_F32_X4},
-           {{bf16, bf16, f32, 6}, KnownDotAlgorithm::BF16_BF16_F32_X6},
-           {{bf16, bf16, f32, 9}, KnownDotAlgorithm::BF16_BF16_F32_X9},
-           {{tf32, tf32, f32, 1}, KnownDotAlgorithm::TF32_TF32_F32},
-diff --ruN a/stablehlo/stablehlo/dialect/Base.h b/stablehlo/stablehlo/dialect/Base.h
---- stablehlo/stablehlo/dialect/Base.h
-+++ stablehlo/stablehlo/dialect/Base.h
-@@ -264,6 +264,8 @@
-   F32_F32_F32 = 11,
-   F64_F64_F64 = 12,
-   BF16_BF16_F32_X9 = 13,
-+  F8E4M3FN_F8E4M3FN_F32_X3 = 14,
-+  F8E4M3FN_F8E4M3FN_F32_X4 = 15,
- };
- 
- FailureOr<KnownDotAlgorithm> getKnownDotAlgorithm(
-diff --ruN a/stablehlo/stablehlo/dialect/ChloBytecode.cpp b/stablehlo/stablehlo/dialect/ChloBytecode.cpp
---- stablehlo/stablehlo/dialect/ChloBytecode.cpp
-+++ stablehlo/stablehlo/dialect/ChloBytecode.cpp
-@@ -111,37 +111,37 @@
- /// This class implements the bytecode interface for the  CHLO dialect.
- class ChloBytecodeInterface : public BytecodeDialectInterface {
-  public:
--  ChloBytecodeInterface(Dialect *dialect) : BytecodeDialectInterface(dialect) {}
-+  ChloBytecodeInterface(Dialect* dialect) : BytecodeDialectInterface(dialect) {}
- 
-   //===--------------------------------------------------------------------===//
-   // Attributes
- 
-   // These methods are invoked by superclass when an attr from  CHLO dialect
-   // is encountered.
--  Attribute readAttribute(DialectBytecodeReader &reader) const override;
-+  Attribute readAttribute(DialectBytecodeReader& reader) const override;
-   LogicalResult writeAttribute(Attribute attr,
--                               DialectBytecodeWriter &writer) const override;
-+                               DialectBytecodeWriter& writer) const override;
- 
-   // TO ADD ATTRIBUTE: Include a read method for each attribute in CHLO
-   // Ex: SomeAttr readSomeAttr(DialectBytecodeReader &reader) const;
-   ComparisonDirectionAttr readComparisonDirectionAttr(
--      DialectBytecodeReader &reader) const;
-+      DialectBytecodeReader& reader) const;
-   ComparisonTypeAttr readComparisonTypeAttr(
--      DialectBytecodeReader &reader) const;
-+      DialectBytecodeReader& reader) const;
- 
-   // TO ADD ATTRIBUTE: Include a write method for each attribute in CHLO
-   // Ex: void write(SomeAttr attr, DialectBytecodeWriter &writer) const;
--  void write(ComparisonDirectionAttr attr, DialectBytecodeWriter &writer) const;
--  void write(ComparisonTypeAttr attr, DialectBytecodeWriter &writer) const;
-+  void write(ComparisonDirectionAttr attr, DialectBytecodeWriter& writer) const;
-+  void write(ComparisonTypeAttr attr, DialectBytecodeWriter& writer) const;
- 
-   //===--------------------------------------------------------------------===//
-   // Types
- 
-   // These methods are invoked by superclass when a type from  CHLO dialect
-   // is encountered.
--  Type readType(DialectBytecodeReader &reader) const override;
-+  Type readType(DialectBytecodeReader& reader) const override;
-   LogicalResult writeType(Type type,
--                          DialectBytecodeWriter &writer) const override;
-+                          DialectBytecodeWriter& writer) const override;
- 
-   // TO ADD TYPE: Include a read method for each type in CHLO
-   // Ex: SomeType readSomeType(DialectBytecodeReader &reader) const;
-@@ -158,7 +158,7 @@
- 
- // TO ADD ATTRIBUTE: Update the switch to include a branch for the attr.
- Attribute ChloBytecodeInterface::readAttribute(
--    DialectBytecodeReader &reader) const {
-+    DialectBytecodeReader& reader) const {
-   uint64_t code;
-   if (failed(reader.readVarInt(code))) return Attribute();
-   switch (code) {
-@@ -173,7 +173,7 @@
- }
- 
- ComparisonDirectionAttr ChloBytecodeInterface::readComparisonDirectionAttr(
--    DialectBytecodeReader &reader) const {
-+    DialectBytecodeReader& reader) const {
-   LOG_READ_CALL;
-   return hlo::bytecode::readEnumAttribute<ComparisonDirectionAttr>(
-       reader, getContext(),
-@@ -181,7 +181,7 @@
- }
- 
- ComparisonTypeAttr ChloBytecodeInterface::readComparisonTypeAttr(
--    DialectBytecodeReader &reader) const {
-+    DialectBytecodeReader& reader) const {
-   LOG_READ_CALL;
-   return hlo::bytecode::readEnumAttribute<ComparisonTypeAttr>(
-       reader, getContext(),
-@@ -195,7 +195,7 @@
- // If this method returns failure, the string serialization is used in the
- // bytecode.
- LogicalResult ChloBytecodeInterface::writeAttribute(
--    Attribute attr, DialectBytecodeWriter &writer) const {
-+    Attribute attr, DialectBytecodeWriter& writer) const {
-   return TypeSwitch<Attribute, LogicalResult>(attr)
-       .Case<ComparisonDirectionAttr, ComparisonTypeAttr>([&](auto attr) {
-         LOG_WRITE_CALL;
-@@ -209,13 +209,13 @@
- }
- 
- void ChloBytecodeInterface::write(ComparisonDirectionAttr attr,
--                                  DialectBytecodeWriter &writer) const {
-+                                  DialectBytecodeWriter& writer) const {
-   writer.writeVarInt(chlo_encoding::kComparisonDirectionAttr);
-   hlo::bytecode::writeEnumAttribute<ComparisonDirection>(attr, writer);
- }
- 
- void ChloBytecodeInterface::write(ComparisonTypeAttr attr,
--                                  DialectBytecodeWriter &writer) const {
-+                                  DialectBytecodeWriter& writer) const {
-   writer.writeVarInt(chlo_encoding::kComparisonTypeAttr);
-   hlo::bytecode::writeEnumAttribute<ComparisonType>(attr, writer);
- }
-@@ -224,7 +224,7 @@
- // Types: Reader
- 
- // TO ADD TYPE: Update the case selection to include the new type.
--Type ChloBytecodeInterface::readType(DialectBytecodeReader &reader) const {
-+Type ChloBytecodeInterface::readType(DialectBytecodeReader& reader) const {
-   uint64_t code;
-   if (failed(reader.readVarInt(code))) return Type();
- 
-@@ -240,7 +240,7 @@
- 
- // TO ADD TYPE: Update the case selection to include the new type.
- LogicalResult ChloBytecodeInterface::writeType(
--    Type type, DialectBytecodeWriter &writer) const {
-+    Type type, DialectBytecodeWriter& writer) const {
-   return TypeSwitch<Type, LogicalResult>(type).Default([&](Type) {
-     LOG_NOT_IMPLEMENTED;
-     return failure();
-@@ -249,7 +249,7 @@
- 
- }  // namespace
- 
--void addBytecodeInterface(ChloDialect *dialect) {
-+void addBytecodeInterface(ChloDialect* dialect) {
-   dialect->addInterfaces<ChloBytecodeInterface>();
- }
- }  // namespace chlo
-diff --ruN a/stablehlo/stablehlo/dialect/ChloBytecode.h b/stablehlo/stablehlo/dialect/ChloBytecode.h
---- stablehlo/stablehlo/dialect/ChloBytecode.h
-+++ stablehlo/stablehlo/dialect/ChloBytecode.h
-@@ -21,7 +21,7 @@
- 
- // Add the interface necessary for encoding and decoding the CHLO dialect
- // components in bytecode.
--void addBytecodeInterface(ChloDialect *dialect);
-+void addBytecodeInterface(ChloDialect* dialect);
- }  // namespace chlo
- }  // namespace mlir
- 
-diff --ruN a/stablehlo/stablehlo/dialect/ChloEnums.td b/stablehlo/stablehlo/dialect/ChloEnums.td
---- stablehlo/stablehlo/dialect/ChloEnums.td
-+++ stablehlo/stablehlo/dialect/ChloEnums.td
-@@ -46,7 +46,9 @@
-   let cppNamespace = "::mlir::chlo";
- }
- 
--def CHLO_ComparisonDirectionAttr : EnumAttr<CHLO_Dialect, CHLO_ComparisonDirection, "comparison_direction">;
-+def CHLO_ComparisonDirectionAttr : EnumAttr<CHLO_Dialect, CHLO_ComparisonDirection, "comparison_direction"> {
-+  let assemblyFormat = "$value";
-+}
- 
- def CHLO_DEFAULT_COMPARISON_TYPE : NativeCodeCall<"::mlir::chlo::ComparisonTypeAttr()">;
- def CHLO_COMPARISON_TYPE_NOTYPE : I32EnumAttrCase<"NOTYPE", 0>;
-@@ -54,6 +56,7 @@
- def CHLO_COMPARISON_TYPE_FLOAT_TOTAL_ORDER : I32EnumAttrCase<"TOTALORDER", 2>;
- def CHLO_COMPARISON_TYPE_SIGNED : I32EnumAttrCase<"SIGNED", 3>;
- def CHLO_COMPARISON_TYPE_UNSIGNED : I32EnumAttrCase<"UNSIGNED", 4>;
-+def CHLO_COMPARISON_TYPE_WEAK_ORDER : I32EnumAttrCase<"WEAKORDER", 5>;
- 
- def CHLO_ComparisonType : I32EnumAttr<"ComparisonType",
-     "Which comparison type to use.",
-@@ -62,13 +65,16 @@
-       CHLO_COMPARISON_TYPE_FLOAT,
-       CHLO_COMPARISON_TYPE_FLOAT_TOTAL_ORDER,
-       CHLO_COMPARISON_TYPE_SIGNED,
--      CHLO_COMPARISON_TYPE_UNSIGNED
-+      CHLO_COMPARISON_TYPE_UNSIGNED,
-+      CHLO_COMPARISON_TYPE_WEAK_ORDER
-     ]> {
-   let genSpecializedAttr = 0;
-   let cppNamespace = "::mlir::chlo";
- }
- 
--def CHLO_ComparisonTypeAttr : EnumAttr<CHLO_Dialect, CHLO_ComparisonType, "comparison_type">;
-+def CHLO_ComparisonTypeAttr : EnumAttr<CHLO_Dialect, CHLO_ComparisonType, "comparison_type"> {
-+  let assemblyFormat = "$value";
-+}
- 
- //===----------------------------------------------------------------------===//
- // Ragged dot op definitions.
-@@ -90,7 +96,9 @@
-   let cppNamespace = "::mlir::chlo";
- }
- 
--def CHLO_PrecisionAttr : EnumAttr<CHLO_Dialect, CHLO_Precision, "precision">;
-+def CHLO_PrecisionAttr : EnumAttr<CHLO_Dialect, CHLO_Precision, "precision"> {
-+  let assemblyFormat = "$value";
-+}
- 
- def CHLO_PrecisionConfigAttr:
-   TypedArrayAttrBase<CHLO_PrecisionAttr, "Precision Config attribute">;
-diff --ruN a/stablehlo/stablehlo/dialect/ChloOps.cpp b/stablehlo/stablehlo/dialect/ChloOps.cpp
---- stablehlo/stablehlo/dialect/ChloOps.cpp
-+++ stablehlo/stablehlo/dialect/ChloOps.cpp
-@@ -827,6 +827,10 @@
-   if (adaptor.getScanDimSize().has_value()) {
-     dimSize = adaptor.getScanDimSize().value();
-   }
-+  int64_t dimBound = ShapedType::kDynamic;
-+  // Prototype encoding attribute used to construct dialect type extensions
-+  // if an output does not already have an encoding.
-+  Attribute prototypeEncoding;
-   for (auto [i, type] : llvm::enumerate(adaptor.getInputs().getTypes())) {
-     auto inputType = dyn_cast<RankedTensorType>(type);
-     if (!inputType) {
-@@ -839,8 +843,27 @@
-     if (dimSize == ShapedType::kDynamic) {
-       dimSize = inputType.getDimSize(dim);
-     }
-+    if (inputType.getEncoding()) {
-+      prototypeEncoding = inputType.getEncoding();
-+    }
-+    ArrayRef<int64_t> bounds = hlo::encodingToBounds(inputType.getEncoding());
-+    if (!bounds.empty() && dim < static_cast<int64_t>(bounds.size())) {
-+      int64_t currentBound = bounds[dim];
-+      if (currentBound != ShapedType::kDynamic) {
-+        if (dimBound == ShapedType::kDynamic) {
-+          dimBound = currentBound;
-+        } else {
-+          dimBound = std::min(dimBound, currentBound);
-+        }
-+      }
-+    }
-   }
- 
-+  // If the scan dimension size is dynamic, propagate any bound inferred from
-+  // the inputs; otherwise static dimensions do not need a bound.
-+  const int64_t insertBound =
-+      (dimSize == ShapedType::kDynamic) ? dimBound : ShapedType::kDynamic;
-+
-   for (auto [i, type] : llvm::enumerate(terminator->getOperands().getTypes())) {
-     auto resultType = dyn_cast<RankedTensorType>(type);
-     if (!resultType) {
-@@ -849,10 +872,34 @@
-     }
-     SmallVector<int64_t> shape(resultType.getShape().begin(),
-                                resultType.getShape().end());
-+    Attribute encoding = resultType.getEncoding();
-+    // The first `numOutputs` terminator operands correspond to scan outputs,
-+    // which have the scan dimension re-inserted. The remaining operands are
-+    // carries whose shapes and encodings are preserved as-is.
-     if (i < numOutputs) {
-       shape.insert(std::next(shape.begin(), dim), dimSize);
-+      SmallVector<int64_t> resultBounds =
-+          llvm::to_vector(hlo::encodingToBounds(encoding));
-+      // Re-insert bounds along the scan dimension if the result has bounds or
-+      // if the scan dimension has a bound.
-+      if (!resultBounds.empty() || insertBound != ShapedType::kDynamic) {
-+        // If the body result did not have bounds, initialize all dimensions to
-+        // dynamic and borrow the prototype encoding to create dialect type
-+        // extensions.
-+        if (resultBounds.empty()) {
-+          resultBounds.assign(resultType.getRank(), ShapedType::kDynamic);
-+          if (!encoding) {
-+            encoding = prototypeEncoding;
-+          }
-+        }
-+        resultBounds.insert(std::next(resultBounds.begin(), dim), insertBound);
-+        if (encoding) {
-+          encoding = hlo::boundsToEncoding(encoding, resultBounds);
-+        }
-+      }
-     }
--    inferredReturnShapes.emplace_back(shape, resultType.getElementType());
-+    inferredReturnShapes.emplace_back(shape, resultType.getElementType(),
-+                                      encoding);
-   }
- 
-   return success();
-@@ -902,7 +949,17 @@
-     if (i < getInputs().size()) {
-       auto argShape = llvm::to_vector(argType.getShape());
-       argShape.erase(std::next(argShape.begin(), dim));
--      argType = argType.clone(argShape);
-+      Attribute encoding = argType.getEncoding();
-+      SmallVector<int64_t> argBounds =
-+          llvm::to_vector(hlo::encodingToBounds(encoding));
-+      if (!argBounds.empty()) {
-+        argBounds.erase(std::next(argBounds.begin(), dim));
-+        encoding = argBounds.empty()
-+                       ? Attribute()
-+                       : hlo::boundsToEncoding(encoding, argBounds);
-+      }
-+      argType =
-+          RankedTensorType::get(argShape, argType.getElementType(), encoding);
-     }
-     if (!hlo::isCompatibleForHloTypeInference(
-             argType, bodyBlock.getArgument(i).getType())) {
-diff --ruN a/stablehlo/stablehlo/dialect/ChloOps.td b/stablehlo/stablehlo/dialect/ChloOps.td
---- stablehlo/stablehlo/dialect/ChloOps.td
-+++ stablehlo/stablehlo/dialect/ChloOps.td
-@@ -54,6 +54,7 @@
-     and provide conversion patterns to fully materialize into lower level
-     dialects.
-   }];
-+  let useStrictPropertiesInAssemblyFormat = 0;
- }
- 
- class CHLO_Op<string mnemonic, list<Trait> traits> :
-diff --ruN a/stablehlo/stablehlo/dialect/Register.cpp b/stablehlo/stablehlo/dialect/Register.cpp
---- stablehlo/stablehlo/dialect/Register.cpp
-+++ stablehlo/stablehlo/dialect/Register.cpp
-@@ -27,7 +27,7 @@
- namespace mlir {
- namespace stablehlo {
- 
--void registerAllDialects(mlir::DialectRegistry &registry) {
-+void registerAllDialects(mlir::DialectRegistry& registry) {
-   // clang-format off
-   registry.insert<mlir::func::FuncDialect,
-                   mlir::quant::QuantDialect,
-diff --ruN a/stablehlo/stablehlo/dialect/Register.h b/stablehlo/stablehlo/dialect/Register.h
---- stablehlo/stablehlo/dialect/Register.h
-+++ stablehlo/stablehlo/dialect/Register.h
-@@ -23,7 +23,7 @@
- namespace stablehlo {
- 
- // Add chlo and stablehlo dialects to the provided registry.
--void registerAllDialects(DialectRegistry &registry);
-+void registerAllDialects(DialectRegistry& registry);
- 
- }  // namespace stablehlo
- }  // namespace mlir
-diff --ruN a/stablehlo/stablehlo/dialect/StablehloBytecode.h b/stablehlo/stablehlo/dialect/StablehloBytecode.h
---- stablehlo/stablehlo/dialect/StablehloBytecode.h
-+++ stablehlo/stablehlo/dialect/StablehloBytecode.h
-@@ -21,7 +21,7 @@
- 
- // Add the interface necessary for encoding and decoding StableHLO dialect
- // components in bytecode.
--void addBytecodeInterface(StablehloDialect *dialect);
-+void addBytecodeInterface(StablehloDialect* dialect);
- }  // namespace stablehlo
- }  // namespace mlir
- 
-diff --ruN a/stablehlo/stablehlo/dialect/StablehloEnums.td b/stablehlo/stablehlo/dialect/StablehloEnums.td
---- stablehlo/stablehlo/dialect/StablehloEnums.td
-+++ stablehlo/stablehlo/dialect/StablehloEnums.td
-@@ -40,7 +40,9 @@
-   let cppNamespace = "::mlir::stablehlo";
- }
- 
--def StableHLO_PrecisionAttr : EnumAttr<StableHLO_Dialect, StableHLO_Precision, "precision">;
-+def StableHLO_PrecisionAttr : EnumAttr<StableHLO_Dialect, StableHLO_Precision, "precision"> {
-+  let assemblyFormat = "$value";
-+}
- 
- // TODO(b/129153247) See if it's possible to also validate the size.
- def StableHLO_PrecisionConfigAttr:
-@@ -91,7 +93,9 @@
-   let cppNamespace = "::mlir::stablehlo";
- }
- 
--def StableHLO_FftTypeAttr : EnumAttr<StableHLO_Dialect, StableHLO_FftType, "fft_type">;
-+def StableHLO_FftTypeAttr : EnumAttr<StableHLO_Dialect, StableHLO_FftType, "fft_type"> {
-+  let assemblyFormat = "$value";
-+}
- 
- //===----------------------------------------------------------------------===//
- // Custom call enum definitions.
-@@ -146,7 +150,9 @@
-   let cppNamespace = "::mlir::stablehlo";
- }
- 
--def StableHLO_ComparisonDirectionAttr : EnumAttr<StableHLO_Dialect, StableHLO_ComparisonDirection, "comparison_direction">;
-+def StableHLO_ComparisonDirectionAttr : EnumAttr<StableHLO_Dialect, StableHLO_ComparisonDirection, "comparison_direction"> {
-+  let assemblyFormat = "$value";
-+}
- 
- def STABLEHLO_DEFAULT_COMPARISON_TYPE : NativeCodeCall<"::mlir::stablehlo::ComparisonTypeAttr()">;
- def STABLEHLO_COMPARISON_TYPE_NOTYPE : I32EnumAttrCase<"NOTYPE", 0>;
-@@ -154,6 +160,7 @@
- def STABLEHLO_COMPARISON_TYPE_FLOAT_TOTAL_ORDER : I32EnumAttrCase<"TOTALORDER", 2>;
- def STABLEHLO_COMPARISON_TYPE_SIGNED : I32EnumAttrCase<"SIGNED", 3>;
- def STABLEHLO_COMPARISON_TYPE_UNSIGNED : I32EnumAttrCase<"UNSIGNED", 4>;
-+def STABLEHLO_COMPARISON_TYPE_WEAK_ORDER : I32EnumAttrCase<"WEAKORDER", 5>;
- 
- def StableHLO_ComparisonType : I32EnumAttr<"ComparisonType",
-     "Which comparison type to use.",
-@@ -162,13 +169,16 @@
-       STABLEHLO_COMPARISON_TYPE_FLOAT,
-       STABLEHLO_COMPARISON_TYPE_FLOAT_TOTAL_ORDER,
-       STABLEHLO_COMPARISON_TYPE_SIGNED,
--      STABLEHLO_COMPARISON_TYPE_UNSIGNED
-+      STABLEHLO_COMPARISON_TYPE_UNSIGNED,
-+      STABLEHLO_COMPARISON_TYPE_WEAK_ORDER
-     ]> {
-   let genSpecializedAttr = 0;
-   let cppNamespace = "::mlir::stablehlo";
- }
- 
--def StableHLO_ComparisonTypeAttr : EnumAttr<StableHLO_Dialect, StableHLO_ComparisonType, "comparison_type">;
-+def StableHLO_ComparisonTypeAttr : EnumAttr<StableHLO_Dialect, StableHLO_ComparisonType, "comparison_type"> {
-+  let assemblyFormat = "$value";
-+}
- 
- // These mirror the XLA Transpose enum in Triangular Solve options.
- def STABLEHLO_TRANSPOSE_INVALID : I32EnumAttrCase<"TRANSPOSE_INVALID", 0>;
-@@ -188,7 +198,9 @@
-   let cppNamespace = "::mlir::stablehlo";
- }
- 
--def StableHLO_TransposeAttr : EnumAttr<StableHLO_Dialect, StableHLO_Transpose, "transpose">;
-+def StableHLO_TransposeAttr : EnumAttr<StableHLO_Dialect, StableHLO_Transpose, "transpose"> {
-+  let assemblyFormat = "$value";
-+}
- 
- def STABLEHLO_RNG_DISTRIBUTION_UNIFORM : I32EnumAttrCase<"UNIFORM", 1>;
- def STABLEHLO_RNG_DISTRIBUTION_NORMAL : I32EnumAttrCase<"NORMAL", 2>;
-@@ -203,7 +215,9 @@
-   let cppNamespace = "::mlir::stablehlo";
- }
- 
--def StableHLO_RngDistributionAttr : EnumAttr<StableHLO_Dialect, StableHLO_RngDistribution, "rng_distribution">;
-+def StableHLO_RngDistributionAttr : EnumAttr<StableHLO_Dialect, StableHLO_RngDistribution, "rng_distribution"> {
-+  let assemblyFormat = "$value";
-+}
- 
- def STABLEHLO_RNG_ALGORITHM_DEFAULT : I32EnumAttrCase<"DEFAULT", 0>;
- def STABLEHLO_RNG_ALGORITHM_THREE_FRY : I32EnumAttrCase<"THREE_FRY", 1>;
-@@ -220,6 +234,8 @@
-   let cppNamespace = "::mlir::stablehlo";
- }
- 
--def StableHLO_RngAlgorithmAttr : EnumAttr<StableHLO_Dialect, StableHLO_RngAlgorithm, "rng_algorithm">;
-+def StableHLO_RngAlgorithmAttr : EnumAttr<StableHLO_Dialect, StableHLO_RngAlgorithm, "rng_algorithm"> {
-+  let assemblyFormat = "$value";
-+}
- 
- #endif // STABLEHLO_DIALECT_STABLEHLO_ENUMS
-diff --ruN a/stablehlo/stablehlo/dialect/StablehloOps.cpp b/stablehlo/stablehlo/dialect/StablehloOps.cpp
---- stablehlo/stablehlo/dialect/StablehloOps.cpp
-+++ stablehlo/stablehlo/dialect/StablehloOps.cpp
-@@ -249,9 +249,11 @@
- INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(DivOp)
- INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(ExpOp)
- INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(Expm1Op)
-+INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(Exp2Op)
- INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(FloorOp)
- INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(LogOp)
- INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(Log1pOp)
-+INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(Log2Op)
- INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(LogisticOp)
- INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(MaxOp)
- INFER_RETURN_TYPE_COMPONENTS_FROM_OPERANDS(MinOp)
-@@ -1129,6 +1131,19 @@
-   return success();
- }
- 
-+// ===---------------------------------------------------------------------===//
-+// Exp2Op
-+//===----------------------------------------------------------------------===//
-+
-+LogicalResult Exp2Op::verify() {
-+  if (auto attr = getResultAccuracyAttr()) {
-+    return ResultAccuracyAttr::verify([&] { return emitError(); },
-+                                      attr.getAtol(), attr.getRtol(),
-+                                      attr.getUlps(), attr.getMode());
-+  }
-+  return success();
-+}
-+
- // ===---------------------------------------------------------------------===//
- // LogOp
- //===----------------------------------------------------------------------===//
-@@ -1155,6 +1170,19 @@
-   return success();
- }
- 
-+// ===---------------------------------------------------------------------===//
-+// Log2Op
-+//===----------------------------------------------------------------------===//
-+
-+LogicalResult Log2Op::verify() {
-+  if (auto attr = getResultAccuracyAttr()) {
-+    return ResultAccuracyAttr::verify([&] { return emitError(); },
-+                                      attr.getAtol(), attr.getRtol(),
-+                                      attr.getUlps(), attr.getMode());
-+  }
-+  return success();
-+}
-+
- // ===---------------------------------------------------------------------===//
- // LogisticOp
- //===----------------------------------------------------------------------===//
-@@ -3283,7 +3311,10 @@
-     RegionRange regions,
-     SmallVectorImpl<ShapedTypeComponents>& inferredReturnShapes) {
-   CompareOp::Adaptor adaptor(operands, attributes, properties, regions);
--  return hlo::inferCompareOp(context, location, adaptor.getLhs(),
-+  std::optional<StringRef> compareType;
-+  if (auto attr = adaptor.getCompareType())
-+    compareType = stringifyComparisonType(*attr);
-+  return hlo::inferCompareOp(context, location, adaptor.getLhs(), compareType,
-                              inferredReturnShapes);
- }
- 
-@@ -4332,6 +4363,23 @@
-   ReturnOp::create(*builder, loc, compare);
- }
- 
-+// For floating-point types, returns a predicate that selects lhs when lhs is
-+// NaN and either rhs is not NaN or lhs_index < rhs_index (lt_index_pred), so
-+// that index selection matches MaxOp/MinOp NaN propagation and tie-breaking.
-+static Value buildNaNLhsSelectionCondition(OpBuilder& builder, Location loc,
-+                                           Value lhs_value, Value rhs_value,
-+                                           Value lt_index_pred) {
-+  auto lhs_is_nan = CompareOp::create(builder, loc, lhs_value, lhs_value,
-+                                      ComparisonDirection::NE)
-+                        .getResult();
-+  auto rhs_not_nan = CompareOp::create(builder, loc, rhs_value, rhs_value,
-+                                       ComparisonDirection::EQ)
-+                         .getResult();
-+  auto nan_lhs_win =
-+      OrOp::create(builder, loc, rhs_not_nan, lt_index_pred).getResult();
-+  return AndOp::create(builder, loc, lhs_is_nan, nan_lhs_win).getResult();
-+}
-+
- void buildMaxAndArgmaxBody(Type elementType, Type indices_type, Region& body,
-                            OpBuilder& builder) {
-   OpBuilder::InsertionGuard guard(builder);
-@@ -4366,15 +4414,18 @@
-   // Final lhs Selection Condition: (gt_pred) OR (tie_breaker_condition)
-   auto final_lhs_condition =
-       OrOp::create(builder, loc, gt_pred, tie_breaker_condition).getResult();
-+  if (isa<FloatType>(elementType)) {
-+    auto nan_lhs_condition = buildNaNLhsSelectionCondition(
-+        builder, loc, lhs_value, rhs_value, lt_index_pred);
-+    final_lhs_condition =
-+        OrOp::create(builder, loc, final_lhs_condition, nan_lhs_condition)
-+            .getResult();
-+  }
- 
--  // Select Final Results:
--  // if final_lhs_condition:
--  //     return (lhs_value, lhs_index)
--  // else:
--  //     return (rhs_value, rhs_index)
-+  // Use MaxOp for the value so that NaNs propagate properly and unused-index
-+  // reductions simplify to kMaximum.
-   auto selected_value =
--      SelectOp::create(builder, loc, final_lhs_condition, lhs_value, rhs_value)
--          .getResult();
-+      MaxOp::create(builder, loc, lhs_value, rhs_value).getResult();
-   auto selected_index =
-       SelectOp::create(builder, loc, final_lhs_condition, lhs_index, rhs_index)
-           .getResult();
-diff --ruN a/stablehlo/stablehlo/dialect/StablehloOps.td b/stablehlo/stablehlo/dialect/StablehloOps.td
---- stablehlo/stablehlo/dialect/StablehloOps.td
-+++ stablehlo/stablehlo/dialect/StablehloOps.td
-@@ -38,6 +38,7 @@
- 
-   let useDefaultAttributePrinterParser = 0;
-   let useDefaultTypePrinterParser = 0;
-+  let useStrictPropertiesInAssemblyFormat = 0;
- }
- 
- class StableHLO_Op<string mnemonic, list<Trait> traits = []> :
-@@ -393,6 +394,33 @@
-   }];
- }
- 
-+def StableHLO_Exp2Op: StableHLO_UnaryElementwiseOp<"exp2",
-+    [HLO_CompatibleOperandsAndResultType],
-+    HLO_FpComplexOrQuantizedIntTensor> {
-+  let summary = "Exp2 operation";
-+  let description = [{
-+    Performs element-wise base-2 exponential operation on `operand` tensor and
-+    produces a `result` tensor.
-+
-+    See:
-+    https://github.com/openxla/stablehlo/blob/main/docs/spec.md#exp2
-+
-+    Example:
-+    ```mlir
-+    %result = stablehlo.exp2 %operand : tensor<2x2xf64>
-+    ```
-+  }];
-+  let arguments = (ins HLO_FpComplexOrQuantizedIntTensor:$operand,
-+                       DefaultValuedOptionalAttr<StableHLO_ResultAccuracyAttr,
-+                       "::mlir::stablehlo::ResultAccuracyMode::DEFAULT">:$result_accuracy);
-+  let results = (outs HLO_FpComplexOrQuantizedIntTensor:$result);
-+  let hasVerifier = 1;
-+
-+  let assemblyFormat = [{
-+    $operand attr-dict `:` custom<SameOperandsAndResultType>(type($operand), type($result))
-+  }];
-+}
-+
- def StableHLO_FloorOp: StableHLO_UnaryElementwiseOp<"floor",
-     [HLO_CompatibleOperandsAndResultType], HLO_FpOrQuantizedIntTensor> {
-   let summary = "Floor operation";
-@@ -506,6 +534,33 @@
-   }];
- }
- 
-+def StableHLO_Log2Op: StableHLO_UnaryElementwiseOp<"log2",
-+    [HLO_CompatibleOperandsAndResultType],
-+    HLO_FpComplexOrQuantizedIntTensor> {
-+  let summary = "Log2 operation";
-+  let description = [{
-+    Performs element-wise base-2 logarithm operation on `operand` tensor and
-+    produces a `result` tensor.
-+
-+    See:
-+    https://github.com/openxla/stablehlo/blob/main/docs/spec.md#log2
-+
-+    Example:
-+    ```mlir
-+    %result = stablehlo.log2 %operand : tensor<2x2xf64>
-+    ```
-+  }];
-+  let arguments = (ins HLO_FpComplexOrQuantizedIntTensor:$operand,
-+                       DefaultValuedOptionalAttr<StableHLO_ResultAccuracyAttr,
-+                       "::mlir::stablehlo::ResultAccuracyMode::DEFAULT">:$result_accuracy);
-+  let results = (outs HLO_FpComplexOrQuantizedIntTensor:$result);
-+  let hasVerifier = 1;
-+
-+  let assemblyFormat = [{
-+    $operand attr-dict `:` custom<SameOperandsAndResultType>(type($operand), type($result))
-+  }];
-+}
-+
- def StableHLO_LogisticOp: StableHLO_UnaryElementwiseOp<"logistic",
-     [HLO_CompatibleOperandsAndResultType /*logistic_c1*/],
-     HLO_FpComplexOrQuantizedIntTensor /*logistic_i1*/> { /*logistic_c1*/
-@@ -2525,8 +2580,9 @@
-       "::mlir::TypeRange":$resultTypes, "::mlir::ValueRange":$inputs,
-       "::mlir::ArrayRef<::mlir::NamedAttribute>":$attributes
-     ), [{
--      build($_builder, $_state, resultTypes, inputs, attributes,
--            /*regionsCount=*/0);
-+      $_state.addOperands(inputs);
-+      buildPropertiesAndDiscardableAttributes($_state, attributes);
-+      $_state.addTypes(resultTypes);
-     }]>,
-     OpBuilder<(ins
-       "::mlir::TypeRange":$resultTypes, "::mlir::ValueRange":$inputs,
-diff --ruN a/stablehlo/stablehlo/dialect/TypeInference.cpp b/stablehlo/stablehlo/dialect/TypeInference.cpp
---- stablehlo/stablehlo/dialect/TypeInference.cpp
-+++ stablehlo/stablehlo/dialect/TypeInference.cpp
-@@ -2032,11 +2032,22 @@
- }
- 
- LogicalResult inferCompareOp(
--    MLIRContext* context, std::optional<Location>, Value lhs,
-+    MLIRContext* context, std::optional<Location> location, Value lhs,
-+    std::optional<StringRef> compareType,
-     SmallVectorImpl<ShapedTypeComponents>& inferredReturnShapes) {
-+  auto argTy = cast<RankedTensorType>(lhs.getType());
-+  // compare_c3
-+  if (compareType && *compareType == "WEAKORDER" &&
-+      !isa<FloatType>(argTy.getElementType())) {
-+    return emitOptionalError(
-+        location,
-+        "WEAKORDER comparison type is only supported for floating-point "
-+        "element types, got ",
-+        argTy.getElementType());
-+  }
-+
-   // compare_c1
-   ShapedTypeComponents& components = inferredReturnShapes.emplace_back();
--  auto argTy = cast<RankedTensorType>(lhs.getType());
-   auto resElementTy = IntegerType::get(context, /*width=*/1);
-   // compare_c2
-   components =
-diff --ruN a/stablehlo/stablehlo/dialect/TypeInference.h b/stablehlo/stablehlo/dialect/TypeInference.h
---- stablehlo/stablehlo/dialect/TypeInference.h
-+++ stablehlo/stablehlo/dialect/TypeInference.h
-@@ -187,7 +187,8 @@
-     SmallVectorImpl<Type>& inferredReturnTypes);
- 
- LogicalResult inferCompareOp(
--    MLIRContext* context, std::optional<Location>, Value lhs,
-+    MLIRContext* context, std::optional<Location> location, Value lhs,
-+    std::optional<StringRef> compareType,
-     SmallVectorImpl<ShapedTypeComponents>& inferredReturnShapes);
- 
- LogicalResult inferComplexOp(std::optional<Location> location, Value lhs,
-diff --ruN a/stablehlo/stablehlo/dialect/Version.h b/stablehlo/stablehlo/dialect/Version.h
---- stablehlo/stablehlo/dialect/Version.h
-+++ stablehlo/stablehlo/dialect/Version.h
-@@ -38,7 +38,7 @@
-   static FailureOr<Version> fromString(llvm::StringRef versionRef);
- 
-   /// Return a Version representing the current VHLO dialect version.
--  static Version getCurrentVersion() { return Version(1, 20, 1); }
-+  static Version getCurrentVersion() { return Version(1, 23, 0); }
- 
-   /// Return a Version representing the minimum supported VHLO dialect version.
-   static Version getMinimumVersion() { return Version(0, 9, 0); }
-diff --ruN a/stablehlo/stablehlo/dialect/VhloBytecode.h b/stablehlo/stablehlo/dialect/VhloBytecode.h
---- stablehlo/stablehlo/dialect/VhloBytecode.h
-+++ stablehlo/stablehlo/dialect/VhloBytecode.h
-@@ -21,7 +21,7 @@
- 
- // Add the interface necessary for encoding and decoding VHLO dialect
- // components in bytecode.
--void addBytecodeInterface(VhloDialect *dialect);
-+void addBytecodeInterface(VhloDialect* dialect);
- }  // namespace vhlo
- }  // namespace mlir
- 
-diff --ruN a/stablehlo/stablehlo/dialect/VhloDialect.td b/stablehlo/stablehlo/dialect/VhloDialect.td
---- stablehlo/stablehlo/dialect/VhloDialect.td
-+++ stablehlo/stablehlo/dialect/VhloDialect.td
-@@ -59,10 +59,14 @@
-       1.18.0: Add `result_tilings` attribute to `custom_call` op.
-       1.19.0: Add CollectiveReduceOp.
-       1.20.0: Add has_dynamic_root to collective_broadcast op. Allow mixed fp8 operands in `convolution` and `dynamic_conv` ops.
-+      1.21.0: Add F8E4M3FN_F8E4M3FN_F32_X3 and F8E4M3FN_F8E4M3FN_F32_X4 dot algorithms.
-+      1.22.0: Add WEAKORDER comparison type to compare op.
-+      1.23.0: Add `exp2` and `log2` ops.
-   }];
- 
-   let useDefaultAttributePrinterParser = 0;
-   let useDefaultTypePrinterParser = 0;
-+  let useStrictPropertiesInAssemblyFormat = 0;
- }
- 
- #endif  // STABLEHLO_DIALECT_VHLO_DIALECT
-diff --ruN a/stablehlo/stablehlo/dialect/VhloEnums.td b/stablehlo/stablehlo/dialect/VhloEnums.td
---- stablehlo/stablehlo/dialect/VhloEnums.td
-+++ stablehlo/stablehlo/dialect/VhloEnums.td
-@@ -41,11 +41,16 @@
-   let cppNamespace = "::mlir::vhlo";
- }
- 
--class VHLO_EnumAttr<EnumAttrInfo enumInfo, string name, string minVersion, string maxVersion>
-+class VHLO_EnumAttr<EnumAttrInfo enumInfo, string name, string minVersion, string maxVersion, code customMinVersion = "">
-   : EnumAttr<VHLO_Dialect, enumInfo, name, [VHLO_VersionedAttrInterface]> {
-+  let assemblyFormat = "$value";
-   let extraClassDeclaration = [{
-     mlir::vhlo::Version getMinVersion() {
--      return mlir::vhlo::Version(}] # !subst(".", ", ", minVersion) # [{);
-+      }] # !if(
-+        !empty(customMinVersion),
-+        [{ return mlir::vhlo::Version(}] # !subst(".", ", ", minVersion) # [{); }],
-+        customMinVersion
-+      ) # [{
-     }
-     mlir::vhlo::Version getMaxVersion() {
-       }] # !if(
-@@ -90,17 +95,31 @@
- def VHLO_COMPARISON_TYPE_V1_FLOAT_TOTAL_ORDER : I32EnumAttrCase<"TOTALORDER", 2>;
- def VHLO_COMPARISON_TYPE_V1_SIGNED : I32EnumAttrCase<"SIGNED", 3>;
- def VHLO_COMPARISON_TYPE_V1_UNSIGNED : I32EnumAttrCase<"UNSIGNED", 4>;
-+def VHLO_COMPARISON_TYPE_V1_WEAK_ORDER : I32EnumAttrCase<"WEAKORDER", 5>;
- 
- def VHLO_ComparisonTypeV1 : VHLO_I32EnumAttr<"ComparisonTypeV1", [
-   VHLO_COMPARISON_TYPE_V1_NOTYPE,
-   VHLO_COMPARISON_TYPE_V1_FLOAT,
-   VHLO_COMPARISON_TYPE_V1_FLOAT_TOTAL_ORDER,
-   VHLO_COMPARISON_TYPE_V1_SIGNED,
--  VHLO_COMPARISON_TYPE_V1_UNSIGNED
-+  VHLO_COMPARISON_TYPE_V1_UNSIGNED,
-+  VHLO_COMPARISON_TYPE_V1_WEAK_ORDER
- ]> {}
- 
- def VHLO_ComparisonTypeAttrV1
--  : VHLO_EnumAttr<VHLO_ComparisonTypeV1, "comparison_type_v1", "0.9.0", "current">;
-+  : VHLO_EnumAttr<VHLO_ComparisonTypeV1, "comparison_type_v1", "0.9.0", "current", [{
-+      switch (getValue()) {
-+        case mlir::vhlo::ComparisonTypeV1::NOTYPE:
-+        case mlir::vhlo::ComparisonTypeV1::FLOAT:
-+        case mlir::vhlo::ComparisonTypeV1::TOTALORDER:
-+        case mlir::vhlo::ComparisonTypeV1::SIGNED:
-+        case mlir::vhlo::ComparisonTypeV1::UNSIGNED:
-+          return mlir::vhlo::Version(0, 9, 0);
-+        case mlir::vhlo::ComparisonTypeV1::WEAKORDER:
-+          return mlir::vhlo::Version(1, 22, 0);
-+      }
-+      llvm_unreachable("Unknown ComparisonTypeV1");
-+    }]>;
- 
- //===----------------------------------------------------------------------===//
- // CustomCallApiVersion
-diff --ruN a/stablehlo/stablehlo/dialect/VhloOps.cpp b/stablehlo/stablehlo/dialect/VhloOps.cpp
---- stablehlo/stablehlo/dialect/VhloOps.cpp
-+++ stablehlo/stablehlo/dialect/VhloOps.cpp
-@@ -382,6 +382,20 @@
-   return success();
- }
- 
-+LogicalResult verifyConstraint_1_21_0(mlir::Operation* op,
-+                                      Version targetVersion) {
-+  auto dotGeneralOp = cast<DotGeneralOpV2>(op);
-+  if (targetVersion < Version(1, 21, 0)) {
-+    auto lhsType = dyn_cast<TypeV1Attr>(dotGeneralOp.getLhsPrecisionType());
-+    auto rhsType = dyn_cast<TypeV1Attr>(dotGeneralOp.getRhsPrecisionType());
-+    if ((lhsType && isa<vhlo::FloatF8E4M3FNV1Type>(lhsType.getValue())) ||
-+        (rhsType && isa<vhlo::FloatF8E4M3FNV1Type>(rhsType.getValue()))) {
-+      return failure();
-+    }
-+  }
-+  return success();
-+}
-+
- }  // namespace
- 
- LogicalResult AllReduceOpV1::validateConstraint(mlir::Operation* op,
-@@ -394,6 +408,11 @@
-   return verifyConstraint_1_20_0(op, targetVersion);
- }
- 
-+LogicalResult DotGeneralOpV2::validateConstraint(mlir::Operation* op,
-+                                                 Version targetVersion) {
-+  return verifyConstraint_1_21_0(op, targetVersion);
-+}
-+
- LogicalResult DynamicConvOpV2::validateConstraint(mlir::Operation* op,
-                                                   Version targetVersion) {
-   return verifyConstraint_1_20_0(op, targetVersion);
-diff --ruN a/stablehlo/stablehlo/dialect/VhloOps.h b/stablehlo/stablehlo/dialect/VhloOps.h
---- stablehlo/stablehlo/dialect/VhloOps.h
-+++ stablehlo/stablehlo/dialect/VhloOps.h
-@@ -31,20 +31,20 @@
- 
- class VhloDialect : public Dialect {
-  public:
--  explicit VhloDialect(MLIRContext *context);
-+  explicit VhloDialect(MLIRContext* context);
-   static StringRef getDialectNamespace() { return "vhlo"; }
- 
-   // Parses a type registered to this dialect.
--  Type parseType(DialectAsmParser &parser) const override;
-+  Type parseType(DialectAsmParser& parser) const override;
- 
-   // Prints a type registered to this dialect.
--  void printType(Type type, DialectAsmPrinter &os) const override;
-+  void printType(Type type, DialectAsmPrinter& os) const override;
- 
-   // Parses an attribute registered to this dialect.
--  Attribute parseAttribute(DialectAsmParser &parser, Type type) const override;
-+  Attribute parseAttribute(DialectAsmParser& parser, Type type) const override;
- 
-   // Prints an attribute registered to this dialect.
--  void printAttribute(Attribute attr, DialectAsmPrinter &os) const override;
-+  void printAttribute(Attribute attr, DialectAsmPrinter& os) const override;
- 
-  private:
-   // Adds VHLO types to this dialect.
-diff --ruN a/stablehlo/stablehlo/dialect/VhloOps.td b/stablehlo/stablehlo/dialect/VhloOps.td
---- stablehlo/stablehlo/dialect/VhloOps.td
-+++ stablehlo/stablehlo/dialect/VhloOps.td
-@@ -512,7 +512,8 @@
-   let results = (outs VHLO_AnyType:$result);
- }
- 
--def VHLO_DotGeneralOpV2 : VHLO_Op<"dot_general_v2", "1.6.0", "current"> {
-+def VHLO_DotGeneralOpV2 : VHLO_Op<"dot_general_v2", "1.6.0", "current",
-+    [DeclareOpInterfaceMethods<VHLO_VersionedOpConstraintInterface>]> {
-   let arguments = (ins
-     VHLO_AnyType:$lhs,
-     VHLO_AnyType:$rhs,
-@@ -702,6 +703,13 @@
-   let results = (outs VHLO_AnyType:$result);
- }
- 
-+def VHLO_Exp2OpV1 : VHLO_Op<"exp2_v1", "1.23.0", "current"> {
-+  let arguments = (ins
-+    VHLO_AnyType:$operand,
-+    VHLO_AnyAttr:$result_accuracy);
-+  let results = (outs VHLO_AnyType:$result);
-+}
-+
- def VHLO_ExpOpV1 : VHLO_Op<"exponential_v1", "0.9.0", "1.8.0"> {
-   let arguments = (ins VHLO_AnyType:$operand);
-   let results = (outs VHLO_AnyType:$result);
-@@ -860,6 +868,13 @@
-   let results = (outs VHLO_AnyType:$result);
- }
- 
-+def VHLO_Log2OpV1 : VHLO_Op<"log2_v1", "1.23.0", "current"> {
-+  let arguments = (ins
-+    VHLO_AnyType:$operand,
-+    VHLO_AnyAttr:$result_accuracy);
-+  let results = (outs VHLO_AnyType:$result);
-+}
-+
- def VHLO_MapOpV1 : VHLO_Op<"map_v1", "0.9.0", "current"> {
-   let arguments = (ins
-     Variadic<VHLO_AnyType>:$inputs,
-diff --ruN a/stablehlo/stablehlo/integrations/c/ChloAttributes.cpp b/stablehlo/stablehlo/integrations/c/ChloAttributes.cpp
---- stablehlo/stablehlo/integrations/c/ChloAttributes.cpp
-+++ stablehlo/stablehlo/integrations/c/ChloAttributes.cpp
-@@ -73,12 +73,12 @@
- 
- MlirAttribute chloRaggedDotDimensionNumbersGet(
-     MlirContext ctx, intptr_t nLhsBatchingDimensions,
--    const int64_t *lhsBatchingDimensions, intptr_t nRhsBatchingDimensions,
--    const int64_t *rhsBatchingDimensions, intptr_t nLhsContractingDimensions,
--    const int64_t *lhsContractingDimensions, intptr_t nRhsContractingDimensions,
--    const int64_t *rhsContractingDimensions, intptr_t nLhsRaggedDimensions,
--    const int64_t *lhsRaggedDimensions, intptr_t nRhsGroupDimensions,
--    const int64_t *rhsGroupDimensions) {
-+    const int64_t* lhsBatchingDimensions, intptr_t nRhsBatchingDimensions,
-+    const int64_t* rhsBatchingDimensions, intptr_t nLhsContractingDimensions,
-+    const int64_t* lhsContractingDimensions, intptr_t nRhsContractingDimensions,
-+    const int64_t* rhsContractingDimensions, intptr_t nLhsRaggedDimensions,
-+    const int64_t* lhsRaggedDimensions, intptr_t nRhsGroupDimensions,
-+    const int64_t* rhsGroupDimensions) {
-   return wrap(mlir::chlo::RaggedDotDimensionNumbersAttr::get(
-       unwrap(ctx),
-       llvm::ArrayRef(lhsBatchingDimensions, nLhsBatchingDimensions),
-diff --ruN a/stablehlo/stablehlo/integrations/c/ChloAttributes.h b/stablehlo/stablehlo/integrations/c/ChloAttributes.h
---- stablehlo/stablehlo/integrations/c/ChloAttributes.h
-+++ stablehlo/stablehlo/integrations/c/ChloAttributes.h
-@@ -53,16 +53,16 @@
- 
- MLIR_CAPI_EXPORTED MlirAttribute chloRaggedDotDimensionNumbersGet(
-     MlirContext ctx,                                                        //
--    intptr_t nLhsBatchingDimensions, const int64_t *lhsBatchingDimensions,  //
--    intptr_t nRhsBatchingDimensions, const int64_t *rhsBatchingDimensions,  //
-+    intptr_t nLhsBatchingDimensions, const int64_t* lhsBatchingDimensions,  //
-+    intptr_t nRhsBatchingDimensions, const int64_t* rhsBatchingDimensions,  //
-     intptr_t nLhsContractingDimensions,                                     //
--    const int64_t *lhsContractingDimensions,                                //
-+    const int64_t* lhsContractingDimensions,                                //
-     intptr_t nRhsContractingDimensions,                                     //
--    const int64_t *rhsContractingDimensions,                                //
-+    const int64_t* rhsContractingDimensions,                                //
-     intptr_t nLhsRaggedDimensions,                                          //
--    const int64_t *lhsRaggedDimensions,                                     //
-+    const int64_t* lhsRaggedDimensions,                                     //
-     intptr_t nRhsGroupDimensions,                                           //
--    const int64_t *rhsGroupDimensions);
-+    const int64_t* rhsGroupDimensions);
- 
- MLIR_CAPI_EXPORTED bool chloAttributeIsARaggedDotDimensionNumbers(
-     MlirAttribute attr);
-diff --ruN a/stablehlo/stablehlo/integrations/c/StablehloDialectApi.cpp b/stablehlo/stablehlo/integrations/c/StablehloDialectApi.cpp
---- stablehlo/stablehlo/integrations/c/StablehloDialectApi.cpp
-+++ stablehlo/stablehlo/integrations/c/StablehloDialectApi.cpp
-@@ -48,18 +48,18 @@
- 
- void stablehloVersionFromCompatibilityRequirement(
-     MlirStablehloCompatibilityRequirement requirement,
--    MlirStringCallback callback, void *userData) {
-+    MlirStringCallback callback, void* userData) {
-   mlir::detail::CallbackOstream stream(callback, userData);
-   stream << mlir::vhlo::Version::fromCompatibilityRequirement(
-       unwrapCompatibilityRequirement(requirement));
- }
- 
--void stablehloGetCurrentVersion(MlirStringCallback callback, void *userData) {
-+void stablehloGetCurrentVersion(MlirStringCallback callback, void* userData) {
-   mlir::detail::CallbackOstream stream(callback, userData);
-   stream << mlir::stablehlo::getCurrentVersion();
- }
- 
--void stablehloGetMinimumVersion(MlirStringCallback callback, void *userData) {
-+void stablehloGetMinimumVersion(MlirStringCallback callback, void* userData) {
-   mlir::detail::CallbackOstream stream(callback, userData);
-   stream << mlir::stablehlo::getMinimumVersion();
- }
-@@ -67,7 +67,7 @@
- MlirLogicalResult stablehloGetSmallerVersion(MlirStringRef version1,
-                                              MlirStringRef version2,
-                                              MlirStringCallback callback,
--                                             void *userData) {
-+                                             void* userData) {
-   mlir::detail::CallbackOstream stream(callback, userData);
-   auto result =
-       mlir::stablehlo::getSmallerVersion(unwrap(version1), unwrap(version2));
-@@ -78,7 +78,7 @@
- 
- MlirLogicalResult stablehloSerializePortableArtifactFromModule(
-     MlirModule moduleStr, MlirStringRef targetVersion,
--    MlirStringCallback callback, void *userData, bool allowOtherDialects) {
-+    MlirStringCallback callback, void* userData, bool allowOtherDialects) {
-   mlir::detail::CallbackOstream stream(callback, userData);
-   if (failed(mlir::stablehlo::serializePortableArtifact(
-           unwrap(moduleStr), unwrap(targetVersion), stream,
-@@ -89,7 +89,7 @@
- 
- MlirLogicalResult stablehloSerializePortableArtifactFromStringRef(
-     MlirStringRef moduleStr, MlirStringRef targetVersion,
--    MlirStringCallback callback, void *userData) {
-+    MlirStringCallback callback, void* userData) {
-   mlir::detail::CallbackOstream stream(callback, userData);
-   if (failed(mlir::stablehlo::serializePortableArtifact(
-           unwrap(moduleStr), unwrap(targetVersion), stream)))
-@@ -98,7 +98,7 @@
- }
- 
- MlirLogicalResult stablehloDeserializePortableArtifact(
--    MlirStringRef artifactStr, MlirStringCallback callback, void *userData) {
-+    MlirStringRef artifactStr, MlirStringCallback callback, void* userData) {
-   mlir::detail::CallbackOstream stream(callback, userData);
-   if (failed(mlir::stablehlo::deserializePortableArtifact(unwrap(artifactStr),
-                                                           stream)))
-diff --ruN a/stablehlo/stablehlo/integrations/c/StablehloUnifiedApi.cpp b/stablehlo/stablehlo/integrations/c/StablehloUnifiedApi.cpp
---- stablehlo/stablehlo/integrations/c/StablehloUnifiedApi.cpp
-+++ stablehlo/stablehlo/integrations/c/StablehloUnifiedApi.cpp
-@@ -32,9 +32,9 @@
- #include "stablehlo/reference/Configuration.h"
- 
- MlirAttribute stablehloEvalModule(MlirModule module, int nArgs,
--                                  MlirAttribute const *args,
--                                  const char *const probeInstrumentationDir,
--                                  int *errorCode) {
-+                                  MlirAttribute const* args,
-+                                  const char* const probeInstrumentationDir,
-+                                  int* errorCode) {
-   std::vector<mlir::DenseElementsAttr> inputs;
-   inputs.reserve(nArgs);
-   for (int i = 0; i < nArgs; ++i) {
-@@ -49,7 +49,7 @@
-     return MlirAttribute{nullptr};
-   }
-   std::vector<MlirAttribute> resultsVec;
--  for (const auto &result : results.value()) {
-+  for (const auto& result : results.value()) {
-     resultsVec.push_back(wrap(result));
-   }
-   return mlirArrayAttrGet(mlirModuleGetContext(module), resultsVec.size(),
-diff --ruN a/stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel b/stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel
---- stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel
-+++ stablehlo/stablehlo/integrations/cpp/builder/BUILD.bazel
-@@ -228,6 +228,9 @@
-         ":func_builder",
-         ":mlir_builder",
-         ":stablehlo_builder",
-+        "//:reference_ops",
-+        "//:reference_tensor",
-+        "//:reference_value",
-         "//:register",
-         "//:stablehlo_ops",
-         "@llvm-project//mlir:IR",
-diff --ruN a/stablehlo/stablehlo/integrations/cpp/builder/StablehloBuilderTest.cpp b/stablehlo/stablehlo/integrations/cpp/builder/StablehloBuilderTest.cpp
---- stablehlo/stablehlo/integrations/cpp/builder/StablehloBuilderTest.cpp
-+++ stablehlo/stablehlo/integrations/cpp/builder/StablehloBuilderTest.cpp
-@@ -15,6 +15,7 @@
- 
- #include <complex>
- #include <cstdint>
-+#include <limits>
- #include <string>
- 
- #include "gtest/gtest.h"
-@@ -29,6 +30,12 @@
- #include "mlir/Support/LLVM.h"
- #include "stablehlo/dialect/Register.h"
- #include "stablehlo/dialect/StablehloOps.h"
-+#include "stablehlo/reference/Ops.h"
-+#include "stablehlo/reference/Tensor.h"
-+#include "stablehlo/reference/Value.h"
-+
-+// Include builder headers after reference/Ops.h so StablehloBuilder.h's
-+// Transpose() function does not hide the Transpose enum under -fno-modules.
- #include "stablehlo/integrations/cpp/builder/AttrTypeBuilderUtil.h"
- #include "stablehlo/integrations/cpp/builder/FuncBuilder.h"
- #include "stablehlo/integrations/cpp/builder/MlirBuilder.h"
-@@ -310,6 +317,75 @@
-   EXPECT_EQ(expected, debugString(*module));
- }
- 
-+TEST(MlirBuilderTest, BuildMaxAndArgmaxBodyUsesMaxOp) {
-+  MLIRContext context;
-+  context.loadDialect<StablehloDialect>();
-+  OpBuilder builder(&context);
-+  OwningOpRef<ModuleOp> module = ModuleOp::create(builder.getUnknownLoc());
-+  builder.setInsertionPointToEnd(module->getBody());
-+
-+  buildMaxAndArgmaxBody(builder.getF32Type(), builder.getI32Type(),
-+                        module->getBodyRegion(), builder);
-+
-+  Operation& returnOp = module->getBody()->back();
-+  EXPECT_TRUE(isa<MaxOp>(returnOp.getOperand(0).getDefiningOp()));
-+  EXPECT_TRUE(isa<SelectOp>(returnOp.getOperand(1).getDefiningOp()));
-+}
-+
-+TEST(MlirBuilderTest, BuildMaxAndArgmaxBodyPropagatesNaN) {
-+  MLIRContext context;
-+  context.loadDialect<StablehloDialect>();
-+  OpBuilder builder(&context);
-+  OwningOpRef<ModuleOp> module = ModuleOp::create(builder.getUnknownLoc());
-+  builder.setInsertionPointToEnd(module->getBody());
-+
-+  buildMaxAndArgmaxBody(builder.getF32Type(), builder.getI32Type(),
-+                        module->getBodyRegion(), builder);
-+
-+  auto makeScalar = [&](TypedAttr attr) {
-+    return InterpreterValue(makeTensor(DenseElementsAttr::get(
-+        RankedTensorType::get({}, attr.getType()), attr)));
-+  };
-+  float nan = std::numeric_limits<float>::quiet_NaN();
-+  // SelectOp would drop NaN when lhs=NaN and rhs=5.0; MaxOp propagates NaN.
-+  auto res =
-+      eval(module->getBodyRegion(), {makeScalar(builder.getF32FloatAttr(nan)),
-+                                     makeScalar(builder.getI32IntegerAttr(0)),
-+                                     makeScalar(builder.getF32FloatAttr(5.0f)),
-+                                     makeScalar(builder.getI32IntegerAttr(1))});
-+  EXPECT_TRUE(res[0].getTensor().get({}).getFloatValue().isNaN());
-+}
-+
-+TEST(MlirBuilderTest, BuildMaxAndArgmaxBodySelectsNaNIndex) {
-+  MLIRContext context;
-+  context.loadDialect<StablehloDialect>();
-+  OpBuilder builder(&context);
-+  OwningOpRef<ModuleOp> module = ModuleOp::create(builder.getUnknownLoc());
-+  builder.setInsertionPointToEnd(module->getBody());
-+
-+  buildMaxAndArgmaxBody(builder.getF32Type(), builder.getI32Type(),
-+                        module->getBodyRegion(), builder);
-+
-+  auto evalIndex = [&](float lhs_val, int32_t lhs_idx, float rhs_val,
-+                       int32_t rhs_idx) {
-+    auto makeScalar = [&](TypedAttr attr) {
-+      return InterpreterValue(makeTensor(DenseElementsAttr::get(
-+          RankedTensorType::get({}, attr.getType()), attr)));
-+    };
-+    auto res = eval(module->getBodyRegion(),
-+                    {makeScalar(builder.getF32FloatAttr(lhs_val)),
-+                     makeScalar(builder.getI32IntegerAttr(lhs_idx)),
-+                     makeScalar(builder.getF32FloatAttr(rhs_val)),
-+                     makeScalar(builder.getI32IntegerAttr(rhs_idx))});
-+    EXPECT_TRUE(res[0].getTensor().get({}).getFloatValue().isNaN());
-+    return res[1].getTensor().get({}).getIntegerValue().getSExtValue();
-+  };
-+  float nan = std::numeric_limits<float>::quiet_NaN();
-+  EXPECT_EQ(evalIndex(nan, 0, 5.0f, 1), 0);
-+  EXPECT_EQ(evalIndex(5.0f, 0, nan, 1), 1);
-+  EXPECT_EQ(evalIndex(nan, 2, nan, 1), 1);
-+}
-+
- TEST(MlirBuilderTest, GatherOp) {
-   std::string expected = R"mlir(module {
-   func.func @main(%arg0: tensor<3xi64>, %arg1: tensor<1x1xi64>) -> tensor<1xi64> {
-@@ -1021,7 +1097,7 @@
- TEST(MlirBuilderTest, ConstantF4E2M1FNArray) {
-   std::string expected = R"mlir(module {
-   func.func @main() -> tensor<10xf4E2M1FN> {
--    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 0.000000e+00, 0.000000e+00, 3.000000e+00, 6.000000e+00, 6.000000e+00, 1.000000e+00, 6.000000e+00]> : tensor<10xf4E2M1FN>
-+    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 0.000000e+00, 0.000000e+00, 3.000000e+00, 6.000000e+00, -6.000000e+00, 5.000000e-01, -5.000000e-01]> : tensor<10xf4E2M1FN>
-     return %cst : tensor<10xf4E2M1FN>
-   }
- })mlir";
-@@ -1031,10 +1107,15 @@
-     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-     auto type = makeTensorType(fb.getContext(), {10}, ElementType::F4E2M1FN);
-+    double posMax = 6.0;
-+    double negMax = -6.0;
-+    double posMinDenorm = 0.5;
-+    double negMinDenorm = -0.5;
-     auto cst = stablehlo::Constant(
--        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
--                                           0x07, 0x0F, 0x01, 0x09}),
--                         type));
-+        fb, makeConstant(
-+                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
-+                                  negMax, posMinDenorm, negMinDenorm}),
-+                type));
-     func::Return(fb, {cst});
-   }
- 
-@@ -1045,7 +1126,7 @@
- TEST(MlirBuilderTest, ConstantF6E2M3FNArray) {
-   std::string expected = R"mlir(module {
-   func.func @main() -> tensor<10xf6E2M3FN> {
--    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.250000e-01, 3.250000e+00, 7.500000e+00, 7.500000e+00, 1.000000e+00, 7.500000e+00]> : tensor<10xf6E2M3FN>
-+    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.250000e-01, 3.250000e+00, 7.500000e+00, -7.500000e+00, 1.250000e-01, -1.250000e-01]> : tensor<10xf6E2M3FN>
-     return %cst : tensor<10xf6E2M3FN>
-   }
- })mlir";
-@@ -1055,10 +1136,15 @@
-     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-     auto type = makeTensorType(fb.getContext(), {10}, ElementType::F6E2M3FN);
-+    double posMax = 7.5;
-+    double negMax = -7.5;
-+    double posMinDenorm = 0x1p-3;
-+    double negMinDenorm = -0x1p-3;
-     auto cst = stablehlo::Constant(
--        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
--                                           0x1F, 0x3F, 0x01, 0x21}),
--                         type));
-+        fb, makeConstant(
-+                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
-+                                  negMax, posMinDenorm, negMinDenorm}),
-+                type));
-     func::Return(fb, {cst});
-   }
- 
-@@ -1069,7 +1155,7 @@
- TEST(MlirBuilderTest, ConstantF6E3M2FNArray) {
-   std::string expected = R"mlir(module {
-   func.func @main() -> tensor<10xf6E3M2FN> {
--    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.250000e-01, 3.000000e+00, 2.800000e+01, 2.800000e+01, 1.000000e+00, 2.800000e+01]> : tensor<10xf6E3M2FN>
-+    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.250000e-01, 3.000000e+00, 2.800000e+01, -2.800000e+01, 6.250000e-02, -6.250000e-02]> : tensor<10xf6E3M2FN>
-     return %cst : tensor<10xf6E3M2FN>
-   }
- })mlir";
-@@ -1079,10 +1165,15 @@
-     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-     auto type = makeTensorType(fb.getContext(), {10}, ElementType::F6E3M2FN);
-+    double posMax = 28.0;
-+    double negMax = -28.0;
-+    double posMinDenorm = 0x1p-4;
-+    double negMinDenorm = -0x1p-4;
-     auto cst = stablehlo::Constant(
--        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
--                                           0x1F, 0x3F, 0x01, 0x21}),
--                         type));
-+        fb, makeConstant(
-+                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
-+                                  negMax, posMinDenorm, negMinDenorm}),
-+                type));
-     func::Return(fb, {cst});
-   }
- 
-@@ -1092,9 +1183,9 @@
- 
- TEST(MlirBuilderTest, ConstantF8E3M4Array) {
-   std::string expected = R"mlir(module {
--  func.func @main() -> tensor<10xf8E3M4> {
--    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.125000e+00, 0x70, 0x70, 1.000000e+00, 0x70]> : tensor<10xf8E3M4>
--    return %cst : tensor<10xf8E3M4>
-+  func.func @main() -> tensor<11xf8E3M4> {
-+    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.125000e+00, 0x70, 0xF0, 0x78, 1.562500e-02, -1.562500e-02]> : tensor<11xf8E3M4>
-+    return %cst : tensor<11xf8E3M4>
-   }
- })mlir";
- 
-@@ -1102,11 +1193,17 @@
-   {
-     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
--    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E3M4);
-+    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E3M4);
-+    double posInf = std::numeric_limits<double>::infinity();
-+    double negInf = -std::numeric_limits<double>::infinity();
-+    double nan = std::numeric_limits<double>::quiet_NaN();
-+    double posMinDenorm = 0x1p-6;
-+    double negMinDenorm = -0x1p-6;
-     auto cst = stablehlo::Constant(
--        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
--                                           0x7F, 0xFF, 0x01, 0x81}),
--                         type));
-+        fb, makeConstant(
-+                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posInf,
-+                                  negInf, nan, posMinDenorm, negMinDenorm}),
-+                type));
-     func::Return(fb, {cst});
-   }
- 
-@@ -1116,9 +1213,9 @@
- 
- TEST(MlirBuilderTest, ConstantF8E4M3B11FNUZArray) {
-   std::string expected = R"mlir(module {
--  func.func @main() -> tensor<10xf8E4M3B11FNUZ> {
--    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 0x80, 0x80, 1.000000e+00, 0x80]> : tensor<10xf8E4M3B11FNUZ>
--    return %cst : tensor<10xf8E4M3B11FNUZ>
-+  func.func @main() -> tensor<11xf8E4M3B11FNUZ> {
-+    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 3.000000e+01, -3.000000e+01, 0x80, 1.220700e-04, -1.220700e-04]> : tensor<11xf8E4M3B11FNUZ>
-+    return %cst : tensor<11xf8E4M3B11FNUZ>
-   }
- })mlir";
- 
-@@ -1127,11 +1224,17 @@
-     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-     auto type =
--        makeTensorType(fb.getContext(), {10}, ElementType::F8E4M3B11FNUZ);
-+        makeTensorType(fb.getContext(), {11}, ElementType::F8E4M3B11FNUZ);
-+    double posMax = 30.0;
-+    double negMax = -30.0;
-+    double nan = std::numeric_limits<double>::quiet_NaN();
-+    double posMinDenorm = 0x1p-13;
-+    double negMinDenorm = -0x1p-13;
-     auto cst = stablehlo::Constant(
--        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
--                                           0x7F, 0xFF, 0x01, 0x81}),
--                         type));
-+        fb, makeConstant(
-+                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
-+                                  negMax, nan, posMinDenorm, negMinDenorm}),
-+                type));
-     func::Return(fb, {cst});
-   }
- 
-@@ -1141,9 +1244,9 @@
- 
- TEST(MlirBuilderTest, ConstantF8E4M3Array) {
-   std::string expected = R"mlir(module {
--  func.func @main() -> tensor<10xf8E4M3> {
--    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 1.280000e+02, 0x78, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E4M3>
--    return %cst : tensor<10xf8E4M3>
-+  func.func @main() -> tensor<11xf8E4M3> {
-+    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 0x78, 0xF8, 0x7C, 1.953130e-03, -1.953130e-03]> : tensor<11xf8E4M3>
-+    return %cst : tensor<11xf8E4M3>
-   }
- })mlir";
- 
-@@ -1151,11 +1254,17 @@
-   {
-     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
--    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E4M3);
-+    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E4M3);
-+    double posInf = std::numeric_limits<double>::infinity();
-+    double negInf = -std::numeric_limits<double>::infinity();
-+    double nan = std::numeric_limits<double>::quiet_NaN();
-+    double posMinDenorm = 0x1p-9;
-+    double negMinDenorm = -0x1p-9;
-     auto cst = stablehlo::Constant(
--        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
--                                           0x7F, 0xFF, 0x01, 0x81}),
--                         type));
-+        fb, makeConstant(
-+                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posInf,
-+                                  negInf, nan, posMinDenorm, negMinDenorm}),
-+                type));
-     func::Return(fb, {cst});
-   }
- 
-@@ -1165,9 +1274,9 @@
- 
- TEST(MlirBuilderTest, ConstantF8E4M3FNArray) {
-   std::string expected = R"mlir(module {
--  func.func @main() -> tensor<10xf8E4M3FN> {
--    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 1.280000e+02, 2.560000e+02, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E4M3FN>
--    return %cst : tensor<10xf8E4M3FN>
-+  func.func @main() -> tensor<11xf8E4M3FN> {
-+    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 4.480000e+02, -4.480000e+02, 0x7F, 1.953130e-03, -1.953130e-03]> : tensor<11xf8E4M3FN>
-+    return %cst : tensor<11xf8E4M3FN>
-   }
- })mlir";
- 
-@@ -1175,11 +1284,17 @@
-   {
-     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
--    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E4M3FN);
-+    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E4M3FN);
-+    double posMax = 448.0;
-+    double negMax = -448.0;
-+    double nan = std::numeric_limits<double>::quiet_NaN();
-+    double posMinDenorm = 0x1p-9;
-+    double negMinDenorm = -0x1p-9;
-     auto cst = stablehlo::Constant(
--        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
--                                           0x7F, 0xFF, 0x01, 0x81}),
--                         type));
-+        fb, makeConstant(
-+                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
-+                                  negMax, nan, posMinDenorm, negMinDenorm}),
-+                type));
-     func::Return(fb, {cst});
-   }
- 
-@@ -1189,9 +1304,9 @@
- 
- TEST(MlirBuilderTest, ConstantF8E4M3FNUZArray) {
-   std::string expected = R"mlir(module {
--  func.func @main() -> tensor<10xf8E4M3FNUZ> {
--    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 1.280000e+02, 0x80, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E4M3FNUZ>
--    return %cst : tensor<10xf8E4M3FNUZ>
-+  func.func @main() -> tensor<11xf8E4M3FNUZ> {
-+    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 2.400000e+02, -2.400000e+02, 0x80, 9.765620e-04, -9.765620e-04]> : tensor<11xf8E4M3FNUZ>
-+    return %cst : tensor<11xf8E4M3FNUZ>
-   }
- })mlir";
- 
-@@ -1199,11 +1314,17 @@
-   {
-     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
--    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E4M3FNUZ);
-+    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E4M3FNUZ);
-+    double posMax = 240.0;
-+    double negMax = -240.0;
-+    double nan = std::numeric_limits<double>::quiet_NaN();
-+    double posMinDenorm = 0x1p-10;
-+    double negMinDenorm = -0x1p-10;
-     auto cst = stablehlo::Constant(
--        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
--                                           0x7F, 0xFF, 0x01, 0x81}),
--                         type));
-+        fb, makeConstant(
-+                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
-+                                  negMax, nan, posMinDenorm, negMinDenorm}),
-+                type));
-     func::Return(fb, {cst});
-   }
- 
-@@ -1213,9 +1334,9 @@
- 
- TEST(MlirBuilderTest, ConstantF8E5M2Array) {
-   std::string expected = R"mlir(module {
--  func.func @main() -> tensor<10xf8E5M2> {
--    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.000000e+00, 1.280000e+02, 2.560000e+02, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E5M2>
--    return %cst : tensor<10xf8E5M2>
-+  func.func @main() -> tensor<11xf8E5M2> {
-+    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.000000e+00, 0x7C, 0xFC, 0x7E, 1.525880e-05, -1.525880e-05]> : tensor<11xf8E5M2>
-+    return %cst : tensor<11xf8E5M2>
-   }
- })mlir";
- 
-@@ -1223,11 +1344,17 @@
-   {
-     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
--    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E5M2);
-+    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E5M2);
-+    double posInf = std::numeric_limits<double>::infinity();
-+    double negInf = -std::numeric_limits<double>::infinity();
-+    double nan = std::numeric_limits<double>::quiet_NaN();
-+    double posMinDenorm = 0x1p-16;
-+    double negMinDenorm = -0x1p-16;
-     auto cst = stablehlo::Constant(
--        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
--                                           0x7F, 0xFF, 0x01, 0x81}),
--                         type));
-+        fb, makeConstant(
-+                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posInf,
-+                                  negInf, nan, posMinDenorm, negMinDenorm}),
-+                type));
-     func::Return(fb, {cst});
-   }
- 
-@@ -1237,9 +1364,9 @@
- 
- TEST(MlirBuilderTest, ConstantF8E5M2FNUZArray) {
-   std::string expected = R"mlir(module {
--  func.func @main() -> tensor<10xf8E5M2FNUZ> {
--    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.000000e+00, 1.280000e+02, 2.560000e+02, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E5M2FNUZ>
--    return %cst : tensor<10xf8E5M2FNUZ>
-+  func.func @main() -> tensor<11xf8E5M2FNUZ> {
-+    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.000000e+00, 5.734400e+04, -5.734400e+04, 0x80, 7.629390e-06, -7.629390e-06]> : tensor<11xf8E5M2FNUZ>
-+    return %cst : tensor<11xf8E5M2FNUZ>
-   }
- })mlir";
- 
-@@ -1247,11 +1374,17 @@
-   {
-     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
--    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E5M2FNUZ);
-+    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E5M2FNUZ);
-+    double posMax = 57344.0;
-+    double negMax = -57344.0;
-+    double nan = std::numeric_limits<double>::quiet_NaN();
-+    double posMinDenorm = 0x1p-17;
-+    double negMinDenorm = -0x1p-17;
-     auto cst = stablehlo::Constant(
--        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
--                                           0x7F, 0xFF, 0x01, 0x81}),
--                         type));
-+        fb, makeConstant(
-+                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
-+                                  negMax, nan, posMinDenorm, negMinDenorm}),
-+                type));
-     func::Return(fb, {cst});
-   }
- 
-@@ -1262,7 +1395,7 @@
- TEST(MlirBuilderTest, ConstantF8E8M0FNUArray) {
-   std::string expected = R"mlir(module {
-   func.func @main() -> tensor<8xf8E8M0FNU> {
--    %cst = stablehlo.constant dense<[5.877470e-39, 1.000000e+00, 1.250000e-01, 1.250000e-01, 4.000000e+00, 5.877470e-39, 1.280000e+02, 2.560000e+02]> : tensor<8xf8E8M0FNU>
-+    %cst = stablehlo.constant dense<[5.877470e-39, 1.000000e+00, 1.250000e-01, 1.250000e-01, 4.000000e+00, 5.877470e-39, 2.000000e+00, 0xFF]> : tensor<8xf8E8M0FNU>
-     return %cst : tensor<8xf8E8M0FNU>
-   }
- })mlir";
-@@ -1272,10 +1405,12 @@
-     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-     auto type = makeTensorType(fb.getContext(), {8}, ElementType::F8E8M0FNU);
-+    double min = 0x1p-127;
-+    double nan = std::numeric_limits<double>::quiet_NaN();
-     auto cst = stablehlo::Constant(
--        fb, makeConstant(ArrayRef<double>(
--                             {0.0, 1.0, 0.125, 0.1, 3.1415, 0x00, 0x80, 0xFF}),
--                         type));
-+        fb, makeConstant(
-+                ArrayRef<double>({0.0, 1.0, 0.125, 0.1, 3.1415, min, 2.0, nan}),
-+                type));
-     func::Return(fb, {cst});
-   }
- 
-@@ -1286,7 +1421,7 @@
- TEST(MlirBuilderTest, ConstantBF16Array) {
-   std::string expected = R"mlir(module {
-   func.func @main() -> tensor<11xbf16> {
--    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000980e-01, 3.140630e+00, 3.264000e+04, 6.553600e+04, 3.276800e+04, 1.000000e+00, 3.276800e+04]> : tensor<11xbf16>
-+    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000980e-01, 3.140630e+00, 0x7F80, 0xFF80, 0x7FC0, 9.183550e-41, -9.183550e-41]> : tensor<11xbf16>
-     return %cst : tensor<11xbf16>
-   }
- })mlir";
-@@ -1296,11 +1431,16 @@
-     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-     auto type = makeTensorType(fb.getContext(), {11}, ElementType::BF16);
-+    double posInf = std::numeric_limits<double>::infinity();
-+    double negInf = -std::numeric_limits<double>::infinity();
-+    double nan = std::numeric_limits<double>::quiet_NaN();
-+    double posMinDenorm = 0x1p-133;
-+    double negMinDenorm = -0x1p-133;
-     auto cst = stablehlo::Constant(
--        fb,
--        makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.140630,
--                                       0x7F80, 0xFF80, 0x7FFF, 0x0001, 0x8001}),
--                     type));
-+        fb, makeConstant(
-+                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.140630, posInf,
-+                                  negInf, nan, posMinDenorm, negMinDenorm}),
-+                type));
-     func::Return(fb, {cst});
-   }
- 
-@@ -1311,7 +1451,7 @@
- TEST(MlirBuilderTest, ConstantF16Array) {
-   std::string expected = R"mlir(module {
-   func.func @main() -> tensor<11xf16> {
--    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.997550e-02, 3.140630e+00, 3.174400e+04, 6.451200e+04, 3.276800e+04, 1.000000e+00, 3.276800e+04]> : tensor<11xf16>
-+    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.997550e-02, 3.140630e+00, 0x7C00, 0xFC00, 0x7E00, 5.960460e-08, -5.960460e-08]> : tensor<11xf16>
-     return %cst : tensor<11xf16>
-   }
- })mlir";
-@@ -1321,11 +1461,16 @@
-     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-     auto type = makeTensorType(fb.getContext(), {11}, ElementType::F16);
-+    double posInf = std::numeric_limits<double>::infinity();
-+    double negInf = -std::numeric_limits<double>::infinity();
-+    double nan = std::numeric_limits<double>::quiet_NaN();
-+    double posMinDenorm = 0x1p-24;
-+    double negMinDenorm = -0x1p-24;
-     auto cst = stablehlo::Constant(
--        fb,
--        makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.140630,
--                                       0x7C00, 0xFC00, 0x7FFF, 0x0001, 0x8001}),
--                     type));
-+        fb, makeConstant(
-+                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.140630, posInf,
-+                                  negInf, nan, posMinDenorm, negMinDenorm}),
-+                type));
-     func::Return(fb, {cst});
-   }
- 
-@@ -1336,7 +1481,7 @@
- TEST(MlirBuilderTest, ConstantF32Array) {
-   std::string expected = R"mlir(module {
-   func.func @main() -> tensor<11xf32> {
--    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000000e-01, 3.14159274, 2.13909504E+9, 4.28657869E+9, 2.14748365E+9, 1.000000e+00, 2.14748365E+9]> : tensor<11xf32>
-+    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000000e-01, 3.14159274, 0x7F800000, 0xFF800000, 0x7FC00000, 1.401300e-45, -1.401300e-45]> : tensor<11xf32>
-     return %cst : tensor<11xf32>
-   }
- })mlir";
-@@ -1346,10 +1491,15 @@
-     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-     auto type = makeTensorType(fb.getContext(), {11}, ElementType::F32);
-+    double posInf = std::numeric_limits<double>::infinity();
-+    double negInf = -std::numeric_limits<double>::infinity();
-+    double nan = std::numeric_limits<double>::quiet_NaN();
-+    double posMinDenorm = std::numeric_limits<float>::denorm_min();
-+    double negMinDenorm = -std::numeric_limits<float>::denorm_min();
-     auto cst = stablehlo::Constant(
-         fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1,
--                                           3.14159274, 0x7F800000, 0xFF800000,
--                                           0x7FFFFFFF, 0x00000001, 0x80000001}),
-+                                           3.14159274, posInf, negInf, nan,
-+                                           posMinDenorm, negMinDenorm}),
-                          type));
-     func::Return(fb, {cst});
-   }
-@@ -1361,7 +1511,7 @@
- TEST(MlirBuilderTest, ConstantF64Array) {
-   std::string expected = R"mlir(module {
-   func.func @main() -> tensor<11xf64> {
--    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000000e-01, 3.1415926535897931, 9.2188684372274053E+18, 1.8442240474082181E+19, 9.2233720368547758E+18, 1.000000e+00, 9.2233720368547758E+18]> : tensor<11xf64>
-+    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000000e-01, 3.1415926535897931, 0x7FF0000000000000, 0xFFF0000000000000, 0x7FF8000000000000, 4.940660e-324, -4.940660e-324]> : tensor<11xf64>
-     return %cst : tensor<11xf64>
-   }
- })mlir";
-@@ -1371,13 +1521,16 @@
-     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-     auto type = makeTensorType(fb.getContext(), {11}, ElementType::F64);
-+    double posInf = std::numeric_limits<double>::infinity();
-+    double negInf = -std::numeric_limits<double>::infinity();
-+    double nan = std::numeric_limits<double>::quiet_NaN();
-+    double posMinDenorm = std::numeric_limits<double>::denorm_min();
-+    double negMinDenorm = -std::numeric_limits<double>::denorm_min();
-     auto cst = stablehlo::Constant(
--        fb, makeConstant(
--                ArrayRef<long double>({0.0, -0.0, 1.0, 0.125, 0.1,
--                                       3.1415926535897931, 0x7FF0000000000000,
--                                       0xFFF0000000000000, 0x7FFFFFFFFFFFFFFF,
--                                       0x0000000000000001, 0x8000000000000001}),
--                type));
-+        fb, makeConstant(ArrayRef<long double>(
-+                             {0.0, -0.0, 1.0, 0.125, 0.1, 3.1415926535897931,
-+                              posInf, negInf, nan, posMinDenorm, negMinDenorm}),
-+                         type));
-     func::Return(fb, {cst});
-   }
- 
-diff --ruN a/stablehlo/stablehlo/integrations/python/StablehloApi.cpp b/stablehlo/stablehlo/integrations/python/StablehloApi.cpp
---- stablehlo/stablehlo/integrations/python/StablehloApi.cpp
-+++ stablehlo/stablehlo/integrations/python/StablehloApi.cpp
-@@ -42,15 +42,15 @@
-   StringWriterHelper() : ss_(s_) {}
- 
-   static MlirStringCallback getMlirStringCallback() {
--    return [](MlirStringRef string_ref, void *user_data) {
--      auto *helper = static_cast<StringWriterHelper *>(user_data);
-+    return [](MlirStringRef string_ref, void* user_data) {
-+      auto* helper = static_cast<StringWriterHelper*>(user_data);
-       helper->ss_ << llvm::StringRef(string_ref.data, string_ref.length);
-     };
-   }
- 
--  void *getUserData() { return static_cast<void *>(this); }
-+  void* getUserData() { return static_cast<void*>(this); }
- 
--  const std::string &toString() {
-+  const std::string& toString() {
-     ss_.flush();
-     return s_;
-   }
-@@ -60,7 +60,7 @@
-   llvm::raw_string_ostream ss_;
- };
- 
--static MlirStringRef toMlirStringRef(const std::string &s) {
-+static MlirStringRef toMlirStringRef(const std::string& s) {
-   return mlirStringRefCreate(s.data(), s.size());
- }
- 
-@@ -68,11 +68,11 @@
-   return mlirStringRefCreate(s.data(), s.size());
- }
- 
--static MlirStringRef toMlirStringRef(const nb::bytes &s) {
--  return mlirStringRefCreate(static_cast<const char *>(s.data()), s.size());
-+static MlirStringRef toMlirStringRef(const nb::bytes& s) {
-+  return mlirStringRefCreate(static_cast<const char*>(s.data()), s.size());
- }
- 
--void AddStablehloApi(nb::module_ &m) {
-+void AddStablehloApi(nb::module_& m) {
-   // Portable API is a subset of StableHLO API
-   AddPortableApi(m);
- 
-@@ -146,8 +146,8 @@
-   //
-   m.def(
-       "eval_module",
--      [](MlirModule module, std::vector<MlirAttribute> &args,
--         const std::string &probe_instrumentation_dir)
-+      [](MlirModule module, std::vector<MlirAttribute>& args,
-+         const std::string& probe_instrumentation_dir)
-           -> std::vector<MlirAttribute> {
-         for (auto arg : args) {
-           if (!mlirAttributeIsADenseElements(arg)) {
-@@ -174,7 +174,7 @@
-       nb::arg("probe_instrumentation_dir") = "");
- }
- 
--void AddPortableApi(nb::module_ &m) {
-+void AddPortableApi(nb::module_& m) {
-   //
-   // Utility APIs.
-   //
-@@ -182,8 +182,8 @@
- 
-   m.def(
-       "get_smaller_version",
--      [](const std::string &version1,
--         const std::string &version2) -> std::string {
-+      [](const std::string& version1,
-+         const std::string& version2) -> std::string {
-         StringWriterHelper accumulator;
-         if (mlirLogicalResultIsFailure(stablehloGetSmallerVersion(
-                 toMlirStringRef(version1), toMlirStringRef(version2),
-@@ -263,7 +263,7 @@
-       nb::arg("artifact_str"));
-   m.def(
-       "deserialize_portable_artifact_str",
--      [](const nb::bytes &artifact) -> nb::bytes {
-+      [](const nb::bytes& artifact) -> nb::bytes {
-         StringWriterHelper accumulator;
-         if (mlirLogicalResultIsFailure(stablehloDeserializePortableArtifact(
-                 toMlirStringRef(artifact), accumulator.getMlirStringCallback(),
-diff --ruN a/stablehlo/stablehlo/reference/Api.cpp b/stablehlo/stablehlo/reference/Api.cpp
---- stablehlo/stablehlo/reference/Api.cpp
-+++ stablehlo/stablehlo/reference/Api.cpp
-@@ -83,7 +83,7 @@
- class DefaultInterpreterFallback : public InterpreterFallback {
-  public:
-   DefaultInterpreterFallback(const InterpreterConfiguration& config)
--      : config(config){};
-+      : config(config) {}
- 
-   virtual llvm::Error operator()(Operation& op, Scope& scope,
-                                  Process* process) final {
-diff --ruN a/stablehlo/stablehlo/reference/Api.h b/stablehlo/stablehlo/reference/Api.h
---- stablehlo/stablehlo/reference/Api.h
-+++ stablehlo/stablehlo/reference/Api.h
-@@ -36,17 +36,17 @@
- /// builtin kernels are matched.
- FailureOr<SmallVector<InterpreterValue>> evalModule(
-     ModuleOp module, ArrayRef<InterpreterValue> inputs,
--    const InterpreterConfiguration &config);
-+    const InterpreterConfiguration& config);
- 
- /// This wrapper is intended to be easily used by the StableHLO Python bindings.
- // It wraps the InterpreterValue API.
- FailureOr<SmallVector<DenseElementsAttr>> evalModule(
-     ModuleOp module, ArrayRef<DenseElementsAttr> inputs,
--    const InterpreterConfiguration &config);
-+    const InterpreterConfiguration& config);
- 
- /// Parses a StableHLO MLIR text program into a ModuleOp.
--FailureOr<OwningOpRef<ModuleOp>> parseStablehloModule(const std::string &mlir,
--                                                      MLIRContext &context);
-+FailureOr<OwningOpRef<ModuleOp>> parseStablehloModule(const std::string& mlir,
-+                                                      MLIRContext& context);
- 
- }  // namespace stablehlo
- }  // namespace mlir
-diff --ruN a/stablehlo/stablehlo/reference/Axes.cpp b/stablehlo/stablehlo/reference/Axes.cpp
---- stablehlo/stablehlo/reference/Axes.cpp
-+++ stablehlo/stablehlo/reference/Axes.cpp
-@@ -21,7 +21,7 @@
- namespace mlir {
- namespace stablehlo {
- 
--raw_ostream &operator<<(raw_ostream &os, const Axes &x) {
-+raw_ostream& operator<<(raw_ostream& os, const Axes& x) {
-   os << "[";
-   llvm::interleave(x, os, ", ");
-   os << "]";
-diff --ruN a/stablehlo/stablehlo/reference/Axes.h b/stablehlo/stablehlo/reference/Axes.h
---- stablehlo/stablehlo/reference/Axes.h
-+++ stablehlo/stablehlo/reference/Axes.h
-@@ -29,8 +29,8 @@
- class Axes : public SmallVector<int64_t> {
-  public:
-   Axes() = default;
--  Axes(const Axes &other) = default;
--  Axes &operator=(const Axes &other) = default;
-+  Axes(const Axes& other) = default;
-+  Axes& operator=(const Axes& other) = default;
- 
-   Axes(std::initializer_list<int64_t> list) : SmallVector(list) {}
-   explicit Axes(size_t size, int64_t element = 0)
-@@ -40,7 +40,7 @@
-       : SmallVector(attr.getValues<int64_t>()) {}
- };
- 
--raw_ostream &operator<<(raw_ostream &os, const Axes &x);
-+raw_ostream& operator<<(raw_ostream& os, const Axes& x);
- 
- }  // namespace stablehlo
- }  // namespace mlir
-diff --ruN a/stablehlo/stablehlo/reference/Configuration.cpp b/stablehlo/stablehlo/reference/Configuration.cpp
---- stablehlo/stablehlo/reference/Configuration.cpp
-+++ stablehlo/stablehlo/reference/Configuration.cpp
-@@ -25,8 +25,8 @@
- namespace mlir {
- namespace stablehlo {
- 
--llvm::Error InterpreterFallback::operator()(Operation &op, Scope &scope,
--                                            Process *process) {
-+llvm::Error InterpreterFallback::operator()(Operation& op, Scope& scope,
-+                                            Process* process) {
-   return stablehlo::invalidArgument("Unsupported op: %s",
-                                     debugString(op).c_str());
- }
-diff --ruN a/stablehlo/stablehlo/reference/Configuration.h b/stablehlo/stablehlo/reference/Configuration.h
---- stablehlo/stablehlo/reference/Configuration.h
-+++ stablehlo/stablehlo/reference/Configuration.h
-@@ -30,7 +30,7 @@
-  public:
-   /// Custom op kernels for any user specified ops not found in the StableHLO
-   /// op dialect or StableHLO interpreter dialect.
--  virtual llvm::Error operator()(Operation &op, Scope &scope, Process *process);
-+  virtual llvm::Error operator()(Operation& op, Scope& scope, Process* process);
- 
-   virtual ~InterpreterFallback() = default;
- };
-diff --ruN a/stablehlo/stablehlo/reference/Element.cpp b/stablehlo/stablehlo/reference/Element.cpp
---- stablehlo/stablehlo/reference/Element.cpp
-+++ stablehlo/stablehlo/reference/Element.cpp
-@@ -844,6 +844,14 @@
-       });
- }
- 
-+Element exp2(const Element& el) {
-+  return mapWithUpcastToDouble(
-+      el, [](double e) { return std::exp2(e); },
-+      [](mlir::Complex<double> e) {
-+        return std::pow(mlir::Complex<double>(2.0, 0.0), e);
-+      });
-+}
-+
- Element floor(const Element& el) {
-   APFloat val = el.getFloatValue();
-   val.roundToIntegral(APFloat::rmTowardNegative);
-@@ -890,6 +898,12 @@
-       });
- }
- 
-+Element log2(const Element& el) {
-+  return mapWithUpcastToDouble(
-+      el, [](double e) { return std::log2(e); },
-+      [](mlir::Complex<double> e) { return std::log(e) / std::log(2.0); });
-+}
-+
- Element logistic(const Element& el) {
-   auto one = convert(el.getType(), 1.0);
-   return one / (one + exponential(-el));
-@@ -1159,6 +1173,40 @@
-       [](mlir::Complex<double> e) { return std::tanh(e); });
- }
- 
-+Element weakOrderEq(const Element& e1, const Element& e2) {
-+  auto type = e1.getType();
-+  auto i1Type = IntegerType::get(type.getContext(), 1);
-+  if (type != e2.getType())
-+    report_fatal_error(invalidArgument("Element types don't match: %s vs %s",
-+                                       debugString(type).c_str(),
-+                                       debugString(e2.getType()).c_str()));
-+  if (!isSupportedFloatType(type))
-+    report_fatal_error(invalidArgument("Unsupported element type: %s",
-+                                       debugString(type).c_str()));
-+  auto floatLhs = e1.getFloatValue();
-+  auto floatRhs = e2.getFloatValue();
-+  return Element(
-+      i1Type, (floatLhs == floatRhs) || (floatLhs.isNaN() && floatRhs.isNaN()));
-+}
-+
-+Element weakOrderLt(const Element& e1, const Element& e2) {
-+  auto type = e1.getType();
-+  auto i1Type = IntegerType::get(type.getContext(), 1);
-+  if (type != e2.getType())
-+    report_fatal_error(invalidArgument("Element types don't match: %s vs %s",
-+                                       debugString(type).c_str(),
-+                                       debugString(e2.getType()).c_str()));
-+  if (!isSupportedFloatType(type))
-+    report_fatal_error(invalidArgument("Unsupported element type: %s",
-+                                       debugString(type).c_str()));
-+  auto floatLhs = e1.getFloatValue();
-+  auto floatRhs = e2.getFloatValue();
-+  bool result =
-+      (!floatLhs.isNaN() && floatRhs.isNaN()) ||
-+      (!floatLhs.isNaN() && !floatRhs.isNaN() && (floatLhs < floatRhs));
-+  return Element(i1Type, result);
-+}
-+
- void Element::print(raw_ostream& os, bool elideType) const {
-   if (isSupportedIntegerType(type_)) {
-     IntegerAttr::get(type_, getIntegerValue()).print(os, elideType);
-diff --ruN a/stablehlo/stablehlo/reference/Element.h b/stablehlo/stablehlo/reference/Element.h
---- stablehlo/stablehlo/reference/Element.h
-+++ stablehlo/stablehlo/reference/Element.h
-@@ -237,6 +237,9 @@
- /// Returns exponential_minus_one of Element object.
- Element exponentialMinusOne(const Element& el);
- 
-+/// Returns exp2 of Element object.
-+Element exp2(const Element& el);
-+
- /// Returns floor of Element object.
- Element floor(const Element& e);
- 
-@@ -253,6 +256,9 @@
- /// Returns log1p of Element object.
- Element logPlusOne(const Element& el);
- 
-+/// Returns log2 of Element object.
-+Element log2(const Element& el);
-+
- /// Returns logistic of Element object.
- Element logistic(const Element& el);
- 
-@@ -314,6 +320,12 @@
- /// Returns tanh of Element object.
- Element tanh(const Element& e);
- 
-+/// Returns weak-order equality of two floating-point Element objects.
-+Element weakOrderEq(const Element& e1, const Element& e2);
-+
-+/// Returns weak-order less-than of two floating-point Element objects.
-+Element weakOrderLt(const Element& e1, const Element& e2);
-+
- /// Print utilities for Element objects.
- inline raw_ostream& operator<<(raw_ostream& os, Element element) {
-   element.print(os, /*elideType=*/true);
-diff --ruN a/stablehlo/stablehlo/reference/Errors.h b/stablehlo/stablehlo/reference/Errors.h
---- stablehlo/stablehlo/reference/Errors.h
-+++ stablehlo/stablehlo/reference/Errors.h
-@@ -26,7 +26,7 @@
- /// error using a format string and a variadic number of arguments to the format
- /// string.
- template <typename... Ts>
--inline llvm::Error invalidArgument(char const *Fmt, const Ts &...Vals) {
-+inline llvm::Error invalidArgument(char const* Fmt, const Ts&... Vals) {
-   return createStringError(llvm::errc::invalid_argument, Fmt, Vals...);
- }
- 
-diff --ruN a/stablehlo/stablehlo/reference/Index.cpp b/stablehlo/stablehlo/reference/Index.cpp
---- stablehlo/stablehlo/reference/Index.cpp
-+++ stablehlo/stablehlo/reference/Index.cpp
-@@ -26,7 +26,7 @@
- namespace mlir {
- namespace stablehlo {
- 
--raw_ostream &operator<<(raw_ostream &os, const Sizes &x) {
-+raw_ostream& operator<<(raw_ostream& os, const Sizes& x) {
-   os << "[";
-   llvm::interleave(x, os, ", ");
-   os << "]";
-@@ -40,7 +40,7 @@
-   return result;
- }
- 
--bool Sizes::inBounds(const Sizes &bounds) const {
-+bool Sizes::inBounds(const Sizes& bounds) const {
-   if (size() != bounds.size()) return false;
-   for (auto [size, bound] : llvm::zip(*this, bounds))
-     if (size < 0 || size >= bound) return false;
-@@ -59,7 +59,7 @@
-   return IndexSpaceIterator(*this);
- }
- 
--Sizes operator+(const Sizes &x, const Sizes &y) {
-+Sizes operator+(const Sizes& x, const Sizes& y) {
-   if (x.size() != y.size()) llvm::report_fatal_error("expected same size");
-   Sizes result(x.size());
-   for (size_t i = 0; i < x.size(); ++i) {
-@@ -68,11 +68,11 @@
-   return result;
- }
- 
--Sizes operator+(const Sizes &x, int64_t y) { return x + Sizes(x.size(), y); }
-+Sizes operator+(const Sizes& x, int64_t y) { return x + Sizes(x.size(), y); }
- 
--Sizes operator+(int64_t x, const Sizes &y) { return y + x; }
-+Sizes operator+(int64_t x, const Sizes& y) { return y + x; }
- 
--Sizes operator-(const Sizes &x, const Sizes &y) {
-+Sizes operator-(const Sizes& x, const Sizes& y) {
-   if (x.size() != y.size()) llvm::report_fatal_error("expected same size");
-   Sizes result(x.size());
-   for (size_t i = 0; i < x.size(); ++i) {
-@@ -81,11 +81,11 @@
-   return result;
- }
- 
--Sizes operator-(const Sizes &x, int64_t y) { return x - Sizes(x.size(), y); }
-+Sizes operator-(const Sizes& x, int64_t y) { return x - Sizes(x.size(), y); }
- 
--Sizes operator-(int64_t x, const Sizes &y) { return Sizes(y.size(), x) - y; }
-+Sizes operator-(int64_t x, const Sizes& y) { return Sizes(y.size(), x) - y; }
- 
--Sizes operator*(const Sizes &x, const Sizes &y) {
-+Sizes operator*(const Sizes& x, const Sizes& y) {
-   if (x.size() != y.size()) llvm::report_fatal_error("expected same size");
-   Sizes result(x.size());
-   for (size_t i = 0; i < x.size(); ++i) {
-@@ -94,23 +94,23 @@
-   return result;
- }
- 
--Sizes operator*(const Sizes &x, int64_t y) { return x * Sizes(x.size(), y); }
-+Sizes operator*(const Sizes& x, int64_t y) { return x * Sizes(x.size(), y); }
- 
--Sizes operator*(int64_t &x, const Sizes &y) { return y + x; }
-+Sizes operator*(int64_t& x, const Sizes& y) { return y + x; }
- 
--Sizes clamp(int64_t min, const Sizes &x, int64_t max) {
-+Sizes clamp(int64_t min, const Sizes& x, int64_t max) {
-   return clamp(Sizes(x.size(), min), x, Sizes(x.size(), max));
- }
- 
--Sizes clamp(int64_t min, const Sizes &x, const Sizes &max) {
-+Sizes clamp(int64_t min, const Sizes& x, const Sizes& max) {
-   return clamp(Sizes(x.size(), min), x, max);
- }
- 
--Sizes clamp(const Sizes &min, const Sizes &x, int64_t max) {
-+Sizes clamp(const Sizes& min, const Sizes& x, int64_t max) {
-   return clamp(min, x, Sizes(x.size(), max));
- }
- 
--Sizes clamp(const Sizes &min, const Sizes &x, const Sizes &max) {
-+Sizes clamp(const Sizes& min, const Sizes& x, const Sizes& max) {
-   if (min.size() != x.size() || x.size() != max.size())
-     llvm::report_fatal_error("expected same size");
-   Sizes result(x.size());
-@@ -120,15 +120,15 @@
-   return result;
- }
- 
--const Index &IndexSpaceIterator::operator*() const {
-+const Index& IndexSpaceIterator::operator*() const {
-   if (!index_)
-     llvm::report_fatal_error("Dereferencing a past-the-end iterator.");
-   return *index_;
- }
- 
--const Index *IndexSpaceIterator::operator->() const { return &(*index_); }
-+const Index* IndexSpaceIterator::operator->() const { return &(*index_); }
- 
--IndexSpaceIterator &IndexSpaceIterator::operator++() {
-+IndexSpaceIterator& IndexSpaceIterator::operator++() {
-   if (!index_)
-     llvm::report_fatal_error("Incrementing a past-the-end iterator.");
- 
-diff --ruN a/stablehlo/stablehlo/reference/Index.h b/stablehlo/stablehlo/reference/Index.h
---- stablehlo/stablehlo/reference/Index.h
-+++ stablehlo/stablehlo/reference/Index.h
-@@ -32,8 +32,8 @@
- class Sizes : public SmallVector<int64_t> {
-  public:
-   Sizes() = default;
--  Sizes(const Sizes &other) = default;
--  Sizes &operator=(const Sizes &other) = default;
-+  Sizes(const Sizes& other) = default;
-+  Sizes& operator=(const Sizes& other) = default;
- 
-   Sizes(std::initializer_list<int64_t> list) : SmallVector(list) {}
-   Sizes(iterator begin, iterator end) : SmallVector(begin, end) {}
-@@ -49,66 +49,66 @@
- 
-   /// Checks if an element `e` at kth axis of `this` object follows
-   /// `0 <= e <= bounds[k]`.
--  bool inBounds(const Sizes &bounds) const;
-+  bool inBounds(const Sizes& bounds) const;
- 
-   /// Iterate over the index space of a Sizes object.
-   IndexSpaceIterator index_begin() const;
-   IndexSpaceIterator index_end() const;
- };
- 
--raw_ostream &operator<<(raw_ostream &os, const Sizes &x);
-+raw_ostream& operator<<(raw_ostream& os, const Sizes& x);
- 
- /// Overloaded add operator to return `Sizes` object `z` such that
- /// `z[k] = x[k] + y[k]` for all axis k.
--Sizes operator+(const Sizes &x, const Sizes &y);
-+Sizes operator+(const Sizes& x, const Sizes& y);
- 
- /// Overloaded add operator to return `Sizes` object `z` such that
- /// `z[k] = x[k] + y` for all axis k.
--Sizes operator+(const Sizes &x, int64_t y);
-+Sizes operator+(const Sizes& x, int64_t y);
- 
- /// Overloaded add operator to return `Sizes` object `z` such that
- /// `z[k] = x + y[k]` for all axis k.
--Sizes operator+(int64_t x, const Sizes &y);
-+Sizes operator+(int64_t x, const Sizes& y);
- 
- /// Overloaded add operator to return `Sizes` object `z` such that
- /// `z[k] = x[k] - y[k]` for all axis k.
--Sizes operator-(const Sizes &x, const Sizes &y);
-+Sizes operator-(const Sizes& x, const Sizes& y);
- 
- /// Overloaded add operator to return `Sizes` object `z` such that
- /// `z[k] = x[k] - y` for all axis k.
--Sizes operator-(const Sizes &x, int64_t y);
-+Sizes operator-(const Sizes& x, int64_t y);
- 
- /// Overloaded add operator to return `Sizes` object `z` such that
- /// `z[k] = x - y[k]` for all axis k.
--Sizes operator-(int64_t x, const Sizes &y);
-+Sizes operator-(int64_t x, const Sizes& y);
- 
- /// Overloaded add operator to return `Sizes` object `z` such that
- /// `z[k] = x[k] * y[k]` for all axis k.
--Sizes operator*(const Sizes &x, const Sizes &y);
-+Sizes operator*(const Sizes& x, const Sizes& y);
- 
- /// Overloaded add operator to return `Sizes` object `z` such that
- /// `z[k] = x[k] * y` for all axis k.
--Sizes operator*(const Sizes &x, int64_t y);
-+Sizes operator*(const Sizes& x, int64_t y);
- 
- /// Overloaded add operator to return `Sizes` object `z` such that
- /// `z[k] = x * y[k]` for all axis k.
--Sizes operator*(int64_t x, const Sizes &y);
-+Sizes operator*(int64_t x, const Sizes& y);
- 
- /// Clamp operator to return `Sizes` object `z` such that
- /// `z[k] = std::min(std::max(x[k], min[k]), max[k])` for all axis k.
--Sizes clamp(const Sizes &min, const Sizes &x, const Sizes &max);
-+Sizes clamp(const Sizes& min, const Sizes& x, const Sizes& max);
- 
- /// Clamp operator to return `Sizes` object `z` such that
- /// `z[k] = std::min(std::max(x[k], min), max)` for all axis k.
--Sizes clamp(int64_t min, const Sizes &x, int64_t max);
-+Sizes clamp(int64_t min, const Sizes& x, int64_t max);
- 
- /// Clamp operator to return `Sizes` object `z` such that
- /// `z[k] = std::min(std::max(x[k], min), max[k])` for all axis k.
--Sizes clamp(int64_t min, const Sizes &x, const Sizes &max);
-+Sizes clamp(int64_t min, const Sizes& x, const Sizes& max);
- 
- /// Clamp operator to return `Sizes` object `z` such that
- /// `z[k] = std::min(std::max(x[k], min[k]), max)` for all axis k.
--Sizes clamp(const Sizes &min, const Sizes &x, int64_t max);
-+Sizes clamp(const Sizes& min, const Sizes& x, int64_t max);
- 
- /// Represents index of a tensor.
- using Index = Sizes;
-@@ -132,23 +132,23 @@
-   /// At any point in time, the iterator can either reference an actual index
-   /// or the past-the-end element in the index space.
-   /// Dereferencing a past-the-end iterator will result in a fatal error.
--  const Index &operator*() const;
--  const Index *operator->() const;
-+  const Index& operator*() const;
-+  const Index* operator->() const;
- 
-   /// Compare the iterator to another iterator.
-   /// Two iterators are equal if they have the same underlying shape and
-   /// reference the same element in the index space.
--  bool operator==(const IndexSpaceIterator &it) {
-+  bool operator==(const IndexSpaceIterator& it) {
-     return shape_ == it.shape_ && index_ == it.index_;
-   }
--  bool operator!=(const IndexSpaceIterator &it) { return !(*this == it); }
-+  bool operator!=(const IndexSpaceIterator& it) { return !(*this == it); }
- 
-   /// Increment to the next index while iterating over the index space
-   /// of a tensor in lexicographical order.
-   /// Incrementing past the last index will result in a past-the-end iterator
-   /// which cannot be dereferenced. Incrementing even further will result in
-   /// a fatal error.
--  IndexSpaceIterator &operator++();
-+  IndexSpaceIterator& operator++();
-   IndexSpaceIterator operator++(int);
- 
-  private:
-diff --ruN a/stablehlo/stablehlo/reference/InterpreterOps.cpp b/stablehlo/stablehlo/reference/InterpreterOps.cpp
---- stablehlo/stablehlo/reference/InterpreterOps.cpp
-+++ stablehlo/stablehlo/reference/InterpreterOps.cpp
-@@ -22,6 +22,7 @@
- #include <mutex>
- #include <queue>
- #include <string>
-+#include <utility>
- 
- #include "llvm/Support/Casting.h"
- #include "llvm/Support/Errc.h"
-@@ -174,37 +175,52 @@
-     ArrayRef<InterpreterValue> inputs, std::queue<StringAttr>& infeed,
-     SmallVector<SmallVector<StringAttr>> programs, SymbolTable& symbolTable,
-     InterpreterFallback* fallback) {
--  llvm::DefaultThreadPool threadPool;
--  SmallVector<std::shared_future<SmallVector<InterpreterValue>>> futures;
--
-   uint32_t numReplicas = programs.size();
-   uint32_t numPartitions = programs[0].size();
-+
-+  // StableHLO collectives are blocking barriers: a process inside rendezvous
-+  // waits for every other process in its group to arrive, so all processes
-+  // must run concurrently. Size the pool from the process count rather than
-+  // the default (hardware_concurrency(), i.e. the host thread count). With a
-+  // default-sized pool, a program with more processes than the host has
-+  // hardware threads leaves the surplus tasks queued and never started while
-+  // the started tasks block forever waiting for them - a deadlock that only
-+  // surfaces as the rendezvous timeout. hardware_concurrency(N) leaves
-+  // ThreadPoolStrategy::Limit false, so the pool gets N threads uncapped by the
-+  // host.
-+  llvm::DefaultThreadPool threadPool(
-+      llvm::hardware_concurrency(numReplicas * numPartitions));
-+  llvm::ThreadPoolTaskGroup taskGroup(threadPool);
-+
-   ProcessGrid processGrid(numReplicas, numPartitions, infeed);
- 
-+  SmallVector<SmallVector<InterpreterValue>> taskOutputs(numReplicas *
-+                                                         numPartitions);
-+
-   auto inputsIt = inputs.begin();
- 
-   for (uint32_t i = 0; i < numReplicas; ++i) {
-     for (uint32_t j = 0; j < numPartitions; ++j) {
-       auto funcName = programs[i][j];
-       auto func = llvm::cast<func::FuncOp>(symbolTable.lookup(funcName));
--      auto evalWrapper = [&](Region& region, ArrayRef<InterpreterValue> args,
--                             ProcessId processId) {
--        Process process{processId, &processGrid};
--        return eval(region, args, fallback, &process,
--                    /*parent=*/nullptr);
--      };
--
-       auto numArgs = func.getBody().front().getArguments().size();
-       SmallVector<InterpreterValue> args(inputsIt, inputsIt + numArgs);
-       inputsIt += numArgs;
- 
--      futures.emplace_back(threadPool.async(
--          evalWrapper, std::ref(func.getBody()), args, ProcessId{i, j}));
-+      uint32_t flatIndex = i * numPartitions + j;
-+      taskGroup.async([&, flatIndex, func, args = std::move(args),
-+                       processId = ProcessId{i, j}]() mutable {
-+        Process process{processId, &processGrid};
-+        taskOutputs[flatIndex] = eval(func.getBody(), args, fallback, &process,
-+                                      /*parent=*/nullptr);
-+      });
-     }
-   }
- 
-+  taskGroup.wait();
-+
-   SmallVector<InterpreterValue> results;
--  for (auto& future : futures) results.append(future.get());
-+  for (auto& output : taskOutputs) results.append(output);
-   // TODO(#1725): Figure out how to test the outfeed queue.
-   return results;
- }
-diff --ruN a/stablehlo/stablehlo/reference/Ops.cpp b/stablehlo/stablehlo/reference/Ops.cpp
---- stablehlo/stablehlo/reference/Ops.cpp
-+++ stablehlo/stablehlo/reference/Ops.cpp
-@@ -568,7 +568,9 @@
-       auto lhs = scope.findTensor(op.getLhs());
-       auto rhs = scope.findTensor(op.getRhs());
-       auto comparisonDirection = op.getComparisonDirection();
--      auto result = compareOp(lhs, rhs, comparisonDirection, op.getType());
-+      auto compareType = op.getCompareType();
-+      auto result =
-+          compareOp(lhs, rhs, comparisonDirection, compareType, op.getType());
-       scope.add(op.getResult(), result);
-     } else if (auto op = dyn_cast<ComplexOp>(operation)) {
-       auto lhs = scope.findTensor(op.getLhs());
-@@ -771,6 +773,10 @@
-       auto operand = scope.findTensor(op.getOperand());
-       auto result = expm1Op(operand, op.getType());
-       scope.add(op.getResult(), result);
-+    } else if (auto op = dyn_cast<Exp2Op>(operation)) {
-+      auto operand = scope.findTensor(op.getOperand());
-+      auto result = exp2Op(operand, op.getType());
-+      scope.add(op.getResult(), result);
-     } else if (auto op = dyn_cast<FftOp>(operation)) {
-       const auto operand = scope.findTensor(op.getOperand());
-       const auto fftType = op.getFftType();
-@@ -828,6 +834,10 @@
-       auto operand = scope.findTensor(op.getOperand());
-       auto result = log1pOp(operand, op.getType());
-       scope.add(op.getResult(), result);
-+    } else if (auto op = dyn_cast<Log2Op>(operation)) {
-+      auto operand = scope.findTensor(op.getOperand());
-+      auto result = log2Op(operand, op.getType());
-+      scope.add(op.getResult(), result);
-     } else if (auto op = dyn_cast<LogOp>(operation)) {
-       auto operand = scope.findTensor(op.getOperand());
-       auto result = logOp(operand, op.getType());
-@@ -1597,27 +1607,56 @@
- 
- Tensor compareOp(const Tensor& lhs, const Tensor& rhs,
-                  ComparisonDirection comparisonDirection,
-+                 std::optional<ComparisonType> compareType,
-                  ShapedType resultType) {
-   Tensor result(resultType);
-+  bool isWeakOrder = compareType == ComparisonType::WEAKORDER;
-   for (auto it = result.index_begin(); it != result.index_end(); ++it) {
-+    auto lhsElem = lhs.get(*it);
-+    auto rhsElem = rhs.get(*it);
-+    if (isWeakOrder) {
-+      switch (comparisonDirection) {
-+        case ComparisonDirection::EQ:
-+          result.set(*it, weakOrderEq(lhsElem, rhsElem));
-+          break;
-+        case ComparisonDirection::NE:
-+          result.set(*it, !weakOrderEq(lhsElem, rhsElem));
-+          break;
-+        case ComparisonDirection::GE:
-+          result.set(*it, weakOrderLt(rhsElem, lhsElem) ||
-+                              weakOrderEq(lhsElem, rhsElem));
-+          break;
-+        case ComparisonDirection::GT:
-+          result.set(*it, weakOrderLt(rhsElem, lhsElem));
-+          break;
-+        case ComparisonDirection::LE:
-+          result.set(*it, weakOrderLt(lhsElem, rhsElem) ||
-+                              weakOrderEq(lhsElem, rhsElem));
-+          break;
-+        case ComparisonDirection::LT:
-+          result.set(*it, weakOrderLt(lhsElem, rhsElem));
-+          break;
-+      }
-+      continue;
-+    }
-     switch (comparisonDirection) {
-       case ComparisonDirection::EQ:
--        result.set(*it, lhs.get(*it) == rhs.get(*it));
-+        result.set(*it, lhsElem == rhsElem);
-         break;
-       case ComparisonDirection::NE:
--        result.set(*it, lhs.get(*it) != rhs.get(*it));
-+        result.set(*it, lhsElem != rhsElem);
-         break;
-       case ComparisonDirection::GE:
--        result.set(*it, lhs.get(*it) >= rhs.get(*it));
-+        result.set(*it, lhsElem >= rhsElem);
-         break;
-       case ComparisonDirection::GT:
--        result.set(*it, lhs.get(*it) > rhs.get(*it));
-+        result.set(*it, lhsElem > rhsElem);
-         break;
-       case ComparisonDirection::LE:
--        result.set(*it, lhs.get(*it) <= rhs.get(*it));
-+        result.set(*it, lhsElem <= rhsElem);
-         break;
-       case ComparisonDirection::LT:
--        result.set(*it, lhs.get(*it) < rhs.get(*it));
-+        result.set(*it, lhsElem < rhsElem);
-         break;
-     }
-   }
-@@ -1777,12 +1816,8 @@
-         dotGeneralOp(reversedLhsWindow, rhs, lhsContractingDimensions,
-                      rhsContractingDimensions, resultType.getElementType());
- 
--    Sizes resultNonSpatialDims;
--    for (auto i = 0; i < result.getRank(); ++i)
--      if (llvm::find(outputSpatialDimensions, i) ==
--          outputSpatialDimensions.end())
--        resultNonSpatialDims.push_back(result.getShape()[i]);
--
-+    Sizes resultNonSpatialDims{result.getShape()[outputBatchDimension],
-+                               result.getShape()[outputFeatureDimension]};
-     Axes resultPermutation;
-     resultPermutation.push_back(outputBatchDimension);
-     resultPermutation.append(outputSpatialDimensions.begin(),
-@@ -1932,6 +1967,13 @@
-   return result;
- }
- 
-+Tensor exp2Op(const Tensor& operand, ShapedType resultType) {
-+  Tensor result(resultType);
-+  for (auto it = result.index_begin(); it != result.index_end(); ++it)
-+    result.set(*it, exp2(operand.get(*it)));
-+  return result;
-+}
-+
- Tensor exponentialOp(const Tensor& operand, ShapedType resultType) {
-   Tensor result(resultType);
-   for (auto it = result.index_begin(); it != result.index_end(); ++it)
-@@ -1942,23 +1984,30 @@
- Tensor __fft_1d(const Tensor& operand, ShapedType resultType,
-                 int64_t dimension) {
-   Tensor result(resultType);
--
-+  auto complexTy = cast<ComplexType>(resultType.getElementType());
-+  const llvm::fltSemantics& sem =
-+      cast<FloatType>(complexTy.getElementType()).getFloatSemantics();
-   for (auto resultIt = result.index_begin(); resultIt != result.index_end();
-        ++resultIt) {
-     auto resultIndex = *resultIt;
-     auto N = operand.getShape()[dimension];
- 
--    Element X(resultType.getElementType(),
--              mlir::Complex<APFloat>(APFloat(0.0), APFloat(0.0)));
-+    Element X(
-+        resultType.getElementType(),
-+        mlir::Complex<APFloat>(APFloat::getZero(sem), APFloat::getZero(sem)));
- 
-     for (int n = 0; n < N; n++) {
-       auto operandIndex = Index(resultIndex);
-       operandIndex[dimension] = n;
- 
--      auto floatType =
--          cast<ComplexType>(operand.getElementType()).getElementType();
--      Element phase(floatType,
--                    APFloat(-2 * M_PI * n * resultIndex[dimension] / N));
-+      auto floatType = cast<FloatType>(
-+          cast<ComplexType>(operand.getElementType()).getElementType());
-+
-+      auto phaseVal = APFloat(-2 * M_PI * n * resultIndex[dimension] / N);
-+      bool losesInfo;
-+      phaseVal.convert(floatType.getFloatSemantics(),
-+                       APFloat::rmNearestTiesToEven, &losesInfo);
-+      Element phase(floatType, phaseVal);
-       auto rotator = complex(cosine(phase), sine(phase));
-       X = X + operand.get(operandIndex) * rotator;
-     }
-@@ -1992,10 +2041,12 @@
-   } else if (fftType == FftType::IFFT) {
-     // compute IFFT as conjugation, FFT, conjugation and normalization
-     Tensor result(operand);
-+    auto complexTy = cast<ComplexType>(resultType.getElementType());
-+    const llvm::fltSemantics& sem =
-+        cast<FloatType>(complexTy.getElementType()).getFloatSemantics();
-     for (auto d = fftDims.begin(); d != fftDims.end(); ++d) {
-       auto divisorValue = mlir::Complex<APFloat>(
--          APFloat(static_cast<double>(resultType.getDimSize(*d))),
--          APFloat(0.0));
-+          APFloat(sem, resultType.getDimSize(*d)), APFloat::getZero(sem));
-       auto divisor =
-           constantOp(mlir::DenseElementsAttr::get(resultType, divisorValue));
-       result = divideOp(conjugate(__fft_1d(conjugate(result), resultType, *d)),
-@@ -2040,6 +2091,8 @@
-     replica = conjugate(reverseOp(replica, Axes{d_irfft}, replica.getType()));
-     auto complexType =
-         ComplexType::get(cast<FloatType>(resultType.getElementType()));
-+    const llvm::fltSemantics& sem =
-+        cast<FloatType>(complexType.getElementType()).getFloatSemantics();
-     result = concatenateOp({result, replica}, Axis{d_irfft},
-                            resultType.cloneWith(std::nullopt, complexType));
- 
-@@ -2048,7 +2101,7 @@
-                     resultType);
- 
-     // normalize result of inverse transformation
--    auto divisorValue = APFloat(static_cast<double>(renormFactor));
-+    auto divisorValue = APFloat(sem, renormFactor);
-     auto divisor =
-         constantOp(mlir::DenseElementsAttr::get(resultType, divisorValue));
-     return divideOp(result, divisor, resultType);
-@@ -2189,6 +2242,13 @@
-   return result;
- }
- 
-+Tensor log2Op(const Tensor& operand, ShapedType resultType) {
-+  Tensor result(resultType);
-+  for (auto it = result.index_begin(); it != result.index_end(); ++it)
-+    result.set(*it, log2(operand.get(*it)));
-+  return result;
-+}
-+
- Tensor logOp(const Tensor& operand, ShapedType resultType) {
-   Tensor result(resultType);
-   for (auto it = result.index_begin(); it != result.index_end(); ++it)
-diff --ruN a/stablehlo/stablehlo/reference/Ops.h b/stablehlo/stablehlo/reference/Ops.h
---- stablehlo/stablehlo/reference/Ops.h
-+++ stablehlo/stablehlo/reference/Ops.h
-@@ -72,6 +72,7 @@
-                            ChannelId channelId, Process* process);
- Tensor compareOp(const Tensor& lhs, const Tensor& rhs,
-                  ComparisonDirection comparisonDirection,
-+                 std::optional<ComparisonType> compareType,
-                  ShapedType resultType);
- Tensor complexOp(const Tensor& lhs, const Tensor& rhs, ShapedType resultType);
- Tensor concatenateOp(ArrayRef<Tensor> inputs, Axis dimension,
-@@ -102,6 +103,7 @@
-                             ArrayRef<Tensor> startIndices,
-                             ShapedType resultType);
- Tensor expm1Op(const Tensor& operand, ShapedType resultType);
-+Tensor exp2Op(const Tensor& operand, ShapedType resultType);
- Tensor exponentialOp(const Tensor& operand, ShapedType resultType);
- Tensor fftOp(const Tensor& operand, const FftType fftType,
-              const ArrayRef<int64_t>& fftLength, ShapedType resultType);
-@@ -124,6 +126,7 @@
- Tensor iotaOp(Axis iotaDimension, ShapedType resultType);
- Tensor isFiniteOp(const Tensor& operand, ShapedType resultType);
- Tensor log1pOp(const Tensor& operand, ShapedType resultType);
-+Tensor log2Op(const Tensor& operand, ShapedType resultType);
- Tensor logOp(const Tensor& operand, ShapedType resultType);
- Tensor logisticOp(const Tensor& operand, ShapedType resultType);
- Tensor mapOp(ArrayRef<Tensor> inputs, Region& computation, Process* process,
-diff --ruN a/stablehlo/stablehlo/reference/Process.cpp b/stablehlo/stablehlo/reference/Process.cpp
---- stablehlo/stablehlo/reference/Process.cpp
-+++ stablehlo/stablehlo/reference/Process.cpp
-@@ -25,7 +25,7 @@
- namespace mlir {
- namespace stablehlo {
- 
--Process::Process(ProcessId id, ProcessGrid *grid) : id_(id), grid_(grid) {}
-+Process::Process(ProcessId id, ProcessGrid* grid) : id_(id), grid_(grid) {}
- 
- ProcessGroups Process::crossPartition(
-     SmallVector<SmallVector<uint32_t>> partitionGroups) {
-diff --ruN a/stablehlo/stablehlo/reference/Process.h b/stablehlo/stablehlo/reference/Process.h
---- stablehlo/stablehlo/reference/Process.h
-+++ stablehlo/stablehlo/reference/Process.h
-@@ -29,7 +29,7 @@
-  public:
-   /// \name Constructors
-   /// @{
--  Process(ProcessId id, ProcessGrid *grid);
-+  Process(ProcessId id, ProcessGrid* grid);
-   /// @}
- 
-   /// See `ProcessGrid::crossPartition`.
-@@ -72,7 +72,7 @@
- 
-   /// See ProcessGrid. The pointer is used to gain access to allow
-   /// synchronization among participating processes.
--  ProcessGrid *grid_;
-+  ProcessGrid* grid_;
- };
- 
- }  // namespace stablehlo
-diff --ruN a/stablehlo/stablehlo/reference/ProcessGrid.cpp b/stablehlo/stablehlo/reference/ProcessGrid.cpp
---- stablehlo/stablehlo/reference/ProcessGrid.cpp
-+++ stablehlo/stablehlo/reference/ProcessGrid.cpp
-@@ -39,16 +39,16 @@
- // ProcessId.
- //===----------------------------------------------------------------------===//
- 
--bool ProcessId::operator!=(const ProcessId &other) const {
-+bool ProcessId::operator!=(const ProcessId& other) const {
-   return !(*this == other);
- }
- 
--bool ProcessId::operator<(const ProcessId &other) const {
-+bool ProcessId::operator<(const ProcessId& other) const {
-   return std::pair<uint32_t, uint32_t>{replicaId, partitionId} <
-          std::pair<uint32_t, uint32_t>{other.replicaId, other.partitionId};
- }
- 
--bool ProcessId::operator==(const ProcessId &other) const {
-+bool ProcessId::operator==(const ProcessId& other) const {
-   return std::pair<uint32_t, uint32_t>{replicaId, partitionId} ==
-          std::pair<uint32_t, uint32_t>{other.replicaId, other.partitionId};
- }
-@@ -70,7 +70,7 @@
- //===----------------------------------------------------------------------===//
- 
- RendezvousResult::RendezvousResult(
--    std::map<ProcessId, SmallVector<Tensor>> const &results)
-+    std::map<ProcessId, SmallVector<Tensor>> const& results)
-     : results_(results) {}
- 
- void RendezvousResult::insert(ProcessId processId,
-@@ -86,12 +86,12 @@
- 
- SmallVector<SmallVector<Tensor>> RendezvousResult::getSortedTensors() const {
-   return llvm::map_to_vector(results_,
--                             [](const auto &pair) { return pair.second; });
-+                             [](const auto& pair) { return pair.second; });
- }
- 
- bool RendezvousResult::hasMatchingOperandsCount() const {
-   auto count = results_.begin()->second.size();
--  for (const auto &it : results_)
-+  for (const auto& it : results_)
-     if (count != it.second.size()) return false;
-   return true;
- }
-@@ -101,7 +101,7 @@
- //===----------------------------------------------------------------------===//
- 
- template <typename K, typename V>
--V &detail::ThreadSafeMap<K, V>::operator[](const K &key) {
-+V& detail::ThreadSafeMap<K, V>::operator[](const K& key) {
-   std::lock_guard<std::mutex> lock(lock_);
-   return map_[key];
- }
-@@ -133,7 +133,7 @@
- //===----------------------------------------------------------------------===//
- 
- template <typename T>
--detail::ThreadSafeQueue<T>::ThreadSafeQueue(const std::queue<T> &queue)
-+detail::ThreadSafeQueue<T>::ThreadSafeQueue(const std::queue<T>& queue)
-     : queue_(queue) {}
- 
- template <typename T>
-@@ -155,7 +155,7 @@
- //===----------------------------------------------------------------------===//
- 
- ProcessGrid::ProcessGrid(uint32_t numReplicas, uint32_t numPartitions,
--                         std::queue<StringAttr> &infeed)
-+                         std::queue<StringAttr>& infeed)
-     : numReplicas_(numReplicas),
-       numPartitions_(numPartitions),
-       infeed_(infeed) {}
-@@ -163,7 +163,7 @@
- ProcessGroups ProcessGrid::crossPartition(
-     SmallVector<SmallVector<uint32_t>> partitionGroups) {
-   ProcessGroups processGroups;
--  for (const auto &partitionGroup : partitionGroups) {
-+  for (const auto& partitionGroup : partitionGroups) {
-     for (uint32_t replicaId = 0; replicaId < numReplicas_; ++replicaId) {
-       ProcessGroup processGroup;
-       for (uint32_t partitionId : partitionGroup)
-@@ -177,7 +177,7 @@
- ProcessGroups ProcessGrid::crossReplica(
-     SmallVector<SmallVector<uint32_t>> replicaGroups) {
-   ProcessGroups processGroups;
--  for (const auto &replicaGroup : replicaGroups) {
-+  for (const auto& replicaGroup : replicaGroups) {
-     for (uint32_t partitionId = 0; partitionId < numPartitions_;
-          ++partitionId) {
-       ProcessGroup processGroup;
-@@ -192,7 +192,7 @@
- ProcessGroups ProcessGrid::crossReplicaAndPartition(
-     SmallVector<SmallVector<uint32_t>> replicaGroups) {
-   ProcessGroups processGroups;
--  for (const auto &replicaGroup : replicaGroups) {
-+  for (const auto& replicaGroup : replicaGroups) {
-     ProcessGroup processGroup;
-     for (uint32_t partitionId = 0; partitionId < numPartitions_; ++partitionId)
-       for (uint32_t replicaId : replicaGroup)
-@@ -205,7 +205,7 @@
- ProcessGroups ProcessGrid::flattenedIds(
-     SmallVector<SmallVector<uint32_t>> flattenedIdGroups) {
-   ProcessGroups processGroups;
--  for (const auto &flattenedIdGroup : flattenedIdGroups) {
-+  for (const auto& flattenedIdGroup : flattenedIdGroups) {
-     ProcessGroup processGroup;
-     for (auto flattenedId : flattenedIdGroup) {
-       uint32_t replicaId = flattenedId / numPartitions_;
-@@ -251,7 +251,7 @@
-   }
- 
-   std::pair<ProcessGroup, ChannelId> channelKey(processGroup, channelId);
--  auto &state = channels_[channelKey];
-+  auto& state = channels_[channelKey];
- 
-   std::unique_lock<std::mutex> lock(state.mutex);
-   state.values[processId] = SmallVector<Tensor>(operands);
-diff --ruN a/stablehlo/stablehlo/reference/ProcessGrid.h b/stablehlo/stablehlo/reference/ProcessGrid.h
---- stablehlo/stablehlo/reference/ProcessGrid.h
-+++ stablehlo/stablehlo/reference/ProcessGrid.h
-@@ -41,7 +41,7 @@
- class RendezvousResult {
-  public:
-   RendezvousResult() = default;
--  RendezvousResult(std::map<ProcessId, SmallVector<Tensor>> const &results);
-+  RendezvousResult(std::map<ProcessId, SmallVector<Tensor>> const& results);
- 
-   /// Iterates through the (ProcessId, SmallVector<Tensor>) map entires and
-   /// returns a vector of Tensors sorted by ProcessId--(replicaId, partitionId)
-@@ -104,7 +104,7 @@
- class ThreadSafeMap {
-  public:
-   /// Returns a reference to the data associated with the `key`.
--  V &operator[](const K &key);
-+  V& operator[](const K& key);
- 
-  private:
-   /// Synchronization primitive used to manage concurrent access to the map.
-@@ -143,7 +143,7 @@
-   /// \name Constructors
-   /// @{
-   ThreadSafeQueue() = default;
--  ThreadSafeQueue(const std::queue<T> &queue);
-+  ThreadSafeQueue(const std::queue<T>& queue);
-   /// @}
- 
-   /// Remove the first element of the queue and return it.
-@@ -173,16 +173,16 @@
-   uint32_t partitionId;
- 
-   /// Overloaded inequality operator.
--  bool operator!=(const ProcessId &other) const;
-+  bool operator!=(const ProcessId& other) const;
- 
-   /// The sort order for ProcessId is not defined in StableHLO, and it's
-   /// internally used in ProcessGrid::rendezvous as part of a sorted key on the
-   /// map. This operator is conveniently used to help define the ordering since
-   /// ordering is defined for StableHLO process group.
--  bool operator<(const ProcessId &other) const;
-+  bool operator<(const ProcessId& other) const;
- 
-   /// Overloaded equality operator.
--  bool operator==(const ProcessId &other) const;
-+  bool operator==(const ProcessId& other) const;
- };
- 
- /// StableHLO `process_group`.
-@@ -190,9 +190,9 @@
-  public:
-   using SmallVector<ProcessId>::SmallVector;
- 
--  bool operator<(const ProcessGroup &other) const {
--    return static_cast<const SmallVector &>(*this) <
--           static_cast<const SmallVector &>(other);
-+  bool operator<(const ProcessGroup& other) const {
-+    return static_cast<const SmallVector&>(*this) <
-+           static_cast<const SmallVector&>(other);
-   }
- };
- 
-@@ -211,7 +211,7 @@
-   /// \name Constructors
-   /// @{
-   ProcessGrid(uint32_t numReplicas, uint32_t numPartitions,
--              std::queue<StringAttr> &infeed);
-+              std::queue<StringAttr>& infeed);
-   /// @}
- 
-   /// StableHLO `cross_partition` communication strategy.
-diff --ruN a/stablehlo/stablehlo/reference/Token.cpp b/stablehlo/stablehlo/reference/Token.cpp
---- stablehlo/stablehlo/reference/Token.cpp
-+++ stablehlo/stablehlo/reference/Token.cpp
-@@ -23,11 +23,11 @@
- namespace mlir {
- namespace stablehlo {
- 
--Token::Token(MLIRContext *context) : type_(TokenType::get(context)) {}
-+Token::Token(MLIRContext* context) : type_(TokenType::get(context)) {}
- 
- TokenType Token::getType() const { return type_; }
- 
--void Token::print(raw_ostream &os) const { getType().print(os); }
-+void Token::print(raw_ostream& os) const { getType().print(os); }
- 
- void Token::dump() const { print(llvm::errs()); }
- 
-diff --ruN a/stablehlo/stablehlo/reference/Token.h b/stablehlo/stablehlo/reference/Token.h
---- stablehlo/stablehlo/reference/Token.h
-+++ stablehlo/stablehlo/reference/Token.h
-@@ -28,14 +28,14 @@
-  public:
-   /// \name Constructors
-   /// @{
--  Token(MLIRContext *context);
-+  Token(MLIRContext* context);
-   /// @}
- 
-   /// Returns the type of the Token object.
-   TokenType getType() const;
- 
-   /// Prints Token object.
--  void print(raw_ostream &os) const;
-+  void print(raw_ostream& os) const;
-   void dump() const;
- 
-  private:
-@@ -43,7 +43,7 @@
- };
- 
- /// Print utilities for Token objects.
--inline raw_ostream &operator<<(raw_ostream &os, Token Token) {
-+inline raw_ostream& operator<<(raw_ostream& os, Token Token) {
-   Token.print(os);
-   return os;
- }
-diff --ruN a/stablehlo/stablehlo/reference/Value.cpp b/stablehlo/stablehlo/reference/Value.cpp
---- stablehlo/stablehlo/reference/Value.cpp
-+++ stablehlo/stablehlo/reference/Value.cpp
-@@ -46,7 +46,7 @@
- 
- TupleType Tuple::getType() const { return type_; }
- 
--void Tuple::print(raw_ostream &os) const {
-+void Tuple::print(raw_ostream& os) const {
-   getType().print(os);
-   os << " (\n";
-   for (size_t i = 0; i < values_.size(); ++i) {
-@@ -63,9 +63,9 @@
- // InterpreterValue.
- //===----------------------------------------------------------------------===//
- 
--InterpreterValue::InterpreterValue(const Tensor &tensor) : value_(tensor) {}
--InterpreterValue::InterpreterValue(const Token &token) : value_(token) {}
--InterpreterValue::InterpreterValue(const Tuple &tuple) : value_(tuple) {}
-+InterpreterValue::InterpreterValue(const Tensor& tensor) : value_(tensor) {}
-+InterpreterValue::InterpreterValue(const Token& token) : value_(token) {}
-+InterpreterValue::InterpreterValue(const Tuple& tuple) : value_(tuple) {}
- 
- Tensor InterpreterValue::getTensor() const {
-   if (!isTensor())
-@@ -108,7 +108,7 @@
-   return std::holds_alternative<Tuple>(value_);
- }
- 
--void InterpreterValue::print(raw_ostream &os) const {
-+void InterpreterValue::print(raw_ostream& os) const {
-   if (isTensor())
-     getTensor().print(os);
-   else if (isToken())
-diff --ruN a/stablehlo/stablehlo/tests/CheckOps.cpp b/stablehlo/stablehlo/tests/CheckOps.cpp
---- stablehlo/stablehlo/tests/CheckOps.cpp
-+++ stablehlo/stablehlo/tests/CheckOps.cpp
-@@ -72,7 +72,7 @@
- // Check Dialect Constructor
- //===----------------------------------------------------------------------===//
- 
--CheckDialect::CheckDialect(MLIRContext *context)
-+CheckDialect::CheckDialect(MLIRContext* context)
-     : Dialect(getDialectNamespace(), context, TypeID::get<CheckDialect>()) {
-   addOperations<
- #define GET_OP_LIST
-@@ -80,13 +80,13 @@
-       >();
- }
- 
--llvm::Error evalExpectAlmostEqConstOp(const Tensor &lhs, ElementsAttr value,
-+llvm::Error evalExpectAlmostEqConstOp(const Tensor& lhs, ElementsAttr value,
-                                       APFloat tolerance) {
-   auto rhs = makeTensor(cast<DenseElementsAttr>(value));
-   return evalExpectAlmostEqOp(lhs, rhs, tolerance);
- }
- 
--llvm::Error evalExpectAlmostEqOp(const Tensor &lhs, const Tensor &rhs,
-+llvm::Error evalExpectAlmostEqOp(const Tensor& lhs, const Tensor& rhs,
-                                  APFloat tolerance) {
-   for (auto lhsIt = lhs.index_begin(), rhsIt = rhs.index_begin();
-        lhsIt != lhs.index_end(); ++lhsIt, ++rhsIt)
-@@ -101,12 +101,12 @@
-   return llvm::Error::success();
- }
- 
--llvm::Error evalExpectEqConstOp(const Tensor &lhs, ElementsAttr value) {
-+llvm::Error evalExpectEqConstOp(const Tensor& lhs, ElementsAttr value) {
-   auto rhs = makeTensor(cast<DenseElementsAttr>(value));
-   return evalExpectEqOp(lhs, rhs);
- }
- 
--llvm::Error evalExpectEqOp(const Tensor &lhs, const Tensor &rhs) {
-+llvm::Error evalExpectEqOp(const Tensor& lhs, const Tensor& rhs) {
-   for (auto lhsIt = lhs.index_begin(), rhsIt = rhs.index_begin();
-        lhsIt != lhs.index_end(); ++lhsIt, ++rhsIt)
-     if ((lhs.get(*lhsIt) != rhs.get(*rhsIt)).getBooleanValue())
-@@ -146,7 +146,7 @@
-   return llvm::errc::bad_address;
- }
- 
--llvm::Error evalExpectSerializedEqOp(const Tensor &expected, StringRef probeId,
-+llvm::Error evalExpectSerializedEqOp(const Tensor& expected, StringRef probeId,
-                                      StringRef probeDir, uint32_t iteration) {
-   auto serializedMetadata =
-       getSerializedTensorMetadata(probeId, probeDir, iteration);
-@@ -206,7 +206,7 @@
-   return std::numeric_limits<uint64_t>::max();
- }
- 
--static uint64_t ULPDifference(const Element &e1, const Element &e2) {
-+static uint64_t ULPDifference(const Element& e1, const Element& e2) {
-   // caller is responsible for ensuring that e1, e2 have both the same
-   // float or complex types
-   if (isSupportedComplexType(e1.getType())) {
-@@ -218,7 +218,7 @@
-   return ULPDifference(e1.getFloatValue(), e2.getFloatValue());
- }
- 
--llvm::Error evalExpectCloseOp(const Tensor &actual, const Tensor &expected,
-+llvm::Error evalExpectCloseOp(const Tensor& actual, const Tensor& expected,
-                               uint64_t min_ulp_difference,
-                               uint64_t max_ulp_difference) {
-   auto type = actual.getElementType();
-diff --ruN a/stablehlo/stablehlo/tests/CheckOps.h b/stablehlo/stablehlo/tests/CheckOps.h
---- stablehlo/stablehlo/tests/CheckOps.h
-+++ stablehlo/stablehlo/tests/CheckOps.h
-@@ -32,20 +32,20 @@
- 
- class CheckDialect : public Dialect {
-  public:
--  explicit CheckDialect(MLIRContext *context);
-+  explicit CheckDialect(MLIRContext* context);
-   static StringRef getDialectNamespace() { return "check"; }
- };
- 
- // The eval functions for the following ops are used only for test harness.
--llvm::Error evalExpectAlmostEqConstOp(const Tensor &lhs, ElementsAttr value,
-+llvm::Error evalExpectAlmostEqConstOp(const Tensor& lhs, ElementsAttr value,
-                                       APFloat tolerance);
--llvm::Error evalExpectAlmostEqOp(const Tensor &lhs, const Tensor &rhs,
-+llvm::Error evalExpectAlmostEqOp(const Tensor& lhs, const Tensor& rhs,
-                                  APFloat tolerance);
--llvm::Error evalExpectEqConstOp(const Tensor &lhs, ElementsAttr value);
--llvm::Error evalExpectEqOp(const Tensor &lhs, const Tensor &rhs);
--llvm::Error evalExpectSerializedEqOp(const Tensor &expected, StringRef probeId,
-+llvm::Error evalExpectEqConstOp(const Tensor& lhs, ElementsAttr value);
-+llvm::Error evalExpectEqOp(const Tensor& lhs, const Tensor& rhs);
-+llvm::Error evalExpectSerializedEqOp(const Tensor& expected, StringRef probeId,
-                                      StringRef probeDir, uint32_t iteration);
--llvm::Error evalExpectCloseOp(const Tensor &actual, const Tensor &expected,
-+llvm::Error evalExpectCloseOp(const Tensor& actual, const Tensor& expected,
-                               uint64_t min_ulp_difference,
-                               uint64_t max_ulp_difference);
- 
-diff --ruN a/stablehlo/stablehlo/tests/CheckOps.td b/stablehlo/stablehlo/tests/CheckOps.td
---- stablehlo/stablehlo/tests/CheckOps.td
-+++ stablehlo/stablehlo/tests/CheckOps.td
-@@ -24,6 +24,7 @@
- 
- def CHECK_Dialect : Dialect {
-   let name = "check";
-+  let useStrictPropertiesInAssemblyFormat = 0;
-   let cppNamespace = "::mlir::stablehlo::check";
- 
-   let description = [{
-diff --ruN a/stablehlo/stablehlo/tests/TestUtils.cpp b/stablehlo/stablehlo/tests/TestUtils.cpp
---- stablehlo/stablehlo/tests/TestUtils.cpp
-+++ stablehlo/stablehlo/tests/TestUtils.cpp
-@@ -60,7 +60,9 @@
- 
-     // Replace with custom call to avoid pattern reapplication
-     auto customCall = stablehlo::CustomCallOp::create(
--        rewriter, op->getLoc(), op->getResultTypes(), *broadcastedOperands);
-+        rewriter, op->getLoc(), op->getResultTypes(), *broadcastedOperands,
-+        typename stablehlo::CustomCallOp::Properties{},
-+        /*discardableAttributes=*/{});
-     customCall.setCallTargetName("numpy_broadcasted");
-     customCall.setHasSideEffect(true);
-     rewriter.replaceOp(op, customCall);
-diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir
---- stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir
-+++ stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir
-@@ -8,10 +8,8 @@
- // CHECK:           %[[ADD_0:.*]] = stablehlo.add %[[CONSTANT_0]], %[[ARG0]] : tensor<bf16>
- // CHECK:           %[[MULTIPLY_0:.*]] = stablehlo.multiply %[[SUBTRACT_0]], %[[ADD_0]] : tensor<bf16>
- // CHECK:           %[[SQRT_0:.*]] = stablehlo.sqrt %[[MULTIPLY_0]] : tensor<bf16>
--// CHECK:           %[[ADD_1:.*]] = stablehlo.add %[[CONSTANT_0]], %[[SQRT_0]] : tensor<bf16>
--// CHECK:           %[[VAL_0:.*]] = stablehlo.atan2 %[[ARG0]], %[[ADD_1]] : tensor<bf16>
--// CHECK:           %[[ADD_2:.*]] = stablehlo.add %[[VAL_0]], %[[VAL_0]] : tensor<bf16>
--// CHECK:           return %[[ADD_2]] : tensor<bf16>
-+// CHECK:           %[[VAL_0:.*]] = stablehlo.atan2 %[[ARG0]], %[[SQRT_0]] : tensor<bf16>
-+// CHECK:           return %[[VAL_0]] : tensor<bf16>
- // CHECK:         }
- func.func @asin_bf16(%arg : tensor<bf16>) -> tensor<bf16> {
-   %result = "chlo.asin"(%arg) : (tensor<bf16>) -> tensor<bf16>
-@@ -27,10 +25,8 @@
- // CHECK:           %[[ADD_0:.*]] = stablehlo.add %[[CONSTANT_0]], %[[ARG0]] : tensor<f16>
- // CHECK:           %[[MULTIPLY_0:.*]] = stablehlo.multiply %[[SUBTRACT_0]], %[[ADD_0]] : tensor<f16>
- // CHECK:           %[[SQRT_0:.*]] = stablehlo.sqrt %[[MULTIPLY_0]] : tensor<f16>
--// CHECK:           %[[ADD_1:.*]] = stablehlo.add %[[CONSTANT_0]], %[[SQRT_0]] : tensor<f16>
--// CHECK:           %[[VAL_0:.*]] = stablehlo.atan2 %[[ARG0]], %[[ADD_1]] : tensor<f16>
--// CHECK:           %[[ADD_2:.*]] = stablehlo.add %[[VAL_0]], %[[VAL_0]] : tensor<f16>
--// CHECK:           return %[[ADD_2]] : tensor<f16>
-+// CHECK:           %[[VAL_0:.*]] = stablehlo.atan2 %[[ARG0]], %[[SQRT_0]] : tensor<f16>
-+// CHECK:           return %[[VAL_0]] : tensor<f16>
- // CHECK:         }
- func.func @asin_f16(%arg : tensor<f16>) -> tensor<f16> {
-   %result = "chlo.asin"(%arg) : (tensor<f16>) -> tensor<f16>
-@@ -46,10 +42,8 @@
- // CHECK:           %[[ADD_0:.*]] = stablehlo.add %[[CONSTANT_0]], %[[ARG0]] : tensor<f32>
- // CHECK:           %[[MULTIPLY_0:.*]] = stablehlo.multiply %[[SUBTRACT_0]], %[[ADD_0]] : tensor<f32>
- // CHECK:           %[[SQRT_0:.*]] = stablehlo.sqrt %[[MULTIPLY_0]] : tensor<f32>
--// CHECK:           %[[ADD_1:.*]] = stablehlo.add %[[CONSTANT_0]], %[[SQRT_0]] : tensor<f32>
--// CHECK:           %[[VAL_0:.*]] = stablehlo.atan2 %[[ARG0]], %[[ADD_1]] : tensor<f32>
--// CHECK:           %[[ADD_2:.*]] = stablehlo.add %[[VAL_0]], %[[VAL_0]] : tensor<f32>
--// CHECK:           return %[[ADD_2]] : tensor<f32>
-+// CHECK:           %[[VAL_0:.*]] = stablehlo.atan2 %[[ARG0]], %[[SQRT_0]] : tensor<f32>
-+// CHECK:           return %[[VAL_0]] : tensor<f32>
- // CHECK:         }
- func.func @asin_f32(%arg : tensor<f32>) -> tensor<f32> {
-   %result = "chlo.asin"(%arg) : (tensor<f32>) -> tensor<f32>
-@@ -65,10 +59,8 @@
- // CHECK:           %[[ADD_0:.*]] = stablehlo.add %[[CONSTANT_0]], %[[ARG0]] : tensor<f64>
- // CHECK:           %[[MULTIPLY_0:.*]] = stablehlo.multiply %[[SUBTRACT_0]], %[[ADD_0]] : tensor<f64>
- // CHECK:           %[[SQRT_0:.*]] = stablehlo.sqrt %[[MULTIPLY_0]] : tensor<f64>
--// CHECK:           %[[ADD_1:.*]] = stablehlo.add %[[CONSTANT_0]], %[[SQRT_0]] : tensor<f64>
--// CHECK:           %[[VAL_0:.*]] = stablehlo.atan2 %[[ARG0]], %[[ADD_1]] : tensor<f64>
--// CHECK:           %[[ADD_2:.*]] = stablehlo.add %[[VAL_0]], %[[VAL_0]] : tensor<f64>
--// CHECK:           return %[[ADD_2]] : tensor<f64>
-+// CHECK:           %[[VAL_0:.*]] = stablehlo.atan2 %[[ARG0]], %[[SQRT_0]] : tensor<f64>
-+// CHECK:           return %[[VAL_0]] : tensor<f64>
- // CHECK:         }
- func.func @asin_f64(%arg : tensor<f64>) -> tensor<f64> {
-   %result = "chlo.asin"(%arg) : (tensor<f64>) -> tensor<f64>
-@@ -2432,9 +2424,10 @@
- // CHECK:           %[[ADD_18:.*]] = stablehlo.add %[[ADD_16]], %[[CONSTANT_1]] : tensor<f32>
- // CHECK:           %[[POWER_10:.*]] = stablehlo.power %[[ADD_18]], %[[NEGATE_0]] : tensor<f32>
- // CHECK:           %[[CONSTANT_2:.*]] = stablehlo.constant dense<1.000000e+00> : tensor<f32>
--// CHECK:           %[[MULTIPLY_0:.*]] = stablehlo.multiply %[[POWER_10]], %[[ADD_18]] : tensor<f32>
- // CHECK:           %[[SUBTRACT_0:.*]] = stablehlo.subtract %[[CONVERT_0]], %[[CONSTANT_2]] : tensor<f32>
--// CHECK:           %[[DIVIDE_0:.*]] = stablehlo.divide %[[MULTIPLY_0]], %[[SUBTRACT_0]] : tensor<f32>
-+// CHECK:           %[[NEGATE_1:.*]] = stablehlo.negate %[[SUBTRACT_0]] : tensor<f32>
-+// CHECK:           %[[POWER_11:.*]] = stablehlo.power %[[ADD_18]], %[[NEGATE_1]] : tensor<f32>
-+// CHECK:           %[[DIVIDE_0:.*]] = stablehlo.divide %[[POWER_11]], %[[SUBTRACT_0]] : tensor<f32>
- // CHECK:           %[[MULTIPLY_1:.*]] = stablehlo.multiply %[[ADD_18]], %[[ADD_18]] : tensor<f32>
- // CHECK:           %[[DIVIDE_1:.*]] = stablehlo.divide %[[CONSTANT_1]], %[[MULTIPLY_1]] : tensor<f32>
- // CHECK:           %[[CONSTANT_3:.*]] = stablehlo.constant dense<2.200000e+01> : tensor<f32>
-@@ -2710,9 +2703,10 @@
- // CHECK:           %[[ADD_40:.*]] = stablehlo.add %[[ADD_38]], %[[CONSTANT_29]] : tensor<f32>
- // CHECK:           %[[POWER_10:.*]] = stablehlo.power %[[ADD_40]], %[[NEGATE_2]] : tensor<f32>
- // CHECK:           %[[CONSTANT_30:.*]] = stablehlo.constant dense<1.000000e+00> : tensor<f32>
--// CHECK:           %[[MULTIPLY_3:.*]] = stablehlo.multiply %[[POWER_10]], %[[ADD_40]] : tensor<f32>
- // CHECK:           %[[SUBTRACT_7:.*]] = stablehlo.subtract %[[ADD_0]], %[[CONSTANT_30]] : tensor<f32>
--// CHECK:           %[[DIVIDE_10:.*]] = stablehlo.divide %[[MULTIPLY_3]], %[[SUBTRACT_7]] : tensor<f32>
-+// CHECK:           %[[NEGATE_4:.*]] = stablehlo.negate %[[SUBTRACT_7]] : tensor<f32>
-+// CHECK:           %[[POWER_11:.*]] = stablehlo.power %[[ADD_40]], %[[NEGATE_4]] : tensor<f32>
-+// CHECK:           %[[DIVIDE_10:.*]] = stablehlo.divide %[[POWER_11]], %[[SUBTRACT_7]] : tensor<f32>
- // CHECK:           %[[MULTIPLY_4:.*]] = stablehlo.multiply %[[ADD_40]], %[[ADD_40]] : tensor<f32>
- // CHECK:           %[[DIVIDE_11:.*]] = stablehlo.divide %[[CONSTANT_29]], %[[MULTIPLY_4]] : tensor<f32>
- // CHECK:           %[[CONSTANT_31:.*]] = stablehlo.constant dense<2.200000e+01> : tensor<f32>
-@@ -3103,9 +3097,10 @@
- // CHECK:           %[[ADD_40:.*]] = stablehlo.add %[[ADD_38]], %[[CONSTANT_29]] : tensor<f64>
- // CHECK:           %[[POWER_10:.*]] = stablehlo.power %[[ADD_40]], %[[NEGATE_2]] : tensor<f64>
- // CHECK:           %[[CONSTANT_30:.*]] = stablehlo.constant dense<1.000000e+00> : tensor<f64>
--// CHECK:           %[[MULTIPLY_3:.*]] = stablehlo.multiply %[[POWER_10]], %[[ADD_40]] : tensor<f64>
- // CHECK:           %[[SUBTRACT_7:.*]] = stablehlo.subtract %[[ADD_0]], %[[CONSTANT_30]] : tensor<f64>
--// CHECK:           %[[DIVIDE_10:.*]] = stablehlo.divide %[[MULTIPLY_3]], %[[SUBTRACT_7]] : tensor<f64>
-+// CHECK:           %[[NEGATE_4:.*]] = stablehlo.negate %[[SUBTRACT_7]] : tensor<f64>
-+// CHECK:           %[[POWER_11:.*]] = stablehlo.power %[[ADD_40]], %[[NEGATE_4]] : tensor<f64>
-+// CHECK:           %[[DIVIDE_10:.*]] = stablehlo.divide %[[POWER_11]], %[[SUBTRACT_7]] : tensor<f64>
- // CHECK:           %[[MULTIPLY_4:.*]] = stablehlo.multiply %[[ADD_40]], %[[ADD_40]] : tensor<f64>
- // CHECK:           %[[DIVIDE_11:.*]] = stablehlo.divide %[[CONSTANT_29]], %[[MULTIPLY_4]] : tensor<f64>
- // CHECK:           %[[CONSTANT_31:.*]] = stablehlo.constant dense<2.200000e+01> : tensor<f64>
-@@ -3498,9 +3493,10 @@
- // CHECK:           %[[ADD_40:.*]] = stablehlo.add %[[ADD_38]], %[[CONSTANT_29]] : tensor<f32>
- // CHECK:           %[[POWER_10:.*]] = stablehlo.power %[[ADD_40]], %[[NEGATE_2]] : tensor<f32>
- // CHECK:           %[[CONSTANT_30:.*]] = stablehlo.constant dense<1.000000e+00> : tensor<f32>
--// CHECK:           %[[MULTIPLY_3:.*]] = stablehlo.multiply %[[POWER_10]], %[[ADD_40]] : tensor<f32>
- // CHECK:           %[[SUBTRACT_7:.*]] = stablehlo.subtract %[[ADD_0]], %[[CONSTANT_30]] : tensor<f32>
--// CHECK:           %[[DIVIDE_10:.*]] = stablehlo.divide %[[MULTIPLY_3]], %[[SUBTRACT_7]] : tensor<f32>
-+// CHECK:           %[[NEGATE_4:.*]] = stablehlo.negate %[[SUBTRACT_7]] : tensor<f32>
-+// CHECK:           %[[POWER_11:.*]] = stablehlo.power %[[ADD_40]], %[[NEGATE_4]] : tensor<f32>
-+// CHECK:           %[[DIVIDE_10:.*]] = stablehlo.divide %[[POWER_11]], %[[SUBTRACT_7]] : tensor<f32>
- // CHECK:           %[[MULTIPLY_4:.*]] = stablehlo.multiply %[[ADD_40]], %[[ADD_40]] : tensor<f32>
- // CHECK:           %[[DIVIDE_11:.*]] = stablehlo.divide %[[CONSTANT_29]], %[[MULTIPLY_4]] : tensor<f32>
- // CHECK:           %[[CONSTANT_31:.*]] = stablehlo.constant dense<2.200000e+01> : tensor<f32>
-@@ -5424,6 +5420,55 @@
- 
- // -----
- 
-+// CHECK-LABEL:   func.func @scan_bounded(
-+// CHECK-SAME:      %[[ARG0:.*]]: tensor<?x?xi32, #stablehlo.bounds<2, 3>>,
-+// CHECK-SAME:      %[[ARG1:.*]]: tensor<?xi32, #stablehlo.bounds<3>>) -> (tensor<?x?xi32, #stablehlo.bounds<2, 3>>, tensor<?xi32, #stablehlo.bounds<3>>) {
-+// CHECK:           %[[GET_DIM_SIZE_0:.*]] = stablehlo.get_dimension_size %[[ARG0]], dim = 0 : (tensor<?x?xi32, #stablehlo.bounds<2, 3>>) -> tensor<i32>
-+// CHECK:           %[[CONVERT_0:.*]] = stablehlo.convert %[[GET_DIM_SIZE_0]] : (tensor<i32>) -> tensor<i64>
-+// CHECK:           %[[C0:.*]] = stablehlo.constant dense<0> : tensor<i64>
-+// CHECK:           %[[C0_I32:.*]] = stablehlo.constant dense<0> : tensor<i32>
-+// CHECK:           %[[SHAPE:.*]] = shape.shape_of %[[ARG0]] : tensor<?x?xi32, #stablehlo.bounds<2, 3>> -> tensor<2xindex>
-+// CHECK:           %[[INIT_OUT:.*]] = stablehlo.dynamic_broadcast_in_dim %[[C0_I32]], %[[SHAPE]], dims = [] : (tensor<i32>, tensor<2xindex>) -> tensor<?x?xi32, #stablehlo.bounds<2, 3>>
-+// CHECK:           %[[WHILE:.*]]:3 = stablehlo.while(%[[ITER_I:.*]] = %[[C0]], %[[ITER_CARRY:.*]] = %[[ARG1]], %[[ITER_OUT:.*]] = %[[INIT_OUT]]) : tensor<i64>, tensor<?xi32, #stablehlo.bounds<3>>, tensor<?x?xi32, #stablehlo.bounds<2, 3>>
-+// CHECK:           cond {
-+// CHECK:             %[[CMP:.*]] = stablehlo.compare LT, %[[ITER_I]], %[[CONVERT_0]] : (tensor<i64>, tensor<i64>) -> tensor<i1>
-+// CHECK:             stablehlo.return %[[CMP]] : tensor<i1>
-+// CHECK:           } do {
-+// CHECK:             %[[C0_1:.*]] = stablehlo.constant dense<0> : tensor<i64>
-+// CHECK:             %[[I_RESHAPED:.*]] = stablehlo.reshape %[[ITER_I]] : (tensor<i64>) -> tensor<1xi64>
-+// CHECK:             %[[C0_RESHAPED:.*]] = stablehlo.reshape %[[C0_1]] : (tensor<i64>) -> tensor<1xi64>
-+// CHECK:             %[[START_INDICES:.*]] = stablehlo.concatenate %[[I_RESHAPED]], %[[C0_RESHAPED]], dim = 0 : (tensor<1xi64>, tensor<1xi64>) -> tensor<2xi64>
-+// CHECK:             %[[C1:.*]] = stablehlo.constant dense<1> : tensor<i64>
-+// CHECK:             %[[I_PLUS_1:.*]] = stablehlo.add %[[ITER_I]], %[[C1]] : tensor<i64>
-+// CHECK:             %[[GET_DIM_SIZE_1:.*]] = stablehlo.get_dimension_size %[[ARG0]], dim = 1 : (tensor<?x?xi32, #stablehlo.bounds<2, 3>>) -> tensor<i32>
-+// CHECK:             %[[CONVERT_1:.*]] = stablehlo.convert %[[GET_DIM_SIZE_1]] : (tensor<i32>) -> tensor<i64>
-+// CHECK:             %[[LIMIT_0:.*]] = stablehlo.reshape %[[I_PLUS_1]] : (tensor<i64>) -> tensor<1xi64>
-+// CHECK:             %[[LIMIT_1:.*]] = stablehlo.reshape %[[CONVERT_1]] : (tensor<i64>) -> tensor<1xi64>
-+// CHECK:             %[[LIMIT_INDICES:.*]] = stablehlo.concatenate %[[LIMIT_0]], %[[LIMIT_1]], dim = 0 : (tensor<1xi64>, tensor<1xi64>) -> tensor<2xi64>
-+// CHECK:             %[[STRIDES:.*]] = stablehlo.constant dense<1> : tensor<2xi64>
-+// CHECK:             %[[SLICE:.*]] = stablehlo.real_dynamic_slice %[[ARG0]], %[[START_INDICES]], %[[LIMIT_INDICES]], %[[STRIDES]] : (tensor<?x?xi32, #stablehlo.bounds<2, 3>>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) -> tensor<1x?xi32, #stablehlo.bounds<?, 3>>
-+// CHECK:             %[[INPUT_ELEM:.*]] = stablehlo.reshape %[[SLICE]] : (tensor<1x?xi32, #stablehlo.bounds<?, 3>>) -> tensor<?xi32, #stablehlo.bounds<3>>
-+// CHECK:             %[[ADD:.*]] = stablehlo.add %[[INPUT_ELEM]], %[[ITER_CARRY]] : tensor<?xi32, #stablehlo.bounds<3>>
-+// CHECK:             %[[OUT_ELEM:.*]] = stablehlo.reshape %[[ADD]] : (tensor<?xi32, #stablehlo.bounds<3>>) -> tensor<1x?xi32, #stablehlo.bounds<?, 3>>
-+// CHECK:             %[[C0_2:.*]] = stablehlo.constant dense<0> : tensor<i64>
-+// CHECK:             %[[UPDATED_OUT:.*]] = stablehlo.dynamic_update_slice %[[ITER_OUT]], %[[OUT_ELEM]], %[[ITER_I]], %[[C0_2]] : (tensor<?x?xi32, #stablehlo.bounds<2, 3>>, tensor<1x?xi32, #stablehlo.bounds<?, 3>>, tensor<i64>, tensor<i64>) -> tensor<?x?xi32, #stablehlo.bounds<2, 3>>
-+// CHECK:             %[[C1_2:.*]] = stablehlo.constant dense<1> : tensor<i64>
-+// CHECK:             %[[NEXT_I:.*]] = stablehlo.add %[[ITER_I]], %[[C1_2]] : tensor<i64>
-+// CHECK:             stablehlo.return %[[NEXT_I]], %[[ADD]], %[[UPDATED_OUT]] : tensor<i64>, tensor<?xi32, #stablehlo.bounds<3>>, tensor<?x?xi32, #stablehlo.bounds<2, 3>>
-+// CHECK:           }
-+// CHECK:           return %[[WHILE]]#2, %[[WHILE]]#1 : tensor<?x?xi32, #stablehlo.bounds<2, 3>>, tensor<?xi32, #stablehlo.bounds<3>>
-+// CHECK:         }
-+func.func @scan_bounded(%arg0: tensor<?x?xi32, #stablehlo.type_extensions<bounds = [2, 3]>>, %arg1: tensor<?xi32, #stablehlo.type_extensions<bounds = [3]>>) -> (tensor<?x?xi32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?xi32, #stablehlo.type_extensions<bounds = [3]>>) {
-+  %0:2 = chlo.scan(%arg0) inits(%arg1) dimension=0 {
-+  ^bb0(%scan_arg0: tensor<?xi32, #stablehlo.type_extensions<bounds = [3]>>, %scan_arg1: tensor<?xi32, #stablehlo.type_extensions<bounds = [3]>>):
-+    %1 = stablehlo.add %scan_arg0, %scan_arg1 : tensor<?xi32, #stablehlo.type_extensions<bounds = [3]>>
-+    stablehlo.return %1, %1 : tensor<?xi32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?xi32, #stablehlo.type_extensions<bounds = [3]>>
-+  } : (tensor<?x?xi32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?xi32, #stablehlo.type_extensions<bounds = [3]>>) -> (tensor<?x?xi32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?xi32, #stablehlo.type_extensions<bounds = [3]>>)
-+  func.return %0#0, %0#1 : tensor<?x?xi32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?xi32, #stablehlo.type_extensions<bounds = [3]>>
-+}
-+
-+// -----
-+
- // CHECK-LABEL:   func.func @mulhi_i32(
- // CHECK-SAME:      %[[ARG0:.*]]: tensor<4xi32>,
- // CHECK-SAME:      %[[ARG1:.*]]: tensor<4xi32>) -> tensor<4xi32> {
-diff --ruN a/stablehlo/stablehlo/tests/interpret/compare.mlir b/stablehlo/stablehlo/tests/interpret/compare.mlir
---- stablehlo/stablehlo/tests/interpret/compare.mlir
-+++ stablehlo/stablehlo/tests/interpret/compare.mlir
-@@ -294,6 +294,68 @@
- 
- // -----
- 
-+func.func @compare_op_test_f64_weakorder_eq() {
-+  // -NaN, -NaN, -Inf, -0.0, 1.0, +Inf, +Inf, -NaN
-+  // -NaN, +NaN, +Inf, +0.0, 2.0, +Inf, -NaN, +Inf
-+  %lhs = stablehlo.constant dense<[0xFFF0000000000001, 0xFFF0000000000001, 0xFFF0000000000000, 0x8000000000000000, 1.0, 0x7FF0000000000000, 0x7FF0000000000000, 0xFFF0000000000001]> : tensor<8xf64>
-+  %rhs = stablehlo.constant dense<[0xFFF0000000000001, 0x7FF0000000000001, 0x7FF0000000000000, 0x0000000000000000, 2.0, 0x7FF0000000000000, 0xFFF0000000000001, 0x7FF0000000000000]> : tensor<8xf64>
-+  %result = stablehlo.compare EQ, %lhs, %rhs, WEAKORDER : (tensor<8xf64>, tensor<8xf64>) -> tensor<8xi1>
-+  check.expect_eq_const %result, dense<[true, true, false, true, false, true, false, false]> : tensor<8xi1>
-+  func.return
-+}
-+
-+// -----
-+
-+func.func @compare_op_test_f64_weakorder_ne() {
-+  %lhs = stablehlo.constant dense<[0xFFF0000000000001, 0xFFF0000000000001, 0xFFF0000000000000, 0x8000000000000000, 1.0, 0x7FF0000000000000, 0x7FF0000000000000, 0xFFF0000000000001]> : tensor<8xf64>
-+  %rhs = stablehlo.constant dense<[0xFFF0000000000001, 0x7FF0000000000001, 0x7FF0000000000000, 0x0000000000000000, 2.0, 0x7FF0000000000000, 0xFFF0000000000001, 0x7FF0000000000000]> : tensor<8xf64>
-+  %result = stablehlo.compare NE, %lhs, %rhs, WEAKORDER : (tensor<8xf64>, tensor<8xf64>) -> tensor<8xi1>
-+  check.expect_eq_const %result, dense<[false, false, true, false, true, false, true, true]> : tensor<8xi1>
-+  func.return
-+}
-+
-+// -----
-+
-+func.func @compare_op_test_f64_weakorder_ge() {
-+  %lhs = stablehlo.constant dense<[0xFFF0000000000001, 0xFFF0000000000001, 0xFFF0000000000000, 0x8000000000000000, 1.0, 0x7FF0000000000000, 0x7FF0000000000000, 0xFFF0000000000001]> : tensor<8xf64>
-+  %rhs = stablehlo.constant dense<[0xFFF0000000000001, 0x7FF0000000000001, 0x7FF0000000000000, 0x0000000000000000, 2.0, 0x7FF0000000000000, 0xFFF0000000000001, 0x7FF0000000000000]> : tensor<8xf64>
-+  %result = stablehlo.compare GE, %lhs, %rhs, WEAKORDER : (tensor<8xf64>, tensor<8xf64>) -> tensor<8xi1>
-+  check.expect_eq_const %result, dense<[true, true, false, true, false, true, false, true]> : tensor<8xi1>
-+  func.return
-+}
-+
-+// -----
-+
-+func.func @compare_op_test_f64_weakorder_gt() {
-+  %lhs = stablehlo.constant dense<[0xFFF0000000000001, 0xFFF0000000000001, 0xFFF0000000000000, 0x8000000000000000, 1.0, 0x7FF0000000000000, 0x7FF0000000000000, 0xFFF0000000000001]> : tensor<8xf64>
-+  %rhs = stablehlo.constant dense<[0xFFF0000000000001, 0x7FF0000000000001, 0x7FF0000000000000, 0x0000000000000000, 2.0, 0x7FF0000000000000, 0xFFF0000000000001, 0x7FF0000000000000]> : tensor<8xf64>
-+  %result = stablehlo.compare GT, %lhs, %rhs, WEAKORDER : (tensor<8xf64>, tensor<8xf64>) -> tensor<8xi1>
-+  check.expect_eq_const %result, dense<[false, false, false, false, false, false, false, true]> : tensor<8xi1>
-+  func.return
-+}
-+
-+// -----
-+
-+func.func @compare_op_test_f64_weakorder_le() {
-+  %lhs = stablehlo.constant dense<[0xFFF0000000000001, 0xFFF0000000000001, 0xFFF0000000000000, 0x8000000000000000, 1.0, 0x7FF0000000000000, 0x7FF0000000000000, 0xFFF0000000000001]> : tensor<8xf64>
-+  %rhs = stablehlo.constant dense<[0xFFF0000000000001, 0x7FF0000000000001, 0x7FF0000000000000, 0x0000000000000000, 2.0, 0x7FF0000000000000, 0xFFF0000000000001, 0x7FF0000000000000]> : tensor<8xf64>
-+  %result = stablehlo.compare LE, %lhs, %rhs, WEAKORDER : (tensor<8xf64>, tensor<8xf64>) -> tensor<8xi1>
-+  check.expect_eq_const %result, dense<[true, true, true, true, true, true, true, false]> : tensor<8xi1>
-+  func.return
-+}
-+
-+// -----
-+
-+func.func @compare_op_test_f64_weakorder_lt() {
-+  %lhs = stablehlo.constant dense<[0xFFF0000000000001, 0xFFF0000000000001, 0xFFF0000000000000, 0x8000000000000000, 1.0, 0x7FF0000000000000, 0x7FF0000000000000, 0xFFF0000000000001]> : tensor<8xf64>
-+  %rhs = stablehlo.constant dense<[0xFFF0000000000001, 0x7FF0000000000001, 0x7FF0000000000000, 0x0000000000000000, 2.0, 0x7FF0000000000000, 0xFFF0000000000001, 0x7FF0000000000000]> : tensor<8xf64>
-+  %result = stablehlo.compare LT, %lhs, %rhs, WEAKORDER : (tensor<8xf64>, tensor<8xf64>) -> tensor<8xi1>
-+  check.expect_eq_const %result, dense<[false, false, true, false, true, false, true, false]> : tensor<8xi1>
-+  func.return
-+}
-+
-+// -----
-+
- func.func @compare_op_test_c128_default() {
-   // (+NaN, +0.0)
-   // (+NaN, -0.0)
-diff --ruN a/stablehlo/stablehlo/tests/interpret/convolution.mlir b/stablehlo/stablehlo/tests/interpret/convolution.mlir
---- stablehlo/stablehlo/tests/interpret/convolution.mlir
-+++ stablehlo/stablehlo/tests/interpret/convolution.mlir
-@@ -134,3 +134,18 @@
-                                          [[21, 21, 29, 29]]]]> : tensor<1x2x1x4xi64>
-   func.return
- }
-+
-+// -----
-+
-+func.func @convolution_op_test_feature_before_batch() {
-+  %lhs = stablehlo.constant dense<[[[1.0]], [[2.0]]]> : tensor<2x1x1xf32>
-+  %rhs = stablehlo.constant dense<[[[3.0, 4.0]]]> : tensor<1x1x2xf32>
-+  %result = stablehlo.convolution(%lhs, %rhs)
-+    dim_numbers = [0, f, b]x[0, i, o]->[0, f, b],
-+    window = {stride = [1], pad = [[0, 0]], rhs_dilate = [1]}
-+    {batch_group_count = 1 : i64, feature_group_count = 1 : i64}
-+    : (tensor<2x1x1xf32>, tensor<1x1x2xf32>) -> tensor<2x2x1xf32>
-+  check.expect_almost_eq_const %result,
-+    dense<[[[3.0], [4.0]], [[6.0], [8.0]]]> : tensor<2x2x1xf32>
-+  func.return
-+}
-diff --ruN a/stablehlo/stablehlo/tests/interpret/fft.mlir b/stablehlo/stablehlo/tests/interpret/fft.mlir
---- stablehlo/stablehlo/tests/interpret/fft.mlir
-+++ stablehlo/stablehlo/tests/interpret/fft.mlir
-@@ -233,3 +233,239 @@
-   > : tensor<2x2x4xf64>
-   func.return
- }
-+
-+// -----
-+
-+func.func @fft_op_test_fft_1d_c64() {
-+  %operand = stablehlo.constant dense<[(1.0, 2.0), (3.0, -4.0)]> : tensor<2xcomplex<f32>>
-+  %result = stablehlo.fft %operand, type = FFT, length = [2] : (tensor<2xcomplex<f32>>) -> tensor<2xcomplex<f32>>
-+  check.expect_almost_eq_const %result, dense<[(4.0, -2.0), (-2.0, 6.0)]> : tensor<2xcomplex<f32>>
-+  func.return
-+}
-+
-+// -----
-+
-+func.func @fft_op_test_fft_2d_c64() {
-+  %operand = stablehlo.constant dense<
-+    [
-+      [(1.0, 2.0), (3.0, -4.0)],
-+      [(-5.0, 6.0), (-7.0, -8.0)]
-+    ]
-+  > : tensor<2x2xcomplex<f32>>
-+  %result = stablehlo.fft %operand, type = FFT, length = [2, 2] : (tensor<2x2xcomplex<f32>>) -> tensor<2x2xcomplex<f32>>
-+  check.expect_almost_eq_const %result, dense<
-+    [
-+      [(-8.0, -4.0), (0.0, 20.0)],
-+      [(16.0, 0.0), (-4.0, -8.0)]
-+    ]
-+  > : tensor<2x2xcomplex<f32>>
-+  func.return
-+}
-+
-+// -----
-+
-+func.func @fft_op_test_fft_3d_c64() {
-+  %operand = stablehlo.constant dense<
-+    [
-+      [
-+        [(1.0, 2.0), (3.0, -4.0)],
-+        [(-5.0, 6.0), (-7.0, -8.0)]
-+      ],[
-+        [(-1.0, -2.0), (-3.0, 4.0)],
-+        [(5.0, -6.0), (7.0, 8.0)]
-+      ]
-+    ]
-+  > : tensor<2x2x2xcomplex<f32>>
-+  %result = stablehlo.fft %operand, type = FFT, length = [2, 2, 2] : (tensor<2x2x2xcomplex<f32>>) -> tensor<2x2x2xcomplex<f32>>
-+  check.expect_almost_eq_const %result, dense<
-+    [
-+      [
-+        [(0.0, 0.0), (0.0, 0.0)],
-+        [(0.0, 0.0), (0.0, 0.0)]
-+      ],[
-+        [(-16.0, -8.0), (0.0, 40.0)],
-+        [(32.0, 0.0), (-8.0, -16.0)]
-+      ]
-+    ]
-+  > : tensor<2x2x2xcomplex<f32>>
-+  func.return
-+}
-+
-+// -----
-+
-+func.func @fft_op_test_ifft_1d_c64() {
-+  %operand = stablehlo.constant dense<[(4.0, -2.0), (-2.0, 6.0)]> : tensor<2xcomplex<f32>>
-+  %result = stablehlo.fft %operand, type = IFFT, length = [2] : (tensor<2xcomplex<f32>>) -> tensor<2xcomplex<f32>>
-+  check.expect_almost_eq_const %result, dense<[(1.0, 2.0), (3.0, -4.0)]> : tensor<2xcomplex<f32>>
-+  func.return
-+}
-+
-+// -----
-+
-+func.func @fft_op_test_ifft_2d_c64() {
-+  %operand = stablehlo.constant dense<
-+    [
-+      [(-8.0, -4.0), (16.0, 0.0)],
-+      [(0.0, 20.0), (-4.0, -8.0)]
-+    ]
-+  > : tensor<2x2xcomplex<f32>>
-+  %result = stablehlo.fft %operand, type = IFFT, length = [2, 2] : (tensor<2x2xcomplex<f32>>) -> tensor<2x2xcomplex<f32>>
-+  check.expect_almost_eq_const %result, dense<
-+    [
-+      [(1.0, 2.0), (-5.0, 6.0)],
-+      [(3.0, -4.0), (-7.0, -8.0)]
-+    ]
-+  > : tensor<2x2xcomplex<f32>>
-+  func.return
-+}
-+
-+// -----
-+
-+func.func @fft_op_test_ifft_3d_c64() {
-+  %operand = stablehlo.constant dense<
-+    [
-+      [
-+        [(0.0, 0.0), (0.0, 0.0)],
-+        [(0.0, 0.0), (0.0, 0.0)]
-+      ],[
-+        [(-16.0, -8.0), (32.0, 0.0)],
-+        [(0.0, 40.0), (-8.0, -16.0)]
-+      ]
-+    ]
-+  > : tensor<2x2x2xcomplex<f32>>
-+  %result = stablehlo.fft %operand, type = IFFT, length = [2, 2, 2] : (tensor<2x2x2xcomplex<f32>>) -> tensor<2x2x2xcomplex<f32>>
-+  check.expect_almost_eq_const %result, dense<
-+    [
-+      [
-+        [(1.0, 2.0), (-5.0, 6.0)],
-+        [(3.0, -4.0), (-7.0, -8.0)]
-+      ],[
-+        [(-1.0, -2.0), (5.0, -6.0)],
-+        [(-3.0, 4.0), (7.0, 8.0)]
-+      ]
-+    ]
-+  > : tensor<2x2x2xcomplex<f32>>
-+  func.return
-+}
-+
-+// -----
-+
-+func.func @fft_op_test_rfft_1d_f32() {
-+  %operand = stablehlo.constant dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
-+  %result = stablehlo.fft %operand, type = RFFT, length = [4] : (tensor<4xf32>) -> tensor<3xcomplex<f32>>
-+  check.expect_almost_eq_const %result, dense<[(10.0, 0.0), (-2.0, 2.0), (-2.0, 0.0)]> : tensor<3xcomplex<f32>>
-+  func.return
-+}
-+
-+// -----
-+
-+func.func @fft_op_test_rfft_2d_f32() {
-+  %operand = stablehlo.constant dense<
-+    [
-+      [1.0, 2.0, 3.0, 4.0],
-+      [5.0, 6.0, 7.0, 8.0],
-+      [9.0, 10.0, 11.0, 12.0],
-+      [13.0, 14.0, 15.0, 16.0]
-+    ]
-+  > : tensor<4x4xf32>
-+  %result = stablehlo.fft %operand, type = RFFT, length = [4, 4] : (tensor<4x4xf32>) -> tensor<4x3xcomplex<f32>>
-+  check.expect_almost_eq_const %result, dense<
-+    [
-+      [(136.0, 0.0), (-8.0, 8.0), (-8.0, 0.0)],
-+      [(-32.0, 32.0), (0.0, 0.0), (0.0, 0.0)],
-+      [(-32.0, 0.0), (0.0, 0.0), (0.0, 0.0)],
-+      [(-32.0, -32.0), (0.0, 0.0), (0.0, 0.0)]
-+    ]
-+  > : tensor<4x3xcomplex<f32>>
-+  func.return
-+}
-+
-+// -----
-+
-+func.func @fft_op_test_rfft_3d_f32() {
-+  %operand = stablehlo.constant dense<
-+    [
-+      [
-+        [1.0, 2.0, 3.0, 4.0],
-+        [5.0, 6.0, 7.0, 8.0]
-+      ],[
-+        [9.0, 10.0, 11.0, 12.0],
-+        [13.0, 14.0, 15.0, 16.0]
-+      ]
-+    ]
-+  > : tensor<2x2x4xf32>
-+  %result = stablehlo.fft %operand, type = RFFT, length = [2, 2, 4] : (tensor<2x2x4xf32>) -> tensor<2x2x3xcomplex<f32>>
-+  check.expect_almost_eq_const %result, dense<
-+    [
-+      [
-+        [(136.0, 0.0), (-8.0, 8.0), (-8.0, 0.0)],
-+        [(-32.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
-+      ],[
-+        [(-64.0, 0.0), (0.0, 0.0), (0.0, 0.0)],
-+        [(0.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
-+      ]
-+    ]
-+  > : tensor<2x2x3xcomplex<f32>>
-+  func.return
-+}
-+
-+// -----
-+
-+func.func @fft_op_test_irfft_1d_c64() {
-+  %operand = stablehlo.constant dense<[(10.0, 0.0), (-2.0, 2.0), (-2.0, 0.0)]> : tensor<3xcomplex<f32>>
-+  %result = stablehlo.fft %operand, type = IRFFT, length = [4] : (tensor<3xcomplex<f32>>) -> tensor<4xf32>
-+  check.expect_almost_eq_const %result, dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
-+  func.return
-+}
-+
-+// -----
-+
-+func.func @fft_op_test_irfft_2d_c64() {
-+  %operand = stablehlo.constant dense<
-+    [
-+      [(136.0, 0.0), (-8.0, 8.0), (-8.0, 0.0)],
-+      [(-32.0, 32.0), (0.0, 0.0), (0.0, 0.0)],
-+      [(-32.0, 0.0), (0.0, 0.0), (0.0, 0.0)],
-+      [(-32.0, -32.0), (0.0, 0.0), (0.0, 0.0)]
-+    ]
-+  > : tensor<4x3xcomplex<f32>>
-+  %result = stablehlo.fft %operand, type = IRFFT, length = [4, 4] : (tensor<4x3xcomplex<f32>>) -> tensor<4x4xf32>
-+  check.expect_almost_eq_const %result, dense<
-+    [
-+      [1.0, 2.0, 3.0, 4.0],
-+      [5.0, 6.0, 7.0, 8.0],
-+      [9.0, 10.0, 11.0, 12.0],
-+      [13.0, 14.0, 15.0, 16.0]
-+    ]
-+  > : tensor<4x4xf32>
-+  func.return
-+}
-+
-+// -----
-+
-+func.func @fft_op_test_irfft_3d_c64() {
-+  %operand = stablehlo.constant dense<
-+    [
-+      [
-+        [(136.0, 0.0), (-8.0, 8.0), (-8.0, 0.0)],
-+        [(-32.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
-+      ],[
-+        [(-64.0, 0.0), (0.0, 0.0), (0.0, 0.0)],
-+        [(0.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
-+      ]
-+    ]
-+  > : tensor<2x2x3xcomplex<f32>>
-+  %result = stablehlo.fft %operand, type = IRFFT, length = [2, 2, 4] : (tensor<2x2x3xcomplex<f32>>) -> tensor<2x2x4xf32>
-+  check.expect_almost_eq_const %result, dense<
-+    [
-+      [
-+        [1.0, 2.0, 3.0, 4.0],
-+        [5.0, 6.0, 7.0, 8.0]
-+      ],[
-+        [9.0, 10.0, 11.0, 12.0],
-+        [13.0, 14.0, 15.0, 16.0]
-+      ]
-+    ]
-+  > : tensor<2x2x4xf32>
-+  func.return
-+}
-diff --ruN a/stablehlo/stablehlo/tests/ops_chlo.mlir b/stablehlo/stablehlo/tests/ops_chlo.mlir
---- stablehlo/stablehlo/tests/ops_chlo.mlir
-+++ stablehlo/stablehlo/tests/ops_chlo.mlir
-@@ -489,6 +489,50 @@
- 
- // -----
- 
-+// CHECK-LABEL: func @scan_bounded_dynamism
-+func.func @scan_bounded_dynamism(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, %arg1: tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>) -> (tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>) {
-+  // CHECK: chlo.scan
-+  // CHECK: } : (tensor<?x?xf32, #stablehlo.bounds<2, 3>>, tensor<?xf32, #stablehlo.bounds<3>>) -> (tensor<?x?xf32, #stablehlo.bounds<2, 3>>, tensor<?xf32, #stablehlo.bounds<3>>)
-+  %0, %1 = chlo.scan (%arg0) inits (%arg1) dimension = 0 {
-+  ^bb0(%input0: tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, %carry0: tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>):
-+    %2 = stablehlo.add %input0, %carry0 : tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>
-+    stablehlo.return %2, %2 : tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>
-+  } : (tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>) -> (tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>)
-+  func.return %0, %1 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>
-+}
-+
-+// -----
-+
-+// CHECK-LABEL: func @scan_bounded_dynamism_multiple_inputs
-+func.func @scan_bounded_dynamism_multiple_inputs(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, %arg1: tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, %arg2: tensor<?x?x?xi32, #stablehlo.type_extensions<bounds = [2, 4, 5]>>, %arg3: tensor<?x?xi32, #stablehlo.type_extensions<bounds = [4, 5]>>) -> (tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?x?x?xi32, #stablehlo.type_extensions<bounds = [2, 4, 5]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?x?xi32, #stablehlo.type_extensions<bounds = [4, 5]>>) {
-+  // CHECK: chlo.scan
-+  // CHECK: } : (tensor<?x?xf32, #stablehlo.bounds<2, 3>>, tensor<?x?x?xi32, #stablehlo.bounds<2, 4, 5>>, tensor<?xf32, #stablehlo.bounds<3>>, tensor<?x?xi32, #stablehlo.bounds<4, 5>>) -> (tensor<?x?xf32, #stablehlo.bounds<2, 3>>, tensor<?x?x?xi32, #stablehlo.bounds<2, 4, 5>>, tensor<?xf32, #stablehlo.bounds<3>>, tensor<?x?xi32, #stablehlo.bounds<4, 5>>)
-+  %0:4 = chlo.scan (%arg0, %arg2) inits (%arg1, %arg3) dimension = 0 {
-+  ^bb0(%input0: tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, %input1: tensor<?x?xi32, #stablehlo.type_extensions<bounds = [4, 5]>>, %carry0: tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, %carry1: tensor<?x?xi32, #stablehlo.type_extensions<bounds = [4, 5]>>):
-+    %1 = stablehlo.add %input0, %carry0 : tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>
-+    %2 = stablehlo.add %input1, %carry1 : tensor<?x?xi32, #stablehlo.type_extensions<bounds = [4, 5]>>
-+    stablehlo.return %1, %2, %1, %2 : tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?x?xi32, #stablehlo.type_extensions<bounds = [4, 5]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?x?xi32, #stablehlo.type_extensions<bounds = [4, 5]>>
-+  } : (tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?x?x?xi32, #stablehlo.type_extensions<bounds = [2, 4, 5]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?x?xi32, #stablehlo.type_extensions<bounds = [4, 5]>>) -> (tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?x?x?xi32, #stablehlo.type_extensions<bounds = [2, 4, 5]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?x?xi32, #stablehlo.type_extensions<bounds = [4, 5]>>)
-+  func.return %0#0, %0#1, %0#2, %0#3 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?x?x?xi32, #stablehlo.type_extensions<bounds = [2, 4, 5]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?x?xi32, #stablehlo.type_extensions<bounds = [4, 5]>>
-+}
-+
-+// -----
-+
-+// CHECK-LABEL: func @scan_bounded_dynamism_multiple_inputs_same_rank
-+func.func @scan_bounded_dynamism_multiple_inputs_same_rank(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, %arg1: tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, %arg2: tensor<?x?xi32, #stablehlo.type_extensions<bounds = [2, 4]>>, %arg3: tensor<?xi32, #stablehlo.type_extensions<bounds = [4]>>) -> (tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?x?xi32, #stablehlo.type_extensions<bounds = [2, 4]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?xi32, #stablehlo.type_extensions<bounds = [4]>>) {
-+  // CHECK: chlo.scan
-+  // CHECK: } : (tensor<?x?xf32, #stablehlo.bounds<2, 3>>, tensor<?x?xi32, #stablehlo.bounds<2, 4>>, tensor<?xf32, #stablehlo.bounds<3>>, tensor<?xi32, #stablehlo.bounds<4>>) -> (tensor<?x?xf32, #stablehlo.bounds<2, 3>>, tensor<?x?xi32, #stablehlo.bounds<2, 4>>, tensor<?xf32, #stablehlo.bounds<3>>, tensor<?xi32, #stablehlo.bounds<4>>)
-+  %0:4 = chlo.scan (%arg0, %arg2) inits (%arg1, %arg3) dimension = 0 {
-+  ^bb0(%input0: tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, %input1: tensor<?xi32, #stablehlo.type_extensions<bounds = [4]>>, %carry0: tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, %carry1: tensor<?xi32, #stablehlo.type_extensions<bounds = [4]>>):
-+    %1 = stablehlo.add %input0, %carry0 : tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>
-+    %2 = stablehlo.add %input1, %carry1 : tensor<?xi32, #stablehlo.type_extensions<bounds = [4]>>
-+    stablehlo.return %1, %2, %1, %2 : tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?xi32, #stablehlo.type_extensions<bounds = [4]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?xi32, #stablehlo.type_extensions<bounds = [4]>>
-+  } : (tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?x?xi32, #stablehlo.type_extensions<bounds = [2, 4]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?xi32, #stablehlo.type_extensions<bounds = [4]>>) -> (tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?x?xi32, #stablehlo.type_extensions<bounds = [2, 4]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?xi32, #stablehlo.type_extensions<bounds = [4]>>)
-+  func.return %0#0, %0#1, %0#2, %0#3 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?x?xi32, #stablehlo.type_extensions<bounds = [2, 4]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?xi32, #stablehlo.type_extensions<bounds = [4]>>
-+}
-+
-+// -----
-+
- func.func @scan_explicit_size_mismatch(%arg0: tensor<2x3xf32>, %arg1: tensor<3xf32>) -> tensor<2x3xf32> {
-   // expected-error @+1 {{'chlo.scan' op invalid scan dimension size of operand 0}}
-   %0, %1 = chlo.scan (%arg0) inits (%arg1) dimension = 0 attributes { scan_dim_size = 3 : i64 } {
-@@ -561,6 +605,30 @@
- 
- // -----
- 
-+func.func @scan_bounded_dynamism_mismatch(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, %arg1: tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>) -> (tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>) {
-+  // expected-error @+1 {{'chlo.scan' op operand and body argument 0 are incompatible}}
-+  %0, %1 = chlo.scan (%arg0) inits (%arg1) dimension = 0 {
-+  ^bb0(%input0: tensor<4xf32>, %carry0: tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>):
-+    stablehlo.return %carry0, %carry0 : tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>
-+  } : (tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>) -> (tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>)
-+  func.return %0, %1 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>
-+}
-+
-+// -----
-+
-+func.func @scan_bounded_dynamism_result_rank_mismatch(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, %arg1: tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>) -> (tensor<?xf32>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>) {
-+  // expected-error @+2 {{op failed to infer returned types}}
-+  // expected-error @+1 {{'chlo.scan' op inferred type(s) 'tensor<?x?xf32, #stablehlo.bounds<2, 3>>', 'tensor<?xf32, #stablehlo.bounds<3>>' are incompatible with return type(s) of operation 'tensor<?xf32>', 'tensor<?xf32, #stablehlo.bounds<3>>'}}
-+  %0, %1 = chlo.scan (%arg0) inits (%arg1) dimension = 0 {
-+  ^bb0(%input0: tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, %carry0: tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>):
-+    %2 = stablehlo.add %input0, %carry0 : tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>
-+    stablehlo.return %2, %2 : tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>
-+  } : (tensor<?x?xf32, #stablehlo.type_extensions<bounds = [2, 3]>>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>) -> (tensor<?xf32>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>)
-+  func.return %0, %1 : tensor<?xf32>, tensor<?xf32, #stablehlo.type_extensions<bounds = [3]>>
-+}
-+
-+// -----
-+
- // CHECK-LABEL: func @mulhi_i32
- func.func @mulhi_i32(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>) -> tensor<4xi32> {
-   %0 = "chlo.mulhi"(%arg0, %arg1) : (tensor<4xi32>, tensor<4xi32>) -> tensor<4xi32>
-diff --ruN a/stablehlo/stablehlo/tests/ops_chlo_roundtrip.mlir b/stablehlo/stablehlo/tests/ops_chlo_roundtrip.mlir
---- stablehlo/stablehlo/tests/ops_chlo_roundtrip.mlir
-+++ stablehlo/stablehlo/tests/ops_chlo_roundtrip.mlir
-@@ -367,12 +367,14 @@
- // CHECK-SAME:  %[[A1:.*]]: tensor<2x3x4xf64>
- // CHECK:       %[[T0:.*]] = chlo.broadcast_compare %[[A0]], %[[A1]] {comparison_direction = #chlo<comparison_direction LT>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
- // CHECK:       %[[T1:.*]] = chlo.broadcast_compare %[[A0]], %[[A1]] {compare_type = #chlo<comparison_type TOTALORDER>, comparison_direction = #chlo<comparison_direction LT>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
--// CHECK:       %[[T2:.*]] = chlo.broadcast_compare %[[A0]], %[[A1]] {broadcast_dimensions = array<i64: 1>, comparison_direction = #chlo<comparison_direction LT>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
-+// CHECK:       %[[T2:.*]] = chlo.broadcast_compare %[[A0]], %[[A1]] {compare_type = #chlo<comparison_type WEAKORDER>, comparison_direction = #chlo<comparison_direction LT>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
-+// CHECK:       %[[T3:.*]] = chlo.broadcast_compare %[[A0]], %[[A1]] {broadcast_dimensions = array<i64: 1>, comparison_direction = #chlo<comparison_direction LT>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
- // CHECK:       return %[[T0]] : tensor<2x3x4xi1>
- func.func @chlo_broadcast_compare(%arg0: tensor<2x3x4xf64>, %arg1: tensor<2x3x4xf64>) -> tensor<2x3x4xi1> {
-   %0 = chlo.broadcast_compare %arg0, %arg1 {comparison_direction = #chlo<comparison_direction LT>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
-   %1 = chlo.broadcast_compare %arg0, %arg1 {comparison_direction = #chlo<comparison_direction LT>, compare_type = #chlo<comparison_type TOTALORDER>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
--  %2 = chlo.broadcast_compare %arg0, %arg1 {comparison_direction = #chlo<comparison_direction LT>, broadcast_dimensions = array<i64: 1>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
-+  %2 = chlo.broadcast_compare %arg0, %arg1 {comparison_direction = #chlo<comparison_direction LT>, compare_type = #chlo<comparison_type WEAKORDER>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
-+  %3 = chlo.broadcast_compare %arg0, %arg1 {comparison_direction = #chlo<comparison_direction LT>, broadcast_dimensions = array<i64: 1>} : (tensor<2x3x4xf64>, tensor<2x3x4xf64>) -> tensor<2x3x4xi1>
-   return %0 : tensor<2x3x4xi1>
- }
- 
-diff --ruN a/stablehlo/stablehlo/tests/ops_dot_general_algorithms.mlir b/stablehlo/stablehlo/tests/ops_dot_general_algorithms.mlir
---- stablehlo/stablehlo/tests/ops_dot_general_algorithms.mlir
-+++ stablehlo/stablehlo/tests/ops_dot_general_algorithms.mlir
-@@ -119,6 +119,40 @@
-   }> : (tensor<2x2x2xbf16>, tensor<2x2x2xbf16>) -> tensor<2x2x2xbf16>  return %0 : tensor<2x2x2xbf16>
- }
- 
-+// CHECK-LABEL: func @dot_algorithm_f8e4m3fn_f8e4m3fn_f32_x3
-+func.func @dot_algorithm_f8e4m3fn_f8e4m3fn_f32_x3(%arg0: tensor<2x2x2xbf16>, %arg1: tensor<2x2x2xbf16>) -> tensor<2x2x2xf32> {
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) <{
-+    dot_dimension_numbers = #stablehlo.dot<lhs_batching_dimensions = [0], rhs_batching_dimensions = [0], lhs_contracting_dimensions = [2], rhs_contracting_dimensions = [1]>,
-+    precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>],
-+    algorithm = #stablehlo.dot_algorithm<
-+      lhs_precision_type = f8E4M3FN,
-+      rhs_precision_type = f8E4M3FN,
-+      accumulation_type = f32,
-+      lhs_component_count = 1,
-+      rhs_component_count = 1,
-+      num_primitive_operations = 3,
-+      allow_imprecise_accumulation = false
-+    >
-+  }> : (tensor<2x2x2xbf16>, tensor<2x2x2xbf16>) -> tensor<2x2x2xf32>  return %0 : tensor<2x2x2xf32>
-+}
-+
-+// CHECK-LABEL: func @dot_algorithm_f8e4m3fn_f8e4m3fn_f32_x4
-+func.func @dot_algorithm_f8e4m3fn_f8e4m3fn_f32_x4(%arg0: tensor<2x2x2xbf16>, %arg1: tensor<2x2x2xbf16>) -> tensor<2x2x2xf32> {
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) <{
-+    dot_dimension_numbers = #stablehlo.dot<lhs_batching_dimensions = [0], rhs_batching_dimensions = [0], lhs_contracting_dimensions = [2], rhs_contracting_dimensions = [1]>,
-+    precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>],
-+    algorithm = #stablehlo.dot_algorithm<
-+      lhs_precision_type = f8E4M3FN,
-+      rhs_precision_type = f8E4M3FN,
-+      accumulation_type = f32,
-+      lhs_component_count = 1,
-+      rhs_component_count = 1,
-+      num_primitive_operations = 4,
-+      allow_imprecise_accumulation = false
-+    >
-+  }> : (tensor<2x2x2xbf16>, tensor<2x2x2xbf16>) -> tensor<2x2x2xf32>  return %0 : tensor<2x2x2xf32>
-+}
-+
- // CHECK-LABEL: func @dot_algorithm_bf16_bf16_f32_x6
- func.func @dot_algorithm_bf16_bf16_f32_x6(%arg0: tensor<2x2x2xbf16>, %arg1: tensor<2x2x2xbf16>) -> tensor<2x2x2xbf16> {
-   %0 = "stablehlo.dot_general"(%arg0, %arg1) <{
-diff --ruN a/stablehlo/stablehlo/tests/ops_stablehlo.mlir b/stablehlo/stablehlo/tests/ops_stablehlo.mlir
---- stablehlo/stablehlo/tests/ops_stablehlo.mlir
-+++ stablehlo/stablehlo/tests/ops_stablehlo.mlir
-@@ -1548,6 +1548,29 @@
- 
- // -----
- 
-+// CHECK-LABEL: func @compare_weakorder_f32
-+func.func @compare_weakorder_f32(%arg0: tensor<3xf32>, %arg1: tensor<3xf32>) -> tensor<3xi1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction LT>,
-+    compare_type = #stablehlo<comparison_type WEAKORDER>
-+  } : (tensor<3xf32>, tensor<3xf32>) -> tensor<3xi1>
-+  func.return %0 : tensor<3xi1>
-+}
-+
-+// -----
-+
-+func.func @compare_weakorder_invalid_i32(%arg0: tensor<3xi32>, %arg1: tensor<3xi32>) -> tensor<3xi1> {
-+  // expected-error@+2 {{failed to infer returned types}}
-+  // expected-error@+1 {{WEAKORDER comparison type is only supported for floating-point element types, got 'i32'}}
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction LT>,
-+    compare_type = #stablehlo<comparison_type WEAKORDER>
-+  } : (tensor<3xi32>, tensor<3xi32>) -> tensor<3xi1>
-+  func.return %0 : tensor<3xi1>
-+}
-+
-+// -----
-+
- // CHECK-LABEL: func @collective_broadcast_c3
- func.func @collective_broadcast_c3(%operand: tensor<16x8xf32>) -> tensor<16x8xf32> {
-   %0 = "stablehlo.collective_broadcast"(%operand) {
-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
-@@ -1,4 +1,4 @@
--// RUN: stablehlo-opt %s -stablehlo-compatibility-expander="target=1.14.0" -allow-unregistered-dialect | FileCheck %s
-+// RUN: stablehlo-opt %s -stablehlo-compatibility-expander="target=1.14.0" | FileCheck %s
- 
- module {
- 
-@@ -40,7 +40,7 @@
- 
-     return %0 : tensor<4xf32>
-   }
--  "sdy.mesh"() {sym_name = "sdy_mesh", stablehlo.mesh = {axes = [{name = "x", size = 2 : i64}, {name = "y", size = 2 : i64}]}} : () -> ()
-+  func.func private @sdy_mesh() attributes {stablehlo.mesh = {axes = [{name = "x", size = 2 : i64}, {name = "y", size = 2 : i64}]}}
- 
-   // CHECK-LABEL: @all_reduce_sdy_mesh
-   func.func @all_reduce_sdy_mesh(%arg0: tensor<4xf32>) -> tensor<4xf32> {
-@@ -55,7 +55,7 @@
-     return %0 : tensor<4xf32>
-   }
- 
--  "sdy.mesh"() {sym_name = "sdy_mesh_dev", stablehlo.mesh = {axes = [{name = "x", size = 2 : i64}, {name = "y", size = 2 : i64}], device_ids = dense<[0, 2, 1, 3]> : tensor<4xi64>}} : () -> ()
-+  func.func private @sdy_mesh_dev() attributes {stablehlo.mesh = {axes = [{name = "x", size = 2 : i64}, {name = "y", size = 2 : i64}], device_ids = dense<[0, 2, 1, 3]> : tensor<4xi64>}}
- 
-   // CHECK-LABEL: @all_reduce_sdy_mesh_dev
-   func.func @all_reduce_sdy_mesh_dev(%arg0: tensor<4xf32>) -> tensor<4xf32> {
-@@ -70,7 +70,7 @@
-     return %0 : tensor<4xf32>
-   }
- 
--  "sdy.mesh"() {sym_name = "sdy_mesh_max", stablehlo.mesh = {axes = [], device_ids = dense<0> : tensor<1xi64>}} : () -> ()
-+  func.func private @sdy_mesh_max() attributes {stablehlo.mesh = {axes = [], device_ids = dense<0> : tensor<1xi64>}}
- 
-   // CHECK-LABEL: @all_reduce_sdy_mesh_max
-   func.func @all_reduce_sdy_mesh_max(%arg0: tensor<4xf32>) -> tensor<4xf32> {
-diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_aggressive_folder.mlir b/stablehlo/stablehlo/tests/transforms/stablehlo_aggressive_folder.mlir
---- stablehlo/stablehlo/tests/transforms/stablehlo_aggressive_folder.mlir
-+++ stablehlo/stablehlo/tests/transforms/stablehlo_aggressive_folder.mlir
-@@ -560,6 +560,31 @@
- 
- // -----
- 
-+// CHECK-LABEL: func.func @compare_fold_weakorder
-+func.func @compare_fold_weakorder()
-+  -> (tensor<i1>, tensor<i1>, tensor<i1>, tensor<i1>, tensor<i1>, tensor<i1>) {
-+  %neg_zero = stablehlo.constant dense<-0.0> : tensor<f32>
-+  %pos_zero = stablehlo.constant dense<0.0> : tensor<f32>
-+  %pos_inf = stablehlo.constant dense<0x7F800000> : tensor<f32>
-+  %neg_nan = stablehlo.constant dense<0xFFC00000> : tensor<f32>
-+  %pos_nan = stablehlo.constant dense<0x7FC00000> : tensor<f32>
-+
-+  %0 = stablehlo.compare LT, %neg_zero, %pos_zero, WEAKORDER : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  %1 = stablehlo.compare EQ, %neg_zero, %pos_zero, WEAKORDER : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  %2 = stablehlo.compare EQ, %neg_nan, %pos_nan, WEAKORDER : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  %3 = stablehlo.compare LT, %neg_nan, %pos_nan, WEAKORDER : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  %4 = stablehlo.compare LT, %pos_inf, %neg_nan, WEAKORDER : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  %5 = stablehlo.compare GT, %neg_nan, %pos_inf, WEAKORDER : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+
-+  // CHECK-DAG:  [[FALSE:%.+]] = stablehlo.constant dense<false> : tensor<i1>
-+  // CHECK-DAG:  [[TRUE:%.+]] = stablehlo.constant dense<true> : tensor<i1>
-+  // CHECK-NEXT: return [[FALSE]], [[TRUE]], [[TRUE]], [[FALSE]], [[TRUE]], [[TRUE]]
-+  return %0, %1, %2, %3, %4, %5 :
-+         tensor<i1>, tensor<i1>, tensor<i1>, tensor<i1>, tensor<i1>, tensor<i1>
-+}
-+
-+// -----
-+
- ////////
- // ConcatenateOp
- 
-diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir b/stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir
---- stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir
-+++ stablehlo/stablehlo/tests/transforms/stablehlo_canonicalize_dynamism.mlir
-@@ -168,6 +168,18 @@
- 
- // -----
- 
-+// CHECK-LABEL: func @dynamic_broadcast_in_dim_discardable_attrs
-+func.func @dynamic_broadcast_in_dim_discardable_attrs(%arg0: tensor<4xf32>) -> tensor<3x4xf32> {
-+  //  CHECK-NOT: stablehlo.dynamic_broadcast_in_dim
-+  //      CHECK: stablehlo.broadcast_in_dim %arg0, dims = [1]
-+  // CHECK-SAME:   {mhlo.frontend_attributes = {MUST_FUSE = "0"}}
-+  %0 = stablehlo.constant dense<[3, 4]> : tensor<2xi64>
-+  %1 = stablehlo.dynamic_broadcast_in_dim %arg0, %0, dims = [1] {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<4xf32>, tensor<2xi64>) -> tensor<3x4xf32>
-+  return %1 : tensor<3x4xf32>
-+}
-+
-+// -----
-+
- // CHECK-LABEL: func @dynamic_broadcast_in_dim_inapplicable_dynamic_operand_type
- func.func @dynamic_broadcast_in_dim_inapplicable_dynamic_operand_type(%arg0: tensor<?xf32>) -> tensor<3x4xf32> {
-   // CHECK: stablehlo.dynamic_broadcast_in_dim
-@@ -225,6 +237,26 @@
- 
- // -----
- 
-+// CHECK-LABEL: func @dynamic_conv_discardable_attrs
-+func.func @dynamic_conv_discardable_attrs(%arg0: tensor<100x26x26x32xf32>, %arg1: tensor<3x3x1x32xf32>) -> tensor<100x28x28x1xf32> {
-+  //  CHECK-NOT: stablehlo.dynamic_conv
-+  //      CHECK: stablehlo.convolution(%arg0, %arg1)
-+  // CHECK-SAME:   mhlo.frontend_attributes = {MUST_FUSE = "0"}
-+  %0 = stablehlo.constant dense<2> : tensor<2x2xi32>
-+  %1 = "stablehlo.dynamic_conv"(%arg0, %arg1, %0) {
-+    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, o, i]->[b, 0, 1, f]>,
-+    window_strides = array<i64: 1, 1>,
-+    lhs_dilation = array<i64: 1, 1>,
-+    rhs_dilation = array<i64: 1, 1>,
-+    feature_group_count = 1 : i64,
-+    batch_group_count = 1 : i64,
-+    mhlo.frontend_attributes = {MUST_FUSE = "0"}
-+  } : (tensor<100x26x26x32xf32>, tensor<3x3x1x32xf32>, tensor<2x2xi32>) -> tensor<100x28x28x1xf32>
-+  return %1 : tensor<100x28x28x1xf32>
-+}
-+
-+// -----
-+
- // CHECK-LABEL: func @dynamic_conv_success_dynamic_result_type
- func.func @dynamic_conv_success_dynamic_result_type(%arg0: tensor<100x26x26x32xf32>, %arg1: tensor<3x3x1x32xf32>) -> tensor<?x28x28x1xf32> {
-   //  CHECK-NOT: stablehlo.dynamic_conv
-@@ -295,6 +327,25 @@
- 
- // -----
- 
-+// CHECK-LABEL: @dynamic_gather_discardable_attrs
-+func.func @dynamic_gather_discardable_attrs(%arg0 : tensor<2x4x9xi32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x8xi32> {
-+  //  CHECK-NOT: stablehlo.dynamic_gather
-+  //      CHECK: "stablehlo.gather"(%arg0, %arg1)
-+  // CHECK-SAME:   {mhlo.frontend_attributes = {MUST_FUSE = "0"}}
-+  %0 = stablehlo.constant dense<[1, 1, 8]> : tensor<3xi32>
-+  %1 = "stablehlo.dynamic_gather"(%arg0, %arg1, %0) <{
-+    dimension_numbers = #stablehlo.gather<
-+      collapsed_slice_dims = [0, 1],
-+      index_vector_dim = 2,
-+      offset_dims = [2],
-+      start_index_map = [0, 1]
-+    >
-+  }> {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<2x4x9xi32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<1x5x8xi32>
-+  return %1 : tensor<1x5x8xi32>
-+}
-+
-+// -----
-+
- // CHECK-LABEL: @dynamic_gather_success_dynamic_result_type
- func.func @dynamic_gather_success_dynamic_result_type(%arg0 : tensor<2x4x9xi32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x?xi32> {
-   //  CHECK-NOT: stablehlo.dynamic_gather
-@@ -348,6 +399,18 @@
- 
- // -----
- 
-+// CHECK-LABEL: func @dynamic_iota_discardable_attrs
-+func.func @dynamic_iota_discardable_attrs() -> tensor<4xf32> {
-+  //  CHECK-NOT: stablehlo.dynamic_iota
-+  //      CHECK: stablehlo.iota
-+  // CHECK-SAME:   {mhlo.frontend_attributes = {MUST_FUSE = "0"}}
-+  %0 = stablehlo.constant dense<4> : tensor<1xi64>
-+  %1 = stablehlo.dynamic_iota %0, dim = 0 {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<1xi64>) -> tensor<4xf32>
-+  return %1 : tensor<4xf32>
-+}
-+
-+// -----
-+
- // CHECK-LABEL: func @dynamic_iota_inapplicable_dynamic_output_shape
- func.func @dynamic_iota_inapplicable_dynamic_output_shape(%arg0: tensor<1xi64>) -> tensor<4xf32> {
-   // CHECK: stablehlo.dynamic_iota
-@@ -380,6 +443,20 @@
- 
- // -----
- 
-+// CHECK-LABEL: func @dynamic_pad_discardable_attrs
-+func.func @dynamic_pad_discardable_attrs(%arg0: tensor<4xf32>, %arg1: tensor<f32>) -> tensor<11xf32> {
-+  //  CHECK-NOT: stablehlo.dynamic_pad
-+  //      CHECK: stablehlo.pad
-+  // CHECK-SAME:   {mhlo.frontend_attributes = {MUST_FUSE = "0"}}
-+  %0 = stablehlo.constant dense<0> : tensor<1xi64>
-+  %1 = stablehlo.constant dense<1> : tensor<1xi64>
-+  %2 = stablehlo.constant dense<2> : tensor<1xi64>
-+  %3 = stablehlo.dynamic_pad %arg0, %arg1, %0, %1, %2 {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<4xf32>, tensor<f32>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) -> tensor<11xf32>
-+  return %3 : tensor<11xf32>
-+}
-+
-+// -----
-+
- // CHECK-LABEL: func @dynamic_pad_success_dynamic_result_type
- func.func @dynamic_pad_success_dynamic_result_type(%arg0: tensor<4xf32>, %arg1: tensor<f32>) -> tensor<?xf32> {
-   // CHECK-NOT: stablehlo.dynamic_pad
-@@ -437,6 +514,17 @@
- 
- // -----
- 
-+// CHECK-LABEL: func @dynamic_reshape_discardable_attrs
-+func.func @dynamic_reshape_discardable_attrs(%arg0: tensor<4xf32>) -> tensor<1x4xf32> {
-+  //  CHECK-NOT: stablehlo.dynamic_reshape
-+  //      CHECK: stablehlo.reshape %arg0 {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<4xf32>) -> tensor<1x4xf32>
-+  %0 = stablehlo.constant dense<[1, 4]> : tensor<2xi64>
-+  %1 = stablehlo.dynamic_reshape %arg0, %0 {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<4xf32>, tensor<2xi64>) -> tensor<1x4xf32>
-+  return %1 : tensor<1x4xf32>
-+}
-+
-+// -----
-+
- // CHECK-LABEL: func @dynamic_reshape_inapplicable_dynamic_output_shape
- func.func @dynamic_reshape_inapplicable_dynamic_output_shape(%arg0: tensor<4xf32>, %arg1: tensor<2xi64>) -> tensor<1x4xf32> {
-   // CHECK: stablehlo.dynamic_reshape
-diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_legalize_deprecated_ops.mlir b/stablehlo/stablehlo/tests/transforms/stablehlo_legalize_deprecated_ops.mlir
---- stablehlo/stablehlo/tests/transforms/stablehlo_legalize_deprecated_ops.mlir
-+++ stablehlo/stablehlo/tests/transforms/stablehlo_legalize_deprecated_ops.mlir
-@@ -10,6 +10,16 @@
- 
- // -----
- 
-+// CHECK-LABEL: op_dot_discardable_attrs
-+func.func @op_dot_discardable_attrs(%arg0: tensor<2x3xf32>,
-+                 %arg1: tensor<3x?xf32>) -> tensor<2x?xf32> {
-+  // CHECK: stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<2x3xf32>, tensor<3x?xf32>) -> tensor<2x?xf32>
-+  %0 = "stablehlo.dot"(%arg0, %arg1) {mhlo.frontend_attributes = {MUST_FUSE = "0"}} : (tensor<2x3xf32>, tensor<3x?xf32>) -> tensor<2x?xf32>
-+  func.return %0 : tensor<2x?xf32>
-+}
-+
-+// -----
-+
- // CHECK-LABEL: op_unary_einsum
- func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
-   // CHECK:      %cst = stablehlo.constant dense<1.000000e+00> : tensor<f32>
-diff --ruN a/stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir b/stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir
---- stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir
-+++ stablehlo/stablehlo/tests/transforms/stablehlo_refine_shapes.mlir
-@@ -907,19 +907,14 @@
-   func.func public @main() -> tensor<?xi32> {
-     // CHECK: %[[SHAPE:.*]] = stablehlo.constant dense<4> : tensor<1xi64>
-     // CHECK-NEXT: %[[CALL_RESULT:.*]] = call @callee()
--    // CHECK-SAME: {arg_attrs = [], no_inline, probe.marker = "keep",
--    // CHECK-SAME: res_attrs = [{probe.result = "keep"}]}
-+    // CHECK-SAME: <arg_attrs = [], res_attrs = [{probe.result = "keep"}], no_inline>
-+    // CHECK-SAME: {probe.marker = "keep"}
-     // CHECK-SAME: : () -> tensor<1xi64>
-     // CHECK-NEXT: %[[RESULT:.*]] = stablehlo.dynamic_iota %[[SHAPE]], dim = 0
-     // CHECK-SAME: -> tensor<4xi32>
-     // CHECK: return %[[RESULT]] : tensor<4xi32>
-     %0 = stablehlo.constant dense<4> : tensor<i64>
--    %1 = call @callee(%0) {
--      arg_attrs = [{probe.argument = "drop"}],
--      no_inline,
--      probe.marker = "keep",
--      res_attrs = [{probe.result = "keep"}]
--    } : (tensor<i64>) -> tensor<1xi64>
-+    %1 = call @callee(%0) <arg_attrs = [{probe.argument = "drop"}], no_inline, res_attrs = [{probe.result = "keep"}]> {probe.marker = "keep"} : (tensor<i64>) -> tensor<1xi64>
-     %2 = stablehlo.dynamic_iota %1, dim = 0 : (tensor<1xi64>) -> tensor<?xi32>
-     return %2 : tensor<?xi32>
-   }
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_21_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_21_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_21_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_21_0.mlir
-@@ -0,0 +1,3425 @@
-+// RUN: stablehlo-opt --mlir-print-op-generic %s.bc | FileCheck %s
-+// RUN: stablehlo-translate --deserialize %s.bc | stablehlo-translate --serialize --target=1.21.0 | stablehlo-opt --mlir-print-op-generic | FileCheck %s
-+// RUN: stablehlo-translate --deserialize %s.bc | stablehlo-opt > %t.0
-+// RUN: stablehlo-opt --strip-debuginfo %s > %t.1
-+// RUN: diff %t.0 %t.1
-+// RUN: stablehlo-translate --serialize --target=1.21.0 --strip-debuginfo %s > %t.2
-+// RUN: diff %s.bc %t.2
-+// RUN: %if asserts %{ stablehlo-opt --stablehlo-legalize-to-vhlo -emit-bytecode -debug-only=vhlo-bytecode %s 2>&1 | FileCheck --check-prefix=CHECK-WARN %s %}
-+// RUN: %if asserts %{ stablehlo-opt --stablehlo-legalize-to-vhlo -emit-bytecode %s | stablehlo-opt -debug-only=vhlo-bytecode 2>&1 | FileCheck --check-prefix=CHECK-WARN %s %}
-+
-+// CHECK-WARN-NOT: Not Implemented
-+
-+func.func private @mesh()
-+
-+// ============ ATTRIBUTES ============
-+
-+// CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-+    comparison_direction = #stablehlo<comparison_direction EQ>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-+    comparison_direction = #stablehlo<comparison_direction NE>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-+    comparison_direction = #stablehlo<comparison_direction GE>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-+    comparison_direction = #stablehlo<comparison_direction GT>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-+    comparison_direction = #stablehlo<comparison_direction LE>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-+    comparison_direction = #stablehlo<comparison_direction LT>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>
-+    // CHECK: compare_type = #vhlo<comparison_type_v1 NOTYPE>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>,
-+    // CHECK: compare_type = #vhlo<comparison_type_v1 FLOAT>,
-+    compare_type = #stablehlo<comparison_type FLOAT>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>,
-+    // CHECK: compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-+    compare_type = #stablehlo<comparison_type TOTALORDER>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>,
-+    // CHECK: compare_type = #vhlo<comparison_type_v1 SIGNED>,
-+    compare_type = #stablehlo<comparison_type SIGNED>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>,
-+    // CHECK: compare_type = #vhlo<comparison_type_v1 UNSIGNED>,
-+    compare_type = #stablehlo<comparison_type UNSIGNED>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// ConvDimensionNumbers aka #stablehlo.conv is covered below.
-+
-+// CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_UNSPECIFIED>
-+    api_version = 0 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>
-+    api_version = 1 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>
-+    api_version = 2 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING_UNIFIED>
-+    api_version = 3 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_dict"
-+// CHECK: #vhlo.dict_v1<{#vhlo.string_v1<"attr1"> = #vhlo.integer_v1<1 : i32>, #vhlo.string_v1<"attr2"> = #vhlo.integer_v1<2 : i32>}
-+func.func @attr_dict() attributes {stablehlo.attr = {attr1 = 1 : i32, attr2 = 2 : i32}} {
-+  return
-+}
-+
-+// CHECK-LABEL: "attr_custom_call_api_version_typed_ffi"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+// CHECK: api_version = #vhlo<api_version_v1 API_VERSION_TYPED_FFI>
-+// CHECK-SAME: backend_config = #vhlo.dict_v1<{#vhlo.string_v1<"bar"> = #vhlo.integer_v1<42 : i32>}>
-+func.func @attr_custom_call_api_version_typed_ffi(%arg0: tensor<f32>) -> tensor<f32> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    backend_config= {bar = 42 : i32},
-+    api_version = 4 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+
-+// CHECK-LABEL: "attr_custom_call_api_version_typed_ffi_no_backend_config"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+// CHECK: api_version = #vhlo<api_version_v1 API_VERSION_TYPED_FFI>
-+// CHECK-SAME: backend_config = #vhlo.dict_v1<{}>
-+func.func @attr_custom_call_api_version_typed_ffi_no_backend_config(%arg0: tensor<f32>) -> tensor<f32> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    api_version = 4 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// DotDimensionNumbers aka #stablehlo.dot is covered below.
-+
-+// CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-+  %0 = "stablehlo.fft"(%arg0) {
-+    // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-+    fft_type = #stablehlo<fft_type FFT>,
-+    fft_length = array<i64: 16>
-+  } : (tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>>
-+  func.return %0 : tensor<16xcomplex<f32>>
-+}
-+
-+// CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-+  %0 = "stablehlo.fft"(%arg0) {
-+    // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-+    fft_type = #stablehlo<fft_type IFFT>,
-+    fft_length = array<i64: 16>
-+  } : (tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>>
-+  func.return %0 : tensor<16xcomplex<f32>>
-+}
-+
-+// CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-+  %0 = "stablehlo.fft"(%arg0) {
-+    // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-+    fft_type = #stablehlo<fft_type RFFT>,
-+    fft_length = array<i64: 16>
-+  } : (tensor<16xf32>) -> tensor<9xcomplex<f32>>
-+  func.return %0 : tensor<9xcomplex<f32>>
-+}
-+
-+// CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-+  %0 = "stablehlo.fft"(%arg0) {
-+    // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-+    fft_type = #stablehlo<fft_type IRFFT>,
-+    fft_length = array<i64: 16>
-+  } : (tensor<9xcomplex<f32>>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "exponential_HIGHEST"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}
-+func.func @exponential_HIGHEST(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
-+  %0 = "stablehlo.exponential"(%arg0) {
-+    // CHECK: result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 HIGHEST>>
-+    result_accuracy = #stablehlo.result_accuracy<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #stablehlo.result_accuracy_mode<HIGHEST>>
-+  } : (tensor<8x16xf32>) -> tensor<8x16xf32>
-+  func.return %0 : tensor<8x16xf32>
-+}
-+
-+// CHECK-LABEL: "exponential_TOLERANCE"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}
-+func.func @exponential_TOLERANCE(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
-+  %0 = "stablehlo.exponential"(%arg0) {
-+    // CHECK: result_accuracy = #vhlo.result_accuracy_v1<atol = 1.000000e-05, rtol = 0.000000e+00, ulps = 1, mode = #vhlo<result_accuracy_mode_v1 TOLERANCE>>
-+    result_accuracy = #stablehlo.result_accuracy<atol = 1.000000e-5, rtol = 0.000000e+00, ulps = 1, mode = #stablehlo.result_accuracy_mode<TOLERANCE>>
-+  } : (tensor<8x16xf32>) -> tensor<8x16xf32>
-+  func.return %0 : tensor<8x16xf32>
-+}
-+
-+// CHECK-LABEL: "exponential_DEFAULT"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}
-+func.func @exponential_DEFAULT(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
-+  %0 = "stablehlo.exponential"(%arg0) {
-+    // CHECK: result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>
-+    result_accuracy = #stablehlo.result_accuracy<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #stablehlo.result_accuracy_mode<DEFAULT>>
-+  } : (tensor<8x16xf32>) -> tensor<8x16xf32>
-+  func.return %0 : tensor<8x16xf32>
-+}
-+
-+// GatherDimensionNumbers aka #stablehlo.gather is covered below.
-+
-+// CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-+  %0 = "stablehlo.dot"(%arg0, %arg1) {
-+    // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-+  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-+  func.return %0 : tensor<8x8xf32>
-+}
-+
-+// CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-+  %0 = "stablehlo.dot"(%arg0, %arg1) {
-+    // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-+    precision_config = [#stablehlo<precision HIGH>, #stablehlo<precision HIGH>]
-+  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-+  func.return %0 : tensor<8x8xf32>
-+}
-+
-+// CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-+  %0 = "stablehlo.dot"(%arg0, %arg1) {
-+    // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-+    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
-+  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-+  func.return %0 : tensor<8x8xf32>
-+}
-+
-+// CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-+  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-+    // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-+    rng_algorithm = #stablehlo<rng_algorithm DEFAULT>
-+  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
-+  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-+  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-+    // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-+    rng_algorithm = #stablehlo<rng_algorithm THREE_FRY>
-+  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
-+  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-+  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-+    // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-+    rng_algorithm = #stablehlo<rng_algorithm PHILOX>
-+  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
-+  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-+  %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-+    // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-+    rng_distribution = #stablehlo<rng_distribution UNIFORM>
-+  } : (tensor<f32>, tensor<f32>, tensor<0xindex>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-+  %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-+    // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-+    rng_distribution = #stablehlo<rng_distribution NORMAL>
-+  } : (tensor<f32>, tensor<f32>, tensor<0xindex>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
-+
-+// CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-+  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-+    left_side = true,
-+    lower = true,
-+    unit_diagonal = true,
-+    // transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-+    transpose_a = #stablehlo<transpose NO_TRANSPOSE>
-+  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-+  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-+    left_side = true,
-+    lower = true,
-+    unit_diagonal = true,
-+    // transpose_a = #vhlo<transpose_v1 TRANSPOSE>,
-+    transpose_a = #stablehlo<transpose TRANSPOSE>
-+  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-+  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-+    left_side = true,
-+    lower = true,
-+    unit_diagonal = true,
-+    // transpose_a = #vhlo<transpose_v1 ADJOINT>,
-+    transpose_a = #stablehlo<transpose ADJOINT>
-+  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
-+
-+// CHECK-LABEL: "attr_type_extensions_bounds"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+  func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
-+}
-+
-+// CHECK-LABEL: "attr_frontend_attributes"
-+func.func @attr_frontend_attributes(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: some.unregistered_attr
-+  %1 = stablehlo.cosine %arg0 {some.unregistered_attr = 1 : i32} : tensor<f32>
-+  return %1 : tensor<f32>
-+}
-+
-+// Builtin attriubute tests
-+
-+// CHECK-LABEL: "byte_packed_boolean"
-+func.func @byte_packed_boolean() -> (tensor<8xi1>, tensor<8xi1>, tensor<4xi1>, tensor<4xi1>, tensor<16xi1>) {
-+  // CHECK: #vhlo.tensor_v1<dense<[true, false, false, false, false, false, false, false]
-+  // CHECK-NEXT: #vhlo.tensor_v1<dense<true> : tensor<8xi1>>
-+  // CHECK-NEXT: #vhlo.tensor_v1<dense<[true, false, false, false]> : tensor<4xi1>>
-+  // CHECK-NEXT: #vhlo.tensor_v1<dense<true> : tensor<4xi1>>
-+  // CHECK-NEXT: #vhlo.tensor_v1<dense<[true, false, false, false, false, false, false, false, true, false, false, false, false, false, false, false]> : tensor<16xi1>>
-+  %c = stablehlo.constant dense<[true, false, false, false, false, false, false, false]> : tensor<8xi1>
-+  %c_0 = stablehlo.constant dense<true> : tensor<8xi1>
-+  %c_1 = stablehlo.constant dense<[true, false, false, false]> : tensor<4xi1>
-+  %c_2 = stablehlo.constant dense<true> : tensor<4xi1>
-+  %c_3 = stablehlo.constant dense<[true, false, false, false, false, false, false, false, true, false, false, false, false, false, false, false]> : tensor<16xi1>
-+  return %c, %c_0, %c_1, %c_2, %c_3 : tensor<8xi1>, tensor<8xi1>, tensor<4xi1>, tensor<4xi1>, tensor<16xi1>
-+}
-+
-+// ============ DEFAULTS ============
-+
-+// CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
-+  //               CHECK: "vhlo.all_gather_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.all_gather"(%arg0) {
-+    all_gather_dim = 1 : i64,
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
-+  } : (tensor<16x8xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "default_all_gather_variadic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_all_gather_variadic(%arg0: tensor<16x8xf32>, %arg1: tensor<16x8xf32>) -> (tensor<16x16xf32>, tensor<16x16xf32>) {
-+  %0:2 = "stablehlo.all_gather"(%arg0, %arg1) {
-+    all_gather_dim = 1 : i64,
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
-+  } : (tensor<16x8xf32>, tensor<16x8xf32>) -> (tensor<16x16xf32>, tensor<16x16xf32>)
-+  func.return %0#0, %0#1 : tensor<16x16xf32>, tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.all_reduce_v2"(%[[ARG0]])
-+  //          CHECK-SAME: <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+
-+  %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 = dense<[[0], [1]]> : tensor<2x1xi64>
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_replica_group_mesh_axes"
-+func.func @attr_replica_group_mesh_axes(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.all_reduce_v2"
-+  // CHECK-SAME: replica_groups = #vhlo.replica_group_mesh_axes_v1<
-+  // CHECK-SAME:   mesh = #vhlo.string_v1<"mesh">,
-+  // CHECK-SAME:   axes = #vhlo.array_v1<[#vhlo.axis_ref_v1<name = #vhlo.string_v1<"x">, sub_axis_info = #vhlo.sub_axis_info_v1<pre_size = 1, size = 2>>, #vhlo.axis_ref_v1<name = #vhlo.string_v1<"y">>]>
-+  // CHECK-SAME: >
-+  %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 = @mesh,
-+      axes = [
-+        #stablehlo.axis_ref<name = "x", sub_axis_info = (1)2>,
-+        #stablehlo.axis_ref<name = "y">
-+      ]
-+    >
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+
-+
-+// CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
-+  //               CHECK: "vhlo.all_to_all_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-+  //          CHECK-SAME:   split_count = #vhlo.integer_v1<4 : i64>
-+  //          CHECK-SAME:   split_dimension = #vhlo.integer_v1<1 : i64>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x4x!vhlo.f32_v1>
-+  %0 = "stablehlo.all_to_all"(%arg0) {
-+    split_dimension = 1 : i64,
-+    concat_dimension = 0 : i64,
-+    split_count = 4 : i64,
-+    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>
-+  } : (tensor<4x16xf32>) -> tensor<16x4xf32>
-+  func.return %0 : tensor<16x4xf32>
-+}
-+
-+// CHECK-LABEL: "default_all_to_all_variadic"
-+func.func @default_all_to_all_variadic(%arg0: tensor<4x16xf32>, %arg1: tensor<5x16xf32>) -> (tensor<16x4xf32>, tensor<20x4xf32>) {
-+  %0:2 = "stablehlo.all_to_all"(%arg0, %arg1) {
-+    split_dimension = 1 : i64,
-+    concat_dimension = 0 : i64,
-+    split_count = 4 : i64,
-+    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>
-+  } : (tensor<4x16xf32>, tensor<5x16xf32>) -> (tensor<16x4xf32>, tensor<20x4xf32>)
-+  func.return %0#0, %0#1 : tensor<16x4xf32>, tensor<20x4xf32>
-+}
-+
-+// CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-+  func.return %0 : tensor<1x16x16xf32>
-+}
-+
-+// CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.collective_permute"(%arg0) {
-+    source_target_pairs = dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>
-+  } : (tensor<16x8xf32>) -> tensor<16x8xf32>
-+  func.return %0 : tensor<16x8xf32>
-+}
-+
-+// CHECK-LABEL: "default_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_collective_broadcast(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
-+  //               CHECK: "vhlo.collective_broadcast_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<false>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.collective_broadcast"(%arg0) {
-+    replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
-+  } : (tensor<16x8xf32>) -> tensor<16x8xf32>
-+  func.return %0 : tensor<16x8xf32>
-+}
-+
-+// CHECK-LABEL: "default_collective_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_collective_reduce(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.collective_reduce_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<false>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>,
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.collective_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 = dense<[[0, 1]]> : tensor<1x2xi64>
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-+  // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "default_composite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_composite(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.composite_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{}>
-+  //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
-+  //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
-+  //          CHECK-SAME:   version = #vhlo.integer_v1<0 : i64>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.composite"(%arg0) {
-+    name = "stablehlo.composite_target",
-+    decomposition = @composite_target
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
-+  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<2x2xi64>>,
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
-+  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<false> : tensor<2xi1>>,
-+  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x6x6x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.convolution"(%arg0, %arg1) {
-+    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
-+    feature_group_count = 1 : i64,
-+    batch_group_count = 1 : i64
-+  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32>
-+  func.return %0 : tensor<1x6x6x16xf32>
-+}
-+
-+// CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
-+  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-+  // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-+  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-+  // CHECK-SAME:   called_computations = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<false>,
-+  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[]>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo"
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
-+  //      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   allow_imprecise_accumulation = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   lhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   num_primitive_operations = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+  // CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   rhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
-+  // CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
-+    dot_dimension_numbers = #stablehlo.dot<
-+      lhs_batching_dimensions = [0],
-+      lhs_contracting_dimensions = [2],
-+      rhs_batching_dimensions = [0],
-+      rhs_contracting_dimensions = [1]
-+    >
-+  } : (tensor<8x8x16xf32>, tensor<8x16x8xf32>) -> tensor<8x8x8xf32>
-+  func.return %0 : tensor<8x8x8xf32>
-+}
-+
-+// CHECK-LABEL: "dot_general_algorithm"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @dot_general_algorithm(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
-+//      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
-+// CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.f32_v1>,
-+// CHECK-SAME:   allow_imprecise_accumulation = #vhlo.bool_v1<false>,
-+// CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.tf31_v1>,
-+// CHECK-SAME:   num_primitive_operations = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+// CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.tf31_v1>
-+// CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
-+    dot_dimension_numbers = #stablehlo.dot<
-+      lhs_batching_dimensions = [0],
-+      lhs_contracting_dimensions = [2],
-+      rhs_batching_dimensions = [0],
-+      rhs_contracting_dimensions = [1]
-+    >,
-+    algorithm = #stablehlo.dot_algorithm<
-+      lhs_precision_type = tf32,
-+      rhs_precision_type = tf32,
-+      accumulation_type = f32,
-+      lhs_component_count = 1,
-+      rhs_component_count = 1,
-+      num_primitive_operations = 1,
-+      allow_imprecise_accumulation = false
-+    >
-+  } : (tensor<8x8x16xf32>, tensor<8x16x8xf32>) -> tensor<8x8x8xf32>
-+  func.return %0 : tensor<8x8x8xf32>
-+}
-+
-+// CHECK-LABEL: "dot_general_algorithm_f8e4m3fn_x3"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @dot_general_algorithm_f8e4m3fn_x3(%arg0: tensor<8x8x16xbf16>, %arg1: tensor<8x16x8xbf16>) -> tensor<8x8x8xf32> {
-+//      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
-+// CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.f32_v1>,
-+// CHECK-SAME:   allow_imprecise_accumulation = #vhlo.bool_v1<false>,
-+// CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>,
-+// CHECK-SAME:   num_primitive_operations = #vhlo.integer_v1<3 : i64>,
-+// CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+// CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>
-+// CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.bf16_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.bf16_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
-+    dot_dimension_numbers = #stablehlo.dot<
-+      lhs_batching_dimensions = [0],
-+      lhs_contracting_dimensions = [2],
-+      rhs_batching_dimensions = [0],
-+      rhs_contracting_dimensions = [1]
-+    >,
-+    algorithm = #stablehlo.dot_algorithm<
-+      lhs_precision_type = f8E4M3FN,
-+      rhs_precision_type = f8E4M3FN,
-+      accumulation_type = f32,
-+      lhs_component_count = 1,
-+      rhs_component_count = 1,
-+      num_primitive_operations = 3,
-+      allow_imprecise_accumulation = false
-+    >
-+  } : (tensor<8x8x16xbf16>, tensor<8x16x8xbf16>) -> tensor<8x8x8xf32>
-+  func.return %0 : tensor<8x8x8xf32>
-+}
-+
-+// CHECK-LABEL: "dot_general_algorithm_f8e4m3fn_x4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @dot_general_algorithm_f8e4m3fn_x4(%arg0: tensor<8x8x16xbf16>, %arg1: tensor<8x16x8xbf16>) -> tensor<8x8x8xf32> {
-+//      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
-+// CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.f32_v1>,
-+// CHECK-SAME:   allow_imprecise_accumulation = #vhlo.bool_v1<false>,
-+// CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>,
-+// CHECK-SAME:   num_primitive_operations = #vhlo.integer_v1<4 : i64>,
-+// CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+// CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>
-+// CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.bf16_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.bf16_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
-+    dot_dimension_numbers = #stablehlo.dot<
-+      lhs_batching_dimensions = [0],
-+      lhs_contracting_dimensions = [2],
-+      rhs_batching_dimensions = [0],
-+      rhs_contracting_dimensions = [1]
-+    >,
-+    algorithm = #stablehlo.dot_algorithm<
-+      lhs_precision_type = f8E4M3FN,
-+      rhs_precision_type = f8E4M3FN,
-+      accumulation_type = f32,
-+      lhs_component_count = 1,
-+      rhs_component_count = 1,
-+      num_primitive_operations = 4,
-+      allow_imprecise_accumulation = false
-+    >
-+  } : (tensor<8x8x16xbf16>, tensor<8x16x8xbf16>) -> tensor<8x8x8xf32>
-+  func.return %0 : tensor<8x8x8xf32>
-+}
-+
-+// CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_broadcast_in_dim"(%arg0, %arg1) {
-+    broadcast_dimensions = array<i64: 0, 1>
-+  } : (tensor<?x?xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-+  func.return %0 : tensor<?x?xf32>
-+}
-+
-+// CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<2x2xi64>) -> tensor<1x?x?x16xf32> {
-+  //      CHECK: "vhlo.dynamic_conv_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
-+  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
-+  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<false> : tensor<2xi1>>,
-+  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x2x!vhlo.i64_v1>) -> !vhlo.tensor_v1<1x?x?x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_conv"(%arg0, %arg1, %arg2) {
-+    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
-+    feature_group_count = 1 : i64,
-+    batch_group_count = 1 : i64
-+  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>, tensor<2x2xi64>) -> tensor<1x?x?x16xf32>
-+  func.return %0 : tensor<1x?x?x16xf32>
-+}
-+
-+// CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
-+  //      CHECK: "vhlo.dynamic_gather_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-+  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>, !vhlo.tensor_v1<3x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_gather"(%arg0, %arg1, %arg2) {
-+    dimension_numbers = #stablehlo.gather<
-+      offset_dims = [2],
-+      collapsed_slice_dims = [0, 1],
-+      start_index_map = [0, 1],
-+      index_vector_dim = 2
-+    >
-+  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<1x5x8xf32>
-+  func.return %0 : tensor<1x5x8xf32>
-+}
-+
-+func.func @default_func(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK:      "vhlo.func_v1"() <{
-+  // CHECK-SAME:   arg_attrs = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>>>,
-+  // CHECK-SAME:   res_attrs = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-+  // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : () -> ()
-+  func.return %arg0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "default_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
-+  //      CHECK: "vhlo.gather_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-+  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<1> : tensor<3xi64>>,
-+  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x1x!vhlo.f32_v1>
-+  %0 = "stablehlo.gather"(%arg0, %arg1) {
-+    dimension_numbers = #stablehlo.gather<
-+      offset_dims = [2],
-+      collapsed_slice_dims = [0, 1],
-+      start_index_map = [0, 1],
-+      index_vector_dim = 2
-+    >,
-+    slice_sizes = array<i64: 1, 1, 1>
-+  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>) -> tensor<1x5x1xf32>
-+  func.return %0 : tensor<1x5x1xf32>
-+}
-+
-+// CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-+  // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-+  //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-+  %0:2 = "stablehlo.infeed"(%arg0) : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
-+  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-+  %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-+  func.return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_recv_with_source_target_pairs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_recv_with_source_target_pairs(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
-+  //      CHECK: "vhlo.recv_v2"(%[[ARG0]]) <{
-+  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>,
-+  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>>
-+  // CHECK-SAME{LITERAL}: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-+  %0:2 = "stablehlo.recv"(%arg0) {
-+    channel_handle = #stablehlo.channel_handle<handle = 0, type = 1>,
-+    source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>
-+  } : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
-+  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
-+  //      CHECK: "vhlo.send_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>,
-+  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>>
-+  // CHECK-SAME{LITERAL}: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-+  %0 = "stablehlo.send"(%arg0, %arg1) {
-+    channel_handle = #stablehlo.channel_handle<handle = 0, type = 1>,
-+    source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>
-+  } : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-+  func.return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.reduce_scatter"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg1, %arg2) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    scatter_dimension = 0 : i64,
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
-+  } : (tensor<16xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-+  //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-+  // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-+  //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-+  //          CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-+  //          CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.maximum_v1"(%[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<2x17x31x7x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<2x16x30x7x!vhlo.f32_v1>
-+  %0 = "stablehlo.reduce_window"(%arg0, %arg1) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %1 = "stablehlo.maximum"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    window_dimensions = array<i64: 1, 2, 2, 1>
-+  } : (tensor<2x17x31x7xf32>, tensor<f32>) -> tensor<2x16x30x7xf32>
-+  func.return %0 : tensor<2x16x30x7xf32>
-+}
-+
-+// CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
-+  //      CHECK: "vhlo.scatter_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-+  // CHECK-SAME:   input_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   scatter_dims_to_operand_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   scatter_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   unique_indices = #vhlo.bool_v1<false>,
-+  // CHECK-SAME:   update_window_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>
-+  %0 = "stablehlo.scatter"(%arg0, %arg1, %arg2) ({
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    scatter_dimension_numbers = #stablehlo.scatter<
-+      update_window_dims = [1],
-+      inserted_window_dims = [0, 1],
-+      scatter_dims_to_operand_dims = [0, 1],
-+      index_vector_dim = 1
-+    >
-+  } : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) -> tensor<200x100x300xf32>
-+  func.return %0 : tensor<200x100x300xf32>
-+}
-+
-+// CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-+  // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-+  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG31:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG41:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL11:.*]] = "vhlo.compare_v1"(%[[ARG31]], %[[ARG41]]) <{compare_type = #vhlo<comparison_type_v1 TOTALORDER>, comparison_direction = #vhlo<comparison_direction_v1 GE>}>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL11]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-+  // CHECK-NEXT: }, {
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG32:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG42:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL12:.*]] = "vhlo.add_v1"(%[[ARG32]], %[[ARG42]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL12]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>, !vhlo.tensor_v1<10x23x23x64x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>
-+  %0 = "stablehlo.select_and_scatter"(%arg0, %arg1, %arg2) ({
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.compare"(%arg3, %arg4) {compare_type = #stablehlo<comparison_type TOTALORDER>, comparison_direction = #stablehlo<comparison_direction GE>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+      "stablehlo.return"(%1) : (tensor<i1>) -> ()
-+  }, {
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    window_dimensions = array<i64: 1, 2, 2, 1>
-+  } : (tensor<10x24x24x64xf32>, tensor<10x23x23x64xf32>, tensor<f32>) -> tensor<10x24x24x64xf32>
-+  func.return %0 : tensor<10x24x24x64xf32>
-+}
-+
-+// CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-+  // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.compare_v1"(%[[ARG1]], %[[ARG2]]) <{compare_type = #vhlo<comparison_type_v1 FLOAT>, comparison_direction = #vhlo<comparison_direction_v1 GT>}>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.sort"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
-+      %1 = "stablehlo.compare"(%arg1, %arg2) {compare_type = #stablehlo<comparison_type FLOAT>, comparison_direction = #stablehlo<comparison_direction GT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+      "stablehlo.return"(%1) : (tensor<i1>) -> ()
-+  }) : (tensor<16xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// ============ OPS ============
-+
-+// CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-+  func.return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
-+  //               CHECK: "vhlo.all_gather_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.all_gather"(%arg0) {
-+    all_gather_dim = 1 : i64,
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+    use_global_device_ids
-+  } : (tensor<16x8xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.all_reduce_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %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 = dense<[[0], [1]]> : tensor<2x1xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+    use_global_device_ids
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_all_reduce_with_promotable_types"
-+func.func @op_all_reduce_with_promotable_types(%operand: tensor<f32>) -> tensor<f64> {
-+  //  CHECK: "vhlo.all_reduce_v2"(%[[ARG0:.*]])
-+  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-+  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-+  //  CHECK: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  %result = "stablehlo.all_reduce"(%operand) ({
-+    ^bb0(%arg0: tensor<f64>, %arg1: tensor<f64>):
-+      %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-+      "stablehlo.return"(%0) : (tensor<f64>) -> ()
-+  }) {
-+    replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+    use_global_device_ids
-+  } : (tensor<f32>) -> tensor<f64>
-+
-+  func.return %result : tensor<f64>
-+}
-+
-+// CHECK-LABEL: "default_all_reduce_variadic"
-+func.func @default_all_reduce_variadic(%arg0: tensor<f32>, %arg1: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-+  %0:2 = "stablehlo.all_reduce"(%arg0, %arg1) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> (tensor<f32>)
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
-+  } : (tensor<f32>, tensor<f32>) -> (tensor<f32>, tensor<f32>)
-+  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
-+  //               CHECK: "vhlo.all_to_all_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-+  //          CHECK-SAME:   split_count = #vhlo.integer_v1<4 : i64>
-+  //          CHECK-SAME:   split_dimension = #vhlo.integer_v1<1 : i64>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x4x!vhlo.f32_v1>
-+  %0 = "stablehlo.all_to_all"(%arg0) {
-+    split_dimension = 1 : i64,
-+    concat_dimension = 0 : i64,
-+    split_count = 4 : i64,
-+    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>
-+  } : (tensor<4x16xf32>) -> tensor<16x4xf32>
-+  func.return %0 : tensor<16x4xf32>
-+}
-+
-+// CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
-+func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-+  // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-+  // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-+  %0:3 = "stablehlo.batch_norm_grad"(%arg0, %arg1, %arg2, %arg3, %arg4) {
-+    epsilon = 0.001 : f32,
-+    feature_index = 0 : i64
-+  } : (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>)
-+  func.return %0#0, %0#1, %0#2 : tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
-+func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-+  // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-+  // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.batch_norm_inference"(%arg0, %arg1, %arg2, %arg3, %arg4) {
-+    epsilon = 0.001 : f32,
-+    feature_index = 0 : i64
-+  } : (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>) -> tensor<16x16x16x16xf32>
-+  func.return %0 : tensor<16x16x16x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-+  // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-+  %0:3 = "stablehlo.batch_norm_training"(%arg0, %arg1, %arg2) {
-+    epsilon = 0.001 : f32,
-+    feature_index = 0 : i64
-+  } : (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>)
-+  func.return %0#0, %0#1, %0#2 : tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-+    broadcast_dimensions = array<i64: 1>
-+  } : (tensor<16xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.broadcast"(%arg0) {
-+    broadcast_sizes = array<i64: 16>
-+  } : (tensor<16xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.case"(%arg0) ({
-+    "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-+  }) : (tensor<i32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.cbrt_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.cholesky"(%arg0) {
-+    lower = true
-+  } : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-+  func.return %0 : tensor<1x16x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "op_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_collective_broadcast(%arg0: tensor<16x8xf32>, %arg1: tensor<1xi32>) -> tensor<16x8xf32> {
-+  //               CHECK: "vhlo.collective_broadcast_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<true>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.collective_broadcast"(%arg0, %arg1) {
-+    replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+    has_dynamic_root
-+  } : (tensor<16x8xf32>, tensor<1xi32>) -> tensor<16x8xf32>
-+  func.return %0 : tensor<16x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_collective_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_collective_reduce(%arg0: tensor<f32>, %arg1: tensor<1xi32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.collective_reduce_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<true>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.collective_reduce"(%arg0, %arg1) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+    use_global_device_ids,
-+    has_dynamic_root
-+  } : (tensor<f32>, tensor<1xi32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.collective_permute"(%arg0) {
-+    source_target_pairs = dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>
-+  } : (tensor<16x8xf32>) -> tensor<16x8xf32>
-+  func.return %0 : tensor<16x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-+  // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>,
-+    compare_type = #stablehlo<comparison_type TOTALORDER>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-+  func.return %0 : tensor<complex<f32>>
-+}
-+
-+// CHECK-LABEL: "op_composite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_composite(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.composite_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{#vhlo.string_v1<"my_int"> = #vhlo.integer_v1<1 : i64>, #vhlo.string_v1<"my_string"> = #vhlo.string_v1<"foo">}>
-+  //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
-+  //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
-+  //          CHECK-SAME:   version = #vhlo.integer_v1<1 : i32>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.composite"(%arg0) {
-+    name = "stablehlo.composite_target",
-+    decomposition = @composite_target,
-+    version = 1 : i32,
-+    composite_attributes = {
-+      my_int = 1 : i64,
-+      my_string = "foo"
-+    }
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "composite_regions"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @composite_regions(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.composite_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{}>
-+  //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
-+  //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
-+  //          CHECK-SAME:   version = #vhlo.integer_v1<1 : i32>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT: ^bb0(%[[ARG1:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.composite"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>):
-+      "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-+  }) {
-+    name = "stablehlo.composite_target",
-+    decomposition = @composite_target,
-+    version = 1 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-+    dimension = 0 : i64
-+  } : (tensor<8xf32>, tensor<8xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.constant_v1"() <{
-+  // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-+  // CHECK-SAME: }> : () -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.constant"() {
-+    value = dense<0.0> : tensor<f32>
-+  } : () -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
-+  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<2x2xi64>>,
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-+  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
-+  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<true> : tensor<2xi1>>,
-+  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x7x7x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.convolution"(%arg0, %arg1) {
-+    window_strides = array<i64: 2, 2>,
-+    padding = dense<1> : tensor<2x2xi64>,
-+    lhs_dilation = array<i64: 2, 2>,
-+    rhs_dilation = array<i64: 2, 2>,
-+    window_reversal = array<i1: true, true>,
-+    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
-+    feature_group_count = 1 : i64,
-+    batch_group_count = 1 : i64,
-+    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
-+  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32>
-+  func.return %0 : tensor<1x7x7x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.cosine_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_create_token"
-+func.func @op_create_token() -> !stablehlo.token {
-+  // CHECK: "vhlo.create_token_v1"() : () -> !vhlo.token_v1
-+  %0 = "stablehlo.create_token"() : () -> !stablehlo.token
-+  func.return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.cross-replica-sum"(%arg0) {
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
-+  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-+  // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-+  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-+  // CHECK-SAME:   called_computations = #vhlo.array_v1<[#vhlo.string_v1<"foo">]>,
-+  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<> : tensor<0xindex>>]>,
-+  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[
-+  // CHECK-SAME:     #vhlo.output_operand_alias_v1<
-+  // CHECK-SAME:       outputTupleIndices = [],
-+  // CHECK-SAME:       operandIndex = 0,
-+  // CHECK-SAME:       operandTupleIndices = []>]>,
-+  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<> : tensor<0xindex>>]>,
-+  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[]>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    has_side_effect = true,
-+    backend_config = "\08\03\1A\02",
-+    api_version = 2 : i32,
-+    called_computations = [@foo],
-+    operand_layouts = [dense<> : tensor<0xindex>],
-+    output_operand_aliases = [
-+      #stablehlo.output_operand_alias<output_tuple_indices = [],
-+                                 operand_index = 0,
-+                                 operand_tuple_indices = []>],
-+    result_layouts = [dense<> : tensor<0xindex>]
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_custom_call_empty_result_layout"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func public @op_custom_call_empty_result_layout(%arg0: tensor<i64>) -> tensor<i64> {
-+  // %0 = "vhlo.custom_call_v2"(%arg0) <{>}> : (!vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tuple_v1<>
-+  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
-+  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-+  // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-+  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"empty_output">,
-+  // CHECK-SAME:   called_computations = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<> : tensor<0xindex>>]>,
-+  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[]>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tuple_v1<>
-+  %0 = "stablehlo.custom_call"(%arg0) <{
-+    api_version = 2 : i32,
-+    call_target_name = "empty_output",
-+    has_side_effect = true,
-+    operand_layouts = [dense<> : tensor<0xindex>],
-+    result_layouts = []
-+  }> : (tensor<i64>) -> tuple<>
-+  return %arg0 : tensor<i64>
-+}
-+
-+// CHECK-LABEL: "op_custom_call_with_result_tilings"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_custom_call_with_result_tilings(%arg0: tensor<64x256xf32>) -> tensor<64x256xf32> {
-+  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
-+  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-+  // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-+  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-+  // CHECK-SAME:   called_computations = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<false>,
-+  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<[1, 0]> : tensor<2xindex>>]>,
-+  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<[1, 0]> : tensor<2xindex>>]>,
-+  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[#vhlo.array_v1<[#vhlo.tensor_v1<dense<[8, 128]> : tensor<2xindex>>]>]>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<64x256x!vhlo.f32_v1>) -> !vhlo.tensor_v1<64x256x!vhlo.f32_v1>
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    api_version = 2 : i32,
-+    call_target_name = "foo",
-+    operand_layouts = [dense<[1, 0]> : tensor<2xindex>],
-+    result_layouts = [dense<[1, 0]> : tensor<2xindex>],
-+    result_tilings = [[dense<[8, 128]> : tensor<2xindex>]]
-+  } : (tensor<64x256xf32>) -> tensor<64x256xf32>
-+  func.return %0 : tensor<64x256xf32>
-+}
-+
-+// CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
-+  //      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   allow_imprecise_accumulation = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   lhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   num_primitive_operations = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-+  // CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   rhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
-+  // CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
-+    dot_dimension_numbers = #stablehlo.dot<
-+      lhs_batching_dimensions = [0],
-+      lhs_contracting_dimensions = [2],
-+      rhs_batching_dimensions = [0],
-+      rhs_contracting_dimensions = [1]
-+    >,
-+    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
-+  } : (tensor<8x8x16xf32>, tensor<8x16x8xf32>) -> tensor<8x8x8xf32>
-+  func.return %0 : tensor<8x8x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot"(%arg0, %arg1) {
-+    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
-+  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-+  func.return %0 : tensor<8x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_broadcast_in_dim"(%arg0, %arg1) {
-+    broadcast_dimensions = array<i64: 0, 1>,
-+    known_expanding_dimensions = array<i64: 0>,
-+    known_nonexpanding_dimensions = array<i64: 1>
-+  } : (tensor<?x?xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-+  func.return %0 : tensor<?x?xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<2x2xi64>) -> tensor<1x?x?x16xf32> {
-+  //      CHECK: "vhlo.dynamic_conv_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
-+  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-+  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
-+  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<true> : tensor<2xi1>>,
-+  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x2x!vhlo.i64_v1>) -> !vhlo.tensor_v1<1x?x?x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_conv"(%arg0, %arg1, %arg2) {
-+    window_strides = array<i64: 2, 2>,
-+    lhs_dilation = array<i64: 2, 2>,
-+    rhs_dilation = array<i64: 2, 2>,
-+    window_reversal = array<i1: true, true>,
-+    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
-+    feature_group_count = 1 : i64,
-+    batch_group_count = 1 : i64,
-+    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
-+  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>, tensor<2x2xi64>) -> tensor<1x?x?x16xf32>
-+  func.return %0 : tensor<1x?x?x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
-+  //      CHECK: "vhlo.dynamic_gather_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>, !vhlo.tensor_v1<3x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_gather"(%arg0, %arg1, %arg2) {
-+    dimension_numbers = #stablehlo.gather<
-+      offset_dims = [2],
-+      collapsed_slice_dims = [0, 1],
-+      start_index_map = [0, 1],
-+      index_vector_dim = 2
-+    >,
-+    indices_are_sorted = true
-+  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<1x5x8xf32>
-+  func.return %0 : tensor<1x5x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_gather_with_batching_dims"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_dynamic_gather_with_batching_dims(%arg0 : tensor<5x2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<4xi32>) -> tensor<1x5x8xf32> {
-+  //      CHECK: "vhlo.dynamic_gather_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<5x2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>, !vhlo.tensor_v1<4x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_gather"(%arg0, %arg1, %arg2) {
-+    dimension_numbers = #stablehlo.gather<
-+      offset_dims = [2],
-+      collapsed_slice_dims = [1, 2],
-+      operand_batching_dims = [0],
-+      start_indices_batching_dims = [1],
-+      start_index_map = [1, 2],
-+      index_vector_dim = 2
-+    >,
-+    indices_are_sorted = true
-+  } : (tensor<5x2x4x9xf32>, tensor<1x5x2xi32>, tensor<4xi32>) -> tensor<1x5x8xf32>
-+  func.return %0 : tensor<1x5x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_iota"(%arg0) {
-+    iota_dimension = 0 : i64
-+  } : (tensor<1xindex>) -> tensor<?xf32>
-+  func.return %0 : tensor<?xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
-+func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-+  func.return %0 : tensor<?xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-+  func.return %0 : tensor<?x?xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-+    slice_sizes = array<i64: 4>
-+  } : (tensor<16xf32>, tensor<i64>) -> tensor<4xf32>
-+  func.return %0 : tensor<4xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.einsum"(%arg0, %arg1) {
-+    einsum_config = "ab,bc->ac"
-+  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-+  func.return %0 : tensor<8x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.exponential_minus_one_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.exponential_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-+  // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-+  %0 = "stablehlo.fft"(%arg0) {
-+    fft_type = #stablehlo<fft_type FFT>,
-+    fft_length = array<i64: 16>
-+  } : (tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>>
-+  func.return %0 : tensor<16xcomplex<f32>>
-+}
-+
-+// CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+func.func private @op_func(%arg0: tensor<f32> {stablehlo.arg = "0"}) -> (tensor<f32> {stablehlo.result = "0"}) {
-+  // CHECK:      "vhlo.func_v1"() <{
-+  // CHECK-SAME:   arg_attrs = #vhlo.array_v1<[#vhlo.dict_v1<{#vhlo.string_v1<"stablehlo.arg"> = #vhlo.string_v1<"0">}>]>,
-+  // CHECK-SAME:   function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>>>,
-+  // CHECK-SAME:   res_attrs = #vhlo.array_v1<[#vhlo.dict_v1<{#vhlo.string_v1<"stablehlo.result"> = #vhlo.string_v1<"0">}>]>,
-+  // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-+  // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : () -> ()
-+
-+  func.return %arg0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
-+  //      CHECK: "vhlo.gather_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<1> : tensor<3xi64>>,
-+  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x1x!vhlo.f32_v1>
-+  %0 = "stablehlo.gather"(%arg0, %arg1) {
-+    dimension_numbers = #stablehlo.gather<
-+      offset_dims = [2],
-+      collapsed_slice_dims = [0, 1],
-+      start_index_map = [0, 1],
-+      index_vector_dim = 2
-+    >,
-+    slice_sizes = array<i64: 1, 1, 1>,
-+    indices_are_sorted = true
-+  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>) -> tensor<1x5x1xf32>
-+  func.return %0 : tensor<1x5x1xf32>
-+}
-+
-+// CHECK-LABEL: "op_gather_with_batching_dims"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_gather_with_batching_dims(%arg0 : tensor<5x2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
-+  //      CHECK: "vhlo.gather_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-+  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<5x2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x1x!vhlo.f32_v1>
-+  %0 = "stablehlo.gather"(%arg0, %arg1) {
-+    dimension_numbers = #stablehlo.gather<
-+      offset_dims = [2],
-+      collapsed_slice_dims = [1, 2],
-+      operand_batching_dims = [0],
-+      start_indices_batching_dims = [1],
-+      start_index_map = [1, 2],
-+      index_vector_dim = 2
-+    >,
-+    slice_sizes = array<i64: 1, 1, 1, 1>,
-+    indices_are_sorted = true
-+  } : (tensor<5x2x4x9xf32>, tensor<1x5x2xi32>) -> tensor<1x5x1xf32>
-+  func.return %0 : tensor<1x5x1xf32>
-+}
-+
-+// CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.get_dimension_size"(%arg0) {
-+    dimension = 0 : i64
-+  } : (tensor<?xf32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-+  // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.get_tuple_element"(%arg0) {
-+    index = 0 : i32
-+  } : (tuple<tensor<f32>, tensor<i32>>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }, {
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.if"(%arg0) ({
-+    "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-+  }, {
-+    "stablehlo.return"(%arg2) : (tensor<f32>) -> ()
-+  }) : (tensor<i1>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-+  // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-+  //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-+  %0:2 = "stablehlo.infeed"(%arg0) {
-+    infeed_config = "foo",
-+    layout = [[]]
-+  } : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
-+  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_iota"
-+func.func @op_iota() -> tensor<16xf32> {
-+  //      CHECK: "vhlo.iota_v1"() <{
-+  // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : () -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.iota"() {
-+    iota_dimension = 0 : i64
-+  } : () -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.log_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.log_plus_one_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.logistic_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.abs_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.map"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>):
-+      %1 = "stablehlo.abs"(%arg1) : (tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    dimensions = array<i64: 0>
-+  } : (tensor<16xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-+  %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-+    outfeed_config = "foo"
-+  } : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-+  func.return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-+  // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-+  // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.pad"(%arg0, %arg1) {
-+    edge_padding_high = array<i64: 4>,
-+    edge_padding_low = array<i64: 4>,
-+    interior_padding = array<i64: 0>
-+  } : (tensor<8xf32>, tensor<f32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
-+func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-+  func.return %0 : tensor<?xf32>
-+}
-+
-+// CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_recv_no_source_target_pairs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_recv_no_source_target_pairs(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
-+  //      CHECK: "vhlo.recv_v2"(%[[ARG0]]) <{
-+  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>,
-+  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<> : tensor<0xi64>
-+  // CHECK-SAME{LITERAL}: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-+  %0:2 = "stablehlo.recv"(%arg0) {
-+    channel_handle = #stablehlo.channel_handle<handle = 0, type = 3>,
-+    is_host_transfer = true
-+  } : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
-+  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-+  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.reduce"(%arg0, %arg1) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    dimensions = array<i64: 0>
-+  } : (tensor<16xf32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-+  // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.reduce_precision"(%arg0) {
-+    exponent_bits = 8 : i32,
-+    mantissa_bits = 10 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK_lABEL: "op_reduce_with_promotable_types"
-+func.func @op_reduce_with_promotable_types(%arg0: tensor<4x4xf32>, %arg1 : tensor<f32>)
-+    -> (tensor<4xf64>) {
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0:.*]], %[[ARG1:.*]])
-+  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-+  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-+  //  CHECK: }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x!vhlo.f64_v1>
-+  %0 = "stablehlo.reduce"(%arg0, %arg1) ({
-+  ^bb0(%arg2: tensor<f64>, %arg3: tensor<f64> ):
-+    %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-+    "stablehlo.return"(%1) : (tensor<f64>) -> ()
-+
-+  }) {dimensions = array<i64: 0>} : (tensor<4x4xf32>, tensor<f32>) -> tensor<4xf64>
-+
-+  func.return %0: tensor<4xf64>
-+}
-+
-+// CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.reduce_scatter"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg1, %arg2) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    scatter_dimension = 0 : i64,
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+    use_global_device_ids
-+  } : (tensor<16xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK_lABEL: "op_reduce_scatter_with_promotable_types"
-+func.func @op_reduce_scatter_with_promotable_types(%data: tensor<4x16xf32>) -> tensor<4x4xf64> {
-+  //  CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0:.*]])
-+  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-+  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-+  //  CHECK: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f64_v1>
-+  %0 = "stablehlo.reduce_scatter"(%data) ({
-+    ^bb0(%arg2: tensor<f64>, %arg3: tensor<f64>):
-+    %1 = stablehlo.add %arg2, %arg3 : tensor<f64>
-+    "stablehlo.return"(%1) : (tensor<f64>) -> ()
-+  }) {replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
-+      scatter_dimension = 1 : i64,
-+      channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+      use_global_device_ids} : (tensor<4x16xf32>) -> tensor<4x4xf64>
-+  func.return %0 : tensor<4x4xf64>
-+}
-+
-+
-+// CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-+  // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-+  //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-+  //          CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-+  //          CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 4, 4, 1]> : tensor<4xi64>>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.maximum_v1"(%[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<2x17x31x7x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<2x9x16x7x!vhlo.f32_v1>
-+  %0 = "stablehlo.reduce_window"(%arg0, %arg1) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %1 = "stablehlo.maximum"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    window_dimensions = array<i64: 1, 2, 2, 1>,
-+    window_strides = array<i64: 1, 4, 4, 1>,
-+    base_dilations = array<i64: 1, 2, 2, 1>,
-+    window_dilations = array<i64: 1, 2, 2, 1>,
-+    padding = dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>
-+  } : (tensor<2x17x31x7xf32>, tensor<f32>) -> tensor<2x9x16x7xf32>
-+  func.return %0 : tensor<2x9x16x7xf32>
-+}
-+
-+// CHECK-LABEL: "op_reduce_window_with_promotable_types"
-+func.func @op_reduce_window_with_promotable_types(%arg0: tensor<4x2xf32>,
-+    %arg1: tensor<4x2xf32>, %init0: tensor<f32>, %init1: tensor<f32>) ->
-+    (tensor<2x2xf64>, tensor<2x2xf32>) {
-+  //  CHECK: "vhlo.reduce_window_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]], %[[ARG3:.*]])
-+  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]], %[[VAL2:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //  CHECK: }) : (!vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<2x2x!vhlo.f64_v1>, !vhlo.tensor_v1<2x2x!vhlo.f32_v1>)
-+  %0:2 = "stablehlo.reduce_window"(%arg0, %arg1, %init0, %init1) ({
-+         ^bb0(%a0: tensor<f64>, %a1: tensor<f32>, %b0: tensor<f64>,
-+                %b1: tensor<f32>):
-+              %2 = stablehlo.add %a0, %b0 : tensor<f64>
-+              %3 = stablehlo.add %a1, %b1 : tensor<f32>
-+              "stablehlo.return"(%2,%3) : (tensor<f64>, tensor<f32>) -> ()
-+            })
-+         { padding = dense<[[2, 2], [0, 0]]> : tensor<2x2xi64>,
-+         window_dimensions = array<i64: 5, 1>,
-+         window_strides = array<i64: 3, 1> }
-+         : (tensor<4x2xf32>, tensor<4x2xf32>, tensor<f32>, tensor<f32>) ->
-+              (tensor<2x2xf64>, tensor<2x2xf32>)
-+  func.return %0#0, %0#1 : tensor<2x2xf64>, tensor<2x2xf32>
-+}
-+
-+// CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_replica_id"
-+func.func @op_replica_id() -> tensor<ui32> {
-+  // CHECK: "vhlo.replica_id_v1"() : () -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  %0 = "stablehlo.replica_id"() : () -> tensor<ui32>
-+  func.return %0 : tensor<ui32>
-+}
-+
-+// CHECK-LABEL: "op_partition_id"
-+func.func @op_partition_id() -> tensor<ui32> {
-+  // CHECK: "vhlo.partition_id_v1"() : () -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  %0 = "stablehlo.partition_id"() : () -> tensor<ui32>
-+  func.return %0 : tensor<ui32>
-+}
-+
-+// CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-+  func.return %0 : tensor<4x4xf32>
-+}
-+
-+// CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.case"(%arg0) ({
-+    "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-+  }) : (tensor<i32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.reverse"(%arg0) {
-+    dimensions = array<i64: 0>
-+  } : (tensor<16xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-+  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-+    rng_algorithm = #stablehlo<rng_algorithm PHILOX>
-+  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
-+  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-+    rng_distribution = #stablehlo<rng_distribution NORMAL>
-+  } : (tensor<f32>, tensor<f32>, tensor<0xindex>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.rsqrt_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
-+  //      CHECK: "vhlo.scatter_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   input_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   scatter_dims_to_operand_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   scatter_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   unique_indices = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   update_window_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>
-+  %0 = "stablehlo.scatter"(%arg0, %arg1, %arg2) ({
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    scatter_dimension_numbers = #stablehlo.scatter<
-+      update_window_dims = [1],
-+      inserted_window_dims = [0, 1],
-+      scatter_dims_to_operand_dims = [0, 1],
-+      index_vector_dim = 1
-+    >,
-+    indices_are_sorted = true,
-+    unique_indices = true
-+  } : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) -> tensor<200x100x300xf32>
-+  func.return %0 : tensor<200x100x300xf32>
-+}
-+
-+// CHECK-LABEL: "op_scatter_with_batching_dims"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_scatter_with_batching_dims(%arg0: tensor<10x200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<10x200x100x300xf32> {
-+  //      CHECK: "vhlo.scatter_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   input_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   scatter_dims_to_operand_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   scatter_indices_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   unique_indices = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   update_window_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<10x200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x200x100x300x!vhlo.f32_v1>
-+  %0 = "stablehlo.scatter"(%arg0, %arg1, %arg2) ({
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    scatter_dimension_numbers = #stablehlo.scatter<
-+      update_window_dims = [1],
-+      inserted_window_dims = [1, 2],
-+      input_batching_dims = [0],
-+      scatter_dims_to_operand_dims = [1, 2],
-+      scatter_indices_batching_dims = [0],
-+      index_vector_dim = 1
-+    >,
-+    indices_are_sorted = true,
-+    unique_indices = true
-+  } : (tensor<10x200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) -> tensor<10x200x100x300xf32>
-+  func.return %0 : tensor<10x200x100x300xf32>
-+}
-+
-+// CHECK_lABEL: "op_scatter_with_promotable_types"
-+func.func @op_scatter_with_promotable_types(%input_tensor: tensor<200x100x300xf32>,
-+    %scatter_indices: tensor<10x2xi32>, %updates: tensor<10x300xf32>) ->
-+      tensor<200x100x300xf64> {
-+  //  CHECK: "vhlo.scatter_v2"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]])
-+  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-+  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-+  //  CHECK: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f64_v1>
-+  %0 = "stablehlo.scatter" (%input_tensor, %scatter_indices, %updates) ({
-+  ^bb0(%lhs: tensor<f64>, %rhs: tensor<f64>):
-+    %add = stablehlo.add %lhs, %rhs : tensor<f64>
-+    "stablehlo.return"(%add) : (tensor<f64>) -> ()
-+  }) {
-+    scatter_dimension_numbers = #stablehlo.scatter<
-+      update_window_dims = [1],
-+      inserted_window_dims = [0, 1],
-+      scatter_dims_to_operand_dims = [0, 1],
-+      index_vector_dim = 1
-+    >,
-+    indices_are_sorted = true,
-+    unique_indices = true
-+  } : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) ->
-+      tensor<200x100x300xf64>
-+  func.return %0 : tensor<200x100x300xf64>
-+}
-+
-+// CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-+  // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-+  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG31:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG41:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL11:.*]] = "vhlo.compare_v1"(%[[ARG31]], %[[ARG41]]) <{compare_type = #vhlo<comparison_type_v1 TOTALORDER>, comparison_direction = #vhlo<comparison_direction_v1 GE>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL11]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-+  // CHECK-NEXT: }, {
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG32:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG42:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL12:.*]] = "vhlo.add_v1"(%[[ARG32]], %[[ARG42]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL12]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>, !vhlo.tensor_v1<12x13x13x66x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>
-+  %0 = "stablehlo.select_and_scatter"(%arg0, %arg1, %arg2) ({
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.compare"(%arg3, %arg4) {compare_type = #stablehlo<comparison_type TOTALORDER>, comparison_direction = #stablehlo<comparison_direction GE>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+      "stablehlo.return"(%1) : (tensor<i1>) -> ()
-+  }, {
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    window_dimensions = array<i64: 1, 2, 2, 1>,
-+    window_strides = array<i64: 1, 2, 2, 1>,
-+    padding = dense<1> : tensor<4x2xi64>
-+  } : (tensor<10x24x24x64xf32>, tensor<12x13x13x66xf32>, tensor<f32>) -> tensor<10x24x24x64xf32>
-+  func.return %0 : tensor<10x24x24x64xf32>
-+}
-+
-+// CHECK-LABEL: "op_select_and_scatter_with_promotable_types"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_select_and_scatter_with_promotable_types(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf64> {
-+  // CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
-+  // CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-+  // CHECK:     %[[VAL:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK:     "vhlo.return_v1"(%[[VAL]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-+  // CHECK: }) : (!vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>, !vhlo.tensor_v1<12x13x13x66x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x24x24x64x!vhlo.f64_v1>
-+  %0 = "stablehlo.select_and_scatter"(%arg0, %arg1, %arg2) ({
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.compare"(%arg3, %arg4) {compare_type = #stablehlo<comparison_type TOTALORDER>, comparison_direction = #stablehlo<comparison_direction GE>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+      "stablehlo.return"(%1) : (tensor<i1>) -> ()
-+  }, {
-+    ^bb0(%arg3: tensor<f64>, %arg4: tensor<f64>):
-+      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-+      "stablehlo.return"(%1) : (tensor<f64>) -> ()
-+  }) {
-+    window_dimensions = array<i64: 1, 2, 2, 1>,
-+    window_strides = array<i64: 1, 2, 2, 1>,
-+    padding = dense<1> : tensor<4x2xi64>
-+  } : (tensor<10x24x24x64xf32>, tensor<12x13x13x66xf32>, tensor<f32>) -> tensor<10x24x24x64xf64>
-+  func.return %0 : tensor<10x24x24x64xf64>
-+}
-+
-+// CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_send_no_source_target_pairs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_send_no_source_target_pairs(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
-+  //      CHECK: "vhlo.send_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>,
-+  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-+  // CHECK-SAME{LITERAL}: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-+  %0 = "stablehlo.send"(%arg0, %arg1) {
-+    channel_handle = #stablehlo.channel_handle<handle = 0, type = 2>,
-+    is_host_transfer = true
-+  } : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-+  func.return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-+    dimension = 0 : i64
-+  } : (tensor<?xf32>, tensor<i32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.sine_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-+  // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-+  %0 = "stablehlo.slice"(%arg0) {
-+    start_indices = array<i64: 0>,
-+    limit_indices = array<i64: 4>,
-+    strides = array<i64: 1>
-+  } : (tensor<16xf32>) -> tensor<4xf32>
-+  func.return %0 : tensor<4xf32>
-+}
-+
-+// CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.compare_v1"(%[[ARG1]], %[[ARG2]]) <{compare_type = #vhlo<comparison_type_v1 FLOAT>, comparison_direction = #vhlo<comparison_direction_v1 GT>}>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.sort"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
-+      %1 = "stablehlo.compare"(%arg1, %arg2) {compare_type = #stablehlo<comparison_type FLOAT>, comparison_direction = #stablehlo<comparison_direction GT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+      "stablehlo.return"(%1) : (tensor<i1>) -> ()
-+  }) {
-+    dimension = 0 : i64,
-+    is_stable = true
-+  } : (tensor<16xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.sqrt_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_tan"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_tan(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.tan_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.tan"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.tanh_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-+  %0 = "stablehlo.torch_index_select"(%arg0, %arg1) {
-+    dim = 0 : i64,
-+    batch_dims = 0 : i64
-+  } : (tensor<5x1x5xf32>, tensor<2xi32>) -> tensor<2x1x5xf32>
-+  func.return %0 : tensor<2x1x5xf32>
-+}
-+
-+// CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.transpose"(%arg0) {
-+    permutation = array<i64: 1, 0>
-+  } : (tensor<16x8xf32>) -> tensor<8x16xf32>
-+  func.return %0 : tensor<8x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-+  // CHECK-SAME:   unit_diagonal = #vhlo.bool_v1<true>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-+    left_side = true,
-+    lower = true,
-+    unit_diagonal = true,
-+    transpose_a = #stablehlo<transpose NO_TRANSPOSE>
-+  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-+  func.return %0 : tuple<tensor<f32>>
-+}
-+
-+// CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-+  %0 = "stablehlo.unary_einsum"(%arg0) {
-+    einsum_config = "ab->a"
-+  } : (tensor<8x16xf32>) -> tensor<8xf32>
-+  func.return %0 : tensor<8xf32>
-+}
-+
-+// CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-+  func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
-+}
-+
-+// CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-+  // CHECK-NEXT:   }, {
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>)
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.while"(%arg0) ({
-+    ^bb0(%arg1: tensor<i1>):
-+      "stablehlo.return"(%arg1) : (tensor<i1>) -> ()
-+    }, {
-+    ^bb0(%arg1: tensor<i1>):
-+      "stablehlo.return"(%arg1) : (tensor<i1>) -> ()
-+  }) : (tensor<i1>) -> tensor<i1>
-+  func.return %0: tensor<i1>
-+}
-+
-+// CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// ============ TYPES ============
-+
-+// CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "type_i2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i2(%arg0: tensor<i2>, %arg1: tensor<i2>) -> tensor<i2> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i2_v1>, !vhlo.tensor_v1<!vhlo.i2_v1>) -> !vhlo.tensor_v1<!vhlo.i2_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i2>, tensor<i2>) -> tensor<i2>
-+  func.return %0 : tensor<i2>
-+}
-+
-+// CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-+  func.return %0 : tensor<i4>
-+}
-+
-+// CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-+  func.return %0 : tensor<i8>
-+}
-+
-+// CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-+  func.return %0 : tensor<i16>
-+}
-+
-+// CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-+  func.return %0 : tensor<i64>
-+}
-+
-+// CHECK-LABEL: "type_ui2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_ui2(%arg0: tensor<ui2>, %arg1: tensor<ui2>) -> tensor<ui2> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui2_v1>, !vhlo.tensor_v1<!vhlo.ui2_v1>) -> !vhlo.tensor_v1<!vhlo.ui2_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui2>, tensor<ui2>) -> tensor<ui2>
-+  func.return %0 : tensor<ui2>
-+}
-+
-+// CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-+  func.return %0 : tensor<ui4>
-+}
-+
-+// CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-+  func.return %0 : tensor<ui8>
-+}
-+
-+// CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-+  func.return %0 : tensor<ui16>
-+}
-+
-+// CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-+  func.return %0 : tensor<ui32>
-+}
-+
-+// CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-+  func.return %0 : tensor<ui64>
-+}
-+
-+// CHECK-LABEL: "type_f4E2M1FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f4E2M1FN(%arg0: tensor<f4E2M1FN>, %arg1: tensor<f4E2M1FN>) -> tensor<f4E2M1FN> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f4E2M1FN_v1>, !vhlo.tensor_v1<!vhlo.f4E2M1FN_v1>) -> !vhlo.tensor_v1<!vhlo.f4E2M1FN_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f4E2M1FN>, tensor<f4E2M1FN>) -> tensor<f4E2M1FN>
-+  func.return %0 : tensor<f4E2M1FN>
-+}
-+
-+// CHECK-LABEL: "type_f6E2M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f6E2M3FN(%arg0: tensor<f6E2M3FN>, %arg1: tensor<f6E2M3FN>) -> tensor<f6E2M3FN> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f6E2M3FN_v1>, !vhlo.tensor_v1<!vhlo.f6E2M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f6E2M3FN_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f6E2M3FN>, tensor<f6E2M3FN>) -> tensor<f6E2M3FN>
-+  func.return %0 : tensor<f6E2M3FN>
-+}
-+
-+// CHECK-LABEL: "type_f6E3M2FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f6E3M2FN(%arg0: tensor<f6E3M2FN>, %arg1: tensor<f6E3M2FN>) -> tensor<f6E3M2FN> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f6E3M2FN_v1>, !vhlo.tensor_v1<!vhlo.f6E3M2FN_v1>) -> !vhlo.tensor_v1<!vhlo.f6E3M2FN_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f6E3M2FN>, tensor<f6E3M2FN>) -> tensor<f6E3M2FN>
-+  func.return %0 : tensor<f6E3M2FN>
-+}
-+
-+// CHECK-LABEL: "type_f8E3M4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E3M4(%arg0: tensor<f8E3M4>, %arg1: tensor<f8E3M4>) -> tensor<f8E3M4> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E3M4_v1>, !vhlo.tensor_v1<!vhlo.f8E3M4_v1>) -> !vhlo.tensor_v1<!vhlo.f8E3M4_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E3M4>, tensor<f8E3M4>) -> tensor<f8E3M4>
-+  func.return %0 : tensor<f8E3M4>
-+}
-+
-+// CHECK-LABEL: "type_f8E4M3"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E4M3(%arg0: tensor<f8E4M3>, %arg1: tensor<f8E4M3>) -> tensor<f8E4M3> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3>, tensor<f8E4M3>) -> tensor<f8E4M3>
-+  func.return %0 : tensor<f8E4M3>
-+}
-+
-+// CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-+  func.return %0 : tensor<f8E4M3FN>
-+}
-+
-+// CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-+  func.return %0 : tensor<f8E5M2>
-+}
-+
-+// CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-+  func.return %0 : tensor<f8E4M3FNUZ>
-+}
-+
-+// CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-+  func.return %0 : tensor<f8E4M3B11FNUZ>
-+}
-+
-+// CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-+  func.return %0 : tensor<f8E5M2FNUZ>
-+}
-+
-+// CHECK-LABEL: "type_f8E8M0FNU"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E8M0FNU(%arg0: tensor<f8E8M0FNU>, %arg1: tensor<f8E8M0FNU>) -> tensor<f8E8M0FNU> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E8M0FNU_v1>, !vhlo.tensor_v1<!vhlo.f8E8M0FNU_v1>) -> !vhlo.tensor_v1<!vhlo.f8E8M0FNU_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E8M0FNU>, tensor<f8E8M0FNU>) -> tensor<f8E8M0FNU>
-+  func.return %0 : tensor<f8E8M0FNU>
-+}
-+
-+// CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-+  func.return %0 : tensor<bf16>
-+}
-+
-+// CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-+  func.return %0 : tensor<f16>
-+}
-+
-+// CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-+  func.return %0 : tensor<f64>
-+}
-+
-+// CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-+  func.return %0 : tensor<complex<f32>>
-+}
-+
-+// CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-+  func.return %0 : tensor<complex<f64>>
-+}
-+
-+// CHECK-LABEL: "type_tf32"
-+// CHECK: #vhlo.type_v1<!vhlo.tf31_v1>
-+func.func @type_tf32() attributes {stablehlo.attr = tf32 } {
-+  return
-+}
-+
-+// CHECK-LABEL: "type_none"
-+// CHECK: #vhlo.type_v1<!vhlo.none_v1>
-+func.func @type_none() attributes {stablehlo.attr = none } {
-+  return
-+}
-+
-+// CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-+  func.return %0 : tensor<?xf32>
-+}
-+
-+// CHECK-LABEL: "type_per_tensor_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_per_tensor_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-+  func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
-+}
-+
-+// CHECK-LABEL: "type_per_axis_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @type_per_axis_quantization(%arg0: tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>) -> tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG0]]) : (!vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, [3.400000e+01, 3.400000e+01], [16, 16], -128:127, 1>>, !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, [3.400000e+01, 3.400000e+01], [16, 16], -128:127, 1>>) -> !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, [3.400000e+01, 3.400000e+01], [16, 16], -128:127, 1>>
-+  %0 = stablehlo.add %arg0, %arg0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
-+  func.return %0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
-+}
-+
-+//       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
-+// CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-+  return %arg0 : !stablehlo.token
-+}
-+
-+//       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
-+// CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-+  %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-+  return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo"
-+  // CHECK: (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>) -> !vhlo.tuple_v1<!vhlo.token_v1>
-+  } : (tuple<tensor<f32>>) -> tuple<!stablehlo.token>
-+  return %0 : tuple<!stablehlo.token>
-+}
-+
-+// CHECK-LABEL: type_buffer_function_input_output
-+// CHECK-NEXT: (%[[ARG0:.*]]: !vhlo.buffer_v1<2x!vhlo.f32_v1>)
-+func.func @type_buffer_function_input_output(%arg0: memref<2xf32>) -> memref<2xf32> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.buffer_v1<2x!vhlo.f32_v1>) -> ()
-+  func.return %arg0 : memref<2xf32>
-+}
-+
-+// CHECK-LABEL: type_buffer_special_custom_calls
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @type_buffer_special_custom_calls(%arg0: tensor<2xf32>) -> tensor<2xf32> {
-+  //               CHECK: %[[CALL0:.*]] = "vhlo.custom_call_v2"(%[[ARG0]])
-+  //          CHECK-SAME: call_target_name = #vhlo.string_v1<"Pin">
-+  //          CHECK-SAME: : (!vhlo.tensor_v1<2x!vhlo.f32_v1>) -> !vhlo.buffer_v1<2x!vhlo.f32_v1>
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "Pin",
-+    api_version = 4 : i32
-+  } : (tensor<2xf32>) -> memref<2xf32>
-+  //               CHECK: %{{.*}} = "vhlo.custom_call_v2"(%[[CALL0]])
-+  //          CHECK-SAME: call_target_name = #vhlo.string_v1<"Unpin">
-+  //          CHECK-SAME: : (!vhlo.buffer_v1<2x!vhlo.f32_v1>) -> !vhlo.tensor_v1<2x!vhlo.f32_v1>
-+  %1 = "stablehlo.custom_call"(%0) {
-+    call_target_name = "Unpin",
-+    api_version = 4 : i32
-+  } : (memref<2xf32>) -> tensor<2xf32>
-+  func.return %1 : tensor<2xf32>
-+}
-+
-+// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>>>
-+// CHECK-LABEL: "type_future"
-+func.func private @type_future() -> !stablehlo.future<tensor<f32>>
-+
-+// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<2x3x!vhlo.f32_v1>>>>
-+// CHECK-LABEL: "type_future_ranked"
-+func.func private @type_future_ranked() -> !stablehlo.future<tensor<2x3xf32>>
-+
-+// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<?x!vhlo.f32_v1>>>>
-+// CHECK-LABEL: "type_future_dynamic"
-+func.func private @type_future_dynamic() -> !stablehlo.future<tensor<?xf32>>
-+
-+// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>>>
-+// CHECK-LABEL: "type_future_multiple"
-+func.func private @type_future_multiple() -> !stablehlo.future<tensor<f32>, tensor<i32>>
-+
-+// CHECK-LABEL: "op_async_start_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_start_all_reduce(%arg0: tensor<4x4xf32>) -> !stablehlo.future<tensor<4x4xf32>> {
-+  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
-+  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.all_reduce_v2"(%[[BARG0]]) <{
-+  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>,
-+  // CHECK-SAME:     use_global_device_ids = #vhlo.bool_v1<false>
-+  // CHECK-SAME:   }> ({
-+  // CHECK-NEXT:     ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:       %[[VAL2:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:       "vhlo.return_v1"(%[[VAL2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<4x4x!vhlo.f32_v1>>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<4x4xf32>):
-+    %1 = "stablehlo.all_reduce"(%barg0) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %2 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%2) : (tensor<f32>) -> ()
-+    }) {replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
-+    "stablehlo.return"(%1) : (tensor<4x4xf32>) -> ()
-+  }) : (tensor<4x4xf32>) -> !stablehlo.future<tensor<4x4xf32>>
-+  func.return %0: !stablehlo.future<tensor<4x4xf32>>
-+}
-+
-+// CHECK-LABEL: "op_async_start_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_start_all_gather(%arg0: tensor<8x2xf32>) -> !stablehlo.future<tensor<8x8xf32>> {
-+  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
-+  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.all_gather_v2"(%[[BARG0]]) <{
-+  // CHECK-SAME:     all_gather_dim = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 2, 4, 6], [1, 3, 5, 7]]> : tensor<2x4xi64>>,
-+  // CHECK-SAME:     use_global_device_ids = #vhlo.bool_v1<false>
-+  // CHECK-SAME:   }> : (!vhlo.tensor_v1<8x2x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<8x8x!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<8x2x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<8x8x!vhlo.f32_v1>>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<8x2xf32>):
-+    %1 = "stablehlo.all_gather"(%barg0) {
-+      all_gather_dim = 1 : i64,
-+      replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7]]> : tensor<2x4xi64>
-+    } : (tensor<8x2xf32>) -> tensor<8x8xf32>
-+    "stablehlo.return"(%1) : (tensor<8x8xf32>) -> ()
-+  }) : (tensor<8x2xf32>) -> !stablehlo.future<tensor<8x8xf32>>
-+  func.return %0: !stablehlo.future<tensor<8x8xf32>>
-+}
-+
-+// CHECK-LABEL: "op_async_start_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_start_all_to_all(%arg0: tensor<4x16xf32>) -> !stablehlo.future<tensor<16x4xf32>> {
-+  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
-+  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.all_to_all_v2"(%[[BARG0]]) <{
-+  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:     concat_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-+  // CHECK-SAME:     split_count = #vhlo.integer_v1<4 : i64>,
-+  // CHECK-SAME:     split_dimension = #vhlo.integer_v1<1 : i64>
-+  // CHECK-SAME:   }> : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x4x!vhlo.f32_v1>
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<16x4x!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<16x4x!vhlo.f32_v1>>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<4x16xf32>):
-+    %1 = "stablehlo.all_to_all"(%barg0) {
-+      split_dimension = 1 : i64,
-+      concat_dimension = 0 : i64,
-+      split_count = 4 : i64,
-+      replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>
-+    } : (tensor<4x16xf32>) -> tensor<16x4xf32>
-+    "stablehlo.return"(%1) : (tensor<16x4xf32>) -> ()
-+  }) : (tensor<4x16xf32>) -> !stablehlo.future<tensor<16x4xf32>>
-+  func.return %0: !stablehlo.future<tensor<16x4xf32>>
-+}
-+
-+// CHECK-LABEL: "op_async_start_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_start_collective_permute(%arg0: tensor<128x32xf32>) -> !stablehlo.future<tensor<128x32xf32>> {
-+  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
-+  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.collective_permute_v1"(%[[BARG0]]) <{
-+  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:     source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-+  // CHECK-SAME:   }> : (!vhlo.tensor_v1<128x32x!vhlo.f32_v1>) -> !vhlo.tensor_v1<128x32x!vhlo.f32_v1>
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<128x32x!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<128x32x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<128x32x!vhlo.f32_v1>>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<128x32xf32>):
-+    %1 = "stablehlo.collective_permute"(%barg0) {
-+      source_target_pairs = dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>
-+    } : (tensor<128x32xf32>) -> tensor<128x32xf32>
-+    "stablehlo.return"(%1) : (tensor<128x32xf32>) -> ()
-+  }) : (tensor<128x32xf32>) -> !stablehlo.future<tensor<128x32xf32>>
-+  func.return %0: !stablehlo.future<tensor<128x32xf32>>
-+}
-+
-+// CHECK-LABEL: "op_async_start_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_start_collective_broadcast(%arg0: tensor<16x8xf32>) -> !stablehlo.future<tensor<16x8xf32>> {
-+  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
-+  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.collective_broadcast_v2"(%[[BARG0]]) <{
-+  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:     has_dynamic_root = #vhlo.bool_v1<false>,
-+  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-+  // CHECK-SAME:   }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<16x8x!vhlo.f32_v1>>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<16x8xf32>):
-+    %1 = "stablehlo.collective_broadcast"(%barg0) {
-+      replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
-+    } : (tensor<16x8xf32>) -> tensor<16x8xf32>
-+    "stablehlo.return"(%1) : (tensor<16x8xf32>) -> ()
-+  }) : (tensor<16x8xf32>) -> !stablehlo.future<tensor<16x8xf32>>
-+  func.return %0: !stablehlo.future<tensor<16x8xf32>>
-+}
-+
-+// CHECK-LABEL: "op_async_start_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_start_reduce_scatter(%arg0: tensor<4x16xf32>) -> !stablehlo.future<tensor<4x4xf32>> {
-+  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
-+  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.reduce_scatter_v1"(%[[BARG0]]) <{
-+  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-+  // CHECK-SAME:     scatter_dimension = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:     use_global_device_ids = #vhlo.bool_v1<false>
-+  // CHECK-SAME:   }> ({
-+  // CHECK-NEXT:     ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:       %[[VAL2:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:       "vhlo.return_v1"(%[[VAL2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<4x4x!vhlo.f32_v1>>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<4x16xf32>):
-+    %1 = "stablehlo.reduce_scatter"(%barg0) ({
-+      ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %2 = stablehlo.add %arg2, %arg3 : tensor<f32>
-+      "stablehlo.return"(%2) : (tensor<f32>) -> ()
-+    }) {replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
-+        scatter_dimension = 1 : i64} : (tensor<4x16xf32>) -> tensor<4x4xf32>
-+    "stablehlo.return"(%1) : (tensor<4x4xf32>) -> ()
-+  }) : (tensor<4x16xf32>) -> !stablehlo.future<tensor<4x4xf32>>
-+  func.return %0: !stablehlo.future<tensor<4x4xf32>>
-+}
-+
-+// CHECK-LABEL: "op_async_done"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_done(%arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
-+  // CHECK: %[[VAL1:.*]] = "vhlo.async_start_v1"(%[[ARG0]])
-+  // CHECK: %[[VAL2:.*]] = "vhlo.async_done_v1"(%[[VAL1]]) : (!vhlo.future_v1<!vhlo.tensor_v1<4x4x!vhlo.f32_v1>>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<4x4xf32>):
-+    %1 = "stablehlo.all_reduce"(%barg0) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %2 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%2) : (tensor<f32>) -> ()
-+    }) {replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
-+    "stablehlo.return"(%1) : (tensor<4x4xf32>) -> ()
-+  }) : (tensor<4x4xf32>) -> !stablehlo.future<tensor<4x4xf32>>
-+  %1 = "stablehlo.async_done"(%0) : (!stablehlo.future<tensor<4x4xf32>>) -> tensor<4x4xf32>
-+  func.return %1: tensor<4x4xf32>
-+}
-+
-+// ============ DEPENDENCIES  ============
-+
-+func.func @composite_target(%arg0: tensor<f32>) -> tensor<f32> {
-+  return %arg0: tensor<f32>
-+}
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_21_0.mlir.bc b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_21_0.mlir.bc
-Binary files /dev/null and stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_21_0.mlir.bc differ
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_22_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_22_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_22_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_22_0.mlir
-@@ -0,0 +1,3436 @@
-+// RUN: stablehlo-opt --mlir-print-op-generic %s.bc | FileCheck %s
-+// RUN: stablehlo-translate --deserialize %s.bc | stablehlo-translate --serialize --target=1.22.0 | stablehlo-opt --mlir-print-op-generic | FileCheck %s
-+// RUN: stablehlo-translate --deserialize %s.bc | stablehlo-opt > %t.0
-+// RUN: stablehlo-opt --strip-debuginfo %s > %t.1
-+// RUN: diff %t.0 %t.1
-+// RUN: stablehlo-translate --serialize --target=1.22.0 --strip-debuginfo %s > %t.2
-+// RUN: diff %s.bc %t.2
-+// RUN: %if asserts %{ stablehlo-opt --stablehlo-legalize-to-vhlo -emit-bytecode -debug-only=vhlo-bytecode %s 2>&1 | FileCheck --check-prefix=CHECK-WARN %s %}
-+// RUN: %if asserts %{ stablehlo-opt --stablehlo-legalize-to-vhlo -emit-bytecode %s | stablehlo-opt -debug-only=vhlo-bytecode 2>&1 | FileCheck --check-prefix=CHECK-WARN %s %}
-+
-+// CHECK-WARN-NOT: Not Implemented
-+
-+func.func private @mesh()
-+
-+// ============ ATTRIBUTES ============
-+
-+// CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-+    comparison_direction = #stablehlo<comparison_direction EQ>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-+    comparison_direction = #stablehlo<comparison_direction NE>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-+    comparison_direction = #stablehlo<comparison_direction GE>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-+    comparison_direction = #stablehlo<comparison_direction GT>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-+    comparison_direction = #stablehlo<comparison_direction LE>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-+    comparison_direction = #stablehlo<comparison_direction LT>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>
-+    // CHECK: compare_type = #vhlo<comparison_type_v1 NOTYPE>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>,
-+    // CHECK: compare_type = #vhlo<comparison_type_v1 FLOAT>,
-+    compare_type = #stablehlo<comparison_type FLOAT>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>,
-+    // CHECK: compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-+    compare_type = #stablehlo<comparison_type TOTALORDER>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>,
-+    // CHECK: compare_type = #vhlo<comparison_type_v1 SIGNED>,
-+    compare_type = #stablehlo<comparison_type SIGNED>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>,
-+    // CHECK: compare_type = #vhlo<comparison_type_v1 UNSIGNED>,
-+    compare_type = #stablehlo<comparison_type UNSIGNED>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_type_weakorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_type_weakorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>,
-+    // CHECK: compare_type = #vhlo<comparison_type_v1 WEAKORDER>,
-+    compare_type = #stablehlo<comparison_type WEAKORDER>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// ConvDimensionNumbers aka #stablehlo.conv is covered below.
-+
-+// CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_UNSPECIFIED>
-+    api_version = 0 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>
-+    api_version = 1 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>
-+    api_version = 2 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING_UNIFIED>
-+    api_version = 3 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_dict"
-+// CHECK: #vhlo.dict_v1<{#vhlo.string_v1<"attr1"> = #vhlo.integer_v1<1 : i32>, #vhlo.string_v1<"attr2"> = #vhlo.integer_v1<2 : i32>}
-+func.func @attr_dict() attributes {stablehlo.attr = {attr1 = 1 : i32, attr2 = 2 : i32}} {
-+  return
-+}
-+
-+// CHECK-LABEL: "attr_custom_call_api_version_typed_ffi"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+// CHECK: api_version = #vhlo<api_version_v1 API_VERSION_TYPED_FFI>
-+// CHECK-SAME: backend_config = #vhlo.dict_v1<{#vhlo.string_v1<"bar"> = #vhlo.integer_v1<42 : i32>}>
-+func.func @attr_custom_call_api_version_typed_ffi(%arg0: tensor<f32>) -> tensor<f32> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    backend_config= {bar = 42 : i32},
-+    api_version = 4 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+
-+// CHECK-LABEL: "attr_custom_call_api_version_typed_ffi_no_backend_config"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+// CHECK: api_version = #vhlo<api_version_v1 API_VERSION_TYPED_FFI>
-+// CHECK-SAME: backend_config = #vhlo.dict_v1<{}>
-+func.func @attr_custom_call_api_version_typed_ffi_no_backend_config(%arg0: tensor<f32>) -> tensor<f32> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    api_version = 4 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// DotDimensionNumbers aka #stablehlo.dot is covered below.
-+
-+// CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-+  %0 = "stablehlo.fft"(%arg0) {
-+    // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-+    fft_type = #stablehlo<fft_type FFT>,
-+    fft_length = array<i64: 16>
-+  } : (tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>>
-+  func.return %0 : tensor<16xcomplex<f32>>
-+}
-+
-+// CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-+  %0 = "stablehlo.fft"(%arg0) {
-+    // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-+    fft_type = #stablehlo<fft_type IFFT>,
-+    fft_length = array<i64: 16>
-+  } : (tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>>
-+  func.return %0 : tensor<16xcomplex<f32>>
-+}
-+
-+// CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-+  %0 = "stablehlo.fft"(%arg0) {
-+    // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-+    fft_type = #stablehlo<fft_type RFFT>,
-+    fft_length = array<i64: 16>
-+  } : (tensor<16xf32>) -> tensor<9xcomplex<f32>>
-+  func.return %0 : tensor<9xcomplex<f32>>
-+}
-+
-+// CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-+  %0 = "stablehlo.fft"(%arg0) {
-+    // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-+    fft_type = #stablehlo<fft_type IRFFT>,
-+    fft_length = array<i64: 16>
-+  } : (tensor<9xcomplex<f32>>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "exponential_HIGHEST"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}
-+func.func @exponential_HIGHEST(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
-+  %0 = "stablehlo.exponential"(%arg0) {
-+    // CHECK: result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 HIGHEST>>
-+    result_accuracy = #stablehlo.result_accuracy<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #stablehlo.result_accuracy_mode<HIGHEST>>
-+  } : (tensor<8x16xf32>) -> tensor<8x16xf32>
-+  func.return %0 : tensor<8x16xf32>
-+}
-+
-+// CHECK-LABEL: "exponential_TOLERANCE"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}
-+func.func @exponential_TOLERANCE(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
-+  %0 = "stablehlo.exponential"(%arg0) {
-+    // CHECK: result_accuracy = #vhlo.result_accuracy_v1<atol = 1.000000e-05, rtol = 0.000000e+00, ulps = 1, mode = #vhlo<result_accuracy_mode_v1 TOLERANCE>>
-+    result_accuracy = #stablehlo.result_accuracy<atol = 1.000000e-5, rtol = 0.000000e+00, ulps = 1, mode = #stablehlo.result_accuracy_mode<TOLERANCE>>
-+  } : (tensor<8x16xf32>) -> tensor<8x16xf32>
-+  func.return %0 : tensor<8x16xf32>
-+}
-+
-+// CHECK-LABEL: "exponential_DEFAULT"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}
-+func.func @exponential_DEFAULT(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
-+  %0 = "stablehlo.exponential"(%arg0) {
-+    // CHECK: result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>
-+    result_accuracy = #stablehlo.result_accuracy<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #stablehlo.result_accuracy_mode<DEFAULT>>
-+  } : (tensor<8x16xf32>) -> tensor<8x16xf32>
-+  func.return %0 : tensor<8x16xf32>
-+}
-+
-+// GatherDimensionNumbers aka #stablehlo.gather is covered below.
-+
-+// CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-+  %0 = "stablehlo.dot"(%arg0, %arg1) {
-+    // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-+  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-+  func.return %0 : tensor<8x8xf32>
-+}
-+
-+// CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-+  %0 = "stablehlo.dot"(%arg0, %arg1) {
-+    // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-+    precision_config = [#stablehlo<precision HIGH>, #stablehlo<precision HIGH>]
-+  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-+  func.return %0 : tensor<8x8xf32>
-+}
-+
-+// CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-+  %0 = "stablehlo.dot"(%arg0, %arg1) {
-+    // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-+    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
-+  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-+  func.return %0 : tensor<8x8xf32>
-+}
-+
-+// CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-+  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-+    // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-+    rng_algorithm = #stablehlo<rng_algorithm DEFAULT>
-+  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
-+  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-+  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-+    // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-+    rng_algorithm = #stablehlo<rng_algorithm THREE_FRY>
-+  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
-+  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-+  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-+    // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-+    rng_algorithm = #stablehlo<rng_algorithm PHILOX>
-+  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
-+  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-+  %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-+    // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-+    rng_distribution = #stablehlo<rng_distribution UNIFORM>
-+  } : (tensor<f32>, tensor<f32>, tensor<0xindex>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-+  %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-+    // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-+    rng_distribution = #stablehlo<rng_distribution NORMAL>
-+  } : (tensor<f32>, tensor<f32>, tensor<0xindex>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
-+
-+// CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-+  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-+    left_side = true,
-+    lower = true,
-+    unit_diagonal = true,
-+    // transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-+    transpose_a = #stablehlo<transpose NO_TRANSPOSE>
-+  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-+  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-+    left_side = true,
-+    lower = true,
-+    unit_diagonal = true,
-+    // transpose_a = #vhlo<transpose_v1 TRANSPOSE>,
-+    transpose_a = #stablehlo<transpose TRANSPOSE>
-+  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-+  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-+    left_side = true,
-+    lower = true,
-+    unit_diagonal = true,
-+    // transpose_a = #vhlo<transpose_v1 ADJOINT>,
-+    transpose_a = #stablehlo<transpose ADJOINT>
-+  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
-+
-+// CHECK-LABEL: "attr_type_extensions_bounds"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+  func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
-+}
-+
-+// CHECK-LABEL: "attr_frontend_attributes"
-+func.func @attr_frontend_attributes(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: some.unregistered_attr
-+  %1 = stablehlo.cosine %arg0 {some.unregistered_attr = 1 : i32} : tensor<f32>
-+  return %1 : tensor<f32>
-+}
-+
-+// Builtin attriubute tests
-+
-+// CHECK-LABEL: "byte_packed_boolean"
-+func.func @byte_packed_boolean() -> (tensor<8xi1>, tensor<8xi1>, tensor<4xi1>, tensor<4xi1>, tensor<16xi1>) {
-+  // CHECK: #vhlo.tensor_v1<dense<[true, false, false, false, false, false, false, false]
-+  // CHECK-NEXT: #vhlo.tensor_v1<dense<true> : tensor<8xi1>>
-+  // CHECK-NEXT: #vhlo.tensor_v1<dense<[true, false, false, false]> : tensor<4xi1>>
-+  // CHECK-NEXT: #vhlo.tensor_v1<dense<true> : tensor<4xi1>>
-+  // CHECK-NEXT: #vhlo.tensor_v1<dense<[true, false, false, false, false, false, false, false, true, false, false, false, false, false, false, false]> : tensor<16xi1>>
-+  %c = stablehlo.constant dense<[true, false, false, false, false, false, false, false]> : tensor<8xi1>
-+  %c_0 = stablehlo.constant dense<true> : tensor<8xi1>
-+  %c_1 = stablehlo.constant dense<[true, false, false, false]> : tensor<4xi1>
-+  %c_2 = stablehlo.constant dense<true> : tensor<4xi1>
-+  %c_3 = stablehlo.constant dense<[true, false, false, false, false, false, false, false, true, false, false, false, false, false, false, false]> : tensor<16xi1>
-+  return %c, %c_0, %c_1, %c_2, %c_3 : tensor<8xi1>, tensor<8xi1>, tensor<4xi1>, tensor<4xi1>, tensor<16xi1>
-+}
-+
-+// ============ DEFAULTS ============
-+
-+// CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
-+  //               CHECK: "vhlo.all_gather_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.all_gather"(%arg0) {
-+    all_gather_dim = 1 : i64,
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
-+  } : (tensor<16x8xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "default_all_gather_variadic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_all_gather_variadic(%arg0: tensor<16x8xf32>, %arg1: tensor<16x8xf32>) -> (tensor<16x16xf32>, tensor<16x16xf32>) {
-+  %0:2 = "stablehlo.all_gather"(%arg0, %arg1) {
-+    all_gather_dim = 1 : i64,
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
-+  } : (tensor<16x8xf32>, tensor<16x8xf32>) -> (tensor<16x16xf32>, tensor<16x16xf32>)
-+  func.return %0#0, %0#1 : tensor<16x16xf32>, tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.all_reduce_v2"(%[[ARG0]])
-+  //          CHECK-SAME: <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+
-+  %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 = dense<[[0], [1]]> : tensor<2x1xi64>
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_replica_group_mesh_axes"
-+func.func @attr_replica_group_mesh_axes(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.all_reduce_v2"
-+  // CHECK-SAME: replica_groups = #vhlo.replica_group_mesh_axes_v1<
-+  // CHECK-SAME:   mesh = #vhlo.string_v1<"mesh">,
-+  // CHECK-SAME:   axes = #vhlo.array_v1<[#vhlo.axis_ref_v1<name = #vhlo.string_v1<"x">, sub_axis_info = #vhlo.sub_axis_info_v1<pre_size = 1, size = 2>>, #vhlo.axis_ref_v1<name = #vhlo.string_v1<"y">>]>
-+  // CHECK-SAME: >
-+  %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 = @mesh,
-+      axes = [
-+        #stablehlo.axis_ref<name = "x", sub_axis_info = (1)2>,
-+        #stablehlo.axis_ref<name = "y">
-+      ]
-+    >
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+
-+
-+// CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
-+  //               CHECK: "vhlo.all_to_all_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-+  //          CHECK-SAME:   split_count = #vhlo.integer_v1<4 : i64>
-+  //          CHECK-SAME:   split_dimension = #vhlo.integer_v1<1 : i64>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x4x!vhlo.f32_v1>
-+  %0 = "stablehlo.all_to_all"(%arg0) {
-+    split_dimension = 1 : i64,
-+    concat_dimension = 0 : i64,
-+    split_count = 4 : i64,
-+    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>
-+  } : (tensor<4x16xf32>) -> tensor<16x4xf32>
-+  func.return %0 : tensor<16x4xf32>
-+}
-+
-+// CHECK-LABEL: "default_all_to_all_variadic"
-+func.func @default_all_to_all_variadic(%arg0: tensor<4x16xf32>, %arg1: tensor<5x16xf32>) -> (tensor<16x4xf32>, tensor<20x4xf32>) {
-+  %0:2 = "stablehlo.all_to_all"(%arg0, %arg1) {
-+    split_dimension = 1 : i64,
-+    concat_dimension = 0 : i64,
-+    split_count = 4 : i64,
-+    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>
-+  } : (tensor<4x16xf32>, tensor<5x16xf32>) -> (tensor<16x4xf32>, tensor<20x4xf32>)
-+  func.return %0#0, %0#1 : tensor<16x4xf32>, tensor<20x4xf32>
-+}
-+
-+// CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-+  func.return %0 : tensor<1x16x16xf32>
-+}
-+
-+// CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.collective_permute"(%arg0) {
-+    source_target_pairs = dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>
-+  } : (tensor<16x8xf32>) -> tensor<16x8xf32>
-+  func.return %0 : tensor<16x8xf32>
-+}
-+
-+// CHECK-LABEL: "default_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_collective_broadcast(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
-+  //               CHECK: "vhlo.collective_broadcast_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<false>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.collective_broadcast"(%arg0) {
-+    replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
-+  } : (tensor<16x8xf32>) -> tensor<16x8xf32>
-+  func.return %0 : tensor<16x8xf32>
-+}
-+
-+// CHECK-LABEL: "default_collective_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_collective_reduce(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.collective_reduce_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<false>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>,
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.collective_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 = dense<[[0, 1]]> : tensor<1x2xi64>
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-+  // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "default_composite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_composite(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.composite_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{}>
-+  //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
-+  //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
-+  //          CHECK-SAME:   version = #vhlo.integer_v1<0 : i64>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.composite"(%arg0) {
-+    name = "stablehlo.composite_target",
-+    decomposition = @composite_target
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
-+  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<2x2xi64>>,
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
-+  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<false> : tensor<2xi1>>,
-+  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x6x6x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.convolution"(%arg0, %arg1) {
-+    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
-+    feature_group_count = 1 : i64,
-+    batch_group_count = 1 : i64
-+  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32>
-+  func.return %0 : tensor<1x6x6x16xf32>
-+}
-+
-+// CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
-+  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-+  // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-+  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-+  // CHECK-SAME:   called_computations = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<false>,
-+  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[]>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo"
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
-+  //      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   allow_imprecise_accumulation = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   lhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   num_primitive_operations = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+  // CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   rhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
-+  // CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
-+    dot_dimension_numbers = #stablehlo.dot<
-+      lhs_batching_dimensions = [0],
-+      lhs_contracting_dimensions = [2],
-+      rhs_batching_dimensions = [0],
-+      rhs_contracting_dimensions = [1]
-+    >
-+  } : (tensor<8x8x16xf32>, tensor<8x16x8xf32>) -> tensor<8x8x8xf32>
-+  func.return %0 : tensor<8x8x8xf32>
-+}
-+
-+// CHECK-LABEL: "dot_general_algorithm"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @dot_general_algorithm(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
-+//      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
-+// CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.f32_v1>,
-+// CHECK-SAME:   allow_imprecise_accumulation = #vhlo.bool_v1<false>,
-+// CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.tf31_v1>,
-+// CHECK-SAME:   num_primitive_operations = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+// CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.tf31_v1>
-+// CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
-+    dot_dimension_numbers = #stablehlo.dot<
-+      lhs_batching_dimensions = [0],
-+      lhs_contracting_dimensions = [2],
-+      rhs_batching_dimensions = [0],
-+      rhs_contracting_dimensions = [1]
-+    >,
-+    algorithm = #stablehlo.dot_algorithm<
-+      lhs_precision_type = tf32,
-+      rhs_precision_type = tf32,
-+      accumulation_type = f32,
-+      lhs_component_count = 1,
-+      rhs_component_count = 1,
-+      num_primitive_operations = 1,
-+      allow_imprecise_accumulation = false
-+    >
-+  } : (tensor<8x8x16xf32>, tensor<8x16x8xf32>) -> tensor<8x8x8xf32>
-+  func.return %0 : tensor<8x8x8xf32>
-+}
-+
-+// CHECK-LABEL: "dot_general_algorithm_f8e4m3fn_x3"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @dot_general_algorithm_f8e4m3fn_x3(%arg0: tensor<8x8x16xbf16>, %arg1: tensor<8x16x8xbf16>) -> tensor<8x8x8xf32> {
-+//      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
-+// CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.f32_v1>,
-+// CHECK-SAME:   allow_imprecise_accumulation = #vhlo.bool_v1<false>,
-+// CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>,
-+// CHECK-SAME:   num_primitive_operations = #vhlo.integer_v1<3 : i64>,
-+// CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+// CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>
-+// CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.bf16_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.bf16_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
-+    dot_dimension_numbers = #stablehlo.dot<
-+      lhs_batching_dimensions = [0],
-+      lhs_contracting_dimensions = [2],
-+      rhs_batching_dimensions = [0],
-+      rhs_contracting_dimensions = [1]
-+    >,
-+    algorithm = #stablehlo.dot_algorithm<
-+      lhs_precision_type = f8E4M3FN,
-+      rhs_precision_type = f8E4M3FN,
-+      accumulation_type = f32,
-+      lhs_component_count = 1,
-+      rhs_component_count = 1,
-+      num_primitive_operations = 3,
-+      allow_imprecise_accumulation = false
-+    >
-+  } : (tensor<8x8x16xbf16>, tensor<8x16x8xbf16>) -> tensor<8x8x8xf32>
-+  func.return %0 : tensor<8x8x8xf32>
-+}
-+
-+// CHECK-LABEL: "dot_general_algorithm_f8e4m3fn_x4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @dot_general_algorithm_f8e4m3fn_x4(%arg0: tensor<8x8x16xbf16>, %arg1: tensor<8x16x8xbf16>) -> tensor<8x8x8xf32> {
-+//      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
-+// CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.f32_v1>,
-+// CHECK-SAME:   allow_imprecise_accumulation = #vhlo.bool_v1<false>,
-+// CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>,
-+// CHECK-SAME:   num_primitive_operations = #vhlo.integer_v1<4 : i64>,
-+// CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+// CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>
-+// CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.bf16_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.bf16_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
-+    dot_dimension_numbers = #stablehlo.dot<
-+      lhs_batching_dimensions = [0],
-+      lhs_contracting_dimensions = [2],
-+      rhs_batching_dimensions = [0],
-+      rhs_contracting_dimensions = [1]
-+    >,
-+    algorithm = #stablehlo.dot_algorithm<
-+      lhs_precision_type = f8E4M3FN,
-+      rhs_precision_type = f8E4M3FN,
-+      accumulation_type = f32,
-+      lhs_component_count = 1,
-+      rhs_component_count = 1,
-+      num_primitive_operations = 4,
-+      allow_imprecise_accumulation = false
-+    >
-+  } : (tensor<8x8x16xbf16>, tensor<8x16x8xbf16>) -> tensor<8x8x8xf32>
-+  func.return %0 : tensor<8x8x8xf32>
-+}
-+
-+// CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_broadcast_in_dim"(%arg0, %arg1) {
-+    broadcast_dimensions = array<i64: 0, 1>
-+  } : (tensor<?x?xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-+  func.return %0 : tensor<?x?xf32>
-+}
-+
-+// CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<2x2xi64>) -> tensor<1x?x?x16xf32> {
-+  //      CHECK: "vhlo.dynamic_conv_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
-+  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
-+  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<false> : tensor<2xi1>>,
-+  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x2x!vhlo.i64_v1>) -> !vhlo.tensor_v1<1x?x?x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_conv"(%arg0, %arg1, %arg2) {
-+    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
-+    feature_group_count = 1 : i64,
-+    batch_group_count = 1 : i64
-+  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>, tensor<2x2xi64>) -> tensor<1x?x?x16xf32>
-+  func.return %0 : tensor<1x?x?x16xf32>
-+}
-+
-+// CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
-+  //      CHECK: "vhlo.dynamic_gather_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-+  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>, !vhlo.tensor_v1<3x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_gather"(%arg0, %arg1, %arg2) {
-+    dimension_numbers = #stablehlo.gather<
-+      offset_dims = [2],
-+      collapsed_slice_dims = [0, 1],
-+      start_index_map = [0, 1],
-+      index_vector_dim = 2
-+    >
-+  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<1x5x8xf32>
-+  func.return %0 : tensor<1x5x8xf32>
-+}
-+
-+func.func @default_func(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK:      "vhlo.func_v1"() <{
-+  // CHECK-SAME:   arg_attrs = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>>>,
-+  // CHECK-SAME:   res_attrs = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-+  // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : () -> ()
-+  func.return %arg0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "default_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
-+  //      CHECK: "vhlo.gather_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-+  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<1> : tensor<3xi64>>,
-+  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x1x!vhlo.f32_v1>
-+  %0 = "stablehlo.gather"(%arg0, %arg1) {
-+    dimension_numbers = #stablehlo.gather<
-+      offset_dims = [2],
-+      collapsed_slice_dims = [0, 1],
-+      start_index_map = [0, 1],
-+      index_vector_dim = 2
-+    >,
-+    slice_sizes = array<i64: 1, 1, 1>
-+  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>) -> tensor<1x5x1xf32>
-+  func.return %0 : tensor<1x5x1xf32>
-+}
-+
-+// CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-+  // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-+  //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-+  %0:2 = "stablehlo.infeed"(%arg0) : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
-+  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-+  %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-+  func.return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_recv_with_source_target_pairs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_recv_with_source_target_pairs(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
-+  //      CHECK: "vhlo.recv_v2"(%[[ARG0]]) <{
-+  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>,
-+  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>>
-+  // CHECK-SAME{LITERAL}: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-+  %0:2 = "stablehlo.recv"(%arg0) {
-+    channel_handle = #stablehlo.channel_handle<handle = 0, type = 1>,
-+    source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>
-+  } : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
-+  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
-+  //      CHECK: "vhlo.send_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>,
-+  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>>
-+  // CHECK-SAME{LITERAL}: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-+  %0 = "stablehlo.send"(%arg0, %arg1) {
-+    channel_handle = #stablehlo.channel_handle<handle = 0, type = 1>,
-+    source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>
-+  } : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-+  func.return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.reduce_scatter"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg1, %arg2) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    scatter_dimension = 0 : i64,
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
-+  } : (tensor<16xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-+  //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-+  // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-+  //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-+  //          CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-+  //          CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.maximum_v1"(%[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<2x17x31x7x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<2x16x30x7x!vhlo.f32_v1>
-+  %0 = "stablehlo.reduce_window"(%arg0, %arg1) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %1 = "stablehlo.maximum"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    window_dimensions = array<i64: 1, 2, 2, 1>
-+  } : (tensor<2x17x31x7xf32>, tensor<f32>) -> tensor<2x16x30x7xf32>
-+  func.return %0 : tensor<2x16x30x7xf32>
-+}
-+
-+// CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
-+  //      CHECK: "vhlo.scatter_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-+  // CHECK-SAME:   input_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   scatter_dims_to_operand_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   scatter_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   unique_indices = #vhlo.bool_v1<false>,
-+  // CHECK-SAME:   update_window_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>
-+  %0 = "stablehlo.scatter"(%arg0, %arg1, %arg2) ({
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    scatter_dimension_numbers = #stablehlo.scatter<
-+      update_window_dims = [1],
-+      inserted_window_dims = [0, 1],
-+      scatter_dims_to_operand_dims = [0, 1],
-+      index_vector_dim = 1
-+    >
-+  } : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) -> tensor<200x100x300xf32>
-+  func.return %0 : tensor<200x100x300xf32>
-+}
-+
-+// CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-+  // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-+  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG31:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG41:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL11:.*]] = "vhlo.compare_v1"(%[[ARG31]], %[[ARG41]]) <{compare_type = #vhlo<comparison_type_v1 TOTALORDER>, comparison_direction = #vhlo<comparison_direction_v1 GE>}>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL11]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-+  // CHECK-NEXT: }, {
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG32:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG42:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL12:.*]] = "vhlo.add_v1"(%[[ARG32]], %[[ARG42]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL12]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>, !vhlo.tensor_v1<10x23x23x64x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>
-+  %0 = "stablehlo.select_and_scatter"(%arg0, %arg1, %arg2) ({
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.compare"(%arg3, %arg4) {compare_type = #stablehlo<comparison_type TOTALORDER>, comparison_direction = #stablehlo<comparison_direction GE>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+      "stablehlo.return"(%1) : (tensor<i1>) -> ()
-+  }, {
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    window_dimensions = array<i64: 1, 2, 2, 1>
-+  } : (tensor<10x24x24x64xf32>, tensor<10x23x23x64xf32>, tensor<f32>) -> tensor<10x24x24x64xf32>
-+  func.return %0 : tensor<10x24x24x64xf32>
-+}
-+
-+// CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-+  // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.compare_v1"(%[[ARG1]], %[[ARG2]]) <{compare_type = #vhlo<comparison_type_v1 FLOAT>, comparison_direction = #vhlo<comparison_direction_v1 GT>}>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.sort"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
-+      %1 = "stablehlo.compare"(%arg1, %arg2) {compare_type = #stablehlo<comparison_type FLOAT>, comparison_direction = #stablehlo<comparison_direction GT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+      "stablehlo.return"(%1) : (tensor<i1>) -> ()
-+  }) : (tensor<16xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// ============ OPS ============
-+
-+// CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-+  func.return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
-+  //               CHECK: "vhlo.all_gather_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.all_gather"(%arg0) {
-+    all_gather_dim = 1 : i64,
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+    use_global_device_ids
-+  } : (tensor<16x8xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.all_reduce_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %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 = dense<[[0], [1]]> : tensor<2x1xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+    use_global_device_ids
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_all_reduce_with_promotable_types"
-+func.func @op_all_reduce_with_promotable_types(%operand: tensor<f32>) -> tensor<f64> {
-+  //  CHECK: "vhlo.all_reduce_v2"(%[[ARG0:.*]])
-+  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-+  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-+  //  CHECK: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  %result = "stablehlo.all_reduce"(%operand) ({
-+    ^bb0(%arg0: tensor<f64>, %arg1: tensor<f64>):
-+      %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-+      "stablehlo.return"(%0) : (tensor<f64>) -> ()
-+  }) {
-+    replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+    use_global_device_ids
-+  } : (tensor<f32>) -> tensor<f64>
-+
-+  func.return %result : tensor<f64>
-+}
-+
-+// CHECK-LABEL: "default_all_reduce_variadic"
-+func.func @default_all_reduce_variadic(%arg0: tensor<f32>, %arg1: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-+  %0:2 = "stablehlo.all_reduce"(%arg0, %arg1) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> (tensor<f32>)
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
-+  } : (tensor<f32>, tensor<f32>) -> (tensor<f32>, tensor<f32>)
-+  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
-+  //               CHECK: "vhlo.all_to_all_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-+  //          CHECK-SAME:   split_count = #vhlo.integer_v1<4 : i64>
-+  //          CHECK-SAME:   split_dimension = #vhlo.integer_v1<1 : i64>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x4x!vhlo.f32_v1>
-+  %0 = "stablehlo.all_to_all"(%arg0) {
-+    split_dimension = 1 : i64,
-+    concat_dimension = 0 : i64,
-+    split_count = 4 : i64,
-+    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>
-+  } : (tensor<4x16xf32>) -> tensor<16x4xf32>
-+  func.return %0 : tensor<16x4xf32>
-+}
-+
-+// CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
-+func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-+  // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-+  // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-+  %0:3 = "stablehlo.batch_norm_grad"(%arg0, %arg1, %arg2, %arg3, %arg4) {
-+    epsilon = 0.001 : f32,
-+    feature_index = 0 : i64
-+  } : (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>)
-+  func.return %0#0, %0#1, %0#2 : tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
-+func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-+  // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-+  // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.batch_norm_inference"(%arg0, %arg1, %arg2, %arg3, %arg4) {
-+    epsilon = 0.001 : f32,
-+    feature_index = 0 : i64
-+  } : (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>) -> tensor<16x16x16x16xf32>
-+  func.return %0 : tensor<16x16x16x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-+  // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-+  %0:3 = "stablehlo.batch_norm_training"(%arg0, %arg1, %arg2) {
-+    epsilon = 0.001 : f32,
-+    feature_index = 0 : i64
-+  } : (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>)
-+  func.return %0#0, %0#1, %0#2 : tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-+    broadcast_dimensions = array<i64: 1>
-+  } : (tensor<16xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.broadcast"(%arg0) {
-+    broadcast_sizes = array<i64: 16>
-+  } : (tensor<16xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.case"(%arg0) ({
-+    "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-+  }) : (tensor<i32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.cbrt_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.cholesky"(%arg0) {
-+    lower = true
-+  } : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-+  func.return %0 : tensor<1x16x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "op_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_collective_broadcast(%arg0: tensor<16x8xf32>, %arg1: tensor<1xi32>) -> tensor<16x8xf32> {
-+  //               CHECK: "vhlo.collective_broadcast_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<true>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.collective_broadcast"(%arg0, %arg1) {
-+    replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+    has_dynamic_root
-+  } : (tensor<16x8xf32>, tensor<1xi32>) -> tensor<16x8xf32>
-+  func.return %0 : tensor<16x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_collective_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_collective_reduce(%arg0: tensor<f32>, %arg1: tensor<1xi32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.collective_reduce_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<true>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.collective_reduce"(%arg0, %arg1) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+    use_global_device_ids,
-+    has_dynamic_root
-+  } : (tensor<f32>, tensor<1xi32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.collective_permute"(%arg0) {
-+    source_target_pairs = dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>
-+  } : (tensor<16x8xf32>) -> tensor<16x8xf32>
-+  func.return %0 : tensor<16x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-+  // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>,
-+    compare_type = #stablehlo<comparison_type TOTALORDER>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-+  func.return %0 : tensor<complex<f32>>
-+}
-+
-+// CHECK-LABEL: "op_composite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_composite(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.composite_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{#vhlo.string_v1<"my_int"> = #vhlo.integer_v1<1 : i64>, #vhlo.string_v1<"my_string"> = #vhlo.string_v1<"foo">}>
-+  //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
-+  //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
-+  //          CHECK-SAME:   version = #vhlo.integer_v1<1 : i32>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.composite"(%arg0) {
-+    name = "stablehlo.composite_target",
-+    decomposition = @composite_target,
-+    version = 1 : i32,
-+    composite_attributes = {
-+      my_int = 1 : i64,
-+      my_string = "foo"
-+    }
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "composite_regions"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @composite_regions(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.composite_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{}>
-+  //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
-+  //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
-+  //          CHECK-SAME:   version = #vhlo.integer_v1<1 : i32>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT: ^bb0(%[[ARG1:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.composite"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>):
-+      "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-+  }) {
-+    name = "stablehlo.composite_target",
-+    decomposition = @composite_target,
-+    version = 1 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-+    dimension = 0 : i64
-+  } : (tensor<8xf32>, tensor<8xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.constant_v1"() <{
-+  // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-+  // CHECK-SAME: }> : () -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.constant"() {
-+    value = dense<0.0> : tensor<f32>
-+  } : () -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
-+  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<2x2xi64>>,
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-+  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
-+  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<true> : tensor<2xi1>>,
-+  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x7x7x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.convolution"(%arg0, %arg1) {
-+    window_strides = array<i64: 2, 2>,
-+    padding = dense<1> : tensor<2x2xi64>,
-+    lhs_dilation = array<i64: 2, 2>,
-+    rhs_dilation = array<i64: 2, 2>,
-+    window_reversal = array<i1: true, true>,
-+    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
-+    feature_group_count = 1 : i64,
-+    batch_group_count = 1 : i64,
-+    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
-+  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32>
-+  func.return %0 : tensor<1x7x7x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.cosine_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_create_token"
-+func.func @op_create_token() -> !stablehlo.token {
-+  // CHECK: "vhlo.create_token_v1"() : () -> !vhlo.token_v1
-+  %0 = "stablehlo.create_token"() : () -> !stablehlo.token
-+  func.return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.cross-replica-sum"(%arg0) {
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
-+  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-+  // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-+  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-+  // CHECK-SAME:   called_computations = #vhlo.array_v1<[#vhlo.string_v1<"foo">]>,
-+  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<> : tensor<0xindex>>]>,
-+  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[
-+  // CHECK-SAME:     #vhlo.output_operand_alias_v1<
-+  // CHECK-SAME:       outputTupleIndices = [],
-+  // CHECK-SAME:       operandIndex = 0,
-+  // CHECK-SAME:       operandTupleIndices = []>]>,
-+  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<> : tensor<0xindex>>]>,
-+  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[]>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    has_side_effect = true,
-+    backend_config = "\08\03\1A\02",
-+    api_version = 2 : i32,
-+    called_computations = [@foo],
-+    operand_layouts = [dense<> : tensor<0xindex>],
-+    output_operand_aliases = [
-+      #stablehlo.output_operand_alias<output_tuple_indices = [],
-+                                 operand_index = 0,
-+                                 operand_tuple_indices = []>],
-+    result_layouts = [dense<> : tensor<0xindex>]
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_custom_call_empty_result_layout"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func public @op_custom_call_empty_result_layout(%arg0: tensor<i64>) -> tensor<i64> {
-+  // %0 = "vhlo.custom_call_v2"(%arg0) <{>}> : (!vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tuple_v1<>
-+  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
-+  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-+  // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-+  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"empty_output">,
-+  // CHECK-SAME:   called_computations = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<> : tensor<0xindex>>]>,
-+  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[]>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tuple_v1<>
-+  %0 = "stablehlo.custom_call"(%arg0) <{
-+    api_version = 2 : i32,
-+    call_target_name = "empty_output",
-+    has_side_effect = true,
-+    operand_layouts = [dense<> : tensor<0xindex>],
-+    result_layouts = []
-+  }> : (tensor<i64>) -> tuple<>
-+  return %arg0 : tensor<i64>
-+}
-+
-+// CHECK-LABEL: "op_custom_call_with_result_tilings"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_custom_call_with_result_tilings(%arg0: tensor<64x256xf32>) -> tensor<64x256xf32> {
-+  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
-+  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-+  // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-+  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-+  // CHECK-SAME:   called_computations = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<false>,
-+  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<[1, 0]> : tensor<2xindex>>]>,
-+  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<[1, 0]> : tensor<2xindex>>]>,
-+  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[#vhlo.array_v1<[#vhlo.tensor_v1<dense<[8, 128]> : tensor<2xindex>>]>]>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<64x256x!vhlo.f32_v1>) -> !vhlo.tensor_v1<64x256x!vhlo.f32_v1>
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    api_version = 2 : i32,
-+    call_target_name = "foo",
-+    operand_layouts = [dense<[1, 0]> : tensor<2xindex>],
-+    result_layouts = [dense<[1, 0]> : tensor<2xindex>],
-+    result_tilings = [[dense<[8, 128]> : tensor<2xindex>]]
-+  } : (tensor<64x256xf32>) -> tensor<64x256xf32>
-+  func.return %0 : tensor<64x256xf32>
-+}
-+
-+// CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
-+  //      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   allow_imprecise_accumulation = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   lhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   num_primitive_operations = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-+  // CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   rhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
-+  // CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
-+    dot_dimension_numbers = #stablehlo.dot<
-+      lhs_batching_dimensions = [0],
-+      lhs_contracting_dimensions = [2],
-+      rhs_batching_dimensions = [0],
-+      rhs_contracting_dimensions = [1]
-+    >,
-+    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
-+  } : (tensor<8x8x16xf32>, tensor<8x16x8xf32>) -> tensor<8x8x8xf32>
-+  func.return %0 : tensor<8x8x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot"(%arg0, %arg1) {
-+    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
-+  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-+  func.return %0 : tensor<8x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_broadcast_in_dim"(%arg0, %arg1) {
-+    broadcast_dimensions = array<i64: 0, 1>,
-+    known_expanding_dimensions = array<i64: 0>,
-+    known_nonexpanding_dimensions = array<i64: 1>
-+  } : (tensor<?x?xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-+  func.return %0 : tensor<?x?xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<2x2xi64>) -> tensor<1x?x?x16xf32> {
-+  //      CHECK: "vhlo.dynamic_conv_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
-+  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-+  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
-+  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<true> : tensor<2xi1>>,
-+  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x2x!vhlo.i64_v1>) -> !vhlo.tensor_v1<1x?x?x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_conv"(%arg0, %arg1, %arg2) {
-+    window_strides = array<i64: 2, 2>,
-+    lhs_dilation = array<i64: 2, 2>,
-+    rhs_dilation = array<i64: 2, 2>,
-+    window_reversal = array<i1: true, true>,
-+    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
-+    feature_group_count = 1 : i64,
-+    batch_group_count = 1 : i64,
-+    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
-+  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>, tensor<2x2xi64>) -> tensor<1x?x?x16xf32>
-+  func.return %0 : tensor<1x?x?x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
-+  //      CHECK: "vhlo.dynamic_gather_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>, !vhlo.tensor_v1<3x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_gather"(%arg0, %arg1, %arg2) {
-+    dimension_numbers = #stablehlo.gather<
-+      offset_dims = [2],
-+      collapsed_slice_dims = [0, 1],
-+      start_index_map = [0, 1],
-+      index_vector_dim = 2
-+    >,
-+    indices_are_sorted = true
-+  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<1x5x8xf32>
-+  func.return %0 : tensor<1x5x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_gather_with_batching_dims"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_dynamic_gather_with_batching_dims(%arg0 : tensor<5x2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<4xi32>) -> tensor<1x5x8xf32> {
-+  //      CHECK: "vhlo.dynamic_gather_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<5x2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>, !vhlo.tensor_v1<4x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_gather"(%arg0, %arg1, %arg2) {
-+    dimension_numbers = #stablehlo.gather<
-+      offset_dims = [2],
-+      collapsed_slice_dims = [1, 2],
-+      operand_batching_dims = [0],
-+      start_indices_batching_dims = [1],
-+      start_index_map = [1, 2],
-+      index_vector_dim = 2
-+    >,
-+    indices_are_sorted = true
-+  } : (tensor<5x2x4x9xf32>, tensor<1x5x2xi32>, tensor<4xi32>) -> tensor<1x5x8xf32>
-+  func.return %0 : tensor<1x5x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_iota"(%arg0) {
-+    iota_dimension = 0 : i64
-+  } : (tensor<1xindex>) -> tensor<?xf32>
-+  func.return %0 : tensor<?xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
-+func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-+  func.return %0 : tensor<?xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-+  func.return %0 : tensor<?x?xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-+    slice_sizes = array<i64: 4>
-+  } : (tensor<16xf32>, tensor<i64>) -> tensor<4xf32>
-+  func.return %0 : tensor<4xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.einsum"(%arg0, %arg1) {
-+    einsum_config = "ab,bc->ac"
-+  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-+  func.return %0 : tensor<8x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.exponential_minus_one_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.exponential_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-+  // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-+  %0 = "stablehlo.fft"(%arg0) {
-+    fft_type = #stablehlo<fft_type FFT>,
-+    fft_length = array<i64: 16>
-+  } : (tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>>
-+  func.return %0 : tensor<16xcomplex<f32>>
-+}
-+
-+// CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+func.func private @op_func(%arg0: tensor<f32> {stablehlo.arg = "0"}) -> (tensor<f32> {stablehlo.result = "0"}) {
-+  // CHECK:      "vhlo.func_v1"() <{
-+  // CHECK-SAME:   arg_attrs = #vhlo.array_v1<[#vhlo.dict_v1<{#vhlo.string_v1<"stablehlo.arg"> = #vhlo.string_v1<"0">}>]>,
-+  // CHECK-SAME:   function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>>>,
-+  // CHECK-SAME:   res_attrs = #vhlo.array_v1<[#vhlo.dict_v1<{#vhlo.string_v1<"stablehlo.result"> = #vhlo.string_v1<"0">}>]>,
-+  // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-+  // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : () -> ()
-+
-+  func.return %arg0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
-+  //      CHECK: "vhlo.gather_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<1> : tensor<3xi64>>,
-+  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x1x!vhlo.f32_v1>
-+  %0 = "stablehlo.gather"(%arg0, %arg1) {
-+    dimension_numbers = #stablehlo.gather<
-+      offset_dims = [2],
-+      collapsed_slice_dims = [0, 1],
-+      start_index_map = [0, 1],
-+      index_vector_dim = 2
-+    >,
-+    slice_sizes = array<i64: 1, 1, 1>,
-+    indices_are_sorted = true
-+  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>) -> tensor<1x5x1xf32>
-+  func.return %0 : tensor<1x5x1xf32>
-+}
-+
-+// CHECK-LABEL: "op_gather_with_batching_dims"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_gather_with_batching_dims(%arg0 : tensor<5x2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
-+  //      CHECK: "vhlo.gather_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-+  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<5x2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x1x!vhlo.f32_v1>
-+  %0 = "stablehlo.gather"(%arg0, %arg1) {
-+    dimension_numbers = #stablehlo.gather<
-+      offset_dims = [2],
-+      collapsed_slice_dims = [1, 2],
-+      operand_batching_dims = [0],
-+      start_indices_batching_dims = [1],
-+      start_index_map = [1, 2],
-+      index_vector_dim = 2
-+    >,
-+    slice_sizes = array<i64: 1, 1, 1, 1>,
-+    indices_are_sorted = true
-+  } : (tensor<5x2x4x9xf32>, tensor<1x5x2xi32>) -> tensor<1x5x1xf32>
-+  func.return %0 : tensor<1x5x1xf32>
-+}
-+
-+// CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.get_dimension_size"(%arg0) {
-+    dimension = 0 : i64
-+  } : (tensor<?xf32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-+  // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.get_tuple_element"(%arg0) {
-+    index = 0 : i32
-+  } : (tuple<tensor<f32>, tensor<i32>>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }, {
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.if"(%arg0) ({
-+    "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-+  }, {
-+    "stablehlo.return"(%arg2) : (tensor<f32>) -> ()
-+  }) : (tensor<i1>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-+  // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-+  //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-+  %0:2 = "stablehlo.infeed"(%arg0) {
-+    infeed_config = "foo",
-+    layout = [[]]
-+  } : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
-+  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_iota"
-+func.func @op_iota() -> tensor<16xf32> {
-+  //      CHECK: "vhlo.iota_v1"() <{
-+  // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : () -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.iota"() {
-+    iota_dimension = 0 : i64
-+  } : () -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.log_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.log_plus_one_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.logistic_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.abs_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.map"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>):
-+      %1 = "stablehlo.abs"(%arg1) : (tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    dimensions = array<i64: 0>
-+  } : (tensor<16xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-+  %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-+    outfeed_config = "foo"
-+  } : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-+  func.return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-+  // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-+  // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.pad"(%arg0, %arg1) {
-+    edge_padding_high = array<i64: 4>,
-+    edge_padding_low = array<i64: 4>,
-+    interior_padding = array<i64: 0>
-+  } : (tensor<8xf32>, tensor<f32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
-+func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-+  func.return %0 : tensor<?xf32>
-+}
-+
-+// CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_recv_no_source_target_pairs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_recv_no_source_target_pairs(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
-+  //      CHECK: "vhlo.recv_v2"(%[[ARG0]]) <{
-+  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>,
-+  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<> : tensor<0xi64>
-+  // CHECK-SAME{LITERAL}: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-+  %0:2 = "stablehlo.recv"(%arg0) {
-+    channel_handle = #stablehlo.channel_handle<handle = 0, type = 3>,
-+    is_host_transfer = true
-+  } : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
-+  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-+  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.reduce"(%arg0, %arg1) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    dimensions = array<i64: 0>
-+  } : (tensor<16xf32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-+  // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.reduce_precision"(%arg0) {
-+    exponent_bits = 8 : i32,
-+    mantissa_bits = 10 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK_lABEL: "op_reduce_with_promotable_types"
-+func.func @op_reduce_with_promotable_types(%arg0: tensor<4x4xf32>, %arg1 : tensor<f32>)
-+    -> (tensor<4xf64>) {
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0:.*]], %[[ARG1:.*]])
-+  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-+  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-+  //  CHECK: }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x!vhlo.f64_v1>
-+  %0 = "stablehlo.reduce"(%arg0, %arg1) ({
-+  ^bb0(%arg2: tensor<f64>, %arg3: tensor<f64> ):
-+    %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-+    "stablehlo.return"(%1) : (tensor<f64>) -> ()
-+
-+  }) {dimensions = array<i64: 0>} : (tensor<4x4xf32>, tensor<f32>) -> tensor<4xf64>
-+
-+  func.return %0: tensor<4xf64>
-+}
-+
-+// CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.reduce_scatter"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg1, %arg2) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    scatter_dimension = 0 : i64,
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+    use_global_device_ids
-+  } : (tensor<16xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK_lABEL: "op_reduce_scatter_with_promotable_types"
-+func.func @op_reduce_scatter_with_promotable_types(%data: tensor<4x16xf32>) -> tensor<4x4xf64> {
-+  //  CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0:.*]])
-+  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-+  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-+  //  CHECK: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f64_v1>
-+  %0 = "stablehlo.reduce_scatter"(%data) ({
-+    ^bb0(%arg2: tensor<f64>, %arg3: tensor<f64>):
-+    %1 = stablehlo.add %arg2, %arg3 : tensor<f64>
-+    "stablehlo.return"(%1) : (tensor<f64>) -> ()
-+  }) {replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
-+      scatter_dimension = 1 : i64,
-+      channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+      use_global_device_ids} : (tensor<4x16xf32>) -> tensor<4x4xf64>
-+  func.return %0 : tensor<4x4xf64>
-+}
-+
-+
-+// CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-+  // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-+  //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-+  //          CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-+  //          CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 4, 4, 1]> : tensor<4xi64>>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.maximum_v1"(%[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<2x17x31x7x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<2x9x16x7x!vhlo.f32_v1>
-+  %0 = "stablehlo.reduce_window"(%arg0, %arg1) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %1 = "stablehlo.maximum"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    window_dimensions = array<i64: 1, 2, 2, 1>,
-+    window_strides = array<i64: 1, 4, 4, 1>,
-+    base_dilations = array<i64: 1, 2, 2, 1>,
-+    window_dilations = array<i64: 1, 2, 2, 1>,
-+    padding = dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>
-+  } : (tensor<2x17x31x7xf32>, tensor<f32>) -> tensor<2x9x16x7xf32>
-+  func.return %0 : tensor<2x9x16x7xf32>
-+}
-+
-+// CHECK-LABEL: "op_reduce_window_with_promotable_types"
-+func.func @op_reduce_window_with_promotable_types(%arg0: tensor<4x2xf32>,
-+    %arg1: tensor<4x2xf32>, %init0: tensor<f32>, %init1: tensor<f32>) ->
-+    (tensor<2x2xf64>, tensor<2x2xf32>) {
-+  //  CHECK: "vhlo.reduce_window_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]], %[[ARG3:.*]])
-+  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]], %[[VAL2:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //  CHECK: }) : (!vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<2x2x!vhlo.f64_v1>, !vhlo.tensor_v1<2x2x!vhlo.f32_v1>)
-+  %0:2 = "stablehlo.reduce_window"(%arg0, %arg1, %init0, %init1) ({
-+         ^bb0(%a0: tensor<f64>, %a1: tensor<f32>, %b0: tensor<f64>,
-+                %b1: tensor<f32>):
-+              %2 = stablehlo.add %a0, %b0 : tensor<f64>
-+              %3 = stablehlo.add %a1, %b1 : tensor<f32>
-+              "stablehlo.return"(%2,%3) : (tensor<f64>, tensor<f32>) -> ()
-+            })
-+         { padding = dense<[[2, 2], [0, 0]]> : tensor<2x2xi64>,
-+         window_dimensions = array<i64: 5, 1>,
-+         window_strides = array<i64: 3, 1> }
-+         : (tensor<4x2xf32>, tensor<4x2xf32>, tensor<f32>, tensor<f32>) ->
-+              (tensor<2x2xf64>, tensor<2x2xf32>)
-+  func.return %0#0, %0#1 : tensor<2x2xf64>, tensor<2x2xf32>
-+}
-+
-+// CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_replica_id"
-+func.func @op_replica_id() -> tensor<ui32> {
-+  // CHECK: "vhlo.replica_id_v1"() : () -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  %0 = "stablehlo.replica_id"() : () -> tensor<ui32>
-+  func.return %0 : tensor<ui32>
-+}
-+
-+// CHECK-LABEL: "op_partition_id"
-+func.func @op_partition_id() -> tensor<ui32> {
-+  // CHECK: "vhlo.partition_id_v1"() : () -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  %0 = "stablehlo.partition_id"() : () -> tensor<ui32>
-+  func.return %0 : tensor<ui32>
-+}
-+
-+// CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-+  func.return %0 : tensor<4x4xf32>
-+}
-+
-+// CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.case"(%arg0) ({
-+    "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-+  }) : (tensor<i32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.reverse"(%arg0) {
-+    dimensions = array<i64: 0>
-+  } : (tensor<16xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-+  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-+    rng_algorithm = #stablehlo<rng_algorithm PHILOX>
-+  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
-+  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-+    rng_distribution = #stablehlo<rng_distribution NORMAL>
-+  } : (tensor<f32>, tensor<f32>, tensor<0xindex>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.rsqrt_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
-+  //      CHECK: "vhlo.scatter_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   input_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   scatter_dims_to_operand_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   scatter_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   unique_indices = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   update_window_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>
-+  %0 = "stablehlo.scatter"(%arg0, %arg1, %arg2) ({
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    scatter_dimension_numbers = #stablehlo.scatter<
-+      update_window_dims = [1],
-+      inserted_window_dims = [0, 1],
-+      scatter_dims_to_operand_dims = [0, 1],
-+      index_vector_dim = 1
-+    >,
-+    indices_are_sorted = true,
-+    unique_indices = true
-+  } : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) -> tensor<200x100x300xf32>
-+  func.return %0 : tensor<200x100x300xf32>
-+}
-+
-+// CHECK-LABEL: "op_scatter_with_batching_dims"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_scatter_with_batching_dims(%arg0: tensor<10x200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<10x200x100x300xf32> {
-+  //      CHECK: "vhlo.scatter_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   input_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   scatter_dims_to_operand_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   scatter_indices_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   unique_indices = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   update_window_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<10x200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x200x100x300x!vhlo.f32_v1>
-+  %0 = "stablehlo.scatter"(%arg0, %arg1, %arg2) ({
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    scatter_dimension_numbers = #stablehlo.scatter<
-+      update_window_dims = [1],
-+      inserted_window_dims = [1, 2],
-+      input_batching_dims = [0],
-+      scatter_dims_to_operand_dims = [1, 2],
-+      scatter_indices_batching_dims = [0],
-+      index_vector_dim = 1
-+    >,
-+    indices_are_sorted = true,
-+    unique_indices = true
-+  } : (tensor<10x200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) -> tensor<10x200x100x300xf32>
-+  func.return %0 : tensor<10x200x100x300xf32>
-+}
-+
-+// CHECK_lABEL: "op_scatter_with_promotable_types"
-+func.func @op_scatter_with_promotable_types(%input_tensor: tensor<200x100x300xf32>,
-+    %scatter_indices: tensor<10x2xi32>, %updates: tensor<10x300xf32>) ->
-+      tensor<200x100x300xf64> {
-+  //  CHECK: "vhlo.scatter_v2"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]])
-+  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-+  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-+  //  CHECK: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f64_v1>
-+  %0 = "stablehlo.scatter" (%input_tensor, %scatter_indices, %updates) ({
-+  ^bb0(%lhs: tensor<f64>, %rhs: tensor<f64>):
-+    %add = stablehlo.add %lhs, %rhs : tensor<f64>
-+    "stablehlo.return"(%add) : (tensor<f64>) -> ()
-+  }) {
-+    scatter_dimension_numbers = #stablehlo.scatter<
-+      update_window_dims = [1],
-+      inserted_window_dims = [0, 1],
-+      scatter_dims_to_operand_dims = [0, 1],
-+      index_vector_dim = 1
-+    >,
-+    indices_are_sorted = true,
-+    unique_indices = true
-+  } : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) ->
-+      tensor<200x100x300xf64>
-+  func.return %0 : tensor<200x100x300xf64>
-+}
-+
-+// CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-+  // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-+  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG31:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG41:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL11:.*]] = "vhlo.compare_v1"(%[[ARG31]], %[[ARG41]]) <{compare_type = #vhlo<comparison_type_v1 TOTALORDER>, comparison_direction = #vhlo<comparison_direction_v1 GE>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL11]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-+  // CHECK-NEXT: }, {
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG32:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG42:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL12:.*]] = "vhlo.add_v1"(%[[ARG32]], %[[ARG42]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL12]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>, !vhlo.tensor_v1<12x13x13x66x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>
-+  %0 = "stablehlo.select_and_scatter"(%arg0, %arg1, %arg2) ({
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.compare"(%arg3, %arg4) {compare_type = #stablehlo<comparison_type TOTALORDER>, comparison_direction = #stablehlo<comparison_direction GE>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+      "stablehlo.return"(%1) : (tensor<i1>) -> ()
-+  }, {
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    window_dimensions = array<i64: 1, 2, 2, 1>,
-+    window_strides = array<i64: 1, 2, 2, 1>,
-+    padding = dense<1> : tensor<4x2xi64>
-+  } : (tensor<10x24x24x64xf32>, tensor<12x13x13x66xf32>, tensor<f32>) -> tensor<10x24x24x64xf32>
-+  func.return %0 : tensor<10x24x24x64xf32>
-+}
-+
-+// CHECK-LABEL: "op_select_and_scatter_with_promotable_types"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_select_and_scatter_with_promotable_types(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf64> {
-+  // CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
-+  // CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-+  // CHECK:     %[[VAL:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK:     "vhlo.return_v1"(%[[VAL]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-+  // CHECK: }) : (!vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>, !vhlo.tensor_v1<12x13x13x66x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x24x24x64x!vhlo.f64_v1>
-+  %0 = "stablehlo.select_and_scatter"(%arg0, %arg1, %arg2) ({
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.compare"(%arg3, %arg4) {compare_type = #stablehlo<comparison_type TOTALORDER>, comparison_direction = #stablehlo<comparison_direction GE>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+      "stablehlo.return"(%1) : (tensor<i1>) -> ()
-+  }, {
-+    ^bb0(%arg3: tensor<f64>, %arg4: tensor<f64>):
-+      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-+      "stablehlo.return"(%1) : (tensor<f64>) -> ()
-+  }) {
-+    window_dimensions = array<i64: 1, 2, 2, 1>,
-+    window_strides = array<i64: 1, 2, 2, 1>,
-+    padding = dense<1> : tensor<4x2xi64>
-+  } : (tensor<10x24x24x64xf32>, tensor<12x13x13x66xf32>, tensor<f32>) -> tensor<10x24x24x64xf64>
-+  func.return %0 : tensor<10x24x24x64xf64>
-+}
-+
-+// CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_send_no_source_target_pairs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_send_no_source_target_pairs(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
-+  //      CHECK: "vhlo.send_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>,
-+  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-+  // CHECK-SAME{LITERAL}: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-+  %0 = "stablehlo.send"(%arg0, %arg1) {
-+    channel_handle = #stablehlo.channel_handle<handle = 0, type = 2>,
-+    is_host_transfer = true
-+  } : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-+  func.return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-+    dimension = 0 : i64
-+  } : (tensor<?xf32>, tensor<i32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.sine_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-+  // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-+  %0 = "stablehlo.slice"(%arg0) {
-+    start_indices = array<i64: 0>,
-+    limit_indices = array<i64: 4>,
-+    strides = array<i64: 1>
-+  } : (tensor<16xf32>) -> tensor<4xf32>
-+  func.return %0 : tensor<4xf32>
-+}
-+
-+// CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.compare_v1"(%[[ARG1]], %[[ARG2]]) <{compare_type = #vhlo<comparison_type_v1 FLOAT>, comparison_direction = #vhlo<comparison_direction_v1 GT>}>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.sort"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
-+      %1 = "stablehlo.compare"(%arg1, %arg2) {compare_type = #stablehlo<comparison_type FLOAT>, comparison_direction = #stablehlo<comparison_direction GT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+      "stablehlo.return"(%1) : (tensor<i1>) -> ()
-+  }) {
-+    dimension = 0 : i64,
-+    is_stable = true
-+  } : (tensor<16xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.sqrt_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_tan"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_tan(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.tan_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.tan"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.tanh_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-+  %0 = "stablehlo.torch_index_select"(%arg0, %arg1) {
-+    dim = 0 : i64,
-+    batch_dims = 0 : i64
-+  } : (tensor<5x1x5xf32>, tensor<2xi32>) -> tensor<2x1x5xf32>
-+  func.return %0 : tensor<2x1x5xf32>
-+}
-+
-+// CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.transpose"(%arg0) {
-+    permutation = array<i64: 1, 0>
-+  } : (tensor<16x8xf32>) -> tensor<8x16xf32>
-+  func.return %0 : tensor<8x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-+  // CHECK-SAME:   unit_diagonal = #vhlo.bool_v1<true>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-+    left_side = true,
-+    lower = true,
-+    unit_diagonal = true,
-+    transpose_a = #stablehlo<transpose NO_TRANSPOSE>
-+  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-+  func.return %0 : tuple<tensor<f32>>
-+}
-+
-+// CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-+  %0 = "stablehlo.unary_einsum"(%arg0) {
-+    einsum_config = "ab->a"
-+  } : (tensor<8x16xf32>) -> tensor<8xf32>
-+  func.return %0 : tensor<8xf32>
-+}
-+
-+// CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-+  func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
-+}
-+
-+// CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-+  // CHECK-NEXT:   }, {
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>)
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.while"(%arg0) ({
-+    ^bb0(%arg1: tensor<i1>):
-+      "stablehlo.return"(%arg1) : (tensor<i1>) -> ()
-+    }, {
-+    ^bb0(%arg1: tensor<i1>):
-+      "stablehlo.return"(%arg1) : (tensor<i1>) -> ()
-+  }) : (tensor<i1>) -> tensor<i1>
-+  func.return %0: tensor<i1>
-+}
-+
-+// CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// ============ TYPES ============
-+
-+// CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "type_i2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i2(%arg0: tensor<i2>, %arg1: tensor<i2>) -> tensor<i2> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i2_v1>, !vhlo.tensor_v1<!vhlo.i2_v1>) -> !vhlo.tensor_v1<!vhlo.i2_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i2>, tensor<i2>) -> tensor<i2>
-+  func.return %0 : tensor<i2>
-+}
-+
-+// CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-+  func.return %0 : tensor<i4>
-+}
-+
-+// CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-+  func.return %0 : tensor<i8>
-+}
-+
-+// CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-+  func.return %0 : tensor<i16>
-+}
-+
-+// CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-+  func.return %0 : tensor<i64>
-+}
-+
-+// CHECK-LABEL: "type_ui2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_ui2(%arg0: tensor<ui2>, %arg1: tensor<ui2>) -> tensor<ui2> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui2_v1>, !vhlo.tensor_v1<!vhlo.ui2_v1>) -> !vhlo.tensor_v1<!vhlo.ui2_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui2>, tensor<ui2>) -> tensor<ui2>
-+  func.return %0 : tensor<ui2>
-+}
-+
-+// CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-+  func.return %0 : tensor<ui4>
-+}
-+
-+// CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-+  func.return %0 : tensor<ui8>
-+}
-+
-+// CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-+  func.return %0 : tensor<ui16>
-+}
-+
-+// CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-+  func.return %0 : tensor<ui32>
-+}
-+
-+// CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-+  func.return %0 : tensor<ui64>
-+}
-+
-+// CHECK-LABEL: "type_f4E2M1FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f4E2M1FN(%arg0: tensor<f4E2M1FN>, %arg1: tensor<f4E2M1FN>) -> tensor<f4E2M1FN> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f4E2M1FN_v1>, !vhlo.tensor_v1<!vhlo.f4E2M1FN_v1>) -> !vhlo.tensor_v1<!vhlo.f4E2M1FN_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f4E2M1FN>, tensor<f4E2M1FN>) -> tensor<f4E2M1FN>
-+  func.return %0 : tensor<f4E2M1FN>
-+}
-+
-+// CHECK-LABEL: "type_f6E2M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f6E2M3FN(%arg0: tensor<f6E2M3FN>, %arg1: tensor<f6E2M3FN>) -> tensor<f6E2M3FN> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f6E2M3FN_v1>, !vhlo.tensor_v1<!vhlo.f6E2M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f6E2M3FN_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f6E2M3FN>, tensor<f6E2M3FN>) -> tensor<f6E2M3FN>
-+  func.return %0 : tensor<f6E2M3FN>
-+}
-+
-+// CHECK-LABEL: "type_f6E3M2FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f6E3M2FN(%arg0: tensor<f6E3M2FN>, %arg1: tensor<f6E3M2FN>) -> tensor<f6E3M2FN> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f6E3M2FN_v1>, !vhlo.tensor_v1<!vhlo.f6E3M2FN_v1>) -> !vhlo.tensor_v1<!vhlo.f6E3M2FN_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f6E3M2FN>, tensor<f6E3M2FN>) -> tensor<f6E3M2FN>
-+  func.return %0 : tensor<f6E3M2FN>
-+}
-+
-+// CHECK-LABEL: "type_f8E3M4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E3M4(%arg0: tensor<f8E3M4>, %arg1: tensor<f8E3M4>) -> tensor<f8E3M4> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E3M4_v1>, !vhlo.tensor_v1<!vhlo.f8E3M4_v1>) -> !vhlo.tensor_v1<!vhlo.f8E3M4_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E3M4>, tensor<f8E3M4>) -> tensor<f8E3M4>
-+  func.return %0 : tensor<f8E3M4>
-+}
-+
-+// CHECK-LABEL: "type_f8E4M3"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E4M3(%arg0: tensor<f8E4M3>, %arg1: tensor<f8E4M3>) -> tensor<f8E4M3> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3>, tensor<f8E4M3>) -> tensor<f8E4M3>
-+  func.return %0 : tensor<f8E4M3>
-+}
-+
-+// CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-+  func.return %0 : tensor<f8E4M3FN>
-+}
-+
-+// CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-+  func.return %0 : tensor<f8E5M2>
-+}
-+
-+// CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-+  func.return %0 : tensor<f8E4M3FNUZ>
-+}
-+
-+// CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-+  func.return %0 : tensor<f8E4M3B11FNUZ>
-+}
-+
-+// CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-+  func.return %0 : tensor<f8E5M2FNUZ>
-+}
-+
-+// CHECK-LABEL: "type_f8E8M0FNU"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E8M0FNU(%arg0: tensor<f8E8M0FNU>, %arg1: tensor<f8E8M0FNU>) -> tensor<f8E8M0FNU> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E8M0FNU_v1>, !vhlo.tensor_v1<!vhlo.f8E8M0FNU_v1>) -> !vhlo.tensor_v1<!vhlo.f8E8M0FNU_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E8M0FNU>, tensor<f8E8M0FNU>) -> tensor<f8E8M0FNU>
-+  func.return %0 : tensor<f8E8M0FNU>
-+}
-+
-+// CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-+  func.return %0 : tensor<bf16>
-+}
-+
-+// CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-+  func.return %0 : tensor<f16>
-+}
-+
-+// CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-+  func.return %0 : tensor<f64>
-+}
-+
-+// CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-+  func.return %0 : tensor<complex<f32>>
-+}
-+
-+// CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-+  func.return %0 : tensor<complex<f64>>
-+}
-+
-+// CHECK-LABEL: "type_tf32"
-+// CHECK: #vhlo.type_v1<!vhlo.tf31_v1>
-+func.func @type_tf32() attributes {stablehlo.attr = tf32 } {
-+  return
-+}
-+
-+// CHECK-LABEL: "type_none"
-+// CHECK: #vhlo.type_v1<!vhlo.none_v1>
-+func.func @type_none() attributes {stablehlo.attr = none } {
-+  return
-+}
-+
-+// CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-+  func.return %0 : tensor<?xf32>
-+}
-+
-+// CHECK-LABEL: "type_per_tensor_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_per_tensor_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-+  func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
-+}
-+
-+// CHECK-LABEL: "type_per_axis_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @type_per_axis_quantization(%arg0: tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>) -> tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG0]]) : (!vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, [3.400000e+01, 3.400000e+01], [16, 16], -128:127, 1>>, !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, [3.400000e+01, 3.400000e+01], [16, 16], -128:127, 1>>) -> !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, [3.400000e+01, 3.400000e+01], [16, 16], -128:127, 1>>
-+  %0 = stablehlo.add %arg0, %arg0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
-+  func.return %0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
-+}
-+
-+//       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
-+// CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-+  return %arg0 : !stablehlo.token
-+}
-+
-+//       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
-+// CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-+  %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-+  return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo"
-+  // CHECK: (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>) -> !vhlo.tuple_v1<!vhlo.token_v1>
-+  } : (tuple<tensor<f32>>) -> tuple<!stablehlo.token>
-+  return %0 : tuple<!stablehlo.token>
-+}
-+
-+// CHECK-LABEL: type_buffer_function_input_output
-+// CHECK-NEXT: (%[[ARG0:.*]]: !vhlo.buffer_v1<2x!vhlo.f32_v1>)
-+func.func @type_buffer_function_input_output(%arg0: memref<2xf32>) -> memref<2xf32> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.buffer_v1<2x!vhlo.f32_v1>) -> ()
-+  func.return %arg0 : memref<2xf32>
-+}
-+
-+// CHECK-LABEL: type_buffer_special_custom_calls
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @type_buffer_special_custom_calls(%arg0: tensor<2xf32>) -> tensor<2xf32> {
-+  //               CHECK: %[[CALL0:.*]] = "vhlo.custom_call_v2"(%[[ARG0]])
-+  //          CHECK-SAME: call_target_name = #vhlo.string_v1<"Pin">
-+  //          CHECK-SAME: : (!vhlo.tensor_v1<2x!vhlo.f32_v1>) -> !vhlo.buffer_v1<2x!vhlo.f32_v1>
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "Pin",
-+    api_version = 4 : i32
-+  } : (tensor<2xf32>) -> memref<2xf32>
-+  //               CHECK: %{{.*}} = "vhlo.custom_call_v2"(%[[CALL0]])
-+  //          CHECK-SAME: call_target_name = #vhlo.string_v1<"Unpin">
-+  //          CHECK-SAME: : (!vhlo.buffer_v1<2x!vhlo.f32_v1>) -> !vhlo.tensor_v1<2x!vhlo.f32_v1>
-+  %1 = "stablehlo.custom_call"(%0) {
-+    call_target_name = "Unpin",
-+    api_version = 4 : i32
-+  } : (memref<2xf32>) -> tensor<2xf32>
-+  func.return %1 : tensor<2xf32>
-+}
-+
-+// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>>>
-+// CHECK-LABEL: "type_future"
-+func.func private @type_future() -> !stablehlo.future<tensor<f32>>
-+
-+// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<2x3x!vhlo.f32_v1>>>>
-+// CHECK-LABEL: "type_future_ranked"
-+func.func private @type_future_ranked() -> !stablehlo.future<tensor<2x3xf32>>
-+
-+// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<?x!vhlo.f32_v1>>>>
-+// CHECK-LABEL: "type_future_dynamic"
-+func.func private @type_future_dynamic() -> !stablehlo.future<tensor<?xf32>>
-+
-+// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>>>
-+// CHECK-LABEL: "type_future_multiple"
-+func.func private @type_future_multiple() -> !stablehlo.future<tensor<f32>, tensor<i32>>
-+
-+// CHECK-LABEL: "op_async_start_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_start_all_reduce(%arg0: tensor<4x4xf32>) -> !stablehlo.future<tensor<4x4xf32>> {
-+  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
-+  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.all_reduce_v2"(%[[BARG0]]) <{
-+  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>,
-+  // CHECK-SAME:     use_global_device_ids = #vhlo.bool_v1<false>
-+  // CHECK-SAME:   }> ({
-+  // CHECK-NEXT:     ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:       %[[VAL2:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:       "vhlo.return_v1"(%[[VAL2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<4x4x!vhlo.f32_v1>>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<4x4xf32>):
-+    %1 = "stablehlo.all_reduce"(%barg0) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %2 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%2) : (tensor<f32>) -> ()
-+    }) {replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
-+    "stablehlo.return"(%1) : (tensor<4x4xf32>) -> ()
-+  }) : (tensor<4x4xf32>) -> !stablehlo.future<tensor<4x4xf32>>
-+  func.return %0: !stablehlo.future<tensor<4x4xf32>>
-+}
-+
-+// CHECK-LABEL: "op_async_start_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_start_all_gather(%arg0: tensor<8x2xf32>) -> !stablehlo.future<tensor<8x8xf32>> {
-+  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
-+  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.all_gather_v2"(%[[BARG0]]) <{
-+  // CHECK-SAME:     all_gather_dim = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 2, 4, 6], [1, 3, 5, 7]]> : tensor<2x4xi64>>,
-+  // CHECK-SAME:     use_global_device_ids = #vhlo.bool_v1<false>
-+  // CHECK-SAME:   }> : (!vhlo.tensor_v1<8x2x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<8x8x!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<8x2x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<8x8x!vhlo.f32_v1>>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<8x2xf32>):
-+    %1 = "stablehlo.all_gather"(%barg0) {
-+      all_gather_dim = 1 : i64,
-+      replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7]]> : tensor<2x4xi64>
-+    } : (tensor<8x2xf32>) -> tensor<8x8xf32>
-+    "stablehlo.return"(%1) : (tensor<8x8xf32>) -> ()
-+  }) : (tensor<8x2xf32>) -> !stablehlo.future<tensor<8x8xf32>>
-+  func.return %0: !stablehlo.future<tensor<8x8xf32>>
-+}
-+
-+// CHECK-LABEL: "op_async_start_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_start_all_to_all(%arg0: tensor<4x16xf32>) -> !stablehlo.future<tensor<16x4xf32>> {
-+  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
-+  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.all_to_all_v2"(%[[BARG0]]) <{
-+  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:     concat_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-+  // CHECK-SAME:     split_count = #vhlo.integer_v1<4 : i64>,
-+  // CHECK-SAME:     split_dimension = #vhlo.integer_v1<1 : i64>
-+  // CHECK-SAME:   }> : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x4x!vhlo.f32_v1>
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<16x4x!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<16x4x!vhlo.f32_v1>>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<4x16xf32>):
-+    %1 = "stablehlo.all_to_all"(%barg0) {
-+      split_dimension = 1 : i64,
-+      concat_dimension = 0 : i64,
-+      split_count = 4 : i64,
-+      replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>
-+    } : (tensor<4x16xf32>) -> tensor<16x4xf32>
-+    "stablehlo.return"(%1) : (tensor<16x4xf32>) -> ()
-+  }) : (tensor<4x16xf32>) -> !stablehlo.future<tensor<16x4xf32>>
-+  func.return %0: !stablehlo.future<tensor<16x4xf32>>
-+}
-+
-+// CHECK-LABEL: "op_async_start_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_start_collective_permute(%arg0: tensor<128x32xf32>) -> !stablehlo.future<tensor<128x32xf32>> {
-+  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
-+  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.collective_permute_v1"(%[[BARG0]]) <{
-+  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:     source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-+  // CHECK-SAME:   }> : (!vhlo.tensor_v1<128x32x!vhlo.f32_v1>) -> !vhlo.tensor_v1<128x32x!vhlo.f32_v1>
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<128x32x!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<128x32x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<128x32x!vhlo.f32_v1>>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<128x32xf32>):
-+    %1 = "stablehlo.collective_permute"(%barg0) {
-+      source_target_pairs = dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>
-+    } : (tensor<128x32xf32>) -> tensor<128x32xf32>
-+    "stablehlo.return"(%1) : (tensor<128x32xf32>) -> ()
-+  }) : (tensor<128x32xf32>) -> !stablehlo.future<tensor<128x32xf32>>
-+  func.return %0: !stablehlo.future<tensor<128x32xf32>>
-+}
-+
-+// CHECK-LABEL: "op_async_start_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_start_collective_broadcast(%arg0: tensor<16x8xf32>) -> !stablehlo.future<tensor<16x8xf32>> {
-+  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
-+  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.collective_broadcast_v2"(%[[BARG0]]) <{
-+  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:     has_dynamic_root = #vhlo.bool_v1<false>,
-+  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-+  // CHECK-SAME:   }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<16x8x!vhlo.f32_v1>>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<16x8xf32>):
-+    %1 = "stablehlo.collective_broadcast"(%barg0) {
-+      replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
-+    } : (tensor<16x8xf32>) -> tensor<16x8xf32>
-+    "stablehlo.return"(%1) : (tensor<16x8xf32>) -> ()
-+  }) : (tensor<16x8xf32>) -> !stablehlo.future<tensor<16x8xf32>>
-+  func.return %0: !stablehlo.future<tensor<16x8xf32>>
-+}
-+
-+// CHECK-LABEL: "op_async_start_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_start_reduce_scatter(%arg0: tensor<4x16xf32>) -> !stablehlo.future<tensor<4x4xf32>> {
-+  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
-+  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.reduce_scatter_v1"(%[[BARG0]]) <{
-+  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-+  // CHECK-SAME:     scatter_dimension = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:     use_global_device_ids = #vhlo.bool_v1<false>
-+  // CHECK-SAME:   }> ({
-+  // CHECK-NEXT:     ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:       %[[VAL2:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:       "vhlo.return_v1"(%[[VAL2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<4x4x!vhlo.f32_v1>>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<4x16xf32>):
-+    %1 = "stablehlo.reduce_scatter"(%barg0) ({
-+      ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %2 = stablehlo.add %arg2, %arg3 : tensor<f32>
-+      "stablehlo.return"(%2) : (tensor<f32>) -> ()
-+    }) {replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
-+        scatter_dimension = 1 : i64} : (tensor<4x16xf32>) -> tensor<4x4xf32>
-+    "stablehlo.return"(%1) : (tensor<4x4xf32>) -> ()
-+  }) : (tensor<4x16xf32>) -> !stablehlo.future<tensor<4x4xf32>>
-+  func.return %0: !stablehlo.future<tensor<4x4xf32>>
-+}
-+
-+// CHECK-LABEL: "op_async_done"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_done(%arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
-+  // CHECK: %[[VAL1:.*]] = "vhlo.async_start_v1"(%[[ARG0]])
-+  // CHECK: %[[VAL2:.*]] = "vhlo.async_done_v1"(%[[VAL1]]) : (!vhlo.future_v1<!vhlo.tensor_v1<4x4x!vhlo.f32_v1>>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<4x4xf32>):
-+    %1 = "stablehlo.all_reduce"(%barg0) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %2 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%2) : (tensor<f32>) -> ()
-+    }) {replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
-+    "stablehlo.return"(%1) : (tensor<4x4xf32>) -> ()
-+  }) : (tensor<4x4xf32>) -> !stablehlo.future<tensor<4x4xf32>>
-+  %1 = "stablehlo.async_done"(%0) : (!stablehlo.future<tensor<4x4xf32>>) -> tensor<4x4xf32>
-+  func.return %1: tensor<4x4xf32>
-+}
-+
-+// ============ DEPENDENCIES  ============
-+
-+func.func @composite_target(%arg0: tensor<f32>) -> tensor<f32> {
-+  return %arg0: tensor<f32>
-+}
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_22_0.mlir.bc b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_22_0.mlir.bc
-Binary files /dev/null and stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_22_0.mlir.bc differ
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_23_0.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_23_0.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_23_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_23_0.mlir
-@@ -0,0 +1,3452 @@
-+// RUN: stablehlo-opt --mlir-print-op-generic %s.bc | FileCheck %s
-+// RUN: stablehlo-translate --deserialize %s.bc | stablehlo-translate --serialize --target=1.23.0 | stablehlo-opt --mlir-print-op-generic | FileCheck %s
-+// RUN: stablehlo-translate --deserialize %s.bc | stablehlo-opt > %t.0
-+// RUN: stablehlo-opt --strip-debuginfo %s > %t.1
-+// RUN: diff %t.0 %t.1
-+// RUN: stablehlo-translate --serialize --target=1.23.0 --strip-debuginfo %s > %t.2
-+// RUN: diff %s.bc %t.2
-+// RUN: %if asserts %{ stablehlo-opt --stablehlo-legalize-to-vhlo -emit-bytecode -debug-only=vhlo-bytecode %s 2>&1 | FileCheck --check-prefix=CHECK-WARN %s %}
-+// RUN: %if asserts %{ stablehlo-opt --stablehlo-legalize-to-vhlo -emit-bytecode %s | stablehlo-opt -debug-only=vhlo-bytecode 2>&1 | FileCheck --check-prefix=CHECK-WARN %s %}
-+
-+// CHECK-WARN-NOT: Not Implemented
-+
-+func.func private @mesh()
-+
-+// ============ ATTRIBUTES ============
-+
-+// CHECK-LABEL: "attr_comparison_direction_eq"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_direction_eq(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 EQ>
-+    comparison_direction = #stablehlo<comparison_direction EQ>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_direction_ne"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_direction_ne(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 NE>
-+    comparison_direction = #stablehlo<comparison_direction NE>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_direction_ge"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_direction_ge(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GE>
-+    comparison_direction = #stablehlo<comparison_direction GE>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_direction_gt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_direction_gt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 GT>
-+    comparison_direction = #stablehlo<comparison_direction GT>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_direction_le"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_direction_le(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LE>
-+    comparison_direction = #stablehlo<comparison_direction LE>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_direction_lt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_direction_lt(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    // CHECK: comparison_direction = #vhlo<comparison_direction_v1 LT>
-+    comparison_direction = #stablehlo<comparison_direction LT>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_type_notype"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_type_notype(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>
-+    // CHECK: compare_type = #vhlo<comparison_type_v1 NOTYPE>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_type_float"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_type_float(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>,
-+    // CHECK: compare_type = #vhlo<comparison_type_v1 FLOAT>,
-+    compare_type = #stablehlo<comparison_type FLOAT>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_type_totalorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_type_totalorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>,
-+    // CHECK: compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-+    compare_type = #stablehlo<comparison_type TOTALORDER>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_type_signed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_type_signed(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>,
-+    // CHECK: compare_type = #vhlo<comparison_type_v1 SIGNED>,
-+    compare_type = #stablehlo<comparison_type SIGNED>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_type_unsigned"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_type_unsigned(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>,
-+    // CHECK: compare_type = #vhlo<comparison_type_v1 UNSIGNED>,
-+    compare_type = #stablehlo<comparison_type UNSIGNED>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "attr_comparison_type_weakorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_type_weakorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>,
-+    // CHECK: compare_type = #vhlo<comparison_type_v1 WEAKORDER>,
-+    compare_type = #stablehlo<comparison_type WEAKORDER>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// ConvDimensionNumbers aka #stablehlo.conv is covered below.
-+
-+// CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_custom_call_api_version_unspecified(%arg0: tensor<f32>) -> tensor<f32> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_UNSPECIFIED>
-+    api_version = 0 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_custom_call_api_version_original"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_custom_call_api_version_original(%arg0: tensor<f32>) -> tensor<f32> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>
-+    api_version = 1 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_custom_call_api_version_status_returning"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_custom_call_api_version_status_returning(%arg0: tensor<f32>) -> tensor<f32> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>
-+    api_version = 2 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_custom_call_api_version_status_returning_unified"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_custom_call_api_version_status_returning_unified(%arg0: tensor<f32>) -> tensor<f32> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    // CHECK: api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING_UNIFIED>
-+    api_version = 3 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_dict"
-+// CHECK: #vhlo.dict_v1<{#vhlo.string_v1<"attr1"> = #vhlo.integer_v1<1 : i32>, #vhlo.string_v1<"attr2"> = #vhlo.integer_v1<2 : i32>}
-+func.func @attr_dict() attributes {stablehlo.attr = {attr1 = 1 : i32, attr2 = 2 : i32}} {
-+  return
-+}
-+
-+// CHECK-LABEL: "attr_custom_call_api_version_typed_ffi"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+// CHECK: api_version = #vhlo<api_version_v1 API_VERSION_TYPED_FFI>
-+// CHECK-SAME: backend_config = #vhlo.dict_v1<{#vhlo.string_v1<"bar"> = #vhlo.integer_v1<42 : i32>}>
-+func.func @attr_custom_call_api_version_typed_ffi(%arg0: tensor<f32>) -> tensor<f32> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    backend_config= {bar = 42 : i32},
-+    api_version = 4 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+
-+// CHECK-LABEL: "attr_custom_call_api_version_typed_ffi_no_backend_config"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+// CHECK: api_version = #vhlo<api_version_v1 API_VERSION_TYPED_FFI>
-+// CHECK-SAME: backend_config = #vhlo.dict_v1<{}>
-+func.func @attr_custom_call_api_version_typed_ffi_no_backend_config(%arg0: tensor<f32>) -> tensor<f32> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    api_version = 4 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// DotDimensionNumbers aka #stablehlo.dot is covered below.
-+
-+// CHECK-LABEL: "attr_fft_type_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_fft_type_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-+  %0 = "stablehlo.fft"(%arg0) {
-+    // CHECK: fft_type = #vhlo<fft_type_v1 FFT>
-+    fft_type = #stablehlo<fft_type FFT>,
-+    fft_length = array<i64: 16>
-+  } : (tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>>
-+  func.return %0 : tensor<16xcomplex<f32>>
-+}
-+
-+// CHECK-LABEL: "attr_fft_type_ifft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_fft_type_ifft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-+  %0 = "stablehlo.fft"(%arg0) {
-+    // CHECK: fft_type = #vhlo<fft_type_v1 IFFT>
-+    fft_type = #stablehlo<fft_type IFFT>,
-+    fft_length = array<i64: 16>
-+  } : (tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>>
-+  func.return %0 : tensor<16xcomplex<f32>>
-+}
-+
-+// CHECK-LABEL: "attr_fft_type_rfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_fft_type_rfft(%arg0: tensor<16xf32>) -> tensor<9xcomplex<f32>> {
-+  %0 = "stablehlo.fft"(%arg0) {
-+    // CHECK: fft_type = #vhlo<fft_type_v1 RFFT>
-+    fft_type = #stablehlo<fft_type RFFT>,
-+    fft_length = array<i64: 16>
-+  } : (tensor<16xf32>) -> tensor<9xcomplex<f32>>
-+  func.return %0 : tensor<9xcomplex<f32>>
-+}
-+
-+// CHECK-LABEL: "attr_fft_type_irfft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_fft_type_irfft(%arg0: tensor<9xcomplex<f32>>) -> tensor<16xf32> {
-+  %0 = "stablehlo.fft"(%arg0) {
-+    // CHECK: fft_type = #vhlo<fft_type_v1 IRFFT>
-+    fft_type = #stablehlo<fft_type IRFFT>,
-+    fft_length = array<i64: 16>
-+  } : (tensor<9xcomplex<f32>>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "exponential_HIGHEST"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}
-+func.func @exponential_HIGHEST(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
-+  %0 = "stablehlo.exponential"(%arg0) {
-+    // CHECK: result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 HIGHEST>>
-+    result_accuracy = #stablehlo.result_accuracy<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #stablehlo.result_accuracy_mode<HIGHEST>>
-+  } : (tensor<8x16xf32>) -> tensor<8x16xf32>
-+  func.return %0 : tensor<8x16xf32>
-+}
-+
-+// CHECK-LABEL: "exponential_TOLERANCE"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}
-+func.func @exponential_TOLERANCE(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
-+  %0 = "stablehlo.exponential"(%arg0) {
-+    // CHECK: result_accuracy = #vhlo.result_accuracy_v1<atol = 1.000000e-05, rtol = 0.000000e+00, ulps = 1, mode = #vhlo<result_accuracy_mode_v1 TOLERANCE>>
-+    result_accuracy = #stablehlo.result_accuracy<atol = 1.000000e-5, rtol = 0.000000e+00, ulps = 1, mode = #stablehlo.result_accuracy_mode<TOLERANCE>>
-+  } : (tensor<8x16xf32>) -> tensor<8x16xf32>
-+  func.return %0 : tensor<8x16xf32>
-+}
-+
-+// CHECK-LABEL: "exponential_DEFAULT"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}
-+func.func @exponential_DEFAULT(%arg0: tensor<8x16xf32>) -> tensor<8x16xf32> {
-+  %0 = "stablehlo.exponential"(%arg0) {
-+    // CHECK: result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>
-+    result_accuracy = #stablehlo.result_accuracy<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #stablehlo.result_accuracy_mode<DEFAULT>>
-+  } : (tensor<8x16xf32>) -> tensor<8x16xf32>
-+  func.return %0 : tensor<8x16xf32>
-+}
-+
-+// GatherDimensionNumbers aka #stablehlo.gather is covered below.
-+
-+// CHECK-LABEL: "attr_precision_config_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_precision_config_default(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-+  %0 = "stablehlo.dot"(%arg0, %arg1) {
-+    // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>
-+  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-+  func.return %0 : tensor<8x8xf32>
-+}
-+
-+// CHECK-LABEL: "attr_precision_config_high"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_precision_config_high(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-+  %0 = "stablehlo.dot"(%arg0, %arg1) {
-+    // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGH>, #vhlo<precision_v1 HIGH>]>
-+    precision_config = [#stablehlo<precision HIGH>, #stablehlo<precision HIGH>]
-+  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-+  func.return %0 : tensor<8x8xf32>
-+}
-+
-+// CHECK-LABEL: "attr_precision_config_highest"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_precision_config_highest(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-+  %0 = "stablehlo.dot"(%arg0, %arg1) {
-+    // CHECK: precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-+    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
-+  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-+  func.return %0 : tensor<8x8xf32>
-+}
-+
-+// CHECK-LABEL: "attr_rng_algorithm_default"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_rng_algorithm_default(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-+  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-+    // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>
-+    rng_algorithm = #stablehlo<rng_algorithm DEFAULT>
-+  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
-+  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_rng_algorithm_three_fry"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_rng_algorithm_three_fry(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-+  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-+    // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 THREE_FRY>
-+    rng_algorithm = #stablehlo<rng_algorithm THREE_FRY>
-+  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
-+  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_rng_algorithm_philox"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_rng_algorithm_philox(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-+  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-+    // CHECK: rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-+    rng_algorithm = #stablehlo<rng_algorithm PHILOX>
-+  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
-+  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_rng_distribution_uniform"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @attr_rng_distribution_uniform(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-+  %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-+    // CHECK: rng_distribution = #vhlo<rng_distribution_v1 UNIFORM>
-+    rng_distribution = #stablehlo<rng_distribution UNIFORM>
-+  } : (tensor<f32>, tensor<f32>, tensor<0xindex>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_rng_distribution_normal"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @attr_rng_distribution_normal(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-+  %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-+    // CHECK: rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-+    rng_distribution = #stablehlo<rng_distribution NORMAL>
-+  } : (tensor<f32>, tensor<f32>, tensor<0xindex>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// ScatterDimensionNumbers aka #stablehlo.scatter is covered below.
-+
-+// CHECK-LABEL: "attr_transpose_no_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_transpose_no_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-+  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-+    left_side = true,
-+    lower = true,
-+    unit_diagonal = true,
-+    // transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-+    transpose_a = #stablehlo<transpose NO_TRANSPOSE>
-+  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "attr_transpose_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_transpose_transpose(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-+  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-+    left_side = true,
-+    lower = true,
-+    unit_diagonal = true,
-+    // transpose_a = #vhlo<transpose_v1 TRANSPOSE>,
-+    transpose_a = #stablehlo<transpose TRANSPOSE>
-+  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "attr_transpose_adjoint"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_transpose_adjoint(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-+  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-+    left_side = true,
-+    lower = true,
-+    unit_diagonal = true,
-+    // transpose_a = #vhlo<transpose_v1 ADJOINT>,
-+    transpose_a = #stablehlo<transpose ADJOINT>
-+  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// TypeExtensionsAttr aka #stablehlo.type_extensions is covered below.
-+
-+// CHECK-LABEL: "attr_type_extensions_bounds"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @attr_type_extensions_bounds(%arg0: tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>) -> tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1, #vhlo.type_extensions_v1<bounds = [16, ?]>>) -> ()
-+  func.return %arg0 : tensor<?x?xf32, #stablehlo.type_extensions<bounds = [16, ?]>>
-+}
-+
-+// CHECK-LABEL: "attr_frontend_attributes"
-+func.func @attr_frontend_attributes(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: some.unregistered_attr
-+  %1 = stablehlo.cosine %arg0 {some.unregistered_attr = 1 : i32} : tensor<f32>
-+  return %1 : tensor<f32>
-+}
-+
-+// Builtin attriubute tests
-+
-+// CHECK-LABEL: "byte_packed_boolean"
-+func.func @byte_packed_boolean() -> (tensor<8xi1>, tensor<8xi1>, tensor<4xi1>, tensor<4xi1>, tensor<16xi1>) {
-+  // CHECK: #vhlo.tensor_v1<dense<[true, false, false, false, false, false, false, false]
-+  // CHECK-NEXT: #vhlo.tensor_v1<dense<true> : tensor<8xi1>>
-+  // CHECK-NEXT: #vhlo.tensor_v1<dense<[true, false, false, false]> : tensor<4xi1>>
-+  // CHECK-NEXT: #vhlo.tensor_v1<dense<true> : tensor<4xi1>>
-+  // CHECK-NEXT: #vhlo.tensor_v1<dense<[true, false, false, false, false, false, false, false, true, false, false, false, false, false, false, false]> : tensor<16xi1>>
-+  %c = stablehlo.constant dense<[true, false, false, false, false, false, false, false]> : tensor<8xi1>
-+  %c_0 = stablehlo.constant dense<true> : tensor<8xi1>
-+  %c_1 = stablehlo.constant dense<[true, false, false, false]> : tensor<4xi1>
-+  %c_2 = stablehlo.constant dense<true> : tensor<4xi1>
-+  %c_3 = stablehlo.constant dense<[true, false, false, false, false, false, false, false, true, false, false, false, false, false, false, false]> : tensor<16xi1>
-+  return %c, %c_0, %c_1, %c_2, %c_3 : tensor<8xi1>, tensor<8xi1>, tensor<4xi1>, tensor<4xi1>, tensor<16xi1>
-+}
-+
-+// ============ DEFAULTS ============
-+
-+// CHECK-LABEL: "default_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
-+  //               CHECK: "vhlo.all_gather_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.all_gather"(%arg0) {
-+    all_gather_dim = 1 : i64,
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
-+  } : (tensor<16x8xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "default_all_gather_variadic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_all_gather_variadic(%arg0: tensor<16x8xf32>, %arg1: tensor<16x8xf32>) -> (tensor<16x16xf32>, tensor<16x16xf32>) {
-+  %0:2 = "stablehlo.all_gather"(%arg0, %arg1) {
-+    all_gather_dim = 1 : i64,
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
-+  } : (tensor<16x8xf32>, tensor<16x8xf32>) -> (tensor<16x16xf32>, tensor<16x16xf32>)
-+  func.return %0#0, %0#1 : tensor<16x16xf32>, tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "default_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.all_reduce_v2"(%[[ARG0]])
-+  //          CHECK-SAME: <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+
-+  %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 = dense<[[0], [1]]> : tensor<2x1xi64>
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "attr_replica_group_mesh_axes"
-+func.func @attr_replica_group_mesh_axes(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.all_reduce_v2"
-+  // CHECK-SAME: replica_groups = #vhlo.replica_group_mesh_axes_v1<
-+  // CHECK-SAME:   mesh = #vhlo.string_v1<"mesh">,
-+  // CHECK-SAME:   axes = #vhlo.array_v1<[#vhlo.axis_ref_v1<name = #vhlo.string_v1<"x">, sub_axis_info = #vhlo.sub_axis_info_v1<pre_size = 1, size = 2>>, #vhlo.axis_ref_v1<name = #vhlo.string_v1<"y">>]>
-+  // CHECK-SAME: >
-+  %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 = @mesh,
-+      axes = [
-+        #stablehlo.axis_ref<name = "x", sub_axis_info = (1)2>,
-+        #stablehlo.axis_ref<name = "y">
-+      ]
-+    >
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+
-+
-+// CHECK-LABEL: "default_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
-+  //               CHECK: "vhlo.all_to_all_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-+  //          CHECK-SAME:   split_count = #vhlo.integer_v1<4 : i64>
-+  //          CHECK-SAME:   split_dimension = #vhlo.integer_v1<1 : i64>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x4x!vhlo.f32_v1>
-+  %0 = "stablehlo.all_to_all"(%arg0) {
-+    split_dimension = 1 : i64,
-+    concat_dimension = 0 : i64,
-+    split_count = 4 : i64,
-+    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>
-+  } : (tensor<4x16xf32>) -> tensor<16x4xf32>
-+  func.return %0 : tensor<16x4xf32>
-+}
-+
-+// CHECK-LABEL: "default_all_to_all_variadic"
-+func.func @default_all_to_all_variadic(%arg0: tensor<4x16xf32>, %arg1: tensor<5x16xf32>) -> (tensor<16x4xf32>, tensor<20x4xf32>) {
-+  %0:2 = "stablehlo.all_to_all"(%arg0, %arg1) {
-+    split_dimension = 1 : i64,
-+    concat_dimension = 0 : i64,
-+    split_count = 4 : i64,
-+    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>
-+  } : (tensor<4x16xf32>, tensor<5x16xf32>) -> (tensor<16x4xf32>, tensor<20x4xf32>)
-+  func.return %0#0, %0#1 : tensor<16x4xf32>, tensor<20x4xf32>
-+}
-+
-+// CHECK-LABEL: "default_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   lower = #vhlo.bool_v1<false>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.cholesky"(%arg0) : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-+  func.return %0 : tensor<1x16x16xf32>
-+}
-+
-+// CHECK-LABEL: "default_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.collective_permute"(%arg0) {
-+    source_target_pairs = dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>
-+  } : (tensor<16x8xf32>) -> tensor<16x8xf32>
-+  func.return %0 : tensor<16x8xf32>
-+}
-+
-+// CHECK-LABEL: "default_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_collective_broadcast(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
-+  //               CHECK: "vhlo.collective_broadcast_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<false>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.collective_broadcast"(%arg0) {
-+    replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
-+  } : (tensor<16x8xf32>) -> tensor<16x8xf32>
-+  func.return %0 : tensor<16x8xf32>
-+}
-+
-+// CHECK-LABEL: "default_collective_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_collective_reduce(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.collective_reduce_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<false>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>,
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.collective_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 = dense<[[0, 1]]> : tensor<1x2xi64>
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "default_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 NOTYPE>,
-+  // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "default_composite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_composite(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.composite_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{}>
-+  //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
-+  //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
-+  //          CHECK-SAME:   version = #vhlo.integer_v1<0 : i64>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.composite"(%arg0) {
-+    name = "stablehlo.composite_target",
-+    decomposition = @composite_target
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "default_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32> {
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
-+  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<2x2xi64>>,
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
-+  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<false> : tensor<2xi1>>,
-+  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x6x6x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.convolution"(%arg0, %arg1) {
-+    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
-+    feature_group_count = 1 : i64,
-+    batch_group_count = 1 : i64
-+  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x6x6x16xf32>
-+  func.return %0 : tensor<1x6x6x16xf32>
-+}
-+
-+// CHECK-LABEL: "default_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
-+  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_ORIGINAL>,
-+  // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-+  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-+  // CHECK-SAME:   called_computations = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<false>,
-+  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[]>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo"
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "default_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
-+  //      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   allow_imprecise_accumulation = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   lhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   num_primitive_operations = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+  // CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   rhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
-+  // CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
-+    dot_dimension_numbers = #stablehlo.dot<
-+      lhs_batching_dimensions = [0],
-+      lhs_contracting_dimensions = [2],
-+      rhs_batching_dimensions = [0],
-+      rhs_contracting_dimensions = [1]
-+    >
-+  } : (tensor<8x8x16xf32>, tensor<8x16x8xf32>) -> tensor<8x8x8xf32>
-+  func.return %0 : tensor<8x8x8xf32>
-+}
-+
-+// CHECK-LABEL: "dot_general_algorithm"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @dot_general_algorithm(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
-+//      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
-+// CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.f32_v1>,
-+// CHECK-SAME:   allow_imprecise_accumulation = #vhlo.bool_v1<false>,
-+// CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.tf31_v1>,
-+// CHECK-SAME:   num_primitive_operations = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+// CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.tf31_v1>
-+// CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
-+    dot_dimension_numbers = #stablehlo.dot<
-+      lhs_batching_dimensions = [0],
-+      lhs_contracting_dimensions = [2],
-+      rhs_batching_dimensions = [0],
-+      rhs_contracting_dimensions = [1]
-+    >,
-+    algorithm = #stablehlo.dot_algorithm<
-+      lhs_precision_type = tf32,
-+      rhs_precision_type = tf32,
-+      accumulation_type = f32,
-+      lhs_component_count = 1,
-+      rhs_component_count = 1,
-+      num_primitive_operations = 1,
-+      allow_imprecise_accumulation = false
-+    >
-+  } : (tensor<8x8x16xf32>, tensor<8x16x8xf32>) -> tensor<8x8x8xf32>
-+  func.return %0 : tensor<8x8x8xf32>
-+}
-+
-+// CHECK-LABEL: "dot_general_algorithm_f8e4m3fn_x3"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @dot_general_algorithm_f8e4m3fn_x3(%arg0: tensor<8x8x16xbf16>, %arg1: tensor<8x16x8xbf16>) -> tensor<8x8x8xf32> {
-+//      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
-+// CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.f32_v1>,
-+// CHECK-SAME:   allow_imprecise_accumulation = #vhlo.bool_v1<false>,
-+// CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>,
-+// CHECK-SAME:   num_primitive_operations = #vhlo.integer_v1<3 : i64>,
-+// CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+// CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>
-+// CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.bf16_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.bf16_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
-+    dot_dimension_numbers = #stablehlo.dot<
-+      lhs_batching_dimensions = [0],
-+      lhs_contracting_dimensions = [2],
-+      rhs_batching_dimensions = [0],
-+      rhs_contracting_dimensions = [1]
-+    >,
-+    algorithm = #stablehlo.dot_algorithm<
-+      lhs_precision_type = f8E4M3FN,
-+      rhs_precision_type = f8E4M3FN,
-+      accumulation_type = f32,
-+      lhs_component_count = 1,
-+      rhs_component_count = 1,
-+      num_primitive_operations = 3,
-+      allow_imprecise_accumulation = false
-+    >
-+  } : (tensor<8x8x16xbf16>, tensor<8x16x8xbf16>) -> tensor<8x8x8xf32>
-+  func.return %0 : tensor<8x8x8xf32>
-+}
-+
-+// CHECK-LABEL: "dot_general_algorithm_f8e4m3fn_x4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @dot_general_algorithm_f8e4m3fn_x4(%arg0: tensor<8x8x16xbf16>, %arg1: tensor<8x16x8xbf16>) -> tensor<8x8x8xf32> {
-+//      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
-+// CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.f32_v1>,
-+// CHECK-SAME:   allow_imprecise_accumulation = #vhlo.bool_v1<false>,
-+// CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>,
-+// CHECK-SAME:   num_primitive_operations = #vhlo.integer_v1<4 : i64>,
-+// CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+// CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>
-+// CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.bf16_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.bf16_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
-+    dot_dimension_numbers = #stablehlo.dot<
-+      lhs_batching_dimensions = [0],
-+      lhs_contracting_dimensions = [2],
-+      rhs_batching_dimensions = [0],
-+      rhs_contracting_dimensions = [1]
-+    >,
-+    algorithm = #stablehlo.dot_algorithm<
-+      lhs_precision_type = f8E4M3FN,
-+      rhs_precision_type = f8E4M3FN,
-+      accumulation_type = f32,
-+      lhs_component_count = 1,
-+      rhs_component_count = 1,
-+      num_primitive_operations = 4,
-+      allow_imprecise_accumulation = false
-+    >
-+  } : (tensor<8x8x16xbf16>, tensor<8x16x8xbf16>) -> tensor<8x8x8xf32>
-+  func.return %0 : tensor<8x8x8xf32>
-+}
-+
-+// CHECK-LABEL: "default_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_broadcast_in_dim"(%arg0, %arg1) {
-+    broadcast_dimensions = array<i64: 0, 1>
-+  } : (tensor<?x?xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-+  func.return %0 : tensor<?x?xf32>
-+}
-+
-+// CHECK-LABEL: "default_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @default_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<2x2xi64>) -> tensor<1x?x?x16xf32> {
-+  //      CHECK: "vhlo.dynamic_conv_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
-+  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>,
-+  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<false> : tensor<2xi1>>,
-+  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<2xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x2x!vhlo.i64_v1>) -> !vhlo.tensor_v1<1x?x?x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_conv"(%arg0, %arg1, %arg2) {
-+    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
-+    feature_group_count = 1 : i64,
-+    batch_group_count = 1 : i64
-+  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>, tensor<2x2xi64>) -> tensor<1x?x?x16xf32>
-+  func.return %0 : tensor<1x?x?x16xf32>
-+}
-+
-+// CHECK-LABEL: "default_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @default_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
-+  //      CHECK: "vhlo.dynamic_gather_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-+  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>, !vhlo.tensor_v1<3x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_gather"(%arg0, %arg1, %arg2) {
-+    dimension_numbers = #stablehlo.gather<
-+      offset_dims = [2],
-+      collapsed_slice_dims = [0, 1],
-+      start_index_map = [0, 1],
-+      index_vector_dim = 2
-+    >
-+  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<1x5x8xf32>
-+  func.return %0 : tensor<1x5x8xf32>
-+}
-+
-+func.func @default_func(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK:      "vhlo.func_v1"() <{
-+  // CHECK-SAME:   arg_attrs = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>>>,
-+  // CHECK-SAME:   res_attrs = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   sym_name = #vhlo.string_v1<"default_func">,
-+  // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"">
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : () -> ()
-+  func.return %arg0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "default_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
-+  //      CHECK: "vhlo.gather_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-+  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<1> : tensor<3xi64>>,
-+  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x1x!vhlo.f32_v1>
-+  %0 = "stablehlo.gather"(%arg0, %arg1) {
-+    dimension_numbers = #stablehlo.gather<
-+      offset_dims = [2],
-+      collapsed_slice_dims = [0, 1],
-+      start_index_map = [0, 1],
-+      index_vector_dim = 2
-+    >,
-+    slice_sizes = array<i64: 1, 1, 1>
-+  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>) -> tensor<1x5x1xf32>
-+  func.return %0 : tensor<1x5x1xf32>
-+}
-+
-+// CHECK-LABEL: "default_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"">,
-+  // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[]>
-+  //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-+  %0:2 = "stablehlo.infeed"(%arg0) : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
-+  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "default_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"">
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-+  %0 = "stablehlo.outfeed"(%arg0, %arg1) : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-+  func.return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_recv_with_source_target_pairs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_recv_with_source_target_pairs(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
-+  //      CHECK: "vhlo.recv_v2"(%[[ARG0]]) <{
-+  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>,
-+  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>>
-+  // CHECK-SAME{LITERAL}: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-+  %0:2 = "stablehlo.recv"(%arg0) {
-+    channel_handle = #stablehlo.channel_handle<handle = 0, type = 1>,
-+    source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>
-+  } : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
-+  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "default_send"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_send(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
-+  //      CHECK: "vhlo.send_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   channel_type = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<false>,
-+  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>>
-+  // CHECK-SAME{LITERAL}: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-+  %0 = "stablehlo.send"(%arg0, %arg1) {
-+    channel_handle = #stablehlo.channel_handle<handle = 0, type = 1>,
-+    source_target_pairs = dense<[[0, 1], [1, 2]]> : tensor<2x2xi64>
-+  } : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-+  func.return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "default_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<false>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.reduce_scatter"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg1, %arg2) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    scatter_dimension = 0 : i64,
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
-+  } : (tensor<16xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "default_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @default_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x16x30x7xf32> {
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]])  <{
-+  //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-+  // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-+  //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-+  //          CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-+  //          CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.maximum_v1"(%[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<2x17x31x7x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<2x16x30x7x!vhlo.f32_v1>
-+  %0 = "stablehlo.reduce_window"(%arg0, %arg1) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %1 = "stablehlo.maximum"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    window_dimensions = array<i64: 1, 2, 2, 1>
-+  } : (tensor<2x17x31x7xf32>, tensor<f32>) -> tensor<2x16x30x7xf32>
-+  func.return %0 : tensor<2x16x30x7xf32>
-+}
-+
-+// CHECK-LABEL: "default_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @default_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
-+  //      CHECK: "vhlo.scatter_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<false>,
-+  // CHECK-SAME:   input_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   scatter_dims_to_operand_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   scatter_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   unique_indices = #vhlo.bool_v1<false>,
-+  // CHECK-SAME:   update_window_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>
-+  %0 = "stablehlo.scatter"(%arg0, %arg1, %arg2) ({
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    scatter_dimension_numbers = #stablehlo.scatter<
-+      update_window_dims = [1],
-+      inserted_window_dims = [0, 1],
-+      scatter_dims_to_operand_dims = [0, 1],
-+      index_vector_dim = 1
-+    >
-+  } : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) -> tensor<200x100x300xf32>
-+  func.return %0 : tensor<200x100x300xf32>
-+}
-+
-+// CHECK-LABEL: "default_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @default_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<10x23x23x64xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<0> : tensor<4x2xi64>>,
-+  // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-+  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG31:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG41:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL11:.*]] = "vhlo.compare_v1"(%[[ARG31]], %[[ARG41]]) <{compare_type = #vhlo<comparison_type_v1 TOTALORDER>, comparison_direction = #vhlo<comparison_direction_v1 GE>}>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL11]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-+  // CHECK-NEXT: }, {
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG32:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG42:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL12:.*]] = "vhlo.add_v1"(%[[ARG32]], %[[ARG42]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL12]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>, !vhlo.tensor_v1<10x23x23x64x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>
-+  %0 = "stablehlo.select_and_scatter"(%arg0, %arg1, %arg2) ({
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.compare"(%arg3, %arg4) {compare_type = #stablehlo<comparison_type TOTALORDER>, comparison_direction = #stablehlo<comparison_direction GE>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+      "stablehlo.return"(%1) : (tensor<i1>) -> ()
-+  }, {
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    window_dimensions = array<i64: 1, 2, 2, 1>
-+  } : (tensor<10x24x24x64xf32>, tensor<10x23x23x64xf32>, tensor<f32>) -> tensor<10x24x24x64xf32>
-+  func.return %0 : tensor<10x24x24x64xf32>
-+}
-+
-+// CHECK-LABEL: "default_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @default_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   dimension = #vhlo.integer_v1<-1 : i64>
-+  // CHECK-SAME:   is_stable = #vhlo.bool_v1<false>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.compare_v1"(%[[ARG1]], %[[ARG2]]) <{compare_type = #vhlo<comparison_type_v1 FLOAT>, comparison_direction = #vhlo<comparison_direction_v1 GT>}>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.sort"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
-+      %1 = "stablehlo.compare"(%arg1, %arg2) {compare_type = #stablehlo<comparison_type FLOAT>, comparison_direction = #stablehlo<comparison_direction GT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+      "stablehlo.return"(%1) : (tensor<i1>) -> ()
-+  }) : (tensor<16xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// ============ OPS ============
-+
-+// CHECK-LABEL: "op_abs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_abs(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.abs"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_add"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_add(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_after_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_after_all(%arg0: !stablehlo.token) -> !stablehlo.token {
-+  // CHECK: "vhlo.after_all_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> !vhlo.token_v1
-+  %0 = "stablehlo.after_all"(%arg0) : (!stablehlo.token) -> !stablehlo.token
-+  func.return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_all_gather(%arg0: tensor<16x8xf32>) -> tensor<16x16xf32> {
-+  //               CHECK: "vhlo.all_gather_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   all_gather_dim = #vhlo.integer_v1<1 : i64>
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.all_gather"(%arg0) {
-+    all_gather_dim = 1 : i64,
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+    use_global_device_ids
-+  } : (tensor<16x8xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_all_reduce(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.all_reduce_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %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 = dense<[[0], [1]]> : tensor<2x1xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+    use_global_device_ids
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_all_reduce_with_promotable_types"
-+func.func @op_all_reduce_with_promotable_types(%operand: tensor<f32>) -> tensor<f64> {
-+  //  CHECK: "vhlo.all_reduce_v2"(%[[ARG0:.*]])
-+  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-+  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-+  //  CHECK: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  %result = "stablehlo.all_reduce"(%operand) ({
-+    ^bb0(%arg0: tensor<f64>, %arg1: tensor<f64>):
-+      %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-+      "stablehlo.return"(%0) : (tensor<f64>) -> ()
-+  }) {
-+    replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+    use_global_device_ids
-+  } : (tensor<f32>) -> tensor<f64>
-+
-+  func.return %result : tensor<f64>
-+}
-+
-+// CHECK-LABEL: "default_all_reduce_variadic"
-+func.func @default_all_reduce_variadic(%arg0: tensor<f32>, %arg1: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-+  %0:2 = "stablehlo.all_reduce"(%arg0, %arg1) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> (tensor<f32>)
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
-+  } : (tensor<f32>, tensor<f32>) -> (tensor<f32>, tensor<f32>)
-+  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_all_to_all(%arg0: tensor<4x16xf32>) -> tensor<16x4xf32> {
-+  //               CHECK: "vhlo.all_to_all_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  //          CHECK-SAME:   concat_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-+  //          CHECK-SAME:   split_count = #vhlo.integer_v1<4 : i64>
-+  //          CHECK-SAME:   split_dimension = #vhlo.integer_v1<1 : i64>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x4x!vhlo.f32_v1>
-+  %0 = "stablehlo.all_to_all"(%arg0) {
-+    split_dimension = 1 : i64,
-+    concat_dimension = 0 : i64,
-+    split_count = 4 : i64,
-+    replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>
-+  } : (tensor<4x16xf32>) -> tensor<16x4xf32>
-+  func.return %0 : tensor<16x4xf32>
-+}
-+
-+// CHECK-LABEL: "op_and"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_and(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "op_atan2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_atan2(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.atan2_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.atan2"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_batch_norm_grad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
-+func.func @op_batch_norm_grad(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
-+  //      CHECK: "vhlo.batch_norm_grad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-+  // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-+  // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-+  %0:3 = "stablehlo.batch_norm_grad"(%arg0, %arg1, %arg2, %arg3, %arg4) {
-+    epsilon = 0.001 : f32,
-+    feature_index = 0 : i64
-+  } : (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16x16x16x16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>)
-+  func.return %0#0, %0#1, %0#2 : tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_batch_norm_inference"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
-+func.func @op_batch_norm_inference(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>, %arg3: tensor<16xf32>, %arg4: tensor<16xf32>) -> tensor<16x16x16x16xf32> {
-+  //      CHECK: "vhlo.batch_norm_inference_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) <{
-+  // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-+  // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.batch_norm_inference"(%arg0, %arg1, %arg2, %arg3, %arg4) {
-+    epsilon = 0.001 : f32,
-+    feature_index = 0 : i64
-+  } : (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>, tensor<16xf32>) -> tensor<16x16x16x16xf32>
-+  func.return %0 : tensor<16x16x16x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_batch_norm_training"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_batch_norm_training(%arg0: tensor<16x16x16x16xf32>, %arg1: tensor<16xf32>, %arg2: tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) {
-+  //      CHECK: "vhlo.batch_norm_training_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   epsilon = #vhlo.float_v1<1.000000e-03 : !vhlo.f32_v1>,
-+  // CHECK-SAME:   feature_index = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>) -> (!vhlo.tensor_v1<16x16x16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x!vhlo.f32_v1>)
-+  %0:3 = "stablehlo.batch_norm_training"(%arg0, %arg1, %arg2) {
-+    epsilon = 0.001 : f32,
-+    feature_index = 0 : i64
-+  } : (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>) -> (tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>)
-+  func.return %0#0, %0#1, %0#2 : tensor<16x16x16x16xf32>, tensor<16xf32>, tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_bitcast_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_bitcast_convert(%arg0: tensor<i32>) -> tensor<f32> {
-+  // CHECK: "vhlo.bitcast_convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.bitcast_convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_broadcast_in_dim(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
-+  //      CHECK: "vhlo.broadcast_in_dim_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.broadcast_in_dim"(%arg0) {
-+    broadcast_dimensions = array<i64: 1>
-+  } : (tensor<16xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_broadcast(%arg0: tensor<16xf32>) -> tensor<16x16xf32> {
-+  //      CHECK: "vhlo.broadcast_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   broadcast_sizes = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.broadcast"(%arg0) {
-+    broadcast_sizes = array<i64: 16>
-+  } : (tensor<16xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_case"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_case(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.case"(%arg0) ({
-+    "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-+  }) : (tensor<i32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_cbrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_cbrt(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.cbrt_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.cbrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_ceil"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_ceil(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.ceil_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.ceil"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_cholesky"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_cholesky(%arg0: tensor<1x16x16xf32>) -> tensor<1x16x16xf32> {
-+  //      CHECK: "vhlo.cholesky_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   lower = #vhlo.bool_v1<true>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.cholesky"(%arg0) {
-+    lower = true
-+  } : (tensor<1x16x16xf32>) -> tensor<1x16x16xf32>
-+  func.return %0 : tensor<1x16x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_clamp"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_clamp(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.clamp_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.clamp"(%arg0, %arg1, %arg2) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_count_leading_zeros"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_count_leading_zeros(%arg0: tensor<i32>) -> tensor<i32> {
-+  // CHECK: "vhlo.count_leading_zeros_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.count_leading_zeros"(%arg0) : (tensor<i32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "op_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_collective_broadcast(%arg0: tensor<16x8xf32>, %arg1: tensor<1xi32>) -> tensor<16x8xf32> {
-+  //               CHECK: "vhlo.collective_broadcast_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<true>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.collective_broadcast"(%arg0, %arg1) {
-+    replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+    has_dynamic_root
-+  } : (tensor<16x8xf32>, tensor<1xi32>) -> tensor<16x8xf32>
-+  func.return %0 : tensor<16x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_collective_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_collective_reduce(%arg0: tensor<f32>, %arg1: tensor<1xi32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.collective_reduce_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  //          CHECK-SAME:   has_dynamic_root = #vhlo.bool_v1<true>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.collective_reduce"(%arg0, %arg1) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+    use_global_device_ids,
-+    has_dynamic_root
-+  } : (tensor<f32>, tensor<1xi32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_collective_permute(%arg0: tensor<16x8xf32>) -> tensor<16x8xf32> {
-+  //               CHECK: "vhlo.collective_permute_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.collective_permute"(%arg0) {
-+    source_target_pairs = dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>
-+  } : (tensor<16x8xf32>) -> tensor<16x8xf32>
-+  func.return %0 : tensor<16x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_compare"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_compare(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  //      CHECK: "vhlo.compare_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   compare_type = #vhlo<comparison_type_v1 TOTALORDER>,
-+  // CHECK-SAME:   comparison_direction = #vhlo<comparison_direction_v1 EQ>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>,
-+    compare_type = #stablehlo<comparison_type TOTALORDER>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "op_complex"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_complex(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<complex<f32>> {
-+  // CHECK: "vhlo.complex_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  %0 = "stablehlo.complex"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<complex<f32>>
-+  func.return %0 : tensor<complex<f32>>
-+}
-+
-+// CHECK-LABEL: "op_composite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_composite(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.composite_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{#vhlo.string_v1<"my_int"> = #vhlo.integer_v1<1 : i64>, #vhlo.string_v1<"my_string"> = #vhlo.string_v1<"foo">}>
-+  //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
-+  //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
-+  //          CHECK-SAME:   version = #vhlo.integer_v1<1 : i32>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.composite"(%arg0) {
-+    name = "stablehlo.composite_target",
-+    decomposition = @composite_target,
-+    version = 1 : i32,
-+    composite_attributes = {
-+      my_int = 1 : i64,
-+      my_string = "foo"
-+    }
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "composite_regions"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @composite_regions(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.composite_v2"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   composite_attributes = #vhlo.dict_v1<{}>
-+  //          CHECK-SAME:   decomposition = #vhlo.string_v1<"composite_target">
-+  //          CHECK-SAME:   name = #vhlo.string_v1<"stablehlo.composite_target">
-+  //          CHECK-SAME:   version = #vhlo.integer_v1<1 : i32>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT: ^bb0(%[[ARG1:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.composite"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>):
-+      "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-+  }) {
-+    name = "stablehlo.composite_target",
-+    decomposition = @composite_target,
-+    version = 1 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_concatenate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_concatenate(%arg0: tensor<8xf32>, %arg1: tensor<8xf32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.concatenate_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.concatenate"(%arg0, %arg1) {
-+    dimension = 0 : i64
-+  } : (tensor<8xf32>, tensor<8xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_constant"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_constant(%arg0: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.constant_v1"() <{
-+  // CHECK-SAME:   value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>
-+  // CHECK-SAME: }> : () -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.constant"() {
-+    value = dense<0.0> : tensor<f32>
-+  } : () -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_convert"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_convert(%arg0: tensor<i32>) -> tensor<f32> {
-+  // CHECK: "vhlo.convert_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.convert"(%arg0) : (tensor<i32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_convolution"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_convolution(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32> {
-+  //      CHECK: "vhlo.convolution_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
-+  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<2x2xi64>>,
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-+  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
-+  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<true> : tensor<2xi1>>,
-+  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<1x7x7x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.convolution"(%arg0, %arg1) {
-+    window_strides = array<i64: 2, 2>,
-+    padding = dense<1> : tensor<2x2xi64>,
-+    lhs_dilation = array<i64: 2, 2>,
-+    rhs_dilation = array<i64: 2, 2>,
-+    window_reversal = array<i1: true, true>,
-+    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
-+    feature_group_count = 1 : i64,
-+    batch_group_count = 1 : i64,
-+    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
-+  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>) -> tensor<1x7x7x16xf32>
-+  func.return %0 : tensor<1x7x7x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_cosine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_cosine(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.cosine_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.cosine"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_create_token"
-+func.func @op_create_token() -> !stablehlo.token {
-+  // CHECK: "vhlo.create_token_v1"() : () -> !vhlo.token_v1
-+  %0 = "stablehlo.create_token"() : () -> !stablehlo.token
-+  func.return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_cross_replica_sum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_cross_replica_sum(%arg0: tensor<f32>) -> tensor<f32> {
-+  //               CHECK: "vhlo.cross-replica-sum_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>
-+  //          CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.cross-replica-sum"(%arg0) {
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_custom_call"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_custom_call(%arg0: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
-+  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-+  // CHECK-SAME:   backend_config = #vhlo.string_v1<"\08\03\1A\02">,
-+  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-+  // CHECK-SAME:   called_computations = #vhlo.array_v1<[#vhlo.string_v1<"foo">]>,
-+  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<> : tensor<0xindex>>]>,
-+  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[
-+  // CHECK-SAME:     #vhlo.output_operand_alias_v1<
-+  // CHECK-SAME:       outputTupleIndices = [],
-+  // CHECK-SAME:       operandIndex = 0,
-+  // CHECK-SAME:       operandTupleIndices = []>]>,
-+  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<> : tensor<0xindex>>]>,
-+  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[]>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo",
-+    has_side_effect = true,
-+    backend_config = "\08\03\1A\02",
-+    api_version = 2 : i32,
-+    called_computations = [@foo],
-+    operand_layouts = [dense<> : tensor<0xindex>],
-+    output_operand_aliases = [
-+      #stablehlo.output_operand_alias<output_tuple_indices = [],
-+                                 operand_index = 0,
-+                                 operand_tuple_indices = []>],
-+    result_layouts = [dense<> : tensor<0xindex>]
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_custom_call_empty_result_layout"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func public @op_custom_call_empty_result_layout(%arg0: tensor<i64>) -> tensor<i64> {
-+  // %0 = "vhlo.custom_call_v2"(%arg0) <{>}> : (!vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tuple_v1<>
-+  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
-+  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-+  // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-+  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"empty_output">,
-+  // CHECK-SAME:   called_computations = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<> : tensor<0xindex>>]>,
-+  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[]>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tuple_v1<>
-+  %0 = "stablehlo.custom_call"(%arg0) <{
-+    api_version = 2 : i32,
-+    call_target_name = "empty_output",
-+    has_side_effect = true,
-+    operand_layouts = [dense<> : tensor<0xindex>],
-+    result_layouts = []
-+  }> : (tensor<i64>) -> tuple<>
-+  return %arg0 : tensor<i64>
-+}
-+
-+// CHECK-LABEL: "op_custom_call_with_result_tilings"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_custom_call_with_result_tilings(%arg0: tensor<64x256xf32>) -> tensor<64x256xf32> {
-+  //      CHECK: "vhlo.custom_call_v2"(%[[ARG0]]) <{
-+  // CHECK-SAME:   api_version = #vhlo<api_version_v1 API_VERSION_STATUS_RETURNING>,
-+  // CHECK-SAME:   backend_config = #vhlo.string_v1<"">,
-+  // CHECK-SAME:   call_target_name = #vhlo.string_v1<"foo">,
-+  // CHECK-SAME:   called_computations = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   has_side_effect = #vhlo.bool_v1<false>,
-+  // CHECK-SAME:   operand_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<[1, 0]> : tensor<2xindex>>]>,
-+  // CHECK-SAME:   output_operand_aliases = #vhlo.array_v1<[]>,
-+  // CHECK-SAME:   result_layouts = #vhlo.array_v1<[#vhlo.tensor_v1<dense<[1, 0]> : tensor<2xindex>>]>,
-+  // CHECK-SAME:   result_tilings = #vhlo.array_v1<[#vhlo.array_v1<[#vhlo.tensor_v1<dense<[8, 128]> : tensor<2xindex>>]>]>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<64x256x!vhlo.f32_v1>) -> !vhlo.tensor_v1<64x256x!vhlo.f32_v1>
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    api_version = 2 : i32,
-+    call_target_name = "foo",
-+    operand_layouts = [dense<[1, 0]> : tensor<2xindex>],
-+    result_layouts = [dense<[1, 0]> : tensor<2xindex>],
-+    result_tilings = [[dense<[8, 128]> : tensor<2xindex>]]
-+  } : (tensor<64x256xf32>) -> tensor<64x256xf32>
-+  func.return %0 : tensor<64x256xf32>
-+}
-+
-+// CHECK-LABEL: "op_divide"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_divide(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.divide_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.divide"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_dot_general"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_dot_general(%arg0: tensor<8x8x16xf32>, %arg1: tensor<8x16x8xf32>) -> tensor<8x8x8xf32> {
-+  //      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   allow_imprecise_accumulation = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   lhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   num_primitive_operations = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-+  // CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   rhs_component_count = #vhlo.type_v1<!vhlo.none_v1>,
-+  // CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
-+  // CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.none_v1>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
-+    dot_dimension_numbers = #stablehlo.dot<
-+      lhs_batching_dimensions = [0],
-+      lhs_contracting_dimensions = [2],
-+      rhs_batching_dimensions = [0],
-+      rhs_contracting_dimensions = [1]
-+    >,
-+    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
-+  } : (tensor<8x8x16xf32>, tensor<8x16x8xf32>) -> tensor<8x8x8xf32>
-+  func.return %0 : tensor<8x8x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_dot"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_dot(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-+  //      CHECK: "vhlo.dot_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot"(%arg0, %arg1) {
-+    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
-+  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-+  func.return %0 : tensor<8x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_broadcast_in_dim"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
-+  //      CHECK: "vhlo.dynamic_broadcast_in_dim_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   broadcast_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   known_expanding_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   known_nonexpanding_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x?x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_broadcast_in_dim"(%arg0, %arg1) {
-+    broadcast_dimensions = array<i64: 0, 1>,
-+    known_expanding_dimensions = array<i64: 0>,
-+    known_nonexpanding_dimensions = array<i64: 1>
-+  } : (tensor<?x?xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-+  func.return %0 : tensor<?x?xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_conv"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_dynamic_conv(%arg0: tensor<1x8x8x207xf32>, %arg1: tensor<3x3x207x16xf32>, %arg2: tensor<2x2xi64>) -> tensor<1x?x?x16xf32> {
-+  //      CHECK: "vhlo.dynamic_conv_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   batch_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   feature_group_count = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   input_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   input_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   input_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   kernel_input_feature_dimension = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   kernel_output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   kernel_spatial_dimensions = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   lhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
-+  // CHECK-SAME:   output_batch_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   output_feature_dimension = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   output_spatial_dimensions = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 HIGHEST>, #vhlo<precision_v1 HIGHEST>]>,
-+  // CHECK-SAME:   rhs_dilation = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>,
-+  // CHECK-SAME:   window_reversal = #vhlo.tensor_v1<dense<true> : tensor<2xi1>>,
-+  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<2> : tensor<2xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x8x8x207x!vhlo.f32_v1>, !vhlo.tensor_v1<3x3x207x16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x2x!vhlo.i64_v1>) -> !vhlo.tensor_v1<1x?x?x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_conv"(%arg0, %arg1, %arg2) {
-+    window_strides = array<i64: 2, 2>,
-+    lhs_dilation = array<i64: 2, 2>,
-+    rhs_dilation = array<i64: 2, 2>,
-+    window_reversal = array<i1: true, true>,
-+    dimension_numbers = #stablehlo.conv<[b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]>,
-+    feature_group_count = 1 : i64,
-+    batch_group_count = 1 : i64,
-+    precision_config = [#stablehlo<precision HIGHEST>, #stablehlo<precision HIGHEST>]
-+  } : (tensor<1x8x8x207xf32>, tensor<3x3x207x16xf32>, tensor<2x2xi64>) -> tensor<1x?x?x16xf32>
-+  func.return %0 : tensor<1x?x?x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_dynamic_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<3xi32>) -> tensor<1x5x8xf32> {
-+  //      CHECK: "vhlo.dynamic_gather_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>, !vhlo.tensor_v1<3x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_gather"(%arg0, %arg1, %arg2) {
-+    dimension_numbers = #stablehlo.gather<
-+      offset_dims = [2],
-+      collapsed_slice_dims = [0, 1],
-+      start_index_map = [0, 1],
-+      index_vector_dim = 2
-+    >,
-+    indices_are_sorted = true
-+  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>, tensor<3xi32>) -> tensor<1x5x8xf32>
-+  func.return %0 : tensor<1x5x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_gather_with_batching_dims"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_dynamic_gather_with_batching_dims(%arg0 : tensor<5x2x4x9xf32>, %arg1 : tensor<1x5x2xi32>, %arg2 : tensor<4xi32>) -> tensor<1x5x8xf32> {
-+  //      CHECK: "vhlo.dynamic_gather_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<5x2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>, !vhlo.tensor_v1<4x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_gather"(%arg0, %arg1, %arg2) {
-+    dimension_numbers = #stablehlo.gather<
-+      offset_dims = [2],
-+      collapsed_slice_dims = [1, 2],
-+      operand_batching_dims = [0],
-+      start_indices_batching_dims = [1],
-+      start_index_map = [1, 2],
-+      index_vector_dim = 2
-+    >,
-+    indices_are_sorted = true
-+  } : (tensor<5x2x4x9xf32>, tensor<1x5x2xi32>, tensor<4xi32>) -> tensor<1x5x8xf32>
-+  func.return %0 : tensor<1x5x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_iota"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_dynamic_iota(%arg0: tensor<1xindex>) -> tensor<?xf32> {
-+  //      CHECK: "vhlo.dynamic_iota_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_iota"(%arg0) {
-+    iota_dimension = 0 : i64
-+  } : (tensor<1xindex>) -> tensor<?xf32>
-+  func.return %0 : tensor<?xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}}, %[[ARG4:.*]]: {{.*}})
-+func.func @op_dynamic_pad(%arg0: tensor<?xf32>, %arg1: tensor<f32>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>, %arg4: tensor<1xindex>) -> tensor<?xf32> {
-+  // CHECK: "vhlo.dynamic_pad_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_pad"(%arg0, %arg1, %arg2, %arg3, %arg4) : (tensor<?xf32>, tensor<f32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-+  func.return %0 : tensor<?xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_dynamic_reshape(%arg0: tensor<16xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
-+  // CHECK: "vhlo.dynamic_reshape_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x?x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_reshape"(%arg0, %arg1) : (tensor<16xf32>, tensor<2xindex>) -> tensor<?x?xf32>
-+  func.return %0 : tensor<?x?xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_dynamic_slice(%arg0: tensor<16xf32>, %arg1: tensor<i64>) -> tensor<4xf32> {
-+  //      CHECK: "vhlo.dynamic_slice_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_slice"(%arg0, %arg1) {
-+    slice_sizes = array<i64: 4>
-+  } : (tensor<16xf32>, tensor<i64>) -> tensor<4xf32>
-+  func.return %0 : tensor<4xf32>
-+}
-+
-+// CHECK-LABEL: "op_dynamic_update_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_dynamic_update_slice(%arg0: tensor<16xf32>, %arg1: tensor<4xf32>, %arg2: tensor<i64>) -> tensor<16xf32> {
-+  // CHECK: "vhlo.dynamic_update_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<16xf32>, tensor<4xf32>, tensor<i64>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_einsum(%arg0: tensor<8x16xf32>, %arg1: tensor<16x8xf32>) -> tensor<8x8xf32> {
-+  //      CHECK: "vhlo.einsum_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab,bc->ac">
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.einsum"(%arg0, %arg1) {
-+    einsum_config = "ab,bc->ac"
-+  } : (tensor<8x16xf32>, tensor<16x8xf32>) -> tensor<8x8xf32>
-+  func.return %0 : tensor<8x8xf32>
-+}
-+
-+// CHECK-LABEL: "op_exponential_minus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_exponential_minus_one(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.exponential_minus_one_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.exponential_minus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_exponential"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_exponential(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.exponential_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.exponential"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_exp2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_exp2(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.exp2_v1"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.exp2"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_fft"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-+  //      CHECK: "vhlo.fft_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   fft_length = #vhlo.tensor_v1<dense<16> : tensor<1xi64>>,
-+  // CHECK-SAME:   fft_type = #vhlo<fft_type_v1 FFT>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<16x!vhlo.complex_v1<!vhlo.f32_v1>>
-+  %0 = "stablehlo.fft"(%arg0) {
-+    fft_type = #stablehlo<fft_type FFT>,
-+    fft_length = array<i64: 16>
-+  } : (tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>>
-+  func.return %0 : tensor<16xcomplex<f32>>
-+}
-+
-+// CHECK-LABEL: "op_floor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_floor(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.floor_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.floor"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+func.func private @op_func(%arg0: tensor<f32> {stablehlo.arg = "0"}) -> (tensor<f32> {stablehlo.result = "0"}) {
-+  // CHECK:      "vhlo.func_v1"() <{
-+  // CHECK-SAME:   arg_attrs = #vhlo.array_v1<[#vhlo.dict_v1<{#vhlo.string_v1<"stablehlo.arg"> = #vhlo.string_v1<"0">}>]>,
-+  // CHECK-SAME:   function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>>>,
-+  // CHECK-SAME:   res_attrs = #vhlo.array_v1<[#vhlo.dict_v1<{#vhlo.string_v1<"stablehlo.result"> = #vhlo.string_v1<"0">}>]>,
-+  // CHECK-SAME:   sym_name = #vhlo.string_v1<"op_func">,
-+  // CHECK-SAME:   sym_visibility = #vhlo.string_v1<"private">
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG0:.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : () -> ()
-+
-+  func.return %arg0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_gather(%arg0 : tensor<2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
-+  //      CHECK: "vhlo.gather_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<1> : tensor<3xi64>>,
-+  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x1x!vhlo.f32_v1>
-+  %0 = "stablehlo.gather"(%arg0, %arg1) {
-+    dimension_numbers = #stablehlo.gather<
-+      offset_dims = [2],
-+      collapsed_slice_dims = [0, 1],
-+      start_index_map = [0, 1],
-+      index_vector_dim = 2
-+    >,
-+    slice_sizes = array<i64: 1, 1, 1>,
-+    indices_are_sorted = true
-+  } : (tensor<2x4x9xf32>, tensor<1x5x2xi32>) -> tensor<1x5x1xf32>
-+  func.return %0 : tensor<1x5x1xf32>
-+}
-+
-+// CHECK-LABEL: "op_gather_with_batching_dims"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_gather_with_batching_dims(%arg0 : tensor<5x2x4x9xf32>, %arg1 : tensor<1x5x2xi32>) -> tensor<1x5x1xf32> {
-+  //      CHECK: "vhlo.gather_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   collapsed_slice_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   offset_dims = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+  // CHECK-SAME:   operand_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   slice_sizes = #vhlo.tensor_v1<dense<1> : tensor<4xi64>>,
-+  // CHECK-SAME:   start_index_map = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   start_indices_batching_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<5x2x4x9x!vhlo.f32_v1>, !vhlo.tensor_v1<1x5x2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<1x5x1x!vhlo.f32_v1>
-+  %0 = "stablehlo.gather"(%arg0, %arg1) {
-+    dimension_numbers = #stablehlo.gather<
-+      offset_dims = [2],
-+      collapsed_slice_dims = [1, 2],
-+      operand_batching_dims = [0],
-+      start_indices_batching_dims = [1],
-+      start_index_map = [1, 2],
-+      index_vector_dim = 2
-+    >,
-+    slice_sizes = array<i64: 1, 1, 1, 1>,
-+    indices_are_sorted = true
-+  } : (tensor<5x2x4x9xf32>, tensor<1x5x2xi32>) -> tensor<1x5x1xf32>
-+  func.return %0 : tensor<1x5x1xf32>
-+}
-+
-+// CHECK-LABEL: "op_get_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_get_dimension_size(%arg0: tensor<?xf32>) -> tensor<i32> {
-+  //      CHECK: "vhlo.get_dimension_size_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.get_dimension_size"(%arg0) {
-+    dimension = 0 : i64
-+  } : (tensor<?xf32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "op_get_tuple_element"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_get_tuple_element(%arg0: tuple<tensor<f32>, tensor<i32>>) -> tensor<f32> {
-+  //      CHECK: "vhlo.get_tuple_element_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   index = #vhlo.integer_v1<0 : i32>
-+  // CHECK-SAME: }> : (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.get_tuple_element"(%arg0) {
-+    index = 0 : i32
-+  } : (tuple<tensor<f32>, tensor<i32>>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_if"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_if(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.if_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }, {
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.if"(%arg0) ({
-+    "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-+  }, {
-+    "stablehlo.return"(%arg2) : (tensor<f32>) -> ()
-+  }) : (tensor<i1>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_imag"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_imag(%arg0: tensor<complex<f32>>) -> tensor<f32> {
-+  // CHECK: "vhlo.imag_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.imag"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_infeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_infeed(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
-+  //               CHECK: "vhlo.infeed_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   infeed_config = #vhlo.string_v1<"foo">,
-+  // CHECK-SAME{LITERAL}:   layout = #vhlo.array_v1<[#vhlo.array_v1<[]>]>
-+  //          CHECK-SAME: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-+  %0:2 = "stablehlo.infeed"(%arg0) {
-+    infeed_config = "foo",
-+    layout = [[]]
-+  } : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
-+  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_iota"
-+func.func @op_iota() -> tensor<16xf32> {
-+  //      CHECK: "vhlo.iota_v1"() <{
-+  // CHECK-SAME:   iota_dimension = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : () -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.iota"() {
-+    iota_dimension = 0 : i64
-+  } : () -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_is_finite"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_is_finite(%arg0: tensor<f32>) -> tensor<i1> {
-+  // CHECK: "vhlo.is_finite_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.is_finite"(%arg0) : (tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "op_log"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_log(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.log_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.log"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_log2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_log2(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.log2_v1"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.log2"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_log_plus_one"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.log_plus_one_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.log_plus_one"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_logistic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_logistic(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.logistic_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.logistic"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_map"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_map(%arg0: tensor<16xf32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.map_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.abs_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.map"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>):
-+      %1 = "stablehlo.abs"(%arg1) : (tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    dimensions = array<i64: 0>
-+  } : (tensor<16xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_maximum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_maximum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.maximum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.maximum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_minimum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_minimum(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.minimum_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.minimum"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_multiply"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_multiply(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.multiply_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.multiply"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_negate"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_negate(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.negate_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.negate"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_not"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_not(%arg0: tensor<i1>) -> tensor<i1> {
-+  // CHECK: "vhlo.not_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.not"(%arg0) : (tensor<i1>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "op_optimization_barrier"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_optimization_barrier(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.optimization_barrier_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.optimization_barrier"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_or"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_or(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
-+  // CHECK: "vhlo.or_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.or"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "op_outfeed"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_outfeed(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
-+  //      CHECK: "vhlo.outfeed_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   outfeed_config = #vhlo.string_v1<"foo">
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-+  %0 = "stablehlo.outfeed"(%arg0, %arg1) {
-+    outfeed_config = "foo"
-+  } : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-+  func.return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_pad"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_pad(%arg0: tensor<8xf32>, %arg1: tensor<f32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.pad_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   edge_padding_high = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-+  // CHECK-SAME:   edge_padding_low = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-+  // CHECK-SAME:   interior_padding = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.pad"(%arg0, %arg1) {
-+    edge_padding_high = array<i64: 4>,
-+    edge_padding_low = array<i64: 4>,
-+    interior_padding = array<i64: 0>
-+  } : (tensor<8xf32>, tensor<f32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_popcnt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_popcnt(%arg0: tensor<i32>) -> tensor<i32> {
-+  // CHECK: "vhlo.popcnt_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.popcnt"(%arg0) : (tensor<i32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "op_power"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_power(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.power_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.power"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_real_dynamic_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}}, %[[ARG3:.*]]: {{.*}})
-+func.func @op_real_dynamic_slice(%arg0: tensor<?xf32>, %arg1: tensor<1xindex>, %arg2: tensor<1xindex>, %arg3: tensor<1xindex>) -> tensor<?xf32> {
-+  // CHECK: "vhlo.real_dynamic_slice_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>, !vhlo.tensor_v1<1x!vhlo.index_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  %0 = "stablehlo.real_dynamic_slice"(%arg0, %arg1, %arg2, %arg3) : (tensor<?xf32>, tensor<1xindex>, tensor<1xindex>, tensor<1xindex>) -> tensor<?xf32>
-+  func.return %0 : tensor<?xf32>
-+}
-+
-+// CHECK-LABEL: "op_real"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_real(%arg0: tensor<complex<f32>>) -> tensor<f32> {
-+  // CHECK: "vhlo.real_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.real"(%arg0) : (tensor<complex<f32>>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_recv_no_source_target_pairs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_recv_no_source_target_pairs(%arg0: !stablehlo.token) -> (tensor<f32>, !stablehlo.token) {
-+  //      CHECK: "vhlo.recv_v2"(%[[ARG0]]) <{
-+  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   channel_type = #vhlo.integer_v1<3 : i64>,
-+  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>,
-+  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<> : tensor<0xi64>
-+  // CHECK-SAME{LITERAL}: }> : (!vhlo.token_v1) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1)
-+  %0:2 = "stablehlo.recv"(%arg0) {
-+    channel_handle = #stablehlo.channel_handle<handle = 0, type = 3>,
-+    is_host_transfer = true
-+  } : (!stablehlo.token) -> (tensor<f32>, !stablehlo.token)
-+  func.return %0#0, %0#1 : tensor<f32>, !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_reduce(%arg0: tensor<16xf32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0]], %[[ARG1]])
-+  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //  CHECK: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.reduce"(%arg0, %arg1) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    dimensions = array<i64: 0>
-+  } : (tensor<16xf32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_reduce_precision"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_reduce_precision(%arg0: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.reduce_precision_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   exponent_bits = #vhlo.integer_v1<8 : i32>
-+  // CHECK-SAME:   mantissa_bits = #vhlo.integer_v1<10 : i32>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.reduce_precision"(%arg0) {
-+    exponent_bits = 8 : i32,
-+    mantissa_bits = 10 : i32
-+  } : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK_lABEL: "op_reduce_with_promotable_types"
-+func.func @op_reduce_with_promotable_types(%arg0: tensor<4x4xf32>, %arg1 : tensor<f32>)
-+    -> (tensor<4xf64>) {
-+  //  CHECK: "vhlo.reduce_v1"(%[[ARG0:.*]], %[[ARG1:.*]])
-+  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-+  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-+  //  CHECK: }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x!vhlo.f64_v1>
-+  %0 = "stablehlo.reduce"(%arg0, %arg1) ({
-+  ^bb0(%arg2: tensor<f64>, %arg3: tensor<f64> ):
-+    %1 = "stablehlo.add"(%arg2, %arg3) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-+    "stablehlo.return"(%1) : (tensor<f64>) -> ()
-+
-+  }) {dimensions = array<i64: 0>} : (tensor<4x4xf32>, tensor<f32>) -> tensor<4xf64>
-+
-+  func.return %0: tensor<4xf64>
-+}
-+
-+// CHECK-LABEL: "op_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_reduce_scatter(%arg0: tensor<16xf32>) -> tensor<16xf32> {
-+  //               CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0]]) <{
-+  //          CHECK-SAME:   channel_id = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME{LITERAL}:   replica_groups = #vhlo.tensor_v1<dense<[[0], [1]]> : tensor<2x1xi64>>,
-+  //          CHECK-SAME:   scatter_dimension = #vhlo.integer_v1<0 : i64>
-+  //          CHECK-SAME:   use_global_device_ids = #vhlo.bool_v1<true>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.reduce_scatter"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg1, %arg2) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    scatter_dimension = 0 : i64,
-+    replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>,
-+    channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+    use_global_device_ids
-+  } : (tensor<16xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK_lABEL: "op_reduce_scatter_with_promotable_types"
-+func.func @op_reduce_scatter_with_promotable_types(%data: tensor<4x16xf32>) -> tensor<4x4xf64> {
-+  //  CHECK: "vhlo.reduce_scatter_v1"(%[[ARG0:.*]])
-+  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-+  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-+  //  CHECK: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f64_v1>
-+  %0 = "stablehlo.reduce_scatter"(%data) ({
-+    ^bb0(%arg2: tensor<f64>, %arg3: tensor<f64>):
-+    %1 = stablehlo.add %arg2, %arg3 : tensor<f64>
-+    "stablehlo.return"(%1) : (tensor<f64>) -> ()
-+  }) {replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
-+      scatter_dimension = 1 : i64,
-+      channel_handle = #stablehlo.channel_handle<handle = 1, type = 0>,
-+      use_global_device_ids} : (tensor<4x16xf32>) -> tensor<4x4xf64>
-+  func.return %0 : tensor<4x4xf64>
-+}
-+
-+
-+// CHECK-LABEL: "op_reduce_window"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_reduce_window(%arg0: tensor<2x17x31x7xf32>, %arg1: tensor<f32>) -> tensor<2x9x16x7xf32> {
-+  //               CHECK: "vhlo.reduce_window_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  //          CHECK-SAME:   base_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-+  // CHECK-SAME{LITERAL}:   padding = #vhlo.tensor_v1<dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>>,
-+  //          CHECK-SAME:   window_dilations = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-+  //          CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-+  //          CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 4, 4, 1]> : tensor<4xi64>>
-+  //          CHECK-SAME: }> ({
-+  //          CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //          CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.maximum_v1"(%[[ARG2]], %[[ARG3]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  //          CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //          CHECK-NEXT: }) : (!vhlo.tensor_v1<2x17x31x7x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<2x9x16x7x!vhlo.f32_v1>
-+  %0 = "stablehlo.reduce_window"(%arg0, %arg1) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %1 = "stablehlo.maximum"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    window_dimensions = array<i64: 1, 2, 2, 1>,
-+    window_strides = array<i64: 1, 4, 4, 1>,
-+    base_dilations = array<i64: 1, 2, 2, 1>,
-+    window_dilations = array<i64: 1, 2, 2, 1>,
-+    padding = dense<[[0, 0], [2, 0], [0, 2], [0, 0]]> : tensor<4x2xi64>
-+  } : (tensor<2x17x31x7xf32>, tensor<f32>) -> tensor<2x9x16x7xf32>
-+  func.return %0 : tensor<2x9x16x7xf32>
-+}
-+
-+// CHECK-LABEL: "op_reduce_window_with_promotable_types"
-+func.func @op_reduce_window_with_promotable_types(%arg0: tensor<4x2xf32>,
-+    %arg1: tensor<4x2xf32>, %init0: tensor<f32>, %init1: tensor<f32>) ->
-+    (tensor<2x2xf64>, tensor<2x2xf32>) {
-+  //  CHECK: "vhlo.reduce_window_v1"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]], %[[ARG3:.*]])
-+  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]], %[[VAL2:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  //  CHECK: }) : (!vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<4x2x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<2x2x!vhlo.f64_v1>, !vhlo.tensor_v1<2x2x!vhlo.f32_v1>)
-+  %0:2 = "stablehlo.reduce_window"(%arg0, %arg1, %init0, %init1) ({
-+         ^bb0(%a0: tensor<f64>, %a1: tensor<f32>, %b0: tensor<f64>,
-+                %b1: tensor<f32>):
-+              %2 = stablehlo.add %a0, %b0 : tensor<f64>
-+              %3 = stablehlo.add %a1, %b1 : tensor<f32>
-+              "stablehlo.return"(%2,%3) : (tensor<f64>, tensor<f32>) -> ()
-+            })
-+         { padding = dense<[[2, 2], [0, 0]]> : tensor<2x2xi64>,
-+         window_dimensions = array<i64: 5, 1>,
-+         window_strides = array<i64: 3, 1> }
-+         : (tensor<4x2xf32>, tensor<4x2xf32>, tensor<f32>, tensor<f32>) ->
-+              (tensor<2x2xf64>, tensor<2x2xf32>)
-+  func.return %0#0, %0#1 : tensor<2x2xf64>, tensor<2x2xf32>
-+}
-+
-+// CHECK-LABEL: "op_remainder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_remainder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.remainder_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.remainder"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_replica_id"
-+func.func @op_replica_id() -> tensor<ui32> {
-+  // CHECK: "vhlo.replica_id_v1"() : () -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  %0 = "stablehlo.replica_id"() : () -> tensor<ui32>
-+  func.return %0 : tensor<ui32>
-+}
-+
-+// CHECK-LABEL: "op_partition_id"
-+func.func @op_partition_id() -> tensor<ui32> {
-+  // CHECK: "vhlo.partition_id_v1"() : () -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  %0 = "stablehlo.partition_id"() : () -> tensor<ui32>
-+  func.return %0 : tensor<ui32>
-+}
-+
-+// CHECK-LABEL: "op_reshape"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_reshape(%arg0: tensor<16xf32>) -> tensor<4x4xf32> {
-+  // CHECK: "vhlo.reshape_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  %0 = "stablehlo.reshape"(%arg0) : (tensor<16xf32>) -> tensor<4x4xf32>
-+  func.return %0 : tensor<4x4xf32>
-+}
-+
-+// CHECK-LABEL: "op_return"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_return(%arg0: tensor<i32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  //      CHECK: "vhlo.case_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.case"(%arg0) ({
-+    "stablehlo.return"(%arg1) : (tensor<f32>) -> ()
-+  }) : (tensor<i32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_reverse"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_reverse(%arg0: tensor<16xf32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.reverse_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.reverse"(%arg0) {
-+    dimensions = array<i64: 0>
-+  } : (tensor<16xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_rng_bit_generator"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_rng_bit_generator(%arg0: tensor<f32>) -> (tensor<f32>, tensor<f32>) {
-+  //      CHECK: "vhlo.rng_bit_generator_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   rng_algorithm = #vhlo<rng_algorithm_v1 PHILOX>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>)
-+  %0:2 = "stablehlo.rng_bit_generator"(%arg0) {
-+    rng_algorithm = #stablehlo<rng_algorithm PHILOX>
-+  } : (tensor<f32>) -> (tensor<f32>, tensor<f32>)
-+  func.return %0#0, %0#1 : tensor<f32>, tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_rng"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_rng(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<0xindex>) -> tensor<f32> {
-+  //      CHECK: "vhlo.rng_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   rng_distribution = #vhlo<rng_distribution_v1 NORMAL>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<0x!vhlo.index_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.rng"(%arg0, %arg1, %arg2) {
-+    rng_distribution = #stablehlo<rng_distribution NORMAL>
-+  } : (tensor<f32>, tensor<f32>, tensor<0xindex>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_round_nearest_afz"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_round_nearest_afz(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.round_nearest_afz_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.round_nearest_afz"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_round_nearest_even"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_round_nearest_even(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.round_nearest_even_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.round_nearest_even"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_rsqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_rsqrt(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.rsqrt_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.rsqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_scatter(%arg0: tensor<200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<200x100x300xf32> {
-+  //      CHECK: "vhlo.scatter_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   input_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   scatter_dims_to_operand_dims = #vhlo.tensor_v1<dense<[0, 1]> : tensor<2xi64>>,
-+  // CHECK-SAME:   scatter_indices_batching_dims = #vhlo.tensor_v1<dense<> : tensor<0xi64>>,
-+  // CHECK-SAME:   unique_indices = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   update_window_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>
-+  %0 = "stablehlo.scatter"(%arg0, %arg1, %arg2) ({
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    scatter_dimension_numbers = #stablehlo.scatter<
-+      update_window_dims = [1],
-+      inserted_window_dims = [0, 1],
-+      scatter_dims_to_operand_dims = [0, 1],
-+      index_vector_dim = 1
-+    >,
-+    indices_are_sorted = true,
-+    unique_indices = true
-+  } : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) -> tensor<200x100x300xf32>
-+  func.return %0 : tensor<200x100x300xf32>
-+}
-+
-+// CHECK-LABEL: "op_scatter_with_batching_dims"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_scatter_with_batching_dims(%arg0: tensor<10x200x100x300xf32>, %arg1: tensor<10x2xi32>, %arg2: tensor<10x300xf32>) -> tensor<10x200x100x300xf32> {
-+  //      CHECK: "vhlo.scatter_v2"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   index_vector_dim = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:   indices_are_sorted = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   input_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   inserted_window_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   scatter_dims_to_operand_dims = #vhlo.tensor_v1<dense<[1, 2]> : tensor<2xi64>>,
-+  // CHECK-SAME:   scatter_indices_batching_dims = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   unique_indices = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   update_window_dims = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG3:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG4:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.add_v1"(%[[ARG3]], %[[ARG4]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<10x200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x200x100x300x!vhlo.f32_v1>
-+  %0 = "stablehlo.scatter"(%arg0, %arg1, %arg2) ({
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    scatter_dimension_numbers = #stablehlo.scatter<
-+      update_window_dims = [1],
-+      inserted_window_dims = [1, 2],
-+      input_batching_dims = [0],
-+      scatter_dims_to_operand_dims = [1, 2],
-+      scatter_indices_batching_dims = [0],
-+      index_vector_dim = 1
-+    >,
-+    indices_are_sorted = true,
-+    unique_indices = true
-+  } : (tensor<10x200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) -> tensor<10x200x100x300xf32>
-+  func.return %0 : tensor<10x200x100x300xf32>
-+}
-+
-+// CHECK_lABEL: "op_scatter_with_promotable_types"
-+func.func @op_scatter_with_promotable_types(%input_tensor: tensor<200x100x300xf32>,
-+    %scatter_indices: tensor<10x2xi32>, %updates: tensor<10x300xf32>) ->
-+      tensor<200x100x300xf64> {
-+  //  CHECK: "vhlo.scatter_v2"(%[[ARG0:.*]], %[[ARG1:.*]], %[[ARG2:.*]])
-+  //  CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-+  //  CHECK:     "vhlo.return_v1"(%[[VAL1:.*]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-+  //  CHECK: }) : (!vhlo.tensor_v1<200x100x300x!vhlo.f32_v1>, !vhlo.tensor_v1<10x2x!vhlo.i32_v1>, !vhlo.tensor_v1<10x300x!vhlo.f32_v1>) -> !vhlo.tensor_v1<200x100x300x!vhlo.f64_v1>
-+  %0 = "stablehlo.scatter" (%input_tensor, %scatter_indices, %updates) ({
-+  ^bb0(%lhs: tensor<f64>, %rhs: tensor<f64>):
-+    %add = stablehlo.add %lhs, %rhs : tensor<f64>
-+    "stablehlo.return"(%add) : (tensor<f64>) -> ()
-+  }) {
-+    scatter_dimension_numbers = #stablehlo.scatter<
-+      update_window_dims = [1],
-+      inserted_window_dims = [0, 1],
-+      scatter_dims_to_operand_dims = [0, 1],
-+      index_vector_dim = 1
-+    >,
-+    indices_are_sorted = true,
-+    unique_indices = true
-+  } : (tensor<200x100x300xf32>, tensor<10x2xi32>, tensor<10x300xf32>) ->
-+      tensor<200x100x300xf64>
-+  func.return %0 : tensor<200x100x300xf64>
-+}
-+
-+// CHECK-LABEL: "op_select_and_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_select_and_scatter(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf32> {
-+  //      CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) <{
-+  // CHECK-SAME:   padding = #vhlo.tensor_v1<dense<1> : tensor<4x2xi64>>,
-+  // CHECK-SAME:   window_dimensions = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>,
-+  // CHECK-SAME:   window_strides = #vhlo.tensor_v1<dense<[1, 2, 2, 1]> : tensor<4xi64>>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG31:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG41:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL11:.*]] = "vhlo.compare_v1"(%[[ARG31]], %[[ARG41]]) <{compare_type = #vhlo<comparison_type_v1 TOTALORDER>, comparison_direction = #vhlo<comparison_direction_v1 GE>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL11]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-+  // CHECK-NEXT: }, {
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG32:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG42:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL12:.*]] = "vhlo.add_v1"(%[[ARG32]], %[[ARG42]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL12]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>, !vhlo.tensor_v1<12x13x13x66x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>
-+  %0 = "stablehlo.select_and_scatter"(%arg0, %arg1, %arg2) ({
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.compare"(%arg3, %arg4) {compare_type = #stablehlo<comparison_type TOTALORDER>, comparison_direction = #stablehlo<comparison_direction GE>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+      "stablehlo.return"(%1) : (tensor<i1>) -> ()
-+  }, {
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+  }) {
-+    window_dimensions = array<i64: 1, 2, 2, 1>,
-+    window_strides = array<i64: 1, 2, 2, 1>,
-+    padding = dense<1> : tensor<4x2xi64>
-+  } : (tensor<10x24x24x64xf32>, tensor<12x13x13x66xf32>, tensor<f32>) -> tensor<10x24x24x64xf32>
-+  func.return %0 : tensor<10x24x24x64xf32>
-+}
-+
-+// CHECK-LABEL: "op_select_and_scatter_with_promotable_types"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_select_and_scatter_with_promotable_types(%arg0: tensor<10x24x24x64xf32>, %arg1: tensor<12x13x13x66xf32>, %arg2: tensor<f32>) -> tensor<10x24x24x64xf64> {
-+  // CHECK: "vhlo.select_and_scatter_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]])
-+  // CHECK:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f64_v1>):
-+  // CHECK:     %[[VAL:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  // CHECK:     "vhlo.return_v1"(%[[VAL]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>) -> ()
-+  // CHECK: }) : (!vhlo.tensor_v1<10x24x24x64x!vhlo.f32_v1>, !vhlo.tensor_v1<12x13x13x66x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<10x24x24x64x!vhlo.f64_v1>
-+  %0 = "stablehlo.select_and_scatter"(%arg0, %arg1, %arg2) ({
-+    ^bb0(%arg3: tensor<f32>, %arg4: tensor<f32>):
-+      %1 = "stablehlo.compare"(%arg3, %arg4) {compare_type = #stablehlo<comparison_type TOTALORDER>, comparison_direction = #stablehlo<comparison_direction GE>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+      "stablehlo.return"(%1) : (tensor<i1>) -> ()
-+  }, {
-+    ^bb0(%arg3: tensor<f64>, %arg4: tensor<f64>):
-+      %1 = "stablehlo.add"(%arg3, %arg4) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-+      "stablehlo.return"(%1) : (tensor<f64>) -> ()
-+  }) {
-+    window_dimensions = array<i64: 1, 2, 2, 1>,
-+    window_strides = array<i64: 1, 2, 2, 1>,
-+    padding = dense<1> : tensor<4x2xi64>
-+  } : (tensor<10x24x24x64xf32>, tensor<12x13x13x66xf32>, tensor<f32>) -> tensor<10x24x24x64xf64>
-+  func.return %0 : tensor<10x24x24x64xf64>
-+}
-+
-+// CHECK-LABEL: "op_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}}, %[[ARG2:.*]]: {{.*}})
-+func.func @op_select(%arg0: tensor<i1>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.select_v1"(%[[ARG0]], %[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.select"(%arg0, %arg1, %arg2) : (tensor<i1>, tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_send_no_source_target_pairs"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_send_no_source_target_pairs(%arg0: tensor<f32>, %arg1: !stablehlo.token) -> !stablehlo.token {
-+  //      CHECK: "vhlo.send_v2"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:   channel_type = #vhlo.integer_v1<2 : i64>,
-+  // CHECK-SAME:   is_host_transfer = #vhlo.bool_v1<true>,
-+  // CHECK-SAME{LITERAL}:   source_target_pairs = #vhlo.tensor_v1<dense<> : tensor<0xi64>>
-+  // CHECK-SAME{LITERAL}: }> : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.token_v1) -> !vhlo.token_v1
-+  %0 = "stablehlo.send"(%arg0, %arg1) {
-+    channel_handle = #stablehlo.channel_handle<handle = 0, type = 2>,
-+    is_host_transfer = true
-+  } : (tensor<f32>, !stablehlo.token) -> !stablehlo.token
-+  func.return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "op_set_dimension_size"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_set_dimension_size(%arg0: tensor<?xf32>, %arg1: tensor<i32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.set_dimension_size_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<?x!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.set_dimension_size"(%arg0, %arg1) {
-+    dimension = 0 : i64
-+  } : (tensor<?xf32>, tensor<i32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_shift_left"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_shift_left(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
-+  // CHECK: "vhlo.shift_left_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.shift_left"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "op_shift_right_arithmetic"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_shift_right_arithmetic(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
-+  // CHECK: "vhlo.shift_right_arithmetic_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.shift_right_arithmetic"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "op_shift_right_logical"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_shift_right_logical(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
-+  // CHECK: "vhlo.shift_right_logical_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.shift_right_logical"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "op_sign"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_sign(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.sign_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.sign"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_sine"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_sine(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.sine_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.sine"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_slice"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_slice(%arg0: tensor<16xf32>) -> tensor<4xf32> {
-+  //      CHECK: "vhlo.slice_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   limit_indices = #vhlo.tensor_v1<dense<4> : tensor<1xi64>>,
-+  // CHECK-SAME:   start_indices = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+  // CHECK-SAME:   strides = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x!vhlo.f32_v1>
-+  %0 = "stablehlo.slice"(%arg0) {
-+    start_indices = array<i64: 0>,
-+    limit_indices = array<i64: 4>,
-+    strides = array<i64: 1>
-+  } : (tensor<16xf32>) -> tensor<4xf32>
-+  func.return %0 : tensor<4xf32>
-+}
-+
-+// CHECK-LABEL: "op_sort"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_sort(%arg0: tensor<16xf32>) -> tensor<16xf32> {
-+  //      CHECK: "vhlo.sort_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   dimension = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME:   is_stable = #vhlo.bool_v1<true>
-+  // CHECK-SAME: }> ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:     %[[VAL1:.*]] = "vhlo.compare_v1"(%[[ARG1]], %[[ARG2]]) <{compare_type = #vhlo<comparison_type_v1 FLOAT>, comparison_direction = #vhlo<comparison_direction_v1 GT>}>
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x!vhlo.f32_v1>
-+  %0 = "stablehlo.sort"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
-+      %1 = "stablehlo.compare"(%arg1, %arg2) {compare_type = #stablehlo<comparison_type FLOAT>, comparison_direction = #stablehlo<comparison_direction GT>} : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+      "stablehlo.return"(%1) : (tensor<i1>) -> ()
-+  }) {
-+    dimension = 0 : i64,
-+    is_stable = true
-+  } : (tensor<16xf32>) -> tensor<16xf32>
-+  func.return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: "op_sqrt"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_sqrt(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.sqrt_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.sqrt"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_subtract"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_subtract(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.subtract_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.subtract"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_tan"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_tan(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.tan_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.tan"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_tanh"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_tanh(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.tanh_v2"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.tanh"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_torch_index_select"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_torch_index_select(%arg0: tensor<5x1x5xf32>, %arg1: tensor<2xi32>) ->  tensor<2x1x5xf32> {
-+  //      CHECK: "vhlo.torch_index_select_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   batch_dims = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME:   dim = #vhlo.integer_v1<0 : i64>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<5x1x5x!vhlo.f32_v1>, !vhlo.tensor_v1<2x!vhlo.i32_v1>) -> !vhlo.tensor_v1<2x1x5x!vhlo.f32_v1>
-+  %0 = "stablehlo.torch_index_select"(%arg0, %arg1) {
-+    dim = 0 : i64,
-+    batch_dims = 0 : i64
-+  } : (tensor<5x1x5xf32>, tensor<2xi32>) -> tensor<2x1x5xf32>
-+  func.return %0 : tensor<2x1x5xf32>
-+}
-+
-+// CHECK-LABEL: "op_transpose"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_transpose(%arg0: tensor<16x8xf32>) ->  tensor<8x16xf32> {
-+  //      CHECK: "vhlo.transpose_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   permutation = #vhlo.tensor_v1<dense<[1, 0]> : tensor<2xi64>>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.transpose"(%arg0) {
-+    permutation = array<i64: 1, 0>
-+  } : (tensor<16x8xf32>) -> tensor<8x16xf32>
-+  func.return %0 : tensor<8x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_triangular_solve"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_triangular_solve(%arg0: tensor<16x16xf32>, %arg1: tensor<16x16xf32>) ->  tensor<16x16xf32> {
-+  //      CHECK: "vhlo.triangular_solve_v1"(%[[ARG0]], %[[ARG1]]) <{
-+  // CHECK-SAME:   left_side = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   lower = #vhlo.bool_v1<true>,
-+  // CHECK-SAME:   transpose_a = #vhlo<transpose_v1 NO_TRANSPOSE>,
-+  // CHECK-SAME:   unit_diagonal = #vhlo.bool_v1<true>
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<16x16x!vhlo.f32_v1>, !vhlo.tensor_v1<16x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x16x!vhlo.f32_v1>
-+  %0 = "stablehlo.triangular_solve"(%arg0, %arg1) {
-+    left_side = true,
-+    lower = true,
-+    unit_diagonal = true,
-+    transpose_a = #stablehlo<transpose NO_TRANSPOSE>
-+  } : (tensor<16x16xf32>, tensor<16x16xf32>) -> tensor<16x16xf32>
-+  func.return %0 : tensor<16x16xf32>
-+}
-+
-+// CHECK-LABEL: "op_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_tuple(%arg0: tensor<f32>) -> tuple<tensor<f32>> {
-+  // CHECK: "vhlo.tuple_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>
-+  %0 = "stablehlo.tuple"(%arg0) : (tensor<f32>) -> tuple<tensor<f32>>
-+  func.return %0 : tuple<tensor<f32>>
-+}
-+
-+// CHECK-LABEL: "op_unary_einsum"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_unary_einsum(%arg0: tensor<8x16xf32>) -> tensor<8xf32> {
-+  //      CHECK: "vhlo.unary_einsum_v1"(%[[ARG0]]) <{
-+  // CHECK-SAME:   einsum_config = #vhlo.string_v1<"ab->a">
-+  // CHECK-SAME: }> : (!vhlo.tensor_v1<8x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x!vhlo.f32_v1>
-+  %0 = "stablehlo.unary_einsum"(%arg0) {
-+    einsum_config = "ab->a"
-+  } : (tensor<8x16xf32>) -> tensor<8xf32>
-+  func.return %0 : tensor<8xf32>
-+}
-+
-+// CHECK-LABEL: "op_uniform_dequantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_uniform_dequantize(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32> {
-+  // CHECK: "vhlo.uniform_dequantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.uniform_dequantize"(%arg0) : (tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "op_uniform_quantize"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_uniform_quantize(%arg0: tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
-+  // CHECK: "vhlo.uniform_quantize_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  %0 = "stablehlo.uniform_quantize"(%arg0) : (tensor<f32>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-+  func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
-+}
-+
-+// CHECK-LABEL: "op_while"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_while(%arg0: tensor<i1>) -> tensor<i1> {
-+  //      CHECK: "vhlo.while_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>):
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-+  // CHECK-NEXT:   }, {
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.bool_v1>)
-+  // CHECK-NEXT:     "vhlo.return_v1"(%[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.while"(%arg0) ({
-+    ^bb0(%arg1: tensor<i1>):
-+      "stablehlo.return"(%arg1) : (tensor<i1>) -> ()
-+    }, {
-+    ^bb0(%arg1: tensor<i1>):
-+      "stablehlo.return"(%arg1) : (tensor<i1>) -> ()
-+  }) : (tensor<i1>) -> tensor<i1>
-+  func.return %0: tensor<i1>
-+}
-+
-+// CHECK-LABEL: "op_xor"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @op_xor(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
-+  // CHECK: "vhlo.xor_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.xor"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// ============ TYPES ============
-+
-+// CHECK-LABEL: "type_i1"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i1(%arg0: tensor<i1>, %arg1: tensor<i1>) -> tensor<i1> {
-+  // CHECK: "vhlo.and_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bool_v1>, !vhlo.tensor_v1<!vhlo.bool_v1>) -> !vhlo.tensor_v1<!vhlo.bool_v1>
-+  %0 = "stablehlo.and"(%arg0, %arg1) : (tensor<i1>, tensor<i1>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
-+// CHECK-LABEL: "type_i2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i2(%arg0: tensor<i2>, %arg1: tensor<i2>) -> tensor<i2> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i2_v1>, !vhlo.tensor_v1<!vhlo.i2_v1>) -> !vhlo.tensor_v1<!vhlo.i2_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i2>, tensor<i2>) -> tensor<i2>
-+  func.return %0 : tensor<i2>
-+}
-+
-+// CHECK-LABEL: "type_i4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i4(%arg0: tensor<i4>, %arg1: tensor<i4>) -> tensor<i4> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i4_v1>, !vhlo.tensor_v1<!vhlo.i4_v1>) -> !vhlo.tensor_v1<!vhlo.i4_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i4>, tensor<i4>) -> tensor<i4>
-+  func.return %0 : tensor<i4>
-+}
-+
-+// CHECK-LABEL: "type_i8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i8(%arg0: tensor<i8>, %arg1: tensor<i8>) -> tensor<i8> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i8_v1>, !vhlo.tensor_v1<!vhlo.i8_v1>) -> !vhlo.tensor_v1<!vhlo.i8_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i8>, tensor<i8>) -> tensor<i8>
-+  func.return %0 : tensor<i8>
-+}
-+
-+// CHECK-LABEL: "type_i16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i16(%arg0: tensor<i16>, %arg1: tensor<i16>) -> tensor<i16> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i16_v1>, !vhlo.tensor_v1<!vhlo.i16_v1>) -> !vhlo.tensor_v1<!vhlo.i16_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i16>, tensor<i16>) -> tensor<i16>
-+  func.return %0 : tensor<i16>
-+}
-+
-+// CHECK-LABEL: "type_i32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<i32> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>) -> !vhlo.tensor_v1<!vhlo.i32_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<i32>
-+  func.return %0 : tensor<i32>
-+}
-+
-+// CHECK-LABEL: "type_i64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_i64(%arg0: tensor<i64>, %arg1: tensor<i64>) -> tensor<i64> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.i64_v1>, !vhlo.tensor_v1<!vhlo.i64_v1>) -> !vhlo.tensor_v1<!vhlo.i64_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<i64>, tensor<i64>) -> tensor<i64>
-+  func.return %0 : tensor<i64>
-+}
-+
-+// CHECK-LABEL: "type_ui2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_ui2(%arg0: tensor<ui2>, %arg1: tensor<ui2>) -> tensor<ui2> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui2_v1>, !vhlo.tensor_v1<!vhlo.ui2_v1>) -> !vhlo.tensor_v1<!vhlo.ui2_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui2>, tensor<ui2>) -> tensor<ui2>
-+  func.return %0 : tensor<ui2>
-+}
-+
-+// CHECK-LABEL: "type_ui4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_ui4(%arg0: tensor<ui4>, %arg1: tensor<ui4>) -> tensor<ui4> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui4_v1>, !vhlo.tensor_v1<!vhlo.ui4_v1>) -> !vhlo.tensor_v1<!vhlo.ui4_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui4>, tensor<ui4>) -> tensor<ui4>
-+  func.return %0 : tensor<ui4>
-+}
-+
-+// CHECK-LABEL: "type_ui8"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_ui8(%arg0: tensor<ui8>, %arg1: tensor<ui8>) -> tensor<ui8> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui8_v1>, !vhlo.tensor_v1<!vhlo.ui8_v1>) -> !vhlo.tensor_v1<!vhlo.ui8_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui8>, tensor<ui8>) -> tensor<ui8>
-+  func.return %0 : tensor<ui8>
-+}
-+
-+// CHECK-LABEL: "type_ui16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_ui16(%arg0: tensor<ui16>, %arg1: tensor<ui16>) -> tensor<ui16> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui16_v1>, !vhlo.tensor_v1<!vhlo.ui16_v1>) -> !vhlo.tensor_v1<!vhlo.ui16_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui16>, tensor<ui16>) -> tensor<ui16>
-+  func.return %0 : tensor<ui16>
-+}
-+
-+// CHECK-LABEL: "type_ui32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_ui32(%arg0: tensor<ui32>, %arg1: tensor<ui32>) -> tensor<ui32> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui32_v1>, !vhlo.tensor_v1<!vhlo.ui32_v1>) -> !vhlo.tensor_v1<!vhlo.ui32_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui32>, tensor<ui32>) -> tensor<ui32>
-+  func.return %0 : tensor<ui32>
-+}
-+
-+// CHECK-LABEL: "type_ui64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_ui64(%arg0: tensor<ui64>, %arg1: tensor<ui64>) -> tensor<ui64> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.ui64_v1>, !vhlo.tensor_v1<!vhlo.ui64_v1>) -> !vhlo.tensor_v1<!vhlo.ui64_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<ui64>, tensor<ui64>) -> tensor<ui64>
-+  func.return %0 : tensor<ui64>
-+}
-+
-+// CHECK-LABEL: "type_f4E2M1FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f4E2M1FN(%arg0: tensor<f4E2M1FN>, %arg1: tensor<f4E2M1FN>) -> tensor<f4E2M1FN> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f4E2M1FN_v1>, !vhlo.tensor_v1<!vhlo.f4E2M1FN_v1>) -> !vhlo.tensor_v1<!vhlo.f4E2M1FN_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f4E2M1FN>, tensor<f4E2M1FN>) -> tensor<f4E2M1FN>
-+  func.return %0 : tensor<f4E2M1FN>
-+}
-+
-+// CHECK-LABEL: "type_f6E2M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f6E2M3FN(%arg0: tensor<f6E2M3FN>, %arg1: tensor<f6E2M3FN>) -> tensor<f6E2M3FN> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f6E2M3FN_v1>, !vhlo.tensor_v1<!vhlo.f6E2M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f6E2M3FN_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f6E2M3FN>, tensor<f6E2M3FN>) -> tensor<f6E2M3FN>
-+  func.return %0 : tensor<f6E2M3FN>
-+}
-+
-+// CHECK-LABEL: "type_f6E3M2FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f6E3M2FN(%arg0: tensor<f6E3M2FN>, %arg1: tensor<f6E3M2FN>) -> tensor<f6E3M2FN> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f6E3M2FN_v1>, !vhlo.tensor_v1<!vhlo.f6E3M2FN_v1>) -> !vhlo.tensor_v1<!vhlo.f6E3M2FN_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f6E3M2FN>, tensor<f6E3M2FN>) -> tensor<f6E3M2FN>
-+  func.return %0 : tensor<f6E3M2FN>
-+}
-+
-+// CHECK-LABEL: "type_f8E3M4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E3M4(%arg0: tensor<f8E3M4>, %arg1: tensor<f8E3M4>) -> tensor<f8E3M4> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E3M4_v1>, !vhlo.tensor_v1<!vhlo.f8E3M4_v1>) -> !vhlo.tensor_v1<!vhlo.f8E3M4_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E3M4>, tensor<f8E3M4>) -> tensor<f8E3M4>
-+  func.return %0 : tensor<f8E3M4>
-+}
-+
-+// CHECK-LABEL: "type_f8E4M3"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E4M3(%arg0: tensor<f8E4M3>, %arg1: tensor<f8E4M3>) -> tensor<f8E4M3> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3>, tensor<f8E4M3>) -> tensor<f8E4M3>
-+  func.return %0 : tensor<f8E4M3>
-+}
-+
-+// CHECK-LABEL: "type_f8E4M3FN"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E4M3FN(%arg0: tensor<f8E4M3FN>, %arg1: tensor<f8E4M3FN>) -> tensor<f8E4M3FN> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FN_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FN>, tensor<f8E4M3FN>) -> tensor<f8E4M3FN>
-+  func.return %0 : tensor<f8E4M3FN>
-+}
-+
-+// CHECK-LABEL: "type_f8E5M2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E5M2(%arg0: tensor<f8E5M2>, %arg1: tensor<f8E5M2>) -> tensor<f8E5M2> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2>, tensor<f8E5M2>) -> tensor<f8E5M2>
-+  func.return %0 : tensor<f8E5M2>
-+}
-+
-+// CHECK-LABEL: "type_f8E4M3FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E4M3FNUZ(%arg0: tensor<f8E4M3FNUZ>, %arg1: tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3FNUZ_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3FNUZ>, tensor<f8E4M3FNUZ>) -> tensor<f8E4M3FNUZ>
-+  func.return %0 : tensor<f8E4M3FNUZ>
-+}
-+
-+// CHECK-LABEL: "type_f8E4M3B11FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E4M3B11FNUZ(%arg0: tensor<f8E4M3B11FNUZ>, %arg1: tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E4M3B11FNUZ_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E4M3B11FNUZ>, tensor<f8E4M3B11FNUZ>) -> tensor<f8E4M3B11FNUZ>
-+  func.return %0 : tensor<f8E4M3B11FNUZ>
-+}
-+
-+// CHECK-LABEL: "type_f8E5M2FNUZ"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E5M2FNUZ(%arg0: tensor<f8E5M2FNUZ>, %arg1: tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>, !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>) -> !vhlo.tensor_v1<!vhlo.f8E5M2FNUZ_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E5M2FNUZ>, tensor<f8E5M2FNUZ>) -> tensor<f8E5M2FNUZ>
-+  func.return %0 : tensor<f8E5M2FNUZ>
-+}
-+
-+// CHECK-LABEL: "type_f8E8M0FNU"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f8E8M0FNU(%arg0: tensor<f8E8M0FNU>, %arg1: tensor<f8E8M0FNU>) -> tensor<f8E8M0FNU> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f8E8M0FNU_v1>, !vhlo.tensor_v1<!vhlo.f8E8M0FNU_v1>) -> !vhlo.tensor_v1<!vhlo.f8E8M0FNU_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f8E8M0FNU>, tensor<f8E8M0FNU>) -> tensor<f8E8M0FNU>
-+  func.return %0 : tensor<f8E8M0FNU>
-+}
-+
-+// CHECK-LABEL: "type_bf16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_bf16(%arg0: tensor<bf16>, %arg1: tensor<bf16>) -> tensor<bf16> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.bf16_v1>, !vhlo.tensor_v1<!vhlo.bf16_v1>) -> !vhlo.tensor_v1<!vhlo.bf16_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<bf16>, tensor<bf16>) -> tensor<bf16>
-+  func.return %0 : tensor<bf16>
-+}
-+
-+// CHECK-LABEL: "type_f16"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f16(%arg0: tensor<f16>, %arg1: tensor<f16>) -> tensor<f16> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f16_v1>, !vhlo.tensor_v1<!vhlo.f16_v1>) -> !vhlo.tensor_v1<!vhlo.f16_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f16>, tensor<f16>) -> tensor<f16>
-+  func.return %0 : tensor<f16>
-+}
-+
-+// CHECK-LABEL: "type_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f32(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
-+// CHECK-LABEL: "type_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_f64(%arg0: tensor<f64>, %arg1: tensor<f64>) -> tensor<f64> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.f64_v1>, !vhlo.tensor_v1<!vhlo.f64_v1>) -> !vhlo.tensor_v1<!vhlo.f64_v1>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<f64>, tensor<f64>) -> tensor<f64>
-+  func.return %0 : tensor<f64>
-+}
-+
-+// CHECK-LABEL: "type_complex_f32"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_complex_f32(%arg0: tensor<complex<f32>>, %arg1: tensor<complex<f32>>) -> tensor<complex<f32>> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f32_v1>>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f32>>, tensor<complex<f32>>) -> tensor<complex<f32>>
-+  func.return %0 : tensor<complex<f32>>
-+}
-+
-+// CHECK-LABEL: "type_complex_f64"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_complex_f64(%arg0: tensor<complex<f64>>, %arg1: tensor<complex<f64>>) -> tensor<complex<f64>> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>, !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>) -> !vhlo.tensor_v1<!vhlo.complex_v1<!vhlo.f64_v1>>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<complex<f64>>, tensor<complex<f64>>) -> tensor<complex<f64>>
-+  func.return %0 : tensor<complex<f64>>
-+}
-+
-+// CHECK-LABEL: "type_tf32"
-+// CHECK: #vhlo.type_v1<!vhlo.tf31_v1>
-+func.func @type_tf32() attributes {stablehlo.attr = tf32 } {
-+  return
-+}
-+
-+// CHECK-LABEL: "type_none"
-+// CHECK: #vhlo.type_v1<!vhlo.none_v1>
-+func.func @type_none() attributes {stablehlo.attr = none } {
-+  return
-+}
-+
-+// CHECK-LABEL: "type_dynamism_ranked"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @type_dynamism_ranked(%arg0: tensor<?xf32>) -> tensor<?xf32> {
-+  // CHECK: "vhlo.abs_v1"(%[[ARG0]]) : (!vhlo.tensor_v1<?x!vhlo.f32_v1>) -> !vhlo.tensor_v1<?x!vhlo.f32_v1>
-+  %0 = "stablehlo.abs"(%arg0) : (tensor<?xf32>) -> tensor<?xf32>
-+  func.return %0 : tensor<?xf32>
-+}
-+
-+// CHECK-LABEL: "type_per_tensor_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @type_per_tensor_quantization(%arg0: tensor<!quant.uniform<i8:f32, 34.0:16>>, %arg1: tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG1]]) : (!vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>, !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>) -> !vhlo.tensor_v1<!vhlo.quant_v1<!vhlo.i8_v1:!vhlo.f32_v1, 3.400000e+01:16, -128:127, 1>>
-+  %0 = "stablehlo.add"(%arg0, %arg1) : (tensor<!quant.uniform<i8:f32, 34.0:16>>, tensor<!quant.uniform<i8:f32, 34.0:16>>) -> tensor<!quant.uniform<i8:f32, 34.0:16>>
-+  func.return %0 : tensor<!quant.uniform<i8:f32, 34.0:16>>
-+}
-+
-+// CHECK-LABEL: "type_per_axis_quantization"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @type_per_axis_quantization(%arg0: tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>) -> tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>> {
-+  // CHECK: "vhlo.add_v1"(%[[ARG0]], %[[ARG0]]) : (!vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, [3.400000e+01, 3.400000e+01], [16, 16], -128:127, 1>>, !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, [3.400000e+01, 3.400000e+01], [16, 16], -128:127, 1>>) -> !vhlo.tensor_v1<2x!vhlo.quant_per_axis_v1<!vhlo.i8_v1:!vhlo.f32_v1, 0, [3.400000e+01, 3.400000e+01], [16, 16], -128:127, 1>>
-+  %0 = stablehlo.add %arg0, %arg0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
-+  func.return %0 : tensor<2x!quant.uniform<i8:f32:0, {34.0:16, 34.0:16}>>
-+}
-+
-+//       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
-+// CHECK-LABEL: "type_token_callee"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @type_token_callee(%arg0: !stablehlo.token) -> !stablehlo.token {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.token_v1) -> ()
-+  return %arg0 : !stablehlo.token
-+}
-+
-+//       CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(!vhlo.token_v1) -> !vhlo.token_v1>>
-+// CHECK-LABEL: "type_token_caller"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @type_token_caller(%arg0: !stablehlo.token) -> !stablehlo.token {
-+  // CHECK:      "vhlo.call_v1"(%[[ARG0]]) <{callee = #vhlo.string_v1<"type_token_callee">}
-+  // CHECK-SAME: (!vhlo.token_v1) -> !vhlo.token_v1
-+  %0 = func.call @type_token_callee(%arg0) : (!stablehlo.token) -> !stablehlo.token
-+  return %0 : !stablehlo.token
-+}
-+
-+// CHECK-LABEL: "type_tuple"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @type_tuple(%arg0: tuple<tensor<f32>>) -> tuple<!stablehlo.token> {
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "foo"
-+  // CHECK: (!vhlo.tuple_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>) -> !vhlo.tuple_v1<!vhlo.token_v1>
-+  } : (tuple<tensor<f32>>) -> tuple<!stablehlo.token>
-+  return %0 : tuple<!stablehlo.token>
-+}
-+
-+// CHECK-LABEL: type_buffer_function_input_output
-+// CHECK-NEXT: (%[[ARG0:.*]]: !vhlo.buffer_v1<2x!vhlo.f32_v1>)
-+func.func @type_buffer_function_input_output(%arg0: memref<2xf32>) -> memref<2xf32> {
-+  // CHECK: "vhlo.return_v1"(%[[ARG0]]) : (!vhlo.buffer_v1<2x!vhlo.f32_v1>) -> ()
-+  func.return %arg0 : memref<2xf32>
-+}
-+
-+// CHECK-LABEL: type_buffer_special_custom_calls
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @type_buffer_special_custom_calls(%arg0: tensor<2xf32>) -> tensor<2xf32> {
-+  //               CHECK: %[[CALL0:.*]] = "vhlo.custom_call_v2"(%[[ARG0]])
-+  //          CHECK-SAME: call_target_name = #vhlo.string_v1<"Pin">
-+  //          CHECK-SAME: : (!vhlo.tensor_v1<2x!vhlo.f32_v1>) -> !vhlo.buffer_v1<2x!vhlo.f32_v1>
-+  %0 = "stablehlo.custom_call"(%arg0) {
-+    call_target_name = "Pin",
-+    api_version = 4 : i32
-+  } : (tensor<2xf32>) -> memref<2xf32>
-+  //               CHECK: %{{.*}} = "vhlo.custom_call_v2"(%[[CALL0]])
-+  //          CHECK-SAME: call_target_name = #vhlo.string_v1<"Unpin">
-+  //          CHECK-SAME: : (!vhlo.buffer_v1<2x!vhlo.f32_v1>) -> !vhlo.tensor_v1<2x!vhlo.f32_v1>
-+  %1 = "stablehlo.custom_call"(%0) {
-+    call_target_name = "Unpin",
-+    api_version = 4 : i32
-+  } : (memref<2xf32>) -> tensor<2xf32>
-+  func.return %1 : tensor<2xf32>
-+}
-+
-+// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<!vhlo.f32_v1>>>>
-+// CHECK-LABEL: "type_future"
-+func.func private @type_future() -> !stablehlo.future<tensor<f32>>
-+
-+// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<2x3x!vhlo.f32_v1>>>>
-+// CHECK-LABEL: "type_future_ranked"
-+func.func private @type_future_ranked() -> !stablehlo.future<tensor<2x3xf32>>
-+
-+// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<?x!vhlo.f32_v1>>>>
-+// CHECK-LABEL: "type_future_dynamic"
-+func.func private @type_future_dynamic() -> !stablehlo.future<tensor<?xf32>>
-+
-+// CHECK: function_type = #vhlo.type_v1<!vhlo.func_v1<(()) -> !vhlo.future_v1<!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.i32_v1>>>>
-+// CHECK-LABEL: "type_future_multiple"
-+func.func private @type_future_multiple() -> !stablehlo.future<tensor<f32>, tensor<i32>>
-+
-+// CHECK-LABEL: "op_async_start_all_reduce"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_start_all_reduce(%arg0: tensor<4x4xf32>) -> !stablehlo.future<tensor<4x4xf32>> {
-+  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
-+  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.all_reduce_v2"(%[[BARG0]]) <{
-+  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>,
-+  // CHECK-SAME:     use_global_device_ids = #vhlo.bool_v1<false>
-+  // CHECK-SAME:   }> ({
-+  // CHECK-NEXT:     ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:       %[[VAL2:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:       "vhlo.return_v1"(%[[VAL2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<4x4x!vhlo.f32_v1>>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<4x4xf32>):
-+    %1 = "stablehlo.all_reduce"(%barg0) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %2 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%2) : (tensor<f32>) -> ()
-+    }) {replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
-+    "stablehlo.return"(%1) : (tensor<4x4xf32>) -> ()
-+  }) : (tensor<4x4xf32>) -> !stablehlo.future<tensor<4x4xf32>>
-+  func.return %0: !stablehlo.future<tensor<4x4xf32>>
-+}
-+
-+// CHECK-LABEL: "op_async_start_all_gather"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_start_all_gather(%arg0: tensor<8x2xf32>) -> !stablehlo.future<tensor<8x8xf32>> {
-+  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
-+  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.all_gather_v2"(%[[BARG0]]) <{
-+  // CHECK-SAME:     all_gather_dim = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 2, 4, 6], [1, 3, 5, 7]]> : tensor<2x4xi64>>,
-+  // CHECK-SAME:     use_global_device_ids = #vhlo.bool_v1<false>
-+  // CHECK-SAME:   }> : (!vhlo.tensor_v1<8x2x!vhlo.f32_v1>) -> !vhlo.tensor_v1<8x8x!vhlo.f32_v1>
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<8x8x!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<8x2x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<8x8x!vhlo.f32_v1>>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<8x2xf32>):
-+    %1 = "stablehlo.all_gather"(%barg0) {
-+      all_gather_dim = 1 : i64,
-+      replica_groups = dense<[[0, 2, 4, 6], [1, 3, 5, 7]]> : tensor<2x4xi64>
-+    } : (tensor<8x2xf32>) -> tensor<8x8xf32>
-+    "stablehlo.return"(%1) : (tensor<8x8xf32>) -> ()
-+  }) : (tensor<8x2xf32>) -> !stablehlo.future<tensor<8x8xf32>>
-+  func.return %0: !stablehlo.future<tensor<8x8xf32>>
-+}
-+
-+// CHECK-LABEL: "op_async_start_all_to_all"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_start_all_to_all(%arg0: tensor<4x16xf32>) -> !stablehlo.future<tensor<16x4xf32>> {
-+  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
-+  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.all_to_all_v2"(%[[BARG0]]) <{
-+  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:     concat_dimension = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-+  // CHECK-SAME:     split_count = #vhlo.integer_v1<4 : i64>,
-+  // CHECK-SAME:     split_dimension = #vhlo.integer_v1<1 : i64>
-+  // CHECK-SAME:   }> : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x4x!vhlo.f32_v1>
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<16x4x!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<16x4x!vhlo.f32_v1>>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<4x16xf32>):
-+    %1 = "stablehlo.all_to_all"(%barg0) {
-+      split_dimension = 1 : i64,
-+      concat_dimension = 0 : i64,
-+      split_count = 4 : i64,
-+      replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>
-+    } : (tensor<4x16xf32>) -> tensor<16x4xf32>
-+    "stablehlo.return"(%1) : (tensor<16x4xf32>) -> ()
-+  }) : (tensor<4x16xf32>) -> !stablehlo.future<tensor<16x4xf32>>
-+  func.return %0: !stablehlo.future<tensor<16x4xf32>>
-+}
-+
-+// CHECK-LABEL: "op_async_start_collective_permute"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_start_collective_permute(%arg0: tensor<128x32xf32>) -> !stablehlo.future<tensor<128x32xf32>> {
-+  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
-+  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.collective_permute_v1"(%[[BARG0]]) <{
-+  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:     source_target_pairs = #vhlo.tensor_v1<dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>>
-+  // CHECK-SAME:   }> : (!vhlo.tensor_v1<128x32x!vhlo.f32_v1>) -> !vhlo.tensor_v1<128x32x!vhlo.f32_v1>
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<128x32x!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<128x32x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<128x32x!vhlo.f32_v1>>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<128x32xf32>):
-+    %1 = "stablehlo.collective_permute"(%barg0) {
-+      source_target_pairs = dense<[[0, 1], [1, 2], [2, 3]]> : tensor<3x2xi64>
-+    } : (tensor<128x32xf32>) -> tensor<128x32xf32>
-+    "stablehlo.return"(%1) : (tensor<128x32xf32>) -> ()
-+  }) : (tensor<128x32xf32>) -> !stablehlo.future<tensor<128x32xf32>>
-+  func.return %0: !stablehlo.future<tensor<128x32xf32>>
-+}
-+
-+// CHECK-LABEL: "op_async_start_collective_broadcast"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_start_collective_broadcast(%arg0: tensor<16x8xf32>) -> !stablehlo.future<tensor<16x8xf32>> {
-+  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
-+  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.collective_broadcast_v2"(%[[BARG0]]) <{
-+  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME:     has_dynamic_root = #vhlo.bool_v1<false>,
-+  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1]]> : tensor<1x2xi64>>
-+  // CHECK-SAME:   }> : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.tensor_v1<16x8x!vhlo.f32_v1>
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<16x8x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<16x8x!vhlo.f32_v1>>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<16x8xf32>):
-+    %1 = "stablehlo.collective_broadcast"(%barg0) {
-+      replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>
-+    } : (tensor<16x8xf32>) -> tensor<16x8xf32>
-+    "stablehlo.return"(%1) : (tensor<16x8xf32>) -> ()
-+  }) : (tensor<16x8xf32>) -> !stablehlo.future<tensor<16x8xf32>>
-+  func.return %0: !stablehlo.future<tensor<16x8xf32>>
-+}
-+
-+// CHECK-LABEL: "op_async_start_reduce_scatter"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_start_reduce_scatter(%arg0: tensor<4x16xf32>) -> !stablehlo.future<tensor<4x4xf32>> {
-+  //      CHECK: "vhlo.async_start_v1"(%[[ARG0]]) ({
-+  // CHECK-NEXT:   ^[[BB:bb.*]](%[[BARG0:.*]]: {{.*}}):
-+  // CHECK-NEXT:   %[[VAL1:.*]] = "vhlo.reduce_scatter_v1"(%[[BARG0]]) <{
-+  // CHECK-SAME:     channel_id = #vhlo.integer_v1<0 : i64>,
-+  // CHECK-SAME{LITERAL}:     replica_groups = #vhlo.tensor_v1<dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>>,
-+  // CHECK-SAME:     scatter_dimension = #vhlo.integer_v1<1 : i64>,
-+  // CHECK-SAME:     use_global_device_ids = #vhlo.bool_v1<false>
-+  // CHECK-SAME:   }> ({
-+  // CHECK-NEXT:     ^[[BB:bb.*]](%[[ARG1:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>, %[[ARG2:arg.*]]: !vhlo.tensor_v1<!vhlo.f32_v1>):
-+  // CHECK-NEXT:       %[[VAL2:.*]] = "vhlo.add_v1"(%[[ARG1]], %[[ARG2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>, !vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  // CHECK-NEXT:       "vhlo.return_v1"(%[[VAL2]]) : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT:   }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  // CHECK-NEXT:   "vhlo.return_v1"(%[[VAL1]]) : (!vhlo.tensor_v1<4x4x!vhlo.f32_v1>) -> ()
-+  // CHECK-NEXT: }) : (!vhlo.tensor_v1<4x16x!vhlo.f32_v1>) -> !vhlo.future_v1<!vhlo.tensor_v1<4x4x!vhlo.f32_v1>>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<4x16xf32>):
-+    %1 = "stablehlo.reduce_scatter"(%barg0) ({
-+      ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %2 = stablehlo.add %arg2, %arg3 : tensor<f32>
-+      "stablehlo.return"(%2) : (tensor<f32>) -> ()
-+    }) {replica_groups = dense<[[0, 1, 2, 3]]> : tensor<1x4xi64>,
-+        scatter_dimension = 1 : i64} : (tensor<4x16xf32>) -> tensor<4x4xf32>
-+    "stablehlo.return"(%1) : (tensor<4x4xf32>) -> ()
-+  }) : (tensor<4x16xf32>) -> !stablehlo.future<tensor<4x4xf32>>
-+  func.return %0: !stablehlo.future<tensor<4x4xf32>>
-+}
-+
-+// CHECK-LABEL: "op_async_done"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_async_done(%arg0: tensor<4x4xf32>) -> tensor<4x4xf32> {
-+  // CHECK: %[[VAL1:.*]] = "vhlo.async_start_v1"(%[[ARG0]])
-+  // CHECK: %[[VAL2:.*]] = "vhlo.async_done_v1"(%[[VAL1]]) : (!vhlo.future_v1<!vhlo.tensor_v1<4x4x!vhlo.f32_v1>>) -> !vhlo.tensor_v1<4x4x!vhlo.f32_v1>
-+  %0 = "stablehlo.async_start"(%arg0) ({
-+    ^bb0(%barg0: tensor<4x4xf32>):
-+    %1 = "stablehlo.all_reduce"(%barg0) ({
-+    ^bb0(%arg2: tensor<f32>, %arg3: tensor<f32>):
-+      %2 = "stablehlo.add"(%arg2, %arg3) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%2) : (tensor<f32>) -> ()
-+    }) {replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>} : (tensor<4x4xf32>) -> tensor<4x4xf32>
-+    "stablehlo.return"(%1) : (tensor<4x4xf32>) -> ()
-+  }) : (tensor<4x4xf32>) -> !stablehlo.future<tensor<4x4xf32>>
-+  %1 = "stablehlo.async_done"(%0) : (!stablehlo.future<tensor<4x4xf32>>) -> tensor<4x4xf32>
-+  func.return %1: tensor<4x4xf32>
-+}
-+
-+// ============ DEPENDENCIES  ============
-+
-+func.func @composite_target(%arg0: tensor<f32>) -> tensor<f32> {
-+  return %arg0: tensor<f32>
-+}
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_23_0.mlir.bc b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_23_0.mlir.bc
-Binary files /dev/null and stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.1_23_0.mlir.bc differ
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir b/stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir
---- stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir
-+++ stablehlo/stablehlo/tests/vhlo/stablehlo_legalize_to_vhlo.mlir
-@@ -128,6 +128,17 @@
-   func.return %0 : tensor<i1>
- }
- 
-+// CHECK-LABEL: "attr_comparison_type_weakorder"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @attr_comparison_type_weakorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+  %0 = "stablehlo.compare"(%arg0, %arg1) {
-+    comparison_direction = #stablehlo<comparison_direction EQ>,
-+    // CHECK: compare_type = #vhlo<comparison_type_v1 WEAKORDER>,
-+    compare_type = #stablehlo<comparison_type WEAKORDER>
-+  } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+  func.return %0 : tensor<i1>
-+}
-+
- // ConvDimensionNumbers aka #stablehlo.conv is covered below.
- 
- // CHECK-LABEL: "attr_custom_call_api_version_unspecified"
-@@ -749,6 +760,80 @@
-   func.return %0 : tensor<8x8x8xf32>
- }
- 
-+// CHECK-LABEL: "dot_general_algorithm_f8e4m3fn_x3"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @dot_general_algorithm_f8e4m3fn_x3(%arg0: tensor<8x8x16xbf16>, %arg1: tensor<8x16x8xbf16>) -> tensor<8x8x8xf32> {
-+//      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
-+// CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.f32_v1>,
-+// CHECK-SAME:   allow_imprecise_accumulation = #vhlo.bool_v1<false>,
-+// CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>,
-+// CHECK-SAME:   num_primitive_operations = #vhlo.integer_v1<3 : i64>,
-+// CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+// CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>
-+// CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.bf16_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.bf16_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
-+    dot_dimension_numbers = #stablehlo.dot<
-+      lhs_batching_dimensions = [0],
-+      lhs_contracting_dimensions = [2],
-+      rhs_batching_dimensions = [0],
-+      rhs_contracting_dimensions = [1]
-+    >,
-+    algorithm = #stablehlo.dot_algorithm<
-+      lhs_precision_type = f8E4M3FN,
-+      rhs_precision_type = f8E4M3FN,
-+      accumulation_type = f32,
-+      lhs_component_count = 1,
-+      rhs_component_count = 1,
-+      num_primitive_operations = 3,
-+      allow_imprecise_accumulation = false
-+    >
-+  } : (tensor<8x8x16xbf16>, tensor<8x16x8xbf16>) -> tensor<8x8x8xf32>
-+  func.return %0 : tensor<8x8x8xf32>
-+}
-+
-+// CHECK-LABEL: "dot_general_algorithm_f8e4m3fn_x4"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
-+func.func @dot_general_algorithm_f8e4m3fn_x4(%arg0: tensor<8x8x16xbf16>, %arg1: tensor<8x16x8xbf16>) -> tensor<8x8x8xf32> {
-+//      CHECK: "vhlo.dot_general_v2"(%[[ARG0]], %[[ARG1]]) <{
-+// CHECK-SAME:   accumulation_type = #vhlo.type_v1<!vhlo.f32_v1>,
-+// CHECK-SAME:   allow_imprecise_accumulation = #vhlo.bool_v1<false>,
-+// CHECK-SAME:   lhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   lhs_contracting_dimensions = #vhlo.tensor_v1<dense<2> : tensor<1xi64>>,
-+// CHECK-SAME:   lhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>,
-+// CHECK-SAME:   num_primitive_operations = #vhlo.integer_v1<4 : i64>,
-+// CHECK-SAME:   precision_config = #vhlo.array_v1<[#vhlo<precision_v1 DEFAULT>, #vhlo<precision_v1 DEFAULT>]>,
-+// CHECK-SAME:   rhs_batching_dimensions = #vhlo.tensor_v1<dense<0> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_component_count = #vhlo.integer_v1<1 : i64>,
-+// CHECK-SAME:   rhs_contracting_dimensions = #vhlo.tensor_v1<dense<1> : tensor<1xi64>>,
-+// CHECK-SAME:   rhs_precision_type = #vhlo.type_v1<!vhlo.f8E4M3FN_v1>
-+// CHECK-SAME: }> : (!vhlo.tensor_v1<8x8x16x!vhlo.bf16_v1>, !vhlo.tensor_v1<8x16x8x!vhlo.bf16_v1>) -> !vhlo.tensor_v1<8x8x8x!vhlo.f32_v1>
-+  %0 = "stablehlo.dot_general"(%arg0, %arg1) {
-+    dot_dimension_numbers = #stablehlo.dot<
-+      lhs_batching_dimensions = [0],
-+      lhs_contracting_dimensions = [2],
-+      rhs_batching_dimensions = [0],
-+      rhs_contracting_dimensions = [1]
-+    >,
-+    algorithm = #stablehlo.dot_algorithm<
-+      lhs_precision_type = f8E4M3FN,
-+      rhs_precision_type = f8E4M3FN,
-+      accumulation_type = f32,
-+      lhs_component_count = 1,
-+      rhs_component_count = 1,
-+      num_primitive_operations = 4,
-+      allow_imprecise_accumulation = false
-+    >
-+  } : (tensor<8x8x16xbf16>, tensor<8x16x8xbf16>) -> tensor<8x8x8xf32>
-+  func.return %0 : tensor<8x8x8xf32>
-+}
-+
- // CHECK-LABEL: "default_dynamic_broadcast_in_dim"
- // CHECK-NEXT: (%[[ARG0:.*]]: {{.*}}, %[[ARG1:.*]]: {{.*}})
- func.func @default_dynamic_broadcast_in_dim(%arg0: tensor<?x?xf32>, %arg1: tensor<2xindex>) -> tensor<?x?xf32> {
-@@ -1826,6 +1911,14 @@
-   func.return %0 : tensor<f32>
- }
- 
-+// CHECK-LABEL: "op_exp2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_exp2(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.exp2_v1"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.exp2"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
- // CHECK-LABEL: "op_fft"
- // CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_fft(%arg0: tensor<16xcomplex<f32>>) -> tensor<16xcomplex<f32>> {
-@@ -2006,6 +2099,14 @@
-   func.return %0 : tensor<f32>
- }
- 
-+// CHECK-LABEL: "op_log2"
-+// CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
-+func.func @op_log2(%arg0: tensor<f32>) -> tensor<f32> {
-+  // CHECK: "vhlo.log2_v1"(%[[ARG0]]) <{result_accuracy = #vhlo.result_accuracy_v1<atol = 0.000000e+00, rtol = 0.000000e+00, ulps = 0, mode = #vhlo<result_accuracy_mode_v1 DEFAULT>>}> : (!vhlo.tensor_v1<!vhlo.f32_v1>) -> !vhlo.tensor_v1<!vhlo.f32_v1>
-+  %0 = "stablehlo.log2"(%arg0) : (tensor<f32>) -> tensor<f32>
-+  func.return %0 : tensor<f32>
-+}
-+
- // CHECK-LABEL: "op_log_plus_one"
- // CHECK-NEXT: (%[[ARG0:.*]]: {{.*}})
- func.func @op_log_plus_one(%arg0: tensor<f32>) -> tensor<f32> {
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_20_0.mlir b/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_20_0.mlir
---- stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_20_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_20_0.mlir
-@@ -0,0 +1,45 @@
-+// RUN: stablehlo-opt --stablehlo-legalize-to-vhlo --vhlo-to-version='target=1.20.0' --verify-diagnostics --split-input-file %s
-+
-+// expected-error @+1 {{failed to convert VHLO to v1.20.0}}
-+module {
-+  func.func @dot_general_algorithm_f8e4m3fn_x3(%arg0: tensor<8x8x16xbf16>, %arg1: tensor<8x16x8xbf16>) -> tensor<8x8x8xf32> {
-+    // expected-error @+1 {{failed to legalize operation 'vhlo.dot_general_v2' that was explicitly marked illegal}}
-+    %0 = "stablehlo.dot_general"(%arg0, %arg1) <{
-+      dot_dimension_numbers = #stablehlo.dot<lhs_batching_dimensions = [0], rhs_batching_dimensions = [0], lhs_contracting_dimensions = [2], rhs_contracting_dimensions = [1]>,
-+      precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>],
-+      algorithm = #stablehlo.dot_algorithm<
-+        lhs_precision_type = f8E4M3FN,
-+        rhs_precision_type = f8E4M3FN,
-+        accumulation_type = f32,
-+        lhs_component_count = 1,
-+        rhs_component_count = 1,
-+        num_primitive_operations = 3,
-+        allow_imprecise_accumulation = false
-+      >
-+    }> : (tensor<8x8x16xbf16>, tensor<8x16x8xbf16>) -> tensor<8x8x8xf32>
-+    func.return %0 : tensor<8x8x8xf32>
-+  }
-+}
-+
-+// -----
-+
-+// expected-error @+1 {{failed to convert VHLO to v1.20.0}}
-+module {
-+  func.func @dot_general_algorithm_f8e4m3fn_x4(%arg0: tensor<8x8x16xbf16>, %arg1: tensor<8x16x8xbf16>) -> tensor<8x8x8xf32> {
-+    // expected-error @+1 {{failed to legalize operation 'vhlo.dot_general_v2' that was explicitly marked illegal}}
-+    %0 = "stablehlo.dot_general"(%arg0, %arg1) <{
-+      dot_dimension_numbers = #stablehlo.dot<lhs_batching_dimensions = [0], rhs_batching_dimensions = [0], lhs_contracting_dimensions = [2], rhs_contracting_dimensions = [1]>,
-+      precision_config = [#stablehlo<precision DEFAULT>, #stablehlo<precision DEFAULT>],
-+      algorithm = #stablehlo.dot_algorithm<
-+        lhs_precision_type = f8E4M3FN,
-+        rhs_precision_type = f8E4M3FN,
-+        accumulation_type = f32,
-+        lhs_component_count = 1,
-+        rhs_component_count = 1,
-+        num_primitive_operations = 4,
-+        allow_imprecise_accumulation = false
-+      >
-+    }> : (tensor<8x8x16xbf16>, tensor<8x16x8xbf16>) -> tensor<8x8x8xf32>
-+    func.return %0 : tensor<8x8x8xf32>
-+  }
-+}
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_21_0.mlir b/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_21_0.mlir
---- stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_21_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_21_0.mlir
-@@ -0,0 +1,27 @@
-+// RUN: stablehlo-opt --stablehlo-legalize-to-vhlo --vhlo-to-version='target=1.21.0' --verify-diagnostics --split-input-file %s
-+
-+// expected-error @+1 {{failed to convert VHLO to v1.21.0}}
-+module {
-+  func.func @compare_weakorder(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<i1> {
-+    // expected-error @+1 {{failed to legalize operation 'vhlo.compare_v1' that was explicitly marked illegal}}
-+    %0 = "stablehlo.compare"(%arg0, %arg1) {
-+      comparison_direction = #stablehlo<comparison_direction LT>,
-+      compare_type = #stablehlo<comparison_type WEAKORDER>
-+    } : (tensor<f32>, tensor<f32>) -> tensor<i1>
-+    func.return %0 : tensor<i1>
-+  }
-+}
-+
-+// -----
-+
-+// expected-error @+1 {{failed to convert VHLO to v1.21.0}}
-+module {
-+  func.func @discardable_weakorder_attr(%arg0: tensor<f32>, %arg1: tensor<f32>) -> tensor<f32> {
-+    // expected-error @+1 {{failed to legalize operation 'vhlo.add_v1' that was explicitly marked illegal}}
-+    %0 = "stablehlo.add"(%arg0, %arg1) {
-+      some_attr = #stablehlo<comparison_type WEAKORDER>
-+    } : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+    func.return %0 : tensor<f32>
-+  }
-+}
-+
-diff --ruN a/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_22_0.mlir b/stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_22_0.mlir
---- stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_22_0.mlir
-+++ stablehlo/stablehlo/tests/vhlo/vhlo_to_version_downgrade_invalid.1_22_0.mlir
-@@ -0,0 +1,21 @@
-+// RUN: stablehlo-opt --stablehlo-legalize-to-vhlo --vhlo-to-version='target=1.22.0' --verify-diagnostics --split-input-file %s
-+
-+// expected-error @+1 {{failed to convert VHLO to v1.22.0}}
-+module {
-+  func.func @exp2(%arg0: tensor<f32>) -> tensor<f32> {
-+    // expected-error @+1 {{failed to legalize operation 'vhlo.exp2_v1' that was explicitly marked illegal}}
-+    %0 = "stablehlo.exp2"(%arg0) : (tensor<f32>) -> tensor<f32>
-+    func.return %0 : tensor<f32>
-+  }
-+}
-+
-+// -----
-+
-+// expected-error @+1 {{failed to convert VHLO to v1.22.0}}
-+module {
-+  func.func @log2(%arg0: tensor<f32>) -> tensor<f32> {
-+    // expected-error @+1 {{failed to legalize operation 'vhlo.log2_v1' that was explicitly marked illegal}}
-+    %0 = "stablehlo.log2"(%arg0) : (tensor<f32>) -> tensor<f32>
-+    func.return %0 : tensor<f32>
-+  }
-+}
-diff --ruN a/stablehlo/stablehlo/tools/StablehloLspServerMain.cpp b/stablehlo/stablehlo/tools/StablehloLspServerMain.cpp
---- stablehlo/stablehlo/tools/StablehloLspServerMain.cpp
-+++ stablehlo/stablehlo/tools/StablehloLspServerMain.cpp
-@@ -19,7 +19,7 @@
- #include "mlir/Tools/mlir-lsp-server/MlirLspServerMain.h"
- #include "stablehlo/dialect/Register.h"
- 
--int main(int argc, char **argv) {
-+int main(int argc, char** argv) {
-   mlir::DialectRegistry registry;
-   mlir::registerAllDialects(registry);
-   mlir::registerAllExtensions(registry);
-diff --ruN a/stablehlo/stablehlo/tools/StablehloOptMain.cpp b/stablehlo/stablehlo/tools/StablehloOptMain.cpp
---- stablehlo/stablehlo/tools/StablehloOptMain.cpp
-+++ stablehlo/stablehlo/tools/StablehloOptMain.cpp
-@@ -28,7 +28,7 @@
- #include "stablehlo/transforms/Passes.h"
- #include "stablehlo/transforms/optimization/Passes.h"
- 
--int main(int argc, char **argv) {
-+int main(int argc, char** argv) {
-   mlir::registerAllPasses();
-   mlir::hlo::registerAllTestPasses();
-   mlir::stablehlo::registerPassPipelines();
-diff --ruN a/stablehlo/stablehlo/tools/StablehloTranslateMain.cpp b/stablehlo/stablehlo/tools/StablehloTranslateMain.cpp
---- stablehlo/stablehlo/tools/StablehloTranslateMain.cpp
-+++ stablehlo/stablehlo/tools/StablehloTranslateMain.cpp
-@@ -92,7 +92,7 @@
- 
- namespace {
- 
--stablehlo::Tensor makeBooleanTensor(MLIRContext *context, bool value) {
-+stablehlo::Tensor makeBooleanTensor(MLIRContext* context, bool value) {
-   auto builder = Builder(context);
-   auto type = RankedTensorType::get({}, builder.getI1Type());
-   auto res = DenseElementsAttr::get(type, builder.getBoolAttr(true));
-@@ -104,7 +104,7 @@
- //   --args=[dense<1> : tensor<2xi32>, ...], where each dense attribute is
- //   interpreted as a tensor.
- mlir::FailureOr<SmallVector<stablehlo::InterpreterValue>>
--parseInterpreterArguments(std::string argsStr, MLIRContext *context) {
-+parseInterpreterArguments(std::string argsStr, MLIRContext* context) {
-   llvm::SmallVector<stablehlo::InterpreterValue> inputs;
-   auto parseError = [&](llvm::StringRef msg) {
-     std::string usage = "--args=[dense<1> : tensor<2xi32>, ...]";
-@@ -129,7 +129,7 @@
- }
- 
- llvm::Error evalCustomCallCheckEq(stablehlo::CustomCallOp op,
--                                  stablehlo::Scope &scope) {
-+                                  stablehlo::Scope& scope) {
-   if (op->getNumOperands() != 2)
-     return stablehlo::invalidArgument("Unsupported op: %s",
-                                       debugString(op).c_str());
-@@ -157,10 +157,10 @@
-     : public stablehlo::InterpreterFallback {
-  public:
-   StablehloTranslateInterpreterFallback(
--      const std::string &probeInstrumentationDir)
-+      const std::string& probeInstrumentationDir)
-       : probeInstrumentationDir(probeInstrumentationDir) {}
--  virtual llvm::Error operator()(Operation &op, stablehlo::Scope &scope,
--                                 stablehlo::Process *process) final {
-+  virtual llvm::Error operator()(Operation& op, stablehlo::Scope& scope,
-+                                 stablehlo::Process* process) final {
-     llvm::StringRef funcName = op.getParentOfType<func::FuncOp>().getSymName();
-     if (auto customCall = dyn_cast<stablehlo::CustomCallOp>(op)) {
-       auto callTarget = customCall.getCallTargetName();
-@@ -271,7 +271,7 @@
- 
- TranslateFromMLIRRegistration interpretRegistration(
-     "interpret", "Interpreter for StableHLO",
--    [](ModuleOp module, raw_ostream &os) -> LogicalResult {
-+    [](ModuleOp module, raw_ostream& os) -> LogicalResult {
-       stablehlo::InterpreterConfiguration config;
-       config.probeInstrumentationDir = probeOutputDir.getValue();
-       config.fallback = std::make_unique<StablehloTranslateInterpreterFallback>(
-@@ -286,21 +286,21 @@
- 
-       if (interpreterPrintDense.getValue()) {
-         DenseElementsAttr denseResult;
--        for (auto &result : *results) {
-+        for (auto& result : *results) {
-           denseResult = makeDenseElementsAttr(result.getTensor());
-           os << denseResult << '\n';
-         }
-         return success();
-       }
- 
--      for (auto &result : *results) {
-+      for (auto& result : *results) {
-         result.print(os);
-         os << '\n';
-       }
- 
-       return success();
-     },
--    [](DialectRegistry &registry) {
-+    [](DialectRegistry& registry) {
-       registry.insert<func::FuncDialect>();
-       registry.insert<quant::QuantDialect>();
-       registry.insert<stablehlo::check::CheckDialect>();
-@@ -310,7 +310,7 @@
- 
- TranslateFromMLIRRegistration serializeRegistration(
-     "serialize", "Serialize StableHLO program into a portable artifact",
--    [](ModuleOp module, raw_ostream &os) -> LogicalResult {
-+    [](ModuleOp module, raw_ostream& os) -> LogicalResult {
-       std::string targetVersion = targetOption.getValue();
-       if (targetVersion == "current")
-         targetVersion = vhlo::Version::getCurrentVersion().toString();
-@@ -325,14 +325,14 @@
-       return stablehlo::serializePortableArtifact(
-           module, targetVersion, os, allowOtherDialectsOption.getValue());
-     },
--    [](DialectRegistry &registry) {
-+    [](DialectRegistry& registry) {
-       mlir::registerAllDialects(registry);
-       mlir::stablehlo::registerAllDialects(registry);
-     });
- 
- TranslateToMLIRRegistration deserializeRegistration(
-     "deserialize", "Deserialize a portable artifact into a StableHLO program",
--    [](llvm::StringRef input, mlir::MLIRContext *context) {
-+    [](llvm::StringRef input, mlir::MLIRContext* context) {
-       if (printStablehloVersion.getValue()) {
-         auto version = stablehlo::getPortableArtifactVersion(input);
-         if (failed(version)) {
-@@ -346,14 +346,14 @@
-       }
-       return stablehlo::deserializePortableArtifact(input, context);
-     },
--    [](DialectRegistry &registry) {
-+    [](DialectRegistry& registry) {
-       mlir::registerAllDialects(registry);
-       mlir::stablehlo::registerAllDialects(registry);
-     });
- 
- }  //  namespace mlir
- 
--int main(int argc, char **argv) {
-+int main(int argc, char** argv) {
-   return failed(
-       mlir::mlirTranslateMain(argc, argv, "StableHLO interpreter driver\n"));
- }
-diff --ruN a/stablehlo/stablehlo/transforms/ChloDecompositionPatternsMath.td b/stablehlo/stablehlo/transforms/ChloDecompositionPatternsMath.td
---- stablehlo/stablehlo/transforms/ChloDecompositionPatternsMath.td
-+++ stablehlo/stablehlo/transforms/ChloDecompositionPatternsMath.td
-@@ -404,23 +404,19 @@
- 
- // Arcus sine on real input:
- //
--//     arcsin(x) = 2 * arctan2(x, 1 + sqrt(1 - x * x))
-+//     arcsin(x) = atan2(x, sqrt((1 - x) * (1 + x)))
- //
- //     To avoid cancellation errors at abs(x) close to 1, we'll use
- //
- //       1 - x * x == (1 - x) * (1 + x)
- //
- def : Pat<(CHLO_AsinOp NonComplexElementType:$x),
--  (StableHLO_AddOp
--    (StableHLO_Atan2Op:$ta
--      $x,
--      (StableHLO_AddOp
--        (StableHLO_ConstantLike<"1">:$one $x),
--        (StableHLO_SqrtOp
--          (StableHLO_MulOp
--            (StableHLO_SubtractOp $one, $x),
--            (StableHLO_AddOp $one, $x)), ConstDefaultResultAccuracyAttr))),
--    $ta)>;
-+  (StableHLO_Atan2Op
-+    $x,
-+    (StableHLO_SqrtOp
-+      (StableHLO_MulOp
-+        (StableHLO_SubtractOp (StableHLO_ConstantLike<"1">:$one $x), $x),
-+        (StableHLO_AddOp $one, $x)), ConstDefaultResultAccuracyAttr))>;
- 
- // Arcus cosine on complex input
- //
-diff --ruN a/stablehlo/stablehlo/transforms/ChloDecompositionUtils.h b/stablehlo/stablehlo/transforms/ChloDecompositionUtils.h
---- stablehlo/stablehlo/transforms/ChloDecompositionUtils.h
-+++ stablehlo/stablehlo/transforms/ChloDecompositionUtils.h
-@@ -22,13 +22,13 @@
- 
- // Utility functions used in the Chlo to stablehlo legalization.
- 
--Value materializeLgamma(OpBuilder &rewriter, Location loc, ValueRange args);
-+Value materializeLgamma(OpBuilder& rewriter, Location loc, ValueRange args);
- 
--Value materializeDigamma(OpBuilder &rewriter, Location loc, ValueRange args);
-+Value materializeDigamma(OpBuilder& rewriter, Location loc, ValueRange args);
- 
--Value materializeZeta(OpBuilder &rewriter, Location loc, ValueRange args);
-+Value materializeZeta(OpBuilder& rewriter, Location loc, ValueRange args);
- 
--Value materializePolygamma(OpBuilder &rewriter, Location loc, ValueRange args);
-+Value materializePolygamma(OpBuilder& rewriter, Location loc, ValueRange args);
- 
- }  // namespace stablehlo
- }  // namespace mlir
-diff --ruN a/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp b/stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp
---- stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp
-+++ stablehlo/stablehlo/transforms/ChloLegalizeToStablehlo.cpp
-@@ -111,6 +111,8 @@
-       return mlir::stablehlo::ComparisonType::SIGNED;
-     case mlir::chlo::ComparisonType::UNSIGNED:
-       return mlir::stablehlo::ComparisonType::UNSIGNED;
-+    case mlir::chlo::ComparisonType::WEAKORDER:
-+      return mlir::stablehlo::ComparisonType::WEAKORDER;
-   }
-   return {};
- }
-@@ -1747,9 +1749,13 @@
-   acc = AddOp::create(rewriter, loc, acc, one);
-   qNegPower = PowOp::create(rewriter, loc, acc, negX);
-   Value oneLikeX = getConstantLike(rewriter, loc, 1.0, x);
-+  // Compute acc^(1-x) directly rather than as acc^-x * acc: the latter
-+  // underflows to 0 for large acc, and is inf * 0 = NaN for acc = inf.
-+  Value xMinusOne = SubtractOp::create(rewriter, loc, x, oneLikeX);
-+  Value accPowOneMinusX = PowOp::create(
-+      rewriter, loc, acc, NegOp::create(rewriter, loc, xMinusOne));
-   Value correctionEulerMaclaurin =
--      DivOp::create(rewriter, loc, MulOp::create(rewriter, loc, qNegPower, acc),
--                    SubtractOp::create(rewriter, loc, x, oneLikeX));
-+      DivOp::create(rewriter, loc, accPowOneMinusX, xMinusOne);
- 
-   // Manual reciprocal of the square root as RsqrtOp produces different results
-   Value rsqrtAcc =
-@@ -2520,7 +2526,7 @@
-     }
-     Value dotGeneral = mlir::stablehlo::DotGeneralOp::create(
-         rewriter, op.getLoc(), resultType, ValueRange{lhsMasked, rhsReshaped},
--        attributes);
-+        typename DotGeneralOp::Properties{}, attributes);
- 
-     if (i == 0) {
-       out = dotGeneral;
-@@ -2672,7 +2678,7 @@
-     }
-     Value dotGeneral = mlir::stablehlo::DotGeneralOp::create(
-         rewriter, op.getLoc(), RankedTensorType::get(resultShape, outDType),
--        {lhsMasked, rhs}, attributes);
-+        {lhsMasked, rhs}, typename DotGeneralOp::Properties{}, attributes);
- 
-     results.push_back(dotGeneral);
-     start = limit;
-@@ -2747,7 +2753,7 @@
-   }
-   rewriter.replaceOpWithNewOp<mlir::stablehlo::DotGeneralOp>(
-       op, op.getResult().getType(), ValueRange{op.getLhs(), op.getRhs()},
--      attributes);
-+      typename DotGeneralOp::Properties{}, attributes);
-   return success();
- }
- 
-@@ -2930,8 +2936,18 @@
-       SmallVector<int64_t> sliceShape(inputType.getShape().begin(),
-                                       inputType.getShape().end());
-       sliceShape[dim] = 1;
--      auto sliceType =
--          RankedTensorType::get(sliceShape, inputType.getElementType());
-+      Attribute sliceEncoding;
-+      ArrayRef<int64_t> inputBounds =
-+          hlo::encodingToBounds(inputType.getEncoding());
-+      if (!inputBounds.empty()) {
-+        SmallVector<int64_t> sliceBounds(inputBounds.begin(),
-+                                         inputBounds.end());
-+        sliceBounds[dim] = ShapedType::kDynamic;
-+        sliceEncoding =
-+            hlo::boundsToEncoding(inputType.getEncoding(), sliceBounds);
-+      }
-+      auto sliceType = RankedTensorType::get(
-+          sliceShape, inputType.getElementType(), sliceEncoding);
- 
-       Value slice;
-       if (inputType.hasStaticShape()) {
-@@ -2980,9 +2996,18 @@
-       for (int64_t d = 0; d < rank; ++d) {
-         if (d != dim) resultShape.push_back(sliceShape[d]);
-       }
-+      Attribute resultEncoding;
-+      if (!inputBounds.empty()) {
-+        SmallVector<int64_t> resultBounds(inputBounds.begin(),
-+                                          inputBounds.end());
-+        resultBounds.erase(std::next(resultBounds.begin(), dim));
-+        resultEncoding =
-+            hlo::boundsToEncoding(inputType.getEncoding(), resultBounds);
-+      }
-       Value reshaped = ReshapeOp::create(
-           rewriter, loc,
--          RankedTensorType::get(resultShape, inputType.getElementType()),
-+          RankedTensorType::get(resultShape, inputType.getElementType(),
-+                                resultEncoding),
-           slice);
-       slicedInputs.push_back(reshaped);
-     }
-@@ -3004,9 +3029,20 @@
-       SmallVector<int64_t> elementShape(outputType.getShape().begin(),
-                                         outputType.getShape().end());
-       elementShape[dim] = 1;
-+      Attribute elementEncoding;
-+      ArrayRef<int64_t> outputBounds =
-+          hlo::encodingToBounds(outputType.getEncoding());
-+      if (!outputBounds.empty()) {
-+        SmallVector<int64_t> elementBounds(outputBounds.begin(),
-+                                           outputBounds.end());
-+        elementBounds[dim] = ShapedType::kDynamic;
-+        elementEncoding =
-+            hlo::boundsToEncoding(outputType.getEncoding(), elementBounds);
-+      }
-       Value reshapedElement = ReshapeOp::create(
-           rewriter, loc,
--          RankedTensorType::get(elementShape, outputType.getElementType()),
-+          RankedTensorType::get(elementShape, outputType.getElementType(),
-+                                elementEncoding),
-           newElement);
- 
-       int64_t outRank = outputType.getRank();
-diff --ruN a/stablehlo/stablehlo/transforms/MapStablehloToVhlo.h b/stablehlo/stablehlo/transforms/MapStablehloToVhlo.h
---- stablehlo/stablehlo/transforms/MapStablehloToVhlo.h
-+++ stablehlo/stablehlo/transforms/MapStablehloToVhlo.h
-@@ -97,6 +97,7 @@
- MAP_STABLEHLO_TO_VHLO(DynamicUpdateSliceOp, V1)
- MAP_STABLEHLO_TO_VHLO(EinsumOp, V1)
- MAP_STABLEHLO_TO_VHLO(Expm1Op, V2)
-+MAP_STABLEHLO_TO_VHLO(Exp2Op, V1)
- MAP_STABLEHLO_TO_VHLO(ExpOp, V2)
- MAP_STABLEHLO_TO_VHLO(FftOp, V1)
- MAP_STABLEHLO_TO_VHLO(FloorOp, V1)
-@@ -109,6 +110,7 @@
- MAP_STABLEHLO_TO_VHLO(IotaOp, V1)
- MAP_STABLEHLO_TO_VHLO(IsFiniteOp, V1)
- MAP_STABLEHLO_TO_VHLO(Log1pOp, V2)
-+MAP_STABLEHLO_TO_VHLO(Log2Op, V1)
- MAP_STABLEHLO_TO_VHLO(LogisticOp, V2)
- MAP_STABLEHLO_TO_VHLO(LogOp, V2)
- MAP_STABLEHLO_TO_VHLO(MapOp, V1)
-diff --ruN a/stablehlo/stablehlo/transforms/PassPipelines.cpp b/stablehlo/stablehlo/transforms/PassPipelines.cpp
---- stablehlo/stablehlo/transforms/PassPipelines.cpp
-+++ stablehlo/stablehlo/transforms/PassPipelines.cpp
-@@ -24,7 +24,7 @@
- namespace mlir {
- namespace stablehlo {
- 
--void createStablehloDeserializePipeline(OpPassManager &pm) {
-+void createStablehloDeserializePipeline(OpPassManager& pm) {
-   // Convert VHLO(version x.y.z) --> VHLO(current).
-   pm.addPass(stablehlo::createVhloToVersionPass(
-       {vhlo::Version::getCurrentVersion().toString()}));
-@@ -33,14 +33,14 @@
-   pm.addPass(stablehlo::createVhloLegalizeToStablehloPass());
- }
- 
--void createChloPreSerializationPipeline(OpPassManager &pm) {
-+void createChloPreSerializationPipeline(OpPassManager& pm) {
-   // CHLO --> StableHLO+Shape --> StableHLO
-   pm.addPass(stablehlo::createChloLegalizeToStablehloPass());
-   pm.addPass(stablehlo::createShapeLegalizeToStablehloPass());
-   pm.addPass(mlir::createReconcileUnrealizedCastsPass());
- }
- 
--void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
-+void createStablehloRemoveDynamismPipeline(OpPassManager& pm,
-                                            TypeRange refinedTypes) {
-   pm.addPass(stablehlo::createStablehloRefineArgumentsPass(refinedTypes));
-   pm.addPass(stablehlo::createStablehloRefineShapesPass());
-@@ -48,7 +48,7 @@
-       stablehlo::createStablehloCanonicalizeDynamismPass());
- }
- 
--void createStablehloLowerQuantPipeline(OpPassManager &pm) {
-+void createStablehloLowerQuantPipeline(OpPassManager& pm) {
-   pm.addNestedPass<mlir::func::FuncOp>(
-       stablehlo::createStablehloLegalizeQuantToMathPass());
-   pm.addNestedPass<mlir::func::FuncOp>(
-diff --ruN a/stablehlo/stablehlo/transforms/PassUtils.h b/stablehlo/stablehlo/transforms/PassUtils.h
---- stablehlo/stablehlo/transforms/PassUtils.h
-+++ stablehlo/stablehlo/transforms/PassUtils.h
-@@ -33,7 +33,7 @@
- // Utility functions common across passes.
- 
- template <typename T>
--Attribute getScalarLike(OpBuilder &b, T constant, Type type) {
-+Attribute getScalarLike(OpBuilder& b, T constant, Type type) {
-   Type element = getElementTypeOrSelf(type);
-   if (isa<IntegerType>(element)) return b.getIntegerAttr(element, constant);
-   if (isa<FloatType>(element)) return b.getFloatAttr(element, constant);
-@@ -47,13 +47,13 @@
- // in `scalar`.
- // Returns stablehlo::ConstantOp if value type if static, else returns
- // chlo::ConstantLikeOp.
--Value getConstantLikeImpl(OpBuilder &b, Location loc, Attribute scalar,
-+Value getConstantLikeImpl(OpBuilder& b, Location loc, Attribute scalar,
-                           Value val);
- 
- // Creates a chlo::ConstantLikeOp using a splat `constant` of the same shape
- // as `val`.
- template <typename T>
--Value getConstantLike(OpBuilder &b, Location loc, T constant, Value val) {
-+Value getConstantLike(OpBuilder& b, Location loc, T constant, Value val) {
-   auto shapedTy = cast<ShapedType>(val.getType());
-   Attribute scalar = getScalarLike(b, constant, shapedTy);
-   return getConstantLikeImpl(b, loc, scalar, val);
-@@ -63,7 +63,7 @@
- // same shape as `val`.
- // The distinction between double and APFloat causes issues so need this
- // explicit template specialization.
--Value getConstantLike(OpBuilder &b, Location loc, const APFloat &constant,
-+Value getConstantLike(OpBuilder& b, Location loc, const APFloat& constant,
-                       Value val);
- 
- // Check if any of the given types are mlir::quant::QuantizedType.
-diff --ruN a/stablehlo/stablehlo/transforms/ShapeLegalizeToStablehlo.cpp b/stablehlo/stablehlo/transforms/ShapeLegalizeToStablehlo.cpp
---- stablehlo/stablehlo/transforms/ShapeLegalizeToStablehlo.cpp
-+++ stablehlo/stablehlo/transforms/ShapeLegalizeToStablehlo.cpp
-@@ -520,6 +520,7 @@
-     }
- 
-     rewriter.replaceOpWithNewOp<OpType>(op, op->getResultTypes(), operandsI32,
-+                                        typename OpType::Properties{},
-                                         op->getAttrs());
-     return success();
-   }
-diff --ruN a/stablehlo/stablehlo/transforms/StablehloCanonicalizeDynamism.cpp b/stablehlo/stablehlo/transforms/StablehloCanonicalizeDynamism.cpp
---- stablehlo/stablehlo/transforms/StablehloCanonicalizeDynamism.cpp
-+++ stablehlo/stablehlo/transforms/StablehloCanonicalizeDynamism.cpp
-@@ -97,8 +97,9 @@
-       }
-       newOperands.push_back(operand.get());
-     }
--    rewriter.replaceOpWithNewOp<CustomCallOp>(op, op.getResultTypes(),
--                                              newOperands, newAttrs);
-+    rewriter.replaceOpWithNewOp<CustomCallOp>(
-+        op, op.getResultTypes(), newOperands,
-+        typename CustomCallOp::Properties{}, newAttrs);
-     return success();
-   }
- };
-@@ -119,8 +120,11 @@
-                                          "expected static output_dimensions");
-     if (!op.getType().hasStaticShape())
-       return rewriter.notifyMatchFailure(op, "expected static result type");
--    rewriter.replaceOpWithNewOp<BroadcastInDimOp>(
--        op, op.getType(), op.getOperand(), op.getBroadcastDimensionsAttr());
-+    auto broadcastInDimOp = BroadcastInDimOp::create(
-+        rewriter, op.getLoc(), op.getType(), op.getOperand(),
-+        op.getBroadcastDimensionsAttr());
-+    broadcastInDimOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
-+    rewriter.replaceOp(op, broadcastInDimOp);
-     return success();
-   }
- };
-@@ -139,12 +143,14 @@
-         RankedTensorType::get({static_cast<int64_t>(padding.size()) / 2, 2},
-                               rewriter.getI64Type()),
-         padding);
--    rewriter.replaceOpWithNewOp<ConvolutionOp>(
--        op, op.getType(), op.getLhs(), op.getRhs(), op.getWindowStridesAttr(),
--        paddingAttr, op.getLhsDilationAttr(), op.getRhsDilationAttr(),
--        op.getWindowReversalAttr(), op.getDimensionNumbers(),
--        op.getFeatureGroupCount(), op.getBatchGroupCount(),
--        op.getPrecisionConfigAttr());
-+    auto convolutionOp = ConvolutionOp::create(
-+        rewriter, op.getLoc(), op.getType(), op.getLhs(), op.getRhs(),
-+        op.getWindowStridesAttr(), paddingAttr, op.getLhsDilationAttr(),
-+        op.getRhsDilationAttr(), op.getWindowReversalAttr(),
-+        op.getDimensionNumbers(), op.getFeatureGroupCount(),
-+        op.getBatchGroupCount(), op.getPrecisionConfigAttr());
-+    convolutionOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
-+    rewriter.replaceOp(op, convolutionOp);
-     return success();
-   }
- };
-@@ -159,10 +165,13 @@
-     SmallVector<int64_t> sliceSizes;
-     if (!succeeded(hlo::matchInts(op.getSliceSizes(), sliceSizes)))
-       return rewriter.notifyMatchFailure(op, "expected static slice_sizes");
--    rewriter.replaceOpWithNewOp<GatherOp>(
--        op, op.getType(), op.getOperand(), op.getStartIndices(),
--        op.getDimensionNumbersAttr(), rewriter.getDenseI64ArrayAttr(sliceSizes),
--        op.getIndicesAreSortedAttr());
-+    auto gatherOp =
-+        GatherOp::create(rewriter, op.getLoc(), op.getType(), op.getOperand(),
-+                         op.getStartIndices(), op.getDimensionNumbersAttr(),
-+                         rewriter.getDenseI64ArrayAttr(sliceSizes),
-+                         op.getIndicesAreSortedAttr());
-+    gatherOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
-+    rewriter.replaceOp(op, gatherOp);
-     return success();
-   }
- };
-@@ -179,8 +188,10 @@
-       return rewriter.notifyMatchFailure(op, "expected static output_shape");
-     if (!op.getType().hasStaticShape())
-       return rewriter.notifyMatchFailure(op, "expected static result type");
--    rewriter.replaceOpWithNewOp<IotaOp>(op, op.getType(),
--                                        op.getIotaDimension());
-+    auto iotaOp = IotaOp::create(rewriter, op.getLoc(), op.getType(),
-+                                 op.getIotaDimension());
-+    iotaOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
-+    rewriter.replaceOp(op, iotaOp);
-     return success();
-   }
- };
-@@ -198,9 +209,11 @@
-       return rewriter.notifyMatchFailure(op, "expected static high");
-     if (!succeeded(hlo::matchInts(op.getInteriorPadding(), interiorPadding)))
-       return rewriter.notifyMatchFailure(op, "expected static interior");
--    rewriter.replaceOpWithNewOp<PadOp>(op, op.getType(), op.getOperand(),
--                                       op.getPaddingValue(), edgePaddingLow,
--                                       edgePaddingHigh, interiorPadding);
-+    auto padOp = PadOp::create(
-+        rewriter, op.getLoc(), op.getType(), op.getOperand(),
-+        op.getPaddingValue(), edgePaddingLow, edgePaddingHigh, interiorPadding);
-+    padOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
-+    rewriter.replaceOp(op, padOp);
-     return success();
-   }
- };
-@@ -216,7 +229,10 @@
-       return rewriter.notifyMatchFailure(op, "expected static output_shape");
-     if (!op.getType().hasStaticShape())
-       return rewriter.notifyMatchFailure(op, "expected static result type");
--    rewriter.replaceOpWithNewOp<ReshapeOp>(op, op.getType(), op.getOperand());
-+    auto reshapeOp =
-+        ReshapeOp::create(rewriter, op.getLoc(), op.getType(), op.getOperand());
-+    reshapeOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
-+    rewriter.replaceOp(op, reshapeOp);
-     return success();
-   }
- };
-diff --ruN a/stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp b/stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp
---- stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp
-+++ stablehlo/stablehlo/transforms/StablehloComplexMathExpander.cpp
-@@ -24,14 +24,14 @@
- 
- namespace {
- 
--static Value getConstantLikeMaxFiniteValue(OpBuilder &b, Location loc,
-+static Value getConstantLikeMaxFiniteValue(OpBuilder& b, Location loc,
-                                            Value val) {
-   auto ty = cast<FloatType>(getElementTypeOrSelf(val.getType()));
-   return getConstantLike(
-       b, loc, llvm::APFloat::getLargest(ty.getFloatSemantics()), val);
- }
- 
--static Value getConstantLikeInfValue(OpBuilder &b, Location loc, Value val,
-+static Value getConstantLikeInfValue(OpBuilder& b, Location loc, Value val,
-                                      bool negative) {
-   auto ty = cast<FloatType>(getElementTypeOrSelf(val.getType()));
-   return getConstantLike(
-@@ -50,7 +50,7 @@
-             StablehloComplexMathExpanderPass>() {}
- 
-  public:
--  LogicalResult initialize(MLIRContext *context) override {
-+  LogicalResult initialize(MLIRContext* context) override {
-     config.setUseTopDownTraversal(true);
-     RewritePatternSet patterns_(context);
-     populateStablehloComplexMathExpanderPatterns(context, &patterns_);
-@@ -76,8 +76,8 @@
- 
- }  // namespace
- 
--void populateStablehloComplexMathExpanderPatterns(MLIRContext *context,
--                                                  RewritePatternSet *patterns) {
-+void populateStablehloComplexMathExpanderPatterns(MLIRContext* context,
-+                                                  RewritePatternSet* patterns) {
-   populateWithGenerated(*patterns);
- }
- 
-diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOps.cpp b/stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOps.cpp
---- stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOps.cpp
-+++ stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOps.cpp
-@@ -132,6 +132,14 @@
-   llvm::report_fatal_error("unsupported type");
- }
- 
-+void copyDiscardableAttrs(Operation* src, Operation* dst) {
-+  dst->setDiscardableAttrs(src->getDiscardableAttrDictionary());
-+}
-+
-+void copyDiscardableAttrs(Value src, Value dst) {
-+  copyDiscardableAttrs(src.getDefiningOp(), dst.getDefiningOp());
-+}
-+
- #include "stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.h.inc"
- }  // namespace
- 
-diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.td b/stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.td
---- stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.td
-+++ stablehlo/stablehlo/transforms/StablehloLegalizeDeprecatedOpsPatterns.td
-@@ -45,9 +45,12 @@
- 
- // def CrossReplicaSumToAllReduce : <Declared in cpp>
- 
-+def CopyDiscardableAttrs : NativeCodeCallVoid<"copyDiscardableAttrs($0, $1)">;
-+
- def DotToDotGeneral : Pat<
--  (StableHLO_DotOp $lhs, $rhs, $precision_config),
--  (StableHLO_DotGeneralOp $lhs, $rhs, (GetDefaultDotDimensionNumbers $lhs), $precision_config, (NativeCodeCall<"Attribute{}">))>;
-+  (StableHLO_DotOp:$src $lhs, $rhs, $precision_config),
-+  (StableHLO_DotGeneralOp:$dst $lhs, $rhs, (GetDefaultDotDimensionNumbers $lhs), $precision_config, (NativeCodeCall<"Attribute{}">)),
-+  [], [(CopyDiscardableAttrs $src, $dst)]>;
- 
- def UnaryEinsumToEinsum : Pat<
-   (StableHLO_UnaryEinsumOp $operand, $equation),
-diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp b/stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp
---- stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp
-+++ stablehlo/stablehlo/transforms/StablehloLegalizeQuantToMath.cpp
-@@ -636,6 +636,7 @@
-   // Execute conversion target op.
-   SmallVector<Value, 2> operands{lhsFloat32Tensor, rhsFloat32Tensor};
-   rewriter.replaceOpWithNewOp<OpType>(op, resFloat32TensorType, operands,
-+                                      typename OpType::Properties{},
-                                       op->getAttrs());
-   return success();
- }
-@@ -865,8 +866,9 @@
-                           Type resultType, Value& lhs, Value& rhs,
-                           ArrayRef<NamedAttribute> attrs,
-                           const DotLikeDimensionNumbers& dims) {
--  return stablehlo::DotGeneralOp::create(builder, loc, resultType,
--                                         ArrayRef<Value>{lhs, rhs}, attrs);
-+  return stablehlo::DotGeneralOp::create(
-+      builder, loc, resultType, ArrayRef<Value>{lhs, rhs},
-+      typename DotGeneralOp::Properties{}, attrs);
- }
- 
- // Template specialization for Convolution op.
-@@ -915,8 +917,9 @@
-       }
-     }
-   }
--  return stablehlo::ConvolutionOp::create(builder, loc, resultType,
--                                          ArrayRef<Value>{lhs, rhs}, newAttrs);
-+  return stablehlo::ConvolutionOp::create(
-+      builder, loc, resultType, ArrayRef<Value>{lhs, rhs},
-+      typename ConvolutionOp::Properties{}, newAttrs);
- }
- 
- template <typename DotLikeOp, typename DotLikeOpAdaptor>
-diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeQuantizedOpToQDQ.cpp b/stablehlo/stablehlo/transforms/StablehloLegalizeQuantizedOpToQDQ.cpp
---- stablehlo/stablehlo/transforms/StablehloLegalizeQuantizedOpToQDQ.cpp
-+++ stablehlo/stablehlo/transforms/StablehloLegalizeQuantizedOpToQDQ.cpp
-@@ -87,8 +87,9 @@
-         llvm::map_to_vector(origOp->getResultTypes(),
-                             [](Type t) { return getQuantExpressedType(t); });
-     auto origAttrs = origOp->getAttrs();
--    auto newOp = StablehloOpType::create(rewriter, op.getLoc(), newResultTypes,
--                                         dequantizedOperands, origAttrs)
-+    auto newOp = StablehloOpType::create(
-+                     rewriter, op.getLoc(), newResultTypes, dequantizedOperands,
-+                     typename StablehloOpType::Properties{}, origAttrs)
-                      .getOperation();
- 
-     SmallVector<Value> quantizedResults;
-diff --ruN a/stablehlo/stablehlo/transforms/StablehloLegalizeToVhlo.cpp b/stablehlo/stablehlo/transforms/StablehloLegalizeToVhlo.cpp
---- stablehlo/stablehlo/transforms/StablehloLegalizeToVhlo.cpp
-+++ stablehlo/stablehlo/transforms/StablehloLegalizeToVhlo.cpp
-@@ -969,8 +969,10 @@
-                 std::is_same<StablehloOpTy, stablehlo::CosineOp>::value ||
-                 std::is_same<StablehloOpTy, stablehlo::ExpOp>::value ||
-                 std::is_same<StablehloOpTy, stablehlo::Expm1Op>::value ||
-+                std::is_same<StablehloOpTy, stablehlo::Exp2Op>::value ||
-                 std::is_same<StablehloOpTy, stablehlo::LogOp>::value ||
-                 std::is_same<StablehloOpTy, stablehlo::Log1pOp>::value ||
-+                std::is_same<StablehloOpTy, stablehlo::Log2Op>::value ||
-                 std::is_same<StablehloOpTy, stablehlo::LogisticOp>::value ||
-                 std::is_same<StablehloOpTy, stablehlo::RsqrtOp>::value ||
-                 std::is_same<StablehloOpTy, stablehlo::SineOp>::value ||
-@@ -1131,15 +1133,17 @@
-     // This can almost be done in a generic fashion, except for
-     // stablehlo.case that uses a variadic number of regions which means an
-     // additional argument for the generic builder.
--    StablehloToVhloOp<StablehloOpTy> vhloOp;
-+    using VhloOpTy = StablehloToVhloOp<StablehloOpTy>;
-+    VhloOpTy vhloOp;
-     if constexpr (StablehloOpTy::template hasTrait<
-                       OpTrait::VariadicRegions>()) {
--      vhloOp = StablehloToVhloOp<StablehloOpTy>::create(
--          rewriter, stablehloOp.getLoc(), vhloTypes, vhloOperands, vhloAttrs,
--          stablehloOp.getNumRegions());
-+      vhloOp = VhloOpTy::create(rewriter, stablehloOp.getLoc(), vhloTypes,
-+                                vhloOperands, typename VhloOpTy::Properties{},
-+                                vhloAttrs, stablehloOp.getNumRegions());
-     } else {
--      vhloOp = StablehloToVhloOp<StablehloOpTy>::create(
--          rewriter, stablehloOp.getLoc(), vhloTypes, vhloOperands, vhloAttrs);
-+      vhloOp = VhloOpTy::create(rewriter, stablehloOp.getLoc(), vhloTypes,
-+                                vhloOperands, typename VhloOpTy::Properties{},
-+                                vhloAttrs);
-     }
- 
-     for (auto [stablehloRegion, vhloRegion] :
-diff --ruN a/stablehlo/stablehlo/transforms/StablehloRefineArguments.cpp b/stablehlo/stablehlo/transforms/StablehloRefineArguments.cpp
---- stablehlo/stablehlo/transforms/StablehloRefineArguments.cpp
-+++ stablehlo/stablehlo/transforms/StablehloRefineArguments.cpp
-@@ -57,7 +57,7 @@
-       builder.getI64TensorAttr(refinedType.getShape()));
-   return stablehlo::CustomCallOp::create(
-       builder, operand.getLoc(), operand.getType(),
--      ValueRange{operand, constant},
-+      ValueRange{operand, constant}, typename CustomCallOp::Properties{},
-       llvm::SmallVector<NamedAttribute>{
-           builder.getNamedAttr(
-               "call_target_name",
-diff --ruN a/stablehlo/stablehlo/transforms/VhloLegalizeToStablehlo.cpp b/stablehlo/stablehlo/transforms/VhloLegalizeToStablehlo.cpp
---- stablehlo/stablehlo/transforms/VhloLegalizeToStablehlo.cpp
-+++ stablehlo/stablehlo/transforms/VhloLegalizeToStablehlo.cpp
-@@ -1045,8 +1045,10 @@
-                 std::is_same<VhloOpTy, vhlo::CosineOpV2>::value ||
-                 std::is_same<VhloOpTy, vhlo::ExpOpV2>::value ||
-                 std::is_same<VhloOpTy, vhlo::Expm1OpV2>::value ||
-+                std::is_same<VhloOpTy, vhlo::Exp2OpV1>::value ||
-                 std::is_same<VhloOpTy, vhlo::LogOpV2>::value ||
-                 std::is_same<VhloOpTy, vhlo::Log1pOpV2>::value ||
-+                std::is_same<VhloOpTy, vhlo::Log2OpV1>::value ||
-                 std::is_same<VhloOpTy, vhlo::LogisticOpV2>::value ||
-                 std::is_same<VhloOpTy, vhlo::RsqrtOpV2>::value ||
-                 std::is_same<VhloOpTy, vhlo::SineOpV2>::value ||
-@@ -1121,15 +1123,17 @@
-     // Convert the VHLO operation to a StableHLO equivalent. This can almost be
-     // done in a generic fashion, except for ops with a variadic number of
-     // regions which means an additional argument for the generic builder.
--    VhloToStablehloOp<VhloOpTy> stablehloOp;
-+    using StablehloOpTy = VhloToStablehloOp<VhloOpTy>;
-+    StablehloOpTy stablehloOp;
-     if constexpr (VhloOpTy::template hasTrait<OpTrait::VariadicRegions>()) {
--      stablehloOp = VhloToStablehloOp<VhloOpTy>::create(
-+      stablehloOp = StablehloOpTy::create(
-           rewriter, vhloOp.getLoc(), stablehloTypes, stablehloOperands,
--          stablehloAttrs, vhloOp.getNumRegions());
-+          typename StablehloOpTy::Properties{}, stablehloAttrs,
-+          vhloOp.getNumRegions());
-     } else {
--      stablehloOp = VhloToStablehloOp<VhloOpTy>::create(
-+      stablehloOp = StablehloOpTy::create(
-           rewriter, vhloOp.getLoc(), stablehloTypes, stablehloOperands,
--          stablehloAttrs);
-+          typename StablehloOpTy::Properties{}, stablehloAttrs);
-     }
- 
-     for (auto [vhloRegion, stablehloRegion] :
-diff --ruN a/stablehlo/stablehlo/transforms/optimization/Passes.h b/stablehlo/stablehlo/transforms/optimization/Passes.h
---- stablehlo/stablehlo/transforms/optimization/Passes.h
-+++ stablehlo/stablehlo/transforms/optimization/Passes.h
-@@ -47,33 +47,33 @@
- 
- /// Collection of canonicalization patterns for StableHLO.
- void populateStablehloCanonicalizationPatterns(
--    MLIRContext *context, RewritePatternSet *patterns,
--    const StablehloAggressiveSimplificationPassOptions &options,
-+    MLIRContext* context, RewritePatternSet* patterns,
-+    const StablehloAggressiveSimplificationPassOptions& options,
-     PatternBenefit benefit = 1);
- 
--void populateStablehloCanonicalizationPatterns(MLIRContext *context,
--                                               RewritePatternSet *patterns,
-+void populateStablehloCanonicalizationPatterns(MLIRContext* context,
-+                                               RewritePatternSet* patterns,
-                                                PatternBenefit benefit = 1);
- 
- /// Collection of folding patterns for StableHLO.
- void populateStablehloAggressiveFolderPatterns(
--    MLIRContext *context, RewritePatternSet *patterns,
--    const StablehloAggressiveFolderPassOptions &options,
-+    MLIRContext* context, RewritePatternSet* patterns,
-+    const StablehloAggressiveFolderPassOptions& options,
-     PatternBenefit benefit = 1);
- 
--void populateStablehloAggressiveFolderPatterns(MLIRContext *context,
--                                               RewritePatternSet *patterns,
-+void populateStablehloAggressiveFolderPatterns(MLIRContext* context,
-+                                               RewritePatternSet* patterns,
-                                                PatternBenefit benefit = 1);
- 
- /// A subset of folding patterns for StableHLO that is necessary for shape
- /// refinement.
- void populateStablehloShapeFolderPatterns(
--    MLIRContext *context, RewritePatternSet *patterns,
--    const StablehloAggressiveFolderPassOptions &options,
-+    MLIRContext* context, RewritePatternSet* patterns,
-+    const StablehloAggressiveFolderPassOptions& options,
-     PatternBenefit benefit = 1);
- 
--void populateStablehloShapeFolderPatterns(MLIRContext *context,
--                                          RewritePatternSet *patterns,
-+void populateStablehloShapeFolderPatterns(MLIRContext* context,
-+                                          RewritePatternSet* patterns,
-                                           PatternBenefit benefit = 1);
- 
- /// Some workloads in XLA import StableHLO from HLO. Since there are a few
-@@ -81,28 +81,28 @@
- /// set of patterns brings the imported HLO back to a more canonical form
- /// without applying a full set of graph simplifications.
- void populateStablehloHloImportCanonicalizationPatterns(
--    MLIRContext *context, RewritePatternSet *patterns,
--    const StablehloAggressiveSimplificationPassOptions &options);
-+    MLIRContext* context, RewritePatternSet* patterns,
-+    const StablehloAggressiveSimplificationPassOptions& options);
- 
- void populateStablehloHloImportCanonicalizationPatterns(
--    MLIRContext *context, RewritePatternSet *patterns);
-+    MLIRContext* context, RewritePatternSet* patterns);
- 
- // If the `options` argument is provided, ensure it's not a temporary.
- void populateStablehloCanonicalizationPatterns(
--    MLIRContext *context, RewritePatternSet *patterns,
--    StablehloAggressiveSimplificationPassOptions &&options,
-+    MLIRContext* context, RewritePatternSet* patterns,
-+    StablehloAggressiveSimplificationPassOptions&& options,
-     PatternBenefit benefit = 1) = delete;
- void populateStablehloShapeFolderPatterns(
--    MLIRContext *context, RewritePatternSet *patterns,
--    StablehloAggressiveFolderPassOptions &&options,
-+    MLIRContext* context, RewritePatternSet* patterns,
-+    StablehloAggressiveFolderPassOptions&& options,
-     PatternBenefit benefit = 1) = delete;
- void populateStablehloAggressiveFolderPatterns(
--    MLIRContext *context, RewritePatternSet *patterns,
--    StablehloAggressiveFolderPassOptions &&options,
-+    MLIRContext* context, RewritePatternSet* patterns,
-+    StablehloAggressiveFolderPassOptions&& options,
-     PatternBenefit benefit = 1) = delete;
- void populateStablehloHloImportCanonicalizationPatterns(
--    MLIRContext *context, RewritePatternSet *patterns,
--    StablehloAggressiveSimplificationPassOptions &&options) = delete;
-+    MLIRContext* context, RewritePatternSet* patterns,
-+    StablehloAggressiveSimplificationPassOptions&& options) = delete;
- 
- }  // namespace stablehlo
- }  // namespace mlir
-diff --ruN a/stablehlo/stablehlo/transforms/optimization/StablehloAggressiveFolder.cpp b/stablehlo/stablehlo/transforms/optimization/StablehloAggressiveFolder.cpp
---- stablehlo/stablehlo/transforms/optimization/StablehloAggressiveFolder.cpp
-+++ stablehlo/stablehlo/transforms/optimization/StablehloAggressiveFolder.cpp
-@@ -743,6 +743,36 @@
-     ComparisonType kind;
- 
-     APInt operator()(APFloat lhs, APFloat rhs) {
-+      if (kind == ComparisonType::WEAKORDER) {
-+        bool eq = (lhs == rhs) || (lhs.isNaN() && rhs.isNaN());
-+        auto lt = [](const APFloat& a, const APFloat& b) {
-+          return (!a.isNaN() && b.isNaN()) ||
-+                 (!a.isNaN() && !b.isNaN() && (a < b));
-+        };
-+        bool result = false;
-+        switch (direction) {
-+          case ComparisonDirection::EQ:
-+            result = eq;
-+            break;
-+          case ComparisonDirection::NE:
-+            result = !eq;
-+            break;
-+          case ComparisonDirection::GE:
-+            result = lt(rhs, lhs) || eq;
-+            break;
-+          case ComparisonDirection::GT:
-+            result = lt(rhs, lhs);
-+            break;
-+          case ComparisonDirection::LE:
-+            result = lt(lhs, rhs) || eq;
-+            break;
-+          case ComparisonDirection::LT:
-+            result = lt(lhs, rhs);
-+            break;
-+        }
-+        return APInt(/*bitwidth=*/1, result);
-+      }
-+
-       if (kind != ComparisonType::FLOAT && kind != ComparisonType::TOTALORDER)
-         llvm::report_fatal_error("invalid float comparison");
- 
-@@ -1661,6 +1691,15 @@
-   }
- };
- 
-+struct FoldExp2OpPattern
-+    : public FoldUnaryOpPattern<FoldExp2OpPattern, Exp2Op> {
-+  using FoldUnaryOpPattern::FoldUnaryOpPattern;
-+
-+  static std::optional<double> EvaluateOp(double operand) {
-+    return std::exp2(operand);
-+  }
-+};
-+
- struct FoldLogOpPattern : public FoldUnaryOpPattern<FoldLogOpPattern, LogOp> {
-   using FoldUnaryOpPattern::FoldUnaryOpPattern;
- 
-@@ -1670,6 +1709,16 @@
-   }
- };
- 
-+struct FoldLog2OpPattern
-+    : public FoldUnaryOpPattern<FoldLog2OpPattern, Log2Op> {
-+  using FoldUnaryOpPattern::FoldUnaryOpPattern;
-+
-+  static std::optional<double> EvaluateOp(double operand) {
-+    if (operand <= 0.0) return std::nullopt;
-+    return std::log2(operand);
-+  }
-+};
-+
- struct FoldLogisticOpPattern
-     : public FoldUnaryOpPattern<FoldLogisticOpPattern, LogisticOp> {
-   using FoldUnaryOpPattern::FoldUnaryOpPattern;
-@@ -2222,8 +2271,10 @@
-                 FoldCosineOpPattern,                  //
-                 FoldErfOpPattern,                     //
-                 FoldExpOpPattern,                     //
-+                FoldExp2OpPattern,                    //
-                 FoldIotaOpPattern,                    //
-                 FoldLogOpPattern,                     //
-+                FoldLog2OpPattern,                    //
-                 FoldLogisticOpPattern,                //
-                 FoldNegOpPattern,                     //
-                 FoldNotOpPattern,                     //
-diff --ruN a/stablehlo/stablehlo/transforms/optimization/StablehloAggressiveSimplification.cpp b/stablehlo/stablehlo/transforms/optimization/StablehloAggressiveSimplification.cpp
---- stablehlo/stablehlo/transforms/optimization/StablehloAggressiveSimplification.cpp
-+++ stablehlo/stablehlo/transforms/optimization/StablehloAggressiveSimplification.cpp
-@@ -1766,6 +1766,7 @@
- void populateStablehloHloImportCanonicalizationPatterns(
-     MLIRContext* context, RewritePatternSet* patterns,
-     const StablehloAggressiveSimplificationPassOptions& options) {
-+  patterns->add<ConcatenateOpNoop, ConcatenateOpRemoveEmpty>(context, options);
-   patterns->add<ReshapeOp_RemoveNoop, GetTupleElementOp_UnpackTuple>(context);
-   patterns->add<TupleIsRepacking, WhileOpImplicitCapture>(context, options);
- }
-diff --ruN a/stablehlo/third_party/BUILD b/stablehlo/third_party/BUILD
---- stablehlo/third_party/BUILD
-+++ stablehlo/third_party/BUILD
-@@ -1 +0,0 @@
--licenses(["notice"])  # Apache 2.0
-\ No newline at end of file
-diff --ruN a/stablehlo/third_party/BUILD.bazel b/stablehlo/third_party/BUILD.bazel
---- stablehlo/third_party/BUILD.bazel
-+++ stablehlo/third_party/BUILD.bazel
-@@ -0,0 +1,12 @@
-+package(
-+    # copybara:uncomment default_applicable_licenses = ["//:license"],
-+    licenses = ["notice"],
-+)
-+
-+# copybara:uncomment_begin
-+# filegroup(
-+#     name = "all_files",
-+#     srcs = glob(["**"]),
-+#     visibility = ["//:__subpackages__"],
-+# )
-+# copybara:uncomment_end
-diff --ruN a/stablehlo/third_party/rules_ml_toolchain/BUILD b/stablehlo/third_party/rules_ml_toolchain/BUILD
---- stablehlo/third_party/rules_ml_toolchain/BUILD
-+++ stablehlo/third_party/rules_ml_toolchain/BUILD
-@@ -1 +0,0 @@
--exports_files(["fix_strip.patch"])
-diff --ruN a/stablehlo/third_party/rules_ml_toolchain/BUILD.bazel b/stablehlo/third_party/rules_ml_toolchain/BUILD.bazel
---- stablehlo/third_party/rules_ml_toolchain/BUILD.bazel
-+++ stablehlo/third_party/rules_ml_toolchain/BUILD.bazel
-@@ -0,0 +1,14 @@
-+package(
-+    # copybara:uncomment default_applicable_licenses = ["//:license"],
-+    licenses = ["notice"],
-+)
-+
-+# copybara:uncomment_begin
-+# filegroup(
-+#     name = "all_files",
-+#     srcs = glob(["**"]),
-+#     visibility = ["//:__subpackages__"],
-+# )
-+# copybara:uncomment_end
-+
-+exports_files(["fix_strip.patch"])
-diff --ruN a/stablehlo/third_party/xz.BUILD b/stablehlo/third_party/xz.BUILD
---- stablehlo/third_party/xz.BUILD
-+++ stablehlo/third_party/xz.BUILD
-@@ -0,0 +1,329 @@
-+# Copyright 2026 The StableHLO 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
-+#
-+#      https://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.
-+
-+load("@bazel_skylib//rules:write_file.bzl", "write_file")
-+load("@rules_cc//cc:cc_library.bzl", "cc_library")
-+
-+package(default_visibility = ["//visibility:private"])
-+
-+licenses(["unencumbered"])
-+
-+exports_files(["COPYING"])
-+
-+cc_library(
-+    name = "lzma",
-+    srcs = glob(
-+        include = [
-+            "src/**/*.h",
-+            "src/liblzma/**/*.c",
-+        ],
-+        exclude = [
-+            "src/**/*_tablegen.c",
-+            "src/liblzma/check/crc_clmul_consts_gen.c",
-+            "src/liblzma/check/sha256.c",
-+            "src/liblzma/check/crc*_small.c",
-+            "src/liblzma/common/stream_decoder_mt.c",
-+            "src/liblzma/common/stream_encoder_mt.c",
-+        ],
-+    ) + select({
-+        "@platforms//os:osx": [],
-+        "@platforms//os:ios": [],
-+        "//conditions:default": [
-+            "src/liblzma/check/sha256.c",
-+        ],
-+    }) + select({
-+        "@platforms//os:none": [],
-+        "//conditions:default": [
-+            "src/common/tuklib_cpucores.c",
-+            "src/common/tuklib_physmem.c",
-+            "src/liblzma/common/stream_decoder_mt.c",
-+            "src/liblzma/common/stream_encoder_mt.c",
-+        ],
-+    }),
-+    hdrs = glob(["src/liblzma/api/**/*.h"]),
-+    copts = select({
-+        "@platforms//os:windows": [],
-+        "//conditions:default": ["-std=c99"],
-+    }),
-+    defines = select({
-+        "@platforms//os:windows": ["LZMA_API_STATIC"],
-+        "//conditions:default": [],
-+    }),
-+    implementation_deps = [
-+        ":config",
-+        ":src_common",
-+        ":src_liblzma",
-+        ":src_liblzma_check",
-+        ":src_liblzma_common",
-+        ":src_liblzma_delta",
-+        ":src_liblzma_lz",
-+        ":src_liblzma_lzma",
-+        ":src_liblzma_rangecoder",
-+        ":src_liblzma_simple",
-+    ],
-+    linkopts = select({
-+        "@platforms//os:android": [],
-+        "@platforms//os:windows": [],
-+        "@platforms//os:none": [],
-+        "//conditions:default": ["-lpthread"],
-+    }),
-+    linkstatic = select({
-+        "@platforms//os:windows": True,
-+        "//conditions:default": False,
-+    }),
-+    local_defines = ["HAVE_CONFIG_H"],
-+    strip_include_prefix = "src/liblzma/api",
-+    visibility = ["//visibility:public"],
-+)
-+
-+write_file(
-+    name = "config_gen",
-+    out = "lzma_config/config.h",
-+    content = [
-+        "#define ASSUME_RAM 128",
-+        "#ifndef _WIN32",
-+        "#define HAVE_BSWAP_16 1",
-+        "#define HAVE_BSWAP_32 1",
-+        "#define HAVE_BSWAP_64 1",
-+        "#endif",
-+        "#if defined(__APPLE__) || defined(_WIN32)",
-+        "#undef HAVE_BYTESWAP_H",
-+        "#else",
-+        "#define HAVE_BYTESWAP_H 1",
-+        "#endif",
-+        "#undef HAVE_CAPSICUM",
-+        "#if defined(__APPLE__)",
-+        "#define HAVE_CC_SHA256_CTX 1",
-+        "#define HAVE_CC_SHA256_INIT 1",
-+        "#define HAVE_COMMONCRYPTO_COMMONDIGEST_H 1",
-+        "#endif",
-+        "#define HAVE_CHECK_CRC32 1",
-+        "#define HAVE_CHECK_CRC64 1",
-+        "#define HAVE_CHECK_SHA256 1",
-+        "#if !defined(__APPLE__)",
-+        "#define HAVE_CLOCK_GETTIME 1",
-+        "#define HAVE_DECL_CLOCK_MONOTONIC 1",
-+        "#endif",
-+        "#if defined(__ANDROID__) || defined(__APPLE__) || defined(_WIN32)",
-+        "#define HAVE_DECL_PROGRAM_INVOCATION_NAME 0",
-+        "#else",
-+        "#define HAVE_DECL_PROGRAM_INVOCATION_NAME 1",
-+        "#endif",
-+        "#define HAVE_DECODERS 1",
-+        "#define HAVE_DECODER_ARM 1",
-+        "#define HAVE_DECODER_ARM64 1",
-+        "#define HAVE_DECODER_ARMTHUMB 1",
-+        "#define HAVE_DECODER_DELTA 1",
-+        "#define HAVE_DECODER_IA64 1",
-+        "#define HAVE_DECODER_LZMA1 1",
-+        "#define HAVE_DECODER_LZMA2 1",
-+        "#define HAVE_DECODER_POWERPC 1",
-+        "#define HAVE_DECODER_SPARC 1",
-+        "#define HAVE_DECODER_X86 1",
-+        "#define HAVE_DLFCN_H 1",
-+        "#define HAVE_ENCODERS 1",
-+        "#define HAVE_ENCODER_ARM 1",
-+        "#define HAVE_ENCODER_ARM64 1",
-+        "#define HAVE_ENCODER_ARMTHUMB 1",
-+        "#define HAVE_ENCODER_DELTA 1",
-+        "#define HAVE_ENCODER_IA64 1",
-+        "#define HAVE_ENCODER_LZMA1 1",
-+        "#define HAVE_ENCODER_LZMA2 1",
-+        "#define HAVE_ENCODER_POWERPC 1",
-+        "#define HAVE_ENCODER_SPARC 1",
-+        "#define HAVE_ENCODER_X86 1",
-+        "#define HAVE_FCNTL_H 1",
-+        "#define HAVE_FUTIMENS 1",
-+        "#define HAVE_GETOPT_H 1",
-+        "#define HAVE_GETOPT_LONG 1",
-+        "#define HAVE_INTTYPES_H 1",
-+        "#define HAVE_LIMITS_H 1",
-+        "#define HAVE_MBRTOWC 1",
-+        "#define HAVE_MEMORY_H 1",
-+        "#define HAVE_MF_BT2 1",
-+        "#define HAVE_MF_BT3 1",
-+        "#define HAVE_MF_BT4 1",
-+        "#define HAVE_MF_HC3 1",
-+        "#define HAVE_MF_HC4 1",
-+        "#if !defined(__APPLE__)",
-+        "#define HAVE_POSIX_FADVISE 1",
-+        "#endif",
-+        "#if !defined(__APPLE__) && !defined(__ANDROID__)",
-+        "#define HAVE_PTHREAD_CONDATTR_SETCLOCK 1",
-+        "#endif",
-+        "#define HAVE_PTHREAD_PRIO_INHERIT 1",
-+        "#undef HAVE_SMALL",
-+        "#define HAVE_STDBOOL_H 1",
-+        "#define HAVE_STDINT_H 1",
-+        "#define HAVE_STDLIB_H 1",
-+        "#define HAVE_STRING_H 1",
-+        "#define HAVE_STRUCT_STAT_ST_ATIM_TV_NSEC 1",
-+        "#undef HAVE_SYS_CAPSICUM_H",
-+        "#ifndef _WIN32",
-+        "#define HAVE_SYS_PARAM_H 1",
-+        "#define HAVE_SYS_STAT_H 1",
-+        "#define HAVE_SYS_TIME_H 1",
-+        "#define HAVE_SYS_TYPES_H 1",
-+        "#endif",
-+        "#define HAVE_UINTPTR_T 1",
-+        "#define HAVE_UNISTD_H 1",
-+        "#define HAVE_VISIBILITY 0",
-+        "#ifndef _WIN32",
-+        "#define HAVE_WCWIDTH 1",
-+        "#endif",
-+        "#define HAVE__BOOL 1",
-+        "#undef HAVE__FUTIME",
-+        "#define LT_OBJDIR \".libs/\"",
-+        "#ifdef _WIN32",
-+        "#define MYTHREAD_VISTA 1",
-+        "#else",
-+        "#define MYTHREAD_POSIX 1",
-+        "#endif",
-+        "#define NDEBUG 1",
-+        "#define PACKAGE \"xz\"",
-+        "#define PACKAGE_BUGREPORT \"lasse.collin@tukaani.org\"",
-+        "#define PACKAGE_NAME \"XZ Utils\"",
-+        "#define PACKAGE_STRING \"XZ Utils 5.8.3\"",
-+        "#define PACKAGE_TARNAME \"xz\"",
-+        "#define PACKAGE_URL \"https://tukaani.org/xz/\"",
-+        "#define PACKAGE_VERSION \"5.8.3\"",
-+        "#define SIZEOF_SIZE_T 8",
-+        "#define STDC_HEADERS 1",
-+        "#define TUKLIB_CPUCORES_SYSCONF 1",
-+        "#define TUKLIB_FAST_UNALIGNED_ACCESS 1",
-+        "#if defined(__APPLE__)",
-+        "#undef TUKLIB_PHYSMEM_SYSCONF",
-+        "#else",
-+        "#define TUKLIB_PHYSMEM_SYSCONF 1",
-+        "#endif",
-+        "#ifndef _ALL_SOURCE",
-+        "#define _ALL_SOURCE 1",
-+        "#endif",
-+        "#ifndef _GNU_SOURCE",
-+        "#define _GNU_SOURCE 1",
-+        "#endif",
-+        "#ifndef _POSIX_PTHREAD_SEMANTICS",
-+        "#define _POSIX_PTHREAD_SEMANTICS 1",
-+        "#endif",
-+        "#ifndef _TANDEM_SOURCE",
-+        "#define _TANDEM_SOURCE 1",
-+        "#endif",
-+        "#ifndef __EXTENSIONS__",
-+        "#define __EXTENSIONS__ 1",
-+        "#endif",
-+        "#define VERSION \"5.8.3\"",
-+        "#if defined AC_APPLE_UNIVERSAL_BUILD",
-+        "#if defined __BIG_ENDIAN__",
-+        "#define WORDS_BIGENDIAN 1",
-+        "#endif",
-+        "#endif",
-+        "#ifndef _DARWIN_USE_64_BIT_INODE",
-+        "#define _DARWIN_USE_64_BIT_INODE 1",
-+        "#endif",
-+        "",
-+    ],
-+)
-+
-+cc_library(
-+    name = "config",
-+    hdrs = ["lzma_config/config.h"],
-+    strip_include_prefix = "lzma_config",
-+)
-+
-+cc_library(
-+    name = "src_common",
-+    srcs = select({
-+        "@platforms//os:none": [],
-+        "//conditions:default": [
-+            "src/common/tuklib_exit.c",
-+            "src/common/tuklib_progname.c",
-+        ],
-+    }),
-+    hdrs = glob(["src/common/*.h"]),
-+    defines = select({
-+        "@platforms//os:windows": ["LZMA_API_STATIC"],
-+        "//conditions:default": [],
-+    }),
-+    local_defines = ["HAVE_CONFIG_H"] + select({
-+        "@platforms//os:windows": ["TUKLIB_GETTEXT=0"],
-+        "//conditions:default": [],
-+    }),
-+    strip_include_prefix = "src/common",
-+    deps = [
-+        ":config",
-+        ":src_liblzma",
-+        ":src_liblzma_api",
-+        ":src_liblzma_common",
-+    ],
-+)
-+
-+cc_library(
-+    name = "src_liblzma",
-+    hdrs = glob([
-+        "src/liblzma/common/*.h",
-+        "src/liblzma/lzma/*.h",
-+    ]),
-+    strip_include_prefix = "src/liblzma",
-+)
-+
-+cc_library(
-+    name = "src_liblzma_api",
-+    hdrs = glob(
-+        include = ["src/liblzma/api/**/*.h"],
-+        exclude = ["src/liblzma/api/lzma.h"],
-+    ),
-+    strip_include_prefix = "src/liblzma/api",
-+)
-+
-+cc_library(
-+    name = "src_liblzma_check",
-+    hdrs = glob(["src/liblzma/check/*.h"]),
-+    strip_include_prefix = "src/liblzma/check",
-+)
-+
-+cc_library(
-+    name = "src_liblzma_common",
-+    hdrs = glob(["src/liblzma/common/*.h"]),
-+    strip_include_prefix = "src/liblzma/common",
-+)
-+
-+cc_library(
-+    name = "src_liblzma_delta",
-+    hdrs = glob(["src/liblzma/delta/*.h"]),
-+    strip_include_prefix = "src/liblzma/delta",
-+)
-+
-+cc_library(
-+    name = "src_liblzma_lz",
-+    hdrs = glob(["src/liblzma/lz/*.h"]),
-+    strip_include_prefix = "src/liblzma/lz",
-+)
-+
-+cc_library(
-+    name = "src_liblzma_lzma",
-+    hdrs = glob(["src/liblzma/lzma/*.h"]),
-+    strip_include_prefix = "src/liblzma/lzma",
-+)
-+
-+cc_library(
-+    name = "src_liblzma_rangecoder",
-+    hdrs = glob(["src/liblzma/rangecoder/*.h"]),
-+    strip_include_prefix = "src/liblzma/rangecoder",
-+)
-+
-+cc_library(
-+    name = "src_liblzma_simple",
-+    hdrs = glob(["src/liblzma/simple/*.h"]),
-+    strip_include_prefix = "src/liblzma/simple",
-+)
diff --git a/third_party/xla/third_party/stablehlo/workspace.bzl b/third_party/xla/third_party/stablehlo/workspace.bzl
index 7c94cab1d31..f4c1641cd39 100644
--- a/third_party/xla/third_party/stablehlo/workspace.bzl
+++ b/third_party/xla/third_party/stablehlo/workspace.bzl
@@ -18,8 +18,8 @@
 load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 
 def repo():
-    STABLEHLO_COMMIT = "b9029ac2228a12e176ea20e56f7d3dcb6bf8a019"
-    STABLEHLO_SHA256 = "16af6c8f0baacda15dd72d98dc295a213be455b8fa1899af46e144d6dae670bd"
+    STABLEHLO_COMMIT = "e7876d86fc84f8d585bc97bb2b269361cd9a2607"
+    STABLEHLO_SHA256 = "7d1dbc3d9bd51494dd5958d5076c31ef6379df254a38fad4ff92500a7bfd71d2"
 
     tf_http_archive(
         name = "stablehlo",

From 362d6d8d1c30c957b3ca905f7f51a30221498d70 Mon Sep 17 00:00:00 2001
From: Dmitri Latushko <latushko@google.com>
Date: Thu, 15 Oct 2026 17:27:26 -0700
Subject: [PATCH 0917/1000] Increase shard_count from 8 to 16 for take_test to
 prevent TSAN timeouts

PiperOrigin-RevId: 1000499925
---
 tensorflow/python/data/kernel_tests/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/python/data/kernel_tests/BUILD b/tensorflow/python/data/kernel_tests/BUILD
index edb0ee1799e..7b03b2ecc33 100644
--- a/tensorflow/python/data/kernel_tests/BUILD
+++ b/tensorflow/python/data/kernel_tests/BUILD
@@ -1251,7 +1251,7 @@ tf_py_strict_test(
     name = "take_test",
     size = "medium",
     srcs = ["take_test.py"],
-    shard_count = 8,
+    shard_count = 16,
     deps = [
         ":checkpoint_test_base",
         ":test_base",

From 3d492f1cde0851d6576ce7fbec9c24af903ad8d6 Mon Sep 17 00:00:00 2001
From: Michael Kuperstein <mkuper@google.com>
Date: Thu, 15 Oct 2026 17:35:51 -0700
Subject: [PATCH 0918/1000] [Shardy] Match tuple and GTE indices in
 `getX64CombineOnFuncResultSharding`

When tracing from an `xla.sdy.FuncResultSharding` custom-call to an `X64Combine` custom-call across intermediate `stablehlo.tuple` and `stablehlo.get_tuple_element` ops, track the `TupleOp` operand index stack and require each `GetTupleElementOp` index to match the corresponding `TupleOp` operand index.

Previously, `getX64CombineOnFuncResultSharding` inspected all `GetTupleElementOp` users of a `TupleOp` regardless of index, which only worked when each `X64Combine` pair had its own 2-element tuple. Matching `TupleOp` operand indices with `GetTupleElementOp` indices allows tracing through flattened multi-result tuples, multi-hop tuple/GTE chains, and nested tuples without cross-contaminating shardings across different results.

PiperOrigin-RevId: 1000504513
---
 .../sdy_round_trip/import_shardy_attrs.cc     |  39 +++--
 ...rt_pipeline_sharding_in_frontend_attr.mlir |  45 +++--
 .../sdy_round_trip_import_shardy_attrs.mlir   | 164 ++++++++++++++++++
 3 files changed, 224 insertions(+), 24 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc
index 58083e905f3..b76583e40ef 100644
--- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc
+++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/import_shardy_attrs.cc
@@ -119,24 +119,43 @@ CustomCallOp getX64CombineUser(Operation* op) {
   return nullptr;
 }
 
-CustomCallOp getX64CombineOnFuncResultSharding(
-    CustomCallOp funcResultSharding) {
-  if (auto combineOp = getX64CombineUser(funcResultSharding)) {
-    return combineOp;
+// Finds the downstream `X64Combine` for a split 32-bit half.
+// Because multiple results may share the same tuple, `tuplePath` records
+// which slot this value entered so we only follow GTEs that extract that slot.
+CustomCallOp getX64CombineOnOp(Operation* op,
+                               llvm::SmallVector<int64_t, 4>& tuplePath) {
+  if (tuplePath.empty()) {
+    if (auto combineOp = getX64CombineUser(op)) {
+      return combineOp;
+    }
   }
-  for (Operation* user : funcResultSharding->getUsers()) {
-    auto tupleOp = mlir::dyn_cast<stablehlo::TupleOp>(user);
-    if (!tupleOp) continue;
-    for (Operation* gteOp : tupleOp->getUsers()) {
-      if (!mlir::isa<stablehlo::GetTupleElementOp>(gteOp)) continue;
-      if (auto combineOp = getX64CombineUser(gteOp)) {
+  for (mlir::OpOperand& use : op->getUses()) {
+    if (auto tupleOp = mlir::dyn_cast<stablehlo::TupleOp>(use.getOwner())) {
+      tuplePath.push_back(use.getOperandNumber());
+      if (CustomCallOp combineOp = getX64CombineOnOp(tupleOp, tuplePath)) {
         return combineOp;
       }
+      tuplePath.pop_back();
+    } else if (auto gteOp =
+                   mlir::dyn_cast<stablehlo::GetTupleElementOp>(use.getOwner());
+               gteOp && !tuplePath.empty() &&
+               gteOp.getIndex() == tuplePath.back()) {
+      int64_t index = tuplePath.pop_back_val();
+      if (CustomCallOp combineOp = getX64CombineOnOp(gteOp, tuplePath)) {
+        return combineOp;
+      }
+      tuplePath.push_back(index);
     }
   }
   return nullptr;
 }
 
+CustomCallOp getX64CombineOnFuncResultSharding(
+    CustomCallOp funcResultSharding) {
+  llvm::SmallVector<int64_t, 4> tuplePath;
+  return getX64CombineOnOp(funcResultSharding, tuplePath);
+}
+
 // TODO(kostiantynl): b/448858211 when API is fixed, use
 // sharding.openShardingDims() instead.
 TensorShardingAttr openShardingDims(TensorShardingAttr sharding) {
diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_sharding_in_frontend_attr.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_sharding_in_frontend_attr.mlir
index c797487cf83..404992919bf 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_sharding_in_frontend_attr.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_pipeline_sharding_in_frontend_attr.mlir
@@ -66,28 +66,45 @@ module @multiple_func_result_shardings attributes {mhlo.frontend_attributes = {x
     return %4 : tensor<16xi64>
   }
 
-  // CHECK-LABEL: func @func_result_sharding_used_by_x64_combine_tuple(%arg0: tensor<16xi64>)
-  // CHECK-SAME:    -> (tensor<16xi64> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}]>}) {
+  // CHECK-LABEL: func @func_result_sharding_used_by_x64_combine_tuple(
+  // CHECK-SAME:    %arg0: tensor<16xi64>, %arg1: tensor<16xi64>)
+  // CHECK-SAME:    -> (tensor<16xi64> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}]>},
+  // CHECK-SAME:        tensor<16xi64> {sdy.sharding = #sdy.sharding<@mesh, [{"b"}]>}) {
   func.func @func_result_sharding_used_by_x64_combine_tuple(
-    %arg0: tensor<16xi64>) -> tensor<16xi64> {
-    // CHECK-NEXT: %[[SPLIT_LOW:.*]] = stablehlo.custom_call @X64SplitLow(%arg0) : (tensor<16xi64>) -> tensor<16xui32>
-    // CHECK-NEXT: %[[SPLIT_HIGH:.*]] = stablehlo.custom_call @X64SplitHigh(%arg0) : (tensor<16xi64>) -> tensor<16xui32>
-    // CHECK-NEXT: %[[COMBINE:.*]] = stablehlo.custom_call @X64Combine(%[[SPLIT_LOW]], %[[SPLIT_HIGH]])
+    %arg0: tensor<16xi64>, %arg1: tensor<16xi64>) -> (tensor<16xi64>, tensor<16xi64>) {
+    // CHECK-NEXT: %[[SPLIT_LOW_0:.*]] = stablehlo.custom_call @X64SplitLow(%arg0) : (tensor<16xi64>) -> tensor<16xui32>
+    // CHECK-NEXT: %[[SPLIT_HIGH_0:.*]] = stablehlo.custom_call @X64SplitHigh(%arg0) : (tensor<16xi64>) -> tensor<16xui32>
+    // CHECK-NEXT: %[[SPLIT_LOW_1:.*]] = stablehlo.custom_call @X64SplitLow(%arg1) : (tensor<16xi64>) -> tensor<16xui32>
+    // CHECK-NEXT: %[[SPLIT_HIGH_1:.*]] = stablehlo.custom_call @X64SplitHigh(%arg1) : (tensor<16xi64>) -> tensor<16xui32>
+    // CHECK-NEXT: %[[COMBINE_0:.*]] = stablehlo.custom_call @X64Combine(%[[SPLIT_LOW_0]], %[[SPLIT_HIGH_0]])
     // CHECK-SAME:   {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}]>]>} : (tensor<16xui32>, tensor<16xui32>) -> tensor<16xi64>
-    // CHECK-NEXT: return %[[COMBINE]]
+    // CHECK-NEXT: %[[COMBINE_1:.*]] = stablehlo.custom_call @X64Combine(%[[SPLIT_LOW_1]], %[[SPLIT_HIGH_1]])
+    // CHECK-SAME:   {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}]>]>} : (tensor<16xui32>, tensor<16xui32>) -> tensor<16xi64>
+    // CHECK-NEXT: return %[[COMBINE_0]], %[[COMBINE_1]]
     %0 = stablehlo.custom_call @X64SplitLow(%arg0) : (tensor<16xi64>) -> tensor<16xui32>
     %1 = stablehlo.custom_call @X64SplitHigh(%arg0) : (tensor<16xi64>) -> tensor<16xui32>
-    %2 = stablehlo.custom_call @xla.sdy.FuncResultSharding(%0)
+    %2 = stablehlo.custom_call @X64SplitLow(%arg1) : (tensor<16xi64>) -> tensor<16xui32>
+    %3 = stablehlo.custom_call @X64SplitHigh(%arg1) : (tensor<16xi64>) -> tensor<16xui32>
+    %4 = stablehlo.custom_call @xla.sdy.FuncResultSharding(%0)
         {has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@mesh, [{\"a\"}]>]>"}} :
         (tensor<16xui32>) -> tensor<16xui32>
-    %3 = stablehlo.custom_call @xla.sdy.FuncResultSharding(%1)
+    %5 = stablehlo.custom_call @xla.sdy.FuncResultSharding(%1)
         {has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@mesh, [{\"a\"}]>]>"}} :
         (tensor<16xui32>) -> tensor<16xui32>
-    %4 = stablehlo.tuple %2, %3 : tuple<tensor<16xui32>, tensor<16xui32>>
-    %5 = stablehlo.get_tuple_element %4[0] : (tuple<tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
-    %6 = stablehlo.get_tuple_element %4[1] : (tuple<tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
-    %7 = stablehlo.custom_call @X64Combine(%5, %6) : (tensor<16xui32>, tensor<16xui32>) -> tensor<16xi64>
-    return %7 : tensor<16xi64>
+    %6 = stablehlo.custom_call @xla.sdy.FuncResultSharding(%2)
+        {has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@mesh, [{\"b\"}]>]>"}} :
+        (tensor<16xui32>) -> tensor<16xui32>
+    %7 = stablehlo.custom_call @xla.sdy.FuncResultSharding(%3)
+        {has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@mesh, [{\"b\"}]>]>"}} :
+        (tensor<16xui32>) -> tensor<16xui32>
+    %8 = stablehlo.tuple %4, %5, %6, %7 : tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>
+    %9 = stablehlo.get_tuple_element %8[0] : (tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %10 = stablehlo.get_tuple_element %8[1] : (tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %11 = stablehlo.custom_call @X64Combine(%9, %10) : (tensor<16xui32>, tensor<16xui32>) -> tensor<16xi64>
+    %12 = stablehlo.get_tuple_element %8[2] : (tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %13 = stablehlo.get_tuple_element %8[3] : (tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %14 = stablehlo.custom_call @X64Combine(%12, %13) : (tensor<16xui32>, tensor<16xui32>) -> tensor<16xi64>
+    return %11, %14 : tensor<16xi64>, tensor<16xi64>
   }
 
   // This might happen due to inlined funcs that originally had result shardings
diff --git a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_shardy_attrs.mlir b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_shardy_attrs.mlir
index dd00f98bb1e..f3661345e07 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_shardy_attrs.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/sdy_round_trip_import_shardy_attrs.mlir
@@ -267,3 +267,167 @@ module @sparse_offload_unreduced_max_module attributes {
   }
 }
 
+// -----
+
+// CHECK-LABEL: module @x64_combine_flat_tuple_func_result_sharding
+// CHECK:         sdy.mesh @mesh = <["a"=2, "b"=2]>
+module @x64_combine_flat_tuple_func_result_sharding attributes {
+  mhlo.frontend_attributes = {
+    xla.sdy.meshes = "{mesh = #sdy.mesh<[\"a\"=2, \"b\"=2]>}"
+  }
+} {
+  // CHECK-LABEL: func @func_result_sharding_used_by_x64_combine_tuple(
+  // CHECK-SAME:      %arg0: tensor<16xi64>, %arg1: tensor<16xi64>)
+  // CHECK-SAME:      -> (tensor<16xi64> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}]>},
+  // CHECK-SAME:          tensor<16xi64> {sdy.sharding = #sdy.sharding<@mesh, [{"b"}]>}) {
+  func.func @func_result_sharding_used_by_x64_combine_tuple(
+      %arg0: tensor<16xi64>, %arg1: tensor<16xi64>) -> (tensor<16xi64>, tensor<16xi64>) {
+    // CHECK-NEXT: %[[SPLIT_LOW_0:.*]] = stablehlo.custom_call @X64SplitLow(%arg0) : (tensor<16xi64>) -> tensor<16xui32>
+    // CHECK-NEXT: %[[SPLIT_HIGH_0:.*]] = stablehlo.custom_call @X64SplitHigh(%arg0) : (tensor<16xi64>) -> tensor<16xui32>
+    // CHECK-NEXT: %[[SPLIT_LOW_1:.*]] = stablehlo.custom_call @X64SplitLow(%arg1) : (tensor<16xi64>) -> tensor<16xui32>
+    // CHECK-NEXT: %[[SPLIT_HIGH_1:.*]] = stablehlo.custom_call @X64SplitHigh(%arg1) : (tensor<16xi64>) -> tensor<16xui32>
+    // CHECK-NEXT: %[[TUPLE:.*]] = stablehlo.tuple %[[SPLIT_LOW_0]], %[[SPLIT_HIGH_0]], %[[SPLIT_LOW_1]], %[[SPLIT_HIGH_1]]
+    // CHECK-NEXT: %[[GTE_0:.*]] = stablehlo.get_tuple_element %[[TUPLE]][0]
+    // CHECK-NEXT: %[[GTE_1:.*]] = stablehlo.get_tuple_element %[[TUPLE]][1]
+    // CHECK-NEXT: %[[COMBINE_0:.*]] = stablehlo.custom_call @X64Combine(%[[GTE_0]], %[[GTE_1]])
+    // CHECK-SAME:   {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}]>]>} : (tensor<16xui32>, tensor<16xui32>) -> tensor<16xi64>
+    // CHECK-NEXT: %[[GTE_2:.*]] = stablehlo.get_tuple_element %[[TUPLE]][2]
+    // CHECK-NEXT: %[[GTE_3:.*]] = stablehlo.get_tuple_element %[[TUPLE]][3]
+    // CHECK-NEXT: %[[COMBINE_1:.*]] = stablehlo.custom_call @X64Combine(%[[GTE_2]], %[[GTE_3]])
+    // CHECK-SAME:   {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}]>]>} : (tensor<16xui32>, tensor<16xui32>) -> tensor<16xi64>
+    // CHECK-NEXT: return %[[COMBINE_0]], %[[COMBINE_1]]
+    %0 = stablehlo.custom_call @X64SplitLow(%arg0) : (tensor<16xi64>) -> tensor<16xui32>
+    %1 = stablehlo.custom_call @X64SplitHigh(%arg0) : (tensor<16xi64>) -> tensor<16xui32>
+    %2 = stablehlo.custom_call @X64SplitLow(%arg1) : (tensor<16xi64>) -> tensor<16xui32>
+    %3 = stablehlo.custom_call @X64SplitHigh(%arg1) : (tensor<16xi64>) -> tensor<16xui32>
+    %4 = stablehlo.custom_call @xla.sdy.FuncResultSharding(%0)
+        {has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@mesh, [{\"a\"}]>]>"}} :
+        (tensor<16xui32>) -> tensor<16xui32>
+    %5 = stablehlo.custom_call @xla.sdy.FuncResultSharding(%1)
+        {has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@mesh, [{\"a\"}]>]>"}} :
+        (tensor<16xui32>) -> tensor<16xui32>
+    %6 = stablehlo.custom_call @xla.sdy.FuncResultSharding(%2)
+        {has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@mesh, [{\"b\"}]>]>"}} :
+        (tensor<16xui32>) -> tensor<16xui32>
+    %7 = stablehlo.custom_call @xla.sdy.FuncResultSharding(%3)
+        {has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@mesh, [{\"b\"}]>]>"}} :
+        (tensor<16xui32>) -> tensor<16xui32>
+    %8 = stablehlo.tuple %4, %5, %6, %7 : tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>
+    %9 = stablehlo.get_tuple_element %8[0] : (tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %10 = stablehlo.get_tuple_element %8[1] : (tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %11 = stablehlo.custom_call @X64Combine(%9, %10) : (tensor<16xui32>, tensor<16xui32>) -> tensor<16xi64>
+    %12 = stablehlo.get_tuple_element %8[2] : (tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %13 = stablehlo.get_tuple_element %8[3] : (tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %14 = stablehlo.custom_call @X64Combine(%12, %13) : (tensor<16xui32>, tensor<16xui32>) -> tensor<16xi64>
+    return %11, %14 : tensor<16xi64>, tensor<16xi64>
+  }
+
+  // CHECK-LABEL: func @func_result_sharding_used_by_x64_combine_multi_hop_tuple(
+  // CHECK-SAME:      %arg0: tensor<16xi64>, %arg1: tensor<16xi64>)
+  // CHECK-SAME:      -> (tensor<16xi64> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}]>},
+  // CHECK-SAME:          tensor<16xi64> {sdy.sharding = #sdy.sharding<@mesh, [{"b"}]>}) {
+  func.func @func_result_sharding_used_by_x64_combine_multi_hop_tuple(
+      %arg0: tensor<16xi64>, %arg1: tensor<16xi64>) -> (tensor<16xi64>, tensor<16xi64>) {
+    // CHECK-NEXT: %[[SPLIT_LOW_0:.*]] = stablehlo.custom_call @X64SplitLow(%arg0) : (tensor<16xi64>) -> tensor<16xui32>
+    // CHECK-NEXT: %[[SPLIT_HIGH_0:.*]] = stablehlo.custom_call @X64SplitHigh(%arg0) : (tensor<16xi64>) -> tensor<16xui32>
+    // CHECK-NEXT: %[[SPLIT_LOW_1:.*]] = stablehlo.custom_call @X64SplitLow(%arg1) : (tensor<16xi64>) -> tensor<16xui32>
+    // CHECK-NEXT: %[[SPLIT_HIGH_1:.*]] = stablehlo.custom_call @X64SplitHigh(%arg1) : (tensor<16xi64>) -> tensor<16xui32>
+    // CHECK-NEXT: %[[T0:.*]] = stablehlo.tuple %[[SPLIT_LOW_0]], %[[SPLIT_HIGH_0]]
+    // CHECK-NEXT: %[[T1:.*]] = stablehlo.tuple %[[SPLIT_LOW_1]], %[[SPLIT_HIGH_1]]
+    // CHECK-NEXT: %[[T0_0:.*]] = stablehlo.get_tuple_element %[[T0]][0]
+    // CHECK-NEXT: %[[T0_1:.*]] = stablehlo.get_tuple_element %[[T0]][1]
+    // CHECK-NEXT: %[[T1_0:.*]] = stablehlo.get_tuple_element %[[T1]][0]
+    // CHECK-NEXT: %[[T1_1:.*]] = stablehlo.get_tuple_element %[[T1]][1]
+    // CHECK-NEXT: %[[FLAT_TUPLE:.*]] = stablehlo.tuple %[[T0_0]], %[[T0_1]], %[[T1_0]], %[[T1_1]]
+    // CHECK-NEXT: %[[GTE_0:.*]] = stablehlo.get_tuple_element %[[FLAT_TUPLE]][0]
+    // CHECK-NEXT: %[[GTE_1:.*]] = stablehlo.get_tuple_element %[[FLAT_TUPLE]][1]
+    // CHECK-NEXT: %[[COMBINE_0:.*]] = stablehlo.custom_call @X64Combine(%[[GTE_0]], %[[GTE_1]])
+    // CHECK-SAME:   {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}]>]>} : (tensor<16xui32>, tensor<16xui32>) -> tensor<16xi64>
+    // CHECK-NEXT: %[[GTE_2:.*]] = stablehlo.get_tuple_element %[[FLAT_TUPLE]][2]
+    // CHECK-NEXT: %[[GTE_3:.*]] = stablehlo.get_tuple_element %[[FLAT_TUPLE]][3]
+    // CHECK-NEXT: %[[COMBINE_1:.*]] = stablehlo.custom_call @X64Combine(%[[GTE_2]], %[[GTE_3]])
+    // CHECK-SAME:   {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}]>]>} : (tensor<16xui32>, tensor<16xui32>) -> tensor<16xi64>
+    // CHECK-NEXT: return %[[COMBINE_0]], %[[COMBINE_1]]
+    %0 = stablehlo.custom_call @X64SplitLow(%arg0) : (tensor<16xi64>) -> tensor<16xui32>
+    %1 = stablehlo.custom_call @X64SplitHigh(%arg0) : (tensor<16xi64>) -> tensor<16xui32>
+    %2 = stablehlo.custom_call @X64SplitLow(%arg1) : (tensor<16xi64>) -> tensor<16xui32>
+    %3 = stablehlo.custom_call @X64SplitHigh(%arg1) : (tensor<16xi64>) -> tensor<16xui32>
+    %4 = stablehlo.custom_call @xla.sdy.FuncResultSharding(%0)
+        {has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@mesh, [{\"a\"}]>]>"}} :
+        (tensor<16xui32>) -> tensor<16xui32>
+    %5 = stablehlo.custom_call @xla.sdy.FuncResultSharding(%1)
+        {has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@mesh, [{\"a\"}]>]>"}} :
+        (tensor<16xui32>) -> tensor<16xui32>
+    %6 = stablehlo.custom_call @xla.sdy.FuncResultSharding(%2)
+        {has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@mesh, [{\"b\"}]>]>"}} :
+        (tensor<16xui32>) -> tensor<16xui32>
+    %7 = stablehlo.custom_call @xla.sdy.FuncResultSharding(%3)
+        {has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@mesh, [{\"b\"}]>]>"}} :
+        (tensor<16xui32>) -> tensor<16xui32>
+    %8 = stablehlo.tuple %4, %5 : tuple<tensor<16xui32>, tensor<16xui32>>
+    %9 = stablehlo.tuple %6, %7 : tuple<tensor<16xui32>, tensor<16xui32>>
+    %10 = stablehlo.get_tuple_element %8[0] : (tuple<tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %11 = stablehlo.get_tuple_element %8[1] : (tuple<tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %12 = stablehlo.get_tuple_element %9[0] : (tuple<tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %13 = stablehlo.get_tuple_element %9[1] : (tuple<tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %14 = stablehlo.tuple %10, %11, %12, %13 : tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>
+    %15 = stablehlo.get_tuple_element %14[0] : (tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %16 = stablehlo.get_tuple_element %14[1] : (tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %17 = stablehlo.custom_call @X64Combine(%15, %16) : (tensor<16xui32>, tensor<16xui32>) -> tensor<16xi64>
+    %18 = stablehlo.get_tuple_element %14[2] : (tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %19 = stablehlo.get_tuple_element %14[3] : (tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %20 = stablehlo.custom_call @X64Combine(%18, %19) : (tensor<16xui32>, tensor<16xui32>) -> tensor<16xi64>
+    return %17, %20 : tensor<16xi64>, tensor<16xi64>
+  }
+
+  // CHECK-LABEL: func @func_result_sharding_used_by_x64_combine_nested_tuple(
+  // CHECK-SAME:      %arg0: tensor<16xi64>, %arg1: tensor<16xi64>)
+  // CHECK-SAME:      -> (tensor<16xi64> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}]>},
+  // CHECK-SAME:          tensor<16xi64> {sdy.sharding = #sdy.sharding<@mesh, [{"b"}]>}) {
+  func.func @func_result_sharding_used_by_x64_combine_nested_tuple(
+      %arg0: tensor<16xi64>, %arg1: tensor<16xi64>) -> (tensor<16xi64>, tensor<16xi64>) {
+    // CHECK-NEXT: %[[SPLIT_LOW_0:.*]] = stablehlo.custom_call @X64SplitLow(%arg0) : (tensor<16xi64>) -> tensor<16xui32>
+    // CHECK-NEXT: %[[SPLIT_HIGH_0:.*]] = stablehlo.custom_call @X64SplitHigh(%arg0) : (tensor<16xi64>) -> tensor<16xui32>
+    // CHECK-NEXT: %[[SPLIT_LOW_1:.*]] = stablehlo.custom_call @X64SplitLow(%arg1) : (tensor<16xi64>) -> tensor<16xui32>
+    // CHECK-NEXT: %[[SPLIT_HIGH_1:.*]] = stablehlo.custom_call @X64SplitHigh(%arg1) : (tensor<16xi64>) -> tensor<16xui32>
+    // CHECK-NEXT: %[[INNER:.*]] = stablehlo.tuple %[[SPLIT_LOW_0]], %[[SPLIT_HIGH_0]], %[[SPLIT_LOW_1]], %[[SPLIT_HIGH_1]]
+    // CHECK-NEXT: %[[OUTER:.*]] = stablehlo.tuple %[[INNER]]
+    // CHECK-NEXT: %[[OUTER_0:.*]] = stablehlo.get_tuple_element %[[OUTER]][0]
+    // CHECK-NEXT: %[[GTE_0:.*]] = stablehlo.get_tuple_element %[[OUTER_0]][0]
+    // CHECK-NEXT: %[[GTE_1:.*]] = stablehlo.get_tuple_element %[[OUTER_0]][1]
+    // CHECK-NEXT: %[[COMBINE_0:.*]] = stablehlo.custom_call @X64Combine(%[[GTE_0]], %[[GTE_1]])
+    // CHECK-SAME:   {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"a"}]>]>} : (tensor<16xui32>, tensor<16xui32>) -> tensor<16xi64>
+    // CHECK-NEXT: %[[GTE_2:.*]] = stablehlo.get_tuple_element %[[OUTER_0]][2]
+    // CHECK-NEXT: %[[GTE_3:.*]] = stablehlo.get_tuple_element %[[OUTER_0]][3]
+    // CHECK-NEXT: %[[COMBINE_1:.*]] = stablehlo.custom_call @X64Combine(%[[GTE_2]], %[[GTE_3]])
+    // CHECK-SAME:   {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"b"}]>]>} : (tensor<16xui32>, tensor<16xui32>) -> tensor<16xi64>
+    // CHECK-NEXT: return %[[COMBINE_0]], %[[COMBINE_1]]
+    %0 = stablehlo.custom_call @X64SplitLow(%arg0) : (tensor<16xi64>) -> tensor<16xui32>
+    %1 = stablehlo.custom_call @X64SplitHigh(%arg0) : (tensor<16xi64>) -> tensor<16xui32>
+    %2 = stablehlo.custom_call @X64SplitLow(%arg1) : (tensor<16xi64>) -> tensor<16xui32>
+    %3 = stablehlo.custom_call @X64SplitHigh(%arg1) : (tensor<16xi64>) -> tensor<16xui32>
+    %4 = stablehlo.custom_call @xla.sdy.FuncResultSharding(%0)
+        {has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@mesh, [{\"a\"}]>]>"}} :
+        (tensor<16xui32>) -> tensor<16xui32>
+    %5 = stablehlo.custom_call @xla.sdy.FuncResultSharding(%1)
+        {has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@mesh, [{\"a\"}]>]>"}} :
+        (tensor<16xui32>) -> tensor<16xui32>
+    %6 = stablehlo.custom_call @xla.sdy.FuncResultSharding(%2)
+        {has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@mesh, [{\"b\"}]>]>"}} :
+        (tensor<16xui32>) -> tensor<16xui32>
+    %7 = stablehlo.custom_call @xla.sdy.FuncResultSharding(%3)
+        {has_side_effect = true, mhlo.frontend_attributes = {xla.sdy.sharding = "#sdy.sharding_per_value<[<@mesh, [{\"b\"}]>]>"}} :
+        (tensor<16xui32>) -> tensor<16xui32>
+    %8 = stablehlo.tuple %4, %5, %6, %7 : tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>
+    %9 = stablehlo.tuple %8 : tuple<tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>>
+    %10 = stablehlo.get_tuple_element %9[0] : (tuple<tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>>) -> tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>
+    %11 = stablehlo.get_tuple_element %10[0] : (tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %12 = stablehlo.get_tuple_element %10[1] : (tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %13 = stablehlo.custom_call @X64Combine(%11, %12) : (tensor<16xui32>, tensor<16xui32>) -> tensor<16xi64>
+    %14 = stablehlo.get_tuple_element %10[2] : (tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %15 = stablehlo.get_tuple_element %10[3] : (tuple<tensor<16xui32>, tensor<16xui32>, tensor<16xui32>, tensor<16xui32>>) -> tensor<16xui32>
+    %16 = stablehlo.custom_call @X64Combine(%14, %15) : (tensor<16xui32>, tensor<16xui32>) -> tensor<16xi64>
+    return %13, %16 : tensor<16xi64>, tensor<16xi64>
+  }
+}
+

From 50c771f8b047b671cada5066165f46cf390b1be0 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Thu, 15 Oct 2026 17:47:29 -0700
Subject: [PATCH 0919/1000] [XLA] Improve float32, float16, and bfloat16
 precision of sinh on CPU and TPU.

- Add a vectorized C++ `sinh` intrinsic (`xla/codegen/intrinsic/cpp/sinh.{h,cc}`) for XLA:CPU using Cody-Waite range reduction and minimax polynomials, achieving 1.0 ULP accuracy across all float32 inputs.
- Add a shared `ElementaryFunctionExpander` HLO pass to expand `kSinh` for `f64`, `c64`, and `c128` (and for `f32`/`f16`/`bf16` on backends without a native `kSinh` emitter).
- Emit `HloOpcode::kSinh` directly from `math.cc` for `f32`, `f16`, and `bf16` on CPU and TPU.

Real ULP bounds on CPU:
- bfloat16 (exhaustive): 0.5 ULP
- float16 (exhaustive):  1.0 ULP
- float32 (exhaustive):  25.0 ULP -> 1.0 ULP

PiperOrigin-RevId: 1000510362
---
 third_party/xla/xla/backends/cpu/nanort/BUILD |   1 +
 .../xla/backends/cpu/nanort/ifrt_client.cc    |   4 +-
 third_party/xla/xla/codegen/BUILD             |   1 +
 .../transforms/lower_xla_intrinsic_lib.cc     |   5 +-
 .../tests/lower_xla_intrinsic_lib.mlir        |  28 +++
 .../xla/xla/codegen/intrinsic/cpp/BUILD       |  52 +++++
 .../intrinsic/cpp/intrinsic_declarations.h    |  25 +-
 .../xla/xla/codegen/intrinsic/cpp/trig.cc     | 143 ++++++++++++
 .../xla/xla/codegen/intrinsic/cpp/trig.h      |  30 +++
 .../xla/codegen/intrinsic/cpp/trig_test.cc    | 161 +++++++++++++
 .../codegen/intrinsic/cpp/vector_ops_test.cc  |  49 ++++
 third_party/xla/xla/codegen/intrinsic_lib.cc  |  10 +
 .../xla/xla/codegen/intrinsic_lib_test.cc     |  26 ++-
 third_party/xla/xla/hlo/builder/lib/math.cc   |  48 +---
 third_party/xla/xla/hlo/builder/lib/math.h    |   3 +-
 third_party/xla/xla/hlo/builder/xla_builder.h |   3 +-
 .../xla/xla/hlo/transforms/expanders/BUILD    |  38 +++
 .../expanders/elementary_function_expander.cc | 157 +++++++++++++
 .../expanders/elementary_function_expander.h  |  51 ++++
 .../elementary_function_expander_test.cc      | 221 ++++++++++++++++++
 .../translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc  |   3 +-
 .../xla/xla/hlo/translate/stablehlo.cc        |   3 +-
 .../chlo_legalize_to_hlo_pass.cc              |  21 +-
 .../xla/xla/mlir_hlo/mhlo/transforms/passes.h |   6 +
 third_party/xla/xla/pjrt/cpu/BUILD            |   1 +
 third_party/xla/xla/pjrt/cpu/cpu_client.cc    |   4 +-
 third_party/xla/xla/python/_ops.pyi           |   2 +-
 third_party/xla/xla/service/cpu/BUILD         |   2 +
 .../xla/xla/service/cpu/cpu_compiler.cc       |   2 +
 .../xla/service/cpu/elemental_ir_emitter.cc   |  16 ++
 .../cpu/tests/cpu_vectorization_test.cc       |  38 +++
 third_party/xla/xla/service/gpu/BUILD         |   1 +
 .../xla/xla/service/gpu/gpu_compiler.cc       |   5 +
 .../exhaustive_unary_test_functions.cc        |  15 +-
 34 files changed, 1103 insertions(+), 72 deletions(-)
 create mode 100644 third_party/xla/xla/codegen/intrinsic/cpp/trig.cc
 create mode 100644 third_party/xla/xla/codegen/intrinsic/cpp/trig.h
 create mode 100644 third_party/xla/xla/codegen/intrinsic/cpp/trig_test.cc
 create mode 100644 third_party/xla/xla/codegen/intrinsic/cpp/vector_ops_test.cc
 create mode 100644 third_party/xla/xla/hlo/transforms/expanders/elementary_function_expander.cc
 create mode 100644 third_party/xla/xla/hlo/transforms/expanders/elementary_function_expander.h
 create mode 100644 third_party/xla/xla/hlo/transforms/expanders/elementary_function_expander_test.cc

diff --git a/third_party/xla/xla/backends/cpu/nanort/BUILD b/third_party/xla/xla/backends/cpu/nanort/BUILD
index 1b8aa55cc85..e673514a9e9 100644
--- a/third_party/xla/xla/backends/cpu/nanort/BUILD
+++ b/third_party/xla/xla/backends/cpu/nanort/BUILD
@@ -179,6 +179,7 @@ cc_library(
         "//xla/hlo/builder:xla_computation",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/ir:hlo_sharding",
+        "//xla/mlir_hlo:mhlo_passes",
         "//xla/pjrt:mlir_to_hlo",
         "//xla/pjrt:pjrt_compiler",
         "//xla/pjrt:pjrt_layout",
diff --git a/third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc b/third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc
index f9ec52343a0..671199e9dc3 100644
--- a/third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc
+++ b/third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc
@@ -52,6 +52,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_sharding.h"
 #include "xla/layout.h"
 #include "xla/layout_util.h"
+#include "xla/mlir_hlo/mhlo/transforms/passes.h"
 #include "xla/pjrt/mlir_to_hlo.h"
 #include "xla/pjrt/pjrt_compiler.h"
 #include "xla/pjrt/pjrt_layout.h"
@@ -913,7 +914,8 @@ class NanoExecutable final
     XlaComputation computation;
     ABSL_RETURN_IF_ERROR(MlirToXlaComputation(
         xla_program->mlir_module(), computation, /*use_tuple_args=*/false,
-        /*return_tuple=*/true, /*exec_build_options=*/nullptr));
+        /*return_tuple=*/true, /*exec_build_options=*/nullptr,
+        mlir::mhlo::getCpuChloToHighLevelMhloOptions()));
     ABSL_ASSIGN_OR_RETURN(auto nano_executable,
                      client->nano_client()->Compile(computation));
 
diff --git a/third_party/xla/xla/codegen/BUILD b/third_party/xla/xla/codegen/BUILD
index 25b28634aa7..2bd73805011 100644
--- a/third_party/xla/xla/codegen/BUILD
+++ b/third_party/xla/xla/codegen/BUILD
@@ -203,6 +203,7 @@ cc_library(
         "//xla/codegen/intrinsic:vec_name_mangler",
         "//xla/codegen/intrinsic/cpp:_cpp_gen_intrinsics",
         "//xla/codegen/intrinsic/cpp:_intrinsic_declarations",
+        "//xla/codegen/intrinsic/cpp:trig_ll",
         "//xla/service/llvm_ir:llvm_util",
         "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/container:flat_hash_map",
diff --git a/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc b/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc
index 8d8f5026931..64316fcb6fa 100644
--- a/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc
+++ b/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc
@@ -344,8 +344,9 @@ class LowerXlaIntrinsicLibPass
                  LowerIntrinsicPattern<ci::Sin, mm::SinOp>,
                  LowerIntrinsicPattern<ci::Cos, mm::CosOp>,
                  LowerIntrinsicPattern<ci::FpTrunc, ma::TruncFOp>,
-                 LowerIntrinsicPattern<ci::Erf, mm::ErfOp>>(context,
-                                                            cpu_features);
+                 LowerIntrinsicPattern<ci::Erf, mm::ErfOp>,
+                 LowerIntrinsicPattern<ci::Sinh, mm::SinhOp>>(context,
+                                                              cpu_features);
     if (mlir::failed(
             mlir::applyPatternsGreedily(module_op, std::move(patterns)))) {
       signalPassFailure();
diff --git a/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir b/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir
index 5d258683b66..0deade04284 100644
--- a/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir
+++ b/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir
@@ -452,3 +452,31 @@ func.func @sin_cos_vector(%arg0: vector<8xf32>, %arg1: vector<4xf64>) -> (vector
 // CHECK: %[[SIN:.*]] = call @xla.sin.v8f32(%arg0) : (vector<8xf32>) -> vector<8xf32>
 // CHECK: %[[COS:.*]] = call @xla.cos.v4f64(%arg1) : (vector<4xf64>) -> vector<4xf64>
 // CHECK: return %[[SIN]], %[[COS]]
+
+// -----
+
+module {
+  func.func @sinh_f32(%arg0: f32) -> f32 {
+    %ret = math.sinh %arg0 : f32
+    return %ret : f32
+  }
+}
+
+// CHECK-LABEL: @sinh_f32
+// CHECK-NOT: math.sinh
+// CHECK: %[[RESULT:.*]] = call @xla.sinh.f32(%arg0) : (f32) -> f32
+// CHECK: return %[[RESULT]] : f32
+
+// -----
+
+module {
+  func.func @sinh_vector(%arg0: vector<8xf32>) -> vector<8xf32> {
+    %ret = math.sinh %arg0 : vector<8xf32>
+    return %ret : vector<8xf32>
+  }
+}
+
+// CHECK-LABEL: @sinh_vector
+// CHECK-NOT: math.sinh
+// CHECK: %[[SINH_CALL:.*]] = call @xla.sinh.v8f32(%arg0) : (vector<8xf32>) -> vector<8xf32>
+// CHECK: return %[[SINH_CALL]]
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/BUILD b/third_party/xla/xla/codegen/intrinsic/cpp/BUILD
index 2c280ebd283..363caf3b3ca 100644
--- a/third_party/xla/xla/codegen/intrinsic/cpp/BUILD
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/BUILD
@@ -107,6 +107,18 @@ cc_library(
     ],
 )
 
+xla_cc_test(
+    name = "vector_ops_test",
+    srcs = ["vector_ops_test.cc"],
+    copts = [
+        "-Wno-psabi",
+    ],
+    deps = [
+        ":vector_ops",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_ir_header(
     name = "tanh_ll",
     src = "tanh.cc",
@@ -353,6 +365,46 @@ xla_cc_test(
     ],
 )
 
+cc_library(
+    name = "trig_lib",
+    srcs = ["trig.cc"],
+    hdrs = ["trig.h"],
+    copts = [
+        "-Wno-psabi",
+    ],
+    deps = [
+        ":vector_ops",
+    ],
+)
+
+cc_ir_header(
+    name = "trig_ll",
+    src = "trig.cc",
+    deps = [
+        ":cpp_context_provider",
+        ":trig_lib",
+        ":vector_ops",
+    ],
+)
+
+xla_cc_test(
+    name = "trig_test",
+    srcs = ["trig_test.cc"],
+    copts = [
+        "-Wno-psabi",
+    ],
+    deps = [
+        ":_cpp_gen_intrinsics",
+        ":trig_lib",
+        ":trig_ll",
+        ":vector_ops",
+        "//xla/codegen/intrinsic:test_matchers",
+        "@com_google_googletest//:gtest_main",
+        "@llvm-project//llvm:Core",
+        "@llvm-project//llvm:Support",
+    ],
+)
+
 cc_library(
     name = "_cpp_gen_intrinsics",
     srcs = ["cpp_gen_intrinsics.cc"],
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/intrinsic_declarations.h b/third_party/xla/xla/codegen/intrinsic/cpp/intrinsic_declarations.h
index 4ca66284dca..d97755e86b0 100644
--- a/third_party/xla/xla/codegen/intrinsic/cpp/intrinsic_declarations.h
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/intrinsic_declarations.h
@@ -110,7 +110,6 @@ class Expm1 : public Intrinsic<Expm1> {
                              UnaryFunctionType(module, type));
   }
 };
-
 class Sin : public Intrinsic<Sin> {
  public:
   static constexpr absl::string_view kName = "sin";
@@ -164,6 +163,30 @@ class Cos : public Intrinsic<Cos> {
                              UnaryFunctionType(module, type));
   }
 };
+
+class Sinh : public Intrinsic<Sinh> {
+ public:
+  static constexpr absl::string_view kName = "sinh";
+
+  static std::vector<std::vector<Type>> SupportedVectorTypes(
+      absl::string_view features) {
+    if (!AreCppGenIntrinsicsAvailable()) {
+      return {};
+    }
+    return {
+        {Type::S(xla::F32)},
+        {Type::V(xla::F32, 4)},
+        {Type::V(xla::F32, 8)},
+        {Type::V(xla::F32, 16)},
+    };
+  }
+
+  static absl::StatusOr<llvm::Function*> CreateDefinition(
+      llvm::Module* module, const IntrinsicOptions& options, Type type) {
+    return GetCppGenFunction(module, Name(type),
+                             UnaryFunctionType(module, type));
+  }
+};
 }  // namespace xla::codegen::intrinsics
 
 #endif  // XLA_CODEGEN_INTRINSIC_CPP_INTRINSIC_DECLARATIONS_H_
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/trig.cc b/third_party/xla/xla/codegen/intrinsic/cpp/trig.cc
new file mode 100644
index 00000000000..f9aabc19aaf
--- /dev/null
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/trig.cc
@@ -0,0 +1,143 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/codegen/intrinsic/cpp/trig.h"
+
+#include <type_traits>
+
+#include "xla/codegen/intrinsic/cpp/vector_ops.h"
+
+namespace xla::codegen {
+
+namespace {
+
+template <typename IntType, typename FloatType>
+inline IntType FloatToInt(FloatType x) {
+  if constexpr (std::is_fundamental_v<FloatType>) {
+    return static_cast<IntType>(x);
+  } else {
+    return __builtin_convertvector(x, IntType);
+  }
+}
+
+// Evaluates float32 sinh(x) = (e^x - e^{-x}) / 2 to within 1 ULP using
+// sinh(-x) = -sinh(x) and splitting |x| at 1.0:
+//
+// - For |x| < 1, avoids cancellation in e^{|x|} - e^{-|x|} via Taylor series:
+//     sinh(|x|) = [(1 + |x| + |x|^2/2! + |x|^3/3! + ...) -
+//                  (1 - |x| + |x|^2/2! - |x|^3/3! + ...)] / 2
+//               = |x| + |x|^3/3! + |x|^5/5! + ...
+//               = |x| + |x|^3 * (1/3! + |x|^2/5! + ...)
+//              ~= |x| + |x|^3 * P_small(|x|^2)
+//
+// - For |x| >= 1, decomposes |x| = n * ln(2) + r_hi + r_lo (r = r_hi + r_lo)
+//   and evaluates:
+//     sinh(|x|) = 2 * (0.25 * e^{|x|} - 0.25 * e^{-|x|})
+//     e^r       = 1 + r + r^2/2! + r^3/3! + ...
+//              ~= 1 + r + r^2 * P_exp(r)
+//               = (1 + r_hi) + (r_lo + r^2 * P_exp(r))
+//     0.25 * e^{|x|} = 2^{n - 2} * e^r
+//                   ~= head + tail
+//       where head = 2^{n - 2} * (1 + r_hi)
+//             tail = 2^{n - 2} * (r_lo + r^2 * P_exp(r))
+//     -0.25 * e^{-|x|} = -0.25^2 / (0.25 * e^{|x|})
+//                     ~= -0.0625 / (head + tail)
+//     sinh(|x|) ~= 2 * ((head + tail) - 0.0625 / (head + tail))
+//                = 2 * (head + (tail - 0.0625 / (head + tail)))
+template <typename T>
+inline T SinhF32(T x) {
+#pragma clang fp contract(fast)
+  using IntType = typename internal::CorrespondingIntVector<T>::type;
+
+  // 1. Absolute value and branch condition:
+  // Split into |x| < 1 (small) and |x| >= 1 (large).
+  T abs_x = BitwiseAbs(x);
+  auto is_small = abs_x < Splat<T>(1.0f);
+
+  // 2. Cody-Waite range reduction for |x| >= 1:
+  // Clamp |x| at 89.5 so that n = round_nearest_even(|x| * log2(e)) <= 129.
+  // This keeps (n + 125) <= 254 within normal float32 exponent range when
+  // constructing scale = 2^(n - 2), while ensuring 2 * head (where
+  // head = scale * (1 + r_hi) ~= 0.25 * exp(|x|) is the leading part of
+  // 0.25 * exp(|x|) computed in step 3) overflows to +inf for |x| > ~89.415985.
+  T clamped_x = abs_x < Splat<T>(89.5f) ? abs_x : Splat<T>(89.5f);
+  T x_log2e = clamped_x * Splat<T>(0x1.715476p+0f);
+  T n_float = __builtin_elementwise_rint(x_log2e);
+  IntType n_int = FloatToInt<IntType>(n_float);
+
+  // Split ln(2) into a 16-bit high part (ln2_hi = 0x1.62e4p-1f) and a low
+  // part (ln2_lo = 0x1.7f7d1cp-20f). Because n <= 129 has at most 8 integer
+  // bits and ln2_hi has 16 bits, n * ln2_hi and r_hi = |x| - n * ln2_hi are
+  // bit-exact without FMA.
+  T r_hi = abs_x - n_float * Splat<T>(0x1.62e4p-1f);
+  T r_lo = -n_float * Splat<T>(0x1.7f7d1cp-20f);
+  T r = r_hi + r_lo;
+  T r2 = r * r;
+
+  // 3. Large-|x| minimax polynomial for (exp(r) - 1 - r) / r^2 on
+  // [-ln(2)/2, ln(2)/2].
+  // Generated in Sollya:
+  //   display = hexadecimal!;
+  //   p_exp = fpminimax((expm1(x) - x) / x, [|1, 2, 3, 4, 5|], [|SG...|],
+  //                     [-log(2)/2, log(2)/2], relative);
+  //   supnorm(1 + x + x * p_exp, exp(x), [-log(2)/2, log(2)/2], relative,
+  //           2^-40);
+  //   // -> 0.95924801632e-8 (~0.0805 ULP in FP32)
+  T p_exp = HornerPoly(r, 0x1.6ca992p-10f, 0x1.120abep-7f, 0x1.55556cp-5f,
+                       0x1.5554dep-3f, 0x1p-1f);
+
+  // Construct scale = 2^(n - 2) via integer exponent bitcast ((n + 125) << 23)
+  // and evaluate:
+  //   head      = scale + scale * r_hi
+  //   tail      = scale * (r_lo + r^2 * P_exp(r))
+  //   half_h    = head + tail                  (~= 0.25 * e^{|x|})
+  //   div_term  = -0.0625 / half_h             (~= -0.25 * e^{-|x|})
+  //   sinh(|x|) ~= 2 * (head + (tail + div_term))
+  T scale = __builtin_bit_cast(T, (n_int + 125u) << 23);
+  T head = scale + scale * r_hi;
+  T tail = scale * (r_lo + r2 * p_exp);
+  T half_h = head + tail;
+  T div_term = Splat<T>(-0.0625f) / half_h;
+  T large_sinh_result = Splat<T>(2.0f) * (head + (tail + div_term));
+
+  // 4. Small-|x| odd minimax polynomial on [-1, 1] for (sinh(|x|) - |x|) /
+  // |x|^3:
+  //   sinh(|x|) = |x| + |x|^3 * (c3 + |x|^2 * (c5 + |x|^2 * (c7 + |x|^2 * c9)))
+  // Computing |x| + |x|^3 * P_small(|x|^2) (rather than
+  // |x| * (1 + |x|^2 * P_small(|x|^2))) avoids an extra rounding error on the
+  // leading |x| term when |x|^2 * P_small(|x|^2) is small. Generated in Sollya:
+  //   display = hexadecimal!;
+  //   p_small = fpminimax(sinh(x), [|3, 5, 7, 9|], [|SG...|], [1b-20, 1],
+  //                       floating, relative, x);
+  //   supnorm(p_small, sinh(x), [1b-20, 1], relative, 2^-40);
+  //   // -> 0.3558484444e-9 (~0.0030 ULP in FP32)
+  T x2 = abs_x * abs_x;
+  T p_small = HornerPoly(x2, 0x1.76771ep-19f, 0x1.a01c22p-13f, 0x1.1110eap-7f,
+                         0x1.555556p-3f);
+  T small_sinh_result = abs_x + (abs_x * x2) * p_small;
+
+  // 5. Select branch and restore sign: sinh(-x) = -sinh(x).
+  T res_mag = is_small ? small_sinh_result : large_sinh_result;
+  return BitwiseCopysign(res_mag, x);
+}
+
+}  // namespace
+
+float sinh_f32(float x) { return SinhF32(x); }
+Vec4f sinh_v4f32(Vec4f x) { return SinhF32(x); }
+Vec8f sinh_v8f32(Vec8f x) { return SinhF32(x); }
+Vec16f sinh_v16f32(Vec16f x) { return SinhF32(x); }
+
+}  // namespace xla::codegen
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/trig.h b/third_party/xla/xla/codegen/intrinsic/cpp/trig.h
new file mode 100644
index 00000000000..5c785bba5a8
--- /dev/null
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/trig.h
@@ -0,0 +1,30 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_CODEGEN_INTRINSIC_CPP_TRIG_H_
+#define XLA_CODEGEN_INTRINSIC_CPP_TRIG_H_
+
+#include "xla/codegen/intrinsic/cpp/vector_ops.h"
+
+namespace xla::codegen {
+
+float sinh_f32(float x) asm("xla.sinh.f32");
+Vec4f sinh_v4f32(Vec4f x) asm("xla.sinh.v4f32");
+Vec8f sinh_v8f32(Vec8f x) asm("xla.sinh.v8f32");
+Vec16f sinh_v16f32(Vec16f x) asm("xla.sinh.v16f32");
+
+}  // namespace xla::codegen
+
+#endif  // XLA_CODEGEN_INTRINSIC_CPP_TRIG_H_
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/trig_test.cc b/third_party/xla/xla/codegen/intrinsic/cpp/trig_test.cc
new file mode 100644
index 00000000000..2e85cadddb7
--- /dev/null
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/trig_test.cc
@@ -0,0 +1,161 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/codegen/intrinsic/cpp/trig.h"
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <cstring>
+#include <limits>
+#include <memory>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/raw_ostream.h"
+#include "xla/codegen/intrinsic/cpp/cpp_gen_intrinsics.h"
+#include "xla/codegen/intrinsic/cpp/trig_ll.h"
+#include "xla/codegen/intrinsic/cpp/vector_ops.h"
+#include "xla/codegen/intrinsic/test_matchers.h"
+
+namespace xla::codegen {
+namespace {
+using ::testing::ContainsRegex;
+using ::xla::codegen::intrinsic::NearUlps;
+
+TEST(TrigTest, SinhF32IsCorrect) {
+  double max_ulp_error = 0.0;
+  float worst_x = 0.0f;
+  int64_t max_int_ulp_diff = 0;
+
+  auto check_point = [&](uint32_t bits) {
+    float x;
+    std::memcpy(&x, &bits, sizeof(x));
+    double expected = std::sinh(static_cast<double>(x));
+    float actual = sinh_f32(x);
+
+    float expected_f32 = static_cast<float>(expected);
+    uint32_t actual_bits, expected_bits;
+    std::memcpy(&actual_bits, &actual, sizeof(actual_bits));
+    std::memcpy(&expected_bits, &expected_f32, sizeof(expected_bits));
+    int64_t int_diff = std::abs(static_cast<int64_t>(actual_bits) -
+                                static_cast<int64_t>(expected_bits));
+    if (int_diff > max_int_ulp_diff) {
+      max_int_ulp_diff = int_diff;
+    }
+
+    int exp;
+    std::frexp(expected, &exp);
+    double ulp = std::ldexp(1.0, std::max(exp - 24, -149));
+    double ulp_error = std::abs(static_cast<double>(actual) - expected) / ulp;
+    if (ulp_error > max_ulp_error) {
+      max_ulp_error = ulp_error;
+      worst_x = x;
+    }
+  };
+
+  // Sweep the full finite positive domain [0.0f, 89.4159851f] with stride 256,
+  // plus exhaustive coverage around the worst-case transition point x
+  // = 1.04064.
+  constexpr uint32_t kMaxFiniteSinhBits = 0x42b2d4fc;
+  for (uint32_t bits = 0; bits <= kMaxFiniteSinhBits; bits += 256) {
+    check_point(bits);
+  }
+  constexpr uint32_t kWorstCaseBits = 0x3f8533bf;  // 1.0406416654586792f
+  for (uint32_t bits = kWorstCaseBits - 4096; bits <= kWorstCaseBits + 4096;
+       ++bits) {
+    check_point(bits);
+  }
+
+  EXPECT_LE(max_int_ulp_diff, 1);
+  EXPECT_LE(max_ulp_error, 0.81) << "worst_x = " << worst_x;
+}
+
+TEST(TrigTest, SinhF32EdgeCases) {
+  EXPECT_EQ(sinh_f32(0.0f), 0.0f);
+  EXPECT_FALSE(std::signbit(sinh_f32(0.0f)));
+  EXPECT_EQ(sinh_f32(-0.0f), -0.0f);
+  EXPECT_TRUE(std::signbit(sinh_f32(-0.0f)));
+
+  // Overflow threshold: sinh(89.4159851f) < FLT_MAX, sinh(89.4159927f) >
+  // FLT_MAX
+  float just_below = 89.4159851f;
+  EXPECT_FALSE(std::isinf(sinh_f32(just_below)));
+  EXPECT_THAT(
+      sinh_f32(just_below),
+      NearUlps(static_cast<float>(std::sinh(static_cast<double>(just_below))),
+               1));
+  EXPECT_FALSE(std::isinf(sinh_f32(-just_below)));
+  EXPECT_THAT(
+      sinh_f32(-just_below),
+      NearUlps(static_cast<float>(std::sinh(static_cast<double>(-just_below))),
+               1));
+
+  float just_above = 89.4159927f;
+  EXPECT_EQ(sinh_f32(just_above), std::numeric_limits<float>::infinity());
+  EXPECT_EQ(sinh_f32(-just_above), -std::numeric_limits<float>::infinity());
+
+  EXPECT_EQ(sinh_f32(std::numeric_limits<float>::infinity()),
+            std::numeric_limits<float>::infinity());
+  EXPECT_EQ(sinh_f32(-std::numeric_limits<float>::infinity()),
+            -std::numeric_limits<float>::infinity());
+  EXPECT_TRUE(std::isnan(sinh_f32(std::numeric_limits<float>::quiet_NaN())));
+}
+
+TEST(TrigTest, SinhF32VectorWidthsAreCorrect) {
+  Vec4f x4 = {-10.0f, -0.5f, 0.5f, 10.0f};
+  Vec4f y4 = sinh_v4f32(x4);
+  for (int i = 0; i < 4; ++i) {
+    EXPECT_THAT(
+        y4[i],
+        NearUlps(static_cast<float>(std::sinh(static_cast<double>(x4[i]))), 1));
+  }
+
+  Vec8f x8 = {-50.0f, -10.0f, -1.0f, -0.1f, 0.1f, 1.0f, 10.0f, 50.0f};
+  Vec8f y8 = sinh_v8f32(x8);
+  for (int i = 0; i < 8; ++i) {
+    EXPECT_THAT(
+        y8[i],
+        NearUlps(static_cast<float>(std::sinh(static_cast<double>(x8[i]))), 1));
+  }
+
+  Vec16f x16 = {-80.0f, -50.0f, -20.0f, -5.0f, -1.5f, -0.5f, -0.01f, -0.0f,
+                0.0f,   0.01f,  0.5f,   1.5f,  5.0f,  20.0f, 50.0f,  80.0f};
+  Vec16f y16 = sinh_v16f32(x16);
+  for (int i = 0; i < 16; ++i) {
+    EXPECT_THAT(
+        y16[i],
+        NearUlps(static_cast<float>(std::sinh(static_cast<double>(x16[i]))),
+                 1));
+  }
+}
+
+TEST(TrigTest, SinhIsVectorized) {
+  llvm::LLVMContext context;
+  std::unique_ptr<llvm::Module> module =
+      ParseEmbeddedBitcode(context, llvm_ir::kTrigLlIr);
+
+  std::string ir;
+  llvm::raw_string_ostream stream(ir);
+  module->print(stream, nullptr);
+  EXPECT_THAT(ir, ContainsRegex("xla.sinh.v16f32"));
+  EXPECT_THAT(ir, ContainsRegex("fmul.*<16 x float>"));
+}
+
+}  // namespace
+}  // namespace xla::codegen
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/vector_ops_test.cc b/third_party/xla/xla/codegen/intrinsic/cpp/vector_ops_test.cc
new file mode 100644
index 00000000000..168e69d0bcd
--- /dev/null
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/vector_ops_test.cc
@@ -0,0 +1,49 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/codegen/intrinsic/cpp/vector_ops.h"
+
+#include <cstdint>
+
+#include <gtest/gtest.h>
+
+namespace xla::codegen {
+namespace {
+
+TEST(VectorOpsTest, SplatPreservesNegativeZero) {
+  Vec4f v = Splat<Vec4f>(-0.0f);
+  Vec4u bits = __builtin_bit_cast(Vec4u, v);
+  for (int i = 0; i < 4; ++i) {
+    EXPECT_EQ(bits[i], 0x80000000u);
+  }
+  float s = Splat<float>(-0.0f);
+  EXPECT_EQ(__builtin_bit_cast(uint32_t, s), 0x80000000u);
+}
+
+TEST(VectorOpsTest, HornerPoly) {
+  // P(x) = 2x^3 - 3x^2 + 5x - 7
+  EXPECT_FLOAT_EQ(HornerPoly(0.0f, 2.0f, -3.0f, 5.0f, -7.0f), -7.0f);
+  EXPECT_FLOAT_EQ(HornerPoly(2.0f, 2.0f, -3.0f, 5.0f, -7.0f), 7.0f);
+
+  Vec4f x = {-1.0f, 0.0f, 1.0f, 2.0f};
+  Vec4f y = HornerPoly(x, 2.0f, -3.0f, 5.0f, -7.0f);
+  EXPECT_FLOAT_EQ(y[0], -17.0f);
+  EXPECT_FLOAT_EQ(y[1], -7.0f);
+  EXPECT_FLOAT_EQ(y[2], -3.0f);
+  EXPECT_FLOAT_EQ(y[3], 7.0f);
+}
+
+}  // namespace
+}  // namespace xla::codegen
diff --git a/third_party/xla/xla/codegen/intrinsic_lib.cc b/third_party/xla/xla/codegen/intrinsic_lib.cc
index 3d278f223ae..6a58824f63e 100644
--- a/third_party/xla/xla/codegen/intrinsic_lib.cc
+++ b/third_party/xla/xla/codegen/intrinsic_lib.cc
@@ -43,6 +43,7 @@ limitations under the License.
 #include "llvm/Support/raw_ostream.h"
 #include "xla/codegen/intrinsic/cpp/cpp_gen_intrinsics.h"
 #include "xla/codegen/intrinsic/cpp/intrinsic_declarations.h"
+#include "xla/codegen/intrinsic/cpp/trig_ll.h"
 #include "xla/codegen/intrinsic/erf.h"
 #include "xla/codegen/intrinsic/exp.h"
 #include "xla/codegen/intrinsic/fptrunc.h"
@@ -125,6 +126,13 @@ class IntrinsicAdapter : public IntrinsicFunction {
 
 IntrinsicFunctionLib::IntrinsicFunctionLib(const IntrinsicOptions& options)
     : options_(options), ir_libraries_(GetCppGenLibraries(options)) {
+  if (options.device_type == intrinsics::DeviceType::kIntelCpu ||
+      options.device_type == intrinsics::DeviceType::kAmdCpu ||
+      options.device_type == intrinsics::DeviceType::kArmCpu) {
+    if (!::llvm_ir::kTrigLlIr.empty()) {
+      ir_libraries_.emplace_back(::llvm_ir::kTrigLlIr, "trig");
+    }
+  }
   intrinsic_functions_.push_back(
       std::make_unique<IntrinsicAdapter<intrinsics::Ldexp>>());
   intrinsic_functions_.push_back(
@@ -147,6 +155,8 @@ IntrinsicFunctionLib::IntrinsicFunctionLib(const IntrinsicOptions& options)
       std::make_unique<IntrinsicAdapter<intrinsics::Sin>>());
   intrinsic_functions_.push_back(
       std::make_unique<IntrinsicAdapter<intrinsics::Cos>>());
+  intrinsic_functions_.push_back(
+      std::make_unique<IntrinsicAdapter<intrinsics::Sinh>>());
 }
 
 namespace {
diff --git a/third_party/xla/xla/codegen/intrinsic_lib_test.cc b/third_party/xla/xla/codegen/intrinsic_lib_test.cc
index 41ae9b5ce97..df6ea4a5546 100644
--- a/third_party/xla/xla/codegen/intrinsic_lib_test.cc
+++ b/third_party/xla/xla/codegen/intrinsic_lib_test.cc
@@ -170,7 +170,6 @@ TEST(IntrinsicLibTest, Expm1Vectorizations) {
                            "xla.expm1.f64:xla.expm1.v4f64:4:_ZGV_LLVM_N4v",
                            "xla.expm1.f64:xla.expm1.v8f64:8:_ZGV_LLVM_N8v"));
 }
-
 TEST(IntrinsicLibTest, SinAndCosVectorizations) {
   auto descs_with_prefix = [](absl::string_view features,
                               absl::string_view prefix) {
@@ -187,20 +186,20 @@ TEST(IntrinsicLibTest, SinAndCosVectorizations) {
   };
 
   EXPECT_THAT(
-      descs_with_prefix("", "xla.sin"),
+      descs_with_prefix("", "xla.sin."),
       UnorderedElementsAre("xla.sin.f32:xla.sin.v4f32:4:_ZGV_LLVM_N4v",
                            "xla.sin.f32:xla.sin.v8f32:8:_ZGV_LLVM_N8v",
                            "xla.sin.f64:xla.sin.v2f64:2:_ZGV_LLVM_N2v",
                            "xla.sin.f64:xla.sin.v4f64:4:_ZGV_LLVM_N4v"));
   EXPECT_THAT(
-      descs_with_prefix("", "xla.cos"),
+      descs_with_prefix("", "xla.cos."),
       UnorderedElementsAre("xla.cos.f32:xla.cos.v4f32:4:_ZGV_LLVM_N4v",
                            "xla.cos.f32:xla.cos.v8f32:8:_ZGV_LLVM_N8v",
                            "xla.cos.f64:xla.cos.v2f64:2:_ZGV_LLVM_N2v",
                            "xla.cos.f64:xla.cos.v4f64:4:_ZGV_LLVM_N4v"));
 
   EXPECT_THAT(
-      descs_with_prefix("+avx512f", "xla.sin"),
+      descs_with_prefix("+avx512f", "xla.sin."),
       UnorderedElementsAre("xla.sin.f32:xla.sin.v4f32:4:_ZGV_LLVM_N4v",
                            "xla.sin.f32:xla.sin.v8f32:8:_ZGV_LLVM_N8v",
                            "xla.sin.f32:xla.sin.v16f32:16:_ZGV_LLVM_N16v",
@@ -208,7 +207,7 @@ TEST(IntrinsicLibTest, SinAndCosVectorizations) {
                            "xla.sin.f64:xla.sin.v4f64:4:_ZGV_LLVM_N4v",
                            "xla.sin.f64:xla.sin.v8f64:8:_ZGV_LLVM_N8v"));
   EXPECT_THAT(
-      descs_with_prefix("+avx512f", "xla.cos"),
+      descs_with_prefix("+avx512f", "xla.cos."),
       UnorderedElementsAre("xla.cos.f32:xla.cos.v4f32:4:_ZGV_LLVM_N4v",
                            "xla.cos.f32:xla.cos.v8f32:8:_ZGV_LLVM_N8v",
                            "xla.cos.f32:xla.cos.v16f32:16:_ZGV_LLVM_N16v",
@@ -217,6 +216,23 @@ TEST(IntrinsicLibTest, SinAndCosVectorizations) {
                            "xla.cos.f64:xla.cos.v8f64:8:_ZGV_LLVM_N8v"));
 }
 
+TEST(IntrinsicLibTest, SinhVectorizations) {
+  IntrinsicOptions options;
+  auto lib = IntrinsicFunctionLib(options);
+  std::vector<llvm::VecDesc> vec_descs = lib.Vectorizations();
+  std::vector<std::string> vec_descs_str;
+  for (const auto& vec_desc : vec_descs) {
+    if (vec_desc.getScalarFnName().starts_with("xla.sinh")) {
+      vec_descs_str.push_back(ToString(vec_desc));
+    }
+  }
+
+  EXPECT_THAT(
+      vec_descs_str,
+      UnorderedElementsAre("xla.sinh.f32:xla.sinh.v4f32:4:_ZGV_LLVM_N4v",
+                           "xla.sinh.f32:xla.sinh.v8f32:8:_ZGV_LLVM_N8v",
+                           "xla.sinh.f32:xla.sinh.v16f32:16:_ZGV_LLVM_N16v"));
+}
 TEST(IntrinsicLibTest, CppGenIntrinsicLibraryPreservesNoInline) {
   llvm::LLVMContext context;
   llvm::Module dst_module("dst_module", context);
diff --git a/third_party/xla/xla/hlo/builder/lib/math.cc b/third_party/xla/xla/hlo/builder/lib/math.cc
index 85e87c861dd..51a629d1c72 100644
--- a/third_party/xla/xla/hlo/builder/lib/math.cc
+++ b/third_party/xla/xla/hlo/builder/lib/math.cc
@@ -1412,51 +1412,9 @@ XlaOp Cosh(XlaOp x, const std::optional<ResultAccuracy>& result_accuracy,
   });
 }
 
-// Sinh(x) = (e^x - e^-x) / 2
-//         = e^(x + log(1/2)) - e^(-x + log(1/2)).
-//
-// The second formulation avoids overflowing when e^x = inf but (e^x)/2 is not
-// inf.
-//
-// This incorrectly overflows to +/-inf for two f32 input values, namely
-// +/-89.4159851, due to rounding error when computing x +/- log(1/2).  The
-// correct answer of 3.40281961e+38 (0x7f7fffec) is very close to max-float, so
-// we deem this acceptable.
-XlaOp Sinh(XlaOp x, const std::optional<ResultAccuracy>& result_accuracy,
-           bool expand) {
-  if (!expand) {
-    return x.builder()->UnaryOp(HloOpcode::kSinh, x, result_accuracy);
-  }
-  XlaBuilder* b = x.builder();
-  auto do_it = [&](XlaOp x) -> absl::StatusOr<XlaOp> {
-    ABSL_ASSIGN_OR_RETURN(auto shape, b->GetShape(x));
-    auto one_half = ScalarLike(x, 0.5);
-    auto log_one_half = Log(ScalarLike(x, 0.5));
-    auto large_sinh_result = Exp(x + log_one_half) - Exp(-x + log_one_half);
-
-    if (primitive_util::IsComplexType(shape.element_type())) {
-      return large_sinh_result;
-    }
-
-    // Here we use e^x = e^(x / 2) * e^(x / 2). This avoids overflow for large
-    // values of x.
-
-    // For smaller x, we get unwanted cancellations of e^x - e^-x, resulting in
-    // 0.
-    // Rewrite this to avoid that. We use expm1(x) because that preserves the
-    // first order term of the taylor series of e^x.
-    // (e^(x) - e^(-x)) / 2. =
-    // (e^(x) - 1 + 1 - e^(-x)) / 2.
-    // (expm1(x) + (e^(x) - 1) / e^x) / 2.
-    // (expm1(x) + expm1(x) / (expm1(x) + 1)) / 2.
-    auto expm1 = Expm1(x);
-    auto one = ScalarLike(x, 1.);
-    auto small_sinh_result = one_half * (expm1 + expm1 / (expm1 + one));
-    return Select(Lt(Abs(x), one), small_sinh_result, large_sinh_result);
-  };
-  return DoWithUpcastToF32(x, {BF16, F16}, [&](XlaOp x) {
-    return b->ReportErrorOrReturn(do_it(x));
-  });
+// Computes the hyperbolic sine of 'x'.
+XlaOp Sinh(XlaOp x, const std::optional<ResultAccuracy>& result_accuracy) {
+  return x.builder()->UnaryOp(HloOpcode::kSinh, x, result_accuracy);
 }
 
 XlaOp MaybeConjugate(XlaOp x, bool conjugate) {
diff --git a/third_party/xla/xla/hlo/builder/lib/math.h b/third_party/xla/xla/hlo/builder/lib/math.h
index 4614c0442aa..238cc3395c5 100644
--- a/third_party/xla/xla/hlo/builder/lib/math.h
+++ b/third_party/xla/xla/hlo/builder/lib/math.h
@@ -121,8 +121,7 @@ XlaOp Cosh(XlaOp x,
 
 // Computes the hyperbolic sine of 'x'.
 XlaOp Sinh(XlaOp x,
-           const std::optional<ResultAccuracy>& result_accuracy = std::nullopt,
-           bool expand = true);
+           const std::optional<ResultAccuracy>& result_accuracy = std::nullopt);
 
 // Applies a complex conjugation operation if 'a' is complex and 'conjugate'
 // is true, otherwise returns its argument.
diff --git a/third_party/xla/xla/hlo/builder/xla_builder.h b/third_party/xla/xla/hlo/builder/xla_builder.h
index de5409b2e78..e689e951ceb 100644
--- a/third_party/xla/xla/hlo/builder/xla_builder.h
+++ b/third_party/xla/xla/hlo/builder/xla_builder.h
@@ -1984,8 +1984,7 @@ class XlaBuilder {
   friend XlaOp Sin(XlaOp operand,
                    const std::optional<ResultAccuracy>& result_accuracy);
   friend XlaOp Sinh(XlaOp x,
-                    const std::optional<ResultAccuracy>& result_accuracy,
-                    bool expand);
+                    const std::optional<ResultAccuracy>& result_accuracy);
   friend XlaOp Tan(XlaOp operand,
                    const std::optional<ResultAccuracy>& result_accuracy);
   friend XlaOp Tanh(XlaOp operand,
diff --git a/third_party/xla/xla/hlo/transforms/expanders/BUILD b/third_party/xla/xla/hlo/transforms/expanders/BUILD
index a3e41f7e6b4..ef8068874c1 100644
--- a/third_party/xla/xla/hlo/transforms/expanders/BUILD
+++ b/third_party/xla/xla/hlo/transforms/expanders/BUILD
@@ -281,6 +281,44 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "elementary_function_expander",
+    srcs = ["elementary_function_expander.cc"],
+    hdrs = ["elementary_function_expander.h"],
+    deps = [
+        ":op_expander_pass",
+        "//xla:shape_util",
+        "//xla:util",
+        "//xla:xla_data_proto_cc",
+        "//xla/hlo/builder:xla_builder",
+        "//xla/hlo/builder:xla_computation",
+        "//xla/hlo/builder/lib:constants",
+        "//xla/hlo/ir:hlo",
+        "//xla/service:call_inliner",
+        "//xla/service:hlo_creation_utils",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+    ],
+)
+
+xla_cc_test(
+    name = "elementary_function_expander_test",
+    srcs = ["elementary_function_expander_test.cc"],
+    deps = [
+        ":elementary_function_expander",
+        "//xla:literal",
+        "//xla:literal_util",
+        "//xla:types",
+        "//xla/hlo/evaluator:hlo_evaluator",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 xla_cc_test(
     name = "logistic_expander_test",
     srcs = ["logistic_expander_test.cc"],
diff --git a/third_party/xla/xla/hlo/transforms/expanders/elementary_function_expander.cc b/third_party/xla/xla/hlo/transforms/expanders/elementary_function_expander.cc
new file mode 100644
index 00000000000..99ad399da94
--- /dev/null
+++ b/third_party/xla/xla/hlo/transforms/expanders/elementary_function_expander.cc
@@ -0,0 +1,157 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/hlo/transforms/expanders/elementary_function_expander.h"
+
+#include <string>
+
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "xla/hlo/builder/lib/constants.h"
+#include "xla/hlo/builder/xla_builder.h"
+#include "xla/hlo/builder/xla_computation.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/primitive_util.h"
+#include "xla/service/call_inliner.h"
+#include "xla/service/hlo_creation_utils.h"
+#include "xla/xla_data.pb.h"
+
+namespace xla {
+namespace {
+
+// For real floating-point types:
+//   Sinh(x) = sign(x) * (t - t^2 / (2 * (1 + t)))                  if |x| < 1
+//           = sign(x) * (e^(|x| - c) +
+//                        (e^(|x| - c) * s_lo - e^(-c - |x|) * s)) if |x| >= 1
+// where t = expm1(|x|), c = 0x1.63p-1 (~0.693359375), s = e^c / 2, and
+// s_lo = s - 1 = 0x1.bd0d1c3050e04p-13. Sub-32-bit floats are upcast to f32.
+//
+// For complex types (c64, c128), the shifted-exponential formula is used for
+// all inputs.
+XlaOp ExpandSinh(XlaOp x, PrimitiveType type) {
+  if (primitive_util::IsFloatingPointType(type) &&
+      primitive_util::BitWidth(type) < 32) {
+    return ConvertElementType(ExpandSinh(ConvertElementType(x, F32), F32),
+                              type);
+  }
+  ResultAccuracy highest_accuracy;
+  highest_accuracy.set_mode(ResultAccuracy::HIGHEST);
+
+  // Computing (e^x - e^-x) / 2 directly overflows for large |x| when e^|x| =
+  // inf but (e^|x|) / 2 is not inf. Shifting the exponent as e^(|x| - ln(2))
+  // avoids that overflow, but rounding error in |x| - ln(2) loses accuracy for
+  // large |x|. Instead, we subtract c = 0x1.63p-1 (~= ln(2)), which has few
+  // significant bits so |x| - c has no rounding error for 1 <= |x| < 1024, and
+  // compensate with scale = e^c / 2 = 1 + scale_lo:
+  //   sinh(|x|) = (e^|x| - e^-|x|) / 2
+  //             = (e^c * e^(|x| - c) - e^c * e^(-c - |x|)) / 2
+  //             = (e^c / 2) * (e^(|x| - c) - e^(-c - |x|))
+  //             = scale * e^(|x| - c) - scale * e^(-c - |x|)
+  //             = e^(|x| - c) + (e^(|x| - c) * scale_lo - e^(-c - |x|) * scale)
+  // Since scale = 1 + scale_lo is very close to 1 (scale_lo ~= 2.12e-4),
+  // storing scale as a single float loses ~12 bits of precision in scale_lo.
+  // Multiplying the large e^(|x| - c) term by 1 + scale_lo keeps those extra
+  // bits. By contrast, for |x| >= 1 the e^(-c - |x|) term is much smaller than
+  // e^(|x| - c) (and underflows to 0 when |x| is large), so multiplying it by
+  // scale directly is accurate enough. We combine the two smaller terms
+  // (e^(|x| - c) * scale_lo - e^(-c - |x|) * scale) before adding them to
+  // e^(|x| - c).
+  auto c = ScalarLike(x, 0x1.63p-1);
+  auto scale = ScalarLike(x, 0x1.000de868e1828p+0);      // e^c / 2
+  auto scale_lo = ScalarLike(x, 0x1.bd0d1c3050e04p-13);  // e^c / 2 - 1
+  XlaOp y = primitive_util::IsComplexType(type) ? x : Abs(x);
+  auto exp_add = Exp(y - c, highest_accuracy);
+  auto exp_sub = Exp(-c - y, highest_accuracy);
+  auto large_sinh_result = exp_add + (exp_add * scale_lo - exp_sub * scale);
+  if (primitive_util::IsComplexType(type)) {
+    return large_sinh_result;
+  }
+
+  // For float64 and |x| < 1, avoid cancellation in e^|x| - e^-|x| near 0 by
+  // writing sinh(|x|) in terms of t = expm1(|x|) = e^|x| - 1 (so e^|x| = 1 +
+  // t):
+  //   sinh(|x|) = (e^|x| - e^-|x|) / 2
+  //             = (e^|x| - 1 / e^|x|) / 2
+  //             = ((1 + t) - 1 / (1 + t)) / 2
+  //             = ((1 + t)^2 - 1) / (2 * (1 + t))
+  //             = (2t + t^2) / (2 * (1 + t))
+  //             = (2t * (1 + t) - t^2) / (2 * (1 + t))
+  //             = t - t^2 / (2 * (1 + t))
+  // Keeping t as the leading term and applying the division only to the O(t^2)
+  // correction damps rounding errors from 1 + t and the division by a factor of
+  // t / (2 * (1 + t)).
+  auto one_half = ScalarLike(x, 0.5);
+  auto one = ScalarLike(y, 1.0);
+  auto t = Expm1(y, highest_accuracy);
+  auto small_sinh_result = t - one_half * (t * t) / (one + t);
+  auto result = Select(Lt(y, one), small_sinh_result, large_sinh_result);
+  return Sign(x) * result;
+}
+
+}  // namespace
+
+bool ElementaryFunctionExpander::DefaultFilter(
+    const HloInstruction* instruction) {
+  switch (instruction->opcode()) {
+    case HloOpcode::kSinh: {
+      PrimitiveType type = instruction->shape().element_type();
+      return type == F64 || type == C64 || type == C128;
+    }
+    default:
+      return false;
+  }
+}
+
+bool ElementaryFunctionExpander::InstructionMatchesPattern(
+    HloInstruction* instruction) {
+  switch (instruction->opcode()) {
+    case HloOpcode::kSinh:
+      return true;
+    default:
+      return false;
+  }
+}
+
+absl::StatusOr<HloInstruction*> ElementaryFunctionExpander::ExpandInstruction(
+    HloInstruction* instruction) {
+  XlaBuilder b(std::string(instruction->name()));
+  XlaOp x = Parameter(&b, 0, instruction->operand(0)->shape(), "x");
+  switch (instruction->opcode()) {
+    case HloOpcode::kSinh:
+      ExpandSinh(x, instruction->shape().element_type());
+      break;
+    default:
+      return absl::InternalError(
+          absl::StrCat("Unsupported instruction: ", instruction->ToString()));
+  }
+  ABSL_ASSIGN_OR_RETURN(XlaComputation xla_computation, b.Build());
+  ABSL_ASSIGN_OR_RETURN(HloComputation * computation,
+                   XlaComputationToHloComputation(xla_computation,
+                                                  instruction->GetModule()));
+  HloInstruction* call =
+      instruction->parent()->AddInstruction(HloInstruction::CreateCall(
+          instruction->shape(), instruction->operands(), computation));
+  ABSL_ASSIGN_OR_RETURN(auto inlined_map, CallInliner::Inline(call));
+  HloInstruction* new_root = inlined_map[computation->root_instruction()];
+  ABSL_RETURN_IF_ERROR(
+      instruction->GetModule()->RemoveEmbeddedComputation(computation));
+  return new_root;
+}
+
+}  // namespace xla
diff --git a/third_party/xla/xla/hlo/transforms/expanders/elementary_function_expander.h b/third_party/xla/xla/hlo/transforms/expanders/elementary_function_expander.h
new file mode 100644
index 00000000000..ecba634b395
--- /dev/null
+++ b/third_party/xla/xla/hlo/transforms/expanders/elementary_function_expander.h
@@ -0,0 +1,51 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_HLO_TRANSFORMS_EXPANDERS_ELEMENTARY_FUNCTION_EXPANDER_H_
+#define XLA_HLO_TRANSFORMS_EXPANDERS_ELEMENTARY_FUNCTION_EXPANDER_H_
+
+#include <utility>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/transforms/expanders/op_expander_pass.h"
+#include "xla/util.h"
+
+namespace xla {
+
+// A pass which expands elementary math functions (such as sinh, cosh, etc.)
+// into other HLO instructions when they are not natively supported or when
+// an algebraic expansion is preferred for the operand type.
+class ElementaryFunctionExpander : public OpExpanderPass {
+ public:
+  static bool DefaultFilter(const HloInstruction* instruction);
+
+  explicit ElementaryFunctionExpander(HloPredicate extra_filter = DefaultFilter)
+      : OpExpanderPass(std::move(extra_filter)) {}
+  ~ElementaryFunctionExpander() override = default;
+  absl::string_view name() const override {
+    return "elementary-function-expander";
+  }
+
+ private:
+  bool InstructionMatchesPattern(HloInstruction* instruction) override;
+  absl::StatusOr<HloInstruction*> ExpandInstruction(
+      HloInstruction* instruction) override;
+};
+
+}  // namespace xla
+
+#endif  // XLA_HLO_TRANSFORMS_EXPANDERS_ELEMENTARY_FUNCTION_EXPANDER_H_
diff --git a/third_party/xla/xla/hlo/transforms/expanders/elementary_function_expander_test.cc b/third_party/xla/xla/hlo/transforms/expanders/elementary_function_expander_test.cc
new file mode 100644
index 00000000000..d8adc1985bb
--- /dev/null
+++ b/third_party/xla/xla/hlo/transforms/expanders/elementary_function_expander_test.cc
@@ -0,0 +1,221 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/hlo/transforms/expanders/elementary_function_expander.h"
+
+#include <cmath>
+#include <complex>
+#include <cstddef>
+#include <limits>
+#include <vector>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "xla/hlo/evaluator/hlo_evaluator.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
+#include "xla/literal.h"
+#include "xla/literal_util.h"
+#include "xla/types.h"
+
+namespace xla {
+namespace {
+
+using ElementaryFunctionExpanderTest = HloHardwareIndependentTestBase;
+
+bool ContainsOpcode(const HloModule& module, HloOpcode opcode) {
+  for (const HloComputation* comp : module.computations()) {
+    for (const HloInstruction* inst : comp->instructions()) {
+      if (inst->opcode() == opcode) {
+        return true;
+      }
+    }
+  }
+  return false;
+}
+
+TEST_F(ElementaryFunctionExpanderTest, SinhF32AndSub32BitTypesNotExpanded) {
+  const char* hlo_string = R"(
+    HloModule module
+
+    ENTRY main {
+      p_f32 = f32[4] parameter(0)
+      p_f16 = f16[4] parameter(1)
+      p_bf16 = bf16[4] parameter(2)
+      s_f32 = f32[4] sinh(p_f32)
+      s_f16 = f16[4] sinh(p_f16)
+      s_bf16 = bf16[4] sinh(p_bf16)
+      ROOT out = (f32[4], f16[4], bf16[4]) tuple(s_f32, s_f16, s_bf16)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  ElementaryFunctionExpander expander;
+  ASSERT_OK_AND_ASSIGN(bool changed, expander.Run(module.get()));
+  EXPECT_FALSE(changed);
+}
+
+TEST_F(ElementaryFunctionExpanderTest,
+       SinhF32AndSub32BitTypesExpandedUnfiltered) {
+  const char* hlo_string = R"(
+    HloModule module
+
+    ENTRY main {
+      p_f32 = f32[4] parameter(0)
+      p_f16 = f16[4] parameter(1)
+      p_bf16 = bf16[4] parameter(2)
+      s_f32 = f32[4] sinh(p_f32)
+      s_f16 = f16[4] sinh(p_f16)
+      s_bf16 = bf16[4] sinh(p_bf16)
+      ROOT out = (f32[4], f16[4], bf16[4]) tuple(s_f32, s_f16, s_bf16)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  ElementaryFunctionExpander expander(/*extra_filter=*/nullptr);
+  ASSERT_OK_AND_ASSIGN(bool changed, expander.Run(module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_FALSE(ContainsOpcode(*module, HloOpcode::kSinh));
+}
+
+TEST_F(ElementaryFunctionExpanderTest, SinhF64ExpandedAndAccurate) {
+  const char* hlo_string = R"(
+    HloModule module
+
+    ENTRY main {
+      p = f64[14] parameter(0)
+      ROOT s = f64[14] sinh(p)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  ElementaryFunctionExpander expander;
+  ASSERT_OK_AND_ASSIGN(bool changed, expander.Run(module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_FALSE(ContainsOpcode(*module, HloOpcode::kSinh));
+
+  const double inf = std::numeric_limits<double>::infinity();
+  const double nan = std::numeric_limits<double>::quiet_NaN();
+  std::vector<double> inputs = {
+      0.0,
+      -0.0,
+      1e-12,
+      -1e-12,
+      0.5,
+      -0.5,
+      1.0,
+      -1.0,
+      10.0,
+      -10.0,
+      // 710.0 > ln(DBL_MAX) (~709.78): exp(710.0) overflows to +inf, but
+      // sinh(710.0) is finite (~1.1169973830808757e308).
+      710.0,
+      -710.0,
+      inf,
+      nan,
+  };
+  Literal input_literal = LiteralUtil::CreateR1<double>(inputs);
+  HloEvaluator evaluator;
+  ASSERT_OK_AND_ASSIGN(Literal result_literal,
+                       evaluator.Evaluate(*module, {&input_literal}));
+  auto results = result_literal.data<double>();
+
+  // Verify signed zero preservation.
+  EXPECT_EQ(results[0], 0.0);
+  EXPECT_FALSE(std::signbit(results[0]));
+  EXPECT_EQ(results[1], -0.0);
+  EXPECT_TRUE(std::signbit(results[1]));
+
+  for (size_t i = 2; i < 12; ++i) {
+    double expected = std::sinh(inputs[i]);
+    double actual = results[i];
+    EXPECT_TRUE(std::isfinite(actual)) << "i=" << i << " x=" << inputs[i];
+    double ulp_tol =
+        2.0 * std::abs(expected) * std::numeric_limits<double>::epsilon();
+    EXPECT_NEAR(actual, expected, ulp_tol) << "i=" << i << " x=" << inputs[i];
+  }
+
+  EXPECT_TRUE(std::isinf(results[12]) && results[12] > 0);
+  EXPECT_TRUE(std::isnan(results[13]));
+}
+
+TEST_F(ElementaryFunctionExpanderTest, SinhComplexExpandedAndAccurate) {
+  const char* hlo_string = R"(
+    HloModule module
+
+    ENTRY main {
+      p64 = c64[3] parameter(0)
+      p128 = c128[3] parameter(1)
+      sinh_c64 = c64[3] sinh(p64)
+      sinh_c128 = c128[3] sinh(p128)
+      ROOT out = (c64[3], c128[3]) tuple(sinh_c64, sinh_c128)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_string));
+  ElementaryFunctionExpander expander;
+  ASSERT_OK_AND_ASSIGN(bool changed, expander.Run(module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_FALSE(ContainsOpcode(*module, HloOpcode::kSinh));
+
+  // Include 89.0 + 0.5i for c64 (where exp(89.0f) overflows float32) and
+  // 710.0 + 0.5i for c128 (where exp(710.0) overflows float64).
+  std::vector<complex64> inputs_c64 = {
+      {0.5f, 0.3f},
+      {-2.0f, 1.5f},
+      {89.0f, 0.5f},
+  };
+  std::vector<complex128> inputs_c128 = {
+      {0.5, 0.3},
+      {-2.0, 1.5},
+      {710.0, 0.5},
+  };
+  Literal input_c64_lit = LiteralUtil::CreateR1<complex64>(inputs_c64);
+  Literal input_c128_lit = LiteralUtil::CreateR1<complex128>(inputs_c128);
+
+  HloEvaluator evaluator;
+  ASSERT_OK_AND_ASSIGN(
+      Literal result_literal,
+      evaluator.Evaluate(*module, {&input_c64_lit, &input_c128_lit}));
+
+  std::vector<Literal> tuple_elements = result_literal.DecomposeTuple();
+  auto results_c64 = tuple_elements[0].data<complex64>();
+  auto results_c128 = tuple_elements[1].data<complex128>();
+
+  for (size_t i = 0; i < inputs_c64.size(); ++i) {
+    std::complex<double> expected = std::sinh(
+        std::complex<double>(inputs_c64[i].real(), inputs_c64[i].imag()));
+    EXPECT_TRUE(std::isfinite(results_c64[i].real())) << "i=" << i;
+    EXPECT_TRUE(std::isfinite(results_c64[i].imag())) << "i=" << i;
+    double tol_r =
+        4.0 * std::abs(expected.real()) * std::numeric_limits<float>::epsilon();
+    double tol_i =
+        4.0 * std::abs(expected.imag()) * std::numeric_limits<float>::epsilon();
+    EXPECT_NEAR(results_c64[i].real(), expected.real(), tol_r) << "i=" << i;
+    EXPECT_NEAR(results_c64[i].imag(), expected.imag(), tol_i) << "i=" << i;
+  }
+
+  for (size_t i = 0; i < inputs_c128.size(); ++i) {
+    std::complex<double> expected = std::sinh(inputs_c128[i]);
+    EXPECT_TRUE(std::isfinite(results_c128[i].real())) << "i=" << i;
+    EXPECT_TRUE(std::isfinite(results_c128[i].imag())) << "i=" << i;
+    double tol_r = 4.0 * std::abs(expected.real()) *
+                   std::numeric_limits<double>::epsilon();
+    double tol_i = 4.0 * std::abs(expected.imag()) *
+                   std::numeric_limits<double>::epsilon();
+    EXPECT_NEAR(results_c128[i].real(), expected.real(), tol_r) << "i=" << i;
+    EXPECT_NEAR(results_c128[i].imag(), expected.imag(), tol_i) << "i=" << i;
+  }
+}
+
+}  // namespace
+}  // namespace xla
diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc
index 4a501b3a688..a0c7bcef615 100644
--- a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc
+++ b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc
@@ -4891,8 +4891,7 @@ LogicalResult ExportXlaOp(SinhOp op, OpLoweringContext ctx) {
   if (failed(GetXlaOp(op.getOperand(), value_map, &operand, op))) {
     return failure();
   }
-  value_map[op] =
-      xla::Sinh(operand, /*result_accuracy=*/std::nullopt, /*expand=*/false);
+  value_map[op] = xla::Sinh(operand, /*result_accuracy=*/std::nullopt);
   return success();
 }
 
diff --git a/third_party/xla/xla/hlo/translate/stablehlo.cc b/third_party/xla/xla/hlo/translate/stablehlo.cc
index e5e58ce57b1..b23ab391815 100644
--- a/third_party/xla/xla/hlo/translate/stablehlo.cc
+++ b/third_party/xla/xla/hlo/translate/stablehlo.cc
@@ -73,7 +73,8 @@ absl::Status StablehloToMhlo(mlir::ModuleOp module, bool run_canonicalizer) {
       mlir::stablehlo_ext::createChloRecomposeOpsPass());
   pm.addPass(mlir::createSymbolDCEPass());
   pm.addNestedPass<mlir::func::FuncOp>(
-      mlir::mhlo::createChloLegalizeToHighLevelMhloPass());
+      mlir::mhlo::createChloLegalizeToHighLevelMhloPass(
+          mlir::mhlo::getDefaultChloToHighLevelMhloOptions()));
   pm.addNestedPass<mlir::func::FuncOp>(
       mlir::stablehlo::createChloLegalizeToStablehloPass());
   if (run_canonicalizer) {
diff --git a/third_party/xla/xla/mlir_hlo/mhlo/transforms/chlo_legalize_to_hlo/chlo_legalize_to_hlo_pass.cc b/third_party/xla/xla/mlir_hlo/mhlo/transforms/chlo_legalize_to_hlo/chlo_legalize_to_hlo_pass.cc
index 8a1ab26eb4b..3e66a53889a 100644
--- a/third_party/xla/xla/mlir_hlo/mhlo/transforms/chlo_legalize_to_hlo/chlo_legalize_to_hlo_pass.cc
+++ b/third_party/xla/xla/mlir_hlo/mhlo/transforms/chlo_legalize_to_hlo/chlo_legalize_to_hlo_pass.cc
@@ -356,7 +356,21 @@ LogicalResult convertMulhiChloToMhlo(chlo::MulhiOp op,
 }  // namespace
 
 ChloLegalizeToHighLevelMhloPassOptions getDefaultChloToHighLevelMhloOptions() {
-  return ChloLegalizeToHighLevelMhloPassOptions();
+  ChloLegalizeToHighLevelMhloPassOptions opts;
+  opts.enable_sinh_ = true;
+  return opts;
+}
+
+ChloLegalizeToHighLevelMhloPassOptions getCpuChloToHighLevelMhloOptions() {
+  ChloLegalizeToHighLevelMhloPassOptions opts;
+  opts.enable_sinh_ = true;
+  return opts;
+}
+
+ChloLegalizeToHighLevelMhloPassOptions getTpuChloToHighLevelMhloOptions() {
+  ChloLegalizeToHighLevelMhloPassOptions opts;
+  opts.enable_sinh_ = true;
+  return opts;
 }
 
 ChloLegalizeToHighLevelMhloPassOptions getGpuChloToHighLevelMhloOptions() {
@@ -412,13 +426,14 @@ void populateChloToHighLevelMhloOpPatterns(
 void populateChloToHighLevelMhloOpPatterns(MLIRContext* context,
                                            RewritePatternSet* patterns) {
   populateChloToHighLevelMhloOpPatterns(
-      context, patterns, mhlo::ChloLegalizeToHighLevelMhloPassOptions());
+      context, patterns, mhlo::getDefaultChloToHighLevelMhloOptions());
 }
 
 void populateChloToHloPatterns(MLIRContext* context,
                                TypeConverter* typeConverter,
                                RewritePatternSet* patterns) {
-  populateChloToHighLevelMhloOpPatterns(context, patterns);
+  populateChloToHighLevelMhloOpPatterns(
+      context, patterns, mhlo::ChloLegalizeToHighLevelMhloPassOptions());
   stablehlo::populateChloToStablehloPatterns(context, patterns);
   stablehlo::populateStablehloToHloPatterns(patterns, typeConverter, context);
 }
diff --git a/third_party/xla/xla/mlir_hlo/mhlo/transforms/passes.h b/third_party/xla/xla/mlir_hlo/mhlo/transforms/passes.h
index a1142a7759b..c0890e261e7 100644
--- a/third_party/xla/xla/mlir_hlo/mhlo/transforms/passes.h
+++ b/third_party/xla/xla/mlir_hlo/mhlo/transforms/passes.h
@@ -41,6 +41,12 @@ namespace mhlo {
 /// options specify the ops that are supported by all XLA backends.
 ChloLegalizeToHighLevelMhloPassOptions getDefaultChloToHighLevelMhloOptions();
 
+/// Returns options for the ChloLegalizeToHighLevelMhloPass for the CPU backend.
+ChloLegalizeToHighLevelMhloPassOptions getCpuChloToHighLevelMhloOptions();
+
+/// Returns options for the ChloLegalizeToHighLevelMhloPass for the TPU backend.
+ChloLegalizeToHighLevelMhloPassOptions getTpuChloToHighLevelMhloOptions();
+
 /// Returns options for the ChloLegalizeToHighLevelMhloPass for the GPU backend.
 ChloLegalizeToHighLevelMhloPassOptions getGpuChloToHighLevelMhloOptions();
 
diff --git a/third_party/xla/xla/pjrt/cpu/BUILD b/third_party/xla/xla/pjrt/cpu/BUILD
index 05fd6604bd7..cf4962d75a2 100644
--- a/third_party/xla/xla/pjrt/cpu/BUILD
+++ b/third_party/xla/xla/pjrt/cpu/BUILD
@@ -204,6 +204,7 @@ cc_library(
         "//xla/client:executable_build_options",
         "//xla/hlo/builder:xla_computation",
         "//xla/hlo/ir:hlo",
+        "//xla/mlir_hlo:mhlo_passes",
         "//xla/pjrt:async_work_runner",
         "//xla/pjrt:common_pjrt_client",
         "//xla/pjrt:compiled_memory_stats",
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
index f6d284af426..f7ab6563594 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
@@ -66,6 +66,7 @@ limitations under the License.
 #include "xla/layout_util.h"
 #include "xla/literal.h"
 #include "xla/literal_util.h"
+#include "xla/mlir_hlo/mhlo/transforms/passes.h"
 #include "xla/pjrt/common_pjrt_client.h"
 #include "xla/pjrt/compiled_memory_stats.h"
 #include "xla/pjrt/cpu/abstract_cpu_buffer.h"
@@ -286,7 +287,8 @@ absl::StatusOr<MlirCompilationSetup> SetupMlirCompilation(
   ABSL_RETURN_IF_ERROR(MlirToXlaComputation(
       module.mlir_module(), xla_computation,
       /*use_tuple_args=*/options.parameter_is_tupled_arguments,
-      /*return_tuple=*/false, &exec_build_options));
+      /*return_tuple=*/false, &exec_build_options,
+      mlir::mhlo::getCpuChloToHighLevelMhloOptions()));
   xla_computation.mutable_proto()->set_pjrt_id(module_id);
   setup.computation = std::move(xla_computation);
 
diff --git a/third_party/xla/xla/python/_ops.pyi b/third_party/xla/xla/python/_ops.pyi
index d8ae9a17dff..2973754c821 100644
--- a/third_party/xla/xla/python/_ops.pyi
+++ b/third_party/xla/xla/python/_ops.pyi
@@ -628,7 +628,7 @@ def Acosh(arg0: XlaOp, arg1: ResultAccuracy | None, arg2: bool, /) -> XlaOp: ...
 def Asinh(arg0: XlaOp, arg1: ResultAccuracy | None, arg2: bool, /) -> XlaOp: ...
 def Atanh(arg0: XlaOp, arg1: ResultAccuracy | None, arg2: bool, /) -> XlaOp: ...
 def Cosh(arg0: XlaOp, arg1: ResultAccuracy | None, arg2: bool, /) -> XlaOp: ...
-def Sinh(arg0: XlaOp, arg1: ResultAccuracy | None, arg2: bool, /) -> XlaOp: ...
+def Sinh(arg0: XlaOp, arg1: ResultAccuracy | None, /) -> XlaOp: ...
 def Real(arg: XlaOp, /) -> XlaOp: ...
 def Imag(arg: XlaOp, /) -> XlaOp: ...
 def Conj(arg: XlaOp, /) -> XlaOp: ...
diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD
index 3a47e7be030..fbe5c41e36c 100644
--- a/third_party/xla/xla/service/cpu/BUILD
+++ b/third_party/xla/xla/service/cpu/BUILD
@@ -225,6 +225,7 @@ cc_library(
         "//xla/hlo/transforms/expanders:dot_decomposer",
         "//xla/hlo/transforms/expanders:dynamic_index_splitter",
         "//xla/hlo/transforms/expanders:eigh_expander",
+        "//xla/hlo/transforms/expanders:elementary_function_expander",
         "//xla/hlo/transforms/expanders:logistic_expander",
         "//xla/hlo/transforms/expanders:optimization_barrier_expander",
         "//xla/hlo/transforms/expanders:permutation_sort_expander",
@@ -1675,6 +1676,7 @@ cc_library(
         "//xla/codegen/intrinsic:exp",
         "//xla/codegen/intrinsic:rsqrt",
         "//xla/codegen/intrinsic:tanh",
+        "//xla/codegen/intrinsic/cpp:_intrinsic_declarations",
         "//xla/hlo/ir:hlo",
         "//xla/service:elemental_ir_emitter",
         "//xla/service/llvm_ir:llvm_util",
diff --git a/third_party/xla/xla/service/cpu/cpu_compiler.cc b/third_party/xla/xla/service/cpu/cpu_compiler.cc
index dccc220b678..aebddc4fda3 100644
--- a/third_party/xla/xla/service/cpu/cpu_compiler.cc
+++ b/third_party/xla/xla/service/cpu/cpu_compiler.cc
@@ -124,6 +124,7 @@ limitations under the License.
 #include "xla/hlo/transforms/expanders/dot_decomposer.h"
 #include "xla/hlo/transforms/expanders/dynamic_index_splitter.h"
 #include "xla/hlo/transforms/expanders/eigh_expander.h"
+#include "xla/hlo/transforms/expanders/elementary_function_expander.h"
 #include "xla/hlo/transforms/expanders/logistic_expander.h"
 #include "xla/hlo/transforms/expanders/optimization_barrier_expander.h"
 #include "xla/hlo/transforms/expanders/permutation_sort_expander.h"
@@ -896,6 +897,7 @@ absl::Status CpuCompiler::RunHloPassesThroughLayoutAssn(
       /*rewrite_inference_op=*/true,
       /*rewrite_grad_op=*/true);
   pipeline.AddPass<LogisticExpander>();
+  pipeline.AddPass<ElementaryFunctionExpander>();
   pipeline.AddPass<ConditionalCanonicalizer>();
   pipeline.AddPass<DynamicDimensionSimplifier>();
 
diff --git a/third_party/xla/xla/service/cpu/elemental_ir_emitter.cc b/third_party/xla/xla/service/cpu/elemental_ir_emitter.cc
index 18a60659d79..71c14c3c583 100644
--- a/third_party/xla/xla/service/cpu/elemental_ir_emitter.cc
+++ b/third_party/xla/xla/service/cpu/elemental_ir_emitter.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/Value.h"
+#include "xla/codegen/intrinsic/cpp/intrinsic_declarations.h"
 #include "xla/codegen/intrinsic/exp.h"
 #include "xla/codegen/intrinsic/intrinsic.h"
 #include "xla/codegen/intrinsic/rsqrt.h"
@@ -90,6 +91,21 @@ absl::StatusOr<llvm::Value*> CpuElementalIrEmitter::EmitCosh(
 
 absl::StatusOr<llvm::Value*> CpuElementalIrEmitter::EmitSinh(
     PrimitiveType prim_type, llvm::Value* value) {
+  if (prim_type == F32) {
+    llvm::Function* sinh =
+        xla::codegen::intrinsics::Sinh::GetOrInsertDeclaration(module(),
+                                                               Type::S(F32));
+    return b()->CreateCall(sinh, value);
+  }
+  if (prim_type == F16 || prim_type == BF16) {
+    llvm::Type* f32_type = b()->getFloatTy();
+    llvm::Value* f32_value = b()->CreateFPCast(value, f32_type, "upcast");
+    llvm::Function* sinh =
+        xla::codegen::intrinsics::Sinh::GetOrInsertDeclaration(module(),
+                                                               Type::S(F32));
+    llvm::Value* f32_result = b()->CreateCall(sinh, f32_value);
+    return b()->CreateFPCast(f32_result, value->getType(), "downcast");
+  }
   return Unimplemented("sinh");
 }
 
diff --git a/third_party/xla/xla/service/cpu/tests/cpu_vectorization_test.cc b/third_party/xla/xla/service/cpu/tests/cpu_vectorization_test.cc
index 5489da59fb8..4f41e5238ad 100644
--- a/third_party/xla/xla/service/cpu/tests/cpu_vectorization_test.cc
+++ b/third_party/xla/xla/service/cpu/tests/cpu_vectorization_test.cc
@@ -481,6 +481,44 @@ TEST_F(Expm1JitVectorizationTest, Expm1F32) {
                                /*match_optimized_ir=*/true));
 }
 
+TEST_F(AtanJitVectorizationTest, SinhF32) {
+  if (!tsl::port::IsX86CPU()) {
+    GTEST_SKIP() << "This feature only works for x86 CPUs.";
+  }
+
+  const absl::string_view hlo_text = R"(
+    HloModule SinhF32
+
+    ENTRY SinhF32 {
+      p0 = f32[1024] parameter(0)
+      ROOT sinh = f32[1024] sinh(p0)
+    }
+  )";
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> hlo_module,
+                       ParseAndReturnVerifiedModule(hlo_text));
+
+  auto compiler = GetCpuCompiler();
+  auto llvm_compiler = absl::down_cast<LLVMCompiler*>(compiler.get());
+  Compiler::CompileOptions compile_options;
+  compile_options.device_allocator = nullptr;
+
+  int num_elements = 16;
+  if (!tsl::port::TestCPUFeature(tsl::port::CPUFeature::AVX512F)) {
+    if (tsl::port::TestCPUFeature(tsl::port::CPUFeature::AVX2)) {
+      num_elements = 8;
+    } else {
+      num_elements = 4;
+    }
+  }
+
+  std::string check_lines =
+      absl::StrFormat("CHECK: fdiv{{.*}}<%d x float>", num_elements);
+
+  ASSERT_OK(CompileAndVerifyIr(llvm_compiler, compile_options,
+                               std::move(hlo_module), check_lines,
+                               /*match_optimized_ir=*/true));
+}
 }  // namespace
 }  // namespace cpu
 }  // namespace xla
diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index c51ad03cf3f..66860a0c093 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -1945,6 +1945,7 @@ cc_library(
         "//xla/hlo/transforms/expanders:dot_decomposer",
         "//xla/hlo/transforms/expanders:dynamic_index_splitter",
         "//xla/hlo/transforms/expanders:eigh_expander",
+        "//xla/hlo/transforms/expanders:elementary_function_expander",
         "//xla/hlo/transforms/expanders:logistic_expander",
         "//xla/hlo/transforms/expanders:optimization_barrier_expander",
         "//xla/hlo/transforms/expanders:permutation_sort_expander",
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 82826a68593..31c6b4659f0 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -191,6 +191,7 @@ limitations under the License.
 #include "xla/hlo/transforms/expanders/dot_decomposer.h"
 #include "xla/hlo/transforms/expanders/dynamic_index_splitter.h"
 #include "xla/hlo/transforms/expanders/eigh_expander.h"
+#include "xla/hlo/transforms/expanders/elementary_function_expander.h"
 #include "xla/hlo/transforms/expanders/logistic_expander.h"
 #include "xla/hlo/transforms/expanders/optimization_barrier_expander.h"
 #include "xla/hlo/transforms/expanders/permutation_sort_expander.h"
@@ -920,6 +921,10 @@ absl::Status RunOptimizationPasses(
       /*rewrite_inference_op=*/true,
       /*rewrite_grad_op=*/true);
 
+  pipeline.AddPass<ElementaryFunctionExpander>(
+      [](const HloInstruction* instruction) {
+        return ShapeUtil::ElementIsComplex(instruction->shape());
+      });
   pipeline.AddPass<LogisticExpander>();
   pipeline.AddPass<ConditionalCanonicalizer>();
   pipeline.AddPass<DynamicDimensionSimplifier>();
diff --git a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_functions.cc b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_functions.cc
index ae6fb272665..4105f995acd 100644
--- a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_functions.cc
+++ b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_functions.cc
@@ -300,13 +300,16 @@ UNARY_TEST(Cosh, {
 UNARY_TEST(Sinh, {
   SinhOp<kT>(this)
       .Error(GetDefaultSpecGenerator())
-      .CpuError(+[](NativeT x) {
-        if constexpr (std::is_same_v<NativeT, tsl::float8_e4m3fn>) {
-          return ErrorSpec::Builder().distance_err(3).build();
-        } else if constexpr (std::is_same_v<NativeT, tsl::float8_e5m2>) {
-          return ErrorSpec::Builder().distance_err(4).build();
+      .CpuError(+[](NativeT) {
+        if constexpr (std::is_same_v<NativeT, bfloat16> ||
+                      std::is_same_v<NativeT, tsl::float8_e4m3fn> ||
+                      std::is_same_v<NativeT, tsl::float8_e5m2>) {
+          return ErrorSpec::Builder().distance_err(0).build();
+        } else if constexpr (std::is_same_v<NativeT, double>) {
+          return ErrorSpec::Builder().distance_err(2).build();
+        } else {
+          return ErrorSpec::Builder().distance_err(1).build();
         }
-        return GetDefaultSpecGenerator()(x);
       })
       .GpuError(+[](NativeT x) {
         if constexpr (std::is_same_v<NativeT, tsl::float8_e5m2>) {

From 2e95144da3c10d36594b39adb598c17265e2a228 Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Thu, 15 Oct 2026 18:09:15 -0700
Subject: [PATCH 0920/1000] [XLA:CPU] Support `math.exp2` and `math.log2` in
 `vectorize_xtile`.

Add `math::Exp2Op` and `math::Log2Op` to the vectorization patterns in `xla/backends/cpu/codegen/tiled/transforms/vectorize_xtile.cc` so tiled CPU fusions can vectorize `exp2` and `log2`.

PiperOrigin-RevId: 1000519987
---
 .../backends/cpu/codegen/tiled/transforms/vectorize_xtile.cc   | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/vectorize_xtile.cc b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/vectorize_xtile.cc
index 7db3664df77..90098ea05f6 100644
--- a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/vectorize_xtile.cc
+++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/vectorize_xtile.cc
@@ -1053,12 +1053,13 @@ class VectorizeXTilePass
       mm::CopySignOp,
       mm::CosOp, mm::CoshOp,
       mm::ErfOp,
+      mm::Exp2Op,
       mm::ExpM1Op,
       mm::ExpOp,
       mm::FloorOp,
       mm::IPowIOp,
       mm::IsFiniteOp,
-      mm::Log1pOp, mm::LogOp,
+      mm::Log1pOp, mm::Log2Op, mm::LogOp,
       mm::PowFOp,
       mm::RoundEvenOp,
       mm::RsqrtOp,

From e3d5ca8b5f29f860d84fceba890906184ae150b5 Mon Sep 17 00:00:00 2001
From: Yash Katariya <yashkatariya@google.com>
Date: Thu, 15 Oct 2026 19:14:22 -0700
Subject: [PATCH 0921/1000] Support non-sum unreduced axes across Shardy
 verification, reshard, and manual export

- Verify unreduced axes on WhileOp (including scalar cond predicate return), CaseOp/IfOp, NamedComputationOp, and ManualComputationOp in VerifyUnreducedAxesPass.
- Preserve reductionOp and unreducedAxes in addFreeAxesToManualComputationSharding, isShardingEquivalentAcrossReshapes, and ReshardToCollectives.
- Exclude outer bound manual axes in createFullyManualComputation and use sdy::getReductionIdentityAttr for non-sum ReplicatedToUnreducedOp in export_manual_reduction_collectives.
- Bump JAX_IFRT_VERSION_NUMBER to 77 and jaxlib_extension_version to 506, and add Shardy/XLA MLIR tests and version-guarded JAX tests.

PiperOrigin-RevId: 1000545472
---
 third_party/xla/xla/python/version.h          |  2 +-
 .../export_manual_reduction_collectives.cc    | 23 +++++++++--
 ...o_export_manual_reduction_collectives.mlir | 40 +++++++++++++++++++
 ..._reduction_collectives_reduce_scatter.mlir | 20 ++++++++++
 .../xla/xla/service/spmd/spmd_partitioner.cc  | 15 +++++--
 5 files changed, 93 insertions(+), 7 deletions(-)

diff --git a/third_party/xla/xla/python/version.h b/third_party/xla/xla/python/version.h
index 06a8c0ccff3..578c87d484f 100644
--- a/third_party/xla/xla/python/version.h
+++ b/third_party/xla/xla/python/version.h
@@ -18,6 +18,6 @@ limitations under the License.
 
 // An increasing version number to protect jax code against breaking changes.
 // In JAX, reference this via jax._src.lib.ifrt_version.
-#define JAX_IFRT_VERSION_NUMBER 76  // while_loop_simplifier GT compare fix
+#define JAX_IFRT_VERSION_NUMBER 77  // Shardy unreduced axes support
 
 #endif  // XLA_PYTHON_VERSION_H_
diff --git a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_manual_reduction_collectives.cc b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_manual_reduction_collectives.cc
index 7fbbe257209..584cd4c777b 100644
--- a/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_manual_reduction_collectives.cc
+++ b/third_party/xla/xla/service/spmd/shardy/stablehlo_round_trip/export_manual_reduction_collectives.cc
@@ -24,10 +24,13 @@ limitations under the License.
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/strings/string_view.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/CommandLine.h"
+#include "mlir/IR/Block.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
@@ -115,14 +118,27 @@ Value createFullyManualComputation(
     std::function<Value(mlir::BlockArgument arg, OpBuilder& blockBuilder,
                         RankedTensorType paddedGlobalType)>
         bodyPopulator) {
+  llvm::SmallDenseSet<StringRef> boundManualAxes;
+  if (mlir::Block* block = builder.getInsertionBlock()) {
+    for (mlir::Operation* parentOp = block->getParentOp(); parentOp != nullptr;
+         parentOp = parentOp->getParentOp()) {
+      if (auto parentManual = mlir::dyn_cast<ManualComputationOp>(parentOp)) {
+        for (StringRef axisName : parentManual.getManualAxes()) {
+          boundManualAxes.insert(axisName);
+        }
+      }
+    }
+  }
   SmallVector<mlir::StringAttr> manualAxes;
   manualAxes.reserve(mesh.getAxes().size());
   for (sdy::MeshAxisAttr axis : mesh.getAxes()) {
-    manualAxes.push_back(builder.getStringAttr(axis.getName()));
+    if (!boundManualAxes.contains(axis.getName())) {
+      manualAxes.push_back(builder.getStringAttr(axis.getName()));
+    }
   }
   TensorShardingAttr inSharding = sdy::getOrCreateSharding(input, mesh);
   if (!unreducedAxes.empty()) {
-    inSharding = inSharding.replaceUnreducedAxes(unreducedAxes);
+    inSharding = inSharding.replaceUnreducedAxes(unreducedAxes, reductionOp);
   }
 
   auto globalType = mlir::dyn_cast<RankedTensorType>(input.getType());
@@ -441,7 +457,8 @@ void syncInOutUnreducedAxes(mlir::Operation* op) {
   }
 
   sdy::setSharding(
-      input, inSharding.replaceUnreducedAxes(outSharding.getUnreducedAxes()));
+      input, inSharding.replaceUnreducedAxes(outSharding.getUnreducedAxes(),
+                                             outSharding.getReductionOp()));
 }
 
 class StablehloExportManualReductionCollectivesPass
diff --git a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_manual_reduction_collectives.mlir b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_manual_reduction_collectives.mlir
index 9ec1beb5b90..722860973f8 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_manual_reduction_collectives.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_manual_reduction_collectives.mlir
@@ -773,6 +773,46 @@ func.func @reduce_scatter_further_shard_non_divisible(%arg0: tensor<2x4xf32> {sd
   return %0 : tensor<2x4xf32>
 }
 
+// -----
+
+sdy.mesh @mesh_x_2_y_2 = <["x"=2, "y"=2]>
+
+// CHECK-LABEL: func @all_reduce_max_inside_manual_computation
+// CHECK:         sdy.manual_computation(%arg0)
+// CHECK-SAME:      in_shardings=[<@mesh_x_2_y_2, [{"y"}, {}], unreduced=max{"x"}>]
+// CHECK-SAME:      out_shardings=[<@mesh_x_2_y_2, [{"y"}, {}]>]
+// CHECK-SAME:      manual_axes={"y"} (%[[IN0:.*]]: tensor<2x4xf32>) {
+// CHECK:           sdy.manual_computation(%[[IN0]])
+// CHECK-SAME:        in_shardings=[<@mesh_x_2_y_2, [{}, {}], unreduced=max{"x"}>]
+// CHECK-SAME:        out_shardings=[<@mesh_x_2_y_2, [{}, {}]>]
+// CHECK-SAME:        manual_axes={"x"} (%[[IN1:.*]]: tensor<2x4xf32>) {
+// CHECK:             "stablehlo.all_reduce"(%[[IN1]])
+// CHECK-SAME:          replica_groups = dense<{{\[\[0, 2\], \[1, 3\]\]}}>
+// CHECK:               stablehlo.maximum
+func.func @all_reduce_max_inside_manual_computation(%arg0: tensor<4x4xf32> {sdy.sharding = #sdy.sharding<@mesh_x_2_y_2, [{"y"}, {}], unreduced=max{"x"}>}) -> tensor<4x4xf32> {
+  %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh_x_2_y_2, [{"y"}, {}], unreduced=max{"x"}>] out_shardings=[<@mesh_x_2_y_2, [{"y"}, {}]>] manual_axes={"y"} (%in0: tensor<2x4xf32>) {
+    %1 = sdy.all_reduce max {"x"} %in0 out_sharding=<@mesh_x_2_y_2, [{}, {}]> : tensor<2x4xf32>
+    sdy.return %1 : tensor<2x4xf32>
+  } : (tensor<4x4xf32>) -> tensor<4x4xf32>
+  return %0 : tensor<4x4xf32>
+}
+
+// -----
+
+sdy.mesh @mesh = <["x"=2, "y"=2]>
+
+// CHECK-LABEL: func @unreduced_max_sharded_transpose_of_sharded_dot_without_unreduced
+func.func @unreduced_max_sharded_transpose_of_sharded_dot_without_unreduced(%arg0: tensor<8x4xf32>, %arg1: tensor<4x2xf32>) -> tensor<2x8xf32> {
+  // CHECK-NEXT: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT]
+  // CHECK-SAME:   {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"x"}], unreduced=max{"y"}>]>}
+  // CHECK-NEXT: %[[TRANSPOSE:.*]] = stablehlo.transpose %[[DOT_GENERAL]], dims = [1, 0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x"}, {}], unreduced=max{"y"}>]>}
+  // CHECK-NEXT: return %[[TRANSPOSE]]
+  %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT]
+      {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{}, {"x"}]>]>} : (tensor<8x4xf32>, tensor<4x2xf32>) -> tensor<8x2xf32>
+  %1 = stablehlo.transpose %0, dims = [1, 0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh, [{"x"}, {}], unreduced=max{"y"}>]>} : (tensor<8x2xf32>) -> tensor<2x8xf32>
+  return %1 : tensor<2x8xf32>
+}
+
 
 
 
diff --git a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_manual_reduction_collectives_reduce_scatter.mlir b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_manual_reduction_collectives_reduce_scatter.mlir
index 5bb63f01b56..fa88262b2cd 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_manual_reduction_collectives_reduce_scatter.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/stablehlo_export_manual_reduction_collectives_reduce_scatter.mlir
@@ -37,6 +37,26 @@ func.func @reduce_scatter(%arg0: tensor<8x8xf32> {sdy.sharding=#sdy.sharding<@me
   return %0 : tensor<8x8xf32>
 }
 
+// CHECK-LABEL: func @reduce_scatter_max_no_unreduced_input
+// NO-EXPORT-LABEL: func @reduce_scatter_max_no_unreduced_input
+func.func @reduce_scatter_max_no_unreduced_input(%arg0: tensor<8x8xf32> {sdy.sharding=#sdy.sharding<@mesh, [{"x"}, {}]>}) -> tensor<8x8xf32> {
+  // CHECK-NEXT: %[[MANUAL_COMP:.*]] = sdy.manual_computation(%arg0)
+  // CHECK-SAME:     in_shardings=[<@mesh, [{"x"}, {}], unreduced=max{"y"}>]
+  // CHECK-SAME:     out_shardings=[<@mesh, [{"x"}, {"y"}]>]
+  // CHECK-SAME:     manual_axes={"x", "y"} (%arg1: tensor<1x8xf32>) {
+  // CHECK-NEXT:   %[[REDUCE_SCATTER:.*]] = "stablehlo.reduce_scatter"(%arg1)
+  // CHECK:          stablehlo.maximum
+  // CHECK:        sdy.return %[[REDUCE_SCATTER]]
+  // CHECK-NEXT: }
+  // CHECK-NEXT: return %[[MANUAL_COMP]]
+
+  // NO-EXPORT-NEXT: %[[RS:.*]] = sdy.reduce_scatter max [{}, {"y"}] %arg0
+  // NO-EXPORT-NEXT: return %[[RS]]
+
+  %0 = sdy.reduce_scatter max [{}, {"y"}] %arg0 out_sharding=<@mesh, [{"x"}, {"y"}]> : tensor<8x8xf32>
+  return %0 : tensor<8x8xf32>
+}
+
 // CHECK-LABEL: func @reduce_scatter_sort_and_merge
 func.func @reduce_scatter_sort_and_merge(%arg0: tensor<8x8x8xf32> {sdy.sharding=#sdy.sharding<@mesh_x_4, [{}, {}, {}], unreduced={"x"}>}) -> tensor<8x8x8xf32> {
   // CHECK-NEXT: %[[MANUAL_COMP:.*]] = sdy.manual_computation(%arg0)
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.cc b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
index 9dff289d687..25d931bb110 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
@@ -7340,6 +7340,9 @@ static HloSharding ResolveReductionOpForSharding(
   } else if (!sharding.IsUnreducedSubgroup()) {
     return sharding;
   }
+  if (sharding.reduction_op() != ReductionOp::kSum) {
+    return sharding;
+  }
 
   const HloInstruction* current = instruction;
   while (current && current->opcode() != HloOpcode::kReduce &&
@@ -7367,9 +7370,15 @@ static HloSharding ResolveReductionOpForSharding(
     if (module && module->has_spmd_parameters_shardings() &&
         current->parameter_number() <
             module->spmd_parameters_shardings().size()) {
-      sharding.set_reduction_op(
-          module->spmd_parameters_shardings()[current->parameter_number()]
-              .reduction_op());
+      const HloSharding& param_sharding =
+          module->spmd_parameters_shardings()[current->parameter_number()];
+      const bool param_has_unreduced =
+          param_sharding.UseNamedShardingLeaf()
+              ? !param_sharding.named_sharding().unreduced_axes().empty()
+              : param_sharding.IsUnreducedSubgroup();
+      if (param_has_unreduced) {
+        sharding.set_reduction_op(param_sharding.reduction_op());
+      }
     }
   }
   return sharding;

From 046efd7c1b15153db16fde64f8b917777ec243ac Mon Sep 17 00:00:00 2001
From: Yash Katariya <yashkatariya@google.com>
Date: Thu, 15 Oct 2026 20:33:15 -0700
Subject: [PATCH 0922/1000] Integrate OpenXLA Shardy at
 c29509bd0c1430beaa7e767c866a131492ab1139

PiperOrigin-RevId: 1000574622
---
 .../xla/third_party/shardy/temporary.patch    | 3958 -----------------
 .../xla/third_party/shardy/workspace.bzl      |    4 +-
 2 files changed, 2 insertions(+), 3960 deletions(-)

diff --git a/third_party/xla/third_party/shardy/temporary.patch b/third_party/xla/third_party/shardy/temporary.patch
index 8a259e4c00e..e69de29bb2d 100644
--- a/third_party/xla/third_party/shardy/temporary.patch
+++ b/third_party/xla/third_party/shardy/temporary.patch
@@ -1,3958 +0,0 @@
-diff --git a/shardy/dialect/sdy/ir/dialect.cc b/shardy/dialect/sdy/ir/dialect.cc
-index 719686a5..fc0561f4 100644
---- a/shardy/dialect/sdy/ir/dialect.cc
-+++ b/shardy/dialect/sdy/ir/dialect.cc
-@@ -969,9 +969,7 @@ Type TensorShardingAttr::getLocalType(Type globalType, MeshAttr mesh,
- RankedTensorType TensorShardingAttr::getLocalTensorType(
-     RankedTensorType globalTensorType, MeshAttr mesh,
-     bool allowNonDivisible) const {
--  // getLocalType should always return something castable to RankedTensorType if
--  // a RankedTensorType is passed in.
--  return cast<RankedTensorType>(
-+  return dyn_cast_or_null<RankedTensorType>(
-       getLocalType(globalTensorType, mesh, allowNonDivisible));
- }
- 
-@@ -1232,16 +1230,17 @@ TensorShardingAttr addFreeAxesToManualComputationSharding(
-   TensorShardingAttr returnedSharding =
-       eraseFreeAxes(outerManualSharding, manualAxes);
- 
--  SmallVector<DimensionShardingAttr> resultDimShardings(
--      returnedSharding.getDimShardings());
--  for (auto [resultDimSharding, newDimSharding] :
--       llvm::zip(resultDimShardings, newSharding.getDimShardings())) {
--    resultDimSharding = DimensionShardingAttr::get(
--        resultDimSharding.getContext(),
--        llvm::to_vector(llvm::concat<const AxisRefAttr>(
--            resultDimSharding.getAxes(), newDimSharding.getAxes())),
--        newDimSharding.getIsClosed(), newDimSharding.getPriority());
--  }
-+  SmallVector<DimensionShardingAttr> resultDimShardings = llvm::map_to_vector(
-+      llvm::zip_equal(returnedSharding.getDimShardings(),
-+                      newSharding.getDimShardings()),
-+      [](auto pair) {
-+        auto [returnedDimSharding, newDimSharding] = pair;
-+        return DimensionShardingAttr::get(
-+            returnedDimSharding.getContext(),
-+            llvm::to_vector(llvm::concat<const AxisRefAttr>(
-+                returnedDimSharding.getAxes(), newDimSharding.getAxes())),
-+            newDimSharding.getIsClosed(), newDimSharding.getPriority());
-+      });
-   return TensorShardingAttr::get(newSharding.getContext(),
-                                  newSharding.getMeshOrRef(), resultDimShardings,
-                                  outerManualSharding.getReplicatedAxes(),
-diff --git a/shardy/dialect/sdy/ir/dialect_test.cc b/shardy/dialect/sdy/ir/dialect_test.cc
-index 1ea9160e..8f149153 100644
---- a/shardy/dialect/sdy/ir/dialect_test.cc
-+++ b/shardy/dialect/sdy/ir/dialect_test.cc
-@@ -624,12 +624,25 @@ TEST_F(DialectTest, TensorShardingAttrGetLocalTensorType) {
-       {createDimSharding({createAxis("y")}), createDimSharding({})});
-   EXPECT_EQ(sharding3.getLocalTensorType(globalType, mesh),
-             RankedTensorType::get({3, 4}, builder.getI32Type()));
-+  EXPECT_EQ(sharding3.getLocalTensorType(globalType, mesh,
-+                                         /*allowNonDivisible=*/false),
-+            nullptr);
- 
-   TensorShardingAttr sharding4 = createTensorSharding(
-       {createDimSharding({createAxis("x"), createAxis("z")}),
-        createDimSharding({createAxis("y")})});
-   EXPECT_EQ(sharding4.getLocalTensorType(globalType, mesh),
-             RankedTensorType::get({1, 2}, builder.getI32Type()));
-+  EXPECT_EQ(sharding4.getLocalTensorType(globalType, mesh,
-+                                         /*allowNonDivisible=*/false),
-+            nullptr);
-+
-+  TensorShardingAttr sharding5 = createTensorSharding(
-+      {createDimSharding({createAxis("x"), createAxis("z")}),
-+       createDimSharding({createAxis("x")})});
-+  EXPECT_EQ(sharding5.getLocalTensorType(globalType, mesh,
-+                                         /*allowNonDivisible=*/false),
-+            RankedTensorType::get({1, 2}, builder.getI32Type()));
- }
- 
- TEST_F(DialectTest, DimensionShardingAttrGetSharded) {
-diff --git a/shardy/dialect/sdy/transforms/export/convert_global_to_local.cc b/shardy/dialect/sdy/transforms/export/convert_global_to_local.cc
-index c6b1a677..c77a0ad4 100644
---- a/shardy/dialect/sdy/transforms/export/convert_global_to_local.cc
-+++ b/shardy/dialect/sdy/transforms/export/convert_global_to_local.cc
-@@ -280,19 +280,6 @@ LogicalResult localizeGenericOp(Operation* op, ValueRange operands,
-   return success();
- }
- 
--void checkUnreducedResultSharding(TensorShardingAttr resSharding,
--                                  ArrayRef<AxisRefAttr> expectedAxes,
--                                  StringRef opName) {
--  SDY_CHECK(resSharding && !resSharding.getUnreducedAxes().empty())
--      << "Expected sharded reduction " << opName.str()
--      << " result to have unreduced axes sharding.";
--  SDY_CHECK(resSharding.getUnreducedAxes() == expectedAxes)
--      << "The unreduced axes of the " << opName.str()
--      << " result sharding should match the sharded reduction dimensions of "
--         "the "
--      << opName.str() << ".";
--}
--
- // Pattern for generic ops that do not require special handling beyond type
- // conversion, such as StableHLO ops, sdy.all_reduce, etc.
- class GenericOpPattern : public ConversionPattern {
-@@ -305,16 +292,15 @@ class GenericOpPattern : public ConversionPattern {
-   LogicalResult matchAndRewrite(
-       Operation* op, ArrayRef<Value> operands,
-       ConversionPatternRewriter& rewriter) const override {
--    Dialect* dialect = op->getDialect();
--    if ((dialect && dialect->getNamespace() != "stablehlo" &&
-+    if ((isa_and_nonnull<SdyDialect, func::FuncDialect>(op->getDialect()) &&
-          !isa<sdy::ReturnOp, func::CallOp, func::ReturnOp>(op)) ||
-         isa<stablehlo::ConcatenateOp, stablehlo::ConvolutionOp,
-             stablehlo::DotGeneralOp, stablehlo::DotOp,
--            stablehlo::DynamicUpdateSliceOp, stablehlo::GatherOp,
--            stablehlo::IotaOp, stablehlo::PadOp, stablehlo::ReduceOp,
--            stablehlo::ReduceWindowOp, stablehlo::RngBitGeneratorOp,
--            stablehlo::ScatterOp, stablehlo::SelectAndScatterOp,
--            stablehlo::SliceOp>(op)) {
-+            stablehlo::DynamicSliceOp, stablehlo::DynamicUpdateSliceOp,
-+            stablehlo::GatherOp, stablehlo::IotaOp, stablehlo::PadOp,
-+            stablehlo::ReduceOp, stablehlo::ReduceWindowOp,
-+            stablehlo::RngBitGeneratorOp, stablehlo::ScatterOp,
-+            stablehlo::SelectAndScatterOp, stablehlo::SliceOp>(op)) {
-       return failure();
-     }
-     return localizeGenericOp(
-@@ -432,23 +418,70 @@ std::pair<Attribute, int64_t> getReplicaGroupsAndSize(
-   return {replicaGroups, groupSize};
- }
- 
--bool hasUnreducedResultSharding(Operation* op,
--                                ArrayRef<AxisRefAttr> reductionAxes,
--                                StringRef opName,
--                                const GlobalToLocalTypeConverter* converter) {
--  bool hasUnreducedAxes = false;
--  for (Value res : op->getResults()) {
--    TensorShardingAttr resSharding = converter->getSharding(res);
--    if (resSharding && !resSharding.getUnreducedAxes().empty()) {
--      checkUnreducedResultSharding(resSharding, reductionAxes, opName);
--      hasUnreducedAxes = true;
--    } else if (llvm::any_of(res.getUsers(), [](Operation* user) {
--                 return isa<sdy::AllReduceOp, sdy::ReduceScatterOp>(user);
--               })) {
--      hasUnreducedAxes = true;
--    }
--  }
--  return hasUnreducedAxes;
-+// Returns whether the op introduces partial results related to the given
-+// reductionAxes for its operands with kReduction factors.
-+//
-+// The op introduces partial results if either:
-+//   - The result's sharding has unreduced axes containing reductionAxes, or
-+//   - the result is fed directly into an AllReduce or ReduceScatter on
-+//     reductionAxes.
-+//
-+// We assume if an op that can introduce partial result and produce multiple
-+// results, all of its results should be sharded the same way. This is the case
-+// with the current StableHLO ops.
-+bool doesOpIntroducePartialResult(Operation* op,
-+                                  ArrayRef<AxisRefAttr> reductionAxes,
-+                                  StringRef opName,
-+                                  const GlobalToLocalTypeConverter* converter) {
-+  auto containsReductionAxes = [&](ArrayRef<AxisRefAttr> axes, MeshAttr mesh) {
-+    if (axes.empty()) {
-+      return false;
-+    }
-+    SmallVector<AxisRefAttr> missingReductionAxes =
-+        getAxisSetDiff(reductionAxes, axes, mesh);
-+    if (missingReductionAxes.empty()) {
-+      return true;
-+    }
-+    // Otherwise, none of the sharded reduction axes of the op should be in
-+    // `axes` (the axes only come from operands' unreduced axes).
-+    SDY_CHECK(ArrayRef<AxisRefAttr>(missingReductionAxes) == reductionAxes)
-+        << "The unreduced or reduction axes for the " << opName.str()
-+        << " result must either contain all or none of the sharded "
-+           "reduction dimensions of the "
-+        << opName.str() << ".";
-+    return false;
-+  };
-+
-+  // All results of a reduction op share the same sharding and unreduced axes,
-+  // so we only need to inspect the first result.
-+  Value res = op->getResult(0);
-+  TensorShardingAttr resSharding = converter->getSharding(res);
-+  const SymbolTable& symbolTable = converter->getSymbolTable();
-+  if (resSharding && containsReductionAxes(resSharding.getUnreducedAxes(),
-+                                           resSharding.getMesh(symbolTable))) {
-+    return true;
-+  }
-+
-+  // Otherwise, if the op is fed directly into an AllReduce or ReduceScatter on
-+  // reductionAxes, we consider it introduce a partial result and the upstream
-+  // passes forgot to properly annotate the result with unreduced axes.
-+  return llvm::any_of(res.getUsers(), [&](Operation* user) {
-+    if (auto allReduce = dyn_cast<sdy::AllReduceOp>(user)) {
-+      return containsReductionAxes(
-+          allReduce.getReductionAxes(),
-+          allReduce.getOutSharding().getMesh(symbolTable));
-+    }
-+    if (auto reduceScatter = dyn_cast<sdy::ReduceScatterOp>(user)) {
-+      SmallVector<AxisRefAttr> reduceScatterAxes;
-+      for (AxisRefListAttr axesList : reduceScatter.getReduceScatterAxes()) {
-+        llvm::append_range(reduceScatterAxes, axesList.getValue());
-+      }
-+      return containsReductionAxes(
-+          reduceScatterAxes,
-+          reduceScatter.getOutSharding().getMesh(symbolTable));
-+    }
-+    return false;
-+  });
- }
- 
- ValueRange createAllReduceForUnreducedAxes(
-@@ -528,18 +561,23 @@ class AllGatherOpPattern : public OpConversionPattern<AllGatherOp> {
-     }
- 
-     SDY_CHECK(gatheringDims.size() > 1);
--
--    // Add a leading dimension of size 1 to gather all partitions at once.
--    SmallVector<int64_t> tmpShape = {1};
--    llvm::append_range(tmpShape, localShape);
-+    int64_t firstGatheringDim = gatheringDims.front();
-+
-+    // Add a dimension of size 1 at `firstGatheringDim` to gather all
-+    // partitions at once, keeping dimensions 0..firstGatheringDim-1 untouched.
-+    SmallVector<int64_t> tmpShape =
-+        llvm::to_vector(localShape.take_front(firstGatheringDim));
-+    tmpShape.push_back(1);
-+    llvm::append_range(tmpShape, localShape.drop_front(firstGatheringDim));
-     Value curInput = stablehlo::ReshapeOp::create(
-         rewriter, op.getLoc(),
-         RankedTensorType::get(tmpShape, localType.getElementType()), input);
- 
--    // Perform All-Gather on dim 0.
-+    // Perform All-Gather on `firstGatheringDim`.
-     Attribute replicaGroups;
--    std::tie(replicaGroups, tmpShape[0]) = getReplicaGroupsAndSize(
--        allGatheringAxes, mesh, meshOrRef, enableRGV3, rewriter);
-+    std::tie(replicaGroups, tmpShape[firstGatheringDim]) =
-+        getReplicaGroupsAndSize(allGatheringAxes, mesh, meshOrRef, enableRGV3,
-+                                rewriter);
-     auto channelHandle = stablehlo::ChannelHandleAttr::get(
-         ctx, /*handle=*/conversionState.getNextChannelId(),
-         kCrossPartitionChannelHandleType);
-@@ -548,13 +586,15 @@ class AllGatherOpPattern : public OpConversionPattern<AllGatherOp> {
-     auto allGather = stablehlo::AllGatherOp::create(
-         rewriter, op.getLoc(),
-         TypeRange{RankedTensorType::get(tmpShape, localType.getElementType())},
--        curInput, /*all_gather_dim=*/0, replicaGroups, channelHandle,
--        /*use_global_device_ids=*/true);
-+        curInput, /*all_gather_dim=*/firstGatheringDim, replicaGroups,
-+        channelHandle, /*use_global_device_ids=*/true);
-     curInput = allGather.getResult(0);
- 
--    // Split the leading dimension into individual gathering factors.
--    SmallVector<int64_t> splitShape = gatheringFactors;
--    llvm::append_range(splitShape, localShape);
-+    // Split `firstGatheringDim` into individual gathering factors.
-+    SmallVector<int64_t> splitShape =
-+        llvm::to_vector(localShape.take_front(firstGatheringDim));
-+    llvm::append_range(splitShape, gatheringFactors);
-+    llvm::append_range(splitShape, localShape.drop_front(firstGatheringDim));
-     curInput = stablehlo::ReshapeOp::create(
-         rewriter, op.getLoc(),
-         RankedTensorType::get(splitShape, localType.getElementType()),
-@@ -562,14 +602,16 @@ class AllGatherOpPattern : public OpConversionPattern<AllGatherOp> {
- 
-     // Transpose factors next to their corresponding dimensions.
-     SmallVector<int64_t> xposePerm;
--    int64_t nextFactorIdx = 0;
--    for (int64_t i = 0; i < rank; ++i) {
-+    xposePerm.reserve(rank + gatheringDims.size());
-+    for (int64_t i = 0; i < firstGatheringDim; ++i) {
-+      xposePerm.push_back(i);
-+    }
-+    int64_t nextFactorIdx = firstGatheringDim;
-+    for (int64_t i = firstGatheringDim; i < rank; ++i) {
-       if (llvm::is_contained(gatheringDims, i)) {
-         xposePerm.push_back(nextFactorIdx++);
--        xposePerm.push_back(i + gatheringDims.size());
--      } else {
--        xposePerm.push_back(i + gatheringDims.size());
-       }
-+      xposePerm.push_back(i + gatheringDims.size());
-     }
-     curInput = stablehlo::TransposeOp::create(
-         rewriter, op.getLoc(), curInput,
-@@ -708,6 +750,24 @@ class AllReduceOpPattern : public OpConversionPattern<sdy::AllReduceOp> {
-   bool enableRGV3;
- };
- 
-+// Looks up a per-device scalar value from a 1D constant `tableAttr` indexed by
-+// `deviceId` (or `getDeviceId` if `deviceId` is null).
-+Value lookupPerDeviceValue(Location loc, DenseElementsAttr tableAttr,
-+                           const ConversionState& conversionState,
-+                           ConversionPatternRewriter& rewriter,
-+                           Value deviceId = nullptr) {
-+  if (!deviceId) {
-+    deviceId = getDeviceId(loc, conversionState, rewriter);
-+  }
-+  Type elemTy = tableAttr.getElementType();
-+  auto tableConst = stablehlo::ConstantOp::create(rewriter, loc, tableAttr);
-+  auto slice = stablehlo::DynamicSliceOp::create(
-+      rewriter, loc, RankedTensorType::get({1}, elemTy), tableConst,
-+      ValueRange{deviceId}, rewriter.getDenseI64ArrayAttr({1}));
-+  return stablehlo::ReshapeOp::create(rewriter, loc,
-+                                      RankedTensorType::get({}, elemTy), slice);
-+}
-+
- // Returns a 0-rank i32 tensor containing the global offset for the given shard
- // axes and local shard size.
- Value getDimensionOffset(Location loc, MeshAttr mesh,
-@@ -716,9 +776,6 @@ Value getDimensionOffset(Location loc, MeshAttr mesh,
-                          ConversionPatternRewriter& rewriter) {
-   int64_t numDevices = mesh.getTotalSize();
-   Type i32Ty = rewriter.getI32Type();
--  auto indexTy = RankedTensorType::get({}, i32Ty);
--
--  Value deviceId = getDeviceId(loc, conversionState, rewriter);
- 
-   // Calculate a compile-time offset table for this dimension.
-   SmallVector<int32_t> offsetsTable = llvm::map_to_vector(
-@@ -728,16 +785,11 @@ Value getDimensionOffset(Location loc, MeshAttr mesh,
-         return static_cast<int32_t>(offset);
-       });
- 
--  // Create the offset table and look up the value for the current device.
--  auto tableConst = stablehlo::ConstantOp::create(
--      rewriter, loc,
-+  return lookupPerDeviceValue(
-+      loc,
-       DenseIntElementsAttr::get(RankedTensorType::get({numDevices}, i32Ty),
--                                offsetsTable));
--  auto offsetSlice = stablehlo::DynamicSliceOp::create(
--      rewriter, loc, RankedTensorType::get({1}, i32Ty), tableConst,
--      ValueRange{deviceId}, rewriter.getDenseI64ArrayAttr({1}));
--
--  return stablehlo::ReshapeOp::create(rewriter, loc, indexTy, offsetSlice);
-+                                offsetsTable),
-+      conversionState, rewriter);
- }
- 
- Value emitDynamicSliceForAxes(Location loc, Value globalTensor, MeshAttr mesh,
-@@ -767,6 +819,43 @@ Value emitDynamicSliceForAxes(Location loc, Value globalTensor, MeshAttr mesh,
-       .getResult();
- }
- 
-+// Masks non-leader shards (where shard index across `axes` != 0) to
-+// `nonLeaderVal` (or zero if `nonLeaderVal` is null).
-+SmallVector<Value> maskNonLeaderShards(
-+    Location loc, ValueRange inputs, MeshAttr mesh, ArrayRef<AxisRefAttr> axes,
-+    const ConversionState& conversionState, ConversionPatternRewriter& rewriter,
-+    Value nonLeaderVal = nullptr, Value deviceId = nullptr) {
-+  int64_t numDevices = mesh.getTotalSize();
-+  SmallVector<bool> leaderTable = llvm::map_to_vector(
-+      llvm::seq<int64_t>(0, numDevices),
-+      [&](int64_t devId) { return getShardIndex(devId, mesh, axes) == 0; });
-+  Value isLeader = lookupPerDeviceValue(
-+      loc,
-+      DenseElementsAttr::get(
-+          RankedTensorType::get({numDevices}, rewriter.getI1Type()),
-+          leaderTable),
-+      conversionState, rewriter, deviceId);
-+  if (!nonLeaderVal) {
-+    auto localType = cast<RankedTensorType>(inputs.front().getType());
-+    nonLeaderVal = createZeroConstant(rewriter, loc, localType);
-+  }
-+  return llvm::map_to_vector(inputs, [&](Value input) -> Value {
-+    return stablehlo::SelectOp::create(rewriter, loc, isLeader, input,
-+                                       nonLeaderVal);
-+  });
-+}
-+
-+// Masks non-leader shards (where shard index across `axes` != 0) to zero.
-+Value maskNonLeaderShardsToZero(Location loc, Value input, MeshAttr mesh,
-+                                ArrayRef<AxisRefAttr> axes,
-+                                const ConversionState& conversionState,
-+                                ConversionPatternRewriter& rewriter,
-+                                Value deviceId = nullptr) {
-+  return maskNonLeaderShards(loc, input, mesh, axes, conversionState, rewriter,
-+                             /*nonLeaderVal=*/nullptr, deviceId)
-+      .front();
-+}
-+
- class AllSliceOpPattern : public OpConversionPattern<AllSliceOp> {
-  public:
-   AllSliceOpPattern(TypeConverter& converter, MLIRContext* ctx,
-@@ -1119,7 +1208,14 @@ class ConstantOpPattern : public OpConversionPattern<sdy::ConstantOp> {
-     auto localType = cast<RankedTensorType>(converter->convertType(op));
-     ElementsAttr elementsAttr = op.getValue();
-     Location loc = op.getLoc();
-+    const SymbolTable& symbolTable = converter->getSymbolTable();
-+    TensorShardingAttr sharding = getSharding(op.getResult());
-+    MeshAttr mesh = sharding ? sharding.getMesh(symbolTable) : nullptr;
-+    if (sharding && !mesh) {
-+      return op.emitOpError("failed to resolve mesh for constant");
-+    }
- 
-+    Value result;
-     // Unsharded or splat constants.
-     if (localType == globalType || elementsAttr.isSplat()) {
-       Attribute newAttr =
-@@ -1128,33 +1224,35 @@ class ConstantOpPattern : public OpConversionPattern<sdy::ConstantOp> {
-               : SplatElementsAttr::get(localType,
-                                        elementsAttr.getSplatValue<Attribute>());
- 
--      rewriter.replaceOp(
--          op, stablehlo::ConstantOp::create(rewriter, loc, localType,
--                                            cast<ElementsAttr>(newAttr)));
--      conversionState.removeToConvertOp(op);
--      return success();
--    }
-+      result = stablehlo::ConstantOp::create(rewriter, loc, localType,
-+                                             cast<ElementsAttr>(newAttr));
-+    } else {
-+      // Sharded dense constants.
-+      auto globalConst = stablehlo::ConstantOp::create(
-+          rewriter, loc, globalType, elementsAttr);
- 
--    // Sharded dense constants.
--    const SymbolTable& symbolTable = converter->getSymbolTable();
--    TensorShardingAttr sharding = getSharding(op.getResult());
--    MeshAttr mesh = sharding.getMesh(symbolTable);
--    if (!mesh) {
--      return op.emitOpError("failed to resolve mesh for constant");
-+      SmallVector<AxisRefListAttr> slicingAxesPerDim;
-+      slicingAxesPerDim.reserve(localType.getRank());
-+      for (DimensionShardingAttr dimSharding : sharding.getDimShardings()) {
-+        slicingAxesPerDim.push_back(
-+            AxisRefListAttr::get(op.getContext(), dimSharding.getAxes()));
-+      }
-+
-+      result =
-+          emitDynamicSliceForAxes(loc, globalConst, mesh, slicingAxesPerDim,
-+                                  localType, conversionState, rewriter);
-     }
--    auto globalConst =
--        stablehlo::ConstantOp::create(rewriter, loc, globalType, elementsAttr);
- 
--    SmallVector<AxisRefListAttr> slicingAxesPerDim;
--    slicingAxesPerDim.reserve(localType.getRank());
--    for (DimensionShardingAttr dimSharding : sharding.getDimShardings()) {
--      slicingAxesPerDim.push_back(
--          AxisRefListAttr::get(op.getContext(), dimSharding.getAxes()));
-+    if (sharding && !sharding.getUnreducedAxes().empty()) {
-+      SDY_CHECK(sharding.getReductionOp() == ReductionOp::SUM)
-+          << "Only SUM reduction is supported for constants with unreduced "
-+             "axes.";
-+      result = maskNonLeaderShardsToZero(loc, result, mesh,
-+                                         sharding.getUnreducedAxes(),
-+                                         conversionState, rewriter);
-     }
- 
--    rewriter.replaceOp(
--        op, emitDynamicSliceForAxes(loc, globalConst, mesh, slicingAxesPerDim,
--                                    localType, conversionState, rewriter));
-+    rewriter.replaceOp(op, result);
-     conversionState.removeToConvertOp(op);
- 
-     return success();
-@@ -1425,6 +1523,97 @@ class ReduceScatterOpPattern : public OpConversionPattern<ReduceScatterOp> {
-   bool enableRGV3;
- };
- 
-+class ReplicatedToUnreducedOpPattern
-+    : public OpConversionPattern<ReplicatedToUnreducedOp> {
-+ public:
-+  ReplicatedToUnreducedOpPattern(TypeConverter& converter, MLIRContext* ctx,
-+                                 ConversionState& state)
-+      : OpConversionPattern<ReplicatedToUnreducedOp>(converter, ctx),
-+        conversionState(state) {}
-+
-+  LogicalResult matchAndRewrite(
-+      ReplicatedToUnreducedOp op, OpAdaptor adaptor,
-+      ConversionPatternRewriter& rewriter) const override {
-+    TensorShardingAttr outSharding = op.getOutSharding();
-+    SDY_CHECK(outSharding.getReductionOp() == ReductionOp::SUM)
-+        << "Only SUM reduction is supported for ReplicatedToUnreducedOp.";
-+    auto* converter =
-+        static_cast<const GlobalToLocalTypeConverter*>(getTypeConverter());
-+    MeshAttr mesh = outSharding.getMesh(converter->getSymbolTable());
-+    if (!mesh) {
-+      return op.emitOpError("failed to resolve mesh");
-+    }
-+
-+    Value result =
-+        maskNonLeaderShardsToZero(op.getLoc(), adaptor.getTensor(), mesh,
-+                                  op.getAxes(), conversionState, rewriter);
-+    rewriter.replaceOp(op, result);
-+    conversionState.removeToConvertOp(op);
-+    return success();
-+  }
-+
-+ private:
-+  ConversionState& conversionState;
-+};
-+
-+class ShardedToUnreducedOpPattern
-+    : public OpConversionPattern<ShardedToUnreducedOp> {
-+ public:
-+  ShardedToUnreducedOpPattern(TypeConverter& converter, MLIRContext* ctx,
-+                              ConversionState& state)
-+      : OpConversionPattern<ShardedToUnreducedOp>(converter, ctx),
-+        conversionState(state) {}
-+
-+  LogicalResult matchAndRewrite(
-+      ShardedToUnreducedOp op, OpAdaptor adaptor,
-+      ConversionPatternRewriter& rewriter) const override {
-+    Location loc = op.getLoc();
-+    TensorShardingAttr outSharding = op.getOutSharding();
-+    SDY_CHECK(outSharding.getReductionOp() == ReductionOp::SUM)
-+        << "Only SUM reduction is supported for ShardedToUnreducedOp.";
-+    auto* converter =
-+        static_cast<const GlobalToLocalTypeConverter*>(getTypeConverter());
-+    MeshAttr mesh = outSharding.getMesh(converter->getSymbolTable());
-+    if (!mesh) {
-+      return op.emitOpError("failed to resolve mesh");
-+    }
-+
-+    Value localInput = adaptor.getTensor();
-+    auto localInputType = cast<RankedTensorType>(localInput.getType());
-+    auto localResultType = cast<RankedTensorType>(converter->convertType(op));
-+    if (!localResultType.hasStaticShape()) {
-+      return op.emitOpError("dynamic shape not supported");
-+    }
-+    Value zeroTensor = createZeroConstant(rewriter, loc, localResultType);
-+
-+    auto indexTy = RankedTensorType::get({}, rewriter.getI32Type());
-+    Value zeroOffset;
-+    SmallVector<Value> startIndices;
-+    startIndices.reserve(localResultType.getRank());
-+    for (int64_t i = 0; i < localResultType.getRank(); ++i) {
-+      ArrayRef<AxisRefAttr> axes = op.getAxes()[i].getValue();
-+      if (axes.empty()) {
-+        if (!zeroOffset) {
-+          zeroOffset = createZeroConstant(rewriter, loc, indexTy);
-+        }
-+        startIndices.push_back(zeroOffset);
-+        continue;
-+      }
-+      startIndices.push_back(getDimensionOffset(loc, mesh, axes,
-+                                                localInputType.getDimSize(i),
-+                                                conversionState, rewriter));
-+    }
-+
-+    rewriter.replaceOpWithNewOp<stablehlo::DynamicUpdateSliceOp>(
-+        op, localResultType, zeroTensor, localInput, startIndices);
-+    conversionState.removeToConvertOp(op);
-+    return success();
-+  }
-+
-+ private:
-+  ConversionState& conversionState;
-+};
-+
- class NamedComputationOpPattern
-     : public OpConversionPattern<NamedComputationOp> {
-  public:
-@@ -1608,11 +1797,11 @@ class StablehloConvolutionOpPattern
-       addRhsReductionAxes(spatialDim);
-     }
- 
--    bool hasUnreducedAxes = true;
-+    bool introducesPartialResult = true;
-     if (!reductionAxes.empty()) {
-       sortAndMergeAxes(reductionAxes, mesh);
--      hasUnreducedAxes = hasUnreducedResultSharding(op, reductionAxes,
--                                                    "convolution", converter);
-+      introducesPartialResult = doesOpIntroducePartialResult(
-+          op, reductionAxes, "convolution", converter);
-     }
- 
-     auto getAxes = [&](TensorShardingAttr sharding, int64_t dim) {
-@@ -1643,7 +1832,7 @@ class StablehloConvolutionOpPattern
-         localBatchGroupCount, op.getPrecisionConfigAttr());
- 
-     ValueRange results = localConv->getResults();
--    if (!reductionAxes.empty() && !hasUnreducedAxes) {
-+    if (!reductionAxes.empty() && !introducesPartialResult) {
-       Attribute meshOrRef =
-           rhsSharding ? rhsSharding.getMeshOrRef() : lhsSharding.getMeshOrRef();
-       results = createAllReduceForUnreducedAxes(
-@@ -1727,8 +1916,8 @@ class StablehloDotOpPattern : public OpConversionPattern<stablehlo::DotOp> {
- 
-         // Verify that an AllReduce has been inserted by a previous pass unless
-         // the result is explicitly unreduced.
--        if (!hasUnreducedResultSharding(op, contractingAxes, "dot",
--                                        converter)) {
-+        if (!doesOpIntroducePartialResult(op, contractingAxes, "dot",
-+                                          converter)) {
-           Operation* localDot = createLocalGenericOp(op, adaptor.getOperands(),
-                                                      rewriter, converter);
-           ValueRange results = createAllReduceForUnreducedAxes(
-@@ -1804,8 +1993,8 @@ class StablehloDotGeneralOpPattern
- 
-         // Verify that an AllReduce has been inserted by a previous pass unless
-         // the result is explicitly unreduced.
--        if (!hasUnreducedResultSharding(op, reductionAxes, "dot_general",
--                                        converter)) {
-+        if (!doesOpIntroducePartialResult(op, reductionAxes, "dot_general",
-+                                          converter)) {
-           Operation* localDotGeneral = createLocalGenericOp(
-               op, adaptor.getOperands(), rewriter, converter);
-           ValueRange results = createAllReduceForUnreducedAxes(
-@@ -1828,6 +2017,71 @@ class StablehloDotGeneralOpPattern
-   bool enableRGV3;
- };
- 
-+class StablehloDynamicSliceOpPattern
-+    : public OpConversionPattern<stablehlo::DynamicSliceOp> {
-+ public:
-+  StablehloDynamicSliceOpPattern(TypeConverter& converter, MLIRContext* ctx,
-+                                 ConversionState& state)
-+      : OpConversionPattern<stablehlo::DynamicSliceOp>(converter, ctx),
-+        conversionState(state) {}
-+
-+  LogicalResult matchAndRewrite(
-+      stablehlo::DynamicSliceOp op, OpAdaptor adaptor,
-+      ConversionPatternRewriter& rewriter) const override {
-+    auto* converter =
-+        static_cast<const GlobalToLocalTypeConverter*>(getTypeConverter());
-+    TensorShardingAttr inSharding = converter->getSharding(op.getOperand());
-+    TensorShardingAttr outSharding = converter->getSharding(op.getResult());
-+
-+    if (isFullyReplicated(inSharding) && isFullyReplicated(outSharding)) {
-+      return localizeGenericOp(op, adaptor.getOperands(), rewriter, converter,
-+                               conversionState);
-+    }
-+
-+    auto globalInputType = op.getOperand().getType();
-+    auto localResultType =
-+        cast<RankedTensorType>(converter->convertType(op.getResult()));
-+    MeshAttr mesh =
-+        inSharding ? inSharding.getMesh(converter->getSymbolTable()) : nullptr;
-+
-+    for (int64_t i = 0; i < globalInputType.getRank(); ++i) {
-+      ArrayRef<AxisRefAttr> inAxes =
-+          inSharding ? inSharding.getDimShardings()[i].getAxes()
-+                     : ArrayRef<AxisRefAttr>();
-+      ArrayRef<AxisRefAttr> outAxes =
-+          outSharding ? outSharding.getDimShardings()[i].getAxes()
-+                      : ArrayRef<AxisRefAttr>();
-+
-+      if (inAxes != outAxes) {
-+        return op.emitOpError()
-+               << "dimension " << i << " has mismatched sharding axes ("
-+               << strippedAttrsString(inAxes) << " vs "
-+               << strippedAttrsString(outAxes)
-+               << "). This requires a reshard which should have been handled "
-+                  "by previous passes.";
-+      }
-+
-+      int64_t globalSliceSize = op.getSliceSizes()[i];
-+      if (!inAxes.empty() && globalSliceSize != globalInputType.getDimSize(i)) {
-+        if (!mesh || inSharding.getDimShardings()[i].getShardedSize(mesh) > 1) {
-+          return op.emitOpError() << "dimension " << i
-+                                  << " is sharded but the slice is not a full "
-+                                     "slice and requires device communication.";
-+        }
-+      }
-+    }
-+
-+    rewriter.replaceOpWithNewOp<stablehlo::DynamicSliceOp>(
-+        op, localResultType, adaptor.getOperand(), adaptor.getStartIndices(),
-+        localResultType.getShape());
-+    conversionState.removeToConvertOp(op);
-+    return success();
-+  }
-+
-+ private:
-+  ConversionState& conversionState;
-+};
-+
- Value createScalarInt(OpBuilder& b, Location loc, RankedTensorType type,
-                       int64_t val) {
-   return stablehlo::ConstantOp::create(
-@@ -2095,7 +2349,7 @@ class StablehloReduceOpPattern
-     }
- 
-     if (!reductionAxes.empty() &&
--        !hasUnreducedResultSharding(op, reductionAxes, "reduce", converter)) {
-+        !doesOpIntroducePartialResult(op, reductionAxes, "reduce", converter)) {
-       auto localReduceOp = cast<stablehlo::ReduceOp>(
-           createLocalGenericOp(op, adaptor.getOperands(), rewriter, converter));
-       if (op.getNumResults() == 1) {
-@@ -2555,8 +2809,8 @@ class StablehloGatherOpPattern
-         mask, bcastDims);
- 
-     TensorShardingAttr resSharding = getSharding(op.getResult());
--    bool hasUnreducedAxes =
--        hasUnreducedResultSharding(op, reductionAxes, "gather", converter);
-+    bool introducesPartialResult =
-+        doesOpIntroducePartialResult(op, reductionAxes, "gather", converter);
-     ReductionOp reductionOp =
-         resSharding ? resSharding.getReductionOp() : ReductionOp::SUM;
- 
-@@ -2573,7 +2827,7 @@ class StablehloGatherOpPattern
-     result = stablehlo::SelectOp::create(rewriter, loc, bcastMask, result,
-                                          identityVal);
- 
--    if (!hasUnreducedAxes) {
-+    if (!introducesPartialResult) {
-       result =
-           createAllReduceForUnreducedAxes(loc, result, reductionAxes, mesh,
-                                           operandSharding.getMeshOrRef(),
-@@ -2747,30 +3001,8 @@ class StablehloRngBitGeneratorOpPattern
-     // Mask non-leader shards (where shard index across dataAxes != 0) to zero
-     // and all-reduce across dataAxes so all partitions receive device 0's
-     // updated output_state.
--    SmallVector<bool> leaderTable = llvm::map_to_vector(
--        llvm::seq<int64_t>(0, numDevices), [&](int64_t devId) {
--          return getShardIndex(devId, mesh, dataAxes) == 0;
--        });
--    auto leaderTableConst = stablehlo::ConstantOp::create(
--        rewriter, loc,
--        DenseElementsAttr::get(
--            RankedTensorType::get({numDevices}, rewriter.getI1Type()),
--            leaderTable));
--    Value isLeader = stablehlo::DynamicSliceOp::create(
--        rewriter, loc, RankedTensorType::get({1}, rewriter.getI1Type()),
--        leaderTableConst, ValueRange{deviceId},
--        rewriter.getDenseI64ArrayAttr({1}));
--    isLeader = stablehlo::ReshapeOp::create(
--        rewriter, loc, RankedTensorType::get({}, rewriter.getI1Type()),
--        isLeader);
--    isLeader = stablehlo::BroadcastInDimOp::create(
--        rewriter, loc,
--        RankedTensorType::get(localStateType.getShape(), rewriter.getI1Type()),
--        isLeader, rewriter.getDenseI64ArrayAttr({}));
--
--    Value zeroState = createZeroConstant(rewriter, loc, localStateType);
--    Value maskedState = stablehlo::SelectOp::create(rewriter, loc, isLeader,
--                                                    finalState, zeroState);
-+    Value maskedState = maskNonLeaderShardsToZero(
-+        loc, finalState, mesh, dataAxes, conversionState, rewriter, deviceId);
- 
-     finalState =
-         createAllReduceForUnreducedAxes(loc, ValueRange{maskedState}, dataAxes,
-@@ -3121,35 +3353,13 @@ class StablehloScatterOpPattern
-       if (!identity) {
-         return op.emitError("failed to get scatter reduction identity.");
-       }
--      // Generate leader mask: (shardIndex(reductionAxes) == 0).
--      int64_t numDevices = mesh.getTotalSize();
--      SmallVector<bool> leaderTable = llvm::map_to_vector(
--          llvm::seq<int64_t>(0, numDevices), [&](int64_t devId) {
--            return getShardIndex(devId, mesh, reductionAxes) == 0;
--          });
--      Value partitionId = getDeviceId(loc, conversionState, rewriter);
--      auto tableConst = stablehlo::ConstantOp::create(
--          rewriter, loc,
--          DenseElementsAttr::get(
--              RankedTensorType::get({numDevices}, rewriter.getI1Type()),
--              leaderTable));
--      Value isLeader =
--          stablehlo::DynamicSliceOp::create(
--              rewriter, loc, RankedTensorType::get({1}, rewriter.getI1Type()),
--              tableConst, {partitionId}, rewriter.getDenseI64ArrayAttr({1}))
--              .getResult();
--
--      isLeader = stablehlo::ReshapeOp::create(
--          rewriter, loc, RankedTensorType::get({}, rewriter.getI1Type()),
--          isLeader);
-       // Broadcast the scalar identity to the shape of the current local input.
-       Value broadcastIdentity = stablehlo::BroadcastInDimOp::create(
-           rewriter, loc, localInputs.front().getType(), identity,
-           rewriter.getDenseI64ArrayAttr({}));
--      llvm::transform(localInputs, localInputs.begin(), [&](Value input) {
--        return stablehlo::SelectOp::create(rewriter, loc, isLeader, input,
--                                           broadcastIdentity);
--      });
-+      localInputs = maskNonLeaderShards(loc, localInputs, mesh, reductionAxes,
-+                                        conversionState, rewriter,
-+                                        broadcastIdentity);
-     }
- 
-     auto globalInputType = cast<RankedTensorType>(input.getType());
-@@ -3207,7 +3417,8 @@ class StablehloScatterOpPattern
- 
-     ValueRange results = localScatter.getResults();
-     if (!reductionAxes.empty() &&
--        !hasUnreducedResultSharding(op, reductionAxes, "scatter", converter)) {
-+        !doesOpIntroducePartialResult(op, reductionAxes, "scatter",
-+                                      converter)) {
-       Attribute meshOrRef = inputSharding ? inputSharding.getMeshOrRef()
-                                           : updateSharding.getMeshOrRef();
-       results = createAllReduceForUnreducedAxes(
-@@ -3503,7 +3714,8 @@ struct ConvertGlobalToLocalPass
-     patterns.add<AllSliceOpPattern, CollectivePermuteOpPattern,
-                  ConstantOpPattern, FuncOpSignaturePattern, GenericOpPattern,
-                  ManualComputationOpPattern, NamedComputationOpPattern,
--                 StablehloConcatenateOpPattern,
-+                 ReplicatedToUnreducedOpPattern, ShardedToUnreducedOpPattern,
-+                 StablehloConcatenateOpPattern, StablehloDynamicSliceOpPattern,
-                  StablehloDynamicUpdateSliceOpPattern, StablehloIotaOpPattern,
-                  StablehloPadOpPattern,
-                  StablehloWindowedOpPattern<stablehlo::ReduceWindowOp>,
-diff --git a/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc b/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc
-index 4cef1c10..7831b6ca 100644
---- a/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc
-+++ b/shardy/dialect/sdy/transforms/export/pad_for_divisibility.cc
-@@ -78,6 +78,20 @@ Value createConstant(OpBuilder& b, Location loc, Type elementType,
-   llvm_unreachable("invalid PaddingValueKind");
- }
- 
-+// Returns the PaddingValueKind of the given scalar constant, or std::nullopt if
-+// not a recognized padding value (e.g. 0 or 1).
-+std::optional<PaddingValueKind> getPaddingValueKind(Value paddingValue) {
-+  if (matchPattern(paddingValue, m_Zero()) ||
-+      matchPattern(paddingValue, m_AnyZeroFloat())) {
-+    return PaddingValueKind::kZero;
-+  }
-+  if (matchPattern(paddingValue, m_One()) ||
-+      matchPattern(paddingValue, m_OneFloat())) {
-+    return PaddingValueKind::kOne;
-+  }
-+  return std::nullopt;
-+}
-+
- // Helper methods for commonly created ops (stablehlo.slice and stablehlo.pad).
- 
- // Slices a tensor to `targetType` along dimensions that need trimming.
-@@ -946,8 +960,9 @@ class StablehloSliceOpPattern : public OpConversionPattern<stablehlo::SliceOp> {
- 
- class StablehloPadOpPattern : public OpConversionPattern<stablehlo::PadOp> {
-  public:
--  explicit StablehloPadOpPattern(TypeConverter& converter, MLIRContext* ctx)
--      : OpConversionPattern(converter, ctx) {}
-+  StablehloPadOpPattern(TypeConverter& converter, MLIRContext* ctx,
-+                        PaddingCache& cache)
-+      : OpConversionPattern(converter, ctx), cache(cache) {}
- 
-   LogicalResult matchAndRewrite(
-       stablehlo::PadOp op, OpAdaptor adaptor,
-@@ -1017,11 +1032,20 @@ class StablehloPadOpPattern : public OpConversionPattern<stablehlo::PadOp> {
-     // Copy sharding attribute to the new result.
-     setSharding(newPadOp.getResult(), getSharding(result));
- 
--    // TODO(b/537380378): We could parse the padding value from the original pad
--    // op to cache the padding kind and avoid downstream selects.
-+    if (std::optional<PaddingValueKind> paddingKind =
-+            getPaddingValueKind(adaptor.getOperands()[1])) {
-+      if (input.getType() == op.getOperand().getType() || highPaddingReduced ||
-+          cache.getPadding(input) == paddingKind) {
-+        cache.setPadding(newPadOp.getResult(), *paddingKind);
-+      }
-+    }
-+
-     rewriter.replaceOp(op, newPadOp.getResult());
-     return success();
-   }
-+
-+ private:
-+  PaddingCache& cache;
- };
- 
- class StablehloDotGeneralOpPattern
-@@ -1406,9 +1430,10 @@ struct PadForDivisibilityPass
-     patterns.add<FuncOpSignaturePattern>(typeConverter, &getContext(),
-                                          paddingCache);
-     patterns.add<GenericOpPattern, StablehloSliceOpPattern,
--                 StablehloGatherOpPattern>(typeConverter, &getContext());
--    patterns.add<StablehloPadOpPattern, StablehloReshapeOpPattern>(
-+                 StablehloGatherOpPattern, StablehloReshapeOpPattern>(
-         typeConverter, &getContext());
-+    patterns.add<StablehloPadOpPattern>(typeConverter, &getContext(),
-+                                        paddingCache);
-     // Sharing the padding cache reference across pattern instances is safe from
-     // data races because pattern application within a function is sequential.
-     // TODO(b/553579414): remove the pattern for stablehlo.ConstantOp.
-diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/generic_ops.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/generic_ops.mlir
-index c840e9f1..f11dcdbe 100644
---- a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/generic_ops.mlir
-+++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/generic_ops.mlir
-@@ -40,6 +40,17 @@ func.func @unknown_dialect_op(
-   // CHECK-NEXT:  return %[[ARG0]] : tensor<8xf32>
-   return %arg0 : tensor<16xf32>
- }
-+
-+// CHECK-LABEL: func.func @registered_non_stablehlo_dialect_op
-+// CHECK-SAME: (%[[ARG0:.*]]: tensor<2x8xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}, {}]>}) -> (tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}]>})
-+func.func @registered_non_stablehlo_dialect_op(
-+  %arg0: tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}, {}]>})
-+  -> (tensor<32xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}]>}) {
-+  // CHECK-NEXT: %[[COLLAPSE:.*]] = tensor.collapse_shape %[[ARG0]] {{\[\[}}0, 1]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_2, [{"x"}]>]>} : tensor<2x8xf32> into tensor<16xf32>
-+  %0 = tensor.collapse_shape %arg0 [[0, 1]] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_2, [{"x"}]>]>} : tensor<4x8xf32> into tensor<32xf32>
-+  // CHECK-NEXT: return %[[COLLAPSE]] : tensor<16xf32>
-+  return %0 : tensor<32xf32>
-+}
- // CHECK-LABEL: func.func @callee_sharded
- // CHECK-SAME: (%arg0: tensor<8xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}]>}) -> (tensor<8xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}]>})
- func.func @callee_sharded(%arg0: tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}]>}) -> (tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2, [{"x"}]>}) {
-diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_all_gather.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_all_gather.mlir
-index d79f9dae..7e6ae755 100644
---- a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_all_gather.mlir
-+++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_all_gather.mlir
-@@ -174,3 +174,27 @@ func.func @two_dims_merge_adjacent_sub_axes(%arg0 : tensor<8x16xf32> {sdy.shardi
-   return %0 : tensor<8x16xf32>
- }
- 
-+// CHECK-LABEL: func @two_dims_non_zero_first_dim
-+// CHECK-SAME: (%[[ARG0:.*]]: tensor<3x5x7x11xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4_2, [{}, {"y"}, {}, {"z"}]>})
-+// CHECK-SAME: -> (tensor<3x20x7x22xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4_2, [{}, {}, {}, {}]>})
-+func.func @two_dims_non_zero_first_dim(%arg0 : tensor<3x20x7x22xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4_2, [{}, {"y"}, {}, {"z"}]>})
-+  -> (tensor<3x20x7x22xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4_2, [{}, {}, {}, {}]>}){
-+
-+  // --- Combined Dimensions All-Gather Strategy (per-dim-all-gather=false) ---
-+  // COMBINED: %[[RESHAPE1:.*]] = stablehlo.reshape %[[ARG0]] : (tensor<3x5x7x11xf32>) -> tensor<3x1x5x7x11xf32>
-+  // COMBINED: %[[GATHER:.*]] = "stablehlo.all_gather"(%[[RESHAPE1]]) <{
-+  // COMBINED-SAME: all_gather_dim = 1 : i64,
-+  // COMBINED-SAME: channel_handle = #stablehlo.channel_handle<handle = 8, type = 1>,
-+  // COMBINED-V1-SAME{LITERAL}: replica_groups = dense<[[0, 1, 2, 3, 4, 5, 6, 7], [8, 9, 10, 11, 12, 13, 14, 15]]>
-+  // COMBINED-V3-SAME: replica_groups = #stablehlo.replica_group_mesh_axes<mesh = @mesh_2_4_2, axes = [#stablehlo.axis_ref<name = "y">, #stablehlo.axis_ref<name = "z">]>
-+  // COMBINED-SAME: }> : (tensor<3x1x5x7x11xf32>) -> tensor<3x8x5x7x11xf32>
-+  // COMBINED: %[[RESHAPE2:.*]] = stablehlo.reshape %[[GATHER]] : (tensor<3x8x5x7x11xf32>) -> tensor<3x4x2x5x7x11xf32>
-+  // COMBINED: %[[TRANSPOSE:.*]] = stablehlo.transpose %[[RESHAPE2]], dims = [0, 1, 3, 4, 2, 5] : (tensor<3x4x2x5x7x11xf32>) -> tensor<3x4x5x7x2x11xf32>
-+  // COMBINED: %[[RESULT:.*]] = stablehlo.reshape %[[TRANSPOSE]] : (tensor<3x4x5x7x2x11xf32>) -> tensor<3x20x7x22xf32>
-+
-+  %0 = sdy.all_gather[{}, {"y"}, {}, {"z"}] %arg0 out_sharding = <@mesh_2_4_2, [{}, {}, {}, {}]> : tensor<3x20x7x22xf32>
-+
-+  // COMBINED: return %[[RESULT]] : tensor<3x20x7x22xf32>
-+  return %0 : tensor<3x20x7x22xf32>
-+}
-+
-diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_constant.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_constant.mlir
-index 2e01e015..15a45782 100644
---- a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_constant.mlir
-+++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_constant.mlir
-@@ -45,3 +45,18 @@ func.func @unsharded_dense() -> (tensor<4x2xi32> {sdy.sharding = #sdy.sharding<@
-   // CHECK-NEXT: return %[[CST]] : tensor<4x2xi32>
-   return %0 : tensor<4x2xi32>
- }
-+
-+// CHECK-LABEL: func.func @unreduced_splat
-+// CHECK-SAME:    -> (tensor<4x4xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {}], unreduced={"x"}>})
-+func.func @unreduced_splat() -> (tensor<4x4xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {}], unreduced={"x"}>}) {
-+  // CHECK-NEXT: %[[CST:.*]] = stablehlo.constant dense<1.000000e+00> : tensor<4x4xf32>
-+  // CHECK-NEXT: %[[PID:.*]] = stablehlo.partition_id : tensor<ui32>
-+  // CHECK-NEXT: %[[TABLE:.*]] = stablehlo.constant dense<[true, true, true, true, false, false, false, false]> : tensor<8xi1>
-+  // CHECK-NEXT: %[[SLICE:.*]] = stablehlo.dynamic_slice %[[TABLE]], %[[PID]], sizes = [1] : (tensor<8xi1>, tensor<ui32>) -> tensor<1xi1>
-+  // CHECK-NEXT: %[[SCALAR:.*]] = stablehlo.reshape %[[SLICE]] : (tensor<1xi1>) -> tensor<i1>
-+  // CHECK-NEXT: %[[ZERO:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<4x4xf32>
-+  // CHECK-NEXT: %[[RES:.*]] = stablehlo.select %[[SCALAR]], %[[CST]], %[[ZERO]] : tensor<i1>, tensor<4x4xf32>
-+  %0 = sdy.constant {sdy.sharding = #sdy.sharding_per_value<[<@mesh_2_4, [{}, {}], unreduced={"x"}>]>} dense<1.0> : tensor<4x4xf32>
-+  // CHECK-NEXT: return %[[RES]] : tensor<4x4xf32>
-+  return %0 : tensor<4x4xf32>
-+}
-diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_replicated_to_unreduced.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_replicated_to_unreduced.mlir
-new file mode 100644
-index 00000000..045684b3
---- /dev/null
-+++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_replicated_to_unreduced.mlir
-@@ -0,0 +1,55 @@
-+// RUN: sdy_opt %s -sdy-convert-global-to-local | FileCheck %s
-+
-+// CHECK: sdy.mesh @mesh_2_4 = <["x"=2, "y"=4]>
-+sdy.mesh @mesh_2_4 = <["x"=2, "y"=4]>
-+
-+// CHECK-LABEL: func.func @partial_mesh_major_axis
-+// CHECK-SAME: (%[[ARG0:.*]]: tensor<4x4xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {"y"}]>})
-+// CHECK-SAME: -> (tensor<4x4xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {"y"}], unreduced={"x"}>})
-+func.func @partial_mesh_major_axis(
-+    %arg0: tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {"y"}]>})
-+    -> (tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {"y"}], unreduced={"x"}>}) {
-+  // CHECK-NEXT: %[[PID:.*]] = stablehlo.partition_id : tensor<ui32>
-+  // CHECK-NEXT: %[[TABLE:.*]] = stablehlo.constant dense<[true, true, true, true, false, false, false, false]> : tensor<8xi1>
-+  // CHECK-NEXT: %[[SLICE:.*]] = stablehlo.dynamic_slice %[[TABLE]], %[[PID]], sizes = [1] : (tensor<8xi1>, tensor<ui32>) -> tensor<1xi1>
-+  // CHECK-NEXT: %[[SCALAR:.*]] = stablehlo.reshape %[[SLICE]] : (tensor<1xi1>) -> tensor<i1>
-+  // CHECK-NEXT: %[[ZERO:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<4x4xf32>
-+  // CHECK-NEXT: %[[RES:.*]] = stablehlo.select %[[SCALAR]], %[[ARG0]], %[[ZERO]] : tensor<i1>, tensor<4x4xf32>
-+  %0 = sdy.replicated_to_unreduced {"x"} %arg0 out_sharding=<@mesh_2_4, [{}, {"y"}], unreduced={"x"}> : tensor<4x16xf32>
-+  // CHECK-NEXT: return %[[RES]] : tensor<4x4xf32>
-+  return %0 : tensor<4x16xf32>
-+}
-+
-+// CHECK-LABEL: func.func @partial_mesh_minor_axis
-+// CHECK-SAME: (%[[ARG0:.*]]: tensor<2x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}]>})
-+// CHECK-SAME: -> (tensor<2x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>})
-+func.func @partial_mesh_minor_axis(
-+    %arg0: tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}]>})
-+    -> (tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>}) {
-+  // CHECK-NEXT: %[[PID:.*]] = stablehlo.partition_id : tensor<ui32>
-+  // CHECK-NEXT: %[[TABLE:.*]] = stablehlo.constant dense<[true, false, false, false, true, false, false, false]> : tensor<8xi1>
-+  // CHECK-NEXT: %[[SLICE:.*]] = stablehlo.dynamic_slice %[[TABLE]], %[[PID]], sizes = [1] : (tensor<8xi1>, tensor<ui32>) -> tensor<1xi1>
-+  // CHECK-NEXT: %[[SCALAR:.*]] = stablehlo.reshape %[[SLICE]] : (tensor<1xi1>) -> tensor<i1>
-+  // CHECK-NEXT: %[[ZERO:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<2x16xf32>
-+  // CHECK-NEXT: %[[RES:.*]] = stablehlo.select %[[SCALAR]], %[[ARG0]], %[[ZERO]] : tensor<i1>, tensor<2x16xf32>
-+  %0 = sdy.replicated_to_unreduced {"y"} %arg0 out_sharding=<@mesh_2_4, [{"x"}, {}], unreduced={"y"}> : tensor<4x16xf32>
-+  // CHECK-NEXT: return %[[RES]] : tensor<2x16xf32>
-+  return %0 : tensor<4x16xf32>
-+}
-+
-+// CHECK-LABEL: func.func @full_mesh_axes
-+// CHECK-SAME: (%[[ARG0:.*]]: tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {}]>})
-+// CHECK-SAME: -> (tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {}], unreduced={"x", "y"}>})
-+func.func @full_mesh_axes(
-+    %arg0: tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {}]>})
-+    -> (tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {}], unreduced={"x", "y"}>}) {
-+  // CHECK-NEXT: %[[PID:.*]] = stablehlo.partition_id : tensor<ui32>
-+  // CHECK-NEXT: %[[TABLE:.*]] = stablehlo.constant dense<[true, false, false, false, false, false, false, false]> : tensor<8xi1>
-+  // CHECK-NEXT: %[[SLICE:.*]] = stablehlo.dynamic_slice %[[TABLE]], %[[PID]], sizes = [1] : (tensor<8xi1>, tensor<ui32>) -> tensor<1xi1>
-+  // CHECK-NEXT: %[[SCALAR:.*]] = stablehlo.reshape %[[SLICE]] : (tensor<1xi1>) -> tensor<i1>
-+  // CHECK-NEXT: %[[ZERO:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<4x16xf32>
-+  // CHECK-NEXT: %[[RES:.*]] = stablehlo.select %[[SCALAR]], %[[ARG0]], %[[ZERO]] : tensor<i1>, tensor<4x16xf32>
-+  %0 = sdy.replicated_to_unreduced {"x", "y"} %arg0 out_sharding=<@mesh_2_4, [{}, {}], unreduced={"x", "y"}> : tensor<4x16xf32>
-+  // CHECK-NEXT: return %[[RES]] : tensor<4x16xf32>
-+  return %0 : tensor<4x16xf32>
-+}
-diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_sharded_to_unreduced.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_sharded_to_unreduced.mlir
-new file mode 100644
-index 00000000..eaf6ba99
---- /dev/null
-+++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/sdy_sharded_to_unreduced.mlir
-@@ -0,0 +1,22 @@
-+// RUN: sdy_opt %s -sdy-convert-global-to-local | FileCheck %s
-+
-+// CHECK: sdy.mesh @mesh_2_4 = <["x"=2, "y"=4]>
-+sdy.mesh @mesh_2_4 = <["x"=2, "y"=4]>
-+
-+// CHECK-LABEL: func.func @sharded_to_unreduced_single_dim
-+// CHECK-SAME: (%[[ARG0:.*]]: tensor<4x4xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {"y"}]>})
-+// CHECK-SAME: -> (tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>})
-+func.func @sharded_to_unreduced_single_dim(
-+    %arg0: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {"y"}]>})
-+    -> (tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>}) {
-+  // CHECK-NEXT: %[[ZERO:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<4x16xf32>
-+  // CHECK-NEXT: %[[OFF0:.*]] = stablehlo.constant dense<0> : tensor<i32>
-+  // CHECK-NEXT: %[[PID:.*]] = stablehlo.partition_id : tensor<ui32>
-+  // CHECK-NEXT: %[[TABLE:.*]] = stablehlo.constant dense<[0, 4, 8, 12, 0, 4, 8, 12]> : tensor<8xi32>
-+  // CHECK-NEXT: %[[DS:.*]] = stablehlo.dynamic_slice %[[TABLE]], %[[PID]], sizes = [1] : (tensor<8xi32>, tensor<ui32>) -> tensor<1xi32>
-+  // CHECK-NEXT: %[[OFF1:.*]] = stablehlo.reshape %[[DS]] : (tensor<1xi32>) -> tensor<i32>
-+  // CHECK-NEXT: %[[RES:.*]] = stablehlo.dynamic_update_slice %[[ZERO]], %[[ARG0]], %[[OFF0]], %[[OFF1]] : (tensor<4x16xf32>, tensor<4x4xf32>, tensor<i32>, tensor<i32>) -> tensor<4x16xf32>
-+  %0 = sdy.sharded_to_unreduced [{}, {"y"}] %arg0 out_sharding=<@mesh_2_4, [{"x"}, {}], unreduced={"y"}> : tensor<8x16xf32>
-+  // CHECK-NEXT: return %[[RES]] : tensor<4x16xf32>
-+  return %0 : tensor<8x16xf32>
-+}
-diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_dynamic_slice.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_dynamic_slice.mlir
-new file mode 100644
-index 00000000..a5aa58eb
---- /dev/null
-+++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_dynamic_slice.mlir
-@@ -0,0 +1,69 @@
-+// RUN: sdy_opt %s -sdy-convert-global-to-local -split-input-file -verify-diagnostics | FileCheck %s
-+
-+sdy.mesh @mesh_4_2 = <["x"=4, "y"=2]>
-+sdy.mesh @mesh_1_2 = <["x"=1, "y"=2]>
-+
-+// CHECK-LABEL: func @replicated
-+// CHECK-SAME:    (%[[ARG0:.*]]: tensor<16x8xf32>, %[[IDX0:.*]]: tensor<i32>, %[[IDX1:.*]]: tensor<i32>) -> tensor<4x8xf32>
-+func.func @replicated(%arg0: tensor<16x8xf32>, %arg1: tensor<i32>, %arg2: tensor<i32>) -> tensor<4x8xf32> {
-+  // CHECK-NEXT: %[[RES:.*]] = stablehlo.dynamic_slice %[[ARG0]], %[[IDX0]], %[[IDX1]], sizes = [4, 8] : (tensor<16x8xf32>, tensor<i32>, tensor<i32>) -> tensor<4x8xf32>
-+  %0 = stablehlo.dynamic_slice %arg0, %arg1, %arg2, sizes = [4, 8] : (tensor<16x8xf32>, tensor<i32>, tensor<i32>) -> tensor<4x8xf32>
-+  // CHECK-NEXT: return %[[RES]] : tensor<4x8xf32>
-+  return %0 : tensor<4x8xf32>
-+}
-+
-+// CHECK-LABEL: func @slicing_dim_not_sharded
-+// CHECK-SAME:    (%[[ARG0:.*]]: tensor<16x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{}, {"x"}]>},
-+// CHECK-SAME:     %[[IDX0:.*]]: tensor<i32>, %[[IDX1:.*]]: tensor<i32>)
-+// CHECK-SAME:    -> (tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{}, {"x"}]>})
-+func.func @slicing_dim_not_sharded(
-+    %arg0: tensor<16x32xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{}, {"x"}]>},
-+    %arg1: tensor<i32>, %arg2: tensor<i32>)
-+    -> (tensor<4x32xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{}, {"x"}]>}) {
-+  // CHECK-NEXT: %[[RES:.*]] = stablehlo.dynamic_slice %[[ARG0]], %[[IDX0]], %[[IDX1]], sizes = [4, 8] : (tensor<16x8xf32>, tensor<i32>, tensor<i32>) -> tensor<4x8xf32>
-+  %0 = stablehlo.dynamic_slice %arg0, %arg1, %arg2, sizes = [4, 32] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4_2, [{}, {"x"}]>]>} : (tensor<16x32xf32>, tensor<i32>, tensor<i32>) -> tensor<4x32xf32>
-+  // CHECK-NEXT: return %[[RES]] : tensor<4x8xf32>
-+  return %0 : tensor<4x32xf32>
-+}
-+
-+// CHECK-LABEL: func @slicing_dim_sharded_on_size_one_axis
-+// CHECK-SAME:    (%[[ARG0:.*]]: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1_2, [{"x"}, {"y"}]>},
-+// CHECK-SAME:     %[[IDX0:.*]]: tensor<i32>, %[[IDX1:.*]]: tensor<i32>)
-+// CHECK-SAME:    -> (tensor<4x8xf32> {sdy.sharding = #sdy.sharding<@mesh_1_2, [{"x"}, {"y"}]>})
-+func.func @slicing_dim_sharded_on_size_one_axis(
-+    %arg0: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_1_2, [{"x"}, {"y"}]>},
-+    %arg1: tensor<i32>, %arg2: tensor<i32>)
-+    -> (tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_1_2, [{"x"}, {"y"}]>}) {
-+  // CHECK-NEXT: %[[RES:.*]] = stablehlo.dynamic_slice %[[ARG0]], %[[IDX0]], %[[IDX1]], sizes = [4, 8] : (tensor<8x8xf32>, tensor<i32>, tensor<i32>) -> tensor<4x8xf32>
-+  %0 = stablehlo.dynamic_slice %arg0, %arg1, %arg2, sizes = [4, 16] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_1_2, [{"x"}, {"y"}]>]>} : (tensor<8x16xf32>, tensor<i32>, tensor<i32>) -> tensor<4x16xf32>
-+  // CHECK-NEXT: return %[[RES]] : tensor<4x8xf32>
-+  return %0 : tensor<4x16xf32>
-+}
-+
-+// -----
-+
-+sdy.mesh @mesh_4_2 = <["x"=4, "y"=2]>
-+
-+func.func @mismatched_sharding_axes(
-+    %arg0: tensor<16x32xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{}, {"x"}]>},
-+    %arg1: tensor<i32>, %arg2: tensor<i32>)
-+    -> (tensor<4x32xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{}, {"y"}]>}) {
-+  // expected-error @+2 {{'stablehlo.dynamic_slice' op dimension 1 has mismatched sharding axes ([axis_ref"x"] vs [axis_ref"y"]). This requires a reshard which should have been handled by previous passes.}}
-+  // expected-error @+1 {{failed to legalize operation 'stablehlo.dynamic_slice' that was explicitly marked illegal}}
-+  %0 = stablehlo.dynamic_slice %arg0, %arg1, %arg2, sizes = [4, 32] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4_2, [{}, {"y"}]>]>} : (tensor<16x32xf32>, tensor<i32>, tensor<i32>) -> tensor<4x32xf32>
-+  return %0 : tensor<4x32xf32>
-+}
-+
-+// -----
-+
-+sdy.mesh @mesh_4_2 = <["x"=4, "y"=2]>
-+
-+func.func @slicing_dim_sharded_non_full_slice(
-+    %arg0: tensor<16x32xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{"x"}, {}]>},
-+    %arg1: tensor<i32>, %arg2: tensor<i32>)
-+    -> (tensor<4x32xf32> {sdy.sharding = #sdy.sharding<@mesh_4_2, [{"x"}, {}]>}) {
-+  // expected-error @+2 {{'stablehlo.dynamic_slice' op dimension 0 is sharded but the slice is not a full slice and requires device communication.}}
-+  // expected-error @+1 {{failed to legalize operation 'stablehlo.dynamic_slice' that was explicitly marked illegal}}
-+  %0 = stablehlo.dynamic_slice %arg0, %arg1, %arg2, sizes = [4, 32] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4_2, [{"x"}, {}]>]>} : (tensor<16x32xf32>, tensor<i32>, tensor<i32>) -> tensor<4x32xf32>
-+  return %0 : tensor<4x32xf32>
-+}
-diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_reduce.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_reduce.mlir
-index 527e6631..4c162c29 100644
---- a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_reduce.mlir
-+++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_reduce.mlir
-@@ -86,3 +86,77 @@ func.func @stablehlo_reduce_multi_result_sharded(
-     }
-   return %0#0, %0#1 : tensor<8xf32>, tensor<8xi32>
- }
-+
-+// CHECK-LABEL: func.func @stablehlo_reduce_unreduced_operand_and_sharded_reduction_dim
-+// CHECK-SAME: (%[[ARG0:.*]]: tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>})
-+// CHECK-SAME: -> (tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}], unreduced={"x", "y"}>})
-+func.func @stablehlo_reduce_unreduced_operand_and_sharded_reduction_dim(
-+    %arg0: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>})
-+    -> (tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}], unreduced={"x", "y"}>}) {
-+  // CHECK-NEXT: %[[INIT:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
-+  // CHECK-NEXT: %[[LOCAL_RED:.*]] = stablehlo.reduce(%[[ARG0]] init: %[[INIT]]) applies stablehlo.add across dimensions = [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_2_4, [{}], unreduced={"x", "y"}>]>} : (tensor<4x16xf32>, tensor<f32>) -> tensor<16xf32>
-+  // CHECK-NEXT: return %[[LOCAL_RED]] : tensor<16xf32>
-+  %init = stablehlo.constant dense<0.0> : tensor<f32>
-+  %0 = stablehlo.reduce(%arg0 init: %init) across dimensions = [0]
-+      {sdy.sharding = #sdy.sharding_per_value<[<@mesh_2_4, [{}], unreduced={"x", "y"}>]>}
-+      : (tensor<8x16xf32>, tensor<f32>) -> tensor<16xf32>
-+    reducer(%lhs: tensor<f32>, %rhs: tensor<f32>) {
-+      %add = stablehlo.add %lhs, %rhs : tensor<f32>
-+      stablehlo.return %add : tensor<f32>
-+    }
-+  return %0 : tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: func.func @stablehlo_reduce_multiple_unreduced_operands_and_sharded_reduction_dim
-+// CHECK-SAME: (%[[ARG0:.*]]: tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>},
-+// CHECK-SAME:  %[[ARG1:.*]]: tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>})
-+// CHECK-SAME: -> (tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}], unreduced={"x", "y"}>},
-+// CHECK-SAME:     tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}], unreduced={"x", "y"}>})
-+func.func @stablehlo_reduce_multiple_unreduced_operands_and_sharded_reduction_dim(
-+    %arg0: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>},
-+    %arg1: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>})
-+    -> (tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}], unreduced={"x", "y"}>},
-+        tensor<16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}], unreduced={"x", "y"}>}) {
-+  // CHECK-NEXT: %[[INIT:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
-+  // CHECK-NEXT: %[[LOCAL_RED:.*]]:2 = stablehlo.reduce(%[[ARG0]] init: %[[INIT]]), (%[[ARG1]] init: %[[INIT]]) across dimensions = [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_2_4, [{}], unreduced={"x", "y"}>, <@mesh_2_4, [{}], unreduced={"x", "y"}>]>} : (tensor<4x16xf32>, tensor<4x16xf32>, tensor<f32>, tensor<f32>) -> (tensor<16xf32>, tensor<16xf32>)
-+  // CHECK:      return %[[LOCAL_RED]]#0, %[[LOCAL_RED]]#1 : tensor<16xf32>, tensor<16xf32>
-+  %init = stablehlo.constant dense<0.0> : tensor<f32>
-+  %0:2 = stablehlo.reduce(%arg0 init: %init), (%arg1 init: %init) across dimensions = [0]
-+      {sdy.sharding = #sdy.sharding_per_value<[<@mesh_2_4, [{}], unreduced={"x", "y"}>, <@mesh_2_4, [{}], unreduced={"x", "y"}>]>}
-+      : (tensor<8x16xf32>, tensor<8x16xf32>, tensor<f32>, tensor<f32>) -> (tensor<16xf32>, tensor<16xf32>)
-+    reducer(%lhs0: tensor<f32>, %rhs0: tensor<f32>) (%lhs1: tensor<f32>, %rhs1: tensor<f32>) {
-+      %add0 = stablehlo.add %lhs0, %rhs0 : tensor<f32>
-+      %add1 = stablehlo.add %lhs1, %rhs1 : tensor<f32>
-+      stablehlo.return %add0, %add1 : tensor<f32>, tensor<f32>
-+    }
-+  return %0#0, %0#1 : tensor<16xf32>, tensor<16xf32>
-+}
-+
-+// CHECK-LABEL: func.func @stablehlo_reduce_unreduced_operand_reduces_newly_introduced_axes
-+// CHECK-SAME: (%[[ARG0:.*]]: tensor<4x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>})
-+// CHECK-SAME: -> tensor<16xf32>
-+func.func @stablehlo_reduce_unreduced_operand_reduces_newly_introduced_axes(
-+    %arg0: tensor<8x16xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}], unreduced={"y"}>})
-+    -> tensor<16xf32> {
-+  // CHECK-NEXT: %[[INIT:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
-+  // CHECK-NEXT: %[[LOCAL_RED:.*]] = stablehlo.reduce(%[[ARG0]] init: %[[INIT]]) applies stablehlo.add across dimensions = [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_2_4, [{}], unreduced={"y"}>]>} : (tensor<4x16xf32>, tensor<f32>) -> tensor<16xf32>
-+  // CHECK:      %[[ALL_RED_X:.*]] = "stablehlo.all_reduce"(%[[LOCAL_RED]])
-+  // CHECK-SAME:   channel_handle = #stablehlo.channel_handle<handle = 4, type = 1>
-+  // CHECK-SAME:   replica_groups = #stablehlo.replica_group_mesh_axes<mesh = @mesh_2_4, axes = [#stablehlo.axis_ref<name = "x">]>
-+  // CHECK-SAME:   use_global_device_ids
-+  // CHECK:      %[[ALL_RED_Y:.*]] = "stablehlo.all_reduce"(%[[ALL_RED_X]])
-+  // CHECK-SAME:   channel_handle = #stablehlo.channel_handle<handle = 5, type = 1>
-+  // CHECK-SAME:   replica_groups = #stablehlo.replica_group_mesh_axes<mesh = @mesh_2_4, axes = [#stablehlo.axis_ref<name = "y">]>
-+  // CHECK-SAME:   use_global_device_ids
-+  // CHECK:      return %[[ALL_RED_Y]] : tensor<16xf32>
-+  %init = stablehlo.constant dense<0.0> : tensor<f32>
-+  %0 = stablehlo.reduce(%arg0 init: %init) across dimensions = [0]
-+      {sdy.sharding = #sdy.sharding_per_value<[<@mesh_2_4, [{}], unreduced={"y"}>]>}
-+      : (tensor<8x16xf32>, tensor<f32>) -> tensor<16xf32>
-+    reducer(%lhs: tensor<f32>, %rhs: tensor<f32>) {
-+      %add = stablehlo.add %lhs, %rhs : tensor<f32>
-+      stablehlo.return %add : tensor<f32>
-+    }
-+  %1 = sdy.all_reduce {"y"} %0 out_sharding=<@mesh_2_4, [{}]> : tensor<16xf32>
-+  return %1 : tensor<16xf32>
-+}
-diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_rng_bit_generator.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_rng_bit_generator.mlir
-index 67ffc38f..54d2a403 100644
---- a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_rng_bit_generator.mlir
-+++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_rng_bit_generator.mlir
-@@ -32,10 +32,9 @@ func.func @sharded_1d_ui64_state(
-   // UNSAFE-NEXT: %[[LOCAL_STATE:.*]], %[[LOCAL_DATA:.*]] = stablehlo.rng_bit_generator %[[ADJ_STATE]], algorithm = DEFAULT : (tensor<2xui64>) -> (tensor<2xui64>, tensor<2x16xf32>)
-   // UNSAFE-NEXT: %[[LEADER_TABLE:.*]] = stablehlo.constant dense<[true, false, false, false]> : tensor<4xi1>
-   // UNSAFE-NEXT: %[[IS_LEADER_SLICE:.*]] = stablehlo.dynamic_slice %[[LEADER_TABLE]], %[[PID]], sizes = [1] : (tensor<4xi1>, tensor<ui32>) -> tensor<1xi1>
--  // UNSAFE-NEXT: %[[IS_LEADER_SCALAR:.*]] = stablehlo.reshape %[[IS_LEADER_SLICE]] : (tensor<1xi1>) -> tensor<i1>
--  // UNSAFE-NEXT: %[[IS_LEADER:.*]] = stablehlo.broadcast_in_dim %[[IS_LEADER_SCALAR]], dims = [] : (tensor<i1>) -> tensor<2xi1>
-+  // UNSAFE-NEXT: %[[IS_LEADER:.*]] = stablehlo.reshape %[[IS_LEADER_SLICE]] : (tensor<1xi1>) -> tensor<i1>
-   // UNSAFE-NEXT: %[[ZERO_STATE:.*]] = stablehlo.constant dense<0> : tensor<2xui64>
--  // UNSAFE-NEXT: %[[MASKED_STATE:.*]] = stablehlo.select %[[IS_LEADER]], %[[LOCAL_STATE]], %[[ZERO_STATE]] : tensor<2xi1>, tensor<2xui64>
-+  // UNSAFE-NEXT: %[[MASKED_STATE:.*]] = stablehlo.select %[[IS_LEADER]], %[[LOCAL_STATE]], %[[ZERO_STATE]] : tensor<i1>, tensor<2xui64>
-   // UNSAFE-NEXT: %[[FINAL_STATE:.*]] = "stablehlo.all_reduce"(%[[MASKED_STATE]])
-   // UNSAFE-SAME{LITERAL}: <{channel_handle = #stablehlo.channel_handle<handle = 1, type = 1>, replica_groups = #stablehlo.replica_group_mesh_axes<mesh = @mesh_4, axes = [#stablehlo.axis_ref<name = "x">]>, use_global_device_ids}> ({
-   // UNSAFE-NEXT: ^bb0(%[[ARG0:.*]]: tensor<ui64>, %[[ARG1:.*]]: tensor<ui64>):
-@@ -76,10 +75,9 @@ func.func @sharded_partial_mesh_ui32_state(
-   // UNSAFE-NEXT: %[[LOCAL_STATE:.*]], %[[LOCAL_DATA:.*]] = stablehlo.rng_bit_generator %[[ADJ_STATE]], algorithm = THREE_FRY : (tensor<4xui32>) -> (tensor<4xui32>, tensor<8x8xf32>)
-   // UNSAFE-NEXT: %[[LEADER_TABLE:.*]] = stablehlo.constant dense<[true, false, true, false]> : tensor<4xi1>
-   // UNSAFE-NEXT: %[[IS_LEADER_SLICE:.*]] = stablehlo.dynamic_slice %[[LEADER_TABLE]], %[[PID]], sizes = [1] : (tensor<4xi1>, tensor<ui32>) -> tensor<1xi1>
--  // UNSAFE-NEXT: %[[IS_LEADER_SCALAR:.*]] = stablehlo.reshape %[[IS_LEADER_SLICE]] : (tensor<1xi1>) -> tensor<i1>
--  // UNSAFE-NEXT: %[[IS_LEADER:.*]] = stablehlo.broadcast_in_dim %[[IS_LEADER_SCALAR]], dims = [] : (tensor<i1>) -> tensor<4xi1>
-+  // UNSAFE-NEXT: %[[IS_LEADER:.*]] = stablehlo.reshape %[[IS_LEADER_SLICE]] : (tensor<1xi1>) -> tensor<i1>
-   // UNSAFE-NEXT: %[[ZERO_STATE:.*]] = stablehlo.constant dense<0> : tensor<4xui32>
--  // UNSAFE-NEXT: %[[MASKED_STATE:.*]] = stablehlo.select %[[IS_LEADER]], %[[LOCAL_STATE]], %[[ZERO_STATE]] : tensor<4xi1>, tensor<4xui32>
-+  // UNSAFE-NEXT: %[[MASKED_STATE:.*]] = stablehlo.select %[[IS_LEADER]], %[[LOCAL_STATE]], %[[ZERO_STATE]] : tensor<i1>, tensor<4xui32>
-   // UNSAFE-NEXT: %[[FINAL_STATE:.*]] = "stablehlo.all_reduce"(%[[MASKED_STATE]])
-   // UNSAFE-SAME{LITERAL}: <{channel_handle = #stablehlo.channel_handle<handle = 2, type = 1>, replica_groups = #stablehlo.replica_group_mesh_axes<mesh = @mesh_2_2, axes = [#stablehlo.axis_ref<name = "y">]>, use_global_device_ids}> ({
-   // UNSAFE-NEXT: ^bb0(%[[ARG0:.*]]: tensor<ui32>, %[[ARG1:.*]]: tensor<ui32>):
-diff --git a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_scatter.mlir b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_scatter.mlir
-index e45919c6..3661799c 100644
---- a/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_scatter.mlir
-+++ b/shardy/dialect/sdy/transforms/export/test/convert_global_to_local/stablehlo_scatter.mlir
-@@ -18,11 +18,11 @@ func.func @input_not_sharded_scatter_indices_update_sharded_on_implicit_batch_di
-     %arg2: tensor<2x3x1xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}, {}]>})
-  -> tensor<3x4x2xf32> {
-   // CHECK-DAG:  %[[CST:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
-+  // CHECK-DAG:  %[[ID_BCAST:.*]] = stablehlo.broadcast_in_dim %[[CST]], dims = []
-   // CHECK-DAG:  %[[PID:.*]] = stablehlo.partition_id : tensor<ui32>
-   // CHECK:      %[[TABLE:.*]] = stablehlo.constant dense<[true, true, true, true, false, false, false, false]> : tensor<8xi1>
-   // CHECK:      %[[IS_LEADER_SLICE:.*]] = stablehlo.dynamic_slice %[[TABLE]], %[[PID]], sizes = [1]
-   // CHECK:      %[[IS_LEADER:.*]] = stablehlo.reshape %[[IS_LEADER_SLICE]] : (tensor<1xi1>) -> tensor<i1>
--  // CHECK:      %[[ID_BCAST:.*]] = stablehlo.broadcast_in_dim %[[CST]], dims = []
-   // CHECK:      %[[INPUT_SEL:.*]] = stablehlo.select %[[IS_LEADER]], %[[ARG0]], %[[ID_BCAST]]
-   // CHECK:      %[[SCATTER:.*]] = "stablehlo.scatter"(%[[INPUT_SEL]], %[[ARG1]], %[[ARG2]])
-   // CHECK-SAME: scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [2], inserted_window_dims = [0, 1], scatter_dims_to_operand_dims = [0, 1], index_vector_dim = 2>
-@@ -63,11 +63,11 @@ func.func @input_scalar_scatter_indices_update_sharded_on_implicit_batch_dim(
-     %arg2: tensor<2xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}]>})
-  -> tensor<f32> {
-   // CHECK-DAG:  %[[CST:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
-+  // CHECK-DAG:  %[[ID_BCAST:.*]] = stablehlo.broadcast_in_dim %[[CST]], dims = []
-   // CHECK-DAG:  %[[PID:.*]] = stablehlo.partition_id : tensor<ui32>
-   // CHECK:      %[[TABLE:.*]] = stablehlo.constant dense<[true, true, true, true, false, false, false, false]> : tensor<8xi1>
-   // CHECK:      %[[IS_LEADER_SLICE:.*]] = stablehlo.dynamic_slice %[[TABLE]], %[[PID]], sizes = [1]
-   // CHECK:      %[[IS_LEADER:.*]] = stablehlo.reshape %[[IS_LEADER_SLICE]] : (tensor<1xi1>) -> tensor<i1>
--  // CHECK:      %[[ID_BCAST:.*]] = stablehlo.broadcast_in_dim %[[CST]], dims = []
-   // CHECK:      %[[INPUT_SEL:.*]] = stablehlo.select %[[IS_LEADER]], %[[ARG0]], %[[ID_BCAST]]
-   // CHECK:      %[[SCATTER:.*]] = "stablehlo.scatter"(%[[INPUT_SEL]], %[[ARG1]], %[[ARG2]])
-   // CHECK-SAME: scatter_dimension_numbers = #stablehlo.scatter<index_vector_dim = 1>
-@@ -113,11 +113,11 @@ func.func @input_sharded_not_on_indexed_dim(
-     %arg2: tensor<2x3x1xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{"x"}, {}, {}]>})
-  -> (tensor<3x4x2xf32> {sdy.sharding = #sdy.sharding<@mesh_2_4, [{}, {}, {"y":(2)2}]>}) {
-   // CHECK-DAG:  %[[CST:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
-+  // CHECK-DAG:  %[[ID_BCAST:.*]] = stablehlo.broadcast_in_dim %[[CST]], dims = []
-   // CHECK-DAG:  %[[PID:.*]] = stablehlo.partition_id : tensor<ui32>
-   // CHECK:      %[[TABLE:.*]] = stablehlo.constant dense<[true, true, true, true, false, false, false, false]> : tensor<8xi1>
-   // CHECK:      %[[IS_LEADER_SLICE:.*]] = stablehlo.dynamic_slice %[[TABLE]], %[[PID]], sizes = [1]
-   // CHECK:      %[[IS_LEADER:.*]] = stablehlo.reshape %[[IS_LEADER_SLICE]] : (tensor<1xi1>) -> tensor<i1>
--  // CHECK:      %[[ID_BCAST:.*]] = stablehlo.broadcast_in_dim %[[CST]], dims = []
-   // CHECK:      %[[INPUT_SEL:.*]] = stablehlo.select %[[IS_LEADER]], %[[ARG0]], %[[ID_BCAST]]
-   // CHECK:      %[[SCATTER:.*]] = "stablehlo.scatter"(%[[INPUT_SEL]], %[[ARG1]], %[[ARG2]])
-   // CHECK-SAME: scatter_dimension_numbers = #stablehlo.scatter<update_window_dims = [2], inserted_window_dims = [0, 1], scatter_dims_to_operand_dims = [0, 1], index_vector_dim = 2>
-diff --git a/shardy/dialect/sdy/transforms/export/test/executable_convert_global_to_local/sdy_all_gather.mlir b/shardy/dialect/sdy/transforms/export/test/executable_convert_global_to_local/sdy_all_gather.mlir
-index 2377e4a7..47c5801a 100644
---- a/shardy/dialect/sdy/transforms/export/test/executable_convert_global_to_local/sdy_all_gather.mlir
-+++ b/shardy/dialect/sdy/transforms/export/test/executable_convert_global_to_local/sdy_all_gather.mlir
-@@ -30,6 +30,14 @@ func.func @all_gather_test(
-   return %0 : tensor<4x4xi32>
- }
- 
-+// This function gathers non-zero dims 1 and 3 of a 2x4x2x2 tensor.
-+func.func @all_gather_test_non_zero_first_dim(
-+  %arg0: tensor<2x4x2x2xi32> {sdy.sharding = #sdy.sharding<@mesh_2_2, [{}, {"x"}, {}, {"y"}]>}
-+  ) -> tensor<2x4x2x2xi32> {
-+  %0 = sdy.all_gather [{}, {"x"}, {}, {"y"}] %arg0 out_sharding=<@mesh_2_2, [{}, {}, {}, {}]> : tensor<2x4x2x2xi32>
-+  return %0 : tensor<2x4x2x2xi32>
-+}
-+
- //--- part2.mlir
- 
- func.func @main() {
-@@ -63,5 +71,26 @@ func.func @main() {
-   "check.expect_eq"(%res_comb#0, %res_pdim#0) : (tensor<4x4xi32>, tensor<4x4xi32>) -> ()
-   "check.expect_eq"(%res_comb#0, %cst) : (tensor<4x4xi32>, tensor<4x4xi32>) -> ()
- 
-+  // Test gathering on non-zero dims 1 and 3 with non-gathering dims 0 and 2.
-+  %iota = stablehlo.iota dim = 0 : tensor<32xi32>
-+  %cst_4d = stablehlo.reshape %iota : (tensor<32xi32>) -> tensor<2x4x2x2xi32>
-+  %s0_4d = stablehlo.slice %cst_4d [0:2, 0:2, 0:2, 0:1] : (tensor<2x4x2x2xi32>) -> tensor<2x2x2x1xi32>
-+  %s1_4d = stablehlo.slice %cst_4d [0:2, 0:2, 0:2, 1:2] : (tensor<2x4x2x2xi32>) -> tensor<2x2x2x1xi32>
-+  %s2_4d = stablehlo.slice %cst_4d [0:2, 2:4, 0:2, 0:1] : (tensor<2x4x2x2xi32>) -> tensor<2x2x2x1xi32>
-+  %s3_4d = stablehlo.slice %cst_4d [0:2, 2:4, 0:2, 1:2] : (tensor<2x4x2x2xi32>) -> tensor<2x2x2x1xi32>
-+
-+  %res_comb_4d:4 = "interpreter.run_parallel"(%s0_4d, %s1_4d, %s2_4d, %s3_4d) {
-+    programs = [[@combined_gather_non_zero_first_dim, @combined_gather_non_zero_first_dim, @combined_gather_non_zero_first_dim, @combined_gather_non_zero_first_dim]]
-+  } : (tensor<2x2x2x1xi32>, tensor<2x2x2x1xi32>, tensor<2x2x2x1xi32>, tensor<2x2x2x1xi32>) ->
-+      (tensor<2x4x2x2xi32>, tensor<2x4x2x2xi32>, tensor<2x4x2x2xi32>, tensor<2x4x2x2xi32>)
-+
-+  %res_pdim_4d:4 = "interpreter.run_parallel"(%s0_4d, %s1_4d, %s2_4d, %s3_4d) {
-+    programs = [[@per_dim_gather_non_zero_first_dim, @per_dim_gather_non_zero_first_dim, @per_dim_gather_non_zero_first_dim, @per_dim_gather_non_zero_first_dim]]
-+  } : (tensor<2x2x2x1xi32>, tensor<2x2x2x1xi32>, tensor<2x2x2x1xi32>, tensor<2x2x2x1xi32>) ->
-+      (tensor<2x4x2x2xi32>, tensor<2x4x2x2xi32>, tensor<2x4x2x2xi32>, tensor<2x4x2x2xi32>)
-+
-+  "check.expect_eq"(%res_comb_4d#0, %res_pdim_4d#0) : (tensor<2x4x2x2xi32>, tensor<2x4x2x2xi32>) -> ()
-+  "check.expect_eq"(%res_comb_4d#0, %cst_4d) : (tensor<2x4x2x2xi32>, tensor<2x4x2x2xi32>) -> ()
-+
-   return
- }
-diff --git a/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/stablehlo_pad.mlir b/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/stablehlo_pad.mlir
-index 59cc6b86..208aa2c2 100644
---- a/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/stablehlo_pad.mlir
-+++ b/shardy/dialect/sdy/transforms/export/test/pad_for_divisibility/stablehlo_pad.mlir
-@@ -56,10 +56,22 @@ func.func @no_op(%arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4, [
-   return %1 : tensor<12x8xf32>
- }
- 
-+// The pad below is padded for divisibility along dim 0 (9 -> 12) with its own
-+// zero padding value, so the trailing rows 9-11 are already zero. The pad
-+// pattern registers this zero padding in the padding cache, so the dot_general
-+// pattern can contract over dim 0 without inserting an iota/compare/select mask
-+// on the lhs. Without caching, the lhs would be masked with a stablehlo.select.
- // CHECK-LABEL: func @pad_then_dot_general
- // CHECK-SAME:    %[[ARG0:.*]]: tensor<7x7xf32>, %[[ARG1:.*]]: tensor<9x5xf32>
--// CHECK:         stablehlo.select
--// CHECK:         stablehlo.dot_general
-+// CHECK-DAG:     %[[CST0:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
-+// CHECK-DAG:     %[[PAD0:.*]] = stablehlo.pad %[[ARG0]], %[[CST0]], low = [1, 0], high = [4, 0], interior = [0, 0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4, [{"x"}, {}]>]>} : (tensor<7x7xf32>, tensor<f32>) -> tensor<12x7xf32>
-+// CHECK-DAG:     %[[CST1:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
-+// CHECK-DAG:     %[[PAD1:.*]] = stablehlo.pad %[[ARG1]], %[[CST1]], low = [0, 0], high = [3, 0], interior = [0, 0] : (tensor<9x5xf32>, tensor<f32>) -> tensor<12x5xf32>
-+// CHECK-DAG:     %[[SLICE:.*]] = sdy.all_slice [{"x"}, {}] %[[PAD1]] out_sharding=<@mesh_4, [{"x"}, {}]> : tensor<12x5xf32>
-+// CHECK-NOT:     stablehlo.select
-+// The dot_general consumes %[[PAD0]] directly, i.e. no mask is applied.
-+// CHECK:         %[[DOT:.*]] = stablehlo.dot_general %[[PAD0]], %[[SLICE]], contracting_dims = [0] x [0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4, [{}, {}]>]>} : (tensor<12x7xf32>, tensor<12x5xf32>) -> tensor<7x5xf32>
-+// CHECK-NEXT:    return %[[DOT]] : tensor<7x5xf32>
- func.func @pad_then_dot_general(%arg0: tensor<7x7xf32>, %arg1: tensor<9x5xf32>) -> tensor<7x5xf32> {
-   %cst = stablehlo.constant dense<0.000000e+00> : tensor<f32>
-   %0 = stablehlo.pad %arg0, %cst, low = [1, 0], high = [1, 0], interior = [0, 0] {sdy.sharding = #sdy.sharding_per_value<[#sdy.sharding<@mesh_4, [{"x"}, {}]>]>} : (tensor<7x7xf32>, tensor<f32>) -> tensor<9x7xf32>
-@@ -68,6 +80,48 @@ func.func @pad_then_dot_general(%arg0: tensor<7x7xf32>, %arg1: tensor<9x5xf32>)
-   return %1 : tensor<7x5xf32>
- }
- 
-+// Same idea for stablehlo.reduce: the pad result is sharded along the
-+// reduction dim and padded for divisibility (7 -> 8) with zero. Zero is the
-+// identity of add, so the reduce consumes the padded tensor directly instead
-+// of masking column 7 with a stablehlo.select.
-+// CHECK-LABEL: func @zero_pad_then_reduce_add
-+// CHECK-SAME:    %[[ARG0:.*]]: tensor<4x5xf32>
-+// CHECK-NEXT:    %[[CST:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
-+// CHECK-NEXT:    %[[PAD:.*]] = stablehlo.pad %[[ARG0]], %[[CST]], low = [0, 0], high = [0, 3], interior = [0, 0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4, [{}, {"x"}]>]>} : (tensor<4x5xf32>, tensor<f32>) -> tensor<4x8xf32>
-+// CHECK-NEXT:    %[[REDUCE:.*]] = stablehlo.reduce(%[[PAD]] init: %[[CST]]) applies stablehlo.add across dimensions = [1]
-+// CHECK-SAME:    : (tensor<4x8xf32>, tensor<f32>) -> tensor<4xf32>
-+// CHECK-NEXT:    return %[[REDUCE]] : tensor<4xf32>
-+func.func @zero_pad_then_reduce_add(%arg0: tensor<4x5xf32>) -> tensor<4xf32> {
-+  %cst = stablehlo.constant dense<0.000000e+00> : tensor<f32>
-+  %0 = stablehlo.pad %arg0, %cst, low = [0, 0], high = [0, 2], interior = [0, 0] {sdy.sharding = #sdy.sharding_per_value<[#sdy.sharding<@mesh_4, [{}, {"x"}]>]>} : (tensor<4x5xf32>, tensor<f32>) -> tensor<4x7xf32>
-+  %1 = stablehlo.reduce(%0 init: %cst) applies stablehlo.add across dimensions = [1] {sdy.sharding = #sdy.sharding_per_value<[#sdy.sharding<@mesh_4, [{}]>]>} : (tensor<4x7xf32>, tensor<f32>) -> tensor<4xf32>
-+  return %1 : tensor<4xf32>
-+}
-+
-+// Negative case: the pad value is 1, so the padded column 7 holds 1, which is
-+// not the identity of add. The mask must stay: column 7 (iota >= 7) is replaced
-+// with the init value 0 before the reduce. The real padding in columns 5-6
-+// (iota < 7) is kept, since those 1s are part of the original tensor.
-+// CHECK-LABEL: func @one_pad_then_reduce_add
-+// CHECK-SAME:    %[[ARG0:.*]]: tensor<4x5xf32>
-+// CHECK-DAG:     %[[ONE:.*]] = stablehlo.constant dense<1.000000e+00> : tensor<f32>
-+// CHECK-DAG:     %[[ZERO:.*]] = stablehlo.constant dense<0.000000e+00> : tensor<f32>
-+// CHECK:         %[[PAD:.*]] = stablehlo.pad %[[ARG0]], %[[ONE]], low = [0, 0], high = [0, 3], interior = [0, 0] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4, [{}, {"x"}]>]>} : (tensor<4x5xf32>, tensor<f32>) -> tensor<4x8xf32>
-+// CHECK:         %[[IOTA:.*]] = stablehlo.iota dim = 1
-+// CHECK:         %[[LIMIT:.*]] = stablehlo.constant dense<7> : tensor<i32>
-+// CHECK:         %[[MASK:.*]] = stablehlo.compare  LT, %[[IOTA]], %{{.*}}
-+// CHECK:         %[[SELECT:.*]] = stablehlo.select %[[MASK]], %[[PAD]], %{{.*}}
-+// CHECK:         %[[REDUCE:.*]] = stablehlo.reduce(%[[SELECT]] init: %[[ZERO]]) applies stablehlo.add across dimensions = [1]
-+// CHECK-SAME:    : (tensor<4x8xf32>, tensor<f32>) -> tensor<4xf32>
-+// CHECK-NEXT:    return %[[REDUCE]] : tensor<4xf32>
-+func.func @one_pad_then_reduce_add(%arg0: tensor<4x5xf32>) -> tensor<4xf32> {
-+  %one = stablehlo.constant dense<1.000000e+00> : tensor<f32>
-+  %zero = stablehlo.constant dense<0.000000e+00> : tensor<f32>
-+  %0 = stablehlo.pad %arg0, %one, low = [0, 0], high = [0, 2], interior = [0, 0] {sdy.sharding = #sdy.sharding_per_value<[#sdy.sharding<@mesh_4, [{}, {"x"}]>]>} : (tensor<4x5xf32>, tensor<f32>) -> tensor<4x7xf32>
-+  %1 = stablehlo.reduce(%0 init: %zero) applies stablehlo.add across dimensions = [1] {sdy.sharding = #sdy.sharding_per_value<[#sdy.sharding<@mesh_4, [{}]>]>} : (tensor<4x7xf32>, tensor<f32>) -> tensor<4xf32>
-+  return %1 : tensor<4xf32>
-+}
-+
- // CHECK-LABEL: func @negative_high_padding(
- // CHECK-SAME:    %[[ARG0:.*]]: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh_4, [{"x"}, {}]>}) -> tensor<7x8xf32> {
- // CHECK-NEXT:    %[[SLICE:.*]] = stablehlo.slice %[[ARG0]] [0:8, 0:8] {sdy.sharding = #sdy.sharding_per_value<[<@mesh_4, [{"x"}, {}]>]>} : (tensor<8x8xf32>) -> tensor<8x8xf32>
-diff --git a/shardy/dialect/sdy/transforms/export/test/resolve_permutation_factors.mlir b/shardy/dialect/sdy/transforms/export/test/resolve_permutation_factors.mlir
-index 162bb3db..b1e2bd3d 100644
---- a/shardy/dialect/sdy/transforms/export/test/resolve_permutation_factors.mlir
-+++ b/shardy/dialect/sdy/transforms/export/test/resolve_permutation_factors.mlir
-@@ -1760,6 +1760,22 @@ func.func @dynamic_update_slice_dynamic_operand_shape(
-   return %0 : tensor<?x8xf32>
- }
- 
-+// CHECK-LABEL: func @dynamic_update_slice_missing_operand_sharding
-+// CHECK-SAME: (%[[ARG0:.*]]: tensor<16x8xf32>, %[[ARG1:.*]]: tensor<2x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>}, %[[ARG2:.*]]: tensor<i32>, %[[ARG3:.*]]: tensor<i32>)
-+// CHECK-NEXT: %[[DUS:.*]] = stablehlo.dynamic_update_slice %[[ARG0]], %[[ARG1]], %[[ARG2]], %[[ARG3]]
-+// CHECK-NEXT: return %[[DUS]]
-+func.func @dynamic_update_slice_missing_operand_sharding(
-+    %arg0: tensor<16x8xf32>,
-+    %arg1: tensor<2x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>},
-+    %arg2: tensor<i32>,
-+    %arg3: tensor<i32>)
-+    -> tensor<16x8xf32> {
-+  %0 = stablehlo.dynamic_update_slice %arg0, %arg1, %arg2, %arg3
-+      : (tensor<16x8xf32>, tensor<2x8xf32>, tensor<i32>, tensor<i32>)
-+      -> tensor<16x8xf32>
-+  return %0 : tensor<16x8xf32>
-+}
-+
- //===----------------------------------------------------------------------===//
- // stablehlo.rng_bit_generator tests
- //===----------------------------------------------------------------------===//
-diff --git a/shardy/dialect/sdy/transforms/export/utils.cc b/shardy/dialect/sdy/transforms/export/utils.cc
-index 92bb4bc3..54e5d6f3 100644
---- a/shardy/dialect/sdy/transforms/export/utils.cc
-+++ b/shardy/dialect/sdy/transforms/export/utils.cc
-@@ -77,6 +77,13 @@ SmallVector<AxisRefAttr> getOrderedAxes(ArrayRef<AxisList> axesPerDim) {
-   return result;
- }
- 
-+// Returns the sharded size of dimension `dim` in `sharding`, or 1 if `sharding`
-+// is null.
-+int64_t getDimShardedSize(TensorShardingAttr sharding, int64_t dim,
-+                          MeshAttr mesh) {
-+  return sharding ? sharding.getDimSharding(dim).getShardedSize(mesh) : 1;
-+}
-+
- }  // namespace
- 
- void alignSubAxesByDecomposition(AxisList& axes,
-@@ -156,7 +163,7 @@ TensorShardingAttr updateSharding(TensorShardingAttr sharding,
- 
- bool isCommunicationFreeSliceDim(int64_t dimIdx, stablehlo::SliceOp sliceOp,
-                                  TensorShardingAttr sharding, MeshAttr mesh) {
--  int64_t shardCount = sharding.getDimShardings()[dimIdx].getShardedSize(mesh);
-+  int64_t shardCount = getDimShardedSize(sharding, dimIdx, mesh);
- 
-   if (shardCount <= 1) {
-     return true;
-@@ -183,7 +190,7 @@ bool isCommunicationFreeSliceDim(int64_t dimIdx, stablehlo::SliceOp sliceOp,
- 
- bool isCommunicationFreePadDim(int64_t dimIdx, stablehlo::PadOp padOp,
-                                TensorShardingAttr sharding, MeshAttr mesh) {
--  int64_t shardCount = sharding.getDimShardings()[dimIdx].getShardedSize(mesh);
-+  int64_t shardCount = getDimShardedSize(sharding, dimIdx, mesh);
-   if (shardCount <= 1) {
-     return true;
-   }
-@@ -229,8 +236,7 @@ bool isCommunicationFreeDynamicUpdateSliceDim(
-     return false;
-   }
- 
--  int64_t shardCount =
--      operandSharding.getDimShardings()[dimIdx].getShardedSize(mesh);
-+  int64_t shardCount = getDimShardedSize(operandSharding, dimIdx, mesh);
-   // Return true if the dimension is unsharded, not sliced, or a single-element
-   // slice.
-   if (shardCount <= 1 || operandDimSize == updateDimSize ||
-@@ -408,10 +414,8 @@ RankedTensorType getDivisiblePaddedType(RankedTensorType globalType,
-       paddedShape.push_back(dimSize);
-       continue;
-     }
--    int64_t inShardSize =
--        inSharding ? inSharding.getDimSharding(dim).getShardedSize(mesh) : 1;
--    int64_t outShardSize =
--        outSharding ? outSharding.getDimSharding(dim).getShardedSize(mesh) : 1;
-+    int64_t inShardSize = getDimShardedSize(inSharding, dim, mesh);
-+    int64_t outShardSize = getDimShardedSize(outSharding, dim, mesh);
-     int64_t divisor = std::lcm(inShardSize, outShardSize);
-     paddedShape.push_back(llvm::alignTo(dimSize, divisor));
-   }
-diff --git a/shardy/dialect/sdy/transforms/import/test/import_pipeline.mlir b/shardy/dialect/sdy/transforms/import/test/import_pipeline.mlir
-index d8200b88..60204b54 100644
---- a/shardy/dialect/sdy/transforms/import/test/import_pipeline.mlir
-+++ b/shardy/dialect/sdy/transforms/import/test/import_pipeline.mlir
-@@ -2,8 +2,8 @@
- 
- sdy.mesh @mesh = <["a"=2]>
- 
--// Verifies that `-apply-sharding-constraints` pass is applied after
--// `-add-data_flow_edges` pass
-+// Verifies that `-apply-sharding-constraints` pass is not applied
-+// before -add-data_flow_edges pass.
- // CHECK-LABEL: func @main
- func.func @main(%arg0: tensor<32x96xf32>) -> tensor<32x96xf32> {
-   // CHECK-NEXT: %[[OPT_BARRIER:.*]] = stablehlo.optimization_barrier %arg0
-@@ -53,11 +53,9 @@ func.func @main(%arg0: tensor<8xf32>) -> tensor<8xf32> {
- 
- sdy.mesh @mesh = <["c"=2, "a"=2, "b"=2]>
- 
--// Due to the in_sharding being fully closed, the in_sharding is added to the
--// func arg but with the manual axis added as replicated.
--// CHECK-LABEL: func @main
--// CHECK-SAME     %arg0: tensor<16x16xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"c"}, {}], replicated={"a"}>}
--// CHECK-SAME     -> tensor<16x16xf32> {
-+// CHECK-LABEL: func @main(
-+// CHECK-SAME:      %arg0: tensor<8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"c"}], replicated={"a"}>})
-+// CHECK-SAME:      -> tensor<8xf32> {
- func.func @main(%arg0: tensor<8xf32>) -> tensor<8xf32> {
-   // CHECK-NEXT: %0 = sdy.manual_computation(%arg0) in_shardings=[<@mesh, [{"c"}], replicated={"a"}>] out_shardings=[<@mesh, [{"c"}], replicated={"a"}>] manual_axes={"c", "a"} (%arg1: tensor<4xf32>) {
-   // CHECK-NEXT:   %2 = sdy.data_flow_edge %arg1 sharding=<@mesh, [{}]>
-diff --git a/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc b/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc
-index 8417874d..48217e35 100644
---- a/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc
-+++ b/shardy/dialect/sdy/transforms/propagation/basic_propagation.cc
-@@ -776,7 +776,9 @@ bool allValidShapes(ModuleOp moduleOp) {
-                   }
-                   if (auto tupleType = dyn_cast<TupleType>(type)) {
-                     op->emitError("Shardy propagation doesn't support tuples: ")
--                        << tupleType;
-+                        << tupleType
-+                        << ". Please flatten or canonicalize tuples before "
-+                           "running Shardy propagation.";
-                     return WalkResult::interrupt();
-                   }
-                 }
-diff --git a/shardy/dialect/sdy/transforms/propagation/debugging/source_sharding.cc b/shardy/dialect/sdy/transforms/propagation/debugging/source_sharding.cc
-index 183f25fd..b2933711 100644
---- a/shardy/dialect/sdy/transforms/propagation/debugging/source_sharding.cc
-+++ b/shardy/dialect/sdy/transforms/propagation/debugging/source_sharding.cc
-@@ -691,7 +691,7 @@ void prepareFuncResultToEdgesHandler(
-   }
- }
- 
--OriginSharding lookUpValueOriginSharding(
-+std::optional<OriginSharding> lookUpValueOriginSharding(
-     Value value, AxisRefAttr axisRef,
-     const ValueToOriginShardingMap& valueToOriginShardingMap) {
-   // NOTE: need to call `getShardableValue` in case the operand/result is
-@@ -700,7 +700,7 @@ OriginSharding lookUpValueOriginSharding(
-   auto valueIt = valueToOriginShardingMap.find(getShardableValue(value));
-   if (valueIt == valueToOriginShardingMap.end()) {
-     SDY_LOG(WARNING) << "Couldn't find sharding origin for value";
--    return {};
-+    return std::nullopt;
-   }
-   const AxisToOriginShardingMap& axisToOriginSharding = valueIt->second;
-   if (auto it = axisToOriginSharding.find(axisRef);
-@@ -717,7 +717,7 @@ OriginSharding lookUpValueOriginSharding(
-     }
-   }
-   SDY_LOG(WARNING) << "Couldn't find sharding origin";
--  return {};
-+  return std::nullopt;
- }
- 
- }  // namespace
-@@ -752,8 +752,9 @@ void SourceShardingHandler::operator()(function_ref<void()> transform,
-   propagationStep++;
-   // If the new and old shardings are different, something was propagated to it.
-   // Find and save it.
--  auto lookUpOriginSharding = [&](EdgeNode edgeNode,
--                                  AxisRefAttr axisRef) -> OriginSharding {
-+  auto lookUpOriginSharding =
-+      [&](EdgeNode edgeNode,
-+          AxisRefAttr axisRef) -> std::optional<OriginSharding> {
-     switch (edgeNode.type) {
-       case EdgeNodeType::OPERAND:
-         return lookUpValueOriginSharding(
-@@ -778,8 +779,11 @@ void SourceShardingHandler::operator()(function_ref<void()> transform,
-       }
-     }
-     if (mappings->debugShardingOrigins) {
--      mappings->valueToOriginShardingMap[getShardableValue(value)].try_emplace(
--          axisRef, lookUpOriginSharding(edge.source, axisRef));
-+      if (std::optional<OriginSharding> originSharding =
-+              lookUpOriginSharding(edge.source, axisRef)) {
-+        mappings->valueToOriginShardingMap[getShardableValue(value)]
-+            .try_emplace(axisRef, *originSharding);
-+      }
-     }
-   };
- 
-diff --git a/shardy/dialect/sdy/transforms/propagation/debugging/test/sharding_origins.mlir b/shardy/dialect/sdy/transforms/propagation/debugging/test/sharding_origins.mlir
-index 16d7d20b..f9c3e2d2 100644
---- a/shardy/dialect/sdy/transforms/propagation/debugging/test/sharding_origins.mlir
-+++ b/shardy/dialect/sdy/transforms/propagation/debugging/test/sharding_origins.mlir
-@@ -790,7 +790,7 @@ sdy.mesh @mesh = <["a"=2, "b"=2, "c"=8]>
- 
- // CHECK-LABEL: func.func @input_already_has_sharding(
- // CHECK-SAME:    %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a", ?}, {"b", ?}]>,
--// CHECK-SAME:                            sdy.sharding_origins = {a = "self", b = "constraint_0"}})
-+// CHECK-SAME:                            sdy.sharding_origins = {b = "constraint_0"}})
- // CHECK-SAME:    -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{?}, {"b", ?}]>,
- // CHECK-SAME:                         sdy.sharding_origins = {b = "constraint_0"}}) {
- func.func @input_already_has_sharding(%arg0: tensor<8x8xf32>) -> tensor<8x8xf32> {
-diff --git a/shardy/dialect/sdy/transforms/propagation/test/basic_propagation.mlir b/shardy/dialect/sdy/transforms/propagation/test/basic_propagation.mlir
-index c5cd8022..6d9f32c2 100644
---- a/shardy/dialect/sdy/transforms/propagation/test/basic_propagation.mlir
-+++ b/shardy/dialect/sdy/transforms/propagation/test/basic_propagation.mlir
-@@ -1104,3 +1104,14 @@ func.func private @foo(%arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mes
-   %0 = stablehlo.abs %arg0 : tensor<8x8xf32>
-   return %0: tensor<8x8xf32>
- }
-+
-+// -----
-+
-+sdy.mesh @mesh = <["a"=2]>
-+
-+func.func @tuple_not_supported(%arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh, [{"a"}, {}]>}) -> tensor<8x8xf32> {
-+  // expected-error @+1 {{Shardy propagation doesn't support tuples: 'tuple<tensor<8x8xf32>>'. Please flatten or canonicalize tuples before running Shardy propagation.}}
-+  %0 = stablehlo.tuple %arg0 : tuple<tensor<8x8xf32>>
-+  %1 = stablehlo.get_tuple_element %0[0] : (tuple<tensor<8x8xf32>>) -> tensor<8x8xf32>
-+  return %1 : tensor<8x8xf32>
-+}
-diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
-index e42d7dbd..caf6ccb5 100644
---- a/third_party/llvm/workspace.bzl
-+++ b/third_party/llvm/workspace.bzl
-@@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
- 
- def repo(name):
-     """Imports LLVM."""
--    LLVM_COMMIT = "2e0d9e12b36e853d9b108c4da09509cd54d1f3f9"
--    LLVM_SHA256 = "8e4a851780af109d623fa6d6db9b1e62d7126272d17268fed1b149f0e737c1d7"
-+    LLVM_COMMIT = "1ee8805f01a663b7f2b09215a8eab5ab2d15ce35"
-+    LLVM_SHA256 = "273fa5e9618fa8085aea51d6eb40129677fea8cf4c8825a2c4ff6ea582706a55"
- 
-     tf_http_archive(
-         name = name,
-diff --git a/third_party/stablehlo/temporary.patch b/third_party/stablehlo/temporary.patch
-index c0bd29a0..d2710427 100644
---- a/third_party/stablehlo/temporary.patch
-+++ b/third_party/stablehlo/temporary.patch
-@@ -15,755 +15,42 @@
- #
- # This file is automatically generated by the `generate_patch` tool.
- # Do not edit directly.
--diff --ruN a/stablehlo/.github/ISSUE_TEMPLATE/bug_report.yml b/stablehlo/.github/ISSUE_TEMPLATE/bug_report.yml
----- stablehlo/.github/ISSUE_TEMPLATE/bug_report.yml
--+++ stablehlo/.github/ISSUE_TEMPLATE/bug_report.yml
--@@ -1,50 +0,0 @@
---# Copyright 2024 The StableHLO Authors.
---# 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.
---
---name: 🐞 Bug Report
---description: Report an issue
---body:
---  - type: markdown
---    attributes:
---      value: |
---        :star2: Thanks for taking the time to report this issue! :star2:
---
---        Please search through [other recent issues](https://github.com/openxla/stablehlo/issues) to see if your report overlaps with an existing issue.
---  - type: textarea
---    id: what-happened
---    attributes:
---      label: What happened?
---      description: |
---        Also tell us, what did you expect to happen?
---
---        Include any MLIR output leading up to whatever error occurred where possible.
---      placeholder: Tell us what you see! Please also copy/paste or link to any relevant log output.
---    validations:
---      required: true
---  - type: textarea
---    id: reproducing
---    attributes:
---      label: Steps to reproduce your issue
---      description: |
---        Include any commands that you ran.
---      value: |
---        1. Go to '...'
---        2. Click on '....'
---        3. Scroll down to '....'
---        4. See error
---  - type: textarea
---    id: version
---    attributes:
---      label: Version information
---      description: What version of our software are you using?
---      placeholder: This could be a git commit hash.
--diff --ruN a/stablehlo/.github/ISSUE_TEMPLATE/config.yml b/stablehlo/.github/ISSUE_TEMPLATE/config.yml
----- stablehlo/.github/ISSUE_TEMPLATE/config.yml
--+++ stablehlo/.github/ISSUE_TEMPLATE/config.yml
--@@ -1,24 +0,0 @@
---# Copyright 2024 The StableHLO Authors.
---# 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.
---
---blank_issues_enabled: true
---contact_links:
---  - name: 📄 Blank Issue
---    url: https://github.com/openxla/stablehlo/issues/new
---    about: If you know what you're doing (especially collaborators)
---  - name: 🗣 Start a discussion on GitHub
---    url: https://github.com/orgs/openxla/discussions/new
---    about: For early feedback on a design or other proposal
---  - name: 💬 Join us on Discord
---    url: https://discord.gg/Wqr8vWZKKJ
---    about: For realtime communication with the team and collaborators
--diff --ruN a/stablehlo/.github/ISSUE_TEMPLATE/feature_request.yml b/stablehlo/.github/ISSUE_TEMPLATE/feature_request.yml
----- stablehlo/.github/ISSUE_TEMPLATE/feature_request.yml
--+++ stablehlo/.github/ISSUE_TEMPLATE/feature_request.yml
--@@ -1,30 +0,0 @@
---# Copyright 2024 The StableHLO Authors.
---# 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.
---
---name: ➕ Feature Request
---description: Ask for a new feature to be added
---body:
---  - type: markdown
---    attributes:
---      value: |
---        :star2: Thanks for taking the time to file this request! :star2:
---
---        We also work closely with upstream communities like the MLIR project, which uses the [MLIR Discourse forum](https://discourse.llvm.org/c/mlir/31) for coordination.
---  - type: textarea
---    id: overview
---    attributes:
---      label: Request description
---      description: |
---        What would you like added or changed?
---    validations:
---      required: true
--diff --ruN a/stablehlo/.github/actions/setup-build/action.yml b/stablehlo/.github/actions/setup-build/action.yml
----- stablehlo/.github/actions/setup-build/action.yml
--+++ stablehlo/.github/actions/setup-build/action.yml
--@@ -1,58 +0,0 @@
---# Copyright 2023 The StableHLO Authors.
---#
---# 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
---#
---#      https://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.
---
---# The setup-build action gets everything needed by buildAndTest into the workspace.
---name: "Setup build environment (ninja, ccache, llvm, clang, lld)"
---description: "Setup build environment (ninja, ccache, llvm, clang, lld)"
---
---inputs:
---  llvm-version:
---    description: |
---      LLVM version for checkout and build. Used for ccache value and checkout.
---    required: true
---  clang-version:
---    description: |
---      Clang and LLD version to use for compilation. Independent of LLVM version.
---    required: false
---    default: '20'
---  python-version:
---    description: |
---      Python version to install.
---    required: false
---    default: '3.11.13'
---runs:
---  # This is a composite action - has a list of steps to execute.
---  using: "composite"
---
---  steps:
---  # Checkout llvm at version specified in input argument.
---  - uses: actions/checkout@v4
---    with:
---      repository: llvm/llvm-project
---      ref: ${{ inputs.llvm-version }}
---      path: llvm-project
---      fetch-depth: 1
---
---  # Setup C++ caching using ccache.
---  # Cache key is a combination of OS arch and LLVM version.
---  - name: Ccache for C++ compilation
---    uses: hendrikmuhs/ccache-action@v1.2
---    with:
---      key: ${{ runner.os }}-stablehlo_build_assets-${{ inputs.llvm-version }}
---      max-size: 4G
---
---  - name: Install MLIR python requirements
---    shell: bash
---    run: |
---      python3 -m pip install -r $GITHUB_WORKSPACE/llvm-project/mlir/python/requirements.txt
- diff --ruN a/stablehlo/.github/workflows/buildAndTestBazel.yml b/stablehlo/.github/workflows/buildAndTestBazel.yml
- --- stablehlo/.github/workflows/buildAndTestBazel.yml
- +++ stablehlo/.github/workflows/buildAndTestBazel.yml
--@@ -1,103 +0,0 @@
---# Copyright 2023 The StableHLO Authors.
---#
---# 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
---#
---#      https://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.
---
---name: Bazel Build
---
---on:
---  pull_request:
---    branches: [ main ]
---  push:
---    branches: [ main ]
---  schedule:
---    # Run once a day
---    - cron:  '0 12 * * *'
---  workflow_dispatch:
---    inputs:
---      halt-for-connection:
---        description: 'Should this workflow run wait for a remote connection?'
---        type: choice
---        required: true
---        default: 'no'
---        options:
---          - 'yes'
---          - 'no'
---
---# Ensure that only a single job or workflow using the same
---# concurrency group will run at a time. This would cancel
---# any in-progress jobs in the same github workflow and github
---# ref (e.g. refs/heads/main or refs/pull/<pr_number>/merge).
---concurrency:
---  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
---  cancel-in-progress: ${{ github.ref != 'main' }}
---
---jobs:
---  bazel-build:
---    strategy:
---      fail-fast: false
-+@@ -45,7 +45,7 @@
-+   bazel-build:
-+     strategy:
-+       fail-fast: false
- -    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n2-16' || 'ubuntu-22.04' }}
---    container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
---
---    steps:
---    # Checkout StableHLO with all commits in the chain
---    - name: Checkout StableHLO
---      uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
---      with:
---        fetch-depth: $(( ${{ github.event_name == 'pull_request' && github.event.pull_request.commits || 0 }} + 1 ))
---        set-safe-directory: true
---
---    - name: Mark project as safe for Git
---      run: git config --global --add safe.directory '*'
---
---    - name: "Wait For Connection"
---      uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
---      with:
---        halt-dispatch-input: ${{ inputs.halt-for-connection }}
---
---    - name: Validate LLVM Commit
---      run: |
---        ./build_tools/github_actions/lint_llvm_commit.sh .
---
---    - name: Get LLVM Version
---      id: llvm-version
---      shell: bash
---      run: |
---        echo "version=$(cat ./build_tools/llvm_version.txt)" >> $GITHUB_OUTPUT
---
---    - name: Build and check docs
---      shell: bash
---      run: |
---        ./build_tools/github_actions/ci_build_docs.sh -c
---
---    - name: Build and Test StableHLO (Diff Only)
---      if: ${{ github.event_name == 'pull_request' }}
---      shell: bash
---      run: |
---        START_HASH=${{ github.event.pull_request.base.sha }}
---        END_HASH=${{ github.event.pull_request.head.sha }}
---        BAZEL_DIFF=/tmp/bazel-diff.jar
---        curl -Lo "$BAZEL_DIFF" https://github.com/Tinder/bazel-diff/releases/download/17.0.2/bazel-diff_deploy.jar
---        ./build_tools/github_actions/ci_build_bazel.sh "$BAZEL_DIFF" "$START_HASH" "$END_HASH"
---
---    - name: Build and Test StableHLO (All)
---      if: ${{ github.event_name != 'pull_request' }}
---      shell: bash
---      run: |
---        ./build_tools/github_actions/ci_build_bazel.sh
---
---    - name: "Wait For Connection - On Error"
---      if: failure()
---      uses: google-ml-infra/actions/ci_connection@7f5ca0c263a81ed09ea276524c1b9192f1304e3c
---      with:
---        halt-dispatch-input: ${{ inputs.halt-for-connection }}
---
-++    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n4-16' || 'ubuntu-22.04' }}
-+     container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest"
-+ 
-+     steps:
- diff --ruN a/stablehlo/.github/workflows/buildAndTestCMake.yml b/stablehlo/.github/workflows/buildAndTestCMake.yml
- --- stablehlo/.github/workflows/buildAndTestCMake.yml
- +++ stablehlo/.github/workflows/buildAndTestCMake.yml
--@@ -1,98 +0,0 @@
---# Copyright 2023 The StableHLO Authors.
---#
---# 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
---#
---#      https://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.
---
---name: CMake Build
---
---on:
---  pull_request:
---    branches: [ main ]
---  push:
---    branches: [ main ]
---  schedule:
---    # Run once every 6hr building against llvm-project@HEAD
---    - cron:  '0 */6 * * *'
---  workflow_dispatch:
---
---# Ensure that only a single job or workflow using the same
---# concurrency group will run at a time. This would cancel
---# any in-progress jobs in the same github workflow and github
---# ref (e.g. refs/heads/main or refs/pull/<pr_number>/merge).
---concurrency:
---  group: ${{ github.workflow }}-${{ github.ref }}
---  cancel-in-progress: true
---
---jobs:
---  cmake-build:
---    # Only run scheduled CI on main repo
---    if: (github.repository == 'openxla/stablehlo' || github.event_name != 'schedule')
---    name: "cmake-build${{ github.event_name == 'schedule' && ' (llvm-project@HEAD)' || ''}}"
---    env:
---      LLVM_PROJECT_DIR: "llvm-project"
---      LLVM_BUILD_DIR: "llvm-build"
---      STABLEHLO_BUILD_DIR: "stablehlo-build"
---      STABLEHLO_PYTHON_BUILD_DIR: "stablehlo-python-build"
---    strategy:
---      fail-fast: false
-+@@ -44,7 +44,7 @@
-+       STABLEHLO_PYTHON_BUILD_DIR: "stablehlo-python-build"
-+     strategy:
-+       fail-fast: false
- -    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n2-16' || 'ubuntu-22.04' }}
---    container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cmake:latest"
---
---    steps:
---    - name: Checkout StableHLO
---      uses: actions/checkout@v4
---
---    - name: Get LLVM Version
---      id: llvm-version
---      shell: bash
---      run: |
---        USE_LLVM_HEAD=${{ github.event_name == 'schedule' }}
---        if [[ $USE_LLVM_HEAD = true ]]; then
---          echo "version=main" >> $GITHUB_OUTPUT
---        else
---          echo "version=$(cat ./build_tools/llvm_version.txt)" >> $GITHUB_OUTPUT
---        fi;
---
---    - name: Setup workspace
---      uses: ./.github/actions/setup-build
---      with:
---        llvm-version: ${{ steps.llvm-version.outputs.version }}
---
---    - name: Configure and Build LLVM
---      shell: bash
---      run: |
---        ./build_tools/github_actions/ci_build_cmake_llvm.sh "$LLVM_PROJECT_DIR" "$LLVM_BUILD_DIR"
---      env:
---          CMAKE_BUILD_TYPE: Release
---          MLIR_ENABLE_BINDINGS_PYTHON: ON
---
---    - name: Build and Test StableHLO (with AddressSanitizer)
---      shell: bash
---      run: |
---        ./build_tools/github_actions/ci_build_cmake.sh "$LLVM_BUILD_DIR" "$STABLEHLO_BUILD_DIR"
---      env:
---          CMAKE_BUILD_TYPE: Release
---          STABLEHLO_ENABLE_BINDINGS_PYTHON: OFF
---          STABLEHLO_ENABLE_SANITIZER: address
---          # TODO: remove this once https://github.com/openxla/stablehlo/pull/2318 is fixed
---          ASAN_OPTIONS: allow_user_poisoning=false
---
---    - name: Build and Test StableHLO (with Python bindings)
---      shell: bash
---      run: |
---        python3 -m pip install tensorflow-cpu nanobind==2.9
---        python3 -m pip install -r "$LLVM_PROJECT_DIR/mlir/python/requirements.txt"
---        ./build_tools/github_actions/ci_build_cmake.sh "$LLVM_BUILD_DIR" "$STABLEHLO_BUILD_DIR"
---      env:
---          CMAKE_BUILD_TYPE: Release
---          STABLEHLO_ENABLE_BINDINGS_PYTHON: ON
---          STABLEHLO_ENABLE_PYTHON_TF_TESTS: ON
--diff --ruN a/stablehlo/.github/workflows/execute_tutorials.yml b/stablehlo/.github/workflows/execute_tutorials.yml
----- stablehlo/.github/workflows/execute_tutorials.yml
--+++ stablehlo/.github/workflows/execute_tutorials.yml
--@@ -1,46 +0,0 @@
---# Copyright 2024 The StableHLO 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
---
---#      https://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.
---
---name: Execute Tutorial Notebooks
---
---on:
---  push:
---    paths:
---      - 'docs/tutorials/**.ipynb'
---      - .github/workflows/execute_tutorials.yml
---  schedule:
---    # Run once a day
---    - cron:  '0 20 * * *'
---  workflow_dispatch:
---
---jobs:
---  test-notebooks:
---    runs-on: ubuntu-latest
---
---    steps:
---      - name: Checkout repository
---        uses: actions/checkout@v4
---
---      - name: Set up Python and install dependencies
---        uses: actions/setup-python@v5
---        with:
---          python-version: '3.11'
---      - run:  |
---          pip install --upgrade pip
---          pip install ipython ipykernel nbclient
---
---      - name: Execute tutorial notebooks
---        run: |
---          jupyter execute docs/tutorials/*.ipynb
---        shell: bash
--diff --ruN a/stablehlo/.github/workflows/lint.yml b/stablehlo/.github/workflows/lint.yml
----- stablehlo/.github/workflows/lint.yml
--+++ stablehlo/.github/workflows/lint.yml
--@@ -1,111 +0,0 @@
---# Copyright 2023 The StableHLO Authors.
---#
---# 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
---#
---#      https://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.
---
---name: Lint
---
---on:
---  pull_request:
---    branches: [ main ]
---  push:
---    branches: [ main ]
---  schedule:
---    # Run once a day
---    - cron:  '0 12 * * *'
---  workflow_dispatch:
---
---env:
---  # Have `go install` place binaries in $PATH
---  GOBIN: "/usr/local/bin"
---
---jobs:
---  clang-format:
---    # This job can only be run on pull_request since GITHUB_BASE_REF is only set on PR.
---    if: "github.event_name == 'pull_request'"
---    runs-on: ubuntu-22.04
---    steps:
---      - name: Checking out repository
---        uses: actions/checkout@v4
---
---      - name: Run clang-format
---        run: |
---          git fetch --no-tags --prune --depth=1 origin "${GITHUB_BASE_REF?}:${GITHUB_BASE_REF?}"
---          ./build_tools/github_actions/lint_clang_format.sh -b "${GITHUB_BASE_REF}"
---
---  whitespace-checks:
---    if: "github.event_name == 'pull_request'"
---    runs-on: ubuntu-22.04
---    steps:
---      - name: Checking out repository
---        uses: actions/checkout@v4
---
---      - name: Run Whitespace Checks
---        run: |
---          git fetch --no-tags --prune --depth=1 origin "${GITHUB_BASE_REF?}:${GITHUB_BASE_REF?}"
---          ./build_tools/github_actions/lint_whitespace_checks.sh
---
---  version-checks:
---    if: "github.event_name == 'pull_request'"
---    runs-on: ubuntu-22.04
---    steps:
---      - name: Checking out repository
---        uses: actions/checkout@v4
---
---      - name: Run Version Checks
---        run: |
---          ./build_tools/github_actions/lint_version.sh
---
---  buildifier-check:
---    if: "github.event_name == 'pull_request'"
---    runs-on: ubuntu-22.04
---    defaults:
---      run:
---        shell: bash
---    steps:
---      - name: Checking out repository
---        uses: actions/checkout@v4
---
---      - name: Install buildifier
---        run: go install github.com/bazelbuild/buildtools/buildifier@433ea85 # 6.4.0
---
---      - name: Run buildifier checks
---        run: ./build_tools/github_actions/lint_buildifier.sh
---
---  license-check:
---    if: "github.event_name == 'pull_request'"
---    runs-on: ubuntu-22.04
---    defaults:
---      run:
---        shell: bash
---    steps:
---      - name: Checking out repository
---        uses: actions/checkout@v4
---
---      - name: Run license checks
---        run: |
---          git fetch --no-tags --prune --depth=1 origin "${GITHUB_BASE_REF?}:${GITHUB_BASE_REF?}"
---          ./build_tools/github_actions/lint_check_license.sh -b "${GITHUB_BASE_REF}"
---
---  shell-script-check:
---    if: "github.event_name == 'pull_request'"
---    runs-on: ubuntu-22.04
---    defaults:
---      run:
---        shell: bash
---    steps:
---      - name: Checking out repository
---        uses: actions/checkout@v4
---
---      # shellcheck is installed by default on Ubuntu runners: https://github.com/actions/runner-images/blob/37a171e44de1f134793d7f0cbff137ca481942ab/images/ubuntu/Ubuntu2204-Readme.md
---      - name: Run shellcheck
---        run: ./build_tools/github_actions/lint_shellcheck.sh
--diff --ruN a/stablehlo/.github/workflows/markdown_lint.yml b/stablehlo/.github/workflows/markdown_lint.yml
----- stablehlo/.github/workflows/markdown_lint.yml
--+++ stablehlo/.github/workflows/markdown_lint.yml
--@@ -1,35 +0,0 @@
---# Copyright 2023 The StableHLO Authors.
---#
---# 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
---#
---#      https://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.
---
---name: Markdown Lint
---
---on:
---  pull_request:
---    branches: [ main ]
---  push:
---    branches: [ main ]
---  schedule:
---    # Run once a day
---    - cron:  '0 12 * * *'
---  workflow_dispatch:
---
---jobs:
---  markdown-lint:
---    runs-on: ubuntu-22.04
---    steps:
---    - name: Checking out repository
---      uses: actions/checkout@v4
---    - name: Running markdownlint
---      run: |
---        ./build_tools/github_actions/lint_markdown.sh --ignore "./docs/generated/**/*.md" "./**/*.md"
-++    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n4-16' || 'ubuntu-22.04' }}
-+     container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cmake:latest"
-+ 
-+     steps:
- diff --ruN a/stablehlo/.github/workflows/publishWheelRelease.yml b/stablehlo/.github/workflows/publishWheelRelease.yml
- --- stablehlo/.github/workflows/publishWheelRelease.yml
- +++ stablehlo/.github/workflows/publishWheelRelease.yml
--@@ -1,136 +0,0 @@
---# Copyright 2024 The StableHLO Authors.
---#
---# 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
---#
---#      https://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.
---
---name: Publish Python Wheel Release
---
---on:
---  workflow_dispatch:
---  schedule:
---    # Runs at 6:00 AM UTC, which is 10:00 PM previous day PST (UTC-8)
---    - cron: '0 6 * * *'
---
---# Ensure that only a single job or workflow using the same
---# concurrency group will run at a time. This would cancel
---# any in-progress jobs in the same github workflow and github
---# ref (e.g. refs/heads/main or refs/pull/<pr_number>/merge).
---concurrency:
---  group: ${{ github.workflow }}-${{ github.ref }}
---  cancel-in-progress: true
---
---jobs:
---  # TODO(fzakaria): Consider running this job only if a new commit for the day has been done
---  cmake-build-wheel:
---    # Only run scheduled CI on main repo
---    if: (github.repository == 'openxla/stablehlo' || github.event_name != 'schedule')
---    name: "cmake-build ${{ github.event_name == 'schedule' && '(llvm-project@HEAD)' || ''}}"
---    env:
---      LLVM_PROJECT_DIR: "llvm-project"
---      LLVM_BUILD_DIR: "llvm-build"
---      # The setup.py script expects the StableHLO build to be in 'build'
---      STABLEHLO_BUILD_DIR: "build"
---    strategy:
---      fail-fast: false
---      matrix:
---        python-version: [ "3.10", "3.11", "3.12" ]
-+@@ -43,7 +43,7 @@
-+       fail-fast: false
-+       matrix:
-+         python-version: [ "3.10", "3.11", "3.12" ]
- -    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n2-16' ||  'ubuntu-22.04'  }}
---    container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cmake:latest"
---
---    steps:
---    - name: Checkout StableHLO
---      uses: actions/checkout@v4
---      # TODO(fzakaria): We need all the commits and the tags for the
---      # git describe` command to work to find the correct version.
---      # This can be removed if the version can be determmined in another way.
---      with:
---        fetch-depth: 0
---    - name: Get LLVM Version
---      id: llvm-version
---      shell: bash
---      run: |
---        echo "version=$(cat ${{ github.workspace }}/build_tools/llvm_version.txt)" >> $GITHUB_OUTPUT
---
---    - name: Setup workspace
---      uses: ./.github/actions/setup-build
---      with:
---        llvm-version: ${{ steps.llvm-version.outputs.version }}
---        python-version: ${{ matrix.python-version }}
---
---    # The actual tool need to execute the python build
---    # TODO(fzakaria): Consider using cibuildwheel to build multiple wheels at
---    # once and mark them as manylinux compatible.
---    - name: Install build dependencies
---      shell: bash
---      run: pip install build
---
---    - name: Configure and Build LLVM
---      shell: bash
---      run: |
---        ./build_tools/github_actions/ci_build_cmake_llvm.sh "$LLVM_PROJECT_DIR" "$LLVM_BUILD_DIR"
---      env:
---          CMAKE_BUILD_TYPE: Release
---          MLIR_ENABLE_BINDINGS_PYTHON: ON
---
---    - name: Build and Test StableHLO with Python
---      shell: bash
---      run: |
---        pip install tensorflow-cpu nanobind==2.4
---        pip install -r "$LLVM_PROJECT_DIR/mlir/python/requirements.txt"
---        ./build_tools/github_actions/ci_build_cmake.sh "$LLVM_BUILD_DIR" "$STABLEHLO_BUILD_DIR"
---      env:
---          CMAKE_BUILD_TYPE: Release
---          STABLEHLO_ENABLE_BINDINGS_PYTHON: ON
---          STABLEHLO_ENABLE_PYTHON_TF_TESTS: ON
---
---    - name: Build and Test Python Wheel
---      shell: bash
---      run: |
---        ./build_tools/github_actions/ci_build_python_wheel.sh "$GITHUB_WORKSPACE/${{ github.sha }}"
---
---    - name: Upload an artifact
---      uses: actions/upload-artifact@v4
---      with:
---        if-no-files-found: error
---        name: build_artifact-${{ matrix.python-version }}
---        path: ${{ github.sha }}
---
---  upload-tarballs:
---    runs-on: 'ubuntu-24.04'
---    needs: cmake-build-wheel
---    strategy:
---      fail-fast: false
---      matrix:
---        python-version: [ "3.10", "3.11", "3.12" ]
---
---    steps:
---      - name: Checkout
---        uses: actions/checkout@v4
---
---      - name: Download artifacts
---        uses: actions/download-artifact@v4
---        with:
---          name: build_artifact-${{ matrix.python-version }}
---          path: ${{ github.sha }}
---
---      - name: Release current commit
---        uses: ncipollo/release-action@v1.13.0
---        with:
---          artifacts: "${{ github.sha }}/*.whl"
---          token: "${{ secrets.PUBLISH_ACCESS_TOKEN }}"
---          # git describe in setup.py explicitly filters out this tag
---          tag: "dev-wheels"
---          name: "dev-wheels"
---          removeArtifacts: false
---          allowUpdates: true
---          replacesArtifacts: true
---          makeLatest: false
--diff --ruN a/stablehlo/.gitignore b/stablehlo/.gitignore
----- stablehlo/.gitignore
--+++ stablehlo/.gitignore
--@@ -1,14 +0,0 @@
---bazel-bin
---bazel-out
---bazel-stablehlo
---bazel-testlogs
---build
---cmake-build
---llvm-project
---llvm-build
---.vscode
---.cache
---.idea/
---**/__pycache__
---*.whl
---**/*.egg-info/
-++    runs-on: ${{ github.repository == 'openxla/stablehlo' && 'linux-x86-n4-16' ||  'ubuntu-22.04'  }}
-+     container: "us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-cmake:latest"
-+ 
-+     steps:
- diff --ruN a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel
- --- stablehlo/BUILD.bazel
- +++ stablehlo/BUILD.bazel
-@@ -3522,6 +2809,479 @@ diff --ruN a/stablehlo/stablehlo/integrations/cpp/builder/StablehloBuilderTest.c
-  TEST(MlirBuilderTest, GatherOp) {
-    std::string expected = R"mlir(module {
-    func.func @main(%arg0: tensor<3xi64>, %arg1: tensor<1x1xi64>) -> tensor<1xi64> {
-+@@ -1021,7 +1097,7 @@
-+ TEST(MlirBuilderTest, ConstantF4E2M1FNArray) {
-+   std::string expected = R"mlir(module {
-+   func.func @main() -> tensor<10xf4E2M1FN> {
-+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 0.000000e+00, 0.000000e+00, 3.000000e+00, 6.000000e+00, 6.000000e+00, 1.000000e+00, 6.000000e+00]> : tensor<10xf4E2M1FN>
-++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 0.000000e+00, 0.000000e+00, 3.000000e+00, 6.000000e+00, -6.000000e+00, 5.000000e-01, -5.000000e-01]> : tensor<10xf4E2M1FN>
-+     return %cst : tensor<10xf4E2M1FN>
-+   }
-+ })mlir";
-+@@ -1031,10 +1107,15 @@
-+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-+     auto type = makeTensorType(fb.getContext(), {10}, ElementType::F4E2M1FN);
-++    double posMax = 6.0;
-++    double negMax = -6.0;
-++    double posMinDenorm = 0.5;
-++    double negMinDenorm = -0.5;
-+     auto cst = stablehlo::Constant(
-+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
-+-                                           0x07, 0x0F, 0x01, 0x09}),
-+-                         type));
-++        fb, makeConstant(
-++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
-++                                  negMax, posMinDenorm, negMinDenorm}),
-++                type));
-+     func::Return(fb, {cst});
-+   }
-+ 
-+@@ -1045,7 +1126,7 @@
-+ TEST(MlirBuilderTest, ConstantF6E2M3FNArray) {
-+   std::string expected = R"mlir(module {
-+   func.func @main() -> tensor<10xf6E2M3FN> {
-+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.250000e-01, 3.250000e+00, 7.500000e+00, 7.500000e+00, 1.000000e+00, 7.500000e+00]> : tensor<10xf6E2M3FN>
-++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.250000e-01, 3.250000e+00, 7.500000e+00, -7.500000e+00, 1.250000e-01, -1.250000e-01]> : tensor<10xf6E2M3FN>
-+     return %cst : tensor<10xf6E2M3FN>
-+   }
-+ })mlir";
-+@@ -1055,10 +1136,15 @@
-+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-+     auto type = makeTensorType(fb.getContext(), {10}, ElementType::F6E2M3FN);
-++    double posMax = 7.5;
-++    double negMax = -7.5;
-++    double posMinDenorm = 0x1p-3;
-++    double negMinDenorm = -0x1p-3;
-+     auto cst = stablehlo::Constant(
-+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
-+-                                           0x1F, 0x3F, 0x01, 0x21}),
-+-                         type));
-++        fb, makeConstant(
-++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
-++                                  negMax, posMinDenorm, negMinDenorm}),
-++                type));
-+     func::Return(fb, {cst});
-+   }
-+ 
-+@@ -1069,7 +1155,7 @@
-+ TEST(MlirBuilderTest, ConstantF6E3M2FNArray) {
-+   std::string expected = R"mlir(module {
-+   func.func @main() -> tensor<10xf6E3M2FN> {
-+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.250000e-01, 3.000000e+00, 2.800000e+01, 2.800000e+01, 1.000000e+00, 2.800000e+01]> : tensor<10xf6E3M2FN>
-++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.250000e-01, 3.000000e+00, 2.800000e+01, -2.800000e+01, 6.250000e-02, -6.250000e-02]> : tensor<10xf6E3M2FN>
-+     return %cst : tensor<10xf6E3M2FN>
-+   }
-+ })mlir";
-+@@ -1079,10 +1165,15 @@
-+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-+     auto type = makeTensorType(fb.getContext(), {10}, ElementType::F6E3M2FN);
-++    double posMax = 28.0;
-++    double negMax = -28.0;
-++    double posMinDenorm = 0x1p-4;
-++    double negMinDenorm = -0x1p-4;
-+     auto cst = stablehlo::Constant(
-+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
-+-                                           0x1F, 0x3F, 0x01, 0x21}),
-+-                         type));
-++        fb, makeConstant(
-++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
-++                                  negMax, posMinDenorm, negMinDenorm}),
-++                type));
-+     func::Return(fb, {cst});
-+   }
-+ 
-+@@ -1092,9 +1183,9 @@
-+ 
-+ TEST(MlirBuilderTest, ConstantF8E3M4Array) {
-+   std::string expected = R"mlir(module {
-+-  func.func @main() -> tensor<10xf8E3M4> {
-+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.125000e+00, 0x70, 0x70, 1.000000e+00, 0x70]> : tensor<10xf8E3M4>
-+-    return %cst : tensor<10xf8E3M4>
-++  func.func @main() -> tensor<11xf8E3M4> {
-++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.125000e+00, 0x70, 0xF0, 0x78, 1.562500e-02, -1.562500e-02]> : tensor<11xf8E3M4>
-++    return %cst : tensor<11xf8E3M4>
-+   }
-+ })mlir";
-+ 
-+@@ -1102,11 +1193,17 @@
-+   {
-+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-+-    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E3M4);
-++    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E3M4);
-++    double posInf = std::numeric_limits<double>::infinity();
-++    double negInf = -std::numeric_limits<double>::infinity();
-++    double nan = std::numeric_limits<double>::quiet_NaN();
-++    double posMinDenorm = 0x1p-6;
-++    double negMinDenorm = -0x1p-6;
-+     auto cst = stablehlo::Constant(
-+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
-+-                                           0x7F, 0xFF, 0x01, 0x81}),
-+-                         type));
-++        fb, makeConstant(
-++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posInf,
-++                                  negInf, nan, posMinDenorm, negMinDenorm}),
-++                type));
-+     func::Return(fb, {cst});
-+   }
-+ 
-+@@ -1116,9 +1213,9 @@
-+ 
-+ TEST(MlirBuilderTest, ConstantF8E4M3B11FNUZArray) {
-+   std::string expected = R"mlir(module {
-+-  func.func @main() -> tensor<10xf8E4M3B11FNUZ> {
-+-    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 0x80, 0x80, 1.000000e+00, 0x80]> : tensor<10xf8E4M3B11FNUZ>
-+-    return %cst : tensor<10xf8E4M3B11FNUZ>
-++  func.func @main() -> tensor<11xf8E4M3B11FNUZ> {
-++    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 3.000000e+01, -3.000000e+01, 0x80, 1.220700e-04, -1.220700e-04]> : tensor<11xf8E4M3B11FNUZ>
-++    return %cst : tensor<11xf8E4M3B11FNUZ>
-+   }
-+ })mlir";
-+ 
-+@@ -1127,11 +1224,17 @@
-+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-+     auto type =
-+-        makeTensorType(fb.getContext(), {10}, ElementType::F8E4M3B11FNUZ);
-++        makeTensorType(fb.getContext(), {11}, ElementType::F8E4M3B11FNUZ);
-++    double posMax = 30.0;
-++    double negMax = -30.0;
-++    double nan = std::numeric_limits<double>::quiet_NaN();
-++    double posMinDenorm = 0x1p-13;
-++    double negMinDenorm = -0x1p-13;
-+     auto cst = stablehlo::Constant(
-+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
-+-                                           0x7F, 0xFF, 0x01, 0x81}),
-+-                         type));
-++        fb, makeConstant(
-++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
-++                                  negMax, nan, posMinDenorm, negMinDenorm}),
-++                type));
-+     func::Return(fb, {cst});
-+   }
-+ 
-+@@ -1141,9 +1244,9 @@
-+ 
-+ TEST(MlirBuilderTest, ConstantF8E4M3Array) {
-+   std::string expected = R"mlir(module {
-+-  func.func @main() -> tensor<10xf8E4M3> {
-+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 1.280000e+02, 0x78, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E4M3>
-+-    return %cst : tensor<10xf8E4M3>
-++  func.func @main() -> tensor<11xf8E4M3> {
-++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 0x78, 0xF8, 0x7C, 1.953130e-03, -1.953130e-03]> : tensor<11xf8E4M3>
-++    return %cst : tensor<11xf8E4M3>
-+   }
-+ })mlir";
-+ 
-+@@ -1151,11 +1254,17 @@
-+   {
-+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-+-    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E4M3);
-++    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E4M3);
-++    double posInf = std::numeric_limits<double>::infinity();
-++    double negInf = -std::numeric_limits<double>::infinity();
-++    double nan = std::numeric_limits<double>::quiet_NaN();
-++    double posMinDenorm = 0x1p-9;
-++    double negMinDenorm = -0x1p-9;
-+     auto cst = stablehlo::Constant(
-+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
-+-                                           0x7F, 0xFF, 0x01, 0x81}),
-+-                         type));
-++        fb, makeConstant(
-++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posInf,
-++                                  negInf, nan, posMinDenorm, negMinDenorm}),
-++                type));
-+     func::Return(fb, {cst});
-+   }
-+ 
-+@@ -1165,9 +1274,9 @@
-+ 
-+ TEST(MlirBuilderTest, ConstantF8E4M3FNArray) {
-+   std::string expected = R"mlir(module {
-+-  func.func @main() -> tensor<10xf8E4M3FN> {
-+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 1.280000e+02, 2.560000e+02, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E4M3FN>
-+-    return %cst : tensor<10xf8E4M3FN>
-++  func.func @main() -> tensor<11xf8E4M3FN> {
-++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 4.480000e+02, -4.480000e+02, 0x7F, 1.953130e-03, -1.953130e-03]> : tensor<11xf8E4M3FN>
-++    return %cst : tensor<11xf8E4M3FN>
-+   }
-+ })mlir";
-+ 
-+@@ -1175,11 +1284,17 @@
-+   {
-+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-+-    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E4M3FN);
-++    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E4M3FN);
-++    double posMax = 448.0;
-++    double negMax = -448.0;
-++    double nan = std::numeric_limits<double>::quiet_NaN();
-++    double posMinDenorm = 0x1p-9;
-++    double negMinDenorm = -0x1p-9;
-+     auto cst = stablehlo::Constant(
-+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
-+-                                           0x7F, 0xFF, 0x01, 0x81}),
-+-                         type));
-++        fb, makeConstant(
-++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
-++                                  negMax, nan, posMinDenorm, negMinDenorm}),
-++                type));
-+     func::Return(fb, {cst});
-+   }
-+ 
-+@@ -1189,9 +1304,9 @@
-+ 
-+ TEST(MlirBuilderTest, ConstantF8E4M3FNUZArray) {
-+   std::string expected = R"mlir(module {
-+-  func.func @main() -> tensor<10xf8E4M3FNUZ> {
-+-    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 1.280000e+02, 0x80, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E4M3FNUZ>
-+-    return %cst : tensor<10xf8E4M3FNUZ>
-++  func.func @main() -> tensor<11xf8E4M3FNUZ> {
-++    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 1.015630e-01, 3.250000e+00, 2.400000e+02, -2.400000e+02, 0x80, 9.765620e-04, -9.765620e-04]> : tensor<11xf8E4M3FNUZ>
-++    return %cst : tensor<11xf8E4M3FNUZ>
-+   }
-+ })mlir";
-+ 
-+@@ -1199,11 +1314,17 @@
-+   {
-+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-+-    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E4M3FNUZ);
-++    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E4M3FNUZ);
-++    double posMax = 240.0;
-++    double negMax = -240.0;
-++    double nan = std::numeric_limits<double>::quiet_NaN();
-++    double posMinDenorm = 0x1p-10;
-++    double negMinDenorm = -0x1p-10;
-+     auto cst = stablehlo::Constant(
-+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
-+-                                           0x7F, 0xFF, 0x01, 0x81}),
-+-                         type));
-++        fb, makeConstant(
-++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
-++                                  negMax, nan, posMinDenorm, negMinDenorm}),
-++                type));
-+     func::Return(fb, {cst});
-+   }
-+ 
-+@@ -1213,9 +1334,9 @@
-+ 
-+ TEST(MlirBuilderTest, ConstantF8E5M2Array) {
-+   std::string expected = R"mlir(module {
-+-  func.func @main() -> tensor<10xf8E5M2> {
-+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.000000e+00, 1.280000e+02, 2.560000e+02, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E5M2>
-+-    return %cst : tensor<10xf8E5M2>
-++  func.func @main() -> tensor<11xf8E5M2> {
-++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.000000e+00, 0x7C, 0xFC, 0x7E, 1.525880e-05, -1.525880e-05]> : tensor<11xf8E5M2>
-++    return %cst : tensor<11xf8E5M2>
-+   }
-+ })mlir";
-+ 
-+@@ -1223,11 +1344,17 @@
-+   {
-+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-+-    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E5M2);
-++    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E5M2);
-++    double posInf = std::numeric_limits<double>::infinity();
-++    double negInf = -std::numeric_limits<double>::infinity();
-++    double nan = std::numeric_limits<double>::quiet_NaN();
-++    double posMinDenorm = 0x1p-16;
-++    double negMinDenorm = -0x1p-16;
-+     auto cst = stablehlo::Constant(
-+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
-+-                                           0x7F, 0xFF, 0x01, 0x81}),
-+-                         type));
-++        fb, makeConstant(
-++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posInf,
-++                                  negInf, nan, posMinDenorm, negMinDenorm}),
-++                type));
-+     func::Return(fb, {cst});
-+   }
-+ 
-+@@ -1237,9 +1364,9 @@
-+ 
-+ TEST(MlirBuilderTest, ConstantF8E5M2FNUZArray) {
-+   std::string expected = R"mlir(module {
-+-  func.func @main() -> tensor<10xf8E5M2FNUZ> {
-+-    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.000000e+00, 1.280000e+02, 2.560000e+02, 1.000000e+00, 1.280000e+02]> : tensor<10xf8E5M2FNUZ>
-+-    return %cst : tensor<10xf8E5M2FNUZ>
-++  func.func @main() -> tensor<11xf8E5M2FNUZ> {
-++    %cst = stablehlo.constant dense<[0.000000e+00, 0.000000e+00, 1.000000e+00, 1.250000e-01, 9.375000e-02, 3.000000e+00, 5.734400e+04, -5.734400e+04, 0x80, 7.629390e-06, -7.629390e-06]> : tensor<11xf8E5M2FNUZ>
-++    return %cst : tensor<11xf8E5M2FNUZ>
-+   }
-+ })mlir";
-+ 
-+@@ -1247,11 +1374,17 @@
-+   {
-+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-+-    auto type = makeTensorType(fb.getContext(), {10}, ElementType::F8E5M2FNUZ);
-++    auto type = makeTensorType(fb.getContext(), {11}, ElementType::F8E5M2FNUZ);
-++    double posMax = 57344.0;
-++    double negMax = -57344.0;
-++    double nan = std::numeric_limits<double>::quiet_NaN();
-++    double posMinDenorm = 0x1p-17;
-++    double negMinDenorm = -0x1p-17;
-+     auto cst = stablehlo::Constant(
-+-        fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415,
-+-                                           0x7F, 0xFF, 0x01, 0x81}),
-+-                         type));
-++        fb, makeConstant(
-++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.1415, posMax,
-++                                  negMax, nan, posMinDenorm, negMinDenorm}),
-++                type));
-+     func::Return(fb, {cst});
-+   }
-+ 
-+@@ -1262,7 +1395,7 @@
-+ TEST(MlirBuilderTest, ConstantF8E8M0FNUArray) {
-+   std::string expected = R"mlir(module {
-+   func.func @main() -> tensor<8xf8E8M0FNU> {
-+-    %cst = stablehlo.constant dense<[5.877470e-39, 1.000000e+00, 1.250000e-01, 1.250000e-01, 4.000000e+00, 5.877470e-39, 1.280000e+02, 2.560000e+02]> : tensor<8xf8E8M0FNU>
-++    %cst = stablehlo.constant dense<[5.877470e-39, 1.000000e+00, 1.250000e-01, 1.250000e-01, 4.000000e+00, 5.877470e-39, 2.000000e+00, 0xFF]> : tensor<8xf8E8M0FNU>
-+     return %cst : tensor<8xf8E8M0FNU>
-+   }
-+ })mlir";
-+@@ -1272,10 +1405,12 @@
-+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-+     auto type = makeTensorType(fb.getContext(), {8}, ElementType::F8E8M0FNU);
-++    double min = 0x1p-127;
-++    double nan = std::numeric_limits<double>::quiet_NaN();
-+     auto cst = stablehlo::Constant(
-+-        fb, makeConstant(ArrayRef<double>(
-+-                             {0.0, 1.0, 0.125, 0.1, 3.1415, 0x00, 0x80, 0xFF}),
-+-                         type));
-++        fb, makeConstant(
-++                ArrayRef<double>({0.0, 1.0, 0.125, 0.1, 3.1415, min, 2.0, nan}),
-++                type));
-+     func::Return(fb, {cst});
-+   }
-+ 
-+@@ -1286,7 +1421,7 @@
-+ TEST(MlirBuilderTest, ConstantBF16Array) {
-+   std::string expected = R"mlir(module {
-+   func.func @main() -> tensor<11xbf16> {
-+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000980e-01, 3.140630e+00, 3.264000e+04, 6.553600e+04, 3.276800e+04, 1.000000e+00, 3.276800e+04]> : tensor<11xbf16>
-++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000980e-01, 3.140630e+00, 0x7F80, 0xFF80, 0x7FC0, 9.183550e-41, -9.183550e-41]> : tensor<11xbf16>
-+     return %cst : tensor<11xbf16>
-+   }
-+ })mlir";
-+@@ -1296,11 +1431,16 @@
-+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-+     auto type = makeTensorType(fb.getContext(), {11}, ElementType::BF16);
-++    double posInf = std::numeric_limits<double>::infinity();
-++    double negInf = -std::numeric_limits<double>::infinity();
-++    double nan = std::numeric_limits<double>::quiet_NaN();
-++    double posMinDenorm = 0x1p-133;
-++    double negMinDenorm = -0x1p-133;
-+     auto cst = stablehlo::Constant(
-+-        fb,
-+-        makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.140630,
-+-                                       0x7F80, 0xFF80, 0x7FFF, 0x0001, 0x8001}),
-+-                     type));
-++        fb, makeConstant(
-++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.140630, posInf,
-++                                  negInf, nan, posMinDenorm, negMinDenorm}),
-++                type));
-+     func::Return(fb, {cst});
-+   }
-+ 
-+@@ -1311,7 +1451,7 @@
-+ TEST(MlirBuilderTest, ConstantF16Array) {
-+   std::string expected = R"mlir(module {
-+   func.func @main() -> tensor<11xf16> {
-+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.997550e-02, 3.140630e+00, 3.174400e+04, 6.451200e+04, 3.276800e+04, 1.000000e+00, 3.276800e+04]> : tensor<11xf16>
-++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 9.997550e-02, 3.140630e+00, 0x7C00, 0xFC00, 0x7E00, 5.960460e-08, -5.960460e-08]> : tensor<11xf16>
-+     return %cst : tensor<11xf16>
-+   }
-+ })mlir";
-+@@ -1321,11 +1461,16 @@
-+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-+     auto type = makeTensorType(fb.getContext(), {11}, ElementType::F16);
-++    double posInf = std::numeric_limits<double>::infinity();
-++    double negInf = -std::numeric_limits<double>::infinity();
-++    double nan = std::numeric_limits<double>::quiet_NaN();
-++    double posMinDenorm = 0x1p-24;
-++    double negMinDenorm = -0x1p-24;
-+     auto cst = stablehlo::Constant(
-+-        fb,
-+-        makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.140630,
-+-                                       0x7C00, 0xFC00, 0x7FFF, 0x0001, 0x8001}),
-+-                     type));
-++        fb, makeConstant(
-++                ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1, 3.140630, posInf,
-++                                  negInf, nan, posMinDenorm, negMinDenorm}),
-++                type));
-+     func::Return(fb, {cst});
-+   }
-+ 
-+@@ -1336,7 +1481,7 @@
-+ TEST(MlirBuilderTest, ConstantF32Array) {
-+   std::string expected = R"mlir(module {
-+   func.func @main() -> tensor<11xf32> {
-+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000000e-01, 3.14159274, 2.13909504E+9, 4.28657869E+9, 2.14748365E+9, 1.000000e+00, 2.14748365E+9]> : tensor<11xf32>
-++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000000e-01, 3.14159274, 0x7F800000, 0xFF800000, 0x7FC00000, 1.401300e-45, -1.401300e-45]> : tensor<11xf32>
-+     return %cst : tensor<11xf32>
-+   }
-+ })mlir";
-+@@ -1346,10 +1491,15 @@
-+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-+     auto type = makeTensorType(fb.getContext(), {11}, ElementType::F32);
-++    double posInf = std::numeric_limits<double>::infinity();
-++    double negInf = -std::numeric_limits<double>::infinity();
-++    double nan = std::numeric_limits<double>::quiet_NaN();
-++    double posMinDenorm = std::numeric_limits<float>::denorm_min();
-++    double negMinDenorm = -std::numeric_limits<float>::denorm_min();
-+     auto cst = stablehlo::Constant(
-+         fb, makeConstant(ArrayRef<double>({0.0, -0.0, 1.0, 0.125, 0.1,
-+-                                           3.14159274, 0x7F800000, 0xFF800000,
-+-                                           0x7FFFFFFF, 0x00000001, 0x80000001}),
-++                                           3.14159274, posInf, negInf, nan,
-++                                           posMinDenorm, negMinDenorm}),
-+                          type));
-+     func::Return(fb, {cst});
-+   }
-+@@ -1361,7 +1511,7 @@
-+ TEST(MlirBuilderTest, ConstantF64Array) {
-+   std::string expected = R"mlir(module {
-+   func.func @main() -> tensor<11xf64> {
-+-    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000000e-01, 3.1415926535897931, 9.2188684372274053E+18, 1.8442240474082181E+19, 9.2233720368547758E+18, 1.000000e+00, 9.2233720368547758E+18]> : tensor<11xf64>
-++    %cst = stablehlo.constant dense<[0.000000e+00, -0.000000e+00, 1.000000e+00, 1.250000e-01, 1.000000e-01, 3.1415926535897931, 0x7FF0000000000000, 0xFFF0000000000000, 0x7FF8000000000000, 4.940660e-324, -4.940660e-324]> : tensor<11xf64>
-+     return %cst : tensor<11xf64>
-+   }
-+ })mlir";
-+@@ -1371,13 +1521,16 @@
-+     Location funcLoc = fileLineColLoc(mb->getContext(), "main.mlir", 1, 1);
-+     func::FunctionBuilder fb(mb.get(), "main", funcLoc);
-+     auto type = makeTensorType(fb.getContext(), {11}, ElementType::F64);
-++    double posInf = std::numeric_limits<double>::infinity();
-++    double negInf = -std::numeric_limits<double>::infinity();
-++    double nan = std::numeric_limits<double>::quiet_NaN();
-++    double posMinDenorm = std::numeric_limits<double>::denorm_min();
-++    double negMinDenorm = -std::numeric_limits<double>::denorm_min();
-+     auto cst = stablehlo::Constant(
-+-        fb, makeConstant(
-+-                ArrayRef<long double>({0.0, -0.0, 1.0, 0.125, 0.1,
-+-                                       3.1415926535897931, 0x7FF0000000000000,
-+-                                       0xFFF0000000000000, 0x7FFFFFFFFFFFFFFF,
-+-                                       0x0000000000000001, 0x8000000000000001}),
-+-                type));
-++        fb, makeConstant(ArrayRef<long double>(
-++                             {0.0, -0.0, 1.0, 0.125, 0.1, 3.1415926535897931,
-++                              posInf, negInf, nan, posMinDenorm, negMinDenorm}),
-++                         type));
-+     func::Return(fb, {cst});
-+   }
-+ 
- diff --ruN a/stablehlo/stablehlo/integrations/python/StablehloApi.cpp b/stablehlo/stablehlo/integrations/python/StablehloApi.cpp
- --- stablehlo/stablehlo/integrations/python/StablehloApi.cpp
- +++ stablehlo/stablehlo/integrations/python/StablehloApi.cpp
-@@ -4182,6 +3942,77 @@ diff --ruN a/stablehlo/stablehlo/reference/Ops.cpp b/stablehlo/stablehlo/referen
-      Axes resultPermutation;
-      resultPermutation.push_back(outputBatchDimension);
-      resultPermutation.append(outputSpatialDimensions.begin(),
-+@@ -1942,23 +1969,30 @@
-+ Tensor __fft_1d(const Tensor& operand, ShapedType resultType,
-+                 int64_t dimension) {
-+   Tensor result(resultType);
-+-
-++  auto complexTy = cast<ComplexType>(resultType.getElementType());
-++  const llvm::fltSemantics& sem =
-++      cast<FloatType>(complexTy.getElementType()).getFloatSemantics();
-+   for (auto resultIt = result.index_begin(); resultIt != result.index_end();
-+        ++resultIt) {
-+     auto resultIndex = *resultIt;
-+     auto N = operand.getShape()[dimension];
-+ 
-+-    Element X(resultType.getElementType(),
-+-              mlir::Complex<APFloat>(APFloat(0.0), APFloat(0.0)));
-++    Element X(
-++        resultType.getElementType(),
-++        mlir::Complex<APFloat>(APFloat::getZero(sem), APFloat::getZero(sem)));
-+ 
-+     for (int n = 0; n < N; n++) {
-+       auto operandIndex = Index(resultIndex);
-+       operandIndex[dimension] = n;
-+ 
-+-      auto floatType =
-+-          cast<ComplexType>(operand.getElementType()).getElementType();
-+-      Element phase(floatType,
-+-                    APFloat(-2 * M_PI * n * resultIndex[dimension] / N));
-++      auto floatType = cast<FloatType>(
-++          cast<ComplexType>(operand.getElementType()).getElementType());
-++
-++      auto phaseVal = APFloat(-2 * M_PI * n * resultIndex[dimension] / N);
-++      bool losesInfo;
-++      phaseVal.convert(floatType.getFloatSemantics(),
-++                       APFloat::rmNearestTiesToEven, &losesInfo);
-++      Element phase(floatType, phaseVal);
-+       auto rotator = complex(cosine(phase), sine(phase));
-+       X = X + operand.get(operandIndex) * rotator;
-+     }
-+@@ -1992,10 +2026,12 @@
-+   } else if (fftType == FftType::IFFT) {
-+     // compute IFFT as conjugation, FFT, conjugation and normalization
-+     Tensor result(operand);
-++    auto complexTy = cast<ComplexType>(resultType.getElementType());
-++    const llvm::fltSemantics& sem =
-++        cast<FloatType>(complexTy.getElementType()).getFloatSemantics();
-+     for (auto d = fftDims.begin(); d != fftDims.end(); ++d) {
-+       auto divisorValue = mlir::Complex<APFloat>(
-+-          APFloat(static_cast<double>(resultType.getDimSize(*d))),
-+-          APFloat(0.0));
-++          APFloat(sem, resultType.getDimSize(*d)), APFloat::getZero(sem));
-+       auto divisor =
-+           constantOp(mlir::DenseElementsAttr::get(resultType, divisorValue));
-+       result = divideOp(conjugate(__fft_1d(conjugate(result), resultType, *d)),
-+@@ -2040,6 +2076,8 @@
-+     replica = conjugate(reverseOp(replica, Axes{d_irfft}, replica.getType()));
-+     auto complexType =
-+         ComplexType::get(cast<FloatType>(resultType.getElementType()));
-++    const llvm::fltSemantics& sem =
-++        cast<FloatType>(complexType.getElementType()).getFloatSemantics();
-+     result = concatenateOp({result, replica}, Axis{d_irfft},
-+                            resultType.cloneWith(std::nullopt, complexType));
-+ 
-+@@ -2048,7 +2086,7 @@
-+                     resultType);
-+ 
-+     // normalize result of inverse transformation
-+-    auto divisorValue = APFloat(static_cast<double>(renormFactor));
-++    auto divisorValue = APFloat(sem, renormFactor);
-+     auto divisor =
-+         constantOp(mlir::DenseElementsAttr::get(resultType, divisorValue));
-+     return divideOp(result, divisor, resultType);
- diff --ruN a/stablehlo/stablehlo/reference/Ops.h b/stablehlo/stablehlo/reference/Ops.h
- --- stablehlo/stablehlo/reference/Ops.h
- +++ stablehlo/stablehlo/reference/Ops.h
-@@ -4828,6 +4659,249 @@ diff --ruN a/stablehlo/stablehlo/tests/interpret/convolution.mlir b/stablehlo/st
- +    dense<[[[3.0], [4.0]], [[6.0], [8.0]]]> : tensor<2x2x1xf32>
- +  func.return
- +}
-+diff --ruN a/stablehlo/stablehlo/tests/interpret/fft.mlir b/stablehlo/stablehlo/tests/interpret/fft.mlir
-+--- stablehlo/stablehlo/tests/interpret/fft.mlir
-++++ stablehlo/stablehlo/tests/interpret/fft.mlir
-+@@ -233,3 +233,239 @@
-+   > : tensor<2x2x4xf64>
-+   func.return
-+ }
-++
-++// -----
-++
-++func.func @fft_op_test_fft_1d_c64() {
-++  %operand = stablehlo.constant dense<[(1.0, 2.0), (3.0, -4.0)]> : tensor<2xcomplex<f32>>
-++  %result = stablehlo.fft %operand, type = FFT, length = [2] : (tensor<2xcomplex<f32>>) -> tensor<2xcomplex<f32>>
-++  check.expect_almost_eq_const %result, dense<[(4.0, -2.0), (-2.0, 6.0)]> : tensor<2xcomplex<f32>>
-++  func.return
-++}
-++
-++// -----
-++
-++func.func @fft_op_test_fft_2d_c64() {
-++  %operand = stablehlo.constant dense<
-++    [
-++      [(1.0, 2.0), (3.0, -4.0)],
-++      [(-5.0, 6.0), (-7.0, -8.0)]
-++    ]
-++  > : tensor<2x2xcomplex<f32>>
-++  %result = stablehlo.fft %operand, type = FFT, length = [2, 2] : (tensor<2x2xcomplex<f32>>) -> tensor<2x2xcomplex<f32>>
-++  check.expect_almost_eq_const %result, dense<
-++    [
-++      [(-8.0, -4.0), (0.0, 20.0)],
-++      [(16.0, 0.0), (-4.0, -8.0)]
-++    ]
-++  > : tensor<2x2xcomplex<f32>>
-++  func.return
-++}
-++
-++// -----
-++
-++func.func @fft_op_test_fft_3d_c64() {
-++  %operand = stablehlo.constant dense<
-++    [
-++      [
-++        [(1.0, 2.0), (3.0, -4.0)],
-++        [(-5.0, 6.0), (-7.0, -8.0)]
-++      ],[
-++        [(-1.0, -2.0), (-3.0, 4.0)],
-++        [(5.0, -6.0), (7.0, 8.0)]
-++      ]
-++    ]
-++  > : tensor<2x2x2xcomplex<f32>>
-++  %result = stablehlo.fft %operand, type = FFT, length = [2, 2, 2] : (tensor<2x2x2xcomplex<f32>>) -> tensor<2x2x2xcomplex<f32>>
-++  check.expect_almost_eq_const %result, dense<
-++    [
-++      [
-++        [(0.0, 0.0), (0.0, 0.0)],
-++        [(0.0, 0.0), (0.0, 0.0)]
-++      ],[
-++        [(-16.0, -8.0), (0.0, 40.0)],
-++        [(32.0, 0.0), (-8.0, -16.0)]
-++      ]
-++    ]
-++  > : tensor<2x2x2xcomplex<f32>>
-++  func.return
-++}
-++
-++// -----
-++
-++func.func @fft_op_test_ifft_1d_c64() {
-++  %operand = stablehlo.constant dense<[(4.0, -2.0), (-2.0, 6.0)]> : tensor<2xcomplex<f32>>
-++  %result = stablehlo.fft %operand, type = IFFT, length = [2] : (tensor<2xcomplex<f32>>) -> tensor<2xcomplex<f32>>
-++  check.expect_almost_eq_const %result, dense<[(1.0, 2.0), (3.0, -4.0)]> : tensor<2xcomplex<f32>>
-++  func.return
-++}
-++
-++// -----
-++
-++func.func @fft_op_test_ifft_2d_c64() {
-++  %operand = stablehlo.constant dense<
-++    [
-++      [(-8.0, -4.0), (16.0, 0.0)],
-++      [(0.0, 20.0), (-4.0, -8.0)]
-++    ]
-++  > : tensor<2x2xcomplex<f32>>
-++  %result = stablehlo.fft %operand, type = IFFT, length = [2, 2] : (tensor<2x2xcomplex<f32>>) -> tensor<2x2xcomplex<f32>>
-++  check.expect_almost_eq_const %result, dense<
-++    [
-++      [(1.0, 2.0), (-5.0, 6.0)],
-++      [(3.0, -4.0), (-7.0, -8.0)]
-++    ]
-++  > : tensor<2x2xcomplex<f32>>
-++  func.return
-++}
-++
-++// -----
-++
-++func.func @fft_op_test_ifft_3d_c64() {
-++  %operand = stablehlo.constant dense<
-++    [
-++      [
-++        [(0.0, 0.0), (0.0, 0.0)],
-++        [(0.0, 0.0), (0.0, 0.0)]
-++      ],[
-++        [(-16.0, -8.0), (32.0, 0.0)],
-++        [(0.0, 40.0), (-8.0, -16.0)]
-++      ]
-++    ]
-++  > : tensor<2x2x2xcomplex<f32>>
-++  %result = stablehlo.fft %operand, type = IFFT, length = [2, 2, 2] : (tensor<2x2x2xcomplex<f32>>) -> tensor<2x2x2xcomplex<f32>>
-++  check.expect_almost_eq_const %result, dense<
-++    [
-++      [
-++        [(1.0, 2.0), (-5.0, 6.0)],
-++        [(3.0, -4.0), (-7.0, -8.0)]
-++      ],[
-++        [(-1.0, -2.0), (5.0, -6.0)],
-++        [(-3.0, 4.0), (7.0, 8.0)]
-++      ]
-++    ]
-++  > : tensor<2x2x2xcomplex<f32>>
-++  func.return
-++}
-++
-++// -----
-++
-++func.func @fft_op_test_rfft_1d_f32() {
-++  %operand = stablehlo.constant dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
-++  %result = stablehlo.fft %operand, type = RFFT, length = [4] : (tensor<4xf32>) -> tensor<3xcomplex<f32>>
-++  check.expect_almost_eq_const %result, dense<[(10.0, 0.0), (-2.0, 2.0), (-2.0, 0.0)]> : tensor<3xcomplex<f32>>
-++  func.return
-++}
-++
-++// -----
-++
-++func.func @fft_op_test_rfft_2d_f32() {
-++  %operand = stablehlo.constant dense<
-++    [
-++      [1.0, 2.0, 3.0, 4.0],
-++      [5.0, 6.0, 7.0, 8.0],
-++      [9.0, 10.0, 11.0, 12.0],
-++      [13.0, 14.0, 15.0, 16.0]
-++    ]
-++  > : tensor<4x4xf32>
-++  %result = stablehlo.fft %operand, type = RFFT, length = [4, 4] : (tensor<4x4xf32>) -> tensor<4x3xcomplex<f32>>
-++  check.expect_almost_eq_const %result, dense<
-++    [
-++      [(136.0, 0.0), (-8.0, 8.0), (-8.0, 0.0)],
-++      [(-32.0, 32.0), (0.0, 0.0), (0.0, 0.0)],
-++      [(-32.0, 0.0), (0.0, 0.0), (0.0, 0.0)],
-++      [(-32.0, -32.0), (0.0, 0.0), (0.0, 0.0)]
-++    ]
-++  > : tensor<4x3xcomplex<f32>>
-++  func.return
-++}
-++
-++// -----
-++
-++func.func @fft_op_test_rfft_3d_f32() {
-++  %operand = stablehlo.constant dense<
-++    [
-++      [
-++        [1.0, 2.0, 3.0, 4.0],
-++        [5.0, 6.0, 7.0, 8.0]
-++      ],[
-++        [9.0, 10.0, 11.0, 12.0],
-++        [13.0, 14.0, 15.0, 16.0]
-++      ]
-++    ]
-++  > : tensor<2x2x4xf32>
-++  %result = stablehlo.fft %operand, type = RFFT, length = [2, 2, 4] : (tensor<2x2x4xf32>) -> tensor<2x2x3xcomplex<f32>>
-++  check.expect_almost_eq_const %result, dense<
-++    [
-++      [
-++        [(136.0, 0.0), (-8.0, 8.0), (-8.0, 0.0)],
-++        [(-32.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
-++      ],[
-++        [(-64.0, 0.0), (0.0, 0.0), (0.0, 0.0)],
-++        [(0.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
-++      ]
-++    ]
-++  > : tensor<2x2x3xcomplex<f32>>
-++  func.return
-++}
-++
-++// -----
-++
-++func.func @fft_op_test_irfft_1d_c64() {
-++  %operand = stablehlo.constant dense<[(10.0, 0.0), (-2.0, 2.0), (-2.0, 0.0)]> : tensor<3xcomplex<f32>>
-++  %result = stablehlo.fft %operand, type = IRFFT, length = [4] : (tensor<3xcomplex<f32>>) -> tensor<4xf32>
-++  check.expect_almost_eq_const %result, dense<[1.0, 2.0, 3.0, 4.0]> : tensor<4xf32>
-++  func.return
-++}
-++
-++// -----
-++
-++func.func @fft_op_test_irfft_2d_c64() {
-++  %operand = stablehlo.constant dense<
-++    [
-++      [(136.0, 0.0), (-8.0, 8.0), (-8.0, 0.0)],
-++      [(-32.0, 32.0), (0.0, 0.0), (0.0, 0.0)],
-++      [(-32.0, 0.0), (0.0, 0.0), (0.0, 0.0)],
-++      [(-32.0, -32.0), (0.0, 0.0), (0.0, 0.0)]
-++    ]
-++  > : tensor<4x3xcomplex<f32>>
-++  %result = stablehlo.fft %operand, type = IRFFT, length = [4, 4] : (tensor<4x3xcomplex<f32>>) -> tensor<4x4xf32>
-++  check.expect_almost_eq_const %result, dense<
-++    [
-++      [1.0, 2.0, 3.0, 4.0],
-++      [5.0, 6.0, 7.0, 8.0],
-++      [9.0, 10.0, 11.0, 12.0],
-++      [13.0, 14.0, 15.0, 16.0]
-++    ]
-++  > : tensor<4x4xf32>
-++  func.return
-++}
-++
-++// -----
-++
-++func.func @fft_op_test_irfft_3d_c64() {
-++  %operand = stablehlo.constant dense<
-++    [
-++      [
-++        [(136.0, 0.0), (-8.0, 8.0), (-8.0, 0.0)],
-++        [(-32.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
-++      ],[
-++        [(-64.0, 0.0), (0.0, 0.0), (0.0, 0.0)],
-++        [(0.0, 0.0), (0.0, 0.0), (0.0, 0.0)]
-++      ]
-++    ]
-++  > : tensor<2x2x3xcomplex<f32>>
-++  %result = stablehlo.fft %operand, type = IRFFT, length = [2, 2, 4] : (tensor<2x2x3xcomplex<f32>>) -> tensor<2x2x4xf32>
-++  check.expect_almost_eq_const %result, dense<
-++    [
-++      [
-++        [1.0, 2.0, 3.0, 4.0],
-++        [5.0, 6.0, 7.0, 8.0]
-++      ],[
-++        [9.0, 10.0, 11.0, 12.0],
-++        [13.0, 14.0, 15.0, 16.0]
-++      ]
-++    ]
-++  > : tensor<2x2x4xf32>
-++  func.return
-++}
- diff --ruN a/stablehlo/stablehlo/tests/ops_chlo.mlir b/stablehlo/stablehlo/tests/ops_chlo.mlir
- --- stablehlo/stablehlo/tests/ops_chlo.mlir
- +++ stablehlo/stablehlo/tests/ops_chlo.mlir
-@@ -13148,363 +13222,375 @@ diff --ruN a/stablehlo/third_party/BUILD b/stablehlo/third_party/BUILD
- @@ -1 +0,0 @@
- -licenses(["notice"])  # Apache 2.0
- \ No newline at end of file
--diff --ruN a/stablehlo/third_party/repo.bzl b/stablehlo/third_party/repo.bzl
----- stablehlo/third_party/repo.bzl
--+++ stablehlo/third_party/repo.bzl
--@@ -1,159 +0,0 @@
---# Copyright 2017 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.
---
---"""Utilities for defining TensorFlow Bazel dependencies."""
---
---def tf_mirror_urls(url):
---    """A helper for generating TF-mirror versions of URLs.
---
---    Given a URL, it returns a list of the TF-mirror cache version of that URL
---    and the original URL, suitable for use in `urls` field of `tf_http_archive`.
---    """
---    if not url.startswith("https://"):
---        return [url]
---    return [
---        "https://storage.googleapis.com/mirror.tensorflow.org/%s" % url[8:],
---        url,
---    ]
---
---def _get_env_var(ctx, name):
---    if name in ctx.os.environ:
---        return ctx.os.environ[name]
---    else:
---        return None
---
---# Checks if we should use the system lib instead of the bundled one
---def _use_system_lib(ctx, name):
---    syslibenv = _get_env_var(ctx, "TF_SYSTEM_LIBS")
---    if not syslibenv:
---        return False
---    return name in [n.strip() for n in syslibenv.split(",")]
---
---def _get_link_dict(ctx, link_files, build_file):
---    link_dict = {ctx.path(v): ctx.path(Label(k)) for k, v in link_files.items()}
---    if build_file:
---        # Use BUILD.bazel because it takes precedence over BUILD.
---        link_dict[ctx.path("BUILD.bazel")] = ctx.path(Label(build_file))
---    return link_dict
---
---def _tf_http_archive_impl(ctx):
---    # Construct all paths early on to prevent rule restart. We want the
---    # attributes to be strings instead of labels because they refer to files
---    # in the TensorFlow repository, not files in repos depending on TensorFlow.
---    # See also https://github.com/bazelbuild/bazel/issues/10515.
---    link_dict = _get_link_dict(ctx, ctx.attr.link_files, ctx.attr.build_file)
---
---    # For some reason, we need to "resolve" labels once before the
---    # download_and_extract otherwise it'll invalidate and re-download the
---    # archive each time.
---    # https://github.com/bazelbuild/bazel/issues/10515
---    patch_files = ctx.attr.patch_file
---    for patch_file in patch_files:
---        if patch_file:
---            ctx.path(Label(patch_file))
---
---    if _use_system_lib(ctx, ctx.attr.name):
---        link_dict.update(_get_link_dict(
---            ctx = ctx,
---            link_files = ctx.attr.system_link_files,
---            build_file = ctx.attr.system_build_file,
---        ))
---    else:
---        ctx.download_and_extract(
---            url = ctx.attr.urls,
---            sha256 = ctx.attr.sha256,
---            type = ctx.attr.type,
---            stripPrefix = ctx.attr.strip_prefix,
---        )
---        if hasattr(ctx.attr, "patch_cmds") and ctx.attr.patch_cmds:
---            for cmd in ctx.attr.patch_cmds:
---                res = ctx.execute(["bash", "-c", cmd])
---                if res.return_code != 0:
---                    fail("patch_cmd failed: %s\n%s" % (cmd, res.stderr))
---        if patch_files:
---            for patch_file in patch_files:
---                patch_file = ctx.path(Label(patch_file)) if patch_file else None
---                if patch_file:
---                    ctx.patch(patch_file, strip = 1)
---
---    for dst, src in link_dict.items():
---        ctx.delete(dst)
---        ctx.symlink(src, dst)
---
---_tf_http_archive = repository_rule(
---    implementation = _tf_http_archive_impl,
---    attrs = {
---        "sha256": attr.string(mandatory = True),
---        "urls": attr.string_list(mandatory = True),
---        "strip_prefix": attr.string(),
---        "type": attr.string(),
---        "patch_file": attr.string_list(),
---        "patch_cmds": attr.string_list(),
---        "build_file": attr.string(),
---        "system_build_file": attr.string(),
---        "link_files": attr.string_dict(),
---        "system_link_files": attr.string_dict(),
---    },
---    environ = ["TF_SYSTEM_LIBS"],
---)
---
---def tf_http_archive(name, sha256, urls, **kwargs):
---    """Downloads and creates Bazel repos for dependencies.
---
---    This is a swappable replacement for both http_archive() and
---    new_http_archive() that offers some additional features. It also helps
---    ensure best practices are followed.
---
---    File arguments are relative to the TensorFlow repository by default. Dependent
---    repositories that use this rule should refer to files either with absolute
---    labels (e.g. '@foo//:bar') or from a label created in their repository (e.g.
---    'str(Label("//:bar"))').
---    """
---    if len(urls) < 2:
---        fail("tf_http_archive(urls) must have redundant URLs.")
---
---    if not any([mirror in urls[0] for mirror in (
---        "mirror.tensorflow.org",
---        "mirror.bazel.build",
---        "storage.googleapis.com",
---    )]):
---        fail("The first entry of tf_http_archive(urls) must be a mirror " +
---             "URL, preferrably mirror.tensorflow.org. Even if you don't have " +
---             "permission to mirror the file, please put the correctly " +
---             "formatted mirror URL there anyway, because someone will come " +
---             "along shortly thereafter and mirror the file.")
---
---    if native.existing_rule(name):
---        return
---
---    _tf_http_archive(
---        name = name,
---        sha256 = sha256,
---        urls = urls,
---        **kwargs
---    )
---
---def _tf_vendored_impl(ctx):
---    ctx.symlink(ctx.path(ctx.attr._root).dirname.get_child(ctx.attr.path), ".")
---
---tf_vendored = repository_rule(
---    implementation = _tf_vendored_impl,
---    doc = "Similar to local_repository, but path is relative to the root of the " +
---          "repository, not the root of the workspace.",
---    attrs = {
---        "_root": attr.label(default = "//:unused"),
---        "path": attr.string(),
---    },
---    local = True,
---)
-+diff --ruN a/stablehlo/third_party/BUILD.bazel b/stablehlo/third_party/BUILD.bazel
-+--- stablehlo/third_party/BUILD.bazel
-++++ stablehlo/third_party/BUILD.bazel
-+@@ -0,0 +1,12 @@
-++package(
-++    # copybara:uncomment default_applicable_licenses = ["//:license"],
-++    licenses = ["notice"],
-++)
-++
-++# copybara:uncomment_begin
-++# filegroup(
-++#     name = "all_files",
-++#     srcs = glob(["**"]),
-++#     visibility = ["//:__subpackages__"],
-++# )
-++# copybara:uncomment_end
- diff --ruN a/stablehlo/third_party/rules_ml_toolchain/BUILD b/stablehlo/third_party/rules_ml_toolchain/BUILD
- --- stablehlo/third_party/rules_ml_toolchain/BUILD
- +++ stablehlo/third_party/rules_ml_toolchain/BUILD
- @@ -1 +0,0 @@
- -exports_files(["fix_strip.patch"])
--diff --ruN a/stablehlo/third_party/rules_ml_toolchain/fix_strip.patch b/stablehlo/third_party/rules_ml_toolchain/fix_strip.patch
----- stablehlo/third_party/rules_ml_toolchain/fix_strip.patch
--+++ stablehlo/third_party/rules_ml_toolchain/fix_strip.patch
--@@ -1,30 +0,0 @@
---diff --git a/cc/impls/linux_x86_64_linux_x86_64/wrappers/BUILD b/cc/impls/linux_x86_64_linux_x86_64/wrappers/BUILD
---index a206127..fa119c10 100644
------ a/cc/impls/linux_x86_64_linux_x86_64/wrappers/BUILD
---+++ b/cc/impls/linux_x86_64_linux_x86_64/wrappers/BUILD
---@@ -24,5 +24,6 @@ filegroup(
---         "idler",
---         "ld",
---         "objcopy",
---+        "strip",
---     ],
---     visibility = ["//visibility:public"],
--- )
---diff --git a/cc/impls/linux_x86_64_linux_x86_64/wrappers/strip b/cc/impls/linux_x86_64_linux_x86_64/wrappers/strip
---new file mode 100755
---index 0000000..e3b0923
------ /dev/null
---+++ b/cc/impls/linux_x86_64_linux_x86_64/wrappers/strip
---@@ -0,0 +1,3 @@
---+#!/bin/bash
---+
---+"$STRIP_PATH" "$@"
---diff --git a/third_party/rules_cc_toolchain/toolchain_config.bzl b/third_party/rules_cc_toolchain/toolchain_config.bzl
---index 7381bbb..23afe63 100644
------ a/third_party/rules_cc_toolchain/toolchain_config.bzl
---+++ b/third_party/rules_cc_toolchain/toolchain_config.bzl
---@@ -58,3 +58,4 @@ ALL_ACTIONS = [
---     ACTION_NAMES.cpp_link_static_library,
---     ACTION_NAMES.clif_match,
---+    ACTION_NAMES.strip,
--- ]
--diff --ruN a/stablehlo/third_party/zlib-ng.BUILD b/stablehlo/third_party/zlib-ng.BUILD
----- stablehlo/third_party/zlib-ng.BUILD
--+++ stablehlo/third_party/zlib-ng.BUILD
--@@ -1,106 +0,0 @@
---# This file is licensed under the Apache License v2.0 with LLVM Exceptions.
---# See https://llvm.org/LICENSE.txt for license information.
---# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
---
---load("@bazel_skylib//rules:copy_file.bzl", "copy_file")
---load("@rules_cc//cc:cc_library.bzl", "cc_library")
---
---package(
---    default_visibility = ["//visibility:public"],
---    # BSD/MIT-like license (for zlib)
---    licenses = ["notice"],
---)
---
---copy_file(
---    # The input template is identical to the CMake output.
---    name = "zconf_gen",
---    src = "zconf.h.in",
---    out = "zconf.h",
---    allow_symlink = True,
---)
---
---cc_library(
---    name = "zlib",
---    srcs = select({
---        "@llvm-project//third-party:llvm_zlib_enabled": [
---            "adler32.c",
---            "adler32_p.h",
---            "chunkset.c",
---            "chunkset_tpl.h",
---            "compare258.c",
---            "compress.c",
---            "crc32.c",
---            "crc32_comb.c",
---            "crc32_comb_tbl.h",
---            "crc32_p.h",
---            "crc32_tbl.h",
---            "deflate.c",
---            "deflate.h",
---            "deflate_fast.c",
---            "deflate_medium.c",
---            "deflate_p.h",
---            "deflate_quick.c",
---            "deflate_slow.c",
---            "fallback_builtins.h",
---            "functable.c",
---            "functable.h",
---            "infback.c",
---            "inffast.c",
---            "inffast.h",
---            "inffixed_tbl.h",
---            "inflate.c",
---            "inflate.h",
---            "inflate_p.h",
---            "inftrees.c",
---            "inftrees.h",
---            "insert_string.c",
---            "insert_string_tpl.h",
---            "match_tpl.h",
---            "trees.c",
---            "trees.h",
---            "trees_emit.h",
---            "trees_tbl.h",
---            "uncompr.c",
---            "zbuild.h",
---            "zendian.h",
---            "zutil.c",
---            "zutil.h",
---            "zutil_p.h",
---        ],
---        "//conditions:default": [],
---    }),
---    hdrs = select({
---        "@llvm-project//third-party:llvm_zlib_enabled": [
---            "zlib.h",
---            ":zconf_gen",
---        ],
---        "//conditions:default": [],
---    }),
---    copts = [
---        "-std=c11",
---        "-DZLIB_COMPAT",
---        "-DWITH_GZFILEOP",
---        "-DWITH_OPTIM",
---        "-DWITH_NEW_STRATEGIES",
---        # For local builds you might want to add "-DWITH_NATIVE_INSTRUCTIONS"
---        # here to improve performance. Native instructions aren't enabled in
---        # the default config for reproducibility.
---    ],
---    defines = select({
---        "@llvm-project//third-party:llvm_zlib_enabled": [
---            "LLVM_ENABLE_ZLIB=1",
---        ],
---        "//conditions:default": [
---            "LLVM_ENABLE_ZLIB=0",
---        ],
---    }),
---    # Clang includes zlib with angled instead of quoted includes, so we need
---    # strip_include_prefix here.
---    strip_include_prefix = ".",
---    visibility = ["//visibility:public"],
---)
---
---alias(
---    name = "zlib-ng",
---    actual = ":zlib",
---)
--diff --ruN a/stablehlo/third_party/zstd.BUILD b/stablehlo/third_party/zstd.BUILD
----- stablehlo/third_party/zstd.BUILD
--+++ stablehlo/third_party/zstd.BUILD
--@@ -1,44 +0,0 @@
---# This file is licensed under the Apache License v2.0 with LLVM Exceptions.
---# See https://llvm.org/LICENSE.txt for license information.
---# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
---load("@rules_cc//cc:cc_library.bzl", "cc_library")
---
---package(
---    default_visibility = ["//visibility:public"],
---    # BSD/MIT-like license (for zstd)
---    licenses = ["notice"],
---)
---
---cc_library(
---    name = "zstd",
---    srcs = select({
---        "@llvm-project//third-party:llvm_zstd_enabled": glob([
---            "lib/common/*.c",
---            "lib/common/*.h",
---            "lib/compress/*.c",
---            "lib/compress/*.h",
---            "lib/decompress/*.c",
---            "lib/decompress/*.h",
---            "lib/decompress/*.S",
---            "lib/dictBuilder/*.c",
---            "lib/dictBuilder/*.h",
---        ]),
---        "//conditions:default": [],
---    }),
---    hdrs = select({
---        "@llvm-project//third-party:llvm_zstd_enabled": [
---            "lib/zdict.h",
---            "lib/zstd.h",
---            "lib/zstd_errors.h",
---        ],
---        "//conditions:default": [],
---    }),
---    defines = select({
---        "@llvm-project//third-party:llvm_zstd_enabled": [
---            "LLVM_ENABLE_ZSTD=1",
---            "ZSTD_MULTITHREAD",
---        ],
---        "//conditions:default": [],
---    }),
---    strip_include_prefix = "lib",
---)
-+diff --ruN a/stablehlo/third_party/rules_ml_toolchain/BUILD.bazel b/stablehlo/third_party/rules_ml_toolchain/BUILD.bazel
-+--- stablehlo/third_party/rules_ml_toolchain/BUILD.bazel
-++++ stablehlo/third_party/rules_ml_toolchain/BUILD.bazel
-+@@ -0,0 +1,14 @@
-++package(
-++    # copybara:uncomment default_applicable_licenses = ["//:license"],
-++    licenses = ["notice"],
-++)
-++
-++# copybara:uncomment_begin
-++# filegroup(
-++#     name = "all_files",
-++#     srcs = glob(["**"]),
-++#     visibility = ["//:__subpackages__"],
-++# )
-++# copybara:uncomment_end
-++
-++exports_files(["fix_strip.patch"])
-+diff --ruN a/stablehlo/third_party/xz.BUILD b/stablehlo/third_party/xz.BUILD
-+--- stablehlo/third_party/xz.BUILD
-++++ stablehlo/third_party/xz.BUILD
-+@@ -0,0 +1,329 @@
-++# Copyright 2026 The StableHLO 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
-++#
-++#      https://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.
-++
-++load("@bazel_skylib//rules:write_file.bzl", "write_file")
-++load("@rules_cc//cc:cc_library.bzl", "cc_library")
-++
-++package(default_visibility = ["//visibility:private"])
-++
-++licenses(["unencumbered"])
-++
-++exports_files(["COPYING"])
-++
-++cc_library(
-++    name = "lzma",
-++    srcs = glob(
-++        include = [
-++            "src/**/*.h",
-++            "src/liblzma/**/*.c",
-++        ],
-++        exclude = [
-++            "src/**/*_tablegen.c",
-++            "src/liblzma/check/crc_clmul_consts_gen.c",
-++            "src/liblzma/check/sha256.c",
-++            "src/liblzma/check/crc*_small.c",
-++            "src/liblzma/common/stream_decoder_mt.c",
-++            "src/liblzma/common/stream_encoder_mt.c",
-++        ],
-++    ) + select({
-++        "@platforms//os:osx": [],
-++        "@platforms//os:ios": [],
-++        "//conditions:default": [
-++            "src/liblzma/check/sha256.c",
-++        ],
-++    }) + select({
-++        "@platforms//os:none": [],
-++        "//conditions:default": [
-++            "src/common/tuklib_cpucores.c",
-++            "src/common/tuklib_physmem.c",
-++            "src/liblzma/common/stream_decoder_mt.c",
-++            "src/liblzma/common/stream_encoder_mt.c",
-++        ],
-++    }),
-++    hdrs = glob(["src/liblzma/api/**/*.h"]),
-++    copts = select({
-++        "@platforms//os:windows": [],
-++        "//conditions:default": ["-std=c99"],
-++    }),
-++    defines = select({
-++        "@platforms//os:windows": ["LZMA_API_STATIC"],
-++        "//conditions:default": [],
-++    }),
-++    implementation_deps = [
-++        ":config",
-++        ":src_common",
-++        ":src_liblzma",
-++        ":src_liblzma_check",
-++        ":src_liblzma_common",
-++        ":src_liblzma_delta",
-++        ":src_liblzma_lz",
-++        ":src_liblzma_lzma",
-++        ":src_liblzma_rangecoder",
-++        ":src_liblzma_simple",
-++    ],
-++    linkopts = select({
-++        "@platforms//os:android": [],
-++        "@platforms//os:windows": [],
-++        "@platforms//os:none": [],
-++        "//conditions:default": ["-lpthread"],
-++    }),
-++    linkstatic = select({
-++        "@platforms//os:windows": True,
-++        "//conditions:default": False,
-++    }),
-++    local_defines = ["HAVE_CONFIG_H"],
-++    strip_include_prefix = "src/liblzma/api",
-++    visibility = ["//visibility:public"],
-++)
-++
-++write_file(
-++    name = "config_gen",
-++    out = "lzma_config/config.h",
-++    content = [
-++        "#define ASSUME_RAM 128",
-++        "#ifndef _WIN32",
-++        "#define HAVE_BSWAP_16 1",
-++        "#define HAVE_BSWAP_32 1",
-++        "#define HAVE_BSWAP_64 1",
-++        "#endif",
-++        "#if defined(__APPLE__) || defined(_WIN32)",
-++        "#undef HAVE_BYTESWAP_H",
-++        "#else",
-++        "#define HAVE_BYTESWAP_H 1",
-++        "#endif",
-++        "#undef HAVE_CAPSICUM",
-++        "#if defined(__APPLE__)",
-++        "#define HAVE_CC_SHA256_CTX 1",
-++        "#define HAVE_CC_SHA256_INIT 1",
-++        "#define HAVE_COMMONCRYPTO_COMMONDIGEST_H 1",
-++        "#endif",
-++        "#define HAVE_CHECK_CRC32 1",
-++        "#define HAVE_CHECK_CRC64 1",
-++        "#define HAVE_CHECK_SHA256 1",
-++        "#if !defined(__APPLE__)",
-++        "#define HAVE_CLOCK_GETTIME 1",
-++        "#define HAVE_DECL_CLOCK_MONOTONIC 1",
-++        "#endif",
-++        "#if defined(__ANDROID__) || defined(__APPLE__) || defined(_WIN32)",
-++        "#define HAVE_DECL_PROGRAM_INVOCATION_NAME 0",
-++        "#else",
-++        "#define HAVE_DECL_PROGRAM_INVOCATION_NAME 1",
-++        "#endif",
-++        "#define HAVE_DECODERS 1",
-++        "#define HAVE_DECODER_ARM 1",
-++        "#define HAVE_DECODER_ARM64 1",
-++        "#define HAVE_DECODER_ARMTHUMB 1",
-++        "#define HAVE_DECODER_DELTA 1",
-++        "#define HAVE_DECODER_IA64 1",
-++        "#define HAVE_DECODER_LZMA1 1",
-++        "#define HAVE_DECODER_LZMA2 1",
-++        "#define HAVE_DECODER_POWERPC 1",
-++        "#define HAVE_DECODER_SPARC 1",
-++        "#define HAVE_DECODER_X86 1",
-++        "#define HAVE_DLFCN_H 1",
-++        "#define HAVE_ENCODERS 1",
-++        "#define HAVE_ENCODER_ARM 1",
-++        "#define HAVE_ENCODER_ARM64 1",
-++        "#define HAVE_ENCODER_ARMTHUMB 1",
-++        "#define HAVE_ENCODER_DELTA 1",
-++        "#define HAVE_ENCODER_IA64 1",
-++        "#define HAVE_ENCODER_LZMA1 1",
-++        "#define HAVE_ENCODER_LZMA2 1",
-++        "#define HAVE_ENCODER_POWERPC 1",
-++        "#define HAVE_ENCODER_SPARC 1",
-++        "#define HAVE_ENCODER_X86 1",
-++        "#define HAVE_FCNTL_H 1",
-++        "#define HAVE_FUTIMENS 1",
-++        "#define HAVE_GETOPT_H 1",
-++        "#define HAVE_GETOPT_LONG 1",
-++        "#define HAVE_INTTYPES_H 1",
-++        "#define HAVE_LIMITS_H 1",
-++        "#define HAVE_MBRTOWC 1",
-++        "#define HAVE_MEMORY_H 1",
-++        "#define HAVE_MF_BT2 1",
-++        "#define HAVE_MF_BT3 1",
-++        "#define HAVE_MF_BT4 1",
-++        "#define HAVE_MF_HC3 1",
-++        "#define HAVE_MF_HC4 1",
-++        "#if !defined(__APPLE__)",
-++        "#define HAVE_POSIX_FADVISE 1",
-++        "#endif",
-++        "#if !defined(__APPLE__) && !defined(__ANDROID__)",
-++        "#define HAVE_PTHREAD_CONDATTR_SETCLOCK 1",
-++        "#endif",
-++        "#define HAVE_PTHREAD_PRIO_INHERIT 1",
-++        "#undef HAVE_SMALL",
-++        "#define HAVE_STDBOOL_H 1",
-++        "#define HAVE_STDINT_H 1",
-++        "#define HAVE_STDLIB_H 1",
-++        "#define HAVE_STRING_H 1",
-++        "#define HAVE_STRUCT_STAT_ST_ATIM_TV_NSEC 1",
-++        "#undef HAVE_SYS_CAPSICUM_H",
-++        "#ifndef _WIN32",
-++        "#define HAVE_SYS_PARAM_H 1",
-++        "#define HAVE_SYS_STAT_H 1",
-++        "#define HAVE_SYS_TIME_H 1",
-++        "#define HAVE_SYS_TYPES_H 1",
-++        "#endif",
-++        "#define HAVE_UINTPTR_T 1",
-++        "#define HAVE_UNISTD_H 1",
-++        "#define HAVE_VISIBILITY 0",
-++        "#ifndef _WIN32",
-++        "#define HAVE_WCWIDTH 1",
-++        "#endif",
-++        "#define HAVE__BOOL 1",
-++        "#undef HAVE__FUTIME",
-++        "#define LT_OBJDIR \".libs/\"",
-++        "#ifdef _WIN32",
-++        "#define MYTHREAD_VISTA 1",
-++        "#else",
-++        "#define MYTHREAD_POSIX 1",
-++        "#endif",
-++        "#define NDEBUG 1",
-++        "#define PACKAGE \"xz\"",
-++        "#define PACKAGE_BUGREPORT \"lasse.collin@tukaani.org\"",
-++        "#define PACKAGE_NAME \"XZ Utils\"",
-++        "#define PACKAGE_STRING \"XZ Utils 5.8.3\"",
-++        "#define PACKAGE_TARNAME \"xz\"",
-++        "#define PACKAGE_URL \"https://tukaani.org/xz/\"",
-++        "#define PACKAGE_VERSION \"5.8.3\"",
-++        "#define SIZEOF_SIZE_T 8",
-++        "#define STDC_HEADERS 1",
-++        "#define TUKLIB_CPUCORES_SYSCONF 1",
-++        "#define TUKLIB_FAST_UNALIGNED_ACCESS 1",
-++        "#if defined(__APPLE__)",
-++        "#undef TUKLIB_PHYSMEM_SYSCONF",
-++        "#else",
-++        "#define TUKLIB_PHYSMEM_SYSCONF 1",
-++        "#endif",
-++        "#ifndef _ALL_SOURCE",
-++        "#define _ALL_SOURCE 1",
-++        "#endif",
-++        "#ifndef _GNU_SOURCE",
-++        "#define _GNU_SOURCE 1",
-++        "#endif",
-++        "#ifndef _POSIX_PTHREAD_SEMANTICS",
-++        "#define _POSIX_PTHREAD_SEMANTICS 1",
-++        "#endif",
-++        "#ifndef _TANDEM_SOURCE",
-++        "#define _TANDEM_SOURCE 1",
-++        "#endif",
-++        "#ifndef __EXTENSIONS__",
-++        "#define __EXTENSIONS__ 1",
-++        "#endif",
-++        "#define VERSION \"5.8.3\"",
-++        "#if defined AC_APPLE_UNIVERSAL_BUILD",
-++        "#if defined __BIG_ENDIAN__",
-++        "#define WORDS_BIGENDIAN 1",
-++        "#endif",
-++        "#endif",
-++        "#ifndef _DARWIN_USE_64_BIT_INODE",
-++        "#define _DARWIN_USE_64_BIT_INODE 1",
-++        "#endif",
-++        "",
-++    ],
-++)
-++
-++cc_library(
-++    name = "config",
-++    hdrs = ["lzma_config/config.h"],
-++    strip_include_prefix = "lzma_config",
-++)
-++
-++cc_library(
-++    name = "src_common",
-++    srcs = select({
-++        "@platforms//os:none": [],
-++        "//conditions:default": [
-++            "src/common/tuklib_exit.c",
-++            "src/common/tuklib_progname.c",
-++        ],
-++    }),
-++    hdrs = glob(["src/common/*.h"]),
-++    defines = select({
-++        "@platforms//os:windows": ["LZMA_API_STATIC"],
-++        "//conditions:default": [],
-++    }),
-++    local_defines = ["HAVE_CONFIG_H"] + select({
-++        "@platforms//os:windows": ["TUKLIB_GETTEXT=0"],
-++        "//conditions:default": [],
-++    }),
-++    strip_include_prefix = "src/common",
-++    deps = [
-++        ":config",
-++        ":src_liblzma",
-++        ":src_liblzma_api",
-++        ":src_liblzma_common",
-++    ],
-++)
-++
-++cc_library(
-++    name = "src_liblzma",
-++    hdrs = glob([
-++        "src/liblzma/common/*.h",
-++        "src/liblzma/lzma/*.h",
-++    ]),
-++    strip_include_prefix = "src/liblzma",
-++)
-++
-++cc_library(
-++    name = "src_liblzma_api",
-++    hdrs = glob(
-++        include = ["src/liblzma/api/**/*.h"],
-++        exclude = ["src/liblzma/api/lzma.h"],
-++    ),
-++    strip_include_prefix = "src/liblzma/api",
-++)
-++
-++cc_library(
-++    name = "src_liblzma_check",
-++    hdrs = glob(["src/liblzma/check/*.h"]),
-++    strip_include_prefix = "src/liblzma/check",
-++)
-++
-++cc_library(
-++    name = "src_liblzma_common",
-++    hdrs = glob(["src/liblzma/common/*.h"]),
-++    strip_include_prefix = "src/liblzma/common",
-++)
-++
-++cc_library(
-++    name = "src_liblzma_delta",
-++    hdrs = glob(["src/liblzma/delta/*.h"]),
-++    strip_include_prefix = "src/liblzma/delta",
-++)
-++
-++cc_library(
-++    name = "src_liblzma_lz",
-++    hdrs = glob(["src/liblzma/lz/*.h"]),
-++    strip_include_prefix = "src/liblzma/lz",
-++)
-++
-++cc_library(
-++    name = "src_liblzma_lzma",
-++    hdrs = glob(["src/liblzma/lzma/*.h"]),
-++    strip_include_prefix = "src/liblzma/lzma",
-++)
-++
-++cc_library(
-++    name = "src_liblzma_rangecoder",
-++    hdrs = glob(["src/liblzma/rangecoder/*.h"]),
-++    strip_include_prefix = "src/liblzma/rangecoder",
-++)
-++
-++cc_library(
-++    name = "src_liblzma_simple",
-++    hdrs = glob(["src/liblzma/simple/*.h"]),
-++    strip_include_prefix = "src/liblzma/simple",
-++)
-diff --git a/third_party/stablehlo/workspace.bzl b/third_party/stablehlo/workspace.bzl
-index 5c603385..7c94cab1 100644
---- a/third_party/stablehlo/workspace.bzl
-+++ b/third_party/stablehlo/workspace.bzl
-@@ -18,10 +18,8 @@
- load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
- 
- def repo():
--    #
-     STABLEHLO_COMMIT = "b9029ac2228a12e176ea20e56f7d3dcb6bf8a019"
-     STABLEHLO_SHA256 = "16af6c8f0baacda15dd72d98dc295a213be455b8fa1899af46e144d6dae670bd"
--    #
- 
-     tf_http_archive(
-         name = "stablehlo",
diff --git a/third_party/xla/third_party/shardy/workspace.bzl b/third_party/xla/third_party/shardy/workspace.bzl
index 152e8c5bffc..cd32b006ffe 100644
--- a/third_party/xla/third_party/shardy/workspace.bzl
+++ b/third_party/xla/third_party/shardy/workspace.bzl
@@ -18,8 +18,8 @@
 load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 
 def repo():
-    SHARDY_COMMIT = "38ab99f4d252c59d491899a694fe8cbadfeccec3"
-    SHARDY_SHA256 = "f8a14581ac86bf170c1aa6814daa000112f9c3a93afa677c28315433e7b2a33f"
+    SHARDY_COMMIT = "c29509bd0c1430beaa7e767c866a131492ab1139"
+    SHARDY_SHA256 = "e67e1c6082302dd7572cc19f2dc53045def9b5eae350ea6d1a41dd236b08d775"
 
     tf_http_archive(
         name = "shardy",

From 2cb3f783da8cac59062b16356bbfeb797b6e0e5e Mon Sep 17 00:00:00 2001
From: Gerardo Carranza <gcarranza@google.com>
Date: Thu, 15 Oct 2026 20:44:17 -0700
Subject: [PATCH 0923/1000] Enforce 16-byte buffer alignment in FlatBuffer
 export.

PiperOrigin-RevId: 1000578242
---
 tensorflow/compiler/mlir/lite/flatbuffer_export.cc        | 8 ++++----
 .../mlir/lite/quantization/lite/quantize_model_test.cc    | 8 ++++++++
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/tensorflow/compiler/mlir/lite/flatbuffer_export.cc b/tensorflow/compiler/mlir/lite/flatbuffer_export.cc
index 06232499077..b1965470b7f 100644
--- a/tensorflow/compiler/mlir/lite/flatbuffer_export.cc
+++ b/tensorflow/compiler/mlir/lite/flatbuffer_export.cc
@@ -1584,10 +1584,9 @@ std::optional<BufferOffset<tflite::Buffer>> Translator::BuildBuffer(
                             /*hash=*/0);
 
     std::string buffer_data = buffer.GetData();
-    if (custom_option_alignment_.has_value()) {
-      builder_.ForceVectorAlignment(buffer_data.size(), sizeof(uint8_t),
-                                    custom_option_alignment_.value());
-    }
+    builder_.ForceVectorAlignment(
+        buffer_data.size(), sizeof(uint8_t),
+        custom_option_alignment_.value_or(kFbAlignment));
     ReleaseBufferData(buffer);
     return tflite::CreateBuffer(
         builder_, builder_.CreateVector(
@@ -3969,6 +3968,7 @@ std::optional<BufferOffset<tflite::SubGraph>> Translator::BuildSubGraph(
 BufferOffset<tflite::Metadata> Translator::BuildMetadata(StringRef name,
                                                          StringRef content) {
   auto buffer_index = buffers_.size();
+  builder_.ForceVectorAlignment(content.size(), sizeof(uint8_t), kFbAlignment);
   auto buffer_data = builder_.CreateVector(
       reinterpret_cast<const uint8_t*>(content.data()), content.size());
   buffers_.push_back(tflite::CreateBuffer(builder_, buffer_data));
diff --git a/tensorflow/compiler/mlir/lite/quantization/lite/quantize_model_test.cc b/tensorflow/compiler/mlir/lite/quantization/lite/quantize_model_test.cc
index 80371ebe694..708e25c4aa0 100644
--- a/tensorflow/compiler/mlir/lite/quantization/lite/quantize_model_test.cc
+++ b/tensorflow/compiler/mlir/lite/quantization/lite/quantize_model_test.cc
@@ -103,6 +103,14 @@ absl::Status QuantizeModel(
 
   auto flatbuffer_model = FlatBufferModelAbslError::BuildFromBuffer(
       output_buffer.data(), output_buffer.size());
+  for (const auto* buffer : *flatbuffer_model->GetModel()->buffers()) {
+    if (buffer != nullptr && buffer->data() != nullptr) {
+      EXPECT_EQ((buffer->data()->Data() -
+                 reinterpret_cast<const uint8_t*>(output_buffer.data())) %
+                    16,
+                0);
+    }
+  }
   *model = UnPackFlatBufferModel(*flatbuffer_model->GetModel());
   return absl::OkStatus();
 }

From 0f316d5b67a7672b28393d780ef05dd9018728ce Mon Sep 17 00:00:00 2001
From: Gerardo Carranza <gcarranza@google.com>
Date: Thu, 15 Oct 2026 21:10:49 -0700
Subject: [PATCH 0924/1000] Ensure natural alignment for read-only tensor
 buffers in `Subgraph::SetTensorParametersReadOnly`

PiperOrigin-RevId: 1000587435
---
 tensorflow/lite/core/subgraph.cc    | 74 +++++++++++++++++++++++++++++
 tensorflow/lite/core/subgraph.h     |  9 ++++
 tensorflow/lite/interpreter_test.cc | 70 +++++++++++++++++++++++++++
 3 files changed, 153 insertions(+)

diff --git a/tensorflow/lite/core/subgraph.cc b/tensorflow/lite/core/subgraph.cc
index 70f2569a07d..0c8bb036d97 100644
--- a/tensorflow/lite/core/subgraph.cc
+++ b/tensorflow/lite/core/subgraph.cc
@@ -23,6 +23,7 @@ limitations under the License.
 #include <cstdlib>
 #include <cstring>
 #include <iterator>
+#include <limits>
 #include <memory>
 #include <numeric>
 #include <string>
@@ -69,6 +70,52 @@ namespace {
 constexpr size_t kExtraNodeCapacity = 10;
 constexpr float kTensorCapacityInc = 1.1f;
 
+size_t GetTypeAlignment(TfLiteType type) {
+  switch (type) {
+    case kTfLiteInt64:
+      return alignof(int64_t);
+    case kTfLiteUInt64:
+      return alignof(uint64_t);
+    case kTfLiteFloat64:
+      return alignof(double);
+    case kTfLiteComplex128:
+      return alignof(TfLiteComplex128);
+    case kTfLiteFloat32:
+      return alignof(float);
+    case kTfLiteInt32:
+      return alignof(int32_t);
+    case kTfLiteUInt32:
+      return alignof(uint32_t);
+    case kTfLiteComplex64:
+      return alignof(TfLiteComplex64);
+    case kTfLiteInt16:
+      return alignof(int16_t);
+    case kTfLiteUInt16:
+      return alignof(uint16_t);
+    case kTfLiteFloat16:
+      return alignof(TfLiteFloat16);
+    case kTfLiteBFloat16:
+      return alignof(TfLiteBFloat16);
+    case kTfLiteInt8:
+      return alignof(int8_t);
+    case kTfLiteUInt8:
+      return alignof(uint8_t);
+    case kTfLiteBool:
+      return alignof(bool);
+    case kTfLiteInt4:
+    case kTfLiteUInt4:
+    case kTfLiteInt2:
+    case kTfLiteFloat8E4M3FN:
+    case kTfLiteFloat8E5M2:
+    case kTfLiteNoType:
+    case kTfLiteString:
+    case kTfLiteResource:
+    case kTfLiteVariant:
+      return 1;
+  }
+  return 1;
+}
+
 struct TfLiteQuantizationDeleter {
   void operator()(TfLiteQuantization* q) {
     if (q) TfLiteQuantizationFree(q);
@@ -1970,6 +2017,33 @@ TfLiteStatus Subgraph::SetTensorParametersReadOnly(
     TF_LITE_ENSURE(&context_, required_bytes <= bytes);
   }
 
+  if (buffer != nullptr && bytes > 0 &&
+      reinterpret_cast<uintptr_t>(buffer) % GetTypeAlignment(type) != 0) {
+    TFLITE_LOG_PROD_ONCE(
+        tflite::TFLITE_LOG_WARNING,
+        "Read-only tensor #%d (%s) has an unaligned buffer (%p, required "
+        "alignment %zu); copying to an aligned heap buffer. Re-export the "
+        "model to align buffers in the FlatBuffer and enable zero-copy mmap.",
+        tensor_index, name ? name : "<unnamed>", buffer,
+        GetTypeAlignment(type));
+    // The extra bytes are required by XNNPack to avoid out-of-bounds
+    // SIMD reads (matching TfLiteTensorAllocationSize in common.cc).
+    constexpr size_t kExtraAllocationBytes = 16;
+    if (bytes > std::numeric_limits<size_t>::max() - kExtraAllocationBytes) {
+      return kTfLiteError;
+    }
+    std::unique_ptr<char, FreeDeleter> aligned_buffer(
+        static_cast<char*>(malloc(bytes + kExtraAllocationBytes)));
+    TF_LITE_ENSURE(&context_, aligned_buffer != nullptr);
+    memcpy(aligned_buffer.get(), buffer, bytes);
+    memset(aligned_buffer.get() + bytes, 0, kExtraAllocationBytes);
+    buffer = aligned_buffer.get();
+    // Clear `allocation` since `buffer` now resides on the heap rather than
+    // inside the original `MMAPAllocation` region.
+    allocation = nullptr;
+    aligned_ro_buffers_.push_back(std::move(aligned_buffer));
+  }
+
   TfLiteTensor& tensor = context_.tensors[tensor_index];
   if (type == tensor.type &&
       EqualArrayAndTfLiteIntArray(tensor.dims, ndims, dims)) {
diff --git a/tensorflow/lite/core/subgraph.h b/tensorflow/lite/core/subgraph.h
index dc39a72b9b1..82339810d65 100644
--- a/tensorflow/lite/core/subgraph.h
+++ b/tensorflow/lite/core/subgraph.h
@@ -20,6 +20,7 @@ limitations under the License.
 
 #include <atomic>
 #include <cstdint>
+#include <cstdlib>
 #include <functional>
 #include <map>
 #include <memory>
@@ -1281,6 +1282,14 @@ class Subgraph {
   // Maps tensor external buffer ids used in the subgraph to a model-wide
   // identifiers.
   std::unordered_map<size_t, size_t> tensor_external_buffer_ids_;
+
+  struct FreeDeleter {
+    void operator()(char* p) const { free(p); }
+  };
+
+  // Holds aligned copies of read-only tensor buffers whose original backing
+  // memory was not naturally aligned for the tensor's element type.
+  std::vector<std::unique_ptr<char, FreeDeleter>> aligned_ro_buffers_;
 };
 
 }  // namespace tflite
diff --git a/tensorflow/lite/interpreter_test.cc b/tensorflow/lite/interpreter_test.cc
index f3fe906cb37..a22cfc158e7 100644
--- a/tensorflow/lite/interpreter_test.cc
+++ b/tensorflow/lite/interpreter_test.cc
@@ -343,6 +343,76 @@ TEST(BasicInterpreter, CheckAlignment) {
   }
 }
 
+struct DummyAllocationPtr {
+  const void* ptr;
+  template <typename T>
+  operator const T*() const {  // NOLINT(google-explicit-constructor)
+    return reinterpret_cast<const T*>(ptr);
+  }
+};
+
+TEST(BasicInterpreter, CheckReadOnlyAlignment) {
+  struct {
+    TfLiteType type;
+    size_t alignment;
+  } cases[] = {
+      {kTfLiteFloat32, alignof(float)},
+      {kTfLiteInt32, alignof(int32_t)},
+      {kTfLiteUInt32, alignof(uint32_t)},
+      {kTfLiteInt64, alignof(int64_t)},
+      {kTfLiteUInt64, alignof(uint64_t)},
+      {kTfLiteFloat64, alignof(double)},
+      {kTfLiteInt16, alignof(int16_t)},
+      {kTfLiteFloat16, alignof(TfLiteFloat16)},
+  };
+
+  alignas(16) char raw_buffer[64] = {};
+  const int64_t expected_i64[2] = {6, -7};
+  // Place expected_i64 at non-overlapping misaligned offsets: an odd offset
+  // (raw_buffer + 1) and a 4-byte-aligned but not 8-byte-aligned offset
+  // (raw_buffer + 20).
+  memcpy(raw_buffer + 1, expected_i64, sizeof(expected_i64));
+  memcpy(raw_buffer + 20, expected_i64, sizeof(expected_i64));
+
+  const DummyAllocationPtr dummy_alloc{raw_buffer};
+
+  for (const auto& test : cases) {
+    Interpreter interpreter;
+    ASSERT_EQ(interpreter.AddTensors(2), kTfLiteOk);
+    TfLiteQuantizationParams quant = {};
+    ASSERT_EQ(
+        interpreter.SetTensorParametersReadOnly(
+            0, test.type, "", {2}, quant, raw_buffer + 1, 2 * 8, dummy_alloc),
+        kTfLiteOk);
+    ASSERT_EQ(
+        interpreter.SetTensorParametersReadOnly(
+            1, test.type, "", {2}, quant, raw_buffer + 20, 2 * 8, dummy_alloc),
+        kTfLiteOk);
+
+    for (int i = 0; i < 2; ++i) {
+      const TfLiteTensor& tensor = *interpreter.tensor(i);
+      EXPECT_EQ(tensor.allocation_type, kTfLiteMmapRo);
+      EXPECT_EQ(reinterpret_cast<uintptr_t>(tensor.data.raw) % test.alignment,
+                0u);
+      if (test.type == kTfLiteInt64) {
+        EXPECT_EQ(tensor.data.i64[0], expected_i64[0]);
+        EXPECT_EQ(tensor.data.i64[1], expected_i64[1]);
+      }
+    }
+    // Tensor 0 (at odd offset +1) is always realigned for align >= 2, so its
+    // `allocation` pointer must be cleared.
+    EXPECT_EQ(interpreter.tensor(0)->allocation, nullptr);
+    // Tensor 1 (at offset +20, i.e. 4 mod 8) is realigned only when
+    // test.alignment > 4; otherwise it stays zero-copy with `dummy_alloc`.
+    if (test.alignment > 4) {
+      EXPECT_EQ(interpreter.tensor(1)->allocation, nullptr);
+    } else {
+      EXPECT_EQ(interpreter.tensor(1)->data.raw, raw_buffer + 20);
+      EXPECT_EQ(interpreter.tensor(1)->allocation, dummy_alloc.ptr);
+    }
+  }
+}
+
 TEST(BasicInterpreter, CheckArenaAllocation) {
   Interpreter interpreter;
   ASSERT_EQ(interpreter.AddTensors(10), kTfLiteOk);

From d294b75f355b3ac2c35e5658ec2851ae17c0046d Mon Sep 17 00:00:00 2001
From: Maddipatla Chatan <chetan121318@gmail.com>
Date: Mon, 12 Oct 2026 18:29:03 +0530
Subject: [PATCH 0925/1000] Refactor _split function to include stack_fn

---
 tensorflow/dtensor/python/numpy_util.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensorflow/dtensor/python/numpy_util.py b/tensorflow/dtensor/python/numpy_util.py
index f534db6f937..e9c2f46b8f2 100644
--- a/tensorflow/dtensor/python/numpy_util.py
+++ b/tensorflow/dtensor/python/numpy_util.py
@@ -51,7 +51,10 @@ def _split(value, splits, axis=0, split_fn=np.split, stack_fn=np.stack):
 
   if len(splits) > 1:
     splits = splits[1:]
-    children = [_split(child, splits, axis + 1) for child in children]
+    children = [
+        _split(child, splits, axis + 1, split_fn=split_fn, stack_fn=stack_fn)
+        for child in children
+    ]
 
   return stack_fn(children)
 

From 43b78a337557f342fe55b9ab333b5366a104df42 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Tue, 13 Oct 2026 20:53:44 +0530
Subject: [PATCH 0926/1000] Address review: add regression test for multi-axis
 sharded pack_tf_tensor

Adds test_pack_tf_tensor_multi_axis_sharded_under_tf_function to
NumpyUtilTest, verifying pack_tf_tensor() on a layout sharded across
2 axes (x, y) under @tf.function tracing -- the exact path that
previously crashed with NotImplementedError before _split() correctly
forwarded split_fn/stack_fn in its recursive call.
---
 .../dtensor/python/tests/numpy_util_test.py     | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/tensorflow/dtensor/python/tests/numpy_util_test.py b/tensorflow/dtensor/python/tests/numpy_util_test.py
index f7101e0c4d1..09e6d6e3841 100644
--- a/tensorflow/dtensor/python/tests/numpy_util_test.py
+++ b/tensorflow/dtensor/python/tests/numpy_util_test.py
@@ -125,6 +125,23 @@ class NumpyUtilTest(test_util.DTensorBaseTest):
         numpy_util.unpacked_to_numpy(tensors, sharded_on_y_x),
         np.arange(16).reshape(4, 4))
 
+  def test_pack_tf_tensor_multi_axis_sharded_under_tf_function(self):
+    from tensorflow.python.eager import polymorphic_function
+    from tensorflow.python.framework import constant_op
+    from tensorflow.python.framework import dtypes
+
+    sharded_layout = layout.Layout([_MESH_DIM_X, _MESH_DIM_Y], mesh=self.mesh)
+    val = constant_op.constant(
+        np.arange(16).reshape(4, 4), dtype=dtypes.float32
+    )
+
+    @polymorphic_function.function
+    def fn(x):
+      return numpy_util.pack_tf_tensor(x, sharded_layout)
+
+    dtensor_out = fn(val)
+    self.assertAllClose(val, numpy_util.to_numpy(dtensor_out))
+
   def test_unpack_uneven_split_raises(self):
     value = np.arange(5)
 

From 8e4a2bfc4422dc6d4c0fdcef7c58419a5c983c4f Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Tue, 13 Oct 2026 23:04:30 +0530
Subject: [PATCH 0927/1000] Address review: move test-only imports to module
 level

polymorphic_function, constant_op and dtypes were imported inside
test_pack_tf_tensor_multi_axis_sharded_under_tf_function, violating the
no-ad-hoc-in-function-imports guideline. Moved to the module-level
import block alongside the existing dtensor/tensorflow imports.
---
 tensorflow/dtensor/python/tests/numpy_util_test.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/tensorflow/dtensor/python/tests/numpy_util_test.py b/tensorflow/dtensor/python/tests/numpy_util_test.py
index 09e6d6e3841..3ce9a904608 100644
--- a/tensorflow/dtensor/python/tests/numpy_util_test.py
+++ b/tensorflow/dtensor/python/tests/numpy_util_test.py
@@ -23,6 +23,9 @@ from tensorflow.dtensor.python import layout
 from tensorflow.dtensor.python import layout as layout_lib
 from tensorflow.dtensor.python import numpy_util
 from tensorflow.dtensor.python.tests import test_util
+from tensorflow.python.eager import polymorphic_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
 from tensorflow.python.platform import test as tf_test
 # pylint: enable=g-direct-tensorflow-import
 
@@ -126,10 +129,6 @@ class NumpyUtilTest(test_util.DTensorBaseTest):
         np.arange(16).reshape(4, 4))
 
   def test_pack_tf_tensor_multi_axis_sharded_under_tf_function(self):
-    from tensorflow.python.eager import polymorphic_function
-    from tensorflow.python.framework import constant_op
-    from tensorflow.python.framework import dtypes
-
     sharded_layout = layout.Layout([_MESH_DIM_X, _MESH_DIM_Y], mesh=self.mesh)
     val = constant_op.constant(
         np.arange(16).reshape(4, 4), dtype=dtypes.float32

From 351ffdebdc503f73a1d6680cf0f07855a305f36e Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 09:20:08 +0530
Subject: [PATCH 0928/1000] Fix polymorphic_function import and untraceable
 pack() test

- Import polymorphic_function as the module (not the namespace-package
  directory), fixing the AttributeError on @polymorphic_function.function.
- pack()/unpack() DTensor device ops must run eagerly, so replace the
  traced test with one that traces numpy_util.unpack() directly (the
  actual _split recursion fixed by this PR) using TF split_fn/stack_fn,
  and reconstructs the original value via unpacked_to_numpy.
- Add a separate eager test for pack_tf_tensor() on a multi-axis sharded
  layout that asserts the packed DTensor's layout via api.fetch_layout.
- Add the BUILD deps needed for the new imports.
---
 tensorflow/dtensor/python/tests/BUILD         |  6 ++++
 .../dtensor/python/tests/numpy_util_test.py   | 29 ++++++++++++++++---
 2 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/tensorflow/dtensor/python/tests/BUILD b/tensorflow/dtensor/python/tests/BUILD
index cdab09849ba..2fdddcb9ad9 100644
--- a/tensorflow/dtensor/python/tests/BUILD
+++ b/tensorflow/dtensor/python/tests/BUILD
@@ -690,8 +690,14 @@ dtensor_test(
     main = "numpy_util_test.py",
     deps = [
         ":test_util",
+        "//tensorflow/dtensor/python:api",
         "//tensorflow/dtensor/python:layout",
         "//tensorflow/dtensor/python:numpy_util",
+        "//tensorflow/python/eager/polymorphic_function",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:array_ops_stack",
         "//tensorflow/python/platform:client_testlib",
         "//third_party/py/numpy",
     ],
diff --git a/tensorflow/dtensor/python/tests/numpy_util_test.py b/tensorflow/dtensor/python/tests/numpy_util_test.py
index 3ce9a904608..61d67500f6f 100644
--- a/tensorflow/dtensor/python/tests/numpy_util_test.py
+++ b/tensorflow/dtensor/python/tests/numpy_util_test.py
@@ -19,13 +19,16 @@ from unittest import mock
 import numpy as np
 
 # pylint: disable=g-direct-tensorflow-import
+from tensorflow.dtensor.python import api
 from tensorflow.dtensor.python import layout
 from tensorflow.dtensor.python import layout as layout_lib
 from tensorflow.dtensor.python import numpy_util
 from tensorflow.dtensor.python.tests import test_util
-from tensorflow.python.eager import polymorphic_function
+from tensorflow.python.eager.polymorphic_function import polymorphic_function
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
 from tensorflow.python.platform import test as tf_test
 # pylint: enable=g-direct-tensorflow-import
 
@@ -128,7 +131,11 @@ class NumpyUtilTest(test_util.DTensorBaseTest):
         numpy_util.unpacked_to_numpy(tensors, sharded_on_y_x),
         np.arange(16).reshape(4, 4))
 
-  def test_pack_tf_tensor_multi_axis_sharded_under_tf_function(self):
+  def test_unpack_multi_axis_sharded_under_tf_function(self):
+    # `pack`/`unpack` the DTensor device op itself must run eagerly, so this
+    # exercises the traceable part of the previously-broken path: `_split`
+    # recursing over 2+ sharded axes with TF (not NumPy) `split_fn`/
+    # `stack_fn` under `tf.function` tracing.
     sharded_layout = layout.Layout([_MESH_DIM_X, _MESH_DIM_Y], mesh=self.mesh)
     val = constant_op.constant(
         np.arange(16).reshape(4, 4), dtype=dtypes.float32
@@ -136,9 +143,23 @@ class NumpyUtilTest(test_util.DTensorBaseTest):
 
     @polymorphic_function.function
     def fn(x):
-      return numpy_util.pack_tf_tensor(x, sharded_layout)
+      return numpy_util.unpack(
+          x, sharded_layout, split_fn=array_ops.split,
+          stack_fn=array_ops_stack.stack)
 
-    dtensor_out = fn(val)
+    shards = fn(val)
+    self.assertAllClose(
+        numpy_util.unpacked_to_numpy(shards, sharded_layout), val)
+
+  def test_pack_tf_tensor_multi_axis_sharded(self):
+    sharded_layout = layout.Layout([_MESH_DIM_X, _MESH_DIM_Y], mesh=self.mesh)
+    val = constant_op.constant(
+        np.arange(16).reshape(4, 4), dtype=dtypes.float32
+    )
+
+    dtensor_out = numpy_util.pack_tf_tensor(val, sharded_layout)
+
+    self.assertEqual(api.fetch_layout(dtensor_out), sharded_layout)
     self.assertAllClose(val, numpy_util.to_numpy(dtensor_out))
 
   def test_unpack_uneven_split_raises(self):

From a6905d257c96accf2828d91184d652532f1e72c0 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 12:53:12 +0530
Subject: [PATCH 0929/1000] Skip test_pack_tf_tensor_multi_axis_sharded on
 <4-device meshes

The test's 2x2 mesh needs 4 devices, but the internal TPU test
environment only provisions 2 TPU cores by default, so
pack_tf_tensor crashed trying to place onto nonexistent
TPU:2/TPU:3. Skip when fewer than 4 devices are available, per
review.
---
 tensorflow/dtensor/python/tests/numpy_util_test.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tensorflow/dtensor/python/tests/numpy_util_test.py b/tensorflow/dtensor/python/tests/numpy_util_test.py
index 61d67500f6f..ab5c947c132 100644
--- a/tensorflow/dtensor/python/tests/numpy_util_test.py
+++ b/tensorflow/dtensor/python/tests/numpy_util_test.py
@@ -152,6 +152,8 @@ class NumpyUtilTest(test_util.DTensorBaseTest):
         numpy_util.unpacked_to_numpy(shards, sharded_layout), val)
 
   def test_pack_tf_tensor_multi_axis_sharded(self):
+    if len(self.mesh.local_device_ids()) < 4:
+      self.skipTest('Test requires at least 4 devices.')
     sharded_layout = layout.Layout([_MESH_DIM_X, _MESH_DIM_Y], mesh=self.mesh)
     val = constant_op.constant(
         np.arange(16).reshape(4, 4), dtype=dtypes.float32

From 56261f094cec4cebf4197034009ec68db81192ab Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 19:08:24 +0530
Subject: [PATCH 0930/1000] Fix dead device-count assertion in TPU test skip

self.mesh is statically configured with a 2x2 layout in setUp(), so
local_device_ids() always returns 4 elements regardless of the actual
hardware. Query the real runtime device count via
tf_config.list_logical_devices(self.mesh.device_type) instead so the
skip actually protects 2-core TPU jobs.
---
 tensorflow/dtensor/python/tests/BUILD              | 1 +
 tensorflow/dtensor/python/tests/numpy_util_test.py | 3 ++-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/tensorflow/dtensor/python/tests/BUILD b/tensorflow/dtensor/python/tests/BUILD
index 2fdddcb9ad9..063042c3dc5 100644
--- a/tensorflow/dtensor/python/tests/BUILD
+++ b/tensorflow/dtensor/python/tests/BUILD
@@ -694,6 +694,7 @@ dtensor_test(
         "//tensorflow/dtensor/python:layout",
         "//tensorflow/dtensor/python:numpy_util",
         "//tensorflow/python/eager/polymorphic_function",
+        "//tensorflow/python/framework:config",
         "//tensorflow/python/framework:constant_op",
         "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/ops:array_ops",
diff --git a/tensorflow/dtensor/python/tests/numpy_util_test.py b/tensorflow/dtensor/python/tests/numpy_util_test.py
index ab5c947c132..bb634d29e89 100644
--- a/tensorflow/dtensor/python/tests/numpy_util_test.py
+++ b/tensorflow/dtensor/python/tests/numpy_util_test.py
@@ -25,6 +25,7 @@ from tensorflow.dtensor.python import layout as layout_lib
 from tensorflow.dtensor.python import numpy_util
 from tensorflow.dtensor.python.tests import test_util
 from tensorflow.python.eager.polymorphic_function import polymorphic_function
+from tensorflow.python.framework import config as tf_config
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
 from tensorflow.python.ops import array_ops
@@ -152,7 +153,7 @@ class NumpyUtilTest(test_util.DTensorBaseTest):
         numpy_util.unpacked_to_numpy(shards, sharded_layout), val)
 
   def test_pack_tf_tensor_multi_axis_sharded(self):
-    if len(self.mesh.local_device_ids()) < 4:
+    if len(tf_config.list_logical_devices(self.mesh.device_type)) < 4:
       self.skipTest('Test requires at least 4 devices.')
     sharded_layout = layout.Layout([_MESH_DIM_X, _MESH_DIM_Y], mesh=self.mesh)
     val = constant_op.constant(

From fca14dc7345eddb350965ef8bedb35b914157cb2 Mon Sep 17 00:00:00 2001
From: MaddipatlaChetan24 <chetan121318@gmail.com>
Date: Wed, 14 Oct 2026 20:22:22 +0530
Subject: [PATCH 0931/1000] Fix uncalled device_type method in TPU device-count
 skip guard

self.mesh.device_type is a bound method on the pywrap Mesh binding
(confirmed: tensorflow/python/_pywrap_dtensor_device.pyi declares
def device_type(self) -> str, no @property), not a property. Passing
it unevaluated to tf_config.list_logical_devices compares the string
device types against the method object itself, which never matches,
so list_logical_devices always returns an empty list per its own
documented behavior for a non-matching argument. That made
len(...) < 4 unconditionally true, silently skipping
test_pack_tf_tensor_multi_axis_sharded on every platform instead of
only on <4-device ones.

Verified directly against the real config.list_logical_devices: an
unevaluated bound method returns [], while calling it returns the
real device list.
---
 tensorflow/dtensor/python/tests/numpy_util_test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/dtensor/python/tests/numpy_util_test.py b/tensorflow/dtensor/python/tests/numpy_util_test.py
index bb634d29e89..8c9cfe9a4b8 100644
--- a/tensorflow/dtensor/python/tests/numpy_util_test.py
+++ b/tensorflow/dtensor/python/tests/numpy_util_test.py
@@ -153,7 +153,7 @@ class NumpyUtilTest(test_util.DTensorBaseTest):
         numpy_util.unpacked_to_numpy(shards, sharded_layout), val)
 
   def test_pack_tf_tensor_multi_axis_sharded(self):
-    if len(tf_config.list_logical_devices(self.mesh.device_type)) < 4:
+    if len(tf_config.list_logical_devices(self.mesh.device_type())) < 4:
       self.skipTest('Test requires at least 4 devices.')
     sharded_layout = layout.Layout([_MESH_DIM_X, _MESH_DIM_Y], mesh=self.mesh)
     val = constant_op.constant(

From 387fbace123d2b8dfcc9d6be3b05ca136acf1cf4 Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Thu, 15 Oct 2026 21:47:05 -0700
Subject: [PATCH 0932/1000] [XLA] Require orthogonal replica groups in
 AllReduceReduceScatterReorder

## Problem

`AllReduceReduceScatterReorder` rewrites `reduce-scatter(all-reduce(x))` into
`all-reduce(reduce-scatter(x))`. That commutation is only valid when the
`all-reduce` and `reduce-scatter` operate on orthogonal device subgroups with
the same associative reduction kind and collective group mode.

Previously, `IsAllReduceReduceScatter` did not check `replica_groups`,
collective group mode, or operand count, and skipped the reduction kind
comparison whenever `MatchReductionComputation(ar->to_apply())` returned
`std::nullopt`. When both operations use the same `replica_groups` (such as
`{{0,1}}`), reordering scatters different slices to devices 0 and 1 first and
then sums those distinct slices across `{0,1}`, corrupting the output on every
device.

## Change

Require both instructions to have a single operand without control
dependencies, match a recognized `ReductionKind`, share the same collective
group mode, and satisfy `ReplicaGroupsOrthogonal`.

## Tests

- `//xla/hlo/transforms/collectives:all_reduce_reduce_scatter_reorder_test`:
  pins identical replica groups, swapped scatter order, mismatched group sizes,
  unrecognized reduction computations, and variadic `reduce-scatter`.

PiperOrigin-RevId: 1000600132
---
 .../all_reduce_reduce_scatter_reorder.cc      |  46 +++++---
 .../all_reduce_reduce_scatter_reorder_test.cc | 105 +++++++++++++++++-
 2 files changed, 135 insertions(+), 16 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/collectives/all_reduce_reduce_scatter_reorder.cc b/third_party/xla/xla/hlo/transforms/collectives/all_reduce_reduce_scatter_reorder.cc
index be49167091b..f4b987ae41d 100644
--- a/third_party/xla/xla/hlo/transforms/collectives/all_reduce_reduce_scatter_reorder.cc
+++ b/third_party/xla/xla/hlo/transforms/collectives/all_reduce_reduce_scatter_reorder.cc
@@ -28,28 +28,46 @@ limitations under the License.
 namespace xla {
 
 namespace {
-// Returns true if
-// 1. `inst` is an all-reduce instruction.
-// 2. `inst` has a single user.
-// 3. The user is a reduce-scatter instruction.
-// 4. The all-reduce and reduce-scatter have the same reduction type.
+// Returns true if:
+// 1. `ar` is a single-operand all-reduce without control dependencies and with
+//    a single user.
+// 2. The user `rs` is a single-operand reduce-scatter without control
+//    dependencies.
+// 3. `ar` and `rs` share the same recognized reduction kind and collective
+//    group mode.
+// 4. `ar` and `rs` have non-empty, orthogonal replica groups with matching
+//    group sizes.
 bool IsAllReduceReduceScatter(const HloInstruction* ar) {
-  if (ar->opcode() != HloOpcode::kAllReduce) {
-    return false;
-  }
-  if (ar->user_count() != 1) {
+  if (ar->opcode() != HloOpcode::kAllReduce || ar->user_count() != 1 ||
+      ar->operand_count() != 1 || ar->HasControlDependencies()) {
     return false;
   }
   const HloInstruction* rs = ar->users().front();
-  if (rs->opcode() != HloOpcode::kReduceScatter) {
+  const auto ar_reduction = MatchReductionComputation(ar->to_apply());
+  if (rs->opcode() != HloOpcode::kReduceScatter || rs->operand_count() != 1 ||
+      rs->HasControlDependencies() || ar->replica_groups().empty() ||
+      rs->replica_groups().empty() || !ar_reduction.has_value() ||
+      ar_reduction != MatchReductionComputation(rs->to_apply())) {
     return false;
   }
-  if (MatchReductionComputation(ar->to_apply()) &&
-      MatchReductionComputation(ar->to_apply()) !=
-          MatchReductionComputation(rs->to_apply())) {
+  absl::StatusOr<CollectiveOpGroupMode> ar_mode = GetCollectiveOpGroupMode(ar);
+  if (!ar_mode.ok() || ar_mode != GetCollectiveOpGroupMode(rs) ||
+      (*ar_mode == CollectiveOpGroupMode::
+                       COLLECTIVE_OP_GROUP_MODE_CROSS_REPLICA_AND_PARTITION &&
+       ar->GetModule()->config().num_partitions() > 1)) {
     return false;
   }
-  return true;
+  for (const ReplicaGroup& g : ar->replica_groups()) {
+    if (g.replica_ids_size() != rs->replica_groups().size()) {
+      return false;
+    }
+  }
+  for (const ReplicaGroup& g : rs->replica_groups()) {
+    if (g.replica_ids_size() != ar->replica_groups().size()) {
+      return false;
+    }
+  }
+  return ReplicaGroupsOrthogonal(ar->replica_groups(), rs->replica_groups());
 }
 
 // Before: operand -> old_ar -> old_rs -> users
diff --git a/third_party/xla/xla/hlo/transforms/collectives/all_reduce_reduce_scatter_reorder_test.cc b/third_party/xla/xla/hlo/transforms/collectives/all_reduce_reduce_scatter_reorder_test.cc
index 0ed1fc95b6b..baaf2fad46d 100644
--- a/third_party/xla/xla/hlo/transforms/collectives/all_reduce_reduce_scatter_reorder_test.cc
+++ b/third_party/xla/xla/hlo/transforms/collectives/all_reduce_reduce_scatter_reorder_test.cc
@@ -116,8 +116,109 @@ TEST_F(AllReduceReduceScatterReorderTest, DifferentReductionFunctions) {
 
   ENTRY main {
     p0 = f32[8] parameter(0)
-    ar = f32[8] all-reduce(p0), replica_groups={}, to_apply=sum
-    ROOT rs = f32[4] reduce-scatter(ar), dimensions={0}, replica_groups={}, to_apply=product
+    ar = f32[8] all-reduce(p0), replica_groups={{0,1}, {2,3}}, to_apply=sum
+    ROOT rs = f32[4] reduce-scatter(ar), dimensions={0}, replica_groups={{0,2}, {1,3}}, to_apply=product
+  }
+  )";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_text));
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass_, module.get()));
+  EXPECT_FALSE(changed);
+}
+
+TEST_F(AllReduceReduceScatterReorderTest, NonOrthogonalReplicaGroups) {
+  absl::string_view hlo_text = R"(
+  sum {
+    a = f32[] parameter(0)
+    b = f32[] parameter(1)
+    ROOT sum = f32[] add(a, b)
+  }
+
+  ENTRY main {
+    p0 = f32[8] parameter(0)
+    ar = f32[8] all-reduce(p0), replica_groups={{0,1}, {2,3}}, to_apply=sum
+    ROOT rs = f32[4] reduce-scatter(ar), dimensions={0}, replica_groups={{0,1}, {2,3}}, to_apply=sum
+  }
+  )";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_text));
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass_, module.get()));
+  EXPECT_FALSE(changed);
+}
+
+TEST_F(AllReduceReduceScatterReorderTest, NonOrthogonalSwappedScatterOrder) {
+  absl::string_view hlo_text = R"(
+  sum {
+    a = f32[] parameter(0)
+    b = f32[] parameter(1)
+    ROOT sum = f32[] add(a, b)
+  }
+
+  ENTRY main {
+    p0 = f32[8] parameter(0)
+    ar = f32[8] all-reduce(p0), replica_groups={{0,1}, {2,3}}, to_apply=sum
+    ROOT rs = f32[4] reduce-scatter(ar), dimensions={0}, replica_groups={{0,2}, {3,1}}, to_apply=sum
+  }
+  )";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_text));
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass_, module.get()));
+  EXPECT_FALSE(changed);
+}
+
+TEST_F(AllReduceReduceScatterReorderTest, UnrecognizedReductionComputation) {
+  absl::string_view hlo_text = R"(
+  sub {
+    a = f32[] parameter(0)
+    b = f32[] parameter(1)
+    ROOT diff = f32[] subtract(a, b)
+  }
+
+  ENTRY main {
+    p0 = f32[8] parameter(0)
+    ar = f32[8] all-reduce(p0), replica_groups={{0,1}, {2,3}}, to_apply=sub
+    ROOT rs = f32[4] reduce-scatter(ar), dimensions={0}, replica_groups={{0,2}, {1,3}}, to_apply=sub
+  }
+  )";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_text));
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass_, module.get()));
+  EXPECT_FALSE(changed);
+}
+
+TEST_F(AllReduceReduceScatterReorderTest, VariadicReduceScatter) {
+  absl::string_view hlo_text = R"(
+  sum {
+    a = f32[] parameter(0)
+    b = f32[] parameter(1)
+    ROOT sum = f32[] add(a, b)
+  }
+
+  ENTRY main {
+    p0 = f32[8] parameter(0)
+    p1 = f32[8] parameter(1)
+    ar = f32[8] all-reduce(p0), replica_groups={{0,1}, {2,3}}, to_apply=sum
+    ROOT rs = (f32[4], f32[4]) reduce-scatter(ar, p1), dimensions={0}, replica_groups={{0,2}, {1,3}}, to_apply=sum
+  }
+  )";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_text));
+  ASSERT_OK_AND_ASSIGN(bool changed, RunHloPass(&pass_, module.get()));
+  EXPECT_FALSE(changed);
+}
+
+TEST_F(AllReduceReduceScatterReorderTest, MismatchedReduceScatterGroupSize) {
+  absl::string_view hlo_text = R"(
+  sum {
+    a = f32[] parameter(0)
+    b = f32[] parameter(1)
+    ROOT sum = f32[] add(a, b)
+  }
+
+  ENTRY main {
+    p0 = f32[8] parameter(0)
+    ar = f32[8] all-reduce(p0), replica_groups={{0}, {1}, {2}, {3}}, to_apply=sum
+    ROOT rs = f32[4] reduce-scatter(ar), dimensions={0}, replica_groups={{0,1}}, to_apply=sum
   }
   )";
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,

From 7425690d0144b8fbed6aa9d33ea6ff5e12709d15 Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Thu, 15 Oct 2026 22:03:42 -0700
Subject: [PATCH 0933/1000] ## Problem

In `RaggedDotMode::kRaggedContracting`, `RaggedToDense` inserts the group
dimension `G` at `new_dim_index` after the batch dimensions `0 .. B-1`.
`CreateRaggedContractingDotDims` places `G` at batch dimension `0` and shifts
every original batch and contracting dimension by `+1`.

| Helper | Before | After |
|---|---|---|
| `TransposeIndexToFront` | Swapped dim `0` and `index` | Moves `index` to `0` and shifts `0 .. index - 1` right by one |
| `FindRhsRaggedDim` | Passed `(std::find(...), begin())` to `std::distance`, returning `-k` | Passes `(begin(), absl::c_find(...))` to `std::distance`, returning `+k` |

## Change

Move `index` to dimension `0` in `TransposeIndexToFront` while preserving the
order of dimensions `0 .. index - 1`. Pass `lhs_contracting_dims.begin()` as
the first argument to `std::distance` with `absl::c_find` in
`FindRhsRaggedDim`.

## Tests

- `//xla/hlo/transforms/expanders:ragged_dot_rewriter_test`:
  pins `kRaggedContracting` with multiple batch dimensions and with multiple
  contracting dimensions where `lhs_ragged_dim` is the second contracting
  dimension.
- `//xla/backends/gpu/tests:ragged_dot_test`:
  `ContractingWithMultipleBatchDims` and `ContractingWithMultipleContractingDims`
  verify numerical execution via `RunAndCompare`.

PiperOrigin-RevId: 1000606393
---
 .../xla/backends/gpu/tests/ragged_dot_test.cc | 53 +++++++++++++++++++
 .../xla/xla/hlo/transforms/expanders/BUILD    |  1 +
 .../expanders/ragged_dot_rewriter.cc          | 22 ++++----
 .../expanders/ragged_dot_rewriter_test.cc     | 50 +++++++++++++++++
 4 files changed, 113 insertions(+), 13 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/tests/ragged_dot_test.cc b/third_party/xla/xla/backends/gpu/tests/ragged_dot_test.cc
index f9011803179..9bae0019ebe 100644
--- a/third_party/xla/xla/backends/gpu/tests/ragged_dot_test.cc
+++ b/third_party/xla/xla/backends/gpu/tests/ragged_dot_test.cc
@@ -130,6 +130,59 @@ ENTRY main {
                             ErrorSpec{0, 0}));
 }
 
+TEST_F(RaggedDotTest, ContractingWithMultipleBatchDims) {
+  const char* hlo_text = R"(
+HloModule m
+
+ENTRY main {
+  p0 = f32[2,3,5,11]{3,2,1,0} parameter(0)
+  p1 = f32[2,3,5,7]{3,2,1,0} parameter(1)
+  p2 = s64[2,3,2]{2,1,0} parameter(2)
+  ROOT ragged-dot = f32[2,2,3,11,7]{4,3,2,1,0} ragged-dot(p0, p1, p2),
+                    lhs_contracting_dims={2}, rhs_contracting_dims={2},
+                    lhs_batch_dims={0,1}, rhs_batch_dims={0,1},
+                    lhs_ragged_dims={2}
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text));
+  FakeArgumentsOptions options;
+  options.max_bits_of_precision = 10;
+  ASSERT_OK_AND_ASSIGN(auto fake_arguments,
+                       MakeFakeArguments(module.get(), options));
+  // Set group sizes to reasonable numbers for ragged_dim_size=5.
+  fake_arguments[2] = LiteralUtil::CreateR3<int64_t>(
+      {{{2, 3}, {1, 4}, {3, 2}}, {{4, 1}, {2, 3}, {1, 4}}});
+  EXPECT_TRUE(RunAndCompare(std::move(module),
+                            LiteralUtil::MakePointers(fake_arguments),
+                            ErrorSpec{0, 0}));
+}
+
+TEST_F(RaggedDotTest, ContractingWithMultipleContractingDims) {
+  const char* hlo_text = R"(
+HloModule m
+
+ENTRY main {
+  p0 = f32[2,11,3,5]{3,2,1,0} parameter(0)
+  p1 = f32[2,7,3,5]{3,2,1,0} parameter(1)
+  p2 = s64[2]{0} parameter(2)
+  ROOT ragged-dot = f32[2,2,11,7]{3,2,1,0} ragged-dot(p0, p1, p2),
+                    lhs_contracting_dims={2,3}, rhs_contracting_dims={2,3},
+                    lhs_batch_dims={0}, rhs_batch_dims={0},
+                    lhs_ragged_dims={3}
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text));
+  FakeArgumentsOptions options;
+  options.max_bits_of_precision = 10;
+  ASSERT_OK_AND_ASSIGN(auto fake_arguments,
+                       MakeFakeArguments(module.get(), options));
+  // Set group sizes to reasonable numbers for ragged_dim_size=5.
+  fake_arguments[2] = LiteralUtil::CreateR1<int64_t>({2, 3});
+  EXPECT_TRUE(RunAndCompare(std::move(module),
+                            LiteralUtil::MakePointers(fake_arguments),
+                            ErrorSpec{0, 0}));
+}
+
 }  // namespace
 }  // namespace gpu
 }  // namespace xla
diff --git a/third_party/xla/xla/hlo/transforms/expanders/BUILD b/third_party/xla/xla/hlo/transforms/expanders/BUILD
index ef8068874c1..436e6c01b95 100644
--- a/third_party/xla/xla/hlo/transforms/expanders/BUILD
+++ b/third_party/xla/xla/hlo/transforms/expanders/BUILD
@@ -477,6 +477,7 @@ cc_library(
         "//xla/stream_executor/cuda:cuda_compute_capability",
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:statusor",
+        "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/status",
diff --git a/third_party/xla/xla/hlo/transforms/expanders/ragged_dot_rewriter.cc b/third_party/xla/xla/hlo/transforms/expanders/ragged_dot_rewriter.cc
index 2d151b6aa55..79d09dc0de2 100644
--- a/third_party/xla/xla/hlo/transforms/expanders/ragged_dot_rewriter.cc
+++ b/third_party/xla/xla/hlo/transforms/expanders/ragged_dot_rewriter.cc
@@ -23,6 +23,7 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
+#include "absl/algorithm/container.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
@@ -207,21 +208,17 @@ HloInstruction* TransposeIndexToFront(HloInstruction* instr, int index) {
   if (index == 0) {
     return instr;
   }
-  auto new_shape = instr->shape();
-  int64_t dim0 = new_shape.dimensions(0);
-  int64_t dimIndex = new_shape.dimensions(index);
-  new_shape.set_dimensions(0, dimIndex);
-  new_shape.set_dimensions(index, dim0);
+  const Shape& old_shape = instr->shape();
   llvm::SmallVector<int64_t> transpose_dims;
-  transpose_dims.reserve(new_shape.dimensions().size());
+  transpose_dims.reserve(old_shape.dimensions().size());
   transpose_dims.push_back(index);
-  for (int i = 1; i < new_shape.dimensions().size(); ++i) {
-    if (i == index) {
-      transpose_dims.push_back(0);
-    } else {
+  for (int i = 0; i < old_shape.dimensions().size(); ++i) {
+    if (i != index) {
       transpose_dims.push_back(i);
     }
   }
+  Shape new_shape =
+      ShapeUtil::PermuteDimensionsIgnoringLayout(transpose_dims, old_shape);
   auto computation = instr->parent();
   return computation->AddInstruction(
       HloInstruction::CreateTranspose(new_shape, instr, transpose_dims));
@@ -251,9 +248,8 @@ RaggedDotMode GetRaggedDotMode(int lhs_ragged_dim,
 int FindRhsRaggedDim(const DotDimensionNumbers& dot_dims, int lhs_ragged_dim) {
   const auto& lhs_contracting_dims = dot_dims.lhs_contracting_dimensions();
   int ragged_contracting_index =
-      std::distance(std::find(lhs_contracting_dims.begin(),
-                              lhs_contracting_dims.end(), lhs_ragged_dim),
-                    lhs_contracting_dims.begin());
+      std::distance(lhs_contracting_dims.begin(),
+                    absl::c_find(lhs_contracting_dims, lhs_ragged_dim));
   return dot_dims.rhs_contracting_dimensions(ragged_contracting_index);
 }
 
diff --git a/third_party/xla/xla/hlo/transforms/expanders/ragged_dot_rewriter_test.cc b/third_party/xla/xla/hlo/transforms/expanders/ragged_dot_rewriter_test.cc
index 11667a62577..00131eaa6a1 100644
--- a/third_party/xla/xla/hlo/transforms/expanders/ragged_dot_rewriter_test.cc
+++ b/third_party/xla/xla/hlo/transforms/expanders/ragged_dot_rewriter_test.cc
@@ -365,6 +365,56 @@ TEST_F(RaggedDotRewriterTest, RaggedContractingWithBatchDims) {
   )");
 }
 
+TEST_F(RaggedDotRewriterTest, RaggedContractingWithMultipleBatchDims) {
+  absl::string_view module_string = R"(
+  HloModule module
+
+  ENTRY main {
+    p0 = f32[2,4,11,5]{3,2,1,0} parameter(0)
+    p1 = f32[2,4,5,7]{3,2,1,0} parameter(1)
+    p2 = s32[2,4,6]{2,1,0} parameter(2)
+    ROOT ragged-dot = f32[6,2,4,11,7]{4,3,2,1,0} ragged-dot(p0, p1, p2),
+                                      lhs_contracting_dims={3},
+                                      rhs_contracting_dims={2},
+                                      lhs_batch_dims={0,1},
+                                      rhs_batch_dims={0,1},
+                                      lhs_ragged_dims={3}
+  })";
+
+  RunAndFilecheckHloRewrite(module_string,
+                            RaggedDotRewriter(GetCudaComputeCapability()), R"(
+  // CHECK: f32[6,2,4,11,5]{4,3,2,1,0} transpose({{.*}}), dimensions={2,0,1,3,4}
+  // CHECK: f32[6,2,4,5,7]{4,3,2,1,0} transpose({{.*}}), dimensions={2,0,1,3,4}
+  // CHECK: ROOT [[dot:%[^ ]+]] = f32[6,2,4,11,7]{4,3,2,1,0}
+  // CHECK-SAME: lhs_batch_dims={0,1,2}, lhs_contracting_dims={4},
+  // CHECK-SAME: rhs_batch_dims={0,1,2}, rhs_contracting_dims={3}
+  )");
+}
+
+TEST_F(RaggedDotRewriterTest, RaggedContractingWithMultipleContractingDims) {
+  absl::string_view module_string = R"(
+  HloModule module
+
+  ENTRY main {
+    p0 = f32[2,11,3,5]{3,2,1,0} parameter(0)
+    p1 = f32[2,7,3,5]{3,2,1,0} parameter(1)
+    p2 = s32[6]{0} parameter(2)
+    ROOT ragged-dot = f32[6,2,11,7]{3,2,1,0} ragged-dot(p0, p1, p2),
+                                      lhs_contracting_dims={2,3},
+                                      rhs_contracting_dims={2,3},
+                                      lhs_batch_dims={0},
+                                      rhs_batch_dims={0},
+                                      lhs_ragged_dims={3}
+  })";
+
+  RunAndFilecheckHloRewrite(module_string,
+                            RaggedDotRewriter(GetCudaComputeCapability()), R"(
+  // CHECK: ROOT [[dot:%[^ ]+]] = f32[6,2,11,7]{3,2,1,0}
+  // CHECK-SAME: lhs_batch_dims={0,1}, lhs_contracting_dims={3,4},
+  // CHECK-SAME: rhs_batch_dims={0,1}, rhs_contracting_dims={3,4}
+  )");
+}
+
 TEST_F(RaggedDotRewriterTest, BatchContracting) {
   absl::string_view module_string = R"(
   HloModule module

From 8309e670bde5921b23f4425599143d9bca481ad4 Mon Sep 17 00:00:00 2001
From: David Dunleavy <ddunleavy@google.com>
Date: Thu, 15 Oct 2026 22:06:04 -0700
Subject: [PATCH 0934/1000] `setHandler` in XLA's `error_util`

I think this was intended in https://github.com/tensorflow/tensorflow/commit/0549fd10f931b49d6186543dced5fbf0800daa4e, and makes the XLA `error_util` functionally the same as TensorFlow's modulo filtering logic

PiperOrigin-RevId: 1000607373
---
 .../compiler/mlir/tensorflow/utils/error_util_test.cc      | 7 -------
 third_party/xla/xla/mlir/utils/error_util.cc               | 4 +++-
 third_party/xla/xla/mlir/utils/error_util_test.cc          | 6 ++++++
 3 files changed, 9 insertions(+), 8 deletions(-)

diff --git a/tensorflow/compiler/mlir/tensorflow/utils/error_util_test.cc b/tensorflow/compiler/mlir/tensorflow/utils/error_util_test.cc
index adfc984a5cd..937ff1becef 100644
--- a/tensorflow/compiler/mlir/tensorflow/utils/error_util_test.cc
+++ b/tensorflow/compiler/mlir/tensorflow/utils/error_util_test.cc
@@ -73,13 +73,6 @@ TEST_F(StatusScopedDiagnosticHandlerTest,
   EXPECT_THAT(s.message(), HasSubstr("Second diagnostic message reported"));
 }
 
-TEST_F(StatusScopedDiagnosticHandlerTest, VerifyThatWarningsAreIgnored) {
-  // Note: this logic is actually implemented in BaseScopedDiagnosticHandler's
-  // handler() function, but only StatusScopedDiagnosticHandler uses it.
-  StatusScopedDiagnosticHandler handler(&context_);
-  emitWarning(loc_) << "Warning message";
-  TF_EXPECT_OK(handler.ConsumeStatus());
-}
 
 TEST(ErrorUtilTest, StatusScopedDiagnosticHandlerWithFilter) {
   // Filtering logic is based on tensorflow::IsInternalFrameForFilename()
diff --git a/third_party/xla/xla/mlir/utils/error_util.cc b/third_party/xla/xla/mlir/utils/error_util.cc
index 10ea77c0791..16cb31514ac 100644
--- a/third_party/xla/xla/mlir/utils/error_util.cc
+++ b/third_party/xla/xla/mlir/utils/error_util.cc
@@ -31,7 +31,9 @@ BaseScopedDiagnosticHandler::BaseScopedDiagnosticHandler(MLIRContext* context,
                                                          bool filter_stack)
     : SourceMgrDiagnosticHandler(source_mgr_, context, diag_stream_),
       diag_stream_(diag_str_),
-      propagate_(propagate) {}
+      propagate_(propagate) {
+  setHandler([this](Diagnostic& diag) { return this->handler(&diag); });
+}
 
 BaseScopedDiagnosticHandler::~BaseScopedDiagnosticHandler() {
   // Verify errors were consumed and re-register old handler.
diff --git a/third_party/xla/xla/mlir/utils/error_util_test.cc b/third_party/xla/xla/mlir/utils/error_util_test.cc
index a64576d432a..c951927f3b2 100644
--- a/third_party/xla/xla/mlir/utils/error_util_test.cc
+++ b/third_party/xla/xla/mlir/utils/error_util_test.cc
@@ -70,5 +70,11 @@ TEST_F(BaseScopedDiagnosticHandlerTest,
       absl::StrContains(s.message(), "Second diagnostic message reported"));
 }
 
+TEST_F(BaseScopedDiagnosticHandlerTest, VerifyThatWarningsAreIgnored) {
+  BaseScopedDiagnosticHandler handler(&context_);
+  emitWarning(loc_) << "Warning message";
+  EXPECT_OK(handler.ConsumeStatus());
+}
+
 }  // namespace
 }  // namespace mlir

From 43e3784d45eeabf2279f595df011593b05e84ddc Mon Sep 17 00:00:00 2001
From: y198nt <bnguyenthai191@gmail.com>
Date: Fri, 16 Oct 2026 14:01:49 +0700
Subject: [PATCH 0935/1000] Fix heap address leak in
 ResourceBase::MakeRefCountingHandleName

MakeRefCountingHandleName() embedded raw this pointer via %p in the handle name. Handle names are exposed to callers through error messages (e.g. Assert failures), leaking the glibc heap address and defeating ASLR.

Since resource_id is already globally unique (atomic counter via GenerateUniqueId()), the this pointer is unnecessary for uniqueness. Remove it entirely.

No format parsers depend on the old "-at-0x..." suffix handle names are used as opaque lookup keys.
---
 tensorflow/core/framework/resource_base.h | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/framework/resource_base.h b/tensorflow/core/framework/resource_base.h
index 18137a69da8..67061531849 100644
--- a/tensorflow/core/framework/resource_base.h
+++ b/tensorflow/core/framework/resource_base.h
@@ -40,8 +40,11 @@ class ResourceBase : public core::WeakRefCounted {
   virtual std::string DebugString() const = 0;
 
   // Returns a name for ref-counting handles.
+  // resource_id is globally unique (atomic counter), so it alone is sufficient
+  // for handle uniqueness. Avoid embedding `this` pointer — handle names are
+  // exposed to callers via error messages and would leak heap addresses.
   virtual std::string MakeRefCountingHandleName(int64_t resource_id) const {
-    return absl::StrFormat("Resource-%d-at-%p", resource_id, this);
+    return absl::StrFormat("Resource-%d", resource_id);
   }
 
   // Returns memory used by this resource.

From 3d897980f0a2e9910aa7995852f353762d0221fb Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 15 Oct 2026 23:47:50 -0700
Subject: [PATCH 0936/1000] Automated Code Change

PiperOrigin-RevId: 1000646186
---
 third_party/xla/xla/backends/gpu/transforms/BUILD              | 1 -
 .../gpu/transforms/fusion_block_level_rewriter_test.cc         | 3 ---
 .../xla/xla/backends/gpu/transforms/gemm_fusion_test.cc        | 2 --
 3 files changed, 6 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD
index 48b69fe4626..99bd41cb23a 100644
--- a/third_party/xla/xla/backends/gpu/transforms/BUILD
+++ b/third_party/xla/xla/backends/gpu/transforms/BUILD
@@ -289,7 +289,6 @@ xla_cc_test(
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest_main",
         "@llvm-project//mlir:IR",
diff --git a/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter_test.cc
index 055647a03e2..235e2955a6f 100644
--- a/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/fusion_block_level_rewriter_test.cc
@@ -17,8 +17,6 @@ License.
 
 #include <memory>
 #include <string>
-#include <tuple>
-#include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
@@ -26,7 +24,6 @@ License.
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
 #include "absl/status/statusor.h"
-#include "absl/strings/str_join.h"
 #include "absl/strings/string_view.h"
 #include "mlir/IR/MLIRContext.h"
 #include "xla/backends/gpu/codegen/emitters/mlir_kernel_emitter.h"
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 4d3eb4bfe45..2fd6de0ff01 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
@@ -18,13 +18,11 @@ limitations under the License.
 #include <memory>
 #include <optional>
 #include <string>
-#include <tuple>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
-#include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "xla/autotuning.pb.h"
 #include "xla/hlo/ir/hlo_instruction.h"

From b3815fd79d87bd1c7e5ddc525f55538a4959f0e1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Longeri?= <tlongeri@google.com>
Date: Fri, 16 Oct 2026 00:13:13 -0700
Subject: [PATCH 0937/1000] [Mosaic] Enable 16-bit find_first_set
 tpu.all_reduce

PiperOrigin-RevId: 1000657170
---
 third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc
index 3f981e0845f..0348b49c9eb 100644
--- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc
+++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc
@@ -3096,10 +3096,10 @@ LogicalResult AllReduceOp::verify() {
 
   if (in_bitwidth == 1) {
     // For mask vectors, the single (semantically scalar) result is broadcast
-    // into a vector of 32-bit ints of whatever shape the target supports (not
+    // into a vector of ints of whatever shape the target supports (not
     // necessarily the same as the input).
-    if (!out_ty.getElementType().isSignlessInteger(32)) {
-      return emitOpError("Vector mask all-reduce must have i32 output");
+    if (!out_ty.getElementType().isSignlessInteger()) {
+      return emitOpError("Vector mask all-reduce must have integer output");
     }
     switch (kind) {
       case ReductionKind::kSum:

From 20c7528f8bca3f37f3eb406237f3087324e70afc Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 16 Oct 2026 00:35:13 -0700
Subject: [PATCH 0938/1000] Automated Code Change

PiperOrigin-RevId: 1000666008
---
 third_party/xla/xla/tsl/platform/cloud/gcs_dns_cache.cc   | 8 ++++++--
 third_party/xla/xla/tsl/platform/cloud/gcs_file_system.cc | 7 +++----
 third_party/xla/xla/tsl/platform/cloud/gcs_throttle.h     | 6 ++----
 third_party/xla/xla/tsl/platform/cloud/http_request.h     | 2 +-
 .../xla/xla/tsl/platform/cloud/ram_file_block_cache.cc    | 7 ++++---
 5 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/third_party/xla/xla/tsl/platform/cloud/gcs_dns_cache.cc b/third_party/xla/xla/tsl/platform/cloud/gcs_dns_cache.cc
index 73af26eed2f..305a952d53b 100644
--- a/third_party/xla/xla/tsl/platform/cloud/gcs_dns_cache.cc
+++ b/third_party/xla/xla/tsl/platform/cloud/gcs_dns_cache.cc
@@ -238,9 +238,13 @@ void GcsDnsCache::WorkerThread() {
     {
       // Don't immediately re-resolve the addresses.
       absl::MutexLock l(mu_);
-      if (cancelled_) return;
+      if (cancelled_) {
+        return;
+      }
       cond_var_.WaitWithTimeout(&mu_, absl::Seconds(refresh_rate_secs_));
-      if (cancelled_) return;
+      if (cancelled_) {
+        return;
+      }
     }
 
     // Resolve DNS values
diff --git a/third_party/xla/xla/tsl/platform/cloud/gcs_file_system.cc b/third_party/xla/xla/tsl/platform/cloud/gcs_file_system.cc
index 817e463edef..8eda7e543e9 100644
--- a/third_party/xla/xla/tsl/platform/cloud/gcs_file_system.cc
+++ b/third_party/xla/xla/tsl/platform/cloud/gcs_file_system.cc
@@ -1514,10 +1514,9 @@ absl::Status GcsFileSystem::FileExists(absl::string_view fname) {
     ABSL_RETURN_IF_ERROR(BucketExists(bucket, &result));
     if (result) {
       return absl::OkStatus();
-    } else {
+    }
       return absl::NotFoundError(
           absl::StrCat("The specified bucket ", fname, " was not found."));
-    }
   }
 
   // Check if the object exists.
@@ -1766,9 +1765,9 @@ absl::Status GcsFileSystem::FolderExists(absl::string_view dirname,
     if (!children.empty()) {
       stat->base = DIRECTORY_STAT;
       return absl::OkStatus();
-    } else {
-      return absl::InvalidArgumentError("Not a directory!");
     }
+    return absl::InvalidArgumentError("Not a directory!");
+    
   };
   GcsFileStat stat;
   absl::Status s = stat_cache_->LookupOrCompute(MaybeAppendSlash(dirname),
diff --git a/third_party/xla/xla/tsl/platform/cloud/gcs_throttle.h b/third_party/xla/xla/tsl/platform/cloud/gcs_throttle.h
index 9fd44e794ba..098367d9f55 100644
--- a/third_party/xla/xla/tsl/platform/cloud/gcs_throttle.h
+++ b/third_party/xla/xla/tsl/platform/cloud/gcs_throttle.h
@@ -112,7 +112,7 @@ class GcsThrottle {
    * purpose of this function is to make available to monitoring or other
    * instrumentation the number of available tokens in the pool.
    */
-  inline int64_t available_tokens() TF_LOCKS_EXCLUDED(mu_) {
+  int64_t available_tokens() TF_LOCKS_EXCLUDED(mu_) {
     absl::MutexLock l(mu_);
     UpdateState();
     return available_tokens_;
@@ -139,9 +139,7 @@ class GcsThrottle {
    */
   void UpdateState() TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);
 
-  inline uint64_t request_bytes_to_tokens(size_t num_bytes) {
-    return num_bytes >> 10;
-  }
+  uint64_t request_bytes_to_tokens(size_t num_bytes) { return num_bytes >> 10; }
 
   absl::Mutex mu_;
 
diff --git a/third_party/xla/xla/tsl/platform/cloud/http_request.h b/third_party/xla/xla/tsl/platform/cloud/http_request.h
index 2a441fee630..cb1cf0dfb4e 100644
--- a/third_party/xla/xla/tsl/platform/cloud/http_request.h
+++ b/third_party/xla/xla/tsl/platform/cloud/http_request.h
@@ -59,7 +59,7 @@ class HttpRequest {
   };
 
   /// RequestMethodName converts a RequestMethod to the canonical method string.
-  inline static const char* RequestMethodName(RequestMethod m) {
+  static const char* RequestMethodName(RequestMethod m) {
     switch (m) {
       case RequestMethod::kGet:
         return "GET";
diff --git a/third_party/xla/xla/tsl/platform/cloud/ram_file_block_cache.cc b/third_party/xla/xla/tsl/platform/cloud/ram_file_block_cache.cc
index d784bbbc1ca..a6d638e3281 100644
--- a/third_party/xla/xla/tsl/platform/cloud/ram_file_block_cache.cc
+++ b/third_party/xla/xla/tsl/platform/cloud/ram_file_block_cache.cc
@@ -33,7 +33,9 @@ bool RamFileBlockCache::BlockNotStale(const std::shared_ptr<Block>& block) {
   if (block->state != FetchState::FINISHED) {
     return true;  // No need to check for staleness.
   }
-  if (max_staleness_ == 0) return true;  // Not enforcing staleness.
+  if (max_staleness_ == 0) {
+    return true;  // Not enforcing staleness.
+  }
   return env_->NowSeconds() - block->timestamp <= max_staleness_;
 }
 
@@ -47,10 +49,9 @@ std::shared_ptr<RamFileBlockCache::Block> RamFileBlockCache::Lookup(
         cache_stats_->RecordCacheHitBlockSize(entry->second->data.size());
       }
       return entry->second;
-    } else {
+    }
       // Remove the stale block and continue.
       RemoveFile_Locked(key.first);
-    }
   }
 
   // Insert a new empty block, setting the bookkeeping to sentinel values

From 08aa694862f88fde5188d13f67973fc18f8ce369 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 16 Oct 2026 01:21:22 -0700
Subject: [PATCH 0939/1000] Bump zizmor from 1.31.0 to 1.31.1 in
 .github/workflows/zizmor.yml.

Upgrade the pinned `zizmor` version in `.github/workflows/zizmor.yml` from `1.31.0` to `1.31.1` (https://github.com/zizmorcore/zizmor/releases/tag/v1.31.1).

PiperOrigin-RevId: 1000684389
---
 third_party/xla/.github/workflows/zizmor.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/third_party/xla/.github/workflows/zizmor.yml b/third_party/xla/.github/workflows/zizmor.yml
index 5767aec9996..7fdd7882e72 100644
--- a/third_party/xla/.github/workflows/zizmor.yml
+++ b/third_party/xla/.github/workflows/zizmor.yml
@@ -37,6 +37,6 @@ jobs:
 
       - name: Run zizmor
         # Pinned version; bumped on a regular schedule by the XLA zizmor release monitor.
-        run:  uvx zizmor@1.31.0 --format=github .github xla
+        run:  uvx zizmor@1.31.1 --format=github .github xla
         env:
           GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}

From 49826a87a96f3f07a7d0c455239e81df6471f565 Mon Sep 17 00:00:00 2001
From: y198nt <bnguyenthai191@gmail.com>
Date: Fri, 16 Oct 2026 14:15:09 +0700
Subject: [PATCH 0940/1000] Fix heap address leak in
 ResourceBase::MakeRefCountingHandleName

MakeRefCountingHandleName() embedded raw this pointer via %p in the handle name. Handle names are exposed to callers through error messages (e.g. Assert failures), leaking the glibc heap address and defeating ASLR.

Since resource_id is already globally unique (atomic counter via GenerateUniqueId()), the this pointer is unnecessary for uniqueness. Remove it entirely.

No format parsers depend on the old "-at-0x..." suffix handle names are used as opaque lookup keys.
---
 .../core/framework/resource_handle_test.cc     | 15 +++++++++++++++
 .../data_structures/lookup_ops_test.py         | 18 ++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/tensorflow/core/framework/resource_handle_test.cc b/tensorflow/core/framework/resource_handle_test.cc
index 022b128c145..68ec490cfe8 100644
--- a/tensorflow/core/framework/resource_handle_test.cc
+++ b/tensorflow/core/framework/resource_handle_test.cc
@@ -65,6 +65,21 @@ TEST_F(ResourceHandleTest, RefCounting) {
   EXPECT_FALSE(alive);
 }
 
+TEST_F(ResourceHandleTest, HandleNameDoesNotLeakAddress) {
+  auto resource = new MockResource(nullptr, 42);
+  auto handle =
+      ResourceHandle::MakeRefCountingHandle(resource, "cpu", {}, {});
+  // Handle name must not contain a raw pointer representation.
+  // Previously the name was "Resource-<id>-at-0x<hex>" which leaked the
+  // heap address of the resource to callers via error messages.
+  EXPECT_EQ(handle.name().find("-at-"), std::string::npos)
+      << "Handle name must not embed raw pointer: " << handle.name();
+  EXPECT_EQ(handle.name().find("0x"), std::string::npos)
+      << "Handle name must not embed hex address: " << handle.name();
+  // Name should start with "Resource-"
+  EXPECT_EQ(handle.name().substr(0, 9), "Resource-");
+}
+
 TEST_F(ResourceHandleTest, SummarizeValue) {
   ResourceHandleProto proto;
   TensorShapeProto shape;
diff --git a/tensorflow/python/kernel_tests/data_structures/lookup_ops_test.py b/tensorflow/python/kernel_tests/data_structures/lookup_ops_test.py
index 9aba606f69b..87c022236bc 100644
--- a/tensorflow/python/kernel_tests/data_structures/lookup_ops_test.py
+++ b/tensorflow/python/kernel_tests/data_structures/lookup_ops_test.py
@@ -4384,6 +4384,24 @@ class MutableHashTableOpTest(test.TestCase):
     self.assertTrue(inferred_shapes[1].is_compatible_with(actual_shapes[1]))
 
 
+class ResourceHandleNameSecurityTest(test.TestCase):
+  """Verify that anonymous resource handle names do not leak heap addresses."""
+
+  @test_util.run_v2_only
+  def testAnonymousHashTableHandleDoesNotLeakAddress(self):
+    # AnonymousHashTable creates a ResourceBase whose handle name was
+    # previously formatted as "Resource-<id>-at-0x<hex_address>", leaking
+    # the glibc heap address to callers via error messages.
+    table = gen_lookup_ops.anonymous_hash_table(
+        key_dtype=dtypes.int64, value_dtype=dtypes.int64)
+    handle_name = str(table)
+    self.assertRegex(handle_name, r'name="Resource-\d+"')
+    self.assertNotIn("-at-", handle_name,
+                     "Handle name must not contain raw pointer")
+    self.assertNotRegex(handle_name, r"0x[0-9a-fA-F]{8,}",
+                        "Handle name must not contain hex address")
+
+
 class MutableHashTableBenchmark(test.Benchmark):
 
   def _create_table(self):

From 0736391f29976322acaeb3c660aa4038cb1d25c2 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel@google.com>
Date: Fri, 16 Oct 2026 01:28:07 -0700
Subject: [PATCH 0941/1000] Support scalar operands in XTile kClamp and kSelect

HLO `kClamp` allows scalar `min` and `max` bounds alongside a non-scalar
`operand`, and `kSelect` allows a scalar `pred` alongside non-scalar `on_true`
and `on_false` operands. Previously, tile propagation propagated the non-scalar
output tile unchanged to all elementwise operands, causing `xtile.extract` on
a 0D buffer to fail verification with a rank mismatch.

Update `PropagateTileToInputForCwiseOp` to propagate a 0D tile to 0D array
operands, and `PropagateTileToOutputForCwiseOp` to delegate 0D-to-ND
propagation to `PropagateTileToOutputForBroadcastOpImpl`. In the XTile
`EmitElementwise`, broadcast scalar `min`/`max` bounds in `kClamp` and scalar
`pred` in `kSelect` to the output tile shape before emitting the elementwise
operations.

Fixes #50116

PiperOrigin-RevId: 1000687194
---
 .../tests/elementwise/clamp_scalar_bounds.hlo | 60 ++++++++++++++
 .../tests/elementwise/select_scalar_pred.hlo  | 58 ++++++++++++++
 .../tiling/experimental/tile_propagation.cc   | 34 ++++++--
 .../experimental/tile_propagation_test.cc     | 80 +++++++++++++++++++
 .../codegen/xtile/codegen/emitter_helpers.cc  | 32 ++++++--
 5 files changed, 248 insertions(+), 16 deletions(-)
 create mode 100644 third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/clamp_scalar_bounds.hlo
 create mode 100644 third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/select_scalar_pred.hlo

diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/clamp_scalar_bounds.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/clamp_scalar_bounds.hlo
new file mode 100644
index 00000000000..986dcc740d9
--- /dev/null
+++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/clamp_scalar_bounds.hlo
@@ -0,0 +1,60 @@
+// RUN: fusion_to_xtile %s --xla_cpu_use_new_xtile_lowering | FileCheck %s
+// RUN: fusion_to_xtile %s --xla_cpu_use_new_xtile_lowering | fusion_compiler_opt \
+// RUN:   --xtile-cpu-new-xtile-to-vector --xla-simplify-arith \
+// RUN:   | FileCheck %s --check-prefix=NEW-VECTOR
+// RUN: fusion_to_xtile %s --xla_cpu_use_new_xtile_lowering | 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
+
+fusion {
+  p0 = f32[] parameter(0)
+  p1 = f32[128,64] parameter(1)
+  c0 = f32[] constant(1.0)
+  ROOT clamp = f32[128,64] clamp(p0, p1, c0)
+}
+
+ENTRY main {
+  p0 = f32[] parameter(0)
+  p1 = f32[128,64] parameter(1)
+  ROOT fusion = f32[128,64] fusion(p0, p1), kind=kLoop, calls=fusion,
+    backend_config={
+      "fusion_backend_config":{
+        "block_level_fusion_config":{
+          "output_tiles":[{"sizes":["4", "4"]}],
+          "num_ctas":"1"
+        }
+      }
+    }
+}
+
+// CHECK:        xtile.entry_func @wrapped_fusion(
+// CHECK-SAME:     %[[ARG0:[^:]+]]: memref<f32>,
+// CHECK-SAME:     %[[ARG1:[^:]+]]: memref<128x64xf32>,
+// CHECK-SAME:     %[[ARG2:[^:]+]]: memref<128x64xf32>,
+// CHECK-SAME:     %{{.*}}: index {xla.range = [0 : index, 0 : index]})
+
+// CHECK:        scf.for %[[IDX:.*]] = %{{.*}} to %{{.*}} step %{{.*}} {
+// CHECK-DAG:      %[[MIN_SCALAR:.*]] = xtile.extract %[[ARG0]][] [] [] : memref<f32> -> tensor<f32>
+// CHECK-DAG:      %[[VAL:.*]] = xtile.extract %[[ARG1]][%{{.*}}, %{{.*}}] [4, 4] [1, 1] : memref<128x64xf32> -> tensor<4x4xf32>
+// CHECK-DAG:      %[[MAX_SCALAR:.*]] = arith.constant dense<1.000000e+00> : tensor<f32>
+// CHECK-DAG:      %[[MIN_BCAST:.*]] = stablehlo.broadcast_in_dim %[[MIN_SCALAR]], dims = [] : (tensor<f32>) -> tensor<4x4xf32>
+// CHECK-DAG:      %[[MAX_BCAST:.*]] = stablehlo.broadcast_in_dim %[[MAX_SCALAR]], dims = [] : (tensor<f32>) -> tensor<4x4xf32>
+// CHECK:          %[[MAX_OP:.*]] = stablehlo.maximum %[[MIN_BCAST]], %[[VAL]] : tensor<4x4xf32>
+// CHECK:          %[[MIN_OP:.*]] = stablehlo.minimum %[[MAX_OP]], %[[MAX_BCAST]] : tensor<4x4xf32>
+// CHECK:          xtile.insert %[[MIN_OP]] into %[[ARG2]][%{{.*}}, %{{.*}}] [4, 4] [1, 1] : tensor<4x4xf32> -> memref<128x64xf32>
+
+// NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl(
+// NEW-VECTOR-SAME:    %[[ARG0:[^:]+]]: memref<f32>,
+// NEW-VECTOR-SAME:    %[[ARG1:[^:]+]]: memref<128x64xf32>,
+// NEW-VECTOR-SAME:    %[[ARG2:[^:]+]]: memref<128x64xf32>,
+// NEW-VECTOR:         scf.for
+// NEW-VECTOR-DAG:       %[[READ_MIN:.*]] = vector.transfer_read %[[ARG0]][], %{{.*}} : memref<f32>, vector<f32>
+// NEW-VECTOR-DAG:       %[[BCAST_MIN:.*]] = vector.broadcast %[[READ_MIN]] : vector<f32> to vector<4x4xf32>
+// NEW-VECTOR-DAG:       %[[READ_VAL:.*]] = vector.transfer_read %[[ARG1]][%{{.*}}, %{{.*}}], %{{.*}} {in_bounds = [true, true]} : memref<128x64xf32>, vector<4x4xf32>
+// NEW-VECTOR:           %[[MAX:.*]] = arith.maximumf %[[BCAST_MIN]], %[[READ_VAL]] : vector<4x4xf32>
+// NEW-VECTOR:           %[[MIN:.*]] = arith.minimumf %[[MAX]], %{{.*}} : vector<4x4xf32>
+// NEW-VECTOR:           vector.transfer_write %[[MIN]], %[[ARG2]][%{{.*}}, %{{.*}}] {in_bounds = [true, true]} : vector<4x4xf32>, memref<128x64xf32>
+
+// 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/tests/elementwise/select_scalar_pred.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/select_scalar_pred.hlo
new file mode 100644
index 00000000000..afc1f8f90e1
--- /dev/null
+++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/elementwise/select_scalar_pred.hlo
@@ -0,0 +1,58 @@
+// RUN: fusion_to_xtile %s --xla_cpu_use_new_xtile_lowering | FileCheck %s
+// RUN: fusion_to_xtile %s --xla_cpu_use_new_xtile_lowering | fusion_compiler_opt \
+// RUN:   --xtile-cpu-new-xtile-to-vector --xla-simplify-arith \
+// RUN:   | FileCheck %s --check-prefix=NEW-VECTOR
+// RUN: fusion_to_xtile %s --xla_cpu_use_new_xtile_lowering | 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
+
+fusion {
+  p0 = pred[] parameter(0)
+  p1 = f32[128,64] parameter(1)
+  p2 = f32[128,64] parameter(2)
+  ROOT select = f32[128,64] select(p0, p1, p2)
+}
+
+ENTRY main {
+  p0 = pred[] parameter(0)
+  p1 = f32[128,64] parameter(1)
+  p2 = f32[128,64] parameter(2)
+  ROOT fusion = f32[128,64] fusion(p0, p1, p2), kind=kLoop, calls=fusion,
+    backend_config={
+      "fusion_backend_config":{
+        "block_level_fusion_config":{
+          "output_tiles":[{"sizes":["4", "4"]}],
+          "num_ctas":"1"
+        }
+      }
+    }
+}
+
+// CHECK:        xtile.entry_func @wrapped_fusion(
+// CHECK-SAME:     %[[ARG0:[^:]+]]: memref<i8>,
+// CHECK-SAME:     %[[ARG1:[^:]+]]: memref<128x64xf32>,
+// CHECK-SAME:     %[[ARG2:[^:]+]]: memref<128x64xf32>,
+// CHECK-SAME:     %[[ARG3:[^:]+]]: memref<128x64xf32>,
+// CHECK-SAME:     %{{.*}}: index {xla.range = [0 : index, 0 : index]})
+
+// CHECK:        scf.for %[[IDX:.*]] = %{{.*}} to %{{.*}} step %{{.*}} {
+// CHECK:          %[[COND_I8:.*]] = xtile.extract %[[ARG0]][] [] [] : memref<i8> -> tensor<i8>
+// CHECK:          %[[COND_I1:.*]] = stablehlo.convert %[[COND_I8]] : (tensor<i8>) -> tensor<i1>
+// CHECK:          %[[ON_TRUE:.*]] = xtile.extract %[[ARG1]][%{{.*}}, %{{.*}}] [4, 4] [1, 1] : memref<128x64xf32> -> tensor<4x4xf32>
+// CHECK:          %[[ON_FALSE:.*]] = xtile.extract %[[ARG2]][%{{.*}}, %{{.*}}] [4, 4] [1, 1] : memref<128x64xf32> -> tensor<4x4xf32>
+// CHECK:          %[[CMP:.*]] = stablehlo.compare NE, %[[COND_I1]], %{{.*}} : (tensor<i1>, tensor<i1>) -> tensor<i1>
+// CHECK:          %[[BCAST:.*]] = stablehlo.broadcast_in_dim %[[CMP]], dims = [] : (tensor<i1>) -> tensor<4x4xi1>
+// CHECK:          %[[SEL:.*]] = arith.select %[[BCAST]], %[[ON_TRUE]], %[[ON_FALSE]] : tensor<4x4xi1>, tensor<4x4xf32>
+// CHECK:          xtile.insert %[[SEL]] into %[[ARG3]][%{{.*}}, %{{.*}}] [4, 4] [1, 1] : tensor<4x4xf32> -> memref<128x64xf32>
+
+// NEW-VECTOR-LABEL: func.func @wrapped_fusion_impl(
+// NEW-VECTOR-SAME:    %[[ARG0:[^:]+]]: memref<i8>,
+// NEW-VECTOR-SAME:    %[[ARG1:[^:]+]]: memref<128x64xf32>,
+// NEW-VECTOR-SAME:    %[[ARG2:[^:]+]]: memref<128x64xf32>,
+// NEW-VECTOR-SAME:    %[[ARG3:[^:]+]]: memref<128x64xf32>,
+// NEW-VECTOR:         scf.for
+// NEW-VECTOR:           arith.select
+
+// NEW-VECTOR-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl
+// NEW-VECTOR-LLVM-NOT: llvm.alloc
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
index cbf1da9cc44..5ce97597980 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
@@ -120,13 +120,19 @@ bool DimIsFullyCovered(const DimTile& tile, int64_t dim) {
 }
 
 Tiles PropagateTileToInputForCwiseOp(const HloInstruction& hlo,
-                                     const Tile& input_tile) {
-  return Tiles(hlo.operand_count(), input_tile);
-}
-
-Tiles PropagateTileToOutputForCwiseOp(const HloInstruction& hlo,
-                                      const Tile& output_tile) {
-  return {output_tile};
+                                     const Tile& output_tile) {
+  Tiles input_tiles;
+  input_tiles.reserve(hlo.operand_count());
+  for (const HloInstruction* operand : hlo.operands()) {
+    // kClamp and kSelect allow scalar operands alongside non-scalar operands.
+    if (operand->shape().IsArray() && operand->shape().dimensions().empty() &&
+        !output_tile.dim_tiles().empty()) {
+      input_tiles.push_back(output_tile.CloneWithNewDims({}));
+    } else {
+      input_tiles.push_back(output_tile);
+    }
+  }
+  return input_tiles;
 }
 
 Tiles PropagateTileToInputForBroadcastOp(const HloBroadcastInstruction& bcast,
@@ -169,6 +175,18 @@ Tiles PropagateTileToOutputForBroadcastOp(const HloBroadcastInstruction& bcast,
       bcast.shape(), bcast.dimensions(), input_tile)};
 }
 
+Tiles PropagateTileToOutputForCwiseOp(const HloInstruction& hlo,
+                                      const Tile& input_tile,
+                                      int64_t input_index) {
+  const Shape& operand_shape = hlo.operand(input_index)->shape();
+  if (operand_shape.IsArray() && operand_shape.dimensions().empty() &&
+      hlo.shape().IsArray() && !hlo.shape().dimensions().empty()) {
+    return {PropagateTileToOutputForBroadcastOpImpl(
+        hlo.shape(), /*bcast_dims=*/{}, input_tile)};
+  }
+  return {input_tile};
+}
+
 absl::Status VerifyConcatenateAlignment(
     const TilingSpace& tiling_space,
     const HloConcatenateInstruction& concatenate, const Tile& output_tile) {
@@ -1924,7 +1942,7 @@ absl::StatusOr<Tiles> PropagateTileToOutput(const TilingSpace& tiling_space,
   VLOG(2) << "tiling_space: " << tiling_space.ToString();
   if (HloInstruction::IsOpElementwise(hlo.opcode()) ||
       hlo.opcode() == HloOpcode::kMap) {
-    return PropagateTileToOutputForCwiseOp(hlo, input_tile);
+    return PropagateTileToOutputForCwiseOp(hlo, input_tile, input_index);
   }
   if (hlo.opcode() == HloOpcode::kBitcast) {
     return PropagateTileToOutputForBitcastOp(hlo, input_tile);
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
index 7737e6f806c..189cd457380 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
@@ -137,6 +137,86 @@ TEST_F(TilePropagationTest, CanPropagateToOutputsOfElementwiseOp) {
   EXPECT_THAT(from_operand_1, MatchToString(kExpected));
 }
 
+TEST_F(TilePropagationTest, CanPropagateToInputsOfClampWithScalarBounds) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    ENTRY e {
+      p0 = f32[] parameter(0)
+      p1 = f32[10,20] parameter(1)
+      p2 = f32[] parameter(2)
+      ROOT clamp0 = f32[10,20] clamp(p0, p1, p2)
+    }
+  )");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TilingSpace> tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                          &mlir_context_));
+  ASSERT_OK_AND_ASSIGN(
+      Tiles tiled_operands,
+      PropagateTileToInput(
+          *tiling_space, *root,
+          GetTestTile(*tiling_space, root->shape().dimensions()), 0));
+  EXPECT_THAT(tiled_operands, MatchToString(R"(
+    0) (tid_0, tid_1)
+      -> offsets []
+         sizes []
+         strides []
+         upper bounds []
+    1) (tid_0, tid_1)
+      -> offsets [tid_0 * ts_0, tid_1 * ts_1]
+         sizes [ts_0, ts_1]
+         strides [1, 2]
+         upper bounds [10, 20]
+    2) (tid_0, tid_1)
+      -> offsets []
+         sizes []
+         strides []
+         upper bounds []
+  )"));
+}
+
+TEST_F(TilePropagationTest, CanPropagateToOutputsOfClampWithScalarBounds) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    ENTRY e {
+      p0 = f32[] parameter(0)
+      p1 = f32[10,20] parameter(1)
+      p2 = f32[10,20] parameter(2)
+      ROOT clamp0 = f32[10,20] clamp(p0, p1, p2)
+    }
+  )");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TilingSpace> tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                          &mlir_context_));
+  ASSERT_OK_AND_ASSIGN(
+      Tiles from_operand_0,
+      PropagateTileToOutput(
+          *tiling_space, *root,
+          GetTestTile(*tiling_space, root->operand(0)->shape().dimensions()),
+          0));
+  EXPECT_THAT(from_operand_0, MatchToString(R"(
+    0) (tid_0, tid_1)
+      -> offsets [0, 0]
+         sizes [16, 32]
+         strides [1, 1]
+         upper bounds [10, 20]
+  )"));
+  ASSERT_OK_AND_ASSIGN(
+      Tiles from_operand_1,
+      PropagateTileToOutput(
+          *tiling_space, *root,
+          GetTestTile(*tiling_space, root->operand(1)->shape().dimensions()),
+          1));
+  EXPECT_THAT(from_operand_1, MatchToString(R"(
+    0) (tid_0, tid_1)
+      -> offsets [tid_0 * ts_0, tid_1 * ts_1]
+         sizes [ts_0, ts_1]
+         strides [1, 2]
+         upper bounds [10, 20]
+  )"));
+}
+
 TEST_F(TilePropagationTest, CanPropagateToInputsOfAllReduceOp) {
   HloInstruction* root = ParseAndGetRoot(R"(
     HloModule m
diff --git a/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.cc b/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.cc
index 39078ba1a47..69dd7d97d14 100644
--- a/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.cc
+++ b/third_party/xla/xla/codegen/xtile/codegen/emitter_helpers.cc
@@ -246,6 +246,17 @@ Value DivideBy(mlir::ImplicitLocOpBuilder& b, Value value,
       b, value, CreateConst(b, value.getType(), elements_per_byte));
 }
 
+Value MaybeBroadcastScalarToShape(mlir::ImplicitLocOpBuilder& b, Value input,
+                                  Type target_type) {
+  ArrayRef<int64_t> target_shape =
+      mlir::cast<ShapedType>(target_type).getShape();
+  if (!target_shape.empty() &&
+      mlir::cast<ShapedType>(input.getType()).getShape().empty()) {
+    return Splat(b, input, target_shape);
+  }
+  return input;
+}
+
 }  // namespace
 
 SmallVector<int64_t> GetPaddedTileSizes(ArrayRef<int64_t> tile_sizes) {
@@ -726,8 +737,13 @@ absl::StatusOr<Value> EmitElementwise(mlir::ImplicitLocOpBuilder& b,
       return Maximum(b, inputs);
     case HloOpcode::kMinimum:
       return Minimum(b, inputs);
-    case HloOpcode::kClamp:
-      return Minimum(b, {Maximum(b, {inputs[0], inputs[1]}), inputs[2]});
+    case HloOpcode::kClamp: {
+      Value min_val =
+          MaybeBroadcastScalarToShape(b, inputs[0], inputs[1].getType());
+      Value max_val =
+          MaybeBroadcastScalarToShape(b, inputs[2], inputs[1].getType());
+      return Minimum(b, {Maximum(b, {min_val, inputs[1]}), max_val});
+    }
     case HloOpcode::kAnd:
       return mlir::stablehlo::AndOp::create(b, inputs[0], inputs[1]);
     case HloOpcode::kOr:
@@ -742,12 +758,12 @@ absl::StatusOr<Value> EmitElementwise(mlir::ImplicitLocOpBuilder& b,
           mlir::stablehlo::symbolizeComparisonDirection(
               ComparisonDirectionToString(hlo.comparison_direction()))
               .value());
-    case HloOpcode::kSelect:
-      return ma::SelectOp::create(
-          b,
-          Compare(b, {inputs[0], ZerosLike(b, inputs[0])},
-                  mlir::stablehlo::ComparisonDirection::NE),
-          inputs[1], inputs[2]);
+    case HloOpcode::kSelect: {
+      Value pred = Compare(b, {inputs[0], ZerosLike(b, inputs[0])},
+                           mlir::stablehlo::ComparisonDirection::NE);
+      pred = MaybeBroadcastScalarToShape(b, pred, inputs[1].getType());
+      return ma::SelectOp::create(b, pred, inputs[1], inputs[2]);
+    }
     case HloOpcode::kReducePrecision:
       return ReducePrecision(b, hlo, inputs[0]);
     case HloOpcode::kAcos:

From 9538096ca0c61aaebc4df93053cbef99fa03361c Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Mon, 12 Oct 2026 09:37:04 -0700
Subject: [PATCH 0942/1000] Keep dynamic dimension sizes in the XLA
 ReverseSequence kernel

ReverseSequence gathers its output with indices and slice sizes built
from each dimension's static bound, so a dynamic dimension of the
input, such as one produced by tf.where and tf.gather, came out at its
bound. The output has the input's dimensions in the same order, so set
each dynamic one to the input's size.
---
 tensorflow/compiler/tests/BUILD               |  1 +
 .../tests/reverse_sequence_op_test.py         | 45 +++++++++++++++++++
 .../tf2xla/kernels/reverse_sequence_op.cc     | 16 ++++++-
 3 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD
index 93071faf080..23bd7314d3e 100644
--- a/tensorflow/compiler/tests/BUILD
+++ b/tensorflow/compiler/tests/BUILD
@@ -1533,6 +1533,7 @@ tf_xla_py_strict_test(
     ],
     deps = [
         ":xla_test",
+        "//tensorflow/python/eager:def_function",
         "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/ops:array_ops",
         "//tensorflow/python/platform:client_testlib",
diff --git a/tensorflow/compiler/tests/reverse_sequence_op_test.py b/tensorflow/compiler/tests/reverse_sequence_op_test.py
index d25cf894995..024cb929388 100644
--- a/tensorflow/compiler/tests/reverse_sequence_op_test.py
+++ b/tensorflow/compiler/tests/reverse_sequence_op_test.py
@@ -17,6 +17,7 @@
 import numpy as np
 
 from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import def_function
 from tensorflow.python.framework import dtypes
 from tensorflow.python.ops import array_ops
 from tensorflow.python.platform import test
@@ -84,6 +85,50 @@ class ReverseSequenceTest(xla_test.XLATestCase):
       for seq_dtype in self.all_types & {np.int32, np.int64}:
         self._testBasic(dtype, seq_dtype)
 
+  def _reverseSelected(self, x, seq_lengths, axis, seq_axis, batch_axis):
+    # Gathers the entries of x along axis where mask is true, which gives that
+    # dimension a size only known at run time, and reverses the result.
+    @def_function.function(jit_compile=True)
+    def f(x, seq_lengths, mask):
+      indices = array_ops.reshape(array_ops.where(mask), [-1])
+      x = array_ops.gather(x, indices, axis=axis)
+      if axis == batch_axis:
+        seq_lengths = array_ops.gather(seq_lengths, indices)
+      return array_ops.reverse_sequence(
+          x, seq_lengths, seq_axis=seq_axis, batch_axis=batch_axis)
+
+    x = np.array(x, dtype=np.int32)
+    seq_lengths = np.array(seq_lengths, dtype=np.int32)
+    with self.session() as session:
+      with self.device_scope():
+        x_ph = array_ops.placeholder(dtypes.int32, x.shape)
+        seq_lengths_ph = array_ops.placeholder(dtypes.int32, seq_lengths.shape)
+        mask_ph = array_ops.placeholder(dtypes.bool, [4])
+        output = f(x_ph, seq_lengths_ph, mask_ph)
+      return session.run(output, {
+          x_ph: x,
+          seq_lengths_ph: seq_lengths,
+          mask_ph: [True, False, True, True]
+      })
+
+  def testDynamicDimensions(self):
+    # Regression test for GitHub issue 128953: the output had the bound of a
+    # dynamic input dimension instead of its size. The mask selects 3 of 4
+    # entries, so the selected dimension has size 3 and bound 4.
+    self.assertAllEqual(
+        self._reverseSelected([[1, 2, 3, 4], [5, 6, 7, 8]], [3, 2],
+                              axis=1, seq_axis=1, batch_axis=0),
+        [[4, 3, 1], [7, 5, 8]])
+    self.assertAllEqual(
+        self._reverseSelected([[1, 5], [2, 6], [3, 7], [4, 8]], [3, 2],
+                              axis=0, seq_axis=0, batch_axis=1),
+        [[4, 7], [3, 5], [1, 8]])
+    self.assertAllEqual(
+        self._reverseSelected([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]],
+                              [3, 1, 2, 3],
+                              axis=0, seq_axis=1, batch_axis=0),
+        [[3, 2, 1], [8, 7, 9], [12, 11, 10]])
+
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc b/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc
index 5c77a4dfe29..277e2e6cd33 100644
--- a/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
 #include "xla/hlo/builder/lib/constants.h"
 #include "xla/hlo/builder/xla_builder.h"
+#include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/xla_data.pb.h"
 #include "tensorflow/core/framework/op_kernel.h"
@@ -129,8 +130,19 @@ class ReverseSequenceOp : public XlaOpKernel {
     slice_sizes[batch_dim_] = 1;
     slice_sizes[seq_dim_] = 1;
 
-    context->SetOutput(0,
-                       xla::Gather(input, start_indices, dnums, slice_sizes));
+    xla::XlaOp output = xla::Gather(input, start_indices, dnums, slice_sizes);
+    // The gather indices and slice sizes come from each dimension's static
+    // bound, so the output drops the input's dynamic sizes. Its dimensions are
+    // the input's, in the same order, so give them the input's sizes.
+    xla::Shape input_xla_shape;
+    OP_REQUIRES_VALUE(input_xla_shape, context, context->InputXlaShape(0));
+    for (int i = 0; i < input_shape.dims(); ++i) {
+      if (input_xla_shape.is_dynamic_dimension(i)) {
+        output =
+            xla::SetDimensionSize(output, xla::GetDimensionSize(input, i), i);
+      }
+    }
+    context->SetOutput(0, output);
   }
 
  private:

From f0ead6259b3125f8bc65af84ecd9587658ac91dc Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Tue, 13 Oct 2026 12:07:20 -0700
Subject: [PATCH 0943/1000] Run the dynamic ReverseSequence test eagerly and
 fix it on GPU and TPU

Use 5-entry inputs, since XLA rewrites a gather from 4 or fewer entries
into selects, which XLA:GPU's dynamic padder rejects: the seq_lengths
gather failed on the CUDA jobs that way. Also run each case in eager
mode and compare, and skip the test on TPU, which fails to compile its
dynamic dimensions with a bitcast shape size mismatch.
---
 .../tests/reverse_sequence_op_test.py         | 31 +++++++++++++------
 1 file changed, 21 insertions(+), 10 deletions(-)

diff --git a/tensorflow/compiler/tests/reverse_sequence_op_test.py b/tensorflow/compiler/tests/reverse_sequence_op_test.py
index 024cb929388..e4fa1813190 100644
--- a/tensorflow/compiler/tests/reverse_sequence_op_test.py
+++ b/tensorflow/compiler/tests/reverse_sequence_op_test.py
@@ -99,34 +99,45 @@ class ReverseSequenceTest(xla_test.XLATestCase):
 
     x = np.array(x, dtype=np.int32)
     seq_lengths = np.array(seq_lengths, dtype=np.int32)
+    mask = np.array([True, False, True, True, False])
     with self.session() as session:
       with self.device_scope():
         x_ph = array_ops.placeholder(dtypes.int32, x.shape)
         seq_lengths_ph = array_ops.placeholder(dtypes.int32, seq_lengths.shape)
-        mask_ph = array_ops.placeholder(dtypes.bool, [4])
+        mask_ph = array_ops.placeholder(dtypes.bool, mask.shape)
         output = f(x_ph, seq_lengths_ph, mask_ph)
-      return session.run(output, {
+      result = session.run(output, {
           x_ph: x,
           seq_lengths_ph: seq_lengths,
-          mask_ph: [True, False, True, True]
+          mask_ph: mask
       })
+    # Eager mode gives the same result.
+    with self.device_scope():
+      self.assertAllEqual(f(x, seq_lengths, mask), result)
+    return result
 
   def testDynamicDimensions(self):
+    if "TPU" in self.device:
+      self.skipTest("XLA:TPU fails to compile these dynamic dimensions with a "
+                    "bitcast shape size mismatch.")
     # Regression test for GitHub issue 128953: the output had the bound of a
-    # dynamic input dimension instead of its size. The mask selects 3 of 4
-    # entries, so the selected dimension has size 3 and bound 4.
+    # dynamic input dimension instead of its size. The mask selects 3 of 5
+    # entries, so the selected dimension has size 3 and bound 5. (XLA turns a
+    # gather from 4 or fewer entries into selects, which XLA:GPU doesn't
+    # support with dynamic shapes.)
     self.assertAllEqual(
-        self._reverseSelected([[1, 2, 3, 4], [5, 6, 7, 8]], [3, 2],
+        self._reverseSelected([[1, 2, 3, 4, 0], [5, 6, 7, 8, 0]], [3, 2],
                               axis=1, seq_axis=1, batch_axis=0),
         [[4, 3, 1], [7, 5, 8]])
     self.assertAllEqual(
-        self._reverseSelected([[1, 5], [2, 6], [3, 7], [4, 8]], [3, 2],
+        self._reverseSelected([[1, 5], [2, 6], [3, 7], [4, 8], [0, 0]], [3, 2],
                               axis=0, seq_axis=0, batch_axis=1),
         [[4, 7], [3, 5], [1, 8]])
     self.assertAllEqual(
-        self._reverseSelected([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]],
-                              [3, 1, 2, 3],
-                              axis=0, seq_axis=1, batch_axis=0),
+        self._reverseSelected(
+            [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12], [0, 0, 0]],
+            [3, 1, 2, 3, 1],
+            axis=0, seq_axis=1, batch_axis=0),
         [[3, 2, 1], [8, 7, 9], [12, 11, 10]])
 
 

From c863dacef33eb7f738029b38be139832889a5b05 Mon Sep 17 00:00:00 2001
From: Vishwak Thatikonda <vishwakt@gmail.com>
Date: Wed, 14 Oct 2026 22:38:09 -0700
Subject: [PATCH 0944/1000] Run the dynamic ReverseSequence test's eager check
 in eager mode

The comparison with an eager call ran in whatever mode the test did,
and in graph mode assertAllEqual evaluated the result through
test_session(), which XLATestCase doesn't support. Also only look for
dynamic dimensions when the input has any.
---
 tensorflow/compiler/tests/BUILD                        |  1 +
 tensorflow/compiler/tests/reverse_sequence_op_test.py  |  6 ++++--
 .../compiler/tf2xla/kernels/reverse_sequence_op.cc     | 10 ++++++----
 3 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/tensorflow/compiler/tests/BUILD b/tensorflow/compiler/tests/BUILD
index 23bd7314d3e..970a177e63b 100644
--- a/tensorflow/compiler/tests/BUILD
+++ b/tensorflow/compiler/tests/BUILD
@@ -1533,6 +1533,7 @@ tf_xla_py_strict_test(
     ],
     deps = [
         ":xla_test",
+        "//tensorflow/python/eager:context",
         "//tensorflow/python/eager:def_function",
         "//tensorflow/python/framework:dtypes",
         "//tensorflow/python/ops:array_ops",
diff --git a/tensorflow/compiler/tests/reverse_sequence_op_test.py b/tensorflow/compiler/tests/reverse_sequence_op_test.py
index e4fa1813190..23abaedbd00 100644
--- a/tensorflow/compiler/tests/reverse_sequence_op_test.py
+++ b/tensorflow/compiler/tests/reverse_sequence_op_test.py
@@ -17,6 +17,7 @@
 import numpy as np
 
 from tensorflow.compiler.tests import xla_test
+from tensorflow.python.eager import context
 from tensorflow.python.eager import def_function
 from tensorflow.python.framework import dtypes
 from tensorflow.python.ops import array_ops
@@ -111,8 +112,9 @@ class ReverseSequenceTest(xla_test.XLATestCase):
           seq_lengths_ph: seq_lengths,
           mask_ph: mask
       })
-    # Eager mode gives the same result.
-    with self.device_scope():
+    # Eager execution of the compiled function gives the same result, also
+    # when the test itself runs in graph mode.
+    with context.eager_mode(), self.device_scope():
       self.assertAllEqual(f(x, seq_lengths, mask), result)
     return result
 
diff --git a/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc b/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc
index 277e2e6cd33..2a88b1d49cd 100644
--- a/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc
+++ b/tensorflow/compiler/tf2xla/kernels/reverse_sequence_op.cc
@@ -136,10 +136,12 @@ class ReverseSequenceOp : public XlaOpKernel {
     // the input's, in the same order, so give them the input's sizes.
     xla::Shape input_xla_shape;
     OP_REQUIRES_VALUE(input_xla_shape, context, context->InputXlaShape(0));
-    for (int i = 0; i < input_shape.dims(); ++i) {
-      if (input_xla_shape.is_dynamic_dimension(i)) {
-        output =
-            xla::SetDimensionSize(output, xla::GetDimensionSize(input, i), i);
+    if (input_xla_shape.is_dynamic()) {
+      for (int i = 0; i < input_shape.dims(); ++i) {
+        if (input_xla_shape.is_dynamic_dimension(i)) {
+          output =
+              xla::SetDimensionSize(output, xla::GetDimensionSize(input, i), i);
+        }
       }
     }
     context->SetOutput(0, output);

From fe714b5102aaead8d28d12b03b1dd531cdddc000 Mon Sep 17 00:00:00 2001
From: Ce Zheng CA <zce@google.com>
Date: Fri, 16 Oct 2026 01:30:53 -0700
Subject: [PATCH 0945/1000] [XLA] Move `HandleSelect` and `HandlePad` from
 `HloEvaluatorTypedVisitor` to `HloEvaluator`

`HloEvaluatorTypedVisitor<ReturnT, ElementwiseT>` is instantiated 32 times across `xla/hlo/evaluator/hlo_evaluator_typed_visitor_*.cc`, even though `HandleSelect` and `HandlePad` perform no arithmetic on `ReturnT` and only copy element bytes based on predicate values or padding index bounds.

Move `HandleSelect` and `HandlePad` onto `HloEvaluator` in `xla/hlo/evaluator/hlo_evaluator.cc` using `primitive_util::ByteWidthTypeSwitch` (`1, 2, 4, 8, 16` bytes) with `PopulateLinearInplaceParallel` and `PopulateInplaceParallel`. This avoids instantiating `HandleSelect`, `HandlePad`, and `ElementwiseTernaryOpImpl<bool, ReturnT, ReturnT>` across all 32 visitor specializations and removes the `EvaluatePad` `absl::FunctionRef` callback indirection from `HloEvaluator`.

**Runtime & Memory Performance Impact:**
- **CPU:** Neutral (`BM_EvaluatePad/4` and `BM_EvaluatePad/64` in `//xla/hlo/evaluator:hlo_evaluator_test` are unchanged within noise, and `HandleSelect` / `HandlePad` inline the fixed-byte-width element copies directly into `PopulateLinearInplaceParallel` and `PopulateInplaceParallel` without `absl::FunctionRef` callback indirection).
- **Memory:** Neutral (zero additional allocations).

PiperOrigin-RevId: 1000688306
---
 .../xla/xla/hlo/evaluator/hlo_evaluator.cc    | 110 ++++++++++++++++--
 .../xla/xla/hlo/evaluator/hlo_evaluator.h     |   9 +-
 .../evaluator/hlo_evaluator_typed_visitor.h   |  58 ---------
 3 files changed, 100 insertions(+), 77 deletions(-)

diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
index ac014cb48ac..b89c6912ff3 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.cc
@@ -3841,7 +3841,59 @@ absl::Status HloEvaluator::HandleSelect(const HloInstruction* select) {
     return absl::OkStatus();
   }
 
-  return DefaultAction(select);
+  CHECK(select->shape().IsArray());
+  Shape shape = GetShapeWithLayout(select->shape());
+  const bool is_on_false_scalar = ShapeUtil::IsScalar(on_false.shape());
+  TF_RET_CHECK(ShapeUtil::SameDimensions(shape, pred.shape()));
+  TF_RET_CHECK(ShapeUtil::SameDimensions(shape, on_true.shape()));
+  TF_RET_CHECK(is_on_false_scalar ||
+               ShapeUtil::SameDimensions(shape, on_false.shape()));
+
+  Literal result(shape);
+  const Layout& pred_layout = pred.shape().layout();
+  const Layout& on_true_layout = on_true.shape().layout();
+  const Layout& on_false_layout = on_false.shape().layout();
+  const bool same_layout = LayoutUtil::Equal(on_true_layout, shape.layout()) &&
+                           LayoutUtil::Equal(pred_layout, shape.layout()) &&
+                           (is_on_false_scalar ||
+                            LayoutUtil::Equal(on_false_layout, shape.layout()));
+
+  ABSL_RETURN_IF_ERROR(primitive_util::ByteWidthTypeSwitch(
+      [&](auto primitive_type_constant) -> absl::Status {
+        using NativeT = NativeTypeOf<primitive_type_constant>;
+        const bool* pred_data = pred.data<bool>().data();
+        const auto* on_true_data =
+            static_cast<const NativeT*>(on_true.untyped_data());
+        const auto* on_false_data =
+            static_cast<const NativeT*>(on_false.untyped_data());
+        if (same_layout) {
+          return result.PopulateLinearInplaceParallel(
+              [&](void* dest, int64_t linear_index, int) {
+                *static_cast<NativeT*>(dest) =
+                    pred_data[linear_index]
+                        ? on_true_data[linear_index]
+                        : on_false_data[is_on_false_scalar ? 0 : linear_index];
+              });
+        }
+        return result.PopulateInplaceParallel(
+            [&](void* dest, absl::Span<const int64_t> multi_index, int) {
+              if (pred.Get<bool>(multi_index)) {
+                *static_cast<NativeT*>(dest) =
+                    on_true_data[IndexUtil::MultidimensionalIndexToLinearIndex(
+                        on_true.shape(), multi_index)];
+              } else {
+                *static_cast<NativeT*>(dest) = on_false_data
+                    [is_on_false_scalar
+                         ? 0
+                         : IndexUtil::MultidimensionalIndexToLinearIndex(
+                               on_false.shape(), multi_index)];
+              }
+            });
+      },
+      shape.element_type()));
+
+  SetEvaluatedLiteralFor(select, std::move(result));
+  return absl::OkStatus();
 }
 
 namespace {
@@ -6064,13 +6116,7 @@ absl::Status HloEvaluator::EvaluateScaledDot(
   return absl::OkStatus();
 }
 
-absl::Status HloEvaluator::EvaluatePad(
-    const HloInstruction* pad,
-    absl::FunctionRef<absl::Status(
-        const Literal& evaluated_operand, const Literal& evaluated_padding,
-        absl::Span<int64_t> target_index, absl::Span<const int64_t> zero_base,
-        absl::Span<const int64_t> step, Literal* result)>
-        evaluate_impl) {
+absl::Status HloEvaluator::HandlePad(const HloInstruction* pad) {
   CHECK(pad->operand(0)->shape().IsArray());
   // Padding value must be scalar.
   CHECK(ShapeUtil::IsScalar(pad->operand(1)->shape()));
@@ -6117,10 +6163,52 @@ absl::Status HloEvaluator::EvaluatePad(
   std::vector<int64_t> zero_base(evaluated_operand.shape().dimensions().size(),
                                  0);
   std::vector<int64_t> step(evaluated_operand.shape().dimensions().size(), 1);
+  const PaddingConfig& pad_config = pad->padding_config();
+
+  ABSL_RETURN_IF_ERROR(primitive_util::ByteWidthTypeSwitch(
+      [&](auto primitive_type_constant) -> absl::Status {
+        using NativeT = NativeTypeOf<primitive_type_constant>;
+        const NativeT scalar =
+            *static_cast<const NativeT*>(evaluated_padding_ptr->untyped_data());
+        ABSL_RETURN_IF_ERROR(result.PopulateLinearInplaceParallel(
+            [&scalar](void* dest, int64_t /*linear_index*/, int /*thread_id*/) {
+              *static_cast<NativeT*>(dest) = scalar;
+            }));
+
+        const auto* operand_data =
+            static_cast<const NativeT*>(evaluated_operand.untyped_data());
+        auto* result_data = static_cast<NativeT*>(result.untyped_data());
+        auto func = [&](absl::Span<const int64_t> input_index) {
+          for (auto i = 0; i < input_index.size(); ++i) {
+            // Interior padding occurs logically before edge padding, so in
+            // the case of negative edge padding elements are removed from the
+            // interior-padded operand.
+            target_index[i] =
+                pad_config.dimensions(i).edge_padding_low() +
+                input_index[i] *
+                    (pad_config.dimensions(i).interior_padding() + 1);
+
+            // Account for negative low and high padding: skip assignment if
+            // the any target index is out of range.
+            if (!(target_index[i] >= 0 &&
+                  target_index[i] < pad->shape().dimensions(i))) {
+              return true;
+            }
+          }
+          result_data[IndexUtil::MultidimensionalIndexToLinearIndex(
+              result.shape(), target_index)] =
+              operand_data[IndexUtil::MultidimensionalIndexToLinearIndex(
+                  evaluated_operand.shape(), input_index)];
+          return true;
+        };
+
+        ShapeUtil::ForEachIndexNoStatus(evaluated_operand.shape(), zero_base,
+                                        evaluated_operand.shape().dimensions(),
+                                        step, func);
+        return absl::OkStatus();
+      },
+      result_type));
 
-  ABSL_RETURN_IF_ERROR(evaluate_impl(evaluated_operand, *evaluated_padding_ptr,
-                                absl::MakeSpan(target_index), zero_base, step,
-                                &result));
   SetEvaluatedLiteralFor(pad, std::move(result));
   return absl::OkStatus();
 }
diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h
index 58e43dbf27b..618bf67a85d 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator.h
@@ -580,14 +580,6 @@ class HloEvaluator : public ConstDfsHloVisitorWithDefault,
           int64_t total_contraction_size, Literal* result)>
           evaluate_impl);
 
-  absl::Status EvaluatePad(
-      const HloInstruction* pad,
-      absl::FunctionRef<absl::Status(
-          const Literal& evaluated_operand, const Literal& evaluated_padding,
-          absl::Span<int64_t> target_index, absl::Span<const int64_t> zero_base,
-          absl::Span<const int64_t> step, Literal* result)>
-          evaluate_impl);
-
   // Make HloEvaluatorTypedVisitor a friend because it is logically part of this
   // class.
   //
@@ -666,6 +658,7 @@ class HloEvaluator : public ConstDfsHloVisitorWithDefault,
   absl::Status HandleCustomCall(const HloInstruction* custom_call) override;
   absl::Status HandleOptimizationBarrier(
       const HloInstruction* optimization_barrier) override;
+  absl::Status HandlePad(const HloInstruction* pad) override;
 
   // Unsupported HLOs, note some of them (such as BatchNorm*) are typically
   // expanded in a semantic-preserving way into other HLOs by adding expansion
diff --git a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h
index 98125e7e21a..c7600c27290 100644
--- a/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h
+++ b/third_party/xla/xla/hlo/evaluator/hlo_evaluator_typed_visitor.h
@@ -891,19 +891,6 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
     return UnsupportedTypeError(clamp);
   }
 
-  absl::Status HandleSelect(const HloInstruction* select) override {
-    CHECK(!ShapeUtil::IsScalar(select->operand(0)->shape()));
-    CHECK(select->shape().IsArray());
-    auto select_op = [](bool pred, ReturnT on_true, ReturnT on_false) {
-      return pred ? on_true : on_false;
-    };
-    ABSL_ASSIGN_OR_RETURN(Literal literal,
-                     (ElementwiseTernaryOp<bool, ReturnT, ReturnT>(
-                         select, std::move(select_op))));
-    parent_->SetEvaluatedLiteralFor(select, std::move(literal));
-    return absl::OkStatus();
-  }
-
   absl::Status HandleConvolutionWithLiterals(
       const HloInstruction* conv, const Literal& lhs_literal,
       const Literal& rhs_literal, const Shape& window_shape,
@@ -1593,51 +1580,6 @@ class HloEvaluatorTypedVisitor : public ConstDfsHloVisitorWithDefault {
   }
 
  public:
-  absl::Status HandlePad(const HloInstruction* pad) override {
-    return parent_->EvaluatePad(
-        pad,
-        [pad](const Literal& evaluated_operand,
-              const Literal& evaluated_padding,
-              absl::Span<int64_t> target_index,
-              absl::Span<const int64_t> zero_base,
-              absl::Span<const int64_t> step, Literal* result) -> absl::Status {
-          const ReturnT scalar = evaluated_padding.Get<ReturnT>({});
-          ABSL_RETURN_IF_ERROR(result->PopulateLinearParallel<ReturnT>(
-              [&scalar](int64_t linear_index, int) { return scalar; }));
-
-          // Loop through each element of the operand, assign them to the
-          // corresponding index of the resulting padded literal.
-          const PaddingConfig& pad_config = pad->padding_config();
-
-          auto func = [&](absl::Span<const int64_t> input_index) {
-            for (auto i = 0; i < input_index.size(); ++i) {
-              // Interior padding occurs logically before edge padding, so in
-              // the case of negative edge padding elements are removed from the
-              // interior-padded operand.
-              target_index[i] =
-                  pad_config.dimensions(i).edge_padding_low() +
-                  input_index[i] *
-                      (pad_config.dimensions(i).interior_padding() + 1);
-
-              // Account for negative low and high padding: skip assignment if
-              // the any target index is out of range.
-              if (!(target_index[i] >= 0 &&
-                    target_index[i] < pad->shape().dimensions(i))) {
-                return true;
-              }
-            }
-            result->Set<ReturnT>(target_index,
-                                 evaluated_operand.Get<ReturnT>(input_index));
-            return true;
-          };
-
-          ShapeUtil::ForEachIndexNoStatus(
-              evaluated_operand.shape(), zero_base,
-              evaluated_operand.shape().dimensions(), step, func);
-          return absl::OkStatus();
-        });
-  }
-
   absl::Status HandleClz(const HloInstruction* clz) override {
     // Enable CLZ only for integer types.
     if constexpr (std::is_integral_v<ElementwiseT> &&

From 7bc5c49c2721fc7c6e0292e9b9ec29d765c8289d Mon Sep 17 00:00:00 2001
From: Keshav <keshavb@nvidia.com>
Date: Fri, 16 Oct 2026 01:34:16 -0700
Subject: [PATCH 0946/1000] PR #50350: [XLA:GPU] Compile cuDNN graphs in
 parallel-thread computations

Imported from GitHub PR https://github.com/openxla/xla/pull/50350

cudnn compilation pass is only invoked for the `main` HLO execution thread. The compute_on wrapper now assigns its computations to `parallel` HLO execution thread, so without this fix any compute_on region with cudnn call(s) fails during cudnn thunk initialization.
Copybara import of the project:

--
2cf70c3d809c6a4f246a5353c49f0b3a8fe18c3e by Keshav <keshavb@nvidia.com>:

[XLA:GPU] Compile cuDNN graphs in parallel-thread computations

Merging this change closes #50350

PiperOrigin-RevId: 1000689850
---
 .../xla/xla/backends/gpu/codegen/BUILD        |  2 +
 .../xla/backends/gpu/codegen/cudnn_test.cc    | 56 +++++++++++++++++++
 third_party/xla/xla/service/gpu/BUILD         |  1 +
 .../xla/xla/service/gpu/nvptx_compiler.cc     | 13 +++--
 4 files changed, 67 insertions(+), 5 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/BUILD b/third_party/xla/xla/backends/gpu/codegen/BUILD
index bf6c3dc3afc..0580f1faff3 100644
--- a/third_party/xla/xla/backends/gpu/codegen/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/BUILD
@@ -96,6 +96,7 @@ xla_test(
         "//xla/service:dump",
         "//xla/service:hlo_module_config",
         "//xla/service:pattern_matcher",
+        "//xla/service/gpu:cublas_cudnn",
         "//xla/service/gpu:cudnn_support_utils",
         "//xla/service/gpu:ir_emission_utils",
         "//xla/stream_executor:device_description",
@@ -107,6 +108,7 @@ xla_test(
         "//xla/tests:hlo_interpreter_reference_mixin",
         "//xla/tsl/platform:env",
         "//xla/tsl/platform:test",
+        "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
         "@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 f83c689216d..dca40ad9e13 100644
--- a/third_party/xla/xla/backends/gpu/codegen/cudnn_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/cudnn_test.cc
@@ -21,6 +21,7 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/algorithm/container.h"
 #include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
@@ -37,12 +38,16 @@ limitations under the License.
 #include "xla/comparison_util.h"
 #include "xla/debug_options_flags.h"
 #include "xla/error_spec.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/testlib/filecheck.h"
 #include "xla/hlo/testlib/pattern_matcher_gmock.h"
 #include "xla/hlo/testlib/verified_hlo_module.h"
 #include "xla/primitive_util.h"
 #include "xla/service/dump.h"
+#include "xla/service/gpu/cublas_cudnn.h"
 #include "xla/service/gpu/cudnn_support_utils.h"
 #include "xla/service/gpu/ir_emission_utils.h"
 #include "xla/service/hlo_module_config.h"
@@ -416,6 +421,57 @@ ENTRY e {
                             ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
 }
 
+class CuDnnFusionStreamAnnotationTest : public CuDnnFusionExecutionTest {
+ public:
+  DebugOptions GetDebugOptionsForTest() const override {
+    DebugOptions debug_options =
+        CuDnnFusionExecutionTest::GetDebugOptionsForTest();
+    debug_options.set_xla_gpu_experimental_stream_annotation(true);
+    return debug_options;
+  }
+};
+
+TEST_F(CuDnnFusionStreamAnnotationTest,
+       FusionInStreamAnnotatedCallExecutesCorrectly) {
+  constexpr absl::string_view kHloText = R"(
+fusion1 {
+  p0 = f32[32,96] parameter(0)
+  p1 = f32[96,64] parameter(1)
+  ROOT r = f32[32,64] dot(p0, p1),
+    lhs_contracting_dims={1}, rhs_contracting_dims={0}
+}
+
+annotated {
+  p0 = f32[32,96] parameter(0)
+  p1 = f32[96,64] parameter(1)
+  ROOT f = f32[32,64] fusion(p0, p1), kind=kCustom, calls=fusion1,
+    backend_config={"fusion_backend_config": {kind: "__cudnn$fusion"}}
+}
+
+ENTRY e {
+  p0 = f32[32,96] parameter(0)
+  p1 = f32[96,64] parameter(1)
+  ROOT c = f32[32,64] call(p0, p1), to_apply=annotated,
+    frontend_attributes={_xla_stream_annotation="1",inlineable="false"}
+})";
+
+  // The stream-annotated call is wrapped into an async computation on the
+  // parallel execution thread, so the cuDNN fusion must be compiled there.
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> optimized,
+                       GetOptimizedModule(kHloText));
+  ASSERT_TRUE(absl::c_any_of(
+      optimized->computations(), [](const HloComputation* computation) {
+        return computation->execution_thread() ==
+                   HloInstruction::kParallelExecutionThread &&
+               absl::c_any_of(computation->instructions(),
+                              [](const HloInstruction* instruction) {
+                                return IsCudnnFusion(*instruction);
+                              });
+      }));
+
+  EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
+}
+
 TEST_F(CuDnnFusionFileCheckTest, VectorTensorMultiplicationWorksCorrectly) {
   const std::string kHloText = R"(
 f {
diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 66860a0c093..f13ebab73d1 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -2498,6 +2498,7 @@ cc_library(
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
diff --git a/third_party/xla/xla/service/gpu/nvptx_compiler.cc b/third_party/xla/xla/service/gpu/nvptx_compiler.cc
index 516681fedb4..401cf3a760b 100644
--- a/third_party/xla/xla/service/gpu/nvptx_compiler.cc
+++ b/third_party/xla/xla/service/gpu/nvptx_compiler.cc
@@ -29,6 +29,7 @@ limitations under the License.
 
 #include "absl/algorithm/container.h"
 #include "absl/base/call_once.h"
+#include "absl/container/flat_hash_set.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
@@ -419,15 +420,17 @@ absl::Status NVPTXCompiler::RunCudnnCompilerPasses(
   });
   const se::DeviceDescription& gpu_device_info =
       gpu_target_config.device_description;
+  // Stream-annotated calls are wrapped into async computations on the parallel
+  // execution thread, and their cuDNN graphs must be compiled as well.
+  const absl::flat_hash_set<absl::string_view> execution_threads = {
+      HloInstruction::kMainExecutionThread,
+      HloInstruction::kParallelExecutionThread};
   CuDnnFusionCompiler fusion_compiler(dnn_support, gpu_device_info,
                                       *dnn_compiled_graphs);
-  ABSL_RETURN_IF_ERROR(
-      fusion_compiler.Run(module, {HloInstruction::kMainExecutionThread})
-          .status());
+  ABSL_RETURN_IF_ERROR(fusion_compiler.Run(module, execution_threads).status());
   CuDnnCustomCallCompiler call_compiler(dnn_support, gpu_device_info,
                                         *dnn_compiled_graphs);
-  return call_compiler.Run(module, {HloInstruction::kMainExecutionThread})
-      .status();
+  return call_compiler.Run(module, execution_threads).status();
 }
 
 namespace {

From 3afba85c61e14621efd791b22c5ad46c0fc8d58c Mon Sep 17 00:00:00 2001
From: Aditya Sharma <adisharma@google.com>
Date: Fri, 16 Oct 2026 02:18:27 -0700
Subject: [PATCH 0947/1000] Hold the GIL only briefly when consuming Python
 tracer events.

`PythonHookContext::Consume()` held the GIL for its whole run while it built every event's name, released each event's Python references, and freed the event queues. With continuous profiling, which calls `Consume()` periodically, draining a few hundred thousand events held the GIL for about 200 ms. That stalled the host's Python threads and, with them, accelerator dispatch.

`Consume()` now holds the GIL only for two short sections:

1. With the GIL: swap each thread's completed events out of the shared buffers. After `Stop()`, the calls still in progress are ended and taken too, as before.
2. Without the GIL: move the events' Python references into one count per object. This reads only pointer values, never the objects.
3. With the GIL: copy the text of each distinct object, then drop all references to it. If the interpreter is shutting down, skip this step; the references leak and the names show `<unknown>`.
4. Without the GIL: build the event names from the copied text and free the queues.

The time spent holding the GIL now grows with the number of distinct strings rather than the number of events, and Python strings are still read only while the GIL is held. `PythonTracer::Consume()` no longer wraps the call in its own GIL acquisition, which would have kept the GIL held throughout. The only change to the recording hot path is one `bool` in existing padding of `PythonTraceEntry`.

PiperOrigin-RevId: 1000710848
---
 .../backends/profiler/cpu/python_tracer.cc    |   2 -
 .../python/profiler/internal/python_hooks.cc  | 202 ++++++++++++------
 .../python/profiler/internal/python_hooks.h   |  17 +-
 3 files changed, 154 insertions(+), 67 deletions(-)

diff --git a/third_party/xla/xla/backends/profiler/cpu/python_tracer.cc b/third_party/xla/xla/backends/profiler/cpu/python_tracer.cc
index 38dce915935..3ef437054d6 100644
--- a/third_party/xla/xla/backends/profiler/cpu/python_tracer.cc
+++ b/third_party/xla/xla/backends/profiler/cpu/python_tracer.cc
@@ -110,9 +110,7 @@ absl::StatusOr<tsl::profiler::ConsumeResult> PythonTracer::Consume() {
     chunk.consumed_data = PythonHooks::GetSingleton()->Consume();
   } else if (context_) {
     if (Py_IsInitialized()) {
-      PyGILState_STATE gil_state = PyGILState_Ensure();
       chunk.consumed_data = context_->Consume();
-      PyGILState_Release(gil_state);
     }
   }
 
diff --git a/third_party/xla/xla/python/profiler/internal/python_hooks.cc b/third_party/xla/xla/python/profiler/internal/python_hooks.cc
index 32c96aec3d3..49a2eadf7ab 100644
--- a/third_party/xla/xla/python/profiler/internal/python_hooks.cc
+++ b/third_party/xla/xla/python/profiler/internal/python_hooks.cc
@@ -14,12 +14,17 @@ limitations under the License.
 ==============================================================================*/
 #include "xla/python/profiler/internal/python_hooks.h"
 
+#include <atomic>
+#include <cstddef>
 #include <cstdint>
 #include <deque>
+#include <initializer_list>
+#include <stack>
 #include <string>
 #include <utility>
 #include <vector>
 
+#include "absl/container/flat_hash_map.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/strings/str_cat.h"
@@ -70,36 +75,80 @@ bool ThreadingSetProfile(PyObject* callback) {
   return res != nullptr;
 }
 
-const char* PyUnicodeToStringOrUnknown(PyObject* obj) {
+// Returns the UTF-8 text of `obj`, or "<unknown>" if `obj` is not a str.
+// Requires the GIL; the result is valid while `obj` is alive.
+absl::string_view PyUnicodeToStringOrUnknown(PyObject* obj) {
   if (obj != nullptr && PyUnicode_Check(obj)) {
-    const char* str = PyUnicode_AsUTF8(obj);
+    Py_ssize_t size = 0;
+    const char* str = PyUnicode_AsUTF8AndSize(obj, &size);
     if (str != nullptr) {
-      return str;
+      return absl::string_view(str, static_cast<size_t>(size));
     }
     PyErr_Clear();
   }
   return "<unknown>";
 }
 
-std::string GetEventName(PyObject* co_filename, PyObject* co_name,
-                         int co_firstlineno) {
-  absl::string_view filename = PyUnicodeToStringOrUnknown(co_filename);
-  absl::string_view function = PyUnicodeToStringOrUnknown(co_name);
-
-  return absl::StrCat("$", tsl::io::Basename(filename), ":", co_firstlineno,
-                      " ", function);
-}
-
-std::string GetEventName(absl::string_view method_name, PyObject* module) {
+// Returns the trace event name of `event`. `to_string` returns the text of one
+// of the event's Python string objects.
+template <typename ToString>
+std::string GetEventName(const PythonTraceEntry& event, ToString to_string) {
+  if (event.co_filename != nullptr) {
+    absl::string_view filename = to_string(event.co_filename);
+    absl::string_view function = to_string(event.co_name);
+    return absl::StrCat("$", tsl::io::Basename(filename), ":",
+                        event.co_firstlineno, " ", function);
+  }
   // Python stack does not have a filename/line_no for native calls.
   // Use module name and function/method name instead.
-  absl::string_view filename = PyUnicodeToStringOrUnknown(module);
-  if (!method_name.empty()) {
-    return absl::StrCat("$", filename, " ", method_name);
+  absl::string_view filename = to_string(event.m_module);
+  if (!event.method_name.empty()) {
+    return absl::StrCat("$", filename, " ", event.method_name);
   }
   return "$<unknown>";
 }
 
+// Ends the calls still open in `active` at `end_time_ns` and moves them to
+// `events`, innermost first.
+void CloseActiveEvents(uint64_t end_time_ns,
+                       std::stack<PythonTraceEntry>& active,
+                       std::deque<PythonTraceEntry>& events) {
+  while (!active.empty()) {
+    PythonTraceEntry& entry = active.top();
+    entry.end_time_ns = end_time_ns;
+    events.push_back(std::move(entry));
+    active.pop();
+  }
+}
+
+// Moves the Python references held by `event` into `ref_counts`. Reads only the
+// pointer values, so it does not need the GIL.
+void TakeReferences(PythonTraceEntry& event,
+                    absl::flat_hash_map<PyObject*, size_t>& ref_counts) {
+  for (PyObject* obj : {event.co_filename, event.co_name, event.m_module}) {
+    if (obj != nullptr) {
+      ++ref_counts[obj];
+    }
+  }
+  event.owns_python_refs = false;
+}
+
+// Drops `count` references to `obj`. Requires the GIL.
+void ReleaseReferences(PyObject* obj, size_t count) {
+#if !defined(Py_GIL_DISABLED) && !defined(Py_LIMITED_API) && \
+    !defined(Py_REF_DEBUG)
+  // Drop all but the last reference in one step, then let Py_DECREF free the
+  // object if that was the last one. Neither call changes immortal objects.
+  Py_SET_REFCNT(obj, Py_REFCNT(obj) - static_cast<Py_ssize_t>(count - 1));
+  Py_DECREF(obj);
+#else
+  // Free-threaded, limited-API, and ref-debug builds: one Py_DECREF each.
+  for (size_t i = 0; i < count; ++i) {
+    Py_DECREF(obj);
+  }
+#endif
+}
+
 void AddEventToXLine(const PythonTraceEntry& event,
                      tsl::profiler::XLineBuilder* line,
                      tsl::profiler::XPlaneBuilder* plane) {
@@ -178,10 +227,7 @@ void ForEachThread(PyThreadState* curr_thread, ForEachThreadFunc&& callback) {
 }
 
 std::string PythonTraceEntry::Name() const {
-  if (co_filename) {
-    return GetEventName(co_filename, co_name, co_firstlineno);
-  }
-  return GetEventName(method_name, m_module);
+  return GetEventName(*this, PyUnicodeToStringOrUnknown);
 }
 
 PythonHooks* PythonHooks::GetSingleton() {
@@ -251,16 +297,28 @@ void PythonHookContext::Stop() {
 }
 
 std::vector<PerThreadConsumeData> PythonHookContext::Consume() {
-  std::vector<PerThreadConsumeData> consumed_data;
+  if (!Py_IsInitialized()) {
+    return {};
+  }
 
+  struct HarvestedEvents {
+    HarvestedEvents() = default;
+    HarvestedEvents(const HarvestedEvents&) = delete;
+    HarvestedEvents& operator=(const HarvestedEvents&) = delete;
+    HarvestedEvents(HarvestedEvents&&) = default;
+    HarvestedEvents& operator=(HarvestedEvents&&) = default;
+
+    int64_t thread_id = 0;
+    std::deque<PythonTraceEntry> events;
+  };
+  std::vector<HarvestedEvents> harvested;
+  const bool stopped = stopped_.load(std::memory_order_acquire);
+
+  // 1. With the GIL: move the per-thread event queues out of `entry_shards_`.
+  // After Stop(), calls still on the stack are ended now and moved out too.
   {
-    PyGILState_STATE gil_state;
-    bool has_gil = false;
-    if (Py_IsInitialized()) {
-      gil_state = PyGILState_Ensure();
-      has_gil = true;
-    }
-
+    PyGILState_STATE gil_state = PyGILState_Ensure();
+    const uint64_t now = tsl::profiler::GetCurrentTimeNanos();
     for (EntryShard& shard : entry_shards_) {
 #ifdef Py_GIL_DISABLED
       absl::MutexLock lock(shard.mu);
@@ -268,46 +326,68 @@ std::vector<PerThreadConsumeData> PythonHookContext::Consume() {
       DCHECK(PyGILState_Check());
 #endif  // Py_GIL_DISABLED
       // NOLINTNEXTLINE
-      for (auto& it : shard.entries) {
-        int64_t thread_id = it.first;
-        PerThreadEvents& thread_events = it.second;
-        VLOG(1) << "Consuming " << thread_events.completed.size() << ":"
-                << thread_events.active.size() << " events on thread "
-                << thread_id;
-
-        PerThreadConsumeData thread_data;
-        thread_data.thread_id = thread_id;
-        thread_data.events.reserve(thread_events.completed.size());
-        for (const PythonTraceEntry& event : thread_events.completed) {
-          thread_data.events.push_back(
-              {event.Name(), event.start_time_ns, event.end_time_ns});
+      for (auto& [thread_id, thread_events] : shard.entries) {
+        HarvestedEvents& h = harvested.emplace_back();
+        h.thread_id = thread_id;
+        h.events.swap(thread_events.completed);
+        if (stopped && options_.include_incomplete_events) {
+          CloseActiveEvents(now, thread_events.active, h.events);
+          CloseActiveEvents(now, thread_events.active_c, h.events);
         }
-        thread_events.completed.clear();
-
-        if (stopped_ && options_.include_incomplete_events) {
-          uint64_t now = tsl::profiler::GetCurrentTimeNanos();
-          while (!thread_events.active.empty()) {
-            PythonTraceEntry& event = thread_events.active.top();
-            thread_data.events.push_back(
-                {event.Name(), event.start_time_ns, now});
-            thread_events.active.pop();
-          }
-          while (!thread_events.active_c.empty()) {
-            PythonTraceEntry& event = thread_events.active_c.top();
-            thread_data.events.push_back(
-                {event.Name(), event.start_time_ns, now});
-            thread_events.active_c.pop();
-          }
-        }
-        consumed_data.push_back(std::move(thread_data));
       }
     }
+    PyGILState_Release(gil_state);
+  }
 
-    if (has_gil) {
-      PyGILState_Release(gil_state);
+  // 2. Without the GIL: take over the events' Python references, counted per
+  // object. This reads pointer values only, never the objects.
+  absl::flat_hash_map<PyObject*, size_t> ref_counts;
+  for (HarvestedEvents& h : harvested) {
+    for (PythonTraceEntry& event : h.events) {
+      TakeReferences(event, ref_counts);
     }
   }
 
+  // 3. With the GIL: copy the text of each distinct object, then drop all of
+  // its references at once. If the interpreter is shutting down, skip this and
+  // leak the references; the names then read "<unknown>".
+  absl::flat_hash_map<PyObject*, std::string> names;
+  if (!ref_counts.empty() && Py_IsInitialized()) {
+    names.reserve(ref_counts.size());
+    PyGILState_STATE gil_state = PyGILState_Ensure();
+    // NOLINTNEXTLINE
+    for (const auto& [obj, count] : ref_counts) {
+      names.try_emplace(obj, PyUnicodeToStringOrUnknown(obj));
+      ReleaseReferences(obj, count);
+    }
+    PyGILState_Release(gil_state);
+  }
+
+  // 4. Without the GIL: build the event names from `names`, then free the
+  // queues. No entry owns a reference any more.
+  auto copied_text = [&names](PyObject* obj) -> absl::string_view {
+    auto it = names.find(obj);
+    if (it == names.end()) {
+      return "<unknown>";
+    }
+    return it->second;
+  };
+  std::vector<PerThreadConsumeData> consumed_data;
+  consumed_data.reserve(harvested.size());
+  for (HarvestedEvents& h : harvested) {
+    // NOLINTNEXTLINE(clang-diagnostic-pre-c++20-compat)
+    VLOG(1) << "Consuming " << h.events.size() << " events on thread "
+            << h.thread_id;
+    PerThreadConsumeData& thread_data = consumed_data.emplace_back();
+    thread_data.thread_id = h.thread_id;
+    thread_data.events.reserve(h.events.size());
+    for (const PythonTraceEntry& event : h.events) {
+      thread_data.events.push_back({GetEventName(event, copied_text),
+                                    event.start_time_ns, event.end_time_ns});
+    }
+    h.events.clear();
+  }
+
   return consumed_data;
 }
 
diff --git a/third_party/xla/xla/python/profiler/internal/python_hooks.h b/third_party/xla/xla/python/profiler/internal/python_hooks.h
index 18f1ade33d3..16e77523c73 100644
--- a/third_party/xla/xla/python/profiler/internal/python_hooks.h
+++ b/third_party/xla/xla/python/profiler/internal/python_hooks.h
@@ -21,6 +21,7 @@ limitations under the License.
 #endif
 
 #include <array>
+#include <atomic>
 #include <cstddef>
 #include <cstdint>
 #include <deque>
@@ -94,15 +95,18 @@ struct PythonTraceEntry {
   }
 
   ~PythonTraceEntry() {
-    Py_XDECREF(co_filename);
-    Py_XDECREF(co_name);
-    Py_XDECREF(m_module);
+    if (owns_python_refs) {
+      Py_XDECREF(co_filename);
+      Py_XDECREF(co_name);
+      Py_XDECREF(m_module);
+    }
   }
 
   PythonTraceEntry(PythonTraceEntry&& other) noexcept {
     start_time_ns = other.start_time_ns;
     end_time_ns = other.end_time_ns;
     co_firstlineno = other.co_firstlineno;
+    owns_python_refs = other.owns_python_refs;
     co_filename = other.co_filename;
     co_name = other.co_name;
     method_name = std::move(other.method_name);
@@ -113,6 +117,7 @@ struct PythonTraceEntry {
     other.m_module = nullptr;
   }
 
+  // Returns the trace event name. Requires the GIL.
   std::string Name() const;
 
   uint64_t start_time_ns;
@@ -120,6 +125,10 @@ struct PythonTraceEntry {
   PyObject* co_filename = nullptr;
   PyObject* co_name = nullptr;
   int co_firstlineno = 0;
+  // Whether this entry holds the references to `co_filename`, `co_name`, and
+  // `m_module`. Consume() takes the references over and then uses the pointers
+  // only as lookup keys, without dereferencing them.
+  bool owns_python_refs = true;
   std::string method_name;
   PyObject* m_module = nullptr;
 
@@ -180,7 +189,7 @@ class PythonHookContext {
   std::array<EntryShard, kNumEntryShards> entry_shards_;
   uint64_t start_timestamp_ns_;
   PythonHooksOptions options_;
-  bool stopped_ = false;
+  std::atomic<bool> stopped_ = false;
   // In end to end mode, Python get uninitialized before Stop()/Finalize(), we
   // need to buffer the result.
   std::optional<tensorflow::profiler::XPlane> end_to_end_xplane_;

From 032f1da60bb4db65cba7dcd4ecc812aa9fdcee44 Mon Sep 17 00:00:00 2001
From: Oleg Shyshkov <shyshkov@google.com>
Date: Fri, 16 Oct 2026 02:19:54 -0700
Subject: [PATCH 0948/1000] [XLA:GPU] Simply support checks in collective
 kernel info.

PiperOrigin-RevId: 1000711464
---
 .../xla/xla/backends/gpu/runtime/BUILD        |  4 -
 .../xla/backends/gpu/runtime/all_gather.cc    | 50 +++++-----
 .../xla/xla/backends/gpu/runtime/all_gather.h | 19 ----
 .../xla/backends/gpu/runtime/all_reduce.cc    | 44 ++++-----
 .../xla/xla/backends/gpu/runtime/all_reduce.h | 26 -----
 .../backends/gpu/runtime/reduce_scatter.cc    | 97 ++++++++++++-------
 .../xla/backends/gpu/runtime/reduce_scatter.h |  1 -
 7 files changed, 103 insertions(+), 138 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index 8eabf313313..ac6992e7a99 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -4182,13 +4182,11 @@ cc_library(
         "//xla:xla_data_proto_cc",
         "//xla/backends/gpu/transforms/collectives:collective_ops_utils",
         "//xla/hlo/ir:hlo",
-        "//xla/service:collective_ops_utils",
         "//xla/service:device_assignment",
         "//xla/service:gpu_topology",
         "//xla/service/gpu:gpu_constants",
         "//xla/service/gpu:launch_dimensions",
         "//xla/stream_executor:device_description",
-        "//xla/stream_executor/gpu:all_reduce_kernel",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
@@ -4245,12 +4243,10 @@ cc_library(
         "//xla/service:device_assignment",
         "//xla/service:gpu_topology",
         "//xla/service/gpu:gpu_constants",
-        "//xla/service/gpu:ir_emission_utils",
         "//xla/service/gpu:launch_dimensions",
         "//xla/stream_executor:device_description",
         "//xla/stream_executor/gpu:all_reduce_kernel",
         "@com_google_absl//absl/algorithm:container",
-        "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_gather.cc b/third_party/xla/xla/backends/gpu/runtime/all_gather.cc
index aa69b32fb5a..98262866205 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_gather.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/all_gather.cc
@@ -15,7 +15,6 @@ limitations under the License.
 
 #include "xla/backends/gpu/runtime/all_gather.h"
 
-#include <algorithm>
 #include <cstdint>
 #include <vector>
 
@@ -31,44 +30,22 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_module.h"
-#include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/primitive_util.h"
-#include "xla/service/collective_ops_utils.h"
 #include "xla/service/device_assignment.h"
 #include "xla/service/gpu/gpu_constants.h"
 #include "xla/service/gpu/launch_dimensions.h"
 #include "xla/service/gpu_topology.h"
 #include "xla/shape_util.h"
 #include "xla/stream_executor/device_description.h"
-#include "xla/stream_executor/gpu/all_reduce_kernel.h"
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
 
 namespace xla::gpu {
 
-absl::Status IsAllGatherKernelSupported(int64_t num_elements,
-                                        int64_t num_devices,
-                                        PrimitiveType element_type) {
-  // Only types in kSupportedAllGatherTypes are allowed. Complex types, tokens,
-  // tuples, and exotic types (e.g. 4-bit, 8-bit floats) are not supported.
-  if (!absl::c_linear_search(kSupportedAllGatherTypes, element_type)) {
-    return absl::UnimplementedError(absl::StrFormat(
-        "Element type %s is not supported for the all-gather kernel. "
-        "Supported types are predicates, 8/16/32/64-bit integers, and "
-        "standard floating-point types; use NCCL/RCCL for other types.",
-        primitive_util::LowercasePrimitiveTypeName(element_type)));
-  }
-
-  const int64_t output_byte_size =
-      num_elements * num_devices * primitive_util::ByteWidth(element_type);
-  if (output_byte_size > kMaxAllGatherSizeBytes) {
-    return absl::UnimplementedError(
-        "Custom all-gather strategy is only supported for small outputs.");
-  }
-
-  return absl::OkStatus();
-}
-
+// A broader check for all-gather kernel support that verifies device, operand
+// count, replica group, gather dimension, and element-type constraints.
+// Returns absl::OkStatus() if supported, or an absl::UnimplementedError
+// explaining why not.
 absl::Status IsAllGatherKernelSupported(
     const se::DeviceDescription& device_info, int32_t num_operands,
     int64_t num_devices, int64_t num_elements, int64_t per_rank_gather_dim_size,
@@ -120,7 +97,24 @@ absl::Status IsAllGatherKernelSupported(
         "dimension to be a power of 2. Got %d.",
         per_rank_gather_dim_size));
   }
-  return IsAllGatherKernelSupported(num_elements, num_devices, element_type);
+  // Only types in kSupportedAllGatherTypes are allowed. Complex types, tokens,
+  // tuples, and exotic types (e.g. 4-bit, 8-bit floats) are not supported.
+  if (!absl::c_linear_search(kSupportedAllGatherTypes, element_type)) {
+    return absl::UnimplementedError(absl::StrFormat(
+        "Element type %s is not supported for the all-gather kernel. "
+        "Supported types are predicates, 8/16/32/64-bit integers, and "
+        "standard floating-point types; use NCCL/RCCL for other types.",
+        primitive_util::LowercasePrimitiveTypeName(element_type)));
+  }
+
+  const int64_t output_byte_size =
+      num_elements * num_devices * primitive_util::ByteWidth(element_type);
+  if (output_byte_size > kMaxAllGatherSizeBytes) {
+    return absl::UnimplementedError(
+        "Custom all-gather strategy is only supported for small outputs.");
+  }
+
+  return absl::OkStatus();
 }
 
 absl::StatusOr<AllGatherInfo> BuildAllGatherInfo(
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_gather.h b/third_party/xla/xla/backends/gpu/runtime/all_gather.h
index ba542216b80..dd942e5b356 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_gather.h
+++ b/third_party/xla/xla/backends/gpu/runtime/all_gather.h
@@ -18,7 +18,6 @@ limitations under the License.
 
 #include <array>
 #include <cstdint>
-#include <vector>
 
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
@@ -28,7 +27,6 @@ limitations under the License.
 #include "xla/service/device_assignment.h"
 #include "xla/service/gpu/launch_dimensions.h"
 #include "xla/service/gpu_topology.h"
-#include "xla/stream_executor/device_description.h"
 #include "xla/xla_data.pb.h"
 
 namespace xla::gpu {
@@ -58,23 +56,6 @@ struct AllGatherInfo {
   PrimitiveType element_type;
 };
 
-// Returns absl::OkStatus() if the all-gather kernel is supported for the given
-// element type, per-rank input element count, and number of devices, or an
-// error status detailing why it is not supported.
-absl::Status IsAllGatherKernelSupported(int64_t num_elements,
-                                        int64_t num_devices,
-                                        PrimitiveType element_type);
-
-// A broader check for all-gather kernel support that verifies device, operand
-// count, replica group, gather dimension, and element-type constraints.
-// Returns absl::OkStatus() if supported, or an absl::UnimplementedError
-// explaining why not.
-absl::Status IsAllGatherKernelSupported(
-    const se::DeviceDescription& device_info, int32_t num_operands,
-    int64_t num_devices, int64_t num_elements, int64_t per_rank_gather_dim_size,
-    PrimitiveType element_type, bool is_local,
-    const std::vector<ReplicaGroup>& replica_groups);
-
 // Constructs an AllGatherInfo object for the given all-gather instruction.
 // Returns absl::UnimplementedError if the Triton all-gather kernel is not
 // supported for this instruction's configuration.
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc b/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc
index be2ec560a24..28fe762f578 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/all_reduce.cc
@@ -15,7 +15,6 @@ limitations under the License.
 
 #include "xla/backends/gpu/runtime/all_reduce.h"
 
-#include <algorithm>
 #include <cstdint>
 #include <optional>
 #include <vector>
@@ -42,13 +41,11 @@ limitations under the License.
 #include "xla/service/collective_ops_utils.h"
 #include "xla/service/device_assignment.h"
 #include "xla/service/gpu/gpu_constants.h"
-#include "xla/service/gpu/ir_emission_utils.h"
 #include "xla/service/gpu/launch_dimensions.h"
 #include "xla/service/gpu_topology.h"
 #include "xla/shape_util.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/gpu/all_reduce_kernel.h"
-#include "xla/types.h"
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
 
@@ -108,28 +105,9 @@ int64_t GetMaxSupportedAllReduceSizeBytes(AllReduceStrategy strategy) {
   }
 }
 
-absl::Status IsAllReduceKernelSupported(int64_t num_ranks, int64_t num_elements,
-                                        PrimitiveType element_type,
-                                        ReductionKind reduction_kind,
-                                        AllReduceStrategy all_reduce_strategy) {
-  if (!IsElementReductionSupported(element_type, reduction_kind)) {
-    return absl::UnimplementedError(absl::StrCat(
-        "Element type (",
-        primitive_util::LowercasePrimitiveTypeName(element_type),
-        ") and reduction kind (", absl::StrFormat("%v", reduction_kind),
-        ") combination is not supported."));
-  }
-  if (num_ranks > stream_executor::gpu::kMaxNumAllReduceInputPtrs) {
-    return absl::UnimplementedError(
-        absl::StrCat("AllReduce XLA implementation does not support more than ",
-                     se::gpu::kMaxNumAllReduceInputPtrs,
-                     " ranks. "
-                     " Required number of ranks: ",
-                     num_ranks));
-  }
-  return absl::OkStatus();
-}
-
+// A broader check overload for all-reduce kernel support.
+// Returns absl::OkStatus() if supported, or an error status detailing why
+// it is not supported.
 absl::Status IsAllReduceKernelSupported(
     const se::DeviceDescription& device_info,        //
     int32_t num_operands,                            //
@@ -204,8 +182,20 @@ absl::Status IsAllReduceKernelSupported(
     return absl::UnimplementedError(
         "Custom all-reduce strategy is only supported for small inputs.");
   }
-  return IsAllReduceKernelSupported(num_devices, num_elements, element_type,
-                                    *reduction_kind, strategy);
+  if (!IsElementReductionSupported(element_type, *reduction_kind)) {
+    return absl::UnimplementedError(absl::StrCat(
+        "Element type (",
+        primitive_util::LowercasePrimitiveTypeName(element_type),
+        ") and reduction kind (", absl::StrFormat("%v", *reduction_kind),
+        ") combination is not supported."));
+  }
+  if (num_devices > stream_executor::gpu::kMaxNumAllReduceInputPtrs) {
+    return absl::UnimplementedError(
+        absl::StrCat("AllReduce XLA implementation does not support more than ",
+                     se::gpu::kMaxNumAllReduceInputPtrs,
+                     " ranks. Required number of ranks: ", num_devices));
+  }
+  return absl::OkStatus();
 }
 
 absl::StatusOr<AllReduceInfo> BuildAllReduceInfo(
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_reduce.h b/third_party/xla/xla/backends/gpu/runtime/all_reduce.h
index db16dd4cd03..d18d9850e84 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_reduce.h
+++ b/third_party/xla/xla/backends/gpu/runtime/all_reduce.h
@@ -18,10 +18,7 @@ limitations under the License.
 
 #include <array>
 #include <cstdint>
-#include <optional>
-#include <vector>
 
-#include "absl/base/nullability.h"
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "absl/types/span.h"
@@ -32,7 +29,6 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/service/gpu/launch_dimensions.h"
-#include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/gpu/all_reduce_kernel.h"
 #include "xla/types.h"  // IWYU pragma: keep
 #include "xla/xla_data.pb.h"
@@ -89,28 +85,6 @@ absl::StatusOr<CollectiveKernelSpec> CreateAllReduceKernelSpec(
     const HloInstruction* instr, const LaunchDimensions& launch_dimensions,
     SymmetricMemoryType scratch_memory_type);
 
-// Returns absl::OkStatus() if supported, or an error status detailing why
-// it is not supported.
-absl::Status IsAllReduceKernelSupported(
-    int64_t num_ranks, int64_t num_elements, PrimitiveType element_type,
-    ReductionKind reduction_kind,
-    se::gpu::AllReduceStrategy all_reduce_strategy);
-
-// A broader check overload for all-reduce kernel support.
-// Returns absl::OkStatus() if supported, or an error status detailing why
-// it is not supported.
-absl::Status IsAllReduceKernelSupported(
-    const se::DeviceDescription& device_info,        //
-    int32_t num_operands,                            //
-    std::optional<ReductionKind> reduction_kind,     //
-    int64_t num_devices,                             //
-    int64_t num_elements,                            //
-    PrimitiveType element_type,                      //
-    bool is_local,                                   //
-    bool is_multimem_enabled,                        //
-    const std::vector<ReplicaGroup>& replica_groups  //
-);
-
 // Constructs an AllReduceInfo object for the given all-reduce instruction.
 // Returns an error status if the all-reduce kernel is not supported.
 absl::StatusOr<AllReduceInfo> BuildAllReduceInfo(
diff --git a/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.cc b/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.cc
index 7e555d76317..9d95b70e271 100644
--- a/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.cc
@@ -15,9 +15,9 @@ limitations under the License.
 
 #include "xla/backends/gpu/runtime/reduce_scatter.h"
 
-#include <algorithm>
 #include <cstdint>
 #include <optional>
+#include <vector>
 
 #include "absl/algorithm/container.h"
 #include "absl/log/log.h"
@@ -85,47 +85,47 @@ bool IsReduceScatterFlattenable(const HloReduceScatterInstruction* rs) {
   return false;
 }
 
-absl::StatusOr<ReduceScatterInfo> BuildReduceScatterInfo(
-    const GpuTopology& gpu_topology,
-    const HloReduceScatterInstruction* reduce_scatter,
-    const DeviceAssignment* device_assignment) {
-  if (!gpu_topology.has_gpu_target_config()) {
-    return absl::InvalidArgumentError(
-        "GpuTopology must have a target config to build ReduceScatterInfo.");
-  }
-  const se::DeviceDescription& device_info =
-      gpu_topology.gpu_target_config().device_description;
+// A broader check for reduce-scatter kernel support that verifies device,
+// operand count, replica group, reduction computation, element-type, layout,
+// and size constraints.
+// Returns absl::OkStatus() if supported, or an absl::UnimplementedError
+// explaining why not.
+absl::Status IsReduceScatterKernelSupported(
+    const se::DeviceDescription& device_info,        //
+    int32_t num_operands,                            //
+    std::optional<ReductionKind> reduction_kind,     //
+    int64_t num_devices,                             //
+    int64_t num_output_elements,                     //
+    PrimitiveType element_type,                      //
+    bool is_flattenable,                             //
+    bool is_local,                                   //
+    const std::vector<ReplicaGroup>& replica_groups  //
+) {
   if (!device_info.cuda_compute_capability().IsAtLeastHopper()) {
     return absl::UnimplementedError(absl::StrFormat(
         "Triton reduce-scatter requires CUDA compute capability >= 9.0. "
         "Got: %s.",
         device_info.gpu_compute_capability().ToString()));
   }
-  if (device_info.device_interconnect_info().active_links <= 0) {
-    return absl::UnimplementedError(
-        "Collective kernels are only supported on devices with NVLink "
-        "support.");
-  }
-  if (reduce_scatter->operand_count() != 1) {
+  // TODO(b/383125489): Support variadic arguments.
+  if (num_operands != 1) {
     return absl::UnimplementedError(
         "Reduce-scatter kernel only supports a single operand.");
   }
-  if (!reduce_scatter->device_list() ||
-      reduce_scatter->replica_groups().empty()) {
+  if (replica_groups.empty()) {
     return absl::UnimplementedError(
         "Replica groups must be explicitly provided for collective kernels.");
   }
-  const int64_t num_devices =
-      reduce_scatter->device_list()->num_devices_per_group();
+  if (!is_local) {
+    return absl::UnimplementedError(
+        "Cross-host symmetric memory collectives are not supported.");
+  }
   if (!llvm::has_single_bit(static_cast<uint64_t>(num_devices))) {
     return absl::UnimplementedError(absl::StrFormat(
         "Collective kernels are only supported for power of 2 number of "
         "devices. Got %d.",
         num_devices));
   }
-  const PrimitiveType element_type = reduce_scatter->shape().element_type();
-  const std::optional<ReductionKind> reduction_kind =
-      MatchReductionComputation(reduce_scatter->to_apply());
   if (!reduction_kind.has_value()) {
     return absl::UnimplementedError(
         "Unsupported reduction computation for the reduce-scatter kernel.");
@@ -141,32 +141,63 @@ absl::StatusOr<ReduceScatterInfo> BuildReduceScatterInfo(
         primitive_util::LowercasePrimitiveTypeName(element_type),
         *reduction_kind));
   }
-  if (!IsReduceScatterFlattenable(reduce_scatter)) {
+  if (!is_flattenable) {
     return absl::UnimplementedError(
         "Reduce-scatter kernel requires the scatter dimension to be the "
         "physically major-most non-trivial dimension and matching layouts.");
   }
-  if (ShapeUtil::ElementsIn(reduce_scatter->shape()) % num_devices != 0) {
+  if (num_output_elements % num_devices != 0) {
     return absl::UnimplementedError(
         "Reduce-scatter kernel requires the output element count to be "
         "divisible by the number of devices.");
   }
-  const int64_t input_bytes =
-      ShapeUtil::ByteSizeOf(reduce_scatter->operand(0)->shape());
+  const int64_t input_bytes = num_output_elements * num_devices *
+                              primitive_util::ByteWidth(element_type);
   if (input_bytes > kMaxReduceScatterSizeBytes) {
     return absl::UnimplementedError(
         "Custom reduce-scatter strategy is only supported for small inputs.");
   }
+  return absl::OkStatus();
+}
+
+absl::StatusOr<ReduceScatterInfo> BuildReduceScatterInfo(
+    const GpuTopology& gpu_topology,
+    const HloReduceScatterInstruction* reduce_scatter,
+    const DeviceAssignment* device_assignment) {
+  if (!gpu_topology.has_gpu_target_config()) {
+    return absl::InvalidArgumentError(
+        "GpuTopology must have a target config to build ReduceScatterInfo.");
+  }
+  const se::DeviceDescription& device_info =
+      gpu_topology.gpu_target_config().device_description;
+  if (!reduce_scatter->device_list()) {
+    return absl::UnimplementedError(
+        "Replica groups must be explicitly provided for collective kernels.");
+  }
+  const int64_t num_devices =
+      reduce_scatter->device_list()->num_devices_per_group();
+  const int32_t num_operands = reduce_scatter->operand_count();
+  const std::optional<ReductionKind> reduction_kind =
+      MatchReductionComputation(reduce_scatter->to_apply());
+  const PrimitiveType element_type = reduce_scatter->shape().element_type();
+  const int64_t num_output_elements =
+      ShapeUtil::ElementsIn(reduce_scatter->shape());
+  const bool is_flattenable = IsReduceScatterFlattenable(reduce_scatter);
+  if (device_info.device_interconnect_info().active_links <= 0) {
+    return absl::UnimplementedError(
+        "Collective kernels are only supported on devices with NVLink "
+        "support.");
+  }
   ABSL_ASSIGN_OR_RETURN(const bool is_local,
                    AreAllReplicasOnSameSlice(gpu_topology, *reduce_scatter,
                                              device_assignment));
-  if (!is_local) {
-    return absl::UnimplementedError(
-        "Cross-host symmetric memory collectives are not supported.");
-  }
+  ABSL_RETURN_IF_ERROR(IsReduceScatterKernelSupported(
+      device_info, num_operands, reduction_kind, num_devices,
+      num_output_elements, element_type, is_flattenable, is_local,
+      reduce_scatter->replica_groups()));
   return ReduceScatterInfo{
       /*.num_devices=*/num_devices,
-      /*.num_output_elements=*/ShapeUtil::ElementsIn(reduce_scatter->shape()),
+      /*.num_output_elements=*/num_output_elements,
   };
 }
 
diff --git a/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.h b/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.h
index 002e93121a8..40a21127470 100644
--- a/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.h
+++ b/third_party/xla/xla/backends/gpu/runtime/reduce_scatter.h
@@ -25,7 +25,6 @@ limitations under the License.
 #include "xla/service/device_assignment.h"
 #include "xla/service/gpu/launch_dimensions.h"
 #include "xla/service/gpu_topology.h"
-#include "xla/stream_executor/device_description.h"
 
 namespace xla::gpu {
 

From c6e8ddad72ac40d5e7bda03cd3e59f51ab5c4809 Mon Sep 17 00:00:00 2001
From: Alexander Lyashuk <crem@google.com>
Date: Fri, 16 Oct 2026 02:40:04 -0700
Subject: [PATCH 0949/1000] [XLA:GPU] SplitkRewriter: use dot operand sizes for
 K padding and follow shape ops at fusion edge.

- Change K padding alignment calculation to use the dot's own operand element
  sizes rather than pre-convert sizes. Neither choice is strictly better or
  worse; on its own this CL is benchmark-neutral, but in follow-up changes that
  remove dot normalization passes, benchmarks show this performs better,
  especially for 4-bit types.
- In the fusion-edge element size heuristic, follow same-type transposes,
  reshapes, and bitcasts in addition to unary elementwise ops.

PiperOrigin-RevId: 1000720251
---
 .../gpu/transforms/splitk_rewriter.cc         | 22 +++--
 .../gpu/transforms/splitk_rewriter_test.cc    | 96 +++++++++++++++++++
 2 files changed, 112 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/transforms/splitk_rewriter.cc b/third_party/xla/xla/backends/gpu/transforms/splitk_rewriter.cc
index d2379c9a3d1..c5961f67046 100644
--- a/third_party/xla/xla/backends/gpu/transforms/splitk_rewriter.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/splitk_rewriter.cc
@@ -98,7 +98,15 @@ DotDimensions GetDotDimensions(const HloInstruction* instr) {
   };
 
   auto get_side_size = [](const HloInstruction* instr) {
-    while (instr->IsElementwise() && instr->operand_count() == 1) {
+    // Follow unary elementwise ops and same-type transposes/reshapes/bitcasts
+    // to find the element size at the fusion edge.
+    while (instr->operand_count() == 1 &&
+           (instr->IsElementwise() ||
+            ((instr->opcode() == HloOpcode::kBitcast ||
+              instr->opcode() == HloOpcode::kReshape ||
+              instr->opcode() == HloOpcode::kTranspose) &&
+             ShapeUtil::SameElementType(instr->shape(),
+                                        instr->operand(0)->shape())))) {
       instr = instr->operand(0);
     }
     return ShapeUtil::ElementSizeInBits(instr->shape());
@@ -246,13 +254,15 @@ HloInstruction* PadInstruction(HloInstruction* instr, int64_t dimension_idx,
       HloInstruction::CreatePad(new_shape, instr, zero, padding_config));
 }
 
-// Returns the padded K dimension so that it is a multiple of split_k and 16B.
+// Returns the padded K dimension so that it is a multiple of split_k and 16B,
+// using the dot's own operand element sizes rather than the pre-convert sizes
+// at the fusion edge.
 int64_t GetPaddedK(HloInstruction& dot, int64_t split_k) {
-  DotDimensions dims = GetDotDimensions(&dot);
   const int64_t alignment_in_bits = 16 * 8;
-  int64_t min_element_size_in_bits = std::min(
-      {alignment_in_bits, dims.lhs_element_bits, dims.rhs_element_bits});
-  return RoundUpTo(dims.k,
+  const int64_t min_element_size_in_bits = std::min(
+      {alignment_in_bits, ShapeUtil::ElementSizeInBits(dot.operand(0)->shape()),
+       ShapeUtil::ElementSizeInBits(dot.operand(1)->shape())});
+  return RoundUpTo(GetDotDimensions(&dot).k,
                    split_k * alignment_in_bits / min_element_size_in_bits);
 }
 
diff --git a/third_party/xla/xla/backends/gpu/transforms/splitk_rewriter_test.cc b/third_party/xla/xla/backends/gpu/transforms/splitk_rewriter_test.cc
index 3e613a40f19..46026b2ab26 100644
--- a/third_party/xla/xla/backends/gpu/transforms/splitk_rewriter_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/splitk_rewriter_test.cc
@@ -215,6 +215,102 @@ CHECK: ROOT {{.*}} = f32[1024,2048]{1,0} reduce
                   .value_or(false));
 }
 
+TEST_F(SplitkRewriterTest,
+       ConvertedSubByteOperandIsNotPaddedForSubByteAlignment) {
+  // K/4 = 2128 is 16B-aligned for the dot's bf16 operand, though not for the
+  // pre-convert s4 source.
+  const char* hlo_string = R"(
+    HloModule module
+
+    ENTRY test {
+      lhs = bf16[64,8512]{1,0} parameter(0)
+      rhs = s4[8512,8512]{1,0:E(4)} parameter(1)
+      rhs_converted = bf16[8512,8512]{1,0} convert(rhs)
+      ROOT dot = bf16[64,8512]{1,0} dot(lhs, rhs_converted),
+                             lhs_contracting_dims={1}, rhs_contracting_dims={1}
+    })";
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_string));
+  module->mutable_config()
+      .mutable_debug_options()
+      .set_xla_gpu_experimental_force_split_k(4);
+
+  ASSERT_OK_AND_ASSIGN(bool changed,
+                       rewriter_.HloModulePass::Run(module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_TRUE(RunFileCheck(module->ToString(), R"(
+CHECK-NOT: pad(
+CHECK: %[[CONVERT:.*]] = bf16[8512,8512]{1,0} convert(%rhs)
+CHECK-NOT: pad(
+CHECK: bf16[8512,4,2128]{2,1,0} reshape(%[[CONVERT]])
+CHECK-NOT: pad(
+  )")
+                  .value_or(false));
+}
+
+TEST_F(SplitkRewriterTest, ConvertedSubByteOperandIsPaddedForOperandAlignment) {
+  // K/4 is not 16B-aligned for bf16, so K is padded for the dot's bf16 operand
+  // alignment (to a multiple of 4 * 8 elements), not the pre-convert s4 size.
+  const char* hlo_string = R"(
+    HloModule module
+
+    ENTRY test {
+      lhs = bf16[64,8500]{1,0} parameter(0)
+      rhs = s4[8512,8500]{1,0:E(4)} parameter(1)
+      rhs_converted = bf16[8512,8500]{1,0} convert(rhs)
+      ROOT dot = bf16[64,8512]{1,0} dot(lhs, rhs_converted),
+                             lhs_contracting_dims={1}, rhs_contracting_dims={1}
+    })";
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_string));
+  module->mutable_config()
+      .mutable_debug_options()
+      .set_xla_gpu_experimental_force_split_k(4);
+
+  ASSERT_OK_AND_ASSIGN(bool changed,
+                       rewriter_.HloModulePass::Run(module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_TRUE(RunFileCheck(module->ToString(), R"(
+CHECK-DAG: bf16[64,8512]{1,0} pad(%lhs, {{.*}}), padding=0_0x0_12
+CHECK-DAG: bf16[8512,8512]{1,0} pad(%rhs_converted, {{.*}}), padding=0_0x0_12
+  )")
+                  .value_or(false));
+}
+
+TEST_F(SplitkRewriterTest, LayoutChangeOfConvertedSubByteOperandIsNotPadded) {
+  // Same as ConvertedSubByteOperandIsNotPaddedForSubByteAlignment, but with
+  // the operand relabeled by a layout-changing bitcast.
+  const char* hlo_string = R"(
+    HloModule module
+
+    ENTRY test {
+      lhs = bf16[64,8512]{1,0} parameter(0)
+      rhs = s4[8512,8512]{1,0:E(4)} parameter(1)
+      rhs_converted = bf16[8512,8512]{1,0} convert(rhs)
+      rhs_bitcast = bf16[8512,8512]{0,1} bitcast(rhs_converted)
+      ROOT dot = bf16[64,8512]{1,0} dot(lhs, rhs_bitcast),
+                             lhs_contracting_dims={1}, rhs_contracting_dims={0}
+    })";
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(hlo_string));
+  module->mutable_config()
+      .mutable_debug_options()
+      .set_xla_gpu_experimental_force_split_k(4);
+
+  ASSERT_OK_AND_ASSIGN(bool changed,
+                       rewriter_.HloModulePass::Run(module.get()));
+  EXPECT_TRUE(changed);
+  EXPECT_TRUE(RunFileCheck(module->ToString(), R"(
+CHECK-NOT: pad(
+CHECK: bf16[4,2128,8512]{1,0,2} reshape(%rhs_bitcast)
+CHECK-NOT: pad(
+  )")
+                  .value_or(false));
+}
+
 }  // namespace
 }  // namespace gpu
 }  // namespace xla

From 27f616ec153946200e03e493b0d064679c179757 Mon Sep 17 00:00:00 2001
From: Levon Ter-Grigoryan <patrios@google.com>
Date: Fri, 16 Oct 2026 03:39:41 -0700
Subject: [PATCH 0950/1000] [XLA:GPU] Add Triton collective emitter and
 lowering for AllToAll.

Implements one and two-shots push-based all-to-all algorithm.

PiperOrigin-RevId: 1000745709
---
 .../gpu/codegen/triton/collective_emitter.cc  | 295 +++++++++++++++++-
 .../gpu/codegen/triton/collective_emitter.h   |   4 +
 .../xla/backends/gpu/codegen/triton/fusion.cc |   6 +-
 .../backends/gpu/codegen/triton/support.cc    |  16 +
 .../transforms/stablehlo_lower_to_triton.cc   |  16 +-
 .../tests/stable_hlo_to_triton_lowering.mlir  |  54 ++++
 .../xla/service/gpu/hlo_fusion_analysis.cc    |  12 +-
 7 files changed, 384 insertions(+), 19 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
index 46164b66cc9..2894232d3c7 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
@@ -48,6 +48,7 @@ limitations under the License.
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypeInterfaces.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Location.h"
@@ -525,15 +526,18 @@ mlir::Value EmitRemoteBufferPtr(mlir::ImplicitLocOpBuilder& b,
 }
 
 // Emits a barrier between this block and the matching blocks of the other
-// ranks. All threads in the block should have completed their writes before we
-// proceed with a block barrier.
+// ranks.
 // Otherwise, remote ranks might start reading the data before it is ready.
 void EmitBlockBarrier(
     mlir::ImplicitLocOpBuilder& b, mlir::Value signal_buffers, mlir::Value rank,
     mlir::Value signal_value, int64_t world_size,
     mlir::Value signal_slot = nullptr, int64_t signal_stride = 0,
-    mtx::BarrierMode barrier_mode = mtx::BarrierMode::kSymmetric) {
-  mlir::triton::gpu::BarrierOp::create(b, mlir::triton::gpu::AddrSpace::Local);
+    mtx::BarrierMode barrier_mode = mtx::BarrierMode::kSymmetric,
+    bool emit_local_barrier = true) {
+  if (emit_local_barrier) {
+    mlir::triton::gpu::BarrierOp::create(b,
+                                         mlir::triton::gpu::AddrSpace::Local);
+  }
   mtx::BlockBarrierOp::create(
       b, signal_buffers, rank, signal_value, signal_slot,
       b.getI32IntegerAttr(static_cast<int32_t>(world_size)),
@@ -1423,6 +1427,7 @@ absl::StatusOr<std::vector<Shape>> GetCollectiveUnmanagedKernelArguments(
   switch (root->opcode()) {
     case HloOpcode::kAllReduce:
     case HloOpcode::kAllGather:
+    case HloOpcode::kAllToAll:
     case HloOpcode::kReduceScatter:
       return GetRemoteBufferUnmanagedKernelArguments(
           computation, Cast<HloCollectiveInstruction>(root));
@@ -1703,10 +1708,11 @@ absl::StatusOr<AllGatherRewriteContext> BuildAllGatherContext(
 // Returns the rank whose output slice contains this program's tile. Programs
 // enumerate output tiles in row-major order (same assumption as
 // `signal_stride` and the asymmetric barrier), so each run of `signal_stride`
-// consecutive programs covers one rank's slice along `gather_dim`.
+// consecutive programs covers one rank's slice along `gather_dim` /
+// `split_dim`.
+template <typename ContextT>
 mlir::Value EmitPeerReplicaId(mlir::ImplicitLocOpBuilder& b,
-                              const AllGatherRewriteContext& ctx,
-                              mlir::Value block_id) {
+                              const ContextT& ctx, mlir::Value block_id) {
   mlir::Value stride_op = arith::ConstantOp::create(
       b, b.getI32IntegerAttr(static_cast<int32_t>(ctx.signal_stride)));
   mlir::Value world_size_op = arith::ConstantOp::create(
@@ -1717,9 +1723,9 @@ mlir::Value EmitPeerReplicaId(mlir::ImplicitLocOpBuilder& b,
 
 // Returns the active slot of `rank`'s remote buffer as a memref with the type
 // of the local input buffer.
+template <typename ContextT>
 mlir::Value EmitRemoteBufferMemref(mlir::ImplicitLocOpBuilder& b,
-                                   const AllGatherRewriteContext& ctx,
-                                   mlir::Value rank,
+                                   const ContextT& ctx, mlir::Value rank,
                                    mlir::Value buffer_offset) {
   const mlir::MemRefType memref_type = ctx.input_buffer.getType();
   mlir::Value ptr =
@@ -1762,6 +1768,206 @@ void CopySubtileToScratch(mlir::ImplicitLocOpBuilder& b,
                               subtile_shape, ctx.input_extract.getStrides());
 }
 
+constexpr int64_t kMaxOneShotAllToAllSizeBytes = 32LL * 1024 * 1024;  // 32 MB
+
+struct AllToAllRewriteContext {
+  mlir::stablehlo::AllToAllOp op;
+  xtile::ExtractTileOp input_extract;
+  // Local input buffer read by `input_extract`.
+  mlir::TypedValue<mlir::MemRefType> input_buffer;
+  mlir::Value rank_arg;
+  mlir::Value signal_buffers_arg;
+  // !tt.ptr<i64> table of the per-rank remote scratch buffers.
+  mlir::Value remote_buffers_arg;
+  int64_t world_size = 0;
+  int64_t split_dim = 0;
+  int64_t shard_size = 0;
+  int64_t signal_stride = 0;
+  bool is_one_shot = true;
+  PrimitiveType element_type = PrimitiveType::PRIMITIVE_TYPE_INVALID;
+  llvm::ArrayRef<int64_t> input_shape;
+  llvm::ArrayRef<int64_t> tile_sizes;
+};
+
+absl::StatusOr<AllToAllRewriteContext> BuildAllToAllContext(
+    mlir::stablehlo::AllToAllOp op) {
+  if (op.getOperands().size() != 1) {
+    return absl::InvalidArgumentError(
+        "AllToAll op must have exactly one operand.");
+  }
+  if (op.getSplitDimension() != op.getConcatDimension()) {
+    return absl::InvalidArgumentError(
+        "AllToAll op split_dimension must equal concat_dimension.");
+  }
+  auto entry_func = op->getParentOfType<xtile::EntryFuncOp>();
+  if (!entry_func) {
+    return absl::InvalidArgumentError(
+        "AllToAll op must be in an XTile entry function.");
+  }
+  if (op->getParentOp() != entry_func.getOperation()) {
+    return absl::InvalidArgumentError(
+        "AllToAll op must be directly nested in the XTile entry function.");
+  }
+  auto num_opaque_attr =
+      entry_func->getAttrOfType<mlir::IntegerAttr>("num_opaque_args");
+  if (!num_opaque_attr ||
+      num_opaque_attr.getInt() <= kNumCollectiveMetadataArgs) {
+    return absl::InvalidArgumentError(
+        "EntryFuncOp must have the collective metadata arguments followed by "
+        "the remote buffer arguments.");
+  }
+
+  // stablehlo.all_to_all requires the split dimension of the operand to be
+  // divisible by split_count. When this is not the case, the XTile fusion
+  // emitter wraps the operand and result in unrealized_conversion_cast with a
+  // padded type.
+  // TODO(patrios): Use xtile.all_to_all to avoid this cast workaround.
+  mlir::Value input_tile = op.getOperand(0);
+  if (auto cast_in =
+          input_tile.getDefiningOp<mlir::UnrealizedConversionCastOp>()) {
+    if (!input_tile.hasOneUse() || cast_in.getInputs().size() != 1) {
+      return absl::InvalidArgumentError(
+          "Expected single-input UnrealizedConversionCastOp wrapper around "
+          "AllToAll operand.");
+    }
+    input_tile = cast_in.getInputs()[0];
+  }
+
+  auto input_extract = llvm::dyn_cast_if_present<xtile::ExtractTileOp>(
+      input_tile.getDefiningOp());
+  if (!input_extract && input_tile.getDefiningOp() &&
+      input_tile.getDefiningOp()->getNumOperands() > 0) {
+    // Workaround(i1_to_i8_workaround): booleans are stored as i8 and cast to
+    // i1 after ExtractTileOp.
+    input_extract = llvm::dyn_cast_if_present<xtile::ExtractTileOp>(
+        input_tile.getDefiningOp()->getOperand(0).getDefiningOp());
+  }
+  if (!input_extract) {
+    return absl::InvalidArgumentError(
+        "AllToAll operand must be defined by xtile::ExtractTileOp.");
+  }
+
+  auto input_arg =
+      mlir::dyn_cast<mlir::BlockArgument>(input_extract.getSource());
+  const int64_t num_parameters =
+      num_opaque_attr.getInt() - kNumCollectiveMetadataArgs;
+  if (!input_arg ||
+      input_arg.getOwner()->getParentOp() != entry_func.getOperation() ||
+      input_arg.getArgNumber() >= num_parameters) {
+    return absl::InvalidArgumentError(
+        "AllToAll operand must be extracted from a fusion parameter.");
+  }
+  if (!input_extract->hasOneUse() || !input_tile.hasOneUse()) {
+    return absl::InvalidArgumentError(
+        "AllToAll input tile must only be used by the AllToAll op.");
+  }
+
+  ABSL_ASSIGN_OR_RETURN(auto replica_groups,
+                   xla::ConvertReplicaGroups(op.getReplicaGroups(), op));
+  if (replica_groups->replica_groups().empty()) {
+    return absl::InvalidArgumentError(
+        "AllToAll replica groups must not be empty.");
+  }
+
+  PrimitiveType element_type = xla::ConvertMlirTypeToPrimitiveType(
+      input_extract.getType().getElementType());
+  if (element_type == PrimitiveType::PRIMITIVE_TYPE_INVALID) {
+    return absl::InvalidArgumentError(
+        "Could not convert AllToAll element type to PrimitiveType.");
+  }
+
+  AllToAllRewriteContext ctx;
+  ctx.op = op;
+  ctx.input_extract = input_extract;
+  ctx.input_buffer = input_extract.getSource();
+  ctx.world_size = replica_groups->num_devices_per_group();
+  ctx.split_dim = op.getSplitDimension();
+  ctx.element_type = element_type;
+  ctx.input_shape = ctx.input_buffer.getType().getShape();
+  ctx.tile_sizes = input_extract.getFullTileShape();
+  if (ctx.world_size <= 0 ||
+      ctx.input_shape[ctx.split_dim] % ctx.world_size != 0) {
+    return absl::InvalidArgumentError(
+        "AllToAll split dimension must be divisible by world size.");
+  }
+  // The threshold is per total transported data, therefore it is divided by the
+  // number of participating devices.
+  ctx.is_one_shot = Product(ctx.input_shape) *
+                        ShapeUtil::ByteSizeOfPrimitiveType(ctx.element_type) <=
+                    kMaxOneShotAllToAllSizeBytes / ctx.world_size;
+  ctx.shard_size = ctx.input_shape[ctx.split_dim] / ctx.world_size;
+  const int64_t tile_size = ctx.tile_sizes[ctx.split_dim];
+  if (ctx.shard_size % tile_size != 0) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        "AllToAll tile size (%d) must divide the per-rank shard size (%d) "
+        "along the split dimension.",
+        tile_size, ctx.shard_size));
+  }
+
+  int32_t total_args = entry_func.getNumArguments();
+  int32_t metadata_args_start =
+      total_args - num_opaque_attr.getInt() - kNumTileIndexArgs;
+  ctx.rank_arg = entry_func.getArgument(metadata_args_start);
+  ctx.signal_buffers_arg = entry_func.getArgument(metadata_args_start + 2);
+  ctx.remote_buffers_arg =
+      entry_func.getArgument(metadata_args_start + kNumCollectiveMetadataArgs +
+                             input_arg.getArgNumber());
+
+  ctx.signal_stride = ctx.shard_size / tile_size;
+  for (int64_t j = ctx.split_dim + 1;
+       j < static_cast<int64_t>(ctx.input_shape.size()); ++j) {
+    ctx.signal_stride *=
+        xla::CeilOfRatio(ctx.input_shape[j], ctx.tile_sizes[j]);
+  }
+  return ctx;
+}
+
+// Emits:
+//   %shard_size = arith.constant <ctx.shard_size> : index
+//   %replica_idx = arith.index_cast %replica_id : i32 to index
+//   %shard_off = arith.muli %replica_idx, %shard_size : index
+//   %split_offset = arith.addi %intra_shard_offset, %shard_off : index
+mlir::Value EmitShardSplitOffset(mlir::ImplicitLocOpBuilder& b,
+                                 AllToAllRewriteContext& ctx,
+                                 mlir::Value replica_id) {
+  mlir::Value shard_size_val =
+      arith::ConstantIndexOp::create(b, ctx.shard_size);
+  mlir::Value replica_idx =
+      arith::IndexCastOp::create(b, b.getIndexType(), replica_id);
+  mlir::Value shard_off = arith::MulIOp::create(b, replica_idx, shard_size_val);
+  return arith::AddIOp::create(b, ctx.input_extract.getOffsets()[ctx.split_dim],
+                               shard_off);
+}
+
+// Pushes this program's tile from the target peer's shard of the local input
+// buffer into this rank's shard of the peer's remote scratch buffer.
+// Needed for push-based algorithms like AllToAll.
+void CopyTileToScratch(mlir::ImplicitLocOpBuilder& b,
+                       AllToAllRewriteContext& ctx, mlir::Value peer_replica_id,
+                       mlir::Value peer_split_offset,
+                       mlir::Value buffer_offset) {
+  llvm::SmallVector<mlir::Value> read_offsets(
+      ctx.input_extract.getOffsets().begin(),
+      ctx.input_extract.getOffsets().end());
+  read_offsets[ctx.split_dim] = peer_split_offset;
+
+  auto tile_tensor_type = mlir::RankedTensorType::get(
+      ctx.tile_sizes, ctx.input_extract.getType().getElementType());
+  auto tile = xtile::ExtractTileOp::create(
+      b, tile_tensor_type, ctx.input_buffer, read_offsets, ctx.tile_sizes,
+      ctx.input_extract.getStrides());
+
+  llvm::SmallVector<mlir::Value> write_offsets(
+      ctx.input_extract.getOffsets().begin(),
+      ctx.input_extract.getOffsets().end());
+  write_offsets[ctx.split_dim] = EmitShardSplitOffset(b, ctx, ctx.rank_arg);
+
+  mlir::Value peer_scratch_slot =
+      EmitRemoteBufferMemref(b, ctx, peer_replica_id, buffer_offset);
+  xtile::InsertTileOp::create(b, tile, peer_scratch_slot, write_offsets,
+                              ctx.tile_sizes, ctx.input_extract.getStrides());
+}
+
 }  // namespace
 
 mlir::LogicalResult RewriteAllGather(mlir::stablehlo::AllGatherOp op,
@@ -1800,6 +2006,77 @@ mlir::LogicalResult RewriteAllGather(mlir::stablehlo::AllGatherOp op,
   return mlir::success();
 }
 
+mlir::LogicalResult RewriteAllToAll(mlir::stablehlo::AllToAllOp op,
+                                    mlir::PatternRewriter& rewriter) {
+  absl::StatusOr<AllToAllRewriteContext> maybe_ctx = BuildAllToAllContext(op);
+  if (!maybe_ctx.ok()) {
+    return rewriter.notifyMatchFailure(op, maybe_ctx.status().message());
+  }
+  AllToAllRewriteContext& ctx = *maybe_ctx;
+
+  mlir::ImplicitLocOpBuilder builder(op.getLoc(), rewriter);
+  builder.setInsertionPoint(ctx.input_extract);
+
+  mlir::Value block_id = ttir::GetProgramIdOp::create(builder, 0);
+  mlir::Value peer_replica_id = EmitPeerReplicaId(builder, ctx, block_id);
+  mlir::Value peer_split_offset =
+      EmitShardSplitOffset(builder, ctx, peer_replica_id);
+  mlir::Value signal_value = EmitDeviceInvocationCount(
+      builder, ctx.signal_buffers_arg, ctx.rank_arg, block_id, ctx.world_size);
+
+  mlir::Value buffer_offset;
+  if (ctx.is_one_shot) {
+    buffer_offset = EmitDoubleBufferOffset(
+        builder, signal_value, Product(ctx.input_shape), ctx.element_type);
+    // Push this program's tile of the local input tensor directly into
+    // `peer_replica_id`'s remote scratch buffer at shard `rank_arg`.
+    CopyTileToScratch(builder, ctx, peer_replica_id, peer_split_offset,
+                      buffer_offset);
+    // Wait until `peer_replica_id` has pushed its corresponding tile into this
+    // rank's scratch buffer.
+    EmitBlockBarrier(builder, ctx.signal_buffers_arg, ctx.rank_arg,
+                     signal_value, ctx.world_size,
+                     /*signal_slot=*/peer_replica_id, ctx.signal_stride,
+                     mtx::BarrierMode::kPairwiseSymmetric);
+  } else {
+    // TODO(patrios): Instead of using scratch buffer, mark AllToAll output
+    // buffer as a collective operation during compilation and write directly to
+    // it. Without this right now the only benefit of two shot is lower memory
+    // usage.
+    buffer_offset = arith::ConstantOp::create(builder, builder.getI64Type(),
+                                              builder.getI64IntegerAttr(0));
+    // Synchronize with `peer_replica_id` before writing into its
+    // single-buffered remote scratch buffer.
+    EmitBlockBarrier(builder, ctx.signal_buffers_arg, ctx.rank_arg,
+                     signal_value, ctx.world_size,
+                     /*signal_slot=*/peer_replica_id, ctx.signal_stride,
+                     mtx::BarrierMode::kPairwiseSymmetric,
+                     /*emit_local_barrier=*/false);
+    CopyTileToScratch(builder, ctx, peer_replica_id, peer_split_offset,
+                      buffer_offset);
+    mlir::Value exit_signal_value = arith::AddIOp::create(
+        builder, signal_value,
+        arith::ConstantOp::create(builder, builder.getI32IntegerAttr(1)));
+    EmitBlockBarrier(builder, ctx.signal_buffers_arg, ctx.rank_arg,
+                     exit_signal_value, ctx.world_size,
+                     /*signal_slot=*/peer_replica_id, ctx.signal_stride,
+                     mtx::BarrierMode::kPairwiseSymmetric);
+  }
+
+  // Read the tile pushed by `peer_replica_id` from this rank's own local
+  // scratch buffer: offset along `split_dim` is `intra_shard_offset +
+  // peer_replica_id * shard_size`, matching `output_insert`.
+  mlir::Value own_scratch_slot =
+      EmitRemoteBufferMemref(builder, ctx, ctx.rank_arg, buffer_offset);
+  rewriter.modifyOpInPlace(ctx.input_extract, [&]() {
+    ctx.input_extract.getSourceMutable().assign(own_scratch_slot);
+    ctx.input_extract.getOffsetsMutable()[ctx.split_dim].assign(
+        peer_split_offset);
+  });
+  rewriter.replaceOp(op, op.getOperand(0));
+  return mlir::success();
+}
+
 absl::Status FlattenReduceScatterFusion(
     HloFusionInstruction* absl_nonnull fusion_instr) {
   HloComputation* parent = fusion_instr->parent();
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h
index 86c93904ec8..7f86608f0ae 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h
@@ -96,6 +96,10 @@ mlir::LogicalResult RewriteAllReduce(mlir::stablehlo::AllReduceOp op,
 mlir::LogicalResult RewriteAllGather(mlir::stablehlo::AllGatherOp op,
                                      mlir::PatternRewriter& rewriter);
 
+// Rewrites stablehlo all-to-all op to a triton implementation.
+mlir::LogicalResult RewriteAllToAll(mlir::stablehlo::AllToAllOp op,
+                                    mlir::PatternRewriter& rewriter);
+
 // Rewrites xtile reduce-scatter op to a triton implementation.
 mlir::LogicalResult RewriteReduceScatter(xla::xtile::ReduceScatterOp op,
                                          mlir::PatternRewriter& rewriter);
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/fusion.cc b/third_party/xla/xla/backends/gpu/codegen/triton/fusion.cc
index 882730ab854..c8cb80a5ae1 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/fusion.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/fusion.cc
@@ -54,6 +54,7 @@ limitations under the License.
 #include "xla/future.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/utils/hlo_traversal.h"
 #include "xla/service/gpu/backend_configs.pb.h"
 #include "xla/service/gpu/gpu_constants.h"
@@ -90,7 +91,10 @@ absl::StatusOr<EmitArgs> EmitCollectiveFusion(
 
   const HloInstructionAdaptor collective_hero = analysis.fusion_hero(0);
 
-  bool use_global_device_ids = collective_hero.use_global_device_ids();
+  std::optional<bool> use_global_device_ids;
+  if (collective_hero.opcode() != HloOpcode::kAllToAll) {
+    use_global_device_ids = collective_hero.use_global_device_ids();
+  }
 
   CollectiveConfig config = GetCollectiveConfig(&collective_hero.instruction(),
                                                 use_global_device_ids);
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/support.cc b/third_party/xla/xla/backends/gpu/codegen/triton/support.cc
index cf5a4976583..2243ca08e5b 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/support.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/support.cc
@@ -497,6 +497,20 @@ CodegenDecision IsTritonSupportedReduceScatter(
   return CodegenDecision::Allow();
 }
 
+CodegenDecision IsTritonSupportedAllToAll(
+    const HloAllToAllInstruction& all_to_all) {
+  if (!all_to_all.shape().IsArray() ||
+      !all_to_all.split_dimension().has_value()) {
+    return CodegenDecision::Forbid(
+        "Only array all-to-alls with split_dimension are supported.");
+  }
+  if (all_to_all.replica_groups().empty()) {
+    return CodegenDecision::Forbid("All-to-all does not have replica groups.");
+  }
+  return CodegenDecision(all_to_all.shape().element_type() != S4,
+                         "S4 is not supported.");
+}
+
 absl::Status CheckSupportedCheckDotDimensions(const HloDotInstruction& dot) {
   const DotDimensionNumbers& dim_numbers = dot.dot_dimension_numbers();
   // Only checking one side of bach and contracting dimensions, since they must
@@ -978,6 +992,8 @@ CodegenDecision IsTritonSupportedInstructionImpl(
     case HloOpcode::kAllGather:
       return CodegenDecision(instr.shape().element_type() != S4,
                              "S4 is not supported.");
+    case HloOpcode::kAllToAll:
+      return IsTritonSupportedAllToAll(*Cast<HloAllToAllInstruction>(&instr));
     case HloOpcode::kReduceScatter:
       return IsTritonSupportedReduceScatter(
           *Cast<HloReduceScatterInstruction>(&instr), gpu_version);
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/stablehlo_lower_to_triton.cc b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/stablehlo_lower_to_triton.cc
index 6531fb2aa49..2c293e1885a 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/stablehlo_lower_to_triton.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/stablehlo_lower_to_triton.cc
@@ -1043,6 +1043,18 @@ class LowerAllGather : public mlir::OpRewritePattern<stablehlo::AllGatherOp> {
   }
 };
 
+class LowerAllToAll : public mlir::OpRewritePattern<stablehlo::AllToAllOp> {
+ public:
+  using OpRewritePattern::OpRewritePattern;
+
+ private:
+  mlir::LogicalResult matchAndRewrite(
+      stablehlo::AllToAllOp op,
+      mlir::PatternRewriter& rewriter) const override {
+    return ::xla::gpu::RewriteAllToAll(op, rewriter);
+  }
+};
+
 class LowerAllReduce : public mlir::OpRewritePattern<stablehlo::AllReduceOp> {
  public:
   using OpRewritePattern::OpRewritePattern;
@@ -1080,8 +1092,8 @@ class StableHLOLowerToTritonPass
     {
       mlir::RewritePatternSet patterns(mlir_context);
       patterns.add<LowerTranspose, LowerIotaToMakeRange, LowerBroadcastInDim,
-                   LowerReduce, LowerReshape, LowerAllGather, LowerAllReduce,
-                   LowerReduceScatter>(mlir_context);
+                   LowerReduce, LowerReshape, LowerAllGather, LowerAllToAll,
+                   LowerAllReduce, LowerReduceScatter>(mlir_context);
       patterns.add<CanonicalizeDotGeneral>(mlir_context);
       patterns.add<LowerDotGeneral>(mlir_context, warp_specialization_allowed_);
       if (mlir::failed(mlir::applyPatternsGreedily(getOperation(),
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/stable_hlo_to_triton_lowering.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/stable_hlo_to_triton_lowering.mlir
index d14ae9fa827..7c4bff43938 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/stable_hlo_to_triton_lowering.mlir
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/stable_hlo_to_triton_lowering.mlir
@@ -492,6 +492,60 @@ xtile.entry_func @reduce_scatter_invalid_tile_size_doesnt_lower(%input: memref<4
   xtile.return
 }
 
+// CHECK-LABEL: xtile.entry_func @all_to_all_one_shot(
+xtile.entry_func @all_to_all_one_shot(%input: memref<128x128xf32>, %output: memref<128x128xf32>, %device_rank: i32, %signal_value: i32, %signal_buffer: !tt.ptr<i64>, %remote_input_buffer: !tt.ptr<i64>, %tile_id: index) attributes {num_opaque_args = 4 : i32} {
+  %c0 = arith.constant 0 : index
+  %tile = xtile.extract %input[%c0, %c0][16, 16][1, 1] : memref<128x128xf32> -> tensor<16x16xf32>
+  // CHECK: triton_xla.ptr_to_memref
+  // CHECK: triton_xla.block_barrier {{.*}} <world_size = 2, signal_stride = 32, barrier_mode = pairwise_symmetric>
+  // CHECK: triton_xla.ptr_to_memref
+  // CHECK-NOT: stablehlo.all_to_all
+  %all_to_all = "stablehlo.all_to_all"(%tile) <{split_dimension = 0 : i64, concat_dimension = 0 : i64, split_count = 2 : i64, replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>}> : (tensor<16x16xf32>) -> tensor<16x16xf32>
+  xtile.insert %all_to_all into %output[%c0, %c0][16, 16][1, 1] : tensor<16x16xf32> -> memref<128x128xf32>
+  xtile.return
+}
+
+// CHECK-LABEL: xtile.entry_func @all_to_all_two_shot(
+xtile.entry_func @all_to_all_two_shot(%input: memref<8192x4096xf32>, %output: memref<8192x4096xf32>, %device_rank: i32, %signal_value: i32, %signal_buffer: !tt.ptr<i64>, %remote_input_buffer: !tt.ptr<i64>, %tile_id: index) attributes {num_opaque_args = 4 : i32} {
+  %c0 = arith.constant 0 : index
+  %tile = xtile.extract %input[%c0, %c0][64, 64][1, 1] : memref<8192x4096xf32> -> tensor<64x64xf32>
+  // CHECK-NOT: ttg.barrier
+  // CHECK: triton_xla.block_barrier {{.*}} <world_size = 2, signal_stride = 4096, barrier_mode = pairwise_symmetric>
+  // CHECK: triton_xla.ptr_to_memref
+  // CHECK: ttg.barrier local
+  // CHECK: triton_xla.block_barrier {{.*}} <world_size = 2, signal_stride = 4096, barrier_mode = pairwise_symmetric>
+  // CHECK-NOT: stablehlo.all_to_all
+  %all_to_all = "stablehlo.all_to_all"(%tile) <{split_dimension = 0 : i64, concat_dimension = 0 : i64, split_count = 2 : i64, replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>}> : (tensor<64x64xf32>) -> tensor<64x64xf32>
+  xtile.insert %all_to_all into %output[%c0, %c0][64, 64][1, 1] : tensor<64x64xf32> -> memref<8192x4096xf32>
+  xtile.return
+}
+
+// CHECK-LABEL: xtile.entry_func @all_to_all_with_unrealized_conversion_cast(
+xtile.entry_func @all_to_all_with_unrealized_conversion_cast(%input: memref<4x128xf32>, %output: memref<4x128xf32>, %device_rank: i32, %signal_value: i32, %signal_buffer: !tt.ptr<i64>, %remote_input_buffer: !tt.ptr<i64>, %tile_id: index) attributes {num_opaque_args = 4 : i32} {
+  %c0 = arith.constant 0 : index
+  %tile = xtile.extract %input[%c0, %c0][1, 16][1, 1] : memref<4x128xf32> -> tensor<1x16xf32>
+  %cast_in = builtin.unrealized_conversion_cast %tile : tensor<1x16xf32> to tensor<2x16xf32>
+  // CHECK: triton_xla.ptr_to_memref
+  // CHECK: triton_xla.block_barrier {{.*}} <world_size = 2, signal_stride = 16, barrier_mode = pairwise_symmetric>
+  // CHECK: triton_xla.ptr_to_memref
+  // CHECK-NOT: builtin.unrealized_conversion_cast
+  // CHECK-NOT: stablehlo.all_to_all
+  %all_to_all = "stablehlo.all_to_all"(%cast_in) <{split_dimension = 0 : i64, concat_dimension = 0 : i64, split_count = 2 : i64, replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>}> : (tensor<2x16xf32>) -> tensor<2x16xf32>
+  %cast_out = builtin.unrealized_conversion_cast %all_to_all : tensor<2x16xf32> to tensor<1x16xf32>
+  xtile.insert %cast_out into %output[%c0, %c0][1, 16][1, 1] : tensor<1x16xf32> -> memref<4x128xf32>
+  xtile.return
+}
+
+// CHECK-LABEL: xtile.entry_func @all_to_all_tile_not_dividing_shard_size_doesnt_lower(
+xtile.entry_func @all_to_all_tile_not_dividing_shard_size_doesnt_lower(%input: memref<96x128xf32>, %output: memref<96x128xf32>, %device_rank: i32, %signal_value: i32, %signal_buffer: !tt.ptr<i64>, %remote_input_buffer: !tt.ptr<i64>, %tile_id: index) attributes {num_opaque_args = 4 : i32} {
+  %c0 = arith.constant 0 : index
+  %tile = xtile.extract %input[%c0, %c0][32, 16][1, 1] : memref<96x128xf32> -> tensor<32x16xf32>
+  // CHECK: stablehlo.all_to_all
+  %all_to_all = "stablehlo.all_to_all"(%tile) <{split_dimension = 0 : i64, concat_dimension = 0 : i64, split_count = 2 : i64, replica_groups = dense<[[0, 1]]> : tensor<1x2xi64>}> : (tensor<32x16xf32>) -> tensor<32x16xf32>
+  xtile.insert %all_to_all into %output[%c0, %c0][32, 16][1, 1] : tensor<32x16xf32> -> memref<96x128xf32>
+  xtile.return
+}
+
 // CHECK: func @lower_dot_with_warp_specialization_to_triton
 func.func @lower_dot_with_warp_specialization_to_triton(
     %arg0: tensor<2x4xf32>,
diff --git a/third_party/xla/xla/service/gpu/hlo_fusion_analysis.cc b/third_party/xla/xla/service/gpu/hlo_fusion_analysis.cc
index 0003e06c35f..81ae2838b80 100644
--- a/third_party/xla/xla/service/gpu/hlo_fusion_analysis.cc
+++ b/third_party/xla/xla/service/gpu/hlo_fusion_analysis.cc
@@ -217,16 +217,14 @@ int SmallestBitWidth(const Container& args) {
 
 std::optional<HloInstructionAdaptor> FindCollectiveHero(
     const HloFusionAdaptor& fusion) {
-  std::optional<HloInstructionAdaptor> collective_hero;
   for (HloInstructionAdaptor instr : fusion.MakeInstructionPostOrder()) {
-    if (instr.opcode() == HloOpcode::kAllGather ||
-        instr.opcode() == HloOpcode::kAllReduce ||
-        instr.opcode() == HloOpcode::kReduceScatter) {
-      collective_hero = instr;
-      break;
+    if (HloPredicateIsOp<HloOpcode::kAllGather, HloOpcode::kAllReduce,
+                         HloOpcode::kAllToAll, HloOpcode::kReduceScatter>(
+            &instr.instruction())) {
+      return instr;
     }
   }
-  return collective_hero;
+  return std::nullopt;
 }
 
 }  // namespace

From 99f81764dcfc4c89b4b8b0674cfe0e236933b1fc Mon Sep 17 00:00:00 2001
From: Akhil Goel <akhil.goel@intel.com>
Date: Fri, 16 Oct 2026 04:07:08 -0700
Subject: [PATCH 0951/1000] PR #50870: [XLA:GPU][oneAPI] Align XPU Triton hook
 with upstream LLVM

Imported from GitHub PR https://github.com/openxla/xla/pull/50870

This PR aligns the XPU Triton hook with https://github.com/llvm/llvm-project/commit/1ee8805f01a663b7f2b09215a8eab5ab2d15ce35. It includes the fix for this breaking LLVM change https://github.com/llvm/llvm-project/commit/80c5d7322718ceed7eb51c05790ea6cdfd373aff that renamed `ArithCommon` to `ArithAttrToLLVM` and https://github.com/llvm/llvm-project/commit/076bd55e20d92c130cf2d50a6324dd457861aaea that moves Ceil/Floor div expansion to ArithCommon.
Copybara import of the project:

--
c4da41239e9aa1305962b602ae962c68592fc0b5 by Akhil Goel <akhil.goel@intel.com>:

Update XPU Triton hook

--
81b2fd498a3374d3400919b7406c33eadfea7d25 by Akhil Goel <akhil.goel@intel.com>:

Include fix for https://github.com/llvm/llvm-project/commit/076bd55e20d92c130cf2d50a6324dd457861aaea that will be part of a future LLVM update

Merging this change closes #50870

PiperOrigin-RevId: 1000758805
---
 third_party/xla/third_party/triton/intel_xpu/workspace.bzl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/third_party/triton/intel_xpu/workspace.bzl b/third_party/xla/third_party/triton/intel_xpu/workspace.bzl
index bf183ad297a..50a6b8150dc 100644
--- a/third_party/xla/third_party/triton/intel_xpu/workspace.bzl
+++ b/third_party/xla/third_party/triton/intel_xpu/workspace.bzl
@@ -1,7 +1,7 @@
 """Intel XPU Triton archive metadata."""
 
-XPU_TRITON_COMMIT = "e45abef5f100e02e2f9e98f5f5a2fd8dfe3c6c34"
-XPU_TRITON_SHA256 = "a0d191feab8e91bc710f96fb0aefde8a22b0c99daac4f2c4153adc8f6c98b26c"
+XPU_TRITON_COMMIT = "491b0483ac16baa85561a499c5ee754bf5e8a96e"
+XPU_TRITON_SHA256 = "a4fc57e44982fbd2eab6019c63b63e24d807b2708bbc299e95aa252d62571c85"
 
 def use_xpu_triton(repository_ctx):
     return repository_ctx.getenv("ENABLE_INTEL_XPU_TRITON", "").strip() == "1"

From 42403c375d497c8c72e1a63615b06cab2d2fd26b Mon Sep 17 00:00:00 2001
From: Alexander Lyashuk <crem@google.com>
Date: Fri, 16 Oct 2026 04:19:09 -0700
Subject: [PATCH 0952/1000] [XLA:GPU] Triton: don't transpose TMA tiles if the
 layout only moves unit dims.

TMA tiles with a non-normalized layout are transposed between physical and
logical order. With batch tile size 1, the transpose only moves unit dims, yet
Triton round-trips the tile through registers instead of feeding shared memory
directly to the dot.

Emit a reshape instead. Rank-reduced stores with a major-to-minor layout now
also use tt.reshape instead of tt.expand_dims; both lower to identical LLVM IR.

E.g. bf16[16,4256,256] = dot(convert(s4[16,8512,4256]), bf16[16,256,8512]{1,0,2})
with 1x64x128 tiles on H100: 900us -> 632us.

No effect while dots are normalized; matters once follow-up changes relax that.

PiperOrigin-RevId: 1000764236
---
 .../triton_xla_extract_insert_to_triton.mlir  | 48 +++++++++++++++++++
 ...riton_xla_extract_insert_to_triton_pass.cc | 44 ++++++++++++-----
 2 files changed, 79 insertions(+), 13 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_extract_insert_to_triton.mlir b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_extract_insert_to_triton.mlir
index a86502850d8..92f6d1daf21 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_extract_insert_to_triton.mlir
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/tests/triton_xla_extract_insert_to_triton.mlir
@@ -607,3 +607,51 @@ module {
 // CHECK: %[[SPLAT_MASK:.*]] = tt.splat %[[MASK]] : i1 -> tensor<8xi1>
 // CHECK: tt.load {{.*}}, %[[SPLAT_MASK]], {{.*}}
 // CHECK: tt.store {{.*}}, %{{.*}} : tensor<8x!tt.ptr<bf16>>
+
+// -----
+
+func.func @extract_insert_with_unit_dim_only_layout_permutation(
+    %arg0: !tt.ptr<bf16>, %arg1: !tt.ptr<bf16>) {
+  %extracted_tensor = triton_xla.extract from %arg0
+      as memref<16x8512x256xbf16, #xtile.layout<[2, 0, 1]>>
+      [0, 0, 0] [1, 64, 128] [1, 1, 1] : tensor<64x128xbf16>
+  triton_xla.insert %extracted_tensor into %arg1
+      as memref<16x8512x256xbf16, #xtile.layout<[2, 0, 1]>>
+      [0, 0, 0] [1, 64, 128] [1, 1, 1] : tensor<64x128xbf16>
+  func.return
+}
+
+// CHECK-TMA-LABEL: tt.func @extract_insert_with_unit_dim_only_layout_permutation
+// CHECK-TMA-SAME:      %arg0: !tt.tensordesc<64x1x128xbf16>
+// CHECK-TMA-SAME:      %arg1: !tt.tensordesc<64x1x128xbf16>
+// CHECK-TMA-NOT:     tt.trans
+// CHECK-TMA:         %[[LOAD:.*]] = tt.descriptor_load %arg0
+// CHECK-TMA-SAME:        -> tensor<64x1x128xbf16>
+// CHECK-TMA-NOT:     tt.trans
+// CHECK-TMA:         %[[RESHAPE0:.*]] = tt.reshape %[[LOAD]]
+// CHECK-TMA-SAME:        tensor<64x1x128xbf16> -> tensor<64x128xbf16>
+// CHECK-TMA-NOT:     tt.trans
+// CHECK-TMA:         %[[RESHAPE1:.*]] = tt.reshape %[[RESHAPE0]]
+// CHECK-TMA-SAME:        tensor<64x128xbf16> -> tensor<64x1x128xbf16>
+// CHECK-TMA-NOT:     tt.trans
+// CHECK-TMA:         tt.descriptor_store %arg1[{{.*}}], %[[RESHAPE1]]
+
+// -----
+
+func.func @extract_with_transposing_layout(
+    %arg0: !tt.ptr<bf16>, %arg1: !tt.ptr<bf16>) {
+  %extracted_tensor = triton_xla.extract from %arg0
+      as memref<16x8512x256xbf16, #xtile.layout<[2, 0, 1]>>
+      [0, 0, 0] [2, 64, 128] [1, 1, 1] : tensor<2x64x128xbf16>
+  triton_xla.insert %extracted_tensor into %arg1
+      as memref<2x64x128xbf16, #xtile.layout<[2, 1, 0]>>
+      [0, 0, 0] [2, 64, 128] [1, 1, 1] : tensor<2x64x128xbf16>
+  func.return
+}
+
+// CHECK-TMA-LABEL: tt.func @extract_with_transposing_layout
+// CHECK-TMA-SAME:      %arg0: !tt.tensordesc<64x2x128xbf16>
+// CHECK-TMA:         %[[LOAD:.*]] = tt.descriptor_load %arg0
+// CHECK-TMA:         %[[TRANS:.*]] = tt.trans %[[LOAD]]
+// CHECK-TMA-SAME:        tensor<64x2x128xbf16> -> tensor<2x64x128xbf16>
+// CHECK-TMA:         tt.descriptor_store %arg1[{{.*}}], %[[TRANS]]
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_extract_insert_to_triton_pass.cc b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_extract_insert_to_triton_pass.cc
index 594f7529f62..3e29789e119 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_extract_insert_to_triton_pass.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/transforms/triton_xla_extract_insert_to_triton_pass.cc
@@ -127,6 +127,16 @@ bool IsOffsetDivisibilityGuaranteed(mlir::Value offset_val,
   return xtriton::IsGuaranteedDivisible(offset_val, divisor);
 }
 
+// Returns true if `layout` is major-to-minor up to unit dimensions of `shape`,
+// i.e. transposing between logical and physical order is a reshape.
+bool IsMajorToMinorIgnoringUnitDims(ArrayRef<int64_t> shape,
+                                    ArrayRef<int64_t> layout) {
+  int64_t last_non_unit_dim = -1;
+  return absl::c_all_of(llvm::reverse(layout), [&](int64_t dim) {
+    return shape[dim] == 1 || std::exchange(last_non_unit_dim, dim) < dim;
+  });
+}
+
 // Limitations of TMA are documented in IsTmaCompatible
 // in third_party/tensorflow/compiler/xla/stream_executor/gpu/tma_metadata.h.
 // Additionally:
@@ -377,13 +387,16 @@ class RewriteExtract : public mlir::OpRewritePattern<ExtractOp> {
           xtriton::IndexCast(builder, builder.getI32Type(), ordered_offsets),
           /*cachePolicy=*/nullptr);
 
-      // Insert a transpose if the layout is not major-to-minor.
-      if (!xtriton::IsMajorToMinorLayout(src_layout)) {
+      // Insert a transpose if the layout is not major-to-minor, unless it only
+      // moves unit dimensions. A reshape suffices then, and lets Triton feed
+      // the TMA-loaded shared memory directly to a dot.
+      if (!IsMajorToMinorIgnoringUnitDims(sizes, src_layout)) {
         result = TransOp::create(
             builder, result, xtriton::GetInverseLayoutPermutation(src_layout));
       }
-      // Insert a reshape if the result is rank-reduced.
-      if (sizes.size() != tile_shape.size()) {
+      // Insert a reshape if the result is rank-reduced or still in physical
+      // order.
+      if (result.getType() != tile_type) {
         result = ReshapeOp::create(builder, tile_shape, result,
                                    /*allowReorder=*/false);
       }
@@ -519,16 +532,21 @@ class RewriteInsert : public mlir::OpRewritePattern<InsertOp> {
           builder, desc_type, op.getDst());
 
       Value src = op.getSrc();
-      // Insert a expand_dims if the source is rank-reduced.
-      for (auto dim : reduced_dims) {
-        src = ExpandDimsOp::create(builder, src, dim);
-      }
-      // Insert a transpose if the layout is not major-to-minor.
-      if (!xtriton::IsMajorToMinorLayout(dst_layout)) {
+      if (!IsMajorToMinorIgnoringUnitDims(sizes, dst_layout)) {
+        // Insert an expand_dims if the source is rank-reduced.
+        for (auto dim : reduced_dims) {
+          src = ExpandDimsOp::create(builder, src, dim);
+        }
         // Transpose to a major-to-minor tensor by simply reversing the layout.
-        auto transpose_order = llvm::to_vector_of<int32_t>(dst_layout);
-        std::reverse(transpose_order.begin(), transpose_order.end());
-        src = TransOp::create(builder, src, transpose_order);
+        src = TransOp::create(
+            builder, src,
+            llvm::to_vector_of<int32_t>(llvm::reverse(dst_layout)));
+      }
+      // Insert a reshape if the source is rank-reduced or still in logical
+      // order.
+      if (src.getType() != result_type) {
+        src = ReshapeOp::create(builder, result_type.getShape(), src,
+                                /*allowReorder=*/false);
       }
 
       auto ordered_offsets = xtriton::GetMajorToMinorOrder(offsets, dst_layout);

From cb82efaca8d43022c398cc48cee4809dd9471689 Mon Sep 17 00:00:00 2001
From: Luca Bruni <lucbruni@amd.com>
Date: Fri, 16 Oct 2026 04:22:05 -0700
Subject: [PATCH 0953/1000] PR #50778: [ROCm] Add gfx1102, gfx1152 and gfx1153
 to supported AMDGPU versions
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50778

📝 Summary of Changes
Add gfx1102, gfx1152 and gfx1153 to `kSupportedGfxVersions`, and to the gfx11 family lists behind `has_hipblaslt()`, matching their already-supported siblings gfx1100/gfx1101/gfx1103/gfx1150/gfx1151.

🎯 Justification
JAX's ROCm build compiles kernels for these targets (jax-ml/jax#38803), but `IsDeviceSupported()` rejects them at runtime, so JAX fails with `no supported devices found for platform ROCM` on e.g. the Radeon 860M (gfx1152). Reported in ROCm/legacy-rocm-build#6694.

🚀 Kind of Contribution
✨ New Feature

🧪 Unit Tests
Extended `//xla/stream_executor:device_description_test`.

🧪 Execution Tests
JAX 0.11.1 plugin rebuilt with this change on ROCm 10.1.0: runs correctly on a Radeon 860M (gfx1152), previously rejected; RX 9070 XT (gfx1201) unchanged.

Copybara import of the project:

--
f78db90a0e9c6f34ada270c049edfaf4785adb2e by Luca Bruni <lucbruni@amd.com>:

[ROCm] Add gfx1102, gfx1152 and gfx1153 to supported AMDGPU versions

JAX's ROCm build compiles kernels for these RDNA3/3.5 targets
(jax-ml/jax#38803), but IsDeviceSupported() rejects them at runtime
because they are missing from kSupportedGfxVersions, so JAX fails with
"no supported devices found for platform ROCM" on e.g. the Radeon 860M
(gfx1152).

Add them to the allowlist, and to the gfx11 family lists that drive
has_hipblaslt() (gfx1102 discrete, gfx1152/gfx1153 APU), matching their
already-supported siblings gfx1100/gfx1101/gfx1103/gfx1150/gfx1151.

Merging this change closes #50778

PiperOrigin-RevId: 1000765538
---
 .../stream_executor/device_description_test.cc   | 13 +++++++++++++
 .../rocm/rocm_compute_capability.h               | 16 ++++++++++++----
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/stream_executor/device_description_test.cc b/third_party/xla/xla/stream_executor/device_description_test.cc
index b09ab9994c8..63acaea73ff 100644
--- a/third_party/xla/xla/stream_executor/device_description_test.cc
+++ b/third_party/xla/xla/stream_executor/device_description_test.cc
@@ -75,6 +75,11 @@ TEST(RocmComputeCapability, IsSupportedGfxVersion) {
   ASSERT_TRUE(RocmComputeCapability{"gfx1201"}.is_supported_gfx_version());
   ASSERT_TRUE(RocmComputeCapability{"gfx942"}.is_supported_gfx_version());
   ASSERT_TRUE(RocmComputeCapability{"gfx1250"}.is_supported_gfx_version());
+  ASSERT_TRUE(RocmComputeCapability{"gfx1102"}.is_supported_gfx_version());
+  ASSERT_TRUE(RocmComputeCapability{"gfx1152"}.is_supported_gfx_version());
+  ASSERT_TRUE(RocmComputeCapability{"gfx1153"}.is_supported_gfx_version());
+  ASSERT_TRUE(
+      RocmComputeCapability{"gfx1152:xnack-"}.is_supported_gfx_version());
   ASSERT_FALSE(RocmComputeCapability{"some_string"}.is_supported_gfx_version());
 }
 
@@ -123,6 +128,14 @@ TEST(RocmComputeCapability, Accessors) {
   EXPECT_TRUE(RocmComputeCapability{"gfx1200"}.has_hipblaslt());
   EXPECT_TRUE(RocmComputeCapability{"gfx1100"}.has_hipblaslt());
   EXPECT_TRUE(RocmComputeCapability{"gfx1103"}.has_hipblaslt());
+  EXPECT_TRUE(RocmComputeCapability{"gfx1102"}.has_hipblaslt());
+  EXPECT_TRUE(RocmComputeCapability{"gfx1152"}.has_hipblaslt());
+  EXPECT_TRUE(RocmComputeCapability{"gfx1153"}.has_hipblaslt());
+  EXPECT_TRUE(RocmComputeCapability{"gfx1102"}.gfx11_discrete());
+  EXPECT_TRUE(RocmComputeCapability{"gfx1152"}.gfx11_apu());
+  EXPECT_TRUE(RocmComputeCapability{"gfx1153"}.gfx11_apu());
+  EXPECT_FALSE(RocmComputeCapability{"gfx1152"}.gfx11_discrete());
+
   EXPECT_TRUE(RocmComputeCapability{"gfx1250"}.has_hipblaslt());
 
   EXPECT_FALSE(RocmComputeCapability{"gfx1250"}.has_nhwc_layout_support());
diff --git a/third_party/xla/xla/stream_executor/rocm/rocm_compute_capability.h b/third_party/xla/xla/stream_executor/rocm/rocm_compute_capability.h
index 1a6cb860c4a..e5b18ace779 100644
--- a/third_party/xla/xla/stream_executor/rocm/rocm_compute_capability.h
+++ b/third_party/xla/xla/stream_executor/rocm/rocm_compute_capability.h
@@ -94,7 +94,14 @@ class RocmComputeCapability {
       "gfx1030",  // RX68xx / RX69xx
       "gfx1100",  // RX7900
       "gfx1101",  // RX7700 / RX7800
-      "gfx1103", "gfx1150", "gfx1151", "gfx1200", "gfx1201",
+      "gfx1102",  // RX7600
+      "gfx1103",  // Radeon 780M
+      "gfx1150",  // Radeon 890M
+      "gfx1151",  // Radeon 8060S
+      "gfx1152",  // Radeon 860M
+      "gfx1153",  // Radeon 820M
+      "gfx1200",  // RX9060
+      "gfx1201",  // RX9070
       "gfx1250",  // MI450
   };
 
@@ -135,11 +142,12 @@ class RocmComputeCapability {
 
   bool gfx11() const { return absl::StartsWith(gfx_version(), "gfx11"); }
 
-  static constexpr absl::string_view kGfx11Discrete[] = {"gfx1100", "gfx1101"};
+  static constexpr absl::string_view kGfx11Discrete[] = {"gfx1100", "gfx1101",
+                                                         "gfx1102"};
   bool gfx11_discrete() const { return IsThisGfxInAnyList(kGfx11Discrete); }
 
-  static constexpr absl::string_view kGfx11Apu[] = {"gfx1103", "gfx1150",
-                                                    "gfx1151"};
+  static constexpr absl::string_view kGfx11Apu[] = {
+      "gfx1103", "gfx1150", "gfx1151", "gfx1152", "gfx1153"};
   bool gfx11_apu() const { return IsThisGfxInAnyList(kGfx11Apu); }
 
   static constexpr absl::string_view kGfx11Rx7900[] = {"gfx1100", "gfx1101",

From 124aa5ebff1fb70c709803e30d288b024ba448ca Mon Sep 17 00:00:00 2001
From: Dmitri Latushko <latushko@google.com>
Date: Fri, 16 Oct 2026 04:28:51 -0700
Subject: [PATCH 0954/1000] Rollback of cl/999614913 and forward fix for
 x86/ARM64

cl/999614913 modified the numerical computation in tf.nn.moments to subtract the squared residual mean before reduction to fix precision issues on XLA x86. However, this broke numerical precision on ARM64 architectures.

This CL rolls back the change in tf.nn.moments to restore ARM64 correctness, and instead fixes the original Kokoro CPU continuous builds failure by relaxing the absolute tolerance in MomentsTest.testIdenticalLargeValues for float32. A variance of 2.0 on inputs of 7.5e10 corresponds to a standard deviation of 1.414, which is a relative error of ~1.9e-11, well within float32 precision.

PiperOrigin-RevId: 1000768225
---
 tensorflow/python/ops/nn_impl.py | 14 +++-----------
 tensorflow/python/ops/nn_test.py |  5 +++--
 2 files changed, 6 insertions(+), 13 deletions(-)

diff --git a/tensorflow/python/ops/nn_impl.py b/tensorflow/python/ops/nn_impl.py
index c18000798a1..ed578011628 100644
--- a/tensorflow/python/ops/nn_impl.py
+++ b/tensorflow/python/ops/nn_impl.py
@@ -1459,17 +1459,9 @@ def moments(
     #       because that gradient is zero
     diff = math_ops.subtract(y, array_ops.stop_gradient(mean))
     variance = math_ops.maximum(
-        math_ops.reduce_mean(
-            math_ops.subtract(
-                math_ops.square(diff),
-                math_ops.square(
-                    array_ops.stop_gradient(
-                        math_ops.reduce_mean(diff, axes, keepdims=True)
-                    )
-                ),
-            ),
-            axes,
-            keepdims=True,
+        math_ops.subtract(
+            math_ops.reduce_mean(math_ops.square(diff), axes, keepdims=True),
+            math_ops.square(math_ops.reduce_mean(diff, axes, keepdims=True)),
         ),
         math_ops.cast(0, y.dtype),
         name="variance",
diff --git a/tensorflow/python/ops/nn_test.py b/tensorflow/python/ops/nn_test.py
index cdbd81da8b1..6f63b1cb0ea 100644
--- a/tensorflow/python/ops/nn_test.py
+++ b/tensorflow/python/ops/nn_test.py
@@ -1274,7 +1274,8 @@ class MomentsTest(test_lib.TestCase):
             t = constant_op.constant([[val, val, val]], dtype=dtype)
             _, var_eager = nn_impl.moments(t, axes=[1], keepdims=keepdims)
             expected = [[0.0]] if keepdims else [0.0]
-            self.assertAllClose(var_eager, expected, atol=1e-5)
+            atol = 2.5 if dtype == dtypes.float32 else 1e-5
+            self.assertAllClose(var_eager, expected, atol=atol)
             try:
               _, var_jit = def_function.function(
                   lambda x, kd=keepdims: nn_impl.moments(
@@ -1282,7 +1283,7 @@ class MomentsTest(test_lib.TestCase):
                   ),
                   jit_compile=True,
               )(t)
-              self.assertAllClose(var_jit, expected, atol=1e-5)
+              self.assertAllClose(var_jit, expected, atol=atol)
             except (errors.UnimplementedError, errors.InvalidArgumentError):
               pass
 

From 9a6d454d18c5e86fc113498f683627101c8d3157 Mon Sep 17 00:00:00 2001
From: Henning Becker <hebecker@google.com>
Date: Fri, 16 Oct 2026 04:29:51 -0700
Subject: [PATCH 0955/1000] Enable RBE for B200 benchmark builds

`XLA_LINUX_X86_GPU_A4_224_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS` only configured `cuda_libraries_from_stubs` and was missing `rbe_linux_cuda_nvcc`, `hermetic_cuda_umd`, and `warnings`. Because of that, B200 benchmark jobs compiled all ~15,500 C++ and CUDA actions locally on the runner (taking 34 to 52 minutes for `bazel build` alone) instead of using RBE like the L4 benchmark build (`XLA_LINUX_X86_GPU_L4_16_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS`).

Therefore this change aligns the B200 benchmark build configuration with the L4 benchmark build so compilation runs on RBE while the resulting benchmark binary continues to run locally on the B200 GPU. In presubmit testing on `nv_maxtext_llama3_8b_1n8g_after_spmd_partitioner`, this reduced `bazel build` time from 33.9 minutes (`15,501 local` actions) to 9.3 minutes (`14,260 remote cache hit, 1,233 remote`, 0 local actions) even with a cold `sm_100` RBE cache, cutting total job time from 44 minutes to 20 minutes.

PiperOrigin-RevId: 1000768642
---
 third_party/xla/build_tools/ci/build.py            | 8 ++++++--
 third_party/xla/build_tools/ci/golden_commands.txt | 4 ++--
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/build_tools/ci/build.py b/third_party/xla/build_tools/ci/build.py
index 7af54c9130c..d8abca54c91 100755
--- a/third_party/xla/build_tools/ci/build.py
+++ b/third_party/xla/build_tools/ci/build.py
@@ -696,14 +696,18 @@ Build(
 Build(
     type_=BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS,
     repo="openxla/xla",
-    configs=("cuda_libraries_from_stubs",),
+    configs=(
+        "warnings",
+        "rbe_linux_cuda_nvcc",
+        "hermetic_cuda_umd",
+        "cuda_libraries_from_stubs",
+    ),
     target_patterns=_XLA_GPU_PRESUBMIT_BENCHMARKS_DEFAULT_TARGET_PATTERNS,
     test_tag_filters=nvidia_single_gpu_test_filters
     + _tag_filters_for_compute_capability(compute_capability=100),
     build_tag_filters=nvidia_single_gpu_build_filters,
     options={
         "run_under": "//build_tools/ci:parallel_gpu_execute",
-        # Use User Mode and Kernel Mode Drivers pre-installed on the system.
         "//xla/tsl:ci_build": True,
         **_DEFAULT_BAZEL_OPTIONS,
     },
diff --git a/third_party/xla/build_tools/ci/golden_commands.txt b/third_party/xla/build_tools/ci/golden_commands.txt
index 65b208da349..785034cec7f 100644
--- a/third_party/xla/build_tools/ci/golden_commands.txt
+++ b/third_party/xla/build_tools/ci/golden_commands.txt
@@ -53,8 +53,8 @@ bazel test --build_tag_filters=-no_oss,-rocm-only,-oneapi-only,multi_gpu,-xla_nv
 # END BuildType.XLA_LINUX_X86_GPU_8X_H100_GITHUB_ACTIONS
 # BEGIN BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS
 nvidia-smi
-parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
-bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
+parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --nobuild -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
+bazel build --build_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu --test_tag_filters=-no_oss,requires-gpu-nvidia,-rocm-only,-oneapi-only,gpu,-multi_gpu,requires-gpu-sm100-only,requires-gpu-sm60,requires-gpu-sm70,requires-gpu-sm80,requires-gpu-sm90,requires-gpu-sm100,-requires-gpu-sm103,-requires-gpu-sm103-only,-requires-gpu-sm107,-requires-gpu-sm107-only,-requires-gpu-sm120,-requires-gpu-sm120-only,-requires-gpu-amd,-requires-gpu-intel --config=warnings --config=rbe_linux_cuda_nvcc --config=hermetic_cuda_umd --config=cuda_libraries_from_stubs --repo_env=TF_CUDA_COMPUTE_CAPABILITIES=10 --run_under=//build_tools/ci:parallel_gpu_execute --//xla/tsl:ci_build --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async -- //xla/tools/multihost_hlo_runner:hlo_runner_main_gpu //xla/tools:compute_xspace_stats_main_gpu
 # END BuildType.XLA_LINUX_X86_GPU_A4_224_VCPU_BENCHMARK_PRESUBMIT_GITHUB_ACTIONS
 # BEGIN BuildType.XLA_LINUX_X86_GPU_HERMETIC_ROCM_GITHUB_ACTIONS
 parallel --ungroup --retries 3 --delay 15 --nonall -- bazel build --build_tag_filters=-no_gpu,-requires-gpu-intel,-requires-gpu-nvidia,-cuda-only,-oneapi-only,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm86,-requires-gpu-sm86-only,-requires-gpu-sm89,-requires-gpu-sm89-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial,gpu --test_tag_filters=-no_gpu,-requires-gpu-intel,-requires-gpu-nvidia,-cuda-only,-oneapi-only,-requires-gpu-sm60,-requires-gpu-sm60-only,-requires-gpu-sm70,-requires-gpu-sm70-only,-requires-gpu-sm80,-requires-gpu-sm80-only,-requires-gpu-sm86,-requires-gpu-sm86-only,-requires-gpu-sm89,-requires-gpu-sm89-only,-requires-gpu-sm90,-requires-gpu-sm90-only,-skip_rocprofiler_sdk,-no_oss,-oss_excluded,-oss_serial,gpu --config=warnings --config=rbe_linux_cpu --config=rocm_clang_hermetic --config=rocm_ci_hermetic --color=yes --test_output=errors --verbose_failures --keep_going --nobuild_tests_only --profile=profile.json.gz --flaky_test_attempts=3 --jobs=150 --bes_upload_mode=fully_async --//xla/tsl:ci_build --nobuild -- //xla/... //build_tools/... @tsl//tsl/...

From 7fd513f228b08f01db70689944e9ae85ea5ecc25 Mon Sep 17 00:00:00 2001
From: Kostiantyn Liepieshov <kostiantynl@google.com>
Date: Fri, 16 Oct 2026 05:16:21 -0700
Subject: [PATCH 0956/1000] [XLA] Relay control dependencies when merging
 sibling fusions

MultiOutputFusion dropped the control edges of the instructions it merged.
CreateFusion replaced the base instruction without relaying its control
dependencies. Fuse also removed the fused sibling while its control edges were
still attached, and HloInstruction::Defuse removed the fusion without relaying
its control edges, so scheduler ordering was lost or the removal failed.

CreateFusion now relays the base's control edges to the new fusion. Fuse moves
the sibling's control predecessors and successors onto the surviving fusion
before merging, which cannot form a cycle because siblings are unconnected.
Defuse relays control predecessors to the first defused instruction and control
successors from the defused root before removing the fusion.

PiperOrigin-RevId: 1000787087
---
 third_party/xla/xla/hlo/ir/hlo_instruction.cc | 17 ++++++++--
 .../xla/xla/hlo/ir/hlo_instruction_test.cc    | 34 +++++++++++++++++++
 .../xla/xla/service/multi_output_fusion.cc    | 10 +++++-
 3 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.cc b/third_party/xla/xla/hlo/ir/hlo_instruction.cc
index cb9317d5877..810fa93619c 100644
--- a/third_party/xla/xla/hlo/ir/hlo_instruction.cc
+++ b/third_party/xla/xla/hlo/ir/hlo_instruction.cc
@@ -3717,6 +3717,7 @@ absl::Status HloInstruction::Defuse() {
   // computation as the fusion instruction itself.
   // TODO(b/68227302): Moving instruction to new computation rather than
   // cloning and deleting.
+  HloInstruction* first_defused = nullptr;
   for (HloInstruction* fused_instruction :
        fused_computation->MakeInstructionPostOrder()) {
     if (fused_instruction->opcode() == HloOpcode::kParameter) {
@@ -3729,11 +3730,23 @@ absl::Status HloInstruction::Defuse() {
     HloInstruction* defused_instruction =
         parent()->AddInstruction(fused_instruction->CloneWithNewOperands(
             fused_instruction->shape(), new_operands));
+    if (first_defused == nullptr) {
+      first_defused = defused_instruction;
+    }
     defused_instructions[fused_instruction] = defused_instruction;
   }
 
-  ABSL_RETURN_IF_ERROR(
-      ReplaceAllUsesWith(defused_instructions.at(fused_expression_root())));
+  HloInstruction* defused_root =
+      defused_instructions.at(fused_expression_root());
+  ABSL_RETURN_IF_ERROR(ReplaceAllUsesWith(defused_root));
+  if (HasControlDependencies()) {
+    if (first_defused != nullptr) {
+      ABSL_RETURN_IF_ERROR(CopyAllControlDepsTo(first_defused, defused_root));
+      ABSL_RETURN_IF_ERROR(DropAllControlDeps());
+    } else {
+      ABSL_RETURN_IF_ERROR(SafelyDropAllControlDependencies());
+    }
+  }
 
   HloModule* module = GetModule();
   ABSL_RETURN_IF_ERROR(parent()->RemoveInstruction(this));
diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc b/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc
index 1f51effb9c0..326e94ea378 100644
--- a/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc
+++ b/third_party/xla/xla/hlo/ir/hlo_instruction_test.cc
@@ -1155,5 +1155,39 @@ ENTRY %main (p0: f32[10]) -> f32[10] {
             async_update->output_operand_aliasing());
 }
 
+TEST_F(HloInstructionTest, DefuseRelaysControlDependencies) {
+  constexpr absl::string_view kHloString = R"(
+HloModule DefuseControlDeps
+
+%fused_comp (p: f32[8]) -> f32[8] {
+  %p = f32[8] parameter(0)
+  %cp = f32[8] copy(%p)
+  ROOT %bc = f32[8] bitcast(%cp)
+}
+
+ENTRY %main (x: f32[8], y: f32[8]) -> f32[8] {
+  %x = f32[8] parameter(0)
+  %y = f32[8] parameter(1)
+  %ctrl_pred = f32[8] negate(%y)
+  %fus = f32[8] fusion(%x), kind=kLoop, calls=%fused_comp, control-predecessors={%ctrl_pred}
+  ROOT %out = f32[8] add(%fus, %ctrl_pred), control-predecessors={%fus}
+}
+)";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloString));
+  HloComputation* entry = module->entry_computation();
+  HloInstruction* ctrl_pred = entry->GetInstructionWithName("ctrl_pred");
+  HloInstruction* fus = entry->GetInstructionWithName("fus");
+  HloInstruction* out = entry->GetInstructionWithName("out");
+  ASSERT_OK(fus->Defuse());
+  HloInstruction* defused_root = out->mutable_operand(0);
+  ASSERT_EQ(defused_root->opcode(), HloOpcode::kBitcast);
+  HloInstruction* defused_copy = defused_root->mutable_operand(0);
+  ASSERT_EQ(defused_copy->opcode(), HloOpcode::kCopy);
+  EXPECT_THAT(defused_copy->control_predecessors(),
+              ::testing::ElementsAre(ctrl_pred));
+  EXPECT_THAT(out->control_predecessors(),
+              ::testing::ElementsAre(defused_root));
+}
+
 }  // namespace
 }  // namespace xla
diff --git a/third_party/xla/xla/service/multi_output_fusion.cc b/third_party/xla/xla/service/multi_output_fusion.cc
index ca6d47574fc..4df429643e4 100644
--- a/third_party/xla/xla/service/multi_output_fusion.cc
+++ b/third_party/xla/xla/service/multi_output_fusion.cc
@@ -105,6 +105,10 @@ HloInstruction* MultiOutputFusion::Fuse(HloInstruction* instr1,
   if (remaining->opcode() != HloOpcode::kFusion) {
     remaining = CreateFusion(remaining, fused);
   }
+  // `fused` now computes as part of `remaining`, so its ordering constraints
+  // move there. The siblings are unconnected, so this adds no cycle.
+  CHECK_OK(remaining->CopyAllControlDepsFrom(fused));
+  CHECK_OK(fused->DropAllControlDeps());
   if (fused->opcode() == HloOpcode::kFusion) {
     remaining->MergeFusionInstructionIntoMultiOutput(
         fused, /*remove_computation=*/false);
@@ -126,7 +130,11 @@ HloInstruction* MultiOutputFusion::CreateFusion(HloInstruction* base,
   candidates_.emplace_back(input_fusion);
   reachability_->Replace(base, input_fusion);
   all_fusion_candidates_.insert(input_fusion);
-  CHECK_OK(computation()->ReplaceInstruction(base, input_fusion));
+  CHECK_OK(computation()
+               ->ReplaceInstruction(base, input_fusion,
+                                    /*preserve_sharding=*/false,
+                                    /*relay_control_dependency=*/true)
+               .status());
   return input_fusion;
 }
 

From 9c72341a2e82d9dc46219f0e1d687225f7291f87 Mon Sep 17 00:00:00 2001
From: Sean Talts <talts@google.com>
Date: Fri, 16 Oct 2026 05:27:59 -0700
Subject: [PATCH 0957/1000] [XLA:CPU] Deduplicate kernels of hoisted small
 while loops.

### Status quo

`SmallWhileLoopHoistingPass` turns small while loops into calls, and each call
is compiled into its own kernel. Unlike loop fusions, these kernels are never
reused, so identical loops are compiled again and again.

### Why

cl/998323502 raised `xla_cpu_small_while_loop_byte_threshold` from 1 KiB to
1 GB, so many more loops are now hoisted. In b/573148551, inlining clones 87
JAX PRNG loops into 1,403 hoisted calls, and compile time went from 68s to
157s.

### Design

`ComputationKernelArguments` lists the buffers a computation kernel uses, in a
deterministic order, and computes a fingerprint of everything the emitted code
depends on: the computation text (including constants), which buffer each value
uses, buffer sizes and overlaps, and constants read from outside. Calls with
equal fingerprints share one kernel in `ThunkEmitter`. Computations with
`rng-get-and-update-state` are never shared, since their kernel owns private
RNG state. `ComputationKernelEmitter` now takes its arguments from this class,
which also makes kernel argument order deterministic.

### Path to MLIR

`ComputationKernelArguments` does not depend on `IrEmitter`. The planned
follow-up emits each hoisted loop as one MLIR kernel: body fusions come from
the existing CPU fusion emitters and the loop becomes `scf.for`/`scf.while`,
behind a flag with fallback to `IrEmitter`. That emitter uses this class for
its kernel arguments and buffer slice indices, and the same kernel cache with
its own key prefix. It must still check that the fingerprint covers
everything it bakes into the code.

### Alternatives considered

*   Lowering the threshold again: the cloned loops in b/573148551 span all
    sizes from under 1 KiB to over 16 MiB, so no threshold fixes compile time
    without losing the runtime wins of cl/998323502 (up to -96%).
*   Reusing `emitters::GetComputationFingerprint`: it only covers the operands
    and outputs of a fusion, not the intermediate buffers of a whole
    computation.
*   Moving these loops to the MLIR emitters first: the MLIR emitters do not
    support while loops yet. This change fixes compile time now.

### Performance

*   b/573148551: 212 kernels instead of 1,403; compile+run 157s -> 61-69s.
*   Fingerprinting costs at most 0.22% of compile time there (about 0.2 ms per
    call), while each reused kernel saves about 75 ms.
*   Other CPU benchmarks: 55 of 85 have no hoisted loops; most others have 1-5,
    where the cost is under 1 ms.
*   Generated code and results are unchanged (bitwise identical).

PiperOrigin-RevId: 1000791507
---
 .../xla/xla/backends/cpu/codegen/BUILD        |  45 +-
 .../codegen/computation_kernel_arguments.cc   | 232 +++++++++
 .../codegen/computation_kernel_arguments.h    |  83 ++++
 .../computation_kernel_arguments_test.cc      | 463 ++++++++++++++++++
 .../cpu/codegen/computation_kernel_emitter.cc | 116 +----
 .../cpu/codegen/computation_kernel_emitter.h  |   7 +-
 .../xla/xla/backends/cpu/testlib/BUILD        |   2 +
 .../cpu/testlib/kernel_runner_extension.cc    |  23 +-
 third_party/xla/xla/service/cpu/BUILD         |   6 +-
 .../cpu/cpu_compiler_internals_test.cc        |  85 ++++
 .../xla/service/cpu/tests/cpu_while_test.cc   |  54 ++
 .../xla/xla/service/cpu/thunk_emitter.cc      |  32 +-
 12 files changed, 1045 insertions(+), 103 deletions(-)
 create mode 100644 third_party/xla/xla/backends/cpu/codegen/computation_kernel_arguments.cc
 create mode 100644 third_party/xla/xla/backends/cpu/codegen/computation_kernel_arguments.h
 create mode 100644 third_party/xla/xla/backends/cpu/codegen/computation_kernel_arguments_test.cc

diff --git a/third_party/xla/xla/backends/cpu/codegen/BUILD b/third_party/xla/xla/backends/cpu/codegen/BUILD
index a3759df9e0c..17c0c63e198 100644
--- a/third_party/xla/xla/backends/cpu/codegen/BUILD
+++ b/third_party/xla/xla/backends/cpu/codegen/BUILD
@@ -644,14 +644,55 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "computation_kernel_arguments",
+    srcs = ["computation_kernel_arguments.cc"],
+    hdrs = ["computation_kernel_arguments.h"],
+    deps = [
+        "//xla:shape_util",
+        "//xla/codegen:kernel_spec",
+        "//xla/hlo/ir:hlo",
+        "//xla/service:buffer_assignment",
+        "//xla/service:shaped_slice",
+        "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+xla_cc_test(
+    name = "computation_kernel_arguments_test",
+    srcs = ["computation_kernel_arguments_test.cc"],
+    deps = [
+        ":computation_kernel_arguments",
+        "//xla/hlo/analysis:alias_info",
+        "//xla/hlo/analysis:hlo_ordering",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+        "//xla/service:buffer_assignment",
+        "//xla/service:buffer_value",
+        "//xla/service:logical_buffer",
+        "//xla/service/cpu:cpu_executable",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "computation_kernel_emitter",
     srcs = ["computation_kernel_emitter.cc"],
     hdrs = ["computation_kernel_emitter.h"],
     deps = [
+        ":computation_kernel_arguments",
         ":kernel_api_ir_builder",
         ":target_machine_features",
-        "//xla:shape_util",
         "//xla:util",
         "//xla/codegen:kernel_definition",
         "//xla/codegen:kernel_emitter",
@@ -667,7 +708,6 @@ cc_library(
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:statusor",
         "@com_google_absl//absl/container:flat_hash_map",
-        "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
@@ -676,7 +716,6 @@ cc_library(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_absl//absl/types:span",
-        "@llvm-project//llvm:Support",
         "@llvm-project//llvm:ir_headers",
     ],
 )
diff --git a/third_party/xla/xla/backends/cpu/codegen/computation_kernel_arguments.cc b/third_party/xla/xla/backends/cpu/codegen/computation_kernel_arguments.cc
new file mode 100644
index 00000000000..611b40e1aaa
--- /dev/null
+++ b/third_party/xla/xla/backends/cpu/codegen/computation_kernel_arguments.cc
@@ -0,0 +1,232 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/cpu/codegen/computation_kernel_arguments.h"
+
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "xla/codegen/kernel_spec.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/hlo/ir/hlo_print_options.h"
+#include "xla/service/buffer_assignment.h"
+#include "xla/service/shaped_slice.h"
+#include "xla/shape.h"
+#include "xla/shape_util.h"
+
+namespace xla::cpu {
+namespace {
+
+struct BufferTable {
+  void Add(const Shape& shape, const BufferAllocation::Slice& slice) {
+    if (seen.insert({shape, slice}).second) {
+      buffers.push_back({slice, shape});
+    }
+  }
+
+  KernelSpec::Buffers buffers;
+  absl::flat_hash_set<const HloInstruction*> seen_instructions;
+  absl::flat_hash_set<std::pair<Shape, BufferAllocation::Slice>> seen;
+};
+
+struct Walk {
+  const BufferAssignment& assignment;
+  BufferTable arguments;
+  BufferTable results;
+  // Every visited instruction occurrence, in visit order.
+  std::vector<const HloInstruction*> instructions;
+  // Allocations of constants defined inside the walk, whose values are part of
+  // the computation text.
+  absl::flat_hash_set<BufferAllocation::Index> walked_constants;
+};
+
+// Adds the slices of all subshapes of `instruction` to `table`.
+absl::Status AddSlices(const HloInstruction* instruction,
+                       const BufferAssignment& assignment, BufferTable& table) {
+  if (!table.seen_instructions.insert(instruction).second) {
+    return absl::OkStatus();
+  }
+  return ShapeUtil::ForEachSubshapeWithStatus(
+      instruction->shape(),
+      [&](const Shape& subshape, const ShapeIndex& index) -> absl::Status {
+        ABSL_ASSIGN_OR_RETURN(BufferAllocation::Slice slice,
+                         assignment.GetUniqueSlice(instruction, index));
+        if (!slice.allocation()->is_thread_local()) {
+          table.Add(subshape, slice);
+        }
+        return absl::OkStatus();
+      });
+}
+
+// Visits instructions in the order HloComputation::Print uses, so that
+// computations with equal text are walked identically.
+absl::Status Visit(const HloComputation& computation, Walk& walk) {
+  for (const HloInstruction* instruction :
+       computation.MakeInstructionPostOrder()) {
+    walk.instructions.push_back(instruction);
+    for (const HloInstruction* operand : instruction->operands()) {
+      ABSL_RETURN_IF_ERROR(AddSlices(operand, walk.assignment, walk.arguments));
+    }
+    // Parameters forward the buffers of their caller.
+    if (instruction->opcode() != HloOpcode::kParameter) {
+      ABSL_RETURN_IF_ERROR(AddSlices(instruction, walk.assignment, walk.results));
+    }
+    if (instruction->opcode() == HloOpcode::kConstant) {
+      ABSL_ASSIGN_OR_RETURN(BufferAllocation::Slice slice,
+                       walk.assignment.GetUniqueTopLevelSlice(instruction));
+      walk.walked_constants.insert(slice.index());
+    }
+    for (const HloComputation* called : instruction->called_computations()) {
+      if (!called->IsFusionComputation()) {
+        ABSL_RETURN_IF_ERROR(Visit(*called, walk));
+      }
+    }
+  }
+  return absl::OkStatus();
+}
+
+// IrEmitter lowers this op to a module-private RNG state global, which a shared
+// kernel would share and race on. Fusion bodies are included.
+bool UsesRngState(const HloComputation& computation) {
+  auto uses = [](const HloComputation* c) {
+    return absl::c_any_of(c->instructions(), [](const HloInstruction* i) {
+      return i->opcode() == HloOpcode::kRngGetAndUpdateState;
+    });
+  };
+  return uses(&computation) ||
+         absl::c_any_of(computation.MakeEmbeddedComputationsList(), uses);
+}
+
+// Returns nullopt if the kernel must not be reused.
+std::optional<std::string> Fingerprint(
+    const HloComputation& computation, const Walk& walk,
+    const absl::flat_hash_map<BufferAllocation::Slice, int64_t>& table_index) {
+  if (UsesRngState(computation)) {
+    return std::nullopt;
+  }
+  // Constants and backend configs are baked into the emitted code.
+  std::string fingerprint =
+      computation.ToString(HloPrintOptions::ModuleFingerprint()
+                               .set_print_only_essential_constants(false)
+                               .set_print_large_constants(true)
+                               .set_print_backend_config(true));
+
+  absl::flat_hash_set<BufferAllocation::Slice> seen_slices;
+  std::vector<BufferAllocation::Slice> distinct_slices;
+  for (const KernelSpec::Buffers* buffers :
+       {&walk.arguments.buffers, &walk.results.buffers}) {
+    for (const ShapedSlice& buffer : *buffers) {
+      const BufferAllocation& allocation = *buffer.slice.allocation();
+      absl::StrAppend(&fingerprint, "\n",
+                      ShapeUtil::HumanStringWithLayout(buffer.shape), " ",
+                      buffer.slice.size(), " ", allocation.size(), " ",
+                      allocation.is_constant(),
+                      allocation.is_parameter_aliased_with_output());
+      if (seen_slices.insert(buffer.slice).second) {
+        distinct_slices.push_back(buffer.slice);
+      }
+    }
+  }
+
+  // Which buffer each instruction uses, so that equal tables are also used
+  // in the same way. Thread-local allocations are numbered by first use, since
+  // the emitter may share one alloca between slices of the same allocation.
+  absl::flat_hash_map<BufferAllocation::Index, int64_t> thread_local_id;
+  for (const HloInstruction* instruction : walk.instructions) {
+    absl::StrAppend(&fingerprint, "\n");
+    absl::Status describable = ShapeUtil::ForEachSubshapeWithStatus(
+        instruction->shape(),
+        [&](const Shape&, const ShapeIndex& shape_index) -> absl::Status {
+          ABSL_ASSIGN_OR_RETURN(
+              BufferAllocation::Slice slice,
+              walk.assignment.GetUniqueSlice(instruction, shape_index));
+          const BufferAllocation& allocation = *slice.allocation();
+          if (auto it = table_index.find(slice); it != table_index.end()) {
+            absl::StrAppend(&fingerprint, " ", it->second);
+            // The literal of a constant from outside the walk is not in the
+            // computation text, and the emitter may bake it into the code.
+            if (allocation.is_constant() &&
+                !walk.walked_constants.contains(allocation.index())) {
+              absl::StrAppend(&fingerprint, ":C", allocation.index());
+            }
+          } else if (allocation.is_thread_local()) {
+            int64_t id =
+                thread_local_id
+                    .try_emplace(allocation.index(), thread_local_id.size())
+                    .first->second;
+            if (seen_slices.insert(slice).second) {
+              distinct_slices.push_back(slice);
+            }
+            absl::StrAppend(&fingerprint, " T", id, ":", slice.offset(), ":",
+                            slice.size(), ":", allocation.size());
+          } else {
+            return absl::UnimplementedError("Buffer outside the kernel table");
+          }
+          return absl::OkStatus();
+        });
+    if (!describable.ok()) {
+      return std::nullopt;
+    }
+  }
+
+  absl::StrAppend(&fingerprint, "\noverlaps:");
+  for (int64_t i = 0; i < distinct_slices.size(); ++i) {
+    for (int64_t j = i + 1; j < distinct_slices.size(); ++j) {
+      if (distinct_slices[i].OverlapsWith(distinct_slices[j])) {
+        absl::StrAppend(&fingerprint, " ", i, ",", j);
+      }
+    }
+  }
+  return fingerprint;
+}
+
+}  // namespace
+
+absl::StatusOr<ComputationKernelArguments> ComputationKernelArguments::Create(
+    const HloComputation& computation, const BufferAssignment& assignment) {
+  Walk walk{assignment};
+  ABSL_RETURN_IF_ERROR(Visit(computation, walk));
+
+  // The emitter stores every argument and result by position, so later
+  // entries win for slices listed more than once.
+  absl::flat_hash_map<BufferAllocation::Slice, int64_t> buffer_table_index;
+  int64_t index = 0;
+  for (const KernelSpec::Buffers* buffers :
+       {&walk.arguments.buffers, &walk.results.buffers}) {
+    for (const ShapedSlice& buffer : *buffers) {
+      buffer_table_index.insert_or_assign(buffer.slice, index++);
+    }
+  }
+
+  std::optional<std::string> fingerprint =
+      Fingerprint(computation, walk, buffer_table_index);
+  return ComputationKernelArguments(
+      std::move(walk.arguments.buffers), std::move(walk.results.buffers),
+      std::move(buffer_table_index), std::move(fingerprint));
+}
+
+}  // namespace xla::cpu
diff --git a/third_party/xla/xla/backends/cpu/codegen/computation_kernel_arguments.h b/third_party/xla/xla/backends/cpu/codegen/computation_kernel_arguments.h
new file mode 100644
index 00000000000..8fedcf4b35c
--- /dev/null
+++ b/third_party/xla/xla/backends/cpu/codegen/computation_kernel_arguments.h
@@ -0,0 +1,83 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_BACKENDS_CPU_CODEGEN_COMPUTATION_KERNEL_ARGUMENTS_H_
+#define XLA_BACKENDS_CPU_CODEGEN_COMPUTATION_KERNEL_ARGUMENTS_H_
+
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <utility>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/statusor.h"
+#include "xla/codegen/kernel_spec.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/service/buffer_assignment.h"
+
+namespace xla::cpu {
+
+// Deterministic argument table for a computation emitted as a single kernel.
+//
+// Operand slices of all instructions become arguments and slices of all
+// non-parameter instructions become results, including nested tuple elements
+// and instructions of nested non-fusion computations. Thread-local slices are
+// skipped. Entries are deduplicated by (shape, slice) in first-visit order.
+//
+// The fingerprint does not use emitters::GetComputationFingerprint, which only
+// describes kernel operands and outputs, not the buffers of intermediate
+// instructions.
+class ComputationKernelArguments {
+ public:
+  // Returns an error if a visited instruction or one of its tuple elements has
+  // no unique slice in `assignment`.
+  static absl::StatusOr<ComputationKernelArguments> Create(
+      const HloComputation& computation, const BufferAssignment& assignment);
+
+  const KernelSpec::Buffers& arguments() const { return arguments_; }
+  const KernelSpec::Buffers& results() const { return results_; }
+
+  // Index of each slice in the buffer table of arguments followed by results.
+  // A slice listed more than once maps to its last position.
+  const absl::flat_hash_map<BufferAllocation::Slice, int64_t>&
+  buffer_table_index() const {
+    return buffer_table_index_;
+  }
+
+  // Equal fingerprints mean a kernel emitted by ComputationKernelEmitter for
+  // one computation is valid for the other's arguments and results. Other
+  // emitters must check what they bake into the code. Unset if the kernel must
+  // not be reused.
+  const std::optional<std::string>& fingerprint() const { return fingerprint_; }
+
+ private:
+  ComputationKernelArguments(
+      KernelSpec::Buffers arguments, KernelSpec::Buffers results,
+      absl::flat_hash_map<BufferAllocation::Slice, int64_t> buffer_table_index,
+      std::optional<std::string> fingerprint)
+      : arguments_(std::move(arguments)),
+        results_(std::move(results)),
+        buffer_table_index_(std::move(buffer_table_index)),
+        fingerprint_(std::move(fingerprint)) {}
+
+  KernelSpec::Buffers arguments_;
+  KernelSpec::Buffers results_;
+  absl::flat_hash_map<BufferAllocation::Slice, int64_t> buffer_table_index_;
+  std::optional<std::string> fingerprint_;
+};
+
+}  // namespace xla::cpu
+
+#endif  // XLA_BACKENDS_CPU_CODEGEN_COMPUTATION_KERNEL_ARGUMENTS_H_
diff --git a/third_party/xla/xla/backends/cpu/codegen/computation_kernel_arguments_test.cc b/third_party/xla/xla/backends/cpu/codegen/computation_kernel_arguments_test.cc
new file mode 100644
index 00000000000..b00aff8a981
--- /dev/null
+++ b/third_party/xla/xla/backends/cpu/codegen/computation_kernel_arguments_test.cc
@@ -0,0 +1,463 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/cpu/codegen/computation_kernel_arguments.h"
+
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status_macros.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_replace.h"
+#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_module.h"
+#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
+#include "xla/service/buffer_assignment.h"
+#include "xla/service/buffer_value.h"
+#include "xla/service/cpu/cpu_executable.h"
+#include "xla/service/logical_buffer.h"
+
+namespace xla::cpu {
+namespace {
+
+using ::testing::Pointwise;
+
+MATCHER(SameShape, "") {
+  return std::get<0>(arg).shape == std::get<1>(arg).shape;
+}
+
+// Two structurally identical computations, call_a and call_b, each running a
+// counter loop. $MUL_A/$MUL_B define the loop multiplier `m` and $ARGS_B the
+// operands of the call to call_b. The copy of the loop counter init mirrors
+// copy insertion; a constant cannot share the loop state buffer.
+constexpr absl::string_view kHloTemplate = R"(
+HloModule m
+
+cond_a {
+  p = (s32[], f32[16]) parameter(0)
+  i = s32[] get-tuple-element(p), index=0
+  n = s32[] constant(5)
+  ROOT lt = pred[] compare(i, n), direction=LT
+}
+
+body_a {
+  p = (s32[], f32[16]) parameter(0)
+  i = s32[] get-tuple-element(p), index=0
+  v = f32[16] get-tuple-element(p), index=1
+  one = s32[] constant(1)
+  i1 = s32[] add(i, one)
+  $MUL_A
+  v1 = f32[16] multiply(v, m)
+  ROOT t = (s32[], f32[16]) tuple(i1, v1)
+}
+
+call_a {
+  p0 = f32[16] parameter(0)
+  p1 = f32[16] parameter(1)
+  s = f32[16] add(p0, p1)
+  zero = s32[] constant(0)
+  zero_copy = s32[] copy(zero)
+  init = (s32[], f32[16]) tuple(zero_copy, s)
+  w = (s32[], f32[16]) while(init), condition=cond_a, body=body_a
+  ROOT r = f32[16] get-tuple-element(w), index=1
+}
+
+cond_b {
+  p = (s32[], f32[16]) parameter(0)
+  i = s32[] get-tuple-element(p), index=0
+  n = s32[] constant(5)
+  ROOT lt = pred[] compare(i, n), direction=LT
+}
+
+body_b {
+  p = (s32[], f32[16]) parameter(0)
+  i = s32[] get-tuple-element(p), index=0
+  v = f32[16] get-tuple-element(p), index=1
+  one = s32[] constant(1)
+  i1 = s32[] add(i, one)
+  $MUL_B
+  v1 = f32[16] multiply(v, m)
+  ROOT t = (s32[], f32[16]) tuple(i1, v1)
+}
+
+call_b {
+  p0 = f32[16] parameter(0)
+  p1 = f32[16] parameter(1)
+  s = f32[16] add(p0, p1)
+  zero = s32[] constant(0)
+  zero_copy = s32[] copy(zero)
+  init = (s32[], f32[16]) tuple(zero_copy, s)
+  w = (s32[], f32[16]) while(init), condition=cond_b, body=body_b
+  ROOT r = f32[16] get-tuple-element(w), index=1
+}
+
+ENTRY e {
+  p0 = f32[16] parameter(0)
+  p1 = f32[16] parameter(1)
+  p2 = f32[16] parameter(2)
+  p3 = f32[16] parameter(3)
+  k0 = f32[16] constant({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16})
+  k1 = f32[16] constant({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16})
+  c0 = f32[16] call($ARGS_A), to_apply=call_a
+  c1 = f32[16] call($ARGS_B), to_apply=call_b
+  ROOT t = (f32[16], f32[16]) tuple(c0, c1)
+}
+)";
+
+constexpr absl::string_view kScalarMul = R"(
+  c = f32[] constant(2.5)
+  m = f32[16] broadcast(c), dimensions={})";
+
+constexpr absl::string_view kLargeMul =
+    "m = f32[16] constant({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, "
+    "16})";
+
+struct Calls {
+  ComputationKernelArguments a;
+  ComputationKernelArguments b;
+};
+
+class ComputationKernelArgumentsTest : public HloHardwareIndependentTestBase {
+ protected:
+  absl::StatusOr<Calls> CreateForCalls(absl::string_view mul_a,
+                                       absl::string_view mul_b,
+                                       absl::string_view args_a = "p0, p1",
+                                       absl::string_view args_b = "p2, p3") {
+    return CreateForModule(
+        absl::StrReplaceAll(kHloTemplate, {{"$MUL_A", mul_a},
+                                           {"$MUL_B", mul_b},
+                                           {"$ARGS_A", args_a},
+                                           {"$ARGS_B", args_b}}));
+  }
+
+  // Assigns buffers for `hlo` and creates arguments for call_a and call_b.
+  absl::StatusOr<Calls> CreateForModule(absl::string_view hlo) {
+    ABSL_ASSIGN_OR_RETURN(module_, ParseAndReturnVerifiedModule(hlo));
+    // As in CpuCompiler::CreateBufferAssignment.
+    BufferAssigner::Options options;
+    options.allocate_buffers_for_constants = true;
+    ABSL_ASSIGN_OR_RETURN(
+        assignment_,
+        BufferAssigner::Run(
+            module_.get(),
+            std::make_unique<DependencyHloOrdering>(module_.get()),
+            [](const BufferValue& buffer) {
+              return CpuExecutable::ShapeSizeBytes(buffer.shape());
+            },
+            &alias_info_, [](LogicalBuffer::Color) { return /*alignment=*/1; },
+            std::move(options)));
+    ABSL_ASSIGN_OR_RETURN(ComputationKernelArguments a, Create("call_a"));
+    ABSL_ASSIGN_OR_RETURN(ComputationKernelArguments b, Create("call_b"));
+    return Calls{std::move(a), std::move(b)};
+  }
+
+  absl::StatusOr<ComputationKernelArguments> Create(absl::string_view name) {
+    return ComputationKernelArguments::Create(
+        *module_->GetComputationWithName(name), *assignment_);
+  }
+
+ private:
+  AliasInfo alias_info_;
+  std::unique_ptr<HloModule> module_;
+  std::unique_ptr<BufferAssignment> assignment_;
+};
+
+TEST_F(ComputationKernelArgumentsTest, IdenticalCallsShareFingerprint) {
+  ASSERT_OK_AND_ASSIGN(Calls calls, CreateForCalls(kScalarMul, kScalarMul));
+
+  ASSERT_TRUE(calls.a.fingerprint().has_value());
+  EXPECT_EQ(calls.a.fingerprint(), calls.b.fingerprint());
+  EXPECT_NE(calls.a.arguments(), calls.b.arguments());
+  EXPECT_THAT(calls.a.arguments(), Pointwise(SameShape(), calls.b.arguments()));
+  EXPECT_THAT(calls.a.results(), Pointwise(SameShape(), calls.b.results()));
+
+  ASSERT_OK_AND_ASSIGN(ComputationKernelArguments again, Create("call_a"));
+  EXPECT_EQ(again.arguments(), calls.a.arguments());
+  EXPECT_EQ(again.results(), calls.a.results());
+  EXPECT_EQ(again.fingerprint(), calls.a.fingerprint());
+}
+
+TEST_F(ComputationKernelArgumentsTest, DifferentCallsNoShare) {
+  const struct {
+    absl::string_view name;
+    std::string mul_a;
+    std::string mul_b;
+    absl::string_view args_b = "p2, p3";
+  } cases[] = {
+      {"float constant", std::string(kScalarMul),
+       absl::StrReplaceAll(kScalarMul, {{"2.5", "3.5"}})},
+      {"large constant", std::string(kLargeMul),
+       absl::StrReplaceAll(kLargeMul, {{"16}", "17}"}})},
+      {"operand sharing", std::string(kScalarMul), std::string(kScalarMul),
+       "p0, p0"},
+  };
+  for (const auto& c : cases) {
+    SCOPED_TRACE(c.name);
+    ASSERT_OK_AND_ASSIGN(Calls calls,
+                         CreateForCalls(c.mul_a, c.mul_b, "p0, p1", c.args_b));
+    ASSERT_TRUE(calls.a.fingerprint().has_value());
+    ASSERT_TRUE(calls.b.fingerprint().has_value());
+    EXPECT_NE(calls.a.fingerprint(), calls.b.fingerprint());
+  }
+}
+
+// k0 and k1 are equal but distinct constant allocations, whose literals the
+// emitter may bake into the code.
+TEST_F(ComputationKernelArgumentsTest, ExternalConstantOperand) {
+  ASSERT_OK_AND_ASSIGN(
+      Calls same, CreateForCalls(kScalarMul, kScalarMul, "k0, p1", "k0, p3"));
+  ASSERT_TRUE(same.a.fingerprint().has_value());
+  EXPECT_EQ(same.a.fingerprint(), same.b.fingerprint());
+
+  ASSERT_OK_AND_ASSIGN(Calls different, CreateForCalls(kScalarMul, kScalarMul,
+                                                       "k0, p1", "k1, p3"));
+  ASSERT_TRUE(different.a.fingerprint().has_value());
+  ASSERT_TRUE(different.b.fingerprint().has_value());
+  EXPECT_NE(different.a.fingerprint(), different.b.fingerprint());
+}
+
+// The sum may reuse the buffer of the temporary operand of call_b, but not the
+// read-only entry parameter of call_a.
+TEST_F(ComputationKernelArgumentsTest, OperandReusePatternNoShare) {
+  constexpr absl::string_view kHlo = R"(
+HloModule m
+
+call_a {
+  p0 = f32[16] parameter(0)
+  p1 = f32[16] parameter(1)
+  ROOT s = f32[16] add(p0, p1)
+}
+
+call_b {
+  p0 = f32[16] parameter(0)
+  p1 = f32[16] parameter(1)
+  ROOT s = f32[16] add(p0, p1)
+}
+
+ENTRY e {
+  p0 = f32[16] parameter(0)
+  p1 = f32[16] parameter(1)
+  p0_copy = f32[16] copy(p0)
+  c0 = f32[16] call(p0, p1), to_apply=call_a
+  c1 = f32[16] call(p0_copy, p1), to_apply=call_b
+  ROOT t = (f32[16], f32[16]) tuple(c0, c1)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(Calls calls, CreateForModule(kHlo));
+
+  // The first argument is p0 and the only result is s.
+  ASSERT_EQ(calls.a.results().size(), 1);
+  ASSERT_EQ(calls.b.results().size(), 1);
+  EXPECT_NE(calls.a.arguments()[0].slice, calls.a.results()[0].slice);
+  EXPECT_EQ(calls.b.arguments()[0].slice, calls.b.results()[0].slice);
+  ASSERT_TRUE(calls.a.fingerprint().has_value());
+  ASSERT_TRUE(calls.b.fingerprint().has_value());
+  EXPECT_NE(calls.a.fingerprint(), calls.b.fingerprint());
+}
+
+// Values of the reducers are thread-local, outside the kernel table.
+TEST_F(ComputationKernelArgumentsTest, ThreadLocalBuffersShareFingerprint) {
+  constexpr absl::string_view kHlo = R"(
+HloModule m
+
+add_a {
+  x = f32[] parameter(0)
+  y = f32[] parameter(1)
+  ROOT s = f32[] add(x, y)
+}
+
+call_a {
+  p0 = f32[16] parameter(0)
+  zero = f32[] constant(0)
+  ROOT r = f32[] reduce(p0, zero), dimensions={0}, to_apply=add_a
+}
+
+add_b {
+  x = f32[] parameter(0)
+  y = f32[] parameter(1)
+  ROOT s = f32[] add(x, y)
+}
+
+call_b {
+  p0 = f32[16] parameter(0)
+  zero = f32[] constant(0)
+  ROOT r = f32[] reduce(p0, zero), dimensions={0}, to_apply=add_b
+}
+
+ENTRY e {
+  p0 = f32[16] parameter(0)
+  p1 = f32[16] parameter(1)
+  c0 = f32[] call(p0), to_apply=call_a
+  c1 = f32[] call(p1), to_apply=call_b
+  ROOT t = (f32[], f32[]) tuple(c0, c1)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(Calls calls, CreateForModule(kHlo));
+  ASSERT_TRUE(calls.a.fingerprint().has_value());
+  EXPECT_EQ(calls.a.fingerprint(), calls.b.fingerprint());
+}
+
+// The emitted kernel owns a private RNG state, which callers must not share.
+TEST_F(ComputationKernelArgumentsTest, RngStateNotReusable) {
+  constexpr absl::string_view kHlo = R"(
+HloModule m
+
+rng_a {
+  ROOT s = u64[2] rng-get-and-update-state(), delta=16
+}
+
+call_a {
+  ROOT c = u64[2] call(), to_apply=rng_a
+}
+
+rng_b {
+  ROOT s = u64[2] rng-get-and-update-state(), delta=16
+}
+
+call_b {
+  ROOT c = u64[2] call(), to_apply=rng_b
+}
+
+ENTRY e {
+  c0 = u64[2] call(), to_apply=call_a
+  c1 = u64[2] call(), to_apply=call_b
+  ROOT t = (u64[2], u64[2]) tuple(c0, c1)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(Calls calls, CreateForModule(kHlo));
+  EXPECT_EQ(calls.a.fingerprint(), std::nullopt);
+  EXPECT_EQ(calls.b.fingerprint(), std::nullopt);
+}
+
+// The RNG state op is in a reducer of a fusion, which the walk does not enter.
+TEST_F(ComputationKernelArgumentsTest, RngStateInFusionNotReusable) {
+  constexpr absl::string_view kHlo = R"(
+HloModule m
+
+add_a {
+  x = f32[] parameter(0)
+  y = f32[] parameter(1)
+  s = u64[2] rng-get-and-update-state(), delta=16
+  ROOT r = f32[] add(x, y)
+}
+
+fused_a {
+  p0 = f32[16] parameter(0)
+  zero = f32[] constant(0)
+  ROOT r = f32[] reduce(p0, zero), dimensions={0}, to_apply=add_a
+}
+
+call_a {
+  p0 = f32[16] parameter(0)
+  ROOT f = f32[] fusion(p0), kind=kLoop, calls=fused_a
+}
+
+add_b {
+  x = f32[] parameter(0)
+  y = f32[] parameter(1)
+  s = u64[2] rng-get-and-update-state(), delta=16
+  ROOT r = f32[] add(x, y)
+}
+
+fused_b {
+  p0 = f32[16] parameter(0)
+  zero = f32[] constant(0)
+  ROOT r = f32[] reduce(p0, zero), dimensions={0}, to_apply=add_b
+}
+
+call_b {
+  p0 = f32[16] parameter(0)
+  ROOT f = f32[] fusion(p0), kind=kLoop, calls=fused_b
+}
+
+ENTRY e {
+  p0 = f32[16] parameter(0)
+  p1 = f32[16] parameter(1)
+  c0 = f32[] call(p0), to_apply=call_a
+  c1 = f32[] call(p1), to_apply=call_b
+  ROOT t = (f32[], f32[]) tuple(c0, c1)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(Calls calls, CreateForModule(kHlo));
+  EXPECT_EQ(calls.a.fingerprint(), std::nullopt);
+  EXPECT_EQ(calls.b.fingerprint(), std::nullopt);
+}
+
+// The inner tuple {0} of p0 is not the top-level buffer of any instruction.
+TEST_F(ComputationKernelArgumentsTest, NestedTupleShareFingerprint) {
+  constexpr absl::string_view kHlo = R"(
+HloModule m
+
+call_a {
+  p0 = ((s32[], f32[16]), f32[16]) parameter(0)
+  ROOT g = f32[16] get-tuple-element(p0), index=1
+}
+
+call_b {
+  p0 = ((s32[], f32[16]), f32[16]) parameter(0)
+  ROOT g = f32[16] get-tuple-element(p0), index=1
+}
+
+ENTRY e {
+  p0 = ((s32[], f32[16]), f32[16]) parameter(0)
+  p1 = ((s32[], f32[16]), f32[16]) parameter(1)
+  c0 = f32[16] call(p0), to_apply=call_a
+  c1 = f32[16] call(p1), to_apply=call_b
+  ROOT t = (f32[16], f32[16]) tuple(c0, c1)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(Calls calls, CreateForModule(kHlo));
+  ASSERT_TRUE(calls.a.fingerprint().has_value());
+  EXPECT_EQ(calls.a.fingerprint(), calls.b.fingerprint());
+}
+
+// The unused parameter p1 has a buffer outside the kernel table.
+TEST_F(ComputationKernelArgumentsTest, UndescribedBufferNotReusable) {
+  constexpr absl::string_view kHlo = R"(
+HloModule m
+
+call_a {
+  p0 = f32[16] parameter(0)
+  p1 = f32[16] parameter(1)
+  ROOT n = f32[16] negate(p0)
+}
+
+call_b {
+  p0 = f32[16] parameter(0)
+  p1 = f32[16] parameter(1)
+  ROOT n = f32[16] negate(p0)
+}
+
+ENTRY e {
+  p0 = f32[16] parameter(0)
+  p1 = f32[16] parameter(1)
+  c0 = f32[16] call(p0, p1), to_apply=call_a
+  c1 = f32[16] call(p1, p0), to_apply=call_b
+  ROOT t = (f32[16], f32[16]) tuple(c0, c1)
+}
+)";
+  ASSERT_OK_AND_ASSIGN(Calls calls, CreateForModule(kHlo));
+  EXPECT_EQ(calls.a.fingerprint(), std::nullopt);
+  EXPECT_EQ(calls.b.fingerprint(), std::nullopt);
+}
+
+}  // namespace
+}  // namespace xla::cpu
diff --git a/third_party/xla/xla/backends/cpu/codegen/computation_kernel_emitter.cc b/third_party/xla/xla/backends/cpu/codegen/computation_kernel_emitter.cc
index c1ec05cb192..22fd5eb7ac1 100644
--- a/third_party/xla/xla/backends/cpu/codegen/computation_kernel_emitter.cc
+++ b/third_party/xla/xla/backends/cpu/codegen/computation_kernel_emitter.cc
@@ -22,7 +22,6 @@ limitations under the License.
 #include <vector>
 
 #include "absl/container/flat_hash_map.h"
-#include "absl/container/flat_hash_set.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
@@ -31,7 +30,6 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
-#include "llvm/ADT/STLExtras.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Function.h"
@@ -41,6 +39,7 @@ limitations under the License.
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Value.h"
 #include "llvm/IR/Verifier.h"
+#include "xla/backends/cpu/codegen/computation_kernel_arguments.h"
 #include "xla/backends/cpu/codegen/kernel_api_ir_builder.h"
 #include "xla/backends/cpu/codegen/target_machine_features.h"
 #include "xla/codegen/kernel_definition.h"
@@ -49,15 +48,12 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_module.h"
-#include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/runtime/work_group.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/cpu/backend_config.pb.h"
 #include "xla/service/cpu/ir_emitter.h"
 #include "xla/service/llvm_ir/ir_array.h"
 #include "xla/service/llvm_ir/llvm_util.h"
-#include "xla/shape.h"
-#include "xla/shape_util.h"
 #include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/statusor.h"
 #include "xla/util.h"
@@ -66,72 +62,26 @@ namespace xla::cpu {
 
 namespace {
 
-// Get the slices for the given instruction, expanding the elements if a tuple.
-absl::Status GetInstructionSlices(
-    const HloInstruction* instruction,
-    const BufferAssignment* buffer_assignment,
-    absl::flat_hash_set<KernelApiIrBuilder::KernelParameter>& parameters) {
-  const Shape& shape = instruction->shape();
-  ABSL_ASSIGN_OR_RETURN(BufferAllocation::Slice slice,
-                   buffer_assignment->GetUniqueTopLevelSlice(instruction));
-  if (slice.allocation()->is_thread_local()) {
-    return absl::OkStatus();
+std::vector<KernelApiIrBuilder::KernelParameter> ToKernelParameters(
+    const KernelSpec::Buffers& buffers) {
+  std::vector<KernelApiIrBuilder::KernelParameter> parameters;
+  parameters.reserve(buffers.size());
+  for (const auto& [slice, shape] : buffers) {
+    parameters.push_back({shape, slice});
   }
-  parameters.insert({shape, std::move(slice)});
-
-  if (shape.IsTuple()) {
-    for (const auto& [leaf_index, leaf_shape] :
-         ShapeUtil::GetLeafShapes(shape)) {
-      ABSL_ASSIGN_OR_RETURN(
-          BufferAllocation::Slice leaf_slice,
-          buffer_assignment->GetUniqueSlice(instruction, {leaf_index}));
-      parameters.insert({leaf_shape, std::move(leaf_slice)});
-    }
-  }
-  return absl::OkStatus();
-}
-
-absl::Status GetAllSlices(
-    const HloComputation* computation,
-    const BufferAssignment* buffer_assignment,
-    absl::flat_hash_set<KernelApiIrBuilder::KernelParameter>& arguments,
-    absl::flat_hash_set<KernelApiIrBuilder::KernelParameter>& results) {
-  for (const HloInstruction* instruction : computation->instructions()) {
-    for (const HloInstruction* operand : instruction->operands()) {
-      ABSL_RETURN_IF_ERROR(
-          GetInstructionSlices(operand, buffer_assignment, arguments));
-    }
-
-    // Parameters just forward the results
-    // TODO(willfroom): Is there a method somewhere to check if an instruction
-    // just forwards the buffer? (e.g get-tuple-arg)
-    if (instruction->opcode() != HloOpcode::kParameter) {
-      ABSL_RETURN_IF_ERROR(
-          GetInstructionSlices(instruction, buffer_assignment, results));
-    }
-
-    for (const HloComputation* nested_computation :
-         instruction->called_computations()) {
-      if (nested_computation->IsFusionComputation()) {
-        continue;
-      }
-
-      ABSL_RETURN_IF_ERROR(GetAllSlices(nested_computation, buffer_assignment,
-                                   arguments, results));
-    }
-  }
-
-  return absl::OkStatus();
+  return parameters;
 }
 
 }  // namespace
 
 ComputationKernelEmitter::ComputationKernelEmitter(
     const HloInstruction* instr, const BufferAssignment* buffer_assignment,
-    const TargetMachineFeatures* target_machine)
+    const TargetMachineFeatures* target_machine,
+    ComputationKernelArguments arguments)
     : instr_(instr),
       buffer_assignment_(buffer_assignment),
-      target_machine_(target_machine) {}
+      target_machine_(target_machine),
+      arguments_(std::move(arguments)) {}
 
 absl::StatusOr<ComputationKernelEmitter::KernelDefinition>
 ComputationKernelEmitter::EmitKernelDefinition() {
@@ -144,11 +94,6 @@ ComputationKernelEmitter::EmitKernelDefinition() {
     return Internal("HloModule is null");
   }
 
-  absl::flat_hash_set<KernelApiIrBuilder::KernelParameter> arguments;
-  absl::flat_hash_set<KernelApiIrBuilder::KernelParameter> results;
-  ABSL_RETURN_IF_ERROR(
-      GetAllSlices(instr_->to_apply(), buffer_assignment_, arguments, results));
-
   // As the computation is a series of operations, buffers are not disjoint.
   KernelApiIrBuilder kernel_api_ir_builder(
       *ctx,
@@ -162,14 +107,12 @@ ComputationKernelEmitter::EmitKernelDefinition() {
   ABSL_ASSIGN_OR_RETURN(std::string kernel_name,
                    kernel_api_ir_builder.GetKernelName(instr_, "_kernel"));
 
-  ABSL_ASSIGN_OR_RETURN(KernelApiIrBuilder::KernelPrototype kernel_prototype,
-                   kernel_api_ir_builder.EmitKernelPrototype(
-                       *llvm_module, kernel_name,
-                       std::vector<KernelApiIrBuilder::KernelParameter>(
-                           arguments.begin(), arguments.end()),
-                       std::vector<KernelApiIrBuilder::KernelParameter>(
-                           results.begin(), results.end()),
-                       BuildModuleMemoryRegionName(name(), instr_)));
+  ABSL_ASSIGN_OR_RETURN(
+      KernelApiIrBuilder::KernelPrototype kernel_prototype,
+      kernel_api_ir_builder.EmitKernelPrototype(
+          *llvm_module, kernel_name, ToKernelParameters(arguments_.arguments()),
+          ToKernelParameters(arguments_.results()),
+          BuildModuleMemoryRegionName(name(), instr_)));
 
   llvm::IRBuilder<> ir_builder(*ctx);
   ir_builder.SetInsertPoint(
@@ -180,32 +123,23 @@ ComputationKernelEmitter::EmitKernelDefinition() {
   llvm::Value* buffer_table = ir_builder.CreateAlloca(
       ir_builder.getPtrTy(), alloca_size, "buffer_table");
 
-  absl::flat_hash_map<BufferAllocation::Slice, int64_t>
-      slice_to_buffer_table_index;
-
   int64_t buffer_table_index = 0;
-  for (const auto& [array, arg] : llvm::zip(
-           kernel_prototype.arguments, kernel_prototype.argument_buffers)) {
-    int64_t index = buffer_table_index++;
-    slice_to_buffer_table_index[arg.slice] = index;
+  for (const auto& array : kernel_prototype.arguments) {
     llvm::Value* buffer_table_ptr = llvm_ir::EmitBufferIndexingGEP(
-        buffer_table, ir_builder.getPtrTy(), index, &ir_builder);
+        buffer_table, ir_builder.getPtrTy(), buffer_table_index++, &ir_builder);
     ir_builder.CreateStore(array.GetBasePointer(), buffer_table_ptr);
   }
-  for (const auto& [array, result] :
-       llvm::zip(kernel_prototype.results, kernel_prototype.result_buffers)) {
-    int64_t index = buffer_table_index++;
-    slice_to_buffer_table_index[result.slice] = index;
+  for (const auto& array : kernel_prototype.results) {
     llvm::Value* buffer_table_ptr = llvm_ir::EmitBufferIndexingGEP(
-        buffer_table, ir_builder.getPtrTy(), index, &ir_builder);
+        buffer_table, ir_builder.getPtrTy(), buffer_table_index++, &ir_builder);
     ir_builder.CreateStore(array.GetBasePointer(), buffer_table_ptr);
   }
 
   ABSL_ASSIGN_OR_RETURN(
       llvm::Function * computation_function,
-      EmitNestedComputation(
-          kernel_prototype.function, kernel_prototype.return_block, ir_builder,
-          *llvm_module, std::move(slice_to_buffer_table_index)));
+      EmitNestedComputation(kernel_prototype.function,
+                            kernel_prototype.return_block, ir_builder,
+                            *llvm_module, arguments_.buffer_table_index()));
 
   ir_builder.SetInsertPoint(
       kernel_prototype.function->getEntryBlock().getTerminator());
diff --git a/third_party/xla/xla/backends/cpu/codegen/computation_kernel_emitter.h b/third_party/xla/xla/backends/cpu/codegen/computation_kernel_emitter.h
index 95c538ad261..5f4807d0ca8 100644
--- a/third_party/xla/xla/backends/cpu/codegen/computation_kernel_emitter.h
+++ b/third_party/xla/xla/backends/cpu/codegen/computation_kernel_emitter.h
@@ -25,6 +25,7 @@ limitations under the License.
 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Module.h"
+#include "xla/backends/cpu/codegen/computation_kernel_arguments.h"
 #include "xla/backends/cpu/codegen/target_machine_features.h"
 #include "xla/codegen/kernel_emitter.h"
 #include "xla/codegen/llvm_kernel_source.h"
@@ -46,9 +47,12 @@ namespace xla::cpu {
 // allocations for small buffers.
 class ComputationKernelEmitter final : public KernelEmitter<LlvmKernelSource> {
  public:
+  // `arguments` must be created from `instr->to_apply()` and
+  // `buffer_assignment`.
   ComputationKernelEmitter(const HloInstruction* instr,
                            const BufferAssignment* buffer_assignment,
-                           const TargetMachineFeatures* target_machine);
+                           const TargetMachineFeatures* target_machine,
+                           ComputationKernelArguments arguments);
 
   absl::string_view name() const final { return "computation_kernel_emitter"; }
   absl::StatusOr<KernelDefinition> EmitKernelDefinition() final;
@@ -65,6 +69,7 @@ class ComputationKernelEmitter final : public KernelEmitter<LlvmKernelSource> {
 
   const BufferAssignment* buffer_assignment_;
   const TargetMachineFeatures* target_machine_;
+  ComputationKernelArguments arguments_;
 };
 
 }  // namespace xla::cpu
diff --git a/third_party/xla/xla/backends/cpu/testlib/BUILD b/third_party/xla/xla/backends/cpu/testlib/BUILD
index 92ed72a1470..7136a71b13e 100644
--- a/third_party/xla/xla/backends/cpu/testlib/BUILD
+++ b/third_party/xla/xla/backends/cpu/testlib/BUILD
@@ -163,11 +163,13 @@ tsl_pybind_extension(
         ":mlir_kernel_emitter",
         # placeholder for index annotation deps  # buildcleaner: keep
         "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
         "@llvm-project//mlir:IR",
         "@nanobind",
         "@rules_python//python/cc:current_py_cc_headers",  # buildcleaner: keep
+        "//xla/backends/cpu/codegen:computation_kernel_arguments",
         "//xla/backends/cpu/codegen:computation_kernel_emitter",
         "//xla/backends/cpu/codegen:fusion_compiler",
         "//xla/backends/cpu/codegen:fusion_emitter",
diff --git a/third_party/xla/xla/backends/cpu/testlib/kernel_runner_extension.cc b/third_party/xla/xla/backends/cpu/testlib/kernel_runner_extension.cc
index 5d433df4bad..d58071e8dcf 100644
--- a/third_party/xla/xla/backends/cpu/testlib/kernel_runner_extension.cc
+++ b/third_party/xla/xla/backends/cpu/testlib/kernel_runner_extension.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include <utility>
 
 #include "absl/log/check.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 #include "mlir/IR/MLIRContext.h"
@@ -29,6 +30,7 @@ limitations under the License.
 #include "nanobind/stl/tuple.h"  // IWYU pragma: keep
 #include "nanobind/stl/unique_ptr.h"  // IWYU pragma: keep
 #include "nanobind/stl/vector.h"  // IWYU pragma: keep
+#include "xla/backends/cpu/codegen/computation_kernel_arguments.h"
 #include "xla/backends/cpu/codegen/computation_kernel_emitter.h"
 #include "xla/backends/cpu/codegen/dot/dot_kernel_emitter.h"
 #include "xla/backends/cpu/codegen/elemental/concatenate_kernel_emitter.h"
@@ -171,10 +173,23 @@ NB_MODULE(_extension, kernel_runner_module) {
 
   nb::class_<ComputationKernelEmitter, KernelEmitter<LlvmKernelSource>>(
       kernel_runner_module, "ComputationKernelEmitter")
-      .def(nb::init<const HloInstruction*, const BufferAssignment*,
-                    const TargetMachineFeatures*>(),
-           nb::keep_alive<1, 2>(), nb::keep_alive<1, 3>(),
-           nb::keep_alive<1, 4>());
+      .def(
+          "__init__",
+          [](ComputationKernelEmitter* self, const HloInstruction* instruction,
+             const BufferAssignment* buffer_assignment,
+             const TargetMachineFeatures* target_machine) {
+            absl::StatusOr<ComputationKernelArguments> arguments =
+                ComputationKernelArguments::Create(*instruction->to_apply(),
+                                                   *buffer_assignment);
+            if (!arguments.ok()) {
+              throw std::runtime_error(arguments.status().ToString());
+            }
+            new (self)
+                ComputationKernelEmitter(instruction, buffer_assignment,
+                                         target_machine, *std::move(arguments));
+          },
+          nb::keep_alive<1, 2>(), nb::keep_alive<1, 3>(),
+          nb::keep_alive<1, 4>());
 
   nb::class_<CpuScatterFusion, KernelEmitter<MlirKernelSource>>(
       kernel_runner_module, "ScatterKernelEmitter")
diff --git a/third_party/xla/xla/service/cpu/BUILD b/third_party/xla/xla/service/cpu/BUILD
index fbe5c41e36c..2900078058b 100644
--- a/third_party/xla/xla/service/cpu/BUILD
+++ b/third_party/xla/xla/service/cpu/BUILD
@@ -810,6 +810,7 @@ cc_library(
         "//xla/backends/cpu:onednn_support",
         "//xla/backends/cpu:ynn_emitter",
         "//xla/backends/cpu:ynn_support",
+        "//xla/backends/cpu/codegen:computation_kernel_arguments",
         "//xla/backends/cpu/codegen:computation_kernel_emitter",
         "//xla/backends/cpu/codegen:fusion_compiler",
         "//xla/backends/cpu/codegen:fusion_emitter",
@@ -817,7 +818,6 @@ cc_library(
         "//xla/backends/cpu/codegen:target_machine_features",
         "//xla/backends/cpu/codegen/dot:dot_kernel_emitter",
         "//xla/backends/cpu/codegen/elemental:concatenate_kernel_emitter",
-        "//xla/backends/cpu/codegen/emitters:cpu_fusion_emitters",
         "//xla/backends/cpu/runtime:all_gather_thunk",
         "//xla/backends/cpu/runtime:all_reduce_thunk",
         "//xla/backends/cpu/runtime:all_to_all_thunk",
@@ -851,7 +851,6 @@ cc_library(
         "//xla/codegen/emitters:computation_fingerprint",
         "//xla/codegen/emitters:kernel_api_builder",
         "//xla/codegen/emitters:kernel_arguments",
-        "//xla/hlo/analysis:symbolic_map",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/utils:sort_utils",
         "//xla/runtime:resource_use",
@@ -878,7 +877,6 @@ cc_library(
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/types:span",
-        "@llvm-project//llvm:JITLink",
         "@llvm-project//llvm:OrcJIT",
         "@llvm-project//llvm:Support",
         "@llvm-project//llvm:ir_headers",
@@ -1877,7 +1875,9 @@ xla_cc_test(
         "//xla/tsl/platform:statusor",
         "//xla/tsl/platform:test",
         "@com_google_absl//absl/base:nullability",
+        "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/synchronization",
         "@llvm-project//llvm:Core",
         "@llvm-project//llvm:Support",
     ],
diff --git a/third_party/xla/xla/service/cpu/cpu_compiler_internals_test.cc b/third_party/xla/xla/service/cpu/cpu_compiler_internals_test.cc
index f7bbeebe76e..8c96ec609ca 100644
--- a/third_party/xla/xla/service/cpu/cpu_compiler_internals_test.cc
+++ b/third_party/xla/xla/service/cpu/cpu_compiler_internals_test.cc
@@ -19,16 +19,22 @@ limitations under the License.
 #include <optional>
 #include <string>
 #include <utility>
+#include <vector>
 
 #include "absl/base/nullability.h"
+#include "absl/strings/match.h"
+#include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
 #include "llvm/IR/Constants.h"
+#include "llvm/IR/Function.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/Casting.h"
 #include "xla/backends/cpu/codegen/emitters/cpu_fusion_emitter_config.h"
 #include "xla/comparison_util.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"
@@ -214,6 +220,85 @@ TEST_F(CpuCompilerInternalsTest, JustOneDylibWithThunks) {
                          << "one dylib is allowed.";
 }
 
+TEST_F(CpuCompilerInternalsTest, IdenticalSmallWhileLoopsShareKernel) {
+  static constexpr absl::string_view kHlo = R"(
+    body {
+      state = (s32[], f32[4]) parameter(0)
+      i = s32[] get-tuple-element(state), index=0
+      x = f32[4] get-tuple-element(state), index=1
+      one = s32[] constant(1)
+      next_i = s32[] add(i, one)
+      c = f32[] constant(2)
+      cb = f32[4] broadcast(c), dimensions={}
+      sum = f32[4] add(x, cb)
+      next_x = f32[4] multiply(sum, cb)
+      ROOT next_state = (s32[], f32[4]) tuple(next_i, next_x)
+    }
+
+    cond {
+      state = (s32[], f32[4]) parameter(0)
+      i = s32[] get-tuple-element(state), index=0
+      limit = s32[] constant(5)
+      ROOT lt = pred[] compare(i, limit), direction=LT
+    }
+
+    ENTRY main {
+      p0 = f32[4] parameter(0)
+      p1 = f32[4] parameter(1)
+      zero = s32[] constant(0)
+      init0 = (s32[], f32[4]) tuple(zero, p0)
+      init1 = (s32[], f32[4]) tuple(zero, p1)
+      while0 = (s32[], f32[4]) while(init0), condition=cond, body=body
+      while1 = (s32[], f32[4]) while(init1), condition=cond, body=body
+      r0 = f32[4] get-tuple-element(while0), index=1
+      r1 = f32[4] get-tuple-element(while1), index=1
+      ROOT r = (f32[4], f32[4]) tuple(r0, r1)
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> hlo_module,
+                       ParseAndReturnVerifiedModule(kHlo));
+
+  CpuCompiler compiler;
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<HloModule> optimized_module,
+      compiler.RunHloPasses(std::move(hlo_module), /*stream_exec=*/nullptr,
+                            /*options=*/{}));
+
+  std::vector<std::string> kernel_suffixes;
+  for (const HloComputation* computation : optimized_module->computations()) {
+    for (const HloInstruction* instruction : computation->instructions()) {
+      if (instruction->opcode() == HloOpcode::kCall &&
+          instruction->get_frontend_attribute("xla_cpu_small_call") == "true") {
+        kernel_suffixes.push_back(absl::StrCat(instruction->name(), "_kernel"));
+      }
+    }
+  }
+  ASSERT_EQ(kernel_suffixes.size(), 2) << optimized_module->ToString();
+
+  absl::Mutex mu;
+  int64_t num_kernels = 0;
+  auto pre_opt_hook = [&](const llvm::Module& llvm_module) {
+    absl::MutexLock lock(mu);
+    for (const llvm::Function& function : llvm_module) {
+      if (function.isDeclaration()) {
+        continue;
+      }
+      for (const std::string& suffix : kernel_suffixes) {
+        if (absl::EndsWith(function.getName(), suffix)) {
+          ++num_kernels;
+        }
+      }
+    }
+  };
+
+  compiler.SetPreOptimizationHook(pre_opt_hook);
+  ASSERT_OK(compiler.RunBackend(std::move(optimized_module),
+                                /*stream_exec=*/nullptr, /*options=*/{}));
+  compiler.RemovePreOptimizationHook();
+
+  EXPECT_EQ(num_kernels, 1);
+}
+
 TEST_F(CpuCompilerInternalsTest, DisablePlatformDependentMathUnified) {
   HloModuleConfig config;
   // Test default: platform dependent math should be enabled.
diff --git a/third_party/xla/xla/service/cpu/tests/cpu_while_test.cc b/third_party/xla/xla/service/cpu/tests/cpu_while_test.cc
index 0ec01ccadf8..1ffe2320c5d 100644
--- a/third_party/xla/xla/service/cpu/tests/cpu_while_test.cc
+++ b/third_party/xla/xla/service/cpu/tests/cpu_while_test.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include <memory>
 #include <string>
 #include <utility>
+#include <vector>
 
 #include <gtest/gtest.h>
 #include "absl/strings/string_view.h"
@@ -327,6 +328,59 @@ TEST_F(CpuWhileTest, FusedDynamicUpdateSliceInSmallWhileLoop) {
       result);
 }
 
+// The two loops share one kernel, which must use each call's own buffers.
+TEST_F(CpuWhileTest, IdenticalSmallWhileLoopsSharedKernelCorrect) {
+  constexpr absl::string_view hlo_text = R"(
+    HloModule IdenticalSmallWhileLoops
+
+    body {
+      state = (s32[], f32[4]) parameter(0)
+      i = s32[] get-tuple-element(state), index=0
+      x = f32[4] get-tuple-element(state), index=1
+      one = s32[] constant(1)
+      next_i = s32[] add(i, one)
+      c = f32[] constant(2)
+      cb = f32[4] broadcast(c), dimensions={}
+      sum = f32[4] add(x, cb)
+      next_x = f32[4] multiply(sum, cb)
+      ROOT next_state = (s32[], f32[4]) tuple(next_i, next_x)
+    }
+
+    cond {
+      state = (s32[], f32[4]) parameter(0)
+      i = s32[] get-tuple-element(state), index=0
+      limit = s32[] constant(5)
+      ROOT lt = pred[] compare(i, limit), direction=LT
+    }
+
+    ENTRY main {
+      p0 = f32[4] parameter(0)
+      p1 = f32[4] parameter(1)
+      zero = s32[] constant(0)
+      init0 = (s32[], f32[4]) tuple(zero, p0)
+      init1 = (s32[], f32[4]) tuple(zero, p1)
+      while0 = (s32[], f32[4]) while(init0), condition=cond, body=body
+      while1 = (s32[], f32[4]) while(init1), condition=cond, body=body
+      r0 = f32[4] get-tuple-element(while0), index=1
+      r1 = f32[4] get-tuple-element(while1), index=1
+      ROOT r = (f32[4], f32[4]) tuple(r0, r1)
+    }
+  )";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(hlo_text));
+  Literal p0 = LiteralUtil::CreateR1<float>({0, 1, 2, 3});
+  Literal p1 = LiteralUtil::CreateR1<float>({10, 20, 30, 40});
+  ASSERT_OK_AND_ASSIGN(Literal result, Execute(std::move(module), {&p0, &p1}));
+
+  // Five iterations of x = (x + 2) * 2 give 32 * x + 124.
+  std::vector<Literal> results = result.DecomposeTuple();
+  ASSERT_EQ(results.size(), 2);
+  LiteralTestUtil::ExpectR1Equal(
+      absl::MakeConstSpan({124.0f, 156.0f, 188.0f, 220.0f}), results[0]);
+  LiteralTestUtil::ExpectR1Equal(
+      absl::MakeConstSpan({444.0f, 764.0f, 1084.0f, 1404.0f}), results[1]);
+}
+
 }  // namespace
 }  // namespace cpu
 }  // namespace xla
diff --git a/third_party/xla/xla/service/cpu/thunk_emitter.cc b/third_party/xla/xla/service/cpu/thunk_emitter.cc
index 8136558339e..8606b8d5561 100644
--- a/third_party/xla/xla/service/cpu/thunk_emitter.cc
+++ b/third_party/xla/xla/service/cpu/thunk_emitter.cc
@@ -39,6 +39,7 @@ limitations under the License.
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/Support/DebugStringHelper.h"
 #include "xla/backends/cpu/alignment.h"
+#include "xla/backends/cpu/codegen/computation_kernel_arguments.h"
 #include "xla/backends/cpu/codegen/computation_kernel_emitter.h"
 #include "xla/backends/cpu/codegen/dot/dot_kernel_emitter.h"
 #include "xla/backends/cpu/codegen/elemental/concatenate_kernel_emitter.h"
@@ -588,8 +589,33 @@ absl::StatusOr<ThunkSequence> ThunkEmitter::EmitCallThunk(
   if (std::optional<std::string> maybe_small_call =
           instruction->get_frontend_attribute("xla_cpu_small_call");
       maybe_small_call.has_value() && *maybe_small_call == "true") {
+    ABSL_ASSIGN_OR_RETURN(ComputationKernelArguments arguments,
+                     ComputationKernelArguments::Create(
+                         *instruction->to_apply(), buffer_assignment_));
+
+    // Hoisted loops are often clones of each other, so compile each distinct
+    // kernel once. The prefix keeps keys apart from fusion fingerprints.
+    std::optional<std::string> cache_key;
+    if (arguments.fingerprint().has_value()) {
+      cache_key = absl::StrCat("small_call:", *arguments.fingerprint());
+      if (const auto itr = kernel_spec_cache_.find(*cache_key);
+          itr != kernel_spec_cache_.end()) {
+        const KernelSpec& cached_spec = itr->second;
+        VLOG(1) << "Reusing kernel: " << cached_spec.name()
+                << " for call: " << instruction->name();
+        // Kernel parameters are positional, matching the order of `arguments`.
+        KernelSpec kernel_spec(cached_spec.name(),
+                               cached_spec.work_dimensions(),
+                               arguments.arguments(), arguments.results(),
+                               cached_spec.invariant_arguments());
+        return MakeKernelThunkSequence(instruction, std::move(kernel_spec),
+                                       /*min_alignment=*/MinAlign());
+      }
+    }
+
     ComputationKernelEmitter emitter(instruction, &buffer_assignment_,
-                                     &target_machine_features_);
+                                     &target_machine_features_,
+                                     std::move(arguments));
     ABSL_ASSIGN_OR_RETURN(KernelDefinition kernel_definition,
                      emitter.EmitKernelDefinition());
 
@@ -599,6 +625,10 @@ absl::StatusOr<ThunkSequence> ThunkEmitter::EmitCallThunk(
     kernels_.push_back(
         {kernel_spec.name(), std::move(kernel_source).thread_safe_module()});
 
+    if (cache_key.has_value()) {
+      kernel_spec_cache_.insert({*std::move(cache_key), kernel_spec});
+    }
+
     return MakeKernelThunkSequence(instruction, std::move(kernel_spec),
                                    /*min_alignment=*/MinAlign());
   } else {

From d6aa13b59ffc6109cf164aefb272b532542b8e33 Mon Sep 17 00:00:00 2001
From: Alexander Lyashuk <crem@google.com>
Date: Fri, 16 Oct 2026 05:54:25 -0700
Subject: [PATCH 0958/1000] [XLA] Support ShapeTracker in DotOperandDims and
 simplify bitcast tracking.

Decompose bitcasts in ShapeTracker as a transpose to physical order, a reshape
to the target physical dimensions, and a transpose back to logical order. This
handles bitcasts whose dimensions do not factor into each other. Also preserve
the input layout in ShapeTracker::GetInverted.

Add DotOperandDims::MapForward, MapBackward, and FoldBackward overloads for
ShapeTracker, plus DotOperandDims::IsSorted. FoldBackward folds a chain of
shape-tracking ops into a dot operand when each dimension category keeps its
sizes and the non-contracting dimensions stay sorted.

PiperOrigin-RevId: 1000802358
---
 third_party/xla/xla/hlo/analysis/BUILD        |   1 +
 .../xla/xla/hlo/analysis/shape_tracker.cc     | 127 ++-----------
 .../xla/hlo/analysis/shape_tracker_test.cc    |  55 +++++-
 third_party/xla/xla/service/BUILD             |   1 +
 .../xla/xla/service/matmul_indexing_utils.cc  |  49 +++++
 .../xla/xla/service/matmul_indexing_utils.h   |  24 +++
 .../xla/service/matmul_indexing_utils_test.cc | 173 ++++++++++++++++++
 7 files changed, 316 insertions(+), 114 deletions(-)

diff --git a/third_party/xla/xla/hlo/analysis/BUILD b/third_party/xla/xla/hlo/analysis/BUILD
index 72797f003cd..1e6799cc5a7 100644
--- a/third_party/xla/xla/hlo/analysis/BUILD
+++ b/third_party/xla/xla/hlo/analysis/BUILD
@@ -823,6 +823,7 @@ xla_cc_test(
         "//xla/service:hlo_module_config",
         "//xla/tsl/platform:test_main",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest",
         "@llvm-project//llvm:Support",
diff --git a/third_party/xla/xla/hlo/analysis/shape_tracker.cc b/third_party/xla/xla/hlo/analysis/shape_tracker.cc
index f1380ea33b9..1454caa1820 100644
--- a/third_party/xla/xla/hlo/analysis/shape_tracker.cc
+++ b/third_party/xla/xla/hlo/analysis/shape_tracker.cc
@@ -59,75 +59,11 @@ limitations under the License.
 namespace xla {
 namespace {
 
-struct PhysicalDimension {
-  int64_t size = 0;
-  int64_t src_logical_idx = -1;
-  int64_t dst_logical_idx = -1;
-  int64_t idx = -1;
-  int64_t src_expanded_idx = -1;
-};
-
-// Returns a common physical decomposition of dimensions for both shapes to
-// facilitate bitcast lowering. The returned dimensions are ordered major to
-// minor.
-absl::StatusOr<std::vector<PhysicalDimension>> BuildPhysicalDimensions(
-    const Shape& src_shape, const Shape& dst_shape) {
-  std::vector<PhysicalDimension> result;
-  int64_t src_idx = 0;
-  int64_t dst_idx = 0;
-  int64_t src_rem = 1;
-  int64_t dst_rem = 1;
-
-  auto get_dim = [](const Shape& shape, int64_t idx) {
-    return shape.dimensions(shape.layout().minor_to_major(idx));
-  };
-  auto get_logical = [](const Shape& shape, int64_t idx) {
-    return shape.layout().minor_to_major(idx);
-  };
-
-  while (src_idx < src_shape.dimensions().size() ||
-         dst_idx < dst_shape.dimensions().size()) {
-    if (src_rem == 1 && src_idx < src_shape.dimensions().size()) {
-      src_rem = get_dim(src_shape, src_idx);
-    }
-    if (dst_rem == 1 && dst_idx < dst_shape.dimensions().size()) {
-      dst_rem = get_dim(dst_shape, dst_idx);
-    }
-
-    int64_t size = std::min(src_rem, dst_rem);
-    PhysicalDimension dim;
-    dim.size = size;
-    dim.src_logical_idx =
-        src_idx < src_shape.dimensions().size()
-            ? get_logical(src_shape, src_idx)
-            : get_logical(src_shape, src_shape.dimensions().size() - 1);
-    dim.dst_logical_idx =
-        dst_idx < dst_shape.dimensions().size()
-            ? get_logical(dst_shape, dst_idx)
-            : get_logical(dst_shape, dst_shape.dimensions().size() - 1);
-    result.push_back(dim);
-
-    if (src_rem % size != 0 || dst_rem % size != 0) {
-      return absl::InvalidArgumentError("Cannot decompose bitcast");
-    }
-
-    src_rem /= size;
-    dst_rem /= size;
-
-    if (src_rem == 1) {
-      src_idx++;
-    }
-    if (dst_rem == 1) {
-      dst_idx++;
-    }
-  }
-
-  absl::c_reverse(result);
-  size_t idx = 0;
-  for (PhysicalDimension& dim : result) {
-    dim.idx = idx++;
-  }
-  return result;
+// Returns the dimensions of `shape` in major-to-minor order, i.e. the
+// permutation that transposes `shape` into its physical order.
+std::vector<int64_t> MajorToMinor(const Shape& shape) {
+  absl::Span<const int64_t> minor_to_major = shape.layout().minor_to_major();
+  return std::vector<int64_t>(minor_to_major.rbegin(), minor_to_major.rend());
 }
 
 void TryFoldProjections(std::vector<ShapeTracker::BufferView>& projections) {
@@ -552,8 +488,9 @@ absl::Status ShapeTracker::AppendReshape(absl::Span<const int64_t> dimensions) {
   return absl::OkStatus();
 }
 
-// Decomposes a bitcast into a reshape-transpose-reshape sequence and appends
-// it to the tracker.
+// A bitcast reinterprets the physical buffer, so it is a transpose of the
+// source into its physical order, a reshape to the destination's physical
+// dimensions and a transpose back into the destination's logical order.
 absl::Status ShapeTracker::AppendBitcast(const xla::Shape& src_shape,
                                          const xla::Shape& dst_shape) {
   if (ShapeUtil::ElementsIn(src_shape) != ShapeUtil::ElementsIn(dst_shape)) {
@@ -566,47 +503,11 @@ absl::Status ShapeTracker::AppendBitcast(const xla::Shape& src_shape,
         "Bitcast operand shape does not match current output shape");
   }
 
-  if (ShapeUtil::ElementsIn(output_shape_) == 1) {
-    return AppendReshape(dst_shape.dimensions());
-  }
-
-  ABSL_ASSIGN_OR_RETURN(std::vector<PhysicalDimension> dims,
-                   BuildPhysicalDimensions(src_shape, dst_shape));
-  absl::c_stable_sort(
-      dims, [](const PhysicalDimension& a, const PhysicalDimension& b) {
-        return a.src_logical_idx < b.src_logical_idx;
-      });
-  size_t idx = 0;
-  for (PhysicalDimension& dim : dims) {
-    dim.src_expanded_idx = idx++;
-  }
-
-  std::vector<int64_t> sizes;
-  sizes.reserve(dims.size());
-  for (const PhysicalDimension& dim : dims) {
-    sizes.push_back(dim.size);
-  }
-  // Decompose bitcast into reshape-transpose-reshape:
-  // 1. Reshape to physical dimensions.
-  ABSL_RETURN_IF_ERROR(AppendReshape(sizes));
-
-  absl::c_sort(dims, [](const PhysicalDimension& a,
-                        const PhysicalDimension& b) { return a.idx < b.idx; });
-  absl::c_stable_sort(
-      dims, [](const PhysicalDimension& a, const PhysicalDimension& b) {
-        return a.dst_logical_idx < b.dst_logical_idx;
-      });
-  std::vector<int64_t> permutation;
-  permutation.reserve(dims.size());
-  for (const PhysicalDimension& dim : dims) {
-    permutation.push_back(dim.src_expanded_idx);
-  }
-  // 2. Transpose to match destination physical layout.
-  ABSL_RETURN_IF_ERROR(AppendTranspose(permutation));
-  // 3. Reshape to final destination shape.
-  ABSL_RETURN_IF_ERROR(AppendReshape(dst_shape.dimensions()));
-
-  return absl::OkStatus();
+  std::vector<int64_t> dst_major_to_minor = MajorToMinor(dst_shape);
+  ABSL_RETURN_IF_ERROR(AppendTranspose(MajorToMinor(src_shape)));
+  ABSL_RETURN_IF_ERROR(
+      AppendReshape(Permute(dst_shape.dimensions(), dst_major_to_minor)));
+  return AppendTranspose(InversePermutation(dst_major_to_minor));
 }
 
 absl::Status ShapeTracker::AppendInstruction(const HloInstruction* inst) {
@@ -742,6 +643,8 @@ absl::StatusOr<ShapeTracker> ShapeTracker::GetInverted() const {
   }
 
   ABSL_RETURN_IF_ERROR(inverted.AppendReshape(input_shape_.dimensions()));
+  // The reshape above resets the layout; restore the original input's one.
+  inverted.output_shape_ = input_shape_;
 
   return inverted;
 }
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 0d5e628d191..c21db498685 100644
--- a/third_party/xla/xla/hlo/analysis/shape_tracker_test.cc
+++ b/third_party/xla/xla/hlo/analysis/shape_tracker_test.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 #include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
 #include "absl/types/span.h"
 #include "llvm/ADT/SmallVector.h"
 #include "xla/hlo/ir/hlo_computation.h"
@@ -145,6 +146,17 @@ TEST(ShapeTrackerTest, GetInvertedSimple) {
   EXPECT_EQ(inverted.output_shape().dimensions(), (std::vector<int64_t>{2, 3}));
 }
 
+TEST(ShapeTrackerTest, GetInvertedPreservesInputLayout) {
+  Shape shape = ShapeUtil::MakeShapeWithDenseLayout(F32, {2, 3, 4}, {0, 1, 2});
+  ShapeTracker tracker(shape);
+  ASSERT_TRUE(tracker.AppendTranspose({2, 0, 1}).ok());
+  ASSERT_TRUE(tracker.AppendReshape({4, 6}).ok());
+
+  auto inverted_or = tracker.GetInverted();
+  ASSERT_TRUE(inverted_or.ok());
+  EXPECT_EQ(inverted_or->output_shape(), shape);
+}
+
 TEST(ShapeTrackerTest, AppendBitcastSimple) {
   Shape src_shape = ShapeUtil::MakeShape(F32, {6});
   Shape dst_shape = ShapeUtil::MakeShape(F32, {2, 3});
@@ -165,12 +177,51 @@ TEST(ShapeTrackerTest, AppendBitcastWithLayout) {
   EXPECT_EQ(tracker.output_shape().dimensions(), (std::vector<int64_t>{6}));
 }
 
-TEST(ShapeTrackerTest, AppendBitcastFail) {
+TEST(ShapeTrackerTest, AppendBitcastNonFactorizable) {
   Shape src_shape = ShapeUtil::MakeShape(F32, {2, 3});
   Shape dst_shape = ShapeUtil::MakeShape(F32, {3, 2});
 
   ShapeTracker tracker(src_shape);
-  EXPECT_FALSE(tracker.AppendBitcast(src_shape, dst_shape).ok());
+  ASSERT_TRUE(tracker.AppendBitcast(src_shape, dst_shape).ok());
+
+  ShapeTracker expected(src_shape);
+  ASSERT_TRUE(expected.AppendReshape({3, 2}).ok());
+  EXPECT_EQ(tracker, expected);
+
+  // Changing the layout between {1,0} and {0,1} reshapes the physical buffer
+  // from [2,3] to [3,2] and transposes it back to [2,3], so individual
+  // dimensions cannot be mapped or narrowed across the bitcast.
+  Shape dst_col_major =
+      ShapeUtil::MakeShapeWithDenseLayout(F32, {2, 3}, {0, 1});
+  ShapeTracker col_major_tracker(src_shape);
+  ASSERT_TRUE(col_major_tracker.AppendBitcast(src_shape, dst_col_major).ok());
+  EXPECT_EQ(col_major_tracker.DebugString(), "[2,3] -> R[3,2] -> T[2,3]");
+  EXPECT_EQ(col_major_tracker.MapInputDimensionsToOutputUnordered({0}),
+            std::nullopt);
+  EXPECT_THAT(
+      col_major_tracker.Narrow({0}),
+      absl_testing::StatusIs(
+          absl::StatusCode::kInvalidArgument,
+          ::testing::HasSubstr("Slice is incompatible with projection")));
+}
+
+TEST(ShapeTrackerTest, AppendBitcastIncompatible) {
+  Shape src_shape = ShapeUtil::MakeShape(F32, {2, 3});
+  ShapeTracker tracker(src_shape);
+
+  EXPECT_THAT(
+      tracker.AppendBitcast(src_shape, ShapeUtil::MakeShape(F32, {2, 4})),
+      absl_testing::StatusIs(
+          absl::StatusCode::kInvalidArgument,
+          ::testing::HasSubstr(
+              "Bitcast must preserve the total number of elements")));
+  EXPECT_THAT(
+      tracker.AppendBitcast(ShapeUtil::MakeShape(F32, {3, 2}),
+                            ShapeUtil::MakeShape(F32, {2, 3})),
+      absl_testing::StatusIs(
+          absl::StatusCode::kInvalidArgument,
+          ::testing::HasSubstr(
+              "Bitcast operand shape does not match current output shape")));
 }
 
 TEST(ShapeTrackerTest, AppendBitcastTricky) {
diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD
index bcf32652adc..74504258ec9 100644
--- a/third_party/xla/xla/service/BUILD
+++ b/third_party/xla/xla/service/BUILD
@@ -6582,6 +6582,7 @@ xla_cc_test(
         "//xla:permutation_util",
         "//xla:shape_util",
         "//xla:xla_data_proto_cc",
+        "//xla/hlo/analysis:shape_tracker",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/parser:hlo_parser",
         "//xla/hlo/testlib:test",
diff --git a/third_party/xla/xla/service/matmul_indexing_utils.cc b/third_party/xla/xla/service/matmul_indexing_utils.cc
index f9e63401689..25d2126ab08 100644
--- a/third_party/xla/xla/service/matmul_indexing_utils.cc
+++ b/third_party/xla/xla/service/matmul_indexing_utils.cc
@@ -21,6 +21,7 @@ limitations under the License.
 #include <numeric>
 #include <optional>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "absl/algorithm/container.h"
@@ -33,6 +34,7 @@ limitations under the License.
 #include "absl/strings/str_join.h"
 #include "absl/types/span.h"
 #include "xla/autotuning.pb.h"
+#include "xla/hlo/analysis/shape_tracker.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/permutation_util.h"
@@ -255,6 +257,10 @@ DotOperandDims DotOperandDims::GetPermuted(
   return copy;
 }
 
+bool DotOperandDims::IsSorted(Category category) const {
+  return absl::c_is_sorted(dim_numbers_[category]);
+}
+
 bool DotOperandDims::IsConsecutive(Category category) const {
   const auto& dims = dim_numbers_[category];
   return absl::c_adjacent_find(dims, [](int64_t a, int64_t b) {
@@ -500,6 +506,49 @@ absl::StatusOr<std::optional<DotOperandDims>> DotOperandDims::MapForward(
   }
 }
 
+absl::StatusOr<std::optional<DotOperandDims>> DotOperandDims::MapForward(
+    const ShapeTracker& tracker) const {
+  TF_RET_CHECK(tracker.input_shape().dimensions() == shape_.dimensions());
+  DotOperandDims dims = *this;
+  for (const ShapeTracker::Step& step : tracker.GetSteps()) {
+    if (step.type == ShapeTracker::Step::Type::kTranspose) {
+      dims.ApplyPermutation(step.dimensions);
+      continue;
+    }
+    ABSL_ASSIGN_OR_RETURN(std::optional<DotOperandDims> reshaped,
+                     dims.Reshape(ShapeUtil::MakeShape(shape_.element_type(),
+                                                       step.dimensions)));
+    if (!reshaped.has_value()) {
+      return std::nullopt;
+    }
+    dims = *std::move(reshaped);
+  }
+  ABSL_RETURN_IF_ERROR(dims.SetShape(tracker.output_shape()));
+  return dims;
+}
+
+absl::StatusOr<std::optional<DotOperandDims>> DotOperandDims::MapBackward(
+    const ShapeTracker& tracker) const {
+  ABSL_ASSIGN_OR_RETURN(ShapeTracker inverted, tracker.GetInverted());
+  return MapForward(inverted);
+}
+
+absl::StatusOr<std::optional<DotOperandDims>> DotOperandDims::FoldBackward(
+    const ShapeTracker& tracker) const {
+  ABSL_ASSIGN_OR_RETURN(std::optional<DotOperandDims> dims, MapBackward(tracker));
+  if (!dims.has_value() || !dims->IsSorted(kNonContracting)) {
+    return std::nullopt;
+  }
+  // Ensure within-category dimension sizes are preserved so the dot's output
+  // shape and dimension pairing with the other operand remain unchanged.
+  for (Category category : {kBatch, kNonContracting, kContracting}) {
+    if (dims->Sizes(category) != Sizes(category)) {
+      return std::nullopt;
+    }
+  }
+  return dims;
+}
+
 absl::StatusOr<int64_t> DotOperandDims::IndexWithinCategory(
     Category category, int64_t global_dim_idx) const {
   const auto& dims = dim_numbers_[category];
diff --git a/third_party/xla/xla/service/matmul_indexing_utils.h b/third_party/xla/xla/service/matmul_indexing_utils.h
index 907d53461da..4f4a8a9a7a9 100644
--- a/third_party/xla/xla/service/matmul_indexing_utils.h
+++ b/third_party/xla/xla/service/matmul_indexing_utils.h
@@ -125,6 +125,9 @@ class DotOperandDims {
   // (if there was any).
   absl::Status CollapseCategory(Category category, bool remove_if_empty);
 
+  // Returns true if the dimensions of the given category are sorted.
+  bool IsSorted(Category category) const;
+
   // Returns true if the dimensions of the given category are sorted and
   // consecutive.
   bool IsConsecutive(Category category) const;
@@ -196,6 +199,27 @@ class DotOperandDims {
   absl::StatusOr<std::optional<DotOperandDims>> MapForward(
       const HloInstruction* inst) const;
 
+  // Same as above, but through all steps of `tracker`: from its input shape to
+  // its output shape (forward) or the other way around (backward). As with the
+  // single-instruction overloads, the categories are tracked, but the dot
+  // output may change (e.g. non-contracting dimensions may be split or
+  // reordered). Returns std::nullopt if the mapping is not possible.
+  absl::StatusOr<std::optional<DotOperandDims>> MapForward(
+      const ShapeTracker& tracker) const;
+  absl::StatusOr<std::optional<DotOperandDims>> MapBackward(
+      const ShapeTracker& tracker) const;
+
+  // Folds `tracker` (whose output is this operand) into the dot: returns the
+  // dims with which the dot computes exactly the same result when reading the
+  // tracker's input directly. This requires every category to keep its
+  // dimension sizes (so dimension pairing and the output shape are unchanged)
+  // and the non-contracting dimensions to stay sorted, as they are in any
+  // operand of an actual dot (so the output is not transposed). Returns
+  // std::nullopt otherwise; in particular, moving degenerate dimensions is a
+  // reshape for `tracker` and therefore not foldable.
+  absl::StatusOr<std::optional<DotOperandDims>> FoldBackward(
+      const ShapeTracker& tracker) const;
+
   // Returns a debug string representation of the DotOperandDims.
   std::string ToString() const;
 
diff --git a/third_party/xla/xla/service/matmul_indexing_utils_test.cc b/third_party/xla/xla/service/matmul_indexing_utils_test.cc
index cc1a22fac94..15e072e4aef 100644
--- a/third_party/xla/xla/service/matmul_indexing_utils_test.cc
+++ b/third_party/xla/xla/service/matmul_indexing_utils_test.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "xla/service/matmul_indexing_utils.h"
 
 #include <cstdint>
+#include <optional>
 #include <vector>
 
 #include <gmock/gmock.h>
@@ -23,6 +24,7 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
 #include "absl/strings/string_view.h"
+#include "xla/hlo/analysis/shape_tracker.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/parser/hlo_parser.h"
 #include "xla/hlo/testlib/test.h"
@@ -627,6 +629,16 @@ TEST(DotOperandDimsTest, ReshapeConsecutiveButNotSorted) {
   EXPECT_FALSE(mapped_dims.has_value());
 }
 
+TEST(DotOperandDimsTest, IsSorted) {
+  ASSERT_OK_AND_ASSIGN(Shape shape, ParseShape("f32[10,20,30,40,50,60]"));
+  DotOperandDims dims(shape, /*batch_dims=*/{},
+                      /*non_contracting_dims=*/{1, 4},
+                      /*contracting_dims=*/{5, 2, 3, 0});
+  EXPECT_TRUE(dims.IsSorted(DotOperandDims::kBatch));
+  EXPECT_TRUE(dims.IsSorted(DotOperandDims::kNonContracting));
+  EXPECT_FALSE(dims.IsSorted(DotOperandDims::kContracting));
+}
+
 TEST(DotOperandDimsTest, IsConsecutive) {
   ASSERT_OK_AND_ASSIGN(Shape shape, ParseShape("f32[10,20,30,40,50,60]"));
 
@@ -745,5 +757,166 @@ TEST(DotOperandDimsTest, ReshapeScalarToDegenerate) {
   }
 }
 
+TEST(DotOperandDimsTest, MapForwardThroughTracker) {
+  ASSERT_OK_AND_ASSIGN(Shape input_shape, ParseShape("f32[10,4,5,30]"));
+  DotOperandDims dims(input_shape, /*batch_dims=*/{0},
+                      /*non_contracting_dims=*/{1, 2},
+                      /*contracting_dims=*/{3});
+  // [10,4,5,30] -> transpose -> [30,10,4,5] -> reshape -> [30,10,20]
+  ShapeTracker tracker(input_shape);
+  ASSERT_OK(tracker.AppendTranspose({3, 0, 1, 2}));
+  ASSERT_OK(tracker.AppendReshape({30, 10, 20}));
+
+  ASSERT_OK_AND_ASSIGN(std::optional<DotOperandDims> mapped,
+                       dims.MapForward(tracker));
+  ASSERT_TRUE(mapped.has_value());
+  EXPECT_THAT(mapped->shape().dimensions(), ElementsAre(30, 10, 20));
+  EXPECT_THAT(mapped->Indices(DotOperandDims::kBatch), ElementsAre(1));
+  EXPECT_THAT(mapped->Indices(DotOperandDims::kNonContracting), ElementsAre(2));
+  EXPECT_THAT(mapped->Indices(DotOperandDims::kContracting), ElementsAre(0));
+}
+
+TEST(DotOperandDimsTest, MapBackwardThroughTracker) {
+  ASSERT_OK_AND_ASSIGN(Shape source_shape,
+                       ParseShape("f32[10,4,5,30]{0,1,2,3}"));
+  // [10,4,5,30] -> transpose -> [30,10,4,5] -> reshape -> [30,10,20]
+  ShapeTracker tracker(source_shape);
+  ASSERT_OK(tracker.AppendTranspose({3, 0, 1, 2}));
+  ASSERT_OK(tracker.AppendReshape({30, 10, 20}));
+  DotOperandDims dims(tracker.output_shape(), /*batch_dims=*/{1},
+                      /*non_contracting_dims=*/{2},
+                      /*contracting_dims=*/{0});
+
+  ASSERT_OK_AND_ASSIGN(std::optional<DotOperandDims> mapped,
+                       dims.MapBackward(tracker));
+  ASSERT_TRUE(mapped.has_value());
+  // The result gets the tracker's input shape, including its layout.
+  EXPECT_EQ(mapped->shape(), source_shape);
+  EXPECT_THAT(mapped->Indices(DotOperandDims::kBatch), ElementsAre(0));
+  EXPECT_THAT(mapped->Indices(DotOperandDims::kNonContracting),
+              ElementsAre(1, 2));
+  EXPECT_THAT(mapped->Indices(DotOperandDims::kContracting), ElementsAre(3));
+}
+
+TEST(DotOperandDimsTest, MapBackwardThroughTrackerReshape) {
+  ASSERT_OK_AND_ASSIGN(Shape source_shape, ParseShape("f32[10,4,5]"));
+  ShapeTracker tracker(source_shape);
+  ASSERT_OK(tracker.AppendReshape({10, 20}));
+  // The reshape merges two source dimensions into one contracting dimension.
+  DotOperandDims dims(tracker.output_shape(), /*batch_dims=*/{0},
+                      /*non_contracting_dims=*/{},
+                      /*contracting_dims=*/{1});
+  ASSERT_OK_AND_ASSIGN(std::optional<DotOperandDims> mapped,
+                       dims.MapBackward(tracker));
+  ASSERT_TRUE(mapped.has_value());
+  EXPECT_THAT(mapped->Indices(DotOperandDims::kContracting), ElementsAre(1, 2));
+
+  // A reshape whose dimensions don't align with the source ones can't be
+  // mapped.
+  ShapeTracker misaligned_tracker(source_shape);
+  ASSERT_OK(misaligned_tracker.AppendReshape({2, 100}));
+  DotOperandDims misaligned_dims(misaligned_tracker.output_shape(),
+                                 /*batch_dims=*/{0},
+                                 /*non_contracting_dims=*/{},
+                                 /*contracting_dims=*/{1});
+  ASSERT_OK_AND_ASSIGN(mapped, misaligned_dims.MapBackward(misaligned_tracker));
+  EXPECT_FALSE(mapped.has_value());
+}
+
+TEST(DotOperandDimsTest, FoldBackwardTranspose) {
+  ASSERT_OK_AND_ASSIGN(Shape source_shape, ParseShape("f32[4,8,16,32]"));
+  ShapeTracker tracker(source_shape);
+  ASSERT_OK(tracker.AppendTranspose({0, 2, 3, 1}));  // -> [4,16,32,8]
+  DotOperandDims dims(tracker.output_shape(), /*batch_dims=*/{0},
+                      /*non_contracting_dims=*/{1, 2},
+                      /*contracting_dims=*/{3});
+
+  ASSERT_OK_AND_ASSIGN(std::optional<DotOperandDims> folded,
+                       dims.FoldBackward(tracker));
+  ASSERT_TRUE(folded.has_value());
+  EXPECT_EQ(folded->shape(), source_shape);
+  EXPECT_THAT(folded->Indices(DotOperandDims::kBatch), ElementsAre(0));
+  EXPECT_THAT(folded->Indices(DotOperandDims::kNonContracting),
+              ElementsAre(2, 3));
+  EXPECT_THAT(folded->Indices(DotOperandDims::kContracting), ElementsAre(1));
+}
+
+TEST(DotOperandDimsTest, FoldBackwardKeepsContractingOrder) {
+  // Swapping contracting dimensions is fine: the dimension numbers keep the
+  // pairing with the other operand.
+  ASSERT_OK_AND_ASSIGN(Shape source_shape, ParseShape("f32[4,8,16]"));
+  ShapeTracker tracker(source_shape);
+  ASSERT_OK(tracker.AppendTranspose({2, 1, 0}));  // -> [16,8,4]
+  DotOperandDims dims(tracker.output_shape(), /*batch_dims=*/{},
+                      /*non_contracting_dims=*/{1},
+                      /*contracting_dims=*/{0, 2});
+
+  ASSERT_OK_AND_ASSIGN(std::optional<DotOperandDims> folded,
+                       dims.FoldBackward(tracker));
+  ASSERT_TRUE(folded.has_value());
+  EXPECT_THAT(folded->Indices(DotOperandDims::kContracting), ElementsAre(2, 0));
+  EXPECT_THAT(folded->Sizes(DotOperandDims::kContracting), ElementsAre(16, 4));
+}
+
+TEST(DotOperandDimsTest, FoldBackwardRejectsNonContractingReorder) {
+  ASSERT_OK_AND_ASSIGN(Shape source_shape, ParseShape("f32[8,16,32]"));
+  ShapeTracker tracker(source_shape);
+  ASSERT_OK(tracker.AppendTranspose({2, 1, 0}));  // -> [32,16,8]
+  DotOperandDims dims(tracker.output_shape(), /*batch_dims=*/{},
+                      /*non_contracting_dims=*/{0, 1},
+                      /*contracting_dims=*/{2});
+  ASSERT_OK_AND_ASSIGN(std::optional<DotOperandDims> folded,
+                       dims.FoldBackward(tracker));
+  EXPECT_FALSE(folded.has_value());
+}
+
+TEST(DotOperandDimsTest, FoldBackwardRejectsReshape) {
+  ASSERT_OK_AND_ASSIGN(Shape source_shape, ParseShape("f32[8,4,8]"));
+  ShapeTracker tracker(source_shape);
+  ASSERT_OK(tracker.AppendReshape({8, 32}));
+  DotOperandDims dims(tracker.output_shape(), /*batch_dims=*/{},
+                      /*non_contracting_dims=*/{1},
+                      /*contracting_dims=*/{0});
+  // Mapping is possible, but the output would have two dimensions instead of
+  // one.
+  ASSERT_OK_AND_ASSIGN(std::optional<DotOperandDims> mapped,
+                       dims.MapBackward(tracker));
+  EXPECT_TRUE(mapped.has_value());
+  ASSERT_OK_AND_ASSIGN(std::optional<DotOperandDims> folded,
+                       dims.FoldBackward(tracker));
+  EXPECT_FALSE(folded.has_value());
+}
+
+TEST(DotOperandDimsTest, FoldBackwardDegenerateDimensions) {
+  // Dropping a degenerate dimension: mapping is possible (the dimension joins
+  // the non-contracting ones), but the output would get an extra dimension.
+  ASSERT_OK_AND_ASSIGN(Shape source_shape, ParseShape("f32[1,8,6]"));
+  ShapeTracker tracker(source_shape);
+  ASSERT_OK(tracker.AppendReshape({8, 6}));
+  DotOperandDims dims(tracker.output_shape(), /*batch_dims=*/{},
+                      /*non_contracting_dims=*/{0},
+                      /*contracting_dims=*/{1});
+  ASSERT_OK_AND_ASSIGN(std::optional<DotOperandDims> mapped,
+                       dims.MapBackward(tracker));
+  ASSERT_TRUE(mapped.has_value());
+  EXPECT_THAT(mapped->Indices(DotOperandDims::kNonContracting),
+              ElementsAre(0, 1));
+  ASSERT_OK_AND_ASSIGN(std::optional<DotOperandDims> folded,
+                       dims.FoldBackward(tracker));
+  EXPECT_FALSE(folded.has_value());
+
+  // Moving a degenerate dimension is a reshape for the tracker, which loses
+  // the dimension's identity. Conservatively rejected.
+  ASSERT_OK_AND_ASSIGN(source_shape, ParseShape("f32[8,1,6]"));
+  ShapeTracker transpose_tracker(source_shape);
+  ASSERT_OK(transpose_tracker.AppendTranspose({1, 0, 2}));  // -> [1,8,6]
+  DotOperandDims transposed_dims(transpose_tracker.output_shape(),
+                                 /*batch_dims=*/{0},
+                                 /*non_contracting_dims=*/{1},
+                                 /*contracting_dims=*/{2});
+  ASSERT_OK_AND_ASSIGN(folded, transposed_dims.FoldBackward(transpose_tracker));
+  EXPECT_FALSE(folded.has_value());
+}
+
 }  // namespace
 }  // namespace xla

From 579f1ed851897e0ab6f6ef0705670cd6d86f69a0 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Fri, 16 Oct 2026 06:14:51 -0700
Subject: [PATCH 0959/1000] [XLA:GPU] Migrate GpuExecutableVaRemapAllocator to
 PerDeviceState.

PiperOrigin-RevId: 1000810350
---
 third_party/xla/xla/service/gpu/BUILD         |   5 +-
 .../xla/xla/service/gpu/gpu_executable.cc     |   3 +-
 .../gpu/gpu_executable_buffer_allocator.cc    |   6 +-
 .../gpu/gpu_executable_buffer_allocator.h     |   2 +-
 .../xla/service/gpu/gpu_executable_test.cc    |   3 +-
 .../gpu/gpu_executable_va_remap_allocator.cc  |  32 ++---
 .../gpu/gpu_executable_va_remap_allocator.h   |  14 +-
 .../gpu_executable_va_remap_allocator_test.cc | 136 ++++++++++++------
 8 files changed, 128 insertions(+), 73 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index f13ebab73d1..fd5f2c696c7 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -721,6 +721,7 @@ cc_library(
         "//xla:util",
         "//xla:xla_proto_cc",
         "//xla/backends/gpu/runtime:command_buffer_thunk",
+        "//xla/backends/gpu/runtime:per_device_state",
         "//xla/backends/gpu/runtime:thunk",
         "//xla/backends/gpu/runtime:thunk_executor",
         "//xla/service:buffer_assignment",
@@ -732,11 +733,9 @@ cc_library(
         "//xla/stream_executor:memory_reservation",
         "//xla/stream_executor:stream",
         "//xla/stream_executor:stream_executor_h",
-        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
-        "@com_google_absl//absl/container:node_hash_map",
         "@com_google_absl//absl/functional:function_ref",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
@@ -803,9 +802,11 @@ xla_cc_test(
         "//xla/stream_executor:mock_stream",
         "//xla/stream_executor:mock_stream_executor",
         "//xla/stream_executor:stream_executor_h",
+        "//xla/tsl/platform:env",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest_main",
     ],
diff --git a/third_party/xla/xla/service/gpu/gpu_executable.cc b/third_party/xla/xla/service/gpu/gpu_executable.cc
index 9667ea18f52..39285ac3a8b 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable.cc
@@ -534,7 +534,8 @@ GpuExecutable::GpuExecutable(
       has_module() ? &module_config().debug_options() : nullptr;
   buffer_allocator_ = GpuExecutableBufferAllocator::Create(
       module_name_, allocation_ptrs_, program_shape_.result(),
-      allocation_debug_options, thunk_executor_.get());
+      allocation_debug_options, thunk_executor_.get(),
+      gpu_topology_.has_value() ? gpu_topology_->num_devices_per_host() : 0);
 }
 
 GpuExecutable::~GpuExecutable() {
diff --git a/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.cc b/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.cc
index a86b4d8cb93..d167b976c4e 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.cc
@@ -96,7 +96,7 @@ GpuExecutableBufferAllocator::Create(
     absl::string_view module_name,
     absl::Span<const BufferAllocation* const> allocations,
     const Shape& result_shape, const DebugOptions* debug_options,
-    ThunkExecutor* thunk_executor) {
+    ThunkExecutor* thunk_executor, int num_devices) {
   const DebugOptions::CommandBufferUpdateMode update_mode =
       debug_options != nullptr
           ? debug_options->xla_gpu_command_buffer_update_mode()
@@ -109,8 +109,8 @@ GpuExecutableBufferAllocator::Create(
     case DebugOptions::SKIP_TEMP:
     case DebugOptions::SKIP_PROFILED:
       return std::make_unique<GpuExecutableVaRemapAllocator>(
-          module_name, allocations, result_shape, debug_options,
-          thunk_executor);
+          module_name, allocations, result_shape, debug_options, thunk_executor,
+          num_devices);
     default:
       LOG(FATAL) << "Unsupported command buffer update mode: " << update_mode;
   }
diff --git a/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.h b/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.h
index a8a8b821e97..c0f8f7959b3 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.h
+++ b/third_party/xla/xla/service/gpu/gpu_executable_buffer_allocator.h
@@ -165,7 +165,7 @@ class GpuExecutableBufferAllocator {
       absl::string_view module_name,
       absl::Span<const BufferAllocation* const> allocations,
       const Shape& result_shape, const DebugOptions* debug_options,
-      ThunkExecutor* thunk_executor);
+      ThunkExecutor* thunk_executor, int num_devices);
 
   GpuExecutableBufferAllocator(
       absl::string_view module_name,
diff --git a/third_party/xla/xla/service/gpu/gpu_executable_test.cc b/third_party/xla/xla/service/gpu/gpu_executable_test.cc
index 4db82daba54..9c3f1b86346 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable_test.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable_test.cc
@@ -444,7 +444,8 @@ TEST_F(GpuExecutableTest, CommandBufferAllocationPolicy) {
     debug_options.set_xla_gpu_command_buffer_update_mode(update_mode);
     std::unique_ptr<GpuExecutableBufferAllocator> buffer_allocator =
         GpuExecutableBufferAllocator::Create("test", allocation_ptrs, shape,
-                                             &debug_options, &thunk_executor);
+                                             &debug_options, &thunk_executor,
+                                             /*num_devices=*/1);
     ServiceExecutableRunOptions run_options;
     ABSL_ASSIGN_OR_RETURN(
         std::unique_ptr<GpuExecutableBufferAllocator::ExecutionScope>
diff --git a/third_party/xla/xla/service/gpu/gpu_executable_va_remap_allocator.cc b/third_party/xla/xla/service/gpu/gpu_executable_va_remap_allocator.cc
index f1816ccd1b6..73b60f2a0ef 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable_va_remap_allocator.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable_va_remap_allocator.cc
@@ -32,6 +32,7 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/gpu/buffer_allocations.h"
 #include "xla/service/gpu/gpu_executable_buffer_allocator.h"
@@ -449,9 +450,10 @@ GpuExecutableVaRemapAllocator::GpuExecutableVaRemapAllocator(
     absl::string_view module_name,
     absl::Span<const BufferAllocation* const> allocations,
     const Shape& result_shape, const DebugOptions* debug_options,
-    ThunkExecutor* thunk_executor)
+    ThunkExecutor* thunk_executor, int num_devices)
     : GpuExecutableBufferAllocator(module_name, allocations, result_shape,
-                                   debug_options, thunk_executor) {
+                                   debug_options, thunk_executor),
+      remappings_(num_devices) {
   update_mode_ = debug_options != nullptr
                      ? debug_options->xla_gpu_command_buffer_update_mode()
                      : DebugOptions::ALWAYS_UPDATE;
@@ -599,20 +601,19 @@ void GpuExecutableVaRemapAllocator::TransitionProfiledRemapping(
 }
 
 GpuExecutableVaRemapAllocator::~GpuExecutableVaRemapAllocator() {
-  absl::MutexLock lock(remappings_mutex_);
-  for (auto& [executor, remapping] : remappings_) {
+  remappings_.ForEach([this](Remapping& remapping) {
     absl::MutexLock remap_lock(remapping.mutex);
     if (remapping.vmm_allocator == nullptr) {
-      continue;
+      return;
     }
     absl::Status status = remapping.vmm_allocator->SynchronizePendingOperations(
-        executor->device_ordinal());
+        remapping.device_ordinal);
     if (!status.ok()) {
       LOG(ERROR) << "Failed to synchronize command buffer VA remapping "
                     "deferred operations for module "
                  << module_name() << ": " << status;
     }
-  }
+  });
 }
 
 absl::StatusOr<std::unique_ptr<GpuExecutableBufferAllocator::ExecutionScope>>
@@ -645,24 +646,21 @@ GpuExecutableVaRemapAllocator::CreateExecutionScope(
         "and execution stream to match");
   }
 
-  Remapping* remapping = nullptr;
-  se::StreamExecutor* executor = run_options->stream()->parent();
-  {
-    absl::MutexLock lock(remappings_mutex_);
-    // This is the lifetime remapping object for this executable/executor. It
-    // owns the VA reservation reused by later ExecuteAsyncOnStream calls.
-    remapping = &remappings_[executor];
-  }
+  // This is the lifetime remapping object for this executable/device. It
+  // owns the VA reservation reused by later ExecuteAsyncOnStream calls.
+  ABSL_ASSIGN_OR_RETURN(Remapping * remapping,
+                   remappings_.GetOrCreate(device_ordinal));
 
   auto remap_lock = std::make_unique<absl::MutexLock>(remapping->mutex);
   if (remapping->vmm_allocator != nullptr &&
       remapping->vmm_allocator != vmm_allocator) {
     return Internal(
         "Command buffer VA remapping for module %s changed VMM allocator for "
-        "executor %p",
-        module_name(), executor);
+        "device %d",
+        module_name(), device_ordinal);
   }
   remapping->vmm_allocator = vmm_allocator;
+  remapping->device_ordinal = device_ordinal;
 
   if (profiled_mode) {
     // The lock is held through the whole execution in the profiling phase as
diff --git a/third_party/xla/xla/service/gpu/gpu_executable_va_remap_allocator.h b/third_party/xla/xla/service/gpu/gpu_executable_va_remap_allocator.h
index b608a519748..5ef77e80905 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable_va_remap_allocator.h
+++ b/third_party/xla/xla/service/gpu/gpu_executable_va_remap_allocator.h
@@ -20,14 +20,13 @@ limitations under the License.
 #include <memory>
 #include <vector>
 
-#include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/container/flat_hash_set.h"
-#include "absl/container/node_hash_map.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
+#include "xla/backends/gpu/runtime/per_device_state.h"
 #include "xla/service/buffer_assignment.h"
 #include "xla/service/gpu/gpu_executable_buffer_allocator.h"
 #include "xla/service/service_executable_run_options.h"
@@ -72,7 +71,7 @@ class GpuExecutableVaRemapAllocator : public GpuExecutableBufferAllocator {
       absl::string_view module_name,
       absl::Span<const BufferAllocation* const> allocations,
       const Shape& result_shape, const DebugOptions* debug_options,
-      ThunkExecutor* thunk_executor);
+      ThunkExecutor* thunk_executor, int num_devices);
   ~GpuExecutableVaRemapAllocator() override;
 
   // Creates the per-run execution scope. When VA remapping applies to this
@@ -122,6 +121,11 @@ class GpuExecutableVaRemapAllocator : public GpuExecutableBufferAllocator {
         allocation_to_mapping_size;
     std::unique_ptr<se::MemoryReservation> va_reservation;
     se::DeviceAddressVmmAllocator* vmm_allocator = nullptr;
+    // -1 is just a sentinel value since state needs to be default
+    // constructible. Always initialized alongside `vmm_allocator` in
+    // `CreateExecutionScope` before any execution or destructor cleanup can
+    // read it.
+    int device_ordinal = -1;
 
     ProfilePhase phase = ProfilePhase::kInactive;
     // Number of completed profiling observations.
@@ -167,9 +171,7 @@ class GpuExecutableVaRemapAllocator : public GpuExecutableBufferAllocator {
   // executions.
   AllocationIndexSet profile_candidate_alloc_indices_;
 
-  absl::Mutex remappings_mutex_;
-  absl::node_hash_map<se::StreamExecutor*, Remapping> remappings_
-      ABSL_GUARDED_BY(remappings_mutex_);
+  PerDeviceState<Remapping> remappings_;
 };
 
 }  // namespace gpu
diff --git a/third_party/xla/xla/service/gpu/gpu_executable_va_remap_allocator_test.cc b/third_party/xla/xla/service/gpu/gpu_executable_va_remap_allocator_test.cc
index 3ac037bd255..f5bc99f2840 100644
--- a/third_party/xla/xla/service/gpu/gpu_executable_va_remap_allocator_test.cc
+++ b/third_party/xla/xla/service/gpu/gpu_executable_va_remap_allocator_test.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "xla/service/gpu/gpu_executable_va_remap_allocator.h"
 
 #include <algorithm>
+#include <array>
 #include <cstddef>
 #include <cstdint>
 #include <memory>
@@ -29,6 +30,8 @@ limitations under the License.
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
+#include "absl/synchronization/blocking_counter.h"
+#include "absl/synchronization/notification.h"
 #include "absl/types/span.h"
 #include "xla/backends/gpu/runtime/thunk.h"
 #include "xla/backends/gpu/runtime/thunk_executor.h"
@@ -47,6 +50,8 @@ limitations under the License.
 #include "xla/stream_executor/mock_stream.h"
 #include "xla/stream_executor/mock_stream_executor.h"
 #include "xla/stream_executor/stream_executor.h"
+#include "xla/tsl/platform/env.h"
+#include "xla/tsl/platform/threadpool.h"
 #include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
 
@@ -325,9 +330,9 @@ TEST_F(GpuExecutableVaRemapAllocatorTest,
   ThunkExecutor thunk_executor{ThunkSequence{}};
   DebugOptions debug_options;
   debug_options.set_xla_gpu_command_buffer_update_mode(DebugOptions::SKIP_TEMP);
-  GpuExecutableVaRemapAllocator allocator("test", allocations,
-                                          ShapeUtil::MakeShape(F32, {256}),
-                                          &debug_options, &thunk_executor);
+  GpuExecutableVaRemapAllocator allocator(
+      "test", allocations, ShapeUtil::MakeShape(F32, {256}), &debug_options,
+      &thunk_executor, /*num_devices=*/1);
   allocator.AddVaRemappedAllocationForTesting(0);
   EXPECT_EQ(allocator.command_buffer_allocation_count(), 1);
 
@@ -440,9 +445,9 @@ TEST_F(GpuExecutableVaRemapAllocatorTest,
   ThunkExecutor thunk_executor{ThunkSequence{}};
   DebugOptions debug_options;
   debug_options.set_xla_gpu_command_buffer_update_mode(DebugOptions::SKIP_TEMP);
-  GpuExecutableVaRemapAllocator allocator("test", allocations,
-                                          ShapeUtil::MakeShape(F32, {256}),
-                                          &debug_options, &thunk_executor);
+  GpuExecutableVaRemapAllocator allocator(
+      "test", allocations, ShapeUtil::MakeShape(F32, {256}), &debug_options,
+      &thunk_executor, /*num_devices=*/1);
   allocator.AddVaRemappedAllocationForTesting(0);
   allocator.AddVaRemappedAllocationForTesting(1);
 
@@ -514,9 +519,9 @@ TEST_F(GpuExecutableVaRemapAllocatorTest,
   ThunkExecutor thunk_executor{ThunkSequence{}};
   DebugOptions debug_options;
   debug_options.set_xla_gpu_command_buffer_update_mode(DebugOptions::SKIP_TEMP);
-  GpuExecutableVaRemapAllocator allocator("test", allocations,
-                                          ShapeUtil::MakeShape(F32, {256}),
-                                          &debug_options, &thunk_executor);
+  GpuExecutableVaRemapAllocator allocator(
+      "test", allocations, ShapeUtil::MakeShape(F32, {256}), &debug_options,
+      &thunk_executor, /*num_devices=*/1);
   allocator.AddVaRemappedAllocationForTesting(0);
 
   ASSERT_OK_AND_ASSIGN(
@@ -582,9 +587,9 @@ TEST_F(GpuExecutableVaRemapAllocatorTest,
   ThunkExecutor thunk_executor{ThunkSequence{}};
   DebugOptions debug_options;
   debug_options.set_xla_gpu_command_buffer_update_mode(DebugOptions::SKIP_TEMP);
-  GpuExecutableVaRemapAllocator allocator("test", allocations,
-                                          ShapeUtil::MakeShape(F32, {250}),
-                                          &debug_options, &thunk_executor);
+  GpuExecutableVaRemapAllocator allocator(
+      "test", allocations, ShapeUtil::MakeShape(F32, {250}), &debug_options,
+      &thunk_executor, /*num_devices=*/1);
   allocator.AddVaRemappedAllocationForTesting(0);
 
   auto get_parameter_buffer = [&](const BufferAllocation& allocation)
@@ -686,9 +691,9 @@ TEST_F(GpuExecutableVaRemapAllocatorTest,
   DebugOptions debug_options;
   debug_options.set_xla_gpu_command_buffer_update_mode(
       DebugOptions::SKIP_PROFILED);
-  GpuExecutableVaRemapAllocator allocator("test", allocations,
-                                          ShapeUtil::MakeShape(F32, {256}),
-                                          &debug_options, &thunk_executor);
+  GpuExecutableVaRemapAllocator allocator(
+      "test", allocations, ShapeUtil::MakeShape(F32, {256}), &debug_options,
+      &thunk_executor, /*num_devices=*/1);
   allocator.AddProfileCandidateAllocationForTesting(0);
 
   // The parameter keeps the same VMM-backed address on every execution, so
@@ -750,9 +755,9 @@ TEST_F(GpuExecutableVaRemapAllocatorTest,
   DebugOptions debug_options;
   debug_options.set_xla_gpu_command_buffer_update_mode(
       DebugOptions::SKIP_PROFILED);
-  GpuExecutableVaRemapAllocator allocator("test", allocations,
-                                          ShapeUtil::MakeShape(F32, {256}),
-                                          &debug_options, &thunk_executor);
+  GpuExecutableVaRemapAllocator allocator(
+      "test", allocations, ShapeUtil::MakeShape(F32, {256}), &debug_options,
+      &thunk_executor, /*num_devices=*/1);
   allocator.AddProfileCandidateAllocationForTesting(0);
 
   ASSERT_OK_AND_ASSIGN(
@@ -816,9 +821,9 @@ TEST_F(GpuExecutableVaRemapAllocatorTest,
   DebugOptions debug_options;
   debug_options.set_xla_gpu_command_buffer_update_mode(
       DebugOptions::SKIP_PROFILED);
-  GpuExecutableVaRemapAllocator allocator("test", allocations,
-                                          ShapeUtil::MakeShape(F32, {256}),
-                                          &debug_options, &thunk_executor);
+  GpuExecutableVaRemapAllocator allocator(
+      "test", allocations, ShapeUtil::MakeShape(F32, {256}), &debug_options,
+      &thunk_executor, /*num_devices=*/1);
   // Simulate automatic constructor classification with an empty test thunk
   // sequence.
   allocator.AddVaRemappedAllocationForTesting(0);
@@ -875,9 +880,9 @@ TEST_F(GpuExecutableVaRemapAllocatorTest,
   DebugOptions debug_options;
   debug_options.set_xla_gpu_command_buffer_update_mode(
       DebugOptions::SKIP_PROFILED);
-  GpuExecutableVaRemapAllocator allocator("test", allocations,
-                                          ShapeUtil::MakeShape(F32, {256}),
-                                          &debug_options, &thunk_executor);
+  GpuExecutableVaRemapAllocator allocator(
+      "test", allocations, ShapeUtil::MakeShape(F32, {256}), &debug_options,
+      &thunk_executor, /*num_devices=*/1);
   // Simulate constructor classification with an empty test thunk sequence:
   // the temp is automatic, while the parameter is profiled.
   allocator.AddVaRemappedAllocationForTesting(0);
@@ -946,9 +951,9 @@ TEST_F(GpuExecutableVaRemapAllocatorTest,
   DebugOptions debug_options;
   debug_options.set_xla_gpu_command_buffer_update_mode(
       DebugOptions::SKIP_PROFILED);
-  GpuExecutableVaRemapAllocator allocator("test", allocations,
-                                          ShapeUtil::MakeShape(F32, {250}),
-                                          &debug_options, &thunk_executor);
+  GpuExecutableVaRemapAllocator allocator(
+      "test", allocations, ShapeUtil::MakeShape(F32, {250}), &debug_options,
+      &thunk_executor, /*num_devices=*/1);
   allocator.AddProfileCandidateAllocationForTesting(0);
 
   auto get_parameter_buffer = [&](const BufferAllocation& allocation)
@@ -1040,9 +1045,9 @@ TEST_F(GpuExecutableVaRemapAllocatorTest,
   DebugOptions debug_options;
   debug_options.set_xla_gpu_command_buffer_update_mode(
       DebugOptions::SKIP_PROFILED);
-  GpuExecutableVaRemapAllocator allocator("test", allocations,
-                                          ShapeUtil::MakeShape(F32, {256}),
-                                          &debug_options, &thunk_executor);
+  GpuExecutableVaRemapAllocator allocator(
+      "test", allocations, ShapeUtil::MakeShape(F32, {256}), &debug_options,
+      &thunk_executor, /*num_devices=*/1);
   allocator.AddProfileCandidateAllocationForTesting(0);
 
   // The parameter address changes between profiling executions.
@@ -1093,9 +1098,9 @@ TEST_F(GpuExecutableVaRemapAllocatorTest,
   DebugOptions debug_options;
   debug_options.set_xla_gpu_command_buffer_update_mode(
       DebugOptions::SKIP_PROFILED);
-  GpuExecutableVaRemapAllocator allocator("test", allocations,
-                                          ShapeUtil::MakeShape(F32, {256}),
-                                          &debug_options, &thunk_executor);
+  GpuExecutableVaRemapAllocator allocator(
+      "test", allocations, ShapeUtil::MakeShape(F32, {256}), &debug_options,
+      &thunk_executor, /*num_devices=*/1);
   allocator.AddProfileCandidateAllocationForTesting(0);
 
   // Stable address, but not an address owned by the VMM allocator, so it
@@ -1139,9 +1144,9 @@ TEST_F(GpuExecutableVaRemapAllocatorTest,
   DebugOptions debug_options;
   debug_options.set_xla_gpu_command_buffer_update_mode(
       DebugOptions::SKIP_PROFILED);
-  GpuExecutableVaRemapAllocator allocator("test", allocations,
-                                          ShapeUtil::MakeShape(F32, {256}),
-                                          &debug_options, &thunk_executor);
+  GpuExecutableVaRemapAllocator allocator(
+      "test", allocations, ShapeUtil::MakeShape(F32, {256}), &debug_options,
+      &thunk_executor, /*num_devices=*/1);
   allocator.AddProfileCandidateAllocationForTesting(0);
   allocator.AddProfileCandidateAllocationForTesting(1);
 
@@ -1185,9 +1190,9 @@ TEST_F(GpuExecutableVaRemapAllocatorTest,
   DebugOptions debug_options;
   debug_options.set_xla_gpu_command_buffer_update_mode(
       DebugOptions::SKIP_PROFILED);
-  GpuExecutableVaRemapAllocator allocator("test", allocations,
-                                          ShapeUtil::MakeShape(F32, {256}),
-                                          &debug_options, &thunk_executor);
+  GpuExecutableVaRemapAllocator allocator(
+      "test", allocations, ShapeUtil::MakeShape(F32, {256}), &debug_options,
+      &thunk_executor, /*num_devices=*/1);
   allocator.AddProfileCandidateAllocationForTesting(0);
 
   ASSERT_OK_AND_ASSIGN(
@@ -1304,9 +1309,9 @@ TEST_F(GpuExecutableVaRemapAllocatorTest,
   DebugOptions debug_options;
   debug_options.set_xla_gpu_command_buffer_update_mode(
       DebugOptions::SKIP_PROFILED);
-  GpuExecutableVaRemapAllocator allocator("test", allocations,
-                                          ShapeUtil::MakeShape(F32, {256}),
-                                          &debug_options, &thunk_executor);
+  GpuExecutableVaRemapAllocator allocator(
+      "test", allocations, ShapeUtil::MakeShape(F32, {256}), &debug_options,
+      &thunk_executor, /*num_devices=*/1);
   allocator.AddProfileCandidateAllocationForTesting(0);
 
   ASSERT_OK_AND_ASSIGN(
@@ -1382,6 +1387,53 @@ TEST_F(GpuExecutableVaRemapAllocatorTest,
   EXPECT_EQ(va_reservation->active_mapping_count(), 0);
 }
 
+TEST_F(GpuExecutableVaRemapAllocatorTest,
+       ConcurrentMultiDeviceExecutionScopes) {
+  constexpr std::array<int, 3> kOrdinals = {0, 1, 17};
+  BufferAllocation temp_alloc(/*index=*/0, /*size=*/1024, /*color=*/0);
+  std::vector<const BufferAllocation*> allocations = {&temp_alloc};
+  ThunkExecutor thunk_executor{ThunkSequence{}};
+  DebugOptions debug_options;
+  debug_options.set_xla_gpu_command_buffer_update_mode(DebugOptions::SKIP_TEMP);
+
+  std::array<NiceMock<se::MockStreamExecutor>, kOrdinals.size()> executors;
+  std::array<NiceMock<se::MockStream>, kOrdinals.size()> streams;
+  std::vector<se::DeviceAddressVmmAllocator::DeviceConfig> device_configs;
+  for (size_t i = 0; i < kOrdinals.size(); ++i) {
+    ON_CALL(executors[i], device_ordinal()).WillByDefault(Return(kOrdinals[i]));
+    ON_CALL(executors[i], SynchronizeAllActivity()).WillByDefault(Return(true));
+    ON_CALL(streams[i], parent()).WillByDefault(Return(&executors[i]));
+    device_configs.push_back({&executors[i], &streams[i]});
+  }
+  ASSERT_OK_AND_ASSIGN(auto vmm_allocator,
+                       TestVmmAllocator::Create(&platform_, device_configs));
+
+  GpuExecutableVaRemapAllocator allocator(
+      "test", allocations, ShapeUtil::MakeShape(F32, {256}), &debug_options,
+      &thunk_executor, /*num_devices=*/2);
+  allocator.AddVaRemappedAllocationForTesting(0);
+
+  absl::BlockingCounter all_scopes_held(kOrdinals.size());
+  absl::Notification release_scopes;
+  tsl::thread::ThreadPool pool(tsl::Env::Default(), "va_remap",
+                               kOrdinals.size());
+  for (size_t i = 0; i < kOrdinals.size(); ++i) {
+    pool.Schedule([&, i]() {
+      ExecutableRunOptions run_opts;
+      run_opts.set_stream(&streams[i]);
+      ServiceExecutableRunOptions service_opts(run_opts);
+      ASSERT_OK_AND_ASSIGN(
+          auto scope, allocator.CreateExecutionScope(
+                          &service_opts, vmm_allocator.get(), kOrdinals[i]));
+      EXPECT_TRUE(scope->va_remap_enabled());
+      all_scopes_held.DecrementCount();
+      release_scopes.WaitForNotification();
+    });
+  }
+  all_scopes_held.Wait();
+  release_scopes.Notify();
+}
+
 }  // namespace
 }  // namespace gpu
 }  // namespace xla

From 1526bb98f0daaf8723e36ff45dd587874d0d9f0c Mon Sep 17 00:00:00 2001
From: Ce Zheng CA <zce@google.com>
Date: Fri, 16 Oct 2026 06:40:03 -0700
Subject: [PATCH 0960/1000] Emit embedded constant buffers as LLVM arrays and
 strip debug info from embedded bitcode.

`AddBufferToLlvmModule` in `xla/util/embedded_constant_buffers.cc` previously initialized embedded constant buffers using `llvm::ConstantDataVector::get`, which emits an LLVM fixed SIMD vector type (`<N x i8>`) instead of an array (`[N x i8]`). Because `llvm::DataLayout::getTypeAllocSize` rounds fixed vector types up to the next power of two, every embedded constant buffer in `.rodata` (including C++ intrinsic bitcode blobs and AOT-compiled `tfcompile` model constants) was zero-padded up to the next power-of-two byte boundary. Switching to `llvm::ConstantDataArray::get` emits exact-sized `[N x i8]` arrays in `.rodata` with no power-of-two rounding padding.

In addition, when builds pass `-gmlt` or `-g` on the compiler command line after rule `copts`, Clang embeds DWARF debug metadata (`!DICompileUnit`, `!DISubprogram`, `!DIFile`, `!DILocation`) into the `.llvmbc` bitcode section of `cc_ir_header` targets (`xla/codegen/intrinsic/cpp/cc_to_llvm_ir.bzl`). Passing `-g0` in `cc_to_llvm_ir.bzl` and calling `llvm::StripDebugInfo` before re-serializing the module in `xla/codegen/intrinsic/cpp/embed_bitcode.cc` removes unused debug metadata from the embedded LLVM bitcode blobs, shrinking the embedded C++ intrinsic bitcode buffers by 81.2% and reducing startup string copy and `ParseEmbeddedBitcode` overhead with no change to generated JIT code.

PiperOrigin-RevId: 1000821353
---
 third_party/xla/xla/codegen/intrinsic/cpp/BUILD |  3 ++-
 .../xla/codegen/intrinsic/cpp/cc_to_llvm_ir.bzl |  1 +
 .../xla/codegen/intrinsic/cpp/embed_bitcode.cc  | 11 ++++++++++-
 .../xla/xla/util/embedded_constant_buffers.cc   |  5 ++++-
 .../xla/util/embedded_constant_buffers_test.cc  | 17 +++++++++++++++++
 5 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/BUILD b/third_party/xla/xla/codegen/intrinsic/cpp/BUILD
index 363caf3b3ca..49f239bd135 100644
--- a/third_party/xla/xla/codegen/intrinsic/cpp/BUILD
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/BUILD
@@ -68,7 +68,8 @@ cc_binary(
         "//xla/util:embedded_constant_buffers",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/types:span",
-        "@llvm-project//llvm:AArch64CodeGen",
+        "@llvm-project//llvm:AArch64CodeGen",  # build_cleaner: keep
+        "@llvm-project//llvm:BitWriter",
         "@llvm-project//llvm:Core",
         "@llvm-project//llvm:IRReader",
         "@llvm-project//llvm:Object",
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/cc_to_llvm_ir.bzl b/third_party/xla/xla/codegen/intrinsic/cpp/cc_to_llvm_ir.bzl
index 89bf66fbc7d..bbb942c3afc 100644
--- a/third_party/xla/xla/codegen/intrinsic/cpp/cc_to_llvm_ir.bzl
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/cc_to_llvm_ir.bzl
@@ -116,6 +116,7 @@ def cc_ir_header(name, src, deps = [], copts = [], **kwargs):
     compile_flags = [
         "-emit-llvm",
         "-O3",
+        "-g0",
         "-DNDEBUG",
         "-mprefer-vector-width=512",
         "-DEIGEN_VECTORIZE_GENERIC",
diff --git a/third_party/xla/xla/codegen/intrinsic/cpp/embed_bitcode.cc b/third_party/xla/xla/codegen/intrinsic/cpp/embed_bitcode.cc
index 4e179d7d452..df3b7c15f76 100644
--- a/third_party/xla/xla/codegen/intrinsic/cpp/embed_bitcode.cc
+++ b/third_party/xla/xla/codegen/intrinsic/cpp/embed_bitcode.cc
@@ -27,7 +27,10 @@ limitations under the License.
 
 #include "absl/status/statusor.h"
 #include "absl/types/span.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Bitcode/BitcodeWriter.h"
+#include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IRReader/IRReader.h"
@@ -296,7 +299,13 @@ int main(int argc, char* argv[]) {
   constant_to_embed.symbol_prefix = args->variable_name;
   std::vector<uint8_t> content_vec;
   if (has_defined_functions) {
-    content_vec.assign(content.begin(), content.end());
+    // Strip any DWARF debug info (e.g. injected by command-line --copt=-gmlt
+    // or -g flags after rule copts) and re-serialize compact bitcode.
+    llvm::StripDebugInfo(*module);
+    llvm::SmallVector<char, 0> stripped_bitcode;
+    llvm::raw_svector_ostream bitcode_stream(stripped_bitcode);
+    llvm::WriteBitcodeToFile(*module, bitcode_stream);
+    content_vec.assign(stripped_bitcode.begin(), stripped_bitcode.end());
   }
   constant_to_embed.SerializeIntoBuffer(content_vec);
 
diff --git a/third_party/xla/xla/util/embedded_constant_buffers.cc b/third_party/xla/xla/util/embedded_constant_buffers.cc
index 40e3e155ae5..1f09b41229e 100644
--- a/third_party/xla/xla/util/embedded_constant_buffers.cc
+++ b/third_party/xla/xla/util/embedded_constant_buffers.cc
@@ -84,7 +84,10 @@ static absl::Status AddBufferToLlvmModule(
 
   absl::Span<const uint8_t> buffer_contents = constant_to_embed.data();
 
-  llvm::Constant* buffer_initializer = llvm::ConstantDataVector::get(
+  // Use ConstantDataArray ([N x i8]) rather than ConstantDataVector (<N x i8>),
+  // because LLVM's DataLayout::getTypeAllocSize rounds fixed SIMD vector types
+  // up to the next power of two and zero-pads them in .rodata.
+  llvm::Constant* buffer_initializer = llvm::ConstantDataArray::get(
       module->getContext(),
       llvm::ArrayRef<uint8_t>(buffer_contents.data(), buffer_contents.size()));
 
diff --git a/third_party/xla/xla/util/embedded_constant_buffers_test.cc b/third_party/xla/xla/util/embedded_constant_buffers_test.cc
index 971bf296a01..5ea93408079 100644
--- a/third_party/xla/xla/util/embedded_constant_buffers_test.cc
+++ b/third_party/xla/xla/util/embedded_constant_buffers_test.cc
@@ -80,5 +80,22 @@ TEST_F(EmbeddedConstantBuffersTest, SerializeEmptyBuffer) {
   EXPECT_FALSE(embedded_constant_buffers.object_file_data.empty());
 }
 
+TEST_F(EmbeddedConstantBuffersTest, DoesNotPadBufferToPowerOfTwo) {
+  ConstantToEmbed constant;
+  constant.symbol_prefix = "non_power_of_two";
+  constant.SerializeIntoBuffer(std::vector<uint8_t>(2000, 0xab));
+  // With the 64-byte size header, the serialized buffer is 2064 bytes. Verify
+  // that the constant is emitted as an exact-sized array rather than an LLVM
+  // fixed vector rounded up to the next power of two (4096 bytes).
+  ASSERT_EQ(constant.data().size(), 2064);
+
+  TF_ASSERT_OK_AND_ASSIGN(
+      EmbeddedConstantBuffers embedded_constant_buffers,
+      CreateEmbeddedConstantBuffers(kTargetTripleForHost,
+                                    absl::MakeSpan(&constant, 1)));
+  EXPECT_GT(embedded_constant_buffers.object_file_data.size(), 2064);
+  EXPECT_LT(embedded_constant_buffers.object_file_data.size(), 4096);
+}
+
 }  // namespace
 }  // namespace xla

From b7fe3cad456b59c05a033fd58d20bbec6f3cb871 Mon Sep 17 00:00:00 2001
From: Levon Ter-Grigoryan <patrios@google.com>
Date: Fri, 16 Oct 2026 07:09:19 -0700
Subject: [PATCH 0961/1000] [XLA:GPU] Unify multi-process tests logic.

PiperOrigin-RevId: 1000833738
---
 third_party/xla/xla/backends/gpu/tests/BUILD  |  81 +++--
 .../all_reduce_multi_process_e2e_test.cc      | 305 +++---------------
 .../gpu/tests/multi_process_gpu_test_base.cc  | 255 +++++++++++++++
 .../gpu/tests/multi_process_gpu_test_base.h   |  97 ++++++
 .../gpu/tests/multi_process_gpu_test_main.cc  | 296 +++++++++++++++++
 5 files changed, 762 insertions(+), 272 deletions(-)
 create mode 100644 third_party/xla/xla/backends/gpu/tests/multi_process_gpu_test_base.cc
 create mode 100644 third_party/xla/xla/backends/gpu/tests/multi_process_gpu_test_base.h
 create mode 100644 third_party/xla/xla/backends/gpu/tests/multi_process_gpu_test_main.cc

diff --git a/third_party/xla/xla/backends/gpu/tests/BUILD b/third_party/xla/xla/backends/gpu/tests/BUILD
index f4409759688..1462837e95a 100644
--- a/third_party/xla/xla/backends/gpu/tests/BUILD
+++ b/third_party/xla/xla/backends/gpu/tests/BUILD
@@ -1745,6 +1745,66 @@ xla_test(
     ],
 )
 
+cc_library(
+    name = "multi_process_gpu_test_base",
+    testonly = True,
+    srcs = ["multi_process_gpu_test_base.cc"],
+    hdrs = ["multi_process_gpu_test_base.h"],
+    deps = [
+        "//xla/pjrt:pjrt_client",
+        "//xla/pjrt/distributed",
+        "//xla/pjrt/distributed:client",
+        "//xla/pjrt/distributed:key_value_store_interface",
+        "//xla/pjrt/plugin/xla_gpu:xla_gpu_allocator_config",
+        "//xla/pjrt/plugin/xla_gpu:xla_gpu_client_options",
+        "//xla/pjrt/plugin/xla_gpu:xla_gpu_pjrt_client",
+        "@com_google_absl//absl/base:no_destructor",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/time",
+        "@com_google_googletest//:gtest_for_library",
+    ],
+)
+
+cc_library(
+    name = "multi_process_gpu_test_main",
+    testonly = True,
+    srcs = ["multi_process_gpu_test_main.cc"],
+    deps = [
+        ":multi_process_gpu_test_base",
+        "//xla:debug_options_flags",
+        "//xla/pjrt/distributed",
+        "//xla/pjrt/distributed:service",
+        "//xla/service:platform_util",
+        "//xla/stream_executor:platform",
+        "//xla/tsl/platform:env",
+        "//xla/tsl/platform:subprocess",
+        "//xla/tsl/platform:test",
+        "//xla/tsl/util:command_line_flags",
+        "@com_google_absl//absl/base:log_severity",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/log:globals",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/synchronization",
+        "@com_google_absl//absl/time",
+        "@com_google_googletest//:gtest_for_library",
+        "@tsl//tsl/platform:path",
+    ],
+    alwayslink = True,
+)
+
 xla_test(
     name = "all_gather_e2e_multiprocess_test",
     srcs = ["all_gather_e2e_multiprocess_test.cc"],
@@ -1814,11 +1874,12 @@ xla_test(
         "gpu",
     ],
     deps = [
+        ":multi_process_gpu_test_base",
+        ":multi_process_gpu_test_main",
         "//xla:debug_options_flags",
         "//xla:literal",
         "//xla:literal_util",
         "//xla:shape_util",
-        "//xla:status_macros",
         "//xla:xla_data_proto_cc",
         "//xla:xla_proto_cc",
         "//xla/hlo/builder:xla_computation",
@@ -1827,14 +1888,7 @@ xla_test(
         "//xla/pjrt:pjrt_client",
         "//xla/pjrt:pjrt_compiler",
         "//xla/pjrt:pjrt_executable",
-        "//xla/pjrt/distributed",
-        "//xla/pjrt/distributed:client",
-        "//xla/pjrt/distributed:service",
-        "//xla/pjrt/plugin/xla_gpu:xla_gpu_allocator_config",
-        "//xla/pjrt/plugin/xla_gpu:xla_gpu_client_options",
-        "//xla/pjrt/plugin/xla_gpu:xla_gpu_pjrt_client",
         "//xla/service:device_assignment",
-        "//xla/service:gpu_topology",
         "//xla/service:platform_util",
         "//xla/service/gpu:backend_configs_cc",
         "//xla/stream_executor:device_description",
@@ -1842,19 +1896,10 @@ xla_test(
         "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor/cuda:cuda_compute_capability",
         "//xla/tests:literal_test_util",
-        "//xla/tsl/platform:env",
-        "//xla/tsl/platform:subprocess",
-        "//xla/tsl/platform:test",
-        "//xla/tsl/util:command_line_flags",
         "@com_google_absl//absl/log",
-        "@com_google_absl//absl/log:check",
-        "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/strings:string_view",
-        "@com_google_absl//absl/time",
         "@com_google_googletest//:gtest",
-        "@tsl//tsl/platform:path",
     ],
 )
 
diff --git a/third_party/xla/xla/backends/gpu/tests/all_reduce_multi_process_e2e_test.cc b/third_party/xla/xla/backends/gpu/tests/all_reduce_multi_process_e2e_test.cc
index 02f19c04d2b..68469be5c84 100644
--- a/third_party/xla/xla/backends/gpu/tests/all_reduce_multi_process_e2e_test.cc
+++ b/third_party/xla/xla/backends/gpu/tests/all_reduce_multi_process_e2e_test.cc
@@ -13,21 +13,15 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include <cstdlib>
 #include <memory>
-#include <optional>
-#include <string>
-#include <utility>
 #include <vector>
 
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
-#include "absl/log/check.h"
 #include "absl/log/log.h"
-#include "absl/status/status.h"
-#include "absl/status/statusor.h"
-#include "absl/strings/str_format.h"
+#include "absl/status/status_matchers.h"
 #include "absl/strings/string_view.h"
-#include "absl/time/time.h"
+#include "xla/backends/gpu/tests/multi_process_gpu_test_base.h"
 #include "xla/debug_options_flags.h"
 #include "xla/hlo/builder/xla_computation.h"
 #include "xla/hlo/ir/hlo_computation.h"
@@ -37,120 +31,42 @@ limitations under the License.
 #include "xla/hlo/parser/hlo_parser.h"
 #include "xla/literal.h"
 #include "xla/literal_util.h"
-#include "xla/pjrt/distributed/client.h"
-#include "xla/pjrt/distributed/distributed.h"
-#include "xla/pjrt/distributed/service.h"
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/pjrt/pjrt_compiler.h"
 #include "xla/pjrt/pjrt_executable.h"
-#include "xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config.h"
-#include "xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h"
-#include "xla/pjrt/plugin/xla_gpu/xla_gpu_pjrt_client.h"
 #include "xla/service/device_assignment.h"
 #include "xla/service/gpu/backend_configs.pb.h"
-#include "xla/service/gpu_topology.h"
 #include "xla/service/platform_util.h"
 #include "xla/shape_util.h"
-#include "xla/status_macros.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/tests/literal_test_util.h"
-#include "xla/tsl/platform/env.h"
-#include "xla/tsl/platform/subprocess.h"
-#include "xla/tsl/platform/test.h"
-#include "xla/tsl/util/command_line_flags.h"
 #include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
-#include "tsl/platform/path.h"
 
 namespace xla {
 namespace {
 
-inline constexpr size_t kMB = 1024LL * 1024LL;
-constexpr int kNumNodes = 2;
-static const char* test_binary_name;
+using AllReduceMultiProcessE2ETest = MultiProcessGpuTest;
 
-struct MultiProcessGpuClientSetup {
-  std::unique_ptr<xla::DistributedRuntimeService> service;
-  std::unique_ptr<PjRtClient> client;
-};
-
-absl::StatusOr<MultiProcessGpuClientSetup> SetUpMultiProcessGpuClient(
-    int rank_id, int num_nodes, int port, absl::string_view log_prefix) {
-  MultiProcessGpuClientSetup prepared_test;
-  std::string coordinator_address = absl::StrFormat("127.0.0.1:%d", port);
-
-  // Rank 0 creates the coordination service.
-  if (rank_id == 0) {
-    LOG(INFO) << log_prefix << ": creating coordination service on "
-              << coordinator_address;
-    xla::CoordinationServiceImpl::Options service_options;
-    service_options.num_nodes = num_nodes;
-    ABSL_ASSIGN_OR_RETURN(prepared_test.service,
-                     xla::GetDistributedRuntimeService(
-                         absl::StrFormat("[::]:%d", port), service_options));
-    LOG(INFO) << log_prefix << ": created coordination service";
-  }
-
-  // Connect to the coordination service.
-  xla::DistributedRuntimeClient::Options distributed_options;
-  distributed_options.node_id = rank_id;
-  distributed_options.init_timeout = absl::Seconds(120);
-  auto distributed_client =
-      GetDistributedRuntimeClient(coordinator_address, distributed_options);
-
-  LOG(INFO) << log_prefix << ": connecting distributed client";
-  ABSL_RETURN_IF_ERROR(distributed_client->Connect());
-  LOG(INFO) << log_prefix << ": distributed client connected";
-
-  // Create the GPU client with a single addressable device per process.
-  GpuClientOptions options;
-  options.node_id = rank_id;
-  options.num_nodes = num_nodes;
-  options.allowed_devices = {rank_id};
-  options.kv_store =
-      GetDistributedKeyValueStore(distributed_client, /*key_prefix=*/"gpu:");
-  options.distributed_client = distributed_client;
-  options.allocator_config.kind = xla::GpuAllocatorConfig::Kind::kBFC;
-  options.allocator_config.gpu_system_memory_size = 32 * kMB;
-  options.allocator_config.collective_memory_size = 0;
-  options.use_tfrt_gpu_client = true;
-
-  LOG(INFO) << log_prefix << ": creating PjRtClient";
-  ABSL_ASSIGN_OR_RETURN(prepared_test.client, GetXlaPjrtGpuClient(options));
-  LOG(INFO) << log_prefix << ": PjRtClient created";
-
-  return prepared_test;
-}
-
-absl::Status AllReduceMultiProcessTestBody(int node_id, int port) {
-  std::string log_prefix = absl::StrFormat("rank_%d", node_id);
-  ABSL_ASSIGN_OR_RETURN(se::Platform * platform, PlatformUtil::GetPlatform("gpu"));
-  ABSL_ASSIGN_OR_RETURN(se::StreamExecutor * executor,
-                   platform->ExecutorForDevice(node_id));
-  const auto& desc = executor->GetDeviceDescription();
+TEST_F(AllReduceMultiProcessE2ETest, OneShotAllReduce2Processes) {
+  ASSERT_OK_AND_ASSIGN(se::Platform * platform,
+                       PlatformUtil::GetPlatform("gpu"));
+  ASSERT_OK_AND_ASSIGN(se::StreamExecutor * executor,
+                       platform->ExecutorForDevice(rank()));
+  const se::DeviceDescription& desc = executor->GetDeviceDescription();
   if (desc.gpu_compute_capability().IsCuda() && !desc.gpu_compute_capability()
                                                      .cuda_compute_capability()
                                                      ->IsAtLeastAmpere()) {
-    LOG(INFO) << log_prefix
-              << ": skipping one-shot all-reduce test, requires Ampere+ for "
-                 "CUDA.";
-    return absl::OkStatus();
+    GTEST_SKIP() << "One-shot all-reduce requires Ampere+ for CUDA.";
   }
 
-  ABSL_ASSIGN_OR_RETURN(
-      MultiProcessGpuClientSetup setup,
-      SetUpMultiProcessGpuClient(node_id, kNumNodes, port, log_prefix));
-  std::unique_ptr<PjRtClient> client = std::move(setup.client);
+  ASSERT_EQ(client()->addressable_device_count(), 1);
+  ASSERT_EQ(client()->device_count(), num_ranks());
 
-  TF_RET_CHECK(client->addressable_device_count() == 1)
-      << "Expected exactly 1 local addressable device per process.";
-  TF_RET_CHECK(client->device_count() == kNumNodes)
-      << "Expected " << kNumNodes << " global devices.";
-
-  constexpr absl::string_view kModuleStr = R"(
+  constexpr absl::string_view kModuleStr = R"hlo(
     HloModule test
 
     apply_op {
@@ -162,33 +78,24 @@ absl::Status AllReduceMultiProcessTestBody(int node_id, int port) {
     ENTRY test_computation {
       param_0 = f32[128] parameter(0)
       ROOT all-reduce = f32[128] all-reduce(param_0), to_apply=apply_op,
-      replica_groups={{0,1}}
+        replica_groups={{0,1}}
     }
-  )";
+  )hlo";
 
-  ABSL_ASSIGN_OR_RETURN(auto hlo_module,
-                   ParseAndReturnUnverifiedModule(kModuleStr, /*config=*/{}));
-  xla::XlaComputation computation(hlo_module->ToProto());
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> hlo_module,
+                       ParseAndReturnUnverifiedModule(kModuleStr));
+  XlaComputation computation(hlo_module->ToProto());
 
-  xla::CompileOptions compile_options;
-  compile_options.executable_build_options.set_num_replicas(kNumNodes);
+  CompileOptions compile_options;
+  compile_options.executable_build_options.set_num_replicas(num_ranks());
   compile_options.executable_build_options.set_num_partitions(1);
-  DeviceAssignment device_assignment(kNumNodes, 1);
-  device_assignment(0, 0) = 0;
-  device_assignment(1, 0) = 1;
+  DeviceAssignment device_assignment(num_ranks(), 1);
+  for (int i = 0; i < num_ranks(); ++i) {
+    device_assignment(i, 0) = i;
+  }
   compile_options.executable_build_options.set_device_assignment(
       device_assignment);
 
-  GpuTopology gpu_topology(
-      /*platform_version=*/"",
-      /*num_partitions=*/1,
-      /*num_hosts_per_partition=*/1,
-      /*num_devices_per_host=*/kNumNodes,
-      /*gpu_target_config=*/std::nullopt,
-      /*host_target_machine_options=*/std::nullopt,
-      /*num_devices_per_process=*/kNumNodes);
-  compile_options.executable_build_options.set_gpu_topology(gpu_topology);
-
   DebugOptions debug_options = GetDebugOptionsFromFlags();
   debug_options.set_xla_gpu_autotune_level(0);
   debug_options.add_xla_gpu_unsupported_use_cross_host_one_shot_kernel(
@@ -196,15 +103,13 @@ absl::Status AllReduceMultiProcessTestBody(int node_id, int port) {
   *compile_options.executable_build_options.mutable_debug_options() =
       debug_options;
 
-  LOG(INFO) << log_prefix << ": compiling HLO module with one-shot all-reduce";
-  ABSL_ASSIGN_OR_RETURN(std::unique_ptr<PjRtLoadedExecutable> executable,
-                   client->CompileAndLoad(computation, compile_options));
-  LOG(INFO) << log_prefix << ": compilation succeeded";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtLoadedExecutable> executable,
+                       client()->CompileAndLoad(computation, compile_options));
 
   // Inspect the collective kernel strategy selected in the optimized HLO.
-  ABSL_ASSIGN_OR_RETURN(auto hlo_modules,
-                   executable->GetExecutable()->GetHloModules());
-  TF_RET_CHECK(!hlo_modules.empty());
+  ASSERT_OK_AND_ASSIGN(auto hlo_modules,
+                       executable->GetExecutable()->GetHloModules());
+  ASSERT_FALSE(hlo_modules.empty());
   const HloModule* optimized_module = hlo_modules.front().get();
   bool found_all_reduce = false;
   for (const HloComputation* comp : optimized_module->computations()) {
@@ -212,156 +117,48 @@ absl::Status AllReduceMultiProcessTestBody(int node_id, int port) {
       if (instr->opcode() == HloOpcode::kAllReduce ||
           instr->opcode() == HloOpcode::kAllReduceStart) {
         found_all_reduce = true;
-        auto gpu_config = instr->backend_config<gpu::GpuBackendConfig>();
-        TF_RET_CHECK(gpu_config.ok())
-            << "Failed to get GpuBackendConfig for " << instr->name() << ": "
-            << gpu_config.status();
+        ASSERT_OK_AND_ASSIGN(auto gpu_config,
+                             instr->backend_config<gpu::GpuBackendConfig>());
         LOG(INFO)
-            << log_prefix << ": AllReduce instruction " << instr->name()
+            << "rank_" << rank() << ": AllReduce instruction " << instr->name()
             << " collective kernel strategy: "
             << gpu::CollectiveBackendConfig::CollectiveKernelStrategy_Name(
-                   gpu_config->collective_backend_config().kernel_strategy());
-        TF_RET_CHECK(
-            gpu_config->collective_backend_config().kernel_strategy() ==
-            gpu::CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT)
-            << "Expected AllReduce to use KERNEL_STRATEGY_TRITON_ONE_SHOT, but "
-               "got: "
-            << gpu::CollectiveBackendConfig::CollectiveKernelStrategy_Name(
-                   gpu_config->collective_backend_config().kernel_strategy());
+                   gpu_config.collective_backend_config().kernel_strategy());
+        EXPECT_EQ(
+            gpu_config.collective_backend_config().kernel_strategy(),
+            gpu::CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT);
       }
     }
   }
-  TF_RET_CHECK(found_all_reduce)
+  EXPECT_TRUE(found_all_reduce)
       << "Expected to find an AllReduce instruction in optimized HLO.";
 
   // Prepare input literal: rank 0 provides 10.0f, rank 1 provides 20.0f.
-  const float input_val = (node_id == 0) ? 10.0f : 20.0f;
+  const float input_val = (rank() == 0) ? 10.0f : 20.0f;
   Literal input_literal =
       LiteralUtil::CreateR1<float>(std::vector<float>(128, input_val));
 
-  ABSL_ASSIGN_OR_RETURN(auto* memory_space,
-                   client->addressable_devices()[0]->default_memory_space());
-  ABSL_ASSIGN_OR_RETURN(auto input_buffer,
-                   client->BufferFromHostLiteral(input_literal, memory_space));
+  ASSERT_OK_AND_ASSIGN(
+      PjRtMemorySpace * memory_space,
+      client()->addressable_devices()[0]->default_memory_space());
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<PjRtBuffer> input_buffer,
+      client()->BufferFromHostLiteral(input_literal, memory_space));
 
   std::vector<std::vector<PjRtBuffer*>> input_ptrs = {{input_buffer.get()}};
-  LOG(INFO) << log_prefix << ": executing one-shot all-reduce";
-  ABSL_ASSIGN_OR_RETURN(auto results,
-                   executable->Execute(input_ptrs, ExecuteOptions()));
-  LOG(INFO) << log_prefix << ": execution finished";
+  ASSERT_OK_AND_ASSIGN(auto results,
+                       executable->Execute(input_ptrs, ExecuteOptions()));
 
-  TF_RET_CHECK(results.size() == 1 && results[0].size() == 1);
+  ASSERT_EQ(results.size(), 1);
+  ASSERT_EQ(results[0].size(), 1);
   Literal result_literal(ShapeUtil::MakeShape(F32, {128}));
-  ABSL_RETURN_IF_ERROR(results[0][0]->ToLiteralSync(&result_literal));
+  ASSERT_OK(results[0][0]->ToLiteralSync(&result_literal));
 
   // Expected result is 10.0f + 20.0f = 30.0f across all ranks.
   Literal expected_literal =
       LiteralUtil::CreateR1<float>(std::vector<float>(128, 30.0f));
-  if (!LiteralTestUtil::Equal(expected_literal, result_literal)) {
-    return absl::InternalError(absl::StrFormat(
-        "Result literal %s does not match expected %s on rank %d",
-        result_literal.ToString(), expected_literal.ToString(), node_id));
-  }
-
-  LOG(INFO) << log_prefix << ": verified result successfully";
-  return absl::OkStatus();
-}
-
-TEST(AllReduceMultiProcessE2ETest, OneShotAllReduce2Processes) {
-  absl::StatusOr<se::Platform*> platform = PlatformUtil::GetPlatform("gpu");
-  if (!platform.ok() || (*platform)->VisibleDeviceCount() < kNumNodes) {
-    GTEST_SKIP() << "Test requires at least " << kNumNodes
-                 << " GPU devices, but found "
-                 << (platform.ok() ? (*platform)->VisibleDeviceCount() : 0);
-  }
-
-  // Re-pass XLA_FLAGS environment variable to child processes.
-  const char* xla_flags = std::getenv("XLA_FLAGS");
-  if (xla_flags != nullptr) {
-    tsl::setenv("XLA_FLAGS", xla_flags, /*overwrite=*/true);
-  }
-
-  int port = tsl::testing::PickUnusedPortOrDie();
-  tsl::SubProcess child[kNumNodes];
-
-  for (int node_id = 0; node_id < kNumNodes; ++node_id) {
-    std::vector<std::string> argv = {
-        test_binary_name,
-        "--test_to_run=AllReduceMultiProcessHelper",
-        absl::StrFormat("--node_id=%d", node_id),
-        absl::StrFormat("--port=%d", port),
-        "--alsologtostderr",
-        "--vmodule=gpu_executable=1,thunk_executor=1,all_reduce_thunk=5,"
-        "collective_kernel_thunk=5,collective_memory=5",
-    };
-    child[node_id].SetProgram(test_binary_name, argv);
-    child[node_id].SetChannelAction(tsl::CHAN_STDOUT, tsl::ACTION_PIPE);
-    child[node_id].SetChannelAction(tsl::CHAN_STDERR, tsl::ACTION_PIPE);
-    ASSERT_TRUE(child[node_id].Start()) << "Failed to start node " << node_id;
-  }
-
-  for (int node_id = 0; node_id < kNumNodes; ++node_id) {
-    std::string stdout_str, stderr_str;
-    int status = child[node_id].Communicate(nullptr, &stdout_str, &stderr_str);
-
-    const char* undeclared_outputs_dir =
-        std::getenv("TEST_UNDECLARED_OUTPUTS_DIR");
-    if (undeclared_outputs_dir != nullptr &&
-        undeclared_outputs_dir[0] != '\0') {
-      std::string stderr_file = tsl::io::JoinPath(
-          undeclared_outputs_dir,
-          absl::StrFormat("subprocess_node_%d_stderr.log", node_id));
-      absl::Status write_status =
-          tsl::WriteStringToFile(tsl::Env::Default(), stderr_file, stderr_str);
-      if (!write_status.ok()) {
-        LOG(WARNING) << "Failed to write stderr to " << stderr_file << ": "
-                     << write_status;
-      }
-    }
-
-    EXPECT_EQ(status, 0) << "node " << node_id << " failed with status "
-                         << status << "\nstdout:\n"
-                         << stdout_str << "\nstderr:\n"
-                         << stderr_str;
-  }
+  EXPECT_TRUE(LiteralTestUtil::Equal(expected_literal, result_literal));
 }
 
 }  // namespace
 }  // namespace xla
-
-int main(int argc, char* argv[]) {
-  std::string test_to_run;
-  int node_id = -1;
-  int port = -1;
-  xla::test_binary_name = argv[0];
-
-  std::vector<tsl::Flag> flag_list = {
-      tsl::Flag("test_to_run", &test_to_run,
-                "The test to run in the child process."),
-      tsl::Flag("node_id", &node_id,
-                "The node id (rank) for the child process."),
-      tsl::Flag("port", &port, "The coordinator port for distributed runtime."),
-  };
-
-  xla::AppendDebugOptionsFlags(&flag_list);
-  tsl::Flags::Parse(&argc, argv, flag_list);
-  testing::InitGoogleTest(&argc, argv);
-
-  if (test_to_run.empty()) {
-    return RUN_ALL_TESTS();
-  }
-
-  absl::Status result = absl::OkStatus();
-  if (test_to_run == "AllReduceMultiProcessHelper") {
-    result = xla::AllReduceMultiProcessTestBody(node_id, port);
-  } else {
-    result = absl::InvalidArgumentError(absl::StrFormat(
-        "Unrecognized multiprocess test name: %s", test_to_run));
-  }
-
-  if (!result.ok()) {
-    LOG(ERROR) << "Child process (node_id " << node_id
-               << ") failed: " << result;
-  }
-  return result.raw_code();
-}
diff --git a/third_party/xla/xla/backends/gpu/tests/multi_process_gpu_test_base.cc b/third_party/xla/xla/backends/gpu/tests/multi_process_gpu_test_base.cc
new file mode 100644
index 00000000000..edd1b9966cb
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/tests/multi_process_gpu_test_base.cc
@@ -0,0 +1,255 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/gpu/tests/multi_process_gpu_test_base.h"
+
+#include <atomic>
+#include <cstddef>
+#include <cstdlib>
+#include <memory>
+#include <optional>
+#include <set>
+#include <string>
+#include <utility>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/base/no_destructor.h"
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
+#include "absl/time/time.h"
+#include "xla/pjrt/distributed/client.h"
+#include "xla/pjrt/distributed/distributed.h"
+#include "xla/pjrt/distributed/key_value_store_interface.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config.h"
+#include "xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h"
+#include "xla/pjrt/plugin/xla_gpu/xla_gpu_pjrt_client.h"
+
+namespace xla {
+namespace {
+
+inline constexpr size_t kDefaultGpuSystemMemoryBytes = 128ULL * 1024 * 1024;
+
+struct WorkerState {
+  absl::Mutex mu;
+  int worker_id = -1;
+  int num_workers = 2;
+  int gpus_per_worker = 1;
+  std::string coordinator_address;
+  std::shared_ptr<DistributedRuntimeClient> distributed_client;
+  std::shared_ptr<KeyValueStoreInterface> kv_store;
+  std::unique_ptr<PjRtClient> shared_client;
+  std::atomic<int> custom_client_counter{0};
+  std::atomic<int> next_test_invocation_id{0};
+  std::atomic<int> current_test_invocation_id{0};
+};
+
+WorkerState& GetWorkerState() {
+  static absl::NoDestructor<WorkerState> state;
+  return *state;
+}
+
+}  // namespace
+
+void ConfigureMultiProcessGpuParent(int num_workers, int gpus_per_worker) {
+  WorkerState& state = GetWorkerState();
+  absl::MutexLock lock(state.mu);
+  state.worker_id = -1;
+  state.num_workers = num_workers;
+  state.gpus_per_worker = gpus_per_worker;
+}
+
+void InitializeMultiProcessGpuWorker(int worker_id, int num_workers,
+                                     int gpus_per_worker,
+                                     absl::string_view coordinator_address) {
+  WorkerState& state = GetWorkerState();
+  absl::MutexLock lock(state.mu);
+  state.worker_id = worker_id;
+  state.num_workers = num_workers;
+  state.gpus_per_worker = gpus_per_worker;
+  state.coordinator_address = std::string(coordinator_address);
+
+  DistributedRuntimeClient::Options distributed_options;
+  distributed_options.node_id = worker_id;
+  distributed_options.init_timeout = absl::Seconds(120);
+  distributed_options.shutdown_timeout = absl::Seconds(10);
+  state.distributed_client = GetDistributedRuntimeClient(
+      state.coordinator_address, distributed_options);
+  CHECK_OK(state.distributed_client->Connect())
+      << "rank_" << worker_id << ": failed to connect DistributedRuntimeClient";
+  state.kv_store = GetDistributedKeyValueStore(state.distributed_client,
+                                               /*key_prefix=*/"shared_gpu:");
+}
+
+void ShutDownMultiProcessGpuWorker() {
+  WorkerState& state = GetWorkerState();
+  absl::MutexLock lock(state.mu);
+  state.shared_client.reset();
+  state.kv_store.reset();
+  if (state.distributed_client != nullptr) {
+    state.distributed_client->Shutdown().IgnoreError();
+    state.distributed_client.reset();
+  }
+}
+
+void MultiProcessGpuTest::SetUp() {
+  WorkerState& state = GetWorkerState();
+  state.current_test_invocation_id.store(
+      state.next_test_invocation_id.fetch_add(1));
+  if (rank() < 0) {
+    GTEST_SKIP() << "Test requires at least " << num_ranks() * gpus_per_worker()
+                 << " GPU devices.";
+  }
+  ASSERT_OK(distributed_client()->WaitAtBarrier(
+      absl::StrCat(CurrentTestId(), "_start"), absl::Seconds(120),
+      std::nullopt));
+}
+
+void MultiProcessGpuTest::TearDown() {
+  if (rank() >= 0 && distributed_client() != nullptr) {
+    const bool has_failure = HasFailure();
+    const bool has_fatal_failure = HasFatalFailure();
+    const std::string fail_key = absl::StrCat(CurrentTestId(), "_failed");
+    if (has_failure) {
+      distributed_client()
+          ->KeyValueSet(fail_key, absl::StrCat(rank()),
+                        /*allow_overwrite=*/true)
+          .IgnoreError();
+    }
+    absl::Status barrier_status = distributed_client()->WaitAtBarrier(
+        absl::StrCat(CurrentTestId(), "_end"),
+        has_fatal_failure ? absl::Seconds(5) : absl::Seconds(120),
+        std::nullopt);
+    if (!barrier_status.ok() && has_fatal_failure) {
+      LOG(ERROR) << "rank_" << rank()
+                 << ": TearDown barrier timed out after fatal test failure ("
+                 << barrier_status
+                 << "); exiting immediately so orchestrator can terminate "
+                    "stuck peers.";
+      std::_Exit(1);
+    }
+    EXPECT_OK(barrier_status);
+    if (!has_failure && distributed_client()->KeyValueTryGet(fail_key).ok()) {
+      ADD_FAILURE() << "Peer rank failed in multi-process test.";
+    }
+  }
+}
+
+int MultiProcessGpuTest::rank() const { return GetWorkerState().worker_id; }
+
+int MultiProcessGpuTest::num_ranks() const {
+  return GetWorkerState().num_workers;
+}
+
+int MultiProcessGpuTest::gpus_per_worker() const {
+  return GetWorkerState().gpus_per_worker;
+}
+
+std::shared_ptr<DistributedRuntimeClient>
+MultiProcessGpuTest::distributed_client() const {
+  return GetWorkerState().distributed_client;
+}
+
+std::shared_ptr<KeyValueStoreInterface> MultiProcessGpuTest::kv_store() const {
+  return GetWorkerState().kv_store;
+}
+
+GpuClientOptions MultiProcessGpuTest::DefaultGpuClientOptions() const {
+  GpuClientOptions options;
+  options.node_id = rank();
+  options.num_nodes = num_ranks();
+  std::set<int> allowed_devices;
+  for (int i = 0; i < gpus_per_worker(); ++i) {
+    allowed_devices.insert(rank() * gpus_per_worker() + i);
+  }
+  options.allowed_devices = std::move(allowed_devices);
+  options.distributed_client = distributed_client();
+  options.kv_store = kv_store();
+  options.allocator_config.kind = GpuAllocatorConfig::Kind::kBFC;
+  options.allocator_config.gpu_system_memory_size =
+      kDefaultGpuSystemMemoryBytes;
+  options.allocator_config.preallocate = false;
+  options.allocator_config.collective_memory_size = 0;
+  return options;
+}
+
+PjRtClient* MultiProcessGpuTest::client() const {
+  WorkerState& state = GetWorkerState();
+  absl::MutexLock lock(state.mu);
+  if (state.shared_client == nullptr) {
+    absl::StatusOr<std::unique_ptr<PjRtClient>> shared_client =
+        GetXlaPjrtGpuClient(DefaultGpuClientOptions());
+    CHECK_OK(shared_client.status())
+        << "rank_" << state.worker_id << ": failed to create shared PjRtClient";
+    state.shared_client = *std::move(shared_client);
+  }
+  return state.shared_client.get();
+}
+
+absl::StatusOr<std::unique_ptr<PjRtClient>>
+MultiProcessGpuTest::CreateCustomPjRtClient(GpuClientOptions options) const {
+  WorkerState& state = GetWorkerState();
+  options.node_id = rank();
+  options.num_nodes = num_ranks();
+  if (!options.allowed_devices.has_value()) {
+    std::set<int> allowed_devices;
+    for (int i = 0; i < gpus_per_worker(); ++i) {
+      allowed_devices.insert(rank() * gpus_per_worker() + i);
+    }
+    options.allowed_devices = std::move(allowed_devices);
+  }
+  if (options.distributed_client == nullptr) {
+    options.distributed_client = distributed_client();
+  }
+  if (options.kv_store == nullptr) {
+    int custom_id = state.custom_client_counter.fetch_add(1);
+    options.kv_store = CreateScopedKeyValueStore(
+        absl::StrFormat("custom_client_%d", custom_id));
+  }
+  if (!options.allocator_config.gpu_system_memory_size.has_value()) {
+    options.allocator_config.gpu_system_memory_size =
+        kDefaultGpuSystemMemoryBytes;
+    options.allocator_config.preallocate = false;
+  }
+  return GetXlaPjrtGpuClient(std::move(options));
+}
+
+std::shared_ptr<KeyValueStoreInterface>
+MultiProcessGpuTest::CreateScopedKeyValueStore(absl::string_view scope) const {
+  return GetDistributedKeyValueStore(
+      distributed_client(),
+      absl::StrFormat("test_%s_%s:", CurrentTestId(), scope));
+}
+
+std::string MultiProcessGpuTest::CurrentTestId() const {
+  const int run_id = GetWorkerState().current_test_invocation_id.load();
+  const ::testing::TestInfo* test_info =
+      ::testing::UnitTest::GetInstance()->current_test_info();
+  if (test_info == nullptr) {
+    return absl::StrFormat("unknown_test_run_%d", run_id);
+  }
+  return absl::StrFormat("%s.%s_run_%d", test_info->test_suite_name(),
+                         test_info->name(), run_id);
+}
+
+}  // namespace xla
diff --git a/third_party/xla/xla/backends/gpu/tests/multi_process_gpu_test_base.h b/third_party/xla/xla/backends/gpu/tests/multi_process_gpu_test_base.h
new file mode 100644
index 00000000000..521aab143f8
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/tests/multi_process_gpu_test_base.h
@@ -0,0 +1,97 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_BACKENDS_GPU_TESTS_MULTI_PROCESS_GPU_TEST_BASE_H_
+#define XLA_BACKENDS_GPU_TESTS_MULTI_PROCESS_GPU_TEST_BASE_H_
+
+#include <memory>
+#include <string>
+
+#include <gtest/gtest.h>
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "xla/pjrt/distributed/client.h"
+#include "xla/pjrt/distributed/key_value_store_interface.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h"
+
+namespace xla {
+
+// Internal hooks called by multi_process_gpu_test_main.
+void ConfigureMultiProcessGpuParent(int num_workers, int gpus_per_worker);
+void InitializeMultiProcessGpuWorker(int worker_id, int num_workers,
+                                     int gpus_per_worker,
+                                     absl::string_view coordinator_address);
+void ShutDownMultiProcessGpuWorker();
+
+// Base fixture for multi-process GPU tests.
+//
+// When linked with `multi_process_gpu_test_main`, the test binary acts as an
+// orchestrator that starts a `DistributedRuntimeService` and spawns N worker
+// subprocesses (one per GPU rank by default). Each worker executes the same
+// GoogleTest suite in lockstep, synchronized via barriers in `SetUp()` and
+// `TearDown()`.
+class MultiProcessGpuTest : public ::testing::Test {
+ protected:
+  void SetUp() override;
+  void TearDown() override;
+
+  // This worker's rank in [0, num_ranks()). Returns -1 if running in the
+  // parent process (e.g. when skipping due to insufficient GPUs).
+  int rank() const;
+  int node_id() const { return rank(); }
+
+  // Total number of worker processes.
+  int num_ranks() const;
+  int num_nodes() const { return num_ranks(); }
+
+  // Number of GPUs assigned to each worker process.
+  int gpus_per_worker() const;
+
+  // Returns a process-wide shared PjRtClient initialized for this worker's
+  // rank and device assignment. Lazily created on first call.
+  PjRtClient* client() const;
+
+  // Returns the process-wide DistributedRuntimeClient connected to the parent
+  // orchestrator's coordination service.
+  std::shared_ptr<DistributedRuntimeClient> distributed_client() const;
+
+  // Returns the distributed KeyValueStore used by the shared PjRtClient.
+  std::shared_ptr<KeyValueStoreInterface> kv_store() const;
+
+  // Returns default GpuClientOptions configured for this worker rank.
+  GpuClientOptions DefaultGpuClientOptions() const;
+
+  // Creates a new PjRtClient with custom options, automatically populating
+  // `node_id`, `num_nodes`, `allowed_devices`, `distributed_client`, and an
+  // isolated `kv_store` if not already set on `options`.
+  absl::StatusOr<std::unique_ptr<PjRtClient>> CreateCustomPjRtClient(
+      GpuClientOptions options = GpuClientOptions()) const;
+
+  // Creates a distributed KeyValueStore scoped to the current test and `scope`.
+  std::shared_ptr<KeyValueStoreInterface> CreateScopedKeyValueStore(
+      absl::string_view scope) const;
+
+  // Returns a unique identifier string for the currently running test case.
+  std::string CurrentTestId() const;
+};
+
+template <typename T>
+class MultiProcessGpuTestWithParam : public MultiProcessGpuTest,
+                                     public ::testing::WithParamInterface<T> {};
+
+}  // namespace xla
+
+#endif  // XLA_BACKENDS_GPU_TESTS_MULTI_PROCESS_GPU_TEST_BASE_H_
diff --git a/third_party/xla/xla/backends/gpu/tests/multi_process_gpu_test_main.cc b/third_party/xla/xla/backends/gpu/tests/multi_process_gpu_test_main.cc
new file mode 100644
index 00000000000..14465c3ac80
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/tests/multi_process_gpu_test_main.cc
@@ -0,0 +1,296 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include <signal.h>
+
+#include <atomic>
+#include <cstdint>
+#include <cstdlib>
+#include <iostream>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include "absl/base/log_severity.h"
+#include "absl/log/check.h"
+#include "absl/log/globals.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/match.h"
+#include "absl/strings/numbers.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
+#include "absl/synchronization/notification.h"
+#include "absl/time/clock.h"
+#include "absl/time/time.h"
+#include "xla/backends/gpu/tests/multi_process_gpu_test_base.h"
+#include "xla/debug_options_flags.h"
+#include "xla/pjrt/distributed/distributed.h"
+#include "xla/pjrt/distributed/service.h"
+#include "xla/service/platform_util.h"
+#include "xla/stream_executor/platform.h"
+#include "xla/tsl/platform/env.h"
+#include "xla/tsl/platform/subprocess.h"
+#include "xla/tsl/platform/test.h"
+#include "xla/tsl/platform/threadpool.h"
+#include "xla/tsl/util/command_line_flags.h"
+#include "tsl/platform/path.h"
+
+namespace {
+
+bool IsWorkerInternalFlag(absl::string_view arg) {
+  return absl::StartsWith(arg, "--multiprocess_worker_id=") ||
+         absl::StartsWith(arg, "--multiprocess_coordinator_address=");
+}
+
+void WriteWorkerLogsToUndeclaredOutputs(int node_id,
+                                        absl::string_view stdout_str,
+                                        absl::string_view stderr_str) {
+  const char* undeclared_outputs_dir =
+      std::getenv("TEST_UNDECLARED_OUTPUTS_DIR");
+  if (undeclared_outputs_dir == nullptr || undeclared_outputs_dir[0] == '\0') {
+    return;
+  }
+  if (!stdout_str.empty()) {
+    std::string stdout_file = tsl::io::JoinPath(
+        undeclared_outputs_dir,
+        absl::StrFormat("subprocess_node_%d_stdout.log", node_id));
+    absl::Status status =
+        tsl::WriteStringToFile(tsl::Env::Default(), stdout_file, stdout_str);
+    if (!status.ok()) {
+      LOG(WARNING) << "Failed to write stdout to " << stdout_file << ": "
+                   << status;
+    }
+  }
+  if (!stderr_str.empty()) {
+    std::string stderr_file = tsl::io::JoinPath(
+        undeclared_outputs_dir,
+        absl::StrFormat("subprocess_node_%d_stderr.log", node_id));
+    absl::Status status =
+        tsl::WriteStringToFile(tsl::Env::Default(), stderr_file, stderr_str);
+    if (!status.ok()) {
+      LOG(WARNING) << "Failed to write stderr to " << stderr_file << ": "
+                   << status;
+    }
+  }
+}
+
+}  // namespace
+
+int main(int argc, char* argv[]) {
+  absl::SetStderrThreshold(absl::LogSeverityAtLeast::kInfo);
+
+  // Snapshot original argv before flag parsers mutate argc/argv so the parent
+  // orchestrator can forward all flags (including XLA and GTest flags) to
+  // worker subprocesses.
+  std::vector<std::string> original_argv;
+  original_argv.reserve(argc);
+  for (int i = 0; i < argc; ++i) {
+    original_argv.emplace_back(argv[i]);
+  }
+
+  int num_workers = 2;
+  int gpus_per_worker = 1;
+  int worker_id = -1;
+  std::string coordinator_address;
+
+  std::vector<tsl::Flag> flag_list = {
+      tsl::Flag("multiprocess_num_workers", &num_workers,
+                "Number of worker processes to launch."),
+      tsl::Flag("multiprocess_gpus_per_worker", &gpus_per_worker,
+                "Number of GPUs assigned per worker process."),
+      tsl::Flag("multiprocess_worker_id", &worker_id,
+                "Internal flag: worker rank in [0, num_workers). -1 for the "
+                "parent orchestrator."),
+      tsl::Flag("multiprocess_coordinator_address", &coordinator_address,
+                "Internal flag: address of the DistributedRuntimeService "
+                "coordinator."),
+  };
+
+  xla::AppendDebugOptionsFlags(&flag_list);
+  tsl::Flags::Parse(&argc, argv, flag_list);
+  testing::InitGoogleTest(&argc, argv);
+
+  // Worker subprocess mode: connect to the coordinator and run tests.
+  if (worker_id >= 0) {
+    const char* undeclared_outputs_dir =
+        std::getenv("TEST_UNDECLARED_OUTPUTS_DIR");
+    if (undeclared_outputs_dir != nullptr &&
+        undeclared_outputs_dir[0] != '\0') {
+      std::string worker_outputs_dir = tsl::io::JoinPath(
+          undeclared_outputs_dir, absl::StrFormat("node_%d", worker_id));
+      absl::Status dir_status =
+          tsl::Env::Default()->RecursivelyCreateDir(worker_outputs_dir);
+      if (!dir_status.ok()) {
+        LOG(WARNING) << "Failed to create worker outputs directory "
+                     << worker_outputs_dir << ": " << dir_status;
+      } else {
+        tsl::setenv("TEST_UNDECLARED_OUTPUTS_DIR", worker_outputs_dir.c_str(),
+                    /*overwrite=*/true);
+      }
+    }
+    xla::InitializeMultiProcessGpuWorker(worker_id, num_workers,
+                                         gpus_per_worker, coordinator_address);
+    int result = RUN_ALL_TESTS();
+    xla::ShutDownMultiProcessGpuWorker();
+    return result;
+  }
+
+  // Parent orchestrator mode.
+  xla::ConfigureMultiProcessGpuParent(num_workers, gpus_per_worker);
+
+  if (GTEST_FLAG_GET(list_tests)) {
+    return RUN_ALL_TESTS();
+  }
+
+  const int required_gpus = num_workers * gpus_per_worker;
+  absl::StatusOr<stream_executor::Platform*> platform =
+      xla::PlatformUtil::GetPlatform("gpu");
+  if (!platform.ok() || (*platform)->VisibleDeviceCount() < required_gpus) {
+    LOG(WARNING) << "Multi-process GPU test requires at least " << required_gpus
+                 << " GPUs, but found "
+                 << (platform.ok() ? (*platform)->VisibleDeviceCount() : 0)
+                 << "; running in-process to skip tests.";
+    return RUN_ALL_TESTS();
+  }
+
+  // Re-export XLA_FLAGS so child processes inherit it.
+  const char* xla_flags = std::getenv("XLA_FLAGS");
+  if (xla_flags != nullptr) {
+    tsl::setenv("XLA_FLAGS", xla_flags, /*overwrite=*/true);
+  }
+
+  // Ensure all workers use the exact same shuffle seed when test shuffling is
+  // enabled without an explicit non-zero seed. Otherwise each worker
+  // independently samples the current millisecond clock and may shuffle tests
+  // into different orders, deadlocking on per-test barriers.
+  if (GTEST_FLAG_GET(shuffle) && GTEST_FLAG_GET(random_seed) == 0) {
+    int32_t seed = 0;
+    const char* env_seed = std::getenv("TEST_RANDOM_SEED");
+    if (env_seed == nullptr || !absl::SimpleAtoi(env_seed, &seed) ||
+        seed <= 0) {
+      seed = static_cast<int32_t>(absl::ToUnixMillis(absl::Now()) % 99999);
+    } else {
+      seed = seed % 99999;
+    }
+    if (seed <= 0) {
+      seed = 1;
+    }
+    const std::string seed_str = absl::StrCat(seed);
+    tsl::setenv("GUNIT_RANDOM_SEED", seed_str.c_str(), /*overwrite=*/true);
+    tsl::setenv("GTEST_RANDOM_SEED", seed_str.c_str(), /*overwrite=*/true);
+  }
+
+  const int port = tsl::testing::PickUnusedPortOrDie();
+  xla::CoordinationServiceImpl::Options service_options;
+  service_options.num_nodes = num_workers;
+  absl::StatusOr<std::unique_ptr<xla::DistributedRuntimeService>> service =
+      xla::GetDistributedRuntimeService(absl::StrFormat("[::]:%d", port),
+                                        service_options);
+  CHECK_OK(service.status()) << "Failed to start DistributedRuntimeService";
+  const std::string resolved_coordinator_address =
+      absl::StrFormat("127.0.0.1:%d", port);
+
+  std::vector<std::unique_ptr<tsl::SubProcess>> children(num_workers);
+  for (int node_id = 0; node_id < num_workers; ++node_id) {
+    if (node_id > 0) {
+      // Only worker 0 should write Bazel/GoogleTest XML output and shard status
+      // files so workers 1..N-1 do not race on or overwrite them.
+      tsl::unsetenv("XML_OUTPUT_FILE");
+      tsl::unsetenv("GUNIT_OUTPUT");
+      tsl::unsetenv("GTEST_OUTPUT");
+      tsl::unsetenv("TEST_SHARD_STATUS_FILE");
+      tsl::unsetenv("TEST_PREMATURE_EXIT_FILE");
+    }
+
+    std::vector<std::string> worker_argv;
+    worker_argv.reserve(original_argv.size() + 3);
+    for (const std::string& arg : original_argv) {
+      if (IsWorkerInternalFlag(arg)) {
+        continue;
+      }
+      if (arg == "--gunit_random_seed=0" || arg == "--gtest_random_seed=0") {
+        continue;
+      }
+      if (node_id > 0 && (absl::StartsWith(arg, "--gtest_output=") ||
+                          absl::StartsWith(arg, "--gunit_output="))) {
+        continue;
+      }
+      worker_argv.push_back(arg);
+    }
+    worker_argv.push_back(
+        absl::StrFormat("--multiprocess_worker_id=%d", node_id));
+    worker_argv.push_back(absl::StrFormat(
+        "--multiprocess_coordinator_address=%s", resolved_coordinator_address));
+
+    children[node_id] = std::make_unique<tsl::SubProcess>();
+    children[node_id]->SetProgram(original_argv[0], worker_argv);
+    children[node_id]->SetChannelAction(tsl::CHAN_STDOUT, tsl::ACTION_PIPE);
+    children[node_id]->SetChannelAction(tsl::CHAN_STDERR, tsl::ACTION_PIPE);
+    CHECK(children[node_id]->Start())
+        << "Failed to start worker process " << node_id;
+  }
+
+  std::vector<int> statuses(num_workers, -1);
+  std::vector<std::string> stdout_strs(num_workers);
+  std::vector<std::string> stderr_strs(num_workers);
+  std::atomic<int> remaining_workers{num_workers};
+  std::atomic<bool> killer_triggered{false};
+  absl::Notification all_done;
+  {
+    tsl::thread::ThreadPool pool(tsl::Env::Default(), "multiprocess_workers",
+                                 num_workers);
+    for (int node_id = 0; node_id < num_workers; ++node_id) {
+      pool.Schedule([&, node_id]() {
+        statuses[node_id] = children[node_id]->Communicate(
+            nullptr, &stdout_strs[node_id], &stderr_strs[node_id]);
+        if (remaining_workers.fetch_sub(1) == 1) {
+          all_done.Notify();
+        } else if (statuses[node_id] != 0 && !killer_triggered.exchange(true)) {
+          if (!all_done.WaitForNotificationWithTimeout(absl::Seconds(5))) {
+            LOG(ERROR) << "Worker " << node_id
+                       << " exited with non-zero status " << statuses[node_id]
+                       << "; killing remaining workers after grace period.";
+            for (int i = 0; i < num_workers; ++i) {
+              children[i]->Kill(SIGKILL);
+            }
+          }
+        }
+      });
+    }
+  }
+
+  if (!stdout_strs.empty() && !stdout_strs[0].empty()) {
+    std::cout << stdout_strs[0] << std::flush;
+  }
+
+  bool any_failed = false;
+  for (int node_id = 0; node_id < num_workers; ++node_id) {
+    WriteWorkerLogsToUndeclaredOutputs(node_id, stdout_strs[node_id],
+                                       stderr_strs[node_id]);
+    if (statuses[node_id] != 0) {
+      any_failed = true;
+      LOG(ERROR) << "Worker " << node_id << " failed with status "
+                 << statuses[node_id] << "\nstdout:\n"
+                 << stdout_strs[node_id] << "\nstderr:\n"
+                 << stderr_strs[node_id];
+    }
+  }
+
+  return any_failed ? EXIT_FAILURE : EXIT_SUCCESS;
+}

From 4ab586c3635e2e5b7697dc8e439bf0e4ee5d927b Mon Sep 17 00:00:00 2001
From: Alexander Lyashuk <crem@google.com>
Date: Fri, 16 Oct 2026 07:25:29 -0700
Subject: [PATCH 0962/1000] [XLA:GPU] Replace TransposeFolding with
 GpuTransposeFolding.

The shared `TransposeFolding` pass is also used by the TPU backend. It is
difficult to add bitcast support to it without changes to TPU behavior. Thus,
this change forks it into a GPU-specific pass.

`GpuTransposeFolding` folds chains of transposes and bitcasts into dot
dimension numbers. A layout-changing bitcast folds as the logical transpose
that it represents. The pass uses `ShapeTracker` to compute the new dimension
numbers, which makes the implementation much simpler. The pass folds only if
GEMM supports the resulting operand layout. For dots that need minor
contracting dimensions (s8 x s8 -> s32, and FP8 before Blackwell), the pass
also keeps the contracting dimensions minor. To share this rule with the pass,
this change moves it from `GpuLayoutAssignment` to new
`DotRequiresMinorContractingDims` and `AreContractingDimsMinor` functions in
`matmul_utils`. The pass also folds transposes into convolution dimension
numbers.

The pass always folds into both dots and convolutions, thus it has no
configuration options. All three `TransposeFolding` passes in `GpuCompiler`
are replaced with `GpuTransposeFolding`. `DotMerger` no longer needs the
`DotDecomposer` normal form, thus the first pass now also folds into dots.

PiperOrigin-RevId: 1000840830
---
 .../codegen/triton/triton_gemm_fusion_test.cc |  19 +-
 .../xla/xla/backends/gpu/transforms/BUILD     |  36 +++
 .../gpu/transforms/gpu_transpose_folding.cc   | 169 +++++++++++++
 .../gpu/transforms/gpu_transpose_folding.h    |  57 +++++
 .../transforms/gpu_transpose_folding_test.cc  | 233 ++++++++++++++++++
 .../gpu/transforms/layout_assignment.cc       |  44 +---
 third_party/xla/xla/service/gpu/BUILD         |   7 +-
 .../xla/xla/service/gpu/gpu_compiler.cc       |  16 +-
 .../xla/xla/service/gpu/matmul_utils.cc       |  23 ++
 .../xla/xla/service/gpu/matmul_utils.h        |  11 +
 .../xla/xla/service/gpu/matmul_utils_test.cc  |  54 ++++
 11 files changed, 613 insertions(+), 56 deletions(-)
 create mode 100644 third_party/xla/xla/backends/gpu/transforms/gpu_transpose_folding.cc
 create mode 100644 third_party/xla/xla/backends/gpu/transforms/gpu_transpose_folding.h
 create mode 100644 third_party/xla/xla/backends/gpu/transforms/gpu_transpose_folding_test.cc

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/triton_gemm_fusion_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/triton_gemm_fusion_test.cc
index 82e44682615..c510c52e7fd 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/triton_gemm_fusion_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/triton_gemm_fusion_test.cc
@@ -742,11 +742,12 @@ ENTRY r {
 })";
 
   MatchOptimizedHlo(kHloText, R"(
-; CHECK: %[[p0:.*]] = f16[10,3,128]{2,0,1} parameter
-; CHECK: %[[cv:.*]] = f32[10,3,128]{2,0,1} convert(%[[p0]])
-; CHECK: %[[tr:.*]] = f32[3,10,128]{2,0,1} transpose(%[[cv]]), dimensions={1,0,2}
+; CHECK-NOT: transpose
+; CHECK: %[[p0:.*]] = f16[3,10,128]{2,1,0} parameter
+; CHECK: %[[cv:.*]] = f32[3,10,128]{2,1,0} convert(%[[p0]])
 ; CHECK: %[[p1:.*]] = f32[3,10,123]{2,1,0} parameter
-; CHECK: f32[3,128,123]{2,1,0} dot(%[[tr]], %[[p1]])
+; CHECK: f32[3,128,123]{2,1,0} dot(%[[cv]], %[[p1]])
+; CHECK-NOT: transpose
 )");
 
   EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/1e-3, /*arel=*/1e-3}));
@@ -1761,10 +1762,12 @@ ENTRY e {
     instr = instr->operand(0);
     ASSERT_TRUE(instr->IsCustomFusion());
   }
-  EXPECT_THAT(
-      instr,
-      GmockMatch(m::Fusion(m::Parameter(), m::Parameter(), m::Parameter())
-                     .WithFusionKind(HloInstruction::FusionKind::kCustom)));
+  EXPECT_THAT(instr, GmockMatch(m::Fusion().WithNumOperands(3).WithFusionKind(
+                         HloInstruction::FusionKind::kCustom)));
+  // Layout-only bitcasts of parameters may be hoisted out of the fusion.
+  EXPECT_THAT(instr->operands(), ::testing::Each(::testing::AnyOf(
+                                     GmockMatch(m::Parameter()),
+                                     GmockMatch(m::Bitcast(m::Parameter())))));
 
   EXPECT_TRUE(RunAndCompare(kHloText, ErrorSpec{/*aabs=*/2e-2, /*arel=*/2e-2}));
 }
diff --git a/third_party/xla/xla/backends/gpu/transforms/BUILD b/third_party/xla/xla/backends/gpu/transforms/BUILD
index 99bd41cb23a..c974d5fec25 100644
--- a/third_party/xla/xla/backends/gpu/transforms/BUILD
+++ b/third_party/xla/xla/backends/gpu/transforms/BUILD
@@ -1509,6 +1509,42 @@ xla_cc_test(
     ],
 )
 
+cc_library(
+    name = "gpu_transpose_folding",
+    srcs = ["gpu_transpose_folding.cc"],
+    hdrs = ["gpu_transpose_folding.h"],
+    deps = [
+        "//xla:xla_data_proto_cc",
+        "//xla/hlo/analysis:shape_tracker",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/pass:hlo_pass",
+        "//xla/service:matmul_indexing_utils",
+        "//xla/service/gpu:matmul_utils",
+        "//xla/stream_executor:device_description",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/types:span",
+    ],
+)
+
+xla_cc_test(
+    name = "gpu_transpose_folding_test",
+    srcs = ["gpu_transpose_folding_test.cc"],
+    deps = [
+        ":gpu_transpose_folding",
+        "//xla/hlo/ir:hlo",
+        "//xla/hlo/testlib:filecheck",
+        "//xla/hlo/testlib:hlo_hardware_independent_test_base",
+        "//xla/stream_executor:device_description",
+        "//xla/stream_executor/cuda:cuda_compute_capability",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "dot_operand_converter",
     srcs = ["dot_operand_converter.cc"],
diff --git a/third_party/xla/xla/backends/gpu/transforms/gpu_transpose_folding.cc b/third_party/xla/xla/backends/gpu/transforms/gpu_transpose_folding.cc
new file mode 100644
index 00000000000..6a906c36121
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/transforms/gpu_transpose_folding.cc
@@ -0,0 +1,169 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/gpu/transforms/gpu_transpose_folding.h"
+
+#include <array>
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <utility>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/status_macros.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "xla/hlo/analysis/shape_tracker.h"
+#include "xla/hlo/ir/hlo_computation.h"
+#include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/service/gpu/matmul_utils.h"
+#include "xla/service/matmul_indexing_utils.h"
+#include "xla/stream_executor/device_description.h"
+#include "xla/xla_data.pb.h"
+
+namespace xla::gpu {
+namespace {
+
+// Skips the chain of transposes and (type-preserving) bitcasts feeding `inst`.
+HloInstruction* SkipTransposesAndBitcasts(HloInstruction* inst) {
+  while (inst->opcode() == HloOpcode::kTranspose ||
+         (inst->opcode() == HloOpcode::kBitcast &&
+          inst->shape().element_type() ==
+              inst->operand(0)->shape().element_type())) {
+    inst = inst->mutable_operand(0);
+  }
+  return inst;
+}
+
+bool IsSupportedByGemm(const DotOperandDims& dims,
+                       bool requires_minor_contracting_dims) {
+  if (requires_minor_contracting_dims &&
+      !AreContractingDimsMinor(dims.shape().layout(),
+                               dims.Indices(DotOperandDims::kContracting))) {
+    return false;
+  }
+  return MatrixLayout::For(dims.shape(), dims.Indices(DotOperandDims::kBatch),
+                           dims.Indices(DotOperandDims::kContracting),
+                           dims.Indices(DotOperandDims::kNonContracting))
+      .ok();
+}
+
+absl::StatusOr<bool> FoldIntoDot(HloInstruction* dot,
+                                 const se::GpuComputeCapability& gpu_version) {
+  ABSL_ASSIGN_OR_RETURN((std::array<DotOperandDims, 2> dims),
+                   DotOperandDims::FromDot(dot));
+  std::array<HloInstruction*, 2> operands = {dot->mutable_operand(0),
+                                             dot->mutable_operand(1)};
+  const bool requires_minor_contracting_dims =
+      DotRequiresMinorContractingDims(*dot, gpu_version);
+  bool changed = false;
+  for (int64_t i = 0; i < 2; ++i) {
+    HloInstruction* source = SkipTransposesAndBitcasts(operands[i]);
+    if (source == operands[i]) {
+      continue;
+    }
+    ABSL_ASSIGN_OR_RETURN(ShapeTracker tracker,
+                     ShapeTracker::FromProducerConsumer(source, operands[i]));
+    ABSL_ASSIGN_OR_RETURN(std::optional<DotOperandDims> source_dims,
+                     dims[i].FoldBackward(tracker));
+    if (!source_dims.has_value() ||
+        !IsSupportedByGemm(*source_dims, requires_minor_contracting_dims)) {
+      continue;
+    }
+    dims[i] = *std::move(source_dims);
+    operands[i] = source;
+    changed = true;
+  }
+  if (!changed) {
+    return false;
+  }
+  ABSL_ASSIGN_OR_RETURN(DotDimensionNumbers dim_nums,
+                   DotOperandDims::CreateDotDimensionNumbers(dims));
+  HloInstruction* new_dot = dot->AddInstruction(
+      HloInstruction::CreateDot(dot->shape(), operands[0], operands[1],
+                                dim_nums, dot->precision_config()));
+  dot->SetupDerivedInstruction(new_dot);
+  ABSL_RETURN_IF_ERROR(dot->parent()->ReplaceInstruction(dot, new_dot));
+  return true;
+}
+
+absl::StatusOr<bool> FoldIntoConvolution(HloInstruction* conv) {
+  ConvolutionDimensionNumbers dim_nums = conv->convolution_dimension_numbers();
+  std::array<HloInstruction*, 2> operands = {conv->mutable_operand(0),
+                                             conv->mutable_operand(1)};
+  bool changed = false;
+  while (operands[0]->opcode() == HloOpcode::kTranspose) {
+    absl::Span<const int64_t> p = operands[0]->dimensions();
+    dim_nums.set_input_batch_dimension(p[dim_nums.input_batch_dimension()]);
+    dim_nums.set_input_feature_dimension(p[dim_nums.input_feature_dimension()]);
+    for (int64_t& dim : *dim_nums.mutable_input_spatial_dimensions()) {
+      dim = p[dim];
+    }
+    operands[0] = operands[0]->mutable_operand(0);
+    changed = true;
+  }
+  while (operands[1]->opcode() == HloOpcode::kTranspose) {
+    absl::Span<const int64_t> p = operands[1]->dimensions();
+    dim_nums.set_kernel_input_feature_dimension(
+        p[dim_nums.kernel_input_feature_dimension()]);
+    dim_nums.set_kernel_output_feature_dimension(
+        p[dim_nums.kernel_output_feature_dimension()]);
+    for (int64_t& dim : *dim_nums.mutable_kernel_spatial_dimensions()) {
+      dim = p[dim];
+    }
+    operands[1] = operands[1]->mutable_operand(0);
+    changed = true;
+  }
+  if (!changed) {
+    return false;
+  }
+  std::unique_ptr<HloInstruction> new_conv =
+      conv->CloneWithNewOperands(conv->shape(), operands);
+  new_conv->set_convolution_dimension_numbers(dim_nums);
+  ABSL_RETURN_IF_ERROR(
+      conv->parent()->ReplaceWithNewInstruction(conv, std::move(new_conv)));
+  return true;
+}
+
+}  // namespace
+
+absl::StatusOr<bool> GpuTransposeFolding::RunImpl(
+    HloModule* module,
+    const absl::flat_hash_set<absl::string_view>& execution_threads) {
+  bool changed = false;
+  for (HloComputation* computation :
+       module->MakeNonfusionComputations(execution_threads)) {
+    for (HloInstruction* instruction :
+         computation->MakeInstructionPostOrder()) {
+      // Skip instructions with extra (e.g. scaling or sparsity) operands.
+      if (instruction->operand_count() != 2) {
+        continue;
+      }
+      if (instruction->opcode() == HloOpcode::kDot) {
+        ABSL_ASSIGN_OR_RETURN(bool folded, FoldIntoDot(instruction, gpu_version_));
+        changed |= folded;
+      } else if (instruction->opcode() == HloOpcode::kConvolution) {
+        ABSL_ASSIGN_OR_RETURN(bool folded, FoldIntoConvolution(instruction));
+        changed |= folded;
+      }
+    }
+  }
+  return changed;
+}
+
+}  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/transforms/gpu_transpose_folding.h b/third_party/xla/xla/backends/gpu/transforms/gpu_transpose_folding.h
new file mode 100644
index 00000000000..eb60110795c
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/transforms/gpu_transpose_folding.h
@@ -0,0 +1,57 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_BACKENDS_GPU_TRANSFORMS_GPU_TRANSPOSE_FOLDING_H_
+#define XLA_BACKENDS_GPU_TRANSFORMS_GPU_TRANSPOSE_FOLDING_H_
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/pass/hlo_pass_interface.h"
+#include "xla/stream_executor/device_description.h"
+
+namespace xla::gpu {
+
+// Folds transposes into dot and convolution dimension numbers.
+//
+// For dots, folds chains of transposes and bitcasts feeding dot operands into
+// the dimension numbers whenever that leaves the dot result unchanged (see
+// DotOperandDims::FoldBackward; layout-changing bitcasts count as the logical
+// transposes they represent) and the source's layout is supported by GEMMs
+// (including the minor contracting dimensions requirement, see
+// DotRequiresMinorContractingDims).
+//
+// For convolutions, folds transpose operands into the input and kernel
+// dimension numbers.
+class GpuTransposeFolding : public HloModulePass {
+ public:
+  explicit GpuTransposeFolding(const se::GpuComputeCapability& gpu_version)
+      : gpu_version_(gpu_version) {}
+
+  absl::string_view name() const override { return "gpu-transpose-folding"; }
+
+ protected:
+  absl::StatusOr<bool> RunImpl(
+      HloModule* module,
+      const absl::flat_hash_set<absl::string_view>& execution_threads) override;
+
+ private:
+  se::GpuComputeCapability gpu_version_;
+};
+
+}  // namespace xla::gpu
+
+#endif  // XLA_BACKENDS_GPU_TRANSFORMS_GPU_TRANSPOSE_FOLDING_H_
diff --git a/third_party/xla/xla/backends/gpu/transforms/gpu_transpose_folding_test.cc b/third_party/xla/xla/backends/gpu/transforms/gpu_transpose_folding_test.cc
new file mode 100644
index 00000000000..bada1b89d88
--- /dev/null
+++ b/third_party/xla/xla/backends/gpu/transforms/gpu_transpose_folding_test.cc
@@ -0,0 +1,233 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/backends/gpu/transforms/gpu_transpose_folding.h"
+
+#include <memory>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/status/status_matchers.h"
+#include "absl/strings/string_view.h"
+#include "xla/hlo/ir/hlo_module.h"
+#include "xla/hlo/testlib/filecheck.h"
+#include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
+#include "xla/stream_executor/cuda/cuda_compute_capability.h"
+#include "xla/stream_executor/device_description.h"
+
+namespace xla::gpu {
+namespace {
+
+using ::absl_testing::IsOkAndHolds;
+
+class GpuTransposeFoldingTest : public HloHardwareIndependentTestBase {
+ protected:
+  void ExpectFolded(absl::string_view hlo, absl::string_view expected) {
+    ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                         ParseAndReturnVerifiedModule(hlo));
+    EXPECT_THAT(GpuTransposeFolding(kHopper).Run(module.get()),
+                IsOkAndHolds(true));
+    EXPECT_THAT(RunFileCheck(module->ToString(), expected), IsOkAndHolds(true));
+  }
+
+  void ExpectNotFolded(absl::string_view hlo) {
+    ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                         ParseAndReturnVerifiedModule(hlo));
+    EXPECT_THAT(GpuTransposeFolding(kHopper).Run(module.get()),
+                IsOkAndHolds(false));
+  }
+
+  const se::GpuComputeCapability kHopper{se::CudaComputeCapability::Hopper()};
+};
+
+TEST_F(GpuTransposeFoldingTest, FoldsTranspose) {
+  ExpectFolded(R"(
+    HloModule m
+    ENTRY e {
+      x = f32[2,3]{1,0} parameter(0)
+      y = f32[2,3]{1,0} parameter(1)
+      t = f32[3,2]{1,0} transpose(y), dimensions={1,0}
+      ROOT dot = f32[2,2]{1,0} dot(x, t), lhs_contracting_dims={1},
+        rhs_contracting_dims={0}
+    })",
+               R"(
+    // CHECK: %[[X:.*]] = {{.*}} parameter(0)
+    // CHECK: %[[Y:.*]] = {{.*}} parameter(1)
+    // CHECK: dot(%[[X]], %[[Y]]), lhs_contracting_dims={1}, rhs_contracting_dims={1}
+    )");
+}
+
+TEST_F(GpuTransposeFoldingTest, FoldsTransposeBitcastChainThatIsIdentity) {
+  // transpose + layout-changing bitcast is logically the identity; only the
+  // layout differs. This is what layout normalization produces from a
+  // layout-changing copy.
+  ExpectFolded(R"(
+    HloModule m
+    ENTRY e {
+      p = f32[8,64,32]{2,1,0} parameter(0)
+      y = f32[8,64,6]{2,1,0} parameter(1)
+      t = f32[8,32,64]{2,1,0} transpose(p), dimensions={0,2,1}
+      b = f32[8,64,32]{1,2,0} bitcast(t)
+      ROOT dot = f32[8,32,6]{2,1,0} dot(b, y), lhs_batch_dims={0},
+        lhs_contracting_dims={1}, rhs_batch_dims={0}, rhs_contracting_dims={1}
+    })",
+               R"(
+    // CHECK: %[[P:.*]] = {{.*}} parameter(0)
+    // CHECK: %[[Y:.*]] = {{.*}} parameter(1)
+    // CHECK-NOT: transpose
+    // CHECK: dot(%[[P]], %[[Y]]), lhs_batch_dims={0}, lhs_contracting_dims={1}, rhs_batch_dims={0}, rhs_contracting_dims={1}
+    )");
+}
+
+TEST_F(GpuTransposeFoldingTest, FoldsLayoutChangingBitcastThatIsTranspose) {
+  // b is logically transpose(p, {0,2,1}).
+  ExpectFolded(R"(
+    HloModule m
+    ENTRY e {
+      p = f32[8,64,32]{2,1,0} parameter(0)
+      y = f32[8,64,6]{2,1,0} parameter(1)
+      b = f32[8,32,64]{1,2,0} bitcast(p)
+      ROOT dot = f32[8,32,6]{2,1,0} dot(b, y), lhs_batch_dims={0},
+        lhs_contracting_dims={2}, rhs_batch_dims={0}, rhs_contracting_dims={1}
+    })",
+               R"(
+    // CHECK: %[[P:.*]] = {{.*}} parameter(0)
+    // CHECK: %[[Y:.*]] = {{.*}} parameter(1)
+    // CHECK: dot(%[[P]], %[[Y]]), lhs_batch_dims={0}, lhs_contracting_dims={1}, rhs_batch_dims={0}, rhs_contracting_dims={1}
+    )");
+}
+
+TEST_F(GpuTransposeFoldingTest, FoldsTransposeKeepingAdjacentDimsTogether) {
+  // Dimensions 2 and 3 of p stay adjacent through the transpose.
+  ExpectFolded(R"(
+    HloModule m
+    ENTRY e {
+      p = f32[4,8,16,32]{3,2,1,0} parameter(0)
+      y = f32[4,8,6]{2,1,0} parameter(1)
+      t = f32[4,16,32,8]{3,2,1,0} transpose(p), dimensions={0,2,3,1}
+      ROOT dot = f32[4,16,32,6]{3,2,1,0} dot(t, y), lhs_batch_dims={0},
+        lhs_contracting_dims={3}, rhs_batch_dims={0}, rhs_contracting_dims={1}
+    })",
+               R"(
+    // CHECK: %[[P:.*]] = {{.*}} parameter(0)
+    // CHECK: %[[Y:.*]] = {{.*}} parameter(1)
+    // CHECK: dot(%[[P]], %[[Y]]), lhs_batch_dims={0}, lhs_contracting_dims={1}, rhs_batch_dims={0}, rhs_contracting_dims={1}
+    )");
+}
+
+TEST_F(GpuTransposeFoldingTest, DoesNotFoldBitcastThatIsReshape) {
+  // Same dimensions as a transpose, but the bitcast is a reshape.
+  ExpectNotFolded(R"(
+    HloModule m
+    ENTRY e {
+      p = f32[8,64,32]{2,1,0} parameter(0)
+      y = f32[8,64,6]{2,1,0} parameter(1)
+      b = f32[8,32,64]{2,1,0} bitcast(p)
+      ROOT dot = f32[8,32,6]{2,1,0} dot(b, y), lhs_batch_dims={0},
+        lhs_contracting_dims={2}, rhs_batch_dims={0}, rhs_contracting_dims={1}
+    })");
+}
+
+TEST_F(GpuTransposeFoldingTest, DoesNotFoldIfBatchDimBecomesMinor) {
+  ExpectNotFolded(R"(
+    HloModule m
+    ENTRY e {
+      p = f32[64,32,8]{2,1,0} parameter(0)
+      y = f32[8,64,6]{2,1,0} parameter(1)
+      t = f32[8,64,32]{2,1,0} transpose(p), dimensions={2,0,1}
+      ROOT dot = f32[8,32,6]{2,1,0} dot(t, y), lhs_batch_dims={0},
+        lhs_contracting_dims={1}, rhs_batch_dims={0}, rhs_contracting_dims={1}
+    })");
+}
+
+TEST_F(GpuTransposeFoldingTest, DoesNotFoldIfNonContractingDimsReorder) {
+  // Folding would swap the order of the non-contracting dims in the output.
+  ExpectNotFolded(R"(
+    HloModule m
+    ENTRY e {
+      p = f32[8,16,32]{2,1,0} parameter(0)
+      y = f32[8,6]{1,0} parameter(1)
+      t = f32[32,16,8]{2,1,0} transpose(p), dimensions={2,1,0}
+      ROOT dot = f32[32,16,6]{2,1,0} dot(t, y), lhs_contracting_dims={2},
+        rhs_contracting_dims={0}
+    })");
+}
+
+TEST_F(GpuTransposeFoldingTest, DoesNotFoldTransposeOfDegenerateDim) {
+  // Moving a size-1 dimension is a reshape, not a transpose: the dimension
+  // can't be tracked back to its category. Conservatively left alone.
+  ExpectNotFolded(R"(
+    HloModule m
+    ENTRY e {
+      p = f32[8,1,6]{2,1,0} parameter(0)
+      y = f32[1,6,5]{2,1,0} parameter(1)
+      t = f32[1,8,6]{2,1,0} transpose(p), dimensions={1,0,2}
+      ROOT dot = f32[1,8,5]{2,1,0} dot(t, y), lhs_batch_dims={0},
+        lhs_contracting_dims={2}, rhs_batch_dims={0}, rhs_contracting_dims={1}
+    })");
+}
+
+TEST_F(GpuTransposeFoldingTest, DoesNotFoldIfS8ContractingDimBecomesMajor) {
+  // This is what layout normalization produces from the copy that layout
+  // assignment inserts to make the contracting dimension of `p` minor.
+  ExpectNotFolded(R"(
+    HloModule m
+    ENTRY e {
+      x = s8[12,4]{1,0} parameter(0)
+      p = s8[4,8]{1,0} parameter(1)
+      t = s8[8,4]{1,0} transpose(p), dimensions={1,0}
+      b = s8[4,8]{0,1} bitcast(t)
+      ROOT dot = s32[12,8]{1,0} dot(x, b), lhs_contracting_dims={1},
+        rhs_contracting_dims={0}
+    })");
+}
+
+TEST_F(GpuTransposeFoldingTest, FoldsS8IfContractingDimStaysMinor) {
+  ExpectFolded(R"(
+    HloModule m
+    ENTRY e {
+      x = s8[12,4]{1,0} parameter(0)
+      p = s8[8,4]{1,0} parameter(1)
+      t = s8[4,8]{1,0} transpose(p), dimensions={1,0}
+      ROOT dot = s32[12,8]{1,0} dot(x, t), lhs_contracting_dims={1},
+        rhs_contracting_dims={0}
+    })",
+               R"(
+    // CHECK: %[[X:.*]] = {{.*}} parameter(0)
+    // CHECK: %[[P:.*]] = {{.*}} parameter(1)
+    // CHECK: dot(%[[X]], %[[P]]), lhs_contracting_dims={1}, rhs_contracting_dims={1}
+    )");
+}
+
+TEST_F(GpuTransposeFoldingTest, FoldsConvolutionOperands) {
+  ExpectFolded(R"(
+    HloModule m
+    ENTRY e {
+      x = f32[3,2,4,5]{3,2,1,0} parameter(0)
+      y = f32[1,2,3,6]{3,2,1,0} parameter(1)
+      tx = f32[2,3,5,4]{3,2,1,0} transpose(x), dimensions={1,0,3,2}
+      ty = f32[3,6,1,2]{3,2,1,0} transpose(y), dimensions={2,3,0,1}
+      ROOT conv = f32[2,6,5,3]{3,2,1,0} convolution(tx, ty),
+        window={size=1x2}, dim_labels=bf01_io01->bf01
+    })",
+               R"(
+    // CHECK: %[[X:.*]] = {{.*}} parameter(0)
+    // CHECK: %[[Y:.*]] = {{.*}} parameter(1)
+    // CHECK: convolution(%[[X]], %[[Y]]), window={size=1x2}, dim_labels=fb10_01io->bf01
+    )");
+}
+
+}  // namespace
+}  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/transforms/layout_assignment.cc b/third_party/xla/xla/backends/gpu/transforms/layout_assignment.cc
index 7736ebe1892..e867f29485f 100644
--- a/third_party/xla/xla/backends/gpu/transforms/layout_assignment.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/layout_assignment.cc
@@ -485,7 +485,6 @@ absl::Status GpuLayoutAssignment::AddDotBackendConstraints(
     const HloInstruction* operand;
     absl::Span<const int64_t> batch_dims;
     absl::Span<const int64_t> contracting_dims;
-    PrimitiveType type;
     std::vector<int64_t> non_contracting_dims;
   };
   auto make_side =
@@ -493,7 +492,6 @@ absl::Status GpuLayoutAssignment::AddDotBackendConstraints(
           absl::Span<const int64_t> contracting_dims) -> absl::StatusOr<Side> {
     Side side = {operand_no, instruction->operand(operand_no), batch_dims,
                  contracting_dims};
-    side.type = side.operand->shape().element_type();
     ABSL_ASSIGN_OR_RETURN(
         side.non_contracting_dims,
         GetNonContractingDims(side.operand->shape(), side.batch_dims,
@@ -508,8 +506,6 @@ absl::Status GpuLayoutAssignment::AddDotBackendConstraints(
                    make_side(1, dot_dims.rhs_batch_dimensions(),
                              dot_dims.rhs_contracting_dimensions()));
 
-  const PrimitiveType& output_type = instruction->shape().element_type();
-
   // Matmuls require the batch dimensions to be in consecutive physical
   // dimensions and likewise for the contracting and non-contracting
   // dimensions. Additionally, no batch dimension can be in the most
@@ -521,17 +517,8 @@ absl::Status GpuLayoutAssignment::AddDotBackendConstraints(
           .debug_options()
           .xla_gpu_experimental_pack_dot_operands_along_k_dimension();
 
-  const bool is_s8_to_s32 = output_type == PrimitiveType::S32 &&
-                            lhs.type == PrimitiveType::S8 &&
-                            rhs.type == PrimitiveType::S8;
-  const bool is_fp8 = (lhs.type == PrimitiveType::F8E4M3FN ||
-                       lhs.type == PrimitiveType::F8E5M2FNUZ) &&
-                      (rhs.type == PrimitiveType::F8E4M3FN ||
-                       rhs.type == PrimitiveType::F8E5M2FNUZ);
-
-  const se::CudaComputeCapability* cc = gpu_version_.cuda_compute_capability();
   const bool both_operands_require_minor_contraction_dims =
-      is_s8_to_s32 || (is_fp8 && !(cc && cc->IsBlackwell()));
+      DotRequiresMinorContractingDims(*instruction, gpu_version_);
 
   for (const Side& side : {lhs, rhs}) {
     if ((IsPackedInstruction(side.operand) && pack_along_contracting_dims) ||
@@ -851,28 +838,15 @@ absl::Status GpuLayoutAssignment::SetDotOperandLayoutToMinorContracting(
     absl::Span<const int64_t> noncontracting_dims) {
   Shape shape = instruction->operand(operand)->shape();
 
+  // If contracting dims are already minor, and the layout is valid, keep it.
   if (shape.has_layout() &&
-      shape.layout().minor_to_major().size() >= contracting_dims.size()) {
-    // Check that the contracting dimensions are physically minor, i.e. check
-    // that minor physical dimensions all point to contracting logical
-    // dimensions.
-    bool contracting_dims_are_minor = true;
-    const auto& minor_to_major = shape.layout().minor_to_major();
-    for (int64_t i = 0; i < contracting_dims.size(); ++i) {
-      if (!absl::c_linear_search(contracting_dims, minor_to_major[i])) {
-        contracting_dims_are_minor = false;
-        break;
-      }
-    }
-
-    // If contracting dims are already minor, and the layout is valid, keep it.
-    if (contracting_dims_are_minor &&
-        MatrixLayout::For(shape, batch_dims, noncontracting_dims,
-                          contracting_dims)
-            .ok()) {
-      // Re-set the operand layout, so it becomes mandatory.
-      return SetOperandLayout(shape, instruction, operand);
-    }
+      shape.layout().minor_to_major().size() >= contracting_dims.size() &&
+      AreContractingDimsMinor(shape.layout(), contracting_dims) &&
+      MatrixLayout::For(shape, batch_dims, noncontracting_dims,
+                        contracting_dims)
+          .ok()) {
+    // Re-set the operand layout, so it becomes mandatory.
+    return SetOperandLayout(shape, instruction, operand);
   }
   return SetOperandMajorToMinorLayout(
       instruction, operand,
diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index fd5f2c696c7..d1b3b7ff95b 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -1300,6 +1300,7 @@ cc_library(
         "//xla/stream_executor:engine_options",
         "//xla/stream_executor:stream",
         "//xla/stream_executor:stream_executor_h",
+        "//xla/stream_executor/cuda:cuda_compute_capability",
         "//xla/stream_executor/gpu:gpu_blas_lt",
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:statusor",
@@ -1325,9 +1326,13 @@ xla_cc_test(
         "//xla/hlo/parser:hlo_parser",
         "//xla/hlo/testlib:hlo_hardware_independent_test_base",
         "//xla/hlo/testlib:test",
+        "//xla/stream_executor:device_description",
+        "//xla/stream_executor/cuda:cuda_compute_capability",
         "//xla/tests:xla_internal_test_main",  # build_cleaner: keep
         "//xla/tsl/platform:statusor",
+        "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest",
     ],
@@ -1859,6 +1864,7 @@ cc_library(
         "//xla/backends/gpu/transforms:gemm_fusion_swap_operands",
         "//xla/backends/gpu/transforms:gemm_rewriter",
         "//xla/backends/gpu/transforms:gpu_copy_async_wrapper",
+        "//xla/backends/gpu/transforms:gpu_transpose_folding",
         "//xla/backends/gpu/transforms:group_collectives_by_key",
         "//xla/backends/gpu/transforms:hoist_fused_bitcasts",
         "//xla/backends/gpu/transforms:layout_assignment",
@@ -2034,7 +2040,6 @@ cc_library(
         "//xla/service:sharding_remover",
         "//xla/service:slow_operation_alarm",
         "//xla/service:topk_rewriter",
-        "//xla/service:transpose_folding",
         "//xla/service:while_loop_constant_sinking",
         "//xla/service:while_loop_simplifier",
         "//xla/service:xla_transform",
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 31c6b4659f0..1db3caaad27 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -123,6 +123,7 @@ limitations under the License.
 #include "xla/backends/gpu/transforms/gemm_fusion_swap_operands.h"
 #include "xla/backends/gpu/transforms/gemm_rewriter.h"
 #include "xla/backends/gpu/transforms/gpu_copy_async_wrapper.h"
+#include "xla/backends/gpu/transforms/gpu_transpose_folding.h"
 #include "xla/backends/gpu/transforms/group_collectives_by_key.h"
 #include "xla/backends/gpu/transforms/hoist_fused_bitcasts.h"
 #include "xla/backends/gpu/transforms/layout_assignment.h"
@@ -323,7 +324,6 @@ limitations under the License.
 #include "xla/service/spmd/schedule_aware_collective_ops_cse.h"
 #include "xla/service/spmd/shardy/shardy_xla_pass.h"
 #include "xla/service/topk_rewriter.h"
-#include "xla/service/transpose_folding.h"
 #include "xla/service/while_loop_constant_sinking.h"
 #include "xla/service/while_loop_simplifier.h"
 #include "xla/service/xla_transform.h"
@@ -1020,12 +1020,7 @@ absl::Status RunOptimizationPasses(
     pipeline.AddPass<ReshapeMover>(reshape_mover_options);
     pipeline.AddPass<HloConstantFolding>();
     pipeline.AddPass<ConditionalSimplifier>();
-    // Do not fold transpose operands into dots yet. This can undo the normal
-    // form established by DotDecomposer, which the DotMerger pass requires.
-    pipeline.AddPass<TransposeFolding>(
-        /*dot_can_fold_transpose_operand=*/
-        [&](const HloInstruction& dot,
-            int64_t operand) -> absl::StatusOr<bool> { return false; });
+    pipeline.AddPass<GpuTransposeFolding>(gpu_version);
     pipeline.AddPass<HloCSE>(/*is_layout_sensitive=*/false);
     pipeline.AddPass<HloDCE>();
   }();
@@ -1044,9 +1039,7 @@ absl::Status RunOptimizationPasses(
   // gpu_version) once Triton emitter supports multiple non-contracting
   // dimensions.
   pipeline.AddPass<DotDimensionNormalizer>(true);
-  // Folding transpose operands into dots can undo the normal form established
-  // by DotDecomposer. Subsequent passes must not rely on it from this point on.
-  pipeline.AddPass<TransposeFolding>(CanFoldTransposeOperandIntoDot);
+  pipeline.AddPass<GpuTransposeFolding>(gpu_version);
 
   // ConvertMover and ReshapeMover fight with each other: ConvertMover wants
   // to move some converts down the graph, but ReshapeMover wants to move them
@@ -2148,8 +2141,7 @@ absl::Status GpuCompiler::OptimizeHloPostLayoutAssignment(
     // GemmRewriter assumes that all transposes are folded into gemms, but,
     // since commit 7d529df, this is not always true at this point.
     // Therefore, rerun transpose folding.
-    pipeline.AddPass<TransposeFolding>(CanFoldTransposeOperandIntoDot,
-                                       TransposeFolding::NeverFoldTranspose);
+    pipeline.AddPass<GpuTransposeFolding>(gpu_version);
 
     pipeline.AddPass<ReshapeDecomposer>();
     pipeline.AddPass<ReduceDecomposer>([&](const HloInstruction* r) {
diff --git a/third_party/xla/xla/service/gpu/matmul_utils.cc b/third_party/xla/xla/service/gpu/matmul_utils.cc
index 9e04e80896c..f54400e6fa2 100644
--- a/third_party/xla/xla/service/gpu/matmul_utils.cc
+++ b/third_party/xla/xla/service/gpu/matmul_utils.cc
@@ -44,6 +44,7 @@ limitations under the License.
 #include "xla/shape_util.h"
 #include "xla/status_macros.h"
 #include "xla/stream_executor/blas.h"
+#include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/device_address.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/engine_options.h"
@@ -1225,6 +1226,28 @@ bool IsDotSupportedByClassicalEmitters(const HloInstruction& dot) {
   }
 }
 
+bool DotRequiresMinorContractingDims(
+    const HloInstruction& dot, const se::GpuComputeCapability& gpu_version) {
+  const PrimitiveType lhs = dot.operand(0)->shape().element_type();
+  const PrimitiveType rhs = dot.operand(1)->shape().element_type();
+  if (dot.shape().element_type() == S32 && lhs == S8 && rhs == S8) {
+    return true;
+  }
+  auto is_fp8 = [](PrimitiveType type) {
+    return type == F8E4M3FN || type == F8E5M2FNUZ;
+  };
+  const se::CudaComputeCapability* cc = gpu_version.cuda_compute_capability();
+  return is_fp8(lhs) && is_fp8(rhs) && !(cc && cc->IsBlackwell());
+}
+
+bool AreContractingDimsMinor(const Layout& layout,
+                             absl::Span<const int64_t> contracting_dims) {
+  return absl::c_all_of(layout.minor_to_major().first(contracting_dims.size()),
+                        [&](int64_t dim) {
+                          return absl::c_linear_search(contracting_dims, dim);
+                        });
+}
+
 PrimitiveType GetGemmAccumulatorType(const HloDotInstruction* dot) {
   // Return the accumulator type if it is explicitly specified as dot algorithm.
   auto accumulator_type = algorithm_util::GetDotAccumulatorType(
diff --git a/third_party/xla/xla/service/gpu/matmul_utils.h b/third_party/xla/xla/service/gpu/matmul_utils.h
index 70b5601f56b..c8fb68f6372 100644
--- a/third_party/xla/xla/service/gpu/matmul_utils.h
+++ b/third_party/xla/xla/service/gpu/matmul_utils.h
@@ -29,6 +29,7 @@ limitations under the License.
 #include "xla/autotuning.pb.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
+#include "xla/layout.h"
 #include "xla/service/gpu/backend_configs.pb.h"
 #include "xla/shape.h"
 #include "xla/stream_executor/blas.h"
@@ -60,6 +61,16 @@ absl::StatusOr<bool> IsMatrixMultiplicationTooSmallForRewriting(
 // so we need to always use cuBLAS or Triton for those.
 bool IsDotSupportedByClassicalEmitters(const HloInstruction& dot);
 
+// Returns true if GEMMs require the contracting dimensions of both operands of
+// `dot` to be physically minor (s8 x s8 -> s32, and FP8 before Blackwell).
+bool DotRequiresMinorContractingDims(
+    const HloInstruction& dot, const se::GpuComputeCapability& gpu_version);
+
+// Returns true if the most minor physical dimensions of `layout` are exactly
+// `contracting_dims`.
+bool AreContractingDimsMinor(const Layout& layout,
+                             absl::Span<const int64_t> contracting_dims);
+
 // Returns the accumulator type for the given dot instruction (either extracted
 // from the dot algorithm or inferred from the output type).
 PrimitiveType GetGemmAccumulatorType(const HloDotInstruction* dot);
diff --git a/third_party/xla/xla/service/gpu/matmul_utils_test.cc b/third_party/xla/xla/service/gpu/matmul_utils_test.cc
index 42c707cb9b6..037bd5cca51 100644
--- a/third_party/xla/xla/service/gpu/matmul_utils_test.cc
+++ b/third_party/xla/xla/service/gpu/matmul_utils_test.cc
@@ -21,14 +21,20 @@ limitations under the License.
 #include <vector>
 
 #include <gtest/gtest.h>
+#include "absl/log/check.h"
 #include "absl/status/status_matchers.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
+#include "absl/strings/substitute.h"
 #include "xla/autotuning.pb.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/parser/hlo_parser.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/hlo/testlib/test.h"
+#include "xla/layout_util.h"
 #include "xla/shape.h"
+#include "xla/stream_executor/cuda/cuda_compute_capability.h"
+#include "xla/stream_executor/device_description.h"
 #include "xla/tsl/platform/statusor.h"
 
 namespace xla {
@@ -391,6 +397,54 @@ TEST(TritonGemmConfigTest, ToStringIncludesWavesPerEu) {
   EXPECT_NE(str.find("waves_per_eu:4"), std::string::npos);
 }
 
+class DotRequiresMinorContractingDimsTest
+    : public HloHardwareIndependentTestBase {
+ protected:
+  bool Requires(absl::string_view lhs, absl::string_view rhs,
+                absl::string_view out, const se::CudaComputeCapability& cc) {
+    absl::StatusOr<std::unique_ptr<HloModule>> module =
+        ParseAndReturnVerifiedModule(absl::Substitute(R"(
+      ENTRY e {
+        x = $0[2,3] parameter(0)
+        y = $1[3,4] parameter(1)
+        ROOT dot = $2[2,4] dot(x, y), lhs_contracting_dims={1},
+          rhs_contracting_dims={0}
+      })",
+                                                      lhs, rhs, out));
+    CHECK_OK(module);
+    return DotRequiresMinorContractingDims(
+        *(*module)->entry_computation()->root_instruction(),
+        se::GpuComputeCapability(cc));
+  }
+};
+
+TEST_F(DotRequiresMinorContractingDimsTest, S8ToS32) {
+  EXPECT_TRUE(Requires("s8", "s8", "s32", se::CudaComputeCapability::Hopper()));
+  EXPECT_FALSE(
+      Requires("s8", "s8", "f32", se::CudaComputeCapability::Hopper()));
+}
+
+TEST_F(DotRequiresMinorContractingDimsTest, Fp8BeforeBlackwell) {
+  EXPECT_TRUE(Requires("f8e4m3fn", "f8e4m3fn", "f32",
+                       se::CudaComputeCapability::Hopper()));
+  EXPECT_FALSE(Requires("f8e4m3fn", "f8e4m3fn", "f32",
+                        se::CudaComputeCapability::Blackwell()));
+}
+
+TEST_F(DotRequiresMinorContractingDimsTest, Float) {
+  EXPECT_FALSE(
+      Requires("f32", "f32", "f32", se::CudaComputeCapability::Hopper()));
+}
+
+TEST(AreContractingDimsMinorTest, Works) {
+  EXPECT_TRUE(AreContractingDimsMinor(LayoutUtil::MakeLayout({2, 1, 0}), {2}));
+  EXPECT_FALSE(AreContractingDimsMinor(LayoutUtil::MakeLayout({2, 1, 0}), {1}));
+  EXPECT_TRUE(
+      AreContractingDimsMinor(LayoutUtil::MakeLayout({1, 2, 0}), {2, 1}));
+  EXPECT_FALSE(
+      AreContractingDimsMinor(LayoutUtil::MakeLayout({1, 0, 2}), {2, 1}));
+}
+
 }  // namespace
 }  // namespace gpu
 }  // namespace xla

From db2d6efccef5c1da12c837d0f46f9a0899a2ceb9 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 16 Oct 2026 07:48:05 -0700
Subject: [PATCH 0963/1000] Automated Code Change

PiperOrigin-RevId: 1000851287
---
 third_party/xla/xla/ffi/ffi_api.cc  | 2 +-
 third_party/xla/xla/ffi/ffi_test.cc | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/ffi/ffi_api.cc b/third_party/xla/xla/ffi/ffi_api.cc
index a6c309b0d3c..72b8bd7658f 100644
--- a/third_party/xla/xla/ffi/ffi_api.cc
+++ b/third_party/xla/xla/ffi/ffi_api.cc
@@ -249,7 +249,7 @@ static XLA_FFI_Error* XLA_FFI_Future_SetError(
     return new XLA_FFI_Error{InvalidArgument("Error must not be null or OK")};
   }
 
-  absl::Status error = TakeStatus(args->error);
+  absl::Status error = TakeError(args->error);
   args->future->async_value.SetError(std::move(error));
 
   return nullptr;
diff --git a/third_party/xla/xla/ffi/ffi_test.cc b/third_party/xla/xla/ffi/ffi_test.cc
index f1cf8c43aa9..e5a5809b663 100644
--- a/third_party/xla/xla/ffi/ffi_test.cc
+++ b/third_party/xla/xla/ffi/ffi_test.cc
@@ -1441,7 +1441,7 @@ TEST(FfiTest, DuplicateHandlerTraits) {
   XLA_FFI_DEFINE_HANDLER(NoOp2, noop, Ffi::Bind(),
                          {Traits::kCmdBufferCompatible});
   XLA_FFI_REGISTER_HANDLER(GetXlaFfiApi(), "duplicate-traits", "Host", NoOp1);
-  auto status = TakeStatus(Ffi::RegisterStaticHandler(
+  auto status = TakeError(Ffi::RegisterStaticHandler(
       GetXlaFfiApi(), "duplicate-traits", "Host", NoOp2));
   EXPECT_TRUE(
       absl::StrContains(status.message(), "Duplicate FFI handler registration"))
@@ -1455,7 +1455,7 @@ TEST(FfiTest, DuplicateHandlerAddress) {
   XLA_FFI_DEFINE_HANDLER(NoOp1, noop1, Ffi::Bind());
   XLA_FFI_DEFINE_HANDLER(NoOp2, noop2, Ffi::Bind());
   XLA_FFI_REGISTER_HANDLER(GetXlaFfiApi(), "duplicate-address", "Host", NoOp1);
-  auto status = TakeStatus(Ffi::RegisterStaticHandler(
+  auto status = TakeError(Ffi::RegisterStaticHandler(
       GetXlaFfiApi(), "duplicate-address", "Host", NoOp2));
   EXPECT_TRUE(
       absl::StrContains(status.message(), "Duplicate FFI handler registration"))
@@ -1468,7 +1468,7 @@ TEST(FfiTest, AllowRegisterDuplicateWhenEqual) {
   XLA_FFI_DEFINE_HANDLER(NoOp, noop, Ffi::Bind());
   XLA_FFI_REGISTER_HANDLER(GetXlaFfiApi(), "duplicate-when-equal", "Host",
                            NoOp);
-  auto status = TakeStatus(Ffi::RegisterStaticHandler(
+  auto status = TakeError(Ffi::RegisterStaticHandler(
       GetXlaFfiApi(), "duplicate-when-equal", "Host", NoOp));
   TF_ASSERT_OK(status);
 }

From 5a9a312ffea37b80feb99298e618faa9375b7f9f Mon Sep 17 00:00:00 2001
From: Levon Ter-Grigoryan <patrios@google.com>
Date: Fri, 16 Oct 2026 07:57:40 -0700
Subject: [PATCH 0964/1000] Migrate multi-process GPU all-gather tests to
 MultiProcessGpuTest.

PiperOrigin-RevId: 1000856121
---
 third_party/xla/xla/backends/gpu/tests/BUILD  |  22 +-
 .../tests/all_gather_e2e_multiprocess_test.cc | 304 +++---------------
 2 files changed, 53 insertions(+), 273 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/tests/BUILD b/third_party/xla/xla/backends/gpu/tests/BUILD
index 1462837e95a..72fac1f7c69 100644
--- a/third_party/xla/xla/backends/gpu/tests/BUILD
+++ b/third_party/xla/xla/backends/gpu/tests/BUILD
@@ -1817,11 +1817,12 @@ xla_test(
         "gpu",
     ],
     deps = [
+        ":multi_process_gpu_test_base",
+        ":multi_process_gpu_test_main",
         "//xla:debug_options_flags",
         "//xla:literal",
         "//xla:literal_util",
         "//xla:shape_util",
-        "//xla:status_macros",
         "//xla:xla_data_proto_cc",
         "//xla:xla_proto_cc",
         "//xla/hlo/builder:xla_computation",
@@ -1830,14 +1831,7 @@ xla_test(
         "//xla/pjrt:pjrt_client",
         "//xla/pjrt:pjrt_compiler",
         "//xla/pjrt:pjrt_executable",
-        "//xla/pjrt/distributed",
-        "//xla/pjrt/distributed:client",
-        "//xla/pjrt/distributed:service",
-        "//xla/pjrt/plugin/xla_gpu:xla_gpu_allocator_config",
-        "//xla/pjrt/plugin/xla_gpu:xla_gpu_client_options",
-        "//xla/pjrt/plugin/xla_gpu:xla_gpu_pjrt_client",
         "//xla/service:device_assignment",
-        "//xla/service:gpu_topology",
         "//xla/service:platform_util",
         "//xla/service/gpu:backend_configs_cc",
         "//xla/stream_executor:device_description",
@@ -1845,20 +1839,10 @@ xla_test(
         "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor/cuda:cuda_compute_capability",
         "//xla/tests:literal_test_util",
-        "//xla/tsl/platform:env",
-        "//xla/tsl/platform:subprocess",
-        "//xla/tsl/platform:test",
-        "//xla/tsl/util:command_line_flags",
         "@com_google_absl//absl/log",
-        "@com_google_absl//absl/log:check",
-        "@com_google_absl//absl/log:flags",
-        "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/strings:string_view",
-        "@com_google_absl//absl/time",
         "@com_google_googletest//:gtest",
-        "@tsl//tsl/platform:path",
     ],
 )
 
diff --git a/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_multiprocess_test.cc b/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_multiprocess_test.cc
index 2a42a496f19..ffed1cc74c9 100644
--- a/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_multiprocess_test.cc
+++ b/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_multiprocess_test.cc
@@ -13,21 +13,15 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include <cstdlib>
 #include <memory>
-#include <optional>
-#include <string>
-#include <utility>
 #include <vector>
 
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
-#include "absl/log/check.h"
 #include "absl/log/log.h"
-#include "absl/status/status.h"
-#include "absl/status/statusor.h"
-#include "absl/strings/str_format.h"
+#include "absl/status/status_matchers.h"
 #include "absl/strings/string_view.h"
-#include "absl/time/time.h"
+#include "xla/backends/gpu/tests/multi_process_gpu_test_base.h"
 #include "xla/debug_options_flags.h"
 #include "xla/hlo/builder/xla_computation.h"
 #include "xla/hlo/ir/hlo_computation.h"
@@ -37,120 +31,42 @@ limitations under the License.
 #include "xla/hlo/parser/hlo_parser.h"
 #include "xla/literal.h"
 #include "xla/literal_util.h"
-#include "xla/pjrt/distributed/client.h"
-#include "xla/pjrt/distributed/distributed.h"
-#include "xla/pjrt/distributed/service.h"
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/pjrt/pjrt_compiler.h"
 #include "xla/pjrt/pjrt_executable.h"
-#include "xla/pjrt/plugin/xla_gpu/xla_gpu_allocator_config.h"
-#include "xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h"
-#include "xla/pjrt/plugin/xla_gpu/xla_gpu_pjrt_client.h"
 #include "xla/service/device_assignment.h"
 #include "xla/service/gpu/backend_configs.pb.h"
-#include "xla/service/gpu_topology.h"
 #include "xla/service/platform_util.h"
 #include "xla/shape_util.h"
-#include "xla/status_macros.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/stream_executor/platform.h"
 #include "xla/stream_executor/stream_executor.h"
 #include "xla/tests/literal_test_util.h"
-#include "xla/tsl/platform/env.h"
-#include "xla/tsl/platform/subprocess.h"
-#include "xla/tsl/platform/test.h"
-#include "xla/tsl/util/command_line_flags.h"
 #include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
-#include "tsl/platform/path.h"
 
 namespace xla {
 namespace {
 
-inline constexpr size_t kMB = 1024LL * 1024LL;
-constexpr int kNumNodes = 2;
-static const char* test_binary_name;
+using AllGatherMultiProcessE2ETest = MultiProcessGpuTest;
 
-struct MultiProcessGpuClientSetup {
-  std::unique_ptr<xla::DistributedRuntimeService> service;
-  std::unique_ptr<PjRtClient> client;
-};
-
-absl::StatusOr<MultiProcessGpuClientSetup> SetUpMultiProcessGpuClient(
-    int rank_id, int num_nodes, int port, absl::string_view log_prefix) {
-  MultiProcessGpuClientSetup prepared_test;
-  std::string coordinator_address = absl::StrFormat("127.0.0.1:%d", port);
-
-  // Rank 0 creates the coordination service.
-  if (rank_id == 0) {
-    LOG(INFO) << log_prefix << ": creating coordination service on "
-              << coordinator_address;
-    xla::CoordinationServiceImpl::Options service_options;
-    service_options.num_nodes = num_nodes;
-    ABSL_ASSIGN_OR_RETURN(prepared_test.service,
-                     xla::GetDistributedRuntimeService(
-                         absl::StrFormat("[::]:%d", port), service_options));
-    LOG(INFO) << log_prefix << ": created coordination service";
-  }
-
-  // Connect to the coordination service.
-  xla::DistributedRuntimeClient::Options distributed_options;
-  distributed_options.node_id = rank_id;
-  distributed_options.init_timeout = absl::Seconds(120);
-  auto distributed_client =
-      GetDistributedRuntimeClient(coordinator_address, distributed_options);
-
-  LOG(INFO) << log_prefix << ": connecting distributed client";
-  ABSL_RETURN_IF_ERROR(distributed_client->Connect());
-  LOG(INFO) << log_prefix << ": distributed client connected";
-
-  // Create the GPU client with a single addressable device per process.
-  GpuClientOptions options;
-  options.node_id = rank_id;
-  options.num_nodes = num_nodes;
-  options.allowed_devices = {rank_id};
-  options.kv_store =
-      GetDistributedKeyValueStore(distributed_client, /*key_prefix=*/"gpu:");
-  options.distributed_client = distributed_client;
-  options.allocator_config.kind = xla::GpuAllocatorConfig::Kind::kBFC;
-  options.allocator_config.gpu_system_memory_size = 32 * kMB;
-  options.allocator_config.collective_memory_size = 0;
-  options.use_tfrt_gpu_client = true;
-
-  LOG(INFO) << log_prefix << ": creating PjRtClient";
-  ABSL_ASSIGN_OR_RETURN(prepared_test.client, GetXlaPjrtGpuClient(options));
-  LOG(INFO) << log_prefix << ": PjRtClient created";
-
-  return prepared_test;
-}
-
-absl::Status AllGatherMultiProcessTestBody(int node_id, int port) {
-  std::string log_prefix = absl::StrFormat("rank_%d", node_id);
-  ABSL_ASSIGN_OR_RETURN(se::Platform * platform, PlatformUtil::GetPlatform("gpu"));
-  ABSL_ASSIGN_OR_RETURN(se::StreamExecutor * executor,
-                   platform->ExecutorForDevice(node_id));
-  const auto& desc = executor->GetDeviceDescription();
+TEST_F(AllGatherMultiProcessE2ETest, OneShotAllGather2Processes) {
+  ASSERT_OK_AND_ASSIGN(se::Platform * platform,
+                       PlatformUtil::GetPlatform("gpu"));
+  ASSERT_OK_AND_ASSIGN(se::StreamExecutor * executor,
+                       platform->ExecutorForDevice(rank()));
+  const se::DeviceDescription& desc = executor->GetDeviceDescription();
   if (desc.gpu_compute_capability().IsCuda() && !desc.gpu_compute_capability()
                                                      .cuda_compute_capability()
                                                      ->IsAtLeastAmpere()) {
-    LOG(INFO) << log_prefix
-              << ": skipping one-shot all-gather test, requires Ampere+ for "
-                 "CUDA.";
-    return absl::OkStatus();
+    GTEST_SKIP() << "One-shot all-gather requires Ampere+ for CUDA.";
   }
 
-  ABSL_ASSIGN_OR_RETURN(
-      MultiProcessGpuClientSetup setup,
-      SetUpMultiProcessGpuClient(node_id, kNumNodes, port, log_prefix));
-  std::unique_ptr<PjRtClient> client = std::move(setup.client);
+  ASSERT_EQ(client()->addressable_device_count(), 1);
+  ASSERT_EQ(client()->device_count(), num_ranks());
 
-  TF_RET_CHECK(client->addressable_device_count() == 1)
-      << "Expected exactly 1 local addressable device per process.";
-  TF_RET_CHECK(client->device_count() == kNumNodes)
-      << "Expected " << kNumNodes << " global devices.";
-
-  constexpr absl::string_view kModuleStr = R"(
+  constexpr absl::string_view kModuleStr = R"hlo(
     HloModule test
 
     ENTRY test_computation {
@@ -158,31 +74,22 @@ absl::Status AllGatherMultiProcessTestBody(int node_id, int port) {
       ROOT all-gather = f32[256] all-gather(param_0), dimensions={0},
         replica_groups={{0,1}}
     }
-  )";
+  )hlo";
 
-  ABSL_ASSIGN_OR_RETURN(auto hlo_module,
-                   ParseAndReturnUnverifiedModule(kModuleStr, /*config=*/{}));
-  xla::XlaComputation computation(hlo_module->ToProto());
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> hlo_module,
+                       ParseAndReturnUnverifiedModule(kModuleStr));
+  XlaComputation computation(hlo_module->ToProto());
 
-  xla::CompileOptions compile_options;
-  compile_options.executable_build_options.set_num_replicas(kNumNodes);
+  CompileOptions compile_options;
+  compile_options.executable_build_options.set_num_replicas(num_ranks());
   compile_options.executable_build_options.set_num_partitions(1);
-  DeviceAssignment device_assignment(kNumNodes, 1);
-  device_assignment(0, 0) = 0;
-  device_assignment(1, 0) = 1;
+  DeviceAssignment device_assignment(num_ranks(), 1);
+  for (int i = 0; i < num_ranks(); ++i) {
+    device_assignment(i, 0) = i;
+  }
   compile_options.executable_build_options.set_device_assignment(
       device_assignment);
 
-  GpuTopology gpu_topology(
-      /*platform_version=*/"",
-      /*num_partitions=*/1,
-      /*num_hosts_per_partition=*/1,
-      /*num_devices_per_host=*/kNumNodes,
-      /*gpu_target_config=*/std::nullopt,
-      /*host_target_machine_options=*/std::nullopt,
-      /*num_devices_per_process=*/kNumNodes);
-  compile_options.executable_build_options.set_gpu_topology(gpu_topology);
-
   DebugOptions debug_options = GetDebugOptionsFromFlags();
   debug_options.set_xla_gpu_autotune_level(0);
   debug_options.add_xla_gpu_experimental_use_collective_kernels(
@@ -192,15 +99,13 @@ absl::Status AllGatherMultiProcessTestBody(int node_id, int port) {
   *compile_options.executable_build_options.mutable_debug_options() =
       debug_options;
 
-  LOG(INFO) << log_prefix << ": compiling HLO module with one-shot all-gather";
-  ABSL_ASSIGN_OR_RETURN(std::unique_ptr<PjRtLoadedExecutable> executable,
-                   client->CompileAndLoad(computation, compile_options));
-  LOG(INFO) << log_prefix << ": compilation succeeded";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtLoadedExecutable> executable,
+                       client()->CompileAndLoad(computation, compile_options));
 
   // Inspect the collective kernel strategy selected in the optimized HLO.
-  ABSL_ASSIGN_OR_RETURN(auto hlo_modules,
-                   executable->GetExecutable()->GetHloModules());
-  TF_RET_CHECK(!hlo_modules.empty());
+  ASSERT_OK_AND_ASSIGN(auto hlo_modules,
+                       executable->GetExecutable()->GetHloModules());
+  ASSERT_FALSE(hlo_modules.empty());
   const HloModule* optimized_module = hlo_modules.front().get();
   bool found_all_gather = false;
   for (const HloComputation* comp : optimized_module->computations()) {
@@ -208,47 +113,42 @@ absl::Status AllGatherMultiProcessTestBody(int node_id, int port) {
       if (instr->opcode() == HloOpcode::kAllGather ||
           instr->opcode() == HloOpcode::kAllGatherStart) {
         found_all_gather = true;
-        auto gpu_config = instr->backend_config<gpu::GpuBackendConfig>();
-        TF_RET_CHECK(gpu_config.ok())
-            << "Failed to get GpuBackendConfig for " << instr->name() << ": "
-            << gpu_config.status();
+        ASSERT_OK_AND_ASSIGN(auto gpu_config,
+                             instr->backend_config<gpu::GpuBackendConfig>());
         LOG(INFO)
-            << log_prefix << ": AllGather instruction " << instr->name()
+            << "rank_" << rank() << ": AllGather instruction " << instr->name()
             << " collective kernel strategy: "
             << gpu::CollectiveBackendConfig::CollectiveKernelStrategy_Name(
-                   gpu_config->collective_backend_config().kernel_strategy());
-        TF_RET_CHECK(
-            gpu_config->collective_backend_config().kernel_strategy() ==
-            gpu::CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT)
-            << "Expected AllGather to use KERNEL_STRATEGY_TRITON_ONE_SHOT, but "
-               "got: "
-            << gpu::CollectiveBackendConfig::CollectiveKernelStrategy_Name(
-                   gpu_config->collective_backend_config().kernel_strategy());
+                   gpu_config.collective_backend_config().kernel_strategy());
+        EXPECT_EQ(
+            gpu_config.collective_backend_config().kernel_strategy(),
+            gpu::CollectiveBackendConfig::KERNEL_STRATEGY_TRITON_ONE_SHOT);
       }
     }
   }
-  TF_RET_CHECK(found_all_gather)
+  EXPECT_TRUE(found_all_gather)
       << "Expected to find an AllGather instruction in optimized HLO.";
 
   // Prepare input literal: rank 0 provides 1.0f, rank 1 provides 2.0f.
-  const float input_val = (node_id == 0) ? 1.0f : 2.0f;
+  const float input_val = (rank() == 0) ? 1.0f : 2.0f;
   Literal input_literal =
       LiteralUtil::CreateR1<float>(std::vector<float>(128, input_val));
 
-  ABSL_ASSIGN_OR_RETURN(auto* memory_space,
-                   client->addressable_devices()[0]->default_memory_space());
-  ABSL_ASSIGN_OR_RETURN(auto input_buffer,
-                   client->BufferFromHostLiteral(input_literal, memory_space));
+  ASSERT_OK_AND_ASSIGN(
+      PjRtMemorySpace * memory_space,
+      client()->addressable_devices()[0]->default_memory_space());
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<PjRtBuffer> input_buffer,
+      client()->BufferFromHostLiteral(input_literal, memory_space));
 
   std::vector<std::vector<PjRtBuffer*>> input_ptrs = {{input_buffer.get()}};
-  LOG(INFO) << log_prefix << ": executing one-shot all-gather";
-  ABSL_ASSIGN_OR_RETURN(auto results,
-                   executable->Execute(input_ptrs, ExecuteOptions()));
-  LOG(INFO) << log_prefix << ": execution finished";
+  ASSERT_OK_AND_ASSIGN(auto results,
+                       executable->Execute(input_ptrs, ExecuteOptions()));
 
-  TF_RET_CHECK(results.size() == 1 && results[0].size() == 1);
+  ASSERT_EQ(results.size(), 1);
+  ASSERT_EQ(results[0].size(), 1);
   Literal result_literal(ShapeUtil::MakeShape(F32, {256}));
-  ABSL_RETURN_IF_ERROR(results[0][0]->ToLiteralSync(&result_literal));
+  ASSERT_OK(results[0][0]->ToLiteralSync(&result_literal));
 
   // Expected result is [1.0f x 128, 2.0f x 128] on all ranks.
   std::vector<float> expected_data;
@@ -260,112 +160,8 @@ absl::Status AllGatherMultiProcessTestBody(int node_id, int port) {
     expected_data.push_back(2.0f);
   }
   Literal expected_literal = LiteralUtil::CreateR1<float>(expected_data);
-
-  if (!LiteralTestUtil::Equal(expected_literal, result_literal)) {
-    return absl::InternalError(absl::StrFormat(
-        "Result literal %s does not match expected %s on rank %d",
-        result_literal.ToString(), expected_literal.ToString(), node_id));
-  }
-
-  LOG(INFO) << log_prefix << ": verified result successfully";
-  return absl::OkStatus();
-}
-
-TEST(AllGatherMultiProcessE2ETest, OneShotAllGather2Processes) {
-  absl::StatusOr<se::Platform*> platform = PlatformUtil::GetPlatform("gpu");
-  if (!platform.ok() || (*platform)->VisibleDeviceCount() < kNumNodes) {
-    GTEST_SKIP() << "Test requires at least " << kNumNodes
-                 << " GPU devices, but found "
-                 << (platform.ok() ? (*platform)->VisibleDeviceCount() : 0);
-  }
-
-  // Re-pass XLA_FLAGS environment variable to child processes.
-  const char* xla_flags = std::getenv("XLA_FLAGS");
-  if (xla_flags != nullptr) {
-    tsl::setenv("XLA_FLAGS", xla_flags, /*overwrite=*/true);
-  }
-
-  int port = tsl::testing::PickUnusedPortOrDie();
-  tsl::SubProcess child[kNumNodes];
-
-  for (int node_id = 0; node_id < kNumNodes; ++node_id) {
-    std::vector<std::string> argv = {
-        test_binary_name,
-        "--test_to_run=AllGatherMultiProcessHelper",
-        absl::StrFormat("--node_id=%d", node_id),
-        absl::StrFormat("--port=%d", port),
-        "--alsologtostderr",
-        "--vmodule=gpu_executable=1,thunk_executor=1,all_gather_thunk=5,"
-        "collective_kernel_thunk=5,collective_memory=5",
-    };
-    child[node_id].SetProgram(test_binary_name, argv);
-    child[node_id].SetChannelAction(tsl::CHAN_STDOUT, tsl::ACTION_PIPE);
-    child[node_id].SetChannelAction(tsl::CHAN_STDERR, tsl::ACTION_PIPE);
-    ASSERT_TRUE(child[node_id].Start()) << "Failed to start node " << node_id;
-  }
-
-  for (int node_id = 0; node_id < kNumNodes; ++node_id) {
-    std::string stdout_str, stderr_str;
-    int status = child[node_id].Communicate(nullptr, &stdout_str, &stderr_str);
-
-    const char* undeclared_outputs_dir =
-        std::getenv("TEST_UNDECLARED_OUTPUTS_DIR");
-    if (undeclared_outputs_dir != nullptr &&
-        undeclared_outputs_dir[0] != '\0') {
-      std::string stderr_file = tsl::io::JoinPath(
-          undeclared_outputs_dir,
-          absl::StrFormat("subprocess_node_%d_stderr.log", node_id));
-      absl::Status write_status =
-          tsl::WriteStringToFile(tsl::Env::Default(), stderr_file, stderr_str);
-      if (!write_status.ok()) {
-        LOG(WARNING) << "Failed to write stderr to " << stderr_file << ": "
-                     << write_status;
-      }
-    }
-
-    EXPECT_EQ(status, 0) << "node " << node_id << " failed with status "
-                         << status << "\nstdout:\n"
-                         << stdout_str << "\nstderr:\n"
-                         << stderr_str;
-  }
+  EXPECT_TRUE(LiteralTestUtil::Equal(expected_literal, result_literal));
 }
 
 }  // namespace
 }  // namespace xla
-
-int main(int argc, char* argv[]) {
-  std::string test_to_run;
-  int node_id = -1;
-  int port = -1;
-  xla::test_binary_name = argv[0];
-
-  std::vector<tsl::Flag> flag_list = {
-      tsl::Flag("test_to_run", &test_to_run,
-                "The test to run in the child process."),
-      tsl::Flag("node_id", &node_id,
-                "The node id (rank) for the child process."),
-      tsl::Flag("port", &port, "The coordinator port for distributed runtime."),
-  };
-
-  xla::AppendDebugOptionsFlags(&flag_list);
-  tsl::Flags::Parse(&argc, argv, flag_list);
-  testing::InitGoogleTest(&argc, argv);
-
-  if (test_to_run.empty()) {
-    return RUN_ALL_TESTS();
-  }
-
-  absl::Status result = absl::OkStatus();
-  if (test_to_run == "AllGatherMultiProcessHelper") {
-    result = xla::AllGatherMultiProcessTestBody(node_id, port);
-  } else {
-    result = absl::InvalidArgumentError(absl::StrFormat(
-        "Unrecognized multiprocess test name: %s", test_to_run));
-  }
-
-  if (!result.ok()) {
-    LOG(ERROR) << "Child process (node_id " << node_id
-               << ") failed: " << result;
-  }
-  return result.raw_code();
-}

From eae7bc861bcf4bca25b982239ed15c21c7672630 Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Fri, 16 Oct 2026 09:02:37 -0700
Subject: [PATCH 0965/1000] SpmdBuilder::SetBroadcastDimsForReshape records
 which dimensions of a partitioned reshape hold replicated broadcast values,
 so that ReshardWithCollectivePermute can skip collective-permute along them.
 It compared local shard shapes instead of unpartitioned base shapes. When a
 reshape merges a sharded source dimension into a dimension whose shard size
 matches a broadcast dimension, that dimension was falsely marked as broadcast
 and resharding emitted a local copy instead of a collective-permute. Compare
 the base shapes of the reshape being visited instead.

PiperOrigin-RevId: 1000891513
---
 .../xla/xla/service/spmd/spmd_partitioner.cc   | 13 +++++++++++--
 .../xla/service/spmd/spmd_partitioner_test.cc  | 18 ++++++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.cc b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
index 25d931bb110..5782a9aae56 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
@@ -344,13 +344,22 @@ void SpmdBuilder::SetBroadcastDimsForReshape(const HloInstruction& hlo) {
   if (it == broadcast_dims_.end()) {
     return;
   }
+  if (visiting_hlo_ == nullptr ||
+      visiting_hlo_->opcode() != HloOpcode::kReshape ||
+      visiting_hlo_->operand(0)->shape().dimensions().size() !=
+          hlo.operand(0)->shape().dimensions().size() ||
+      visiting_hlo_->shape().dimensions().size() !=
+          hlo.shape().dimensions().size()) {
+    return;
+  }
   std::vector<int64_t> iota_dims(hlo.shape().dimensions().size());
   absl::c_iota(iota_dims, 0);
   absl::flat_hash_set<int64_t> reshape_broadcast_dims(iota_dims.begin(),
                                                       iota_dims.end());
 
-  absl::Span<const int64_t> operand_dims = hlo.operand(0)->shape().dimensions();
-  absl::Span<const int64_t> hlo_dims = hlo.shape().dimensions();
+  absl::Span<const int64_t> operand_dims =
+      visiting_hlo_->operand(0)->shape().dimensions();
+  absl::Span<const int64_t> hlo_dims = visiting_hlo_->shape().dimensions();
   std::vector<int64_t> before_dim_size_stack(operand_dims.rbegin(),
                                              operand_dims.rend());
   std::vector<int64_t> after_dim_size_stack(hlo_dims.rbegin(), hlo_dims.rend());
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
index 017a60bbcf4..64b84218163 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
@@ -16951,6 +16951,24 @@ ENTRY entry {
                   op::DynamicSlice(op::Pad(_, op::Constant()), _, _, _, _), _));
 }
 
+TEST_P(SpmdPartitioningTest,
+       ReshapeOfBroadcastDoesNotFalselyMarkMergedShardedDimAsBroadcast) {
+  constexpr absl::string_view hlo_string = R"(
+HloModule module
+
+ENTRY entry {
+  %p = f32[4] parameter(0), sharding={devices=[4]<=[4]}
+  %br = f32[4,2] broadcast(%p), dimensions={0}, sharding={devices=[4,1]<=[4]}
+  %reshape = f32[2,4] reshape(%br), sharding={devices=[2,2]<=[4]}
+  ROOT %copy = f32[2,4] copy(%reshape), sharding={devices=[2,2]1,0,3,2}
+})";
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       PartitionComputation(hlo_string, /*num_devices=*/4));
+  EXPECT_THAT(module->entry_computation()->root_instruction(),
+              op::Copy(op::CollectivePermute(_)));
+}
+
 TEST_P(SpmdPartitioningTest, KeepPartitionedNonSlicedDimension) {
   const char* const hlo_string = R"(
 HloModule module

From 5194b35901595807079b0ea467c9afd61d93e6ff Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Fri, 16 Oct 2026 09:16:13 -0700
Subject: [PATCH 0966/1000] [XLA:GPU]: Prepare collective fusion and runtime
 for variadic all-gather.

Wrap tuple collective fusions in get-tuple-element and tuple instructions with tuple simplification disabled, and generalize all-gather runtime info and kernel spec for multiple operands.

PiperOrigin-RevId: 1000899539
---
 .../xla/backends/gpu/codegen/triton/fusion.cc |  23 +++-
 .../xla/xla/backends/gpu/runtime/BUILD        |   1 +
 .../xla/backends/gpu/runtime/all_gather.cc    | 107 +++++++++++-------
 .../xla/xla/backends/gpu/runtime/all_gather.h |  18 +--
 .../backends/gpu/transforms/collectives/BUILD |   1 +
 .../collectives/collective_fusion.cc          |  17 ++-
 .../simplifiers/tuple_simplifier.cc           |   4 +
 .../transforms/simplifiers/tuple_simplifier.h |   3 +
 .../simplifiers/tuple_simplifier_test.cc      |  27 +++++
 9 files changed, 142 insertions(+), 59 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/fusion.cc b/third_party/xla/xla/backends/gpu/codegen/triton/fusion.cc
index c8cb80a5ae1..0bf031ca39e 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/fusion.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/fusion.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 #include "xla/backends/gpu/codegen/triton/fusion.h"
 
+#include <algorithm>
 #include <cstdint>
 #include <memory>
 #include <optional>
@@ -30,6 +31,7 @@ limitations under the License.
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
 #include "absl/types/span.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/IR/Function.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
@@ -382,15 +384,26 @@ std::optional<TritonFusion::LaunchConfig> TritonFusion::GetLaunchConfig(
       BlockLevelParameters::FromBlockLevelFusionConfig(
           analysis->fusion_backend_config().block_level_fusion_config());
 
-  // We expect all roots to have the same number of blocks. Otherwise we
-  // cannot codegen it.
+  CHECK_EQ(analysis->fusion_root_count(),
+           block_level_parameters.output_tile_sizes.size());
   int64_t num_blocks =
       GetNumberOfBlocks(analysis->fusion_root(0).shape().dimensions(),
                         block_level_parameters.output_tile_sizes[0]);
   for (int64_t i = 1; i < analysis->fusion_root_count(); ++i) {
-    CHECK_EQ(GetNumberOfBlocks(analysis->fusion_root(i).shape().dimensions(),
-                               block_level_parameters.output_tile_sizes[i]),
-             num_blocks);
+    int64_t root_blocks =
+        GetNumberOfBlocks(analysis->fusion_root(i).shape().dimensions(),
+                          block_level_parameters.output_tile_sizes[i]);
+    if (analysis->fusion_backend_config().kind() ==
+        kTritonCollectiveFusionKind) {
+      // Variadic collective fusion: each output has its own tile size, and
+      // smaller outputs are padded with masked-out blocks to match the largest
+      // output's block count. So the launch grid is the max across outputs.
+      num_blocks = std::max(num_blocks, root_blocks);
+    } else {
+      // We expect all roots to have the same number of blocks. Otherwise we
+      // cannot codegen it.
+      CHECK_EQ(root_blocks, num_blocks);
+    }
   }
 
   LaunchConfig launch_config;
diff --git a/third_party/xla/xla/backends/gpu/runtime/BUILD b/third_party/xla/xla/backends/gpu/runtime/BUILD
index ac6992e7a99..49f6a829056 100644
--- a/third_party/xla/xla/backends/gpu/runtime/BUILD
+++ b/third_party/xla/xla/backends/gpu/runtime/BUILD
@@ -4178,6 +4178,7 @@ cc_library(
     deps = [
         ":collective_params",
         "//xla:shape_util",
+        "//xla:status_macros",
         "//xla:util",
         "//xla:xla_data_proto_cc",
         "//xla/backends/gpu/transforms/collectives:collective_ops_utils",
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_gather.cc b/third_party/xla/xla/backends/gpu/runtime/all_gather.cc
index 98262866205..c7335a890c8 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_gather.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/all_gather.cc
@@ -15,6 +15,7 @@ limitations under the License.
 
 #include "xla/backends/gpu/runtime/all_gather.h"
 
+#include <algorithm>
 #include <cstdint>
 #include <vector>
 
@@ -36,6 +37,7 @@ limitations under the License.
 #include "xla/service/gpu/launch_dimensions.h"
 #include "xla/service/gpu_topology.h"
 #include "xla/shape_util.h"
+#include "xla/status_macros.h"
 #include "xla/stream_executor/device_description.h"
 #include "xla/util.h"
 #include "xla/xla_data.pb.h"
@@ -91,6 +93,8 @@ absl::Status IsAllGatherKernelSupported(
         "devices. Got %d.",
         num_devices));
   }
+  // Triton pads tile sizes to powers of 2, so each rank's slice along the
+  // gather dimension must be a power of 2.
   if (!llvm::has_single_bit(static_cast<uint64_t>(per_rank_gather_dim_size))) {
     return absl::UnimplementedError(absl::StrFormat(
         "All-gather kernel requires the per-rank size along the gather "
@@ -132,13 +136,6 @@ absl::StatusOr<AllGatherInfo> BuildAllGatherInfo(
   }
   const int64_t num_devices =
       all_gather->device_list()->num_devices_per_group();
-  const int64_t num_elements =
-      ShapeUtil::ElementsIn(all_gather->operand(0)->shape());
-  const int64_t per_rank_gather_dim_size =
-      all_gather->operand(0)->shape().dimensions(
-          all_gather->all_gather_dimension());
-  const PrimitiveType element_type =
-      all_gather->operand(0)->shape().element_type();
   const int32_t num_operands = all_gather->operand_count();
   if (device_info.device_interconnect_info().active_links <= 0) {
     return absl::UnimplementedError(
@@ -148,14 +145,22 @@ absl::StatusOr<AllGatherInfo> BuildAllGatherInfo(
   ABSL_ASSIGN_OR_RETURN(
       const bool is_local,
       AreAllReplicasOnSameSlice(gpu_topology, *all_gather, device_assignment));
-  ABSL_RETURN_IF_ERROR(IsAllGatherKernelSupported(
-      device_info, num_operands, num_devices, num_elements,
-      per_rank_gather_dim_size, element_type, is_local,
-      all_gather->replica_groups()));
+
+  int64_t num_elements = 0;
+  for (const HloInstruction* operand : all_gather->operands()) {
+    const int64_t operand_elements = ShapeUtil::ElementsIn(operand->shape());
+    const int64_t per_rank_gather_dim_size =
+        operand->shape().dimensions(all_gather->all_gather_dimension());
+    ABSL_RETURN_IF_ERROR(IsAllGatherKernelSupported(
+        device_info, num_operands, num_devices, operand_elements,
+        per_rank_gather_dim_size, operand->shape().element_type(), is_local,
+        all_gather->replica_groups()));
+    num_elements = std::max(num_elements, operand_elements);
+  }
   return AllGatherInfo{
       /*.num_devices =*/num_devices,
       /*.num_elements =*/num_elements,
-      /*.element_type =*/element_type,
+      /*.element_type =*/all_gather->operand(0)->shape().element_type(),
   };
 }
 
@@ -167,43 +172,57 @@ absl::StatusOr<CollectiveKernelSpec> CreateAllGatherKernelSpec(
       instr->replica_groups()[0].replica_ids_size() > 0) {
     group_size = instr->replica_groups()[0].replica_ids_size();
   }
+  const int64_t num_operands = instr->operand_count();
+  TF_RET_CHECK(num_operands >= 1) << "all-gather requires at least one operand";
 
-  // The symmetric scratch buffer holds one rank's slice (= input size).
-  // Double-buffering is handled by the should_double_buffer flag.
-  const int64_t input_size_bytes =
-      ShapeUtil::ByteSizeOf(instr->operand(0)->shape());
   const int64_t num_signal_flags = group_size * launch_dimensions.num_blocks();
   const int64_t signal_size = xla::RoundUpTo<uint64_t>(
       num_signal_flags * sizeof(int32_t), kXlaAllocatedBufferAlignBytes);
-  // Each rank allocates its own input-sized symmetric buffer; pointers are
-  // exchanged via kXlaRendezvous so every rank can read from each peer.
-  const int64_t remote_size =
-      xla::RoundUpTo<uint64_t>(input_size_bytes, kXlaAllocatedBufferAlignBytes);
 
-  CollectiveKernelSpec kernel_spec = {
-      /* .codegen_config= */ {
-          /* .copy_input_to_scratch= */ false,
-          /* .input_buffer_specs= */
-          {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}},
-          /* .output_buffer_specs= */
-          {{/*requires_multimem=*/false, SymmetricMemoryType::kNone}},
-          /* .argument_descriptors= */
-          {{KernelArgType::kInputBuffer, /*index=*/0},
-           {KernelArgType::kOutputBuffer, /*index=*/0},
-           {KernelArgType::kRuntimeRank},
-           {KernelArgType::kInvocationCount},
-           {KernelArgType::kScratchBuffer,
-            /*index=*/0},  // signal buffers
-           {KernelArgType::kScratchBuffer,
-            /*index=*/1}},  // remote scratch buffers
-          /* .sync_count_increment= */ 1u},
-      /* .scratch_buffers= */
-      {{signal_size, /*requires_multimem=*/false, scratch_memory_type,
-        /*should_memzero=*/true,
-        /*should_double_buffer=*/true},
-       {remote_size, /*requires_multimem=*/false, scratch_memory_type,
-        /*should_memzero=*/false,
-        /*should_double_buffer=*/true}}};
+  CollectiveKernelSpec kernel_spec;
+  CollectiveCodegenConfig& codegen_config = kernel_spec.codegen_config;
+  codegen_config.copy_input_to_scratch = false;
+  codegen_config.sync_count_increment = 1u;
+  // Argument layout (must match the Triton entry function): inputs, outputs,
+  // rank, invocation count, signal flags, one remote scratch table per operand.
+  for (int32_t i = 0; i < num_operands; ++i) {
+    codegen_config.input_buffer_specs.push_back(
+        {/*requires_multimem=*/false, SymmetricMemoryType::kNone});
+    codegen_config.argument_descriptors.push_back(
+        {KernelArgType::kInputBuffer, /*index=*/i});
+  }
+  for (int32_t i = 0; i < num_operands; ++i) {
+    codegen_config.output_buffer_specs.push_back(
+        {/*requires_multimem=*/false, SymmetricMemoryType::kNone});
+    codegen_config.argument_descriptors.push_back(
+        {KernelArgType::kOutputBuffer, /*index=*/i});
+  }
+  codegen_config.argument_descriptors.push_back({KernelArgType::kRuntimeRank});
+  codegen_config.argument_descriptors.push_back(
+      {KernelArgType::kInvocationCount});
+  // Signal buffers.
+  codegen_config.argument_descriptors.push_back(
+      {KernelArgType::kScratchBuffer, /*index=*/0});
+  kernel_spec.scratch_buffers.push_back({signal_size,
+                                         /*requires_multimem=*/false,
+                                         scratch_memory_type,
+                                         /*should_memzero=*/true,
+                                         /*should_double_buffer=*/true});
+  // One symmetric scratch buffer per operand that holds one rank's slice
+  // (= input size). Double-buffering is handled by the should_double_buffer
+  // flag.
+  for (int32_t i = 0; i < num_operands; ++i) {
+    const int64_t remote_size = xla::RoundUpTo<uint64_t>(
+        ShapeUtil::ByteSizeOf(instr->operand(i)->shape()),
+        kXlaAllocatedBufferAlignBytes);
+    codegen_config.argument_descriptors.push_back(
+        {KernelArgType::kScratchBuffer, /*index=*/1 + i});
+    kernel_spec.scratch_buffers.push_back({remote_size,
+                                           /*requires_multimem=*/false,
+                                           scratch_memory_type,
+                                           /*should_memzero=*/false,
+                                           /*should_double_buffer=*/true});
+  }
   return kernel_spec;
 }
 
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_gather.h b/third_party/xla/xla/backends/gpu/runtime/all_gather.h
index dd942e5b356..1bb5599616a 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_gather.h
+++ b/third_party/xla/xla/backends/gpu/runtime/all_gather.h
@@ -64,14 +64,16 @@ absl::StatusOr<AllGatherInfo> BuildAllGatherInfo(
     const DeviceAssignment* device_assignment);
 
 // Creates a CollectiveKernelSpec describing the resource requirements of a
-// Triton all-gather kernel.  The kernel argument layout matches all-reduce:
-//   [0] local input buffer (kInputBuffer, index 0)
-//   [1] output buffer (kOutputBuffer, index 0)
-//   [2] runtime rank  (kRuntimeRank)
-//   [3] invocation count (kInvocationCount)
-//   [4] signal flags (kScratchBuffer, index 0)
-//   [5] remote scratch buffer pointer table (kScratchBuffer, index 1)
-// `scratch_memory_type` is used for both scratch buffers.
+// Triton all-gather kernel with N operands. The kernel argument layout matches
+// all-reduce:
+//   [0, N)      local input buffers (kInputBuffer, index i)
+//   [N, 2N)     output buffers (kOutputBuffer, index i)
+//   [2N]        runtime rank  (kRuntimeRank)
+//   [2N + 1]    invocation count (kInvocationCount)
+//   [2N + 2]    signal flags (kScratchBuffer, index 0)
+//   [2N + 3 + i] remote scratch buffer pointer table of operand i
+//               (kScratchBuffer, index 1 + i)
+// `scratch_memory_type` is used for all scratch buffers.
 absl::StatusOr<CollectiveKernelSpec> CreateAllGatherKernelSpec(
     const HloInstruction* instr, const LaunchDimensions& launch_dimensions,
     SymmetricMemoryType scratch_memory_type);
diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/BUILD b/third_party/xla/xla/backends/gpu/transforms/collectives/BUILD
index 704ab31abd7..8d37e294167 100644
--- a/third_party/xla/xla/backends/gpu/transforms/collectives/BUILD
+++ b/third_party/xla/xla/backends/gpu/transforms/collectives/BUILD
@@ -971,6 +971,7 @@ cc_library(
         "//xla/backends/gpu/runtime:all_reduce",
         "//xla/hlo/ir:hlo",
         "//xla/hlo/pass:hlo_pass",
+        "//xla/hlo/transforms/simplifiers:tuple_simplifier",
         "//xla/service:device_assignment",
         "//xla/service:gpu_topology",
         "//xla/service/gpu:backend_configs_cc",
diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion.cc b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion.cc
index 31577cf9150..9289e60e0cc 100644
--- a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion.cc
@@ -38,6 +38,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_opcode.h"
+#include "xla/hlo/transforms/simplifiers/tuple_simplifier.h"
 #include "xla/primitive_util.h"
 #include "xla/service/device_assignment.h"
 #include "xla/service/gpu/backend_configs.pb.h"
@@ -70,8 +71,20 @@ HloFusionInstruction* CreateCollectiveFusionInstruction(
     parameters.push_back(builder.AddInstruction(HloInstruction::CreateParameter(
         i, operand->shape(), absl::StrCat("param_", i))));
   }
-  HloInstruction* fused_root = builder.AddInstruction(
+  HloInstruction* collective = builder.AddInstruction(
       instr->CloneWithNewOperands(instr->shape(), parameters));
+  HloInstruction* fused_root = collective;
+  if (collective->shape().IsTuple()) {
+    std::vector<HloInstruction*> gtes;
+    gtes.reserve(collective->shape().tuple_shapes().size());
+    for (int64_t i = 0; i < collective->shape().tuple_shapes().size(); ++i) {
+      gtes.push_back(builder.AddInstruction(
+          HloInstruction::CreateGetTupleElement(collective, i)));
+    }
+    fused_root = builder.AddInstruction(HloInstruction::CreateTuple(gtes));
+    fused_root->add_frontend_attribute(kDisableTupleSimplifierAttribute,
+                                       "true");
+  }
   HloComputation* fused_computation =
       module->AddEmbeddedComputation(builder.Build(fused_root));
   HloInstruction* fusion_instr =
@@ -81,7 +94,7 @@ HloFusionInstruction* CreateCollectiveFusionInstruction(
           fused_computation, absl::StrCat(instr->name(), "-")));
   // Propagate metadata.
   fusion_instr->set_metadata(instr->metadata());
-  fusion_instr->set_frontend_attributes(fused_root->frontend_attributes());
+  fusion_instr->set_frontend_attributes(collective->frontend_attributes());
   fusion_instr->CopyBackendConfigFrom(instr);
   return Cast<HloFusionInstruction>(fusion_instr);
 }
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/tuple_simplifier.cc b/third_party/xla/xla/hlo/transforms/simplifiers/tuple_simplifier.cc
index c07682580dd..1cd9e91014d 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/tuple_simplifier.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/tuple_simplifier.cc
@@ -36,6 +36,10 @@ TupleSimplifier::TupleSimplifier(bool exclude_entry_computation)
 
 absl::StatusOr<HloInstruction*> TupleSimplifier::RemoveWholeTuple(
     HloInstruction* tuple) {
+  if (tuple->get_frontend_attribute(kDisableTupleSimplifierAttribute) ==
+      "true") {
+    return nullptr;
+  }
   HloInstruction* top_tuple = nullptr;
   for (int64_t operand_number = 0; operand_number < tuple->operand_count();
        ++operand_number) {
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/tuple_simplifier.h b/third_party/xla/xla/hlo/transforms/simplifiers/tuple_simplifier.h
index e63a12f1fed..b24b7a20338 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/tuple_simplifier.h
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/tuple_simplifier.h
@@ -25,6 +25,9 @@ limitations under the License.
 
 namespace xla {
 
+inline constexpr absl::string_view kDisableTupleSimplifierAttribute =
+    "disable_tuple_simplifier";
+
 // A pass which simplifies patterns of Tuple and GetTupleElement instructions in
 // the module.
 class TupleSimplifier : public HloModulePass {
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/tuple_simplifier_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/tuple_simplifier_test.cc
index e6e74b1b6da..e3c37076f0e 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/tuple_simplifier_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/tuple_simplifier_test.cc
@@ -441,5 +441,32 @@ ENTRY main {
   EXPECT_TRUE(ordering.ExecutesBefore(op_start_1, op_done_1));
 }
 
+TEST_F(TupleSimplifierTest, DoesNotRemoveTupleWithDisableAttribute) {
+  constexpr absl::string_view kHlo = R"(
+HloModule m
+
+fused_computation {
+  p0 = f32[4] parameter(0)
+  p1 = f32[8] parameter(1)
+  ag = (f32[8], f32[16]) all-gather(p0, p1), dimensions={0}
+  gte.0 = f32[8] get-tuple-element(ag), index=0
+  gte.1 = f32[16] get-tuple-element(ag), index=1
+  ROOT tuple = (f32[8], f32[16]) tuple(gte.0, gte.1), frontend_attributes={disable_tuple_simplifier="true"}
+}
+
+ENTRY main {
+  p0 = f32[4] parameter(0)
+  p1 = f32[8] parameter(1)
+  ROOT fusion = (f32[8], f32[16]) fusion(p0, p1), kind=kCustom, calls=fused_computation
+}
+  )";
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kHlo));
+  Run(module.get(), /*change_expected=*/false);
+  EXPECT_THAT(
+      FindComputation(module.get(), "fused_computation")->root_instruction(),
+      op::Tuple(op::GetTupleElement(), op::GetTupleElement()));
+}
+
 }  // namespace
 }  // namespace xla

From c4d1e6df33b58c237c1a7c4d8c536910d25ab808 Mon Sep 17 00:00:00 2001
From: Nicholas Levin <nglevin@google.com>
Date: Fri, 16 Oct 2026 10:24:31 -0700
Subject: [PATCH 0967/1000] Automated Code Change

PiperOrigin-RevId: 1000942430
---
 .../lite/objc/apps/TestApp/TestApp/Base.lproj/Main.storyboard   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/lite/objc/apps/TestApp/TestApp/Base.lproj/Main.storyboard b/tensorflow/lite/objc/apps/TestApp/TestApp/Base.lproj/Main.storyboard
index 4f233d9ec68..0654b62e090 100644
--- a/tensorflow/lite/objc/apps/TestApp/TestApp/Base.lproj/Main.storyboard
+++ b/tensorflow/lite/objc/apps/TestApp/TestApp/Base.lproj/Main.storyboard
@@ -51,7 +51,7 @@
                                     </barButtonItem>
                                 </items>
                             </toolbar>
-                            <textView clipsSubviews="YES" multipleTouchEnabled="YES" contentMode="scaleToFill" misplaced="YES" editable="NO" adjustsFontForContentSizeCategory="YES" selectable="NO" translatesAutoresizingMaskIntoConstraints="NO" id="7Ws-3t-76I">
+                            <textView clipsSubviews="YES" multipleTouchEnabled="YES" contentMode="scaleToFill" misplaced="YES" editable="NO" selectable="NO" translatesAutoresizingMaskIntoConstraints="NO" id="7Ws-3t-76I">
                                 <rect key="frame" x="0.0" y="194" width="375" height="488"/>
                                 <color key="backgroundColor" red="0.12820077356385221" green="0.40366933178860925" blue="0.96080166101455688" alpha="1" colorSpace="custom" customColorSpace="displayP3"/>
                                 <color key="textColor" systemColor="tableCellGroupedBackgroundColor"/>

From 0ca2a04db3263f961b79e7c437ddd73b248455e6 Mon Sep 17 00:00:00 2001
From: Jane Liu <janeliu@nvidia.com>
Date: Fri, 16 Oct 2026 11:17:47 -0700
Subject: [PATCH 0968/1000] PR #50881: [XLA:GPU] Add missing CUPTI V2
 dropped-records stub
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Imported from GitHub PR https://github.com/openxla/xla/pull/50881

📝 Summary of Changes

Add `cuptiActivityGetNumDroppedRecords_v2` to the CUPTI forwarding-stub symbol list. This lets JAX Mosaic GPU resolve a required API for per-subscriber profiling through XLA’s existing CUPTI loader without `LD_PRELOAD`.

🎯 Justification

JAX Mosaic GPU requires all seven CUPTI V2 entry points to enable per-subscriber profiling. Without this stub, `cuptiActivityGetNumDroppedRecords_v2` remains an unresolved weak reference in `_mosaic_gpu_ext.so`. Mosaic therefore falls back to V1, where subscription fails if another CUPTI profiler, such as Nsight Systems, is already active.

🚀 Kind of Contribution

🐛 Bug Fix

📊 Benchmark (for Performance Improvements)

N/A

🧪 Unit Tests:

No new unit tests added.

🧪 Execution Tests:

No new execution tests added.
Copybara import of the project:

--
738fed52b523c39cbaa5bc4283efa4bb958b118a by Jane Liu <janeliu@nvidia.com>:

[GPU] Add missing CUPTI V2 dropped-records stub

--
49a0e61f27e207dd88c1446a867092fdada9b67a by Jane Liu <janeliu@nvidia.com>:

[GPU] Remove CUPTI stub symbol test

Merging this change closes #50881

PiperOrigin-RevId: 1000979018
---
 third_party/xla/xla/tsl/cuda/cupti.symbols | 1 +
 1 file changed, 1 insertion(+)

diff --git a/third_party/xla/xla/tsl/cuda/cupti.symbols b/third_party/xla/xla/tsl/cuda/cupti.symbols
index 535be58d290..c39328b4a78 100644
--- a/third_party/xla/xla/tsl/cuda/cupti.symbols
+++ b/third_party/xla/xla/tsl/cuda/cupti.symbols
@@ -20,6 +20,7 @@ cuptiActivityGetAttribute
 cuptiActivityGetNextRecord
 cuptiActivityGetNextRecord_v2
 cuptiActivityGetNumDroppedRecords
+cuptiActivityGetNumDroppedRecords_v2
 cuptiActivityGetNvtxExtPayloadAttr
 cuptiActivityGetNvtxExtPayloadEntryTypeInfo
 cuptiActivityPopExternalCorrelationId

From 3fa835e278875694c5e5397c3e4d80a4add81988 Mon Sep 17 00:00:00 2001
From: Bill Varcho <varcho@google.com>
Date: Fri, 16 Oct 2026 11:18:20 -0700
Subject: [PATCH 0969/1000] [XLA:SPMD][unreduced][cleanup] Remove obsolete
 ResolveReductionOpForSharding helper.

I believe this function was a temporary hueristic to infer reductionOp. Now SDY propagates reductionOp directly on TensorShardingAttr, so just read it from there.

PiperOrigin-RevId: 1000979439
---
 .../xla/xla/service/spmd/spmd_partitioner.cc  | 81 +------------------
 1 file changed, 1 insertion(+), 80 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.cc b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
index 5782a9aae56..8018a29b374 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
@@ -7316,83 +7316,6 @@ int64_t SpmdPartitioner::CommunicationCostInBytes(HloInstruction* hlo) {
   }
 }
 
-static HloSharding ResolveReductionOpForSharding(
-    const HloInstruction* instruction, HloSharding sharding) {
-  if (sharding.IsTuple()) {
-    std::vector<HloSharding> new_elements;
-    new_elements.reserve(sharding.tuple_elements().size());
-    for (int i = 0; i < sharding.tuple_elements().size(); ++i) {
-      const HloInstruction* element_instr = nullptr;
-      if (instruction->opcode() == HloOpcode::kTuple &&
-          i < instruction->operand_count()) {
-        element_instr = instruction->operand(i);
-      } else {
-        element_instr = instruction;
-      }
-      new_elements.push_back(ResolveReductionOpForSharding(
-          element_instr, sharding.tuple_elements()[i]));
-    }
-    Shape tuple_shape;
-    if (instruction->shape().IsTuple()) {
-      tuple_shape = instruction->shape();
-    } else {
-      std::vector<Shape> shapes(new_elements.size(), instruction->shape());
-      tuple_shape = ShapeUtil::MakeTupleShape(shapes);
-    }
-    return HloSharding::Tuple(tuple_shape, new_elements);
-  }
-
-  if (sharding.UseNamedShardingLeaf()) {
-    if (sharding.named_sharding().unreduced_axes().empty()) {
-      return sharding;
-    }
-  } else if (!sharding.IsUnreducedSubgroup()) {
-    return sharding;
-  }
-  if (sharding.reduction_op() != ReductionOp::kSum) {
-    return sharding;
-  }
-
-  const HloInstruction* current = instruction;
-  while (current && current->opcode() != HloOpcode::kReduce &&
-         current->opcode() != HloOpcode::kAllReduce &&
-         current->opcode() != HloOpcode::kReduceScatter &&
-         current->operand_count() > 0) {
-    current = current->operand(0);
-  }
-  if (current && (current->opcode() == HloOpcode::kReduce ||
-                  current->opcode() == HloOpcode::kAllReduce ||
-                  current->opcode() == HloOpcode::kReduceScatter)) {
-    const HloComputation* computation = current->to_apply();
-    if (computation) {
-      const HloInstruction* root = computation->root_instruction();
-      ReductionOp reduction_op = ReductionOp::kSum;
-      if (root && root->opcode() == HloOpcode::kMaximum) {
-        reduction_op = ReductionOp::kMax;
-      } else if (root && root->opcode() == HloOpcode::kMinimum) {
-        reduction_op = ReductionOp::kMin;
-      }
-      sharding.set_reduction_op(reduction_op);
-    }
-  } else if (current && current->opcode() == HloOpcode::kParameter) {
-    const HloModule* module = current->parent()->parent();
-    if (module && module->has_spmd_parameters_shardings() &&
-        current->parameter_number() <
-            module->spmd_parameters_shardings().size()) {
-      const HloSharding& param_sharding =
-          module->spmd_parameters_shardings()[current->parameter_number()];
-      const bool param_has_unreduced =
-          param_sharding.UseNamedShardingLeaf()
-              ? !param_sharding.named_sharding().unreduced_axes().empty()
-              : param_sharding.IsUnreducedSubgroup();
-      if (param_has_unreduced) {
-        sharding.set_reduction_op(param_sharding.reduction_op());
-      }
-    }
-  }
-  return sharding;
-}
-
 /* static */ void SpmdPartitioner::RecordInputsOutputsSharding(
     HloModule* module) {
   // Add the parameters' and output's shardings to the module.
@@ -7407,9 +7330,7 @@ static HloSharding ResolveReductionOpForSharding(
   module->set_spmd_parameters_shardings(entry_params_shardings);
   auto entry_root = module->entry_computation()->root_instruction();
   CHECK(entry_root->has_sharding()) << "Missing sharding in entry root.";
-  HloSharding output_sharding =
-      ResolveReductionOpForSharding(entry_root, entry_root->sharding());
-  module->set_spmd_output_sharding(std::move(output_sharding));
+  module->set_spmd_output_sharding(entry_root->sharding());
 }
 
 absl::StatusOr<bool> SpmdPartitioner::RunImpl(

From 96a5acc9707d027af4ffd796493216de9fc953ab Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 16 Oct 2026 11:30:51 -0700
Subject: [PATCH 0970/1000] Wait for all pending transfers and execution
 futures on error in FunctionalHloRunner

`CopyArgumentsToDevice` issues asynchronous `BufferFromHostLiteral` transfers
that read directly from the caller's host `Literal`s until each buffer's ready
future completes. Previously, `CopyArgumentsToDevice` returned immediately if
creating a subsequent argument buffer failed or if the first
`GetReadyFuture().Await()` returned an error, allowing callers such as
`CreateArgumentsOnDevice`, `Run`, and `LoadAndRun` to destroy their stack-local
`PerDeviceLiteralVecType` while background H2D transfer threads were still
reading the freed host `Literal` buffers.

Similarly, `RunInternal`, `CreateUninitializedArgumentsOnDevice`, and the
non-returning output path in `FetchAndLogOutput` returned on the first failed
future instead of awaiting all in-flight per-device execution or buffer-ready
futures before unwinding.

Update `CopyArgumentsToDevice`, `RunInternal`,
`CreateUninitializedArgumentsOnDevice`, and `FetchAndLogOutput` to await all
already-issued futures via `absl::Status::Update` before propagating errors.

PiperOrigin-RevId: 1000987968
---
 .../xla/xla/tools/multihost_hlo_runner/BUILD  |   4 +
 .../functional_hlo_runner.cc                  |  85 ++++++-----
 .../functional_hlo_runner_test.cc             | 136 ++++++++++++++++++
 3 files changed, 190 insertions(+), 35 deletions(-)

diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/BUILD b/third_party/xla/xla/tools/multihost_hlo_runner/BUILD
index 6703a79dc9b..e9043fd62ac 100644
--- a/third_party/xla/xla/tools/multihost_hlo_runner/BUILD
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/BUILD
@@ -290,6 +290,7 @@ xla_test(
         ":hlo_input_output_format",
         ":profiler_interface",
         "//xla:debug_options_flags",
+        "//xla:future",
         "//xla:literal",
         "//xla:shape_layout",
         "//xla:shape_util",
@@ -301,6 +302,7 @@ xla_test(
         "//xla/hlo/testlib:filecheck",
         "//xla/pjrt:maybe_owning_mlir_module",
         "//xla/pjrt:mlir_to_hlo",
+        "//xla/pjrt:mock_pjrt_client",
         "//xla/pjrt:pjrt_client",
         "//xla/pjrt:pjrt_compiler",
         "//xla/pjrt:pjrt_executable",
@@ -322,6 +324,7 @@ xla_test(
         "//xla/tsl/platform:test",
         "//xla/tsl/testing:temporary_directory",
         "//xla/tsl/util:command_line_flags",
+        "@com_google_absl//absl/functional:any_invocable",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
@@ -330,6 +333,7 @@ xla_test(
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
         "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest",
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
index ac8cce570ff..45c25fa5e69 100644
--- a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
@@ -387,6 +387,7 @@ absl::StatusOr<PerDeviceLiteralVecType> FetchAndLogOutput(
   size_t num_pending_transfers = 0;
 
   absl::Status issue_status = [&]() -> absl::Status {
+    absl::Status ready_status;
     for (int i = 0; i < output_buffers.size(); ++i) {
       if (output_buffers[i].empty()) {
         continue;
@@ -424,11 +425,11 @@ absl::StatusOr<PerDeviceLiteralVecType> FetchAndLogOutput(
                 "All outputs from a given vector of outputs should be for the "
                 "same device");
           }
-          ABSL_RETURN_IF_ERROR(buffer->GetReadyFuture().Await());
+          ready_status.Update(buffer->GetReadyFuture().Await());
         }
       }
     }
-    return absl::OkStatus();
+    return ready_status;
   }();
 
   // The ToLiteral callbacks above reference `outputs`, `mu`, `status` and
@@ -632,9 +633,11 @@ absl::StatusOr<PerDeviceLiteralVecType> RunInternal(
       ABSL_ASSIGN_OR_RETURN(
           output_buffers,
           executable->Execute(argument_ptrs, execute_options, futures));
+      absl::Status execute_status;
       for (auto& future : *futures) {
-        ABSL_RETURN_IF_ERROR(future.Await());
+        execute_status.Update(future.Await());
       }
+      ABSL_RETURN_IF_ERROR(execute_status);
 
       const bool upload_active_profiler_session =
           running_options.recreate_profiler_session_between_repeats ||
@@ -762,46 +765,56 @@ CopyArgumentsToDevice(PjRtClient& client,
   ABSL_ASSIGN_OR_RETURN(std::vector<std::shared_ptr<HloModule>> hlo_modules,
                    executable->GetHloModules());
 
-  for (int i = 0; i < num_addressable_devices; ++i) {
-    PjRtDevice* curr_device = addressable_devices[i];
-    int curr_device_id = curr_device->id();
-    // 'source_device' determines where we get the input literal from.
-    PjRtDevice* source_device =
-        addressable_devices[clone_device0_arguments ? 0 : i];
-    int source_device_id = source_device->id();
-    if (!arguments.contains(source_device_id)) {
-      return InvalidArgument(
-          "The provided argument map does not contain arguments "
-          "for device: %d",
-          curr_device_id);
-    }
+  absl::Status status = [&]() -> absl::Status {
+    for (int i = 0; i < num_addressable_devices; ++i) {
+      PjRtDevice* curr_device = addressable_devices[i];
+      int curr_device_id = curr_device->id();
+      // 'source_device' determines where we get the input literal from.
+      PjRtDevice* source_device =
+          addressable_devices[clone_device0_arguments ? 0 : i];
+      int source_device_id = source_device->id();
+      auto it = arguments.find(source_device_id);
+      if (it == arguments.end()) {
+        return InvalidArgument(
+            "The provided argument map does not contain arguments "
+            "for device: %d",
+            curr_device_id);
+      }
 
-    const std::vector<Literal>& curr_device_arguments =
-        arguments.at(source_device_id);
+      const std::vector<Literal>& curr_device_arguments = it->second;
 
-    int executable_idx = hlo_modules.size() == 1
-                             ? 0
-                             : addressable_device_logical_ids[i].partition;
-    HloModule* module = hlo_modules[executable_idx].get();
+      int executable_idx = hlo_modules.size() == 1
+                               ? 0
+                               : addressable_device_logical_ids[i].partition;
+      HloModule* module = hlo_modules[executable_idx].get();
 
-    argument_buffers[i].reserve(curr_device_arguments.size());
-    for (int arg_i = 0; arg_i < curr_device_arguments.size(); ++arg_i) {
-      const Literal& literal = curr_device_arguments[arg_i];
-      if (log_input) {
-        LOG(INFO) << "device_id=" << curr_device_id
-                  << ", input = " << literal.ToString();
+      argument_buffers[i].reserve(curr_device_arguments.size());
+      for (int arg_i = 0; arg_i < curr_device_arguments.size(); ++arg_i) {
+        const Literal& literal = curr_device_arguments[arg_i];
+        if (log_input) {
+          LOG(INFO) << "device_id=" << curr_device_id
+                    << ", input = " << literal.ToString();
+        }
+        ABSL_ASSIGN_OR_RETURN(
+            std::unique_ptr<PjRtBuffer> argument_buffer,
+            buffer_from_host_literal(module, curr_device, arg_i, literal));
+        argument_buffers[i].push_back(std::move(argument_buffer));
       }
-      ABSL_ASSIGN_OR_RETURN(
-          std::unique_ptr<PjRtBuffer> argument_buffer,
-          buffer_from_host_literal(module, curr_device, arg_i, literal));
-      argument_buffers[i].push_back(std::move(argument_buffer));
     }
-  }
+    return absl::OkStatus();
+  }();
+
+  // BufferFromHostLiteral transfers read directly from the host `Literal`s in
+  // `arguments` until each buffer's ready future completes, so every created
+  // buffer must be awaited before returning (including on error, since callers
+  // such as CreateArgumentsOnDevice own `arguments` in their local stack
+  // frame).
   for (const auto& device_argument_buffers : argument_buffers) {
     for (const auto& device_buffer : device_argument_buffers) {
-      ABSL_RETURN_IF_ERROR(device_buffer->GetReadyFuture().Await());
+      status.Update(device_buffer->GetReadyFuture().Await());
     }
   }
+  ABSL_RETURN_IF_ERROR(status);
   return argument_buffers;
 }
 
@@ -883,11 +896,13 @@ CreateUninitializedArgumentsOnDevice(PjRtClient& client,
   }
   LOG(INFO) << "Allocated argument buffers: " << buffer_count;
 
+  absl::Status status;
   for (const auto& argument_buffers : argument_buffers_per_device) {
     for (const auto& buffer : argument_buffers) {
-      ABSL_RETURN_IF_ERROR(buffer->GetReadyFuture().Await());
+      status.Update(buffer->GetReadyFuture().Await());
     }
   }
+  ABSL_RETURN_IF_ERROR(status);
   LOG(INFO) << "Argument buffers are ready.";
 
   return argument_buffers_per_device;
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner_test.cc b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner_test.cc
index 64c1d0e6588..357bc76a31f 100644
--- a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner_test.cc
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner_test.cc
@@ -15,6 +15,8 @@ limitations under the License.
 
 #include "xla/tools/multihost_hlo_runner/functional_hlo_runner.h"
 
+#include <atomic>
+#include <cstddef>
 #include <cstdint>
 #include <cstdlib>
 #include <memory>
@@ -26,6 +28,7 @@ limitations under the License.
 #include "xla/tests/xla_test_backend_predicates.h"
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/functional/any_invocable.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
@@ -37,6 +40,7 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/string_view.h"
+#include "absl/synchronization/notification.h"
 #include "absl/time/time.h"
 #include "absl/types/span.h"
 #include "mlir/IR/BuiltinOps.h"
@@ -44,6 +48,7 @@ limitations under the License.
 #include "mlir/IR/OwningOpRef.h"
 #include "xla/client/executable_build_options.h"
 #include "xla/debug_options_flags.h"
+#include "xla/future.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
 #include "xla/hlo/testlib/filecheck.h"
@@ -51,6 +56,7 @@ limitations under the License.
 #include "xla/literal.h"
 #include "xla/pjrt/maybe_owning_mlir_module.h"
 #include "xla/pjrt/mlir_to_hlo.h"
+#include "xla/pjrt/mock_pjrt_client.h"
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/pjrt/pjrt_compiler.h"
 #include "xla/pjrt/pjrt_executable.h"
@@ -1462,6 +1468,136 @@ TEST_F(FunctionalHloRunnerTest, SingleDeviceHloWithRandomData) {
       {GetHloPath("single_device.hlo")}, InputFormat::kText));
 }
 
+class FakeReadyFuturePjRtBuffer : public PjRtBuffer {
+ public:
+  explicit FakeReadyFuturePjRtBuffer(Future<> ready_future)
+      : ready_future_(std::move(ready_future)) {}
+
+  const Shape& on_device_shape() const override { return shape_; }
+  PjRtMemorySpace* memory_space() const override { return nullptr; }
+  PjRtDevice* device() const override { return nullptr; }
+  PjRtClient* client() const override { return nullptr; }
+  absl::StatusOr<std::unique_ptr<ExternalReference>> AcquireExternalReference()
+      override {
+    return absl::UnimplementedError("Not implemented");
+  }
+  Future<> ToLiteral(MutableLiteralBase* literal) override {
+    return Future<>(absl::UnimplementedError("Not implemented"));
+  }
+  Future<> LazyToLiteral(
+      absl::AnyInvocable<Future<MutableLiteralBase*>() &&> generator) override {
+    return Future<>(absl::UnimplementedError("Not implemented"));
+  }
+  absl::StatusOr<size_t> GetOnDeviceSizeInBytes() const override { return 0; }
+  Future<> CopyRawToHost(void* dst, int64_t offset,
+                         int64_t transfer_size) override {
+    return Future<>(absl::UnimplementedError("Not implemented"));
+  }
+  void Delete() override {}
+  absl::StatusOr<std::unique_ptr<ExternalReference>>
+  ReleaseDeviceMemoryOwnership(bool wait_for_operations_to_complete) override {
+    return absl::UnimplementedError("Not implemented");
+  }
+  bool IsDeleted() const override { return false; }
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> CopyToMemorySpace(
+      PjRtMemorySpace* dst_memory_space) override {
+    return absl::UnimplementedError("Not implemented");
+  }
+  void CopyToRemoteDevice(Future<std::string> serialized_descriptor,
+                          RemoteSendCallback on_done) override {}
+  absl::StatusOr<std::unique_ptr<PjRtBuffer>> Bitcast(
+      xla::PrimitiveType element_type, absl::Span<const int64_t> dims,
+      const Layout* device_layout) override {
+    return absl::UnimplementedError("Not implemented");
+  }
+  Future<> GetReadyFuture() override { return ready_future_; }
+  bool IsOnCpu() const override { return true; }
+
+ private:
+  Shape shape_;
+  Future<> ready_future_;
+};
+
+class HostLiteralMockPjRtClient : public MockPjRtClient {
+ public:
+  MOCK_METHOD(absl::StatusOr<std::unique_ptr<PjRtBuffer>>,
+              BufferFromHostLiteral,
+              (const LiteralSlice& literal, PjRtMemorySpace* memory_space,
+               const Layout* device_layout),
+              (override));
+};
+
+TEST_F(FunctionalHloRunnerTest,
+       CopyArgumentsToDeviceWaitsForPendingTransfersOnError) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<xla::PjRtClient> real_client,
+                       GetPjRtClient());
+  constexpr absl::string_view kTwoArgStablehlo = R"mlir(
+    module {
+      func.func @main(%arg0: tensor<4xi32>, %arg1: tensor<4xi32>)
+          -> tensor<4xi32> {
+        %0 = stablehlo.add %arg0, %arg1 : tensor<4xi32>
+        func.return %0 : tensor<4xi32>
+      }
+    }
+  )mlir";
+  auto context = std::make_unique<mlir::MLIRContext>();
+  ASSERT_OK_AND_ASSIGN(mlir::OwningOpRef<mlir::ModuleOp> stablehlo_module,
+                       xla::ParseMlirModuleString(kTwoArgStablehlo, *context));
+  xla::CompileOptions compile_options;
+  *compile_options.executable_build_options.mutable_debug_options() =
+      GetDebugOptionsFromFlags();
+  compile_options.executable_build_options.set_num_replicas(1);
+  compile_options.executable_build_options.set_num_partitions(1);
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtLoadedExecutable> executable,
+                       real_client->CompileAndLoad(
+                           MaybeOwningMlirModule(std::move(context),
+                                                 std::move(stablehlo_module)),
+                           compile_options));
+
+  HostLiteralMockPjRtClient mock_client;
+  absl::Notification second_arg_reached;
+  std::atomic<bool> first_arg_transfer_finished{false};
+  std::unique_ptr<tsl::Thread> transfer_thread;
+
+  EXPECT_CALL(mock_client, BufferFromHostLiteral)
+      .WillOnce([&](const LiteralSlice& literal, PjRtMemorySpace* memory_space,
+                    const Layout* device_layout)
+                    -> absl::StatusOr<std::unique_ptr<PjRtBuffer>> {
+        absl::Span<const int32_t> host_data = literal.data<int32_t>();
+        auto [promise, future] = MakePromise<>();
+        transfer_thread.reset(tsl::Env::Default()->StartThread(
+            tsl::ThreadOptions(), "async_h2d_transfer",
+            [&second_arg_reached, &first_arg_transfer_finished, host_data,
+             promise = std::move(promise)]() mutable {
+              second_arg_reached.WaitForNotification();
+              // Read from `host_data` to verify the caller kept the host
+              // Literal alive until this transfer's ready future completed.
+              int32_t sum = 0;
+              for (int32_t v : host_data) {
+                sum += v;
+              }
+              CHECK_GE(sum, 0);
+              first_arg_transfer_finished.store(true);
+              promise.Set(absl::OkStatus());
+            }));
+        return std::make_unique<FakeReadyFuturePjRtBuffer>(std::move(future));
+      })
+      .WillOnce([&](const LiteralSlice& literal, PjRtMemorySpace* memory_space,
+                    const Layout* device_layout)
+                    -> absl::StatusOr<std::unique_ptr<PjRtBuffer>> {
+        second_arg_reached.Notify();
+        return absl::ResourceExhaustedError("Staging buffer allocation failed");
+      });
+
+  FunctionalHloRunner::RunningOptions running_options;
+  running_options.module_argument_mode =
+      FunctionalHloRunner::ModuleArgumentMode::kUseZerosAsInput;
+  EXPECT_THAT(FunctionalHloRunner::Run(mock_client, executable.get(),
+                                       /*arguments=*/{}, running_options),
+              absl_testing::StatusIs(absl::StatusCode::kResourceExhausted));
+  EXPECT_TRUE(first_arg_transfer_finished.load());
+}
+
 }  // namespace
 }  // namespace xla
 

From aede00d206ebc4f0c3eebe95a1fe21c404249e5c Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Fri, 16 Oct 2026 11:46:42 -0700
Subject: [PATCH 0971/1000] [XLA:GPU]: Refactor collective emitter helpers for
 variadic collectives.

- Consolidate anonymous namespaces
- extract shared helpers (`GetCollectiveHero`, `GetAllGatherOutputTile`, `EmitDoubleBufferIndex`, and `ExtractCollectiveOperand`)

Functional changes are:
  - GetCollectiveHero allows for epilogues.
  - GreedyPowerOfTwoTiles for all-gather calculates tile sizes based on input
    instead of output.

PiperOrigin-RevId: 1000998488
---
 .../gpu/codegen/triton/collective_emitter.cc  | 748 ++++++++++--------
 .../gpu/codegen/triton/collective_emitter.h   |   5 +
 2 files changed, 411 insertions(+), 342 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
index 2894232d3c7..49ee9d6d4c1 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
@@ -35,6 +35,7 @@ limitations under the License.
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/str_join.h"
+#include "absl/strings/string_view.h"
 #include "absl/types/span.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
@@ -471,19 +472,26 @@ mlir::Value EmitDeviceInvocationCount(mlir::ImplicitLocOpBuilder& b,
       b, counter, arith::ConstantOp::create(b, b.getI32IntegerAttr(1)));
 }
 
-// Returns the element offset of the active half of a double buffer that holds
-// `num_elements` elements of `element_type` per half. The lowest bit of
-// `signal_value` selects the half.
-mlir::Value EmitDoubleBufferOffset(mlir::ImplicitLocOpBuilder& b,
-                                   mlir::Value signal_value,
-                                   int64_t num_elements,
-                                   PrimitiveType element_type) {
+// Picks the active half (0 or 1) of a double buffer from the lowest bit of
+// `signal_value`.
+mlir::Value EmitDoubleBufferIndex(mlir::ImplicitLocOpBuilder& b,
+                                  mlir::Value signal_value) {
   const mlir::Type i64_type = b.getI64Type();
   mlir::Value signal_value_i64 =
       arith::ExtSIOp::create(b, i64_type, signal_value);
   mlir::Value one =
       arith::ConstantOp::create(b, i64_type, b.getI64IntegerAttr(1));
-  mlir::Value buffer_index = arith::AndIOp::create(b, signal_value_i64, one);
+  return arith::AndIOp::create(b, signal_value_i64, one);
+}
+
+// Returns the element offset of the active half (`buffer_index` in {0, 1}) of
+// a double buffer that holds `num_elements` elements of `element_type` per
+// half.
+mlir::Value EmitDoubleBufferOffsetFromIndex(mlir::ImplicitLocOpBuilder& b,
+                                            mlir::Value buffer_index,
+                                            int64_t num_elements,
+                                            PrimitiveType element_type) {
+  const mlir::Type i64_type = b.getI64Type();
   // The allocated double buffer size in bytes is aligned to
   // kXlaAllocatedBufferAlignBytes. To get the offset to the second buffer, we
   // divide the buffer size by 2 and then divide by the element size to get
@@ -498,6 +506,18 @@ mlir::Value EmitDoubleBufferOffset(mlir::ImplicitLocOpBuilder& b,
   return arith::MulIOp::create(b, buffer_index, elements_per_buffer_op);
 }
 
+// Returns the element offset of the active half of a double buffer that holds
+// `num_elements` elements of `element_type` per half. The lowest bit of
+// `signal_value` selects the half.
+mlir::Value EmitDoubleBufferOffset(mlir::ImplicitLocOpBuilder& b,
+                                   mlir::Value signal_value,
+                                   int64_t num_elements,
+                                   PrimitiveType element_type) {
+  mlir::Value buffer_index = EmitDoubleBufferIndex(b, signal_value);
+  return EmitDoubleBufferOffsetFromIndex(b, buffer_index, num_elements,
+                                         element_type);
+}
+
 // Emits instructions to get the pointer to the remote buffer of the given
 // rank.
 // We have a !tt.ptr<i64> pointing to the base of the remote
@@ -547,6 +567,70 @@ void EmitBlockBarrier(
       mtx::BarrierModeAttr::get(b.getContext(), barrier_mode));
 }
 
+struct CollectiveOperandExtractInfo {
+  xtile::TensorValue input_tile;
+  xtile::ExtractTileOp input_extract;
+  mlir::TypedValue<mlir::MemRefType> input_buffer;
+  mlir::Value remote_buffers_arg;
+  PrimitiveType element_type = PrimitiveType::PRIMITIVE_TYPE_INVALID;
+};
+
+absl::StatusOr<CollectiveOperandExtractInfo> ExtractCollectiveOperand(
+    mlir::Value operand, int64_t operand_index, absl::string_view op_name,
+    xtile::EntryFuncOp entry_func, int64_t num_parameters,
+    int32_t remote_buffers_start, bool require_single_use = true) {
+  auto input_tile = mlir::cast<xtile::TensorValue>(operand);
+  auto input_extract = llvm::dyn_cast_if_present<xtile::ExtractTileOp>(
+      input_tile.getDefiningOp());
+  if (!input_extract && input_tile.getDefiningOp() &&
+      input_tile.getDefiningOp()->getNumOperands() > 0) {
+    // Workaround(i1_to_i8_workaround): booleans are stored as i8 and cast to
+    // i1 after ExtractTileOp.
+    input_extract = llvm::dyn_cast_if_present<xtile::ExtractTileOp>(
+        input_tile.getDefiningOp()->getOperand(0).getDefiningOp());
+  }
+  if (!input_extract) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        "%s operand %d must be defined by xtile::ExtractTileOp.", op_name,
+        operand_index));
+  }
+
+  // The collective operand must be read from a fusion parameter. Its argument
+  // number selects the matching remote buffer argument.
+  auto input_arg =
+      mlir::dyn_cast<mlir::BlockArgument>(input_extract.getSource());
+  if (!input_arg ||
+      input_arg.getOwner()->getParentOp() != entry_func.getOperation() ||
+      input_arg.getArgNumber() >= num_parameters) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        "%s operand %d must be extracted from a fusion parameter.", op_name,
+        operand_index));
+  }
+  if (require_single_use &&
+      (!input_extract->hasOneUse() || !input_tile.hasOneUse())) {
+    return absl::InvalidArgumentError(
+        absl::StrFormat("%s input tile %d must only be used by the %s op.",
+                        op_name, operand_index, op_name));
+  }
+
+  PrimitiveType element_type = xla::ConvertMlirTypeToPrimitiveType(
+      input_extract.getType().getElementType());
+  if (element_type == PrimitiveType::PRIMITIVE_TYPE_INVALID) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        "Could not convert %s operand %d element type to PrimitiveType.",
+        op_name, operand_index));
+  }
+
+  CollectiveOperandExtractInfo info;
+  info.input_tile = input_tile;
+  info.input_extract = input_extract;
+  info.input_buffer = input_extract.getSource();
+  info.remote_buffers_arg =
+      entry_func.getArgument(remote_buffers_start + input_arg.getArgNumber());
+  info.element_type = element_type;
+  return info;
+}
+
 class ReductionEmitter {
  protected:
   ReductionEmitter(ReductionEmitterContext& ctx,
@@ -1192,53 +1276,53 @@ class ReduceScatterEmitter : public ReductionEmitter {
   ReduceScatterEmitterContext& ctx_;
 };
 
-}  // namespace
-
-llvm::SmallVector<int64_t> GreedyPowerOfTwoTiles(const Shape& output_shape,
-                                                 int32_t num_blocks) {
-  CHECK_GT(num_blocks, 0) << "num_blocks must be positive. Was " << num_blocks;
-  // Rank fits in int32_t.
-  const auto rank = static_cast<int32_t>(output_shape.dimensions().size());
-  const llvm::ArrayRef<const int64_t> minor_to_major =
-      LayoutUtil::MinorToMajor(output_shape);
-  llvm::SmallVector<int64_t, 4> tile_sizes(rank);
-  // NB: Unsigned because llvm::bit_<> functions expect unsigned.
-  uint64_t remaining_blocks = num_blocks;
-  // Iterate from most major to most minor to keep memory contiguous for each
-  // block.
-  for (int32_t i = rank - 1; i >= 0; --i) {
-    const auto dim = static_cast<int32_t>(minor_to_major[i]);
-    const uint64_t dim_size = output_shape.dimensions(dim);
-    // Largest power-of-two k <= std::min(dim_size, remaining_blocks).
-    const uint64_t k = std::max(
-        uint64_t{1}, llvm::bit_floor(std::min(remaining_blocks, dim_size)));
-    // Round up number of tiles in this dimension to power of two.
-    tile_sizes[dim] =
-        static_cast<int64_t>(llvm::bit_ceil(xla::CeilOfRatio(dim_size, k)));
-    const uint64_t blocks_used =
-        xla::CeilOfRatio(dim_size, static_cast<uint64_t>(tile_sizes[dim]));
-    remaining_blocks = xla::FloorOfRatio(remaining_blocks, blocks_used);
-  }
-  return tile_sizes;
-}
-
-namespace {
-
 // Heuristic maximum warps per block across collective kernels.
 // Triton expects a power of 2 for num_warps.
 constexpr int64_t kMaxWarpsPerBlock = 16;
 
 BlockLevelFusionConfig CreateBlockLevelFusionConfig(
-    int64_t num_warps, absl::Span<const int64_t> tile_sizes) {
+    int64_t num_warps,
+    absl::Span<const llvm::SmallVector<int64_t>> output_tiles) {
   BlockLevelFusionConfig config;
   config.set_num_warps(num_warps);
   config.set_num_ctas(1);    // No block-level clustering.
   config.set_num_stages(1);  // No pipelining of loops.
-  config.add_output_tiles()->mutable_sizes()->Assign(tile_sizes.begin(),
-                                                     tile_sizes.end());
+  for (const llvm::SmallVector<int64_t>& tile_sizes : output_tiles) {
+    config.add_output_tiles()->mutable_sizes()->Assign(tile_sizes.begin(),
+                                                       tile_sizes.end());
+  }
   return config;
 }
 
+BlockLevelFusionConfig CreateBlockLevelFusionConfig(
+    int64_t num_warps, const llvm::SmallVector<int64_t>& tile_sizes) {
+  return CreateBlockLevelFusionConfig(
+      num_warps, absl::MakeConstSpan(&tile_sizes, /*size=*/1));
+}
+
+// Computes the output tile for an all-gather output of `output_shape`.
+// We divide `gather_dim` by `num_devices` to get one rank's slice and tile
+// that across `blocks_per_peer` programs, so each tile stays within a single
+// peer.
+absl::StatusOr<llvm::SmallVector<int64_t>> GetAllGatherOutputTile(
+    const Shape& output_shape, int64_t gather_dim, int64_t num_devices,
+    int32_t blocks_per_peer) {
+  TF_RET_CHECK(output_shape.dimensions(gather_dim) % num_devices == 0)
+      << "All-gather output " << output_shape.ToString()
+      << " must be divisible by " << num_devices
+      << " along the gather dimension " << gather_dim;
+  const int64_t per_rank_size =
+      output_shape.dimensions(gather_dim) / num_devices;
+  Shape per_peer_shape = output_shape;
+  per_peer_shape.set_dimensions(gather_dim, per_rank_size);
+  llvm::SmallVector<int64_t> tile_sizes =
+      GreedyPowerOfTwoTiles(per_peer_shape, blocks_per_peer);
+  tile_sizes[gather_dim] = std::min(tile_sizes[gather_dim],
+                                    static_cast<int64_t>(llvm::bit_floor(
+                                        static_cast<uint64_t>(per_rank_size))));
+  return tile_sizes;
+}
+
 absl::StatusOr<BlockLevelFusionConfig> GetBlockLevelConfigForAllReduce(
     const GpuTopology& gpu_topology, const HloAllReduceInstruction* instr,
     const DeviceAssignment* device_assignment) {
@@ -1284,8 +1368,6 @@ absl::StatusOr<BlockLevelFusionConfig> GetBlockLevelConfigForAllGather(
   ABSL_ASSIGN_OR_RETURN(const AllGatherInfo all_gather_info,
                    BuildAllGatherInfo(gpu_topology, instr, device_assignment));
 
-  const Shape& output_shape = instr->shape();
-
   const int64_t elements_per_warp =
       WarpSize(gpu_topology.gpu_target_config().device_description) *
       se::gpu::kNumElementsPerThread;
@@ -1300,19 +1382,12 @@ absl::StatusOr<BlockLevelFusionConfig> GetBlockLevelConfigForAllGather(
   const int64_t num_blocks =
       std::min(kAllGatherMaxBlocksPerGrid, CeilOfRatio(total_warps, num_warps));
 
-  llvm::SmallVector<int64_t> tile_sizes =
-      GreedyPowerOfTwoTiles(output_shape, num_blocks);
-
-  // For AllGather, the tile on the gather dimension must not exceed the
-  // per-rank size. Otherwise a single tile would span multiple replicas,
-  // but the tiling framework assigns a single replica_id per tile.
-  const int64_t gather_dim = instr->all_gather_dimension();
-  const int64_t per_rank_size =
-      output_shape.dimensions(gather_dim) / all_gather_info.num_devices;
-  tile_sizes[gather_dim] = std::min(tile_sizes[gather_dim],
-                                    static_cast<int64_t>(llvm::bit_floor(
-                                        static_cast<uint64_t>(per_rank_size))));
-
+  const int32_t blocks_per_peer = std::max<int32_t>(
+      1, static_cast<int32_t>(num_blocks / all_gather_info.num_devices));
+  ABSL_ASSIGN_OR_RETURN(
+      llvm::SmallVector<int64_t> tile_sizes,
+      GetAllGatherOutputTile(instr->shape(), instr->all_gather_dimension(),
+                             all_gather_info.num_devices, blocks_per_peer));
   return CreateBlockLevelFusionConfig(num_warps, tile_sizes);
 }
 
@@ -1359,197 +1434,6 @@ absl::StatusOr<BlockLevelFusionConfig> GetBlockLevelConfigForReduceScatter(
   return CreateBlockLevelFusionConfig(num_warps, tile_sizes);
 }
 
-}  // namespace
-
-absl::StatusOr<BlockLevelFusionConfig> GetCollectiveBlockLevelFusionConfig(
-    const GpuTopology& gpu_topology, const HloFusionInstruction* fusion_instr,
-    const DeviceAssignment* device_assignment) {
-  const HloInstruction* instr = fusion_instr->fused_expression_root();
-  ABSL_ASSIGN_OR_RETURN(GpuBackendConfig gpu_config,
-                   instr->backend_config<GpuBackendConfig>());
-  if (!IsTritonCollectiveKernel(
-          gpu_config.collective_backend_config().kernel_strategy())) {
-    return absl::InvalidArgumentError(absl::StrFormat(
-        "Instruction was not annotated with Triton strategy but still called "
-        "for Triton codegen. %s",
-        instr->ToString()));
-  }
-  if (!gpu_topology.has_gpu_target_config()) {
-    return absl::InvalidArgumentError(
-        "GpuTopology must have a target config to build "
-        "BlockLevelFusionConfig.");
-  }
-
-  if (const HloAllReduceInstruction* all_reduce =
-          DynCast<HloAllReduceInstruction>(instr)) {
-    return GetBlockLevelConfigForAllReduce(gpu_topology, all_reduce,
-                                           device_assignment);
-  }
-
-  if (const HloAllGatherInstruction* all_gather =
-          DynCast<HloAllGatherInstruction>(instr)) {
-    return GetBlockLevelConfigForAllGather(gpu_topology, all_gather,
-                                           device_assignment);
-  }
-
-  if (const HloReduceScatterInstruction* reduce_scatter =
-          DynCast<HloReduceScatterInstruction>(instr)) {
-    return GetBlockLevelConfigForReduceScatter(gpu_topology, reduce_scatter,
-                                               device_assignment);
-  }
-
-  return absl::UnimplementedError(
-      absl::StrFormat("Unsupported collective opcode for Triton codegen: %s",
-                      HloOpcodeString(instr->opcode())));
-}
-
-absl::Status TrySetGpuBackendConfigForCollective(
-    const GpuTopology& gpu_topology, HloFusionInstruction* fusion_instr,
-    const DeviceAssignment* device_assignment) {
-  ABSL_ASSIGN_OR_RETURN(BlockLevelFusionConfig block_config,
-                   GetCollectiveBlockLevelFusionConfig(
-                       gpu_topology, fusion_instr, device_assignment));
-  ABSL_ASSIGN_OR_RETURN(GpuBackendConfig gpu_backend_config,
-                   fusion_instr->backend_config<GpuBackendConfig>());
-  gpu_backend_config.mutable_fusion_backend_config()->set_kind(
-      kTritonCollectiveFusionKind);
-  *gpu_backend_config.mutable_fusion_backend_config()
-       ->mutable_block_level_fusion_config() = std::move(block_config);
-  ABSL_RETURN_IF_ERROR(
-      fusion_instr->set_backend_config(std::move(gpu_backend_config)));
-  return absl::OkStatus();
-}
-
-absl::StatusOr<std::vector<Shape>> GetCollectiveUnmanagedKernelArguments(
-    const HloFusionInstruction* fusion) {
-  const HloComputation* computation = fusion->fused_instructions_computation();
-  const HloInstruction* root = computation->root_instruction();
-  switch (root->opcode()) {
-    case HloOpcode::kAllReduce:
-    case HloOpcode::kAllGather:
-    case HloOpcode::kAllToAll:
-    case HloOpcode::kReduceScatter:
-      return GetRemoteBufferUnmanagedKernelArguments(
-          computation, Cast<HloCollectiveInstruction>(root));
-    default:
-      return std::vector<Shape>();
-  }
-}
-
-absl::StatusOr<int32_t> AddCollectiveMetadataArguments(
-    llvm::SmallVector<mlir::Type>& fn_arg_types, mlir::ImplicitLocOpBuilder& b,
-    const HloComputation* hlo_computation) {
-  const bool has_invocation_count =
-      hlo_computation->root_instruction()->opcode() !=
-      HloOpcode::kReduceScatter;
-  // rank: i32
-  fn_arg_types.push_back(b.getI32Type());
-  if (has_invocation_count) {
-    // signal_value: i32
-    fn_arg_types.push_back(b.getI32Type());
-  }
-  // signal_buffers: !tt.ptr<i64>
-  fn_arg_types.push_back(
-      ttir::PointerType::get(b.getI64Type(), kGlobalAddressSpace));
-
-  for (HloInstruction* p : hlo_computation->parameter_instructions()) {
-    (void)p;
-    // Also add the remote/scratch buffers for collectives.
-    // !tt.ptr<i64>
-    fn_arg_types.push_back(
-        ttir::PointerType::get(b.getI64Type(), kGlobalAddressSpace));
-  }
-  const int32_t num_metadata_args = has_invocation_count
-                                        ? kNumCollectiveMetadataArgs
-                                        : kNumCollectiveMetadataArgs - 1;
-  return hlo_computation->num_parameters() + num_metadata_args;
-}
-
-mlir::LogicalResult RewriteAllReduce(mlir::stablehlo::AllReduceOp op,
-                                     mlir::PatternRewriter& rewriter) {
-  const mlir::Location loc = op->getLoc();
-  absl::StatusOr<AllReduceEmitterContext> maybe_context =
-      CreateAllReduceEmitterContext(op);
-  if (!maybe_context.ok()) {
-    VLOG(3) << "Failed to create AllReduceEmitterContext: "
-            << maybe_context.status().message();
-    return rewriter.notifyMatchFailure(
-        loc, absl::StrCat("Failed to create AllReduceEmitterContext: ",
-                          maybe_context.status().message()));
-  }
-  VLOG(3) << "AllReduceEmitter::Emit using strategy: "
-          << maybe_context->strategy;
-  return AllReduceEmitter::Emit(*std::move(maybe_context), rewriter);
-}
-
-mlir::LogicalResult RewriteReduceScatter(xtile::ReduceScatterOp op,
-                                         mlir::PatternRewriter& rewriter) {
-  const mlir::Location loc = op->getLoc();
-  absl::StatusOr<ReduceScatterEmitterContext> maybe_context =
-      CreateReduceScatterEmitterContext(op);
-  if (!maybe_context.ok()) {
-    VLOG(3) << "Failed to create ReduceScatterEmitterContext: "
-            << maybe_context.status().message();
-    return rewriter.notifyMatchFailure(
-        loc, absl::StrCat("Failed to create ReduceScatterEmitterContext: ",
-                          maybe_context.status().message()));
-  }
-  return ReduceScatterEmitter::Emit(*std::move(maybe_context), rewriter);
-}
-
-absl::StatusOr<SymmetricMemoryType> GetSymmetricMemoryType(
-    const GpuTopology& gpu_topology, const HloInstruction& instr,
-    const DeviceAssignment* device_assignment) {
-  const HloInstruction* collective = &instr;
-  if (instr.opcode() == HloOpcode::kFusion) {
-    collective = instr.fused_instructions_computation()->root_instruction();
-  }
-  ABSL_ASSIGN_OR_RETURN(
-      const bool is_collective_single_host,
-      IsCollectiveSingleHost(gpu_topology, *collective, device_assignment));
-  ABSL_ASSIGN_OR_RETURN(const bool is_cross_host_possible,
-                   IsCrossHostCollectiveKernelPossible(
-                       gpu_topology, *collective, device_assignment));
-  const bool is_lsa_possible = IsLsaPossible(gpu_topology);
-
-  if (!is_collective_single_host) {
-    if (is_cross_host_possible && is_lsa_possible) {
-      return SymmetricMemoryType::kLoadStoreAccessible;
-    }
-    return absl::UnimplementedError(
-        "Cross-host symmetric memory collectives are not supported.");
-  }
-  return is_lsa_possible ? SymmetricMemoryType::kLoadStoreAccessible
-                         : SymmetricMemoryType::kXlaRendezvous;
-}
-
-absl::StatusOr<CollectiveKernelSpec> CreateCollectiveKernelSpec(
-    const HloInstruction* instr, const LaunchDimensions& launch_dimensions,
-    SymmetricMemoryType scratch_memory_type) {
-  const HloInstruction* collective = instr;
-  if (instr->opcode() == HloOpcode::kFusion) {
-    collective = instr->fused_instructions_computation()->root_instruction();
-  }
-  switch (collective->opcode()) {
-    case HloOpcode::kAllReduce:
-      return CreateAllReduceKernelSpec(collective, launch_dimensions,
-                                       scratch_memory_type);
-    case HloOpcode::kAllGather:
-      return CreateAllGatherKernelSpec(collective, launch_dimensions,
-                                       scratch_memory_type);
-    case HloOpcode::kReduceScatter:
-      return CreateReduceScatterKernelSpec(collective, launch_dimensions,
-                                           scratch_memory_type);
-    default:
-      return absl::UnimplementedError(
-          absl::StrFormat("CollectiveKernelSpec creation not implemented for "
-                          "opcode: %s",
-                          HloOpcodeString(collective->opcode())));
-  }
-}
-
-namespace {
-
 struct AllGatherRewriteContext {
   mlir::stablehlo::AllGatherOp op;
   xtile::ExtractTileOp input_extract;
@@ -1596,40 +1480,6 @@ absl::StatusOr<AllGatherRewriteContext> BuildAllGatherContext(
         "the remote buffer arguments.");
   }
 
-  mlir::Value input_tile = op.getOperand(0);
-  auto input_extract = llvm::dyn_cast_if_present<xtile::ExtractTileOp>(
-      input_tile.getDefiningOp());
-  if (!input_extract && input_tile.getDefiningOp() &&
-      input_tile.getDefiningOp()->getNumOperands() > 0) {
-    // Workaround(i1_to_i8_workaround): booleans are stored as i8 and cast to
-    // i1 after ExtractTileOp.
-    input_extract = llvm::dyn_cast_if_present<xtile::ExtractTileOp>(
-        input_tile.getDefiningOp()->getOperand(0).getDefiningOp());
-  }
-  if (!input_extract) {
-    return absl::InvalidArgumentError(
-        "AllGather operand must be defined by xtile::ExtractTileOp.");
-  }
-
-  // The gathered operand must be read from a fusion parameter. Its argument
-  // number selects the matching remote buffer argument.
-  auto input_arg =
-      mlir::dyn_cast<mlir::BlockArgument>(input_extract.getSource());
-  const int64_t num_parameters =
-      num_opaque_attr.getInt() - kNumCollectiveMetadataArgs;
-  if (!input_arg ||
-      input_arg.getOwner()->getParentOp() != entry_func.getOperation() ||
-      input_arg.getArgNumber() >= num_parameters) {
-    return absl::InvalidArgumentError(
-        "AllGather operand must be extracted from a fusion parameter.");
-  }
-  // `input_extract` is redirected to the peer's buffer below, so the local
-  // tile must not have other users.
-  if (!input_extract->hasOneUse() || !input_tile.hasOneUse()) {
-    return absl::InvalidArgumentError(
-        "AllGather input tile must only be used by the AllGather op.");
-  }
-
   ABSL_ASSIGN_OR_RETURN(auto replica_groups,
                    xla::ConvertReplicaGroups(op.getReplicaGroups(), op));
   if (replica_groups->replica_groups().empty()) {
@@ -1637,22 +1487,27 @@ absl::StatusOr<AllGatherRewriteContext> BuildAllGatherContext(
         "AllGather replica groups must not be empty.");
   }
 
-  PrimitiveType element_type = xla::ConvertMlirTypeToPrimitiveType(
-      input_extract.getType().getElementType());
-  if (element_type == PrimitiveType::PRIMITIVE_TYPE_INVALID) {
-    return absl::InvalidArgumentError(
-        "Could not convert AllGather element type to PrimitiveType.");
-  }
+  const int64_t num_parameters =
+      num_opaque_attr.getInt() - kNumCollectiveMetadataArgs;
+  int32_t total_args = entry_func.getNumArguments();
+  int32_t metadata_args_start =
+      total_args - num_opaque_attr.getInt() - kNumTileIndexArgs;
+  ABSL_ASSIGN_OR_RETURN(
+      CollectiveOperandExtractInfo extract_info,
+      ExtractCollectiveOperand(op.getOperand(0), /*operand_index=*/0,
+                               "AllGather", entry_func, num_parameters,
+                               metadata_args_start + kNumCollectiveMetadataArgs,
+                               /*require_single_use=*/true));
 
   AllGatherRewriteContext ctx;
   ctx.op = op;
-  ctx.input_extract = input_extract;
-  ctx.input_buffer = input_extract.getSource();
+  ctx.input_extract = extract_info.input_extract;
+  ctx.input_buffer = extract_info.input_buffer;
   ctx.world_size = replica_groups->num_devices_per_group();
   ctx.gather_dim = op.getAllGatherDim();
-  ctx.element_type = element_type;
+  ctx.element_type = extract_info.element_type;
   ctx.per_rank_shape = ctx.input_buffer.getType().getShape();
-  ctx.tile_sizes = input_extract.getFullTileShape();
+  ctx.tile_sizes = extract_info.input_extract.getFullTileShape();
 
   // Every tile must be gathered from a single rank (see EmitPeerReplicaId).
   const int64_t per_rank_size = ctx.per_rank_shape[ctx.gather_dim];
@@ -1664,17 +1519,12 @@ absl::StatusOr<AllGatherRewriteContext> BuildAllGatherContext(
         tile_size, per_rank_size));
   }
 
-  int32_t total_args = entry_func.getNumArguments();
-  int32_t metadata_args_start =
-      total_args - num_opaque_attr.getInt() - kNumTileIndexArgs;
   // Layout: opaque[0]=rank, opaque[1]=invocation count (unused, the signal
   // value comes from a counter in device memory), opaque[2]=signal_buffers,
   // opaque[3 + i]=remote buffers of parameter i (same as AllReduce).
   ctx.rank_arg = entry_func.getArgument(metadata_args_start);
   ctx.signal_buffers_arg = entry_func.getArgument(metadata_args_start + 2);
-  ctx.remote_buffers_arg =
-      entry_func.getArgument(metadata_args_start + kNumCollectiveMetadataArgs +
-                             input_arg.getArgNumber());
+  ctx.remote_buffers_arg = extract_info.remote_buffers_arg;
 
   // Find a dimension to split the tile size across `world_size` to balance
   // the in-kernel D2D subtile copy work across participating ranks.
@@ -1833,34 +1683,17 @@ absl::StatusOr<AllToAllRewriteContext> BuildAllToAllContext(
     input_tile = cast_in.getInputs()[0];
   }
 
-  auto input_extract = llvm::dyn_cast_if_present<xtile::ExtractTileOp>(
-      input_tile.getDefiningOp());
-  if (!input_extract && input_tile.getDefiningOp() &&
-      input_tile.getDefiningOp()->getNumOperands() > 0) {
-    // Workaround(i1_to_i8_workaround): booleans are stored as i8 and cast to
-    // i1 after ExtractTileOp.
-    input_extract = llvm::dyn_cast_if_present<xtile::ExtractTileOp>(
-        input_tile.getDefiningOp()->getOperand(0).getDefiningOp());
-  }
-  if (!input_extract) {
-    return absl::InvalidArgumentError(
-        "AllToAll operand must be defined by xtile::ExtractTileOp.");
-  }
-
-  auto input_arg =
-      mlir::dyn_cast<mlir::BlockArgument>(input_extract.getSource());
   const int64_t num_parameters =
       num_opaque_attr.getInt() - kNumCollectiveMetadataArgs;
-  if (!input_arg ||
-      input_arg.getOwner()->getParentOp() != entry_func.getOperation() ||
-      input_arg.getArgNumber() >= num_parameters) {
-    return absl::InvalidArgumentError(
-        "AllToAll operand must be extracted from a fusion parameter.");
-  }
-  if (!input_extract->hasOneUse() || !input_tile.hasOneUse()) {
-    return absl::InvalidArgumentError(
-        "AllToAll input tile must only be used by the AllToAll op.");
-  }
+  int32_t total_args = entry_func.getNumArguments();
+  int32_t metadata_args_start =
+      total_args - num_opaque_attr.getInt() - kNumTileIndexArgs;
+  ABSL_ASSIGN_OR_RETURN(
+      CollectiveOperandExtractInfo extract_info,
+      ExtractCollectiveOperand(input_tile, /*operand_index=*/0, "AllToAll",
+                               entry_func, num_parameters,
+                               metadata_args_start + kNumCollectiveMetadataArgs,
+                               /*require_single_use=*/true));
 
   ABSL_ASSIGN_OR_RETURN(auto replica_groups,
                    xla::ConvertReplicaGroups(op.getReplicaGroups(), op));
@@ -1869,22 +1702,15 @@ absl::StatusOr<AllToAllRewriteContext> BuildAllToAllContext(
         "AllToAll replica groups must not be empty.");
   }
 
-  PrimitiveType element_type = xla::ConvertMlirTypeToPrimitiveType(
-      input_extract.getType().getElementType());
-  if (element_type == PrimitiveType::PRIMITIVE_TYPE_INVALID) {
-    return absl::InvalidArgumentError(
-        "Could not convert AllToAll element type to PrimitiveType.");
-  }
-
   AllToAllRewriteContext ctx;
   ctx.op = op;
-  ctx.input_extract = input_extract;
-  ctx.input_buffer = input_extract.getSource();
+  ctx.input_extract = extract_info.input_extract;
+  ctx.input_buffer = extract_info.input_buffer;
   ctx.world_size = replica_groups->num_devices_per_group();
   ctx.split_dim = op.getSplitDimension();
-  ctx.element_type = element_type;
+  ctx.element_type = extract_info.element_type;
   ctx.input_shape = ctx.input_buffer.getType().getShape();
-  ctx.tile_sizes = input_extract.getFullTileShape();
+  ctx.tile_sizes = extract_info.input_extract.getFullTileShape();
   if (ctx.world_size <= 0 ||
       ctx.input_shape[ctx.split_dim] % ctx.world_size != 0) {
     return absl::InvalidArgumentError(
@@ -1904,14 +1730,9 @@ absl::StatusOr<AllToAllRewriteContext> BuildAllToAllContext(
         tile_size, ctx.shard_size));
   }
 
-  int32_t total_args = entry_func.getNumArguments();
-  int32_t metadata_args_start =
-      total_args - num_opaque_attr.getInt() - kNumTileIndexArgs;
   ctx.rank_arg = entry_func.getArgument(metadata_args_start);
   ctx.signal_buffers_arg = entry_func.getArgument(metadata_args_start + 2);
-  ctx.remote_buffers_arg =
-      entry_func.getArgument(metadata_args_start + kNumCollectiveMetadataArgs +
-                             input_arg.getArgNumber());
+  ctx.remote_buffers_arg = extract_info.remote_buffers_arg;
 
   ctx.signal_stride = ctx.shard_size / tile_size;
   for (int64_t j = ctx.split_dim + 1;
@@ -1970,6 +1791,249 @@ void CopyTileToScratch(mlir::ImplicitLocOpBuilder& b,
 
 }  // namespace
 
+llvm::SmallVector<int64_t> GreedyPowerOfTwoTiles(const Shape& output_shape,
+                                                 int32_t num_blocks) {
+  CHECK_GT(num_blocks, 0) << "num_blocks must be positive. Was " << num_blocks;
+  // Rank fits in int32_t.
+  const auto rank = static_cast<int32_t>(output_shape.dimensions().size());
+  const llvm::ArrayRef<const int64_t> minor_to_major =
+      LayoutUtil::MinorToMajor(output_shape);
+  llvm::SmallVector<int64_t, 4> tile_sizes(rank);
+  // NB: Unsigned because llvm::bit_<> functions expect unsigned.
+  uint64_t remaining_blocks = num_blocks;
+  // Iterate from most major to most minor to keep memory contiguous for each
+  // block.
+  for (int32_t i = rank - 1; i >= 0; --i) {
+    const auto dim = static_cast<int32_t>(minor_to_major[i]);
+    const uint64_t dim_size = output_shape.dimensions(dim);
+    // Largest power-of-two k <= std::min(dim_size, remaining_blocks).
+    const uint64_t k = std::max(
+        uint64_t{1}, llvm::bit_floor(std::min(remaining_blocks, dim_size)));
+    // Round up number of tiles in this dimension to power of two.
+    tile_sizes[dim] =
+        static_cast<int64_t>(llvm::bit_ceil(xla::CeilOfRatio(dim_size, k)));
+    const uint64_t blocks_used =
+        xla::CeilOfRatio(dim_size, static_cast<uint64_t>(tile_sizes[dim]));
+    remaining_blocks = xla::FloorOfRatio(remaining_blocks, blocks_used);
+  }
+  return tile_sizes;
+}
+
+const HloInstruction* GetCollectiveHero(const HloComputation* computation) {
+  if (computation == nullptr) {
+    return nullptr;
+  }
+  auto is_collective = [](const HloInstruction* instr) {
+    return instr->opcode() == HloOpcode::kAllReduce ||
+           instr->opcode() == HloOpcode::kAllGather ||
+           instr->opcode() == HloOpcode::kAllToAll ||
+           instr->opcode() == HloOpcode::kReduceScatter;
+  };
+  const HloInstruction* root = computation->root_instruction();
+  if (is_collective(root)) {
+    return root;
+  }
+  for (const HloInstruction* instr : computation->instructions()) {
+    if (is_collective(instr)) {
+      return instr;
+    }
+  }
+  return nullptr;
+}
+
+absl::StatusOr<BlockLevelFusionConfig> GetCollectiveBlockLevelFusionConfig(
+    const GpuTopology& gpu_topology, const HloFusionInstruction* fusion_instr,
+    const DeviceAssignment* device_assignment) {
+  const HloInstruction* instr =
+      GetCollectiveHero(fusion_instr->fused_instructions_computation());
+  if (instr == nullptr) {
+    return absl::UnimplementedError(absl::StrFormat(
+        "No collective hero found for Triton codegen: %s",
+        fusion_instr->fused_instructions_computation()->ToString()));
+  }
+  ABSL_ASSIGN_OR_RETURN(GpuBackendConfig gpu_config,
+                   instr->backend_config<GpuBackendConfig>());
+  if (!IsTritonCollectiveKernel(
+          gpu_config.collective_backend_config().kernel_strategy())) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        "Instruction was not annotated with Triton strategy but still called "
+        "for Triton codegen. %s",
+        instr->ToString()));
+  }
+  if (!gpu_topology.has_gpu_target_config()) {
+    return absl::InvalidArgumentError(
+        "GpuTopology must have a target config to build "
+        "BlockLevelFusionConfig.");
+  }
+
+  if (const HloAllReduceInstruction* all_reduce =
+          DynCast<HloAllReduceInstruction>(instr)) {
+    return GetBlockLevelConfigForAllReduce(gpu_topology, all_reduce,
+                                           device_assignment);
+  }
+
+  if (const HloAllGatherInstruction* all_gather =
+          DynCast<HloAllGatherInstruction>(instr)) {
+    return GetBlockLevelConfigForAllGather(gpu_topology, all_gather,
+                                           device_assignment);
+  }
+
+  if (const HloReduceScatterInstruction* reduce_scatter =
+          DynCast<HloReduceScatterInstruction>(instr)) {
+    return GetBlockLevelConfigForReduceScatter(gpu_topology, reduce_scatter,
+                                               device_assignment);
+  }
+
+  return absl::UnimplementedError(
+      absl::StrFormat("Unsupported collective opcode for Triton codegen: %s",
+                      HloOpcodeString(instr->opcode())));
+}
+
+absl::Status TrySetGpuBackendConfigForCollective(
+    const GpuTopology& gpu_topology, HloFusionInstruction* fusion_instr,
+    const DeviceAssignment* device_assignment) {
+  ABSL_ASSIGN_OR_RETURN(BlockLevelFusionConfig block_config,
+                   GetCollectiveBlockLevelFusionConfig(
+                       gpu_topology, fusion_instr, device_assignment));
+  ABSL_ASSIGN_OR_RETURN(GpuBackendConfig gpu_backend_config,
+                   fusion_instr->backend_config<GpuBackendConfig>());
+  gpu_backend_config.mutable_fusion_backend_config()->set_kind(
+      kTritonCollectiveFusionKind);
+  *gpu_backend_config.mutable_fusion_backend_config()
+       ->mutable_block_level_fusion_config() = std::move(block_config);
+  ABSL_RETURN_IF_ERROR(
+      fusion_instr->set_backend_config(std::move(gpu_backend_config)));
+  return absl::OkStatus();
+}
+
+absl::StatusOr<std::vector<Shape>> GetCollectiveUnmanagedKernelArguments(
+    const HloFusionInstruction* fusion) {
+  const HloComputation* computation = fusion->fused_instructions_computation();
+  const HloInstruction* collective = GetCollectiveHero(computation);
+  if (collective == nullptr) {
+    return std::vector<Shape>();
+  }
+  return GetRemoteBufferUnmanagedKernelArguments(
+      computation, Cast<HloCollectiveInstruction>(collective));
+}
+
+absl::StatusOr<int32_t> AddCollectiveMetadataArguments(
+    llvm::SmallVector<mlir::Type>& fn_arg_types, mlir::ImplicitLocOpBuilder& b,
+    const HloComputation* hlo_computation) {
+  const HloInstruction* collective = GetCollectiveHero(hlo_computation);
+  const bool has_invocation_count =
+      collective == nullptr ||
+      collective->opcode() != HloOpcode::kReduceScatter;
+  // rank: i32
+  fn_arg_types.push_back(b.getI32Type());
+  if (has_invocation_count) {
+    // signal_value: i32
+    fn_arg_types.push_back(b.getI32Type());
+  }
+  // signal_buffers: !tt.ptr<i64>
+  fn_arg_types.push_back(
+      ttir::PointerType::get(b.getI64Type(), kGlobalAddressSpace));
+
+  for (HloInstruction* p : hlo_computation->parameter_instructions()) {
+    (void)p;
+    // Also add the remote/scratch buffers for collectives.
+    // !tt.ptr<i64>
+    fn_arg_types.push_back(
+        ttir::PointerType::get(b.getI64Type(), kGlobalAddressSpace));
+  }
+  const int32_t num_metadata_args = has_invocation_count
+                                        ? kNumCollectiveMetadataArgs
+                                        : kNumCollectiveMetadataArgs - 1;
+  return hlo_computation->num_parameters() + num_metadata_args;
+}
+
+mlir::LogicalResult RewriteAllReduce(mlir::stablehlo::AllReduceOp op,
+                                     mlir::PatternRewriter& rewriter) {
+  const mlir::Location loc = op->getLoc();
+  absl::StatusOr<AllReduceEmitterContext> maybe_context =
+      CreateAllReduceEmitterContext(op);
+  if (!maybe_context.ok()) {
+    VLOG(3) << "Failed to create AllReduceEmitterContext: "
+            << maybe_context.status().message();
+    return rewriter.notifyMatchFailure(
+        loc, absl::StrCat("Failed to create AllReduceEmitterContext: ",
+                          maybe_context.status().message()));
+  }
+  VLOG(3) << "AllReduceEmitter::Emit using strategy: "
+          << maybe_context->strategy;
+  return AllReduceEmitter::Emit(*std::move(maybe_context), rewriter);
+}
+
+mlir::LogicalResult RewriteReduceScatter(xtile::ReduceScatterOp op,
+                                         mlir::PatternRewriter& rewriter) {
+  const mlir::Location loc = op->getLoc();
+  absl::StatusOr<ReduceScatterEmitterContext> maybe_context =
+      CreateReduceScatterEmitterContext(op);
+  if (!maybe_context.ok()) {
+    VLOG(3) << "Failed to create ReduceScatterEmitterContext: "
+            << maybe_context.status().message();
+    return rewriter.notifyMatchFailure(
+        loc, absl::StrCat("Failed to create ReduceScatterEmitterContext: ",
+                          maybe_context.status().message()));
+  }
+  return ReduceScatterEmitter::Emit(*std::move(maybe_context), rewriter);
+}
+
+absl::StatusOr<SymmetricMemoryType> GetSymmetricMemoryType(
+    const GpuTopology& gpu_topology, const HloInstruction& instr,
+    const DeviceAssignment* device_assignment) {
+  const HloInstruction* collective = &instr;
+  if (instr.opcode() == HloOpcode::kFusion) {
+    collective = GetCollectiveHero(instr.fused_instructions_computation());
+    TF_RET_CHECK(collective != nullptr)
+        << "No collective hero found in fusion: " << instr.ToString();
+  }
+  ABSL_ASSIGN_OR_RETURN(
+      const bool is_collective_single_host,
+      IsCollectiveSingleHost(gpu_topology, *collective, device_assignment));
+  ABSL_ASSIGN_OR_RETURN(const bool is_cross_host_possible,
+                   IsCrossHostCollectiveKernelPossible(
+                       gpu_topology, *collective, device_assignment));
+  const bool is_lsa_possible = IsLsaPossible(gpu_topology);
+
+  if (!is_collective_single_host) {
+    if (is_cross_host_possible && is_lsa_possible) {
+      return SymmetricMemoryType::kLoadStoreAccessible;
+    }
+    return absl::UnimplementedError(
+        "Cross-host symmetric memory collectives are not supported.");
+  }
+  return is_lsa_possible ? SymmetricMemoryType::kLoadStoreAccessible
+                         : SymmetricMemoryType::kXlaRendezvous;
+}
+
+absl::StatusOr<CollectiveKernelSpec> CreateCollectiveKernelSpec(
+    const HloInstruction* instr, const LaunchDimensions& launch_dimensions,
+    SymmetricMemoryType scratch_memory_type) {
+  const HloInstruction* collective = instr;
+  if (instr->opcode() == HloOpcode::kFusion) {
+    collective = GetCollectiveHero(instr->fused_instructions_computation());
+    TF_RET_CHECK(collective != nullptr)
+        << "No collective hero found in fusion: " << instr->ToString();
+  }
+  switch (collective->opcode()) {
+    case HloOpcode::kAllReduce:
+      return CreateAllReduceKernelSpec(collective, launch_dimensions,
+                                       scratch_memory_type);
+    case HloOpcode::kAllGather:
+      return CreateAllGatherKernelSpec(collective, launch_dimensions,
+                                       scratch_memory_type);
+    case HloOpcode::kReduceScatter:
+      return CreateReduceScatterKernelSpec(collective, launch_dimensions,
+                                           scratch_memory_type);
+    default:
+      return absl::UnimplementedError(
+          absl::StrFormat("CollectiveKernelSpec creation not implemented for "
+                          "opcode: %s",
+                          HloOpcodeString(collective->opcode())));
+  }
+}
+
 mlir::LogicalResult RewriteAllGather(mlir::stablehlo::AllGatherOp op,
                                      mlir::PatternRewriter& rewriter) {
   absl::StatusOr<AllGatherRewriteContext> maybe_ctx = BuildAllGatherContext(op);
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h
index 7f86608f0ae..126bf761a3f 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.h
@@ -125,5 +125,10 @@ absl::StatusOr<CollectiveKernelSpec> CreateCollectiveKernelSpec(
 absl::Status FlattenReduceScatterFusion(
     HloFusionInstruction* absl_nonnull fusion_instr);
 
+// Returns the collective instruction (all-reduce, all-gather, all-to-all, or
+// reduce-scatter) of a collective fusion computation, or nullptr if there is
+// none.
+const HloInstruction* GetCollectiveHero(const HloComputation* computation);
+
 }  // namespace xla::gpu
 #endif  // XLA_BACKENDS_GPU_CODEGEN_TRITON_COLLECTIVE_EMITTER_H_

From e50b01d31cc4a970e93df949d2632a20b14479c4 Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Fri, 16 Oct 2026 12:51:56 -0700
Subject: [PATCH 0972/1000] In `spmd_partitioner.cc`, `HandleScatter` pads
 `indices` with `-1` before dispatching to scatter partitioning helpers only
 when `indices` is already unevenly sharded. When `indices` is initially
 `{replicated}` and `updates` is unevenly sharded on an index passthrough or
 trivial sliced operand dimension,
 `PartitionScatterIndexPassthroughDimensions` and
 `PartitionScatterTrivialSlicedOperandDimensions` called
 `indices.Reshard(...)` without `pad_value = -1`. This defaulted to `0` and
 scattered into index `0` on the padded shard.

Pass `-1` as `pad_value` when resharding `indices` and forward `pad_value` through `PerGroupSliceFromReplicated`.

PiperOrigin-RevId: 1001040218
---
 .../service/spmd/gather_scatter_handler.cc    |  8 +-
 .../xla/xla/service/spmd/spmd_partitioner.cc  |  2 +-
 .../xla/service/spmd/spmd_partitioner_test.cc | 81 +++++++++++++++++++
 .../xla/service/spmd/spmd_partitioner_util.cc |  5 +-
 .../xla/service/spmd/spmd_partitioner_util.h  |  3 +-
 5 files changed, 93 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/gather_scatter_handler.cc b/third_party/xla/xla/service/spmd/gather_scatter_handler.cc
index 10089972997..385e83151f6 100644
--- a/third_party/xla/xla/service/spmd/gather_scatter_handler.cc
+++ b/third_party/xla/xla/service/spmd/gather_scatter_handler.cc
@@ -1708,7 +1708,9 @@ absl::StatusOr<HloInstruction*> PartitionScatterIndexPassthroughDimensions(
               updates[0].sharding(), index_passthrough_dims.output_dims,
               index_passthrough_dims.indices_dims, indices.num_dimensions()),
           &indices_sharding)) {
-    indices = indices.Reshard(indices_sharding);
+    indices = indices.Reshard(indices_sharding,
+                              /*pad_value=*/LiteralUtil::CreateR0(
+                                  indices.base_shape().element_type(), -1));
   }
   const HloSharding original_indices_sharding = indices.sharding();
 
@@ -1861,7 +1863,9 @@ absl::StatusOr<HloInstruction*> PartitionScatterTrivialSlicedOperandDimensions(
     const GroupedSharding& output_grouped = operand_grouped;
     // Reshard indices to its intended sharding before adjusting.
     indices =
-        indices.Reshard(hlo_sharding_util::UngroupSharding(indices_grouped));
+        indices.Reshard(hlo_sharding_util::UngroupSharding(indices_grouped),
+                        /*pad_value=*/LiteralUtil::CreateR0(
+                            indices.base_shape().element_type(), -1));
     HloInstruction* indices_min;
     std::tie(indices_min, std::ignore) =
         IndexBoundsForGatherScatterOperandPartitionedOnTrivialSliceDims(
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner.cc b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
index 8018a29b374..03c7ea7cdbc 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner.cc
@@ -708,7 +708,7 @@ PartitionedHlo PartitionedHlo::ReshardNoCache(
         hlo_sharding_util::GroupShardingOnDims(target_v2, group_dims);
     auto partially_sharded = PerGroupSliceFromReplicated(
         hlo_, state_.partition_id, target_grouped.device_groups, group_dims,
-        target_grouped.group_dim_sizes, state_.b);
+        target_grouped.group_dim_sizes, state_.b, std::move(pad_value));
     partially_sharded->set_sharding(target);
     return PartitionedHlo(partially_sharded, base_shape(), state_);
   }
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
index 64b84218163..2bdcfc9fe3b 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner_test.cc
@@ -16969,6 +16969,87 @@ ENTRY entry {
               op::Copy(op::CollectivePermute(_)));
 }
 
+TEST_P(SpmdPartitioningTest,
+       ScatterIndexPassthroughUnevenUpdatesPadsIndicesWithMinusOne) {
+  constexpr absl::string_view hlo_string = R"(
+HloModule module
+
+add {
+  lhs = f32[] parameter(0)
+  rhs = f32[] parameter(1)
+  ROOT add = f32[] add(lhs, rhs)
+}
+
+ENTRY entry {
+  %operand = f32[4,2] parameter(0), sharding={replicated}
+  %indices = s32[3,1] parameter(1), sharding={replicated}
+  %updates = f32[3,2] parameter(2), sharding={devices=[2,1]0,1}
+  ROOT %scatter = f32[4,2] scatter(%operand, %indices, %updates),
+      update_window_dims={1},
+      inserted_window_dims={0},
+      scatter_dims_to_operand_dims={0},
+      index_vector_dim=1,
+      to_apply=add,
+      sharding={replicated}
+})";
+
+  ASSERT_OK_AND_ASSIGN(auto module,
+                       PartitionComputation(hlo_string, /*num_devices=*/2));
+  EXPECT_THAT(module->entry_computation()->root_instruction(),
+              op::AllReduce(op::Scatter(
+                  _,
+                  op::DynamicSlice(
+                      op::Pad(op::Parameter(1),
+                              op::Constant(LiteralUtil::CreateR0<int32_t>(-1))),
+                      _, _),
+                  _)));
+}
+
+TEST_P(SpmdPartitioningTest,
+       ScatterTrivialSlicedOperandUnevenUpdatesPadsIndicesWithMinusOne) {
+  constexpr absl::string_view hlo_string = R"(
+HloModule module
+
+add {
+  lhs = f32[] parameter(0)
+  rhs = f32[] parameter(1)
+  ROOT add = f32[] add(lhs, rhs)
+}
+
+ENTRY entry {
+  %operand = f32[4,2] parameter(0),
+      sharding={devices=[2,1,2]0,1,2,3 last_tile_dim_replicate}
+  %indices = s32[3,1] parameter(1), sharding={replicated}
+  %updates = f32[3,2] parameter(2),
+      sharding={devices=[2,1,2]0,2,1,3 last_tile_dim_replicate}
+  ROOT %scatter = f32[4,2] scatter(%operand, %indices, %updates),
+      update_window_dims={1},
+      inserted_window_dims={0},
+      scatter_dims_to_operand_dims={0},
+      index_vector_dim=1,
+      to_apply=add,
+      sharding={devices=[2,1,2]0,1,2,3 last_tile_dim_replicate}
+})";
+
+  SpmdPartitionerOptions options;
+  options.preferred_scatter_partition_methods = {
+      GatherScatterPartitioningMethod::kTrivialSlicedOperand};
+  ASSERT_OK_AND_ASSIGN(
+      auto module,
+      PartitionComputation(hlo_string, /*num_devices=*/4, options));
+  EXPECT_THAT(
+      module->entry_computation()->root_instruction(),
+      op::AllReduce(op::Scatter(
+          _,
+          op::Subtract(
+              op::DynamicSlice(
+                  op::Pad(op::Parameter(1),
+                          op::Constant(LiteralUtil::CreateR0<int32_t>(-1))),
+                  _, _),
+              _),
+          _)));
+}
+
 TEST_P(SpmdPartitioningTest, KeepPartitionedNonSlicedDimension) {
   const char* const hlo_string = R"(
 HloModule module
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner_util.cc b/third_party/xla/xla/service/spmd/spmd_partitioner_util.cc
index de57f8dc4b6..b24104039e3 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner_util.cc
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner_util.cc
@@ -2538,7 +2538,8 @@ HloInstruction* PerGroupSliceFromReplicated(
     HloInstruction* replicated, HloInstruction* partition_id,
     const hlo_sharding_util::DeviceGroupTileAssignment& device_groups,
     absl::Span<const int64_t> group_dims,
-    absl::Span<const int64_t> group_dim_sizes, SpmdBuilder* b) {
+    absl::Span<const int64_t> group_dim_sizes, SpmdBuilder* b,
+    std::optional<Literal> pad_value) {
   std::vector<uint32_t> group_ids(device_groups.num_total_devices());
   for (int64_t g = 0; g < device_groups.num_groups(); ++g) {
     for (int64_t d = 0; d < device_groups.num_devices_per_group(); ++d) {
@@ -2568,7 +2569,7 @@ HloInstruction* PerGroupSliceFromReplicated(
   }();
   auto group_level_sharding = HloSharding::Tile(std::move(group_level_tile));
   auto padded_hlo = PadBaseShapeBeforeUnevenTiledSharding(
-      replicated, group_level_sharding, b);
+      replicated, group_level_sharding, b, std::move(pad_value));
   auto shard_shape =
       MakePartitionedShape(replicated->shape(), group_level_sharding);
   return b->AddInstruction(HloInstruction::CreateDynamicSlice(
diff --git a/third_party/xla/xla/service/spmd/spmd_partitioner_util.h b/third_party/xla/xla/service/spmd/spmd_partitioner_util.h
index d144073d5cc..c1e82bbde42 100644
--- a/third_party/xla/xla/service/spmd/spmd_partitioner_util.h
+++ b/third_party/xla/xla/service/spmd/spmd_partitioner_util.h
@@ -487,7 +487,8 @@ HloInstruction* PerGroupSliceFromReplicated(
     HloInstruction* replicated, HloInstruction* partition_id,
     const hlo_sharding_util::DeviceGroupTileAssignment& device_groups,
     absl::Span<const int64_t> group_dims,
-    absl::Span<const int64_t> group_dim_sizes, SpmdBuilder* b);
+    absl::Span<const int64_t> group_dim_sizes, SpmdBuilder* b,
+    std::optional<Literal> pad_value = std::nullopt);
 
 // Pad the shape from partial replicate shape for `dst_sharding`.
 // If dst_sharding needs more padding and per_shard_size increased in

From 15376e0d28198a4d8e50ba2ac37f70bf89de1f86 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Fri, 16 Oct 2026 13:11:06 -0700
Subject: [PATCH 0973/1000] [XLA:GPU]: Preserve padding on all-gather input
 tile propagation.

There was a bit of an issue with input_offset reading stuff even though it wasn't writing anything.

PiperOrigin-RevId: 1001052098
---
 .../tiling/experimental/tile_propagation.cc   | 13 +++-
 .../experimental/tile_propagation_test.cc     | 71 +++++++++++++++----
 2 files changed, 70 insertions(+), 14 deletions(-)

diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
index 5ce97597980..74c5e7e6ad7 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation.cc
@@ -1688,8 +1688,9 @@ absl::StatusOr<Tiles> PropagateTileToInputForAllGatherOp(
   // with the same index.
   const Shape& input_shape = hlo.operand(output_index)->shape();
   const Shape& output_shape = GetFirstShape(&hlo, output_index);
+  int64_t output_dim_size = output_shape.dimensions(gather_dim);
   int64_t local_size = input_shape.dimensions(gather_dim);
-  int64_t num_replicas = output_shape.dimensions(gather_dim) / local_size;
+  int64_t num_replicas = output_dim_size / local_size;
 
   const DimTile& output_dim_tile = output_tile.dim_tiles()[gather_dim];
 
@@ -1708,7 +1709,15 @@ absl::StatusOr<Tiles> PropagateTileToInputForAllGatherOp(
   }
 
   SymbolicExpr replica_id = output_dim_tile.offset / local_size;
-  SymbolicExpr input_offset = output_dim_tile.offset % local_size;
+  // `offset % local_size` maps an output offset into `[0, local_size)`.
+  // However, if the variadic schedule padded this tile past `output_dim_size`
+  // (e.g., offset = 5 when local_size = 2 and output_dim_size = 4), `5 % 2 = 1`
+  // would wrap it back in-bounds. Adding
+  // `(offset / output_dim_size) * local_size` is 0 for valid tiles and pushes
+  // padded tiles to `>= local_size` so they get masked out.
+  SymbolicExpr input_offset =
+      (output_dim_tile.offset % local_size) +
+      (output_dim_tile.offset / output_dim_size) * local_size;
 
   llvm::SmallVector<DimTile> input_dim_tiles;
   input_dim_tiles.reserve(output_tile.num_dim_tiles());
diff --git a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
index 189cd457380..dede140a180 100644
--- a/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
+++ b/third_party/xla/xla/codegen/tiling/experimental/tile_propagation_test.cc
@@ -268,7 +268,7 @@ TEST_F(TilePropagationTest, CanPropagateToInputsOfAllGatherOp) {
           GetTestTile(*tiling_space, root->shape().dimensions()), 0));
   EXPECT_THAT(tiled_operands, MatchToString(R"(
     0) (tid_0, tid_1)
-      -> offsets [(tid_0 * ts_0) mod 64, tid_1 * ts_1]
+      -> offsets [(tid_0 * ts_0) mod 64 + ((tid_0 * ts_0) / 128) * 64, tid_1 * ts_1]
          sizes [ts_0, ts_1]
          strides [1, 2]
          upper bounds [64, 256]
@@ -309,18 +309,18 @@ TEST_F(TilePropagationTest, CanPropagateToInputsOfVariadicAllGatherOp) {
       root->fused_instructions_computation()->GetInstructionWithName("ag");
   ASSERT_EQ(tiling_space->tiled_roots().size(), 2);
 
-  // Output `k` maps directly to operand `k`.
-  // Since the gather dimension size is 1, the local offset is always 0 and the
-  // replica ID directly equals the output offset. Each output gets distinct
-  // tile IDs (e.g., tid_0 and tid_4) that the scheduler will later assign to
-  // the same peer.
+  // Output `k` comes from operand `k` (`p0` has gather size 1 -> output 4).
+  // For valid tiles (`tid_0 * ts_0 < 4`), the input offset `(tid_0 * ts_0) / 4`
+  // is 0; if the schedule pads a tile past 4, it evaluates to >= 1 so the read
+  // is masked out. Output 0 uses `tid_0` and output 1 uses `tid_4`, which the
+  // scheduler later maps to the same peer.
   ASSERT_OK_AND_ASSIGN(
       Tiles operand_0_tiles,
       PropagateTileToInput(*tiling_space, *ag, tiling_space->tiled_roots()[0],
                            /*output_index=*/0));
   EXPECT_THAT(operand_0_tiles, MatchToString(R"(
     0) (tid_0, tid_1, tid_2, tid_3, tid_4, tid_5)
-      -> offsets [0, tid_1 * ts_1, tid_2 * ts_2, tid_3 * ts_3]
+      -> offsets [(tid_0 * ts_0) / 4, tid_1 * ts_1, tid_2 * ts_2, tid_3 * ts_3]
          sizes [ts_0, ts_1, ts_2, ts_3]
          strides [1, 1, 1, 1]
          upper bounds [1, 2, 8, 16]
@@ -338,7 +338,7 @@ TEST_F(TilePropagationTest, CanPropagateToInputsOfVariadicAllGatherOp) {
                            /*output_index=*/1));
   EXPECT_THAT(operand_1_tiles, MatchToString(R"(
     0) (tid_0, tid_1, tid_2, tid_3, tid_4, tid_5)
-      -> offsets [0, tid_5 * ts_5]
+      -> offsets [(tid_4 * ts_4) / 4, tid_5 * ts_5]
          sizes [ts_4, ts_5]
          strides [1, 1]
          upper bounds [1, 2]
@@ -400,7 +400,7 @@ TEST_F(TilePropagationTest, CanPropagateFromVariadicAllGatherToDotProducers) {
       PropagateTileToInput(*tiling_space, *ag->operand(0), ag_op0[0], 0));
   EXPECT_THAT(dot0_operands, MatchToString(R"(
     0) (tid_0, tid_1, tid_2, tid_3, tid_4, tid_5)
-      -> offsets [(tid_0 * ts_0) mod 2, tid_5 * ts_5]
+      -> offsets [(tid_0 * ts_0) mod 2 + ((tid_0 * ts_0) / 8) * 2, tid_5 * ts_5]
          sizes [ts_0, ts_5]
          strides [1, 1]
          upper bounds [2, 8]
@@ -432,7 +432,7 @@ TEST_F(TilePropagationTest, CanPropagateFromVariadicAllGatherToDotProducers) {
       PropagateTileToInput(*tiling_space, *ag->operand(1), ag_op1[0], 0));
   EXPECT_THAT(dot1_operands, MatchToString(R"(
     0) (tid_0, tid_1, tid_2, tid_3, tid_4, tid_5)
-      -> offsets [(tid_2 * ts_2) mod 4, tid_4 * ts_4]
+      -> offsets [(tid_2 * ts_2) mod 4 + ((tid_2 * ts_2) / 16) * 4, tid_4 * ts_4]
          sizes [ts_2, ts_4]
          strides [1, 1]
          upper bounds [4, 8]
@@ -456,6 +456,53 @@ TEST_F(TilePropagationTest, CanPropagateFromVariadicAllGatherToDotProducers) {
   )"));
 }
 
+TEST_F(TilePropagationTest,
+       AllGatherInputOffsetPreservesSchedulePaddingOutOfBounds) {
+  HloInstruction* root = ParseAndGetRoot(R"(
+    HloModule m
+    f {
+      p0 = f32[2,16] parameter(0)
+      p1 = pred[2] parameter(1)
+      ag = (f32[4,16], pred[4]) all-gather(p0, p1),
+        replica_groups={{0,1}}, dimensions={0}
+      gte0 = f32[4,16] get-tuple-element(ag), index=0
+      gte1 = pred[4] get-tuple-element(ag), index=1
+      ROOT t = (f32[4,16], pred[4]) tuple(gte0, gte1)
+    }
+
+    ENTRY e {
+      p0 = f32[2,16] parameter(0)
+      p1 = pred[2] parameter(1)
+      ROOT fusion = (f32[4,16], pred[4]) fusion(p0, p1), kind=kCustom, calls=f
+    }
+  )");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<TilingSpace> tiling_space,
+      TilingSpace::Create(*HloFusionAdaptor::ForInstruction(root),
+                          &mlir_context_));
+  ASSERT_OK(tiling_space->AssignTileSizes({1, 4, 1}));
+  const HloInstruction* ag =
+      root->fused_instructions_computation()->GetInstructionWithName("ag");
+
+  // Simulate a padded block for output 1 (`pred[4]` from `pred[2]`) by adding
+  // `output_dim_size = 4` to its gather-dim offset (`tid_2 + 4`).
+  Tile padded_output_tile = tiling_space->tiled_roots()[1];
+  SmallVector<DimTile> dim_tiles(padded_output_tile.dim_tiles());
+  dim_tiles[0].offset = dim_tiles[0].offset + 4;
+  ASSERT_OK_AND_ASSIGN(
+      Tiles operand_1_tiles,
+      PropagateTileToInput(*tiling_space, *ag,
+                           padded_output_tile.CloneWithNewDims(dim_tiles),
+                           /*output_index=*/1));
+  ASSERT_EQ(operand_1_tiles.size(), 1);
+  // At `tid_2 = 1` (padded offset = 5), `5 % 2 = 1` alone would wrap back
+  // inside `pred[2]`. With the overflow term, the input offset is
+  // `(5 % 2) + (5 / 4) * 2 = 3 >= 2`, keeping it out of bounds.
+  EXPECT_EQ(operand_1_tiles[0].offsets()[0].Evaluate({0, 0, 1}), 3);
+  EXPECT_GE(operand_1_tiles[0].offsets()[0].Evaluate({0, 0, 1}),
+            operand_1_tiles[0].upper_bounds()[0].Evaluate({0, 0, 1}));
+}
+
 TEST_F(TilePropagationTest, VariadicAllGatherRejectsOutOfRangeOutputIndex) {
   HloInstruction* root = ParseAndGetRoot(R"(
     HloModule m
@@ -1546,7 +1593,7 @@ TEST_F(TilePropagationTest, CanPropagateReplicaIdThroughBroadcast) {
 
   EXPECT_THAT(tiled_ag_operands, MatchToString(R"(
     0) (tid_0, tid_1, tid_2)
-      -> offsets [tid_0 * ts_0, (tid_1 * ts_1) mod 32, tid_2 * ts_2]
+      -> offsets [tid_0 * ts_0, (tid_1 * ts_1) mod 32 + ((tid_1 * ts_1) / 64) * 32, tid_2 * ts_2]
          sizes [ts_0, ts_1, ts_2]
          strides [1, 2, 3]
          upper bounds [10, 32, 5]
@@ -1564,7 +1611,7 @@ TEST_F(TilePropagationTest, CanPropagateReplicaIdThroughBroadcast) {
                                             tiled_ag_operands[0], 0));
   EXPECT_THAT(tiled_broadcast_operands, MatchToString(R"(
     0) (tid_0, tid_1, tid_2)
-      -> offsets [tid_0 * ts_0, (tid_1 * ts_1) mod 32]
+      -> offsets [tid_0 * ts_0, (tid_1 * ts_1) mod 32 + ((tid_1 * ts_1) / 64) * 32]
          sizes [ts_0, ts_1]
          strides [1, 2]
          upper bounds [10, 32]

From d7a966cb268fe34e2faceacd7c928ab99a9b7d89 Mon Sep 17 00:00:00 2001
From: Dmitri Latushko <latushko@google.com>
Date: Fri, 16 Oct 2026 13:12:15 -0700
Subject: [PATCH 0974/1000] Increase shard_count for remote_test and
 batch_resource_base_test

- Increase `shard_count` from 2 to 4 on `//tensorflow/python/eager:remote_test` to reduce per-shard in-process `Server` thread accumulation across test methods.
- Add `shard_count = 4` to `//tensorflow/core/kernels/batching_util:batch_resource_base_test` to distribute its parameterized test suites across 4 shards and prevent sanitizer timeouts.

PiperOrigin-RevId: 1001052968
---
 tensorflow/core/kernels/batching_util/BUILD | 1 +
 tensorflow/python/eager/BUILD               | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/tensorflow/core/kernels/batching_util/BUILD b/tensorflow/core/kernels/batching_util/BUILD
index 1ec12530c8c..2d33c9dd5aa 100644
--- a/tensorflow/core/kernels/batching_util/BUILD
+++ b/tensorflow/core/kernels/batching_util/BUILD
@@ -542,6 +542,7 @@ cc_library(
 tf_cc_test(
     name = "batch_resource_base_test",
     srcs = ["batch_resource_base_test.cc"],
+    shard_count = 4,
     deps = [
         ":batch_resource_base",
         ":batch_scheduler_hdrs",
diff --git a/tensorflow/python/eager/BUILD b/tensorflow/python/eager/BUILD
index 3d1d87b4bcb..d99bf8e9634 100644
--- a/tensorflow/python/eager/BUILD
+++ b/tensorflow/python/eager/BUILD
@@ -1157,7 +1157,7 @@ cuda_py_strict_test(
         "//conditions:default": None,
     }),
     grpc_enabled = True,
-    shard_count = 2,
+    shard_count = 4,
     tags = [
         "no_oss",  # This test launches local server.
         "nofastbuild",  # times out

From f09012d726dc4963cd4d7347f2d79203928fa692 Mon Sep 17 00:00:00 2001
From: Vrishab Madduri <vrishab@google.com>
Date: Fri, 16 Oct 2026 13:28:04 -0700
Subject: [PATCH 0975/1000] No longer needed because host<->device transfers
 can happen directly under GPU Confidential Computing.

Reverts fc058b6345f5d1b61c6da47a3bf06134df69f5f6

PiperOrigin-RevId: 1001063070
---
 third_party/xla/xla/pjrt/gpu/BUILD            |   3 -
 .../xla/pjrt/gpu/se_gpu_pjrt_client_test.cc   | 158 ------------------
 third_party/xla/xla/pjrt/se/BUILD             |   1 -
 .../xla/xla/pjrt/se/local_device_state.cc     |  16 --
 .../xla/xla/pjrt/se/local_device_state.h      |  10 --
 third_party/xla/xla/pjrt/se/se_raw_buffer.cc  | 109 ++----------
 6 files changed, 16 insertions(+), 281 deletions(-)

diff --git a/third_party/xla/xla/pjrt/gpu/BUILD b/third_party/xla/xla/pjrt/gpu/BUILD
index 4a2935e79be..44ffae3ae0e 100644
--- a/third_party/xla/xla/pjrt/gpu/BUILD
+++ b/third_party/xla/xla/pjrt/gpu/BUILD
@@ -451,10 +451,7 @@ xla_test(
         "//xla/service/gpu:gpu_executable_run_options",
         "//xla/service/gpu:gpu_memory_space_assignment",
         "//xla/stream_executor:device_address_allocator",
-        "//xla/stream_executor:event",
-        "//xla/stream_executor:stream",
         "//xla/stream_executor:stream_executor_address_allocator",
-        "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor/integrations:tf_allocator_adapter",
         "//xla/tests:literal_test_util",
         "//xla/tsl/concurrency:async_value",
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
index 557fb946447..50815923322 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
@@ -110,10 +110,7 @@ limitations under the License.
 #endif  // GOOGLE_CUDA
 #include "xla/pjrt/gpu/se_gpu_pjrt_client_test_helper.h"
 #include "xla/stream_executor/device_address_allocator.h"
-#include "xla/stream_executor/event.h"
 #include "xla/stream_executor/integrations/tf_allocator_adapter.h"
-#include "xla/stream_executor/stream.h"
-#include "xla/stream_executor/stream_executor.h"
 #include "xla/stream_executor/stream_executor_address_allocator.h"
 #include "xla/tests/literal_test_util.h"
 #include "xla/tsl/concurrency/async_value_ref.h"
@@ -1793,161 +1790,6 @@ TEST(StreamExecutorGpuClientTest, ChunkedStagingWithNonzeroOffset) {
                                 10 * kTestStagingChunkSize + 45);
 }
 
-// Forces chunked staging (the transfer path used under GPU Confidential
-// Computing) on `client`.
-PjRtStreamExecutorRawClient* ForceChunkedStaging(CommonPjRtClient* client) {
-  auto* raw_client =
-      absl::down_cast<PjRtStreamExecutorRawClient*>(client->raw_client());
-  raw_client->SetStagingChunkingForTesting(/*force=*/true,
-                                           kTestStagingChunkSize);
-  return raw_client;
-}
-
-void ExpectDeviceBytes(PjRtRawBufferInterface& raw_buffer,
-                       const std::vector<uint8_t>& expected) {
-  std::vector<uint8_t> actual(expected.size(), 0);
-  TF_ASSERT_OK(
-      raw_buffer.CopyRawDeviceToHost(actual.data(), 0, actual.size()).Await());
-  EXPECT_EQ(actual, expected);
-}
-
-// A chunked staged H2D transfer into a buffer whose allocation is not yet
-// ready on the compute stream waits on the host, not on the shared
-// host-to-device stream, so transfers into ready buffers do not queue behind
-// compute work.
-TEST(StreamExecutorGpuClientTest, ChunkedStagingH2DDoesNotQueueBehindCompute) {
-  GpuClientOptions options = GetTestGpuClientOptions();
-  options.should_stage_host_to_device_transfers = true;
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtClient> pjrt_client,
-                          GetStreamExecutorGpuClient(options));
-  auto* client = absl::down_cast<CommonPjRtClient*>(pjrt_client.get());
-  PjRtStreamExecutorRawClient* raw_client = ForceChunkedStaging(client);
-  PjRtDevice* device = client->addressable_devices()[0];
-  LocalDeviceState* local_device_state =
-      raw_client->device_state(device->local_device_id());
-  TF_ASSERT_OK_AND_ASSIGN(PjRtMemorySpace * memory_space,
-                          device->default_memory_space());
-  const int64_t size = 3 * kTestStagingChunkSize;
-  auto allocate = [&]() {
-    return client->AllocateRawBuffer(memory_space, size, /*retry_on_oom=*/true,
-                                     /*allocate_after=*/{});
-  };
-  const std::vector<uint8_t> src = MakeStagingTestPattern(size);
-  const std::vector<uint8_t> src2(src.rbegin(), src.rend());
-
-  // `ready`'s allocation is complete before the compute stream stalls.
-  TF_ASSERT_OK_AND_ASSIGN(PjRtRawBufferRef ready, allocate());
-  TF_ASSERT_OK(ready->CopyRawHostToDevice(src.data(), 0, size).Await());
-
-  auto release_compute = std::make_shared<absl::Notification>();
-  absl::Cleanup release = [&] {
-    if (!release_compute->HasBeenNotified()) {
-      release_compute->Notify();
-    }
-  };
-  TF_ASSERT_OK(local_device_state->compute_stream()->DoHostCallback(
-      [release_compute] { release_compute->WaitForNotification(); }));
-  // `pending`'s allocation event is recorded behind the stalled compute work.
-  TF_ASSERT_OK_AND_ASSIGN(PjRtRawBufferRef pending, allocate());
-  Future<> pending_h2d = pending->CopyRawHostToDevice(src.data(), 0, size);
-  Future<> ready_h2d = ready->CopyRawHostToDevice(src2.data(), 0, size);
-  // The transfers run asynchronously and expose no signal for "copies issued",
-  // so give both time to reach the host-to-device stream. Without this, the
-  // event below could be recorded before either transfer enqueues work on that
-  // stream and pass trivially. Their relative order does not matter.
-  absl::SleepFor(absl::Milliseconds(200));
-  // Host functions in a context run serially on one driver thread, so
-  // `ready_h2d`'s completion callback cannot run during the stall. Instead,
-  // poll an event recorded on the host-to-device stream.
-  se::Stream* h2d_stream = local_device_state->host_to_device_stream();
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<se::Event> h2d_done,
-                          h2d_stream->parent()->CreateEvent());
-  TF_ASSERT_OK(h2d_stream->RecordEvent(h2d_done.get()));
-  const absl::Time deadline = absl::Now() + absl::Seconds(10);
-  while (h2d_done->PollForStatus() == se::Event::Status::kPending &&
-         absl::Now() < deadline) {
-    absl::SleepFor(absl::Milliseconds(1));
-  }
-  EXPECT_EQ(h2d_done->PollForStatus(), se::Event::Status::kComplete);
-  EXPECT_FALSE(pending_h2d.IsReady());
-
-  std::move(release).Invoke();
-  TF_ASSERT_OK(pending_h2d.Await());
-  TF_ASSERT_OK(ready_h2d.Await());
-  ExpectDeviceBytes(*pending, src);
-  ExpectDeviceBytes(*ready, src2);
-}
-
-// Many concurrent chunked staged H2D transfers into buffers that may reuse the
-// memory of inputs still being read by queued executions complete without
-// deadlock, and neither the transfers nor the executions see corrupted data.
-TEST(StreamExecutorGpuClientTest, ChunkedStagingConcurrentH2DAfterExecute) {
-  GpuClientOptions options = GetTestGpuClientOptions();
-  options.should_stage_host_to_device_transfers = true;
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtClient> pjrt_client,
-                          GetStreamExecutorGpuClient(options));
-  auto* client = absl::down_cast<CommonPjRtClient*>(pjrt_client.get());
-  ForceChunkedStaging(client);
-  PjRtDevice* device = client->addressable_devices()[0];
-  TF_ASSERT_OK_AND_ASSIGN(PjRtMemorySpace * memory_space,
-                          device->default_memory_space());
-  static constexpr char const* kAddOneProgram = R"(
-HloModule add_one, entry_computation_layout={(f32[4096])->f32[4096]}
-
-ENTRY main {
-  x = f32[4096] parameter(0)
-  one = f32[] constant(1)
-  ones = f32[4096] broadcast(one), dimensions={}
-  ROOT result = f32[4096] add(x, ones)
-}
-)";
-  TF_ASSERT_OK_AND_ASSIGN(auto executable,
-                          CompileExecutable(kAddOneProgram, *client));
-  Shape shape = ShapeUtil::MakeShape(F32, {4096});
-  const int64_t size = ShapeUtil::ByteSizeOf(shape);
-  constexpr int kIterations = 4;
-  constexpr int kTransfers = 8;
-  for (int iter = 0; iter < kIterations; ++iter) {
-    std::vector<float> data(4096, static_cast<float>(iter));
-    TF_ASSERT_OK_AND_ASSIGN(
-        auto input,
-        client->BufferFromHostBuffer(
-            data.data(), shape.element_type(), shape.dimensions(),
-            /*byte_strides=*/std::nullopt,
-            PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall,
-            /*on_done_with_host_buffer=*/nullptr, memory_space,
-            /*device_layout=*/nullptr));
-    TF_ASSERT_OK_AND_ASSIGN(
-        auto result, executable->Execute({{input.get()}}, /*options=*/{}));
-    // Free the input while the execution may still be reading it, so the
-    // buffers below may reuse its memory.
-    input.reset();
-
-    std::vector<std::vector<uint8_t>> srcs;
-    std::vector<PjRtRawBufferRef> buffers;
-    std::vector<Future<>> transfers;
-    for (int i = 0; i < kTransfers; ++i) {
-      srcs.push_back(MakeStagingTestPattern(size));
-      srcs.back()[0] = static_cast<uint8_t>(iter * kTransfers + i);
-      TF_ASSERT_OK_AND_ASSIGN(
-          PjRtRawBufferRef buffer,
-          client->AllocateRawBuffer(memory_space, size, /*retry_on_oom=*/true,
-                                    /*allocate_after=*/{}));
-      transfers.push_back(
-          buffer->CopyRawHostToDevice(srcs.back().data(), 0, size));
-      buffers.push_back(std::move(buffer));
-    }
-    for (int i = 0; i < kTransfers; ++i) {
-      TF_ASSERT_OK(transfers[i].Await());
-      ExpectDeviceBytes(*buffers[i], srcs[i]);
-    }
-    TF_ASSERT_OK_AND_ASSIGN(auto literal, result[0][0]->ToLiteral().Await());
-    EXPECT_TRUE(LiteralTestUtil::Equal(
-        LiteralUtil::CreateR1<float>(std::vector<float>(4096, iter + 1)),
-        *literal));
-  }
-}
-
 TEST(StreamExecutorGpuClientTest, BufferFromHostBufferPinnedMemory) {
   TF_ASSERT_OK_AND_ASSIGN(
       auto client, GetStreamExecutorGpuClient(GetTestGpuClientOptions()));
diff --git a/third_party/xla/xla/pjrt/se/BUILD b/third_party/xla/xla/pjrt/se/BUILD
index 1b16f4ae26e..9356860796a 100644
--- a/third_party/xla/xla/pjrt/se/BUILD
+++ b/third_party/xla/xla/pjrt/se/BUILD
@@ -80,7 +80,6 @@ cc_library(
         "//xla/pjrt:device_event",
         "//xla/pjrt:pjrt_common",
         "//xla/pjrt:semaphore",
-        "//xla/pjrt:thread_pool_async_work_runner",
         "//xla/pjrt:worker_thread",
         "//xla/stream_executor:device_address",
         "//xla/stream_executor:event",
diff --git a/third_party/xla/xla/pjrt/se/local_device_state.cc b/third_party/xla/xla/pjrt/se/local_device_state.cc
index e0685a7995f..13ab5fee8c0 100644
--- a/third_party/xla/xla/pjrt/se/local_device_state.cc
+++ b/third_party/xla/xla/pjrt/se/local_device_state.cc
@@ -23,7 +23,6 @@ limitations under the License.
 #include <utility>
 #include <vector>
 
-#include "absl/base/call_once.h"
 #include "absl/base/casts.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/functional/any_invocable.h"
@@ -37,7 +36,6 @@ limitations under the License.
 #include "xla/client/local_client.h"
 #include "xla/pjrt/async_work_runner.h"
 #include "xla/pjrt/se/buffer_sequencing_event.h"
-#include "xla/pjrt/thread_pool_async_work_runner.h"
 #include "xla/pjrt/worker_thread.h"
 #include "xla/stream_executor/device_address.h"
 #include "xla/stream_executor/platform.h"
@@ -160,9 +158,6 @@ LocalDeviceState::LocalDeviceState(se::StreamExecutor* executor,
 }
 
 LocalDeviceState::~LocalDeviceState() {
-  // Finish staged host-to-device copies before synchronizing the streams they
-  // use.
-  staged_host_to_device_runner_.reset();
   absl::Status status = SynchronizeAllActivity();
   if (!status.ok()) {
     LOG(ERROR) << "Error when closing device: " << status;
@@ -245,17 +240,6 @@ absl::Status LocalDeviceState::SynchronizeAllActivity() {
   return status;
 }
 
-AsyncWorkRunner* LocalDeviceState::staged_host_to_device_runner() {
-  absl::call_once(staged_host_to_device_runner_once_, [this] {
-    tsl::ThreadOptions thread_options;
-    thread_options.numa_node = executor_->numa_node();
-    staged_host_to_device_runner_ =
-        MakeThreadPoolAsyncWorkRunner(tsl::Env::Default(), "py_xla_staged_h2d",
-                                      /*num_threads=*/1, thread_options);
-  });
-  return staged_host_to_device_runner_.get();
-}
-
 absl::Status LocalDeviceState::ThenMemcpyDeviceToDevice(
     se::Stream* transfer_stream, se::Stream* dst_stream,
     se::DeviceAddressBase src_buffer, se::DeviceAddressBase dst_buffer) {
diff --git a/third_party/xla/xla/pjrt/se/local_device_state.h b/third_party/xla/xla/pjrt/se/local_device_state.h
index 407eb593e3b..1f8abb435e7 100644
--- a/third_party/xla/xla/pjrt/se/local_device_state.h
+++ b/third_party/xla/xla/pjrt/se/local_device_state.h
@@ -24,7 +24,6 @@ limitations under the License.
 #include <stack>
 #include <vector>
 
-#include "absl/base/call_once.h"
 #include "absl/base/thread_annotations.h"
 #include "absl/container/flat_hash_map.h"
 #include "absl/functional/any_invocable.h"
@@ -214,12 +213,6 @@ class LocalDeviceState {
         nullptr, "ThenRelease");
   }
 
-  // Returns a single-threaded runner on which chunked staged host-to-device
-  // transfers (used under GPU Confidential Computing) issue their copies, so
-  // at most one staged transfer at a time issues copies to this device's
-  // `host_to_device_stream()`. Created on first use.
-  AsyncWorkRunner* staged_host_to_device_runner();
-
   std::optional<Semaphore>& compute_semaphore() { return compute_semaphore_; }
 
 
@@ -283,9 +276,6 @@ class LocalDeviceState {
   std::stack<std::unique_ptr<se::Stream>> usage_stream_pool_
       ABSL_GUARDED_BY(stream_pool_mu_);
 
-  absl::once_flag staged_host_to_device_runner_once_;
-  std::unique_ptr<AsyncWorkRunner> staged_host_to_device_runner_;
-
   // Callback map pairs callback stream with a device stream and is used for
   // running short host-side callbacks after device side events, without
   // preventing the device-side stream from doing useful work.
diff --git a/third_party/xla/xla/pjrt/se/se_raw_buffer.cc b/third_party/xla/xla/pjrt/se/se_raw_buffer.cc
index 757aae63525..23401ed95b3 100644
--- a/third_party/xla/xla/pjrt/se/se_raw_buffer.cc
+++ b/third_party/xla/xla/pjrt/se/se_raw_buffer.cc
@@ -19,7 +19,6 @@ limitations under the License.
 #include <cstdint>
 #include <cstring>
 #include <memory>
-#include <optional>
 #include <utility>
 #include <vector>
 
@@ -81,12 +80,10 @@ namespace {
 constexpr int64_t kMaxStagingChunksInFlight = 4;
 
 // A pinned host staging buffer for one chunk, and an event recorded after the
-// most recent DMA that used it. The event comes from the device's EventPool and
-// returns to it when the slot is released, so staged transfers do not create or
-// destroy device events on the hot path.
+// most recent DMA that used it.
 struct StagingSlot {
   std::shared_ptr<void> buffer;
-  std::optional<EventPool::Handle> dma_done;
+  std::unique_ptr<se::Event> dma_done;
 };
 
 // Allocates the staging slots for a chunked transfer of `transfer_size` bytes.
@@ -112,9 +109,7 @@ absl::StatusOr<std::vector<StagingSlot>> AllocateStagingSlots(
           "XLA_PJRT_GPU_CC_STAGING_CHUNK_MB).",
           slot_size);
     }
-    ABSL_ASSIGN_OR_RETURN(slot.dma_done,
-                     local_device->event_pool().AllocateEvent(
-                         client->async_work_runner(), stream->parent()));
+    ABSL_ASSIGN_OR_RETURN(slot.dma_done, stream->parent()->CreateEvent());
   }
   return slots;
 }
@@ -130,12 +125,6 @@ absl::StatusOr<std::vector<StagingSlot>> AllocateStagingSlots(
 // event is recorded after already-enqueued stream work, so it completes on GPU
 // progress alone and never waits on work from this runner. Clients that supply
 // a bounded runner would hold one thread per in-flight staged transfer.
-// Chunked staged transfers run only once the device has passed the
-// destination or source allocation event (see `ExecuteWhenAllocationReady`),
-// so the transfer streams never wait on compute work. Chunked H2D transfers to
-// a device run on its single-threaded
-// `LocalDeviceState::staged_host_to_device_runner()`, which only ever waits for
-// its own slots' previous DMAs on `host_to_device_stream()`.
 
 // Copies `transfer_size` bytes from host `src` to device `dst` in
 // `chunk_size`-byte chunks staged through `slots`. Host copies into the
@@ -154,10 +143,7 @@ absl::Status ChunkedStagedHostToDevice(se::Stream* stream, const void* src,
     const int64_t size = std::min(chunk_size, transfer_size - offset);
     StagingSlot& slot = slots[i % num_slots];
     if (i >= num_slots) {
-      // Safe to wait in the driver: the slot event only follows this
-      // transfer's own copies on the host-to-device stream, which no longer
-      // waits on compute work.
-      ABSL_RETURN_IF_ERROR(slot.dma_done->event()->Synchronize());
+      ABSL_RETURN_IF_ERROR(slot.dma_done->Synchronize());
     }
     {
       tsl::profiler::TraceMe trace("H2D Copy To Staging Chunk");
@@ -166,7 +152,7 @@ absl::Status ChunkedStagedHostToDevice(se::Stream* stream, const void* src,
     }
     se::DeviceAddressBase dst_chunk = dst.GetByteSlice(offset, size);
     ABSL_RETURN_IF_ERROR(stream->Memcpy(&dst_chunk, slot.buffer.get(), size));
-    ABSL_RETURN_IF_ERROR(stream->RecordEvent(slot.dma_done->event()));
+    ABSL_RETURN_IF_ERROR(stream->RecordEvent(slot.dma_done.get()));
   }
   return absl::OkStatus();
 }
@@ -188,7 +174,7 @@ absl::Status ChunkedStagedDeviceToHost(se::Stream* stream,
     StagingSlot& slot = slots[i % num_slots];
     ABSL_RETURN_IF_ERROR(stream->Memcpy(slot.buffer.get(),
                                    src.GetByteSlice(offset, size), size));
-    return stream->RecordEvent(slot.dma_done->event());
+    return stream->RecordEvent(slot.dma_done.get());
   };
   for (int64_t i = 0; i < num_slots; ++i) {
     ABSL_RETURN_IF_ERROR(enqueue_dma(i));
@@ -197,7 +183,7 @@ absl::Status ChunkedStagedDeviceToHost(se::Stream* stream,
     const int64_t offset = i * chunk_size;
     const int64_t size = std::min(chunk_size, transfer_size - offset);
     StagingSlot& slot = slots[i % num_slots];
-    ABSL_RETURN_IF_ERROR(slot.dma_done->event()->Synchronize());
+    ABSL_RETURN_IF_ERROR(slot.dma_done->Synchronize());
     {
       tsl::profiler::TraceMe trace("D2H Copy From Staging Chunk");
       std::memcpy(static_cast<char*>(dst) + offset, slot.buffer.get(), size);
@@ -218,40 +204,6 @@ absl::Status DrainStreamOnError(se::Stream* stream, absl::Status status) {
   return status;
 }
 
-// Executes `task` on `runner` once `dependencies` are ready and the device has
-// passed `buffer`'s allocation event, without making a stream or a thread wait
-// for compute work: the allocation event becomes ready from the host callback
-// that `LocalDeviceState::AllocateAndRecordEvent` enqueued when recording it.
-// As in `WaitForAllocation`, an allocation event that failed is not waited
-// for. Sets `device_event` to an error if the allocation event can't be
-// obtained.
-void ExecuteWhenAllocationReady(PjRtStreamExecutorRawClient* client,
-                                tsl::AsyncValueRef<RawSEDeviceMemory> buffer,
-                                BufferSequencingEventRef device_event,
-                                const PjRtDeviceEventRefVector& dependencies,
-                                AsyncWorkRunner* runner,
-                                AsyncWorkRunner::Task task) {
-  ExecuteWhenReady(dependencies, client->async_work_runner(),
-                   [client, buffer = std::move(buffer),
-                    device_event = std::move(device_event), runner,
-                    task = std::move(task)]() mutable {
-                     absl::StatusOr<BufferSequencingEventRef> allocation_event =
-                         buffer->GetDefinitionEvent(client->async_work_runner(),
-                                                    /*nullptr_if_past=*/true);
-                     if (!allocation_event.ok()) {
-                       client->SetEventAsError(device_event,
-                                               allocation_event.status());
-                       return;
-                     }
-                     if (!*allocation_event) {
-                       runner->Execute(std::move(task));
-                       return;
-                     }
-                     runner->ExecuteWhenReady({allocation_event->CopyRCRef()},
-                                              std::move(task));
-                   });
-}
-
 }  // namespace
 
 PjRtStreamExecutorDeviceEventPromise::PjRtStreamExecutorDeviceEventPromise(
@@ -322,11 +274,7 @@ PjRtStreamExecutorRawBuffer::CopyRawHostToDeviceAndReturnEvent(
     // the slots outlive the drain below if AllocateAndRecordEvent fails.
     std::shared_ptr<std::vector<StagingSlot>> staging_slots;
     auto status = [&]() -> absl::Status {
-      // Chunked staged transfers run once their allocation is ready; see
-      // ExecuteWhenAllocationReady below.
-      if (!client->ShouldUseChunkedStaging()) {
-        ABSL_RETURN_IF_ERROR(client->WaitForAllocation(stream, *buf));
-      }
+      ABSL_RETURN_IF_ERROR(client->WaitForAllocation(stream, *buf));
       if (transfer_size > 0) {
         if (client->ShouldStageHostToDeviceTransfers(src, transfer_size)) {
           if (client->GetHostMemoryAllocator() == nullptr) {
@@ -341,16 +289,9 @@ PjRtStreamExecutorRawBuffer::CopyRawHostToDeviceAndReturnEvent(
                                                   transfer_size, chunk_size));
             staging_slots =
                 std::make_shared<std::vector<StagingSlot>>(std::move(slots));
-            ABSL_RETURN_IF_ERROR(ChunkedStagedHostToDevice(
-                stream, src, sub_buffer, transfer_size, chunk_size,
-                absl::MakeSpan(*staging_slots)));
-            // Nothing waits on the slot events after the copies are issued, so
-            // return them to the pool for the next transfer to reuse. The slot
-            // buffers stay alive until the transfer's completion event.
-            for (StagingSlot& slot : *staging_slots) {
-              slot.dma_done.reset();
-            }
-            return absl::OkStatus();
+            return ChunkedStagedHostToDevice(stream, src, sub_buffer,
+                                             transfer_size, chunk_size,
+                                             absl::MakeSpan(*staging_slots));
           }
           HostMemoryAllocator::AllocateOptions alloc_opts;
           alloc_opts.numa_node = stream->parent()->numa_node();
@@ -404,17 +345,8 @@ PjRtStreamExecutorRawBuffer::CopyRawHostToDeviceAndReturnEvent(
   // `run_transfer` may be deferred until `dependencies` are ready, so record
   // the allocation event now.
   client_->MaterializeAllocationEvent(*this);
-  if (client_->ShouldUseChunkedStaging()) {
-    // Issue the copies on the device's staged host-to-device runner, one
-    // transfer at a time. A transfer whose allocation is not ready yet does
-    // not hold up later transfers to the same device.
-    ExecuteWhenAllocationReady(
-        client_, device_buffer_, device_event, dependencies,
-        local_device_->staged_host_to_device_runner(), std::move(run_transfer));
-  } else {
-    ExecuteWhenReady(dependencies, client_->async_work_runner(),
-                     std::move(run_transfer));
-  }
+  ExecuteWhenReady(dependencies, client_->async_work_runner(),
+                   std::move(run_transfer));
 
   return PjRtDeviceEventRef(std::move(device_event));
 }
@@ -444,10 +376,7 @@ PjRtStreamExecutorRawBuffer::CopyRawDeviceToHostAndReturnEvent(
     }
     std::shared_ptr<void> staging_buffer;
     auto status = [&]() -> absl::Status {
-      // See CopyRawHostToDeviceAndReturnEvent.
-      if (!client->ShouldUseChunkedStaging()) {
-        ABSL_RETURN_IF_ERROR(client->WaitForAllocation(stream, *buf));
-      }
+      ABSL_RETURN_IF_ERROR(client->WaitForAllocation(stream, *buf));
       if (transfer_size > 0) {
         if (client->ShouldStageHostToDeviceTransfers(dst, transfer_size)) {
           if (client->GetHostMemoryAllocator() == nullptr) {
@@ -513,14 +442,8 @@ PjRtStreamExecutorRawBuffer::CopyRawDeviceToHostAndReturnEvent(
   // `run_transfer` may be deferred until `dependencies` are ready, so record
   // the allocation event now.
   client_->MaterializeAllocationEvent(*this);
-  if (client_->ShouldUseChunkedStaging()) {
-    ExecuteWhenAllocationReady(client_, device_buffer_, device_event,
-                               dependencies, client_->async_work_runner(),
-                               std::move(run_transfer));
-  } else {
-    ExecuteWhenReady(dependencies, client_->async_work_runner(),
-                     std::move(run_transfer));
-  }
+  ExecuteWhenReady(dependencies, client_->async_work_runner(),
+                   std::move(run_transfer));
 
   return PjRtDeviceEventRef(std::move(device_event));
 }

From 2f4d3e9fb6ccfec7b5ad21137bec40e59c15a8a8 Mon Sep 17 00:00:00 2001
From: Sam Ginzburg <samginzburg@google.com>
Date: Fri, 16 Oct 2026 13:44:57 -0700
Subject: [PATCH 0976/1000] [XLA] Fix kAnd init, U64 dynamic slice, and Pattern
 3 in RecognizeReduceWindow

`RecognizeReduceWindow` has three bugs when folding slices and reductions into `reduce-window` instructions. First, `CreateInitValue` returns `LiteralUtil::One(type)` for `HloOpcode::kAnd`, which clears upper bits on integer bitwise AND reductions instead of using all ones. Second, `RunOnComputation` dereferences `GetFirstInteger()` on constant `dynamic-slice` indices without checking `has_value()`, which aborts on `U64` constants at or above `2^63`. Third, Pattern 3 merges `op(reduce_window, slice)` without checking that the initial value matches the identity or that the window dimensions are compatible.

Compute the `HloOpcode::kAnd` identity constant as the bitwise NOT of `LiteralUtil::Zero(type)` via `HloEvaluator` so both `PRED` and integer types receive all ones. Check `GetFirstInteger().has_value()` before folding constant `dynamic-slice` operands into static slices. In Pattern 3, require trivial other window dimensions and forbid window padding, stride, base dilation, and reversal. Also verify that `rw->operand(1)` is identical to `CreateInitValue` before merging the reduction.

Tests:
- `//third_party/tensorflow/compiler/xla/hlo/transforms/simplifiers:recognize_reduce_window_test`: `BitwiseAndInitAndLargeU64DynamicSliceAndPattern3Checks` covers integer `kAnd` reduction, large `U64` `dynamic-slice` indices, initial values that are not the identity, and incompatible window dimensions.
PiperOrigin-RevId: 1001073620
---
 .../simplifiers/recognize_reduce_window.cc    | 20 +++-
 .../recognize_reduce_window_test.cc           | 99 +++++++++++++++++++
 2 files changed, 115 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/recognize_reduce_window.cc b/third_party/xla/xla/hlo/transforms/simplifiers/recognize_reduce_window.cc
index 23d4bdca8e4..5b481003970 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/recognize_reduce_window.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/recognize_reduce_window.cc
@@ -75,8 +75,9 @@ HloInstruction* CreateInitValue(HloComputation* comp, HloOpcode opcode,
       return comp->AddInstruction(
           HloInstruction::CreateConstant(LiteralUtil::MaxValue(type)));
     case HloOpcode::kAnd:
-      return comp->AddInstruction(
-          HloInstruction::CreateConstant(LiteralUtil::One(type)));
+      return comp->AddInstruction(HloInstruction::CreateConstant(
+          *HloEvaluator{}.EvaluateElementwiseUnaryOp(HloOpcode::kNot,
+                                                     LiteralUtil::Zero(type))));
     case HloOpcode::kOr:
       return comp->AddInstruction(
           HloInstruction::CreateConstant(LiteralUtil::Zero(type)));
@@ -360,7 +361,8 @@ absl::StatusOr<bool> RunOnComputation(HloComputation* computation) {
     if (inst->opcode() == HloOpcode::kDynamicSlice) {
       bool all_static = true;
       for (int i = 1; i < inst->operand_count(); ++i) {
-        if (inst->operand(i)->opcode() != HloOpcode::kConstant) {
+        if (inst->operand(i)->opcode() != HloOpcode::kConstant ||
+            !inst->operand(i)->literal().GetFirstInteger().has_value()) {
           all_static = false;
           break;
         }
@@ -487,7 +489,9 @@ absl::StatusOr<bool> RunOnComputation(HloComputation* computation) {
                 break;
               }
               if (get_start(rw_slice, i) == get_start(other, i) &&
-                  get_limit(rw_slice, i) == get_limit(other, i)) {
+                  get_limit(rw_slice, i) == get_limit(other, i) &&
+                  window_util::IsTrivialWindowDimension(
+                      rw->window().dimensions(i))) {
                 new_starts[i] = get_start(rw_slice, i);
                 new_limits[i] = get_limit(rw_slice, i);
               } else if (dim != -1) {
@@ -520,6 +524,14 @@ absl::StatusOr<bool> RunOnComputation(HloComputation* computation) {
                 current_window_size = size + 1;
               }
             }
+            HloInstruction* init = CreateInitValue(
+                computation, inst->opcode(), inst->shape().element_type());
+            matched &= init && rw->operand(1)->Identical(*init) &&
+                       !window_util::HasPadding(rw->window()) &&
+                       !window_util::HasStride(rw->window()) &&
+                       !window_util::HasBaseDilation(rw->window()) &&
+                       !window_util::HasWindowReversal(rw->window());
+            if (init) ABSL_RETURN_IF_ERROR(computation->RemoveInstruction(init));
             if (matched && dim != -1) {
               if (rw_slice->opcode() == HloOpcode::kSlice) {
                 std::vector<int64_t> strides(rank, 1);
diff --git a/third_party/xla/xla/hlo/transforms/simplifiers/recognize_reduce_window_test.cc b/third_party/xla/xla/hlo/transforms/simplifiers/recognize_reduce_window_test.cc
index dfc48aa56d1..912387cf613 100644
--- a/third_party/xla/xla/hlo/transforms/simplifiers/recognize_reduce_window_test.cc
+++ b/third_party/xla/xla/hlo/transforms/simplifiers/recognize_reduce_window_test.cc
@@ -970,4 +970,103 @@ ENTRY main {
                              /*optimization_level=*/2);
 }
 
+TEST_F(RecognizeReduceWindowTest,
+       BitwiseAndInitAndLargeU64DynamicSliceAndPattern3Checks) {
+  const absl::string_view and_hlo = R"(
+HloModule BitwiseAndReduceWindow
+
+and_reducer {
+  lhs = s32[] parameter(0)
+  rhs = s32[] parameter(1)
+  ROOT res = s32[] and(lhs, rhs)
+}
+
+ENTRY main {
+  src = s32[4] parameter(0)
+  slice_rw = s32[2] slice(src), slice={[0:2]}
+  init = s32[] constant(-1)
+  rw = s32[1] reduce-window(slice_rw, init), window={size=2}, to_apply=and_reducer
+  slice_o = s32[1] slice(src), slice={[2:3]}
+  ROOT res = s32[1] and(rw, slice_o)
+}
+)";
+  CheckRecognizeReduceWindow(and_hlo, R"(
+// CHECK: [[INIT:%.*]] = s32[] constant(-1)
+// CHECK: ROOT {{.*}} = s32[1]{0} reduce-window({{.*}}, [[INIT]]), window={size=3}
+)");
+
+  const absl::string_view u64_ds_hlo = R"(
+HloModule DynamicSliceLargeU64Index
+
+ENTRY main {
+  x = s32[10] parameter(0)
+  idx = u64[] constant(9223372036854775808)
+  ds = s32[6] dynamic-slice(x, idx), dynamic_slice_sizes={6}
+  slice_2 = s32[6] slice(x), slice={[2:8]}
+  ROOT add = s32[6] add(ds, slice_2)
+}
+)";
+  CheckRecognizeReduceWindow(u64_ds_hlo, std::nullopt);
+
+  const absl::string_view non_identity_init_hlo = R"(
+HloModule Pattern3NonIdentityInit
+
+add_reducer {
+  lhs = f32[] parameter(0)
+  rhs = f32[] parameter(1)
+  ROOT add = f32[] add(lhs, rhs)
+}
+
+ENTRY main {
+  src = f32[4] parameter(0)
+  slice_rw = f32[2] slice(src), slice={[0:2]}
+  init = f32[] constant(5)
+  rw = f32[1] reduce-window(slice_rw, init), window={size=2}, to_apply=add_reducer
+  slice_o = f32[1] slice(src), slice={[2:3]}
+  ROOT add = f32[1] add(rw, slice_o)
+}
+)";
+  CheckRecognizeReduceWindow(non_identity_init_hlo, std::nullopt);
+
+  const absl::string_view non_trivial_window_hlo = R"(
+HloModule Pattern3NonTrivialOtherWindowDim
+
+add_reducer {
+  lhs = f32[] parameter(0)
+  rhs = f32[] parameter(1)
+  ROOT add = f32[] add(lhs, rhs)
+}
+
+ENTRY main {
+  src = f32[2,4] parameter(0)
+  slice_rw = f32[2,2] slice(src), slice={[0:2], [0:2]}
+  init = f32[] constant(0)
+  rw = f32[2,1] reduce-window(slice_rw, init), window={size=2x2 pad=1_0x0_0}, to_apply=add_reducer
+  slice_o = f32[2,1] slice(src), slice={[0:2], [2:3]}
+  ROOT add = f32[2,1] add(rw, slice_o)
+}
+)";
+  CheckRecognizeReduceWindow(non_trivial_window_hlo, std::nullopt);
+
+  const absl::string_view padded_reduced_dim_hlo = R"(
+HloModule Pattern3PaddedReducedWindowDim
+
+add_reducer {
+  lhs = f32[] parameter(0)
+  rhs = f32[] parameter(1)
+  ROOT add = f32[] add(lhs, rhs)
+}
+
+ENTRY main {
+  src = f32[4] parameter(0)
+  slice_rw = f32[2] slice(src), slice={[0:2]}
+  init = f32[] constant(0)
+  rw = f32[1] reduce-window(slice_rw, init), window={size=2 stride=2 pad=1_0}, to_apply=add_reducer
+  slice_o = f32[1] slice(src), slice={[2:3]}
+  ROOT add = f32[1] add(rw, slice_o)
+}
+)";
+  CheckRecognizeReduceWindow(padded_reduced_dim_hlo, std::nullopt);
+}
+
 }  // namespace xla

From 6cec3fb88157b36a787059f5937459ef6d3a63af Mon Sep 17 00:00:00 2001
From: Tzu-Wei Sung <twsung@google.com>
Date: Fri, 16 Oct 2026 13:53:53 -0700
Subject: [PATCH 0977/1000] [Pallas:benchmark] Force large second minor only if
 applicable.

PiperOrigin-RevId: 1001078521
---
 .../benchmarks/pallas_microbenchmarks/memory_utils.py  | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/memory_utils.py b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/memory_utils.py
index dbba163b515..ca1bdc2c130 100644
--- a/third_party/xla/xla/benchmarks/pallas_microbenchmarks/memory_utils.py
+++ b/third_party/xla/xla/benchmarks/pallas_microbenchmarks/memory_utils.py
@@ -64,17 +64,21 @@ def copy_to_vmem(arr: jax.Array) -> jax.Array:
 
 
 def with_large_2nd_minor_layout(arr: jax.Array) -> jax.Array:
-  """Utility function to apply a layout with a large 2nd minor."""
+  """Utility function to apply a layout with a large 2nd minor if applicable."""
   dtype = arr.dtype
   bits = jax.dtypes.itemsize_bits(dtype)
   elems_per_word = 32 // bits
   lanes = pltpu.get_tpu_info().num_lanes
   sublanes = pltpu.get_tpu_info().num_sublanes
+  large_second_minor = elems_per_word * sublanes
+  if arr.ndim < 2 or arr.shape[-2] < large_second_minor:
+    # No-op if the array is less than 2D or the second minor is too small to
+    # adopt the large 2nd minor layout.
+    return arr
   return layout.with_layout_constraint(
       arr,
       layout.Layout(
           major_to_minor=tuple(range(0, arr.ndim)),
-          tiling=((elems_per_word * sublanes, lanes), (elems_per_word, 1)),
+          tiling=((large_second_minor, lanes), (elems_per_word, 1)),
       ),
   )
-

From 195fe4137184511e59a8555fd016168355d58cf3 Mon Sep 17 00:00:00 2001
From: Joshua Lang <joshualang@google.com>
Date: Fri, 16 Oct 2026 13:56:32 -0700
Subject: [PATCH 0978/1000] Add CudaGraphAnnotationRegistry and turn
 ScopedCommandBufferAnnotation into a joined annotation path

PiperOrigin-RevId: 1001080558
---
 .../xla/xla/stream_executor/cuda/BUILD        |  35 +
 .../cuda/cuda_graph_annotation_registry.cc    | 332 +++++++++
 .../cuda/cuda_graph_annotation_registry.h     | 270 +++++++
 .../cuda_graph_annotation_registry_test.cc    | 701 ++++++++++++++++++
 third_party/xla/xla/stream_executor/gpu/BUILD |   1 +
 .../gpu/scoped_command_buffer_annotation.cc   |  72 +-
 .../gpu/scoped_command_buffer_annotation.h    |  88 ++-
 .../scoped_command_buffer_annotation_test.cc  | 173 ++++-
 8 files changed, 1605 insertions(+), 67 deletions(-)
 create mode 100644 third_party/xla/xla/stream_executor/cuda/cuda_graph_annotation_registry.cc
 create mode 100644 third_party/xla/xla/stream_executor/cuda/cuda_graph_annotation_registry.h
 create mode 100644 third_party/xla/xla/stream_executor/cuda/cuda_graph_annotation_registry_test.cc

diff --git a/third_party/xla/xla/stream_executor/cuda/BUILD b/third_party/xla/xla/stream_executor/cuda/BUILD
index 9131900c9e7..4c87437919f 100644
--- a/third_party/xla/xla/stream_executor/cuda/BUILD
+++ b/third_party/xla/xla/stream_executor/cuda/BUILD
@@ -3478,3 +3478,38 @@ xla_test(
         "@local_config_cuda//cuda:cuda_headers",
     ],
 )
+
+cc_library(
+    name = "cuda_graph_annotation_registry",
+    srcs = ["cuda_graph_annotation_registry.cc"],
+    hdrs = ["cuda_graph_annotation_registry.h"],
+    deps = [
+        "//xla/stream_executor:semantic_version",
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/base:no_destructor",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/container:inlined_vector",
+        "@com_google_absl//absl/functional:function_ref",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/synchronization",
+    ],
+)
+
+xla_cc_test(
+    name = "cuda_graph_annotation_registry_test",
+    srcs = ["cuda_graph_annotation_registry_test.cc"],
+    deps = [
+        ":cuda_graph_annotation_registry",
+        "//xla/stream_executor:semantic_version",
+        "//xla/tsl/platform:env",
+        "@com_google_absl//absl/base:log_severity",
+        "@com_google_absl//absl/log:scoped_mock_log",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/time",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_graph_annotation_registry.cc b/third_party/xla/xla/stream_executor/cuda/cuda_graph_annotation_registry.cc
new file mode 100644
index 00000000000..5f6d2ce634c
--- /dev/null
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_graph_annotation_registry.cc
@@ -0,0 +1,332 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/stream_executor/cuda/cuda_graph_annotation_registry.h"
+
+#include <atomic>
+#include <cstddef>
+#include <cstdint>
+#include <initializer_list>
+#include <string>
+#include <utility>
+
+#include "absl/base/no_destructor.h"
+#include "absl/functional/function_ref.h"
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
+#include "xla/stream_executor/semantic_version.h"
+
+namespace stream_executor::gpu {
+
+CudaGraphAnnotationRegistry::HarvestHold::~HarvestHold() {
+  if (registry_ != nullptr) {
+    registry_->ReleaseHarvestHold();
+  }
+}
+
+CudaGraphAnnotationRegistry& CudaGraphAnnotationRegistry::Default() {
+  static absl::NoDestructor<CudaGraphAnnotationRegistry> instance;
+  return *instance;
+}
+
+CudaGraphAnnotationRegistry::CudaGraphAnnotationRegistry(
+    size_t max_parked_bytes)
+    : max_parked_bytes_(max_parked_bytes) {}
+
+CudaGraphAnnotationRegistry::TrackResult
+CudaGraphAnnotationRegistry::TrackGraph(uint32_t graph_id, TrackMode mode,
+                                        KeySpace space) {
+  const GraphKey graph_key = MakeGraphKey(space, graph_id);
+  // Declared before the lock, so destroyed after it is released.
+  Garbage garbage;
+  absl::MutexLock lock(mu_);
+  auto [it, inserted] = live_.try_emplace(graph_key);
+  if (!inserted) {
+    if (mode == TrackMode::kGetOrCreate) {
+      return {it->second.token, /*newly_tracked=*/false};
+    }
+    // The caller just created this graph, so the entry belongs to a destroyed
+    // graph with the same ID whose cleanup hasn't run yet. Once it runs, its
+    // token no longer matches.
+    RetireLocked(it, garbage);
+    // `RetireLocked` erased `it`, so insert a fresh entry for the new graph.
+    it = live_.try_emplace(graph_key).first;
+  }
+  const Token token = next_token_++;
+  it->second.token = token;
+  UpdateNumGraphsLocked();
+  return {token, /*newly_tracked=*/true};
+}
+
+bool CudaGraphAnnotationRegistry::IsGraphTracked(uint32_t graph_id,
+                                                 KeySpace space) const {
+  const GraphKey graph_key = MakeGraphKey(space, graph_id);
+  absl::ReaderMutexLock lock(mu_);
+  return live_.contains(graph_key);
+}
+
+bool CudaGraphAnnotationRegistry::RegisterNodeAnnotation(
+    uint32_t graph_id, Token token, uint64_t node_id,
+    absl::string_view annotation, KeySpace space) {
+  if (annotation.empty()) {
+    return false;
+  }
+  const GraphKey graph_key = MakeGraphKey(space, graph_id);
+  absl::MutexLock lock(mu_);
+  GraphEntry* entry = FindLiveEntryLocked(graph_key, token);
+  if (entry == nullptr) {
+    return false;
+  }
+  if (!entry->nodes.try_emplace(node_id, annotation).second) {
+    return false;
+  }
+  AddLiveNodeLocked(*entry, annotation.size());
+  return true;
+}
+
+bool CudaGraphAnnotationRegistry::CopyNodeAnnotation(
+    uint32_t graph_id, Token token, uint64_t node_id, uint32_t source_graph_id,
+    uint64_t source_node_id, absl::string_view fallback, KeySpace space) {
+  const GraphKey graph_key = MakeGraphKey(space, graph_id);
+  const GraphKey source_graph_key =
+      MakeGraphKey(KeySpace::kTemplateGraph, source_graph_id);
+  absl::MutexLock lock(mu_);
+  GraphEntry* entry = FindLiveEntryLocked(graph_key, token);
+  if (entry == nullptr || entry->nodes.contains(node_id)) {
+    return false;
+  }
+  const std::string* source = FindNodeLocked(source_graph_key, source_node_id);
+  // Copy before inserting: the source may be in `entry->nodes` itself, which
+  // the insertion may rehash.
+  std::string annotation = source != nullptr ? *source : std::string(fallback);
+  if (annotation.empty()) {
+    return false;
+  }
+  const size_t size = annotation.size();
+  entry->nodes.emplace(node_id, std::move(annotation));
+  AddLiveNodeLocked(*entry, size);
+  return true;
+}
+
+bool CudaGraphAnnotationRegistry::VisitNodeAnnotation(
+    uint32_t graph_id, uint64_t node_id,
+    absl::FunctionRef<void(absl::string_view)> visitor, KeySpace space) const {
+  const GraphKey graph_key = MakeGraphKey(space, graph_id);
+  absl::ReaderMutexLock lock(mu_);
+  const std::string* annotation = FindNodeLocked(graph_key, node_id);
+  if (annotation == nullptr) {
+    return false;
+  }
+  visitor(*annotation);
+  return true;
+}
+
+void CudaGraphAnnotationRegistry::UnregisterGraph(uint32_t graph_id,
+                                                  Token token, KeySpace space) {
+  const GraphKey graph_key = MakeGraphKey(space, graph_id);
+  // Declared before the lock, so destroyed after it is released: freeing a
+  // large graph's strings must not lengthen the driver thread's critical
+  // section.
+  Garbage garbage;
+  absl::MutexLock lock(mu_);
+  auto it = live_.find(graph_key);
+  if (it == live_.end() || it->second.token != token) {
+    // `TrackGraph(kNewGraph)` or `Clear()` already removed this graph's entry,
+    // and the ID may now belong to a newer graph.
+    return;
+  }
+  RetireLocked(it, garbage);
+  UpdateNumGraphsLocked();
+}
+
+CudaGraphAnnotationRegistry::HarvestHold
+CudaGraphAnnotationRegistry::AcquireHarvestHold() {
+  absl::MutexLock lock(mu_);
+  if (active_holds_++ == 0) {
+    park_dropped_in_session_ = 0;
+  }
+  return HarvestHold(this);
+}
+
+void CudaGraphAnnotationRegistry::ReleaseHarvestHold() {
+  // Declared before the lock, so destroyed after it is released.
+  GraphMap purged;
+  uint64_t dropped = 0;
+  {
+    absl::MutexLock lock(mu_);
+    DCHECK_GT(active_holds_, 0);
+    if (--active_holds_ > 0) {
+      return;
+    }
+    purged.swap(parked_);
+    parked_bytes_ = 0;
+    dropped = std::exchange(park_dropped_in_session_, 0);
+    UpdateNumGraphsLocked();
+  }
+  // Warn here, on the thread that ends the harvest: `UnregisterGraph`, which
+  // drops the entries, runs on CUDA driver threads and must not log.
+  if (dropped > 0) {
+    LOG(WARNING) << "Command buffer telemetry dropped the annotations of "
+                 << dropped
+                 << " graph(s) destroyed during the profiling session (parked "
+                    "annotations are capped at "
+                 << max_parked_bytes_
+                 << " bytes). Their CUPTI events keep the command buffer "
+                    "launch annotation.";
+  }
+}
+
+bool CudaGraphAnnotationRegistry::empty() const {
+  return num_graphs_.load(std::memory_order_acquire) == 0;
+}
+
+CudaGraphAnnotationRegistry::TelemetryStrategy
+CudaGraphAnnotationRegistry::strategy() const {
+  return strategy_.load(std::memory_order_acquire);
+}
+
+void CudaGraphAnnotationRegistry::SetStrategy(TelemetryStrategy strategy) {
+  strategy_.store(strategy, std::memory_order_release);
+}
+
+CudaGraphAnnotationRegistry::Stats CudaGraphAnnotationRegistry::GetStats()
+    const {
+  absl::ReaderMutexLock lock(mu_);
+  Stats stats;
+  stats.live_graphs = live_.size();
+  stats.parked_graphs = parked_.size();
+  stats.parked_bytes = parked_bytes_;
+  stats.active_holds = active_holds_;
+  // Order-independent integer sums, so iteration order doesn't matter.
+  // NOLINTNEXTLINE(*-custom-deterministic-iteration-order)
+  for (const auto& [id, entry] : live_) {
+    stats.live_nodes += entry.nodes.size();
+    stats.live_bytes += entry.bytes;
+  }
+  // Order-independent integer sums, so iteration order doesn't matter.
+  // NOLINTNEXTLINE(*-custom-deterministic-iteration-order)
+  for (const auto& [id, entry] : parked_) {
+    stats.parked_nodes += entry.nodes.size();
+  }
+  return stats;
+}
+
+void CudaGraphAnnotationRegistry::ClearForTesting() {
+  // Declared before the lock, so destroyed after it is released.
+  GraphMap live;
+  GraphMap parked;
+  absl::MutexLock lock(mu_);
+  live.swap(live_);
+  parked.swap(parked_);
+  parked_bytes_ = 0;
+  park_dropped_in_session_ = 0;
+  // `next_token_` keeps counting, so that the cleanups of graphs tracked before
+  // the `Clear()` stay stale.
+  UpdateNumGraphsLocked();
+  strategy_.store(TelemetryStrategy::kDisabled, std::memory_order_release);
+}
+
+CudaGraphAnnotationRegistry::GraphEntry*
+CudaGraphAnnotationRegistry::FindLiveEntryLocked(GraphKey graph_key,
+                                                 Token token) {
+  auto it = live_.find(graph_key);
+  if (it == live_.end() || it->second.token != token) {
+    return nullptr;
+  }
+  return &it->second;
+}
+
+const std::string* CudaGraphAnnotationRegistry::FindNodeLocked(
+    GraphKey graph_key, uint64_t node_id) const {
+  // Live entries first: a parked entry with the same key belongs to an older
+  // graph.
+  for (const GraphMap* graphs : {&live_, &parked_}) {
+    auto graph_it = graphs->find(graph_key);
+    if (graph_it == graphs->end()) {
+      continue;
+    }
+    auto node_it = graph_it->second.nodes.find(node_id);
+    if (node_it != graph_it->second.nodes.end()) {
+      return &node_it->second;
+    }
+  }
+  return nullptr;
+}
+
+void CudaGraphAnnotationRegistry::AddLiveNodeLocked(GraphEntry& entry,
+                                                    size_t size) {
+  entry.bytes += size;
+}
+
+void CudaGraphAnnotationRegistry::RetireLocked(GraphMap::iterator it,
+                                               Garbage& garbage) {
+  const GraphKey graph_key = it->first;
+  GraphEntry entry = std::move(it->second);
+  live_.erase(it);
+  // Without a hold nobody will resolve the entry's nodes any more, and an
+  // entry without nodes has nothing to resolve.
+  if (active_holds_ == 0 || entry.nodes.empty()) {
+    garbage.push_back(std::move(entry));
+    return;
+  }
+  // Check the cap after accounting for any older parked entry with the same
+  // key that this entry would replace. If the new entry still exceeds the cap,
+  // drop the new entry and keep the older parked entry intact.
+  auto parked_it = parked_.find(graph_key);
+  const size_t replaced_bytes =
+      parked_it != parked_.end() ? parked_it->second.bytes : 0;
+  if (parked_bytes_ - replaced_bytes + entry.bytes > max_parked_bytes_) {
+    ++park_dropped_in_session_;
+    garbage.push_back(std::move(entry));
+    return;
+  }
+  // Read before `entry` is moved from below.
+  const size_t entry_bytes = entry.bytes;
+  if (parked_it != parked_.end()) {
+    // Replace the older parked entry in place; it goes to `garbage`, so its
+    // strings are freed after `mu_` is released.
+    parked_bytes_ -= replaced_bytes;
+    ++park_dropped_in_session_;
+    garbage.push_back(std::move(parked_it->second));
+    parked_it->second = std::move(entry);
+  } else {
+    parked_.try_emplace(graph_key, std::move(entry));
+  }
+  parked_bytes_ += entry_bytes;
+}
+
+void CudaGraphAnnotationRegistry::UpdateNumGraphsLocked() {
+  num_graphs_.store(live_.size() + parked_.size(), std::memory_order_release);
+}
+
+CudaGraphAnnotationRegistry::TelemetryStrategy
+CudaGraphAnnotationRegistry::SelectTelemetryStrategy(
+    SemanticVersion compiled_toolkit, SemanticVersion driver,
+    SemanticVersion runtime) {
+  SemanticVersion effective = driver < runtime ? driver : runtime;
+  if (compiled_toolkit < SemanticVersion{13, 2, 0} ||
+      effective < SemanticVersion{13, 2, 0}) {
+    return TelemetryStrategy::kDisabled;
+  }
+  if (compiled_toolkit >= SemanticVersion{13, 4, 0} &&
+      effective >= SemanticVersion{13, 4, 0}) {
+    return TelemetryStrategy::kCuda134SourceGraph;
+  }
+  return TelemetryStrategy::kPre134ExecGraph;
+}
+
+}  // namespace stream_executor::gpu
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_graph_annotation_registry.h b/third_party/xla/xla/stream_executor/cuda/cuda_graph_annotation_registry.h
new file mode 100644
index 00000000000..0094673bb4d
--- /dev/null
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_graph_annotation_registry.h
@@ -0,0 +1,270 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_STREAM_EXECUTOR_CUDA_CUDA_GRAPH_ANNOTATION_REGISTRY_H_
+#define XLA_STREAM_EXECUTOR_CUDA_CUDA_GRAPH_ANNOTATION_REGISTRY_H_
+
+#include <atomic>
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <utility>
+
+#include "absl/base/thread_annotations.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/container/inlined_vector.h"
+#include "absl/functional/function_ref.h"
+#include "absl/strings/str_format.h"
+#include "absl/strings/string_view.h"
+#include "absl/synchronization/mutex.h"
+#include "xla/stream_executor/semantic_version.h"
+
+namespace stream_executor::gpu {
+
+// Process-wide map from CUDA graph telemetry IDs to the HLO annotation that was
+// active when XLA created each graph node. `CudaCommandBuffer` writes it while
+// command buffer annotation recording is on; `CuptiTraceCollector` reads it
+// when harvesting activity records. Entries are keyed by `(KeySpace, graph_id)`
+// because CUDA numbers template and executable graphs independently.
+//
+// Each entry is removed by a `CUuserObject` destructor that CUDA runs
+// asynchronously on a driver thread. Therefore:
+// - Writes and `UnregisterGraph` present the entry's `Token`, so a late cleanup
+//   can never remove a newer graph's entry if an ID is reused.
+// - While a `HarvestHold` is alive, entries of destroyed graphs are parked
+//   rather than erased: records are resolved only at harvest, but graphs are
+//   destroyed mid-session. Parked bytes are capped at `max_parked_bytes`.
+//
+// `mu_` is a leaf lock: no CUDA calls or logging while it is held, and removed
+// entries are destroyed after it is released.
+class CudaGraphAnnotationRegistry {
+ public:
+  // Identifies one tracked lifetime of a graph ID. Tokens are issued in
+  // increasing order and are never reused, not even after `Clear()`.
+  using Token = uint64_t;
+  static constexpr Token kInvalidToken = 0;
+
+  // Disjoint key spaces for template (`CUgraph`) and pre-13.4 executable
+  // (`CUgraphExec`) graph IDs, which the CUDA driver numbers independently.
+  enum class KeySpace : uint8_t {
+    kTemplateGraph = 0,
+    // TODO(b/483450622): Delete once CUDA < 13.4 support is retired.
+    kPre134ExecGraph = 1,
+  };
+
+  // Which CUPTI record IDs are resolved, and in which KeySpace.
+  enum class TelemetryStrategy : uint8_t {
+    kDisabled = 0,
+    // CUDA 13.2-13.3: `(graphId, graphNodeId)` in `kPre134ExecGraph`.
+    // TODO(b/483450622): Delete once CUDA < 13.4 support is retired.
+    kPre134ExecGraph,
+    // CUDA 13.4+: `(sourceGraphId, sourceGraphNodeId)` in `kTemplateGraph`.
+    // P2P memcpy records lack these fields and keep the launch annotation.
+    kCuda134SourceGraph,
+  };
+
+  // Picks the telemetry strategy from the CUDA toolkit XLA was compiled with
+  // and the driver and runtime it runs on. A strategy depends only on the
+  // compile-time toolkit version, the driver version, and the runtime version.
+  static TelemetryStrategy SelectTelemetryStrategy(
+      SemanticVersion compiled_toolkit, SemanticVersion driver,
+      SemanticVersion runtime);
+
+  enum class TrackMode {
+    // The caller just created the graph, so a live entry with the same ID
+    // belongs to a destroyed graph whose cleanup hasn't run yet. Retire it.
+    kNewGraph,
+    // The graph may be tracked already (for example, a graph that XLA did not
+    // create). Reuse its live entry if it has one.
+    kGetOrCreate,
+  };
+
+  struct TrackResult {
+    Token token = kInvalidToken;
+    // True if this call created the entry. The caller must then attach the
+    // cleanup that calls `UnregisterGraph(graph_id, token, space)`.
+    bool newly_tracked = false;
+  };
+
+  // Byte counts are sums of annotation sizes, without container overhead.
+  // Walks every entry; for diagnostics only.
+  struct Stats {
+    size_t live_graphs = 0;
+    size_t live_nodes = 0;
+    size_t live_bytes = 0;
+    size_t parked_graphs = 0;
+    size_t parked_nodes = 0;
+    size_t parked_bytes = 0;
+    int active_holds = 0;
+
+    template <typename Sink>
+    friend void AbslStringify(Sink& sink, const Stats& stats) {
+      absl::Format(
+          &sink,
+          "live_graphs=%d live_nodes=%d live_bytes=%d parked_graphs=%d "
+          "parked_nodes=%d parked_bytes=%d active_holds=%d",
+          stats.live_graphs, stats.live_nodes, stats.live_bytes,
+          stats.parked_graphs, stats.parked_nodes, stats.parked_bytes,
+          stats.active_holds);
+    }
+  };
+
+  // Keeps entries of graphs destroyed during a profiling session until the
+  // session's records are harvested. Move-only; releases the hold when
+  // destroyed.
+  class HarvestHold {
+   public:
+    HarvestHold(HarvestHold&& other) noexcept
+        : registry_(std::exchange(other.registry_, nullptr)) {}
+    HarvestHold& operator=(HarvestHold&& other) = delete;
+    HarvestHold(const HarvestHold&) = delete;
+    HarvestHold& operator=(const HarvestHold&) = delete;
+    ~HarvestHold();
+
+   private:
+    friend class CudaGraphAnnotationRegistry;
+    explicit HarvestHold(CudaGraphAnnotationRegistry* registry)
+        : registry_(registry) {}
+
+    CudaGraphAnnotationRegistry* registry_ = nullptr;
+  };
+
+  // Default for `max_parked_bytes`: 256 MiB.
+  static constexpr size_t kDefaultMaxParkedBytes = size_t{256} << 20;
+
+  // Returns the process-wide registry. It is never destroyed, so cleanups that
+  // run during process exit can still use it.
+  static CudaGraphAnnotationRegistry& Default();
+
+  explicit CudaGraphAnnotationRegistry(
+      size_t max_parked_bytes = kDefaultMaxParkedBytes);
+
+  CudaGraphAnnotationRegistry(const CudaGraphAnnotationRegistry&) = delete;
+  CudaGraphAnnotationRegistry& operator=(const CudaGraphAnnotationRegistry&) =
+      delete;
+
+  // Starts tracking `graph_id` in `space`, or with `kGetOrCreate` returns its
+  // live entry. `kNewGraph` retires a live entry with the same `(space,
+  // graph_id)` as `UnregisterGraph` would (parking it while a hold is active).
+  TrackResult TrackGraph(uint32_t graph_id, TrackMode mode,
+                         KeySpace space = KeySpace::kTemplateGraph)
+      ABSL_LOCKS_EXCLUDED(mu_);
+
+  // Returns whether `graph_id` has a live entry in `space`. Parked entries
+  // don't count: they belong to destroyed graphs.
+  bool IsGraphTracked(uint32_t graph_id,
+                      KeySpace space = KeySpace::kTemplateGraph) const
+      ABSL_LOCKS_EXCLUDED(mu_);
+
+  // First registration wins. Returns false if empty, untracked, stale token, or
+  // present.
+  bool RegisterNodeAnnotation(uint32_t graph_id, Token token, uint64_t node_id,
+                              absl::string_view annotation,
+                              KeySpace space = KeySpace::kTemplateGraph)
+      ABSL_LOCKS_EXCLUDED(mu_);
+
+  // Like RegisterNodeAnnotation, with a copy of a template node's annotation
+  // (live or parked), or fallback. One lock, so the source can't be purged in
+  // between.
+  bool CopyNodeAnnotation(uint32_t graph_id, Token token, uint64_t node_id,
+                          uint32_t source_graph_id, uint64_t source_node_id,
+                          absl::string_view fallback,
+                          KeySpace space = KeySpace::kTemplateGraph)
+      ABSL_LOCKS_EXCLUDED(mu_);
+
+  // Calls visitor with the annotation (live, then parked) under the reader
+  // lock; the visitor must copy and must not re-enter.
+  bool VisitNodeAnnotation(uint32_t graph_id, uint64_t node_id,
+                           absl::FunctionRef<void(absl::string_view)> visitor,
+                           KeySpace space = KeySpace::kTemplateGraph) const
+      ABSL_LOCKS_EXCLUDED(mu_);
+
+  // Cleanup callback on a CUDA driver thread: no CUDA calls, no logging. Parks
+  // or erases the entry if token matches.
+
+  void UnregisterGraph(uint32_t graph_id, Token token,
+                       KeySpace space = KeySpace::kTemplateGraph)
+      ABSL_LOCKS_EXCLUDED(mu_);
+
+  // Takes a harvest hold. Holds are counted, so with overlapping sessions the
+  // parked entries are purged only when the last hold is released.
+  [[nodiscard]] HarvestHold AcquireHarvestHold() ABSL_LOCKS_EXCLUDED(mu_);
+
+  // Returns whether there are neither live nor parked entries. Takes no lock,
+  // so consumers can skip lookups cheaply while telemetry is off.
+  bool empty() const;
+
+  // Returns or sets the active telemetry lookup strategy. Lock-free.
+  TelemetryStrategy strategy() const;
+  void SetStrategy(TelemetryStrategy strategy);
+
+  Stats GetStats() const ABSL_LOCKS_EXCLUDED(mu_);
+
+  // For tests: drops all live and parked entries and resets `strategy()` to
+  // `kDisabled`, but keeps the active holds.
+  void ClearForTesting() ABSL_LOCKS_EXCLUDED(mu_);
+
+ private:
+  using GraphKey = uint64_t;
+  static constexpr GraphKey MakeGraphKey(KeySpace space, uint32_t graph_id) {
+    return (static_cast<uint64_t>(space) << 32) | graph_id;
+  }
+
+  struct GraphEntry {
+    Token token = kInvalidToken;
+    // Sum of the annotation sizes.
+    size_t bytes = 0;
+    absl::flat_hash_map<uint64_t, std::string> nodes;
+  };
+  using GraphMap = absl::flat_hash_map<GraphKey, GraphEntry>;
+  // Entries removed under `mu_`, to be destroyed once it is released. One call
+  // removes at most two: a retired entry and the parked entry it replaces.
+  using Garbage = absl::InlinedVector<GraphEntry, 2>;
+
+  // Decrements the hold count. The last release purges the parked entries and
+  // warns if the session dropped any.
+  void ReleaseHarvestHold() ABSL_LOCKS_EXCLUDED(mu_);
+
+  // Returns the live entry of `graph_key` if its token is `token`.
+  GraphEntry* FindLiveEntryLocked(GraphKey graph_key, Token token)
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
+  // Returns the annotation of the node, from a live entry, else a parked one.
+  const std::string* FindNodeLocked(GraphKey graph_key, uint64_t node_id) const
+      ABSL_SHARED_LOCKS_REQUIRED(mu_);
+  // Accounts for a node of `size` bytes just added to the live `entry`.
+  void AddLiveNodeLocked(GraphEntry& entry, size_t size)
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
+  // Removes the live entry at `it`, parking it or moving it to `garbage`.
+  void RetireLocked(GraphMap::iterator it, Garbage& garbage)
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
+  void UpdateNumGraphsLocked() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
+
+  const size_t max_parked_bytes_;
+
+  mutable absl::Mutex mu_;
+  GraphMap live_ ABSL_GUARDED_BY(mu_);
+  GraphMap parked_ ABSL_GUARDED_BY(mu_);
+  Token next_token_ ABSL_GUARDED_BY(mu_) = 1;
+  int active_holds_ ABSL_GUARDED_BY(mu_) = 0;
+  size_t parked_bytes_ ABSL_GUARDED_BY(mu_) = 0;
+  uint64_t park_dropped_in_session_ ABSL_GUARDED_BY(mu_) = 0;
+  // `live_.size() + parked_.size()`, readable without the lock by `empty()`.
+  std::atomic<size_t> num_graphs_{0};
+  std::atomic<TelemetryStrategy> strategy_{TelemetryStrategy::kDisabled};
+};
+
+}  // namespace stream_executor::gpu
+
+#endif  // XLA_STREAM_EXECUTOR_CUDA_CUDA_GRAPH_ANNOTATION_REGISTRY_H_
diff --git a/third_party/xla/xla/stream_executor/cuda/cuda_graph_annotation_registry_test.cc b/third_party/xla/xla/stream_executor/cuda/cuda_graph_annotation_registry_test.cc
new file mode 100644
index 00000000000..58006102aa5
--- /dev/null
+++ b/third_party/xla/xla/stream_executor/cuda/cuda_graph_annotation_registry_test.cc
@@ -0,0 +1,701 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/stream_executor/cuda/cuda_graph_annotation_registry.h"
+
+#include <atomic>
+#include <cstdint>
+#include <optional>
+#include <string>
+#include <utility>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/base/log_severity.h"
+#include "absl/log/scoped_mock_log.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/time/clock.h"
+#include "absl/time/time.h"
+#include "xla/stream_executor/semantic_version.h"
+#include "xla/tsl/platform/env.h"
+#include "xla/tsl/platform/threadpool.h"
+
+namespace stream_executor::gpu {
+namespace {
+
+using ::testing::_;
+using ::testing::Eq;
+using ::testing::HasSubstr;
+using ::testing::Optional;
+using Registry = CudaGraphAnnotationRegistry;
+using TrackMode = CudaGraphAnnotationRegistry::TrackMode;
+
+constexpr uint32_t kGraph = 10;
+constexpr uint32_t kOtherGraph = 20;
+constexpr uint64_t kNode = 1001;
+constexpr uint64_t kNode2 = 1002;
+constexpr uint64_t kOtherNode = 2001;
+
+// Returns a copy of the node's annotation, if the registry has one.
+std::optional<std::string> Lookup(
+    const Registry& registry, uint32_t graph_id, uint64_t node_id,
+    Registry::KeySpace space = Registry::KeySpace::kTemplateGraph) {
+  std::optional<std::string> annotation;
+  registry.VisitNodeAnnotation(
+      graph_id, node_id,
+      [&annotation](absl::string_view a) { annotation.emplace(a); }, space);
+  return annotation;
+}
+
+// Tracks a newly created graph and returns its token.
+Registry::Token Track(Registry& registry, uint32_t graph_id) {
+  return registry.TrackGraph(graph_id, TrackMode::kNewGraph).token;
+}
+
+TEST(CudaGraphAnnotationRegistryTest, GetOrCreateReturnsExistingToken) {
+  Registry registry;
+  EXPECT_TRUE(registry.empty());
+  EXPECT_FALSE(registry.IsGraphTracked(kGraph));
+
+  Registry::TrackResult first =
+      registry.TrackGraph(kGraph, TrackMode::kGetOrCreate);
+  EXPECT_NE(first.token, Registry::kInvalidToken);
+  EXPECT_TRUE(first.newly_tracked);
+  EXPECT_TRUE(registry.IsGraphTracked(kGraph));
+  EXPECT_FALSE(registry.empty());
+
+  Registry::TrackResult again =
+      registry.TrackGraph(kGraph, TrackMode::kGetOrCreate);
+  EXPECT_EQ(again.token, first.token);
+  EXPECT_FALSE(again.newly_tracked);
+
+  Registry::TrackResult other =
+      registry.TrackGraph(kOtherGraph, TrackMode::kNewGraph);
+  EXPECT_NE(other.token, Registry::kInvalidToken);
+  EXPECT_NE(other.token, first.token);
+  EXPECT_TRUE(other.newly_tracked);
+  EXPECT_EQ(registry.GetStats().live_graphs, 2);
+}
+
+TEST(CudaGraphAnnotationRegistryTest, InsertIfAbsent) {
+  Registry registry;
+  Registry::Token token = Track(registry, kGraph);
+
+  EXPECT_FALSE(registry.RegisterNodeAnnotation(kGraph, token, kNode, ""));
+  EXPECT_TRUE(registry.RegisterNodeAnnotation(kGraph, token, kNode, "first"));
+  EXPECT_FALSE(registry.RegisterNodeAnnotation(kGraph, token, kNode, "second"));
+  EXPECT_THAT(Lookup(registry, kGraph, kNode), Optional(Eq("first")));
+
+  // Writes need a live entry and its token.
+  EXPECT_FALSE(registry.RegisterNodeAnnotation(kGraph, token + 1, kNode2, "x"));
+  EXPECT_FALSE(registry.RegisterNodeAnnotation(kGraph, Registry::kInvalidToken,
+                                               kNode2, "x"));
+  EXPECT_FALSE(
+      registry.RegisterNodeAnnotation(kOtherGraph, token, kOtherNode, "x"));
+  EXPECT_EQ(Lookup(registry, kGraph, kNode2), std::nullopt);
+  EXPECT_FALSE(registry.IsGraphTracked(kOtherGraph));
+  EXPECT_EQ(registry.GetStats().live_graphs, 1);
+}
+
+TEST(CudaGraphAnnotationRegistryTest, TokenMismatchNoOp) {
+  Registry registry;
+  Registry::Token token = Track(registry, kGraph);
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, token, kNode, "a"));
+
+  registry.UnregisterGraph(kGraph, token + 1);
+  EXPECT_THAT(Lookup(registry, kGraph, kNode), Optional(Eq("a")));
+
+  // An untracked ID matches no token either.
+  registry.UnregisterGraph(kOtherGraph, token);
+  EXPECT_THAT(Lookup(registry, kGraph, kNode), Optional(Eq("a")));
+
+  registry.UnregisterGraph(kGraph, token);
+  EXPECT_EQ(Lookup(registry, kGraph, kNode), std::nullopt);
+  EXPECT_TRUE(registry.empty());
+
+  // The old token can't bring the entry back or unregister it twice.
+  EXPECT_FALSE(registry.RegisterNodeAnnotation(kGraph, token, kNode, "a"));
+  registry.UnregisterGraph(kGraph, token);
+  EXPECT_TRUE(registry.empty());
+}
+
+TEST(CudaGraphAnnotationRegistryTest, NewGraphRetiresStale) {
+  Registry registry;
+  Registry::Token t1 = Track(registry, kGraph);
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, t1, kNode, "old"));
+
+  // A new graph got the ID before the old graph's cleanup ran.
+  Registry::TrackResult second =
+      registry.TrackGraph(kGraph, TrackMode::kNewGraph);
+  Registry::Token t2 = second.token;
+  EXPECT_TRUE(second.newly_tracked);
+  EXPECT_NE(t2, t1);
+  EXPECT_EQ(Lookup(registry, kGraph, kNode), std::nullopt);
+  EXPECT_FALSE(registry.RegisterNodeAnnotation(kGraph, t1, kNode, "old"));
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, t2, kNode, "new"));
+
+  // The old graph's late cleanup must leave the new graph's entry alone.
+  registry.UnregisterGraph(kGraph, t1);
+  EXPECT_TRUE(registry.IsGraphTracked(kGraph));
+  EXPECT_THAT(Lookup(registry, kGraph, kNode), Optional(Eq("new")));
+  EXPECT_EQ(registry.GetStats().live_graphs, 1);
+}
+
+TEST(CudaGraphAnnotationRegistryTest, NewGraphParksStaleEntryWhileHeld) {
+  Registry registry;
+  Registry::HarvestHold hold = registry.AcquireHarvestHold();
+  Registry::Token t1 = Track(registry, kGraph);
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, t1, kNode, "old"));
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, t1, kNode2, "old2"));
+
+  // The retired entry takes the park path, so the old graph's events resolve.
+  Registry::Token t2 = Track(registry, kGraph);
+  EXPECT_EQ(registry.GetStats().parked_graphs, 1);
+  EXPECT_THAT(Lookup(registry, kGraph, kNode), Optional(Eq("old")));
+
+  // Parking the newer entry replaces the older parked one.
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, t2, kNode, "new"));
+  registry.UnregisterGraph(kGraph, t2);
+  Registry::Stats stats = registry.GetStats();
+  EXPECT_EQ(stats.parked_graphs, 1);
+  EXPECT_EQ(stats.parked_bytes, 3);
+  EXPECT_THAT(Lookup(registry, kGraph, kNode), Optional(Eq("new")));
+  EXPECT_EQ(Lookup(registry, kGraph, kNode2), std::nullopt);
+}
+
+TEST(CudaGraphAnnotationRegistryTest, ParkThenPurge) {
+  Registry registry;
+  Registry::Token token = Track(registry, kGraph);
+  ASSERT_TRUE(
+      registry.RegisterNodeAnnotation(kGraph, token, kNode, "annotation"));
+  {
+    Registry::HarvestHold hold = registry.AcquireHarvestHold();
+    registry.UnregisterGraph(kGraph, token);
+    Registry::Stats stats = registry.GetStats();
+    EXPECT_EQ(stats.live_graphs, 0);
+    EXPECT_EQ(stats.parked_graphs, 1);
+    EXPECT_EQ(stats.parked_bytes, 10);
+    EXPECT_EQ(stats.active_holds, 1);
+    EXPECT_FALSE(registry.IsGraphTracked(kGraph));
+    EXPECT_FALSE(registry.empty());
+    // The destroyed graph's events still resolve until the harvest ends.
+    EXPECT_THAT(Lookup(registry, kGraph, kNode), Optional(Eq("annotation")));
+    // But the graph is dead: its token no longer admits writes.
+    EXPECT_FALSE(registry.RegisterNodeAnnotation(kGraph, token, kNode2, "x"));
+  }
+  Registry::Stats stats = registry.GetStats();
+  EXPECT_EQ(stats.parked_graphs, 0);
+  EXPECT_EQ(stats.parked_bytes, 0);
+  EXPECT_EQ(stats.active_holds, 0);
+  EXPECT_EQ(Lookup(registry, kGraph, kNode), std::nullopt);
+  EXPECT_TRUE(registry.empty());
+}
+
+TEST(CudaGraphAnnotationRegistryTest, OverlappingHolds) {
+  Registry registry;
+  Registry::Token token = Track(registry, kGraph);
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, token, kNode, "a"));
+  std::optional<Registry::HarvestHold> first(registry.AcquireHarvestHold());
+  Registry::HarvestHold second = registry.AcquireHarvestHold();
+  EXPECT_EQ(registry.GetStats().active_holds, 2);
+  registry.UnregisterGraph(kGraph, token);
+
+  // Moving a hold transfers it without releasing it.
+  std::optional<Registry::HarvestHold> moved(std::move(second));
+  EXPECT_EQ(registry.GetStats().active_holds, 2);
+
+  // Releasing one of two holds keeps the parked entry.
+  first.reset();
+  EXPECT_EQ(registry.GetStats().active_holds, 1);
+  EXPECT_THAT(Lookup(registry, kGraph, kNode), Optional(Eq("a")));
+
+  // Releasing the last hold purges the parked entry.
+  moved.reset();
+  Registry::Stats stats = registry.GetStats();
+  EXPECT_EQ(stats.active_holds, 0);
+  EXPECT_EQ(stats.parked_graphs, 0);
+  EXPECT_EQ(Lookup(registry, kGraph, kNode), std::nullopt);
+}
+
+TEST(CudaGraphAnnotationRegistryTest, EmptyEntryNotParked) {
+  Registry registry;
+  Registry::HarvestHold hold = registry.AcquireHarvestHold();
+  Registry::Token token = Track(registry, kGraph);
+  registry.UnregisterGraph(kGraph, token);
+
+  EXPECT_EQ(registry.GetStats().parked_graphs, 0);
+  EXPECT_TRUE(registry.empty());
+}
+
+TEST(CudaGraphAnnotationRegistryTest, ParkCapDrops) {
+  // Room for "12345", but not for "123456" as well.
+  Registry registry(/*max_parked_bytes=*/10);
+  Registry::Token t1 = Track(registry, kGraph);
+  Registry::Token t2 = Track(registry, kOtherGraph);
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, t1, kNode, "12345"));
+  ASSERT_TRUE(
+      registry.RegisterNodeAnnotation(kOtherGraph, t2, kOtherNode, "123456"));
+
+  absl::ScopedMockLog log(absl::MockLogDefault::kIgnoreUnexpected);
+  // Only the last release warns, and only about the drops.
+  EXPECT_CALL(log, Log(absl::LogSeverity::kWarning, _,
+                       HasSubstr("dropped the annotations of 1 graph")))
+      .Times(1);
+  log.StartCapturingLogs();
+  {
+    Registry::HarvestHold hold = registry.AcquireHarvestHold();
+    registry.UnregisterGraph(kGraph, t1);
+    registry.UnregisterGraph(kOtherGraph, t2);
+    Registry::Stats stats = registry.GetStats();
+    EXPECT_EQ(stats.parked_graphs, 1);
+    EXPECT_EQ(stats.parked_bytes, 5);
+    EXPECT_THAT(Lookup(registry, kGraph, kNode), Optional(Eq("12345")));
+    EXPECT_EQ(Lookup(registry, kOtherGraph, kOtherNode), std::nullopt);
+  }
+  EXPECT_TRUE(registry.empty());
+}
+
+TEST(CudaGraphAnnotationRegistryTest,
+     ParkReplacementOverCapKeepsOldParkedEntry) {
+  // Room for "1234" (4 bytes), not for "123456" (6 bytes).
+  Registry registry(/*max_parked_bytes=*/5);
+  Registry::HarvestHold hold = registry.AcquireHarvestHold();
+
+  Registry::Token t1 = Track(registry, kGraph);
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, t1, kNode, "1234"));
+  registry.UnregisterGraph(kGraph, t1);
+  EXPECT_EQ(registry.GetStats().parked_graphs, 1);
+  EXPECT_EQ(registry.GetStats().parked_bytes, 4);
+
+  // A new graph with the same ID has a 6-byte annotation that exceeds the cap
+  // even after subtracting the old 4-byte entry. Retiring it must drop the new
+  // entry and preserve the old parked entry.
+  Registry::Token t2 = Track(registry, kGraph);
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, t2, kNode, "123456"));
+  registry.UnregisterGraph(kGraph, t2);
+
+  Registry::Stats stats = registry.GetStats();
+  EXPECT_EQ(stats.parked_graphs, 1);
+  EXPECT_EQ(stats.parked_bytes, 4);
+  EXPECT_THAT(Lookup(registry, kGraph, kNode), Optional(Eq("1234")));
+}
+
+TEST(CudaGraphAnnotationRegistryTest, ParkReplacementUnderCapReplacesOldEntry) {
+  // Room for both the old 6-byte entry and the new 7-byte entry.
+  Registry registry(/*max_parked_bytes=*/20);
+  absl::ScopedMockLog log(absl::MockLogDefault::kIgnoreUnexpected);
+  // The replaced parked entry counts as one dropped graph.
+  EXPECT_CALL(log, Log(absl::LogSeverity::kWarning, _,
+                       HasSubstr("dropped the annotations of 1 graph")))
+      .Times(1);
+  log.StartCapturingLogs();
+  {
+    Registry::HarvestHold hold = registry.AcquireHarvestHold();
+
+    Registry::Token t1 = Track(registry, kGraph);
+    ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, t1, kNode, "1234"));
+    ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, t1, kNode2, "56"));
+    registry.UnregisterGraph(kGraph, t1);
+    Registry::Stats stats = registry.GetStats();
+    EXPECT_EQ(stats.parked_graphs, 1);
+    EXPECT_EQ(stats.parked_nodes, 2);
+    EXPECT_EQ(stats.parked_bytes, 6);
+
+    // A new graph with the same ID is retired too. Its entry fits under the
+    // cap, so it replaces the older parked entry.
+    Registry::Token t2 = Track(registry, kGraph);
+    ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, t2, kNode, "abcdefg"));
+    registry.UnregisterGraph(kGraph, t2);
+
+    stats = registry.GetStats();
+    EXPECT_EQ(stats.live_graphs, 0);
+    EXPECT_EQ(stats.parked_graphs, 1);
+    EXPECT_EQ(stats.parked_nodes, 1);
+    EXPECT_EQ(stats.parked_bytes, 7);
+    EXPECT_THAT(Lookup(registry, kGraph, kNode), Optional(Eq("abcdefg")));
+    EXPECT_EQ(Lookup(registry, kGraph, kNode2), std::nullopt);
+  }
+  EXPECT_TRUE(registry.empty());
+}
+
+TEST(CudaGraphAnnotationRegistryTest,
+     TemplateAndPre134ExecKeySpacesAreDisjoint) {
+  constexpr uint32_t kSameGraphId = 7;
+  Registry registry;
+
+  // Track template graph 7 and register an annotation.
+  Registry::TrackResult tmpl = registry.TrackGraph(
+      kSameGraphId, TrackMode::kNewGraph, Registry::KeySpace::kTemplateGraph);
+  ASSERT_TRUE(tmpl.newly_tracked);
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(
+      kSameGraphId, tmpl.token, kNode, "template_op",
+      Registry::KeySpace::kTemplateGraph));
+
+  // Track pre-13.4 exec graph 7 with the same integer ID and kNewGraph; it must
+  // NOT retire or overwrite the live template graph 7.
+  Registry::TrackResult exec = registry.TrackGraph(
+      kSameGraphId, TrackMode::kNewGraph, Registry::KeySpace::kPre134ExecGraph);
+  ASSERT_TRUE(exec.newly_tracked);
+  EXPECT_TRUE(registry.IsGraphTracked(kSameGraphId,
+                                      Registry::KeySpace::kTemplateGraph));
+  EXPECT_TRUE(registry.IsGraphTracked(kSameGraphId,
+                                      Registry::KeySpace::kPre134ExecGraph));
+  EXPECT_EQ(registry.GetStats().live_graphs, 2);
+
+  // Copy from template graph 7 into exec graph 7, plus a direct exec node.
+  EXPECT_TRUE(registry.CopyNodeAnnotation(
+      kSameGraphId, exec.token, kOtherNode, kSameGraphId, kNode, "fallback",
+      Registry::KeySpace::kPre134ExecGraph));
+  EXPECT_TRUE(registry.RegisterNodeAnnotation(
+      kSameGraphId, exec.token, kNode, "exec_only_op",
+      Registry::KeySpace::kPre134ExecGraph));
+
+  // Each key space sees only its own entries.
+  EXPECT_THAT(
+      Lookup(registry, kSameGraphId, kNode, Registry::KeySpace::kTemplateGraph),
+      Optional(Eq("template_op")));
+  EXPECT_EQ(Lookup(registry, kSameGraphId, kOtherNode,
+                   Registry::KeySpace::kTemplateGraph),
+            std::nullopt);
+  EXPECT_THAT(Lookup(registry, kSameGraphId, kNode,
+                     Registry::KeySpace::kPre134ExecGraph),
+              Optional(Eq("exec_only_op")));
+  EXPECT_THAT(Lookup(registry, kSameGraphId, kOtherNode,
+                     Registry::KeySpace::kPre134ExecGraph),
+              Optional(Eq("template_op")));
+
+  // Unregistering exec graph 7 leaves template graph 7 live and intact.
+  registry.UnregisterGraph(kSameGraphId, exec.token,
+                           Registry::KeySpace::kPre134ExecGraph);
+  EXPECT_TRUE(registry.IsGraphTracked(kSameGraphId,
+                                      Registry::KeySpace::kTemplateGraph));
+  EXPECT_FALSE(registry.IsGraphTracked(kSameGraphId,
+                                       Registry::KeySpace::kPre134ExecGraph));
+  EXPECT_THAT(
+      Lookup(registry, kSameGraphId, kNode, Registry::KeySpace::kTemplateGraph),
+      Optional(Eq("template_op")));
+  EXPECT_EQ(Lookup(registry, kSameGraphId, kNode,
+                   Registry::KeySpace::kPre134ExecGraph),
+            std::nullopt);
+}
+
+TEST(CudaGraphAnnotationRegistryTest, LiveBeforeParked) {
+  Registry registry;
+  Registry::HarvestHold hold = registry.AcquireHarvestHold();
+  Registry::Token t1 = Track(registry, kGraph);
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, t1, kNode, "parked"));
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, t1, kNode2, "parked2"));
+  registry.UnregisterGraph(kGraph, t1);
+
+  // A newer graph with the same ID. Parked entries don't count as tracked.
+  Registry::TrackResult t2 =
+      registry.TrackGraph(kGraph, TrackMode::kGetOrCreate);
+  EXPECT_TRUE(t2.newly_tracked);
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, t2.token, kNode, "live"));
+  EXPECT_THAT(Lookup(registry, kGraph, kNode), Optional(Eq("live")));
+  EXPECT_THAT(Lookup(registry, kGraph, kNode2), Optional(Eq("parked2")));
+}
+
+TEST(CudaGraphAnnotationRegistryTest, CopyFromLive) {
+  Registry registry;
+  Registry::Token source = Track(registry, kOtherGraph);
+  Registry::Token clone = Track(registry, kGraph);
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kOtherGraph, source, kOtherNode,
+                                              "original"));
+
+  EXPECT_TRUE(registry.CopyNodeAnnotation(kGraph, clone, kNode, kOtherGraph,
+                                          kOtherNode, "fallback"));
+  EXPECT_THAT(Lookup(registry, kGraph, kNode), Optional(Eq("original")));
+}
+
+TEST(CudaGraphAnnotationRegistryTest, CopyFromParked) {
+  Registry registry;
+  Registry::HarvestHold hold = registry.AcquireHarvestHold();
+  Registry::Token source = Track(registry, kOtherGraph);
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kOtherGraph, source, kOtherNode,
+                                              "original"));
+  registry.UnregisterGraph(kOtherGraph, source);
+  Registry::Token clone = Track(registry, kGraph);
+
+  EXPECT_TRUE(registry.CopyNodeAnnotation(kGraph, clone, kNode, kOtherGraph,
+                                          kOtherNode, "fallback"));
+  EXPECT_THAT(Lookup(registry, kGraph, kNode), Optional(Eq("original")));
+}
+
+TEST(CudaGraphAnnotationRegistryTest, CopyFallback) {
+  Registry registry;
+  Registry::Token clone = Track(registry, kGraph);
+
+  EXPECT_FALSE(registry.CopyNodeAnnotation(kGraph, clone, kNode, kOtherGraph,
+                                           kOtherNode, /*fallback=*/""));
+  EXPECT_EQ(Lookup(registry, kGraph, kNode), std::nullopt);
+  EXPECT_TRUE(registry.CopyNodeAnnotation(kGraph, clone, kNode, kOtherGraph,
+                                          kOtherNode, "fallback"));
+  EXPECT_THAT(Lookup(registry, kGraph, kNode), Optional(Eq("fallback")));
+  // Copies need the live entry's token too.
+  EXPECT_FALSE(registry.CopyNodeAnnotation(kGraph, clone + 1, kNode2,
+                                           kOtherGraph, kOtherNode, "x"));
+  EXPECT_EQ(Lookup(registry, kGraph, kNode2), std::nullopt);
+}
+
+TEST(CudaGraphAnnotationRegistryTest, CopyNoOverwrite) {
+  Registry registry;
+  Registry::Token source = Track(registry, kOtherGraph);
+  Registry::Token clone = Track(registry, kGraph);
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kOtherGraph, source, kOtherNode,
+                                              "original"));
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, clone, kNode, "precise"));
+
+  EXPECT_FALSE(registry.CopyNodeAnnotation(kGraph, clone, kNode, kOtherGraph,
+                                           kOtherNode, "fallback"));
+  EXPECT_THAT(Lookup(registry, kGraph, kNode), Optional(Eq("precise")));
+
+  // The source may be a node of the destination graph itself.
+  EXPECT_TRUE(
+      registry.CopyNodeAnnotation(kGraph, clone, kNode2, kGraph, kNode, ""));
+  EXPECT_THAT(Lookup(registry, kGraph, kNode2), Optional(Eq("precise")));
+}
+
+TEST(CudaGraphAnnotationRegistryTest, VisitorNotCalledOnMiss) {
+  Registry registry;
+  Registry::Token token = Track(registry, kGraph);
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, token, kNode, "a"));
+  int calls = 0;
+  auto visitor = [&calls](absl::string_view) { ++calls; };
+
+  EXPECT_FALSE(registry.VisitNodeAnnotation(kGraph, kNode2, visitor));
+  EXPECT_FALSE(registry.VisitNodeAnnotation(kOtherGraph, kNode, visitor));
+  EXPECT_EQ(calls, 0);
+  EXPECT_TRUE(registry.VisitNodeAnnotation(kGraph, kNode, visitor));
+  EXPECT_EQ(calls, 1);
+}
+
+TEST(CudaGraphAnnotationRegistryTest, ClearKeepsHolds) {
+  Registry registry;
+  Registry::HarvestHold hold = registry.AcquireHarvestHold();
+  Registry::Token live = Track(registry, kGraph);
+  Registry::Token parked = Track(registry, kOtherGraph);
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, live, kNode, "a"));
+  ASSERT_TRUE(
+      registry.RegisterNodeAnnotation(kOtherGraph, parked, kOtherNode, "b"));
+  registry.UnregisterGraph(kOtherGraph, parked);
+
+  registry.ClearForTesting();
+  EXPECT_TRUE(registry.empty());
+  Registry::Stats stats = registry.GetStats();
+  EXPECT_EQ(stats.live_graphs, 0);
+  EXPECT_EQ(stats.parked_graphs, 0);
+  EXPECT_EQ(stats.parked_bytes, 0);
+  EXPECT_EQ(stats.active_holds, 1);
+
+  // Tokens keep increasing, so a cleared graph's late cleanup stays stale.
+  Registry::Token token = Track(registry, kGraph);
+  EXPECT_GT(token, parked);
+  registry.UnregisterGraph(kGraph, live);
+  EXPECT_TRUE(registry.IsGraphTracked(kGraph));
+
+  // The hold survived, so unregistered entries are still parked.
+  ASSERT_TRUE(registry.RegisterNodeAnnotation(kGraph, token, kNode, "c"));
+  registry.UnregisterGraph(kGraph, token);
+  EXPECT_EQ(registry.GetStats().parked_graphs, 1);
+}
+
+// Writers track, register, copy and unregister graphs while visitors read and
+// another thread keeps taking overlapping holds. Run it under TSAN.
+TEST(CudaGraphAnnotationRegistryTest, Concurrency) {
+  constexpr int kNumWriters = 8;
+  constexpr int kGraphsPerWriter = 200;
+  constexpr int kNodesPerGraph = 4;
+  constexpr int kNumGraphs = kNumWriters * kGraphsPerWriter;
+  // Each writer owns a disjoint range of graph IDs, starting at 1.
+  auto annotation = [](uint32_t graph_id, int node) {
+    return absl::StrCat("graph", graph_id, "_node", node);
+  };
+  Registry registry;
+  std::atomic<int> writers_running = kNumWriters;
+  std::atomic<int> visits_hit = 0;
+  {
+    tsl::thread::ThreadPool pool(tsl::Env::Default(), "registry_concurrency",
+                                 kNumWriters + 3);
+    for (int w = 0; w < kNumWriters; ++w) {
+      pool.Schedule([&, w] {
+        for (int g = 0; g < kGraphsPerWriter; ++g) {
+          const uint32_t graph_id = w * kGraphsPerWriter + g + 1;
+          Registry::Token token = Track(registry, graph_id);
+          Registry::TrackResult again =
+              registry.TrackGraph(graph_id, TrackMode::kGetOrCreate);
+          EXPECT_EQ(again.token, token);
+          EXPECT_FALSE(again.newly_tracked);
+          for (int n = 0; n + 1 < kNodesPerGraph; ++n) {
+            EXPECT_TRUE(registry.RegisterNodeAnnotation(
+                graph_id, token, static_cast<uint64_t>(n),
+                annotation(graph_id, n)));
+          }
+          // The last node copies node 0 of the same graph.
+          EXPECT_TRUE(registry.CopyNodeAnnotation(
+              graph_id, token, /*node_id=*/3, graph_id,
+              /*source_node_id=*/0, "fallback"));
+          // Unregister every other graph, the second time with a stale token.
+          if (g % 2 == 0) {
+            registry.UnregisterGraph(graph_id, token);
+            registry.UnregisterGraph(graph_id, token);
+          }
+        }
+        writers_running.fetch_sub(1, std::memory_order_release);
+      });
+    }
+    for (int v = 0; v < 2; ++v) {
+      pool.Schedule([&] {
+        while (writers_running.load(std::memory_order_acquire) > 0) {
+          for (uint32_t graph_id = 1; graph_id <= kNumGraphs; graph_id += 7) {
+            constexpr uint64_t kNodeId = 3;
+            if (registry.VisitNodeAnnotation(
+                    graph_id, kNodeId, [&](absl::string_view a) {
+                      EXPECT_EQ(a, annotation(graph_id, 0));
+                    })) {
+              visits_hit.fetch_add(1, std::memory_order_relaxed);
+            }
+          }
+        }
+      });
+    }
+    pool.Schedule([&] {
+      while (writers_running.load(std::memory_order_acquire) > 0) {
+        Registry::HarvestHold outer = registry.AcquireHarvestHold();
+        {
+          Registry::HarvestHold inner = registry.AcquireHarvestHold();
+          absl::SleepFor(absl::Microseconds(50));
+        }
+        EXPECT_GE(registry.GetStats().active_holds, 1);
+      }
+    });
+  }  // Joins all threads.
+
+  Registry::Stats stats = registry.GetStats();
+  EXPECT_EQ(stats.active_holds, 0);
+  EXPECT_EQ(stats.live_graphs, kNumGraphs / 2);
+  // The last release purged everything that was parked.
+  EXPECT_EQ(stats.parked_graphs, 0);
+  EXPECT_EQ(stats.parked_bytes, 0);
+  EXPECT_FALSE(registry.empty());
+  for (uint32_t graph_id = 1; graph_id <= kNumGraphs; ++graph_id) {
+    const bool unregistered = (graph_id - 1) % kGraphsPerWriter % 2 == 0;
+    EXPECT_EQ(registry.IsGraphTracked(graph_id), !unregistered);
+    if (unregistered) {
+      EXPECT_EQ(Lookup(registry, graph_id, 0), std::nullopt);
+    } else {
+      EXPECT_THAT(Lookup(registry, graph_id, 0),
+                  Optional(Eq(annotation(graph_id, 0))));
+    }
+  }
+}
+
+TEST(CudaGraphAnnotationRegistryTest, StatsReportLiveAndParkedSizes) {
+  Registry& registry = Registry::Default();
+  registry.ClearForTesting();
+
+  auto token1 = Track(registry, 1);
+  registry.RegisterNodeAnnotation(1, token1, 10, "node_a");
+  registry.RegisterNodeAnnotation(1, token1, 11, "node_b");
+
+  auto token2 = Track(registry, 2);
+  registry.RegisterNodeAnnotation(2, token2, 20, "node_c");
+
+  {
+    Registry::Stats stats = registry.GetStats();
+    EXPECT_EQ(stats.live_graphs, 2);
+    EXPECT_EQ(stats.live_nodes, 3);
+    EXPECT_EQ(stats.live_bytes, 18);
+  }
+
+  Registry::HarvestHold hold = registry.AcquireHarvestHold();
+  registry.UnregisterGraph(1, token1);
+
+  {
+    Registry::Stats stats = registry.GetStats();
+    EXPECT_EQ(stats.live_graphs, 1);
+    EXPECT_EQ(stats.live_nodes, 1);
+    EXPECT_EQ(stats.live_bytes, 6);
+    EXPECT_EQ(stats.parked_graphs, 1);
+    EXPECT_EQ(stats.parked_nodes, 2);
+    EXPECT_EQ(stats.parked_bytes, 12);
+
+    std::string str = absl::StrCat(stats);
+    EXPECT_THAT(str, ::testing::HasSubstr("live_bytes=6"));
+    EXPECT_THAT(str, ::testing::HasSubstr("parked_nodes=2"));
+  }
+}
+
+TEST(SelectTelemetryStrategyTest, EvaluatesRulesCorrectly) {
+  using TelemetryStrategy = CudaGraphAnnotationRegistry::TelemetryStrategy;
+
+  struct TestCase {
+    unsigned compiled_major;
+    unsigned compiled_minor;
+    unsigned drv_major;
+    unsigned drv_minor;
+    unsigned runtime_major;
+    unsigned runtime_minor;
+    TelemetryStrategy expected;
+  } test_cases[] = {
+      // zero versions
+      {0, 0, 13, 5, 13, 5, TelemetryStrategy::kDisabled},
+      {13, 4, 0, 0, 13, 5, TelemetryStrategy::kDisabled},
+      {13, 4, 13, 5, 0, 0, TelemetryStrategy::kDisabled},
+
+      // compiled 13.2 / 13.3 / 13.4
+      // driver 13.1 / 13.2 / 13.3 / 13.4 / 13.5
+      // runtime lower than the driver, and the reverse
+
+      // compiled 13.2, effective < 13.2 (disabled)
+      {13, 2, 13, 5, 13, 1, TelemetryStrategy::kDisabled},
+      {13, 2, 13, 1, 13, 5, TelemetryStrategy::kDisabled},
+
+      // compiled < 13.2 (disabled)
+      {13, 1, 13, 5, 13, 5, TelemetryStrategy::kDisabled},
+
+      // compiled 13.2, effective 13.2
+      {13, 2, 13, 2, 13, 2, TelemetryStrategy::kPre134ExecGraph},
+      // compiled 13.2, effective 13.5
+      {13, 2, 13, 5, 13, 5, TelemetryStrategy::kPre134ExecGraph},
+      // compiled 13.3, effective 13.5
+      {13, 3, 13, 5, 13, 5, TelemetryStrategy::kPre134ExecGraph},
+
+      // compiled 13.4, effective 13.3
+      {13, 4, 13, 5, 13, 3, TelemetryStrategy::kPre134ExecGraph},
+      {13, 4, 13, 3, 13, 5, TelemetryStrategy::kPre134ExecGraph},
+      // compiled 13.4, effective 13.4
+      {13, 4, 13, 4, 13, 4, TelemetryStrategy::kCuda134SourceGraph},
+      // compiled 13.4, effective 13.5
+      {13, 4, 13, 5, 13, 5, TelemetryStrategy::kCuda134SourceGraph},
+  };
+
+  for (const auto& test_case : test_cases) {
+    SemanticVersion compiled_toolkit{test_case.compiled_major,
+                                     test_case.compiled_minor, 0};
+    SemanticVersion driver{test_case.drv_major, test_case.drv_minor, 0};
+    SemanticVersion runtime{test_case.runtime_major, test_case.runtime_minor,
+                            0};
+
+    EXPECT_EQ(CudaGraphAnnotationRegistry::SelectTelemetryStrategy(
+                  compiled_toolkit, driver, runtime),
+              test_case.expected)
+        << "compiled_toolkit: " << compiled_toolkit << " driver: " << driver
+        << " runtime: " << runtime;
+  }
+}
+}  // namespace
+}  // namespace stream_executor::gpu
diff --git a/third_party/xla/xla/stream_executor/gpu/BUILD b/third_party/xla/xla/stream_executor/gpu/BUILD
index 4c759fc0f0d..9704ed2cdab 100644
--- a/third_party/xla/xla/stream_executor/gpu/BUILD
+++ b/third_party/xla/xla/stream_executor/gpu/BUILD
@@ -1175,6 +1175,7 @@ xla_cc_test(
     deps = [
         ":scoped_command_buffer_annotation",
         "//xla/tsl/platform:env",
+        "@com_google_absl//absl/strings",
         "@com_google_googletest//:gtest_main",
     ],
 )
diff --git a/third_party/xla/xla/stream_executor/gpu/scoped_command_buffer_annotation.cc b/third_party/xla/xla/stream_executor/gpu/scoped_command_buffer_annotation.cc
index da8ea91846b..dc88b6cc691 100644
--- a/third_party/xla/xla/stream_executor/gpu/scoped_command_buffer_annotation.cc
+++ b/third_party/xla/xla/stream_executor/gpu/scoped_command_buffer_annotation.cc
@@ -15,38 +15,82 @@ limitations under the License.
 
 #include "xla/stream_executor/gpu/scoped_command_buffer_annotation.h"
 
+#include <cstddef>
 #include <string>
-#include <utility>
 #include <vector>
 
+#include "absl/strings/string_view.h"
+
 namespace stream_executor {
 namespace {
 
-std::vector<std::string>& GetThreadLocalAnnotations() {
-  thread_local std::vector<std::string> annotations;
-  return annotations;
+struct ThreadState {
+  // The number of enabled roots open on the thread.
+  int depth = 0;
+  std::string joined;
+  // The size of `joined` before each pushed frame.
+  std::vector<size_t> prev_sizes;
+};
+
+ThreadState& state() {
+  thread_local ThreadState s;
+  return s;
 }
 
 }  // namespace
 
+ScopedCommandBufferAnnotation ScopedCommandBufferAnnotation::Root(
+    bool enabled, absl::string_view annotation) {
+  return ScopedCommandBufferAnnotation(RootTag{}, enabled, annotation);
+}
+
 ScopedCommandBufferAnnotation::ScopedCommandBufferAnnotation(
-    std::string annotation) {
-  GetThreadLocalAnnotations().push_back(std::move(annotation));
+    RootTag, bool enabled, absl::string_view annotation) {
+  if (!enabled) {
+    return;
+  }
+  ++state().depth;
+  is_root_ = true;
+  Push(annotation);
+}
+
+ScopedCommandBufferAnnotation::ScopedCommandBufferAnnotation(
+    absl::string_view annotation) {
+  if (IsActive()) {
+    Push(annotation);
+  }
 }
 
 ScopedCommandBufferAnnotation::~ScopedCommandBufferAnnotation() {
-  auto& annotations = GetThreadLocalAnnotations();
-  if (!annotations.empty()) {
-    annotations.pop_back();
+  ThreadState& s = state();
+  if (pushed_ && !s.prev_sizes.empty()) {
+    s.joined.resize(s.prev_sizes.back());
+    s.prev_sizes.pop_back();
+  }
+  if (is_root_ && --s.depth == 0) {
+    // Recording is off again: drop any leftover text.
+    s.joined.clear();
+    s.prev_sizes.clear();
   }
 }
 
-std::string ScopedCommandBufferAnnotation::GetCurrentAnnotation() {
-  const auto& annotations = GetThreadLocalAnnotations();
-  if (annotations.empty()) {
-    return "";
+void ScopedCommandBufferAnnotation::Push(absl::string_view annotation) {
+  if (annotation.empty()) {
+    return;
   }
-  return annotations.back();
+  ThreadState& s = state();
+  s.prev_sizes.push_back(s.joined.size());
+  if (!s.joined.empty()) {
+    s.joined.append("::");
+  }
+  s.joined.append(annotation.data(), annotation.size());
+  pushed_ = true;
+}
+
+bool ScopedCommandBufferAnnotation::IsActive() { return state().depth > 0; }
+
+absl::string_view ScopedCommandBufferAnnotation::GetCurrentAnnotation() {
+  return state().joined;
 }
 
 }  // namespace stream_executor
diff --git a/third_party/xla/xla/stream_executor/gpu/scoped_command_buffer_annotation.h b/third_party/xla/xla/stream_executor/gpu/scoped_command_buffer_annotation.h
index 75c08c9c631..9667135cc4c 100644
--- a/third_party/xla/xla/stream_executor/gpu/scoped_command_buffer_annotation.h
+++ b/third_party/xla/xla/stream_executor/gpu/scoped_command_buffer_annotation.h
@@ -16,40 +16,90 @@ limitations under the License.
 #ifndef XLA_STREAM_EXECUTOR_GPU_SCOPED_COMMAND_BUFFER_ANNOTATION_H_
 #define XLA_STREAM_EXECUTOR_GPU_SCOPED_COMMAND_BUFFER_ANNOTATION_H_
 
-#include <string>
+#include <type_traits>
 
 #include "absl/strings/string_view.h"
 
 namespace stream_executor {
 
-// A thread-local scope used to track the current HLO instruction being
-// recorded into a command buffer.
+// Thread-local stack of HLO annotations for the GPU command buffer nodes
+// created while command buffer annotation recording is on.
 //
-// Example usage:
-//   void RecordThunk(Thunk* thunk) {
-//     // Push the thunk's annotation onto the thread-local stack.
-//     ScopedCommandBufferAnnotation annotation(thunk->profile_annotation());
+// `CommandBufferThunk` opens the outermost scope with `Root()`, which turns
+// recording on only when command buffer telemetry is enabled. Inside it,
+// `CommandExecutor` pushes one scope per command with the lazy constructor, and
+// stream_executor reads `GetCurrentAnnotation()` and `IsActive()` while it
+// creates graph nodes. Like `tsl::profiler::ScopedAnnotation` while profiling
+// is off, scopes do nothing while recording is off, so producers do no
+// telemetry work unless telemetry is on.
 //
-//     // Any command buffer node creation calls nested within this scope
-//     // (even in lower-level libraries) can retrieve this annotation.
-//     RecordCommands();
-//   }
+// Path semantics: the retrieved `::`-joined string matches the
+// `tsl::profiler::AnnotationStack` format. It keeps the outer module prefix
+// (e.g. `XlaModule:#module=...#`) even for nodes within nested conditional or
+// while loop bodies.
 class ScopedCommandBufferAnnotation {
  public:
-  // Pushes the given `annotation` onto the thread-local annotation stack.
-  explicit ScopedCommandBufferAnnotation(std::string annotation);
+  // Opens the outermost scope of a command buffer recording. If `enabled`,
+  // annotation recording is on for the calling thread for the lifetime of the
+  // returned scope and `annotation` is its outermost frame (an empty
+  // `annotation` adds no frame); otherwise the returned scope does nothing.
+  // This is the only way to turn recording on. Roots nest.
+  //
+  // Returned by value through C++17 guaranteed copy elision: the class is
+  // neither copyable nor movable, so bind the result directly, e.g.
+  // `auto root = ScopedCommandBufferAnnotation::Root(...);`.
+  static ScopedCommandBufferAnnotation Root(bool enabled,
+                                            absl::string_view annotation);
+
+  // Appends `annotation` (separated by `"::"`) to the calling thread's stack
+  // for the lifetime of this scope while recording is on; does nothing
+  // otherwise, so a producer can never turn recording on. An empty
+  // `annotation` adds nothing.
+  explicit ScopedCommandBufferAnnotation(absl::string_view annotation);
+
+  // Like the above, but calls `generator` (returning anything convertible to
+  // `absl::string_view`, e.g. `std::string` or `absl::string_view`) only while
+  // recording is on, so producers do no work when it is off.
+  template <typename Generator,
+            std::enable_if_t<std::is_invocable_v<Generator&>, int> = 0>
+  explicit ScopedCommandBufferAnnotation(Generator&& generator) {
+    if (IsActive()) {
+      // Bound to a reference so that a returned `std::string` outlives the
+      // append.
+      const auto& annotation = generator();
+      Push(annotation);
+    }
+  }
 
-  // Pops the annotation from the thread-local annotation stack.
   ~ScopedCommandBufferAnnotation();
 
-  // Returns the annotation at the top of the thread-local stack,
-  // or empty string if the stack is empty.
-  static std::string GetCurrentAnnotation();
-
- private:
   ScopedCommandBufferAnnotation(const ScopedCommandBufferAnnotation&) = delete;
   ScopedCommandBufferAnnotation& operator=(
       const ScopedCommandBufferAnnotation&) = delete;
+
+  // Returns true while an enabled root is open on the calling thread. Readers
+  // that do telemetry work without pushing (e.g. graph finalization) check
+  // this.
+  static bool IsActive();
+
+  // Returns the calling thread's `"::"`-joined annotation stack, which is empty
+  // while recording is off. The returned view points into the thread-local
+  // buffer and is valid only until the next push or pop on this thread.
+  static absl::string_view GetCurrentAnnotation();
+
+ private:
+  struct RootTag {};
+
+  ScopedCommandBufferAnnotation(RootTag, bool enabled,
+                                absl::string_view annotation);
+
+  // Appends a non-empty `annotation` as a new frame and sets `pushed_`.
+  void Push(absl::string_view annotation);
+
+  // An enabled root: decrements the thread's root depth on destruction.
+  bool is_root_ = false;
+  // Pops its frame on destruction.
+  bool pushed_ = false;
 };
 
 }  // namespace stream_executor
diff --git a/third_party/xla/xla/stream_executor/gpu/scoped_command_buffer_annotation_test.cc b/third_party/xla/xla/stream_executor/gpu/scoped_command_buffer_annotation_test.cc
index 845f7555b3e..4503eed1ea3 100644
--- a/third_party/xla/xla/stream_executor/gpu/scoped_command_buffer_annotation_test.cc
+++ b/third_party/xla/xla/stream_executor/gpu/scoped_command_buffer_annotation_test.cc
@@ -15,56 +15,161 @@ limitations under the License.
 
 #include "xla/stream_executor/gpu/scoped_command_buffer_annotation.h"
 
+#include <memory>
+#include <string>
+
 #include <gtest/gtest.h>
+#include "absl/strings/string_view.h"
 #include "xla/tsl/platform/env.h"
-#include "xla/tsl/platform/threadpool.h"
 
 namespace stream_executor {
 namespace {
 
-TEST(ScopedCommandBufferAnnotationTest, EmptyStack) {
-  EXPECT_EQ(ScopedCommandBufferAnnotation::GetCurrentAnnotation(), "");
-}
+using Scope = ScopedCommandBufferAnnotation;
 
-TEST(ScopedCommandBufferAnnotationTest, SingleScope) {
+TEST(ScopedCommandBufferAnnotationTest, EnabledRootRecordsNestedPushes) {
+  EXPECT_FALSE(Scope::IsActive());
+  EXPECT_TRUE(Scope::GetCurrentAnnotation().empty());
   {
-    ScopedCommandBufferAnnotation annotation("layer1");
-    EXPECT_EQ(ScopedCommandBufferAnnotation::GetCurrentAnnotation(), "layer1");
-  }
-  EXPECT_EQ(ScopedCommandBufferAnnotation::GetCurrentAnnotation(), "");
-}
-
-TEST(ScopedCommandBufferAnnotationTest, NestedScopes) {
-  {
-    ScopedCommandBufferAnnotation annotation1("outer");
-    EXPECT_EQ(ScopedCommandBufferAnnotation::GetCurrentAnnotation(), "outer");
+    auto root = Scope::Root(/*enabled=*/true, "XlaModule:#m#");
+    EXPECT_TRUE(Scope::IsActive());
+    EXPECT_EQ(Scope::GetCurrentAnnotation(), "XlaModule:#m#");
     {
-      ScopedCommandBufferAnnotation annotation2("inner");
-      EXPECT_EQ(ScopedCommandBufferAnnotation::GetCurrentAnnotation(), "inner");
+      Scope inner("Thunk:#hlo_op=fusion.1#");
+      EXPECT_EQ(Scope::GetCurrentAnnotation(),
+                "XlaModule:#m#::Thunk:#hlo_op=fusion.1#");
+      {
+        // An empty push adds nothing.
+        Scope empty("");
+        EXPECT_EQ(Scope::GetCurrentAnnotation(),
+                  "XlaModule:#m#::Thunk:#hlo_op=fusion.1#");
+      }
     }
-    EXPECT_EQ(ScopedCommandBufferAnnotation::GetCurrentAnnotation(), "outer");
+    EXPECT_EQ(Scope::GetCurrentAnnotation(), "XlaModule:#m#");
   }
-  EXPECT_EQ(ScopedCommandBufferAnnotation::GetCurrentAnnotation(), "");
+  EXPECT_FALSE(Scope::IsActive());
+  EXPECT_TRUE(Scope::GetCurrentAnnotation().empty());
 }
 
-TEST(ScopedCommandBufferAnnotationTest, MultiThreadedIndependence) {
+TEST(ScopedCommandBufferAnnotationTest, DisabledRootDoesNothing) {
+  auto root = Scope::Root(/*enabled=*/false, "XlaModule:#m#");
+  EXPECT_FALSE(Scope::IsActive());
+  EXPECT_TRUE(Scope::GetCurrentAnnotation().empty());
+  Scope inner("Thunk:#hlo_op=fusion.1#");
+  EXPECT_FALSE(Scope::IsActive());
+  EXPECT_TRUE(Scope::GetCurrentAnnotation().empty());
+}
+
+TEST(ScopedCommandBufferAnnotationTest, PushWithoutRootDoesNothing) {
+  Scope eager("XlaModule:#m#");
+  EXPECT_FALSE(Scope::IsActive());
+  EXPECT_TRUE(Scope::GetCurrentAnnotation().empty());
+
+  int calls = 0;
+  Scope lazy([&] {
+    ++calls;
+    return std::string("Thunk:#hlo_op=fusion.1#");
+  });
+  EXPECT_EQ(calls, 0);
+  EXPECT_TRUE(Scope::GetCurrentAnnotation().empty());
+}
+
+TEST(ScopedCommandBufferAnnotationTest, LazyGeneratorCalledOnceInsideRoot) {
+  auto root = Scope::Root(/*enabled=*/true, "XlaModule:#m#");
+  int string_calls = 0;
+  int view_calls = 0;
   {
-    tsl::thread::ThreadPool pool(tsl::Env::Default(), "test_pool", 2);
-
-    pool.Schedule([] {
-      ScopedCommandBufferAnnotation annotation("thread1_annotation");
-      EXPECT_EQ(ScopedCommandBufferAnnotation::GetCurrentAnnotation(),
-                "thread1_annotation");
-    });
-
-    pool.Schedule([] {
-      ScopedCommandBufferAnnotation annotation("thread2_annotation");
-      EXPECT_EQ(ScopedCommandBufferAnnotation::GetCurrentAnnotation(),
-                "thread2_annotation");
+    Scope from_string([&] {
+      ++string_calls;
+      return std::string("Thunk:#hlo_op=fusion.1#");
     });
+    EXPECT_EQ(string_calls, 1);
+    EXPECT_EQ(Scope::GetCurrentAnnotation(),
+              "XlaModule:#m#::Thunk:#hlo_op=fusion.1#");
+    {
+      Scope from_view([&] {
+        ++view_calls;
+        return absl::string_view("Thunk:#hlo_op=fusion.2#");
+      });
+      EXPECT_EQ(view_calls, 1);
+      EXPECT_EQ(Scope::GetCurrentAnnotation(),
+                "XlaModule:#m#::Thunk:#hlo_op=fusion.1#::"
+                "Thunk:#hlo_op=fusion.2#");
+    }
   }
-  // Pool destruction waits for workers to finish.
-  EXPECT_EQ(ScopedCommandBufferAnnotation::GetCurrentAnnotation(), "");
+  EXPECT_EQ(string_calls, 1);
+  EXPECT_EQ(view_calls, 1);
+  EXPECT_EQ(Scope::GetCurrentAnnotation(), "XlaModule:#m#");
+}
+
+TEST(ScopedCommandBufferAnnotationTest, NestedRootsKeepRecordingOn) {
+  {
+    auto outer = Scope::Root(/*enabled=*/true, "outer");
+    {
+      auto inner = Scope::Root(/*enabled=*/true, "inner");
+      EXPECT_TRUE(Scope::IsActive());
+      EXPECT_EQ(Scope::GetCurrentAnnotation(), "outer::inner");
+    }
+    // The outer root keeps recording on after the inner one closes.
+    EXPECT_TRUE(Scope::IsActive());
+    EXPECT_EQ(Scope::GetCurrentAnnotation(), "outer");
+    {
+      // A disabled inner root does not turn recording off.
+      auto disabled = Scope::Root(/*enabled=*/false, "disabled");
+      EXPECT_TRUE(Scope::IsActive());
+      EXPECT_EQ(Scope::GetCurrentAnnotation(), "outer");
+    }
+  }
+  EXPECT_FALSE(Scope::IsActive());
+  EXPECT_TRUE(Scope::GetCurrentAnnotation().empty());
+}
+
+TEST(ScopedCommandBufferAnnotationTest, EmptyRootEnablesRecording) {
+  auto root = Scope::Root(/*enabled=*/true, "");
+  EXPECT_TRUE(Scope::IsActive());
+  EXPECT_TRUE(Scope::GetCurrentAnnotation().empty());
+  {
+    // No leading "::" after an empty root.
+    Scope inner("Thunk:#hlo_op=fusion.1#");
+    EXPECT_EQ(Scope::GetCurrentAnnotation(), "Thunk:#hlo_op=fusion.1#");
+  }
+  EXPECT_TRUE(Scope::GetCurrentAnnotation().empty());
+}
+
+TEST(ScopedCommandBufferAnnotationTest, StringArgumentsUseEagerConstructor) {
+  // `const char*` and `std::string` are not invocable, so they bind to the
+  // `absl::string_view` constructor rather than the lazy one.
+  auto root = Scope::Root(/*enabled=*/true, "XlaModule:#m#");
+  {
+    Scope literal("literal");
+    EXPECT_EQ(Scope::GetCurrentAnnotation(), "XlaModule:#m#::literal");
+  }
+  {
+    const std::string annotation = "string";
+    Scope from_string(annotation);
+    EXPECT_EQ(Scope::GetCurrentAnnotation(), "XlaModule:#m#::string");
+  }
+  EXPECT_EQ(Scope::GetCurrentAnnotation(), "XlaModule:#m#");
+}
+
+TEST(ScopedCommandBufferAnnotationTest, ThreadIsolation) {
+  auto main_root = Scope::Root(/*enabled=*/true, "XlaModule:#module=main#");
+  EXPECT_EQ(Scope::GetCurrentAnnotation(), "XlaModule:#module=main#");
+
+  bool active_on_other_thread = true;
+  bool current_empty_on_other_thread = false;
+  std::unique_ptr<tsl::Thread> other(
+      tsl::Env::Default()->StartThread(tsl::ThreadOptions(), "other", [&] {
+        active_on_other_thread = Scope::IsActive();
+        current_empty_on_other_thread = Scope::GetCurrentAnnotation().empty();
+      }));
+  other.reset();
+
+  EXPECT_FALSE(active_on_other_thread);
+  EXPECT_TRUE(current_empty_on_other_thread);
+
+  EXPECT_TRUE(Scope::IsActive());
+  EXPECT_EQ(Scope::GetCurrentAnnotation(), "XlaModule:#module=main#");
 }
 
 }  // namespace

From a9c3219d2e71c0fbe70eac650bfaa29e7782665d Mon Sep 17 00:00:00 2001
From: Shyamli Agrawal <ashyamli@google.com>
Date: Fri, 16 Oct 2026 13:58:39 -0700
Subject: [PATCH 0979/1000] Add xla_autotune_cache_require_strict_codegen_match
 flag to control TieredCache KeyMatchingMode.

When true, autotune cache lookup requires the entire codegen environment to match (KeyMatchingMode::kStrict).
When false (default), only the codegen version of the optimal backend is required to match (KeyMatchingMode::kLoose).

PiperOrigin-RevId: 1001081814
---
 .../xla/xla/backends/gpu/autotuner/BUILD        | 17 -----------------
 .../backends/gpu/autotuner/autotuner_main.cc    |  9 ++++++---
 third_party/xla/xla/debug_options_flags.cc      | 11 +++++++++++
 third_party/xla/xla/debug_options_flags_test.cc | 14 ++++++++++++++
 .../gpu/autotuning/config_assigner_pass.cc      |  9 ++++++---
 third_party/xla/xla/xla.proto                   |  8 +++++++-
 6 files changed, 44 insertions(+), 24 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/autotuner/BUILD b/third_party/xla/xla/backends/gpu/autotuner/BUILD
index 8003f53fc7d..bd7d33371e8 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/BUILD
+++ b/third_party/xla/xla/backends/gpu/autotuner/BUILD
@@ -1035,19 +1035,15 @@ xla_cc_binary(
         "gpu",
     ],
     deps = [
-        ":factory",
         ":gpu_profiler",
-        ":legacy_cache",
         "//xla:debug_options_flags",
         "//xla:status_macros",
         "//xla:xla_proto_cc",
         "//xla/backends/autotuner",
-        "//xla/backends/autotuner:autotune_cache_store",
         "//xla/backends/autotuner:autotuner_cache_interface",
         "//xla/backends/autotuner:backends_proto_cc",
         "//xla/backends/autotuner:codegen_backend",
         "//xla/backends/autotuner:codegen_orchestrator",
-        "//xla/backends/autotuner:config_assigner",
         "//xla/backends/autotuner:directory_store",
         "//xla/backends/autotuner:hlo_extractor",
         "//xla/backends/autotuner:in_memory_store",
@@ -1056,36 +1052,24 @@ xla_cc_binary(
         "//xla/backends/gpu/codegen/emitters:mlir_kernel_emitter",
         "//xla/hlo/analysis:symbolic_map",
         "//xla/hlo/ir:hlo",
-        "//xla/hlo/parser:hlo_parser",
         "//xla/service:compiler",
         "//xla/service:gpu_plugin",
         "//xla/service:platform_util",
-        "//xla/service/gpu:alias_info",
-        "//xla/service/gpu:backend_configs_cc",
-        "//xla/service/gpu:cublas_cudnn",
         "//xla/service/gpu:gpu_compiler",
-        "//xla/service/gpu:ir_emission_utils",
         "//xla/service/gpu:mlir_context_pool",
         "//xla/service/gpu/autotuning:config_assigner_pass",
         "//xla/stream_executor:device_address_allocator",
         "//xla/stream_executor:platform",
         "//xla/stream_executor:platform_manager",
-        "//xla/stream_executor:stream_executor_address_allocator",
-        "//xla/stream_executor:stream_executor_memory_allocator",
-        "//xla/stream_executor/cuda:cuda_platform_id",
         "//xla/stream_executor/integrations:device_mem_allocator",
         "//xla/stream_executor/integrations:tf_allocator_adapter",
-        "//xla/stream_executor/platform:platform_object_registry",
         "//xla/tools:hlo_module_loader",
         "//xla/tsl/concurrency:executor",
         "//xla/tsl/framework:bfc_allocator",
         "//xla/tsl/framework:device_id",
         "//xla/tsl/platform:env",
-        "//xla/tsl/platform:errors",
-        "//xla/tsl/platform:statusor",
         "//xla/tsl/util:command_line_flags",
         "@com_google_absl//absl/base",
-        "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status",
@@ -1097,7 +1081,6 @@ xla_cc_binary(
         "@com_google_absl//absl/time",
         "@com_google_absl//absl/types:span",
         "@llvm-project//mlir:IR",
-        "@tsl//tsl/platform:casts",
         "@tsl//tsl/platform:platform_port",
     ] + if_cuda_is_configured([
         ":factory_cuda",
diff --git a/third_party/xla/xla/backends/gpu/autotuner/autotuner_main.cc b/third_party/xla/xla/backends/gpu/autotuner/autotuner_main.cc
index 327b80aed5e..4d939e13908 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/autotuner_main.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/autotuner_main.cc
@@ -336,9 +336,12 @@ absl::Status RunAutotuning(const std::vector<std::string>& hlo_files,
     auto primary = std::make_unique<InMemoryStore>();
     auto secondary = std::make_unique<DirectoryStore>(std::string(cache_dir),
                                                       CacheMode::kReadWrite);
-    autotuner_cache =
-        std::make_unique<TieredCache>(env.cache_ctx, KeyMatchingMode::kLoose,
-                                      std::move(primary), std::move(secondary));
+    KeyMatchingMode matching_mode =
+        debug_options.xla_autotune_cache_require_strict_codegen_match()
+            ? KeyMatchingMode::kStrict
+            : KeyMatchingMode::kLoose;
+    autotuner_cache = std::make_unique<TieredCache>(
+        env.cache_ctx, matching_mode, std::move(primary), std::move(secondary));
   } else {
     autotuner_cache = std::make_unique<PrintingAutotunerCache>();
   }
diff --git a/third_party/xla/xla/debug_options_flags.cc b/third_party/xla/xla/debug_options_flags.cc
index 77c317693c7..2b23fce7087 100644
--- a/third_party/xla/xla/debug_options_flags.cc
+++ b/third_party/xla/xla/debug_options_flags.cc
@@ -490,6 +490,7 @@ DebugOptions DefaultDebugOptionsIgnoringFlags() {
   opts.set_xla_gpu_per_fusion_autotune_cache_dir("");
 
   opts.set_xla_gpu_use_new_autotune_cache_format(true);
+  opts.set_xla_autotune_cache_require_strict_codegen_match(false);
 
   opts.set_xla_compile_all_supported_configs(false);
   opts.set_xla_deduplicate_backend_configs_min_size(
@@ -3145,6 +3146,16 @@ void MakeDebugOptionsFlags(std::vector<tsl::Flag>* flag_list,
       "Whether to use the new protos for the autotune cache"
       " (xla.autotuner.AutotuneCache rather than xla.AutotuneResults."));
 
+  flag_list->push_back(tsl::Flag(
+      "xla_autotune_cache_require_strict_codegen_match",
+      bool_setter_for(
+          &DebugOptions::set_xla_autotune_cache_require_strict_codegen_match),
+      debug_options->xla_autotune_cache_require_strict_codegen_match(),
+      "When true, autotune cache lookup requires the entire codegen "
+      "environment to match. When false, only the codegen version of the "
+      "optimal backend is required to match (e.g., allowing Triton cached "
+      "entries to be reused even when the cuDNN version changes)."));
+
   flag_list->push_back(tsl::Flag(
       "xla_compile_all_supported_configs",
       bool_setter_for(&DebugOptions::set_xla_compile_all_supported_configs),
diff --git a/third_party/xla/xla/debug_options_flags_test.cc b/third_party/xla/xla/debug_options_flags_test.cc
index 86eba8bf450..c52480633aa 100644
--- a/third_party/xla/xla/debug_options_flags_test.cc
+++ b/third_party/xla/xla/debug_options_flags_test.cc
@@ -530,5 +530,19 @@ TEST(DebugOptions, DeterministicOpsTriton) {
   EXPECT_FALSE(parse("invalid"));
 }
 
+TEST(DebugOptions, AutotuneCacheRequireStrictCodegenMatchDefaultIsFalse) {
+  DebugOptions opts = DefaultDebugOptionsIgnoringFlags();
+  EXPECT_FALSE(opts.xla_autotune_cache_require_strict_codegen_match());
+}
+
+TEST(DebugOptions, AutotuneCacheRequireStrictCodegenMatchFlagsParsing) {
+  DebugOptions opts;
+  std::vector<tsl::Flag> flags;
+  MakeDebugOptionsFlags(&flags, &opts);
+  std::vector<std::string> flag_args = {
+      "--xla_autotune_cache_require_strict_codegen_match=true"};
+  EXPECT_TRUE(tsl::Flags::Parse(flag_args, flags));
+  EXPECT_TRUE(opts.xla_autotune_cache_require_strict_codegen_match());
+}
 }  // namespace
 }  // namespace xla
diff --git a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc
index 146cc60d759..40d1bf5a3bd 100644
--- a/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc
+++ b/third_party/xla/xla/service/gpu/autotuning/config_assigner_pass.cc
@@ -273,9 +273,12 @@ std::unique_ptr<AutotunerCacheInterface> CreateAutotunerCache(
         cache_dir,
         GetCacheMode(debug_options.xla_gpu_experimental_autotune_cache_mode()));
   }
-  return std::make_unique<TieredCache>(cache_ctx, KeyMatchingMode::kLoose,
-                                       std::move(primary),
-                                       std::move(secondary));
+  KeyMatchingMode matching_mode =
+      debug_options.xla_autotune_cache_require_strict_codegen_match()
+          ? KeyMatchingMode::kStrict
+          : KeyMatchingMode::kLoose;
+  return std::make_unique<TieredCache>(
+      cache_ctx, matching_mode, std::move(primary), std::move(secondary));
 }
 
 }  // namespace
diff --git a/third_party/xla/xla/xla.proto b/third_party/xla/xla/xla.proto
index 6aa03f34c95..867f2c106de 100644
--- a/third_party/xla/xla/xla.proto
+++ b/third_party/xla/xla/xla.proto
@@ -1961,7 +1961,13 @@ message DebugOptions {
   // dimension are tiled for now; extend to the other cases.
   optional bool xla_gpu_experimental_enable_reduction_dimension_tiling = 558;
 
-  // Next id: 559
+  // When true, autotune cache lookup requires the entire codegen environment
+  // to match. When false, only the codegen version of the optimal backend is
+  // required to match (e.g., allowing Triton cached entries to be reused even
+  // when the cuDNN version changes).
+  optional bool xla_autotune_cache_require_strict_codegen_match = 559;
+
+  // Next id: 560
 
   // Extra options to pass to the compilation backend (e.g. LLVM); specific
   // interpretation of these values is left to the backend.

From c2f595dd7bf5e1dc3f505c75632ca89c8ac84c62 Mon Sep 17 00:00:00 2001
From: Ce Zheng CA <zce@google.com>
Date: Fri, 16 Oct 2026 13:59:34 -0700
Subject: [PATCH 0980/1000] Replace broad `core_cpu_lib` and
 `core_cpu_internal` deps in `tensor_bundle` and `tf2xla:common` with granular
 targets

1. `//tensorflow/core/util/tensor_bundle:tensor_bundle` depended on `//tensorflow/core:core_cpu_lib` solely to satisfy `layering_check` for `#include "tensorflow/core/framework/versions.h"`, whose implementation (`versions.cc`) is already compiled in `//tensorflow/core:framework_internal`. Exposing `//tensorflow/core/framework:versions.h` in `framework_internal_public_headers` allows removing `core_cpu_lib` from `tensor_bundle`.
2. `//tensorflow/compiler/tf2xla:common` only includes `tensorflow/core/common_runtime/dma_helper.h`; replace `//tensorflow/core:core_cpu_internal` with `//tensorflow/core/common_runtime:dma_helper`.

Runtime and memory performance impact: Zero CPU runtime or memory regression; avoids pulling `core_cpu_impl` and `core_cpu_internal` (`alwayslink = 1` graph execution and optimization passes) into targets that only depend on `tensor_bundle` or `tf2xla:common`.

PiperOrigin-RevId: 1001082343
---
 tensorflow/compiler/tf2xla/BUILD         | 2 +-
 tensorflow/core/BUILD                    | 1 +
 tensorflow/core/util/tensor_bundle/BUILD | 1 -
 3 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tensorflow/compiler/tf2xla/BUILD b/tensorflow/compiler/tf2xla/BUILD
index e3e2a078863..042a4f3a373 100644
--- a/tensorflow/compiler/tf2xla/BUILD
+++ b/tensorflow/compiler/tf2xla/BUILD
@@ -1012,10 +1012,10 @@ cc_library(
     ],
     visibility = [":friends"],
     deps = [
-        "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
         "//tensorflow/core:protos_all_cc",
+        "//tensorflow/core/common_runtime:dma_helper",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD
index 69bc1c4662d..824404a69b8 100644
--- a/tensorflow/core/BUILD
+++ b/tensorflow/core/BUILD
@@ -1694,6 +1694,7 @@ filegroup(
         "//tensorflow/core/framework:tensor_reference.h",
         "//tensorflow/core/framework:tracking_allocator.h",  # only needed for tests
         "//tensorflow/core/framework:variant.h",
+        "//tensorflow/core/framework:versions.h",
         "//tensorflow/core/util:framework_internal_public_hdrs",
         "@xla//xla/tsl/framework:tracking_allocator.h",  # only needed for tests
     ],
diff --git a/tensorflow/core/util/tensor_bundle/BUILD b/tensorflow/core/util/tensor_bundle/BUILD
index 29df74fd208..8c14ccb0f49 100644
--- a/tensorflow/core/util/tensor_bundle/BUILD
+++ b/tensorflow/core/util/tensor_bundle/BUILD
@@ -73,7 +73,6 @@ cc_library(
     deps = [
         ":byteswaptensor",
         ":naming",
-        "//tensorflow/core:core_cpu_lib",
         "//tensorflow/core:framework",
         "//tensorflow/core:framework_internal",
         "//tensorflow/core:lib",

From 181a7b273699ad1492e66b25692f3a1511396430 Mon Sep 17 00:00:00 2001
From: Sohaib Iftikhar <sohaibiftikhar@google.com>
Date: Fri, 16 Oct 2026 14:03:12 -0700
Subject: [PATCH 0981/1000] [XLA:GPU]: E2E variadic all-gather support.

PiperOrigin-RevId: 1001084563
---
 .../gpu/codegen/triton/collective_emitter.cc  | 331 +++++++++++++-----
 .../codegen/triton/collective_emitter_test.cc | 122 +++++++
 .../xla/backends/gpu/runtime/all_gather.cc    |  20 +-
 .../gpu/runtime/all_gather_build_info_test.cc |  39 +++
 .../backends/gpu/tests/all_gather_e2e_test.cc | 172 +++++++++
 .../collectives/collective_fusion_test.cc     |  42 +++
 6 files changed, 616 insertions(+), 110 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
index 49ee9d6d4c1..81757c3ab54 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter.cc
@@ -52,6 +52,7 @@ limitations under the License.
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypeInterfaces.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/PatternMatch.h"
@@ -72,6 +73,7 @@ limitations under the License.
 #include "xla/codegen/xtile/ir/xtile_ops.h"
 #include "xla/codegen/xtile/xtile_config.pb.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_opcode.h"
@@ -1323,6 +1325,31 @@ absl::StatusOr<llvm::SmallVector<int64_t>> GetAllGatherOutputTile(
   return tile_sizes;
 }
 
+// Computes one output tile per operand of an all-gather. Each output is tiled
+// independently using `num_blocks / num_devices` programs per peer; for
+// variadic all-gathers, smaller outputs get padded with masked-out blocks by
+// `GetIndexWiseVariadicSchedule`.
+absl::StatusOr<std::vector<llvm::SmallVector<int64_t>>> GetAllGatherOutputTiles(
+    const HloAllGatherInstruction* all_gather, int64_t num_devices,
+    int32_t num_blocks) {
+  const int64_t gather_dim = all_gather->all_gather_dimension();
+  const int32_t blocks_per_peer =
+      std::max<int32_t>(1, num_blocks / static_cast<int32_t>(num_devices));
+  const absl::Span<const Shape> output_shapes =
+      all_gather->shape().IsTuple()
+          ? absl::MakeConstSpan(all_gather->shape().tuple_shapes())
+          : absl::MakeConstSpan(&all_gather->shape(), 1);
+  std::vector<llvm::SmallVector<int64_t>> tiles;
+  tiles.reserve(output_shapes.size());
+  for (const Shape& output_shape : output_shapes) {
+    ABSL_ASSIGN_OR_RETURN(llvm::SmallVector<int64_t> tile_sizes,
+                     GetAllGatherOutputTile(output_shape, gather_dim,
+                                            num_devices, blocks_per_peer));
+    tiles.push_back(std::move(tile_sizes));
+  }
+  return tiles;
+}
+
 absl::StatusOr<BlockLevelFusionConfig> GetBlockLevelConfigForAllReduce(
     const GpuTopology& gpu_topology, const HloAllReduceInstruction* instr,
     const DeviceAssignment* device_assignment) {
@@ -1382,13 +1409,10 @@ absl::StatusOr<BlockLevelFusionConfig> GetBlockLevelConfigForAllGather(
   const int64_t num_blocks =
       std::min(kAllGatherMaxBlocksPerGrid, CeilOfRatio(total_warps, num_warps));
 
-  const int32_t blocks_per_peer = std::max<int32_t>(
-      1, static_cast<int32_t>(num_blocks / all_gather_info.num_devices));
   ABSL_ASSIGN_OR_RETURN(
-      llvm::SmallVector<int64_t> tile_sizes,
-      GetAllGatherOutputTile(instr->shape(), instr->all_gather_dimension(),
-                             all_gather_info.num_devices, blocks_per_peer));
-  return CreateBlockLevelFusionConfig(num_warps, tile_sizes);
+      std::vector<llvm::SmallVector<int64_t>> output_tiles,
+      GetAllGatherOutputTiles(instr, all_gather_info.num_devices, num_blocks));
+  return CreateBlockLevelFusionConfig(num_warps, output_tiles);
 }
 
 absl::StatusOr<BlockLevelFusionConfig> GetBlockLevelConfigForReduceScatter(
@@ -1434,29 +1458,103 @@ absl::StatusOr<BlockLevelFusionConfig> GetBlockLevelConfigForReduceScatter(
   return CreateBlockLevelFusionConfig(num_warps, tile_sizes);
 }
 
-struct AllGatherRewriteContext {
-  mlir::stablehlo::AllGatherOp op;
+// Per-operand state of an all-gather rewrite.
+struct AllGatherOperandContext {
+  // Tile of the local input that feeds the all-gather. For booleans this is
+  // the i8 -> i1 cast of `input_extract` (see i1_to_i8_workaround).
+  mlir::Value input_tile;
   xtile::ExtractTileOp input_extract;
   // Local input buffer read by `input_extract`.
   mlir::TypedValue<mlir::MemRefType> input_buffer;
-  mlir::Value rank_arg;
-  mlir::Value signal_buffers_arg;
   // !tt.ptr<i64> table of the gathered parameter's per-rank remote buffers.
   mlir::Value remote_buffers_arg;
-  int64_t world_size = 0;
-  int64_t gather_dim = 0;
   int64_t split_dim = -1;
-  int64_t signal_stride = 0;
   PrimitiveType element_type = PrimitiveType::PRIMITIVE_TYPE_INVALID;
   llvm::ArrayRef<int64_t> per_rank_shape;
   llvm::ArrayRef<int64_t> tile_sizes;
 };
 
+struct AllGatherRewriteContext {
+  mlir::stablehlo::AllGatherOp op;
+  llvm::SmallVector<AllGatherOperandContext, 2> operands;
+  mlir::Value rank_arg;
+  mlir::Value signal_buffers_arg;
+  int64_t world_size = 0;
+  int64_t gather_dim = 0;
+  // Number of consecutive programs that read from the same peer. See
+  // EmitPeerReplicaId.
+  int64_t signal_stride = 0;
+};
+
+// Returns the number of programs needed to cover one rank's slice of an
+// operand, i.e. the product of the per-dimension tile counts.
+int64_t ProgramsPerRank(const AllGatherOperandContext& operand) {
+  int64_t programs = 1;
+  for (auto [dim_size, tile_size] :
+       llvm::zip(operand.per_rank_shape, operand.tile_sizes)) {
+    programs *= xla::CeilOfRatio(dim_size, tile_size);
+  }
+  return programs;
+}
+
+absl::StatusOr<AllGatherOperandContext> BuildAllGatherOperandContext(
+    mlir::stablehlo::AllGatherOp op, int64_t operand_index,
+    xtile::EntryFuncOp entry_func, int64_t num_parameters,
+    int32_t metadata_args_start, int64_t world_size, int64_t gather_dim) {
+  ABSL_ASSIGN_OR_RETURN(
+      CollectiveOperandExtractInfo operand_info,
+      ExtractCollectiveOperand(
+          op.getOperand(operand_index), operand_index, "AllGather", entry_func,
+          num_parameters, metadata_args_start + kNumCollectiveMetadataArgs));
+  AllGatherOperandContext ctx;
+  ctx.input_tile = operand_info.input_tile;
+  ctx.input_extract = operand_info.input_extract;
+  ctx.input_buffer = operand_info.input_buffer;
+  ctx.remote_buffers_arg = operand_info.remote_buffers_arg;
+  ctx.element_type = operand_info.element_type;
+  ctx.per_rank_shape = ctx.input_buffer.getType().getShape();
+  ctx.tile_sizes = ctx.input_extract.getFullTileShape();
+
+  // Every tile must be gathered from a single rank (see EmitPeerReplicaId).
+  const int64_t per_rank_size = ctx.per_rank_shape[gather_dim];
+  const int64_t tile_size = ctx.tile_sizes[gather_dim];
+  if (per_rank_size % tile_size != 0) {
+    return absl::InvalidArgumentError(absl::StrFormat(
+        "AllGather tile size (%d) must divide the per-rank size (%d) along the "
+        "gather dimension.",
+        tile_size, per_rank_size));
+  }
+
+  // Find a dimension to split the tile size across `world_size` to balance
+  // the in-kernel D2D subtile copy work across participating ranks.
+  // We prefer splitting along `gather_dim`, but fallback to the innermost
+  // cleanly divisible dimension. If none are found, split_dim remains -1
+  // (unsplit).
+  if (ctx.tile_sizes[gather_dim] >= world_size &&
+      ctx.tile_sizes[gather_dim] % world_size == 0) {
+    ctx.split_dim = gather_dim;
+  } else {
+    for (int64_t d = static_cast<int64_t>(ctx.tile_sizes.size()) - 1; d >= 0;
+         --d) {
+      if (ctx.tile_sizes[d] >= world_size &&
+          ctx.tile_sizes[d] % world_size == 0) {
+        ctx.split_dim = d;
+        break;
+      }
+    }
+  }
+  return ctx;
+}
+
 absl::StatusOr<AllGatherRewriteContext> BuildAllGatherContext(
     mlir::stablehlo::AllGatherOp op) {
-  if (op.getOperands().size() != 1) {
+  if (op.getOperands().empty()) {
     return absl::InvalidArgumentError(
-        "AllGather op must have exactly one operand.");
+        "AllGather op must have at least one operand.");
+  }
+  if (op.getNumOperands() != op.getNumResults()) {
+    return absl::InvalidArgumentError(
+        "AllGather op must have one result per operand.");
   }
   auto entry_func = op->getParentOfType<xtile::EntryFuncOp>();
   if (!entry_func) {
@@ -1479,6 +1577,11 @@ absl::StatusOr<AllGatherRewriteContext> BuildAllGatherContext(
         "EntryFuncOp must have the collective metadata arguments followed by "
         "the remote buffer arguments.");
   }
+  const int64_t num_parameters =
+      num_opaque_attr.getInt() - kNumCollectiveMetadataArgs;
+  const int32_t total_args = entry_func.getNumArguments();
+  const int32_t metadata_args_start =
+      total_args - num_opaque_attr.getInt() - kNumTileIndexArgs;
 
   ABSL_ASSIGN_OR_RETURN(auto replica_groups,
                    xla::ConvertReplicaGroups(op.getReplicaGroups(), op));
@@ -1487,71 +1590,47 @@ absl::StatusOr<AllGatherRewriteContext> BuildAllGatherContext(
         "AllGather replica groups must not be empty.");
   }
 
-  const int64_t num_parameters =
-      num_opaque_attr.getInt() - kNumCollectiveMetadataArgs;
-  int32_t total_args = entry_func.getNumArguments();
-  int32_t metadata_args_start =
-      total_args - num_opaque_attr.getInt() - kNumTileIndexArgs;
-  ABSL_ASSIGN_OR_RETURN(
-      CollectiveOperandExtractInfo extract_info,
-      ExtractCollectiveOperand(op.getOperand(0), /*operand_index=*/0,
-                               "AllGather", entry_func, num_parameters,
-                               metadata_args_start + kNumCollectiveMetadataArgs,
-                               /*require_single_use=*/true));
-
   AllGatherRewriteContext ctx;
   ctx.op = op;
-  ctx.input_extract = extract_info.input_extract;
-  ctx.input_buffer = extract_info.input_buffer;
   ctx.world_size = replica_groups->num_devices_per_group();
   ctx.gather_dim = op.getAllGatherDim();
-  ctx.element_type = extract_info.element_type;
-  ctx.per_rank_shape = ctx.input_buffer.getType().getShape();
-  ctx.tile_sizes = extract_info.input_extract.getFullTileShape();
-
-  // Every tile must be gathered from a single rank (see EmitPeerReplicaId).
-  const int64_t per_rank_size = ctx.per_rank_shape[ctx.gather_dim];
-  const int64_t tile_size = ctx.tile_sizes[ctx.gather_dim];
-  if (per_rank_size % tile_size != 0) {
-    return absl::InvalidArgumentError(absl::StrFormat(
-        "AllGather tile size (%d) must divide the per-rank size (%d) along the "
-        "gather dimension.",
-        tile_size, per_rank_size));
-  }
-
   // Layout: opaque[0]=rank, opaque[1]=invocation count (unused, the signal
   // value comes from a counter in device memory), opaque[2]=signal_buffers,
   // opaque[3 + i]=remote buffers of parameter i (same as AllReduce).
   ctx.rank_arg = entry_func.getArgument(metadata_args_start);
   ctx.signal_buffers_arg = entry_func.getArgument(metadata_args_start + 2);
-  ctx.remote_buffers_arg = extract_info.remote_buffers_arg;
 
-  // Find a dimension to split the tile size across `world_size` to balance
-  // the in-kernel D2D subtile copy work across participating ranks.
-  // We prefer splitting along `gather_dim`, but fallback to the innermost
-  // cleanly divisible dimension. If none are found, split_dim remains -1
-  // (unsplit).
-  if (ctx.tile_sizes[ctx.gather_dim] >= ctx.world_size &&
-      ctx.tile_sizes[ctx.gather_dim] % ctx.world_size == 0) {
-    ctx.split_dim = ctx.gather_dim;
-  } else {
-    for (int64_t d = static_cast<int64_t>(ctx.tile_sizes.size()) - 1; d >= 0;
-         --d) {
-      if (ctx.tile_sizes[d] >= ctx.world_size &&
-          ctx.tile_sizes[d] % ctx.world_size == 0) {
-        ctx.split_dim = d;
-        break;
-      }
+  for (int64_t i = 0; i < op.getNumOperands(); ++i) {
+    ABSL_ASSIGN_OR_RETURN(AllGatherOperandContext operand,
+                     BuildAllGatherOperandContext(
+                         op, i, entry_func, num_parameters, metadata_args_start,
+                         ctx.world_size, ctx.gather_dim));
+    ctx.operands.push_back(std::move(operand));
+  }
+
+  if (op.getNumOperands() == 1) {
+    // Single operand: programs enumerate the output tiles in row-major order,
+    // so `signal_stride` consecutive programs cover one rank's slice along
+    // `gather_dim`.
+    const AllGatherOperandContext& operand = ctx.operands.front();
+    ctx.signal_stride = operand.per_rank_shape[ctx.gather_dim] /
+                        operand.tile_sizes[ctx.gather_dim];
+    for (int64_t j = ctx.gather_dim + 1;
+         j < static_cast<int64_t>(operand.per_rank_shape.size()); ++j) {
+      ctx.signal_stride *=
+          xla::CeilOfRatio(operand.per_rank_shape[j], operand.tile_sizes[j]);
     }
+    return ctx;
   }
 
-  ctx.signal_stride =
-      ctx.per_rank_shape[ctx.gather_dim] / ctx.tile_sizes[ctx.gather_dim];
-  for (int64_t j = ctx.gather_dim + 1;
-       j < static_cast<int64_t>(ctx.per_rank_shape.size()); ++j) {
-    ctx.signal_stride *=
-        xla::CeilOfRatio(ctx.per_rank_shape[j], ctx.tile_sizes[j]);
+  // Each program handles one peer's tile. `signal_stride` is the max number of
+  // programs per peer across operands; the barrier and `EmitPeerReplicaId` step
+  // by `signal_stride` to find which peer a program belongs to.
+  int64_t signal_stride = 0;
+  for (const AllGatherOperandContext& operand : ctx.operands) {
+    signal_stride = std::max(signal_stride, ProgramsPerRank(operand));
   }
+  ctx.signal_stride = signal_stride;
   return ctx;
 }
 
@@ -1584,38 +1663,68 @@ mlir::Value EmitRemoteBufferMemref(mlir::ImplicitLocOpBuilder& b,
   return mtx::PtrToMemrefOp::create(b, memref_type, ptr);
 }
 
+// Copies this program's slice of `operand`'s local input into this rank's
+// scratch buffer. If `split_dim != -1`, the tile is split across `world_size`
+// ranks so each peer copies `1 / world_size` of it; otherwise (`split_dim ==
+// -1`) the whole tile is copied.
 void CopySubtileToScratch(mlir::ImplicitLocOpBuilder& b,
-                          AllGatherRewriteContext& ctx,
+                          const AllGatherRewriteContext& ctx,
+                          const AllGatherOperandContext& operand,
                           mlir::Value peer_replica_id,
+                          mlir::Value& peer_replica_idx,
                           mlir::Value buffer_offset) {
+  // Op handles are value types; MLIR accessors are not const-qualified.
+  xtile::ExtractTileOp input_extract = operand.input_extract;
   llvm::SmallVector<mlir::Value> stage_offsets(
-      ctx.input_extract.getOffsets().begin(),
-      ctx.input_extract.getOffsets().end());
-  llvm::SmallVector<int64_t> subtile_shape(ctx.tile_sizes.begin(),
-                                           ctx.tile_sizes.end());
-  if (ctx.split_dim != -1) {
-    const int64_t subtile_size = ctx.tile_sizes[ctx.split_dim] / ctx.world_size;
+      input_extract.getOffsets().begin(), input_extract.getOffsets().end());
+  llvm::SmallVector<int64_t> subtile_shape(operand.tile_sizes.begin(),
+                                           operand.tile_sizes.end());
+  if (operand.split_dim != -1) {
+    const int64_t subtile_size =
+        operand.tile_sizes[operand.split_dim] / ctx.world_size;
     mlir::Value subtile_size_val =
         arith::ConstantIndexOp::create(b, subtile_size);
-    mlir::Value peer_replica_idx =
-        arith::IndexCastOp::create(b, b.getIndexType(), peer_replica_id);
+    if (!peer_replica_idx) {
+      peer_replica_idx =
+          arith::IndexCastOp::create(b, b.getIndexType(), peer_replica_id);
+    }
     mlir::Value subtile_off =
         arith::MulIOp::create(b, peer_replica_idx, subtile_size_val);
-    stage_offsets[ctx.split_dim] =
-        arith::AddIOp::create(b, stage_offsets[ctx.split_dim], subtile_off);
-    subtile_shape[ctx.split_dim] = subtile_size;
+    stage_offsets[operand.split_dim] =
+        arith::AddIOp::create(b, stage_offsets[operand.split_dim], subtile_off);
+    subtile_shape[operand.split_dim] = subtile_size;
   }
 
   auto subtile_tensor_type = mlir::RankedTensorType::get(
-      subtile_shape, ctx.input_extract.getType().getElementType());
+      subtile_shape, input_extract.getType().getElementType());
   auto subtile = xtile::ExtractTileOp::create(
-      b, subtile_tensor_type, ctx.input_buffer, stage_offsets, subtile_shape,
-      ctx.input_extract.getStrides());
+      b, subtile_tensor_type, operand.input_buffer, stage_offsets,
+      subtile_shape, input_extract.getStrides());
 
   mlir::Value own_scratch_slot =
-      EmitRemoteBufferMemref(b, ctx, ctx.rank_arg, buffer_offset);
+      EmitRemoteBufferMemref(b, operand, ctx.rank_arg, buffer_offset);
   xtile::InsertTileOp::create(b, subtile, own_scratch_slot, stage_offsets,
-                              subtile_shape, ctx.input_extract.getStrides());
+                              subtile_shape, input_extract.getStrides());
+}
+
+// Emits the tile of `operand` read from the peer's remote buffer: a clone of
+// the local `input_extract` (and of the i1 cast, if any) that reads from
+// `pull_scratch_slot` instead.
+mlir::Value EmitPullTileFromPeer(mlir::PatternRewriter& rewriter,
+                                 const AllGatherOperandContext& operand,
+                                 mlir::Value pull_scratch_slot) {
+  xtile::ExtractTileOp input_extract = operand.input_extract;
+  mlir::IRMapping mapping;
+  mapping.map(operand.input_buffer, pull_scratch_slot);
+  mlir::Operation* peer_extract = rewriter.clone(*input_extract, mapping);
+  if (operand.input_tile.getDefiningOp() == input_extract) {
+    return peer_extract->getResult(0);
+  }
+  // Workaround(i1_to_i8_workaround): re-apply the i8 -> i1 cast.
+  mapping.map(input_extract.getResult(), peer_extract->getResult(0));
+  mlir::Operation* peer_cast =
+      rewriter.clone(*operand.input_tile.getDefiningOp(), mapping);
+  return peer_cast->getResult(0);
 }
 
 constexpr int64_t kMaxOneShotAllToAllSizeBytes = 32LL * 1024 * 1024;  // 32 MB
@@ -1865,6 +1974,13 @@ absl::StatusOr<BlockLevelFusionConfig> GetCollectiveBlockLevelFusionConfig(
         "GpuTopology must have a target config to build "
         "BlockLevelFusionConfig.");
   }
+  if (instr->shape().IsTuple()) {
+    const HloInstruction* root = fusion_instr->fused_expression_root();
+    TF_RET_CHECK(root->opcode() == HloOpcode::kTuple &&
+                 root->operand_count() == instr->shape().tuple_shapes().size())
+        << "Expected a tuple fusion root for variadic collective: "
+        << fusion_instr->fused_instructions_computation()->ToString();
+  }
 
   if (const HloAllReduceInstruction* all_reduce =
           DynCast<HloAllReduceInstruction>(instr)) {
@@ -2042,30 +2158,53 @@ mlir::LogicalResult RewriteAllGather(mlir::stablehlo::AllGatherOp op,
   }
   AllGatherRewriteContext& ctx = *maybe_ctx;
 
+  // All local tiles (and their offsets) are defined before `op`, so the
+  // rewrite is emitted right before it.
   mlir::ImplicitLocOpBuilder builder(op.getLoc(), rewriter);
-  builder.setInsertionPoint(ctx.input_extract);
+  builder.setInsertionPoint(op);
 
   mlir::Value block_id = ttir::GetProgramIdOp::create(builder, 0);
   mlir::Value peer_replica_id = EmitPeerReplicaId(builder, ctx, block_id);
   mlir::Value signal_value = EmitDeviceInvocationCount(
       builder, ctx.signal_buffers_arg, ctx.rank_arg, block_id, ctx.world_size,
       /*signal_slot=*/peer_replica_id, ctx.signal_stride);
-  mlir::Value buffer_offset = EmitDoubleBufferOffset(
-      builder, signal_value, Product(ctx.per_rank_shape), ctx.element_type);
 
-  // Copy this program's subtile of the local tile to this rank's remote buffer.
-  CopySubtileToScratch(builder, ctx, peer_replica_id, buffer_offset);
-  // Wait until the peer has staged the whole tile.
+  // Copy this program's subtile of every local tile to this rank's remote
+  // buffers.
+  mlir::Value buffer_index = EmitDoubleBufferIndex(builder, signal_value);
+  mlir::Value peer_replica_idx;
+  llvm::SmallVector<mlir::Value, 2> buffer_offsets;
+  buffer_offsets.reserve(ctx.operands.size());
+  for (const AllGatherOperandContext& operand : ctx.operands) {
+    buffer_offsets.push_back(EmitDoubleBufferOffsetFromIndex(
+        builder, buffer_index, Product(operand.per_rank_shape),
+        operand.element_type));
+    CopySubtileToScratch(builder, ctx, operand, peer_replica_id,
+                         peer_replica_idx, buffer_offsets.back());
+  }
+  // Wait until the peer has staged the whole tile of every operand. A single
+  // barrier covers all operands because every program stages all of them.
   EmitBlockBarrier(builder, ctx.signal_buffers_arg, ctx.rank_arg, signal_value,
                    ctx.world_size, /*signal_slot=*/peer_replica_id,
                    ctx.signal_stride, mtx::BarrierMode::kConsumerSymmetric);
-  // Pull the tile from the peer's remote buffer.
-  mlir::Value pull_scratch_slot =
-      EmitRemoteBufferMemref(builder, ctx, peer_replica_id, buffer_offset);
-  rewriter.modifyOpInPlace(ctx.input_extract, [&]() {
-    ctx.input_extract.getSourceMutable().assign(pull_scratch_slot);
-  });
-  rewriter.replaceOp(op, op.getOperand(0));
+  // Pull the tiles from the peer's remote buffers.
+  llvm::SmallVector<mlir::Value, 2> peer_tiles;
+  peer_tiles.reserve(ctx.operands.size());
+  for (const auto& [operand, buffer_offset] :
+       llvm::zip(ctx.operands, buffer_offsets)) {
+    mlir::Value pull_scratch_slot = EmitRemoteBufferMemref(
+        builder, operand, peer_replica_id, buffer_offset);
+    peer_tiles.push_back(
+        EmitPullTileFromPeer(rewriter, operand, pull_scratch_slot));
+  }
+  rewriter.replaceOp(op, peer_tiles);
+  // The local tiles were only used by `op`.
+  for (const AllGatherOperandContext& operand : ctx.operands) {
+    if (operand.input_tile.getDefiningOp() != operand.input_extract) {
+      rewriter.eraseOp(operand.input_tile.getDefiningOp());
+    }
+    rewriter.eraseOp(operand.input_extract);
+  }
 
   return mlir::success();
 }
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter_test.cc
index fef68f8dbec..1519179863d 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/collective_emitter_test.cc
@@ -161,6 +161,35 @@ class CollectiveBlockLevelConfigTest : public HloHardwareIndependentTestBase {
     return ModuleWithFusion{std::move(module_with_fusion)};
   }
 
+  // Wraps the (variadic) all-gather of `module_str` into a custom fusion with
+  // a get-tuple-element + tuple root.
+  absl::StatusOr<ModuleWithFusion> BuildVariadicAllGatherModule(
+      absl::string_view module_str, int64_t replica_count) const {
+    ABSL_ASSIGN_OR_RETURN(std::unique_ptr<HloModule> module,
+                     ParseAndReturnVerifiedModule(module_str, replica_count,
+                                                  /*num_partitions=*/1));
+    const HloInstruction* all_gather = hlo_query::GetFirstInstructionWithOpcode(
+        *module->entry_computation(), HloOpcode::kAllGather);
+    TF_RET_CHECK(all_gather != nullptr);
+    std::unique_ptr<HloModule> fused_module =
+        NewModuleWithFusion(all_gather, HloInstruction::FusionKind::kCustom);
+    HloComputation* fused_comp = fused_module->entry_computation()
+                                     ->root_instruction()
+                                     ->fused_instructions_computation();
+    HloInstruction* fused_ag = fused_comp->root_instruction();
+    if (fused_ag->shape().IsTuple()) {
+      std::vector<HloInstruction*> gtes;
+      gtes.reserve(fused_ag->shape().tuple_shapes().size());
+      for (int64_t i = 0; i < fused_ag->shape().tuple_shapes().size(); ++i) {
+        gtes.push_back(fused_comp->AddInstruction(
+            HloInstruction::CreateGetTupleElement(fused_ag, i)));
+      }
+      fused_comp->set_root_instruction(
+          fused_comp->AddInstruction(HloInstruction::CreateTuple(gtes)));
+    }
+    return ModuleWithFusion{std::move(fused_module)};
+  }
+
  protected:
   static std::string GetModuleStr(const Shape& shape,
                                   absl::string_view replica_groups = "{0,1}",
@@ -874,6 +903,99 @@ TEST_F(CollectiveEmitterTest,
           .Await());
 }
 
+constexpr absl::string_view kVariadicAllGatherTwoDevicesHloStr = R"(
+  HloModule test
+  ENTRY test_computation {
+    param_0 = f32[4,4096] parameter(0)
+    param_1 = pred[2] parameter(1)
+    ROOT ag = (f32[8,4096], pred[4]) all-gather(param_0, param_1),
+      dimensions={0}, replica_groups={{0,1}},
+      backend_config={"collective_backend_config":{"kernel_strategy":"KERNEL_STRATEGY_TRITON_ONE_SHOT"}}
+  }
+)";
+
+// Every output is tiled independently for `num_blocks / num_devices` programs
+// per peer: 32768 output elements -> 16 blocks -> 8 blocks per peer. The f32
+// output gets 8 tiles per peer, the pred output only 2; the schedule pads the
+// latter so the grid is 2 * 8 programs.
+TEST_F(CollectiveBlockLevelConfigTest,
+       VariadicAllGatherBlockLevelConfigTilesOutputsIndependently) {
+  ASSERT_OK_AND_ASSIGN(
+      ModuleWithFusion module_with_fusion,
+      BuildVariadicAllGatherModule(kVariadicAllGatherTwoDevicesHloStr,
+                                   /*replica_count=*/2));
+  HloFusionInstruction* fusion = module_with_fusion.MutableFusionInstr();
+  ASSERT_OK(TrySetGpuBackendConfigForCollective(*gpu_topology_, fusion));
+  ASSERT_OK_AND_ASSIGN(
+      BlockLevelFusionConfig block_level_config,
+      GetCollectiveBlockLevelFusionConfig(*gpu_topology_, fusion));
+  ASSERT_EQ(block_level_config.output_tiles_size(), 2);
+  EXPECT_THAT(block_level_config.output_tiles(0).sizes(), ElementsAre(1, 2048));
+  EXPECT_THAT(block_level_config.output_tiles(1).sizes(), ElementsAre(1));
+
+  HloFusionAnalysis analysis = HloFusionAnalysis::Create(*fusion, device_info_);
+  std::optional<TritonFusion::LaunchConfig> launch_config =
+      TritonFusion::GetLaunchConfig(&analysis);
+  ASSERT_TRUE(launch_config.has_value());
+  EXPECT_EQ(launch_config->launch_dimensions.num_blocks(), 16);
+}
+
+// Shape from `all-gather.170` on 16 ranks.
+// Total output is ~10.5 MB -> 64 blocks (4 blocks per peer).
+// Each output is tiled independently:
+// - Output 0 (`f32[32,8,10240]`, per-peer `[2,8,10240]`): 4 tiles of
+//   `[1,4,16384]`.
+// - Output 3 (`pred[32]`, per-peer `[2]`): 2 tiles of `[1]` (padded to 4 by
+//   the schedule).
+TEST_F(CollectiveBlockLevelConfigTest,
+       VariadicAllGatherBlockLevelConfigSmallOutputsDoNotLimitGrid) {
+  constexpr absl::string_view kHloStr = R"(
+    HloModule test
+    ENTRY test_computation {
+      param_0 = f32[2,8,10240] parameter(0)
+      param_1 = pred[2,128] parameter(1)
+      param_2 = s32[2,256] parameter(2)
+      param_3 = pred[2] parameter(3)
+      param_4 = pred[2] parameter(4)
+      param_5 = s32[8,2] parameter(5)
+      ROOT ag = (f32[32,8,10240], pred[32,128], s32[32,256], pred[32], pred[32],
+                 s32[128,2]) all-gather(param_0, param_1, param_2, param_3,
+                 param_4, param_5), dimensions={0},
+        replica_groups={{0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15}},
+        backend_config={"collective_backend_config":{"kernel_strategy":"KERNEL_STRATEGY_TRITON_ONE_SHOT"}}
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(
+      ModuleWithFusion module_with_fusion,
+      BuildVariadicAllGatherModule(kHloStr, /*replica_count=*/16));
+  HloFusionInstruction* fusion = module_with_fusion.MutableFusionInstr();
+  ASSERT_OK_AND_ASSIGN(GpuTopology topology_16,
+                       GetGpuTopologyForPlatform("nvidia_h100", 1, 1, 16));
+  ASSERT_OK(TrySetGpuBackendConfigForCollective(topology_16, fusion));
+  ASSERT_OK_AND_ASSIGN(
+      BlockLevelFusionConfig block_level_config,
+      GetCollectiveBlockLevelFusionConfig(topology_16, fusion));
+  ASSERT_EQ(block_level_config.output_tiles_size(), 6);
+  // 4 tiles per peer: 128 elements per thread with 512 threads, the same as
+  // the single-operand all-gather of f32[2,8,10240].
+  EXPECT_THAT(block_level_config.output_tiles(0).sizes(),
+              ElementsAre(1, 4, 16384));
+  EXPECT_THAT(block_level_config.output_tiles(1).sizes(), ElementsAre(1, 64));
+  EXPECT_THAT(block_level_config.output_tiles(2).sizes(), ElementsAre(1, 128));
+  // 2 tiles per peer each; padded to 4 by the schedule.
+  EXPECT_THAT(block_level_config.output_tiles(3).sizes(), ElementsAre(1));
+  EXPECT_THAT(block_level_config.output_tiles(4).sizes(), ElementsAre(1));
+  // 4 tiles per peer.
+  EXPECT_THAT(block_level_config.output_tiles(5).sizes(), ElementsAre(2, 2));
+
+  HloFusionAnalysis analysis = HloFusionAnalysis::Create(*fusion, device_info_);
+  std::optional<TritonFusion::LaunchConfig> launch_config =
+      TritonFusion::GetLaunchConfig(&analysis);
+  ASSERT_TRUE(launch_config.has_value());
+  EXPECT_EQ(launch_config->launch_dimensions.num_blocks(), 64);
+  EXPECT_EQ(launch_config->launch_dimensions.num_threads_per_block(), 512);
+}
+
 }  // namespace
 
 }  // namespace xla::gpu
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_gather.cc b/third_party/xla/xla/backends/gpu/runtime/all_gather.cc
index c7335a890c8..01670df5a39 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_gather.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/all_gather.cc
@@ -44,13 +44,13 @@ limitations under the License.
 
 namespace xla::gpu {
 
-// A broader check for all-gather kernel support that verifies device, operand
-// count, replica group, gather dimension, and element-type constraints.
+// A broader check for all-gather kernel support that verifies device, replica
+// group, gather dimension, and element-type constraints.
 // Returns absl::OkStatus() if supported, or an absl::UnimplementedError
 // explaining why not.
 absl::Status IsAllGatherKernelSupported(
-    const se::DeviceDescription& device_info, int32_t num_operands,
-    int64_t num_devices, int64_t num_elements, int64_t per_rank_gather_dim_size,
+    const se::DeviceDescription& device_info, int64_t num_devices,
+    int64_t num_elements, int64_t per_rank_gather_dim_size,
     PrimitiveType element_type, bool is_local,
     const std::vector<ReplicaGroup>& replica_groups) {
   // Check if the device supports Triton collective codegen:
@@ -72,13 +72,6 @@ absl::Status IsAllGatherKernelSupported(
         absl::StrCat("Collective kernels are not supported on ",
                      rocm_cc->gfx_version(), "."));
   }
-  // TODO(b/383125489): Support variadic arguments.
-  if (num_operands != 1) {
-    return absl::UnimplementedError(absl::StrFormat(
-        "Collective kernel is not supported for number of operands not equal "
-        "to 1. Got %d.",
-        num_operands));
-  }
   if (replica_groups.empty()) {
     return absl::UnimplementedError(
         "Replica groups must be explicitly provided for collective kernels.");
@@ -136,7 +129,6 @@ absl::StatusOr<AllGatherInfo> BuildAllGatherInfo(
   }
   const int64_t num_devices =
       all_gather->device_list()->num_devices_per_group();
-  const int32_t num_operands = all_gather->operand_count();
   if (device_info.device_interconnect_info().active_links <= 0) {
     return absl::UnimplementedError(
         "Collective kernels are only supported on devices with NVLink/UALink "
@@ -152,8 +144,8 @@ absl::StatusOr<AllGatherInfo> BuildAllGatherInfo(
     const int64_t per_rank_gather_dim_size =
         operand->shape().dimensions(all_gather->all_gather_dimension());
     ABSL_RETURN_IF_ERROR(IsAllGatherKernelSupported(
-        device_info, num_operands, num_devices, operand_elements,
-        per_rank_gather_dim_size, operand->shape().element_type(), is_local,
+        device_info, num_devices, operand_elements, per_rank_gather_dim_size,
+        operand->shape().element_type(), is_local,
         all_gather->replica_groups()));
     num_elements = std::max(num_elements, operand_elements);
   }
diff --git a/third_party/xla/xla/backends/gpu/runtime/all_gather_build_info_test.cc b/third_party/xla/xla/backends/gpu/runtime/all_gather_build_info_test.cc
index ecc3654f72a..2bfccf96234 100644
--- a/third_party/xla/xla/backends/gpu/runtime/all_gather_build_info_test.cc
+++ b/third_party/xla/xla/backends/gpu/runtime/all_gather_build_info_test.cc
@@ -249,6 +249,45 @@ TEST_F(BuildAllGatherInfoTest, SucceedsFor16MBOutput) {
                          Field(&AllGatherInfo::element_type, F32))));
 }
 
+TEST_F(BuildAllGatherInfoTest, VariadicChecksSizePerOperandAndUsesMaxElements) {
+  // Each operand's output is <= 16 MB (16 MB and 8 MB), while the combined
+  // output is 24 MB. Since each operand tile is processed sequentially, the
+  // byte limit applies per operand and `num_elements` is the maximum across
+  // operands.
+  constexpr absl::string_view kModuleStr = R"(
+  HloModule test
+  ENTRY test_computation {
+    param_0 = f32[2097152] parameter(0)
+    param_1 = f32[1048576] parameter(1)
+    ROOT all-gather = (f32[4194304], f32[2097152]) all-gather(param_0, param_1),
+        dimensions={0}, replica_groups={{0,1}}
+  }
+  )";
+  se::DeviceDescription device_info = TestGpuDeviceInfo::H100SXMDeviceInfo();
+  stream_executor::GpuTargetConfigProto target_config_proto;
+  *target_config_proto.mutable_gpu_device_info() = device_info.ToProto();
+  target_config_proto.mutable_gpu_device_info()
+      ->mutable_device_interconnect_info()
+      ->set_active_links(18);
+  target_config_proto.set_platform_name("CUDA");
+  ASSERT_OK_AND_ASSIGN(gpu::GpuTargetConfig target_config,
+                       gpu::GpuTargetConfig::FromProto(target_config_proto));
+  GpuTopology gpu_topology("platform_version", /*num_partitions=*/1,
+                           /*num_hosts_per_partition=*/1,
+                           /*num_devices_per_host=*/2, target_config);
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                       ParseAndReturnVerifiedModule(kModuleStr, 2));
+  const HloInstruction* hlo_instr =
+      HloHardwareIndependentTestBase::FindInstruction(module.get(),
+                                                      HloOpcode::kAllGather);
+  EXPECT_THAT(
+      BuildAllGatherInfo(gpu_topology, Cast<HloAllGatherInstruction>(hlo_instr),
+                         /*device_assignment=*/nullptr),
+      IsOkAndHolds(AllOf(Field(&AllGatherInfo::num_devices, 2),
+                         Field(&AllGatherInfo::num_elements, 2 * 1024 * 1024),
+                         Field(&AllGatherInfo::element_type, F32))));
+}
+
 TEST_F(BuildAllGatherInfoTest,
        CreateAllGatherKernelSpecMatchesAllReduceArgumentLayout) {
   constexpr absl::string_view kModuleStr = R"(
diff --git a/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_test.cc b/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_test.cc
index 60431fb8e9c..479e37bc448 100644
--- a/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_test.cc
+++ b/third_party/xla/xla/backends/gpu/tests/all_gather_e2e_test.cc
@@ -22,6 +22,7 @@ limitations under the License.
 #include <gtest/gtest.h>
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
+#include "absl/strings/str_join.h"
 #include "absl/strings/string_view.h"
 #include "xla/backends/gpu/runtime/all_gather.h"
 #include "xla/backends/gpu/tests/collective_ops_e2e_test_base.h"
@@ -33,6 +34,7 @@ limitations under the License.
 #include "xla/literal_util.h"
 #include "xla/primitive_util.h"
 #include "xla/service/gpu/backend_configs.pb.h"
+#include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/tests/literal_test_util.h"
 #include "xla/tsl/platform/test.h"
@@ -820,5 +822,175 @@ INSTANTIATE_TEST_SUITE_P(
           primitive_util::LowercasePrimitiveTypeName(info.param));
     });
 
+struct VariadicAllGatherTestCase {
+  std::string name;
+  // Per-rank operand shapes.
+  std::vector<Shape> operand_shapes;
+  int64_t gather_dim;
+};
+
+class VariadicAllGatherTest
+    : public AllGatherTest,
+      public ::testing::WithParamInterface<VariadicAllGatherTestCase> {};
+
+TEST_P(VariadicAllGatherTest, TwoGpus) {
+  constexpr int32_t kNumReplicas = 2;
+  if (!CheckDeviceCount(kNumReplicas)) {
+    return;
+  }
+
+  const VariadicAllGatherTestCase& tc = GetParam();
+  const int64_t num_operands = tc.operand_shapes.size();
+
+  std::vector<Shape> output_shapes;
+  std::vector<std::string> param_lines;
+  std::vector<std::string> operand_names;
+  for (int64_t k = 0; k < num_operands; ++k) {
+    const Shape& in = tc.operand_shapes[k];
+    Shape out = in;
+    out.set_dimensions(tc.gather_dim,
+                       in.dimensions(tc.gather_dim) * kNumReplicas);
+    output_shapes.push_back(out);
+    param_lines.push_back(absl::StrFormat("param_%d = %s parameter(%d)", k,
+                                          ShapeUtil::HumanString(in), k));
+    operand_names.push_back(absl::StrFormat("param_%d", k));
+  }
+  const std::string module_str = absl::StrFormat(
+      R"(
+  HloModule test
+  ENTRY test_computation {
+    %s
+    ROOT all-gather = %s all-gather(%s), dimensions={%d},
+      replica_groups={{0,1}}
+  }
+  )",
+      absl::StrJoin(param_lines, "\n    "),
+      ShapeUtil::HumanString(ShapeUtil::MakeTupleShape(output_shapes)),
+      absl::StrJoin(operand_names, ", "), tc.gather_dim);
+  SCOPED_TRACE(module_str);
+
+  ASSERT_OK_AND_ASSIGN(auto module,
+                       ParseAndReturnVerifiedModule(module_str, kNumReplicas));
+
+  // Rank r contributes the r-th slice of each expected output along the
+  // gather dimension.
+  std::vector<Literal> expected_outputs;
+  std::vector<std::vector<Literal>> inputs(kNumReplicas);
+  for (int64_t k = 0; k < num_operands; ++k) {
+    ASSERT_OK_AND_ASSIGN(Literal expected, MakeFakeLiteral(output_shapes[k]));
+    const int64_t per_rank = tc.operand_shapes[k].dimensions(tc.gather_dim);
+    for (int r = 0; r < kNumReplicas; ++r) {
+      std::vector<int64_t> start(output_shapes[k].dimensions().size(), 0);
+      std::vector<int64_t> limit(output_shapes[k].dimensions().begin(),
+                                 output_shapes[k].dimensions().end());
+      start[tc.gather_dim] = r * per_rank;
+      limit[tc.gather_dim] = (r + 1) * per_rank;
+      inputs[r].push_back(expected.Slice(start, limit));
+    }
+    expected_outputs.push_back(std::move(expected));
+  }
+  std::vector<std::vector<Literal*>> args(kNumReplicas);
+  for (int r = 0; r < kNumReplicas; ++r) {
+    for (Literal& input : inputs[r]) {
+      args[r].push_back(&input);
+    }
+  }
+
+  ASSERT_OK_AND_ASSIGN(ExecutionResult result,
+                       ExecuteReplicated(std::move(module), args));
+
+  VerifyOneShotAllGather(result.optimized_module);
+
+  Literal expected_tuple =
+      LiteralUtil::MakeTupleOwned(std::move(expected_outputs));
+  ASSERT_EQ(result.results.size(), kNumReplicas);
+  for (int i = 0; i < kNumReplicas; ++i) {
+    EXPECT_TRUE(LiteralTestUtil::Equal(expected_tuple, result.results[i]))
+        << "Mismatch at replica " << i;
+  }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    VariadicAllGather, VariadicAllGatherTest,
+    ::testing::ValuesIn<VariadicAllGatherTestCase>({
+        {"TwoF32SameShape",
+         {ShapeUtil::MakeShape(F32, {256}), ShapeUtil::MakeShape(F32, {256})},
+         /*gather_dim=*/0},
+        {"TwoF32DifferentSizes",
+         {ShapeUtil::MakeShape(F32, {128}), ShapeUtil::MakeShape(F32, {64})},
+         /*gather_dim=*/0},
+        {"MixedTypesAndRanks",
+         {ShapeUtil::MakeShape(F32, {16, 32}), ShapeUtil::MakeShape(S32, {8})},
+         /*gather_dim=*/0},
+        {"ThreeOperands",
+         {ShapeUtil::MakeShape(F32, {64}), ShapeUtil::MakeShape(BF16, {128}),
+          ShapeUtil::MakeShape(S32, {32})},
+         /*gather_dim=*/0},
+        {"GatherDim1",
+         {ShapeUtil::MakeShape(F32, {4, 8}),
+          ShapeUtil::MakeShape(F32, {4, 16})},
+         /*gather_dim=*/1},
+        {"NonPowerOfTwoNonGatherDim",
+         {ShapeUtil::MakeShape(S32, {2, 7}), ShapeUtil::MakeShape(F32, {4, 5})},
+         /*gather_dim=*/0},
+        {"NonPowerOfTwoNonGatherDimLarge",
+         {ShapeUtil::MakeShape(S32, {2, 529}),
+          ShapeUtil::MakeShape(BF16, {1, 300})},
+         /*gather_dim=*/0},
+        {"Large",
+         {ShapeUtil::MakeShape(F32, {4096}),
+          ShapeUtil::MakeShape(BF16, {8192})},
+         /*gather_dim=*/0},
+        // blocks_per_peer > 1 with a different number of tiles per peer on
+        // each root (4 and 2) and a private non-gather dimension.
+        {"MixedTilesPerPeer",
+         {ShapeUtil::MakeShape(F32, {4, 2048}),
+          ShapeUtil::MakeShape(S32, {2, 1024})},
+         /*gather_dim=*/0},
+        // Regression test for schedule padding on variadic all-gather:
+        // - Output 0 (`f32[8,4096]` from `f32[4,4096]`) uses 8 blocks per peer.
+        // - Output 1 (`pred[4]` from `pred[2]`) only needs 2 blocks per peer,
+        //   so the remaining blocks per peer are padded out of bounds and must
+        //   be masked out.
+        {"PaddedSmallOutput",
+         {ShapeUtil::MakeShape(F32, {4, 4096}),
+          ShapeUtil::MakeShape(PRED, {2})},
+         /*gather_dim=*/0},
+        // Same with 4 tiles per peer on the small output and a non-power-of-two
+        // minor dimension, so masking covers both padding and partial tiles.
+        {"PaddedSmallOutputNonPowerOfTwo",
+         {ShapeUtil::MakeShape(F32, {4, 4096}),
+          ShapeUtil::MakeShape(S32, {2, 3})},
+         /*gather_dim=*/0},
+        // Real variadic all-gathers of the jit_generate decode step (operands
+        // grouped by source op).
+        {"JitGenerateStart2",
+         {ShapeUtil::MakeShape(F32, {2, 4096}), ShapeUtil::MakeShape(S32, {2})},
+         /*gather_dim=*/0},
+        {"JitGenerateStart3",
+         {ShapeUtil::MakeShape(S32, {2}), ShapeUtil::MakeShape(PRED, {2, 128}),
+          ShapeUtil::MakeShape(S32, {2, 256}), ShapeUtil::MakeShape(PRED, {2})},
+         /*gather_dim=*/0},
+        {"JitGenerateStart4",
+         {ShapeUtil::MakeShape(S32, {64, 32}),
+          ShapeUtil::MakeShape(F32, {64, 32}),
+          ShapeUtil::MakeShape(U32, {2, 2}), ShapeUtil::MakeShape(S32, {2, 4}),
+          ShapeUtil::MakeShape(U32, {2}),
+          ShapeUtil::MakeShape(U32, {2, 512, 2}),
+          ShapeUtil::MakeShape(S32, {2})},
+         /*gather_dim=*/0},
+        {"JitGenerateStart25",
+         {ShapeUtil::MakeShape(S32, {64, 8, 32}),
+          ShapeUtil::MakeShape(F32, {64, 8, 32})},
+         /*gather_dim=*/0},
+        {"JitGenerateStart28",
+         {ShapeUtil::MakeShape(S32, {2}), ShapeUtil::MakeShape(S32, {2, 8}),
+          ShapeUtil::MakeShape(S32, {2, 529})},
+         /*gather_dim=*/0},
+    }),
+    [](const ::testing::TestParamInfo<VariadicAllGatherTestCase>& info) {
+      return info.param.name;
+    });
+
 }  // namespace
 }  // namespace xla
diff --git a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion_test.cc b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion_test.cc
index 8b0f22c873d..ead6de7d25f 100644
--- a/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/collectives/collective_fusion_test.cc
@@ -359,5 +359,47 @@ TEST_F(CollectiveFusionTest, FusesAndFlattensReduceScatterOneShot) {
                             /*after_pass_checks=*/nullptr, &config);
 }
 
+TEST_F(CollectiveFusionTest, FusesVariadicAllGatherIntoGteTupleRoot) {
+  static constexpr absl::string_view kVariadicAllGatherHlo = R"(
+    HloModule variadic_all_gather_test
+
+    ENTRY e {
+      p0 = f32[4,4096]{1,0} parameter(0)
+      p1 = pred[2]{0} parameter(1)
+      ROOT ag = (f32[8,4096]{1,0}, pred[4]{0}) all-gather(p0, p1),
+          replica_groups={{0,1}},
+          dimensions={0}
+    }
+  )";
+  static constexpr absl::string_view kExpected = R"(
+    // CHECK: %[[FUSION_COMPUTATION:.*]] (param_0: f32[4,4096], param_1: pred[2]) -> (f32[8,4096], pred[4]) {
+    // CHECK:   %[[P0:.*]] = f32[4,4096]{1,0} parameter(0)
+    // CHECK:   %[[P1:.*]] = pred[2]{0} parameter(1)
+    // CHECK:   %[[AG:.*]] = (f32[8,4096]{1,0}, pred[4]{0}) all-gather(%[[P0]], %[[P1]]),
+    // CHECK-SAME: dimensions={0}
+    // CHECK:   %[[GTE0:.*]] = f32[8,4096]{1,0} get-tuple-element(%[[AG]]), index=0
+    // CHECK:   %[[GTE1:.*]] = pred[4]{0} get-tuple-element(%[[AG]]), index=1
+    // CHECK:   ROOT {{.*}} = (f32[8,4096]{1,0}, pred[4]{0}) tuple(%[[GTE0]], %[[GTE1]]), frontend_attributes={disable_tuple_simplifier="true"}
+    // CHECK: }
+    //
+    // CHECK: ENTRY %e (p0: f32[4,4096], p1: pred[2]) -> (f32[8,4096], pred[4]) {
+    // CHECK:   %[[P0_ENTRY:.*]] = f32[4,4096]{1,0} parameter(0)
+    // CHECK:   %[[P1_ENTRY:.*]] = pred[2]{0} parameter(1)
+    // CHECK:   ROOT {{.*}} = (f32[8,4096]{1,0}, pred[4]{0}) fusion(%[[P0_ENTRY]], %[[P1_ENTRY]]),
+    // CHECK-SAME: kind=kCustom,
+    // CHECK-SAME: calls=%[[FUSION_COMPUTATION]],
+    // CHECK-SAME: backend_config={
+    // CHECK-SAME: "kind":"__triton_collective"
+    // CHECK-SAME: "block_level_fusion_config"
+  )";
+  HloModuleConfig config = GetModuleConfigForTest(/*replica_count=*/2);
+  config.mutable_debug_options()
+      .add_xla_gpu_experimental_use_collective_kernels(
+          DebugOptions::COLLECTIVE_KERNEL_ALL_GATHER);
+  RunAndFilecheckHloRewrite(kVariadicAllGatherHlo,
+                            AnnotateAndFusePipeline(*gpu_topology_), kExpected,
+                            /*after_pass_checks=*/nullptr, &config);
+}
+
 }  // namespace
 }  // namespace xla::gpu

From 787cb8ac51ab246b84025bbb8e49a2c4a0094736 Mon Sep 17 00:00:00 2001
From: Tori Baker <vwbaker@google.com>
Date: Fri, 16 Oct 2026 14:03:27 -0700
Subject: [PATCH 0982/1000] Reland cl/999678120 with fixes for the ROCm and JAX
 failures.

- ROCm: allow F8E8M0FNU data-movement ops in Triton on ROCm, so GemmFusionV2 fuses scale broadcasts there too.
- A100: reject F8E4M3FN scaled dots in Triton below sm89, where Triton cannot lower F8E4M3FN conversions and crashed. ScaledDotRewriter decomposes them instead.

Reverts 4cda1bb6103d48107da783f79c2880edc57df5f3

PiperOrigin-RevId: 1001084735
---
 .../backends/gpu/codegen/triton/support.cc    |  13 +-
 .../gpu/codegen/triton/support_test.cc        |  35 ++++
 .../scaled_dot_online_quantization.hlo        |   4 +-
 ...scaled_dot_online_quantization_bitcast.hlo |   4 +-
 .../triton/tests/scaled_dot_device_test.cc    |   4 +-
 .../backends/gpu/transforms/gemm_fusion.cc    | 160 ++++++++----------
 .../gpu/transforms/gemm_fusion_test.cc        |  34 ++++
 7 files changed, 159 insertions(+), 95 deletions(-)

diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/support.cc b/third_party/xla/xla/backends/gpu/codegen/triton/support.cc
index 2243ca08e5b..40a95376490 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/support.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/support.cc
@@ -107,7 +107,7 @@ CodegenDecision IsTritonSupportedDataTypeForOpcode(
     PrimitiveType type, HloOpcode opcode,
     const se::GpuComputeCapability& gpu_version) {
   // F8E8M0FNU is lowered as i8, so it supports ops that don't change values.
-  if (type == F8E8M0FNU && gpu_version.IsCuda()) {
+  if (type == F8E8M0FNU && (gpu_version.IsCuda() || gpu_version.IsRocm())) {
     switch (opcode) {
       case HloOpcode::kBitcast:
       case HloOpcode::kBroadcast:
@@ -772,6 +772,17 @@ CodegenDecision IsTritonSupportedScaledDot(
   }
   PrimitiveType lhs_scale_type = dot.operand(2)->shape().element_type();
   PrimitiveType rhs_scale_type = dot.operand(3)->shape().element_type();
+  // Lowering tt.dot_scaled with F8E4M3FN operands or scales needs F8E4M3FN
+  // conversions, which Triton only supports from Ada onwards.
+  if (const auto* cc = gpu_version.cuda_compute_capability();
+      cc && !cc->IsAtLeastAda() &&
+      absl::c_linear_search(
+          std::vector<PrimitiveType>{lhs_type, rhs_type, lhs_scale_type,
+                                     rhs_scale_type},
+          F8E4M3FN)) {
+    return CodegenDecision::Forbid(
+        "Scaled dot with F8E4M3FN is not supported by Triton before Ada.");
+  }
   std::vector<PrimitiveType> supported_scale_types = {F8E4M3FN, F8E5M2,
                                                       F8E8M0FNU, S8};
   // tt.dot_scaled only dequantizes some operand types; a scale on any other
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc
index a874d18b8da..221d4da393a 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/support_test.cc
@@ -2532,6 +2532,41 @@ ENTRY triton_computation {
       *dot, se::GpuComputeCapability(se::CudaComputeCapability::Volta())));
 }
 
+// Triton cannot lower F8E4M3FN conversions before sm89.
+class ScaledDotF8E4M3FNTest
+    : public SupportTest,
+      public ::testing::WithParamInterface<se::GpuComputeCapability> {};
+
+TEST_P(ScaledDotF8E4M3FNTest, MixedWithBF16) {
+  const std::string kHloTestTemplate = R"(
+HloModule ScaledDotF8E4M3FNMixedWithBF16
+
+ENTRY triton_computation {
+  lhs = f8e4m3fn[16, 32] parameter(0)
+  lhs_scale = f8e8m0fnu[16, 1] parameter(1)
+  rhs = $0[32, 16] parameter(2)
+  rhs_scale = f8e8m0fnu[] constant(1)
+  rhs_scale_broadcast = f8e8m0fnu[1, 16] broadcast(rhs_scale), dimensions={}
+  ROOT dot = f32[16, 16] scaled-dot(lhs, rhs, lhs_scale, rhs_scale_broadcast),
+      lhs_contracting_dims={1},
+      rhs_contracting_dims={0},
+      backend_config={sizes:[16]}
+}
+)";
+  ASSERT_OK_AND_ASSIGN(TestedInstruction ti,
+                       ParseTemplateAndGetInstruction(kHloTestTemplate, BF16,
+                                                      HloOpcode::kScaledDot));
+  RunSupportTest(std::move(ti), /*output_tile_sizes=*/{16, 16}, GetParam(),
+                 ExpectedFailMode::kFailOrCrash);
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    ScaledDotF8E4M3FNTestSuite, ScaledDotF8E4M3FNTest,
+    ::testing::Values(
+        se::GpuComputeCapability(se::CudaComputeCapability::Ampere()),
+        se::GpuComputeCapability(se::CudaComputeCapability(8, 9)),
+        se::GpuComputeCapability(se::CudaComputeCapability::Hopper())));
+
 TEST_F(SupportTest, NestedFusionsAreRejected) {
   // Nested fusions are not supported by xtile emitter.
   absl::string_view hlo_text = R"(
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization.hlo
index 5619eee8d23..f0f2196d3ed 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization.hlo
@@ -46,12 +46,12 @@ ENTRY %main.4 (a.1: f8e4m3fn[2,256], b.1: f8e4m3fn[256,256], a_scale.1: f8e8m0fn
   ROOT %tuple.3 = (f8e3m4[2,256]{1,0}, f8e8m0fnu[2,8]{1,0}) tuple(%jit_online_quantize_.4, %jit_online_quantize_.5)
 }
 
-// CHECK: %fusion_scaled-dot (parameter
+// CHECK: %[[FUSION_DOT:[^ ]*scaled-dot[^ ]*]] (param
 // CHECK: %fused_reduce (param
 // CHECK: %fused_convert (param
 // CHECK: %fused_convert.1 (param
 // CHECK: ENTRY
-// CHECK: %[[fusion_dot:.*]] = f32[2,256]{1,0} fusion({{.*}}), kind=kCustom, calls=%fusion_scaled-dot
+// CHECK: %[[fusion_dot:.*]] = f32[2,256]{1,0} fusion({{.*}}), kind=kCustom, calls=%[[FUSION_DOT]]
 // CHECK: %[[fusion_reduce:.*]] = f32[2,8]{1,0} fusion(%[[fusion_dot]]), kind=kInput, calls=%fused_reduce
 // CHECK: %[[fusion_convert1:.*]] = f8e3m4[2,256]{1,0} fusion(%[[fusion_dot]], %[[fusion_reduce]]), kind=kLoop, calls=%fused_convert
 // CHECK: %[[fusion_convert2:.*]] = f8e8m0fnu[2,8]{1,0} fusion(%[[fusion_reduce]]), kind=kLoop, calls=%fused_convert.1
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization_bitcast.hlo b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization_bitcast.hlo
index 20fd0b6c248..ff21e7c54c4 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization_bitcast.hlo
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot/scaled_dot_online_quantization_bitcast.hlo
@@ -59,12 +59,12 @@ ENTRY %main.5 (a.1: f8e4m3fn[2,256], b.1: f8e4m3fn[256,256], a_scale.1: f8e8m0fn
   ROOT %tuple.0 = (f8e3m4[2,256]{1,0}, f8e8m0fnu[2,8]{1,0}) tuple(%convert_element_type.1, %convert_element_type.2)
 }
 
-// CHECK: %fusion_scaled-dot (parameter
+// CHECK: %[[FUSION_DOT:[^ ]*scaled-dot[^ ]*]] (param
 // CHECK: %fused_reduce (param
 // CHECK: %fused_convert (param
 // CHECK: %fused_convert.1 (param
 // CHECK: ENTRY
-// CHECK: %[[fusion_dot:.*]] = f32[2,256]{1,0} fusion({{.*}}), kind=kCustom, calls=%fusion_scaled-dot
+// CHECK: %[[fusion_dot:.*]] = f32[2,256]{1,0} fusion({{.*}}), kind=kCustom, calls=%[[FUSION_DOT]]
 // CHECK: %[[fusion_reduce:.*]] = f32[2,8]{1,0} fusion(%[[fusion_dot]]), kind=kInput, calls=%fused_reduce
 // CHECK: %[[fusion_convert1:.*]] = f8e3m4[2,256]{1,0} fusion(%[[fusion_dot]], %[[fusion_reduce]]), kind=kLoop, calls=%fused_convert
 // CHECK: %[[fusion_convert2:.*]] = f8e8m0fnu[2,8]{1,0} fusion(%[[fusion_reduce]]), kind=kLoop, calls=%fused_convert.1
diff --git a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc
index abf8ca66f8f..ee848028d19 100644
--- a/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/triton/tests/scaled_dot_device_test.cc
@@ -706,7 +706,7 @@ ENTRY e {
   ASSERT_OK_AND_ASSIGN(auto optimized_module,
                        GetOptimizedModule(kHloTextTemplate));
   constexpr absl::string_view kExpectedOptimizedHLO = R"(
-    CHECK: %fusion
+    CHECK: %{{.*}} (param
     CHECK: %{{.*}} = f8e8m0fnu[3,128,4]{2,1,0} broadcast(%{{.*}}), dimensions={0,1}
     CHECK: %{{.*}} = f8e8m0fnu[3,128,4]{2,1,0} broadcast(%{{.*}}), dimensions={0,1}
     CHECK: ROOT {{.*}} scaled-dot
@@ -894,7 +894,7 @@ ENTRY e {
   ASSERT_OK_AND_ASSIGN(auto optimized_module,
                        GetOptimizedModule(kHloTextTemplate));
   constexpr absl::string_view kExpectedOptimizedHLO = R"(
-    CHECK: %[[fusion_name:.*]] (parameter
+    CHECK: %[[fusion_name:.*]] (param
     CHECK: %[[scaled_dot:.*]] = bf16[3,128,128]{2,1,0} scaled-dot
     CHECK: %[[global_scaler:.*]] = bf16[3,128,128]{2,1,0} broadcast
     CHECK: ROOT %{{.*}} = bf16[3,128,128]{2,1,0} multiply(%[[scaled_dot]], %[[global_scaler]])
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 7e9797180ec..10cf219cfaa 100644
--- a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc
@@ -718,8 +718,9 @@ bool AllowedInGemmFusion(const HloInstruction& instr) {
     // the moment. Includes kFusion, kReduce, kAllReduce, etc.
     return false;
   }
-  return HloPredicateIsNotOp<HloOpcode::kDot, HloOpcode::kParameter,
-                             HloOpcode::kAllGather, HloOpcode::kAllGatherStart,
+  return HloPredicateIsNotOp<HloOpcode::kDot, HloOpcode::kScaledDot,
+                             HloOpcode::kParameter, HloOpcode::kAllGather,
+                             HloOpcode::kAllGatherStart,
                              HloOpcode::kAllGatherDone>(&instr);
 }
 
@@ -1265,7 +1266,8 @@ FusionDecision CanUnpackS4ParametersInFusion(
   }
 
   if (tiled_computation->roots().size() != 1 ||
-      tiled_computation->roots()[0]->hlo()->opcode() != HloOpcode::kDot) {
+      !HloPredicateIsOp<HloOpcode::kDot, HloOpcode::kScaledDot>(
+          tiled_computation->roots()[0]->hlo())) {
     return FusionDecision::Forbid(
         "Expected tiled computation root to be a single dot instruction.");
   }
@@ -1309,7 +1311,8 @@ FusionDecision CanTile(mlir::MLIRContext& mlir_context,
   // a tile size to be divisible by 3, but all Triton tiles are powers of 2.
   std::optional<HloInstructionAdaptor> dot =
       HloBfsFindIf(fusion.GetRoots(), fusion, [](HloInstructionAdaptor node) {
-        return node.opcode() == HloOpcode::kDot;
+        return HloPredicateIsOp<HloOpcode::kDot, HloOpcode::kScaledDot>(
+              &node.instruction());
       });
   if (!dot.has_value()) {
     return FusionDecision::Forbid("No dot instruction found in the fusion.");
@@ -1655,22 +1658,26 @@ FusionDecision ShouldFuseUserTranspose(const HloInstruction& transpose,
                                        const HloInstruction& original_transpose,
                                        const HloInstruction& fusion,
                                        const ShapeTracker& tracker) {
-  const HloInstruction* dot = hlo_query::FindInstruction(
-      fusion.fused_instructions_computation(), HloOpcode::kDot);
+  const HloInstruction* dot = hlo_query::GetFirstInstructionWithOpcode(
+      *fusion.fused_instructions_computation(),
+      {HloOpcode::kDot, HloOpcode::kScaledDot});
   if (dot == nullptr) {
     return FusionDecision::Forbid("Dot not found in fusion.");
   }
 
-  absl::StatusOr<std::array<DotOperandDims, 2>> dot_dims =
-      DotOperandDims::FromDot(dot);
-  if (!dot_dims.ok()) {
+  // The output of a scaled dot has the same layout as the output of a dot of
+  // its LHS and RHS, so the dims of operands 0 and 1 are sufficient.
+  absl::StatusOr<DotOperandDims> lhs_dims =
+      DotOperandDims::FromDotOperand(dot, 0);
+  absl::StatusOr<DotOperandDims> rhs_dims =
+      DotOperandDims::FromDotOperand(dot, 1);
+  if (!lhs_dims.ok() || !rhs_dims.ok()) {
     return FusionDecision::Forbid("Failed to get dot operand dims.");
   }
-  const auto& [lhs_dims, rhs_dims] = *dot_dims;
 
-  int64_t num_batch = lhs_dims.Rank(DotOperandDims::kBatch);
-  int64_t num_lhs_nc = lhs_dims.Rank(DotOperandDims::kNonContracting);
-  int64_t num_rhs_nc = rhs_dims.Rank(DotOperandDims::kNonContracting);
+  int64_t num_batch = lhs_dims->Rank(DotOperandDims::kBatch);
+  int64_t num_lhs_nc = lhs_dims->Rank(DotOperandDims::kNonContracting);
+  int64_t num_rhs_nc = rhs_dims->Rank(DotOperandDims::kNonContracting);
 
   // Dot result always maps to [batch..., lhs_nc..., rhs_nc...].
   absl::InlinedVector<int64_t, 4> batch_dims(num_batch);
@@ -1897,18 +1904,22 @@ absl::StatusOr<std::variant<Fusion, FusionDecision>> CreateTileableFusion(
     return FusionDecision::Forbid("Cannot tile the dot instruction.");
   }
 
-  // Initialize queue with trackers for dot operands.
-  absl::StatusOr<std::array<DotOperandDims, 2>> dot_dims =
-      DotOperandDims::FromDot(dot);
-  if (!dot_dims.ok()) {
-    return FusionDecision::Forbid("Failed to get dot operand dims.");
+  // Initialize queue with trackers for the LHS and RHS. Scaled dot scales get
+  // no tracker, so transposes and concatenates are not fused into them.
+  std::queue<std::pair<HloInstruction*, std::optional<TrackerInfo>>> queue;
+  for (int64_t i = 0; i < dot->operand_count(); ++i) {
+    HloInstruction* operand = dot->mutable_operand(i);
+    std::optional<TrackerInfo> tracker;
+    if (i < 2) {
+      absl::StatusOr<DotOperandDims> dims =
+          DotOperandDims::FromDotOperand(dot, i);
+      if (!dims.ok()) {
+        return FusionDecision::Forbid("Failed to get dot operand dims.");
+      }
+      tracker = TrackerInfo{ShapeTracker(operand->shape()), i, *dims};
+    }
+    queue.push({operand, std::move(tracker)});
   }
-  std::queue<std::pair<HloInstruction*, std::optional<TrackerInfo>>> queue(
-      {{dot->mutable_operand(0),
-        TrackerInfo{ShapeTracker(dot->operand(0)->shape()), 0, (*dot_dims)[0]}},
-       {dot->mutable_operand(1),
-        TrackerInfo{ShapeTracker(dot->operand(1)->shape()), 1,
-                    (*dot_dims)[1]}}});
 
   // Start with a fusion containing only the dot instruction.
   auto entry = fusion_search_space.entry();
@@ -2000,9 +2011,13 @@ absl::StatusOr<std::variant<Fusion, FusionDecision>> CreateTileableFusion(
                 fusion_computation->Clone(std::string(name)), original_output};
 }
 
+// Creates a GEMM fusion around a dot-like instruction (kDot or kScaledDot)
+// using the GemmFusionV2 algorithm.
 absl::StatusOr<std::variant<Fusion, FusionDecision>> CreateDotFusionV2(
-    HloDotInstruction& dot, const se::DeviceDescription& device_description,
+    HloInstruction& dot, const se::DeviceDescription& device_description,
     absl::string_view name) {
+  CHECK((HloPredicateIsOp<HloOpcode::kDot, HloOpcode::kScaledDot>(&dot)))
+      << dot.ToString();
   VLOG(3) << "Creating dot fusion v2 around dot: " << dot.ToString();
   FusionSearchSpace fusion_search_space(&dot, device_description);
   VLOG(3) << "Found fusion search space: \n"
@@ -2116,8 +2131,7 @@ absl::StatusOr<std::variant<Fusion, FusionDecision>> CreateDotFusion(
 class GemmFusionVisitor : public DfsHloRewriteVisitor {
  public:
   explicit GemmFusionVisitor(const se::DeviceDescription& device_description)
-      : device_description_(device_description),
-        gpu_version_(device_description.gpu_compute_capability()) {}
+      : device_description_(device_description) {}
   // Checks that a dot() should be targeting the triton GEMM emitter;
   // if so - fuses all its compatible inputs and outputs as a new computation
   // and replaces the original dot() with a call to the computation.
@@ -2149,18 +2163,43 @@ class GemmFusionVisitor : public DfsHloRewriteVisitor {
       return absl::OkStatus();
     }
 
-    Fusion fusion = std::get<Fusion>(std::move(fusion_or_decision));
+    return InsertGemmFusion(
+        dot, std::get<Fusion>(std::move(fusion_or_decision)), fusion_name);
+  }
 
+  absl::Status HandleScaledDot(HloInstruction* scaled_dot) override {
+    CHECK_EQ(scaled_dot->opcode(), HloOpcode::kScaledDot);
+    std::string fusion_name = absl::StrCat("gemm_fusion_", scaled_dot->name());
+    ABSL_ASSIGN_OR_RETURN(
+        auto fusion_or_decision,
+        CreateDotFusionV2(*scaled_dot, device_description_,
+                          absl::StrCat(fusion_name, "_computation")));
+    if (std::holds_alternative<FusionDecision>(fusion_or_decision)) {
+      VLOG(3) << "Not fusing: "
+              << std::get<FusionDecision>(fusion_or_decision).Explain();
+      return absl::OkStatus();
+    }
+    return InsertGemmFusion(scaled_dot,
+                            std::get<Fusion>(std::move(fusion_or_decision)),
+                            fusion_name);
+  }
+
+ private:
+  // Wraps `fusion` into a fusion instruction and replaces `fusion.output` in
+  // the original computation with it. `dot_like` is the dot-like instruction
+  // the fusion was created around.
+  absl::Status InsertGemmFusion(HloInstruction* dot_like, Fusion fusion,
+                                absl::string_view fusion_name) {
     HloComputation* computation =
-        dot->GetModule()->AddComputationAndUnifyNamesAndIds(
+        dot_like->GetModule()->AddComputationAndUnifyNamesAndIds(
             std::move(fusion.computation), /*is_entry=*/false);
     HloInstruction* dot_fusion =
-        dot->parent()->AddInstruction(HloInstruction::CreateFusion(
+        dot_like->parent()->AddInstruction(HloInstruction::CreateFusion(
             computation->root_instruction()->shape(),
             HloInstruction::FusionKind::kCustom, fusion.inputs, computation));
-    // Copy the metadata of the `dot` to the newly created `fusion` op. This
+    // Copy the metadata of `dot_like` to the newly created `fusion` op. This
     // is convenient for handling metadata in split-k rewriting subsequently.
-    dot_fusion->set_metadata(dot->metadata());
+    dot_fusion->set_metadata(dot_like->metadata());
     dot_fusion->GetModule()->SetAndUniquifyInstrName(dot_fusion, fusion_name);
 
     ABSL_ASSIGN_OR_RETURN(auto gpu_config,
@@ -2172,7 +2211,7 @@ class GemmFusionVisitor : public DfsHloRewriteVisitor {
 
     HloInstruction* replacement = dot_fusion;
     if (fusion.output->shape() != dot_fusion->shape()) {
-      replacement = dot->parent()->AddInstruction(
+      replacement = dot_like->parent()->AddInstruction(
           HloInstruction::CreateBitcast(fusion.output->shape(), dot_fusion));
     }
 
@@ -2189,62 +2228,7 @@ class GemmFusionVisitor : public DfsHloRewriteVisitor {
     return absl::OkStatus();
   }
 
-  absl::Status HandleScaledDot(HloInstruction* scaled_dot) override {
-    CHECK_EQ(scaled_dot->opcode(), HloOpcode::kScaledDot);
-    HloComputation::Builder builder(
-        absl::StrCat("fusion_", scaled_dot->name()));
-
-    std::vector<HloInstruction*> fusion_inputs;
-
-    std::vector<HlosAndRequirements> hlos_and_reqs;
-    hlos_and_reqs.reserve(scaled_dot->operand_count());
-    ABSL_ASSIGN_OR_RETURN(HlosAndRequirements lhs_hlos_and_reqs,
-                     FuseDotOperand(*scaled_dot, /*operand_index=*/0,
-                                    gpu_version_, builder, fusion_inputs));
-    hlos_and_reqs.push_back(lhs_hlos_and_reqs);
-    ABSL_ASSIGN_OR_RETURN(HlosAndRequirements rhs_hlos_and_reqs,
-                     FuseDotOperand(*scaled_dot, /*operand_index=*/1,
-                                    gpu_version_, builder, fusion_inputs));
-    hlos_and_reqs.push_back(rhs_hlos_and_reqs);
-    ABSL_ASSIGN_OR_RETURN(HlosAndRequirements lhs_scale_hlos_and_reqs,
-                     FuseDotOperand(*scaled_dot, /*operand_index=*/2,
-                                    gpu_version_, builder, fusion_inputs));
-    hlos_and_reqs.push_back(lhs_scale_hlos_and_reqs);
-    ABSL_ASSIGN_OR_RETURN(HlosAndRequirements rhs_scale_hlos_and_reqs,
-                     FuseDotOperand(*scaled_dot, /*operand_index=*/3,
-                                    gpu_version_, builder, fusion_inputs));
-    hlos_and_reqs.push_back(rhs_scale_hlos_and_reqs);
-
-    HloInstruction& fused_dot = FuseDot(*scaled_dot, hlos_and_reqs, builder);
-
-    HlosAndRequirements fused_output_and_reqs =
-        FuseDotOutput(*scaled_dot, fused_dot, gpu_version_,
-                      lhs_hlos_and_reqs.requirements, builder, fusion_inputs);
-    HloComputation* computation =
-        scaled_dot->GetModule()->AddComputationAndUnifyNamesAndIds(
-            builder.Build(),
-            /*is_entry=*/false);
-
-    HloInstruction* fusion =
-        scaled_dot->parent()->AddInstruction(HloInstruction::CreateFusion(
-            computation->root_instruction()->shape(),
-            HloInstruction::FusionKind::kCustom, fusion_inputs, computation));
-
-    ABSL_ASSIGN_OR_RETURN(auto gpu_config,
-                     fusion->backend_config<GpuBackendConfig>());
-    FusionBackendConfig& backend_config =
-        *gpu_config.mutable_fusion_backend_config();
-    backend_config.set_kind(kTritonGemmFusionKind);
-    ABSL_RETURN_IF_ERROR(fusion->set_backend_config(gpu_config));
-    HloInstruction* fusion_output = fused_output_and_reqs.original_hlo;
-    ABSL_RETURN_IF_ERROR(ReplaceInstruction(fusion_output, fusion));
-    MarkAsChanged();
-    return absl::OkStatus();
-  }
-
- private:
   const se::DeviceDescription& device_description_;
-  se::GpuComputeCapability gpu_version_;
 };
 
 absl::StatusOr<bool> RunOnComputation(
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 2fd6de0ff01..e72b3e6c135 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
@@ -2052,6 +2052,40 @@ TEST_F(GemmFusionTest, ScaledDotIsFused) {
   MatchHloModule(*module, kExpectedHloText);
 }
 
+TEST_P(GemmFusionTestVersioned, ScaledDotScaleBroadcastsAreFusedOnRocm) {
+  constexpr absl::string_view kHloText = R"(
+    HloModule m
+
+    ENTRY e {
+      lhs = f8e4m3fn[3,128,128] parameter(0)
+      rhs = f8e4m3fn[3,128,128] parameter(1)
+      lhs_scale = f8e8m0fnu[3,128] parameter(2)
+      lhs_scale_broadcast = f8e8m0fnu[3,128,4] broadcast(lhs_scale),
+          dimensions={0,1}
+      rhs_scale = f8e8m0fnu[3,128] parameter(3)
+      rhs_scale_broadcast = f8e8m0fnu[3,128,4] broadcast(rhs_scale),
+          dimensions={0,1}
+      ROOT dot = bf16[3,128,128] scaled-dot(lhs, rhs, lhs_scale_broadcast,
+          rhs_scale_broadcast),
+        lhs_batch_dims={0}, rhs_batch_dims={0},
+        lhs_contracting_dims={2}, rhs_contracting_dims={2}
+    }
+  )";
+
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(kHloText));
+  EXPECT_THAT(GemmFusion(TestGpuDeviceInfo::AMDMI350DeviceInfo())
+                  .Run(module.get()),
+              IsOkAndHolds(true));
+  MatchHloModule(*module, R"(
+    CHECK: f8e8m0fnu[3,128,4]{2,1,0} broadcast(
+    CHECK: f8e8m0fnu[3,128,4]{2,1,0} broadcast(
+    CHECK: ROOT {{.*}} scaled-dot(
+    CHECK: ENTRY
+    CHECK-NOT: broadcast(
+    CHECK: ROOT {{.*}} fusion(
+  )");
+}
+
 TEST_P(GemmFusionTestVersioned, SkipNonFusibleComputations) {
   const char* hlo_text = R"(
 HloModule module

From 2f5acabb5f02600b29cbbd8d61b79d59006bfe45 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 16 Oct 2026 14:04:23 -0700
Subject: [PATCH 0983/1000] Support multi-pass profiling in OpenXLA
 FunctionalHloRunner.

Add SimpleMultiPassTracer test tracer (simple_multipass_tracer_for_test.{h,cc})
and extract profiler classes (XSpaceProfilerInterface, HLORunnerProfiler, and
multi-pass runners) into hlo_runner_profiler.{h,cc} with AllReduceOr collective
synchronization.

PiperOrigin-RevId: 1001085342
---
 .../third_party/tsl/tsl/profiler/lib/BUILD    |   3 +
 .../lib/multi_pass_profiler_controller.cc     |  12 +-
 .../lib/multi_pass_profiler_controller.h      |   3 +-
 .../tsl/profiler/lib/profiler_factory_test.cc |   5 +-
 .../tsl/tsl/profiler/lib/profiler_interface.h |   2 +-
 .../tsl/tsl/profiler/lib/profiler_passes.cc   |  11 +-
 .../tsl/tsl/profiler/lib/profiler_passes.h    |   3 +-
 .../xla/xla/tools/multihost_hlo_runner/BUILD  | 109 +++++-
 .../functional_hlo_runner.cc                  | 157 ++++-----
 .../functional_hlo_runner.h                   |  59 +---
 .../hlo_runner_profiler.cc                    | 322 +++++++++++++++++
 .../hlo_runner_profiler.h                     | 116 +++++++
 .../multipass_functional_hlo_runner_test.cc   | 327 ++++++++++++++++++
 .../multihost_hlo_runner/profiler_interface.h |  18 +
 .../simple_multipass_tracer_for_test.cc       | 167 +++++++++
 .../simple_multipass_tracer_for_test.h        |  30 ++
 16 files changed, 1185 insertions(+), 159 deletions(-)
 create mode 100644 third_party/xla/xla/tools/multihost_hlo_runner/hlo_runner_profiler.cc
 create mode 100644 third_party/xla/xla/tools/multihost_hlo_runner/hlo_runner_profiler.h
 create mode 100644 third_party/xla/xla/tools/multihost_hlo_runner/multipass_functional_hlo_runner_test.cc
 create mode 100644 third_party/xla/xla/tools/multihost_hlo_runner/simple_multipass_tracer_for_test.cc
 create mode 100644 third_party/xla/xla/tools/multihost_hlo_runner/simple_multipass_tracer_for_test.h

diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD b/third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD
index 062921a2f27..c97b6ae8dfc 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/BUILD
@@ -132,6 +132,7 @@ tsl_cc_test(
         "//tsl/profiler/protobuf:profiler_options_proto_cc",
         "//tsl/profiler/protobuf:xplane_proto_cc",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings:string_view",
         "@com_google_googletest//:gtest_main",
         "@xla//xla/tsl/platform:macros",
@@ -216,6 +217,7 @@ cc_library(
         ":profiler_interface",
         "//tsl/profiler/protobuf:xplane_proto_cc",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@xla//xla/tsl/platform:logging",
     ],
@@ -239,6 +241,7 @@ cc_library(
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
         "@xla//xla/tsl/platform:env",
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/multi_pass_profiler_controller.cc b/third_party/xla/third_party/tsl/tsl/profiler/lib/multi_pass_profiler_controller.cc
index 803a7d0a576..6a4d2b3a1be 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/multi_pass_profiler_controller.cc
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/multi_pass_profiler_controller.cc
@@ -19,6 +19,7 @@ limitations under the License.
 #include <utility>
 
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "xla/tsl/platform/logging.h"
 #include "tsl/profiler/lib/profiler_interface.h"
@@ -43,12 +44,17 @@ MultiPassProfilerController::~MultiPassProfilerController() {
   }
 }
 
-bool MultiPassProfilerController::NeedMorePasses() {
+absl::StatusOr<bool> MultiPassProfilerController::NeedMorePasses() {
   if (state_ == MultiPassProfilerState::kPassStopped ||
       state_ == MultiPassProfilerState::kStarted ||
       state_ == MultiPassProfilerState::kInit) {
-    if (!status_.ok()) return false;
-    return profiler_->NeedMorePasses();
+    if (!status_.ok()) return status_;
+    absl::StatusOr<bool> need_more = profiler_->NeedMorePasses();
+    if (!need_more.ok()) {
+      status_ = need_more.status();
+      LOG(ERROR) << status_;
+    }
+    return need_more;
   }
   return false;
 }
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/multi_pass_profiler_controller.h b/third_party/xla/third_party/tsl/tsl/profiler/lib/multi_pass_profiler_controller.h
index 9cd8a55c0fa..41a1b560daa 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/multi_pass_profiler_controller.h
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/multi_pass_profiler_controller.h
@@ -19,6 +19,7 @@ limitations under the License.
 #include <memory>
 
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "tsl/profiler/lib/profiler_interface.h"
 #include "tsl/profiler/protobuf/xplane.pb.h"
@@ -42,7 +43,7 @@ class MultiPassProfilerController : public MultiPassProfilerInterface {
   ~MultiPassProfilerController() override;
 
   // MultiPassProfilerInterface methods
-  bool NeedMorePasses() override;
+  absl::StatusOr<bool> NeedMorePasses() override;
   absl::Status StartPass() override;
   absl::Status PushRange(absl::string_view name) override;
   absl::Status PopRange() override;
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_factory_test.cc b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_factory_test.cc
index 38977cca4b7..36edd6be731 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_factory_test.cc
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_factory_test.cc
@@ -18,6 +18,7 @@ limitations under the License.
 #include <utility>
 
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "xla/tsl/platform/macros.h"
 #include "xla/tsl/platform/test.h"
@@ -102,7 +103,7 @@ TEST(ProfilerFactoryTest, FactoryClassCapturedByLambda) {
 
 class TestMultiPassProfiler : public MultiPassProfilerInterface {
  public:
-  bool NeedMorePasses() override { return false; }
+  absl::StatusOr<bool> NeedMorePasses() override { return false; }
   absl::Status StartPass() override { return absl::OkStatus(); }
   absl::Status PushRange(absl::string_view name) override {
     return absl::OkStatus();
@@ -168,7 +169,7 @@ class TrackingMultiPassProfiler : public MultiPassProfilerInterface {
   explicit TrackingMultiPassProfiler(Tracker* tracker = nullptr)
       : tracker_(tracker) {}
 
-  bool NeedMorePasses() override { return false; }
+  absl::StatusOr<bool> NeedMorePasses() override { return false; }
   absl::Status StartPass() override {
     start_pass_called_++;
     if (tracker_) tracker_->start_pass_called++;
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_interface.h b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_interface.h
index ed93bd28f64..c99fba38b6d 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_interface.h
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_interface.h
@@ -75,7 +75,7 @@ class ProfilerInterface {
 class MultiPassProfilerInterface : public ProfilerInterface {
  public:
   // Returns true if there are more passes to profile.
-  virtual bool NeedMorePasses() = 0;
+  virtual absl::StatusOr<bool> NeedMorePasses() = 0;
 
   // Starts a new profiling pass.
   virtual absl::Status StartPass() = 0;
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_passes.cc b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_passes.cc
index a341a0e2647..fb084f3a99e 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_passes.cc
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_passes.cc
@@ -20,6 +20,7 @@ limitations under the License.
 
 #include "absl/memory/memory.h"
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
 #include "xla/tsl/platform/logging.h"
@@ -67,14 +68,18 @@ absl::Status ProfilerPasses::Status() {
   return status_;
 }
 
-bool ProfilerPasses::NeedMorePasses() {
+absl::StatusOr<bool> ProfilerPasses::NeedMorePasses() {
   absl::MutexLock l(&mutex_);
-  if (!status_.ok()) return false;
+  if (!status_.ok()) return status_;
 #if !defined(IS_MOBILE_PLATFORM)
   if (multi_pass_profiler_ == nullptr) {
     return false;
   }
-  return multi_pass_profiler_->NeedMorePasses();
+  absl::StatusOr<bool> need_more = multi_pass_profiler_->NeedMorePasses();
+  if (!need_more.ok()) {
+    status_ = need_more.status();
+  }
+  return need_more;
 #else
   return false;
 #endif
diff --git a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_passes.h b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_passes.h
index f56046328c3..e9eb3363b38 100644
--- a/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_passes.h
+++ b/third_party/xla/third_party/tsl/tsl/profiler/lib/profiler_passes.h
@@ -21,6 +21,7 @@ limitations under the License.
 
 #include "absl/base/thread_annotations.h"
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
 #include "tsl/profiler/protobuf/profiler_options.pb.h"
@@ -68,7 +69,7 @@ class ProfilerPasses {
 
   absl::Status Status() ABSL_LOCKS_EXCLUDED(mutex_);
 
-  bool NeedMorePasses() ABSL_LOCKS_EXCLUDED(mutex_);
+  absl::StatusOr<bool> NeedMorePasses() ABSL_LOCKS_EXCLUDED(mutex_);
 
   absl::Status StartPass() ABSL_LOCKS_EXCLUDED(mutex_);
 
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/BUILD b/third_party/xla/xla/tools/multihost_hlo_runner/BUILD
index e9043fd62ac..f882f6bbd6b 100644
--- a/third_party/xla/xla/tools/multihost_hlo_runner/BUILD
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/BUILD
@@ -45,6 +45,7 @@ cc_library(
     hdrs = ["profiler_interface.h"],
     deps = [
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
     ],
 )
 
@@ -181,12 +182,36 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "hlo_runner_profiler",
+    srcs = ["hlo_runner_profiler.cc"],
+    hdrs = ["hlo_runner_profiler.h"],
+    deps = [
+        ":profiler_interface",
+        "//xla/pjrt/distributed:key_value_store_interface",
+        "//xla/tsl/platform:env",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_macros",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/time",
+        "@tsl//tsl/profiler/lib:profiler_factory_impl",
+        "@tsl//tsl/profiler/lib:profiler_passes",
+        "@tsl//tsl/profiler/lib:profiler_session",
+        "@tsl//tsl/profiler/lib:profiler_session_impl",
+        "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
+    ],
+)
+
 cc_library(
     name = "functional_hlo_runner",
     srcs = ["functional_hlo_runner.cc"],
     hdrs = ["functional_hlo_runner.h"],
     deps = [
         ":hlo_input_output_format",
+        ":hlo_runner_profiler",
         ":profiler_interface",
         "//xla:debug_options_flags",
         "//xla:future",
@@ -226,10 +251,9 @@ cc_library(
         "//xla/tests:test_utils",
         "//xla/tools:hlo_control_flow_flattening",
         "//xla/tsl/platform:env",
-        "//xla/tsl/platform:errors",
-        "//xla/tsl/platform:statusor",
         "//xla/tsl/util:fixed_option_set_flag",
         "@com_google_absl//absl/algorithm:container",
+        "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/functional:function_ref",
         "@com_google_absl//absl/log",
@@ -245,11 +269,6 @@ cc_library(
         "@com_google_absl//absl/types:span",
         "@llvm-project//mlir:FuncExtensions",
         "@llvm-project//mlir:IR",
-        "@tsl//tsl/profiler/lib:profiler_factory_impl",
-        "@tsl//tsl/profiler/lib:profiler_session",
-        "@tsl//tsl/profiler/lib:profiler_session_impl",
-        "@tsl//tsl/profiler/protobuf:profiler_options_proto_cc",
-        "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
     ],
 )
 
@@ -346,6 +365,82 @@ xla_test(
     ]),
 )
 
+cc_library(
+    name = "simple_multipass_tracer_for_test",
+    testonly = True,
+    srcs = ["simple_multipass_tracer_for_test.cc"],
+    hdrs = ["simple_multipass_tracer_for_test.h"],
+    deps = [
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings:string_view",
+        "@tsl//tsl/profiler/lib:multi_pass_profiler_controller",
+        "@tsl//tsl/profiler/lib:profiler_collection",
+        "@tsl//tsl/profiler/lib:profiler_factory",
+        "@tsl//tsl/profiler/lib:profiler_interface",
+        "@tsl//tsl/profiler/protobuf:profiler_options_proto_cc",
+        "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
+    ],
+)
+
+xla_test(
+    name = "multipass_functional_hlo_runner_test",
+    srcs = ["multipass_functional_hlo_runner_test.cc"],
+    backend_args = if_google(
+        {
+            "gpu": ["--heap_check="],
+            "a100": ["--heap_check="],
+        },
+        {},
+    ),
+    backend_tags = {
+        "gpu": [
+            "nomsan",
+            "requires-gpu-sm80-only",
+        ],
+        "nvgpu_any": [
+            "broken",
+            "no_oss",
+        ],
+        "amdgpu_any": [
+            "broken",
+            "no_oss",
+        ],
+    },
+    backends = [
+        "cpu",
+        "gpu",
+    ],
+    data = [":hlo_files"],
+    tags = [
+        "no_mac",
+    ],
+    deps = [
+        ":create_client",
+        ":functional_hlo_runner",
+        ":hlo_input_output_format",
+        ":hlo_runner_profiler",
+        ":simple_multipass_tracer_for_test",
+        "//xla/pjrt:pjrt_client",
+        "//xla/pjrt/distributed:in_memory_key_value_store",
+        "//xla/tests:xla_test_backend_predicates",
+        "//xla/tsl/platform:env",
+        "//xla/tsl/platform:test",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:status_matchers",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_absl//absl/time",
+        "@com_google_googletest//:gtest",
+        "@tsl//tsl/platform:path",
+        "@tsl//tsl/profiler/lib:profiler_factory",
+        "@tsl//tsl/profiler/protobuf:xplane_proto_cc",
+    ],
+)
+
 tsl_pybind_extension(
     name = "py_hlo_multihost_runner",
     srcs = ["python_hlo_runner.cc"],
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
index 45c25fa5e69..c7fd524dcbb 100644
--- a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.cc
@@ -26,6 +26,7 @@ limitations under the License.
 #include <vector>
 
 #include "absl/algorithm/container.h"
+#include "absl/cleanup/cleanup.h"
 #include "absl/container/btree_map.h"
 #include "absl/functional/function_ref.h"
 #include "absl/log/check.h"
@@ -33,6 +34,7 @@ limitations under the License.
 #include "absl/random/random.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
 #include "absl/strings/str_join.h"
@@ -85,18 +87,13 @@ limitations under the License.
 #include "xla/tools/hlo_control_flow_flattening.h"
 #include "xla/tools/multihost_hlo_runner/hlo_input_output_format.h"
 #include "xla/tsl/platform/env.h"
-#include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/file_system.h"
 #include "xla/tsl/platform/file_system_helper.h"
-#include "xla/tsl/platform/statusor.h"
 #include "xla/tsl/platform/threadpool.h"
 #include "xla/tsl/util/fixed_option_set_flag.h"
 #include "xla/util.h"
 #include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
-#include "tsl/profiler/lib/profiler_session.h"
-#include "tsl/profiler/protobuf/profiler_options.pb.h"
-#include "tsl/profiler/protobuf/xplane.pb.h"
 
 namespace xla {
 namespace FunctionalHloRunner {
@@ -597,13 +594,41 @@ absl::StatusOr<PerDeviceLiteralVecType> RunInternal(
   std::vector<std::vector<std::unique_ptr<PjRtBuffer>>> device_buffers;
   std::vector<std::vector<PjRtBuffer*>> argument_ptrs;
 
+  auto refresh_argument_ptrs = [&]() -> absl::Status {
+    if (running_options.recreate_buffers_between_repeats) {
+      device_buffers.clear();
+      argument_ptrs.clear();
+      ABSL_ASSIGN_OR_RETURN(device_buffers,
+                       create_argument_buffers_on_device(flatten_arguments));
+      argument_ptrs = CreateArgumentPointersFromDeviceBuffers(device_buffers);
+      return absl::OkStatus();
+    }
+    switch (parameter_type) {
+      case ParameterType::kOneTupleOfArrays:
+        argument_ptrs = CreateArgumentPointersBasedOnAliasing(
+            output_buffers, device_buffers,
+            get_output_index_for_one_tuple_of_arrays);
+        break;
+      case ParameterType::kOneListOfArrays:
+        argument_ptrs = CreateArgumentPointersBasedOnAliasing(
+            output_buffers, device_buffers,
+            get_output_index_for_one_list_of_arrays);
+        break;
+      case ParameterType::kOther:
+        argument_ptrs = CreateArgumentPointersFromDeviceBuffers(device_buffers);
+        break;
+    }
+    return absl::OkStatus();
+  };
+
   bool has_active_profiler_session = false;
   for (int repeat = 0; repeat < running_options.num_repeats; ++repeat) {
     const bool is_last_repeat = (repeat == running_options.num_repeats - 1);
     const bool profile_current_repeat =
         (running_options.profiler != nullptr) &&
         (repeat >= running_options.num_repeats -
-                       running_options.num_repeats_with_profiler);
+                       running_options.num_repeats_with_profiler) &&
+        (!running_options.multipass_profiling || is_last_repeat);
 
     VLOG(1) << "FunctionalHloRunner: ExecuteOnDevices started (repeat = "
             << repeat << ").";
@@ -629,16 +654,47 @@ absl::StatusOr<PerDeviceLiteralVecType> RunInternal(
         running_options.profiler->CreateSession();
         has_active_profiler_session = true;
       }
-      futures->clear();
-      ABSL_ASSIGN_OR_RETURN(
-          output_buffers,
-          executable->Execute(argument_ptrs, execute_options, futures));
-      absl::Status execute_status;
-      for (auto& future : *futures) {
-        execute_status.Update(future.Await());
-      }
-      ABSL_RETURN_IF_ERROR(execute_status);
+      auto execute_once = [&]() -> absl::Status {
+        futures->clear();
+        ABSL_ASSIGN_OR_RETURN(
+            std::vector<std::vector<std::unique_ptr<PjRtBuffer>>>
+                new_output_buffers,
+            executable->Execute(argument_ptrs, execute_options, futures));
+        absl::Status execute_status;
+        for (auto& future : *futures) {
+          execute_status.Update(future.Await());
+        }
+        ABSL_RETURN_IF_ERROR(execute_status);
+        output_buffers = std::move(new_output_buffers);
+        return absl::OkStatus();
+      };
 
+      if (running_options.multipass_profiling && profile_current_repeat) {
+        bool executed_in_multipass = false;
+        while (true) {
+          ABSL_ASSIGN_OR_RETURN(bool need_more_pass,
+                           running_options.profiler->NeedMorePasses());
+          if (!need_more_pass) {
+            break;
+          }
+          if (executed_in_multipass) {
+            ABSL_RETURN_IF_ERROR(refresh_argument_ptrs());
+          }
+          running_options.profiler->StartPass();
+          running_options.profiler->PushRange("hlo_runner_pass");
+          absl::Cleanup end_pass = [&] {
+            running_options.profiler->PopRange();
+            running_options.profiler->StopPass();
+          };
+          ABSL_RETURN_IF_ERROR(execute_once());
+          executed_in_multipass = true;
+        }
+        if (!executed_in_multipass) {
+          ABSL_RETURN_IF_ERROR(execute_once());
+        }
+      } else {
+        ABSL_RETURN_IF_ERROR(execute_once());
+      }
       const bool upload_active_profiler_session =
           running_options.recreate_profiler_session_between_repeats ||
           is_last_repeat;
@@ -650,23 +706,8 @@ absl::StatusOr<PerDeviceLiteralVecType> RunInternal(
     }
     VLOG(1) << "FunctionalHloRunner: ExecuteOnDevices succeeded (repeat = "
             << repeat << ")";
-    if (!is_last_repeat) {
-      switch (parameter_type) {
-        case ParameterType::kOneTupleOfArrays:
-          argument_ptrs = CreateArgumentPointersBasedOnAliasing(
-              output_buffers, device_buffers,
-              get_output_index_for_one_tuple_of_arrays);
-          break;
-        case ParameterType::kOneListOfArrays:
-          argument_ptrs = CreateArgumentPointersBasedOnAliasing(
-              output_buffers, device_buffers,
-              get_output_index_for_one_list_of_arrays);
-          break;
-        case ParameterType::kOther:
-          argument_ptrs =
-              CreateArgumentPointersFromDeviceBuffers(device_buffers);
-          break;
-      }
+    if (!is_last_repeat && !running_options.recreate_buffers_between_repeats) {
+      ABSL_RETURN_IF_ERROR(refresh_argument_ptrs());
     }
   }
 
@@ -1827,58 +1868,6 @@ absl::StatusOr<ResolveTopologyResult> ResolveTopology(
 
 }  // namespace FunctionalHloRunner
 
-HLORunnerProfiler::HLORunnerProfiler(absl::string_view dump_path,
-                                     bool keep_xspace)
-    : dump_path_(dump_path), keep_xspace_(keep_xspace) {}
-
-absl::StatusOr<std::unique_ptr<HLORunnerProfiler>> HLORunnerProfiler::Create(
-    absl::string_view dump_path, bool keep_xspace) {
-  if (dump_path.empty()) {
-    return absl::InvalidArgumentError(
-        "Please provide a valid dump path to save XSpace results to disk.");
-  }
-  return std::make_unique<HLORunnerProfiler>(dump_path, keep_xspace);
-}
-
-void HLORunnerProfiler::CreateSession() {
-  auto options = tsl::ProfilerSession::DefaultOptions();
-  session_ = tsl::ProfilerSession::Create(options);
-}
-
-void HLORunnerProfiler::UploadSession() {
-  xspace_ = std::make_unique<tensorflow::profiler::XSpace>();
-  // Stops the ProfilerSession
-  CHECK_OK(session_->CollectData(xspace_.get()));
-
-  CHECK(!dump_path_.empty());
-
-  std::string unique_dump_path = dump_path_;
-  if (session_index_ > 0) {
-    absl::string_view stem = dump_path_;
-    absl::string_view suffix = "";
-    const std::string::size_type dot_pos = dump_path_.rfind('.');
-    const std::string::size_type slash_pos = dump_path_.rfind('/');
-    if (dot_pos != std::string::npos &&
-        (slash_pos == std::string::npos || dot_pos > slash_pos)) {
-      suffix = stem.substr(dot_pos);
-      stem = stem.substr(0, dot_pos);
-    }
-    unique_dump_path = absl::StrCat(stem, "_", session_index_, suffix);
-  }
-  ++session_index_;
-
-  LOG(INFO) << "Saving xspace result to " << unique_dump_path;
-  // Save in binary format to create xprof sessions and extract device stats.
-  CHECK_OK(WriteBinaryProto(tsl::Env::Default(), unique_dump_path, *xspace_));
-  if (!keep_xspace_) {
-    xspace_ = nullptr;
-  }
-}
-
-const tensorflow::profiler::XSpace* HLORunnerProfiler::GetXSpace() {
-  return xspace_.get();
-}
-
 void AddShardingAnnotationsToSpmdPartitionedModule(HloModule* hlo_module) {
   auto set_manual_sharding = [](HloInstruction* hlo) {
     if (!hlo->has_sharding()) {
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.h b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.h
index 5e36ee2e9ce..5b446ee3a8b 100644
--- a/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.h
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/functional_hlo_runner.h
@@ -39,68 +39,12 @@ limitations under the License.
 #include "xla/pjrt/pjrt_executable.h"
 #include "xla/shape.h"
 #include "xla/tools/multihost_hlo_runner/hlo_input_output_format.h"
+#include "xla/tools/multihost_hlo_runner/hlo_runner_profiler.h"
 #include "xla/tools/multihost_hlo_runner/profiler_interface.h"
 #include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
-#include "tsl/profiler/lib/profiler_session.h"
-#include "tsl/profiler/protobuf/xplane.pb.h"
 
 namespace xla {
-// Interface that may optionally returns an XSpace proto after UploadSession()
-// is called. This can be used by caller to get a programmatic handler of the
-// profile data.
-class XSpaceProfilerInterface : public ProfilerInterface {
- public:
-  virtual const tensorflow::profiler::XSpace* GetXSpace() = 0;
-};
-
-// HLORunnerProfiler is a profiler plugin that using tsl::ProfilerSession to
-// profile CPU/GPU execution and allows programmable control of
-// profiling sessions for the MultihostHloRunner. It needs to be created after
-// PJRT client is initialized. Example usage:
-//
-//   ABSL_ASSIGN_OR_RETURN(
-//       env, xla::GetPjRtEnvironmentForGpu(...)));
-//   if (env.client != nullptr) {
-//     ABSL_ASSIGN_OR_RETURN(auto profiler, HLORunnerProfiler::Create());
-//   }
-//   profiler.CreateSession();
-//   ...
-//   profiler.UploadSession();
-class HLORunnerProfiler : public XSpaceProfilerInterface {
- public:
-  // Factory method to create a GPURunnerProfiler with profile result dump path.
-  // If keep_xspace is true, the XSpace proto can be retrieved
-  // by GetXSpace() after UploadSession() is called, which can be used by
-  // caller to get a programmatic handler of the profile data and create XProf.
-  static absl::StatusOr<std::unique_ptr<HLORunnerProfiler>> Create(
-      absl::string_view dump_path, bool keep_xspace = false);
-
-  // Default ctor.
-  explicit HLORunnerProfiler(absl::string_view dump_path, bool keep_xspace);
-
-  // Start a new profiling session.
-  void CreateSession() override;
-
-  // Stop the current profiling session.
-  void UploadSession() override;
-
-  // Returns the XSpace proto.
-  const tensorflow::profiler::XSpace* GetXSpace() override;
-
- private:
-  // The file path to dump the profiling result.
-  std::string dump_path_;
-  // Whether to keep the XSpace proto after UploadSession() is called.
-  bool keep_xspace_;
-  // The profiler session.
-  std::unique_ptr<tsl::ProfilerSession> session_;
-  // The XSpace proto to be returned by GetXSpace().
-  std::unique_ptr<tensorflow::profiler::XSpace> xspace_;
-  // Session counter to uniquely name dump paths when multiple sessions are
-  // uploaded.
-  int session_index_ = 0;
-};
 
 // FunctionalHloRunner takes an HLO module as input and runs the HLO module
 // on a single or multiple hosts with various options (e.g. SPMD). The HLO
@@ -295,6 +239,7 @@ struct RunningOptions {
   // This indicates whether we log the inputs and outputs to stderr.
   LogOutputMode log_input_output_mode = LogOutputMode::kNotLogOutput;
   const MultiSliceConfig* multi_slice_config = nullptr;
+  bool multipass_profiling = false;
   ProfilerInterface* profiler = nullptr;
   // If not null, profiles will be stored for this run, one per repeat.
   // Note that the first repeat is a warmup run, and uses less precise
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/hlo_runner_profiler.cc b/third_party/xla/xla/tools/multihost_hlo_runner/hlo_runner_profiler.cc
new file mode 100644
index 00000000000..5e069da3946
--- /dev/null
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/hlo_runner_profiler.cc
@@ -0,0 +1,322 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include "xla/tools/multihost_hlo_runner/hlo_runner_profiler.h"
+
+#include <algorithm>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/status_macros.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/time/clock.h"
+#include "absl/time/time.h"
+#include "xla/pjrt/distributed/key_value_store_interface.h"
+#include "xla/tsl/platform/env.h"
+#include "tsl/profiler/lib/profiler_passes.h"
+#include "tsl/profiler/lib/profiler_session.h"
+#include "tsl/profiler/protobuf/xplane.pb.h"
+
+namespace xla {
+namespace {
+
+constexpr absl::string_view kAllReduceErrorToken = "ERROR";
+
+// Multi-pass profiler implementation for single-host HLO runner execution.
+class MultiPassHloProfiler : public HLORunnerProfiler {
+ public:
+  MultiPassHloProfiler(absl::string_view dump_path, bool keep_xspace)
+      : HLORunnerProfiler(dump_path, keep_xspace) {}
+
+  void CreateSession() override {
+    multi_passes_ =
+        tsl::ProfilerPasses::Create(tsl::ProfilerPasses::DefaultOptions());
+  }
+
+  bool IsMultipassEnabled() const override { return true; }
+
+  absl::StatusOr<bool> NeedMorePasses() override {
+    if (multi_passes_ == nullptr) {
+      return false;
+    }
+    return multi_passes_->NeedMorePasses();
+  }
+
+  void StartPass() override {
+    if (multi_passes_ != nullptr) {
+      absl::Status status = multi_passes_->StartPass();
+      if (!status.ok()) {
+        LOG(ERROR) << "Failed to start pass: " << status;
+      }
+    }
+  }
+
+  void PushRange(absl::string_view name) override {
+    if (multi_passes_ != nullptr) {
+      absl::Status status = multi_passes_->PushRange(name);
+      if (!status.ok()) {
+        LOG(ERROR) << "Failed to push range: " << status;
+      }
+    }
+  }
+
+  void PopRange() override {
+    if (multi_passes_ != nullptr) {
+      absl::Status status = multi_passes_->PopRange();
+      if (!status.ok()) {
+        LOG(ERROR) << "Failed to pop range: " << status;
+      }
+    }
+  }
+
+  void StopPass() override {
+    if (multi_passes_ != nullptr) {
+      absl::Status status = multi_passes_->StopPass();
+      if (!status.ok()) {
+        LOG(ERROR) << "Failed to stop pass: " << status;
+      }
+    }
+  }
+
+  void UploadSession() override {
+    auto local_xspace = std::make_unique<tensorflow::profiler::XSpace>();
+    if (multi_passes_ != nullptr) {
+      CHECK_OK(multi_passes_->CollectData(local_xspace.get()));
+    }
+    SaveXSpace(std::move(local_xspace));
+  }
+
+ protected:
+  std::unique_ptr<tsl::ProfilerPasses> multi_passes_;
+};
+
+// Multi-pass profiler implementation for multi-host HLO runner execution.
+// Uses KeyValueStoreInterface to synchronize pass progression across ranks.
+class CollectiveMultipassHloProfiler : public MultiPassHloProfiler {
+ public:
+  CollectiveMultipassHloProfiler(
+      absl::string_view dump_path, bool keep_xspace,
+      std::shared_ptr<KeyValueStoreInterface> multipass_synchronizer,
+      int task_id, int num_nodes,
+      std::optional<int> collective_id = std::nullopt)
+      : MultiPassHloProfiler(dump_path, keep_xspace),
+        multipass_synchronizer_(std::move(multipass_synchronizer)),
+        task_id_(task_id),
+        num_nodes_(num_nodes),
+        collective_id_(collective_id) {}
+
+  void CreateSession() override {
+    MultiPassHloProfiler::CreateSession();
+    session_count_++;
+    pass_idx_ = 0;
+  }
+
+  absl::StatusOr<bool> NeedMorePasses() override {
+    bool local_need_more = false;
+    if (multi_passes_ != nullptr) {
+      ABSL_ASSIGN_OR_RETURN(local_need_more, multi_passes_->NeedMorePasses());
+    }
+
+    std::string keystr =
+        collective_id_.has_value()
+            ? absl::StrCat("hlo_runner_need_more_pass_", *collective_id_, "_",
+                           session_count_, "_", pass_idx_)
+            : absl::StrCat("hlo_runner_need_more_pass_", session_count_, "_",
+                           pass_idx_);
+    ABSL_ASSIGN_OR_RETURN(
+        bool global_need_more,
+        AllReduceOr(multipass_synchronizer_.get(), task_id_, num_nodes_,
+                    local_need_more, keystr, absl::Seconds(30)));
+
+    local_need_more_for_current_pass_ = local_need_more;
+    return global_need_more;
+  }
+
+  void StartPass() override {
+    if (local_need_more_for_current_pass_) {
+      MultiPassHloProfiler::StartPass();
+    }
+  }
+
+  void PushRange(absl::string_view name) override {
+    if (local_need_more_for_current_pass_) {
+      MultiPassHloProfiler::PushRange(name);
+    }
+  }
+
+  void PopRange() override {
+    if (local_need_more_for_current_pass_) {
+      MultiPassHloProfiler::PopRange();
+    }
+  }
+
+  void StopPass() override {
+    if (local_need_more_for_current_pass_) {
+      MultiPassHloProfiler::StopPass();
+    }
+    pass_idx_++;
+  }
+
+ private:
+  std::shared_ptr<KeyValueStoreInterface> multipass_synchronizer_;
+  int task_id_ = 0;
+  int num_nodes_ = 1;
+  std::optional<int> collective_id_ = std::nullopt;
+  int session_count_ = 0;
+  int pass_idx_ = 0;
+  bool local_need_more_for_current_pass_ = true;
+};
+
+}  // namespace
+
+absl::StatusOr<bool> AllReduceOr(KeyValueStoreInterface* kv_store, int rank_id,
+                                 int num_ranks, bool local_value,
+                                 absl::string_view keystr,
+                                 absl::Duration timeout) {
+  if (kv_store == nullptr || num_ranks <= 1) {
+    return local_value;
+  }
+  if (rank_id < 0 || rank_id >= num_ranks) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Invalid rank_id ", rank_id, " for num_ranks ", num_ranks));
+  }
+
+  std::string my_key = absl::StrCat(keystr, "_", rank_id);
+  ABSL_RETURN_IF_ERROR(kv_store->Set(my_key, local_value ? "1" : "0"));
+
+  bool global_or = false;
+  std::string decision_key = absl::StrCat(keystr, "_decision");
+
+  if (rank_id == 0) {
+    global_or = local_value;
+    const absl::Time deadline = absl::Now() + timeout;
+    for (int i = 1; i < num_ranks; ++i) {
+      std::string peer_key = absl::StrCat(keystr, "_", i);
+      const absl::Duration remaining =
+          std::max(absl::ZeroDuration(), deadline - absl::Now());
+      absl::StatusOr<std::string> val_str = kv_store->Get(peer_key, remaining);
+      if (!val_str.ok()) {
+        // Publish an error token so peer ranks waiting on `decision_key`
+        // unblock promptly rather than waiting out the full timeout.
+        kv_store->Set(decision_key, kAllReduceErrorToken).IgnoreError();
+        return val_str.status();
+      }
+      if (*val_str == "1") {
+        global_or = true;
+      }
+    }
+    ABSL_RETURN_IF_ERROR(kv_store->Set(decision_key, global_or ? "1" : "0"));
+  } else {
+    ABSL_ASSIGN_OR_RETURN(std::string decision_str,
+                     kv_store->Get(decision_key, timeout));
+    if (decision_str == kAllReduceErrorToken) {
+      return absl::InternalError(
+          "Rank 0 failed to gather peer values during AllReduceOr.");
+    }
+    global_or = (decision_str == "1");
+  }
+
+  return global_or;
+}
+
+HLORunnerProfiler::HLORunnerProfiler(absl::string_view dump_path,
+                                     bool keep_xspace)
+    : dump_path_(dump_path), keep_xspace_(keep_xspace) {}
+
+HLORunnerProfiler::~HLORunnerProfiler() = default;
+
+absl::StatusOr<std::unique_ptr<HLORunnerProfiler>> HLORunnerProfiler::Create(
+    absl::string_view dump_path, bool keep_xspace,
+    bool enable_multipass_profiling,
+    std::shared_ptr<KeyValueStoreInterface> kv_store, int task_id,
+    int num_nodes, std::optional<int> collective_id) {
+  if (dump_path.empty()) {
+    return absl::InvalidArgumentError(
+        "Please provide a valid dump path to save XSpace results to disk.");
+  }
+  if (enable_multipass_profiling) {
+    if (num_nodes > 1) {
+      if (kv_store == nullptr) {
+        return absl::InvalidArgumentError(
+            "Multi-host multi-pass profiling requires a non-null "
+            "KeyValueStoreInterface for collective synchronization.");
+      }
+      if (task_id < 0 || task_id >= num_nodes) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Invalid task_id ", task_id, " for num_nodes ", num_nodes));
+      }
+      return std::make_unique<CollectiveMultipassHloProfiler>(
+          dump_path, keep_xspace, std::move(kv_store), task_id, num_nodes,
+          collective_id);
+    }
+    LOG(INFO) << "Creating Single-host MultiPassHloProfiler";
+    return std::make_unique<MultiPassHloProfiler>(dump_path, keep_xspace);
+  }
+  return std::unique_ptr<HLORunnerProfiler>(
+      new HLORunnerProfiler(dump_path, keep_xspace));
+}
+
+void HLORunnerProfiler::CreateSession() {
+  auto options = tsl::ProfilerSession::DefaultOptions();
+  session_ = tsl::ProfilerSession::Create(options);
+}
+
+void HLORunnerProfiler::UploadSession() {
+  auto local_xspace = std::make_unique<tensorflow::profiler::XSpace>();
+  if (session_ != nullptr) {
+    CHECK_OK(session_->CollectData(local_xspace.get()));
+  }
+  SaveXSpace(std::move(local_xspace));
+}
+
+void HLORunnerProfiler::SaveXSpace(
+    std::unique_ptr<tensorflow::profiler::XSpace> local_xspace) {
+  xspace_ = std::move(local_xspace);
+  CHECK(!dump_path_.empty());
+  std::string unique_dump_path = dump_path_;
+  if (session_index_ > 0) {
+    absl::string_view stem = dump_path_;
+    absl::string_view suffix = "";
+    const std::string::size_type dot_pos = dump_path_.rfind('.');
+    const std::string::size_type slash_pos = dump_path_.rfind('/');
+    if (dot_pos != std::string::npos &&
+        (slash_pos == std::string::npos || dot_pos > slash_pos)) {
+      suffix = stem.substr(dot_pos);
+      stem = stem.substr(0, dot_pos);
+    }
+    unique_dump_path = absl::StrCat(stem, "_", session_index_, suffix);
+  }
+  ++session_index_;
+
+  LOG(INFO) << "Saving xspace result to " << unique_dump_path;
+  // Save in binary format to create xprof sessions and extract device stats.
+  CHECK_OK(WriteBinaryProto(tsl::Env::Default(), unique_dump_path, *xspace_));
+  if (!keep_xspace_) {
+    xspace_ = nullptr;
+  }
+}
+
+const tensorflow::profiler::XSpace* HLORunnerProfiler::GetXSpace() {
+  return xspace_.get();
+}
+
+}  // namespace xla
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/hlo_runner_profiler.h b/third_party/xla/xla/tools/multihost_hlo_runner/hlo_runner_profiler.h
new file mode 100644
index 00000000000..786e3f59b83
--- /dev/null
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/hlo_runner_profiler.h
@@ -0,0 +1,116 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_TOOLS_MULTIHOST_HLO_RUNNER_HLO_RUNNER_PROFILER_H_
+#define XLA_TOOLS_MULTIHOST_HLO_RUNNER_HLO_RUNNER_PROFILER_H_
+
+#include <memory>
+#include <optional>
+#include <string>
+
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "absl/time/time.h"
+#include "xla/pjrt/distributed/key_value_store_interface.h"
+#include "xla/tools/multihost_hlo_runner/profiler_interface.h"
+#include "tsl/profiler/lib/profiler_session.h"
+#include "tsl/profiler/protobuf/xplane.pb.h"
+
+namespace xla {
+
+// Performs an all-reduce OR across all ranks using the KeyValueStoreInterface.
+// In multipass profiling, this is used to reach consensus on whether any host
+// requires an additional profiling pass (`NeedMorePasses()`), ensuring all
+// ranks remain synchronized and execute the same number of passes.
+//
+// If `num_ranks <= 1` or `kv_store == nullptr` (single rank / single host
+// mode), no communication is needed and the local value is returned directly.
+//
+// Callers must provide a unique `keystr` for each reduction step (e.g., by
+// appending the pass index) to avoid collisions with prior passes or concurrent
+// reductions in the key-value store.
+absl::StatusOr<bool> AllReduceOr(KeyValueStoreInterface* kv_store, int rank_id,
+                                 int num_ranks, bool local_value,
+                                 absl::string_view keystr,
+                                 absl::Duration timeout);
+
+// Interface that may optionally return an XSpace proto after UploadSession()
+// is called. This can be used by callers to get a programmatic handle of the
+// profile data.
+class XSpaceProfilerInterface : public ProfilerInterface {
+ public:
+  virtual const tensorflow::profiler::XSpace* GetXSpace() = 0;
+};
+
+// HLORunnerProfiler is a profiler plugin that uses tsl::ProfilerSession to
+// profile CPU/GPU execution and allows programmable control of
+// profiling sessions for the MultihostHloRunner. It needs to be created after
+// PJRT client is initialized. Example usage:
+//
+//   ABSL_ASSIGN_OR_RETURN(
+//       env, xla::GetPjRtEnvironmentForGpu(...));
+//   if (env.client != nullptr) {
+//     ABSL_ASSIGN_OR_RETURN(auto profiler, HLORunnerProfiler::Create(...));
+//   }
+//   profiler->CreateSession();
+//   ...
+//   profiler->UploadSession();
+class HLORunnerProfiler : public XSpaceProfilerInterface {
+ public:
+  // Factory method to create an HLORunnerProfiler with profile result dump
+  // path.
+  // If keep_xspace is true, the XSpace proto can be retrieved
+  // by GetXSpace() after UploadSession() is called, which can be used by
+  // callers to get a programmatic handle of the profile data and create XProf.
+  static absl::StatusOr<std::unique_ptr<HLORunnerProfiler>> Create(
+      absl::string_view dump_path, bool keep_xspace = false,
+      bool enable_multipass_profiling = false,
+      std::shared_ptr<KeyValueStoreInterface> kv_store = nullptr,
+      int task_id = 0, int num_nodes = 1,
+      std::optional<int> collective_id = std::nullopt);
+  ~HLORunnerProfiler() override;
+
+  // Start a new profiling session.
+  void CreateSession() override;
+
+  // Stop the current profiling session.
+  void UploadSession() override;
+
+  // Returns the XSpace proto.
+  const tensorflow::profiler::XSpace* GetXSpace() override;
+
+ protected:
+  explicit HLORunnerProfiler(absl::string_view dump_path, bool keep_xspace);
+
+  void SaveXSpace(std::unique_ptr<tensorflow::profiler::XSpace> xspace);
+
+  // The file path to dump the profiling result.
+  std::string dump_path_;
+  // Whether to keep the XSpace proto after UploadSession() is called.
+  bool keep_xspace_;
+  // The XSpace proto to be returned by GetXSpace().
+  std::unique_ptr<tensorflow::profiler::XSpace> xspace_;
+
+ private:
+  // The profiler session.
+  std::unique_ptr<tsl::ProfilerSession> session_;
+  // Session counter to uniquely name dump paths when multiple sessions are
+  // uploaded.
+  int session_index_ = 0;
+};
+
+}  // namespace xla
+
+#endif  // XLA_TOOLS_MULTIHOST_HLO_RUNNER_HLO_RUNNER_PROFILER_H_
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/multipass_functional_hlo_runner_test.cc b/third_party/xla/xla/tools/multihost_hlo_runner/multipass_functional_hlo_runner_test.cc
new file mode 100644
index 00000000000..1aef8f1b792
--- /dev/null
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/multipass_functional_hlo_runner_test.cc
@@ -0,0 +1,327 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "xla/tests/xla_test_backend_predicates.h"
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/status_matchers.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/time/time.h"
+#include "xla/pjrt/distributed/in_memory_key_value_store.h"
+#include "xla/pjrt/pjrt_client.h"
+#include "xla/tools/multihost_hlo_runner/create_client.h"
+#include "xla/tools/multihost_hlo_runner/functional_hlo_runner.h"
+#include "xla/tools/multihost_hlo_runner/hlo_input_output_format.h"
+#include "xla/tools/multihost_hlo_runner/hlo_runner_profiler.h"
+#include "xla/tools/multihost_hlo_runner/simple_multipass_tracer_for_test.h"
+#include "xla/tsl/platform/env.h"
+#include "xla/tsl/platform/test.h"
+#include "tsl/platform/path.h"
+#include "tsl/profiler/lib/profiler_factory.h"
+#include "tsl/profiler/protobuf/xplane.pb.h"
+
+namespace xla {
+namespace {
+
+using ::absl_testing::IsOkAndHolds;
+using ::absl_testing::StatusIs;
+using ::testing::IsEmpty;
+using ::testing::Not;
+using ::testing::UnorderedElementsAreArray;
+
+std::string GetHloPath(absl::string_view file_name) {
+  return tsl::io::JoinPath(tsl::testing::XlaSrcRoot(), "tools",
+                           "multihost_hlo_runner", "data", file_name);
+}
+
+absl::StatusOr<std::unique_ptr<xla::PjRtClient>> GetPjRtClient() {
+  if (test::DeviceTypeIs(test::kCpu)) {
+    LOG(INFO) << "Running on CPU";
+    return CreateHostClient();
+  }
+  LOG(INFO) << "Running on GPU";
+  return CreateGpuClient({});
+}
+
+class MultiPassFunctionalHloRunnerTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    xla::profiler::RegisterSimpleMultiPassTracerForTest();
+  }
+
+  void TearDown() override {
+    tsl::profiler::ClearRegisteredMultiPassProfilersForTest();
+  }
+};
+
+TEST_F(MultiPassFunctionalHloRunnerTest, SimpleMultiPassTracerTest) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<xla::PjRtClient> client,
+                       GetPjRtClient());
+
+  // 1. Run non-multipass profiling session.
+  std::string profile_dump_path_non_multipass =
+      tsl::io::JoinPath(testing::TempDir(), "xspace_non_multipass.pb");
+  FunctionalHloRunner::RunningOptions running_options_non_multipass;
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<HLORunnerProfiler> profiler_non_multipass,
+      HLORunnerProfiler::Create(profile_dump_path_non_multipass,
+                                /*keep_xspace=*/true,
+                                /*enable_multipass_profiling=*/false));
+  running_options_non_multipass.profiler = profiler_non_multipass.get();
+  running_options_non_multipass.multipass_profiling = false;
+
+  // We must use a valid HLO file to run.
+  EXPECT_OK(FunctionalHloRunner::LoadAndRunAndDump(
+      *client,
+      /* preproc_options= */ {},
+      /* raw_compile_options = */ {}, running_options_non_multipass,
+      GetHloPath("single_device.hlo"), InputFormat::kText));
+
+  const tensorflow::profiler::XSpace* xspace_non_multipass =
+      profiler_non_multipass->GetXSpace();
+  ASSERT_NE(xspace_non_multipass, nullptr);
+
+  // 2. Run multipass profiling session.
+  std::string profile_dump_path_multipass =
+      tsl::io::JoinPath(testing::TempDir(), "xspace_multipass.pb");
+  FunctionalHloRunner::RunningOptions running_options_multipass;
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HLORunnerProfiler> profiler_multipass,
+                       HLORunnerProfiler::Create(
+                           profile_dump_path_multipass, /*keep_xspace=*/true,
+                           /*enable_multipass_profiling=*/true));
+  running_options_multipass.profiler = profiler_multipass.get();
+  running_options_multipass.multipass_profiling = true;
+
+  EXPECT_OK(FunctionalHloRunner::LoadAndRunAndDump(
+      *client,
+      /* preproc_options= */ {},
+      /* raw_compile_options = */ {}, running_options_multipass,
+      GetHloPath("single_device.hlo"), InputFormat::kText));
+
+  const tensorflow::profiler::XSpace* xspace_multipass =
+      profiler_multipass->GetXSpace();
+  ASSERT_NE(xspace_multipass, nullptr);
+
+  // 3. Verify that the planes collected in both modes are identical.
+  std::vector<std::string> plane_names_non_multipass;
+  for (const auto& plane : xspace_non_multipass->planes()) {
+    LOG(INFO) << "Non multipass profiling results plane: " << plane.name();
+    plane_names_non_multipass.push_back(plane.name());
+  }
+
+  std::vector<std::string> plane_names_multipass;
+  for (const auto& plane : xspace_multipass->planes()) {
+    LOG(INFO) << "Multipass profiling results plane: " << plane.name();
+    plane_names_multipass.push_back(plane.name());
+  }
+
+  EXPECT_THAT(plane_names_multipass, Not(IsEmpty()));
+  EXPECT_THAT(plane_names_multipass,
+              UnorderedElementsAreArray(plane_names_non_multipass));
+}
+
+TEST_F(MultiPassFunctionalHloRunnerTest,
+       CreateFailsWhenMultiNodeWithoutKvStore) {
+  std::string profile_dump_path =
+      tsl::io::JoinPath(testing::TempDir(), "xspace_multinode.pb");
+  EXPECT_THAT(HLORunnerProfiler::Create(profile_dump_path, /*keep_xspace=*/true,
+                                        /*enable_multipass_profiling=*/true,
+                                        /*kv_store=*/nullptr, /*task_id=*/0,
+                                        /*num_nodes=*/2),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+TEST_F(MultiPassFunctionalHloRunnerTest, CreateFailsWhenTaskIdOutOfBounds) {
+  auto kv_store = std::make_shared<InMemoryKeyValueStore>();
+  std::string profile_dump_path =
+      tsl::io::JoinPath(testing::TempDir(), "xspace_invalid_task_id.pb");
+  EXPECT_THAT(HLORunnerProfiler::Create(profile_dump_path, /*keep_xspace=*/true,
+                                        /*enable_multipass_profiling=*/true,
+                                        kv_store, /*task_id=*/-1,
+                                        /*num_nodes=*/2),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+  EXPECT_THAT(HLORunnerProfiler::Create(profile_dump_path, /*keep_xspace=*/true,
+                                        /*enable_multipass_profiling=*/true,
+                                        kv_store, /*task_id=*/2,
+                                        /*num_nodes=*/2),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+TEST_F(MultiPassFunctionalHloRunnerTest,
+       UploadSessionWithoutCreateSessionDoesNotCrash) {
+  std::string profile_dump_path =
+      tsl::io::JoinPath(testing::TempDir(), "xspace_empty_upload.pb");
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<HLORunnerProfiler> profiler,
+      HLORunnerProfiler::Create(profile_dump_path, /*keep_xspace=*/true,
+                                /*enable_multipass_profiling=*/true));
+  profiler->UploadSession();
+  EXPECT_NE(profiler->GetXSpace(), nullptr);
+}
+
+TEST_F(MultiPassFunctionalHloRunnerTest, AllReduceOrSingleRankAndBoundsCheck) {
+  InMemoryKeyValueStore kv_store;
+  EXPECT_THAT(AllReduceOr(/*kv_store=*/nullptr, /*rank_id=*/0, /*num_ranks=*/2,
+                          /*local_value=*/true, "key", absl::Seconds(1)),
+              IsOkAndHolds(true));
+  EXPECT_THAT(AllReduceOr(&kv_store, /*rank_id=*/0, /*num_ranks=*/1,
+                          /*local_value=*/false, "key", absl::Seconds(1)),
+              IsOkAndHolds(false));
+  EXPECT_THAT(AllReduceOr(&kv_store, /*rank_id=*/-1, /*num_ranks=*/2,
+                          /*local_value=*/false, "key", absl::Seconds(1)),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+  EXPECT_THAT(AllReduceOr(&kv_store, /*rank_id=*/2, /*num_ranks=*/2,
+                          /*local_value=*/false, "key", absl::Seconds(1)),
+              StatusIs(absl::StatusCode::kInvalidArgument));
+}
+
+TEST_F(MultiPassFunctionalHloRunnerTest, AllReduceOrMultiRankConsensus) {
+  InMemoryKeyValueStore kv_store;
+
+  auto run_two_ranks = [&](bool rank0_val, bool rank1_val,
+                           absl::string_view key) {
+    absl::StatusOr<bool> res0;
+    absl::StatusOr<bool> res1;
+    std::unique_ptr<tsl::Thread> t1(
+        tsl::Env::Default()->StartThread(tsl::ThreadOptions(), "rank1", [&]() {
+          res1 = AllReduceOr(&kv_store, /*rank_id=*/1, /*num_ranks=*/2,
+                             rank1_val, key, absl::Seconds(5));
+        }));
+    res0 = AllReduceOr(&kv_store, /*rank_id=*/0, /*num_ranks=*/2, rank0_val,
+                       key, absl::Seconds(5));
+    t1.reset();
+    EXPECT_EQ(res0, res1);
+    return res0;
+  };
+
+  EXPECT_THAT(run_two_ranks(false, false, "step_0"), IsOkAndHolds(false));
+  EXPECT_THAT(run_two_ranks(false, true, "step_1"), IsOkAndHolds(true));
+  EXPECT_THAT(run_two_ranks(true, false, "step_2"), IsOkAndHolds(true));
+  EXPECT_THAT(run_two_ranks(true, true, "step_3"), IsOkAndHolds(true));
+}
+
+TEST_F(MultiPassFunctionalHloRunnerTest,
+       AllReduceOrTimeoutPropagatesErrorTokenToPeer) {
+  InMemoryKeyValueStore kv_store;
+  absl::StatusOr<bool> res0;
+  absl::StatusOr<bool> res2;
+
+  // 3 ranks, where rank 1 never participates. Rank 0 times out waiting for
+  // rank 1 and publishes kAllReduceErrorToken so rank 2 unblocks promptly
+  // instead of waiting for its 10-second timeout.
+  std::unique_ptr<tsl::Thread> t2(
+      tsl::Env::Default()->StartThread(tsl::ThreadOptions(), "rank2", [&]() {
+        res2 = AllReduceOr(&kv_store, /*rank_id=*/2, /*num_ranks=*/3,
+                           /*local_value=*/false, "timeout_key",
+                           absl::Seconds(10));
+      }));
+  res0 =
+      AllReduceOr(&kv_store, /*rank_id=*/0, /*num_ranks=*/3,
+                  /*local_value=*/true, "timeout_key", absl::Milliseconds(50));
+  t2.reset();
+
+  EXPECT_FALSE(res0.ok());
+  EXPECT_THAT(res2, StatusIs(absl::StatusCode::kInternal));
+}
+
+TEST_F(MultiPassFunctionalHloRunnerTest,
+       CollectiveMultipassHloProfilerSynchronizesPasses) {
+  auto kv_store = std::make_shared<InMemoryKeyValueStore>();
+  std::string dump_path_0 =
+      tsl::io::JoinPath(testing::TempDir(), "xspace_collective_rank0.pb");
+
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<HLORunnerProfiler> profiler_0,
+      HLORunnerProfiler::Create(dump_path_0, /*keep_xspace=*/true,
+                                /*enable_multipass_profiling=*/true, kv_store,
+                                /*task_id=*/0, /*num_nodes=*/2,
+                                /*collective_id=*/7));
+
+  profiler_0->CreateSession();
+
+  // SimpleMultiPassTracer on Rank 0 requires 2 local passes (passes 0 and 1).
+  // Simulate Peer Rank 1 requiring 3 passes (passes 0, 1, and 2) and false on
+  // pass 3. CollectiveMultipassHloProfiler on Rank 0 must synchronize and run
+  // all 3 passes, then stop on pass 3.
+  std::unique_ptr<tsl::Thread> peer_rank(
+      tsl::Env::Default()->StartThread(tsl::ThreadOptions(), "rank1", [&]() {
+        for (int pass_idx = 0; pass_idx < 4; ++pass_idx) {
+          bool peer_need_more = (pass_idx < 3);
+          std::string keystr =
+              absl::StrCat("hlo_runner_need_more_pass_7_1_", pass_idx);
+          EXPECT_THAT(
+              AllReduceOr(kv_store.get(), /*rank_id=*/1, /*num_ranks=*/2,
+                          peer_need_more, keystr, absl::Seconds(5)),
+              IsOkAndHolds(pass_idx < 3));
+        }
+      }));
+
+  int passes_rank_0 = 0;
+  while (true) {
+    ASSERT_OK_AND_ASSIGN(bool need_more_pass, profiler_0->NeedMorePasses());
+    if (!need_more_pass) {
+      break;
+    }
+    profiler_0->StartPass();
+    profiler_0->PushRange("pass");
+    profiler_0->PopRange();
+    profiler_0->StopPass();
+    passes_rank_0++;
+  }
+  peer_rank.reset();
+
+  EXPECT_EQ(passes_rank_0, 3);
+
+  profiler_0->UploadSession();
+  EXPECT_NE(profiler_0->GetXSpace(), nullptr);
+}
+
+TEST_F(MultiPassFunctionalHloRunnerTest,
+       CollectiveMultipassHloProfilerPropagatesErrorFromAllReduceOr) {
+  auto kv_store = std::make_shared<InMemoryKeyValueStore>();
+  std::string dump_path_1 =
+      tsl::io::JoinPath(testing::TempDir(), "xspace_collective_rank1.pb");
+
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<HLORunnerProfiler> profiler_1,
+      HLORunnerProfiler::Create(dump_path_1, /*keep_xspace=*/true,
+                                /*enable_multipass_profiling=*/true, kv_store,
+                                /*task_id=*/1, /*num_nodes=*/2,
+                                /*collective_id=*/9));
+
+  profiler_1->CreateSession();
+
+  // Simulate rank 0 publishing an error token during AllReduceOr.
+  ASSERT_OK(kv_store->Set("hlo_runner_need_more_pass_9_1_0_decision", "ERROR"));
+
+  EXPECT_THAT(profiler_1->NeedMorePasses(),
+              StatusIs(absl::StatusCode::kInternal));
+}
+
+}  // namespace
+}  // namespace xla
+
+int main(int argc, char* argv[]) {
+  testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/profiler_interface.h b/third_party/xla/xla/tools/multihost_hlo_runner/profiler_interface.h
index 4c3316a975e..a4a24672ac4 100644
--- a/third_party/xla/xla/tools/multihost_hlo_runner/profiler_interface.h
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/profiler_interface.h
@@ -16,6 +16,9 @@ limitations under the License.
 #ifndef XLA_TOOLS_MULTIHOST_HLO_RUNNER_PROFILER_INTERFACE_H_
 #define XLA_TOOLS_MULTIHOST_HLO_RUNNER_PROFILER_INTERFACE_H_
 
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+
 namespace xla {
 
 // Interface for profiler plugins. If being set in RunningOptions, profiling
@@ -27,7 +30,22 @@ class ProfilerInterface {
   virtual void CreateSession() = 0;
   // Uploads profiling session data after finishing running HLO module.
   virtual void UploadSession() = 0;
+
+  // Multipass profiling interface methods.
+  // Returns true if multipass profiling is enabled for this profiler.
+  virtual bool IsMultipassEnabled() const { return false; }
+  // Returns true if another profiling pass is required.
+  virtual absl::StatusOr<bool> NeedMorePasses() { return false; }
+  // Starts a profiling pass.
+  virtual void StartPass() {}
+  // Pushes a named range to delimit profiling regions within the active pass.
+  virtual void PushRange(absl::string_view name) {}
+  // Pops the innermost named range within the active pass.
+  virtual void PopRange() {}
+  // Stops the active profiling pass.
+  virtual void StopPass() {}
 };
+
 }  // namespace xla
 
 #endif  // XLA_TOOLS_MULTIHOST_HLO_RUNNER_PROFILER_INTERFACE_H_
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/simple_multipass_tracer_for_test.cc b/third_party/xla/xla/tools/multihost_hlo_runner/simple_multipass_tracer_for_test.cc
new file mode 100644
index 00000000000..b1669781ede
--- /dev/null
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/simple_multipass_tracer_for_test.cc
@@ -0,0 +1,167 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+// Simple multipass tracer wrapping all registered single-pass profilers.
+#include "xla/tools/multihost_hlo_runner/simple_multipass_tracer_for_test.h"
+
+#include <memory>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "tsl/profiler/lib/multi_pass_profiler_controller.h"
+#include "tsl/profiler/lib/profiler_collection.h"
+#include "tsl/profiler/lib/profiler_factory.h"
+#include "tsl/profiler/lib/profiler_interface.h"
+#include "tsl/profiler/protobuf/profiler_options.pb.h"
+#include "tsl/profiler/protobuf/xplane.pb.h"
+
+namespace xla {
+namespace profiler {
+namespace {
+
+using ::tensorflow::ProfileOptions;
+using ::tensorflow::profiler::XSpace;
+using ::tsl::profiler::MultiPassProfilerController;
+using ::tsl::profiler::MultiPassProfilerInterface;
+using ::tsl::profiler::ProfilerCollection;
+
+class SimpleMultiPassTracer : public MultiPassProfilerInterface {
+ public:
+  explicit SimpleMultiPassTracer(const ProfileOptions& options)
+      : options_(options) {
+    VLOG(3) << "SimpleMultiPassTracer created for device "
+            << options_.device_type();
+  }
+  ~SimpleMultiPassTracer() override {
+    VLOG(3) << "SimpleMultiPassTracer destroyed.";
+  }
+
+  // MultiPassProfilerInterface methods
+  absl::StatusOr<bool> NeedMorePasses() override {
+    VLOG(3) << "SimpleMultiPassTracer::NeedMorePasses. pass_count="
+            << pass_count_;
+    return pass_count_ < 2;
+  }
+
+  absl::Status StartPass() override {
+    VLOG(3) << "SimpleMultiPassTracer::StartPass. pass_count=" << pass_count_;
+    if (IsRealPass()) {
+      if (profilers_ != nullptr) {
+        VLOG(3) << "Starting real profilers for Pass " << pass_count_;
+        return profilers_->Start();
+      }
+      LOG(WARNING) << "Real profilers are null, cannot start them.";
+    }
+    return absl::OkStatus();
+  }
+
+  absl::Status PushRange(absl::string_view name) override {
+    VLOG(3) << "SimpleMultiPassTracer::PushRange: " << name;
+    return absl::OkStatus();
+  }
+
+  absl::Status PopRange() override {
+    VLOG(3) << "SimpleMultiPassTracer::PopRange";
+    return absl::OkStatus();
+  }
+
+  absl::Status StopPass() override {
+    VLOG(3) << "SimpleMultiPassTracer::StopPass. pass_count=" << pass_count_;
+    absl::Status status = absl::OkStatus();
+    if (IsRealPass()) {
+      if (profilers_ != nullptr) {
+        VLOG(3) << "Stopping real profilers for Pass " << pass_count_;
+        status = profilers_->Stop();
+      }
+    }
+    pass_count_++;
+    return status;
+  }
+
+  // ProfilerInterface methods
+  absl::Status Start() override {
+    VLOG(1) << "SimpleMultiPassTracer::Start";
+    pass_count_ = 0;
+
+    ProfileOptions options = options_;
+    options.set_enable_multipass(false);
+
+    profilers_ = std::make_unique<ProfilerCollection>(
+        tsl::profiler::CreateProfilers(options));
+
+    return absl::OkStatus();
+  }
+
+  absl::Status Stop() override {
+    VLOG(1) << "SimpleMultiPassTracer::Stop";
+    return absl::OkStatus();
+  }
+
+  absl::Status CollectData(XSpace* space) override {
+    VLOG(1) << "SimpleMultiPassTracer::CollectData";
+    if (profilers_ != nullptr) {
+      return profilers_->CollectData(space);
+    }
+    return absl::OkStatus();
+  }
+
+ private:
+  bool IsRealPass() const {
+    if (options_.device_type() == ProfileOptions::GPU ||
+        options_.device_type() == ProfileOptions::UNSPECIFIED) {
+      // GPU and UNSPECIFIED (default): Pass 0 is dummy (warmup), Pass 1 is
+      // real.
+      return pass_count_ == 1;
+    }
+    // TPU: Pass 0 is real, Pass 1 is dummy.
+    return pass_count_ == 0;
+  }
+
+  ProfileOptions options_;
+  int pass_count_ = 0;
+  std::unique_ptr<ProfilerCollection> profilers_;
+};
+
+std::unique_ptr<MultiPassProfilerInterface> CreateSimpleMultiPassTracer(
+    const ProfileOptions& options) {
+  VLOG(1) << "CreateSimpleMultiPassTracer called. device_type="
+          << options.device_type()
+          << ", enable_multipass=" << options.enable_multipass();
+  if (!options.enable_multipass()) {
+    VLOG(1) << "enable_multipass is false, not creating SimpleMultiPassTracer";
+    return nullptr;
+  }
+  // We support GPU, TPU, and UNSPECIFIED.
+  if (options.device_type() != ProfileOptions::GPU &&
+      options.device_type() != ProfileOptions::TPU &&
+      options.device_type() != ProfileOptions::UNSPECIFIED) {
+    VLOG(1) << "Device is not supported, not creating SimpleMultiPassTracer";
+    return nullptr;
+  }
+  return std::make_unique<MultiPassProfilerController>(
+      std::make_unique<SimpleMultiPassTracer>(options));
+}
+
+}  // namespace
+
+void RegisterSimpleMultiPassTracerForTest() {
+  VLOG(1) << "Registering SimpleMultiPassTracer factory.";
+  tsl::profiler::RegisterMultiPassProfilerFactory(&CreateSimpleMultiPassTracer);
+}
+
+}  // namespace profiler
+}  // namespace xla
diff --git a/third_party/xla/xla/tools/multihost_hlo_runner/simple_multipass_tracer_for_test.h b/third_party/xla/xla/tools/multihost_hlo_runner/simple_multipass_tracer_for_test.h
new file mode 100644
index 00000000000..5ce8d35e181
--- /dev/null
+++ b/third_party/xla/xla/tools/multihost_hlo_runner/simple_multipass_tracer_for_test.h
@@ -0,0 +1,30 @@
+/* Copyright 2026 The OpenXLA Authors.
+
+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.
+==============================================================================*/
+
+#ifndef XLA_TOOLS_MULTIHOST_HLO_RUNNER_SIMPLE_MULTIPASS_TRACER_FOR_TEST_H_
+#define XLA_TOOLS_MULTIHOST_HLO_RUNNER_SIMPLE_MULTIPASS_TRACER_FOR_TEST_H_
+
+namespace xla {
+namespace profiler {
+
+// Registers the SimpleMultiPassTracer factory with the TSL profiler registry.
+// SimpleMultiPassTracer coordinates multi-pass profiling for workloads on
+// accelerators.
+void RegisterSimpleMultiPassTracerForTest();
+
+}  // namespace profiler
+}  // namespace xla
+
+#endif  // XLA_TOOLS_MULTIHOST_HLO_RUNNER_SIMPLE_MULTIPASS_TRACER_FOR_TEST_H_

From f935e86c346737a17f81ad06dcbbda8251d0ba60 Mon Sep 17 00:00:00 2001
From: Vrishab Madduri <vrishab@google.com>
Date: Fri, 16 Oct 2026 14:10:20 -0700
Subject: [PATCH 0984/1000] No longer needed because host<->device transfers
 can happen directly under GPU Confidential Computing.

Reverts 5daf5f25b82bf37d136b7a60ee2afa4b586676f9

PiperOrigin-RevId: 1001089354
---
 .../xla/pjrt/gpu/se_gpu_pjrt_client_test.cc   | 103 -----------
 third_party/xla/xla/pjrt/se/BUILD             |   3 -
 .../pjrt/se/pjrt_stream_executor_client.cc    |  29 ---
 .../xla/pjrt/se/pjrt_stream_executor_client.h |  36 +---
 third_party/xla/xla/pjrt/se/se_raw_buffer.cc  | 172 +-----------------
 5 files changed, 5 insertions(+), 338 deletions(-)

diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
index 50815923322..acd040d5fea 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
@@ -1656,32 +1656,6 @@ TEST(StreamExecutorGpuClientTest, ShouldStageHostToDeviceTransfersSetToTrue) {
       *literal, LiteralUtil::CreateR1<float>(std::vector<float>(1024, 1.0f))));
 }
 
-TEST(StreamExecutorGpuClientTest, ShouldStageLargeTransfersWhenChunking) {
-  GpuClientOptions options = GetTestGpuClientOptions();
-  options.should_stage_host_to_device_transfers = true;
-  TF_ASSERT_OK_AND_ASSIGN(auto client, GetStreamExecutorGpuClient(options));
-  auto* raw_client = absl::down_cast<PjRtStreamExecutorRawClient*>(
-      absl::down_cast<CommonPjRtClient*>(client.get())->raw_client());
-
-  // Only the pointer value is used; no 1 GiB host buffer is allocated.
-  std::vector<char> host(16);
-  constexpr int64_t kOneGiB = int64_t{1} << 30;
-  ASSERT_FALSE(raw_client->IsDmaMapped(host.data(), kOneGiB));
-
-  // Without chunking, large transfers are not staged (unchanged behavior).
-  EXPECT_FALSE(
-      raw_client->ShouldStageHostToDeviceTransfers(host.data(), kOneGiB));
-  EXPECT_FALSE(
-      raw_client->ShouldStageHostToDeviceTransfers(host.data(), kOneGiB + 1));
-
-  // With chunking, staging uses bounded chunk buffers, so any size is staged.
-  raw_client->SetStagingChunkingForTesting(/*force=*/true, /*chunk_size=*/4096);
-  EXPECT_TRUE(
-      raw_client->ShouldStageHostToDeviceTransfers(host.data(), kOneGiB));
-  EXPECT_TRUE(
-      raw_client->ShouldStageHostToDeviceTransfers(host.data(), kOneGiB + 1));
-}
-
 TEST(StreamExecutorGpuClientTest, ShouldStageHostToDeviceTransfersSetToFalse) {
   GpuClientOptions options_no_staging = GetTestGpuClientOptions();
   options_no_staging.should_stage_host_to_device_transfers = false;
@@ -1713,83 +1687,6 @@ TEST(StreamExecutorGpuClientTest, ShouldStageHostToDeviceTransfersSetToFalse) {
       *literal, LiteralUtil::CreateR1<float>(std::vector<float>(1024, 1.0f))));
 }
 
-constexpr int64_t kTestStagingChunkSize = 4 * 1024;
-
-// Returns `size` bytes whose content differs between any two chunks, so that
-// misplaced or reordered chunks are detected.
-std::vector<uint8_t> MakeStagingTestPattern(int64_t size) {
-  std::vector<uint8_t> pattern(size);
-  for (int64_t i = 0; i < size; ++i) {
-    pattern[i] = static_cast<uint8_t>((i * 131) ^ (i >> 8));
-  }
-  return pattern;
-}
-
-// Copies `transfer_size` bytes host->device->host at `offset` of a device
-// buffer with chunked host staging forced on, and checks the bytes survive.
-void ExpectChunkedStagingRoundTrip(int64_t offset, int64_t transfer_size) {
-  GpuClientOptions options = GetTestGpuClientOptions();
-  options.should_stage_host_to_device_transfers = true;
-  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtClient> pjrt_client,
-                          GetStreamExecutorGpuClient(options));
-  auto* client = absl::down_cast<CommonPjRtClient*>(pjrt_client.get());
-  auto* raw_client =
-      absl::down_cast<PjRtStreamExecutorRawClient*>(client->raw_client());
-  raw_client->SetStagingChunkingForTesting(/*force=*/true,
-                                           kTestStagingChunkSize);
-
-  std::vector<uint8_t> src = MakeStagingTestPattern(transfer_size);
-  ASSERT_TRUE(
-      raw_client->ShouldStageHostToDeviceTransfers(src.data(), transfer_size));
-  TF_ASSERT_OK_AND_ASSIGN(
-      PjRtMemorySpace * memory_space,
-      client->addressable_devices()[0]->default_memory_space());
-  TF_ASSERT_OK_AND_ASSIGN(
-      PjRtRawBufferRef raw_buffer,
-      client->AllocateRawBuffer(memory_space, offset + transfer_size,
-                                /*retry_on_oom=*/true,
-                                /*allocate_after=*/{}));
-
-  TF_ASSERT_OK(
-      raw_buffer->CopyRawHostToDevice(src.data(), offset, transfer_size)
-          .Await());
-  std::vector<uint8_t> dst(transfer_size, 0);
-  TF_ASSERT_OK(
-      raw_buffer->CopyRawDeviceToHost(dst.data(), offset, transfer_size)
-          .Await());
-  EXPECT_EQ(dst, src);
-}
-
-TEST(StreamExecutorGpuClientTest, ChunkedStagingZeroBytes) {
-  ExpectChunkedStagingRoundTrip(/*offset=*/0, 0);
-}
-
-TEST(StreamExecutorGpuClientTest, ChunkedStagingOneByte) {
-  ExpectChunkedStagingRoundTrip(/*offset=*/0, 1);
-}
-
-TEST(StreamExecutorGpuClientTest, ChunkedStagingSmallerThanChunk) {
-  ExpectChunkedStagingRoundTrip(/*offset=*/0, kTestStagingChunkSize / 2);
-}
-
-TEST(StreamExecutorGpuClientTest, ChunkedStagingEqualToChunk) {
-  ExpectChunkedStagingRoundTrip(/*offset=*/0, kTestStagingChunkSize);
-}
-
-TEST(StreamExecutorGpuClientTest, ChunkedStagingMultipleOfChunk) {
-  // More chunks than staging slots, so slots are reused.
-  ExpectChunkedStagingRoundTrip(/*offset=*/0, 8 * kTestStagingChunkSize);
-}
-
-TEST(StreamExecutorGpuClientTest, ChunkedStagingNonMultipleOfChunk) {
-  ExpectChunkedStagingRoundTrip(/*offset=*/0, 10 * kTestStagingChunkSize + 123);
-}
-
-TEST(StreamExecutorGpuClientTest, ChunkedStagingWithNonzeroOffset) {
-  ExpectChunkedStagingRoundTrip(/*offset=*/123,
-                                10 * kTestStagingChunkSize + 45);
-}
-
 TEST(StreamExecutorGpuClientTest, BufferFromHostBufferPinnedMemory) {
   TF_ASSERT_OK_AND_ASSIGN(
       auto client, GetStreamExecutorGpuClient(GetTestGpuClientOptions()));
diff --git a/third_party/xla/xla/pjrt/se/BUILD b/third_party/xla/xla/pjrt/se/BUILD
index 9356860796a..47d24622f1a 100644
--- a/third_party/xla/xla/pjrt/se/BUILD
+++ b/third_party/xla/xla/pjrt/se/BUILD
@@ -275,9 +275,7 @@ cc_library(
         "//xla/service/gpu:gpu_executable_run_options",
         "//xla/stream_executor:device_address",
         "//xla/stream_executor:device_address_allocator",
-        "//xla/stream_executor:event",
         "//xla/stream_executor:stream",
-        "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor/abi:executable_abi_version",
         "//xla/stream_executor/integrations:tf_allocator_adapter",
         "//xla/tsl/concurrency:async_value",
@@ -287,7 +285,6 @@ cc_library(
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:logging",
         "//xla/tsl/platform:statusor",
-        "//xla/tsl/util:env_var",
         "//xla/util/split_proto:split_executable_and_options_writer",
         "//xla/util/split_proto:split_proto_reader",
         "@com_google_absl//absl/algorithm:container",
diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc
index 693fb557690..65ec1b2e934 100644
--- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc
+++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.cc
@@ -71,7 +71,6 @@ limitations under the License.
 #include <cstring>
 #include <functional>
 #include <initializer_list>
-#include <limits>
 #include <memory>
 #include <optional>
 #include <string>
@@ -175,7 +174,6 @@ limitations under the License.
 #include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/statusor.h"
 #include "xla/tsl/platform/threadpool.h"
-#include "xla/tsl/util/env_var.h"
 #include "xla/util.h"
 #include "xla/util/split_proto/split_executable_and_options_writer.h"
 #include "xla/util/split_proto/split_proto_reader.h"
@@ -188,30 +186,6 @@ limitations under the License.
 
 namespace xla {
 
-namespace {
-
-constexpr int64_t kDefaultStagingChunkSizeMb = 16;
-
-// Returns the chunk size, in bytes, used for chunked staging of host<->device
-// transfers. Overridable via the XLA_PJRT_GPU_CC_STAGING_CHUNK_MB environment
-// variable.
-int64_t GetStagingChunkSizeFromEnv() {
-  constexpr absl::string_view kEnvVar = "XLA_PJRT_GPU_CC_STAGING_CHUNK_MB";
-  int64_t chunk_size_mb;
-  absl::Status status = tsl::ReadInt64FromEnvVar(
-      kEnvVar, kDefaultStagingChunkSizeMb, &chunk_size_mb);
-  if (!status.ok() || chunk_size_mb <= 0 ||
-      chunk_size_mb > (std::numeric_limits<int64_t>::max() >> 20)) {
-    LOG(WARNING) << "Ignoring invalid " << kEnvVar
-                 << " (must be a positive integer number of MiB); using the "
-                 << "default of " << kDefaultStagingChunkSizeMb << " MiB.";
-    chunk_size_mb = kDefaultStagingChunkSizeMb;
-  }
-  return chunk_size_mb << 20;
-}
-
-}  // namespace
-
 template <typename MemorySpaceKind>
 static bool IsMemorySpaceKind(const PjRtMemorySpace* memory_space) {
   return memory_space->kind_id() == MemorySpaceKind::kKindId;
@@ -323,9 +297,6 @@ PjRtStreamExecutorRawClient::PjRtStreamExecutorRawClient(
       should_stage_host_to_device_transfers_(
           should_stage_host_to_device_transfers),
       confidential_computing_enabled_(confidential_computing_enabled),
-      staging_chunk_size_(confidential_computing_enabled
-                              ? GetStagingChunkSizeFromEnv()
-                              : kDefaultStagingChunkSizeMb << 20),
       executor_(executor),
       gpu_run_options_(std::move(gpu_run_options)),
       compile_thread_pool_(
diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h
index be720193ab8..071209962ea 100644
--- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h
+++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h
@@ -218,36 +218,11 @@ class PjRtStreamExecutorRawClient : public PjRtRawClient {
     if (confidential_computing_enabled_) {
       return true;
     }
-    // Allocating a multi-gigabyte pinned staging buffer can be very slow, in
-    // which case staging is probably worse than not staging. This cap does not
-    // apply to chunked staging, which only uses staging_chunk_size() buffers.
-    // TODO(phawkins): add chunking for transfers when chunked staging is off.
+    // Allocating multi-gigabyte pinned buffers can be very slow. In that case,
+    // using a staging buffer is probably worse than not using one.
+    // TODO(phawkins): add chunking for transfers.
     return should_stage_host_to_device_transfers_ &&
-           (size < (int64_t{1} << 30) || ShouldUseChunkedStaging()) &&
-           !IsDmaMapped(data, size);
-  }
-
-  // Returns true if staged transfers should be split into at most
-  // `staging_chunk_size()`-byte chunks that are pipelined through a bounded
-  // number of staging buffers, with host copies done on the transfer's work
-  // thread instead of in stream host callbacks. Applies to staged transfers of
-  // every size (a transfer no larger than a chunk is a single chunk). Only
-  // enabled in Confidential Computing mode (where every transfer is staged),
-  // or when forced for testing.
-  bool ShouldUseChunkedStaging() const {
-    return confidential_computing_enabled_ ||
-           force_staging_chunking_for_testing_;
-  }
-
-  int64_t staging_chunk_size() const { return staging_chunk_size_; }
-
-  // Test-only: enables chunked staging (when `force` is true) even when
-  // Confidential Computing is disabled, using `chunk_size`-byte chunks. Must be
-  // called before any transfer is issued on this client.
-  void SetStagingChunkingForTesting(bool force, int64_t chunk_size) {
-    CHECK_GT(chunk_size, 0);
-    force_staging_chunking_for_testing_ = force;
-    staging_chunk_size_ = chunk_size;
+           size < (int64_t{1} << 30) && !IsDmaMapped(data, size);
   }
 
   tsl::AsyncValueRef<PjRtExecutable> ToAsyncExecutable(
@@ -399,9 +374,6 @@ class PjRtStreamExecutorRawClient : public PjRtRawClient {
   // transfer via pinned memory.
   bool should_stage_host_to_device_transfers_;
   bool confidential_computing_enabled_ = false;
-  // Chunk size for chunked staging; see ShouldUseChunkedStaging().
-  int64_t staging_chunk_size_;
-  bool force_staging_chunking_for_testing_ = false;
 
   se::StreamExecutor* absl_nonnull executor_;
   std::unique_ptr<gpu::GpuExecutableRunOptions> gpu_run_options_;
diff --git a/third_party/xla/xla/pjrt/se/se_raw_buffer.cc b/third_party/xla/xla/pjrt/se/se_raw_buffer.cc
index 23401ed95b3..0ac1b71c73f 100644
--- a/third_party/xla/xla/pjrt/se/se_raw_buffer.cc
+++ b/third_party/xla/xla/pjrt/se/se_raw_buffer.cc
@@ -15,19 +15,16 @@ limitations under the License.
 
 #include "xla/pjrt/se/se_raw_buffer.h"
 
-#include <algorithm>
 #include <cstdint>
 #include <cstring>
 #include <memory>
 #include <utility>
-#include <vector>
 
 #include "absl/functional/any_invocable.h"
 #include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
-#include "absl/types/span.h"
 #include "xla/pjrt/async_work_runner.h"
 #include "xla/pjrt/common_pjrt_client.h"
 #include "xla/pjrt/device_event.h"
@@ -41,9 +38,7 @@ limitations under the License.
 #include "xla/pjrt/se/pjrt_stream_executor_client.h"
 #include "xla/pjrt/se/tracked_device_buffer.h"
 #include "xla/stream_executor/device_address.h"
-#include "xla/stream_executor/event.h"
 #include "xla/stream_executor/stream.h"
-#include "xla/stream_executor/stream_executor.h"
 #include "xla/tsl/concurrency/async_value.h"
 #include "xla/tsl/concurrency/async_value_ref.h"
 #include "xla/tsl/concurrency/ref_count.h"
@@ -74,138 +69,6 @@ GetOrCreateAllocationReadyEvent(const xla::RawSEDeviceMemory& device_buffer,
 
 namespace xla {
 
-namespace {
-
-// Maximum number of chunks of a single chunked staged transfer in flight.
-constexpr int64_t kMaxStagingChunksInFlight = 4;
-
-// A pinned host staging buffer for one chunk, and an event recorded after the
-// most recent DMA that used it.
-struct StagingSlot {
-  std::shared_ptr<void> buffer;
-  std::unique_ptr<se::Event> dma_done;
-};
-
-// Allocates the staging slots for a chunked transfer of `transfer_size` bytes.
-// Each slot holds one chunk, so transfers smaller than `chunk_size` only pin
-// `transfer_size` bytes.
-absl::StatusOr<std::vector<StagingSlot>> AllocateStagingSlots(
-    PjRtStreamExecutorRawClient* client, LocalDeviceState* local_device,
-    se::Stream* stream, int64_t transfer_size, int64_t chunk_size) {
-  HostMemoryAllocator::AllocateOptions alloc_opts;
-  alloc_opts.numa_node = stream->parent()->numa_node();
-  alloc_opts.local_device_id = local_device->local_device_id();
-  const int64_t slot_size = std::min(chunk_size, transfer_size);
-  std::vector<StagingSlot> slots(std::min(
-      kMaxStagingChunksInFlight, CeilOfRatio(transfer_size, chunk_size)));
-  for (StagingSlot& slot : slots) {
-    slot.buffer =
-        client->GetHostMemoryAllocator()->Allocate(slot_size, alloc_opts);
-    if (slot.buffer == nullptr) {
-      return ResourceExhausted(
-          "Failed to allocate a %d-byte pinned host staging chunk for a "
-          "chunked host<->device transfer. The pinned host pool may be "
-          "exhausted or fragmented (see XLA_PJRT_GPU_HOST_MEMORY_LIMIT_GB and "
-          "XLA_PJRT_GPU_CC_STAGING_CHUNK_MB).",
-          slot_size);
-    }
-    ABSL_ASSIGN_OR_RETURN(slot.dma_done, stream->parent()->CreateEvent());
-  }
-  return slots;
-}
-
-// Threading: the chunked helpers below run inside `run_transfer` on the
-// client's `async_work_runner()` and do the host<->staging copies on that
-// thread instead of in stream host callbacks. They block that thread on
-// `se::Event`s while chunk DMAs complete (for D2H, on every transfer). This is
-// safe because (1) chunking is only enabled for GPU Confidential Computing (or
-// forced in tests), and the SE GPU client always uses the default
-// `MakeUnboundedAsyncWorkRunner`, which starts a new thread when none are
-// idle, so a blocked transfer cannot starve other work; and (2) every awaited
-// event is recorded after already-enqueued stream work, so it completes on GPU
-// progress alone and never waits on work from this runner. Clients that supply
-// a bounded runner would hold one thread per in-flight staged transfer.
-
-// Copies `transfer_size` bytes from host `src` to device `dst` in
-// `chunk_size`-byte chunks staged through `slots`. Host copies into the
-// staging slots run on the calling thread and overlap with the DMAs of
-// previously enqueued chunks. Blocks the calling thread until each slot's
-// previous DMA completed before reusing it.
-absl::Status ChunkedStagedHostToDevice(se::Stream* stream, const void* src,
-                                       se::DeviceAddressBase dst,
-                                       int64_t transfer_size,
-                                       int64_t chunk_size,
-                                       absl::Span<StagingSlot> slots) {
-  const int64_t num_chunks = CeilOfRatio(transfer_size, chunk_size);
-  const int64_t num_slots = slots.size();
-  for (int64_t i = 0; i < num_chunks; ++i) {
-    const int64_t offset = i * chunk_size;
-    const int64_t size = std::min(chunk_size, transfer_size - offset);
-    StagingSlot& slot = slots[i % num_slots];
-    if (i >= num_slots) {
-      ABSL_RETURN_IF_ERROR(slot.dma_done->Synchronize());
-    }
-    {
-      tsl::profiler::TraceMe trace("H2D Copy To Staging Chunk");
-      std::memcpy(slot.buffer.get(), static_cast<const char*>(src) + offset,
-                  size);
-    }
-    se::DeviceAddressBase dst_chunk = dst.GetByteSlice(offset, size);
-    ABSL_RETURN_IF_ERROR(stream->Memcpy(&dst_chunk, slot.buffer.get(), size));
-    ABSL_RETURN_IF_ERROR(stream->RecordEvent(slot.dma_done.get()));
-  }
-  return absl::OkStatus();
-}
-
-// Copies `transfer_size` bytes from device `src` to host `dst` in
-// `chunk_size`-byte chunks staged through `slots`. Keeps up to `slots.size()`
-// chunk DMAs in flight while the calling thread copies completed chunks out of
-// the staging slots. Returns once `dst` is fully written.
-absl::Status ChunkedStagedDeviceToHost(se::Stream* stream,
-                                       se::DeviceAddressBase src, void* dst,
-                                       int64_t transfer_size,
-                                       int64_t chunk_size,
-                                       absl::Span<StagingSlot> slots) {
-  const int64_t num_chunks = CeilOfRatio(transfer_size, chunk_size);
-  const int64_t num_slots = slots.size();
-  auto enqueue_dma = [&](int64_t i) -> absl::Status {
-    const int64_t offset = i * chunk_size;
-    const int64_t size = std::min(chunk_size, transfer_size - offset);
-    StagingSlot& slot = slots[i % num_slots];
-    ABSL_RETURN_IF_ERROR(stream->Memcpy(slot.buffer.get(),
-                                   src.GetByteSlice(offset, size), size));
-    return stream->RecordEvent(slot.dma_done.get());
-  };
-  for (int64_t i = 0; i < num_slots; ++i) {
-    ABSL_RETURN_IF_ERROR(enqueue_dma(i));
-  }
-  for (int64_t i = 0; i < num_chunks; ++i) {
-    const int64_t offset = i * chunk_size;
-    const int64_t size = std::min(chunk_size, transfer_size - offset);
-    StagingSlot& slot = slots[i % num_slots];
-    ABSL_RETURN_IF_ERROR(slot.dma_done->Synchronize());
-    {
-      tsl::profiler::TraceMe trace("D2H Copy From Staging Chunk");
-      std::memcpy(static_cast<char*>(dst) + offset, slot.buffer.get(), size);
-    }
-    if (i + num_slots < num_chunks) {
-      ABSL_RETURN_IF_ERROR(enqueue_dma(i + num_slots));
-    }
-  }
-  return absl::OkStatus();
-}
-
-// On failure, chunk DMAs may still be accessing the staging slots, so drain
-// `stream` before the slots can be released.
-absl::Status DrainStreamOnError(se::Stream* stream, absl::Status status) {
-  if (!status.ok()) {
-    stream->BlockHostUntilDone().IgnoreError();
-  }
-  return status;
-}
-
-}  // namespace
-
 PjRtStreamExecutorDeviceEventPromise::PjRtStreamExecutorDeviceEventPromise(
     PjRtStreamExecutorRawClient* client, LocalDeviceState* local_device,
     AsyncWorkRunner* async_work_runner)
@@ -270,9 +133,6 @@ PjRtStreamExecutorRawBuffer::CopyRawHostToDeviceAndReturnEvent(
       sub_buffer = sub_buffer.GetByteSlice(offset, transfer_size);
     }
     std::shared_ptr<void> staging_buffer;
-    // Only set for chunked staging. Shared with the cleanup callback so that
-    // the slots outlive the drain below if AllocateAndRecordEvent fails.
-    std::shared_ptr<std::vector<StagingSlot>> staging_slots;
     auto status = [&]() -> absl::Status {
       ABSL_RETURN_IF_ERROR(client->WaitForAllocation(stream, *buf));
       if (transfer_size > 0) {
@@ -282,17 +142,6 @@ PjRtStreamExecutorRawBuffer::CopyRawHostToDeviceAndReturnEvent(
                 "host_memory_allocator should be initialized for "
                 "staging buffer transfer.");
           }
-          if (client->ShouldUseChunkedStaging()) {
-            const int64_t chunk_size = client->staging_chunk_size();
-            ABSL_ASSIGN_OR_RETURN(std::vector<StagingSlot> slots,
-                             AllocateStagingSlots(client, local_device, stream,
-                                                  transfer_size, chunk_size));
-            staging_slots =
-                std::make_shared<std::vector<StagingSlot>>(std::move(slots));
-            return ChunkedStagedHostToDevice(stream, src, sub_buffer,
-                                             transfer_size, chunk_size,
-                                             absl::MakeSpan(*staging_slots));
-          }
           HostMemoryAllocator::AllocateOptions alloc_opts;
           alloc_opts.numa_node = stream->parent()->numa_node();
           alloc_opts.local_device_id = local_device->local_device_id();
@@ -326,18 +175,11 @@ PjRtStreamExecutorRawBuffer::CopyRawHostToDeviceAndReturnEvent(
       status = client->AllocateAndRecordEvent(
           device_event, local_device, stream,
           "PjRtStreamExecutorRawBuffer::CopyRawHostToDevice",
-          [staging_buffer = std::move(staging_buffer),
-           staging_slots]() mutable {
+          [staging_buffer = std::move(staging_buffer)]() mutable {
             staging_buffer.reset();
-            staging_slots.reset();
           });
     }
     if (!status.ok()) {
-      // Chunk DMAs may still be reading the staging slots, so drain before
-      // releasing them.
-      if (staging_slots != nullptr) {
-        stream->BlockHostUntilDone().IgnoreError();
-      }
       client->SetEventAsError(device_event, status);
     }
   };
@@ -384,18 +226,6 @@ PjRtStreamExecutorRawBuffer::CopyRawDeviceToHostAndReturnEvent(
                 "host_memory_allocator should be initialized for "
                 "staging buffer transfer.");
           }
-          if (client->ShouldUseChunkedStaging()) {
-            // All chunk DMAs and host copies complete before this returns, so
-            // the staging slots can be released when it goes out of scope.
-            const int64_t chunk_size = client->staging_chunk_size();
-            ABSL_ASSIGN_OR_RETURN(std::vector<StagingSlot> staging_slots,
-                             AllocateStagingSlots(client, local_device, stream,
-                                                  transfer_size, chunk_size));
-            return DrainStreamOnError(
-                stream, ChunkedStagedDeviceToHost(
-                            stream, sub_buffer, dst, transfer_size, chunk_size,
-                            absl::MakeSpan(staging_slots)));
-          }
           HostMemoryAllocator::AllocateOptions alloc_opts;
           alloc_opts.numa_node = stream->parent()->numa_node();
           alloc_opts.local_device_id = local_device->local_device_id();

From 9dd1241c2a15417ac0c9eb3940f67bd7639b4646 Mon Sep 17 00:00:00 2001
From: Viktor Moros <vgm@google.com>
Date: Fri, 16 Oct 2026 14:14:48 -0700
Subject: [PATCH 0985/1000] [XLA Highway] Add 8x32 byte transpose microkernel

PiperOrigin-RevId: 1001092197
---
 third_party/xla/xla/pjrt/transpose_kernels.cc | 50 ++++++++++++++++++-
 1 file changed, 49 insertions(+), 1 deletion(-)

diff --git a/third_party/xla/xla/pjrt/transpose_kernels.cc b/third_party/xla/xla/pjrt/transpose_kernels.cc
index 224ef3b6c52..af51aae8c72 100644
--- a/third_party/xla/xla/pjrt/transpose_kernels.cc
+++ b/third_party/xla/xla/pjrt/transpose_kernels.cc
@@ -394,6 +394,51 @@ HWY_INLINE void BuildVecs256RectCPS(D d, DHalf d_half, const char* a,
   BuildVecs256RectCPS<0, T, bs>(d, d_half, a, lda, cont);
 }
 
+// Flatten inner Unpack/StoreLane helpers into this function while keeping it
+// out-of-line so calling it 4 times in TransposeMicroKernel32x32 doesn't spill
+// vector registers.
+template <typename T>
+HWY_NOINLINE HWY_FLATTEN void TransposeMicroKernel8x32(
+    const char* HWY_RESTRICT a, int64_t lda, char* HWY_RESTRICT b,
+    int64_t ldb) {
+  using D = hn::FixedTag<uint8_t, 32>;
+  const D d;
+  using D64 = hn::Repartition<uint64_t, D>;
+  const D64 d64;
+  LoadRowsCPS<0, T, 8, D, 32>(d, a, lda, [&](auto... loads) HWY_ATTR {
+    InterleaveWithinBlocksCPS<1, 1, 8>(
+        [&](auto... interleaved) HWY_ATTR {
+          ForEachIndexed(
+              [&](auto idx_c, auto&& vec) HWY_ATTR {
+                constexpr size_t i = decltype(idx_c)::value;
+                auto v64 = hn::BitCast(d64, vec);
+                StoreLane<D64, sizeof(uint64_t), 0>(d64, b + (2 * i) * ldb,
+                                                    v64);
+                StoreLane<D64, sizeof(uint64_t), 1>(d64, b + (2 * i + 1) * ldb,
+                                                    v64);
+                StoreLane<D64, sizeof(uint64_t), 2>(d64, b + (2 * i + 16) * ldb,
+                                                    v64);
+                StoreLane<D64, sizeof(uint64_t), 3>(d64, b + (2 * i + 17) * ldb,
+                                                    v64);
+              },
+              interleaved...);
+        },
+        loads...);
+  });
+}
+
+// Transposes a 32x32 block of bytes as four 8x32 row strips.
+template <typename T>
+HWY_FLATTEN void TransposeMicroKernel32x32(const char* HWY_RESTRICT a,
+                                           int64_t lda, char* HWY_RESTRICT b,
+                                           int64_t ldb) {
+  static_assert(sizeof(T) == 1);
+  TransposeMicroKernel8x32<T>(a, lda, b, ldb);
+  TransposeMicroKernel8x32<T>(a + 8 * lda, lda, b + 8 * sizeof(T), ldb);
+  TransposeMicroKernel8x32<T>(a + 16 * lda, lda, b + 16 * sizeof(T), ldb);
+  TransposeMicroKernel8x32<T>(a + 24 * lda, lda, b + 24 * sizeof(T), ldb);
+}
+
 template <typename T, int bs>
 HWY_FLATTEN void TransposeMicroKernel256Rect(const char* HWY_RESTRICT a,
                                              int64_t lda, char* HWY_RESTRICT b,
@@ -440,7 +485,10 @@ HWY_FLATTEN void TransposeMicroKernel(const char* HWY_RESTRICT a, int64_t lda,
     TransposeMicroKernelScalar<T, bs>(a, lda, b, ldb);
   } else {
 #if HWY_MIN_BYTES >= 32
-    if constexpr (row_bytes == 2 * kBlockBytes) {
+    if constexpr (sizeof(T) == 1 && bs == 32) {
+      TransposeMicroKernel32x32<T>(a, lda, b, ldb);
+      return;
+    } else if constexpr (row_bytes == 2 * kBlockBytes) {
       TransposeMicroKernel256Square<T, bs>(a, lda, b, ldb);
       return;
     } else if constexpr (row_bytes == kBlockBytes) {

From f6a32b7a75fe2e9c61ef4d8f7be75d115a201876 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Longeri?= <tlongeri@google.com>
Date: Fri, 16 Oct 2026 15:29:57 -0700
Subject: [PATCH 0986/1000] Roll forward of "[Mosaic] Fix
 tpu.memref_squeeze..." with relaxed verifier

The verifier now ignores layout mismatch if the input or result layout is strided so that older serialized artifacts don't fail. This is a short-term temporary hack that we will remove once artifacts are updated.

Note: The artifacts in question were outside the Mosaic compatibility window
https://docs.jax.dev/en/latest/export/export.html#compatibility-guarantees

Reverts a636085aba4bc50fb13a3113322b7396ca49d945

PiperOrigin-RevId: 1001139025
---
 .../xla/xla/mosaic/dialect/tpu/tpu_dialect.cc |   4 +-
 .../xla/xla/mosaic/dialect/tpu/tpu_ops.cc     | 156 +++++++++++-------
 .../xla/xla/mosaic/dialect/tpu/tpu_ops.td     |   5 +-
 .../xla/xla/mosaic/dialect/tpu/util.cc        |   8 +-
 third_party/xla/xla/mosaic/dialect/tpu/util.h |   2 +-
 5 files changed, 107 insertions(+), 68 deletions(-)

diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.cc b/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.cc
index b8ad6c6bf88..6bb8efcaae3 100644
--- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.cc
+++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_dialect.cc
@@ -260,11 +260,11 @@ struct MemRefDimOfSqueeze : public OpRewritePattern<memref::DimOp> {
     }
     MemRefType source_type = squeeze_op.getInput().getType();
     FAILUREOR_ASSIGN_OR_RETURN(
-        SmallVector<int> squeezed,
+        SmallVector<int64_t> squeezed,
         computeSqueezedDimsChecked(squeeze_op, source_type.getShape(),
                                    result_type.getShape()));
     int64_t source_dim = dim;
-    for (int squeezed_dim : squeezed) {
+    for (int64_t squeezed_dim : squeezed) {
       if (squeezed_dim <= source_dim) {
         ++source_dim;
       }
diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc
index 0348b49c9eb..d71c11febc3 100644
--- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc
+++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.cc
@@ -29,6 +29,7 @@ limitations under the License.
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/FloatingPointMode.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Casting.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
@@ -45,6 +46,7 @@ limitations under the License.
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/IRMapping.h"
+#include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OperationSupport.h"
@@ -661,87 +663,121 @@ FailureOr<TiledLayoutAttr> SharedMemRefSliceOp::inferResultLayout(
 }
 
 LogicalResult MemRefSqueezeOp::verify() {
-  MemRefType source_type = getInput().getType();
-  MemRefType target_type = getType();
+  MemRefType input_type = getInput().getType();
+  MemRefType result_type = getType();
 
-  if (target_type.getMemorySpace() != source_type.getMemorySpace()) {
+  if (result_type.getMemorySpace() != input_type.getMemorySpace()) {
     return emitOpError("Memory spaces do not match.");
   }
 
-  if (target_type.getElementType() != source_type.getElementType()) {
+  if (result_type.getElementType() != input_type.getElementType()) {
     return emitOpError("Element types don't match.");
   }
 
-  auto source_shape = source_type.getShape();
-  auto target_shape = target_type.getShape();
+  const ArrayRef<int64_t> input_shape = input_type.getShape();
+  const ArrayRef<int64_t> result_shape = result_type.getShape();
+  // NOTE: In cases where there is flexibility on which dimension to squeeze,
+  //   such as 1x1x2 -> 1x2 (can squeeze dimension 1 or 2), this may choose
+  //   dimensions that have padding and fail in inferResultLayout.
+  // TODO(tlongeri): Unify logic in inferMemRefReshape such that we don't
+  //  need computeSqueezedDimsChecked at all.
   FAILUREOR_ASSIGN_OR_RETURN(
-      auto squeezed,
-      computeSqueezedDimsChecked(*this, source_shape, target_shape));
-  if (squeezed.empty() && source_shape != target_shape) {
+      const SmallVector<int64_t> squeezed,
+      computeSqueezedDimsChecked(*this, input_shape, result_shape));
+  if (squeezed.empty() && input_shape != result_shape) {
     return emitOpError(
         "Source and target shapes must be the same if no dimensions are "
         "squeezed.");
   }
 
-  auto source_layout = source_type.getLayout();
-  auto target_layout = target_type.getLayout();
-  bool has_tiled_layout = isa<TiledLayoutAttr>(source_layout);
-  if (has_tiled_layout != isa<TiledLayoutAttr>(target_layout)) {
-    return emitOpError(
-        "Either both src and dst or none of them should have a tiled layout");
+  if (isa<StridedLayoutAttr>(input_type.getLayout()) ||
+      isa<StridedLayoutAttr>(result_type.getLayout())) {
+    // TODO(b/568373460): This is for compatibility with older serialized
+    //   artifacts outside of the guaranteed compatibility window. Remove once
+    //   all artifacts are updated.
+    return success();
   }
-  if (has_tiled_layout) {
-    return verifyTiling();
+
+  FAILUREOR_ASSIGN_OR_RETURN(
+      const MemRefLayoutAttrInterface expected_result_layout,
+      inferResultLayout(input_type.getLayout(), squeezed,
+                        [&]() { return emitOpError(); }));
+  if (result_type.getLayout() != expected_result_layout) {
+    return emitOpError("Expected result layout to be ")
+           << expected_result_layout;
   }
   return success();
 }
 
-mlir::InFlightDiagnostic MemRefSqueezeOp::verifyTiling() {
-  MemRefType source_type = getInput().getType();
-  auto source_shape = source_type.getShape();
-  auto target_shape = getType().getShape();
-  auto squeezed_or =
-      computeSqueezedDimsChecked(*this, source_shape, target_shape);
-  if (failed(squeezed_or)) {
-    return {};
-  }
-  auto& squeezed = squeezed_or.value();
+FailureOr<MemRefLayoutAttrInterface> MemRefSqueezeOp::inferResultLayout(
+    const MemRefLayoutAttrInterface input_layout,
+    const ArrayRef<int64_t> squeezed,
+    const function_ref<InFlightDiagnostic()> emit_error) {
+  // TODO(tlongeri): Make MemRefReshapeOp::inferResultLayout more general and
+  // use that instead.
+  MLIRContext* const ctx = input_layout.getContext();
+  if (auto tiled_layout = dyn_cast<TiledLayoutAttr>(input_layout)) {
+    const int64_t input_rank = tiled_layout.getRank();
+    const int64_t result_rank = input_rank - squeezed.size();
 
-  auto tiles = cast<TiledLayoutAttr>(source_type.getLayout()).getTiles();
-  switch (tiles.size()) {
-    case 0:
-      break;
-    case 1: {
-      auto tile = tiles.front();
-      auto tile_dims = tile.dimensions();
-      int first_tiled = source_shape.size() - tile_dims.size();
-      for (int dim : squeezed) {
-        if (dim >= first_tiled) {
-          int tile_idx = dim - first_tiled;
-          if (tile_idx < 0 || tile_idx >= static_cast<int>(tile_dims.size())) {
-            return emitOpError() << "Internal error: tile index out of bounds.";
-          }
-          if (tile_dims[tile_idx] != 1) {
-            return emitOpError()
-                   << "All tiled squeezed dimensions must be of size 1.";
-          }
-        }
-      }
-      break;
-    }
-    default: {
-      auto first_tile = tiles.front();
-      for (int dim : squeezed) {
-        int first_tiled = source_shape.size() - first_tile.dimensions().size();
-        if (dim >= first_tiled) {
-          return emitOpError() << "When multiple tiles are present, no tiled "
-                                  "dimensions can be squeezed.";
-        }
+    SmallVector<int64_t> tile_strides;
+    tile_strides.reserve(result_rank);
+    for (int64_t i = 0; i < input_rank; ++i) {
+      if (!llvm::is_contained(squeezed, i)) {
+        tile_strides.push_back(tiled_layout.getTileStrides()[i]);
       }
     }
-  }
 
-  return {};
+    const ArrayRef<xla::Tile> tiles = tiled_layout.getTiles();
+
+    if (tiles.size() == 1 && tiles[0].dimensions().size() == 2 &&
+        tiles[0].dimension(0) == 1 &&
+        !llvm::is_contained(squeezed, input_rank - 1) &&
+        llvm::is_contained(squeezed, input_rank - 2) && result_rank >= 2) {
+      // For legacy reasons, for T(1, B) that squeezes the 2nd minor, infer
+      // T(1, B) instead of T(B) like in the code below.
+      return MemRefLayoutAttrInterface(TiledLayoutAttr::get(
+          ctx, {xla::Tile({1, tiles[0].dimension(1)})}, tile_strides));
+    }
+
+    SmallVector<xla::Tile> result_tiles;
+    result_tiles.reserve(tiles.size());
+    // Perform expansion as in TiledLayoutAttr::getExpandedShape, maintaining a
+    // mapping from expanded shape dimension to corresponding input dimension.
+    SmallVector<int64_t> expanded_dims = llvm::to_vector(
+        llvm::iota_range<int64_t>(0, input_rank, /*Inclusive=*/false));
+    for (const xla::Tile& input_tile : tiles) {
+      const int64_t tile_rank = input_tile.dimensions().size();
+      const int64_t expanded_rank = expanded_dims.size();
+      SmallVector<int64_t> tile;
+      for (int64_t i = 0; i < tile_rank; ++i) {
+        const int64_t input_dim = expanded_dims[expanded_rank - tile_rank + i];
+        if (llvm::is_contained(squeezed, input_dim)) {
+          if (input_tile.dimension(i) != 1) {
+            return emit_error() << "Dimension " << input_dim
+                                << " is padded but is squeezed.";
+          }
+        } else {
+          tile.push_back(input_tile.dimension(i));
+        }
+        expanded_dims.push_back(input_dim);
+      }
+      if (!tile.empty()) {
+        result_tiles.emplace_back(tile);
+      }
+    }
+    return MemRefLayoutAttrInterface(
+        TiledLayoutAttr::get(ctx, result_tiles, tile_strides));
+  }
+  if (auto affine_map_attr = dyn_cast<AffineMapAttr>(input_layout);
+      affine_map_attr && affine_map_attr.isIdentity()) {
+    const int64_t input_rank = affine_map_attr.getValue().getNumInputs();
+    const int64_t result_rank = input_rank - squeezed.size();
+    // Untiled squeeze
+    return MemRefLayoutAttrInterface(AffineMapAttr::get(
+        AffineMap::getMultiDimIdentityMap(result_rank, ctx)));
+  }
+  return emit_error() << "Only tiled or identity layouts supported.";
 }
 
 // Rewrites
@@ -783,7 +819,7 @@ struct MemRefSqueezeFoldCast : public OpRewritePattern<MemRefSqueezeOp> {
 
     MemRefType result_type = op.getType();
     FAILUREOR_ASSIGN_OR_RETURN(
-        SmallVector<int> squeezed,
+        SmallVector<int64_t> squeezed,
         computeSqueezedDimsChecked(op, cast_result_type.getShape(),
                                    result_type.getShape()));
 
diff --git a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td
index 4786a651980..ef66888d3a5 100644
--- a/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td
+++ b/third_party/xla/xla/mosaic/dialect/tpu/tpu_ops.td
@@ -1422,7 +1422,10 @@ def TPU_MemRefSqueezeOp : TPU_Op<"memref_squeeze", [Pure]> {
     $input attr-dict `:` type($input) `->` type($result)
   }];
   let extraClassDeclaration = [{
-    mlir::InFlightDiagnostic verifyTiling();
+    static ::mlir::FailureOr<::mlir::MemRefLayoutAttrInterface> inferResultLayout(
+        ::mlir::MemRefLayoutAttrInterface source_layout,
+        ::mlir::ArrayRef<int64_t> squeezed,
+        ::mlir::function_ref<::mlir::InFlightDiagnostic()> emit_error);
   }];
   let hasVerifier = 1;
   let hasCanonicalizer = 1;
diff --git a/third_party/xla/xla/mosaic/dialect/tpu/util.cc b/third_party/xla/xla/mosaic/dialect/tpu/util.cc
index 209de7fbd05..600780f8135 100644
--- a/third_party/xla/xla/mosaic/dialect/tpu/util.cc
+++ b/third_party/xla/xla/mosaic/dialect/tpu/util.cc
@@ -35,12 +35,12 @@ limitations under the License.
 
 namespace mlir::tpu {
 
-FailureOr<SmallVector<int>> computeSqueezedDimsChecked(
+FailureOr<SmallVector<int64_t>> computeSqueezedDimsChecked(
     Operation* op, ArrayRef<int64_t> source_shape,
     ArrayRef<int64_t> target_shape) {
-  SmallVector<int> squeezed;
-  int source_index = source_shape.size() - 1;
-  int target_index = target_shape.size() - 1;
+  SmallVector<int64_t> squeezed;
+  int64_t source_index = source_shape.size() - 1;
+  int64_t target_index = target_shape.size() - 1;
 
   while (source_index >= 0 || target_index >= 0) {
     int64_t target_dim = (target_index >= 0) ? target_shape[target_index] : -1;
diff --git a/third_party/xla/xla/mosaic/dialect/tpu/util.h b/third_party/xla/xla/mosaic/dialect/tpu/util.h
index 18aab9e1b65..ec3896f5017 100644
--- a/third_party/xla/xla/mosaic/dialect/tpu/util.h
+++ b/third_party/xla/xla/mosaic/dialect/tpu/util.h
@@ -167,7 +167,7 @@ std::string shapeToString(const T& shape) {
 
 // Computes the dimensions that were squeezed from the source shape to match the
 // target shape. Returns the dimensions in increasing order.
-FailureOr<SmallVector<int>> computeSqueezedDimsChecked(
+FailureOr<SmallVector<int64_t>> computeSqueezedDimsChecked(
     Operation* op, ArrayRef<int64_t> source_shape,
     ArrayRef<int64_t> target_shape);
 

From 55310b93b77ae7741a26b9b8a463e2b8ef11f45e Mon Sep 17 00:00:00 2001
From: Parker Schuh <parkers@google.com>
Date: Fri, 16 Oct 2026 15:33:28 -0700
Subject: [PATCH 0987/1000] Add an executable checks option to deserialization
 and add LaunchOnDevice.

PiperOrigin-RevId: 1001140969
---
 .../xla/xla/pjrt/common_pjrt_client.cc        | 12 +++++++---
 third_party/xla/xla/pjrt/common_pjrt_client.h | 22 +++++++++++++------
 .../xla/xla/pjrt/cpu/cpu_pjrt_compiler.cc     |  3 ++-
 .../xla/xla/pjrt/cpu/cpu_pjrt_compiler.h      |  3 ++-
 third_party/xla/xla/pjrt/gpu/BUILD            |  2 --
 .../xla/xla/pjrt/gpu/se_gpu_pjrt_compiler.cc  |  3 ++-
 .../xla/xla/pjrt/gpu/se_gpu_pjrt_compiler.h   |  3 ++-
 third_party/xla/xla/pjrt/pjrt_compiler.cc     |  8 ++++---
 third_party/xla/xla/pjrt/pjrt_compiler.h      |  6 +++--
 third_party/xla/xla/pjrt/raw_pjrt_client.h    |  9 ++++++++
 .../se/pjrt_stream_executor_client_test.cc    |  3 ++-
 11 files changed, 52 insertions(+), 22 deletions(-)

diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.cc b/third_party/xla/xla/pjrt/common_pjrt_client.cc
index 578aad5f869..2e84481de57 100644
--- a/third_party/xla/xla/pjrt/common_pjrt_client.cc
+++ b/third_party/xla/xla/pjrt/common_pjrt_client.cc
@@ -513,7 +513,7 @@ CommonPjRtClient::DeserializeExecutable(absl::string_view serialized,
                    GetTopologyDescription());
   return PjRtDeserializeExecutable(
       *topology, std::make_unique<riegeli::StringReader<>>(serialized),
-      std::move(options));
+      std::move(options), executable_compatibility_check_on_deserialization());
 }
 absl::StatusOr<std::unique_ptr<PjRtExecutable>>
 CommonPjRtClient::DeserializeExecutable(const absl::Cord& serialized,
@@ -524,7 +524,7 @@ CommonPjRtClient::DeserializeExecutable(const absl::Cord& serialized,
                    GetTopologyDescription());
   return PjRtDeserializeExecutable(
       *topology, std::make_unique<riegeli::CordReader<>>(&serialized),
-      std::move(options));
+      std::move(options), executable_compatibility_check_on_deserialization());
 }
 
 absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
@@ -600,7 +600,10 @@ CommonPjRtClient::LoadInternal(std::shared_ptr<PjRtExecutable> executable,
                                const LoadOptions& load_options, bool dump) {
   ABSL_ASSIGN_OR_RETURN(CompileOptions compile_options,
                    executable->GetCompileOptions());
-  ABSL_RETURN_IF_ERROR(compile_options.ApplyAllOptionOverrides());
+  // TODO(parkers): Apply options overrides consistently.
+  if (!IsTpuId(platform_id())) {
+    ABSL_RETURN_IF_ERROR(compile_options.ApplyAllOptionOverrides());
+  }
   return LoadInternal(executable, std::move(compile_options), nullptr,
                       load_options, dump);
 }
@@ -4716,6 +4719,9 @@ CommonPjRtClientImpl::CommonPjRtClientImpl(
   set_bool_attr_from_plugin_attrs("dump_on_deserialize", dump_on_deserialize_);
   set_bool_attr_from_plugin_attrs("should_stage_host_to_device_transfers",
                                   should_stage_host_to_device_transfers_);
+  set_bool_attr_from_plugin_attrs(
+      "executable_compatibility_check_on_deserialization",
+      executable_compatibility_check_on_deserialization_);
 }
 
 void CommonPjRtClientImpl::AttachDevices(
diff --git a/third_party/xla/xla/pjrt/common_pjrt_client.h b/third_party/xla/xla/pjrt/common_pjrt_client.h
index 1c615e4486f..14d0e704d94 100644
--- a/third_party/xla/xla/pjrt/common_pjrt_client.h
+++ b/third_party/xla/xla/pjrt/common_pjrt_client.h
@@ -84,6 +84,9 @@ class CommonPjRtClient : public PjRtClient {
   virtual bool supports_two_phase_launch() const { return true; }
   virtual bool dump_on_deserialize() const { return false; }
   virtual bool should_stage_host_to_device_transfers() const { return false; }
+  virtual bool executable_compatibility_check_on_deserialization() const {
+    return true;
+  }
   // Returns true if we should skip the staging buffer during ToLiteral.
   bool ShouldDoDirectTransfer(const MutableLiteralBase& literal,
                               const Shape& shape,
@@ -113,21 +116,22 @@ class CommonPjRtClient : public PjRtClient {
   }
 
   // Backend specific handlers for when an oom is detected during execute.
-  virtual void CallOomHandlers() const {}
+  void CallOomHandlers() const { raw_client()->CallOomHandlers(); }
 
   virtual PjRtDynamicShapeKind GetDynamicShapeKind(
       int memory_space_kind_id) const;
 
-  virtual void LaunchOnDevice(PjRtDevice* device,
-                              absl::AnyInvocable<void()> execute_fn) const {
+  void LaunchOnDevice(PjRtDevice* device,
+                      absl::AnyInvocable<void()> execute_fn) const {
     raw_client()->LaunchOnDevice(device->local_device_id(),
                                  std::move(execute_fn));
   }
 
-  virtual bool ShouldRetryOnOom(int attempts, PjRtDevice* device,
-                                PjRtExecutableLoadState* load_state,
-                                absl::Status prepare_status) {
-    return false;
+  bool ShouldRetryOnOom(int attempts, PjRtDevice* device,
+                        PjRtExecutableLoadState* load_state,
+                        absl::Status prepare_status) {
+    return raw_client()->ShouldRetryOnOom(attempts, device, load_state,
+                                          std::move(prepare_status));
   }
 
   virtual absl::StatusOr<PjRtExecutableLoadState::DeviceAndAssignment>
@@ -1192,6 +1196,9 @@ class CommonPjRtClientImpl : public CommonPjRtClient {
   bool should_stage_host_to_device_transfers() const override {
     return should_stage_host_to_device_transfers_;
   }
+  bool executable_compatibility_check_on_deserialization() const override {
+    return executable_compatibility_check_on_deserialization_;
+  }
 
  private:
   const PjRtPlatformId platform_id_;
@@ -1228,6 +1235,7 @@ class CommonPjRtClientImpl : public CommonPjRtClient {
   bool use_stream_based_compaction_ = false;
   bool dump_on_deserialize_ = false;
   bool should_stage_host_to_device_transfers_ = false;
+  bool executable_compatibility_check_on_deserialization_ = true;
 };
 
 // A common base class for PjRtDevice implementations that delegate to
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 1973950c439..f726b800f4f 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.cc
+++ b/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.cc
@@ -107,7 +107,8 @@ absl::StatusOr<std::unique_ptr<PjRtExecutable>>
 CpuPjRtCompiler::DeserializeExecutable(
     const PjRtTopologyDescription& topology,
     riegeli::Any<riegeli::Reader*> reader,
-    std::optional<CompileOptions>&& options) {
+    std::optional<CompileOptions>&& options,
+    bool executable_compatibility_check_on_deserialization) {
   ABSL_ASSIGN_OR_RETURN(auto* cpu_topology, GetCpuTopology(topology));
   return PjRtCpuExecutable::Deserialize(std::move(reader), *cpu_topology,
                                         std::move(options));
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 6cb1307d826..1cd0ff8cd2b 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.h
+++ b/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.h
@@ -47,7 +47,8 @@ class CpuPjRtCompiler : public PjRtCompiler {
   absl::StatusOr<std::unique_ptr<PjRtExecutable>> DeserializeExecutable(
       const PjRtTopologyDescription& topology,
       riegeli::Any<riegeli::Reader*> reader,
-      std::optional<CompileOptions>&& options) override;
+      std::optional<CompileOptions>&& options,
+      bool executable_compatibility_check_on_deserialization) override;
 
   // Deserializes a PjRtTopologyDescription from a string.
   absl::StatusOr<std::unique_ptr<PjRtTopologyDescription>>
diff --git a/third_party/xla/xla/pjrt/gpu/BUILD b/third_party/xla/xla/pjrt/gpu/BUILD
index 44ffae3ae0e..0b8d6fdfddc 100644
--- a/third_party/xla/xla/pjrt/gpu/BUILD
+++ b/third_party/xla/xla/pjrt/gpu/BUILD
@@ -892,10 +892,8 @@ cc_library(
         "//xla/stream_executor:platform",
         "//xla/stream_executor:platform_id",
         "@com_google_absl//absl/base:core_headers",
-        "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/synchronization",
-        "@llvm-project//mlir:IR",
         "@riegeli//riegeli/base:any",
         "@riegeli//riegeli/bytes:reader",
     ] + if_rocm([
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler.cc
index ce562a9a963..4e24a5d2f76 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler.cc
@@ -503,7 +503,8 @@ absl::StatusOr<std::unique_ptr<PjRtExecutable>>
 StreamExecutorGpuCompiler::DeserializeExecutable(
     const PjRtTopologyDescription& topology,
     riegeli::Any<riegeli::Reader*> reader,
-    std::optional<CompileOptions>&& options) {
+    std::optional<CompileOptions>&& options,
+    bool executable_compatibility_check_on_deserialization) {
   return StreamExecutorExecutable::Deserialize(std::move(reader), topology,
                                                std::move(options));
 }
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler.h b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler.h
index 75ab759fb96..bc3e52134b1 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler.h
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler.h
@@ -74,7 +74,8 @@ class StreamExecutorGpuCompiler : public PjRtCompiler {
   absl::StatusOr<std::unique_ptr<PjRtExecutable>> DeserializeExecutable(
       const PjRtTopologyDescription& topology,
       riegeli::Any<riegeli::Reader*> reader,
-      std::optional<CompileOptions>&& options) override;
+      std::optional<CompileOptions>&& options,
+      bool executable_compatibility_check_on_deserialization) override;
 
   PjRtPlatformId pjrt_platform_id() const { return pjrt_platform_id_; }
 
diff --git a/third_party/xla/xla/pjrt/pjrt_compiler.cc b/third_party/xla/xla/pjrt/pjrt_compiler.cc
index ac790d4b5be..7fd1f8f1ddc 100644
--- a/third_party/xla/xla/pjrt/pjrt_compiler.cc
+++ b/third_party/xla/xla/pjrt/pjrt_compiler.cc
@@ -298,10 +298,12 @@ absl::StatusOr<std::unique_ptr<PjRtExecutable>> PjRtCompile(
 absl::StatusOr<std::unique_ptr<PjRtExecutable>> PjRtDeserializeExecutable(
     const PjRtTopologyDescription& topology,
     riegeli::Any<riegeli::Reader*> reader,
-    std::optional<CompileOptions> options) {
+    std::optional<CompileOptions> options,
+    bool executable_compatibility_check_on_deserialization) {
   ABSL_ASSIGN_OR_RETURN(PjRtCompiler * compiler, GetPjRtCompiler(topology));
-  return compiler->DeserializeExecutable(topology, std::move(reader),
-                                         std::move(options));
+  return compiler->DeserializeExecutable(
+      topology, std::move(reader), std::move(options),
+      executable_compatibility_check_on_deserialization);
 }
 
 absl::Status PjRtPhaseCompiler::RegisterPhase(
diff --git a/third_party/xla/xla/pjrt/pjrt_compiler.h b/third_party/xla/xla/pjrt/pjrt_compiler.h
index d87254e633f..8f0685c9583 100644
--- a/third_party/xla/xla/pjrt/pjrt_compiler.h
+++ b/third_party/xla/xla/pjrt/pjrt_compiler.h
@@ -558,7 +558,8 @@ class PjRtCompiler {
   virtual absl::StatusOr<std::unique_ptr<PjRtExecutable>> DeserializeExecutable(
       const PjRtTopologyDescription& topology,
       riegeli::Any<riegeli::Reader*> reader,
-      std::optional<CompileOptions>&& options) {
+      std::optional<CompileOptions>&& options,
+      bool executable_compatibility_check_on_deserialization) {
     return absl::UnimplementedError(
         absl::StrCat("DeserializeExecutable is not implemented for: ",
                      topology.platform_name(), "."));
@@ -612,7 +613,8 @@ absl::StatusOr<std::unique_ptr<PjRtExecutable>> PjRtCompile(
 absl::StatusOr<std::unique_ptr<PjRtExecutable>> PjRtDeserializeExecutable(
     const PjRtTopologyDescription& topology,
     riegeli::Any<riegeli::Reader*> reader,
-    std::optional<CompileOptions> options = std::nullopt);
+    std::optional<CompileOptions> options = std::nullopt,
+    bool executable_compatibility_check_on_deserialization = true);
 
 // Stores a compilation phase's compiler and validator functions.
 // This struct bundles the essential functional components required to define
diff --git a/third_party/xla/xla/pjrt/raw_pjrt_client.h b/third_party/xla/xla/pjrt/raw_pjrt_client.h
index 3bf8b9c9b3d..e880bbee2fb 100644
--- a/third_party/xla/xla/pjrt/raw_pjrt_client.h
+++ b/third_party/xla/xla/pjrt/raw_pjrt_client.h
@@ -267,6 +267,15 @@ class PjRtRawClient {
     async_work_runner()->Execute(std::move(execute_fn));
   }
 
+  // Backend specific handlers for when an oom is detected during execute.
+  virtual void CallOomHandlers() const {}
+
+  virtual bool ShouldRetryOnOom(int attempts, PjRtDevice* device,
+                                PjRtExecutableLoadState* load_state,
+                                absl::Status prepare_status) {
+    return false;
+  }
+
   virtual void UpdateGlobalProcessInfo(
       absl::Span<xla::coordination::TaskInfo> infos) {
     LOG(WARNING) << "UpdateGlobalProcessInfo is not supported.";
diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc
index e0d68dbcc12..862a4b3f5f0 100644
--- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc
+++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client_test.cc
@@ -123,7 +123,8 @@ class StreamExecutorCpuCompiler : public PjRtCompiler {
   absl::StatusOr<std::unique_ptr<PjRtExecutable>> DeserializeExecutable(
       const PjRtTopologyDescription& topology,
       riegeli::Any<riegeli::Reader*> reader,
-      std::optional<CompileOptions>&& options) override {
+      std::optional<CompileOptions>&& options,
+      bool executable_compatibility_check_on_deserialization) override {
     return StreamExecutorExecutable::Deserialize(std::move(reader), topology,
                                                  std::move(options));
   }

From 77ab63cfaae1589716aacfc7e191de0e6edbac8f Mon Sep 17 00:00:00 2001
From: Viktor Moros <vgm@google.com>
Date: Fri, 16 Oct 2026 15:39:12 -0700
Subject: [PATCH 0988/1000] [XLA Highway] Switch XLA transpose kernels to
 Highway & clean up legacy intrinsics

PiperOrigin-RevId: 1001144549
---
 third_party/xla/xla/pjrt/BUILD               |   2 +-
 third_party/xla/xla/pjrt/transpose.cc        |  32 +-
 third_party/xla/xla/pjrt/transpose_kernels.h | 722 +------------------
 third_party/xla/xla/pjrt/transpose_test.cc   |  55 +-
 4 files changed, 30 insertions(+), 781 deletions(-)

diff --git a/third_party/xla/xla/pjrt/BUILD b/third_party/xla/xla/pjrt/BUILD
index 81e4f44e1e6..628f2991029 100644
--- a/third_party/xla/xla/pjrt/BUILD
+++ b/third_party/xla/xla/pjrt/BUILD
@@ -950,7 +950,7 @@ cc_library(
     srcs = ["transpose_kernels.cc"],
     hdrs = ["transpose_kernels.h"],
     deps = [
-        "//xla:compiler_macros",
+        "@com_google_absl//absl/base:core_headers",
         "@highway//:hwy",
     ],
 )
diff --git a/third_party/xla/xla/pjrt/transpose.cc b/third_party/xla/xla/pjrt/transpose.cc
index 9868c7817e9..4a7a3e78418 100644
--- a/third_party/xla/xla/pjrt/transpose.cc
+++ b/third_party/xla/xla/pjrt/transpose.cc
@@ -112,9 +112,7 @@ namespace xla {
 
 namespace {
 #ifdef __AVX__
-static constexpr int kMaxInnerBlockSizeBytes = sizeof(__m256i);
-#elif defined(XLA_HAS_VEC128)
-static constexpr int kMaxInnerBlockSizeBytes = sizeof(Vec128);
+static constexpr int kMaxInnerBlockSizeBytes = 32;
 #else
 static constexpr int kMaxInnerBlockSizeBytes = 16;
 #endif
@@ -336,13 +334,8 @@ void MacroKernel(const char* __restrict a, int64_t lda, int outer_bs_a,
     ldb = outer_bs_b * inner_bs * sizeof(T);
   }
 
-  for (int i = 0; i < outer_bs_a; ++i) {
-    for (int j = 0; j < outer_bs_b; ++j) {
-      TransposeMicroKernel<T, inner_bs>::Apply(
-          a + inner_bs * j * lda + i * inner_bs * sizeof(T), lda,
-          b + inner_bs * i * ldb + j * inner_bs * sizeof(T), ldb);
-    }
-  }
+  TransposeMacroKernelDispatch<T, inner_bs>(a, lda, outer_bs_a, b, ldb,
+                                            outer_bs_b);
 
   if constexpr (transformation == TransposePlan::Transformation::kPackSubbyte) {
     int elements_per_byte = 8 / bits_per_element;
@@ -662,6 +655,10 @@ void TransposePlan::ExecuteTyped(const char* a, char* b,
           handle_inner_block_elems(std::integral_constant<int, 16>{},
                                    const_kind);
           break;
+        case 32:
+          handle_inner_block_elems(std::integral_constant<int, 32>{},
+                                   const_kind);
+          break;
         default:
           LOG(FATAL) << "Invalid inner_block_elems_ " << inner_block_elems_;
       }
@@ -1525,6 +1522,8 @@ absl::Status TransposePlan::Initialize() {
   } else {
     // What are the smallest and largest block sizes for which we have a
     // vectorized kernel for this element size?
+    const int max_outer_block_elems =
+        kMaxOuterBlockSizeBytes / elem_size_in_bytes_;
     int min_inner_block_elems;
     int max_inner_block_elems;
     switch (elem_size_in_bytes_) {
@@ -1533,8 +1532,9 @@ absl::Status TransposePlan::Initialize() {
       case 4:
       case 8:
         min_inner_block_elems = 1;
-        max_inner_block_elems = std::min<int>(
-            kMaxOuterBlockElems, kMaxInnerBlockSizeBytes / elem_size_in_bytes_);
+        max_inner_block_elems =
+            std::min<int>(max_outer_block_elems,
+                          kMaxInnerBlockSizeBytes / elem_size_in_bytes_);
         break;
       case 16:
         min_inner_block_elems = 1;
@@ -1546,13 +1546,13 @@ absl::Status TransposePlan::Initialize() {
     if (inner_kernel_kind_ == InnerKernelKind::kInterleave) {
       inner_block_elems_ = 1;
       outer_block_elems_a_ = std::min<int64_t>(
-          large_loop_size, max_inner_block_elems * kMaxOuterBlockElems);
+          large_loop_size, max_inner_block_elems * max_outer_block_elems);
       outer_block_elems_b_ = 1;
     } else if (inner_kernel_kind_ == InnerKernelKind::kDeinterleave) {
       inner_block_elems_ = 1;
       outer_block_elems_a_ = 1;
       outer_block_elems_b_ = std::min<int64_t>(
-          large_loop_size, max_inner_block_elems * kMaxOuterBlockElems);
+          large_loop_size, max_inner_block_elems * max_outer_block_elems);
     } else {
       inner_block_elems_ = max_inner_block_elems;
       while (inner_block_elems_ > std::min(a_stride1_size, b_stride1_size)) {
@@ -1564,11 +1564,11 @@ absl::Status TransposePlan::Initialize() {
         inner_block_elems_ = 1;
       }
       outer_block_elems_a_ = FloorOfRatio<int64_t>(
-          std::min<int64_t>(kMaxOuterBlockElems, a_stride1_size),
+          std::min<int64_t>(max_outer_block_elems, a_stride1_size),
           inner_block_elems_);
       outer_block_elems_a_ = std::max<int64_t>(outer_block_elems_a_, 1);
       outer_block_elems_b_ = FloorOfRatio<int64_t>(
-          std::min<int64_t>(kMaxOuterBlockElems, b_stride1_size),
+          std::min<int64_t>(max_outer_block_elems, b_stride1_size),
           inner_block_elems_);
       outer_block_elems_b_ = std::max<int64_t>(outer_block_elems_b_, 1);
     }
diff --git a/third_party/xla/xla/pjrt/transpose_kernels.h b/third_party/xla/xla/pjrt/transpose_kernels.h
index 9b47abd0e5c..577c91b4800 100644
--- a/third_party/xla/xla/pjrt/transpose_kernels.h
+++ b/third_party/xla/xla/pjrt/transpose_kernels.h
@@ -16,17 +16,20 @@ limitations under the License.
 #ifndef XLA_PJRT_TRANSPOSE_KERNELS_H_
 #define XLA_PJRT_TRANSPOSE_KERNELS_H_
 
-#include <array>
-#include <cstddef>
 #include <cstdint>
 #include <cstring>
-#include <type_traits>
-#include <utility>
 
-#include "xla/compiler_macros.h"
+#include "absl/base/optimization.h"
 
 namespace xla {
 
+// Maximum byte size along each dimension of a macroblock.
+inline constexpr int kMaxOuterBlockSizeBytes = ABSL_CACHELINE_SIZE;
+
+// Maximum input byte stride (lda) for which the 128-bit square microkernel
+// is preferred over the 256-bit rectangular microkernel when lda >= ldb.
+inline constexpr int kMaxSquare128StrideBytes = ABSL_CACHELINE_SIZE;
+
 using TransposeMacroKernelFn = void (*)(const char* a, int64_t lda,
                                         int outer_bs_a, char* b, int64_t ldb,
                                         int outer_bs_b);
@@ -51,715 +54,6 @@ void TransposeMicroKernelDispatch(const char* a, int64_t lda, char* b,
                                       /*outer_bs_b=*/1);
 }
 
-}  // namespace xla
-
-#ifdef XLA_HAS_SSE2
-#include <immintrin.h>  // IWYU pragma: keep
-#endif
-
-#ifdef XLA_HAS_ARM_NEON
-#include <arm_neon.h>
-#endif  // XLA_HAS_ARM_NEON
-
-#if defined(XLA_HAS_SSE2) || defined(XLA_HAS_ARM_NEON)
-#define XLA_HAS_VEC128
-#endif  // defined(XLA_HAS_SSE2) || defined(XLA_HAS_ARM_NEON)
-
-namespace xla {
-
-// Maximum number of elements along each dimension of a macroblock. When the
-// inner microkernel block size `bs` is at least `kMaxOuterBlockElems`, the
-// outer block dimensions (`outer_bs_a` and `outer_bs_b`) are always 1.
-inline constexpr int kMaxOuterBlockElems = 16;
-
-// Maximum input byte stride (`lda`) for which the 128-bit square microkernel
-// is preferred over the 256-bit rectangular microkernel when `lda >= ldb`.
-inline constexpr int kMaxSquare128StrideBytes = 64;
-
-// The transpose microkernels use a general approach of zipping elements from
-// different rows together. We start zipping together elements of size 1, size 2
-// and so-on until we have achieved our transpose. As we increase the number of
-// consecutive elements we are zipping together, we also increase the distance
-// between the rows which are getting zipped.
-//
-// For example, let's say we want to transpose a 4x4 matrix:
-//
-// row 0: w0 w1 w2 w3
-// row 1: x0 x1 x2 x3
-// row 2: y0 y1 y2 y3
-// row 3: z0 z1 z2 z3
-//
-// We first zip groups of single elements; we will zip row 0 and row 1 together
-// and row 2 and row 3 together:
-//
-// row 0: w0 x0 w1 x1
-// row 1: w2 x2 w3 x3
-// row 2: y0 z0 y1 z1
-// row 3: y2 z2 y3 z3
-//
-// Then, we zip groups of two elements; we will zip row 0 and row 2 together and
-// row 1 and row 3 together:
-//
-// row 0: w0 x0 y0 z0
-// row 1: w1 x1 y1 z1
-// row 2: w2 x2 y2 z2
-// row 3: w3 x3 y3 z3
-//
-// Note that as we double the number of elements we are zipping, we are also
-// doubling the distance between rows which get zipped together.
-//
-// This standard algorithm gets slightly tweaked for what we will call
-// "rectangular" transposes. Such transposes are trying to use vectors larger
-// than the row length to speed up the transpose. This is accomplished by doing
-// the outermost transpose first via loads and stores. If we go back to our
-// original example, we would have:
-//
-// row 0: w0 w1 w2 w3 y0 y1 y2 y3
-// row 1: x0 x1 x2 x3 z0 z1 z2 z3
-//
-// Now, we do a half zip of our elements:
-//
-// row 0: w0 x0 w1 x1 y0 z0 y1 z1
-// row 1: w2 x2 w3 x3 y2 z2 y3 z3
-//
-// We can see that we have w{0-3} and y{0-3} in row 0 and x{0-3} and z{0-3} in
-// row 1 but they are not in the right order. We need to shuffle them once to
-// get them in the right order:
-//
-// row 0: w0 x0 y0 z0 w2 x2 y2 z2
-// row 1: w1 x1 y1 z1 w3 x3 y3 z3
-//
-// Now, we can extract two rows of 4 elements from row 0 and two rows of 4
-// elements from row 1 to store into memory.
-
-enum class Extract { kLo, kHi };
-
-#ifdef __AVX__
-template <size_t element_size, Extract>
-__m256i Unpack(__m256i a, __m256i b);
-
-#if defined(__AVX2__)
-template <>
-inline __m256i Unpack<1, Extract::kLo>(__m256i a, __m256i b) {
-  return _mm256_unpacklo_epi8(a, b);
-}
-template <>
-inline __m256i Unpack<1, Extract::kHi>(__m256i a, __m256i b) {
-  return _mm256_unpackhi_epi8(a, b);
-}
-
-template <>
-inline __m256i Unpack<2, Extract::kLo>(__m256i a, __m256i b) {
-  return _mm256_unpacklo_epi16(a, b);
-}
-template <>
-inline __m256i Unpack<2, Extract::kHi>(__m256i a, __m256i b) {
-  return _mm256_unpackhi_epi16(a, b);
-}
-
-template <>
-inline __m256i Unpack<4, Extract::kLo>(__m256i a, __m256i b) {
-  return _mm256_unpacklo_epi32(a, b);
-}
-template <>
-inline __m256i Unpack<4, Extract::kHi>(__m256i a, __m256i b) {
-  return _mm256_unpackhi_epi32(a, b);
-}
-
-template <>
-inline __m256i Unpack<8, Extract::kLo>(__m256i a, __m256i b) {
-  return _mm256_unpacklo_epi64(a, b);
-}
-template <>
-inline __m256i Unpack<8, Extract::kHi>(__m256i a, __m256i b) {
-  return _mm256_unpackhi_epi64(a, b);
-}
-#else
-template <>
-inline __m256i Unpack<1, Extract::kLo>(__m256i a, __m256i b) {
-  __m128i a_hi = _mm256_extractf128_si256(a, 1);
-  __m128i b_hi = _mm256_extractf128_si256(b, 1);
-  __m128i a_lo = _mm256_castsi256_si128(a);
-  __m128i b_lo = _mm256_castsi256_si128(b);
-  __m128i hi = _mm_unpacklo_epi8(a_hi, b_hi);
-  __m128i lo = _mm_unpacklo_epi8(a_lo, b_lo);
-  return _mm256_set_m128i(hi, lo);
-}
-template <>
-inline __m256i Unpack<1, Extract::kHi>(__m256i a, __m256i b) {
-  __m128i a_hi = _mm256_extractf128_si256(a, 1);
-  __m128i b_hi = _mm256_extractf128_si256(b, 1);
-  __m128i a_lo = _mm256_castsi256_si128(a);
-  __m128i b_lo = _mm256_castsi256_si128(b);
-  __m128i hi = _mm_unpackhi_epi8(a_hi, b_hi);
-  __m128i lo = _mm_unpackhi_epi8(a_lo, b_lo);
-  return _mm256_set_m128i(hi, lo);
-}
-
-template <>
-inline __m256i Unpack<2, Extract::kLo>(__m256i a, __m256i b) {
-  __m128i a_hi = _mm256_extractf128_si256(a, 1);
-  __m128i b_hi = _mm256_extractf128_si256(b, 1);
-  __m128i a_lo = _mm256_castsi256_si128(a);
-  __m128i b_lo = _mm256_castsi256_si128(b);
-  __m128i hi = _mm_unpacklo_epi16(a_hi, b_hi);
-  __m128i lo = _mm_unpacklo_epi16(a_lo, b_lo);
-  return _mm256_set_m128i(hi, lo);
-}
-template <>
-inline __m256i Unpack<2, Extract::kHi>(__m256i a, __m256i b) {
-  __m128i a_hi = _mm256_extractf128_si256(a, 1);
-  __m128i b_hi = _mm256_extractf128_si256(b, 1);
-  __m128i a_lo = _mm256_castsi256_si128(a);
-  __m128i b_lo = _mm256_castsi256_si128(b);
-  __m128i hi = _mm_unpackhi_epi16(a_hi, b_hi);
-  __m128i lo = _mm_unpackhi_epi16(a_lo, b_lo);
-  return _mm256_set_m128i(hi, lo);
-}
-
-template <>
-inline __m256i Unpack<4, Extract::kLo>(__m256i a, __m256i b) {
-  return _mm256_castps_si256(
-      _mm256_unpacklo_ps(_mm256_castsi256_ps(a), _mm256_castsi256_ps(b)));
-}
-template <>
-inline __m256i Unpack<4, Extract::kHi>(__m256i a, __m256i b) {
-  return _mm256_castps_si256(
-      _mm256_unpackhi_ps(_mm256_castsi256_ps(a), _mm256_castsi256_ps(b)));
-}
-
-template <>
-inline __m256i Unpack<8, Extract::kLo>(__m256i a, __m256i b) {
-  return _mm256_castpd_si256(
-      _mm256_unpacklo_pd(_mm256_castsi256_pd(a), _mm256_castsi256_pd(b)));
-}
-template <>
-inline __m256i Unpack<8, Extract::kHi>(__m256i a, __m256i b) {
-  return _mm256_castpd_si256(
-      _mm256_unpackhi_pd(_mm256_castsi256_pd(a), _mm256_castsi256_pd(b)));
-}
-#endif
-#endif
-
-#ifdef XLA_HAS_SSE2
-template <size_t element_size, Extract>
-__m128i Unpack(__m128i a, __m128i b);
-
-template <>
-inline __m128i Unpack<1, Extract::kLo>(__m128i a, __m128i b) {
-  return _mm_unpacklo_epi8(a, b);
-}
-template <>
-inline __m128i Unpack<1, Extract::kHi>(__m128i a, __m128i b) {
-  return _mm_unpackhi_epi8(a, b);
-}
-
-template <>
-inline __m128i Unpack<2, Extract::kLo>(__m128i a, __m128i b) {
-  return _mm_unpacklo_epi16(a, b);
-}
-template <>
-inline __m128i Unpack<2, Extract::kHi>(__m128i a, __m128i b) {
-  return _mm_unpackhi_epi16(a, b);
-}
-
-template <>
-inline __m128i Unpack<4, Extract::kLo>(__m128i a, __m128i b) {
-  return _mm_unpacklo_epi32(a, b);
-}
-template <>
-inline __m128i Unpack<4, Extract::kHi>(__m128i a, __m128i b) {
-  return _mm_unpackhi_epi32(a, b);
-}
-
-template <>
-inline __m128i Unpack<8, Extract::kLo>(__m128i a, __m128i b) {
-  return _mm_unpacklo_epi64(a, b);
-}
-template <>
-inline __m128i Unpack<8, Extract::kHi>(__m128i a, __m128i b) {
-  return _mm_unpackhi_epi64(a, b);
-}
-
-using Vec128 = __m128i;
-
-template <size_t bytes>
-__m128i LoadElementIntoVec128(const void* p);
-
-template <>
-inline __m128i LoadElementIntoVec128</*bytes=*/sizeof(uint16_t)>(
-    const void* p) {
-  // Note: We would ideally use `_mm_loadu_si16` here but older compilers do
-  // not support it. However, we can replicate it using a sequence such that
-  // even older compilers will turn this into a single movd instruction.
-  // memcpy is used because `p` is not guaranteed to be aligned to a 2-byte
-  // address.
-  uint16_t load;
-  memcpy(&load, p, sizeof(load));
-  return _mm_cvtsi32_si128(load);
-}
-
-template <>
-inline __m128i LoadElementIntoVec128</*bytes=*/sizeof(uint32_t)>(
-    const void* p) {
-  // Note: We would ideally use `_mm_loadu_si32` here but older compilers do
-  // not support it. However, we can replicate it using a sequence such that
-  // even older compilers will turn this into a single movd instruction.
-  // memcpy is used because `p` is not guaranteed to be aligned to a 4-byte
-  // address.
-  uint32_t load;
-  memcpy(&load, p, sizeof(load));
-  return _mm_cvtsi32_si128(load);
-}
-
-template <>
-inline __m128i LoadElementIntoVec128</*bytes=*/sizeof(uint64_t)>(
-    const void* p) {
-  // Note: We would ideally use `_mm_loadu_si64` here but older compilers do
-  // not support it. However, we can replicate it using a sequence such that
-  // even older compilers will turn this into a single movd instruction.
-  // memcpy is used because `p` is not guaranteed to be aligned to a 8-byte
-  // address.
-  uint64_t load;
-  memcpy(&load, p, sizeof(load));
-  return _mm_cvtsi64_si128(load);
-}
-
-template <>
-inline __m128i LoadElementIntoVec128</*bytes=*/sizeof(__m128i)>(const void* p) {
-  return _mm_loadu_si128(reinterpret_cast<const __m128i*>(p));
-}
-
-template <size_t bytes, int lane>
-inline void StoreElementFromVec128(void* p, __m128i v) {
-  constexpr size_t element_start = bytes * lane;
-  constexpr size_t element_end = element_start + bytes;
-  static_assert(element_start >= 0);
-  static_assert(element_end <= sizeof(Vec128));
-  constexpr bool halfway = element_start == sizeof(Vec128) / 2;
-  if constexpr (bytes == sizeof(uint16_t)) {
-    // Note: We would ideally use `_mm_storeu_si16` here but older compilers do
-    // not support it. However, we can replicate it using a sequence such that
-    // even older compilers will turn this into a single movd instruction.
-    // memcpy is used because `p` is not guaranteed to be aligned to a 4-byte
-    // address.
-    const uint16_t scalar = _mm_extract_epi16(v, lane);
-    memcpy(p, &scalar, bytes);
-  } else if constexpr (bytes == sizeof(uint32_t)) {
-    if constexpr (halfway) {
-      v = Unpack<bytes, Extract::kHi>(v, v);
-    } else if constexpr (lane != 0) {
-      v = _mm_shuffle_epi32(v, _MM_SHUFFLE(lane, lane, lane, lane));
-    }
-    // Note: We would ideally use `_mm_storeu_si32` here but older compilers do
-    // not support it. However, we can replicate it using a sequence such that
-    // even older compilers will turn this into a single movd instruction.
-    // memcpy is used because `p` is not guaranteed to be aligned to a 4-byte
-    // address.
-    const uint32_t scalar = _mm_cvtsi128_si32(v);
-    memcpy(p, &scalar, bytes);
-  } else if constexpr (bytes == sizeof(uint64_t)) {
-    if constexpr (halfway) {
-      v = Unpack<bytes, Extract::kHi>(v, v);
-    } else {
-      static_assert(lane == 0);
-    }
-    // Note: We would ideally use `_mm_storeu_si64` here but older compilers do
-    // not support it. However, we can replicate it using a sequence such that
-    // even older compilers will turn this into a single movd instruction.
-    // memcpy is used because `p` is not guaranteed to be aligned to a 8-byte
-    // address.
-    const uint64_t scalar = _mm_cvtsi128_si64(v);
-    memcpy(p, &scalar, bytes);
-  } else if constexpr (bytes == sizeof(__m128i)) {
-    static_assert(lane == 0);
-    _mm_storeu_si128(reinterpret_cast<__m128i*>(p), v);
-  } else {
-    static_assert(bytes == 0);
-  }
-}
-#endif
-
-#ifdef XLA_HAS_ARM_NEON
-template <size_t element_size, Extract>
-uint64x2_t Unpack(uint64x2_t a, uint64x2_t b);
-
-template <>
-inline uint64x2_t Unpack<1, Extract::kLo>(uint64x2_t a, uint64x2_t b) {
-  return vreinterpretq_u64_u8(
-      vzipq_u8(vreinterpretq_u8_u64(a), vreinterpretq_u8_u64(b)).val[0]);
-}
-template <>
-inline uint64x2_t Unpack<1, Extract::kHi>(uint64x2_t a, uint64x2_t b) {
-  return vreinterpretq_u64_u8(
-      vzipq_u8(vreinterpretq_u8_u64(a), vreinterpretq_u8_u64(b)).val[1]);
-}
-
-template <>
-inline uint64x2_t Unpack<2, Extract::kLo>(uint64x2_t a, uint64x2_t b) {
-  return vreinterpretq_u64_u16(
-      vzipq_u16(vreinterpretq_u16_u64(a), vreinterpretq_u16_u64(b)).val[0]);
-}
-template <>
-inline uint64x2_t Unpack<2, Extract::kHi>(uint64x2_t a, uint64x2_t b) {
-  return vreinterpretq_u64_u16(
-      vzipq_u16(vreinterpretq_u16_u64(a), vreinterpretq_u16_u64(b)).val[1]);
-}
-
-template <>
-inline uint64x2_t Unpack<4, Extract::kLo>(uint64x2_t a, uint64x2_t b) {
-  return vreinterpretq_u64_u32(
-      vzipq_u32(vreinterpretq_u32_u64(a), vreinterpretq_u32_u64(b)).val[0]);
-}
-template <>
-inline uint64x2_t Unpack<4, Extract::kHi>(uint64x2_t a, uint64x2_t b) {
-  return vreinterpretq_u64_u32(
-      vzipq_u32(vreinterpretq_u32_u64(a), vreinterpretq_u32_u64(b)).val[1]);
-}
-
-template <>
-inline uint64x2_t Unpack<8, Extract::kLo>(uint64x2_t a, uint64x2_t b) {
-  uint64x1_t a_lo = vget_low_u64(a);
-  uint64x1_t b_lo = vget_low_u64(b);
-  return vcombine_u64(a_lo, b_lo);
-}
-template <>
-inline uint64x2_t Unpack<8, Extract::kHi>(uint64x2_t a, uint64x2_t b) {
-  uint64x1_t a_hi = vget_high_u64(a);
-  uint64x1_t b_hi = vget_high_u64(b);
-  return vcombine_u64(a_hi, b_hi);
-}
-
-using Vec128 = uint64x2_t;
-
-template <size_t>
-uint64x2_t LoadElementIntoVec128(const void* p);
-
-template <>
-inline uint64x2_t LoadElementIntoVec128</*bytes=*/sizeof(uint16_t)>(
-    const void* p) {
-  // Ideally, we would use `vld1q_lane_u16` but it assumes that its input is
-  // aligned to a 16-bit boundary. We can only promise 8-bit aligned. That said,
-  // this sequence will compile to `ldr St, [Xn]` but without an alignment hint.
-  uint16_t x;
-  memcpy(&x, p, sizeof(x));
-  return vreinterpretq_u64_u16(vsetq_lane_u16(x, vdupq_n_u16(0), 0));
-}
-
-template <>
-inline uint64x2_t LoadElementIntoVec128</*bytes=*/sizeof(uint32_t)>(
-    const void* p) {
-  // Ideally, we would use `vld1q_lane_u32` but it assumes that its input is
-  // aligned to a 32-bit boundary. We can only promise 8-bit aligned. That said,
-  // this sequence will compile to `ldr St, [Xn]` but without an alignment hint.
-  uint32_t x;
-  memcpy(&x, p, sizeof(x));
-  return vreinterpretq_u64_u32(vsetq_lane_u32(x, vdupq_n_u32(0), 0));
-}
-
-template <>
-inline uint64x2_t LoadElementIntoVec128</*bytes=*/sizeof(uint64_t)>(
-    const void* p) {
-  // Ideally, we would use `vld1q_lane_u64` but it assumes that its input is
-  // aligned to a 64-bit boundary. We can only promise 8-bit aligned. That said,
-  // this sequence will compile to `ldr Dt, [Xn]` but without an alignment hint.
-  return vreinterpretq_u64_u8(
-      vcombine_u8(vld1_u8(reinterpret_cast<const uint8_t*>(p)), vdup_n_u8(0)));
-}
-
-template <>
-inline uint64x2_t LoadElementIntoVec128</*bytes=*/sizeof(uint64x2_t)>(
-    const void* p) {
-  return vreinterpretq_u64_u8(vld1q_u8(reinterpret_cast<const uint8_t*>(p)));
-}
-
-template <size_t bytes, int lane>
-inline void StoreElementFromVec128(void* p, uint64x2_t v) {
-  static_assert(bytes * (lane + 1) <= sizeof(uint64x2_t));
-  if constexpr (bytes == sizeof(uint64x2_t)) {
-    static_assert(lane == 0);
-    vst1q_u8(reinterpret_cast<uint8_t*>(p), vreinterpretq_u8_u64(v));
-  } else {
-    if constexpr (bytes == sizeof(uint64_t)) {
-      // Ideally, we would use `vst1q_lane_u64` but it assumes that its input is
-      // aligned to a 64-bit boundary. We can only promise 8-bit aligned. That
-      // said, this sequence will compile to `st1 {vt.d}[lane], [xn]` but
-      // without an alignment hint.
-      uint64_t extracted = vgetq_lane_u64(v, lane);
-      memcpy(p, &extracted, sizeof(extracted));
-    } else if constexpr (bytes == sizeof(uint32_t)) {
-      // Ideally, we would use `vst1q_lane_u32` but it assumes that its input is
-      // aligned to a 32-bit boundary. We can only promise 8-bit aligned. That
-      // said, this sequence will compile to `st1 {vt.s}[lane], [xn]` but
-      // without an alignment hint.
-      uint32_t extracted = vgetq_lane_u32(vreinterpretq_u32_u64(v), lane);
-      memcpy(p, &extracted, sizeof(extracted));
-    } else if constexpr (bytes == sizeof(uint16_t)) {
-      // Ideally, we would use `vst1q_lane_u16` but it assumes that its input is
-      // aligned to a 16-bit boundary. We can only promise 8-bit aligned. That
-      // said, this sequence will compile to `st1 {vt.h}[lane], [xn]` but
-      // without an alignment hint.
-      uint16_t extracted = vgetq_lane_u16(vreinterpretq_u16_u64(v), lane);
-      memcpy(p, &extracted, sizeof(extracted));
-    } else {
-      static_assert(bytes == 0);
-    }
-  }
-}
-#endif
-
-#ifdef XLA_HAS_VEC128
-template <size_t element_size, size_t step_size, typename T, size_t N>
-inline std::array<T, N> UnpackStep(const std::array<T, N>& last_transpose) {
-  static_assert(N % (step_size * 2) == 0);
-  std::array<T, N> unpack;
-  XLA_UNROLL
-  for (int i = 0; i < N; i += step_size * 2) {
-    XLA_UNROLL
-    for (int j = 0; j < step_size; ++j) {
-      unpack[i + 2 * j + 0] = Unpack<element_size * step_size, Extract::kLo>(
-          last_transpose[i + j], last_transpose[i + j + step_size]);
-      unpack[i + 2 * j + 1] = Unpack<element_size * step_size, Extract::kHi>(
-          last_transpose[i + j], last_transpose[i + j + step_size]);
-    }
-  }
-  return unpack;
-}
-
-template <size_t element_size, size_t step_size, size_t unpack_limit,
-          typename T, size_t N>
-inline std::array<T, N> UnpackSequence(const std::array<T, N>& last_transpose) {
-  if constexpr (element_size * step_size < unpack_limit) {
-    std::array<T, N> unpack =
-        UnpackStep<element_size, step_size>(last_transpose);
-    return UnpackSequence<element_size, step_size * 2, unpack_limit>(unpack);
-  }
-  return last_transpose;
-}
-
-template <size_t element_size, size_t bs, typename T, size_t N>
-inline auto UnpackLowSequence(const std::array<T, N>& last_transpose) {
-  if constexpr (N > 1 && element_size * bs < sizeof(T)) {
-    static_assert(N % 2 == 0);
-    std::array<T, N / 2> unpack;
-    for (int i = 0; i < N; i += 2) {
-      unpack[i / 2] = Unpack<element_size, Extract::kLo>(last_transpose[i],
-                                                         last_transpose[i + 1]);
-    }
-    return UnpackLowSequence<element_size * 2, bs>(unpack);
-  } else {
-    return last_transpose;
-  }
-}
-
-template <size_t bytes, size_t... lane>
-inline void StoreElementsFromVec128(char* b, int64_t ldb, Vec128 x, size_t i,
-                                    std::index_sequence<lane...>) {
-  (StoreElementFromVec128</*bytes=*/bytes, lane>(b + ldb * (i + lane), x), ...);
-}
-
-template <typename T, int bs>
-struct Vec128RectangularTransposeMicroKernelImpl {
-  XLA_FLATTEN static void Apply(const char* __restrict a, int64_t lda,
-                                char* __restrict b, int64_t ldb) {
-    constexpr size_t element_size = sizeof(T);
-    static_assert(sizeof(Vec128) % element_size == 0);
-    std::array<Vec128, bs> loads;
-    // `loads`:
-    // [  0,  1,  2,  3 ]
-    // [  4,  5,  6,  7 ]
-    // [  8,  9, 10, 11 ]
-    // [ 12, 13, 14, 15 ]
-    XLA_UNROLL
-    for (int i = 0; i < bs; ++i) {
-      loads[i] = LoadElementIntoVec128<element_size * bs>(a + lda * i);
-    }
-
-    // Each load may not have filled the entire vector. Combine rows to get a
-    // fully populated vector.
-    auto last_transpose = UnpackLowSequence<element_size, bs>(loads);
-
-    constexpr int kBytesInMatrix = element_size * bs * bs;
-    static_assert(kBytesInMatrix <= sizeof(Vec128) * last_transpose.size());
-    static_assert(bs % last_transpose.size() == 0);
-    constexpr int kStoresPerCombinedRow = bs / last_transpose.size();
-
-    // Each row of the matrix originally occupied some fraction of a vector.
-    // We may need to finish the transpose if `last_transpose.size() > 1`.
-    // `last_transpose`:
-    // [ 0,  4,  1,  5,  2,  6,  3,  7 ]
-    // [ 8, 12,  9, 13, 10, 14, 11, 15 ]
-    last_transpose =
-        UnpackSequence<kBytesInMatrix / (bs * last_transpose.size()),
-                       /*step_size=*/1,
-                       /*unpack_limit=*/kBytesInMatrix / bs>(last_transpose);
-    // `last_transpose`:
-    // [ 0,  4,  8, 12,  1,  5,  9, 13 ]
-    // [ 2,  6, 10, 14,  3,  7, 11, 15 ]
-    XLA_UNROLL
-    for (size_t i = 0; i < last_transpose.size(); ++i) {
-      StoreElementsFromVec128<element_size * bs>(
-          b, ldb, last_transpose[i], i * kStoresPerCombinedRow,
-          std::make_index_sequence<kStoresPerCombinedRow>{});
-    }
-  }
-};
-
-#endif
-
-#ifdef __AVX__
-template <typename T, int bs>
-struct AvxSquareTransposeMicroKernelImpl {
-  XLA_FLATTEN static void Apply(const char* __restrict a, int64_t lda,
-                                char* __restrict b, int64_t ldb) {
-    constexpr size_t element_size = sizeof(T);
-    static_assert(element_size <= sizeof(__m128i));
-    static_assert(sizeof(__m128i) % element_size == 0);
-    static_assert(bs % 2 == 0);
-    static_assert(element_size * bs == sizeof(__m256i));
-    std::array<__m256i, bs> last_transpose;
-    XLA_UNROLL
-    for (int i = 0; i < bs / 2; ++i) {
-      auto* row0_low = reinterpret_cast<const __m128i*>(a + lda * (i + 0));
-      auto* row0_high = row0_low + 1;
-      auto* row1_low = reinterpret_cast<const __m128i*>(a + lda * (i + bs / 2));
-      auto* row1_high = row1_low + 1;
-
-      last_transpose[i] = _mm256_set_m128i(_mm_loadu_si128(row1_low),
-                                           _mm_loadu_si128(row0_low));
-      last_transpose[i + bs / 2] = _mm256_set_m128i(_mm_loadu_si128(row1_high),
-                                                    _mm_loadu_si128(row0_high));
-    }
-
-    last_transpose =
-        UnpackSequence<element_size, /*step_size=*/1,
-                       /*unpack_limit=*/sizeof(__m128i)>(last_transpose);
-
-    XLA_UNROLL
-    for (int i = 0; i < bs; ++i) {
-      _mm256_storeu_si256(reinterpret_cast<__m256i*>(b + ldb * i),
-                          last_transpose[i]);
-    }
-  }
-};
-#endif
-
-#ifdef __AVX__
-template <typename T, int bs>
-struct AvxRectangularTransposeMicroKernelImpl {
-  XLA_FLATTEN static void Apply(const char* __restrict a, int64_t lda,
-                                char* __restrict b, int64_t ldb) {
-    constexpr size_t element_size = sizeof(T);
-    static_assert(element_size <= sizeof(__m128i));
-    static_assert(sizeof(__m128i) % element_size == 0);
-    static_assert(bs % 2 == 0);
-    static_assert(element_size * bs * 2 == sizeof(__m256i));
-    std::array<__m256i, bs / 2> last_transpose;
-    XLA_UNROLL
-    for (int i = 0; i < bs / 2; ++i) {
-      auto* lo = reinterpret_cast<const __m128i*>(a + lda * (i + 0));
-      auto* hi = reinterpret_cast<const __m128i*>(a + lda * (i + bs / 2));
-      last_transpose[i] =
-          _mm256_set_m128i(_mm_loadu_si128(hi), _mm_loadu_si128(lo));
-    }
-
-    last_transpose =
-        UnpackSequence<element_size, /*step_size=*/1,
-                       /*unpack_limit=*/sizeof(__m128i) / 2>(last_transpose);
-
-    if constexpr (element_size <= sizeof(__m128i) / 2) {
-      XLA_UNROLL
-      for (int i = 0; i < bs / 2; ++i) {
-#if defined(__AVX2__)
-        last_transpose[i] = _mm256_permute4x64_epi64(last_transpose[i],
-                                                     _MM_SHUFFLE(3, 1, 2, 0));
-#else
-        auto a = last_transpose[i];
-        auto hi = _mm256_permute2f128_si256(a, a, 0b0001'0001);
-        auto lo = _mm256_insertf128_si256(a, _mm256_castsi256_si128(a), 1);
-        last_transpose[i] = _mm256_castpd_si256(_mm256_shuffle_pd(
-            _mm256_castsi256_pd(lo), _mm256_castsi256_pd(hi), 0b1100));
-#endif
-      }
-    }
-
-    XLA_UNROLL
-    for (int i = 0; i < bs / 2; ++i) {
-      auto* lo = reinterpret_cast<__m128i*>(b + ldb * (i * 2 + 0));
-      auto* hi = reinterpret_cast<__m128i*>(b + ldb * (i * 2 + 1));
-
-      _mm_storeu_si128(lo, _mm256_castsi256_si128(last_transpose[i]));
-      _mm_storeu_si128(hi, _mm256_extractf128_si256(last_transpose[i], 1));
-    }
-  }
-};
-
-template <typename T, int bs>
-struct SseSquareTransposeMicroKernelImpl {
-  XLA_FLATTEN static void Apply(const char* __restrict a, int64_t lda,
-                                char* __restrict b, int64_t ldb) {
-    constexpr size_t element_size = sizeof(T);
-    static_assert(element_size * bs == sizeof(__m128i));
-    std::array<__m128i, bs> last_transpose;
-    XLA_UNROLL
-    for (int i = 0; i < bs; ++i) {
-      last_transpose[i] =
-          _mm_loadu_si128(reinterpret_cast<const __m128i*>(a + lda * i));
-    }
-
-    last_transpose = UnpackSequence<element_size, /*step_size=*/1,
-                                    /*unpack_limit=*/16>(last_transpose);
-
-    XLA_UNROLL
-    for (int i = 0; i < bs; ++i) {
-      _mm_storeu_si128(reinterpret_cast<__m128i*>(b + ldb * i),
-                       last_transpose[i]);
-    }
-  }
-};
-#endif
-
-// The transpose kernel requires its input to be contiguous in one of the two
-// dimensions being transposed, and the output to be contiguous in the other
-// dimension.
-//
-// lda, ldb are strides in bytes.
-template <typename T, int bs>
-struct TransposeMicroKernel {
-  static void Apply(const char* __restrict a, int64_t lda, char* __restrict b,
-                    int64_t ldb) {
-    if constexpr (bs % 2 == 0) {
-#ifdef __AVX__
-      if constexpr (sizeof(T) * bs == sizeof(__m256i)) {
-        return AvxSquareTransposeMicroKernelImpl<T, bs>::Apply(a, lda, b, ldb);
-      } else if constexpr (sizeof(T) * bs == sizeof(__m128i)) {
-        // Prefer SseSquare when gathering from memory (lda >= ldb) with small
-        // strides, rather than dealing with lo/hi packing.
-        if (lda >= ldb && lda <= kMaxSquare128StrideBytes) {
-          return SseSquareTransposeMicroKernelImpl<T, bs>::Apply(a, lda, b,
-                                                                 ldb);
-        }
-        return AvxRectangularTransposeMicroKernelImpl<T, bs>::Apply(a, lda, b,
-                                                                    ldb);
-      }
-#endif
-#ifdef XLA_HAS_VEC128
-      if constexpr (sizeof(T) * bs <= sizeof(Vec128)) {
-        return Vec128RectangularTransposeMicroKernelImpl<T, bs>::Apply(a, lda,
-                                                                       b, ldb);
-      }
-#endif
-    }
-
-    for (int i = 0; i < bs; ++i) {
-      for (int j = 0; j < bs; ++j) {
-        std::memcpy(b + i * ldb + j * sizeof(T), a + j * lda + i * sizeof(T),
-                    sizeof(T));
-      }
-    }
-  }
-};
-
 // Kernels for InnerKernelKind::kInterleave and kDeinterleave over `n` elements
 // of each of the K rows. They are written as scalar loops with a compile-time
 // K for clang to vectorize, e.g. with NEON structured loads/stores (ld2/ld3,
diff --git a/third_party/xla/xla/pjrt/transpose_test.cc b/third_party/xla/xla/pjrt/transpose_test.cc
index b804505c96b..94013a47c67 100644
--- a/third_party/xla/xla/pjrt/transpose_test.cc
+++ b/third_party/xla/xla/pjrt/transpose_test.cc
@@ -71,9 +71,7 @@ namespace xla {
 // macro-kernel test and benchmark.
 template <typename T, int bs>
 constexpr int BenchmarkOuterBlockSize() {
-  int max_by_elems = kMaxOuterBlockElems / bs;
-  int max_by_bytes = kMaxSquare128StrideBytes / (bs * sizeof(T));
-  return std::max(1, std::min(max_by_elems, max_by_bytes));
+  return std::max<int>(1, kMaxOuterBlockSizeBytes / (bs * sizeof(T)));
 }
 
 // Checks that `kernel(a, lda, b, ldb)` transposes an n x n matrix of T, where
@@ -134,16 +132,6 @@ void TestMacroKernelEquivalence() {
       });
 }
 
-// TODO(b/543906735): Remove along with the legacy SIMD intrinsic kernels once
-// transpose.cc uses the Highway kernels.
-template <typename T, int bs>
-void TestLegacyMicroKernelEquivalence() {
-  CheckTransposeKernel<T, bs>(
-      /*outer_bs=*/1, [](const char* a, int64_t lda, char* b, int64_t ldb) {
-        TransposeMicroKernel<T, bs>::Apply(a, lda, b, ldb);
-      });
-}
-
 TEST(TransposeMicroKernelTest, ExactEquivalence) {
   for (int64_t target : hwy::SupportedAndGeneratedTargets()) {
     hwy::SetSupportedTargetsForTest(target);
@@ -212,39 +200,6 @@ TEST(TransposeMacroKernelTest, ExactEquivalence) {
   hwy::SetSupportedTargetsForTest(0);
 }
 
-// TODO(b/543906735): Remove along with the legacy SIMD intrinsic kernels once
-// transpose.cc uses the Highway kernels.
-TEST(TransposeMicroKernelTest, LegacyExactEquivalence) {
-  // AvxSquareTransposeMicroKernelImpl is triggered when a logical row of the
-  // tile (bs * sizeof(T)) is exactly 256 bits to fit in __m256i.
-  TestLegacyMicroKernelEquivalence<int8_t, 32>();
-  TestLegacyMicroKernelEquivalence<int16_t, 16>();
-  TestLegacyMicroKernelEquivalence<float, 8>();
-  TestLegacyMicroKernelEquivalence<int64_t, 4>();
-  TestLegacyMicroKernelEquivalence<absl::uint128, 2>();
-
-  // SseSquareTransposeMicroKernelImpl or AvxRectangularTransposeMicroKernelImpl
-  // is triggered when a logical row of the tile (bs * sizeof(T)) is exactly
-  // 128 bits. SseSquare operates directly on __m128i when gathering from memory
-  // (lda >= ldb && lda <= 64B) while AvxRectangular packs two rows into __m256i
-  // when scattering to memory (ldb > lda) or for large strides.
-  TestLegacyMicroKernelEquivalence<int8_t, 16>();
-  TestLegacyMicroKernelEquivalence<int16_t, 8>();
-  TestLegacyMicroKernelEquivalence<bfloat16, 8>();
-  TestLegacyMicroKernelEquivalence<float, 4>();
-  TestLegacyMicroKernelEquivalence<int64_t, 2>();
-
-  // Smaller or larger cases fall back to either Vec128 or a for loop.
-  TestLegacyMicroKernelEquivalence<int8_t, 8>();
-  TestLegacyMicroKernelEquivalence<int16_t, 4>();
-  TestLegacyMicroKernelEquivalence<bfloat16, 4>();
-  TestLegacyMicroKernelEquivalence<float, 2>();
-  TestLegacyMicroKernelEquivalence<int8_t, 4>();
-  TestLegacyMicroKernelEquivalence<int16_t, 2>();
-  TestLegacyMicroKernelEquivalence<int8_t, 2>();
-  TestLegacyMicroKernelEquivalence<int8_t, 64>();
-}
-
 class TestTransposePlan : public TransposePlan {
  public:
   using Loop = TransposePlan::Loop;
@@ -1806,10 +1761,10 @@ TEST(TransposeTest, PackInt4_2D) {
                TransposePlan::ChunkContiguity::kOutput,
                /*expected_parallelism=*/3,
                /*expected_inner_kernel_is_memcpy=*/true);
-  TestPackIntN(4, /*dims=*/{64, 1}, /*permutation=*/{1, 0},
-               /*input_tiling=*/{}, /*output_tiling=*/{},
-               /*input_striding=*/{2, 1}, /*num_threads=*/8,
-               TransposePlan::ChunkContiguity::kOutput,
+  TestPackIntN(4, /*dims=*/{4 * kMaxOuterBlockSizeBytes, 1},
+               /*permutation=*/{1, 0}, /*input_tiling=*/{},
+               /*output_tiling=*/{}, /*input_striding=*/{2, 1},
+               /*num_threads=*/8, TransposePlan::ChunkContiguity::kOutput,
                /*expected_parallelism=*/4,
                /*expected_inner_kernel_is_memcpy=*/false);
 }

From 25ca0f0b536e4ccfe1e3151193f46e66a4825d32 Mon Sep 17 00:00:00 2001
From: Junwhan Ahn <junwhan@google.com>
Date: Fri, 16 Oct 2026 15:56:57 -0700
Subject: [PATCH 0989/1000] Add protobuf serialization and deserialization to
 tensorflow::RequestCost.

Define `tensorflow.RequestCostProto` and add `RequestCost::ToProto` and `RequestCost::FromProto` methods next to `RequestCost` so that callers can serialize and deserialize request costs and metrics without maintaining separate conversion logic when new fields are added.

PiperOrigin-RevId: 1001155709
---
 tensorflow/core/BUILD                         |   2 +
 tensorflow/core/common_runtime/BUILD          |  10 ++
 .../core/common_runtime/request_cost.cc       |  98 +++++++++++++
 tensorflow/core/common_runtime/request_cost.h |  10 ++
 .../core/common_runtime/request_cost.proto    |  61 ++++++++
 .../core/common_runtime/request_cost_test.cc  | 130 ++++++++++++++++++
 6 files changed, 311 insertions(+)
 create mode 100644 tensorflow/core/common_runtime/request_cost.proto

diff --git a/tensorflow/core/BUILD b/tensorflow/core/BUILD
index 824404a69b8..3ccb707ba62 100644
--- a/tensorflow/core/BUILD
+++ b/tensorflow/core/BUILD
@@ -192,6 +192,7 @@ tf_proto_library(
     make_default_target_header_only = True,
     visibility = ["//visibility:public"],
     deps = [
+        "//tensorflow/core/common_runtime:request_cost_proto",
         "//tensorflow/core/example:protos_all",
         "//tensorflow/core/framework:protos_all",
         "//tensorflow/core/grappler/costs:op_performance_data",
@@ -1158,6 +1159,7 @@ cc_library(
     name = "mobile_additional_lib_deps",
     deps = tf_additional_lib_deps() + [
         ":platform_base",
+        "//tensorflow/core/common_runtime:request_cost_proto_cc",
         "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:flat_hash_set",
diff --git a/tensorflow/core/common_runtime/BUILD b/tensorflow/core/common_runtime/BUILD
index eb6715d4603..1cb98e0075b 100644
--- a/tensorflow/core/common_runtime/BUILD
+++ b/tensorflow/core/common_runtime/BUILD
@@ -37,6 +37,7 @@ load("//tensorflow:tensorflow.default.bzl", "cc_header_only_library", "filegroup
 load(
     "//tensorflow/core/platform:build_config.bzl",
     "tf_additional_core_deps",
+    "tf_proto_library",
     "tf_protos_all",
     "tf_protos_grappler",
 )
@@ -2462,17 +2463,26 @@ cc_library(
     alwayslink = 1,
 )
 
+tf_proto_library(
+    name = "request_cost_proto",
+    srcs = ["request_cost.proto"],
+    visibility = ["//visibility:public"],
+)
+
 cc_library(
     name = "request_cost",
     srcs = ["request_cost.cc"],
     hdrs = ["request_cost.h"],
     visibility = ["//visibility:public"],
     deps = [
+        ":request_cost_proto_cc",
         "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/functional:overload",
         "@com_google_absl//absl/strings",
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
+        "@com_google_absl//absl/types:span",
     ],
 )
 
diff --git a/tensorflow/core/common_runtime/request_cost.cc b/tensorflow/core/common_runtime/request_cost.cc
index bd0b25d615d..8e2ff1df21f 100644
--- a/tensorflow/core/common_runtime/request_cost.cc
+++ b/tensorflow/core/common_runtime/request_cost.cc
@@ -17,12 +17,15 @@ limitations under the License.
 
 #include <string>
 #include <utility>
+#include <variant>
 #include <vector>
 
 #include "absl/container/flat_hash_map.h"
+#include "absl/functional/overload.h"
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/time/time.h"
+#include "absl/types/span.h"
 
 namespace tensorflow {
 
@@ -103,4 +106,99 @@ std::vector<RequestCost::BatchMetrics> RequestCost::GetBatchMetrics() const {
   return batch_metrics_;
 }
 
+void RequestCost::ToProto(RequestCostProto& proto) const {
+  absl::MutexLock lock(mutex_);
+  proto.Clear();
+  for (const auto& [cost_type, cost] : cost_map_) {
+    proto.mutable_costs()->insert(
+        {cost_type, absl::ToDoubleMicroseconds(cost)});
+  }
+  for (const auto& [cost_type, dims] : structured_cost_map_) {
+    auto& structured = (*proto.mutable_structured_costs())[cost_type];
+    for (const auto& [dim, value] : dims) {
+      structured.mutable_dimensions()->insert(
+          {dim, absl::ToDoubleMicroseconds(value)});
+    }
+  }
+  for (const auto& [metric_type, metric_value] : metric_map_) {
+    proto.mutable_metrics()->insert({metric_type, metric_value});
+  }
+  for (const auto& [name, structured] : structured_metric_map_) {
+    RequestCostProto::StructuredMetric& structured_proto =
+        (*proto.mutable_structured_metrics())[name];
+    std::visit(absl::Overload{
+                   [&structured_proto](absl::Span<const double> values) {
+                     structured_proto.mutable_double_values()->Assign(
+                         values.begin(), values.end());
+                   },
+                   [&structured_proto](absl::Span<const std::string> values) {
+                     structured_proto.mutable_bytes_values()->Assign(
+                         values.begin(), values.end());
+                   },
+               },
+               structured.values);
+  }
+  for (const auto& bm : batch_metrics_) {
+    RequestCostProto::BatchMetrics* bm_proto = proto.add_batch_metrics();
+    bm_proto->set_processed_size(bm.processed_size);
+    bm_proto->set_input_size(bm.input_size);
+    bm_proto->set_padding_size(bm.padding_size);
+    for (const auto& [cost_type, cost] : bm.batch_costs) {
+      bm_proto->mutable_batch_costs()->insert(
+          {cost_type, absl::ToDoubleMicroseconds(cost)});
+    }
+  }
+}
+
+RequestCostProto RequestCost::ToProto() const {
+  RequestCostProto proto;
+  ToProto(proto);
+  return proto;
+}
+
+void RequestCost::FromProto(const RequestCostProto& proto) {
+  absl::MutexLock lock(mutex_);
+  cost_map_.clear();
+  structured_cost_map_.clear();
+  metric_map_.clear();
+  structured_metric_map_.clear();
+  batch_metrics_.clear();
+
+  for (const auto& [name, value] : proto.costs()) {
+    cost_map_[name] = absl::Microseconds(value);
+  }
+  for (const auto& [name, structured] : proto.structured_costs()) {
+    StructuredCost dims;
+    for (const auto& [dim, value] : structured.dimensions()) {
+      dims[dim] = absl::Microseconds(value);
+    }
+    structured_cost_map_[name] = std::move(dims);
+  }
+  for (const auto& [name, value] : proto.metrics()) {
+    metric_map_[name] = value;
+  }
+  for (const auto& [name, src] : proto.structured_metrics()) {
+    StructuredMetric metric;
+    if (!src.bytes_values().empty()) {
+      metric.values.emplace<1>(src.bytes_values().begin(),
+                               src.bytes_values().end());
+    } else {
+      metric.values.emplace<0>(src.double_values().begin(),
+                               src.double_values().end());
+    }
+    structured_metric_map_[name] = std::move(metric);
+  }
+  batch_metrics_.reserve(proto.batch_metrics_size());
+  for (const auto& bm_proto : proto.batch_metrics()) {
+    BatchMetrics bm;
+    bm.processed_size = bm_proto.processed_size();
+    bm.input_size = bm_proto.input_size();
+    bm.padding_size = bm_proto.padding_size();
+    for (const auto& [cost_type, cost] : bm_proto.batch_costs()) {
+      bm.batch_costs[cost_type] = absl::Microseconds(cost);
+    }
+    batch_metrics_.push_back(std::move(bm));
+  }
+}
+
 }  // namespace tensorflow
diff --git a/tensorflow/core/common_runtime/request_cost.h b/tensorflow/core/common_runtime/request_cost.h
index 15418b02d39..cac400b25c7 100644
--- a/tensorflow/core/common_runtime/request_cost.h
+++ b/tensorflow/core/common_runtime/request_cost.h
@@ -27,6 +27,7 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/time/time.h"
+#include "tensorflow/core/common_runtime/request_cost.pb.h"
 
 namespace tensorflow {
 
@@ -115,6 +116,15 @@ class RequestCost {
   // rpc request, when all batch processing has completed.
   std::vector<BatchMetrics> GetBatchMetrics() const;
 
+  // Serializes the costs and metrics into `proto`.
+  // It's thread-safe.
+  void ToProto(RequestCostProto& proto) const;
+  RequestCostProto ToProto() const;
+
+  // Replaces the costs and metrics with those from `proto`.
+  // It's thread-safe.
+  void FromProto(const RequestCostProto& proto);
+
  private:
   mutable absl::Mutex mutex_;
 
diff --git a/tensorflow/core/common_runtime/request_cost.proto b/tensorflow/core/common_runtime/request_cost.proto
new file mode 100644
index 00000000000..64a4db54eb2
--- /dev/null
+++ b/tensorflow/core/common_runtime/request_cost.proto
@@ -0,0 +1,61 @@
+// 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.
+// ==============================================================================
+
+syntax = "proto3";
+
+package tensorflow;
+
+// Proto representation of `tensorflow::RequestCost`.
+message RequestCostProto {
+  // A cost entry broken down by named dimensions.
+  // Values are in microseconds.
+  message StructuredCost {
+    map<string, double> dimensions = 1;
+  }
+
+  // A metric entry carrying an array of doubles or an array of byte strings.
+  message StructuredMetric {
+    // Exactly one of the following fields will be set.
+    repeated double double_values = 1;
+    repeated bytes bytes_values = 2;
+  }
+
+  // Metrics of a batch that processes this rpc request.
+  message BatchMetrics {
+    // Size of the batch.
+    int64 processed_size = 1;
+    // In this batch, input size from this rpc request.
+    int64 input_size = 2;
+    // In this batch, the padding amount.
+    int64 padding_size = 3;
+    // Costs for processing this batch in microseconds.
+    map<string, double> batch_costs = 4;
+  }
+
+  // Accumulated costs in microseconds.
+  map<string, double> costs = 1;
+
+  // Accumulated structured costs in microseconds, keyed by cost type name.
+  map<string, StructuredCost> structured_costs = 2;
+
+  // User-defined metrics.
+  map<string, double> metrics = 3;
+
+  // User-defined structured metrics.
+  map<string, StructuredMetric> structured_metrics = 4;
+
+  // Metrics of batches that process this rpc request.
+  repeated BatchMetrics batch_metrics = 5;
+}
diff --git a/tensorflow/core/common_runtime/request_cost_test.cc b/tensorflow/core/common_runtime/request_cost_test.cc
index 8859491f03b..ef4b5171399 100644
--- a/tensorflow/core/common_runtime/request_cost_test.cc
+++ b/tensorflow/core/common_runtime/request_cost_test.cc
@@ -154,5 +154,135 @@ TEST(RequestCostTest, RecordStructuredCosts) {
   EXPECT_EQ(costs["cost_b"]["dim1"], absl::Milliseconds(3));
 }
 
+TEST(RequestCostTest, ToProtoAndFromProtoRoundTrip) {
+  RequestCost orig_cost;
+  orig_cost.RecordCost(
+      {{"tpu_v1", absl::Milliseconds(11)}, {"cpu_v1", absl::Milliseconds(30)}});
+  orig_cost.RecordStructuredCosts(
+      {{"tpu_v1",
+        {{"dim1", absl::Milliseconds(5)}, {"dim2", absl::Milliseconds(6)}}}});
+  orig_cost.RecordMetrics({{"latency", 100.0}, {"throughput", 2.0}});
+  orig_cost.RecordStructuredMetrics(
+      {{"blockwise_cache_read", {.values = std::vector<double>{10.0, 20.0}}},
+       {"blockwise_cache_ids",
+        {.values = std::vector<std::string>{"id1", "id2"}}}});
+  orig_cost.RecordBatchMetrics(RequestCost::BatchMetrics{
+      /*processed_size=*/8,
+      /*input_size=*/4,
+      /*padding_size=*/1,
+      {{"gcu", absl::Milliseconds(80)}, {"tpu", absl::Milliseconds(160)}}});
+
+  RequestCostProto proto = orig_cost.ToProto();
+
+  EXPECT_EQ(proto.costs().at("tpu_v1"), 11000.0);
+  EXPECT_EQ(proto.costs().at("cpu_v1"), 30000.0);
+  EXPECT_EQ(proto.structured_costs().at("tpu_v1").dimensions().at("dim1"),
+            5000.0);
+  EXPECT_EQ(proto.structured_costs().at("tpu_v1").dimensions().at("dim2"),
+            6000.0);
+  EXPECT_EQ(proto.metrics().at("latency"), 100.0);
+  EXPECT_EQ(proto.metrics().at("throughput"), 2.0);
+  EXPECT_THAT(
+      proto.structured_metrics().at("blockwise_cache_read").double_values(),
+      ElementsAre(10.0, 20.0));
+  EXPECT_THAT(
+      proto.structured_metrics().at("blockwise_cache_ids").bytes_values(),
+      ElementsAre("id1", "id2"));
+  ASSERT_EQ(proto.batch_metrics_size(), 1);
+  EXPECT_EQ(proto.batch_metrics(0).processed_size(), 8);
+  EXPECT_EQ(proto.batch_metrics(0).input_size(), 4);
+  EXPECT_EQ(proto.batch_metrics(0).padding_size(), 1);
+  EXPECT_EQ(proto.batch_metrics(0).batch_costs().at("gcu"), 80000.0);
+  EXPECT_EQ(proto.batch_metrics(0).batch_costs().at("tpu"), 160000.0);
+
+  // Deserialize into new RequestCost.
+  RequestCost new_cost;
+  new_cost.FromProto(proto);
+
+  EXPECT_THAT(new_cost.GetCosts(),
+              UnorderedElementsAre(Pair("tpu_v1", absl::Milliseconds(11)),
+                                   Pair("cpu_v1", absl::Milliseconds(30))));
+  EXPECT_THAT(new_cost.GetStructuredCosts(),
+              UnorderedElementsAre(Pair(
+                  "tpu_v1",
+                  UnorderedElementsAre(Pair("dim1", absl::Milliseconds(5)),
+                                       Pair("dim2", absl::Milliseconds(6))))));
+  EXPECT_THAT(
+      new_cost.GetMetrics(),
+      UnorderedElementsAre(Pair("latency", 100.0), Pair("throughput", 2.0)));
+
+  auto structured_metrics = new_cost.GetStructuredMetrics();
+  EXPECT_EQ(structured_metrics.size(), 2);
+  ASSERT_TRUE(std::holds_alternative<std::vector<double>>(
+      structured_metrics["blockwise_cache_read"].values));
+  EXPECT_THAT(std::get<std::vector<double>>(
+                  structured_metrics["blockwise_cache_read"].values),
+              ElementsAre(10.0, 20.0));
+  ASSERT_TRUE(std::holds_alternative<std::vector<std::string>>(
+      structured_metrics["blockwise_cache_ids"].values));
+  EXPECT_THAT(std::get<std::vector<std::string>>(
+                  structured_metrics["blockwise_cache_ids"].values),
+              ElementsAre("id1", "id2"));
+
+  EXPECT_THAT(new_cost.GetBatchMetrics(),
+              ElementsAre(FieldsAre(
+                  8, 4, 1,
+                  UnorderedElementsAre(Pair("gcu", absl::Milliseconds(80)),
+                                       Pair("tpu", absl::Milliseconds(160))))));
+}
+
+TEST(RequestCostTest, FromProtoReplacesExistingEntries) {
+  RequestCost cost;
+  cost.RecordCost({{"old_cost", absl::Milliseconds(10)}});
+  cost.RecordStructuredCosts(
+      {{"old_structured", {{"dim", absl::Milliseconds(1)}}}});
+  cost.RecordMetrics({{"old_metric", 1.0}});
+  cost.RecordStructuredMetrics(
+      {{"old_struct_metric", {.values = std::vector<double>{1.0}}}});
+  cost.RecordBatchMetrics(
+      RequestCost::BatchMetrics{/*processed_size=*/1,
+                                /*input_size=*/1,
+                                /*padding_size=*/0,
+                                {{"old_batch", absl::Milliseconds(10)}}});
+
+  RequestCostProto proto;
+  proto.mutable_costs()->insert({"new_cost", 20000.0});
+  (*proto.mutable_structured_costs())["new_structured"]
+      .mutable_dimensions()
+      ->insert({"new_dim", 3000.0});
+  proto.mutable_metrics()->insert({"new_metric", 42.0});
+  (*proto.mutable_structured_metrics())["new_struct_metric"]
+      .mutable_bytes_values()
+      ->Add("str_val");
+  auto* bm = proto.add_batch_metrics();
+  bm->set_processed_size(2);
+  bm->set_input_size(2);
+  bm->set_padding_size(0);
+  bm->mutable_batch_costs()->insert({"new_batch", 50000.0});
+
+  cost.FromProto(proto);
+
+  EXPECT_THAT(cost.GetCosts(),
+              UnorderedElementsAre(Pair("new_cost", absl::Milliseconds(20))));
+  EXPECT_THAT(cost.GetStructuredCosts(),
+              UnorderedElementsAre(Pair(
+                  "new_structured", UnorderedElementsAre(Pair(
+                                        "new_dim", absl::Milliseconds(3))))));
+  EXPECT_THAT(cost.GetMetrics(),
+              UnorderedElementsAre(Pair("new_metric", 42.0)));
+  auto structured_metrics = cost.GetStructuredMetrics();
+  EXPECT_EQ(structured_metrics.size(), 1);
+  ASSERT_TRUE(std::holds_alternative<std::vector<std::string>>(
+      structured_metrics["new_struct_metric"].values));
+  EXPECT_THAT(std::get<std::vector<std::string>>(
+                  structured_metrics["new_struct_metric"].values),
+              ElementsAre("str_val"));
+  EXPECT_THAT(
+      cost.GetBatchMetrics(),
+      ElementsAre(FieldsAre(
+          2, 2, 0,
+          UnorderedElementsAre(Pair("new_batch", absl::Milliseconds(50))))));
+}
+
 }  // namespace
 }  // namespace tensorflow

From 883ec7a57c2ad3a7eb2bb344bdc8553af940bd35 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 16 Oct 2026 16:13:23 -0700
Subject: [PATCH 0990/1000] Integrate LLVM at llvm/llvm-project@4fa92e8e51e6

Updates LLVM usage to match
[4fa92e8e51e6](https://github.com/llvm/llvm-project/commit/4fa92e8e51e6833f8d70ab9fa136d899e497f135)

PiperOrigin-RevId: 1001165416
---
 .../cuda_tile/patches/llvm_cl1000840178.patch | 105 ++++++++++++++++++
 .../xla/third_party/cuda_tile/patches/series  |   1 +
 .../xla/third_party/llvm/workspace.bzl        |   4 +-
 .../xla/third_party/shardy/temporary.patch    |  15 +++
 .../xla/third_party/shardy/workspace.bzl      |   4 +-
 .../triton/common/llvm_cl1000840178.patch     |  17 +++
 .../xla/third_party/triton/common/series      |   1 +
 .../triton/oss_only/build_files.patch         |   3 +-
 8 files changed, 145 insertions(+), 5 deletions(-)
 create mode 100755 third_party/xla/third_party/cuda_tile/patches/llvm_cl1000840178.patch
 create mode 100755 third_party/xla/third_party/triton/common/llvm_cl1000840178.patch

diff --git a/third_party/xla/third_party/cuda_tile/patches/llvm_cl1000840178.patch b/third_party/xla/third_party/cuda_tile/patches/llvm_cl1000840178.patch
new file mode 100755
index 00000000000..fea316933e8
--- /dev/null
+++ b/third_party/xla/third_party/cuda_tile/patches/llvm_cl1000840178.patch
@@ -0,0 +1,105 @@
+diff --git a/tools/cuda-tile-tblgen/AttrCodeGen.cpp b/tools/cuda-tile-tblgen/AttrCodeGen.cpp
+--- a/tools/cuda-tile-tblgen/AttrCodeGen.cpp
++++ b/tools/cuda-tile-tblgen/AttrCodeGen.cpp
+@@ -110,12 +110,12 @@
+ 
+ static const Record *findEnumByClassName(const RecordKeeper &records,
+                                          StringRef className) {
+-  for (const auto &def : records.getDefs()) {
+-    if (!isCudaTileEnumDef(records, def.second.get())) {
++  for (const Record* def : records.getDefs()) {
++    if (!isCudaTileEnumDef(records, def)) {
+       continue;
+     }
+-    if (def.second->getValueAsString("className") == className) {
+-      return def.second.get();
++    if (def->getValueAsString("className") == className) {
++      return def;
+     }
+   }
+   return nullptr;
+@@ -134,9 +134,9 @@
+   os << "#endif\n\n";
+ 
+   std::vector<EnumAttrInfo> infos;
+-  for (const auto &def : records.getDefs()) {
+-    if (isCudaTileEnumAttrWrapper(records, def.second.get())) {
+-      if (auto info = extractEnumAttrInfo(records, def.second.get())) {
++  for (const Record* def : records.getDefs()) {
++    if (isCudaTileEnumAttrWrapper(records, def)) {
++      if (auto info = extractEnumAttrInfo(records, def)) {
+         infos.push_back(*info);
+       }
+     }
+diff --git a/tools/cuda-tile-tblgen/SpecGen.cpp b/tools/cuda-tile-tblgen/SpecGen.cpp
+--- a/tools/cuda-tile-tblgen/SpecGen.cpp
++++ b/tools/cuda-tile-tblgen/SpecGen.cpp
+@@ -84,18 +84,15 @@
+ 
+   Regex includeRegex(opIncFilter), excludeRegex(opExcFilter);
+   std::vector<const Record *> defs;
+-  for (const auto &def : records.getDefs()) {
+-    if (!def.second->isSubClassOf(classDef))
+-      continue;
++  for (const Record* def : records.getDefs()) {
++    if (!def->isSubClassOf(classDef)) continue;
+     // Include if no include filter or include filter matches.
+-    if (!opIncFilter.empty() &&
+-        !includeRegex.match(getOperationName(*def.second)))
++    if (!opIncFilter.empty() && !includeRegex.match(getOperationName(*def)))
+       continue;
+     // Unless there is an exclude filter and it matches.
+-    if (!opExcFilter.empty() &&
+-        excludeRegex.match(getOperationName(*def.second)))
++    if (!opExcFilter.empty() && excludeRegex.match(getOperationName(*def)))
+       continue;
+-    defs.push_back(def.second.get());
++    defs.push_back(def);
+   }
+ 
+   return defs;
+@@ -565,11 +562,11 @@
+                            StringRef opClassName = "") {
+   // Find the HintKey enum.
+   const Record *hintKeyEnum = nullptr;
+-  for (const auto &def : records.getDefs()) {
++  for (const Record* def : records.getDefs()) {
+     const Record *cls = records.getClass("CudaTileI32EnumAttr");
+-    if (cls && def.second->isSubClassOf(cls) &&
+-        def.second->getValueAsString("className") == "HintKey") {
+-      hintKeyEnum = def.second.get();
++    if (cls && def->isSubClassOf(cls) &&
++        def->getValueAsString("className") == "HintKey") {
++      hintKeyEnum = def;
+       break;
+     }
+   }
+@@ -580,11 +577,11 @@
+   // Collect family -> SM prefix mapping for arch override docs.
+   std::map<std::string, std::vector<std::string>> familyPrefixes;
+   const Record *gpuArchEnum = nullptr;
+-  for (const auto &def : records.getDefs()) {
++  for (const Record* def : records.getDefs()) {
+     const Record *cls = records.getClass("CudaTileI32EnumAttr");
+-    if (cls && def.second->isSubClassOf(cls) &&
+-        def.second->getValueAsString("className") == "GpuArchitecture") {
+-      gpuArchEnum = def.second.get();
++    if (cls && def->isSubClassOf(cls) &&
++        def->getValueAsString("className") == "GpuArchitecture") {
++      gpuArchEnum = def;
+       break;
+     }
+   }
+@@ -888,9 +885,9 @@
+   auto orderedSections = splitBySections(records);
+ 
+   std::vector<const Record *> attrDefs;
+-  for (const auto &def : records.getDefs()) {
+-    if (def.second->isSubClassOf("AttrDef")) {
+-      attrDefs.push_back(def.second.get());
++  for (const Record* def : records.getDefs()) {
++    if (def->isSubClassOf("AttrDef")) {
++      attrDefs.push_back(def);
+     }
+   }
+ 
diff --git a/third_party/xla/third_party/cuda_tile/patches/series b/third_party/xla/third_party/cuda_tile/patches/series
index 46a712b02b5..19350d3efc7 100644
--- a/third_party/xla/third_party/cuda_tile/patches/series
+++ b/third_party/xla/third_party/cuda_tile/patches/series
@@ -2,3 +2,4 @@ constructor.patch
 unused_variable.patch
 pow_op_alias.patch
 llvm_cl_989877215.patch
+llvm_cl1000840178.patch
diff --git a/third_party/xla/third_party/llvm/workspace.bzl b/third_party/xla/third_party/llvm/workspace.bzl
index caf6ccb58d8..8bde547d57e 100644
--- a/third_party/xla/third_party/llvm/workspace.bzl
+++ b/third_party/xla/third_party/llvm/workspace.bzl
@@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
 
 def repo(name):
     """Imports LLVM."""
-    LLVM_COMMIT = "1ee8805f01a663b7f2b09215a8eab5ab2d15ce35"
-    LLVM_SHA256 = "273fa5e9618fa8085aea51d6eb40129677fea8cf4c8825a2c4ff6ea582706a55"
+    LLVM_COMMIT = "4fa92e8e51e6833f8d70ab9fa136d899e497f135"
+    LLVM_SHA256 = "3e33d6506a6d9e779dda03c355310e58269b3c90aa70bb2de9a1203c449aedff"
 
     tf_http_archive(
         name = name,
diff --git a/third_party/xla/third_party/shardy/temporary.patch b/third_party/xla/third_party/shardy/temporary.patch
index e69de29bb2d..09c2e1b5f72 100644
--- a/third_party/xla/third_party/shardy/temporary.patch
+++ b/third_party/xla/third_party/shardy/temporary.patch
@@ -0,0 +1,15 @@
+diff --git a/third_party/llvm/workspace.bzl b/third_party/llvm/workspace.bzl
+index caf6ccb5..8bde547d 100644
+--- a/third_party/llvm/workspace.bzl
++++ b/third_party/llvm/workspace.bzl
+@@ -19,8 +19,8 @@ load("//third_party:repo.bzl", "tf_http_archive")
+ 
+ def repo(name):
+     """Imports LLVM."""
+-    LLVM_COMMIT = "1ee8805f01a663b7f2b09215a8eab5ab2d15ce35"
+-    LLVM_SHA256 = "273fa5e9618fa8085aea51d6eb40129677fea8cf4c8825a2c4ff6ea582706a55"
++    LLVM_COMMIT = "4fa92e8e51e6833f8d70ab9fa136d899e497f135"
++    LLVM_SHA256 = "3e33d6506a6d9e779dda03c355310e58269b3c90aa70bb2de9a1203c449aedff"
+ 
+     tf_http_archive(
+         name = name,
diff --git a/third_party/xla/third_party/shardy/workspace.bzl b/third_party/xla/third_party/shardy/workspace.bzl
index cd32b006ffe..732b54292a9 100644
--- a/third_party/xla/third_party/shardy/workspace.bzl
+++ b/third_party/xla/third_party/shardy/workspace.bzl
@@ -18,8 +18,8 @@
 load("//third_party:repo.bzl", "tf_http_archive", "tf_mirror_urls")
 
 def repo():
-    SHARDY_COMMIT = "c29509bd0c1430beaa7e767c866a131492ab1139"
-    SHARDY_SHA256 = "e67e1c6082302dd7572cc19f2dc53045def9b5eae350ea6d1a41dd236b08d775"
+    SHARDY_COMMIT = "8a1a867c0e9cc44abaeddc83515d1ac621d375ee"
+    SHARDY_SHA256 = "37822de88945d9d20339bbcbbf12a2967a64b599f3bc3eda1a68f777e32644b4"
 
     tf_http_archive(
         name = "shardy",
diff --git a/third_party/xla/third_party/triton/common/llvm_cl1000840178.patch b/third_party/xla/third_party/triton/common/llvm_cl1000840178.patch
new file mode 100755
index 00000000000..f61de2c41f1
--- /dev/null
+++ b/third_party/xla/third_party/triton/common/llvm_cl1000840178.patch
@@ -0,0 +1,17 @@
+--- a/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp
++++ b/third_party/nvidia/lib/TritonNVIDIAGPUToLLVM/TritonGPUToLLVM.cpp
+@@ -1,12 +1,12 @@
+ #include "Dialect/NVGPU/IR/Dialect.h"
+ #include "PatternTritonGPUOpToLLVM.h"
+ #include "TritonNVIDIAGPUToLLVM/Passes.h"
++#include "mlir/Conversion/ArithCommon/CeilFloorDivExpansion.h"
+ #include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"
+ #include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"
+ #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
+ #include "mlir/Conversion/MathToLLVM/MathToLLVM.h"
+ #include "mlir/Conversion/UBToLLVM/UBToLLVM.h"
+-#include "mlir/Dialect/Arith/Transforms/Passes.h"
+ #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
+ #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+ #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
+
diff --git a/third_party/xla/third_party/triton/common/series b/third_party/xla/third_party/triton/common/series
index 7661facc3a2..ba6b5bee674 100644
--- a/third_party/xla/third_party/triton/common/series
+++ b/third_party/xla/third_party/triton/common/series
@@ -38,3 +38,4 @@ fix_load_ops.patch
 cxx17_support.patch
 llvm_cl996124567.patch
 llvm_cl998243985.patch
+llvm_cl1000840178.patch
diff --git a/third_party/xla/third_party/triton/oss_only/build_files.patch b/third_party/xla/third_party/triton/oss_only/build_files.patch
index 5519e5cace5..d6bd46b4e88 100644
--- a/third_party/xla/third_party/triton/oss_only/build_files.patch
+++ b/third_party/xla/third_party/triton/oss_only/build_files.patch
@@ -2854,7 +2854,7 @@ diff --git b/third_party/nvidia/BUILD b/third_party/nvidia/BUILD
 new file mode 100644
 --- /dev/null
 +++ b/third_party/nvidia/BUILD
-@@ -0,0 +1,568 @@
+@@ -0,0 +1,569 @@
 +load("@rules_cc//cc:cc_library.bzl", "cc_library")
 +load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
 +
@@ -2986,6 +2986,7 @@ new file mode 100644
 +        "@llvm-project//llvm:Support",
 +        "@llvm-project//llvm:ir_headers",
 +        "@llvm-project//mlir:Analysis",
++        "@llvm-project//mlir:ArithCommonConversion",
 +        "@llvm-project//mlir:ArithToLLVM",
 +        "@llvm-project//mlir:ArithTransforms",
 +        "@llvm-project//mlir:ControlFlowDialect",

From ae7e665e5f29fda9ed1548d523ca7f16c0121e2c Mon Sep 17 00:00:00 2001
From: Ziyin Huang <ziyinh@google.com>
Date: Fri, 16 Oct 2026 16:20:55 -0700
Subject: [PATCH 0991/1000] Add `spanner_fdo_key` field to
 `SparseCoreTableLayout` proto for non-abbreviated Spanner FDO table keys.

PiperOrigin-RevId: 1001169412
---
 tensorflow/core/tpu/kernels/sparse_core_layout.proto | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tensorflow/core/tpu/kernels/sparse_core_layout.proto b/tensorflow/core/tpu/kernels/sparse_core_layout.proto
index 81764e3ec18..79af61c7828 100644
--- a/tensorflow/core/tpu/kernels/sparse_core_layout.proto
+++ b/tensorflow/core/tpu/kernels/sparse_core_layout.proto
@@ -57,6 +57,10 @@ message SparseCoreTableLayout {
   // The key used to look up the shared FDO config for this table.
   string fdo_key = 12;
 
+  // The non-abbreviated key used to look up the shared FDO config for this
+  // table in Spanner.
+  string spanner_fdo_key = 13;
+
   // The name of the stack it wound up in. The sparse core python library uses
   // this name as a variable name.
   string stacked_table_name = 2;

From 3b5883f49798646a739d15961856dd4a1314aa53 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 16 Oct 2026 16:52:03 -0700
Subject: [PATCH 0992/1000] Use relative tolerance in
 `DefaultSchedulerCore::AddOccupierToResource` latency check.

PiperOrigin-RevId: 1001186899
---
 .../xla/xla/service/latency_hiding_scheduler.cc      |  6 ++++--
 .../xla/xla/service/latency_hiding_scheduler_test.cc | 12 ++++++++++++
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/third_party/xla/xla/service/latency_hiding_scheduler.cc b/third_party/xla/xla/service/latency_hiding_scheduler.cc
index e819190eeab..0bf60560f52 100644
--- a/third_party/xla/xla/service/latency_hiding_scheduler.cc
+++ b/third_party/xla/xla/service/latency_hiding_scheduler.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "xla/service/latency_hiding_scheduler.h"
 
 #include <algorithm>
+#include <cmath>
 #include <cstdint>
 #include <cstdlib>
 #include <cstring>
@@ -2423,8 +2424,9 @@ bool DefaultSchedulerCore::AddOccupierToResource(
   // Since it points to the newly inserted element, increment it
   it++;
   accumulated_delay += new_edge_remaining;
-  CHECK(new_edge.OriginalLatency() - 0.0001 < accumulated_delay &&
-        accumulated_delay < new_edge.OriginalLatency() + 0.0001);
+  const HloGraphNode::TimeCost tolerance =
+      std::max(1e-4, 1e-9 * new_edge.OriginalLatency());
+  CHECK_LE(std::abs(accumulated_delay - new_edge.OriginalLatency()), tolerance);
   for (; it != occupiers.end(); it++) {
     it->second += accumulated_delay;
   }
diff --git a/third_party/xla/xla/service/latency_hiding_scheduler_test.cc b/third_party/xla/xla/service/latency_hiding_scheduler_test.cc
index 297e5ac6a83..4c6ee4ad58d 100644
--- a/third_party/xla/xla/service/latency_hiding_scheduler_test.cc
+++ b/third_party/xla/xla/service/latency_hiding_scheduler_test.cc
@@ -3475,6 +3475,18 @@ TEST_F(LatencyHidingSchedulerTest, AddDeleteOccupierForSharedResource) {
   DefaultSchedulerCore::AddOccupierToResource(0, edge3, occupiers);
   DefaultSchedulerCore::DeleteOccupierFromResource(4, edge3, occupiers);
   CHECK(check_eq({3, 4.5}));
+
+  // Large latencies where 0.0001 is smaller than 0.5 ULP of OriginalLatency().
+  occupiers.clear();
+  HloEdge large_edge1(1e13, nullptr);
+  HloEdge large_edge2(2e13, nullptr);
+  HloEdge large_edge3(3e13, nullptr);
+  EXPECT_TRUE(
+      DefaultSchedulerCore::AddOccupierToResource(0, large_edge1, occupiers));
+  EXPECT_TRUE(
+      DefaultSchedulerCore::AddOccupierToResource(0, large_edge2, occupiers));
+  EXPECT_TRUE(
+      DefaultSchedulerCore::AddOccupierToResource(0, large_edge3, occupiers));
 }
 
 TEST_F(LatencyHidingSchedulerTest, DepthPressureReduction) {

From 6d648305df6b05e4d62651f0ca1fead0654374dd Mon Sep 17 00:00:00 2001
From: Sean Talts <talts@google.com>
Date: Fri, 16 Oct 2026 17:04:35 -0700
Subject: [PATCH 0993/1000] [XLA:CPU] Fold sin(-x) and cos(-x) before intrinsic
 lowering.

Once `math.sin` and `math.cos` are lowered to calls to `xla.sin.*` and `xla.cos.*` in `xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc`, LLVM `InstCombine` no longer recognizes them as odd and even functions and cannot fold `sin(-x) -> -sin(x)` or `cos(-x) -> cos(x)`.

Apply these parity folds in `LowerXlaIntrinsicLibPass` before lowering `math.sin` and `math.cos` to `xla.sin.*` and `xla.cos.*`. Both folds are exact for the XLA implementations, which reduce `|x|` and apply the sign of `x` for `sin`. This preserves CSE across `sin(x)`/`sin(-x)` and `cos(x)`/`cos(-x)` and keeps downstream expression canonicalization consistent with `@llvm.sin` and `@llvm.cos`.

PiperOrigin-RevId: 1001193850
---
 .../transforms/lower_xla_intrinsic_lib.cc     | 34 +++++++++++++++++++
 .../tests/lower_xla_intrinsic_lib.mlir        | 14 ++++++++
 2 files changed, 48 insertions(+)

diff --git a/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc b/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc
index 64316fcb6fa..63fda921bed 100644
--- a/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc
+++ b/third_party/xla/xla/codegen/emitters/transforms/lower_xla_intrinsic_lib.cc
@@ -325,6 +325,39 @@ class LowerIntrinsicPattern : public mlir::OpRewritePattern<Op> {
   std::string cpu_features_;
 };
 
+// Once math.sin / math.cos are lowered to calls to `xla.sin.*` / `xla.cos.*`,
+// LLVM no longer knows that sin is odd and cos is even, so it cannot apply the
+// folds InstCombine does for `llvm.sin` / `llvm.cos` (sin(-x) -> -sin(x),
+// cos(-x) -> cos(x)). Apply them here, before lowering. They are exact for the
+// XLA implementations, which reduce |x| and then apply the sign of x for sin.
+struct FoldSinOfNeg : public mlir::OpRewritePattern<mm::SinOp> {
+  using OpRewritePattern::OpRewritePattern;
+  mlir::LogicalResult matchAndRewrite(
+      mm::SinOp op, mlir::PatternRewriter& rewriter) const override {
+    auto neg = op.getOperand().getDefiningOp<ma::NegFOp>();
+    if (!neg) {
+      return rewriter.notifyMatchFailure(op, "operand is not negf");
+    }
+    Value sin = mm::SinOp::create(rewriter, op.getLoc(), neg.getOperand(),
+                                  op.getFastmathAttr());
+    rewriter.replaceOpWithNewOp<ma::NegFOp>(op, sin);
+    return mlir::success();
+  }
+};
+
+struct FoldCosOfNeg : public mlir::OpRewritePattern<mm::CosOp> {
+  using OpRewritePattern::OpRewritePattern;
+  mlir::LogicalResult matchAndRewrite(
+      mm::CosOp op, mlir::PatternRewriter& rewriter) const override {
+    auto neg = op.getOperand().getDefiningOp<ma::NegFOp>();
+    if (!neg) {
+      return rewriter.notifyMatchFailure(op, "operand is not negf");
+    }
+    rewriter.modifyOpInPlace(op, [&] { op->setOperand(0, neg.getOperand()); });
+    return mlir::success();
+  }
+};
+
 class LowerXlaIntrinsicLibPass
     : public impl::LowerXlaIntrinsicLibPassBase<LowerXlaIntrinsicLibPass> {
  public:
@@ -347,6 +380,7 @@ class LowerXlaIntrinsicLibPass
                  LowerIntrinsicPattern<ci::Erf, mm::ErfOp>,
                  LowerIntrinsicPattern<ci::Sinh, mm::SinhOp>>(context,
                                                               cpu_features);
+    patterns.add<FoldSinOfNeg, FoldCosOfNeg>(context, /*benefit=*/2);
     if (mlir::failed(
             mlir::applyPatternsGreedily(module_op, std::move(patterns)))) {
       signalPassFailure();
diff --git a/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir b/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir
index 0deade04284..2e13a338a03 100644
--- a/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir
+++ b/third_party/xla/xla/codegen/emitters/transforms/tests/lower_xla_intrinsic_lib.mlir
@@ -480,3 +480,17 @@ module {
 // CHECK-NOT: math.sinh
 // CHECK: %[[SINH_CALL:.*]] = call @xla.sinh.v8f32(%arg0) : (vector<8xf32>) -> vector<8xf32>
 // CHECK: return %[[SINH_CALL]]
+
+// -----
+
+func.func @sin_cos_of_neg(%arg0: f32) -> (f32, f32) {
+  %n = arith.negf %arg0 : f32
+  %s = math.sin %n : f32
+  %c = math.cos %n : f32
+  return %s, %c : f32, f32
+}
+// CHECK-LABEL: @sin_cos_of_neg
+// CHECK-DAG: %[[SIN:.*]] = call @xla.sin.f32(%arg0) : (f32) -> f32
+// CHECK-DAG: %[[NEG_SIN:.*]] = arith.negf %[[SIN]] : f32
+// CHECK-DAG: %[[COS:.*]] = call @xla.cos.f32(%arg0) : (f32) -> f32
+// CHECK: return %[[NEG_SIN]], %[[COS]]

From cae52ab735293bcf2743802662108207ee86402e Mon Sep 17 00:00:00 2001
From: Junwhan Ahn <junwhan@google.com>
Date: Fri, 16 Oct 2026 18:17:02 -0700
Subject: [PATCH 0994/1000] Wait for receiver allocation sync points before
 cross-host GPU recv.

Under `LocalDeviceState::kComputeSynchronized`, `BFCAllocator` reuses device memory from completed host launches while prior kernels may still be executing on the receiver's compute stream. `StreamExecutorGpuRawClient::CrossHostReceiveBuffersInto` previously called `WaitForAllocation` on the receiver's device-to-device stream before immediately exporting fabric handles and invoking `notifier`. Because the actual cross-host NVLink `Memcpy` executes on the sender's device-to-device stream in `ScheduleRemoteSend` as soon as the descriptor arrives, enqueuing a wait on the receiver's device-to-device stream did not prevent the sender from writing into recycled HBM while prior kernels on the receiver's compute stream were still accessing it.

Wait for `transfer_dependency_avs` and each buffer's materialized compute-stream allocation event to become ready on the host before exporting fabric handles and invoking `notifier`, and report early errors to `notifier`. Also retain the imported `dst_base` allocation in `ScheduleRemoteSend` until `usage_event` completes so that uncached fabric handles (`!preallocate`) are not unmapped while the async `Memcpy` is still in flight.

PiperOrigin-RevId: 1001224843
---
 third_party/xla/xla/pjrt/gpu/BUILD            |  97 -------
 .../xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc    | 269 +++++++++++-------
 .../se_gpu_pjrt_cross_host_transfer_test.cc   | 138 +++++++++
 3 files changed, 305 insertions(+), 199 deletions(-)

diff --git a/third_party/xla/xla/pjrt/gpu/BUILD b/third_party/xla/xla/pjrt/gpu/BUILD
index 0b8d6fdfddc..08ebdd03c3f 100644
--- a/third_party/xla/xla/pjrt/gpu/BUILD
+++ b/third_party/xla/xla/pjrt/gpu/BUILD
@@ -113,8 +113,6 @@ cc_library(
         "//xla:debug_options_flags",
         "//xla:executable_run_options",
         "//xla:future",
-        "//xla:literal",
-        "//xla:shape_tree",
         "//xla:shape_util",
         "//xla:status_macros",
         "//xla:util",
@@ -128,7 +126,6 @@ cc_library(
         "//xla/backends/gpu/collectives:gpu_collectives",
         "//xla/backends/gpu/collectives:gpu_communicator",
         "//xla/backends/gpu/target_config",
-        "//xla/client:client_library",
         "//xla/client:local_client",
         "//xla/core/collectives",
         "//xla/core/collectives:clique_id",
@@ -143,7 +140,6 @@ cc_library(
         "//xla/pjrt:host_memory_allocator",
         "//xla/pjrt:host_memory_spaces",
         "//xla/pjrt:maybe_owning_mlir_module",
-        "//xla/pjrt:mlir_to_hlo",
         "//xla/pjrt:pjrt_abi_version",
         "//xla/pjrt:pjrt_client",
         "//xla/pjrt:pjrt_common",
@@ -151,9 +147,7 @@ cc_library(
         "//xla/pjrt:pjrt_device_description",
         "//xla/pjrt:pjrt_executable",
         "//xla/pjrt:raw_buffer",
-        "//xla/pjrt:utils",
         "//xla/pjrt:worker_thread",
-        "//xla/pjrt/distributed:client",
         "//xla/pjrt/distributed:in_memory_key_value_store",
         "//xla/pjrt/distributed:key_value_store_interface",
         "//xla/pjrt/distributed:protocol_proto_cc",
@@ -162,24 +156,16 @@ cc_library(
         "//xla/pjrt/plugin/xla_gpu:xla_gpu_client_options",
         "//xla/pjrt/proto:compile_options_proto_cc",
         "//xla/pjrt/proto:topology_description_proto_cc",
-        "//xla/pjrt/se:event_pool",
         "//xla/pjrt/se:local_device_state",
         "//xla/pjrt/se:pjrt_stream_executor_client",
-        "//xla/pjrt/se:stream_executor_executable",
         "//xla/pjrt/se:stream_executor_executable_proto_cc",
-        "//xla/pjrt/se:stream_executor_pjrt_abi_version",
         "//xla/runtime:device_id",
         "//xla/runtime:hang_watchdog",
         "//xla/runtime:process_id",
         "//xla/service:buffer_assignment",
         "//xla/service:compiler",
-        "//xla/service:executable",
         "//xla/service:gpu_topology",
         "//xla/service:gpu_topology_proto_cc",
-        "//xla/service:pjrt_gpu_utils",
-        "//xla/service:platform_util",
-        "//xla/service:shaped_buffer",
-        "//xla/service:transfer_manager",
         "//xla/service/gpu:buffer_allocations",
         "//xla/service/gpu:gpu_constants",
         "//xla/service/gpu:gpu_executable_run_options",
@@ -197,27 +183,19 @@ cc_library(
         "//xla/stream_executor:stream_executor_address_allocator",
         "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor/gpu:gpu_executor_header",
-        "//xla/stream_executor/integrations:device_mem_allocator",
         "//xla/stream_executor/integrations:tf_allocator_adapter",
         "//xla/tsl/concurrency:async_value",
         "//xla/tsl/concurrency:ref_count",
-        "//xla/tsl/distributed_runtime/coordination:coordination_service",
         "//xla/tsl/framework:allocator",
-        "//xla/tsl/framework:bfc_allocator",
-        "//xla/tsl/framework:device_id",
-        "//xla/tsl/framework:device_id_impl",
         "//xla/tsl/framework:scoped_allocation_trace",
-        "//xla/tsl/lib/strings:proto_serialization",
         "//xla/tsl/platform:env",
         "//xla/tsl/platform:errors",
-        "//xla/tsl/platform:status",
         "//xla/tsl/platform:statusor",
         "//xla/tsl/protobuf:coordination_service_proto_cc",
         "//xla/tsl/util:env_var",
         "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base",
         "@com_google_absl//absl/base:core_headers",
-        "@com_google_absl//absl/cleanup",
         "@com_google_absl//absl/container:btree",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/container:inlined_vector",
@@ -225,7 +203,6 @@ cc_library(
         "@com_google_absl//absl/functional:bind_front",
         "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
-        "@com_google_absl//absl/memory",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:statusor",
@@ -235,17 +212,10 @@ cc_library(
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
         "@com_google_absl//absl/types:span",
-        "@llvm-project//mlir:IR",
-        "@tsl//tsl/platform:casts",
-        "@tsl//tsl/platform:env",
-        "@tsl//tsl/platform:errors",
         "@tsl//tsl/platform:fingerprint",
         "@tsl//tsl/platform:platform_port",
         "@tsl//tsl/platform:protobuf",
         "@tsl//tsl/platform:random",
-        "@tsl//tsl/platform:status",
-        "@tsl//tsl/platform:statusor",
-        "@tsl//tsl/profiler/lib:connected_traceme",
         "@tsl//tsl/profiler/lib:nvtx_utils",
         "@tsl//tsl/profiler/lib:traceme",
     ] + if_cuda_or_rocm([
@@ -636,92 +606,25 @@ xla_test(
         ":se_gpu_pjrt_client",
         ":se_gpu_pjrt_client_test_lib",
         ":se_gpu_topology_description",
-        "//xla:debug_options_flags",
         "//xla:future",
         "//xla:literal",
         "//xla:literal_util",
         "//xla:shape_util",
-        "//xla:status_macros",
-        "//xla:types",
-        "//xla:util",
-        "//xla:xla_data_proto_cc",
-        "//xla:xla_proto_cc",
-        "//xla/backends/gpu:ffi",
-        "//xla/ffi",
-        "//xla/ffi:ffi_api",
-        "//xla/hlo/builder:xla_computation",
-        "//xla/hlo/parser:hlo_parser",
-        "//xla/hlo/testlib:test",
         "//xla/pjrt:common_pjrt_client",
-        "//xla/pjrt:device_event",
-        "//xla/pjrt:host_memory_spaces",
-        "//xla/pjrt:maybe_owning_mlir_module",
-        "//xla/pjrt:mlir_to_hlo",
-        "//xla/pjrt:pjrt_abi_version",
         "//xla/pjrt:pjrt_client",
-        "//xla/pjrt:pjrt_common",
-        "//xla/pjrt:pjrt_compiler",
-        "//xla/pjrt:pjrt_device_description",
-        "//xla/pjrt:pjrt_executable",
-        "//xla/pjrt:raw_buffer",
-        "//xla/pjrt/distributed",
-        "//xla/pjrt/distributed:client",
-        "//xla/pjrt/distributed:in_memory_key_value_store",
-        "//xla/pjrt/distributed:service",
         "//xla/pjrt/plugin/xla_gpu:xla_gpu_client_options",
-        "//xla/pjrt/profiling:device_time_measurement",
-        "//xla/pjrt/proto:compile_options_proto_cc",
         "//xla/pjrt/se:local_device_state",
         "//xla/pjrt/se:pjrt_stream_executor_client",
-        "//xla/runtime:device_id",
-        "//xla/service:gpu_topology",
-        "//xla/service:gpu_topology_proto_cc",
-        "//xla/service:platform_util",
-        "//xla/service/gpu:gpu_memory_space_assignment",
-        "//xla/stream_executor:device_address",
-        "//xla/stream_executor:stream",
         "//xla/tests:literal_test_util",
-        "//xla/tsl/concurrency:async_value",
-        "//xla/tsl/concurrency:ref_count",
-        "//xla/tsl/lib/core:status_test_util",
-        "//xla/tsl/platform:env",
-        "//xla/tsl/platform:errors",
-        "//xla/tsl/platform:status",
-        "//xla/tsl/platform:statusor",
-        "//xla/tsl/platform:subprocess",
         "//xla/tsl/platform:test_main",
-        "//xla/tsl/util:command_line_flags",
-        "//xla/util/split_proto:split_executable_and_options_writer",
-        "//xla/util/split_proto:split_proto_reader",
-        "@com_google_absl//absl/algorithm:container",
         "@com_google_absl//absl/base",
-        "@com_google_absl//absl/base:log_severity",
-        "@com_google_absl//absl/cleanup",
-        "@com_google_absl//absl/container:flat_hash_map",
-        "@com_google_absl//absl/log",
         "@com_google_absl//absl/log:check",
-        "@com_google_absl//absl/log:globals",
-        "@com_google_absl//absl/log:scoped_mock_log",
         "@com_google_absl//absl/status",
-        "@com_google_absl//absl/status:status_macros",
         "@com_google_absl//absl/status:status_matchers",
         "@com_google_absl//absl/status:statusor",
-        "@com_google_absl//absl/strings",
-        "@com_google_absl//absl/strings:str_format",
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
-        "@com_google_absl//absl/types:span",
         "@com_google_googletest//:gtest",
-        "@com_google_protobuf//:protobuf",
-        "@llvm-project//mlir:IR",
-        "@riegeli//riegeli/bytes:reader",
-        "@riegeli//riegeli/bytes:string_reader",
-        "@riegeli//riegeli/bytes:string_writer",
-        "@tsl//tsl/platform",
-        "@tsl//tsl/platform:casts",
-        "@tsl//tsl/platform:platform_port",
-        "@tsl//tsl/platform:status_matchers",
-        "@tsl//tsl/platform:test",
     ] + if_cuda([
         "//xla/stream_executor/cuda:cuda_compute_capability",
         "//xla/stream_executor/cuda:cuda_device_address_vmm_allocator",
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
index f2f65fe0d3e..a25c260df6d 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_client.cc
@@ -874,6 +874,7 @@ void StreamExecutorGpuRawClient::ScheduleRemoteSend(
                   ABSL_ASSIGN_OR_RETURN(
                       std::shared_ptr<se::DeviceAddressBase> dst_base,
                       GetOrImportFabricHandle(executor, desc.buffer_handle()));
+                  usage_event.AndThen([dst_base]() {});
 
                   // Copy the source buffer into the destination buffer.
                   se::DeviceAddressBase dst(
@@ -1014,11 +1015,17 @@ StreamExecutorGpuRawClient::CrossHostReceiveBuffersInto(
           "same device");
     }
     raw_buffers.push_back(raw_buffer);
-    // `WaitForAllocation` is deferred to `async_work_runner()` below, so record
-    // the allocation event now.
+    // Obtaining the allocation event is deferred to `async_work_runner()`
+    // below, so record the allocation event now.
     MaterializeAllocationEvent(*raw_buffer);
   }
 
+  PjRtDeviceEventRefVector transfer_dependencies;
+  transfer_dependencies.reserve(transfer_dependency_avs.size());
+  for (size_t i = 0; i < transfer_dependency_avs.size(); ++i) {
+    transfer_dependencies.push_back(transfer_dependency_avs[i].CopyRef());
+  }
+
   // All buffers are on the same device.
   LocalDeviceState* local_device =
       raw_buffers[0]->down_cast<PjRtStreamExecutorRawBuffer>()->local_device();
@@ -1031,107 +1038,165 @@ StreamExecutorGpuRawClient::CrossHostReceiveBuffersInto(
         BufferSequencingEvent::Create(this->async_work_runner()));
   }
 
-  auto recv = [this, raw_buffers = std::move(raw_buffers),
-               buffer_sequencing_events, notifier = std::move(notifier),
-               local_device, stream]() mutable {
-    const uint64_t transfer_id = tsl::random::New64();
-    const std::string transfer_key =
-        absl::StrCat("pjrt_gpu_remote_recv_", transfer_id);
-
-    auto results = [&]() -> absl::StatusOr<std::vector<Future<>>> {
-      auto* executor =
-          absl::down_cast<se::gpu::GpuExecutor*>(local_device->executor());
-
-      StreamExecutorGpuCrossHostRecvDescriptor desc;
-      std::vector<PjRtCrossHostRecvDescriptors> serialized_descriptors;
-      serialized_descriptors.reserve(raw_buffers.size());
-      std::vector<std::string> rendezvous_keys;
-      rendezvous_keys.reserve(raw_buffers.size());
-
-      for (int i = 0; i < raw_buffers.size(); ++i) {
-        auto* se_raw_buffer =
-            raw_buffers[i]->down_cast<PjRtStreamExecutorRawBuffer>();
-
-        tsl::AsyncValueRef<RawSEDeviceMemory> mem =
-            se_raw_buffer->device_buffer();
-        ABSL_RETURN_IF_ERROR(WaitForAllocation(stream, *raw_buffers[i]));
-
-        // Keep mem alive until the Recv has finished executing.
-        buffer_sequencing_events[i].AndThen([mem]() {});
-
-        std::string rendezvous_key = absl::StrCat(transfer_key, "/", i);
-        rendezvous_keys.push_back(rendezvous_key);
-
-        // Export the buffer fabric handle and use it as a descriptor to be sent
-        // to the sender.
-        desc.Clear();
-        desc.set_buffer_index(i);
-        if (mem->mem().size() > 0) {
-          ABSL_ASSIGN_OR_RETURN(auto range,
-                           executor->GetAllocationRange(mem->mem().opaque()));
-          ABSL_ASSIGN_OR_RETURN(*desc.mutable_buffer_handle(),
-                           GetOrExportFabricHandle(executor, range.opaque()));
-          desc.set_buffer_offset(
-              reinterpret_cast<intptr_t>(mem->mem().opaque()) -
-              reinterpret_cast<intptr_t>(range.opaque()));
-        }
-        desc.set_rendezvous_key(rendezvous_key);
-
-        serialized_descriptors.push_back(
-            PjRtCrossHostRecvDescriptors{{desc.SerializeAsString()}});
-      }
-
-      auto state = std::make_shared<CrossHostRecvState>(
-          std::move(rendezvous_keys), kv_store_);
-
-      // Notify the receiver of the descriptors and cancellation callback. The
-      // caller is responsible for sending the descriptors to the sender and/or
-      // cancelling the transfer if needed.
-      VLOG(3) << "Notifying receiver of descriptors for cross-host recv of "
-              << serialized_descriptors.size() << " buffers";
-      notifier(PjRtCrossHostRecvState{
-          /*descriptors=*/std::move(serialized_descriptors),
-          /*cancel_notifier=*/
-          absl::bind_front(&CrossHostRecvState::NotifyCancellation, state),
-      });
-
-      return state->futures();
-    }();
-    if (!results.ok()) {
-      VLOG(2) << "CrossHostReceiveBuffersInto failed: " << results.status();
-      for (const auto& buffer_sequencing_event : buffer_sequencing_events) {
-        SetEventAsError(buffer_sequencing_event, results.status());
-      }
-      return;
-    }
-
-    CHECK_EQ(results->size(), buffer_sequencing_events.size());
-    for (int i = 0; i < buffer_sequencing_events.size(); ++i) {
-      (*results)[i].OnReady(
-          [this, local_device, stream,
-           event = buffer_sequencing_events[i]](const absl::Status& status) {
-            if (status.ok()) {
-              absl::Status s = AllocateAndRecordEvent(
-                  event, local_device, stream, "CrossHostReceiveBuffersInto");
-              if (!s.ok()) {
-                SetEventAsError(event, s);
-              }
-            } else {
-              SetEventAsError(event, status);
-            }
-          });
-    }
-
-    JoinFutures(*results).OnReady(
-        *async_work_runner(),
-        [kv_store = kv_store_, transfer_key](const absl::Status& /*status*/) {
-          if (absl::Status s = kv_store->Delete(transfer_key); !s.ok()) {
-            LOG(DFATAL) << "Failed to delete rendezvous keys for transfer "
-                        << transfer_key << ": " << s;
+  ExecuteWhenReady(
+      transfer_dependency_avs, async_work_runner(),
+      [this, raw_buffers = std::move(raw_buffers),
+       transfer_dependencies = std::move(transfer_dependencies),
+       buffer_sequencing_events, notifier = std::move(notifier), local_device,
+       stream]() mutable {
+        auto fail_recv = [&](const absl::Status& status) {
+          VLOG(2) << "CrossHostReceiveBuffersInto failed: " << status;
+          for (const auto& buffer_sequencing_event : buffer_sequencing_events) {
+            SetEventAsError(buffer_sequencing_event, status);
           }
-        });
-  };
-  async_work_runner()->Execute(std::move(recv));
+          notifier(status);
+        };
+
+        if (absl::Status status = GetErrors(transfer_dependencies);
+            !status.ok()) {
+          fail_recv(status);
+          return;
+        }
+
+        // The cross-host copy is executed on the sender's device-to-device
+        // stream once `notifier` delivers the exported fabric handles, so wait
+        // on the host for the receiver's compute stream to pass each buffer's
+        // allocation event before exporting descriptors and notifying the
+        // sender.
+        absl::InlinedVector<tsl::RCReference<tsl::AsyncValue>, 4>
+            allocation_events;
+        for (const PjRtRawBufferRef& raw_buffer : raw_buffers) {
+          auto* se_raw_buffer =
+              raw_buffer->down_cast<PjRtStreamExecutorRawBuffer>();
+          absl::StatusOr<BufferSequencingEventRef> allocation_event =
+              se_raw_buffer->device_buffer()->GetDefinitionEvent(
+                  async_work_runner(), /*nullptr_if_past=*/true);
+          if (!allocation_event.ok()) {
+            fail_recv(allocation_event.status());
+            return;
+          }
+          if (*allocation_event && !allocation_event->IsAvailable()) {
+            allocation_events.push_back(allocation_event->CopyRCRef());
+          }
+        }
+
+        auto recv = [this, raw_buffers = std::move(raw_buffers),
+                     buffer_sequencing_events =
+                         std::move(buffer_sequencing_events),
+                     notifier = std::move(notifier), local_device,
+                     stream]() mutable {
+          const uint64_t transfer_id = tsl::random::New64();
+          const std::string transfer_key =
+              absl::StrCat("pjrt_gpu_remote_recv_", transfer_id);
+
+          auto results = [&]() -> absl::StatusOr<std::vector<Future<>>> {
+            auto* executor = absl::down_cast<se::gpu::GpuExecutor*>(
+                local_device->executor());
+
+            StreamExecutorGpuCrossHostRecvDescriptor desc;
+            std::vector<PjRtCrossHostRecvDescriptors> serialized_descriptors;
+            serialized_descriptors.reserve(raw_buffers.size());
+            std::vector<std::string> rendezvous_keys;
+            rendezvous_keys.reserve(raw_buffers.size());
+
+            for (int i = 0; i < raw_buffers.size(); ++i) {
+              auto* se_raw_buffer =
+                  raw_buffers[i]->down_cast<PjRtStreamExecutorRawBuffer>();
+
+              tsl::AsyncValueRef<RawSEDeviceMemory> mem =
+                  se_raw_buffer->device_buffer();
+
+              // Keep mem alive until the Recv has finished executing.
+              buffer_sequencing_events[i].AndThen([mem]() {});
+
+              std::string rendezvous_key = absl::StrCat(transfer_key, "/", i);
+              rendezvous_keys.push_back(rendezvous_key);
+
+              // Export the buffer fabric handle and use it as a descriptor to
+              // be sent to the sender.
+              desc.Clear();
+              desc.set_buffer_index(i);
+              if (mem->mem().size() > 0) {
+                ABSL_ASSIGN_OR_RETURN(
+                    se::DeviceAddressBase range,
+                    executor->GetAllocationRange(mem->mem().opaque()));
+                ABSL_ASSIGN_OR_RETURN(
+                    *desc.mutable_buffer_handle(),
+                    GetOrExportFabricHandle(executor, range.opaque()));
+                desc.set_buffer_offset(
+                    reinterpret_cast<intptr_t>(mem->mem().opaque()) -
+                    reinterpret_cast<intptr_t>(range.opaque()));
+              }
+              desc.set_rendezvous_key(rendezvous_key);
+
+              serialized_descriptors.push_back(
+                  PjRtCrossHostRecvDescriptors{{desc.SerializeAsString()}});
+            }
+
+            auto state = std::make_shared<CrossHostRecvState>(
+                std::move(rendezvous_keys), kv_store_);
+
+            // Notify the receiver of the descriptors and cancellation callback.
+            // The caller is responsible for sending the descriptors to the
+            // sender and/or cancelling the transfer if needed.
+            VLOG(3) << "Notifying receiver of descriptors for cross-host recv "
+                       "of "
+                    << serialized_descriptors.size() << " buffers";
+            notifier(PjRtCrossHostRecvState{
+                /*descriptors=*/std::move(serialized_descriptors),
+                /*cancel_notifier=*/
+                absl::bind_front(&CrossHostRecvState::NotifyCancellation,
+                                 state),
+            });
+
+            return state->futures();
+          }();
+          if (!results.ok()) {
+            VLOG(2) << "CrossHostReceiveBuffersInto failed: "
+                    << results.status();
+            for (const auto& buffer_sequencing_event :
+                 buffer_sequencing_events) {
+              SetEventAsError(buffer_sequencing_event, results.status());
+            }
+            notifier(results.status());
+            return;
+          }
+
+          CHECK_EQ(results->size(), buffer_sequencing_events.size());
+          for (int i = 0; i < buffer_sequencing_events.size(); ++i) {
+            (*results)[i].OnReady([this, local_device, stream,
+                                   event = buffer_sequencing_events[i]](
+                                      const absl::Status& status) {
+              if (status.ok()) {
+                absl::Status s = AllocateAndRecordEvent(
+                    event, local_device, stream, "CrossHostReceiveBuffersInto");
+                if (!s.ok()) {
+                  SetEventAsError(event, s);
+                }
+              } else {
+                SetEventAsError(event, status);
+              }
+            });
+          }
+
+          JoinFutures(*results).OnReady(
+              *async_work_runner(), [kv_store = kv_store_, transfer_key](
+                                        const absl::Status& /*status*/) {
+                if (absl::Status s = kv_store->Delete(transfer_key); !s.ok()) {
+                  LOG(DFATAL) << "Failed to delete rendezvous keys for "
+                                 "transfer "
+                              << transfer_key << ": " << s;
+                }
+              });
+        };
+
+        if (allocation_events.empty()) {
+          recv();
+        } else {
+          async_work_runner()->ExecuteWhenReady(allocation_events,
+                                                std::move(recv));
+        }
+      });
 
   PjRtDeviceEventRefVector definition_events;
   for (const auto& buffer_sequencing_event : buffer_sequencing_events) {
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_cross_host_transfer_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_cross_host_transfer_test.cc
index 5659b44704b..31e35f4bdd4 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_cross_host_transfer_test.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_cross_host_transfer_test.cc
@@ -13,6 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <atomic>
+#include <cstddef>
 #include <memory>
 #include <string>
 #include <tuple>
@@ -21,18 +23,23 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/base/casts.h"
 #include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/status/status_matchers.h"
 #include "absl/status/statusor.h"
 #include "absl/synchronization/notification.h"
+#include "absl/time/time.h"
 #include "xla/future.h"
 #include "xla/literal.h"
 #include "xla/literal_util.h"
+#include "xla/pjrt/common_pjrt_client.h"
 #include "xla/pjrt/gpu/se_gpu_pjrt_client.h"
 #include "xla/pjrt/gpu/se_gpu_pjrt_client_test_helper.h"
 #include "xla/pjrt/pjrt_client.h"
 #include "xla/pjrt/plugin/xla_gpu/xla_gpu_client_options.h"
+#include "xla/pjrt/se/local_device_state.h"
+#include "xla/pjrt/se/pjrt_stream_executor_client.h"
 #include "xla/shape.h"
 #include "xla/shape_util.h"
 #include "xla/tests/literal_test_util.h"
@@ -546,5 +553,136 @@ TEST(StreamExecutorGpuClientTest, CrossHostTransferRelayEcho) {
   EXPECT_TRUE(LiteralTestUtil::Equal(input_literal, *result_literal));
 }
 
+TEST(StreamExecutorGpuClientTest,
+     CrossHostReceiveWaitsForComputeStreamAllocationSyncPoint) {
+  GpuClientOptions options = GetTestGpuClientOptions(2);
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtClient> pjrt_client,
+                       GetStreamExecutorGpuClient(options));
+  ASSERT_GE(pjrt_client->devices().size(), 2);
+
+  PjRtDevice* const src_device = pjrt_client->addressable_devices()[0];
+  PjRtDevice* const dst_device = pjrt_client->addressable_devices()[1];
+
+  auto* common_client = absl::down_cast<CommonPjRtClient*>(pjrt_client.get());
+  auto* raw_client = absl::down_cast<PjRtStreamExecutorRawClient*>(
+      common_client->raw_client());
+  ASSERT_OK_AND_ASSIGN(
+      LocalDeviceState* const dst_local_device_state,
+      raw_client->GetLocalDeviceState(dst_device->local_device_id()));
+  ASSERT_EQ(dst_local_device_state->allocation_model(),
+            LocalDeviceState::kComputeSynchronized);
+
+  Literal input_literal =
+      LiteralUtil::CreateR1<float>({1.0f, 2.0f, 3.0f, 4.0f});
+  ASSERT_OK_AND_ASSIGN(PjRtMemorySpace* const src_memory_space,
+                       src_device->default_memory_space());
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<PjRtBuffer> src_buffer,
+      pjrt_client->BufferFromHostLiteral(input_literal, src_memory_space));
+
+  // Block dst_device's compute stream to simulate an in-flight kernel whose
+  // freed buffer is reused by MakeCrossHostReceiveBuffers.
+  absl::Notification compute_callback_started;
+  absl::Notification unblock_compute_stream;
+  std::atomic<bool> compute_stream_finished = false;
+  ASSERT_OK(dst_local_device_state->compute_stream()->DoHostCallback([&]() {
+    compute_callback_started.Notify();
+    unblock_compute_stream.WaitForNotification();
+    compute_stream_finished.store(true);
+  }));
+  compute_callback_started.WaitForNotification();
+
+  const size_t sync_point =
+      dst_local_device_state->GetNextComputeStreamSyncPoint();
+  auto [desc_promise, desc_future] = MakePromise<std::string>();
+  auto desc_promise_ptr =
+      std::make_shared<Promise<std::string>>(std::move(desc_promise));
+  absl::Notification notifier_called;
+  std::atomic<bool> compute_finished_before_notifier = false;
+
+  ASSERT_OK_AND_ASSIGN(
+      std::vector<std::unique_ptr<PjRtBuffer>> recv_buffers,
+      pjrt_client->MakeCrossHostReceiveBuffers(
+          {input_literal.shape()}, dst_device,
+          [&, desc_promise_ptr](
+              absl::StatusOr<PjRtCrossHostRecvState> descriptors) {
+            CHECK_OK(descriptors);
+            CHECK_EQ(descriptors->descriptors.size(), 1);
+            compute_finished_before_notifier.store(
+                compute_stream_finished.load());
+            desc_promise_ptr->Set(
+                descriptors->descriptors[0].serialized_descriptors[0]);
+            notifier_called.Notify();
+          }));
+  ASSERT_EQ(recv_buffers.size(), 1);
+
+  // The allocation sync point event is recorded synchronously when
+  // MakeCrossHostReceiveBuffers is scheduled, and the notifier must not fire
+  // while prior work on dst_device's compute stream is still in flight.
+  EXPECT_EQ(dst_local_device_state->GetNextComputeStreamSyncPoint(),
+            sync_point + 1);
+  EXPECT_FALSE(
+      notifier_called.WaitForNotificationWithTimeout(absl::Milliseconds(50)));
+
+  src_buffer->CopyToRemoteDevice(desc_future,
+                                 [](absl::Status s, bool dispatched) {
+                                   ASSERT_TRUE(dispatched);
+                                   ASSERT_OK(s);
+                                 });
+
+  unblock_compute_stream.Notify();
+  notifier_called.WaitForNotification();
+  EXPECT_TRUE(compute_finished_before_notifier.load());
+
+  ASSERT_OK_AND_ASSIGN(std::shared_ptr<Literal> recv_literal,
+                       recv_buffers[0]->ToLiteral().Await());
+  EXPECT_TRUE(LiteralTestUtil::Equal(input_literal, *recv_literal));
+}
+
+TEST(StreamExecutorGpuClientTest, CrossHostTransferWithoutPreallocate) {
+  GpuClientOptions options = GetTestGpuClientOptions(2);
+  options.allocator_config.preallocate = false;
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtClient> client,
+                       GetStreamExecutorGpuClient(options));
+  ASSERT_GE(client->devices().size(), 2);
+
+  PjRtDevice* const src_device = client->addressable_devices()[0];
+  PjRtDevice* const dst_device = client->addressable_devices()[1];
+
+  Literal input_literal =
+      LiteralUtil::CreateR1<float>({1.0f, 2.0f, 3.0f, 4.0f});
+  ASSERT_OK_AND_ASSIGN(PjRtMemorySpace* const src_memory_space,
+                       src_device->default_memory_space());
+  ASSERT_OK_AND_ASSIGN(
+      std::unique_ptr<PjRtBuffer> src_buffer,
+      client->BufferFromHostLiteral(input_literal, src_memory_space));
+
+  auto [desc_promise, desc_future] = MakePromise<std::string>();
+  auto desc_promise_ptr =
+      std::make_shared<Promise<std::string>>(std::move(desc_promise));
+  ASSERT_OK_AND_ASSIGN(
+      std::vector<std::unique_ptr<PjRtBuffer>> recv_buffers,
+      client->MakeCrossHostReceiveBuffers(
+          {input_literal.shape()}, dst_device,
+          [desc_promise_ptr](
+              absl::StatusOr<PjRtCrossHostRecvState> descriptors) {
+            CHECK_OK(descriptors);
+            CHECK_EQ(descriptors->descriptors.size(), 1);
+            desc_promise_ptr->Set(
+                descriptors->descriptors[0].serialized_descriptors[0]);
+          }));
+  ASSERT_EQ(recv_buffers.size(), 1);
+
+  src_buffer->CopyToRemoteDevice(desc_future,
+                                 [](absl::Status s, bool dispatched) {
+                                   ASSERT_TRUE(dispatched);
+                                   ASSERT_OK(s);
+                                 });
+
+  ASSERT_OK_AND_ASSIGN(std::shared_ptr<Literal> recv_literal,
+                       recv_buffers[0]->ToLiteral().Await());
+  EXPECT_TRUE(LiteralTestUtil::Equal(input_literal, *recv_literal));
+}
+
 }  // namespace
 }  // namespace xla

From 0bb599ee9143fea0ee0d4469d07ca62eaf9a5375 Mon Sep 17 00:00:00 2001
From: Vrishab Madduri <vrishab@google.com>
Date: Fri, 16 Oct 2026 20:26:58 -0700
Subject: [PATCH 0995/1000] Copy host<->device transfers directly under GPU
 Confidential Computing.

Under GPU Confidential Computing, memory allocated with cuMemHostAlloc is driver-managed rather than pinned, and the driver already copies pageable host memory through its own shared bounce buffers. Staging transfers through PJRT's cuMemHostAlloc buffers therefore only adds a host memcpy and driver overhead.

In CC mode, ShouldStageHostToDeviceTransfers now returns false, so host-to-device and device-to-host transfers copy directly from and to the caller's host memory. Behavior with CC disabled is unchanged.

On a B200 CC disaggregated serving workload, KV cache transfer time drops 9-13% compared to chunked staging through PJRT buffers, and cold starts are stable.

PiperOrigin-RevId: 1001269978
---
 third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h
index 071209962ea..d65667edb66 100644
--- a/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h
+++ b/third_party/xla/xla/pjrt/se/pjrt_stream_executor_client.h
@@ -213,10 +213,8 @@ class PjRtStreamExecutorRawClient : public PjRtRawClient {
   }
 
   bool ShouldStageHostToDeviceTransfers(const void* data, int64_t size) const {
-    // In Confidential Computing VMs, transfers must always be staged onto
-    // host memory allocated via cuMemHostAlloc.
     if (confidential_computing_enabled_) {
-      return true;
+      return false;
     }
     // Allocating multi-gigabyte pinned buffers can be very slow. In that case,
     // using a staging buffer is probably worse than not using one.

From ec591b48038d4e5f0a2123d5216dc9439e54eb2a Mon Sep 17 00:00:00 2001
From: Bixia Zheng <bixia@google.com>
Date: Fri, 16 Oct 2026 21:50:49 -0700
Subject: [PATCH 0996/1000] Fix two problems in handling sub axes in
 rewriteReplicaGroupAttr.

Currently, we don't merge consecutive sub-axes and can produce illegal group like ["a":(1)2, "a":(2)2, "b"].

When rewriting an sub-axis, we drop its pre-size and size, and would rewrite"y":1(2) into ["a":(2)2, "b"]`.

These two cases are demonstrated in the added FileCheck tests.

To fix the problems, we first use addOrMergeNewAxisRefAttr to build the sdy AxisRefVector representing the rewrite, which accounts for sub-axes from the old attribute and merges consecutive sub-axes when building the new attribute. Then we use convertAxisRefAttr utility to convert the sdy AxisRefAttr to stablehlo attribute.

PiperOrigin-RevId: 1001296897
---
 .../service/spmd/shardy/sdy_round_trip/BUILD  |  1 +
 .../shardy/sdy_round_trip/dedup_meshes.cc     | 43 ++++++++++---------
 .../spmd/shardy/test/dedup_meshes.mlir        | 28 ++++++++++++
 3 files changed, 52 insertions(+), 20 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD
index 8b05f417940..942c1abffad 100644
--- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD
+++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/BUILD
@@ -145,6 +145,7 @@ cc_library(
         "@llvm-project//mlir:Support",
         "@shardy//shardy/dialect/sdy/ir:dialect",
         "@shardy//shardy/dialect/sdy/transforms/common:sharding_walker",
+        "@shardy//shardy/dialect/sdy/transforms/export:utils",
         "@stablehlo//:stablehlo_ops",
     ],
 )
diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/dedup_meshes.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/dedup_meshes.cc
index 3a6c6fce006..ffade0f4e43 100644
--- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/dedup_meshes.cc
+++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/dedup_meshes.cc
@@ -43,6 +43,7 @@ limitations under the License.
 #include "shardy/dialect/sdy/ir/dialect.h"
 #include "shardy/dialect/sdy/ir/utils.h"
 #include "shardy/dialect/sdy/transforms/common/sharding_walker.h"
+#include "shardy/dialect/sdy/transforms/export/utils.h"
 #include "stablehlo/dialect/StablehloOps.h"
 
 namespace xla {
@@ -401,7 +402,7 @@ void replaceManualAxes(sdy::ManualComputationOp manualComputation,
 }
 
 Attribute rewriteReplicaGroupsAttr(
-    Attribute attr, MLIRContext* context,
+    Attribute attr, MLIRContext* context, const SymbolTable& symbolTable,
     const MeshToAxisMap& duplicateMeshesToAxisMap) {
   auto rgv3 = mlir::dyn_cast<mlir::stablehlo::ReplicaGroupMeshAxesAttr>(attr);
   if (!rgv3) {
@@ -421,39 +422,41 @@ Attribute rewriteReplicaGroupsAttr(
   if (meshNameAndAxisMap == duplicateMeshesToAxisMap.end()) {
     return attr;
   }
-  auto [mainMeshName, axisMap] = meshNameAndAxisMap->getSecond();
+  const auto& [mainMeshName, axisMap] = meshNameAndAxisMap->getSecond();
+  MeshAttr mainMesh = sdy::getMeshOp(symbolTable, mainMeshName).getMesh();
 
-  SmallVector<mlir::Attribute> newAxes;
+  AxisRefVector newSdyAxes;
+  newSdyAxes.reserve(rgv3.getAxes().size());
   for (mlir::Attribute axis : rgv3.getAxes()) {
     auto axisRef = mlir::cast<mlir::stablehlo::AxisRefAttr>(axis);
-    StringRef axisName = axisRef.getName();
-    auto it = axisMap.find(axisName);
-    if (it == axisMap.end()) {
-      continue;
-    }
-    for (AxisRefAttr mainAxisRef : it->second) {
-      mlir::stablehlo::SubAxisInfoAttr subAxisInfo = nullptr;
-      if (auto sdySubAxisInfo = mainAxisRef.getSubAxisInfo()) {
-        subAxisInfo = mlir::stablehlo::SubAxisInfoAttr::get(
-            context, sdySubAxisInfo.getPreSize(), sdySubAxisInfo.getSize());
-      }
-      newAxes.push_back(mlir::stablehlo::AxisRefAttr::get(
-          context, mainAxisRef.getName(), subAxisInfo));
+    SubAxisInfoAttr sdySubAxisInfo;
+    if (auto subAxisInfo = axisRef.getSubAxisInfo()) {
+      sdySubAxisInfo = SubAxisInfoAttr::get(context, subAxisInfo.getPreSize(),
+                                            subAxisInfo.getSize());
     }
+    auto oldSdyAxisRef =
+        AxisRefAttr::get(context, axisRef.getName(), sdySubAxisInfo);
+    addOrMergeNewAxisRefAttr(oldSdyAxisRef, newSdyAxes, context, mainMesh,
+                             axisMap);
   }
 
+  SmallVector<mlir::Attribute> newAxes;
+  newAxes.reserve(newSdyAxes.size());
+  llvm::transform(newSdyAxes, std::back_inserter(newAxes),
+                  sdy::convertAxisRefAttr);
+
   return mlir::stablehlo::ReplicaGroupMeshAxesAttr::get(
       context, mlir::FlatSymbolRefAttr::get(context, mainMeshName),
       mlir::ArrayAttr::get(context, newAxes));
 }
 
-void replaceReplicaGroups(ModuleOp moduleOp,
+void replaceReplicaGroups(ModuleOp moduleOp, const SymbolTable& symbolTable,
                           const MeshToAxisMap& duplicateMeshesToAxisMap) {
   MLIRContext* context = moduleOp.getContext();
   moduleOp.walk([&](Operation* op) {
     if (auto attr = op->getAttr("replica_groups")) {
-      Attribute newAttr =
-          rewriteReplicaGroupsAttr(attr, context, duplicateMeshesToAxisMap);
+      Attribute newAttr = rewriteReplicaGroupsAttr(attr, context, symbolTable,
+                                                   duplicateMeshesToAxisMap);
       if (newAttr != attr) {
         op->setAttr("replica_groups", newAttr);
       }
@@ -483,7 +486,7 @@ void dedupMeshes(ModuleOp moduleOp, const SymbolTable& symbolTable,
                             duplicateMeshesToAxisMap);
         }
       });
-  replaceReplicaGroups(moduleOp, duplicateMeshesToAxisMap);
+  replaceReplicaGroups(moduleOp, symbolTable, duplicateMeshesToAxisMap);
   for (const auto& [targetMesh, mainMeshAndMap] : duplicateMeshesToAxisMap) {
     StringRef mainMeshName = mainMeshAndMap.first;
     if (failed(SymbolTable::replaceAllSymbolUses(
diff --git a/third_party/xla/xla/service/spmd/shardy/test/dedup_meshes.mlir b/third_party/xla/xla/service/spmd/shardy/test/dedup_meshes.mlir
index d16715b4f20..0af2d20fe21 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/dedup_meshes.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/dedup_meshes.mlir
@@ -374,3 +374,31 @@ func.func @dedup_stablehlo_replica_group_mesh_axes(%arg0: tensor<32xi32>) -> ten
 
   return %0 : tensor<32xi32>
 }
+
+// -----
+
+// CHECK: sdy.mesh @mesh = <["a"=4, "b"=2]>
+// CHECK-NOT: sdy.mesh @mesh_0 = <["x"=2, "y"=4]>
+sdy.mesh @mesh = <["a"=4, "b"=2]>
+sdy.mesh @mesh_0 = <["x"=2, "y"=4]>
+
+// CHECK-LABEL: func @dedup_stablehlo_replica_group_mesh_axes_merge_sub_axes
+func.func @dedup_stablehlo_replica_group_mesh_axes_merge_sub_axes(%arg0: tensor<32xi32>) -> tensor<32xi32> {
+  // CHECK-NEXT: "stablehlo.all_reduce"
+  // CHECK-SAME: replica_groups = #stablehlo.replica_group_mesh_axes<mesh = @mesh, axes = [#stablehlo.axis_ref<name = "a">, #stablehlo.axis_ref<name = "b">]>
+  %0 = "stablehlo.all_reduce"(%arg0) ({
+  ^bb0(%arg1: tensor<i32>, %arg2: tensor<i32>):
+    %1 = stablehlo.add %arg1, %arg2 : tensor<i32>
+    stablehlo.return %1 : tensor<i32>
+  }) {replica_groups = #stablehlo.replica_group_mesh_axes<mesh = @mesh_0, axes = [#stablehlo.axis_ref<name = "x">, #stablehlo.axis_ref<name = "y">]>} : (tensor<32xi32>) -> tensor<32xi32>
+
+  // CHECK: "stablehlo.all_reduce"
+  // CHECK-SAME: replica_groups = #stablehlo.replica_group_mesh_axes<mesh = @mesh, axes = [#stablehlo.axis_ref<name = "a", sub_axis_info = (2)2>]>
+  %2 = "stablehlo.all_reduce"(%0) ({
+  ^bb0(%arg1: tensor<i32>, %arg2: tensor<i32>):
+    %3 = stablehlo.add %arg1, %arg2 : tensor<i32>
+    stablehlo.return %3 : tensor<i32>
+  }) {replica_groups = #stablehlo.replica_group_mesh_axes<mesh = @mesh_0, axes = [#stablehlo.axis_ref<name = "y", sub_axis_info = (1)2>]>} : (tensor<32xi32>) -> tensor<32xi32>
+
+  return %2 : tensor<32xi32>
+}

From f601a92cdb5b7a38d85d7733b655f1215ae1d8bc Mon Sep 17 00:00:00 2001
From: Bixia Zheng <bixia@google.com>
Date: Fri, 16 Oct 2026 23:01:37 -0700
Subject: [PATCH 0997/1000] Clean up dedup_meshes.cc and preserve reductionOp
 in replaceAxesInSharding. 1. Bind [mainMeshName, axisMap] by const reference
 in replaceAxesInSharding and    replaceManualAxes to avoid copying AxisMap by
 value. 2. Reuse the buildNewAxisRefList lambda for dimension shardings in   
 replaceAxesInSharding instead of duplicating the loop. 3. Extract
 addOrMergeFullAxis to share full-axis remapping between   
 addOrMergeNewAxisRefAttr and replaceManualAxes without copying AxisRefVector.
 4. Preserve oldSharding.getReductionOp() when rebuilding TensorShardingAttr
 in    replaceAxesInSharding. 5. Early-return in replaceManualAxes when
 manual_axes is empty, use    llvm::hasSingleElement in runOnOperation, and
 remove unused includes/usings    and redundant mlir:: qualifiers.

PiperOrigin-RevId: 1001322103
---
 .../shardy/sdy_round_trip/dedup_meshes.cc     | 94 +++++++++----------
 .../spmd/shardy/test/dedup_meshes.mlir        |  4 +-
 2 files changed, 44 insertions(+), 54 deletions(-)

diff --git a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/dedup_meshes.cc b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/dedup_meshes.cc
index ffade0f4e43..e7f5150bf16 100644
--- a/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/dedup_meshes.cc
+++ b/third_party/xla/xla/service/spmd/shardy/sdy_round_trip/dedup_meshes.cc
@@ -33,7 +33,6 @@ limitations under the License.
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
-#include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/DialectRegistry.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/Pass/Pass.h"
@@ -55,8 +54,6 @@ using ::llvm::DenseMapInfo;
 using ::llvm::SmallDenseMap;
 using ::mlir::ArrayRef;
 using ::mlir::Attribute;
-using ::mlir::DenseMap;
-using ::mlir::DenseSet;
 using ::mlir::MLIRContext;
 using ::mlir::ModuleOp;
 using ::mlir::Operation;
@@ -120,7 +117,7 @@ struct MeshDeviceIdentifierInfo : public DenseMapInfo<MeshDeviceIdentifier> {
 // successful, `targetToMainAxisMap` will be populated.
 bool mapTargetAxesToMainAxes(MeshOp targetMesh, MeshOp mainMesh,
                              AxisMap& targetToMainAxisMap,
-                             mlir::MLIRContext* context) {
+                             MLIRContext* context) {
   ArrayRef<MeshAxisAttr> mainAxes = mainMesh.getMesh().getAxes();
   ArrayRef<MeshAxisAttr> targetAxes = targetMesh.getMesh().getAxes();
   struct AxisMappingState {
@@ -192,25 +189,31 @@ using MeshIdentifierToMainMeshesMap =
 // map of old axes to the (sub)axes in the main mesh.
 using MeshToAxisMap = SmallDenseMap<StringRef, std::pair<StringRef, AxisMap>>;
 
+void addOrMergeFullAxis(StringRef oldAxisName, AxisRefVector& newAxisRefs,
+                        const MeshAttr& mainMesh, const AxisMap& axisMap) {
+  if (auto it = axisMap.find(oldAxisName); it != axisMap.end()) {
+    for (AxisRefAttr mainAxisRef : it->second) {
+      sdy::addAxisOrMerge(newAxisRefs, mainAxisRef, mainMesh);
+    }
+  }
+}
+
 // Takes an old axis ref, find the mapped new (sub-)axes, add them to the new
 // axis ref list and merge consecutive sub-axes.
 void addOrMergeNewAxisRefAttr(AxisRefAttr oldAxisRef,
-                              AxisRefVector& newAxisRefs,
-                              mlir::MLIRContext* context,
+                              AxisRefVector& newAxisRefs, MLIRContext* context,
                               const MeshAttr& mainMesh,
                               const AxisMap& axisMap) {
+  if (!oldAxisRef.getSubAxisInfo()) {
+    addOrMergeFullAxis(oldAxisRef.getName(), newAxisRefs, mainMesh, axisMap);
+    return;
+  }
   auto it = axisMap.find(oldAxisRef.getName());
   if (it == axisMap.end()) {
     // The old axis is of size 1, skip it.
     return;
   }
   ArrayRef<AxisRefAttr> mainAxisRefs = it->second;
-  if (!oldAxisRef.getSubAxisInfo()) {
-    for (AxisRefAttr mainAxisRef : mainAxisRefs) {
-      sdy::addAxisOrMerge(newAxisRefs, mainAxisRef, mainMesh);
-    }
-    return;
-  }
   // TODO(zenong): We can iterate the list once for sub-axes of a target axis.
   int64_t remainingOldPreSize = oldAxisRef.getSubAxisInfo().getPreSize();
   int64_t remainingOldSize = oldAxisRef.getSubAxisInfo().getSize();
@@ -306,7 +309,7 @@ MeshToAxisMap buildDuplicateMeshesToAxisMap(ModuleOp moduleOp) {
 TensorShardingAttr replaceAxesInSharding(
     TensorShardingAttr oldSharding, const SymbolTable& symbolTable,
     const MeshToAxisMap& duplicateMeshesToAxisMap) {
-  mlir::MLIRContext* context = oldSharding.getContext();
+  MLIRContext* context = oldSharding.getContext();
   if (mlir::isa<MeshAttr>(oldSharding.getMeshOrRef())) {
     // Skip shardings with inlined meshes.
     return oldSharding;
@@ -318,21 +321,8 @@ TensorShardingAttr replaceAxesInSharding(
   if (meshNameAndAxisMap == duplicateMeshesToAxisMap.end()) {
     return oldSharding;
   }
-  auto [mainMeshName, axisMap] = meshNameAndAxisMap->getSecond();
+  const auto& [mainMeshName, axisMap] = meshNameAndAxisMap->getSecond();
   MeshAttr mainMesh = sdy::getMeshOp(symbolTable, mainMeshName).getMesh();
-  SmallVector<DimensionShardingAttr> newDimShardings;
-  newDimShardings.reserve(oldSharding.getDimShardings().size());
-  for (DimensionShardingAttr oldDimSharding : oldSharding.getDimShardings()) {
-    AxisRefVector newAxisRefs;
-    newAxisRefs.reserve(oldDimSharding.getAxes().size());
-    for (AxisRefAttr oldAxisRef : oldDimSharding.getAxes()) {
-      addOrMergeNewAxisRefAttr(oldAxisRef, newAxisRefs, context, mainMesh,
-                               axisMap);
-    }
-    newDimShardings.push_back(DimensionShardingAttr::get(
-        context, newAxisRefs, oldDimSharding.getIsClosed(),
-        oldDimSharding.getPriority()));
-  }
   auto buildNewAxisRefList =
       [&, &axisMap = axisMap](ArrayRef<AxisRefAttr> oldAxisRefs) {
         AxisRefVector newAxisRefs;
@@ -343,12 +333,22 @@ TensorShardingAttr replaceAxesInSharding(
         }
         return newAxisRefs;
       };
+  SmallVector<DimensionShardingAttr> newDimShardings;
+  newDimShardings.reserve(oldSharding.getDimShardings().size());
+  for (DimensionShardingAttr oldDimSharding : oldSharding.getDimShardings()) {
+    newDimShardings.push_back(DimensionShardingAttr::get(
+        context, buildNewAxisRefList(oldDimSharding.getAxes()),
+        oldDimSharding.getIsClosed(), oldDimSharding.getPriority()));
+  }
   AxisRefVector newReplicatedAxes =
       buildNewAxisRefList(oldSharding.getReplicatedAxes());
   AxisRefVector newUnreducedAxes =
       buildNewAxisRefList(oldSharding.getUnreducedAxes());
   return TensorShardingAttr::get(context, mainMeshName, newDimShardings,
-                                 newReplicatedAxes, newUnreducedAxes);
+                                 newReplicatedAxes, newUnreducedAxes,
+                                 newUnreducedAxes.empty()
+                                     ? sdy::ReductionOp::SUM
+                                     : oldSharding.getReductionOp());
 }
 
 // Replaces the manual axes in `manualComputation`, if it refers to some mesh
@@ -357,11 +357,12 @@ TensorShardingAttr replaceAxesInSharding(
 void replaceManualAxes(sdy::ManualComputationOp manualComputation,
                        const SymbolTable& symbolTable,
                        const MeshToAxisMap& duplicateMeshesToAxisMap) {
-  if (manualComputation.getInShardings().empty() &&
-      manualComputation.getOutShardings().empty()) {
+  if (manualComputation.getManualAxes().empty() ||
+      (manualComputation.getInShardings().empty() &&
+       manualComputation.getOutShardings().empty())) {
     return;
   }
-  mlir::Attribute meshOrRef = getCommonMeshOrRef(
+  Attribute meshOrRef = getCommonMeshOrRef(
       manualComputation.getInShardings().getShardings(),
       manualComputation.getOutShardings().getShardings(), symbolTable);
   CHECK(meshOrRef) << "no common mesh found for ManualComputationOp";
@@ -375,19 +376,13 @@ void replaceManualAxes(sdy::ManualComputationOp manualComputation,
   if (meshNameAndAxisMap == duplicateMeshesToAxisMap.end()) {
     return;
   }
-  auto [mainMeshName, axisMap] = meshNameAndAxisMap->getSecond();
+  const auto& [mainMeshName, axisMap] = meshNameAndAxisMap->getSecond();
   MeshAttr mainMesh = sdy::getMeshOp(symbolTable, mainMeshName).getMesh();
   AxisRefVector newAxisRefs;
   newAxisRefs.reserve(manualComputation.getManualAxes().size());
-  llvm::transform(manualComputation.getManualAxes(),
-                  sdy::AddAxisOrMergeInserter(&newAxisRefs, &mainMesh),
-                  [&axisMap = axisMap](StringAttr manualAxis) -> AxisRefVector {
-                    auto it = axisMap.find(manualAxis.getValue());
-                    if (it == axisMap.end()) {
-                      return {};
-                    }
-                    return it->second;
-                  });
+  for (StringAttr manualAxis : manualComputation.getManualAxes()) {
+    addOrMergeFullAxis(manualAxis.getValue(), newAxisRefs, mainMesh, axisMap);
+  }
   SmallVector<StringAttr> newManualAxes;
   newManualAxes.reserve(newAxisRefs.size());
   llvm::transform(
@@ -408,7 +403,7 @@ Attribute rewriteReplicaGroupsAttr(
   if (!rgv3) {
     return attr;
   }
-  mlir::Attribute meshAttr = rgv3.getMesh();
+  Attribute meshAttr = rgv3.getMesh();
   StringRef oldMeshName;
   if (auto symbolRef = mlir::dyn_cast<mlir::FlatSymbolRefAttr>(meshAttr)) {
     oldMeshName = symbolRef.getValue();
@@ -427,7 +422,7 @@ Attribute rewriteReplicaGroupsAttr(
 
   AxisRefVector newSdyAxes;
   newSdyAxes.reserve(rgv3.getAxes().size());
-  for (mlir::Attribute axis : rgv3.getAxes()) {
+  for (Attribute axis : rgv3.getAxes()) {
     auto axisRef = mlir::cast<mlir::stablehlo::AxisRefAttr>(axis);
     SubAxisInfoAttr sdySubAxisInfo;
     if (auto subAxisInfo = axisRef.getSubAxisInfo()) {
@@ -440,7 +435,7 @@ Attribute rewriteReplicaGroupsAttr(
                              axisMap);
   }
 
-  SmallVector<mlir::Attribute> newAxes;
+  SmallVector<Attribute> newAxes;
   newAxes.reserve(newSdyAxes.size());
   llvm::transform(newSdyAxes, std::back_inserter(newAxes),
                   sdy::convertAxisRefAttr);
@@ -479,7 +474,7 @@ void dedupMeshes(ModuleOp moduleOp, const SymbolTable& symbolTable,
         return replaceAxesInSharding(oldSharding, symbolTable,
                                      duplicateMeshesToAxisMap);
       },
-      [&](mlir::Operation* op) {
+      [&](Operation* op) {
         if (auto manualComputation =
                 mlir::dyn_cast<sdy::ManualComputationOp>(op)) {
           replaceManualAxes(manualComputation, symbolTable,
@@ -514,13 +509,8 @@ class SdyRoundTripDedupMeshesPass
     ModuleOp moduleOp = getOperation();
 
     // Exit early if there are no meshes or just one mesh.
-    auto meshIter = moduleOp.getOps<MeshOp>();
-    if (meshIter.empty()) {
-      return;
-    }
-    auto meshIterBegin = meshIter.begin();
-    std::advance(meshIterBegin, 1);
-    if (meshIterBegin == meshIter.end()) {
+    auto meshOps = moduleOp.getOps<MeshOp>();
+    if (meshOps.empty() || llvm::hasSingleElement(meshOps)) {
       return;
     }
 
diff --git a/third_party/xla/xla/service/spmd/shardy/test/dedup_meshes.mlir b/third_party/xla/xla/service/spmd/shardy/test/dedup_meshes.mlir
index 0af2d20fe21..28fe85289d6 100644
--- a/third_party/xla/xla/service/spmd/shardy/test/dedup_meshes.mlir
+++ b/third_party/xla/xla/service/spmd/shardy/test/dedup_meshes.mlir
@@ -97,10 +97,10 @@ sdy.mesh @mesh_mapping_consecutive_sub_axes = <["x"=4, "y"=8, "z"=7]>
 
 // CHECK-LABEL: @full_axes
 // CHECK-SAME:  %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh1, [{"a", ?}, {?}]>}
-// CHECK-SAME:  -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh1, [{"b", ?}, {?}], unreduced={"a"}>}) {
+// CHECK-SAME:  -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh1, [{"b", ?}, {?}], unreduced=max{"a"}>}) {
 func.func @full_axes(
   %arg0: tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh1, [{"a", ?}, {?}]>}
-) -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh2, [{"model", ?}, {?}], unreduced={"data"}>}) {
+) -> (tensor<8x8xf32> {sdy.sharding = #sdy.sharding<@mesh2, [{"model", ?}, {?}], unreduced=max{"data"}>}) {
   // CHECK-NEXT: stablehlo.add
   // CHECK-SAME{LITERAL}: #sdy.sharding_per_value<[<@mesh1, [{"a", ?}p1, {}], replicated={"b"}>]>
   %0 = stablehlo.add %arg0, %arg0 {sdy.sharding = #sdy.sharding_per_value<[<@mesh3, [{"x", ?}p1, {}], replicated={"y"}>]>} : tensor<8x8xf32>

From b9b04b174ab779066680ea3fe65f056e5f88854e Mon Sep 17 00:00:00 2001
From: Ce Zheng CA <zce@google.com>
Date: Sat, 17 Oct 2026 00:31:51 -0700
Subject: [PATCH 0998/1000] [XLA:TSL] Deduplicate Eigen spatial convolution
 contraction mappers across alignment and contiguity parameters

In `xla/tsl/framework/convolution/eigen_spatial_convolutions-inl.h`, the partial specializations of `TensorContractionInputMapper` and `TensorContractionSubMapper` for `TensorReshapingOp<..., TensorImagePatchOp<...>>` ignore `inner_dim_contiguous`, `inner_dim_reordered`, and `Alignment` in their method bodies (all packet loads use `Unaligned`). However, `gemm_pack_rhs` and `gemm_pack_colmajor_block` are instantiated once per distinct `SubMapper` type across 4 combinations of `(C = inner_dim_contiguous, R = inner_dim_reordered, A = Alignment)`: `(false, false, 0)`, `(false, false, Aligned)`, `(true, false, 0)`, and `(true, false, Aligned)`.

Before:
- `TensorContractionInputMapper<..., C, R, A>` (4 independent specializations)
  `-> SubMapper = TensorContractionSubMapper<..., C, R, A>` (4 duplicate `gemm_pack_rhs` instantiations)

After:
- `TensorContractionInputMapper<..., C, R, A> : TensorContractionInputMapper<..., false, false, 0>`
  `-> SubMapper = TensorContractionSubMapper<..., false, false, 0>` (1 shared `gemm_pack_rhs` instantiation)

Zero runtime or memory overhead (identical machine code, 3 fewer copies of the RHS packing kernels per convolution instantiation).

PiperOrigin-RevId: 1001353631
---
 .../eigen_spatial_convolutions-inl.h          | 74 ++++++++++++++-----
 .../eigen_spatial_convolutions_test.cc        |  8 +-
 2 files changed, 57 insertions(+), 25 deletions(-)

diff --git a/third_party/xla/xla/tsl/framework/convolution/eigen_spatial_convolutions-inl.h b/third_party/xla/xla/tsl/framework/convolution/eigen_spatial_convolutions-inl.h
index 82d9baa425b..92c9bce7ad6 100644
--- a/third_party/xla/xla/tsl/framework/convolution/eigen_spatial_convolutions-inl.h
+++ b/third_party/xla/xla/tsl/framework/convolution/eigen_spatial_convolutions-inl.h
@@ -55,18 +55,24 @@ namespace internal {
 // TODO(ezhulenev): Consolidate this part of the code with the image patch
 // extraction code since they are both very similar.
 
+// Image patch mappers do not depend on `inner_dim_contiguous`,
+// `inner_dim_reordered`, or `Alignment` (all packet loads use `Unaligned`).
+// Canonicalize the implementation and `SubMapper` typedef to `<..., false,
+// false, 0>` so `TensorContractionEvaluatorBase::evalProduct` does not
+// instantiate duplicate copies of `gemm_pack_rhs` and
+// `gemm_pack_colmajor_block` for each `(inner_dim_contiguous,
+// inner_dim_reordered, Alignment)` combination.
 template <typename NewDimension, Index Rows, Index Cols, typename ArgType,
           typename Device, typename Scalar_, typename Index,
-          typename nocontract_t, typename contract_t, int Side, int packet_size,
-          bool inner_dim_contiguous, bool inner_dim_reordered, int Alignment>
+          typename nocontract_t, typename contract_t, int Side, int packet_size>
 class TensorContractionInputMapper<
     Scalar_, Index, Side,
     TensorEvaluator<
         const TensorReshapingOp<NewDimension,
                                 const TensorImagePatchOp<Rows, Cols, ArgType> >,
         Device>,
-    nocontract_t, contract_t, packet_size, inner_dim_contiguous,
-    inner_dim_reordered, Alignment> {
+    nocontract_t, contract_t, packet_size, /*inner_dim_contiguous=*/false,
+    /*inner_dim_reordered=*/false, /*Alignment=*/0> {
  public:
   typedef Scalar_ Scalar;
 
@@ -76,8 +82,8 @@ class TensorContractionInputMapper<
           const TensorReshapingOp<
               NewDimension, const TensorImagePatchOp<Rows, Cols, ArgType> >,
           Device>,
-      nocontract_t, contract_t, packet_size, inner_dim_contiguous,
-      inner_dim_reordered, Alignment>
+      nocontract_t, contract_t, packet_size, /*inner_dim_contiguous=*/false,
+      /*inner_dim_reordered=*/false, /*Alignment=*/0>
       Self;
 
   typedef TensorContractionSubMapper<
@@ -86,8 +92,8 @@ class TensorContractionInputMapper<
           const TensorReshapingOp<
               NewDimension, const TensorImagePatchOp<Rows, Cols, ArgType> >,
           Device>,
-      nocontract_t, contract_t, packet_size, inner_dim_contiguous,
-      inner_dim_reordered, Alignment>
+      nocontract_t, contract_t, packet_size, /*inner_dim_contiguous=*/false,
+      /*inner_dim_reordered=*/false, /*Alignment=*/0>
       SubMapper;
 
   typedef SubMapper VectorMapper;
@@ -281,8 +287,8 @@ class TensorContractionInputMapper<
           const TensorReshapingOp<
               NewDimension, const TensorImagePatchOp<Rows, Cols, ArgType> >,
           Device>,
-      nocontract_t, contract_t, packet_size, inner_dim_contiguous,
-      inner_dim_reordered, Alignment>;
+      nocontract_t, contract_t, packet_size, /*inner_dim_contiguous=*/false,
+      /*inner_dim_reordered=*/false, /*Alignment=*/0>;
 
   // Load coefficient from a patch specified by the "within patch offset"
   // (patchId) and the precomputed indices of the first element of the patch.
@@ -677,14 +683,47 @@ template <typename NewDimension, Index Rows, Index Cols, typename ArgType,
           typename Device, typename Scalar, typename Index,
           typename nocontract_t, typename contract_t, int Side, int packet_size,
           bool inner_dim_contiguous, bool inner_dim_reordered, int Alignment>
-class TensorContractionSubMapper<
+class TensorContractionInputMapper<
     Scalar, Index, Side,
     TensorEvaluator<
         const TensorReshapingOp<NewDimension,
                                 const TensorImagePatchOp<Rows, Cols, ArgType> >,
         Device>,
     nocontract_t, contract_t, packet_size, inner_dim_contiguous,
-    inner_dim_reordered, Alignment> {
+    inner_dim_reordered, Alignment>
+    : public TensorContractionInputMapper<
+          Scalar, Index, Side,
+          TensorEvaluator<
+              const TensorReshapingOp<
+                  NewDimension, const TensorImagePatchOp<Rows, Cols, ArgType> >,
+              Device>,
+          nocontract_t, contract_t, packet_size,
+          /*inner_dim_contiguous=*/false, /*inner_dim_reordered=*/false,
+          /*Alignment=*/0> {
+ public:
+  using Base = TensorContractionInputMapper<
+      Scalar, Index, Side,
+      TensorEvaluator<
+          const TensorReshapingOp<
+              NewDimension, const TensorImagePatchOp<Rows, Cols, ArgType> >,
+          Device>,
+      nocontract_t, contract_t, packet_size,
+      /*inner_dim_contiguous=*/false, /*inner_dim_reordered=*/false,
+      /*Alignment=*/0>;
+  using Base::Base;
+};
+
+template <typename NewDimension, Index Rows, Index Cols, typename ArgType,
+          typename Device, typename Scalar, typename Index,
+          typename nocontract_t, typename contract_t, int Side, int packet_size>
+class TensorContractionSubMapper<
+    Scalar, Index, Side,
+    TensorEvaluator<
+        const TensorReshapingOp<NewDimension,
+                                const TensorImagePatchOp<Rows, Cols, ArgType> >,
+        Device>,
+    nocontract_t, contract_t, packet_size, /*inner_dim_contiguous=*/false,
+    /*inner_dim_reordered=*/false, /*Alignment=*/0> {
  public:
   typedef typename packet_traits<Scalar>::type Packet;
   typedef typename packet_traits<Scalar>::half HalfPacket;
@@ -695,8 +734,8 @@ class TensorContractionSubMapper<
           const TensorReshapingOp<
               NewDimension, const TensorImagePatchOp<Rows, Cols, ArgType> >,
           Device>,
-      nocontract_t, contract_t, packet_size, inner_dim_contiguous,
-      inner_dim_reordered, Alignment>
+      nocontract_t, contract_t, packet_size, /*inner_dim_contiguous=*/false,
+      /*inner_dim_reordered=*/false, /*Alignment=*/0>
       ParentMapper;
 
   typedef TensorContractionSubMapper<
@@ -705,8 +744,8 @@ class TensorContractionSubMapper<
           const TensorReshapingOp<
               NewDimension, const TensorImagePatchOp<Rows, Cols, ArgType> >,
           Device>,
-      nocontract_t, contract_t, packet_size, inner_dim_contiguous,
-      inner_dim_reordered, Alignment>
+      nocontract_t, contract_t, packet_size, /*inner_dim_contiguous=*/false,
+      /*inner_dim_reordered=*/false, /*Alignment=*/0>
       Self;
 
   typedef Self LinearMapper;
@@ -746,8 +785,7 @@ class TensorContractionSubMapper<
   template <typename PacketT = Packet>
   EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE PacketT loadPacket(Index i,
                                                            Index j) const {
-    return m_base_mapper.template loadPacket<Alignment>(i + m_depth_offset,
-                                                        j + m_col_offset);
+    return m_base_mapper.loadPacket(i + m_depth_offset, j + m_col_offset);
   }
   EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE Scalar
   loadCoeffStandard(Index i) const {
diff --git a/third_party/xla/xla/tsl/framework/convolution/eigen_spatial_convolutions_test.cc b/third_party/xla/xla/tsl/framework/convolution/eigen_spatial_convolutions_test.cc
index ffb67d90290..01b78132544 100644
--- a/third_party/xla/xla/tsl/framework/convolution/eigen_spatial_convolutions_test.cc
+++ b/third_party/xla/xla/tsl/framework/convolution/eigen_spatial_convolutions_test.cc
@@ -924,13 +924,7 @@ static void PackRhsHelper(::testing::benchmark::State& state,
       /*inner_dim_reordered*/ false,              //
       /*Alignment*/ 0>;
 
-  using SubMapper = Eigen::internal::TensorContractionSubMapper<
-      T, Index, Eigen::internal::Rhs, Evaluator,  //
-      nocontract_t, contract_t,                   //
-      packet_size,                                //
-      /*inner_dim_contiguous*/ true,              //
-      /*inner_dim_reordered*/ false,              //
-      /*Alignment*/ 0>;
+  using SubMapper = typename InputMapper::SubMapper;
 
 #if defined(TENSORFLOW_USE_MKLDNN_CONTRACTION_KERNEL)
   using PackRhsImpl =

From 9223b75468a239ec00058eead2b9b8206ccbaebe Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Sat, 17 Oct 2026 01:06:51 -0700
Subject: [PATCH 0999/1000] [XLA] Remove deprecated
 HloPassInterface::Run(std::unique_ptr<HloModule>&).

No passes override `RunImpl(std::unique_ptr<HloModule>&)` anymore; every HLO pass updates its `HloModule` in place via `Run(HloModule*)`.

Delete `HloPassInterface::Run(std::unique_ptr<HloModule>&)` and `RunImpl(std::unique_ptr<HloModule>&)`, along with the matching `HloPassPipeline` override, and remove the `HloT` templates on `HloPassPipeline` helpers now that they only take `HloModule*`.

PiperOrigin-RevId: 1001364016
---
 third_party/xla/xla/hlo/pass/BUILD            |  1 -
 .../xla/xla/hlo/pass/hlo_pass_interface.cc    | 14 -------
 .../xla/xla/hlo/pass/hlo_pass_interface.h     | 20 ---------
 .../xla/xla/hlo/pass/hlo_pass_pipeline.cc     | 42 ++++---------------
 .../xla/xla/hlo/pass/hlo_pass_pipeline.h      | 24 ++++-------
 5 files changed, 16 insertions(+), 85 deletions(-)

diff --git a/third_party/xla/xla/hlo/pass/BUILD b/third_party/xla/xla/hlo/pass/BUILD
index acc79a485ca..8a41c45567b 100644
--- a/third_party/xla/xla/hlo/pass/BUILD
+++ b/third_party/xla/xla/hlo/pass/BUILD
@@ -40,7 +40,6 @@ cc_library(
         "//xla/tsl/platform:errors",
         "//xla/tsl/platform:logging",
         "//xla/tsl/platform:statusor",
-        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/container:flat_hash_set",
         "@com_google_absl//absl/hash",
         "@com_google_absl//absl/log",
diff --git a/third_party/xla/xla/hlo/pass/hlo_pass_interface.cc b/third_party/xla/xla/hlo/pass/hlo_pass_interface.cc
index 381799d81db..2cc24995dc6 100644
--- a/third_party/xla/xla/hlo/pass/hlo_pass_interface.cc
+++ b/third_party/xla/xla/hlo/pass/hlo_pass_interface.cc
@@ -67,18 +67,4 @@ absl::StatusOr<bool> HloPassInterface::Run(
   return RunImpl(module, execution_threads);
 }
 
-absl::StatusOr<bool> HloPassInterface::Run(
-    std::unique_ptr<HloModule>& module,
-    const absl::flat_hash_set<absl::string_view>& execution_threads) {
-  if (!ShouldRunPass(module.get(), name(), IsPassPipeline())) {
-    return false;
-  }
-
-  auto* env = tsl::Env::Default();
-  std::unique_ptr<tsl::ThreadNote> thread_note;
-  thread_note = env->AddThreadNote(absl::StrCat("Running HLO pass on module ",
-                                                module->name(), ": ", name()));
-  return RunImpl(module, execution_threads);
-}
-
 }  // namespace xla
diff --git a/third_party/xla/xla/hlo/pass/hlo_pass_interface.h b/third_party/xla/xla/hlo/pass/hlo_pass_interface.h
index 1fcdf4058c6..23a4e9439b0 100644
--- a/third_party/xla/xla/hlo/pass/hlo_pass_interface.h
+++ b/third_party/xla/xla/hlo/pass/hlo_pass_interface.h
@@ -17,10 +17,8 @@ limitations under the License.
 #define XLA_HLO_PASS_HLO_PASS_INTERFACE_H_
 
 #include <cstdint>
-#include <memory>
 #include <string>
 
-#include "absl/base/attributes.h"
 #include "absl/container/flat_hash_set.h"
 #include "absl/log/log.h"
 #include "absl/status/status.h"
@@ -83,17 +81,6 @@ class HloPassInterface {
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads = {});
 
-  // Same as above, except that this API allows the pass to return a *different*
-  // module, rather than modifying the module in-place.
-  // TODO(b/454418238): Remove this function (and the corresponding RunImpl).
-  ABSL_DEPRECATED(
-      "This interface exists for passes which create an entire new HloModule "
-      "instead of updating the existing one in place. This is not the norm and "
-      "should not be done.")
-  absl::StatusOr<bool> Run(
-      std::unique_ptr<HloModule>& module_ptr,
-      const absl::flat_hash_set<absl::string_view>& execution_threads = {});
-
   // Run the pass on computation on changed computations from last iteration in
   // given HLO module for specified execution_threads, with caller provided
   // RunState which holds the state information across multiple iterations.
@@ -132,13 +119,6 @@ class HloPassInterface {
   virtual absl::StatusOr<bool> RunImpl(
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads) = 0;
-
-  // TODO(b/454418238): Remove this function.
-  virtual absl::StatusOr<bool> RunImpl(
-      std::unique_ptr<HloModule>& module_ptr,
-      const absl::flat_hash_set<absl::string_view>& execution_threads) {
-    return RunImpl(module_ptr.get(), execution_threads);
-  }
 };
 
 // Base class for passes which are module-scoped.
diff --git a/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc b/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc
index 5cf8c571280..2cdd722e1d3 100644
--- a/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc
+++ b/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.cc
@@ -80,15 +80,14 @@ void RecordPassEndMetadata(HloModule& module, absl::string_view pass_name,
 
 }  // namespace
 
-template <typename HloT>
 absl::Status HloPassPipeline::RunInvariantCheckers(
-    HloT hlo, absl::string_view after_pass_name,
+    HloModule* hlo, absl::string_view after_pass_name,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   tsl::profiler::TraceMe traceme("RunInvariantCheckers");
   for (auto& invariant_checker : invariant_checkers_) {
     VLOG(1) << "    Invariant checker " << invariant_checker->name();
     absl::StatusOr<bool> changed_status =
-        RunHelper<HloT>(invariant_checker.get(), hlo, execution_threads);
+        RunHelper(invariant_checker.get(), hlo, execution_threads);
     VLOG(1) << "    Invariant checker done " << invariant_checker->name();
     if (!changed_status.ok()) {
       VLOG(2) << "Failed invariant check:";
@@ -109,8 +108,7 @@ std::string UniqueId(const HloModule& mod) {
   return std::to_string(mod.unique_id());
 }
 
-template <typename HloT>
-static void VerifyPassChangedReport(const HloT hlo, bool pass_changed,
+static void VerifyPassChangedReport(const HloModule* hlo, bool pass_changed,
                                     const DebugOptions& debug_options,
                                     absl::string_view pass_name,
                                     absl::string_view pipeline_name,
@@ -137,9 +135,8 @@ static void VerifyPassChangedReport(const HloT hlo, bool pass_changed,
 
 }  // namespace
 
-template <typename HloT>
 absl::StatusOr<bool> HloPassPipeline::RunPassesInternal(
-    HloT hlo, const DebugOptions& debug_options,
+    HloModule* hlo, const DebugOptions& debug_options,
     const absl::flat_hash_set<absl::string_view>& execution_threads) {
   auto* env = tsl::Env::Default();
   std::unique_ptr<tsl::ThreadNote> thread_note;
@@ -160,8 +157,7 @@ absl::StatusOr<bool> HloPassPipeline::RunPassesInternal(
                            pipeline_name, hlo->name(), UniqueId(*hlo));
   }};
 
-  ABSL_RETURN_IF_ERROR(
-      RunInvariantCheckers<HloT>(hlo, kPipelineStart, execution_threads));
+  ABSL_RETURN_IF_ERROR(RunInvariantCheckers(hlo, kPipelineStart, execution_threads));
 
   RecordPassStartMetadata(*hlo, std::string(kPipelineStart), pipeline_name);
   MaybeDumpHloAndSaveFilenames(*hlo,
@@ -245,15 +241,15 @@ absl::StatusOr<bool> HloPassPipeline::RunPassesInternal(
     if (!pass->IsPassPipeline()) {
       compilation_stats_->StartPass(pass_name);
     }
-    auto status_or_changed = RunHelper<HloT>(pass, hlo, execution_threads);
+    auto status_or_changed = RunHelper(pass, hlo, execution_threads);
     if (auto status = status_or_changed.status(); !status.ok()) {
       compilation_stats_->RecordPassError(
           pass_name, absl::StatusCodeToString(status.code()));
     }
     ABSL_ASSIGN_OR_RETURN(bool pass_changed, status_or_changed);
     if (verify_pass_changed_report) {
-      VerifyPassChangedReport<HloT>(hlo, pass_changed, debug_options, pass_name,
-                                    pipeline_name, hash_before.value());
+      VerifyPassChangedReport(hlo, pass_changed, debug_options, pass_name,
+                              pipeline_name, hash_before.value());
     }
     if (!dump_regex.empty() && (pass_changed || dump_regex != ".*")) {
       MaybeDumpHloAndSaveFilenames(*hlo,
@@ -266,8 +262,7 @@ absl::StatusOr<bool> HloPassPipeline::RunPassesInternal(
     changed |= pass_changed;
     if (pass_changed) {
       VLOG(3) << "  Pass caused changes " << pass_name;
-      auto status =
-          RunInvariantCheckers<HloT>(hlo, pass_name, execution_threads);
+      auto status = RunInvariantCheckers(hlo, pass_name, execution_threads);
       if (!status.ok()) {
         compilation_stats_->RecordPassError(
             pass_name, absl::StatusCodeToString(status.code()));
@@ -312,23 +307,4 @@ absl::StatusOr<bool> HloPassPipeline::RunImpl(
   return RunPassesInternal(module, debug_options, execution_threads);
 }
 
-absl::StatusOr<bool> HloPassPipeline::RunImpl(
-    std::unique_ptr<HloModule>& module,
-    const absl::flat_hash_set<absl::string_view>& execution_threads) {
-  run_called_ = true;
-
-  VLOG(1) << "Running HLO pass pipeline on module " << module->name() << ": "
-          << name();
-
-  tsl::profiler::TraceMe traceme([&] {
-    return tsl::profiler::TraceMeEncode(
-        absl::StrCat(name(), " (", module->name(), ")"),
-        {{"module", module->name()}});
-  });
-  // Copy debug options by value as passes may modify module config.
-  DebugOptions debug_options = module->config().debug_options();
-  return RunPassesInternal<std::unique_ptr<HloModule>&>(module, debug_options,
-                                                        execution_threads);
-}
-
 }  // namespace xla
diff --git a/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.h b/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.h
index e28d08bcad4..989055bd8ab 100644
--- a/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.h
+++ b/third_party/xla/xla/hlo/pass/hlo_pass_pipeline.h
@@ -98,9 +98,6 @@ class HloPassPipeline : public HloPassInterface {
   absl::StatusOr<bool> RunImpl(
       HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads) override;
-  absl::StatusOr<bool> RunImpl(
-      std::unique_ptr<HloModule>& module,
-      const absl::flat_hash_set<absl::string_view>& execution_threads) override;
 
  private:
   // Maybe dumps the given module depending on flag values contained in
@@ -113,31 +110,24 @@ class HloPassPipeline : public HloPassInterface {
   // Runs the invariant checker on the given HLO for specified
   // `execution_threads`. Empty `execution_threads` means all execution threads
   // are included.
-
-  template <typename HloT>
-  absl::Status RunInvariantCheckers(HloT hlo,
+  absl::Status RunInvariantCheckers(HloModule* hlo,
                                     absl::string_view after_pass_name) {
     return RunInvariantCheckers(hlo, after_pass_name, /*execution_threads=*/{});
   }
-  template <typename HloT>
   absl::Status RunInvariantCheckers(
-      HloT hlo, absl::string_view after_pass_name,
+      HloModule* hlo, absl::string_view after_pass_name,
       const absl::flat_hash_set<absl::string_view>& execution_threads);
 
   // Helper which runs the given pass on the given HLO.
-  template <typename HloT>
   absl::StatusOr<bool> RunPassesInternal(
-      HloT hlo, const DebugOptions& debug_options,
+      HloModule* hlo, const DebugOptions& debug_options,
       const absl::flat_hash_set<absl::string_view>& execution_threads);
 
-  // Helpers which run the given passes on the given HLO construct. Only
-  // computations with specified `execution_threads` are considered by the pass,
-  // empty thread list means all `execution_threads` are considered. These
-  // helpers enable templating of the core of the pipeline logic by providing
-  // HloModule specific methods with the same name.
-  template <typename HloT>
+  // Helper which runs the given pass on the given HLO module. Only
+  // computations with specified `execution_threads` are considered by the pass;
+  // an empty thread list means all `execution_threads` are considered.
   static absl::StatusOr<bool> RunHelper(
-      HloPassInterface* pass, HloT module,
+      HloPassInterface* pass, HloModule* module,
       const absl::flat_hash_set<absl::string_view>& execution_threads) {
     ABSL_ASSIGN_OR_RETURN(bool changed, pass->Run(module, execution_threads));
     module->Cleanup();

From 8329f1bc2bdffaecb9ae935aec723f89da67a50e Mon Sep 17 00:00:00 2001
From: Peter Hawkins <phawkins@google.com>
Date: Sat, 17 Oct 2026 03:15:31 -0700
Subject: [PATCH 1000/1000] Lower `lax.exp2_p` and `lax.log2_p` directly to
 `stablehlo.exp2` and `stablehlo.log2`

Update JAX primitives `lax.exp2_p` and `lax.log2_p` to lower directly to `stablehlo.exp2` and `stablehlo.log2` when not in forward-compatibility mode.

PiperOrigin-RevId: 1001403771
---
 .../hlo/analysis/stablehlo_indexing_analysis.cc  |  4 ++--
 .../hlo_to_mhlo/hlo_function_importer.cc         |  2 ++
 .../hlo/translate/hlo_to_mhlo/tests/import.hlo   | 16 ++++++++++++++++
 third_party/xla/xla/python/version.h             |  2 +-
 4 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/third_party/xla/xla/hlo/analysis/stablehlo_indexing_analysis.cc b/third_party/xla/xla/hlo/analysis/stablehlo_indexing_analysis.cc
index e0c3a6a5cbf..220114ab062 100644
--- a/third_party/xla/xla/hlo/analysis/stablehlo_indexing_analysis.cc
+++ b/third_party/xla/xla/hlo/analysis/stablehlo_indexing_analysis.cc
@@ -1011,8 +1011,8 @@ HloInstructionIndexing ComputeOutputToInputIndexing(Operation* op,
           // Elementwise identity operations, all operands use identity mapping.
           .Case<AddOp, SubtractOp, MulOp, DivOp, RemOp, MaxOp, MinOp, AndOp,
                 OrOp, XorOp, AbsOp, NegOp, SignOp, CosineOp, SineOp, TanhOp,
-                SqrtOp, RsqrtOp, ExpOp, Expm1Op, LogOp, Log1pOp, FloorOp,
-                CeilOp, ConvertOp, SelectOp, ClampOp, CompareOp,
+                SqrtOp, RsqrtOp, ExpOp, Exp2Op, Expm1Op, LogOp, Log1pOp, Log2Op,
+                FloorOp, CeilOp, ConvertOp, SelectOp, ClampOp, CompareOp,
                 PopulationCountOp, NotOp, IsFiniteOp, RoundNearestEvenOp,
                 OptimizationBarrierOp, MapOp, SortOp>([&](Operation* op) {
             if (!dyn_cast<RankedTensorType>(op->getResult(0).getType())) {
diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc
index 8537e6e6c82..290036173cc 100644
--- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc
+++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/hlo_function_importer.cc
@@ -2441,9 +2441,11 @@ absl::StatusOr<mlir::Operation*> HloFunctionImporter::ImportInstructionImpl(
       RESULT_ACCURACY_CASE(kCbrt, CbrtOp);
       RESULT_ACCURACY_CASE(kCos, CosineOp);
       RESULT_ACCURACY_CASE(kExp, ExpOp);
+      RESULT_ACCURACY_CASE(kExp2, Exp2Op);
       RESULT_ACCURACY_CASE(kExpm1, Expm1Op);
       RESULT_ACCURACY_CASE(kLog, LogOp);
       RESULT_ACCURACY_CASE(kLog1p, Log1pOp);
+      RESULT_ACCURACY_CASE(kLog2, Log2Op);
       RESULT_ACCURACY_CASE(kLogistic, LogisticOp);
       RESULT_ACCURACY_CASE(kRsqrt, RsqrtOp);
       RESULT_ACCURACY_CASE(kSin, SineOp);
diff --git a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import.hlo b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import.hlo
index 26fae8d6955..7ce74fa6509 100644
--- a/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import.hlo
+++ b/third_party/xla/xla/hlo/translate/hlo_to_mhlo/tests/import.hlo
@@ -946,6 +946,14 @@ add {
   ROOT %exp.2 = f32[16] exponential(f32[16] %arg0.1)
 }
 
+// CHECK-LABEL:  func private @test_exp2(%arg0: tensor<16xf32>) -> tensor<16xf32>
+%test_exp2 (arg0.1: f32[16]) -> f32[16] {
+  %arg0.1 = f32[16] parameter(0)
+
+  // CHECK-NEXT:  mhlo.exp2 %arg0 : tensor<16xf32>
+  ROOT %exp2.2 = f32[16] exp2(f32[16] %arg0.1)
+}
+
 // CHECK-LABEL:  func private @test_expm1(%arg0: tensor<16xf32>) -> tensor<16xf32>
 %test_expm1 (arg0.1: f32[16]) -> f32[16] {
   %arg0.1 = f32[16] parameter(0)
@@ -1104,6 +1112,14 @@ add {
   ROOT %log1p.2 = f32[16] log-plus-one(f32[16] %arg0.1)
 }
 
+// CHECK-LABEL:  func private @test_log2(%arg0: tensor<16xf32>) -> tensor<16xf32>
+%test_log2 (arg0.1: f32[16]) -> f32[16] {
+  %arg0.1 = f32[16] parameter(0)
+
+  // CHECK:  mhlo.log2 %arg0 : tensor<16xf32>
+  ROOT %log2.2 = f32[16] log2(f32[16] %arg0.1)
+}
+
 // Test mhlo.map
 %map_computation {
   lhs = f32[] parameter(0)
diff --git a/third_party/xla/xla/python/version.h b/third_party/xla/xla/python/version.h
index 578c87d484f..a8fc1a982e4 100644
--- a/third_party/xla/xla/python/version.h
+++ b/third_party/xla/xla/python/version.h
@@ -18,6 +18,6 @@ limitations under the License.
 
 // An increasing version number to protect jax code against breaking changes.
 // In JAX, reference this via jax._src.lib.ifrt_version.
-#define JAX_IFRT_VERSION_NUMBER 77  // Shardy unreduced axes support
+#define JAX_IFRT_VERSION_NUMBER 78  // exp2 and log2 StableHLO import
 
 #endif  // XLA_PYTHON_VERSION_H_